Process Network

This module provides functions for processing power networks and scenarios.

pf_preprocessing

Set variables to the results of OPF.

Copies the complete OPF operating point needed to initialize PF:

• gen.pg / gen.qg: active and reactive generator dispatch
• bus.vm / bus.va: bus voltage magnitude and angle
• gen.vg: generator voltage setpoints, synced to the OPF bus voltage so gen.vg == bus.vm in the written case and PowerModels does not emit voltage-setpoint mismatch warnings.

Parameters:

Name	Type	Description	Default
net	Network	The power network to preprocess.	required
res	Dict[str, Any]	OPF result dictionary containing solution data.	required

Returns:

Type	Description
Network	Updated network with OPF results applied.

Source code in gridfm_datakit/process/process_network.py

def pf_preprocessing(net: Network, res: Dict[str, Any]) -> Network:
    """Set variables to the results of OPF.

    Copies the complete OPF operating point needed to initialize PF:

    - gen.pg / gen.qg: active and reactive generator dispatch
    - bus.vm / bus.va: bus voltage magnitude and angle
    - gen.vg: generator voltage setpoints, synced to the OPF bus voltage so
      gen.vg == bus.vm in the written case and PowerModels does not emit
      voltage-setpoint mismatch warnings.

    Args:
        net: The power network to preprocess.
        res: OPF result dictionary containing solution data.

    Returns:
        Updated network with OPF results applied.
    """
    _, gen_pq, bus_vmva = _solution_arrays(res, net)

    net.Pg_gen = gen_pq[:, 0] * net.baseMVA
    net.Qg_gen = gen_pq[:, 1] * net.baseMVA
    net.Vm = bus_vmva[:, 0]
    # PowerModels reports radians; MATPOWER stores degrees.
    net.Va = np.rad2deg(bus_vmva[:, 1])

    # Sync each in-service generator's voltage setpoint (VG) to the OPF bus
    # voltage at its terminal bus, so gen.vg == bus.vm in the written .m file.
    for gen_idx in net.idx_gens_in_service:
        bus_idx = int(net.gens[gen_idx, GEN_BUS])
        net.gens[gen_idx, VG] = bus_vmva[bus_idx, 0]

    return net

pf_post_processing

Post-process solved network results into numpy arrays for CSV export.

This function extracts power flow results and builds four arrays matching the column schemas defined in gridfm_datakit.utils.column_names:

• Bus data with BUS_COLUMNS (+ DC_BUS_COLUMNS if include_dc_res=True)
• Generator data with GEN_COLUMNS
• Branch data with BRANCH_COLUMNS
• Y-bus nonzero entries with [index1, index2, G, B]

Parameters:

Name	Type	Description	Default
net	Network	The power network to process (must have solved power flow results).	required
res	Dict[str, Any]	Power flow result dictionary containing solution data.	required
include_dc_res	bool	If True, include DC power flow voltage magnitude/angle (Vm_dc, Va_dc).	required

Returns:

Type	Description
Dict[str, ndarray]	Dictionary containing:
Dict[str, ndarray]	"bus": np.ndarray with bus-level features
Dict[str, ndarray]	"gen": np.ndarray with generator features
Dict[str, ndarray]	"branch": np.ndarray with branch features and admittances
Dict[str, ndarray]	"Y_bus": np.ndarray with nonzero Y-bus entries

Source code in gridfm_datakit/process/process_network.py

