AtomArray Direct Loading

molecularnodes/entities/molecule/array.py

@@ -0,0 +1,52 @@
+import uuid
+from biotite.structure import AtomArray
+
+from .base import Molecule
+from .pdb import _comp_secondary_structure
+from ..base import EntityType
+
+
+class Array(Molecule):
+    def __init__(self, array: AtomArray, name="AtomArray"):
+        self.file_path = None
+        self.file = "ARRAY_LOADED_DIRECTLY"
+        self._frames_collection = None
+        self._entity_type = EntityType.MOLECULE
+        self._assemblies = lambda: None
+        self._uuid = str(uuid.uuid4())
+        self.array = self._validate_structure(array)
+        self.n_atoms = self.array.array_length()    # note: I am not sure if this is consistent
+        self.create_object(name=name)               # Create the Blender object
+
+    def read(self, file_path):
+        pass
+
+    def _validate_structure(self, array: AtomArray):
+        # TODO: implement entity ID, sec_struct for PDB files
+        # extra_fields = ["b_factor", "occupancy", "charge", "atom_id"]
+        #
+        # https://github.com/biotite-dev/biotite/blob/main/src/biotite/structure/io/pdb/file.py#L331
+        # for field in extra_fields:
+        #     if field == "atom_id":
+        #         # Copy is necessary to avoid double masking in
+        #         # later altloc ID filtering
+        #         # array.set_annotation("atom_id", atom_id.copy())
+        #         pass
+        #     elif field == "charge":
+        #         charge = np.array(charge_raw)
+        #         array.set_annotation(
+        #             "charge", np.where(charge == "  ", "0", charge).astype(int)
+        #         )
+        #     elif field == "occupancy":
+        #         array.set_annotation("occupancy", occupancy)
+        #     elif field == "b_factor":
+        #         array.set_annotation("b_factor", b_factor)
+        #     else:
+        #         raise ValueError(f"Unknown extra field: {field}")
+
+        sec_struct = _comp_secondary_structure(array)
+        array.set_annotation("sec_struct", sec_struct)
+        return array
+
+    def _assemblies(self):
+        return None

molecularnodes/entities/molecule/base.py

@@ -5,7 +5,6 @@
 from pathlib import Path
 from typing import Optional, Tuple, Union
 import json
-
 import biotite.structure as struc
 import bpy
 import numpy as np
@@ -65,6 +64,12 @@ def __init__(self, file_path: Union[str, Path, io.BytesIO]):
         self._frames_collection: str | None
         self._entity_type = EntityType.MOLECULE
 
+
+    @classmethod
+    def from_array(cls, array: struc.AtomArray, name: str = "FromArray"):
+        from .array import Array
+        return Array(array, name=name)
+
     @property
     def frames(self) -> bpy.types.Collection:
         """

molecularnodes/entities/molecule/ui.py

@@ -21,7 +21,6 @@
 
 from ...pref import addon_preferences
 
-
 def parse(filepath) -> Molecule:
     """Parse a molecular structure file into a Molecule object.
 

tests/test_load_biotite.py

@@ -0,0 +1,29 @@
+import bpy
+import numpy as np
+import pytest
+import itertools
+import molecularnodes as mn
+import databpy as db
+from biotite.structure import io
+from biotite.structure import AtomArray, AtomArrayStack
+from .constants import data_dir, codes, attributes
+from .utils import NumpySnapshotExtension
+
+
+
+
+def test_loading():
+    test_file =  data_dir / "1f2n.bcif"
+    arr = io.load_structure(test_file, template=None)
+
+    assert isinstance(arr, AtomArray)
+
+    # use the class method
+    mol = mn.entities.Molecule.from_array(arr)
+    assert isinstance(mol,  mn.entities.Molecule)
+    assert mol.file_path == None
+    assert mol.file == "ARRAY_LOADED_DIRECTLY"
+    assert mol._frames_collection == None
+    #assert mol._entity_type == EntityType.MOLECULE
+    assert mol._assemblies() == None
+    assert mol.n_atoms == 4730