DemograhpicsTemplates: Remove CrudeRate/YearlyRate/DtkRate (#40)

ckirkman-IDM · web-flow · commit 8a5b665e40ae · 2025-11-05T11:45:36.000-08:00
diff --git a/emod_api/demographics/demographics_base.py b/emod_api/demographics/demographics_base.py
@@ -14,7 +14,6 @@
 
 from emod_api.demographics import demographics_templates as DT
 from emod_api.demographics.base_input_file import BaseInputFile
-from emod_api.demographics.demographics_templates import CrudeRate, YearlyRate
 from emod_api.demographics.node import Node
 from emod_api.demographics.age_distribution_old import AgeDistributionOld as AgeDistribution
 from emod_api.demographics.demographic_exceptions import InvalidNodeIdException
@@ -400,38 +399,40 @@ def SetMinimalNodeAttributes(self):
     # DTK is births per person per day.
     def SetBirthRate(self,
                      birth_rate: float,
-                     node_ids: List = None):
+                     node_ids: List[int] = None) -> None:
         """
         Set Default birth rate to birth_rate. Turn on Vital Dynamics and Births implicitly.
+
+        Args:
+            birth_rate: (float) The birth rate in units of births/year/1000-women
+            node_ids: a list of node_ids. None or 0 indicates the default node.
+
+        Returns:
+            Nothing
         """
         warnings.warn('SetBirthRate() is deprecated. Default nodes should now be represented by Node '
                       'objects and passed to the Demographics object during the constructor call. They can be modified '
                       'afterward, if needed.',
                       DeprecationWarning, stacklevel=2)
-        if type(birth_rate) is float or type(birth_rate) is int:
-            birth_rate = CrudeRate(birth_rate)
-        dtk_birthrate = birth_rate.get_dtk_rate()
+        dtk_birthrate = birth_rate / 365 / 1000
+
         if node_ids is None:
             self.raw['Defaults']['NodeAttributes'].update({
                 "BirthRate": dtk_birthrate
             })
         else:
-            for node_id in node_ids:
-                self.get_node_by_id(node_id=node_id).birth_rate = dtk_birthrate
+            nodes = self.get_nodes_by_id(node_ids=node_ids)
+            for _, node in nodes.items():
+                node.birth_rate = dtk_birthrate
         self.implicits.append(DT._set_population_dependent_birth_rate)
 
     def SetMortalityRate(self,
-                         mortality_rate: CrudeRate, node_ids: List[int] = None):
+                         mortality_rate: float, node_ids: List[int] = None):
         """
         Set constant mortality rate to mort_rate. Turn on Enable_Natural_Mortality implicitly.
         """
         warnings.warn('SetMortalityRate() is deprecated. Please use the emodpy Demographics method: '
                       'set_mortality_distribution()', DeprecationWarning, stacklevel=2)
-
-        # yearly_mortality_rate = YearlyRate(mortality_rate)
-        if type(mortality_rate) is float or type(mortality_rate) is int:
-            mortality_rate = CrudeRate(mortality_rate)
-        mortality_rate = mortality_rate.get_dtk_rate()
         if node_ids is None:
             # setting = {"MortalityDistribution": DT._ConstantMortality(yearly_mortality_rate).to_dict()}
             setting = {"MortalityDistribution": DT._ConstantMortality(mortality_rate).to_dict()}
@@ -643,7 +644,8 @@ def SetDefaultNodeAttributes(self, birth=True):
                                                   "Region": 1,
                                                   "Seaport": 1}
         if birth:
-            self.SetBirthRate(YearlyRate(math.log(1.03567)))
+            # raise Exception("This will be removed in a new issue/PR shortly. Do not use.")
+            self.SetBirthRate(birth_rate=math.log(1.03567))
 
     def SetDefaultProperties(self):
         """
@@ -674,26 +676,36 @@ def SetDefaultFromTemplate(self, template, setter_fn=None):
 
     # TODO: is this useful in a way that warrants a special-case function in emodpy built around set_age_distribution?
     #  https://github.com/InstituteforDiseaseModeling/emod-api-old/issues/788
-    def SetEquilibriumAgeDistFromBirthAndMortRates(self, CrudeBirthRate=CrudeRate(40), CrudeMortRate=CrudeRate(20),
-                                                   node_ids=None):
+    def SetEquilibriumAgeDistFromBirthAndMortRates(self,
+                                                   birth_rate: float = 40.0,
+                                                   mortality_rate: float = 20.0,
+                                                   node_ids: List[int] = None):
         """
