guts-base 2.0.0b0__py3-none-any.whl

guts_base/prob/conditional_binom_mv.py

@@ -0,0 +1,233 @@
+from scipy.stats._multivariate import multinomial_frozen, multi_rv_generic
+from scipy.stats._discrete_distns import binom
+import numpy as np
+
+# TODO: Update example
+# TODO: Use forward filling again to deal with nans (these were automatically solved by using rv_discrete)
+# TODO: When all is complete delete conditional_bionm.py
+
+def conditional_prob_neglogexp(p, p_init=1.0, eps=1e-12):
+    p_ = np.column_stack([p_init, p])
+    # p needs to be clipped zero division does not occurr, if the last p values are zero
+    p_clipped = np.clip(p_, eps, 1.0)
+    # convert to logscale
+    neg_log_p = -np.log(p_clipped)
+    # exponent substraction is numerically more stable than division
+    return np.exp(neg_log_p[:, :-1] - neg_log_p[:, 1:])
+    
+
+def conditional_prob(p, p_init=1.0):
+    p_ = np.concatenate([p_init, p[:]])
+    # divide later though previous probability
+    return p_[1:] / p_[:-1]
+    
+
+def conditional_prob_from_neglogp(p, p_init=1.0):
+    p_ = np.concatenate([[p_init], p[:]])
+    return np.exp(p_[:-1] - p_[1:])
+    
+
+class conditional_binomial(multi_rv_generic):
+    """
+    A scipy distribution for a conditional survival probability distribution
+
+    Parameters
+    ----------
+    k: Array
+        Number of repeated positive observations of a quantity that 
+        can only decrease. k must be monotonically decreasing (e.g. survivors)
+    
+    p: Array
+        survival function of the repeated observation. 
+        p must be monotonically decreasing
+
+    n: int
+        The starting number of positive observations (e.g. initial number of organisms
+        in a survival trial)
+
+    Example
+    -------
+    Define a survival function (using a beta cdf and use it to make multinomial draws)
+    to simulate survivals from repeated observations
+    >>> n = 100
+    >>> B = stats.beta(5, 5)
+    >>> p = 1 - B.cdf(np.linspace(0, 1))
+    >>> s = stats.multinomial(n, p=np.diff(p)*-1).rvs()[0]
+    >>> s = n - s.cumsum()
+
+    construct a frozen distribution
+    >>> from guts_base.prob import conditional_survival
+    >>> S = conditional_survival(n=n, p=p[1:])
+
+    Compute the pmf
+    >>> S.pmf(s)
+
+    Compute the logpmf
+    >>> S.logpmf(s)
+
+    Draw random samples
+    >>> samples = S.rvs(size=(1000, 49))
+
+    Plot the observational variation of a given survival function under repeated
+    Observations
+    >>> plt.plot(samples.T, color="black", alpha=.02)
+
+
+    """
+    p_init = 1.0
+    eps = 1e-12
+    shapes = ("n", "p")
+    def __call__(self, n, p, seed=None):
+        """Create a frozen multinomial distribution.
+
+        See `multinomial_frozen` for more information.
+        """
+        return conditional_binomial_frozen(n, p, seed)
+    
+    def _preprocess_params(self, x, n, p):
+        x = np.array(x, ndmin=2)
+        p = np.array(p, ndmin=2)
+        p_init = np.broadcast_to(self.p_init, p[:, [0]].shape)
+        n_init = np.broadcast_to(n, x[:,[0]].shape)
+        # nan filling is not necessary, because nans are thrown out this shifts the 
+        # p vector to where it belongs
+        n_ = np.column_stack([n_init[:,[0]], x[:, :-1]])
+        p_conditional = conditional_prob_neglogexp(p, p_init=p_init, eps=self.eps)
+        return x, n_, p_conditional
+
+    def logpmf(self, x, n, p):
+        x, n, p = self._preprocess_params(x, n, p)
+        return binom._logpmf(x, n, p)
+
+    def pmf(self, x, n, p):
+        x, n, p = self._preprocess_params(x, n, p)
+        return binom._pmf(x, n, p)
+
+    def rvs(self, n, p, size=None, random_state=None):
+        r"""
+        Generate random samples from the conditional binomial distribution.
+
+        The random generation process is described by the following equations:
+        
+        The lethality matrix is initialized with
+        \[
+        L \leftarrow \mathbf{0} \in \mathbb{N}_0^{K \times T}
+        \]
+
+        where $K$ is the number of treatments and $T$ is the number of observations.
+        For each observation $t$ in \( t = 1, \ldots, T \):
+        \[
+        L_{k,t} \sim \text{Binomial}\left(n=N_{k} - \sum_{t'=0}^{t-1} L_{k,t'},~p=1 - \frac{S_{k,t}}{S_{k,t-1}}\right)
+        \]
+        Where:
+        - \( L_{k,t} \) is the number of organisms dying between $t$ and $t-1$ in each treatment $k$.
+        - \( N_{k} \) is the initial number of alive organism in each treatment \( k \) at the start of the observations $t=0$.
+        - \( S_{k,t} \) is survival function (the probability of being alive at a given time) computed for each observations $t$ in $T$.
+
+        Finally:
+        \[
+        S_{k,t}^{obs} = N_{k} - \sum_{t'=0}^{t} L_{k,t'}
+        \]
+
+
+        Parameters
+        ----------
+        n: int or Array of int
+            The initial number of positive observations (e.g. initial number of organisms
+            in a survival trial).
+        p: Array
+            Survival function of the repeated observation.
+            p must be monotonically decreasing.
+        size: tuple, optional
+            Shape of the random variates to generate.
+        random_state: RandomState or int, optional
+            If seed is not None, it will be used by the RandomState to generate
+            random time steps.
+
+        Returns
+        -------
+        ndarray
+            Random samples from the conditional binomial distribution representing the
+            number of entities surviving at each time step.
+
+
+        """
+        random_state = self._get_random_state(random_state)
+
+        p = np.array(p, ndmin=2)
+        p_init = np.broadcast_to(self.p_init, p[:, [0]].shape)
+        n_init = np.broadcast_to(n, p.shape)
+
+        p_conditional = conditional_prob_neglogexp(p, p_init=p_init, eps=self.eps)
+    
+        if size is None:
+            size = p.shape
+
+        # axis-0 is the batch dimension
+        # axis-1 is the time dimension (probability)
+        L = np.zeros(shape=size)
+
+        for i in range(L.shape[1]):
+            # calculate the binomial response of the conditional survival
+            # i.e. the probability to die within an interval conditional on 
+            # having survived until the beginning of that interval
+            L[..., i] = random_state.binomial(
+                p=1 - p_conditional[:, i], 
+                n=n_init[:, i]-L.sum(axis=-1).astype(int),
+                size=size[slice(len(size)-1)]
+            )
+
+        return n_init-L.cumsum(axis=-1)
+    
+conditional_survival = conditional_binomial()
+
+
+class conditional_binomial_frozen(multinomial_frozen):
+
+    def __init__(self, n, p, seed=None):
+        self._dist = conditional_binomial(seed)
+        self.n, self.p = n, p
+        # self.n, self.p, self.npcond = self._dist._process_parameters(n, p)
+
+        # # monkey patch self._dist
+        # def _process_parameters(n, p):
+        #     return self.n, self.p, self.npcond
+
+        # self._dist._process_parameters = _process_parameters
+
+    def logpmf(self, x):
+        return self._dist.logpmf(x, n=self.n, p=self.p)
+
+    def pmf(self, x):
+        return self._dist.pmf(x, n=self.n, p=self.p)
+
+    def rvs(self, n=None, size=None, random_state=None):
+        if n is None:
+            n = self.n
+        return self._dist.rvs(n=self.n, p=self.p, size=size, random_state=random_state)
+
+
+if __name__ == "__main__":
+
+    S = conditional_survival(n=10, p=[[0.8,0.4,0.2],[0.8,0.4,0.2]])
+
+
+    prob = S.pmf([[10,5,2], [7,2,0]])
+    prob = S.pmf([10,5,2])
+    sample = S.rvs()
+    sample = S.rvs(size=(10, 2, 3))
+
+
+    S = conditional_survival(n=10, p=[0.8,0.4,0.2])
+    prob = S.pmf([[10,5,2], [7,2,0]])
+    prob = S.pmf([10,5,2])
+    S.rvs(size=(10,3))
+    S.rvs()
+
+    sample
+
+    S = conditional_survival(n=[[10],[2000]], p=[[0.8,0.4,0.2],[0.8,0.4,0.2]])
+    prob = S.pmf([[10,5,2], [7,2,0]])
+    S.rvs()
+
+

