gemlib 0.9.2__py3-none-any.whl

gemlib/mcmc/experimental/composable_kernel.py

@@ -0,0 +1,270 @@
+"""Implementation of Metropolis-within-Gibbs framework"""
+
+from collections import ChainMap, namedtuple
+from typing import AnyStr, Callable, List, Optional, Tuple
+
+import tensorflow_probability as tfp
+
+from .mcmc_base import (
+    ChainState,
+    KernelInfo,
+    KernelState,
+    Position,
+    SamplingAlgorithm,
+)
+
+split_seed = tfp.random.split_seed
+
+__all__ = ["Step"]
+
+
+def _maybe_list(x):
+    if isinstance(x, list):
+        return x
+    return [x]
+
+
+def _project_position(
+    position: Position, varnames: List[AnyStr]
+) -> Tuple[Position, Position]:
+    """Splits `position` into `position[varnames]` and
+    `position[~varnames]`
+    """
+    if varnames is None:
+        return (position, ())
+
+    target = {k: v for k, v in position._asdict().items() if k in varnames}
+    target_compl = {
+        k: v for k, v in position._asdict().items() if k not in varnames
+    }
+
+    return (
+        namedtuple("Target", target.keys())(**target),
+        namedtuple("TargetCompl", target_compl.keys())(**target_compl),
+    )
+
+
+def _join_dicts(a: dict, b: dict):
+    """Joins two dictionaries `a` and `b`"""
+    return dict(ChainMap(a, b))
+
+
+def _maybe_flatten(x: List):
+    """Flatten a list if `len(x) <= 1`"""
+    if len(x) == 0:
+        return None
+    if len(x) == 1:
+        return x[0]
+    return x
+
+
+class KernelInitMonad:
+    """KernelInitMonad is a Writer monad allowing us to build an initial
+    state tuple for a Metropolis-within-Gibbs algorithm
+    """
+
+    def __init__(self, fn: SamplingAlgorithm):
+        """The monad 'unit' function"""
+        self.__fn = fn
+
+    def __call__(self, *args, **kwargs):
+        """Monad ``run'' function"""
+        return self.__fn(*args, **kwargs)
+
+    def fish(self, next_kernel):
+        """Monad combination, i.e. Haskell fish operator"""
+
+        @KernelInitMonad
+        def compound_init_fn(
+            target_log_prob_fn: Callable[[Position], float],
+            initial_position: ChainState,
+        ):
+            _, self_kernel_state = self(target_log_prob_fn, initial_position)
+            next_chain_state, next_kernel_state = next_kernel(
+                target_log_prob_fn, initial_position
+            )
+
+            return (
+                next_chain_state,
+                _maybe_list(self_kernel_state) + _maybe_list(next_kernel_state),
+            )
+
+        return compound_init_fn
+
+    def __rshift__(self, next_kernel):
+        return self.fish(next_kernel)
+
+
+class KernelStepMonad:
+    """StepMonad is a state monad that allows us to chain MCMC kernels
+    together.
+    """
+
+    def __init__(self, fn: SamplingAlgorithm):
+        """The monad 'unit' function"""
+        self.__fn = fn  # Make private
+
+    def __call__(self, *args, **kwargs):
+        """Apply the state transformer computation to a state."""
+        return self.__fn(*args, **kwargs)
+
+    def __rshift__(self, next_kernel_fn):
+        """The monad 'bind' operator which allows chaining.
+        ma >> f :: ma -> (a -> mb) -> mb
+        """
+
+        def compound_step_kernel(
+            target_log_prob_fn: Callable[[Position], float],
+            chain_and_kernel_state: Tuple[ChainState, KernelState],
+            seed,
+        ):
+            first_seed, second_seed = split_seed(seed)
+            # Pre-order recursive descent
+            chain_state, kernel_state = chain_and_kernel_state
+
+            self_kernel_state = kernel_state[:-1]
+            next_kernel_state = kernel_state[-1]
+
+            # Execute self kernel step
+            (chain_state, self_kernel_state), self_info = self.__fn(
+                target_log_prob_fn,
+                (chain_state, _maybe_flatten(self_kernel_state)),
+                seed=first_seed,
+            )
+
+            # Descend right hand branch
+            (chain_state, next_kernel_state), next_info = next_kernel_fn(
+                target_log_prob_fn,
+                (chain_state, next_kernel_state),
+                seed=second_seed,
+            )
+
+            return (
+                (
+                    chain_state,
+                    _maybe_list(self_kernel_state)
+                    + _maybe_list(next_kernel_state),
+                ),
+                self_info + next_info,
+            )
+
+        return KernelStepMonad(compound_step_kernel)
+
+
+class ComposableSamplingAlgorithm:
+    """ComposableSamplingAlgorithm implements a Metropolis-within-Gibbs step"""
+
+    def __init__(self, init: KernelInitMonad, step: KernelStepMonad):
+        self.__initialize = init
+        self.__step = step
+
+    @property
+    def init(self):
+        """Initialize an MCMC chain"""
+        return self.__initialize
+
+    @property
+    def step(self):
+        """Function to invoke the MCMC kernel"""
+        return self.__step
+
+    def __rshift__(self, rhs):
+        """Combinator"""
+        return ComposableSamplingAlgorithm(
+            init=self.init >> rhs.init, step=self.step >> rhs.step
+        )
+
+
+class Step:  # pylint: disable=too-few-public-methods
+    """A Metropolis-within-Gibbs step"""
+
+    def __new__(
+        cls,
+        sampling_algorithm: SamplingAlgorithm,
+        target_names: Optional[List[str]] = None,
+        step_kwargs_fn: Callable[[Position], dict] = lambda _: {},
+    ):
+        """Transforms a base kernel to operate on a substate of a Markov chain.
+
+        Args:
+        ----
+        sampling_algorithm: a named tuple containing the generic kernel `init`
+                            and `step` function.
+        target_log_prob_fn: a list of variable names on which the
+                            Metropolis-within-Gibbs step is to operate
+        step_kwargs_fn: a callable taking the chain position as an argument,
+                        and returning a dictionary of extra kwargs to
+                        `sampling_algorithm.step`.
+
+        Returns:
+        -------
+        A monad of type StepMonad (State -> (State, Info))
+
+        """
+
+        @KernelInitMonad
+        def init(
+            target_log_prob_fn: Callable[[Position], float],
+            initial_position: Position,
+        ):
+            target, target_compl = _project_position(
+                initial_position, target_names
+            )
+
+            def conditional_tlp(*args):
+                state = _join_dicts(
+                    dict(zip(target._fields, args)),
+                    target_compl._asdict(),
+                )
+                return target_log_prob_fn(**state)
+
+            kernel_state = sampling_algorithm.init(conditional_tlp, target)
+
+            chain_state = ChainState(
+                position=initial_position,
+                log_density=kernel_state[0].log_density,
+                log_density_grad=kernel_state[0].log_density_grad,
+            )
+
+            return chain_state, kernel_state[1]
+
+        @KernelStepMonad
+        def step(
+            target_log_prob_fn: Callable[[Position], float],
+            chain_and_kernel_state: Tuple[ChainState, KernelState],
+            seed,
+        ) -> Tuple[Tuple[ChainState, KernelState], KernelInfo]:
+            chain_state, kernel_state = chain_and_kernel_state
+
+            # Split global state and generate conditional density
+            target, target_compl = _project_position(
+                chain_state.position, target_names
+            )
+
+            # Calculate the conditional log density
+            def conditional_tlp(*args):
+                state = _join_dicts(
+                    dict(zip(target._fields, args)),
+                    target_compl._asdict(),
+                )
+                return target_log_prob_fn(**state)
+
+            # Invoke the kernel on the target state
+            (new_target, new_kernel_state), info = sampling_algorithm.step(
+                conditional_tlp,
+                (chain_state._replace(position=target), kernel_state),
+                seed,
+                **step_kwargs_fn(chain_state.position),
+            )
+
+            # Stitch the global position back together
+            new_global_position = chain_state.position.__class__(
+                **new_target.position._asdict(), **target_compl._asdict()
+            )
+            new_global_state = new_target._replace(
+                position=new_global_position,
+            )
+
+            return (new_global_state, new_kernel_state), info
+
+        return ComposableSamplingAlgorithm(init, step)

