guts-base 2.0.0b0__py3-none-any.whl

guts_base/mod.py

@@ -0,0 +1,332 @@
+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 guts_base.prob import likelihood
+
+from pymob.solvers.base import mappar
+from pymob.solvers.symbolic import (
+    PiecewiseSymbolicODESolver, FunctionPythonCode, get_return_arguments, dX_dt2X
+)
+
+_params_info_defaults = {
+    "initial": 1.0, 
+    "name": None,
+    "min": None,
+    "max": None,
+    "prior": None,
+    "dims": None,
+    "vary": True,
+    "module": None,
+    "unit": None,
+}
+
+class RED_SD:
+    """Simplest guts model, mainly for testing"""
+    extra_dim = "substance"
+    _likelihood_func_jax = likelihood
+    _it_model = False
+    _params_info_defaults = _params_info_defaults
+
+    params_info = {
+        "hb": dict(name="hb", min=None, max=None, initial=1.0, vary=True, prior="lognorm(scale=1,s=5)", dims=None, unit="1/{T}", module="background-mortality"),
+        "kd": dict(name="kd", min=None, max=None, initial=1.0, vary=True, prior="lognorm(scale=1,s=5)", dims=None, unit="1/{T}", module="tktd"),
+        "m":  dict(name="m", min=None, max=None, initial=1.0, vary=True, prior="lognorm(scale=1,s=5)", dims=None, unit="{X}", module="tktd"),
+        "b":  dict(name="b", min=None, max=None, initial=1.0, vary=True, prior="lognorm(scale=1,s=5)", dims=None, unit="1/{T}/{X}", module="tktd"),
+    }
+
+    state_variables = {
+        "exposure": dict(dimensions=["id", "time", "substance"], observed=False),
+        "D": dict(dimensions=["id", "time"], observed=False, y0=[0.0]),
+        "H": dict(dimensions=["id", "time"], observed=False, y0=[0.0]),
+        "survival": dict(dimensions=["id", "time"], observed=True)
+    }
+
+    @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
+
+class RED_IT:
+    """Simplest guts model, mainly for testing"""
+    extra_dim = "substance"
+    _likelihood_func_jax = likelihood
+    _it_model = True
+    _params_info_defaults = _params_info_defaults
+
+    params_info = {
+        "hb":   dict(name="hb",   min=None, max=None, initial=1.0, vary=True, prior="lognorm(scale=1,s=5)", dims=None, module="background-mortality"),
+        "kd":   dict(name="kd",   min=None, max=None, initial=1.0, vary=True, prior="lognorm(scale=1,s=5)", dims=None, module="tktd"),
+        "m":    dict(name="m",    min=None, max=None, initial=1.0, vary=True, prior="lognorm(scale=1,s=5)", dims=None, module="tktd"),
+        "beta": dict(name="beta", min=None, max=None, initial=1.0, vary=True, prior="lognorm(scale=1,s=5)", dims=None, module="tktd"),
+    }
+
+    state_variables = {
+        "exposure": dict(dimensions=["id", "time", "substance"], observed=False),
+        "D": dict(dimensions=["id", "time"], observed=False, y0=[0.0]),
+        "H": dict(dimensions=["id", "time"], observed=False),
+        "survival": dict(dimensions=["id", "time"], observed=True)
+    }
+
+    @staticmethod
+    def _rhs_jax(t, y, x_in, kd):
+        D, = y
+        C = x_in.evaluate(t)
+
+        dD_dt = kd * (C - D)
+
+        return (dD_dt, )
+
+    @staticmethod
+    def _solver_post_processing(results, t, interpolation, m, beta, hb, eps):
+        """
+        TODO: Try alternative formulation. This is computationally simpler and numerically
+        more stable:
+        log S = log 1.0 + log (1.0 - F) + log exp -hb * t = 0 + log (1.0 - F) - hb * t
+        """
+
+        d_max = jnp.squeeze(jnp.array([jnp.max(results["D"][:i+1])+eps for i in range(len(t))]))
+        F = jnp.where(d_max > 0, 1.0 / (1.0 + (d_max / m) ** -beta), 0)
+        S = 1.0 * (jnp.array([1.0], dtype=float) - F) * jnp.exp(-hb * t)
+        results["H"] = - jnp.log(S)
+        results["survival"] = S
+        results["exposure"] = jax.vmap(interpolation.evaluate)(t)
+
+        return results
+
+
+class RED_SD_DA:
+    """Simplest guts model, mainly for testing"""
+    extra_dim = "substance"
+    _likelihood_func_jax = likelihood
+    _it_model = False
+    _params_info_defaults = _params_info_defaults
+    params_info = {}
+
+    def __init__(self, num_expos = 1):
+        for i in range(num_expos):
+            self.params_info[f'kd{i+1}'] = {
+                'name':f'kd{i+1}', 'min':1.0e-3, 'max':1.0e3, 'initial':1.0, 
+                'vary':True, "dims": "substance", "module": "tktd",
+                "unit": "1/{T}"
+            }
+            if i == 0:
+                self.params_info[f'w{i+1}'] = {
+                    'name':f'w{i+1}', 'initial':1.0, 'vary':False, 
+                    "dims": "substance", "module": "tktd",
+                    "unit": "{X}/{X_i}",
+                }
+            else:
+                self.params_info[f'w{i+1}'] = {
+                    'name':f'w{i+1}', 'min':1.0e-3, 'max':1.0e3, 'initial':1.0, 
+                    'vary':True, "dims": "substance", "module": "tktd",
+                    "unit": "{X}/{X_i}"
+                }
+
+        self.params_info.update({
+            "hb": dict(name="hb", min=None, max=None, initial=1.0, vary=True, prior="lognorm(scale=1,s=5)", dims=None, unit="1/{T}", module="background-mortality"),
+            "m":  dict(name="m", min=None, max=None, initial=1.0, vary=True, prior="lognorm(scale=1,s=5)",  dims=None, unit="{X}", module="tktd"),
+            "b":  dict(name="b", min=None, max=None, initial=1.0, vary=True, prior="lognorm(scale=1,s=5)",  dims=None, unit="1/{T}/{X}", module="tktd"),
+        })
+
+        self.state_variables = {
+            "exposure": dict(dimensions=["id", "time", "substance"], unit="{Xi}", observed=False),
+            "D": {"dimensions": ["id", "time", "substance"], "observed": False, "y0": [0.0] * num_expos},
+            "H": dict(dimensions=["id", "time"], observed=False, y0=[0.0]),
+            "survival": dict(dimensions=["id", "time"], observed=True)
+        }
+
+    @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
+
+    @staticmethod
+    def _solver_post_processing(results, t, interpolation):
+        results["survival"] = jnp.exp(-results["H"])
+        results["exposure"] = jax.vmap(interpolation.evaluate)(t)
+        return results
+
+class RED_SD_explicit_units(RED_SD):
+    def __init__(self):
+        self.params_info["hb"]["unit"] = "1/day"
+        self.params_info["b"]["unit"] = "1/mg/day"
+        self.params_info["kd"]["unit"] = "1/day"
+        self.params_info["m"]["unit"] = "mg"
+
+red_sd = RED_SD._rhs_jax
+red_sd_post_processing = RED_SD._solver_post_processing
+
+
+red_sd_da = RED_SD_DA._rhs_jax
+red_sd_da_post_processing = RED_SD_DA._solver_post_processing
+
+
+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)
+    
+    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 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,201 @@
+from cycler import cycler
+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: SimulationBase, 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: SimulationBase, results, ncols=6, title=lambda _id: _id, filename="survival_multipanel"):
+
+    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)
+
+        ax.set_xlabel(f"Time [{sim.config.guts_base.unit_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) # type: ignore
+        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}/{filename}.png"
+    fig.tight_layout()
+    fig.savefig(out)
+
+    return out
+
+def plot_exposure_multipanel(sim: SimulationBase, results, ncols=6, title=lambda _id: _id, filename="exposure_multipanel"):
+
+    n_panels = results.sizes["id"]
+
+    nrows = int(np.ceil(n_panels / ncols))
+
+    fig, axes = plt.subplots(nrows, ncols, sharex=True, figsize=(ncols*2+2, nrows*1.5+2))
+    axes = axes.flatten()
+    mean = results
+
+    plot_kwargs = {"color": "black"}
+    custom_cycler = (
+        cycler(ls=["-", "--", ":", "-."])
+    )
+
+    labels = {}
+    for _id, ax in zip(sim.observations.id.values, axes):
+        ax.set_prop_cycle(custom_cycler)
+
+        ax.set_xlabel(f"Time [{sim.config.guts_base.unit_time}]")
+        ax.set_ylabel("Exposure")
+        for expo in sim.coordinates[sim._exposure_dimension]: # type: ignore
+            line, = ax.plot(
+                mean.time, mean.sel({"id":_id, sim._exposure_dimension: expo}).exposure, # type: ignore
+                **plot_kwargs, label=f"Exposure: {expo}"
+            )
+            labels.update({f"Exposure: {expo}": line})
+        ax.set_title(title(_id))
+
+    fig.legend(labels.values(), labels.keys(), loc='lower center', fontsize=10, frameon=False)
+    out = f"{sim.output_path}/{filename}.png"
+    fig.tight_layout(rect=[0,0.05,1.0,1.0], ) # type: ignore
+    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])
+
+    plot_results(results=results, batch_dim=sim.config.simulation.batch_dimension)
+
+def plot_results(results, batch_dim="id", axes=None, **plot_kwargs):
+    datavars = list(results.data_vars.keys())
+    if axes is None:
+        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", batch_dim)
+            for r in res_:
+                ax.plot(r.time, r, **plot_kwargs)
+        else:
+            ax.plot(results.time, res.transpose("time", batch_dim), **plot_kwargs)
+
+        ax.set_ylabel(dv)