def pf_post_processing(
    scenario_index: int,
    net: Network,
    res: Dict[str, Any],
    res_dc: Dict[str, Any],
    include_dc_res: bool,
) -> Dict[str, np.ndarray]:
    """Post-process solved network results into numpy arrays for CSV export.

    This function extracts power flow results and builds four arrays matching
    the column schemas defined in `gridfm_datakit.utils.column_names`:

    - Bus data with BUS_COLUMNS (+ DC_BUS_COLUMNS if include_dc_res=True)
    - Generator data with GEN_COLUMNS
    - Branch data with BRANCH_COLUMNS
    - Y-bus nonzero entries with [index1, index2, G, B]

    Args:
        net: The power network to process (must have solved power flow results).
        res: Power flow result dictionary containing solution data.
        include_dc_res: If True, include DC power flow voltage magnitude/angle (Vm_dc, Va_dc).

    Returns:
        Dictionary containing:
        - "bus": np.ndarray with bus-level features
        - "gen": np.ndarray with generator features
        - "branch": np.ndarray with branch features and admittances
        - "Y_bus": np.ndarray with nonzero Y-bus entries
    """

    # --- Edge (branch) info ---
    n_branches = net.branches.shape[0]
    n_cols = (
        len(BRANCH_COLUMNS) + len(DC_BRANCH_COLUMNS)
        if include_dc_res
        else len(BRANCH_COLUMNS)
    )
    X_branch = np.zeros((n_branches, n_cols))
    X_branch[:, 0] = scenario_index
    X_branch[:, 1] = np.arange(n_branches)
    X_branch[:, 2] = np.real(net.branches[:, F_BUS])
    X_branch[:, 3] = np.real(net.branches[:, T_BUS])

    # pf, qf, pt, qt
    if res["solution"]["pf"]:
        # when solving pf, the flow of all branches is computed, so the number of branches in solution should match the number of branches in network
        assert len(res["solution"]["branch"]) == n_branches, (
            "Number of branches in solution should match number of branches in network"
        )
    else:
        # when solving opf, the flow of only the in-service branches is computed, so the number of branches in solution should match the number of in-service branches in network
        assert len(res["solution"]["branch"]) == len(net.idx_branches_in_service), (
            "Number of branches in solution should match number of branches in network"
        )

    branch_flows, gen_pq, bus_vmva = _solution_arrays(res, net)
    if len(net.idx_branches_in_service) > 0:
        X_branch[net.idx_branches_in_service, 4:8] = branch_flows * net.baseMVA

    X_branch[:, 8] = net.branches[:, BR_R]
    X_branch[:, 9] = net.branches[:, BR_X]
    X_branch[:, 10] = net.branches[:, BR_B]

    # admittances
    Ytt, Yff, Yft, Ytf = branch_vectors(net.branches, net.branches.shape[0])
    X_branch[:, 11] = np.real(Yff)
    X_branch[:, 12] = np.imag(Yff)
    X_branch[:, 13] = np.real(Yft)
    X_branch[:, 14] = np.imag(Yft)
    X_branch[:, 15] = np.real(Ytf)
    X_branch[:, 16] = np.imag(Ytf)
    X_branch[:, 17] = np.real(Ytt)
    X_branch[:, 18] = np.imag(Ytt)

    X_branch[:, 19] = net.branches[:, TAP]
    # assign 1 to tap = 0
    X_branch[net.branches[:, TAP] == 0, 19] = 1

    X_branch[:, 20] = net.branches[:, SHIFT]
    X_branch[:, 21] = net.branches[:, ANGMIN]
    X_branch[:, 22] = net.branches[:, ANGMAX]
    X_branch[:, 23] = net.branches[:, RATE_A]
    X_branch[:, 24] = net.branches[:, BR_STATUS]

    if include_dc_res:
        if res_dc is not None:
            dc_flows, gen_pq_dc, bus_vmva_dc = _solution_arrays(res_dc, net)
            dc_flows = dc_flows * net.baseMVA
            pf_dc = dc_flows[:, 0]
            pt_dc = dc_flows[:, 2]
            X_branch[net.idx_branches_in_service, 25] = pf_dc
            X_branch[net.idx_branches_in_service, 26] = pt_dc
        else:
            X_branch[net.idx_branches_in_service, 25] = np.nan
            X_branch[net.idx_branches_in_service, 26] = np.nan

    # --- Bus data ---
    n_buses = net.buses.shape[0]
    n_cols = (
        len(BUS_COLUMNS) + len(DC_BUS_COLUMNS) if include_dc_res else len(BUS_COLUMNS)
    )
    X_bus = np.zeros((n_buses, n_cols))

    # --- Loads ---
    X_bus[:, 0] = scenario_index
    X_bus[:, 1] = net.buses[:, BUS_I]  # bus
    X_bus[:, 2] = net.buses[:, PD]
    X_bus[:, 3] = net.buses[:, QD]

    # --- Generator injections
    assert len(res["solution"]["gen"]) == len(net.idx_gens_in_service), (
        "Number of generators in solution should match number of generators in network"
    )
    gen_pq = gen_pq * net.baseMVA
    pg_gen = gen_pq[:, 0]
    qg_gen = gen_pq[:, 1]
    gen_bus = net.gens[net.idx_gens_in_service, GEN_BUS].astype(int)
    Pg_bus = np.bincount(gen_bus, weights=pg_gen, minlength=n_buses)
    Qg_bus = np.bincount(gen_bus, weights=qg_gen, minlength=n_buses)

    # Pg_bus / Qg_bus are indexed by 0-based bus INDEX; net.buses rows may be
    # in a different order (pypowsybl does not guarantee sorted bus export).
    # Build a bus-type array indexed by bus index so the mask aligns correctly.
    bus_type_by_idx = np.zeros(n_buses)
    bus_type_by_idx[net.buses[:, BUS_I].astype(int)] = net.buses[:, BUS_TYPE]

    assert np.all(Pg_bus[bus_type_by_idx == PQ] == 0)
    assert np.all(Qg_bus[bus_type_by_idx == PQ] == 0)

    if include_dc_res:
        if res_dc is not None:
            # check if "gen" key is in res_dc["solution"]
            if "gen" in res_dc["solution"]:
                pg_gen_dc = gen_pq_dc[:, 0] * net.baseMVA
            else:
                pg_gen_dc = apply_slack_single_gen(net, pg_gen, Pg_bus, pf_dc, pt_dc)
            Pg_bus_dc = np.bincount(gen_bus, weights=pg_gen_dc, minlength=n_buses)
            assert np.all(Pg_bus_dc[bus_type_by_idx == PQ] == 0)

    # Reindex Pg/Qg from bus-index order to bus-row order for X_bus assignment.
    bus_row_idx = net.buses[:, BUS_I].astype(int)
    X_bus[:, 4] = Pg_bus[bus_row_idx]
    X_bus[:, 5] = Qg_bus[bus_row_idx]

    # Voltage. Extraction (_solution_arrays) raises on any missing expected
    # bus key, so together with this length check the key sets must match.
    assert len(res["solution"]["bus"]) == n_buses, (
        "Buses in solution should match buses in network"
    )

    X_bus[:, 6] = bus_vmva[:, 0]
    va = np.rad2deg(bus_vmva[:, 1])

    # convert to range [-180, 180]
    va = (va + 180) % 360 - 180
    X_bus[:, 7] = va

    # one-hot encoding of bus type
    assert np.all(np.isin(net.buses[:, BUS_TYPE], [PQ, PV, REF])), (
        "Bus type should be PQ, PV, or REF, no disconnected buses (4)"
    )

    X_bus[np.arange(n_buses), 8 + net.buses[:, BUS_TYPE].astype(int) - 1] = (
        1  # because type is 1, 2, 3, not 0, 1, 2
    )

    # base_kv, min_vm_pu, max_vm_pu
    X_bus[:, 11] = net.buses[:, BASE_KV]
    X_bus[:, 12] = net.buses[:, VMIN]
    X_bus[:, 13] = net.buses[:, VMAX]

    X_bus[:, 14] = net.buses[:, GS] / net.baseMVA
    X_bus[:, 15] = net.buses[:, BS] / net.baseMVA

    if include_dc_res:
        if res_dc is not None:
            va = np.rad2deg(bus_vmva_dc[:, 1])
            # convert to range [-180, 180]
            va = (va + 180) % 360 - 180
            X_bus[:, 16] = va
            X_bus[:, 17] = Pg_bus_dc[bus_row_idx]
        else:
            X_bus[:, 16] = np.nan
            X_bus[:, 17] = np.nan

    # --- Generator data ---

    n_cost = net.gencosts[0, NCOST]
    assert np.all(net.gencosts[:, NCOST] == n_cost), (
        "NCOST should be the same for all generators"
    )
    n_gens = net.gens.shape[0]
    n_cols = (
        len(GEN_COLUMNS) + len(DC_GEN_COLUMNS) if include_dc_res else len(GEN_COLUMNS)
    )

    X_gen = np.zeros((n_gens, n_cols))
    X_gen[:, 0] = scenario_index
    X_gen[:, 1] = np.arange(n_gens)
    X_gen[:, 2] = net.gens[:, GEN_BUS]
    X_gen[net.idx_gens_in_service, 3] = pg_gen  # 0 if not in service
    X_gen[net.idx_gens_in_service, 4] = qg_gen  # 0 if not in service
    X_gen[:, 5] = net.gens[:, PMIN]
    X_gen[:, 6] = net.gens[:, PMAX]
    X_gen[:, 7] = net.gens[:, QMIN]
    X_gen[:, 8] = net.gens[:, QMAX]

    if n_cost == 3:  # order in .m file is c2, c1, c0
        X_gen[:, 9] = net.gencosts[:, COST + 2]
        X_gen[:, 10] = net.gencosts[:, COST + 1]
        X_gen[:, 11] = net.gencosts[:, COST]

    if n_cost == 2:  # order in .m file is c1, c0, and there is no cp2 cost
        X_gen[:, 9] = net.gencosts[:, COST + 1]
        X_gen[:, 10] = net.gencosts[:, COST]
        X_gen[:, 11] = 0  # no cp2 cost for linear cost function

    if n_cost == 1:  # order in .m file is c0, and there is no cp1 or cp2 cost
        X_gen[:, 9] = net.gencosts[:, COST]
        X_gen[:, 10] = 0  # no cp1 cost for constant cost function
        X_gen[:, 11] = 0  # no cp2 cost for constant cost function

    X_gen[net.idx_gens_in_service, 12] = 1

    # slack gen (can be any generator connected to the ref node)
    slack_gen_idx = np.where(net.gens[:, GEN_BUS] == net.ref_bus_idx)[0]
    X_gen[slack_gen_idx, 13] = 1

    if include_dc_res:
        if res_dc is not None:
            X_gen[net.idx_gens_in_service, 14] = pg_gen_dc
        else:
            X_gen[net.idx_gens_in_service, 14] = np.nan

    # --- Y-bus ---
    Y_bus, Yf, Yt = makeYbus(net.baseMVA, net.buses, net.branches)

    i, j = np.nonzero(Y_bus)
    # note that Y_bus[i,j] can be != 0 even if a branch from i to j is not in service because there might be other branches connected to the same buses

    s = Y_bus[i, j]
    G = np.real(s)
    B = np.imag(s)

    edge_index = np.column_stack((i, j))
    edge_attr = np.stack((G, B)).T
    Y_bus = np.zeros(
        (edge_index.shape[0], edge_attr.shape[1] + edge_index.shape[1] + 1),
    )
    Y_bus[:, 0] = scenario_index
    Y_bus[:, 1:] = np.column_stack((edge_index, edge_attr))

    # ---- runtime data ----
    n_cols = (
        len(RUNTIME_COLUMNS) + len(DC_RUNTIME_COLUMNS)
        if include_dc_res
        else len(RUNTIME_COLUMNS)
    )
    X_runtime = np.zeros((1, n_cols))
    X_runtime[0, 0] = scenario_index
    X_runtime[0, 1] = res["solve_time"]
    if include_dc_res:
        if res_dc is not None:
            X_runtime[0, 2] = res_dc["solve_time"]
        else:
            X_runtime[0, 2] = np.nan
    return {
        "bus": X_bus,
        "gen": X_gen,
        "branch": X_branch,
        "Y_bus": Y_bus,
        "runtime": X_runtime,
    }

