From 4038a4b8c0a68599477ec4fe4ac1f065fe2fc775 Mon Sep 17 00:00:00 2001 From: Sam Bessey Date: Mon, 6 Dec 2021 13:29:37 -0500 Subject: [PATCH 1/2] refactor(distributions): move poisson wrapper to distributions.py Poisson wrapper was located in utils. That and the order of checking for distribution location used np poisson for safe_dist. --- titan/distributions.py | 9 +++++++++ titan/interactions/injection.py | 4 ++-- titan/interactions/sex.py | 4 ++-- titan/population.py | 5 +++-- titan/utils.py | 11 ++--------- 5 files changed, 18 insertions(+), 15 deletions(-) diff --git a/titan/distributions.py b/titan/distributions.py index 98bdf388..2d36a808 100644 --- a/titan/distributions.py +++ b/titan/distributions.py @@ -56,3 +56,12 @@ def weibull_modified(np_random, shape, scale): """ random_number = np_random.random() return scale * (-log(1 - random_number)) ** (1 / shape) + + +def poisson(np_rand, mu: float): + """ + Mirrors scipy poisson.rvs function as used in code + """ + if mu < 0: + return 0 + return np_rand.poisson(mu) diff --git a/titan/interactions/injection.py b/titan/interactions/injection.py index 30a56c37..28ff4e88 100644 --- a/titan/interactions/injection.py +++ b/titan/interactions/injection.py @@ -1,8 +1,8 @@ from . import base_interaction -from .. import utils from .. import features from .. import model from .. import agent +from ..distributions import poisson class Injection(base_interaction.BaseInteraction): @@ -39,7 +39,7 @@ def get_num_acts(cls, model: "model.TITAN", rel: "agent.Relationship") -> int: min(agent_params.num_acts, partner_params.num_acts) * model.calibration.injection.act ) - share_acts = utils.poisson(mean_num_acts, model.np_random) + share_acts = poisson(mean_num_acts, model.np_random) if share_acts < 1: return 0 diff --git a/titan/interactions/sex.py b/titan/interactions/sex.py index fe005c47..dc0a25f8 100644 --- a/titan/interactions/sex.py +++ b/titan/interactions/sex.py @@ -1,5 +1,5 @@ from . import base_interaction -from .. import utils +from ..distributions import poisson from .. import model from .. import agent @@ -21,7 +21,7 @@ def get_num_acts(cls, model: "model.TITAN", rel: "agent.Relationship") -> int: mean_sex_acts = ( rel.get_number_of_sex_acts(model.np_random) * model.calibration.sex.act ) - total_sex_acts = utils.poisson(mean_sex_acts, model.np_random) + total_sex_acts = poisson(model.np_random, mean_sex_acts) # Get condom usage p_safe_sex = ( diff --git a/titan/population.py b/titan/population.py index 4ea2ae57..44b1d9f2 100644 --- a/titan/population.py +++ b/titan/population.py @@ -19,6 +19,7 @@ from . import utils from . import features from . import exposures +from .distributions import poisson class Population: @@ -431,8 +432,8 @@ def update_partner_targets(self): """ for a in self.all_agents: for bond in self.params.classes.bond_types: - a.target_partners[bond] = utils.poisson( - a.mean_num_partners[bond], self.np_random + a.target_partners[bond] = poisson( + self.np_random, a.mean_num_partners[bond] ) self.update_partnerability(a) diff --git a/titan/utils.py b/titan/utils.py index 0afceabf..fc35d929 100644 --- a/titan/utils.py +++ b/titan/utils.py @@ -143,11 +143,11 @@ def parse_var(dist_value, dist_type): def get_dist(rand_gen, dist_type): if dist_type == "randint": return lambda *args: safe_random_int(*args, rand_gen) - elif hasattr(rand_gen, dist_type): - return getattr(rand_gen, dist_type) elif hasattr(distributions, dist_type): dist = getattr(distributions, dist_type) return lambda *args: dist(rand_gen, *args) + elif hasattr(rand_gen, dist_type): + return getattr(rand_gen, dist_type) else: raise AttributeError(f"Distribution type {dist_type} not found!") @@ -184,13 +184,6 @@ def binom_0(n: int, p: float): return (1 - p) ** n -def poisson(mu: float, np_rand): - """ - Mirrors scipy poisson.rvs function as used in code - """ - return np_rand.poisson(mu) - - def get_param_from_path(params: ObjMap, param_path: str, delimiter: str): """ Given a params object and a delimited path, get the leaf of the params tree From ad3e7343b6b13c8dae07796fb8ae835f9caa80d5 Mon Sep 17 00:00:00 2001 From: Sam Bessey Date: Tue, 21 Dec 2021 11:12:05 -0500 Subject: [PATCH 2/2] fix: update remaining utils.poisson to distributions.poisson --- titan/features/high_risk.py | 5 +++-- titan/interactions/injection.py | 2 +- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/titan/features/high_risk.py b/titan/features/high_risk.py index 15c048b8..7277904d 100644 --- a/titan/features/high_risk.py +++ b/titan/features/high_risk.py @@ -7,6 +7,7 @@ from .. import agent from .. import population from .. import model +from ..distributions import poisson class HighRisk(base_feature.BaseFeature): @@ -170,7 +171,7 @@ def update_partner_numbers(self, pop: "population.Population", amount: int): """ for bond in self.agent.location.params.high_risk.partnership_types: self.agent.mean_num_partners[bond] += amount # could be negative - self.agent.target_partners[bond] = utils.poisson( - self.agent.mean_num_partners[bond], pop.np_random + self.agent.target_partners[bond] = poisson( + pop.np_random, self.agent.mean_num_partners[bond] ) pop.update_partnerability(self.agent) diff --git a/titan/interactions/injection.py b/titan/interactions/injection.py index 28ff4e88..909ae7a5 100644 --- a/titan/interactions/injection.py +++ b/titan/interactions/injection.py @@ -39,7 +39,7 @@ def get_num_acts(cls, model: "model.TITAN", rel: "agent.Relationship") -> int: min(agent_params.num_acts, partner_params.num_acts) * model.calibration.injection.act ) - share_acts = poisson(mean_num_acts, model.np_random) + share_acts = poisson(model.np_random, mean_num_acts) if share_acts < 1: return 0