torchref.refinement.weighting package

Weighting schemes for loss component aggregation.

This module provides weighting schemes that adjust the relative importance of different loss components during refinement.

All weighting schemes inherit from BaseWeighting and receive data through LossState rather than direct refinement references.

class torchref.refinement.weighting.BaseWeighting(device=None, **kwargs)

Bases: DeviceMixin, Module, ABC

Abstract base class for weighting schemes using LossState.

Weighting schemes: - Are set up without state (only device and hyperparameters) - Receive state only when computing weights via forward() - Return weights dict (do not modify state directly)

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

Parameters:

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

name

Unique name for this weighting scheme.

Type:

str

device

Computation device.

Type:

torch.device

name: str = 'base_weighting'

__init__(device=None, **kwargs)

abstractmethod forward(state)

Compute weights from the current LossState.

Access data via state[“key”] or state.get(“key”, default). Meta data (rwork, rfree, etc.) is in state.meta. Cached losses are in state._losses.

Parameters:

state (LossState) – Current loss state with meta and _losses populated.

Returns:

Dictionary mapping component names to weights.

Return type:

Dict[str, float]

stats(state=None)

Return statistics for reporting.

Parameters:

state (LossState, optional) – If provided, can pull data from LossState.

Returns:

Statistics dictionary with StatEntry objects.

Return type:

Dict[str, StatEntry]

torchref.refinement.weighting.WeightingScheme

alias of BaseWeighting

class torchref.refinement.weighting.ResolutionWeighting(device=None, base_w_geometry=1.0, base_w_adp=1.0, d_ref=2.0, alpha=0.0)

Bases: BaseWeighting

Base prior strength with optional resolution-dependent correction.

With proper NLL sums and perfectly calibrated sigmas, w=1 would be the pure Bayesian answer. However empirical sweep on 50 structures (1.15-3.0 A) showed the monomer library sigmas lead to too-loose geometry at w=1. The optimal base is:

w_geometry ~= 10 (from median Rfree minimum) w_adp ~= 1 (lower is better for most resolutions)

Geometry and ADP have different base weights because the monomer library geometry sigmas are tighter than the ADP restraint sigmas are loose, so they need different compensation.

Optional resolution dependence:

w_geometry = base_w_geometry * (d_min / d_ref) ^ alpha w_adp = base_w_adp * (d_min / d_ref) ^ alpha

alpha=0 disables resolution correction. The sweep found the resolution dependence was weak (<20% effect), so alpha=0 is the default.

Parameters:

• device (torch.device, optional) – Computation device.

• base_w_geometry (float, optional) – Base geometry weight. Default 10.0 (from empirical sweep).

• base_w_adp (float, optional) – Base ADP weight. Default 1.0 (from empirical sweep).

• d_ref (float, optional) – Reference resolution for the optional power-law correction. Default 2.0 A.

• alpha (float, optional) – Resolution sensitivity. Default 0.0 (disabled).

name: str = 'resolution_weighting'

__init__(device=None, base_w_geometry=1.0, base_w_adp=1.0, d_ref=2.0, alpha=0.0)

forward(state)

Compute base weights with optional resolution correction.

stats(state=None)

Return statistics for reporting.

Parameters:

state (LossState, optional) – If provided, can pull data from LossState.

Returns:

Statistics dictionary with StatEntry objects.

Return type:

Dict[str, StatEntry]

class torchref.refinement.weighting.OverfittingWeighting(device=None, target_gap=0.05, min_weight=1.0, sharpness=30.0, geom_share=1.0, smoothing=0.8)

Bases: BaseWeighting

Dynamic overfitting correction based on Rfree - Rwork gap.

Uses R-factors (scale-invariant, per-reflection-normalized) rather than NLL values, which are incomparable between work/test sets after switching to summed NLLs.

When the Rfree-Rwork gap exceeds target_gap, exponentially increases regularization. The correction is applied primarily to ADP weights and only weakly to geometry weights: in crystallographic refinement, overfitting is typically driven by B-factors (which have one parameter per atom and relatively weak restraints) rather than coordinates (which are held tightly by the geometry prior).

Effective correction:

factor = min_weight + exp(sharpness * (gap - target_gap)) w_adp *= factor w_geometry *= 1 + geom_share * (factor - 1)