gemlib/mcmc/experimental/discrete_time_state_transition_model/__init__.py

@@ -0,0 +1 @@
+"""DiscreteTimeStateTransitionModel-related MCMC samplers"""

gemlib/mcmc/experimental/discrete_time_state_transition_model/left_censored_events_mh.py

@@ -0,0 +1,149 @@
+"""Left-censored events MCMC kernel for DiscreteTimeStateTransitionModel"""
+
+from functools import partial
+from typing import NamedTuple, Optional, Tuple
+
+import tensorflow as tf
+import tensorflow_probability as tfp
+
+from gemlib.mcmc.discrete_time import UncalibratedLeftCensoredEventTimesUpdate
+from gemlib.mcmc.experimental.mcmc_base import ChainState, SamplingAlgorithm
+
+
+class LeftCensoredEventsState(NamedTuple):
+    max_timepoint: int
+    max_events: int
+
+
+class LeftCensoredEventsInfo(NamedTuple):
+    is_accepted: bool
+    log_acceptance_correction: float
+    target_log_prob: float
+    unit: int
+    timepoint: int
+    direction: int
+    num_events: int
+    seed: Tuple[int, int]
+
+
+class LeftCensoredEventsPosition(NamedTuple):
+    initial_conditions: tf.Tensor
+    events: tf.Tensor
+
+
+def _get_state_tuple(
+    initial_conditions_varname: str,
+    events_varname: str,
+    position: NamedTuple,
+):
+    return LeftCensoredEventsPosition(
+        getattr(position, initial_conditions_varname),
+        getattr(position, events_varname),
+    )
+
+
+def _repack_state_tuple(
+    initial_conditions_varname: str,
+    events_varname: str,
+    new_structure: LeftCensoredEventsPosition,
+    original_structure: NamedTuple,
+):
+    return original_structure.__class__(
+        **{
+            initial_conditions_varname: new_structure.initial_conditions,
+            events_varname: new_structure.events,
+        }
+    )
+
+
+def left_censored_events_mh(
+    transition_index: int,
+    max_timepoint: int,
+    max_events: int,
+    stoichiometry: tf.Tensor,
+    events_varname: str = "events",
+    initial_conditions_varname: str = "initial_conditions",
+    name: Optional[str] = None,
+):
+    """Left censored event times Metropolis Hastings
+
+    This MCMC kernel updates left-censored event times in a
+    DiscreteTimeStateTransitionModel realisation.
+
+    Args
+    ----
+    transition_index: the index of the transition to update
+    max_timepoint: max timepoint up to which to propose moves
+    max_events: max number of events per unit/timepoint to move
+
+    Returns
+    -------
+    A instance of SamplingAlgorithm
+    """
+
+    flatten_state = partial(
+        _get_state_tuple, initial_conditions_varname, events_varname
+    )
+    repack_state = partial(
+        _repack_state_tuple, initial_conditions_varname, events_varname
+    )
+
+    def _build_kernel(target_log_prob_fn):
+        return tfp.mcmc.MetropolisHastings(
+            inner_kernel=UncalibratedLeftCensoredEventTimesUpdate(
+                target_log_prob_fn,
+                transition_index,
+                stoichiometry,
+                max_timepoint,
+                max_events,
+                name,
+            )
+        )
+
+    def init_fn(target_log_prob_fn, target_state):
+        kernel = _build_kernel(target_log_prob_fn)
+        results = kernel.bootstrap_results(flatten_state(target_state))
+        chain_state = ChainState(
+            position=target_state,
+            log_density=results.accepted_results.target_log_prob,
+        )
+        kernel_state = LeftCensoredEventsState(
+            max_timepoint=max_timepoint, max_events=max_events
+        )
+
+        return chain_state, kernel_state
+
+    def step_fn(target_log_prob_fn, target_and_kernel_state, seed):
+        kernel = _build_kernel(target_log_prob_fn)
+        target_chain_state, kernel_state = target_and_kernel_state
+
+        new_target_position, results = kernel.one_step(
+            flatten_state(target_chain_state.position),
+            kernel.bootstrap_results(
+                flatten_state(target_chain_state.position)
+            ),
+            seed=seed,
+        )
+
+        new_chain_and_kernel_state = (
+            ChainState(
+                position=repack_state(
+                    new_target_position, target_chain_state.position
+                ),
+                log_density=results.accepted_results.target_log_prob,
+            ),
+            kernel_state,
+        )
+
+        return new_chain_and_kernel_state, LeftCensoredEventsInfo(
+            is_accepted=results.is_accepted,
+            log_acceptance_correction=results.proposed_results.log_acceptance_correction,
+            target_log_prob=results.proposed_results.target_log_prob,
+            unit=results.proposed_results.unit,
+            timepoint=results.proposed_results.timepoint,
+            direction=results.proposed_results.direction,
+            num_events=results.proposed_results.num_events,
+            seed=seed,
+        )
+
+    return SamplingAlgorithm(init_fn, step_fn)

