fix test 'remove_transcription_site'

bigfish/stack/tests/test_postprocess.py

@@ -49,22 +49,53 @@ def test_identify_objects_in_region(ndim, mask_dtype):
 
 
 @pytest.mark.parametrize("ndim", [2, 3])
-def test_remove_transcription_site(ndim):
-    # simulate coordinates or rna and transcription sites
-    ts = [[0, 0, 0, 5, 0], [0, 0, 0, 5, 1]]
+@pytest.mark.parametrize("mask_dtype", [
+    np.uint8, np.uint16, np.int64, bool])
+def test_remove_transcription_site(ndim, mask_dtype):
+    # simulate mask and coordinates
+    nuc_mask = np.zeros((10, 10), dtype=mask_dtype)
+    nuc_mask[1:4, 1:5] = np.ones((3, 4), dtype=mask_dtype)
+
+    spots_in_nuc_out_ts = [[2, 1, 1, -1], [5, 2, 2, -1], [2, 3, 1, -1]]
+    spots_in_nuc_out_ts = np.array(spots_in_nuc_out_ts, dtype=np.int64)
+
+    spots_in_nuc_in_ts = [[2, 3, 3, 3], [5, 2, 4, 2], [2, 1, 3, 3]]
+    spots_in_nuc_in_ts = np.array(spots_in_nuc_in_ts, dtype=np.int64)
+
+    spots_out_nuc_out_foci = [[1, 0, 0, -1], [3, 7, 2, -1], [2, 1, 8, -1]]
+    spots_out_nuc_out_foci = np.array(spots_out_nuc_out_foci, dtype=np.int64)
+
+    spots_out_nuc_in_foci = [[1, 0, 4, 0], [3, 7, 7, 0], [2, 1, 5, 1]]
+    spots_out_nuc_in_foci = np.array(spots_out_nuc_in_foci, dtype=np.int64)
+
+    if ndim == 2:
+        spots_in_nuc_out_ts = spots_in_nuc_out_ts[:, 1:]
+        spots_in_nuc_in_ts = spots_in_nuc_in_ts[:, 1:]
+        spots_out_nuc_out_foci = spots_out_nuc_out_foci[:, 1:]
+        spots_out_nuc_in_foci = spots_out_nuc_in_foci[:, 1:]
+    spots = np.concatenate((spots_in_nuc_out_ts, spots_in_nuc_in_ts,
+                            spots_out_nuc_out_foci, spots_out_nuc_in_foci))
+    spots_out_ts = np.concatenate((spots_in_nuc_out_ts,
+                                   spots_out_nuc_out_foci,
+                                   spots_out_nuc_in_foci))
+
+    ts = [[2, 2, 4, 1, 2], [4, 2, 3, 2, 3]]
     ts = np.array(ts, dtype=np.int64)
-    rna_out = [[0, 0, 0, -1], [0, 0, 0, -1], [0, 0, 0, -1],
-               [0, 0, 0, -1], [0, 0, 0, 3]]
-    rna_out = np.array(rna_out, dtype=np.int64)
-    rna_in = [[0, 0, 0, 0], [0, 0, 0, 0], [0, 0, 0, 1]]
-    rna_in = np.array(rna_in, dtype=np.int64)
+
+    foci = [[0, 6, 7, 2, 0], [1, 2, 5, 1, 1]]
+    foci = np.array(foci, dtype=np.int64)
+
     if ndim == 2:
         ts = ts[:, 1:]
-        rna_in = rna_in[:, 1:]
-        rna_out = rna_out[:, 1:]
-    rna = np.concatenate((rna_in, rna_out))
+        foci = foci[:, 1:]
+    all_foci = np.concatenate((ts, foci))
 
     # test
-    rna_out_ = stack.remove_transcription_site_rna(rna, ts)
-    assert_array_equal(rna_out_, rna_out)
-    assert rna_out_.dtype == rna_out.dtype
+    spots_out_ts_, foci_, ts_ = stack.remove_transcription_site(
+        spots, all_foci, nuc_mask, ndim)
+    assert_array_equal(spots_out_ts_, spots_out_ts)
+    assert spots_out_ts_.dtype == spots_out_ts.dtype
+    assert_array_equal(foci_, foci)
+    assert foci_.dtype == foci.dtype
+    assert_array_equal(ts_, ts)
+    assert ts_.dtype == ts.dtype