torchref.kinetic package

Kinetic refinement for time-resolved crystallography.

This submodule provides tools for refining multiple structural models against multiple time-resolved datasets with kinetics-constrained occupancy fractions.

Key Classes

ModelCollection

Named dictionary of MixedModel instances at different timepoints, sharing the same base structural models.

KineticRefinement

Orchestrator combining DatasetCollection + ModelCollection with difference and ML targets, geometry/ADP restraints, and optional kinetic prior regularization.

KineticModel

PyTorch ODE solver for kinetic schemes (matrix exponential).

occupancies_kinetics

Kinetics-constrained occupancy model wrapping KineticModel.

class torchref.kinetic.KineticModel(flow_chart, timepoints, rate_constants=None, efficiencies=None, instrument_function='gaussian', instrument_width=10, initial_state=None, light_activated=False, activation_level=0.5, verbose=1)

Bases: DeviceMixin, Module

Configurable PyTorch module for fitting kinetic behavior.

Supports arbitrary kinetic schemes defined by relational strings: - “A->B,B->C” (sequential) - “A->B,B->A,B->C” (with back reactions) - “A->B,A->C,B->D,C->D” (parallel pathways) - “A->B,B->C,D” (D is non-reactive state)

Each transition has TWO parameters: - Reactivity constant k (rate) - Reaction efficiency η (0 to 1, controls maximum conversion)

States can have baseline occupancy offsets (default: 0, not refined).

The initial transfer is driven by photoabsorption with quasi-instant conversion around zero, with spread accounted for by an instrument function.

Parameters:

• flow_chart (str) – Relational string describing the kinetic scheme using comma-separated transitions. Standalone states (non-reactive) can be included without transitions. Example: “A->B,B->C” or “A->B,B->A,B->C,D” (D is non-reactive)

• timepoints (torch.Tensor or array-like) – Time points at which to evaluate the kinetics

• rate_constants (dict or list, optional) – Initial rate constants. Can be: - Dict mapping “A->B” to float value - List of floats (same order as transitions in flow_chart) - None (random initialization)

• efficiencies (dict or list, optional) – Initial reaction efficiencies (0-1). Same format as rate_constants. Default: all 1.0 (100% efficient)

• instrument_function (str, optional) – Type of instrument response function. Options: ‘gaussian’, ‘none’ Default: ‘gaussian’

• instrument_width (float, optional) – Width parameter for the instrument function (e.g., sigma for gaussian) Default: 0.1

• initial_state (str, optional) – Which state starts with population 1. Default: first state in flow chart

• light_activated (bool, optional) – If True, treats this as a light-activated reaction where the initial photoexcitation can only happen once. Products returning to the initial state become inactive (A*) and cannot undergo photoactivation again. Default: False

• verbose (int, optional) – Verbosity level. Default: 1

__init__(flow_chart, timepoints, rate_constants=None, efficiencies=None, instrument_function='gaussian', instrument_width=10, initial_state=None, light_activated=False, activation_level=0.5, verbose=1)

forward()

Forward pass: compute populations at all timepoints.

Returns:

populations – Population of each state at each timepoint Shape: (n_timepoints, n_states)

Return type:

torch.Tensor

set_baseline(state, occupancy, refinable=False)

Set baseline occupancy offset for a state.

Baseline occupancies are constant offsets added to the population of a state. This is useful for non-reactive background states.

Parameters:

• state (str) – Name of the state

• occupancy (float) – Baseline occupancy value (offset)

• refinable (bool, optional) – If True, this baseline becomes a refinable parameter. If False (default), it remains constant.

Examples

>>> model.set_baseline('D', 0.1, refinable=False)  # Constant 10% background
>>> model.set_baseline('E', 0.05, refinable=True)  # Refinable baseline

get_baselines()

Get current baseline occupancies for all states.

Returns:

baselines – Dictionary mapping state names to baseline occupancies

