Collapse classes for Molecule

docs/api_introduction.qmd

@@ -31,18 +31,18 @@ mn.template.install()
 
 
 # create a 'Molecule' object, by fetching a structure and parsing it into the scene
-mol = mn.fetch("6N2Y", style = 'ribbon')
+mol = mn.Molecule.fetch("6N2Y").add_style('ribbon')
 ```
 
 ## The Molecule Object
 
 The `Molecule` object has the original data, as well as the Blender object associated with.
 
-The different methods that are associated mostly interact with the Blender object, which is accessible via the `mol.object`, and the data is accessible via `mol.array`, which is the `biotite.AtomArrayStack` object.
+The different methods that are associated mostly interact with the Blender object, which is accessible via the `mol.object`, and the data is accessible via `mol.atom_array`, which is a `biotite.AtomArray` object.
 ```{python}
 print(f"{len(mol)=}")
 print(f"{mol.object=}")
-print(f"{mol.array[0][:10]=}")
+print(f"{mol.atom_array[0][:10]=}")
 ```
 
 ```{python}

molecularnodes/__init__.py

@@ -1,7 +1,8 @@
 from . import assets, blender, color, session, ui
 from .assets import template
-from .entities import fetch, parse
+from .entities import Molecule, Trajectory
 from .ui.addon import register, unregister
+from . import download
 
 try:
     from .scene import Canvas

molecularnodes/blender/nodes.py

