biotite 0.40.0__cp311-cp311-win_amd64.whl → 0.41.0__cp311-cp311-win_amd64.whl

biotite/structure/io/general.py

@@ -11,9 +11,10 @@ __name__ = "biotite.structure.io"
 __author__ = "Patrick Kunzmann"
 __all__ = ["load_structure", "save_structure"]
 
+import datetime
 import os.path
 import io
-from ..atoms import AtomArray, AtomArrayStack
+from ..atoms import AtomArrayStack
 
 
 def load_structure(file_path, template=None, **kwargs):
@@ -61,72 +62,83 @@ def load_structure(file_path, template=None, **kwargs):
 
     # We only need the suffix here
     _, suffix = os.path.splitext(file_path)
-    if suffix == ".pdb":
-        from .pdb import PDBFile
-        file = PDBFile.read(file_path)
-        array = file.get_structure(**kwargs)
-        return _as_single_model_if_possible(array)
-    elif suffix == ".pdbqt":
-        from .pdbqt import PDBQTFile
-        file = PDBQTFile.read(file_path)
-        array = file.get_structure(**kwargs)
-        return _as_single_model_if_possible(array)
-    elif suffix == ".cif" or suffix == ".pdbx":
-        from .pdbx import CIFFile, get_structure
-        file = CIFFile.read(file_path)
-        array = get_structure(file, **kwargs)
-        return _as_single_model_if_possible(array)
-    elif suffix == ".bcif":
-        from .pdbx import BinaryCIFFile, get_structure
-        file = BinaryCIFFile.read(file_path)
-        array = get_structure(file, **kwargs)
-        return _as_single_model_if_possible(array)
-    elif suffix == ".gro":
-        from .gro import GROFile
-        file = GROFile.read(file_path)
-        array = file.get_structure(**kwargs)
-        return _as_single_model_if_possible(array)
-    elif suffix == ".mmtf":
-        from .mmtf import MMTFFile, get_structure
-        file = MMTFFile.read(file_path)
-        array = get_structure(file, **kwargs)
-        return _as_single_model_if_possible(array)
-    elif suffix == ".npz":
-        from .npz import NpzFile
-        file = NpzFile.read(file_path)
-        array = file.get_structure(**kwargs)
-        return _as_single_model_if_possible(array)
-    elif suffix == ".mol" or suffix == ".sdf":
-        from .mol import MOLFile
-        file = MOLFile.read(file_path)
-        array = file.get_structure(**kwargs)
-        # MOL files only contain a single model
-        return array
-    elif suffix in [".trr", ".xtc", ".tng", ".dcd", ".netcdf"]:
-        if template is None:
-            raise TypeError("Template must be specified for trajectory files")
-        # filter template for atom ids if it is an unfiltered template
-        if "atom_i" in kwargs and template.shape[-1] != len(kwargs["atom_i"]):
-            template = template[..., kwargs["atom_i"]]
-        from .trr import TRRFile
-        from .xtc import XTCFile
-        from .tng import TNGFile
-        from .dcd import DCDFile
-        from .netcdf import NetCDFFile
-        if suffix == ".trr":
-            traj_file_cls = TRRFile
-        if suffix == ".xtc":
-            traj_file_cls = XTCFile
-        if suffix == ".tng":
-            traj_file_cls = TNGFile
-        if suffix == ".dcd":
-            traj_file_cls = DCDFile
-        if suffix == ".netcdf":
-            traj_file_cls = NetCDFFile
-        file = traj_file_cls.read(file_path, **kwargs)
-        return file.get_structure(template)
-    else:
-        raise ValueError(f"Unknown file format '{suffix}'")
+    match suffix:
+        case ".pdb":
+            from .pdb import PDBFile
+            file = PDBFile.read(file_path)
+            array = file.get_structure(**kwargs)
+            return _as_single_model_if_possible(array)
+        case ".pdbqt":
+            from .pdbqt import PDBQTFile
+            file = PDBQTFile.read(file_path)
+            array = file.get_structure(**kwargs)
+            return _as_single_model_if_possible(array)
+        case ".cif" | ".pdbx":
+            from .pdbx import CIFFile, get_structure
+            file = CIFFile.read(file_path)
+            array = get_structure(file, **kwargs)
+            return _as_single_model_if_possible(array)
+        case ".bcif":
+            from .pdbx import BinaryCIFFile, get_structure
+            file = BinaryCIFFile.read(file_path)
+            array = get_structure(file, **kwargs)
+            return _as_single_model_if_possible(array)
+        case ".gro":
+            from .gro import GROFile
+            file = GROFile.read(file_path)
+            array = file.get_structure(**kwargs)
+            return _as_single_model_if_possible(array)
+        case ".mmtf":
+            from .mmtf import MMTFFile, get_structure
+            file = MMTFFile.read(file_path)
+            array = get_structure(file, **kwargs)
+            return _as_single_model_if_possible(array)
+        case ".npz":
+            from .npz import NpzFile
+            file = NpzFile.read(file_path)
+            array = file.get_structure(**kwargs)
+            return _as_single_model_if_possible(array)
+        case ".mol":
+            from .mol import MOLFile
+            file = MOLFile.read(file_path)
+            array = file.get_structure(**kwargs)
+            # MOL and SDF files only contain a single model
+            return array
+        case ".sdf" | ".sd":
+            from .mol import SDFile, get_structure
+            file = SDFile.read(file_path)
+            array = get_structure(file, **kwargs)
+            return array
+        case ".trr" | ".xtc" | ".tng" | ".dcd" | ".netcdf":
+            if template is None:
+                raise TypeError(
+                    "Template must be specified for trajectory files"
+                )
+            # Filter template for atom ids, if an unfiltered template
+            if (
+                "atom_i" in kwargs
+                and template.shape[-1] != len(kwargs["atom_i"])
+            ):
+                template = template[..., kwargs["atom_i"]]
+            from .trr import TRRFile
+            from .xtc import XTCFile
+            from .tng import TNGFile
+            from .dcd import DCDFile
+            from .netcdf import NetCDFFile
+            if suffix == ".trr":
+                traj_file_cls = TRRFile
+            if suffix == ".xtc":
+                traj_file_cls = XTCFile
+            if suffix == ".tng":
+                traj_file_cls = TNGFile
+            if suffix == ".dcd":
+                traj_file_cls = DCDFile
+            if suffix == ".netcdf":
+                traj_file_cls = NetCDFFile
+            file = traj_file_cls.read(file_path, **kwargs)
+            return file.get_structure(template)
+        case unknown_suffix:
+            raise ValueError(f"Unknown file format '{unknown_suffix}'")
 
 
 def save_structure(file_path, array, **kwargs):