process_scenario_opf_mode

Processes a load scenario in OPF mode

In OPF mode, perturbations are applied first, then OPF is run to get generator setpoints that account for the perturbed topology. This ensures all constraints are satisfied in the final operating point.

Parameters:

Name	Type	Description	Default
net	Network	The power network.	required
scenarios	ndarray	Array of load scenarios with shape (n_loads, n_scenarios, 2).	required
scenario_index	int	Index of the current scenario to process.	required
topology_generator	TopologyGenerator	Generator for topology perturbations (line/transformer outages).	required
generation_generator	GenerationGenerator	Generator for generation cost perturbations.	required
admittance_generator	AdmittanceGenerator	Generator for line admittance perturbations.	required
local_processed_data	List[ndarray]	List to accumulate processed data tuples.	required
error_log_file	str	Path to error log file for recording failures.	required
include_dc_res	bool	Whether to include DC power flow results in output.	required
jl	Any	Julia interface object for running power flow calculations.	required
scenario_data_index	Optional[int]	Optional local index into a sliced scenarios tensor. Defaults to scenario_index.	None

Returns:

Type	Description
List[ndarray]	Updated list of processed data (bus, gen, branch, Y_bus arrays)

Note

Random seed is controlled by the calling context (process_scenario_chunk).

Source code in gridfm_datakit/process/process_network.py

