diff --git a/molecularnodes/entities/molecule/array.py b/molecularnodes/entities/molecule/array.py new file mode 100644 index 00000000..8f787226 --- /dev/null +++ b/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 diff --git a/molecularnodes/entities/molecule/base.py b/molecularnodes/entities/molecule/base.py index 5ea47f14..0af7ee52 100644 --- a/molecularnodes/entities/molecule/base.py +++ b/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: """ diff --git a/molecularnodes/entities/molecule/ui.py b/molecularnodes/entities/molecule/ui.py index 8fde2c75..d368e08e 100644 --- a/molecularnodes/entities/molecule/ui.py +++ b/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. diff --git a/tests/test_load_biotite.py b/tests/test_load_biotite.py new file mode 100644 index 00000000..814ffaaf --- /dev/null +++ b/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