torchref.restraints.builders module

Restraint Builder Classes for Crystallographic Refinement

This module provides a hierarchy of builder classes for constructing geometry restraints efficiently. The base class handles common functionality while specialized child classes implement restraint-type-specific logic.

The builder pattern allows: - Single-pass iteration over residues - Easy unit testing of individual components - Memory-efficient accumulation and finalization - Sorted indices for cache-friendly access

Classes

ResidueIterator: Efficient iterator over residues with pre-grouped data.
RestraintBuilder: Abstract base class for all restraint builders.
BondRestraintBuilder: Builder for bond length restraints.
AngleRestraintBuilder: Builder for angle restraints.
TorsionRestraintBuilder: Builder for torsion angle restraints.
PlaneRestraintBuilder: Builder for planarity restraints.
ChiralRestraintBuilder: Builder for chiral volume restraints.

class torchref.restraints.builders.ResidueIterator(pdb, filter_atom_type=None)[source]

Bases: object

Efficient iterator over residues with pre-grouped data.

Pre-groups PDB data by (chainid, resseq) once at initialization, avoiding O(N) DataFrame filtering per residue. This reduces the complexity from O(N×R) to O(N log N) where N is the number of atoms and R is the number of residues.

Parameters:

pdb (pd.DataFrame) – PDB DataFrame with columns ‘chainid’, ‘resseq’, ‘resname’, ‘name’, ‘index’.
filter_atom_type (str, optional) – If provided, only include atoms with this ATOM type (e.g., ‘ATOM’ for protein).

pdb

Reference to the input DataFrame.

Type:: pd.DataFrame

groups

List of (chainid, resseq) tuples for iteration.

Type:: list of tuple

Examples

iterator = ResidueIterator(model.pdb)
for chain_id, resseq, residue_df in iterator:
    print(f"Processing {chain_id}:{resseq}")

__init__(pdb, filter_atom_type=None)[source]

Initialize with pre-grouping of residues.

__iter__()[source]

Iterate over residues.

Yields:

chain_id (str) – Chain identifier.
resseq (int) – Residue sequence number.
residue (pd.DataFrame) – DataFrame containing all atoms in this residue.

__len__()[source]

Return number of residues.

get_consecutive_pairs()[source]

Iterate over consecutive residue pairs within each chain.

Useful for building inter-residue restraints (peptide bonds, etc.).

Yields:

residue_i (pd.DataFrame) – First residue in the pair.
residue_next (pd.DataFrame) – Next consecutive residue (resseq + 1).

class torchref.restraints.builders.RestraintBuilder(verbose=0)[source]

Bases: ABC

Abstract base class for restraint builders.

Provides common functionality for accumulating restraint data during residue iteration and finalizing to sorted tensors. Child classes implement the specific logic for each restraint type.

Parameters:: verbose (int, default 0) – Verbosity level for debug output.

_indices

Accumulated index arrays.

Type:: list of np.ndarray

_references

Accumulated reference value arrays.

Type:: list of np.ndarray

_sigmas

Accumulated sigma arrays.

Type:: list of np.ndarray

_count

Total number of restraints added.

Type:: int

restraint_type: str = 'base'

atom_columns: List[str] = []

n_atoms: int = 0

__init__(verbose=0)[source]

Initialize empty accumulator lists.

reset()[source]: Clear all accumulated data.

abstractmethod process_residue(residue, cif_restraints)[source]

Process a single residue, matching atoms to CIF restraints.

Parameters:

residue (pd.DataFrame) – Residue atoms with ‘name’ and ‘index’ columns.
cif_restraints (pd.DataFrame) – Restraints from CIF dictionary for this residue type.

Returns:

Number of restraints added.

Return type:

int

finalize(device, sort_indices=True, min_sigma=0.0001)[source]

Convert accumulated data to sorted tensors.

Parameters:

device (torch.device) – Target device for tensors.
sort_indices (bool, default True) – Whether to sort by first atom index for memory locality.
min_sigma (float, default 1e-4) – Minimum sigma value to avoid division by zero.

Returns:

Dictionary with ‘indices’, ‘references’, ‘sigmas’ tensors, or None if no restraints were accumulated.

Return type:

dict or None

property count: int: Return total number of restraints accumulated.

class torchref.restraints.builders.BondRestraintBuilder(verbose=0)[source]

Bases: RestraintBuilder

Builder for bond length restraints.

Bond restraints define expected distances between pairs of bonded atoms. Each restraint specifies: atom1, atom2, target_distance, sigma.

Examples

builder = BondRestraintBuilder()
for residue in residues:
    builder.process_residue(residue, cif_bonds)
result = builder.finalize(device)
print(result['indices'].shape)  # (n_bonds, 2)

