torchref.refinement package

Refinement module for crystallographic structure refinement.

This module provides the core refinement framework including: - Refinement classes for running optimization - Target (loss) functions for X-ray, geometry, and ADP restraints - Weighting schemes for balancing loss components - Loss aggregation and state tracking

Example

Basic refinement:

from torchref.refinement import LBFGSRefinement

refinement = LBFGSRefinement(
    data_file='reflections.mtz',
    pdb='structure.pdb',
)
refinement.refine_everything(macro_cycles=5)

Access targets and weighting schemes:

from torchref.refinement.targets import XrayTarget, BondTarget
from torchref.refinement.weighting import ComponentWeighting

class torchref.refinement.Refinement(data_file=None, pdb=None, cif=None, verbose=1, max_res=None, device=None, nbins=10, manual_weights=None, component_weights=None, column_names=None)

Bases: DeviceMixin, DebugMixin, Module

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')
   

Parameters:

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

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

• cif (str, optional) – Path to CIF file for restraints.

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

• max_res (float, optional) – Maximum resolution for reflections.

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

• weighter (LossWeightingModule, optional) – Loss weighting module. Creates default if None.

• nbins (int, optional) – Number of resolution bins. Default is 10.

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

weighter

Loss weighting module.

Type:

LossWeightingModule

__init__(data_file=None, pdb=None, cif=None, verbose=1, max_res=None, device=None, nbins=10, manual_weights=None, component_weights=None, column_names=None)

Initialize Refinement.

If data_file and pdb are provided, fully initializes the refinement. If not provided (empty init), creates a shell with empty submodules ready for load_state_dict().

Parameters:

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

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

• cif (str, optional) – Path to CIF file for restraints.

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

• max_res (float, optional) – Maximum resolution for reflections.

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

• weighter (LossWeightingModule, optional) – Loss weighting module. Creates default if None.

• nbins (int, optional) – Number of resolution bins. Default is 10.

set_xray_target_mode(mode)

Change the X-ray target mode.

Parameters:

mode (str) – X-ray target mode. Options are ‘gaussian’, ‘ls’, or ‘ml’.

property data

Expose reflection_data as ‘data’ for weighting module compatibility.

Returns:

The reflection data container.

Return type:

ReflectionData

property loss_state: LossState

Get or create the persistent LossState.

The LossState is created once and reused across refinement cycles. Targets are registered once; weights are updated each cycle.

Returns:

The persistent loss state with targets registered.

Return type:

LossState

property logger: Logger

Get or create the Logger for this refinement.

Returns:

Logger instance linked to the persistent LossState.

Return type:

Logger

reset_loss_state()

Reset the persistent LossState and Logger.

Call this if targets need to be re-registered (e.g., after changing target modes or reinitializing targets).

get_scales()

setup_scaler()

parameters(recurse=True)

Return unique parameters from this module and all submodules.

Uses the default Module.parameters() to gather parameters, then removes duplicates while preserving order to avoid passing the same tensor multiple times to the optimizer.

Parameters:

recurse (bool, optional) – If True, yields parameters of this module and all submodules. Default is True.

Returns:

List of unique parameter tensors.

Return type:

list

get_fcalc(hkl=None, recalc=False)

get_fcalc_scaled(hkl=None, recalc=False)

adp_loss()

Compute total ADP loss using TotalADPTarget.

This combines:

• Bond-based similarity (SIMU-like)

• Spread control (tighter than KL)

• Bounds penalty

Returns:

Total ADP loss value.

Return type:

torch.Tensor

get_F_calc(hkl=None, recalc=False)

get_F_calc_scaled(hkl=None, recalc=False)

nll_xray()

Compute X-ray negative log-likelihood for work and test sets.

Returns:

Tuple of (work_nll, test_nll) tensors.

Return type:

tuple of torch.Tensor

xray_loss_work()

Compute X-ray loss on work set using instantiated target.

Returns:

X-ray loss on work set.

Return type:

torch.Tensor

xray_loss_test()

Compute X-ray loss on test set using instantiated target.

Returns:

X-ray loss on test set.

Return type:

torch.Tensor

bond_loss()

Compute bond length NLL via geometry_target.

Returns:

Bond length NLL loss.

Return type:

torch.Tensor

angle_loss()

Compute angle NLL via geometry_target.

Returns:

Angle NLL loss.

Return type:

torch.Tensor

torsion_loss()

Compute torsion angle NLL via geometry_target.

Returns:

Torsion angle NLL loss.

Return type:

torch.Tensor

geometry_loss()

Compute total geometry NLL using TotalGeometryTarget.

Returns:

Total geometry NLL loss.

Return type:

torch.Tensor

loss()

Compute total loss using LossState pipeline.

Creates a LossState, populates meta, caches losses, updates weights, and returns the aggregated weighted loss.

Returns:

Total weighted loss.

Return type:

torch.Tensor

setup_component_weighting()

Set up component weighting with ResolutionWeighting + OverfittingWeighting.

populate_state_meta(state)

Populate LossState.meta with all model-level data.

Called once per macro cycle before weighting schemes are applied. This is the single location where refinement data is extracted into state.

Parameters:

state (LossState) – State to populate with meta data.

Returns:

State with meta populated.

Return type:

LossState

update_weights(state, multiply=False)

Compute weights from component_weighting and update state. Weights are clipped to [0.01, 100.0] to avoid extreme values.

Parameters:

• state (LossState) – State with meta populated.

• multiply (bool, optional) – If True, multiply existing weights by computed weights. If False, replace existing weights with computed weights.

Returns:

State with weights updated.

Return type:

LossState

create_loss_state()

Create a configured LossState for optimization.

Deprecated since version Use: the loss_state property instead for the persistent state. This method is kept for backwards compatibility.

Sets up a LossState with all targets registered as callables with hierarchical naming (e.g., ‘geometry/bond’, ‘adp/simu’). Weights are applied from component_weighting.

Usage:

from torchref.utils import validate_loss

state = refinement.create_loss_state() params = list(refinement.parameters())

# Log initial state state.aggregate(log_values=True)

# In an LBFGS closure, wrap with validate_loss so non-finite # losses warn + reject the step instead of poisoning the run. def closure():

optimizer.zero_grad() loss = state.aggregate() loss.backward() ok = validate_loss(

loss, state=state, parameters=params, context=”my_refinement”, raise_on_fail=False,

) if not ok:

for p in params:

if p.grad is not None:

p.grad.zero_()

return torch.full_like(loss.detach(), float(“inf”))

return loss

optimizer.step(closure)

# Log final state state.new_entry() state.aggregate(log_values=True)

Returns:

Configured LossState with targets and weights.

Return type:

LossState

complete_loss_state()

Update and return the persistent LossState.

Updates the persistent LossState with current meta, target info, cached losses, and weights. The state is reused across cycles.

The cached active-parameter leaf set is not refreshed here. Stale leaves are not a correctness hazard: a leaf that’s in the set but whose Parameter object was replaced externally (e.g. by Model.freeze) just gets ignored by _freeze_graph_extras, which costs a marginal amount of wasted backward work but never produces wrong answers. If you do call Model.freeze / Model.unfreeze between LossState creation and a refinement step, call state.refresh_loss_leaves() explicitly.

