torchref.refinement.targets.combined module

Combined targets for refinement (e.g., geometry + ADP).

Integrate multiple component targets into a single combined target using nn.ModuleDict for clean organization and easy access.

Integrate into lossState via add_to_state

class torchref.refinement.targets.combined.CombinedTargets(verbose=0)[source]

Bases: Target

Base class for combined targets.

Uses nn.ModuleDict to store component targets for clean organization and easy access via dictionary-style notation.

Subclasses should override _create_targets() to define their component targets.

Parameters:

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

_targets

Dictionary of component targets.

Type:

nn.ModuleDict

__init__(verbose=0)[source]

Initialize CombinedTargets.

Parameters:

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

targets()[source]

Return registered sub-targets as ModuleDict.

__getitem__(key)[source]

Get a target by name using dictionary-style access.

__contains__(key)[source]

Check if a target exists.

keys()[source]

Return target names.

values()[source]

Return target instances.

items()[source]

Return (name, target) pairs.

target_losses()[source]

Get individual component losses (without weights).

forward()[source]

Compute total combined target loss.

stats()[source]

Get statistics from all registered targets.

get()[source]

Get individual component losses.

add_to_state(state)[source]

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

class torchref.refinement.targets.combined.CombinedModelTargets(model=None, verbose=0)[source]

Bases: ModelTarget

Base class for combined targets that only need Model (geometry/ADP targets).

Uses nn.ModuleDict to store component targets for clean organization and easy access via dictionary-style notation.

Subclasses should override _create_targets() to define their component targets.

Parameters:

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

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

_targets

Dictionary of component targets.

Type:

nn.ModuleDict

__init__(model=None, verbose=0)[source]

Initialize CombinedModelTargets.

Parameters:

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

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

targets()[source]

Return registered sub-targets as ModuleDict.

__getitem__(key)[source]

Get a target by name using dictionary-style access.

__contains__(key)[source]

Check if a target exists.

keys()[source]

Return target names.

values()[source]

Return target instances.

items()[source]

Return (name, target) pairs.

target_losses()[source]

Get individual component losses (without weights).

forward()[source]

Compute total combined target loss.

stats()[source]

Get statistics from all registered targets.

get()[source]

Get individual component losses.

add_to_state(state)[source]

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

class torchref.refinement.targets.combined.TotalGeometryTarget(model=None, verbose=0)[source]

Bases: CombinedModelTargets

Computes weighted sum of all geometry restraint NLLs.

Uses nn.ModuleDict to store component targets: - ‘bond’: BondTarget - ‘angle’: AngleTarget - ‘torsion’: TorsionTarget - ‘planarity’: PlanarityTarget - ‘chiral’: ChiralTarget - ‘nonbonded’: NonBondedHTarget (includes riding hydrogen VDW)

The torsion weight is reduced because:

  1. Protein torsions naturally deviate from ideal (Ramachandran plot)

  2. Side chain rotamers have discrete populations, not single ideals

  3. High torsion weight can over-constrain the structure

The nonbonded weight is very low because:

  1. PROLSQ repulsion is already steep (E ~ violation^4)

  2. Most contacts should be satisfied by covalent geometry

  3. High VDW weight can prevent proper packing

Set weight to 0 to disable a component.

Parameters:

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

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

Examples

geom_target = TotalGeometryTarget(model)
loss = geom_target()
bond_loss = geom_target['bond']()
for name, target in geom_target.items():
    print(f"{name}: {target()}")
get_metrics(verbosity=2)[source]

Get all geometry metrics as a flat dictionary for logging/reporting.

Parameters:

verbosity (int, optional) – Verbosity level for filtering. Default is VERBOSITY_DETAILED.

Returns:

Dictionary with validation metrics from all component targets. All values are Python floats (not tensors).

Return type:

dict

print_statistics()[source]

Print REFMAC-style geometry statistics with losses.

class torchref.refinement.targets.combined.TotalADPTarget(model=None, verbose=0)[source]

Bases: CombinedModelTargets

Total ADP restraint target combining global, similarity, and local components.

Uses nn.ModuleDict to store component targets: - ‘simu’: ADPSimilarityTarget (SIMU-like bond similarity) - ‘locality’: ADPLocalityTarget (spatial smoothness) - ‘KL’: ADPEntropyTarget (KL divergence regularization)

B-factors follow a LOG-NORMAL distribution (B > 0, right-skewed). If B ~ LogNormal(μ, σ), then log(B) ~ Normal(μ, σ).

This target combines:

  1. Similarity restraint (SIMU-like): Bond-based B-factor similarity - Enforces bonded atoms have similar B-factors - Based on covalent bond topology (strongest local constraint)

  2. Locality restraint: Spatial smoothness - nearby atoms should have similar B - Uses K-NN with distance-based sigma (d² scaling) - Medium-range spatial correlation

  3. KL divergence: Controls the spread of B-factor distribution - Prevents overfitting by controlling distribution width

Log-normal distribution properties:

  • If log(B) ~ N(μ, σ), then:

  • Mean of B: exp(μ + σ²/2)

  • Mode of B: exp(μ - σ²)

  • For typical proteins: σ_logB ≈ 0.3-0.5 (in log space)

Parameters:

  • model (Model) – Reference to Model object.

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

Examples

adp_target = TotalADPTarget(model)
loss = adp_target()
simu_loss = adp_target['simu']()
for name, target in adp_target.items():
    print(f"{name}: {target()}")
print_statistics()[source]

Print comprehensive ADP restraint statistics.

Displays statistics from all registered ADP targets.

get_metrics(verbosity=2)[source]

Get all ADP metrics as a flat dictionary for logging/reporting.

Parameters:

verbosity (int, optional) – Verbosity level for filtering. Default is VERBOSITY_DETAILED.

Returns:

Dictionary with validation metrics from all component targets. All values are Python floats (not tensors).

Return type:

dict