restraint_type: str = 'bond'

atom_columns: List[str] = ['atom1', 'atom2']

n_atoms: int = 2

process_residue(residue, cif_restraints)[source]

Process bond restraints for a single residue.

Parameters:

residue (pd.DataFrame) – Residue atoms with ‘name’ and ‘index’ columns.
cif_restraints (pd.DataFrame) – Bond restraints with ‘atom1’, ‘atom2’, ‘value’, ‘sigma’ columns.

Returns:

Number of bond restraints added.

Return type:

int

class torchref.restraints.builders.AngleRestraintBuilder(verbose=0)[source]

Bases: RestraintBuilder

Builder for angle restraints.

Angle restraints define expected angles between three atoms (1-2-3). The angle is measured at the central atom (atom2).

Examples

builder = AngleRestraintBuilder()
for residue in residues:
    builder.process_residue(residue, cif_angles)
result = builder.finalize(device)
print(result['indices'].shape)  # (n_angles, 3)

restraint_type: str = 'angle'

atom_columns: List[str] = ['atom1', 'atom2', 'atom3']

n_atoms: int = 3

process_residue(residue, cif_restraints)[source]

Process angle restraints for a single residue.

Parameters:

residue (pd.DataFrame) – Residue atoms with ‘name’ and ‘index’ columns.
cif_restraints (pd.DataFrame) – Angle restraints with ‘atom1’, ‘atom2’, ‘atom3’, ‘value’, ‘sigma’ columns.

Returns:

Number of angle restraints added.

Return type:

int

class torchref.restraints.builders.TorsionRestraintBuilder(verbose=0)[source]

Bases: RestraintBuilder

Builder for torsion angle restraints.

Torsion restraints define expected dihedral angles between four atoms. Includes periodicity information for symmetric torsions (e.g., period=3 for methyl groups).

_periods

Accumulated periodicity arrays.

Type:: list of np.ndarray

Examples

builder = TorsionRestraintBuilder()
for residue in residues:
    builder.process_residue(residue, cif_torsions)
result = builder.finalize(device)
print(result['indices'].shape)  # (n_torsions, 4)
print(result['periods'].shape)  # (n_torsions,)

restraint_type: str = 'torsion'

atom_columns: List[str] = ['atom1', 'atom2', 'atom3', 'atom4']

n_atoms: int = 4

__init__(verbose=0)[source]

Initialize with additional periods accumulator.

reset()[source]: Clear all accumulated data including periods.

process_residue(residue, cif_restraints)[source]

Process torsion restraints for a single residue.

Filters out torsions with sigma=0 (undefined/flexible torsions).

Parameters:

residue (pd.DataFrame) – Residue atoms with ‘name’ and ‘index’ columns.
cif_restraints (pd.DataFrame) – Torsion restraints with ‘atom1’-‘atom4’, ‘value’, ‘sigma’, and optionally ‘periodicity’ columns.

Returns:

Number of torsion restraints added.

Return type:

int

finalize(device, sort_indices=True, min_sigma=0.0001)[source]

Convert accumulated data to sorted tensors, including periods.

Parameters:

device (torch.device) – Target device for tensors.
sort_indices (bool, default True) – Whether to sort by first atom index for memory locality.
min_sigma (float, default 1e-4) – Minimum sigma value.

Returns:

Dictionary with ‘indices’, ‘references’, ‘sigmas’, ‘periods’ tensors.

Return type:

dict or None

class torchref.restraints.builders.PlaneRestraintBuilder(verbose=0)[source]

Bases: RestraintBuilder

Builder for planarity restraints.

Plane restraints define groups of atoms that should be coplanar. Unlike other restraints, planes have variable atom counts (3-10 atoms). Results are grouped by atom count for efficient tensor operations.

_planes_by_count

Dictionary mapping atom_count -> {‘indices’: […], ‘sigmas’: […]}.

Type:: dict

Examples

builder = PlaneRestraintBuilder()
for residue in residues:
    builder.process_residue(residue, cif_planes)
result = builder.finalize(device)
# Returns dict like {'4_atoms': {...}, '6_atoms': {...}}

restraint_type: str = 'plane'

atom_columns: List[str] = ['atom']

n_atoms: int = None

__init__(verbose=0)[source]

Initialize with planes-by-count dictionary.

reset()[source]: Clear all accumulated data.

process_residue(residue, cif_restraints)[source]

Process plane restraints for a single residue.

Groups plane atoms by plane_id and stores separately by atom count.

Parameters:

residue (pd.DataFrame) – Residue atoms with ‘name’ and ‘index’ columns.
cif_restraints (pd.DataFrame) – Plane restraints with ‘plane_id’, ‘atom’, ‘sigma’ columns.

Returns:

