downweight extrapolated edges in model2comps

ratt-ru · Oct 16, 2024 · 6590170 · 6590170
1 parent 3acfded
commit 6590170
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 12 deletions.
diff --git a/pfb/operators/gridder.py b/pfb/operators/gridder.py
@@ -770,15 +770,12 @@ def compute_residual(dsl,
     tassign = time() - ti
 
     # we only need to write MODEL and RESIDUAL
-    # ds = ds[['RESIDUAL','MODEL']]
-    for var in ds.data_vars:
-        if var not in ['RESIDUAL', 'MODEL']:
-            ds = ds.drop_vars(var)
+    ds = ds[['RESIDUAL','MODEL']]
 
     # save
     ti = time()
     with cf.ThreadPoolExecutor(max_workers=1) as executor:
-        future = executor.submit(dataset_to_zarr, ds, output_name)
+        future = executor.submit(ds.to_zarr, output_name, mode='a')
     twrite = time() - ti
 
     ttot = time() - tii

diff --git a/pfb/utils/stokes2im.py b/pfb/utils/stokes2im.py
@@ -249,19 +249,16 @@ def single_stokes_image(
                                  usign=1.0 if flip_u else -1.0,
                                  vsign=1.0 if flip_v else -1.0)
 
-        imwgt = counts_to_weights(
+        weight = counts_to_weights(
             counts,
             uvw,
             freq,
+            weight,
             nx, ny,
             cell_rad, cell_rad,
             opts.robustness,
             usign=1.0 if flip_u else -1.0,
             vsign=1.0 if flip_v else -1.0)
-        if weight is not None:
-            weight *= imwgt
-        else:
-            weight = imwgt
 
     wsum = weight[~flag].sum()
 

diff --git a/pfb/workers/model2comps.py b/pfb/workers/model2comps.py
@@ -186,7 +186,7 @@ def _model2comps(**kw):
             model = np.concatenate((mlow[:, None], model), axis=1)
             mfreqs = np.concatenate((np.array((flow,)), mfreqs))
             # TODO - duplicate first non-null value?
-            wsums = np.concatenate((wsums[:, Ilow][:, None], wsums),
+            wsums = np.concatenate((0.5*wsums[:, Ilow][:, None], wsums),
                                    axis=1)
             nband = mfreqs.size
         if fhigh > mfreqs.max():
@@ -203,7 +203,7 @@ def _model2comps(**kw):
             mhigh = slopes * fhigh + intercepts
             model = np.concatenate((model, mhigh[:, None]), axis=1)
             mfreqs = np.concatenate((mfreqs, np.array((fhigh,))))
-            wsums = np.concatenate((wsums, wsums[:, Ihigh][:, None]),
+            wsums = np.concatenate((wsums, 0.5*wsums[:, Ihigh][:, None]),
                                    axis=1)
             nband = mfreqs.size