torchref.restraints.builders_fast module

Fast Restraint Builder Classes with Internal Build Logic

This module provides high-performance restraint builders that handle the entire build process internally. No external looping required.

Usage:

from torchref.restraints.builders_fast import (

BondRestraintBuilder, AngleRestraintBuilder, TorsionRestraintBuilder, PlaneRestraintBuilder, ChiralRestraintBuilder,

)

# Simple API - just call build() once bond_builder = BondRestraintBuilder() bond_restraints = bond_builder.build(pdb, cif_dict, device)

# Or use the all-in-one function from torchref.restraints.builders_fast import build_all_restraints restraints = build_all_restraints(pdb, cif_dict, device)

class torchref.restraints.builders_fast.PreprocessedPDB(pdb)

Bases: object

Pre-processed PDB data as NumPy arrays for fast iteration.

Converts DataFrame to arrays once, computes residue boundaries, enabling O(1) access to residue data without DataFrame operations.

Supports altloc expansion: residues with alternate conformations are expanded into multiple conformations, each with common atoms plus the specific altloc atoms.

__init__(pdb)

Initialize from PDB DataFrame.

Parameters:

pdb (pd.DataFrame) – PDB DataFrame with standard columns.

get_residue_data(residue_idx)

Get atom data for a residue.

Returns:

• atom_names (np.ndarray) – Atom names for this residue.

• atom_indices (np.ndarray) – Global atom indices.

• resname (str) – Residue name.

Return type:

Tuple[ndarray, ndarray, str]

has_duplicate_atoms(residue_idx)

Check if residue has duplicate atom names (altlocs).

has_altlocs(residue_idx)

Check if residue has any alternate conformations.

get_altloc_conformations(residue_idx)

Iterate over altloc conformations for a residue.

For residues without altlocs, yields once with all atoms. For residues with altlocs, yields once per unique altloc, each time with common atoms (no altloc) + altloc-specific atoms.

Yields:

• atom_names (np.ndarray) – Atom names for this conformation.

• atom_indices (np.ndarray) – Global atom indices.

• resname (str) – Residue name.

class torchref.restraints.builders_fast.PreprocessedCIF(cif_dict)

Bases: object

Pre-processed CIF restraints as NumPy arrays per residue type.

__init__(cif_dict)

Initialize from CIF dictionary.

Parameters:

cif_dict (dict) – CIF dictionary with restraints per residue type.

class torchref.restraints.builders_fast.RestraintBuilder(verbose=0)

Bases: ABC

Abstract base class for restraint builders.

All builders share the same API:

builder = SomeRestraintBuilder(verbose=0) result = builder.build(pdb, cif_dict, device)

__init__(verbose=0)

Initialize builder.

abstractmethod build(pdb, cif_dict, device, sort_indices=True)

Build restraints from PDB and CIF data.

Parameters:

• pdb (pd.DataFrame) – PDB DataFrame with atom data.

• cif_dict (dict) – CIF dictionary with restraints per residue type.

• device (torch.device) – Target device for output tensors.

• sort_indices (bool, default True) – Whether to sort by first atom index for cache efficiency.

Returns:

Dictionary with restraint tensors, or None if no restraints found.

Return type:

dict or None

class torchref.restraints.builders_fast.BondRestraintBuilder(verbose=0)

Bases: RestraintBuilder

Fast bond restraint builder.

Usage:

builder = BondRestraintBuilder() result = builder.build(pdb, cif_dict, device) # result = {‘indices’: tensor, ‘references’: tensor, ‘sigmas’: tensor}

build(pdb, cif_dict, device, sort_indices=True)

Build all bond restraints.

class torchref.restraints.builders_fast.AngleRestraintBuilder(verbose=0)

Bases: RestraintBuilder

Fast angle restraint builder.

Usage:

builder = AngleRestraintBuilder() result = builder.build(pdb, cif_dict, device)

build(pdb, cif_dict, device, sort_indices=True)

Build all angle restraints.

class torchref.restraints.builders_fast.TorsionRestraintBuilder(verbose=0)

Bases: RestraintBuilder

Fast torsion restraint builder.

Usage:

builder = TorsionRestraintBuilder() result = builder.build(pdb, cif_dict, device) # result includes ‘periods’ tensor

build(pdb, cif_dict, device, sort_indices=True)

Build all torsion restraints.

class torchref.restraints.builders_fast.PlaneRestraintBuilder(verbose=0)

Bases: RestraintBuilder

Fast plane restraint builder.

Returns planes grouped by atom count (e.g., ‘4_atoms’, ‘5_atoms’).

Usage:

builder = PlaneRestraintBuilder() result = builder.build(pdb, cif_dict, device) # result = {‘4_atoms’: {‘indices’: …, ‘sigmas’: …}, ‘5_atoms’: {…}}

build(pdb, cif_dict, device, sort_indices=True)

Build all plane restraints, grouped by atom count.

class torchref.restraints.builders_fast.ChiralRestraintBuilder(verbose=0)

Bases: RestraintBuilder

Fast chiral restraint builder.

Usage:

builder = ChiralRestraintBuilder() result = builder.build(pdb, cif_dict, device) # result includes ‘ideal_volumes’ tensor

build(pdb, cif_dict, device, sort_indices=True)

Build all chiral restraints.

torchref.restraints.builders_fast.build_all_restraints(pdb, cif_dict, device, verbose=0)

Build all intra-residue restraints at once.

Parameters:

• pdb (pd.DataFrame) – PDB DataFrame with atom data.

• cif_dict (dict) – CIF dictionary with restraints per residue type.

• device (torch.device) – Target device for output tensors.

• verbose (int) – Verbosity level.