def process_scenario_opf_mode(
    net: Network,
    scenarios: np.ndarray,
    scenario_index: int,
    topology_generator: TopologyGenerator,
    generation_generator: GenerationGenerator,
    admittance_generator: AdmittanceGenerator,
    local_processed_data: List[np.ndarray],
    error_log_file: str,
    include_dc_res: bool,
    jl: Any,
    *,
    scenario_data_index: Optional[int] = None,
) -> List[np.ndarray]:
    """Processes a load scenario in OPF mode

    In OPF mode, perturbations are applied first, then OPF is run to get
    generator setpoints that account for the perturbed topology. This ensures
    all constraints are satisfied in the final operating point.

    Args:
        net: The power network.
        scenarios: Array of load scenarios with shape (n_loads, n_scenarios, 2).
        scenario_index: Index of the current scenario to process.
        topology_generator: Generator for topology perturbations (line/transformer outages).
        generation_generator: Generator for generation cost perturbations.
        admittance_generator: Generator for line admittance perturbations.
        local_processed_data: List to accumulate processed data tuples.
        error_log_file: Path to error log file for recording failures.
        include_dc_res: Whether to include DC power flow results in output.
        jl: Julia interface object for running power flow calculations.
        scenario_data_index: Optional local index into a sliced ``scenarios``
            tensor. Defaults to ``scenario_index``.

    Returns:
        Updated list of processed data (bus, gen, branch, Y_bus arrays)

    Note:
        Random seed is controlled by the calling context (process_scenario_chunk).
    """

    # apply the load scenario to the network
    data_index = scenario_index if scenario_data_index is None else scenario_data_index
    net.Pd = scenarios[:, data_index, 0]
    net.Qd = scenarios[:, data_index, 1]

    # Generate perturbed topologies
    perturbations = topology_generator.generate(net)

    # Apply generation perturbations
    perturbations = generation_generator.generate(perturbations)

    # Apply admittance perturbations
    perturbations = admittance_generator.generate(perturbations)

    for perturbation in (
        perturbations
    ):  # (that returns copies of the network with the topology perturbation applied)
        res_dcopf = None
        if include_dc_res:
            try:
                res_dcopf = run_dcopf(perturbation, jl)
            except Exception as e:
                with open(error_log_file, "a") as f:
                    f.write(
                        f"Caught an exception at scenario {scenario_index} in run_dcopf function: {e}\n",
                    )
        try:
            # run OPF to get the gen set points. Here the set points account for the topology perturbation.
            res = run_opf(perturbation, jl)
        except Exception as e:
            with open(error_log_file, "a") as f:
                f.write(
                    f"Caught an exception at scenario {scenario_index} in run_opf function: {e}\n",
                )
            continue

        # Append processed power flow data
        pf_data = pf_post_processing(
            scenario_index,
            perturbation,
            res,
            res_dcopf,
            include_dc_res,
        )
        local_processed_data.append(
            (
                pf_data["bus"],
                pf_data["gen"],
                pf_data["branch"],
                pf_data["Y_bus"],
                pf_data["runtime"],
            ),
        )
    return local_processed_data

