jact 0.1.1__py3-none-any.whl

jact/__init__.py

@@ -0,0 +1,76 @@
+"""jact — JAX-based transition probability computation for multi-state models.
+
+A framework for computing transition probabilities in semi-Markov
+multi-state models with duration-dependent transition intensities.
+
+Example
+-------
+>>> import jact
+>>> import jax.numpy as jnp
+>>>
+>>> state_space = jact.StateSpace(
+...     states=["healthy", "disabled", "dead"],
+...     transitions=[
+...         ("healthy", "disabled"),
+...         ("healthy", "dead"),
+...         ("disabled", "dead"),
+...     ],
+... )
+>>>
+>>> model = state_space.build(
+...     transitions={
+...         ("healthy", "disabled"): onset_fn,
+...         ("healthy", "dead"): mortality_fn,
+...         ("disabled", "dead"): disabled_mort_fn,
+...     }
+... )
+>>>
+>>> result = model.solve(
+...     initial="healthy",
+...     horizon=10,
+...     steps_per_unit=12,
+...     age=ages,
+... )
+"""
+
+from importlib.metadata import PackageNotFoundError, version
+
+from . import callbacks
+from .cashflows import (
+    ByKind,
+    ByState,
+    CashflowDeclaration,
+    Group,
+    Raw,
+    ScheduledEvent,
+    StateRate,
+    Total,
+    TransitionLump,
+)
+from .initial_distribution import InitialDistribution
+from .model import Model
+from .solver import solve
+from .state_space import StateSpace
+
+try:
+    __version__ = version("jact")
+except PackageNotFoundError:
+    __version__ = "0.1.0"
+
+__all__ = [
+    "__version__",
+    "StateSpace",
+    "Model",
+    "InitialDistribution",
+    "solve",
+    "callbacks",
+    "StateRate",
+    "TransitionLump",
+    "ScheduledEvent",
+    "Raw",
+    "Group",
+    "Total",
+    "ByState",
+    "ByKind",
+    "CashflowDeclaration",
+]

jact/callbacks.py