Return type:

Dict[str, float]

get_rate_constants()

Get current rate constants (k) as a dictionary.

Returns:

rate_dict – Dictionary mapping transition strings to rate constants

Return type:

Dict[str, float]

set_rate_constant(transition, value)

Set rate constant for a specific transition.

Parameters:

• transition (str) – Transition string in the format “A->B”

• value (float) – New rate constant value (must be positive)

get_efficiencies()

Get current reaction efficiencies (η) as a dictionary.

Returns:

eff_dict – Dictionary mapping transition strings to efficiencies (0-1)

Return type:

Dict[str, float]

get_effective_rates()

Get effective rates (k * η) as a dictionary.

Returns:

eff_rate_dict – Dictionary mapping transition strings to effective rates

Return type:

Dict[str, float]

get_time_constants()

Get time constants (1/k_eff) for each transition.

Returns:

time_dict – Dictionary mapping transition strings to time constants

Return type:

Dict[str, float]

parameters()

Get all flexible (learnable) parameters as a dictionary.

Returns:

params – Dictionary mapping parameter names to their tensors: - ‘log_rate_constants’: log-transformed rate constants - ‘log_instrument_width’: log-transformed instrument width (if refinable) - ‘baseline_{state}’: refinable baseline for specific states (if any)

Return type:

Dict[str, torch.Tensor]

Examples

>>> model = KineticModel(...)
>>> params = model.parameters()
>>> print(params.keys())
>>> # Use with optimizer: optimizer = torch.optim.Adam(params.values(), lr=0.01)

print_parameters()

Print current model parameters.

plot_occupancies(outpath, times=None, log=False, figsize=(10, 6), dpi=150, title=None)

Plot state occupancies over time and save to file.

Parameters:

• outpath (str) – Path to save the plot (e.g., ‘kinetics.png’)

• log (bool, optional) – If True, use log scale for x-axis. Default: False

• figsize (Tuple[int, int], optional) – Figure size (width, height). Default: (10, 6)

• dpi (int, optional) – DPI for saving figure. Default: 150

• title (str, optional) – Custom title for the plot. If None, uses flow chart string

visualize(outpath, **kwargs)

Alias for plot_occupancies for convenience.

Parameters:

• outpath (str) – Path to save the plot

• **kwargs – Additional arguments passed to plot_occupancies

class torchref.kinetic.occupancy_unrestrained(nstates, time)

Bases: DeviceMixin, Module

Unrestrained occupancy model where each state at each timepoint is independent.

This is the most flexible model but may lead to physically unrealistic solutions since occupancies at different timepoints are not coupled.

Parameters:

• nstates (int) – Number of structural states

• time (list or array-like) – Time points for the experiment

__init__(nstates, time)

forward()

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

plot_occupancies(path, log_scale=False, figsize=(10, 6))

class torchref.kinetic.occupancies_kinetics(flow_chart, time, rate_constants=None, efficiencies=None, instrument_function='none', instrument_width=10.0, light_activated=False, state_mapping=None, regularization=None, verbose=1)

Bases: DeviceMixin, Module

Kinetics-constrained occupancy model.

This model uses a kinetic scheme to constrain occupancies at different timepoints. Instead of independent parameters for each timepoint, the occupancies are derived from rate constants, efficiencies, and a kinetic flow chart.

This provides several advantages over unrestrained refinement: 1. Physical constraints: occupancies follow kinetic laws 2. Reduced parameters: n_rates instead of n_states * n_timepoints 3. Extrapolation: can predict occupancies at unmeasured timepoints 4. Interpretability: rate constants have physical meaning

Refinement Considerations

Kinetic refinement has unique challenges compared to normal refinement:

1. Parameter Scale Separation: Rate constants can span many orders of magnitude (e.g., ps to ms). Use log-parameterization and consider different learning rates.

