Use expressions instead of _number

openfisca_survey_manager/calibration.py

@@ -52,9 +52,8 @@ def __init__(self, simulation, target_margins, period, target_entity_count = Non
         variable_instance_by_variable_name = simulation.tax_benefit_system.variables
         entities = set(
             variable_instance_by_variable_name[variable].entity.key
-            for variable in margin_variables if not variable.endswith("_number")
+            for variable in margin_variables if not len(variable.split(" ")) > 1
             )
-        assert all([(len(variable.split("_number")) < 3) for variable in margin_variables]), "invalid name of a variable for calibration, with '_number' not only in last position"
         if entity is not None:
             entities.add(entity)
         self.entities = list(entities)
@@ -132,10 +131,10 @@ def _build_calmar_data(self) -> dict:
         data[self.target_entity][self._initial_weight_name] = self.initial_weight * self.filter_by
         period = self.period
         for variable in self.margins_by_variable:
-            assert variable in self.simulation.tax_benefit_system.variables or variable.split("_number")[0] in self.simulation.tax_benefit_system.variables
-            if variable.endswith("_number"):
-                assert self.simulation.tax_benefit_system.variables[variable.split("_number")[0]].value_type in [int, float], "variable with suffix 'number' should be float-like so we can separate the positive values"
-                value = (self.simulation.adaptative_calculate_variable(variable.split("_number")[0], period = period) > 0).astype(int)
+            assert variable in self.simulation.tax_benefit_system.variables or variable.split(" ")[0] in self.simulation.tax_benefit_system.variables
+            if len(variable.split(" ")) > 1:
+                var_to_eval = simulation.adaptative_calculate_variable(variable.split(" ")[0], period = period)
+                value = eval('var_to_eval' + ''.join(variable.split(" ")[1: ])).astype(float)
             else:
                 value = self.simulation.adaptative_calculate_variable(variable, period = period)
             data[self.simulation.tax_benefit_system.variables[variable].entity.key][variable] = value
@@ -215,8 +214,8 @@ def set_target_margin(self, variable, target):
         """
         simulation = self.simulation
         period = self.period
-        assert variable in simulation.tax_benefit_system.variables or variable.split("_number")[0] in simulation.tax_benefit_system.variables
-        variable_instance = simulation.tax_benefit_system.variables[variable.split("_number")[0]]
+        assert variable in simulation.tax_benefit_system.variables or variable.split(" ")[0] in simulation.tax_benefit_system.variables
+        variable_instance = simulation.tax_benefit_system.variables[variable.split(" ")[0]]
 
         filter_by = self.filter_by
         target_by_category = None
@@ -281,13 +280,14 @@ def _update_margins(self):
             filter_by = self.filter_by
             initial_weight = self.initial_weight
 
-            if variable.endswith("_number"):
-                value = (simulation.adaptative_calculate_variable(variable.split("_number")[0], period = period) < 0).astype(int)
+            if len(variable.split(" ")) > 1:
+                var_to_eval = simulation.adaptative_calculate_variable(variable.split(" ")[0], period = period)
+                value = eval( 'var_to_eval' + ''.join(variable.split(" ")[1: ])).astype(float)
             else:
                 value = simulation.adaptative_calculate_variable(variable, period = period)
             weight_variable = simulation.weight_variable_by_entity[target_entity]
 
-            if len(self.entities) == 2 and simulation.tax_benefit_system.variables[variable.split("_number")[0]].entity.key != self.target_entity:
+            if len(self.entities) == 2 and simulation.tax_benefit_system.variables[variable.split(" ")[0]].entity.key != self.target_entity:
                 value_df = pd.DataFrame(value)
                 id_variable = self.parameters["id_variable_link"]
                 value_df[id_variable] = simulation.adaptative_calculate_variable(id_variable, period = period)
@@ -306,7 +306,7 @@ def _update_margins(self):
                 ('initial', initial_weight),
                 ]
 
-            variable_instance = simulation.tax_benefit_system.get_variable(variable.split("_number")[0])
+            variable_instance = simulation.tax_benefit_system.get_variable(variable.split(" ")[0])
             assert variable_instance is not None
             if variable_instance.value_type in [bool, Enum]:
                 margin_items.append(('category', value))