process_scenario_pf_mode

Processes a load scenario in PF mode.

In PF mode, OPF is run first to get generator setpoints, then topology perturbations are applied. This can lead to constraint violations (overloads, voltage violations) since the setpoints are not re-optimized for the new topology.

Parameters

net: The base power network (copied internally before mutation). scenarios: Array of load scenarios with shape (n_loads, n_scenarios, 2). scenario_index: Global index of the current scenario to process and write. topology_generator: Generator for topology perturbations (line/transformer outages). generation_generator: Generator for generation cost perturbations. admittance_generator: Generator for line admittance perturbations. local_processed_data: List to accumulate processed data tuples. error_log_file: Path to error log file for recording failures. include_dc_res: Whether to include DC power flow results in output. pf_fast: Whether to use the fast AC PF solver (compute_ac_pf from PowerModels.jl). Only consulted when pf_solver='powermodel'. dcpf_fast: Whether to use the fast DC PF solver (compute_dc_pf from PowerModels.jl). Only consulted when pf_solver='powermodel'. jl: Julia interface object. Always required — even when pf_solver='powsybl' Julia is used for the OPF step that produces the generator set-points before topology perturbation. pf_solver: Which engine to use for the power flow solve after topology perturbation. Must be 'powermodel' (default) or 'powsybl'. OPF is always solved by PowerModels regardless of this value.

Keyword-only arguments (only required when pf_solver='powsybl')

meta: Optional dictionary containing metadata for PowSyBl processing, with keys: - pp_net: the PowSyBl network. - mapping_p2g: dictionary mapping from PowSyBl to GFM. scenario_data_index: Optional local index into a sliced scenarios tensor. Defaults to scenario_index for callers that pass the full tensor.

Returns

List[np.ndarray] Updated local_processed_data list with one tuple (bus, gen, branch, Y_bus, runtime) appended per successfully solved perturbation.

Note

Random seed is controlled by the calling context (process_scenario_chunk or generate_power_flow_data).

Source code in gridfm_datakit/process/process_network.py

