diff --git a/src/estimation.py b/src/estimation.py index a70cd18..e345a11 100644 --- 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 index ef94b55..41ace8c 100644 --- 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 index c82cbe0..96d4325 100644 --- 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 index 382624a..c61a63f 100644 --- 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 index a406f41..48d3112 100644 --- 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 index 3773a7f..2c2c2fe 100644 --- 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 index 9e94149..a93ae24 100644 --- 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 index bf1c89d..bfd8a3b 100644 --- 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 index a08a6ab..abf6b69 100644 --- 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]})