Skip to content

Commit

Permalink
Merge pull request radio-astro-tools#765 from preshanth/stokes2feed
Browse files Browse the repository at this point in the history 
Attribute stokes_type
  • Loading branch information
@e-koch
e-koch authored Sep 26, 2022
2 parents de76f8a + 4a1e474 commit 1484b14
Show file tree
Hide file tree
Showing 2 changed files with 138 additions and 10 deletions.
99 changes: 96 additions & 3 deletions spectral_cube/stokes_spectral_cube.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,9 +12,13 @@
__all__ = ['StokesSpectralCube']

VALID_STOKES = ['I', 'Q', 'U', 'V', 'RR', 'LL', 'RL', 'LR', 'XX', 'XY', 'YX', 'YY',
'RX', 'RY', 'LX', 'LY', 'XR,', 'XL', 'YR', 'YL', 'PP', 'PQ', 'QP', 'QQ',
'RX', 'RY', 'LX', 'LY', 'XR', 'XL', 'YR', 'YL', 'PP', 'PQ', 'QP', 'QQ',
'RCircular', 'LCircular', 'Linear', 'Ptotal', 'Plinear', 'PFtotal',
'PFlinear', 'Pangle']
STOKES_SKY = ['I','Q','U','V']
STOKES_FEED_LINEAR = ['XX', 'XY', 'YX', 'YY']
STOKES_FEED_CIRCULAR = ['RR', 'RL', 'LR', 'LL']
STOKES_FEED_GENERIC = ['PP', 'PQ', 'QP', 'QQ']


class StokesSpectralCube(object):
Expand Down Expand Up @@ -45,8 +49,9 @@ def __init__(self, stokes_data, mask=None, meta=None, fill_value=None):
"should have the same WCS")

if component not in VALID_STOKES:
raise ValueError("Invalid Stokes component: {0} - should be "
"one of I, Q, U, V, RR, LL, RL, LR, XX, XY, YX, YY, RX, RY, LX, LY, XR, XL, YR, YL, PP, PQ, QP, QQ, RCircular, LCircular, Linear, Ptotal, Plinear, PFtotal, PFlinear, Pangle".format(component))
raise ValueError("Invalid Stokes component: {0} - should be one of I, Q, U, V, RR, LL, RL, LR, XX, XY, YX, YY, \
RX, RY, LX, LY, XR, XL, YR, YL, PP, PQ, QP, QQ, \
RCircular, LCircular, Linear, Ptotal, Plinear, PFtotal, PFlinear, Pangle".format(component))

if stokes_data[component].shape != stokes_data[reference].shape:
raise ValueError("All spectral cubes should have the same shape")
Expand All @@ -59,6 +64,17 @@ def __init__(self, stokes_data, mask=None, meta=None, fill_value=None):
raise ValueError("Mask shape is not broadcastable to data shape:"
" {0} vs {1}".format(mask.shape, self._shape))

if set(stokes_data).issubset(set(STOKES_SKY)):
self._stokes_type = 'SKY_STOKES'
elif set(stokes_data).issubset(set(STOKES_FEED_LINEAR)):
self._stokes_type = 'FEED_LINEAR'
elif set(stokes_data).issubset(set(STOKES_FEED_CIRCULAR)):
self._stokes_type = 'FEED_CIRCULAR'
elif set(stokes_data).issubset(set(STOKES_FEED_GENERIC)):
self._stokes_type = 'FEED_GENERIC'
elif set(stokes_data).issubset(set(VALID_STOKES)):
self._stokes_type = 'VALID_STOKES'

self._mask = mask

def __getitem__(self, key):
Expand Down Expand Up @@ -106,6 +122,17 @@ def __dir__(self):
def components(self):
return list(self._stokes_data.keys())

