Refactor compute target

src/andromede/hydro_heuristic/data.py

@@ -10,111 +10,133 @@
 #
 # This file is part of the Antares project.
 
+from dataclasses import dataclass
 from pathlib import Path
 from typing import List, Optional
 
 import numpy as np
 
 
+@dataclass
+class ReservoirParameters:
+    capacity: float
+    initial_level: float
+    folder_name: str
+    scenario: int
+
+
+@dataclass
+class HydroHeuristicParameters:
+    inter_breakdown: int = 1
+    total_target: Optional[float] = None
+
+
+@dataclass
+class DataAggregatorParameters:
+    hours_aggregated_time_steps: List[int]
+    timesteps: List[int]
+
+
 def get_number_of_days_in_month(month: int) -> int:
     number_day_month = [31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31][month]
     return number_day_month
 
 
 class RawHydroData:
+    def __init__(self, folder_name: str, scenario: int) -> None:
+        self.folder_name = folder_name
+        self.scenario = scenario
 
-    def __init__(self, name: str, folder_name: str, scenario: int) -> None:
-
-        name_file = {
+        self.name_file = {
             "demand": "load",
             "inflow": "mod",
             "lower_rule_curve": "reservoir",
             "upper_rule_curve": "reservoir",
             "max_generating": "maxpower",
-        }[name]
+        }
 
-        column = {
-            "demand": scenario,
-            "inflow": scenario,
+        self.column = {
+            "demand": self.scenario,
+            "inflow": self.scenario,
             "lower_rule_curve": 0,
             "upper_rule_curve": 2,
             "max_generating": 0,
-        }[name]
+        }
 
+    def read_data(self, name: str) -> list[float]:
         hours_input = 1 if name == "demand" else 24
 
         data = np.loadtxt(
             Path(__file__).parent
             / (
                 "../../../tests/functional/data/"
-                + folder_name
+                + self.folder_name
                 + "/"
-                + name_file
+                + self.name_file[name]
                 + ".txt"
             )
         )
-        data = data[:, column]
+        data = data[:, self.column[name]]
         data = np.repeat(data, hours_input)
-        if name_file == "mod":
+        if self.name_file[name] == "mod":
             data = data / hours_input
 
-        self.time_series = data
+        return list(data)
 
 
 class HydroHeuristicData:
     def __init__(
         self,
-        scenario: int,
-        hours_aggregated_time_steps: List[int],
-        folder_name: str,
-        timesteps: List[int],
-        capacity: float,
-        initial_level: float,
+        data_aggregator_parameters: DataAggregatorParameters,
+        reservoir_data: ReservoirParameters,
     ):
-        self.folder_name = folder_name
+        self.reservoir_data = reservoir_data
 
-        self.capacity = capacity
-        self.initial_level = initial_level
+        data_aggregator = DataAggregator(
+            data_aggregator_parameters.hours_aggregated_time_steps,
+            data_aggregator_parameters.timesteps,
+        )
 
-        data_aggregator = DataAggregator(hours_aggregated_time_steps, timesteps)
+        raw_data_reader = RawHydroData(
+            reservoir_data.folder_name, reservoir_data.scenario
+        )
 
         self.demand = data_aggregator.aggregate_data(
             operator="sum",
-            raw_data=RawHydroData("demand", folder_name, scenario),
+            data=raw_data_reader.read_data("demand"),
         )
         self.inflow = data_aggregator.aggregate_data(
             operator="sum",
-            raw_data=RawHydroData("inflow", folder_name, scenario),
+            data=raw_data_reader.read_data("inflow"),
         )
         self.lower_rule_curve = data_aggregator.aggregate_data(
             operator="lag_first_element",
-            raw_data=RawHydroData("lower_rule_curve", folder_name, scenario),
+            data=raw_data_reader.read_data("lower_rule_curve"),
         )
         self.upper_rule_curve = data_aggregator.aggregate_data(
             operator="lag_first_element",
-            raw_data=RawHydroData("upper_rule_curve", folder_name, scenario),
+            data=raw_data_reader.read_data("upper_rule_curve"),
         )
         self.max_generating = data_aggregator.aggregate_data(
             operator="sum",
-            raw_data=RawHydroData("max_generating", folder_name, scenario),
+            data=raw_data_reader.read_data("max_generating"),
         )
 
