guts-base 0.8.2__py3-none-any.whl

guts_base/sim/ecx.py

@@ -0,0 +1,357 @@
+import warnings
+from functools import partial
+import numpy as np
+import xarray as xr
+from typing import Literal, Optional, Dict, List
+import pandas as pd
+from scipy.optimize import minimize
+from matplotlib import pyplot as plt
+from tqdm import tqdm
+
+from pymob import SimulationBase
+
+class ECxEstimator:
+    """Estimates the exposure level that corresponds to a given effect. The algorithm 
+    operates by varying a given exposure profile (x_in)
+    """
+    _name = "EC"
+
+    def __init__(
+        self, 
+        sim: SimulationBase, 
+        effect: str, 
+        x: float=0.5,
+        id: Optional[str]=None, 
+        time: Optional[float]=None, 
+        x_in: Optional[xr.Dataset]=None, 
+    ):
+        self.sim = sim.copy()
+        self.time = time
+        self.x = x
+        self.id = id
+        self.effect = effect
+        self._mode = None
+
+        if id is None:
+            self.sim.coordinates["id"] = [self.sim.coordinates["id"][0]]
+        else:
+            self.sim.coordinates["id"] = [id]
+
+        self.sim.model_parameters["x_in"] = x_in
+
+        # self.sim.observations = self.sim.expand_batch_like_coordinate_to_new_dimension(
+        #     coordinate="exposure_path",
+        #     variables=["Flupyradifurone"]
+        # )
+
+        # self.sim.config.data_structure.remove("Flupyradifurone")
+
+        # # TODO: COnstruct a sim if the input dims change
+        # self.sim.config.data_structure.exposure.dimensions = ["id", "time", "exposure_path"]
+        self.sim.config.data_structure.survival.observed = False
+        self.sim.observations = self.sim.observations.sel(id=self.sim.coordinates["id"])
+
+        self.sim.model_parameters["y0"] = self.sim.parse_input("y0", drop_dims="time")
+        self.sim.dispatch_constructor()
+
+
+    
+    def _evaluate(self, factor, theta):
+        evaluator = self.sim.dispatch(
+            theta=theta, 
+            x_in=self.sim.validate_model_input(self.sim.model_parameters["x_in"] * factor)
+        )
+        evaluator()
+        return evaluator
+
+    def _loss(self, log_factor, theta):
+        # exponentiate the log factor
+        factor = np.exp(log_factor)
+
+        e = self._evaluate(factor, theta)
+        s = e.results.sel(time=self.time)[self.effect].values
+
+        return (s - (1 - self.x)) ** 2
+
+    def _posterior_mean(self):
+        mean = self.sim.inferer.idata.posterior.mean(("chain", "draw"))
+        mean = {k: v["data"] for k, v in mean.to_dict()["data_vars"].items()}
+        return mean
+
+    def _posterior_sample(self, i):
+        posterior_stacked = self.sim.inferer.idata.posterior.stack(
+            sample=("chain", "draw")
+        )
+        sample = posterior_stacked.isel(sample=i)
+        sample = {k: v["data"] for k, v in sample.to_dict()["data_vars"].items()}
+        return sample
+
+    def plot_loss_curve(self):
+        posterior_mean = self._posterior_mean()
+
+        factor = np.linspace(-2,2, 100)
+        y = list(map(partial(self._loss, theta=posterior_mean), factor))
+
+        fig, ax = plt.subplots(1,1, sharey=True, figsize=(4, 3))
+        ax.plot(
+            np.exp(factor), y, 
+            color="black",
+            label=f"$\ell = S(t={self.time},x_{{in}}=C_{{ext}} \cdot \phi) - {self.x}$"
+        )
+        ax.set_ylabel("Loss ($\ell$)")
+        ax.set_xlabel("Multiplication factor ($\phi$)")
+        ax.set_title(f"ID: {self.sim.coordinates['id'][0]}")
+        ax.set_ylim(0, np.max(y) * 1.25)
+        ax.legend(frameon=False)
+        fig.tight_layout()
+
+    def estimate(
+        self, 
+        mode: Literal["draws", "mean", "manual"] = "draws", 
+        draws: Optional[int] = None, 
+        parameters: Optional[Dict[str,float|List[float]]] = None,
+        log_x0: float = 0.0, 
+        accept_tol: float = 1e-5, 
+        optimizer_tol: float = 1e-5,
+        method: str = "cobyla", 
+        **optimizer_kwargs
+    ):
+        """The minimizer for the EC_x operates on the unbounded linear scale, estimating 
+        the log-modification factor. Converted to the linear scale by factor=exp(x), the 
+        profile modification factor is obtained.
+
+        Using x0=0.0 means optimization will start on the linear scale at the unmodified 
+        exposure profile. Using the log scale for optimization will provide much smoother
+        optimization performance because multiplicative steps on the log scale require 
+        much less adaptation.
+
+        Parameters
+        ----------
+
+        mode : Literal['draws', 'mean', 'manual']
+            mode of estimation. mode='mean' takes the mean of the posterior and estimate
+            the ECx for this singular value. mode='draws' takes samples from the posterior
+            and estimate the ECx for each of the parameter draws. mode='manual' takes
+            a parameter set (Dict) in the parameters argument and uses that for estimation. 
+            Default: 'draws'
+        
+        draws : int
+            Number of draws to take from the posterior. Only takes effect if mode='draw'.
+            Raises an exception if draws < 100, because this is insufficient for a 
+            reasonable uncertainty estimate. Default: None (using all samples from the
+            posterior)
+        
+        parameters : Dict[str,float|list[float]]
+            a parameter dictionary passed used as model parameters for finding the ECx
+            value. Default: None
+
+        log_x0 : float
+            the starting value for the multiplication factor of the exposure profile for
+            the minimization algorithm. This value is on the log scale. This means, 
+            exp(log_x0=0.0) = 1.0, which means that the log_x0=0.0 will start at an
+            unmodified exposure profile. Default: 0.0
+        
+        accept_tol : float
+            After optimization is finished, accept_tol is used to assess if the loss
+            function for the individual draws exceed a tolerance. These results are
+            discarded and a warning is emitted. This is to assert that no faulty optimization
+            results enter the estimate. Default: 1e-5
+        
+        optimizer_tol : float
+            Tolerance limit for the minimzer to stop optimization. Default 1e-5
+
+        method : str
+            Minization algorithm. See: https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.minimize.html
+            Default: 'cobyla'
+        
+        optimizer_kwargs :
+            Additional arguments to pass to the optimizer
+
+        """
+        x0_tries = np.array([0.0, -1.0, 1.0, -2.0, 2.0]) + log_x0
+
+        if mode == "draws":
+            if draws is None:
+                draws = (
+                    self.sim.inferer.idata.posterior.sizes["chain"] * 
+                    self.sim.inferer.idata.posterior.sizes["draw"]
+                )
+            elif draws < 100:
+                raise ValueError(
+                    "draws must be larger than 100. Preferably > 1000. "
+                    f"If you don't want uncertainty assessment of the {self._name} "
+                    "estimates, use mode='mean'"
+                )
+            else:
+                pass
+
+        elif mode == "mean":
+            draws = 1
+        elif mode == "manual":
+            draws = 1
+            if parameters is None:
+                raise ValueError(
+                    "parameters need to be provided if mode='manual'"
+                )
+        else:
+            raise NotImplementedError(
+                f"Bad mode: {mode}. Mode must be one 'mean' or 'draws'"
+            )
+
+        self._mode = mode
+        mult_factor = []
+        loss = []
+        iterations = []
+        for i in tqdm(range(draws)):
+            if mode == "draws":
+                sample = self._posterior_sample(i)
+            elif mode == "mean":
+                sample = self._posterior_mean()
+            elif mode == "manual":
+                sample = parameters
+            else: 
+                raise NotImplementedError(
+                    f"Bad mode: {mode}. Mode must be one 'mean' or 'draws'"
+                )
+
+            success = False
+            iteration = 0
+            while not success and iteration < len(x0_tries):
+                opt_res = minimize(
+                    self._loss, x0=x0_tries[iteration], 
+                    method=method,
+                    tol=optimizer_tol,
+                    args=(sample,),
+                    **optimizer_kwargs
+                ) 
+
+                success = opt_res.fun < accept_tol
+
+            # convert to linear scale from log scale
+            factor = np.exp(opt_res.x)
+
+            mult_factor.extend(factor)
+            iterations.append(iteration)
+            loss.append(opt_res.fun)
+
+        res_full = pd.DataFrame(dict(factor = mult_factor, loss=loss, retries=iterations))
+        if sum(res_full.loss >= accept_tol) > 0:
+            warnings.warn(
+                f"Not all optimizations converged on the {self._name}_{self.x}. " +
+                "Adjust starting values and method")
+            print(res_full)
+        
+        res = res_full.loc[res_full.loss < accept_tol,:]
+
+        summary = {
+            "mean": np.round(np.mean(res.factor.values), 4),
+            "q05": np.round(np.quantile(res.factor.values, 0.05), 4),
+            "q95": np.round(np.quantile(res.factor.values, 0.95), 4),
+            "std": np.round(np.std(res.factor.values), 4),
+            "cv": np.round(np.std(res.factor.values)/np.mean(res.factor.values), 2),
+        }
+
+        self.results = pd.Series(summary)
+        self.results_full = res_full
+
+        print("{name}_{x}".format(name=self._name, x=int(self.x * 100),))
+        print(self.results)
+        print("\n")
+
+    def plot_profile_and_effect(
+        self,
+        parameters: Optional[Dict[str,float|List[float]]] = None
+    ):
+        coordinates_backup = self.sim.coordinates["time"].copy()
+
+        self.sim.coordinates["time"] = np.linspace(0, self.time, 100)
+        self.sim.dispatch_constructor()
+
+        if self._mode is None:
+            raise RuntimeError(
+                "Run .estimate() before plot_profile_and_effect()"
+            )
+        elif self._mode == "mean" or self._mode == "draws":
+            e_new = self._evaluate(factor=self.results["mean"], theta=self._posterior_mean())
+            e_old = self._evaluate(factor=1.0, theta=self._posterior_mean())
+        elif self._mode == "manual":
+            if parameters is None:
+                raise RuntimeError(
+                    f"If {self._name}_x was estimated using manual mode, parameters must "+
+                    "also be provided here."
+                )
+            e_new = self._evaluate(factor=self.results["mean"], theta=parameters)
+            e_old = self._evaluate(factor=1.0, theta=parameters)
+    
+        extra_dim = [k for k in list(e_old.results.coords.keys()) if k not in ["time", "id"]]
+
+        if len(extra_dim) > 0:
+            labels_old = [
+                f"{l} (original)" for l 
+                in e_old.results.coords[extra_dim[0]].values
+            ]
+            labels_new = [
+                f"{l} (modified)" for l 
+                in e_new.results.coords[extra_dim[0]].values
+            ]
+        else:
+            labels_old = "original"
+            labels_new = "modified"
+
+
+
+        fig, (ax1, ax2) = plt.subplots(2,1, height_ratios=[1,3], sharex=True)
+        ax1.plot(
+            e_old.results.time, e_old.results.exposure.isel(id=0), 
+            ls="--", label=labels_old,
+        )
+        ax1.set_prop_cycle(None)
+        ax1.plot(
+            e_new.results.time, e_new.results.exposure.isel(id=0), 
+            label=labels_new
+        )
+
+
+        ax2.plot(
+            e_new.results.time, e_new.results.survival.isel(id=0), 
+            color="black", ls="--", label="modified"
+        )
+        ax1.set_prop_cycle(None)
+
+        ax2.plot(
+            e_old.results.time, e_old.results.survival.isel(id=0), 
+            color="black", ls="-", label="original"
+        )
+        ax2.hlines(self.x, e_new.results.time[0], self.time, color="grey")
+        ax1.set_ylabel("Exposure")
+        ax2.set_ylabel("Survival")
+        ax2.set_xlabel("Time")
+        ax1.legend()
+        ax2.legend()
+        ax2.set_xlim(0, None)
+        ax1.set_ylim(0, None)
+        ax2.set_ylim(0, None)
+        fig.tight_layout()
+
+        self.sim.coordinates["time"] = coordinates_backup
+        self.sim.dispatch_constructor()
+
+    
+
+class LPxEstimator(ECxEstimator):
+    """
+    the LPx is computed, using the existing exposure profile for 
+    the specified ID and estimating the multiplication factor for the profile that results
+    in an effect of X %
+    """
+    _name = "LP"
+
+    def __init__(
+        self, 
+        sim: SimulationBase, 
+        id: str,
+        x: float=0.5
+    ):
+        x_in = sim.model_parameters["x_in"].sel(id=[id])
+        time = sim.coordinates["time"][-1]
+        super().__init__(sim=sim, effect="survival", x=x, id=id, time=time, x_in=x_in)

