diff --git a/misc/Dockerfile b/misc/Dockerfile
index d72533347a656e80e8c0b30bded81874284c9da2..82e729f4537d3adc4c7f7217b6e5ed59d0bbe3c3 100644
--- a/misc/Dockerfile
+++ b/misc/Dockerfile
@@ -2,8 +2,8 @@ FROM kernsuite/base:latest
ENV DEBIAN_FRONTEND noninteractive
RUN docker-apt-install dysco python3-casacore
RUN apt-get update -qq && apt-get install -qq python3-pip python3-matplotlib git
-RUN pip3 install scipy pybind11 git+https://gitlab.mpcdf.mpg.de/ift/nifty.git@NIFTy_7
-RUN pip3 install git+https://gitlab.mpcdf.mpg.de/mtr/ducc.git@ducc0
+RUN pip3 install scipy pybind11 git+https://gitlab.mpcdf.mpg.de/ift/nifty.git@NIFTy_8
+RUN pip3 install ducc0
COPY . /resolve
RUN cd resolve && python3 setup.py install
ENTRYPOINT ["python3", "resolve/misc/ms2npz.py"]
diff --git a/misc/ms2npz.py b/misc/ms2npz.py
old mode 100644
new mode 100755
index a9e6ee779e24909cb4e98cad2ea38135d6ce5ab3..7d8f427ca75c1072fdcbada8b1f50a77bcc9c8b4
--- a/misc/ms2npz.py
+++ b/misc/ms2npz.py
@@ -1,11 +1,13 @@
+#!/usr/bin/python3
# SPDX-License-Identifier: GPL-3.0-or-later
# Copyright(C) 2019-2020 Max-Planck-Society
# Author: Philipp Arras
import argparse
-from os.path import join, split, splitext
from os import makedirs
+from os.path import join, split, splitext
+import numpy as np
import resolve as rve
if __name__ == "__main__":
@@ -18,12 +20,50 @@ if __name__ == "__main__":
parser.add_argument("--ch-begin", type=int)
parser.add_argument("--ch-end", type=int)
parser.add_argument("--ch-jump", type=int)
+ parser.add_argument("--freq-begin")
+ parser.add_argument("--freq-end")
+ parser.add_argument("--ignore-flags", action="store_true")
+ parser.add_argument("--autocorrelations-only", action="store_true",
+ help=("If this flag is set, all autocorrelations are "
+ "imported irrespective of whether they are "
+ "flagged or not."))
parser.add_argument("ms")
parser.add_argument(
"polarization_mode",
help="Can be 'stokesiavg', 'stokesi', 'all', or something like 'LL' or 'XY'.",
)
args = parser.parse_args()
+ if args.freq_begin is not None and args.freq_end is not None:
+ assert args.ch_begin is None
+ assert args.ch_jump is None
+ assert args.ch_end is None
+ assert args.spectral_window == 0
+
+ f0 = rve.str2val(args.freq_begin)
+ f1 = rve.str2val(args.freq_end)
+
+ # Determine spectral window
+ with rve.ms_table(join(args.ms, "SPECTRAL_WINDOW")) as t:
+ for ii in range(t.nrows()):
+ c = t.getcol("CHAN_FREQ", startrow=ii, nrow=1)[0]
+ fmin, fmax = min(c), max(c)
+ if fmin <= f0 <= fmax and fmin <= f1 <= fmax:
+ break
+ print(f"Load spectral window {ii}")
+
+ # Determine channels
+ if np.all(np.diff(c) > 0):
+ begin = np.searchsorted(c, f0)
+ end = np.searchsorted(c, f1)
+ elif np.all(np.diff(c) < 0):
+ begin = c.size - np.searchsorted(c[::-1], f0)
+ end = c.size - np.searchsorted(c[::-1], f1)
+ else:
+ raise RuntimeError("Channels are not sorted")
+ channels = slice(begin, end)
+ else:
+ channels = slice(args.ch_begin, args.ch_end, args.ch_jump)
+
makedirs(args.output_folder, exist_ok=True)
name = splitext(split(args.ms)[1])[0]
nspec = rve.ms_n_spectral_windows(args.ms)
@@ -34,9 +74,19 @@ if __name__ == "__main__":
args.include_calibration_info,
args.spectral_window,
args.polarization_mode,
- slice(args.ch_begin, args.ch_end, args.ch_jump),
+ channels,
+ args.ignore_flags
)
for ifield, oo in enumerate(obs):
- fname = join(args.output_folder, f"{name}field{ifield}.npz")
+ if oo is None:
+ continue
+ if args.autocorrelations_only:
+ oo = oo.restrict_to_autocorrelations()
+ auto = "autocorrelationsonly"
+ else:
+ auto = ""
+ fname = join(args.output_folder, f"{name}field{ifield}{args.data_column}{auto}.npz")
print(f"Save {fname}")
oo.save(fname, args.compress)
+ if oo.vis.size == 0:
+ print(f"WARNING: {fname} is empty")
diff --git a/misc/ms_summary.py b/misc/ms_summary.py
old mode 100644
new mode 100755
index c36bed83ae6292a6772d8951b96d4c04922fb95f..7a1734dd50b34604d7f87ca14454bed24f90c92f
--- a/misc/ms_summary.py
+++ b/misc/ms_summary.py
@@ -1,8 +1,10 @@
+#!/usr/bin/python3
# SPDX-License-Identifier: GPL-3.0-or-later
# Copyright(C) 2021 Max-Planck-Society
# Author: Philipp Arras
import argparse
+from os.path import join
import resolve as rve
@@ -13,6 +15,29 @@ if __name__ == "__main__":
nspec = rve.ms_n_spectral_windows(args.ms)
from casacore.tables import tablesummary
tablesummary(args.ms)
- print()
- print(f"The data set has {nspec} spectral windows.")
- print()
+
+ with rve.ms_table(join(args.ms, "SPECTRAL_WINDOW")) as t:
+ for ii in range(nspec):
+ print(f"Spectral window #{ii}")
+ chans = t.getcol("CHAN_FREQ", startrow=ii, nrow=1)
+
+ print(f"Shape: {chans.shape}")
+ print(f"f1-f0: {(chans[0][1]-chans[0][0])/1e6} MHz")
+ print("Frequencies (GHz)")
+ print(chans/1e9)
+ print()
+
+ with rve.ms_table(join(args.ms, "FIELD")) as t:
+ name = t.getcol("NAME")
+ refdir = t.getcol("REFERENCE_DIR")
+ deldir = t.getcol("DELAY_DIR")
+ phdir = t.getcol("PHASE_DIR")
+
+ print("NAME REFERNCE_DIR DELAY_DIR PHASE_DIR")
+ for nn, rd, dd, pd in zip(name, refdir, deldir, phdir):
+ print(nn, rd, dd, pd)
+
+ with rve.ms_table(join(args.ms, "POINTING")) as t:
+ print(t.getcol("ANTENNA_ID"))
+ print(t.getcol("TIME"))
+ print(t.getcol("DIRECTION").shape)
diff --git a/resolve/__init__.py b/resolve/__init__.py
index 4a1ade5b438f0cae57c7cee0a4e78a7ded42e6c3..68cdefef795f56efc4bf8836d86271b2eab1df6b 100644
--- a/resolve/__init__.py
+++ b/resolve/__init__.py
@@ -23,6 +23,7 @@ from .plotter import MfPlotter, Plotter
from .points import PointInserter
from .data.polarization import Polarization
from .polarization_matrix_exponential import *
+from .polarization_space import *
from .response import MfResponse, ResponseDistributor, StokesIResponse, SingleResponse
from .simple_operators import *
from .util import *
diff --git a/resolve/constants.py b/resolve/constants.py
index ffd0999edd9f0396b6df23374b125ae56d8766f5..0a537533ff8e89947947f3bac2d32bf109acabaa 100644
--- a/resolve/constants.py
+++ b/resolve/constants.py
@@ -38,3 +38,30 @@ def str2rad(s):
if unit == kk:
return float(s[:nn]) * c[kk]
raise RuntimeError("Unit not understood")
+
+
+def str2val(s):
+ """Convert string of number and unit to value.
+
+ Support the following keys: p n mu m (nothing) k M G T
+
+ Parameters
+ ----------
+ s : str
+ TODO
+
+ """
+ c = {
+ "p": 1e-12,
+ "n": 1e-9,
+ "mu": 1e-6,
+ "m": 1e-3,
+ "k": 1e3,
+ "M": 1e6,
+ "G": 1e9,
+ "T": 1e12
+ }
+ keys = set(c.keys())
+ if s[-1] in keys:
+ return float(s[:-1]) * c[s[-1]]
+ return float(s)
diff --git a/resolve/data/auxiliary_table.py b/resolve/data/auxiliary_table.py
new file mode 100644
index 0000000000000000000000000000000000000000..a2afcb71b2b7535b8baae84827ca43816eaeb0ed
--- /dev/null
+++ b/resolve/data/auxiliary_table.py
@@ -0,0 +1,55 @@
+# SPDX-License-Identifier: GPL-3.0-or-later
+# Copyright(C) 2021 Max-Planck-Society
+# Author: Philipp Arras
+
+import numpy as np
+
+from ..util import compare_attributes, my_assert, my_assert_isinstance, my_asserteq
+
+
+class AuxiliaryTable:
+ def __init__(self, inp_dict):
+ my_assert_isinstance(inp_dict, dict)
+ nrow = None
+ for kk, lst in inp_dict.items():
+ my_assert_isinstance(kk, str)
+ if not isinstance(lst, (list, np.ndarray)):
+ raise RuntimeError(f"{kk} neither list nor np.ndarray")
+ if nrow is None:
+ nrow = len(lst)
+ my_asserteq(nrow, len(lst))
+ my_asserteq(type(elem) for elem in lst)
+ self._dct = inp_dict
+ self._eq_attributes = "_dct",
+
+ def __eq__(self, other):
+ if not isinstance(other, type(self)):
+ return False
+ return compare_attributes(self, other, self._eq_attributes)
+
+ def __str__(self):
+ s = ["AuxiliaryTable:"]
+ for kk, lst in self._dct.items():
+ if isinstance(lst, np.ndarray):
+ s.append(f" {kk:<20} {str(lst.shape):>10} {str(lst.dtype):>15}")
+ else:
+ s.append(f" {kk:<20} {len(lst):>10} {str(type(lst[0])):>15}")
+ return "\n".join(s)
+
+ def __getitem__(self, key):
+ return self._dct[key]
+
+ def __contains__(self, key):
+ return key in self._dct
+
+ def row(self, i):
+ return AuxiliaryTable({kk: vv[i:i+1] for kk, vv in self._dct.items()})
+
+ def to_list(self):
+ return [list(self._dct.keys())] + [ss for ss in self._dct.values()]
+
+ @staticmethod
+ def from_list(lst):
+ keys = lst[0]
+ dct = {kk: lst[ii+1] for ii, kk in enumerate(keys)}
+ return AuxiliaryTable(dct)
diff --git a/resolve/data/ms_import.py b/resolve/data/ms_import.py
index 34105deb1dc3bf8022cca0ef8b13a6537a6ee45a..f16ce7339c3b48363dbf64f11aabda8591781550 100644
--- a/resolve/data/ms_import.py
+++ b/resolve/data/ms_import.py
@@ -1,15 +1,16 @@
# SPDX-License-Identifier: GPL-3.0-or-later
# Copyright(C) 2019-2021 Max-Planck-Society
-from os.path import isdir, join, splitext
+import os
+from os.path import isdir, join
import numpy as np
+from ..util import my_assert, my_assert_isinstance, my_asserteq
from .antenna_positions import AntennaPositions
-from .direction import Direction
+from .auxiliary_table import AuxiliaryTable
from .observation import Observation
from .polarization import Polarization
-from ..util import my_assert, my_asserteq, my_assert_isinstance
def ms_table(path):
@@ -17,14 +18,15 @@ def ms_table(path):
return table(path, readonly=True, ack=False)
-def ms2observations(
- ms,
- data_column,
- with_calib_info,
- spectral_window,
- polarizations="all",
- channel_slice=slice(None),
-):
+def _pol_id(ms_path, spectral_window):
+ """Return id for indexing polarization table for a given spectral window."""
+ with ms_table(join(ms_path, "DATA_DESCRIPTION")) as t:
+ polid = t.getcol("POLARIZATION_ID")[spectral_window]
+ return polid
+
+
+def ms2observations(ms, data_column, with_calib_info, spectral_window,
+ polarizations="all", channels=slice(None), ignore_flags=False):
"""Read and convert a given measurement set into an array of :class:`Observation`
If WEIGHT_SPECTRUM is available this column is used for weighting.
@@ -41,15 +43,18 @@ def ms2observations(
Reads in all information necessary for calibration, if True. If only
imaging shall be performed, this can be set to False.
spectral_window : int
- Index of spectral window which shall be imported.
+ Index of spectral window that shall be imported.
polarizations
"all": All polarizations are imported.
"stokesi": Only LL/RR or XX/YY polarizations are imported.
"stokesiavg": Only LL/RR or XX/YY polarizations are imported and averaged on the fly.
List of strings: Strings can be "XX", "YY", "XY", "YX", "LL", "LR", "RL", "RR". The respective polarizations are loaded.
- channel_slice : slice
- Slice of selected channels. Default: select all channels
- FIXME Select channels by indices
+ channels : slice or list
+ Select channels. Can be either a slice object or an index list. Default:
+ select all channels
+ ignore_flags : bool
+ If True, the whole measurement set is imported irrespective of the
+ flags. Default is false.
Returns
-------
@@ -58,85 +63,73 @@ def ms2observations(
Note
----
- We cannot import multiple spectral windows into one Observation instance
+ Multiple spectral windows are not imported into one Observation instance
because it is not guaranteed by the measurement set data structure that all
baselines are present in all spectral windows.
"""
+ # Input checks
if ms[-1] == "/":
ms = ms[:-1]
- if not isdir(ms) or splitext(ms)[1].lower() != ".ms":
+ if not isdir(ms):
raise RuntimeError
- if isinstance(polarizations, str):
- polarizations = [polarizations]
- if (
- "stokesiavg" in polarizations
- or "stokesi" in polarizations
- or "all" in polarizations
- ) and len(polarizations) > 1:
- raise ValueError
-
- # Spectral windows
- my_assert_isinstance(channel_slice, slice)
+ if ms == ".":
+ ms = os.getcwd()
+ if isinstance(polarizations, list):
+ for ll in polarizations:
+ my_assert(ll in ["XX", "YY", "XY", "YX", "LL", "LR", "RL", "RR"])
+ my_asserteq(len(set(polarizations)), len(polarizations))
+ else:
+ my_assert(polarizations in ["stokesi", "stokesiavg", "all"])
+ my_assert_isinstance(channels, (slice, list))
my_assert_isinstance(spectral_window, int)
my_assert(spectral_window >= 0)
my_assert(spectral_window < ms_n_spectral_windows(ms))
- with ms_table(join(ms, "SPECTRAL_WINDOW")) as t:
- freq = t.getcol("CHAN_FREQ")[spectral_window]
+ # /Input checks
# Polarization
with ms_table(join(ms, "POLARIZATION")) as t:
- pol = t.getcol("CORR_TYPE")
- my_asserteq(pol.ndim, 2)
- my_asserteq(pol.shape[0], 1)
- pol = list(pol[0]) # Not clear what the first dimension is used for
+ pol = t.getcol("CORR_TYPE", startrow=_pol_id(ms, spectral_window), nrow=1)[0]
polobj = Polarization(pol)
- if polarizations[0] == "stokesi":
- polarizations = ["LL", "RR"] if polobj.circular() else ["XX", "YY"]
- if polarizations[0] == "stokesiavg":
+ if polarizations == "stokesiavg":
pol_ind = polobj.stokes_i_indices()
polobj = Polarization.trivial()
pol_summation = True
- elif polarizations[0] == "all":
+ elif polarizations == "all":
pol_ind = None
pol_summation = False
else:
+ if polarizations == "stokesi":
+ polarizations = ["LL", "RR"] if polobj.circular() else ["XX", "YY"]
polobj = polobj.restrict_by_name(polarizations)
pol_ind = [polobj.to_str_list().index(ii) for ii in polarizations]
pol_summation = False
- # Field
- with ms_table(join(ms, "FIELD")) as t:
- equinox = t.coldesc("REFERENCE_DIR")["desc"]["keywords"]["MEASINFO"]["Ref"]
- equinox = str(equinox)[1:]
- # FIXME Put proper support for equinox here
- if equinox == "1950_VLA":
- equinox = 1950
- dirs = []
- for pc in t.getcol("REFERENCE_DIR"):
- my_asserteq(pc.shape, (1, 2))
- dirs.append(Direction(pc[0], equinox))
- dirs = tuple(dirs)
-
- # FIXME Determine which observation is calibration observation
- # FIXME Import name of source
observations = []
- for ifield, direction in enumerate(dirs):
- uvw, ant1, ant2, time, freq_out, vis, wgt = read_ms_i(
- ms,
- data_column,
- freq,
- ifield,
- spectral_window,
- pol_ind,
- pol_summation,
- with_calib_info,
- channel_slice,
- )
- vis[wgt == 0] = 0.0
- vis = _ms2resolve_transpose(vis)
- wgt = _ms2resolve_transpose(wgt)
- antpos = AntennaPositions(uvw, ant1, ant2, time)
- observations.append(Observation(antpos, vis, wgt, polobj, freq_out, direction))
+ for ifield in range(ms_n_fields(ms)):
+ auxtables = {}
+ auxtables["ANTENNA"] = _import_aux_table(ms, "ANTENNA")
+ auxtables["STATE"] = _import_aux_table(ms, "STATE")
+ auxtables["SPECTRAL_WINDOW"] = _import_aux_table(ms, "SPECTRAL_WINDOW", row=spectral_window, skip=["ASSOC_NATURE"])
+ sf = _source_and_field_table(ms, spectral_window, ifield)
+ if sf is None:
+ print(f"Field {ifield} cannot be found in SOURCE table")
+ observations.append(None)
+ continue
+ auxtables = {**auxtables, **sf}
+ print(f"Work on Field {ifield}: {auxtables['SOURCE']['NAME'][0]}")
+
+ mm = read_ms_i(ms, data_column, ifield, spectral_window, pol_ind, pol_summation,
+ with_calib_info, channels, ignore_flags)
+ if mm is None:
+ print(f"{ms}, field #{ifield} is empty or fully flagged")
+ observations.append(None)
+ continue
+ if mm["ptg"] is not None:
+ raise NotImplementedError
+ antpos = AntennaPositions(mm["uvw"], mm["ant1"], mm["ant2"], mm["time"])
+ obs = Observation(antpos, mm["vis"], mm["wgt"], polobj, mm["freq"],
+ auxiliary_tables=auxtables)
+ observations.append(obs)
return observations
@@ -145,103 +138,52 @@ def _ms2resolve_transpose(arr):
return np.ascontiguousarray(np.transpose(arr, (2, 0, 1)))
-def _determine_weighting(t):
- fullwgt = False
- weightcol = "WEIGHT"
- try:
- t.getcol("WEIGHT_SPECTRUM", startrow=0, nrow=1)
- weightcol = "WEIGHT_SPECTRUM"
- fullwgt = True
- except RuntimeError:
- pass
+def _determine_weighting(ms):
+ with ms_table(ms) as t:
+ if "WEIGHT_SPECTRUM" in t.colnames():
+ weightcol = "WEIGHT_SPECTRUM"
+ fullwgt = True
+ else:
+ weightcol = "WEIGHT"
+ fullwgt = False
return fullwgt, weightcol
-def read_ms_i(
- name,
- data_column,
- freq,
- field,
- spectral_window,
- pol_indices,
- pol_summation,
- with_calib_info,
- channel_slice,
-):
- freq = np.array(freq)
- my_asserteq(freq.ndim, 1)
- my_assert(len(freq) > 0)
- nchan = len(freq)
+def read_ms_i(name, data_column, field, spectral_window, pol_indices, pol_summation,
+ with_calib_info, channels, ignore_flags):
assert pol_indices is None or isinstance(pol_indices, list)
if pol_indices is None:
pol_indices = slice(None)
if pol_summation:
my_asserteq(len(pol_indices), 2)
+ # Check if data column is available and get shape
with ms_table(name) as t:
- # FIXME Get rid of fullwgt
- fullwgt, weightcol = _determine_weighting(t)
+ nmspol = t.getcol(data_column, startrow=0, nrow=3).shape[2]
nrow = t.nrows()
- nmspol = t.getcol("FLAG", startrow=0, nrow=1).shape[2]
- print("Measurement set visibilities:")
- print(f" shape: ({nrow}, {nchan}, {nmspol})")
- active_rows = np.ones(nrow, dtype=np.bool)
- active_channels = np.zeros(nchan, dtype=np.bool)
- step = max(1, nrow // 100) # how many rows to read in every step
-
- # Check if data column is available
- t.getcol(data_column, startrow=0, nrow=10)
+ print("Measurement set visibilities:")
+ print(f" shape: ({nrow}, {_ms_nchannels(name, spectral_window)}, {nmspol})")
- # Determine which subset of rows/channels we need to input
- start = 0
- while start < nrow:
- print("First pass:", f"{(start/nrow*100):.1f}%", end="\r")
- stop = min(nrow, start + step)
- tflags = t.getcol("FLAG", startrow=start, nrow=stop - start)[
- ..., pol_indices
- ]
- twgt = t.getcol(weightcol, startrow=start, nrow=stop - start)[
- ..., pol_indices
- ]
- if channel_slice != slice(None):
- tchslcflags = np.ones_like(tflags)
- tchslcflags[:, channel_slice] = False
- tflags = np.logical_or(tflags, tchslcflags)
- if not fullwgt:
- twgt = np.repeat(twgt[:, None], nchan, axis=1)
- my_asserteq(twgt.ndim, 3)
- npol = tflags.shape[2]
- if pol_summation:
- tflags = np.any(tflags.astype(np.bool), axis=2)[..., None]
- twgt = np.sum(twgt, axis=2)[..., None]
- tflags[twgt == 0] = True
-
- # Select field and spectral window
- tfieldid = t.getcol("FIELD_ID", startrow=start, nrow=stop - start)
- tflags[tfieldid != field] = True
- tspw = t.getcol("DATA_DESC_ID", startrow=start, nrow=stop - start)
- tflags[tspw != spectral_window] = True
+ active_rows, active_channels = _first_pass(name, field, spectral_window, channels, pol_indices,
+ pol_summation, ignore_flags)
+ nrealrows, nrealchan = np.sum(active_rows), np.sum(active_channels)
+ if nrealrows == 0 or nrealchan == 0:
+ return None
- # Inactive if all polarizations are flagged
- assert tflags.ndim == 3
- tflags = np.all(tflags, axis=2)
- active_rows[start:stop] = np.invert(np.all(tflags, axis=1))
- active_channels = np.logical_or(
- active_channels, np.invert(np.all(tflags, axis=0))
- )
- start = stop
- nrealrows, nrealchan = np.sum(active_rows), np.sum(active_channels)
- if nrealrows == 0 or nrealchan == 0:
- raise RuntimeError("Empty data set")
+ # Freq
+ freq = _ms_channels(name, spectral_window)[active_channels]
- # Create output arrays
+ # Vis, wgt, (flags)
+ fullwgt, weightcol = _determine_weighting(name)
+ nchan = _ms_nchannels(name, spectral_window)
+ with ms_table(name) as t:
if pol_summation:
npol = 1
else:
- if pol_indices != slice(None):
- npol = len(pol_indices)
+ if pol_indices == slice(None):
+ npol = nmspol
else:
- npol = npol
+ npol = len(pol_indices)
shp = (nrealrows, nrealchan, npol)
vis = np.empty(shp, dtype=np.complex64)
wgt = np.empty(shp, dtype=np.float32)
@@ -250,14 +192,13 @@ def read_ms_i(
start, realstart = 0, 0
while start < nrow:
print("Second pass:", f"{(start/nrow*100):.1f}%", end="\r")
- stop = min(nrow, start + step)
+ stop = _ms_stop(start, nrow)
realstop = realstart + np.sum(active_rows[start:stop])
if realstop > realstart:
allrows = stop - start == realstop - realstart
- twgt = t.getcol(weightcol, startrow=start, nrow=stop - start)[
- ..., pol_indices
- ]
+ # Weights
+ twgt = t.getcol(weightcol, startrow=start, nrow=stop - start)[..., pol_indices]
assert twgt.dtype == np.float32
if not fullwgt:
twgt = np.repeat(twgt[:, None], nchan, axis=1)
@@ -265,24 +206,24 @@ def read_ms_i(
twgt = twgt[active_rows[start:stop]]
twgt = twgt[:, active_channels]
- tvis = t.getcol(data_column, startrow=start, nrow=stop - start)[
- ..., pol_indices
- ]
+ # Vis
+ tvis = t.getcol(data_column, startrow=start, nrow=stop - start)[..., pol_indices]
assert tvis.dtype == np.complex64
if not allrows:
tvis = tvis[active_rows[start:stop]]
tvis = tvis[:, active_channels]
- tflags = t.getcol("FLAG", startrow=start, nrow=stop - start)[
- ..., pol_indices
- ]
+ # Flags
+ tflags = _conditional_flags(t, start, stop, pol_indices, ignore_flags)
if not allrows:
tflags = tflags[active_rows[start:stop]]
tflags = tflags[:, active_channels]
+ # Polarization summation
assert twgt.ndim == tflags.ndim == 3
assert tflags.dtype == np.bool
- twgt = twgt * (~tflags)
+ if not ignore_flags:
+ twgt = twgt * (~tflags)
if pol_summation:
assert twgt.shape[2] == 2
# Noise-weighted average
@@ -294,31 +235,162 @@ def read_ms_i(
wgt[realstart:realstop] = twgt
start, realstart = stop, realstop
- uvw = t.getcol("UVW")[active_rows]
- if with_calib_info:
+ print("Selected:", 10 * " ")
+ print(f" shape: {vis.shape}")
+ print(f" flagged: {(1.0-np.sum(wgt!=0)/wgt.size)*100:.1f} %")
+
+ # UVW
+ with ms_table(name) as t:
+ uvw = np.ascontiguousarray(t.getcol("UVW")[active_rows])
+
+ # Calibration info
+ if with_calib_info:
+ with ms_table(name) as t:
ant1 = np.ascontiguousarray(t.getcol("ANTENNA1")[active_rows])
ant2 = np.ascontiguousarray(t.getcol("ANTENNA2")[active_rows])
time = np.ascontiguousarray(t.getcol("TIME")[active_rows])
+ else:
+ ant1 = ant2 = time = None
+
+ # Pointing
+ with ms_table(join(name, "POINTING")) as t:
+ if t.nrows() == 0:
+ ptg = None
else:
- ant1 = ant2 = time = None
- print("Selected:", 10 * " ")
- print(f" shape: {vis.shape}")
- print(f" flagged: {(1.0-np.sum(wgt!=0)/wgt.size)*100:.1f} %")
- freq = freq[active_channels]
+ ptg = np.empty((nrealrows, 1, 2), dtype=np.float64)
+ start, realstart = 0, 0
+ while start < nrow:
+ print("Second pass:", f"{(start/nrow*100):.1f}%", end="\r")
+ stop = _ms_stop(start, nrow)
+ realstop = realstart + np.sum(active_rows[start:stop])
+ if realstop > realstart:
+ allrows = stop - start == realstop - realstart
+ tptg = t.getcol("DIRECTION", startrow=start, nrow=stop - start)
+ tptg = tptg[active_rows[start:stop]]
+ ptg[realstart:realstop] = tptg
+ start, realstart = stop, realstop
my_asserteq(wgt.shape, vis.shape)
- return (
- np.ascontiguousarray(uvw),
- ant1,
- ant2,
- time,
- np.ascontiguousarray(freq),
- np.ascontiguousarray(vis),
- np.ascontiguousarray(wgt),
- )
+ vis = np.ascontiguousarray(_ms2resolve_transpose(vis))
+ wgt = np.ascontiguousarray(_ms2resolve_transpose(wgt))
+ vis[wgt == 0] = 0.0
+ return {"uvw": uvw, "ant1": ant1, "ant2": ant2, "time": time, "freq": freq, "vis": vis,
+ "wgt": wgt, "ptg": ptg}
+
+
+def _first_pass(ms, field, spectral_window, channels, pol_indices, pol_summation, ignore_flags):
+ """Go through measurement set and determine which rows and which channels are active for a given
+ field and a given spectral window.
+ """
+ fullwgt, weightcol = _determine_weighting(ms)
+ nchan = _ms_nchannels(ms, spectral_window)
+ with ms_table(ms) as t:
+ nrow = t.nrows()
+ active_rows = np.ones(nrow, dtype=np.bool)
+ active_channels = np.zeros(nchan, dtype=np.bool)
+
+ # Determine which subset of rows/channels we need to input
+ start = 0
+ while start < nrow:
+ print("First pass:", f"{(start/nrow*100):.1f}%", end="\r")
+ stop = _ms_stop(start, nrow)
+ tflags = _conditional_flags(t, start, stop, pol_indices, ignore_flags)
+ twgt = t.getcol(weightcol, startrow=start, nrow=stop - start)[..., pol_indices]
+
+ tchslcflags = np.ones_like(tflags)
+ tchslcflags[:, channels] = False
+ tflags = np.logical_or(tflags, tchslcflags)
+
+ if not fullwgt:
+ twgt = np.repeat(twgt[:, None], nchan, axis=1)
+ my_asserteq(twgt.ndim, 3)
+ if pol_summation:
+ tflags = np.any(tflags.astype(np.bool), axis=2)[..., None]
+ twgt = np.sum(twgt, axis=2)[..., None]
+ tflags[twgt == 0] = True
+
+ # Select field and spectral window
+ tfieldid = t.getcol("FIELD_ID", startrow=start, nrow=stop - start)
+ tflags[tfieldid != field] = True
+ tspw = t.getcol("DATA_DESC_ID", startrow=start, nrow=stop - start)
+ tflags[tspw != spectral_window] = True
+
+ # Inactive if all polarizations are flagged
+ assert tflags.ndim == 3
+ tflags = np.all(tflags, axis=2)
+ active_rows[start:stop] = np.invert(np.all(tflags, axis=1))
+ active_channels = np.logical_or(active_channels, np.invert(np.all(tflags, axis=0)))
+ start = stop
+ return active_rows, active_channels
+
+
+def _ms_stop(start, nrow):
+ """Compute sensible step size for going through the rows of a measurement set and return stop
+ index.
+ """
+ step = max(1, nrow // 100)
+ return min(nrow, start + step)
+
+
+def _ms_nchannels(ms, spectral_window):
+ """Return number of channels in a given spectral window of a measurement set.
+ """
+ return len(_ms_channels(ms, spectral_window))
+
+
+def _ms_channels(ms, spectral_window):
+ """Return frequencies of channels in a given spectral window of a measurement set.
+ """
+ with ms_table(join(ms, "SPECTRAL_WINDOW")) as t:
+ freq = t.getcol("CHAN_FREQ", startrow=spectral_window, nrow=1)[0]
+ my_asserteq(freq.ndim, 1)
+ my_assert(len(freq) > 0)
+ return freq
def ms_n_spectral_windows(ms):
with ms_table(join(ms, "SPECTRAL_WINDOW")) as t:
- freq = t.getcol("CHAN_FREQ")
- return freq.shape[0]
+ n_spectral_windows = t.nrows()
+ return n_spectral_windows
+
+
+def ms_n_fields(ms):
+ with ms_table(join(ms, "FIELD")) as t:
+ n = t.nrows()
+ return n
+
+
+def _conditional_flags(table, start, stop, pol_indices, ignore):
+ tflags = table.getcol("FLAG", startrow=start, nrow=stop - start)[..., pol_indices]
+ if ignore:
+ tflags = np.zeros_like(tflags)
+ return tflags
+
+
+def _import_aux_table(ms, table_name, row=None, skip=[]):
+ with ms_table(join(ms, table_name)) as t:
+ keys = filter(lambda s: s not in skip, t.colnames())
+ if row is None:
+ dct = {kk: t.getcol(kk) for kk in keys}
+ else:
+ dct = {kk: t.getcol(kk, startrow=row, nrow=1) for kk in keys}
+ aux_table = AuxiliaryTable(dct)
+ return aux_table
+
+
+def _source_and_field_table(ms, spectral_window, ifield):
+ source_table = _import_aux_table(ms, "SOURCE", skip=["POSITION", "TRANSITION", "REST_FREQUENCY",
+ "SYSVEL", "SOURCE_MODEL", "PULSAR_ID"])
+ field_table = _import_aux_table(ms, "FIELD", row=ifield)
+ source_id = field_table["SOURCE_ID"][0]
+ ind = np.where(np.logical_and(source_table["SOURCE_ID"] == source_id,
+ np.logical_or(source_table["SPECTRAL_WINDOW_ID"] == spectral_window,
+ source_table["SPECTRAL_WINDOW_ID"] == -1)))[0]
+ if len(ind) == 0:
+ return None
+ elif len(ind) == 1:
+ ind = ind[0]
+ else:
+ raise RuntimeError
+
+ return {"SOURCE": source_table.row(ind), "FIELD": field_table}
diff --git a/resolve/data/observation.py b/resolve/data/observation.py
index 2dbe6c1543f5aa70e5bfcda5277efd09d7f6fc51..21d3bc4821cf6fb1d3162efcb506de657fa62f49 100644
--- a/resolve/data/observation.py
+++ b/resolve/data/observation.py
@@ -2,22 +2,24 @@
# Copyright(C) 2019-2021 Max-Planck-Society
# Author: Philipp Arras
-import numpy as np
-
import nifty8 as ift
+import numpy as np
-from .antenna_positions import AntennaPositions
from ..constants import SPEEDOFLIGHT
-from .direction import Direction, Directions
from ..mpi import onlymaster
+from ..util import (compare_attributes, my_assert, my_assert_isinstance,
+ my_asserteq)
+from .antenna_positions import AntennaPositions
+from .auxiliary_table import AuxiliaryTable
+from .direction import Direction, Directions
from .polarization import Polarization
-from ..util import compare_attributes, my_assert, my_assert_isinstance, my_asserteq
class BaseObservation:
@property
def _dom(self):
- dom = [ift.UnstructuredDomain(ss) for ss in self._vis.shape]
+ pol_dom = self.polarization.space
+ dom = [pol_dom] + [ift.UnstructuredDomain(ss) for ss in self._vis.shape[1:]]
return ift.makeDomain(dom)
@property
@@ -157,7 +159,7 @@ class BaseObservation:
return NotImplementedError
def __eq__(self, other):
- if not isinstance(other, self.__class__):
+ if not isinstance(other, type(self)):
return False
if (
self._vis.dtype != other._vis.dtype
@@ -266,17 +268,16 @@ class Observation(BaseObservation):
Polarization information of the data set.
freq : numpy.ndarray
Contains the measured frequencies. Shape (n_channels)
- direction : Direction
- Direction information of the data set.
+ auxiliary_tables : dict
+ Dictionary of auxiliary tables. Default: None.
Note
----
vis and weight must have the same shape.
"""
- def __init__(self, antenna_positions, vis, weight, polarization, freq, direction):
+ def __init__(self, antenna_positions, vis, weight, polarization, freq, auxiliary_tables=None):
nrows = len(antenna_positions)
- my_assert_isinstance(direction, Direction)
my_assert_isinstance(polarization, Polarization)
my_assert_isinstance(antenna_positions, AntennaPositions)
my_asserteq(weight.shape, vis.shape)
@@ -294,6 +295,8 @@ class Observation(BaseObservation):
my_assert(np.all(np.isfinite(vis[weight > 0.])))
my_assert(np.all(np.isfinite(weight)))
+ my_assert(np.all(np.diff(freq)))
+
vis.flags.writeable = False
weight.flags.writeable = False
@@ -302,10 +305,16 @@ class Observation(BaseObservation):
self._weight = weight
self._polarization = polarization
self._freq = freq
- self._direction = direction
- self._eq_attributes = "_direction", "_polarization", "_freq", "_antpos", "_vis", "_weight"
+ if auxiliary_tables is not None:
+ my_assert_isinstance(auxiliary_tables, dict)
+ for kk, vv in auxiliary_tables.items():
+ my_assert_isinstance(vv, AuxiliaryTable)
+ my_assert_isinstance(kk, str)
+ self._auxiliary_tables = auxiliary_tables
+ self._eq_attributes = ("_polarization", "_freq", "_antpos", "_vis", "_weight",
+ "_auxiliary_tables")
@onlymaster
def save(self, file_name, compress):
@@ -314,12 +323,15 @@ class Observation(BaseObservation):
weight=self._weight,
freq=self._freq,
polarization=self._polarization.to_list(),
- direction=self._direction.to_list(),
)
for ii, vv in enumerate(self._antpos.to_list()):
if vv is None:
vv = np.array([])
dct[f"antpos{ii}"] = vv
+ if self._auxiliary_tables is not None:
+ for kk, auxtable in self._auxiliary_tables.items():
+ for ii, elem in enumerate(auxtable.to_list()):
+ dct[f"auxtable_{kk}_{ii:>04}"] = elem
f = np.savez_compressed if compress else np.savez
f(file_name, **dct)
@@ -332,8 +344,23 @@ class Observation(BaseObservation):
if val.size == 0:
val = None
antpos.append(val)
+
+ # Load auxtables
+ keys = set(kk[9:-5] for kk in dct.keys() if kk[:9] == "auxtable_")
+ if len(keys) == 0:
+ auxtables = None
+ else:
+ auxtables = {}
+ for kk in keys:
+ ii = 0
+ inp = []
+ while f"auxtable_{kk}_{ii:>04}" in dct.keys():
+ inp.append(dct[f"auxtable_{kk}_{ii:>04}"])
+ ii += 1
+ auxtables[kk] = AuxiliaryTable.from_list(inp)
+ # /Load auxtables
+
pol = Polarization.from_list(dct["polarization"])
- direction = Direction.from_list(dct["direction"])
slc = slice(None) if lo_hi_index is None else slice(*lo_hi_index)
# FIXME Put barrier here that makes sure that only one full Observation is loaded at a time
return Observation(
@@ -342,7 +369,7 @@ class Observation(BaseObservation):
dct["weight"][..., slc],
pol,
dct["freq"][slc],
- direction,
+ auxtables
)
def flags_to_nan(self):
@@ -351,7 +378,7 @@ class Observation(BaseObservation):
vis = self._vis.copy()
vis[self.flags.val] = np.nan
return Observation(self._antpos, vis, self._weight, self._polarization, self._freq,
- self._direction)
+ self._auxiliary_tables)
@staticmethod
def load_mf(file_name, n_imaging_bands, comm=None):
@@ -382,13 +409,14 @@ class Observation(BaseObservation):
return obs_list, nu0
def __getitem__(self, slc):
+ # FIXME Do I need to change something in self._auxiliary_tables?
return Observation(
self._antpos[slc],
self._vis[:, slc],
self._weight[:, slc],
self._polarization,
self._freq,
- self._direction,
+ self._auxiliary_tables,
)
def get_freqs(self, frequency_list):
@@ -404,16 +432,18 @@ class Observation(BaseObservation):
return self.get_freqs_by_slice(mask)
def get_freqs_by_slice(self, slc):
+ # FIXME Do I need to change something in self._auxiliary_tables?
return Observation(
self._antpos,
self._vis[..., slc],
self._weight[..., slc],
self._polarization,
self._freq[slc],
- self._direction,
+ self._auxiliary_tables,
)
def average_stokesi(self):
+ # FIXME Do I need to change something in self._auxiliary_tables?
my_asserteq(self._vis.shape[0], 2)
my_asserteq(self._polarization.restrict_to_stokes_i(), self._polarization)
vis = np.sum(self._weight * self._vis, axis=0)[None]
@@ -422,18 +452,42 @@ class Observation(BaseObservation):
vis /= wgt + np.ones_like(wgt) * invmask
vis[invmask] = 0.0
return Observation(
- self._antpos, vis, wgt, Polarization.trivial(), self._freq, self._direction
+ self._antpos, vis, wgt, Polarization.trivial(), self._freq, self._auxiliary_tables
)
+ def restrict_by_time(self, tmin, tmax, with_index=False):
+ my_assert(all(np.diff(self.time) >= 0))
+ start, stop = np.searchsorted(self.time, [tmin, tmax])
+ ind = slice(start, stop)
+ res = self[ind]
+ if with_index:
+ return res, ind
+ return res
+
+ def restrict_by_freq(self, fmin, fmax, with_index=False):
+ my_assert(all(np.diff(self.freq) > 0))
+ start, stop = np.searchsorted(self.freq, [fmin, fmax])
+ ind = slice(start, stop)
+ res = self.get_freqs_by_slice(ind)
+ if with_index:
+ return res, ind
+ return res
+
def restrict_to_stokesi(self):
+ # FIXME Do I need to change something in self._auxiliary_tables?
my_asserteq(self._vis.shape[0], 4)
ind = self._polarization.stokes_i_indices()
vis = self._vis[ind]
wgt = self._weight[ind]
pol = self._polarization.restrict_to_stokes_i()
- return Observation(self._antpos, vis, wgt, pol, self._freq, self._direction)
+ return Observation(self._antpos, vis, wgt, pol, self._freq, self._auxiliary_tables)
+
+ def restrict_to_autocorrelations(self):
+ slc = self._antpos.ant1 == self._antpos.ant2
+ return self[slc]
def move_time(self, t0):
+ # FIXME Do I need to change something in self._auxiliary_tables?
antpos = self._antpos.move_time(t0)
return Observation(
antpos,
@@ -441,9 +495,56 @@ class Observation(BaseObservation):
self._weight,
self._polarization,
self._freq,
- self._direction,
+ self._auxiliary_tables
)
+ def time_average(self, list_of_timebins):
+ # FIXME check that timebins do not overlap
+ # time, ant1, ant2
+ ts = self._antpos.time
+ row_to_bin_map = np.empty(ts.shape)
+ row_to_bin_map[:] = np.nan
+
+ for ii, (lo, hi) in enumerate(list_of_timebins):
+ ind = np.logical_and(ts >= lo, ts < hi)
+ assert np.all(np.isnan(row_to_bin_map[ind]))
+ row_to_bin_map[ind] = ii
+
+ assert np.all(~np.isnan(row_to_bin_map))
+ row_to_bin_map = row_to_bin_map.astype(int)
+
+ ant1 = self._antpos.ant1
+ ant2 = self._antpos.ant2
+ assert np.max(ant1) < 1000
+ assert np.max(ant2) < 1000
+ assert np.max(row_to_bin_map) < np.iinfo(np.dtype("int64")).max / 1000000
+ atset = set(zip(ant1, ant2, row_to_bin_map))
+ dct = {aa: ii for ii, aa in enumerate(atset)}
+ dct_inv = {yy: xx for xx, yy in dct.items()}
+ masterindex = np.array([dct[(a1, a2, tt)] for a1, a2, tt in zip(ant1, ant2, row_to_bin_map)])
+
+ vis, wgt = self.vis.val, self.weight.val
+ new_vis = np.empty((self.npol, len(atset), self.nfreq), dtype=self.vis.dtype)
+ new_wgt = np.empty((self.npol, len(atset), self.nfreq), dtype=self.weight.dtype)
+ for pol in range(self.npol):
+ for freq in range(self.nfreq):
+ enum = np.bincount(masterindex, weights=vis[pol, :, freq].real*wgt[pol, :, freq])
+ enum = enum + 1j*np.bincount(masterindex, weights=vis[pol, :, freq].imag*wgt[pol, :, freq])
+ denom = np.bincount(masterindex, weights=wgt[pol, :, freq])
+ new_vis[pol, :, freq] = enum/denom
+ new_wgt[pol, :, freq] = denom
+
+ new_times = np.array([dct_inv[ii][2] for ii in range(len(atset))])
+ new_times = np.mean(np.array(list_of_timebins), axis=1)[new_times]
+
+ new_ant1 = np.array([dct_inv[ii][0] for ii in range(len(atset))])
+ new_ant2 = np.array([dct_inv[ii][1] for ii in range(len(atset))])
+ new_uvw = ap.uvw
+ # FIXME Find correct uvw
+ ap = self._antpos
+ ap = AntennaPositions(new_uvw, new_ant1, new_ant2, new_times)
+ return Observation(ap, new_vis, new_wgt, self._polarization, self._freq, self._auxiliary_tables)
+
@property
def uvw(self):
return self._antpos.uvw
@@ -452,6 +553,30 @@ class Observation(BaseObservation):
def antenna_positions(self):
return self._antpos
+ @property
+ def ant1(self):
+ return self._antpos.ant1
+
+ @property
+ def ant2(self):
+ return self._antpos.ant2
+
+ @property
+ def time(self):
+ return self._antpos.time
+
+ @property
+ def direction(self):
+ if self._auxiliary_tables is not None and "FIELD" in self._auxiliary_tables:
+ equinox = 2000 # FIXME Figure out how to extract this from a measurement set
+ refdir = self._auxiliary_tables["FIELD"]["REFERENCE_DIR"][0]
+ my_asserteq(refdir.shape[0] == 0)
+ return Direction(refdir[0], equinox)
+ return None
+
+ def auxiliary_table(self, name):
+ return self._auxiliary_tables[name]
+
def effective_uvw(self):
out = np.einsum("ij,k->jik", self.uvw, self._freq / SPEEDOFLIGHT)
my_asserteq(out.shape, (3, self.nrow, self.nfreq))
diff --git a/resolve/data/polarization.py b/resolve/data/polarization.py
index 3bbbd199cef3a65a20fa97e74e49fdfaf9cba295..48b2c79c645c055242cca75795836cee1305b938 100644
--- a/resolve/data/polarization.py
+++ b/resolve/data/polarization.py
@@ -68,6 +68,15 @@ class Polarization:
def to_str_list(self):
return [TABLE[ii] for ii in self._ind]
+ @property
+ def space(self):
+ from ..polarization_space import PolarizationSpace
+ if self == Polarization.trivial():
+ x = "I"
+ else:
+ x = [TABLE[ii] for ii in self._ind]
+ return PolarizationSpace(x)
+
@staticmethod
def from_list(lst):
return Polarization(lst)
diff --git a/resolve/likelihood.py b/resolve/likelihood.py
index 7e75481a576326cbf3bfcfb1dd21dfd9348d389c..54c9ed462895fd652a9118a32abd05cbb5cd7967 100644
--- a/resolve/likelihood.py
+++ b/resolve/likelihood.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: GPL-3.0-or-later
-# Copyright(C) 2020 Max-Planck-Society
+# Copyright(C) 2020-2021 Max-Planck-Society
# Author: Philipp Arras
from functools import reduce
diff --git a/resolve/multi_frequency/irg_space.py b/resolve/multi_frequency/irg_space.py
index 70a07f7126b0e22df548d08dac98ee6c96bd7da0..a2a23bc76611820836e7c0562ebb1d8fe2b9d942 100644
--- a/resolve/multi_frequency/irg_space.py
+++ b/resolve/multi_frequency/irg_space.py
@@ -12,8 +12,6 @@
# along with this program. If not, see <http://www.gnu.org/licenses/>.
#
# Copyright(C) 2013-2020 Max-Planck-Society
-#
-# NIFTy is being developed at the Max-Planck-Institut fuer Astrophysik.
import numpy as np
diff --git a/resolve/multi_frequency/operators.py b/resolve/multi_frequency/operators.py
index 583e67077fc8293311c819367662ae2516d3a960..b77f98b2ecd8e21dacd600472e5297b722d63c76 100644
--- a/resolve/multi_frequency/operators.py
+++ b/resolve/multi_frequency/operators.py
@@ -1,29 +1,36 @@
# SPDX-License-Identifier: GPL-3.0-or-later
-# Copyright(C) 2013-2020 Max-Planck-Society
+# Copyright(C) 2013-2021 Max-Planck-Society
# Authors: Philipp Frank, Philipp Arras, Philipp Haim
-import numpy as np
-
import nifty8 as ift
+import numpy as np
-from ..util import my_asserteq
+from ..polarization_space import PolarizationSpace
+from ..util import my_assert, my_assert_isinstance, my_asserteq
+from .irg_space import IRGSpace
class WienerIntegrations(ift.LinearOperator):
- def __init__(self, freqdomain, imagedomain):
- # FIXME Write interface checks
- self._target = ift.makeDomain((freqdomain, imagedomain))
- dom = list(self._target)
- dom = ift.UnstructuredDomain((2, freqdomain.size - 1)), imagedomain
+ """Operator that performs the integrations necessary for an integrated
+ Wiener process.
+
+ Parameters
+ ----------
+ time_domain : IRGSpace
+ Domain that contains the temporal information of the process.
+
+ remaining_domain : DomainTuple or Domain
+ All integrations are handled independently for this domain.
+ """
+ def __init__(self, time_domain, remaining_domain):
+ my_assert_isinstance(time_domain, IRGSpace)
+ self._target = ift.makeDomain((time_domain, remaining_domain))
+ dom = ift.UnstructuredDomain((2, time_domain.size - 1)), remaining_domain
self._domain = ift.makeDomain(dom)
- self._volumes = freqdomain.dvol[:, None, None]
+ self._volumes = time_domain.dvol[:, None, None]
self._capability = self.TIMES | self.ADJOINT_TIMES
def apply(self, x, mode):
- # FIXME If it turns out that this operator is a performance
- # bottleneck we can try implement it in parallel in C++. But
- # it may be hard to achieve good scaling because I think it
- # becomes memory bound quickly.
self._check_input(x, mode)
first, second = (0,), (1,)
from_second = (slice(1, None),)
@@ -50,7 +57,7 @@ class WienerIntegrations(ift.LinearOperator):
def IntWProcessInitialConditions(a0, b0, wfop):
for op in [a0, b0, wfop]:
- ift.is_operator(op)
+ my_assert(ift.is_operator(op))
my_asserteq(a0.target, b0.target, ift.makeDomain(wfop.target[1]))
bc = _FancyBroadcast(wfop.target)
factors = ift.full(wfop.target[0], 0)
@@ -84,13 +91,13 @@ class MfWeightingInterpolation(ift.LinearOperator):
my_asserteq(domain.shape[0], eff_uvw.shape[2]) # freqaxis
self._domain = domain
nrow, nfreq = eff_uvw.shape[1:]
- tgt = [ift.UnstructuredDomain(aa) for aa in [1, nrow, nfreq]]
+ tgt = [PolarizationSpace("I")] + [ift.UnstructuredDomain(aa) for aa in [nrow, nfreq]]
self._target = ift.DomainTuple.make(tgt)
self._capability = self.TIMES | self.ADJOINT_TIMES
# FIXME Try to unify all those operators which loop over freq dimension
self._ops = []
for ii in range(nfreq):
- op = ift.LinearInterpolator(domain[1], eff_uvw[:, :, ii])
+ op = ift.LinearInterpolator(domain[1], eff_uvw.val[:, :, ii])
self._ops.append(op)
my_asserteq(self.target.shape[0], 1)
diff --git a/resolve/plotter.py b/resolve/plotter.py
index 598a70817bb817a88162db045872aaa734eabe84..9d4de3443ba1dee9c7d3dec22c5a0755ee9c13de 100644
--- a/resolve/plotter.py
+++ b/resolve/plotter.py
@@ -203,7 +203,8 @@ def _plot_nifty(state, op, kwargs, fname):
sc.add(op.force(ss))
p = ift.Plot()
p.add(sc.mean, **kwargs)
- p.add(sc.var.sqrt() / sc.mean)
+ if len(state) > 1:
+ p.add(sc.var.sqrt() / sc.mean)
p.output(nx=2, ny=1, xsize=2 * UNIT, ysize=UNIT, name=fname)
else:
pos = state if isinstance(state, ift.MultiField) else state.mean
diff --git a/resolve/polarization_space.py b/resolve/polarization_space.py
new file mode 100644
index 0000000000000000000000000000000000000000..dad3e786189775d14b398350c3d5e0c1402f902a
--- /dev/null
+++ b/resolve/polarization_space.py
@@ -0,0 +1,34 @@
+# SPDX-License-Identifier: GPL-3.0-or-later
+# Copyright(C) 2021 Max-Planck-Society
+# Author: Philipp Arras, Jakob Knollmüller
+
+import nifty8 as ift
+
+from .util import my_assert
+
+
+class PolarizationSpace(ift.UnstructuredDomain):
+ """
+
+ Parameters
+ ----------
+ coordinates : np.ndarray
+ Must be sorted and strictly ascending.
+ """
+
+ _needed_for_hash = ["_shape", "_lbl"]
+
+ def __init__(self, polarization_labels):
+ if isinstance(polarization_labels, str):
+ polarization_labels = (polarization_labels,)
+ self._lbl = tuple(polarization_labels)
+ for ll in self._lbl:
+ my_assert(ll in ["I", "Q", "U", "V", "LL", "LR", "RL", "RR", "XX", "XY", "YX", "YY"])
+ super(PolarizationSpace, self).__init__(len(self._lbl))
+
+ def __repr__(self):
+ return f"PolarizationSpace(polarization_labels={self._lbl})"
+
+ @property
+ def labels(self):
+ return self._lbl
diff --git a/resolve/response.py b/resolve/response.py
index c9d66ee3a5910ed8efb30b39266ddcef7b7052b8..fcb097b37d095eb5b611d7fe30a470a2e5084dfd 100644
--- a/resolve/response.py
+++ b/resolve/response.py
@@ -41,7 +41,7 @@ def StokesIResponse(observation, domain):
return contr.adjoint @ sr0
elif npol == 2:
sr1 = SingleResponse(domain, observation.uvw, observation.freq, mask[1])
- return ResponseDistributor(sr0, sr1)
+ return ResponseDistributor(observation.polarization.space, sr0, sr1)
raise RuntimeError
@@ -186,7 +186,7 @@ class MfResponse(ift.LinearOperator):
class ResponseDistributor(ift.LinearOperator):
- def __init__(self, *ops):
+ def __init__(self, iterating_target_space, *ops):
dom, tgt = ops[0].domain, ops[0].target
cap = self.TIMES | self.ADJOINT_TIMES
for op in ops:
@@ -194,8 +194,10 @@ class ResponseDistributor(ift.LinearOperator):
my_assert(dom is op.domain)
my_assert(tgt is op.target)
my_assert(self.TIMES & op.capability, self.ADJOINT_TIMES & op.capability)
+ my_asserteq(len(iterating_target_space.shape), 1)
+ my_asserteq(len(ops), iterating_target_space.shape[0])
self._domain = ift.makeDomain(dom)
- self._target = ift.makeDomain((ift.UnstructuredDomain(len(ops)), *tgt))
+ self._target = ift.makeDomain((iterating_target_space, *tgt))
self._capability = cap
self._ops = ops
@@ -222,7 +224,7 @@ class FullResponse(ift.LinearOperator):
class SingleResponse(ift.LinearOperator):
- def __init__(self, domain, uvw, freq, mask, facets=(1, 1)):
+ def __init__(self, domain, uvw, freq, mask=None, facets=(1, 1)):
my_assert_isinstance(facets, tuple)
for ii in range(2):
if domain.shape[ii] % facets[ii] != 0:
@@ -234,7 +236,6 @@ class SingleResponse(ift.LinearOperator):
self._args = {
"uvw": uvw,
"freq": freq,
- "mask": mask.astype(np.uint8),
"pixsize_x": self._domain[0].distances[0],
"pixsize_y": self._domain[0].distances[1],
"epsilon": epsilon(),
@@ -242,6 +243,8 @@ class SingleResponse(ift.LinearOperator):
"nthreads": nthreads(),
"flip_v": True,
}
+ if mask is not None:
+ self._args["mask"] = mask.astype(np.uint8)
self._vol = self._domain[0].scalar_dvol
self._target_dtype = np_dtype(True)
self._domain_dtype = np_dtype(False)
diff --git a/resolve/util.py b/resolve/util.py
index 4c34445507d89ae621eff0e4b74bb3c03db8e8ab..8093ef2e7395ca5b156d221fb1732122731a128e 100644
--- a/resolve/util.py
+++ b/resolve/util.py
@@ -28,20 +28,46 @@ def my_assert_isinstance(*args):
def compare_attributes(obj0, obj1, attribute_list):
- for a in attribute_list:
- compare = getattr(obj0, a) != getattr(obj1, a)
- if type(compare) is bool:
- if compare:
- return False
- continue
- if isinstance(compare, ift.MultiField):
- for vv in compare.val.values():
- if vv.any():
- return False
- continue
- if compare.any():
- return False
- return True
+ return all(_fancy_equal(getattr(obj0, a), getattr(obj1, a))
+ for a in attribute_list)
+
+
+def _fancy_equal(o1, o2):
+ if not _types_equal(o1, o2):
+ return False
+
+ # Turn MultiField into dict
+ if isinstance(o1, ift.MultiField):
+ o1, o2 = o1.val, o2.val
+
+ # Compare dicts
+ if isinstance(o1, dict):
+ return _deep_equal(o1, o2)
+
+ # Compare simple objects and np.ndarrays
+ return _compare_simple_or_array(o1, o2)
+
+
+def _deep_equal(a, b):
+ if not isinstance(a, dict) or not isinstance(b, dict):
+ raise TypeError
+
+ if a.keys() != b.keys():
+ return False
+
+ return all(_compare_simple_or_array(a[kk], b[kk]) for kk in a.keys())
+
+
+def _compare_simple_or_array(a, b):
+ equal = a == b
+ if isinstance(equal, np.ndarray):
+ return np.all(equal)
+ assert isinstance(equal, bool)
+ return equal
+
+
+def _types_equal(a, b):
+ return type(a) == type(b)
class Reshaper(ift.LinearOperator):
diff --git a/test/test_general.py b/test/test_general.py
index ec799a032b6adbf3a118bb8bfb758743fb1ec113..4c91daf9e53afe24a80456ca15d1891def4b38b6 100644
--- a/test/test_general.py
+++ b/test/test_general.py
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: GPL-3.0-or-later
-# Copyright(C) 2020 Max-Planck-Society
+# Copyright(C) 2020-2021 Max-Planck-Society
# Author: Philipp Arras
from os.path import join
@@ -42,10 +42,7 @@ freqdomain = rve.IRGSpace(np.linspace(1000, 1050, num=10))
FACETS = [(1, 1), (2, 2), (2, 1), (1, 4)]
-@pmp(
- "ms",
- ("CYG-ALL-2052-2MHZ", "CYG-D-6680-64CH-10S", "AM754_A030124_flagged", "1052735056"),
-)
+@pmp("ms", ("CYG-ALL-2052-2MHZ", "CYG-D-6680-64CH-10S", "AM754_A030124_flagged"))
@pmp("with_calib_info", (False, True))
@pmp("compress", (False, True))
def test_save_and_load_observation(ms, with_calib_info, compress):
@@ -58,11 +55,11 @@ def test_save_and_load_observation(ms, with_calib_info, compress):
assert ob == ob1
-@pmp("slc", [slice(3), slice(14, 15), slice(None, None, None), slice(None, None, 5)])
+@pmp("slc", [slice(3), slice(14, 15), [14, 15], slice(None, None, None), slice(None, None, 5)])
def test_sliced_readin(slc):
ms = f"{direc}CYG-D-6680-64CH-10S.ms"
obs0 = rve.ms2observations(ms, "DATA", False, 0)[0]
- obs = rve.ms2observations(ms, "DATA", False, 0, channel_slice=slc)[0]
+ obs = rve.ms2observations(ms, "DATA", False, 0, channels=slc)[0]
ch0 = obs0.weight.val[..., slc]
ch = obs.weight.val
assert ch0.ndim == 3
@@ -149,7 +146,7 @@ def test_calibration_likelihood(time_mode):
npol, nfreq = obs[0].npol, obs[0].nfreq
total_N = npol * nants * nfreq
dom = [
- ift.UnstructuredDomain(npol),
+ obs[0].polarization.space,
ift.UnstructuredDomain(len(uants)),
time_domain,
ift.UnstructuredDomain(nfreq),
@@ -216,10 +213,8 @@ def test_calibration_distributor(obs):
tgt = obs.vis.domain
utimes = rve.unique_times(obs)
uants = obs.antenna_positions.unique_antennas()
- dom = [
- ift.UnstructuredDomain(nn)
- for nn in [obs.npol, len(uants), len(utimes), obs.nfreq]
- ]
+ dom = [obs.polarization.space] + \
+ [ift.UnstructuredDomain(nn) for nn in [len(uants), len(utimes), obs.nfreq]]
uants = rve.unique_antennas(obs)
time_dct = {aa: ii for ii, aa in enumerate(utimes)}
antenna_dct = {aa: ii for ii, aa in enumerate(uants)}
@@ -262,7 +257,7 @@ def test_response_distributor():
dom = ift.UnstructuredDomain(2), ift.UnstructuredDomain(4)
op0 = ift.makeOp(ift.makeField(dom, np.arange(8).reshape(2, -1)))
op1 = ift.makeOp(ift.makeField(dom, 2 * np.arange(8).reshape(2, -1)))
- op = rve.response.ResponseDistributor(op0, op1)
+ op = rve.response.ResponseDistributor(ift.UnstructuredDomain(2), op0, op1)
ift.extra.check_linear_operator(op)
@@ -309,9 +304,11 @@ def test_randomonmaster():
def test_mfweighting():
nrow = 100
nchan = 4
- effuv = ift.random.current_rng().random((nrow, nchan))
- dom = ift.UnstructuredDomain(nchan), ift.RGSpace(npix, 2 * np.max(effuv) / npix)
- op = rve.MfWeightingInterpolation(effuv[None], dom)
+ npol = 1
+ effuv = ift.random.current_rng().random((npol, nrow, nchan))
+ effuv = ift.makeField([ift.UnstructuredDomain(ii) for ii in effuv.shape], effuv)
+ dom = ift.UnstructuredDomain(nchan), ift.RGSpace(npix, 2 * np.max(effuv.val) / npix)
+ op = rve.MfWeightingInterpolation(effuv, dom)
ift.extra.check_linear_operator(op)