lanfactory.derive

Integrate a trained LAN's density over reaction time so that choice probabilities and deadline (omission) mass can be derived from the LAN instead of simulated.

lanfactory.derive.integrate

Integrate a trained LAN's density over reaction time.

A LAN approximates log p(rt, choice | theta). Integrating the density over rt gives each choice's probability and the probability of responding before a deadline: the labels of the auxiliary networks (CPNs, OPNs), here obtained from the LAN instead of from a fresh simulation.

Grid policy

The density is integrated with the trapezoid rule over the LANs' training support [t_min, max_t] (max_t = 20 s, ssms' default) and is never evaluated past it. It rises from zero within milliseconds of its onset, so the one grid, :class:OnsetGrid, is refined per parameter vector around the onset: the non-decision time t, or t - scale * theta[param] for a model whose ssms config declares an ndt_edge_shift. On the Hub ddm LAN its 1160 points agree with a 16 000-point uniform grid to 1e-3 (p99; 3.5e-3 at worst), where a uniform 1000-point grid is off by up to 25 % at a < 0.5. The grid starts at t_min, not at the onset: the mass a LAN leaks below its onset is part of what it predicts. A model without an onset parameter passes onset = 0.

Tail policy

Masses are not renormalised: they keep the LAN's own scale error and whatever density it puts past max_t. :attr:ChoiceMass.total exposes the deficit; renormalising is the corpus's job.

lanfactory.derive.load_onnx_predictor

load_onnx_predictor(path: str | Path) -> OnnxPredictor

Load a (1, D) ONNX network as a batched :class:OnnxPredictor.

The ecosystem's ONNX graphs are exported with a concrete (1, D) input (see lanfactory.onnx.contract), which onnxruntime enforces. The graph's batch dimension (dim 0 of every graph input and output) is set to the symbolic "N" on the in-memory protobuf, and the session is built from that copy. The file on disk is read, never written, so the single-trial contract the artifact was published under is untouched.

Parameters:

• path (str | Path) –

  Path to the ONNX artifact.

Returns:

• OnnxPredictor –

  Callable on (N, D) float32 batches.

Raises:

• ValueError –

  If the graph's first input is not rank 2. The sbi and bayesflow exporters trace rank-1 (D,) graphs; those have no batch axis to widen, and this loader does not wrap them.

lanfactory.derive.OnnxPredictor

OnnxPredictor(session: InferenceSession)

A (1, D) ONNX network callable on (N, D) batches.

Build with :func:load_onnx_predictor.

lanfactory.derive.Predictor

Bases: Protocol

Maps (N, input_width) float32 rows to (N,) log-densities.

lanfactory.derive.OnsetGrid dataclass

OnsetGrid(n_pre: int = 32, n_knee: int = 128, knee: float = 0.05, n_onset: int = 600, onset_window: float = 1.0, n_tail: int = 400, max_t: float = 20.0, t_min: float = 0.0001)

Reaction-time grid refined, per parameter vector, around the onset.

Each row concatenates four linspace segments around its own onset t: [t_min, knee), [knee, t), [t, t + onset_window) and [t + onset_window, max_t]. The points concentrate where the density moves while the whole support is still integrated. Every row has :attr:n_points points, contains t_min, t and max_t exactly and is strictly increasing.

Parameters:

• n_pre (int, default: 32 ) –

  Points per segment, in that order; only the tail includes its endpoint.

• n_knee (int, default: 32 ) –

  Points per segment, in that order; only the tail includes its endpoint.

• n_onset (int, default: 32 ) –

  Points per segment, in that order; only the tail includes its endpoint.

• n_tail (int, default: 32 ) –

  Points per segment, in that order; only the tail includes its endpoint.

• knee (float, default: 0.05 ) –

  Width in seconds of the segment just below t, where the LAN blurs the onset; half of [t_min, t] instead when that is shorter.

• onset_window (float, default: 1.0 ) –

  Width in seconds of the segment just above t.

• max_t (float, default: 20.0 ) –

  Upper edge in seconds: the LANs' training support.

• t_min (float, default: 0.0001 ) –

  Lower edge in seconds, positive so no network is fed rt = 0.

Raises:

• ValueError –

  If a segment has no points, if knee, onset_window or t_min is not positive, or if max_t is not finite or leaves no room for onset_window above t_min.

Notes

Onsets are clipped to [t_min + 1e-3, max_t - onset_window - 1e-3] so that every segment has positive length.

Methods:

• for_theta –

  Build one grid row per onset time.

Attributes:

• n_points (int) –

  Points per row: n_pre + n_knee + n_onset + n_tail.

lanfactory.derive.OnsetGrid.n_points property

n_points: int

Points per row: n_pre + n_knee + n_onset + n_tail.

lanfactory.derive.OnsetGrid.for_theta

for_theta(onset: ArrayLike) -> NDArray[np.float64]

