From 3c5e00aa075d5d29da1bb8c31d102b8f77fe6cef Mon Sep 17 00:00:00 2001
From: gabed <78431504+gdolsten@users.noreply.github.com>
Date: Wed, 30 Aug 2023 15:19:12 -0400
Subject: [PATCH] Insulation score should weight by # of NaNs in each diagonal

I had trouble fully understanding the insul_diamond function but I thought we could add something like this so that insulation score at regions with NaN covering some of the diamond are not punished
---
 cooltools/api/insulation.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/cooltools/api/insulation.py b/cooltools/api/insulation.py
index 2c2a87e8..d17527cb 100644
--- a/cooltools/api/insulation.py
+++ b/cooltools/api/insulation.py
@@ -115,7 +115,12 @@ def insul_diamond(
 
         i = diag_pixels.bin1_id.values - lo_bin_id
         j = diag_pixels.bin2_id.values - lo_bin_id
-
+        ####
+        # Weight each bin by the fraction of NaNs at each diagonal
+        diag_pixels['distance_from_diag'] = (diag_pixels['bin1_id'] - diag_pixels['bin2_id']).abs()
+        diag_pixels['valid_pixel_mask'] = ~diag_pixels["balanced"].isnull().values
+        diag_weighting_factor = diag_pixels.groupby(['bin1_id', 'distance_from_diag'])['valid_pixel_mask'].mean()
+        ####
         for i_shift in range(0, window):
             for j_shift in range(0, window):
                 if i_shift + j_shift < ignore_diags:
@@ -132,9 +137,10 @@ def insul_diamond(
                 )
 
                 if clr_weight_name:
+                    ### Weight sum_balanced by the number of non-Nans in diagonal
                     sum_balanced += np.bincount(
                         i[mask & valid_pixel_mask] + i_shift,
-                        diag_pixels["balanced"].values[mask & valid_pixel_mask],
+                        diag_pixels["balanced"].values[mask & valid_pixel_mask] / diag_weighting_factor.loc[[i, i_shift + j_shift]],
                         minlength=N,
                     )