Skip to content

synth_afm.io

The io module handles loading protein structure files and assigning per-atom van der Waals radii for use with the AFM simulation kernels.

load_coords_and_radii(file_path, filter_protein=False)

Loads coordinates and assigns vdW radii from a PDB/mmCIF file.

Parameters:

Name Type Description Default
file_path str

Path to the structure file.

 required
filter_protein bool

If True, filters only amino acid atoms.

 False

Returns:

Name Type Description
tuple Tuple[ndarray, ndarray]

(coords (N, 3), radii (N,)) as JAX arrays.
Source code in synth_afm/io.py 
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
def load_coords_and_radii(
    file_path: str, filter_protein: bool = False
) -> Tuple[jnp.ndarray, jnp.ndarray]:
    """
    Loads coordinates and assigns vdW radii from a PDB/mmCIF file.

    Args:
        file_path: Path to the structure file.
        filter_protein: If True, filters only amino acid atoms.

    Returns:
        tuple: (coords (N, 3), radii (N,)) as JAX arrays.
    """
    atom_array = strucio.load_structure(file_path)

    if filter_protein:
        atom_array = atom_array[struc.filter_amino_acids(atom_array)]

    coords = jnp.array(atom_array.coord)

    # Assign radii based on element
    elements = atom_array.element
    radii_list = []
    for elem in elements:
        radii_list.append(VDW_RADII.get(elem.upper(), DEFAULT_ATOM_RADIUS))

    radii = jnp.array(radii_list)

    return coords, radii