Returns:

Complete LossState with targets, meta, losses, and weights.

Return type:

LossState

xray_loss()

Compute X-ray loss on work set.

Returns:

X-ray loss on work set.

Return type:

torch.Tensor

restraints_loss()

Compute total geometry restraints loss.

Returns:

Total geometry restraints loss.

Return type:

torch.Tensor

collect_metrics()

Collect all metrics from component_weighting.stats().

This is the standard method for gathering refinement metrics for logging. Uses the centralized component_weighting module for all statistics. Returns full unfiltered stats - filtering is done at display time.

Returns:

Dictionary with all metrics (unfiltered, with StatEntry objects).

Return type:

dict

add_target_info_to_state(state)

Add target information from geometry and ADP targets to LossState.meta.

Deprecated since version This: method is no longer needed. Use complete_loss_state() instead, which handles all state setup in one call.

Parameters:

state (LossState) – Current loss state. Meta will be updated with target info.

Returns:

Updated loss state (unchanged).

Return type:

LossState

get_rfactor()

update_outliers(z_threshold=4.0)

plot_fcalc_vs_fobs(outpath='fcalc_vs_fobs.png')

write_out_mtz(out_mtz_path='refined_output.mtz')

collect_deposition_metadata(metadata=None)

Collect refinement statistics into a RefinementMetadata object.

Reuses existing statistics from collect_metrics(), get_rfactor(), and reflection data attributes.

Parameters:

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

Returns:

Metadata populated with final refinement statistics.

Return type:

RefinementMetadata

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)

Save the complete state of the refinement to a file.

Parameters:

path (str) – Path to save the state dictionary to.

load_state(path, strict=True)

Load the complete state of the refinement from a file.

Parameters:

• path (str) – Path to load the state dictionary from.

• strict (bool, optional) – Whether to strictly enforce that keys match. Default is True.

classmethod create_from_state_dict(state_dict, device=device(type='cpu'), verbose=1)

Create a fully initialized Refinement from a state dictionary.

This is the recommended way to restore a Refinement from a saved state. It creates the proper submodules using their respective create_from_state_dict methods, then calls PyTorch’s default load_state_dict.

Parameters:

• state_dict (dict) – State dictionary from torch.save(refinement.state_dict(), …) or from loading a checkpoint file.

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

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

Returns:

Fully initialized instance with restored state.

Return type:

Refinement

Examples

Save and load refinement state:

# Save
torch.save(refinement.state_dict(), 'refinement.pt')

# Load
state = torch.load('refinement.pt')
refinement = Refinement.create_from_state_dict(state)

# Continue refinement
rwork, rfree = refinement.get_rfactor()
print(f"Restored at R-work={rwork:.4f}, R-free={rfree:.4f}")

class torchref.refinement.LBFGSRefinement(*args, target_mode='bhattacharyya', sigma_m_scale=1.0, use_lossstate_scaler=True, **kwargs)

Bases: Refinement

LBFGS-based refinement subclass using the L-BFGS optimizer for fast convergence.

L-BFGS (Limited-memory BFGS) is a quasi-Newton optimization method that approximates the Hessian matrix, leading to much faster convergence than first-order methods.

Key advantages:

• Converges in 1-2 macro cycles (vs 5+ for Adam)

• Better final R-factors

• More stable convergence

• Automatically handles step size via line search

Parameters:

• target_mode (str, optional) – X-ray target mode (‘gaussian’, ‘ls’, or ‘ml’). Default is ‘ml’.

• *args – Passed to parent Refinement class.

• **kwargs – Passed to parent Refinement class.

target_mode

Current X-ray target mode.

Type:

str

Examples

Basic usage:

from torchref.refinement import LBFGSRefinement

refinement = LBFGSRefinement(
    data_file='data.mtz',
    pdb='model.pdb',
    target_mode='ml'
)
refinement.refine(macro_cycles=2)

LBFGS_DEFAULTS = {'history_size': 100, 'line_search_fn': 'strong_wolfe', 'lr': 1.0, 'max_iter': 20}

__init__(*args, target_mode='bhattacharyya', sigma_m_scale=1.0, use_lossstate_scaler=True, **kwargs)

Initialize LBFGS refinement.

Parameters:

• target_mode (str, optional) – X-ray target mode (‘gaussian’, ‘ls’, ‘ml’, ‘bhattacharyya’). Default is ‘bhattacharyya’.

• sigma_m_scale (float, optional) – Global multiplier for σ_m in the Bhattacharyya target only. Ignored for other target modes. Default 1.0.

• use_lossstate_scaler (bool, optional) – If True (default), refine_scaler() uses the full LossState with the body’s x-ray target — so scaler and body steps share one consistent loss. If False, falls back to Scaler.refine_lbfgs which minimises a standalone nll_xray and can pull scales in a different direction than the body optimization.

• *args – Passed to parent Refinement class.

• **kwargs – Passed to parent Refinement class.

xray_loss()

Compute X-ray loss using the instantiated target.

Returns:

X-ray loss on work set.

Return type:

torch.Tensor

refine_scaler()

Refine scaler parameters against the full refinement loss.

Builds the body LossState via complete_loss_state(), constructs a fresh LBFGS optimizer over list(self.scaler.parameters()), and delegates to LossState.step(). Because state.step disables requires_grad on every loss leaf outside the optimizer’s intent set, xyz / adp / u / occupancy are pinned for the duration — only scaler parameters move.

The critical property is that the x-ray target used here is the same one the body refine_xyz() and refine_adp() see. The legacy Scaler.refine_lbfgs() minimises a standalone nll_xray + U^2 penalty, which can pull scales in a different direction than a bhattacharyya or ml body loss and leaves the body to chase a scaler that disagrees with its own objective.

When use_lossstate_scaler is False, fall back to the legacy Scaler.refine_lbfgs() path.

Returns:

LossState with history if use_lossstate_scaler is True, otherwise the metrics dict from Scaler.refine_lbfgs().

Return type:

LossState or dict

refine_xyz()

Refine Cartesian coordinates jointly with scaler parameters.

Scaler parameters (log_scale, U, solvent terms) are included in the same LBFGS call as xyz. The joint curvature lets xyz steps see the scaler as an anchor — residuals the scaler can absorb do not have to be chased by atomic motion — and the adp/scaler_U and adp/scaler_log_scale priors bite on every step, so nothing in the scaler drifts between refine_xyz and refine_adp calls.

Returns:

State with history containing before/after loss values.

Return type:

LossState

refine_adp()

Refine ADP / U / occupancy jointly with scaler parameters.

Scaler parameters (log_scale, U, solvent terms) are included in the same LBFGS call as the ADP-block body parameters so the joint curvature can slide along the atomic-B / scaler-U degeneracy ridge together with the adp/scaler_U regularizer. XYZ is left frozen.

Returns:

State with history containing before/after loss values.

Return type:

LossState

refine_joint()

Joint LBFGS over every refinable parameter in one step.