2. Identifiability: Some parameters may be correlated: - Rate and efficiency can compensate for each other - Back-reactions can create degenerate solutions Consider using efficiency constraints or regularization.

3. Local Minima: The optimization landscape can have multiple minima corresponding to different kinetic interpretations. Initialize carefully.

4. Gradient Flow: Matrix exponentials can have vanishing/exploding gradients. The implementation clips extreme values for stability.

5. Regularization: Consider adding priors on: - Rate constants (log-normal centered on expected timescales) - Efficiencies (Beta distribution favoring high efficiency) - Smoothness of rate changes if doing temperature-dependent fitting

param flow_chart:

Kinetic scheme, e.g., “A->B,B->C,C->D” or “A->B,B->A,B->C” (with back-reaction)

type flow_chart:

str

param time:

Time points at which to evaluate occupancies

type time:

list or tensor

param rate_constants:

Initial rate constants as {“A->B”: value, …}. If None, uses smart initialization.

type rate_constants:

dict or None, optional

param efficiencies:

Initial efficiencies as {“A->B”: value, …}. Default: all 1.0

type efficiencies:

dict or None, optional

param instrument_width:

Instrument response function width (Gaussian sigma). Default: 10

type instrument_width:

float, optional

param light_activated:

If True, products returning to ground state become inactive. Default: False

type light_activated:

bool, optional

param state_mapping:

Mapping from kinetic states to structural model indices. E.g., {“A”: 0, “B”: 1, “C”: 2, “D”: 3} or {“A”: 0, “B”: 1, “C”: 1, “D”: 2} The latter allows multiple kinetic states to map to the same structure. If None, assumes sequential mapping (A=0, B=1, …).

type state_mapping:

dict or None, optional

param regularization:

Regularization settings: - ‘rate_prior_weight’: weight for log-normal prior on rates - ‘rate_prior_mean’: mean of log-rate prior (default: based on time range) - ‘rate_prior_std’: std of log-rate prior (default: 2.0, allows ~2 orders of magnitude) - ‘efficiency_prior_weight’: weight for efficiency prior (favoring 1.0)

type regularization:

dict or None, optional

param verbose:

Verbosity level. Default: 1

type verbose:

int, optional

Examples

>>> # Simple sequential kinetics: A -> B -> C -> D
>>> occ = occupancies_kinetics(
...     flow_chart="A->B,B->C,C->D",
...     time=torch.linspace(0, 100, 50),
...     rate_constants={"A->B": 1.0, "B->C": 0.1, "C->D": 0.01}
... )
>>> occupancies = occ()  # Shape: [n_states, n_timepoints]

>>> # With back-reaction
>>> occ = occupancies_kinetics(
...     flow_chart="A->B,B->A,B->C",
...     time=times,
...     light_activated=True  # Products returning to A become inactive
... )

>>> # Mapping multiple kinetic states to same structure
>>> occ = occupancies_kinetics(
...     flow_chart="A->B,B->C,C->D",
...     time=times,
...     state_mapping={"A": 0, "B": 1, "C": 1, "D": 0}  # B and C share structure
... )

__init__(flow_chart, time, rate_constants=None, efficiencies=None, instrument_function='none', instrument_width=10.0, light_activated=False, state_mapping=None, regularization=None, verbose=1)

forward()

Compute occupancies at all timepoints.

Returns:

occupancies – Occupancy of each structural state at each timepoint. Shape: [n_structural_states, n_timepoints]

Return type:

torch.Tensor

get_regularization_loss()

Compute regularization loss for kinetic parameters.

This implements prior distributions on the parameters: - Log-normal prior on rate constants - Beta-like prior on efficiencies (favoring values near 1)

Returns:

reg_loss – Regularization loss to be added to the main loss

Return type:

torch.Tensor

get_rate_constants()

Get current rate constants.

get_efficiencies()

Get current efficiencies.

get_time_constants()

Get current time constants (1/k_eff).

