From d9ff273e9d06653351e1d0cd05e667675305aa68 Mon Sep 17 00:00:00 2001
From: jarodhanko-crafco <114525548+jarodhanko-crafco@users.noreply.github.com>
Date: Tue, 3 Oct 2023 08:07:49 -0400
Subject: [PATCH] cache skeleton_image shape for path_label_image (#210)

It may be possible to update prune_paths to use this as well if the
skeleton_image is None (basically create new image and add paths that
are not to be pruned), but for now now that is being left as is.

Closes #208

---------

Co-authored-by: Juan Nunez-Iglesias <jni@fastmail.com>
---
 src/skan/csr.py           |  3 ++-
 src/skan/test/test_csr.py | 14 ++++++++++++++
 2 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/src/skan/csr.py b/src/skan/csr.py
index b3c2ead4..fac6b883 100644
--- a/src/skan/csr.py
+++ b/src/skan/csr.py
@@ -513,6 +513,7 @@ def __init__(
         self.distances = np.empty(self.n_paths, dtype=float)
         self._distances_initialized = False
         self.skeleton_image = None
+        self.skeleton_shape = skeleton_image.shape
         self.source_image = None
         self.degrees = np.diff(self.graph.indptr)
         self.spacing = (
@@ -618,7 +619,7 @@ def path_label_image(self):
             Image of the same shape as self.skeleton_image where each pixel
             has the value of its branch id + 1.
         """
-        image_out = np.zeros(self.skeleton_image.shape, dtype=int)
+        image_out = np.zeros(self.skeleton_shape, dtype=int)
         for i in range(self.n_paths):
             coords_to_wipe = self.path_coordinates(i)
             coords_idxs = tuple(np.round(coords_to_wipe).astype(int).T)
diff --git a/src/skan/test/test_csr.py b/src/skan/test/test_csr.py
index 5de462a4..b7370125 100644
--- a/src/skan/test/test_csr.py
+++ b/src/skan/test/test_csr.py
@@ -188,6 +188,7 @@ def test_fast_graph_center_idx():
     i = csr._fast_graph_center_idx(s)
     assert i == 1
 
+
 def test_sholl():
     s = csr.Skeleton(skeleton0)
     c, r, counts = csr.sholl_analysis(s, shells=np.arange(0, 5, 1.5))
@@ -241,3 +242,16 @@ def test_zero_degree_nodes():
     np.testing.assert_equal(
             np.ravel(s.coordinates), [0, 1, 3, 5, 7, 9, 11, 12]
             )
+
+
+def test_skeleton_path_image_no_keep_image():
+    """See #208: "Skeleton.path_label_image requires keep_images=True."
+
+    Before PR #210, it was an implicit requirement of path_label_image
+    to create a Skeleton with keep_images=True. However, we only needed
+    the shape of the image. This makes sure that the method works even
+    when keep_images=False.
+    """
+    s = csr.Skeleton(skeleton2, keep_images=False)
+    pli = s.path_label_image()
+    assert np.max(pli) == s.n_paths