Number of plane restraints added.

Return type:

int

finalize(device, sort_indices=True, min_sigma=0.0001)[source]

Convert accumulated data to tensors grouped by atom count.

Parameters:

device (torch.device) – Target device for tensors.
sort_indices (bool, default True) – Whether to sort by first atom index.
min_sigma (float, default 1e-4) – Minimum sigma value.

Returns:

Dictionary with keys like ‘4_atoms’, ‘6_atoms’, each containing ‘indices’ and ‘sigmas’ tensors.

Return type:

dict or None

class torchref.restraints.builders.ChiralRestraintBuilder(ideal_volume=2.5, sigma=0.2, verbose=0)[source]

Bases: RestraintBuilder

Builder for chiral volume restraints.

Chiral restraints define the expected handedness (R/S configuration) of tetrahedral centers. Each restraint specifies a center atom and three neighbors, with a signed ideal volume.

_ideal_volumes

Accumulated ideal volume arrays (signed based on chirality).

Type:: list of np.ndarray

ideal_volume

Magnitude of ideal chiral volume in Å³.

Type:: float

sigma

Standard deviation for restraint in Å³.

Type:: float

Examples

builder = ChiralRestraintBuilder(ideal_volume=2.5, sigma=0.2)
for residue in residues:
    builder.process_residue(residue, cif_chirals)
result = builder.finalize(device)
print(result['indices'].shape)  # (n_chirals, 4)
print(result['ideal_volumes'].shape)  # (n_chirals,)

restraint_type: str = 'chiral'

atom_columns: List[str] = ['atom_centre', 'atom1', 'atom2', 'atom3']

n_atoms: int = 4

__init__(ideal_volume=2.5, sigma=0.2, verbose=0)[source]

Initialize with chiral volume parameters.

Parameters:

ideal_volume (float, default 2.5) – Magnitude of ideal chiral volume in Å³.
sigma (float, default 0.2) – Standard deviation for restraint in Å³.
verbose (int, default 0) – Verbosity level.

reset()[source]: Clear all accumulated data.

process_residue(residue, cif_restraints)[source]

Process chiral restraints for a single residue.

Parameters:

residue (pd.DataFrame) – Residue atoms with ‘name’ and ‘index’ columns.
cif_restraints (pd.DataFrame) – Chiral restraints with ‘atom_centre’, ‘atom1’, ‘atom2’, ‘atom3’, ‘volume_sign’ columns.

Returns:

Number of chiral restraints added.

Return type:

int

finalize(device, sort_indices=True, min_sigma=0.0001)[source]

Convert accumulated data to sorted tensors.

Parameters:

device (torch.device) – Target device for tensors.
sort_indices (bool, default True) – Whether to sort by first atom index.
min_sigma (float, default 1e-4) – Minimum sigma value.

Returns:

Dictionary with ‘indices’, ‘ideal_volumes’, ‘sigmas’ tensors.

Return type:

dict or None

class torchref.restraints.builders.InterResidueBondBuilder(verbose=0)[source]

Bases: object

Builder for inter-residue bond restraints (peptide bonds, disulfide bonds).

Unlike intra-residue builders, this handles bonds between atoms in different residues. It processes residue pairs and link definitions.

Parameters:: verbose (int, default 0) – Verbosity level for debug output.

_indices

Accumulated bond index arrays.

Type:: list of np.ndarray

_references

Accumulated reference distance arrays.

Type:: list of np.ndarray

_sigmas

Accumulated sigma arrays.

Type:: list of np.ndarray

Examples

builder = InterResidueBondBuilder()
for res_i, res_next in iterator.get_consecutive_pairs():
    builder.process_peptide_bond(res_i, res_next, trans_link)
result = builder.finalize(device)

__init__(verbose=0)[source]

Initialize empty accumulator lists.

reset()[source]: Clear all accumulated data.

process_peptide_bond(residue_i, residue_next, link_bonds)[source]

Process peptide bond restraints between consecutive residues.

Parameters:

residue_i (pd.DataFrame) – First residue (C-terminal of bond).
residue_next (pd.DataFrame) – Second residue (N-terminal of bond).
link_bonds (pd.DataFrame) – Bond definitions from link dictionary with columns: ‘atom_1_comp_id’, ‘atom1’, ‘atom_2_comp_id’, ‘atom2’, ‘value’, ‘sigma’.

Returns:

Number of bond restraints added.

Return type:

int

process_disulfide_bond(sg1_idx, sg2_idx, bond_length, bond_sigma)[source]

Process a single disulfide bond restraint.

Parameters:

sg1_idx (int) – Index of first SG atom.
sg2_idx (int) – Index of second SG atom.
bond_length (float) – Target bond length in Å.
bond_sigma (float) – Sigma for restraint in Å.

