Skip to content

Allow writing multi-model structure files with different atoms in each model #921

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
padix-key merged 1 commit into from
Jul 2, 2026
Merged

Allow writing multi-model structure files with different atoms in each model #921

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions src/biotite/structure/atoms.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -598,7 +598,7 @@ def __copy_create__(self) -> "Atom":
return Atom(self.coord, **self._annot)


class AtomArray(_AtomArrayBase, Generic[N]):
class AtomArray(_AtomArrayBase, Sequence["Atom"], Generic[N]):
"""
An array representation of a model consisting of multiple atoms.

Expand Down Expand Up @@ -891,7 +891,7 @@ def __copy_create__(self) -> "AtomArray[Any]":
return AtomArray(self.array_length())


class AtomArrayStack(_AtomArrayBase, Generic[M, N]):
class AtomArrayStack(_AtomArrayBase, Sequence["AtomArray"], Generic[M, N]):
"""
A collection of multiple :class:`AtomArray` instances, where each
atom array has equal annotation arrays.
Expand Down
212 changes: 99 additions & 113 deletions src/biotite/structure/io/gro/file.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
__author__ = "Daniel Bauer, Patrick Kunzmann"
__all__ = ["GROFile"]

import copy
from collections.abc import Iterable
from datetime import datetime
from typing import Any, overload
import numpy as np
Expand Down Expand Up @@ -217,134 +217,120 @@ def _set_box_dimen(

return array

def set_structure(self, array: AtomArray[N] | AtomArrayStack[M, N]) -> None:
def set_structure(
self,
array: AtomArray[N] | AtomArrayStack[M, N] | Iterable[AtomArray[N]],
) -> None:
"""
Set the :class:`AtomArray` or :class:`AtomArrayStack` for the
file.

Parameters
----------
array : AtomArray or AtomArrayStack
The array or stack to be saved into this file. If a stack
is given, each array in the stack is saved as separate
model.
array : AtomArray or AtomArrayStack or iterable of AtomArray
The structure to be written.
If a stack or an iterable of :class:`AtomArray` is given, multiple models
will be written.
"""

def get_box_dimen(array: AtomArray[Any]) -> str:
"""
GRO files have the box dimensions as last line for each
model.
In case, the `box` attribute of the atom array is
`None`, we simply use the min and max coordinates in xyz
to get the correct size

Parameters
----------
array : AtomArray
The atom array to get the box dimensions from.

Returns
-------
box : str
The box, properly formatted for GRO files.
"""
if array.box is None:
coord = array.coord
bx, by, bz = (coord.max(axis=0) - coord.min(axis=0)) / 10
return f"{bx:>8.3f} {by:>8.3f} {bz:>8.3f}"
else:
box = array.box
if is_orthogonal(box):
bx, by, bz = np.diag(box) / 10
return f"{bx:>9.5f} {by:>9.5f} {bz:>9.5f}"
else:
box = box / 10
box_elements = (
box[0, 0],
box[1, 1],
box[2, 2],
box[0, 1],
box[0, 2],
box[1, 0],
box[1, 2],
box[2, 0],
box[2, 1],
)
return " ".join([f"{e:>9.5f}" for e in box_elements])

if "atom_id" in array.get_annotation_categories():
atom_id = array.atom_id
else:
atom_id = np.arange(1, array.array_length() + 1)
# Atom IDs are supported up to 99999,
# but negative IDs are also possible
gro_atom_id = np.where(atom_id > 0, ((atom_id - 1) % 99999) + 1, atom_id)
# Residue IDs are supported up to 9999,
# but negative IDs are also possible
gro_res_id = np.where(
array.res_id > 0, ((array.res_id - 1) % 99999) + 1, array.res_id
)

if isinstance(array, AtomArray):
new_lines: list[str] = [""] * (array.array_length() + 3)

# Write header lines
new_lines[0] = f"Generated by Biotite at {datetime.now()}"
new_lines[1] = str(array.array_length())

# Write atom lines
fmt = "{:>5d}{:5s}{:>5s}{:>5d}{:>8.3f}{:>8.3f}{:>8.3f}"
for i in range(array.array_length()):
models = [array]
else:
models = list(array)
# Only multi-model files get the model number in the comment line
multiple_models = len(models) > 1

self.lines = []
fmt = "{:>5d}{:5s}{:>5s}{:>5d}{:>8.3f}{:>8.3f}{:>8.3f}"
for model_i, model in enumerate(models):
gro_res_id, gro_atom_id = _get_gro_ids(model)
header = f"Generated by Biotite at {datetime.now()}"
if multiple_models:
header += f", model={model_i + 1}"
self.lines.append(header)
self.lines.append(str(model.array_length()))
for i in range(model.array_length()):
# gro format is in nm -> multiply coords by 10
new_lines[i + 2] = fmt.format(
gro_res_id[i],
array.res_name[i],
array.atom_name[i],
gro_atom_id[i],
array.coord[i, 0] / 10,
array.coord[i, 1] / 10,
array.coord[i, 2] / 10,
)
# Write box lines
new_lines[-1] = get_box_dimen(array)
self.lines = new_lines
elif isinstance(array, AtomArrayStack):
self.lines = []
# The entire information, but the coordinates,
# is equal for each model
# Therefore template lines are created
# which are afterwards applied for each model
templines: list[str] = [""] * array.array_length()
fmt = "{:>5d}{:5s}{:>5s}{:5d}"
for i in range(array.array_length()):
templines[i] = fmt.format(
gro_res_id[i], array.res_name[i], array.atom_name[i], gro_atom_id[i]
)

for i in range(array.stack_depth()):
self.lines.append(
f"Generated by Biotite at {datetime.now()}, model={i + 1}"
)
self.lines.append(str(array.array_length()))

# Fill in coordinates for each model
modellines = copy.copy(templines)
for j, line in enumerate(modellines):
# Insert coordinates
line = line + "{:>8.3f}{:>8.3f}{:>8.3f}".format(
array.coord[i, j, 0] / 10,
array.coord[i, j, 1] / 10,
array.coord[i, j, 2] / 10,
fmt.format(
gro_res_id[i],
model.res_name[i],
model.atom_name[i],
gro_atom_id[i],
model.coord[i, 0] / 10,
model.coord[i, 1] / 10,
model.coord[i, 2] / 10,
)
modellines[j] = line
self.lines.extend(modellines)
self.lines.append(get_box_dimen(array[i]))
else:
raise TypeError("An atom array or stack must be provided")
)
self.lines.append(_get_box_dimen(model))
# Add terminal newline, since PyMOL requires it
self.lines.append("")


def _get_box_dimen(array: AtomArray[Any]) -> str:
"""
GRO files have the box dimensions as last line for each
model.
In case, the `box` attribute of the atom array is
`None`, we simply use the min and max coordinates in xyz
to get the correct size

Parameters
----------
array : AtomArray
The atom array to get the box dimensions from.

Returns
-------
box : str
The box, properly formatted for GRO files.
"""
if array.box is None:
coord = array.coord
bx, by, bz = (coord.max(axis=0) - coord.min(axis=0)) / 10
return f"{bx:>8.3f} {by:>8.3f} {bz:>8.3f}"
else:
box = array.box
if is_orthogonal(box):
bx, by, bz = np.diag(box) / 10
return f"{bx:>9.5f} {by:>9.5f} {bz:>9.5f}"
else:
box = box / 10
box_elements = (
box[0, 0],
box[1, 1],
box[2, 2],
box[0, 1],
box[0, 2],
box[1, 0],
box[1, 2],
box[2, 0],
box[2, 1],
)
return " ".join([f"{e:>9.5f}" for e in box_elements])


def _get_gro_ids(
model: AtomArray[Any],
) -> tuple[NDArray1[Any, np.integer], NDArray1[Any, np.integer]]:
"""
Get the GRO-compatible residue and atom IDs for a model.
"""
if "atom_id" in model.get_annotation_categories():
atom_id = model.atom_id
else:
atom_id = np.arange(1, model.array_length() + 1)
# Atom IDs are supported up to 99999,
# but negative IDs are also possible
gro_atom_id = np.where(atom_id > 0, ((atom_id - 1) % 99999) + 1, atom_id)
# Residue IDs are supported up to 9999,
# but negative IDs are also possible
gro_res_id = np.where(
model.res_id > 0, ((model.res_id - 1) % 99999) + 1, model.res_id
)
return gro_res_id, gro_atom_id


def _is_int(string: str) -> bool:
"""
Return ``True`, if the string can be parsed to an int.
Expand Down
9 changes: 5 additions & 4 deletions src/biotite/structure/io/pdb/convert.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -123,7 +123,7 @@ def get_structure(

def set_structure(
pdb_file: PDBFile,
array: AtomArray[N] | AtomArrayStack[M, N],
array: AtomArray[N] | AtomArrayStack[M, N] | Iterable[AtomArray[N]],
hybrid36: bool = False,
) -> None:
"""
Expand All @@ -142,9 +142,10 @@ def set_structure(
----------
pdb_file : PDBFile
The file object.
array : AtomArray or AtomArrayStack
The structure to be written. If a stack is given, each array in
the stack will be in a separate model.
array : AtomArray or AtomArrayStack or iterable of AtomArray
The structure to be written.
If a stack or an iterable of :class:`AtomArray` is given, multiple models will
be written.
hybrid36 : boolean, optional
Defines whether the file should be written in hybrid-36 format.

Expand Down
62 changes: 41 additions & 21 deletions src/biotite/structure/io/pdb/file.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -169,29 +169,21 @@ def get_structure(

def set_structure(
self,
atoms: AtomArray[N] | AtomArrayStack[M, N],
atoms: AtomArray[N] | AtomArrayStack[M, N] | Iterable[AtomArray[N]],
hybrid36: bool = False,
) -> None:
_check_pdb_compatibility(atoms, hybrid36)

# PDB files only contains ``CONECT`` records for bonds between non-water hetero
# residues and inter-residue bonds
# -> Preprocess `AtomArray` to remove those bonds
if atoms.bonds is not None:
original_bonds = atoms.bonds
# We only replace the BondList -> shallow copy is enough
atoms = copy.copy(atoms)
hetero_indices = np.where(atoms.hetero & ~filter_solvent(atoms))[0]
bond_array = original_bonds.as_array()
bond_array = bond_array[
np.isin(bond_array[:, 0], hetero_indices)
| np.isin(bond_array[:, 1], hetero_indices)
| (atoms.res_id[bond_array[:, 0]] != atoms.res_id[bond_array[:, 1]])
| (atoms.chain_id[bond_array[:, 0]] != atoms.chain_id[bond_array[:, 1]])
]
atoms.bonds = BondList(atoms.array_length(), bond_array)

super().set_structure(atoms, hybrid36)
if isinstance(atoms, (AtomArray, AtomArrayStack)):
# A single `AtomArray` is written as a single model
_check_pdb_compatibility(atoms, hybrid36)
atoms = _remove_non_conect_bonds(atoms)
super().set_structure(atoms, hybrid36)
elif isinstance(atoms, Iterable):
models = [_remove_non_conect_bonds(model) for model in atoms]
if len(models) == 0:
raise BadStructureError("Structure must not be empty")
for model in models:
_check_pdb_compatibility(model, hybrid36)
super().set_structure(models, hybrid36)

def get_space_group(self) -> SpaceGroupInfo:
"""
Expand Down Expand Up @@ -748,6 +740,34 @@ def _apply_transformations(
return assembly


def _remove_non_conect_bonds(
atoms: AtomArray[Any] | AtomArrayStack[Any, Any],
) -> AtomArray[Any] | AtomArrayStack[Any, Any]:
"""
Remove all bonds that are not written as ``CONECT`` records.

PDB files only contain ``CONECT`` records for bonds between non-water
hetero residues and for inter-residue bonds.
If `atoms` has no associated :class:`BondList`, it is returned
unchanged.
"""
if atoms.bonds is None:
return atoms
original_bonds = atoms.bonds
# We only replace the BondList -> shallow copy is enough
atoms = copy.copy(atoms)
hetero_indices = np.where(atoms.hetero & ~filter_solvent(atoms))[0]
bond_array = original_bonds.as_array()
bond_array = bond_array[
np.isin(bond_array[:, 0], hetero_indices)
| np.isin(bond_array[:, 1], hetero_indices)
| (atoms.res_id[bond_array[:, 0]] != atoms.res_id[bond_array[:, 1]])
| (atoms.chain_id[bond_array[:, 0]] != atoms.chain_id[bond_array[:, 1]])
]
atoms.bonds = BondList(atoms.array_length(), bond_array)
return atoms


def _check_pdb_compatibility(
array: AtomArray[Any] | AtomArrayStack[Any, Any], hybrid36: bool
) -> None:
Expand Down
Loading
Loading