def test_build_dynamic__with_microdistancing(monkeypatch):
"""
Ensure dynamic mixing matrix can use microdistancing set by the user.
"""
monkeypatch.setattr(location_adjuster, "get_country_mixing_matrix",
_get_country_mixing_matrix)
mobility_params = {
"mixing": {},
"age_mixing": None,
"microdistancing": {
"foo": {
"function_type": "empiric",
"parameters": {
"max_effect": 0.6,
"times": [0, 365],
"values": [0, 0.1],
},
"locations": ["work"],
}
},
"square_mobility_effect": False,
**UNTESTED_PARAMS,
}
mm_func = build_dynamic_mixing_matrix(
base_matrix=MM,
country=Country(iso3="AUS"),
mobility=Mobility(**mobility_params),
)
mm = mm_func(0)
assert_allclose(mm, MM, atol=0.01, verbose=True)
mm = mm_func(365)
expected_mm = MM.copy() + ((1 - 0.06) - 1) * WORK_MM
assert_allclose(mm, expected_mm, atol=0.01, verbose=True)
def build_victorian_mixing_matrix_func(
static_mixing_matrix,
metro_mobility,
regional_mobility,
country,
intercluster_mixing_matrix,
):
# Having a consistent ordering of cluster names when building the matrix is critical, or else
# we will mix up our clutsers and people in Gippsland will be telepathically infecting people in the Grampians.
metro_clusters = [r.replace("-", "_") for r in Region.VICTORIA_METRO]
regional_clusters = [region.replace("-", "_") for region in Region.VICTORIA_RURAL]
all_clusters = [region.replace("-", "_") for region in Region.VICTORIA_SUBREGIONS]
# Collate the cluster-specific mixing matrices
cluster_age_mm_funcs = []
for region in all_clusters:
# Get the mobility parameters out for the cluster of interest
if region in metro_clusters:
cluster_mobility = deepcopy(metro_mobility)
elif region in regional_clusters:
cluster_mobility = deepcopy(regional_mobility)
else:
raise ValueError("Mobility region not found")
cluster_mobility.region = region.upper()
# Build the cluster-specific dynamic mixing matrix
cluster_age_mm_func = build_dynamic_mixing_matrix(
static_mixing_matrix,
cluster_mobility,
country,
)
cluster_age_mm_funcs.append(cluster_age_mm_func)
def get_mixing_matrix(self: StratifiedModel, time: float):
# Collate the within-cluster mixing matrices
cluster_age_mms = np.zeros((len(cluster_age_mm_funcs), *static_mixing_matrix.shape))
for idx, func in enumerate(cluster_age_mm_funcs):
cluster_age_mms[idx] = func(time)
# Pre-allocate
static_matrix_size = len(static_mixing_matrix)
num_clusters = len(all_clusters)
combined_matrix_size = static_matrix_size * num_clusters
super_matrix = np.zeros((combined_matrix_size, combined_matrix_size))
_set_mixing_matrix(
super_matrix,
intercluster_mixing_matrix,
cluster_age_mms,
static_matrix_size,
num_clusters,
)
return super_matrix
return get_mixing_matrix
def test_build_dynamic__with_age_mobility_data(monkeypatch):
"""
Ensure dynamic mixing matrix can use age-based mobility data set by the user.
"""
monkeypatch.setattr(location_adjuster, "get_country_mixing_matrix",
_get_country_mixing_matrix)
mobility_params = {
"mixing": {},
"age_mixing": {
"0": {
"times": [0, 1, 2, 3],
"values": [2, 2, 2, 2],
},
"5": {
"times": [0, 1, 2, 3],
"values": [3, 3, 3, 3],
},
},
"microdistancing": {},
"square_mobility_effect": False,
**UNTESTED_PARAMS,
}
mm_func = build_dynamic_mixing_matrix(
base_matrix=MM,
country=Country(iso3="AUS"),
mobility=Mobility(**mobility_params),
)
expected_mm = MM.copy()
# Add adjustment of 2 to the 1st row and col for the 0-5 age bracket
expected_mm[0, :] *= 2
expected_mm[:, 0] *= 2
# Add adjustment of 3 to the 2nd row and col for the 5-10 age bracket
expected_mm[1, :] *= 3
expected_mm[:, 1] *= 3
mm = mm_func(0)
assert_allclose(mm, expected_mm, atol=0.01, verbose=True)
mm = mm_func(2)
assert_allclose(mm, expected_mm, atol=0.01, verbose=True)
def test_build_dynamic__with_no_changes(monkeypatch):
"""
Ensure dynamic mixing matrix has no change over time, if no changes are supplied.
"""
monkeypatch.setattr(location_adjuster, "get_country_mixing_matrix",
_get_country_mixing_matrix)
mobility_params = {
"mixing": {},
"age_mixing": None,
"microdistancing": {},
"square_mobility_effect": False,
**UNTESTED_PARAMS,
}
mm_func = build_dynamic_mixing_matrix(base_matrix=MM,
country=Country(iso3="AUS"),
mobility=Mobility(**mobility_params))
mm = mm_func(0)
assert_allclose(mm, MM, atol=0.01, verbose=True)
mm = mm_func(111)
assert_allclose(mm, MM, atol=0.01, verbose=True)
def test_build_dynamic__with_location_mobility_data(monkeypatch):
"""
Ensure dynamic mixing matrix can use location-based mobility data set by the user + Google.
"""
def get_fake_mobility_data(*args, **kwargs):
vals = {"work": [1, 1.5, 1.3, 1.1]}
days = [0, 1, 2, 3]
return vals, days
monkeypatch.setattr(mobility, "get_mobility_data", get_fake_mobility_data)
monkeypatch.setattr(location_adjuster, "get_country_mixing_matrix",
_get_country_mixing_matrix)
mobility_params = {
"mixing": {
"school": {
"append": False,
"times": get_date_from_base([0, 1, 2, 3]),
"values": [1, 0.5, 0.3, 0.1],
}
},
"age_mixing": None,
"microdistancing": {},
"square_mobility_effect": False,
**UNTESTED_PARAMS,
}
mm_func = build_dynamic_mixing_matrix(
base_matrix=MM,
country=Country(iso3="AUS"),
mobility=Mobility(**mobility_params),
)
mm = mm_func(0)
assert_allclose(mm, MM, atol=0.01, verbose=True)
mm = mm_func(2)
expected_mm = MM.copy() + (0.3 - 1) * SCHOOL_MM + (1.3 - 1) * WORK_MM
assert_allclose(mm, expected_mm, atol=0.01, verbose=True)
def test_build_dynamic__with_everything(monkeypatch):
"""
Ensure dynamic mixing matrix can use:
- microdistancing set by the user
- user provided mobility data
- google mobility data
- age based mixing
"""
def get_fake_mobility_data(*args, **kwargs):
vals = {"work": [1, 1.5, 1.3, 1.1]}
days = [0, 1, 2, 3]
return vals, days
monkeypatch.setattr(mobility, "get_mobility_data", get_fake_mobility_data)
monkeypatch.setattr(location_adjuster, "get_country_mixing_matrix",
_get_country_mixing_matrix)
mobility_params = {
"mixing": {
"school": {
"append": False,
"times": get_date_from_base([0, 1, 2, 3]),
"values": [1, 0.5, 0.3, 0.1],
}
},
"age_mixing": {
"0": {
"times": [0, 1, 2, 3],
"values": [0.5, 0.5, 0.5, 0.5],
},
"5": {
"times": [0, 1, 2, 3],
"values": [0.3, 0.3, 0.3, 0.3],
},
},
"microdistancing": {
"foo": {
"function_type": "empiric",
"parameters": {
"max_effect": 0.6,
"times": [0, 1, 2, 3],
"values": [0, 0.1, 0.2, 0.3],
},
"locations": ["work"],
}
},
"square_mobility_effect": True,
**UNTESTED_PARAMS,
}
mm_func = build_dynamic_mixing_matrix(
base_matrix=MM,
country=Country(iso3="AUS"),
mobility=Mobility(**mobility_params),
)
# Expect only age-based mixing to occur.
mm = mm_func(0)
expected_mm = MM.copy()
# Add adjustment of 2 to the 1st row and col for the 0-5 age bracket
expected_mm[0, :] *= 0.5
expected_mm[:, 0] *= 0.5
# Add adjustment of 3 to the 2nd row and col for the 5-10 age bracket
expected_mm[1, :] *= 0.3
expected_mm[:, 1] *= 0.3
assert_allclose(mm, expected_mm, atol=0.01, verbose=True)
# Expect age based, microdistancing, google mobility and user-specified mobility to be used.
mm = mm_func(3)
work_microdistancing = (1 - 0.3 * 0.6)**2
work_google_mobility = 1.1**2
work_factor = work_microdistancing * work_google_mobility - 1
school_factor = 0.1**2 - 1
expected_mm = MM.copy() + work_factor * WORK_MM + school_factor * SCHOOL_MM
# Add adjustment of 2 to the 1st row and col for the 0-5 age bracket
expected_mm[0, :] *= 0.5
expected_mm[:, 0] *= 0.5
# Add adjustment of 3 to the 2nd row and col for the 5-10 age bracket
expected_mm[1, :] *= 0.3
expected_mm[:, 1] *= 0.3
assert_allclose(mm, expected_mm, atol=0.01, verbose=True)