@@ -155,67 +167,76 @@ def save_structure(file_path, array, **kwargs):
     """
     # We only need the suffix here
     _, suffix = os.path.splitext(file_path)
-    if suffix == ".pdb":
-        from .pdb import PDBFile
-        file = PDBFile()
-        file.set_structure(array, **kwargs)
-        file.write(file_path)
-    elif suffix == ".pdbqt":
-        from .pdbqt import PDBQTFile
-        file = PDBQTFile()
-        file.set_structure(array, **kwargs)
-        file.write(file_path)
-    elif suffix == ".cif" or suffix == ".pdbx":
-        from .pdbx import CIFFile, set_structure
-        file = CIFFile()
-        set_structure(file, array, **kwargs)
-        file.write(file_path)
-    elif suffix == ".bcif":
-        from .pdbx import BinaryCIFFile, set_structure
-        file = BinaryCIFFile()
-        set_structure(file, array, **kwargs)
-        file.write(file_path)
-    elif suffix == ".gro":
-        from .gro import GROFile
-        file = GROFile()
-        file.set_structure(array, **kwargs)
-        file.write(file_path)
-    elif suffix == ".mmtf":
-        from .mmtf import MMTFFile, set_structure
-        file = MMTFFile()
-        set_structure(file, array, **kwargs)
-        file.write(file_path)
-    elif suffix == ".npz":
-        from .npz import NpzFile
-        file = NpzFile()
-        file.set_structure(array, **kwargs)
-        file.write(file_path)
-    elif suffix == ".mol" or suffix == ".sdf":
-        from .mol import MOLFile
-        file = MOLFile()
-        file.set_structure(array, **kwargs)
-        file.write(file_path)
-    elif suffix in [".trr", ".xtc", ".tng", ".dcd", ".netcdf"]:
-        from .trr import TRRFile
-        from .xtc import XTCFile
-        from .tng import TNGFile
-        from .dcd import DCDFile
-        from .netcdf import NetCDFFile
-        if suffix == ".trr":
-            traj_file_cls = TRRFile
-        if suffix == ".xtc":
-            traj_file_cls = XTCFile
-        if suffix == ".tng":
-            traj_file_cls = TNGFile
-        if suffix == ".dcd":
-            traj_file_cls = DCDFile
-        if suffix == ".netcdf":
-            traj_file_cls = NetCDFFile
-        file = traj_file_cls()
-        file.set_structure(array, **kwargs)
-        file.write(file_path)
-    else:
-        raise ValueError(f"Unknown file format '{suffix}'")
+    match suffix:
+        case ".pdb":
+            from .pdb import PDBFile
+            file = PDBFile()
+            file.set_structure(array, **kwargs)
+            file.write(file_path)
+        case ".pdbqt":
+            from .pdbqt import PDBQTFile
+            file = PDBQTFile()
+            file.set_structure(array, **kwargs)
+            file.write(file_path)
+        case ".cif" | ".pdbx":
+            from .pdbx import CIFFile, set_structure
+            file = CIFFile()
+            set_structure(file, array, **kwargs)
+            file.write(file_path)
+        case ".bcif":
+            from .pdbx import BinaryCIFFile, set_structure
+            file = BinaryCIFFile()
+            set_structure(file, array, **kwargs)
+            file.write(file_path)
+        case ".gro":
+            from .gro import GROFile
+            file = GROFile()
+            file.set_structure(array, **kwargs)
+            file.write(file_path)
+        case ".mmtf":
+            from .mmtf import MMTFFile, set_structure
+            file = MMTFFile()
+            set_structure(file, array, **kwargs)
+            file.write(file_path)
+        case ".npz":
+            from .npz import NpzFile
+            file = NpzFile()
+            file.set_structure(array, **kwargs)
+            file.write(file_path)
+        case ".mol":
+            from .mol import MOLFile
+            file = MOLFile()
+            file.set_structure(array, **kwargs)
+            file.header = _mol_header()
+            file.write(file_path)
+        case ".sdf" | ".sd":
+            from .mol import SDFile, SDRecord, set_structure
+            record = SDRecord()
+            record.set_structure(array, **kwargs)
+            record.header = _mol_header()
+            file = SDFile({"Molecule": record})
+            file.write(file_path)
+        case ".trr" | ".xtc" | ".tng" | ".dcd" | ".netcdf":
+            from .trr import TRRFile
+            from .xtc import XTCFile
+            from .tng import TNGFile
+            from .dcd import DCDFile
+            from .netcdf import NetCDFFile
+            if suffix == ".trr":
+                traj_file_cls = TRRFile
+            if suffix == ".xtc":
+                traj_file_cls = XTCFile
+            if suffix == ".tng":
+                traj_file_cls = TNGFile
+            if suffix == ".dcd":
+                traj_file_cls = DCDFile
+            if suffix == ".netcdf":
+                traj_file_cls = NetCDFFile
+            file = traj_file_cls()
+            file.set_structure(array, **kwargs)
+            file.write(file_path)
+        case unknown_suffix:
+            raise ValueError(f"Unknown file format '{unknown_suffix}'")
 
 
 def _as_single_model_if_possible(atoms):
@@ -226,40 +247,11 @@ def _as_single_model_if_possible(atoms):
         return atoms
 
 
-# Helper function to estimate elements from atom names
-_elements = [elem.upper() for elem in
-["H", "He", "Li", "Be", "B", "C", "N", "O", "F", "Ne", "Na", "Mg",
-"Al", "Si", "P", "S", "Cl", "Ar", "K", "Ca", "Sc", "Ti", "V", "Cr", "Mn", "Fe",
-"Co", "Ni", "Cu", "Zn", "Ga", "Ge", "As", "Se", "Br", "Kr", "Rb", "Sr", "Y",
-"Zr", "Nb", "Mo", "Tc", "Ru", "Rh", "Pd", "Ag", "Cd", "In", "Sn", "Sb", "Te",
-"I", "Xe", "Cs", "Ba", "La", "Ce", "Pr", "Nd", "Pm", "Sm", "Eu", "Gd", "Tb",
-"Dy", "Ho", "Er", "Tm", "Yb", "Lu", "Hf", "Ta", "W", "Re", "Os", "Ir", "Pt",
-"Au", "Hg", "Tl", "Pb", "Bi", "Po", "At", "Rn", "Fr", "Ra", "Ac", "Th", "Pa",
-"U", "Np", "Pu", "Am", "Cm", "Bk", "Cf", "Es", "Fm", "Md", "No", "Lr", "Rf",
-"Db", "Sg", "Bh", "Hs", "Mt", "Ds", "Rg", "Cn", "Nh", "Fl", "Mc", "Lv", "Ts",
-"Og"]
-]
-def _guess_element(atom_name):
-    # remove digits (1H -> H)
-    elem = "".join([i for i in atom_name if not i.isdigit()])
-    elem = elem.upper()
-    if len(elem) == 0:
-        return ""
-
-    # Some often used elements for biomolecules
-    if elem.startswith("C") or elem.startswith("N") or \
-        elem.startswith("O") or elem.startswith("S") or \
-        elem.startswith("H"):
-        return elem[0]
-
-    # Exactly match element abbreviations
-    try:
-        return _elements[_elements.index(elem[:2])]
-    except ValueError:
-        try:
-            return _elements[_elements.index(elem[0])]
-        except ValueError:
-            pass
-
-    return ""
-
+def _mol_header():
+    from .mol import Header
+    return Header(
+        mol_name="Molecule",
+        program="Biotite",
+        time=datetime.datetime.now(),
+        dimensions="3D",
+    )

biotite/structure/io/gro/file.py

@@ -10,7 +10,7 @@ import numpy as np
 from ...atoms import AtomArray, AtomArrayStack
 from ...box import is_orthogonal
 from ....file import TextFile, InvalidFileError
-from ..general import _guess_element as guess_element
+from ...repair import infer_elements
 from ...error import BadStructureError
 import copy
 from datetime import datetime
@@ -38,7 +38,7 @@ class GROFile(TextFile):
     --------
     Load a `\\*.gro` file, modify the structure and save the new
     structure into a new file:
-    
+
     >>> import os.path
     >>> file = GROFile.read(os.path.join(path_to_structures, "1l2y.gro"))
     >>> array_stack = file.get_structure()
@@ -46,7 +46,7 @@ class GROFile(TextFile):
     >>> file = GROFile()
     >>> file.set_structure(array_stack_mod)
     >>> file.write(os.path.join(path_to_directory, "1l2y_mod.gro"))
-    
+
     """
     def get_model_count(self):
         """
@@ -68,7 +68,7 @@ class GROFile(TextFile):
         """
         Get an :class:`AtomArray` or :class:`AtomArrayStack` from the
         GRO file.