def process_scenario_pf_mode(
    net: Network,
    scenarios: np.ndarray,
    scenario_index: int,
    topology_generator: TopologyGenerator,
    generation_generator: GenerationGenerator,
    admittance_generator: AdmittanceGenerator,
    local_processed_data: List[np.ndarray],
    error_log_file: str,
    include_dc_res: bool,
    pf_fast: bool,
    dcpf_fast: bool,
    jl: Any,
    pf_solver: str = "powermodel",
    *,
    meta: Optional[Dict] = None,
    scenario_data_index: Optional[int] = None,
) -> List[np.ndarray]:
    """Processes a load scenario in PF mode.

    In PF mode, OPF is run first to get generator setpoints, then topology
    perturbations are applied. This can lead to constraint violations (overloads,
    voltage violations) since the setpoints are not re-optimized for the new topology.

    Parameters
    ----------
    net:
        The base power network (copied internally before mutation).
    scenarios:
        Array of load scenarios with shape ``(n_loads, n_scenarios, 2)``.
    scenario_index:
        Global index of the current scenario to process and write.
    topology_generator:
        Generator for topology perturbations (line/transformer outages).
    generation_generator:
        Generator for generation cost perturbations.
    admittance_generator:
        Generator for line admittance perturbations.
    local_processed_data:
        List to accumulate processed data tuples.
    error_log_file:
        Path to error log file for recording failures.
    include_dc_res:
        Whether to include DC power flow results in output.
    pf_fast:
        Whether to use the fast AC PF solver (``compute_ac_pf`` from
        PowerModels.jl).  Only consulted when ``pf_solver='powermodel'``.
    dcpf_fast:
        Whether to use the fast DC PF solver (``compute_dc_pf`` from
        PowerModels.jl).  Only consulted when ``pf_solver='powermodel'``.
    jl:
        Julia interface object.  Always required — even when
        ``pf_solver='powsybl'`` Julia is used for the OPF step that
        produces the generator set-points before topology perturbation.
    pf_solver:
        Which engine to use for the power flow solve after topology
        perturbation.  Must be ``'powermodel'`` (default) or
        ``'powsybl'``.  OPF is always solved by PowerModels regardless
        of this value.

    Keyword-only arguments (only required when ``pf_solver='powsybl'``)
    -------------------------------------------------------------------
    meta:
    Optional dictionary containing metadata for PowSyBl processing, with keys:
        - pp_net: the PowSyBl network.
        - mapping_p2g: dictionary mapping from PowSyBl to GFM.
    scenario_data_index:
        Optional local index into a sliced ``scenarios`` tensor. Defaults to
        ``scenario_index`` for callers that pass the full tensor.

    Returns
    -------
    List[np.ndarray]
        Updated ``local_processed_data`` list with one tuple
        ``(bus, gen, branch, Y_bus, runtime)`` appended per successfully
        solved perturbation.

    Note
    ----
    Random seed is controlled by the calling context
    (``process_scenario_chunk`` or ``generate_power_flow_data``).
    """
    net = net.copy_for_perturbation()

    # apply the load scenario to the network
    data_index = scenario_index if scenario_data_index is None else scenario_data_index
    net.Pd = scenarios[:, data_index, 0]
    net.Qd = scenarios[:, data_index, 1]

    # Apply generation perturbations before OPF.
    perturbations = generation_generator.generate((x for x in [net]))

    # Apply admittance perturbations
    perturbations = admittance_generator.generate(perturbations)

    net = next(perturbations)

    # first run OPF to get the gen set points
    try:
        res = run_opf(net, jl)
    except Exception as e:
        with open(error_log_file, "a") as f:
            f.write(
                f"Caught an exception at scenario {scenario_index} in run_opf function: {e}\n",
            )
        return local_processed_data

    net_pf = pf_preprocessing(net, res)

    # Generate perturbed topologies
    perturbations = topology_generator.generate(net_pf)

    if pf_solver == "powsybl":
        powsybl.check_powsybl_available()
        if meta is None or "pp_net" not in meta or "mapping_p2g" not in meta:
            raise ValueError("Network seems to not be initialized for PowSyBl solver")
        pp_net = meta["pp_net"]
        mapping_p2g = meta["mapping_p2g"]
        base_variant_id = pp_net.get_working_variant_id()
        lf_params = powsybl.get_default_lf_params()

    # to get PF points that can violate some OPF inequality constraints (to train PF solvers that can handle points outside of normal operating limits), we apply the topology perturbation after OPF.
    # The setpoints are then no longer adapted to the new topology, and might lead to e.g. abranch overload or a voltage magnitude violation once we drop an element.
    for pert_index, perturbation in enumerate(perturbations):
        if pf_solver == "powermodel":
            res_dcpf = None
            if include_dc_res:
                try:
                    res_dcpf = run_dcpf(perturbation, jl, fast=dcpf_fast)

                except Exception as e:
                    with open(error_log_file, "a") as f:
                        f.write(
                            f"Caught an exception at scenario {scenario_index} when solving dcpf function: {e}\n",
                        )
            try:
                res = run_pf(perturbation, jl, fast=pf_fast)
            except Exception as e:
                with open(error_log_file, "a") as f:
                    f.write(
                        f"Caught an exception at scenario {scenario_index} when solving in run_pf function: {e}\n",
                    )
                continue

        if pf_solver == "powsybl":
            variant_id = f"scenario_{scenario_index}_perturbation_{pert_index}"
            pp_net.clone_variant(base_variant_id, variant_id)
            pp_net.set_working_variant(variant_id)
            try:
                powsybl.update_powsybl(pp_net, perturbation, mapping_p2g)

                res_dcpf = None
                if include_dc_res:
                    try:
                        start_time = time.perf_counter()
                        dcpf_metadata = powsybl.pypowsybl.loadflow.run_dc(
                            pp_net,
                            lf_params,
                        )
                        end_time = time.perf_counter()
                        solve_time = end_time - start_time
                        res_dcpf = powsybl.get_pf_res(
                            pp_net,
                            solve_time,
                            dcpf_metadata,
                            mapping_p2g,
                        )

                    except Exception as e:
                        with open(error_log_file, "a") as f:
                            f.write(
                                f"Caught an exception at scenario {scenario_index} when solving dcpf function with PowSyBl solver: {e}\n",
                            )
                try:
                    start_time = time.perf_counter()
                    pf_metadata = powsybl.pypowsybl.loadflow.run_ac(pp_net, lf_params)
                    end_time = time.perf_counter()
                    solve_time = end_time - start_time
                    res = powsybl.get_pf_res(
                        pp_net,
                        solve_time,
                        pf_metadata,
                        mapping_p2g,
                    )
                except Exception as e:
                    with open(error_log_file, "a") as f:
                        f.write(
                            f"Caught an exception at scenario {scenario_index} when solving in run_pf function with PowSyBl solver: {e}\n",
                        )
                    continue
            finally:
                pp_net.set_working_variant(base_variant_id)
                pp_net.remove_variant(variant_id)

        # Append processed power flow data
        pf_data = pf_post_processing(
            scenario_index,
            perturbation,
            res,
            res_dcpf,
            include_dc_res,
        )
        local_processed_data.append(
            (
                pf_data["bus"],
                pf_data["gen"],
                pf_data["branch"],
                pf_data["Y_bus"],
                pf_data["runtime"],
            ),
        )
    return local_processed_data