Build one grid row per onset time.

Parameters:

• onset (ArrayLike) –

  (n_theta,) onset times in seconds.

Returns:

• NDArray[float64] –

  (n_theta, n_points) strictly increasing rows.

Raises:

• ValueError –

  If onset is not 1-D or not finite.

lanfactory.derive.choice_mass

choice_mass(predictor: Predictor, theta: ArrayLike, choices: Sequence[float] | NDArray, onset: ArrayLike, grid: OnsetGrid = OnsetGrid()) -> ChoiceMass

Integrate a LAN's density over reaction time for every choice.

Rows are laid out as [theta..., rt, choice], the input layout the LANs are trained on, and evaluated 512 parameter vectors at a time, which bounds the row buffer handed to the network.

Parameters:

• predictor (Predictor) –

  Maps (N, n_params + 2) rows to (N,) log-densities; typically an :class:OnnxPredictor.

• theta (ArrayLike) –

  (n_theta, n_params) parameter vectors (a single vector is promoted).

• choices (Sequence[float] | NDArray) –

  The choice codes the LAN was trained on, e.g. [-1, 1].

• onset (ArrayLike) –

  Scalar or (n_theta,) onset (non-decision) times the grid is built around, e.g. theta[:, t_idx], less the config's ndt_edge_shift where it declares one; 0 for a model without an onset.

• grid (OnsetGrid, default: OnsetGrid() ) –

  Reaction-time grid (see the module's grid policy).

Returns:

• ChoiceMass –

Raises:

• ValueError –

  If predictor.input_width != n_params + 2, the usual symptom of pairing a parameter set with the wrong LAN, or if onset is not finite or does not broadcast to (n_theta,).

lanfactory.derive.ChoiceMass dataclass

ChoiceMass(t: NDArray[float64], cdf: NDArray[float64], by_choice: NDArray[float64])

A LAN's density integrated over reaction time, per parameter vector.

Notes

Masses are not renormalised; see the module's tail policy.

Methods:

• mass_before –

  Mass accumulated before deadline, linearly interpolated on the grid.

• quantile –

  Time at which the mass reaches u of :attr:total.

lanfactory.derive.ChoiceMass.total property

total: NDArray[float64]

(n_theta,) mass on [t_min, max_t], summed over choices.

lanfactory.derive.ChoiceMass.mass_before

mass_before(deadline: ArrayLike) -> NDArray[np.float64]

Mass accumulated before deadline, linearly interpolated on the grid.

Parameters:

• deadline (ArrayLike) –

  Scalar or (n_theta,) deadlines in seconds. Values below t_min give 0 and values above max_t give :attr:total.

Returns:

• NDArray[float64] –

  (n_theta,).

lanfactory.derive.ChoiceMass.quantile

quantile(u: ArrayLike) -> NDArray[np.float64]

Time at which the mass reaches u of :attr:total.

The inverse of :meth:mass_before, the same interpolation with the axes swapped: mass_before(quantile(u)) is u * total.

Parameters:

• u (ArrayLike) –

  Fraction in (0, 1); scalar or (n_theta,).

Returns:

• NDArray[float64] –

  (n_theta,) times in seconds; NaN where there is no mass to invert (a zero or non-finite total).

lanfactory.derive.survey

survey(predictor: Predictor, param_bounds: dict[str, tuple[float, float]], params: Sequence[str], choices: Sequence[float] | NDArray, *, n_theta: int = 20000, seed: int = 0, onset_param: str | None = 't', ndt_edge_shift: Mapping[str, Any] | None = None) -> dict

Whole-box statistics of a LAN's integrated mass.

Draws n_theta parameter vectors uniformly inside param_bounds, integrates each with :func:choice_mass (streaming: only the totals are kept) and summarises how far the total mass strays from one and where in the box it does so. The output is plain numbers (json.dumps-able).

Parameters:

• predictor (Predictor) –

  The LAN, e.g. from :func:load_onnx_predictor.

• param_bounds (dict[str, tuple[float, float]]) –

  {name: (lo, hi)} for every name in params; the box is sampled uniformly (the training box of the LAN, typically).

• params (Sequence[str]) –

  Parameter names in the order the LAN's rows expect them.

• choices (Sequence[float] | NDArray) –

  The choice codes the LAN was trained on.

• n_theta (int, default: 20000 ) –

  Number of parameter vectors drawn.

• seed (int, default: 0 ) –

  Seed of the numpy.random.default_rng draw.

• onset_param (str | None, default: 't' ) –

  Name of the non-decision-time parameter: the onset the grid is built around and the edge leak_below_onset is measured at. None builds the grid around onset 0 and reports no leak.

• ndt_edge_shift (Mapping[str, Any] | None, default: None ) –

  The model config's ndt_edge_shift, {"param": name, "scale": s}, for a model whose reaction times start below t: the onset becomes t - s * theta[name]. None for every other model.

Returns:

• dict –

  n_theta, grid (its repr), seconds (wall time), total (mean/min/max of the total and quantiles and exceedance fractions of |total - 1|; a NaN total counts as an infinite deviation), leak_below_onset (mean/p99/max of :meth:ChoiceMass.mass_before at the onset, or None) and by_param ({name: [10 bins of lo/hi/mean_dev/max_abs_dev/n]}; empty bins carry None).

Raises:

• ValueError –

  If params is empty or names a parameter without bounds, if onset_param is not in params, if ndt_edge_shift comes without onset_param or names a parameter not in params, if a bound has hi < lo, or if n_theta < 1.

Derived corpora

lanfactory.derive.corpus turns those masses into training corpora the trainers read unchanged ({type}_data, {type}_labels, generator_config, model_config per pickle), plus a derive_manifest.json sidecar. The derive-aux command wraps derive_aux_corpus; the network types page explains what is derived, the row layouts, the renormalisation and the fallback, the recorded statistics and provenance, and how to train on the result.

A run trained on a derived corpus logs its provenance to MLflow: see Derived corpora.

lanfactory.derive.derive_aux_corpus

derive_aux_corpus(onnx_path: str | Path, model: str, network_type: str, out_folder: str | Path, *, n_files: int = 100, n_theta_per_file: int = 4096, fallback_n_sim: int = 20000, seed: int = 0, source_run_uuid: str | None = None, source_hf_revision: str | None = None) -> list[Path]

Write a cpn / opn training corpus derived from a LAN.

Each file holds n_theta_per_file parameter vectors sampled as ssms would (:func:sample_theta), their masses under the LAN on :data:GRID, and the rows and labels of :func:cpn_labels or :func:opn_labels. The pickles carry {type}_data, {type}_labels, generator_config and model_config — the layout DatasetTorch and the training CLIs read — so the result trains with the trainers unchanged.

Labels are the LAN's masses renormalised by its own total. A parameter vector whose total lies outside :data:FALLBACK_WINDOW is one the LAN gets wrong beyond a scale error, so its labels come from :func:simulate_labels on ssms instead. opn deadlines are drawn by :func:sample_deadlines with :data:DEADLINE_QUANTILE_FRAC on ssms' deadline box.

Every pickle's generator_config is one corpus record (model, type, sizes, settings under derive, provenance under source) plus its file_index. The :data:MANIFEST_NAME sidecar is that record plus the corpus-wide derive_stats (total-mass statistics, the share that fell back, and the simulation's blind spot past max_t) and the file names.

