feat: add method to rotate cell

pyneuroml/utils/__init__.py

@@ -6,13 +6,16 @@
 Copyright 2023 NeuroML Contributors
 """
 
-import typing
+import math
+import copy
 import logging
 import re
+import numpy
 import neuroml
 
 
 logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
 
 
 def extract_position_info(
@@ -133,3 +136,124 @@ def convert_case(name):
     """Converts from camelCase to under_score"""
     s1 = re.sub("(.)([A-Z][a-z]+)", r"\1_\2", name)
     return re.sub("([a-z0-9])([A-Z])", r"\1_\2", s1).lower()
+
+
+def rotate_cell(
+    cell: neuroml.Cell,
+    x: float = 0,
+    y: float = 0,
+    z: float = 0,
+    order: str = "xyz",
+    relative_to_soma: bool = False
+) -> neuroml.Cell:
+    """Return a new cell object rotated in the provided order along the
+    provided angles (in radians) relative to the soma position.
+
+    :param cell: cell object to rotate
+    :type cell: neuroml.Cell
+    :param x: angle to rotate in x direction, in radians
+    :type x: float
+    :param y: angle to rotate in y direction, in radians
+    :type y: float
+    :param z: angle to rotate in z direction, in radians
+    :type z: float
+    :param order: rotation order in terms of x, y,  and z
+    :type order: str
+    :param relative_to_soma: whether rotation is relative to soma
+    :type relative_to_soma: bool
+    :returns: new neuroml.Cell object
+    :rtype: neuroml.Cell
+
+    Derived from LFPy's implementation:
+    https://github.com/LFPy/LFPy/blob/master/LFPy/cell.py#L1600
+    """
+
+    valid_orders = [
+        "xyz", "yzx", "zxy", "xzy", "yxz", "zyx"
+    ]
+    if order not in valid_orders:
+        raise ValueError(f"order must be one of {valid_orders}")
+
+    soma_seg_id = cell.get_morphology_root()
+    soma_seg = cell.get_segment(soma_seg_id)
+    cell_origin = numpy.array([soma_seg.proximal.x, soma_seg.proximal.y, soma_seg.proximal.z])
+    newcell = copy.deepcopy(cell)
+    print(f"Rotating {newcell.id} by {x}, {y}, {z}")
+
+    # calculate rotations
+    if x != 0:
+        anglex = x
+        rotation_x = numpy.array([[1, 0, 0],
+                                  [0, math.cos(anglex), -math.sin(anglex)],
+                                  [0, math.sin(anglex), math.cos(anglex)]
+                                  ])
+        logger.debug(f"x matrix is: {rotation_x}")
+
+    if y != 0:
+        angley = y
+        rotation_y = numpy.array([[math.cos(angley), 0, math.sin(angley)],
+                                  [0, 1, 0],
+                                  [-math.sin(angley), 0, math.cos(angley)]
+                                  ])
+        logger.debug(f"y matrix is: {rotation_y}")
+
+    if z != 0:
+        anglez = z
+        rotation_z = numpy.array([[math.cos(anglez), -math.sin(anglez), 0],
+                                  [math.sin(anglez), math.cos(anglez), 0],
+                                  [0, 0, 1]
+                                  ])
+        logger.debug(f"z matrix is: {rotation_z}")
+
+    # rotate each segment
+    for aseg in newcell.morphology.segments:
+        prox = dist = numpy.array([])
+        # may not have a proximal
+        try:
+            prox = numpy.array([aseg.proximal.x, aseg.proximal.y, aseg.proximal.z])
+        except AttributeError:
+            pass
+
+        # must have distal
+        dist = numpy.array([aseg.distal.x, aseg.distal.y, aseg.distal.z])
+
+        if relative_to_soma:
+            if prox.any():
+                prox = prox - cell_origin
+            dist = dist - cell_origin
+
+        # rotate
+        for axis in order:
+            if axis == 'x' and x != 0:
+                if prox.any():
+                    prox = numpy.dot(prox, rotation_x)
+                dist = numpy.dot(dist, rotation_x)
+
+            if axis == 'y' and y != 0:
+                if prox.any():
+                    prox = numpy.dot(prox, rotation_y)
+                dist = numpy.dot(dist, rotation_y)
+
+            if axis == 'z' and z != 0:
+                if prox.any():
+                    prox = numpy.dot(prox, rotation_z)
+                dist = numpy.dot(dist, rotation_z)
+
+        if relative_to_soma:
+            if prox.any():
+                prox = prox + cell_origin
+            dist = dist + cell_origin
+
+        if prox.any():
+            aseg.proximal.x = prox[0]
+            aseg.proximal.y = prox[1]
+            aseg.proximal.z = prox[2]
+
+        aseg.distal.x = dist[0]
+        aseg.distal.y = dist[1]
+        aseg.distal.z = dist[2]
+
+        logger.debug(f"prox is: {aseg.proximal}")
+        logger.debug(f"distal is: {aseg.distal}")
+
+    return newcell

tests/utils/test_utils.py

@@ -10,9 +10,11 @@
 
 import logging
 import pathlib as pl
+import math
 
-from pyneuroml.pynml import read_neuroml2_file
-from pyneuroml.utils import extract_position_info
+import neuroml
+from pyneuroml.pynml import read_neuroml2_file, write_neuroml2_file
+from pyneuroml.utils import extract_position_info, rotate_cell
 
 from .. import BaseTestCase
 
@@ -46,3 +48,46 @@ def test_extract_position_info(self):
 
             for c in ["HL23PV", "HL23PYR", "HL23VIP", "HL23SST"]:
                 self.assertIn(c, cell_id_vs_cell.keys())
+
+    def test_rotate_cell(self):
+        """Test rotate_cell"""
+        acell = neuroml.utils.component_factory("Cell", id="test_cell", validate=False)  # type: neuroml.Cell
+        soma = acell.add_segment(prox=[0, 0, 0, 1], dist=[0, 0, 0, 1], seg_id=0,
+                                 use_convention=False, reorder_segment_groups=False,
+                                 optimise_segment_groups=False)
+        acell.add_segment(prox=[0, 0, 0, 1], dist=[5, 0, 0, 1], seg_id=1,
+                          use_convention=False, reorder_segment_groups=False,
+                          optimise_segment_groups=False, parent=soma)
+
+        acell.add_segment(prox=[0, 0, 0, 10], dist=[0, 10, 0, 10], seg_id=2,
+                          use_convention=False, reorder_segment_groups=False,
+                          optimise_segment_groups=False, parent=soma)
+
+        acell.add_segment(prox=[0, 0, 0, 25], dist=[0, 0, 25, 25], seg_id=3,
+                          use_convention=False, reorder_segment_groups=False,
+                          optimise_segment_groups=False, parent=soma)
+
+        print(acell)
+
+        rotated_cell = rotate_cell(acell, x=math.pi / 2, y=0, z=0, order="xyz")
+        rotated_cell.id = "test_rotated_cell"
+        print(rotated_cell)
+
+        newdoc = neuroml.utils.component_factory("NeuroMLDocument",
+                                                 id="test_doc")  # type: neuroml.NeuroMLDocument
+        newdoc.add(acell)
+        newdoc.add(rotated_cell)
+
+        net = newdoc.add("Network", id="test_net", validate=False)
+        pop1 = net.add("Population", id="test_pop1", size=1, component=acell.id,
+                       type="populationList", validate=False)
+        pop1.add("Instance", id=0, location=pop1.component_factory("Location", x=0, y=0, z=0))
+
+        pop2 = net.add("Population", id="test_pop2", size=1,
+                       component=rotated_cell.id,
+                       type="populationList", validate=False)
+        pop2.add("Instance", id=0, location=pop1.component_factory("Location",
+                                                                   x=50, y=0, z=0))
+
+        newdoc.validate(recursive=True)
+        write_neuroml2_file(newdoc, "test_rotation.net.nml", validate=True)