torchref.refinement.base_refinement

Base class for crystallographic refinement.

Classes

Refinement([data_file, pdb, cif, verbose, ...])

Refinement class to handle the overall crystallographic refinement process.

class torchref.refinement.base_refinement.Refinement(data_file=None, pdb=None, cif=None, verbose=1, max_res=None, device=None, nbins=10, n_iso_coeff=6, column_names=None, wavelength=1.0, anomalous_threshold=0.5, french_wilson=True, anomalous=None, adp_mode='isotropic', xray_mode='ml', sigma_a_max=0.99, shrink=True, scale_target='ls', aniso_selection=None)

Refinement class to handle the overall crystallographic refinement process.

Supports two initialization patterns:

1. Empty initialization (for state_dict loading):

   refinement = Refinement()  # Creates empty shell with submodules
   refinement.load_state_dict(torch.load('refinement.pt'))
   

2. Full initialization with file paths:

   refinement = Refinement(data_file='data.mtz', pdb='model.pdb')
   

Constructor parameters are documented on __init__().

device

Computation device.

Type:

torch.device

verbose

Verbosity level.

Type:

int

reflection_data

Reflection data container.

Type:

ReflectionData

model

Structure factor model (includes lazy restraints via model.restraints).

Type:

ModelFT

scaler

Scale factor calculator.

Type:

Scaler

weighting

Loss weighting scheme holding the data/prior group weights. Defaults to ManualWeighting(DEFAULT_GROUP_WEIGHTS); reassign to change the scheme.

Type:

BaseWeighting

weighter

Vestigial state-dict placeholder, always None; the live weighting knob is weighting.

Type:

None

__init__(data_file=None, pdb=None, cif=None, verbose=1, max_res=None, device=None, nbins=10, n_iso_coeff=6, column_names=None, wavelength=1.0, anomalous_threshold=0.5, french_wilson=True, anomalous=None, adp_mode='isotropic', xray_mode='ml', sigma_a_max=0.99, shrink=True, scale_target='ls', aniso_selection=None)

Initialize Refinement, fully if data_file and pdb are given.

Without them this is an empty init: a shell with empty submodules, ready for load_state_dict().

Parameters:

• data_file (str, optional) – Path to the MTZ or CIF file holding reflection data.

• pdb (str, optional) – Path to the PDB or CIF file holding the initial model.

• cif (str, optional) – Path to a CIF file of restraints (monomer library).

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

• max_res (float, optional) – High-resolution cutoff; defaults to the data’s own limit.

• device (torch.device, optional) – Computation device. Defaults to the configured default device.

• nbins (int, optional) – Number of resolution bins used to seed the scaler’s scale. Default 10.

• n_iso_coeff (int, optional) – Number of Chebyshev terms in the scaler’s isotropic scale. Default 6.

• column_names (dict, optional) – Mapping of logical column roles to MTZ column labels.

• wavelength (float, optional) – X-ray wavelength in Angstroms for the anomalous (f’/f’’) correction. 0 means “no anomalous refinement”: it disables the correction and forces a Friedel-merged read, overriding anomalous to False.

• anomalous_threshold (float, optional) – Threshold controlling anomalous data handling. Default 0.5.

• french_wilson (bool, optional) – Derive amplitudes from intensities via French-Wilson. Set False to use existing F/SIGF columns when the MTZ also carries intensities.

• anomalous (bool, optional) – Anomalous (Bijvoet) load preference. None auto-detects F(+)/F(-) (or I(+)/I(-)) and loads Friedel pairs when present, enabling the model’s f’’ term; True forces it, False forces a merged load.

• adp_mode (str, optional) – ADP parametrization: "isotropic" (default) refines a per-atom B-factor, "anisotropic" a 6-component U tensor for the atoms selected by aniso_selection (see Model.set_adp_mode()).

• xray_mode (str, optional) – X-ray target taxonomy row; see set_xray_target_mode().

