add stepwise functionality, tests, notebook

Author: kjohnsen

notebooks/stepwise.ipynb

@@ -37,11 +37,12 @@ nbdev = "^2.3.12"
 brian2 = "^2.5.4"
 nbmake = "^1.5.0"
 jupyter = "^1.0.0"
+pytest-xdist = "^3.5.0"
 
 [build-system]
 requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"
 
 [tool.pytest.ini_options]
 addopts = "-q --nbmake"
-testpaths = ["tests", "README.ipynb"]
+testpaths = ["tests", "README.ipynb", "notebooks/real_time.ipynb"]

poetry.lock

@@ -26,6 +26,9 @@
         ([1, 5, 9, 12, 17], [1, 3, 5, 9, 12], [], False),
         ([2, 4, 6, 23, 25], [7, 10, 15, 21, 25], [12, 20, 24], True),
         ([2, 3, 4, 5, 6], [50, 60, 70, 80, 90], [70, 100], False),
+        ([0], [6], [6], True),
+        ([6], [6], [6], False),
+        ([6], [0], [6], False),
     ],
 )
 def test_calculate(t_ampa, t_gaba, t_eval, success, n_elec, seed):
@@ -38,8 +41,6 @@ def test_calculate(t_ampa, t_gaba, t_eval, success, n_elec, seed):
     I_gaba = -np.arange(len(t_gaba)) - 2
     I_gaba = I_gaba[:, np.newaxis] - np.arange(n_src)
 
-    # xs = ys = zs = np.array([1, 2, 3])
-    # source_coords = np.column_stack((xs, ys, zs))
     source_coords = rng.uniform(-10, 10, (n_src, 3))
     elec_coords = rng.uniform(-500, 500, (n_elec, 3))
 
@@ -56,7 +57,7 @@ def test_calculate(t_ampa, t_gaba, t_eval, success, n_elec, seed):
             )
         ), "LFP increasing order does not match expectation"
     else:
-        with pytest.raises(ValueError):
+        with pytest.raises(Exception):
             calc.calculate(t_eval, t_ampa, I_ampa, t_gaba, I_gaba)
 
 

pyproject.toml

@@ -80,6 +80,8 @@ def n_sources(self) -> int:
     def _interp_currents(self, t_ms, I, delay_ms, t_eval_ms):
         if not np.all(np.diff(t_ms) > 0):
             raise ValueError("t_ms must be monotonically increasing")
+        if len(t_ms) == 0:
+            return np.zeros((len(t_eval_ms), self.n_sources))
 
         t_eval_delayed = np.subtract(t_eval_ms, delay_ms)
         t_needed = (np.min(t_eval_delayed), np.max(t_eval_delayed))
@@ -99,7 +101,12 @@ def _interp_currents(self, t_ms, I, delay_ms, t_eval_ms):
                     f"Needed [{t_needed[0]}, {t_needed[1]}] ms, "
                     f"provided [{t_provided[0]}, {t_provided[1]}] ms."
                 )
-        I_interp = PchipInterpolator(t_ms, I, extrapolate=False)(t_eval_delayed)
+        if len(t_ms) > 1:
+            interpolator = PchipInterpolator(t_ms, I, extrapolate=False)
+        elif len(t_ms) == 1:
+            interpolator = lambda t_eval: (t_eval == t_ms[0]) * I[0:1]
+
+        I_interp = interpolator(t_eval_delayed)
         assert I_interp.shape == (len(t_eval_ms), self.n_sources)
         I_interp[np.isnan(I_interp)] = 0
         return I_interp
@@ -112,6 +119,7 @@ def calculate(
         t_gaba_ms: np.ndarray,
         I_gaba: np.ndarray,
         normalize: bool = True,
+        wsum_mean_std_for_norm: tuple[np.ndarray, np.ndarray] = None,
     ) -> np.ndarray:
         """Calculate WSLFP at requested times for initialized coordinates given currents
 
@@ -124,8 +132,12 @@ def calculate(
                 in milliseconds
             I_gaba (np.ndarray): GABAergic currents, shape (len(t_ampa_ms), n_sources)
             normalize (bool, optional): Whether to normalize to mean of 0 and variance of 1.
-                The main reason not to normalize is if you are computing one time step at a time.
-                Defaults to True.
+                The main reason not to normalize is if you are computing one time step at a time
+                (see `notebooks/stepwise.ipynb`).  Defaults to True.
+            wsum_mean_std_for_norm (tuple[np.ndarray, np.ndarray], optional): If provided,
+                the mean and standard deviation of the weighted sum term are used to normalize
+                before a realistic amplitude is applied. (see `notebooks/stepwise.ipynb`).
+                Defaults to None.
 
         Returns:
             np.ndarray: (len(t_eval_ms), n_elec) array of WSLFP at requested times for
@@ -153,18 +165,33 @@ def calculate(
         assert wsum.shape == (len(t_eval_ms), self.n_elec, self.n_sources)
         wsum = np.sum(wsum, axis=2)
         assert wsum.shape == (len(t_eval_ms), self.n_elec)
-        if normalize:
-            wsum = wsum - np.mean(wsum, axis=0)
-            if len(t_eval_ms) > 1:
-                wsum /= np.std(wsum, axis=0)
-            assert wsum.shape == (len(t_eval_ms), self.n_elec)
-
-        assert np.allclose(wsum.mean(axis=0), 0)
-        if len(t_eval_ms) > 1:
-            assert np.allclose(wsum.std(axis=0), 1)
 
-        wsum *= np.abs(self.amp_uV.mean(axis=1))
-        return wsum
+        if normalize:
+            if wsum_mean_std_for_norm:
+                wsum_mean, wsum_std = wsum_mean_std_for_norm
+            else:
+                wsum_mean = np.mean(wsum, axis=0)
+                if len(t_eval_ms) > 1:
+                    wsum_std = np.std(wsum, axis=0)
+                else:
+                    wsum_std = 1
+
+            lfp = (wsum - wsum_mean) / wsum_std
+            assert lfp.shape == (len(t_eval_ms), self.n_elec)
+
+            if not wsum_mean_std_for_norm:
+                assert np.allclose(lfp.mean(axis=0), 0)
+                if len(t_eval_ms) > 1:
+                    assert np.allclose(lfp.std(axis=0), 1)
+
+            lfp *= np.abs(self.amp_uV.mean(axis=1))
+        else:
+            lfp = wsum
+
+        if normalize and wsum_mean_std_for_norm:
+            return lfp, wsum
+        else:
+            return lfp
 
 
 def from_rec_radius_depth(