ALLEN INSTITUTE

# Source code for bmtk.simulator.bionet.morphology

    # Copyright 2022. Allen Institute. All rights reserved
    #
    # Redistribution and use in source and binary forms, with or without modification, are permitted provided that the
    # following conditions are met:
    #
    # 1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following
    # disclaimer.
    #
    # 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following
    # disclaimer in the documentation and/or other materials provided with the distribution.
    #
    # 3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote
    # products derived from this software without specific prior written permission.
    #
    # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES,
    # INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
    # DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
    # SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
    # SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
    # WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
    # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
    #
    import re
    from collections import namedtuple
    import warnings
    import numpy as np
    import pandas as pd
    from neuron import h
    import neuron

    from bmtk.simulator.bionet.io_tools import io

    pc = h.ParallelContext()  # object to access MPI methods

    class _LazySegmentProps(object):
        """A Struct for storing coordinate invariant properties of each segment of a cell."""
        def __init__(self, morphology):
            self._morphology = morphology
            self._hobj = morphology.hobj
            self.sec_id = None
            self.type = None
            self.area = None
            self.x = None
            self.x0 = None
            self.x1 = None
            self.dist = None
            self.length = None
            self.dist0 = None
            self.dist1 = None

        def get(self):
            if self.sec_id is None:
                seg_type = []
                seg_area = []
                seg_x = []
                seg_x0 = []
                seg_x1 = []
                seg_dist = []
                seg_length = []
                seg_sec_id = []

                h.distance(sec=self._morphology.soma[0])  # measure distance relative to the soma

                for sec_id, sec in enumerate(self._hobj.all):
                    fullsecname = sec.name()
                    sec_type = fullsecname.split(".")[-1][:4]  # get sec name type without the cell name
                    sec_type_swc = Morphology.sec_type_swc[sec_type]  # convert to swc code
                    x_range = 1.0 / (sec.nseg * 2)  # used to calculate [x0, x1]

                    for seg in sec:
                        seg_area.append(h.area(seg.x))
                        seg_x.append(seg.x)
                        seg_x0.append(seg.x - x_range)
                        seg_x1.append(seg.x + x_range)
                        seg_length.append(sec.L / sec.nseg)
                        seg_type.append(sec_type_swc)  # record section type in a list
                        # seg_dist.append(h.distance(seg.x))  # distance to the center of the segment
                        seg_dist.append(h.distance(seg))
                        seg_sec_id.append(sec_id)

                length_arr = np.array(seg_length)
                dist_arr = np.array(seg_dist)
                dist0_arr = dist_arr - length_arr / 2.0
                dist1_arr = dist_arr + length_arr / 2.0

                self.type = np.array(seg_type)
                self.area = np.array(seg_area)
                self.x = np.array(seg_x)
                self.x0 = np.array(seg_x0)
                self.x1 = np.array(seg_x1)
                self.dist = dist_arr
                self.length = length_arr
                self.dist0 = dist0_arr
                self.dist1 = dist1_arr
                self.sec_id = np.array(seg_sec_id)

            return self

    class _LazySegmentCoords(object):
        """A Struct for storing properties of each segment of a cell that vary by soma position and rotation."""
        def __init__(self, morphology):
            self._hobj = morphology.hobj
            self.p0 = None
            self.p1 = None
            self.p05 = None
            self.soma_pos = None
            self.d = None

        def get(self):
            if self.p0 is None:
                # Iterates through each segment (one section may contain one or more segments). We use NEURON's
                # segment.x to get the mid-point and lenght of each segment we we can use to find and store the beginning
                # (p0) and end (p1) of each segment. We also shift so the middle of the soma is at the origin.
                nseg = np.sum([s.nseg for s in self._hobj.all])

                n3dsoma = 0
                soma_pos = np.zeros(3)
                for sec in self._hobj.soma:
                    n3d = int(h.n3d(sec=sec))  # get number of n3d points in each section
                    n3dsoma += n3d
                    for i in range(n3d):
                        soma_pos[0] += h.x3d(i, sec=sec)
                        soma_pos[1] += h.y3d(i, sec=sec)
                        soma_pos[2] += h.z3d(i, sec=sec)

                soma_pos /= n3dsoma

                ix = 0  # segment index
                self.p0 = np.zeros((3, nseg))  # hold the coordinates of segment starting points
                self.p1 = np.zeros((3, nseg))  # hold the coordinates of segment end points
                self.p05 = np.zeros((3, nseg))
                for sec in self._hobj.all:
                    n3d = int(h.n3d(sec=sec))  # get number of n3d points in each section
                    p3d = np.zeros((3, n3d))  # to hold locations of 3D morphology for the current section
                    l3d = np.zeros(n3d)  # to hold locations of 3D morphology for the current section
                    # diam3d = np.zeros(n3d)  # to diameters, at the moment we don't need to keep track of diameter

                    for i in range(n3d):
                        p3d[0, i] = h.x3d(i, sec=sec) - soma_pos[0]  # shift coordinates such to place soma at the origin.
                        p3d[1, i] = h.y3d(i, sec=sec) - soma_pos[1]
                        p3d[2, i] = h.z3d(i, sec=sec) - soma_pos[2]
                        # diam3d[i] = h.diam3d(i, sec=sec)
                        l3d[i] = h.arc3d(i, sec=sec)

                    l3d /= sec.L  # normalize
                    nseg = sec.nseg

                    l0 = np.zeros(nseg)  # keep range of segment starting point
                    l1 = np.zeros(nseg)  # keep range of segment ending point
                    l05 = np.zeros(nseg)

                    for iseg, seg in enumerate(sec):
                        l0[iseg] = seg.x - 0.5 * 1 / nseg  # x (normalized distance along the section) for the beginning of the segment
                        l1[iseg] = seg.x + 0.5 * 1 / nseg  # x for the end of the segment
                        l05[iseg] = seg.x

                    self.p0[0, ix:ix + nseg] = np.interp(l0, l3d, p3d[0, :])
                    self.p0[1, ix:ix + nseg] = np.interp(l0, l3d, p3d[1, :])
                    self.p0[2, ix:ix + nseg] = np.interp(l0, l3d, p3d[2, :])

                    self.p1[0, ix:ix + nseg] = np.interp(l1, l3d, p3d[0, :])
                    self.p1[1, ix:ix + nseg] = np.interp(l1, l3d, p3d[1, :])
                    self.p1[2, ix:ix + nseg] = np.interp(l1, l3d, p3d[2, :])

                    self.p05[0, ix:ix + nseg] = np.interp(l05, l3d, p3d[0, :])
                    self.p05[1, ix:ix + nseg] = np.interp(l05, l3d, p3d[1, :])
                    self.p05[2, ix:ix + nseg] = np.interp(l05, l3d, p3d[2, :])

                    # x0[ix:ix + nseg] = l0
                    # x1[ix:ix + nseg] = l1

                    ix += nseg

                # TODO: keep track of diameter only when minimum_distance='auto' in EcpMod
                self.d = np.array([seg.diam for sec in self._hobj.all for seg in sec])

                # Also calculate the middle of the shifted soma, NOTE: in theory this should be (0, 0, 0), but due to
                # percision the actual soma center might be a little off.
                n3dsoma = 0
                r3dsoma = np.zeros(3)
                for sec in self._hobj.soma:
                    n3d = int(h.n3d(sec=sec))  # get number of n3d points in each section
                    n3dsoma += n3d
                    for i in range(n3d):
                        r3dsoma[0] += h.x3d(i, sec=sec) - soma_pos[0]
                        r3dsoma[1] += h.y3d(i, sec=sec) - soma_pos[1]
                        r3dsoma[2] += h.z3d(i, sec=sec) - soma_pos[2]

                r3dsoma /= n3dsoma
                self.soma_pos = r3dsoma

            return self

    # For a lot of different cells created from the same morphology and model_processing function, the core segment
    # properties will be the same (even if coordinates, synapses, hobjs, etc. are different). In such cases we want
    # to calculate and store seg_props only once for all equivelent cells. We can do this by caching the seg_props
    # helper function pointer, then all equivelent morphologies (as defined by __hash__ function) use the same
    # SegmentProps object
    morphology_cache = {}

    def rotation_matrix(axis, theta):
        """Return the rotation matrix associated with counterclockwise rotation about the given axis by theta radians.
        """
        axis = np.asarray(axis)
        theta = np.asarray(theta)
        axis = axis/np.sqrt(np.dot(axis, axis))
        a = np.cos(theta/2.0)
        b, c, d = -axis*np.sin(theta/2.0)
        aa, bb, cc, dd = a*a, b*b, c*c, d*d
        bc, ad, ac, ab, bd, cd = b*c, a*d, a*c, a*b, b*d, c*d

        return np.array([
            [aa+bb-cc-dd, 2*(bc+ad), 2*(bd-ac)],
            [2*(bc-ad), aa+cc-bb-dd, 2*(cd+ab)],
            [2*(bd+ac), 2*(cd-ab), aa+dd-bb-cc]
        ])

    class Morphology(object):
        sec_type_swc = {
            'soma': 1, 'somatic': 1,
            'axon': 2, 'axonal': 2,
            'dend': 3, 'basal': 3,
            'apic': 4, 'apical': 4
        }

        def __init__(self, hobj, rng_seed=None, swc_path=None):
            self.hobj = hobj  # h.Biophys1(swc_path)
            self.rng_seed = rng_seed
            self.swc_path = swc_path

            self._prng = np.random.RandomState(self.rng_seed)
            self._sections = None
            self._segments = None
            self._seg_props = _LazySegmentProps(self)  # None
            self._seg_coords = _LazySegmentCoords(self)  # None
            self._nseg = None
            self._swc_map = None
            self._soma = None #

            # Used by find_sections() and other methods when building edges/synapses. Should make it faster to look-up
            # cooresponding segments for a large number of syns that target the same area of a cell
            self._trg_segs_cache = {}

        @property
        def soma(self):
            if self._soma is None:
                if isinstance(self.hobj.soma, (neuron.nrn.Section, neuron.nrn.Segment)):
                    self._soma = [self.hobj.soma]
                else:
                    self._soma = self.hobj.soma

            return self._soma

        def _copy(self):
            new_morph = Morphology(hobj=self.hobj, swc_path=self.swc_path)
            return new_morph

        def _soma_position_orig(self):
            n3dsoma = 0
            r3dsoma = np.zeros(3)
            for sec in self.hobj.soma:
                n3d = int(h.n3d(sec=sec))  # get number of n3d points in each section
                # r3d = np.zeros((3, n3d))  # to hold locations of 3D morphology for the current section
                n3dsoma += n3d

                for i in range(n3d):
                    r3dsoma[0] += h.x3d(i, sec=sec)
                    r3dsoma[1] += h.y3d(i, sec=sec)
                    r3dsoma[2] += h.z3d(i, sec=sec)

            r3dsoma /= n3dsoma
            return r3dsoma

        @property
        def segments(self):
            if self._segments is None:
                self._segments = []
                for sec in self.hobj.all:
                    for seg in sec:
                        self._segments.append(seg)

            return self._segments

        @property
        def sections(self):
            if self._sections is None:
                self._sections = []
                for sec in self.hobj.all:
                    self._sections.append(sec)

            return self._sections

        @property
        def n_sections(self):
            return len(self.sections)

        @property
        def seg_props(self):
            return self._seg_props.get()

        @property
        def seg_coords(self):
            """Get the coordinates of all segments of the cells. Each segment is defined by three values, the beginning
            of the segment (seg_coords.p0), the middle of the segment(seg_coords.p05), and the end (seg_coords.p1)."""
            return self._seg_coords.get()

        @property
        def nseg(self):
            if self._nseg is None:
                nseg = 0
                for sec in self.hobj.all:
                    nseg += sec.nseg  # get the total # of segments in the cell
                self._nseg = nseg

            return self._nseg

        @property
        def soma_position(self):
            return self.seg_coords.soma_pos

        @property
        def swc_map(self):
            if self._swc_map is None:
                swc_df = pd.read_csv(self.swc_path, sep=' ', names=['id', 'type', 'x', 'y', 'z', 'r', 'pid'], comment='#')
                swc_df = swc_df[['id', 'type']]

                name_indices = []
                for ln in range(len(swc_df)):
                    # in read_swc.hoc all non-bifurcating section of rows are assigned a global section-id. But for each
                    # type (soma, dend, etc), there is also a nameindex id which is unique within the type, used to assign
                    # sections to their name. So lines with sec_id=10 may have nameindex=3 (Biophys1[0].dend[3])
                    sec_id = int(self.hobj.nl.point2sec[ln])
                    # swc_id = int(self.hobj.nl.pt2id(ln))
                    # swc_type = int(self.hobj.nl.type[ln])
                    name_index = int(self.hobj.nl.sections.object(sec_id).nameindex)
                    name_indices.append(name_index)

                swc_df['nameindex'] = name_indices
                self._swc_map = swc_df

            return self._swc_map

        def get_section(self, section_idx):
            return self.sections[section_idx]

        def set_segment_dl(self, dl):
            """Define number of segments in a cell"""
            self._nseg = 0
            for sec in self.hobj.all:
                sec.nseg = 1 + 2 * int(sec.L/(2*dl))
                self._nseg += sec.nseg  # get the total number of segments in the cell

        def choose_sections(self, section_names, distance_range, n_sections=1, cache=True):
            """Similar to find_sections, but will only N=n_section number of sections_ids/x values randomly selected (may
            return less if there aren't as many sections

            :param section_names: 'soma', 'dend', 'apic', 'axon'
            :param distance_range: [float, float]: distance range of sections from the soma, in um.
            :param n_sections: int: maximum number of sections to select
            :param cache: caches the segments + probs so that next time the same section_names/distance_range values are
                called it will return the same values without recalculating. default True.
            :return: [int], [float]: A list of all section_ids and a list of all segment_x values (as defined by NEURON)
                that meet the given critera.
            """
            secs, probs = self.find_sections(section_names, distance_range, cache=cache)
            secs_ix = self._prng.choice(secs, n_sections, p=probs)
            return secs_ix, self.seg_props.x[secs_ix]

        def find_sections(self, section_names, distance_range, cache=True):
            """Retrieves a list of sections ids and section x's given a section name/type (eg axon, soma, apic, dend) and
            the distance from the soma.

            :param section_names: A list of sections to target, 'soma', 'dend', 'apic', 'axon'
            :param distance_range: [float, float]: distance range of sections from the soma, in um.
            :param cache: caches the segments + probs so that next time the same section_names/distance_range values are
                called it will return the same values without recalculating. default True.
            :return: [float], [float]: A list of all section_ids and a list of all segment_x values (as defined by NEURON)
                that meet the given critera.
            """
            cache_key = (tuple(section_names), tuple(distance_range))
            if cache and cache_key in self._trg_segs_cache:
                return self._trg_segs_cache[cache_key]

            dmin, dmax = distance_range[0], distance_range[1]

            seg_d0 = self.seg_props.dist0
            seg_d1 = self.seg_props.dist1
            seg_length = self.seg_props.length
            seg_type = self.seg_props.type

            # Find the fractional overlap between the segment and the distance range:
            # this is done by finding the overlap between [d0,d1] and [dmin,dmax]
            # np.minimum(seg_d1,dmax) find the smaller of the two end locations
            # np.maximum(seg_d0,dmin) find the larger of the two start locations
            # np.maximum(0,overlap) is used to return zero when segments do not overlap
            # and then dividing by the segment length
            frac_overlap = np.maximum(0, (np.minimum(seg_d1, dmax) - np.maximum(seg_d0, dmin))) / seg_length
            ix_drange = np.where(frac_overlap > 0)  # find indexes with non-zero overlap
            ix_labels = np.array([], dtype=int)

            for tar_sec_label in section_names:  # find indexes within sec_labels
                sec_type = self.sec_type_swc[tar_sec_label]  # get swc code for the section label
                ix_label = np.where(seg_type == sec_type)
                ix_labels = np.append(ix_labels, ix_label)  # target segment indexes

            tar_seg_ix = np.intersect1d(ix_drange, ix_labels)  # find intersection between indexes for range and labels
            tar_seg_length = seg_length[tar_seg_ix] * frac_overlap[tar_seg_ix]  # weighted length of targeted segments
            tar_seg_prob = tar_seg_length / np.sum(tar_seg_length)  # probability of targeting segments

            if cache:
                self._trg_segs_cache[cache_key] = (tar_seg_ix, tar_seg_prob)

            return tar_seg_ix, tar_seg_prob

        def move_and_rotate(self, soma_coords=None, rotation_angles=None, inplace=False):
            old_seg_coords = self.seg_coords
            new_p0 = old_seg_coords.p0.copy()
            new_p1 = old_seg_coords.p1.copy()
            new_p05 = old_seg_coords.p05.copy()
            new_soma_pos = self.soma_position.copy()

            if rotation_angles is not None:
                assert(len(rotation_angles) == 3)
                # try to calc the euler rotation matrix around an arbitary point was causing problems, instead move coords
                # so the the soma center (p05[0]) is at the origin
                soma_pos_mat = new_soma_pos.reshape((3, 1))
                new_p0 = new_p0 - soma_pos_mat
                new_p1 = new_p1 - soma_pos_mat
                new_p05 = new_p05 - soma_pos_mat

                rotx_mat = rotation_matrix([1, 0, 0], rotation_angles[0])
                roty_mat = rotation_matrix([0, 1, 0], rotation_angles[1])
                rotz_mat = rotation_matrix([0, 0, 1], rotation_angles[2])
                rotxyz_mat = np.dot(rotx_mat, roty_mat.dot(rotz_mat))

                new_p0 = np.dot(rotxyz_mat, new_p0) + soma_pos_mat
                new_p1 = np.dot(rotxyz_mat, new_p1) + soma_pos_mat
                new_p05 = np.dot(rotxyz_mat, new_p05) + soma_pos_mat

            if soma_coords is not None:
                assert(len(soma_coords) == 3)
                soma_coords = soma_coords if isinstance(soma_coords, np.ndarray) else np.array(soma_coords)
                displacement = soma_coords - self.soma_position
                displacement = displacement.reshape((3, 1))  # Req to allow adding vector to matrix
                new_p0 += displacement
                new_p1 += displacement
                new_p05 += displacement
                new_soma_pos = soma_coords

            # new_seg_coords = SegmentCoords(p0=new_p0, p1=new_p1, p05=new_p05, soma_pos=new_soma_pos)
            new_seg_coords = _LazySegmentCoords(self)
            new_seg_coords.p0 = new_p0
            new_seg_coords.p05 = new_p05
            new_seg_coords.p1 = new_p1
            new_seg_coords.soma_pos = new_soma_pos
            new_seg_coords.d = old_seg_coords.d  # diameter won't change in new position

            if inplace:
                self._seg_coords = new_seg_coords
                return self
            else:
                new_morph = self._copy()
                new_morph._seg_props = self.seg_props
                new_morph._seg_coords = new_seg_coords
                new_morph._soma_pos = new_soma_pos
                return new_morph

        def get_coords(self, sec_id, sec_x):
            segs_indices = np.argwhere(self.seg_props.sec_id == sec_id).flatten()
            return self.seg_coords.p05[:, segs_indices][:, 0]

        def get_swc_id(self, sec_id, sec_x):
            # use sec type and nameindex to find all rows in the swc that correspond to sec_id
            sec = self.sections[sec_id]
            sec_nameindex = self._get_sec_nameindex(sec)
            sec_type = self._get_sec_type(sec)
            filtered_swc = self.swc_map[(self.swc_map['type'] == sec_type) & (self.swc_map['nameindex'] == sec_nameindex)]
            swc_ids = filtered_swc['id'].values

            # use sec_x, a value between [0, 1], to estimate which swc_id/line to choose.
            # NOTE: At the moment it assumes each line in the swc is the same distance apart, making estimating the sec_x
            #  location easy by the number it appears in the squence
            if len(swc_ids) == 0:
                return -1, 0.0
            elif len(swc_ids) == 1:
                return swc_ids[0], sec_x
            else:
                swc_place = np.max((0.0, sec_x*len(swc_ids) - 1.0))
                swc_indx = int(np.ceil(swc_place))
                swc_id = swc_ids[swc_indx]
                swc_dist = swc_indx - swc_place
                return swc_id, swc_dist

        def _get_sec_type(self, sec):
            sec_name = sec.hname()
            sec_type_str = sec_name.split('.')[-1].split('[')[0]
            type_str = self.sec_type_swc[sec_type_str]
            return int(type_str)

        def _get_sec_nameindex(self, sec):
            sec_name = sec.hname()
            sec_type_name = sec_name.split('.')[-1]
            nameindex_str = re.search(r'\[(\d+)\]', sec_type_name).group(1)
            return int(nameindex_str)

        def __eq__(self, other):
            if self.swc_path is None or other.swc_path is None:
                return False
            else:
                return hash(self) == hash(other)

        def __hash__(self):
            """Used for comparing two Morphology objects and for storing as dict keys in seg_props_cache.

            We assume two Morphologies are the same if they have the same morphology file path + same # of segments. This
            is not a good way of determining if two Morphologies are equal and in actuality we would need to compare
            seg_props on an segment-by-segment basis. However that would be comp. expensive for large networks and instead
            do a quick spot test."""
            comp_vals = (self.swc_path if self.swc_path else np.random.randint, self.nseg)
            return hash(comp_vals)

        @classmethod
        def load(cls, hobj=None, morphology_file=None, rng_seed=None, cache_seg_props=True):
            """Factory method, perfered way to create a Morphology object

            :param hobj:
            :param morphology_file:
            :param rng_seed:
            :param cache_seg_props:
            :return:
            """
            if hobj is None and morphology_file is None:
                io.log_exception('Unable to create Morphology, no valid HOC object or morphology file specified')

            elif morphology_file is not None and hobj is None:
                # If the HOC object hasn't already been created try creating it from the morphology
                hobj = h.Biophys1(morphology_file)

            morph = cls(hobj=hobj, swc_path=morphology_file, rng_seed=rng_seed)

            if cache_seg_props and morphology_file is not None:
                if morph in morphology_cache:
                    morph._seg_props = morphology_cache[morph][0]
                    morph._seg_coords = morphology_cache[morph][1]
                else:
                    morphology_cache[morph] = (morph._seg_props, morph._seg_coords)

            return morph