-        Set the inital ages of the population to a sensible equilibrium profile based on the specified input birth and
-        death rates. Note this does not set the fertility and mortality rates.
+            Set age distribution based on birth and death rates. Implicit function.
+
+            Args:
+                birth_rate: (float) The birth rate in units of births/year/1000-women
+                mortality_rate: (float) The mortality rate in units of deaths/year/1000 people
+                node_ids: a list of node_ids. None or 0 indicates the default node.
+            Returns:
+                Nothing
         """
-        warnings.warn('SetEquilibriumAgeDistFromBirthAndMortRates() is deprecated. Please use the emodpy Demographics method: '
-                      'set_age_distribution()', DeprecationWarning, stacklevel=2)
+        warnings.warn(
+            'SetEquilibriumAgeDistFromBirthAndMortRates() is deprecated. Please use the emodpy Demographics method: '
+            'set_age_distribution()', DeprecationWarning, stacklevel=2)
 
-        yearly_birth_rate = YearlyRate(CrudeBirthRate)
-        yearly_mortality_rate = YearlyRate(CrudeMortRate)
-        dist = DT._EquilibriumAgeDistFromBirthAndMortRates(yearly_birth_rate, yearly_mortality_rate)
+        dist = DT._EquilibriumAgeDistFromBirthAndMortRates(birth_rate=birth_rate,
+                                                           mortality_rate=mortality_rate)
         setter_fn = DT._set_age_complex
+
         if node_ids is None:
             self.SetDefaultFromTemplate(dist, setter_fn)
         else:
             new_dist = AgeDistribution()
             dist = new_dist.from_dict(dist["AgeDistribution"])
-            for node in node_ids:
-                self.get_node_by_id(node_id=node)._set_age_complex_distribution(dist)
+            nodes = self.get_nodes_by_id(node_ids=node_ids)
+            for _, node in nodes.items():
+                node._set_age_complex_distribution(dist)
             self.implicits.append(setter_fn)
 
     def SetInitialAgeExponential(self, rate=0.0001068, description=""):
diff --git a/emod_api/demographics/demographics_templates.py b/emod_api/demographics/demographics_templates.py
@@ -25,34 +25,6 @@ class DemographicsTemplatesConstants:
                                        32849.6, 34679.5, 34679.6, 36509.5, 36509.6, 38339.5]
 
 
-class CrudeRate:  # would like to derive from float
-    def __init__(self, init_rate):
-        self._time_units = 365
-        self._people_units = 1000
-        self._rate = init_rate
-
-    def get_dtk_rate(self):
-        return self._rate / self._time_units / self._people_units
-
-
-class YearlyRate(CrudeRate):  # would like to derive from float
-    def __init__(self, init_rate):
-        self._time_units = 365
-        self._people_units = 1
-        if type(init_rate) is CrudeRate:
-            self._rate = init_rate._rate / 1000.
-        else:
-            self._rate = init_rate
-
-
-class DtkRate(CrudeRate):
-    def __init__(self, init_rate):
-        super().__init__(init_rate)
-        self._time_units = 1
-        self._people_units = 1
-        self._rate = init_rate
-
-
 # Migration
 def _set_migration_model_fixed_rate(config):
     config.parameters.Migration_Model = "FIXED_RATE_MIGRATION"
@@ -557,24 +529,30 @@ def AgeStructureUNWPP(demog):
     demog.SetDefaultFromTemplate(setting, _set_age_complex)
 
 
-def _EquilibriumAgeDistFromBirthAndMortRates(birth_rate=YearlyRate(40 / 1000.), mort_rate=YearlyRate(20 / 1000.)):
+def _EquilibriumAgeDistFromBirthAndMortRates(birth_rate=40.0, mortality_rate=20.0):
     """
-    Set age distribution based on birth and death rates.
+    Set age distribution based on birth and death rates. Implicit function.
 
     Args:
-        birth_rate: births per person per year.
-        mort_rate: deaths per person per year.
+        # birth_rate: births per person per year.
+        birth_rate: (float) The birth rate in units of births/year/1000-women
+        # mort_rate: deaths per person per year.
+        mortality_rate: (float) The mortality rate in units of deaths/year/1000 people
 
     Returns:
         dictionary which can be inserted into demographics object.
 
     """
-    BirthRate = math.log(1 + birth_rate.get_dtk_rate())
-    MortRate = -1 * math.log(1 - mort_rate.get_dtk_rate())
+    # convert to daily rate per person, EMOD units
+    birth_rate = (birth_rate / 1000) / 365  # what is actually used below
+    mortality_rate = (mortality_rate / 1000) / 365  # what is actually used below
+
+    birth_rate = math.log(1 + birth_rate)
+    mortality_rate = -1 * math.log(1 - mortality_rate)
 
     # It is important for the age distribution computation that the age-spacing be very fine; I've used 30 days here.
     # With coarse spacing, the computation in practice doesn't work as well.
-    ageDist = _computeAgeDist(BirthRate, [i * 30 for i in range(1200)], 1200 * [MortRate], 12 * [1.0])
+    ageDist = _computeAgeDist(birth_rate, [i * 30 for i in range(1200)], 1200 * [mortality_rate], 12 * [1.0])
 
     # The final demographics file, though, can use coarser binning interpolated from the finely-spaced computed distribution.
     EMODAgeBins = list(range(16)) + [20 + 5 * i for i in range(14)]
diff --git a/tests/test_demog.py b/tests/test_demog.py
@@ -264,7 +264,7 @@ def test_template_age_structure_UNWPP(self):
 
     def test_template_equilibrium_age_dist_from_birth_and_mort_rates(self):
         demog = Demographics.from_template_node()
-        demog.SetEquilibriumAgeDistFromBirthAndMortRates(CrudeBirthRate=20 / 1000, CrudeMortRate=10 / 1000)
+        demog.SetEquilibriumAgeDistFromBirthAndMortRates(birth_rate=20, mortality_rate=10)
         self.assertIn('AgeDistribution', demog.raw['Defaults']['IndividualAttributes'])
         self.assertEqual(len(demog.implicits), 2)
         print(demog.raw)
@@ -297,13 +297,14 @@ def test_set_default_from_template_mortality_rate_by_age(self):
         self.assertEqual(len(mort_rate), len(mort_dist['PopulationGroups'][1]))
         self.assertIn('MortalityDistribution', demog.raw['Defaults']['IndividualAttributes'])  # Can't use set_default_from_template_test since template is implicit
 
+    # TODO: restore/refactor after moving new distribution classes down into emod-api? Or is this duplicative?
     def test_set_default_from_template_constant_mortality(self):
         demog = Demographics.from_template_node()
         demog.implicits = []
-        mortality_rate = DT.DtkRate(0.0001)
+        mortality_rate = 0.0001
         demog.SetMortalityRate(mortality_rate=mortality_rate)  # ca
         self.assertIn('MortalityDistribution', demog.raw['Defaults']['IndividualAttributes'])  # Can't use set_default_from_template_test since template is implicit
-        expected_rate = [[-1 * (math.log(1 - mortality_rate.get_dtk_rate()))]] * 2
+        expected_rate = [[-1 * (math.log(1 - mortality_rate))]] * 2
         demog_rate = demog.raw['Defaults']['IndividualAttributes']['MortalityDistribution']['ResultValues']
         self.assertListEqual(expected_rate, demog_rate)
 
@@ -1088,13 +1089,13 @@ def test_set_predefined_mortality_distribution(self):
     def test_mortality_rate_with_node_ids(self):
         input_file = os.path.join(manifest.demo_folder, 'nodes.csv')
         demog = Demographics.from_csv(input_file)
-        mortality_rate = 0.1234  # CrudeRate
+        mortality_rate = 0.1234 / 365 / 1000
         node_ids = [97, 99]
         demog.SetMortalityRate(mortality_rate, node_ids)
 
         set_mortality_dist_97 = demog.get_node_by_id(node_id=97).individual_attributes.mortality_distribution.to_dict()
         set_mortality_dist_99 = demog.get_node_by_id(node_id=99).individual_attributes.mortality_distribution.to_dict()
-        expected_mortality_dist = DT._ConstantMortality(DT.CrudeRate(mortality_rate).get_dtk_rate()).to_dict()
+        expected_mortality_dist = DT._ConstantMortality(mortality_rate).to_dict()
         self.assertDictEqual(set_mortality_dist_99, expected_mortality_dist)
         self.assertDictEqual(set_mortality_dist_97, expected_mortality_dist)
         self.assertIsNone(demog.get_node_by_id(node_id=96).individual_attributes.mortality_distribution)
