Update rxd.MultipleGeometry to support 3D and hybrid models. (#3389)

* Update rxd.MultipleGeometry to support 3D and hybrid models.

* Made corrections suggested by @ramcdougal.

Author: adamjhn

share/lib/python/neuron/rxd/geometry.py

@@ -3,6 +3,7 @@
 import numpy
 from neuron import h, nrn
 from .rxdException import RxDException
+from itertools import groupby
 
 try:
     from neuron.rxd import geometry3d
@@ -151,6 +152,8 @@ def scale_by_constant(scale, f):
 inside = RxDGeometry()
 if has_geometry3d:
     inside.volumes3d = geometry3d.voxelize2
+    inside._scale_internal_volumes3d = 1
+    inside._scale_surface_volumes3d = 1
 # neighbor_area_fraction can be a constant or a function
 inside.neighbor_area_fraction = 1
 inside.volumes1d = _volumes1d
@@ -201,6 +204,8 @@ class DistributedBoundary(RxDGeometry):
 
     def __init__(self, area_per_vol, perim_per_area=0):
         self._area_per_vol = area_per_vol
+        self._scale_internal_volumes3d = area_per_vol
+        self._scale_surface_volumes3d = area_per_vol
         self._perim_per_area = 0
 
         self.surface_areas1d = _always_0
@@ -268,6 +273,9 @@ def __init__(
         self.is_volume = _always_true
         self.is_area = _always_false
         self._volume_fraction = volume_fraction
+        self._scale_internal_volumes3d = volume_fraction
+        self._scale_surface_volumes3d = volume_fraction
+
         self._surface_fraction = surface_fraction
         self._neighbor_areas_fraction = neighbor_areas_fraction
         # TODO: does the else case ever make sense?
@@ -585,6 +593,36 @@ def _get_geo(self, sec):
                 )
         return geo
 
+    def volumes3d(self, source, dx=0.25):
+        geos = {}
+        for sec in source:
+            geo = self._get_geo(sec)
+            geos.setdefault(geo, []).append(sec)
+        internal_voxels, surface_voxels, mesh_grid = geometry3d.voxelize2(source, dx)
+        internal_key = lambda itm: itm[1][1].sec
+        surface_key = lambda itm: itm[1][2].sec
+        internal_by_sec = groupby(internal_voxels.items(), key=internal_key)
+        surface_by_sec = groupby(surface_voxels.items(), key=surface_key)
+        for geo, secs in geos.items():
+            # only scale if not multiplying by 1
+            if geo._scale_internal_volumes3d != 1:
+                for sec, voxels in internal_by_sec:
+                    # only scale if sec is in this geometry
+                    if sec in secs:
+                        # scale the internal volume
+                        for key, _ in voxels:
+                            internal_voxels[key][0] *= geo._scale_internal_volumes3d
+            # only scale if not multiplying by 1
+            if geo._scale_surface_volumes3d != 1:
+                for sec, voxels in surface_by_sec:
+                    # only scale if sec is in this geometry
+                    if sec in secs:
+                        # scale the surface volume
+                        for key, _ in voxels:
+                            surface_voxels[key][0] *= geo._scale_surface_volumes3d
+
+        return internal_voxels, surface_voxels, mesh_grid
+
     def volumes1d(self, sec):
         return self._get_geo(sec).volumes1d(sec)
 

test/rxd/3d/test_multigeometry.py

@@ -0,0 +1,41 @@
+import pytest
+import numpy
+from testutils import compare_data, tol
+
+
+@pytest.mark.parametrize("hybrid", [True, False])
+def test_multi_geometry3d(neuron_nosave_instance, hybrid):
+    """A model using multiple different geometries in a 3D model"""
+
+    h, rxd, save_path = neuron_nosave_instance
+    dendA = h.Section("dendA")
+    dendA.diam = 1
+    dendA.nseg = 3
+    dendA.L = 5
+    dendB = h.Section("dendB")
+    dendB.diam = 1
+    dendB.nseg = 3
+    dendB.L = 5
+
+    if hybrid:
+        rxd.set_solve_type([dendA], dimension=3)
+    else:
+        rxd.set_solve_type(dimension=3)
+
+    geoA = rxd.inside
+    geoB = rxd.FractionalVolume(volume_fraction=0.5, surface_fraction=1)
+    geometry = rxd.MultipleGeometry(secs=[dendA, dendB], geos=[geoA, geoB])
+    r = rxd.Region(h.allsec(), dx=0.75, geometry=geometry)
+    ca = rxd.Species(r, name="ca", charge=2)
+
+    if hybrid:
+        assert len(ca.nodes(dendA)) == 28
+        assert (len(ca.nodes(dendB))) == 3
+        assert sum(ca.nodes(dendA).volume) == 2.671875
+        assert sum(ca.nodes(dendB).volume) == 1.9634954084936205
+
+    else:
+        assert len(ca.nodes(dendA)) == 28
+        assert len(ca.nodes(dendA)) == 28
+        assert sum(ca.nodes(dendA).volume) == 2.671875
+        assert sum(ca.nodes(dendB).volume) == 1.3359375