Added:

- Defined explicit data types for inputs and outputs of functions for better type checking and readability.
- Added a test script for `prepshot.utils`.

Fixed:

- Added `pyoptinterface._src.core_ext` to Pylint's extension package allow list to resolve cpp-extension-no-member warning.

Changed:

- Updated `model.py` to keep necessary decision variables and use expressions for intermediate variables instead of direct determination.
- Refactored `extract_results_non_hydro` in `output_data.py` to extract common features for different variables, simplifying the code.
- Removed definitions of complex sets and opted for simple sets wherever possible to streamline the code.

.pylintrc

@@ -0,0 +1,2 @@
+[MASTER]
+extension-pkg-allow-list=pyoptinterface._src.core_ext

prepshot/_model/co2.py

@@ -4,29 +4,41 @@
 """This module contains constraints related to carbon emissions. 
 """
 
+from typing import Union
+
 import pyoptinterface as poi
 import numpy as np
 
 class AddCo2EmissionConstraints:
     """Class for carbon emission constraints and calculations.
     """
-    def __init__(self, model):
+    def __init__(self,
+        model : Union[
+            poi._src.highs.Model,
+            poi._src.gurobi.Model,
+            poi._src.mosek.Model,
+            poi._src.copt.Model
+        ],
+    ) -> None:
         """Initialize the class.
 
         Parameters
         ----------
-        model : pyoptinterface._src.solver.Model
-            Model index.
-        params : dict
-            Dictionary containing parameters.
+        model : Union[
+            poi._src.highs.Model,
+            poi._src.gurobi.Model,
+            poi._src.mosek.Model,
+            poi._src.copt.Model
+        ]
+            Model object depending on the solver.
         """
         self.model = model
         model.carbon_breakdown = poi.make_tupledict(
-            model.year_zone_tech_tuples,
+            model.year, model.zone, model.tech,
             rule=self.carbon_breakdown
         )
         model.carbon_capacity = poi.make_tupledict(
-            model.year_zone_tuples,
+            model.year, model.zone,
             rule=self.emission_calc_by_zone_rule
         )
         model.carbon = poi.make_tupledict(
@@ -38,7 +50,10 @@ def __init__(self, model):
             rule=self.emission_limit_rule
         )
 
-    def emission_limit_rule(self, y):
+    def emission_limit_rule(
+        self,
+        y : int
+    ) -> poi._src.core_ext.ConstraintIndex:
         """Annual carbon emission limits across all zones and technologies.
         
         Parameters
@@ -58,7 +73,10 @@ def emission_limit_rule(self, y):
         lhs = model.carbon[y] - limit[y]
         return model.add_linear_constraint(lhs, poi.Leq, 0)
 
-    def emission_calc_rule(self, y):
+    def emission_calc_rule(
+        self,
+        y : int
+    ) -> poi._src.core_ext.ConstraintIndex:
         """Calculation of annual carbon emission across all zones and
         technologies.
 
@@ -78,7 +96,11 @@ def emission_calc_rule(self, y):
             for z in model.zone
         )
 
-    def emission_calc_by_zone_rule(self, y, z):
+    def emission_calc_by_zone_rule(
+        self,
+        y : int,
+        z : str
+    ) -> poi._src.core_ext.ConstraintIndex:
         """Calculation of annual carbon emissions by zone.
 
         Parameters
@@ -99,7 +121,12 @@ def emission_calc_by_zone_rule(self, y, z):
             for te in model.tech
         )
 
-    def carbon_breakdown(self, y, z, te):
+    def carbon_breakdown(
+        self,
+        y : int,
+        z : str,
+        te : str
+    ) -> poi._src.core_ext.ExprBuilder:
         """Carbon emission cost breakdown.
 
         Parameters
@@ -121,5 +148,5 @@ def carbon_breakdown(self, y, z, te):
         dt = model.params['dt']
         return poi.quicksum(
             ef * model.gen[h, m, y, z, te] * dt
-            for h, m in model.hour_month_tuples
+            for h in model.hour for m in model.month
         )

prepshot/_model/cost.py

@@ -4,12 +4,21 @@
 """This module contains objective functions for the model.
 """
 
+from typing import Union
+
 import pyoptinterface as poi
 
 class AddCostObjective:
     """Objective function class to determine the total cost of the model.
     """
-    def __init__(self, model):
+    def __init__(
+        self, model : Union[
+            poi._src.highs.Model,
+            poi._src.gurobi.Model,
+            poi._src.mosek.Model,
+            poi._src.copt.Model
+        ]
+    ) -> None:
         """The constructor for objective functions class.
 
         Parameters
@@ -20,7 +29,7 @@ def __init__(self, model):
         self.model = model
         self.define_objective()
 
-    def define_objective(self):
+    def define_objective(self) -> None:
         """Objective function of the model, to minimize total cost.
         """
         model = self.model
@@ -33,7 +42,11 @@ def define_objective(self):
             + model.cost_fix + model.cost_newline - model.income
         model.set_objective(model.cost, sense=poi.ObjectiveSense.Minimize)
 
-    def fuel_cost_breakdown(self, y, z, te):
+    def fuel_cost_breakdown(self,
+            y : int,
+            z : str,
+            te : str
+        ) -> poi._src.core_ext.ExprBuilder:
         """Fuel cost breakdown of technologies.
 
         Parameters