gemlib/mcmc/experimental/discrete_time_state_transition_model/move_events.py

@@ -0,0 +1,71 @@
+"""Move partially-censored events in DiscreteTimeStateTransitionModel"""
+
+from typing import NamedTuple
+
+import tensorflow_probability as tfp
+
+from gemlib.mcmc.discrete_time_state_transition_model import (
+    TransitionTopology,
+    UncalibratedEventTimesUpdate,
+)
+from gemlib.mcmc.experimental.mcmc_base import ChainState, SamplingAlgorithm
+
+
+class MoveEventsState(NamedTuple):
+    dmax: int
+    mmax: int
+    nmax: int
+
+
+class MoveEventsInfo(NamedTuple):
+    is_accepted: bool
+
+
+def move_events(topology: TransitionTopology, dmax, mmax, nmax):
+    def _build_kernel(target_log_prob_fn, initial_conditions):
+        return tfp.mcmc.MetropolisHastings(
+            inner_kernel=UncalibratedEventTimesUpdate(
+                target_log_prob_fn,
+                topology,
+                initial_conditions,
+                dmax,
+                mmax,
+                nmax,
+            )
+        )
+
+    def init_fn(target_log_prob_fn, target_state, initial_conditions):
+        kernel = _build_kernel(target_log_prob_fn, initial_conditions)
+        results = kernel.bootstrap_results(target_state)
+        chain_state = ChainState(
+            position=target_state,
+            log_density=results.accepted_results.target_log_prob,
+        )
+        kernel_state = MoveEventsState(dmax=dmax, mmax=mmax, nmax=nmax)
+
+        return chain_state, kernel_state
+
+    def step_fn(
+        target_log_prob_fn, target_and_kernel_state, seed, initial_conditions
+    ):
+        kernel = _build_kernel(target_log_prob_fn, initial_conditions)
+
+        target_chain_state, kernel_state = target_and_kernel_state
+
+        new_target_position, results = kernel.one_step(
+            target_chain_state.position,
+            kernel.bootstrap_results(target_chain_state.position),
+            seed=seed,
+        )
+
+        new_chain_and_kernel_state = (
+            ChainState(
+                position=new_target_position,
+                log_density=results.accepted_results.target_log_prob,
+            ),
+            kernel_state,
+        )
+
+        return new_chain_and_kernel_state, MoveEventsInfo(results.is_accepted)
+
+    return SamplingAlgorithm(init_fn, step_fn)