• sigma_a_max (optional) – sigma_A estimator knobs; see torchref.refinement.model_error_estimation.sigma_a.

• shrink (optional) – sigma_A estimator knobs; see torchref.refinement.model_error_estimation.sigma_a.

• scale_target (str, optional) – Objective the scale fit minimises; see torchref.scaling.scaler_base.ScalerBase.refine_lbfgs().

• aniso_selection (str, optional) – Phenix-style selection of atoms refined anisotropically when adp_mode="anisotropic". Defaults to all non-water heavy atoms.

set_xray_target_mode(mode)

Change the X-ray target mode.

Parameters:

mode (str) – X-ray target mode: ‘ml’ (default), ‘ml_noalpha’, ‘ml_full’, ‘nll_beta’, ‘nll’, ‘ls’, ‘ls_wunit_k1’. See torchref.refinement.targets.xray._specs for the taxonomy.

property data

reflection_data under the name the weighting modules expect.

property loss_state: LossState

The persistent LossState, created on first access and reused across refinement cycles (targets registered once, weights re-applied).

property logger: Logger

The Logger bound to loss_state, created on first access.

reset_loss_state()

Drop the persistent LossState and Logger so targets are re-registered.

Required after changing target modes or reinitializing targets; without it the stale state keeps serving the old targets and weights.

refine_scaler()

Refit the scaler against the current model.

The only place a scale gets fitted. Every driver routes here, so the scale the in-run R-factor is computed from is the one an external 0-cycle score would produce.

Warm: the existing scale is the starting point and carries across macrocycles. For a cold start – fresh c_iso, rebuilt solvent and anisotropy – call get_scales().

The objective is self.scale_target, not the body target: scaling is a nuisance-magnitude fit that need not carry a model-error term, alpha is degenerate with the scale being fitted, and for ml_full the body target would put a 32-node quadrature inside every line-search evaluation. The fit runs on the same LossState machinery as the body steps, differing only in the loss and in exposing only the scaler’s parameters to the optimizer; see refine_lbfgs().

No-op when the scaler is ``None`` – targets such as ls_wunit_k1 in binwise_optimal mode compute their own scale and leave it unset.

Returns:

The scaler’s per-step metrics, or None when there is no scaler.

Return type:

dict or None

get_scales()

Cold-start the scaler against the current model: initialize() then refine_scaler().

initialize() replaces c_iso with a fresh parameter and rebuilds the solvent and anisotropy terms, so this discards a refined scale. Use it at construction, when the model has been swapped, or for a one-shot 0-cycle score; inside a macrocycle loop call refine_scaler() instead.

setup_scaler()

Construct self.scaler from self._scaler_class (default Scaler), wired to the current model, data, nbins, n_iso_coeff and device.

parameters(recurse=True)

Unique parameters of this module and, with recurse, its submodules.

Deduplicates Module.parameters() in order, so a tensor shared between two submodules is not handed to the optimizer twice. Returns a list, not a generator.

get_fcalc(hkl=None, recalc=False)

Complex F_calc per reflection, from the model.

Parameters:

• hkl (array_like, optional) – Reflection indices. None evaluates on the data’s own reflections via reflection_data.structure_factors, returning the canonical-ASU convention. An explicit hkl is used as given.

• recalc (bool, optional) – Force recomputation rather than reusing the cached SF.

get_fcalc_scaled(hkl=None, recalc=False)

scaler(F_calc); see get_fcalc() for hkl and recalc.

adp_loss()

Total ADP loss: bond-based B similarity, locality smoothness, and the shifted inverse-gamma distribution prior registered by TotalADPTarget.

get_F_calc(hkl=None, recalc=False)

|F_calc|; see get_fcalc() for hkl and recalc.

get_F_calc_scaled(hkl=None, recalc=False)

|scaler(F_calc)|; see get_fcalc() for hkl and recalc.

nll_xray()

(work_nll, test_nll) from the two instantiated x-ray targets.

xray_loss_work()