guts_base/prob/predictions.py

@@ -0,0 +1,164 @@
+from functools import partial
+import numpy as np
+import jax
+import jax.numpy as jnp
+import numpyro
+from numpyro.infer import Predictive
+from typing import Literal
+
+
+def survival_predictions(
+    probs, n_trials, 
+    eps: float = 0.0,
+    seed=1, 
+    mode: Literal["survival", "lethality", "deaths"] = "survival"
+):
+    """Generate predictions for survival based on a multinomial survival distribution
+
+    Parameters
+    ----------
+    probs : ArrayLike
+        2D Array denoting the multinomial probabilities of deaths for each time interval 
+        per experiment 
+        dims=(experiment, time)
+    n_trials : ArrayLike
+        1D Array denoting the number of organisms at the beginning in each experiment
+        dims = (experiment,)
+    seed : int, optional
+        Seed for the random number generator, by default 1
+    mode : str, optional
+        How should the random draws be returned? 
+        - survival: Decreasing from n_trials to 0
+        - lethality: Increasing from 0 to n_trials
+        - deaths: Between 0 and n_trials in each interval. Summing to n_trials
+    """
+
+    def survival_to_death_probs(pr_survival):
+        # truncate here, because numeric errors below the solver tolerance can
+        # lead to negative values in the difference. This needs to be cured here
+        pr_survival_ = np.trunc(pr_survival / eps) * eps
+        pr_death = pr_survival_[:-1] - pr_survival_[1:]
+
+        pr_death = np.concatenate([
+            # concatenate a zero at the beginning in order to "simulate" no
+            # deaths at T = 0
+            jnp.zeros((1,)),
+            # Delta S
+            pr_death,
+            # The remaining mortility as T -> infinity 
+            jnp.ones((1,))-pr_death.sum()
+        ])
+
+        # make sure the vector is not zero or 1 (this is always problematic for 
+        # probabilities) and make sure the vector sums to 1
+        pr_death = np.clip(pr_death, eps, 1-eps)
+        pr_death = pr_death / pr_death.sum()
+        return pr_death
+
+    rng = np.random.default_rng(seed)
+    deaths = jnp.array(list(map(
+        lambda n, pr_survival: rng.multinomial(
+            n=n, pvals=survival_to_death_probs(pr_survival)
+        ), 
+        n_trials, 
+        probs
+    )))
+
+    # remove the last observations to trim off the simulated unobserved mortality
+    deaths = deaths[:, :-1]
+
+    if mode == "deaths":
+        return deaths
+    elif mode == "lethality":
+        return deaths.cumsum(axis=1)
+    elif mode == "survival":
+        return np.expand_dims(n_trials, axis=1) - deaths.cumsum(axis=1)
+    else:
+        raise NotImplementedError(
+            f"Mode {mode} is not implemented."+
+            "Use one of 'survival', 'lethality', or 'deaths'."
+        )
+    
+def posterior_predictions(sim, idata, seed=None):
+    """Make posterior predictions for survival data"""
+    if seed is None:
+        seed = sim.config.simulation.seed
+
+    n = idata.posterior.dims["draw"]
+    c = idata.posterior.dims["chain"]
+
+    key = jax.random.PRNGKey(seed)
+
+    obs, masks = sim.inferer.observation_parser()
+
+    model_kwargs = sim.inferer.preprocessing(obs=obs, masks=masks)
+    if sim.config.inference_numpyro.user_defined_error_model is None:
+        model_kwargs["obs"]["survival"] = None
+
+    # prepare model
+    model = partial(
+        sim.inferer.inference_model, 
+        solver=sim.inferer.evaluator, 
+        **model_kwargs
+    )    
+
+    posterior_samples = {
+        k: np.array(v["data"]) for k, v 
+        in idata.unconstrained_posterior.to_dict()["data_vars"].items()
+    }
+
+    predictive = Predictive(
+        model, posterior_samples=posterior_samples,
+        num_samples=n, batch_ndims=2
+    )
+
+    samples = predictive(key)
+
+    if sim.config.inference_numpyro.user_defined_error_model is not None:
+        chains = []
+        for i in range(c):
+            # TODO: Switch to vmap and jit, but it did not work, so if you do it TEST IT!!!!!!!
+            predictions = list(map(
+                partial(
+                    survival_predictions,
+                    n_trials=obs["survival"][:, 0].astype(int),
+                    eps=obs["eps"],
+                    seed=seed,
+                    mode="survival",
+                ),
+                samples["survival"][i]
+            ))
+            chains.append(predictions)
+
+        posterior_predictive = {"survival_obs": np.array(chains)}
+    else:
+        posterior_predictive = {"survival_obs": samples.pop("survival_obs")}
+
+
+    new_idata = sim.inferer.to_arviz_idata(
+        posterior=samples,
+        posterior_predictive=posterior_predictive,
+        n_draws=n,
+        n_chains=c
+    )
+
+    # update chain names in case they were subselected (clustering)
+    # new_idata = new_idata.assign_coords({"chain": idata.posterior.chain.values})
+    new_idata = new_idata.assign_coords({"chain": idata.posterior.chain.values})
+
+    # assert the new posterior matches the old posterior
+    tol = sim.config.jaxsolver.atol * 100
+    abs_diff_posterior = np.abs(idata.posterior - new_idata.posterior)
+    np.testing.assert_array_less(abs_diff_posterior.mean().to_array(), tol)
+
+    fit_tol = tol * sim.coordinates["time"].max()
+    abs_diff_fits = np.abs(new_idata.posterior_model_fits - idata.posterior_model_fits)
+    np.testing.assert_array_less(abs_diff_fits.mean().to_array(), fit_tol)
+
+    idata.posterior_model_fits = new_idata.posterior_model_fits
+    idata.posterior_predictive = new_idata.posterior_predictive
+
+    return idata
+
+
+

guts_base/sim/__init__.py

@@ -0,0 +1,28 @@
+from . import base
+from . import ecx
+from . import report
+from . import utils
+from . import transformer
+from . import config
+from . import units
+
+from .base import (
+    GutsBase,
+    GutsSimulationConstantExposure,
+    GutsSimulationVariableExposure
+)
+
+from .ecx import ECxEstimator, LPxEstimator
+from .report import GutsReport, ParameterConverter
+
+from .mempy import PymobSimulator
+from .utils import (
+    GutsBaseError
+)
+
+from .constructors import (
+    construct_sim_from_config, 
+    load_idata
+)
+
+from .config import GutsBaseConfig