gemlib/mcmc/experimental/discrete_time_state_transition_model/move_events_test.py

@@ -0,0 +1,48 @@
+"""Test partially censored events move for DiscreteTimeStateTransitionModel"""
+
+import pytest
+import tensorflow as tf
+
+from gemlib.mcmc.discrete_time_state_transition_model import TransitionTopology
+
+from .move_events import move_events
+
+
+@pytest.fixture
+def random_events():
+    """SEIR model with prescribed starting conditions"""
+    events = tf.random.uniform(
+        [10, 10, 3], minval=0, maxval=100, dtype=tf.float64, seed=0
+    )
+    return events
+
+
+@pytest.fixture
+def initial_state():
+    popsize = tf.fill([10], tf.constant(100.0, tf.float64))
+    initial_state = tf.stack(
+        [
+            popsize,
+            tf.ones_like(popsize),
+            tf.zeros_like(popsize),
+            tf.zeros_like(popsize),
+        ],
+        axis=-1,
+    )
+    return initial_state
+
+
+def test_move_events(random_events, initial_state):
+    def tlp(x):
+        return tf.constant(0.0, dtype=tf.float64)
+
+    kernel = move_events(
+        topology=TransitionTopology(0, 1, 2),
+        dmax=5,
+        mmax=1,
+        nmax=10,
+    )
+
+    state = kernel.init(tlp, random_events, initial_state)
+    seed = [0, 0]
+    new_state, results = kernel.step(tlp, state, seed, initial_state)

