"""Treatment regimens: what a unit would have been given at every time point.
A *regime* (:mod:`cleverly.interventions`) says what one treatment decision would have
been. A **regimen** says what the whole sequence would have been --
:math:`\\bar a = (a_1, \\ldots, a_T)` -- and it is the thing a longitudinal fit's
parameters are indexed by. The two words are one letter apart and name different
objects, which is why this module spells the distinction out rather than reusing
:class:`~cleverly.interventions.base.Intervention`: an intervention is a density over
arms at *one* node, and a regimen is a plan across nodes.
Two kinds live here, and the difference between them is the whole reason this module is
not a tuple of floats. A :class:`Regimen` assigns the same arm to everybody at each
node. A :class:`DynamicRegimen` assigns node :math:`t`'s arm by a rule
:math:`d_t(H_t)` -- treat once the biomarker crosses a threshold, say -- so its
**followers are a covariate-dependent set that differs at every node**, and the rows each
sequential regression is fitted on move with the data rather than being a fixed slice.
Both answer :meth:`assignment`, which returns the same ``(n, T)`` object: for a static
regimen a *broadcast view* of its plan, costing nothing. Everything downstream reads
that matrix and nothing reads a scalar arm, which is what keeps the static path
bit-for-bit what it was before rules existed rather than a second implementation of it.
A rule is handed :meth:`~cleverly.longitudinal.data.LongitudinalData.history_frame` --
:math:`[W, L_1, \\ldots, L_t]`, in the backend the data came from, and nothing else.
Not the outcome, because reading it is not an intervention; not the earlier treatments,
because under the regimen those are what the rule itself assigned, and passing them would
let a rule read the treatment of a unit that *deviated*. The same restriction, enforced
the same way, as :func:`cleverly.interventions.base._covariate_frame` at one time point.
A rule can close over any estimate, such as a threshold at a sample mean, and no code can inspect a
closure. So a :class:`DynamicRegimen` declares ``rule_kind``, and :func:`refuse_regimen_rules`
refuses a plan with a callable node unless it is ``"known"``. The declaration is the one that
:class:`~cleverly.interventions.Rule` carries. A callable written inline in a ``regimens=`` mapping
carries no declaration, so a fit refuses it.
A node can also hold a **known policy density** :math:`q_t(\\cdot \\mid H_t)`, written as a
:class:`~cleverly.interventions.Stochastic` node. Such a node draws the arm rather than
assigning it. Its density function is handed
:meth:`~cleverly.longitudinal.data.LongitudinalData.policy_frame`, which adds the earlier
treatments to the history frame, because under a policy the earlier arms vary among the
units that remain on the plan. The density is evaluated once, in :func:`resolve_plans`,
and a :class:`Plan` carries it. At a policy node the plan's ``values`` column holds the
*observed* arm, so the mechanism is evaluated at the observed history, and
:meth:`Plan.masks` keeps the rows whose observed arm has positive policy probability. A
policy density that is one-hot on every reachable row is the rule it equals, and
:func:`resolve_plans` resolves it as that rule.
A node can hold an **odds tilt** of the treatment mechanism, written as an
:class:`~cleverly.interventions.Incremental` node (Kennedy 2019, Equation (1), and its
categorical form). Its density
:math:`q_t(a \\mid H_t) = \\delta_{t,a} \\pi_t(a \\mid H_t) / D_t` reads the fitted
mechanism :math:`\\pi_t`, so :func:`resolve_plans` records only the factors
:math:`\\delta_{t,a}`, resolved against the node's own levels, and :meth:`Plan.with_tilts`
attaches the density once a mechanism is fitted. At a tilt node the plan's ``values`` column
holds the observed arm, as at a policy node, and every row with a treatment there stays on
the plan. Only :class:`~cleverly.longitudinal.IPSI` fits such a plan.
"""
from __future__ import annotations
from collections.abc import Callable, Iterator, Mapping, Sequence
from dataclasses import dataclass
from functools import partial
from typing import TYPE_CHECKING, Any, TypeAlias, cast
import numpy as np
from .._declarations import FunctionKind
from .._typing import BoolArray, FloatArray, IntArray
from ..exceptions import CapabilityError, DataError, Refusal, RefusalReason
from ..interventions.base import (
_RULE_DECLARATION,
Stochastic,
_as_array,
check_regime_density,
refuse_regime_densities,
)
from ..interventions.incremental import Incremental
from ..interventions.policy import (
POLICY_TYPES,
ForbiddenRead,
JointComponent,
JointPolicy,
check_continuous_policy,
component_source,
describe_policy,
discrete_assignments,
joint_assignments,
joint_discrete_assignments,
lazy_frame,
policy_branches,
refuse_policy_declarations,
resolve_tuned_policy,
tuned_cap_parts,
tuned_maps,
with_resolved_component,
)
from ..interventions.tuning import (
FOLD_AVERAGE_ONE_FOLD,
DoseQuantile,
FittedTuner,
SupportBoundary,
fit_tuner,
refuse_outcome_tuners,
tuned_cap_at_atoms,
tuner_influence,
)
from ..learners.crossfit import Folds
if TYPE_CHECKING: # pragma: no cover - import cycle avoidance, types only
from ..interventions.policy import Policy
from .data import LongitudinalData, RegimenMasks
__all__ = [
"DynamicRegimen",
"Plan",
"Regimen",
"RegimenSpec",
"declares_mtp",
"declares_policy",
"declares_tilt",
"refuse_regimen_rules",
"resolve_plans",
"resolve_regimens",
"resolve_tuned_plans",
]
#: One node of a plan: a categorical label, a rule reading that node's history, a known
#: policy density over the node's levels, a modified treatment policy that reads the
#: unit's own treatment at the node, or an odds tilt of the node's treatment mechanism.
TreatmentLabel: TypeAlias = "str | np.str_ | bool | int | float | np.number"
RuleNode: TypeAlias = "TreatmentLabel | Callable[[Any], Any] | Stochastic | Policy | Incremental"
def _is_mtp(node: object) -> bool:
"""Whether a plan node is a modified treatment policy."""
return isinstance(node, POLICY_TYPES)
def _is_tilt(node: object) -> bool:
"""Whether a plan node is an odds tilt of the treatment mechanism."""
return isinstance(node, Incremental)
def _is_drawn(node: object) -> bool:
"""Whether a plan node draws or moves the arm rather than assigning it."""
return isinstance(node, Stochastic) or _is_mtp(node) or _is_tilt(node)
@dataclass(frozen=True)
class Regimen:
"""A static treatment plan: one arm per time point, under the user's own label.
Attributes
----------
label:
What the reported parameter is named by -- ``ey_regimen[always]``. Part of
the estimand rather than decoration: two regimens are two different
parameters, and the report has to be able to say which is which.
values:
The treatment label assigned at each time point, one per node.
"""
label: str
values: tuple[object, ...]
def __post_init__(self) -> None:
if not self.values:
raise DataError(f"regimen {self.label!r} assigns no treatment at any time point")
if any(_is_drawn(value) for value in self.values):
raise DataError(
f"regimen {self.label!r} holds a policy node. Write a plan with a "
"policy node as DynamicRegimen(label, plan), or pass it in regimens= as a "
"mapping value"
)
@property
def n_times(self) -> int:
return len(self.values)
def at(self, time: int) -> object:
"""The arm assigned at ``time``, counted from one."""
return self.values[time - 1]
def assignment(self, data: LongitudinalData) -> Any:
"""The ``(n, T)`` matrix of assigned *labels*, as a broadcast view of the plan.
A view rather than a copy, so reading a static regimen through the same matrix
interface a rule needs allocates nothing. ``object`` rather than ``float64``
because a label is whatever the analyst's treatment column held -- a string as
readily as a number -- and the dense codes a fit runs on are produced from this
by :meth:`~cleverly.longitudinal.data.LongitudinalData.encode_assignment`, which
reproduces the old float path exactly for a 0/1 node. That is where "a static
binary fit is unchanged bit for bit" is now delivered.
"""
return np.broadcast_to(np.asarray(self.values, dtype=object), (data.n, self.n_times))
def __repr__(self) -> str: # pragma: no cover - debugging aid
plan = "/".join(str(value) for value in self.values)
return f"Regimen({self.label!r}, {plan})"
[docs]
@dataclass(frozen=True)
class DynamicRegimen:
"""A plan whose nodes may be *rules* rather than constants.
``rule_kind="known"`` declares that every callable node is a fixed rowwise function of
one unit's available history, chosen independently of the analysis sample. A rule
that estimates a threshold, or otherwise aggregates across the sample, defines a
different, data-adaptive target, and inference for it needs conditions this path does
not check. No code can inspect a closure, so the declaration is the check.
:func:`refuse_regimen_rules` refuses ``None`` and ``"estimated"`` when the regimen is
built, and ``LTMLE.fit`` refuses them again before any learner.
.. code-block:: python
DynamicRegimen("treat once L2 rises", (0, lambda h: h["L2"] > 0), rule_kind="known")
A node can hold a known policy density instead, written as a
:class:`~cleverly.interventions.Stochastic` node. The unit then draws its arm at that
node with probability :math:`q_t(a \\mid H_t)`:
.. code-block:: python
DynamicRegimen("mix", (Stochastic(q1, "q1", density_kind="known"), "low"))
A rule node is handed the history frame, which holds no earlier treatment. A rule that
must read an arm drawn at an earlier policy node, such as "continue the arm drawn at the
first node", is written as a one-hot :class:`~cleverly.interventions.Stochastic` node.
That node reads the policy frame, and :func:`resolve_plans` resolves a one-hot density
as the rule it equals.
Parameters
----------
label : str
What the reported parameter is named by, exactly as for :class:`Regimen`.
plan : tuple
One entry per time point. An entry is a categorical treatment label, meaning that
label for everybody at that node, a callable :math:`d_t(H_t)` handed that node's
history frame and returning one arm per row, or a
:class:`~cleverly.interventions.Stochastic` node declared
``density_kind="known"``, whose density function is handed that node's policy frame
and returns one probability per level. Mixing them is the ordinary case: "treat at
the first node, then keep treating only while the biomarker stays high" is a
constant followed by a rule. The regimen stores the plan as a tuple, so a list or
an iterator is read once, when it is built. A single label, callable or policy is
not a plan and raises :class:`~cleverly.exceptions.DataError`.
rule_kind : {"known", "estimated"} or None
The declaration that every callable node of ``plan`` is a known function. One
declaration covers every node. ``"known"`` is the one value a fit accepts for a
plan with a callable node. ``None``, the default, and ``"estimated"`` raise
:class:`~cleverly.exceptions.CapabilityError` for such a plan, and a plan of
labels and policy nodes alone is exempt. Each policy node carries its own
``density_kind``. Any other value raises
:class:`~cleverly.exceptions.DataError`. It is the last field, so
``DynamicRegimen(label, plan)`` keeps its positional order.
Attributes
----------
n_times : int
has_policy : bool
"""
label: str
plan: tuple[RuleNode, ...]
rule_kind: FunctionKind | None = None
def __post_init__(self) -> None:
# The plan is stored as a tuple before any check, so the check reads the nodes that
# every later reader reads. An iterator is read here, once, and never again.
plan: Any = self.plan
nodes = _plan_nodes(self.label, tuple(plan) if isinstance(plan, Iterator) else plan)
if nodes is None:
raise DataError(
f"regimen {self.label!r} needs a plan with one entry per treatment node; got "
f"{plan!r}. Write one rule for every node as (rule,) * T, with T the number "
"of nodes"
)
object.__setattr__(self, "plan", nodes)
if not self.plan:
raise DataError(f"regimen {self.label!r} assigns no treatment at any time point")
refuse_regimen_rules(self)
@property
def n_times(self) -> int:
"""The number of treatment nodes, one for each entry of ``plan``."""
return len(self.plan)
[docs]
def is_rule(self, time: int) -> bool:
"""Whether ``time``'s arm is decided by a rule rather than declared.
Parameters
----------
time : int
The treatment node, counted from one.
Returns
-------
bool
``True`` when the entry of ``plan`` at ``time`` is callable.
"""
node = self.plan[time - 1]
return callable(node) and not _is_mtp(node)
[docs]
def is_mtp(self, time: int) -> bool:
"""Whether ``time``'s arm is a modified treatment policy of the unit's own treatment.
Parameters
----------
time : int
The treatment node, counted from one.
Returns
-------
bool
``True`` when the entry of ``plan`` at ``time`` is a policy object, such as
:class:`~cleverly.interventions.Shift` or
:class:`~cleverly.interventions.ModifiedPolicy`.
"""
return _is_mtp(self.plan[time - 1])
[docs]
def is_policy(self, time: int) -> bool:
"""Whether ``time``'s arm is drawn from a known policy density.
Parameters
----------
time : int
The treatment node, counted from one.
Returns
-------
bool
``True`` when the entry of ``plan`` at ``time`` is a
:class:`~cleverly.interventions.Stochastic` node.
"""
return isinstance(self.plan[time - 1], Stochastic)
[docs]
def is_tilt(self, time: int) -> bool:
"""Whether ``time``'s arm is drawn from an odds tilt of the treatment mechanism.
Parameters
----------
time : int
The treatment node, counted from one.
Returns
-------
bool
``True`` when the entry of ``plan`` at ``time`` is an
:class:`~cleverly.interventions.Incremental` node.
"""
return _is_tilt(self.plan[time - 1])
@property
def has_policy(self) -> bool:
"""Whether any node of ``plan`` is a stochastic policy or a modified treatment policy."""
return any(isinstance(node, Stochastic) or _is_mtp(node) for node in self.plan)
@property
def has_tilt(self) -> bool:
"""Whether any node of ``plan`` is an odds tilt of the treatment mechanism."""
return any(_is_tilt(node) for node in self.plan)
@property
def has_rule(self) -> bool:
"""Whether any node of ``plan`` is a rule :math:`d_t(H_t)`."""
return any(callable(node) and not _is_drawn(node) for node in self.plan)
@property
def has_mtp(self) -> bool:
"""Whether any node of ``plan`` is a modified treatment policy."""
return any(_is_mtp(node) for node in self.plan)
[docs]
def assignment(self, data: LongitudinalData) -> Any:
"""Evaluate every node's rule and return the ``(n, T)`` arm matrix.
Called **once** per fit, in :meth:`cleverly.longitudinal.LTMLE.fit`, and the
matrix is what every mask, mechanism design and clever covariate then reads.
Evaluating it once rather than at each use is not only cheaper: a rule that is
not a deterministic function of the frame would otherwise let the follower masks
disagree with the designs the mechanism was evaluated at, and the fit would be
answering for no single regimen at all. The callable receives a frame for ergonomic
vectorization, but its supported contract is rowwise: it must not read sample summaries
or learn a rule from those rows.
A rule is asked for an arm on every row that is still *in the study* before the
node -- uncensored through ``t - 1``, and on a survival fit event-free through
``t - 1`` as well, since a unit that has had the event has no treatment decision
at ``t`` for a rule to make. Off that set the history is the zero fill
:meth:`~cleverly.longitudinal.data.LongitudinalData.covariate_history` puts there,
so whatever the rule returns is meaningless; it is replaced by a sentinel rather than
validated, because such a row is masked out of every regression and every
influence curve, and the only way it could still matter is by putting a ``nan``
into a design matrix that a learner is called on.
A policy node assigns no arm. Its column holds the label each unit was *observed*
to receive, because the mechanism of a policy plan is evaluated at the observed
history. :meth:`policy_density` evaluates the policy itself.
It runs :func:`refuse_regimen_rules` before it calls any rule, because a modified
regimen can reach it directly with a declaration this version refuses.
Parameters
----------
data : LongitudinalData
The validated panel, which supplies each node's history frame and the rows
still in the study.
Returns
-------
ndarray
``(n, T)`` object array of treatment labels, one column per node. A row that
has left the study before a node holds ``None`` at that node.
Raises
------
CapabilityError
If a node is callable and ``rule_kind`` is ``None`` or ``"estimated"``, or a
policy node does not declare ``density_kind="known"``.
DataError
If ``rule_kind`` is not one of the three states, or a rule raises or returns
other than one arm per row.
"""
refuse_regimen_rules(self)
columns = []
for time, node in enumerate(self.plan, start=1):
reachable = data.uncensored_through(time - 1) & data.event_free_through(time - 1)
if _is_drawn(node):
arms = data.observed_labels(time)
elif callable(node):
arms = self._evaluate(node, data, time, reachable)
else:
arms = np.full(data.n, node, dtype=object)
columns.append(np.where(reachable, arms, None))
return np.column_stack(columns)
[docs]
def policy_density(self, data: LongitudinalData, time: int) -> FloatArray:
"""Evaluate the policy node at ``time`` and check it before it is used.
The density function is handed
:meth:`~cleverly.longitudinal.data.LongitudinalData.policy_frame`. It returns an
``(n, K)`` array whose columns follow ``data.treatment_levels[time - 1]``, or a
dataframe whose columns are exactly those levels, which this reorders by level. The
label route exists because a permutation of the columns of an array is silent.
The check reads the rows still in the study before the node: uncensored and
event-free through ``time - 1``. Every other row is set to a zero row, which no
regression and no influence curve reads.
Parameters
----------
data : LongitudinalData
The validated panel.
time : int
The policy node, counted from one.
Returns
-------
ndarray
``(n, K)`` policy probabilities, zero on the rows that left the study.
Raises
------
CapabilityError
If the node does not declare ``density_kind="known"``.
DataError
If ``time`` is not a policy node, if the density function raises, or if it
returns other than a probability simplex per reachable row over the levels.
"""
node = self.plan[time - 1]
if not isinstance(node, Stochastic):
raise DataError(f"node {time} of regimen {self.label!r} is not a policy node")
refuse_regime_densities((node,))
label = f"the policy at time {time} of regimen {self.label!r}"
levels = tuple(data.treatment_levels[time - 1])
name = data.treatment_names[time - 1]
frame = data.policy_frame(time)
try:
returned = node.density_fn(frame)
except Exception as error:
raise DataError(
f"{label} raised {type(error).__name__}: {error}. It is handed "
"[W, L_1, ..., L_t] and the earlier treatments, which at time "
f"{time} are {list(data.policy_names(time))}. A policy reading a covariate "
"measured later can only be used from that node on, so pass a plan with one "
"entry per node."
) from error
if hasattr(returned, "columns") and not isinstance(returned, np.ndarray):
returned = _frame_by_level(returned, levels, label=label, name=name)
try:
values = np.asarray(_as_array(returned), dtype=float)
except (TypeError, ValueError) as error:
raise DataError(f"{label} returned a value that is not a probability") from error
expected = (data.n, len(levels))
if values.shape != expected:
raise DataError(
f"{label} returned shape {values.shape}; expected {expected}. One column per "
f"level of treatment column {name!r}, in the order {levels!r}."
)
reachable = data.uncensored_through(time - 1) & data.event_free_through(time - 1)
check_regime_density(values[reachable], label=label)
return np.where(reachable[:, None], values, 0.0)
def _evaluate(
self, rule: Callable[[Any], Any], data: LongitudinalData, time: int, reachable: Any
) -> Any:
"""One rule, called on its node's history frame and checked before it is used."""
names = data.history_names(time)
try:
returned = rule(data.history_frame(time))
except Exception as error:
raise DataError(
f"the rule at time {time} of regimen {self.label!r} raised "
f"{type(error).__name__}: {error}. It is handed [W, L_1, ..., L_t], which "
f"at time {time} is {list(names)} -- a rule reading a covariate measured "
"later can only be used from that node on, so pass a plan with one entry "
"per node rather than a single rule for all of them"
) from error
arms = np.asarray(_as_array(returned), dtype=object)
if arms.ndim != 1 or arms.shape[0] != data.n:
raise DataError(
f"the rule at time {time} of regimen {self.label!r} returned "
f"{arms.shape} assignments for {data.n} rows; it must return one arm per row"
)
return arms
def __repr__(self) -> str: # pragma: no cover - debugging aid
return f"DynamicRegimen({self.label!r}, {describe_plan(self)})"
def _frame_by_level(frame: Any, levels: tuple[object, ...], *, label: str, name: str) -> Any:
"""A policy's dataframe as an ``(n, K)`` array in level order, or a ``DataError``.
The columns must be exactly the node's levels, matched by ``==`` on the analyst's own
labels. A frame with any other column set is refused, because no order can be read
from a column the node does not have.
"""
columns = list(frame.columns)
order: list[int] = []
for level in levels:
matches = [index for index, column in enumerate(columns) if _same_label(column, level)]
if len(matches) != 1:
break
order.append(matches[0])
if len(columns) != len(levels) or len(order) != len(levels) or len(set(order)) != len(levels):
raise DataError(
f"{label} returned columns {columns!r}. A frame must have exactly the levels "
f"{list(levels)!r} of treatment column {name!r}, or return an array in that "
"column order."
)
return np.column_stack([_as_array(frame[columns[index]]) for index in order])
def _same_label(column: object, level: object) -> bool:
"""Whether a frame column names a treatment level, by the label's own ``==``."""
try:
return bool(column == level)
except (TypeError, ValueError):
return False
#: A regimen of either kind. A closed union rather than a Protocol: :func:`resolve_regimens`
#: is the only constructor, so there is no third-party regimen type to admit -- and an open
#: one would invite into this module the very extensions it refuses by name elsewhere.
RegimenSpec: TypeAlias = "Regimen | DynamicRegimen"
def describe_plan(regimen: RegimenSpec) -> str:
"""A plan as ``1/0`` for constants, ``d`` for a rule and ``q`` for a policy."""
if isinstance(regimen, Regimen):
return "/".join(str(value) for value in regimen.values)
return "/".join(_describe_node(node) for node in regimen.plan)
def _describe_node(node: RuleNode) -> str:
"""One node of a plan: its arm, or the rule's or the policy's own name.
A ``def``\\ -ed rule carries the name the analyst gave it, and reporting ``d:responds``
rather than a bare ``d`` costs nothing and says which of two rules was run. A lambda
has no such name -- ``__name__`` is ``"<lambda>"``, which names nothing -- so it falls
back to ``d`` and the plan fingerprint on the config is what tells two of them apart.
A policy node prints ``q:`` and the name its :class:`~cleverly.interventions.Stochastic`
carries. An odds tilt node prints ``o:`` and the name its
:class:`~cleverly.interventions.Incremental` carries.
"""
if isinstance(node, Stochastic):
return f"q:{node.name}" if node.name else "q"
if _is_tilt(node):
return f"o:{getattr(node, 'name', '')}"
if _is_mtp(node):
return f"m:{getattr(node, 'name', '')}"
if not callable(node):
return str(node)
name = getattr(node, "__name__", "<lambda>")
return "d" if name == "<lambda>" else f"d:{name}"
def tilted_density_binary(pi: FloatArray, delta: float) -> FloatArray:
r"""Kennedy's tilted density :math:`(1 - q, q)` on two levels, from :math:`\pi(1 \mid H)`.
Parameters
----------
pi : FloatArray
``(n,)`` mechanism probability of the second level.
delta : float
The odds multiplier :math:`\delta = \delta_1/\delta_0`.
Returns
-------
FloatArray
``(n, 2)`` density, :math:`q = \delta\pi/(\delta\pi + 1 - \pi)`.
"""
tilted = delta * pi / (delta * pi + 1.0 - pi)
return np.column_stack([1.0 - tilted, tilted])
def tilted_density_general(g: FloatArray, factors: FloatArray) -> FloatArray:
r"""The tilted density :math:`q(a) = \delta_a g(a)/\sum_b \delta_b g(b)` for any ``K``.
Parameters
----------
g : FloatArray
``(n, K)`` mechanism.
factors : FloatArray
``(K,)`` odds factors.
Returns
-------
FloatArray
``(n, K)`` density.
"""
return np.asarray(g * factors[None, :] / (g @ factors)[:, None])
@dataclass(frozen=True)
class Plan:
"""A regimen together with the ``(n, T)`` arms it assigns *this* sample.
The pair travels as one object so that nothing downstream can reach a rule and call
it a second time: :func:`~cleverly.longitudinal.sequential.fit_mechanism` and
:func:`~cleverly.longitudinal.sequential.fit_regimen` see arms, never callables.
``values`` holds the *dense codes* the container assigned, not the analyst's labels.
The labels are recoverable from
:attr:`~cleverly.longitudinal.data.LongitudinalData.treatment_levels`, which every
reader of this plan already holds, and carrying a second ``(n, T)`` object array
beside the codes would be a copy that only some of them kept in step.
At a policy node ``values`` holds the code each unit was observed to receive, and
``policy`` holds the evaluated density. So the mechanism is evaluated at the observed
history, :meth:`arm` is the observed arm, and :meth:`intervention_density` reads the
policy at that arm. A plan with no policy node has ``policy=()``, and every reader then
runs the code a deterministic plan runs.
Parameters
----------
regimen : Regimen or DynamicRegimen
The resolved regimen.
values : FloatArray
``(n, T)`` dense codes: the assigned arm at a label or rule node, and the observed
arm at a policy node.
policy : tuple of FloatArray or None
One entry per node: the ``(n, K_t)`` policy density at a policy node and ``None``
elsewhere. Empty when the plan has no policy node.
point_mass_nodes : tuple of int
The policy nodes whose density was one-hot on every reachable row, and which were
therefore resolved as the rule they equal.
mtp_nodes : tuple of int
The nodes that hold a modified treatment policy. At such a node ``values`` holds
the observed code or dose, and ``policy`` holds the ``(n, J)`` weights of the
predictions the node carries: the levels the policy sends a unit to at a
categorical node, and the randomizer's probabilities at a continuous or mixed one.
policy_targets : tuple
Per node, the ``(n,)`` dose of each branch at a continuous policy node, the
``(n, d)`` component values of each branch at a mixed node, and ``None``
elsewhere. Empty without one.
mtp_assignments : tuple
Per node, each branch of a categorical policy node as ``(probability, codes)``,
with ``codes`` the ``(n, K)`` level map of the discrete formula, and ``None``
elsewhere. Empty without one.
mtp_numerators : tuple
Per node, the ratio numerator of a policy node once the mechanism is fitted, and
``None`` elsewhere. Empty until :func:`attach_policy_numerators` sets it.
node_factors : tuple
Per node, the ``(K_t,)`` tuple of odds factors :math:`\\delta_{t,a}` of a tilt node,
in arm-code order, and ``None`` elsewhere. Empty without a tilt node. At a tilt
node ``values`` holds the observed code, and ``policy`` holds the tilted density once
:meth:`with_tilts` attaches it.
Attributes
----------
label : str
has_policy : bool
has_mtp : bool
has_tilt : bool
"""
regimen: RegimenSpec
values: FloatArray
policy: tuple[FloatArray | None, ...] = ()
point_mass_nodes: tuple[int, ...] = ()
#: The nodes, counted from one, that hold a modified treatment policy.
mtp_nodes: tuple[int, ...] = ()
#: Per node, the ``(n,)`` dose each branch of a continuous node's policy assigns, and
#: ``None`` elsewhere. Empty without a continuous policy node.
policy_targets: tuple[tuple[FloatArray, ...] | None, ...] = ()
#: Per node, each branch of a categorical node's policy as ``(probability, codes)``
#: with ``codes`` the ``(n, K)`` level map of the discrete formula, and ``None``
#: elsewhere. Empty without a categorical policy node.
mtp_assignments: tuple[tuple[tuple[float, IntArray], ...] | None, ...] = ()
#: Per node, the ``(n,)`` ratio numerator of a policy node once the mechanism is fitted:
#: :math:`g^d_t(A_t \mid H_t)` at a categorical node and :math:`r_t = g^d_t / g_t` at a
#: continuous one. Empty until :func:`attach_policy_numerators` sets it.
mtp_numerators: tuple[FloatArray | None, ...] = ()
#: Per node, the ``(K_t,)`` odds factors of a tilt node in arm-code order, and ``None``
#: elsewhere. Empty without a tilt node.
node_factors: tuple[tuple[float, ...] | None, ...] = ()
#: The tuned nodes of a plan whose caps a fit read from the treatment law (X40 part
#: (e)), in time order. Empty on every other plan.
tuning: tuple[NodeTuning, ...] = ()
#: Per outer fold, the plan at the tuners fitted on that fold's training rows. Empty
#: in sample and on a plan with no tuned node.
fold_plans: tuple[Plan, ...] = ()
#: The outer split the fold plans belong to, or ``None``.
folds: Folds | None = None
#: Whether the plan's target is the fold average: then row ``i``'s targets are those of
#: the fold that holds row ``i`` out, and each fold is targeted on its own rows.
fold_average: bool = False
@property
def label(self) -> str:
return self.regimen.label
@property
def has_tilt(self) -> bool:
"""Whether any node of this plan is an odds tilt of the treatment mechanism."""
return any(factors is not None for factors in self.node_factors)
def is_tilt_node(self, time: int) -> bool:
"""Whether node ``time``, counted from one, is an odds tilt.
Parameters
----------
time : int
The treatment node, counted from one.
Returns
-------
bool
``True`` at a tilt node, whether or not its density is attached.
"""
return bool(self.node_factors) and self.node_factors[time - 1] is not None
def tilt_factors(self, time: int) -> FloatArray:
"""The odds factors of tilt node ``time``, one per level in arm-code order.
Parameters
----------
time : int
A tilt node, counted from one.
Returns
-------
FloatArray
``(K_t,)`` strictly positive factors :math:`\\delta_{t,a}`. A scalar declaration
on two levels reads ``(1, delta)``.
Raises
------
ValueError
If ``time`` is not a tilt node.
"""
factors = self.node_factors[time - 1] if self.node_factors else None
if factors is None:
raise ValueError(f"node {time} of regimen {self.label!r} is not a tilt node")
return np.asarray(factors, dtype=float)
def with_tilts(self, propensity: Mapping[int, FloatArray]) -> Plan:
r"""This plan with each tilt node's density, from the mechanism :math:`\pi_t`.
At a tilt node the density is
:math:`q_t(a \mid H_t) = \delta_{t,a}\,\pi_t(a \mid H_t) / D_t` with
:math:`D_t = \sum_b \delta_{t,b}\,\pi_t(b \mid H_t)`. On two levels it is evaluated
in the form of Kennedy (2019), Equation (1), on :math:`\pi_t(1 \mid H_t)` with
:math:`\delta = \delta_{t,1}/\delta_{t,0}`, so a binary node is bit-identical to the
scalar form. The node then reads as a policy node: the regression holds the
current arm as a column, and the node carries :math:`\sum_a q_t(a) m_t(a)`.
Parameters
----------
propensity : mapping of int to FloatArray
By tilt node, counted from one, the ``(n, K_t)`` mechanism in arm-code order.
Returns
-------
Plan
The plan with its tilt densities attached.
"""
from dataclasses import replace
n_times = int(self.values.shape[1])
policy = list(self.policy) if self.policy else [None] * n_times
for time in range(1, n_times + 1):
if not self.is_tilt_node(time):
continue
g = np.clip(np.asarray(propensity[time], dtype=float), 0.0, 1.0)
factors = self.tilt_factors(time)
if g.shape[1] == 2:
policy[time - 1] = tilted_density_binary(g[:, 1], factors[1] / factors[0])
else:
policy[time - 1] = tilted_density_general(g, factors)
return replace(self, policy=tuple(policy))
@property
def has_mtp(self) -> bool:
"""Whether any node of this plan is a modified treatment policy."""
return bool(self.mtp_nodes)
def is_mtp_node(self, time: int) -> bool:
"""Whether node ``time``, counted from one, is a modified treatment policy.
Parameters
----------
time : int
The treatment node, counted from one.
Returns
-------
bool
``True`` at a modified treatment policy node.
"""
return time in self.mtp_nodes
def carry_targets(self, time: int) -> tuple[FloatArray, ...] | None:
"""The values a continuous or mixed policy node's regression is predicted at, by branch.
Parameters
----------
time : int
The treatment node, counted from one.
Returns
-------
tuple of FloatArray or None
One ``(n,)`` dose per branch at a continuous policy node, one ``(n, d)`` block of
component values per branch at a mixed node, and ``None`` elsewhere.
"""
if not self.policy_targets:
return None
return self.policy_targets[time - 1]
@property
def has_policy(self) -> bool:
"""Whether any node of this plan draws its arm from an evaluated policy density."""
return any(entry is not None for entry in self.policy)
def is_policy_node(self, time: int) -> bool:
"""Whether node ``time``, counted from one, carries an evaluated policy density.
Parameters
----------
time : int
The treatment node, counted from one.
Returns
-------
bool
``True`` at a policy node that was not resolved as a rule.
"""
return bool(self.policy) and self.policy[time - 1] is not None
def policy_at(self, time: int) -> FloatArray:
"""The ``(n, K_t)`` policy density at node ``time``.
Parameters
----------
time : int
A policy node, counted from one.
Returns
-------
FloatArray
The density evaluated in :func:`resolve_plans`.
"""
density = self.policy[time - 1] if self.policy else None
if density is None:
raise ValueError(f"node {time} of regimen {self.label!r} is not a policy node")
return density
def arm(self, time: int) -> FloatArray:
"""The arm this plan assigns each unit at ``time``, counted from one."""
return self.values[:, time - 1]
def intervention_density(self, data: LongitudinalData) -> FloatArray:
r"""``(n, T)`` :math:`\pi_t(A_t \mid H_t)`, the plan's density at the observed arm.
At a label or rule node it is :math:`1\{A_t = d_t(H_t)\}`. At a policy node it is
:math:`q_t(A_t \mid H_t)`. A row with no treatment at a node holds zero there.
Parameters
----------
data : LongitudinalData
The panel the plan was resolved against.
Returns
-------
FloatArray
One column per node.
"""
return np.column_stack(
[_node_numerator(self, data, time) for time in range(1, data.n_times + 1)]
)
def support(self, data: LongitudinalData) -> BoolArray:
r"""``(n, T)`` indicator that a row stays on the plan at each node.
At a label or rule node it is :math:`1\{A_t = d_t(H_t)\}`, at a stochastic policy
node :math:`q_t(A_t \mid H_t) > 0`, and at a modified treatment policy node every
row with a treatment there: the policy reads the unit's own treatment, so no row
leaves the plan, and a row whose ratio is zero keeps a zero term. A tilt node also
keeps every row with a treatment there: its factor
:math:`(\delta_t A_t + 1 - A_t) / (\delta_t \pi_t + 1 - \pi_t)` is positive.
Parameters
----------
data : LongitudinalData
The panel the plan was resolved against.
Returns
-------
BoolArray
Booleans as one column per node.
"""
columns = []
for time in range(1, data.n_times + 1):
if self.is_mtp_node(time) or self.is_tilt_node(time):
columns.append(data.node_present(time))
else:
columns.append(_node_numerator(self, data, time) > 0.0)
return np.column_stack(columns)
def cumulative_numerator(self, data: LongitudinalData) -> FloatArray | None:
r"""``(n, T)`` running product :math:`\prod_{s \le t} \pi_s(A_s \mid H_s)`, or ``None``.
The numerator of the cumulative ratio :math:`R_t`. ``None`` on a plan without a
policy node, whose ratio numerator is the indicator that the masks already apply.
Parameters
----------
data : LongitudinalData
The panel the plan was resolved against.
Returns
-------
FloatArray or None
One column per node, or ``None`` without a policy node.
"""
if not self.has_policy:
return None
density = self.intervention_density(data)
running = np.ones(data.n)
columns = []
for time in range(data.n_times):
running = running * density[:, time]
columns.append(running)
return np.column_stack(columns)
def masks(self, data: LongitudinalData) -> RegimenMasks:
"""The plan's prefix masks: support-consistent, uncensored and event-free.
Without a policy node this is ``data.regimen_masks(values)`` itself, so a
deterministic plan keeps its exact arrays. With one, a row follows the plan
through ``t`` when its observed arm has positive intervention density at every
node up to ``t``. At a label or rule node that is the follower indicator, and at
a policy node it keeps every row whose observed arm the policy can draw.
Parameters
----------
data : LongitudinalData
The panel the plan was resolved against.
Returns
-------
RegimenMasks
The prefix scans every node of the recursion reads.
"""
if not self.has_policy:
return data.regimen_masks(self.values)
return data.regimen_masks(self.values, support=self.support(data))
def _node_numerator(plan: Plan, data: LongitudinalData, time: int) -> FloatArray:
r""":math:`\pi_t(A_t \mid H_t)` at the observed arm, the ratio numerator at node ``time``.
The one place the numerator is read. At a policy node it reads the evaluated density
at the observed code, and elsewhere the indicator that the observed arm is the assigned
one. A row with no treatment at the node, censored or past an event, holds zero.
Parameters
----------
plan : Plan
The resolved plan.
data : LongitudinalData
The panel the plan was resolved against.
time : int
The treatment node, counted from one.
Returns
-------
FloatArray
One value per row.
"""
present = data.node_present(time)
if plan.is_mtp_node(time):
numerator = plan.mtp_numerators[time - 1] if plan.mtp_numerators else None
if numerator is None:
raise ValueError(
f"node {time} of regimen {plan.label!r} is a modified treatment policy, whose "
"ratio numerator reads the fitted mechanism; attach it first"
)
return np.where(present, numerator, 0.0)
if plan.is_tilt_node(time) and not plan.is_policy_node(time):
raise ValueError(
f"node {time} of regimen {plan.label!r} is an odds tilt, whose density reads the "
"fitted mechanism; attach it first with Plan.with_tilts"
)
observed = data.node_codes(time)
if plan.is_policy_node(time):
codes = np.nan_to_num(observed, nan=0.0).astype(np.int64)
density = plan.policy_at(time)
return np.where(present, density[np.arange(data.n), codes], 0.0)
return np.asarray(observed == plan.values[:, time - 1], dtype=float)
def resolve_plans(
regimens: Sequence[RegimenSpec],
data: LongitudinalData,
*,
_collapse_point_masses: bool = True,
allow_set_dose: bool = False,
) -> tuple[Plan, ...]:
"""Evaluate every regimen against ``data``, once, before any nuisance is fitted.
Each rule and each policy density is called here and nowhere else. A policy density
that is one-hot on every reachable row, by exact equality, is resolved as the rule it
equals: its ``values`` column holds the argmax codes, and its ``policy`` entry is
``None``. The node is recorded on :attr:`Plan.point_mass_nodes`. A plan whose every
policy node collapses runs exactly the code its rule regimen runs.
Parameters
----------
regimens : sequence of Regimen or DynamicRegimen
The resolved regimens, in report order.
data : LongitudinalData
The validated panel.
_collapse_point_masses : bool, default=True
Private. ``False`` keeps a one-hot policy on the policy path, which a test uses to
check that path against the rule it equals.
allow_set_dose : bool, default=False
Whether a fixed dose or a known rule that sets a dose, at a continuous node, is
resolved. Only a model-based estimator accepts one, because the nonparametric target is
not pathwise differentiable.
Returns
-------
tuple of Plan
One per regimen, in the order given.
"""
return tuple(
_resolve_plan(regimen, data, collapse=_collapse_point_masses, allow_set_dose=allow_set_dose)
for regimen in regimens
)
def _resolve_plan(
regimen: RegimenSpec, data: LongitudinalData, *, collapse: bool, allow_set_dose: bool = False
) -> Plan:
"""One regimen's plan, with its policy densities evaluated and point masses collapsed."""
refuse_continuous_assignments(regimen, data, allow_set_dose=allow_set_dose)
values = data.encode_assignment(regimen.assignment(data), regimen.label)
if not (isinstance(regimen, DynamicRegimen) and (regimen.has_policy or regimen.has_tilt)):
return Plan(regimen, values)
values = np.array(values, dtype=float, copy=True)
policy: list[FloatArray | None] = []
collapsed: list[int] = []
mtp_nodes: list[int] = []
targets: list[tuple[FloatArray, ...] | None] = []
assignments: list[tuple[tuple[float, IntArray], ...] | None] = []
factors: list[tuple[float, ...] | None] = []
for time in range(1, data.n_times + 1):
if data.is_held_node(time):
# An identity node of a held design: never a policy or MTP node.
policy.append(None)
targets.append(None)
assignments.append(None)
factors.append(None)
continue
if regimen.is_tilt(time):
node = regimen.plan[time - 1]
assert isinstance(node, Incremental)
policy.append(None)
targets.append(None)
assignments.append(None)
# Per-node resolution: the declared factors read against this node's own levels,
# so a scalar on a node with other than two levels (R1) and a key that is not a
# level of this node (R2) refuse here, before any learner. A mixed or continuous
# node was refused above, by refuse_continuous_assignments.
odds = node.factors(
data.treatment_levels[time - 1], subject=f"node {data.treatment_names[time - 1]!r}"
)
factors.append(tuple(float(value) for value in odds))
continue
factors.append(None)
if regimen.is_mtp(time):
carry, target, assigned = _resolve_mtp_node(regimen, data, time)
policy.append(carry)
targets.append(target)
assignments.append(assigned)
mtp_nodes.append(time)
continue
targets.append(None)
assignments.append(None)
if not regimen.is_policy(time):
policy.append(None)
continue
density = regimen.policy_density(data, time)
reachable = data.uncensored_through(time - 1) & data.event_free_through(time - 1)
if collapse and _is_point_mass(density[reachable]):
values[:, time - 1] = np.where(reachable, np.argmax(density, axis=1), 0.0)
policy.append(None)
collapsed.append(time)
else:
policy.append(density)
kept = tuple(policy) if any(entry is not None for entry in policy) else ()
return Plan(
regimen,
values,
kept,
tuple(collapsed),
mtp_nodes=tuple(mtp_nodes),
policy_targets=tuple(targets) if any(t is not None for t in targets) else (),
mtp_assignments=tuple(assignments) if any(a is not None for a in assignments) else (),
node_factors=tuple(factors) if any(f is not None for f in factors) else (),
)
@dataclass(frozen=True)
class NodeTuning:
"""One tuned node of a plan: the tuner fitted on the at-risk rows, and its curve parts.
Parameters
----------
time : int
The node, counted from one.
policy : object
The declared tuned :class:`~cleverly.interventions.Shift` or
:class:`~cleverly.interventions.Scale`.
fitted : FittedTuner
The tuner fitted on every row at risk at the node: :math:`\\theta_n`.
fold_fits : tuple of FittedTuner
Per outer fold, the tuner fitted on that fold's training rows at risk. Empty in
sample.
influence : FloatArray or None
``(n, p)`` the tuner's influence function on the population target, ``None`` on the
fold-average target.
boundary : FloatArray
``(n,)`` :math:`c_t(H_t)` at :math:`\\theta_n`.
edge : FloatArray
``(n,)`` :math:`e_t(H_t)`, the last dose the cap lets move.
gradient : FloatArray
``(n, p)`` :math:`\\partial c_t / \\partial\\theta_t`.
slope : float
:math:`|e_t'|`.
at_risk : BoolArray
``(n,)`` the rows at risk at the node.
component : str or None
At a mixed node, the tuned dose component of the node's
:class:`~cleverly.interventions.JointPolicy`; ``None`` at a node of one dose.
slots : tuple of tuple
``(branch, probability)`` of each branch of the node's policy in which this map
sets the dose. The curve carries one derivative term per slot, weighted by the
branch probability. ``((0, 1.0),)`` at a deterministic node.
"""
time: int
policy: Any
fitted: FittedTuner
fold_fits: tuple[FittedTuner, ...]
influence: FloatArray | None
boundary: FloatArray
edge: FloatArray
gradient: FloatArray
slope: float
at_risk: BoolArray
component: str | None = None
slots: tuple[tuple[int, float], ...] = ((0, 1.0),)
#: What a fold-average plan appends to its regimen's label in every reported name.
FOLD_AVERAGE_SUFFIX = ", fold average"
def declares_tuned(regimens: Sequence[RegimenSpec]) -> bool:
"""Whether any regimen holds a map whose cap is a tuner.
Parameters
----------
regimens : sequence of Regimen or DynamicRegimen
The resolved regimens.
Returns
-------
bool
``True`` when some node is a tuned :class:`~cleverly.interventions.Shift` or
:class:`~cleverly.interventions.Scale`, or a
:class:`~cleverly.interventions.JointPolicy` with a tuned component.
"""
return any(_tuned(regimen) for regimen in regimens)
def _tuned(regimen: RegimenSpec) -> bool:
return isinstance(regimen, DynamicRegimen) and any(tuned_maps(node) for node in regimen.plan)
def _regimen_target(regimen: DynamicRegimen) -> str | None:
targets = {str(tuned.target) for node in regimen.plan for _, tuned, _ in tuned_maps(node)}
if len(targets) > 1:
raise CapabilityError(
f"regimen {regimen.label!r} declares the targets {sorted(targets)} on its tuned "
"maps. One regimen is one parameter, so its tuned maps declare one target",
reason=RefusalReason.INVALID_DECLARATION,
)
return next(iter(targets)) if targets else None
def resolve_tuned_plans(
regimens: Sequence[RegimenSpec],
data: LongitudinalData,
folds: Folds,
) -> tuple[Plan, ...]:
"""Resolve every regimen, fitting each tuned cap before any learner (X40 part (e)).
A regimen without a tuned node is :func:`resolve_plans` exactly. At a tuned node the
tuner is fitted on the rows at risk at the node, on their observed doses and policy
frame. ``target="population"`` resolves the plan at that fit, :math:`\\theta_n`, and
records the tuner's influence function. On a cross-fitted fit each outer fold also
gets the plan at its training rows' fit, :math:`\\theta^{(-k)}`, which that fold's
untargeted recursion and its classifier read (``test_t_e7_each_fold_trains_at_its_own_tuner``).
``target="fold_average"`` resolves one plan per outer fold, at its training rows' fit,
and the returned plan holds, at row ``i``, the doses of the fold that holds ``i`` out.
Each fold plan runs the declared-support and atom checks over every row, since a fold's
training rows take its map in their pseudo-outcomes.
Parameters
----------
regimens : sequence of Regimen or DynamicRegimen
The resolved regimens, in report order.
data : LongitudinalData
The validated panel.
folds : Folds
The outer split.
Returns
-------
tuple of Plan
One per regimen, in the order given.
Raises
------
CapabilityError
R-e3 for a tuner that reads an outcome column, R-e4 for the fold-average target on
one fold, a tuned cap at a node that does not hold one dose or a dose component, a
tuned cap at a dose with declared atoms, and the joint compositions of
:func:`~cleverly.interventions.policy.tuned_components`.
"""
plans = []
for regimen in regimens:
if not _tuned(regimen):
plans.append(_resolve_plan(regimen, data, collapse=True))
continue
assert isinstance(regimen, DynamicRegimen)
plans.append(_resolve_tuned_plan(regimen, data, folds))
return tuple(plans)
def _outcome_columns(data: LongitudinalData) -> tuple[str, ...]:
names = [data.outcome_name] if getattr(data, "outcome_name", None) else []
names.extend(getattr(data, "event_names", None) or ())
return tuple(str(name) for name in names if name)
def _resolve_tuned_plan(regimen: DynamicRegimen, data: LongitudinalData, folds: Folds) -> Plan:
from dataclasses import replace
target = _regimen_target(regimen)
if target == "fold_average":
# The fold-average target is a different parameter from the population one, so its
# report names it: ``ey_regimen[<label>, fold average]``.
regimen = replace(regimen, label=f"{regimen.label}{FOLD_AVERAGE_SUFFIX}")
tuned_nodes = [
(time, node, component, tuned, slots)
for time, node in enumerate(regimen.plan, start=1)
for component, tuned, slots in tuned_maps(node)
]
refuse_outcome_tuners(
[
(tuned.name, cast(DoseQuantile | SupportBoundary, tuned.cap))
for _, _, _, tuned, _ in tuned_nodes
],
_outcome_columns(data),
)
if target == "fold_average" and folds.is_single:
raise CapabilityError(FOLD_AVERAGE_ONE_FOLD, reason=RefusalReason.INVALID_DECLARATION)
sources = {
(time, component): _tuner_source(regimen, data, time, node, component)
for time, node, component, _, _ in tuned_nodes
}
weights = data.weights
tunings: list[NodeTuning] = []
in_sample: dict[int, RuleNode] = {}
by_fold: list[dict[int, RuleNode]] = [
{} for _ in range(0 if folds.is_single else folds.n_folds)
]
for time, node, component, tuned, slots in tuned_nodes:
tuner = cast(DoseQuantile | SupportBoundary, tuned.cap)
doses, frame = sources[time, component]
reachable = data.uncensored_through(time - 1) & data.event_free_through(time - 1)
fitted = _guarded(
regimen, time, frame, partial(fit_tuner, tuner, doses, frame, reachable, weights)
)
fold_fits: list[FittedTuner] = []
if not folds.is_single:
for k, (train, _) in enumerate(folds):
rows = np.zeros(data.n, dtype=bool)
rows[train] = True
fold_fit = fit_tuner(tuner, doses, frame, reachable & rows, weights)
fold_fits.append(fold_fit)
by_fold[k][time] = with_resolved_component(
by_fold[k].get(time, node),
component,
tuned,
resolve_tuned_policy(tuned, fold_fit),
)
in_sample[time] = with_resolved_component(
in_sample.get(time, node), component, tuned, resolve_tuned_policy(tuned, fitted)
)
influence = (
tuner_influence(fitted, doses, frame, reachable, weights)
if target == "population"
else None
)
boundary, edge, gradient, slope = tuned_cap_parts(tuned, fitted, frame(), data.n)
tunings.append(
NodeTuning(
time=time,
policy=tuned,
fitted=fitted,
fold_fits=tuple(fold_fits),
influence=influence,
boundary=boundary,
edge=edge,
gradient=gradient,
slope=slope,
at_risk=reachable,
component=component,
slots=slots,
)
)
def at(resolved: Mapping[int, RuleNode]) -> DynamicRegimen:
plan = tuple(resolved.get(time, node) for time, node in enumerate(regimen.plan, start=1))
return replace(regimen, plan=plan)
fold_plans = tuple(_resolve_plan(at(nodes), data, collapse=True) for nodes in by_fold)
if target == "population":
main = _resolve_plan(at(in_sample), data, collapse=True)
return replace(
main,
tuning=tuple(tunings),
fold_plans=fold_plans,
folds=None if folds.is_single else folds,
)
return _stitched_plan(regimen, fold_plans, folds, tuple(tunings))
def _tuner_source(
regimen: DynamicRegimen, data: LongitudinalData, time: int, node: object, name: str | None
) -> tuple[FloatArray, Callable[[], Any]]:
"""The natural values of a node's tuned dose and the frame its tuner reads.
At a node of one dose the tuner reads the node's dose and policy frame. At a mixed
node it reads the tuned component's natural dose and the frame that component's map
reads: the natural categorical components and the doses before it in the policy's
order (:func:`~cleverly.interventions.policy.component_source`). The tuner is then the
component's conditional quantile given what its map may read, and the cap stays a
triangular map of the node.
"""
where = f"the node {data.decision_name(time)!r}"
if data.node_kind(time) == "mixed" and isinstance(node, JointPolicy):
assert name is not None
components = data.node_components(time)
names = [component.name for component in components]
if name not in names or not components[names.index(name)].continuous:
raise CapabilityError(
f"regimen {regimen.label!r} tunes the cap of {name!r} at {where}, whose "
f"continuous components are "
f"{[c.name for c in components if c.continuous]}. A tuned cap reads a "
"quantile of a dose's law, so it applies to a continuous component",
reason=RefusalReason.INVALID_DECLARATION,
)
atoms = data.component_atoms(time).get(name, ())
if atoms:
raise CapabilityError.from_refusal(
tuned_cap_at_atoms(f"component {name!r} of {where}", atoms)
)
natural = np.nan_to_num(data.node_values(time), nan=0.0)
source = component_source(
node,
name,
components,
natural,
lambda extra: data.policy_frame(time, extra),
where,
)
return natural[:, names.index(name)], source
if data.node_kind(time) != "continuous":
tuned = name if name is not None else str(getattr(node, "name", node))
raise CapabilityError(
f"regimen {regimen.label!r} tunes the cap of {tuned!r} at {where}, which is "
"not one continuous dose "
"or a dose component of a mixed node. A tuned cap reads a quantile of the dose "
"law, which has a density only at a dose",
reason=RefusalReason.INVALID_DECLARATION,
)
atoms = data.node_atoms(time)
if atoms:
raise CapabilityError.from_refusal(tuned_cap_at_atoms(where, atoms))
doses = np.nan_to_num(data.node_codes(time), nan=0.0)
return doses, lazy_frame(partial(data.policy_frame, time))
def _guarded(
regimen: DynamicRegimen, time: int, frame: Any, fit: Callable[[], FittedTuner]
) -> FittedTuner:
"""Fit a tuner, and refuse a tuner that reads a component its map may not read."""
try:
fitted = fit()
check = getattr(frame, "check", None)
if check is not None:
check()
except ForbiddenRead as error:
raise DataError(
f"the tuned cap of regimen {regimen.label!r} at time {time} reads {error}, which "
"the tuned component's map may not read. The tuner reads what the map reads: "
"the natural categorical components and the doses before the component in the "
"policy's order. Declare an order that puts the read dose first, or drop the "
"feature"
) from error
return fitted
def _stitched_plan(
regimen: DynamicRegimen,
fold_plans: tuple[Plan, ...],
folds: Folds,
tunings: tuple[NodeTuning, ...],
) -> Plan:
"""The fold-average plan: at row ``i``, the doses of the fold that holds ``i`` out."""
from dataclasses import replace
first = fold_plans[0]
owner = np.asarray(folds.assignment)
targets: list[tuple[FloatArray, ...] | None] = []
for time in range(1, int(first.values.shape[1]) + 1):
branches = first.carry_targets(time)
if branches is None:
targets.append(None)
continue
stitched = []
for index in range(len(branches)):
column = np.array(branches[index], dtype=float, copy=True)
for k, plan in enumerate(fold_plans):
own = plan.carry_targets(time)
assert own is not None
column[owner == k] = np.asarray(own[index])[owner == k]
stitched.append(column)
targets.append(tuple(stitched))
return replace(
first,
regimen=regimen,
policy_targets=tuple(targets),
tuning=tunings,
fold_plans=fold_plans,
folds=folds,
fold_average=True,
)
def _resolve_mtp_node(
regimen: DynamicRegimen, data: LongitudinalData, time: int
) -> tuple[FloatArray, tuple[FloatArray, ...] | None, tuple[tuple[float, IntArray], ...] | None]:
r"""Evaluate a modified treatment policy at one node, once, before any learner.
Returns the ``(n, J)`` carry weights, and either the ``(n,)`` dose of each branch at a
continuous node, the ``(n, d)`` component values of each branch at a mixed node, or
each branch's ``(n, K)`` level map at a categorical one. The policy reads the policy
frame, :math:`[W, L_1, \ldots, L_t]` and the earlier treatments, which hold the observed
values on which the regression conditions. Rows that left the study before the node
hold zero weight and a zero dose.
"""
node = regimen.plan[time - 1]
name = data.treatment_names[time - 1]
where = f"node {name!r} at time {time} of regimen {regimen.label!r}"
reachable = data.uncensored_through(time - 1) & data.event_free_through(time - 1)
if data.node_kind(time) == "mixed":
assert isinstance(node, JointPolicy) # refuse_continuous_assignments refuses the rest
joint = joint_assignments(
node,
data.node_components(time),
np.nan_to_num(data.node_values(time), nan=0.0),
lambda extra: data.policy_frame(time, extra),
reachable,
where=where,
supports=data.component_supports(time),
regimen=regimen.label,
atoms=data.component_atoms(time),
)
weights = np.array([probability for probability, _ in joint])
carry = np.where(reachable[:, None], np.tile(weights, (data.n, 1)), 0.0)
return carry, tuple(values for _, values in joint), None
if isinstance(node, JointPolicy):
# An all-categorical vector node: the joint policy is a map of the observed tuples,
# read by the discrete formula as a label map of the tuples is.
levels = data.treatment_levels[time - 1]
components = tuple(
JointComponent(column, _component_labels(levels, index))
for index, column in enumerate(data.component_names(time))
)
assigned_codes = joint_discrete_assignments(
node,
components,
levels, # type: ignore[arg-type]
lambda extra: data.policy_frame(time, extra),
reachable,
where=where,
)
return _categorical_carry(data, time, reachable, assigned_codes)
frame = lazy_frame(lambda: data.policy_frame(time))
branches = policy_branches(node)
probabilities = np.array([probability for probability, _ in branches])
if data.is_continuous_node(time):
observed = np.nan_to_num(data.node_codes(time), nan=0.0)
check_continuous_policy(
node,
observed,
frame,
reachable,
where=where,
support=data.node_support(time),
atoms=data.node_atoms(time),
regimen=regimen.label,
node=name,
)
doses = []
for _, branch in branches:
assigned = np.asarray(branch.assign(observed, frame), dtype=float).reshape(-1)
doses.append(np.where(reachable, assigned, 0.0))
carry = np.where(reachable[:, None], np.tile(probabilities, (data.n, 1)), 0.0)
return carry, tuple(doses), None
levels = data.treatment_levels[time - 1]
assigned_codes = discrete_assignments(node, levels, frame, reachable, where=where)
return _categorical_carry(data, time, reachable, assigned_codes)
def _categorical_carry(
data: LongitudinalData,
time: int,
reachable: BoolArray,
assigned_codes: tuple[tuple[float, IntArray], ...],
) -> tuple[FloatArray, None, tuple[tuple[float, IntArray], ...]]:
"""The ``(n, K)`` carry weights of a categorical policy node, from its level maps."""
observed_codes = np.nan_to_num(data.node_codes(time), nan=0.0).astype(np.int64)
carry = np.zeros((data.n, len(data.treatment_levels[time - 1])))
rows = np.arange(data.n)
for probability, codes in assigned_codes:
np.add.at(carry, (rows, codes[rows, observed_codes]), probability)
carry = np.where(reachable[:, None], carry, 0.0)
return carry, None, assigned_codes
def _component_labels(levels: Sequence[Any], index: int) -> tuple[object, ...]:
"""The observed labels of one component of a categorical vector node, sorted."""
values = {level[index] for level in levels}
try:
return tuple(sorted(values))
except TypeError:
return tuple(sorted(values, key=repr))
def refuse_continuous_assignments(
regimen: RegimenSpec, data: LongitudinalData, *, allow_set_dose: bool = False
) -> None:
"""Refuse a plan node the treatment node it sits at cannot take, before any learner.
A static dose on a continuous treatment is not pathwise differentiable (Díaz, Williams,
Hoffman and Schenck 2023, Section 4), and neither is a rule that sets one, so a
continuous node takes a modified treatment policy only. A mixed node, a vector with a
continuous component, takes a :class:`~cleverly.interventions.JointPolicy` only, which
says what happens to each component. A joint policy needs a vector node.
Parameters
----------
regimen : Regimen or DynamicRegimen
The resolved regimen.
data : LongitudinalData
The validated panel.
allow_set_dose : bool, default=False
Accept a fixed dose, a number, or a known rule that sets a dose from the history, at a
continuous node. A parametric model gives either a finite-dimensional target.
Raises
------
CapabilityError
If a continuous node holds anything other than a policy object, a mixed node holds
anything other than a joint policy, or a joint policy sits at a node of one column.
"""
nodes = regimen.values if isinstance(regimen, Regimen) else regimen.plan
for time, node in enumerate(nodes, start=1):
kind = data.node_kind(time)
name = data.decision_name(time)
if (
allow_set_dose
and kind == "continuous"
and not _is_drawn(node)
and (
callable(node)
or (
isinstance(node, (int, float, np.integer, np.floating))
and not isinstance(node, bool)
)
)
):
continue
if isinstance(node, JointPolicy):
if kind == "mixed" or len(data.component_names(time)) > 1:
continue
raise CapabilityError(
f"regimen {regimen.label!r} declares the JointPolicy {node.name!r} at the "
f"node {name!r}, which is one column. A JointPolicy declares one policy per "
"component of a vector node.",
reason=RefusalReason.INVALID_DECLARATION,
alternative="Declare the component's policy itself.",
)
if kind == "mixed" and _is_tilt(node):
raise CapabilityError.from_refusal(tilt_mixed_refusal(regimen.label, name))
if kind == "continuous" and _is_tilt(node):
raise CapabilityError.from_refusal(
tilt_continuous_refusal(regimen.label, data.treatment_names[time - 1])
)
if kind == "mixed":
doses = [
component.name for component in data.node_components(time) if component.continuous
]
what = (
describe_policy(node)
if _is_mtp(node)
else "a known policy density"
if isinstance(node, Stochastic)
else "a rule"
if callable(node)
else f"the label {node!r}"
)
raise CapabilityError(
f"regimen {regimen.label!r} declares {what} at the vector node {name!r}, "
f"which holds the continuous component(s) {doses}. A vector node with a "
"continuous component takes a JointPolicy, one policy per component: a "
"static or dynamic value is not pathwise differentiable (Díaz et al. 2023, "
"Section 4), and a policy of one column does not say which component it "
"moves.",
reason=RefusalReason.DIFFERENT_ESTIMAND,
alternative=(
"Declare JointPolicy(name, components={...}, policy_kind='known'). It "
"answers a different question: the mean outcome when the policy moves "
"each received component."
),
)
if kind == "continuous" and not _is_mtp(node):
raise CapabilityError(
f"regimen {regimen.label!r} sets the continuous node "
f"{data.treatment_names[time - 1]!r} to a fixed dose or rule. A static dose on "
"a continuous treatment is not pathwise differentiable (Díaz et al. 2023, "
"Section 4), so no regular estimator of it exists.",
reason=RefusalReason.DIFFERENT_ESTIMAND,
alternative=(
"Declare a modified treatment policy at that node. It answers a "
"different question: the mean outcome when the policy moves each "
"received dose."
),
)
def attach_policy_numerators(
plans: Sequence[Plan], numerators: Mapping[str, dict[int, FloatArray]]
) -> tuple[Plan, ...]:
"""Each plan with its policy nodes' ratio numerators, once the mechanism is fitted.
Parameters
----------
plans : sequence of Plan
The resolved plans.
numerators : mapping
By plan label, by node counted from one, the ``(n,)`` numerator.
Returns
-------
tuple of Plan
The plans, unchanged where they hold no modified treatment policy.
"""
from dataclasses import replace
out = []
for plan in plans:
if not plan.has_mtp:
out.append(plan)
continue
by_node = numerators[plan.label]
n_times = int(plan.values.shape[1])
out.append(
replace(
plan,
mtp_numerators=tuple(by_node.get(time) for time in range(1, n_times + 1)),
)
)
return tuple(out)
def _is_point_mass(density: FloatArray) -> bool:
"""Whether every row is one-hot, by exact equality and with no tolerance."""
ones = density == 1.0
return bool(np.all(ones | (density == 0.0)) and np.all(ones.sum(axis=1) == 1))
def refuse_regimen_rules(regimens: Any) -> None:
"""Raise unless every plan among ``regimens`` with a callable node declares it known.
A rule can close over any estimate, and no code can inspect a closure, so the status
of each callable node is the ``rule_kind`` of its :class:`DynamicRegimen`. The table
gives the plans this reads from each shape of the ``regimens=`` argument. Any other
value holds no plan, and :func:`resolve_regimens` raises its own error for it.
============================================== ============
``regimens`` plans
============================================== ============
a :class:`Regimen` or :class:`DynamicRegimen` that regimen
a mapping its values
a sequence other than a ``str`` its items
============================================== ============
The declaration of a plan is the ``rule_kind`` of a :class:`DynamicRegimen`, and
``None`` for any other plan. So a callable written inline in a mapping, or held by a
:class:`Regimen`, is undeclared. The check of each plan runs in order:
1. A plan that is an iterator, such as a generator, is a
:class:`~cleverly.exceptions.DataError`. Reading it would consume it, and a fit
reads its ``regimens=`` each time it runs.
2. A declaration outside ``"known"``, ``"estimated"`` and ``None`` is a
:class:`~cleverly.exceptions.DataError`, whatever the plan holds.
3. A plan with a callable node and a declaration of ``None`` or ``"estimated"`` is a
:class:`~cleverly.exceptions.CapabilityError`. A plan of labels alone passes.
This check and :func:`resolve_regimens` read the shape of a plan through one helper,
so they cannot read one plan two ways. :class:`DynamicRegimen` runs this when it is
built and before :meth:`DynamicRegimen.assignment` calls a rule. ``LTMLE.fit`` runs
it on the raw ``regimens=`` before any other check of the data and before any learner,
and :func:`~cleverly.longitudinal.estimator.longitudinal_truncation_curve` runs it on
the resolved regimens of a result. A regimen changed with ``object.__setattr__`` can
carry a declaration this version refuses.
Parameters
----------
regimens : Any
The ``regimens=`` argument of a fit, or the resolved regimens of a result.
Raises
------
DataError
If a plan is an iterator, or a declaration is not one of the three states.
CapabilityError
If a plan with a callable node declares ``None`` or ``"estimated"``.
"""
for label, plan in _plans(regimens):
# Typed ``object`` on purpose: this checks what a modified regimen holds at run
# time, which its annotations do not guarantee.
kind: object = plan.rule_kind if isinstance(plan, DynamicRegimen) else None
nodes = _plan_nodes(label, plan)
entries = (plan,) if nodes is None else nodes
if any(callable(node) and not _is_mtp(node) for node in entries):
# ``refuse`` runs ``check`` first, so every refusal of a rule is the shared one.
_RULE_DECLARATION.refuse(kind)
else:
_RULE_DECLARATION.check(kind)
# Each policy node carries its own declaration, which the point-treatment
# ``Stochastic`` check reads, so the two cannot refuse one policy differently.
refuse_regime_densities(node for node in entries if isinstance(node, Stochastic))
# A modified treatment policy is known by its class, or declares policy_kind.
refuse_policy_declarations(node for node in entries if _is_mtp(node))
def declares_policy(regimens: Any) -> bool:
"""Whether any plan of a ``regimens=`` argument holds a known policy node.
Reads the plans by the table of :func:`refuse_regimen_rules`, and refuses nothing: a
plan this cannot read holds no policy node.
Parameters
----------
regimens : Any
The ``regimens=`` argument of a fit or of a longitudinal estimand.
Returns
-------
bool
``True`` when some node of some plan is a
:class:`~cleverly.interventions.Stochastic` node.
"""
for label, plan in _plans(regimens):
try:
nodes = _plan_nodes(label, plan)
except DataError:
continue
if any(isinstance(node, Stochastic) for node in ((plan,) if nodes is None else nodes)):
return True
return False
def declares_tilt(regimens: Any) -> bool:
"""Whether any plan of a ``regimens=`` argument holds an odds tilt node.
Parameters
----------
regimens : Any
The ``regimens=`` argument of a fit or of a longitudinal estimand.
Returns
-------
bool
``True`` when some node of some plan is an
:class:`~cleverly.interventions.Incremental` node.
"""
for label, plan in _plans(regimens):
try:
nodes = _plan_nodes(label, plan)
except DataError:
continue
if any(_is_tilt(node) for node in ((plan,) if nodes is None else nodes)):
return True
return False
def tilt_nodes(regimens: Any, n_times: int) -> frozenset[int]:
"""The nodes, counted from one, that some plan of ``regimens=`` sets by an odds tilt.
A plan of one :class:`~cleverly.interventions.Incremental` entry tilts every node. A
plan this cannot read contributes no node, and the fit refuses it.
Parameters
----------
regimens : Any
The ``regimens=`` argument of a fit or of a longitudinal estimand.
n_times : int
The number of treatment nodes.
Returns
-------
frozenset of int
The tilt nodes.
"""
times: set[int] = set()
for label, plan in _plans(regimens):
try:
nodes = _plan_nodes(label, plan)
except DataError:
continue
if nodes is None:
if _is_tilt(plan):
times.update(range(1, n_times + 1))
continue
times.update(time for time, node in enumerate(nodes, 1) if _is_tilt(node))
return frozenset(times)
def tilt_mix_refusal(label: str, node: str, kind: str) -> Refusal:
"""The refusal of a tilt plan that also sets a node by a policy :class:`IPSI` cannot fit.
Parameters
----------
label : str
The regimen label.
node : str
The treatment column of the node.
kind : str
What sets the node, such as ``"a modified treatment policy"``.
Returns
-------
Refusal
An :attr:`~cleverly._refusal.RefusalReason.IMPLEMENTATION_GAP` record.
"""
return Refusal(
RefusalReason.IMPLEMENTATION_GAP,
f"regimen {label!r} sets node {node!r} by {kind}. IPSI estimates no density ratio "
"for a modified treatment policy or a risk-ratio tilt, so it has no estimator for a "
"plan with such a node.",
alternative="Fit that plan without an odds tilt by LTMLE or SDR.",
roadmap="F33",
)
def tilt_mixed_refusal(label: str, node: str) -> Refusal:
"""The refusal of an odds tilt declared on a whole mixed vector node.
Parameters
----------
label : str
The regimen label.
node : str
The decision name of the vector node.
Returns
-------
Refusal
An :attr:`~cleverly._refusal.RefusalReason.INVALID_DECLARATION` record.
"""
return Refusal(
RefusalReason.INVALID_DECLARATION,
f"regimen {label!r} sets the vector node {node!r}, which holds a continuous "
"component, by an odds tilt. An odds tilt multiplies the odds of one categorical "
"node, Kennedy (2019), Equation (1), and a node with a continuous component is not "
"one categorical node: its joint density has no odds to multiply.",
alternative=(
"Declare the component's policy inside a JointPolicy at that node, with the "
"randomizer on the JointPolicy and a label map per branch for a randomized "
"categorical component, and fit it by LTMLE."
),
)
def tilt_continuous_refusal(label: str, node: str) -> Refusal:
"""The refusal of an odds tilt at a continuous node.
Parameters
----------
label : str
The regimen label.
node : str
The treatment column of the node.
Returns
-------
Refusal
An :attr:`~cleverly._refusal.RefusalReason.INVALID_DECLARATION` record.
"""
return Refusal(
RefusalReason.INVALID_DECLARATION,
f"regimen {label!r} sets the continuous node {node!r} by an odds tilt. An odds tilt "
"multiplies the odds of one of two arms, and a continuous dose has no such odds.",
alternative="Declare a modified treatment policy at that node and fit it by LTMLE or SDR.",
)
def declares_mtp(regimens: Any) -> bool:
"""Whether any plan of a ``regimens=`` argument holds a modified treatment policy.
Parameters
----------
regimens : Any
The ``regimens=`` argument of a fit or of a longitudinal estimand.
Returns
-------
bool
``True`` when some node of some plan is a policy object, such as
:class:`~cleverly.interventions.Shift`.
"""
for label, plan in _plans(regimens):
try:
nodes = _plan_nodes(label, plan)
except DataError:
continue
if any(_is_mtp(node) for node in ((plan,) if nodes is None else nodes)):
return True
return False
def _plans(regimens: Any) -> tuple[tuple[object, Any], ...]:
"""The labelled plans of ``regimens=``, by the table of :func:`refuse_regimen_rules`.
A mapping labels each plan by its key. Any other plan carries its own label, or
``None`` when it is not a regimen, and :func:`resolve_regimens` refuses that plan.
"""
if isinstance(regimens, (Regimen, DynamicRegimen)):
return ((regimens.label, regimens),)
if isinstance(regimens, Mapping):
return tuple(regimens.items())
if isinstance(regimens, Sequence) and not isinstance(regimens, (str, bytes)):
return tuple((getattr(plan, "label", None), plan) for plan in regimens)
return ()
#: The scalar types that one entry broadcasts across every node, as one treatment label.
_LABEL_TYPES = (bool, int, float, str, np.str_, np.number)
def _plan_nodes(label: object, plan: Any) -> tuple[Any, ...] | None:
"""Read the shape of one plan: its nodes, or ``None`` for a plan of one entry.
:class:`DynamicRegimen`, :func:`refuse_regimen_rules` and :func:`resolve_regimens`
read the shape of a plan here and nowhere else, so the check and the resolver cannot
disagree about the nodes of one plan.
====================================================== ===========================
``plan`` reads as
====================================================== ===========================
a :class:`Regimen` its ``values``
a :class:`DynamicRegimen` its ``plan``
a callable, a treatment label, or no ``__iter__`` ``None``
an iterator, such as a generator a ``DataError``
any other iterable a tuple of its items
====================================================== ===========================
``None`` leaves the one entry to the caller, which broadcasts a callable or a label
across the nodes and refuses anything else. An iterator is refused, not read: reading
it would consume it, and a fit reads its ``regimens=`` again at the next call.
:class:`DynamicRegimen` reads an iterator into a tuple before it calls this, because it
stores its plan.
"""
if isinstance(plan, Regimen):
return tuple(plan.values)
if isinstance(plan, DynamicRegimen):
return tuple(plan.plan)
if isinstance(plan, Mapping):
raise DataError(
f"the plan of regimen {label!r} is a mapping. Pass its treatment nodes as a "
"tuple in time order, or use DynamicRegimen(label, ordered_nodes, rule_kind='known') "
"for a plan with a rule"
)
if isinstance(plan, np.ndarray) and plan.ndim == 0:
raise DataError(
f"the plan of regimen {label!r} is a zero-dimensional array. Pass a treatment "
"label to assign it at every node, a sequence with one entry per node, or "
"DynamicRegimen(label, ordered_nodes, rule_kind='known') for a plan with a rule"
)
if (
callable(plan)
or _is_mtp(plan)
or isinstance(plan, _LABEL_TYPES)
or not hasattr(plan, "__iter__")
):
return None
if isinstance(plan, Iterator):
raise DataError(
f"the plan of regimen {label!r} is an iterator, {type(plan).__name__}. A fit reads "
"regimens= each time it runs, and an iterator is empty after its first read. "
"Pass the plan as a tuple, or as a DynamicRegimen, which stores its plan"
)
return tuple(plan)
def resolve_regimens(
spec: Any, n_times: int, *, held_nodes: int | None = None
) -> tuple[RegimenSpec, ...]:
"""Turn a user's ``regimens=`` argument into an ordered tuple of regimens.
Accepts a mapping from label to plan, where a plan is a single arm meaning "that arm
at every node", or a sequence of ``n_times`` arms. A mapping value may also be a
:class:`Regimen` or a :class:`DynamicRegimen`, and the key is then its label. A
sequence of :class:`Regimen` or :class:`DynamicRegimen` objects passes through. A
plan that is an iterator, such as a generator, is refused: a fit reads its
``regimens=`` each time it runs, and an iterator is empty after its first read.
A plan with a callable node must be a :class:`DynamicRegimen` declared
``rule_kind="known"``. A callable written inline in a mapping, as one rule for every
node or as a node of a sequence, carries no declaration. Its regimen is built with
``rule_kind=None``, and :func:`refuse_regimen_rules` refuses it.
A resolved :class:`DynamicRegimen` keeps the ``rule_kind`` of the one it was given.
A :class:`~cleverly.interventions.Stochastic` node is a known policy at its node, and a
mapping value that is one such node is that policy at every node. Each one carries its own
``density_kind``, which :func:`refuse_regimen_rules` checks.
A plan with no rule and no policy in it comes back a :class:`Regimen`, which is what keeps
a static fit on exactly the code path it was on before rules existed.
Order is preserved, because the first regimen is the one contrasts are taken
against by default and so is part of what the fit reports.
Parameters
----------
spec : Any
The ``regimens=`` argument: a mapping from label to plan, one regimen, or a
sequence of regimens.
n_times : int
The number of treatment nodes, which a single arm is broadcast across. On a held
design it is the number of decisions, one.
held_nodes : int or None, default=None
The node count of a design that holds one baseline decision over every node, and
``None`` for any other design. A plan of one label repeated over ``held_nodes``
nodes is then read as that label, and a repeat of another length raises
:class:`DataError`. Any other plan with more than one node raises
:class:`ValueError`, because the later nodes keep the baseline value.
Returns
-------
tuple of Regimen or DynamicRegimen
The regimens in the order given.
Raises
------
DataError
If ``spec`` is missing, empty or malformed, if a label repeats, if a plan is an
iterator or has the wrong number of nodes, or if a declaration is not one of the
three states.
CapabilityError
If a plan with a callable node is not declared ``rule_kind="known"``.
"""
if spec is None:
raise DataError(
"a longitudinal fit needs regimens= : the parameter is the mean outcome under "
"a treatment plan, and there is no default plan to fall back on. Pass, for "
"example, regimens={'always': 1, 'never': 0}"
)
if isinstance(spec, (Regimen, DynamicRegimen)):
spec = (spec,)
if isinstance(spec, Mapping):
items: list[tuple[str, Any]] = list(spec.items())
elif isinstance(spec, Sequence) and not isinstance(spec, (str, bytes)):
items = []
for entry in spec:
if not isinstance(entry, (Regimen, DynamicRegimen)):
raise DataError(
"a sequence of regimens must hold Regimen or DynamicRegimen objects; "
"pass a mapping {label: plan} to name plans inline"
)
items.append((entry.label, entry))
else:
raise DataError(f"regimens= must be a mapping or a sequence of Regimen; got {spec!r}")
if not items:
raise DataError("regimens= is empty; a fit with no regimen reports no parameter")
resolved: list[RegimenSpec] = []
seen: set[str] = set()
for label, plan in items:
name = str(label)
if name in seen:
raise DataError(f"regimen label {name!r} appears twice; labels name parameters")
seen.add(name)
resolved.append(
_resolve_one(
name,
plan if held_nodes is None else _held_plan(name, plan, held_nodes),
n_times,
)
)
return tuple(resolved)
def _held_plan(label: str, plan: Any, n_times: int) -> Any:
"""One plan of a held design, read as its single decision.
A plan of one entry passes. A static plan that repeats one label over all ``n_times``
nodes is that label, since the later nodes keep the baseline value whatever the plan
says there. A repeated label of another length gets the length refusal of an unheld
design. Any other plan of more than one node is refused before any learner.
"""
nodes = _plan_nodes(label, plan)
if nodes is None or len(nodes) == 1:
return plan
static = not isinstance(plan, DynamicRegimen) and not any(
callable(node) or _is_drawn(node) for node in nodes
)
if static and all(_same_label(node, nodes[0]) for node in nodes):
if len(nodes) != n_times:
raise DataError(
f"regimen {label!r} assigns {len(nodes)} arm(s) but the data has {n_times} "
"treatment node(s); a plan must say what happens at every one of them"
)
return Regimen(label, (nodes[0],)) if isinstance(plan, Regimen) else nodes[0]
raise CapabilityError(
f"regimen {label!r} declares {len(nodes)} nodes, but this design has one treatment "
"decision, at baseline, held over every node.",
reason=RefusalReason.DIFFERENT_ESTIMAND,
alternative=(
"Declare one arm, or a one-node rule, policy or modified treatment policy. The "
"later nodes keep the baseline value."
),
)
def _resolve_one(label: str, plan: Any, n_times: int) -> RegimenSpec:
"""Read one plan into the regimen kind it describes.
A rebuilt :class:`DynamicRegimen` carries the ``rule_kind`` of the one it was given,
and ``None`` for any other plan, so its construction checks the declaration.
"""
kind: FunctionKind | None = plan.rule_kind if isinstance(plan, DynamicRegimen) else None
nodes = _nodes(label, plan, n_times)
if any(callable(node) or _is_drawn(node) for node in nodes):
return DynamicRegimen(label, nodes, rule_kind=kind)
return Regimen(label, nodes)
def _nodes(label: str, plan: Any, n_times: int) -> tuple[RuleNode, ...]:
"""Read one plan into one entry per node, broadcasting a single entry across them."""
nodes = _plan_nodes(label, plan)
if nodes is None:
# A numpy array and a pandas Series are plans by every reading except
# ``isinstance(..., Sequence)``, which neither registers for. ``_plan_nodes`` tests
# for the iteration protocol instead, so the message about rules is not aimed at
# an array whose diagnosis it gets wrong.
if not (callable(plan) or _is_drawn(plan) or isinstance(plan, _LABEL_TYPES)):
raise DataError(
f"regimen {label!r} must be a treatment label, a rule d_t(H_t), a known "
f"policy Stochastic(q, name, density_kind='known'), a modified treatment "
f"policy such as Shift(0.5, cap=4.0), an odds tilt Incremental(2.0), or a "
f"sequence of {n_times} of these; got {plan!r}"
)
# A rule or a policy gets the broadcast that a scalar arm gets. So a rule that reads a
# late-measured covariate is diagnosed at evaluation rather than here: whether
# ``lambda h: h["L2"] > 0`` is usable at node 1 is a question about the data.
return (plan,) * n_times
if len(nodes) != n_times:
raise DataError(
f"regimen {label!r} assigns {len(nodes)} arm(s) but the data has {n_times} "
"treatment node(s); a plan must say what happens at every one of them"
)
return nodes