xarray complaining when trying to combine WEIGHT and IMAGING_WEIGHT_SPECTRUM columns

landmanbester commented 3 years ago

I am trying to add an option to apply additional imaging weights as requested in https://github.com/ratt-ru/pfb-clean/issues/30. I am running into a problem when trying to combine WEIGHT and IMAGING_WEIGHT_SPECTRUM because they do not have the same dimensions. When I try to broadcasts WEIGHTS over frequency axis and combine them with IMAGING_WEIGHT_SPECTRUM, i.e. doing something like

                weights = getattr(ds, self.weight_column).data
                if len(weights.shape) < 3:
                    weights = da.broadcast_to(weights[:, None, :], flag.shape, chunks=flag.chunks)

                if self.imaging_weight_column is not None:
                    imaging_weights = getattr(ds, self.imaging_weight_column)
                    if len(imaging_weights.shape) < 3:
                        imaging_weights = da.broadcast_to(imaging_weights[:, None, :], data.shape, chunks=data.chunks)

                    print(weights.shape, imaging_weights.shape)
                    weightsxx = imaging_weights[:, : 0] * weights[:, : 0]
                    weightsyy = imaging_weights[:, : -1] * weights[:, : -1]
                else:
                    weightsxx = weights[:, :, 0]
                    weightsyy = weights[:, :, -1]

results in the following error

Traceback (most recent call last):
  File "/home/landman/venvs/pfbenv/bin/ssclean.py", line 7, in <module>
    exec(compile(f.read(), __file__, 'exec'))
  File "/home/landman/Projects/pfb-clean/scripts/ssclean.py", line 457, in <module>
    main(args)
  File "/home/landman/Projects/pfb-clean/scripts/ssclean.py", line 170, in main
    psf_array = R.make_psf()
  File "/home/landman/Projects/pfb-clean/pfb/operators/gridder.py", line 349, in make_psf
    weights = weightsxx + weightsyy
  File "/home/landman/venvs/pfbenv/lib/python3.6/site-packages/xarray/core/dataarray.py", line 2685, in func
    self, other = align(self, other, join=align_type, copy=False)
  File "/home/landman/venvs/pfbenv/lib/python3.6/site-packages/xarray/core/alignment.py", line 327, in align
    % (dim, sizes)
ValueError: arguments without labels along dimension 'IMAGING_WEIGHT_SPECTRUM-1' cannot be aligned because they have different dimension sizes: {0, 31}

which I guess happens because broadcast_to doesn't actually make a copy for each channel. Any ideas on this @sjperkins?

sjperkins commented 3 years ago

Probably caused because weights is a dask array while imaging_weights is a xarray DataArray.

Also IMAGING_WEIGHT is getting default dimension labels (IMAGING_WEIGHT-1, IMAGING_WEIGHT-2) which don't correspond to the usual (row, chan, corr).

I would recommend dealing with it all in dask or xarray land.

landmanbester commented 3 years ago

Sorry, I missed the .data on

imaging_weights = getattr(ds, self.imaging_weight_column).data

I now get

Traceback (most recent call last):
  File "/home/landman/venvs/pfbenv/bin/ssclean.py", line 7, in <module>
    exec(compile(f.read(), __file__, 'exec'))
  File "/home/landman/Projects/pfb-clean/scripts/ssclean.py", line 457, in <module>
    main(args)
  File "/home/landman/Projects/pfb-clean/scripts/ssclean.py", line 170, in main
    psf_array = R.make_psf()
  File "/home/landman/Projects/pfb-clean/pfb/operators/gridder.py", line 348, in make_psf
    weights = weightsxx + weightsyy
  File "/home/landman/venvs/pfbenv/lib/python3.6/site-packages/dask/array/core.py", line 1814, in __add__
    return elemwise(operator.add, self, other)
  File "/home/landman/venvs/pfbenv/lib/python3.6/site-packages/dask/array/core.py", line 3886, in elemwise
    broadcast_shapes(*shapes)
  File "/home/landman/venvs/pfbenv/lib/python3.6/site-packages/dask/array/core.py", line 3847, in broadcast_shapes
    "shapes {0}".format(" ".join(map(str, shapes)))
ValueError: operands could not be broadcast together with shapes (758160, 0, 4) (758160, 31, 4)

Any quick fixes you can recommend?

sjperkins commented 3 years ago

Its probably failing because of the 0 in dimension 2 of the first shape. It should be a 1,