diff --git a/src/functions_iterative.py b/src/functions_iterative.py
index 41ace8c..404d0a5 100644
--- a/src/functions_iterative.py
+++ b/src/functions_iterative.py
@@ -315,7 +315,7 @@ def itr_control(
             list_models=list_models,
             G=G,
             i=i,
-            name_reservoir=reservoir_management.reservoir.area,
+            name_reservoir=reservoir_management.reservoir.area.area,
         )
         itr_tot.append(current_itr)
 
@@ -334,14 +334,14 @@ def itr_control(
             list_models=list_models,
             V={
                 week: UniVariateEstimator(
-                    {reservoir_management.reservoir.area: V[week]}
+                    {reservoir_management.reservoir.area.area: V[week]}
                 )
                 for week in range(param.len_week + 1)
             },
             stock_discretization=StockDiscretization(
-                {AreaIndex(reservoir_management.reservoir.area): X}
+                {reservoir_management.reservoir.area: X}
             ),
-            reward_approximation={reservoir_management.reservoir.area: G},
+            reward_approximation={reservoir_management.reservoir.area.area: G},
         )
         itr_tot.append(current_itr)
         controls_upper.append(ctr)
diff --git a/src/launch_calculation.py b/src/launch_calculation.py
index 5ad5ee2..102241d 100644
--- a/src/launch_calculation.py
+++ b/src/launch_calculation.py
@@ -48,7 +48,7 @@ def calculate_bellman_values(
             param=param,
             multi_stock_management=MultiStockManagement([reservoir_management]),
             output_path=output_path,
-            X={AreaIndex(reservoir_management.reservoir.area): X},
+            X={reservoir_management.reservoir.area: X},
             solver=solver,
             univariate=True,
         )
