Read .ndjson CryoET Files

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

@@ -1,8 +1,12 @@
 import numpy as np
 import bpy
+import json
+
+from mathutils import Matrix
 from .ensemble import Ensemble
 from ... import blender as bl
 
+
 @bpy.app.handlers.persistent
 def _rehydrate_ensembles(scene):
     for obj in bpy.data.objects:
@@ -13,11 +17,11 @@ def _rehydrate_ensembles(scene):
                     bpy.types.Scene.MN_starfile_ensembles = []
                 bpy.types.Scene.MN_starfile_ensembles.append(ensemble)
 
+
 class StarFile(Ensemble):
     def __init__(self, file_path):
         super().__init__(file_path)
-
-
+
     @classmethod
     def from_starfile(cls, file_path):
         self = cls(file_path)
@@ -28,7 +32,7 @@ def from_starfile(cls, file_path):
         self._create_mn_columns()
         self.n_images = self._n_images()
         return self
-    
+
     @classmethod
     def from_blender_object(cls, blender_object):
         import bpy
@@ -42,9 +46,9 @@ def from_blender_object(cls, blender_object):
         self.current_image = -1
         self._create_mn_columns()
         self.n_images = self._n_images()
-        bpy.app.handlers.depsgraph_update_post.append(self._update_micrograph_texture)
+        bpy.app.handlers.depsgraph_update_post.append(
+            self._update_micrograph_texture)
         return self
-
 
     def _read(self):
         import starfile
@@ -69,7 +73,8 @@ def _create_mn_columns(self):
 
         # Get absolute position and orientations
         if self.star_type == 'relion':
-            df = self.data['particles'].merge(self.data['optics'], on='rlnOpticsGroup')
+            df = self.data['particles'].merge(
+                self.data['optics'], on='rlnOpticsGroup')
 
             # get necessary info from dataframes
             # Standard cryoEM starfile don't have rlnCoordinateZ. If this column is not present
@@ -78,7 +83,7 @@ def _create_mn_columns(self):
                 df['rlnCoordinateZ'] = 0
 
             self.positions = df[['rlnCoordinateX', 'rlnCoordinateY',
-                      'rlnCoordinateZ']].to_numpy()
+                                 'rlnCoordinateZ']].to_numpy()
             pixel_size = df['rlnImagePixelSize'].to_numpy().reshape((-1, 1))
             self.positions = self.positions * pixel_size
             shift_column_names = ['rlnOriginXAngst',
@@ -99,7 +104,7 @@ def _create_mn_columns(self):
                         'category').cat.codes.to_numpy()
                 except KeyError:
                     df['MNImageId'] = 0.0
-            
+
             self.data = df
 
         elif self.star_type == 'cistem':
@@ -109,34 +114,38 @@ def _create_mn_columns(self):
             df['cisTEMZFromDefocus'] = df['cisTEMZFromDefocus'] - \
                 df['cisTEMZFromDefocus'].median()
             self.positions = df[['cisTEMOriginalXPosition',
-                      'cisTEMOriginalYPosition', 'cisTEMZFromDefocus']].to_numpy()
+                                 'cisTEMOriginalYPosition', 'cisTEMZFromDefocus']].to_numpy()
             df['MNAnglePhi'] = df['cisTEMAnglePhi']
             df['MNAngleTheta'] = df['cisTEMAngleTheta']
             df['MNAnglePsi'] = df['cisTEMAnglePsi']
             df['MNPixelSize'] = df['cisTEMPixelSize']
             df['MNImageId'] = df['cisTEMOriginalImageFilename'].astype(
                 'category').cat.codes.to_numpy()
-    
+
     def _convert_mrc_to_tiff(self):
         import mrcfile
         from pathlib import Path
         if self.star_type == 'relion':
             micrograph_path = self.object['rlnMicrographName_categories'][self.star_node.inputs['Image'].default_value - 1]
         elif self.star_type == 'cistem':
-            micrograph_path = self.object['cisTEMOriginalImageFilename_categories'][self.star_node.inputs['Image'].default_value - 1].strip("'")
+            micrograph_path = self.object['cisTEMOriginalImageFilename_categories'][
+                self.star_node.inputs['Image'].default_value - 1].strip("'")
         else:
             return False
-        
+
         # This could be more elegant
         if not Path(micrograph_path).exists():
             pot_micrograph_path = Path(self.file_path).parent / micrograph_path
             if not pot_micrograph_path.exists():
                 if self.star_type == 'relion':