@@ -409,6 +409,10 @@ def create_starting_node_tree(
     node_input.location = [0, 0]
     node_output.location = [700, 0]
 
+    if style is None:
+        link(node_input.outputs[0], node_output.inputs[0])
+        return tree
+
     node_style = add_custom(tree, styles_mapping[style], [450, 0], material=material)
     link(node_style.outputs[0], node_output.inputs[0])
     link(node_input.outputs[0], node_style.inputs[0])

molecularnodes/download.py

@@ -56,6 +56,7 @@ def download(code, format="cif", cache=CACHE_DIR, database="rcsb"):
     ValueError
         If the specified format is not supported.
     """
+    format = format.strip(".")
     supported_formats = ["cif", "pdb", "bcif"]
     if format not in supported_formats:
         raise ValueError(f"File format '{format}' not in: {supported_formats=}")

molecularnodes/entities/__init__.py

@@ -1,9 +1,3 @@
-from . import molecule, trajectory
-from .base import EntityType, MolecularEntity
 from .ensemble import CellPack, Ensemble, StarFile
-from .molecule import BCIF, CIF, PDB, SDF, Molecule
-from .molecule.pdb import PDB
-from .molecule.pdbx import BCIF, CIF
-from .molecule.sdf import SDF
-from .molecule.io import fetch, load_local, parse
+from .molecule.base import Molecule
 from .trajectory import OXDNA, Trajectory

molecularnodes/entities/base.py

@@ -23,6 +23,7 @@ def __init__(self) -> None:
         super().__init__(obj=None)
         self._entity_type: EntityType
         self._register_with_session()
+        self._world_scale = 0.01
 
     @property
     def bob(self) -> BlenderObject:

molecularnodes/entities/ensemble/cellpack.py

@@ -6,9 +6,11 @@
 from databpy import AttributeTypes, BlenderObject, store_named_attribute
 
 from .base import Ensemble
+from ..molecule.base import Molecule
 from ... import blender as bl
 from ... import color
-from ..molecule.base import _create_object
+
+# from ..molecule.base import _create_object
 from .reader import CellPackReader
 from biotite.structure import AtomArray
 
@@ -19,9 +21,12 @@ def __init__(self, file_path):
         self.file_type = self._file_type()
         self.file = CellPackReader(file_path)
         self.file.get_molecules()
-        self.transformations = self.file.assemblies(as_array=True)
+        self.transformations = self.file.get_assemblies()
         self.color_entity = {}
         self._color_palette_path = Path(file_path).parent / "color_palette.json"
+        self.object = self._create_data_object(name=f"{Path(file_path).name}")
+        self._create_object_instances(name=self.object.name, node_setup=False)
+        self._setup_node_tree(fraction=0.1)
         # self._setup_colors()
 
     def _setup_colors(self):
@@ -44,19 +49,6 @@ def _setup_colors(self):
     def molecules(self):
         return self.file.molecules
 
-    def create_object(
-        self,
-        name="CellPack",
-        node_setup: bool = True,
-        world_scale: float = 0.01,
-        fraction: float = 1.0,
-        simplify=False,
-    ):
-        self.object = self._create_data_object(name=f"{name}")
-        self._create_object_instances(name=self.object.name, node_setup=node_setup)
-        self._setup_node_tree(fraction=fraction)
-        return self.object
-
     def _file_type(self):
         return Path(self.file_path).suffix.strip(".")
 
@@ -87,8 +79,8 @@ def _create_object_instances(
             array = self.molecules[mol_id]
             chain_name = array.asym_id[0]
 
-            obj, coll_none = _create_object(
-                array=array,
+            obj = Molecule._create_object(
+                atom_array=array,
                 name=mol_id,
                 collection=collection,
             )

molecularnodes/entities/ensemble/reader.py

@@ -1,19 +1,17 @@
-from io import BytesIO
 from pathlib import Path
-
+from ..molecule.pdbx import PDBXReader
+from ...utils import array_quaternions_from_dict
 import numpy as np
 from biotite import InvalidFileError
 from biotite import structure as struc
 from biotite.structure.io import pdbx
 
-from ..molecule.pdbx import PDBX
-
 
 # For reading cellpack files, we override the CIFFile from biotite. The only change we
 # implement is that we add the `line.strip()` when initially getting the lines from the
 # text in the `deserialize` method. This fixes the reading of the PETWORLD files, and we
 # don't have to write-out a modifier version of the file before reading it back in
-class PetworldCIFFileReader(pdbx.CIFFile):
+class CellPackCIFFileReader(pdbx.CIFFile):
     @classmethod
     def read(cls, file):
         # File name
@@ -43,91 +41,85 @@ def deserialize(text):
         )
 
 
-class CellPackReader(PDBX):
+class CellPackReader(PDBXReader):
     def __init__(self, file_path):
-        super().__init__(file_path)
-        self._extra_annotations["asym_id"] = self._get_asym_id
-        self._extra_annotations["pdb_model_num"] = self._get_pdb_model_num
+        self._extra_annotations = {
+            "sec_struct": self._get_secondary_structure,
+            "entity_id": self._get_entity_id,
+            "asym_id": self._get_asym_id,
+            "pdb_model_num": self._get_pdb_model_num,
+        }
+        self._extra_fields = ["b_factor", "occupancy", "atom_id"]
         self.file_path = file_path
-        self.file: pdbx.BinaryCIFFile | pdbx.CIFFile = self._read()
+        self.file = self.read(file_path)
         self.n_molecules: int = pdbx.get_model_count(self.file)
-        self.molecules: dict[str, stuc.AtomArray] = {}
+        self.molecules: dict[str, struc.AtomArray] = self.get_molecules()
 
     @property
     def mol_ids(self) -> np.ndarray:
         return np.unique(list(self.molecules.keys()))
 
-    def _parse_filepath(self, file_path: Path | str | BytesIO) -> None:
-        pass
-
-    def _read(self):
-        suffix = Path(self.file_path).suffix
+    def read(self, file_path):
+        suffix = Path(file_path).suffix
         if suffix == ".bcif":
             return pdbx.BinaryCIFFile.read(self.file_path)
         elif suffix == ".cif":
-            return PetworldCIFFileReader.read(self.file_path)
+            return CellPackCIFFileReader.read(self.file_path)
         else:
             raise ValueError(f"Invalid file format: '{suffix}")
 
-    def _get_asym_id(self, array, file) -> np.ndarray:
+    @staticmethod
+    def _get_asym_id(array, file) -> np.ndarray:
         return array.chain_id
 
-    def _get_pdb_model_num(self, array, file) -> np.ndarray:
+    @staticmethod
+    def _get_pdb_model_num(array, file) -> np.ndarray:
         return (
-            list(self.file.values())[0]["atom_site"]["pdbx_PDB_model_num"]
+            list(file.values())[0]["atom_site"]["pdbx_PDB_model_num"]
             .as_array()
             .astype(int)
         )
 
-    def get_structure(
-        self,
-        extra_fields=["b_factor", "occupancy", "atom_id"],
-        bonds=True,
-        model: int | None = None,
-    ) -> struc.AtomArray:
-        array: struc.AtomArray = pdbx.get_structure(
-            self.file, model=model, extra_fields=extra_fields
-        )  # type: ignore
-        array = self.set_extra_annotations(array, self.file)
-
-        if not array.bonds and bonds:
-            array.bonds = struc.bonds.connect_via_residue_names(
-                array, inter_residue=True
-            )
-        return array
-
     @property
     def blocks(self):
         return list(self.file.values())[0]
 
-    def get_molecules(
-        self, extra_fields=["b_factor", "occupancy", "atom_id"], bonds=True
-    ):
+    def get_assemblies(self):
+        return array_quaternions_from_dict(self._assemblies())
+
+    def get_molecules(self) -> dict[str, struc.AtomArray]:
         self._is_petworld = False
 
         if "PDB_model_num" in self.blocks["pdbx_struct_assembly_gen"]:
             self._is_petworld = True
 
+        molecules = {}
+
         try:
-            array = self.get_structure(extra_fields, bonds)
+            array = self.get_structure()
             if isinstance(array, struc.AtomArrayStack):
                 array = array[0]
 
-            # if self._is_petworld:
-            #     array.set_annotation(
-            #         "chain_id", np.char.rjust(array.pdb_model_num, 4, "0")
-            #     )
-
-            self.molecules = {
+            molecules = {
                 c: array[array.chain_id == c] for c in np.unique(array.chain_id)
             }
+
         except InvalidFileError:
             self._is_petworld = True
             for i in range(self.n_molecules):
-                array = self.get_structure(extra_fields, bonds, model=i + 1)
+                array = self.get_structure(model=int(i + 1))
                 array.set_annotation("pdbx_PDB_model_num", np.repeat(i + 1, len(array)))
                 array.chain_id = array.pdbx_PDB_model_num
                 chain_name = "{}_{}".format(
                     str(i).rjust(4, "0"), str(array.chain_id[0])
                 )
-                self.molecules[chain_name] = array
+                molecules[chain_name] = array
+
+        for name, array in molecules.items():
+            array = self.set_extra_annotations(
+                self._extra_annotations, array, self.file
+            )
+            array = self.set_standard_annotations(array)
+            molecules[name] = array
+
+        return molecules