guts-base 0.8.6__py3-none-any.whl → 1.0.0__py3-none-any.whl

guts_base/__init__.py

@@ -4,11 +4,12 @@ from . import data
 from . import prob
 from . import plot
 
-__version__ = "0.8.6"
+__version__ = "1.0.0"
 
 from .sim import (
     GutsBase,
     PymobSimulator,
     ECxEstimator,
     LPxEstimator, 
+    GutsBaseError,
 )

guts_base/data/__init__.py

@@ -21,7 +21,7 @@ from .survival import (
     generate_survival_repeated_observations
 )
 
-from .generator import create_artificial_data, design_exposure_scenario
+from .generator import create_artificial_data, design_exposure_scenario, ExposureDataDict
 
 from .expydb import (
     to_dataset,

guts_base/data/generator.py

@@ -11,7 +11,8 @@ def create_artificial_data(
     t_max, 
     dt, 
     exposure_paths=["oral", "topical", "contact"],
-    intensity=[0.1, 0.5, 0.05]
+    intensity=[0.1, 0.5, 0.05],
+    seed=1,
 ):
     rng = np.random.default_rng(1)
     time = np.arange(0, t_max, step=dt)  # daily time resolution

guts_base/data/survival.py

@@ -40,6 +40,12 @@ def prepare_survival_data_for_conditional_binomial(observations: xr.Dataset) ->
             nsurv.isel(time=list(range(0, len(nsurv.time)-1))).values
     ]).astype(int))})
 
+    observations = observations.assign_coords({
+        "survivors_at_start": (("id", "time"), np.broadcast_to(
+            nsurv.isel(time=0).values.reshape(-1,1), 
+            shape=nsurv.shape
+    ).astype(int))})
+
     return observations
 
 

guts_base/mod.py

@@ -1,10 +1,3 @@
-"""
-GUTS models
-
-TODO: Import guts models from mempy and update to work well with jax
-TODO: Based on this implement the bufferguts model in mempy and import in the
-      bufferguts case-study
-"""
 from functools import partial
 from pathlib import Path
 
@@ -19,79 +12,37 @@ from pymob.solvers.symbolic import (
     PiecewiseSymbolicODESolver, FunctionPythonCode, get_return_arguments, dX_dt2X
 )
 
-from mempy.model import (
-    RED_IT, 
-    RED_SD,
-    RED_IT_DA,
-    RED_SD_DA,
-    RED_IT_IA,
-    RED_SD_IA,
-    BufferGUTS_IT,
-    BufferGUTS_IT_CA,
-    BufferGUTS_IT_DA
-)
+class RED_SD:
+    """Simplest guts model, mainly for testing"""
+    @staticmethod
+    def _rhs_jax(t, y, x_in, kd, b, m, hb):
+        D, H = y
+        dD_dt = kd * (x_in.evaluate(t) - D)
+        dH_dt = b * jnp.maximum(0.0, D - m) + hb
+        return dD_dt, dH_dt
+
+    @staticmethod
+    def _solver_post_processing(results, t, interpolation):
+        results["survival"] = jnp.exp(-results["H"])
+        results["exposure"] = jax.vmap(interpolation.evaluate)(t)
+        return results
 
 red_sd = RED_SD._rhs_jax
 red_sd_post_processing = RED_SD._solver_post_processing
 
-red_it = RED_IT._rhs_jax
-red_it_post_processing = RED_IT._solver_post_processing
+class RED_SD_DA(RED_SD): 
+    @staticmethod
+    def _rhs_jax(t, y, x_in, kd, w, b, m, hb):
+        D, H = y
+        dD_dt = kd * (x_in.evaluate(t) - D)
+        dH_dt = b * jnp.maximum(0.0, jnp.sum(w * D) - m) + hb
+        return dD_dt, dH_dt
 