guts_base/prob/__init__.py

@@ -0,0 +1,13 @@
+from . import conditional_binom
+from . import conditional_binom_mv
+from . import predictions
+from . import binom
+
+from .binom import likelihood
+
+from .predictions import (
+    survival_predictions,
+    posterior_predictions,
+)
+
+from .conditional_binom_mv import conditional_survival

guts_base/prob/binom.py

@@ -0,0 +1,18 @@
+import numpyro
+
+def likelihood(theta, simulation_results, indices, observations, masks, make_predictions):
+    """Uses lookup and error model from the local function context"""
+    if make_predictions:
+        obs = None
+    else:
+        obs = observations["survival"]
+    
+    _ = numpyro.sample(
+        name="survival" + "_obs",
+        fn=numpyro.distributions.Binomial(
+            probs=simulation_results["survival"],
+            total_count=observations["survivors_at_start"],
+        ).mask(masks["survival"]),
+        obs=obs
+    )
+

guts_base/prob/conditional_binom.py

@@ -0,0 +1,118 @@
+from guts_base.data.survival import generate_survival_repeated_observations
+from scipy.stats import rv_discrete
+from scipy.stats._discrete_distns import _isintegral
+import numpy as np
+import scipy.stats._boost as _boost
+from scipy import stats
+from matplotlib import pyplot as plt
+
+def conditional_prob_neglogexp(p, p_init=1.0, eps=1e-12):
+    p_ = np.concatenate([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_survival_gen(rv_discrete):
+    """
+    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_init: int
+        The starting number of positive observations (e.g. initial number of organisms
+        in a survival trial)
+
+    p_init: float
+        The starting survival probability in an experiment
+
+    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(p=p[1:], n_init=[n], p_init=[1.0], eps=[1e-12])
+
+    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)
+
+
+    """
+    def __init__(self, **kwargs):
+        # Initialize your custom parameters here
+        super().__init__(**kwargs)
+        # Set up any custom state needed for sampling
+
+    def _argcheck(self, p, n_init, p_init, eps):
+        return (n_init >= 0) & _isintegral(n_init) & (p >= 0) & (p <= 1)# & (np.diff(np.concatenate([p_init, p])) <= 0)
+
+    def _pmf(self, x, p, n_init, p_init, eps):
+        # nan filling is not necessary, because nans are thrown out this shifts the 
+        # p vector to where it belongs
+        n_ = np.concatenate([n_init[[0]], x[:-1]])
+        p_conditional = conditional_prob_neglogexp(p, p_init=p_init[[0]], eps=eps[[0]])
+        return _boost._binom_pdf(x, n_, p_conditional)
+        
+    def _rvs(self, p, n_init=10, p_init=[1.0], eps=[1e-12], size=1, random_state=None):
+        p_conditional = conditional_prob_neglogexp(p, p_init=p_init[[0]], eps=eps[[0]])
+    
+        # axis-0 is the batch dimension
+        # axis-1 is the time dimension (probability)
+        L = np.zeros(shape=(*size,))
+        L = np.array(L, ndmin=2)
+
+        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)
+            )
+
+        return n_init-L.cumsum(axis=1)
+
+conditional_survival = conditional_survival_gen(name="conditional_survival", )
+
+
+