Topology Optimizer

toop_engine_topology_optimizer.interfaces.messages.commands

Defines the interaction between the optimizer and a backend.

The backend will send commands in the form of messages to the optimizer, which will trigger a certain behaviour. The Optimizer will respond with results, for this see results.py

StartOptimizationCommand

Bases: BaseModel

Command with parameters for starting an optimization run.

message_type class-attribute instance-attribute

message_type = 'start_optimization'

The command type for deserialization, don't change this

dc_params class-attribute instance-attribute

dc_params = DCOptimizerParameters()

The parameters for the DC optimizer

ac_params class-attribute instance-attribute

ac_params = ACOptimizerParameters()

The parameters for the AC optimizer

grid_files instance-attribute

grid_files

The grid files to load, where each gridfile represents one timestep. The grid files also include coupling information for the timesteps.

optimization_id instance-attribute

optimization_id

The id of the optimization run, used to identify the optimization run in the results. Should stay the same for the whole optimization run and should be equal to the kafka event key

ShutdownCommand

Bases: BaseModel

Command to shutdown the worker.

message_type class-attribute instance-attribute

message_type = 'shutdown'

The command type for deserialization, don't change this

exit_code class-attribute instance-attribute

exit_code = 0

The exit code to exit with

Command

Bases: BaseModel

Base class for all commands to the optimizer.

command class-attribute instance-attribute

command = Field(discriminator='message_type')

The actual command

uuid class-attribute instance-attribute

uuid = Field(default_factory=lambda: str(uuid4()))

A unique identifier for the command message, used to avoid duplicate processing on optimizer side

timestamp class-attribute instance-attribute

timestamp = Field(default_factory=lambda: str(now()))

When the command was sent

validate_first_gridfile_uncoupled

validate_first_gridfile_uncoupled(val)

Check that the first gridfile is uncoupled

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/interfaces/messages/commands.py

def validate_first_gridfile_uncoupled(val: list[GridFile]) -> list[GridFile]:
    """Check that the first gridfile is uncoupled"""
    if val:
        if val[0].coupling != "none":
            raise ValueError("The first gridfile must be uncoupled")
    return val

toop_engine_topology_optimizer.interfaces.messages.dc_params

The DC optimizer is the GPU stage where massive amounts of topologies are being checked.

This holds the parameters to start the optimization. Some parameters can not be changed (mainly the names of the kafka streams) and are included in the command line start parameters instead.

BatchedMEParameters

Bases: BaseModel

Parameters for starting the batched genetic algorithm (In this case Map-Elites)

runtime_seconds class-attribute instance-attribute

runtime_seconds = 60

The runtime in seconds

iterations_per_epoch class-attribute instance-attribute

iterations_per_epoch = 100

The number of iterations per epoch

random_seed class-attribute instance-attribute

random_seed = 42

The random seed to use for reproducibility

enable_nodal_inj_optim class-attribute instance-attribute

enable_nodal_inj_optim = False

Whether to enable the nodal injection optimization stage. This can optimize PSTs (currently) and soon HVDC and potentially even redispatch clusters. Using this will increase runtime.

plot class-attribute instance-attribute

plot = False

Whether to plot the repertoire

n_worst_contingencies class-attribute instance-attribute

n_worst_contingencies = 20

The number of worst contingencies to consider in the scoring function. This is used to determine the worst cases for overloads.

enable_bb_outage class-attribute instance-attribute

enable_bb_outage = False

Whether the optimizer should include busbar outage effects in scoring. If the preprocessed grid file does not contain busbar outage data, this flag is ignored.

bb_outage_as_nminus1 class-attribute instance-attribute

bb_outage_as_nminus1 = True

Whether busbar outages are handled as additional N-1 cases. If this is True, they are just part of the N-1 matrix and will contribute to normal overload energy, max load, etc. If this is False, a separate busbar outage penalty is computed and added to the scores. If no busbar outage data is provided in the grid model, this parameter will be ignored.

clip_bb_outage_penalty class-attribute instance-attribute

clip_bb_outage_penalty = False

Whether busbar outage penalties are clipped at 0. This is only relevant in case of bb_outage_as_nminus1=False, where busbar outage penalties are added to the scores as a separate term. If this is True, the busbar outage penalty will be clipped at 0, meaning that topologies that improve the busbar outage penalty will not be rewarded for it. If this is False, topologies that improve the busbar outage penalty will receive a negative penalty, which can lead to higher scores and thus be rewarded by the optimizer. If no busbar outage data is provided in the grid model or if bb_outage_as_nminus1 is True, this parameter will be ignored.

bb_outage_more_islands_penalty class-attribute instance-attribute

bb_outage_more_islands_penalty = 0.0

Islanding penalty used for busbar outage baseline comparisons. This is only relevant in case of bb_outage_as_nminus1=False, where busbar outage penalties are added to the scores as a separate term. If a busbar outage computation fails due to grid splits/islanding, this penalty will be applied to the score proportional to the number of islands. If no busbar outage data is provided in the grid model or if bb_outage_as_nminus1 is True, this parameter will be ignored.

mutation_repetition class-attribute instance-attribute

mutation_repetition = 1

More chance to get unique mutations by mutating multiple copies of the repertoire

random_topo_prob class-attribute instance-attribute

random_topo_prob = 0.1

The probability to create a random topology instead of mutating the existing one. This can help to escape local minima.

n_subs_mutated_lambda class-attribute instance-attribute

n_subs_mutated_lambda = 2.0

The number of substations to mutate in a single iteration is drawn from a poisson with this lambda

add_split_prob class-attribute instance-attribute

add_split_prob = 0.25

The probability to add an additional split to a substation, applied after the number of substations to mutate is drawn

change_split_prob class-attribute instance-attribute

change_split_prob = 0.3

The probability to change an existing split in a substation

remove_split_prob class-attribute instance-attribute

remove_split_prob = 0.25

The probability to remove a split from a substation

add_disconnection_prob class-attribute instance-attribute

add_disconnection_prob = 0.2

The probability to add a disconnection, applied after the number of disconnections to mutate is drawn

change_disconnection_prob class-attribute instance-attribute

change_disconnection_prob = 0.2

The probability to change an existing disconnection

remove_disconnection_prob class-attribute instance-attribute

remove_disconnection_prob = 0.2

The probability to remove an existing disconnection

pst_mutation_sigma class-attribute instance-attribute

pst_mutation_sigma = 0.0

The sigma to use for the normal distribution that mutates the PST taps. The mutation is applied by adding a random value drawn from this distribution to the current tap position. A value of 0.0 means no PST mutation.

pst_mutation_probability class-attribute instance-attribute

pst_mutation_probability = 0.2

The probability for an individual PST to be selected for mutation. A value of 0.0 means no PST mutation. A value of 1.0 means all PSTs will be mutated.

pst_reset_probability class-attribute instance-attribute

pst_reset_probability = 0.0

The probability for an individual PST to be reverted to its initial set point. A value of 0.0 means no reset. A value of 1.0 means all PSTs will be reset.

proportion_crossover class-attribute instance-attribute

proportion_crossover = 0.1

The proportion of the first topology to take in the crossover

crossover_mutation_ratio class-attribute instance-attribute

crossover_mutation_ratio = 0.5

The ratio of crossovers to mutations

target_metrics class-attribute instance-attribute

target_metrics = (('overload_energy_n_1', 1.0),)

The list of metrics to optimize for with their weights

observed_metrics class-attribute instance-attribute

observed_metrics = (
    "max_flow_n_0",
    "overload_energy_n_0",
    "overload_energy_limited_n_0",
    "max_flow_n_1",
    "overload_energy_n_1",
    "overload_energy_limited_n_1",
    "split_subs",
    "switching_distance",
)

The observed metrics, i.e. which metrics are to be computed for logging purposes. The target_metrics and me_descriptors must be included in the observed metrics and will be added automatically by the validator if they are missing

me_descriptors class-attribute instance-attribute

me_descriptors = (
    DescriptorDef(metric="split_subs", num_cells=5),
    DescriptorDef(
        metric="switching_distance", num_cells=45
    ),
)

The descriptors to use for MAP-Elites. This includes a metric that determines the cell index and a number of cells. If the metric exceeds the number of cells, it will be clipped to the largest cell index. Currently, this must be integer metrics

cell_depth class-attribute instance-attribute

cell_depth = 1

When applicable, each cell contains cell_depth unique topologies. Use 1 to retain the original map-elites behaviour

infer_missing_observed_metrics

infer_missing_observed_metrics()

Add potentially missing target and descriptor metrics to the observed metrics.

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/interfaces/messages/dc_params.py

@model_validator(mode="after")
def infer_missing_observed_metrics(self) -> "BatchedMEParameters":
    """Add potentially missing target and descriptor metrics to the observed metrics."""
    for metric, _ in self.target_metrics:
        if metric not in self.observed_metrics:
            self.observed_metrics += (metric,)
    for descriptor in self.me_descriptors:
        if descriptor.metric not in self.observed_metrics:
            self.observed_metrics += (descriptor.metric,)
    return self

probabilities_less_than_one

probabilities_less_than_one()

Check that the mutation probabilities are not larger than 1.

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/interfaces/messages/dc_params.py

@model_validator(mode="after")
def probabilities_less_than_one(self) -> "BatchedMEParameters":
    """Check that the mutation probabilities are not larger than 1."""
    if self.add_split_prob + self.change_split_prob + self.remove_split_prob > 1.0:
        raise ValueError("The sum of the substation mutation probabilities cannot be larger than 1.")
    if self.add_disconnection_prob + self.change_disconnection_prob + self.remove_disconnection_prob > 1.0:
        raise ValueError("The sum of the disconnection mutation probabilities cannot be larger than 1.")
    if self.random_topo_prob > 1.0:
        raise ValueError("The random topology probability cannot be larger than 1.")
    return self

me_descriptors_cannot_be_empty

me_descriptors_cannot_be_empty()

Check that MeDescriptors tuple is not empty.

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/interfaces/messages/dc_params.py

@model_validator(mode="after")
def me_descriptors_cannot_be_empty(self) -> "BatchedMEParameters":
    """Check that MeDescriptors tuple is not empty."""
    if not self.me_descriptors:
        raise ValueError(
            "Map-elites config used an empty tuple for MeDescriptors. "
            "Set parameter or remove the empty tuple to use the default values."
        )
    return self

LoadflowSolverParameters

Bases: BaseModel

Parameters for the loadflow solver.

max_num_splits class-attribute instance-attribute

max_num_splits = 4

The maximum number of splits per topology

max_num_disconnections class-attribute instance-attribute

max_num_disconnections = 0

The maximum number of disconnections to apply per topology

batch_size class-attribute instance-attribute

batch_size = 8

The batch size for the genetic algorithm

distributed class-attribute instance-attribute

distributed = False

Whether to run the genetic algorithm distributed over multiple devices

cross_coupler_flow class-attribute instance-attribute

cross_coupler_flow = False

Whether to compute cross-coupler flows

DCOptimizerParameters

Bases: BaseModel

The set of parameters that are used in the DC optimizer only

ga_config class-attribute instance-attribute

ga_config = BatchedMEParameters()

The configuration options for the genetic algorithm

loadflow_solver_config class-attribute instance-attribute

loadflow_solver_config = LoadflowSolverParameters()

The configuration options for the loadflow solver

double_limits class-attribute instance-attribute

double_limits = None

The double limits for the optimization, if they should be updated

summary_frequency class-attribute instance-attribute

summary_frequency = 10

The frequency to push back results, based on number of iterations. Default is after every 10 iterations.

check_command_frequency class-attribute instance-attribute

check_command_frequency = 10

The frequency to check for new commands, based on number of iterations. Should be a multiple of summary_frequency

max_num_splits_unequal_zero_if_descriptor_exists

max_num_splits_unequal_zero_if_descriptor_exists()

Check that max_num_splits is larger than the MeDescriptor size-1 of split_subs.

If the map elites descriptor contains split_subs, max_num_splits must be larger than 0. Otherwise, the computation of the BSDF will fail.

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/interfaces/messages/dc_params.py

@model_validator(mode="after")
def max_num_splits_unequal_zero_if_descriptor_exists(self) -> "DCOptimizerParameters":
    """Check that `max_num_splits` is larger than the MeDescriptor size-1 of `split_subs`.

    If the map elites descriptor contains `split_subs`, `max_num_splits` must be larger
    than 0. Otherwise, the computation of the BSDF will fail.
    """
    max_split_subs = (desc.num_cells - 1 for desc in self.ga_config.me_descriptors if desc.metric == "split_subs")
    # If the max is 0, there is likely no topol. optimization happening (e.g. only node injection optimization)
    max_value = max(max_split_subs, default=0)
    max_num_splits = self.loadflow_solver_config.max_num_splits
    if max_num_splits < max_value:
        raise ValueError(
            f"LoadflowParameters.max_num_splits ({max_num_splits}) "
            f"is smaller than the largest cell of MeDescriptor `splits_subs` ({max_value}). "
        )
    return self

toop_engine_topology_optimizer.interfaces.messages.ac_params

The parameters for the AC optimizer.

On AC, some subtelties are different to the DC optimization such as that the optimization is not batched, and the parameters are slightly different.

ACGAParameters

Bases: BaseModel

Parameters for the AC genetic algorithm

runtime_seconds class-attribute instance-attribute

runtime_seconds = 180

The maximum runtime of the AC optimization in seconds

n_worst_contingencies class-attribute instance-attribute

n_worst_contingencies = 20

How many worst contingencies to consider for the initial metrics, i.e. the top k contingencies that are used to compute the initial metrics. This is used to compute the top_k_overloads_n_1

seed class-attribute instance-attribute

seed = 42

The seed for the random number generator

runner_processes class-attribute instance-attribute

runner_processes = 1

How many processes to spawn for computing the N-1 cases in each timestep in parallel. Note that this multiplies with contingency_processes and you might run out of memory if you set both too high

contingency_processes class-attribute instance-attribute

contingency_processes = 1

How many processes to spawn for computing the contingencies of each strategy in parallel. Note that this multiplies with runner_processes and you might run out of memory if you set both too high

worst_k_runner_processes class-attribute instance-attribute

worst_k_runner_processes = 1

How many processes to spawn per topology during the worst-k stage.

worst_k_contingency_processes class-attribute instance-attribute

worst_k_contingency_processes = 1

How many processes to spawn for computing the contingencies of each strategy in parallel during the worst-k stage.

remaining_loadflow_wait_seconds class-attribute instance-attribute

remaining_loadflow_wait_seconds = 30.0

Maximum time to keep collecting non-rejected strategies before starting the remaining contingency evaluation, even if the survivor threshold has not been reached.

filter_strategy class-attribute instance-attribute

filter_strategy = None

The filter strategy to use for the optimization, used to filter out strategies based on the discriminator, median or dominator filter.

enable_ac_rejection class-attribute instance-attribute

enable_ac_rejection = True

Whether to enable the AC rejection, i.e. no messages will be sent to the results topic in case of non-acceptance.

reject_convergence_threshold class-attribute instance-attribute

reject_convergence_threshold = 1.0

The rejection threshold for the convergence rate, i.e. the split case must have at most the same amount of non converging loadflows as the unsplit case or it will be rejected.

reject_overload_threshold class-attribute instance-attribute

reject_overload_threshold = 0.95

The rejection threshold for the overload energy improvement, i.e. the split case must have at least 5% lower overload energy than the unsplit case or it will be rejected.

reject_critical_branch_threshold class-attribute instance-attribute

reject_critical_branch_threshold = 1.1

The rejection threshold for the critical branches increase, i.e. the split case must have less than 10% more critical branches than the unsplit case or it will be rejected.

early_stop_validation class-attribute instance-attribute

early_stop_validation = True

Whether to enable early stopping during the optimization process.

early_stopping_non_convergence_percentage_threshold class-attribute instance-attribute

early_stopping_non_convergence_percentage_threshold = 0.1

The threshold for the early stopping criterion, i.e. if the percentage of non-converging cases is greater than this value, the ac validation will be stopped early.

max_initial_wait_seconds class-attribute instance-attribute

max_initial_wait_seconds = 60

The maximum amount of seconds to wait for the initial DC results. If no results have arrived within this time, we assume the DC optimizer had some problem and abort the optimization run.

ACOptimizerParameters

Bases: BaseModel

The set of parameters that are used in the AC optimizer only

initial_loadflow class-attribute instance-attribute

initial_loadflow = None

If an initial AC loadflow was computed before the start of the optimization run, this can be passed and will be used e.g. to compute double limits. It will be sent back through the initial topology push.

ga_config class-attribute instance-attribute

ga_config = ACGAParameters()

The genetic algorithm configuration

AC Validation

toop_engine_topology_optimizer.ac.evolution_functions

Implements an adjusted AC evolution

Instead of the original GA evolution which only knows mutate and crossover, we introduce the following operations: - The pull operator will take a promising topology from the DC repertoire and re-evaluate it on AC. The notion of promising is defined through a interest-scoring function which tries to balance the explore/exploit trade-off.

logger module-attribute

logger = get_logger(__name__)

INF_FITNESS module-attribute

INF_FITNESS = 9999999.0

select_repertoire

select_repertoire(
    optimization_id,
    optimizer_type,
    without_children_on,
    session,
)

Select the topologies that are suitable for mutation and crossover

In this case, all topologies that satisfy the filter criteria are suitable, however a check after the mutate is necessary to ensure that the topology is not already in the database.

The unsplit strategy is never selected for mutation or crossover.

PARAMETER	DESCRIPTION
optimization_id	The optimization ID to filter for TYPE: str
optimizer_type	The optimizer types to filter for (whitelist) TYPE: list[OptimizerType]
without_children_on	If the topology already spawned a child on any of these optimizer types, it will be filtered out. TYPE: list[OptimizerType]
session	The database session to use TYPE: Session

RETURNS	DESCRIPTION
list[ACOptimTopology]	The topologies that are suitable for mutation and crossover

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/ac/evolution_functions.py

def select_repertoire(
    optimization_id: str,
    optimizer_type: list[OptimizerType],
    without_children_on: list[OptimizerType],
    session: Session,
) -> list[ACOptimTopology]:
    """Select the topologies that are suitable for mutation and crossover

    In this case, all topologies that satisfy the filter criteria are suitable, however a check after
    the mutate is necessary to ensure that the topology is not already in the database.

    The unsplit strategy is never selected for mutation or crossover.

    Parameters
    ----------
    optimization_id : str
        The optimization ID to filter for
    optimizer_type : list[OptimizerType]
        The optimizer types to filter for (whitelist)
    without_children_on : list[OptimizerType]
        If the topology already spawned a child on any of these optimizer types, it will be filtered out.

    session : Session
        The database session to use

    Returns
    -------
    list[ACOptimTopology]
        The topologies that are suitable for mutation and crossover
    """
    # Query to select topologies suitable for mutation and crossover
    mutate_query = select(ACOptimTopology).where(
        ACOptimTopology.optimization_id == optimization_id,
        ACOptimTopology.optimizer_type.in_(optimizer_type),
        ACOptimTopology.unsplit == False,  # noqa: E712
    )

    # Filter out topologies whose parent has the specified optimizer types
    # (i.e., topologies that already spawned children on those optimizer types)
    if without_children_on:
        child = aliased(ACOptimTopology)
        mutate_query = mutate_query.where(
            ~exists(
                select(1)
                .select_from(child)
                .where(
                    child.parent_id == ACOptimTopology.id,
                    child.optimizer_type.in_(without_children_on),
                    child.optimization_id == optimization_id,
                )
            )
        )

    # Execute the query and get the results
    suitable_topologies = session.exec(mutate_query).all()

    return suitable_topologies

get_unsplit_ac_topology

get_unsplit_ac_topology(optimization_id, session)

Get the unsplit AC topology for the given optimization ID

PARAMETER	DESCRIPTION
optimization_id	The optimization ID to filter for TYPE: str
session	The database session to use TYPE: Session

RETURNS	DESCRIPTION
ACOptimTopology	The unsplit AC topology for the given optimization ID, or None if not found

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/ac/evolution_functions.py

def get_unsplit_ac_topology(
    optimization_id: str,
    session: Session,
) -> Optional[ACOptimTopology]:
    """Get the unsplit AC topology for the given optimization ID

    Parameters
    ----------
    optimization_id : str
        The optimization ID to filter for
    session : Session
        The database session to use

    Returns
    -------
    ACOptimTopology
        The unsplit AC topology for the given optimization ID, or None if not found
    """
    query = select(ACOptimTopology).where(
        ACOptimTopology.optimization_id == optimization_id,
        ACOptimTopology.unsplit == True,  # noqa: E712
        ACOptimTopology.optimizer_type == OptimizerType.AC,
    )
    return session.exec(query).first()

default_scorer

default_scorer(metrics)

Return raw fitness values for the default lower-is-better AC selection.

PARAMETER	DESCRIPTION
metrics	The metrics DataFrame to score TYPE: DataFrame

RETURNS	DESCRIPTION
Series	The topology fitness values.

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/ac/evolution_functions.py

def default_scorer(metrics: pd.DataFrame) -> pd.Series:
    """Return raw fitness values for the default lower-is-better AC selection.

    Parameters
    ----------
    metrics : pd.DataFrame
        The metrics DataFrame to score

    Returns
    -------
    pd.Series
        The topology fitness values.
    """
    return metrics["fitness"]

get_contingency_indices_from_ids

get_contingency_indices_from_ids(
    case_ids, n_minus1_definition
)

Map contingency ids to their indices in the N-1 definition.

This is a helper method used in update_initial_metrics_with_worst_k_contingencies method.

PARAMETER	DESCRIPTION
case_ids	A list of contingency ids for a specific topology. TYPE: Sequence[str]
n_minus1_definition	The N-1 definition containing the contingencies. TYPE: Nminus1Definition

RETURNS	DESCRIPTION
list[int]	A list of indices of the contingencies in the N-1 definition. If a contingency id is not found, it will be skipped.

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/ac/evolution_functions.py

def get_contingency_indices_from_ids(case_ids: Collection[str], n_minus1_definition: Nminus1Definition) -> list[int]:
    """Map contingency ids to their indices in the N-1 definition.

    This is a helper method used in update_initial_metrics_with_worst_k_contingencies
    method.

    Parameters
    ----------
    case_ids : Sequence[str]
        A list of contingency ids for a specific topology.
    n_minus1_definition : Nminus1Definition
        The N-1 definition containing the contingencies.

    Returns
    -------
    list[int]
        A list of indices of the contingencies in the N-1 definition. If a contingency id is not found, it
        will be skipped.
    """
    id_to_index = {cont.id: idx for idx, cont in enumerate(n_minus1_definition.contingencies)}
    case_indices = [id_to_index[case_id] for case_id in case_ids if case_id in id_to_index]
    return case_indices

pull

pull(selected_strategy, session=None)

Pull a promising topology from the DC repertoire to AC

This only copies the topology without any changes other than setting the optimizer type to AC. This function takes a list of selected DC topologies and creates corresponding AC topologies by copying relevant attributes, setting the optimizer type to AC, and merging contingency case indices from both the DC and unsplit AC topologies.

PARAMETER	DESCRIPTION
selected_strategy	The selected strategy to pull TYPE: ACStrategy
session	The database session to use, by default None. The session object is used to fetch the unsplit AC topology which is then used to add the critical contingency cases to the pulled strategy. These critical cases can then be used for early stopping of AC N-1 contingency analysis. TYPE: Session DEFAULT: None

RETURNS	DESCRIPTION
ACStrategy	The a copy of the input topologies with the optimizer type set to AC

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/ac/evolution_functions.py

def pull(
    selected_strategy: ACStrategy,
    session: Session = None,
) -> ACStrategy:
    """Pull a promising topology from the DC repertoire to AC

    This only copies the topology without any changes other than setting the optimizer type to AC.
    This function takes a list of selected DC topologies and creates corresponding AC topologies by copying
    relevant attributes, setting the optimizer type to AC, and merging contingency case indices from both
    the DC and unsplit AC topologies.

    Parameters
    ----------
    selected_strategy : ACStrategy
        The selected strategy to pull
    session : Session, optional
        The database session to use, by default None. The session object is used to fetch the unsplit AC topology
        which is then used to add the critical contingency cases to the pulled strategy. These critical cases
        can then be used for early stopping of AC N-1 contingency analysis.

    Returns
    -------
    ACStrategy
        The a copy of the input topologies with the optimizer type set to AC
    """
    if not selected_strategy:
        return []

    optimization_id = selected_strategy[0].optimization_id

    # Unsplit AC topology will be used to merge critical contingency cases. This topology is pushed in the repertoire
    # in the initialization phase of the AC optimizer.
    unsplit_ac_topo = get_unsplit_ac_topology(optimization_id=optimization_id, session=session) if session else None
    worst_k_cont_ids_unsplit = set(unsplit_ac_topo.worst_k_contingency_cases) if unsplit_ac_topo else set()
    pulled_strategy = []
    for topo in selected_strategy:
        # Merge case_ids from DC and unsplit AC
        merged_ids = sorted(set(topo.worst_k_contingency_cases) | worst_k_cont_ids_unsplit)

        data = topo.model_dump(
            include=[
                "actions",
                "disconnections",
                "pst_setpoints",
                "unsplit",
                "timestep",
                "optimization_id",
                "strategy_hash",
            ]
        )
        metrics = topo.metrics
        metrics["fitness_dc"] = topo.fitness

        new_topo = ACOptimTopology(
            **data,
            optimizer_type=OptimizerType.AC,
            fitness=-INF_FITNESS,
            parent_id=topo.id,
            metrics=metrics,
            worst_k_contingency_cases=merged_ids,
        )
        new_topo.metrics["top_k_overloads_n_1"] = (
            unsplit_ac_topo.metrics.get("top_k_overloads_n_1", None) if unsplit_ac_topo else None
        )
        pulled_strategy.append(new_topo)

    return pulled_strategy

evolution

evolution(
    rng,
    session,
    optimization_id,
    max_retries,
    batch_size,
    filter_strategy=None,
)

Perform the AC evolution.

PARAMETER	DESCRIPTION
rng	The random number generator to use TYPE: Generator
session	The database session to use, will write the new topologies to the database TYPE: Session
optimization_id	The optimization ID to filter for TYPE: str
max_retries	The maximum number of retries to perform if a strategy is already in the database TYPE: int
batch_size	Number of unevaluated topologies to sample and convert to AC in one try. TYPE: int
filter_strategy	The filter strategy to use for the optimization, used to filter out strategies that are too far away from the original topology. TYPE: Optional[FilterStrategy] DEFAULT: None

RETURNS	DESCRIPTION
ACStrategy	The strategy that was created during the evolution or an empty list if something went wrong at all retries

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/ac/evolution_functions.py

def evolution(
    rng: Rng,
    session: Session,
    optimization_id: str,
    max_retries: int,
    batch_size: int,
    filter_strategy: Optional[FilterStrategy] = None,
) -> list[ACOptimTopology]:
    """Perform the AC evolution.

    Parameters
    ----------
    rng : Rng
        The random number generator to use
    session : Session
        The database session to use, will write the new topologies to the database
    optimization_id : str
        The optimization ID to filter for
    max_retries : int
        The maximum number of retries to perform if a strategy is already in the database
    batch_size : int
        Number of unevaluated topologies to sample and convert to AC in one try.
    filter_strategy : Optional[FilterStrategy]
        The filter strategy to use for the optimization,
        used to filter out strategies that are too far away from the original topology.

    Returns
    -------
    ACStrategy
        The strategy that was created during the evolution or an empty list if something went
        wrong at all retries
    """
    for _try in range(max_retries):
        new_topo_batch = evolution_try(
            rng=rng,
            session=session,
            optimization_id=optimization_id,
            batch_size=batch_size,
            filter_strategy=filter_strategy,
        )
        if len(new_topo_batch):
            return new_topo_batch
    return []

evolution_try

evolution_try(
    rng,
    session,
    optimization_id,
    batch_size,
    filter_strategy=None,
)

Perform a single try of the AC evolution.

PARAMETER	DESCRIPTION
rng	The random number generator to use TYPE: Generator
session	The database session to use, will write the new topologies to the database TYPE: Session
optimization_id	The optimization ID to filter for TYPE: str
batch_size	Number of unevaluated topologies to sample and convert to AC. TYPE: int
filter_strategy	The filter strategy to use for the optimization, used to filter out strategies that are too far away from the original topology. TYPE: Optional[FilterStrategy] DEFAULT: None

RETURNS	DESCRIPTION
list[ACOptimTopology]	The list of topologies that were created during the evolution or an empty list if something went wrong during the try

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/ac/evolution_functions.py

def evolution_try(
    rng: Rng,
    session: Session,
    optimization_id: str,
    batch_size: int,
    filter_strategy: Optional[FilterStrategy] = None,
) -> list[ACOptimTopology]:
    """Perform a single try of the AC evolution.

    Parameters
    ----------
    rng : Rng
        The random number generator to use
    session : Session
        The database session to use, will write the new topologies to the database
    optimization_id : str
        The optimization ID to filter for
    batch_size : int
        Number of unevaluated topologies to sample and convert to AC.
    filter_strategy : Optional[FilterStrategy]
        The filter strategy to use for the optimization,
        used to filter out strategies that are too far away from the original topology.

    Returns
    -------
    list[ACOptimTopology]
        The list of topologies that were created during the evolution or an empty list if something went
        wrong during the try
    """
    selected_topologies = select_strategy(
        rng=rng,
        repertoire=select_repertoire(
            optimization_id=optimization_id,
            optimizer_type=[OptimizerType.DC, OptimizerType.AC],
            without_children_on=[],
            session=session,
        ),
        candidates=select_repertoire(
            optimization_id=optimization_id,
            optimizer_type=[OptimizerType.DC],
            without_children_on=[OptimizerType.AC],
            session=session,
        ),
        interest_scorer=default_scorer,
        batch_size=batch_size,
        lower_scores_are_better=True,
        filter_strategy=filter_strategy,
    )
    new_strategy = pull(selected_strategy=selected_topologies, session=session)
    # Something went wrong during the mutation
    if new_strategy == []:
        return []

    # Insert the new strategies into the database
    for topo in new_strategy:
        session.add(topo)

    try:
        session.commit()
        for topo in new_strategy:
            session.refresh(topo)
    except IntegrityError:
        session.rollback()
        return []
    return new_strategy

toop_engine_topology_optimizer.ac.listener

A listener that listens for new topologies on a kafka result stream and saves them to db

This is intended for usage both in the AC optimizer and the backend, as they both listen for topologies on the result kafka stream. The backend has some further logic as it watches all optimizations and might change the state of the optimization to "running" if it receives a start optimization result and "stopped" if a stop optimization result is received.

One of the requirements is that the AC listener will only want to listen to a single optimization_id and not to messages by itself. Hence a filtering mechanic is added.

finish_optimization

finish_optimization(db, optimization_id, optimizer_type)

Mark an optimization as finished in the database.

If it was a DC optimization, this function just saves an entry to DB. If it was an AC optimization, this function also deletes all topologies from the database that belong to this optimization, as they are no longer relevant and we don't want to keep them around.

PARAMETER	DESCRIPTION
db	The database session to use for saving the finished optimization TYPE: Session
optimization_id	The optimization ID of the finished optimization as sent via kafka TYPE: str
optimizer_type	Which optimizer type has finished TYPE: OptimizerType

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/ac/listener.py

def finish_optimization(db: Session, optimization_id: str, optimizer_type: OptimizerType) -> None:
    """Mark an optimization as finished in the database.

    If it was a DC optimization, this function just saves an entry to DB.
    If it was an AC optimization, this function also deletes all topologies from the database that belong to this
    optimization, as they are no longer relevant and we don't want to keep them around.

    Parameters
    ----------
    db : Session
        The database session to use for saving the finished optimization
    optimization_id : str
        The optimization ID of the finished optimization as sent via kafka
    optimizer_type : OptimizerType
        Which optimizer type has finished
    """
    if optimizer_type == OptimizerType.AC:
        # For AC optimizations, we delete all topologies that belong to this optimization, as they are no longer relevant
        db.exec(delete(ACOptimTopology).where(ACOptimTopology.optimization_id == optimization_id))
        db.commit()

    # For both AC and DC optimizations, we save an entry to the finished optimizations table to prevent picking up old
    # topologies from previous optimizations
    try:
        finished_optimization = FinishedOptimizations(
            optimization_id=optimization_id,
            optimizer_type=optimizer_type,
        )
        db.add(finished_optimization)
        db.commit()
    except IntegrityError:
        db.rollback()

poll_results_topic

poll_results_topic(db, consumer, first_poll=True)

Poll the results topic for new topologies to store in the DB

We store topologies from all optimization jobs, as it could be that we will later optimize the same job in this worker.

PARAMETER	DESCRIPTION
db	The database session to use for saving the topologies TYPE: Session
consumer	The kafka consumer to poll messages from. It should already be subscribed to the result topic. TYPE: LongRunningKafkaConsumer
first_poll	If True, we assume the optimizatin has just started and we can afford to wait for the first DC results for a longer time (30 seconds). If False, we assume the optimization is already running and we don't want to block the consumer for too long (100 ms), by default True TYPE: bool DEFAULT: True

RETURNS	DESCRIPTION
dict[str, int]	A dictionary with the number of topologies added to the database for each optimization ID
list[str]	A list of optimization IDs for which a stop optimization result was received.

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/ac/listener.py

def poll_results_topic(
    db: Session,
    consumer: LongRunningKafkaConsumer,
    first_poll: bool = True,
) -> tuple[dict[str, int], list[str]]:
    """Poll the results topic for new topologies to store in the DB

    We store topologies from all optimization jobs, as it could be that we will later optimize the same job in this worker.

    Parameters
    ----------
    db : Session
        The database session to use for saving the topologies
    consumer : LongRunningKafkaConsumer
        The kafka consumer to poll messages from. It should already be subscribed to the result
        topic.
    first_poll : bool, optional
        If True, we assume the optimizatin has just started and we can afford to wait for the first DC results for
        a longer time (30 seconds). If False, we assume the optimization is already running and we don't want to block
        the consumer for too long (100 ms), by default True

    Returns
    -------
    dict[str, int]
        A dictionary with the number of topologies added to the database for each optimization ID
    list[str]
        A list of optimization IDs for which a stop optimization result was received.
    """
    added_topos = {}
    finished_optimizations: list[Result] = []
    messages = consumer.consume(timeout=30.0 if first_poll else 0.1, num_messages=10000)
    for message in messages:
        result = Result.model_validate_json(deserialize_message(message.value()))

        strategy = None
        if isinstance(result.result, TopologyPushResult):
            strategy = result.result.strategy
        elif isinstance(result.result, OptimizationStartedResult):
            strategy = result.result.initial_topology
        elif isinstance(result.result, OptimizationStoppedResult):
            finished_optimizations.append(result)
            continue
        else:
            continue

        topologies = convert_message_topo_to_db_topo(strategy, result.optimization_id, result.optimizer_type)

        # Push the topologies to the database, ignoring duplicates
        # Duplicates will trigger an IntegrityError
        for topo in topologies:
            try:
                db.add(topo)
                db.commit()
                added_topos[result.optimization_id] = added_topos.get(result.optimization_id, 0) + 1
            except IntegrityError:
                db.rollback()
                pass

    for result in finished_optimizations:
        finish_optimization(db, result.optimization_id, result.optimizer_type)
        if result.optimization_id in added_topos:
            del added_topos[result.optimization_id]

    return added_topos, [result.optimization_id for result in finished_optimizations]

toop_engine_topology_optimizer.ac.optimizer

Implements initialize and run_epoch functions for the AC optimizer

logger module-attribute

logger = get_logger(__name__)

AcNotConvergedError

Bases: Exception

An exception that is raised when the AC optimization did not converge in the base grid

update_initial_metrics_with_worst_k_contingencies

update_initial_metrics_with_worst_k_contingencies(
    initial_loadflow, initial_metrics, worst_k
)

Update the initial metrics with the worst k contingencies.

This function computes the worst k contingencies for each timestep in the initial loadflow results and updates the initial metrics with the case ids and the top k overloads. This way, a baseline for the worst k contingencies to compare to is established, i.e. in the initial loadflow results these are the reference, disregarding the worst k for the specific strategy.

PARAMETER	DESCRIPTION
initial_loadflow	The initial loadflow results containing the branch results. TYPE: LoadflowResultsPolars
initial_metrics	The initial metrics for each timestep. TYPE: Metrics
worst_k	The number of worst contingencies to consider for the initial metrics. TYPE: int

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/ac/optimizer.py

def update_initial_metrics_with_worst_k_contingencies(
    initial_loadflow: LoadflowResultsPolars,
    initial_metrics: Metrics,
    worst_k: int,
) -> None:
    """Update the initial metrics with the worst k contingencies.

    This function computes the worst k contingencies for each timestep in the initial loadflow results
    and updates the initial metrics with the case ids and the top k overloads.
    This way, a baseline for the worst k contingencies to compare to is established, i.e. in the initial loadflow results
    these are the reference, disregarding the worst k for the specific strategy.

    Parameters
    ----------
    initial_loadflow : LoadflowResultsPolars
        The initial loadflow results containing the branch results.
    initial_metrics : Metrics
        The initial metrics for each timestep.
    worst_k : int
        The number of worst contingencies to consider for the initial metrics.
    """
    case_ids, top_k_overloads_n_1 = get_worst_k_contingencies_ac(
        cast("patpl.LazyFrame[BranchResultSchemaPolars]", initial_loadflow.branch_results),
        k=worst_k,
    )

    # case_ids is an empty list if the loadflow didn't converge -> the initial_loadflow.branch_results is full of NaNs
    if len(case_ids) == 0:
        logger.warning("No worst case ids found as the loadflow didn't converge")
        top_k_overloads_n_1 = [0.0]
        case_ids = [[]]

    initial_metrics.extra_scores.update(
        {
            "top_k_overloads_n_1": top_k_overloads_n_1[0],
        }
    )
    initial_metrics.worst_k_contingency_cases = case_ids[0]

make_runner

make_runner(
    action_set,
    nminus1_definition,
    grid_file,
    n_processes,
    batch_size,
    processed_gridfile_fs,
    lf_params=None,
)

Initialize a runner for a gridfile, action set and n-1 def

PARAMETER	DESCRIPTION
action_set	The action set to use TYPE: ActionSet
nminus1_definition	The N-1 definition to use TYPE: Nminus1Definition
grid_file	The grid file to use TYPE: GridFile
n_processes	The number of processes to use, from the ACGAParameters TYPE: int
batch_size	The batch size to use, if any, from the ACGAParameters TYPE: Optional[int]
processed_gridfile_fs	The target filesystem for the preprocessing worker. This contains all processed grid files. During the import job, a new folder import_results.data_folder was created which will be completed with the preprocess call to this function. Internally, only the data folder is passed around as a dirfs. Note that the unprocessed_gridfile_fs is not needed here anymore, as all preprocessing steps that need the unprocessed gridfiles were already done. TYPE: AbstractFileSystem
lf_params	The loadflow parameters to use for the runner, if any. This is passed in the preprocessing results and can be used to run the loadflows with the same parameters as the initial loadflow in the preprocessing step. If None, the runner will use default parameters. TYPE: Parameters | dict | None DEFAULT: None

RETURNS	DESCRIPTION
AbstractLoadflowRunner	The initialized loadflow runner, either Pandapower or Powsybl

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/ac/optimizer.py

def make_runner(
    action_set: ActionSet,
    nminus1_definition: Nminus1Definition,
    grid_file: GridFile,
    n_processes: int,
    batch_size: Optional[int],
    processed_gridfile_fs: AbstractFileSystem,
    lf_params: pypowsybl.loadflow.Parameters | dict | None = None,
) -> AbstractLoadflowRunner:
    """Initialize a runner for a gridfile, action set and n-1 def

    Parameters
    ----------
    action_set : ActionSet
        The action set to use
    nminus1_definition : Nminus1Definition
        The N-1 definition to use
    grid_file : GridFile
        The grid file to use
    n_processes : int
        The number of processes to use, from the ACGAParameters
    batch_size : Optional[int]
        The batch size to use, if any, from the ACGAParameters
    processed_gridfile_fs: AbstractFileSystem
        The target filesystem for the preprocessing worker. This contains all processed grid files.
        During the import job,  a new folder import_results.data_folder was created
        which will be completed with the preprocess call to this function.
        Internally, only the data folder is passed around as a dirfs.
        Note that the unprocessed_gridfile_fs is not needed here anymore, as all preprocessing steps that need the
        unprocessed gridfiles were already done.
    lf_params: Optional[pypowsybl.loadflow.Parameters | dict]
        The loadflow parameters to use for the runner, if any. This is passed in the preprocessing results
        and can be used to run the loadflows with the same parameters
        as the initial loadflow in the preprocessing step.
        If None, the runner will use default parameters.

    Returns
    -------
    AbstractLoadflowRunner
        The initialized loadflow runner, either Pandapower or Powsybl
    """
    logger.debug(
        "Initializing loadflow runner "
        f"framework={grid_file.framework}, grid_folder={grid_file.grid_folder}, "
        f"n_processes={n_processes}, batch_size={batch_size}"
    )

    if grid_file.framework == Framework.PANDAPOWER:
        runner = PandapowerRunner(
            n_processes=n_processes, batch_size=batch_size, lf_params=lf_params if isinstance(lf_params, dict) else None
        )
        grid_file_path = Path(grid_file.grid_folder) / PREPROCESSING_PATHS["grid_file_path_pandapower"]
    elif grid_file.framework == Framework.PYPOWSYBL:
        runner = PowsyblRunner(
            n_processes=n_processes,
            batch_size=batch_size,
            lf_params=lf_params if isinstance(lf_params, pypowsybl.loadflow.Parameters) else None,
        )
        grid_file_path = Path(grid_file.grid_folder) / PREPROCESSING_PATHS["grid_file_path_powsybl"]
    else:
        raise ValueError(f"Unknown framework {grid_file.framework}")
    runner.load_base_grid_fs(filesystem=processed_gridfile_fs, grid_path=grid_file_path)
    runner.store_action_set(action_set)
    runner.store_nminus1_definition(nminus1_definition)
    logger.debug(
        f"Runner prepared for {grid_file_path}: n_actions={len(action_set.local_actions)}, "
        f"n_contingencies={len(nminus1_definition.contingencies)}"
    )
    return runner

initialize_optimization

initialize_optimization(
    params,
    session,
    optimization_id,
    grid_file,
    loadflow_result_fs,
    processed_gridfile_fs,
)

Initialize an optimization run for the AC optimizer

PARAMETER	DESCRIPTION
params	The parameters for the AC optimizer TYPE: ACOptimizerParameters
session	The database session to use for storing topologies TYPE: Session
optimization_id	The ID of the optimization run TYPE: str
grid_file	The grid file to optimize on TYPE: GridFile
loadflow_result_fs	A filesystem where the loadflow results are stored. Loadflows will be stored here using the uuid generation process and passed as a StoredLoadflowReference which contains the subfolder in this filesystem. TYPE: AbstractFileSystem
processed_gridfile_fs	The target filesystem for the preprocessing worker. This contains all processed grid files. During the import job, a new folder import_results.data_folder was created which will be completed with the preprocess call to this function. Internally, only the data folder is passed around as a dirfs. Note that the unprocessed_gridfile_fs is not needed here anymore, as all preprocessing steps that need the unprocessed gridfiles were already done. TYPE: AbstractFileSystem

RETURNS	DESCRIPTION
OptimizerData	The initial optimizer data
Strategy	The initial strategy

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/ac/optimizer.py

def initialize_optimization(
    params: ACOptimizerParameters,
    session: Session,
    optimization_id: str,
    grid_file: GridFile,
    loadflow_result_fs: AbstractFileSystem,
    processed_gridfile_fs: AbstractFileSystem,
) -> tuple[OptimizerData, Strategy]:
    """Initialize an optimization run for the AC optimizer

    Parameters
    ----------
    params : ACOptimizerParameters
        The parameters for the AC optimizer
    session : Session
        The database session to use for storing topologies
    optimization_id : str
        The ID of the optimization run
    grid_file : GridFile
        The grid file to optimize on
    loadflow_result_fs: AbstractFileSystem
        A filesystem where the loadflow results are stored. Loadflows will be stored here using the uuid generation process
        and passed as a StoredLoadflowReference which contains the subfolder in this filesystem.
    processed_gridfile_fs: AbstractFileSystem
        The target filesystem for the preprocessing worker. This contains all processed grid files.
        During the import job,  a new folder import_results.data_folder was created
        which will be completed with the preprocess call to this function.
        Internally, only the data folder is passed around as a dirfs.
        Note that the unprocessed_gridfile_fs is not needed here anymore, as all preprocessing steps that need the
        unprocessed gridfiles were already done.

    Returns
    -------
    OptimizerData
        The initial optimizer data
    Strategy
        The initial strategy
    """
    ga_config = params.ga_config
    logger.info("Initializing AC optimization", framework={grid_file.framework}, seed={ga_config.seed})
    # Load the network datas
    action_set = load_action_set_fs(
        filesystem=processed_gridfile_fs,
        json_file_path=grid_file.action_set_file,
        diff_file_path=grid_file.action_set_diff_file,
    )
    nminus1_definition = load_pydantic_model_fs(
        filesystem=processed_gridfile_fs, file_path=grid_file.nminus1_definition_file, model_class=Nminus1Definition
    )

    lf_params = load_lf_params_from_fs(filesystem=processed_gridfile_fs, file_path=grid_file.loadflow_parameters_file)

    logger.debug("Loaded preprocessing inputs")

    base_case_id = getattr(nminus1_definition.base_case, "id", None)

    # Prepare the loadflow runners
    def build_runner_group(n_topo_processes: int, n_contingency_processes: int) -> RunnerGroup:
        return [
            make_runner(
                action_set,
                nminus1_definition,
                grid_file,
                n_processes=n_contingency_processes,
                batch_size=None,
                processed_gridfile_fs=processed_gridfile_fs,
                lf_params=lf_params,
            )
            for _ in range(n_topo_processes)
        ]

    worst_k_runner_group = build_runner_group(ga_config.worst_k_runner_processes, ga_config.worst_k_contingency_processes)
    logger.debug(f"Prepared {len(worst_k_runner_group)} runner(s) for Early Stopping AC optimization")

    runner_group = build_runner_group(ga_config.runner_processes, ga_config.contingency_processes)
    logger.debug(f"Prepared {len(runner_group)} runner(s) for AC optimization")

    # Prepare the evolution function
    rng = np.random.default_rng(ga_config.seed)
    evolution_fn = partial(
        evolution,
        rng=rng,
        session=session,
        optimization_id=optimization_id,
        max_retries=10,
        batch_size=ga_config.worst_k_runner_processes,
        filter_strategy=ga_config.filter_strategy,
    )

    # Prepare the initial strategy
    unsplit_topology = ACOptimTopology(
        actions=[],
        disconnections=[],
        pst_setpoints=None,
        timestep=0,
        fitness=0,
        unsplit=True,
        strategy_hash=b"willbeupdated",
        optimization_id=optimization_id,
        optimizer_type=OptimizerType.AC,
    )
    initial_hash = hash_topologies([unsplit_topology])
    unsplit_topology.strategy_hash = initial_hash

    def store_loadflow(loadflow: LoadflowResultsPolars) -> StoredLoadflowReference:
        return save_loadflow_results_polars(
            loadflow_result_fs, f"{optimization_id}-{loadflow.job_id}-{datetime.now()}", loadflow
        )

    def loadflow_ref(loadflow: StoredLoadflowReference) -> LoadflowResultsPolars:
        return load_loadflow_results_polars(loadflow_result_fs, reference=loadflow)

    # This requires a full loadflow computation if the loadflow results are not passed in
    initial_loadflow_reference = params.initial_loadflow
    if initial_loadflow_reference is None:
        logger.info("No initial loadflow provided, computing initial AC loadflow")
        initial_loadflow, _, initial_metrics = compute_loadflow_and_metrics(
            topology=unsplit_topology,
            runner=runner_group[0],
            base_case_id=base_case_id,
        )
        initial_loadflow_reference = store_loadflow(initial_loadflow)
        logger.debug(f"Initial AC loadflow computed and stored under reference={initial_loadflow_reference}")
    else:
        logger.info(f"Using precomputed initial loadflow reference={initial_loadflow_reference}")
        # If the initial loadflow is passed in, we load it from the database
        initial_loadflow = loadflow_ref(initial_loadflow_reference)
        # Compute the metrics for the initial loadflow
        initial_metrics = compute_metrics_single_timestep(
            actions=unsplit_topology.actions,
            disconnections=unsplit_topology.disconnections,
            loadflow=initial_loadflow,
            additional_info=None,
            base_case_id=base_case_id,
        )
        logger.debug("Computed initial metrics from provided loadflow")

    # Update the initial metrics with the worst k contingencies
    update_initial_metrics_with_worst_k_contingencies(
        initial_loadflow, initial_metrics, params.ga_config.n_worst_contingencies
    )

    logger.debug(
        "Initial convergence summary: "
        f"non_converging={initial_metrics.extra_scores.get('non_converging_loadflows', 0)}, "
        f"allowed={len(nminus1_definition.contingencies) / 2}"
    )
    if initial_metrics.extra_scores.get("non_converging_loadflows", 0) > len(nminus1_definition.contingencies) / 2:
        raise AcNotConvergedError(
            "Too many non-converging loadflows in initial loadflow: "
            f"{initial_metrics.extra_scores.get('non_converging_loadflows', 0)} > "
            f"{len(nminus1_definition.contingencies) / 2}"
        )

    unsplit_topology.fitness = initial_metrics.fitness
    unsplit_topology.metrics = initial_metrics.extra_scores
    unsplit_topology.worst_k_contingency_cases = initial_metrics.worst_k_contingency_cases
    unsplit_topology.set_loadflow_reference(initial_loadflow_reference)
    session.add(unsplit_topology)
    session.commit()
    logger.debug("Stored initial AC strategy in DB")

    # As we have the initial loadflows, we can now define a scoring+acceptance function
    scoring_params = ACScoringParameters(
        base_case_id=base_case_id,
        early_stop_validation=ga_config.early_stop_validation,
        reject_convergence_threshold=ga_config.reject_convergence_threshold,
        reject_overload_threshold=ga_config.reject_overload_threshold,
        reject_critical_branch_threshold=ga_config.reject_critical_branch_threshold,
    )

    def scoring_fn(
        topologies: list[ACOptimTopology],
        early_stage_results: Optional[list[EarlyStoppingStageResult]] = None,
    ) -> list[TopologyScoringResult]:
        return score_topology_batch(
            topologies,
            runner_group=runner_group,
            metrics_unsplit=initial_metrics,
            scoring_params=scoring_params,
            early_stage_results=early_stage_results,
        )

    worst_k_scoring_fn = partial(
        score_strategy_worst_k_batch,
        worst_k_runner_groups=worst_k_runner_group,
        metrics_unsplit=initial_metrics,
        loadflow_results_unsplit=initial_loadflow,
        scoring_params=scoring_params,
    )
    # Convert the initial strategy to a message strategy
    initial_strategy_message = convert_db_topo_to_message_topo([unsplit_topology])
    assert len(initial_strategy_message) == 1

    logger.info(
        f"Initialization completed, metrics: {initial_metrics.extra_scores}, fitness: {initial_metrics.fitness}, "
        f"worst_k_contingency_cases: {initial_metrics.worst_k_contingency_cases}. Waiting for DC results..."
    )

    return (
        OptimizerData(
            params=params,
            session=session,
            evolution_fn=evolution_fn,
            scoring_fn=scoring_fn,
            worst_k_scoring_fn=worst_k_scoring_fn,
            store_loadflow_fn=store_loadflow,
            load_loadflow_fn=loadflow_ref,
            rng=rng,
            framework=grid_file.framework,
            runners=runner_group,
            worst_k_runner_groups=worst_k_runner_group,
            action_set=action_set,
        ),
        initial_strategy_message[0],
    )

wait_for_first_dc_results

wait_for_first_dc_results(
    results_consumer,
    session,
    max_wait_time,
    optimization_id,
    heartbeat_fn,
)

Wait an initial period for DC results to arrive before proceeding with the optimization.

Call this after initialize optimization and before run epoch to ensure that the DC optimizer has started, and avoid the AC optimizer idling while waiting for the first DC results to arrive.

PARAMETER	DESCRIPTION
results_consumer	The consumer where to listen for DC results TYPE: LongRunningKafkaConsumer
session	The database session to use for storing topologies TYPE: Session
max_wait_time	The maximum time to wait for DC results, in seconds TYPE: PositiveInt
optimization_id	The ID of the optimization run, used to filter the incoming topologies and only proceed when DC results from the correct optimization run arrive. Note that other DC runs could be active. TYPE: str
heartbeat_fn	A function to send heartbeats while waiting for DC results, as this wait time can be relatively long. TYPE: Callable[[], None]

RAISES	DESCRIPTION
TimeoutError	If no DC results arrive within the maximum wait time

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/ac/optimizer.py

def wait_for_first_dc_results(
    results_consumer: LongRunningKafkaConsumer,
    session: Session,
    max_wait_time: PositiveInt,
    optimization_id: str,
    heartbeat_fn: Callable[[], None],
) -> None:
    """Wait an initial period for DC results to arrive before proceeding with the optimization.

    Call this after initialize optimization and before run epoch to ensure that the DC optimizer has started, and avoid the
    AC optimizer idling while waiting for the first DC results to arrive.

    Parameters
    ----------
    results_consumer : LongRunningKafkaConsumer
        The consumer where to listen for DC results
    session : Session
        The database session to use for storing topologies
    max_wait_time : PositiveInt
        The maximum time to wait for DC results, in seconds
    optimization_id : str
        The ID of the optimization run, used to filter the incoming topologies and only proceed when DC results from
        the correct optimization run arrive. Note that other DC runs could be active.
    heartbeat_fn : Callable[[], None]
        A function to send heartbeats while waiting for DC results, as this wait time can be relatively long.

    Raises
    ------
    TimeoutError
        If no DC results arrive within the maximum wait time
    """
    existing_dc_topology_id = session.exec(
        select(ACOptimTopology.id)
        .where(ACOptimTopology.optimization_id == optimization_id)
        .where(ACOptimTopology.optimizer_type == OptimizerType.DC)
    ).first()
    if isinstance(existing_dc_topology_id, int):
        logger.info("DC results were already available in the database, proceeding with optimization")
        return

    start_wait = datetime.now()
    poll_iteration = 0
    while datetime.now() - start_wait < timedelta(seconds=max_wait_time):
        poll_iteration += 1
        elapsed_seconds = (datetime.now() - start_wait).total_seconds()
        added_topos, stopped_optimization_ids = poll_results_topic(db=session, consumer=results_consumer, first_poll=True)
        logger.debug(
            f"DC poll iteration={poll_iteration}, elapsed_seconds={elapsed_seconds:.2f}, "
            f"received_topologies={len(added_topos)}"
        )
        new_topos_for_optimization = added_topos.get(optimization_id, 0)
        if new_topos_for_optimization > 0:
            logger.info(f"Received {new_topos_for_optimization} topologies from DC results, proceeding with optimization")
            return
        if optimization_id in stopped_optimization_ids:
            logger.warning("Received DC optimization stopped message before receiving any DC results, stop optimization")
            break
        heartbeat_fn()
    raise TimeoutError(f"Did not receive DC results within {max_wait_time} seconds, cannot proceed with optimization")

persist_topology

persist_topology(topology, scoring_result, optimizer_data)

Persist a topology and return the payload needed for later emission.

This function stores the topology in the database, including the loadflow results if available, and returns the final message payload and scoring result. Emission is handled separately so callers can persist without necessarily sending a result immediately.

PARAMETER	DESCRIPTION
topology	The topology to persist and send TYPE: ACOptimTopology
scoring_result	The scoring result for the topology, containing the metrics, loadflow results and rejection reason if any TYPE: TopologyScoringResult
optimizer_data	The optimizer data containing the session and the loadflow storage function TYPE: OptimizerData

RETURNS	DESCRIPTION
tuple[Topology, TopologyScoringResult]	The message payload to emit and the final scoring result after persistence handling.

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/ac/optimizer.py

def persist_topology(
    topology: ACOptimTopology,
    scoring_result: TopologyScoringResult,
    optimizer_data: OptimizerData,
) -> tuple[Topology, TopologyScoringResult]:
    """Persist a topology and return the payload needed for later emission.

    This function stores the topology in the database, including the loadflow results if available,
    and returns the final message payload and scoring result. Emission is handled separately so callers
    can persist without necessarily sending a result immediately.

    Parameters
    ----------
    topology : ACOptimTopology
        The topology to persist and send
    scoring_result : TopologyScoringResult
        The scoring result for the topology, containing the metrics, loadflow results and rejection reason if any
    optimizer_data : OptimizerData
        The optimizer data containing the session and the loadflow storage function

    Returns
    -------
    tuple[Topology, TopologyScoringResult]
        The message payload to emit and the final scoring result after persistence handling.
    """
    rejection_reason = scoring_result.rejection_reason
    loadflow_result_reference = None
    if scoring_result.loadflow_results is not None:
        try:
            loadflow_result_reference = optimizer_data.store_loadflow_fn(scoring_result.loadflow_results)
        except Exception as exc:
            logger.error("Error while storing loadflow results", error=str(exc))
            rejection_reason = TopologyRejectionReason(
                criterion="error", description=str(exc), value_after=1.0, value_before=0.0, early_stopping=False
            )
            scoring_result = TopologyScoringResult(
                loadflow_results=None,
                metrics=Metrics(fitness=INF_FITNESS, extra_scores={}),
                rejection_reason=rejection_reason,
            )

    if "fitness_dc" in topology.metrics:
        scoring_result.metrics.extra_scores["fitness_dc"] = topology.metrics["fitness_dc"]
    topology.metrics = scoring_result.metrics.extra_scores
    topology.fitness = scoring_result.metrics.fitness
    topology.acceptance = rejection_reason is None
    topology.set_loadflow_reference(loadflow_result_reference)

    optimizer_data.session.add(topology)

    message_topology = Topology(
        actions=topology.actions,
        pst_setpoints=topology.pst_setpoints,
        disconnections=topology.disconnections,
        loadflow_results=loadflow_result_reference,
        metrics=scoring_result.metrics,
    )

    optimizer_data.session.commit()

    return message_topology, scoring_result

send_topology_result

send_topology_result(
    topology_message,
    scoring_result,
    epoch,
    enable_ac_rejection,
    send_result_fn,
)

Send the persisted topology result to the result topic.

This function sends either a TopologyPushResult or a TopologyRejectionResult depending on the scoring result and the AC rejection settings.

PARAMETER	DESCRIPTION
topology_message	The topology message to send, containing the actions, metrics and loadflow result reference TYPE: Topology
scoring_result	The scoring result for the topology, containing the metrics, loadflow results and rejection reason if any. The rejection reason will be used to determine whether to send a push result or a rejection result TYPE: TopologyScoringResult
epoch	The current epoch number. TYPE: int
enable_ac_rejection	Whether to enable AC rejection. If True, topologies with a rejection reason in the scoring result will be sent as rejections. If False, all topologies will be sent as push results regardless of the scoring result. TYPE: bool
send_result_fn	The function to send results to the result topic, used for sending either the push result or the rejection result depending on the scoring result and AC rejection settings. TYPE: Callable[[ResultUnion], None]

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/ac/optimizer.py

def send_topology_result(
    topology_message: Topology,
    scoring_result: TopologyScoringResult,
    epoch: int,
    enable_ac_rejection: bool,
    send_result_fn: Callable[[ResultUnion], None],
) -> None:
    """Send the persisted topology result to the result topic.

    This function sends either a TopologyPushResult or a TopologyRejectionResult depending on the scoring result and the AC
    rejection settings.

    Parameters
    ----------
    topology_message : Topology
        The topology message to send, containing the actions, metrics and loadflow result reference
    scoring_result : TopologyScoringResult
        The scoring result for the topology, containing the metrics, loadflow results and rejection reason if any. The
        rejection reason will be used to determine whether to send a push result or a rejection result
    epoch : int
        The current epoch number.
    enable_ac_rejection : bool
        Whether to enable AC rejection. If True, topologies with a rejection reason in the scoring result will be sent as
        rejections.
        If False, all topologies will be sent as push results regardless of the scoring result.
    send_result_fn : Callable[[ResultUnion], None]
        The function to send results to the result topic, used for sending either the push result or the rejection result
        depending on the scoring result and AC rejection settings.
    """
    rejection_reason = scoring_result.rejection_reason
    local_enable_ac_rejection = enable_ac_rejection or (
        rejection_reason is not None and rejection_reason.criterion == "error"
    )

    if not local_enable_ac_rejection or rejection_reason is None:
        send_result_fn(TopologyPushResult(strategy=Strategy(timesteps=[topology_message]), epoch=epoch))
    else:
        send_result_fn(
            TopologyRejectionResult(
                reason=rejection_reason,
                strategy=Strategy(timesteps=[topology_message]),
                epoch=epoch,
            )
        )

run_fast_failing_epoch

run_fast_failing_epoch(optimizer_data)

Run one epoch of fast-failing AC evaluation only.

This imports new DC topologies, pulls up to topology_batch_size strategies from the evolution function, evaluates them with the worst-k fast-failing scorer, and returns the evaluated strategies together with their early-stage scoring results.

PARAMETER	DESCRIPTION
optimizer_data	The optimizer data containing the evolution and fast-failing scoring functions. TYPE: OptimizerData

RETURNS	DESCRIPTION
list[ACOptimTopology]	The strategies that were pulled and evaluated in the fast-failing stage.
list[EarlyStoppingStageResult]	The corresponding fast-failing scoring results.

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/ac/optimizer.py

def run_fast_failing_epoch(
    optimizer_data: OptimizerData,
) -> tuple[list[ACOptimTopology], list[EarlyStoppingStageResult]]:
    """Run one epoch of fast-failing AC evaluation only.

    This imports new DC topologies, pulls up to ``topology_batch_size`` strategies from the
    evolution function, evaluates them with the worst-k fast-failing scorer, and returns the
    evaluated strategies together with their early-stage scoring results.

    Parameters
    ----------
    optimizer_data : OptimizerData
        The optimizer data containing the evolution and fast-failing scoring functions.

    Returns
    -------
    list[ACOptimTopology]
        The strategies that were pulled and evaluated in the fast-failing stage.
    list[EarlyStoppingStageResult]
        The corresponding fast-failing scoring results.
    """
    topologies = optimizer_data.evolution_fn()
    n_topologies = len(topologies)
    if n_topologies == 0:
        logger.debug("Evolution returned no new strategy")
        return [], []
    logger.debug("Running early-stopping batch", n_topologies=n_topologies)
    fast_failing_results = optimizer_data.worst_k_scoring_fn(topologies)

    return topologies, fast_failing_results

process_fast_failing_results

process_fast_failing_results(
    optimizer_data,
    topologies,
    fast_failing_results,
    epoch,
    send_result_fn,
)

Process the fast-failing stage, emitting only rejected strategies.

Rejected strategies are persisted immediately and emitted as rejections. Surviving strategies are returned so they can be fully evaluated on the remaining contingencies before any accepted push result is sent.

PARAMETER	DESCRIPTION
optimizer_data	The optimizer data containing the session and the loadflow storage function. TYPE: OptimizerData
topologies	The topologies that were evaluated in the fast-failing stage. TYPE: list[ACOptimTopology]
fast_failing_results	The corresponding fast-failing scoring results, containing the rejection reason if any. TYPE: list[EarlyStoppingStageResult]
epoch	The current epoch number, used for logging and for sending in the result messages. TYPE: int
send_result_fn	The function to send results to the result topic. TYPE: Callable[[ResultUnion], None]

RETURNS	DESCRIPTION
list[ACOptimTopology]	The topologies that passed the fast-failing stage and can proceed to the remaining-contingency evaluation.
list[EarlyStoppingStageResult]	The corresponding fast-failing scoring results for the surviving topologies.

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/ac/optimizer.py

def process_fast_failing_results(
    optimizer_data: OptimizerData,
    topologies: list[ACOptimTopology],
    fast_failing_results: list[EarlyStoppingStageResult],
    epoch: int,
    send_result_fn: Callable[[ResultUnion], None],
) -> tuple[list[ACOptimTopology], list[EarlyStoppingStageResult]]:
    """Process the fast-failing stage, emitting only rejected strategies.

    Rejected strategies are persisted immediately and emitted as rejections. Surviving strategies are
    returned so they can be fully evaluated on the remaining contingencies before any accepted push
    result is sent.

    Parameters
    ----------
    optimizer_data : OptimizerData
        The optimizer data containing the session and the loadflow storage function.
    topologies : list[ACOptimTopology]
        The topologies that were evaluated in the fast-failing stage.
    fast_failing_results : list[EarlyStoppingStageResult]
        The corresponding fast-failing scoring results, containing the rejection reason if any.
    epoch : int
        The current epoch number, used for logging and for sending in the result messages.
    send_result_fn : Callable[[ResultUnion], None]
        The function to send results to the result topic.

    Returns
    -------
    list[ACOptimTopology]
        The topologies that passed the fast-failing stage and can proceed to the remaining-contingency evaluation.
    list[EarlyStoppingStageResult]
        The corresponding fast-failing scoring results for the surviving topologies.
    """
    logger.debug("Validating early-stopping batch", n_topologies=len(topologies))

    survivor_topologies: list[ACOptimTopology] = []
    survivor_early_results: list[EarlyStoppingStageResult] = []
    for topology, early_stop_result in zip(topologies, fast_failing_results, strict=True):
        if early_stop_result.rejection_reason is None:
            # If the topology is not rejected we want to fully evaluate it later
            survivor_topologies.append(topology)
            survivor_early_results.append(early_stop_result)
            continue

        topology_message, persisted_result = persist_topology(
            topology=topology,
            scoring_result=early_stop_result,
            optimizer_data=optimizer_data,
        )
        send_topology_result(
            topology_message=topology_message,
            scoring_result=persisted_result,
            epoch=epoch,
            enable_ac_rejection=True,
            send_result_fn=send_result_fn,
        )
    return survivor_topologies, survivor_early_results

run_remaining_epoch

run_remaining_epoch(
    optimizer_data, topologies, early_stage_results
)

Run one epoch of remaining-contingency AC evaluation only.

This evaluates the full remaining-loadflow stage for strategies that already passed the fast-failing worst-k evaluation and returns the same strategies together with their final scoring results.

PARAMETER	DESCRIPTION
optimizer_data	The optimizer data containing the remaining-stage scoring function. TYPE: OptimizerData
topologies	The survivor topologies to evaluate in the remaining stage. TYPE: list[ACOptimTopology]
early_stage_results	The corresponding fast-failing stage results for each topology. TYPE: list[EarlyStoppingStageResult]

RETURNS	DESCRIPTION
list[ACOptimTopology]	The topologies that were evaluated in the remaining stage.
list[TopologyScoringResult]	The corresponding final scoring results.

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/ac/optimizer.py

def run_remaining_epoch(
    optimizer_data: OptimizerData,
    topologies: list[ACOptimTopology],
    early_stage_results: list[EarlyStoppingStageResult],
) -> tuple[list[ACOptimTopology], list[TopologyScoringResult]]:
    """Run one epoch of remaining-contingency AC evaluation only.

    This evaluates the full remaining-loadflow stage for strategies that already passed the
    fast-failing worst-k evaluation and returns the same strategies together with their final
    scoring results.

    Parameters
    ----------
    optimizer_data : OptimizerData
        The optimizer data containing the remaining-stage scoring function.
    topologies : list[ACOptimTopology]
        The survivor topologies to evaluate in the remaining stage.
    early_stage_results : list[EarlyStoppingStageResult]
        The corresponding fast-failing stage results for each topology.

    Returns
    -------
    list[ACOptimTopology]
        The topologies that were evaluated in the remaining stage.
    list[TopologyScoringResult]
        The corresponding final scoring results.
    """
    if not topologies:
        logger.debug("No topologies provided for remaining contingencies")
        return [], []

    logger.debug("Running remaining contingencies", survivor_count=len(topologies))
    full_results = optimizer_data.scoring_fn(topologies, early_stage_results)
    return topologies, full_results

process_remaining_results

process_remaining_results(
    optimizer_data,
    topologies,
    full_results,
    epoch,
    send_result_fn,
)

Process the results of the remaining-contingency stage and send the strategies to the result topic.

PARAMETER	DESCRIPTION
optimizer_data	The optimizer data containing the session and the loadflow storage function. TYPE: OptimizerData
topologies	The topologies that were evaluated in the remaining stage. TYPE: list[ACOptimTopology]
full_results	The corresponding full scoring results, containing the final metrics, loadflow results and rejection reason if any. TYPE: list[TopologyScoringResult]
epoch	The current epoch number, used for logging and for sending in the result messages. TYPE: int
send_result_fn	The function to send results to the result topic. TYPE: Callable[[ResultUnion], None]

RETURNS	DESCRIPTION
list[ACOptimTopology]	The topologies that were evaluated in the remaining stage, with their metrics and fitness updated based on the full scoring results.
list[TopologyScoringResult]	The corresponding full scoring results for each topology, which can be used for further processing in the next epoch, e.g. for the evolution function.

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/ac/optimizer.py

def process_remaining_results(
    optimizer_data: OptimizerData,
    topologies: list[ACOptimTopology],
    full_results: list[TopologyScoringResult],
    epoch: int,
    send_result_fn: Callable[[ResultUnion], None],
) -> tuple[list[ACOptimTopology], list[TopologyScoringResult]]:
    """Process the results of the remaining-contingency stage and send the strategies to the result topic.

    Parameters
    ----------
    optimizer_data : OptimizerData
        The optimizer data containing the session and the loadflow storage function.
    topologies : list[ACOptimTopology]
        The topologies that were evaluated in the remaining stage.
    full_results : list[TopologyScoringResult]
        The corresponding full scoring results, containing the final metrics,
        loadflow results and rejection reason if any.
    epoch : int
        The current epoch number, used for logging and for sending in the result messages.
    send_result_fn : Callable[[ResultUnion], None]
        The function to send results to the result topic.

    Returns
    -------
    list[ACOptimTopology]
        The topologies that were evaluated in the remaining stage,
        with their metrics and fitness updated based
        on the full scoring results.
    list[TopologyScoringResult]
        The corresponding full scoring results for each topology,
        which can be used for further processing in the next epoch, e.g. for the evolution function.
    """
    for topology, full_result in zip(topologies, full_results, strict=True):
        topology_message, persisted_result = persist_topology(
            topology=topology,
            scoring_result=full_result,
            optimizer_data=optimizer_data,
        )
        send_topology_result(
            topology_message=topology_message,
            scoring_result=persisted_result,
            epoch=epoch,
            enable_ac_rejection=optimizer_data.params.ga_config.enable_ac_rejection,
            send_result_fn=send_result_fn,
        )
    return topologies, full_results

evaluate_remaining_contingencies

evaluate_remaining_contingencies(
    send_result_fn,
    optimizer_data,
    epoch,
    survivor_topologies,
    survivor_early_results,
)

Evaluate the remaining contingencies for the strategies that passed the fast-failing stage.

PARAMETER	DESCRIPTION
survivor_topologies	The topologies that passed the fast-failing stage and need to be evaluated on the remaining contingencies. TYPE: list[ACOptimTopology]
survivor_early_results	The early results from the fast-failing stage for the topologies that passed. TYPE: list[EarlyStoppingStageResult]
send_result_fn	The function to send results to the result topic, used for sending the final results after evaluating the remaining contingencies. TYPE: Callable[[ResultUnion], None]
optimizer_data	The optimizer data containing the scoring function for the remaining contingencies and other necessary data for the evaluation. TYPE: OptimizerData
epoch	The current epoch number, used for logging and for sending in the result messages. TYPE: int

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/ac/optimizer.py

def evaluate_remaining_contingencies(
    send_result_fn: Callable[[ResultUnion], None],
    optimizer_data: OptimizerData,
    epoch: int,
    survivor_topologies: list[ACOptimTopology],
    survivor_early_results: list[EarlyStoppingStageResult],
) -> None:
    """Evaluate the remaining contingencies for the strategies that passed the fast-failing stage.

    Parameters
    ----------
    survivor_topologies : list[ACOptimTopology]
        The topologies that passed the fast-failing stage and need to be evaluated
        on the remaining contingencies.
    survivor_early_results : list[EarlyStoppingStageResult]
        The early results from the fast-failing stage for the topologies that passed.
    send_result_fn : Callable[[ResultUnion], None]
        The function to send results to the result topic, used for sending the final results
        after evaluating the remaining contingencies.
    optimizer_data : OptimizerData
        The optimizer data containing the scoring function for the remaining contingencies and other necessary data for the
        evaluation.
    epoch : int
        The current epoch number, used for logging and for sending in the result messages.
    """
    topologies, full_results = run_remaining_epoch(
        topologies=survivor_topologies,
        early_stage_results=survivor_early_results,
        optimizer_data=optimizer_data,
    )
    process_remaining_results(
        optimizer_data=optimizer_data,
        topologies=topologies,
        full_results=full_results,
        send_result_fn=send_result_fn,
        epoch=epoch,
    )
    logger.info(
        "Completed AC evaluation",
        accepted={sum(1 for result in full_results if result.rejection_reason is None)},
        rejected=[result.rejection_reason.criterion for result in full_results if result.rejection_reason is not None],
    )

toop_engine_topology_optimizer.ac.scoring_functions

Scoring functions for the AC optimizer - in this case this runs an N-1 and computes metrics for it

logger module-attribute

logger = get_logger(__name__)

ACScoringParameters dataclass

ACScoringParameters(
    reject_convergence_threshold,
    reject_overload_threshold,
    reject_critical_branch_threshold,
    base_case_id,
    early_stop_validation,
)

Parameters for ac scoring

This is a subset of all ac parameters and grouped to shorten the signature of the scoring and acceptance functions.

reject_convergence_threshold instance-attribute

reject_convergence_threshold

The rejection threshold for the convergence rate, i.e. the split case must have at most the same amount of non converging loadflows as the unsplit case or it will be rejected.

reject_overload_threshold instance-attribute

reject_overload_threshold

The rejection threshold for the overload energy improvement, i.e. the split case must have at least 5% lower overload energy than the unsplit case or it will be rejected.

reject_critical_branch_threshold instance-attribute

reject_critical_branch_threshold

The rejection threshold for the critical branches increase, i.e. the split case must have less than 10% more critical branches than the unsplit case or it will be rejected.

base_case_id instance-attribute

base_case_id

The base case id for the loadflow runner (used to separately compute the N-0 flows).

early_stop_validation instance-attribute

early_stop_validation

Whether to enable early stopping during the optimization process.

get_early_stopping_contingency_ids

get_early_stopping_contingency_ids(
    topology, base_case_id=None
)

Extract the contingency ids for early stopping from a list of ACOptimTopology strategies.

This function extracts the worst k contingency case ids from each topology's worst_k_contingency_cases attribute for each timestep. These ids are used to determine which contingencies to include in the N-1 analysis for early stopping.

PARAMETER	DESCRIPTION
topology	An ACOptimTopology object containing a worst_k_contingency_cases attribute with contingency case ids. TYPE: ACOptimTopology
base_case_id	An optional base case id to include in the early stopping subset. If provided, this will be added to the list of contingency case ids for each timestep. TYPE: Optional[str] DEFAULT: None

RETURNS	DESCRIPTION
Optional[list[str]]	A list of contingency case IDs, or None if any required metric is missing.

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/ac/scoring_functions.py

def get_early_stopping_contingency_ids(
    topology: ACOptimTopology,
    base_case_id: Optional[str] = None,
) -> Optional[list[str]]:
    """Extract the contingency ids for early stopping from a list of ACOptimTopology strategies.

    This function extracts the worst k contingency case ids from each topology's worst_k_contingency_cases
    attribute for each timestep. These ids are used to determine which contingencies to include in the N-1
    analysis for early stopping.

    Parameters
    ----------
    topology : ACOptimTopology
        An ACOptimTopology object containing a worst_k_contingency_cases attribute with contingency case ids.
    base_case_id : Optional[str]
        An optional base case id to include in the early stopping subset. If provided, this will be added to the
        list of contingency case ids for each timestep.

    Returns
    -------
    Optional[list[str]]
        A list of contingency case IDs, or None if any required metric is missing.
    """
    worst_k_contingency_cases = deepcopy(topology.worst_k_contingency_cases)
    if len(worst_k_contingency_cases) == 0:
        # TODO Does this make sense?
        # Shouldnt we just continue?
        logger.warning(
            f"No overload threshold or case ids found in the topologies worst contingency_cases: {worst_k_contingency_cases}"
        )
        return None
    if base_case_id is not None:
        worst_k_contingency_cases.append(base_case_id)
    return worst_k_contingency_cases

update_runner_nminus1

update_runner_nminus1(runner, nminus1_def, case_ids_all_t)

Update the N-1 definitions in the runners to only include the worst k contingencies.

This modifies the N-1 definitions in the runners to only include the contingencies at the given indices.

PARAMETER	DESCRIPTION
runner	The loadflow runner to update. TYPE: AbstractLoadflowRunner
nminus1_def	The original N-1 definition to use as a template. TYPE: Nminus1Definition
case_ids_all_t	A list of contingency ids for each runner, indicating which contingencies to keep in the N-1 definition. Each element should be an index to the contingencies in the original N-1 definition. TYPE: Sequence[Collection[str]]

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/ac/scoring_functions.py

def update_runner_nminus1(
    runner: AbstractLoadflowRunner,
    nminus1_def: Nminus1Definition,
    case_ids_all_t: Collection[str],
) -> None:
    """Update the N-1 definitions in the runners to only include the worst k contingencies.

    This modifies the N-1 definitions in the runners to only include the contingencies at the given indices.

    Parameters
    ----------
    runner : AbstractLoadflowRunner
        The loadflow runner to update.
    nminus1_def : Nminus1Definition
        The original N-1 definition to use as a template.
    case_ids_all_t : Sequence[Collection[str]]
        A list of contingency ids for each runner, indicating which contingencies to keep in the N-1 definition.
        Each element should be an index to the contingencies in the original N-1 definition.
    """
    case_indices = get_contingency_indices_from_ids(case_ids_all_t, n_minus1_definition=nminus1_def)
    contingencies = np.array(nminus1_def.contingencies)[list(case_indices)]
    n1_def_copy = nminus1_def.model_copy()
    n1_def_copy.contingencies = contingencies.tolist()
    runner.store_nminus1_definition(n1_def_copy)

compute_loadflow_and_metrics

compute_loadflow_and_metrics(
    runner, topology, base_case_id, cases_subset=None
)

Compute loadflow results and associated metrics for a given set of strategies.

This function runs loadflow simulations for each provided strategy using the specified runners, then computes additional metrics based on the simulation results.

PARAMETER	DESCRIPTION
runner	The loadflow runner to use for simulations. TYPE: AbstractLoadflowRunner
topology	The topology to evaluate. TYPE: ACOptimTopology
base_case_id	The base case identifier for the topology. Can be None. TYPE: Optional[str]
cases_subset	Subset of contingency cases to use for loadflow computation. If None, all available contingencies are used. TYPE: Optional[Collection[str]] DEFAULT: None

RETURNS	DESCRIPTION
lfs	The results of the loadflow simulations. TYPE: LoadflowResultsPolars
additional_info	Additional information for the actions taken in the topology. TYPE: Optional[AdditionalActionInfo]
metrics	Computed metrics for the topology. TYPE: Metrics

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/ac/scoring_functions.py

def compute_loadflow_and_metrics(
    runner: AbstractLoadflowRunner,
    topology: ACOptimTopology,
    base_case_id: Optional[str],
    cases_subset: Optional[Collection[str]] = None,
) -> tuple[LoadflowResultsPolars, Optional[AdditionalActionInfo], Metrics]:
    """Compute loadflow results and associated metrics for a given set of strategies.

    This function runs loadflow simulations for each provided strategy using the specified runners,
    then computes additional metrics based on the simulation results.

    Parameters
    ----------
    runner : AbstractLoadflowRunner
        The loadflow runner to use for simulations.
    topology : ACOptimTopology
        The topology to evaluate.
    base_case_id : Optional[str]
        The base case identifier for the topology. Can be None.
    cases_subset : Optional[Collection[str]]
        Subset of contingency cases to use for loadflow computation. If None, all available contingencies are used.

    Returns
    -------
    lfs : LoadflowResultsPolars
        The results of the loadflow simulations.
    additional_info : Optional[AdditionalActionInfo]
        Additional information for the actions taken in the topology.
    metrics : Metrics
        Computed metrics for the topology.
    """
    original_n_minus1_def = runner.get_nminus1_definition()
    if cases_subset is not None:
        update_runner_nminus1(runner, original_n_minus1_def, cases_subset)

    lfs, additional_info = compute_loadflow(
        actions=topology.actions,
        disconnections=topology.disconnections,
        pst_setpoints=topology.pst_setpoints,
        runner=runner,
    )
    metrics = compute_metrics_single_timestep(
        actions=topology.actions,
        disconnections=topology.disconnections,
        loadflow=lfs,
        additional_info=additional_info,
        base_case_id=base_case_id,
    )

    if cases_subset is not None:
        # Restore the original N-1 definitions in the runners
        runner.store_nminus1_definition(original_n_minus1_def)

    return lfs, additional_info, metrics

extract_switching_distance

extract_switching_distance(additional_info)

Extract the switching distance from the additional action info

PARAMETER	DESCRIPTION
additional_info	The additional action info containing the switching distance TYPE: AdditionalActionInfo

RETURNS	DESCRIPTION
int	The switching distance, or 0 if not available

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/ac/scoring_functions.py

def extract_switching_distance(additional_info: AdditionalActionInfo) -> int:
    """Extract the switching distance from the additional action info

    Parameters
    ----------
    additional_info : AdditionalActionInfo
        The additional action info containing the switching distance

    Returns
    -------
    int
        The switching distance, or 0 if not available
    """
    if isinstance(additional_info, RealizedTopology):
        return len(additional_info.reassignment_diff)
    if isinstance(additional_info, pd.DataFrame):
        return len(additional_info)
    raise ValueError("Unknown format for additional info")

compute_metrics_single_timestep

compute_metrics_single_timestep(
    actions,
    disconnections,
    loadflow,
    additional_info,
    base_case_id=None,
)

Compute the metrics for a single timestep

PARAMETER	DESCRIPTION
actions	The reconfiguration assignment for the timestep TYPE: list[int]
disconnections	The disconnections for the timestep TYPE: list[int]
loadflow	The loadflow results for the timestep, use select_timestep to get the results for a specific timestep TYPE: LoadflowResults
additional_info	Additional information about the actions taken, such as switching distance or other metrics. TYPE: Optional[AdditionalActionInfo]
base_case_id	The base case id from the nminus1 definition, to separate N-0 flows from N-1 TYPE: Optional[str] DEFAULT: None

RETURNS	DESCRIPTION
Metrics	The metrics for the timestep

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/ac/scoring_functions.py

def compute_metrics_single_timestep(
    actions: list[int],
    disconnections: list[int],
    loadflow: LoadflowResultsPolars,
    additional_info: Optional[AdditionalActionInfo],
    base_case_id: Optional[str] = None,
) -> Metrics:
    """Compute the metrics for a single timestep

    Parameters
    ----------
    actions : list[int]
        The reconfiguration assignment for the timestep
    disconnections : list[int]
        The disconnections for the timestep
    loadflow : LoadflowResults
        The loadflow results for the timestep, use select_timestep to get the results for a specific timestep
    additional_info : Optional[AdditionalActionInfo]
        Additional information about the actions taken, such as switching distance or other metrics.
    base_case_id: Optional[str]
        The base case id from the nminus1 definition, to separate N-0 flows from N-1

    Returns
    -------
    Metrics
        The metrics for the timestep
    """
    metrics = compute_metrics_lfs(loadflow_results=loadflow, base_case_id=base_case_id)
    metrics = {
        key: (0.0 if value is None else np.nan_to_num(value, nan=0, posinf=INF_FITNESS, neginf=-INF_FITNESS).item())
        for key, value in metrics.items()
    }
    non_successful_states = [
        ConvergenceStatus.FAILED.value,
        ConvergenceStatus.MAX_ITERATION_REACHED.value,
        ConvergenceStatus.NO_CALCULATION.value,
    ]
    metrics.update(
        {
            "split_subs": len(actions),
            "disconnected_branches": len(disconnections),
            "non_converging_loadflows": loadflow.converged.filter(pl.col("status").is_in(non_successful_states))
            .select(pl.len())
            .collect()
            .item(),
        }
    )
    if additional_info is not None:
        metrics["switching_distance"] = extract_switching_distance(additional_info)
    worst_k_contingent_cases = metrics.pop("worst_k_contingent_cases", None)
    return Metrics(
        fitness=metrics["overload_energy_n_1"], extra_scores=metrics, worst_k_contingent_cases=worst_k_contingent_cases
    )

compute_loadflow

compute_loadflow(
    actions, disconnections, pst_setpoints, runner
)

Compute the loadflow for a given strategy

PARAMETER	DESCRIPTION
actions	The reconfiguration actions for the timestep TYPE: list[int]
disconnections	The disconnections for the timestep TYPE: list[int]
pst_setpoints	The PST setpoints for the topology, or None if PST taps are not part of the topology. TYPE: Optional[list[int]]
runner	The loadflow runner to use TYPE: AbstractLoadflowRunner

RETURNS	DESCRIPTION
LoadflowResultsPolars	The loadflow results for all timesteps in the strategy
list[Optional[AdditionalActionInfo]]	Additional information about the actions taken, such as switching distance or other metrics. The length of the list is n_timesteps.

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/ac/scoring_functions.py

def compute_loadflow(
    actions: list[int],
    disconnections: list[int],
    pst_setpoints: Optional[list[int]],
    runner: AbstractLoadflowRunner,
) -> tuple[LoadflowResultsPolars, Optional[AdditionalActionInfo]]:
    """Compute the loadflow for a given strategy

    Parameters
    ----------
    actions : list[int]
        The reconfiguration actions for the timestep
    disconnections : list[int]
        The disconnections for the timestep
    pst_setpoints : Optional[list[int]]
        The PST setpoints for the topology, or None if PST taps are not part of the topology.
    runner : AbstractLoadflowRunner
        The loadflow runner to use

    Returns
    -------
    LoadflowResultsPolars
        The loadflow results for all timesteps in the strategy
    list[Optional[AdditionalActionInfo]]
        Additional information about the actions taken, such as switching distance or other metrics. The length of
        the list is n_timesteps.
    """
    loadflow = runner.run_ac_loadflow(actions, disconnections, pst_setpoints)
    additional_information = runner.get_last_action_info()

    return loadflow, additional_information

evaluate_acceptance

evaluate_acceptance(
    metrics_split,
    metrics_unsplit,
    reject_convergence_threshold=1.0,
    reject_overload_threshold=0.95,
    reject_critical_branch_threshold=1.1,
    early_stopping=False,
)

Evaluate if the split loadflow results are acceptable compared to the unsplit results.

Compares the unsplit metrics * the thresholds to the split metrics. If all split metrics are better than the unsplit metrics * thresholds, the split results are accepted.

Checked metrics are: non_converging_loadflows: the number of non-converging loadflows should be less than or equal to reject_convergence_threshold * unsplit.extra_scores.get("non_converging_loadflows", 0) overload_energy_n_1: the overload energy should be less than or equal to reject_overload_threshold * unsplit.extra_scores.get("overload_energy_n_1", 0) critical_branch_count_n_1: the number of critical branches should be less than or equal to reject_critical_branch_threshold * unsplit.extra_scores.get("critical_branch_count_n_1", 0) TODO: Check Voltage Jumps between N0 and N1

PARAMETER	DESCRIPTION
metrics_split	The metrics for the split case. TYPE: Metrics
metrics_unsplit	The metrics for the unsplit case. TYPE: Metrics
reject_convergence_threshold	The threshold for the convergence rate, by default 1. (i.e. the split case must have at most the same amount of nonconverging loadflows as the unsplit case.) TYPE: float DEFAULT: 1.0
reject_overload_threshold	The threshold for the overload energy improvement, by default 0.95 (i.e. the split case must have at least 5% lower overload energy than the unsplit case). TYPE: float DEFAULT: 0.95
reject_critical_branch_threshold	The threshold for the critical branches increase, by default 1.1 (i.e. the split case must not have more than 110 % of the critical branches in the unsplit case). TYPE: float DEFAULT: 1.1
early_stopping	Whether the acceptance is computed as part of an early stopping criterion, will set the early_stopping field in the TopologyRejectionReason TYPE: bool DEFAULT: False

RETURNS	DESCRIPTION
Optional[TopologyRejectionReason]	A TopologyRejectionReason if the split results are rejected, None if accepted.

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/ac/scoring_functions.py

def evaluate_acceptance(
    metrics_split: Metrics,
    metrics_unsplit: Metrics,
    reject_convergence_threshold: float = 1.0,
    reject_overload_threshold: float = 0.95,
    reject_critical_branch_threshold: float = 1.1,
    early_stopping: bool = False,
) -> Optional[TopologyRejectionReason]:
    """Evaluate if the split loadflow results are acceptable compared to the unsplit results.

    Compares the unsplit metrics * the thresholds to the split metrics. If all split metrics are better than
    the unsplit metrics * thresholds, the split results are accepted.

    Checked metrics are:
        non_converging_loadflows: the number of non-converging loadflows should be less than or equal to
            reject_convergence_threshold * unsplit.extra_scores.get("non_converging_loadflows", 0)
        overload_energy_n_1: the overload energy should be less than or equal to
            reject_overload_threshold * unsplit.extra_scores.get("overload_energy_n_1", 0)
        critical_branch_count_n_1: the number of critical branches should be less than or equal
            to reject_critical_branch_threshold * unsplit.extra_scores.get("critical_branch_count_n_1", 0)
        TODO: Check Voltage Jumps between N0 and N1

    Parameters
    ----------
    metrics_split : Metrics
        The metrics for the split case.
    metrics_unsplit : Metrics
        The metrics for the unsplit case.
    reject_convergence_threshold : float, optional
        The threshold for the convergence rate, by default 1.
        (i.e. the split case must have at most the same amount of nonconverging loadflows as the unsplit case.)
    reject_overload_threshold : float, optional
        The threshold for the overload energy improvement, by default 0.95
        (i.e. the split case must have at least 5% lower overload energy than the unsplit case).
    reject_critical_branch_threshold : float, optional
        The threshold for the critical branches increase, by default 1.1
        (i.e. the split case must not have more than 110 % of the critical branches in the unsplit case).
    early_stopping : bool, optional
        Whether the acceptance is computed as part of an early stopping criterion, will set the early_stopping field in the
        TopologyRejectionReason

    Returns
    -------
    Optional[TopologyRejectionReason]
        A TopologyRejectionReason if the split results are rejected, None if accepted.
    """
    n_non_converged_unsplit = np.array(
        [
            metrics_unsplit.extra_scores.get("non_converging_loadflows", 0)
            - metrics_unsplit.extra_scores.get("disconnected_branches", 0)
        ]
    )
    n_non_converged_split = np.array(
        [
            metrics_split.extra_scores.get("non_converging_loadflows", 0)
            - metrics_split.extra_scores.get("disconnected_branches", 0)
        ]
    )
    convergence_acceptable = np.all(n_non_converged_split <= n_non_converged_unsplit * reject_convergence_threshold)
    if not convergence_acceptable:
        return TopologyRejectionReason(
            criterion="convergence",
            value_after=float(n_non_converged_split.sum()),
            value_before=float(n_non_converged_unsplit.sum()),
            threshold=reject_convergence_threshold,
            early_stopping=early_stopping,
        )

    unsplit_overload = np.array([metrics_unsplit.extra_scores.get("overload_energy_n_1", 0)])
    split_overload = np.array([metrics_split.extra_scores.get("overload_energy_n_1", 99999)])
    overload_improvement = np.all(split_overload <= unsplit_overload * reject_overload_threshold)
    if not overload_improvement:
        return TopologyRejectionReason(
            criterion="overload-energy",
            value_after=float(split_overload.sum()),
            value_before=float(unsplit_overload.sum()),
            threshold=reject_overload_threshold,
            early_stopping=early_stopping,
        )

    unsplit_critical_branches = np.array([metrics_unsplit.extra_scores.get("critical_branch_count_n_1", 999)], dtype=float)
    split_critical_branches = np.array([metrics_split.extra_scores.get("critical_branch_count_n_1", 0)], dtype=float)

    critical_branches_acceptable = np.all(
        split_critical_branches <= unsplit_critical_branches * reject_critical_branch_threshold
    )
    if not critical_branches_acceptable:
        return TopologyRejectionReason(
            criterion="critical-branch-count",
            value_after=float(split_critical_branches.sum()),
            value_before=float(unsplit_critical_branches.sum()),
            threshold=reject_critical_branch_threshold,
            early_stopping=early_stopping,
        )

    return None

compute_remaining_loadflows

compute_remaining_loadflows(
    runner,
    topology,
    base_case_id,
    loadflows_subset,
    cases_subset,
)

Compute the loadflows for the remaining contingencies that were not included in the early stopping subset.

This function is called after the early stopping loadflows have been computed and accepted. It computes the loadflows for the remaining contingencies that were not included in the early stopping subset, and then computes the metrics for the full set of loadflows.

PARAMETER	DESCRIPTION
runner	The loadflow runners to use, length n_timesteps. TYPE: AbstractLoadflowRunner
topology	The topology to score, length n_timesteps TYPE: ACOptimTopology
base_case_id	The base case id for the loadflow runners, used to separately compute the N-0 flows. TYPE: Optional[str]
loadflows_subset	The loadflow results for the early stopping subset, used to avoid recomputing these loadflows. TYPE: LoadflowResultsPolars
cases_subset	The contingency case ids that were included in the early stopping subset for each timestep. This could be extracted from the loadflows_subset but as it is available it is faster to pass it in. TYPE: list[str]

RETURNS	DESCRIPTION
LoadflowResultsPolars	The loadflow results for all contingencies, including those from the early stopping subset.
Metrics	The metrics for the full set of loadflows.

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/ac/scoring_functions.py

def compute_remaining_loadflows(
    runner: AbstractLoadflowRunner,
    topology: ACOptimTopology,
    base_case_id: Optional[str],
    loadflows_subset: LoadflowResultsPolars,
    cases_subset: list[str],
) -> tuple[LoadflowResultsPolars, Metrics]:
    """Compute the loadflows for the remaining contingencies that were not included in the early stopping subset.

    This function is called after the early stopping loadflows have been computed and accepted. It computes the loadflows
    for the remaining contingencies that were not included in the early stopping subset, and then computes the metrics for
    the full set of loadflows.

    Parameters
    ----------
    runner : AbstractLoadflowRunner
        The loadflow runners to use, length n_timesteps.
    topology : ACOptimTopology
        The topology to score, length n_timesteps
    base_case_id : Optional[str]
        The base case id for the loadflow runners, used to separately compute the N-0 flows.
    loadflows_subset : LoadflowResultsPolars
        The loadflow results for the early stopping subset, used to avoid recomputing these loadflows.
    cases_subset : list[str]
        The contingency case ids that were included in the early stopping subset for each timestep. This could be extracted
        from the loadflows_subset but as it is available it is faster to pass it in.

    Returns
    -------
    LoadflowResultsPolars
        The loadflow results for all contingencies, including those from the early stopping subset.
    Metrics
        The metrics for the full set of loadflows.
    """
    original_n_minus1_def = runner.get_nminus1_definition()
    all_cases = [contingency.id for contingency in original_n_minus1_def.contingencies]

    # Remove the already computed contingencies so we do not re-compute them
    remaining_cases = set(all_cases) - set(cases_subset)

    logger.info(f"Running N-1 analysis with {len(remaining_cases)} non-critical contingencies per timestep.")
    update_runner_nminus1(runner, original_n_minus1_def, remaining_cases)

    lfs_remaining, additional_info_remaining = compute_loadflow(
        actions=topology.actions,
        disconnections=topology.disconnections,
        pst_setpoints=topology.pst_setpoints,
        runner=runner,
    )

    lfs = concatenate_loadflow_results_polars([loadflows_subset, lfs_remaining])

    # We can pass the additional info from either critical or non critical contingencies as they are the same
    metrics = compute_metrics_single_timestep(
        actions=topology.actions,
        disconnections=topology.disconnections,
        loadflow=lfs,
        additional_info=additional_info_remaining,
        base_case_id=base_case_id,
    )

    # Restore the original N-1 definitions in the runners
    runner.store_nminus1_definition(original_n_minus1_def)

    return lfs, metrics

score_strategy_worst_k

score_strategy_worst_k(
    topology,
    runner,
    loadflow_results_unsplit,
    metrics_unsplit,
    scoring_params,
)

Evaluate only the worst-k stage for a single strategy.

PARAMETER	DESCRIPTION
topology	The topology to evaluate, length n_timesteps TYPE: ACOptimTopology
runner	The loadflow runner to use for the evaluation of the strategy TYPE: AbstractLoadflowRunner
loadflow_results_unsplit	The loadflow results for the unsplit case, used for comparison in the acceptance evaluation. TYPE: LoadflowResultsPolars
metrics_unsplit	The metrics for the unsplit case, used for comparison in the acceptance evaluation. TYPE: Metrics
scoring_params	The parameters for scoring, including thresholds for acceptance and early stopping settings. TYPE: ACScoringParameters

RETURNS	DESCRIPTION
EarlyStoppingStageResult	The result of the worst-k stage evaluation, including loadflow results, metrics and rejection reason if rejected. If early stopping is enabled and the strategy is rejected based on the worst-k contingencies, the early_stopping flag will be set.

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/ac/scoring_functions.py

def score_strategy_worst_k(
    topology: ACOptimTopology,
    runner: AbstractLoadflowRunner,
    loadflow_results_unsplit: LoadflowResultsPolars,
    metrics_unsplit: Metrics,
    scoring_params: ACScoringParameters,
) -> EarlyStoppingStageResult:
    """Evaluate only the worst-k stage for a single strategy.

    Parameters
    ----------
    topology : ACOptimTopology
        The topology to evaluate, length n_timesteps
    runner : AbstractLoadflowRunner
        The loadflow runner to use for the evaluation of the strategy
    loadflow_results_unsplit : LoadflowResultsPolars
        The loadflow results for the unsplit case, used for comparison in the acceptance evaluation.
    metrics_unsplit : Metrics
        The metrics for the unsplit case, used for comparison in the acceptance evaluation.
    scoring_params : ACScoringParameters
        The parameters for scoring, including thresholds for acceptance and early stopping settings.

    Returns
    -------
    EarlyStoppingStageResult
        The result of the worst-k stage evaluation, including loadflow results, metrics
        and rejection reason if rejected. If early stopping is enabled and the strategy
        is rejected based on the worst-k contingencies, the early_stopping flag will be set.
    """
    if scoring_params.early_stop_validation:
        cases_subset = get_early_stopping_contingency_ids(topology, base_case_id=scoring_params.base_case_id)
        assert cases_subset is not None, (
            "Early stopping enabled but no contingency case ids found for early stopping."
            "This might happen when the DC optimizer pushes topologies without worst_k entries."
        )
        lfs_early_stop, additional_info, metrics_early_stop = compute_loadflow_and_metrics(
            runner=runner,
            topology=topology,
            base_case_id=scoring_params.base_case_id,
            cases_subset=cases_subset,
        )
        lfs_early_stop_unsplit = subset_contingencies_polars(loadflow_results_unsplit, cases_subset)
        metrics_early_stop_unsplit = compute_metrics_single_timestep(
            actions=topology.actions,
            disconnections=topology.disconnections,
            loadflow=lfs_early_stop_unsplit,
            additional_info=additional_info,
            base_case_id=scoring_params.base_case_id,
        )
        rejection_reason = evaluate_acceptance(
            metrics_split=metrics_early_stop,
            metrics_unsplit=metrics_early_stop_unsplit,
            reject_convergence_threshold=scoring_params.reject_convergence_threshold,
            reject_overload_threshold=scoring_params.reject_overload_threshold,
            reject_critical_branch_threshold=scoring_params.reject_critical_branch_threshold,
            early_stopping=True,
        )
        return EarlyStoppingStageResult(
            loadflow_results=lfs_early_stop,
            metrics=metrics_early_stop,
            rejection_reason=rejection_reason,
            cases_subset=cases_subset,
        )

    lfs, _, metrics = compute_loadflow_and_metrics(
        runner=runner,
        topology=topology,
        base_case_id=scoring_params.base_case_id,
    )
    rejection_reason = evaluate_acceptance(
        metrics_split=metrics,
        metrics_unsplit=metrics_unsplit,
        reject_convergence_threshold=scoring_params.reject_convergence_threshold,
        reject_overload_threshold=scoring_params.reject_overload_threshold,
        reject_critical_branch_threshold=scoring_params.reject_critical_branch_threshold,
        early_stopping=False,
    )
    return EarlyStoppingStageResult(
        loadflow_results=lfs, metrics=metrics, rejection_reason=rejection_reason, cases_subset=None
    )

score_strategy_worst_k_batch

score_strategy_worst_k_batch(
    topologies,
    worst_k_runner_groups,
    loadflow_results_unsplit,
    metrics_unsplit,
    scoring_params,
)

Evaluate the worst-k stage for a batch of strategies.

PARAMETER	DESCRIPTION
topologies	The topologies to evaluate, length n_strategies. TYPE: list[ACOptimTopology]
worst_k_runner_groups	The loadflow runner groups to use for the evaluation of the strategies, length n_strategies. TYPE: RunnerGroup
loadflow_results_unsplit	The loadflow results for the unsplit case, used for comparison in the acceptance evaluation. TYPE: LoadflowResultsPolars
metrics_unsplit	The metrics for the unsplit case, used for comparison in the acceptance evaluation. TYPE: Metrics
scoring_params	The parameters for scoring, including thresholds for acceptance and early stopping settings. TYPE: ACScoringParameters

RETURNS	DESCRIPTION
list[EarlyStoppingStageResult]	The results of the worst-k stage evaluation for each strategy, including loadflow results, metrics and rejection reason if rejected. If early stopping is enabled and a strategy is rejected based on the worst-k contingencies, the early_stopping flag will be set in the rejection reason.

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/ac/scoring_functions.py

def score_strategy_worst_k_batch(
    topologies: list[ACOptimTopology],
    worst_k_runner_groups: RunnerGroup,
    loadflow_results_unsplit: LoadflowResultsPolars,
    metrics_unsplit: Metrics,
    scoring_params: ACScoringParameters,
) -> list[EarlyStoppingStageResult]:
    """Evaluate the worst-k stage for a batch of strategies.

    Parameters
    ----------
    topologies : list[ACOptimTopology]
        The topologies to evaluate, length n_strategies.
    worst_k_runner_groups : RunnerGroup
        The loadflow runner groups to use for the evaluation of the strategies, length n_strategies.
    loadflow_results_unsplit : LoadflowResultsPolars
        The loadflow results for the unsplit case, used for comparison in the acceptance evaluation.
    metrics_unsplit : Metrics
        The metrics for the unsplit case, used for comparison in the acceptance evaluation.
    scoring_params : ACScoringParameters
        The parameters for scoring, including thresholds for acceptance and early stopping settings.

    Returns
    -------
    list[EarlyStoppingStageResult]
        The results of the worst-k stage evaluation for each strategy, including loadflow results, metrics
        and rejection reason if rejected. If early stopping is enabled and a strategy is rejected based
        on the worst-k contingencies, the early_stopping flag will be set in the rejection reason.
    """
    if not topologies:
        return []
    if len(topologies) > len(worst_k_runner_groups):
        raise ValueError("Not enough worst-k runner groups configured for the requested strategy batch")

    worst_stage_results: list[Optional[EarlyStoppingStageResult]] = [
        _error_result_for_topology("Initial error", early_stopping=True)
    ] * len(topologies)
    with ThreadPoolExecutor(max_workers=len(topologies)) as executor:
        future_to_index = {
            executor.submit(
                score_strategy_worst_k,
                topology,
                runner,
                loadflow_results_unsplit,
                metrics_unsplit,
                scoring_params,
            ): index
            for index, (topology, runner) in enumerate(
                zip(topologies, worst_k_runner_groups[: len(topologies)], strict=True)
            )
        }
        for future in as_completed(future_to_index):
            index = future_to_index[future]
            try:
                worst_stage_results[index] = future.result()
            except Exception as exc:  # pragma: no cover - defensive guard
                logger.exception("Worst-k stage failed")
                final_result = _error_result_for_topology(str(exc), early_stopping=True)
                worst_stage_results[index] = EarlyStoppingStageResult(
                    loadflow_results=loadflow_results_unsplit,
                    metrics=final_result.metrics,
                    rejection_reason=final_result.rejection_reason,
                    cases_subset=None,
                )

    return [result for result in worst_stage_results if result is not None]

score_topology_remaining

score_topology_remaining(
    topology,
    runner,
    metrics_unsplit,
    scoring_params,
    early_stage_result,
)

Evaluate the remaining contingencies for a surviving strategy.

This function is called for strategies that survived the worst-k stage (i.e. were not rejected based on the worst-k contingencies). It computes the loadflows for the remaining contingencies that were not included in the worst-k stage, and then computes the metrics for the full set of loadflows. Finally, it evaluates the acceptance of the full set of loadflows compared to the unsplit case.

PARAMETER	DESCRIPTION
topology	The topology to evaluate, length n_timesteps TYPE: ACOptimTopology
runner	The loadflow runner to use for the evaluation of the strategy TYPE: AbstractLoadflowRunner
metrics_unsplit	The metrics for the unsplit case, used for comparison in the acceptance evaluation. TYPE: Metrics
scoring_params	The parameters for scoring, including thresholds for acceptance and early stopping settings. TYPE: ACScoringParameters
early_stage_result	The result from the worst-k stage evaluation, including loadflow results, metrics and cases subset used for early stopping. This is used to compute the remaining loadflows and metrics without recomputing the early stopping subset. TYPE: EarlyStoppingStageResult

RETURNS	DESCRIPTION
TopologyScoringResult	The result of the full evaluation, including loadflow results, metrics and rejection reason if rejected.

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/ac/scoring_functions.py

def score_topology_remaining(
    topology: ACOptimTopology,
    runner: AbstractLoadflowRunner,
    metrics_unsplit: Metrics,
    scoring_params: ACScoringParameters,
    early_stage_result: EarlyStoppingStageResult,
) -> TopologyScoringResult:
    """Evaluate the remaining contingencies for a surviving strategy.

    This function is called for strategies that survived the worst-k stage
    (i.e. were not rejected based on the worst-k contingencies).
    It computes the loadflows for the remaining contingencies that were not included in the worst-k stage,
    and then computes the metrics for the full set of loadflows.
    Finally, it evaluates the acceptance of the full set of loadflows compared to the unsplit case.

    Parameters
    ----------
    topology : ACOptimTopology
        The topology to evaluate, length n_timesteps
    runner : AbstractLoadflowRunner
        The loadflow runner to use for the evaluation of the strategy
    metrics_unsplit : Metrics
        The metrics for the unsplit case, used for comparison in the acceptance evaluation.
    scoring_params : ACScoringParameters
        The parameters for scoring, including thresholds for acceptance and early stopping settings.
    early_stage_result : EarlyStoppingStageResult
        The result from the worst-k stage evaluation, including loadflow results, metrics and cases
        subset used for early stopping. This is used to compute the remaining loadflows and metrics
        without recomputing the early stopping subset.

    Returns
    -------
    TopologyScoringResult
        The result of the full evaluation, including loadflow results, metrics
        and rejection reason if rejected.
    """
    if scoring_params.early_stop_validation:
        assert early_stage_result.cases_subset is not None
        lfs, metrics = compute_remaining_loadflows(
            runner=runner,
            topology=topology,
            base_case_id=scoring_params.base_case_id,
            loadflows_subset=early_stage_result.loadflow_results,
            cases_subset=early_stage_result.cases_subset,
        )
    else:
        lfs = early_stage_result.loadflow_results
        metrics = early_stage_result.metrics

    rejection_reason = evaluate_acceptance(
        metrics_split=metrics,
        metrics_unsplit=metrics_unsplit,
        reject_convergence_threshold=scoring_params.reject_convergence_threshold,
        reject_overload_threshold=scoring_params.reject_overload_threshold,
        reject_critical_branch_threshold=scoring_params.reject_critical_branch_threshold,
        early_stopping=False,
    )
    return TopologyScoringResult(loadflow_results=lfs, metrics=metrics, rejection_reason=rejection_reason)

score_strategy_full

score_strategy_full(
    topology, runner, metrics_unsplit, scoring_params
)

Evaluate a strategy on the full set of contingencies in one pass.

PARAMETER	DESCRIPTION
topology	The topology to evaluate, length n_timesteps TYPE: ACOptimTopology
runner	The loadflow runner to use for the evaluation of the strategy TYPE: AbstractLoadflowRunner
metrics_unsplit	The metrics for the unsplit case, used for comparison in the acceptance evaluation. TYPE: Metrics
scoring_params	The parameters for scoring, including thresholds for acceptance and early stopping settings. TYPE: ACScoringParameters

RETURNS	DESCRIPTION
TopologyScoringResult	The result of the full evaluation, including loadflow results, metrics and rejection reason if rejected.

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/ac/scoring_functions.py

def score_strategy_full(
    topology: ACOptimTopology,
    runner: AbstractLoadflowRunner,
    metrics_unsplit: Metrics,
    scoring_params: ACScoringParameters,
) -> TopologyScoringResult:
    """Evaluate a strategy on the full set of contingencies in one pass.

    Parameters
    ----------
    topology : ACOptimTopology
        The topology to evaluate, length n_timesteps
    runner : AbstractLoadflowRunner
        The loadflow runner to use for the evaluation of the strategy
    metrics_unsplit : Metrics
        The metrics for the unsplit case, used for comparison in the acceptance evaluation.
    scoring_params : ACScoringParameters
        The parameters for scoring, including thresholds for acceptance and early stopping settings.

    Returns
    -------
    TopologyScoringResult
        The result of the full evaluation, including loadflow results, metrics and rejection reason if rejected.
    """
    lfs, _, metrics = compute_loadflow_and_metrics(
        runner=runner,
        topology=topology,
        base_case_id=scoring_params.base_case_id,
    )
    rejection_reason = evaluate_acceptance(
        metrics_split=metrics,
        metrics_unsplit=metrics_unsplit,
        reject_convergence_threshold=scoring_params.reject_convergence_threshold,
        reject_overload_threshold=scoring_params.reject_overload_threshold,
        reject_critical_branch_threshold=scoring_params.reject_critical_branch_threshold,
        early_stopping=False,
    )
    return TopologyScoringResult(loadflow_results=lfs, metrics=metrics, rejection_reason=rejection_reason)

score_strategy_full_batch

score_strategy_full_batch(
    topologies,
    runner_groups,
    metrics_unsplit,
    scoring_params,
)

Evaluate a batch of topologies on the full set of contingencies.

PARAMETER	DESCRIPTION
topologies	The topologies to evaluate, length n_strategies. TYPE: list[ACOptimTopology]
runner_groups	The loadflow runner groups to use for the evaluation of the strategies, length n_strategies. TYPE: RunnerGroup
metrics_unsplit	The metrics for the unsplit case, used for comparison in the acceptance evaluation. TYPE: Metrics
scoring_params	The parameters for scoring, including thresholds for acceptance and early stopping settings. TYPE: ACScoringParameters

RETURNS	DESCRIPTION
list[TopologyScoringResult]	The results of the full evaluation for each strategy, including loadflow results, metrics and rejection reason if rejected.

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/ac/scoring_functions.py

def score_strategy_full_batch(
    topologies: list[ACOptimTopology],
    runner_groups: RunnerGroup,
    metrics_unsplit: Metrics,
    scoring_params: ACScoringParameters,
) -> list[TopologyScoringResult]:
    """Evaluate a batch of topologies on the full set of contingencies.

    Parameters
    ----------
    topologies : list[ACOptimTopology]
        The topologies to evaluate, length n_strategies.
    runner_groups : RunnerGroup
        The loadflow runner groups to use for the evaluation of the strategies, length n_strategies.
    metrics_unsplit : Metrics
        The metrics for the unsplit case, used for comparison in the acceptance evaluation.
    scoring_params : ACScoringParameters
        The parameters for scoring, including thresholds for acceptance and early stopping settings.

    Returns
    -------
    list[TopologyScoringResult]
        The results of the full evaluation for each strategy, including loadflow results,
        metrics and rejection reason if rejected.
    """
    if not topologies:
        return []
    if len(runner_groups) == 0:
        raise ValueError("At least one runner group is required for full-contingency evaluation")

    results: list[Optional[TopologyScoringResult]] = [
        _error_result_for_topology("Initial error", early_stopping=False)
    ] * len(topologies)
    max_parallel = min(len(topologies), len(runner_groups))

    with ThreadPoolExecutor(max_workers=max_parallel) as executor:
        future_to_assignment: dict = {}
        next_topology_index = 0

        def submit(index: int, runner_index: int) -> None:
            topology_for_scoring = ACOptimTopology(**topologies[index].model_dump())
            future = executor.submit(
                score_strategy_full,
                topology_for_scoring,
                runner_groups[runner_index],
                metrics_unsplit,
                scoring_params,
            )
            future_to_assignment[future] = (index, runner_index)

        for runner_index in range(max_parallel):
            submit(next_topology_index, runner_index)
            next_topology_index += 1

        while future_to_assignment:
            future = next(as_completed(tuple(future_to_assignment)))
            index, runner_index = future_to_assignment.pop(future)
            try:
                results[index] = future.result()
            except Exception as exc:  # pragma: no cover - defensive guard
                logger.exception("Full-contingency stage failed")
                results[index] = _error_result_for_topology(str(exc), early_stopping=False)

            if next_topology_index < len(topologies):
                submit(next_topology_index, runner_index)
                next_topology_index += 1

    return [result for result in results if result is not None]

score_remaining_contingency_batch

score_remaining_contingency_batch(
    topologies,
    early_stage_results,
    runner_group,
    metrics_unsplit,
    scoring_params,
)

Evaluate the remaining contingencies for a batch of surviving topologies.

This function is called for strategies that survived the worst-k stage (i.e. were not rejected based on the worst-k contingencies). It computes the loadflows for the remaining contingencies that were not included in the worst-k stage, and then computes the metrics for the full set of loadflows.

PARAMETER	DESCRIPTION
topologies	The topologies to evaluate, length n_strategies. TYPE: list[ACOptimTopology]
early_stage_results	The results from the worst-k stage evaluation for each topology, including loadflow results, metrics and cases subset used for early stopping. This is used to compute the remaining loadflows and metrics without recomputing the early stopping subset. TYPE: list[EarlyStoppingStageResult]
runner_group	The loadflow runner group to use for the evaluation of the strategies, length n_strategies. TYPE: RunnerGroup
metrics_unsplit	The metrics for the unsplit case, used for comparison in the acceptance evaluation. TYPE: Metrics
scoring_params	The parameters for scoring, including thresholds for acceptance and early stopping settings. TYPE: ACScoringParameters

RETURNS	DESCRIPTION
list[TopologyScoringResult]	The results of the full evaluation for each strategy, including loadflow results, metrics and rejection reason if rejected.

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/ac/scoring_functions.py

def score_remaining_contingency_batch(
    topologies: list[ACOptimTopology],
    early_stage_results: list[EarlyStoppingStageResult],
    runner_group: RunnerGroup,
    metrics_unsplit: Metrics,
    scoring_params: ACScoringParameters,
) -> list[TopologyScoringResult]:
    """Evaluate the remaining contingencies for a batch of surviving topologies.

    This function is called for strategies that survived the worst-k stage
    (i.e. were not rejected based on the worst-k contingencies).
    It computes the loadflows for the remaining contingencies that were not included in the worst-k stage,
    and then computes the metrics for the full set of loadflows.

    Parameters
    ----------
    topologies : list[ACOptimTopology]
        The topologies to evaluate, length n_strategies.
    early_stage_results : list[EarlyStoppingStageResult]
        The results from the worst-k stage evaluation for each topology, including loadflow results, metrics
        and cases subset used for early stopping. This is used to compute the remaining loadflows and metrics
        without recomputing the early stopping subset.
    runner_group : RunnerGroup
        The loadflow runner group to use for the evaluation of the strategies, length n_strategies.
    metrics_unsplit : Metrics
        The metrics for the unsplit case, used for comparison in the acceptance evaluation.
    scoring_params : ACScoringParameters
        The parameters for scoring, including thresholds for acceptance and early stopping settings.

    Returns
    -------
    list[TopologyScoringResult]
        The results of the full evaluation for each strategy, including loadflow results,
        metrics and rejection reason if rejected.
    """
    if not topologies:
        return []
    if len(runner_group) == 0:
        raise ValueError("At least one remaining-stage runner group is required")
    if len(topologies) != len(early_stage_results):
        raise ValueError("Topologies and early-stage results must have the same length")

    results: list[Optional[TopologyScoringResult]] = [
        _error_result_for_topology("Initial error", early_stopping=False)
    ] * len(topologies)
    max_parallel = min(len(topologies), len(runner_group))

    with ThreadPoolExecutor(max_workers=max_parallel) as executor:
        future_to_assignment: dict = {}
        next_topology_index = 0

        def submit(index: int, runner_index: int) -> None:
            future = executor.submit(
                score_topology_remaining,
                topologies[index],
                runner_group[runner_index],
                metrics_unsplit,
                scoring_params,
                early_stage_results[index],
            )
            future_to_assignment[future] = (index, runner_index)

        for runner_index in range(max_parallel):
            submit(next_topology_index, runner_index)
            next_topology_index += 1

        while future_to_assignment:
            future = next(as_completed(tuple(future_to_assignment)))
            index, runner_index = future_to_assignment.pop(future)
            try:
                results[index] = future.result()
            except Exception as exc:  # pragma: no cover - defensive guard
                logger.exception("Remaining-contingency stage failed")
                results[index] = _error_result_for_topology(str(exc), early_stopping=False)

            if next_topology_index < len(topologies):
                submit(next_topology_index, runner_index)
                next_topology_index += 1

    return [result for result in results if result is not None]

score_topology_batch

score_topology_batch(
    topologies,
    runner_group,
    metrics_unsplit,
    scoring_params,
    early_stage_results=None,
)

Score a batch of topologies in two stages.

PARAMETER	DESCRIPTION
topologies	The topologies to be scored. TYPE: list[ACOptimTopology]
runner_group	The group of runners to use for scoring. TYPE: RunnerGroup
metrics_unsplit	The metrics to be used for scoring. TYPE: Metrics
scoring_params	The parameters for scoring, including thresholds for acceptance and early stopping settings. TYPE: ACScoringParameters
early_stage_results	The results from the early stage, by default None. TYPE: Optional[list[EarlyStoppingStageResult]] DEFAULT: None

RETURNS	DESCRIPTION
list[TopologyScoringResult]	The results of the full evaluation for each strategy, including loadflow results, metrics and rejection reason if rejected.

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/ac/scoring_functions.py

def score_topology_batch(
    topologies: list[ACOptimTopology],
    runner_group: RunnerGroup,
    metrics_unsplit: Metrics,
    scoring_params: ACScoringParameters,
    early_stage_results: Optional[list[EarlyStoppingStageResult]] = None,
) -> list[TopologyScoringResult]:
    """Score a batch of topologies in two stages.

    Parameters
    ----------
    topologies : list[ACOptimTopology]
        The topologies to be scored.
    runner_group : RunnerGroup
        The group of runners to use for scoring.
    metrics_unsplit : Metrics
        The metrics to be used for scoring.
    scoring_params : ACScoringParameters
        The parameters for scoring, including thresholds for acceptance and early stopping settings.
    early_stage_results : Optional[list[EarlyStoppingStageResult]], optional
        The results from the early stage, by default None.

    Returns
    -------
    list[TopologyScoringResult]
        The results of the full evaluation for each strategy, including loadflow results,
        metrics and rejection reason if rejected.
    """
    if early_stage_results is None:
        return score_strategy_full_batch(
            topologies=topologies,
            runner_groups=runner_group,
            metrics_unsplit=metrics_unsplit,
            scoring_params=scoring_params,
        )
    if len(topologies) != len(early_stage_results):
        raise ValueError("Topologies and early-stage results must have the same length")
    results = score_remaining_contingency_batch(
        topologies=topologies,
        early_stage_results=early_stage_results,
        runner_group=runner_group,
        metrics_unsplit=metrics_unsplit,
        scoring_params=scoring_params,
    )

    return results

toop_engine_topology_optimizer.ac.select_strategy

Selection strategy for AC optimization topologies.

select_strategy

select_strategy(
    rng,
    repertoire,
    candidates,
    interest_scorer,
    batch_size=1,
    lower_scores_are_better=False,
    filter_strategy=None,
)

Select promising unevaluated topologies from the candidate pool.

Make sure the repertoire only contains topologies with the right optimizer type and optimization id as select_strategy will not filter for this.

PARAMETER	DESCRIPTION
rng	The random number generator to use TYPE: Generator
repertoire	The filtered repertoire with all individuals of the optimization in all optimizer types TYPE: list[BaseDBTopology]
candidates	Candidates which have not yet been evaluated. For a pull operation this will only include DC candidates without an AC parent. TYPE: list[BaseDBTopology]
interest_scorer	The function to score the topologies in the candidate pool. The higher the score, the more interesting the topology is. Eventually, the topology will be selected with a probability proportional to its score. TYPE: Callable[[DataFrame], Series]
batch_size	Number of topologies to sample without replacement. If more topologies are requested than available candidates, all available candidates are returned. TYPE: int DEFAULT: 1
lower_scores_are_better	Whether lower values returned by interest_scorer should be preferred when converting scores into sampling weights. TYPE: bool DEFAULT: False
filter_strategy	Whether to filter the candidates based on discriminator, median or dominator filter. TYPE: Optional[FilterStrategy] DEFAULT: None

RETURNS	DESCRIPTION
list[BaseDBTopology]	The selected topologies sampled from the candidate pool. If no topology could be selected because the candidate pool is empty, return an empty list

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/ac/select_strategy.py

def select_strategy(
    rng: Rng,
    repertoire: list[BaseDBTopology],
    candidates: list[BaseDBTopology],
    interest_scorer: Callable[[pd.DataFrame], pd.Series],
    batch_size: int = 1,
    lower_scores_are_better: bool = False,
    filter_strategy: Optional[FilterStrategy] = None,
) -> list[BaseDBTopology]:
    """Select promising unevaluated topologies from the candidate pool.

    Make sure the repertoire only contains topologies with the right optimizer type and optimization id
    as select_strategy will not filter for this.

    Parameters
    ----------
    rng : Rng
        The random number generator to use
    repertoire : list[BaseDBTopology]
        The filtered repertoire with all individuals of the optimization in all optimizer types
    candidates : list[BaseDBTopology]
        Candidates which have not yet been evaluated. For a pull operation this will only include DC candidates without an
        AC parent.
    interest_scorer : Callable[[pd.DataFrame], pd.Series]
        The function to score the topologies in the candidate pool. The higher the score, the more
        interesting the topology is. Eventually, the topology will be selected with a probability
        proportional to its score.
    batch_size : int, optional
        Number of topologies to sample without replacement. If more topologies are requested than available
        candidates, all available candidates are returned.
    lower_scores_are_better : bool, optional
        Whether lower values returned by ``interest_scorer`` should be preferred when converting
        scores into sampling weights.
    filter_strategy : Optional[FilterStrategy], optional
        Whether to filter the candidates based on discriminator, median or dominator filter.

    Returns
    -------
    list[BaseDBTopology]
        The selected topologies sampled from the candidate pool.
        If no topology could be selected because the candidate pool is empty,
        return an empty list
    """
    if len(candidates) == 0 or batch_size <= 0:
        return []

    metrics = metrics_dataframe(candidates)
    if metrics.empty:
        return []

    if filter_strategy is not None:
        if filter_strategy.filter_dominator_metrics_target is not None:
            repo_metrics = metrics_dataframe(repertoire)
            discriminator_df = get_discriminator_df(
                repo_metrics[repo_metrics["optimizer_type"] == OptimizerType.AC.value],
                filter_strategy.filter_dominator_metrics_target,
            )
        else:
            discriminator_df = pd.DataFrame()
        metrics = filter_metrics_df(
            metrics_df=metrics,
            discriminator_df=discriminator_df,
            filter_strategy=filter_strategy,
        )
    if metrics.empty:
        return []

    scores = interest_scorer(metrics).astype(float).fillna(0.0).to_numpy(copy=True)
    if lower_scores_are_better and len(scores):
        scores = scores.max() - scores
    else:
        min_score = scores.min(initial=0.0)
        if min_score < 0.0:
            scores -= min_score
    if len(scores):
        scores += np.finfo(float).eps
    score_sum = scores.sum()
    if np.isclose(score_sum, 0.0):
        probabilities = np.full(len(metrics), 1.0 / len(metrics))
    else:
        probabilities = scores / score_sum

    sample_size = min(batch_size, len(metrics))
    selected_ids = rng.choice(metrics.index.to_numpy(), size=sample_size, replace=False, p=probabilities)
    selected_topologies_by_id = {topology.id: topology for topology in candidates}
    return [selected_topologies_by_id[int(topology_id)] for topology_id in np.atleast_1d(selected_ids).tolist()]

filter_metrics_df

filter_metrics_df(
    metrics_df, discriminator_df, filter_strategy
)

Get a mask for the metrics DataFrame that filters out rows based on discriminator and median masks.

This function applies a discriminator, median and dominator mask.

PARAMETER	DESCRIPTION
metrics_df	The DataFrame containing the metrics to filter. This is typically the DC repertoire from which new results shall be pulled. TYPE: DataFrame
discriminator_df	The DataFrame containing the discriminator metrics. These are topologies that have previously been AC validated TYPE: DataFrame
filter_strategy	The filter strategy to use for the optimization, used to filter out strategies based on the discriminator, median or dominator filter. TYPE: FilterStrategy

RETURNS	DESCRIPTION
DataFrame	The filtered metrics DataFrame with similar topologies removed. If all topologies are filtered out, the original metrics DataFrame is returned.

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/ac/select_strategy.py

def filter_metrics_df(
    metrics_df: pd.DataFrame,
    discriminator_df: pd.DataFrame,
    filter_strategy: FilterStrategy,
) -> pd.DataFrame:
    """Get a mask for the metrics DataFrame that filters out rows based on discriminator and median masks.

    This function applies a discriminator, median and dominator mask.

    Parameters
    ----------
    metrics_df : pd.DataFrame
        The DataFrame containing the metrics to filter. This is typically the DC repertoire from which new results shall
        be pulled.
    discriminator_df : pd.DataFrame
        The DataFrame containing the discriminator metrics. These are topologies that have previously been AC validated
    filter_strategy : FilterStrategy
        The filter strategy to use for the optimization,
        used to filter out strategies based on the discriminator, median or dominator filter.

    Returns
    -------
    pd.DataFrame
        The filtered metrics DataFrame with similar topologies removed.
        If all topologies are filtered out, the original metrics DataFrame is returned.
    """
    repertoire_mask = get_repertoire_filter_mask(
        metrics_df=metrics_df,
        discriminator_df=discriminator_df,
        filter_strategy=filter_strategy,
    )
    # make sure that the metrics_df is not empty after filtering
    if len(metrics_df[repertoire_mask]) != 0:
        metrics_df = metrics_df[repertoire_mask]
    return metrics_df

get_repertoire_filter_mask

get_repertoire_filter_mask(
    metrics_df, discriminator_df, filter_strategy
)

Get a mask for the metrics DataFrame that filters out rows based on discriminator and median masks.

This function applies a discriminator, median and dominator mask.

PARAMETER	DESCRIPTION
metrics_df	The DataFrame containing the metrics to filter. TYPE: DataFrame
discriminator_df	The DataFrame containing the discriminator metrics. TYPE: DataFrame
filter_strategy	The filter strategy to use for the optimization, used to filter out strategies based on the discriminator, median or dominator filter. TYPE: FilterStrategy

RETURNS	DESCRIPTION
Bool[ndarray, ' metrics_df.shape[0]']	A boolean mask where True indicates the row is not filtered out.

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/ac/select_strategy.py

def get_repertoire_filter_mask(
    metrics_df: pd.DataFrame,
    discriminator_df: pd.DataFrame,
    filter_strategy: FilterStrategy,
) -> np.ndarray:
    """Get a mask for the metrics DataFrame that filters out rows based on discriminator and median masks.

    This function applies a discriminator, median and dominator mask.

    Parameters
    ----------
    metrics_df : pd.DataFrame
        The DataFrame containing the metrics to filter.
    discriminator_df : pd.DataFrame
        The DataFrame containing the discriminator metrics.
    filter_strategy : FilterStrategy
        The filter strategy to use for the optimization,
        used to filter out strategies based on the discriminator, median or dominator filter.

    Returns
    -------
    Bool[np.ndarray, " metrics_df.shape[0]"]
        A boolean mask where True indicates the row is not filtered out.
    """
    # set the dominator metrics observed if not provided
    if (
        filter_strategy.filter_dominator_metrics_observed is None
        and filter_strategy.filter_dominator_metrics_target is not None
    ):
        # set a to target
        filter_strategy.filter_dominator_metrics_observed = filter_strategy.filter_dominator_metrics_target

    # get the discriminator filter mask if the config is provided
    if (
        not discriminator_df.empty
        and filter_strategy.filter_discriminator_metric_distances is not None
        and filter_strategy.filter_discriminator_metric_multiplier is not None
    ):
        discriminator_filter_mask = get_discriminator_mask(
            metrics_df=metrics_df,
            discriminator_df=discriminator_df,
            metric_distances=filter_strategy.filter_discriminator_metric_distances,
            metric_multiplier=filter_strategy.filter_discriminator_metric_multiplier,
        )
    else:
        discriminator_filter_mask = np.ones(len(metrics_df), dtype=bool)

    # get the median filter mask if the config is provided
    if filter_strategy.filter_median_metric is not None:
        median_filter_mask = get_median_mask(
            metrics_df=metrics_df,
            target_metrics=filter_strategy.filter_median_metric,
        )
    else:
        median_filter_mask = np.ones(len(metrics_df), dtype=bool)

    # apply the discriminator and median filter masks
    metrics_df_filtered = metrics_df[discriminator_filter_mask & median_filter_mask]

    # get the dominator filter mask if the config is provided
    if (
        filter_strategy.filter_dominator_metrics_target is not None
        and filter_strategy.filter_dominator_metrics_observed is not None
    ):
        dominator_filter_mask = get_dominator_mask(
            metrics_df=metrics_df_filtered,
            target_metrics=filter_strategy.filter_dominator_metrics_target,
            observed_metrics=filter_strategy.filter_dominator_metrics_observed,
        )
    else:
        dominator_filter_mask = np.ones(len(metrics_df_filtered), dtype=bool)

    # apply the dominator filter mask
    filtered_index = metrics_df_filtered[dominator_filter_mask].index

    return np.isin(metrics_df.index, filtered_index)

get_median_mask

get_median_mask(
    metrics_df, target_metrics, fitness_col="fitness"
)

Get a mask for fitness values below the median for each discrete value of the target metrics.

Note: expects the target metrics to be discrete values, not continuous.

PARAMETER	DESCRIPTION
metrics_df	The DataFrame containing the metrics to filter. TYPE: DataFrame
target_metrics	A list of metrics with discrete values to consider for filtering. example: ["split_subs"]. TYPE: Sequence[MetricType]
fitness_col	The column name that contains the fitness values. Defaults to "fitness". TYPE: Optional[OperationMetric] DEFAULT: 'fitness'

RETURNS	DESCRIPTION
filter_mask	A boolean mask where True indicates the row is not below the median for any of the target metrics. TYPE: Bool[ndarray, ' metrics_df.shape[0]']

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/ac/select_strategy.py

def get_median_mask(
    metrics_df: pd.DataFrame, target_metrics: Sequence[MetricType], fitness_col: Optional[OperationMetric] = "fitness"
) -> np.ndarray:
    """Get a mask for fitness values below the median for each discrete value of the target metrics.

    Note: expects the target metrics to be discrete values, not continuous.

    Parameters
    ----------
    metrics_df : pd.DataFrame
        The DataFrame containing the metrics to filter.
    target_metrics : Sequence[MetricType]
        A list of metrics with discrete values to consider for filtering.
        example: ["split_subs"].
    fitness_col : Optional[OperationMetric], optional
        The column name that contains the fitness values. Defaults to "fitness".

    Returns
    -------
    filter_mask : Bool[np.ndarray, " metrics_df.shape[0]"]
        A boolean mask where True indicates the row is not below the median for any of the target metrics.
    """
    filter_mask = np.zeros(len(metrics_df), dtype=bool)
    for target_metric in target_metrics:
        for discrete_value in metrics_df[target_metric].unique():
            col_mask = (metrics_df[target_metric] == discrete_value).values
            metric_df = metrics_df[col_mask][fitness_col]
            # remove median
            median_fitness = metric_df.median()
            filter_mask |= col_mask & (metrics_df[fitness_col] < median_fitness).values
    # return the filter mask, that removes rows below the median
    filter_mask = ~filter_mask
    return filter_mask

get_dominator_mask

get_dominator_mask(
    metrics_df,
    target_metrics,
    observed_metrics,
    fitness_col="fitness",
)

Get a mask for rows from a DataFrame that are dominated by other rows based on specified metrics.

A metric entry if there is any other metric entry with a better fitness, in respect to the distance to the original topology. The distance in measured by the metric, assuming that lower values are better.

The target metric is used to fix the discrete value for which the dominance is checked. The fitness column is used to determine the fitness of the rows. Each observed metric is checked against the minimum fitness of all discrete target values.

Intended use is target_metrics = ["switching_distance", "split_subs"] and observed_metrics = Any additional metric to the target metrics, e.g. "disconnections"

PARAMETER	DESCRIPTION
metrics_df	The DataFrame to filter. TYPE: DataFrame
target_metrics	A list of metrics to consider for dominance. A target metric is expected to have discrete values (e.g. not fitness, overload_energy, or max_flow) If None, defaults to ["switching_distance", "split_subs"]. TYPE: Sequence[MetricType]
observed_metrics	A list of metrics to observe for dominance. If None, defaults to ["switching_distance", "split_subs"]. TYPE: Sequence[MetricType]
fitness_col	The column name that contains the fitness values. Defaults to "fitness". Note: the values are expected to be negative, best fitness converges to zero. TYPE: Optional[OperationMetric] DEFAULT: 'fitness'

RETURNS	DESCRIPTION
filter_mask	A boolean mask where True indicates the row is not dominated by another row. TYPE: Bool[ndarray, ' metrics_df.shape[0]']

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/ac/select_strategy.py

def get_dominator_mask(
    metrics_df: pd.DataFrame,
    target_metrics: Sequence[MetricType],
    observed_metrics: Sequence[MetricType],
    fitness_col: Optional[OperationMetric] = "fitness",
) -> np.ndarray:
    """Get a mask for rows from a DataFrame that are dominated by other rows based on specified metrics.

    A metric entry if there is any other metric entry with a better fitness,
    in respect to the distance to the original topology.
    The distance in measured by the metric, assuming that lower values are better.

    The target metric is used to fix the discrete value for which the dominance is checked.
    The fitness column is used to determine the fitness of the rows.
    Each observed metric is checked against the minimum fitness of all discrete target values.

    Intended use is target_metrics = ["switching_distance", "split_subs"]
    and observed_metrics = Any additional metric to the target metrics, e.g. "disconnections"


    Parameters
    ----------
    metrics_df : pd.DataFrame
        The DataFrame to filter.
    target_metrics : Sequence[MetricType]
        A list of metrics to consider for dominance.
        A target metric is expected to have discrete values (e.g. not fitness, overload_energy, or max_flow)
        If None, defaults to ["switching_distance", "split_subs"].
    observed_metrics : Sequence[MetricType]
        A list of metrics to observe for dominance.
        If None, defaults to ["switching_distance", "split_subs"].
    fitness_col : Optional[OperationMetric], optional
        The column name that contains the fitness values. Defaults to "fitness".
        Note: the values are expected to be negative, best fitness converges to zero.

    Returns
    -------
    filter_mask : Bool[np.ndarray, " metrics_df.shape[0]"]
        A boolean mask where True indicates the row is not dominated by another row.

    """
    filter_mask = np.zeros(len(metrics_df), dtype=bool)
    for target_metric in target_metrics:
        for discrete_value in metrics_df[target_metric].unique():
            # get columns mask
            col_mask = (metrics_df[target_metric] == discrete_value).values
            if (col_mask & ~filter_mask).sum() == 0:
                # all the elements have been filtered out already
                # -> no element is dominated by an already eliminated element
                continue
            max_idx = metrics_df[col_mask & ~filter_mask][fitness_col].idxmax()

            # get fitness mask
            fitness_mask = (metrics_df[fitness_col] < metrics_df[fitness_col].loc[max_idx]).values

            # apply dominator condition
            for col in observed_metrics:
                if col == fitness_col:
                    continue
                filter_mask |= (metrics_df[col] > metrics_df[col].loc[max_idx]).values & fitness_mask & col_mask

    # retrun the filter mask, that removes dominated rows
    filter_mask = ~filter_mask

    return np.asarray(filter_mask)

get_discriminator_mask

get_discriminator_mask(
    metrics_df,
    discriminator_df,
    metric_distances,
    metric_multiplier=None,
)

Get a mask for rows in metrics_df that are within a certain distance from the discriminator_df.

The distance is defined by the metric_distances dictionary, which contains the metrics and their respective distances. Use the use_split_sub_multiplier flag to apply a multiplier in respect to the split_subs_col. e.g. use_split_sub_multiplier=False, the metric_distances is applied directly to the metrics_df. If use_split_sub_multiplier=True, the metric_distances are multiplied by the split_subs, leading to a larger distance for larger split_subs values.

PARAMETER	DESCRIPTION
metrics_df	The DataFrame containing the metrics to filter. TYPE: DataFrame
discriminator_df	The DataFrame containing the discriminator metrics. TYPE: DataFrame
metric_distances	A dictionary defining the metric distances for filtering. The keys are metric names and the values are sets of distances. example: metric_distances = { "split_subs": {0}, "switching_distance": {-0.9, 0.9}, "fitness": {-0.1, 0.1}, } Note: the fitness is treated as a percentage. TYPE: dict[str, set[float | int]]
metric_multiplier	A dictionary defining multiplier for the metric distances. The keys are metric names and the values are multipliers. If None, defaults to an empty dictionary. Multiple values are added by: distance_multiplier = ( metric_multiplier[metric1] * discriminator_df[metric1] + metric_multiplier[metric2] * discriminator_df[metric2] + ... ) example: {"split_subs": 0.5} The discriminator_df will be multiplied by the split_subs value. In the case of the metric_distances values will be multiplied by the split_subs value and the metric_multiplier. -> metric_distances["switching_distance"] * 0.5 * split_subs_col (e.g. 4 splits) -> the metric distance is increased by this metric multiplier. TYPE: Optional[dict[str, float | int]] DEFAULT: None

RETURNS	DESCRIPTION
filter_mask	A boolean mask where True indicates the row is not within the distance defined by the discriminator_df and metric_distances. TYPE: Bool[ndarray, ' metrics_df.shape[0]']

RAISES	DESCRIPTION
ValueError	If not all metric distances are present in the discriminator DataFrame columns. If the metric_distances is None, it will use default values based on the use_split_sub_multiplier flag.

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/ac/select_strategy.py

def get_discriminator_mask(
    metrics_df: pd.DataFrame,
    discriminator_df: pd.DataFrame,
    metric_distances: dict[str, set[float | int]],
    metric_multiplier: Optional[dict[str, float | int]] = None,
) -> np.ndarray:
    """Get a mask for rows in metrics_df that are within a certain distance from the discriminator_df.

    The distance is defined by the metric_distances dictionary, which contains the metrics and their respective distances.
    Use the `use_split_sub_multiplier` flag to apply a multiplier in respect to the split_subs_col.
    e.g. use_split_sub_multiplier=False, the metric_distances is applied directly to the metrics_df.
    If use_split_sub_multiplier=True, the metric_distances are multiplied by the split_subs, leading to a
    larger distance for larger split_subs values.

    Parameters
    ----------
    metrics_df : pd.DataFrame
        The DataFrame containing the metrics to filter.
    discriminator_df : pd.DataFrame
        The DataFrame containing the discriminator metrics.
    metric_distances : dict[str, set[float | int]]
        A dictionary defining the metric distances for filtering.
        The keys are metric names and the values are sets of distances.
        example:
        metric_distances = {
            "split_subs": {0},
            "switching_distance": {-0.9, 0.9},
            "fitness": {-0.1, 0.1},
        }
        Note: the fitness is treated as a percentage.
    metric_multiplier : Optional[dict[str, float | int]], optional
        A dictionary defining multiplier for the metric distances.
        The keys are metric names and the values are multipliers.
        If None, defaults to an empty dictionary.
        Multiple values are added by:
        distance_multiplier = (
           `metric_multiplier[metric1]` * `discriminator_df[metric1]` +
           `metric_multiplier[metric2]` * `discriminator_df[metric2]` + ...
        )
        example: {"split_subs": 0.5}
                 The discriminator_df will be multiplied by the split_subs value.
                 In the case of the metric_distances values will be multiplied by the
                 split_subs value and the metric_multiplier.
                 -> metric_distances["switching_distance"] * 0.5 * split_subs_col (e.g. 4 splits)
                 -> the metric distance is increased by this metric multiplier.


    Returns
    -------
    filter_mask : Bool[np.ndarray, " metrics_df.shape[0]"]
        A boolean mask where True indicates the row is not within the distance defined by the discriminator_df
        and metric_distances.

    Raises
    ------
    ValueError
        If not all metric distances are present in the discriminator DataFrame columns.
        If the metric_distances is None, it will use default values based on the use_split_sub_multiplier flag.
    """
    if metric_multiplier is None:
        metric_multiplier = {}

    if not set(metric_distances.keys()).issubset(discriminator_df.columns):
        raise ValueError(
            f"Not all metric distances {metric_distances.keys()} are "
            f"present in the discriminator DataFrame {discriminator_df.columns}. "
        )
    if not set(metric_multiplier.keys()).issubset(discriminator_df.columns):
        raise ValueError(
            f"Not all metric multipliers {metric_multiplier.keys()} are "
            f"present in the discriminator DataFrame {discriminator_df.columns}. "
        )

    discriminator_df["fitness_save"] = discriminator_df["fitness"].copy()  # save original fitness for later use
    discriminator_df["fitness"] = discriminator_df["fitness"].abs()  # ensure fitness is
    metrics_df["fitness_save"] = metrics_df["fitness"].copy()  # save original fitness for later use
    metrics_df["fitness"] = metrics_df["fitness"].abs()  # ensure fitness is positive for the discriminator

    # filter mask for discriminator
    filter_mask = np.zeros(len(metrics_df), dtype=bool)
    for _idx, row in discriminator_df.iterrows():
        mask_metrics = np.ones(len(metrics_df), dtype=bool)
        for metric, distance in metric_distances.items():
            # get distance multiplier
            distance_multiplier = 0
            for metric_multiplier_key, metric_multiplier_value in metric_multiplier.items():
                distance_multiplier += metric_multiplier_value * row[metric_multiplier_key]
            if distance_multiplier == 0:
                # if no multiplier is defined, fall back to 1
                distance_multiplier = 1

            # apply the distance to the metrics_df
            if metric != "fitness":
                min_condition = row[metric] + min(distance) * distance_multiplier
                max_condition = row[metric] + max(distance) * distance_multiplier

            else:
                min_condition = row[metric] * (1 + min(distance) * distance_multiplier)
                max_condition = row[metric] * (1 + max(distance) * distance_multiplier)

            mask_metrics = mask_metrics & (metrics_df[metric] >= min_condition) & (metrics_df[metric] <= max_condition)

        filter_mask += mask_metrics

    # restore original fitness
    metrics_df["fitness"] = metrics_df["fitness_save"]
    metrics_df.drop(columns=["fitness_save"], inplace=True)

    # return the filter mask, that removes discriminated rows
    filter_mask = ~filter_mask

    return np.asarray(filter_mask)

get_discriminator_df

get_discriminator_df(metrics_df, target_metrics)

Get a discriminator DataFrame from the metrics DataFrame.

The discriminator DataFrame is a subset of the metrics DataFrame that contains only the target metrics. It is used to filter out similar topologies from the metrics DataFrame.

PARAMETER	DESCRIPTION
metrics_df	The DataFrame containing the metrics to filter. Note: expects the metrics_df to contain only AC topologies that have as a metric the "fitness_dc" column. TYPE: DataFrame
target_metrics	A list of metrics to consider for filtering. TYPE: list[str]

RETURNS	DESCRIPTION
DataFrame	A DataFrame containing only the target metrics.

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/ac/select_strategy.py

def get_discriminator_df(metrics_df: pd.DataFrame, target_metrics: Sequence[str]) -> pd.DataFrame:
    """Get a discriminator DataFrame from the metrics DataFrame.

    The discriminator DataFrame is a subset of the metrics DataFrame that contains only the target metrics.
    It is used to filter out similar topologies from the metrics DataFrame.

    Parameters
    ----------
    metrics_df : pd.DataFrame
        The DataFrame containing the metrics to filter.
        Note: expects the metrics_df to contain only AC topologies
        that have as a metric the "fitness_dc" column.
    target_metrics : list[str]
        A list of metrics to consider for filtering.

    Returns
    -------
    pd.DataFrame
        A DataFrame containing only the target metrics.
    """
    if metrics_df.empty:
        return pd.DataFrame()
    discriminator_df = metrics_df[[*target_metrics, "fitness_dc"]]
    discriminator_df.rename(columns={"fitness_dc": "fitness"}, inplace=True)
    return discriminator_df

toop_engine_topology_optimizer.ac.storage

The database models to store topologies in the AC optimizer

ACOptimTopology

Bases: BaseDBTopology

Inherits from the base topology to make a database table for AC optimizer topologies

This can include both AC and DC topologies, with the specific needs of the AC optimizer

__tablename__ class-attribute instance-attribute

__tablename__ = 'ac_optim_topology'

parent_id class-attribute instance-attribute

parent_id = Field(
    foreign_key="ac_optim_topology.id",
    nullable=True,
    default=None,
)

The mutation parent id, i.e. the topology that this topology was mutated from. This is mostly for debugging purposes to see where a topology came from and understand the mutation process.

parent class-attribute instance-attribute

parent = Relationship()

The mutation parent

acceptance class-attribute instance-attribute

acceptance = Field(default=None)

Whether the strategy was accepted or not.

id class-attribute instance-attribute

id = Field(default=None, primary_key=True)

The table primary key

actions class-attribute instance-attribute

actions = Field(sa_type=JSON)

The branch/injection reconfiguration actions as indices into the action set.

disconnections class-attribute instance-attribute

disconnections = Field(sa_type=JSON)

A list of disconnections, indexing into the disconnectable branches set in the action set.

pst_setpoints class-attribute instance-attribute

pst_setpoints = Field(sa_type=JSON, default=None)

The setpoints for the PSTs if they have been computed. This is an index into the range of pst taps, i.e. the smallest tap is 0 and the neutral tap somewhere in the middle of the range. The tap range is defined in the action set. The list always has the same length, i.e. the number of controllable PSTs in the system, and each entry corresponds to the PST at the same position in the action set.

unsplit instance-attribute

unsplit

Whether all topologies in the strategy including this one have no branch assignments, disconnections or injections.

timestep instance-attribute

timestep

The timestep of this topology, starting at 0

strategy_hash instance-attribute

strategy_hash

The hash of the strategy - this hashes actions, disconnections and pst_setpoints for all timesteps in the strategy, making it possible to form a unique constraint on the strategy. This value will be set to the same for all topologies in the same strategy, furthermore making it possible to group timesteps.

strategy_hash_str property

strategy_hash_str

The strategy hash as a string, for human readability

optimization_id instance-attribute

optimization_id

The optimization ID this topology belongs to

optimizer_type instance-attribute

optimizer_type

Which optimizer created this topology

fitness instance-attribute

fitness

The fitness of this topology

metrics class-attribute instance-attribute

metrics = Field(default_factory=lambda: {}, sa_type=JSON)

The metrics of this topology

worst_k_contingency_cases class-attribute instance-attribute

worst_k_contingency_cases = Field(
    default_factory=lambda: [], sa_type=JSON
)

The worst k contingency case IDs for the topology.

created_at class-attribute instance-attribute

created_at = Field(default=now(), nullable=False)

The time the topology was recorded in the database

stored_loadflow_reference class-attribute instance-attribute

stored_loadflow_reference = None

The file reference for the loadflow results of this topology/strategy, if they were computed. Multiple topologies belonging to the same strategy will have the same serialized loadflow results object as there is a timestep notion in the loadflow results. To obtain the correct loadflow results, use the timestep attribute. This is stored as a json serialized StoredLoadflowReference object

__table_args__ class-attribute instance-attribute

__table_args__ = (
    UniqueConstraint(
        "optimization_id",
        "optimizer_type",
        "strategy_hash",
        "timestep",
        name="topo_unique",
    ),
)

get_loadflow_reference

get_loadflow_reference()

Get the loadflow reference as a StoredLoadflowReference object

RETURNS	DESCRIPTION
Optional[StoredLoadflowReference]	The loadflow reference, or None if it is not set

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/interfaces/models/base_storage.py

def get_loadflow_reference(self) -> Optional[StoredLoadflowReference]:
    """Get the loadflow reference as a StoredLoadflowReference object

    Returns
    -------
    Optional[StoredLoadflowReference]
        The loadflow reference, or None if it is not set
    """
    if self.stored_loadflow_reference is None:
        return None
    return StoredLoadflowReference.model_validate_json(self.stored_loadflow_reference)

set_loadflow_reference

set_loadflow_reference(loadflow_reference)

Set the loadflow reference from a StoredLoadflowReference object

PARAMETER	DESCRIPTION
loadflow_reference	The loadflow reference to set, or None to unset it TYPE: Optional[StoredLoadflowReference]

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/interfaces/models/base_storage.py

def set_loadflow_reference(self, loadflow_reference: Optional[StoredLoadflowReference]) -> "BaseDBTopology":
    """Set the loadflow reference from a StoredLoadflowReference object

    Parameters
    ----------
    loadflow_reference : Optional[StoredLoadflowReference]
        The loadflow reference to set, or None to unset it
    """
    if loadflow_reference is None:
        self.stored_loadflow_reference = None
    else:
        self.stored_loadflow_reference = loadflow_reference.model_dump_json()
    return self

FinishedOptimizations

Bases: SQLModel

A table to store finished optimizations, to prevent picking up old topologies from previous optimizations

__tablename__ class-attribute instance-attribute

__tablename__ = 'finished_optimizations'

optimization_id class-attribute instance-attribute

optimization_id = Field(primary_key=True)

The optimization ID of the finished optimization as sent via kafka

optimizer_type class-attribute instance-attribute

optimizer_type = Field(primary_key=True)

Which optimizer type has finished

finished_at class-attribute instance-attribute

finished_at = Field(default_factory=now)

When the optimization was marked as finished in the database

convert_single_topology

convert_single_topology(
    topology,
    optimization_id,
    optimizer_type,
    timestep,
    strategy_hash,
    unsplit,
)

Convert a single Topology to a ACOptimTopology

PARAMETER	DESCRIPTION
topology	The topology to convert TYPE: Topology
optimization_id	The optimization ID to assign to the db topology TYPE: str
optimizer_type	The optimizer type to assign to the db topology TYPE: OptimizerType
timestep	The timestep of the topology TYPE: int
strategy_hash	The hash of the strategy computed through hash_strategy TYPE: bytes
unsplit	Whether the strategy is unsplit TYPE: bool

RETURNS	DESCRIPTION
ACOptimTopology	The converted topology

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/ac/storage.py

def convert_single_topology(
    topology: Topology,
    optimization_id: str,
    optimizer_type: OptimizerType,
    timestep: int,
    strategy_hash: bytes,
    unsplit: bool,
) -> ACOptimTopology:
    """Convert a single Topology to a ACOptimTopology

    Parameters
    ----------
    topology : Topology
        The topology to convert
    optimization_id : str
        The optimization ID to assign to the db topology
    optimizer_type : OptimizerType
        The optimizer type to assign to the db topology
    timestep : int
        The timestep of the topology
    strategy_hash : bytes
        The hash of the strategy computed through hash_strategy
    unsplit : bool
        Whether the strategy is unsplit

    Returns
    -------
    ACOptimTopology
        The converted topology
    """
    return ACOptimTopology(
        actions=topology.actions,
        disconnections=topology.disconnections,
        pst_setpoints=topology.pst_setpoints,
        unsplit=unsplit,
        timestep=timestep,
        strategy_hash=strategy_hash,
        optimization_id=optimization_id,
        optimizer_type=optimizer_type,
        fitness=topology.metrics.fitness,
        metrics=topology.metrics.extra_scores,
        worst_k_contingency_cases=topology.metrics.worst_k_contingency_cases,
    ).set_loadflow_reference(topology.loadflow_results)

convert_message_topo_to_db_topo

convert_message_topo_to_db_topo(
    message_strategy, optimization_id, optimizer_type
)

Convert a TopologyPushResult to a list of ACOptimTopology

PARAMETER	DESCRIPTION
message_strategy	The strategy to convert, usually from a TopologyPushResult or OptimizationStartedResult. Strategy timesteps are flattened into the list of topologies that are being returned. TYPE: Strategy
optimization_id	The optimization ID to assign to the topologies. This was sent through with the parent result message and is not part of TopologyPushResult TYPE: str
optimizer_type	The optimizer type to assign. This was sent through with the parent result message and is not part of TopologyPushResult TYPE: OptimizerType

RETURNS	DESCRIPTION
list[ACOptimTopology]	A list of converted topologies where for each timestep a new ACOptimTopology instance is created.

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/ac/storage.py

def convert_message_topo_to_db_topo(
    message_strategy: Strategy, optimization_id: str, optimizer_type: OptimizerType
) -> list[ACOptimTopology]:
    """Convert a TopologyPushResult to a list of ACOptimTopology

    Parameters
    ----------
    message_strategy : Strategy
        The strategy to convert, usually from a TopologyPushResult or OptimizationStartedResult. Strategy timesteps
        are flattened into the list of topologies that are being returned.
    optimization_id : str
        The optimization ID to assign to the topologies. This was sent through with the parent
        result message and is not part of TopologyPushResult
    optimizer_type : OptimizerType
        The optimizer type to assign. This was sent through with the parent result message and is not
        part of TopologyPushResult

    Returns
    -------
    list[ACOptimTopology]
        A list of converted topologies where for each timestep a new
        ACOptimTopology instance is created.
    """
    converted = []
    # Hash the strategy to get a unique global identifier
    strategy_hash = hash_strategy(message_strategy)
    unsplit = is_unsplit_strategy(message_strategy)
    for time_id, timestep_topo in enumerate(message_strategy.timesteps):
        converted.append(
            convert_single_topology(
                topology=timestep_topo,
                optimization_id=optimization_id,
                optimizer_type=optimizer_type,
                timestep=time_id,
                strategy_hash=strategy_hash,
                unsplit=unsplit,
            )
        )
    return converted

create_session

create_session()

Create a thread-safe shared in-memory SQLite session for the AC worker.

Using StaticPool plus check_same_thread=False keeps the worker on a single shared in-memory database connection.

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/ac/storage.py

def create_session() -> Session:
    """Create a thread-safe shared in-memory SQLite session for the AC worker.

    Using ``StaticPool`` plus ``check_same_thread=False`` keeps the
    worker on a single shared in-memory database connection.
    """
    engine = create_engine(
        "sqlite://",
        connect_args={"check_same_thread": False},
        poolclass=StaticPool,
    )
    SQLModel.metadata.create_all(engine, tables=[ACOptimTopology.__table__, FinishedOptimizations.__table__])
    return Session(engine, expire_on_commit=False)

scrub_db

scrub_db(session, max_age_seconds=86400)

Scrub the database of the worker to prevent memory issues.

The AC worker has to store topologies from all runs, not only the current one, as it might have to pick up an optimization later. Meaning, the AC worker busily collects all topologies into memory. To prevent memory leak issues, we scrub the database of topologies that are older than a day as we do not expect to ever have to go back to those.

PARAMETER	DESCRIPTION
session	The database session. The session object will be modified in-place TYPE: Session
max_age_seconds	The maximum age of topologies to keep in the database, in seconds. Topologies older than this will be removed. TYPE: int DEFAULT: 86400

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/ac/storage.py

def scrub_db(session: Session, max_age_seconds: int = 86400) -> None:
    """Scrub the database of the worker to prevent memory issues.

    The AC worker has to store topologies from all runs, not only the current one, as it might have to pick up an
    optimization later. Meaning, the AC worker busily collects all topologies into memory. To prevent memory leak issues,
    we scrub the database of topologies that are older than a day as we do not expect to ever have to go back to those.

    Parameters
    ----------
    session : Session
        The database session. The session object will be modified in-place
    max_age_seconds : int
        The maximum age of topologies to keep in the database, in seconds. Topologies older than this will be removed.
    """
    cutoff_time = datetime.now() - timedelta(seconds=max_age_seconds)
    session.exec(delete(ACOptimTopology).where(ACOptimTopology.created_at < cutoff_time))
    session.commit()

toop_engine_topology_optimizer.ac.worker

The AC worker that listens to the kafka topics, organizes optimization runs, etc.

logger module-attribute

logger = get_logger(__name__)

Args

Bases: Args

Command line arguments for the AC worker.

Mostly the same as the DC worker except for an additional loadflow results folder

kafka_broker class-attribute instance-attribute

kafka_broker = 'localhost:9092'

The Kafka broker to connect to.

optimizer_command_topic class-attribute instance-attribute

optimizer_command_topic = 'commands'

The Kafka topic to listen for commands on.

optimizer_results_topic class-attribute instance-attribute

optimizer_results_topic = 'results'

The topic to push results to.

optimizer_heartbeat_topic class-attribute instance-attribute

optimizer_heartbeat_topic = 'heartbeat'

The topic to push heartbeats to.

heartbeat_interval_ms class-attribute instance-attribute

heartbeat_interval_ms = 1000

The interval in milliseconds to send heartbeats.

max_command_age_hours class-attribute instance-attribute

max_command_age_hours = 3.0

The maximum age of a command in hours. If a command is received that is older than this, the command will be ignored.

WorkerData dataclass

WorkerData(command_consumer, result_consumer, producer, db)

Data that is stored across optimization runs

command_consumer instance-attribute

command_consumer

A kafka consumer listening in for optimization commands

result_consumer instance-attribute

result_consumer

A kafka consumer listening on the results topic, constantly writing results to the database. This is polled both during the optimization and idle loop to keep the database up to date.

producer instance-attribute

producer

A kafka producer to send heartbeats and results

db instance-attribute

db

An initialized database session to an in-memory sqlite database.

initialize_optimization_run

initialize_optimization_run(
    ac_params,
    grid_file,
    worker_data,
    send_result_fn,
    send_heartbeat_fn,
    optimization_id,
    loadflow_result_fs,
    processed_gridfile_fs,
)

Initialize the AC optimization and wait for the first DC results.

Parameters are identical to optimization_loop plus the bound logger.

RETURNS	DESCRIPTION
tuple[OptimizerData, Strategy]	The initialized optimizer data and the initial topology message.

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/ac/worker.py

def initialize_optimization_run(
    ac_params: ACOptimizerParameters,
    grid_file: GridFile,
    worker_data: WorkerData,
    send_result_fn: Callable[[ResultUnion], None],
    send_heartbeat_fn: Callable[[HeartbeatUnion], None],
    optimization_id: str,
    loadflow_result_fs: AbstractFileSystem,
    processed_gridfile_fs: AbstractFileSystem,
) -> tuple[OptimizerData, Strategy]:
    """Initialize the AC optimization and wait for the first DC results.

    Parameters are identical to `optimization_loop` plus the bound logger.

    Returns
    -------
    tuple[OptimizerData, Strategy]
        The initialized optimizer data and the initial topology message.
    """
    send_heartbeat_fn(
        OptimizationStartedHeartbeat(
            optimization_id=optimization_id,
        )
    )
    optimizer_data, initial_topology = initialize_optimization(
        session=worker_data.db,
        params=ac_params,
        optimization_id=optimization_id,
        grid_file=grid_file,
        loadflow_result_fs=loadflow_result_fs,
        processed_gridfile_fs=processed_gridfile_fs,
    )
    wait_for_first_dc_results(
        results_consumer=worker_data.result_consumer,
        session=worker_data.db,
        max_wait_time=ac_params.ga_config.max_initial_wait_seconds,
        optimization_id=optimization_id,
        heartbeat_fn=partial(
            send_heartbeat_fn,
            OptimizationStatsHeartbeat(
                optimization_id=optimization_id,
                wall_time=0,
                iteration=0,
                num_branch_topologies_tried=0,
                num_injection_topologies_tried=0,
            ),
        ),
    )
    send_result_fn(
        OptimizationStartedResult(
            initial_topology=initial_topology,
        )
    )
    return optimizer_data, initial_topology

run_optimization_epochs

run_optimization_epochs(
    ac_params,
    optimizer_data,
    worker_data,
    send_result_fn,
    send_heartbeat_fn,
    optimization_id,
)

Run the iterative AC optimization phase.

PARAMETER	DESCRIPTION
ac_params	The parameters for the AC optimizer TYPE: ACOptimizerParameters
optimizer_data	The initialized optimizer data containing the curried functions and database session TYPE: OptimizerData
worker_data	The dataclass with the results consumer and database TYPE: WorkerData
send_result_fn	The function to send results TYPE: Callable[[ResultUnion], None]
send_heartbeat_fn	The function to send heartbeats TYPE: Callable[[HeartbeatUnion], None]
optimization_id	The ID of the optimization run TYPE: str

RETURNS	DESCRIPTION
None

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/ac/worker.py

def run_optimization_epochs(
    ac_params: ACOptimizerParameters,
    optimizer_data: OptimizerData,
    worker_data: WorkerData,
    send_result_fn: Callable[[ResultUnion], None],
    send_heartbeat_fn: Callable[[HeartbeatUnion], None],
    optimization_id: str,
) -> None:
    """Run the iterative AC optimization phase.

    Parameters
    ----------
    ac_params : ACOptimizerParameters
        The parameters for the AC optimizer
    optimizer_data : OptimizerData
        The initialized optimizer data containing the curried functions and database session
    worker_data : WorkerData
        The dataclass with the results consumer and database
    send_result_fn : Callable[[ResultUnion], None]
        The function to send results
    send_heartbeat_fn : Callable[[HeartbeatUnion], None]
        The function to send heartbeats
    optimization_id : str
        The ID of the optimization run

    Returns
    -------
    None

    """
    start_time = time.time()
    last_full_run = start_time
    survivor_batch_size = ac_params.ga_config.runner_processes
    epoch = 1
    evaluated_topologies = 0
    survivor_topologies = []
    survivor_early_results = []

    while True:
        with structlog.contextvars.bound_contextvars(epoch=epoch):
            added_topos, _ = poll_results_topic(
                db=optimizer_data.session, consumer=worker_data.result_consumer, first_poll=epoch == 1
            )
            logger.debug("Imported topologies from result stream", imported_topology_count=len(added_topos))

            # Even though the separation into fast-failing and remaining contingencies is not strictly necessary when
            # enable_ac_rejection is False, we still run the N-1 analysis in two steps to keep the logic similar to the
            # default enable_ac_rejection=True case and avoid having too many if statements in the code.
            topologies, worst_k_results = run_fast_failing_epoch(
                optimizer_data=optimizer_data,
            )
            if ac_params.ga_config.enable_ac_rejection:
                success_topologies, success_early_stop_results = process_fast_failing_results(
                    optimizer_data=optimizer_data,
                    topologies=topologies,
                    fast_failing_results=worst_k_results,
                    send_result_fn=send_result_fn,
                    epoch=epoch,
                )
                evaluated_topologies += len(topologies)
                survivor_topologies.extend(success_topologies)
                survivor_early_results.extend(success_early_stop_results)
            else:
                survivor_topologies.extend(topologies)
                survivor_early_results.extend(worst_k_results)
                evaluated_topologies += len(topologies)

            enough_survivors = len(survivor_topologies) >= survivor_batch_size
            runtime_exceeded_since_last_full_run = (
                time.time() - last_full_run
            ) > ac_params.ga_config.remaining_loadflow_wait_seconds
            if enough_survivors or (runtime_exceeded_since_last_full_run and len(survivor_topologies) > 0):
                logger.debug(
                    f"Collected {len(survivor_topologies)} survivor topologies, running remaining contingencies evaluation"
                )
                evaluate_remaining_contingencies(
                    send_result_fn,
                    optimizer_data,
                    epoch,
                    survivor_topologies[:survivor_batch_size],
                    survivor_early_results[:survivor_batch_size],
                )
                survivor_topologies = survivor_topologies[survivor_batch_size:]
                survivor_early_results = survivor_early_results[survivor_batch_size:]
                epoch += 1
                last_full_run = time.time()

            send_heartbeat_fn(
                OptimizationStatsHeartbeat(
                    optimization_id=optimization_id,
                    wall_time=time.time() - start_time,
                    iteration=epoch,
                    num_branch_topologies_tried=evaluated_topologies - len(survivor_topologies),
                    num_injection_topologies_tried=0,
                )
            )

            if time.time() - start_time > ac_params.ga_config.runtime_seconds:
                if len(survivor_topologies) > 0:
                    logger.info(
                        f"Stopping optimization {optimization_id} at epoch {epoch} due to runtime limit"
                        f" with survivor strategies still present"
                        f" Running remaining contingencies evaluation before stopping"
                    )
                    evaluate_remaining_contingencies(
                        send_result_fn,
                        optimizer_data,
                        epoch,
                        survivor_topologies,
                        survivor_early_results,
                    )
                else:
                    logger.info(f"Stopping optimization at epoch {epoch} due to runtime limit with no survivor strategies")
                send_result_fn(OptimizationStoppedResult(epoch=epoch, reason="converged", message="runtime limit"))
                return

summarize_optimization_run

summarize_optimization_run(
    optimization_id,
    grid_file,
    worker_data,
    optimizer_data,
    processed_gridfile_fs,
)

Write the optimization summary artifacts.

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/ac/worker.py

def summarize_optimization_run(
    optimization_id: str,
    grid_file: GridFile,
    worker_data: WorkerData,
    optimizer_data: OptimizerData,
    processed_gridfile_fs: AbstractFileSystem,
) -> None:
    """Write the optimization summary artifacts."""
    logger.info(f"Writing summary for optimization {optimization_id}")
    write_summary(
        grid_file=grid_file,
        db=worker_data.db,
        processed_gridfile_fs=processed_gridfile_fs,
        optimization_id=optimization_id,
        action_set=optimizer_data.action_set,
    )

optimization_loop

optimization_loop(
    ac_params,
    grid_file,
    worker_data,
    send_result_fn,
    send_heartbeat_fn,
    optimization_id,
    loadflow_result_fs,
    processed_gridfile_fs,
)

Run the main loop for the AC optimizer.

This function will run the AC optimizer on the given grid files with the given parameters.

PARAMETER	DESCRIPTION
ac_params	The parameters for the AC optimizer TYPE: ACOptimizerParameters
grid_file	The grid file to optimize on TYPE: GridFile
worker_data	The dataclass with the results consumer and database TYPE: WorkerData
send_result_fn	The function to send results TYPE: Callable[[ResultUnion], None]
send_heartbeat_fn	The function to send heartbeats TYPE: Callable[[HeartbeatUnion], None]
optimization_id	The ID of the optimization run TYPE: str
loadflow_result_fs	A filesystem where the loadflow results are stored. Loadflows will be stored here using the uuid generation process and passed as a StoredLoadflowReference which contains the subfolder in this filesystem. TYPE: AbstractFileSystem
processed_gridfile_fs	The target filesystem for the preprocessing worker. This contains all processed grid files. During the import job, a new folder import_results.data_folder was created which will be completed with the preprocess call to this function. Internally, only the data folder is passed around as a dirfs. Note that the unprocessed_gridfile_fs is not needed here anymore, as all preprocessing steps that need the unprocessed gridfiles were already done. TYPE: AbstractFileSystem

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/ac/worker.py

def optimization_loop(
    ac_params: ACOptimizerParameters,
    grid_file: GridFile,
    worker_data: WorkerData,
    send_result_fn: Callable[[ResultUnion], None],
    send_heartbeat_fn: Callable[[HeartbeatUnion], None],
    optimization_id: str,
    loadflow_result_fs: AbstractFileSystem,
    processed_gridfile_fs: AbstractFileSystem,
) -> None:
    """Run the main loop for the AC optimizer.

    This function will run the AC optimizer on the given grid files with the given parameters.

    Parameters
    ----------
    ac_params : ACOptimizerParameters
        The parameters for the AC optimizer
    grid_file : GridFile
        The grid file to optimize on
    worker_data : WorkerData
        The dataclass with the results consumer and database
    send_result_fn : Callable[[ResultUnion], None]
        The function to send results
    send_heartbeat_fn : Callable[[HeartbeatUnion], None]
        The function to send heartbeats
    optimization_id : str
        The ID of the optimization run
    loadflow_result_fs: AbstractFileSystem
        A filesystem where the loadflow results are stored. Loadflows will be stored here using the uuid generation process
        and passed as a StoredLoadflowReference which contains the subfolder in this filesystem.
    processed_gridfile_fs: AbstractFileSystem
        The target filesystem for the preprocessing worker. This contains all processed grid files.
        During the import job,  a new folder import_results.data_folder was created
        which will be completed with the preprocess call to this function.
        Internally, only the data folder is passed around as a dirfs.
        Note that the unprocessed_gridfile_fs is not needed here anymore, as all preprocessing steps that need the
        unprocessed gridfiles were already done.
    """
    with structlog.contextvars.bound_contextvars(optimization_id=optimization_id):
        logger.info("Initializing optimization")

        try:
            optimizer_data, _ = initialize_optimization_run(
                ac_params=ac_params,
                grid_file=grid_file,
                worker_data=worker_data,
                send_result_fn=send_result_fn,
                send_heartbeat_fn=send_heartbeat_fn,
                optimization_id=optimization_id,
                loadflow_result_fs=loadflow_result_fs,
                processed_gridfile_fs=processed_gridfile_fs,
            )
        except AcNotConvergedError as e:
            # If the AC optimization did not converge in the base grid, we send a special message
            # to indicate that the optimization cannot be run.
            send_result_fn(OptimizationStoppedResult(reason="ac-not-converged", message=str(e)))
            logger.error(f"AC optimization {optimization_id} did not converge in the base grid: {e}")
            return
        except TimeoutError as e:
            # If the DC results did not arrive in time, we assume a failure on DC side and abandon the optimization
            send_result_fn(OptimizationStoppedResult(reason="dc-not-started", message=str(e)))
            logger.error(f"DC results for optimization {optimization_id} did not arrive in time: {e}")
            return
        except Exception as e:
            send_result_fn(OptimizationStoppedResult(reason="error", message=str(e)))
            logger.error(f"Error during initialization of optimization {optimization_id}: {e}")
            return

        logger.info(f"Starting optimization {optimization_id}")

        try:
            run_optimization_epochs(
                ac_params=ac_params,
                optimizer_data=optimizer_data,
                worker_data=worker_data,
                send_result_fn=send_result_fn,
                send_heartbeat_fn=send_heartbeat_fn,
                optimization_id=optimization_id,
            )
        except Exception as e:
            # Send a stop message to the results
            send_result_fn(OptimizationStoppedResult(reason="error", message=str(e)))
            logger.error(f"Error during optimization {optimization_id}: {e}")
            logger.error(f"Stack trace: {traceback.format_exc()}")
            return

        try:
            summarize_optimization_run(
                optimization_id=optimization_id,
                grid_file=grid_file,
                worker_data=worker_data,
                optimizer_data=optimizer_data,
                processed_gridfile_fs=processed_gridfile_fs,
            )
        except Exception as e:
            logger.error(f"Error while writing summary for optimization {optimization_id}: {e}")
            logger.error(f"Stack trace: {traceback.format_exc()}")
            return

idle_loop

idle_loop(
    worker_data,
    send_heartbeat_fn,
    send_result_fn,
    heartbeat_interval_ms,
    max_command_age_hours,
)

Run idle loop of the AC optimizer worker.

This will be running when the worker is currently not optimizing This will wait until a StartOptimizationCommand is received and return it. In case a ShutdownCommand is received, the worker will exit with the exit code provided in the command.

PARAMETER	DESCRIPTION
worker_data	The dataclass with the command consumer, results consumer and database TYPE: WorkerData
send_heartbeat_fn	A function to call when there were no messages received for a while. TYPE: Callable[[HeartbeatUnion], None]
send_result_fn	A function to call to send results back to the results topic, used to send a message in case a command is too old. TYPE: Callable[[ResultUnion, str], None]
heartbeat_interval_ms	The time to wait for a new command in milliseconds. If no command has been received, a heartbeat will be sent and then the receiver will wait for commands again. TYPE: int
max_command_age_hours	The maximum age of a command in hours. If a command is received that is older than this, the command will be ignored and a message will be sent to the results topic. TYPE: float

RETURNS	DESCRIPTION
StartOptimizationCommand	The start optimization command to start the optimization run with

RAISES	DESCRIPTION
SystemExit	If a ShutdownCommand is received

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/ac/worker.py

def idle_loop(
    worker_data: WorkerData,
    send_heartbeat_fn: Callable[[HeartbeatUnion], None],
    send_result_fn: Callable[[ResultUnion, str], None],
    heartbeat_interval_ms: int,
    max_command_age_hours: float,
) -> StartOptimizationCommand:
    """Run idle loop of the AC optimizer worker.

    This will be running when the worker is currently not optimizing
    This will wait until a StartOptimizationCommand is received and return it. In case a
    ShutdownCommand is received, the worker will exit with the exit code provided in the command.

    Parameters
    ----------
    worker_data : WorkerData
        The dataclass with the command consumer, results consumer and database
    send_heartbeat_fn : Callable[[HeartbeatUnion], None]
        A function to call when there were no messages received for a while.
    send_result_fn : Callable[[ResultUnion, str], None]
        A function to call to send results back to the results topic,
        used to send a message in case a command is too old.
    heartbeat_interval_ms : int
        The time to wait for a new command in milliseconds. If no command has been received, a
        heartbeat will be sent and then the receiver will wait for commands again.
    max_command_age_hours: float
        The maximum age of a command in hours.
        If a command is received that is older than this, the command will be ignored
        and a message will be sent to the results topic.


    Returns
    -------
    StartOptimizationCommand
        The start optimization command to start the optimization run with

    Raises
    ------
    SystemExit
        If a ShutdownCommand is received
    """
    send_heartbeat_fn(IdleHeartbeat())
    logger.info("Entering idle loop")
    while True:
        message = worker_data.command_consumer.poll(heartbeat_interval_ms / 1000)

        # Wait timeout exceeded - send a heartbeat and poll the results topic
        if not message:
            send_heartbeat_fn(IdleHeartbeat())
            poll_results_topic(
                db=worker_data.db,
                consumer=worker_data.result_consumer,
                first_poll=False,
            )
            continue

        message_value = message.value()
        if message_value is None:
            logger.warning("Received command without payload, dropping message")
            worker_data.command_consumer.commit()
            continue

        command = Command.model_validate_json(deserialize_message(message_value))

        if isinstance(command.command, ShutdownCommand):
            logger.info("Shutting down due to ShutdownCommand")
            worker_data.command_consumer.close()
            worker_data.result_consumer.close()
            raise SystemExit(command.command.exit_code)

        if isinstance(command.command, StartOptimizationCommand):
            time_of_command = datetime.fromisoformat(command.timestamp)
            if time_of_command < datetime.now() - timedelta(hours=max_command_age_hours):
                logger.warning(
                    f"Received command with timestamp from the past (timestamp: {time_of_command}, "
                    f"now: {datetime.now()}), skipping command"
                )
                send_result_fn(
                    OptimizationStoppedResult(
                        reason="command-too-old", message=f"Received outdated command: {command}. Skipping.."
                    ),
                    command.command.optimization_id,
                )
                worker_data.command_consumer.commit()
                continue
            with structlog.contextvars.bound_contextvars(
                optimization_id=command.command.optimization_id,
            ):
                return command.command

        # If we are here, we received a command that we do not know
        logger.warning(f"Received unknown command, dropping: {command} / {message.value}")
        worker_data.command_consumer.commit()

warmup_result_storage

warmup_result_storage(
    worker_data, heartbeat_fn, heartbeat_interval_ms=5000
)

Go through the results topic and stores all published topology results in the database

This way, when an optimization run starts, the local in-memory results storage is already populated with results, so no time is wasted in spooling through the results topic.

PARAMETER	DESCRIPTION
worker_data	The dataclass with the results consumer and database TYPE: WorkerData
heartbeat_fn	A function to call to send heartbeats during the warmup loop, as it can take a while if there are many results to spool through. TYPE: Callable[[HeartbeatUnion], None]
heartbeat_interval_ms	The time interval in milliseconds to send heartbeats during the warmup loop, by default 5000 ms TYPE: int DEFAULT: 5000

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/ac/worker.py

def warmup_result_storage(
    worker_data: WorkerData,
    heartbeat_fn: Callable[[HeartbeatUnion], None],
    heartbeat_interval_ms: int = 5000,
) -> None:
    """Go through the results topic and stores all published topology results in the database

    This way, when an optimization run starts, the local in-memory results storage is already populated with results, so
    no time is wasted in spooling through the results topic.

    Parameters
    ----------
    worker_data : WorkerData
        The dataclass with the results consumer and database
    heartbeat_fn : Callable[[HeartbeatUnion], None]
        A function to call to send heartbeats during the warmup loop, as it can take a while if there are many results to
        spool through.
    heartbeat_interval_ms : int
        The time interval in milliseconds to send heartbeats during the warmup loop, by default 5000 ms
    """
    logger.info("Starting warmup loop to spool through results topic and populate database")
    first_poll = True
    last_heartbeat_time = time.time()
    while True:
        added_topos, _ = poll_results_topic(
            db=worker_data.db,
            consumer=worker_data.result_consumer,
            first_poll=first_poll,
        )
        first_poll = False

        if added_topos == {}:
            break

        if time.time() - last_heartbeat_time > heartbeat_interval_ms / 1000:
            heartbeat_fn(StartupHeartbeat())
            last_heartbeat_time = time.time()

main

main(
    args,
    loadflow_result_fs,
    processed_gridfile_fs,
    producer,
    command_consumer,
    result_consumer,
)

Run the main AC worker loop.

PARAMETER	DESCRIPTION
args	The command line arguments TYPE: Args
loadflow_result_fs	A filesystem where the loadflow results are stored. Loadflows will be stored here using the uuid generation process and passed as a StoredLoadflowReference which contains the subfolder in this filesystem. TYPE: AbstractFileSystem
processed_gridfile_fs	The target filesystem for the preprocessing worker. This contains all processed grid files. During the import job, a new folder import_results.data_folder was created which will be completed with the preprocess call to this function. Internally, only the data folder is passed around as a dirfs. Note that the unprocessed_gridfile_fs is not needed here anymore, as all preprocessing steps that need the unprocessed gridfiles were already done. TYPE: AbstractFileSystem
producer	A kafka producer to send heartbeats and results TYPE: Producer
command_consumer	A kafka consumer listening in for optimization commands TYPE: LongRunningKafkaConsumer
result_consumer	A kafka consumer listening in for results TYPE: LongRunningKafkaConsumer

RAISES	DESCRIPTION
SystemExit	If the worker receives a ShutdownCommand

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/ac/worker.py

def main(
    args: Args,
    loadflow_result_fs: AbstractFileSystem,
    processed_gridfile_fs: AbstractFileSystem,
    producer: Producer,
    command_consumer: LongRunningKafkaConsumer,
    result_consumer: LongRunningKafkaConsumer,
) -> None:
    """Run the main AC worker loop.

    Parameters
    ----------
    args : Args
        The command line arguments
    loadflow_result_fs: AbstractFileSystem
        A filesystem where the loadflow results are stored. Loadflows will be stored here using the uuid generation process
        and passed as a StoredLoadflowReference which contains the subfolder in this filesystem.
    processed_gridfile_fs: AbstractFileSystem
        The target filesystem for the preprocessing worker. This contains all processed grid files.
        During the import job,  a new folder import_results.data_folder was created
        which will be completed with the preprocess call to this function.
        Internally, only the data folder is passed around as a dirfs.
        Note that the unprocessed_gridfile_fs is not needed here anymore, as all preprocessing steps that need the
        unprocessed gridfiles were already done.
    producer : Producer
        A kafka producer to send heartbeats and results
    command_consumer : LongRunningKafkaConsumer
        A kafka consumer listening in for optimization commands
    result_consumer : LongRunningKafkaConsumer
        A kafka consumer listening in for results

    Raises
    ------
    SystemExit
        If the worker receives a ShutdownCommand
    """
    instance_id = str(uuid4())
    logger.info(
        f"Starting AC worker {instance_id} with config {args}, spooling through results topic to warm up result storage"
    )

    # We create two separate consumers for the command and result topics as we don't want to
    # catch results during the idle loop.
    worker_data = WorkerData(
        command_consumer=command_consumer,
        # Create a results consumer that will listen to results from any DC optimizers
        # Make sure to use a unique group.id for each instance to avoid conflicts
        result_consumer=result_consumer,
        producer=producer,
        db=create_session(),
    )

    def send_heartbeat(message: HeartbeatUnion, ping_commands: bool) -> None:
        logger.debug(f"Sending heartbeat: {message}", message_type=type(message).__name__)
        heartbeat = Heartbeat(
            optimizer_type=OptimizerType.AC,
            instance_id=instance_id,
            message=message,
        )
        Heartbeat.model_validate(heartbeat)  # Validate the heartbeat message
        worker_data.producer.produce(
            args.optimizer_heartbeat_topic,
            value=serialize_message(heartbeat.model_dump_json()),
            key=heartbeat.instance_id.encode(),
        )
        worker_data.producer.flush()
        if ping_commands:
            worker_data.command_consumer.heartbeat()

    def send_result(message: ResultUnion, optimization_id: str) -> None:
        logger.info(
            f"Sending result for optimization {optimization_id}: {message}",
            optimization_id=optimization_id,
            result_type=type(message).__name__,
        )
        result = Result(
            result=message,
            optimization_id=optimization_id,
            optimizer_type=OptimizerType.AC,
            instance_id=instance_id,
        )
        Result.model_validate(result)  # Validate the result message
        worker_data.producer.produce(
            args.optimizer_results_topic,
            value=serialize_message(result.model_dump_json()),
            key=result.optimization_id.encode(),
        )
        worker_data.producer.flush()

    send_heartbeat(StartupHeartbeat(), ping_commands=False)
    worker_data.command_consumer.start_processing()
    warmup_result_storage(
        worker_data=worker_data,
        heartbeat_fn=partial(send_heartbeat, ping_commands=True),
        heartbeat_interval_ms=args.heartbeat_interval_ms,
    )
    worker_data.command_consumer.stop_processing()

    logger.info("Finished warmup loop, entering main loop to wait for commands and run optimizations")

    while True:
        # During the idle loop, the result consumer is paused and only the command consumer is active
        command = idle_loop(
            worker_data=worker_data,
            send_heartbeat_fn=partial(send_heartbeat, ping_commands=False),
            send_result_fn=send_result,
            heartbeat_interval_ms=args.heartbeat_interval_ms,
            max_command_age_hours=args.max_command_age_hours,
        )

        # During the optimization loop, the command consumer is paused and the result consumer is active
        worker_data.command_consumer.start_processing()
        assert len(command.grid_files) == 1, "Exactly one grid file should be provided for the AC optimizer"
        optimization_loop(
            ac_params=command.ac_params,
            grid_file=command.grid_files[0],
            worker_data=worker_data,
            send_result_fn=partial(send_result, optimization_id=command.optimization_id),
            send_heartbeat_fn=partial(send_heartbeat, ping_commands=True),
            optimization_id=command.optimization_id,
            loadflow_result_fs=loadflow_result_fs,
            processed_gridfile_fs=processed_gridfile_fs,
        )
        worker_data.command_consumer.stop_processing()
        scrub_db(worker_data.db)

Benchmarks

toop_engine_topology_optimizer.benchmark.benchmark

Run benchmark.

This module provides functionality to run benchmark tasks based on a provided configuration. It supports grid search for sweeping through the values of number of CUDA devices or a single parameter not defined inside lf_config and ga_config. Multiple parameter sweeps are not supported yet.

FUNCTION	DESCRIPTION
main

main

main(cfg)

Run benchmark tasks based on the provided configuration.

PARAMETER	DESCRIPTION
cfg	Configuration object containing benchmark tasks and output settings. TYPE: DictConfig

RETURNS	DESCRIPTION
None

Notes

The function performs the following steps: 1. Iterates over the benchmark tasks specified in the configuration. 2. Checks for grid search configurations and generates new configurations for each value in the grid. 3. Sets environment variables for each benchmark task. 4. Runs each benchmark task in a separate process to avoid side-effects. 5. Collects the results and logs them into a JSON file specified in the configuration.

Grid search currently supports sweeping through the values of number of CUDA devices or a single parameter not defined inside lf_config and ga_config. Multiple parameter sweeps are not supported yet.

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/benchmark/benchmark.py

@hydra.main(config_path="configs", config_name="ms_benchmark.yaml", version_base="1.2")
def main(cfg: DictConfig) -> None:  # FIXME: Does not work for ToOp engine
    """Run benchmark tasks based on the provided configuration.

    Parameters
    ----------
    cfg : DictConfig
        Configuration object containing benchmark tasks and output settings.

    Returns
    -------
    None

    Notes
    -----
    The function performs the following steps:
    1. Iterates over the benchmark tasks specified in the configuration.
    2. Checks for grid search configurations and generates new configurations for each value in the grid.
    3. Sets environment variables for each benchmark task.
    4. Runs each benchmark task in a separate process to avoid side-effects.
    5. Collects the results and logs them into a JSON file specified in the configuration.

    Grid search currently supports sweeping through the values of number of CUDA devices or
    a single parameter not defined inside `lf_config` and `ga_config`. Multiple parameter sweeps are not supported yet.
    """
    benchmarks = []
    for task_cfg in cfg.benchmark_tasks:
        task_cfg_composed = compose(config_name=task_cfg.task_config)["task"]

        # check for grid searches
        if "grid_search" in task_cfg_composed:
            # Note: the grid search only works for sweeping through the values of
            # number of cuda devices or a single parameter not defined inside lf_config and ga_config.
            # The support to sweep multiple paramters are not there yet.
            key = next(iter(task_cfg_composed["grid_search"].keys()))
            values = next(iter(task_cfg_composed["grid_search"].values()))
            for val in values:
                new_config = deepcopy(task_cfg_composed)
                new_config[key] = val

                #  edit task name
                new_config["task_name"] = new_config["task_name"] + "_" + key + "_" + str(val)
                benchmarks.append(new_config)
        else:
            benchmarks.append(task_cfg_composed)

    results = []
    for benchmark in tqdm(benchmarks):
        set_environment_variables(benchmark)

        # Run the benchmark in a clean process to avoid side-effects from previous benchmarks
        ctx = mp.get_context("spawn")
        parent_conn, child_conn = ctx.Pipe()
        process = Process(target=run_task_process, args=(benchmark, child_conn))
        process.start()
        res = parent_conn.recv()
        process.join()
        results.append(res)

    #  Log the results as a json file
    with open(cfg.output_json, "w", encoding="utf-8") as f:
        json.dump(results, f, indent=2)

DC Optimizer

toop_engine_topology_optimizer.dc.ga_helpers

Helper functions for genetic algorithms.

MixingEmitterState

Bases: EmitterState

The state of a MixingEmitter.

It extends the EmitterState with additional fields to track the number of branch and injection combinations and splits.

total_branch_combis instance-attribute

total_branch_combis

total_inj_combis instance-attribute

total_inj_combis

total_num_splits instance-attribute

total_num_splits

TrackingMixingEmitter

Bases: MixingEmitter

A MixingEmitter that tracks the number of branch and injection combinations and splits.

init

init(random_key, init_genotypes)

Overwrite the Emitter.init function to seed an EmitterState.

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/ga_helpers.py

def init(
    self,
    random_key: PRNGKeyArray,
    init_genotypes: Optional[Genotype],  # noqa: ARG002
) -> tuple[EmitterState, PRNGKeyArray]:
    """Overwrite the Emitter.init function to seed an EmitterState."""
    return MixingEmitterState(
        total_branch_combis=jnp.array(0, dtype=int),
        total_inj_combis=jnp.array(0, dtype=int),
        total_num_splits=jnp.array(0, dtype=int),
    ), random_key

state_update

state_update(
    emitter_state,
    repertoire,
    genotypes,
    fitnesses,
    descriptors,
    extra_scores,
)

Overwrite the state update to store information for the running means.

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/ga_helpers.py

def state_update(
    self,
    emitter_state: Optional[EmitterState],
    repertoire: Optional[Repertoire | DiscreteMapElitesRepertoire],  # noqa: ARG002
    genotypes: Optional[Genotype],  # noqa: ARG002
    fitnesses: Optional[Fitness],  # noqa: ARG002
    descriptors: Optional[Descriptor],  # noqa: ARG002
    extra_scores: ExtraScores,
) -> EmitterState:
    """Overwrite the state update to store information for the running means."""
    assert emitter_state is not None
    assert extra_scores is not None
    return MixingEmitterState(
        total_branch_combis=emitter_state.total_branch_combis
        + extra_scores.get("n_branch_combis", jnp.array(0, dtype=int)).astype(int),
        total_inj_combis=emitter_state.total_inj_combis
        + extra_scores.get("n_inj_combis", jnp.array(0, dtype=int)).astype(int),
        total_num_splits=emitter_state.total_num_splits
        + extra_scores.get("n_split_grids", jnp.array(0, dtype=int)).astype(int),
    )

RunningMeans dataclass

RunningMeans(
    start_time,
    last_time,
    time_step,
    total_branch_combis,
    total_inj_combis,
    br_per_sec,
    inj_per_sec,
    split_per_iter,
    last_emitter_state,
    n_outages,
    n_devices,
)

A dataclass to hold the running means estimations for the TQDM progress bar.

start_time instance-attribute

start_time

last_time instance-attribute

last_time

time_step instance-attribute

time_step

total_branch_combis instance-attribute

total_branch_combis

total_inj_combis instance-attribute

total_inj_combis

br_per_sec instance-attribute

br_per_sec

inj_per_sec instance-attribute

inj_per_sec

split_per_iter instance-attribute

split_per_iter

last_emitter_state instance-attribute

last_emitter_state

n_outages instance-attribute

n_outages

n_devices instance-attribute

n_devices

init_running_means

init_running_means(n_outages, n_devices)

Initialize an empty RunningMeans object.

PARAMETER	DESCRIPTION
n_outages	The number of outages TYPE: int
n_devices	The number of devices TYPE: int

RETURNS	DESCRIPTION
RunningMeans	The initialized running means

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/ga_helpers.py

def init_running_means(n_outages: int, n_devices: int) -> RunningMeans:
    """Initialize an empty RunningMeans object.

    Parameters
    ----------
    n_outages : int
        The number of outages
    n_devices : int
        The number of devices

    Returns
    -------
    RunningMeans
        The initialized running means
    """
    now = time.time()
    return RunningMeans(
        start_time=now,
        last_time=now,
        time_step=0.0,
        total_branch_combis=0,
        total_inj_combis=0,
        br_per_sec=EWMA(),
        inj_per_sec=EWMA(),
        split_per_iter=EWMA(),
        last_emitter_state=None,
        n_outages=n_outages,
        n_devices=n_devices,
    )

update_running_means

update_running_means(running_means, emitter_state)

Aggregate the emitter state statistics into the running means.

PARAMETER	DESCRIPTION
running_means	The running means to be updated TYPE: RunningMeans
emitter_state	The emitter state of the current iteration TYPE: EmitterState

RETURNS	DESCRIPTION
RunningMeans	The updated running means

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/ga_helpers.py

def update_running_means(running_means: RunningMeans, emitter_state: EmitterState) -> RunningMeans:
    """Aggregate the emitter state statistics into the running means.

    Parameters
    ----------
    running_means : RunningMeans
        The running means to be updated
    emitter_state : EmitterState
        The emitter state of the current iteration

    Returns
    -------
    RunningMeans
        The updated running means
    """
    now = time.time()
    running_means = deepcopy(running_means)
    emitter_state = jax.tree_util.tree_map(lambda x: x * running_means.n_devices, emitter_state)
    last_emitter_state = (
        running_means.last_emitter_state
        if running_means.last_emitter_state is not None
        else MixingEmitterState(
            total_branch_combis=jnp.array(0, dtype=int),
            total_inj_combis=jnp.array(0, dtype=int),
            total_num_splits=jnp.array(0, dtype=int),
        )
    )

    branch_diff = emitter_state.total_branch_combis.item() - last_emitter_state.total_branch_combis.item()
    inj_diff = emitter_state.total_inj_combis.item() - last_emitter_state.total_inj_combis.item()
    split_diff = emitter_state.total_num_splits.item() - last_emitter_state.total_num_splits.item()

    # Clip diffs due to sporadic integer overflows
    branch_diff = max(branch_diff, 0)
    inj_diff = max(inj_diff, 0)
    split_diff = max(split_diff, 0)

    running_means.total_branch_combis += branch_diff
    running_means.total_inj_combis += inj_diff

    running_means.time_step = now - running_means.last_time

    running_means.br_per_sec.update(branch_diff / running_means.time_step)
    running_means.inj_per_sec.update(inj_diff / running_means.time_step)
    running_means.split_per_iter.update(split_diff)

    running_means.last_time = now
    running_means.last_emitter_state = emitter_state

    return running_means

toop_engine_topology_optimizer.dc.main

Launcher for the Map-Elites optimizer.

example args:

--fixed_files /workspaces/AICoE_HPC_RL_Optimizer/data/static_information.hdf5 --stats_dir /workspaces/AICoE_HPC_RL_Optimizer/stats/ --tensorboard_dir /workspaces/AICoE_HPC_RL_Optimizer/stats/tensorboard/ --ga_config.target_metrics overload_energy_n_1 1.0 # The metrics to optimize with their weight --ga_config.me_descriptors.0.metric split_subs --ga_config.me_descriptors.0.num_cells 5 --ga_config.me_descriptors.1.metric switching_distance --ga_config.me_descriptors.1.num_cells 45 # The metrics to use as descriptors along with their maximum value --ga_config.observed_metrics overload_energy_n_1 split_subs switching_distance # All the relevant metrics, including target, descriptors and extra metrics you want to see in the report --ga_config.substation_unsplit_prob 0.5 # Instinctively better suited for mapelites --ga_config.substation_split_prob 0.5 # Instinctively better suited for mapelites --ga_config.plot # Enable plot generation and saving in stats_dir/plots/ --ga_config.iterations_per_epoch 50 # Basically how often you wanna get a report. Suggested range : 50 - 1000 --ga_config.runtime_seconds 300 # How many seconds to run the optimization for The first descriptor metric is the vertical axis, the second the horizontal axis.

logger module-attribute

logger = get_logger(__name__)

args module-attribute

args = cli(CLIArgs)

file_system module-attribute

file_system = LocalFileSystem()

CLIArgs

Bases: BaseModel

The arguments for a CLI invocation, mostly equal to OptimizationStartCommand.

ga_config class-attribute instance-attribute

ga_config = Field(default_factory=BatchedMEParameters)

The configuration for the genetic algorithm

lf_config class-attribute instance-attribute

lf_config = Field(default_factory=LoadflowSolverParameters)

The configuration for the loadflow solver

fixed_files class-attribute instance-attribute

fixed_files = ()

The file containing the static information. You can pass multiple files here

tensorboard_dir class-attribute instance-attribute

tensorboard_dir = 'tensorboard'

The directory to store the tensorboard logs

stats_dir class-attribute instance-attribute

stats_dir = 'stats'

The directory to store the json summaries

summary_frequency class-attribute instance-attribute

summary_frequency = 10

How often to write tensorboard summaries

checkpoint_frequency class-attribute instance-attribute

checkpoint_frequency = 100

How often to write json summaries

double_limits class-attribute instance-attribute

double_limits = (1.0, 1.0)

The double limits to use in the form (lower_limit, upper_limit). lower_limit: float The relative lower limit to set, for branches whose n-1 flows are below the lower limit upper_limit: float The relative upper_limit determining at what relative load a branch is considered overloaded. Branches in the band between lower and upper limit are considered overloaded if more load is added.

log_tensorboard

log_tensorboard(fitness, metrics, iteration, writer)

Log fitness and metrics to tensorboard.

PARAMETER	DESCRIPTION
fitness	The fitness value TYPE: float
metrics	The metrics to log TYPE: dict[MetricType, float]
iteration	The current iteration number TYPE: int
writer	The tensorboard writer TYPE: SummaryWriter

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/main.py

def log_tensorboard(
    fitness: float,
    metrics: dict[MetricType, float],
    iteration: int,
    writer: SummaryWriter,
) -> None:
    """Log fitness and metrics to tensorboard.

    Parameters
    ----------
    fitness : float
        The fitness value
    metrics : dict[MetricType, float]
        The metrics to log
    iteration : int
        The current iteration number
    writer : SummaryWriter
        The tensorboard writer
    """
    writer.add_scalar("fitness", fitness, iteration)
    for key, value in metrics.items():
        writer.add_scalar(key, value, iteration)

write_summary

write_summary(
    optimizer_data,
    repertoire,
    emitter_state,
    iteration,
    folder,
    cli_args,
    processed_gridfile_fs,
    final_results=False,
)

Write a summary to a json file.

Watch out : here, the optimizer data is out of sync with the jax_data

PARAMETER	DESCRIPTION
optimizer_data	The optimizer data of the optimization. This includes a jax_data, however as the jax_data is updated per-iteration and the optimizer_data only per-epoch, we need to pass the jax_data separately TYPE: OptimizerData
repertoire	The current repertoire for this iteration, will be used instead of the one in the optimizer_data TYPE: DiscreteMapElitesRepertoire
emitter_state	The emitter state for this iteration, will be used instead of the one in the optimizer_data TYPE: EmitterState
iteration	The iteration number TYPE: int
folder	The folder to write the summary to, relative to the processed_gridfile_fs TYPE: str
cli_args	The arguments used for invocation, will be added to the summary for documentation purposes TYPE: CLIArgs
processed_gridfile_fs	The target filesystem for the preprocessing worker. This contains all processed grid files. During the import job, a new folder import_results.data_folder was created which will be completed with the preprocess call to this function. Internally, only the data folder is passed around as a dirfs. Note that the unprocessed_gridfile_fs is not needed here anymore, as all preprocessing steps that need the unprocessed gridfiles were already done. TYPE: AbstractFileSystem
final_results	Whether this is the final results summary "res.json" or an intermediate one "res_{iteration}.json" TYPE: bool DEFAULT: False

RETURNS	DESCRIPTION
dict	The summary that was written

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/main.py

def write_summary(
    optimizer_data: OptimizerData,
    repertoire: DiscreteMapElitesRepertoire,
    emitter_state: EmitterState,
    iteration: int,
    folder: str,
    cli_args: CLIArgs,
    processed_gridfile_fs: AbstractFileSystem,
    final_results: bool = False,
) -> dict:
    """Write a summary to a json file.

    Watch out : here, the optimizer data is out of sync with the jax_data

    Parameters
    ----------
    optimizer_data : OptimizerData
        The optimizer data of the optimization. This includes a jax_data, however as the jax_data
        is updated per-iteration and the optimizer_data only per-epoch, we need to pass the
        jax_data separately
    repertoire : DiscreteMapElitesRepertoire
        The current repertoire for this iteration, will be used instead of the one in the optimizer_data
    emitter_state : EmitterState
        The emitter state for this iteration, will be used instead of the one in the optimizer_data
    iteration : int
        The iteration number
    folder : str
        The folder to write the summary to, relative to the processed_gridfile_fs
    cli_args : CLIArgs
        The arguments used for invocation, will be added to the summary for
        documentation purposes
    processed_gridfile_fs: AbstractFileSystem
        The target filesystem for the preprocessing worker. This contains all processed grid files.
        During the import job,  a new folder import_results.data_folder was created
        which will be completed with the preprocess call to this function.
        Internally, only the data folder is passed around as a dirfs.
        Note that the unprocessed_gridfile_fs is not needed here anymore, as all preprocessing steps that need the
        unprocessed gridfiles were already done.
    final_results : bool
        Whether this is the final results summary "res.json" or an intermediate one "res_{iteration}.json"

    Returns
    -------
    dict
        The summary that was written
    """
    processed_gridfile_fs.makedirs(folder, exist_ok=True)

    ga_config = cli_args.ga_config
    n_cells_per_dim = tuple(desc.num_cells for desc in ga_config.me_descriptors)
    descriptor_metrics = tuple(desc.metric for desc in ga_config.me_descriptors)
    args_dict = cli_args.model_dump()

    # Here we assume that contingency_ids are the same for all topos in the repertoire
    contingency_ids = optimizer_data.solver_configs[0].contingency_ids
    summary = summarize(
        repertoire=repertoire,
        emitter_state=emitter_state,
        initial_fitness=optimizer_data.initial_fitness,
        initial_metrics=optimizer_data.initial_metrics,
        contingency_ids=contingency_ids,
        grid_model_low_tap=optimizer_data.jax_data.dynamic_informations[0].nodal_injection_information.grid_model_low_tap
        if optimizer_data.jax_data.dynamic_informations[0].nodal_injection_information is not None
        else None,
    )
    summary.update(
        {
            "args": args_dict,
            "iteration": iteration,
        }
    )
    filename = "res.json" if final_results else f"res_{iteration}.json"
    with processed_gridfile_fs.open(os.path.join(folder, filename), "w") as f:
        json.dump(summary, f)
    if ga_config.plot:
        plot_repertoire(
            repertoire.fitnesses[: np.prod(n_cells_per_dim)],
            iteration,
            folder,
            n_cells_per_dim=n_cells_per_dim,
            descriptor_metrics=descriptor_metrics,
            save_plot=ga_config.plot,
        )
    return summary

main

main(args, processed_gridfile_fs)

Run main optimization function for CLI execution.

PARAMETER	DESCRIPTION
args	The arguments for the optimization TYPE: CLIArgs
processed_gridfile_fs	The target filesystem for the preprocessing worker. This contains all processed grid files. During the import job, a new folder import_results.data_folder was created which will be completed with the preprocess call to this function. Internally, only the data folder is passed around as a dirfs. Note that the unprocessed_gridfile_fs is not needed here anymore, as all preprocessing steps that need the unprocessed gridfiles were already done. TYPE: AbstractFileSystem

RETURNS	DESCRIPTION
dict	The final results of the optimization

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/main.py

def main(
    args: CLIArgs,
    processed_gridfile_fs: AbstractFileSystem,
) -> dict:
    """Run main optimization function for CLI execution.

    Parameters
    ----------
    args : CLIArgs
        The arguments for the optimization
    processed_gridfile_fs: AbstractFileSystem
        The target filesystem for the preprocessing worker. This contains all processed grid files.
        During the import job,  a new folder import_results.data_folder was created
        which will be completed with the preprocess call to this function.
        Internally, only the data folder is passed around as a dirfs.
        Note that the unprocessed_gridfile_fs is not needed here anymore, as all preprocessing steps that need the
        unprocessed gridfiles were already done.

    Returns
    -------
    dict
        The final results of the optimization
    """
    jax.config.update("jax_enable_x64", True)

    logger.info(f"Starting with config {args}")

    args_dict = args.model_dump()

    partial_write_summary = partial(
        write_summary,
        folder=args.stats_dir,
        cli_args=args,
        processed_gridfile_fs=processed_gridfile_fs,
    )

    optimizer_data, stats, initial_topology = initialize_optimization(
        params=DCOptimizerParameters(
            ga_config=args.ga_config,
            loadflow_solver_config=args.lf_config,
            summary_frequency=args.summary_frequency,
            check_command_frequency=args.summary_frequency,
            double_limits=DoubleLimitsSetpoint(lower=args.double_limits[0], upper=args.double_limits[1])
            if args.double_limits != (1.0, 1.0)
            else None,
        ),
        optimization_id="CLI",
        static_information_files=args.fixed_files,
        processed_gridfile_fs=processed_gridfile_fs,
    )

    running_means = init_running_means(
        n_outages=optimizer_data.jax_data.dynamic_informations[0].n_nminus1_cases,
        n_devices=len(jax.devices()) if args.lf_config.distributed else 1,
    )

    logger.info(f"Optimization started: {stats}")

    writer = SummaryWriter(f"{args.tensorboard_dir}/{datetime.datetime.now()}")
    writer.add_hparams(args_dict, {})

    # Log initial results
    log_tensorboard(
        fitness=initial_topology.timesteps[0].metrics.fitness,
        metrics=initial_topology.timesteps[0].metrics.extra_scores,
        iteration=0,
        writer=writer,
    )

    epoch = 0
    with tqdm(
        total=args.ga_config.runtime_seconds,
        bar_format="{l_bar}{bar}[{elapsed}<{remaining}]{postfix}",
        postfix={
            "f0": optimizer_data.initial_fitness,
            "f": optimizer_data.initial_fitness,
            "br/s": 0,
            "inj/s": 0,
            "lfs": 0,
            "epoch": 0,
        },
    ) as pbar:
        while time.time() - running_means.start_time < args.ga_config.runtime_seconds:
            optimizer_data = run_epoch(optimizer_data)

            with jax.default_device(jax.devices("cpu")[0]):
                repertoire = (
                    jax.tree_util.tree_map(lambda x: x[0], optimizer_data.jax_data.repertoire)
                    if args.lf_config.distributed
                    else optimizer_data.jax_data.repertoire
                )
                emitter_state = (
                    jax.tree_util.tree_map(lambda x: x[0], optimizer_data.jax_data.emitter_state)
                    if args.lf_config.distributed
                    else optimizer_data.jax_data.emitter_state
                )

                fitness, metrics = get_repertoire_metrics(repertoire, args.ga_config.observed_metrics)
                log_tensorboard(fitness, metrics, epoch, writer)
                partial_write_summary(
                    optimizer_data,
                    repertoire,
                    emitter_state,
                    iteration=epoch,
                    final_results=False,
                )

                running_means = update_running_means(running_means=running_means, emitter_state=emitter_state)
            pbar.update(time.time() - running_means.start_time - pbar.n)
            pbar.set_postfix(
                {
                    "f0": optimizer_data.initial_fitness,
                    "f": fitness,
                    "br/s": running_means.br_per_sec.get(),
                    "inj/s": running_means.inj_per_sec.get(),
                    "lfs": running_means.total_inj_combis * running_means.n_outages,
                    "epoch": epoch,
                }
            )
            epoch += 1

    final_results = partial_write_summary(
        optimizer_data,
        jax.tree_util.tree_map(lambda x: x[0], optimizer_data.jax_data.repertoire)
        if args.lf_config.distributed
        else optimizer_data.jax_data.repertoire,
        jax.tree_util.tree_map(lambda x: x[0], optimizer_data.jax_data.emitter_state)
        if args.lf_config.distributed
        else optimizer_data.jax_data.emitter_state,
        iteration=epoch,
        final_results=True,
    )
    return final_results

Genetic Algorithm Functions

toop_engine_topology_optimizer.dc.genetic_functions.genotype

Dataclass for the genotype used in the genetic algorithm.

Genotype

Bases: Module

A single genome in the repertoire representing a topology.

action_index instance-attribute

action_index

An action index into the action set

disconnections instance-attribute

disconnections

The disconnections to apply, padded with int_max for disconnection slots that are unused. These are indices into the disconnectable branches set.

nodal_injections_optimized instance-attribute

nodal_injections_optimized

The results of the nodal injection optimization, if any was performed.

deduplicate_genotypes

deduplicate_genotypes(genotypes, desired_size=None)

Deduplicate the genotypes in the repertoire.

This version is jittable because we set the size

PARAMETER	DESCRIPTION
genotypes	The genotypes to deduplicate. These should be sorted by action id already. TYPE: Genotype
desired_size	How many unique values you are expecting. If not given, this is not jittable TYPE: Optional[int] DEFAULT: None

RETURNS	DESCRIPTION
Genotype	The deduplicated genotypes
Int[Array, ' n_unique']	The indices of the unique genotypes

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/genetic_functions/genotype.py

def deduplicate_genotypes(
    genotypes: Genotype,
    desired_size: Optional[int] = None,
) -> tuple[Genotype, Int[Array, " n_unique"]]:
    """Deduplicate the genotypes in the repertoire.

    This version is jittable because we set the size

    Parameters
    ----------
    genotypes : Genotype
        The genotypes to deduplicate. These should be sorted by action id already.
    desired_size : Optional[int]
        How many unique values you are expecting. If not given, this is not jittable

    Returns
    -------
    Genotype
        The deduplicated genotypes
    Int[Array, " n_unique"]
        The indices of the unique genotypes
    """
    # Use the action indices and the disconnections for the uniqueness check.
    genotype_parts = [
        genotypes.action_index,
        genotypes.disconnections,
    ]
    # Include nodal_injections_optimized (PST taps) in deduplication when present
    if genotypes.nodal_injections_optimized is not None:
        # Flatten the nodal injection optimization results into the comparison
        # Shape: (batch_size, n_timesteps, n_controllable_pst)
        # -> (batch_size, n_timesteps * n_controllable_pst)
        pst_taps_flat = genotypes.nodal_injections_optimized.pst_tap_idx.reshape(
            genotypes.nodal_injections_optimized.pst_tap_idx.shape[0], -1
        )
        genotype_parts.append(pst_taps_flat)

    genotype_flat = jnp.concatenate(genotype_parts, axis=1)

    _, indices = jnp.unique(
        genotype_flat,
        axis=0,
        return_index=True,
        size=desired_size,
        # fill_value takes the minimum flattened topology by default
        # it also corresponds to the first index in the list
    )
    unique_genotypes = jax.tree_util.tree_map(lambda x: x[indices], genotypes)
    return unique_genotypes, indices

fix_dtypes

fix_dtypes(genotypes)

Fix the dtypes of the genotypes to their native type.

For some reason, qdax aggressively converts everything to float

PARAMETER	DESCRIPTION
genotypes	The genotypes to fix TYPE: Genotype

RETURNS	DESCRIPTION
Genotype	The genotypes with fixed dtypes

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/genetic_functions/genotype.py

def fix_dtypes(genotypes: Genotype) -> Genotype:
    """Fix the dtypes of the genotypes to their native type.

    For some reason, qdax aggressively converts everything to float

    Parameters
    ----------
    genotypes : Genotype
        The genotypes to fix

    Returns
    -------
    Genotype
        The genotypes with fixed dtypes
    """
    return Genotype(
        action_index=genotypes.action_index.astype(int),
        disconnections=genotypes.disconnections.astype(int),
        nodal_injections_optimized=genotypes.nodal_injections_optimized,
    )

empty_repertoire

empty_repertoire(
    batch_size,
    max_num_splits,
    max_num_disconnections,
    n_timesteps,
    starting_taps=None,
)

Create an initial genotype repertoire with all zeros for all entries and int_max for all subs.

PARAMETER	DESCRIPTION
batch_size	The batch size TYPE: int
max_num_splits	The maximum number of splits per topology TYPE: int
max_num_disconnections	The maximum number of disconnections as topological measures per topology TYPE: int
n_timesteps	The number of timesteps in the optimization horizon, used for the nodal injection optimization results TYPE: int
starting_taps	The starting taps for the psts. If None, no nodal inj optimization will be enabled and nodal_injections_optimized will be set to None. If provided, nodal_injections_optimized will be initialized with these starting taps. TYPE: Optional[Int[Array, ' n_controllable_pst']] DEFAULT: None

RETURNS	DESCRIPTION
Genotype	The initial genotype

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/genetic_functions/genotype.py

def empty_repertoire(
    batch_size: int,
    max_num_splits: int,
    max_num_disconnections: int,
    n_timesteps: int,
    starting_taps: Optional[Int[Array, " n_controllable_pst"]] = None,
) -> Genotype:
    """Create an initial genotype repertoire with all zeros for all entries and int_max for all subs.

    Parameters
    ----------
    batch_size : int
        The batch size
    max_num_splits : int
        The maximum number of splits per topology
    max_num_disconnections : int
        The maximum number of disconnections as topological measures per topology
    n_timesteps : int
        The number of timesteps in the optimization horizon, used for the nodal injection optimization results
    starting_taps : Optional[Int[Array, " n_controllable_pst"]]
        The starting taps for the psts. If None, no nodal inj optimization will be enabled and nodal_injections_optimized
        will be set to None. If provided, nodal_injections_optimized will be initialized with these starting taps.

    Returns
    -------
    Genotype
        The initial genotype
    """
    if starting_taps is not None:
        nodal_injections_optimized = NodalInjOptimResults(
            pst_tap_idx=jnp.tile(starting_taps[None, None, :], (batch_size, n_timesteps, 1))
        )
    else:
        nodal_injections_optimized = None

    return Genotype(
        action_index=jnp.full((batch_size, max_num_splits), int_max(), dtype=int),
        disconnections=jnp.full((batch_size, max_num_disconnections), int_max(), dtype=int),
        nodal_injections_optimized=nodal_injections_optimized,
    )

toop_engine_topology_optimizer.dc.genetic_functions.initialization

Initialization of the genetic algorithm for branch and injection choice optimization.

logger module-attribute

logger = get_logger(__name__)

JaxOptimizerData

Bases: Module

The part of the optimizer data that lives on GPU.

If distributed is enabled, every item will have a leading device dimension.

repertoire instance-attribute

repertoire

The repertoire object

emitter_state instance-attribute

emitter_state

The emitter state object

dynamic_informations instance-attribute

dynamic_informations

The list containing the dynamic information objects

random_key instance-attribute

random_key

The random key

latest_iteration instance-attribute

latest_iteration

The iteration that this emitter_state/repertoire belong to

update_max_mw_flows_according_to_double_limits

update_max_mw_flows_according_to_double_limits(
    dynamic_informations,
    solver_configs,
    lower_limit,
    upper_limit,
)

Update all dynamic informations max mw loads.

Runs an initial n-1 analysis to determine limits in mw.

PARAMETER	DESCRIPTION
dynamic_informations	List of static informations to calculate with max_mw_flow limits set at 1.0 TYPE: tuple[DynamicInformation, ...]
solver_configs	List of solver configurations to use for the loadflow TYPE: tuple[SolverConfig, ...]
lower_limit	The relative lower limit to set, for branches whose n-1 flows are below the lower limit TYPE: float
upper_limit	The relative upper_limit determining at what relative load a branch is considered overloaded. Branches in the band between lower and upper limit are considered overloaded if more load is added. TYPE: float

RETURNS	DESCRIPTION
tuple[DynamicInformation, ...]	The updated dynamic informations with new limits set.

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/genetic_functions/initialization.py

def update_max_mw_flows_according_to_double_limits(
    dynamic_informations: tuple[DynamicInformation, ...],
    solver_configs: tuple[SolverConfig, ...],
    lower_limit: float,
    upper_limit: float,
) -> tuple[DynamicInformation, ...]:
    """Update all dynamic informations max mw loads.

    Runs an initial n-1 analysis to determine limits in mw.

    Parameters
    ----------
    dynamic_informations: tuple[DynamicInformation, ...]
        List of static informations to calculate with max_mw_flow limits set at 1.0
    solver_configs: tuple[SolverConfig, ...]
        List of solver configurations to use for the loadflow
    lower_limit: float
        The relative lower limit to set, for branches whose n-1 flows are below the lower limit
    upper_limit: float
        The relative upper_limit determining at what relative load a branch is considered overloaded.
        Branches in the band between lower and upper limit are considered overloaded if more load is added.

    Returns
    -------
    tuple[DynamicInformation, ...]
        The updated dynamic informations with new limits set.

    """
    if lower_limit > upper_limit:
        raise ValueError(f"Lower limit {lower_limit} must be smaller than upper limit {upper_limit}")

    updated_dynamic_informations = []
    for dynamic_information, solver_config in zip(dynamic_informations, solver_configs, strict=True):
        solver_config_local = replace(solver_config, batch_size_bsdf=1)
        lf_res, success = compute_symmetric_batch(
            topology_batch=default_topology(solver_config_local),
            disconnection_batch=None,
            injections=None,
            nodal_inj_start_options=None,
            dynamic_information=dynamic_information,
            solver_config=solver_config_local,
        )
        assert jnp.all(success)
        # We will always have N-1 limits, so we compute the N-0 limits on the N-0 loadflow results
        # However, the N-0 results lack a dimension, so we need to add a virtual "failure" dim
        limited_max_mw_flow = compute_double_limits(
            lf_res.n_0_matrix[0, :, None, :],
            dynamic_information.branch_limits.max_mw_flow,
            lower_limit=lower_limit,
            upper_limit=upper_limit,
        )

        limited_max_mw_flow_n_1 = compute_double_limits(
            lf_res.n_1_matrix[0],
            dynamic_information.branch_limits.max_mw_flow_n_1
            if dynamic_information.branch_limits.max_mw_flow_n_1 is not None
            else dynamic_information.branch_limits.max_mw_flow,
            lower_limit=lower_limit,
            upper_limit=upper_limit,
        )

        updated_dynamic_informations.append(
            replace(
                dynamic_information,
                branch_limits=replace(
                    dynamic_information.branch_limits,
                    max_mw_flow_limited=limited_max_mw_flow,
                    max_mw_flow_n_1_limited=limited_max_mw_flow_n_1,
                ),
            )
        )

    return tuple(updated_dynamic_informations)

verify_static_information

verify_static_information(
    static_informations,
    max_num_disconnections,
    enable_nodal_inj_optim,
)

Verify the static information.

This function will be called after loading the static information. It should be used to verify that the static information is correct and can be used for the optimization run.

PARAMETER	DESCRIPTION
static_informations	The static information to verify TYPE: Iterable[StaticInformation]
max_num_disconnections	The maximum number of disconnections that can be made TYPE: int
enable_nodal_inj_optim	Whether to enable the nodal injection optimization. If so, all static informations are verified to contain nodal injection information with at least one controllable PST. TYPE: bool

RAISES	DESCRIPTION
AssertionError	If the static information is not correct

RETURNS	DESCRIPTION
None

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/genetic_functions/initialization.py

def verify_static_information(
    static_informations: Iterable[StaticInformation],
    max_num_disconnections: int,
    enable_nodal_inj_optim: bool,
) -> None:
    """Verify the static information.

    This function will be called after loading the static information. It should be used to verify
    that the static information is correct and can be used for the optimization run.

    Parameters
    ----------
    static_informations : Iterable[StaticInformation]
        The static information to verify
    max_num_disconnections : int
        The maximum number of disconnections that can be made
    enable_nodal_inj_optim: bool
        Whether to enable the nodal injection optimization. If so, all static informations are verified to contain
        nodal injection information with at least one controllable PST.

    Raises
    ------
    AssertionError
        If the static information is not correct

    Returns
    -------
    None
    """
    first_static_information = next(iter(static_informations))

    assert all(
        [
            jnp.array_equal(
                static_information.solver_config.branches_per_sub.val,
                first_static_information.solver_config.branches_per_sub.val,
            )
            for static_information in static_informations
        ]
    )
    first_set = first_static_information.dynamic_information.action_set
    if first_set is not None:
        assert all(
            static_information.dynamic_information.action_set is not None for static_information in static_informations
        )
        assert all(
            [(static_information.dynamic_information.action_set == first_set) for static_information in static_informations]
        ), "All static informations must have the same branch actions"
        assert all(
            [
                jnp.array_equal(
                    static_information.dynamic_information.action_set.n_actions_per_sub,
                    first_set.n_actions_per_sub,
                )
                for static_information in static_informations
            ]
        ), "All static informations must have the same number of branch actions"

    assert first_static_information.dynamic_information.disconnectable_branches.shape[0] >= max_num_disconnections, (
        "Not enough disconnectable branches for the maximum number of disconnections, "
        + f"got {first_static_information.dynamic_information.disconnectable_branches.shape[0]} and {max_num_disconnections}"
    )
    if first_static_information.dynamic_information.disconnectable_branches.shape[0] > 0:
        assert all(
            [
                jnp.array_equal(
                    first_static_information.dynamic_information.disconnectable_branches,
                    static_information.dynamic_information.disconnectable_branches,
                )
                for static_information in static_informations
            ]
        ), "All static informations must have the same disconnectable branches"
    if enable_nodal_inj_optim:
        assert first_static_information.dynamic_information.nodal_injection_information is not None, (
            "Nodal injection opt. is enabled, but the first static information does not contain nodal injection info. "
            "For nodal injection optimization, we require at least one controllable PST in the nodal injection info. "
        )
        assert all(
            [
                static_information.dynamic_information.nodal_injection_information is not None
                and static_information.dynamic_information.nodal_injection_information.controllable_pst_indices.shape[0] > 0
                for static_information in static_informations
            ]
        ), (
            "Nodal injection optimization is enabled, but some static/nodal injection info does not contain controll. PSTs. "
            "This requires at least one controllable PST in the nodal injection information. "
            "Disable nodal injection optimization or provide correct static information. "
        )

update_static_information

update_static_information(
    static_informations,
    batch_size,
    enable_nodal_inj_optim,
    enable_bb_outage,
    bb_outage_as_nminus1,
    clip_bb_outage_penalty,
    bb_outage_more_islands_penalty,
)

Perform any necessary preprocessing on the static information.

This mainly applies updated optimization parameters such as batch size, busbar settings, etc.

PARAMETER	DESCRIPTION
static_informations	The list of static informations to preprocess TYPE: list[StaticInformation]
batch_size	The batch size to use, will replace the batch size in the solver config TYPE: int
enable_nodal_inj_optim	Whether to enable the nodal injection optimization, if False, nodal_inj_optim related information will be removed from the dynamic information to save GPU memory. TYPE: bool
enable_bb_outage	Whether the optimizer should include busbar outage effects. This is only enabled when the loaded static information contains busbar outage data. TYPE: bool
bb_outage_as_nminus1	Whether busbar outages are handled as additional N-1 cases. This value is always written into the solver config. TYPE: bool
clip_bb_outage_penalty	Whether busbar outage penalties are clipped at 0. This value is always written into the solver config. TYPE: bool
bb_outage_more_islands_penalty	The islanding penalty stored in the busbar outage baseline. This value is always written into the loaded baseline analysis when present. TYPE: float

RETURNS	DESCRIPTION
tuple[StaticInformation, ...]	The updated static informations

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/genetic_functions/initialization.py

def update_static_information(
    static_informations: tuple[StaticInformation, ...],
    batch_size: int,
    enable_nodal_inj_optim: bool,
    enable_bb_outage: bool,
    bb_outage_as_nminus1: bool,
    clip_bb_outage_penalty: bool,
    bb_outage_more_islands_penalty: float,
) -> tuple[StaticInformation, ...]:
    """Perform any necessary preprocessing on the static information.

    This mainly applies updated optimization parameters such as batch size, busbar settings, etc.

    Parameters
    ----------
    static_informations : list[StaticInformation]
        The list of static informations to preprocess
    batch_size : int
        The batch size to use, will replace the batch size in the solver config
    enable_nodal_inj_optim: bool
        Whether to enable the nodal injection optimization, if False, nodal_inj_optim related information will be removed
        from the dynamic information to save GPU memory.
    enable_bb_outage : bool
        Whether the optimizer should include busbar outage effects. This is only enabled when
        the loaded static information contains busbar outage data.
    bb_outage_as_nminus1 : bool
        Whether busbar outages are handled as additional N-1 cases. This value is always written into
        the solver config.
    clip_bb_outage_penalty : bool
        Whether busbar outage penalties are clipped at 0. This value is always written into the
        solver config.
    bb_outage_more_islands_penalty : float
        The islanding penalty stored in the busbar outage baseline. This value is always written into
        the loaded baseline analysis when present.

    Returns
    -------
    tuple[StaticInformation, ...]
        The updated static informations
    """
    updated_pairs = []
    for static_information in static_informations:
        solver_config, dynamic_information = update_single_pair_branch_limit_information(
            solver_config=static_information.solver_config,
            dynamic_information=static_information.dynamic_information,
            batch_size=batch_size,
            enable_nodal_inj_optim=enable_nodal_inj_optim,
        )
        solver_config, dynamic_information = update_single_pair_bb_outage_information(
            solver_config=solver_config,
            dynamic_information=dynamic_information,
            enable_bb_outage=enable_bb_outage,
            bb_outage_as_nminus1=bb_outage_as_nminus1,
            clip_bb_outage_penalty=clip_bb_outage_penalty,
            bb_outage_more_islands_penalty=bb_outage_more_islands_penalty,
        )
        updated_pairs.append((solver_config, dynamic_information))

    static_informations = [
        replace(
            static_information,
            solver_config=solver_config,
            dynamic_information=dynamic_information,
        )
        for static_information, (solver_config, dynamic_information) in zip(static_informations, updated_pairs, strict=True)
    ]

    return tuple(static_informations)

update_single_pair_branch_limit_information

update_single_pair_branch_limit_information(
    solver_config,
    dynamic_information,
    batch_size,
    enable_nodal_inj_optim,
)

Normalize branch-limit and nodal-injection data for one timestep.

This helper applies the optimizer batch size to the solver config and fills branch-limit fields that the optimizer assumes are always present at runtime. It also drops nodal injection payloads when nodal injection optimization is disabled to avoid carrying unused GPU data.

PARAMETER	DESCRIPTION
solver_config	The solver configuration for one static information object. TYPE: SolverConfig
dynamic_information	The runtime data paired with solver_config. TYPE: DynamicInformation
batch_size	The batch size to write into both batch dimensions of the solver config. TYPE: int
enable_nodal_inj_optim	Whether nodal injection optimization is enabled. TYPE: bool

RETURNS	DESCRIPTION
tuple[SolverConfig, DynamicInformation]	The updated solver config and dynamic information pair.

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/genetic_functions/initialization.py

def update_single_pair_branch_limit_information(
    solver_config: SolverConfig,
    dynamic_information: DynamicInformation,
    batch_size: int,
    enable_nodal_inj_optim: bool,
) -> tuple[SolverConfig, DynamicInformation]:
    """Normalize branch-limit and nodal-injection data for one timestep.

    This helper applies the optimizer batch size to the solver config and fills branch-limit fields
    that the optimizer assumes are always present at runtime. It also drops nodal injection payloads
    when nodal injection optimization is disabled to avoid carrying unused GPU data.

    Parameters
    ----------
    solver_config : SolverConfig
        The solver configuration for one static information object.
    dynamic_information : DynamicInformation
        The runtime data paired with ``solver_config``.
    batch_size : int
        The batch size to write into both batch dimensions of the solver config.
    enable_nodal_inj_optim : bool
        Whether nodal injection optimization is enabled.

    Returns
    -------
    tuple[SolverConfig, DynamicInformation]
        The updated solver config and dynamic information pair.
    """
    updated_solver_config = replace(
        solver_config,
        batch_size_bsdf=batch_size,
        batch_size_injection=batch_size,
    )
    updated_dynamic_information = replace(
        dynamic_information,
        branch_limits=replace(
            dynamic_information.branch_limits,
            max_mw_flow_limited=(
                dynamic_information.branch_limits.max_mw_flow
                if dynamic_information.branch_limits.max_mw_flow_limited is None
                else dynamic_information.branch_limits.max_mw_flow_limited
            ),
            n0_n1_max_diff=(
                jnp.zeros_like(dynamic_information.branch_limits.max_mw_flow)
                if dynamic_information.branch_limits.n0_n1_max_diff is None
                else dynamic_information.branch_limits.n0_n1_max_diff
            ),
        ),
        nodal_injection_information=(dynamic_information.nodal_injection_information if enable_nodal_inj_optim else None),
    )
    return updated_solver_config, updated_dynamic_information

update_single_pair_bb_outage_information

update_single_pair_bb_outage_information(
    solver_config,
    dynamic_information,
    enable_bb_outage,
    bb_outage_as_nminus1,
    clip_bb_outage_penalty,
    bb_outage_more_islands_penalty,
)

Apply runtime busbar-outage configuration for one timestep.

Busbar outage data is persisted during preprocessing so it can be reused at optimization time. This helper decides whether that payload is actually needed for the current run. When busbar outages are disabled, the associated runtime data is stripped from the dynamic information to reduce memory use. When enabled, the solver config is updated and the baseline penalty analysis is refreshed to match the requested islanding penalty.

PARAMETER	DESCRIPTION
solver_config	The solver configuration for one static information object. TYPE: SolverConfig
dynamic_information	The runtime data paired with solver_config. TYPE: DynamicInformation
enable_bb_outage	Whether the optimizer should use busbar outage data when available. TYPE: bool
bb_outage_as_nminus1	Whether busbar outages should be treated as additional N-1 contingencies. TYPE: bool
clip_bb_outage_penalty	Whether the busbar outage penalty should be clipped at zero. TYPE: bool
bb_outage_more_islands_penalty	Penalty applied when a busbar outage creates more islands than the unsplit baseline. TYPE: float

RETURNS	DESCRIPTION
tuple[SolverConfig, DynamicInformation]	The updated solver config and dynamic information pair.

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/genetic_functions/initialization.py

def update_single_pair_bb_outage_information(
    solver_config: SolverConfig,
    dynamic_information: DynamicInformation,
    enable_bb_outage: bool,
    bb_outage_as_nminus1: bool,
    clip_bb_outage_penalty: bool,
    bb_outage_more_islands_penalty: float,
) -> tuple[SolverConfig, DynamicInformation]:
    """Apply runtime busbar-outage configuration for one timestep.

    Busbar outage data is persisted during preprocessing so it can be reused at optimization time.
    This helper decides whether that payload is actually needed for the current run. When busbar
    outages are disabled, the associated runtime data is stripped from the dynamic information to
    reduce memory use. When enabled, the solver config is updated and the baseline penalty analysis
    is refreshed to match the requested islanding penalty.

    Parameters
    ----------
    solver_config : SolverConfig
        The solver configuration for one static information object.
    dynamic_information : DynamicInformation
        The runtime data paired with ``solver_config``.
    enable_bb_outage : bool
        Whether the optimizer should use busbar outage data when available.
    bb_outage_as_nminus1 : bool
        Whether busbar outages should be treated as additional N-1 contingencies.
    clip_bb_outage_penalty : bool
        Whether the busbar outage penalty should be clipped at zero.
    bb_outage_more_islands_penalty : float
        Penalty applied when a busbar outage creates more islands than the unsplit baseline.

    Returns
    -------
    tuple[SolverConfig, DynamicInformation]
        The updated solver config and dynamic information pair.
    """
    has_rel_bb_outage_data = dynamic_information.action_set.rel_bb_outage_data is not None
    has_monitored_branches = dynamic_information.branches_monitored.size > 0
    has_stored_bb_outage_baseline = dynamic_information.bb_outage_baseline_analysis is not None
    has_non_rel_bb_outage_data = dynamic_information.non_rel_bb_outage_data is not None

    # Busbar outages can only be enabled when some preprocessed outage payload exists.
    has_bb_outage_data = has_stored_bb_outage_baseline or has_non_rel_bb_outage_data or has_rel_bb_outage_data
    should_enable_bb_outage = enable_bb_outage and has_bb_outage_data
    # A fresh baseline is only needed when busbar outages are scored as a separate penalty.
    needs_penalty_baseline = (
        should_enable_bb_outage and not bb_outage_as_nminus1 and has_rel_bb_outage_data and has_monitored_branches
    )
    # Keep or create a baseline only when busbar outages stay enabled for this run.
    should_keep_or_create_baseline = should_enable_bb_outage and (has_stored_bb_outage_baseline or needs_penalty_baseline)

    updated_solver_config = replace(
        solver_config,
        enable_bb_outages=should_enable_bb_outage,
        bb_outage_as_nminus1=bb_outage_as_nminus1,
        clip_bb_outage_penalty=clip_bb_outage_penalty,
    )
    updated_action_set = (
        dynamic_information.action_set
        if should_enable_bb_outage
        else replace(dynamic_information.action_set, rel_bb_outage_data=None)
    )
    updated_dynamic_information = replace(
        dynamic_information,
        action_set=updated_action_set,
        bb_outage_baseline_analysis=(
            replace(
                dynamic_information.bb_outage_baseline_analysis
                if dynamic_information.bb_outage_baseline_analysis is not None
                else get_bb_outage_baseline_analysis(dynamic_information, bb_outage_more_islands_penalty),
                more_splits_penalty=jnp.array(bb_outage_more_islands_penalty),
            )
            if should_keep_or_create_baseline
            else None
        ),
        non_rel_bb_outage_data=(dynamic_information.non_rel_bb_outage_data if should_enable_bb_outage else None),
    )
    return updated_solver_config, updated_dynamic_information

initialize_genetic_algorithm

initialize_genetic_algorithm(
    batch_size,
    max_num_splits,
    max_num_disconnections,
    static_informations,
    target_metrics,
    mutation_config,
    action_set,
    proportion_crossover,
    crossover_mutation_ratio,
    random_seed,
    observed_metrics,
    me_descriptors,
    distributed,
    devices=None,
    cell_depth=1,
    n_worst_contingencies=10,
)

Initialize the mapelites algorithm.

PARAMETER	DESCRIPTION
batch_size	The batch size to use TYPE: int
max_num_splits	The maximum number of substations that can be split TYPE: int
max_num_disconnections	The maximum number of disconnections that can be made TYPE: int
static_informations	The static information to use for the optimization run TYPE: list[StaticInformation]
target_metrics	The target metrics to use for the optimization run TYPE: tuple[tuple[MetricType, float], ...]
mutation_config	The mutation configuration to use for the optimization run TYPE: MutationConfig
action_set	The action set to use for mutations TYPE: ActionSet
proportion_crossover	The proportion of crossover to mutation TYPE: float
crossover_mutation_ratio	The ratio of crossover to mutation TYPE: float
random_seed	The random seed to use for reproducibility TYPE: int
observed_metrics	The observed metrics, i.e. which metrics are to be computed for logging purposes. TYPE: tuple[MetricType, ...]
me_descriptors	The descriptors to use for map elites TYPE: tuple[DescriptorDef, ...]
distributed	Whether to run the optimization on multiple devices TYPE: bool
devices	The devices to run the optimization on, if distributed TYPE: Optional[list[Device]] DEFAULT: None
cell_depth	The cell depth to use if applicable TYPE: int DEFAULT: 1
n_worst_contingencies	The number of worst contingencies to consider in the scoring function for calculating top_k_overloads_n_1. TYPE: int DEFAULT: 10

RETURNS	DESCRIPTION
DiscreteMapElites	The genetic algorithm object including scoring, mutate and crossover functions
JaxOptimizerData	The initialized jax dataclass

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/genetic_functions/initialization.py

def initialize_genetic_algorithm(
    batch_size: int,
    max_num_splits: int,
    max_num_disconnections: int,
    static_informations: tuple[StaticInformation, ...],
    target_metrics: tuple[tuple[MetricType, float], ...],
    mutation_config: MutationConfig,
    action_set: ActionSet,
    proportion_crossover: float,
    crossover_mutation_ratio: float,
    random_seed: int,
    observed_metrics: tuple[MetricType, ...],
    me_descriptors: tuple[DescriptorDef, ...],
    distributed: bool,
    devices: Optional[list[jax.Device]] = None,
    cell_depth: int = 1,
    n_worst_contingencies: int = 10,
) -> tuple[DiscreteMapElites, JaxOptimizerData]:
    """Initialize the mapelites algorithm.

    Parameters
    ----------
    batch_size : int
        The batch size to use
    max_num_splits : int
        The maximum number of substations that can be split
    max_num_disconnections : int
        The maximum number of disconnections that can be made
    static_informations : list[StaticInformation]
        The static information to use for the optimization run
    target_metrics : tuple[tuple[MetricType, float], ...]
        The target metrics to use for the optimization run
    mutation_config : MutationConfig
        The mutation configuration to use for the optimization run
    action_set : ActionSet
        The action set to use for mutations
    proportion_crossover : float
        The proportion of crossover to mutation
    crossover_mutation_ratio : float
        The ratio of crossover to mutation
    random_seed: int
        The random seed to use for reproducibility
    observed_metrics: tuple[MetricType, ...]
        The observed metrics, i.e. which metrics are to be computed for logging purposes.
    me_descriptors: tuple[Descriptor, ...]
        The descriptors to use for map elites
    distributed: bool
        Whether to run the optimization on multiple devices
    devices: Optional[list[jax.Device]]
        The devices to run the optimization on, if distributed
    cell_depth: int
        The cell depth to use if applicable
    n_worst_contingencies: int
        The number of worst contingencies to consider in the scoring function for calculating
        top_k_overloads_n_1.

    Returns
    -------
    DiscreteMapElites
        The genetic algorithm object including scoring, mutate and crossover functions
    JaxOptimizerData
        The initialized jax dataclass
    """
    assert max_num_splits <= static_informations[0].dynamic_information.n_sub_relevant, (
        "The maximum number of splits cannot be larger than the number of relevant substations"
    )

    assert max_num_disconnections <= static_informations[0].dynamic_information.disconnectable_branches.shape[0], (
        "The maximum number of disconnections cannot be larger than the number of disconnectable branches"
    )

    n_devices = len(jax.devices()) if distributed else 1

    dynamic_informations = tuple([static_information.dynamic_information for static_information in static_informations])
    solver_configs = tuple(
        [replace(static_information.solver_config, batch_size_bsdf=batch_size) for static_information in static_informations]
    )

    initial_topologies = empty_repertoire(
        batch_size=batch_size * n_devices,
        max_num_splits=max_num_splits,
        max_num_disconnections=max_num_disconnections,
        n_timesteps=dynamic_informations[0].n_timesteps,
        starting_taps=dynamic_informations[0].nodal_injection_information.starting_tap_idx
        if dynamic_informations[0].nodal_injection_information is not None
        else None,
    )

    scoring_function_partial = partial(
        scoring_function,
        solver_configs=solver_configs,
        target_metrics=target_metrics,
        observed_metrics=observed_metrics,
        descriptor_metrics=tuple([desc.metric for desc in me_descriptors]),
        n_worst_contingencies=n_worst_contingencies,
    )

    mutate_partial = partial(
        mutate,
        mutation_config=mutation_config,
        action_set=action_set,
    )
    crossover_partial = partial(crossover, action_set=action_set, prob_take_a=proportion_crossover)

    emitter = TrackingMixingEmitter(
        mutate_partial,
        crossover_partial,
        crossover_mutation_ratio,
        batch_size,
    )
    algo = DiscreteMapElites(
        scoring_function=scoring_function_partial,
        emitter=emitter,
        metrics_function=default_ga_metrics,  # TODO: Why do we set this to default and not observed?
        distributed=distributed,
        n_cells_per_dim=tuple([desc.num_cells for desc in me_descriptors]),
        cell_depth=cell_depth,
    )

    random_key = jax.random.PRNGKey(random_seed)
    latest_iteration = jnp.array(1, dtype=int)

    init_fn = algo.init
    # If we are running on multiple devices, we need to replicate the data so it lives on every
    # device. The only exception is the random key, where we want a different one on every device
    if distributed:
        initial_topologies = jax.tree_util.tree_map(
            lambda x: jnp.reshape(
                x,
                (
                    n_devices,
                    batch_size,
                )
                + x.shape[1:],
            ),
            initial_topologies,
        )
        random_key = jax.random.split(random_key, n_devices)
        dynamic_informations = jax.tree_util.tree_map(
            lambda x: jax.device_put_replicated(x, devices),
            dynamic_informations,
        )
        latest_iteration = jax.device_put_replicated(latest_iteration, devices)

        init_fn = jax.pmap(
            init_fn,
            axis_name="p",
            in_axes=(
                jax.tree_util.tree_map(lambda _x: 0, initial_topologies),
                0,
                jax.tree_util.tree_map(
                    lambda _x: 0,
                    dynamic_informations,
                ),
            ),
        )

    repertoire, emitter_state, random_key = init_fn(initial_topologies, random_key, dynamic_informations)

    jax_data = JaxOptimizerData(
        repertoire=repertoire,
        emitter_state=emitter_state,
        dynamic_informations=dynamic_informations,
        random_key=random_key,
        latest_iteration=latest_iteration,
    )
    return algo, jax_data

get_repertoire_metrics

get_repertoire_metrics(repertoire, observed_metrics)

Get the metrics of the best individual in the Mapelites repertoire.

PARAMETER	DESCRIPTION
repertoire	The repertoire TYPE: DiscreteMapElitesRepertoire
observed_metrics	The metrics to observe (max_flow_n_0, median_flow_n_0 ...) TYPE: tuple[MetricType, ...]

RETURNS	DESCRIPTION
float	The fitness
dict[MetricType, float]	The metrics as defined in METRICS

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/genetic_functions/initialization.py

def get_repertoire_metrics(
    repertoire: DiscreteMapElitesRepertoire, observed_metrics: tuple[MetricType, ...]
) -> tuple[float, dict[MetricType, float]]:
    """Get the metrics of the best individual in the Mapelites repertoire.

    Parameters
    ----------
    repertoire : DiscreteMapElitesRepertoire
        The repertoire

    observed_metrics : tuple[MetricType, ...]
        The metrics to observe (max_flow_n_0, median_flow_n_0 ...)

    Returns
    -------
    float
        The fitness
    dict[MetricType, float]
        The metrics as defined in METRICS
    """
    distributed = len(repertoire.fitnesses.shape) > 1
    repertoire = jax.tree_util.tree_map(lambda x: x[0], repertoire) if distributed else repertoire

    fitnesses = repertoire.fitnesses
    # Get best individual and its metrics
    best_idx = jnp.argsort(fitnesses, descending=True)
    metrics = jax.tree_util.tree_map(lambda x: x[best_idx], repertoire.extra_scores)
    # only keep metrics in observed_metrics
    metrics = {key: metrics[key] for key in observed_metrics}
    fitnesses = fitnesses[best_idx]

    best_individual_fitness = fitnesses[0].item()
    # best_individual_metrics corresponds to the first element of each observed metric
    best_individual_metrics = {key: value[0].item() for key, value in metrics.items()}

    return best_individual_fitness, best_individual_metrics  # , descriptors[0]

algo_setup

algo_setup(
    ga_args,
    lf_args,
    double_limits,
    static_information_files,
    processed_gridfile_fs,
)

Set up the genetic algorithm run.

PARAMETER	DESCRIPTION
ga_args	The genetic algorithm parameters TYPE: GeneticAlgorithParameters
lf_args	The loadflow solver parameters TYPE: LoadflowSolverParameters
double_limits	The lower and upper limit for the relative max mw flow if double limits are used TYPE: Optional[tuple[float, float]]
static_information_files	A list of files with static information to load TYPE: Sequence[str | Path]
processed_gridfile_fs	The target filesystem for the preprocessing worker. This contains all processed grid files. During the import job, a new folder import_results.data_folder was created which will be completed with the preprocess call to this function. Internally, only the data folder is passed around as a dirfs. Note that the unprocessed_gridfile_fs is not needed here anymore, as all preprocessing steps that need the unprocessed gridfiles were already done. TYPE: AbstractFileSystem

RETURNS	DESCRIPTION
DiscreteMapElites	The initialized genetic algorithm object, can be used to update the optimization run
JaxOptimizerData	The jax dataclass of all GPU data including dynamic information and the GA data
tuple[SolverConfig, ...]	The solver configurations for every timestep (the dynamic information is part of the jax dataclass)
float	The initial fitness, for logging purposes
dict	The initial metrics, for logging purposes
list[StaticInformationDescription]	Some statistics on the static information dataclasses that were loaded

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/genetic_functions/initialization.py

def algo_setup(
    ga_args: BatchedMEParameters,
    lf_args: LoadflowSolverParameters,
    double_limits: Optional[tuple[float, float]],
    static_information_files: Sequence[str | Path],
    processed_gridfile_fs: AbstractFileSystem,
) -> tuple[
    DiscreteMapElites,
    JaxOptimizerData,
    tuple[SolverConfig, ...],
    float,
    dict,
    list[StaticInformationStats],
]:
    """Set up the genetic algorithm run.

    Parameters
    ----------
    ga_args : GeneticAlgorithParameters
        The genetic algorithm parameters
    lf_args : LoadflowSolverParameters
        The loadflow solver parameters
    double_limits: Optional[tuple[float, float]]
        The lower and upper limit for the relative max mw flow if double limits are used
    static_information_files : Sequence[str | Path]
        A list of files with static information to load
    processed_gridfile_fs: AbstractFileSystem
        The target filesystem for the preprocessing worker. This contains all processed grid files.
        During the import job,  a new folder import_results.data_folder was created
        which will be completed with the preprocess call to this function.
        Internally, only the data folder is passed around as a dirfs.
        Note that the unprocessed_gridfile_fs is not needed here anymore, as all preprocessing steps that need the
        unprocessed gridfiles were already done.

    Returns
    -------
    DiscreteMapElites
        The initialized genetic algorithm object, can be used to update the optimization run
    JaxOptimizerData
        The jax dataclass of all GPU data including dynamic information and the GA data
    tuple[SolverConfig, ...]
        The solver configurations for every timestep (the dynamic information is part of the jax
        dataclass)
    float
        The initial fitness, for logging purposes
    dict
        The initial metrics, for logging purposes
    list[StaticInformationDescription]
        Some statistics on the static information dataclasses that were loaded
    """
    n_devices = len(jax.devices()) if lf_args.distributed else 1

    static_informations = tuple(
        [load_static_information_fs(filesystem=processed_gridfile_fs, filename=str(f)) for f in static_information_files]
    )

    logger.info(f"Running {n_devices} GPUs with config {ga_args}, {lf_args}")

    verify_static_information(
        static_informations, lf_args.max_num_disconnections, enable_nodal_inj_optim=ga_args.enable_nodal_inj_optim
    )

    static_informations = update_static_information(
        static_informations,
        lf_args.batch_size,
        enable_nodal_inj_optim=ga_args.enable_nodal_inj_optim,
        enable_bb_outage=ga_args.enable_bb_outage,
        bb_outage_as_nminus1=ga_args.bb_outage_as_nminus1,
        clip_bb_outage_penalty=ga_args.clip_bb_outage_penalty,
        bb_outage_more_islands_penalty=ga_args.bb_outage_more_islands_penalty,
    )

    if double_limits is not None:
        logger.info(f"Updating double limits to {double_limits}")
        dynamic_infos = update_max_mw_flows_according_to_double_limits(
            dynamic_informations=tuple(s.dynamic_information for s in static_informations),
            solver_configs=tuple(s.solver_config for s in static_informations),
            lower_limit=double_limits[0],
            upper_limit=double_limits[1],
        )
        static_informations = tuple(
            [
                replace(static_information, dynamic_information=dynamic_info)
                for static_information, dynamic_info in zip(static_informations, dynamic_infos, strict=True)
            ]
        )

    pst_metrics_without_optimization = {
        metric
        for metric, _ in ga_args.target_metrics
        if metric in {"pst_switching_distance", "pst_switching_distance_squared", "pst_activated"}
    }
    if not ga_args.enable_nodal_inj_optim and pst_metrics_without_optimization:
        logger.warning(
            (
                f"The target metrics include {pst_metrics_without_optimization} but nodal injection optimization "
                "is disabled. This will lead to these metrics being always 0 and not optimized for. "
                "Consider enabling nodal injection optimization or removing these metrics from the target metrics. "
            )
        )
    n_rel_subs = static_informations[0].dynamic_information.n_sub_relevant
    n_disconnectable_branches = len(static_informations[0].dynamic_information.disconnectable_branches)
    mutation_config = MutationConfig(
        mutation_repetition=ga_args.mutation_repetition,
        random_topo_prob=ga_args.random_topo_prob,
        substation_mutation_config=SubstationMutationConfig(
            n_subs_mutated_lambda=ga_args.n_subs_mutated_lambda,
            add_split_prob=ga_args.add_split_prob,
            change_split_prob=ga_args.change_split_prob,
            remove_split_prob=ga_args.remove_split_prob,
            n_rel_subs=n_rel_subs,
        ),
        disconnection_mutation_config=DisconnectionMutationConfig(
            add_disconnection_prob=ga_args.add_disconnection_prob,
            change_disconnection_prob=ga_args.change_disconnection_prob,
            remove_disconnection_prob=ga_args.remove_disconnection_prob,
            n_disconnectable_branches=n_disconnectable_branches,
        ),
        nodal_injection_mutation_config=NodalInjectionMutationConfig(
            pst_mutation_sigma=ga_args.pst_mutation_sigma,
            pst_mutation_probability=ga_args.pst_mutation_probability,
            pst_reset_probability=ga_args.pst_reset_probability,
            pst_n_taps=static_informations[0].dynamic_information.nodal_injection_information.pst_n_taps,
            pst_start_tap_idx=static_informations[0].dynamic_information.nodal_injection_information.starting_tap_idx,
        )
        if static_informations[0].dynamic_information.nodal_injection_information is not None
        else None,
    )
    algo, jax_data = initialize_genetic_algorithm(
        batch_size=lf_args.batch_size,
        max_num_splits=min(n_rel_subs, lf_args.max_num_splits),
        max_num_disconnections=min(n_disconnectable_branches, lf_args.max_num_disconnections),
        static_informations=static_informations,
        target_metrics=ga_args.target_metrics,
        action_set=static_informations[0].dynamic_information.action_set,
        mutation_config=mutation_config,
        proportion_crossover=ga_args.proportion_crossover,
        crossover_mutation_ratio=ga_args.crossover_mutation_ratio,
        random_seed=ga_args.random_seed,
        observed_metrics=ga_args.observed_metrics,
        distributed=lf_args.distributed,
        devices=jax.devices() if lf_args.distributed else None,
        me_descriptors=ga_args.me_descriptors,
        cell_depth=ga_args.cell_depth,
        n_worst_contingencies=ga_args.n_worst_contingencies,
    )

    initial_fitness, initial_metrics = get_repertoire_metrics(
        jax.tree_util.tree_map(lambda x: x[0], jax_data.repertoire) if lf_args.distributed else jax_data.repertoire,
        ga_args.observed_metrics,
    )

    static_information_descriptions = [
        extract_static_information_stats(
            static_information=static_information,
            overload_n0=initial_metrics.get("overload_energy_n_0", 0.0),
            overload_n1=initial_metrics.get("overload_energy_n_1", 0.0),
            time="",
        )
        for static_information in static_informations
    ]

    for desc in static_information_descriptions:
        logger.info(f"Starting optimization with static information: {desc}")

    return (
        algo,
        jax_data,
        tuple([static_information.solver_config for static_information in static_informations]),
        initial_fitness,
        initial_metrics,
        static_information_descriptions,
    )

toop_engine_topology_optimizer.dc.genetic_functions.scoring_functions

Create scoring functions for the genetic algorithm.

METRICS module-attribute

METRICS = {
    "max_flow_n_0": maximum,
    "median_flow_n_0": maximum,
    "overload_energy_n_0": add,
    "overload_energy_limited_n_0": add,
    "underload_energy_n_0": add,
    "transport_n_0": add,
    "exponential_overload_energy_n_0": add,
    "exponential_overload_energy_limited_n_0": add,
    "critical_branch_count_n_0": maximum,
    "critical_branch_count_limited_n_0": maximum,
    "max_flow_n_1": maximum,
    "median_flow_n_1": maximum,
    "overload_energy_n_1": add,
    "overload_energy_limited_n_1": add,
    "underload_energy_n_1": add,
    "transport_n_1": add,
    "exponential_overload_energy_n_1": add,
    "exponential_overload_energy_limited_n_1": add,
    "critical_branch_count_n_1": maximum,
    "critical_branch_count_limited_n_1": maximum,
    "n0_n1_delta": add,
    "cross_coupler_flow": add,
    "switching_distance": maximum,
    "split_subs": maximum,
    "n_2_penalty": add,
    "pst_switching_distance": maximum,
    "pst_switching_distance_squared": maximum,
    "pst_activated": maximum,
}

compute_overloads

compute_overloads(
    topologies,
    dynamic_information,
    solver_config,
    observed_metrics,
    n_worst_contingencies=10,
)

Compute the overloads for a single timestep by invoking the solver and aggregating the results.

PARAMETER	DESCRIPTION
topologies	The topologies to score, where the first max_num_splits entries are the substations, the second max_num_splits entries are the branch topos and the last max_num_splits entries are the injection topos TYPE: Genotype
dynamic_information	The dynamic information of the grid TYPE: DynamicInformation
solver_config	The static solver configuration TYPE: SolverConfig
observed_metrics	The metrics to observe TYPE: tuple[MetricType, ...]
n_worst_contingencies	The number of worst contingencies to return, by default 10 TYPE: int DEFAULT: 10

RETURNS	DESCRIPTION
dict[str, Float[Array, ' batch_size']]	A dictionary with the overload energy, transport, max flow and other metrics, from aggregate_to_metric_batched
NodalInjOptimResults	The results of the nodal injection optimization
Bool[Array, ' batch_size']	Whether the topologies were successfully solved

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/genetic_functions/scoring_functions.py

def compute_overloads(
    topologies: Genotype,
    dynamic_information: DynamicInformation,
    solver_config: SolverConfig,
    observed_metrics: tuple[MetricType, ...],
    n_worst_contingencies: int = 10,
) -> tuple[dict[str, Float[Array, " batch_size"]], Optional[NodalInjOptimResults], Bool[Array, " batch_size"]]:
    """Compute the overloads for a single timestep by invoking the solver and aggregating the results.

    Parameters
    ----------
    topologies : Genotype
        The topologies to score, where the first max_num_splits entries are the substations, the
        second max_num_splits entries are the branch topos and the last max_num_splits entries are
        the injection topos
    dynamic_information : DynamicInformation
        The dynamic information of the grid
    solver_config : SolverConfig
        The static solver configuration
    observed_metrics : tuple[MetricType, ...]
        The metrics to observe
    n_worst_contingencies : int, optional
        The number of worst contingencies to return, by default 10

    Returns
    -------
    dict[str, Float[Array, " batch_size"]]
        A dictionary with the overload energy, transport, max flow and other metrics, from aggregate_to_metric_batched
    NodalInjOptimResults
        The results of the nodal injection optimization
    Bool[Array, " batch_size"]
        Whether the topologies were successfully solved
    """
    topo_comp, disconnections, nodal_inj_start = translate_topology(topologies)

    lf_res, success = compute_symmetric_batch(
        topology_batch=topo_comp,
        disconnection_batch=disconnections,
        injections=None,  # Use from action set
        nodal_inj_start_options=nodal_inj_start,
        dynamic_information=dynamic_information,
        solver_config=solver_config,
    )

    # Extract initial PST tap indices if PST optimization is enabled
    initial_pst_tap_idx = None
    if dynamic_information.nodal_injection_information is not None:
        initial_pst_tap_idx = dynamic_information.nodal_injection_information.starting_tap_idx

    aggregates = {}
    for metric_name in observed_metrics:
        metric = aggregate_to_metric_batched(
            lf_res_batch=lf_res,
            branch_limits=dynamic_information.branch_limits,
            reassignment_distance=dynamic_information.action_set.reassignment_distance,
            n_relevant_subs=dynamic_information.n_sub_relevant,
            metric=metric_name,
            initial_pst_tap_idx=initial_pst_tap_idx,
        )
        metric = jnp.where(success, metric, jnp.inf)
        aggregates[metric_name] = metric

    # Note: compute_overloads is called for each timestep separately, so the results are not batched.
    # As we don't have multi timestep optimisation support yet, we compute the worst k contingencies
    # sequentially one timestep at a time. This means that the timestep dimension will always be 1.
    #  TODO This is a temporary solution until we have multi timestep support.
    worst_k_res = jax.vmap(get_worst_k_contingencies, in_axes=(None, 0, None))(
        n_worst_contingencies, lf_res.n_1_matrix, dynamic_information.branch_limits.max_mw_flow
    )
    aggregates["top_k_overloads_n_1"] = worst_k_res.top_k_overloads[:, 0]  # Take the first timestep only
    aggregates["case_indices"] = worst_k_res.case_indices[:, 0, :]

    return aggregates, lf_res.nodal_injections_optimized, success

scoring_function

scoring_function(
    topologies,
    random_key,
    dynamic_informations,
    solver_configs,
    target_metrics,
    observed_metrics,
    descriptor_metrics,
    n_worst_contingencies=10,
)

Create scoring function for the genetic algorithm.

PARAMETER	DESCRIPTION
topologies	The topologies to score TYPE: Genotype
random_key	The random key to use for the scoring (currently not used) TYPE: PRNGKeyArray
dynamic_informations	The dynamic information of the grid for every timestep TYPE: tuple[DynamicInformation, ...]
solver_configs	The solver configuration for every timestep TYPE: tuple[SolverConfig, ...]
target_metrics	The list of metrics to optimize for with their weights TYPE: tuple[tuple[MetricType, float], ...]
observed_metrics	The observed metrics TYPE: tuple[MetricType, ...]
descriptor_metrics	The metrics to use as descriptors TYPE: tuple[MetricType, ...]
n_worst_contingencies	The number of worst contingencies to consider for calculating top_k_overloads_n_1 TYPE: int DEFAULT: 10

RETURNS	DESCRIPTION
Float[Array, ' batch_size']	The fitness of the topologies, calculated as the weighted sum of the target metrics
Int[Array, ' batch_size n_dims']	The descriptors of the topologies
dict	The metrics of the topologies, including the target metrics and any other observed metrics
ExtraScores	Emitter Information
PRNGKeyArray	The random key that was passed in, unused
Genotype	The genotypes that were passed in, but updated to account for in-the-loop optimizations such as the nodal injection optimization.

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/genetic_functions/scoring_functions.py

def scoring_function(
    topologies: Genotype,
    random_key: PRNGKeyArray,
    dynamic_informations: tuple[DynamicInformation, ...],
    solver_configs: tuple[SolverConfig, ...],
    target_metrics: tuple[tuple[MetricType, float], ...],
    observed_metrics: tuple[MetricType, ...],
    descriptor_metrics: tuple[MetricType, ...],
    n_worst_contingencies: int = 10,
) -> tuple[
    Float[Array, " batch_size"],
    Int[Array, " batch_size n_dims"],
    dict,
    ExtraScores,
    PRNGKeyArray,
    Genotype,
]:
    """Create scoring function for the genetic algorithm.

    Parameters
    ----------
    topologies : Genotype
        The topologies to score
    random_key : PRNGKeyArray
        The random key to use for the scoring (currently not used)
    dynamic_informations : tuple[DynamicInformation, ...]
        The dynamic information of the grid for every timestep
    solver_configs : tuple[SolverConfig, ...]
        The solver configuration for every timestep
    target_metrics : tuple[tuple[MetricType, float], ...]
        The list of metrics to optimize for with their weights
    observed_metrics : tuple[MetricType, ...]
        The observed metrics
    descriptor_metrics : tuple[MetricType, ...]
        The metrics to use as descriptors
    n_worst_contingencies : int, optional
        The number of worst contingencies to consider for calculating
        top_k_overloads_n_1

    Returns
    -------
    Float[Array, " batch_size"]
        The fitness of the topologies, calculated as the weighted sum of the target metrics
    Int[Array, " batch_size n_dims"]
        The descriptors of the topologies
    dict
        The metrics of the topologies, including the target metrics and any other observed metrics
    ExtraScores
        Emitter Information
    PRNGKeyArray
        The random key that was passed in, unused
    Genotype
        The genotypes that were passed in, but updated to account for in-the-loop optimizations such as the nodal
        injection optimization.
    """
    n_topologies = len(topologies.action_index)

    metrics, nodal_injections_optimized, success = compute_overloads(
        topologies=topologies,
        dynamic_information=dynamic_informations[0],
        solver_config=solver_configs[0],
        observed_metrics=observed_metrics,
        n_worst_contingencies=n_worst_contingencies,
    )
    # Sequentially compute each subsequent timestep
    for dynamic_information, solver_config in zip(dynamic_informations[1:], solver_configs[1:], strict=True):
        metrics_local, _nodal_injections_optimized_local, success_local = compute_overloads(
            topologies=topologies,
            dynamic_information=dynamic_information,
            solver_config=solver_config,
            observed_metrics=observed_metrics,
        )
        success = success & success_local

        # TODO figure out how to stack nodal_inj optim results for multiple timesteps
        for key in observed_metrics:
            combine_fn = METRICS[key]
            metrics[key] = combine_fn(metrics[key], metrics_local[key])

    fitness = sum(-metrics[metric_name] * weight for metric_name, weight in target_metrics)

    emitter_info = {
        "n_branch_combis": jnp.array(n_topologies, dtype=int),
        "n_inj_combis": jnp.array(n_topologies, dtype=int),
        "n_split_grids": jnp.sum(~success),
    }

    descriptors = jnp.stack([metrics[key].astype(int) for key in descriptor_metrics], axis=1)

    topologies = replace(topologies, nodal_injections_optimized=nodal_injections_optimized)

    return (
        fitness,
        descriptors,
        metrics,
        emitter_info,
        random_key,
        topologies,
    )

translate_topology

translate_topology(topology)

Translate the topology into the format used by the solver.

PARAMETER	DESCRIPTION
topology	The topology in genotype form TYPE: Genotype

RETURNS	DESCRIPTION
ActionIndexComputations	The topology computations
Int[Array, ' batch_size max_num_disconnections']	The branch disconnections to apply
NodalInjStartOptions | None	The nodal injection optimization start options containing pst taps, or None if no PST optimization

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/genetic_functions/scoring_functions.py

def translate_topology(
    topology: Genotype,
) -> tuple[
    ActionIndexComputations,
    Int[Array, " batch_size max_num_disconnections"],
    NodalInjStartOptions | None,
]:
    """Translate the topology into the format used by the solver.

    Parameters
    ----------
    topology : Genotype
        The topology in genotype form

    Returns
    -------
    ActionIndexComputations
        The topology computations
    Int[Array, " batch_size max_num_disconnections"]
        The branch disconnections to apply
    NodalInjStartOptions | None
        The nodal injection optimization start options containing pst taps, or None if no PST optimization
    """
    batch_size, _max_num_splits = topology.action_index.shape

    topology = fix_dtypes(topology)

    # Branch actions can be read straight out of the branch actions array
    topo_comp = ActionIndexComputations(
        action=topology.action_index,
        pad_mask=jnp.ones((batch_size,), dtype=bool),
    )

    nodal_inj_start = make_start_options(topology.nodal_injections_optimized)

    return topo_comp, topology.disconnections, nodal_inj_start

filter_repo

filter_repo(repertoire, initial_fitness)

Reduce the repertoire to only valid and deduplicated topologies.

This will not return any topologies that are worse than the initial fitness and deduplicate the topologies. The function is currently not jax jit compatible.

PARAMETER	DESCRIPTION
repertoire	The repertoire to reduce TYPE: DiscreteMapElitesRepertoire
initial_fitness	The initial fitness of the grid. This is used to filter out topologies that are worse than this fitness. TYPE: float

RETURNS	DESCRIPTION
DiscreteMapElitesRepertoire	The reduced repertoire with only valid and deduplicated topologies.

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/genetic_functions/scoring_functions.py

def filter_repo(
    repertoire: DiscreteMapElitesRepertoire,
    initial_fitness: float,
) -> DiscreteMapElitesRepertoire:
    """Reduce the repertoire to only valid and deduplicated topologies.

    This will not return any topologies that are worse than the initial fitness and deduplicate the topologies.
    The function is currently not jax jit compatible.

    Parameters
    ----------
    repertoire : DiscreteMapElitesRepertoire
        The repertoire to reduce
    initial_fitness : float
        The initial fitness of the grid. This is used to filter out topologies that are worse than this fitness.

    Returns
    -------
    DiscreteMapElitesRepertoire
        The reduced repertoire with only valid and deduplicated topologies.
    """
    distributed = len(repertoire.fitnesses.shape) > 1
    if distributed:
        repertoire = repertoire[0]

    assert len(repertoire.fitnesses.shape) == 1, "Wrong shape on repertoire"
    # Deduplicate the repertoire and remove invalid entries
    valid_mask = jnp.isfinite(repertoire.fitnesses) & (repertoire.fitnesses > initial_fitness)
    repertoire = repertoire[valid_mask]

    _, indices = deduplicate_genotypes(repertoire.genotypes)
    repertoire = repertoire[indices]

    return repertoire

convert_to_topologies

convert_to_topologies(
    repertoire, contingency_ids, grid_model_low_tap=None
)

Take a repertoire and convert it to a list of kafka-sendable topologies.

PARAMETER	DESCRIPTION
repertoire	The repertoire to convert. You might want to filter it using filter_repo first. TYPE: DiscreteMapElitesRepertoire
contingency_ids	The contingency IDs for each topology TYPE: list[str]
grid_model_low_tap	The lowest tap value in the grid model, used to convert the relative tap values in the genotype to absolute tap values that can be sent to the kafka topics. This will only be read if nodal_injection results are present in the genotype. TYPE: Int[Array, ' n_controllable_psts'] | None DEFAULT: None

RETURNS	DESCRIPTION
list[Topology]	The list of topologies in the format used by the kafka topics.

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/genetic_functions/scoring_functions.py

def convert_to_topologies(
    repertoire: DiscreteMapElitesRepertoire,
    contingency_ids: list[str],
    grid_model_low_tap: Int[Array, " n_controllable_psts"] | None = None,
) -> list[Topology]:
    """Take a repertoire and convert it to a list of kafka-sendable topologies.

    Parameters
    ----------
    repertoire : DiscreteMapElitesRepertoire
        The repertoire to convert. You might want to filter it using filter_repo first.
    contingency_ids : list[str]
        The contingency IDs for each topology
    grid_model_low_tap : Int[Array, " n_controllable_psts"] | None
        The lowest tap value in the grid model, used to convert the relative tap values in the genotype to absolute tap
        values that can be sent to the kafka topics. This will only be read if nodal_injection results are present
        in the genotype.

    Returns
    -------
    list[Topology]
        The list of topologies in the format used by the kafka topics.
    """
    topologies = []

    for i in range(len(repertoire.fitnesses)):
        iter_repertoire = repertoire[i]

        action_indices = [int(act) for act in iter_repertoire.genotypes.action_index if act != int_max()]

        disconnections = [int(disc) for disc in iter_repertoire.genotypes.disconnections if disc != int_max()]

        nodal_inj = iter_repertoire.genotypes.nodal_injections_optimized
        pst_setpoints = None
        if nodal_inj is not None:
            assert grid_model_low_tap is not None, (
                "grid_model_low_tap must be provided if nodal_injections_optimized is present"
            )
            assert len(nodal_inj.pst_tap_idx.shape) == 2
            assert nodal_inj.pst_tap_idx.shape[0] == 1, "Only one timestep is supported, but found shape " + str(
                nodal_inj.pst_tap_idx.shape
            )
            tap_array = nodal_inj.pst_tap_idx[0].astype(int) + grid_model_low_tap
            pst_setpoints = tap_array.tolist()

        case_indices = iter_repertoire.extra_scores.pop("case_indices", [])
        case_ids = np.array(contingency_ids)[case_indices].tolist()
        metrics = Metrics(
            fitness=float(iter_repertoire.fitnesses),
            extra_scores={key: value.item() for key, value in iter_repertoire.extra_scores.items()},
            worst_k_contingency_cases=case_ids,
        )

        topologies.append(
            Topology(
                actions=action_indices,
                disconnections=disconnections,
                pst_setpoints=pst_setpoints,
                metrics=metrics,
            )
        )
    return topologies

summarize_repo

summarize_repo(
    repertoire,
    initial_fitness,
    contingency_ids,
    grid_model_low_tap=None,
)

Summarize the repertoire into a list of topologies.

PARAMETER	DESCRIPTION
repertoire	The repertoire to summarize TYPE: DiscreteMapElitesRepertoire
initial_fitness	The initial fitness of the grid TYPE: float
contingency_ids	The contingency IDs for each topology. Here we assume that this list is common for all the topologies in the repertoire. TODO: Fix me to have per topology contingency ids if needed TYPE: list[str]
grid_model_low_tap	The lowest tap value in the grid model, from nodal_injection_information.grid_model_low_tap. TYPE: Int[Array, ' n_controllable_psts'] | None DEFAULT: None

RETURNS	DESCRIPTION
list[Topology]	The summarized topologies

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/genetic_functions/scoring_functions.py

def summarize_repo(
    repertoire: DiscreteMapElitesRepertoire,
    initial_fitness: float,
    contingency_ids: list[str],
    grid_model_low_tap: Int[Array, " n_controllable_psts"] | None = None,
) -> list[Topology]:
    """Summarize the repertoire into a list of topologies.

    Parameters
    ----------
    repertoire : DiscreteMapElitesRepertoire
        The repertoire to summarize
    initial_fitness : float
        The initial fitness of the grid
    contingency_ids : list[str]
        The contingency IDs for each topology. Here we assume that this list is common for all the topologies
        in the repertoire.
        TODO: Fix me to have per topology contingency ids if needed
    grid_model_low_tap : Int[Array, " n_controllable_psts"] | None
        The lowest tap value in the grid model, from nodal_injection_information.grid_model_low_tap.

    Returns
    -------
    list[Topology]
        The summarized topologies
    """
    with jax.default_device(jax.devices("cpu")[0]):
        best_repo = filter_repo(
            repertoire=repertoire,
            initial_fitness=initial_fitness,
        )

        topologies = convert_to_topologies(best_repo, contingency_ids, grid_model_low_tap=grid_model_low_tap)

    return topologies

summarize

summarize(
    repertoire,
    emitter_state,
    initial_fitness,
    initial_metrics,
    contingency_ids,
    grid_model_low_tap=None,
)

Summarize the repertoire and emitter state into a serializable dictionary.

PARAMETER	DESCRIPTION
repertoire	The repertoire to summarize TYPE: DiscreteMapElitesRepertoire
emitter_state	The emitter state to summarize TYPE: EmitterState
initial_fitness	The initial fitness of the grid TYPE: float
initial_metrics	The initial metrics of the grid TYPE: dict
contingency_ids	A list of contingency ids. Here we assume that the list of contingency ids is common for all the topologies TYPE: list[str]
grid_model_low_tap	The lowest tap value in the grid model, from nodal_injection_information.grid_model_low_tap. TYPE: Int[Array, ' n_controllable_psts'] | None DEFAULT: None

RETURNS	DESCRIPTION
dict	The summarized dictionary, json serializable

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/genetic_functions/scoring_functions.py

def summarize(
    repertoire: DiscreteMapElitesRepertoire,
    emitter_state: EmitterState,
    initial_fitness: float,
    initial_metrics: dict,
    contingency_ids: list[str],
    grid_model_low_tap: Int[Array, " n_controllable_psts"] | None = None,
) -> dict:
    """Summarize the repertoire and emitter state into a serializable dictionary.

    Parameters
    ----------
    repertoire : DiscreteMapElitesRepertoire
        The repertoire to summarize
    emitter_state : EmitterState
        The emitter state to summarize
    initial_fitness : float
        The initial fitness of the grid
    initial_metrics : dict
        The initial metrics of the grid
    contingency_ids : list[str]
        A list of contingency ids. Here we assume that the list of contingency ids is common for all the topologies
    grid_model_low_tap : Int[Array, " n_controllable_psts"] | None
        The lowest tap value in the grid model, from nodal_injection_information.grid_model_low_tap.

    Returns
    -------
    dict
        The summarized dictionary, json serializable
    """
    topologies = summarize_repo(
        repertoire=repertoire,
        initial_fitness=initial_fitness,
        contingency_ids=contingency_ids,
        grid_model_low_tap=grid_model_low_tap,
    )
    max_fitness = max(t.metrics.fitness for t in topologies) if len(topologies) > 0 else initial_fitness

    # Store the topologies
    best_topos = [t.model_dump() for t in topologies]
    retval = {k: v.item() for k, v in emitter_state.__dict__.items()}
    retval.update(
        {
            "max_fitness": max_fitness,
            "best_topos": best_topos,
            "initial_fitness": initial_fitness,
            "initial_metrics": initial_metrics,
        }
    )
    return retval

toop_engine_topology_optimizer.dc.genetic_functions.crossover

Contains the genetic operations for the topologies.

This module contains the functions to perform the genetic operations of mutation and crossover on the topologies of the grid. The topologies are represented as Genotype dataclasses, which contain the substation ids, the branch topology, the injection topology and the disconnections.

sample_unique_from_array

sample_unique_from_array(
    random_key, sample_pool, sample_probs, n_samples
)

Sample n unique elements from an array (returning indices), counting int_max as always unique.

PARAMETER	DESCRIPTION
random_key	The random key to use for the sampling TYPE: PRNGKeyArray
sample_pool	The array to sample from. Only unique elements are sampled, however int_max entries are not checked for uniqueness TYPE: Int[Array, ' n']
sample_probs	The probabilities to sample each element TYPE: Float[Array, ' n']
n_samples	The number of samples to take, should be less than the number of non-int_max entries in array. TYPE: int

RETURNS	DESCRIPTION
Int[Array, ' n_samples']	The sampled indices into sample_pool

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/genetic_functions/crossover.py

def sample_unique_from_array(
    random_key: PRNGKeyArray,
    sample_pool: Int[Array, " n"],
    sample_probs: Float[Array, " n"],
    n_samples: int,
) -> Int[Array, " n_samples"]:
    """Sample n unique elements from an array (returning indices), counting int_max as always unique.

    Parameters
    ----------
    random_key : PRNGKeyArray
        The random key to use for the sampling
    sample_pool : Int[Array, " n"]
        The array to sample from. Only unique elements are sampled, however int_max entries are
        not checked for uniqueness
    sample_probs : Float[Array, " n"]
        The probabilities to sample each element
    n_samples : int
        The number of samples to take, should be less than the number of non-int_max entries in
        array.

    Returns
    -------
    Int[Array, " n_samples"]
        The sampled indices into sample_pool
    """
    subkeys = jax.random.split(random_key, n_samples)

    def _body_fn(
        i: Int[ArrayLike, " "],
        entries_sampled: tuple[Int[Array, " max_num_splits"], Bool[Array, " n_subs_rel"]],
    ) -> tuple[Int[Array, " max_num_splits"], Bool[Array, " n_subs_rel"]]:
        indices_sampled, choice_mask = entries_sampled

        probs = sample_probs * choice_mask
        probs = probs / jnp.sum(probs)

        current_index = jax.random.choice(subkeys[i], jnp.arange(sample_pool.shape[0]), shape=(1,), p=probs)[0]

        # Update the choice mask
        # If a substation has been sampled, we can't sample this substation again
        # If a no-split (int_max) has been sampled, we only mask out the sampled index
        choice_mask = jnp.where(
            sample_pool[current_index] == int_max(),
            choice_mask.at[current_index].set(False, mode="promise_in_bounds"),
            jnp.where(
                sample_pool == sample_pool[current_index],
                False,
                choice_mask,
            ),
        )

        # Update the indices sampled
        indices_sampled = indices_sampled.at[i].set(current_index, mode="promise_in_bounds")

        return (indices_sampled, choice_mask)

    indices_sampled, _choice_mask = jax.lax.fori_loop(
        lower=0,
        upper=n_samples,
        body_fun=_body_fn,
        init_val=(
            jnp.full((n_samples,), int_max(), dtype=int),
            jnp.ones(sample_pool.shape, dtype=bool),
        ),
        unroll=True,
    )

    return indices_sampled

crossover_unbatched

crossover_unbatched(
    topologies_a,
    topologies_b,
    random_key,
    action_set,
    prob_take_a,
)

Crossover two topologies while making sure that no substation is present twice.

This version is unbatched, i.e. it only works on a single topology. Use crossover for batched inputs.

PARAMETER	DESCRIPTION
topologies_a	The first topology TYPE: Genotype
topologies_b	The second topology TYPE: Genotype
random_key	The random key to use for the crossover TYPE: PRNGKeyArray
action_set	The branch action set containing available actions on a per-substation basis. This is needed to resolve the sub_ids for the branch actions TYPE: ActionSet
prob_take_a	The probability to take the value from topology_a, otherwise the value from topology_b is taken TYPE: float

RETURNS	DESCRIPTION
Genotype	The new topology
PRNGKeyArray	The new random key

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/genetic_functions/crossover.py

def crossover_unbatched(
    topologies_a: Genotype,
    topologies_b: Genotype,
    random_key: PRNGKeyArray,
    action_set: ActionSet,
    prob_take_a: float,
) -> tuple[Genotype, PRNGKeyArray]:
    """Crossover two topologies while making sure that no substation is present twice.

    This version is unbatched, i.e. it only works on a single topology. Use crossover for batched
    inputs.

    Parameters
    ----------
    topologies_a : Genotype
        The first topology
    topologies_b : Genotype
        The second topology
    random_key : PRNGKeyArray
        The random key to use for the crossover
    action_set : ActionSet
        The branch action set containing available actions on a per-substation basis.
        This is needed to resolve the sub_ids for the branch actions
    prob_take_a : float
        The probability to take the value from topology_a, otherwise the value from topology_b is
        taken

    Returns
    -------
    Genotype
        The new topology
    PRNGKeyArray
        The new random key
    """
    # The tricky part in the crossover is that both topologies could have the same sub-id on
    # different indices. We need to make sure that we don't end up with the same substation twice
    # in the new topology

    max_num_splits = topologies_a.action_index.shape[0]
    max_num_disconnections = topologies_a.disconnections.shape[0]
    sample_bus_key, sample_disc_key, random_key = jax.random.split(random_key, 3)

    # The probability to take the value from a is prob_take_a
    # The probability to take the value from b is 1 - prob_take_a
    base_sample_probs = jnp.ones(2 * max_num_splits)
    base_sample_probs = base_sample_probs.at[:max_num_splits].set(prob_take_a)
    base_sample_probs = base_sample_probs.at[max_num_splits:].set(1 - prob_take_a)

    topologies_concatenated: Genotype = jax.tree_util.tree_map(
        lambda x, y: jnp.concatenate([x, y], axis=0),
        topologies_a,
        topologies_b,
    )
    sub_ids_concatenated = extract_sub_ids(topologies_concatenated.action_index, action_set)

    indices_sampled = sample_unique_from_array(
        random_key=sample_bus_key,
        sample_pool=sub_ids_concatenated,
        sample_probs=base_sample_probs,
        n_samples=max_num_splits,
    )

    actions = topologies_concatenated.action_index.at[indices_sampled].get(mode="fill", fill_value=int_max())

    # Sample disconnection choices
    if max_num_disconnections != 0:
        base_sample_probs = jnp.ones(2 * max_num_disconnections)
        base_sample_probs = base_sample_probs.at[:max_num_disconnections].set(prob_take_a)
        base_sample_probs = base_sample_probs.at[max_num_disconnections:].set(1 - prob_take_a)

        disconnections_concatenated: Int[Array, " 2*max_num_disconnections"] = topologies_concatenated.disconnections
        indices_sampled = sample_unique_from_array(
            random_key=sample_disc_key,
            sample_pool=disconnections_concatenated,
            sample_probs=base_sample_probs,
            n_samples=max_num_disconnections,
        )

        disconnections = disconnections_concatenated.at[indices_sampled].get(mode="fill", fill_value=int_max())
    else:
        disconnections = jnp.array([], dtype=int)

    return Genotype(
        action_index=actions,
        disconnections=disconnections,
        nodal_injections_optimized=topologies_a.nodal_injections_optimized,
    ), random_key

crossover

crossover(
    topologies_a,
    topologies_b,
    random_key,
    action_set,
    prob_take_a,
)

Crossover two topologies while making sure that no substation is present twice.

PARAMETER	DESCRIPTION
topologies_a	The first topology TYPE: Genotype
topologies_b	The second topology TYPE: Genotype
random_key	The random key to use for the crossover TYPE: PRNGKeyArray
action_set	The branch action set containing available actions on a per-substation basis. TYPE: ActionSet
prob_take_a	The probability to take the value from topology_a, otherwise the value from topology_b is taken TYPE: float

RETURNS	DESCRIPTION
Genotype	The new topology
PRNGKeyArray	The new random key

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/genetic_functions/crossover.py

def crossover(
    topologies_a: Genotype,
    topologies_b: Genotype,
    random_key: PRNGKeyArray,
    action_set: ActionSet,
    prob_take_a: float,
) -> tuple[Genotype, PRNGKeyArray]:
    """Crossover two topologies while making sure that no substation is present twice.

    Parameters
    ----------
    topologies_a : Genotype
        The first topology
    topologies_b : Genotype
        The second topology
    random_key : PRNGKeyArray
        The random key to use for the crossover
    action_set : ActionSet
        The branch action set containing available actions on a per-substation basis.
    prob_take_a : float
        The probability to take the value from topology_a, otherwise the value from topology_b is
        taken

    Returns
    -------
    Genotype
        The new topology
    PRNGKeyArray
        The new random key
    """
    batch_size = topologies_a.action_index.shape[0]
    crossover_keys = jax.random.split(random_key, batch_size)
    topo, random_keys = jax.vmap(crossover_unbatched, in_axes=(0, 0, 0, None, None))(
        topologies_a, topologies_b, crossover_keys, action_set, prob_take_a
    )
    return topo, random_keys[-1]

Mutation

toop_engine_topology_optimizer.dc.genetic_functions.mutation.config

Mutation configuration classes for the genetic algorithm.

SubstationMutationConfig

Bases: Module

Configuration for the substation mutation operation.

The sum of all probabilities should be <= 1.0, the remaining probability will be used for no mutation. If all substations are already split, the add_split_prob will be ignored and the probabilities will be renormalized to sum to 1.

n_subs_mutated_lambda class-attribute instance-attribute

n_subs_mutated_lambda = field(static=True)

add_split_prob class-attribute instance-attribute

add_split_prob = field(static=True)

The probability to split an additional substation. In case all substations are already split, this probability is ignored.

change_split_prob class-attribute instance-attribute

change_split_prob = field(static=True)

The probability to change an already split substation to a different split in the same or another substation.

remove_split_prob class-attribute instance-attribute

remove_split_prob = field(static=True)

The probability to reset a substation to the unsplit state.

n_rel_subs class-attribute instance-attribute

n_rel_subs = field(static=True)

The number of substations in the topology, used to determine the valid range of substation ids.

DisconnectionMutationConfig

Bases: Module

Configuration for the disconnection mutation operation.

Holds the random parameters used during the mutation of the disconnections, which are applied after the substation mutation.

add_disconnection_prob class-attribute instance-attribute

add_disconnection_prob = field(static=True)

The probability to disconnect an additional branch.

change_disconnection_prob class-attribute instance-attribute

change_disconnection_prob = field(static=True)

The probability to change a disconnected branch to another one.

remove_disconnection_prob class-attribute instance-attribute

remove_disconnection_prob = field(static=True)

The probability to remove a reconnect a disconnected branch.

n_disconnectable_branches class-attribute instance-attribute

n_disconnectable_branches = field(static=True)

The number of disconnectable branches in the topology, used to determine the valid range of branch ids.

NodalInjectionMutationConfig

Bases: Module

Configuration for the nodal injection mutation operation.

Holds the random parameters used during the mutation of the nodal injection optimization results, which are applied after the substation mutation.

pst_mutation_sigma class-attribute instance-attribute

pst_mutation_sigma = field(static=True)

The sigma to use for the normal distribution to sample the PST tap mutation from.

pst_mutation_probability class-attribute instance-attribute

pst_mutation_probability = field(static=True)

The probability for an individual PST to be selected for mutation.

pst_reset_probability class-attribute instance-attribute

pst_reset_probability = field(static=True)

The probability for an individual PST to be reverted to its initial set point. A value of 0.0 means no reset. A value of 1.0 means all PSTs will be reset.

pst_n_taps instance-attribute

pst_n_taps

The number of taps for each PST, used to determine the valid range of tap positions for mutation.

pst_start_tap_idx instance-attribute

pst_start_tap_idx

The starting tap position as an index into the tap range of each controllable PST

MutationConfig

Bases: Module

Configuration for the mutation operation.

mutation_repetition class-attribute instance-attribute

mutation_repetition = field(static=True)

More chance to get unique mutations by mutating mutation_repetition copies of the repertoire. The repertoire is repeated x times before mutation and deduplicated after mutation.

random_topo_prob class-attribute instance-attribute

random_topo_prob = field(static=True)

The probability to apply a completely random topology, instead of mutating. This is added to increase the exploration capabilities of the mutation operator, but should be used with care as it can easily lead to a decrease in performance if the random topologies are of low quality.

This does not include PSTs for now.

substation_mutation_config class-attribute instance-attribute

substation_mutation_config = field(static=True)

The configuration for the substation mutation operation.

disconnection_mutation_config class-attribute instance-attribute

disconnection_mutation_config = field(static=True)

The configuration for the disconnection mutation operation.

nodal_injection_mutation_config instance-attribute

nodal_injection_mutation_config

The configuration for the nodal injection mutation operation.

toop_engine_topology_optimizer.dc.genetic_functions.mutation.mutate

Mutation functions for the genetic algorithm.

mutate_topology

mutate_topology(
    random_key,
    sub_ids,
    disconnections_topo,
    action,
    mutate_config,
    action_set,
)

Mutate the topology by changing the substation splits and disconnections according to the mutation configuration.

PARAMETER	DESCRIPTION
random_key	The random key to use for the mutation TYPE: PRNGKeyArray
sub_ids	The substation ids before the mutation TYPE: Int[Array, ' max_num_splits']
disconnections_topo	The disconnections before the mutation TYPE: Int[Array, ' max_num_disconnections']
action	The actions before the mutation TYPE: Int[Array, ' max_num_splits']
mutate_config	The mutation configuration TYPE: MutationConfig
action_set	The set of possible actions TYPE: ActionSet

RETURNS	DESCRIPTION
(Int[Array, ' max_num_splits'],)	The splits substation ids after the mutation
(Int[Array, ' max_num_splits'],)	The action ids after the mutation
(Int[Array, ' max_num_disconnections'],)	The disconnections after the mutation
PRNGKeyArray	The mutated substation ids, actions, disconnections, and the updated random key

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/genetic_functions/mutation/mutate.py

def mutate_topology(
    random_key: PRNGKeyArray,
    sub_ids: Int[Array, " max_num_splits"],
    disconnections_topo: Int[Array, " max_num_disconnections"],
    action: Int[Array, " max_num_splits"],
    mutate_config: MutationConfig,
    action_set: ActionSet,
) -> tuple[
    Int[Array, " max_num_splits"], Int[Array, " max_num_splits"], Int[Array, " max_num_disconnections"], PRNGKeyArray
]:
    """Mutate the topology by changing the substation splits and disconnections according to the mutation configuration.

    Parameters
    ----------
    random_key : PRNGKeyArray
        The random key to use for the mutation
    sub_ids : Int[Array, " max_num_splits"]
        The substation ids before the mutation
    disconnections_topo : Int[Array, " max_num_disconnections"]
        The disconnections before the mutation
    action : Int[Array, " max_num_splits"]
        The actions before the mutation
    mutate_config : MutationConfig
        The mutation configuration
    action_set : ActionSet
        The set of possible actions

    Returns
    -------
    Int[Array, " max_num_splits"],
        The splits substation ids after the mutation

    Int[Array, " max_num_splits"],
        The action ids after the mutation

    Int[Array, " max_num_disconnections"],
        The disconnections after the mutation
    PRNGKeyArray
        The mutated substation ids, actions, disconnections, and the updated random key
    """
    random_key, substation_key, disconnection_key = jax.random.split(random_key, 3)
    max_num_splits = sub_ids.shape[0]
    # Randomly decide how many mutation steps should be applied to this substation, at least 1 and at most max_num_splits
    n_subs_mutated = jax.random.poisson(
        substation_key, lam=mutate_config.substation_mutation_config.n_subs_mutated_lambda, shape=()
    )
    n_subs_mutated = jnp.clip(n_subs_mutated, 1, max_num_splits)

    # Mutate sub_ids, action n_subs_mutated times
    sub_ids, action, random_key = jax.lax.fori_loop(
        0,
        n_subs_mutated,
        lambda _i, args: mutate_sub_splits(
            sub_ids=args[0],
            action=args[1],
            random_key=args[2],
            sub_mutate_config=mutate_config.substation_mutation_config,
            action_set=action_set,
        ),
        (
            sub_ids,
            action,
            substation_key,
        ),
    )
    # mutate disconnections after the substations,
    # so that the disconnection mutation can take into account the mutated substation topology.
    if (disconnections_topo.size > 0) and mutate_config.disconnection_mutation_config.n_disconnectable_branches > 0:
        disconnections_topo = mutate_disconnections(
            random_key=disconnection_key,
            sub_ids=sub_ids,
            disconnections=disconnections_topo,
            disconnection_mutation_config=mutate_config.disconnection_mutation_config,
        )
    return sub_ids, action, disconnections_topo, random_key

create_random_topology

create_random_topology(
    random_key,
    sub_ids,
    disconnections,
    action_set,
    n_rel_subs,
    n_disconnectable_branches,
)

Create a random topology by sampling random substation splits, branch topologies and disconnections.

PARAMETER	DESCRIPTION
random_key	The random key to use for the mutation TYPE: PRNGKeyArray
sub_ids	The substation ids before the mutation, used to determine the number of splits to sample TYPE: Int[Array, ' max_num_splits']
disconnections	The disconnections before the mutation, used to determine the number of disconnections to sample TYPE: Int[Array, ' max_num_disconnections']
action_set	The set of possible actions, used to sample valid actions for the new topology TYPE: ActionSet
n_rel_subs	The number of relevant substations in the grid, used to determine the valid range of substation ids for mutation. TYPE: int
n_disconnectable_branches	The number of disconnectable branches in the grid, used to determine the valid range of branch ids for mutation. TYPE: int

RETURNS	DESCRIPTION
(Int[Array, ' max_num_splits'],)	The splits substation ids after the random mutation
(Int[Array, ' max_num_splits'],)	The action ids after the random mutation
(Int[Array, ' max_num_disconnections'],)	The disconnections after the random mutation
PRNGKeyArray	The mutated substation ids, actions, disconnections, and the updated random key

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/genetic_functions/mutation/mutate.py

def create_random_topology(
    random_key: PRNGKeyArray,
    sub_ids: Int[Array, " max_num_splits"],
    disconnections: Int[Array, " max_num_disconnections"],
    action_set: ActionSet,
    n_rel_subs: int,
    n_disconnectable_branches: int,
) -> tuple[
    Int[Array, " max_num_splits"], Int[Array, " max_num_splits"], Int[Array, " max_num_disconnections"], PRNGKeyArray
]:
    """Create a random topology by sampling random substation splits, branch topologies and disconnections.

    Parameters
    ----------
    random_key : PRNGKeyArray
        The random key to use for the mutation
    sub_ids : Int[Array, " max_num_splits"]
        The substation ids before the mutation, used to determine the number of splits to sample
    disconnections : Int[Array, " max_num_disconnections"]
        The disconnections before the mutation, used to determine the number of disconnections to sample
    action_set : ActionSet
        The set of possible actions, used to sample valid actions for the new topology
    n_rel_subs : int
        The number of relevant substations in the grid, used to determine the valid range of substation ids for mutation.
    n_disconnectable_branches: int
        The number of disconnectable branches in the grid, used to determine the valid range of branch ids for mutation.

    Returns
    -------
    Int[Array, " max_num_splits"],
        The splits substation ids after the random mutation
    Int[Array, " max_num_splits"],
        The action ids after the random mutation
    Int[Array, " max_num_disconnections"],
        The disconnections after the random mutation
    PRNGKeyArray
        The mutated substation ids, actions, disconnections, and the updated random key
    """
    unsplit_key, subs_key, actions_key, not_disc_key, disc_key, random_key = jax.random.split(random_key, 6)

    # Extract the number of splits and disconnections to sample from the input,
    # so we can create a random topology with the same number of splits and disconnections as the input topology.
    max_num_splits = sub_ids.shape[0]
    max_disconnections = disconnections.shape[0]

    # Sample random substation splits first, override some with int_max to not split all
    unsplit_mask = jax.random.bernoulli(unsplit_key, p=0.5, shape=(max_num_splits,))
    random_subs = jax.random.choice(subs_key, shape=(max_num_splits,), a=jnp.arange(n_rel_subs), replace=False)
    sub_ids = jnp.where(unsplit_mask, int_max(), random_subs)

    # Sample actions for the chosen substations
    action = jax.vmap(lambda sub_id, key: sample_action_index_from_branch_actions(key, sub_id, action_set))(
        sub_ids,
        jax.random.split(actions_key, max_num_splits),
    )

    # Sample disconnections
    not_disconnected_mask = jax.random.bernoulli(not_disc_key, p=0.5, shape=(max_disconnections,))
    random_disconnections = jax.random.choice(
        disc_key, shape=(max_disconnections,), a=jnp.arange(n_disconnectable_branches), replace=False
    )
    disconnections = jnp.where(not_disconnected_mask, int_max(), random_disconnections)

    return sub_ids, action, disconnections, random_key

mutate

mutate(topologies, random_key, mutation_config, action_set)

Mutate the topologies by splitting substations, disconnecting branches and mutating nodal injections (PSTs).

Makes sure that at all times, a substation is split at most once and that all branch actions are in range of the available actions for the substation. If a substation is not split, this is indicated by the value int_max in the substation, branch.

We mutate mutation_repetition copies of the initial repertoire to increase the chance of getting unique mutations.

PARAMETER	DESCRIPTION
topologies	The topologies to mutate TYPE: Genotype
random_key	The random key to use for the mutation TYPE: PRNGKeyArray
mutation_config	The mutation configuration containing the probabilities for the different mutation operations and their parameters TYPE: MutationConfig
action_set	The action set containing available actions on a per-substation basis. TYPE: ActionSet

RETURNS	DESCRIPTION
Genotype	The mutated topologies
PRNGKeyArray	The new random key

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/genetic_functions/mutation/mutate.py

def mutate(
    topologies: Genotype,
    random_key: PRNGKeyArray,
    mutation_config: MutationConfig,
    action_set: ActionSet,
) -> tuple[Genotype, PRNGKeyArray]:
    """Mutate the topologies by splitting substations, disconnecting branches and mutating nodal injections (PSTs).

    Makes sure that at all times, a substation is split at most once and that all branch actions
    are in range of the available actions for the substation. If a substation is
    not split, this is indicated by the value int_max in the substation, branch.

    We mutate mutation_repetition copies of the initial repertoire to increase the chance of getting
    unique mutations.

    Parameters
    ----------
    topologies : Genotype
        The topologies to mutate
    random_key : PRNGKeyArray
        The random key to use for the mutation
    mutation_config : MutationConfig
        The mutation configuration containing the probabilities for the different mutation operations and their parameters
    action_set : ActionSet
        The action set containing available actions on a per-substation basis.

    Returns
    -------
    Genotype
        The mutated topologies
    PRNGKeyArray
        The new random key
    """
    topologies = fix_dtypes(topologies)
    batch_size = topologies.action_index.shape[0]

    # Repeat the topologies to increase the chance of getting unique mutations
    repeated_topologies = repeat_topologies(topologies, batch_size, mutation_config.mutation_repetition)
    n_mutations = batch_size * mutation_config.mutation_repetition
    mutation_key, replacement_key, pst_key, random_key = jax.random.split(random_key, 4)
    sub_ids = extract_sub_ids(
        repeated_topologies.action_index,
        action_set,
    )
    n_random_topologies = round(mutation_config.random_topo_prob * n_mutations)

    # Setup batch functions for mutation and random topology creation
    mutate_topologies_batch = jax.vmap(
        lambda sub_id, action_single, disconnection_single, key: mutate_topology(
            random_key=key,
            sub_ids=sub_id,
            disconnections_topo=disconnection_single,
            action=action_single,
            mutate_config=mutation_config,
            action_set=action_set,
        )
    )

    create_random_topologies_batch = jax.vmap(
        lambda sub_id, disconnection_single, key: create_random_topology(
            random_key=key,
            sub_ids=sub_id,
            disconnections=disconnection_single,
            action_set=action_set,
            n_rel_subs=mutation_config.substation_mutation_config.n_rel_subs,
            n_disconnectable_branches=mutation_config.disconnection_mutation_config.n_disconnectable_branches,
        )
    )
    # If n_random_topologies is 0, we only mutate.
    if n_random_topologies == 0:
        sub_ids, action, disconnections_topo, _ = mutate_topologies_batch(
            sub_ids,
            repeated_topologies.action_index,
            repeated_topologies.disconnections,
            jax.random.split(mutation_key, n_mutations),
        )
    # If n_random_topologies is equal to n_mutations, we only create random topologies.
    elif n_random_topologies == n_mutations:
        sub_ids, action, disconnections_topo, _ = create_random_topologies_batch(
            sub_ids,
            repeated_topologies.disconnections,
            jax.random.split(replacement_key, n_mutations),
        )
    # If n_random_topologies is between 0 and n_mutations, we mutate all topologies and
    # then replace a random subset of them with random topologies.
    else:
        sub_ids, action, disconnections_topo, _ = mutate_topologies_batch(
            sub_ids,
            repeated_topologies.action_index,
            repeated_topologies.disconnections,
            jax.random.split(mutation_key, n_mutations),
        )

        replacement_idx_key, replacement_topology_key = jax.random.split(replacement_key)
        replacement_indices = jax.random.choice(
            replacement_idx_key,
            n_mutations,
            shape=(n_random_topologies,),
            replace=False,
        )

        random_sub_ids, random_action, random_disconnections, _ = create_random_topologies_batch(
            sub_ids[replacement_indices],
            disconnections_topo[replacement_indices],
            jax.random.split(replacement_topology_key, n_random_topologies),
        )

        sub_ids = sub_ids.at[replacement_indices].set(random_sub_ids)
        action = action.at[replacement_indices].set(random_action)
        disconnections_topo = disconnections_topo.at[replacement_indices].set(random_disconnections)

    nodal_injections_optimized = mutate_nodal_injections(
        random_key=pst_key,
        nodal_inj_info=repeated_topologies.nodal_injections_optimized,
        nodal_mutation_config=mutation_config.nodal_injection_mutation_config,
    )

    # Sort all action_idx and disconnections, so that the order does not matter for the uniqueness check.
    # We can sort the action_idx and disconnections because the order of the splits and disconnections does not matter
    topologies_mutated = Genotype(
        action_index=jnp.sort(action, axis=1),
        disconnections=jnp.sort(disconnections_topo, axis=1),
        nodal_injections_optimized=nodal_injections_optimized,
    )
    if mutation_config.mutation_repetition > 1:
        topologies_mutated, _ = deduplicate_genotypes(
            topologies_mutated,
            desired_size=batch_size,
        )

    return topologies_mutated, random_key

repeat_topologies

repeat_topologies(
    topologies, batch_size, mutation_repetition
)

Repeat the topologies mutation_repetition times to increase the chance of getting unique mutations.

PARAMETER	DESCRIPTION
topologies	The topologies to repeat TYPE: Genotype
batch_size	The batch size of the original topologies, used to determine the total size after repetition TYPE: int
mutation_repetition	The number of times to repeat the topologies TYPE: int

RETURNS	DESCRIPTION
Genotype	The repeated topologies

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/genetic_functions/mutation/mutate.py

def repeat_topologies(topologies: Genotype, batch_size: int, mutation_repetition: int) -> Genotype:
    """Repeat the topologies mutation_repetition times to increase the chance of getting unique mutations.

    Parameters
    ----------
    topologies : Genotype
        The topologies to repeat
    batch_size : int
        The batch size of the original topologies, used to determine the total size after repetition
    mutation_repetition : int
        The number of times to repeat the topologies

    Returns
    -------
    Genotype
        The repeated topologies
    """
    if mutation_repetition == 1:
        repeated_topologies = topologies
    else:
        repeated_topologies: Genotype = jax.tree.map(
            lambda x: jnp.repeat(
                x, repeats=mutation_repetition, axis=0, total_repeat_length=batch_size * mutation_repetition
            ),
            topologies,
        )

    return repeated_topologies

toop_engine_topology_optimizer.dc.genetic_functions.mutation.mutate_substations

Mutation functions for substations in the genetic algorithm.

change_split_substation

change_split_substation(
    random_key, sub_ids, n_rel_subs, int_max_value
)

Change a split to a different one.

Either in the same substation or in a different substation. This assumes, that there is a split already.

PARAMETER	DESCRIPTION
random_key	The random key to use for the mutation TYPE: PRNGKeyArray
sub_ids	The substation ids before mutation TYPE: Int[Array, ' max_num_splits']
n_rel_subs	Number of relevant substations that may be selected. TYPE: int
int_max_value	The value that indicates an unsplit substation, used to determine which substations are currently split and to sample a new substation id from the valid range TYPE: int

RETURNS	DESCRIPTION
Int[Array, ' max_num_splits']	The mutated substation ids, where one split substation is changed to a different one. If there are no split substations, the input sub_ids are returned unchanged
Int[Array, ' ']	The index of the substation that was mutated. If no substation was mutated, this is set to int_max

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/genetic_functions/mutation/mutate_substations.py

def change_split_substation(
    random_key: PRNGKeyArray,
    sub_ids: Int[Array, " max_num_splits"],
    n_rel_subs: int,
    int_max_value: int,
) -> tuple[Int[Array, " "], Int[Array, " "]]:
    """Change a split to a different one.

    Either in the same substation or in a different substation.
    This assumes, that there is a split already.

    Parameters
    ----------
    random_key : PRNGKeyArray
        The random key to use for the mutation
    sub_ids : Int[Array, " max_num_splits"]
        The substation ids before mutation
    n_rel_subs : int
        Number of relevant substations that may be selected.
    int_max_value : int
        The value that indicates an unsplit substation, used to determine which substations
        are currently split and to sample a new substation id from the valid range

    Returns
    -------
    Int[Array, " max_num_splits"]
        The mutated substation ids, where one split substation is changed to a different one.
        If there are no split substations, the input sub_ids are returned unchanged
    Int[Array, " "]
        The index of the substation that was mutated. If no substation was mutated, this is set to int_max
    """
    split_key, sub_key = jax.random.split(random_key, 2)
    is_split = sub_ids != int_max_value
    split_idx = get_random_true_idx(split_key, is_split, int_max_value)
    new_substation_idx = sample_new_id(sub_key, sub_ids, n_rel_subs, split_idx)
    return split_idx, new_substation_idx

unsplit_substation

unsplit_substation(random_key, sub_ids, int_max_value)

Reset a split substation to the unsplit state, with a certain probability.

PARAMETER	DESCRIPTION
random_key	The random key to use for the reset TYPE: PRNGKey
sub_ids	The substation ids before the reset TYPE: Int[Array, ' max_num_splits']
int_max_value	The value that indicates an unsplit substation, used to determine which substations are currently split and to sample a new substation id from the valid range TYPE: int

RETURNS	DESCRIPTION
Int[Array, ' max_num_splits']	The substation ids after the reset. If a reset occurred, the substation id at split_idx will be set to int_max
Int[Array, ' max_num_splits']	The topology action after the reset. If a reset occurred, the topology at split_idx will be set to int_max

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/genetic_functions/mutation/mutate_substations.py

def unsplit_substation(
    random_key: PRNGKeyArray,
    sub_ids: Int[Array, " max_num_splits"],
    int_max_value: int,
) -> tuple[
    Int[Array, " "],
    Int[Array, " "],
]:
    """Reset a split substation to the unsplit state, with a certain probability.

    Parameters
    ----------
    random_key : jax.random.PRNGKey
        The random key to use for the reset
    sub_ids : Int[Array, " max_num_splits"]
        The substation ids before the reset
    int_max_value : int
        The value that indicates an unsplit substation, used to determine which substations are currently
        split and to sample a new substation id from the valid range

    Returns
    -------
    Int[Array, " max_num_splits"]
        The substation ids after the reset. If a reset occurred, the substation id at split_idx
        will be set to int_max
    Int[Array, " max_num_splits"]
        The topology action after the reset. If a reset occurred, the topology at split_idx
        will be set to int_max
    """
    already_split = sub_ids != int_max_value
    split_idx = get_random_true_idx(random_key, already_split, int_max_value)
    return split_idx, jnp.array(int_max_value)

split_additional_sub

split_additional_sub(
    random_key, sub_ids, n_rel_subs, int_max_value
)

Mutate the substation ids of a single topology.

PARAMETER	DESCRIPTION
random_key	The random key to use for the mutation TYPE: PRNGKeyArray
sub_ids	The substation ids before mutation TYPE: Int[Array, ' max_num_splits']
n_rel_subs	Number of relevant substations that may be selected. TYPE: int
int_max_value	The value to use for the mutation, which should be the maximum integer value used to indicate an empty slot in the genotype. TYPE: int

RETURNS	DESCRIPTION
Int[Array, ' max_num_splits']	The mutated substation ids
Int[Array, ' ']	The index of the substation that was mutated. If no substation was mutated, this is set to a random substation that was already split, so branch and injection mutations can still happen despite no new split. If no substation was split yet and no split happened, this is set to int_max

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/genetic_functions/mutation/mutate_substations.py

def split_additional_sub(
    random_key: PRNGKeyArray,
    sub_ids: Int[Array, " max_num_splits"],
    n_rel_subs: int,
    int_max_value: int,
) -> tuple[Int[Array, " "], Int[Array, " "]]:
    """Mutate the substation ids of a single topology.

    Parameters
    ----------
    random_key : PRNGKeyArray
        The random key to use for the mutation
    sub_ids : Int[Array, " max_num_splits"]
        The substation ids before mutation
    n_rel_subs : int
        Number of relevant substations that may be selected.
    int_max_value : int
        The value to use for the mutation, which should be the maximum integer value used
        to indicate an empty slot in the genotype.

    Returns
    -------
    Int[Array, " max_num_splits"]
        The mutated substation ids
    Int[Array, " "]
        The index of the substation that was mutated. If no substation was mutated, this is set to
        a random substation that was already split, so branch and injection mutations can still
        happen despite no new split. If no substation was split yet and no split happened, this is
        set to int_max
    """
    split_key, sub_key = jax.random.split(random_key, 2)
    non_split = sub_ids == int_max_value
    split_idx = get_random_true_idx(split_key, non_split, int_max_value)
    new_substation_idx = sample_new_id(sub_key, sub_ids, n_rel_subs, int_max_value)
    return split_idx, new_substation_idx

mutate_sub_splits

mutate_sub_splits(
    sub_ids,
    action,
    random_key,
    sub_mutate_config,
    action_set,
)

Mutate a single substation, changing the sub_ids, branch and inj topos.

The sub-ids are implicit to the branch topo action index, however we pass them explicitely to aid substation mutation. Impossible mutations (e.g. adding a split when there is no room to add one, or removing a split when there are none) are handled by setting the probabilities for these mutations to zero, and normalising the remaining probabilities to sum to 1.

PARAMETER	DESCRIPTION
sub_ids	The substation ids before mutation TYPE: Int[Array, ' max_num_splits']
action	The branch topology before mutation TYPE: Int[Array, ' max_num_splits']
random_key	The random key to use for the mutation TYPE: PRNGKeyArray
sub_mutate_config	The configuration for the substation mutation, containing the probabilities for the different mutation types and the number of relevant substations in the grid, which is needed to determine the valid range of substation ids for mutation. TYPE: SubstationMutationConfig
action_set	The actions for every substation. If not provided, will sample from all possible actions per substation TYPE: ActionSet

RETURNS	DESCRIPTION
sub_ids	The substation ids after mutation TYPE: Int[Array, ' max_num_splits']
action	The branch topology after mutation TYPE: Int[Array, ' max_num_splits']
random_key	The random key used for the mutation TYPE: PRNGKeyArray

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/genetic_functions/mutation/mutate_substations.py

def mutate_sub_splits(
    sub_ids: Int[Array, " max_num_splits"],
    action: Int[Array, " max_num_splits"],
    random_key: PRNGKeyArray,
    sub_mutate_config: SubstationMutationConfig,
    action_set: ActionSet,
) -> tuple[
    Int[Array, " max_num_splits"],
    Int[Array, " max_num_splits"],
    PRNGKeyArray,
]:
    """Mutate a single substation, changing the sub_ids, branch and inj topos.

    The sub-ids are implicit to the branch topo action index, however we pass them explicitely
    to aid substation mutation.
    Impossible mutations (e.g. adding a split when there is no room to add one,
    or removing a split when there are none) are handled by setting the probabilities
    for these mutations to zero, and normalising the remaining probabilities to sum to 1.

    Parameters
    ----------
    sub_ids : Int[Array, " max_num_splits"]
        The substation ids before mutation
    action : Int[Array, " max_num_splits"]
        The branch topology before mutation
    random_key : PRNGKeyArray
        The random key to use for the mutation
    sub_mutate_config : SubstationMutationConfig
        The configuration for the substation mutation,
        containing the probabilities for the different mutation types and the
        number of relevant substations in the grid, which is needed to determine
        the valid range of substation ids for mutation.
    action_set : ActionSet
        The actions for every substation. If not provided, will sample from all possible
        actions per substation

    Returns
    -------
    sub_ids : Int[Array, " max_num_splits"]
        The substation ids after mutation
    action : Int[Array, " max_num_splits"]
        The branch topology after mutation
    random_key : PRNGKeyArray
        The random key used for the mutation
    """
    int_max_value = int_max()
    operation_key, sub_key, action_key, random_key = jax.random.split(random_key, 4)

    # Gather probabilities for the different mutation operations
    add_split_prob = sub_mutate_config.add_split_prob
    change_split_prob = sub_mutate_config.change_split_prob
    remove_split_prob = sub_mutate_config.remove_split_prob
    prob_remain = 1 - add_split_prob - change_split_prob - remove_split_prob

    # Gather config values for the mutation operations
    n_rel_subs = sub_mutate_config.n_rel_subs
    n_max_splits = sub_ids.shape[0]
    is_split = sub_ids != int_max_value
    n_splits = jnp.sum(is_split)

    # Determine which mutation operations are allowed based on the current topology, and adjust probabilities accordingly
    allow_add = n_splits < n_max_splits
    # We can only remove a split if there is at least one split to remove
    allow_remove = n_splits > 0
    # We can only replace a split if there is at least one split to replace
    allow_replace = n_splits > 0
    # We can always choose to remain unchanged
    allow_remain = True

    probs = jnp.array([add_split_prob, remove_split_prob, change_split_prob, prob_remain], dtype=float)
    # assert jnp.isclose(probs.sum(), 1.0), f"Probabilities must sum to 1, but got {probs}"
    allowed = jnp.array([allow_add, allow_remove, allow_replace, allow_remain], dtype=bool)
    probs = jnp.where(allowed, probs, 0.0)
    probs = jnp.where(jnp.sum(probs) > 0, probs / jnp.sum(probs), jnp.array([0.0, 0.0, 0.0, 1.0]))

    # Pick one of the mutations at random
    chosen_op = jax.random.choice(a=probs.shape[0], shape=(), p=probs, key=operation_key)

    changed_indices, new_sub_ids = jax.lax.switch(
        chosen_op,
        [
            lambda _sub_ids: split_additional_sub(sub_key, _sub_ids, n_rel_subs, int_max_value),
            lambda _sub_ids: unsplit_substation(sub_key, _sub_ids, int_max_value),
            lambda _sub_ids: change_split_substation(sub_key, _sub_ids, n_rel_subs, int_max_value),
            lambda _sub_ids: do_nothing(int_max_value),  # remain unchanged
        ],
        sub_ids,
    )
    sub_ids = sub_ids.at[changed_indices].set(new_sub_ids, mode="drop")

    # Update the action for the changed substation
    new_action = sample_action_index_from_branch_actions(
        rng_key=action_key,
        sub_id=new_sub_ids,
        branch_action_set=action_set,
    )

    action = action.at[changed_indices].set(new_action, mode="drop")
    return sub_ids, action, random_key

toop_engine_topology_optimizer.dc.genetic_functions.mutation.mutate_disconnections

Mutation functions for the disconnections in the genetic algorithm.

change_disconnected_branch

change_disconnected_branch(
    random_key, disconnections, n_disconnectable_branches
)

Change a disconnected branch in the topology to a different one.

This assumes, that one branch is already disconnected, otherwise there would be nothing to change.

PARAMETER	DESCRIPTION
random_key	The random key to use for the mutation TYPE: PRNGKey
disconnections	The disconnections before mutation of one individual TYPE: Int[Array, ' max_num_disconnections']
n_disconnectable_branches	The number of disconnectable branches in the action set. It is not necessary to know the contents of the action set here because we only sample an index into the action set. TYPE: int

RETURNS	DESCRIPTION
Int[Array, ' ']	The index that is changed
Int[Array, ' ']	The new disconnection that is added. If no disconnection was changed, this is set to int_max

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/genetic_functions/mutation/mutate_disconnections.py

def change_disconnected_branch(
    random_key: PRNGKeyArray, disconnections: Int[Array, " max_num_disconnections"], n_disconnectable_branches: int
) -> tuple[Int[Array, " "], Int[Array, " "]]:
    """Change a disconnected branch in the topology to a different one.

    This assumes, that one branch is already disconnected, otherwise there would be nothing to change.

    Parameters
    ----------
    random_key : jax.random.PRNGKey
        The random key to use for the mutation
    disconnections : Int[Array, " max_num_disconnections"]
        The disconnections before mutation of one individual
    n_disconnectable_branches: int
        The number of disconnectable branches in the action set. It is not necessary to know the
        contents of the action set here because we only sample an index into the action set.

    Returns
    -------
    Int[Array, " "]
        The index that is changed
    Int[Array, " "]
        The new disconnection that is added. If no disconnection was changed, this is set to int_max
    """
    int_max_value = int_max()
    random_index_key, random_disc_key = jax.random.split(random_key)
    is_disconnected = disconnections != int_max_value
    disc_idx = get_random_true_idx(random_index_key, is_disconnected, int_max_value)

    new_disc_id = sample_new_id(random_disc_key, disconnections, n_disconnectable_branches, disc_idx)
    return disc_idx, new_disc_id

reconnect_disconnected_branch

reconnect_disconnected_branch(random_key, disconnections)

Reconnect a disconnected branch in the topology.

PARAMETER	DESCRIPTION
random_key	The random key to use for the mutation TYPE: PRNGKey
disconnections	The disconnections before mutation of one individual TYPE: Int[Array, ' max_num_disconnections']

RETURNS	DESCRIPTION
Int[Array, ' ']	The index that is changed
Int[Array, ' ']	The new disconnection that is added. This is always int_max to indicate that the slot is now empty

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/genetic_functions/mutation/mutate_disconnections.py

def reconnect_disconnected_branch(
    random_key: PRNGKeyArray, disconnections: Int[Array, " max_num_disconnections"]
) -> tuple[Int[Array, " "], Int[Array, " "]]:
    """Reconnect a disconnected branch in the topology.

    Parameters
    ----------
    random_key : jax.random.PRNGKey
        The random key to use for the mutation
    disconnections : Int[Array, " max_num_disconnections"]
        The disconnections before mutation of one individual

    Returns
    -------
    Int[Array, " "]
        The index that is changed
    Int[Array, " "]
        The new disconnection that is added. This is always int_max to indicate that the slot is now empty
    """
    int_max_value = int_max()
    random_index_key = random_key
    is_disconnected = disconnections != int_max_value
    disc_idx = get_random_true_idx(random_index_key, is_disconnected, int_max_value)
    return disc_idx, jnp.array(int_max_value, dtype=int)

disconnect_additional_branch

disconnect_additional_branch(
    random_key, disconnections, n_disconnectable_branches
)

Add a new disconnection to the topology.

This assumes, that there is still room to add a disconnection, otherwise there would be no valid branch to add.

PARAMETER	DESCRIPTION
random_key	The random key to use for the mutation TYPE: PRNGKey
disconnections	The disconnections before mutation of one individual TYPE: Int[Array, ' max_num_disconnections']
n_disconnectable_branches	The number of disconnectable branches in the action set. It is not necessary to know the contents of the action set here because we only sample an index into the action set. TYPE: int

RETURNS	DESCRIPTION
Int[Array, ' ']	The index that is changed
Int[Array, ' ']	The new disconnection that is added. If no disconnection was added, this is set to int_max

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/genetic_functions/mutation/mutate_disconnections.py

def disconnect_additional_branch(
    random_key: PRNGKeyArray, disconnections: Int[Array, " max_num_disconnections"], n_disconnectable_branches: int
) -> tuple[Int[Array, " "], Int[Array, " "]]:
    """Add a new disconnection to the topology.

    This assumes, that there is still room to add a disconnection, otherwise there would be no valid branch to add.

    Parameters
    ----------
    random_key : jax.random.PRNGKey
        The random key to use for the mutation
    disconnections : Int[Array, " max_num_disconnections"]
        The disconnections before mutation of one individual
    n_disconnectable_branches: int
        The number of disconnectable branches in the action set. It is not necessary to know the
        contents of the action set here because we only sample an index into the action set.

    Returns
    -------
    Int[Array, " "]
        The index that is changed
    Int[Array, " "]
        The new disconnection that is added. If no disconnection was added, this is set to int_max
    """
    int_max_value = int_max()
    random_index_key, random_disc_key = jax.random.split(random_key)
    # List available disconnections
    is_disconnectable = disconnections == int_max_value
    disc_idx = get_random_true_idx(random_index_key, is_disconnectable, int_max_value)

    new_disc_id = sample_new_id(random_disc_key, disconnections, n_disconnectable_branches, int_max_value)
    return disc_idx, new_disc_id

mutate_disconnections

mutate_disconnections(
    random_key,
    sub_ids,
    disconnections,
    disconnection_mutation_config,
)

Mutate the disconnections of a single topology.

Impossible mutations (e.g. adding a disconnection when there is no room to add one, or removing a disconnection when there are none) are handled by setting the probabilities for these mutations to zero, and adding them to the remain probability.

PARAMETER	DESCRIPTION
random_key	The random key to use for the mutation TYPE: PRNGKey
sub_ids	The substation ids of the topology, which are needed to determine whether certain disconnection mutations are allowed TYPE: Int[Array, ' max_num_splits']
disconnections	The disconnections before mutation of one individual TYPE: Int[Array, ' max_num_disconnections']
disconnection_mutation_config	The configuration for the disconnection mutation, containing the probabilities for the different mutation types and the number of disconnectable branches in the grid, which is needed to determine the valid range of branch ids for mutation. TYPE: DisconnectionMutationConfig

RETURNS	DESCRIPTION
Int[Array, ' max_num_disconnections']	The mutated disconnections

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/genetic_functions/mutation/mutate_disconnections.py

def mutate_disconnections(
    random_key: PRNGKeyArray,
    sub_ids: Int[Array, " max_num_splits"],
    disconnections: Int[Array, " max_num_disconnections"],
    disconnection_mutation_config: DisconnectionMutationConfig,
) -> Int[Array, " max_num_disconnections"]:
    """Mutate the disconnections of a single topology.

    Impossible mutations (e.g. adding a disconnection when there is no room to add one,
    or removing a disconnection when there are none) are handled by setting the probabilities
    for these mutations to zero, and adding them to the remain probability.

    Parameters
    ----------
    random_key : jax.random.PRNGKey
        The random key to use for the mutation
    sub_ids : Int[Array, " max_num_splits"]
        The substation ids of the topology,
        which are needed to determine whether certain disconnection mutations are allowed
    disconnections : Int[Array, " max_num_disconnections"]
        The disconnections before mutation of one individual
    disconnection_mutation_config : DisconnectionMutationConfig
        The configuration for the disconnection mutation,
        containing the probabilities for the different mutation types and the
        number of disconnectable branches in the grid,
        which is needed to determine the valid range of branch ids for mutation.

    Returns
    -------
    Int[Array, " max_num_disconnections"]
        The mutated disconnections
    """
    int_max_value = int_max()
    operation_key, disc_key, random_key = jax.random.split(random_key, 3)

    # Gather probabilities and config values from the mutation config
    add_disconnection_prob = disconnection_mutation_config.add_disconnection_prob
    change_disconnection_prob = disconnection_mutation_config.change_disconnection_prob
    remove_disconnection_prob = disconnection_mutation_config.remove_disconnection_prob

    n_disconnectable_branches = disconnection_mutation_config.n_disconnectable_branches

    # Gather info about the current topology
    has_splits = jnp.any(sub_ids != int_max_value)
    n_disconnections = jnp.sum(disconnections != int_max_value)
    max_num_disconnections = disconnections.shape[0]

    # Check which actions are allowed.
    # We only allow to add a disconnection if there are less disconnections than the maximum number of disconnections.
    allow_add = (n_disconnections < max_num_disconnections) & (add_disconnection_prob > 0.0)
    # We only allow to remove a disconnection if there is at least one disconnection,
    # and we don't want to end up with zero disconnections if there are no splits in the topology,
    # because then we would always end up with the same unsplit topology after mutation.
    allow_remove = (has_splits & (n_disconnections == 1)) | (n_disconnections > 1)

    # We only allow to change a disconnection if there is at least one disconnection,
    # otherwise there is nothing to change
    allow_replace = n_disconnections > 0

    allow_remain = True  # We can always choose to remain unchanged
    # We temporarily set the remain probability to 0, because we will add the probabilities
    # of the illegal actions to the remain probability later, after we set the illegal action probabilities to 0.
    temp_remain_prob = 0.0
    # Create an array of the probabilities for the different operations,
    # and set the probabilities to 0 for the operations that are not allowed.
    probs = jnp.array(
        [add_disconnection_prob, remove_disconnection_prob, change_disconnection_prob, temp_remain_prob], dtype=float
    )
    allowed = jnp.array([allow_add, allow_remove, allow_replace, allow_remain], dtype=bool)
    probs = jnp.where(allowed, probs, 0.0)

    # Replace all "illegal" operations with "remain unchanged".
    prob_sum = jnp.sum(probs)
    # If there are no splits, always add a disconnection
    # Otherwise, normalise the allowed probabilities to sum to 1
    # If probs are negative, only the remain option is considered
    probs = jnp.where(
        (~has_splits) & allow_add & (n_disconnections == 0),
        jnp.array([1.0, 0.0, 0.0, 0.0]),
        probs.at[3].set(1.0 - prob_sum),
    )

    # Randomly choose which operation to perform based on the probabilities
    # At least one of the operations is executed
    chosen_op = jax.random.choice(operation_key, jnp.arange(4), p=probs)
    disc_idx, new_branch_id = jax.lax.switch(
        chosen_op,
        [
            lambda _disconnections: disconnect_additional_branch(disc_key, _disconnections, n_disconnectable_branches),
            lambda _disconnections: reconnect_disconnected_branch(disc_key, _disconnections),
            lambda _disconnections: change_disconnected_branch(disc_key, _disconnections, n_disconnectable_branches),
            lambda _disconnections: do_nothing(int_max_value),  # remain unchanged
        ],
        disconnections,
    )
    # Update the disconnections with the new branch id at the chosen index
    disconnections = disconnections.at[disc_idx].set(new_branch_id, mode="drop")
    return disconnections

toop_engine_topology_optimizer.dc.genetic_functions.mutation.mutate_nodal_inj

Mutation functions for the nodal injections in the genetic algorithm.

mutate_psts

mutate_psts(
    random_key,
    pst_taps,
    pst_n_taps,
    pst_starting_taps,
    pst_mutation_sigma,
    pst_mutation_probability=0.2,
    pst_reset_probability=0.1,
)

Mutate the PST taps of a single topology.

PARAMETER	DESCRIPTION
random_key	The random key to use for the mutation TYPE: PRNGKeyArray
pst_taps	The PST tap positions before mutation TYPE: Int[Array, ' n_controllable_pst']
pst_n_taps	The number of taps for each PST. If a PST has N taps in this array, then it is assumed that all taps from 0 to N-1 are valid tap positions. Output taps will be clipped to this range. TYPE: Int[Array, ' n_controllable_pst']
pst_starting_taps	The initial PST taps of each controllable PST. TYPE: Int[Array, ' n_controllable_pst']
pst_mutation_sigma	The sigma to use for the normal distribution to sample the mutation from. The mutation will be sampled as an integer from a normal distribution with mean 0 and sigma pst_mutation_sigma. TYPE: float | int
pst_mutation_probability	The probability for an individual PST to be selected for mutation. 1.0 indicates that all PSTs will be mutated, 0.0 indicates that no PSTs will be mutated. Default 0.2 TYPE: float DEFAULT: 0.2
pst_reset_probability	The probability for an individual PST to be reverted to its initial set point. A value of 0.0 means no reset. A value of 1.0 means all PSTs will be reset. Default 0.1 TYPE: float DEFAULT: 0.1

RETURNS	DESCRIPTION
Int[Array, ' n_controllable_pst']	The mutated PST tap positions, clipped to the valid range of taps for each PST.

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/genetic_functions/mutation/mutate_nodal_inj.py

def mutate_psts(
    random_key: PRNGKeyArray,
    pst_taps: Int[Array, " n_controllable_pst"],
    pst_n_taps: Int[Array, " n_controllable_pst"],
    pst_starting_taps: Int[Array, " n_controllable_pst"],
    pst_mutation_sigma: float | int,
    pst_mutation_probability: float = 0.2,
    pst_reset_probability: float = 0.1,
) -> Int[Array, " n_controllable_pst"]:
    """Mutate the PST taps of a single topology.

    Parameters
    ----------
    random_key : PRNGKeyArray
        The random key to use for the mutation
    pst_taps : Int[Array, " n_controllable_pst"]
        The PST tap positions before mutation
    pst_n_taps : Int[Array, " n_controllable_pst"]
        The number of taps for each PST. If a PST has N taps in this array, then it is assumed that all taps from
        0 to N-1 are valid tap positions. Output taps will be clipped to this range.
    pst_starting_taps: Int[Array, " n_controllable_pst"]
        The initial PST taps of each controllable PST.
    pst_mutation_sigma : float | int
        The sigma to use for the normal distribution to sample the mutation from. The mutation will be sampled as an
        integer from a normal distribution with mean 0 and sigma pst_mutation_sigma.
    pst_mutation_probability: float
        The probability for an individual PST to be selected for mutation. 1.0 indicates that all PSTs will be mutated,
        0.0 indicates that no PSTs will be mutated. Default 0.2
    pst_reset_probability: float
        The probability for an individual PST to be reverted to its initial set point. A value of 0.0 means no reset. A
        value of 1.0 means all PSTs will be reset. Default 0.1

    Returns
    -------
    Int[Array, " n_controllable_pst"]
        The mutated PST tap positions, clipped to the valid range of taps for each PST.
    """
    # Sample number of PSTs to adjust from a n_controllable_pst-dimensional uniform distribution
    key, key_mutate, key_reset = jax.random.split(random_key, 3)

    pst_indices_to_mutate = jax.random.bernoulli(key=key, p=pst_mutation_probability, shape=pst_taps.shape)

    # Keep the sample shape static so this function can run under vmap/jit.
    mutation_samples = jax.random.normal(key_mutate, shape=pst_taps.shape) * pst_mutation_sigma
    mutation = jnp.where(pst_indices_to_mutate, mutation_samples, 0.0)
    mutation = jnp.round(mutation).astype(int)
    new_pst_taps = pst_taps + mutation

    # Reset random PSTs
    pst_indices_to_reset = jax.random.bernoulli(key=key_reset, p=pst_reset_probability, shape=pst_taps.shape)
    new_pst_taps = jnp.where(pst_indices_to_reset, pst_starting_taps, new_pst_taps)

    new_pst_taps = jnp.clip(new_pst_taps, a_min=0, a_max=pst_n_taps - 1)
    return new_pst_taps

mutate_nodal_injections

mutate_nodal_injections(
    random_key, nodal_inj_info, nodal_mutation_config
)

Mutate the nodal injection optimization results, currently only the PST taps.

PARAMETER	DESCRIPTION
random_key	The random key to use for the mutation TYPE: PRNGKeyArray
nodal_inj_info	The nodal injection optimization results before mutation. If None, no mutation is performed and None is returned. TYPE: Optional[NodalInjOptimResults]
nodal_mutation_config	The configuration for the nodal injection mutation. If None, no mutation is performed TYPE: Optional[NodalInjectionMutationConfig]

RETURNS	DESCRIPTION
Optional[NodalInjOptimResults]	The mutated nodal injection optimization results. If nodal_inj_info was None, returns None.

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/genetic_functions/mutation/mutate_nodal_inj.py

def mutate_nodal_injections(
    random_key: PRNGKeyArray,
    nodal_inj_info: Optional[NodalInjOptimResults],
    nodal_mutation_config: Optional[NodalInjectionMutationConfig],
) -> Optional[NodalInjOptimResults]:
    """Mutate the nodal injection optimization results, currently only the PST taps.

    Parameters
    ----------
    random_key : PRNGKeyArray
        The random key to use for the mutation
    nodal_inj_info : Optional[NodalInjOptimResults]
        The nodal injection optimization results before mutation. If None, no mutation is performed and None is returned.
    nodal_mutation_config : Optional[NodalInjectionMutationConfig]
        The configuration for the nodal injection mutation. If None, no mutation is performed

    Returns
    -------
    Optional[NodalInjOptimResults]
        The mutated nodal injection optimization results. If nodal_inj_info was None, returns None.
    """
    if nodal_inj_info is None:
        return None

    if nodal_mutation_config is None or nodal_mutation_config.pst_mutation_sigma <= 0:
        return nodal_inj_info

    batch_size = nodal_inj_info.pst_tap_idx.shape[0]
    n_timesteps = nodal_inj_info.pst_tap_idx.shape[1]
    random_key = jax.random.split(random_key, (batch_size, n_timesteps))

    # vmap to mutate the PST taps for each timestep + batch independently
    new_pst_taps = jax.vmap(
        jax.vmap(
            partial(
                mutate_psts,
                pst_n_taps=nodal_mutation_config.pst_n_taps,
                pst_starting_taps=nodal_mutation_config.pst_start_tap_idx,
                pst_mutation_sigma=nodal_mutation_config.pst_mutation_sigma,
                pst_mutation_probability=nodal_mutation_config.pst_mutation_probability,
                pst_reset_probability=nodal_mutation_config.pst_reset_probability,
            )
        )
    )(
        random_key=random_key,
        pst_taps=nodal_inj_info.pst_tap_idx.astype(int),
    )

    return NodalInjOptimResults(pst_tap_idx=new_pst_taps)

toop_engine_topology_optimizer.dc.genetic_functions.mutation.utils

Mutation utility functions for the genetic algorithm.

sample_new_id

sample_new_id(
    random_key,
    already_used_ids,
    n_available_ids,
    ignored_idx,
)

Sample a relevant id that is not already used in the active split slots.

PARAMETER	DESCRIPTION
random_key	Random key used for sampling. TYPE: PRNGKeyArray
already_used_ids	Current ids that should not be sampled again. TYPE: Int[Array, ' max_num_splits']
n_available_ids	Number of available ids that may be selected. TYPE: int
ignored_idx	Index in already_used_ids to ignore while checking duplicates. This is used for branch replacement, where the currently replaced split must not block its own id. TYPE: Int[Array, ' ']

RETURNS	DESCRIPTION
Int[Array, ' ']	A valid new substation id, or int_max() if none are available.

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/genetic_functions/mutation/utils.py

def sample_new_id(
    random_key: PRNGKeyArray,
    already_used_ids: Int[Array, " n_splits_or_disconnections"],
    n_available_ids: int,
    ignored_idx: Int[ArrayLike, " "],
) -> Int[Array, " "]:
    """Sample a relevant id that is not already used in the active split slots.

    Parameters
    ----------
    random_key : PRNGKeyArray
        Random key used for sampling.
    already_used_ids : Int[Array, " max_num_splits"]
        Current ids that should not be sampled again.
    n_available_ids : int
        Number of available ids that may be selected.
    ignored_idx : Int[Array, " "]
        Index in ``already_used_ids`` to ignore while checking duplicates. This is used for branch replacement,
        where the currently replaced split must not block its own id.

    Returns
    -------
    Int[Array, " "]
        A valid new substation id, or ``int_max()`` if none are available.
    """
    int_max_value = int_max()
    available_mask = jnp.ones((n_available_ids,), dtype=bool).at[already_used_ids].set(False, mode="drop")
    available_mask = available_mask.at[already_used_ids.at[ignored_idx].get(mode="fill", fill_value=int_max_value)].set(
        True, mode="drop"
    )
    return get_random_true_idx(random_key, available_mask, int_max_value)

get_random_true_idx

get_random_true_idx(
    random_key, boolean_array, int_max_value
)

Return random index of a True entry, or int_max_value if none are True.

PARAMETER	DESCRIPTION
random_key	The random key to use for the sampling TYPE: PRNGKeyArray
boolean_array	The boolean array to sample from. Should have shape (n_possibilities,). Is true where the index is a valid choice to sample, and false where it is not. If all entries are false, int_max_value is returned. TYPE: Bool[Array, ' n_possibilities']
int_max_value	The value to return if all entries in boolean_array are False. This should be the maximum integer value used to indicate an empty slot in the genotype. TYPE: int

RETURNS	DESCRIPTION
Int[Array, ' ']	A random index of a True entry in boolean_array, or int_max_value if all entries are False.

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/genetic_functions/mutation/utils.py

def get_random_true_idx(
    random_key: PRNGKeyArray, boolean_array: Bool[Array, " n_possibilities"], int_max_value: int
) -> Int[Array, " "]:
    """Return random index of a True entry, or int_max_value if none are True.

    Parameters
    ----------
    random_key : PRNGKeyArray
        The random key to use for the sampling
    boolean_array : Bool[Array, " n_possibilities"]
        The boolean array to sample from. Should have shape (n_possibilities,).
        Is true where the index is a valid choice to sample, and false where it is not.
        If all entries are false, int_max_value is returned.
    int_max_value : int
        The value to return if all entries in boolean_array are False. This should be the maximum integer value used
        to indicate an empty slot in the genotype.

    Returns
    -------
    Int[Array, " "]
        A random index of a True entry in boolean_array, or int_max_value if all entries
        are False.
    """
    n_possibilities = boolean_array.shape[0]
    candidate_indices = jnp.nonzero(boolean_array, size=n_possibilities, fill_value=int_max_value)[0]
    n_candidates = jnp.sum(boolean_array)
    safe_n_candidates = jnp.maximum(n_candidates, 1)
    sampled_position = jax.random.randint(random_key, shape=(), minval=0, maxval=safe_n_candidates)
    sampled_index = candidate_indices[sampled_position]
    return jnp.where(n_candidates > 0, sampled_index, jnp.array(int_max_value))

do_nothing

do_nothing(int_max_value)

Return a no-op mutation index and value.

PARAMETER	DESCRIPTION
int_max_value	The value to use for the no-op mutation, which should be the maximum integer value used to indicate an empty slot in the genotype. TYPE: int

RETURNS	DESCRIPTION
Int[Array, ' ']	The index that is changed. This is set to int_max_value to indicate that no index is changed.
Int[Array, ' ']	The new value that is added. This is set to int_max_value to indicate that no value is added.

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/genetic_functions/mutation/utils.py

def do_nothing(int_max_value: int) -> tuple[Int[Array, " "], Int[Array, " "]]:
    """Return a no-op mutation index and value.

    Parameters
    ----------
    int_max_value : int
        The value to use for the no-op mutation, which should be the maximum integer value used
        to indicate an empty slot in the genotype.

    Returns
    -------
    Int[Array, " "]
        The index that is changed. This is set to int_max_value to indicate that no index is changed.
    Int[Array, " "]
        The new value that is added. This is set to int_max_value to indicate that no value is added.
    """
    return jnp.array(int_max_value), jnp.array(int_max_value)

Repertoire

toop_engine_topology_optimizer.dc.repertoire.discrete_map_elites

Core components of the MAP-Elites algorithm.

Adapted from QDax (https://github.com/adaptive-intelligent-robotics/QDax)

EmitterScores module-attribute

EmitterScores = PyTree

DiscreteMapElites

DiscreteMapElites(
    scoring_function,
    emitter,
    metrics_function,
    n_cells_per_dim,
    cell_depth=1,
    distributed=False,
)

Discrete MAP-Elites algorithm.

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/repertoire/discrete_map_elites.py

def __init__(
    self,
    scoring_function: Callable[
        [Genotype, PRNGKeyArray, PyTree],
        Tuple[Fitness, Descriptor, ExtraScores, EmitterScores, PRNGKeyArray, Genotype],
    ],
    emitter: Emitter,
    metrics_function: Callable[[DiscreteMapElitesRepertoire], Metrics],
    n_cells_per_dim: tuple[int, ...],
    cell_depth: PositiveInt = 1,
    distributed: bool = False,
) -> None:
    self._scoring_function = scoring_function
    self._emitter = emitter
    self._metrics_function = metrics_function
    self._distributed = distributed
    self._n_cells_per_dim = n_cells_per_dim
    self._cell_depth = cell_depth

init

init(genotypes, random_key, scoring_data)

Initialize a Map-Elites repertoire with an initial population of genotypes.

Requires the definition of centroids that can be computed with any method such as CVT or Euclidean mapping.

Before the repertoire is initialised, individuals are gathered from all the devices.

Args: genotypes: initial genotypes, pytree in which leaves have shape (batch_size, num_features) random_key: a random key used for stochastic operations. n_cells_per_dim: number of cells per dimension in the repertoire

RETURNS	DESCRIPTION
An initialized MAP-Elite repertoire with the initial state of the emitter,	and a random key.

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/repertoire/discrete_map_elites.py

def init(
    self,
    genotypes: Genotype,
    random_key: PRNGKeyArray,
    scoring_data: PyTree,
) -> Tuple[DiscreteMapElitesRepertoire, Optional[EmitterState], PRNGKeyArray]:
    """Initialize a Map-Elites repertoire with an initial population of genotypes.

    Requires the definition of centroids that can be computed with any method
    such as CVT or Euclidean mapping.

    Before the repertoire is initialised, individuals are gathered from all the
    devices.

    Args:
        genotypes: initial genotypes, pytree in which leaves
            have shape (batch_size, num_features)
        random_key: a random key used for stochastic operations.
        n_cells_per_dim: number of cells per dimension in the repertoire

    Returns
    -------
        An initialized MAP-Elite repertoire with the initial state of the emitter,
        and a random key.
    """
    # score initial genotypes
    (
        fitnesses,
        descriptors,
        extra_scores,
        emitter_scores,
        random_key,
        genotypes,
    ) = self._scoring_function(genotypes, random_key, scoring_data)

    # gather across all devices
    if self._distributed:
        (
            genotypes,
            fitnesses,
            descriptors,
            extra_scores,
        ) = jax.tree_util.tree_map(
            lambda x: jnp.concatenate(jax.lax.all_gather(x, axis_name="p"), axis=0),
            (genotypes, fitnesses, descriptors, extra_scores),
        )

    # init the repertoire
    repertoire = init_repertoire(
        genotypes=genotypes,
        descriptors=descriptors,
        fitnesses=fitnesses,
        extra_scores=extra_scores,
        n_cells_per_dim=self._n_cells_per_dim,
        cell_depth=self._cell_depth,
    )

    # get initial state of the emitter
    emitter_state, random_key = self._emitter.init(
        random_key=random_key,
        init_genotypes=genotypes,
    )

    # update emitter state
    emitter_state = self._emitter.state_update(
        emitter_state=emitter_state,
        repertoire=repertoire,
        genotypes=genotypes,
        fitnesses=fitnesses,
        descriptors=descriptors,
        extra_scores=emitter_scores,
    )

    return repertoire, emitter_state, random_key

update

update(repertoire, emitter_state, random_key, scoring_data)

Perform one iteration of the MAP-Elites algorithm.

1. A batch of genotypes is sampled in the repertoire and the genotypes are copied.
2. The copies are mutated and crossed-over
3. The obtained offsprings are scored and then added to the repertoire.

Before the repertoire is updated, individuals are gathered from all the devices.

Args: repertoire: the MAP-Elites repertoire emitter_state: state of the emitter random_key: a jax PRNG random key

RETURNS	DESCRIPTION
the updated MAP-Elites repertoire	the updated (if needed) emitter state metrics about the updated repertoire a new jax PRNG key

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/repertoire/discrete_map_elites.py

@partial(jax.jit, static_argnames=("self",))
def update(
    self,
    repertoire: DiscreteMapElitesRepertoire,
    emitter_state: Optional[EmitterState],
    random_key: PRNGKeyArray,
    scoring_data: PyTree,
) -> Tuple[DiscreteMapElitesRepertoire, Optional[EmitterState], Metrics, PRNGKeyArray]:
    """Perform one iteration of the MAP-Elites algorithm.

    1. A batch of genotypes is sampled in the repertoire and the genotypes
        are copied.
    2. The copies are mutated and crossed-over
    3. The obtained offsprings are scored and then added to the repertoire.

    Before the repertoire is updated, individuals are gathered from all the
    devices.

    Args:
        repertoire: the MAP-Elites repertoire
        emitter_state: state of the emitter
        random_key: a jax PRNG random key

    Returns
    -------
        the updated MAP-Elites repertoire
        the updated (if needed) emitter state
        metrics about the updated repertoire
        a new jax PRNG key
    """
    # generate offsprings with the emitter
    genotypes, _extra_info, random_key = self._emitter.emit(repertoire, emitter_state, random_key)
    # scores the offsprings
    (
        fitnesses,
        descriptors,
        extra_scores,
        emitter_info,
        random_key,
        genotypes,
    ) = self._scoring_function(genotypes, random_key, scoring_data)

    # gather across all devices
    if self._distributed:
        (
            genotypes,
            fitnesses,
            descriptors,
            extra_scores,
        ) = jax.tree_util.tree_map(
            lambda x: jnp.concatenate(jax.lax.all_gather(x, axis_name="p"), axis=0),
            (genotypes, fitnesses, descriptors, extra_scores),
        )

    # add genotypes in the repertoire
    repertoire = add_to_repertoire(
        repertoire=repertoire,
        batch_of_genotypes=genotypes,
        batch_of_descriptors=descriptors,
        batch_of_fitnesses=fitnesses,
        batch_of_extra_scores=extra_scores,
    )

    # update emitter state after scoring is made
    emitter_state = self._emitter.state_update(
        emitter_state=emitter_state,
        repertoire=repertoire,
        genotypes=genotypes,
        fitnesses=fitnesses,
        descriptors=descriptors,
        extra_scores=emitter_info,
    )

    # update the metrics
    metrics = self._metrics_function(repertoire)

    return repertoire, emitter_state, metrics, random_key

toop_engine_topology_optimizer.dc.repertoire.discrete_me_repertoire

Contains the DiscreteMapElitesRepertoire class and utility functions.

This file contains util functions and a class to define a repertoire, used to store individuals in the MAP-Elites algorithm as well as several variants. Adapted from QDax (https://github.com/adaptive-intelligent-robotics/QDax)

DiscreteMapElitesRepertoire

Bases: Module

A class to store the MAP-Elites repertoire.

genotypes instance-attribute

genotypes

The genotypes in the repertoire.

The PyTree can be a simple Jax array or a more complex nested structure such as to represent parameters of neural network in Flax.

fitnesses instance-attribute

fitnesses

The fitness of solutions in each cell of the repertoire.

descriptors instance-attribute

descriptors

The descriptors of solutions in each cell of the repertoire.

extra_scores instance-attribute

extra_scores

The extra scores of solutions in each cell of the repertoire. Usually the metrics

n_cells_per_dim class-attribute instance-attribute

n_cells_per_dim = field(static=True)

The number of cells per dimension.

cell_depth class-attribute instance-attribute

cell_depth = field(static=True, default=1)

Each cell contains cell_depth unique individuals

sample

sample(random_key, num_samples)

Sample elements in the repertoire.

PARAMETER	DESCRIPTION
random_key	the random key to be used for sampling TYPE: PRNGKeyArray
num_samples	the number of samples to be drawn from the repertoire TYPE: int

RETURNS	DESCRIPTION
samples	the sampled genotypes TYPE: Genotype
random_key	the updated random key TYPE: jax PRNG key

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/repertoire/discrete_me_repertoire.py

@partial(jax.jit, static_argnames=("num_samples",))
def sample(
    self: "DiscreteMapElitesRepertoire", random_key: PRNGKeyArray, num_samples: int
) -> Tuple[Genotype, PRNGKeyArray]:
    """Sample elements in the repertoire.

    Parameters
    ----------
    random_key: jax PRNG key
        the random key to be used for sampling
    num_samples: int
        the number of samples to be drawn from the repertoire

    Returns
    -------
    samples: Genotype
        the sampled genotypes
    random_key: jax PRNG key
        the updated random key
    """
    repertoire_empty = self.fitnesses == -jnp.inf
    p = (1.0 - repertoire_empty) / jnp.sum(1.0 - repertoire_empty)

    random_key, subkey = jax.random.split(random_key)
    samples = jax.tree_util.tree_map(
        lambda x: jax.random.choice(subkey, x, shape=(num_samples,), p=p),
        self.genotypes,
    )

    return samples, random_key

__getitem__

__getitem__(index)

Get a slice of the repertoire.

PARAMETER	DESCRIPTION
index	the index of the elements to be selected TYPE: Union[int, slice, ndarray]

RETURNS	DESCRIPTION
a new repertoire with the selected elements

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/repertoire/discrete_me_repertoire.py

def __getitem__(self, index: Union[int, slice, jnp.ndarray]) -> "DiscreteMapElitesRepertoire":
    """Get a slice of the repertoire.

    Parameters
    ----------
    index: int, slice or jnp.ndarray
        the index of the elements to be selected

    Returns
    -------
        a new repertoire with the selected elements
    """
    return DiscreteMapElitesRepertoire(
        genotypes=jax.tree_util.tree_map(lambda x: x[index], self.genotypes),
        fitnesses=self.fitnesses[index],
        descriptors=self.descriptors[index],
        extra_scores=jax.tree_util.tree_map(lambda x: x[index], self.extra_scores),
        n_cells_per_dim=self.n_cells_per_dim,
        cell_depth=self.cell_depth,
    )

get_cells_indices

get_cells_indices(descriptors, n_cells_per_dim)

Compute the index for many descriptors at once.

Args: descriptors: an array that contains the descriptors of the solutions. Its shape is (batch_size, num_descriptors) n_cells_per_dim: the number of cells per dimension

RETURNS	DESCRIPTION
The index of the cell in which each descriptor belongs.

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/repertoire/discrete_me_repertoire.py

@partial(jax.jit, static_argnames=("n_cells_per_dim",))
def get_cells_indices(
    descriptors: Int[Array, " batch_size *dim"],
    n_cells_per_dim: tuple[int, ...],
) -> Int[Array, " batch_size"]:
    """Compute the index for many descriptors at once.

    Args:
        descriptors: an array that contains the descriptors of the
            solutions. Its shape is (batch_size, num_descriptors)
        n_cells_per_dim: the number of cells per dimension

    Returns
    -------
        The index of the cell in which each descriptor belongs.
    """
    return jax.vmap(partial(get_cell_index, n_cells_per_dim=n_cells_per_dim))(descriptors)

get_cell_index

get_cell_index(descriptor, n_cells_per_dim)

Compute the index of the cell in which each descriptor belongs.

The index is effectively like the reshape operation, spreading multiple dimensions to a flat single dimension.

Args: descriptor: an array that contains the descriptors. There is one descriptor per dimension n_cells_per_dim: the number of cells per dimension

RETURNS	DESCRIPTION
The index of the cell in which each descriptor belongs.

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/repertoire/discrete_me_repertoire.py

@partial(jax.jit, static_argnames=("n_cells_per_dim",))
def get_cell_index(
    descriptor: Int[Array, " n_dims"],
    n_cells_per_dim: tuple[int, ...],
) -> Int[Array, " "]:
    """Compute the index of the cell in which each descriptor belongs.

    The index is effectively like the reshape operation, spreading multiple dimensions to a flat
    single dimension.

    Args:
        descriptor: an array that contains the descriptors. There is one descriptor per dimension
        n_cells_per_dim: the number of cells per dimension

    Returns
    -------
        The index of the cell in which each descriptor belongs.
    """
    assert len(descriptor) == len(n_cells_per_dim)
    assert all(c > 0 for c in n_cells_per_dim)

    return jnp.ravel_multi_index(
        multi_index=descriptor,
        dims=n_cells_per_dim,
        mode="clip",
    )  # jittable thanks to clip mode

add_to_repertoire

add_to_repertoire(
    repertoire,
    batch_of_genotypes,
    batch_of_descriptors,
    batch_of_fitnesses,
    batch_of_extra_scores=None,
)

Add a batch of elements to the repertoire.

Args: repertoire: the MAP-Elites repertoire to which the elements will be added. batch_of_genotypes: a batch of genotypes to be added to the repertoire. Similarly to the repertoire.genotypes argument, this is a PyTree in which the leaves have a shape (batch_size, num_features) batch_of_descriptors: an array that contains the descriptors of the aforementioned genotypes. Its shape is (batch_size, num_descriptors). This will determine the cell in which the genotype will be stored, all descriptors are clipped to be within the bounds of n_cells_per_dim. batch_of_fitnesses: an array that contains the fitnesses of the aforementioned genotypes. Its shape is (batch_size,) batch_of_extra_scores: tree that contains the extra_scores of aforementioned genotypes.

RETURNS	DESCRIPTION
The updated MAP-Elites repertoire.

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/repertoire/discrete_me_repertoire.py

@eqx.filter_jit  # TODO. Why did this not fail before?
def add_to_repertoire(
    repertoire: DiscreteMapElitesRepertoire,
    batch_of_genotypes: Genotype,
    batch_of_descriptors: Int[Array, " batch_size n_dims"],
    batch_of_fitnesses: Float[Array, " batch_size"],
    batch_of_extra_scores: Optional[ExtraScores] = None,
) -> DiscreteMapElitesRepertoire:
    """Add a batch of elements to the repertoire.

    Args:
        repertoire: the MAP-Elites repertoire to which the elements will be added.
        batch_of_genotypes: a batch of genotypes to be added to the repertoire.
            Similarly to the repertoire.genotypes argument, this is a PyTree in which
            the leaves have a shape (batch_size, num_features)
        batch_of_descriptors: an array that contains the descriptors of the
            aforementioned genotypes. Its shape is (batch_size, num_descriptors). This will
            determine the cell in which the genotype will be stored, all descriptors are clipped
            to be within the bounds of n_cells_per_dim.
        batch_of_fitnesses: an array that contains the fitnesses of the
            aforementioned genotypes. Its shape is (batch_size,)
        batch_of_extra_scores: tree that contains the extra_scores of
            aforementioned genotypes.

    Returns
    -------
        The updated MAP-Elites repertoire.
    """
    if repertoire.cell_depth > 1:
        return add_to_repertoire_with_cell_depth(
            repertoire,
            batch_of_genotypes,
            batch_of_descriptors,
            batch_of_fitnesses,
            batch_of_extra_scores,
        )
    return add_to_repertoire_without_cell_depth(
        repertoire,
        batch_of_genotypes,
        batch_of_descriptors,
        batch_of_fitnesses,
        batch_of_extra_scores,
    )

add_to_repertoire_without_cell_depth

add_to_repertoire_without_cell_depth(
    repertoire,
    batch_of_genotypes,
    batch_of_descriptors,
    batch_of_fitnesses,
    batch_of_extra_scores=None,
)

Add a batch of elements to the repertoire.

PARAMETER	DESCRIPTION
repertoire	the MAP-Elites repertoire to which the elements will be added. TYPE: DiscreteMapElitesRepertoire
batch_of_genotypes	a batch of genotypes to be added to the repertoire. TYPE: Genotype
batch_of_descriptors	an array that contains the descriptors of the genotypes. TYPE: Int[Array, ' batch_size n_dims']
batch_of_fitnesses	an array that contains the fitnesses of the genotypes. TYPE: Float[Array, ' batch_size']
batch_of_extra_scores	tree that contains the extra_scores of genotypes. TYPE: Optional[ExtraScores] DEFAULT: None

RETURNS	DESCRIPTION
DiscreteMapElitesRepertoire	The updated MAP-Elites repertoire.

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/repertoire/discrete_me_repertoire.py

@jax.jit
def add_to_repertoire_without_cell_depth(
    repertoire: DiscreteMapElitesRepertoire,
    batch_of_genotypes: Genotype,
    batch_of_descriptors: Int[Array, " batch_size n_dims"],
    batch_of_fitnesses: Float[Array, " batch_size"],
    batch_of_extra_scores: Optional[ExtraScores] = None,
) -> DiscreteMapElitesRepertoire:
    """Add a batch of elements to the repertoire.

    Parameters
    ----------
    repertoire: DiscreteMapElitesRepertoire
        the MAP-Elites repertoire to which the elements will be added.
    batch_of_genotypes: Genotype
        a batch of genotypes to be added to the repertoire.
    batch_of_descriptors: Int[Array, " batch_size n_dims"]
        an array that contains the descriptors of the genotypes.
    batch_of_fitnesses: Float[Array, " batch_size"]
        an array that contains the fitnesses of the genotypes.
    batch_of_extra_scores: Optional[ExtraScores]
        tree that contains the extra_scores of genotypes.

    Returns
    -------
    DiscreteMapElitesRepertoire
        The updated MAP-Elites repertoire.
    """
    batch_of_indices = get_cells_indices(batch_of_descriptors, repertoire.n_cells_per_dim)
    batch_of_indices = jnp.expand_dims(batch_of_indices, axis=-1)
    batch_of_fitnesses = jnp.expand_dims(batch_of_fitnesses, axis=-1)
    repertoire_size = repertoire.fitnesses.shape[0]

    # get fitness segment max
    # Necessary because there could be multiple scores with the same descriptor
    best_fitnesses = jax.ops.segment_max(
        batch_of_fitnesses,
        batch_of_indices.astype(jnp.int32).squeeze(axis=-1),
        num_segments=repertoire_size,
    )

    cond_values = jnp.take_along_axis(best_fitnesses, batch_of_indices, 0)

    # put dominated fitness to -jnp.inf
    batch_of_fitnesses = jnp.where(batch_of_fitnesses == cond_values, batch_of_fitnesses, -jnp.inf)

    # get addition condition
    repertoire_fitnesses = jnp.expand_dims(repertoire.fitnesses, axis=-1)
    current_fitnesses = jnp.take_along_axis(repertoire_fitnesses, batch_of_indices, 0)
    addition_condition = batch_of_fitnesses > current_fitnesses

    # assign fake position when relevant : num_centroids is out of bound
    batch_of_indices = jnp.where(addition_condition, batch_of_indices, repertoire_size).squeeze(axis=-1)

    # create new repertoire
    new_repertoire_genotypes = jax.tree_util.tree_map(
        lambda repertoire_genotypes, new_genotypes: repertoire_genotypes.at[batch_of_indices].set(
            new_genotypes, mode="drop"
        ),
        repertoire.genotypes,
        batch_of_genotypes,
    )

    # compute new fitness and descriptors
    new_fitnesses = repertoire.fitnesses.at[batch_of_indices].set(batch_of_fitnesses.squeeze(axis=-1), mode="drop")
    new_descriptors = repertoire.descriptors.at[batch_of_indices].set(batch_of_descriptors, mode="drop")
    if batch_of_extra_scores is None:
        new_extra_scores = {}
    else:
        new_extra_scores = jax.tree_util.tree_map(
            lambda repertoire_extra_scores, new_extra_scores: repertoire_extra_scores.at[batch_of_indices].set(
                new_extra_scores, mode="drop"
            ),
            repertoire.extra_scores,
            batch_of_extra_scores,
        )

    return DiscreteMapElitesRepertoire(
        genotypes=new_repertoire_genotypes,
        fitnesses=new_fitnesses,
        descriptors=new_descriptors,
        extra_scores=new_extra_scores,
        n_cells_per_dim=repertoire.n_cells_per_dim,
        cell_depth=repertoire.cell_depth,
    )

add_to_repertoire_with_cell_depth

add_to_repertoire_with_cell_depth(
    repertoire,
    batch_of_genotypes,
    batch_of_descriptors,
    batch_of_fitnesses,
    batch_of_extra_scores=None,
)

Add a batch of elements to a repertoire with depth.

Assumes the repertoire puts elements on a same depth level next to another (layer1 layer1 layer2 layer2 ...) Manipulates abstract indices instead of moving genotypes directly.

PARAMETER	DESCRIPTION
repertoire	the MAP-Elites repertoire to which the elements will be added. TYPE: DiscreteMapElitesRepertoire
batch_of_genotypes	a batch of genotypes to be added to the repertoire. TYPE: Genotype
batch_of_descriptors	an array that contains the descriptors of the genotypes. TYPE: Int[Array, ' batch_size n_dims']
batch_of_fitnesses	an array that contains the fitnesses of the genotypes. TYPE: Float[Array, ' batch_size']
batch_of_extra_scores	tree that contains the extra_scores of genotypes. TYPE: Optional[ExtraScores] DEFAULT: None

RETURNS	DESCRIPTION
DiscreteMapElitesRepertoire	The updated MAP-Elites repertoire.

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/repertoire/discrete_me_repertoire.py

@jax.jit
def add_to_repertoire_with_cell_depth(
    repertoire: DiscreteMapElitesRepertoire,
    batch_of_genotypes: Genotype,
    batch_of_descriptors: Int[Array, " batch_size n_dims"],
    batch_of_fitnesses: Float[Array, " batch_size"],
    batch_of_extra_scores: Optional[ExtraScores] = None,
) -> DiscreteMapElitesRepertoire:
    """Add a batch of elements to a repertoire with depth.

    Assumes the repertoire puts elements on a same depth level next to another (layer1 layer1 layer2 layer2 ...)
    Manipulates abstract indices instead of moving genotypes directly.

    Parameters
    ----------
    repertoire: DiscreteMapElitesRepertoire
        the MAP-Elites repertoire to which the elements will be added.
    batch_of_genotypes: Genotype
        a batch of genotypes to be added to the repertoire.
    batch_of_descriptors: Int[Array, " batch_size n_dims"]
        an array that contains the descriptors of the genotypes.
    batch_of_fitnesses: Float[Array, " batch_size"]
        an array that contains the fitnesses of the genotypes.
    batch_of_extra_scores: Optional[ExtraScores]
        tree that contains the extra_scores of genotypes.

    Returns
    -------
    DiscreteMapElitesRepertoire
        The updated MAP-Elites repertoire.
    """
    cell_depth = repertoire.cell_depth
    num_cells = np.prod(np.array(repertoire.n_cells_per_dim)).item()
    repertoire_size = cell_depth * num_cells
    batch_size = batch_of_fitnesses.shape[0]

    # Work on indices in order to abstract the genotype
    abstract_current_genotypes = jnp.arange(repertoire_size)
    abstract_new_genotypes = jnp.arange(start=repertoire_size, stop=repertoire_size + batch_size)

    """
    Part 1 : Sort
    """

    # rearrange repertoire fitnesses so cell fitnesses are in the same dimension
    current_fitnesses_per_cell: Float[Array, "num_cells cell_depth"] = jnp.reshape(
        repertoire.fitnesses,
        (-1, num_cells),  # one dimension per cell
    ).T  # transpose to have cells in the first dimension

    # calculate new indices for each layer
    new_indices: Int[Array, "batch_size"] = get_cells_indices(
        batch_of_descriptors, repertoire.n_cells_per_dim
    )  # example : [0, 0, 0, 1]

    # repeat to get :
    # [[0, 0, 0, 1],
    #  [0, 0, 0, 1],
    #  [0, 0, 0, 1]]
    # shape (num_cells, batch_size)
    repeated_indices: Int[Array, "num_cells batch_size"] = jnp.tile(new_indices, reps=(num_cells, 1))

    # Create an array of shape (num_cells, batch_size) where the nth element contains the number n
    # example : [[0, 0, 0, 0],
    #            [1, 1, 1, 1],
    #            [2, 2, 2, 2]]
    cell_number: Int[Array, "num_cells"] = jnp.arange(num_cells)
    repeated_cell_number: Int[Array, "num_cells batch_size"] = jnp.repeat(
        cell_number.reshape(-1, 1),
        repeats=batch_size,
        axis=-1,
        total_repeat_length=batch_size,  # to make it jittable
    )

    # Repeat new fitnesses
    repeated_new_fitnesses: Float[Array, "num_cells batch_size"] = jnp.tile(batch_of_fitnesses, reps=(num_cells, 1))

    # condition is true if the fitness is at the right index
    # example : [[True , True , True , False],
    #            [False, False, False, True ],
    #            [False, False, False, False]]
    belongs_to_right_cell = repeated_cell_number == repeated_indices

    filtered_new_fitnesses_per_cell: Float[Array, "num_cells batch_size"] = jnp.where(
        belongs_to_right_cell,
        repeated_new_fitnesses,
        -jnp.inf,
    )

    # extend current fitnesses per cell to welcome new ones
    all_fitnesses_per_cell: Float[Array, "num_cells max_size_per_cell"] = jnp.concatenate(
        [current_fitnesses_per_cell, filtered_new_fitnesses_per_cell], axis=1
    )

    sorted_fitness_indices_per_cell: Int[Array, "num_cells max_size_per_cell"] = jnp.argsort(
        all_fitnesses_per_cell, axis=-1, descending=True
    )

    # we only want the first cell_depth elements of each cell
    cropped_best_fitness_indices_per_cell: Int[Array, "num_cells cell_depth"] = sorted_fitness_indices_per_cell[
        :, :cell_depth
    ]

    # adapt indices to select genotypes : must be laid flat with first num_cells belonging to the first depth etc
    # adapted_indices_for_selection = cropped_best_fitness_indices_per_cell.T.reshape(-1)

    # Shape genotypes per-cell to use the indices
    current_genotypes_per_cell: Int[Array, "num_cells cell_depth"] = abstract_current_genotypes.reshape(-1, num_cells).T
    repeated_batch_of_genotypes: Int[Array, "num_cells batch_size"] = (
        jnp.repeat(abstract_new_genotypes, repeats=num_cells).reshape(-1, num_cells).T
    )
    all_genotypes_per_cell: Int[Array, "num_cells max_size_per_cell"] = jnp.concatenate(
        [current_genotypes_per_cell, repeated_batch_of_genotypes],
        axis=1,  # concatenate on cells
    )

    # select the genotypes by order of fitness. If a cell doesn't belong there,
    # it will be placed last (-inf fitness) and thus cropped out
    cropped_selected_genotypes_per_cell: Int[Array, "num_cells cell_depth"] = jnp.take_along_axis(
        all_genotypes_per_cell, cropped_best_fitness_indices_per_cell, axis=1
    )

    # reshape to the original shape to apply on a concat of repertoire and new batch
    final_indices_selection: Int[Array, "repertoire_size"] = cropped_selected_genotypes_per_cell.T.reshape(
        num_cells * cell_depth
    )

    """
    Part 2 : Selection
    """

    # genotypes
    selected_genotypes = jax.tree.map(
        lambda x, y: jnp.concat([x, y])[final_indices_selection],
        repertoire.genotypes,
        batch_of_genotypes,
    )

    # fitness
    selected_fitness = jnp.concat([repertoire.fitnesses, batch_of_fitnesses])[final_indices_selection]

    # descriptors
    selected_descriptors = jnp.concat([repertoire.descriptors, batch_of_descriptors], axis=0)[final_indices_selection]

    # extra scores
    if batch_of_extra_scores is None:
        # TODO IS THIS INTENDED
        selected_extra_scores = {}
    else:
        selected_extra_scores = jax.tree_util.tree_map(
            lambda x, y: jnp.concatenate([x, y], axis=0)[final_indices_selection],
            repertoire.extra_scores,
            batch_of_extra_scores,
        )

    return DiscreteMapElitesRepertoire(
        genotypes=selected_genotypes,
        fitnesses=selected_fitness,
        descriptors=selected_descriptors,
        extra_scores=selected_extra_scores,
        n_cells_per_dim=repertoire.n_cells_per_dim,
        cell_depth=cell_depth,
    )

init_repertoire

init_repertoire(
    genotypes,
    fitnesses,
    descriptors,
    extra_scores,
    n_cells_per_dim,
    cell_depth,
)

Initialize a Map-Elites repertoire with an initial population of genotypes.

Requires the definition of centroids that can be computed with any method such as CVT or Euclidean mapping.

Note: this function has been kept outside of the object MapElites, so it can be called easily called from other modules.

Args: genotypes: initial genotypes, pytree in which leaves have shape (batch_size, num_features) fitnesses: fitness of the initial genotypes of shape (batch_size,) descriptors: descriptors of the initial genotypes of shape (batch_size, num_descriptors) extra_scores: the observed load flow metrics n_cells_per_dim: the number of cells per dimension cell_depth: the number of topologies per cell

RETURNS	DESCRIPTION
an initialized MAP-Elite repertoire

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/repertoire/discrete_me_repertoire.py

def init_repertoire(
    genotypes: Genotype,
    fitnesses: Fitness,
    descriptors: Descriptor,
    extra_scores: Optional[ExtraScores],
    n_cells_per_dim: tuple[int, ...],
    cell_depth: Static[int],
) -> DiscreteMapElitesRepertoire:
    """Initialize a Map-Elites repertoire with an initial population of genotypes.

    Requires the definition of centroids that can be computed with any method
    such as CVT or Euclidean mapping.

    Note: this function has been kept outside of the object MapElites, so it can
    be called easily called from other modules.

    Args:
        genotypes: initial genotypes, pytree in which leaves
            have shape (batch_size, num_features)
        fitnesses: fitness of the initial genotypes of shape (batch_size,)
        descriptors: descriptors of the initial genotypes
            of shape (batch_size, num_descriptors)
        extra_scores: the observed load flow metrics
        n_cells_per_dim: the number of cells per dimension
        cell_depth: the number of topologies per cell

    Returns
    -------
        an initialized MAP-Elite repertoire
    """
    # retrieve one genotype from the population
    (first_genotype, first_extra_score) = jax.tree_util.tree_map(
        lambda x: x[0] if x is not None else None, (genotypes, extra_scores)
    )

    # create a repertoire with default values
    repertoire = _init_default(
        genotype=first_genotype,
        extra_scores=first_extra_score,
        n_cells_per_dim=n_cells_per_dim,
        cell_depth=cell_depth,
    )

    # add initial population to the repertoire
    new_repertoire = add_to_repertoire(repertoire, genotypes, descriptors, fitnesses, extra_scores)

    return new_repertoire  # type: ignore

toop_engine_topology_optimizer.dc.repertoire.plotting

Plotting functions for the Map-Elites algorithm.

plot_repertoire_1d

plot_repertoire_1d(
    fitnesses, n_cells_per_dim, descriptor_metrics
)

Plot a bar chart of the repertoire's fitnesses.

PARAMETER	DESCRIPTION
fitnesses	The fitnesses of the repertoire TYPE: Fitness
n_cells_per_dim	The number of cells per dimension TYPE: tuple[int, ...]
descriptor_metrics	The descriptor metrics TYPE: tuple[str, ...]

RETURNS	DESCRIPTION
Figure	The plot

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/repertoire/plotting.py

def plot_repertoire_1d(
    fitnesses: Fitness,
    n_cells_per_dim: tuple[int, ...],
    descriptor_metrics: tuple[str, ...],
) -> Figure:
    """Plot a bar chart of the repertoire's fitnesses.

    Parameters
    ----------
    fitnesses : Fitness
        The fitnesses of the repertoire
    n_cells_per_dim : tuple[int, ...]
        The number of cells per dimension
    descriptor_metrics : tuple[str, ...]
        The descriptor metrics

    Returns
    -------
    plt.Figure
        The plot
    """
    plt.bar(
        x=range(n_cells_per_dim[0]),
        height=fitnesses,
    )
    plt.xlabel(descriptor_metrics[0])
    plt.ylabel("Fitness")
    return plt.gcf()

plot_repertoire_2d

plot_repertoire_2d(
    fitnesses, n_cells_per_dim, descriptor_metrics
)

Plot a heatmap of the repertoire's fitnesses.

PARAMETER	DESCRIPTION
fitnesses	The fitnesses of the repertoire TYPE: Fitness
n_cells_per_dim	The number of cells per dimension TYPE: tuple[int, ...]
descriptor_metrics	The descriptor metrics TYPE: tuple[str, ...]

RETURNS	DESCRIPTION
Axes	The plot

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/repertoire/plotting.py

def plot_repertoire_2d(
    fitnesses: Fitness,
    n_cells_per_dim: tuple[int, ...],
    descriptor_metrics: tuple[str, ...],
) -> Axes:
    """Plot a heatmap of the repertoire's fitnesses.

    Parameters
    ----------
    fitnesses : Fitness
        The fitnesses of the repertoire
    n_cells_per_dim : tuple[int, ...]
        The number of cells per dimension
    descriptor_metrics : tuple[str, ...]
        The descriptor metrics

    Returns
    -------
    Axes
        The plot
    """
    reshaped_fitnesses = fitnesses.reshape(n_cells_per_dim)
    ax = sns.heatmap(
        reshaped_fitnesses,
        cbar_kws={"label": "Fitness"},
        linewidths=0.01,
        linecolor=(0.3, 0.3, 0.3, 0.3),
    )  # vmin=-15000, vmax=-5000
    ax.invert_yaxis()
    ax.set_xlabel(descriptor_metrics[1])  # first axis is vertical for some reason
    ax.set_ylabel(descriptor_metrics[0])
    return ax

plot_repertoire

plot_repertoire(
    fitnesses,
    iteration,
    folder,
    n_cells_per_dim,
    descriptor_metrics,
    save_plot,
)

Plot the repertoire (1D or 2D) and saves the figure.

PARAMETER	DESCRIPTION
fitnesses	The fitnesses of the repertoire TYPE: Fitness
iteration	The current iteration number TYPE: Optional[int]
folder	The folder to save the plot. Will create a "plots" folder inside. TYPE: str
n_cells_per_dim	The number of cells per dimension TYPE: tuple[int, ...]
descriptor_metrics	The descriptor metrics TYPE: tuple[str, ...]
save_plot	Whether to save the plot TYPE: bool

RETURNS	DESCRIPTION
Axes | Figure	The plot. Axes for heatmap, plt.Figure for barchart.

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/repertoire/plotting.py

def plot_repertoire(
    fitnesses: Fitness,
    iteration: Optional[int],
    folder: Optional[str],
    n_cells_per_dim: tuple[int, ...],
    descriptor_metrics: tuple[str, ...],
    save_plot: bool,
) -> Axes | plt.Figure:
    """Plot the repertoire (1D or 2D) and saves the figure.

    Parameters
    ----------
    fitnesses : Fitness
        The fitnesses of the repertoire
    iteration : Optional[int]
        The current iteration number
    folder : str
        The folder to save the plot. Will create a "plots" folder inside.
    n_cells_per_dim : tuple[int, ...]
        The number of cells per dimension
    descriptor_metrics : tuple[str, ...]
        The descriptor metrics
    save_plot : bool
        Whether to save the plot

    Returns
    -------
    Axes | plt.Figure
        The plot. Axes for heatmap, plt.Figure for barchart.
    """
    plt.clf()  # clear or it plots on top of itself

    # Prepare fitnesses : convert neginf to minimum fitness * weight
    weight = 1.05  # Any number greater than 1, can fine-tune to your liking
    minimum_fitness = jnp.min(fitnesses, where=jnp.isfinite(fitnesses), initial=0)
    fitnesses = jnp.nan_to_num(fitnesses, neginf=weight * minimum_fitness)

    # plot in N dimensions
    if len(n_cells_per_dim) == 1:  # can replace with plotting_fn = plot_repertoire_1d
        plot = plot_repertoire_1d(fitnesses, n_cells_per_dim, descriptor_metrics)
    elif len(n_cells_per_dim) == 2:
        plot = plot_repertoire_2d(fitnesses, n_cells_per_dim, descriptor_metrics)
    else:
        raise ValueError(
            "Only 1D and 2D plots are supported"
        )  # TODO : 3D heatmaps https://www.geeksforgeeks.org/3d-heatmap-in-python/

    # add iteration to the title
    if iteration is not None:
        plt.title(f"Map-Elites repertoire iteration {iteration}")
    else:
        plt.title("Map-Elites repertoire")

    # save the figure
    if save_plot:
        folder = os.path.join(folder, "plots")
        os.makedirs(folder, exist_ok=True)
        filename = "repertoire.png" if iteration is None else f"repertoire_{iteration}.png"
        plt.savefig(os.path.join(folder, filename))
    return plot

Worker

toop_engine_topology_optimizer.dc.worker.optimizer

Callable functions for the optimizer worker.

OptimizerData dataclass

OptimizerData(
    start_params,
    optimization_id,
    solver_configs,
    algo,
    initial_fitness,
    initial_metrics,
    jax_data,
    start_time,
    sent_topologies,
)

A wrapper dataclass for all the data stored by the optimizer.

Because this dataclass holds irrelevant information for the GPU optimization, it is split into two dataclasses. The parent (this one) is used to store all the information and the child (JaxOptimizerData) is used to store the information that is needed on the GPU.

start_params instance-attribute

start_params

The initial args for the optimization run

optimization_id instance-attribute

optimization_id

The id of the optimization run

solver_configs instance-attribute

solver_configs

The solver config for every timestep

algo instance-attribute

algo

The genetic algorithm object

initial_fitness instance-attribute

initial_fitness

The initial fitness value

initial_metrics instance-attribute

initial_metrics

The initial metrics

jax_data instance-attribute

jax_data

Everything that needs to live on GPU. This dataclass is updated per-iteration while OptimizerData is only updated per-epoch, hence there are points in the command flow where this variable is out of sync. At the end of an epoch, this dataclass is updated to match the state in the optimization.

start_time instance-attribute

start_time

The time the optimization run started

sent_topologies instance-attribute

sent_topologies

A set of strategy hashes that have already been sent to the results topic, to avoid sending duplicates.

initialize_optimization

initialize_optimization(
    params,
    optimization_id,
    static_information_files,
    processed_gridfile_fs,
)

Initialize the optimization run.

This function will be called at the start of the optimization run. It should be used to load the static information files and set up the optimization run.

PARAMETER	DESCRIPTION
params	The parameters for the optimization run TYPE: DCOptimizerParameters
optimization_id	The id of the optimization run, used to annotate results and heartbeats TYPE: str
static_information_files	The paths to the static information files to load TYPE: tuple[str | Path, ...]
processed_gridfile_fs	The target filesystem for the preprocessing worker. This contains all processed grid files. During the import job, a new folder import_results.data_folder was created which will be completed with the preprocess call to this function. Internally, only the data folder is passed around as a dirfs. Note that the unprocessed_gridfile_fs is not needed here anymore, as all preprocessing steps that need the unprocessed gridfiles were already done. TYPE: AbstractFileSystem

RETURNS	DESCRIPTION
OptimizerData	The data to store for the optimization run
list[StaticInformationStats]	The static information descriptions, will be sent via the heartbeats channel
Strategy	The initial strategy (unsplit) for the grid, including the initial fitness and metrics

RAISES	DESCRIPTION
Exception	If an error occurs during the initialization. It will be caught by the worker and sent back to the results topic

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/worker/optimizer.py

def initialize_optimization(
    params: DCOptimizerParameters,
    optimization_id: str,
    static_information_files: tuple[str | Path, ...],
    processed_gridfile_fs: AbstractFileSystem,
) -> tuple[OptimizerData, list[StaticInformationStats], Strategy]:
    """Initialize the optimization run.

    This function will be called at the start of the optimization run. It should be used to load
    the static information files and set up the optimization run.

    Parameters
    ----------
    params : DCOptimizerParameters
        The parameters for the optimization run
    optimization_id : str
        The id of the optimization run, used to annotate results and heartbeats
    static_information_files : tuple[str | Path, ...]
        The paths to the static information files to load
    processed_gridfile_fs: AbstractFileSystem
        The target filesystem for the preprocessing worker. This contains all processed grid files.
        During the import job,  a new folder import_results.data_folder was created
        which will be completed with the preprocess call to this function.
        Internally, only the data folder is passed around as a dirfs.
        Note that the unprocessed_gridfile_fs is not needed here anymore, as all preprocessing steps that need the
        unprocessed gridfiles were already done.


    Returns
    -------
    OptimizerData
        The data to store for the optimization run
    list[StaticInformationStats]
        The static information descriptions, will be sent via the heartbeats channel
    Strategy
        The initial strategy (unsplit) for the grid, including the initial fitness and metrics

    Raises
    ------
    Exception
        If an error occurs during the initialization. It will be caught by the worker and sent back to
        the results topic
    """
    (
        algo,
        jax_data,
        solver_configs,
        initial_fitness,
        initial_metrics,
        static_information_descriptions,
    ) = algo_setup(
        ga_args=params.ga_config,
        lf_args=params.loadflow_solver_config,
        double_limits=(
            params.double_limits.lower,
            params.double_limits.upper,
        )
        if params.double_limits is not None
        else None,
        static_information_files=static_information_files,
        processed_gridfile_fs=processed_gridfile_fs,
    )

    metrics = convert_metrics(initial_fitness, initial_metrics)

    # For now we send None as initial topology PST setpoints, as the AC solver can
    # not distinguish a topology with taps set to default from a topology without taps.
    initial_topology = Strategy(
        timesteps=[
            Topology(
                actions=[],
                disconnections=[],
                # pst_setpoints=di.nodal_injection_information.starting_tap_idx.tolist()
                # if di.nodal_injection_information is not None
                # else None,
                pst_setpoints=None,
                metrics=metrics,
            )
            for _di in jax_data.dynamic_informations
        ]
    )

    optimization_data = OptimizerData(
        start_params=params,
        optimization_id=optimization_id,
        solver_configs=solver_configs,
        algo=algo,
        jax_data=jax_data,
        initial_fitness=initial_fitness,
        initial_metrics=initial_metrics,
        start_time=time.time(),
        sent_topologies=set(),
    )

    return optimization_data, static_information_descriptions, initial_topology

convert_metrics

convert_metrics(fitness, metrics_dict)

Convert a metrics dictionary to a Metrics dataclass.

PARAMETER	DESCRIPTION
fitness	The fitness value TYPE: float
metrics_dict	The metrics dictionary TYPE: dict[MetricType, float]

RETURNS	DESCRIPTION
Metrics	The metrics dataclass

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/worker/optimizer.py

def convert_metrics(fitness: float, metrics_dict: dict[MetricType, float]) -> Metrics:
    """Convert a metrics dictionary to a Metrics dataclass.

    Parameters
    ----------
    fitness : float
        The fitness value

    metrics_dict : dict[MetricType, float]
        The metrics dictionary

    Returns
    -------
    Metrics
        The metrics dataclass
    """
    case_indices = metrics_dict.pop("case_indices", None)
    metrics = Metrics(
        fitness=fitness,
        extra_scores=metrics_dict,
        worst_k_contingency_cases=case_indices,
    )

    return metrics

run_single_iteration

run_single_iteration(_i, jax_data, update_fn)

Run a single iteration of the optimization.

This involves updating the genetic algorithm and calling the metrics callback

PARAMETER	DESCRIPTION
_i	The iteration number, will be ignored. Its only purpose is to make the function signature compatible with lax.fori_loop TYPE: Int[Array, '']
jax_data	The data stored for the optimization run from the last epoch or from initialize_optimization TYPE: JaxOptimizerData
update_fn	[GARepertoire, EmitterState, jax.random.PRNGKey, Any], tuple[GARepertoire, EmitterState, Any, jax.random.PRNGKey] )] The update function of the genetic algorithm TYPE: Callable[(

RETURNS	DESCRIPTION
JaxOptimizerData	The updated data for the optimization run

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/worker/optimizer.py

@partial(
    jax.jit,
    static_argnames=("update_fn"),
    donate_argnames=("jax_data",),
)
def run_single_iteration(
    _i: Int[Array, ""],
    jax_data: JaxOptimizerData,
    update_fn: Callable[
        [DiscreteMapElitesRepertoire, EmitterState, jax.random.PRNGKey, Any],
        tuple[DiscreteMapElitesRepertoire, EmitterState, Any, jax.random.PRNGKey],
    ],
) -> JaxOptimizerData:
    """Run a single iteration of the optimization.

    This involves updating the genetic algorithm and calling the metrics callback

    Parameters
    ----------
    _i : Int[Array, ""]
        The iteration number, will be ignored. Its only purpose is to make the function signature
        compatible with lax.fori_loop
    jax_data : JaxOptimizerData
        The data stored for the optimization run from the last epoch or from initialize_optimization
    update_fn : Callable[(
                        [GARepertoire, EmitterState, jax.random.PRNGKey, Any],
                        tuple[GARepertoire, EmitterState, Any, jax.random.PRNGKey]
                    )]
        The update function of the genetic algorithm

    Returns
    -------
    JaxOptimizerData
        The updated data for the optimization run
    """
    repertoire, emitter_state, _metrics, random_key = update_fn(
        jax_data.repertoire,
        jax_data.emitter_state,
        jax_data.random_key,
        jax_data.dynamic_informations,
    )

    jax_data = replace(
        jax_data,
        repertoire=repertoire,
        emitter_state=emitter_state,
        random_key=random_key,
        latest_iteration=jax_data.latest_iteration + 1,
    )

    return jax_data

run_single_device_epoch

run_single_device_epoch(
    jax_data, iterations_per_epoch, update_fn
)

Run one epoch of the optimization on a single device.

Basically this is just a fori loop over the iterations_per_epoch, calling run_single_iteration This can be used to pass to pmap

PARAMETER	DESCRIPTION
jax_data	The data stored for the optimization run from the last epoch or from initialize_optimization TYPE: JaxOptimizerData
iterations_per_epoch	The number of iterations to run in this epoch TYPE: int
update_fn	[GARepertoire, EmitterState, jax.random.PRNGKey, Any], tuple[GARepertoire, EmitterState, Any, jax.random.PRNGKey] )] The update function of the genetic algorithm TYPE: Callable[(

RETURNS	DESCRIPTION
JaxOptimizerData	The updated data for the optimization run

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/worker/optimizer.py

@partial(
    jax.jit,
    static_argnames=("iterations_per_epoch", "update_fn"),
    # donate_argnames=("jax_data",),
)
def run_single_device_epoch(
    jax_data: JaxOptimizerData,
    iterations_per_epoch: int,
    update_fn: Callable[
        [DiscreteMapElitesRepertoire, EmitterState, jax.random.PRNGKey, Any],
        tuple[DiscreteMapElitesRepertoire, EmitterState, Any, jax.random.PRNGKey],
    ],
) -> JaxOptimizerData:
    """Run one epoch of the optimization on a single device.

    Basically this is just a fori loop over the iterations_per_epoch, calling run_single_iteration
    This can be used to pass to pmap

    Parameters
    ----------
    jax_data : JaxOptimizerData
        The data stored for the optimization run from the last epoch or from initialize_optimization
    iterations_per_epoch : int
        The number of iterations to run in this epoch
    update_fn : Callable[(
                        [GARepertoire, EmitterState, jax.random.PRNGKey, Any],
                        tuple[GARepertoire, EmitterState, Any, jax.random.PRNGKey]
                    )]
        The update function of the genetic algorithm

    Returns
    -------
    JaxOptimizerData
        The updated data for the optimization run
    """
    return lax.fori_loop(
        0,
        iterations_per_epoch,
        partial(run_single_iteration, update_fn=update_fn),
        jax_data,
    )

run_epoch

run_epoch(optimizer_data)

Run one epoch of the optimization.

This function will be called repeatedly by the worker. It should include multiple iterations, according to the configuration of the optimizer. Furthermore it should call the metrics_callback during the epoch, at least at the beginning. This could happen through a jax callback to not block the main thread.

PARAMETER	DESCRIPTION
optimizer_data	The data stored for the optimization run from the last epoch or from initialize_optimization TYPE: OptimizerData

RETURNS	DESCRIPTION
OptimizerData	The updated data for the optimization run

RAISES	DESCRIPTION
Exception	If an error occurs during the epoch. It will be caught by the worker and sent back to the results topic

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/worker/optimizer.py

def run_epoch(
    optimizer_data: OptimizerData,
) -> OptimizerData:
    """Run one epoch of the optimization.

    This function will be called repeatedly by the worker. It should include multiple iterations,
    according to the configuration of the optimizer. Furthermore it should call the metrics_callback
    during the epoch, at least at the beginning. This could happen through a jax callback to not
    block the main thread.

    Parameters
    ----------
    optimizer_data : OptimizerData
        The data stored for the optimization run from the last epoch or from initialize_optimization

    Returns
    -------
    OptimizerData
        The updated data for the optimization run

    Raises
    ------
    Exception
        If an error occurs during the epoch. It will be caught by the worker and sent back to the
        results topic
    """
    epoch_fn = run_single_device_epoch
    if optimizer_data.start_params.loadflow_solver_config.distributed:
        epoch_fn = jax.pmap(
            epoch_fn,
            axis_name="p",
            static_broadcasted_argnums=(1, 2),
            donate_argnums=(0,),
        )

    jax_data = epoch_fn(
        optimizer_data.jax_data,
        optimizer_data.start_params.ga_config.iterations_per_epoch,
        optimizer_data.algo.update,
    )

    return replace(
        optimizer_data,
        jax_data=jax_data,
    )

extract_topologies

extract_topologies(optimizer_data)

Extract unique and not sent topologies from the repertoire.

Will return a list of topologies that have not been sent to the results topic yet. Also it will update the optimizer data sent set in place to include those topologies.

PARAMETER	DESCRIPTION
optimizer_data	The data stored for the optimization run, the sent_topologies set will be updated in place TYPE: OptimizerData

RETURNS	DESCRIPTION
list[Topology]	The topologies to send to the results topic, in message topo format

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/worker/optimizer.py

def extract_topologies(optimizer_data: OptimizerData) -> list[Topology]:
    """Extract unique and not sent topologies from the repertoire.

    Will return a list of topologies that have not been sent to the results topic yet. Also it will update the
    optimizer data sent set in place to include those topologies.

    Parameters
    ----------
    optimizer_data : OptimizerData
        The data stored for the optimization run, the sent_topologies set will be updated in place

    Returns
    -------
    list[Topology]
        The topologies to send to the results topic, in message topo format
    """
    # Assuming that contingency_ids stay the same for all timesteps
    contingency_ids = optimizer_data.solver_configs[0].contingency_ids

    # Get grid_model_low_tap if nodal injection information is available
    nodal_inj_info = optimizer_data.jax_data.dynamic_informations[0].nodal_injection_information
    grid_model_low_tap = nodal_inj_info.grid_model_low_tap if nodal_inj_info is not None else None

    topologies = summarize_repo(
        optimizer_data.jax_data.repertoire,
        initial_fitness=optimizer_data.initial_fitness,
        contingency_ids=contingency_ids,
        grid_model_low_tap=grid_model_low_tap,
    )

    # Filter out topologies that have already been sent
    new_topologies = []
    new_topology_hashes = []
    for topo in topologies:
        topology_hash = hash_strategy(Strategy(timesteps=[topo]))
        if topology_hash not in optimizer_data.sent_topologies:
            new_topologies.append(topo)
            new_topology_hashes.append(topology_hash)
    optimizer_data.sent_topologies.update(new_topology_hashes)

    return new_topologies

convert_topologies_to_messages

convert_topologies_to_messages(topologies, epoch)

Convert a list of topologies in message format to a list of TopologyPushResult.

PARAMETER	DESCRIPTION
topologies	The topologies in message format TYPE: list[Topology]
epoch	The epoch number, used to annotate the results TYPE: int

RETURNS	DESCRIPTION
list[TopologyPushResult]	The topologies in TopologyPushResult format, ready to be sent to the results topic

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/worker/optimizer.py

def convert_topologies_to_messages(topologies: list[Topology], epoch: int) -> list[TopologyPushResult]:
    """Convert a list of topologies in message format to a list of TopologyPushResult.

    Parameters
    ----------
    topologies : list[Topology]
        The topologies in message format
    epoch : int
        The epoch number, used to annotate the results

    Returns
    -------
    list[TopologyPushResult]
        The topologies in TopologyPushResult format, ready to be sent to the results topic
    """
    topo_messages = []
    for topo in topologies:
        topo_message = TopologyPushResult(
            strategy=Strategy(timesteps=[topo]),
            epoch=epoch,
        )
        topo_messages.append(topo_message)

    return topo_messages

toop_engine_topology_optimizer.dc.worker.worker

Kafka worker for the genetic algorithm optimization.

logger module-attribute

logger = get_logger(__name__)

Args

Bases: BaseModel

Launch arguments for the worker, which can not be changed during the optimization run.

kafka_broker class-attribute instance-attribute

kafka_broker = 'localhost:9092'

The Kafka broker to connect to.

optimizer_command_topic class-attribute instance-attribute

optimizer_command_topic = 'commands'

The Kafka topic to listen for commands on.

optimizer_results_topic class-attribute instance-attribute

optimizer_results_topic = 'results'

The topic to push results to.

optimizer_heartbeat_topic class-attribute instance-attribute

optimizer_heartbeat_topic = 'heartbeat'

The topic to push heartbeats to.

heartbeat_interval_ms class-attribute instance-attribute

heartbeat_interval_ms = 1000

The interval in milliseconds to send heartbeats.

max_command_age_hours class-attribute instance-attribute

max_command_age_hours = 3.0

The maximum age of a command in hours. If a command is received that is older than this, the command will be ignored.

idle_loop

idle_loop(
    consumer,
    send_heartbeat_fn,
    send_result_fn,
    heartbeat_interval_ms,
    max_command_age_hours,
)

Run idle loop of the worker.

This will be running when the worker is currently not optimizing This will wait until a StartOptimizationCommand is received and return it. In case a ShutdownCommand is received, the worker will exit with the exit code provided in the command.

PARAMETER	DESCRIPTION
consumer	The initialized Kafka consumer to listen for commands on. TYPE: LongRunningKafkaConsumer
send_heartbeat_fn	A function to call when there were no messages received for a while. TYPE: Callable[[HeartbeatUnion], None]
send_result_fn	A function to call to send results back to the results topic, used to send a message in case a command is too old. TYPE: Callable[[ResultUnion, str], None]
heartbeat_interval_ms	The time to wait for a new command in milliseconds. If no command has been received, a heartbeat will be sent and then the receiver will wait for commands again. TYPE: int
max_command_age_hours	The maximum age of a command in hours. If a command is received that is older than this, the command will be ignored and a message will be sent to the results topic. TYPE: float

RETURNS	DESCRIPTION
(StartOptimizationCommand,)	The start optimization command to start the optimization run with

RAISES	DESCRIPTION
SystemExit	If a ShutdownCommand is received

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/worker/worker.py

def idle_loop(
    consumer: LongRunningKafkaConsumer,
    send_heartbeat_fn: Callable[[HeartbeatUnion], None],
    send_result_fn: Callable[[ResultUnion, str], None],
    heartbeat_interval_ms: int,
    max_command_age_hours: float,
) -> StartOptimizationCommand:
    """Run idle loop of the worker.

    This will be running when the worker is currently not optimizing
    This will wait until a StartOptimizationCommand is received and return it. In case a
    ShutdownCommand is received, the worker will exit with the exit code provided in the command.

    Parameters
    ----------
    consumer : LongRunningKafkaConsumer
        The initialized Kafka consumer to listen for commands on.
    send_heartbeat_fn : Callable[[HeartbeatUnion], None]
        A function to call when there were no messages received for a while.
    send_result_fn : Callable[[ResultUnion, str], None]
        A function to call to send results back to the results topic, used to send a message in case a command is too old.
    heartbeat_interval_ms : int
        The time to wait for a new command in milliseconds. If no command has been received, a
        heartbeat will be sent and then the receiver will wait for commands again.
    max_command_age_hours: float
        The maximum age of a command in hours.
        If a command is received that is older than this, the command will be ignored
        and a message will be sent to the results topic.

    Returns
    -------
    StartOptimizationCommand,
        The start optimization command to start the optimization run with

    Raises
    ------
    SystemExit
        If a ShutdownCommand is received
    """
    send_heartbeat_fn(IdleHeartbeat())
    logger.info("Entering idle loop")
    while True:
        message = consumer.poll(timeout=heartbeat_interval_ms / 1000)

        # Wait timeout exceeded
        if not message:
            send_heartbeat_fn(IdleHeartbeat())
            continue

        command = Command.model_validate_json(deserialize_message(message.value()))

        if isinstance(command.command, ShutdownCommand):
            logger.info("Shutting down due to ShutdownCommand")
            consumer.commit()
            consumer.consumer.close()
            raise SystemExit(command.command.exit_code)
        if isinstance(command.command, StartOptimizationCommand):
            time_of_command = datetime.fromisoformat(command.timestamp)
            if time_of_command < datetime.now() - timedelta(hours=max_command_age_hours):
                logger.warning(
                    f"Received command with timestamp from the past (timestamp: {time_of_command}, "
                    f"now: {datetime.now()}), skipping command"
                )
                send_result_fn(
                    OptimizationStoppedResult(
                        reason="command-too-old", message=f"Received outdated command: {command}. Skipping.."
                    ),
                    command.command.optimization_id,
                )
                consumer.commit()
                continue
            with structlog.contextvars.bound_contextvars(
                optimization_id=command.command.optimization_id,
            ):
                return command.command

        # If we are here, we received a command that we do not know
        logger.warning(f"Received unknown command, dropping: {command} / {message.value}")
        consumer.commit()

push_topologies

push_topologies(optimizer_data, epoch, send_result_fn)

Push topologies to the results topic.

PARAMETER	DESCRIPTION
optimizer_data	The data of the optimizer, containing the repertoire with topologies to push. TYPE: OptimizerData
epoch	The current epoch, used for logging and to include in the messages. TYPE: int
send_result_fn	A function to call to send results back to the results topic. TYPE: Callable[[ResultUnion], None]

RETURNS	DESCRIPTION
int	The number of topologies pushed

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/worker/worker.py

def push_topologies(optimizer_data: OptimizerData, epoch: int, send_result_fn: Callable[[ResultUnion], None]) -> int:
    """Push topologies to the results topic.

    Parameters
    ----------
    optimizer_data : OptimizerData
        The data of the optimizer, containing the repertoire with topologies to push.
    epoch : int
        The current epoch, used for logging and to include in the messages.
    send_result_fn : Callable[[ResultUnion], None]
        A function to call to send results back to the results topic.

    Returns
    -------
    int
        The number of topologies pushed
    """
    with jax.default_device(jax.devices("cpu")[0]):
        push_results = convert_topologies_to_messages(extract_topologies(optimizer_data), epoch)
        for push_result in push_results:
            send_result_fn(push_result)
        return len(push_results)

optimization_loop

optimization_loop(
    dc_params,
    grid_files,
    send_result_fn,
    flush_result_fn,
    send_heartbeat_fn,
    optimization_id,
    processed_gridfile_fs,
)

Run an optimization until the optimization has converged

PARAMETER	DESCRIPTION
dc_params	The parameters for the optimization run, usually from the start command TYPE: DCOptimizerParameters
grid_files	The grid files to load, where each gridfile represents one timestep. TYPE: list[GridFile]
send_result_fn	A function to queue results for the results topic. This callback is not expected to flush and will be called multiple times in every epoch, for every topology discovered. After all topologies have been sent, the flush_result_fn will be called to flush the results to Kafka. TYPE: Callable[[ResultUnion], None]
flush_result_fn	A function to flush queued results to Kafka after one or more calls to send_result_fn. TYPE: Callable[[], None]
send_heartbeat_fn	A function to call after every epoch to signal that the worker is still alive. TYPE: Callable[[HeartbeatUnion], None]
optimization_id	The id of the optimization run. This will be used to identify the optimization run in the results. Should stay the same for the whole optimization run and should be equal to the kafka event key. TYPE: str
processed_gridfile_fs	The target filesystem for the preprocessing worker. This contains all processed grid files. During the import job, a new folder import_results.data_folder was created which will be completed with the preprocess call to this function. Internally, only the data folder is passed around as a dirfs. Note that the unprocessed_gridfile_fs is not needed here anymore, as all preprocessing steps that need the unprocessed gridfiles were already done. TYPE: AbstractFileSystem

RAISES	DESCRIPTION
SystemExit	If a ShutdownCommand is received

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/worker/worker.py

def optimization_loop(
    dc_params: DCOptimizerParameters,
    grid_files: list[GridFile],
    send_result_fn: Callable[[ResultUnion], None],
    flush_result_fn: Callable[[], None],
    send_heartbeat_fn: Callable[[HeartbeatUnion], None],
    optimization_id: str,
    processed_gridfile_fs: AbstractFileSystem,
) -> None:
    """Run an optimization until the optimization has converged

    Parameters
    ----------
    dc_params : DCOptimizerParameters
        The parameters for the optimization run, usually from the start command
    grid_files : list[GridFile]
        The grid files to load, where each gridfile represents one timestep.
    send_result_fn : Callable[[ResultUnion], None]
        A function to queue results for the results topic. This callback is not expected to flush and will be called
        multiple times in every epoch, for every topology discovered. After all topologies have been sent, the
        flush_result_fn will be called to flush the results to Kafka.
    flush_result_fn : Callable[[], None]
        A function to flush queued results to Kafka after one or more calls to ``send_result_fn``.
    send_heartbeat_fn : Callable[[HeartbeatUnion], None]
        A function to call after every epoch to signal that the worker is still alive.
    optimization_id : str
        The id of the optimization run. This will be used to identify the optimization run in the
        results. Should stay the same for the whole optimization run and should be equal to the kafka
        event key.
    processed_gridfile_fs: AbstractFileSystem
        The target filesystem for the preprocessing worker. This contains all processed grid files.
        During the import job,  a new folder import_results.data_folder was created
        which will be completed with the preprocess call to this function.
        Internally, only the data folder is passed around as a dirfs.
        Note that the unprocessed_gridfile_fs is not needed here anymore, as all preprocessing steps that need the
        unprocessed gridfiles were already done.

    Raises
    ------
    SystemExit
        If a ShutdownCommand is received
    """
    logger.info(f"Initializing DC optimization {optimization_id}")

    try:
        send_heartbeat_fn(
            OptimizationStartedHeartbeat(
                optimization_id=optimization_id,
            )
        )
        optimizer_data, stats, initial_strategy = initialize_optimization(
            params=dc_params,
            optimization_id=optimization_id,
            static_information_files=tuple([gf.static_information_file for gf in grid_files]),
            processed_gridfile_fs=processed_gridfile_fs,
        )
        send_result_fn(
            OptimizationStartedResult(
                initial_topology=initial_strategy,
                initial_stats=stats,
            )
        )
        flush_result_fn()

    except Exception as e:
        send_result_fn(OptimizationStoppedResult(reason="error", message=str(e)))
        flush_result_fn()
        logger.error(f"Error during initialization of optimization {optimization_id}: {e}")
        return

    logger.info(f"Starting optimization {optimization_id}")
    epoch = 1
    running = True
    start_time = time.time()
    while running:
        try:
            optimizer_data = run_epoch(optimizer_data)
            n_pushes = push_topologies(optimizer_data, epoch, send_result_fn)
            if n_pushes > 0:
                flush_result_fn()
            logger.info(
                f"Sent {n_pushes} strategies to results topic,"
                f" best repofitness: {optimizer_data.jax_data.repertoire.fitnesses.max().item()}, epoch: {epoch}"
            )
        except Exception as e:
            # Send a stop message to the results
            send_result_fn(OptimizationStoppedResult(reason="error", message=str(e)))
            flush_result_fn()

            logger.error(f"Error during optimization {optimization_id}, epoch {epoch}: {e}")
            return
        epoch += 1

        send_heartbeat_fn(
            OptimizationStatsHeartbeat(
                optimization_id=optimization_id,
                wall_time=time.time() - start_time,
                iteration=epoch,
                num_branch_topologies_tried=optimizer_data.jax_data.emitter_state.total_branch_combis.sum().item(),
                num_injection_topologies_tried=optimizer_data.jax_data.emitter_state.total_inj_combis.sum().item(),
            )
        )

        if time.time() - start_time > dc_params.ga_config.runtime_seconds:
            logger.info(f"Stopping optimization {optimization_id} at epoch {epoch} due to runtime limit")
            send_result_fn(OptimizationStoppedResult(epoch=epoch, reason="converged", message="runtime limit"))
            flush_result_fn()
            running = False
            break

main

main(
    args, processed_gridfile_fs, producer, command_consumer
)

Run the main DC worker loop.

PARAMETER	DESCRIPTION
args	The command line arguments TYPE: Args
processed_gridfile_fs	The target filesystem for the preprocessing worker. This contains all processed grid files. During the import job, a new folder import_results.data_folder was created which will be completed with the preprocess call to this function. Internally, only the data folder is passed around as a dirfs. Note that the unprocessed_gridfile_fs is not needed here anymore, as all preprocessing steps that need the unprocessed gridfiles were already done. TYPE: AbstractFileSystem
producer	The Kafka producer to send results and heartbeats with. TYPE: Producer
command_consumer	The Kafka consumer to receive commands with. TYPE: LongRunningKafkaConsumer

RAISES	DESCRIPTION
SystemExit	If the worker receives a ShutdownCommand

Source code in packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/worker/worker.py

def main(
    args: Args,
    processed_gridfile_fs: AbstractFileSystem,
    producer: Producer,
    command_consumer: LongRunningKafkaConsumer,
) -> None:
    """Run the main DC worker loop.

    Parameters
    ----------
    args : Args
        The command line arguments
    processed_gridfile_fs: AbstractFileSystem
        The target filesystem for the preprocessing worker. This contains all processed grid files.
        During the import job,  a new folder import_results.data_folder was created
        which will be completed with the preprocess call to this function.
        Internally, only the data folder is passed around as a dirfs.
        Note that the unprocessed_gridfile_fs is not needed here anymore, as all preprocessing steps that need the
        unprocessed gridfiles were already done.
    producer : Producer
        The Kafka producer to send results and heartbeats with.
    command_consumer : LongRunningKafkaConsumer
        The Kafka consumer to receive commands with.

    Raises
    ------
    SystemExit
        If the worker receives a ShutdownCommand
    """
    instance_id = str(uuid4())
    logger.info(f"Starting DC worker {instance_id} with config {args}")
    jax.config.update("jax_enable_x64", True)
    jax.config.update("jax_logging_level", "INFO")

    def send_heartbeat(message: HeartbeatUnion, ping_consumer: bool) -> None:
        logger.debug(f"Sending heartbeat: {message}", message_type=type(message).__name__)
        heartbeat = Heartbeat(
            optimizer_type=OptimizerType.DC,
            instance_id=instance_id,
            message=message,
        )
        producer.produce(
            args.optimizer_heartbeat_topic,
            value=serialize_message(heartbeat.model_dump_json()),
            key=heartbeat.instance_id.encode(),
        )
        producer.flush()
        if ping_consumer:
            command_consumer.heartbeat()

    def send_result(message: ResultUnion, optimization_id: str) -> None:
        logger.info(
            f"Sending result for optimization {optimization_id}: {message}",
            optimization_id=optimization_id,
            result_type=type(message).__name__,
        )
        result = Result(
            result=message,
            optimization_id=optimization_id,
            optimizer_type=OptimizerType.DC,
            instance_id=instance_id,
        )
        producer.produce(
            args.optimizer_results_topic,
            value=serialize_message(result.model_dump_json()),
            key=optimization_id.encode(),
        )

    def send_result_and_flush(message: ResultUnion, optimization_id: str) -> None:
        send_result(message=message, optimization_id=optimization_id)
        producer.flush()

    while True:
        command = idle_loop(
            consumer=command_consumer,
            send_heartbeat_fn=partial(send_heartbeat, ping_consumer=False),
            send_result_fn=send_result_and_flush,
            heartbeat_interval_ms=args.heartbeat_interval_ms,
            max_command_age_hours=args.max_command_age_hours,
        )

        command_consumer.start_processing()
        optimization_loop(
            dc_params=command.dc_params,
            grid_files=command.grid_files,
            send_result_fn=partial(send_result, optimization_id=command.optimization_id),
            flush_result_fn=producer.flush,
            send_heartbeat_fn=partial(send_heartbeat, ping_consumer=True),
            optimization_id=command.optimization_id,
            processed_gridfile_fs=processed_gridfile_fs,
        )
        command_consumer.stop_processing()