Visualize PPCA covs

cwindolf · Nov 21, 2024 · b3d55eb · b3d55eb
1 parent 64721b4
commit b3d55eb
Showing 1 changed file with 11 additions and 4 deletions.
diff --git a/src/dartsort/cluster/gaussian_mixture.py b/src/dartsort/cluster/gaussian_mixture.py
@@ -1671,10 +1671,17 @@ def pick_channels(self, active_chans, aweights_sum):
 
         assert False
 
-    def marginal_covariance(self, channels, cache_key=None, device=None):
-        ncov = self.noise.marginal_covariance(
-            channels, cache_key=cache_key, device=device
-        )
+    def marginal_covariance(
+        self, channels, cache_key=None, device=None, signal_only=False
+    ):
+        if signal_only:
+            ncov = operators.ZeroLinearOperator(
+                2 * [channels.numel() * self.noise.rank]
+            )
+        else:
+            ncov = self.noise.marginal_covariance(
+                channels, cache_key=cache_key, device=device
+            )
         zero_signal = (
             self.cov_kind == "zero" or self.cov_kind == "ppca" and not self.ppca_rank
         )