X-ray loss on the work set.

xray_loss_test()

X-ray loss on the test set.

bond_loss()

Bond-length NLL component of the geometry target.

angle_loss()

Bond-angle NLL component of the geometry target.

torsion_loss()

Torsion-angle NLL component of the geometry target.

geometry_loss()

Total geometry NLL (all components of TotalGeometryTarget).

create_loss_state()

A fresh configured LossState; prefer the persistent loss_state.

complete_loss_state()

Refresh the persistent LossState’s cached losses and return it.

The cached active-parameter leaf set is not refreshed. A stale leaf is only wasted backward work, never a wrong answer – but after calling Model.freeze/unfreeze mid-run, call state.refresh_loss_leaves() yourself.

xray_loss()

Alias for xray_loss_work().

restraints_loss()

Alias for geometry_loss().

collect_metrics()

R-factors, geometry and ADP stats for logging, unfiltered.

Filtering by verbosity happens at display time, so the returned dict holds StatEntry objects rather than plain values.

add_target_info_to_state(state)

Deprecated no-op that returns state unchanged; use complete_loss_state(), which does all state setup in one call.

get_rfactor()

(R_work, R_free) for the current model.

Delegates to the work X-ray target, the single source of truth: R is computed from exactly the scaled |F_calc| the target’s loss sees (the scaler’s scaling, or the target’s own closed-form per-bin scale for binwise_optimal). See XrayTarget.get_rfactor().

plot_fcalc_vs_fobs(outpath='fcalc_vs_fobs.png')

Scatter-plot calculated vs observed amplitudes, saved as a PNG at outpath.

write_out_mtz(out_mtz_path='refined_output.mtz', anomalous=None)

Write refined map coefficients to an MTZ file.

Parameters:

• out_mtz_path (str) – Output MTZ path.

• anomalous (bool, optional) – True emits a phenix-style anomalous MTZ: maps and merged columns in the canonical ASU (Friedel mates merged by mean amplitude) plus unstacked F-obs(+/-) / F-model(+/-) on the same ASU index. False writes the per-row layout. None picks anomalous when the data were loaded as Bijvoet pairs (reflection_data.friedel_merged is False).

collect_deposition_metadata(metadata=None)

Collect refinement statistics into a RefinementMetadata.

Parameters:

metadata (RefinementMetadata, optional) – Existing metadata to merge with (e.g. input-file pass-through). Refinement statistics take precedence over pass-through values.

write_out_pdb(out_pdb_path='refined_output.pdb', metadata=None)

Write refined PDB with optional metadata header.

Parameters:

• out_pdb_path (str) – Output PDB file path.

• metadata (RefinementMetadata, optional) – Metadata for PDB header. If None, auto-collected from refinement.

write_out_cif(out_cif_path='refined_output.cif', metadata=None)

Write refined coordinates as mmCIF with metadata.

Parameters:

• out_cif_path (str) – Output mmCIF file path.

• metadata (RefinementMetadata, optional) – Metadata for mmCIF categories. If None, auto-collected from refinement.

save_state(path)

torch.save the full refinement state dict to path.

load_state(path, strict=True)

Load a saved state dict from path into this instance.

Requires submodules that already match the checkpoint’s structure; to build one from scratch use create_from_state_dict(). strict enforces an exact key match.

classmethod create_from_state_dict(state_dict, device=None, verbose=1)

Rebuild a fully initialized Refinement from a saved state dict.

The recommended restore path: it rebuilds reflection data, model and scaler through their own factories before calling load_state_dict, which load_state() cannot do. Restraints are normally lazy via model.restraints; the standalone handling here is a legacy state-dict path and does not make them a first-class persisted submodule.

Parameters:

• state_dict (dict) – From torch.save(refinement.state_dict(), ...) or a checkpoint file.

• device (torch.device, optional) – Device to place tensors on. Defaults to the configured default device.

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

Returns:

Fully initialized instance with restored state.

Return type:

Refinement