-red_sd_ia = RED_SD_IA._rhs_jax
-red_sd_ia_post_processing = RED_SD_IA._solver_post_processing
+red_sd_da = RED_SD_DA._rhs_jax
+red_sd_da_post_processing = RED_SD_DA._solver_post_processing
 
 
-def p_survival(results, t, interpolation, z, k_k, h_b):
-    """Computes the stochastic death survival probability after computing
-    
-    """
-    # calculate survival 
-    results["exposure"] = jax.vmap(interpolation.evaluate)(t)
-    p_surv = survival_jax(t, results["D"], z, k_k, h_b)
-    results["survival"] = p_surv
-    results["lethality"] = 1 - p_surv
-    return results
-
-def it_post_processing(results, t, interpolation, alpha, beta, h_b, eps):
-    results["exposure"] = jax.vmap(interpolation.evaluate)(t)
-    p_surv = survival_IT_jax(t, results["D"], alpha, beta, h_b, eps)
-    results["survival"] = p_surv
-    results["H"] = - jnp.log(p_surv)
-    return results
-
-def post_exposure(results, t, interpolation):
-    results["survival"] = jnp.exp(-results["H"])
-    results["exposure"] = jax.vmap(interpolation.evaluate)(t)
-    return results
-
-
-def no_post_processing(results):
-    return results
-
-
-@jax.jit
-def survival_jax(t, damage, z, kk, h_b):
-    """
-    survival probability derived from hazard 
-    first calculate cumulative Hazard by integrating hazard cumulatively over t
-    then calculate the resulting survival probability
-    It was checked that `survival_jax` behaves exactly the same as `survival`
-    """
-    hazard = kk * jnp.where(damage - z < 0, 0, damage - z) + h_b
-    H = jnp.array([jax.scipy.integrate.trapezoid(hazard[:i+1], t[:i+1], axis=0) for i in range(len(t))])
-    # H = jnp.array([jnp.trapz(hazard[:i+1], t[:i+1], axis=0) for i in range(len(t))])
-    S = jnp.exp(-H)
-
-    return S
-
-@jax.jit
-def survival_IT_jax(t, damage, alpha, beta, h_b, eps):
-    d_max = jnp.squeeze(jnp.array([jnp.max(damage[:i+1])+eps for i in range(len(t))]))
-    F = jnp.where(d_max > 0, 1.0 / (1.0 + (d_max / alpha) ** -beta), 0)
-    S = 1.0 * (jnp.array([1.0], dtype=float) - F) * jnp.exp(-h_b * t)
-    return S
-
-def guts_jax(t, y, C_0, k_d, z, b, h_b):
+def guts_constant_exposure(t, y, C_0, k_d, z, b, h_b):
     # for constant exposure
     D, H, S = y
     dD_dt = k_d * (C_0 - D)
@@ -103,16 +54,6 @@ def guts_jax(t, y, C_0, k_d, z, b, h_b):
 
     return dD_dt, dH_dt, dS_dt
 
-
-def RED_IT(t, y, x_in, kd):
-    D, = y
-    C = x_in.evaluate(t)
-
-    dD_dt = kd * (C - D)
-
-    return (dD_dt, )
-    
-
 def guts_variable_exposure(t, y, x_in, k_d, z, b, h_b):
     # for constant exposure
     D, H, S = y

guts_base/prob.py

@@ -2,270 +2,10 @@ 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
 
 
-maximum = jnp.frompyfunc(jnp.maximum, nin=2, nout=1, identity=None) # type: ignore
-
-
-@jax.jit
-def ffill_na(x, mask):
-    """Forward-fill nan values in x. If a mask is provided, assume 
-
-    Parameters
-    ----------
-    x : _type_
-        _description_
-    mask : _type_
-        _description_
-
-    Returns
-    -------
-    _type_
-        _description_
-    """
-    if mask is None:
-        mask = jnp.logical_not(jnp.isnan(x))
-    idx = jnp.where(mask,jnp.arange(mask.shape[1]),0)
-    idx_ = maximum.accumulate(jnp.array(idx), axis=1).astype(int)
-    return x[jnp.arange(idx.shape[0])[:,None], idx_]
-
-
-def conditional_survival_from_hazard(x, mask):
-    """Calculates the conditional survival from cumulative hazard values. 
-    This equation is used when survival is repeatedly observed over time and
-
-    Parameters
-    ----------
-    x : np.ndarray[I,T, float]
-        A 2-dimensional I x T array of cumulative hazards defined as H = -ln(S).
-        I is the batch dimension and T is the time dimension
-
-    mask : np.ndarray[I,T, bool]
-        A 2-dimensional array of the same shape as x, taking True if the survival
-        was observed for the given index (i,t) and taking False if survival was
-        not observed for the given index (i,t).
-    
-    Returns
-    -------
-    out : np.ndarray[I,T, float]
-        A matrix with conditional probabilities and nans in place where the 
-        mask has nans. Output has the same shape as input.
-
-    Example
-    -------
-    Calculation example from survival probabilities to conditional survival
-    probabilities given some masked values.
-    
-    >>> S_i = np.array([
-    >>>     [1.  , 0.75, 0.1 , 0.05, 0.01, 0.  ],
-    >>>     [1.  , 0.75, 0.1 , 0.05, 0.01, 0.  ],
-    >>>     [1.  , 0.75, 0.1 , 0.05, 0.01, 0.  ],
-    >>>     [1.  , 0.75, 0.1 , 0.05, 0.01, 0.  ],
-    >>>     [1.  , 0.75, 0.1 , 0.05, 0.01, 0.  ],
-    >>> ])
-
-    >>> mask_obs = np.array([
-    >>>     [ True,  True,  True,  True,  True,  True],
-    >>>     [ True, False,  True,  True,  True,  True],
-    >>>     [ True, False, False,  True,  True,  True],
-    >>>     [ True,  True, False,  True,  True,  True],
-    >>>     [ True, False,  True, False,  True,  True],
-    >>> ])
-
-    >>> conditional_survival_from_hazard(-jnp.log(S_i), mask_obs)
-    array([
-       [1.0  0.75  0.133  0.5     0.2  0.0]
-       [1.0   nan  0.1    0.5     0.2  0.0]
-       [1.0   nan  nan    0.05    0.2  0.0]
-       [1.0  0.75  nan    0.0667  0.2  0.0]
-       [1.0   nan  0.1    nan     0.1  0.0]
-    ])
-    """
-
-    # Append zeros (hazard) to the beginning of the array (this aligns with the
-    # safe assumption that before the zeroth observation S(t=-1) = 1.0)
-    x_ = jnp.column_stack([jnp.zeros_like(x[:, 0]), x])
-
-    # also mask needs to be expanded accordingly
-    mask_ = jnp.column_stack([jnp.ones_like(mask[:, 0]), mask])
-
-    # fill NaNs with forward
-    x_filled = ffill_na(x_, mask_)
-
-    # calculate the conditional survival.
-    conditional_survival = jnp.exp(x_filled[:, :-1] - (x_[:, 1:]))
-
-    # add nans and return
-    return jnp.where(
-        mask, conditional_survival, jnp.nan
-    )
-
-
-def conditional_survival_hazard_error_model(theta, simulation_results, observations, indices, masks, make_predictions):
-    """Computes the likelihood of observing K survivors at time t given that
-    N survivors were alive at the previous observation t-1
-
-    This is achieved by using the Cumulative hazard directly in the calculation
-    of the conditional survival probability. Note that this equation can easily
-    be adapted to compute the conditional lethality probability.
-
-    $$\\Pr(t < T~|~t_0 < T) = \\frac{e^{-H(t)}}{e^{-H(t_0)}} = e^{-H(t) - (- H(t_0))} = e^{-H(t) + H(t_0)} = e^{-(H(t) - H(t_0))} = e^{H(t_0) - H(t)}$$
-
-    Example
-    -------
-
-                          t0    t1   t2     t3  tinf 
-    observations          10     8    5      2     0
-    Pr(T > t)            1.0   0.8  0.5    0.2   0.0
-    Pr(T > t | T > t-1) (1.0)  0.8  0.625  0.4   0.0
-
-
-    If the experiment ends before the last organism has died, these information
-    have to be included in he computation of the likelihood. For multinomial
-    distributions this is done by computing the remaining (ubobserved) lethality
-    until the end of the experiment, whihc is just the inverse of the number of
-    survivors at the end of the experiment. Therefore, for conditional survival 
-    the information about the number of alive organisms
-    at the end of the experiment is contained in the last observation.      
-
-    We compute an additional interval until the start of the experiment in order
-    to align the results with the observations (which include t=0). Otherwise
-    we would only compute T-1 intervals if we have T observations (including t=0).
-    """
-
-    H = simulation_results["H"]
-    n_surv = observations["survivors_before_t"]
-    S_mask = masks["survival"]
-    obs_survival = observations["survival"]
-    S_conditional = conditional_survival_from_hazard(H, S_mask)
-
-    if make_predictions:
-        obs_survival = None
-
-    ID, T = S_mask.shape
-
-    with numpyro.plate("time", size=T):
-        with numpyro.plate("id", size=ID):
-            numpyro.sample(
-                "survival_obs", numpyro.distributions.Binomial(
-                    probs=S_conditional, 
-                    total_count=n_surv
-                ).mask(S_mask), 
-                obs=obs_survival
-            )
-
-
-def conditional_survival_error_model(theta, simulation_results, observations, indices, masks, make_predictions):
-    """Note that the conditional survival error model will not work to generate
-    posterior predictions for observational uncertainty out of the box
-    
-    This model can be used for SD and IT models, but for SD models, the 
-    conditional_survival_hazard_model is preferrable.
-    """
-    # error model
-    EPS = observations["eps"]
-    S = jnp.clip(simulation_results["survival"], EPS, 1 - EPS) 
-
-    S_mask = masks["survival"]
-    n_surv = observations["survivors_before_t"]
-    obs_survival = observations["survival"]
-    
-    S_ = ffill_na(S, S_mask)
-    S_cond_na = S_[:, 1:] / S_[:, :-1]
-    S_cond_na_ = jnp.column_stack([jnp.ones_like(S[:,[0]]), S_cond_na])
-    S_cond_na__ = jnp.clip(S_cond_na_, EPS, 1 - EPS) 
-
-    if make_predictions:
-        obs_survival = None
-
-    ID, T = S_mask.shape
-
-    with numpyro.plate("time", size=T):
-        with numpyro.plate("id", size=ID):
-            numpyro.sample(
-                "survival_obs", numpyro.distributions.Binomial(
-                    probs=S_cond_na__, 
-                    total_count=n_surv
-                ).mask(S_mask), 
-                obs=obs_survival
-            )
-
-
-def conditional_lethality_error_model(theta, simulation_results, observations, indices, masks, make_predictions):
-    raise NotImplementedError(
-        "Conditional lethality error model needs to forward fill survival " +
-        "probabilities for missing (observation) values"
-        ""
-    )
-
-    # error model
-    EPS = observations["eps"]
-    S = jnp.clip(simulation_results["S"], EPS, 1 - EPS) 
-
-    # TODO work on this
-    # Conditional survival may be wrong, but also maybe the data generation
-    # fct may be wrong
-    # a test showed that it looks preddy darn good. I guess this is because
-    # it is the conditional probability to survive. Not the conditional
-    # probability to die.
-    S_cond = (S[:, :-1] - S[:, 1:]) / S[:, :-1]
-    S_cond_ = jnp.column_stack([jnp.zeros_like(S[:,[0]]), S_cond])
-    S_cond__ = jnp.clip(S_cond_, EPS, 1 - EPS) 
-    
-    n_surv = observations["survivors_before_t"]
-    L_mask = masks["L"]
-    L = observations["L"]
-
-    if make_predictions:
-        L = None
-
-    numpyro.sample(
-        "L_obs", numpyro.distributions.Binomial(
-            probs=S_cond__, 
-            total_count=n_surv
-        ).mask(L_mask), 
-        obs=L
-    )
-
-
-def multinomial_error_model(theta, simulation_results, observations, indices, masks, make_predictions):
-    raise NotImplementedError(
-        "Multinomial error model needs to forward fill survival " +
-        "probabilities for missing (observation) values"
-        ""
-    )
-    # error model
-    EPS = observations["eps"]
-    # TODO This should not be done before calculating the difference
-    S = jnp.clip(simulation_results["S"], EPS, 1 - EPS) 
-
-    # TODO work on this
-    # This is not working with replicates
-    s_probs = S[:, :-1] - S[:, 1:]
-    s_probs_ = jnp.column_stack([jnp.zeros_like(S[:,[0]]), s_probs])
-    
-    L_mask = masks["L"]
-    L = observations["L"]
-    N = observations["n_subjects"]
-    # N_ = jnp.broadcast_to(jnp.expand_dims(N, 1), L.shape)
-
-    if make_predictions:
-        L = None
-
-    ID, T = L_mask.shape
-
-    with numpyro.plate("id", ID):
-        numpyro.sample(
-            "L_obs", numpyro.distributions.Multinomial(
-                probs=s_probs_, 
-                total_count=N
-            ).mask(L_mask), 
-            obs=L
-        )
-
 def survival_predictions(
     probs, n_trials, 
     eps: float = 0.0,
@@ -351,6 +91,8 @@ def posterior_predictions(sim, idata, seed=None):
     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(
@@ -371,24 +113,30 @@ def posterior_predictions(sim, idata, seed=None):
 
     samples = predictive(key)
 
-    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)
+    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={"survival_obs": np.array(chains)},
+        posterior_predictive=posterior_predictive,
         n_draws=n,
         n_chains=c
     )

guts_base/sim/__init__.py

@@ -1,6 +1,7 @@
 from . import base
 from . import ecx
 from . import report
+from . import utils
 
 from .base import (
     GutsBase,
@@ -11,4 +12,12 @@ from .base import (
 from .ecx import ECxEstimator, LPxEstimator
 from .report import GutsReport
 
-from .mempy import PymobSimulator
+from .mempy import PymobSimulator
+from .utils import (
+    GutsBaseError
+)
+
+from .constructors import (
+    construct_sim_from_config, 
+    load_idata
+)