feat: Simple bruteforce worker for full action set evaluation

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
 

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()

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,25 +18,20 @@
 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,
     extract_topologies,
     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.
 

packages/topology_optimizer_pkg/src/toop_engine_topology_optimizer/dc_bruteforce/__init__.py

@@ -0,0 +1,13 @@
+# 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.
+
+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.
+"""