process_scenario_chunk

Process a chunk of scenarios for distributed processing.

This function processes multiple scenarios in a single worker process, accumulating results before returning them to the main process.

Parameters:

Name	Type	Description	Default
mode	str	Processing mode ("opf" or "pf").	required
start_idx	int	Starting scenario index (inclusive).	required
end_idx	int	Ending scenario index (exclusive).	required
scenarios	ndarray	Array of load scenarios with shape (n_loads, n_scenarios, 2).	required
net	Network	The power network.	required
progress_queue	Optional[Any]	Queue for reporting progress to main process.	required
topology_generator	TopologyGenerator	Generator for topology perturbations.	required
generation_generator	GenerationGenerator	Generator for generation cost perturbations.	required
admittance_generator	AdmittanceGenerator	Generator for line admittance perturbations.	required
error_log_path	str	Path to error log file for recording failures.	required
include_dc_res	bool	Whether to include DC power flow results in output.	required
pf_fast	bool	Whether to use fast AC PF solver.	required
dcpf_fast	bool	Whether to use fast DC PF solver.	required
solver_log_dir	str	Directory for solver logs.	required
max_iter	int	Maximum iterations for the solver.	required
seed	int	Global random seed for reproducibility.	required
pf_solver	str	PF solver to use in pf mode; either 'powermodel' or 'powsybl'. OPF is always solved by PowerModels regardless of this value.	'powermodel'
meta	Optional[Dict]	metadata dict; when pf_solver='powsybl', must contain 'network_path' and 'mapping_p2g'. 'pp_net' is loaded fresh per worker from 'network_path'.	None
scenario_index_offset	int	Global index added to local indices when scenarios contains only the current chunk.	0

Returns:

Type	Description
Tuple[Union[None, Exception], Union[None, str], Optional[List[ndarray]]]	Tuple containing: - Exception object (None if successful) - Traceback string (None if successful) - List of processed data tuples (bus, gen, branch, Y_bus arrays)

Source code in gridfm_datakit/process/process_network.py

