rechunk beam coords on read

ratt-ru · Dec 11, 2023 · ebd49bc · ebd49bc
1 parent 9467ca7
commit ebd49bc
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 4 deletions.
diff --git a/pfb/operators/gridder.py b/pfb/operators/gridder.py
@@ -609,7 +609,12 @@ def image_data_products(uvw,
         wcount = mask.sum()
         if wcount:
             ovar = ressq.sum()/wcount
-            wgt = (l2reweight_dof + 1)/(l2reweight_dof + ressq/ovar)/ovar
+            l2wgt = (l2reweight_dof + 1)/(l2reweight_dof + ressq/ovar)/ovar
+        else:
+            l2wgt = None
+    else:
+        l2wgt = None
+
 
     # we usually want to re-evaluate this since the robustness may change
     if robustness is not None:
@@ -624,6 +629,12 @@ def image_data_products(uvw,
             cellx, celly,
             robustness)
 
+        # this is necessitated by the way the weighting is done i.e.
+        # wgt applied to vis@init so only imwgt*l2wgt needs to be applied
+        if l2wgt is not None:
+            imwgt *= l2wgt
+
+        # wgt*imwgt*l2wgt required for PSF
         wgt *= imwgt
 
     if do_weight:
@@ -636,7 +647,7 @@ def image_data_products(uvw,
         uvw=uvw,
         freq=freq,
         vis=vis,
-        wgt=imwgt,  # data have already been naturally weighted
+        wgt=imwgt,  # data already naturally weighted
         mask=mask,
         npix_x=nx, npix_y=ny,
         pixsize_x=cellx, pixsize_y=celly,

diff --git a/pfb/workers/grid.py b/pfb/workers/grid.py
@@ -143,9 +143,15 @@ def _grid(xdsi=None, **kw):
     if xdsi is not None:
         xds = []
         for ds in xdsi:
-            xds.append(ds.chunk({'row':-1}))
+            xds.append(ds.chunk({'row':-1,
+                                 'chan': -1,
+                                 'l_beam': -1,
+                                 'm_beam': -1}))
     else:
-        xds = xds_from_zarr(xds_name, chunks={'row': -1})
+        xds = xds_from_zarr(xds_name, chunks={'row': -1,
+                                              'chan': -1,
+                                              'l_beam': -1,
+                                              'm_beam': -1})
     # dds contains image space products including imaging weights and uvw
     dds_name = f'{basename}_{opts.postfix}.dds'