Optimizes xyz, adp, u, occupancy, and every scaler parameter (log_scale, anisotropic U, solvent terms) in a single LBFGS call. The joint curvature couples all of them through the same x-ray target and through the adp/scaler_U / adp/scaler_log_scale priors — unlike alternating refine_xyz → refine_adp, there’s no “frozen partner” in either half that could lock the step into a locally bad direction.

Returns:

State with history containing before/after loss values.

Return type:

LossState

run_training_trajectory(policy_weighting, n_steps=10, pdb_id='', structure_path='', sf_path='', seed=None, policy_version=None)

Run a training trajectory with policy-guided refinement.

This method runs a sequence of refinement steps using a policy to select component weights. It records state-action-reward tuples for training the policy with AWR or similar algorithms.

Parameters:

• policy_weighting (PolicyComponentWeighting) – Policy weighting scheme (should be in training mode with sampling).

• n_steps (int, optional) – Number of refinement steps in the trajectory (default: 10).

• pdb_id (str, optional) – PDB identifier for recording.

• structure_path (str, optional) – Path to structure file for recording.

• sf_path (str, optional) – Path to structure factors file for recording.

• seed (int, optional) – Random seed for reproducibility.

• policy_version (str, optional) – Version identifier of the policy being used.

Returns:

Complete trajectory with state-action-reward tuples.

Return type:

TrajectoryData

run_training_trajectory_joint(policy_weighting, n_steps=10, pdb_id='', structure_path='', sf_path='', seed=None, policy_version=None)

Run a training trajectory with joint XYZ+ADP refinement.

Similar to run_training_trajectory() but refines xyz, adp, u, and occupancy together in each step. The LBFGS curvature history is reset at the start of each policy step because the weight updates invalidate any prior Hessian approximation.

Parameters:

• policy_weighting (PolicyComponentWeighting) – Policy weighting scheme (should be in training mode).

• n_steps (int, optional) – Number of refinement steps (default: 10).

• pdb_id (str, optional) – Identifiers for trajectory recording.

• structure_path (str, optional) – Identifiers for trajectory recording.

• sf_path (str, optional) – Identifiers for trajectory recording.

• seed (int, optional) – Random seed for reproducibility.

• policy_version (str, optional) – Policy version identifier.

Returns:

Complete trajectory with state-action-reward tuples.

Return type:

TrajectoryData

refine(macro_cycles=5)

Run full LBFGS refinement cycle (ADP + XYZ).

Parameters:

macro_cycles (int, optional) – Number of refinement cycles to perform. Default is 5.

Returns:

History dictionary with all metrics per cycle (hierarchical structure).

Return type:

dict

refine_everything(macro_cycles=5)

Run full LBFGS refinement cycle (ADP + XYZ) without weight screening.

Parameters:

macro_cycles (int, optional) – Number of refinement cycles to perform. Default is 5.

Returns:

History dictionary with all metrics per cycle (hierarchical structure).

Return type:

dict

class torchref.refinement.LossState(device=<factory>, targets=<factory>, weights=<factory>, history=<factory>, _losses=<factory>, _compilable=<factory>, _compiled_aggregate=None, _loss_leaves=<factory>, _resettable_modules=<factory>, meta=<factory>)

Bases: DeviceMixin

Hierarchical loss state with lazy evaluation.

device

Computation device.

Type:

torch.device

targets

Target functions keyed by hierarchical name (e.g., ‘geometry/bond’).

Type:

Dict[str, Callable]

weights

Weights keyed by name. Can be group weights (‘geometry’) or component weights (‘geometry/bond’).

Type:

Dict[str, float]

history

Log of computed values per aggregation call.

Type:

List[Dict]

meta

Model-level data (rwork, rfree, n_atoms, etc.) populated by refinement.

Type:

Dict[str, Any]

device: device

targets: Dict[str, Callable]

weights: Dict[str, float]

history: List[Dict[str, Any]]

meta: Dict[str, Any]

__getitem__(key)

Get value from meta or _losses by key.

Parameters:

key (str) – Key to look up. Checks meta first, then _losses.

Returns:

Value from meta or _losses.

Return type:

Any

Raises:

KeyError – If key not found in either dict.

__contains__(key)

Check if key exists in meta or _losses.

get(key, default=None)

Get value with default fallback.

Parameters:

• key (str) – Key to look up.

• default (Any) – Value to return if key not found.

Returns:

Value from meta, _losses, or default.

Return type:

Any

cache_losses(force=False)

Cache all target losses.

Evaluates all registered targets and stores results in _losses.

Parameters:

force (bool) – If True, re-evaluate all targets even if already cached.

Returns:

Self for chaining.

Return type:

LossState

update_meta(data)

Update meta dict with model-level data.

Parameters:

data (Dict[str, Any]) – Data to add to meta.

Returns:

Self for chaining.

Return type:

LossState

register_target(name, target, prefix=None, compile=False, probe=True)

Register a target function.

Automatically detects combined targets (like TotalGeometryTarget, TotalADPTarget) and expands them into their component targets.

Parameters:

• name (str) – Hierarchical name (e.g., ‘geometry/bond’, ‘adp/simu’).

• target (Callable) – Function that returns a loss tensor when called. Can also be a combined target with .items() method, which will be auto-expanded.

• prefix (str, optional) – Prefix to prepend to the name (e.g., ‘model1’ -> ‘model1/geometry/bond’). Useful for registering targets from multiple models in the same state.

• compile (bool) – If True, mark this target (or all its sub-targets if combined) as eligible for the compiled aggregate closure built by compile_aggregate().

• probe (bool) – If True (default), run the target’s forward once, walk the autograd graph, and merge the resulting leaf set into self._loss_leaves. The target’s dependencies (model loaded, data attached, etc.) must therefore be in place before registration. Set probe=False to skip — the leaf-set entry for this target will be empty, so step()/run() will not auto-disable any leaves on its account. Useful only for targets whose forward genuinely cannot be called at registration time.

Returns:

Self for chaining.

Return type:

LossState

register_targets(targets, prefix=None, compile=False, probe=True)

Register multiple targets from a component target or dict.

For targets with a .name attribute, uses target.name as the key. For plain callables, uses the dict key.

Parameters:

• targets (dict) – Dictionary of name -> target mappings.

• prefix (str, optional) – Prefix to prepend to all target names.

• compile (bool) – If True, propagate the compile flag to all sub-targets.

• probe (bool) – Forwarded to register_target().

set_weight(name, weight)

Set a weight value.

Parameters:

• name (str) – Weight name. Can be a group (‘geometry’) or component (‘geometry/bond’).

• weight (float) – Weight value.

Returns:

Self for chaining.

Return type:

LossState

set_weights(weights)

Set multiple weights.

get_weight(name, default=1.0)

Get a weight value.

get_effective_weight(name)

Get effective weight for a target, including group weights.

For ‘geometry/bond’, returns: weights[‘geometry’] * weights[‘geometry/bond’] Missing weights default to 1.0.

Parameters:

name (str) – Target name (e.g., ‘geometry/bond’).

Returns:

Product of all hierarchical weights.

Return type:

float

mark_compilable(names)

Mark already-registered targets as eligible for the compiled aggregate.