@@ -0,0 +1,185 @@
+"""Probability callbacks for extracting results from the solver state."""
+
+from __future__ import annotations
+
+from typing import Any, Callable, NamedTuple, Union
+
+import jax
+import jax.numpy as jnp
+
+__all__ = [
+    "PointMass",
+    "StateCarry",
+    "density",
+    "density_probability",
+    "full",
+    "marginal_components",
+    "none_callback",
+    "point_mass",
+    "state_probability",
+    "resolve_callback",
+]
+
+
+def _validate_shape(
+    name: str,
+    value: jnp.ndarray,
+    expected_shape: tuple[int, ...],
+) -> None:
+    shape = jnp.shape(value)
+    if shape != expected_shape:
+        raise ValueError(
+            f"{name} must have shape {expected_shape}, got {shape}."
+        )
+
+
+def _validate_non_negative_if_concrete(value: jnp.ndarray) -> None:
+    try:
+        arr = jnp.asarray(value)
+        if bool(jnp.any(arr < 0)):
+            raise ValueError("value must be non-negative.")
+    except Exception as exc:  # pragma: no cover - tracer path
+        if "tracer" not in type(exc).__name__.lower():
+            try:
+                message = str(exc).lower()
+            except Exception:  # pragma: no cover
+                message = ""
+            if "tracer" not in message and "concret" not in message:
+                raise
+
+
+@jax.tree_util.register_pytree_node_class
+class PointMass:
+    """Per-individual point mass carried along a characteristic."""
+
+    __slots__ = ("value", "d_0", "log_value")
+
+    def __init__(
+        self,
+        value: jnp.ndarray,
+        d_0: jnp.ndarray,
+        log_value: jnp.ndarray | None = None,
+    ):
+        value_shape = jnp.shape(value)
+        _validate_shape("d_0", d_0, value_shape)
+        if log_value is not None:
+            _validate_shape("log_value", log_value, value_shape)
+        _validate_non_negative_if_concrete(value)
+        self.value = value
+        self.d_0 = d_0
+        self.log_value = (
+            jnp.where(value > 0, jnp.log(value), -jnp.inf)
+            if log_value is None
+            else log_value
+        )
+
+    def tree_flatten(self):
+        return (self.value, self.d_0, self.log_value), None
+
+    @classmethod
+    def tree_unflatten(cls, _, children):
+        value, d_0, log_value = children
+        self = cls.__new__(cls)
+        self.value = value
+        self.d_0 = d_0
+        self.log_value = log_value
+        return self
+
+
+class StateCarry(NamedTuple):
+    """Per-state solver carry."""
+
+    density: jnp.ndarray
+    point_mass: PointMass | None
+
+
+@jax.jit
+def none_callback(state: tuple[StateCarry, ...]):
+    """Return nothing. Use when only side effects of the solve matter."""
+    return None
+
+
+@jax.jit
+def full(state: tuple[StateCarry, ...]):
+    """Return the full pytree state unchanged."""
+    return state
+
+
+@jax.jit
+def marginal_components(state: tuple[StateCarry, ...]):
+    """Marginalize density over duration while keeping point masses explicit."""
+    return tuple(
+        StateCarry(
+            density=jnp.sum(carry.density, axis=-1),
+            point_mass=carry.point_mass,
+        )
+        for carry in state
+    )
+
+
+@jax.jit
+def state_probability(state: tuple[StateCarry, ...]):
+    """Return total state occupancy after marginalizing over duration."""
+    return jnp.stack(
+        tuple(
+            jnp.sum(carry.density, axis=-1)
+            if carry.point_mass is None
+            else jnp.sum(carry.density, axis=-1) + carry.point_mass.value
+            for carry in state
+        ),
+        axis=-1,
+    )
+
+
+@jax.jit
+def point_mass(state: tuple[StateCarry, ...]):
+    """Return only the point-mass component per state."""
+    return tuple(carry.point_mass for carry in state)
+
+
+@jax.jit
+def density(state: tuple[StateCarry, ...]):
+    """Return only the absolutely continuous density per state."""
+    return tuple(carry.density for carry in state)
+
+
+@jax.jit
+def density_probability(state: tuple[StateCarry, ...]):
+    """Return the duration-marginal density, restacked across states."""
+    return jnp.stack(
+        tuple(jnp.sum(carry.density, axis=-1) for carry in state),
+        axis=-1,
+    )
+
+
+_CALLBACKS = {
+    "none": none_callback,
+    "full": full,
+    "marginal_components": marginal_components,
+    "state_probability": state_probability,
+    "point_mass": point_mass,
+    "density": density,
+    "density_probability": density_probability,
+}
+
+
+def resolve_callback(
+    callback: Union[None, str, Callable[[tuple[StateCarry, ...]], Any]],
+) -> Callable[[tuple[StateCarry, ...]], Any]:
+    """Resolve a callback specification to a callable."""
+    if callback is None:
+        return none_callback
+    if callable(callback):
+        return callback
+    if isinstance(callback, str):
+        if callback not in _CALLBACKS:
+            available = ", ".join(sorted(_CALLBACKS.keys()))
+            raise ValueError(
+                f"Unknown callback '{callback}'. "
+                f"Available callbacks: {available}"
+            )
+        return _CALLBACKS[callback]
+    raise TypeError(
+        "callback must be None, a string, or a callable, "
+        f"got {type(callback)}"
+    )

jact/cashflows.py