def process_scenario_chunk(
    mode: str,
    start_idx: int,
    end_idx: int,
    scenarios: np.ndarray,
    net: Network,
    progress_queue: Optional[Any],
    topology_generator: TopologyGenerator,
    generation_generator: GenerationGenerator,
    admittance_generator: AdmittanceGenerator,
    error_log_path: str,
    include_dc_res: bool,
    pf_fast: bool,
    dcpf_fast: bool,
    solver_log_dir: str,
    max_iter: int,
    seed: int,
    pf_solver: str = "powermodel",
    meta: Optional[Dict] = None,
    *,
    scenario_index_offset: int = 0,
    opf_formulation: str = "polar",
) -> Tuple[
    Union[None, Exception],
    Union[None, str],
    Optional[List[np.ndarray]],
]:
    """Process a chunk of scenarios for distributed processing.

    This function processes multiple scenarios in a single worker process,
    accumulating results before returning them to the main process.

    Args:
        mode: Processing mode ("opf" or "pf").
        start_idx: Starting scenario index (inclusive).
        end_idx: Ending scenario index (exclusive).
        scenarios: Array of load scenarios with shape (n_loads, n_scenarios, 2).
        net: The power network.
        progress_queue: Queue for reporting progress to main process.
        topology_generator: Generator for topology perturbations.
        generation_generator: Generator for generation cost perturbations.
        admittance_generator: Generator for line admittance perturbations.
        error_log_path: Path to error log file for recording failures.
        include_dc_res: Whether to include DC power flow results in output.
        pf_fast: Whether to use fast AC PF solver.
        dcpf_fast: Whether to use fast DC PF solver.
        solver_log_dir: Directory for solver logs.
        max_iter: Maximum iterations for the solver.
        seed: Global random seed for reproducibility.
        pf_solver: PF solver to use in pf mode; either 'powermodel' or 'powsybl'.
            OPF is always solved by PowerModels regardless of this value.
        meta: metadata dict; when pf_solver='powsybl', must contain 'network_path'
            and 'mapping_p2g'. 'pp_net' is loaded fresh per worker from 'network_path'.
        scenario_index_offset: Global index added to local indices when ``scenarios``
            contains only the current chunk.

    Returns:
        Tuple containing:
            - Exception object (None if successful)
            - Traceback string (None if successful)
            - List of processed data tuples (bus, gen, branch, Y_bus arrays)
    """

    global _worker_jl
    completed_scenarios = 0
    try:
        if _worker_jl is None:
            _worker_jl = init_julia(
                max_iter,
                solver_log_dir,
                opf_formulation=opf_formulation,
            )
        jl = _worker_jl

        # In distributed (spawn) workers pp_net is not passed; reload it here.
        if (
            pf_solver == "powsybl"
            and meta
            and "network_path" in meta
            and "pp_net" not in meta
        ):
            import gridfm_datakit.powsybl as _powsybl

            loaded_net = _powsybl.load_net(meta["network_path"])
            meta["pp_net"] = loaded_net.pp_net

        local_processed_data = []

        global_start_idx = start_idx + scenario_index_offset

        # Use custom_seed to set seed based on the global start index for this chunk
        # This ensures each chunk gets a unique but deterministic seed
        # we multiply by 20_000 to ensure there is no collision with other runs where the seed would be close to each other
        # example (assuming we have chunks of length 1, hence an increment of 1 between start indices)
        # Run A: base seed = 42 → scenario seeds = 42, 43, 44, …, 10041 (for 10,000 scenarios)
        # Run B: base seed = 120 → scenario seeds = 120, 121, 122, …, 10119
        # These sets overlap on seeds 120..10041 (so 9,922 overlapping seeds).
        # we also add 1 in case the seed is 0, to not have collision witht he seed used for the load perturbations
        with custom_seed(seed * 20_000 + global_start_idx + 1):
            for scenario_data_index in range(start_idx, end_idx):
                scenario_index = scenario_data_index + scenario_index_offset
                if mode == "opf":
                    local_processed_data = process_scenario_opf_mode(
                        net,
                        scenarios,
                        scenario_index,
                        topology_generator,
                        generation_generator,
                        admittance_generator,
                        local_processed_data,
                        error_log_path,
                        include_dc_res,
                        jl,
                        scenario_data_index=scenario_data_index,
                    )
                elif mode == "pf":
                    local_processed_data = process_scenario_pf_mode(
                        net,
                        scenarios,
                        scenario_index,
                        topology_generator,
                        generation_generator,
                        admittance_generator,
                        local_processed_data,
                        error_log_path,
                        include_dc_res,
                        pf_fast,
                        dcpf_fast,
                        jl,
                        pf_solver,
                        meta=meta,
                        scenario_data_index=scenario_data_index,
                    )
                else:
                    raise ValueError(f"Invalid mode: {mode!r}")

                completed_scenarios += 1
                if progress_queue is not None:
                    progress_queue.put(1)

        return (
            None,
            None,
            local_processed_data,
        )
    except Exception as e:
        with open(error_log_path, "a") as f:
            f.write(f"Caught an exception in process_scenario_chunk function: {e}\n")
            f.write(traceback.format_exc())
            f.write("\n")
        if progress_queue is not None:
            for _ in range(end_idx - start_idx - completed_scenarios):
                progress_queue.put(1)
        return e, traceback.format_exc(), None