fix for add_poisson_drive not working with int when cell_specific=True

dylansdaniels · Jul 10, 2024 · c79813c · c79813c
1 parent 1aa6cbd
commit c79813c
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/hnn_core/network.py b/hnn_core/network.py
@@ -745,7 +745,7 @@ def add_poisson_drive(self, name, *, tstart=0, tstop=None, rate_constant,
                                  f"n_drive_cells='n_cells'. Got cell_specific"
                                  f" cell_specific={cell_specific} and "
                                  f"n_drive_cells={n_drive_cells}.")
-        elif isinstance(rate_constant, float):
+        elif isinstance(rate_constant, (float, int)):
             if cell_specific:
                 rate_constant = {cell_type: rate_constant for cell_type in
                                  target_populations}

diff --git a/hnn_core/tests/test_drives.py b/hnn_core/tests/test_drives.py
@@ -469,6 +469,7 @@ def test_drive_random_state():
 
 @pytest.mark.parametrize("rate_constant,cell_specific,n_drive_cells",
                          [(2, False, 1), (2.0, False, 1),
+                          (2, False, 1), (2.0, False, 1),
                           ])
 def test_add_poisson_drive(setup_net, rate_constant, cell_specific,
                            n_drive_cells):