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'])