importing via UI successfully

molecularnodes/io/parse/__init__.py

@@ -7,7 +7,7 @@
 # from .cif import CIF
 from .pdb import PDB
 from .cellpack import CellPack
-from .star import StarFile
+from .star import StarFile, NDJSON
 from .sdf import SDF
 from .mda import MDAnalysisSession
 from .mrc import MRC

molecularnodes/io/parse/star.py

@@ -2,7 +2,7 @@
 import bpy
 import json
 
-
+from mathutils import Matrix
 from .ensemble import Ensemble
 from ... import blender as bl
 
@@ -230,22 +230,64 @@ def create_model(self, name='StarFileObject', node_setup=True, world_scale=0.01)
         return blender_object
 
 
-def read_ndjson(file):
-    with open(file, 'r') as f:
-        lines = f.readlines()
-
-    has_rotation = bool(json.loads(lines[0]).get('xyz_rotation_matrix'))
-
-    arr = np.zeros((len(lines), 4, 4), float)
-
-    for i, line in enumerate(lines):
-        matrix = np.identity(4, float)
-        data = json.loads(line)
-        pos = [data['location'][axis] for axis in 'xyz']
+class NDJSON(Ensemble):
+    def __init__(self, file_path):
+        super().__init__(file_path)
+        self.scale = 10
 
-        matrix[:3, 3] = pos
-        if has_rotation:
-            matrix[:3, :3] = data['xyz_rotation_matrix']
-        arr[i] = matrix
+    @classmethod
+    def from_ndjson(cls, file_path):
+        self = cls(file_path)
+        self.data = self._read()
+        return self
 
-    return arr
+    def _read(self):
+        with open(self.file_path, 'r') as f:
+            lines = f.readlines()
+
+        has_rotation = bool(json.loads(lines[0]).get('xyz_rotation_matrix'))
+
+        arr = np.zeros((len(lines), 4, 4), float)
+
+        for i, line in enumerate(lines):
+            matrix = np.identity(4, float)
+            data = json.loads(line)
+            pos = [data['location'][axis] for axis in 'xyz']
+
+            matrix[:3, 3] = pos
+            if has_rotation:
+                matrix[:3, :3] = data['xyz_rotation_matrix']
+            arr[i] = matrix
+
+        # returns a (n, 4, 4) matrix, where the 4x4 rotation matrix is returned for
+        # each point from the ndjson file
+        # this currently doesn't handle where there might be different points or different
+        # proteins being instanced on those points, at which point we will have to change
+        # what kind of array we are returning
+        return arr
+
+    def create_model(self, name='NewInstances', world_scale=0.01, node_setup=True):
+        n_points = len(self.data)
+        data = np.zeros((n_points, 7), float)
+
+        for i in range(n_points):
+            translation, rotation, scale = Matrix(self.data[i]).decompose()
+            data[i, :3] = translation
+            data[i, 3:] = rotation
+
+        # use the positions to create the object
+        bob = bl.obj.create_object(
+            vertices=data[:, :3] * world_scale * self.scale,
+            collection=bl.coll.mn(),
+            name=name
+        )
+        bob.mn['molecule_type'] = 'ndjson'
+
+        bl.obj.set_attribute(
+            object=bob,
+            name='rotation',
+            data=data[:, 3:],
+            type='QUATERNION'
+        )
+
+        self.object = bob

molecularnodes/io/star.py

@@ -1,6 +1,8 @@
 import bpy
+from pathlib import Path
 from . import parse
 
+
 bpy.types.Scene.MN_import_star_file_path = bpy.props.StringProperty(
     name='File',
     description='File path for the `.star` file to import.',
@@ -21,10 +23,18 @@ def load(
     node_setup=True,
     world_scale=0.01
 ):
+    suffix = Path(file_path).suffix
+    parser = {
+        '.star': parse.StarFile.from_starfile,
+        '.ndjson': parse.NDJSON.from_ndjson
+    }
 
-    ensemble = parse.StarFile.from_starfile(file_path)
-    ensemble.create_model(name=name, node_setup=node_setup,
-                          world_scale=world_scale)
+    ensemble = parser[suffix](file_path)
+    ensemble.create_model(
+        name=name,
+        node_setup=node_setup,
+        world_scale=world_scale
+    )
 
     return ensemble