stock_discretization

AntaresSimulatorTeam · Jan 8, 2025 · 1748b26 · 1748b26
1 parent 40a4387
commit 1748b26
Show file tree

Hide file tree

Showing 9 changed files with 40 additions and 29 deletions.
diff --git a/src/estimation.py b/src/estimation.py
@@ -99,7 +99,7 @@ def get_value(self, x: Dict[str, float]) -> float:
                     [
                         self.V[f"slope_{area}"][idx]
                         * (
-                            x[area]
+                            x[area.area]
                             - self.discretization.list_discretization[area][idx[i]]
                         )
                         for i, area in enumerate(

diff --git a/src/functions_iterative.py b/src/functions_iterative.py
@@ -13,7 +13,16 @@
 from read_antares_data import TimeScenarioIndex, TimeScenarioParameter
 from reservoir_management import MultiStockManagement
 from stock_discretization import StockDiscretization
-from type_definition import Array1D, Array2D, Array3D, Array4D, Dict, List, Optional
+from type_definition import (
+    AreaIndex,
+    Array1D,
+    Array2D,
+    Array3D,
+    Array4D,
+    Dict,
+    List,
+    Optional,
+)
 
 
 def compute_x_multi_scenario(
@@ -330,7 +339,7 @@ def itr_control(
                 for week in range(param.len_week + 1)
             },
             stock_discretization=StockDiscretization(
-                {reservoir_management.reservoir.area: X}
+                {AreaIndex(reservoir_management.reservoir.area): X}
             ),
             reward_approximation={reservoir_management.reservoir.area: G},
         )

diff --git a/src/optimization.py b/src/optimization.py
@@ -774,9 +774,9 @@ def build_multivariate_bellman_constraints(
                             - sum(
                                 [
                                     V[f"slope_{area}"][idx]
-                                    * stock_discretization.list_discretization[
-                                        area.area
-                                    ][idx[i]]
+                                    * stock_discretization.list_discretization[area][
+                                        idx[i]
+                                    ]
                                     for i, area in enumerate(self.range_reservoir)
                                 ]
                             )
@@ -804,9 +804,9 @@ def build_multivariate_bellman_constraints(
                                     self.stored_variables_and_constraints[area.area][
                                         "final_level"
                                     ]
-                                    - stock_discretization.list_discretization[
-                                        area.area
-                                    ][idx[i]]
+                                    - stock_discretization.list_discretization[area][
+                                        idx[i]
+                                    ]
                                 )
                                 for i, area in enumerate(self.range_reservoir)
                             ]

diff --git a/src/simple_bellman_value_calculation.py b/src/simple_bellman_value_calculation.py
@@ -17,7 +17,7 @@
 from read_antares_data import TimeScenarioIndex, TimeScenarioParameter
 from reservoir_management import MultiStockManagement
 from stock_discretization import StockDiscretization
-from type_definition import Array1D, Array2D, Dict, List
+from type_definition import AreaIndex, Array1D, Array2D, Dict, List
 
 
 def calculate_complete_reward(
@@ -131,7 +131,7 @@ def calculate_bellman_value_with_precalculated_reward(
     upper_bound, control_ub, current_itr = compute_upper_bound(
         multi_stock_management=multi_stock_management,
         stock_discretization=StockDiscretization(
-            {reservoir_management.reservoir.area: X}
+            {AreaIndex(reservoir_management.reservoir.area): X}
         ),
         param=param,
         list_models=list_models,
@@ -197,7 +197,7 @@ def calculate_bellman_value_directly(
             )
             list_models[TimeScenarioIndex(week, scenario)] = m
 
-    stock_discretization = StockDiscretization(X)
+    stock_discretization = StockDiscretization({AreaIndex(a): x for a, x in X.items()})
 
     V: Dict[int, Estimator] = {}
     if univariate:
@@ -240,7 +240,7 @@ def calculate_bellman_value_directly(
                 _, _, Vu, slope, _, _, _ = m.solve_problem_with_bellman_values(
                     V=V[week + 1],
                     level_i={
-                        area.area: stock_discretization.list_discretization[area.area][
+                        area.area: stock_discretization.list_discretization[area][
                             idx[i]
                         ]
                         for i, area in enumerate(m.range_reservoir)

diff --git a/src/stock_discretization.py b/src/stock_discretization.py
@@ -1,11 +1,11 @@
 from itertools import product
 
-from type_definition import Array1D, Dict, Iterable
+from type_definition import AreaIndex, Array1D, Dict, Iterable
 
 
 class StockDiscretization:
 
-    def __init__(self, list_discretization: Dict[str, Array1D]) -> None:
+    def __init__(self, list_discretization: Dict[AreaIndex, Array1D]) -> None:
         self.list_discretization = list_discretization
 
     def get_product_stock_discretization(self) -> Iterable:

diff --git a/tests/test_basis.py b/tests/test_basis.py
@@ -13,6 +13,7 @@
 from read_antares_data import Reservoir
 from reservoir_management import MultiStockManagement
 from stock_discretization import StockDiscretization
+from type_definition import AreaIndex
 
 
 def test_basis_with_xpress() -> None:
@@ -103,7 +104,7 @@ def test_basis_with_upper_bound() -> None:
         upper_bound_1, _, _ = compute_upper_bound(
             param=param,
             multi_stock_management=MultiStockManagement([reservoir_management]),
-            stock_discretization=StockDiscretization({reservoir.area: X}),
+            stock_discretization=StockDiscretization({AreaIndex(reservoir.area): X}),
             list_models=list_models,
             V={
                 week: UniVariateEstimator({reservoir.area: V[week]})
@@ -118,7 +119,7 @@ def test_basis_with_upper_bound() -> None:
         upper_bound_2, _, itr_with_basis = compute_upper_bound(
             param=param,
             multi_stock_management=MultiStockManagement([reservoir_management]),
-            stock_discretization=StockDiscretization({reservoir.area: X}),
+            stock_discretization=StockDiscretization({AreaIndex(reservoir.area): X}),
             list_models=list_models,
             V={
                 week: UniVariateEstimator({reservoir.area: V[week]})

diff --git a/tests/test_bellman_value_exact_two_stocks.py b/tests/test_bellman_value_exact_two_stocks.py
@@ -10,6 +10,7 @@
 from reservoir_management import MultiStockManagement
 from simple_bellman_value_calculation import calculate_bellman_value_directly
 from stock_discretization import StockDiscretization
+from type_definition import AreaIndex
 
 
 def test_iterate_over_stock_discretization() -> None:
@@ -22,7 +23,7 @@ def test_iterate_over_stock_discretization() -> None:
     x_2 = np.linspace(0, reservoir_2.capacity, num=20)
 
     stock_discretization = StockDiscretization(
-        {reservoir_1.area: x_1, reservoir_2.area: x_2}
+        {AreaIndex(reservoir_1.area): x_1, AreaIndex(reservoir_2.area): x_2}
     )
 
     assert [idx for idx in stock_discretization.get_product_stock_discretization()][
@@ -61,7 +62,7 @@ def test_solve_with_bellman_multi_stock() -> None:
 
     x_1 = np.linspace(0, reservoir_1.capacity, num=5)
     x_2 = np.linspace(0, reservoir_2.capacity, num=5)
-    X = {"area_1": x_1, "area_2": x_2}
+    X = {AreaIndex("area_1"): x_1, AreaIndex("area_2"): x_2}
 
     m = AntaresProblem(
         scenario=0,
@@ -154,9 +155,7 @@ def test_solve_with_bellman_multi_stock() -> None:
 
     _, _, Vu, slope, _, xf, _ = m.solve_problem_with_bellman_values(
         multi_stock_management=multi_stock_management,
-        stock_discretization=StockDiscretization(
-            {area.area: X[area.area] for area in m.range_reservoir}
-        ),
+        stock_discretization=StockDiscretization(X),
         V=BellmanValueEstimation(V, StockDiscretization(X)),
         level_i={
             area.area: multi_stock_management.dict_reservoirs[

diff --git a/tests/test_read_optimization_problem.py b/tests/test_read_optimization_problem.py
@@ -8,6 +8,7 @@
 from read_antares_data import Reservoir
 from reservoir_management import MultiStockManagement
 from stock_discretization import StockDiscretization
+from type_definition import AreaIndex
 
 
 def test_create_and_modify_weekly_problem() -> None:
@@ -147,7 +148,7 @@ def test_create_and_modify_weekly_problem_with_bellman_values() -> None:
     problem.set_constraints_initial_level_and_bellman_values(
         UniVariateEstimator({"area": V[1]}),
         {"area": reservoir.initial_level},
-        StockDiscretization({"area": X}),
+        StockDiscretization({AreaIndex("area"): X}),
     )
 
     lp = problem.solver.ExportModelAsLpFormat(False)
@@ -157,7 +158,7 @@ def test_create_and_modify_weekly_problem_with_bellman_values() -> None:
 
     _, _, cout, _, optimal_controls, _, _ = problem.solve_problem_with_bellman_values(
         reservoir_management,
-        StockDiscretization({"area": X}),
+        StockDiscretization({AreaIndex("area"): X}),
         UniVariateEstimator({"area": V[1]}),
         {"area": reservoir.initial_level},
         True,

diff --git a/tests/test_upper_bound.py b/tests/test_upper_bound.py
@@ -13,6 +13,7 @@
 from read_antares_data import Reservoir
 from reservoir_management import MultiStockManagement
 from stock_discretization import StockDiscretization
+from type_definition import AreaIndex
 
 bellman_values = np.array(
     [
@@ -208,7 +209,7 @@ def test_upper_bound() -> None:
     upper_bound, controls, _ = compute_upper_bound(
         param=param,
         multi_stock_management=MultiStockManagement([reservoir_management]),
-        stock_discretization=StockDiscretization({reservoir.area: X}),
+        stock_discretization=StockDiscretization({AreaIndex(reservoir.area): X}),
         list_models=list_models,
         V={
             week: UniVariateEstimator({reservoir.area: V[week]})
@@ -225,7 +226,7 @@ def test_upper_bound() -> None:
     upper_bound, controls, _ = compute_upper_bound(
         param=param,
         multi_stock_management=MultiStockManagement([reservoir_management]),
-        stock_discretization=StockDiscretization({reservoir.area: X}),
+        stock_discretization=StockDiscretization({AreaIndex(reservoir.area): X}),
         list_models=list_models,
         V={
             week: UniVariateEstimator({reservoir.area: V[week]})
@@ -268,7 +269,7 @@ def test_upper_bound_with_bellman_values() -> None:
     upper_bound, controls, _ = compute_upper_bound(
         param=param,
         multi_stock_management=MultiStockManagement([reservoir_management]),
-        stock_discretization=StockDiscretization({reservoir.area: X}),
+        stock_discretization=StockDiscretization({AreaIndex(reservoir.area): X}),
         list_models=list_models,
         V={
             week: UniVariateEstimator({reservoir.area: V[week]})
@@ -320,7 +321,7 @@ def test_upper_bound_with_xpress() -> None:
         upper_bound, controls, _ = compute_upper_bound(
             param=param,
             multi_stock_management=MultiStockManagement([reservoir_management]),
-            stock_discretization=StockDiscretization({reservoir.area: X}),
+            stock_discretization=StockDiscretization({AreaIndex(reservoir.area): X}),
             list_models=list_models,
             V={
                 week: UniVariateEstimator({reservoir.area: V[week]})
@@ -339,7 +340,7 @@ def test_upper_bound_with_xpress() -> None:
         upper_bound, controls, _ = compute_upper_bound(
             param=param,
             multi_stock_management=MultiStockManagement([reservoir_management]),
-            stock_discretization=StockDiscretization({reservoir.area: X}),
+            stock_discretization=StockDiscretization({AreaIndex(reservoir.area): X}),
             list_models=list_models,
             V={
                 week: UniVariateEstimator({reservoir.area: V[week]})