src.create_initial_states.create_group_transition_probs

Functions to work with transition matrices for assortative matching.

This module is copied from sid commit 206886a14eeb3257deb71db91aba4e7fb2385fc2.

Module Contents

Functions

create_group_transition_probs(states, assort_by, params, model_name)	Create a transition matrix for groups.
_get_transition_matrix_from_params(params, states, variable, model_name)	Extract transition matrix for one assort_by variable from params.
_create_transition_matrix_from_own_prob(own_prob, group_names=None)	Create a transition matrix.
_join_transition_matrices(trans_mats)	Join several transition matrices into one, assuming independence.
_einsum_kronecker_product(*trans_mats)	Compute a Kronecker product of multiple matrices with numpy.einsum().
_generate_einsum_signature(n_trans_prob)	Generate the signature for numpy.einsum() to compute a Kronecker product.

create_group_transition_probs(states, assort_by, params, model_name)

Create a transition matrix for groups.

Parameters

• states (pandas.DataFrame) – The states.

• assort_by (list) – List of variables that influence matching probabilities.

• params (pandas.DataFrame) – The parameters.

• model_name (str) – name of the contact model.

Returns

cum_probs (numpy.ndarray): Array of shape n_group, n_groups. cum_probs[i, j]

is the probability that an individual from group i meets someone from group j or lower.

_get_transition_matrix_from_params(params, states, variable, model_name)

Extract transition matrix for one assort_by variable from params.

Parameters

• params (pd.DataFrame) – The parameters.

• states (pd.DataFrame) – The states.

• variable (str) – Name of the assort by variable

• model_name (str) – Name of the contact model in which variable is used.

Returns

The transition matrix.

Return type

pd.DataFrame

_create_transition_matrix_from_own_prob(own_prob, group_names=None)

Create a transition matrix.

The matrix is calculated from the probability of staying inside a group, spreading the remaining probability mass uniformly across other groups.

Parameters

• own_prob (float or pd.Series) – Probability of staying inside the own group, either as scalar or as pandas.Series with one entry per group.

• group_names (list) – List group codes. Mandatory if own_group is a scalar.

Returns

Transition matrix as square DataFrame. The index

and columns are the group_names.

Return type

pd.DataFrame

Example

>>> _create_transition_matrix_from_own_prob(0.6, ["a", "b"])
     a    b
a  0.6  0.4
b  0.4  0.6

>>> op = pd.Series([0.6, 0.7], index=["a", "b"])
>>> _create_transition_matrix_from_own_prob(op)
     a    b
a  0.6  0.4
b  0.3  0.7

_join_transition_matrices(trans_mats)

Join several transition matrices into one, assuming independence.

Parameters

trans_mats (list) – List of square DataFrames. The index and columns are the group_names.

Returns

Joined transition matrix. The index and columns

are the cartesian product of all individual group names in the same order as trans_mats.

Return type

pd.DataFrame

_einsum_kronecker_product(*trans_mats)

Compute a Kronecker product of multiple matrices with numpy.einsum().

The reshape is necessary because numpy.einsum() produces a matrix with as many dimensions as transition probability matrices. Each dimension has as many values as rows or columns in the transition matrix.

The ordering of letters in the numpy.einsum() signature for the result ensure that the reshape to a two-dimensional matrix does produce the correct Kronecker product.

_generate_einsum_signature(n_trans_prob)

Generate the signature for numpy.einsum() to compute a Kronecker product.

The ordering of the letters in the result is necessary so that the reshape to a square matrix does not fail.

Example

>>> _generate_einsum_signature(2)
'ab, cd -> acbd'

>>> _generate_einsum_signature(3)
'ab, cd, ef -> acebdf'

>>> _generate_einsum_signature(4)
'ab, cd, ef, gh -> acegbdfh'