torchref.restraints.builders_numba module

Fast Restraint Builder Classes using NumPy and Numba

This module provides optimized builder classes that avoid Pandas operations in the hot loop.

torchref.restraints.builders_numba.find_atom_index(atom_names, target)

Find index of target atom name in array.

Returns -1 if not found.

torchref.restraints.builders_numba.match_bonds_numba(residue_atom_names, residue_atom_indices, bond_atom1, bond_atom2, bond_values, bond_sigmas, out_idx1, out_idx2, out_refs, out_sigmas)

Match bond restraints for a single residue.

Returns number of matched bonds.

torchref.restraints.builders_numba.match_angles_numba(residue_atom_names, residue_atom_indices, angle_atom1, angle_atom2, angle_atom3, angle_values, angle_sigmas, out_idx1, out_idx2, out_idx3, out_refs, out_sigmas)

Match angle restraints for a single residue.

torchref.restraints.builders_numba.match_torsions_numba(residue_atom_names, residue_atom_indices, torsion_atom1, torsion_atom2, torsion_atom3, torsion_atom4, torsion_values, torsion_sigmas, torsion_periods, out_idx1, out_idx2, out_idx3, out_idx4, out_refs, out_sigmas, out_periods)

Match torsion restraints for a single residue.

torchref.restraints.builders_numba.match_chirals_numba(residue_atom_names, residue_atom_indices, chiral_center, chiral_atom1, chiral_atom2, chiral_atom3, chiral_volume_signs, chiral_sigmas, out_center, out_idx1, out_idx2, out_idx3, out_signs, out_sigmas)

Match chiral restraints for a single residue.

class torchref.restraints.builders_numba.PreprocessedPDB(pdb)

Bases: object

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

Converts DataFrame operations to array lookups.

__init__(pdb)

Initialize from PDB DataFrame.

Parameters:

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

get_residue_atoms(residue_idx)

Get atom names and indices for a residue.

Parameters:

residue_idx (int) – Index into residue arrays.

Returns:

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

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

Return type:

Tuple[ndarray, ndarray]

class torchref.restraints.builders_numba.PreprocessedCIF(cif_dict)

Bases: object

Pre-processed CIF restraints as NumPy arrays.

Stores restraint data per residue type for fast lookup.

__init__(cif_dict)

Initialize from CIF dictionary.

Parameters:

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

class torchref.restraints.builders_numba.FastResidueIterator(preprocessed, filter_protein=False)

Bases: object

Fast residue iterator using pre-processed PDB data.

__init__(preprocessed, filter_protein=False)

Initialize from preprocessed PDB.

Parameters:

• preprocessed (PreprocessedPDB) – Pre-processed PDB data.

• filter_protein (bool) – If True, only iterate over ATOM records.

__iter__()

Iterate over residues.

Yields:

• residue_idx (int) – Index of residue.

• chain_id (str) – Chain identifier.

• resseq (int) – Residue sequence number.

• resname (str) – Residue name.

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

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

class torchref.restraints.builders_numba.FastBondBuilder(verbose=0)

Bases: object

Fast bond restraint builder using Numba.

__init__(verbose=0)

process_residue(atom_names, atom_indices, cif_bonds)

Process bonds for a single residue.

Parameters:

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

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

• cif_bonds (dict) – Pre-processed bond data from CIF.

Returns:

Number of bonds added.

Return type:

int

finalize(device, sort_indices=True, min_sigma=0.0001)

Convert accumulated data to tensors.

reset()

class torchref.restraints.builders_numba.FastAngleBuilder(verbose=0)

Bases: object

Fast angle restraint builder using Numba.

__init__(verbose=0)

process_residue(atom_names, atom_indices, cif_angles)

finalize(device, sort_indices=True, min_sigma=0.0001)

reset()

class torchref.restraints.builders_numba.FastTorsionBuilder(verbose=0)

Bases: object

Fast torsion restraint builder using Numba.

__init__(verbose=0)

process_residue(atom_names, atom_indices, cif_torsions)

finalize(device, sort_indices=True, min_sigma=0.0001)

reset()

class torchref.restraints.builders_numba.FastPlaneBuilder(verbose=0)

Bases: object

Fast plane restraint builder.

__init__(verbose=0)

process_residue(atom_names, atom_indices, cif_planes)

Process plane restraints for a residue.

finalize(device, sort_indices=True, min_sigma=0.0001)

Finalize plane restraints, grouped by atom count.

reset()

class torchref.restraints.builders_numba.FastChiralBuilder(verbose=0)

Bases: object

Fast chiral restraint builder using Numba.

__init__(verbose=0)

process_residue(atom_names, atom_indices, cif_chirals)

finalize(device, sort_indices=True, min_sigma=0.0001)

reset()

torchref.restraints.builders_numba.build_all_restraints_fast(pdb, cif_dict, device, verbose=0)

Build all restraints using the fast implementation.

Parameters:

• pdb (pd.DataFrame) – PDB DataFrame.

• cif_dict (dict) – CIF dictionary with restraints.

• device (torch.device) – Target device.

• verbose (int) – Verbosity level.

Returns:

Dictionary with all restraint tensors.

Return type:

dict