fix(GaussianNoiseDataset): explicitly operate on local_array

radiocosmology · Aug 16, 2024 · bae3d8c · bae3d8c
1 parent f3c2650
commit bae3d8c
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/draco/synthesis/noise.py b/draco/synthesis/noise.py
@@ -102,16 +102,17 @@ def process(self, data):
         data.redistribute("freq")
 
         # Replace visibilities with noise
-        dset = data[dataset_name][:]
+        dset = data[dataset_name][:].local_array
+        weight = data.weight[:].local_array
         if np.iscomplexobj(dset):
             random.complex_normal(
-                scale=tools.invert_no_zero(data.weight[:]) ** 0.5,
+                scale=tools.invert_no_zero(weight) ** 0.5,
                 out=dset,
                 rng=self.rng,
             )
         else:
             self.rng.standard_normal(out=dset)
-            dset *= tools.invert_no_zero(data.weight[:]) ** 0.5
+            dset *= tools.invert_no_zero(weight) ** 0.5
 
         # We need to loop to ensure the autos are real and have the correct variance
         if dataset_name == "vis":