Avoid np.where (#171)

* got rid of some use of np.where

* python_requires='>=3.7'

* fancy indexing

* WHY?

* install ipympl

* action version bump

* be nice

* clean up

* formatting, fix tests

* Fixes LFPy's example_MEA.py

.github/workflows/coveralls.yml

@@ -19,20 +19,20 @@ jobs:
             python-version: [3.8]
 
     steps:
-    - uses: actions/checkout@v2
+    - uses: actions/checkout@v3
       with:
         github-token: ${{ secrets.GITHUB_TOKEN }}
     - name: Set up Python 3
-      uses: actions/setup-python@v2
+      uses: actions/setup-python@v3
       with:
         python-version: ${{ matrix.python-version }}
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
         pip install flake8 pytest wheel coverage pytest-cov coveralls
         if [ -f requirements.txt ]; then pip install -r requirements.txt; fi
-        if [[ "${{ matrix.python-version }}" < "3.9" ]]; then pip install neuron; fi
         pip install .
+        pip install neuron
     - name: Test with pytest and coveralls
       env:
         COVERALLS_REPO_TOKEN: ${{ secrets.COVERALLS_REPO_TOKEN }}

.github/workflows/flake8.yml

@@ -8,9 +8,9 @@ jobs:
     name: Lint
     steps:
       - name: Check out source repository
-        uses: actions/checkout@v2
+        uses: actions/checkout@v3
       - name: Set up Python environment
-        uses: actions/setup-python@v2
+        uses: actions/setup-python@v3
         with:
           python-version: "3.8"
       - name: flake8 Lint

.github/workflows/python-app.yml

@@ -19,9 +19,9 @@ jobs:
             python-version: ["3.7", "3.8", "3.9", "3.10"]
 
     steps:
-    - uses: actions/checkout@v2
+    - uses: actions/checkout@v3
     - name: Set up Python 3
-      uses: actions/setup-python@v2
+      uses: actions/setup-python@v3
       with:
         python-version: ${{ matrix.python-version }}
     - name: Install dependencies

.github/workflows/python-publish.yml

@@ -13,9 +13,9 @@ jobs:
     runs-on: ubuntu-latest
 
     steps:
-    - uses: actions/checkout@v2
+    - uses: actions/checkout@v3
     - name: Set up Python
-      uses: actions/setup-python@v2
+      uses: actions/setup-python@v3
       with:
         python-version: '3.x'
     - name: Install dependencies

examples/Example_SphericallySymmetricVolCondMEG_vs_InfiniteHomogeneousVolCondMEG.ipynb

@@ -14,6 +14,18 @@
   {
    "cell_type": "code",
    "execution_count": 1,
+   "id": "40203944-3ed4-46b9-94fa-208638c3880d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# This example uses matplotlib widgets. Uncomment and run either line. A kernel restart may be needed\n",
+    "# !pip install ipympl  # or\n",
+    "# !conda install ipympl -y"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
    "id": "f5e2386c",
    "metadata": {},
    "outputs": [],
@@ -23,7 +35,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 3,
    "id": "816eaac5",
    "metadata": {},
    "outputs": [],
@@ -35,7 +47,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 4,
    "id": "30dc67a2",
    "metadata": {},
    "outputs": [],
@@ -59,7 +71,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 5,
    "id": "a9d06789",
    "metadata": {},
    "outputs": [],
@@ -73,7 +85,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 6,
    "id": "d8623405",
    "metadata": {},
    "outputs": [],
@@ -87,7 +99,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 7,
    "id": "c83a674f",
    "metadata": {},
    "outputs": [],
@@ -99,7 +111,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 8,
    "id": "beb6b872",
    "metadata": {},
    "outputs": [],
@@ -110,7 +122,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 9,
    "id": "f4bad94f",
    "metadata": {
     "scrolled": true,
@@ -123,14 +135,14 @@
        "Text(0.5, 0, 'z (µm)')"
       ]
      },
-     "execution_count": 8,
+     "execution_count": 9,
      "metadata": {},
      "output_type": "execute_result"
     },
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "2fe6924c71fc4dbe804d10153e4a9368",
+       "model_id": "6e46dc7fa71240d695428c89fbd440cd",
        "version_major": 2,
        "version_minor": 0
       },
