torchref.model.model_collection module

Model collection for time-resolved kinetic refinement.

Provides ModelCollection — a named dictionary of MixedModel instances at different timepoints that share the same base structural models (ModelFT). Keys match DatasetCollection keys so targets can automatically pair them.

class torchref.model.model_collection.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).