-        
+
         Parameters
         ----------
         model : int, optional
@@ -80,13 +80,13 @@ class GROFile(TextFile):
             If this parameter is omitted, an :class:`AtomArrayStack`
             containing all models will be returned, even if the
             structure contains only one model.
-        
+
         Returns
         -------
         array : AtomArray or AtomArrayStack
             The return type depends on the `model` parameter.
         """
-        
+
         def get_atom_line_i(model_start_i, model_atom_counts):
             """
             Helper function to get the indices of all atoms for a model
@@ -94,7 +94,7 @@ class GROFile(TextFile):
             return np.arange(
                 model_start_i+1, model_start_i+1+model_atom_counts
             )
-        
+
         def set_box_dimen(box_param):
             """
             Helper function to create the box vectors from the values
@@ -104,7 +104,7 @@ class GROFile(TextFile):
             ----------
             box_param : list of float
                 The box dimensions in the GRO file.
-            
+
             Returns
             -------
             box_vectors : ndarray, dtype=float, shape=(3,3)
@@ -171,7 +171,7 @@ class GROFile(TextFile):
             array.res_id[i] = int(line[0:5])
             array.res_name[i] = line[5:10].strip()
             array.atom_name[i] = line[10:15].strip()
-            array.element[i] = guess_element(line[10:15].strip())
+        array.element = infer_elements(array.atom_name)
 
         # Fill in coordinates and boxes
         if isinstance(array, AtomArray):
@@ -186,7 +186,7 @@ class GROFile(TextFile):
             box_i = atom_i[-1] + 1
             box_param = [float(e)*10 for e in self.lines[box_i].split()]
             array.box = set_box_dimen(box_param)
-        
+
         elif isinstance(array, AtomArrayStack):
             for m in range(len(model_start_i)):
                 atom_i = get_atom_line_i(
@@ -204,18 +204,18 @@ class GROFile(TextFile):
                 # Create a box in the stack if not already existing
                 # and the box is not a dummy
                 if box is not None:
-                    if array.box is None: 
+                    if array.box is None:
                         array.box = np.zeros((array.stack_depth(), 3, 3))
                     array.box[m] = box
-                    
+
         return array
 
-            
+
     def set_structure(self, array):
         """
         Set the :class:`AtomArray` or :class:`AtomArrayStack` for the
         file.
-        
+
         Parameters
         ----------
         array : AtomArray or AtomArrayStack
@@ -235,7 +235,7 @@ class GROFile(TextFile):
             ----------
             array : AtomArray
                 The atom array to get the box dimensions from.
-            
+
             Returns
             -------
             box : str
@@ -259,7 +259,7 @@ class GROFile(TextFile):
                         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: