From e6435429755193273e51de839ba863ab4de61054 Mon Sep 17 00:00:00 2001 From: Nico Westerbeck Date: Sun, 5 Jul 2026 10:16:08 +0000 Subject: [PATCH 1/7] feat: Simple bruteforce worker for full action set evaluation Signed-off-by: Nico Westerbeck --- .../dc/worker/idle_loop.py | 99 +++ .../dc/worker/worker.py | 91 +-- .../dc_bruteforce/__init__.py | 1 + .../dc_bruteforce/generator.py | 137 ++++ .../dc_bruteforce/optimizer.py | 678 ++++++++++++++++++ .../dc_bruteforce/worker.py | 262 +++++++ .../benchmark/test_dc_bruteforce_generator.py | 38 + .../dc/worker/test_dc_optimizer_worker.py | 108 +-- .../tests/dc/worker/test_idle_loop.py | 126 ++++ .../tests/dc_bruteforce/test_generator.py | 120 ++++ .../tests/dc_bruteforce/test_optimizer.py | 109 +++ .../tests/dc_bruteforce/test_worker.py | 66 ++ 12 files changed, 1640 insertions(+), 195 deletions(-) create mode 100644 packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/worker/idle_loop.py create mode 100644 packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc_bruteforce/__init__.py create mode 100644 packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc_bruteforce/generator.py create mode 100644 packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc_bruteforce/optimizer.py create mode 100644 packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc_bruteforce/worker.py create mode 100644 packages/topology_optimizer_pkg/tests/benchmark/test_dc_bruteforce_generator.py create mode 100644 packages/topology_optimizer_pkg/tests/dc/worker/test_idle_loop.py create mode 100644 packages/topology_optimizer_pkg/tests/dc_bruteforce/test_generator.py create mode 100644 packages/topology_optimizer_pkg/tests/dc_bruteforce/test_optimizer.py create mode 100644 packages/topology_optimizer_pkg/tests/dc_bruteforce/test_worker.py diff --git a/packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/worker/idle_loop.py b/packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/worker/idle_loop.py new file mode 100644 index 000000000..cf1e15800 --- /dev/null +++ b/packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/worker/idle_loop.py @@ -0,0 +1,99 @@ +# Copyright 2026 50Hertz Transmission GmbH and Elia Transmission Belgium SA/NV +# +# This Source Code Form is subject to the terms of the Mozilla Public License, v. 2.0. +# If a copy of the MPL was not distributed with this file, +# you can obtain one at https://mozilla.org/MPL/2.0/. +# Mozilla Public License, version 2.0 + +"""Shared idle loop for DC worker processes.""" + +from datetime import datetime, timedelta + +import structlog +from beartype.typing import Callable +from toop_engine_contingency_analysis.ac_loadflow_service.kafka_client import LongRunningKafkaConsumer +from toop_engine_interfaces.messages.protobuf_message_factory import deserialize_message +from toop_engine_topology_optimizer.interfaces.messages.commands import Command, ShutdownCommand, StartOptimizationCommand +from toop_engine_topology_optimizer.interfaces.messages.heartbeats import HeartbeatUnion, IdleHeartbeat +from toop_engine_topology_optimizer.interfaces.messages.results import OptimizationStoppedResult, ResultUnion + +logger = structlog.get_logger(__name__) + + +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) + + 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 + + logger.warning(f"Received unknown command, dropping: {command} / {message.value}") + consumer.commit() diff --git a/packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/worker/worker.py b/packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/worker/worker.py index 7cf71af5d..5ca4e74a9 100644 --- a/packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/worker/worker.py +++ b/packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/worker/worker.py @@ -8,7 +8,6 @@ """Kafka worker for the genetic algorithm optimization.""" import time -from datetime import datetime, timedelta from functools import partial from uuid import uuid4 @@ -19,7 +18,8 @@ from fsspec import AbstractFileSystem from pydantic import BaseModel from toop_engine_contingency_analysis.ac_loadflow_service.kafka_client import LongRunningKafkaConsumer -from toop_engine_interfaces.messages.protobuf_message_factory import deserialize_message, serialize_message +from toop_engine_interfaces.messages.protobuf_message_factory import serialize_message +from toop_engine_topology_optimizer.dc.worker.idle_loop import idle_loop from toop_engine_topology_optimizer.dc.worker.optimizer import ( OptimizerData, convert_topologies_to_messages, @@ -27,17 +27,11 @@ initialize_optimization, run_epoch, ) -from toop_engine_topology_optimizer.interfaces.messages.commands import ( - Command, - ShutdownCommand, - StartOptimizationCommand, -) from toop_engine_topology_optimizer.interfaces.messages.commons import GridFile, OptimizerType from toop_engine_topology_optimizer.interfaces.messages.dc_params import DCOptimizerParameters from toop_engine_topology_optimizer.interfaces.messages.heartbeats import ( Heartbeat, HeartbeatUnion, - IdleHeartbeat, OptimizationStartedHeartbeat, OptimizationStatsHeartbeat, ) @@ -74,87 +68,6 @@ class Args(BaseModel): the command will be ignored.""" -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() - - def push_topologies(optimizer_data: OptimizerData, epoch: int, send_result_fn: Callable[[ResultUnion], None]) -> int: """Push topologies to the results topic. diff --git a/packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc_bruteforce/__init__.py b/packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc_bruteforce/__init__.py new file mode 100644 index 000000000..e15459408 --- /dev/null +++ b/packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc_bruteforce/__init__.py @@ -0,0 +1 @@ +"""Bruteforce DC optimizer components.""" diff --git a/packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc_bruteforce/generator.py b/packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc_bruteforce/generator.py new file mode 100644 index 000000000..0a2bde704 --- /dev/null +++ b/packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc_bruteforce/generator.py @@ -0,0 +1,137 @@ +# Copyright 2026 50Hertz Transmission GmbH and Elia Transmission Belgium SA/NV +# +# This Source Code Form is subject to the terms of the Mozilla Public License, v. 2.0. +# If a copy of the MPL was not distributed with this file, +# you can obtain one at https://mozilla.org/MPL/2.0/. +# Mozilla Public License, version 2.0 + +"""Lazy workset generation for the DC bruteforce optimizer.""" + +from dataclasses import dataclass +from itertools import combinations, islice, product +from math import comb + +from beartype.typing import Iterator, Sequence + + +@dataclass(frozen=True) +class WorksetEntry: + """One exhaustively enumerated topology candidate. + + Empty action/disconnection slots are not represented as the empty action but just by a shorter tuple as that is easier + to enumerate on using itertools.combinations. + """ + + action_indices: tuple[int, ...] + """Split-action indices into the action set. Each index belongs to a different substation. + """ + + disconnections: tuple[int, ...] + """Disconnection indices into ``disconnectable_branches``.""" + + +def count_workset_size( + split_action_groups: Sequence[Sequence[int]], + max_num_splits: int, + n_disconnectable_branches: int, + max_num_disconnections: int, +) -> int: + """Count the total number of bruteforce candidates. + + Parameters + ---------- + split_action_groups : Sequence[Sequence[int]] + Split-action indices grouped by substation. + max_num_splits : int + Maximum number of simultaneously split substations. + n_disconnectable_branches : int + Number of disconnectable branches. + max_num_disconnections : int + Maximum number of simultaneous disconnections. + + Returns + ------- + int + The exact total number of topology candidates the workset iterator will emit. + """ + limited_num_splits = min(max_num_splits, len(split_action_groups)) + limited_num_disconnections = min(max_num_disconnections, n_disconnectable_branches) + disconnection_factor = sum(comb(n_disconnectable_branches, n_disc) for n_disc in range(limited_num_disconnections + 1)) + + total_split_combinations = 0 + for n_splits in range(limited_num_splits + 1): + for chosen_groups in combinations(split_action_groups, n_splits): + split_combination_count = 1 + for action_group in chosen_groups: + split_combination_count *= len(action_group) + total_split_combinations += split_combination_count + + return total_split_combinations * disconnection_factor + + +def iter_workset( + action_start_indices: Sequence[int], + n_actions_per_sub: Sequence[int], + max_num_splits: int, + n_disconnectable_branches: int, + max_num_disconnections: int, +) -> Iterator[WorksetEntry]: + """Yield topology candidates lazily from contiguous per-substation action blocks. + + Parameters + ---------- + action_start_indices : Sequence[int] + Start index of each substation's contiguous action block in the flattened action set. + n_actions_per_sub : Sequence[int] + Number of actions available for each substation. The bruteforce path assumes every relevant + substation has at least one action. + max_num_splits : int + Maximum number of simultaneously split substations. + n_disconnectable_branches : int + Number of disconnectable branches. + max_num_disconnections : int + Maximum number of simultaneous disconnections. + + Yields + ------ + WorksetEntry + The next lazily generated topology candidate. + """ + split_action_groups = tuple( + tuple(range(int(start), int(start) + int(count))) + for start, count in zip(action_start_indices, n_actions_per_sub, strict=True) + ) + limited_num_splits = min(max_num_splits, len(split_action_groups)) + limited_num_disconnections = min(max_num_disconnections, n_disconnectable_branches) + disconnectable_indices = tuple(range(n_disconnectable_branches)) + + for n_splits in range(limited_num_splits + 1): + for chosen_groups in combinations(split_action_groups, n_splits): + for chosen_actions in product(*chosen_groups) if chosen_groups else [()]: + action_indices = tuple(int(action_index) for action_index in chosen_actions) + for n_disconnections in range(limited_num_disconnections + 1): + for chosen_disconnections in combinations(disconnectable_indices, n_disconnections): + yield WorksetEntry( + action_indices=action_indices, + disconnections=tuple(int(disconnection) for disconnection in chosen_disconnections), + ) + + +def take_workset_chunk(workset: Iterator[WorksetEntry], chunk_size: int) -> list[WorksetEntry]: + """Take the next chunk from a lazy workset iterator. + + Parameters + ---------- + workset : Iterator[WorksetEntry] + The lazy workset iterator. + chunk_size : int + Maximum number of entries to retrieve. + + Returns + ------- + list[WorksetEntry] + The next chunk from the iterator. + """ + if chunk_size <= 0: + raise ValueError(f"chunk_size must be positive, got {chunk_size}") + return list(islice(workset, chunk_size)) diff --git a/packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc_bruteforce/optimizer.py b/packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc_bruteforce/optimizer.py new file mode 100644 index 000000000..1f15788d8 --- /dev/null +++ b/packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc_bruteforce/optimizer.py @@ -0,0 +1,678 @@ +# Copyright 2026 50Hertz Transmission GmbH and Elia Transmission Belgium SA/NV +# +# This Source Code Form is subject to the terms of the Mozilla Public License, v. 2.0. +# If a copy of the MPL was not distributed with this file, +# you can obtain one at https://mozilla.org/MPL/2.0/. +# Mozilla Public License, version 2.0 + +"""Standalone exhaustive DC bruteforce optimizer runtime.""" + +import time +from dataclasses import dataclass, field +from functools import partial +from pathlib import Path + +import jax +import jax.numpy as jnp +import numpy as np +from beartype.typing import Any, Callable, Iterator, Sequence +from fsspec import AbstractFileSystem +from jax_dataclasses import replace +from jaxtyping import Array, ArrayLike, Float, Int, PRNGKeyArray +from qdax.core.emitters.standard_emitters import ExtraScores +from toop_engine_dc_solver.jax.inputs import load_static_information_fs +from toop_engine_dc_solver.jax.types import DynamicInformation, SolverConfig, StaticInformation, int_max +from toop_engine_dc_solver.preprocess.convert_to_jax import StaticInformationStats, extract_static_information_stats +from toop_engine_interfaces.types import MetricType +from toop_engine_topology_optimizer.dc.genetic_functions.genotype import Genotype +from toop_engine_topology_optimizer.dc.genetic_functions.initialization import ( + update_max_mw_flows_according_to_double_limits, + update_static_information, + verify_static_information, +) +from toop_engine_topology_optimizer.dc.genetic_functions.scoring_functions import scoring_function +from toop_engine_topology_optimizer.dc_bruteforce.generator import ( + WorksetEntry, + count_workset_size, + iter_workset, + take_workset_chunk, +) +from toop_engine_topology_optimizer.interfaces.messages.dc_params import DCOptimizerParameters +from toop_engine_topology_optimizer.interfaces.messages.results import Metrics, Strategy, Topology, TopologyPushResult +from toop_engine_topology_optimizer.interfaces.models.base_storage import hash_strategy + +ScoreChunkFn = Callable[ + [Genotype, PRNGKeyArray], + tuple[ + Float[Array, " batch_size"], + Int[Array, " batch_size n_dims"], + dict[str, Any], + ExtraScores, + PRNGKeyArray, + Genotype, + ], +] + + +@dataclass +class BruteforceRuntimeState: + """State kept across bruteforce epochs.""" + + dynamic_informations: tuple[DynamicInformation, ...] + """Dynamic solver inputs for each timestep.""" + + solver_configs: tuple[SolverConfig, ...] + """Static solver configuration for each timestep.""" + + score_chunk_fn: ScoreChunkFn + """Jitted scoring function that evaluates a chunk of bruteforce candidates.""" + + workset: Iterator[WorksetEntry] + """Lazy iterator over the remaining bruteforce candidates.""" + + total_workset_size: int + """Exact number of candidates represented by the bruteforce workset.""" + + total_branch_topologies_tried: int = 0 + """Number of bruteforce candidates scored so far.""" + + exhausted: bool = False + """Whether the lazy workset has been fully consumed.""" + + best_fitness: float = float("-inf") + """Best fitness observed so far across all evaluated candidates.""" + + +@dataclass +class OptimizerData: + """Outer state for the standalone bruteforce worker.""" + + start_params: DCOptimizerParameters + """Optimization parameters used to initialize the bruteforce run.""" + + optimization_id: str + """Identifier of the optimization run, propagated to emitted messages.""" + + initial_fitness: float + """Fitness of the unsplit baseline topology.""" + + initial_metrics: dict[MetricType, float] + """Observed metrics for the unsplit baseline topology.""" + + runtime_state: BruteforceRuntimeState + """Mutable bruteforce state updated after every epoch.""" + + start_time: float + """Wall-clock timestamp when the bruteforce run started.""" + + sent_topologies: set[bytes] + """Hashes of topologies already emitted to avoid duplicate result messages.""" + + latest_topologies: list[Topology] = field(default_factory=list) + """Improved topologies produced by the most recent epoch.""" + + +def initialize_optimization( + params: DCOptimizerParameters, + optimization_id: str, + static_information_files: Sequence[str | Path], + processed_gridfile_fs: AbstractFileSystem, +) -> tuple[OptimizerData, list[StaticInformationStats], Strategy]: + """Initialize the bruteforce optimization run. + + Parameters + ---------- + params : DCOptimizerParameters + Parameters controlling the DC bruteforce optimization. + optimization_id : str + Identifier of the optimization run, used in emitted messages. + static_information_files : Sequence[str | Path] + Paths to the preprocessed static-information files for all timesteps. + processed_gridfile_fs : AbstractFileSystem + Filesystem containing the preprocessed grid data. + + Returns + ------- + tuple[OptimizerData, list[StaticInformationStats], Strategy] + The initialized optimizer state, static-information descriptions, and initial unsplit + strategy. + """ + topologies_per_epoch = _get_topologies_per_epoch(params) + + static_informations = _load_and_prepare_static_informations( + params=params, + static_information_files=static_information_files, + processed_gridfile_fs=processed_gridfile_fs, + batch_size=topologies_per_epoch, + ) + _validate_bruteforce_inputs(params, topologies_per_epoch, static_informations[0].dynamic_information) + runtime_state, initial_fitness, initial_metrics, initial_case_ids = _initialize_runtime_state( + params=params, + static_informations=static_informations, + topologies_per_epoch=topologies_per_epoch, + ) + initial_strategy = build_initial_strategy( + n_timesteps=len(runtime_state.dynamic_informations), + fitness=initial_fitness, + initial_metrics=initial_metrics, + initial_case_ids=initial_case_ids, + ) + + 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 + ] + + return ( + OptimizerData( + start_params=params, + optimization_id=optimization_id, + initial_fitness=initial_fitness, + initial_metrics=initial_metrics, + runtime_state=runtime_state, + start_time=time.time(), + sent_topologies=set(), + ), + static_information_descriptions, + initial_strategy, + ) + + +def build_initial_strategy( + n_timesteps: int, + fitness: float, + initial_metrics: dict[MetricType, float], + initial_case_ids: list[str], +) -> Strategy: + """Build the initial unsplit strategy message. + + Parameters + ---------- + n_timesteps : int + Number of timesteps for which the initial strategy must be emitted. + fitness : float + Fitness of the unsplit baseline topology. + initial_metrics : dict[MetricType, float] + Observed metrics of the unsplit baseline topology. + initial_case_ids : list[str] + Worst contingency case identifiers for the baseline topology. + + Returns + ------- + Strategy + The initial strategy containing the unsplit topology for every timestep. + """ + metrics = Metrics( + fitness=fitness, + extra_scores=dict(initial_metrics), + worst_k_contingency_cases=initial_case_ids, + ) + return Strategy( + timesteps=[ + Topology( + actions=[], + disconnections=[], + pst_setpoints=None, + metrics=metrics, + ) + for _ in range(n_timesteps) + ] + ) + + +def run_epoch(optimizer_data: OptimizerData) -> OptimizerData: + """Run one bruteforce epoch. + + Parameters + ---------- + optimizer_data : OptimizerData + Current state of the bruteforce optimization. + + Returns + ------- + OptimizerData + Updated optimizer state after evaluating one chunk of the workset. + """ + runtime_state = optimizer_data.runtime_state + topologies_per_epoch = _get_topologies_per_epoch(optimizer_data.start_params) + if runtime_state.exhausted: + return replace(optimizer_data, latest_topologies=[]) + + chunk = take_workset_chunk(runtime_state.workset, topologies_per_epoch) + if not chunk: + return replace( + optimizer_data, + runtime_state=replace(runtime_state, exhausted=True), + latest_topologies=[], + ) + + max_num_splits = optimizer_data.start_params.loadflow_solver_config.max_num_splits + max_num_disconnections = optimizer_data.start_params.loadflow_solver_config.max_num_disconnections + + genotype = _chunk_to_genotype( + chunk=chunk, + batch_size=topologies_per_epoch, + max_num_splits=max_num_splits, + max_num_disconnections=max_num_disconnections, + ) + fitness, _descriptors, metrics, _emitter_info, _unused_random_key, scored_genotypes = runtime_state.score_chunk_fn( + genotype, jax.random.PRNGKey(0) + ) + + evaluated_count = len(chunk) + evaluated_fitness = np.asarray(fitness[:evaluated_count]) + improved_indices = np.flatnonzero(np.isfinite(evaluated_fitness) & (evaluated_fitness > optimizer_data.initial_fitness)) + topologies = _convert_improved_topologies( + genotypes=scored_genotypes, + fitness=evaluated_fitness, + metrics=metrics, + observed_metrics=optimizer_data.start_params.ga_config.observed_metrics, + contingency_ids=list(runtime_state.solver_configs[0].contingency_ids), + survivor_indices=improved_indices, + ) + + finite_fitness = evaluated_fitness[np.isfinite(evaluated_fitness)] + best_fitness = runtime_state.best_fitness + if finite_fitness.size > 0: + best_fitness = max(best_fitness, float(finite_fitness.max())) + + return replace( + optimizer_data, + runtime_state=replace( + runtime_state, + total_branch_topologies_tried=runtime_state.total_branch_topologies_tried + evaluated_count, + exhausted=evaluated_count < topologies_per_epoch, + best_fitness=best_fitness, + ), + latest_topologies=topologies, + ) + + +def _get_topologies_per_epoch(params: DCOptimizerParameters) -> int: + """Compute how many topologies the bruteforce optimizer evaluates per epoch. + + Parameters + ---------- + params : DCOptimizerParameters + Optimization parameters controlling the bruteforce runtime. + + Returns + ------- + int + Number of topologies evaluated in a single bruteforce epoch. + """ + return params.ga_config.iterations_per_epoch * params.loadflow_solver_config.batch_size + + +def extract_topologies(optimizer_data: OptimizerData) -> list[Topology]: + """Return deduplicated topologies discovered in the latest epoch. + + Parameters + ---------- + optimizer_data : OptimizerData + Current state of the bruteforce optimization. + + Returns + ------- + list[Topology] + Newly discovered topologies that have not yet been emitted. + """ + new_topologies = [] + new_hashes = [] + for topology in optimizer_data.latest_topologies: + topology_hash = hash_strategy(Strategy(timesteps=[topology])) + if topology_hash not in optimizer_data.sent_topologies: + new_topologies.append(topology) + new_hashes.append(topology_hash) + optimizer_data.sent_topologies.update(new_hashes) + return new_topologies + + +def convert_topologies_to_messages(topologies: list[Topology], epoch: int) -> list[TopologyPushResult]: + """Convert topologies to result messages. + + Parameters + ---------- + topologies : list[Topology] + Topologies to emit. + epoch : int + Epoch in which the topologies were discovered. + + Returns + ------- + list[TopologyPushResult] + Result messages ready to be sent to Kafka. + """ + return [TopologyPushResult(strategy=Strategy(timesteps=[topology]), epoch=epoch) for topology in topologies] + + +def get_num_branch_topologies_tried(optimizer_data: OptimizerData) -> int: + """Return the number of branch topologies evaluated so far. + + Parameters + ---------- + optimizer_data : OptimizerData + Current state of the bruteforce optimization. + + Returns + ------- + int + Number of scored bruteforce candidates. + """ + return optimizer_data.runtime_state.total_branch_topologies_tried + + +def is_exhausted(optimizer_data: OptimizerData) -> bool: + """Check whether the workset has been fully consumed. + + Parameters + ---------- + optimizer_data : OptimizerData + Current state of the bruteforce optimization. + + Returns + ------- + bool + ``True`` if no more candidates remain to be evaluated. + """ + return optimizer_data.runtime_state.exhausted + + +def _validate_bruteforce_inputs( + params: DCOptimizerParameters, topologies_per_epoch: int, dynamic_information: DynamicInformation +) -> None: + """Validate standalone bruteforce constraints. + + Parameters + ---------- + params : DCOptimizerParameters + Optimization parameters to validate. + topologies_per_epoch : int + Number of bruteforce candidates to evaluate per epoch. + dynamic_information : DynamicInformation + Dynamic solver inputs for the first timestep, used to validate available actions. + + Raises + ------ + ValueError + If the chunk size is invalid or if unsupported optimizer features are enabled. + """ + if topologies_per_epoch <= 0: + raise ValueError(f"topologies_per_epoch must be positive, got {topologies_per_epoch}") + if params.ga_config.enable_nodal_inj_optim: + raise ValueError("Bruteforce optimizer keeps PSTs untouched and does not support nodal injection optimization.") + if params.ga_config.enable_parallel_pst_group_optim: + raise ValueError("Bruteforce optimizer does not support parallel PST group optimization.") + if params.loadflow_solver_config.distributed: + raise ValueError("Bruteforce optimizer does not support distributed execution.") + + action_set = dynamic_information.action_set + assert action_set is not None, "Bruteforce optimization requires an action set." + assert params.loadflow_solver_config.max_num_splits <= len(action_set.action_start_indices), ( + "Bruteforce optimizer requires max_num_splits to be smaller than or equal to the number of substations " + "with available actions." + ) + assert params.loadflow_solver_config.max_num_disconnections <= dynamic_information.n_disconnectable_branches, ( + "Bruteforce optimizer requires max_num_disconnections to be smaller than or equal to the number of " + "disconnectable branches." + ) + + +def _load_and_prepare_static_informations( + params: DCOptimizerParameters, + static_information_files: Sequence[str | Path], + processed_gridfile_fs: AbstractFileSystem, + batch_size: int, +) -> tuple[StaticInformation, ...]: + """Load and preprocess static information for bruteforce execution. + + Parameters + ---------- + params : DCOptimizerParameters + Optimization parameters controlling preprocessing options. + static_information_files : Sequence[str | Path] + Paths to the static-information files for all timesteps. + processed_gridfile_fs : AbstractFileSystem + Filesystem containing the preprocessed grid data. + batch_size : int + Batch size to configure in the updated static information. + + Returns + ------- + tuple[StaticInformation, ...] + Prepared static-information objects for all timesteps. + """ + static_informations = tuple( + load_static_information_fs(filesystem=processed_gridfile_fs, filename=str(filename)) + for filename in static_information_files + ) + verify_static_information( + static_informations, + params.loadflow_solver_config.max_num_disconnections, + enable_nodal_inj_optim=False, + enable_parallel_pst_group_optim=False, + ) + static_informations = update_static_information( + static_informations, + batch_size=batch_size, + enable_nodal_inj_optim=False, + enable_parallel_pst_group_optim=False, + enable_bb_outage=params.ga_config.enable_bb_outage, + bb_outage_as_nminus1=params.ga_config.bb_outage_as_nminus1, + clip_bb_outage_penalty=params.ga_config.clip_bb_outage_penalty, + bb_outage_more_islands_penalty=params.ga_config.bb_outage_more_islands_penalty, + ) + + if params.double_limits is not None: + dynamic_infos = update_max_mw_flows_according_to_double_limits( + dynamic_informations=tuple(static_information.dynamic_information for static_information in static_informations), + solver_configs=tuple(static_information.solver_config for static_information in static_informations), + lower_limit=params.double_limits.lower, + upper_limit=params.double_limits.upper, + ) + static_informations = tuple( + replace(static_information, dynamic_information=dynamic_info) + for static_information, dynamic_info in zip(static_informations, dynamic_infos, strict=True) + ) + + return static_informations + + +def _initialize_runtime_state( + params: DCOptimizerParameters, + static_informations: tuple[StaticInformation, ...], + topologies_per_epoch: int, +) -> tuple[BruteforceRuntimeState, float, dict[MetricType, float], list[str]]: + """Create the JAX scoring state and evaluate the unsplit baseline. + + Parameters + ---------- + params : DCOptimizerParameters + Optimization parameters for the bruteforce run. + static_informations : tuple[StaticInformation, ...] + Prepared static-information objects for all timesteps. + topologies_per_epoch : int + Number of bruteforce candidates to evaluate per epoch. + + Returns + ------- + tuple[BruteforceRuntimeState, float, dict[MetricType, float], list[str]] + Runtime state, baseline fitness, baseline metrics, and baseline worst contingency ids. + """ + dynamic_informations = tuple(static_information.dynamic_information for static_information in static_informations) + solver_configs = tuple(static_information.solver_config for static_information in static_informations) + action_set = dynamic_informations[0].action_set + assert action_set is not None, "Bruteforce optimization requires an action set." + + split_action_groups = tuple( + tuple(range(int(start), int(start) + int(count))) + for start, count in zip(action_set.action_start_indices.tolist(), action_set.n_actions_per_sub.tolist(), strict=True) + ) + n_disconnectable_branches = dynamic_informations[0].n_disconnectable_branches + max_num_splits = params.loadflow_solver_config.max_num_splits + max_num_disconnections = params.loadflow_solver_config.max_num_disconnections + descriptor_metrics: tuple[MetricType, ...] = tuple(desc.metric for desc in params.ga_config.me_descriptors) + score_chunk_fn = jax.jit( + partial( + scoring_function, + dynamic_informations=dynamic_informations, + solver_configs=solver_configs, + target_metrics=params.ga_config.target_metrics, + observed_metrics=params.ga_config.observed_metrics, + descriptor_metrics=descriptor_metrics, + n_worst_contingencies=params.ga_config.n_worst_contingencies, + ) + ) + + initial_genotype = _chunk_to_genotype( + chunk=[WorksetEntry(action_indices=(), disconnections=())], + batch_size=topologies_per_epoch, + max_num_splits=max_num_splits, + max_num_disconnections=max_num_disconnections, + ) + fitness, _descriptors, metrics, _emitter_info, _unused_random_key, _ = score_chunk_fn( + initial_genotype, jax.random.PRNGKey(0) + ) + initial_fitness = float(np.asarray(fitness[0]).item()) + initial_metrics: dict[MetricType, float] = { + metric_name: float(np.asarray(metrics[metric_name][0]).item()) for metric_name in params.ga_config.observed_metrics + } + initial_case_ids = _case_ids_from_metric(metrics["case_indices"][0], solver_configs[0].contingency_ids) + + runtime_state = BruteforceRuntimeState( + dynamic_informations=dynamic_informations, + solver_configs=solver_configs, + score_chunk_fn=score_chunk_fn, + workset=iter_workset( + action_start_indices=action_set.action_start_indices.tolist(), + n_actions_per_sub=action_set.n_actions_per_sub.tolist(), + max_num_splits=max_num_splits, + n_disconnectable_branches=n_disconnectable_branches, + max_num_disconnections=max_num_disconnections, + ), + total_workset_size=count_workset_size( + split_action_groups=split_action_groups, + max_num_splits=max_num_splits, + n_disconnectable_branches=n_disconnectable_branches, + max_num_disconnections=max_num_disconnections, + ), + best_fitness=initial_fitness, + ) + return runtime_state, initial_fitness, initial_metrics, initial_case_ids + + +def _chunk_to_genotype( + chunk: list[WorksetEntry], + batch_size: int, + max_num_splits: int, + max_num_disconnections: int, +) -> Genotype: + """Convert a lazy workset chunk into a fixed-size genotype batch. + + Parameters + ---------- + chunk : list[WorksetEntry] + Bruteforce candidates to convert. + batch_size : int + Fixed batch size expected by the scoring function. Usually this is equal to len(chunk) except on the last batch where + the chunk might be shorter. We still pad to a fixed shape here so the JIT-compiled scoring path can reuse the same + array shapes for every epoch. + max_num_splits : int + Maximum number of split actions per candidate. + max_num_disconnections : int + Maximum number of disconnections per candidate. + + Returns + ------- + Genotype + Padded genotype batch suitable for the shared DC scoring function. + """ + action_index = np.full((batch_size, max_num_splits), int_max(), dtype=int) + disconnections = np.full((batch_size, max_num_disconnections), int_max(), dtype=int) + + for row_index, entry in enumerate(chunk): + if entry.action_indices: + action_index[row_index, : len(entry.action_indices)] = np.asarray(entry.action_indices, dtype=int) + if entry.disconnections: + disconnections[row_index, : len(entry.disconnections)] = np.asarray(entry.disconnections, dtype=int) + + return Genotype( + action_index=jnp.asarray(action_index), + disconnections=jnp.asarray(disconnections), + nodal_injections_optimized=None, + ) + + +def _convert_improved_topologies( + genotypes: Genotype, + fitness: np.ndarray, + metrics: dict[str, Any], + observed_metrics: tuple[MetricType, ...], + contingency_ids: list[str], + survivor_indices: np.ndarray, +) -> list[Topology]: + """Convert improved chunk members into result topologies. + + Parameters + ---------- + genotypes : Genotype + Scored genotype batch. + fitness : np.ndarray + Fitness values for the valid rows in the current chunk. + metrics : dict[str, Any] + Metrics returned by the shared scoring function. + observed_metrics : tuple[MetricType, ...] + Metrics that should be copied into the result payload. + contingency_ids : list[str] + Contingency identifiers used to resolve worst-case indices. + survivor_indices : np.ndarray + Indices of candidates that improve on the baseline fitness. + + Returns + ------- + list[Topology] + Result topologies corresponding to the improved candidates. + """ + topologies = [] + for row_index in survivor_indices.tolist(): + action_indices = [int(value) for value in np.asarray(genotypes.action_index[row_index]) if int(value) != int_max()] + disconnections = [int(value) for value in np.asarray(genotypes.disconnections[row_index]) if int(value) != int_max()] + topologies.append( + Topology( + actions=action_indices, + disconnections=disconnections, + pst_setpoints=None, + metrics=Metrics( + fitness=float(fitness[row_index]), + extra_scores={ + metric_name: float(np.asarray(metrics[metric_name][row_index]).item()) + for metric_name in observed_metrics + }, + worst_k_contingency_cases=_case_ids_from_metric(metrics["case_indices"][row_index], contingency_ids), + ), + ) + ) + return topologies + + +def _case_ids_from_metric(case_indices: Int[ArrayLike, " n_cases"], contingency_ids: Sequence[str]) -> list[str]: + """Resolve worst-case contingency indices to contingency ids. + + Parameters + ---------- + case_indices : Int[ArrayLike, " n_cases"] + Indices of the worst contingencies as returned by the scoring function. + contingency_ids : Sequence[str] + Ordered contingency identifiers corresponding to solver output indices. + + Returns + ------- + list[str] + Contingency identifiers referenced by ``case_indices``. + """ + return np.asarray(contingency_ids)[np.asarray(case_indices).astype(int)].tolist() diff --git a/packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc_bruteforce/worker.py b/packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc_bruteforce/worker.py new file mode 100644 index 000000000..65bdf40f3 --- /dev/null +++ b/packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc_bruteforce/worker.py @@ -0,0 +1,262 @@ +# Copyright 2026 50Hertz Transmission GmbH and Elia Transmission Belgium SA/NV +# +# This Source Code Form is subject to the terms of the Mozilla Public License, v. 2.0. +# If a copy of the MPL was not distributed with this file, +# you can obtain one at https://mozilla.org/MPL/2.0/. +# Mozilla Public License, version 2.0 + +"""Standalone Kafka worker for DC bruteforce optimization.""" + +import time +from functools import partial +from uuid import uuid4 + +import jax +import structlog +from beartype.typing import Callable +from confluent_kafka import Producer +from fsspec import AbstractFileSystem +from pydantic import BaseModel +from toop_engine_contingency_analysis.ac_loadflow_service.kafka_client import LongRunningKafkaConsumer +from toop_engine_interfaces.messages.protobuf_message_factory import serialize_message +from toop_engine_topology_optimizer.dc.worker.idle_loop import idle_loop +from toop_engine_topology_optimizer.dc_bruteforce.optimizer import ( + OptimizerData, + convert_topologies_to_messages, + extract_topologies, + get_num_branch_topologies_tried, + initialize_optimization, + is_exhausted, + run_epoch, +) +from toop_engine_topology_optimizer.interfaces.messages.commons import GridFile, OptimizerType +from toop_engine_topology_optimizer.interfaces.messages.dc_params import DCOptimizerParameters +from toop_engine_topology_optimizer.interfaces.messages.heartbeats import ( + Heartbeat, + HeartbeatUnion, + OptimizationStartedHeartbeat, + OptimizationStatsHeartbeat, +) +from toop_engine_topology_optimizer.interfaces.messages.results import ( + OptimizationStartedResult, + OptimizationStoppedResult, + Result, + ResultUnion, +) + +logger = structlog.get_logger(__name__) + + +class Args(BaseModel): + """Launch arguments for the standalone bruteforce worker.""" + + kafka_broker: str = "localhost:9092" + """The Kafka broker to connect to.""" + + optimizer_command_topic: str = "commands" + """The Kafka topic to listen for commands on.""" + + optimizer_results_topic: str = "results" + """The topic to push results to.""" + + optimizer_heartbeat_topic: str = "heartbeat" + """The topic to push heartbeats to.""" + + heartbeat_interval_ms: int = 1000 + """The interval in milliseconds to send heartbeats.""" + + max_command_age_hours: float = 3.0 + """The maximum age of a command in hours before it is ignored.""" + + +def push_topologies(optimizer_data: OptimizerData, epoch: int, send_result_fn: Callable[[ResultUnion], None]) -> int: + """Push new bruteforce topologies to Kafka. + + Parameters + ---------- + optimizer_data : OptimizerData + The optimizer state containing the latest improved topologies to emit. + epoch : int + The current epoch, used in the emitted result messages. + send_result_fn : Callable[[ResultUnion], None] + A function to call to queue results for the Kafka results topic. + + Returns + ------- + int + The number of topology messages queued for emission. + """ + 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) + + +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 the standalone bruteforce optimization loop. + + Parameters + ---------- + dc_params : DCOptimizerParameters + Parameters controlling the bruteforce optimization run. + grid_files : list[GridFile] + Grid files to load, where each file 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 may be called + multiple times within a single epoch. + 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 + Identifier of the optimization run. This value is propagated into emitted results. + processed_gridfile_fs : AbstractFileSystem + Filesystem containing the preprocessed grid data for all timesteps. + """ + logger.info(f"Initializing DC bruteforce 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 exc: + send_result_fn(OptimizationStoppedResult(reason="error", message=str(exc))) + flush_result_fn() + logger.error(f"Error during bruteforce initialization {optimization_id}: {exc}") + return + + epoch = 1 + start_time = time.time() + while True: + 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} bruteforce strategies to results topic," + f" best fitness: {optimizer_data.runtime_state.best_fitness}, epoch: {epoch}" + f" progress: {get_num_branch_topologies_tried(optimizer_data)} " + f"/ {optimizer_data.runtime_state.total_workset_size}" + ) + except Exception as exc: + send_result_fn(OptimizationStoppedResult(reason="error", message=str(exc))) + flush_result_fn() + logger.error(f"Error during bruteforce optimization {optimization_id}, epoch {epoch}: {exc}") + return + + epoch += 1 + send_heartbeat_fn( + OptimizationStatsHeartbeat( + optimization_id=optimization_id, + wall_time=time.time() - start_time, + iteration=epoch, + num_branch_topologies_tried=get_num_branch_topologies_tried(optimizer_data), + num_injection_topologies_tried=get_num_branch_topologies_tried(optimizer_data), + ) + ) + + if is_exhausted(optimizer_data): + send_result_fn(OptimizationStoppedResult(epoch=epoch, reason="converged", message="workset exhausted")) + flush_result_fn() + return + + if time.time() - start_time > dc_params.ga_config.runtime_seconds: + send_result_fn(OptimizationStoppedResult(epoch=epoch, reason="converged", message="runtime limit")) + flush_result_fn() + return + + +def main( + args: Args, + processed_gridfile_fs: AbstractFileSystem, + producer: Producer, + command_consumer: LongRunningKafkaConsumer, +) -> None: + """Run the standalone bruteforce worker loop. + + Parameters + ---------- + args : Args + Worker launch arguments. + processed_gridfile_fs : AbstractFileSystem + Filesystem containing the preprocessed grid data. + producer : Producer + Kafka producer used to send heartbeats and results. + command_consumer : LongRunningKafkaConsumer + Kafka consumer used to receive optimization commands. + """ + instance_id = str(uuid4()) + logger.info(f"Starting DC bruteforce 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: + 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: + 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() + + def flush_results() -> None: + 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=flush_results, + send_heartbeat_fn=partial(send_heartbeat, ping_consumer=True), + optimization_id=command.optimization_id, + processed_gridfile_fs=processed_gridfile_fs, + ) + command_consumer.stop_processing() diff --git a/packages/topology_optimizer_pkg/tests/benchmark/test_dc_bruteforce_generator.py b/packages/topology_optimizer_pkg/tests/benchmark/test_dc_bruteforce_generator.py new file mode 100644 index 000000000..8c9bfcef3 --- /dev/null +++ b/packages/topology_optimizer_pkg/tests/benchmark/test_dc_bruteforce_generator.py @@ -0,0 +1,38 @@ +# Copyright 2026 50Hertz Transmission GmbH and Elia Transmission Belgium SA/NV +# +# This Source Code Form is subject to the terms of the Mozilla Public License, v. 2.0. +# If a copy of the MPL was not distributed with this file, +# you can obtain one at https://mozilla.org/MPL/2.0/. +# Mozilla Public License, version 2.0 + +from time import perf_counter + +from toop_engine_topology_optimizer.dc_bruteforce.generator import iter_workset, take_workset_chunk + + +def test_dc_bruteforce_generator_benchmark(record_property) -> None: + split_action_groups = tuple(tuple(range(group_start, group_start + 8)) for group_start in range(0, 40 * 8, 8)) + action_start_indices = tuple(group_start for group_start in range(0, 40 * 8, 8)) + n_actions_per_sub = (8,) * 40 + chunk_size = 50_000 + + start = perf_counter() + chunk = take_workset_chunk( + iter_workset( + action_start_indices=action_start_indices, + n_actions_per_sub=n_actions_per_sub, + max_num_splits=3, + n_disconnectable_branches=20, + max_num_disconnections=2, + ), + chunk_size, + ) + elapsed_seconds = perf_counter() - start + throughput = chunk_size / elapsed_seconds + + record_property("dc_bruteforce_generator_chunk_size", chunk_size) + record_property("dc_bruteforce_generator_elapsed_seconds", elapsed_seconds) + record_property("dc_bruteforce_generator_topologies_per_second", throughput) + + assert len(chunk) == chunk_size + assert elapsed_seconds < 2.0 diff --git a/packages/topology_optimizer_pkg/tests/dc/worker/test_dc_optimizer_worker.py b/packages/topology_optimizer_pkg/tests/dc/worker/test_dc_optimizer_worker.py index 79b74c639..7820b99e7 100644 --- a/packages/topology_optimizer_pkg/tests/dc/worker/test_dc_optimizer_worker.py +++ b/packages/topology_optimizer_pkg/tests/dc/worker/test_dc_optimizer_worker.py @@ -6,9 +6,8 @@ # Mozilla Public License, version 2.0 import logging -from datetime import datetime, timedelta from pathlib import Path -from unittest.mock import Mock, patch +from unittest.mock import patch import pytest from beartype.typing import Literal, Union @@ -16,12 +15,10 @@ from fsspec.implementations.dirfs import DirFileSystem from toop_engine_contingency_analysis.ac_loadflow_service.kafka_client import LongRunningKafkaConsumer from toop_engine_interfaces.messages.protobuf_message_factory import deserialize_message, serialize_message -from toop_engine_topology_optimizer.dc.worker.worker import Args, idle_loop, main, optimization_loop +from toop_engine_topology_optimizer.dc.worker.worker import Args, main, optimization_loop from toop_engine_topology_optimizer.interfaces.messages.commands import ( Command, DCOptimizerParameters, - ShutdownCommand, - StartOptimizationCommand, ) from toop_engine_topology_optimizer.interfaces.messages.commons import Framework, GridFile from toop_engine_topology_optimizer.interfaces.messages.dc_params import BatchedMEParameters, LoadflowSolverParameters @@ -34,8 +31,6 @@ TopologyPushResult, ) -from packages.topology_optimizer_pkg.tests.fake_kafka import FakeMessage - # Ensure that tests using Kafka are not run in parallel with each other pytestmark = pytest.mark.xdist_group("kafka") @@ -77,105 +72,6 @@ def create_consumer( return consumer -@pytest.mark.timeout(60) -def test_idle_loop( - kafka_command_topic: str, - kafka_connection_str: str, -) -> None: - producer = Producer( - { - "bootstrap.servers": kafka_connection_str, - "log_level": 2, - } - ) - - command = Command( - command=StartOptimizationCommand( - dc_params=DCOptimizerParameters( - summary_frequency=1, - check_command_frequency=1, - ), - grid_files=[ - GridFile( - framework=Framework.PANDAPOWER, - grid_folder="child_folder", - ) - ], - optimization_id="test", - ) - ) - - producer.produce(kafka_command_topic, value=serialize_message(command.model_dump_json())) - producer.flush() - - consumer = LongRunningKafkaConsumer( - topic=kafka_command_topic, - bootstrap_servers=kafka_connection_str, - group_id="test_idle_loop", - client_id="test_idle_loop_client", - ) - - parsed = idle_loop( - consumer, lambda _: None, lambda _result, _optim: None, heartbeat_interval_ms=100, max_command_age_hours=2.0 - ) - assert parsed.optimization_id == "test" - assert tuple(gf.grid_folder for gf in parsed.grid_files) == ("child_folder",) - assert consumer.last_msg is not None - consumer.commit() - - command = Command(command=ShutdownCommand()) - producer.produce(kafka_command_topic, value=serialize_message(command.model_dump_json())) - producer.flush() - - with pytest.raises(SystemExit) as excinfo: - idle_loop( - consumer, lambda _: None, lambda _result, _optim: None, heartbeat_interval_ms=100, max_command_age_hours=2.0 - ) - assert excinfo.value.code == 0 - consumer.consumer.close() - - -def test_idle_loop_optimization_started_command_too_old() -> None: - mock_consumer = Mock(spec=LongRunningKafkaConsumer) - mock_consumer.consumer = Mock(spec=Consumer) - shutdown_command = Command(command=ShutdownCommand()) - # Make poll() return an OptimizationStartedCommand message - start_command = Command( - command=StartOptimizationCommand( - dc_params=DCOptimizerParameters(), - grid_files=[GridFile(framework=Framework.PYPOWSYBL, grid_folder="not/exist")], - optimization_id="test", - ), - timestamp=(datetime.now() - timedelta(hours=1)).isoformat(), - ) - start_message = FakeMessage( - value_bytes=serialize_message(start_command.model_dump_json()), - ) - shutdown_message = FakeMessage( - value_bytes=serialize_message(shutdown_command.model_dump_json()), - ) - - mock_consumer.poll.side_effect = [start_message, shutdown_message] - results = [] - heartbeats = [] - - def send_result_fn(result: ResultUnion, optimization_id: str) -> None: - results.append((result, optimization_id)) - - with pytest.raises(SystemExit): - idle_loop( - consumer=mock_consumer, - send_heartbeat_fn=lambda hb: heartbeats.append(hb), - send_result_fn=send_result_fn, - heartbeat_interval_ms=100, - max_command_age_hours=0.5, - ) - assert len(results) == 1 - assert isinstance(results[0][0], OptimizationStoppedResult) - assert results[0][1] == "test" - assert results[0][0].reason == "command-too-old" - - @pytest.mark.timeout(60) def test_main_simple( kafka_command_topic: str, diff --git a/packages/topology_optimizer_pkg/tests/dc/worker/test_idle_loop.py b/packages/topology_optimizer_pkg/tests/dc/worker/test_idle_loop.py new file mode 100644 index 000000000..017bde782 --- /dev/null +++ b/packages/topology_optimizer_pkg/tests/dc/worker/test_idle_loop.py @@ -0,0 +1,126 @@ +# Copyright 2026 50Hertz Transmission GmbH and Elia Transmission Belgium SA/NV +# +# This Source Code Form is subject to the terms of the Mozilla Public License, v. 2.0. +# If a copy of the MPL was not distributed with this file, +# you can obtain one at https://mozilla.org/MPL/2.0/. +# Mozilla Public License, version 2.0 + +from datetime import datetime, timedelta +from unittest.mock import Mock + +import pytest +from confluent_kafka import Consumer, Producer +from toop_engine_contingency_analysis.ac_loadflow_service.kafka_client import LongRunningKafkaConsumer +from toop_engine_interfaces.messages.protobuf_message_factory import serialize_message +from toop_engine_topology_optimizer.dc.worker.idle_loop import idle_loop +from toop_engine_topology_optimizer.interfaces.messages.commands import ( + Command, + DCOptimizerParameters, + ShutdownCommand, + StartOptimizationCommand, +) +from toop_engine_topology_optimizer.interfaces.messages.commons import Framework, GridFile +from toop_engine_topology_optimizer.interfaces.messages.results import OptimizationStoppedResult, ResultUnion + +from packages.topology_optimizer_pkg.tests.fake_kafka import FakeMessage + +# Ensure that tests using Kafka are not run in parallel with each other +pytestmark = pytest.mark.xdist_group("kafka") + + +@pytest.mark.timeout(60) +def test_idle_loop( + kafka_command_topic: str, + kafka_connection_str: str, +) -> None: + producer = Producer( + { + "bootstrap.servers": kafka_connection_str, + "log_level": 2, + } + ) + + command = Command( + command=StartOptimizationCommand( + dc_params=DCOptimizerParameters( + summary_frequency=1, + check_command_frequency=1, + ), + grid_files=[ + GridFile( + framework=Framework.PANDAPOWER, + grid_folder="child_folder", + ) + ], + optimization_id="test", + ) + ) + + producer.produce(kafka_command_topic, value=serialize_message(command.model_dump_json())) + producer.flush() + + consumer = LongRunningKafkaConsumer( + topic=kafka_command_topic, + bootstrap_servers=kafka_connection_str, + group_id="test_idle_loop", + client_id="test_idle_loop_client", + ) + + parsed = idle_loop( + consumer, lambda _: None, lambda _result, _optim: None, heartbeat_interval_ms=100, max_command_age_hours=2.0 + ) + assert parsed.optimization_id == "test" + assert tuple(gf.grid_folder for gf in parsed.grid_files) == ("child_folder",) + assert consumer.last_msg is not None + consumer.commit() + + command = Command(command=ShutdownCommand()) + producer.produce(kafka_command_topic, value=serialize_message(command.model_dump_json())) + producer.flush() + + with pytest.raises(SystemExit) as excinfo: + idle_loop( + consumer, lambda _: None, lambda _result, _optim: None, heartbeat_interval_ms=100, max_command_age_hours=2.0 + ) + assert excinfo.value.code == 0 + consumer.consumer.close() + + +def test_idle_loop_optimization_started_command_too_old() -> None: + mock_consumer = Mock(spec=LongRunningKafkaConsumer) + mock_consumer.consumer = Mock(spec=Consumer) + shutdown_command = Command(command=ShutdownCommand()) + start_command = Command( + command=StartOptimizationCommand( + dc_params=DCOptimizerParameters(), + grid_files=[GridFile(framework=Framework.PYPOWSYBL, grid_folder="not/exist")], + optimization_id="test", + ), + timestamp=(datetime.now() - timedelta(hours=1)).isoformat(), + ) + start_message = FakeMessage( + value_bytes=serialize_message(start_command.model_dump_json()), + ) + shutdown_message = FakeMessage( + value_bytes=serialize_message(shutdown_command.model_dump_json()), + ) + + mock_consumer.poll.side_effect = [start_message, shutdown_message] + results = [] + heartbeats = [] + + def send_result_fn(result: ResultUnion, optimization_id: str) -> None: + results.append((result, optimization_id)) + + with pytest.raises(SystemExit): + idle_loop( + consumer=mock_consumer, + send_heartbeat_fn=lambda hb: heartbeats.append(hb), + send_result_fn=send_result_fn, + heartbeat_interval_ms=100, + max_command_age_hours=0.5, + ) + assert len(results) == 1 + assert isinstance(results[0][0], OptimizationStoppedResult) + assert results[0][1] == "test" + assert results[0][0].reason == "command-too-old" diff --git a/packages/topology_optimizer_pkg/tests/dc_bruteforce/test_generator.py b/packages/topology_optimizer_pkg/tests/dc_bruteforce/test_generator.py new file mode 100644 index 000000000..061483e04 --- /dev/null +++ b/packages/topology_optimizer_pkg/tests/dc_bruteforce/test_generator.py @@ -0,0 +1,120 @@ +# Copyright 2026 50Hertz Transmission GmbH and Elia Transmission Belgium SA/NV +# +# This Source Code Form is subject to the terms of the Mozilla Public License, v. 2.0. +# If a copy of the MPL was not distributed with this file, +# you can obtain one at https://mozilla.org/MPL/2.0/. +# Mozilla Public License, version 2.0 + +from pathlib import Path + +import numpy as np +import pytest +from toop_engine_dc_solver.jax.inputs import load_static_information +from toop_engine_interfaces.folder_structure import PREPROCESSING_PATHS +from toop_engine_topology_optimizer.dc_bruteforce.generator import ( + WorksetEntry, + count_workset_size, + iter_workset, + take_workset_chunk, +) + + +def test_iter_workset_from_groups_exhaustive_count() -> None: + split_action_groups = ((0, 1), (2, 3, 4), (5,)) + action_start_indices = (0, 2, 5) + n_actions_per_sub = (2, 3, 1) + + workset = list( + iter_workset( + action_start_indices=action_start_indices, + n_actions_per_sub=n_actions_per_sub, + max_num_splits=2, + n_disconnectable_branches=2, + max_num_disconnections=1, + ) + ) + + expected_count = count_workset_size( + split_action_groups=split_action_groups, + max_num_splits=2, + n_disconnectable_branches=2, + max_num_disconnections=1, + ) + assert len(workset) == expected_count + assert workset[0] == WorksetEntry(action_indices=(), disconnections=()) + assert WorksetEntry(action_indices=(0, 2), disconnections=(1,)) in workset + assert all(len(entry.action_indices) <= 2 for entry in workset) + + +def test_take_workset_chunk_advances_iterator() -> None: + workset = iter_workset( + action_start_indices=(0, 2), + n_actions_per_sub=(2, 2), + max_num_splits=2, + n_disconnectable_branches=1, + max_num_disconnections=0, + ) + + first_chunk = take_workset_chunk(workset, 3) + second_chunk = take_workset_chunk(workset, 3) + + assert first_chunk == [ + WorksetEntry(action_indices=(), disconnections=()), + WorksetEntry(action_indices=(0,), disconnections=()), + WorksetEntry(action_indices=(1,), disconnections=()), + ] + assert second_chunk == [ + WorksetEntry(action_indices=(2,), disconnections=()), + WorksetEntry(action_indices=(3,), disconnections=()), + WorksetEntry(action_indices=(0, 2), disconnections=()), + ] + + +def test_iter_workset_includes_all_single_split_actions(_grid_folder: Path) -> None: + static_information = load_static_information( + _grid_folder / "complex_grid" / PREPROCESSING_PATHS["static_information_file_path"] + ) + action_set = static_information.dynamic_information.action_set + + chunk = take_workset_chunk( + iter_workset( + action_start_indices=action_set.action_start_indices.tolist(), + n_actions_per_sub=action_set.n_actions_per_sub.tolist(), + max_num_splits=1, + n_disconnectable_branches=0, + max_num_disconnections=0, + ), + len(action_set) + 1, + ) + + assert chunk[0] == WorksetEntry(action_indices=(), disconnections=()) + generated_single_split_actions = {entry.action_indices[0] for entry in chunk[1:]} + assert generated_single_split_actions == set(range(len(action_set))) + + +def test_iter_workset_uses_action_set_boundaries(_grid_folder: Path) -> None: + static_information = load_static_information( + _grid_folder / "complex_grid" / PREPROCESSING_PATHS["static_information_file_path"] + ) + action_set = static_information.dynamic_information.action_set + + chunk = take_workset_chunk( + iter_workset( + action_start_indices=action_set.action_start_indices.tolist(), + n_actions_per_sub=action_set.n_actions_per_sub.tolist(), + max_num_splits=2, + n_disconnectable_branches=int(static_information.dynamic_information.n_disconnectable_branches), + max_num_disconnections=1, + ), + 128, + ) + + substation_correspondence = np.asarray(action_set.substation_correspondence).astype(int).tolist() + for entry in chunk: + chosen_substations = {substation_correspondence[action_index] for action_index in entry.action_indices} + assert len(chosen_substations) == len(entry.action_indices) + + +def test_take_workset_chunk_rejects_non_positive_size() -> None: + with pytest.raises(ValueError, match="chunk_size must be positive"): + take_workset_chunk(iter(()), 0) diff --git a/packages/topology_optimizer_pkg/tests/dc_bruteforce/test_optimizer.py b/packages/topology_optimizer_pkg/tests/dc_bruteforce/test_optimizer.py new file mode 100644 index 000000000..191bdfabb --- /dev/null +++ b/packages/topology_optimizer_pkg/tests/dc_bruteforce/test_optimizer.py @@ -0,0 +1,109 @@ +# Copyright 2026 50Hertz Transmission GmbH and Elia Transmission Belgium SA/NV +# +# This Source Code Form is subject to the terms of the Mozilla Public License, v. 2.0. +# If a copy of the MPL was not distributed with this file, +# you can obtain one at https://mozilla.org/MPL/2.0/. +# Mozilla Public License, version 2.0 + +from fsspec.implementations.dirfs import DirFileSystem +from toop_engine_interfaces.messages.preprocess.preprocess_results import StaticInformationStats +from toop_engine_topology_optimizer.dc_bruteforce.optimizer import ( + OptimizerData, + convert_topologies_to_messages, + extract_topologies, + get_num_branch_topologies_tried, + initialize_optimization, + is_exhausted, + run_epoch, +) +from toop_engine_topology_optimizer.interfaces.messages.commands import StartOptimizationCommand +from toop_engine_topology_optimizer.interfaces.messages.commons import Framework, GridFile +from toop_engine_topology_optimizer.interfaces.messages.dc_params import ( + BatchedMEParameters, + DCOptimizerParameters, + DescriptorDef, + LoadflowSolverParameters, +) +from toop_engine_topology_optimizer.interfaces.messages.results import Strategy, Topology, TopologyPushResult + + +def _ga_config(runtime_seconds: float = 5, enable_nodal_inj_optim: bool = False) -> BatchedMEParameters: + return BatchedMEParameters( + runtime_seconds=runtime_seconds, + enable_nodal_inj_optim=enable_nodal_inj_optim, + iterations_per_epoch=8, + me_descriptors=( + DescriptorDef(metric="split_subs", num_cells=2), + DescriptorDef(metric="switching_distance", num_cells=45), + ), + ) + + +def test_initialize_and_run_epoch(grid_folder: str) -> None: + start_opt_command = StartOptimizationCommand( + dc_params=DCOptimizerParameters( + summary_frequency=1, + check_command_frequency=1, + ga_config=_ga_config(runtime_seconds=5), + loadflow_solver_config=LoadflowSolverParameters(max_num_splits=1, max_num_disconnections=0), + ), + grid_files=[GridFile(framework=Framework.PANDAPOWER, grid_folder="case14")], + optimization_id="test", + ) + + processed_gridfile_fs = DirFileSystem(str(grid_folder)) + topologies_per_epoch = ( + start_opt_command.dc_params.ga_config.iterations_per_epoch + * start_opt_command.dc_params.loadflow_solver_config.batch_size + ) + optimizer_data, stats, initial_strategy = initialize_optimization( + params=start_opt_command.dc_params, + optimization_id="test123", + static_information_files=tuple(gf.static_information_file for gf in start_opt_command.grid_files), + processed_gridfile_fs=processed_gridfile_fs, + ) + + assert isinstance(optimizer_data, OptimizerData) + assert isinstance(stats[0], StaticInformationStats) + assert isinstance(initial_strategy, Strategy) + + optimizer_data = run_epoch(optimizer_data) + assert get_num_branch_topologies_tried(optimizer_data) == min( + topologies_per_epoch, optimizer_data.runtime_state.total_workset_size + ) + + topologies = extract_topologies(optimizer_data) + assert isinstance(topologies, list) + if topologies: + assert isinstance(topologies[0], Topology) + + messages = convert_topologies_to_messages(topologies, epoch=1) + if messages: + assert isinstance(messages[0], TopologyPushResult) + assert isinstance(messages[0].strategy, Strategy) + assert len(messages[0].strategy.timesteps) == 1 + + assert extract_topologies(optimizer_data) == [] + assert is_exhausted(optimizer_data) is (optimizer_data.runtime_state.total_workset_size <= topologies_per_epoch) + + +def test_initialize_rejects_pst_optimization(grid_folder: str) -> None: + processed_gridfile_fs = DirFileSystem(str(grid_folder)) + params = DCOptimizerParameters( + ga_config=_ga_config(enable_nodal_inj_optim=True), + loadflow_solver_config=LoadflowSolverParameters(max_num_splits=1), + ) + + try: + initialize_optimization( + params=params, + optimization_id="test", + static_information_files=( + GridFile(framework=Framework.PANDAPOWER, grid_folder="case14").static_information_file, + ), + processed_gridfile_fs=processed_gridfile_fs, + ) + except ValueError as exc: + assert "PSTs untouched" in str(exc) + else: + raise AssertionError("Expected initialize_optimization to reject nodal injection optimization.") diff --git a/packages/topology_optimizer_pkg/tests/dc_bruteforce/test_worker.py b/packages/topology_optimizer_pkg/tests/dc_bruteforce/test_worker.py new file mode 100644 index 000000000..ca3f8970f --- /dev/null +++ b/packages/topology_optimizer_pkg/tests/dc_bruteforce/test_worker.py @@ -0,0 +1,66 @@ +# Copyright 2026 50Hertz Transmission GmbH and Elia Transmission Belgium SA/NV +# +# This Source Code Form is subject to the terms of the Mozilla Public License, v. 2.0. +# If a copy of the MPL was not distributed with this file, +# you can obtain one at https://mozilla.org/MPL/2.0/. +# Mozilla Public License, version 2.0 + +from fsspec.implementations.dirfs import DirFileSystem +from toop_engine_topology_optimizer.dc_bruteforce.worker import optimization_loop +from toop_engine_topology_optimizer.interfaces.messages.commons import Framework, GridFile +from toop_engine_topology_optimizer.interfaces.messages.dc_params import ( + BatchedMEParameters, + DCOptimizerParameters, + DescriptorDef, + LoadflowSolverParameters, +) +from toop_engine_topology_optimizer.interfaces.messages.heartbeats import ( + OptimizationStartedHeartbeat, + OptimizationStatsHeartbeat, +) +from toop_engine_topology_optimizer.interfaces.messages.results import ( + OptimizationStartedResult, + OptimizationStoppedResult, + ResultUnion, + TopologyPushResult, +) + + +def _ga_config() -> BatchedMEParameters: + return BatchedMEParameters( + runtime_seconds=30, + iterations_per_epoch=32, + me_descriptors=( + DescriptorDef(metric="split_subs", num_cells=2), + DescriptorDef(metric="switching_distance", num_cells=45), + ), + ) + + +def test_optimization_loop_emits_start_and_stop(grid_folder: str) -> None: + results: list[ResultUnion] = [] + heartbeats = [] + + optimization_loop( + dc_params=DCOptimizerParameters( + summary_frequency=1, + check_command_frequency=1, + ga_config=_ga_config(), + loadflow_solver_config=LoadflowSolverParameters(max_num_splits=1, max_num_disconnections=0), + ), + grid_files=[GridFile(framework=Framework.PANDAPOWER, grid_folder="case14")], + send_result_fn=results.append, + flush_result_fn=lambda: None, + send_heartbeat_fn=heartbeats.append, + optimization_id="test-bruteforce", + processed_gridfile_fs=DirFileSystem(str(grid_folder)), + ) + + assert any(isinstance(heartbeat, OptimizationStartedHeartbeat) for heartbeat in heartbeats) + assert any(isinstance(heartbeat, OptimizationStatsHeartbeat) for heartbeat in heartbeats) + assert isinstance(results[0], OptimizationStartedResult) + assert isinstance(results[-1], OptimizationStoppedResult) + assert results[-1].reason == "converged" + assert all( + isinstance(result, (OptimizationStartedResult, OptimizationStoppedResult, TopologyPushResult)) for result in results + ) From 4eecad13f36b2baa5605b86b1efb1b23331df476 Mon Sep 17 00:00:00 2001 From: Nico Westerbeck Date: Sun, 5 Jul 2026 21:00:10 +0000 Subject: [PATCH 2/7] refactor: Use topology looper Signed-off-by: Nico Westerbeck --- .../dc/genetic_functions/scoring_functions.py | 91 +++-- .../dc_bruteforce/optimizer.py | 336 +++++++++++------- .../tests/dc_bruteforce/test_optimizer.py | 96 +++++ 3 files changed, 374 insertions(+), 149 deletions(-) diff --git a/packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/genetic_functions/scoring_functions.py b/packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/genetic_functions/scoring_functions.py index f95455731..4a0f250bb 100644 --- a/packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/genetic_functions/scoring_functions.py +++ b/packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc/genetic_functions/scoring_functions.py @@ -23,6 +23,7 @@ NodalInjOptimResults, NodalInjStartOptions, SolverConfig, + SolverLoadflowResults, int_max, ) from toop_engine_interfaces.types import MetricType @@ -69,6 +70,64 @@ } +def get_aggregate_metrics( + lf_res: SolverLoadflowResults, + success: Bool[Array, " batch_size"], + dynamic_information: DynamicInformation, + observed_metrics: tuple[MetricType, ...], + n_worst_contingencies: int, +) -> dict[str, Array]: + """Aggregate one scored batch of loadflow results into optimizer metrics. + + Parameters + ---------- + lf_res : SolverLoadflowResults + Batched loadflow results returned by the symmetric solver path. + success : Bool[Array, " batch_size"] + Per-topology success mask used to invalidate failed candidates. + dynamic_information : DynamicInformation + Dynamic grid information required by the aggregation metrics. + observed_metrics : tuple[MetricType, ...] + Metrics that should be preserved for scoring and result emission. + n_worst_contingencies : int + Number of worst contingency indices to keep for each topology. + + Returns + ------- + dict[str, Array] + Aggregated metrics for the full batch, including ``case_indices`` and + ``top_k_overloads_n_1``. + """ + # 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 + + # The scoring function runs the loadflow and computes the overload energy def compute_overloads( topologies: Genotype, @@ -114,33 +173,13 @@ def compute_overloads( 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 = get_aggregate_metrics( + lf_res=lf_res, + success=success, + dynamic_information=dynamic_information, + observed_metrics=observed_metrics, + n_worst_contingencies=n_worst_contingencies, ) - 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 diff --git a/packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc_bruteforce/optimizer.py b/packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc_bruteforce/optimizer.py index 1f15788d8..4746dc529 100644 --- a/packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc_bruteforce/optimizer.py +++ b/packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc_bruteforce/optimizer.py @@ -9,28 +9,35 @@ import time from dataclasses import dataclass, field -from functools import partial from pathlib import Path import jax import jax.numpy as jnp import numpy as np -from beartype.typing import Any, Callable, Iterator, Sequence +from beartype.typing import Any, Iterator, Sequence from fsspec import AbstractFileSystem from jax_dataclasses import replace -from jaxtyping import Array, ArrayLike, Float, Int, PRNGKeyArray -from qdax.core.emitters.standard_emitters import ExtraScores +from jaxtyping import Array, ArrayLike, Float, Int from toop_engine_dc_solver.jax.inputs import load_static_information_fs -from toop_engine_dc_solver.jax.types import DynamicInformation, SolverConfig, StaticInformation, int_max +from toop_engine_dc_solver.jax.topology_looper import run_solver_symmetric +from toop_engine_dc_solver.jax.types import ( + ActionIndexComputations, + DynamicInformation, + SolverConfig, + SolverLoadflowResults, + StaticInformation, + int_max, +) from toop_engine_dc_solver.preprocess.convert_to_jax import StaticInformationStats, extract_static_information_stats from toop_engine_interfaces.types import MetricType -from toop_engine_topology_optimizer.dc.genetic_functions.genotype import Genotype from toop_engine_topology_optimizer.dc.genetic_functions.initialization import ( update_max_mw_flows_according_to_double_limits, update_static_information, verify_static_information, ) -from toop_engine_topology_optimizer.dc.genetic_functions.scoring_functions import scoring_function +from toop_engine_topology_optimizer.dc.genetic_functions.scoring_functions import ( + get_aggregate_metrics, +) from toop_engine_topology_optimizer.dc_bruteforce.generator import ( WorksetEntry, count_workset_size, @@ -41,31 +48,16 @@ from toop_engine_topology_optimizer.interfaces.messages.results import Metrics, Strategy, Topology, TopologyPushResult from toop_engine_topology_optimizer.interfaces.models.base_storage import hash_strategy -ScoreChunkFn = Callable[ - [Genotype, PRNGKeyArray], - tuple[ - Float[Array, " batch_size"], - Int[Array, " batch_size n_dims"], - dict[str, Any], - ExtraScores, - PRNGKeyArray, - Genotype, - ], -] - @dataclass class BruteforceRuntimeState: """State kept across bruteforce epochs.""" - dynamic_informations: tuple[DynamicInformation, ...] - """Dynamic solver inputs for each timestep.""" + dynamic_information: DynamicInformation + """Dynamic solver inputs for the single supported timestep.""" - solver_configs: tuple[SolverConfig, ...] - """Static solver configuration for each timestep.""" - - score_chunk_fn: ScoreChunkFn - """Jitted scoring function that evaluates a chunk of bruteforce candidates.""" + solver_config: SolverConfig + """Static solver configuration for the single supported timestep.""" workset: Iterator[WorksetEntry] """Lazy iterator over the remaining bruteforce candidates.""" @@ -138,21 +130,20 @@ def initialize_optimization( strategy. """ topologies_per_epoch = _get_topologies_per_epoch(params) + solver_batch_size = params.loadflow_solver_config.batch_size - static_informations = _load_and_prepare_static_informations( + static_information = _load_and_prepare_static_informations( params=params, static_information_files=static_information_files, processed_gridfile_fs=processed_gridfile_fs, - batch_size=topologies_per_epoch, + batch_size=solver_batch_size, ) - _validate_bruteforce_inputs(params, topologies_per_epoch, static_informations[0].dynamic_information) + _validate_bruteforce_inputs(params, topologies_per_epoch, static_information.dynamic_information) runtime_state, initial_fitness, initial_metrics, initial_case_ids = _initialize_runtime_state( params=params, - static_informations=static_informations, - topologies_per_epoch=topologies_per_epoch, + static_information=static_information, ) initial_strategy = build_initial_strategy( - n_timesteps=len(runtime_state.dynamic_informations), fitness=initial_fitness, initial_metrics=initial_metrics, initial_case_ids=initial_case_ids, @@ -165,7 +156,6 @@ def initialize_optimization( overload_n1=initial_metrics.get("overload_energy_n_1", 0.0), time="", ) - for static_information in static_informations ] return ( @@ -184,7 +174,6 @@ def initialize_optimization( def build_initial_strategy( - n_timesteps: int, fitness: float, initial_metrics: dict[MetricType, float], initial_case_ids: list[str], @@ -193,8 +182,6 @@ def build_initial_strategy( Parameters ---------- - n_timesteps : int - Number of timesteps for which the initial strategy must be emitted. fitness : float Fitness of the unsplit baseline topology. initial_metrics : dict[MetricType, float] @@ -220,7 +207,6 @@ def build_initial_strategy( pst_setpoints=None, metrics=metrics, ) - for _ in range(n_timesteps) ] ) @@ -253,26 +239,34 @@ def run_epoch(optimizer_data: OptimizerData) -> OptimizerData: max_num_splits = optimizer_data.start_params.loadflow_solver_config.max_num_splits max_num_disconnections = optimizer_data.start_params.loadflow_solver_config.max_num_disconnections - - genotype = _chunk_to_genotype( + evaluated_count = len(chunk) + topology_chunk, disconnection_chunk = _chunk_to_topologies( chunk=chunk, - batch_size=topologies_per_epoch, + chunk_size=topologies_per_epoch, max_num_splits=max_num_splits, max_num_disconnections=max_num_disconnections, ) - fitness, _descriptors, metrics, _emitter_info, _unused_random_key, scored_genotypes = runtime_state.score_chunk_fn( - genotype, jax.random.PRNGKey(0) + + fitness, metrics = _score_chunk( + topology_chunk=topology_chunk, + disconnection_chunk=disconnection_chunk, + evaluated_count=evaluated_count, + dynamic_information=runtime_state.dynamic_information, + solver_config=runtime_state.solver_config, + target_metrics=optimizer_data.start_params.ga_config.target_metrics, + observed_metrics=optimizer_data.start_params.ga_config.observed_metrics, + n_worst_contingencies=optimizer_data.start_params.ga_config.n_worst_contingencies, ) - evaluated_count = len(chunk) evaluated_fitness = np.asarray(fitness[:evaluated_count]) improved_indices = np.flatnonzero(np.isfinite(evaluated_fitness) & (evaluated_fitness > optimizer_data.initial_fitness)) - topologies = _convert_improved_topologies( - genotypes=scored_genotypes, + topologies_out = _convert_improved_topologies( + topology_actions=topology_chunk.action, + disconnection_chunk=disconnection_chunk, fitness=evaluated_fitness, - metrics=metrics, + metrics={metric_name: np.asarray(metric_values[:evaluated_count]) for metric_name, metric_values in metrics.items()}, observed_metrics=optimizer_data.start_params.ga_config.observed_metrics, - contingency_ids=list(runtime_state.solver_configs[0].contingency_ids), + contingency_ids=list(runtime_state.solver_config.contingency_ids), survivor_indices=improved_indices, ) @@ -281,15 +275,17 @@ def run_epoch(optimizer_data: OptimizerData) -> OptimizerData: if finite_fitness.size > 0: best_fitness = max(best_fitness, float(finite_fitness.max())) + exhausted = evaluated_count < topologies_per_epoch + return replace( optimizer_data, runtime_state=replace( runtime_state, total_branch_topologies_tried=runtime_state.total_branch_topologies_tried + evaluated_count, - exhausted=evaluated_count < topologies_per_epoch, + exhausted=exhausted, best_fitness=best_fitness, ), - latest_topologies=topologies, + latest_topologies=topologies_out, ) @@ -428,15 +424,15 @@ def _load_and_prepare_static_informations( static_information_files: Sequence[str | Path], processed_gridfile_fs: AbstractFileSystem, batch_size: int, -) -> tuple[StaticInformation, ...]: - """Load and preprocess static information for bruteforce execution. +) -> StaticInformation: + """Load and preprocess the single static information used by bruteforce. Parameters ---------- params : DCOptimizerParameters Optimization parameters controlling preprocessing options. static_information_files : Sequence[str | Path] - Paths to the static-information files for all timesteps. + Paths to the preprocessed static-information files. Bruteforce supports exactly one. processed_gridfile_fs : AbstractFileSystem Filesystem containing the preprocessed grid data. batch_size : int @@ -444,21 +440,24 @@ def _load_and_prepare_static_informations( Returns ------- - tuple[StaticInformation, ...] - Prepared static-information objects for all timesteps. + StaticInformation + Prepared static-information object for the single supported timestep. """ - static_informations = tuple( - load_static_information_fs(filesystem=processed_gridfile_fs, filename=str(filename)) - for filename in static_information_files + assert len(static_information_files) == 1, "Bruteforce optimizer supports exactly one static information file." + static_information = load_static_information_fs( + filesystem=processed_gridfile_fs, + filename=str(static_information_files[0]), ) + assert static_information.dynamic_information.n_timesteps == 1, "Bruteforce optimizer supports exactly one timestep." + verify_static_information( - static_informations, + (static_information,), params.loadflow_solver_config.max_num_disconnections, enable_nodal_inj_optim=False, enable_parallel_pst_group_optim=False, ) - static_informations = update_static_information( - static_informations, + static_information = update_static_information( + (static_information,), batch_size=batch_size, enable_nodal_inj_optim=False, enable_parallel_pst_group_optim=False, @@ -466,27 +465,23 @@ def _load_and_prepare_static_informations( bb_outage_as_nminus1=params.ga_config.bb_outage_as_nminus1, clip_bb_outage_penalty=params.ga_config.clip_bb_outage_penalty, bb_outage_more_islands_penalty=params.ga_config.bb_outage_more_islands_penalty, - ) + )[0] if params.double_limits is not None: - dynamic_infos = update_max_mw_flows_according_to_double_limits( - dynamic_informations=tuple(static_information.dynamic_information for static_information in static_informations), - solver_configs=tuple(static_information.solver_config for static_information in static_informations), + dynamic_information = update_max_mw_flows_according_to_double_limits( + dynamic_informations=(static_information.dynamic_information,), + solver_configs=(static_information.solver_config,), lower_limit=params.double_limits.lower, upper_limit=params.double_limits.upper, - ) - static_informations = tuple( - replace(static_information, dynamic_information=dynamic_info) - for static_information, dynamic_info in zip(static_informations, dynamic_infos, strict=True) - ) + )[0] + static_information = replace(static_information, dynamic_information=dynamic_information) - return static_informations + return static_information def _initialize_runtime_state( params: DCOptimizerParameters, - static_informations: tuple[StaticInformation, ...], - topologies_per_epoch: int, + static_information: StaticInformation, ) -> tuple[BruteforceRuntimeState, float, dict[MetricType, float], list[str]]: """Create the JAX scoring state and evaluate the unsplit baseline. @@ -494,60 +489,54 @@ def _initialize_runtime_state( ---------- params : DCOptimizerParameters Optimization parameters for the bruteforce run. - static_informations : tuple[StaticInformation, ...] - Prepared static-information objects for all timesteps. - topologies_per_epoch : int - Number of bruteforce candidates to evaluate per epoch. + static_information : StaticInformation + Prepared static-information object for the single supported timestep. Returns ------- tuple[BruteforceRuntimeState, float, dict[MetricType, float], list[str]] Runtime state, baseline fitness, baseline metrics, and baseline worst contingency ids. """ - dynamic_informations = tuple(static_information.dynamic_information for static_information in static_informations) - solver_configs = tuple(static_information.solver_config for static_information in static_informations) - action_set = dynamic_informations[0].action_set + dynamic_information = static_information.dynamic_information + solver_config = static_information.solver_config + action_set = dynamic_information.action_set assert action_set is not None, "Bruteforce optimization requires an action set." split_action_groups = tuple( tuple(range(int(start), int(start) + int(count))) for start, count in zip(action_set.action_start_indices.tolist(), action_set.n_actions_per_sub.tolist(), strict=True) ) - n_disconnectable_branches = dynamic_informations[0].n_disconnectable_branches + n_disconnectable_branches = dynamic_information.n_disconnectable_branches max_num_splits = params.loadflow_solver_config.max_num_splits max_num_disconnections = params.loadflow_solver_config.max_num_disconnections - descriptor_metrics: tuple[MetricType, ...] = tuple(desc.metric for desc in params.ga_config.me_descriptors) - score_chunk_fn = jax.jit( - partial( - scoring_function, - dynamic_informations=dynamic_informations, - solver_configs=solver_configs, - target_metrics=params.ga_config.target_metrics, - observed_metrics=params.ga_config.observed_metrics, - descriptor_metrics=descriptor_metrics, - n_worst_contingencies=params.ga_config.n_worst_contingencies, - ) - ) - - initial_genotype = _chunk_to_genotype( + solver_batch_size = params.loadflow_solver_config.batch_size + # Pad the initial topology to the size of one batch and run it through the scoring function + # Reusing chunking functions. + initial_topology_chunk, initial_disconnection_chunk = _chunk_to_topologies( chunk=[WorksetEntry(action_indices=(), disconnections=())], - batch_size=topologies_per_epoch, + chunk_size=solver_batch_size, max_num_splits=max_num_splits, max_num_disconnections=max_num_disconnections, ) - fitness, _descriptors, metrics, _emitter_info, _unused_random_key, _ = score_chunk_fn( - initial_genotype, jax.random.PRNGKey(0) + fitness, metrics = _score_chunk( + topology_chunk=initial_topology_chunk, + disconnection_chunk=initial_disconnection_chunk, + evaluated_count=1, + dynamic_information=dynamic_information, + solver_config=solver_config, + target_metrics=params.ga_config.target_metrics, + observed_metrics=params.ga_config.observed_metrics, + n_worst_contingencies=params.ga_config.n_worst_contingencies, ) initial_fitness = float(np.asarray(fitness[0]).item()) initial_metrics: dict[MetricType, float] = { metric_name: float(np.asarray(metrics[metric_name][0]).item()) for metric_name in params.ga_config.observed_metrics } - initial_case_ids = _case_ids_from_metric(metrics["case_indices"][0], solver_configs[0].contingency_ids) + initial_case_ids = _case_ids_from_metric(metrics["case_indices"][0], solver_config.contingency_ids) runtime_state = BruteforceRuntimeState( - dynamic_informations=dynamic_informations, - solver_configs=solver_configs, - score_chunk_fn=score_chunk_fn, + dynamic_information=dynamic_information, + solver_config=solver_config, workset=iter_workset( action_start_indices=action_set.action_start_indices.tolist(), n_actions_per_sub=action_set.n_actions_per_sub.tolist(), @@ -566,21 +555,118 @@ def _initialize_runtime_state( return runtime_state, initial_fitness, initial_metrics, initial_case_ids -def _chunk_to_genotype( +class _AggregateMetricsOutputFn: + """Adapt the shared batch metric helper to the solver's per-topology callback API.""" + + def __init__( + self, + dynamic_information: DynamicInformation, + observed_metrics: tuple[MetricType, ...], + n_worst_contingencies: int, + fixed_hash: int, + ) -> None: + """Create a per-topology aggregation callback for ``run_solver_symmetric``. + + Parameters + ---------- + dynamic_information : DynamicInformation + Dynamic grid information required by the metric aggregation. + observed_metrics : tuple[MetricType, ...] + Metrics that should be returned for each topology. + n_worst_contingencies : int + Number of worst contingency indices to retain. + fixed_hash : int + Stable hash used for JAX static-argument caching. + """ + self.dynamic_information = dynamic_information + self.observed_metrics = observed_metrics + self.n_worst_contingencies = n_worst_contingencies + self.fixed_hash = fixed_hash + + def __call__(self, lf_res: SolverLoadflowResults) -> dict[str, Array]: + """Aggregate one topology's solver output into the bruteforce metrics payload.""" + lf_res_batch = jax.tree_util.tree_map( + lambda leaf: jnp.expand_dims(leaf, axis=0) if leaf is not None else None, + lf_res, + ) + success = ( + jnp.expand_dims(jnp.all(lf_res.contingency_success), axis=0) + if lf_res.contingency_success is not None + else jnp.ones((1,), dtype=bool) + ) + metrics = get_aggregate_metrics( + lf_res=lf_res_batch, + success=success, + dynamic_information=self.dynamic_information, + observed_metrics=self.observed_metrics, + n_worst_contingencies=self.n_worst_contingencies, + ) + return jax.tree_util.tree_map(lambda leaf: leaf[0], metrics) + + def __hash__(self) -> int: + """Return a stable hash so JAX can cache the compiled solver loop.""" + return self.fixed_hash + + def __eq__(self, other: object) -> bool: + """Compare callbacks by hash to avoid unnecessary recompilation.""" + if not isinstance(other, _AggregateMetricsOutputFn): + return False + return hash(self) == hash(other) + + +def _score_chunk( + topology_chunk: ActionIndexComputations, + disconnection_chunk: Int[Array, " chunk_size n_disconnections"], + evaluated_count: int, + dynamic_information: DynamicInformation, + solver_config: SolverConfig, + target_metrics: tuple[tuple[MetricType, float], ...], + observed_metrics: tuple[MetricType, ...], + n_worst_contingencies: int, +) -> tuple[Float[Array, " chunk_size"], dict[str, Array]]: + """Score one padded bruteforce chunk directly in solver input format. + + The chunk can span multiple solver mini-batches; ``run_solver_symmetric`` handles + the internal mini-batching. + """ + metrics, _contingency_success = run_solver_symmetric( + topologies=topology_chunk, + disconnections=disconnection_chunk, + injections=None, + dynamic_information=dynamic_information, + solver_config=solver_config, + aggregate_output_fn=_AggregateMetricsOutputFn( + dynamic_information=dynamic_information, + observed_metrics=observed_metrics, + n_worst_contingencies=n_worst_contingencies, + fixed_hash=hash((solver_config, observed_metrics, n_worst_contingencies)), + ), + ) + + fitness = sum(-metrics[metric_name] * weight for metric_name, weight in target_metrics) + invalid_rows = jnp.arange(len(topology_chunk)) >= evaluated_count + fitness = jnp.where(invalid_rows, -jnp.inf, fitness) + for metric_name in observed_metrics: + metrics[metric_name] = jnp.where(invalid_rows, jnp.nan, metrics[metric_name]) + metrics["case_indices"] = jnp.where(invalid_rows[:, None], -1, metrics["case_indices"]) + return fitness, metrics + + +def _chunk_to_topologies( chunk: list[WorksetEntry], - batch_size: int, + chunk_size: int, max_num_splits: int, max_num_disconnections: int, -) -> Genotype: - """Convert a lazy workset chunk into a fixed-size genotype batch. +) -> tuple[ActionIndexComputations, Int[Array, " chunk_size n_disconnections"]]: + """Convert a lazy workset chunk into solver-ready topology arrays. Parameters ---------- chunk : list[WorksetEntry] Bruteforce candidates to convert. - batch_size : int - Fixed batch size expected by the scoring function. Usually this is equal to len(chunk) except on the last batch where - the chunk might be shorter. We still pad to a fixed shape here so the JIT-compiled scoring path can reuse the same + chunk_size : int + Fixed chunk size for one bruteforce epoch. Usually this is equal to ``len(chunk)`` except on the last epoch where + the chunk might be shorter. We still pad to a fixed shape here so the solver-facing chunk path can reuse the same array shapes for every epoch. max_num_splits : int Maximum number of split actions per candidate. @@ -589,11 +675,11 @@ def _chunk_to_genotype( Returns ------- - Genotype - Padded genotype batch suitable for the shared DC scoring function. + tuple[ActionIndexComputations, Int[Array, " chunk_size n_disconnections"]] + Padded action-index topologies and aligned disconnections. """ - action_index = np.full((batch_size, max_num_splits), int_max(), dtype=int) - disconnections = np.full((batch_size, max_num_disconnections), int_max(), dtype=int) + action_index = np.full((chunk_size, max_num_splits), int_max(), dtype=int) + disconnections = np.full((chunk_size, max_num_disconnections), int_max(), dtype=int) for row_index, entry in enumerate(chunk): if entry.action_indices: @@ -601,15 +687,15 @@ def _chunk_to_genotype( if entry.disconnections: disconnections[row_index, : len(entry.disconnections)] = np.asarray(entry.disconnections, dtype=int) - return Genotype( - action_index=jnp.asarray(action_index), - disconnections=jnp.asarray(disconnections), - nodal_injections_optimized=None, - ) + return ActionIndexComputations( + action=jnp.asarray(action_index), + pad_mask=jnp.ones((chunk_size,), dtype=bool), + ), jnp.asarray(disconnections) def _convert_improved_topologies( - genotypes: Genotype, + topology_actions: Int[ArrayLike, " batch_size n_splits"], + disconnection_chunk: Int[ArrayLike, " chunk_size n_disconnections"], fitness: np.ndarray, metrics: dict[str, Any], observed_metrics: tuple[MetricType, ...], @@ -620,8 +706,10 @@ def _convert_improved_topologies( Parameters ---------- - genotypes : Genotype - Scored genotype batch. + topology_actions : Int[ArrayLike, " batch_size n_splits"] + Scored topology actions for the current bruteforce chunk in action-index format. + disconnection_chunk : Int[ArrayLike, " chunk_size n_disconnections"] + Disconnections aligned with ``topology_actions`` for the current chunk. fitness : np.ndarray Fitness values for the valid rows in the current chunk. metrics : dict[str, Any] @@ -640,12 +728,14 @@ def _convert_improved_topologies( """ topologies = [] for row_index in survivor_indices.tolist(): - action_indices = [int(value) for value in np.asarray(genotypes.action_index[row_index]) if int(value) != int_max()] - disconnections = [int(value) for value in np.asarray(genotypes.disconnections[row_index]) if int(value) != int_max()] + action_indices = [int(value) for value in np.asarray(topology_actions[row_index]) if int(value) != int_max()] + disconnection_values = [ + int(value) for value in np.asarray(disconnection_chunk[row_index]) if int(value) != int_max() + ] topologies.append( Topology( actions=action_indices, - disconnections=disconnections, + disconnections=disconnection_values, pst_setpoints=None, metrics=Metrics( fitness=float(fitness[row_index]), diff --git a/packages/topology_optimizer_pkg/tests/dc_bruteforce/test_optimizer.py b/packages/topology_optimizer_pkg/tests/dc_bruteforce/test_optimizer.py index 191bdfabb..c45af7250 100644 --- a/packages/topology_optimizer_pkg/tests/dc_bruteforce/test_optimizer.py +++ b/packages/topology_optimizer_pkg/tests/dc_bruteforce/test_optimizer.py @@ -5,8 +5,12 @@ # you can obtain one at https://mozilla.org/MPL/2.0/. # Mozilla Public License, version 2.0 +import numpy as np +import toop_engine_topology_optimizer.dc_bruteforce.optimizer as bruteforce_optimizer from fsspec.implementations.dirfs import DirFileSystem from toop_engine_interfaces.messages.preprocess.preprocess_results import StaticInformationStats +from toop_engine_topology_optimizer.dc.genetic_functions.genotype import Genotype +from toop_engine_topology_optimizer.dc.genetic_functions.scoring_functions import scoring_function from toop_engine_topology_optimizer.dc_bruteforce.optimizer import ( OptimizerData, convert_topologies_to_messages, @@ -66,6 +70,10 @@ def test_initialize_and_run_epoch(grid_folder: str) -> None: assert isinstance(optimizer_data, OptimizerData) assert isinstance(stats[0], StaticInformationStats) assert isinstance(initial_strategy, Strategy) + assert ( + optimizer_data.runtime_state.solver_config.batch_size_bsdf + == start_opt_command.dc_params.loadflow_solver_config.batch_size + ) optimizer_data = run_epoch(optimizer_data) assert get_num_branch_topologies_tried(optimizer_data) == min( @@ -107,3 +115,91 @@ def test_initialize_rejects_pst_optimization(grid_folder: str) -> None: assert "PSTs untouched" in str(exc) else: raise AssertionError("Expected initialize_optimization to reject nodal injection optimization.") + + +def test_initialize_rejects_multiple_static_information_files(grid_folder: str) -> None: + processed_gridfile_fs = DirFileSystem(str(grid_folder)) + static_information_file = GridFile(framework=Framework.PANDAPOWER, grid_folder="case14").static_information_file + + try: + initialize_optimization( + params=DCOptimizerParameters( + ga_config=_ga_config(), + loadflow_solver_config=LoadflowSolverParameters(max_num_splits=1), + ), + optimization_id="test", + static_information_files=(static_information_file, static_information_file), + processed_gridfile_fs=processed_gridfile_fs, + ) + except AssertionError as exc: + assert "exactly one static information file" in str(exc) + else: + raise AssertionError("Expected initialize_optimization to reject multiple static-information files.") + + +def test_chunk_scoring_matches_genotype_scoring(grid_folder: str) -> None: + start_opt_command = StartOptimizationCommand( + dc_params=DCOptimizerParameters( + summary_frequency=1, + check_command_frequency=1, + ga_config=_ga_config(runtime_seconds=5), + loadflow_solver_config=LoadflowSolverParameters(max_num_splits=1, max_num_disconnections=0), + ), + grid_files=[GridFile(framework=Framework.PANDAPOWER, grid_folder="case14")], + optimization_id="test", + ) + + processed_gridfile_fs = DirFileSystem(str(grid_folder)) + optimizer_data, _stats, _initial_strategy = initialize_optimization( + params=start_opt_command.dc_params, + optimization_id="test123", + static_information_files=tuple(gf.static_information_file for gf in start_opt_command.grid_files), + processed_gridfile_fs=processed_gridfile_fs, + ) + + action_set = optimizer_data.runtime_state.dynamic_information.action_set + assert action_set is not None + topology_action = int(action_set.action_start_indices[0]) + topology_chunk, disconnection_chunk = bruteforce_optimizer._chunk_to_topologies( + chunk=[bruteforce_optimizer.WorksetEntry(action_indices=(topology_action,), disconnections=())], + chunk_size=optimizer_data.runtime_state.solver_config.batch_size_bsdf, + max_num_splits=start_opt_command.dc_params.loadflow_solver_config.max_num_splits, + max_num_disconnections=start_opt_command.dc_params.loadflow_solver_config.max_num_disconnections, + ) + + chunk_fitness, chunk_metrics = bruteforce_optimizer._score_chunk( + topology_chunk=topology_chunk, + disconnection_chunk=disconnection_chunk, + evaluated_count=1, + dynamic_information=optimizer_data.runtime_state.dynamic_information, + solver_config=optimizer_data.runtime_state.solver_config, + target_metrics=start_opt_command.dc_params.ga_config.target_metrics, + observed_metrics=start_opt_command.dc_params.ga_config.observed_metrics, + n_worst_contingencies=start_opt_command.dc_params.ga_config.n_worst_contingencies, + ) + + genotype = Genotype( + action_index=topology_chunk.action, + disconnections=disconnection_chunk, + nodal_injections_optimized=None, + ) + genotype_fitness, _descriptors, genotype_metrics, _emitter_info, _random_key, _genotype = scoring_function( + genotype, + bruteforce_optimizer.jax.random.PRNGKey(0), + dynamic_informations=(optimizer_data.runtime_state.dynamic_information,), + solver_configs=(optimizer_data.runtime_state.solver_config,), + target_metrics=start_opt_command.dc_params.ga_config.target_metrics, + observed_metrics=start_opt_command.dc_params.ga_config.observed_metrics, + descriptor_metrics=tuple(desc.metric for desc in start_opt_command.dc_params.ga_config.me_descriptors), + n_worst_contingencies=start_opt_command.dc_params.ga_config.n_worst_contingencies, + ) + + assert np.isclose(float(np.asarray(chunk_fitness[0])), float(np.asarray(genotype_fitness[0]))) + for metric_name in start_opt_command.dc_params.ga_config.observed_metrics: + assert np.isclose( + float(np.asarray(chunk_metrics[metric_name][0])), + float(np.asarray(genotype_metrics[metric_name][0])), + ) + np.testing.assert_array_equal( + np.asarray(chunk_metrics["case_indices"][0]), np.asarray(genotype_metrics["case_indices"][0]) + ) From 9ad5d88d693994fc4965167a7fda5993e95d3191 Mon Sep 17 00:00:00 2001 From: Nico Westerbeck Date: Sun, 5 Jul 2026 21:01:53 +0000 Subject: [PATCH 3/7] chore: Add license headers Signed-off-by: Nico Westerbeck --- .../dc_bruteforce/__init__.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc_bruteforce/__init__.py b/packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc_bruteforce/__init__.py index e15459408..e634a1204 100644 --- a/packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc_bruteforce/__init__.py +++ b/packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc_bruteforce/__init__.py @@ -1 +1,8 @@ +# Copyright 2026 50Hertz Transmission GmbH and Elia Transmission Belgium SA/NV +# +# This Source Code Form is subject to the terms of the Mozilla Public License, v. 2.0. +# If a copy of the MPL was not distributed with this file, +# you can obtain one at https://mozilla.org/MPL/2.0/. +# Mozilla Public License, version 2.0 + """Bruteforce DC optimizer components.""" From e459b9fb1e098c89d1efaddb2768bb9bc52b051a Mon Sep 17 00:00:00 2001 From: Nico Westerbeck Date: Sun, 5 Jul 2026 21:03:37 +0000 Subject: [PATCH 4/7] refactor: Rename tests for uniqueness Signed-off-by: Nico Westerbeck --- .../dc_bruteforce/{test_generator.py => test_brute_generator.py} | 0 .../dc_bruteforce/{test_optimizer.py => test_brute_optimizer.py} | 0 .../tests/dc_bruteforce/{test_worker.py => test_brute_worker.py} | 0 3 files changed, 0 insertions(+), 0 deletions(-) rename packages/topology_optimizer_pkg/tests/dc_bruteforce/{test_generator.py => test_brute_generator.py} (100%) rename packages/topology_optimizer_pkg/tests/dc_bruteforce/{test_optimizer.py => test_brute_optimizer.py} (100%) rename packages/topology_optimizer_pkg/tests/dc_bruteforce/{test_worker.py => test_brute_worker.py} (100%) diff --git a/packages/topology_optimizer_pkg/tests/dc_bruteforce/test_generator.py b/packages/topology_optimizer_pkg/tests/dc_bruteforce/test_brute_generator.py similarity index 100% rename from packages/topology_optimizer_pkg/tests/dc_bruteforce/test_generator.py rename to packages/topology_optimizer_pkg/tests/dc_bruteforce/test_brute_generator.py diff --git a/packages/topology_optimizer_pkg/tests/dc_bruteforce/test_optimizer.py b/packages/topology_optimizer_pkg/tests/dc_bruteforce/test_brute_optimizer.py similarity index 100% rename from packages/topology_optimizer_pkg/tests/dc_bruteforce/test_optimizer.py rename to packages/topology_optimizer_pkg/tests/dc_bruteforce/test_brute_optimizer.py diff --git a/packages/topology_optimizer_pkg/tests/dc_bruteforce/test_worker.py b/packages/topology_optimizer_pkg/tests/dc_bruteforce/test_brute_worker.py similarity index 100% rename from packages/topology_optimizer_pkg/tests/dc_bruteforce/test_worker.py rename to packages/topology_optimizer_pkg/tests/dc_bruteforce/test_brute_worker.py From aa7f6ca2aad4feda6b0efc3d120b7bd8931fd79c Mon Sep 17 00:00:00 2001 From: Nico Westerbeck Date: Sun, 5 Jul 2026 21:44:17 +0000 Subject: [PATCH 5/7] refactor: Remove some unnecessary logic Signed-off-by: Nico Westerbeck --- .../dc_bruteforce/__init__.py | 7 ++++- .../dc_bruteforce/optimizer.py | 29 ------------------- .../dc_bruteforce/worker.py | 4 +-- .../dc_bruteforce/test_brute_optimizer.py | 4 +-- 4 files changed, 9 insertions(+), 35 deletions(-) diff --git a/packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc_bruteforce/__init__.py b/packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc_bruteforce/__init__.py index e634a1204..adb18de77 100644 --- a/packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc_bruteforce/__init__.py +++ b/packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc_bruteforce/__init__.py @@ -5,4 +5,9 @@ # you can obtain one at https://mozilla.org/MPL/2.0/. # Mozilla Public License, version 2.0 -"""Bruteforce DC optimizer components.""" +"""Bruteforce DC optimizer components. + +The idea of the bruteforce optimizer is to walk through all combinations of the action set in dc exhaustively, or at least +up to the part that is doable within the runtime. The topologies are sent the same way to kafka for AC evaluation. The +optimizer will print on command line how many topologies are still remaining. +""" diff --git a/packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc_bruteforce/optimizer.py b/packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc_bruteforce/optimizer.py index 4746dc529..5faa3ffed 100644 --- a/packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc_bruteforce/optimizer.py +++ b/packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc_bruteforce/optimizer.py @@ -46,7 +46,6 @@ ) from toop_engine_topology_optimizer.interfaces.messages.dc_params import DCOptimizerParameters from toop_engine_topology_optimizer.interfaces.messages.results import Metrics, Strategy, Topology, TopologyPushResult -from toop_engine_topology_optimizer.interfaces.models.base_storage import hash_strategy @dataclass @@ -97,9 +96,6 @@ class OptimizerData: start_time: float """Wall-clock timestamp when the bruteforce run started.""" - sent_topologies: set[bytes] - """Hashes of topologies already emitted to avoid duplicate result messages.""" - latest_topologies: list[Topology] = field(default_factory=list) """Improved topologies produced by the most recent epoch.""" @@ -166,7 +162,6 @@ def initialize_optimization( initial_metrics=initial_metrics, runtime_state=runtime_state, start_time=time.time(), - sent_topologies=set(), ), static_information_descriptions, initial_strategy, @@ -305,30 +300,6 @@ def _get_topologies_per_epoch(params: DCOptimizerParameters) -> int: return params.ga_config.iterations_per_epoch * params.loadflow_solver_config.batch_size -def extract_topologies(optimizer_data: OptimizerData) -> list[Topology]: - """Return deduplicated topologies discovered in the latest epoch. - - Parameters - ---------- - optimizer_data : OptimizerData - Current state of the bruteforce optimization. - - Returns - ------- - list[Topology] - Newly discovered topologies that have not yet been emitted. - """ - new_topologies = [] - new_hashes = [] - for topology in optimizer_data.latest_topologies: - topology_hash = hash_strategy(Strategy(timesteps=[topology])) - if topology_hash not in optimizer_data.sent_topologies: - new_topologies.append(topology) - new_hashes.append(topology_hash) - optimizer_data.sent_topologies.update(new_hashes) - return new_topologies - - def convert_topologies_to_messages(topologies: list[Topology], epoch: int) -> list[TopologyPushResult]: """Convert topologies to result messages. diff --git a/packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc_bruteforce/worker.py b/packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc_bruteforce/worker.py index 65bdf40f3..ec92cf15c 100644 --- a/packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc_bruteforce/worker.py +++ b/packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc_bruteforce/worker.py @@ -23,7 +23,6 @@ from toop_engine_topology_optimizer.dc_bruteforce.optimizer import ( OptimizerData, convert_topologies_to_messages, - extract_topologies, get_num_branch_topologies_tried, initialize_optimization, is_exhausted, @@ -87,7 +86,7 @@ def push_topologies(optimizer_data: OptimizerData, epoch: int, send_result_fn: C The number of topology messages queued for emission. """ with jax.default_device(jax.devices("cpu")[0]): - push_results = convert_topologies_to_messages(extract_topologies(optimizer_data), epoch) + push_results = convert_topologies_to_messages(optimizer_data.latest_topologies, epoch) for push_result in push_results: send_result_fn(push_result) return len(push_results) @@ -134,6 +133,7 @@ def optimization_loop( ) send_result_fn(OptimizationStartedResult(initial_topology=initial_strategy, initial_stats=stats)) flush_result_fn() + logger.info(f"Starting optimization with initial fitness {initial_strategy.timesteps[0].metrics.fitness}") except Exception as exc: send_result_fn(OptimizationStoppedResult(reason="error", message=str(exc))) flush_result_fn() diff --git a/packages/topology_optimizer_pkg/tests/dc_bruteforce/test_brute_optimizer.py b/packages/topology_optimizer_pkg/tests/dc_bruteforce/test_brute_optimizer.py index c45af7250..e48fa9c05 100644 --- a/packages/topology_optimizer_pkg/tests/dc_bruteforce/test_brute_optimizer.py +++ b/packages/topology_optimizer_pkg/tests/dc_bruteforce/test_brute_optimizer.py @@ -14,7 +14,6 @@ from toop_engine_topology_optimizer.dc_bruteforce.optimizer import ( OptimizerData, convert_topologies_to_messages, - extract_topologies, get_num_branch_topologies_tried, initialize_optimization, is_exhausted, @@ -80,7 +79,7 @@ def test_initialize_and_run_epoch(grid_folder: str) -> None: topologies_per_epoch, optimizer_data.runtime_state.total_workset_size ) - topologies = extract_topologies(optimizer_data) + topologies = optimizer_data.latest_topologies assert isinstance(topologies, list) if topologies: assert isinstance(topologies[0], Topology) @@ -91,7 +90,6 @@ def test_initialize_and_run_epoch(grid_folder: str) -> None: assert isinstance(messages[0].strategy, Strategy) assert len(messages[0].strategy.timesteps) == 1 - assert extract_topologies(optimizer_data) == [] assert is_exhausted(optimizer_data) is (optimizer_data.runtime_state.total_workset_size <= topologies_per_epoch) From d7d053fe643e6d121ff89979141874aa9bc687a6 Mon Sep 17 00:00:00 2001 From: Nico Westerbeck Date: Tue, 7 Jul 2026 07:24:56 +0000 Subject: [PATCH 6/7] fix: Fix imports Signed-off-by: Nico Westerbeck --- .../tests/dc/worker/test_dc_optimizer_worker.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/packages/topology_optimizer_pkg/tests/dc/worker/test_dc_optimizer_worker.py b/packages/topology_optimizer_pkg/tests/dc/worker/test_dc_optimizer_worker.py index 7820b99e7..c7fcc26d1 100644 --- a/packages/topology_optimizer_pkg/tests/dc/worker/test_dc_optimizer_worker.py +++ b/packages/topology_optimizer_pkg/tests/dc/worker/test_dc_optimizer_worker.py @@ -19,6 +19,8 @@ from toop_engine_topology_optimizer.interfaces.messages.commands import ( Command, DCOptimizerParameters, + ShutdownCommand, + StartOptimizationCommand, ) from toop_engine_topology_optimizer.interfaces.messages.commons import Framework, GridFile from toop_engine_topology_optimizer.interfaces.messages.dc_params import BatchedMEParameters, LoadflowSolverParameters From 27f9439f6f597b834de6c9ea1afbddd52060267b Mon Sep 17 00:00:00 2001 From: Nico Westerbeck Date: Wed, 8 Jul 2026 09:17:31 +0000 Subject: [PATCH 7/7] chore: Improve logging Signed-off-by: Nico Westerbeck --- .../src/toop_engine_topology_optimizer/dc_bruteforce/worker.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc_bruteforce/worker.py b/packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc_bruteforce/worker.py index ec92cf15c..cf66bad5b 100644 --- a/packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc_bruteforce/worker.py +++ b/packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc_bruteforce/worker.py @@ -172,11 +172,13 @@ def optimization_loop( ) if is_exhausted(optimizer_data): + logger.info(f"Ending bruteforce search due to workset exhaustion after {time.time() - start_time}s") send_result_fn(OptimizationStoppedResult(epoch=epoch, reason="converged", message="workset exhausted")) flush_result_fn() return if time.time() - start_time > dc_params.ga_config.runtime_seconds: + logger.info(f"Ending due to runtime limit after {time.time() - start_time}s") send_result_fn(OptimizationStoppedResult(epoch=epoch, reason="converged", message="runtime limit")) flush_result_fn() return