guts-base 0.8.6__tar.gz → 1.0.0__tar.gz

{guts_base-0.8.6 → guts_base-1.0.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: guts_base
-Version: 0.8.6
+Version: 1.0.0
 Summary: Basic GUTS model implementation in pymob
 Author-email: Florian Schunck <fluncki@protonmail.com>
 License:                     GNU GENERAL PUBLIC LICENSE
@@ -692,8 +692,7 @@ License-File: LICENSE
 Requires-Dist: openpyxl>=3.1.3
 Requires-Dist: Bottleneck>=1.5.0
 Requires-Dist: expydb>=0.6.0
-Requires-Dist: mempyguts>=1.5.0
-Requires-Dist: pymob[interactive,numpyro]<0.6.0,>=0.4.1
+Requires-Dist: pymob[interactive,numpyro]<0.6.0,>=0.5.10
 Provides-Extra: dev
 Requires-Dist: pytest>=7.3; extra == "dev"
 Requires-Dist: bumpver; extra == "dev"

{guts_base-0.8.6 → guts_base-1.0.0}/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-0.8.6 → guts_base-1.0.0}/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-0.8.6 → guts_base-1.0.0}/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-0.8.6 → guts_base-1.0.0}/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-0.8.6 → guts_base-1.0.0}/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-1.0.0/guts_base/prob.py

@@ -0,0 +1,160 @@
+from functools import partial
+import numpy as np
+import jax
+import jax.numpy as jnp
+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-0.8.6 → guts_base-1.0.0}/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
+)