set_rate_constant(transition, value)

Set a specific rate constant.

freeze_rates(transitions=None)

Freeze rate constants (exclude from optimization).

Parameters:

transitions (list of str or None) – Transitions to freeze (e.g., [“A->B”]). If None, freezes all.

unfreeze_rates()

Unfreeze all rate constants.

freeze_efficiencies()

No-op: efficiencies are always frozen at 1.0.

unfreeze_efficiencies()

No-op: efficiencies are always frozen at 1.0.

freeze_instrument()

Freeze instrument function width.

unfreeze_instrument()

Unfreeze instrument function width.

get_parameter_groups(base_lr=0.001)

Get parameter groups with appropriate learning rates.

Kinetic parameters often benefit from different learning rates: - Rate constants (log-space): can use larger steps - Efficiencies: moderate steps - Instrument width: small steps (often well-constrained)

Parameters:

base_lr (float) – Base learning rate

Returns:

param_groups – Parameter groups suitable for torch optimizers

Return type:

list of dict

print_parameters()

Print current kinetic parameters.

plot_occupancies(path, log_scale=True, show_kinetic_states=False, figsize=(10, 6), title=None)

Plot state occupancies over time.

Parameters:

• path (str) – Output path for the plot

• log_scale (bool) – Whether to use log scale for time axis

• show_kinetic_states (bool) – If True, shows all kinetic states. If False, shows mapped structural states.

• figsize (tuple) – Figure size

• title (str or None) – Custom title

plot_comparison(target_occupancies, path, log_scale=True, figsize=(12, 5))

Plot comparison between current and target occupancies.

Useful for debugging refinement by comparing to known ground truth or to unrestrained refinement results.

Parameters:

• target_occupancies (torch.Tensor) – Target occupancies, shape [n_states, n_timepoints]

• path (str) – Output path for the plot

• log_scale (bool) – Whether to use log scale for time axis

• figsize (tuple) – Figure size

state_dict_kinetics()

Get state dict with kinetic parameters for saving/loading.

classmethod load_from_kinetics_state(state, verbose=1)

Load model from saved kinetic state.

class torchref.kinetic.occupancies_kinetics_multiexperiment(flow_chart, experiments, shared_rates=None, verbose=1)

Bases: DeviceMixin, Module

Kinetics-constrained occupancy model for multiple experiments.

This handles the case where you have multiple datasets with different conditions (e.g., different temperatures, different excitation wavelengths) but want to share some kinetic parameters across them.

Parameters:

• flow_chart (str) – Kinetic scheme (shared across experiments)

• experiments (list of dict) – Each dict contains: - ‘time’: time points for this experiment - ‘rate_constants’: experiment-specific rate constants (or None to share) - ‘name’: optional name for the experiment

• shared_rates (list of str or None) – List of transitions that should share rates across experiments. E.g., [“A->B”, “C->D”]. If None, all rates are experiment-specific.

• verbose (int) – Verbosity level

__init__(flow_chart, experiments, shared_rates=None, verbose=1)

forward(experiment_idx)

Get occupancies for a specific experiment.

Parameters:

experiment_idx (int) – Index of the experiment

Returns:

occupancies – Shape [n_states, n_timepoints]

Return type:

torch.Tensor

forward_all()

Get occupancies for all experiments.

Returns:

occupancies_list – List of occupancy tensors, one per experiment

Return type:

list of torch.Tensor

class torchref.kinetic.ModelCollection(base_models, dark_key='dark', verbose=0)

Bases: DeviceMixin, Module

Named dictionary of MixedModel instances at different timepoints.

All timepoint models share the same base structural models (ModelFT objects stored once in an nn.ModuleList). Each timepoint gets its own independent fraction parameters via _SharedMixedModel.

Keys should match DatasetCollection keys so that collection-aware targets can automatically pair datasets with models.

Parameters:

• base_models (List[ModelFT]) – The K shared structural models (e.g., ground state + intermediates).

• dark_key (str) – Key for the dark / reference entry. Default "dark".

• verbose (int) – Verbosity level.

Examples

models = ModelCollection([model_dark, model_light])
models.add_dark()                             # fractions=[1, 0]
models.add_timepoint("1ps", [0.9, 0.1])
models.add_timepoint("5ps", [0.7, 0.3])

# Access
mixed = models["1ps"]
fcalc = mixed(hkl)
print(mixed.fractions)

__init__(base_models, dark_key='dark', verbose=0)

add_timepoint(name, fractions=None, frozen_fractions=False)

Add a timepoint with given initial fractions.

Parameters:

• name (str) – Timepoint identifier (should match DatasetCollection key).

• fractions (List[float], optional) – Initial population fractions. If None, uses equal fractions.

• frozen_fractions (bool) – If True, fractions are not updated during optimization.

Returns:

Self, for method chaining.

Return type:

ModelCollection

add_dark(fractions=None)

Add the dark / reference entry.

Default fractions: [1, 0, 0, …] (100 % ground state).

Parameters:

fractions (List[float], optional) – Override dark fractions. Default is pure ground state.

Returns:

Self, for method chaining.

Return type:

ModelCollection

classmethod from_kinetics(base_models, occ_model, timepoint_names, dark_key='dark', verbose=0)

Create a ModelCollection from a kinetics occupancy model.

Parameters:

• base_models (List[ModelFT]) – Shared structural models.

• occ_model (occupancies_kinetics) – Kinetic occupancy model whose forward() returns shape [n_states, n_timepoints].

• timepoint_names (List[str]) – Names for each timepoint column (excluding dark).

• dark_key (str) – Key for the dark entry.

• verbose (int) – Verbosity level.

Return type:

ModelCollection

classmethod from_ihm(filepath, max_res=1.5, radius_angstrom=4.0, device=None, verbose=0)

Load a ModelCollection from an IHM mmCIF file.

Requires the optional python-ihm dependency.

Parameters:

• filepath (str) – Path to IHM mmCIF file.

• max_res (float) – Maximum resolution for FFT grid setup.

• radius_angstrom (float) – Radius for electron density calculation.

• device (torch.device, optional) – Device for model tensors.

• verbose (int) – Verbosity level.

Return type:

tuple of (ModelCollection, IHMEnsembleMapping)

write_ihm(filepath, mapping=None, datasets=None)

Write this ModelCollection to IHM mmCIF format.

Requires the optional python-ihm dependency.

Parameters:

• filepath (str) – Output file path.

• mapping (IHMEnsembleMapping, optional) – Mapping with metadata for round-tripping. If None, a minimal mapping is created from the collection structure.

• datasets (dict of str -> ReflectionData, optional) – Per-timepoint reflection data to embed in the CIF. Each key should match a timepoint name.

keys()

values()

items()

get(name, default=None)

property dark_key: str

property dark_model: _SharedMixedModel

Shortcut for self[dark_key].

property base_models: ModuleList

The shared structural models (owned by this collection).

property n_base_models: int

property timepoint_names: List[str]

All keys except the dark key.

property cell

property spacegroup

property device

get_all_fractions()

Current fractions for each timepoint (including dark).

get_fractions_matrix()

All fractions as a matrix [n_timepoints, n_models].

Rows are ordered by self._order (i.e. insertion order).

freeze_all_fractions()

Freeze fractions at all timepoints.

unfreeze_all_fractions()

Unfreeze fractions at all timepoints (except dark).

freeze_structures()

Freeze xyz and adp on all base models.

unfreeze_structures()

Unfreeze xyz and adp on all base models.

write_pdbs(outdir)

Write each base model to a PDB file in outdir.

Files are named base_model_0.pdb, base_model_1.pdb, etc.

Parameters:

outdir (str) – Directory to write PDB files into (must exist).