@@ -56,9 +69,13 @@ def fuel_cost_breakdown(self, y, z, te):
         vf = model.params['var_factor'][y]
         w = model.params['weight']
         return 1 / w * fp * sum(model.gen[h, m, y, z, te]
-            for (h,m) in model.hour_month_tuples) * dt * vf
+            for h in model.hour for m in model.month) * dt * vf
 
-    def cost_var_line_breakdown(self, y, z, z1):
+    def cost_var_line_breakdown(self,
+            y : int,
+            z : str,
+            z1 : str
+        ) -> poi._src.core_ext.ExprBuilder:
         """Variable operation and maintenance cost breakdown of transmission 
         lines.
 
@@ -84,9 +101,13 @@ def cost_var_line_breakdown(self, y, z, z1):
         w = model.params['weight']
         return 0.5 / w * lvc * dt * vf * \
             sum(model.trans_export[h, m, y, z, z1]
-            for (h,m) in model.hour_month_tuples)
+            for h in model.hour for m in model.month)
 
-    def cost_var_tech_breakdown(self, y, z, te):
+    def cost_var_tech_breakdown(self,
+        y : int,
+        z : str,
+        te : str
+    ) -> poi._src.core_ext.ExprBuilder:
         """Variable operation and maintenance cost breakdown.
 
         Parameters
@@ -110,9 +131,13 @@ def cost_var_tech_breakdown(self, y, z, te):
         vf = model.params['var_factor'][y]
         w = model.params['weight']
         return 1 / w * tvc * sum(model.gen[h, m, y, z, te]
-            for (h,m) in model.hour_month_tuples) * dt * vf
+            for h in model.hour for m in model.month) * dt * vf
 
-    def cost_fix_line_breakdown(self, y, z, z1):
+    def cost_fix_line_breakdown(self,
+        y : int,
+        z : str,
+        z1 : str
+    ) -> poi._src.core_ext.ExprBuilder:
         """Fixed operation and maintenance cost breakdown of transmission 
         lines.
 
@@ -136,7 +161,9 @@ def cost_fix_line_breakdown(self, y, z, z1):
         ff = model.params['fix_factor']
         return lfc[z, z1] * model.cap_lines_existing[y, z, z1] * ff[y] * 0.5
 
-    def cost_fix_tech_breakdown(self, y, z, te):
+    def cost_fix_tech_breakdown(self,
+        y : int, z : str, te : str
+    ) -> poi._src.core_ext.ExprBuilder:
         """Fixed operation and maintenance cost breakdown.
 
         Parameters
@@ -159,7 +186,9 @@ def cost_fix_tech_breakdown(self, y, z, te):
         ff = model.params['fix_factor'][y]
         return  tfc * model.cap_existing[y, z, te] * ff
 
-    def cost_newtech_breakdown(self, y, z, te):
+    def cost_newtech_breakdown(self,
+        y : int, z : str, te : str
+    ) -> poi._src.core_ext.ExprBuilder:
         """New technology investment cost breakdown.
 
         Parameters
@@ -182,7 +211,9 @@ def cost_newtech_breakdown(self, y, z, te):
         ivf = model.params['inv_factor'][te, y]
         return tic * model.cap_newtech[y, z, te] * ivf
 
-    def cost_newline_breakdown(self, y, z, z1):
+    def cost_newline_breakdown(
+        self, y : int, z : str, z1 : str
+    ) -> poi._src.core_ext.ExprBuilder:
         """New transmission line investment cost breakdown.
 
         Parameters
@@ -207,7 +238,7 @@ def cost_newline_breakdown(self, y, z, z1):
         capacity_invested_line = model.cap_newline[y, z, z1]
         return lic * capacity_invested_line * d * ivf * 0.5
 
-    def income_rule(self):
+    def income_rule(self) -> poi._src.core_ext.ExprBuilder:
         """Income from water withdrawal.
         Reference: https://www.nature.com/articles/s44221-023-00126-0
 
@@ -221,13 +252,16 @@ def income_rule(self):
             coef = 3600 * model.params['dt'] * model.params['price']
             income = sum(
                 model.withdraw[s, h, m, y] * coef
-                for s, h, m, y in model.station_hour_month_year_tuples
+                for s in model.station
+                for h in model.hour
+                for m in model.month
+                for y in model.year
             )
             return income
 
         return poi.ExprBuilder(0)
 
-    def var_cost_rule(self):
+    def var_cost_rule(self) -> poi._src.core_ext.ExprBuilder:
         """Calculate total variable cost, which is sum of the fuel cost of 
         technologies and variable Operation and maintenance (O&M) cost of 
         technologies and transmission lines.
