torchref.kinetic.occupancies module

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