@@ -57,7 +57,7 @@ def calculate_bellman_values(
             [
                 [
                     intermediate_vb[week].get_value(
-                        {reservoir_management.reservoir.area: x}
+                        {reservoir_management.reservoir.area.area: x}
                     )
                     for x in X
                 ]
diff --git a/src/multi_stock_bellman_value_calculation.py b/src/multi_stock_bellman_value_calculation.py
index 50efb4d..f7518b0 100644
--- a/src/multi_stock_bellman_value_calculation.py
+++ b/src/multi_stock_bellman_value_calculation.py
@@ -20,7 +20,15 @@
 )
 from read_antares_data import TimeScenarioIndex, TimeScenarioParameter
 from reservoir_management import MultiStockManagement
-from type_definition import Array1D, Dict, List, Optional
+from type_definition import (
+    Array2D,
+    Dict,
+    List,
+    Optional,
+    area_value_to_array,
+    array_to_area_value,
+    list_area_value_to_array,
+)
 
 jl = juliacall.Main
 
@@ -148,10 +156,10 @@ def get_bellman_values_from_costs(
     n_scenarios = param.len_scenario
 
     # Initializing controls, costs and duals
-    controls = []
-    costs = []
-    duals = []
-    all_levels = []
+    controls: List[np.ndarray] = []
+    costs: List[np.ndarray] = []
+    duals: List[np.ndarray] = []
+    all_levels: List[Array2D] = []
 
     # Starting from last week dynamically solving the optimal control problem (from every starting level)
     week_range = range(n_weeks - 1, -1, -1)
@@ -162,16 +170,18 @@ def get_bellman_values_from_costs(
         levels = multi_stock_management.get_disc(
             week=week,
             xNsteps=nSteps_bellman,
-            reference_pt=np.mean(trajectory[week], axis=1),
+            reference_pt=array_to_area_value(
+                np.mean(trajectory[week], axis=1), multi_stock_management.areas
+            ),
             correlation_matrix=correlations,
             method=method,
         )
-        all_levels.append(levels)
+        all_levels.append(list_area_value_to_array(levels))
 
         # Preparing receiving arrays
-        controls_w = np.zeros(list(levels.shape) + [n_scenarios])
-        costs_w = np.zeros(list(levels.shape)[:-1])
-        duals_w = np.zeros(list(levels.shape))
+        controls_w = np.zeros([len(levels), len(levels[0]), n_scenarios])
+        costs_w = np.zeros(len(levels))
+        duals_w = np.zeros([len(levels), len(levels[0])])
         for lvl_id, lvl_init in enumerate(levels):
 
             # Remove previous constraints / vars
@@ -181,7 +191,7 @@ def get_bellman_values_from_costs(
             # Rewrite problem
             problem.write_problem(
                 week=week,
-                level_init=lvl_init,
+                level_init=area_value_to_array(lvl_init),
                 future_costs_estimation=future_costs_approx,
             )
 
@@ -209,7 +219,9 @@ def get_bellman_values_from_costs(
         # Updating the future estimator
         # costs_w - np.min(costs_w)
         future_costs_approx = LinearInterpolator(
-            controls=levels, costs=costs_w - np.min(costs_w), duals=duals_w
+            controls=list_area_value_to_array(levels),
+            costs=costs_w - np.min(costs_w),
+            duals=duals_w,
         )
         future_costs_approx_l.insert(0, future_costs_approx)
     return (
diff --git a/src/read_antares_data.py b/src/read_antares_data.py
index 920b889..bd749fe 100644
--- a/src/read_antares_data.py
+++ b/src/read_antares_data.py
@@ -5,6 +5,8 @@
 
 import numpy as np
 
+from type_definition import AreaIndex
+
 
 @dataclass
 class TimeScenarioParameter:
@@ -60,7 +62,7 @@ def __init__(
         None
         """
 
-        self.area = name_area
+        self.area = AreaIndex(name_area)
 
         hydro_ini_file = self.get_hydro_ini_file(dir_study=dir_study)
 
@@ -118,12 +120,12 @@ def get_hydro_ini_file(self, dir_study: str) -> ConfigParser:
 
     def read_capacity(self, hydro_ini_file: ConfigParser) -> None:
 
-        capacity = hydro_ini_file.getfloat("reservoir capacity", self.area)
+        capacity = hydro_ini_file.getfloat("reservoir capacity", self.area.area)
 
         self.capacity = capacity
 
     def read_efficiency(self, hydro_ini_file: ConfigParser) -> None:
-        efficiency = hydro_ini_file.getfloat("pumping efficiency", self.area)
+        efficiency = hydro_ini_file.getfloat("pumping efficiency", self.area.area)
         self.efficiency = efficiency
 
 
diff --git a/src/reservoir_management.py b/src/reservoir_management.py
index 7a83919..20e7e73 100644
--- a/src/reservoir_management.py
+++ b/src/reservoir_management.py
@@ -4,7 +4,15 @@
 from scipy.interpolate import interp1d
 
 from read_antares_data import Reservoir
-from type_definition import AreaIndex, Array1D, Callable, Dict, List, Optional
+from type_definition import (
+    AreaIndex,
+    Callable,
+    Dict,
+    List,
+    Optional,
+    area_value_to_array,
+    array_to_list_area_value,
+)
 
 
 class ReservoirManagement:
@@ -104,7 +112,7 @@ def __init__(self, list_reservoirs: List[ReservoirManagement]) -> None:
         """
         self.dict_reservoirs: Dict[AreaIndex, ReservoirManagement] = {}
         for res in list_reservoirs:
-            self.dict_reservoirs[AreaIndex(res.reservoir.area)] = res
+            self.dict_reservoirs[res.reservoir.area] = res
         self.areas: List[AreaIndex] = [a for a in self.dict_reservoirs.keys()]
 
     def get_disc(
@@ -112,14 +120,12 @@ def get_disc(
         method: str,
         week: int,
         xNsteps: int,
-        reference_pt: np.ndarray,
+        reference_pt: Dict[AreaIndex, float],
         correlation_matrix: np.ndarray,
         alpha: float = 1.0,
         in_out_ratio: float = 3.0,
-    ) -> Array1D:
+    ) -> List[Dict[AreaIndex, float]]:
         n_reservoirs = len(self.dict_reservoirs)
-        if len(reference_pt.shape) > 1:
-            reference_pt = np.mean(reference_pt, axis=1)
         lbs = np.array(
             [
                 mng.reservoir.bottom_rule_curve[week] * alpha
@@ -183,12 +189,13 @@ def get_disc(
                     for r in range(n_reservoirs)
                 ]
             ).T
-            diffs_to_ref = all_pts[:, None] - reference_pt[None, :]  # Disc * R
+            array_ref_pt = area_value_to_array(reference_pt)
+            diffs_to_ref = all_pts[:, None] - array_ref_pt[None, :]  # Disc * R
             diffs_to_ref = (
                 diffs_to_ref[:, :, None] * np.eye(n_reservoirs)[None, :, :]
             )  # Disc * R * R
             diffs_to_ref = np.dot(diffs_to_ref, correlation_matrix)  # Disc * R * R
-            new_levels = reference_pt[None, None, :] + diffs_to_ref  # Disc * R * R
+            new_levels = array_ref_pt[None, None, :] + diffs_to_ref  # Disc * R * R
             new_levels = np.maximum(
                 new_levels, empty[None, None, :]
             )  # Do or do not make other points leave guiding curves ?
@@ -198,6 +205,7 @@ def get_disc(
             levels = np.reshape(
                 new_levels, (xNsteps * n_reservoirs, n_reservoirs)
             )  # (Disc * R) * R
+            dict_levels = array_to_list_area_value(levels, self.areas)
         else:
             # Listing all levels
             levels_discretization = product(
@@ -213,4 +221,4 @@ def get_disc(
                 ]
             )
             levels = np.array([level for level in levels_discretization])
-        return levels
+        return dict_levels
diff --git a/src/simple_bellman_value_calculation.py b/src/simple_bellman_value_calculation.py
index feadff0..2941de0 100644
--- a/src/simple_bellman_value_calculation.py
+++ b/src/simple_bellman_value_calculation.py
@@ -139,12 +139,14 @@ 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(
-            {AreaIndex(reservoir_management.reservoir.area): X}
+            {reservoir_management.reservoir.area: X}
         ),
         param=param,
         list_models=list_models,
         V={
-            week: UniVariateEstimator({reservoir_management.reservoir.area: V[week]})
+            week: UniVariateEstimator(
+                {reservoir_management.reservoir.area.area: V[week]}
+            )
             for week in range(param.len_week + 1)
         },
     )
diff --git a/src/type_definition.py b/src/type_definition.py
index 993c36a..6f52ed6 100644
--- a/src/type_definition.py
+++ b/src/type_definition.py
@@ -16,3 +16,26 @@ class AreaIndex:
 
     def __str__(self) -> str:
         return self.area
+
+
+def area_value_to_array(x: Dict[AreaIndex, float]) -> Array1D:
+    return np.array([y for y in x.values()])
+
+
+def list_area_value_to_array(x: List[Dict[AreaIndex, float]]) -> Array2D:
+    return np.array([[z for z in y.values()] for y in x])
+
+
+def array_to_area_value(
+    x: Array1D, list_areas: List[AreaIndex]
+) -> Dict[AreaIndex, float]:
+    assert len(x.shape) == 1
+    assert len(x) == len(list_areas)
+    return {a: x[i] for i, a in enumerate(list_areas)}
+
+
+def array_to_list_area_value(
+    x: Array2D, list_areas: List[AreaIndex]
+) -> List[Dict[AreaIndex, float]]:
+    assert len(x.shape) == 2
+    return [array_to_area_value(y, list_areas) for y in x]
diff --git a/tests/test_basis.py b/tests/test_basis.py
index 2c2c2fe..c0f9ed4 100644
--- a/tests/test_basis.py
+++ b/tests/test_basis.py
@@ -104,10 +104,10 @@ def test_basis_with_upper_bound() -> None:
         upper_bound_1, _, _ = compute_upper_bound(
             param=param,
             multi_stock_management=MultiStockManagement([reservoir_management]),
-            stock_discretization=StockDiscretization({AreaIndex(reservoir.area): X}),
+            stock_discretization=StockDiscretization({reservoir.area: X}),
             list_models=list_models,
             V={
-                week: UniVariateEstimator({reservoir.area: V[week]})
+                week: UniVariateEstimator({reservoir.area.area: V[week]})
                 for week in range(param.len_week + 1)
             },
         )
@@ -119,10 +119,10 @@ 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({AreaIndex(reservoir.area): X}),
+            stock_discretization=StockDiscretization({reservoir.area: X}),
             list_models=list_models,
             V={
-                week: UniVariateEstimator({reservoir.area: V[week]})
+                week: UniVariateEstimator({reservoir.area.area: V[week]})
                 for week in range(param.len_week + 1)
             },
         )
diff --git a/tests/test_bellman_value_exact.py b/tests/test_bellman_value_exact.py
index dc2c235..9f9f7b9 100644
--- a/tests/test_bellman_value_exact.py
+++ b/tests/test_bellman_value_exact.py
@@ -356,7 +356,7 @@ def test_bellman_value_exact() -> None:
         param=param,
         multi_stock_management=MultiStockManagement([reservoir_management]),
         output_path="test_data/one_node",
-        X={AreaIndex(reservoir.area): X},
+        X={reservoir.area: X},
         univariate=True,
     )
 
@@ -366,7 +366,7 @@ def test_bellman_value_exact() -> None:
 
     assert np.transpose(
         [
-            [vb[week].get_value({reservoir.area: x}) for x in X]
+            [vb[week].get_value({reservoir.area.area: x}) for x in X]
             for week in range(param.len_week + 1)
         ]
     ) == pytest.approx(expected_vb)
@@ -400,7 +400,7 @@ def test_bellman_value_exact_with_multi_stock() -> None:
 
     computed_vb = np.transpose(
         [
-            [vb[week].get_value({reservoir.area: x}) for x in X]
+            [vb[week].get_value({reservoir.area.area: x}) for x in X]
             for week in range(param.len_week + 1)
         ]
     )
diff --git a/tests/test_bellman_value_exact_two_stocks.py b/tests/test_bellman_value_exact_two_stocks.py
index 6c539c1..b8965d1 100644
--- a/tests/test_bellman_value_exact_two_stocks.py
+++ b/tests/test_bellman_value_exact_two_stocks.py
@@ -23,7 +23,7 @@ def test_iterate_over_stock_discretization() -> None:
     x_2 = np.linspace(0, reservoir_2.capacity, num=20)
 
     stock_discretization = StockDiscretization(
-        {AreaIndex(reservoir_1.area): x_1, AreaIndex(reservoir_2.area): x_2}
+        {reservoir_1.area: x_1, reservoir_2.area: x_2}
     )
 
     assert [idx for idx in stock_discretization.get_product_stock_discretization()][
diff --git a/tests/test_bellman_value_precalculated_reward.py b/tests/test_bellman_value_precalculated_reward.py
index 55c5f56..f5e49bb 100644
--- a/tests/test_bellman_value_precalculated_reward.py
+++ b/tests/test_bellman_value_precalculated_reward.py
@@ -226,9 +226,9 @@ def test_bellman_value_precalculated_reward() -> None:
         (-0.0004060626000000001, -38705645.55951345),
     ]
     for i, cut in enumerate(true_list_cut):
-        assert G[AreaIndex(reservoir.area)][TimeScenarioIndex(0, 0)].list_cut[
-            i
-        ] == pytest.approx(cut)
+        assert G[reservoir.area][TimeScenarioIndex(0, 0)].list_cut[i] == pytest.approx(
+            cut
+        )
 
     true_breaking_point = [
         -8400000.0,
@@ -254,7 +254,7 @@ def test_bellman_value_precalculated_reward() -> None:
         8400000.0,
     ]
     for i, pt in enumerate(true_breaking_point):
-        assert G[AreaIndex(reservoir.area)][TimeScenarioIndex(0, 0)].breaking_point[
+        assert G[reservoir.area][TimeScenarioIndex(0, 0)].breaking_point[
             i
         ] == pytest.approx(pt, 1e-5)
 
diff --git a/tests/test_upper_bound.py b/tests/test_upper_bound.py
index abf6b69..25e9a7f 100644
--- a/tests/test_upper_bound.py
+++ b/tests/test_upper_bound.py
@@ -209,16 +209,18 @@ def test_upper_bound() -> None:
     upper_bound, controls, _ = compute_upper_bound(
         param=param,
         multi_stock_management=MultiStockManagement([reservoir_management]),
-        stock_discretization=StockDiscretization({AreaIndex(reservoir.area): X}),
+        stock_discretization=StockDiscretization({reservoir.area: X}),
         list_models=list_models,
         V={
-            week: UniVariateEstimator({reservoir.area: V[week]})
+            week: UniVariateEstimator({reservoir.area.area: V[week]})
             for week in range(param.len_week + 1)
         },
     )
 
     assert upper_bound == pytest.approx(380492940.000565)
-    assert controls[TimeScenarioIndex(0, 0)][reservoir.area] == pytest.approx(4482011.0)
+    assert controls[TimeScenarioIndex(0, 0)][reservoir.area.area] == pytest.approx(
+        4482011.0
+    )
     assert len(problem.solver.constraints()) == 3555
     assert len(problem.solver.variables()) == 3533
 
@@ -226,16 +228,18 @@ def test_upper_bound() -> None:
     upper_bound, controls, _ = compute_upper_bound(
         param=param,
         multi_stock_management=MultiStockManagement([reservoir_management]),
-        stock_discretization=StockDiscretization({AreaIndex(reservoir.area): X}),
+        stock_discretization=StockDiscretization({reservoir.area: X}),
         list_models=list_models,
         V={
-            week: UniVariateEstimator({reservoir.area: V[week]})
+            week: UniVariateEstimator({reservoir.area.area: V[week]})
             for week in range(param.len_week + 1)
         },
     )
 
     assert upper_bound == pytest.approx(5046992854.133574)
-    assert controls[TimeScenarioIndex(0, 0)][reservoir.area] == pytest.approx(1146984.0)
+    assert controls[TimeScenarioIndex(0, 0)][reservoir.area.area] == pytest.approx(
+        1146984.0
+    )
     assert len(problem.solver.constraints()) == 3555
     assert len(problem.solver.variables()) == 3533
 
@@ -269,15 +273,17 @@ 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({AreaIndex(reservoir.area): X}),
+        stock_discretization=StockDiscretization({reservoir.area: X}),
         list_models=list_models,
         V={
-            week: UniVariateEstimator({reservoir.area: V[week]})
+            week: UniVariateEstimator({reservoir.area.area: V[week]})
             for week in range(param.len_week + 1)
         },
     )
 
-    assert controls[TimeScenarioIndex(0, 0)][reservoir.area] == pytest.approx(133776)
+    assert controls[TimeScenarioIndex(0, 0)][reservoir.area.area] == pytest.approx(
+        133776
+    )
 
     control = 123864.0
     vb = V[1](reservoir.initial_level + reservoir.inflow[0, 0] - control)
@@ -321,16 +327,16 @@ def test_upper_bound_with_xpress() -> None:
         upper_bound, controls, _ = compute_upper_bound(
             param=param,
             multi_stock_management=MultiStockManagement([reservoir_management]),
-            stock_discretization=StockDiscretization({AreaIndex(reservoir.area): X}),
+            stock_discretization=StockDiscretization({reservoir.area: X}),
             list_models=list_models,
             V={
-                week: UniVariateEstimator({reservoir.area: V[week]})
+                week: UniVariateEstimator({reservoir.area.area: V[week]})
                 for week in range(param.len_week + 1)
             },
         )
 
         assert upper_bound == pytest.approx(380492940.000565)
-        assert controls[TimeScenarioIndex(0, 0)][reservoir.area] == pytest.approx(
+        assert controls[TimeScenarioIndex(0, 0)][reservoir.area.area] == pytest.approx(
             4482011.0
         )
         assert len(problem.solver.constraints()) == 3555
@@ -340,16 +346,16 @@ def test_upper_bound_with_xpress() -> None:
         upper_bound, controls, _ = compute_upper_bound(
             param=param,
             multi_stock_management=MultiStockManagement([reservoir_management]),
-            stock_discretization=StockDiscretization({AreaIndex(reservoir.area): X}),
+            stock_discretization=StockDiscretization({reservoir.area: X}),
             list_models=list_models,
             V={
-                week: UniVariateEstimator({reservoir.area: V[week]})
+                week: UniVariateEstimator({reservoir.area.area: V[week]})
                 for week in range(param.len_week + 1)
             },
         )
 
         assert upper_bound == pytest.approx(5046992854.133574)
-        assert controls[TimeScenarioIndex(0, 0)][reservoir.area] == pytest.approx(
+        assert controls[TimeScenarioIndex(0, 0)][reservoir.area.area] == pytest.approx(
             1146984.0
         )
         assert len(problem.solver.constraints()) == 3555