File i is generated from np.random.default_rng([seed, i]), so a file's content depends on seed and its index only — the simulator seeds of the fallback are drawn from the same generator.

Parameters:

• onnx_path (str | Path) –

  The LAN, a (1, n_params + 2) ONNX graph.

• model (str) –

  ssms model name the LAN was trained for, e.g. "ddm".

• network_type (str) –

  One of :data:NETWORK_TYPES.

• out_folder (str | Path) –

  Destination directory; created when the first file is written.

• n_files (int, default: 100 ) –

  Number of pickles; at least 2, since the trainers split files into train and validation sets.

• n_theta_per_file (int, default: 4096 ) –

  Parameter vectors per file. Rows per file are this for opn and this times the number of choices for cpn; the trainers require the batch size to divide the row count.

• fallback_n_sim (int, default: 20000 ) –

  Trials simulated per fallback theta.

• seed (int, default: 0 ) –

  Base seed.

• source_run_uuid (str | None, default: None ) –

  Training run uuid of the LAN; parsed from the trainer's file name conventions when omitted.

• source_hf_revision (str | None, default: None ) –

  Hub revision (commit, tag, or branch) the LAN was downloaded at, if any; recorded as source_lan_hf_commit.

Returns:

• list[Path] –

  The pickles written, in file order.

Raises:

• ValueError –

  For an unknown network_type, n_files < 2, n_theta_per_file < 1, fallback_n_sim < 1, a model without a non-decision time t, or a LAN whose input width is not n_params + 2 for model.

lanfactory.derive.NETWORK_TYPES module-attribute

NETWORK_TYPES: tuple[str, ...] = ('cpn', 'opn')

Auxiliary network types a corpus can be derived for.

lanfactory.derive.AUX_CATEGORY module-attribute

AUX_CATEGORY: dict[str, str] = {'cpn': 'choice', 'opn': 'omission'}

What each network type's label is the probability of.

lanfactory.derive.DERIVATION_METHOD module-attribute

DERIVATION_METHOD = 'derived-from-lan'

Value of the derivation_method provenance key.

lanfactory.derive.MANIFEST_NAME module-attribute

MANIFEST_NAME = 'derive_manifest.json'

Sidecar written next to the pickles by :func:derive_aux_corpus.