diff --git a/molecularnodes/io/parse/__init__.py b/molecularnodes/io/parse/__init__.py index e19866d0..1c2f6e96 100644 --- a/molecularnodes/io/parse/__init__.py +++ b/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 diff --git a/molecularnodes/io/parse/star.py b/molecularnodes/io/parse/star.py index 1b73aff1..d50276f4 100644 --- a/molecularnodes/io/parse/star.py +++ b/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 diff --git a/molecularnodes/io/star.py b/molecularnodes/io/star.py index 5b6743ca..02a5e6f2 100644 --- a/molecularnodes/io/star.py +++ b/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