@@ -239,17 +273,17 @@ def var_cost_rule(self):
         """
         model = self.model
         model.cost_var_tech_breakdown = poi.make_tupledict(
-            model.year_zone_tech_tuples,
+            model.year, model.zone, model.tech,
             rule=self.cost_var_tech_breakdown
         )
 
         model.cost_fuel_breakdown = poi.make_tupledict(
-            model.year_zone_tech_tuples,
+            model.year, model.zone, model.tech,
             rule=self.fuel_cost_breakdown
         )
 
         model.cost_var_line_breakdown = poi.make_tupledict(
-            model.year_zone_zone_tuples,
+            model.year, model.zone, model.zone,
             rule=self.cost_var_line_breakdown
         )
         cost_var = poi.ExprBuilder()
@@ -258,7 +292,7 @@ def var_cost_rule(self):
         cost_var += poi.quicksum(model.cost_var_line_breakdown)
         return cost_var
 
-    def newtech_cost_rule(self):
+    def newtech_cost_rule(self) -> poi._src.core_ext.ExprBuilder:
         """Total investment cost of new technologies.
         
         Returns
@@ -269,12 +303,12 @@ def newtech_cost_rule(self):
         """
         model = self.model
         model.cost_newtech_breakdown = poi.make_tupledict(
-            model.year_zone_tech_tuples,
+            model.year, model.zone, model.tech,
             rule=self.cost_newtech_breakdown
         )
         return poi.quicksum(model.cost_newtech_breakdown)
 
-    def newline_cost_rule(self):
+    def newline_cost_rule(self) -> poi._src.core_ext.ExprBuilder:
         """Total investment cost of new transmission lines.
 
         Returns
@@ -285,12 +319,12 @@ def newline_cost_rule(self):
         """
         model = self.model
         model.cost_newline_breakdown = poi.make_tupledict(
-            model.year_zone_zone_tuples,
+            model.year, model.zone, model.zone,
             rule=self.cost_newline_breakdown
         )
         return poi.quicksum(model.cost_newline_breakdown)
 
-    def fix_cost_rule(self):
+    def fix_cost_rule(self) -> poi._src.core_ext.ExprBuilder:
         """Fixed O&M cost of technologies and transmission lines.
         
         Returns
@@ -301,11 +335,11 @@ def fix_cost_rule(self):
         """
         model = self.model
         model.cost_fix_tech_breakdown = poi.make_tupledict(
-            model.year_zone_tech_tuples,
+            model.year, model.zone, model.tech,
             rule=self.cost_fix_tech_breakdown
         )
         model.cost_fix_line_breakdown = poi.make_tupledict(
-            model.year_zone_zone_tuples,
+            model.year, model.zone, model.zone,
             rule=self.cost_fix_line_breakdown
         )
         return poi.quicksum(model.cost_fix_tech_breakdown) + \

prepshot/_model/demand.py

@@ -3,20 +3,31 @@
 
 """This module contains constraints related to demand. 
 """
+from typing import Union
 
 import pyoptinterface as poi
 
 class AddDemandConstraints:
     """This class contains demand constraints.
     """
-    def __init__(self, model):
+    def __init__(self,
+        model : Union[
+            poi._src.highs.Model,
+            poi._src.gurobi.Model,
+            poi._src.mosek.Model,
+            poi._src.copt.Model
+        ]
+    ) -> None:
         """Initialize the class and add constraints.
         """
         self.model = model
         model.power_balance_cons = poi.make_tupledict(
-            model.hour_month_year_zone_tuples, rule=self.power_balance_rule
+            model.hour, model.month, model.year, model.zone,
+            rule=self.power_balance_rule
         )
-    def power_balance_rule(self, h, m, y, z):
+    def power_balance_rule(self,
+        h : int, m : int, y : int, z : str
+    ) -> poi._src.core_ext.ConstraintIndex:
         """Nodal power balance. The total electricity demand for each time 
         period and in each zone should be met by the following.
         
@@ -42,15 +53,12 @@ def power_balance_rule(self, h, m, y, z):
             Constraint index of the model.
         """
         model = self.model
-        lc = model.params['transmission_line_existing_capacity']
         load = model.params['demand']
         imp_z = poi.quicksum(
-            model.trans_import[h, m, y, z1, z]
-            for z1 in model.zone if (z, z1) in lc.keys()
+            model.trans_import[h, m, y, z1, z] for z1 in model.zone
         )
         exp_z = poi.quicksum(
-            model.trans_export[h, m, y, z, z1]
-            for z1 in model.zone if (z, z1) in lc.keys()
+            model.trans_export[h, m, y, z, z1] for z1 in model.zone
         )
         gen_z = poi.quicksum(
             model.gen[h, m, y, z, te] for te in model.tech