diff --git a/pyneuroml/utils/__init__.py b/pyneuroml/utils/__init__.py index 044c3781..38cbeb4e 100644 --- a/pyneuroml/utils/__init__.py +++ b/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 diff --git a/tests/utils/test_utils.py b/tests/utils/test_utils.py index ff0dff68..e2024eec 100644 --- a/tests/utils/test_utils.py +++ b/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)