Returns:

Always returns 1.

Return type:

int

finalize(device, sort_indices=True, min_sigma=0.0001)[source]

Convert accumulated data to sorted tensors.

Parameters:

device (torch.device) – Target device for tensors.
sort_indices (bool, default True) – Whether to sort by first atom index.
min_sigma (float, default 1e-4) – Minimum sigma value.

Returns:

Dictionary with ‘indices’, ‘references’, ‘sigmas’ tensors.

Return type:

dict or None

property count: int: Return total number of restraints accumulated.

class torchref.restraints.builders.InterResidueAngleBuilder(verbose=0)[source]

Bases: object

Builder for inter-residue angle restraints (peptide angles).

Handles angles that span two consecutive residues.

Parameters:: verbose (int, default 0) – Verbosity level for debug output.

__init__(verbose=0)[source]

Initialize empty accumulator lists.

reset()[source]: Clear all accumulated data.

process_peptide_angles(residue_i, residue_next, link_angles)[source]

Process peptide angle restraints between consecutive residues.

Parameters:

residue_i (pd.DataFrame) – First residue.
residue_next (pd.DataFrame) – Second residue.
link_angles (pd.DataFrame) – Angle definitions from link dictionary.

Returns:

Number of angle restraints added.

Return type:

int

process_disulfide_angles(res1_atoms, res2_atoms, link_angles)[source]

Process disulfide angle restraints.

Parameters:

res1_atoms (pd.DataFrame) – First cysteine residue atoms.
res2_atoms (pd.DataFrame) – Second cysteine residue atoms.
link_angles (pd.DataFrame) – Angle definitions from disulfide link.

Returns:

Number of angle restraints added.

Return type:

int

finalize(device, sort_indices=True, min_sigma=0.0001)[source]

Convert accumulated data to sorted tensors.

property count: int: Return total number of restraints accumulated.

class torchref.restraints.builders.InterResidueTorsionBuilder(verbose=0)[source]

Bases: object

Builder for inter-residue torsion restraints (phi, psi, omega).

Handles backbone torsion angles that span consecutive residues. Separates phi, psi, and omega angles for different treatment.

Parameters:: verbose (int, default 0) – Verbosity level for debug output.

__init__(verbose=0)[source]

Initialize empty accumulator lists for each torsion type.

reset()[source]: Clear all accumulated data.

process_peptide_torsions(residue_i, residue_next, link_torsions)[source]

Process backbone torsion restraints between consecutive residues.

Parameters:

residue_i (pd.DataFrame) – First residue.
residue_next (pd.DataFrame) – Second residue.
link_torsions (pd.DataFrame) – Torsion definitions from link dictionary with ‘id’ column indicating ‘phi’, ‘psi’, or ‘omega’.

Returns:

Number of (phi, psi, omega) torsions added.

Return type:

tuple of (int, int, int)

process_disulfide_torsions(res1_atoms, res2_atoms, link_torsions)[source]

Process disulfide torsion restraints.

Parameters:

res1_atoms (pd.DataFrame) – First cysteine residue atoms.
res2_atoms (pd.DataFrame) – Second cysteine residue atoms.
link_torsions (pd.DataFrame) – Torsion definitions from disulfide link.

Returns:

Number of torsion restraints added.

Return type:

int

finalize_phi(device, sort_indices=True)[source]

Finalize phi angle indices and periods.

finalize_psi(device, sort_indices=True)[source]

Finalize psi angle indices and periods.

finalize_omega(device, sort_indices=True)[source]

Finalize omega angle restraints.

finalize_disulfide(device, sort_indices=True)[source]

Finalize disulfide torsion restraints.

class torchref.restraints.builders.InterResiduePlaneBuilder(verbose=0)[source]

Bases: object

Builder for inter-residue plane restraints (peptide planes).

Handles planes that span consecutive residues (e.g., the planar peptide bond group).

Parameters:: verbose (int, default 0) – Verbosity level for debug output.

__init__(verbose=0)[source]

Initialize with planes-by-count dictionary.

reset()[source]: Clear all accumulated data.

process_peptide_planes(residue_i, residue_next, link_planes)[source]

Process peptide plane restraints between consecutive residues.

Parameters:

residue_i (pd.DataFrame) – First residue.
residue_next (pd.DataFrame) – Second residue.
link_planes (pd.DataFrame) – Plane definitions from link dictionary with ‘plane_id’, ‘atom_comp_id’, ‘atom’, ‘sigma’ columns.

Returns:

Number of plane restraints added.

Return type:

int

finalize(device, sort_indices=True, min_sigma=0.0001)[source]

Convert accumulated data to tensors grouped by atom count.

property count: int: Return total number of restraints accumulated.