[Python] appropriate type hint for generic/frozen multivariate distributions in scipy.stats

For my application, I need to differentiate between 4 different sets of probability distributions from the scipy.stats module: generic univariate, frozen univariate, generic multivariate, & frozen multivariate. Throughout the application, I would like to add type hints for these 4 sets.

For the univariate cases, mypy has no gripes with type hints like in this MWE:

from typing import Any

from scipy.stats import norm
from scipy.stats._distn_infrastructure import rv_generic, rv_frozen


def sample_frozen_univariate(n_samples: int, dist: rv_frozen):
return dist.rvs(n_samples)


def sample_generic_univariate(n_samples: int, dist: rv_generic,
*distparams: Any):
return dist(*distparams).rvs(n_samples)


n_samples = 4
frozen_dist = norm()
print(sample_frozen_univariate(n_samples, frozen_dist))

n_samples = 4
generic_dist = norm
loc, scale = 1, 2
print(sample_generic_univariate(n_samples, generic_dist, loc, scale))


Now here's an MWE (working in the sense that executing the script gives the expected result) for what I've tried for the multivariate cases:

from typing import Any

from scipy.stats import multivariate_normal
from scipy.stats._multivariate import multi_rv_generic, multi_rv_frozen


def sample_frozen_multivariate(n_samples: int, n_variates: int, dist: multi_rv_frozen):
if dist.dim != n_variates:
msg = 'distribution dimension %s inconsistent with n_variates=%s'
raise ValueError(msg % (dist.dim, n_variates))
sample = dist.rvs(n_samples)
return sample


def sample_generic_multivariate(n_samples: int, n_variates: int, dist: multi_rv_generic,
*distparams: Any):
sample = dist(*distparams).rvs(n_samples)
if sample.shape[1] != n_variates:
msg = 'sample dimension %s inconsistent with n_variates=%s'
raise ValueError(msg % (sample.shape[1], n_variates))
return sample


n_samples = 4
n_variates = 2
frozen_dist = multivariate_normal(mean=[0, 0], cov=[[1, 0], [0, 1]])
print(sample_frozen_multivariate(n_samples, n_variates, frozen_dist))

n_samples = 4
n_variates = 2
generic_dist = multivariate_normal
mean, cov = [-1, 1], [[1, 0], [0, 1]]
print(sample_generic_multivariate(n_samples, n_variates, generic_dist, mean, cov))


However, mypy reports the following issues:

example_2.py:8: error: "multi_rv_frozen[multi_rv_generic]" has no attribute "dim" [attr-defined]
example_2.py:10: error: "multi_rv_frozen[multi_rv_generic]" has no attribute "dim" [attr-defined]
example_2.py:11: error: "multi_rv_frozen[multi_rv_generic]" has no attribute "rvs" [attr-defined]
example_2.py:17: error: "multi_rv_generic" not callable [operator]
Found 4 errors in 1 file (checked 1 source file)


I understand the issues that mypy is reporting, just not how to properly address them. What approach(es) to type hints for generic and frozen multivariate distributions from scipy.stats would appease mypy? (As an aside, I am aware of the problem with relying on types from "private" modules in scipy.stats. Though I understand that is a separate issue, I welcome solutions that simultaneously address that problem.) Thanks in advance.

Continue reading...