def test_group_mixins():
"""
Test grouping a table with mixin columns
"""
# Setup mixins
idx = np.arange(4)
x = np.array([3., 1., 2., 1.])
q = x * u.m
lon = coordinates.Longitude(x * u.deg)
lat = coordinates.Latitude(x * u.deg)
# For Time do J2000.0 + few * 0.1 ns (this requires > 64 bit precision)
tm = time.Time(2000, format='jyear') + time.TimeDelta(x * 1e-10,
format='sec')
sc = coordinates.SkyCoord(ra=lon, dec=lat)
aw = table_helpers.ArrayWrapper(x)
nd = np.array([(3, 'c'), (1, 'a'), (2, 'b'), (1, 'a')],
dtype='<i4,|S1').view(NdarrayMixin)
qt = QTable([idx, x, q, lon, lat, tm, sc, aw, nd],
names=['idx', 'x', 'q', 'lon', 'lat', 'tm', 'sc', 'aw', 'nd'])
# Test group_by with each supported mixin type
mixin_keys = ['x', 'q', 'lon', 'lat', 'tm', 'sc', 'aw', 'nd']
for key in mixin_keys:
qtg = qt.group_by(key)
# Test that it got the sort order correct
assert np.all(qtg['idx'] == [1, 3, 2, 0])
# Test that the groups are right
# Note: skip testing SkyCoord column because that doesn't have equality
for name in ['x', 'q', 'lon', 'lat', 'tm', 'aw', 'nd']:
assert np.all(qt[name][[1, 3]] == qtg.groups[0][name])
assert np.all(qt[name][[2]] == qtg.groups[1][name])
assert np.all(qt[name][[0]] == qtg.groups[2][name])
# Test that unique also works with mixins since most of the work is
# done with group_by(). This is using *every* mixin as key.
uqt = unique(qt, keys=mixin_keys)
assert len(uqt) == 3
assert np.all(uqt['idx'] == [1, 2, 0])
assert np.all(uqt['x'] == [1., 2., 3.])
# Column group_by() with mixins
idxg = qt['idx'].group_by(qt[mixin_keys])
assert np.all(idxg == [1, 3, 2, 0])
time.Time([2000, 2001, 2002, 2003], format='jyear'),
'skycoord':
coordinates.SkyCoord(ra=[0, 1, 2, 3] * u.deg, dec=[0, 1, 2, 3] * u.deg),
'sphericalrep':
coordinates.SphericalRepresentation([0, 1, 2, 3] * u.deg,
[0, 1, 2, 3] * u.deg, 1 * u.kpc),
'cartesianrep':
coordinates.CartesianRepresentation([0, 1, 2, 3] * u.pc,
[4, 5, 6, 7] * u.pc,
[9, 8, 8, 6] * u.pc),
'sphericaldiff':
coordinates.SphericalCosLatDifferential([0, 1, 2, 3] * u.mas / u.yr,
[0, 1, 2, 3] * u.mas / u.yr,
10 * u.km / u.s),
'arraywrap':
table_helpers.ArrayWrapper([0, 1, 2, 3]),
'ndarray':
np.array([(7, 'a'), (8, 'b'), (9, 'c'), (9, 'c')],
dtype='<i4,|S1').view(table.NdarrayMixin),
}
MIXIN_COLS['earthlocation'] = coordinates.EarthLocation(
lon=MIXIN_COLS['longitude'],
lat=MIXIN_COLS['latitude'],
height=MIXIN_COLS['quantity'])
MIXIN_COLS['sphericalrepdiff'] = coordinates.SphericalRepresentation(
MIXIN_COLS['sphericalrep'], differentials=MIXIN_COLS['sphericaldiff'])
@pytest.fixture(params=sorted(MIXIN_COLS))
def mixin_cols(request):
"""
request.param
return MaskedTable
except AttributeError:
return table.Table
# Stuff for testing mixin columns
MIXIN_COLS = {'quantity': [0, 1, 2, 3] * u.m,
'longitude': coordinates.Longitude([0., 1., 5., 6.]*u.deg,
wrap_angle=180.*u.deg),
'latitude': coordinates.Latitude([5., 6., 10., 11.]*u.deg),
'time': time.Time([2000, 2001, 2002, 2003], format='jyear'),
'skycoord': coordinates.SkyCoord(ra=[0, 1, 2, 3] * u.deg,
dec=[0, 1, 2, 3] * u.deg),
'arraywrap': table_helpers.ArrayWrapper([0, 1, 2, 3]),
'ndarray': np.array([(7, 'a'), (8, 'b'), (9, 'c'), (9, 'c')],
dtype='<i4,|S1').view(table.NdarrayMixin),
}
MIXIN_COLS['earthlocation'] = coordinates.EarthLocation(
lon=MIXIN_COLS['longitude'], lat=MIXIN_COLS['latitude'],
height=MIXIN_COLS['quantity'])
@pytest.fixture(params=sorted(MIXIN_COLS))
def mixin_cols(request):
"""
Fixture to return a set of columns for mixin testing which includes
an index column 'i', two string cols 'a', 'b' (for joins etc), and
one of the available mixin column types.
"""