@property
def stokes_type(self):
"""
Defines the type of stokes that has been setup. `stokes_type` can be sky, linear, circular, generic, or other.
* `Sky` refers to stokes in sky basis, I,Q,U,V
* `Linear` refers to the stokes in linear feed basis, XX, XY, YX, YY
* `Circular` refers to stokes in circular feed basis, RR, RL, LR, LL
* `Generic` refers to the general four orthogonal components, PP, PQ, QP, QQ
"""
return self._stokes_type

def __getattr__(self, attribute):
"""
Descriptor to return the Stokes cubes
Expand Down Expand Up @@ -170,6 +197,72 @@ def _new_cube_with(self, stokes_data=None,

return cube

def transform_basis(self, stokes_basis=''):
"""
The goal of this function is to transform the stokes basis of the cube to
the one specified if possible. In principle one needs at least two related
stokes of a given basis for the transformation to be possible. At this moment
we are limiting it to the cases where all four stokes parameters in a given
basis be available within the cube. The transformations are as follows
Linear to Sky
I = (XX + YY) / 2
Q = (XX - YY) / 2
U = (XY + YX) / 2
V = 1j(XY - YX) / 2
Circular to Sky
I = (RR + LL) / 2
Q = (RL + LR) / 2
U = 1j*(RL - LR) / 2
V = (RR - LL) / 2
"""
if self.shape[0] < 4:
errmsg = "Transformation of a subset of Stokes axes is not yet supported."
errmsg_template = "Transformation from {} to {} is not yet supported."

raise NotImplementedError(errmsg)
elif stokes_basis == "Generic":
data = self._stokes_data
elif self.stokes_type == "Linear":
if stokes_basis == "Circular":
errmsg = errmsg_template.format("Linear", "Circular")
raise NotImplementedError(errmsg)
elif stokes_basis == "Sky":
data = dict(I = self.__class__(self._stokes_data['XX'] + self._stokes_data['YY'], wcs = self.wcs, mask=self._mask['XX']&self._mask['YY']),
Q = self.__class__(self._stokes_data['XX'] - self._stokes_data['YY'], wcs = self.wcs, mask=self._mask['XX']&self._mask['YY']),
U = self.__class__(self._stokes_data['XY'] + self._stokes_data['YX'], wcs = self.wcs, mask=self._mask['XY']&self._mask['YX']),
V = self.__class__(1j*(self._stokes_data['XY'] - self._stokes_data['YX']), wcs = self.wcs, mask=self._mask['XY']&self._mask['YX']))
elif stokes_basis == "Linear":
data = self._stokes_data

elif self.stokes_type == "Circular":
if stokes_basis == "Linear":
errmsg = errmsg_template.format("Circular", "Linear")
raise NotImplementedError(errmsg)
elif stokes_basis == "Sky":
data = dict(I = self.__class__(self._stokes_data['RR'] + self._stokes_data['LL'], wcs = self.wcs, mask=self._mask['RR']&self._mask['LL']),
Q = self.__class__(self._stokes_data['RL'] + self._stokes_data['RL'], wcs = self.wcs, mask=self._mask['RL']&self._mask['LR']),
U = self.__class__(1j*(self._stokes_data['RL'] - self._stokes_data['LR']), wcs = self.wcs, mask=self._mask['RL']&self._mask['LR']),
V = self.__class__(self._stokes_data['RR'] - self._stokes_data['LL'], wcs = self.wcs, mask=self._mask['RR']&self._mask['LL']))
elif stokes_basis == "Circular":
data = self._stokes_data

elif self.stokes_type == "Sky":
if stokes_basis == "Linear":
data = dict(XX = self.__class__(0.5*(self._stokes_data['I'] + self._stokes_data['Q']), wcs = self.wcs, mask=self._mask['I']&self._mask['Q']),
XY = self.__class__(0.5*(self._stokes_data['U'] + 1j*self._stokes_data['V']), wcs = self.wcs, mask=self._mask['U']&self._mask['V']),
YX = self.__class__(0.5*(self._stokes_data['U'] - 1j*self._stokes_data['V']), wcs = self.wcs, mask=self._mask['U']&self._mask['V']),
YY = self.__class__(0.5*(self._stokes_data['I'] - self._stokes_data['Q']), wcs = self.wcs, mask=self._mask['I']&self._mask['Q']))
elif stokes_basis == "Circular":
data = dict(RR = self.__class__(0.5*(self._stokes_data['I'] + self._stokes_data['V']), wcs = self.wcs, mask=self._mask['I']&self._mask['V']),
RL = self.__class__(0.5*(self._stokes_data['Q'] + 1j*self._stokes_data['U']), wcs = self.wcs, mask=self._mask['Q']&self._mask['U']),
LR = self.__class__(0.5*(self._stokes_data['Q'] - 1j*self._stokes_data['U']), wcs = self.wcs, mask=self._mask['Q']&self._mask['U']),
LL = self.__class__(0.5*(self._stokes_data['I'] - self._stokes_data['V']), wcs = self.wcs, mask=self._mask['I']&self._mask['V']))
elif stokes_basis == "Sky":
data = self._stokes_data
return self._new_cube_with(stokes_data=data)


def with_spectral_unit(self, unit, **kwargs):

stokes_data = {k: self._stokes_data[k].with_spectral_unit(unit, **kwargs)
Expand Down
49 changes: 42 additions & 7 deletions spectral_cube/tests/test_stokes_spectral_cube.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@
from ..spectral_cube import SpectralCube
from ..stokes_spectral_cube import StokesSpectralCube
from ..masks import BooleanArrayMask

from ..stokes_spectral_cube import VALID_STOKES

class TestStokesSpectralCube():

Expand Down Expand Up @@ -45,10 +45,7 @@ def test_direct_init_invalid_shape(self, use_dask):
cube = StokesSpectralCube(stokes_data)
assert exc.value.args[0] == "All spectral cubes should have the same shape"

@pytest.mark.parametrize('component', ('I', 'Q', 'U', 'V', 'RR', 'LL', 'RL', 'LR', 'XX', 'XY', 'YX', 'YY',
'RX', 'RY', 'LX', 'LY', 'XR,', 'XL', 'YR', 'YL', 'PP', 'PQ', 'QP', 'QQ',
'RCircular', 'LCircular', 'Linear', 'Ptotal', 'Plinear', 'PFtotal',
'PFlinear', 'Pangle'))
@pytest.mark.parametrize('component', VALID_STOKES)

def test_valid_component_name(self, component, use_dask):
stokes_data = {component: SpectralCube(self.data[0], wcs=self.wcs, use_dask=use_dask)}
Expand All @@ -60,7 +57,9 @@ def test_invalid_component_name(self, component, use_dask):
stokes_data = {component: SpectralCube(self.data[0], wcs=self.wcs, use_dask=use_dask)}
with pytest.raises(ValueError) as exc:
cube = StokesSpectralCube(stokes_data)
assert exc.value.args[0] == "Invalid Stokes component: {0} - should be one of I, Q, U, V, RR, LL, RL, LR, XX, XY, YX, YY, RX, RY, LX, LY, XR, XL, YR, YL, PP, PQ, QP, QQ, RCircular, LCircular, Linear, Ptotal, Plinear, PFtotal, PFlinear, Pangle".format(component)
assert exc.value.args[0] == "Invalid Stokes component: {0} - should be one of I, Q, U, V, RR, LL, RL, LR, XX, XY, YX, YY, \
RX, RY, LX, LY, XR, XL, YR, YL, PP, PQ, QP, QQ, \
RCircular, LCircular, Linear, Ptotal, Plinear, PFtotal, PFlinear, Pangle".format(component)

def test_invalid_wcs(self, use_dask):
wcs2 = WCS(naxis=3)
Expand All @@ -84,6 +83,40 @@ def test_attributes(self, use_dask):
assert_allclose(cube.V.unmasked_data[...], 3)
assert cube.components == ['I', 'Q', 'U', 'V']

def test_stokes_type_sky(self, use_dask):
stokes_data = OrderedDict()
stokes_data['I'] = SpectralCube(self.data[0], wcs=self.wcs, use_dask=use_dask)
stokes_data['Q'] = SpectralCube(self.data[1], wcs=self.wcs, use_dask=use_dask)
stokes_data['U'] = SpectralCube(self.data[2], wcs=self.wcs, use_dask=use_dask)
stokes_data['V'] = SpectralCube(self.data[3], wcs=self.wcs, use_dask=use_dask)
cube = StokesSpectralCube(stokes_data)
assert cube.stokes_type == "SKY_STOKES"

def test_stokes_type_feed_circular(self, use_dask):
stokes_data = OrderedDict()
stokes_data['RR'] = SpectralCube(self.data[0], wcs=self.wcs, use_dask=use_dask)
stokes_data['RL'] = SpectralCube(self.data[1], wcs=self.wcs, use_dask=use_dask)
stokes_data['LR'] = SpectralCube(self.data[2], wcs=self.wcs, use_dask=use_dask)
stokes_data['LL'] = SpectralCube(self.data[3], wcs=self.wcs, use_dask=use_dask)
cube = StokesSpectralCube(stokes_data)
assert cube.stokes_type == "FEED_CIRCULAR"

def test_stokes_type_feed_linear(self, use_dask):
stokes_data = OrderedDict()
stokes_data['XX'] = SpectralCube(self.data[0], wcs=self.wcs, use_dask=use_dask)
stokes_data['XY'] = SpectralCube(self.data[1], wcs=self.wcs, use_dask=use_dask)
stokes_data['YX'] = SpectralCube(self.data[2], wcs=self.wcs, use_dask=use_dask)
stokes_data['YY'] = SpectralCube(self.data[3], wcs=self.wcs, use_dask=use_dask)
cube = StokesSpectralCube(stokes_data)
assert cube.stokes_type == "FEED_LINEAR"

def test_stokes_type_feed_linear_partial(self, use_dask):
stokes_data = OrderedDict()
stokes_data['XX'] = SpectralCube(self.data[0], wcs=self.wcs, use_dask=use_dask)
stokes_data['YY'] = SpectralCube(self.data[1], wcs=self.wcs, use_dask=use_dask)
cube = StokesSpectralCube(stokes_data)
assert cube.stokes_type == "FEED_LINEAR"

def test_dir(self, use_dask):
stokes_data = dict(I=SpectralCube(self.data[0], wcs=self.wcs, use_dask=use_dask),
Q=SpectralCube(self.data[1], wcs=self.wcs, use_dask=use_dask),
Expand Down Expand Up @@ -119,7 +152,9 @@ def test_mask_invalid_component_name(self, use_dask):
stokes_data = {'BANANA': SpectralCube(self.data[0], wcs=self.wcs, use_dask=use_dask)}
with pytest.raises(ValueError) as exc:
cube = StokesSpectralCube(stokes_data)
assert exc.value.args[0] == "Invalid Stokes component: BANANA - should be one of I, Q, U, V, RR, LL, RL, LR, XX, XY, YX, YY, RX, RY, LX, LY, XR, XL, YR, YL, PP, PQ, QP, QQ, RCircular, LCircular, Linear, Ptotal, Plinear, PFtotal, PFlinear, Pangle"
assert exc.value.args[0] == "Invalid Stokes component: BANANA - should be one of I, Q, U, V, RR, LL, RL, LR, XX, XY, YX, YY, \
RX, RY, LX, LY, XR, XL, YR, YL, PP, PQ, QP, QQ, \
RCircular, LCircular, Linear, Ptotal, Plinear, PFtotal, PFlinear, Pangle"

def test_mask_invalid_shape(self, use_dask):
stokes_data = dict(I=SpectralCube(self.data[0], wcs=self.wcs, use_dask=use_dask),
Expand Down

0 comments on commit 1484b14

Please sign in to comment.