Parameters:

names (List[str]) – Target keys to mark (must already be registered).

Returns:

Self for chaining.

Return type:

LossState

compile_aggregate(**compile_kwargs)

Build and cache a torch.compile’d closure over all compilable targets.

Must be called after all targets and weights have been registered. Re-call if weights or compilable targets change (or call reset_compiled_aggregate()).

Parameters:

**compile_kwargs – Keyword arguments forwarded to torch.compile. Defaults to fullgraph=False so partial-graph fallback is allowed.

Returns:

Self for chaining.

Return type:

LossState

reset_compiled_aggregate()

Clear the cached compiled closure (e.g. after changing weights).

log(name, value)

Log a value to the current history entry.

Creates a new history entry if needed.

Parameters:

• name (str) – Key for the logged value.

• value (Any) – Value to log. Tensors are converted to Python floats.

new_entry()

Start a new history entry.

get_history(name)

Get all logged values for a key across history.

aggregate(log_values=False)

Evaluate all targets and compute weighted sum.

When compile_aggregate() has been called and log_values=False, the compilable targets are evaluated through a single torch.compile’d closure for improved performance. With log_values=True all targets run eagerly so per-target losses are available in _losses.

Parameters:

log_values (bool) – If True, log all losses, weights, and total to history.

Returns:

Total weighted loss.

Return type:

torch.Tensor

get_loss(name)

Get a cached loss value (after aggregate() was called).

Parameters:

name (str) – Target name.

Returns:

Cached loss, or None if not computed.

Return type:

torch.Tensor or None

active_parameters()

Return the set of leaf ``nn.Parameter``s that registered targets’ backward passes will accumulate gradient into.

Populated incrementally by register_target() via a one-shot probe forward + autograd graph walk — calling this method does not run any forward, walk any graph, or evaluate any target. The result is conservative: a target whose weight is later set to 0 still contributes its leaves here, which is harmless for the freezing logic in step() (it can only over-freeze, never under-freeze).

refresh_loss_leaves()

Re-probe every registered target and rebuild _loss_leaves and the resettable-modules cache.

Use this after external code has replaced parameter identity on the underlying model — for example after Model.freeze() / Model.unfreeze() (which rebuild refinable_params tensors). Under normal step()/run() usage no parameter identity ever changes, so this method is rarely needed.

reset_caches()

Call reset_cache() on every registered target’s submodules that expose one. Invoked automatically at the end of step().

restore_loss_leaf_grads()

Unconditionally re-enable requires_grad on every leaf in self._loss_leaves. Called at the end of step() so the next call sees a clean, fully-differentiable model regardless of what state the previous step (or external code) left things in.

run(optimizer, log=False, nsteps=1, *, context='loss_state.step')

Run a single optimizer.step(closure).

Builds the closure, validates each loss for finiteness via torchref.utils.validate_loss(), and on failure zeros the gradients and returns +inf so the strong-Wolfe line search backtracks. Automatically disables requires_grad on any leaf that the loss touches but the optimizer was not constructed with — autograd then prunes those subgraphs from the backward pass.

Technically this should work with all optimzers in pytorch that support closures but it has only been tested for LBFGS so far. The closure is built to be as general as possible, so if you have a custom optimizer that supports closures it should “just work” with this method.

Every collected reset_cache-bearing submodule is reset before the optimizer step so the closure’s first forward sees a clean cache (a previous rejected closure may have stored a NaN/inf forward result that the fingerprint would happily serve again if parameter values haven’t changed).

After the the run we call maintenance on all targets.

On exit, requires_grad=True is unconditionally re-enabled on every leaf in self._loss_leaves — defending against state bleeding between successive refinement methods.

Parameters:

• optimizer (torch.optim.Optimizer) – Optimizer to step. Its param_groups define the intent — the leaves the caller actually wants to update.

• log (bool) – If True, calls aggregate(log_values=True) before and after the optimization loop

• nsteps (int) – Number of steps to run (default 1). Only the first step’s closure caching is enabled between multiple steps. If you want to run truly independent steps, call this method multiple times with nsteps=1. This adds overhead but might be desirable if the overhead is negligible anyway.

• context (str) – Diagnostic label forwarded to validate_loss.

Returns:

The loss tensor from the last accepted closure call, or None if no closure call succeeded (every call produced non-finite loss).

Return type:

torch.Tensor or None

step(optimizer, *args, **kwargs)

Convenience method that calls run() with 1 step.

Parameters:

• optimizer (torch.optim.Optimizer) – Optimizer to run.

• *args – Forwarded to run().

• **kwargs – Forwarded to run().

get_breakdown()

Get breakdown of losses by group.

Returns:

Nested dict: {group: {component: {‘loss’: …, ‘weight’: …, ‘weighted’: …}}}

Return type:

Dict

get_group_totals()

Get total weighted loss per group.

Returns:

{group_name: total_weighted_loss}

Return type:

Dict[str, float]

format_breakdown()

Return per-target loss / weight / weighted / finite as a string.

One row per target currently in self._losses (populated by the most recent eager aggregate() call). Used by both summary() and torchref.utils.validate_loss() so the diagnostic format does not drift.

summary()

Print a per-target loss breakdown to stdout.

to(*args, **kwargs)

Move via DeviceMixin; honour an explicit device when no tensors exist yet.

clear()

Clear cached losses (not targets or weights).

clear_history()

Clear history log.

__init__(device=<factory>, targets=<factory>, weights=<factory>, history=<factory>, _losses=<factory>, _compilable=<factory>, _compiled_aggregate=None, _loss_leaves=<factory>, _resettable_modules=<factory>, meta=<factory>)

class torchref.refinement.Logger(state, verbose=1, pattern='.*', _records=<factory>, _labels=<factory>)

Bases: object

Refinement logging with verbosity-aware stat reporting.

Integrates with LossState to record, compare, and display refinement stats. Supports regex-based filtering of target names.

Parameters:

• state (LossState) – The loss state to monitor.

• verbose (int) – Verbosity level (0=essential, 1=standard, 2=detailed, 3=debug).

• pattern (str) – Regex pattern for filtering target names. Default “.*” matches all. Examples: “xray.*” for X-ray targets only, “geometry/bond” for specific target.

Examples

Basic usage:

from torchref.refinement import LossState, Logger

state = LossState(device=device)
state.register_target("xray/work", xray_target)
state.register_target("geometry/bond", bond_target)

logger = Logger(state, verbose=1)

# Record before refinement
logger.record(label="before_xyz")

# ... run refinement ...

# Record after and compare
logger.record(label="after_xyz")
logger.compare(title="XYZ Refinement")

Filtering by pattern:

# Show only X-ray targets
logger.current(pattern="xray.*")

# Create a logger that only tracks geometry
geom_logger = Logger(state, pattern="geometry.*")

state: LossState

verbose: int = 1

pattern: str = '.*'

record(label=None)

Record current refinement state.

Gathers stats from all targets in LossState, stores in history. Uses the instance’s pattern filter.

Parameters:

label (str, optional) – Label for this record (e.g., “before_xyz”, “after_xyz”).

Returns:

