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/data/observation.py b/resolve/data/observation.py
index fd44cde7b90aac90940664e30da2689aa37f6fd4..b4c7934d80e8cd6791bda4fc348faeebe0a6310d 100644
--- a/resolve/data/observation.py
+++ b/resolve/data/observation.py
@@ -18,7 +18,8 @@ from .polarization import Polarization
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
@@ -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
@@ -476,6 +479,24 @@ class Observation(BaseObservation):
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)
@@ -501,6 +522,53 @@ class Observation(BaseObservation):
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
@@ -509,6 +577,18 @@ 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:
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 11aae3dbb3e0e2ca7fc867643daed1c265bb87ad..b77f98b2ecd8e21dacd600472e5297b722d63c76 100644
--- a/resolve/multi_frequency/operators.py
+++ b/resolve/multi_frequency/operators.py
@@ -2,10 +2,10 @@
# 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 ..polarization_space import PolarizationSpace
from ..util import my_assert, my_assert_isinstance, my_asserteq
from .irg_space import IRGSpace
@@ -91,7 +91,7 @@ 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
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/test/test_general.py b/test/test_general.py
index 44e255b9da498142fa55c73bd4b7686d93227bb5..5ec89ebc45b01fdde333752a1f16128cd1dbb3f8 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
@@ -150,7 +150,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),
@@ -217,10 +217,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)}
@@ -263,7 +261,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)