diff --git a/pyproject.toml b/pyproject.toml index ad4c61ea..d4cd32e4 100644 --- a/pyproject.toml +++ b/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 ", "Mary McGrath "] diff --git a/tests/agent_test.py b/tests/agent_test.py index 9f75198c..a87906d1 100644 --- a/tests/agent_test.py +++ b/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 diff --git a/tests/params/basic.yml b/tests/params/basic.yml index 537e2cd1..60443957 100644 --- a/tests/params/basic.yml +++ b/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 diff --git a/tests/partnering_test.py b/tests/partnering_test.py index 35190481..41bf2798 100644 --- a/tests/partnering_test.py +++ b/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() diff --git a/titan/agent.py b/titan/agent.py index 9b0ac045..7b8db7e4 100644 --- a/titan/agent.py +++ b/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 diff --git a/titan/params/assort_mix.yml b/titan/params/assort_mix.yml index 4601f9e5..0c5b0734 100644 --- a/titan/params/assort_mix.yml +++ b/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." diff --git a/titan/params/classes.yml b/titan/params/classes.yml index f55563a8..1243c161 100644 --- a/titan/params/classes.yml +++ b/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