guts_base/sim/mempy.py

@@ -0,0 +1,252 @@
+import pathlib
+from typing import Dict, Optional, Literal
+import re
+
+import numpy as np
+import pandas as pd
+import xarray as xr
+from pymob import SimulationBase
+from pymob.sim.config import Config, DataVariable, Datastructure
+from pymob.sim.parameters import Param
+from guts_base.sim import GutsBase
+from mempy.model import (
+    Model, 
+    RED_IT, 
+    RED_SD,
+    RED_IT_DA,
+    RED_SD_DA,
+    RED_IT_IA,
+    RED_SD_IA,
+    BufferGUTS_IT,
+    BufferGUTS_IT_CA,
+    BufferGUTS_IT_DA
+)
+
+__all__ = [
+    "PymobSimulator",
+]
+
+class PymobSimulator(GutsBase):
+
+    @classmethod
+    def from_mempy(
+        cls,
+        exposure_data: Dict,
+        survival_data: Dict,
+        model: Model,
+        info_dict: Dict = {},
+        pymob_config: Optional[Config] = None,
+        output_directory: str|pathlib.Path = pathlib.Path("output/pymob"),
+        default_prior: Literal["uniform", "lognorm"] = "lognorm",
+    ) -> SimulationBase:
+        """Construct a PymobSimulator from the 
+        """
+
+        if pymob_config is None:
+            cfg = Config()
+            # Configure: The configuration can be overridden in a subclass to override the 
+            # configuration
+            cls.configure(config=cfg)
+        else:
+            cfg = pymob_config
+
+        if isinstance(output_directory, str):
+            output_directory = pathlib.Path(output_directory)
+
+        cfg.case_study.output = str(output_directory)
+
+        # parse observations
+        # obs can be simply subset by selection obs.sel(substance="Exposure-Dime")
+        observations = xr.combine_by_coords([
+            cls._exposure_data_to_xarray(exposure_data, dim=model.extra_dim),
+            cls._survival_data_to_xarray(survival_data)
+        ])
+
+        # configure model and likelihood function
+        cfg.simulation.model = type(model).__name__
+        cfg.inference_numpyro.user_defined_error_model = str(model._likelihood_func_jax.__name__)
+
+        # derive data structure and params from the model instance
+        cls._set_data_structure(config=cfg, model=model)
+        cls._set_params(config=cfg, model=model, default_prior=default_prior)
+
+        # configure starting values and input
+        cfg.simulation.x_in = ["exposure=exposure"]
+        cfg.simulation.y0 = [f"{k}={v['y0']}" for k, v in model.state_variables.items() if "y0" in v]
+
+        # create a simulation object
+        sim = cls(config=cfg)
+        sim.config.create_directory(directory="results", force=True)
+
+        # initialize
+        sim.load_modules()
+        sim.set_logger()
+
+        sim.initialize(input={"observations": observations, "model": model})
+        
+        sim.validate()
+        sim.dispatch_constructor()
+
+
+        return sim
+
+    def initialize(self, input=None):
+        self.model = input["model"]._rhs_jax
+        self.solver_post_processing = input["model"]._solver_post_processing
+
+        super().initialize(input=input)
+
+
+    @classmethod
+    def configure(cls, config: Config):
+        """This is normally set in the configuration file passed to a SimulationBase class.
+        Since the mempy to pymob converter initializes pymob.SimulationBase from scratch
+        (without using a config file), the necessary settings have to be specified here.
+        """
+        config.case_study.output = "results"
+
+        config.simulation.x_dimension = "time"
+        config.simulation.batch_dimension = "id"
+        config.simulation.solver_post_processing = None
+        config.simulation.unit_time = "day"
+        config.simulation.n_reindexed_x = 100
+        config.simulation.forward_interpolate_exposure_data = True
+        
+        config.inference.extra_vars = ["eps", "survivors_before_t"]
+        config.inference.n_predictions = 100
+
+        config.jaxsolver.diffrax_solver = "Tsit5"
+        config.jaxsolver.rtol = 1e-10
+        config.jaxsolver.atol = 1e-12
+        config.jaxsolver.throw_exception = True
+        config.jaxsolver.pcoeff = 0.3
+        config.jaxsolver.icoeff = 0.3
+        config.jaxsolver.dcoeff = 0.0
+        config.jaxsolver.max_steps = 1000000
+        config.jaxsolver.throw_exception = True
+
+
+        config.inference_numpyro.gaussian_base_distribution = True
+        config.inference_numpyro.kernel = "svi"
+        config.inference_numpyro.init_strategy = "init_to_median"
+        config.inference_numpyro.svi_iterations = 10_000
+        config.inference_numpyro.svi_learning_rate = 0.001
+
+    @staticmethod
+    def _exposure_data_to_xarray(exposure_data: Dict[str, pd.DataFrame], dim: str):
+        """
+        TODO: Currently no rect interpolation
+        """
+        arrays = {}
+        for key, df in exposure_data.items():
+            # this override is necessary to make all dimensions work out
+            df.index.name = "time"
+            arrays.update({
+                key: df.to_xarray().to_dataarray(dim="id", name=key)
+            }) 
+
+        exposure_array = xr.Dataset(arrays).to_array(dim=dim, name="exposure")
+        exposure_array = exposure_array.transpose("id", "time", ...)
+        return xr.Dataset({"exposure": exposure_array})
+
+    @staticmethod
+    def _survival_data_to_xarray(survival_data: pd.DataFrame):
+        # TODO: survival name is currently not kept because the raw data is not transferred from the survival
+        survival_data.index.name = "time"
+        
+        survival_array = survival_data.to_xarray().to_dataarray(dim="id", name="survival")
+        survival_array = survival_array.transpose("id", "time", ...)
+        arrays = {"survival": survival_array}
+        return xr.Dataset(arrays)
+
+    @classmethod
+    def _set_data_structure(cls, config: Config, model: Model):
+        """Takes a dictionary that is specified in the model and uses only keys that
+        are fields of the DataVariable config-model"""
+        
+        state_dict = model.state_variables
+
+        config.data_structure = Datastructure(**{
+            key: DataVariable(**{
+                k: v for k, v in state_info.items()
+                if k in DataVariable.model_fields
+            })
+            for key, state_info in state_dict.items()
+        })
+
+
+    @classmethod
+    def _set_params(cls, config: Config, model: Model, default_prior: str):
+        params_info = model.params_info
+
+        if isinstance(model, (
+            RED_IT, RED_IT_DA, RED_IT_IA, 
+            BufferGUTS_IT, BufferGUTS_IT_CA, BufferGUTS_IT_DA
+        )):
+            eps = config.jaxsolver.atol * 10
+            params_info["eps"] = {'name':'eps', 'initial':eps, 'vary':False}
+
+
+        for par, param_dict in params_info.items():
+            for k, v in model._params_info_defaults.items():
+                if k not in param_dict:
+                    param_dict.update({k:v})
+
+        param_df = pd.DataFrame(params_info).T
+        param_df["param_index"] = param_df.name.apply(lambda x: re.findall(r"\d+", x))
+        param_df["param_index"] = param_df.param_index.apply(lambda x: int(x[0])-1 if len(x) == 1 else None)
+        param_df["name"] = param_df.name.apply(lambda x: re.sub(r"\d+", "", x).strip("_"))
+
+        for (param_name, ), group in param_df.groupby(["name"]):
+
+            dims = list(dict.fromkeys(group["dims"]))
+            dims = tuple([]) if dims == [None] else tuple(dims)
+
+            prior = list(dict.fromkeys(group["prior"]))
+            prior = prior[0] if len(prior) == 1 else prior
+            
+            _min = np.min(np.ma.masked_invalid(group["min"].values.astype(float)))
+            _max = np.max(np.ma.masked_invalid(group["max"].values.astype(float)))
+            _init = group["initial"].values.astype(float)
+            _free = group["vary"].values
+
+            # TODO: allow for parsing one N-D prior from multiple priors
+            # TODO: Another choice would be to parse vary=False priors as deterministic
+            #       and use a composite prior from a deterministic and a free prior as
+            #       the input into the model
+
+            if prior is None:
+                if default_prior == "uniform":
+                    _loc = _init * np.logical_not(_free) + _min * _free - config.jaxsolver.atol * 10 * np.logical_not(_free)
+                    _scale = _init * np.logical_not(_free) + _max * _free + config.jaxsolver.atol * 10 * np.logical_not(_free)
+                    _loc = _loc[0] if len(_loc) == 1 else _loc
+                    _scale = _scale[0] if len(_scale) == 1 else _scale
+                    prior = f"uniform(loc={_loc},scale={_scale})"
+                elif default_prior == "lognorm":
+                    _s = 3 * _free + config.jaxsolver.atol * 10 * np.logical_not(_free)
+                    _init = _init[0] if len(_init) == 1 else _init
+                    _s = _s[0] if len(_s) == 1 else _s
+
+                    prior = f"lognorm(scale={_init},s={_s})"
+                else:
+                    raise ValueError(
+                        f"Default prior: '{default_prior}' is not implemented. "+
+                        "Use one of 'uniform', 'lognorm' or specify priors for each "+
+                        "parameter directly with: "+
+                        f"`model.params_dict['prior'] = {default_prior}(...)`"
+                    )
+
+                prior = prior.replace(" ", ",")
+
+            # if isinstance(value,float):
+            param = Param(
+                value=_init,
+                free=np.max(_free),
+                min=_min,
+                max=_max,
+                prior=prior,
+                dims=dims
+            )
+
+            setattr(config.model_parameters, param_name, param)
+