class torchref.kinetic.KineticRefinement(dataset_collection, model_collection, xray_weight_difference=2.0, xray_weight_ml=1.0, geometry_weight=10.0, adp_weight=3.0, kinetic_prior_weight=0.0, device=None, verbose=1)

Bases: DeviceMixin, Module

Orchestrator for kinetic refinement of time-resolved data.

Manages LossState registration, scaler initialization, and optimization loops for alternating structure / fraction refinement.

Parameters:

• dataset_collection (DatasetCollection) – Collection of reflection datasets keyed by timepoint name.

• model_collection (ModelCollection) – Collection of mixed models keyed by timepoint name.

• xray_weight_difference (float) – Weight for the difference X-ray target.

• xray_weight_ml (float) – Weight for the ML amplitude target.

• geometry_weight (float) – Weight for geometry restraints.

• adp_weight (float) – Weight for ADP restraints.

• kinetic_prior_weight (float) – Weight for kinetic prior regularization (0 to disable).

• device (torch.device, optional) – Computation device. If None, inferred from model collection.

• verbose (int) – Verbosity level.

__init__(dataset_collection, model_collection, xray_weight_difference=2.0, xray_weight_ml=1.0, geometry_weight=10.0, adp_weight=3.0, kinetic_prior_weight=0.0, device=None, verbose=1)

setup(cif_paths=None, kinetic_model=None, timepoints_map=None)

One-shot initialization: scalers, restraints, targets, LossState.

Parameters:

• cif_paths (List[str], optional) – Paths to CIF restraint dictionaries for ligands.

• kinetic_model (occupancies_kinetics, optional) – Kinetic model for prior regularization. If None, the kinetic prior target is not created.

• timepoints_map (Dict[str, int], optional) – Maps timepoint names to indices into the kinetic model’s time axis. Required when kinetic_model is provided.

set_weights(**kwargs)

Set target weights by name.

Parameters:

**kwargs – Keyword arguments mapping target paths to weights. E.g. set_weights(geometry=5.0, adp=2.0).

get_loss(log_values=False)

Evaluate all targets and return weighted total loss.

print_loss_summary()

Print breakdown of current losses.

refine(macro_cycles=5, niter=10, max_iter=50)

Full refinement: structures + fractions jointly.

Parameters:

• macro_cycles (int) – Number of macro-cycles.

• niter (int) – LBFGS outer iterations per macro-cycle.

• max_iter (int) – LBFGS inner iterations per step.

refine_structures(niter=10, max_iter=50)

Refine base model structures only (xyz, adp).

Freezes fractions during optimization.

refine_fractions(niter=10, max_iter=50)

Refine per-timepoint fractions only.

Freezes structures during optimization.

refine_alternating(n_cycles=5, niter_structures=10, niter_fractions=5, refit_prior_every=2, max_iter=50)

Alternating refinement: structures → fractions → refit prior.

Parameters:

• n_cycles (int) – Number of alternating cycles.

• niter_structures (int) – LBFGS iterations for structure refinement.

• niter_fractions (int) – LBFGS iterations for fraction refinement.

• refit_prior_every (int) – Refit kinetic prior every N cycles (0 to disable).

• max_iter (int) – LBFGS inner iterations per step.

refit_kinetic_prior(niter=50, lr=0.01)

Refit the kinetic model to match current free fractions.

This is the M-step in the EM-style alternation.

refine_kinetics(niter=200, lr=0.01)

Optimize kinetic model parameters directly against the X-ray targets.

Component structure factors (base ModelFT) are frozen. The kinetic model predictions are injected as fraction overrides into each timepoint’s mixed model, so the gradient path is:

kinetic params → occupancies → fractions → F_calc → X-ray loss

After optimization, the free fraction parameters are updated to match the final kinetic predictions.

Parameters:

• niter (int) – Number of Adam optimizer steps.

• lr (float) – Learning rate.

