From 5f16d6fa02687b8c69e51c8ec2f43fe14c2c604f Mon Sep 17 00:00:00 2001
From: alex-l-kong <31424707+alex-l-kong@users.noreply.github.com>
Date: Thu, 21 Nov 2024 12:44:24 -0800
Subject: [PATCH] Normalize by maximum pixel in case of Rosetta tiling for
 sparse signal  (#472)

* Correct normalization for sparse signal in Rosetta tiling

* Ensure perc_source does not attempt to check == 0 if percent_norm not set (meaning perc_source is None)

---------

Co-authored-by: Alex Kong <alkong@stanford.edu>
---
 src/toffy/rosetta.py  |  9 +++++++
 tests/rosetta_test.py | 58 +++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 67 insertions(+)

diff --git a/src/toffy/rosetta.py b/src/toffy/rosetta.py
index 50728aa6..9c32584d 100644
--- a/src/toffy/rosetta.py
+++ b/src/toffy/rosetta.py
@@ -558,6 +558,10 @@ def add_source_channel_to_tiled_image(
     # get percentile of source row if percent_norm set, otherwise leave unset
     perc_source = np.percentile(source_scaled, percent_norm) if percent_norm else None
 
+    # normalize by max if percent_norm is 0 for source in case of sparse signal
+    if percent_norm and perc_source == 0:
+        perc_source = np.percentile(source_scaled, 100)
+
     # confirm tiled images have expected shape
     tiled_images = io_utils.list_files(tiled_img_dir)
     test_file = io.imread(os.path.join(tiled_img_dir, tiled_images[0]))
@@ -577,6 +581,11 @@ def add_source_channel_to_tiled_image(
         perc_ratio = 1
         if percent_norm:
             perc_tile = np.percentile(current_tile, percent_norm)
+
+            # normalize by max if percent_norm is 0 for tile in case of sparse signal
+            if perc_tile == 0:
+                perc_tile = np.percentile(current_tile, 100)
+
             perc_ratio = perc_source / perc_tile
 
         rescaled_source = source_scaled / perc_ratio
diff --git a/tests/rosetta_test.py b/tests/rosetta_test.py
index e518c13a..0f34bd8e 100644
--- a/tests/rosetta_test.py
+++ b/tests/rosetta_test.py
@@ -501,6 +501,7 @@ def test_create_tiled_comparison(dir_num, channel_subset, img_size_scale):
 @parametrize("percent_norm", [98, None])
 @parametrize("img_scale", [1, 0.25])
 def test_add_source_channel_to_tiled_image(percent_norm, img_scale):
+    # test compensation without percent norm correction
     with tempfile.TemporaryDirectory() as top_level_dir:
         num_fovs = 5
         num_chans = 4
@@ -547,6 +548,63 @@ def test_add_source_channel_to_tiled_image(percent_norm, img_scale):
                 img_scale * tiled_shape[1],
             )
 
+    # test compensation with percent norm correction
+    with tempfile.TemporaryDirectory() as top_level_dir:
+        num_fovs = 1
+        num_chans = 2
+        im_size = 20
+
+        # create directory containing raw images
+        raw_dir = os.path.join(top_level_dir, "raw_dir")
+        os.makedirs(raw_dir)
+
+        fovs, chans = test_utils.gen_fov_chan_names(num_fovs=num_fovs, num_chans=num_chans)
+        filelocs, data_xr = test_utils.create_paired_xarray_fovs(
+            raw_dir, fovs, chans, img_shape=(im_size, im_size), fills=True
+        )
+
+        # create sparse signal source
+        for fov in data_xr.fovs.values:
+            fileloc_ind = 0
+            for chan in data_xr.channels.values:
+                data_xr_sub = np.zeros(data_xr.loc[fov, ..., chan].shape)
+                data_xr_sub[0, 0] = 1e-20
+                io.imsave(filelocs[fov][fileloc_ind] + ".tiff", data_xr_sub)
+                fileloc_ind += 1
+
+        # create directory containing sparse stitched images
+        tiled_shape = (im_size * 3, im_size * num_fovs)
+        tiled_dir = os.path.join(top_level_dir, "tiled_dir")
+        os.makedirs(tiled_dir)
+        for i in range(2):
+            vals = np.zeros(tiled_shape)
+            vals[0, 0] = 1e-20
+            if img_scale != 1:
+                vals = rescale_images(vals, scale=img_scale)
+            fname = os.path.join(tiled_dir, f"tiled_image_{i}.tiff")
+            image_utils.save_image(fname, vals)
+
+        output_dir = os.path.join(top_level_dir, "output_dir")
+        os.makedirs(output_dir)
+        rosetta.add_source_channel_to_tiled_image(
+            raw_img_dir=raw_dir,
+            tiled_img_dir=tiled_dir,
+            output_dir=output_dir,
+            source_channel="chan1",
+            max_img_size=im_size,
+            percent_norm=percent_norm,
+            img_size_scale=img_scale,
+        )
+
+        # each image should now have an extra row added on top
+        tiled_images = io_utils.list_files(output_dir)
+        for im_name in tiled_images:
+            image = io.imread(os.path.join(output_dir, im_name))
+            assert image.shape == (
+                img_scale * (tiled_shape[0] + im_size),
+                img_scale * tiled_shape[1],
+            )
+
 
 @parametrize("fovs", [None, ["fov1"]])
 @parametrize("replace", [True, False])