SYMBOL INDEX (1085 symbols across 92 files)

FILE: src/atomica/data/converter/atom_blocks_to_frag_blocks.py
  function atom_blocks_to_frag_blocks (line 9) | def atom_blocks_to_frag_blocks(

FILE: src/atomica/data/converter/ligand_utils.py
  function _load_smiles_data (line 11) | def _load_smiles_data():
  function _load_inchikey_data (line 27) | def _load_inchikey_data():
  function get_ligand_smiles (line 43) | def get_ligand_smiles(ligand_code: str) -> str:
  function get_ligand_inchikey (line 57) | def get_ligand_inchikey(ligand_code: str) -> str:

FILE: src/atomica/data/converter/mol2_to_blocks.py
  function mol2_to_blocks (line 11) | def mol2_to_blocks(mol2_file: str, using_hydrogen: bool = False, molecul...

FILE: src/atomica/data/converter/pdb_lig_to_blocks.py
  function extract_pdb_ligand (line 7) | def extract_pdb_ligand(pdb, lig_code, chain_id, smiles, lig_idx:int=None...

FILE: src/atomica/data/converter/pdb_to_list_blocks.py
  function pdb_to_list_blocks (line 14) | def pdb_to_list_blocks(pdb: str, selected_chains: Optional[List[str]]=None,
  function atoms_array_to_blocks (line 124) | def atoms_array_to_blocks(atoms_array: AtomArray) -> List[Block]:
  function get_residues (line 150) | def get_residues(atoms_array: AtomArray) -> Tuple[np.ndarray, List[Tuple...
  function pdb_to_list_blocks_and_atom_array (line 159) | def pdb_to_list_blocks_and_atom_array(pdb: str, selected_chains: Optiona...

FILE: src/atomica/data/converter/sm_pdb_to_blocks.py
  function sm_pdb_to_blocks (line 10) | def sm_pdb_to_blocks(ligand_path, fragment=None, mol_idx=0):

FILE: src/atomica/data/dataset.py
  class Block (line 27) | class Block:
    method __init__ (line 28) | def __init__(self, symbol: str, units: List[Atom]) -> None:
    method __len__ (line 32) | def __len__(self):
    method __iter__ (line 35) | def __iter__(self):
    method coords (line 39) | def coords(self):
    method to_data (line 42) | def to_data(self):
  function blocks_to_data (line 53) | def blocks_to_data(*blocks_list: List[List[Block]]):
  function data_to_blocks (line 99) | def data_to_blocks(data, fragmentation_method=None):
  function blocks_to_coords (line 132) | def blocks_to_coords(blocks: List[Block]):
  function df_to_blocks (line 148) | def df_to_blocks(df, key_residue='residue', key_insertion_code='insertio...
  function blocks_interface (line 186) | def blocks_interface(blocks1, blocks2, dist_th, return_indexes=False):
  function item_to_pdb_file (line 205) | def item_to_pdb_file(item, output_pdb_file):
  function filter_pdb_with_item (line 263) | def filter_pdb_with_item(item, pdb_path):
  function filter_for_segment (line 289) | def filter_for_segment(data, keep_segment):
  function randomly_sample_blocks (line 315) | def randomly_sample_blocks(data, p_keep=1.0, rng=None, p_none=0.0):
  class BlockGeoAffDataset (line 386) | class BlockGeoAffDataset(torch.utils.data.Dataset):
    method __init__ (line 388) | def __init__(self, data_file, database=None, dist_th=6, n_cpu=4, suffi...
    method _load_data_file (line 428) | def _load_data_file(self):
    method _post_process (line 433) | def _post_process(self, items, processed_data):
    method _preprocess (line 438) | def _preprocess(self, item):
    method __len__ (line 448) | def __len__(self):
    method __getitem__ (line 451) | def __getitem__(self, idx):
    method collate_fn (line 468) | def collate_fn(cls, batch):
  class PDBBindBenchmark (line 483) | class PDBBindBenchmark(torch.utils.data.Dataset):
    method __init__ (line 485) | def __init__(self, data_file):
    method __len__ (line 490) | def __len__(self):
    method __getitem__ (line 493) | def __getitem__(self, idx):
    method collate_fn (line 513) | def collate_fn(cls, batch):
  class PDBDataset (line 538) | class PDBDataset(torch.utils.data.Dataset):
    method __init__ (line 540) | def __init__(self, data_file):
    method __len__ (line 545) | def __len__(self):
    method __getitem__ (line 548) | def __getitem__(self, idx):
    method collate_fn (line 566) | def collate_fn(cls, batch):
  class ProtInterfaceDataset (line 580) | class ProtInterfaceDataset(PDBDataset):
    method __init__ (line 581) | def __init__(self, data_file):
    method __getitem__ (line 588) | def __getitem__(self, idx):
    method collate_fn (line 611) | def collate_fn(cls, batch):
  class LabelledPDBDataset (line 622) | class LabelledPDBDataset(torch.utils.data.Dataset):
    method __init__ (line 624) | def __init__(self, data_file, random_block_sampling=False, p_keep=1.0,...
    method __len__ (line 638) | def __len__(self):
    method __getitem__ (line 641) | def __getitem__(self, idx):
    method collate_fn (line 672) | def collate_fn(cls, batch):
  class MultiClassLabelledPDBDataset (line 694) | class MultiClassLabelledPDBDataset(torch.utils.data.Dataset):
    method __init__ (line 696) | def __init__(self, data_file, random_block_sampling=False, p_keep=1.0,...
    method __len__ (line 710) | def __len__(self):
    method __getitem__ (line 713) | def __getitem__(self, idx):
    method collate_fn (line 741) | def collate_fn(cls, batch):
  class MixDatasetWrapper (line 763) | class MixDatasetWrapper(torch.utils.data.Dataset):
    method __init__ (line 764) | def __init__(self, *datasets) -> None:
    method __len__ (line 774) | def __len__(self):
    method __getitem__ (line 777) | def __getitem__(self, idx):
    method _get_raw_item (line 785) | def _get_raw_item(self, idx):
  class DynamicBatchWrapper (line 795) | class DynamicBatchWrapper(torch.utils.data.Dataset):
    method __init__ (line 796) | def __init__(self, dataset, max_n_vertex_per_batch, max_n_vertex_per_i...
    method _form_batch (line 810) | def _form_batch(self):
    method __len__ (line 846) | def __len__(self):
    method __getitem__ (line 849) | def __getitem__(self, idx):
    method collate_fn (line 852) | def collate_fn(self, batched_batch):
  class BalancedDynamicBatchWrapper (line 858) | class BalancedDynamicBatchWrapper(DynamicBatchWrapper):
    method __init__ (line 859) | def __init__(self, dataset, max_n_vertex_per_batch, max_n_vertex_per_i...
    method _form_batch (line 866) | def _form_batch(self):
  class PretrainBalancedDynamicBatchWrapper (line 872) | class PretrainBalancedDynamicBatchWrapper(DynamicBatchWrapper):
    method __init__ (line 873) | def __init__(self, dataset, max_n_vertex_per_batch, max_n_vertex_per_i...
    method _form_batch (line 881) | def _form_batch(self):
  class DistillationDatasetWrapper (line 888) | class DistillationDatasetWrapper(torch.utils.data.Dataset):
    method __init__ (line 897) | def __init__(self, base_dataset, teacher_logits_file):
    method __len__ (line 920) | def __len__(self):
    method __getitem__ (line 923) | def __getitem__(self, idx):
    method collate_fn (line 943) | def collate_fn(cls, batch):
  class ResidueDistillationDatasetWrapper (line 969) | class ResidueDistillationDatasetWrapper(torch.utils.data.Dataset):
    method __init__ (line 982) | def __init__(self, base_dataset, teacher_logits_file):
    method __len__ (line 1010) | def __len__(self):
    method __getitem__ (line 1013) | def __getitem__(self, idx):
    method collate_fn (line 1054) | def collate_fn(cls, batch):
  class PocketEmbeddingDatasetWrapper (line 1081) | class PocketEmbeddingDatasetWrapper(torch.utils.data.Dataset):
    method __init__ (line 1093) | def __init__(self, base_dataset, embeddings_file):
    method __len__ (line 1123) | def __len__(self):
    method __getitem__ (line 1126) | def __getitem__(self, idx):
    method collate_fn (line 1136) | def collate_fn(cls, batch):
  function parse (line 1163) | def parse():
  function dataset_to_compressed_jsonl (line 1169) | def dataset_to_compressed_jsonl(dataset, output_file):
  function compressed_jsonl_to_dataset (line 1176) | def compressed_jsonl_to_dataset(input_file):
  function _maybe_convert_numpy_to_list (line 1187) | def _maybe_convert_numpy_to_list(item):
  function open_data_file (line 1197) | def open_data_file(data_file):

FILE: src/atomica/data/dataset_pretrain.py
  class PretrainMaskedDataset (line 11) | class PretrainMaskedDataset(torch.utils.data.Dataset):
    method __init__ (line 12) | def __init__(self, data_file, mask_proportion, mask_token, atom_mask_t...
    method set_crop (line 24) | def set_crop(self, max_n_vertex_per_item, fragmentation_method):
    method preprocess (line 33) | def preprocess(self):
    method __len__ (line 48) | def __len__(self):
    method _get_raw_item (line 51) | def _get_raw_item(self, idx):
    method __getitem__ (line 55) | def __getitem__(self, idx):
    method collate_fn (line 114) | def collate_fn(cls, batch):
  class PretrainTorsionDataset (line 140) | class PretrainTorsionDataset(torch.utils.data.Dataset):
    method __init__ (line 142) | def __init__(self, data_file):
    method preprocess (line 150) | def preprocess(self):
    method _can_apply_torsion_noise (line 162) | def _can_apply_torsion_noise(cls, data, chosen_segment):
    method set_torsion_noise (line 165) | def set_torsion_noise(self, noise_level):
    method set_translation_noise (line 168) | def set_translation_noise(self, noise_level):
    method set_rotation_noise (line 171) | def set_rotation_noise(self, noise_level, max_theta):
    method set_crop (line 174) | def set_crop(self, max_n_vertex_per_item, fragmentation_method):
    method __len__ (line 184) | def __len__(self):
    method _get_raw_item (line 187) | def _get_raw_item(self, idx):
    method __getitem__ (line 191) | def __getitem__(self, idx):
    method collate_fn (line 256) | def collate_fn(cls, batch):
  class PretrainMaskedTorsionDataset (line 335) | class PretrainMaskedTorsionDataset(PretrainTorsionDataset):
    method __init__ (line 336) | def __init__(self, data_file, mask_proportion, mask_token, atom_mask_t...
    method preprocess (line 348) | def preprocess(self):
    method get_mask_for_item (line 360) | def get_mask_for_item(self, data):
    method __getitem__ (line 368) | def __getitem__(self, idx):
    method collate_fn (line 439) | def collate_fn(cls, batch):
  class PretrainAtomDataset (line 450) | class PretrainAtomDataset(torch.utils.data.Dataset):
    method __init__ (line 452) | def __init__(self, data_file):
    method set_atom_noise (line 458) | def set_atom_noise(self, noise_level):
    method set_translation_noise (line 461) | def set_translation_noise(self, noise_level):
    method set_rotation_noise (line 464) | def set_rotation_noise(self, noise_level, max_theta):
    method set_crop (line 467) | def set_crop(self, max_n_vertex_per_item, fragmentation_method):
    method __len__ (line 475) | def __len__(self):
    method _get_raw_item (line 478) | def _get_raw_item(self, idx):
    method __getitem__ (line 482) | def __getitem__(self, idx):
    method collate_fn (line 547) | def collate_fn(cls, batch):
  class NoisyNodesTorsionDataset (line 570) | class NoisyNodesTorsionDataset(PretrainTorsionDataset):
    method __init__ (line 571) | def __init__(self, data_file):
    method __getitem__ (line 574) | def __getitem__(self, idx):

FILE: src/atomica/data/distributed_sampler.py
  class DistributedSamplerResume (line 16) | class DistributedSamplerResume(Sampler[T_co]):
    method __init__ (line 71) | def __init__(self, dataset: Dataset, num_replicas: Optional[int] = None,
    method __iter__ (line 106) | def __iter__(self) -> Iterator[T_co]:
    method __len__ (line 136) | def __len__(self) -> int:
    method set_epoch (line 139) | def set_epoch(self, epoch: int, resume_index: Optional[int] = None) ->...

FILE: src/atomica/data/get_torsion_mask.py
  function process_one (line 14) | def process_one(item):
  function main (line 53) | def main(args):
  function parse (line 64) | def parse():

FILE: src/atomica/data/mmap_dataset.py
  function compress (line 13) | def compress(x):
  function decompress (line 22) | def decompress(compressed_x):
  function create_mmap (line 30) | def create_mmap(iterator, out_dir, total_len=None, commit_batch=10000):
  class MMAPDataset (line 55) | class MMAPDataset(torch.utils.data.Dataset):
    method __init__ (line 57) | def __init__(self, mmap_dir: str) -> None:
    method __del__ (line 72) | def __del__(self):
    method __len__ (line 76) | def __len__(self):
    method __getitem__ (line 79) | def __getitem__(self, idx: int):
    method collate_fn (line 91) | def collate_fn(cls, batch):

FILE: src/atomica/data/pdb_utils.py
  function format_atom_element (line 70) | def format_atom_element(atom: str):
  class Vocab (line 77) | class Vocab:
    method __init__ (line 79) | def __init__(self):
    method load_tokenizer (line 143) | def load_tokenizer(self, method: Optional[str]):
    method abrv_to_symbol (line 149) | def abrv_to_symbol(self, abrv):
    method symbol_to_abrv (line 153) | def symbol_to_abrv(self, symbol):
    method abrv_to_idx (line 157) | def abrv_to_idx(self, abrv):
    method symbol_to_idx (line 160) | def symbol_to_idx(self, symbol):
    method idx_to_symbol (line 163) | def idx_to_symbol(self, idx):
    method idx_to_abrv (line 166) | def idx_to_abrv(self, idx):
    method get_pad_idx (line 169) | def get_pad_idx(self):
    method get_mask_idx (line 172) | def get_mask_idx(self):
    method get_special_mask (line 175) | def get_special_mask(self):
    method get_atom_pad_idx (line 178) | def get_atom_pad_idx(self):
    method get_atom_mask_idx (line 181) | def get_atom_mask_idx(self):
    method get_atom_global_idx (line 184) | def get_atom_global_idx(self):
    method get_atom_pos_pad_idx (line 187) | def get_atom_pos_pad_idx(self):
    method get_atom_pos_mask_idx (line 190) | def get_atom_pos_mask_idx(self):
    method get_atom_pos_global_idx (line 193) | def get_atom_pos_global_idx(self):
    method idx_to_atom (line 196) | def idx_to_atom(self, idx):
    method atom_to_idx (line 199) | def atom_to_idx(self, atom):
    method idx_to_atom_pos (line 202) | def idx_to_atom_pos(self, idx):
    method atom_pos_to_idx (line 205) | def atom_pos_to_idx(self, atom_pos):
    method get_num_atom_type (line 208) | def get_num_atom_type(self):
    method get_num_atom_pos (line 211) | def get_num_atom_pos(self):
    method get_num_amino_acid_type (line 214) | def get_num_amino_acid_type(self):
    method __len__ (line 217) | def __len__(self):
  function format_aa_abrv (line 224) | def format_aa_abrv(abrv):  # special cases
  class Atom (line 230) | class Atom:
    method __init__ (line 231) | def __init__(self, atom_name: str, coordinate: List, element: str, pos...
    method get_element (line 242) | def get_element(self):
    method get_coord (line 245) | def get_coord(self):
    method get_pos_code (line 248) | def get_pos_code(self):
    method __str__ (line 251) | def __str__(self) -> str:
  class Residue (line 255) | class Residue:
    method __init__ (line 256) | def __init__(self, symbol: str, atom_map: Dict, _id: Tuple, sidechain:...
    method get_symbol (line 267) | def get_symbol(self):
    method get_coord (line 270) | def get_coord(self, atom_name):
    method get_coord_map (line 273) | def get_coord_map(self) -> Dict[str, List]:
    method get_backbone_coord_map (line 276) | def get_backbone_coord_map(self) -> Dict[str, List]:
    method get_sidechain_coord_map (line 279) | def get_sidechain_coord_map(self) -> Dict[str, List]:
    method get_atom_names (line 286) | def get_atom_names(self):
    method get_atom (line 289) | def get_atom(self, atom_name):
    method get_id (line 292) | def get_id(self):
    method has_atom (line 295) | def has_atom(self, atom_name):
    method set_symbol (line 298) | def set_symbol(self, symbol):
    method set_atom_map (line 302) | def set_atom_map(self, atom_map):
    method dist_to (line 305) | def dist_to(self, residue):  # measured by nearest atoms
    method to_bio (line 313) | def to_bio(self):
    method __iter__ (line 336) | def __iter__(self):
    method __len__ (line 339) | def __len__(self):
  class Peptide (line 343) | class Peptide:
    method __init__ (line 344) | def __init__(self, _id, residues: List[Residue]):
    method set_id (line 351) | def set_id(self, _id):
    method get_id (line 354) | def get_id(self):
    method get_seq (line 357) | def get_seq(self):
    method get_span (line 360) | def get_span(self, i, j):  # [i, j)
    method get_residue (line 368) | def get_residue(self, i):
    method get_ca_pos (line 371) | def get_ca_pos(self, i):
    method get_cb_pos (line 374) | def get_cb_pos(self, i):
    method set_residue_symbol (line 377) | def set_residue_symbol(self, i, symbol):
    method set_residue (line 381) | def set_residue(self, i, symbol, coord):
    method to_bio (line 385) | def to_bio(self):
    method __iter__ (line 391) | def __iter__(self):
    method __len__ (line 394) | def __len__(self):
    method __str__ (line 397) | def __str__(self):
  class Protein (line 401) | class Protein:
    method __init__ (line 402) | def __init__(self, pdb_id, peptides):
    method from_pdb (line 407) | def from_pdb(cls, pdb_path, include_all=False):
    method get_id (line 466) | def get_id(self):
    method num_chains (line 469) | def num_chains(self):
    method get_chain (line 472) | def get_chain(self, name):
    method get_chain_names (line 478) | def get_chain_names(self):
    method to_bio (line 481) | def to_bio(self):
    method to_pdb (line 489) | def to_pdb(self, path, atoms=None):
    method __iter__ (line 507) | def __iter__(self):
    method __eq__ (line 510) | def __eq__(self, other):
    method __str__ (line 520) | def __str__(self):
  class Complex (line 527) | class Complex(Protein):
    method __init__ (line 528) | def __init__(self, pdb_id, peptides, receptor_chains=None, ligand_chai...
    method from_pdb (line 546) | def from_pdb(cls, pdb_path, receptor_chains=None, ligand_chains=None, ...
    method get_interacting_residues (line 550) | def get_interacting_residues(self, dist_th=10):
    method __str__ (line 572) | def __str__(self):
  function atom_features_from_residues (line 589) | def atom_features_from_residues(residue_list):
  function coords_from_residues (line 611) | def coords_from_residues(residue_list):
  function dist_matrix_from_coords (line 638) | def dist_matrix_from_coords(coords1, masks1, coords2, masks2):
  function dist_matrix_from_residues (line 645) | def dist_matrix_from_residues(residue_list1, residue_list2):

FILE: src/atomica/data/process_PeSTo_results.py
  function parse_args (line 14) | def parse_args():
  function process_one (line 24) | def process_one(atom_array: bs.AtomArray):

FILE: src/atomica/data/process_QBioLiP_parallel.py
  function pmap_multi (line 18) | def pmap_multi(pickleable_fn, data, n_jobs=None, verbose=1, desc=None, *...
  function parse (line 54) | def parse():
  function residue_to_pd_rows (line 78) | def residue_to_pd_rows(chain: str, residue: Residue):
  function process_one_PP (line 94) | def process_one_PP(protein_file_name, data_dir_rec, data_dir_lig, interf...
  function process_one_complex (line 135) | def process_one_complex(complex_file_name, data_dir_rec, data_dir_lig, i...
  function filter_PP_indexes (line 218) | def filter_PP_indexes(args, start, end):
  function filter_complex_indexes (line 242) | def filter_complex_indexes(args, start, end):
  function process_shard (line 286) | def process_shard(params):
  function main (line 323) | def main(args):

FILE: src/atomica/data/process_csd.py
  function process_crystals (line 29) | def process_crystals(
  function get_search_settings (line 79) | def get_search_settings(search_settings_param):
  function get_crystal_features_pool (line 90) | def get_crystal_features_pool(params):
  function save_results (line 111) | def save_results(results, csd_data_dir, processed_dir, worker_id):
  function blocks_from_molecule (line 119) | def blocks_from_molecule(molecule):
  function process_one (line 142) | def process_one(data_idx, molecule1, molecule2, interface_dist_th):
  function process_crystal_entry (line 165) | def process_crystal_entry(crystal, intermolecular_cutoff):
  function check_valid_molecule (line 211) | def check_valid_molecule(molecule):
  function get_crystal_from_identifier (line 218) | def get_crystal_from_identifier(csd_data_directory, identifier):

FILE: src/atomica/data/process_pdbs.py
  function parse_args (line 15) | def parse_args():
  function process_PL_pdb (line 33) | def process_PL_pdb(pdb_file, pdb_id, rec_chain, lig_code, lig_chain, smi...
  function group_chains (line 59) | def group_chains(list_chain_blocks, list_chain_pdb_indexes, group1, grou...
  function process_pdb (line 73) | def process_pdb(pdb_file, pdb_id, group1_chains, group2_chains, dist_th):
  function main (line 90) | def main(args):

FILE: src/atomica/data/tokenize_CSD.py
  function pmap_multi (line 10) | def pmap_multi(pickleable_fn, data, n_jobs=None, verbose=1, desc=None, *...
  function process_one (line 46) | def process_one(item, fragmentation_method):
  function main (line 62) | def main(args):
  function parse (line 79) | def parse():

FILE: src/atomica/data/tokenizer/chem_utils.py
  function smi2mol (line 8) | def smi2mol(smiles: str, kekulize=False, sanitize=True):
  function mol2smi (line 16) | def mol2smi(mol, canonical=True):
  function get_submol (line 20) | def get_submol(mol, atom_indices, kekulize=False):
  function get_submol_atom_map (line 38) | def get_submol_atom_map(mol, submol, group, kekulize=False):
  function cnt_atom (line 69) | def cnt_atom(smi, return_dict=False):

FILE: src/atomica/data/tokenizer/mol_atom_match.py
  function _mol_to_topology (line 109) | def _mol_to_topology(mol: Union[RDKitMol, str], include_Hs: bool=False):
  function get_atom_map (line 135) | def get_atom_map(
  function struct_to_topology (line 148) | def struct_to_topology(
  function struct_to_bonds (line 173) | def struct_to_bonds(

FILE: src/atomica/data/tokenizer/mol_bpe.py
  class MolInSubgraph (line 16) | class MolInSubgraph:
    method __init__ (line 17) | def __init__(self, mol, kekulize=False):
    method get_nei_subgraphs (line 36) | def get_nei_subgraphs(self):
    method get_nei_smis (line 56) | def get_nei_smis(self):
    method merge (line 71) | def merge(self, smi):
    method get_smis_subgraphs (line 93) | def get_smis_subgraphs(self):
  function freq_cnt (line 104) | def freq_cnt(mol):
  function graph_bpe (line 113) | def graph_bpe(fname, vocab_len, vocab_path, cpus, kekulize):
  class Tokenizer (line 174) | class Tokenizer:
    method __init__ (line 175) | def __init__(self, vocab_path):
    method tokenize (line 195) | def tokenize(self, mol):
    method idx_to_subgraph (line 231) | def idx_to_subgraph(self, idx):
    method subgraph_to_idx (line 234) | def subgraph_to_idx(self, subgraph):
    method atom_vocab (line 237) | def atom_vocab(self):
    method num_subgraph_type (line 240) | def num_subgraph_type(self):
    method atom_pos_pad_idx (line 243) | def atom_pos_pad_idx(self):
    method atom_pos_start_idx (line 246) | def atom_pos_start_idx(self):
    method __call__ (line 249) | def __call__(self, mol):
    method __len__ (line 252) | def __len__(self):
  function parse (line 255) | def parse():

FILE: src/atomica/data/tokenizer/molecule.py
  class SubgraphNode (line 16) | class SubgraphNode:
    method __init__ (line 20) | def __init__(self, smiles: str, pos: int, atom_mapping: dict, kekulize...
    method get_mol (line 27) | def get_mol(self):
    method get_atom_mapping (line 31) | def get_atom_mapping(self):
    method __str__ (line 34) | def __str__(self):
  class SubgraphEdge (line 42) | class SubgraphEdge:
    method __init__ (line 46) | def __init__(self, src: int, dst: int, edges: list):
    method get_edges (line 54) | def get_edges(self):
    method get_num_edges (line 57) | def get_num_edges(self):
    method __str__ (line 60) | def __str__(self):
  class Molecule (line 67) | class Molecule(nx.Graph):
    method __init__ (line 70) | def __init__(self, mol: Union[str, RDKitMol]=None, groups: list=None, ...
    method from_nx_graph (line 115) | def from_nx_graph(cls, graph: nx.Graph, deepcopy=True):
    method merge (line 122) | def merge(cls, mol0, mol1, edge=None):
    method get_edge (line 145) | def get_edge(self, i, j) -> SubgraphEdge:
    method get_node (line 148) | def get_node(self, i) -> SubgraphNode:
    method add_edge (line 151) | def add_edge(self, edge: SubgraphEdge) -> None:
    method add_node (line 155) | def add_node(self, node: SubgraphNode) -> None:
    method subgraph (line 159) | def subgraph(self, nodes: list):
    method to_rdkit_mol (line 164) | def to_rdkit_mol(self):
    method to_SVG (line 210) | def to_SVG(self, path: str, size: tuple=(200, 200), add_idx=False) -> ...
    method to_smiles (line 251) | def to_smiles(self):
    method __str__ (line 255) | def __str__(self):

FILE: src/atomica/data/tokenizer/singleton.py
  function singleton (line 3) | def singleton(cls):

FILE: src/atomica/data/tokenizer/tokenize_3d.py
  class TokenizerWrapper (line 22) | class TokenizerWrapper:
    method __init__ (line 24) | def __init__(self, method=None):
    method load (line 27) | def load(self, method: Optional[str]):
    method __call__ (line 38) | def __call__(self, mol):
    method get_frag_smiles (line 41) | def get_frag_smiles(self):
    method __len__ (line 48) | def __len__(self):
  function clean_bonds (line 57) | def clean_bonds(atoms: List[str], bonds: List[Tuple[int, int, int]]):
  function format_atom (line 120) | def format_atom(atom):
  function tokenize_3d (line 126) | def tokenize_3d(

FILE: src/atomica/data/torsion.py
  function get_torsion_mask (line 8) | def get_torsion_mask(atoms, coords):
  function modify_conformer_torsion_angles (line 55) | def modify_conformer_torsion_angles(coords, rotateable_edges, mask_rotat...
  function get_backbone_mask (line 82) | def get_backbone_mask(atoms, atom_positions, is_nucleotide: bool = False):
  function get_segment_torsion_mask (line 103) | def get_segment_torsion_mask(blocks):
  function get_side_chain_torsion_mask (line 125) | def get_side_chain_torsion_mask(blocks):
  function get_side_chain_torsion_mask_block (line 203) | def get_side_chain_torsion_mask_block(blocks):

FILE: src/atomica/get_embeddings.py
  function main (line 12) | def main(args):
  function parse_args (line 169) | def parse_args():
  function cli (line 181) | def cli():

FILE: src/atomica/interaction_profiler/detection.py
  class SegmentType (line 19) | class SegmentType(Enum):
  class Mol (line 23) | class Mol:
    method __init__ (line 24) | def __init__(self, mapper, segment_id):
    method hydrophobic_atoms (line 35) | def hydrophobic_atoms(self, all_atoms):
    method find_hba (line 48) | def find_hba(self, all_atoms):
    method find_hbd (line 60) | def find_hbd(self, all_atoms, hydroph_atoms):
    method find_rings (line 81) | def find_rings(self, mol, all_atoms):
    method get_hydrophobic_atoms (line 120) | def get_hydrophobic_atoms(self):
    method get_hba (line 123) | def get_hba(self):
    method get_hbd (line 126) | def get_hbd(self):
    method get_weak_hbd (line 129) | def get_weak_hbd(self):
    method get_pos_charged (line 132) | def get_pos_charged(self):
    method get_neg_charged (line 135) | def get_neg_charged(self):
  class Ligand (line 139) | class Ligand(Mol):
    method __init__ (line 140) | def __init__(self, mapper, segment_id):
    method get_canonical_num (line 162) | def get_canonical_num(self, atomnum):
    method is_functional_group (line 167) | def is_functional_group(atom, group):
    method find_hal_acc (line 209) | def find_hal_acc(self, atoms):
    method find_hal_don (line 223) | def find_hal_don(self, atoms):
    method find_charged (line 237) | def find_charged(self, all_atoms):
    method find_metal_binding (line 317) | def find_metal_binding(self, lig_atoms): #, water_oxygens):
  class BindingSite (line 387) | class BindingSite(Mol):
    method __init__ (line 388) | def __init__(self, mapper, segment_id):
    method find_hal (line 400) | def find_hal(self, atoms):
    method find_charged (line 413) | def find_charged(self, mol):
    method find_metal_binding (line 463) | def find_metal_binding(self, mol):
  class DataMapper (line 500) | class DataMapper:
    method __init__ (line 503) | def __init__(self, data, pybel_mols):
    method mapid (line 530) | def mapid(self, segment_id, pybabel_idx):  # Mapping to original IDs
    method id_to_atom (line 533) | def id_to_atom(self, original_idx):
    method id_to_block (line 543) | def id_to_block(self, original_idx):
  function saltbridge (line 548) | def saltbridge(poscenter, negcenter, protispos, noprot):
  class InteractionProfile (line 573) | class InteractionProfile:
    method __init__ (line 574) | def __init__(self, item, pybel_mols, *segment_types: Tuple[SegmentType]):
    method __str__ (line 636) | def __str__(self):
    method _convert_data (line 644) | def _convert_data(out_dict):
    method summary (line 675) | def summary(self):

FILE: src/atomica/interaction_profiler/interact_score.py
  function mask_block (line 15) | def mask_block(data, block_idx):
  function get_residue_model_scores (line 30) | def get_residue_model_scores(model, data):
  function get_residue_model_score (line 40) | def get_residue_model_score(model, data, block_idx):
  function parse_args (line 51) | def parse_args():

FILE: src/atomica/interaction_profiler/preparation.py
  function generated_to_xyz (line 10) | def generated_to_xyz(data):
  function generated_to_sdf (line 19) | def generated_to_sdf(data):
  function visualize_mol_sdf (line 29) | def visualize_mol_sdf(node_type, coords, save_sdf_path=None):
  function visualize_data_sdf (line 39) | def visualize_data_sdf(data, save_dir=None):
  function addh_to_mol (line 62) | def addh_to_mol(input_sdf_path, output_sdf_path):
  function build_pybel_mols (line 97) | def build_pybel_mols(data, with_hydrogens=True, tmpfile_dir='./'):

FILE: src/atomica/models/__init__.py
  function create_model (line 8) | def create_model(args):

FILE: src/atomica/models/affinity_predictor.py
  class AffinityPredictor (line 12) | class AffinityPredictor(PredictionModel):
    method __init__ (line 14) | def __init__(self, num_affinity_pred_layers, nonlinearity, affinity_pr...
    method init_block_embedding (line 73) | def init_block_embedding(self, nonlinearity: nn.Module, block_embeddin...
    method _load_from_pretrained (line 86) | def _load_from_pretrained(cls, pretrained_model, **kwargs):
    method get_config (line 135) | def get_config(self):
    method forward (line 151) | def forward(self, Z, B, A, block_lengths, lengths, segment_ids, label,...
    method infer (line 222) | def infer(self, batch):

FILE: src/atomica/models/atomica/atomica.py
  class InteractionModule (line 9) | class InteractionModule(torch.nn.Module):
    method __init__ (line 10) | def __init__(
    method remove_torsion_denoiser (line 150) | def remove_torsion_denoiser(self):
    method forward (line 157) | def forward(self, node_attr, coords, batch_id, perturb_mask, edges, ed...
    method build_tor_edges (line 273) | def build_tor_edges(self, tor_bonds, coords, node_embeddings, batch_id...

FILE: src/atomica/models/atomica/encoder.py
  class ATOMICAEncoder (line 10) | class ATOMICAEncoder(nn.Module):
    method __init__ (line 11) | def __init__(self, hidden_size, edge_size, n_layers=3, return_atom_noi...
    method forward (line 21) | def forward(self, H, Z, batch_id, perturb_mask, edges, edge_attr, tor_...
  class AttentionPooling (line 34) | class AttentionPooling(nn.Module):
    method __init__ (line 35) | def __init__(self, hidden_size, num_heads=4, dropout=0.0, num_layers=4):
    method forward (line 43) | def forward(self, block_repr, batch_id):

FILE: src/atomica/models/atomica/utils.py
  class EquivariantLayerNorm (line 14) | class EquivariantLayerNorm(nn.Module):
    method __init__ (line 15) | def __init__(self, irreps, eps=1e-5, affine=True, normalization="compo...
    method __repr__ (line 38) | def __repr__(self):
    method forward (line 42) | def forward(self, node_input, **kwargs):
  class TensorProductConvLayer (line 118) | class TensorProductConvLayer(torch.nn.Module):
    method __init__ (line 119) | def __init__(
    method forward (line 162) | def forward(
  class GaussianEmbedding (line 200) | class GaussianEmbedding(torch.nn.Module):
    method __init__ (line 203) | def __init__(self, start=0.0, stop=5.0, num_gaussians=32):
    method forward (line 210) | def forward(self, dist):
  class AtomEncoder (line 217) | class AtomEncoder(nn.Module):
    method __init__ (line 219) | def __init__(self, dim, trainable=True):
    method forward (line 229) | def forward(self, elems):
  class SphericalHarmonicEdgeAttrs (line 236) | class SphericalHarmonicEdgeAttrs(nn.Module):
    method __init__ (line 237) | def __init__(
    method forward (line 249) | def forward(self, edge_vec):
  function batchify (line 253) | def batchify(tensor, batch_ids, max_seq_len=None):
  function unbatchify (line 273) | def unbatchify(batchified_tensor, batchify_mask):

FILE: src/atomica/models/classifier_model.py
  class ClassifierModel (line 7) | class ClassifierModel(PredictionModel):
    method __init__ (line 9) | def __init__(self, num_pred_layers, nonlinearity, pred_dropout, pred_h...
    method _load_from_pretrained (line 22) | def _load_from_pretrained(cls, pretrained_model, **kwargs):
    method get_config (line 65) | def get_config(self):
    method forward (line 75) | def forward(self, Z, B, A, block_lengths, lengths, segment_ids, label,...
    method infer (line 81) | def infer(self, batch, extra_info=False):
  class MultiClassClassifierModel (line 96) | class MultiClassClassifierModel(PredictionModel):
    method __init__ (line 98) | def __init__(self, num_classes, loss_type='cross_entropy', focal_alpha...
    method set_class_weights (line 131) | def set_class_weights(self, class_weights):
    method _load_from_pretrained (line 147) | def _load_from_pretrained(cls, pretrained_model, **kwargs):
    method get_config (line 211) | def get_config(self):
    method forward (line 220) | def forward(self, Z, B, A, block_lengths, lengths, segment_ids, label,...
    method infer (line 256) | def infer(self, batch, extra_info=False):
  class RegressionPredictor (line 286) | class RegressionPredictor(PredictionModel):
    method __init__ (line 288) | def __init__(self, **kwargs) -> None:
    method _load_from_pretrained (line 304) | def _load_from_pretrained(cls, pretrained_model, **kwargs):
    method forward (line 311) | def forward(self, Z, B, A, block_lengths, lengths, segment_ids, label,...
    method infer (line 316) | def infer(self, batch, extra_info=False):
  class ResidueClassifierModel (line 330) | class ResidueClassifierModel(PredictionModel):
    method __init__ (line 332) | def __init__(self, num_pred_layers, nonlinearity, pred_dropout, pred_h...
    method _load_from_pretrained (line 345) | def _load_from_pretrained(cls, pretrained_model, **kwargs):
    method get_config (line 388) | def get_config(self):
    method forward (line 398) | def forward(self, Z, B, A, block_lengths, lengths, segment_ids, label,...
    method infer (line 406) | def infer(self, batch, extra_info=False):
  class MultiLabelClassifierModel (line 422) | class MultiLabelClassifierModel(MultiClassClassifierModel):
    method __init__ (line 424) | def __init__(self, num_classes, loss_type='binary_cross_entropy', foca...
    method forward (line 432) | def forward(self, Z, B, A, block_lengths, lengths, segment_ids, label,...
    method infer (line 469) | def infer(self, batch, extra_info=False):

FILE: src/atomica/models/deprecated/DimeNet/dimenet.py
  function stable_norm (line 28) | def stable_norm(input, *args, **kwargs):
  function sph1 (line 47) | def sph1(x, theta, sph_harm_form, modules):
  function sph (line 51) | def sph(x, theta, sph_harm_form, modules):
  class SphericalBasisLayer (line 56) | class SphericalBasisLayer(torch.nn.Module):
    method __init__ (line 57) | def __init__(
    method forward (line 98) | def forward(self, dist: Tensor, angle: Tensor, idx_kj: Tensor) -> Tensor:
  class EmbeddingBlock (line 109) | class EmbeddingBlock(torch.nn.Module):
    method __init__ (line 110) | def __init__(self, num_radial: int, hidden_channels: int, act: Callable):
    method reset_parameters (line 119) | def reset_parameters(self):
    method forward (line 123) | def forward(self, x: Tensor, rbf: Tensor, i: Tensor, j: Tensor) -> Ten...
  class InteractionBlock (line 128) | class InteractionBlock(torch.nn.Module):
    method __init__ (line 129) | def __init__(
    method reset_parameters (line 163) | def reset_parameters(self):
    method forward (line 178) | def forward(self, x: Tensor, rbf: Tensor, sbf: Tensor, idx_kj: Tensor,
  class InteractionPPBlock (line 199) | class InteractionPPBlock(torch.nn.Module):
    method __init__ (line 200) | def __init__(
    method reset_parameters (line 241) | def reset_parameters(self):
    method forward (line 262) | def forward(self, x: Tensor, rbf: Tensor, sbf: Tensor, idx_kj: Tensor,
  class OutputBlock (line 295) | class OutputBlock(torch.nn.Module):
    method __init__ (line 296) | def __init__(
    method reset_parameters (line 315) | def reset_parameters(self):
    method forward (line 322) | def forward(self, x: Tensor, rbf: Tensor, i: Tensor,
  class OutputPPBlock (line 331) | class OutputPPBlock(torch.nn.Module):
    method __init__ (line 332) | def __init__(
    method reset_parameters (line 355) | def reset_parameters(self):
    method forward (line 363) | def forward(self, x: Tensor, rbf: Tensor, i: Tensor,
  function triplets (line 373) | def triplets(
  class DimeNet (line 399) | class DimeNet(torch.nn.Module):
    method __init__ (line 440) | def __init__(
    method reset_parameters (line 490) | def reset_parameters(self):
    method forward (line 499) | def forward(
  class DimeNetPlusPlus (line 546) | class DimeNetPlusPlus(DimeNet):
    method __init__ (line 583) | def __init__(

FILE: src/atomica/models/deprecated/DimeNet/encoder.py
  function stable_norm (line 10) | def stable_norm(input, *args, **kwargs):
  class DimeNetEncoder (line 13) | class DimeNetEncoder(nn.Module):
    method __init__ (line 14) | def __init__(self, hidden_size, n_layers=3) -> None:
    method forward (line 33) | def forward(self, H, Z, block_id, batch_id, edges, edge_attr=None):

FILE: src/atomica/models/deprecated/EGNN/egnn.py
  class E_GCL (line 12) | class E_GCL(nn.Module):
    method __init__ (line 18) | def __init__(self, input_nf, output_nf, hidden_nf, edges_in_d=0, act_f...
    method edge_model (line 56) | def edge_model(self, source, target, radial, edge_attr):
    method node_model (line 67) | def node_model(self, x, edge_index, edge_attr, node_attr):
    method coord_model (line 79) | def coord_model(self, coord, edge_index, coord_diff, edge_feat):
    method coord2radial (line 91) | def coord2radial(self, edge_index, coord):
    method forward (line 102) | def forward(self, h, edge_index, coord, edge_attr=None, node_attr=None):
  class EGNN (line 113) | class EGNN(nn.Module):
    method __init__ (line 114) | def __init__(self, in_node_nf, hidden_nf, out_node_nf, in_edge_nf=0, a...
    method forward (line 145) | def forward(self, h, x, edges, edge_attr):
  function unsorted_segment_sum (line 153) | def unsorted_segment_sum(data, segment_ids, num_segments):
  function unsorted_segment_mean (line 161) | def unsorted_segment_mean(data, segment_ids, num_segments):

FILE: src/atomica/models/deprecated/EGNN/encoder.py
  class EGNNEncoder (line 9) | class EGNNEncoder(nn.Module):
    method __init__ (line 10) | def __init__(self, hidden_size, edge_size, n_layers=3) -> None:
    method forward (line 20) | def forward(self, H, Z, block_id, batch_id, edges, edge_attr=None):

FILE: src/atomica/models/deprecated/GET/encoder.py
  class GETEncoder (line 13) | class GETEncoder(nn.Module):
    method __init__ (line 14) | def __init__(self, hidden_size, radial_size, n_channel,
    method forward (line 27) | def forward(self, H, Z, block_id, batch_id, edges, edge_attr=None):

FILE: src/atomica/models/deprecated/GET/model.py
  class GETEncoder (line 20) | class GETEncoder(nn.Module):
    method __init__ (line 21) | def __init__(self, hidden_size, n_channel, radial_size=16,
    method message_passing (line 50) | def message_passing(self, Z, B, A, atom_positions, block_lengths, leng...
    method forward (line 81) | def forward(self, Z, B, A, atom_positions, block_lengths, lengths, seg...
    method infer (line 85) | def infer(self, batch):

FILE: src/atomica/models/deprecated/GET/modules/get.py
  function stable_norm (line 15) | def stable_norm(input, *args, **kwargs):
  class GET (line 18) | class GET(nn.Module):
    method __init__ (line 21) | def __init__(self, d_hidden, d_radial, n_channel, n_rbf, cutoff=7.0, d...
    method recover_scale (line 68) | def recover_scale(self, Z, block_id, batch_id, record_scale):
    method forward (line 79) | def forward(self, H, Z, block_id, batch_id, edges, edge_attr=None, cac...
  class GETLayer (line 115) | class GETLayer(nn.Module):
    method __init__ (line 120) | def __init__(self, d_hidden, d_radial, n_channel, n_rbf, cutoff=7.0,
    method attention (line 203) | def attention(self, H, Z, edges, edge_attr, cached_unit_edge_info):
    method invariant_update (line 255) | def invariant_update(self, H_v, H, alpha, beta, D, cached_unit_edge_in...
    method equivariant_update (line 274) | def equivariant_update(self, H_v, Z, alpha, beta, D, dZ, cached_unit_e...
    method forward (line 297) | def forward(self, H, Z, block_id, edges, edge_attr=None, cached_unit_e...
  class EquivariantFFN (line 321) | class EquivariantFFN(nn.Module):
    method __init__ (line 322) | def __init__(self, d_in, d_hidden, d_out, n_channel, n_rbf=16, act_fn=...
    method forward (line 370) | def forward(self, H, Z, block_id):
    method _radial (line 404) | def _radial(self, Z, block_id):
  class EquivariantLayerNorm (line 421) | class EquivariantLayerNorm(nn.Module):
    method __init__ (line 423) | def __init__(self, d_hidden, n_channel, n_rbf=16, cutoff=7.0, act_fn=n...
    method forward (line 440) | def forward(self, H, Z, block_id, batch_id):

FILE: src/atomica/models/deprecated/GET/modules/radial_basis.py
  class GaussianSmearing (line 16) | class GaussianSmearing(torch.nn.Module):
    method __init__ (line 17) | def __init__(
    method forward (line 28) | def forward(self, dist: Tensor) -> Tensor:
  class PolynomialEnvelope (line 33) | class PolynomialEnvelope(torch.nn.Module):
    method __init__ (line 43) | def __init__(self, exponent):
    method forward (line 51) | def forward(self, d_scaled):
  class ExponentialEnvelope (line 61) | class ExponentialEnvelope(torch.nn.Module):
    method __init__ (line 69) | def __init__(self):
    method forward (line 72) | def forward(self, d_scaled):
  class SphericalBesselBasis (line 79) | class SphericalBesselBasis(torch.nn.Module):
    method __init__ (line 91) | def __init__(
    method forward (line 108) | def forward(self, d_scaled):
  class BernsteinBasis (line 116) | class BernsteinBasis(torch.nn.Module):
    method __init__ (line 133) | def __init__(
    method forward (line 157) | def forward(self, d_scaled):
  class RadialBasis (line 165) | class RadialBasis(torch.nn.Module):
    method __init__ (line 180) | def __init__(
    method forward (line 219) | def forward(self, d):

FILE: src/atomica/models/deprecated/GET/modules/tools.py
  function _unit_edges_from_block_edges (line 11) | def _unit_edges_from_block_edges(unit_block_id, block_src_dst, Z=None, k...
  function _block_edge_dist (line 96) | def _block_edge_dist(X, block_id, src_dst):
  class KNNBatchEdgeConstructor (line 115) | class KNNBatchEdgeConstructor(BatchEdgeConstructor):
    method __init__ (line 116) | def __init__(self, k_neighbors, global_message_passing=True, global_no...
    method _construct_intra_edges (line 121) | def _construct_intra_edges(self, S, batch_id, segment_ids, **kwargs):
    method _construct_inter_edges (line 132) | def _construct_inter_edges(self, S, batch_id, segment_ids, **kwargs):
    method _construct_global_edges (line 143) | def _construct_global_edges(self, S, batch_id, segment_ids, **kwargs):
    method _construct_seq_edges (line 149) | def _construct_seq_edges(self, S, batch_id, segment_ids, **kwargs):
  class BlockEmbedding (line 154) | class BlockEmbedding(nn.Module):
    method __init__ (line 158) | def __init__(self, num_block_type, num_atom_type, num_atom_position, e...
    method forward (line 166) | def forward(self, B, A, atom_positions, block_id):

FILE: src/atomica/models/deprecated/GET/pool_encoder.py
  class GETPoolEncoder (line 13) | class GETPoolEncoder(nn.Module):
    method __init__ (line 14) | def __init__(self, hidden_size, radial_size, n_channel,
    method forward (line 27) | def forward(self, H, Z, block_id, batch_id, edges, edge_attr=None):

FILE: src/atomica/models/deprecated/SchNet/encoder.py
  class SchNetEncoder (line 9) | class SchNetEncoder(nn.Module):
    method __init__ (line 10) | def __init__(self, hidden_size, edge_size, n_layers=3) -> None:
    method forward (line 22) | def forward(self, H, Z, block_id, batch_id, edges, edge_attr=None):

FILE: src/atomica/models/deprecated/SchNet/schnet.py
  class SchNet (line 40) | class SchNet(torch.nn.Module):
    method __init__ (line 98) | def __init__(
    method reset_parameters (line 150) | def reset_parameters(self):
    method forward (line 161) | def forward(self, h: Tensor, pos: Tensor,
    method __repr__ (line 185) | def __repr__(self) -> str:
  class RadiusInteractionGraph (line 194) | class RadiusInteractionGraph(torch.nn.Module):
    method __init__ (line 206) | def __init__(self, cutoff: float = 10.0, max_num_neighbors: int = 32):
    method forward (line 211) | def forward(self, pos: Tensor, batch: Tensor) -> Tuple[Tensor, Tensor]:
  class InteractionBlock (line 227) | class InteractionBlock(torch.nn.Module):
    method __init__ (line 228) | def __init__(self, hidden_channels: int, num_gaussians: int,
    method reset_parameters (line 243) | def reset_parameters(self):
    method forward (line 252) | def forward(self, x: Tensor, edge_index: Tensor, edge_weight: Tensor,
  class CFConv (line 260) | class CFConv(MessagePassing):
    method __init__ (line 261) | def __init__(
    method reset_parameters (line 277) | def reset_parameters(self):
    method forward (line 282) | def forward(self, x: Tensor, edge_index: Tensor, edge_weight: Tensor,
    method message (line 292) | def message(self, x_j: Tensor, W: Tensor) -> Tensor:
  class GaussianSmearing (line 296) | class GaussianSmearing(torch.nn.Module):
    method __init__ (line 297) | def __init__(
    method forward (line 308) | def forward(self, dist: Tensor) -> Tensor:
  class ShiftedSoftplus (line 313) | class ShiftedSoftplus(torch.nn.Module):
    method __init__ (line 314) | def __init__(self):
    method forward (line 318) | def forward(self, x: Tensor) -> Tensor:

FILE: src/atomica/models/deprecated/TorchMD/encoder.py
  function stable_norm (line 10) | def stable_norm(input, *args, **kwargs):
  class TorchMDEncoder (line 14) | class TorchMDEncoder(nn.Module):
    method __init__ (line 15) | def __init__(self, hidden_size, edge_size, n_layers=3) -> None:
    method forward (line 29) | def forward(self, H, Z, block_id, batch_id, edges, edge_attr=None):

FILE: src/atomica/models/deprecated/TorchMD/torchmd_et.py
  class TorchMD_ET (line 18) | class TorchMD_ET(nn.Module):
    method __init__ (line 57) | def __init__(
    method reset_parameters (line 142) | def reset_parameters(self):
    method forward (line 150) | def forward(
    method __repr__ (line 188) | def __repr__(self):
  class EquivariantMultiHeadAttention (line 206) | class EquivariantMultiHeadAttention(MessagePassing):
    method __init__ (line 207) | def __init__(
    method reset_parameters (line 252) | def reset_parameters(self):
    method forward (line 270) | def forward(self, x, vec, edge_index, r_ij, f_ij, d_ij):
    method message (line 312) | def message(self, q_i, k_j, v_j, vec_j, dk, dv, r_ij, d_ij):
    method aggregate (line 335) | def aggregate(
    method update (line 347) | def update(

FILE: src/atomica/models/deprecated/TorchMD/utils.py
  function visualize_basis (line 13) | def visualize_basis(basis_type, num_rbf=50, cutoff_lower=0, cutoff_upper...
  class NeighborEmbedding (line 44) | class NeighborEmbedding(MessagePassing):
    method __init__ (line 45) | def __init__(self, hidden_channels, num_rbf, cutoff_lower, cutoff_uppe...
    method reset_parameters (line 54) | def reset_parameters(self):
    method forward (line 61) | def forward(self, z, x, edge_index, edge_weight, edge_attr):
    method message (line 78) | def message(self, x_j, W):
  class GaussianSmearing (line 82) | class GaussianSmearing(nn.Module):
    method __init__ (line 83) | def __init__(self, cutoff_lower=0.0, cutoff_upper=5.0, num_rbf=50, tra...
    method _initial_params (line 98) | def _initial_params(self):
    method reset_parameters (line 103) | def reset_parameters(self):
    method forward (line 108) | def forward(self, dist):
  class ExpNormalSmearing (line 113) | class ExpNormalSmearing(nn.Module):
    method __init__ (line 114) | def __init__(self, cutoff_lower=0.0, cutoff_upper=5.0, num_rbf=50, tra...
    method _initial_params (line 132) | def _initial_params(self):
    method reset_parameters (line 144) | def reset_parameters(self):
    method forward (line 149) | def forward(self, dist):
  class ShiftedSoftplus (line 157) | class ShiftedSoftplus(nn.Module):
    method __init__ (line 158) | def __init__(self):
    method forward (line 162) | def forward(self, x):
  class CosineCutoff (line 166) | class CosineCutoff(nn.Module):
    method __init__ (line 167) | def __init__(self, cutoff_lower=0.0, cutoff_upper=5.0):
    method forward (line 172) | def forward(self, distances):
  class Distance (line 197) | class Distance(nn.Module):
    method __init__ (line 198) | def __init__(
    method forward (line 213) | def forward(self, pos, batch):
  class GatedEquivariantBlock (line 258) | class GatedEquivariantBlock(nn.Module):
    method __init__ (line 263) | def __init__(
    method reset_parameters (line 289) | def reset_parameters(self):
    method forward (line 297) | def forward(self, x, v):

FILE: src/atomica/models/deprecated/_binary_predictors.py
  class BinaryPredictor (line 16) | class BinaryPredictor(nn.Module):
    method __init__ (line 18) | def __init__(self, atom_hidden_size, block_hidden_size, edge_size, k_n...
    method load_from_pretrained (line 80) | def load_from_pretrained(cls, pretrain_ckpt, **kwargs):
    method forward_one_encoder (line 124) | def forward_one_encoder(self, Z, B, A, block_lengths, lengths, segment...
    method forward (line 139) | def forward(self, Z0, B0, A0, block_lengths0, lengths0, segment_ids0,
    method infer (line 179) | def infer(self, batch):
  class BinaryPredictorMSP (line 197) | class BinaryPredictorMSP(nn.Module):
    method __init__ (line 199) | def __init__(self, atom_hidden_size, block_hidden_size, edge_size, k_n...
    method load_from_pretrained (line 269) | def load_from_pretrained(cls, pretrain_ckpt, **kwargs):
    method forward_one_encoder_bottom (line 316) | def forward_one_encoder_bottom(self, Z, B, A, atom_repr, block_lengths...
    method forward_one_encoder_bottom_to_top (line 349) | def forward_one_encoder_bottom_to_top(self, Z, B, A, atom_repr, block_...
    method forward_one_encoder_top (line 386) | def forward_one_encoder_top(self, top_H_0, top_Z, batch_id, edges, edg...
    method forward (line 396) | def forward(self, Z0, B0, A0, block_lengths0, lengths0, segment_ids0, ...
    method infer (line 519) | def infer(self, batch):
  class BinaryPredictorMSP2 (line 539) | class BinaryPredictorMSP2(nn.Module):
    method __init__ (line 541) | def __init__(self, atom_hidden_size, block_hidden_size, edge_size, k_n...
    method load_from_pretrained (line 628) | def load_from_pretrained(cls, pretrain_ckpt, **kwargs):
    method forward_one_encoder_bottom (line 669) | def forward_one_encoder_bottom(self, Z, B, A, atom_repr, block_lengths...
    method forward_one_encoder_bottom_to_top (line 697) | def forward_one_encoder_bottom_to_top(self, Z, B, A, atom_repr, block_...
    method forward_one_encoder_top (line 729) | def forward_one_encoder_top(self, top_H_0, top_Z, batch_id, edges, edg...
    method forward (line 734) | def forward(self, Z0, B0, A0, block_lengths0, lengths0, segment_ids0, ...
    method infer (line 803) | def infer(self, batch):

FILE: src/atomica/models/deprecated/_ddG_predictor.py
  class DDGPredictor (line 12) | class DDGPredictor(PredictionModel):
    method __init__ (line 14) | def __init__(self, **kwargs) -> None:
    method load_from_pretrained (line 66) | def load_from_pretrained(cls, pretrain_ckpt, **kwargs):
    method get_pred (line 73) | def get_pred(self, B, top_Z, esm_embeddings, lengths, segment_ids, mt_...
    method forward (line 120) | def forward(self, data, ddg, mt_block_indexes) -> PredictionReturnValue:
    method infer (line 130) | def infer(self, batch):
  class GLOFPredictor (line 143) | class GLOFPredictor(PredictionModel):
    method __init__ (line 145) | def __init__(self, **kwargs) -> None:
    method load_from_pretrained (line 174) | def load_from_pretrained(cls, pretrain_ckpt, **kwargs):
    method get_pred (line 181) | def get_pred(self, B, top_Z, lengths, segment_ids, mt_block_indexes):
    method forward (line 213) | def forward(self, data, labels, mt_block_indexes) -> PredictionReturnV...
    method infer (line 222) | def infer(self, batch):

FILE: src/atomica/models/deprecated/get.py
  function stable_norm (line 15) | def stable_norm(input, *args, **kwargs):
  class GaussianSmearing (line 19) | class GaussianSmearing(torch.nn.Module):
    method __init__ (line 20) | def __init__(
    method forward (line 31) | def forward(self, dist: Tensor) -> Tensor:
  class PolynomialEnvelope (line 36) | class PolynomialEnvelope(torch.nn.Module):
    method __init__ (line 46) | def __init__(self, exponent):
    method forward (line 54) | def forward(self, d_scaled):
  class ExponentialEnvelope (line 64) | class ExponentialEnvelope(torch.nn.Module):
    method __init__ (line 72) | def __init__(self):
    method forward (line 75) | def forward(self, d_scaled):
  class SphericalBesselBasis (line 82) | class SphericalBesselBasis(torch.nn.Module):
    method __init__ (line 94) | def __init__(
    method forward (line 111) | def forward(self, d_scaled):
  class BernsteinBasis (line 119) | class BernsteinBasis(torch.nn.Module):
    method __init__ (line 136) | def __init__(
    method forward (line 160) | def forward(self, d_scaled):
  class RadialBasis (line 168) | class RadialBasis(torch.nn.Module):
    method __init__ (line 183) | def __init__(
    method forward (line 222) | def forward(self, d):
  class GET (line 231) | class GET(nn.Module):
    method __init__ (line 234) | def __init__(self, d_hidden, d_radial, n_channel, n_rbf, cutoff=7.0, d...
    method recover_scale (line 281) | def recover_scale(self, Z, block_id, batch_id, record_scale):
    method forward (line 292) | def forward(self, H, Z, block_id, batch_id, edges, edge_attr=None, cac...
  class GETLayer (line 328) | class GETLayer(nn.Module):
    method __init__ (line 333) | def __init__(self, d_hidden, d_radial, n_channel, n_rbf, cutoff=7.0,
    method attention (line 416) | def attention(self, H, Z, edges, edge_attr, cached_unit_edge_info):
    method invariant_update (line 468) | def invariant_update(self, H_v, H, alpha, beta, D, cached_unit_edge_in...
    method equivariant_update (line 487) | def equivariant_update(self, H_v, Z, alpha, beta, D, dZ, cached_unit_e...
    method forward (line 510) | def forward(self, H, Z, block_id, edges, edge_attr=None, cached_unit_e...
  class EquivariantFFN (line 534) | class EquivariantFFN(nn.Module):
    method __init__ (line 535) | def __init__(self, d_in, d_hidden, d_out, n_channel, n_rbf=16, act_fn=...
    method forward (line 583) | def forward(self, H, Z, block_id):
    method _radial (line 617) | def _radial(self, Z, block_id):
  class EquivariantLayerNorm (line 634) | class EquivariantLayerNorm(nn.Module):
    method __init__ (line 636) | def __init__(self, d_hidden, n_channel, n_rbf=16, cutoff=7.0, act_fn=n...
    method forward (line 653) | def forward(self, H, Z, block_id, batch_id):
  class GETEncoder (line 676) | class GETEncoder(nn.Module):
    method __init__ (line 677) | def __init__(self, hidden_size, radial_size, n_channel,
    method forward (line 690) | def forward(self, H, Z, block_id, batch_id, edges, edge_attr=None):

FILE: src/atomica/models/deprecated/graph_classifier.py
  class GraphClassifier (line 9) | class GraphClassifier(PredictionModel):
    method __init__ (line 11) | def __init__(self, model_type, hidden_size, n_channel, num_class, **kw...
    method forward (line 18) | def forward(self, Z, B, A, atom_positions, block_lengths, lengths, seg...
    method infer (line 24) | def infer(self, batch):

FILE: src/atomica/models/deprecated/graph_multi_binary_classifier.py
  class GraphMultiBinaryClassifier (line 10) | class GraphMultiBinaryClassifier(PredictionModel):
    method __init__ (line 11) | def __init__(self, n_task, model_type, hidden_size, n_channel, **kwarg...
    method forward (line 21) | def forward(self, Z, B, A, atom_positions, block_lengths, lengths, seg...

FILE: src/atomica/models/deprecated/graph_pair_classifier.py
  class GraphPairClassifier (line 10) | class GraphPairClassifier(PredictionModel):
    method __init__ (line 12) | def __init__(self, model_type, hidden_size, n_channel, num_class, **kw...
    method forward (line 19) | def forward(self, Z, B, A, atom_positions, block_lengths, lengths, seg...
    method infer (line 26) | def infer(self, batch):

FILE: src/atomica/models/masking_model.py
  class MaskedNodeModel (line 11) | class MaskedNodeModel(DenoisePretrainModel):
    method __init__ (line 13) | def __init__(self, atom_hidden_size, block_hidden_size, edge_size, k_n...
    method _load_from_pretrained (line 25) | def _load_from_pretrained(cls, pretrained_model, **kwargs):
    method get_config (line 56) | def get_config(self):
    method load_from_pretrained (line 72) | def load_from_pretrained(cls, pretrain_ckpt, **kwargs):
    method load_from_config_and_weights (line 77) | def load_from_config_and_weights(cls, config_path, weights_path, **kwa...
    method forward (line 93) | def forward(self, Z, B, A, block_lengths, lengths, segment_ids, masked...
    method infer (line 146) | def infer(self, batch):

FILE: src/atomica/models/prediction_model.py
  class PredictionModel (line 15) | class PredictionModel(DenoisePretrainModel):
    method __init__ (line 16) | def __init__(self, atom_hidden_size, block_hidden_size, edge_size, k_n...
    method _load_from_pretrained (line 27) | def _load_from_pretrained(cls, pretrained_model: DenoisePretrainModel,...
    method get_config (line 62) | def get_config(self):
    method load_from_pretrained (line 77) | def load_from_pretrained(cls, pretrain_ckpt, **kwargs):
    method load_from_config_and_weights (line 82) | def load_from_config_and_weights(cls, config_path, weights_path, **kwa...
    method forward (line 119) | def forward(self, Z, B, A, block_lengths, lengths, segment_ids, return...
    method infer (line 185) | def infer(self, batch):

FILE: src/atomica/models/pretrain_model.py
  function construct_edges (line 28) | def construct_edges(edge_constructor, B, batch_id, segment_ids, X, block...
  class DenoisePretrainModel (line 81) | class DenoisePretrainModel(nn.Module):
    method __init__ (line 83) | def __init__(self, atom_hidden_size, block_hidden_size, edge_size=16, ...
    method get_config (line 199) | def get_config(self):
    method load_from_config_and_weights (line 225) | def load_from_config_and_weights(cls, config_path, weights_path):
    method get_edges (line 235) | def get_edges(self, B, batch_id, segment_ids, Z, block_id, global_mess...
    method forward (line 257) | def forward(self, Z, B, A, block_lengths, lengths, segment_ids,
  class DenoisePretrainModelWithBlockEmbedding (line 416) | class DenoisePretrainModelWithBlockEmbedding(DenoisePretrainModel):
    method __init__ (line 417) | def __init__(self, atom_hidden_size, block_hidden_size, edge_size=16, ...
    method get_config (line 456) | def get_config(self):
    method load_from_config_and_weights (line 469) | def load_from_config_and_weights(cls, config_path, weights_path):
    method init_block_embedding (line 478) | def init_block_embedding(self, nonlinearity: nn.Module, block_embeddin...
    method forward (line 491) | def forward(self, Z, B, A, block_lengths, lengths, segment_ids,

FILE: src/atomica/models/prot_interface_model.py
  class ProteinInterfaceModel (line 9) | class ProteinInterfaceModel(nn.Module):
    method __init__ (line 10) | def __init__(self, model: PredictionModel) -> None:
    method _load_from_pretrained (line 34) | def _load_from_pretrained(cls, pretrained_model, **kwargs):
    method get_config (line 70) | def get_config(self):
    method load_from_pretrained (line 77) | def load_from_pretrained(cls, pretrain_ckpt, **kwargs):
    method load_from_config_and_weights (line 84) | def load_from_config_and_weights(cls, config_path, weights_path, **kwa...
    method forward (line 106) | def forward(self, batch_cmplx, batch_prot) -> PredictionReturnValue:
    method infer (line 124) | def infer(self, batch_prot, extra_info=False):
  function calculate_contrastive_loss (line 138) | def calculate_contrastive_loss(z, augmented_z, temperature=0.5, device="...

FILE: src/atomica/models/tools.py
  function _unit_edges_from_block_edges (line 11) | def _unit_edges_from_block_edges(unit_block_id, block_src_dst, Z=None, k...
  function _block_edge_dist (line 96) | def _block_edge_dist(X, block_id, src_dst):
  function sequential_and (line 114) | def sequential_and(*tensors):
  function sequential_or (line 121) | def sequential_or(*tensors):
  function graph_to_batch (line 128) | def graph_to_batch(tensor, batch_id, padding_value=0, mask_is_pad=True):
  function _knn_edges (line 148) | def _knn_edges(dist, src_dst, k_neighbors, batch_info):
  function _radial_edges (line 180) | def _radial_edges(dist, src_dst, dist_cut_off):
  class BatchEdgeConstructor (line 191) | class BatchEdgeConstructor:
    method __init__ (line 201) | def __init__(self, global_node_id_vocab=[], delete_self_loop=True) -> ...
    method _reset_buffer (line 208) | def _reset_buffer(self):
    method get_batch_edges (line 221) | def get_batch_edges(self, batch_id):
    method _prepare (line 244) | def _prepare(self, S, batch_id, segment_ids) -> None:
    method _construct_intra_edges (line 265) | def _construct_intra_edges(self, S, batch_id, segment_ids, **kwargs):
    method _construct_inter_edges (line 272) | def _construct_inter_edges(self, S, batch_id, segment_ids, **kwargs):
    method _construct_global_edges (line 279) | def _construct_global_edges(self, S, batch_id, segment_ids, **kwargs):
    method _construct_seq_edges (line 289) | def _construct_seq_edges(self, S, batch_id, segment_ids, **kwargs):
    method __call__ (line 301) | def __call__(self, S, batch_id, segment_ids, **kwargs):
  class KNNBatchEdgeConstructor (line 325) | class KNNBatchEdgeConstructor(BatchEdgeConstructor):
    method __init__ (line 326) | def __init__(self, k_neighbors, global_message_passing=True, global_no...
    method _construct_intra_edges (line 331) | def _construct_intra_edges(self, S, batch_id, segment_ids, **kwargs):
    method _construct_inter_edges (line 342) | def _construct_inter_edges(self, S, batch_id, segment_ids, **kwargs):
    method _construct_global_edges (line 353) | def _construct_global_edges(self, S, batch_id, segment_ids, **kwargs):
    method _construct_seq_edges (line 359) | def _construct_seq_edges(self, S, batch_id, segment_ids, **kwargs):
  class BlockEmbedding (line 364) | class BlockEmbedding(nn.Module):
    method __init__ (line 368) | def __init__(self, num_block_type, num_atom_type, atom_embed_size, blo...
    method forward (line 375) | def forward(self, B, A, block_id):
  class CrossAttention (line 408) | class CrossAttention(nn.Module):
    method __init__ (line 409) | def __init__(self, dim_query, dim_kv, dim_out, num_heads, dropout):
    method forward (line 421) | def forward(self, query, key_value):
  class CrossAttentionWithSpatialEncoding (line 457) | class CrossAttentionWithSpatialEncoding(nn.Module):
    method __init__ (line 458) | def __init__(self, dim_query, dim_kv, dim_out, num_heads, dropout):
    method forward (line 480) | def forward(self, query, key_value, pairwise_distance, pairwise_mask):

FILE: src/atomica/train.py
  function parse (line 27) | def parse():
  function create_dataset (line 152) | def create_dataset(task, path, path2=None, path3=None, fragment=None, ra...
  function set_noise (line 263) | def set_noise(dataset, args):
  function create_trainer (line 291) | def create_trainer(model, train_loader, valid_loader, config, resume_sta...
  function main (line 324) | def main(args):
  function cli (line 628) | def cli():

FILE: src/atomica/trainers/abs_trainer.py
  class TrainConfig (line 13) | class TrainConfig:
    method __init__ (line 14) | def __init__(self, save_dir, lr, max_epoch,
    method add_parameter (line 33) | def add_parameter(self, **kwargs):
    method __str__ (line 36) | def __str__(self):
  class LearningRateWarmup (line 39) | class LearningRateWarmup(object):
    method __init__ (line 41) | def __init__(self, optimizer, warmup_iteration, start_lr, target_lr, a...
    method warmup_learning_rate (line 49) | def warmup_learning_rate(self, cur_iteration):
    method step (line 54) | def step(self, cur_iteration):
    method load_state_dict (line 60) | def load_state_dict(self, state_dict):
    method get_last_lr (line 63) | def get_last_lr(self):
  class Trainer (line 66) | class Trainer:
    method __init__ (line 67) | def __init__(self, model, train_loader, valid_loader, config):
    method to_device (line 104) | def to_device(cls, data, device):
    method _is_main_proc (line 115) | def _is_main_proc(self):
    method _get_version (line 118) | def _get_version(self):
    method _before_train_epoch_start (line 127) | def _before_train_epoch_start(self):
    method _train_epoch (line 130) | def _train_epoch(self, device):
    method _valid_epoch (line 184) | def _valid_epoch(self, device):
    method _metric_better (line 243) | def _metric_better(self, new):
    method _maintain_topk_checkpoint (line 252) | def _maintain_topk_checkpoint(self, valid_metric, ckpt_path):
    method _maintain_topk_weights (line 278) | def _maintain_topk_weights(self, valid_metric, weights_path):
    method train (line 304) | def train(self, device_ids, local_rank, use_wandb=False, use_raytune=F...
    method log (line 336) | def log(self, name, value, step, val=False):
    method get_optimizer (line 348) | def get_optimizer(self):
    method get_scheduler (line 353) | def get_scheduler(self, optimizer):
    method train_step (line 362) | def train_step(self, batch, batch_idx):
    method valid_step (line 368) | def valid_step(self, batch, batch_idx):

FILE: src/atomica/trainers/affinity_trainer.py
  class AffinityTrainer (line 18) | class AffinityTrainer(Trainer):
    method __init__ (line 20) | def __init__(self, model, train_loader, valid_loader, config):
    method get_optimizer (line 27) | def get_optimizer(self):
    method get_scheduler (line 31) | def get_scheduler(self, optimizer):
    method lr_weight (line 40) | def lr_weight(self, step):
    method train_step (line 45) | def train_step(self, batch, batch_idx):
    method valid_step (line 49) | def valid_step(self, batch, batch_idx):
    method _before_train_epoch_start (line 52) | def _before_train_epoch_start(self):
    method share_step (line 57) | def share_step(self, batch, batch_idx, val=False):
    method _valid_epoch (line 80) | def _valid_epoch(self, device):
  class ClassifierTrainer (line 149) | class ClassifierTrainer(Trainer):
    method __init__ (line 151) | def __init__(self, model, train_loader, valid_loader, config):
    method get_optimizer (line 159) | def get_optimizer(self):
    method get_scheduler (line 163) | def get_scheduler(self, optimizer):
    method lr_weight (line 172) | def lr_weight(self, step):
    method train_step (line 177) | def train_step(self, batch, batch_idx):
    method valid_step (line 181) | def valid_step(self, batch, batch_idx):
    method _before_train_epoch_start (line 184) | def _before_train_epoch_start(self):
    method share_step (line 189) | def share_step(self, batch, batch_idx, val=False):
    method _valid_epoch (line 278) | def _valid_epoch(self, device):
  class MultiClassClassifierTrainer (line 354) | class MultiClassClassifierTrainer(Trainer):
    method __init__ (line 356) | def __init__(self, model, train_loader, valid_loader, config):
    method get_optimizer (line 363) | def get_optimizer(self):
    method get_scheduler (line 367) | def get_scheduler(self, optimizer):
    method lr_weight (line 376) | def lr_weight(self, step):
    method train_step (line 381) | def train_step(self, batch, batch_idx):
    method valid_step (line 385) | def valid_step(self, batch, batch_idx):
    method _before_train_epoch_start (line 388) | def _before_train_epoch_start(self):
    method share_step (line 393) | def share_step(self, batch, batch_idx, val=False):
    method _valid_epoch (line 480) | def _valid_epoch(self, device):

FILE: src/atomica/trainers/deprecated/EAB_trainer.py
  class EABTrainer (line 8) | class EABTrainer(Trainer):
    method __init__ (line 12) | def __init__(self, model, train_loader, valid_loader, config):
    method get_optimizer (line 19) | def get_optimizer(self):
    method get_scheduler (line 23) | def get_scheduler(self, optimizer):
    method lr_weight (line 32) | def lr_weight(self, step):
    method train_step (line 37) | def train_step(self, batch, batch_idx):
    method valid_step (line 40) | def valid_step(self, batch, batch_idx):
    method share_step (line 45) | def share_step(self, batch, batch_idx, val=False):

FILE: src/atomica/trainers/deprecated/binary_trainer.py
  class BinaryPredictorTrainer (line 21) | class BinaryPredictorTrainer(Trainer):
    method __init__ (line 25) | def __init__(self, model, train_loader, valid_loader, config):
    method get_optimizer (line 32) | def get_optimizer(self):
    method get_scheduler (line 36) | def get_scheduler(self, optimizer):
    method lr_weight (line 45) | def lr_weight(self, step):
    method train_step (line 50) | def train_step(self, batch, batch_idx):
    method valid_step (line 54) | def valid_step(self, batch, batch_idx):
    method _before_train_epoch_start (line 57) | def _before_train_epoch_start(self):
    method share_step (line 64) | def share_step(self, batch, batch_idx, val=False):
    method _valid_epoch (line 91) | def _valid_epoch(self, device):

FILE: src/atomica/trainers/deprecated/ddG_trainer.py
  class DDGTrainer (line 20) | class DDGTrainer(Trainer):
    method __init__ (line 24) | def __init__(self, model, train_loader, valid_loader, config):
    method get_optimizer (line 31) | def get_optimizer(self):
    method get_scheduler (line 35) | def get_scheduler(self, optimizer):
    method lr_weight (line 44) | def lr_weight(self, step):
    method train_step (line 49) | def train_step(self, batch, batch_idx):
    method valid_step (line 53) | def valid_step(self, batch, batch_idx):
    method _before_train_epoch_start (line 56) | def _before_train_epoch_start(self):
    method share_step (line 63) | def share_step(self, batch, batch_idx, val=False):
    method _valid_epoch (line 77) | def _valid_epoch(self, device):
  class GLOFTrainer (line 135) | class GLOFTrainer(Trainer):
    method __init__ (line 139) | def __init__(self, model, train_loader, valid_loader, config):
    method get_optimizer (line 146) | def get_optimizer(self):
    method get_scheduler (line 150) | def get_scheduler(self, optimizer):
    method lr_weight (line 159) | def lr_weight(self, step):
    method train_step (line 164) | def train_step(self, batch, batch_idx):
    method valid_step (line 168) | def valid_step(self, batch, batch_idx):
    method _before_train_epoch_start (line 171) | def _before_train_epoch_start(self):
    method share_step (line 178) | def share_step(self, batch, batch_idx, val=False):
    method _valid_epoch (line 192) | def _valid_epoch(self, device):

FILE: src/atomica/trainers/deprecated/ec_trainer.py
  class ECTrainer (line 11) | class ECTrainer(Trainer):
    method __init__ (line 15) | def __init__(self, model, train_loader, valid_loader, config):
    method get_optimizer (line 24) | def get_optimizer(self):
    method get_scheduler (line 28) | def get_scheduler(self, optimizer):
    method lr_weight (line 56) | def lr_weight(self, step):
    method train_step (line 61) | def train_step(self, batch, batch_idx):
    method valid_step (line 64) | def valid_step(self, batch, batch_idx):
    method share_step (line 70) | def share_step(self, batch, batch_idx, val=False):

FILE: src/atomica/trainers/deprecated/graph_classification_trainer.py
  class GraphClassificationTrainer (line 8) | class GraphClassificationTrainer(Trainer):
    method __init__ (line 12) | def __init__(self, model, train_loader, valid_loader, config):
    method get_optimizer (line 19) | def get_optimizer(self):
    method get_scheduler (line 23) | def get_scheduler(self, optimizer):
    method lr_weight (line 32) | def lr_weight(self, step):
    method train_step (line 37) | def train_step(self, batch, batch_idx):
    method valid_step (line 40) | def valid_step(self, batch, batch_idx):
    method _before_train_epoch_start (line 43) | def _before_train_epoch_start(self):
    method share_step (line 50) | def share_step(self, batch, batch_idx, val=False):

FILE: src/atomica/trainers/deprecated/graph_pair_classification_trainer.py
  class GraphPairClassificationTrainer (line 8) | class GraphPairClassificationTrainer(Trainer):
    method __init__ (line 12) | def __init__(self, model, train_loader, valid_loader, config):
    method get_optimizer (line 19) | def get_optimizer(self):
    method get_scheduler (line 23) | def get_scheduler(self, optimizer):
    method lr_weight (line 32) | def lr_weight(self, step):
    method train_step (line 37) | def train_step(self, batch, batch_idx):
    method valid_step (line 40) | def valid_step(self, batch, batch_idx):
    method _before_train_epoch_start (line 43) | def _before_train_epoch_start(self):
    method share_step (line 50) | def share_step(self, batch, batch_idx, val=False):

FILE: src/atomica/trainers/masking_trainer.py
  class MaskingTrainer (line 12) | class MaskingTrainer(Trainer):
    method __init__ (line 13) | def __init__(self, model, train_loader, valid_loader, config, resume_s...
    method get_optimizer (line 34) | def get_optimizer(self):
    method get_scheduler (line 38) | def get_scheduler(self, optimizer):
    method lr_weight (line 45) | def lr_weight(self, step):
    method train_step (line 50) | def train_step(self, batch, batch_idx):
    method valid_step (line 54) | def valid_step(self, batch, batch_idx):
    method _before_train_epoch_start (line 57) | def _before_train_epoch_start(self):
    method share_step (line 62) | def share_step(self, batch, batch_idx, val=False):
    method _train_epoch (line 94) | def _train_epoch(self, device):
    method _get_training_state (line 154) | def _get_training_state(self):
    method _valid_epoch (line 163) | def _valid_epoch(self, device):

FILE: src/atomica/trainers/pretrain_trainer.py
  class PretrainTrainer (line 12) | class PretrainTrainer(Trainer):
    method __init__ (line 14) | def __init__(self, model, train_loader, valid_loader, config, resume_s...
    method get_optimizer (line 35) | def get_optimizer(self):
    method get_scheduler (line 39) | def get_scheduler(self, optimizer):
    method lr_weight (line 46) | def lr_weight(self, step):
    method train_step (line 51) | def train_step(self, batch, batch_idx):
    method valid_step (line 54) | def valid_step(self, batch, batch_idx):
    method _before_train_epoch_start (line 57) | def _before_train_epoch_start(self):
    method share_step (line 63) | def share_step(self, batch, batch_idx, val=False):
    method _train_epoch (line 104) | def _train_epoch(self, device):
    method _get_training_state (line 191) | def _get_training_state(self):
    method _valid_epoch (line 200) | def _valid_epoch(self, device):
  class PretrainMaskingNoisingTrainer (line 275) | class PretrainMaskingNoisingTrainer(PretrainTrainer):
    method __init__ (line 277) | def __init__(self, model, train_loader, valid_loader, config, resume_s...
    method share_step (line 280) | def share_step(self, batch, batch_idx, val=False):
    method _train_epoch (line 322) | def _train_epoch(self, device):
    method _valid_epoch (line 407) | def _valid_epoch(self, device):
  class PretrainMaskingNoisingTrainerWithBlockEmbedding (line 490) | class PretrainMaskingNoisingTrainerWithBlockEmbedding(PretrainMaskingNoi...
    method __init__ (line 491) | def __init__(self, model, train_loader, valid_loader, config, resume_s...
    method share_step (line 494) | def share_step(self, batch, batch_idx, val=False):

FILE: src/atomica/trainers/prot_interface_trainer.py
  class ProtInterfaceTrainer (line 11) | class ProtInterfaceTrainer(Trainer):
    method __init__ (line 13) | def __init__(self, model, train_loader, valid_loader, config):
    method get_optimizer (line 20) | def get_optimizer(self):
    method get_scheduler (line 24) | def get_scheduler(self, optimizer):
    method lr_weight (line 31) | def lr_weight(self, step):
    method train_step (line 36) | def train_step(self, batch, batch_idx):
    method valid_step (line 40) | def valid_step(self, batch, batch_idx):
    method _before_train_epoch_start (line 43) | def _before_train_epoch_start(self):
    method share_step (line 47) | def share_step(self, batch, batch_idx, val=False):
    method _valid_epoch (line 64) | def _valid_epoch(self, device):

FILE: src/atomica/utils/logger.py
  function init_map (line 12) | def init_map():
  function get_prio (line 19) | def get_prio(level):
  function print_log (line 26) | def print_log(s, level='INFO', end='\n', no_prefix=False):

FILE: src/atomica/utils/losses.py
  class FocalLoss (line 10) | class FocalLoss(nn.Module):
    method __init__ (line 38) | def __init__(self, alpha=None, gamma=2.0, reduction='mean'):
    method forward (line 50) | def forward(self, inputs, targets):
  class MultiLabelFocalLoss (line 94) | class MultiLabelFocalLoss(nn.Module):
    method __init__ (line 117) | def __init__(self, alpha=None, gamma=2.0, reduction='mean'):
    method forward (line 129) | def forward(self, inputs, targets):

FILE: src/atomica/utils/nn_utils.py
  function print_cuda_memory (line 3) | def print_cuda_memory():
  function count_parameters (line 10) | def count_parameters(model):

FILE: src/atomica/utils/noise_transforms.py
  function modify_conformer_torsion_angles (line 13) | def modify_conformer_torsion_angles(coords, rotateable_edges, mask_rotat...
  function rigid_transform_Kabsch_3D (line 43) | def rigid_transform_Kabsch_3D(A, B):
  class CropTransform (line 79) | class CropTransform:
    method __init__ (line 80) | def __init__(self, max_blocks, fragmentation_method):
    method __call__ (line 86) | def __call__(self, data):
  class TorsionNoiseTransform (line 136) | class TorsionNoiseTransform:
    method __init__ (line 137) | def __init__(self, tor_sigma):
    method __call__ (line 140) | def __call__(self, data, chosen_segment):
  class GaussianNoiseTransform (line 224) | class GaussianNoiseTransform:
    method __init__ (line 225) | def __init__(self, sigma):
    method __call__ (line 228) | def __call__(self, data, chosen_segment):
  class GlobalTranslationTransform (line 261) | class GlobalTranslationTransform:
    method __init__ (line 262) | def __init__(self, tr_sigma):
    method __call__ (line 265) | def __call__(self, data, chosen_segment):
  function _expansion (line 296) | def _expansion(theta, sigma, L=2000):  # the summation term only
  function _density (line 302) | def _density(expansion, theta):
  function _score (line 307) | def _score(exp, theta, sigma, L=2000):
  class GlobalRotationTransform (line 318) | class GlobalRotationTransform:
    method __init__ (line 319) | def __init__(self, rot_sigma, max_theta):
    method __call__ (line 329) | def __call__(self, data, chosen_segment):

FILE: src/atomica/utils/random_seed.py
  function setup_seed (line 9) | def setup_seed(seed):

FILE: src/atomica/utils/torus.py
  function p_torus (line 13) | def p_torus(x, sigma, N=10):
  function grad (line 20) | def grad(x, sigma, N=10):
  function _load_torus_data (line 34) | def _load_torus_data():
  function score (line 60) | def score(x, sigma):
  function p (line 73) | def p(x, sigma):
  function sample (line 85) | def sample(sigma):
  function _load_score_norm (line 94) | def _load_score_norm():
  function score_norm (line 110) | def score_norm(sigma):

FILE: src/atomica/utils/visualize.py
  class MolGraph (line 100) | class MolGraph:
    method __init__ (line 114) | def __init__(self):
    method read_xyz (line 124) | def read_xyz(self, file_path: str) -> None:
    method read_coords (line 139) | def read_coords(self, atoms, coords):
    method _generate_adjacency_list (line 147) | def _generate_adjacency_list(self):
    method edges (line 169) | def edges(self):
    method __len__ (line 180) | def __len__(self):
    method __getitem__ (line 183) | def __getitem__(self, position):
  function to_plotly_figure (line 227) | def to_plotly_figure(graph: MolGraph) -> go.Figure:
  function to_networkx_graph (line 375) | def to_networkx_graph(graph: MolGraph) -> nx.Graph:

FILE: tests/test_import.py
  function test_import_atomica (line 3) | def test_import_atomica():
  function test_import_models (line 9) | def test_import_models():
  function test_import_data (line 15) | def test_import_data():
  function test_cli_entry_points (line 21) | def test_cli_entry_points():

FILE: tutorials/3_rna_structure_function/multiclass_metrics.py
  class MetricsResult (line 39) | class MetricsResult:
    method to_dict (line 69) | def to_dict(self) -> Dict[str, Any]:
  function _normalize_rows (line 74) | def _normalize_rows(mat: np.ndarray, eps: float = 1e-12) -> np.ndarray:
  function _ensure_label_space (line 83) | def _ensure_label_space(y_true: np.ndarray, y_pred: np.ndarray, labels: ...
  function compute_multiclass_metrics (line 91) | def compute_multiclass_metrics(

FILE: tutorials/3_rna_structure_function/multilabel_metrics.py
  class MultilabelMetricsResult (line 39) | class MultilabelMetricsResult:
    method to_dict (line 72) | def to_dict(self) -> Dict[str, Any]:
  function _validate_and_prepare (line 76) | def _validate_and_prepare(
  function compute_multilabel_metrics (line 118) | def compute_multilabel_metrics(

FILE: tutorials/3_rna_structure_function/tutorial.py
  function load_model (line 88) | def load_model(task_name: str, seed: int):
  function run_inference (line 99) | def run_inference(task_name: str, split: str, model, device: str) -> pd....
  function _stack_probs (line 148) | def _stack_probs(series: pd.Series) -> np.ndarray:
  function metrics_rnago (line 152) | def metrics_rnago(per_seed_test: List[pd.DataFrame]) -> Tuple[pd.DataFra...
  function metrics_rna_ligand (line 172) | def metrics_rna_ligand(per_seed_test: List[pd.DataFrame]) -> Tuple[pd.Da...
  function _binary_residue_metrics (line 196) | def _binary_residue_metrics(
  function run_task (line 240) | def run_task(task_name: str, device: str) -> None:
  function main (line 318) | def main():

FILE: tutorials/4_atomica_masif_benchmark/tutorial.py
  function load_model (line 43) | def load_model(seed: int) -> Tuple[MultiClassClassifierModel, str]:
  function run_inference (line 55) | def run_inference(model, device: str) -> pd.DataFrame:
  function average_seed_predictions (line 79) | def average_seed_predictions(per_seed: List[pd.DataFrame]) -> pd.DataFrame:
  function main (line 92) | def main():

FILE: tutorials/5_ppi_and_inhibitors/prepare_data.py
  function _load_structure_any (line 58) | def _load_structure_any(path: str, model: int = 1):
  function resolve_structure_path (line 68) | def resolve_structure_path(cif_dir: Path, pdb_id: str) -> Path | None:
  function check_chain_exists (line 76) | def check_chain_exists(path: Path, chain: str):
  function get_lig_resi (line 82) | def get_lig_resi(path: Path, chain: str, lig_code: str) -> List[int]:
  function infill_chain1 (line 90) | def infill_chain1(row):
  function build_inhibitors_index (line 103) | def build_inhibitors_index(cif_dir: Path) -> pd.DataFrame:
  function build_peptide_partners (line 141) | def build_peptide_partners(cif_dir: Path):
  function build_protein_partner_patches (line 176) | def build_protein_partner_patches(cif_dir: Path, mesh_dir: Path,
  function closest_residue_on_chain (line 246) | def closest_residue_on_chain(structure_path, xyz, chain, atom_filter="ca"):
  function run (line 260) | def run(cmd: List[str]):
  function process_and_embed (line 265) | def process_and_embed(ckpt_dir: Path):
  function _chain_ca (line 295) | def _chain_ca(arr, chain):
  function _seq_resids (line 300) | def _seq_resids(ca):
  function _kabsch_refine (line 310) | def _kabsch_refine(P, Q, cutoff=2.0, cycles=5):
  function align_chains (line 334) | def align_chains(pdb1, chain1, pdb2, chain2):
  function get_block_coords (line 351) | def get_block_coords(X, block_lengths):
  function build_peptide_inhibitor_block_results (line 359) | def build_peptide_inhibitor_block_results(cif_dir: Path):
  function main (line 438) | def main():

FILE: tutorials/5_ppi_and_inhibitors/surface_sampler/compute_surface_mesh.py
  function read_msms (line 23) | def read_msms(file_root):
  function compute_msms (line 86) | def compute_msms(pdb_file, density=3.0, probe_radius=1.5, tmp_dir=None):
  function save_mesh_ply (line 157) | def save_mesh_ply(vertices, faces, normals, output_path):

FILE: tutorials/5_ppi_and_inhibitors/surface_sampler/extract_chain.py
  class NotDisordered (line 14) | class NotDisordered(Select):
    method accept_atom (line 16) | def accept_atom(self, atom):
  function find_modified_amino_acids (line 20) | def find_modified_amino_acids(path):
  function extract_chain (line 40) | def extract_chain(infilename, outfilename, chain_ids=None):

FILE: tutorials/5_ppi_and_inhibitors/surface_sampler/pdb_to_xyzrn.py
  function pdb_to_xyzrn (line 45) | def pdb_to_xyzrn(pdb_filename, xyzrn_filename):

FILE: tutorials/5_ppi_and_inhibitors/surface_sampler/sample_surface_points.py
  function compute_triangle_areas (line 8) | def compute_triangle_areas(vertices, faces):
  function sample_point_in_triangle (line 35) | def sample_point_in_triangle(v0, v1, v2):
  function sample_surface_points (line 62) | def sample_surface_points(vertices, faces, num_points, seed=None):
  function load_ply_mesh (line 107) | def load_ply_mesh(ply_path):
  function save_points_xyz (line 156) | def save_points_xyz(points, output_path):

FILE: tutorials/5_ppi_and_inhibitors/surface_sampler/surface_sampler.py
  function get_mesh (line 19) | def get_mesh(pdb_cif_path, chain_id, output_path=None, density=3.0, prob...
  function sample_points (line 80) | def sample_points(mesh_path, num_points, output_path=None, seed=None):
  function get_mesh_and_sample (line 118) | def get_mesh_and_sample(pdb_cif_path, chain_id, num_points,

FILE: tutorials/5_ppi_and_inhibitors/tutorial_protein_peptide_inhibitors.py
  function load_metadata (line 56) | def load_metadata():
  function load_embeddings (line 65) | def load_embeddings(df_2p2idb):
  function precision_at_k (line 86) | def precision_at_k(E, G, k=K, geom_thresh=GEOM_THRESH):
  function enrichment_at_k (line 96) | def enrichment_at_k(E, G, k=K, geom_thresh=GEOM_THRESH):
  function load_block_results (line 105) | def load_block_results():
  function report_stats (line 122) | def report_stats(br, families):
  function plot_stripplot (line 162) | def plot_stripplot(br, value_col, ylabel, out_name, yticks=None):
  function pick_entry_by_name (line 191) | def pick_entry_by_name(br, family, lig_code=None, inhibitor_pdb=None):
  function plot_swarm_by_name (line 206) | def plot_swarm_by_name(br, family, lig_code, inhibitor_pdb):
  function main (line 261) | def main():

FILE: tutorials/5_ppi_and_inhibitors/tutorial_protein_protein_inhibitors.py
  function load_metadata (line 57) | def load_metadata():
  function load_inhibitor_embeddings (line 66) | def load_inhibitor_embeddings(df_2p2idb):
  function load_ppi_patches (line 85) | def load_ppi_patches(df_2p2idb):
  function precision_at_k (line 100) | def precision_at_k(E, G, k=K, geom_thresh=GEOM_THRESH):
  function enrichment_at_k (line 110) | def enrichment_at_k(E, G, k=K, geom_thresh=GEOM_THRESH):
  function run_retrieval (line 119) | def run_retrieval(inh_emb, ppi_emb, dist_df):
  function report_stats (line 165) | def report_stats(inh_emb):
  function plot_stripplot (line 189) | def plot_stripplot(df, value_col, ylabel, out_name, yticks=None):
  function pick_inhibitor_by_name (line 218) | def pick_inhibitor_by_name(inh_emb, family, pdb_code):
  function plot_swarm_by_name (line 230) | def plot_swarm_by_name(inh_emb, ppi_points, family, pdb_code, seed=0):
  function main (line 287) | def main():