With geom_share = 0.2, only 20% of the overfitting correction is applied to geometry, keeping most of the effect on ADP.

Tunable parameters (as buffers): - target_gap: R-factor gap threshold. Default 0.05 (5%). - min_weight: base correction factor. Default 1.0. - sharpness: exponential response steepness. Default 30.0. - geom_share: fraction of correction applied to geometry. Default 1.0. - smoothing: EMA smoothing factor (0-1). Default 0.8.

name: str = 'overfitting_weighting'

__init__(device=None, target_gap=0.05, min_weight=1.0, sharpness=30.0, geom_share=1.0, smoothing=0.8)

forward(state)

Compute overfitting correction weights from R-factor gap.

stats(state=None)

Return statistics for reporting.

Parameters:

state (LossState, optional) – If provided, can pull data from LossState.

Returns:

Statistics dictionary with StatEntry objects.

Return type:

Dict[str, StatEntry]

class torchref.refinement.weighting.ManualWeighting(weights, device=None)

Bases: BaseWeighting

Apply fixed manual weights.

This scheme doesn’t need any state data - just returns the present weights.

name: str = 'manual_weighting'

__init__(weights, device=None)

forward(state)

Return manual weights (state is not used).

class torchref.refinement.weighting.ComponentWeighting(device=None, weights=None, component_weights=None, schemes=None, initial_xray_loss=None)

Bases: DeviceMixin, Module

Combines multiple weighting schemes using nn.ModuleDict.

Holds weighting schemes but does NOT hold a refinement reference. Weighting is computed via forward(state) which receives a LossState.

Default schemes: - ‘resolution’: ResolutionWeighting - resolution-dependent prior strength - ‘overfitting’: OverfittingWeighting - prevents overfitting via Rfree gap

Parameters:

• device (torch.device, optional) – Computation device.

• weights (dict, optional) – Manual weight overrides.

• component_weights (dict, optional) – Manual weight overrides for specific components.

• schemes (list of BaseWeighting, optional) – Additional custom weighting schemes.

schemes

Dictionary of weighting schemes.

Type:

nn.ModuleDict

__init__(device=None, weights=None, component_weights=None, schemes=None, initial_xray_loss=None)

__getitem__(key)

Get a scheme by name using dictionary-style access.

__contains__(key)

Check if a scheme exists.

keys()

Return scheme names.

values()

Return scheme instances.

items()

Return (name, scheme) pairs.

add_scheme(name, scheme)

Add a new weighting scheme.

forward(state)

Compute weights from all schemes.

Returns combined weights (multiplicative for shared keys). Does NOT modify state - just returns the computed weights.

total_loss_from_state(state)

Compute total weighted loss from a LossState.

stats(state=None)

Return statistics for reporting.

Submodules

• torchref.refinement.weighting.base_weighting module
  • BaseWeighting
    • BaseWeighting.name
    • BaseWeighting.device
    • BaseWeighting.name
    • BaseWeighting.__init__()
    • BaseWeighting.forward()
    • BaseWeighting.stats()
• torchref.refinement.weighting.component_weighting module
  • 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()
• torchref.refinement.weighting.es_policy_training module
  • ESConfig
    • ESConfig.population_size
    • ESConfig.sigma
    • ESConfig.learning_rate
    • ESConfig.n_structures_per_gen
    • ESConfig.n_steps
    • ESConfig.normalize_returns
    • ESConfig.antithetic
    • ESConfig.__init__()
  • ESResult
    • ESResult.pdb_id
    • ESResult.perturbation_idx
    • ESResult.epsilon
    • ESResult.final_rfree
    • ESResult.initial_rfree
    • ESResult.delta_rfree
    • ESResult.success
    • ESResult.error
    • ESResult.__init__()
  • GenerationResult
    • GenerationResult.generation
    • GenerationResult.results
    • GenerationResult.mean_final_rfree
    • GenerationResult.best_final_rfree
    • GenerationResult.gradient_norm
    • GenerationResult.param_update_norm
    • GenerationResult.__init__()
  • ESPolicyTrainer
    • ESPolicyTrainer.__init__()
    • ESPolicyTrainer.history
    • ESPolicyTrainer.get_flat_params()
    • ESPolicyTrainer.set_flat_params()
    • ESPolicyTrainer.sample_perturbations()
    • ESPolicyTrainer.create_perturbed_policy()
    • ESPolicyTrainer.compute_advantages()
    • ESPolicyTrainer.compute_gradient()
    • ESPolicyTrainer.update()
    • ESPolicyTrainer.evaluate_generation()
    • ESPolicyTrainer.save_checkpoint()
    • ESPolicyTrainer.load_checkpoint()
  • create_policy_network()
  • initialize_policy_sensibly()
