torchref.kinetic.targets module

Collection-aware targets for kinetic refinement.

All targets extend torchref.refinement.targets.base.Target and operate on paired DatasetCollection + ModelCollection instances. Keys are matched automatically so that each timepoint dataset is paired with its corresponding mixed model.

Targets

CollectionDifferenceTarget

Multi-timepoint difference target (primary optimization driver).

CollectionMLTarget

Multi-timepoint maximum-likelihood amplitude target.

MultiModelGeometryTarget

Geometry restraints applied to the shared base models.

MultiModelADPTarget

ADP restraints applied to the shared base models.

KineticPriorTarget

Regularizes per-timepoint fractions towards a kinetic model.

class torchref.kinetic.targets.CollectionDifferenceTarget(dataset_collection, model_collection, scaler=None, normalize=True, use_work_set=True, verbose=0)

Bases: Target

Mean-based difference target using DatasetCollection + ModelCollection.

Computes differences relative to the mean across all N datasets (dark + timepoints), with proper error propagation accounting for the covariance between each dataset and the mean:

F_mean(h) = (1/N) Σ_i F_obs_i(h)
ΔF_obs_i  = F_obs_i - F_mean
ΔF_calc_i = |F_calc_i| - F_calc_mean

Var(F_i - F_mean) = σ_i²·(1 - 2/N) + (Σ_j σ_j²)/N²

For N=2 (dark + one timepoint) this gives identical gradients to the direct dark-reference subtraction. For N>2 the mean reference has lower noise.

All computation is vectorized on stacked (N, n_hkl) tensors — no Python loops over datasets.

Parameters:

• dataset_collection (DatasetCollection)

• model_collection (ModelCollection)

• scaler (ScalerBase) – Single scaler applied to all F_calc (uses forward_mixed with per-model fractions when available).

• normalize (bool) – If True, divide total NLL by number of datasets.

• use_work_set (bool) – If True, compute loss only on the work set (rfree_flags=True).

• verbose (int) – Verbosity level.

name: str = 'difference_xray'

__init__(dataset_collection, model_collection, scaler=None, normalize=True, use_work_set=True, verbose=0)

Initialize target.

Parameters:

verbose (int, optional) – Verbosity level. Default is 0.

forward()

Compute and return the loss. Override in subclasses.

class torchref.kinetic.targets.CollectionMLTarget(dataset_collection, model_collection, scaler=None, normalize=True, use_work_set=True, verbose=0)

Bases: Target

Multi-timepoint maximum-likelihood amplitude target.

Computes Rice-distribution NLL (acentric) and the corresponding centric NLL for each timepoint, with proper validity masking and NaN/Inf protection. Vectorized on stacked (N_tp, n_hkl) tensors.

Parameters:

• dataset_collection (DatasetCollection)

• model_collection (ModelCollection)

• scaler (ScalerBase, optional) – Single scaler applied to each timepoint’s F_calc.

• normalize (bool) – Divide total NLL by number of matched timepoints.

• use_work_set (bool) – Compute loss only on work set.

• verbose (int) – Verbosity level.

name: str = 'collection_ml_xray'

__init__(dataset_collection, model_collection, scaler=None, normalize=True, use_work_set=True, verbose=0)

Initialize target.

Parameters:

verbose (int, optional) – Verbosity level. Default is 0.

forward()

Compute and return the loss. Override in subclasses.

class torchref.kinetic.targets.MultiModelGeometryTarget(model_collection, verbose=0)

Bases: Target

Geometry restraints for the shared base models in a ModelCollection.

Creates a TotalGeometryTarget for each base model and sums them. Since models are shared across timepoints, restraints only need to be computed once per base model (not per timepoint).

Parameters:

• model_collection (ModelCollection)

• verbose (int) – Verbosity level.

name: str = 'multi_model_geometry'

__init__(model_collection, verbose=0)

Initialize target.

Parameters:

verbose (int, optional) – Verbosity level. Default is 0.

forward()

Compute and return the loss. Override in subclasses.

register_to_state(state)

Register each base model’s geometry sub-targets individually into a LossState with hierarchical naming.

Parameters:

state (LossState) – The loss state to register targets into.

items()

Expose sub-targets for LossState auto-expansion.

class torchref.kinetic.targets.MultiModelADPTarget(model_collection, verbose=0)

Bases: Target

ADP restraints for the shared base models in a ModelCollection.

Same pattern as MultiModelGeometryTarget but using TotalADPTarget.

Parameters:

• model_collection (ModelCollection)

• verbose (int) – Verbosity level.

name: str = 'multi_model_adp'

__init__(model_collection, verbose=0)

Initialize target.

Parameters:

verbose (int, optional) – Verbosity level. Default is 0.

forward()

Compute and return the loss. Override in subclasses.

register_to_state(state)

Register per-model ADP sub-targets into LossState.

items()

class torchref.kinetic.targets.KineticPriorTarget(model_collection, kinetic_model, timepoints_map, verbose=0)

Bases: Target

Regularize per-timepoint fractions towards a kinetic model.

The kinetic model provides a smooth prior over how population fractions should evolve over time. The fractions in ModelCollection are free parameters; this target penalizes deviation from the kinetic prediction.

Periodically call refit_prior() to update the kinetic model to match the current free fractions (EM-style alternation).

Parameters:

• model_collection (ModelCollection)

• kinetic_model (occupancies_kinetics) – The kinetic occupancy model whose forward() returns shape [n_states, n_timepoints].

• timepoints_map (Dict[str, int]) – Maps timepoint names to indices into the kinetic model’s time axis. E.g. {"1ps": 0, "5ps": 1, "10ps": 2}.

• verbose (int) – Verbosity level.

name: str = 'kinetic_prior'

__init__(model_collection, kinetic_model, timepoints_map, verbose=0)

Initialize target.

Parameters:

verbose (int, optional) – Verbosity level. Default is 0.

forward()

Squared difference between current fractions and kinetic predictions.

refit_prior(niter=50, lr=0.01)

Refit kinetic model to match current free fractions (M-step).

Freezes model fractions, optimizes kinetic model parameters to minimize prediction error against current fractions.

Parameters:

• niter (int) – Number of optimizer steps.

• lr (float) – Learning rate for Adam optimizer.