@@ -173,7 +185,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": 10,
    "id": "cf08c0e6",
    "metadata": {
     "tags": []
@@ -185,14 +197,14 @@
        "Text(0.5, 0, 'z (µm)')"
       ]
      },
-     "execution_count": 9,
+     "execution_count": 10,
      "metadata": {},
      "output_type": "execute_result"
     },
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
-       "model_id": "b5b7f041f4644987b0049ab1866096e9",
+       "model_id": "d676b69090e14c319ef6a429669a6ba3",
        "version_major": 2,
        "version_minor": 0
       },
@@ -266,7 +278,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.10"
+   "version": "3.10.2"
   }
  },
  "nbformat": 4,

lfpykit/lfpcalc.py

@@ -117,6 +117,7 @@ def calc_lfp_linesource_anisotropic(cell_x, cell_y, cell_z,
          sigma[0] * sigma[2] * (y - ystart)**2 +
          sigma[0] * sigma[1] * (z - zstart)**2)
 
+    # this looks not optimal:
     for idx in np.where(rs < r_limit)[0]:
         r = rs[idx]
         closest_point = closest_points[:, idx]
@@ -162,7 +163,7 @@ def calc_lfp_linesource_anisotropic(cell_x, cell_y, cell_z,
     if (i + iia + iib + iii + iiii).sum() != xstart.size:
         print(a, b, c)
         print(i, iia, iib, iii, iiii)
-        raise RuntimeError
+        raise RuntimeError  # WHY?
 
     mapping = np.zeros(xstart.size)
     mapping[i] = _anisotropic_line_source_case_i(a[i], c[i])
@@ -230,6 +231,7 @@ def calc_lfp_root_as_point_anisotropic(cell_x, cell_y, cell_z,
          sigma[0] * sigma[2] * (y - ystart)**2 +
          sigma[0] * sigma[1] * (z - zstart)**2)
 
+    # this looks not optimal:
     for idx in np.where(rs < r_limit)[0]:
         r = rs[idx]
         closest_point = closest_points[:, idx]
@@ -252,9 +254,8 @@ def calc_lfp_root_as_point_anisotropic(cell_x, cell_y, cell_z,
             p_[:] = pos + (pos - closest_point) * (r_limit[idx] - r) / r
 
         if np.sqrt(np.sum((p_ - closest_point)**2)) - r_limit[idx] > 1e-9:
-            print(p_, closest_point)
-
-            raise RuntimeError("Segment adjustment not working")
+            raise RuntimeError(f"Adjustment failed for segment {idx}",
+                               p_, closest_point)
 
         b[idx] = -2 * ((sigma[1] * sigma[2] * (p_[0] - xstart[idx])
                         * (xend[idx] - xstart[idx])) +
@@ -275,7 +276,7 @@ def calc_lfp_root_as_point_anisotropic(cell_x, cell_y, cell_z,
     if (i + iia + iib + iii + iiii).sum() != xstart.size:
         print(a, b, c)
         print(i, iia, iib, iii, iiii)
-        raise RuntimeError
+        raise RuntimeError  # WHY?
 
     mapping = np.zeros(xstart.size)
     mapping[i] = _anisotropic_line_source_case_i(a[i], c[i])
@@ -301,19 +302,17 @@ def calc_lfp_root_as_point_anisotropic(cell_x, cell_y, cell_z,
     r2_root = dx2_root + dy2_root + dz2_root
 
     # Go through and correct all (if any) root idxs that are too close
-    for close_idx in np.where(np.abs(r2_root) < 1e-6)[0]:
-        dx2_root[close_idx] += 0.001
-        r2_root[close_idx] += 0.001
-
-    for close_idx in np.where(r2_root < r_limit[rootinds]**2)[0]:
-        # For anisotropic media, the direction in which to move points matter.
-        # Radial distance between point source and electrode is scaled to r_lim
-        r2_scale_factor = r_limit[rootinds[close_idx]
-                                  ] * r_limit[rootinds[close_idx]
-                                              ] / r2_root[close_idx]
-        dx2_root[close_idx] *= r2_scale_factor
-        dy2_root[close_idx] *= r2_scale_factor
-        dz2_root[close_idx] *= r2_scale_factor
+    close_idx = r2_root < 1e-6
+    dx2_root[close_idx] += 0.001
+    r2_root[close_idx] += 0.001
+
+    close_idx = r2_root < r_limit[rootinds]**2
+    r2_scale_factor = r_limit[rootinds[close_idx]
+                              ] * r_limit[rootinds[close_idx]
+                                          ] / r2_root[close_idx]
+    dx2_root[close_idx] *= r2_scale_factor
+    dy2_root[close_idx] *= r2_scale_factor
+    dz2_root[close_idx] *= r2_scale_factor
 
     mapping[rootinds] = 1 / np.sqrt(sigma[1] * sigma[2] * dx2_root
                                     + sigma[0] * sigma[2] * dy2_root
@@ -476,7 +475,8 @@ def calc_lfp_root_as_point(cell_x, cell_y, cell_z, x, y, z, sigma, r_limit,
         r_root[r_root < r_limit[rootinds]
                ] = r_limit[rootinds][r_root < r_limit[rootinds]]
 
-    too_close_idxs = np.where(r2 < r_limit * r_limit)[0]
+    # avoid denominator approaching 0
+    too_close_idxs = r2 < (r_limit * r_limit)
     r2[too_close_idxs] = r_limit[too_close_idxs]**2
     l_ = h + deltaS
 
@@ -724,7 +724,7 @@ def calc_lfp_pointsource_moi(cell_x, cell_y, cell_z,
     dz2 = (z_ - cell_z_mid)**2
 
     dL2 = dx2 + dy2
-    # inds = np.where(dL2 + dz2 < r_limit * r_limit)[0]
+    # avoid denominator approaching 0
     inds = (dL2 + dz2) < (r_limit * r_limit)
     dL2[inds] = r_limit[inds] * r_limit[inds] - dz2[inds]
 
@@ -815,7 +815,7 @@ def calc_lfp_linesource_moi(cell_x, cell_y, cell_z,
     rs, _ = return_dist_from_segments(xstart, ystart, zstart,
                                       xend, yend, zend, pos)
     z0_ = z0.copy()
-    # z0_[np.where(rs < r_limit)] = r_limit[np.where(rs < r_limit)]
+    # avoid denominator approaching 0
     inds = rs < r_limit
     z0_[inds] = r_limit[inds]
 
@@ -986,8 +986,9 @@ def _omega(a_z):
     dz2 = (z_ - cell_z[rootinds, :].mean(axis=-1))**2
 
     dL2 = dx2 + dy2
-    inds = np.where(dL2 + dz2 < r_limit[rootinds] * r_limit[rootinds])[0]
-    dL2[inds] = r_limit[inds] * r_limit[inds] - dz2[inds]
+    # avoid denominator approaching 0
+    inds = (dL2 + dz2) < (r_limit[rootinds] * r_limit[rootinds])
+    dL2[inds] = r_limit[rootinds][inds] * r_limit[rootinds][inds] - dz2[inds]
 
     def _omega(dz):
         return 1 / np.sqrt(dL2 + dz * dz)

lfpykit/models.py

@@ -1263,7 +1263,8 @@ def _squeeze_cell_in_depth_direction(self):
             bad_comps, reason = self._return_comp_outside_slice()
             msg = ("Compartments {} of cell ({}) has cell.{} slice. "
                    "Increase squeeze_cell_factor, move or rotate cell."
-                   ).format(bad_comps, self.cell, reason)
+                   ).format(np.arange(self.cell.totnsegs)[bad_comps],
+                            self.cell, reason)
 
             raise RuntimeError(msg)
 
@@ -1281,18 +1282,18 @@ def _return_comp_outside_slice(self):
             Numpy array with the compartments that are outside the slice,
             and a string with additional information on the problem.
         """
-        zstart_above = np.where(self.cell.z[:, 0] > self.z_shift + self.h)[0]
-        zend_above = np.where(self.cell.z[:, -1] > self.z_shift + self.h)[0]
-        zend_below = np.where(self.cell.z[:, -1] < self.z_shift)[0]
-        zstart_below = np.where(self.cell.z[:, 0] < self.z_shift)[0]
+        zstart_above = self.cell.z[:, 0] > (self.z_shift + self.h)
+        zend_above = self.cell.z[:, -1] > (self.z_shift + self.h)
+        zend_below = self.cell.z[:, -1] < self.z_shift
+        zstart_below = self.cell.z[:, 0] < self.z_shift
 
-        if len(zstart_above) > 0:
+        if zstart_above.sum() > 0:
             return zstart_above, "zstart above"
-        if len(zstart_below) > 0:
+        if zstart_below.sum() > 0:
             return zstart_below, "zstart below"
-        if len(zend_above) > 0:
+        if zend_above.sum() > 0:
             return zend_above, "zend above"
-        if len(zend_below) > 0:
+        if zend_below.sum() > 0:
             return zend_below, "zend below"
         raise RuntimeError("This function should only be called if cell"
                            "extends outside slice")
@@ -1321,7 +1322,8 @@ def _test_cell_extent(self):
                 bad_comps, reason = self._return_comp_outside_slice()
                 msg = ("Compartments {} of cell ({}) has cell.{} slice "
                        "and argument squeeze_cell_factor is None."
-                       ).format(bad_comps, self.cell, reason)
+                       ).format(np.arange(self.cell.totnsegs)[bad_comps],
+                                self.cell, reason)
                 raise RuntimeError(msg)
         else:
             if self.verbose:

lfpykit/tests/test_module.py

@@ -581,25 +581,29 @@ def test_RecMEAElectrode_02(self):
         stick.z[true_bad_comp, 0] = 1000
         bad_comp, reason = MEA._return_comp_outside_slice()
         np.testing.assert_equal(reason, "zstart above")
-        np.testing.assert_equal(true_bad_comp, bad_comp)
+        np.testing.assert_equal(true_bad_comp,
+                                np.arange(stick.totnsegs)[bad_comp])
         stick.z[true_bad_comp, 0] = 100
 
         stick.z[true_bad_comp, 0] = -1000
         bad_comp, reason = MEA._return_comp_outside_slice()
         np.testing.assert_equal(reason, "zstart below")
-        np.testing.assert_equal(true_bad_comp, bad_comp)
+        np.testing.assert_equal(true_bad_comp,
+                                np.arange(stick.totnsegs)[bad_comp])
         stick.z[true_bad_comp, 0] = 100
 
         stick.z[true_bad_comp, -1] = 1000
         bad_comp, reason = MEA._return_comp_outside_slice()
         np.testing.assert_equal(reason, "zend above")
-        np.testing.assert_equal(true_bad_comp, bad_comp)
+        np.testing.assert_equal(true_bad_comp,
+                                np.arange(stick.totnsegs)[bad_comp])
         stick.z[true_bad_comp, -1] = 100
 
         stick.z[true_bad_comp, -1] = -1000
         bad_comp, reason = MEA._return_comp_outside_slice()
         np.testing.assert_equal(reason, "zend below")
-        np.testing.assert_equal(true_bad_comp, bad_comp)
+        np.testing.assert_equal(true_bad_comp,
+                                np.arange(stick.totnsegs)[bad_comp])
         stick.z[true_bad_comp, -1] = 100
 
     def test_RecMEAElectrode_03(self):