torchref.model.closed_segmented_internal_coordinates module
Closed segmented internal coordinate parametrization for atomic structures.
This module provides the ClosedSegmentedInternalCoordinateTensor class which extends the segmented approach with analytical chain closure. Between larger segments (~18 residues each), 3-residue “junctions” are placed whose backbone torsions (phi, psi) are slave DOFs solved by Newton’s method. Backward gradients flow via the Implicit Function Theorem, giving exact gradients without differentiating through the solver.
Key differences from SegmentedInternalCoordinateTensor: - Larger segments (default 18 residues vs 3) - Junction residues between segments maintain chain continuity - Junction phi/psi are solved, not free parameters - IFT provides exact gradients through the closure constraint - Junctions preferentially placed in loop regions
- class torchref.model.closed_segmented_internal_coordinates.ClosedSegmentedInternalCoordinateTensor(initial_xyz, pdb, n_aa_per_segment=18, junction_size=3, bond_cutoff=2.0, cif_dict=None, prefer_loops=True, requires_grad=True, dtype=None, device=None)[source]
Bases:
DeviceMixin,CachedForwardMixin,ModuleParameter wrapper using segmented internal coordinates with chain closure.
Stores: per-segment bond_lengths, angles, torsions, segment_positions, segment_orientations. Junction backbone torsions are slave DOFs solved by Newton’s method with IFT gradients.
- Parameters:
initial_xyz (torch.Tensor) – Initial Cartesian coordinates of shape (N, 3).
pdb (pd.DataFrame) – PDB DataFrame with columns ‘chainid’, ‘resseq’, ‘name’, ‘index’, ‘resname’.
n_aa_per_segment (int, optional) – Number of amino acids per segment. Default is 18.
junction_size (int, optional) – Number of residues per junction. Default is 3.
bond_cutoff (float, optional) – Distance cutoff for bond detection in Angstroms. Default is 2.0.
cif_dict (dict, optional) – CIF dictionary containing bond definitions per residue type.
prefer_loops (bool, optional) – If True, slide junctions to prefer loop regions. Default is True.
requires_grad (bool, optional) – Whether parameters should have gradients. Default is True.
dtype (torch.dtype, optional) – Data type for tensors.
device (torch.device, optional) – Device for tensors.
- __init__(initial_xyz, pdb, n_aa_per_segment=18, junction_size=3, bond_cutoff=2.0, cif_dict=None, prefer_loops=True, requires_grad=True, dtype=None, device=None)[source]
- property dtype
- property device
- forward()[source]
Reconstruct Cartesian xyz from internal coordinates with chain closure.
Steps: 1. Place segment atoms (existing NeRF pipeline) 2. Solve junction closures (Newton + IFT) 3. Place junction backbone atoms 4. Place junction sidechain atoms 5. Apply frozen overlay
- Returns:
Reconstructed Cartesian coordinates of shape (N, 3).
- Return type:
- shake(magnitude=0.1)[source]
Add Gaussian noise to internal parameters.
Perturbs torsions, bond lengths, and bond angles. Segment positions and orientations are NOT perturbed because random independent translation of segments creates gaps that exceed the junction chain’s reach. During optimization, the optimizer adjusts these rigid-body DOFs smoothly.