guts_base/sim/report.py

@@ -0,0 +1,72 @@
+import os
+import itertools as it
+import pandas as pd
+
+from pymob import SimulationBase
+from pymob.sim.report import Report, reporting
+
+from guts_base.plot import plot_survival_multipanel
+from guts_base.sim.ecx import ECxEstimator
+
+class GutsReport(Report):
+
+    def additional_reports(self, sim: "SimulationBase"):
+        super().additional_reports(sim=sim)
+        self.model_fits(sim)
+        self.LCx_estimates(sim)
+
+    @reporting
+    def model_fits(self, sim: SimulationBase):
+        self._write("### Survival model fits")
+        
+        out_mp = plot_survival_multipanel(
+            sim=sim,
+            results=sim.inferer.idata.posterior_model_fits,
+            ncols=6,
+        )
+
+        lab = self._label.format(placeholder='survival_fits')
+        self._write(f"![Surival model fits.\label{{{lab}}}]({os.path.basename(out_mp)})")
+
+        return out_mp
+
+
+    @reporting
+    def LCx_estimates(self, sim):
+        X = [0.1, 0.25, 0.5, 0.75, 0.9]
+        T = [1, 2]
+        P = sim.predefined_scenarios
+
+        estimates = pd.DataFrame(
+            it.product(X, T, P.keys()), 
+            columns=["x", "time", "scenario"]
+        )
+
+        ecx = []
+
+        for i, row in estimates.iterrows():
+            ecx_estimator = ECxEstimator(
+                sim=sim,
+                effect="survival", 
+                x=row.x,
+                id=None, 
+                time=row.time, 
+                x_in=P[row.scenario], 
+            )
+            
+            ecx_estimator.estimate(
+                mode=sim.ecx_mode,
+                draws=250,
+            )
+
+            ecx.append(ecx_estimator.results)
+
+        results = pd.DataFrame(ecx)
+        estimates[results.columns] = results
+
+
+        estimates.to_csv()
+        file = os.path.join(sim.output_path, "lcx_estimates.csv")
+        lab = self._label.format(placeholder='$LC_x$ estimates')
+        self._write_table(tab=estimates, label_insert=f"$LC_x$ estimates \label{{{lab}}}]({os.path.basename(file)})")
+