def test_smartTimeTicksSubDay(self):
"""smartTimeTicks should give known output (regression)"""
# hits all the different cases
# else
t1 = tb.linspace(datetime.datetime(2000, 1, 1),
datetime.datetime(2000, 1, 10), 20)
Mtick, mtick, fmt = spacepy.plot.utils.smartTimeTicks(t1)
self.assertEqual('%d %b', fmt.fmt)
# elif nHours < 4:
t1 = tb.linspace(datetime.datetime(2000, 1, 1),
datetime.datetime(2000, 1, 1, 1), 20)
Mtick, mtick, fmt = spacepy.plot.utils.smartTimeTicks(t1)
self.assertEqual('%H:%M UT', fmt.fmt)
# elif nHours < 24:
t1 = tb.linspace(datetime.datetime(2000, 1, 1),
datetime.datetime(2000, 1, 1, 13), 20)
Mtick, mtick, fmt = spacepy.plot.utils.smartTimeTicks(t1)
self.assertEqual('%H:%M UT', fmt.fmt)
# elif nHours < 12:
t1 = tb.linspace(datetime.datetime(2000, 1, 1),
datetime.datetime(2000, 1, 1, 11), 20)
Mtick, mtick, fmt = spacepy.plot.utils.smartTimeTicks(t1)
self.assertEqual('%H:%M UT', fmt.fmt)
# if nHours < 1:
t1 = tb.linspace(datetime.datetime(2000, 1, 1),
datetime.datetime(2000, 1, 1, 0, 30), 20)
Mtick, mtick, fmt = spacepy.plot.utils.smartTimeTicks(t1)
self.assertEqual('%H:%M UT', fmt.fmt)
def test_linspace(self):
"""linspace should return know answer for known input"""
try:
from matplotlib.dates import date2num
except ImportError: # just pass if matplotlib is not installed
return
# should exactly match here since it is the same
numpy.testing.assert_equal(tb.linspace(10, 100, 200), numpy.linspace(10,100, 200))
t1 = datetime.datetime(2000, 1, 1)
t2 = datetime.datetime(2000, 1, 10)
real_ans = [datetime.datetime(2000, 1, 1, 0, 0),
datetime.datetime(2000, 1, 3, 6, 0),
datetime.datetime(2000, 1, 5, 12, 0),
datetime.datetime(2000, 1, 7, 18, 0),
datetime.datetime(2000, 1, 10, 0, 0)]
ans = tb.linspace(t1, t2, 5)
numpy.testing.assert_almost_equal(date2num(real_ans), date2num(ans) , 4)
def test_linspace(self):
"""linspace should return know answer for known input"""
try:
from matplotlib.dates import date2num
except ImportError: # just pass if matplotlib is not installed
return
# should exactly match here since it is the same
numpy.testing.assert_equal(tb.linspace(10, 100, 200), numpy.linspace(10,100, 200))
t1 = datetime.datetime(2000, 1, 1)
t2 = datetime.datetime(2000, 1, 10)
real_ans = [datetime.datetime(2000, 1, 1, 0, 0),
datetime.datetime(2000, 1, 3, 6, 0),
datetime.datetime(2000, 1, 5, 12, 0),
datetime.datetime(2000, 1, 7, 18, 0),
datetime.datetime(2000, 1, 10, 0, 0)]
ans = tb.linspace(t1, t2, 5)
numpy.testing.assert_almost_equal(date2num(real_ans), date2num(ans) , 4)
def test_smartTimeTicksLonger(self):
"""smartTimeTicks should give known output (regression)"""
# elif nHours < 48:
t1 = tb.linspace(datetime.datetime(2000, 1, 1), datetime.datetime(2000, 1, 2, 0, 30), 20)
Mtick, mtick, fmt = spacepy.plot.utils.smartTimeTicks(t1)
self.assertEqual('%H:%M UT', fmt.fmt)
# elif nDays < 32:
t1 = tb.linspace(datetime.datetime(2000, 1, 2), datetime.datetime(2000, 2, 1, 0, 30), 20)
Mtick, mtick, fmt = spacepy.plot.utils.smartTimeTicks(t1)
self.assertEqual('%d %b', fmt.fmt)
# elif nDays < 731:
t1 = tb.linspace(datetime.datetime(2000, 1, 2), datetime.datetime(2001, 2, 1, 0, 30), 20)
Mtick, mtick, fmt = spacepy.plot.utils.smartTimeTicks(t1)
self.assertEqual('%b %Y', fmt.fmt)
# elif nDays >= 731:
t1 = tb.linspace(datetime.datetime(2000, 1, 2), datetime.datetime(2010, 2, 1, 0, 30), 20)
Mtick, mtick, fmt = spacepy.plot.utils.smartTimeTicks(t1)
self.assertEqual('%Y', fmt.fmt)
def test_smartTimeTicksLonger(self):
"""smartTimeTicks should give known output (regression)"""
# elif nHours < 48:
t1 = tb.linspace(datetime.datetime(2000, 1, 1), datetime.datetime(2000, 1, 2, 0, 30), 20)
Mtick, mtick, fmt = spacepy.plot.utils.smartTimeTicks(t1)
self.assertEqual('%H:%M UT', fmt.fmt)
# elif nDays < 32:
t1 = tb.linspace(datetime.datetime(2000, 1, 2), datetime.datetime(2000, 2, 1, 0, 30), 20)
Mtick, mtick, fmt = spacepy.plot.utils.smartTimeTicks(t1)
self.assertEqual('%d %b', fmt.fmt)
# elif nDays < 731:
t1 = tb.linspace(datetime.datetime(2000, 1, 2), datetime.datetime(2001, 2, 1, 0, 30), 20)
Mtick, mtick, fmt = spacepy.plot.utils.smartTimeTicks(t1)
self.assertEqual('%b %Y', fmt.fmt)
# elif nDays >= 731:
t1 = tb.linspace(datetime.datetime(2000, 1, 2), datetime.datetime(2010, 2, 1, 0, 30), 20)
Mtick, mtick, fmt = spacepy.plot.utils.smartTimeTicks(t1)
self.assertEqual('%Y', fmt.fmt)
def test_smartTimeTicksSubDay(self):
"""smartTimeTicks should give known output (regression)"""
# hits all the different cases
# else
t1 = tb.linspace(datetime.datetime(2000, 1, 1), datetime.datetime(2000, 1, 10), 20)
Mtick, mtick, fmt = spacepy.plot.utils.smartTimeTicks(t1)
self.assertEqual('%d %b', fmt.fmt)
# elif nHours < 4:
t1 = tb.linspace(datetime.datetime(2000, 1, 1), datetime.datetime(2000, 1, 1, 1), 20)
Mtick, mtick, fmt = spacepy.plot.utils.smartTimeTicks(t1)
self.assertEqual('%H:%M UT', fmt.fmt)
# elif nHours < 24:
t1 = tb.linspace(datetime.datetime(2000, 1, 1), datetime.datetime(2000, 1, 1, 13), 20)
Mtick, mtick, fmt = spacepy.plot.utils.smartTimeTicks(t1)
self.assertEqual('%H:%M UT', fmt.fmt)
# elif nHours < 12:
t1 = tb.linspace(datetime.datetime(2000, 1, 1), datetime.datetime(2000, 1, 1, 11), 20)
Mtick, mtick, fmt = spacepy.plot.utils.smartTimeTicks(t1)
self.assertEqual('%H:%M UT', fmt.fmt)
# if nHours < 1:
t1 = tb.linspace(datetime.datetime(2000, 1, 1), datetime.datetime(2000, 1, 1, 0, 30), 20)
Mtick, mtick, fmt = spacepy.plot.utils.smartTimeTicks(t1)
self.assertEqual('%H:%M UT', fmt.fmt)
def test_linspace_bug(self):
"""This catches a linspace datetime bug with 0-d arrays (regression)"""
try:
from matplotlib.dates import date2num
except ImportError: # just pass if matplotlib is not installed
return
t1 = numpy.array(datetime.datetime(2000, 1, 1))
t2 = numpy.array(datetime.datetime(2000, 1, 10))
real_ans = [datetime.datetime(2000, 1, 1, 0, 0),
datetime.datetime(2000, 1, 3, 6, 0),
datetime.datetime(2000, 1, 5, 12, 0),
datetime.datetime(2000, 1, 7, 18, 0),
datetime.datetime(2000, 1, 10, 0, 0)]
ans = tb.linspace(t1, t2, 5)
numpy.testing.assert_almost_equal(date2num(real_ans), date2num(ans) , 4)
def test_linspace_bug(self):
"""This catches a linspace datetime bug with 0-d arrays (regression)"""
try:
from matplotlib.dates import date2num
except ImportError: # just pass if matplotlib is not installed
return
t1 = numpy.array(datetime.datetime(2000, 1, 1))
t2 = numpy.array(datetime.datetime(2000, 1, 10))
real_ans = [datetime.datetime(2000, 1, 1, 0, 0),
datetime.datetime(2000, 1, 3, 6, 0),
datetime.datetime(2000, 1, 5, 12, 0),
datetime.datetime(2000, 1, 7, 18, 0),
datetime.datetime(2000, 1, 10, 0, 0)]
ans = tb.linspace(t1, t2, 5)
numpy.testing.assert_almost_equal(date2num(real_ans), date2num(ans) , 4)
def __init__(self, data, **kwargs):
"""
"""
super(spectrogram, self).__init__()
## setup a default dictionary to step through to set values from kwargs
self.specSettings = dm.SpaceData()
self.specSettings['variables'] = ['Epoch', 'Energy', 'Flux']
self.specSettings['bins'] = None # this is the linspace over the range with sqrt() of the len bins
self.specSettings['xlim'] = None
self.specSettings['ylim'] = None
self.specSettings['zlim'] = None
self.specSettings['extended_out'] = True
self.specSettings['axisDates'] = False
# if the key exists in kwargs replace setting with it, otherwise its an error
for key in kwargs:
if key not in self.specSettings:
raise(KeyError('Invalid keyword specified ' + str(key)))
self.specSettings[key] = kwargs[key]
# check to see if the variables are in the spacedata
for var in self.specSettings['variables']:
if not var in data: # TODO could check other capitalization
raise(KeyError('"{0}" not found in the input data'.format(var) ))
# if the variables are empty error and quit
if len(data[self.specSettings['variables'][0]]) == 0:
raise(ValueError('No {0} data passed in'.format(self.specSettings['variables'][0])))
if len(data[self.specSettings['variables'][1]]) == 0:
raise(ValueError('No {0} data passed in'.format(self.specSettings['variables'][1])))
if len(data[self.specSettings['variables'][2]]) == 0:
raise(ValueError('No {0} data passed in'.format(self.specSettings['variables'][2])))
# set limits, keywords override those in the data
#if (self.specSettings['xlim'] is None) and (self.specSettings['bins'] is None):
if self.specSettings['xlim'] is None:
try:
if 'lim' in data[self.specSettings['variables'][0]].attrs:
self.specSettings['xlim'] = (data[self.specSettings['variables'][0]].attrs['lim'][0],
data[self.specSettings['variables'][0]].attrs['lim'][1])
else:
dum1 = np.min(data[self.specSettings['variables'][0]]).tolist() #TODO: tolist here is a workaround for a bug in
dum2 = np.max(data[self.specSettings['variables'][0]]).tolist() # datamodel's min method
self.specSettings['xlim'] = (dum1, dum2)
except AttributeError: # was a numpy array not dmarray
self.specSettings['xlim'] = (np.min(data[self.specSettings['variables'][0]]),
np.max(data[self.specSettings['variables'][0]]))
if self.specSettings['ylim'] is None:
try:
if 'lim' in data[self.specSettings['variables'][1]].attrs:
self.specSettings['ylim'] = (data[self.specSettings['variables'][1]].attrs['lim'][0],
data[self.specSettings['variables'][1]].attrs['lim'][1])
else:
self.specSettings['ylim'] = (np.min(data[self.specSettings['variables'][1]]),
np.max(data[self.specSettings['variables'][1]]))
except AttributeError: # was a numpy array not dmarray
self.specSettings['ylim'] = (np.min(data[self.specSettings['variables'][1]]),
np.max(data[self.specSettings['variables'][1]]))
if self.specSettings['zlim'] is None:
try:
if 'lim' in data[self.specSettings['variables'][2]].attrs:
self.specSettings['zlim'] = (data[self.specSettings['variables'][2]].attrs['lim'][0],
data[self.specSettings['variables'][2]].attrs['lim'][1])
else:
self.specSettings['zlim'] = (np.min(data[self.specSettings['variables'][2]]),
np.max(data[self.specSettings['variables'][2]]))
except AttributeError: # was a numpy array not dmarray
self.specSettings['zlim'] = (np.min(data[self.specSettings['variables'][2]]),
np.max(data[self.specSettings['variables'][2]]))
# are the axes dates?
forcedate = [False] * 2
if isinstance(data[self.specSettings['variables'][0]][0], datetime.datetime):
forcedate[0] = True
if isinstance(data[self.specSettings['variables'][1]][0], datetime.datetime):
forcedate[1] = True
self.specSettings['axisDates'] = forcedate
# set default bins
if self.specSettings['bins'] is None:
# since it is not set by keyword was it set in the datamodel?
attr_bins = ['bins' in data[var].attrs for var in self.specSettings['variables']]
if dm.dmarray(attr_bins[0:2]).all():
self.specSettings['bins'] = [dm.dmarray(data[self.specSettings['variables'][0]].attrs['bins']),
dm.dmarray(data[self.specSettings['variables'][1]].attrs['bins']),]
# TODO this is not a hard extension to doing one with bins and one default
else:
# use the toolbox version of linspace so it works on dates
self.specSettings['bins'] = [dm.dmarray(tb.linspace(self.specSettings['xlim'][0],
self.specSettings['xlim'][1],
np.sqrt(len(data[self.specSettings['variables'][0]])))),
dm.dmarray(tb.linspace(self.specSettings['ylim'][0],
self.specSettings['ylim'][1],
np.sqrt(len(data[self.specSettings['variables'][1]]))))]
# copy all the used keys
for key in self.specSettings['variables']:
self[key] = data[key]
try:
self[key].attrs = data[key].attrs
except:
pass
# do the spectrogram
self._computeSpec()
def __init__(self, data, **kwargs):
"""
"""
super(spectrogram, self).__init__()
## setup a default dictionary to step through to set values from kwargs
self.specSettings = dm.SpaceData()
self.specSettings['variables'] = ['Epoch', 'Energy', 'Flux']
self.specSettings[
'bins'] = None # this is the linspace over the range with sqrt() of the len bins
self.specSettings['xlim'] = None
self.specSettings['ylim'] = None
self.specSettings['zlim'] = None
self.specSettings['extended_out'] = True
self.specSettings['axisDates'] = False
# if the key exists in kwargs replace setting with it, otherwise its an error
for key in kwargs:
if key not in self.specSettings:
raise (KeyError('Invalid keyword specified ' + str(key)))
self.specSettings[key] = kwargs[key]
# check to see if the variables are in the spacedata
for var in self.specSettings['variables']:
if not var in data: # TODO could check other capitalization
raise (KeyError(
'"{0}" not found in the input data'.format(var)))
# if the variables are empty error and quit
if len(data[self.specSettings['variables'][0]]) == 0:
raise (ValueError('No {0} data passed in'.format(
self.specSettings['variables'][0])))
if len(data[self.specSettings['variables'][1]]) == 0:
raise (ValueError('No {0} data passed in'.format(
self.specSettings['variables'][1])))
if len(data[self.specSettings['variables'][2]]) == 0:
raise (ValueError('No {0} data passed in'.format(
self.specSettings['variables'][2])))
# set limits, keywords override those in the data
#if (self.specSettings['xlim'] is None) and (self.specSettings['bins'] is None):
if self.specSettings['xlim'] is None:
try:
if 'lim' in data[self.specSettings['variables'][0]].attrs:
self.specSettings['xlim'] = (
data[self.specSettings['variables']
[0]].attrs['lim'][0],
data[self.specSettings['variables'][0]].attrs['lim'][1]
)
else:
dum1 = np.min(
data[self.specSettings['variables'][0]]).tolist(
) #TODO: tolist here is a workaround for a bug in
dum2 = np.max(
data[self.specSettings['variables']
[0]]).tolist() # datamodel's min method
self.specSettings['xlim'] = (dum1, dum2)
except AttributeError: # was a numpy array not dmarray
self.specSettings['xlim'] = (
np.min(data[self.specSettings['variables'][0]]),
np.max(data[self.specSettings['variables'][0]]))
if self.specSettings['ylim'] is None:
try:
if 'lim' in data[self.specSettings['variables'][1]].attrs:
self.specSettings['ylim'] = (
data[self.specSettings['variables']
[1]].attrs['lim'][0],
data[self.specSettings['variables'][1]].attrs['lim'][1]
)
else:
self.specSettings['ylim'] = (
np.min(data[self.specSettings['variables'][1]]),
np.max(data[self.specSettings['variables'][1]]))
except AttributeError: # was a numpy array not dmarray
self.specSettings['ylim'] = (
np.min(data[self.specSettings['variables'][1]]),
np.max(data[self.specSettings['variables'][1]]))
if self.specSettings['zlim'] is None:
try:
if 'lim' in data[self.specSettings['variables'][2]].attrs:
self.specSettings['zlim'] = (
data[self.specSettings['variables']
[2]].attrs['lim'][0],
data[self.specSettings['variables'][2]].attrs['lim'][1]
)
else:
self.specSettings['zlim'] = (
np.min(data[self.specSettings['variables'][2]]),
np.max(data[self.specSettings['variables'][2]]))
except AttributeError: # was a numpy array not dmarray
self.specSettings['zlim'] = (
np.min(data[self.specSettings['variables'][2]]),
np.max(data[self.specSettings['variables'][2]]))
# are the axes dates?
forcedate = [False] * 2
if isinstance(data[self.specSettings['variables'][0]][0],
datetime.datetime):
forcedate[0] = True
if isinstance(data[self.specSettings['variables'][1]][0],
datetime.datetime):
forcedate[1] = True
self.specSettings['axisDates'] = forcedate
# set default bins
if self.specSettings['bins'] is None:
# since it is not set by keyword was it set in the datamodel?
attr_bins = [
'bins' in data[var].attrs
for var in self.specSettings['variables']
]
if dm.dmarray(attr_bins[0:2]).all():
self.specSettings['bins'] = [
dm.dmarray(
data[self.specSettings['variables'][0]].attrs['bins']),
dm.dmarray(
data[self.specSettings['variables'][1]].attrs['bins']),
]
# TODO this is not a hard extension to doing one with bins and one default
else:
# use the toolbox version of linspace so it works on dates
self.specSettings['bins'] = [
dm.dmarray(
tb.linspace(
self.specSettings['xlim'][0],
self.specSettings['xlim'][1],
int(
np.sqrt(
len(data[self.specSettings['variables']
[0]]))))),
dm.dmarray(
tb.linspace(
self.specSettings['ylim'][0],
self.specSettings['ylim'][1],
int(
np.sqrt(
len(data[self.specSettings['variables']
[1]])))))
]
# copy all the used keys
for key in self.specSettings['variables']:
self[key] = data[key]
try:
self[key].attrs = data[key].attrs
except:
pass
# do the spectrogram
self._computeSpec()
def setUp(self):
super(PAmodelTests, self).setUp()
self.PA = tb.linspace(0.01,179.99,20000)
def vampolaPA(omniflux, **kwargs):
'''Pitch angle model of sin^n form
Parameters
==========
omniflux : arraylike or float
omnidirectional number flux data
order : integer or float (optional)
order of sin^n functional form for distribution (default=2)
alphas : arraylike (optional)
pitch angles at which to evaluate the differential number flux
(default is 5 to 90 degrees in 36 steps)
Returns
=======
dnflux : array
differential number flux corresponding to pitch angles alphas
alphas : array
pitch angles at which the differential number flux was evaluated
Examples
========
Omnidirectional number flux of [3000, 6000]
>>> from spacepy.empiricals import vampolaPA
>>> vampolaPA(3000, alpha=[45, 90])
(array([ 179.04931098, 358.09862196]), [45, 90])
>>> data, pas = vampolaPA([3000, 6000], alpha=[45, 90])
>>> pas
[45, 90]
>>> data
array([[ 179.04931098, 358.09862196],
[ 358.09862196, 716.19724391]])
Notes
=====
Directional number flux integrated over pitch angle from 0 to 90 degrees
is a factor of 4*pi lower than omnidirectional number flux.
'''
defaults = {'order': 2, 'alpha': tb.linspace(5, 90, 18)}
if hasattr(omniflux, '__iter__'):
omniflux = np.asanyarray(omniflux)
else:
omniflux = np.asanyarray([omniflux])
#substitute defaults
for key in defaults:
if key not in kwargs:
kwargs[key] = defaults[key]
if hasattr(kwargs['order'], '__iter__'):
try:
assert len(kwargs['order']) == len(omniflux)
except AssertionError:
raise ValueError(
'order must be either single-valued or the same length as omniflux'
)
else:
kwargs['order'] = np.asanyarray([kwargs['order']] * len(omniflux))
normfac = np.empty(len(kwargs['order']), dtype=float)
def sinfunc(x, order=kwargs['order']): #define distribution function
dum = np.sin(x)
return dum**order
for idx, tmporder in enumerate(kwargs['order']):
sinfunc_o = partial(sinfunc, order=tmporder + 1)
normfac[idx] = omniflux[idx] / (2 * np.pi *
integ.quad(sinfunc_o, 0, np.pi)[0])
#now make the differential number flux
dnflux = np.zeros((len(kwargs['alpha']), len(omniflux))).squeeze()
for i, a_val in enumerate(np.deg2rad(kwargs['alpha'])):
dnflux[i] = normfac * sinfunc(a_val)
return dnflux, kwargs['alpha']
def vampolaPA(omniflux, **kwargs):
'''Pitch angle model of sin^n form
Parameters
==========
omniflux : arraylike or float
omnidirectional number flux data
order : integer or float (optional)
order of sin^n functional form for distribution (default=2)
alphas : arraylike (optional)
pitch angles at which to evaluate the differential number flux
(default is 5 to 90 degrees in 36 steps)
Returns
=======
dnflux : array
differential number flux corresponding to pitch angles alphas
alphas : array
pitch angles at which the differential number flux was evaluated
Examples
========
Omnidirectional number flux of [3000, 6000]
>>> from spacepy.empiricals import vampolaPA
>>> vampolaPA(3000, alpha=[45, 90])
(array([ 75.99088773, 151.98177546]), [45, 90])
>>> data, pas = vampolaPA([3000, 6000], alpha=[45, 90])
>>> pas
[45, 90]
>>> data
array([[ 75.99088773, 151.98177546],
[ 151.98177546, 303.96355093]])
Notes
=====
Directional number flux integrated over pitch angle from 0 to 90 degrees
is a factor of 4*pi lower than omnidirectional number flux.
'''
defaults = {'order': 2,
'alpha': tb.linspace(5,90,18)}
if hasattr(omniflux, '__iter__'):
omniflux = np.asanyarray(omniflux)
else:
omniflux = np.asanyarray([omniflux])
#substitute defaults
for key in defaults:
if key not in kwargs:
kwargs[key] = defaults[key]
if hasattr(kwargs['order'], '__iter__'):
try:
assert len(kwargs['order'])==len(omniflux)
except AssertionError:
raise ValueError('order must be either single-valued or the same length as omniflux')
else:
kwargs['order'] = np.asanyarray([kwargs['order']]*len(omniflux))
normfac = np.empty(len(kwargs['order']), dtype=float)
def sinfunc(x, order=kwargs['order']): #define distribution function
dum = np.sin(x)
return dum**order
for idx, tmporder in enumerate(kwargs['order']):
#def partial function so that order is fixed
sinfunc_o = partial(sinfunc, order=tmporder)
normfac[idx] = 4*np.pi*integ.quad(sinfunc_o, 0, np.pi)[0] #quad returns (val, accuracy)
#now make the differential number flux
dnflux = np.zeros((len(kwargs['alpha']), len(omniflux))).squeeze()
for i, a_val in enumerate(np.deg2rad(kwargs['alpha'])):
dnflux[i] = omniflux * sinfunc(a_val) / normfac
return dnflux, kwargs['alpha']