The recorded stats dictionary.

Return type:

dict

compare(label_before=None, label_after=None, pattern=None, title='Refinement Comparison')

Print comparison between two recorded states.

If labels not provided, compares last two records.

Parameters:

• label_before (str, optional) – Label of “before” state. Default: second-to-last record.

• label_after (str, optional) – Label of “after” state. Default: last record.

• pattern (str, optional) – Regex to filter which targets to display. Default: use instance pattern.

• title (str, optional) – Title for the comparison output. Default: “Refinement Comparison”.

current(pattern=None, title='Current State')

Print the current refinement state.

Uses latest recorded state, or records new one if none exist.

Parameters:

• pattern (str, optional) – Regex to filter which targets to display. Default: use instance pattern.

• title (str, optional) – Title for the output. Default: “Current State”.

get_record(label)

Get a specific recorded state by label.

Parameters:

label (str) – The label to look up.

Returns:

The recorded stats dictionary, or None if label not found.

Return type:

dict or None

clear()

Clear all recorded history.

property history: List[Dict[str, Any]]

Access full recording history.

__init__(state, verbose=1, pattern='.*', _records=<factory>, _labels=<factory>)

class torchref.refinement.Target(verbose=0, **kwargs)

Bases: DeviceMixin, Module

Abstract base class for all target functions.

All tunable parameters should be registered as buffers using register_buffer() so they can be accessed/modified via state_dict notation.

Supports empty initialization for state_dict loading:

target = Target()  # Creates empty shell
target.load_state_dict(torch.load('target.pt'))

LossState Integration:

Targets can work with LossState for the new pipeline:

state = target.add_to_state(state)  # Adds loss to state

Parameters:

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

name

Unique name for this target (used as loss key in LossState).

Type:

str

verbose

Verbosity level.

Type:

int

name: str = 'base_target'

__init__(verbose=0, **kwargs)

Initialize target.

Parameters:

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

forward()

Compute and return the loss. Override in subclasses.

add_to_state(state)

Compute loss and add it to the LossState.

This method enables the new LossState pipeline pattern where targets receive a state object, compute their loss, add it to the state, and return the state for chaining.

Parameters:

state (LossState) – Current loss state with computed data.

Returns:

State with this target’s loss added.

Return type:

LossState

maintenance()

Between-step housekeeping hook (no-op by default).

LossState calls this on every registered target after each successful outer optimizer step returns. Targets override this to rebuild stale internal state (VDW pair lists, solvent masks, etc.) based on how far parameters have drifted since the last refresh.

Contract

• Must be idempotent: calling it multiple times in a row on an unchanged model should not mutate the target.

• Fast path first: cheap staleness check up front, expensive rebuild only when strictly necessary. LossState calls this every outer step — the happy-path cost is paid every time.

• Must not raise on routine drift. If a rebuild fails, let the exception propagate — that’s a real bug.

class torchref.refinement.DataTarget(data=None, model=None, scaler=None, verbose=0, **kwargs)

Bases: Target

Base class for targets that need ReflectionData and optionally Model/Scaler.

This class provides a flexible interface for X-ray targets that can work in two modes:

1. With Model: Computes F_calc from the model on each forward pass

2. Without Model: Uses pre-computed F_calc passed directly

This decoupling allows targets to be used for: - Standard refinement (with model) - Analysis/scoring of pre-computed structure factors (without model) - Testing and validation workflows

All objects (model, data, scaler) are registered as proper submodules, allowing PyTorch to handle device movement and state_dict operations.

Parameters:

• data (ReflectionData, optional) – Reference to the ReflectionData object. Required for forward().

• model (Model or ModelFT, optional) – Reference to a Model object for F_calc computation. If None, F_calc must be provided to forward().

• scaler (Scaler, optional) – Reference to the Scaler object for scaling F_calc.

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

• target_value (float, optional) – Target value for this loss. Default is 0.0.

• sigma (float, optional) – Sigma parameter for weighting. Default is 0.5.

name

Unique name for this target (used as loss key in LossState).

Type:

str

_model

Reference to the model object (registered as submodule).

Type:

Model

_data

Reference to the reflection data object (registered as submodule).

Type:

ReflectionData

_scaler

Reference to the scaler object (registered as submodule).

Type:

Scaler

verbose

Verbosity level.

Type:

int

name: str = 'data_target'

__init__(data=None, model=None, scaler=None, verbose=0, **kwargs)

Initialize data target.

Parameters:

• data (ReflectionData, optional) – Reference to the ReflectionData object. Required for forward().

• model (Model or ModelFT, optional) – Reference to Model object for F_calc computation. If None, F_calc must be provided when calling forward().

• scaler (Scaler, optional) – Reference to the Scaler object.

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

property model: Model

Access the model object.

property data: ReflectionData

Access the reflection data object.

property scaler: Scaler

Access the scaler object.

property has_model: bool

Check if a model is available for F_calc computation.

get_fcalc(hkl=None, recalc=False)

Compute structure factors from model.

Parameters:

• hkl (torch.Tensor, optional) – Miller indices. If None, uses data’s hkl.

• recalc (bool, optional) – Force recalculation. Default is False.

Returns:

Complex structure factors.

Return type:

torch.Tensor

Raises:

RuntimeError – If no model is set.

get_fcalc_scaled(hkl=None, recalc=False, fcalc=None)

Compute or scale structure factors.

Parameters:

• hkl (torch.Tensor, optional) – Miller indices. If None, uses data’s hkl.

• recalc (bool, optional) – Force recalculation. Default is False.

• fcalc (torch.Tensor, optional) – Pre-computed structure factors. If provided, skips model computation.

Returns:

Scaled complex structure factors.

Return type:

torch.Tensor

get_F_calc_scaled(hkl=None, recalc=False, fcalc=None)

Compute scaled structure factor amplitudes.

Parameters:

• hkl (torch.Tensor, optional) – Miller indices. If None, uses data’s hkl.

• recalc (bool, optional) – Force recalculation. Default is False.

• fcalc (torch.Tensor, optional) – Pre-computed structure factors. If provided, skips model computation.

Returns:

Scaled structure factor amplitudes |F_calc|.

Return type:

torch.Tensor

get_rfactor()

Compute R-factors using scaler.

Returns:

(R_work, R_free) values.

Return type:

tuple

Raises:

RuntimeError – If no scaler is set.

class torchref.refinement.ModelTarget(model=None, verbose=0, **kwargs)

Bases: Target

Base class for targets that only need a Model reference.

This class provides a simpler interface for geometry and ADP targets that don’t need access to reflection data or refinement machinery. Targets inherit from this class when they only need the atomic model.

The model is registered as a proper submodule, allowing PyTorch to handle device movement and state_dict operations automatically.

Parameters:

• model (Model, optional) – Reference to the Model object.

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

• target_value (float, optional) – Target value for this loss. Default is 0.0.

• sigma (float, optional) – Sigma parameter for weighting. Default is 0.5.

name

Unique name for this target (used as loss key in LossState).

Type:

str

_model

Reference to the model object (registered as submodule).

Type:

Model

verbose

Verbosity level.

Type:

int

