From 7e71f66389e7667426964bc65910f8a3ee77c4ce Mon Sep 17 00:00:00 2001
From: jnsbck <65561470+jnsbck@users.noreply.github.com>
Date: Thu, 21 Nov 2024 16:55:48 +0100
Subject: [PATCH] Morph cell pickling (#525)

* mv: add io module and move swc things there

* fix: isort and import tests fix

* fix: fix module prefix for Cell

* fix: fix circular imports

* enh: add morph test case to pickling

* fix: make tests pass
---
 jaxley/io/swc.py           |  3 ++-
 jaxley/utils/cell_utils.py | 35 ++++++++++++++++++++++++-----------
 tests/test_pickle.py       |  7 +++++--
 3 files changed, 31 insertions(+), 14 deletions(-)

diff --git a/jaxley/io/swc.py b/jaxley/io/swc.py
index ff131e62..c1198451 100644
--- a/jaxley/io/swc.py
+++ b/jaxley/io/swc.py
@@ -12,6 +12,7 @@
 from jaxley.utils.cell_utils import (
     _build_parents,
     _compute_pathlengths,
+    _padded_radius_generating_fn,
     _radius_generating_fns,
     _split_into_branches_and_sort,
     build_radiuses_from_xyzr,
@@ -69,7 +70,7 @@ def swc_to_jaxley(
         parents[1:] += 1
         parents = parents.tolist()
         pathlengths = [0.1] + pathlengths
-        radius_fns = [lambda x: content[0, 5] * np.ones_like(x)] + radius_fns
+        radius_fns = [_padded_radius_generating_fn(content[0, 5])] + radius_fns
         sorted_branches = [[0]] + sorted_branches
 
         # Type of padded section is assumed to be of `custom` type:
diff --git a/jaxley/utils/cell_utils.py b/jaxley/utils/cell_utils.py
index d47ab6ea..ba055eab 100644
--- a/jaxley/utils/cell_utils.py
+++ b/jaxley/utils/cell_utils.py
@@ -1,6 +1,7 @@
 # This file is part of Jaxley, a differentiable neuroscience simulator. Jaxley is
 # licensed under the Apache License Version 2.0, see <https://www.apache.org/licenses/>
 
+from functools import partial
 from math import pi
 from typing import Callable, Dict, List, Optional, Tuple, Union
 from warnings import warn
@@ -186,6 +187,28 @@ def _radius_generating_fns(
     return radius_fns
 
 
+def _padded_radius(loc: float, radiuses: np.ndarray) -> float:
+    return radiuses * np.ones_like(loc)
+
+
+def _radius(loc: float, cutoffs: np.ndarray, radiuses: np.ndarray) -> float:
+    """Function which returns the radius via linear interpolation.
+
+    Defined outside of `_radius_generating_fns` to allow for pickling of the resulting
+    Cell object."""
+    index = np.digitize(loc, cutoffs, right=False)
+    left_rad = radiuses[index - 1]
+    right_rad = radiuses[index]
+    left_loc = cutoffs[index - 1]
+    right_loc = cutoffs[index]
+    loc_within_bin = (loc - left_loc) / (right_loc - left_loc)
+    return left_rad + (right_rad - left_rad) * loc_within_bin
+
+
+def _padded_radius_generating_fn(radiuses: np.ndarray) -> Callable:
+    return partial(_padded_radius, radiuses=radiuses)
+
+
 def _radius_generating_fn(radiuses: np.ndarray, each_length: np.ndarray) -> Callable:
     # Avoid division by 0 with the `summed_len` below.
     each_length[each_length < 1e-8] = 1e-8
@@ -201,17 +224,7 @@ def _radius_generating_fn(radiuses: np.ndarray, each_length: np.ndarray) -> Call
     if len(radiuses) == 1:
         radiuses = np.tile(radiuses, 2)
 
-    def radius(loc: float) -> float:
-        """Function which returns the radius via linear interpolation."""
-        index = np.digitize(loc, cutoffs, right=False)
-        left_rad = radiuses[index - 1]
-        right_rad = radiuses[index]
-        left_loc = cutoffs[index - 1]
-        right_loc = cutoffs[index]
-        loc_within_bin = (loc - left_loc) / (right_loc - left_loc)
-        return left_rad + (right_rad - left_rad) * loc_within_bin
-
-    return radius
+    return partial(_radius, cutoffs=cutoffs, radiuses=radiuses)
 
 
 def _compute_pathlengths(
diff --git a/tests/test_pickle.py b/tests/test_pickle.py
index a8652312..d12407eb 100644
--- a/tests/test_pickle.py
+++ b/tests/test_pickle.py
@@ -1,6 +1,7 @@
 # This file is part of Jaxley, a differentiable neuroscience simulator. Jaxley is
 # licensed under the Apache License Version 2.0, see <https://www.apache.org/licenses/>
 
+import os
 import pickle
 
 import pytest
@@ -9,11 +10,13 @@
 from jaxley.channels import HH
 from jaxley.synapses import IonotropicSynapse
 
-# create modules
+# create modules (cannot use fixtures for pickling, since they rely on local func defs)
 comp = jx.Compartment()
 branch = jx.Branch(comp, 4)
 cell = jx.Cell([branch] * 3, [-1, 0, 0])
 net = jx.Network([cell] * 2)
+fname = os.path.join(os.path.dirname(__file__), "swc_files", "morph.swc")
+morph_cell = jx.read_swc(fname, nseg=1, max_branch_len=2_000, assign_groups=True)
 
 # insert mechanisms
 net.cell(0).branch("all").insert(HH())
@@ -24,7 +27,7 @@
 
 
 @pytest.mark.parametrize(
-    "module", [comp, branch, cell, net], ids=lambda x: x.__class__.__name__
+    "module", [comp, branch, cell, morph_cell, net], ids=lambda x: x.__class__.__name__
 )
 def test_pickle(module):
     pickled = pickle.dumps(module)