diff --git a/src/biotite/structure/atoms.py b/src/biotite/structure/atoms.py index 7eb876399..8c6973dcd 100644 --- a/src/biotite/structure/atoms.py +++ b/src/biotite/structure/atoms.py @@ -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. @@ -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. diff --git a/src/biotite/structure/io/gro/file.py b/src/biotite/structure/io/gro/file.py index 0802902f9..4b9475022 100644 --- a/src/biotite/structure/io/gro/file.py +++ b/src/biotite/structure/io/gro/file.py @@ -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 @@ -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. diff --git a/src/biotite/structure/io/pdb/convert.py b/src/biotite/structure/io/pdb/convert.py index d1a0035cf..454ce477b 100644 --- a/src/biotite/structure/io/pdb/convert.py +++ b/src/biotite/structure/io/pdb/convert.py @@ -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: """ @@ -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. diff --git a/src/biotite/structure/io/pdb/file.py b/src/biotite/structure/io/pdb/file.py index 14fcffbec..3a37489ff 100644 --- a/src/biotite/structure/io/pdb/file.py +++ b/src/biotite/structure/io/pdb/file.py @@ -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: """ @@ -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: diff --git a/src/biotite/structure/io/pdbx/convert.py b/src/biotite/structure/io/pdbx/convert.py index 48e4b15dd..208932d0b 100644 --- a/src/biotite/structure/io/pdbx/convert.py +++ b/src/biotite/structure/io/pdbx/convert.py @@ -966,7 +966,7 @@ def _get_chem_comp_bond_type(bond_type: BondType) -> tuple[str, str]: def set_structure( pdbx_file: _PDBxFile, - array: AtomArray[N] | AtomArrayStack[M, N], + array: AtomArray[N] | AtomArrayStack[M, N] | Iterable[AtomArray[N]], # pyright: ignore[reportRedeclaration] data_block: str | None = None, include_bonds: bool = False, extra_fields: Iterable[str] = [], @@ -988,9 +988,10 @@ def set_structure( ---------- pdbx_file : CIFFile or CIFBlock or BinaryCIFFile or BinaryCIFBlock 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. data_block : str, optional The name of the data block. Default is the first (and most times only) data block of the @@ -1027,6 +1028,21 @@ def set_structure( DeprecationWarning, ) + if not isinstance(array, (AtomArray, AtomArrayStack)): + # An iterable of `AtomArray` objects represents multiple models that - + # unlike an `AtomArrayStack` - may contain different atoms in each model + # -> concatenate them into a single flat `AtomArray` and store the model + # membership in the `pdbx_PDB_model_num` annotation, which is written below + models: list[AtomArray[Any]] = list(array) + if len(models) == 0: + raise BadStructureError("Structure must not be empty") + model_num = np.repeat( + np.arange(1, len(models) + 1, dtype=np.int32), + [model.array_length() for model in models], + ) + array: AtomArray[Any] = concatenate(models) + array.set_annotation("pdbx_PDB_model_num", model_num) + _check_non_empty(array) block = _get_or_create_block(pdbx_file, data_block) @@ -1127,7 +1143,16 @@ def set_structure( atom_site["Cartn_x"] = np.copy(np.ravel(array.coord[..., 0])) atom_site["Cartn_y"] = np.copy(np.ravel(array.coord[..., 1])) atom_site["Cartn_z"] = np.copy(np.ravel(array.coord[..., 2])) - atom_site["pdbx_PDB_model_num"] = np.ones(array.array_length(), dtype=np.int32) + # An iterable of `AtomArray` was concatenated into a single flat `AtomArray`, + # whose `pdbx_PDB_model_num` annotation stores the model membership + if "pdbx_PDB_model_num" in annot_categories: + atom_site["pdbx_PDB_model_num"] = np.copy( + array.get_annotation("pdbx_PDB_model_num") + ) + else: + atom_site["pdbx_PDB_model_num"] = np.ones( + array.array_length(), dtype=np.int32 + ) # In case of multiple models repeat annotations # and use model specific coordinates else: diff --git a/src/rust/structure/io/pdb/file.rs b/src/rust/structure/io/pdb/file.rs index 3df6108dd..c4444ba40 100644 --- a/src/rust/structure/io/pdb/file.rs +++ b/src/rust/structure/io/pdb/file.rs @@ -1,10 +1,7 @@ //! Low-level PDB file parsing and writing. use num_traits::{Float, FloatConst}; -use numpy::ndarray::Axis; -use numpy::ndarray::{ - Array, Array1, Array2, ArrayView1, ArrayView3, ArrayViewD, ArrayViewMut1, Ix3, -}; +use numpy::ndarray::{Array, Array1, Array2, ArrayView1, ArrayView2, ArrayViewMut1, Ix2}; use numpy::PyArrayMethods; use numpy::{IntoPyArray, PyArray1, PyArray2, PyArrayDyn, PyReadonlyArray1, PyReadonlyArray2}; use pyo3::exceptions; @@ -702,10 +699,13 @@ impl PDBFile { /// /// 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. + /// atoms : AtomArray or AtomArrayStack or list of AtomArray + /// The structure to be written. + /// If a stack is given, each model in the stack is written as a + /// separate model. + /// If a list of :class:`AtomArray` is given, each array is written + /// as a separate model, which - unlike a stack - may contain + /// different atoms. /// hybrid36 : bool, optional /// Defines wether the file should be written in hybrid-36 /// format. @@ -721,159 +721,39 @@ impl PDBFile { atoms: Bound<'_, PyAny>, hybrid36: bool, ) -> PyResult<()> { - // Write the CRYST1 record if box is present - let box_vectors = atoms.getattr("box")?; - if !box_vectors.is_none() { - // For AtomArrayStack, the box is 3D - use first model's box - let box_ndim: usize = box_vectors.getattr("ndim")?.extract()?; - let box_to_write = match box_ndim { - 3 => box_vectors.get_item(0)?, - 2 => box_vectors, - _ => Err(biotite::BadStructureError::new_err( - "Invalid box dimensions", - ))?, - }; - self.write_box(py, box_to_write)?; - } - - // Create a binding first to avoid borrowing a dropped array - let coord_array = atoms - .getattr("coord")? - .extract::>>()? - .readonly(); - let coord: ArrayViewD = coord_array.as_array(); - let coord: ArrayView3 = match coord.ndim() { - 2 => coord.insert_axis(Axis(0)), - 3 => coord, - _ => Err(biotite::BadStructureError::new_err( - "Invalid dimensions of coordinates", - ))?, - } - .into_dimensionality::() - .unwrap(); - let is_multi_model = coord.shape()[0] > 1; - - let chain_id = get_annotation_as_strings(&atoms, "chain_id")?; - let res_id_numpy = get_annotation_as_type(&atoms, "res_id", "int64")?.readonly(); - let res_id: ArrayView1 = res_id_numpy.as_array(); - let ins_code = get_annotation_as_strings(&atoms, "ins_code")?; - let res_name = get_annotation_as_strings(&atoms, "res_name")?; - let hetero_numpy = get_annotation_as_type(&atoms, "hetero", "bool")?.readonly(); - let hetero: ArrayView1 = hetero_numpy.as_array(); - let atom_name = get_annotation_as_strings(&atoms, "atom_name")?; - let element = get_annotation_as_strings(&atoms, "element")?; - - let atom_id_numpy: Option> = if atoms.hasattr("atom_id")? { - Some(get_annotation_as_type(&atoms, "atom_id", "int64")?.readonly()) - } else { - None - }; - let atom_id = atom_id_numpy.as_ref().map(|arr| arr.as_array()); - let b_factor_numpy: Option> = if atoms.hasattr("b_factor")? { - Some(get_annotation_as_type(&atoms, "b_factor", "float64")?.readonly()) + let atom_array_class = py.import("biotite.structure")?.getattr("AtomArray")?; + let models: Vec> = if atoms.is_instance(&atom_array_class)? { + vec![atoms] } else { - None + atoms.try_iter()?.collect::>>()? }; - let b_factor = b_factor_numpy.as_ref().map(|arr| arr.as_array()); - let occupancy_numpy: Option> = if atoms.hasattr("occupancy")? { - Some(get_annotation_as_type(&atoms, "occupancy", "float64")?.readonly()) - } else { - None - }; - let occupancy = occupancy_numpy.as_ref().map(|arr| arr.as_array()); - let charge_numpy: Option> = if atoms.hasattr("charge")? { - Some(get_annotation_as_type(&atoms, "charge", "int64")?.readonly()) - } else { - None - }; - let charge = charge_numpy.as_ref().map(|arr| arr.as_array()); - - // These will contain the ATOM records for each atom - // These are reused in every model by adding the coordinates to the string - // This procedure aims to increase the performance is repetitive formatting is omitted - let mut prefix: Vec = Vec::new(); - let mut suffix: Vec = Vec::new(); - - for i in 0..coord.shape()[1] { - let element_i = &element[i]; - let atom_name_i = &atom_name[i]; - - let raw_atom_id_i = atom_id.as_ref().map_or((i + 1) as i64, |arr| arr[i]); - let atom_id_i = if hybrid36 { - encode_hybrid36(raw_atom_id_i, 5)? - } else { - truncate_id(raw_atom_id_i, 99999).to_string() - }; - - let res_id_i = if hybrid36 { - encode_hybrid36(res_id[i], 4)? - } else { - truncate_id(res_id[i], 9999).to_string() - }; - - prefix.push(format!( - "{:6}{:>5} {:4} {:>3} {:1}{:>4}{:1} ", - if hetero[i] { "HETATM" } else { "ATOM" }, - atom_id_i, - if element_i.len() == 1 && atom_name_i.len() < 4 { - format!(" {}", atom_name_i) - } else { - atom_name_i.to_string() - }, - res_name[i], - chain_id[i], - res_id_i, - ins_code[i], + if models.is_empty() { + return Err(biotite::BadStructureError::new_err( + "Structure must not be empty", )); + } + let is_multi_model = models.len() > 1; - suffix.push(format!( - "{:>6.2}{:>6.2} {:>2}{}", - occupancy.as_ref().map_or(1f64, |arr| arr[i]), - b_factor.as_ref().map_or(0f64, |arr| arr[i]), - element_i, - charge.as_ref().map_or(String::from(" "), |arr| { - let c = arr[i]; - match c.cmp(&0) { - Ordering::Greater => format!("{:1}+", c), - Ordering::Less => format!("{:1}-", -c), - Ordering::Equal => String::from(" "), - } - }), - )); + // Write the CRYST1 record from the first model's box, as a PDB file + // can store only a single box for all models + let box_vectors = models[0].getattr("box")?; + if !box_vectors.is_none() { + self.write_cryst_record(py, box_vectors)?; } - for model_i in 0..coord.shape()[0] { + for (model_i, model) in models.iter().enumerate() { if is_multi_model { self.lines.push(format!("MODEL {:>8}", model_i + 1)); } - for atom_i in 0..coord.shape()[1] { - let coord_string = format!( - "{:>8.3}{:>8.3}{:>8.3}", - coord[[model_i, atom_i, 0]], - coord[[model_i, atom_i, 1]], - coord[[model_i, atom_i, 2]], - ); - self.lines - .push(prefix[atom_i].clone() + &coord_string + &suffix[atom_i]); - } + self.write_atom_records(model, hybrid36)?; if is_multi_model { self.lines.push(String::from("ENDMDL")); } } - // Write CONECT records if bonds are present - let bonds = atoms.getattr("bonds")?; - if !bonds.is_none() { - // Bond type is unused since PDB does not support bond orders - let (bonds_array, _bond_types): (PyReadonlyArray2<'_, i32>, PyReadonlyArray2<'_, i8>) = - bonds.call_method0("get_all_bonds")?.extract()?; - let atom_id_for_bonds: Bound<'_, PyArray1> = match &atom_id_numpy { - Some(arr) => arr.as_array().to_owned().into_pyarray(py), - // If atom_id annotation is not present, generate sequential IDs - None => PyArray1::from_iter(py, (1..=coord.shape()[1] as i64).collect::>()), - }; - self.write_bonds(bonds_array, atom_id_for_bonds.readonly())?; - } + // The CONECT records are taken from the first model, as a PDB file + // can store only a single connectivity for all models + self.write_conect_records(py, &models[0])?; self._index_models_and_atoms(); Ok(()) @@ -1083,7 +963,11 @@ impl PDBFile { } /// Write the `CRYST1` record to this [`PDBFile`] based on the given box vectors. - fn write_box(&mut self, py: Python<'_>, box_vectors: Bound<'_, PyAny>) -> PyResult<()> { + fn write_cryst_record( + &mut self, + py: Python<'_>, + box_vectors: Bound<'_, PyAny>, + ) -> PyResult<()> { let struc = PyModule::import(py, "biotite.structure")?; let (len_a, len_b, len_c, alpha, beta, gamma): (f64, f64, f64, f64, f64, f64) = struc .getattr("unitcell_from_vectors")? @@ -1101,19 +985,30 @@ impl PDBFile { Ok(()) } - /// Write `CONECT` records to this [`PDBFile`] based on the given `bonds` - /// array containing indices pointing to bonded atoms in the `AtomArray`. - /// The `atom_id` annotation array is required to map the atom IDs in `CONECT` records - /// to atom indices. - fn write_bonds<'py>( - &mut self, - bonds: PyReadonlyArray2<'py, i32>, - atom_id: PyReadonlyArray1<'py, i64>, - ) -> PyResult<()> { - let bonds = bonds.as_array(); - let atom_id = atom_id.as_array(); + /// Write the ``CONECT`` records for the bonds of the given model, if any. + /// + /// The `atom_id` annotation is used to map the atom indices in the bonds + /// to the atom IDs in the ``CONECT`` records; if it is not present, + /// sequential IDs are generated. + fn write_conect_records(&mut self, py: Python<'_>, model: &Bound<'_, PyAny>) -> PyResult<()> { + let bonds = model.getattr("bonds")?; + if bonds.is_none() { + return Ok(()); + } + // Bond type is unused since PDB does not support bond orders + let (bonds_array, _bond_types): (PyReadonlyArray2<'_, i32>, PyReadonlyArray2<'_, i8>) = + bonds.call_method0("get_all_bonds")?.extract()?; + let atom_id_array: Bound<'_, PyArray1> = if model.hasattr("atom_id")? { + get_annotation_as_type(model, "atom_id", "int64")? + } else { + // If the atom_id annotation is not present, generate sequential IDs + let n_atoms: usize = model.call_method0("array_length")?.extract()?; + PyArray1::from_iter(py, 1..=n_atoms as i64) + }; + let atom_id_array = atom_id_array.readonly(); + let atom_id = atom_id_array.as_array(); - for (center_i, bonded_indices) in bonds.outer_iter().enumerate() { + for (center_i, bonded_indices) in bonds_array.as_array().outer_iter().enumerate() { let mut n_added: usize = 0; let mut line: String = String::new(); for bonded_i in bonded_indices.iter() { @@ -1149,6 +1044,105 @@ impl PDBFile { } Ok(()) } + + /// Write the ``ATOM``/``HETATM`` records for a single `AtomArray` model. + /// The ``MODEL``/``ENDMDL`` records are written by the caller. + fn write_atom_records(&mut self, model: &Bound<'_, PyAny>, hybrid36: bool) -> PyResult<()> { + let chain_id = get_annotation_as_strings(model, "chain_id")?; + let res_id_numpy = get_annotation_as_type(model, "res_id", "int64")?.readonly(); + let res_id: ArrayView1 = res_id_numpy.as_array(); + let ins_code = get_annotation_as_strings(model, "ins_code")?; + let res_name = get_annotation_as_strings(model, "res_name")?; + let hetero_numpy = get_annotation_as_type(model, "hetero", "bool")?.readonly(); + let hetero: ArrayView1 = hetero_numpy.as_array(); + let atom_name = get_annotation_as_strings(model, "atom_name")?; + let element = get_annotation_as_strings(model, "element")?; + + let atom_id_numpy: Option> = if model.hasattr("atom_id")? { + Some(get_annotation_as_type(model, "atom_id", "int64")?.readonly()) + } else { + None + }; + let atom_id = atom_id_numpy.as_ref().map(|arr| arr.as_array()); + let b_factor_numpy: Option> = if model.hasattr("b_factor")? { + Some(get_annotation_as_type(model, "b_factor", "float64")?.readonly()) + } else { + None + }; + let b_factor = b_factor_numpy.as_ref().map(|arr| arr.as_array()); + let occupancy_numpy: Option> = if model.hasattr("occupancy")? { + Some(get_annotation_as_type(model, "occupancy", "float64")?.readonly()) + } else { + None + }; + let occupancy = occupancy_numpy.as_ref().map(|arr| arr.as_array()); + let charge_numpy: Option> = if model.hasattr("charge")? { + Some(get_annotation_as_type(model, "charge", "int64")?.readonly()) + } else { + None + }; + let charge = charge_numpy.as_ref().map(|arr| arr.as_array()); + + // Create a binding first to avoid borrowing a dropped array + let coord_array = model + .getattr("coord")? + .extract::>>()? + .readonly(); + let coord: ArrayView2 = coord_array + .as_array() + .into_dimensionality::() + .map_err(|_| { + biotite::BadStructureError::new_err("Invalid dimensions of coordinates") + })?; + + for i in 0..coord.shape()[0] { + let element_i = &element[i]; + let atom_name_i = &atom_name[i]; + + let raw_atom_id_i = atom_id.as_ref().map_or((i + 1) as i64, |arr| arr[i]); + let atom_id_i = if hybrid36 { + encode_hybrid36(raw_atom_id_i, 5)? + } else { + truncate_id(raw_atom_id_i, 99999).to_string() + }; + + let res_id_i = if hybrid36 { + encode_hybrid36(res_id[i], 4)? + } else { + truncate_id(res_id[i], 9999).to_string() + }; + + self.lines.push(format!( + "{:6}{:>5} {:4} {:>3} {:1}{:>4}{:1} {:>8.3}{:>8.3}{:>8.3}{:>6.2}{:>6.2} {:>2}{}", + if hetero[i] { "HETATM" } else { "ATOM" }, + atom_id_i, + if element_i.len() == 1 && atom_name_i.len() < 4 { + format!(" {}", atom_name_i) + } else { + atom_name_i.to_string() + }, + res_name[i], + chain_id[i], + res_id_i, + ins_code[i], + coord[[i, 0]], + coord[[i, 1]], + coord[[i, 2]], + occupancy.as_ref().map_or(1f64, |arr| arr[i]), + b_factor.as_ref().map_or(0f64, |arr| arr[i]), + element_i, + charge.as_ref().map_or(String::from(" "), |arr| { + let c = arr[i]; + match c.cmp(&0) { + Ordering::Greater => format!("{:1}+", c), + Ordering::Less => format!("{:1}-", -c), + Ordering::Equal => String::from(" "), + } + }), + )); + } + Ok(()) + } } // Get an `AtomArray` annotation as a NumPy array with the given dtype diff --git a/tests/structure/io/test_gro.py b/tests/structure/io/test_gro.py index 806c65a0d..1794ca2a4 100644 --- a/tests/structure/io/test_gro.py +++ b/tests/structure/io/test_gro.py @@ -160,3 +160,27 @@ def test_gro_no_box(): # Assert no box with 0 dimension assert s.box is None + + +def test_multiple_atom_array(): + """ + Setting the structure with an :class:`AtomArrayStack` and with an equivalent list of + :class:`AtomArray` objects should result in the same file. + """ + stack = gro.GROFile.read(join(data_dir("structure"), "1l2y.gro")).get_structure() + + ref_file = gro.GROFile() + ref_file.set_structure(stack) + + test_file = gro.GROFile() + test_file.set_structure(list(stack)) + + # The header line of each model contains a timestamp, which may + # differ between the two files -> ignore these lines in comparison + ref_lines = [ + line for line in ref_file.lines if not line.startswith("Generated by Biotite") + ] + test_lines = [ + line for line in test_file.lines if not line.startswith("Generated by Biotite") + ] + assert test_lines == ref_lines diff --git a/tests/structure/io/test_pdb.py b/tests/structure/io/test_pdb.py index 67cb3b632..5c5e91bdd 100644 --- a/tests/structure/io/test_pdb.py +++ b/tests/structure/io/test_pdb.py @@ -603,3 +603,19 @@ def test_hetatm_intra_residue_bonds(): actual_bonds = structure.bonds.as_array() np.testing.assert_array_equal(actual_bonds, expected_bonds) + + +def test_multiple_atom_array(): + """ + Setting the structure with an :class:`AtomArrayStack` and with an equivalent list of + :class:`AtomArray` objects should result in the same file. + """ + stack = pdb.PDBFile.read(join(data_dir("structure"), "1l2y.pdb")).get_structure() + + ref_file = pdb.PDBFile() + ref_file.set_structure(stack) + + test_file = pdb.PDBFile() + test_file.set_structure(list(stack)) + + assert test_file.lines == ref_file.lines diff --git a/tests/structure/io/test_pdbx.py b/tests/structure/io/test_pdbx.py index 47eee810d..c68f00eb6 100644 --- a/tests/structure/io/test_pdbx.py +++ b/tests/structure/io/test_pdbx.py @@ -1188,3 +1188,27 @@ def test_get_structure_missing_ins_code(format): assert atoms.array_length() > 0 assert np.all(atoms.ins_code == "") + + +@pytest.mark.parametrize("format", ["cif", "bcif"]) +def test_multiple_atom_array(format): + """ + Setting the structure with an :class:`AtomArrayStack` and with an equivalent list of + :class:`AtomArray` objects should result in an equivalent output file. + """ + base_path = join(data_dir("structure"), "1gya") + if format == "cif": + File = pdbx.CIFFile + else: + File = pdbx.BinaryCIFFile + stack = pdbx.get_structure(File.read(f"{base_path}.{format}"), include_bonds=True) + + ref_file = File() + pdbx.set_structure(ref_file, stack) + + test_file = File() + pdbx.set_structure(test_file, list(stack)) + + assert pdbx.get_structure(test_file, include_bonds=True) == pdbx.get_structure( + ref_file, include_bonds=True + )