name: str = 'model_target'

__init__(model=None, verbose=0, **kwargs)

Initialize model target.

Parameters:

• model (Model, optional) – Reference to the Model object (optional for empty init).

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

property model: Model

Access the model object.

property restraints

Access model’s restraints (built lazily on first access).

Subpackages

• torchref.refinement.optimizers package
  • AdamWithAdaptiveNoise
    • AdamWithAdaptiveNoise.alpha
    • AdamWithAdaptiveNoise.eps
    • AdamWithAdaptiveNoise.noise_scale
    • AdamWithAdaptiveNoise.update_weight
    • AdamWithAdaptiveNoise.__init__()
    • AdamWithAdaptiveNoise.inject_noise()
    • AdamWithAdaptiveNoise.step()
    • AdamWithAdaptiveNoise.update_noise_scale()
  • ExploratoryLBFGS
    • ExploratoryLBFGS.__init__()
    • ExploratoryLBFGS.phase
    • ExploratoryLBFGS.step()
  • LangevinSA
    • LangevinSA.__init__()
    • LangevinSA.temperature
    • LangevinSA.current_step
    • LangevinSA.total_steps
    • LangevinSA.kinetic_energy
    • LangevinSA.calibrate()
    • LangevinSA.step()
  • MomentumStochasticSA
    • MomentumStochasticSA.__init__()
    • MomentumStochasticSA.step()
  • Submodules
    • torchref.refinement.optimizers.adam_noise module
      • AdamWithAdaptiveNoise
    • torchref.refinement.optimizers.exploratory_lbfgs module
      • OptimizerPhase
      • Mode
      • ScanPoint
      • Basin
      • ParameterGroup
      • ConvergenceTracker
      • ExplorationResult
      • LanczosError
      • ExploratoryLBFGS
    • torchref.refinement.optimizers.internal_coord_sa module
      • optimize_internal_coord_sa()
      • optimize_gradient_sa()
      • refine_sa_lbfgs()
    • torchref.refinement.optimizers.langevin_sa module
      • LangevinSA
    • torchref.refinement.optimizers.momentum_sa module
      • MomentumStochasticSA
    • torchref.refinement.optimizers.simulated_annealing module
      • optimize_simulated_annealing()