-    def compute_target(  # TODO : rajouter un test avec vraiment très peu de données
-        self, total_target: Optional[float], inter_breakdown: int = 1
-    ) -> None:
-        if total_target is None:
+    def compute_target(self, heuristic_parameters: HydroHeuristicParameters) -> None:
+        if heuristic_parameters.total_target is None:
             total_target = sum(self.inflow)
+        else:
+            total_target = heuristic_parameters.total_target
         target = (
             total_target
-            * np.power(self.demand, inter_breakdown)
-            / sum(np.power(self.demand, inter_breakdown))
+            * np.power(self.demand, heuristic_parameters.inter_breakdown)
+            / sum(np.power(self.demand, heuristic_parameters.inter_breakdown))
         )
 
         self.target = list(target)
 
 
 class DataAggregator:
-
     def __init__(
         self,
         hours_aggregated_time_steps: List[int],
@@ -123,8 +145,7 @@ def __init__(
         self.hours_aggregated_time_steps = hours_aggregated_time_steps
         self.timesteps = timesteps
 
-    def aggregate_data(self, operator: str, raw_data: RawHydroData) -> List[float]:
-        data = raw_data.time_series
+    def aggregate_data(self, operator: str, data: list[float]) -> List[float]:
         aggregated_data: List[float] = []
         hour = 0
         for time_step, hours_time_step in enumerate(self.hours_aggregated_time_steps):

src/andromede/hydro_heuristic/heuristic_model.py

@@ -27,7 +27,6 @@
 
 
 class HeuristicHydroModelBuilder:
-
     def __init__(
         self,
         hydro_model: Model,

src/andromede/hydro_heuristic/problem.py

@@ -10,12 +10,19 @@
 #
 # This file is part of the Antares project.
 
-from typing import Optional, List
+from dataclasses import dataclass
+from typing import List, Optional
 
 import ortools.linear_solver.pywraplp as pywraplp
 import pandas as pd
 
-from andromede.hydro_heuristic.data import HydroHeuristicData
+from andromede.hydro_heuristic.data import (
+    DataAggregatorParameters,
+    HydroHeuristicData,
+    HydroHeuristicParameters,
+    ReservoirParameters,
+)
+from andromede.model import Model
 from andromede.simulation import (
     BlockBorderManagement,
     OutputValues,
@@ -31,7 +38,18 @@
     TimeSeriesData,
     create_component,
 )
-from andromede.model import Model
+
+
+@dataclass
+class SolvingOutput:
+    status: str
+    objective: float
+
+
+@dataclass
+class OutputHeuristic:
+    generating: list[float]
+    level: float
 
 
 class HydroHeuristicProblem:
@@ -40,9 +58,6 @@ def __init__(self, hydro_data: HydroHeuristicData, heuristic_model: Model) -> No
         self.id = "H"
         database = self.generate_database()
 
-        time_block = TimeBlock(1, [i for i in range(len(hydro_data.target))])
-        scenarios = 1
-
         hydro = create_component(
             model=heuristic_model,
             id=self.id,
@@ -54,14 +69,14 @@ def __init__(self, hydro_data: HydroHeuristicData, heuristic_model: Model) -> No
         problem = build_problem(
             network,
             database,
-            time_block,
-            scenarios,
+            TimeBlock(1, [i for i in range(len(hydro_data.target))]),
+            1,
             border_management=(BlockBorderManagement.CYCLE),
         )
 
         self.problem = problem
 
-    def solve_hydro_problem(self) -> tuple[int, float, list[float], float]:
+    def solve_hydro_problem(self) -> tuple[SolvingOutput, OutputHeuristic]:
         parameters = pywraplp.MPSolverParameters()
         parameters.SetIntegerParam(parameters.PRESOLVE, parameters.PRESOLVE_OFF)
         parameters.SetIntegerParam(parameters.SCALING, 0)
@@ -71,20 +86,25 @@ def solve_hydro_problem(self) -> tuple[int, float, list[float], float]:
         output = OutputValues(self.problem)
 
         return (
-            status,
-            self.problem.solver.Objective().Value(),
-            output.component(self.id).var("generating").value[0],  # type:ignore
-            output.component(self.id).var("level").value[0][-1],  # type:ignore
+            SolvingOutput(status, self.problem.solver.Objective().Value()),
+            OutputHeuristic(
+                output.component(self.id).var("generating").value[0],  # type:ignore
+                output.component(self.id).var("level").value[0][-1],  # type:ignore
+            ),
         )
 
     def generate_database(
         self,
     ) -> DataBase:
         database = DataBase()
 
-        database.add_data(self.id, "capacity", ConstantData(self.hydro_data.capacity))
         database.add_data(
-            self.id, "initial_level", ConstantData(self.hydro_data.initial_level)
+            self.id, "capacity", ConstantData(self.hydro_data.reservoir_data.capacity)
+        )
+        database.add_data(
+            self.id,
+            "initial_level",
+            ConstantData(self.hydro_data.reservoir_data.initial_level),
         )
 
         inflow_data = pd.DataFrame(
@@ -112,7 +132,7 @@ def generate_database(
             TimeSeriesData(
                 {
                     TimeIndex(i): self.hydro_data.lower_rule_curve[i]
-                    * self.hydro_data.capacity
+                    * self.hydro_data.reservoir_data.capacity
                     for i in range(len(self.hydro_data.lower_rule_curve))
                 }
             ),
@@ -123,7 +143,7 @@ def generate_database(
             TimeSeriesData(
                 {
                     TimeIndex(i): self.hydro_data.upper_rule_curve[i]
-                    * self.hydro_data.capacity
+                    * self.hydro_data.reservoir_data.capacity
                     for i in range(len(self.hydro_data.lower_rule_curve))
                 }
             ),
@@ -146,7 +166,9 @@ def generate_database(
             "max_epsilon",
             TimeSeriesData(
                 {
-                    TimeIndex(i): self.hydro_data.capacity if i == 0 else 0
+                    TimeIndex(i): (
+                        self.hydro_data.reservoir_data.capacity if i == 0 else 0
+                    )
                     for i in range(len(self.hydro_data.max_generating))
                 }
             ),
@@ -156,33 +178,24 @@ def generate_database(
 
 
 def optimize_target(
-    inter_breakdown: int,
-    folder_name: str,
-    capacity: float,
-    scenario: int,
-    initial_level: float,
-    hours_aggregated_time_steps: List[int],
-    timesteps: list[int],
-    total_target: Optional[float],
+    heuristic_parameters: HydroHeuristicParameters,
+    data_aggregator_parameters: DataAggregatorParameters,
+    reservoir_data: ReservoirParameters,
     heuristic_model: Model,
-) -> tuple[float, int, float, list[float]]:
+) -> tuple[SolvingOutput, OutputHeuristic]:
     # Récupération des données
     data = HydroHeuristicData(
-        scenario,
-        hours_aggregated_time_steps=hours_aggregated_time_steps,
-        folder_name=folder_name,
-        timesteps=timesteps,
-        capacity=capacity,
-        initial_level=initial_level,
+        data_aggregator_parameters,
+        reservoir_data,
     )
     # Calcul de la préallocation
-    data.compute_target(total_target, inter_breakdown=inter_breakdown)
+    data.compute_target(heuristic_parameters)
 
     # Ajustement de la réapartition
     heuristic_problem = HydroHeuristicProblem(
         hydro_data=data, heuristic_model=heuristic_model
     )
 
-    status, obj, generation, initial_level = heuristic_problem.solve_hydro_problem()
+    solving_output, heuristic_output = heuristic_problem.solve_hydro_problem()
 
-    return initial_level, status, obj, generation
+    return solving_output, heuristic_output