def test_resample1(self):
'''resample should give consistent results'''
ans = dm.SpaceData()
ans.attrs['foo'] = 'bar'
ans['a'] = [ 0.5, 2.5, 4.5, 6.5, 8.5]
ans['b'] = dm.dmarray([4.5, 6.5, 8.5, 10.5, 12.5])
ans['b'].attrs['marco'] = 'polo'
ans['Epoch'] = [datetime.datetime(2010, 1, 1, 1, 0),
datetime.datetime(2010, 1, 1, 3, 0),
datetime.datetime(2010, 1, 1, 5, 0),
datetime.datetime(2010, 1, 1, 7, 0),
datetime.datetime(2010, 1, 1, 9, 0)]
a = dm.SpaceData()
a['a'] = dm.dmarray(range(10)) #sequence 0 through 9
a['b'] = dm.dmarray(range(10)) + 4 #sequence 4 through 13
a['b'].attrs['marco'] = 'polo'
a['c'] = dm.dmarray(range(3)) + 10
a.attrs['foo'] = 'bar'
times = [datetime.datetime(2010, 1, 1) + datetime.timedelta(hours=i) for i in range(10)]
out = dm.resample(a, times, winsize=datetime.timedelta(hours=2), overlap=datetime.timedelta(hours=0))
#starting from element 0, step two hours (points) with no overlap and average
#therefore each resulting value should be the mean of a pair of points
#that is, ans['a'][0] should be the mean of 0 and 1
for k, v in out.items():
np.testing.assert_equal(v, ans[k])
self.assertEqual(ans.attrs, out.attrs)
self.assertEqual(ans['b'].attrs['marco'], 'polo')
self.assertTrue(out['b'].attrs['DEPEND_0'], 'Epoch')
self.assertFalse('c' in out)
def test_resample1(self):
'''resample should give consistent results'''
ans = dm.SpaceData()
ans.attrs['foo'] = 'bar'
ans['a'] = [ 1., 3., 5., 7.]
ans['b'] = dm.dmarray([5., 7., 9., 11.])
ans['b'].attrs['marco'] = 'polo'
ans['Epoch'] = [datetime.datetime(2010, 1, 1, 1, 0),
datetime.datetime(2010, 1, 1, 3, 0),
datetime.datetime(2010, 1, 1, 5, 0),
datetime.datetime(2010, 1, 1, 7, 0)]
a = dm.SpaceData()
a['a'] = dm.dmarray(range(10))
a['b'] = dm.dmarray(range(10)) + 4
a['b'].attrs['marco'] = 'polo'
a['c'] = dm.dmarray(range(3)) + 10
a.attrs['foo'] = 'bar'
times = [datetime.datetime(2010, 1, 1) + datetime.timedelta(hours=i) for i in range(10)]
out = dm.resample(a, times, winsize=datetime.timedelta(hours=2), overlap=datetime.timedelta(hours=0))
for k, v in out.items():
np.testing.assert_equal(v, ans[k])
self.assertEqual(ans.attrs, out.attrs)
self.assertEqual(ans['b'].attrs['marco'], 'polo')
self.assertTrue(out['b'].attrs['DEPEND_0'], 'Epoch')
self.assertFalse('c' in out)
def test_unflatten_function_nonflat(self):
"""unflatten will error for a nested input"""
a = dm.SpaceData()
a['1'] = dm.SpaceData(
dog=5, pig=dm.SpaceData(fish=dm.SpaceData(a='carp', b='perch')))
a['4'] = dm.SpaceData(cat='kitty')
a['5'] = 4
with self.assertRaises(TypeError):
dm.unflatten(a)
def test_HDF5roundtrip2GZIP(self):
"""Data can go to hdf without altering datetimes in the datamodel with compression"""
a = dm.SpaceData()
a['foo'] = dm.SpaceData()
dm.toHDF5(self.testfile, a, compression='gzip')
newobj = dm.fromHDF5(self.testfile)
self.assertEqual(a['foo'], newobj['foo'])
a['bar'] = dm.dmarray([datetime.datetime(2000, 1, 1)])
dm.toHDF5(self.testfile, a, compression='gzip')
self.assertEqual(a['bar'], dm.dmarray([datetime.datetime(2000, 1, 1)]))
def test_copy_toHDF(self):
"""Removed code in #404 for an old python bug, issue4380, make sure it works"""
a = dm.SpaceData({'dat': dm.dmarray([1, 2, 3])})
a2 = copy.deepcopy(a)
a2.toHDF5(self.testfile, mode='a')
newobj = dm.fromHDF5(self.testfile)
np.testing.assert_array_equal([1, 2, 3], newobj['dat'])
def test_resample3(self):
'''resample should give consistent results (2d)'''
ans = {}
ans['a'] = [[1., 2.], [5., 6.], [9., 10.], [13., 14.], [17., 18.]]
ans['b'] = [4.5, 6.5, 8.5, 10.5, 12.5]
ans['Epoch'] = [
datetime.datetime(2010, 1, 1, 1, 0),
datetime.datetime(2010, 1, 1, 3, 0),
datetime.datetime(2010, 1, 1, 5, 0),
datetime.datetime(2010, 1, 1, 7, 0),
datetime.datetime(2010, 1, 1, 9, 0)
]
a = dm.SpaceData()
a['a'] = dm.dmarray(range(10 * 2)).reshape(10, 2)
a['b'] = dm.dmarray(range(10)) + 4
a['c'] = dm.dmarray(range(3)) + 10
times = [
datetime.datetime(2010, 1, 1) + datetime.timedelta(hours=i)
for i in range(10)
]
out = dm.resample(a,
times,
winsize=datetime.timedelta(hours=2),
overlap=datetime.timedelta(hours=0))
for k, v in out.items():
np.testing.assert_equal(v, ans[k])
def test_resample2(self):
'''resample should give consistent results (ticktock)'''
ans = {}
ans['a'] = [0.5, 2.5, 4.5, 6.5, 8.5]
ans['b'] = [4.5, 6.5, 8.5, 10.5, 12.5]
ans['Epoch'] = [
datetime.datetime(2010, 1, 1, 1, 0),
datetime.datetime(2010, 1, 1, 3, 0),
datetime.datetime(2010, 1, 1, 5, 0),
datetime.datetime(2010, 1, 1, 7, 0),
datetime.datetime(2010, 1, 1, 9, 0)
]
#For justification of test results see test above (test_resample1)
a = dm.SpaceData()
a['a'] = dm.dmarray(range(10))
a['b'] = dm.dmarray(range(10)) + 4
a['c'] = dm.dmarray(range(3)) + 10
times = spt.Ticktock([
datetime.datetime(2010, 1, 1) + datetime.timedelta(hours=i)
for i in range(10)
])
out = dm.resample(a,
times,
winsize=datetime.timedelta(hours=2),
overlap=datetime.timedelta(hours=0))
for k, v in out.items():
np.testing.assert_equal(v, ans[k])
def readIMP8plasmafile(fname):
"""
ftp://space.mit.edu/pub/plasma/imp/fine_res/1989/
yr doy hh mm ss sc decimal yr rg md xse yse zse ysm zsm speed thermal speed density E/W angle N/S angle
mom nonlin mom nonlin mom nonlin mom best thresh threshs
"""
header = []
with open(fname, 'r') as fh:
while True:
pos = fh.tell()
line = fh.readline().strip()
if not line:
# empty line, skip
continue
if not line[0].isdigit():
# this is header, save it
header.append(line)
else:
# first line of data, roll back to start and pass to numpy
fh.seek(pos)
break
data = np.loadtxt(fh)
def toDate(yr, doy, hh, mm, ss):
MM, DD = spt.doy2date(yr, doy)
dates = dm.dmfilled(len(yr), fillval=None, dtype=object)
for idx, (mon, day) in enumerate(zip(MM, DD)):
dates[idx] = dt.datetime(yr[idx], mon, day, hh[idx], mm[idx],
ss[idx])
return dates
region = data[:, 7]
outdata = dm.SpaceData(attrs={'header': header, 'fname': fname})
outdata['time'] = toDate(data[:, 0].astype(int), data[:, 1].astype(int),
data[:, 2].astype(int), data[:, 3].astype(int),
data[:, 4].astype(int))
outdata['region'] = dm.dmarray(region)
outdata['pos_gse'] = dm.dmarray(data[:, 9:12], attrs={'coord_sys': 'gse'})
outdata['pos_gsm'] = dm.dmfilled(outdata['pos_gse'].shape,
fillval=0,
dtype=float,
attrs={'coord_sys': 'gsm'})
outdata['pos_gsm'][:, 0] = data[:, 9]
outdata['pos_gsm'][:, 1:] = data[:, 12:14]
outdata['speed'] = dm.dmarray(data[:, 14],
attrs={'description': 'speed from moments'})
# outdata['speed'][region > 2] = np.nan # region 3 is sheath
outdata['speed_nl'] = dm.dmarray(data[:, 15])
vmask = outdata['speed_nl'] >= 9000
# outdata['speed_nl'][region > 2] = np.nan # region 3 is sheath
outdata['speed_nl'][vmask] = np.nan # region 3 is sheath
outdata['n_dens'] = dm.dmarray(
data[:, 18], attrs={'description': 'number density from moments'})
outdata['n_dens_nl'] = dm.dmarray(data[:, 19])
outdata['temp'] = 60.5 * dm.dmarray(data[:, 16])**2
outdata['temp_nl'] = 60.5 * dm.dmarray(data[:, 17])**2
outdata['data'] = data
return outdata
def test_toJSONheadedASCII(self):
"""Write known datamodel to JSON-headed ASCII and ensure it has right stuff added"""
a = dm.SpaceData()
a.attrs['Global'] = 'A global attribute'
a['Var1'] = dm.dmarray([1, 2, 3, 4, 5],
attrs={'Local1': 'A local attribute'})
a['Var2'] = dm.dmarray([[8, 9], [9, 1], [3, 4], [8, 9], [7, 8]])
a['MVar'] = dm.dmarray([7.8], attrs={'Note': 'Metadata'})
t_file = tempfile.NamedTemporaryFile(delete=False)
t_file.close()
dm.toJSONheadedASCII(t_file.name,
a,
depend0='Var1',
order=['Var1', 'Var2'])
dat2 = dm.readJSONheadedASCII(t_file.name)
#test global attr
self.assertTrue(a.attrs == dat2.attrs)
#test that metadata is back and all original keys are present
for key in a['MVar'].attrs:
self.assertTrue(key in dat2['MVar'].attrs)
np.testing.assert_array_equal(a['MVar'], dat2['MVar'])
#test vars are right
np.testing.assert_almost_equal(a['Var1'], dat2['Var1'])
np.testing.assert_almost_equal(a['Var2'], dat2['Var2'])
#test for added dimension and start col
self.assertTrue(dat2['Var1'].attrs['DIMENSION'] == [1])
self.assertTrue(dat2['Var2'].attrs['DIMENSION'] == [2])
os.remove(t_file.name)
def hslice(self, value):
"""
slice a spectrogram at a given position along the y axis, maintains
variable names from spectrogram
Parameters
==========
value : float or datetime.datetime
the value to slice the spectrogram at
Returns
=======
out : datamodel.SpaceData
spacedata containing the slice
"""
# using bisect find the index of the spectrogram to use
if isinstance(value, datetime.datetime):
value = date2num(value)
ind = bisect.bisect_right(self['spectrogram']['yedges'], value)
ans = dm.SpaceData()
ans[self['spectrogram'].attrs['variables']
[0]] = tb.bin_edges_to_center(self['spectrogram']['xedges'])
ans['yedges'] = self['spectrogram']['yedges'][ind:ind + 2].copy()
ans['xedges'] = self['spectrogram']['xedges'].copy()
ans[self['spectrogram'].attrs['variables']
[2]] = self['spectrogram']['spectrogram'][ind, :]
return ans
def test_toHDF5ListString(self):
"""Convert to HDF5, including a list of string in attributes"""
a = dm.SpaceData()
a.attrs['foo'] = ['hi']
dm.toHDF5(self.testfile, a, mode='a')
newobj = dm.fromHDF5(self.testfile)
self.assertEqual(a.attrs['foo'], newobj.attrs['foo'])
def forecast():
RICE_URL_1_last = 'http://mms.rice.edu/realtime/Predictions_1.last'
RICE_URL_3_last = 'http://mms.rice.edu/realtime/Predictions_3.last'
RICE_Boyle_all = 'http://mms.rice.edu/realtime/File1.txt'
RICE_1hr_Kp_Dst = 'http://mms.rice.edu/realtime/File2.txt'
RICE_3hr_Kp_Dst = 'http://mms.rice.edu/realtime/File3.txt'
# grab all the 1 hour data
hr1 = u.urlopen(RICE_1hr_Kp_Dst)
data = hr1.readlines()
hr1.close()
dd1 = _parseRICE(data, '1')
# grab all the 3 hour data
hr3 = u.urlopen(RICE_3hr_Kp_Dst)
data = hr3.readlines()
hr3.close()
dd3 = _parseRICE(data, '3')
dd = dm.SpaceData()
dd.attrs['URL_1hr'] = RICE_1hr_Kp_Dst
dd.attrs['URL_3hr'] = RICE_3hr_Kp_Dst
dd.attrs['retrive_time'] = datetime.datetime.now()
for key1, key3 in zip(dd1, dd3):
dd[key1] = dd1[key1]
dd[key3] = dd3[key3]
return dd
def test_HDF5Exceptions(self):
"""HDF5 has warnings and exceptions"""
dm.toHDF5(self.testfile, self.SDobj)
self.assertRaises(IOError, dm.toHDF5, self.testfile, self.SDobj, overwrite=False)
a = dm.SpaceData()
a['foo'] = 'bar' # not an allowed type for data
self.assertRaises(dm.DMWarning, dm.toHDF5, self.testfile, a)
def setUp(self):
super(converterTestsCDF, self).setUp()
self.SDobj = dm.SpaceData(attrs={'global': 'test'})
self.SDobj['var'] = dm.dmarray([1, 2, 3], attrs={'a': 'a'})
self.testdir = tempfile.mkdtemp()
self.testfile = os.path.join(self.testdir, 'test.cdf')
warnings.simplefilter('error', dm.DMWarning)
def test_writemeta(self):
"""Check on writing to the meta property"""
a = dm.SpaceData(attrs={'a': 1, 'b': 2})
a.meta['c'] = 3
self.assertEqual(3, a.attrs['c'])
a.meta['a'] = 99
self.assertEqual(99, a.attrs['a'])
def test_unflatten_function(self):
"""Unflatten should unflatten a flattened SpaceData"""
a = dm.SpaceData()
a['1'] = dm.SpaceData(dog = 5, pig = dm.SpaceData(fish=dm.SpaceData(a='carp', b='perch')))
a['4'] = dm.SpaceData(cat = 'kitty')
a['5'] = 4
a[9] = dm.dmarray([1,2,3])
b = dm.flatten(a)
c = dm.unflatten(b)
self.assertTrue(9 in a.keys())
self.assertTrue(9 in c.keys())
del a[9]
del c[9]
self.assertEqual(sorted(a.keys()), sorted(c.keys()))
self.assertEqual(sorted(a['1'].keys()), sorted(c['1'].keys()))
self.assertEqual(sorted(a['1']['pig'].keys()), sorted(c['1']['pig'].keys()))
self.assertEqual(sorted(a['1']['pig']['fish'].keys()), sorted(c['1']['pig']['fish'].keys()))
def setUp(self):
super(spectrogramTests, self).setUp()
self.kwargs = {}
self.kwargs['variables'] = ['xval', 'yval', 'zval']
np.random.seed(8675309)
self.data = dm.SpaceData(xval=dm.dmarray(np.random.random_sample(200)),
yval=dm.dmarray(np.random.random_sample(200)),
zval=dm.dmarray(np.random.random_sample(200)))
def test_pickle(self):
"""Make sure that SpaceData objects behave with pickle"""
a = dm.SpaceData({'a': [1, 2, 3]})
a.attrs['FILLVAL'] = 123
p_str = pickle.dumps(a)
ret = pickle.loads(p_str)
assert a == ret
assert a.attrs == ret.attrs
def test_marshal(self):
"""SpaceData objects don't marshal, note that here"""
a = dm.SpaceData({'a': [1, 2, 3]})
a.attrs['FILLVAL'] = 123
p_str = marshal.dumps(a)
ret = marshal.loads(p_str)
assert a == ret
assert a.attrs == ret.attrs
def setUp(self):
super(spectrogramDateTests, self).setUp()
self.kwargs = {'variables': ['xval', 'yval', 'zval']}
np.random.seed(8675309)
self.data = dm.SpaceData(
xval=dm.dmarray([datetime.datetime(2000, 1, 1) + datetime.timedelta(days=nn) for nn in range(200)]),
yval=dm.dmarray(np.random.random_sample(200)),
zval=dm.dmarray(np.random.random_sample(200)))
def test_resample_shape(self):
'''resample should give consistent results, 1d or 2d'''
a = dm.SpaceData()
a['a'] = dm.dmarray(range(10*3*4)).reshape(10,3,4)
a['b'] = dm.dmarray(range(10)) + 4
a['c'] = dm.dmarray(range(3)) + 10
times = [datetime.datetime(2010, 1, 1) + datetime.timedelta(hours=i) for i in range(10)]
self.assertRaises(IndexError, dm.resample, a, times, datetime.timedelta(hours=2), datetime.timedelta(hours=0))
def test_dateToISOunaltered_SD(self):
"""Test to check that _dateToISO doesn't change datatypes, input SpaceData"""
data = dm.SpaceData()
data['Epoch'] = spt.tickrange('20200101', '20200102',
deltadays=datetime.timedelta(hours=1)).UTC
exptype = type(data['Epoch'][0])
newdata = dm._dateToISO(data)
restype = type(data['Epoch'][0])
self.assertEqual(exptype, restype)
def test_flatten_function(self):
"""Flatten should flatten a nested SpaceData"""
a = dm.SpaceData()
a['1'] = dm.SpaceData(dog = 5, pig = dm.SpaceData(fish=dm.SpaceData(a='carp', b='perch')))
a['4'] = dm.SpaceData(cat = 'kitty')
a['5'] = 4
self.assertEqual(a['1']['dog'], 5)
b = dm.flatten(a)
try:
b['1']['dog']
except KeyError:
pass
else:
self.fail('KeyError not raised')
# might be possible that list order is not preserved and this fails,
# if so change to a bunch of self.assertTrue and in statements
self.assertEqual(sorted(b.keys()),
sorted(['1<--pig<--fish<--a', '4<--cat', '1<--dog', '1<--pig<--fish<--b', '5']))
def test_toRecArray_contents(self):
'''a record array can be created from a SpaceData, keys and values equal'''
sd = dm.SpaceData()
sd['x'] = dm.dmarray([1.0, 2.0])
sd['y'] = dm.dmarray([2,4])
ra = dm.toRecArray(sd)
np.testing.assert_equal(ra['x'], [1.0, 2.0])
np.testing.assert_equal(ra['y'], [2, 4])
self.assertEqual(['x', 'y'], sorted(ra.dtype.fields))
def test_toRecArray_dtypes1(self):
'''recarray created from dmarray preserves data types (32-bit)'''
sd = dm.SpaceData()
sd['x'] = dm.dmarray([1.0, 2.0], dtype=np.float32)
sd['y'] = dm.dmarray([2,4], dtype=np.int32)
ra = dm.toRecArray(sd)
expected = [sd[key].dtype for key in sd]
got = [ra.dtype[name] for name in ra.dtype.names]
self.assertEqual(expected, got)
def test_toRecArray(self):
'''a record array can be created from a SpaceData'''
sd = dm.SpaceData()
sd['x'] = dm.dmarray([1.0, 2.0])
sd['y'] = dm.dmarray([2,4])
ra = dm.toRecArray(sd)
np.testing.assert_equal(ra['x'], [1.0, 2.0])
np.testing.assert_equal(ra['y'], [2, 4])
self.assertEqual(ra.dtype, np.dtype((np.record, [('x', '<f8'), ('y', '<i8'), ])))
def test_flatten_method(self):
"""Flatten should flatted a nested dict"""
a = dm.SpaceData()
a['1'] = dm.SpaceData(dog = 5, pig = dm.SpaceData(fish=dm.SpaceData(a='carp', b='perch')))
a['4'] = dm.SpaceData(cat = 'kitty')
a['5'] = 4
self.assertEqual(a['1']['dog'], 5)
a.flatten()
try:
a['1']['dog']
except KeyError:
pass
else:
self.fail('KeyError not raised')
ans = ['4<--cat', '1<--dog', '5', '1<--pig<--fish<--a', '1<--pig<--fish<--b']
ans.sort()
val = sorted(a.keys())
self.assertEqual(val, ans)
def plotSpectrogram(self, ecol=0, **kwargs):
'''
Plot a spectrogram of the flux along the requested orbit, as a function of Lm and time
Other Parameters
----------------
zlim : list
2-element list with upper and lower bounds for color scale
colorbar_label : string
text to appear next to colorbar (default is 'Flux' plus the units)
ylabel : string
text to label y-axis (default is 'Lm' plus the field model name)
title : string
text to appear above spectrogram (default is climatology model name, data type and energy)
'''
import spacepy.plot as splot
if 'Lm' not in self:
self.getLm()
sd = dm.SpaceData()
sd['Lm'] = self['Lm']
#filter any bad Lm
goodidx = sd['Lm'] > 1
sd['Lm'] = sd['Lm'][goodidx]
Lm_lim = [2.0, 8.0]
#TODO: allow user-definition of bins in time and Lm
varname = self.attrs['varname']
sd['Epoch'] = dm.dmcopy(
self['Epoch'])[goodidx] #TODO: assumes 1 pitch angle, generalize
sd['1D_dataset'] = self[varname][
goodidx, ecol] #TODO: assumes 1 pitch angle, generalize
spec = splot.spectrogram(sd,
variables=['Epoch', 'Lm', '1D_dataset'],
ylim=Lm_lim)
if 'zlim' not in kwargs:
zmax = 10**(int(np.log10(max(sd['1D_dataset']))) + 1)
idx = np.logical_and(sd['1D_dataset'] > 0, sd['Lm'] > Lm_lim[0])
idx = np.logical_and(idx, sd['Lm'] <= Lm_lim[1])
zmin = 10**int(np.log10(min(sd['1D_dataset'][idx])))
kwargs['zlim'] = [zmin, zmax]
if 'colorbar_label' not in kwargs:
flux_units = self[varname].attrs['UNITS']
kwargs['colorbar_label'] = '{0} ['.format(varname) + re.sub(
'(\^[\d|-]*)+', _grp2mathmode, flux_units) + ']'
if 'ylabel' not in kwargs:
kwargs['ylabel'] = 'L$_M$' + ' ' + '[{0}]'.format(
self['Lm'].attrs['MODEL'])
if 'title' not in kwargs:
kwargs['title'] = '{model_type} {varname}: '.format(**self.attrs) + \
'{0} {1}'.format(self['Energy'][ecol], self['Energy'].attrs['UNITS'])
reset_shrink = splot.mpl.mathtext.SHRINK_FACTOR
splot.mpl.mathtext.SHRINK_FACTOR = 0.85
splot.mpl.mathtext.GROW_FACTOR = 1 / 0.85
ax = spec.plot(cmap='plasma', **kwargs)
splot.mpl.mathtext.SHRINK_FACTOR = reset_shrink
splot.mpl.mathtext.GROW_FACTOR = 1 / reset_shrink
return ax
def test_toRecArray_dtypes2(self):
'''recarray created from dmarray preserves data types (16-bit+str)'''
sd = dm.SpaceData()
sd['x'] = dm.dmarray([1.0, 2.0], dtype=np.float16)
sd['y'] = dm.dmarray([2,4], dtype=np.int16)
sd['str'] = dm.dmarray(['spam', 'eggs'], dtype='|S5')
ra = dm.toRecArray(sd)
expected = [sd[key].dtype for key in sd]
got = [ra.dtype[name] for name in ra.dtype.names]
self.assertEqual(expected, got)
def test_multiget(self):
'''Allow for multiple keys to be specified'''
a = dm.SpaceData()
a['a'] = dm.dmarray([1,2,3])
a['b'] = dm.dmarray([1,2,3])
a['c'] = dm.dmarray([1,2,3])
a.attrs['foo']='bar'
np.testing.assert_equal(a['a'], dm.dmarray([1,2,3]))
np.testing.assert_equal(a[['a', 'b']], {'a': dm.dmarray([1, 2, 3]), 'b': dm.dmarray([1, 2, 3])})
self.assertRaises(KeyError, a.__getitem__, 'NotAkey')
self.assertRaises(KeyError, a.__getitem__, ['a', 'nokey'])
