Merge pull request #224 from pph-collective/develop

add age bins

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "titan-model"
-version = "3.0.0"
+version = "3.1.0"
 description = "TITAN Agent Based Model"
 license = "GPL-3.0-only"
 authors = ["Sam Bessey <[email protected]>", "Mary McGrath <[email protected]>"]

tests/agent_test.py

@@ -257,3 +257,18 @@ def test_clear_set(make_agent):
     assert s.members == set()
     assert a not in s
     assert s.num_members() == 0
+
+
+@pytest.mark.unit
+def test_age_bin(make_agent):
+    a = make_agent()
+    a.age = 30
+    assert a.age_bin == "mid_adult"
+
+
+@pytest.mark.unit
+def test_age_bin_error(make_agent):
+    a = make_agent()
+    a.age = -1
+    with pytest.raises(ValueError):
+        a.age_bin

tests/params/basic.yml

@@ -391,3 +391,16 @@ classes:
     new_ag:
       enter_type: new_agent
       prob: 1.0
+  age_bins:
+    fallback_young:
+      min_age: 0
+      max_age: 15
+    young_adult:
+      min_age: 16
+      max_age: 24
+    mid_adult:
+      min_age: 25
+      max_age: 44
+    fallback_old:
+      min_age: 45
+      max_age: 999

tests/partnering_test.py

@@ -583,6 +583,91 @@ def test_get_assort_hiv_prep(make_population, make_agent, params):
     assert partner == p2
 
 
+@pytest.mark.unit
+def test_get_assort_age_bin_same(make_population, make_agent, params):
+    pop = make_population()
+    a = make_agent(age=30)
+    p1 = make_agent(age=30)
+    p2 = make_agent(age=66)
+    for bond in params.classes.bond_types:
+        a.target_partners[bond] = 1
+        p1.target_partners[bond] = 1
+        p2.target_partners[bond] = 1
+        a.partners[bond] = set()
+        p1.partners[bond] = set()
+        p2.partners[bond] = set()
+    pop.add_agent(a)
+    pop.add_agent(p1)
+    pop.add_agent(p2)
+
+    params.features.assort_mix = True
+
+    test_rule = ObjMap(
+        {
+            "attribute": "age_bin",
+            "partner_attribute": "__agent__",
+            "bond_types": ["Sex"],
+            "agent_value": "__any__",
+            "partner_values": {"__other__": 0.1, "__same__": 0.9},
+        }
+    )
+    params.assort_mix["test_rule"] = test_rule
+
+    partner = select_partner(
+        a,
+        pop.all_agents.members,
+        pop.sex_partners,
+        pop.pwid_agents,
+        params,
+        FakeRandom(0.5),
+        "Sex",
+    )
+
+    assert partner == p1
+
+
+@pytest.mark.unit
+def test_get_assort_age_bin_other(make_population, make_agent, params):
+    pop = make_population()
+    a = make_agent(age=30)
+    p1 = make_agent(age=30)
+    p2 = make_agent(age=16)
+    for bond in params.classes.bond_types:
+        a.target_partners[bond] = 1
+        p1.target_partners[bond] = 1
+        p2.target_partners[bond] = 1
+        a.partners[bond] = set()
+        p1.partners[bond] = set()
+        p2.partners[bond] = set()
+    pop.add_agent(a)
+    pop.add_agent(p1)
+    pop.add_agent(p2)
+
+    params.features.assort_mix = True
+
+    test_rule = ObjMap(
+        {
+            "attribute": "age_bin",
+            "partner_attribute": "__agent__",
+            "bond_types": ["Sex"],
+            "agent_value": "__any__",
+            "partner_values": {"__other__": 0.9, "__same__": 0.1},
+        }
+    )
+    params.assort_mix["test_rule"] = test_rule
+
+    partner = select_partner(
+        a,
+        pop.all_agents.members,
+        pop.sex_partners,
+        pop.pwid_agents,
+        params,
+        FakeRandom(0.5),
+        "Sex",
+    )
+    assert partner == p2
+
+
 @pytest.mark.unit
 def test_get_assort_bond_type(make_population, make_agent, params):
     pop = make_population()

titan/agent.py

@@ -109,6 +109,15 @@ def __ne__(self, other) -> bool:
     def __hash__(self) -> int:
         return self.id
 
+    @property
+    def age_bin(self):
+        for key, val in self.location.params.classes.age_bins.items():
+            if self.age >= val.min_age and self.age <= val.max_age:
+                return key
+        raise ValueError(
+            f"Agent age ({self.age}) must be within an age bin [params.classes.age_bins]"
+        )
+
     def iter_partners(self) -> Iterator["Agent"]:
         """
         Get an iterator over an agent's partners

titan/params/assort_mix.yml

@@ -15,7 +15,7 @@ assort_mix:
       class: bond_types
       default: []
     agent_value:
-      description: "Value of agent's attribute which triggers using this rule.  Can use the value `__any__` to then allow assoring on partner_values of `__same__` and `__other__`."
+      description: "Value of agent's attribute which triggers using this rule.  Can use the value `__any__` to then allow assorting on partner_values of `__same__` and `__other__`."
       type: none
     partner_values:
       description: "Which partners to select if assortative mixing happens. Expects a sub-key of an attribute value with a value of the probability. `__other__` can be used as an attribute value as a fallthrough option.  If the `partner_attribute` is `location`, it is also possible to use the key `__neighbor__` here to indicate a probability of assorting with agents from neighboring locations (vs. `__same__` or `__other__`).  All of the probabilities should add to one."

titan/params/classes.yml

@@ -185,3 +185,25 @@ classes:
     default:
       replace:
         enter_type: replace
+  age_bins:
+    type: definition
+    description: Bin of ages into categories for use in assortative mixing by age. Will error if the agent's age is not in a bin.
+    fields:
+      min_age:
+        description: Lower bound (inclusive) of this "bin".
+        type: int
+        min: 0
+      max_age:
+        description: Upper bound (inclusive) of this "bin".
+        type: int
+        min: 0
+    default:
+      0:
+        min_age: 0
+        max_age: 24
+      1:
+        min_age: 25
+        max_age: 54
+      2:
+        min_age: 55
+        max_age: 9999