From afb93a914bd6c1f1331cc75173091c955b427e46 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?upx3=E2=80=94TM=20=28CFA=29?= <127630341+AFg6K7h4fhy2@users.noreply.github.com> Date: Mon, 22 Apr 2024 12:20:19 -0400 Subject: [PATCH] First Pass Documentation (docstrings) For Top Level Pyrenew Files (#89) * numpydoc changes (testing) * numpydoc changes (testing) * poetry lock and precommit for top-level docs * docs for convolve; poetry lock; ignore TODOs * docs for convolve; poetry lock; ignore TODOs * distutil * math docs * minor mcmcutils edit * update poetry lock file * top level file modifications * pre-commit edit --- model/src/pyrenew/convolve.py | 70 ++++++++++++++++++-- model/src/pyrenew/distutil.py | 38 ++++++++++- model/src/pyrenew/math.py | 53 +++++++++------ model/src/pyrenew/mcmcutils.py | 5 +- model/src/pyrenew/metaclass.py | 49 ++++++++++++-- model/src/pyrenew/regression.py | 30 ++++++++- model/src/pyrenew/transform.py | 111 +++++++++++++++++++++++++++++--- 7 files changed, 311 insertions(+), 45 deletions(-) diff --git a/model/src/pyrenew/convolve.py b/model/src/pyrenew/convolve.py index 6c894c20..2dc869c9 100755 --- a/model/src/pyrenew/convolve.py +++ b/model/src/pyrenew/convolve.py @@ -10,12 +10,44 @@ jax.lax.scan. Factories generate functions that can be passed to scan with an appropriate array to scan. + +Notes +----- +TODO: Look into adding blocks for Functions and Examples in this +docstring. """ +from typing import Callable, Tuple + import jax.numpy as jnp +from jax.typing import ArrayLike + + +def new_convolve_scanner(discrete_dist_flipped: ArrayLike) -> Callable: + """ + Factory function to create a scanner function for convolving a discrete distribution + over a time series data subset. + + Parameters + ---------- + discrete_dist_flipped : ArrayLike + A 1D jax array representing the discrete distribution flipped for convolution. + Returns + ------- + Callable + A scanner function that can be used with jax.lax.scan for convolution. + This function takes a history subset and a multiplier, computes the dot product, + then updates and returns the new history subset and the convolution result. -def new_convolve_scanner(discrete_dist_flipped): - def _new_scanner(history_subset, multiplier): + Notes + ----- + TODO: Add Example. + TODO: Clarification on Returns description. + """ + + def _new_scanner( + history_subset: ArrayLike, multiplier: float + ) -> Tuple[ArrayLike, float]: new_val = multiplier * jnp.dot(discrete_dist_flipped, history_subset) latest = jnp.hstack([history_subset[1:], new_val]) return latest, new_val @@ -23,15 +55,43 @@ def _new_scanner(history_subset, multiplier): return _new_scanner -def new_double_scanner(dists, transforms): +def new_double_scanner( + dists: Tuple[ArrayLike, ArrayLike], transforms: Tuple[Callable, Callable] +) -> Callable: + """ + Factory function to create a scanner function that applies two sequential transformations + and convolutions using two discrete distributions. + + Parameters + ---------- + dists : Tuple[ArrayLike, ArrayLike] + A tuple of two 1D jax arrays, each representing a discrete distribution for the + two stages of convolution. + transforms : Tuple[Callable, Callable] + A tuple of two functions, each transforming the output of the dot product at each + convolution stage. + + Returns + ------- + Callable + A scanner function that applies two sequential convolutions and transformations. + It takes a history subset and a tuple of multipliers, computes the transformations + and dot products, and returns the updated history subset and a tuple of new values. + + Notes + ----- + TODO: Add Example + """ d1, d2 = dists t1, t2 = transforms - def _new_scanner(history_subset, multipliers): + def _new_scanner( + history_subset: jnp.ndarray, multipliers: Tuple[float, float] + ) -> (jnp.ndarray, Tuple[float, float]): m1, m2 = multipliers m_net1 = t1(m1 * jnp.dot(d1, history_subset)) new_val = t2(m2 * m_net1 * jnp.dot(d2, history_subset)) latest = jnp.hstack([history_subset[1:], new_val]) - return (latest, (new_val, m_net1)) + return latest, (new_val, m_net1) return _new_scanner diff --git a/model/src/pyrenew/distutil.py b/model/src/pyrenew/distutil.py index a718b56b..315b3405 100755 --- a/model/src/pyrenew/distutil.py +++ b/model/src/pyrenew/distutil.py @@ -9,15 +9,37 @@ such as discrete time-to-event distributions """ import jax.numpy as jnp +from jax.typing import ArrayLike def validate_discrete_dist_vector( - discrete_dist: jnp.ndarray, tol: float = 1e-20 -) -> bool: + discrete_dist: ArrayLike, tol: float = 1e-20 +) -> ArrayLike: """ Validate that a vector represents a discrete probability distribution to within a specified tolerance, raising a ValueError if not. + + Parameters + ---------- + discrete_dist : ArrayLike + An jax array containing non-negative values that + represent a discrete probability distribution. The values + must sum to 1 within the specified tolerance. + tol : float, optional + The tolerance within which the sum of the distribution must + be 1. Defaults to 1e-20. + + Returns + ------- + ArrayLike + The normalized distribution array if the input is valid. + + Raises + ------ + ValueError + If any value in discrete_dist is negative or if the sum of the + distribution does not equal 1 within the specified tolerance. """ discrete_dist = discrete_dist.flatten() if not jnp.all(discrete_dist >= 0): @@ -39,10 +61,20 @@ def validate_discrete_dist_vector( return discrete_dist / dist_norm -def reverse_discrete_dist_vector(dist): +def reverse_discrete_dist_vector(dist: ArrayLike) -> ArrayLike: """ Reverse a discrete distribution vector (useful for discrete time-to-event distributions). + + Parameters + ---------- + dist : ArrayLike + A discrete distribution vector (likely discrete time-to-event distribution) + + Returns + ------- + ArrayLike + A reversed (jnp.flip) discrete distribution vector """ return jnp.flip(dist) diff --git a/model/src/pyrenew/math.py b/model/src/pyrenew/math.py index 4a4d59de..91b35284 100755 --- a/model/src/pyrenew/math.py +++ b/model/src/pyrenew/math.py @@ -6,11 +6,13 @@ a given renewal process. """ + import jax.numpy as jnp +from jax.typing import ArrayLike from pyrenew.distutil import validate_discrete_dist_vector -def get_leslie_matrix(R, generation_interval_pmf): +def get_leslie_matrix(R: float, generation_interval_pmf: ArrayLike) -> float: """ Create the Leslie matrix corresponding to a basic @@ -28,7 +30,7 @@ def get_leslie_matrix(R, generation_interval_pmf): mass vector of the renewal process Returns - -------- + ------- The Leslie matrix for the renewal process, as a jax array. """ @@ -44,7 +46,9 @@ def get_leslie_matrix(R, generation_interval_pmf): return jnp.vstack([R * generation_interval_pmf, aging_matrix]) -def get_asymptotic_growth_rate_and_age_dist(R, generation_interval_pmf): +def get_asymptotic_growth_rate_and_age_dist( + R: float, generation_interval_pmf: ArrayLike +) -> tuple[float, ArrayLike]: """ Get the asymptotic per-timestep growth rate of the renewal process (the dominant @@ -52,6 +56,7 @@ def get_asymptotic_growth_rate_and_age_dist(R, generation_interval_pmf): associated stable age distribution (a normalized eigenvector associated to that eigenvalue). + Parameters ---------- R : float @@ -61,11 +66,17 @@ def get_asymptotic_growth_rate_and_age_dist(R, generation_interval_pmf): mass vector of the renewal process Returns - -------- - A tuple consisting of the asymptotic growth rate of - the process, as jax float, and the stable age distribution - of the process, as a jax array probability vector of the - same shape as the generation interval probability vector. + ------- + tuple[float, ArrayLike] + A tuple consisting of the asymptotic growth rate of + the process, as jax float, and the stable age distribution + of the process, as a jax array probability vector of the + same shape as the generation interval probability vector. + + Raises + ------ + ValueError + If an age distribution vector with non-zero imaginary part is produced. """ L = get_leslie_matrix(R, generation_interval_pmf) eigenvals, eigenvecs = jnp.linalg.eig(L) @@ -92,7 +103,9 @@ def get_asymptotic_growth_rate_and_age_dist(R, generation_interval_pmf): return d_val_real, d_vec_norm -def get_stable_age_distribution(R, generation_interval_pmf): +def get_stable_age_distribution( + R: float, generation_interval_pmf: ArrayLike +) -> ArrayLike: """ Get the stable age distribution for a renewal process with a given value of @@ -114,18 +127,21 @@ def get_stable_age_distribution(R, generation_interval_pmf): mass vector of the renewal process Returns - -------- - The stable age distribution for the - process, as a jax array probability vector of - the same shape as the generation interval - probability vector. + ------- + ArrayLike + The stable age distribution for the + process, as a jax array probability vector of + the same shape as the generation interval + probability vector. """ return get_asymptotic_growth_rate_and_age_dist(R, generation_interval_pmf)[ 1 ] -def get_asymptotic_growth_rate(R, generation_interval_pmf): +def get_asymptotic_growth_rate( + R: float, generation_interval_pmf: ArrayLike +) -> float: """ Get the asymptotic per timestep growth rate for a renewal process with a given value of @@ -145,9 +161,10 @@ def get_asymptotic_growth_rate(R, generation_interval_pmf): mass vector of the renewal process Returns - -------- - The asymptotic growth rate of the renewal process, - as a jax float. + ------- + float + The asymptotic growth rate of the renewal process, + as a jax float. """ return get_asymptotic_growth_rate_and_age_dist(R, generation_interval_pmf)[ 0 diff --git a/model/src/pyrenew/mcmcutils.py b/model/src/pyrenew/mcmcutils.py index 3b081958..733d7c75 100644 --- a/model/src/pyrenew/mcmcutils.py +++ b/model/src/pyrenew/mcmcutils.py @@ -10,7 +10,7 @@ def spread_draws( posteriors: dict, - variables_names: list, + variables_names: list[str] | list[tuple], ) -> pl.DataFrame: """Get nicely shaped draws from the posterior @@ -29,7 +29,8 @@ def spread_draws( Returns ------- - polars.DataFrame + pl.DataFrame + A dataframe of draw-indexed """ for i_var, v in enumerate(variables_names): diff --git a/model/src/pyrenew/metaclass.py b/model/src/pyrenew/metaclass.py index d00c5ace..5b51584d 100644 --- a/model/src/pyrenew/metaclass.py +++ b/model/src/pyrenew/metaclass.py @@ -17,19 +17,26 @@ def _assert_sample_and_rtype( ) -> None: """Return type-checking for RandomVariable's sample function - Objects passed as `RandomVariable` should (a) have a `sample()` method that + Objects passed as `RandomVariable` should (a) have a sample() method that (b) returns either a tuple or a named tuple. Parameters ---------- rp : RandomVariable Random variable to check. - skip_if_none: bool - When `True` it returns if `rp` is None. + skip_if_none: bool, optional + When `True` it returns if `rp` is None. Defaults to True Returns ------- None + + Raises + ------ + Exception + If rp is not a RandomVariable, does not have a sample function, or + does not return a tuple. Also occurs if rettype does not initialized + properly. """ # Addressing the None case @@ -101,7 +108,7 @@ def sample( ---------- **kwargs : dict, optional Additional keyword arguments passed through to internal `sample()` - calls, if any + calls, should there be any. Notes ----- @@ -117,6 +124,9 @@ def sample( @staticmethod @abstractmethod def validate(**kwargs) -> None: + """ + Validation of kwargs to be implemented in subclasses. + """ pass @@ -149,7 +159,7 @@ def sample( ---------- **kwargs : dict, optional Additional keyword arguments passed through to internal `sample()` - calls, if any + calls, should there be any. Notes ----- @@ -244,9 +254,34 @@ def print_summary( prob: float = 0.9, exclude_deterministic: bool = True, ) -> None: - """A wrapper of MCMC.print_summary""" + """ + A wrapper of MCMC.print_summary + + Parameters + ---------- + prob : float, optional + The acceptance probability of print_summary. Defaults to 0.9 + exclude_deterministic : bool, optional. + Whether to print deterministic variables in the summary. + Defaults to True. + + Returns + ------- + None + """ return self.mcmc.print_summary(prob, exclude_deterministic) def spread_draws(self, variables_names: list) -> pl.DataFrame: - """A wrapper of mcmcutils.spread_draws""" + """A wrapper of mcmcutils.spread_draws + + Parameters + ---------- + variables_names : list + A list of variable names to create a table of samples. + + Returns + ------- + pl.DataFrame + """ + return spread_draws(self.mcmc.get_samples(), variables_names) diff --git a/model/src/pyrenew/regression.py b/model/src/pyrenew/regression.py index 67772bc6..d3e6789e 100755 --- a/model/src/pyrenew/regression.py +++ b/model/src/pyrenew/regression.py @@ -78,11 +78,11 @@ def __init__( If `None`, use an identity transform. Default `None`. - intercept_suffix : str + intercept_suffix : str, optional Suffix for naming the intercept random variable in class to numpyro.sample(). Default `"_intercept"`. - coefficient_suffix : str + coefficient_suffix : str, optional Suffix for naming the regression coefficient random variables in calls to numpyro.sample(). Default `"_coefficients"`. @@ -99,12 +99,38 @@ class to numpyro.sample(). Default `"_intercept"`. self.coefficient_suffix = coefficient_suffix def predict(self, intercept, coefficients): + """ + Generates a transformed prediction w/ intercept, coefficients, and + fixed predictor values + + Parameters + ---------- + intercept : ArrayLike + Sampled numpyro distribution generated from intercept priors. + coefficients : ArrayLike + Sampled prediction coefficients distribution generated + from coefficients priors. + + Returns + ------- + ArrayLike + Array of transformed predictions. + """ transformed_prediction = ( intercept + self.fixed_predictor_values @ coefficients ) return self.transform.inverse(transformed_prediction) def sample(self): + """ + Sample generalized linear model + + Returns + ------- + dict + A dictionary containing transformed predictions, and + the intercept and coefficients sample distributions. + """ intercept = numpyro.sample( self.name + self.intercept_suffix, self.intercept_prior ) diff --git a/model/src/pyrenew/transform.py b/model/src/pyrenew/transform.py index 0701eb2f..bca9a792 100755 --- a/model/src/pyrenew/transform.py +++ b/model/src/pyrenew/transform.py @@ -7,6 +7,7 @@ import jax import jax.numpy as jnp +from jax.typing import ArrayLike class AbstractTransform(metaclass=ABCMeta): @@ -19,10 +20,16 @@ def __call__(self, x): @abstractmethod def transform(self, x): + """ + Transform generated predictions + """ pass @abstractmethod def inverse(self, x): + """ + Take the inverse of transformed predictions + """ pass @@ -35,10 +42,32 @@ class IdentityTransform(AbstractTransform): f^-1(x) = x """ - def transform(self, x): + def transform(self, x: any): + """ + Parameters + ---------- + x : any + Input, usually ArrayLike + + Returns + ------- + any + The same object that was inputted. + """ return x - def inverse(self, x): + def inverse(self, x: any): + """ + Parameters + ---------- + x : any + Input, usually ArrayLike + + Returns + ------- + any + The same object that was inputted. + """ return x @@ -51,10 +80,32 @@ class LogTransform(AbstractTransform): f^-1(x) = exp(x) """ - def transform(self, x): + def transform(self, x: ArrayLike): + """ + Parameters + ---------- + x : ArrayLike + Input, usually predictions array.. + + Returns + ------- + ArrayLike + Log-transformed input + """ return jnp.log(x) - def inverse(self, x): + def inverse(self, x: ArrayLike): + """ + Parameters + ---------- + x : ArrayLike + Input, usually log-scale predictions array. + + Returns + ------- + ArrayLike + Exponentiated input + """ return jnp.exp(x) @@ -68,10 +119,32 @@ class LogitTransform(AbstractTransform): f^-1(x) = 1 / (1 + exp(-x)) """ - def transform(self, x): + def transform(self, x: ArrayLike): + """ + Parameters + ---------- + x : ArrayLike + Input, usually predictions array. + + Returns + ------- + ArrayLike + Logit transformed input. + """ return jax.scipy.special.logit(x) - def inverse(self, x): + def inverse(self, x: ArrayLike): + """ + Parameters + ---------- + x : ArrayLike + Input, usually logit-transformed predictions array. + + Returns + ------- + ArrayLike + Inversed logit transformed input. + """ return jax.scipy.special.expit(x) @@ -97,8 +170,30 @@ def __init__(self, x_max: float): """ self.x_max = x_max - def transform(self, x): + def transform(self, x: ArrayLike): + """ + Parameters + ---------- + x : ArrayLike + Input, usually predictions array. + + Returns + ------- + ArrayLike + x_max scaled logit transformed input. + """ return jax.scipy.special.logit(x / self.x_max) - def inverse(self, x): + def inverse(self, x: ArrayLike): + """ + Parameters + ---------- + x : ArrayLike + Input, usually scaled logit predictions array. + + Returns + ------- + ArrayLike + Inverse of x_max scaled logit transformed input. + """ return self.x_max * jax.scipy.special.expit(x)