• torchref.refinement.targets package
  • Target
    • Target.name
    • Target.verbose
    • Target.name
    • Target.__init__()
    • Target.forward()
    • Target.add_to_state()
    • Target.maintenance()
  • ModelTarget
    • ModelTarget.name
    • ModelTarget._model
    • ModelTarget.verbose
    • ModelTarget.name
    • ModelTarget.__init__()
    • ModelTarget.model
    • ModelTarget.restraints
  • DataTarget
    • DataTarget.name
    • DataTarget._model
    • DataTarget._data
    • DataTarget._scaler
    • DataTarget.verbose
    • DataTarget.name
    • DataTarget.__init__()
    • DataTarget.model
    • DataTarget.data
    • DataTarget.scaler
    • DataTarget.has_model
    • DataTarget.get_fcalc()
    • DataTarget.get_fcalc_scaled()
    • DataTarget.get_F_calc_scaled()
    • DataTarget.get_rfactor()
  • gaussian_nll()
  • von_mises_nll()
  • adp_similarity_nll()
  • XrayTarget
    • XrayTarget.use_work_set
    • XrayTarget.__init__()
    • XrayTarget.name
    • XrayTarget.reset_get_data_cache()
    • XrayTarget.get_data()
    • XrayTarget.stats()
  • GaussianXrayTarget
    • GaussianXrayTarget.target_value
    • GaussianXrayTarget.forward()
  • MaximumLikelihoodXrayTarget
    • MaximumLikelihoodXrayTarget.forward()
  • LeastSquaresXrayTarget
    • LeastSquaresXrayTarget.__init__()
    • LeastSquaresXrayTarget.forward()
  • create_xray_target()
  • DifferenceXrayTarget
    • DifferenceXrayTarget.name
    • DifferenceXrayTarget.__init__()
    • DifferenceXrayTarget.dataset_collection
    • DifferenceXrayTarget.data_light
    • DifferenceXrayTarget.data_dark
    • DifferenceXrayTarget.model_light
    • DifferenceXrayTarget.model_dark
    • DifferenceXrayTarget.scaler_light
    • DifferenceXrayTarget.scaler_dark
    • DifferenceXrayTarget.hkl
    • DifferenceXrayTarget.get_delta_F_obs()
    • DifferenceXrayTarget.get_delta_F_calc()
    • DifferenceXrayTarget.forward()
    • DifferenceXrayTarget.stats()
  • PhaseInformedDifferenceTarget
    • PhaseInformedDifferenceTarget.name
    • PhaseInformedDifferenceTarget.__init__()
    • PhaseInformedDifferenceTarget.hkl
    • PhaseInformedDifferenceTarget.forward()
    • PhaseInformedDifferenceTarget.stats()
  • RiceDifferenceTarget
    • RiceDifferenceTarget.name
    • RiceDifferenceTarget.__init__()
    • RiceDifferenceTarget.hkl
    • RiceDifferenceTarget.forward()
    • RiceDifferenceTarget.compute_free_metrics()
    • RiceDifferenceTarget.stats()
  • TaylorCorrectedDifferenceTarget
    • TaylorCorrectedDifferenceTarget.name
    • TaylorCorrectedDifferenceTarget.__init__()
    • TaylorCorrectedDifferenceTarget.hkl
    • TaylorCorrectedDifferenceTarget.forward()
    • TaylorCorrectedDifferenceTarget.compute_free_metrics()
    • TaylorCorrectedDifferenceTarget.stats()
  • GeometryTarget
    • GeometryTarget.__init__()
    • GeometryTarget.stats()
  • BondTarget
    • BondTarget.name
    • BondTarget.__init__()
    • BondTarget.forward()
    • BondTarget.stats()
  • AngleTarget
    • AngleTarget.name
    • AngleTarget.__init__()
    • AngleTarget.forward()
    • AngleTarget.stats()
  • TorsionTarget
    • TorsionTarget.name
    • TorsionTarget.__init__()
    • TorsionTarget.forward()
    • TorsionTarget.stats()
  • PlanarityTarget
    • PlanarityTarget.name
    • PlanarityTarget.__init__()
    • PlanarityTarget.forward()
    • PlanarityTarget.stats()
  • ChiralTarget
    • ChiralTarget.name
    • ChiralTarget.__init__()
    • ChiralTarget.forward()
    • ChiralTarget.get_violations()
    • ChiralTarget.stats()
  • NonBondedTarget
    • NonBondedTarget.name
    • NonBondedTarget.__init__()
    • NonBondedTarget.c_rep
    • NonBondedTarget.sigma_vdw
    • NonBondedTarget.r_exp
    • NonBondedTarget.buffer
    • NonBondedTarget.maintenance()
    • NonBondedTarget.forward()
    • NonBondedTarget.get_violations()
    • NonBondedTarget.stats()
  • NonBondedHTarget
    • NonBondedHTarget.name
    • NonBondedHTarget.__init__()
    • NonBondedHTarget.forward()
    • NonBondedHTarget.get_violations()
    • NonBondedHTarget.stats()
  • RamachandranTarget
    • RamachandranTarget.name
    • RamachandranTarget.__init__()
    • RamachandranTarget.forward()
    • RamachandranTarget.stats()
  • ADPTarget
    • ADPTarget.__init__()
  • ADPSimilarityTarget
    • ADPSimilarityTarget.name
    • ADPSimilarityTarget.__init__()
    • ADPSimilarityTarget.simu_sigma
    • ADPSimilarityTarget.forward()
    • ADPSimilarityTarget.stats()
  • RigidBondTarget
    • RigidBondTarget.name
    • RigidBondTarget.__init__()
    • RigidBondTarget.forward()
    • RigidBondTarget.get_delta_z_stats()
    • RigidBondTarget.stats()
  • ADPEntropyTarget
    • ADPEntropyTarget.name
    • ADPEntropyTarget.__init__()
    • ADPEntropyTarget.forward()
    • ADPEntropyTarget.stats()
  • ADPLocalityTarget
    • ADPLocalityTarget.name
    • ADPLocalityTarget.__init__()
    • ADPLocalityTarget.k_neighbors
    • ADPLocalityTarget.correlation_length
    • ADPLocalityTarget.scale
    • ADPLocalityTarget.forward()
    • ADPLocalityTarget.stats()
  • CombinedTargets
    • CombinedTargets._targets
    • CombinedTargets.__init__()
    • CombinedTargets.targets()
    • CombinedTargets.__getitem__()
    • CombinedTargets.__contains__()
    • CombinedTargets.keys()
    • CombinedTargets.values()
    • CombinedTargets.items()
    • CombinedTargets.target_losses()
    • CombinedTargets.forward()
    • CombinedTargets.stats()
    • CombinedTargets.get()
    • CombinedTargets.add_to_state()
  • TotalGeometryTarget
    • TotalGeometryTarget.get_metrics()
    • TotalGeometryTarget.print_statistics()
  • TotalADPTarget
    • TotalADPTarget.print_statistics()
    • TotalADPTarget.get_metrics()
  • ForceFieldTarget
    • ForceFieldTarget.name
    • ForceFieldTarget.__init__()
    • ForceFieldTarget.forward()
    • ForceFieldTarget.stats()
  • AmberTarget
    • AmberTarget.name
    • AmberTarget.__init__()
    • AmberTarget.forward()
    • AmberTarget.stats()
  • OccupancyFloorDiagnostic
    • OccupancyFloorDiagnostic.__init__()
    • OccupancyFloorDiagnostic.compute_density_at_positions()
    • OccupancyFloorDiagnostic.analyze_at_dark_positions()
    • OccupancyFloorDiagnostic.estimate_alpha_floor_from_difference_map()
  • NegativeDensityPenalty
    • NegativeDensityPenalty.__init__()
    • NegativeDensityPenalty.forward()
  • DisplacementRegularizer
    • DisplacementRegularizer.__init__()
    • DisplacementRegularizer.forward()
  • DifferenceAmplitudeRegularizer
    • DifferenceAmplitudeRegularizer.__init__()
    • DifferenceAmplitudeRegularizer.hkl
    • DifferenceAmplitudeRegularizer.forward()
  • SampledMLPhaseTarget
    • SampledMLPhaseTarget.name
    • SampledMLPhaseTarget.__init__()
    • SampledMLPhaseTarget.name
    • SampledMLPhaseTarget.n_samples
    • SampledMLPhaseTarget.sigma_model_log
    • SampledMLPhaseTarget.use_analytical
    • SampledMLPhaseTarget.use_antithetic
    • SampledMLPhaseTarget.french_wilson_moments()
    • SampledMLPhaseTarget.compute_sigma_phi()
    • SampledMLPhaseTarget.forward()
    • SampledMLPhaseTarget.stats()
  • SampledMLDifferenceTarget
    • SampledMLDifferenceTarget.name
    • SampledMLDifferenceTarget.__init__()
    • SampledMLDifferenceTarget.hkl
    • SampledMLDifferenceTarget.n_samples
    • SampledMLDifferenceTarget.sigma_model_log
    • SampledMLDifferenceTarget.forward()
    • SampledMLDifferenceTarget.stats()
  • create_sampled_ml_target()
  • create_sampled_ml_difference_target()
  • RealSpaceTarget
    • RealSpaceTarget.VALID_MAP_TYPES
    • RealSpaceTarget.__init__()
    • RealSpaceTarget.update_mask()
  • RealSpaceCorrelationTarget
    • RealSpaceCorrelationTarget.name
    • RealSpaceCorrelationTarget.__init__()
    • RealSpaceCorrelationTarget.forward()
    • RealSpaceCorrelationTarget.stats()
  • RealSpaceDifferenceTarget
    • RealSpaceDifferenceTarget.name
    • RealSpaceDifferenceTarget.__init__()
    • RealSpaceDifferenceTarget.forward()
    • RealSpaceDifferenceTarget.stats()
  • RealSpaceExtrapolatedTarget
    • RealSpaceExtrapolatedTarget.name
    • RealSpaceExtrapolatedTarget.__init__()
    • RealSpaceExtrapolatedTarget.forward()
    • RealSpaceExtrapolatedTarget.stats()
  • CoordinateSimilarityTarget
    • CoordinateSimilarityTarget.name
    • CoordinateSimilarityTarget.__init__()
    • CoordinateSimilarityTarget.model_dark
    • CoordinateSimilarityTarget.model_light
    • CoordinateSimilarityTarget.alpha
    • CoordinateSimilarityTarget.forward()
    • CoordinateSimilarityTarget.stats()
  • Subpackages
    • torchref.refinement.targets.adp package
      • ADPTarget
      • ADPSimilarityTarget
      • RigidBondTarget
      • ADPEntropyTarget
      • ADPLocalityTarget
      • ScalerURegularizationTarget
      • ScalerLogScaleTrendTarget
      • Submodules
    • torchref.refinement.targets.geometry package
      • GeometryTarget
      • BondTarget
      • AngleTarget
      • TorsionTarget
      • PlanarityTarget
      • ChiralTarget
      • NonBondedTarget
      • NonBondedHTarget
      • RamachandranTarget
      • Submodules
    • torchref.refinement.targets.xray package
      • XrayTarget
      • GaussianXrayTarget
      • LeastSquaresXrayTarget
      • MaximumLikelihoodXrayTarget
      • BhattacharyyaXrayTarget
      • create_xray_target()
      • Submodules
  • Submodules
    • torchref.refinement.targets.amber_target module
      • Performance note
      • Design notes
      • AMBER14_STANDARD
      • AmberTarget
    • torchref.refinement.targets.base module
      • Target
      • ModelTarget
      • DataTarget
      • gaussian_nll()
      • von_mises_nll()
      • adp_similarity_nll()
      • detach_phases()
    • torchref.refinement.targets.combined module
      • CombinedTargets
      • CombinedModelTargets
      • TotalGeometryTarget
      • TotalADPTarget
    • torchref.refinement.targets.difference module
      • DifferenceXrayTarget
      • PhaseInformedDifferenceTarget
      • TaylorCorrectedDifferenceTarget
      • RiceDifferenceTarget
    • torchref.refinement.targets.forcefield_target module
      • ForceFieldTarget
    • torchref.refinement.targets.occupancy_floor_diagnostic module
      • OccupancyFloorDiagnostic
      • NegativeDensityPenalty
      • DisplacementRegularizer
      • DifferenceAmplitudeRegularizer
    • torchref.refinement.targets.realspace module
      • RealSpaceTarget
      • RealSpaceCorrelationTarget
      • RealSpaceDifferenceTarget
      • RealSpaceExtrapolatedTarget
    • torchref.refinement.targets.sampled_ml_phase_target module
      • SampledMLPhaseTarget
      • SampledMLDifferenceTarget
      • create_sampled_ml_target()
      • create_sampled_ml_difference_target()
    • torchref.refinement.targets.similarity module
      • CoordinateSimilarityTarget
