From cdfa35b12a817b133c4bcb79d075eb52df587b5b Mon Sep 17 00:00:00 2001
From: Richard Shaw <richard@phas.ubc.ca>
Date: Sat, 7 Mar 2020 17:52:18 -0800
Subject: [PATCH] feat(sensitivity): add frac lost to SystemSensitivity

---
 draco/analysis/sensitivity.py |  9 ++++++---
 draco/core/containers.py      | 10 ++++++++++
 2 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/draco/analysis/sensitivity.py b/draco/analysis/sensitivity.py
index 4c8422529..e6d5b0278 100644
--- a/draco/analysis/sensitivity.py
+++ b/draco/analysis/sensitivity.py
@@ -225,11 +225,11 @@ def process(self, data):
         dnu = np.median(data.index_map["freq"]["width"]) * 1e6
 
         if ("flags" in data) and ("frac_lost" in data["flags"]):
-            frac_lost = data["flags"]["frac_lost"][:][:, np.newaxis, :]
+            frac_lost = data["flags"]["frac_lost"][:]
         else:
-            frac_lost = 0.0
+            frac_lost = np.zeros((1, 1), dtype=np.float32)
 
-        nint = dnu * tint * (1.0 - frac_lost)
+        nint = dnu * tint * (1.0 - frac_lost[:, np.newaxis, :])
 
         # Normalize by the number of independent samples
         # and the total number of baselines squared
@@ -257,4 +257,7 @@ def process(self, data):
         # Save the total number of baselines that were averaged in the weight dataset
         metrics.weight[:] = counter
 
+        # Save the fraction of missing samples
+        metrics.frac_lost[:] = frac_lost
+
         return metrics
diff --git a/draco/core/containers.py b/draco/core/containers.py
index bf4061ac9..3beb340a3 100644
--- a/draco/core/containers.py
+++ b/draco/core/containers.py
@@ -796,6 +796,12 @@ class SystemSensitivity(TODContainer):
             "initialise": True,
             "distributed": True,
         },
+        "frac_lost": {
+            "axes": ["freq", "time"],
+            "dtype": np.float32,
+            "initialise": True,
+            "distributed": True,
+        },
     }
 
     @property
@@ -810,6 +816,10 @@ def radiometer(self):
     def weight(self):
         return self.datasets["weight"]
 
+    @property
+    def frac_lost(self):
+        return self.datasets["frac_lost"]
+
     @property
     def freq(self):
         return self.index_map["freq"]["centre"]