1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
|
from collections import OrderedDict
import numpy as np
from numpy.testing import assert_allclose, assert_equal
from astropy.wcs import WCS
from astropy.tests.helper import pytest
from astropy.utils import NumpyRNGContext
from ..spectral_cube import SpectralCube
from ..stokes_spectral_cube import StokesSpectralCube
from ..masks import BooleanArrayMask
from ..stokes_spectral_cube import VALID_STOKES
class TestStokesSpectralCube():
def setup_class(self):
self.wcs = WCS(naxis=3)
self.wcs.wcs.ctype = ['RA---TAN', 'DEC--TAN', 'FREQ']
self.data = np.arange(4)[:, None, None, None] * np.ones((5, 20, 30))
def test_direct_init(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),
U=SpectralCube(self.data[2], wcs=self.wcs, use_dask=use_dask),
V=SpectralCube(self.data[3], wcs=self.wcs, use_dask=use_dask))
cube = StokesSpectralCube(stokes_data)
def test_direct_init_invalid_type(self, use_dask):
stokes_data = dict(I=self.data[0],
Q=self.data[1],
U=self.data[2],
V=self.data[3])
with pytest.raises(TypeError) as exc:
cube = StokesSpectralCube(stokes_data)
assert exc.value.args[0] == "stokes_data should be a dictionary of SpectralCube objects"
def test_direct_init_invalid_shape(self, use_dask):
stokes_data = dict(I=SpectralCube(np.ones((6, 2, 30)), wcs=self.wcs, use_dask=use_dask),
Q=SpectralCube(self.data[1], wcs=self.wcs, use_dask=use_dask),
U=SpectralCube(self.data[2], wcs=self.wcs, use_dask=use_dask),
V=SpectralCube(self.data[3], wcs=self.wcs, use_dask=use_dask))
with pytest.raises(ValueError) as exc:
cube = StokesSpectralCube(stokes_data)
assert exc.value.args[0] == "All spectral cubes should have the same shape"
@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)}
cube = StokesSpectralCube(stokes_data)
assert cube.components == [component]
@pytest.mark.parametrize('component', ('A', 'B', 'IQUV'))
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)
def test_invalid_wcs(self, use_dask):
wcs2 = WCS(naxis=3)
wcs2.wcs.ctype = ['GLON-CAR', 'GLAT-CAR', 'FREQ']
stokes_data = dict(I=SpectralCube(self.data[0], wcs=self.wcs, use_dask=use_dask),
Q=SpectralCube(self.data[1], wcs2))
with pytest.raises(ValueError) as exc:
cube = StokesSpectralCube(stokes_data)
assert exc.value.args[0] == "All spectral cubes in stokes_data should have the same WCS"
def test_attributes(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_allclose(cube.I.unmasked_data[...], 0)
assert_allclose(cube.Q.unmasked_data[...], 1)
assert_allclose(cube.U.unmasked_data[...], 2)
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),
U=SpectralCube(self.data[2], wcs=self.wcs, use_dask=use_dask))
cube = StokesSpectralCube(stokes_data)
attributes = dir(cube)
for stokes in 'IQU':
assert stokes in attributes
assert 'V' not in attributes
assert 'mask' in attributes
assert 'wcs' in attributes
assert 'shape' in attributes
def test_mask(self, use_dask):
with NumpyRNGContext(12345):
mask1 = BooleanArrayMask(np.random.random((5, 20, 30)) > 0.2, wcs=self.wcs)
# Deliberately don't use a BooleanArrayMask to check auto-conversion
mask2 = np.random.random((5, 20, 30)) > 0.4
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),
U=SpectralCube(self.data[2], wcs=self.wcs, use_dask=use_dask),
V=SpectralCube(self.data[3], wcs=self.wcs, use_dask=use_dask))
cube1 = StokesSpectralCube(stokes_data, mask=mask1)
cube2 = cube1.with_mask(mask2)
assert_equal(cube2.mask.include(), (mask1).include() & mask2)
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"
def test_mask_invalid_shape(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),
U=SpectralCube(self.data[2], wcs=self.wcs, use_dask=use_dask),
V=SpectralCube(self.data[3], wcs=self.wcs, use_dask=use_dask))
mask1 = BooleanArrayMask(np.random.random((5, 20, 15)) > 0.2, wcs=self.wcs)
with pytest.raises(ValueError) as exc:
cube1 = StokesSpectralCube(stokes_data, mask=mask1)
assert exc.value.args[0] == "Mask shape is not broadcastable to data shape: (5, 20, 15) vs (5, 20, 30)"
def test_separate_mask(self, use_dask):
with NumpyRNGContext(12345):
mask1 = BooleanArrayMask(np.random.random((5, 20, 30)) > 0.2, wcs=self.wcs)
mask2 = [BooleanArrayMask(np.random.random((5, 20, 30)) > 0.4, wcs=self.wcs) for i in range(4)]
mask3 = BooleanArrayMask(np.random.random((5, 20, 30)) > 0.2, wcs=self.wcs)
stokes_data = dict(I=SpectralCube(self.data[0], wcs=self.wcs, mask=mask2[0], use_dask=use_dask),
Q=SpectralCube(self.data[1], wcs=self.wcs, mask=mask2[1], use_dask=use_dask),
U=SpectralCube(self.data[2], wcs=self.wcs, mask=mask2[2], use_dask=use_dask),
V=SpectralCube(self.data[3], wcs=self.wcs, mask=mask2[3], use_dask=use_dask))
cube1 = StokesSpectralCube(stokes_data, mask=mask1)
assert_equal(cube1.I.mask.include(), (mask1 & mask2[0]).include())
assert_equal(cube1.Q.mask.include(), (mask1 & mask2[1]).include())
assert_equal(cube1.U.mask.include(), (mask1 & mask2[2]).include())
assert_equal(cube1.V.mask.include(), (mask1 & mask2[3]).include())
cube2 = cube1.I.with_mask(mask3)
assert_equal(cube2.mask.include(), (mask1 & mask2[0] & mask3).include())
def test_key_access_valid(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_equal(cube['I'],cube._stokes_data['I'])
assert_equal(cube['Q'],cube._stokes_data['Q'])
assert_equal(cube['U'],cube._stokes_data['U'])
assert_equal(cube['V'],cube._stokes_data['V'])
|