gemlib/mcmc/experimental/discrete_time_state_transition_model/right_censored_events_mh.py

@@ -0,0 +1,89 @@
+"""Right-censored events MCMC kernel for DiscreteTimeStateTransitionModel"""
+
+from typing import NamedTuple, Optional, Tuple
+
+import tensorflow as tf
+import tensorflow_probability as tfp
+
+from gemlib.mcmc.discrete_time_state_transition_model import (
+    TransitionTopology,
+    UncalibratedOccultUpdate,
+)
+from gemlib.mcmc.experimental.mcmc_base import ChainState, SamplingAlgorithm
+
+
+class RightCensoredEventsState(NamedTuple):
+    nmax: int
+    t_range: Tuple[int, int]
+
+
+class RightCensoredEventsInfo(NamedTuple):
+    is_accepted: bool
+    m: int
+    t: int
+    delta: int
+    x_star: int
+    proposed_state: tf.Tensor
+    seed: Tuple[int, int]
+
+
+def right_censored_events_mh(
+    topology: TransitionTopology,
+    nmax: int = 1,
+    t_range: Optional[Tuple[int, int]] = None,
+    name: Optional[str] = None,
+):
+    def _build_kernel(target_log_prob_fn, initial_conditions):
+        return tfp.mcmc.MetropolisHastings(
+            inner_kernel=UncalibratedOccultUpdate(
+                target_log_prob_fn,
+                topology,
+                initial_conditions,
+                nmax,
+                t_range,
+                name,
+            ),
+        )
+
+    def init_fn(target_log_prob_fn, target_state, initial_conditions):
+        kernel = _build_kernel(target_log_prob_fn, initial_conditions)
+        results = kernel.bootstrap_results(target_state)
+        chain_state = ChainState(
+            position=target_state,
+            log_density=results.accepted_results.target_log_prob,
+        )
+        kernel_state = RightCensoredEventsState(nmax=nmax, t_range=t_range)
+
+        return chain_state, kernel_state
+
+    def step_fn(
+        target_log_prob_fn, target_and_kernel_state, seed, initial_conditions
+    ):
+        kernel = _build_kernel(target_log_prob_fn, initial_conditions)
+        target_chain_state, kernel_state = target_and_kernel_state
+
+        new_target_position, results = kernel.one_step(
+            target_chain_state.position,
+            kernel.bootstrap_results(target_chain_state.position),
+            seed=seed,
+        )
+
+        new_chain_and_kernel_state = (
+            ChainState(
+                position=new_target_position,
+                log_density=results.accepted_results.target_log_prob,
+            ),
+            kernel_state,
+        )
+
+        return new_chain_and_kernel_state, RightCensoredEventsInfo(
+            results.is_accepted,
+            m=results.proposed_results.m,
+            t=results.proposed_results.t,
+            delta=results.proposed_results.delta_t,
+            x_star=results.proposed_results.x_star,
+            proposed_state=results.proposed_state,
+            seed=seed,
+        )
+
+    return SamplingAlgorithm(init_fn, step_fn)

gemlib/mcmc/experimental/discrete_time_state_transition_model/right_censored_events_mh_test.py

@@ -0,0 +1,28 @@
+"""Right-censored events MCMC test"""
+
+import pytest
+import tensorflow as tf
+
+from gemlib.mcmc.discrete_time_state_transition_model import TransitionTopology
+
+from .right_censored_events_mh import right_censored_events_mh
+
+
+@pytest.mark.skip("Work in progress")
+def test_right_censored_events_mh(random_events, initial_state):
+    def tlp(_):
+        return tf.constant(0.0, tf.float64)
+
+    kernel = right_censored_events_mh(
+        topology=TransitionTopology(0, 1, 2),
+        nmax=10,
+        t_range=(6, 10),
+    )
+
+    state = kernel.init(tlp, random_events, initial_state)
+    seed = [0, 0]
+    new_state, info = kernel.step(tlp, state, seed, initial_state)
+
+    print(info)
+
+    raise AssertionError()