Returns:

Dictionary with all restraint types: {

’bond’: {‘indices’: …, ‘references’: …, ‘sigmas’: …}, ‘angle’: {…}, ‘torsion’: {…, ‘periods’: …}, ‘plane’: {‘4_atoms’: {…}, ‘5_atoms’: {…}, …}, ‘chiral’: {…, ‘ideal_volumes’: …}

}

Return type:

dict

class torchref.restraints.builders_fast.PreprocessedLinkData(link_dict)

Bases: object

Pre-processed link restraint data for fast inter-residue matching.

Converts link DataFrames to NumPy arrays for efficient access.

__init__(link_dict)

Initialize from link dictionary.

Parameters:

link_dict (dict) – Dictionary with link restraints (e.g., TRANS, disulf).

class torchref.restraints.builders_fast.InterResidueBondBuilder(verbose=0)

Bases: object

Fast builder for inter-residue bond restraints.

Usage:

builder = InterResidueBondBuilder() result = builder.build(pdb, link_dict, device)

# Or for disulfides (incremental): builder = InterResidueBondBuilder() for sg1_idx, sg2_idx, length, sigma in disulfide_pairs:

builder.process_disulfide_bond(sg1_idx, sg2_idx, length, sigma)

result = builder.finalize(device)

__init__(verbose=0)

Initialize builder with accumulators for disulfide bonds.

reset()

Clear all accumulated data.

process_disulfide_bond(sg1_idx, sg2_idx, bond_length, bond_sigma)

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)

Convert accumulated disulfide bond data to 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.

build(pdb, link_dict, device, filter_atom_type='ATOM', sort_indices=True)

Build all inter-residue bond restraints.

Parameters:

• pdb (pd.DataFrame) – PDB DataFrame.

• link_dict (dict) – Link dictionary with ‘bonds’ DataFrame.

• device (torch.device) – Target device.

• filter_atom_type (str, optional) – Filter to only this atom type (e.g., ‘ATOM’ for protein).

• sort_indices (bool) – Whether to sort output by first atom index.

Returns:

Dictionary with restraint tensors.

Return type:

dict or None

class torchref.restraints.builders_fast.InterResidueAngleBuilder(verbose=0)

Bases: object

Fast builder for inter-residue angle restraints.

Usage:

builder = InterResidueAngleBuilder() result = builder.build(pdb, link_dict, device)

# Or for disulfides (incremental): builder = InterResidueAngleBuilder() builder.process_disulfide_angles(res1_atoms, res2_atoms, link_angles) result = builder.finalize(device)

__init__(verbose=0)

Initialize builder with accumulators for disulfide angles.

reset()

Clear all accumulated data.

process_disulfide_angles(res1_atoms, res2_atoms, link_angles)

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)

Convert accumulated data to sorted tensors.

property count: int

Return total number of restraints accumulated.

build(pdb, link_dict, device, filter_atom_type='ATOM', sort_indices=True, next_resname_filter=None, exclude_next_resname=None)

Build all inter-residue angle restraints.

Parameters:

• pdb (pd.DataFrame) – Atom DataFrame.

• link_dict (Dict) – Link definition dictionary containing angle parameters.

• device (torch.device) – Target device for tensors.

• filter_atom_type (str, optional) – Filter to this ATOM type (default “ATOM”).

• sort_indices (bool, optional) – Sort output by first atom index (default True).

• next_resname_filter (str, optional) – If set, only build angles for pairs where the second (next) residue has this residue name (e.g. “PRO” for proline links).

• exclude_next_resname (str, optional) – If set, skip pairs where the second (next) residue has this residue name. Useful for excluding PRO from TRANS angles when PTRANS is handled separately.

class torchref.restraints.builders_fast.InterResidueTorsionBuilder(verbose=0)

Bases: object

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

Usage:

builder = InterResidueTorsionBuilder() result = builder.build(pdb, link_dict, device) # result = {‘phi’: {…}, ‘psi’: {…}, ‘omega’: {…}}

# Or for disulfides (incremental): builder = InterResidueTorsionBuilder() builder.process_disulfide_torsions(res1_atoms, res2_atoms, link_torsions) result = builder.finalize_disulfide(device)

__init__(verbose=0)

Initialize builder with accumulators for disulfide torsions.

reset()

Clear all accumulated disulfide data.

process_disulfide_torsions(res1_atoms, res2_atoms, link_torsions)

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_disulfide(device, sort_indices=True)

Finalize disulfide torsion restraints.

property disulfide_count: int

Return total number of disulfide torsion restraints accumulated.

build(pdb, link_dict, device, filter_atom_type='ATOM', sort_indices=True)

Build all inter-residue torsion restraints.

Returns separate phi, psi, omega, and ramachandran results.

class torchref.restraints.builders_fast.InterResiduePlaneBuilder(verbose=0)

Bases: object

Fast builder for inter-residue plane restraints (peptide planes).

Usage:

builder = InterResiduePlaneBuilder() result = builder.build(pdb, link_dict, device)

__init__(verbose=0)

Initialize builder.

build(pdb, link_dict, device, filter_atom_type='ATOM', sort_indices=True)

Build all inter-residue plane restraints, grouped by atom count.

class torchref.restraints.builders_fast.ResidueIterator(pdb, filter_atom_type=None)

Bases: object

Efficient iterator over residues.

Provided for compatibility - prefer using build_all_restraints() or the individual builder.build() methods instead.

__init__(pdb, filter_atom_type=None)

Initialize with pre-grouping of residues.

__iter__()

Iterate over residues.

__len__()

Return number of residues.

get_consecutive_pairs()

Iterate over consecutive residue pairs within each chain.