• torchref.refinement.weighting package
  • BaseWeighting
    • BaseWeighting.name
    • BaseWeighting.device
    • BaseWeighting.name
    • BaseWeighting.__init__()
    • BaseWeighting.forward()
    • BaseWeighting.stats()
  • WeightingScheme
  • ResolutionWeighting
    • ResolutionWeighting.name
    • ResolutionWeighting.__init__()
    • ResolutionWeighting.forward()
    • ResolutionWeighting.stats()
  • OverfittingWeighting
    • OverfittingWeighting.name
    • OverfittingWeighting.__init__()
    • OverfittingWeighting.forward()
    • OverfittingWeighting.stats()
  • ManualWeighting
    • ManualWeighting.name
    • ManualWeighting.__init__()
    • ManualWeighting.forward()
  • ComponentWeighting
    • ComponentWeighting.schemes
    • ComponentWeighting.__init__()
    • ComponentWeighting.__getitem__()
    • ComponentWeighting.__contains__()
    • ComponentWeighting.keys()
    • ComponentWeighting.values()
    • ComponentWeighting.items()
    • ComponentWeighting.add_scheme()
    • ComponentWeighting.forward()
    • ComponentWeighting.total_loss_from_state()
    • ComponentWeighting.stats()
  • Submodules
    • torchref.refinement.weighting.base_weighting module
      • BaseWeighting
    • torchref.refinement.weighting.component_weighting module
      • BaseWeighting
      • WeightingScheme
      • ResolutionWeighting
      • OverfittingWeighting
      • ManualWeighting
      • ComponentWeighting
    • torchref.refinement.weighting.es_policy_training module
      • ESConfig
      • ESResult
      • GenerationResult
      • ESPolicyTrainer
      • create_policy_network()
      • initialize_policy_sensibly()
    • torchref.refinement.weighting.policy_weighting module
      • PolicyComponentWeighting
      • StepState
      • StepRecord
      • TrajectoryData
      • trajectory_to_dict()
    • torchref.refinement.weighting.random_weighting module
      • Weight Sampling Strategy
      • RandomWeightingScheme
      • RandomComponentWeighting

Submodules

• torchref.refinement.base_refinement module
  • Refinement
    • Refinement.device
    • Refinement.verbose
    • Refinement.reflection_data
    • Refinement.model
    • Refinement.scaler
    • Refinement.weighter
    • Refinement.__init__()
    • Refinement.set_xray_target_mode()
    • Refinement.data
    • Refinement.loss_state
    • Refinement.logger
    • Refinement.reset_loss_state()
    • Refinement.get_scales()
    • Refinement.setup_scaler()
    • Refinement.parameters()
    • Refinement.get_fcalc()
    • Refinement.get_fcalc_scaled()
    • Refinement.adp_loss()
    • Refinement.get_F_calc()
    • Refinement.get_F_calc_scaled()
    • Refinement.nll_xray()
    • Refinement.xray_loss_work()
    • Refinement.xray_loss_test()
    • Refinement.bond_loss()
    • Refinement.angle_loss()
    • Refinement.torsion_loss()
    • Refinement.geometry_loss()
    • Refinement.loss()
    • Refinement.setup_component_weighting()
    • Refinement.populate_state_meta()
    • Refinement.update_weights()
    • Refinement.create_loss_state()
    • Refinement.complete_loss_state()
    • Refinement.xray_loss()
    • Refinement.restraints_loss()
    • Refinement.collect_metrics()
    • Refinement.add_target_info_to_state()
    • Refinement.get_rfactor()
    • Refinement.update_outliers()
    • Refinement.plot_fcalc_vs_fobs()
    • Refinement.write_out_mtz()
    • Refinement.collect_deposition_metadata()
    • Refinement.write_out_pdb()
    • Refinement.write_out_cif()
    • Refinement.save_state()
    • Refinement.load_state()
    • Refinement.create_from_state_dict()
• torchref.refinement.lbfgs_refinement module
  • LBFGSRefinement
    • LBFGSRefinement.target_mode
    • LBFGSRefinement.LBFGS_DEFAULTS
    • LBFGSRefinement.__init__()
    • LBFGSRefinement.xray_loss()
    • LBFGSRefinement.refine_scaler()
    • LBFGSRefinement.refine_xyz()
    • LBFGSRefinement.refine_adp()
    • LBFGSRefinement.refine_joint()
    • LBFGSRefinement.run_training_trajectory()
    • LBFGSRefinement.run_training_trajectory_joint()
    • LBFGSRefinement.refine()
    • LBFGSRefinement.refine_everything()
• torchref.refinement.logger module
  • Logger
    • Logger.state
    • Logger.verbose
    • Logger.pattern
    • Logger.record()
    • Logger.compare()
    • Logger.current()
    • Logger.get_record()
    • Logger.clear()
    • Logger.history
    • Logger.__init__()
• torchref.refinement.loss_state module
  • LossState
    • LossState.device
    • LossState.targets
    • LossState.weights
    • LossState.history
    • LossState.meta
    • LossState.device
    • LossState.targets
    • LossState.weights
    • LossState.history
    • LossState.meta
    • LossState.__getitem__()
    • LossState.__contains__()
    • LossState.get()
    • LossState.cache_losses()
    • LossState.update_meta()
    • LossState.register_target()
    • LossState.register_targets()
    • LossState.set_weight()
    • LossState.set_weights()
    • LossState.get_weight()
    • LossState.get_effective_weight()
    • LossState.mark_compilable()
    • LossState.compile_aggregate()
    • LossState.reset_compiled_aggregate()
    • LossState.log()
    • LossState.new_entry()
    • LossState.get_history()
    • LossState.aggregate()
    • LossState.get_loss()
    • LossState.active_parameters()
    • LossState.refresh_loss_leaves()
    • LossState.reset_caches()
    • LossState.restore_loss_leaf_grads()
    • LossState.run()
    • LossState.step()
    • LossState.get_breakdown()
    • LossState.get_group_totals()
    • LossState.format_breakdown()
    • LossState.summary()
    • LossState.to()
    • LossState.clear()
    • LossState.clear_history()
    • LossState.__init__()
  • create_loss_state()