-                    pot_micrograph_path = Path(self.file_path).parent.parent.parent / micrograph_path
+                    pot_micrograph_path = Path(
+                        self.file_path).parent.parent.parent / micrograph_path
                     if not pot_micrograph_path.exists():
-                        raise FileNotFoundError(f"Micrograph file {micrograph_path} not found")
+                        raise FileNotFoundError(
+                            f"Micrograph file {micrograph_path} not found")
                 else:
-                    raise FileNotFoundError(f"Micrograph file {micrograph_path} not found")
+                    raise FileNotFoundError(
+                        f"Micrograph file {micrograph_path} not found")
             micrograph_path = pot_micrograph_path
 
         tiff_path = Path(micrograph_path).with_suffix('.tiff')
@@ -148,21 +157,23 @@ def _convert_mrc_to_tiff(self):
             if micrograph_data.ndim == 3:
                 micrograph_data = np.sum(micrograph_data, axis=0)
             # Normalize the data to 0-1
-            micrograph_data = (micrograph_data - micrograph_data.min()) / (micrograph_data.max() - micrograph_data.min())
-
+            micrograph_data = (micrograph_data - micrograph_data.min()) / \
+                (micrograph_data.max() - micrograph_data.min())
+
             if micrograph_data.dtype != np.float32:
                 micrograph_data = micrograph_data.astype(np.float32)
             from PIL import Image
             # Need to invert in Y to generate the correct tiff
-            Image.fromarray(micrograph_data[::-1,:]).save(tiff_path)
+            Image.fromarray(micrograph_data[::-1, :]).save(tiff_path)
         return tiff_path
-    
+
     def _update_micrograph_texture(self, *_):
         try:
             show_micrograph = self.star_node.inputs['Show Micrograph']
             _ = self.object['mn']
         except ReferenceError:
-            bpy.app.handlers.depsgraph_update_post.remove(self._update_micrograph_texture)
+            bpy.app.handlers.depsgraph_update_post.remove(
+                self._update_micrograph_texture)
             return
         if self.star_node.inputs['Image'].default_value == self.current_image:
             return
@@ -180,15 +191,13 @@ def _update_micrograph_texture(self, *_):
             self.micrograph_material.node_tree.nodes['Image Texture'].image = image_obj
             self.star_node.inputs['Micrograph'].default_value = image_obj
 
-
-
     def create_model(self, name='StarFileObject', node_setup=True, world_scale=0.01):
         from molecularnodes.blender.nodes import get_star_node, MN_micrograph_material
         blender_object = bl.obj.create_object(
             self.positions * world_scale, collection=bl.coll.mn(), name=name)
 
         blender_object.mn['molecule_type'] = 'star'
-        
+
         # create attribute for every column in the STAR file
         for col in self.data.columns:
             col_type = self.data[col].dtype
@@ -201,7 +210,8 @@ def create_model(self, name='StarFileObject', node_setup=True, world_scale=0.01)
             elif col_type == object:
                 codes = self.data[col].astype(
                     'category').cat.codes.to_numpy().reshape(-1)
-                bl.obj.set_attribute(blender_object, col, codes, 'INT', 'POINT')
+                bl.obj.set_attribute(blender_object, col,
+                                     codes, 'INT', 'POINT')
                 # Add the category names as a property to the blender object
                 blender_object[f'{col}_categories'] = list(
                     self.data[col].astype('category').cat.categories)
@@ -215,5 +225,72 @@ def create_model(self, name='StarFileObject', node_setup=True, world_scale=0.01)
         self.object = blender_object
         self.star_node = get_star_node(self.object)
         self.micrograph_material = MN_micrograph_material()
-        bpy.app.handlers.depsgraph_update_post.append(self._update_micrograph_texture)
+        bpy.app.handlers.depsgraph_update_post.append(
+            self._update_micrograph_texture)
         return blender_object
+
+
+class NDJSON(Ensemble):
+    def __init__(self, file_path):
+        super().__init__(file_path)
+        self.scale = 10
+
+    @classmethod
+    def from_ndjson(cls, file_path):
+        self = cls(file_path)
+        self.data = self._read()
+        return self
+
+    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']
+                # fixes orientation issues for how the matrices are stored
+                # https://github.com/BradyAJohnston/MolecularNodes/pull/493#issuecomment-2127461280
+                matrix[:3, :3] = np.flip(matrix[:3, :3])
+            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