write_pdbs(outdir)

Write base model PDBs to outdir.

class torchref.kinetic.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.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.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.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.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.

Submodules

• torchref.kinetic.kinetics module
  • KineticModel
    • KineticModel.__init__()
    • KineticModel.forward()
    • KineticModel.set_baseline()
    • KineticModel.get_baselines()
    • KineticModel.get_rate_constants()
    • KineticModel.set_rate_constant()
    • KineticModel.get_efficiencies()
    • KineticModel.get_effective_rates()
    • KineticModel.get_time_constants()
    • KineticModel.parameters()
    • KineticModel.print_parameters()
    • KineticModel.plot_occupancies()
    • KineticModel.visualize()
• torchref.kinetic.occupancies module
  • occupancy_unrestrained
    • occupancy_unrestrained.__init__()
    • occupancy_unrestrained.forward()
    • occupancy_unrestrained.plot_occupancies()
  • occupancies_kinetics
    • occupancies_kinetics.__init__()
    • occupancies_kinetics.forward()
    • occupancies_kinetics.get_regularization_loss()
    • occupancies_kinetics.get_rate_constants()
    • occupancies_kinetics.get_efficiencies()
    • occupancies_kinetics.get_time_constants()
    • occupancies_kinetics.set_rate_constant()
    • occupancies_kinetics.freeze_rates()
    • occupancies_kinetics.unfreeze_rates()
    • occupancies_kinetics.freeze_efficiencies()
    • occupancies_kinetics.unfreeze_efficiencies()
    • occupancies_kinetics.freeze_instrument()
    • occupancies_kinetics.unfreeze_instrument()
    • occupancies_kinetics.get_parameter_groups()
    • occupancies_kinetics.print_parameters()
    • occupancies_kinetics.plot_occupancies()
    • occupancies_kinetics.plot_comparison()
    • occupancies_kinetics.state_dict_kinetics()
    • occupancies_kinetics.load_from_kinetics_state()
  • occupancies_kinetics_multiexperiment
    • occupancies_kinetics_multiexperiment.__init__()
    • occupancies_kinetics_multiexperiment.forward()
    • occupancies_kinetics_multiexperiment.forward_all()
• torchref.kinetic.refinement module
  • KineticRefinement
    • KineticRefinement.__init__()
    • KineticRefinement.setup()
    • KineticRefinement.set_weights()
    • KineticRefinement.get_loss()
    • KineticRefinement.print_loss_summary()
    • KineticRefinement.refine()
    • KineticRefinement.refine_structures()
    • KineticRefinement.refine_fractions()
    • KineticRefinement.refine_alternating()
    • KineticRefinement.refit_kinetic_prior()
    • KineticRefinement.refine_kinetics()
    • KineticRefinement.write_pdbs()
• torchref.kinetic.targets module
  • Targets
  • CollectionDifferenceTarget
    • CollectionDifferenceTarget.name
    • CollectionDifferenceTarget.__init__()
    • CollectionDifferenceTarget.forward()
  • CollectionMLTarget
    • CollectionMLTarget.name
    • CollectionMLTarget.__init__()
    • CollectionMLTarget.forward()
  • MultiModelGeometryTarget
    • MultiModelGeometryTarget.name
    • MultiModelGeometryTarget.__init__()
    • MultiModelGeometryTarget.forward()
    • MultiModelGeometryTarget.register_to_state()
    • MultiModelGeometryTarget.items()
  • MultiModelADPTarget
    • MultiModelADPTarget.name
    • MultiModelADPTarget.__init__()
    • MultiModelADPTarget.forward()
    • MultiModelADPTarget.register_to_state()
    • MultiModelADPTarget.items()
  • KineticPriorTarget
    • KineticPriorTarget.name
    • KineticPriorTarget.__init__()
    • KineticPriorTarget.forward()
    • KineticPriorTarget.refit_prior()