@@ -0,0 +1,276 @@
+"""Cashflow declarations and solve-time views."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from numbers import Number
+from typing import Any, Callable, Mapping, Sequence
+
+import jax.numpy as jnp
+
+Scalar = int | float
+
+__all__ = [
+    "ByKind",
+    "ByState",
+    "CashflowDeclaration",
+    "Group",
+    "Raw",
+    "ScheduledEvent",
+    "StateRate",
+    "Total",
+    "TransitionLump",
+]
+
+
+@dataclass(frozen=True)
+class StateRate:
+    """Payment-rate callables attached to occupied states."""
+
+    payments: Mapping[str, Callable[..., jnp.ndarray]]
+
+
+@dataclass(frozen=True)
+class TransitionLump:
+    """Lump-sum payment callables attached to transitions."""
+
+    payments: Mapping[tuple[str, str], Callable[..., jnp.ndarray]]
+
+
+@dataclass(frozen=True)
+class ScheduledEvent:
+    """State-conditioned payments at deterministic event times."""
+
+    when: Callable[..., jnp.ndarray]
+    payments: Mapping[str, Callable[..., jnp.ndarray]]
+
+
+@dataclass(frozen=True)
+class Raw:
+    """Return one raw component or all raw components."""
+
+    name: str | None = None
+    weight: Callable[..., jnp.ndarray] | Scalar | None = None
+    terminal: bool = False
+
+
+@dataclass(frozen=True)
+class Group:
+    """Return the sum of selected raw components."""
+
+    members: Sequence[str]
+    weight: Callable[..., jnp.ndarray] | Scalar | None = None
+    terminal: bool = False
+
+
+@dataclass(frozen=True)
+class Total:
+    """Return the sum of all raw components."""
+
+    weight: Callable[..., jnp.ndarray] | Scalar | None = None
+    terminal: bool = False
+
+
+@dataclass(frozen=True)
+class ByState:
+    """Return cashflows split by reachable state."""
+
+    weight: Callable[..., jnp.ndarray] | Scalar | None = None
+    terminal: bool = False
+
+
+@dataclass(frozen=True)
+class ByKind:
+    """Return cashflows split by component kind."""
+
+    weight: Callable[..., jnp.ndarray] | Scalar | None = None
+    terminal: bool = False
+
+
+@dataclass(frozen=True)
+class CashflowDeclaration:
+    """Validated cashflow components bound to a state-space topology."""
+
+    state_space: Any
+    components: tuple[tuple[str, StateRate | TransitionLump | ScheduledEvent], ...]
+
+    @property
+    def names(self) -> tuple[str, ...]:
+        return tuple(name for name, _ in self.components)
+
+    def component(self, name: str) -> StateRate | TransitionLump | ScheduledEvent:
+        for component_name, component in self.components:
+            if component_name == name:
+                return component
+        raise ValueError(f"Unknown cashflow component '{name}'.")
+
+
+def _check_component_name(name: Any) -> None:
+    if not isinstance(name, str) or not name:
+        raise ValueError("cashflow component names must be non-empty strings.")
+
+
+def _check_callable(value: Any, field: str) -> None:
+    if not callable(value):
+        raise TypeError(f"{field} must be callable.")
+
+
+def _validate_payment_mapping(payments: Any, field: str) -> Mapping[Any, Any]:
+    if not isinstance(payments, Mapping) or not payments:
+        raise ValueError(f"{field} must be a non-empty mapping.")
+    for fn in payments.values():
+        _check_callable(fn, f"{field} values")
+    return dict(payments)
+
+
+def _is_scalar_array_like(value: Any) -> bool:
+    if value is None:
+        return False
+    try:
+        return bool(jnp.asarray(value).ndim == 0)
+    except Exception:
+        return False
+
+
+def _normalise_weight(weight: Any) -> Any:
+    if weight is None:
+        return None
+    if _is_scalar_array_like(weight):
+        return jnp.asarray(weight).item()
+    return weight
+
+
+def validate_cashflow_components(
+    state_space: Any,
+    components: Mapping[str, StateRate | TransitionLump | ScheduledEvent],
+) -> CashflowDeclaration:
+    """Validate and freeze a component mapping for a state space."""
+    if not isinstance(components, Mapping) or not components:
+        raise ValueError("cashflows() requires a non-empty component mapping.")
+
+    frozen: list[tuple[str, StateRate | TransitionLump | ScheduledEvent]] = []
+    seen: set[str] = set()
+    for name, component in components.items():
+        _check_component_name(name)
+        if name in seen:
+            raise ValueError(f"Duplicate cashflow component name '{name}'.")
+        seen.add(name)
+
+        if isinstance(component, StateRate):
+            payments = _validate_payment_mapping(
+                component.payments,
+                f"StateRate('{name}').payments",
+            )
+            for state in payments:
+                state_space._check_state(state)
+            frozen_component = StateRate(payments=payments)
+        elif isinstance(component, TransitionLump):
+            payments = _validate_payment_mapping(
+                component.payments,
+                f"TransitionLump('{name}').payments",
+            )
+            for transition in payments:
+                if (
+                    not isinstance(transition, tuple)
+                    or len(transition) != 2
+                    or not state_space.has_transition(*transition)
+                ):
+                    raise ValueError(
+                        f"TransitionLump('{name}') references unknown "
+                        f"transition {transition!r}."
+                    )
+            frozen_component = TransitionLump(payments=payments)
+        elif isinstance(component, ScheduledEvent):
+            _check_callable(component.when, f"ScheduledEvent('{name}').when")
+            payments = _validate_payment_mapping(
+                component.payments,
+                f"ScheduledEvent('{name}').payments",
+            )
+            for state in payments:
+                state_space._check_state(state)
+            frozen_component = ScheduledEvent(
+                when=component.when,
+                payments=payments,
+            )
+        else:
+            raise TypeError(
+                "cashflow components must be StateRate, TransitionLump, "
+                f"or ScheduledEvent; got {type(component)}."
+            )
+        frozen.append((name, frozen_component))
+
+    return CashflowDeclaration(state_space=state_space, components=tuple(frozen))
+
+
+def _validate_view_common(view: Any) -> None:
+    if not isinstance(view.terminal, bool):
+        raise TypeError("cashflow view terminal must be a bool.")
+    weight = view.weight
+    if weight is None or callable(weight) or isinstance(weight, Number):
+        return
+    if _is_scalar_array_like(weight):
+        return
+    raise TypeError("cashflow view weight must be None, a scalar, or callable.")
+
+
+def validate_cashflow_views(
+    declaration: CashflowDeclaration,
+    views: Mapping[str, Raw | Group | Total | ByState | ByKind] | None,
+) -> tuple[tuple[str, Raw | Group | Total | ByState | ByKind], ...]:
+    """Validate and freeze solve-time cashflow views."""
+    if views is None:
+        views = {"raw": Raw()}
+    if not isinstance(views, Mapping):
+        raise TypeError("cashflow_views must be a mapping or None.")
+
+    component_names = set(declaration.names)
+    frozen: list[tuple[str, Raw | Group | Total | ByState | ByKind]] = []
+    seen: set[str] = set()
+    for name, view in views.items():
+        if not isinstance(name, str) or not name:
+            raise ValueError("cashflow view names must be non-empty strings.")
+        if name in seen:
+            raise ValueError(f"Duplicate cashflow view name '{name}'.")
+        seen.add(name)
+
+        if isinstance(view, Raw):
+            _validate_view_common(view)
+            view = Raw(
+                name=view.name,
+                weight=_normalise_weight(view.weight),
+                terminal=view.terminal,
+            )
+            if view.name is not None and view.name not in component_names:
+                raise ValueError(
+                    f"Raw view '{name}' references unknown component "
+                    f"'{view.name}'."
+                )
+        elif isinstance(view, Group):
+            _validate_view_common(view)
+            if isinstance(view.members, str) or not view.members:
+                raise ValueError(f"Group view '{name}' requires members.")
+            members = tuple(view.members)
+            for member in members:
+                if member not in component_names:
+                    raise ValueError(
+                        f"Group view '{name}' references unknown component "
+                        f"'{member}'."
+                    )
+            view = Group(
+                members=members,
+                weight=_normalise_weight(view.weight),
+                terminal=view.terminal,
+            )
+        elif isinstance(view, (Total, ByState, ByKind)):
+            _validate_view_common(view)
+            view = type(view)(
+                weight=_normalise_weight(view.weight),
+                terminal=view.terminal,
+            )
+        else:
+            raise TypeError(
+                "cashflow views must be Raw, Group, Total, ByState, or ByKind; "
+                f"got {type(view)}."
+            )
+        frozen.append((name, view))
+    return tuple(frozen)