• torchref.refinement.weighting.policy_weighting module
  • PolicyComponentWeighting
    • PolicyComponentWeighting.policy
    • PolicyComponentWeighting.sample
    • PolicyComponentWeighting.temperature
    • PolicyComponentWeighting.current_step
    • PolicyComponentWeighting.last_log_weights
    • PolicyComponentWeighting.last_weights
    • PolicyComponentWeighting.trajectory
    • PolicyComponentWeighting.name
    • PolicyComponentWeighting.__init__()
    • PolicyComponentWeighting.set_training_mode()
    • PolicyComponentWeighting.set_eval_mode()
    • PolicyComponentWeighting.reset_trajectory()
    • PolicyComponentWeighting.start_recording()
    • PolicyComponentWeighting.stop_recording()
    • PolicyComponentWeighting.increment_step()
    • PolicyComponentWeighting.extract_state_features()
    • PolicyComponentWeighting.forward()
    • PolicyComponentWeighting.compute_weights()
    • PolicyComponentWeighting.apply_to_state()
    • PolicyComponentWeighting.stats()
  • StepState
    • StepState.step
    • StepState.progress
    • StepState.improvement_rate
    • StepState.rwork
    • StepState.rfree
    • StepState.rfree_gap
    • StepState.delta_rfree
    • StepState.xray_loss
    • StepState.xray_loss_test
    • StepState.xray_work_test_ratio
    • StepState.geometry_loss
    • StepState.adp_loss
    • StepState.bond_rmsd
    • StepState.angle_rmsd
    • StepState.mean_adp_normalized
    • StepState.adp_std_normalized
    • StepState.component_losses
    • StepState.component_weights
    • StepState.__init__()
  • StepRecord
    • StepRecord.state
    • StepRecord.log_weights
    • StepRecord.weights
    • StepRecord.reward
    • StepRecord.__init__()
  • TrajectoryData
    • TrajectoryData.pdb_id
    • TrajectoryData.structure_path
    • TrajectoryData.sf_path
    • TrajectoryData.steps
    • TrajectoryData.initial_rfree
    • TrajectoryData.final_rfree
    • TrajectoryData.initial_rwork
    • TrajectoryData.final_rwork
    • TrajectoryData.total_time
    • TrajectoryData.random_seed
    • TrajectoryData.policy_version
    • TrajectoryData.success
    • TrajectoryData.error_message
    • TrajectoryData.__init__()
  • trajectory_to_dict()
• torchref.refinement.weighting.random_weighting module
  • Weight Sampling Strategy
  • RandomWeightingScheme
    • RandomWeightingScheme.base_log_weights
    • RandomWeightingScheme.current_log_weights
    • RandomWeightingScheme.current_weights
    • RandomWeightingScheme.sample_count
    • RandomWeightingScheme.name
    • RandomWeightingScheme.__init__()
    • RandomWeightingScheme.apply_step_perturbation()
    • RandomWeightingScheme.resample_trajectory()
    • RandomWeightingScheme.resample_weights()
    • RandomWeightingScheme.forward()
    • RandomWeightingScheme.sample_count
    • RandomWeightingScheme.get_log_weights()
    • RandomWeightingScheme.get_base_log_weights()
    • RandomWeightingScheme.get_weights()
    • RandomWeightingScheme.stats()
  • RandomComponentWeighting
    • RandomComponentWeighting.__init__()
    • RandomComponentWeighting.random_scheme
    • RandomComponentWeighting.resample_trajectory()
    • RandomComponentWeighting.resample_weights()
    • RandomComponentWeighting.compute_weights()
    • RandomComponentWeighting.get_sampled_log_weights()
    • RandomComponentWeighting.get_base_log_weights()
    • RandomComponentWeighting.get_sampled_weights()
    • RandomComponentWeighting.stats()