guts-base 0.8.2__py3-none-any.whl

guts_base/mod.py

@@ -0,0 +1,251 @@
+"""
+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
+
+import jax
+import jax.numpy as jnp
+import sympy as sp
+import numpy as np
+from scipy.interpolate import interp1d
+
+from pymob.solvers.base import mappar
+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
+)
+
+red_sd = RED_SD._rhs_jax
+red_it = RED_IT._rhs_jax
+
+
+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):
+    # for constant exposure
+    D, H, S = y
+    dD_dt = k_d * (C_0 - D)
+    
+    switchDS = 0.5 + (1 / jnp.pi) * jnp.arctan(1e16 * (D - z))
+    dH_dt = (b * switchDS * (D - z) + h_b)
+
+    dS_dt = -dH_dt * S
+
+    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
+    C = x_in.evaluate(t)
+    dD_dt = k_d * (C - D)
+    
+    switchDS = 0.5 + (1 / jnp.pi) * jnp.arctan(1e16 * (D - z))
+    dH_dt = (b * switchDS * (D - z) + h_b)
+
+    dS_dt = -dH_dt * S
+
+    return dD_dt, dH_dt, dS_dt
+
+
+class Interpolation:
+    def __init__(self, xs, ys, method="previous") -> None:
+        self.f = interp1d(
+            x=xs,
+            y=ys,
+            axis=0,
+            kind=method
+        )
+
+    def evaluate(self, t) -> np.ndarray:
+        return self.f(t)
+
+
+class PiecewiseSymbolicSolver(PiecewiseSymbolicODESolver):
+    interpolation_type = "previous"
+
+    def t_jump(self, func_name, compiled_functions={}):
+        t, Y, Y_0, theta = self.define_symbols()
+        
+        D_t = compiled_functions["F"]["algebraic_solutions"]["D"]
+        z = theta["z"]
+        eq = sp.Eq(D_t.rhs, z).expand()
+
+        t_0 = sp.solve(eq, t)
+
+        assert len(t_0) == 1
+        func = sp.simplify(t_0[0])
+
+        python_code = FunctionPythonCode(
+            func_name=func_name,
+            lhs_0=("Y_0", tuple(Y_0.keys())),
+            theta=("theta", tuple(theta.keys())),
+            lhs=("Y",),
+            rhs=(func,),
+            expand_arguments=False,
+            modules=("numpy","scipy"),
+            docstring=""
+        )
+        
+        tex = self.to_latex(solutions=[sp.Eq(sp.Symbol(func_name), func)])
+        code_file = Path(self.output_path, f"{func_name}.tex")
+        with open(code_file, "w") as f:
+            f.writelines(tex)
+
+        return func, python_code
+
+    def define_symbols(self):
+        """Define the necessary symbols solely based on the function"""
+        thetanames = mappar(
+            self.model, {}, 
+            exclude=["t", "dt", "y", "x_in", "Y", "X"], 
+            to="names"
+        )
+        ynames = [dX_dt2X(a) for a in get_return_arguments(self.model)]
+
+        # define symbols for t, Y, Y_0 and theta
+        t = sp.Symbol("t", positive=True, real=True)
+        Y = {y: sp.Function(y, positive=True, real=True) for y in ynames}
+        Y_0 = {
+            f"{y}_0": sp.Symbol(f"{y}_0", positive=True, real=True) 
+            for y in ynames
+        }
+        theta = {p: sp.Symbol(p, positive=True, real=True) for p in thetanames}
+        
+        symbols = (t, Y, Y_0, theta)
+
+        return symbols
+        
+    @property
+    def compiler_recipe(self):
+        return {
+            "F":self.compile_model,
+            "t_jump":self.t_jump,
+            "F_piecewise":partial(self.jump_solution, funcnames="F t_jump")
+        }
+
+    def solve(self, parameters, y0, x_in):
+        odeargs = mappar(
+            self.model, 
+            parameters, 
+            exclude=["t", "dt", "y", "x_in"], 
+            to="dict"
+        )
+
+        F_piecewise = self.compiled_functions["F_piecewise"]["compiled_function"]
+
+        # get arguments
+        time = np.array(self.x)
+        Y_0_values = [v for v in y0.values()]
+
+        # handle interpolation
+        if "exposure" in self.coordinates_input_vars["x_in"]:
+            exposure_interpolation = Interpolation(
+                xs=self.coordinates_input_vars["x_in"]["exposure"][self.x_dim],
+                ys=x_in["exposure"],
+                method=self.interpolation_type
+            )
+        else:
+            exposure_interpolation = Interpolation(
+                xs=self.x,
+                ys=np.full_like(self.x, odeargs["C_0"]), #type: ignore
+                method=self.interpolation_type
+            )
+
+        # run main loop
+        sol = [Y_0_values]
+        for i in range(1, len(time)):
+            # parse arguments
+            C_0 = exposure_interpolation.evaluate(time[i-1])
+            odeargs.update({"C_0": C_0}) # type: ignore
+            dt = time[i] - time[i-1]
+
+            # call piecewise function
+            y_t = F_piecewise(
+                t=dt, 
+                Y_0=sol[i-1], 
+                θ=tuple(odeargs.values()), # type: ignore
+                ε=1e-14
+            ) 
+            
+            sol.append(y_t)
+
+        Y_t = np.array(sol)
+
+        results = {k:y for k, y in zip(y0.keys(), Y_t.T)}
+        results["exposure"] = exposure_interpolation.evaluate(np.array(self.x))
+
+        return results

guts_base/plot.py

@@ -0,0 +1,162 @@
+import numpy as np
+import arviz as az
+from matplotlib import pyplot as plt
+
+from pymob.sim.plot import SimulationPlot
+from pymob import SimulationBase
+
+def plot_prior_predictions(
+    sim: SimulationBase, 
+    data_vars=["survival"],
+    title_func=lambda sp, c: f"{sp.observations.id.values[0]}"
+):
+    idata = sim.inferer.prior_predictions() # type: ignore
+    def plot_survival_data_probs(idata):
+        return idata["survival"] * sim.observations.sel(id=[id_], time=0).survival
+    
+    fig, axes = plt.subplots(
+        ncols=sim.observations.dims["id"], nrows=len(data_vars), 
+        sharex=True, sharey="row", figsize=(30,10), squeeze=False)
+    for i, id_ in enumerate(sim.coordinates["id"]):
+        simplot = SimulationPlot(
+            observations=sim.observations.sel(id=[id_]),
+            idata=idata.sel(id=[id_]),  # type: ignore
+            rows=data_vars,
+            coordinates=sim.coordinates,
+            config=sim.config,
+            obs_idata_map={
+                "survival": plot_survival_data_probs
+            },
+            idata_groups=["prior_model_fits"],  # type: ignore
+        )
+        # replace simplot axis
+        for j, k in enumerate(simplot.rows):
+            simplot.axes_map[k]["all"] = axes[j][i]
+
+        simplot.plot_data_variables()
+        simplot.set_titles(title_func)
+        for j, k in enumerate(simplot.rows):
+            if i != 0:
+                simplot.axes_map[k]["all"].set_ylabel("")
+            if j != 3:
+                simplot.axes_map[k]["all"].set_xlabel("")
+
+        simplot.close()
+    
+    fig.tight_layout()
+    fig.savefig(f"{sim.output_path}/combined_prior_predictions.png")
+
+
+def plot_posterior_predictions(
+    sim: SimulationBase, 
+    data_vars=["survival"],
+    title_func=lambda sp, c: f"{sp.observations.id.values[0]}",
+    groups=["posterior_model_fits", "posterior_predictive"],
+):
+    fig, axes = plt.subplots(
+        ncols=sim.observations.dims["id"], nrows=len(data_vars), 
+        sharex=True, sharey="row", figsize=(30,10), squeeze=False
+    )
+
+    def plot_survival_data_probs_and_preds(dataset):
+        if dataset.attrs["group"] == "posterior_model_fits":
+            return dataset["survival"] * sim.observations.sel(id=[id_], time=0).survival
+        if dataset.attrs["group"] == "posterior_predictive":
+            return dataset["survival"]
+
+    for i, id_ in enumerate(sim.coordinates["id"]):
+        simplot = SimulationPlot(
+            observations=sim.observations.sel(id=[id_]),
+            idata=sim.inferer.idata.sel(id=[id_]),  # type: ignore
+            rows=data_vars,
+            coordinates=sim.coordinates,
+            config=sim.config,
+            obs_idata_map={
+                "survival": plot_survival_data_probs_and_preds,
+            },
+            idata_groups=groups,  # type: ignore
+        )
+        
+        # replace simplot axis
+        for j, k in enumerate(simplot.rows):
+            simplot.axes_map[k]["all"] = axes[j][i]
+
+        simplot.plot_data_variables()
+        simplot.set_titles(title_func)
+        for j, k in enumerate(simplot.rows):
+            if i != 0:
+                simplot.axes_map[k]["all"].set_ylabel("")
+            if j != 3:
+                simplot.axes_map[k]["all"].set_xlabel("")
+
+        simplot.close()
+    
+    fig.tight_layout()
+    fig.savefig(f"{sim.output_path}/combined_posterior_predictions.png")
+
+
+def plot_survival(sim, results):
+    fig, ax = plt.subplots(1,1)
+    obs = sim.observations.survival / sim.observations.subject_count
+    ax.plot(sim.observations.time, obs.T, marker="o", color="black")
+    ax.plot(results.time, results.survival.T, color="black")
+
+
+def plot_survival_multipanel(sim, results, ncols=6, title=lambda _id: _id):
+
+    n_panels = results.sizes["id"]
+
+    nrows = int(np.ceil(n_panels / ncols))
+
+    fig, axes = plt.subplots(nrows, ncols, sharex=True, sharey=True, figsize=(ncols*2+2, nrows*1.5+2))
+    axes = axes.flatten()
+    mean = results.mean(("chain", "draw"))
+    hdi = az.hdi(results, 0.95)
+    survival = sim.observations.survival / sim.observations.survival.isel(time=0)
+
+    plot_kwargs = {"color": "black"}
+    # param_cycler = plt.rcParams['axes.prop_cycle']
+    for _id, ax in zip(sim.observations.id.values, axes):
+        ax.set_ylim(-0.05,1.05)
+
+        # TODO: use time unit from observations (?)
+        ax.set_xlabel("Time")
+        ax.set_ylabel("Survival")
+        ax.plot(mean.time, mean.sel(id=_id).survival.T, **plot_kwargs)
+        ax.fill_between(hdi.time, *hdi.sel(id=_id).survival.T, alpha=.5, **plot_kwargs)
+        ax.plot(survival.time, survival.sel(id=_id).T, ls="", marker="o", alpha=.5, **plot_kwargs)
+        ax.set_title(title(_id))
+
+    out = f"{sim.output_path}/survival_multipanel.png"
+    fig.tight_layout()
+    fig.savefig(out)
+
+    return out
+
+
+def multipanel_title(sim, _id): 
+    oid = sim.observations.sel(id=_id)
+    exposure_path = oid.exposure_path.values
+    rac = np.round(oid.concentration_closer_x_rac.max().values  * 100, 3)
+    return "{ep}\n{c} %RAC".format(ep=str(exposure_path), c=str(rac))
+
+def plot_intermediate_results(sim: SimulationBase, id=0):
+    e = sim.dispatch()
+    e()
+    e.results
+
+    results = e.results.isel(id=[id])
+
+    datavars = sim.config.data_structure.data_variables
+    fig, axes = plt.subplots(len(datavars), 1)
+
+    for ax, dv in zip(axes, datavars):
+        res = results[dv]
+        if len(res.shape) > 2:
+            res_ = res.transpose(..., "time", sim.config.simulation.batch_dimension)
+            for r in res_:
+                ax.plot(r.time, r, alpha=.5)
+        else:
+            ax.plot(results.time, res.transpose("time", sim.config.simulation.batch_dimension), alpha=.5)
+
+        ax.set_ylabel(dv)