def test_applySmartTimeTicks(self):
"""applySmartTimeTicks should have known behaviour"""
plt.ion()
ticks = st.tickrange('2002-02-01T00:00:00', '2002-02-07T00:00:00', deltadays=1)
y = list(range(len(ticks)))
fig = plt.figure()
ax = fig.add_subplot(111)
line = ax.plot(ticks.UTC, y)
spacepy.plot.utils.applySmartTimeTicks(ax, ticks.UTC)
plt.draw()
plt.draw()
# should not have moved the ticks
real_ans = numpy.array([ 730882., 730883., 730884., 730885., 730886., 730887.,
730888.])
numpy.testing.assert_almost_equal(real_ans, ax.get_xticks())
# should have named them 01 Feb, 02 Feb etc
try:
real_ans = ['{0:02d} Feb'.format(i+1).decode() for i in range(7)]
except AttributeError: #Py3k
real_ans = ['{0:02d} Feb'.format(i+1) for i in range(7)]
ans = [t.get_text()
for t in ax.xaxis.get_majorticklabels()]
numpy.testing.assert_array_equal(real_ans, ans)
plt.close()
plt.ioff()
def runQs():
#make figure
fig = plt.figure(figsize=(17, 21))
#set up test runs
nvals, rvals = 5, [2, 3, 4, 5, 6]
dates = spt.tickrange('2001-04-22T12:00:00', '2001-04-23T23:00:00', 1/24)
alpha = 45
#loop over radial positions, dates and quality flags
for rval in rvals:
Lstar = [[]]*nvals
for date, qual in itertools.product(dates.UTC, range(nvals)):
pos = dm.dmarray([-1*rval, 0, 0], attrs={'sys': 'SM'})
data = lgmpy.get_Lstar(pos, date, alpha=alpha, coord_system='SM', Bfield='Lgm_B_cdip', LstarThresh=20.0, extended_out=True, LstarQuality=qual)
try:
Lstar[qual].extend(data[alpha]['Lstar'])
except TypeError:
Lstar[qual].append(data[alpha]['Lstar'].tolist())
print('Did [-%d,0,0] for all qualities' % rval)
#make plots
fstr = '%d1%d' % (len(rvals), rval-rvals[0]+1)
ax = fig.add_subplot(fstr)
ax.boxplot(Lstar)
#ax = plt.gca()
ax.set_title('LGM - Centred dipole [-%d, 0, 0]$_{SM}$; PA=%d$^{o}$' % (rval, alpha))
ax.set_ylabel('L* (LanlGeoMag)')
ax.set_xlabel('Quality Flag')
ax.set_xticklabels([str(n) for n in range(5)])
tb.savepickle('lgm_cdip_lstar%d_alpha%d.pkl' % (rval, alpha), {'Lstar': Lstar})
plt.ylim([rval-0.04, rval+0.04])
plt.savefig('lgm_cdip_verify%d.png' % alpha, dpi=300)
def runPAdungey(quality):
#test for range of pitch angles
ticks = spt.tickrange('2002-04-18', '2002-04-19', 1)
loci = spc.Coords([[-4, 0, 0], [-5, 0, 0]], 'SM', 'car')
vals = []
for pp in range(1, 91):
dum = lgmpy.get_Lstar(loci.data[1],
ticks.UTC[1],
alpha=pp,
coord_system='SM',
Bfield='Lgm_B_Dungey',
extended_out=True,
LstarQuality=quality)
try:
vals.extend(dum[pp]['Lstar'])
except (IndexError,
TypeError): #get_Lstar returns a 0-D nan in some cases...
vals.extend([dum[pp]['Lstar'].tolist()])
fig = plt.figure()
expect = dungeyLeq(tb.hypot(loci.data[1]))
#print(expect)
plt.plot(100 * (dm.dmarray(vals) - expect) / expect, drawstyle='steps-mid')
plt.ylim([-0.01, 0.01])
plt.xlabel('Pitch Angle')
plt.ylabel('100*(L* - L*$_{exp}$)/L*$_{exp}$')
plt.title('Dungey model [-5,0,0]$_{SM}$' +
'; Quality ({0})'.format(str(quality)))
figname = 'MagEphem_Dungey_test_q{0}'.format(
str(quality)) #'LGM_L*_vs_PA_cdip_q{0}'.format(str(quality))
plt.savefig(''.join([figname, '.png']), dpi=300)
plt.savefig(''.join([figname, '.pdf']))
def test_applySmartTimeTicks(self):
"""applySmartTimeTicks should have known behaviour"""
plt.ion()
ticks = st.tickrange('2002-02-01T00:00:00',
'2002-02-07T00:00:00',
deltadays=1)
y = list(range(len(ticks)))
fig = plt.figure()
ax = fig.add_subplot(111)
line = ax.plot(ticks.UTC, y)
spacepy.plot.utils.applySmartTimeTicks(ax, ticks.UTC)
plt.draw()
plt.draw()
# should not have moved the ticks
real_ans = numpy.array(
[730882., 730883., 730884., 730885., 730886., 730887., 730888.])
numpy.testing.assert_almost_equal(real_ans, ax.get_xticks())
# should have named them 01 Feb, 02 Feb etc
try:
real_ans = ['{0:02d} Feb'.format(i + 1).decode() for i in range(7)]
except AttributeError: #Py3k
real_ans = ['{0:02d} Feb'.format(i + 1) for i in range(7)]
ans = [t.get_text() for t in ax.xaxis.get_majorticklabels()]
numpy.testing.assert_array_equal(real_ans, ans)
plt.close()
plt.ioff()
def test_tickrange2(self):
"""tickrange should return a known value for known input (timedelta)"""
inval = (
("2002-02-01T00:00:00", "2002-02-04T00:00:00", datetime.timedelta(days=1)),
("2002-02-01T00:00:00", "2002-02-04T00:00:00", datetime.timedelta(hours=12)),
)
strarray_dtype = numpy.array("x" * 19).dtype
real_ans = (
numpy.array(
["2002-02-01T00:00:00", "2002-02-02T00:00:00", "2002-02-03T00:00:00", "2002-02-04T00:00:00"],
dtype=strarray_dtype,
),
numpy.array(
[
"2002-02-01T00:00:00",
"2002-02-01T12:00:00",
"2002-02-02T00:00:00",
"2002-02-02T12:00:00",
"2002-02-03T00:00:00",
"2002-02-03T12:00:00",
"2002-02-04T00:00:00",
],
dtype=strarray_dtype,
),
)
for i, val in enumerate(inval):
ans = t.tickrange(*val)
numpy.testing.assert_equal(real_ans[i], ans.ISO)
def test_dateToISOunaltered_dm(self):
"""Test to check that _dateToISO doesn't change datatypes, input dmarray"""
data = spt.tickrange('20200101', '20200102',
deltadays=datetime.timedelta(hours=1)).UTC
exptype = type(data[0])
newdata = dm._dateToISO(data)
restype = type(data[0])
self.assertEqual(exptype, restype)
def test_tickrange(self):
"""tickrange should return a known value for known input"""
inval = (('2002-02-01T00:00:00', '2002-02-04T00:00:00', 1),
('2002-02-01T00:00:00', '2002-02-04T00:00:00', 0.5))
strarray_dtype = numpy.array('x' * 19).dtype
real_ans = (numpy.array(['2002-02-01T00:00:00', '2002-02-02T00:00:00', '2002-02-03T00:00:00',
'2002-02-04T00:00:00'], dtype=strarray_dtype),
numpy.array(['2002-02-01T00:00:00', '2002-02-01T12:00:00', '2002-02-02T00:00:00',
'2002-02-02T12:00:00', '2002-02-03T00:00:00', '2002-02-03T12:00:00',
'2002-02-04T00:00:00'], dtype=strarray_dtype))
for i, val in enumerate(inval):
ans = t.tickrange(*val)
numpy.testing.assert_equal(real_ans[i], ans.ISO)
def setup_ticks(self, start, end, delta, dtype='ISO'):
"""
Add time information to the simulation by specifying
a start and end time, timestep, and time type (optional).
Examples
========
>>> start = datetime.datetime(2003,10,14)
>>> end = datetime.datetime(2003,12,26)
>>> delta = datetime.timedelta(hours=1)
>>> rmod.setup_ticks(start, end, delta, dtype='UTC')
"""
self.ticks = st.tickrange(start, end, delta, dtype)
def test_tickrange2(self):
"""tickrange should return a known value for known input (timedelta)"""
inval = (('2002-02-01T00:00:00', '2002-02-04T00:00:00', datetime.timedelta(days=1)),
('2002-02-01T00:00:00', '2002-02-04T00:00:00', datetime.timedelta(hours=12)))
strarray_dtype = numpy.array('x' * 19).dtype
real_ans = (numpy.array(['2002-02-01T00:00:00', '2002-02-02T00:00:00', '2002-02-03T00:00:00',
'2002-02-04T00:00:00'], dtype=strarray_dtype),
numpy.array(['2002-02-01T00:00:00', '2002-02-01T12:00:00', '2002-02-02T00:00:00',
'2002-02-02T12:00:00', '2002-02-03T00:00:00', '2002-02-03T12:00:00',
'2002-02-04T00:00:00'], dtype=strarray_dtype))
for i, val in enumerate(inval):
ans = t.tickrange(*val)
numpy.testing.assert_equal(real_ans[i], ans.ISO)
def setup_ticks(self, start, end, delta, dtype='ISO'):
"""
Add time information to the simulation by specifying
a start and end time, timestep, and time type (optional).
Examples
========
>>> start = datetime.datetime(2003,10,14)
>>> end = datetime.datetime(2003,12,26)
>>> delta = datetime.timedelta(hours=1)
>>> rmod.setup_ticks(start, end, delta, dtype='UTC')
"""
self.ticks = st.tickrange(start, end, delta, dtype)
def solarRotationPlot(ticks, data, targ_ax=None, rtype='bartels', nbins=27):
'''Plots a 1-D time series as a Carrington or Bartels plot
'''
SRlength = {'bartels': 27.0, 'carrington': 27.2753}
if not targ_ax:
#setup axes
fig, targ_ax = plt.subplots()
#data are 1-D, get range of SR numbers for Y-axis
min_sr = emp.getSolarRotation(ticks[0], rtype=rtype)
max_sr = emp.getSolarRotation(ticks[-1], rtype=rtype)
n_sr = 1 + max_sr - min_sr
sr_values = np.linspace(min_sr, max_sr, n_sr)
#get bin size in time
start_time = emp.getSolarRotation(
min_sr, rtype=rtype, reverse=True
)[0] #convert back from Solar Rotation to date, so we get the start time of the rotation
end_time = emp.getSolarRotation(max_sr + 1, rtype=rtype, reverse=True)[
0] #add 1 to last solar rotation number then convert back to date
binned_times = spt.tickrange(start_time, end_time,
SRlength[rtype.lower()] / nbins)
#now bin data on new time grid -- TODO: use windowMean from toolbox??
digital = np.digitize(ticks.RDT, binned_times.RDT)
bin_means = np.zeros(len(binned_times) - 1)
bin_means.fill(np.nan)
for i in range(1, len(binned_times)):
try:
bin_means[i] = data[digital == i].mean()
except (
ZeroDivisionError, IndexError
): #ZDE happens when no data for that bin, IE is when bins extend beyond data
pass
plot_view = np.array(bin_means).reshape((n_sr, nbins))
im = targ_ax.imshow(plot_view,
interpolation='nearest',
vmin=data.min(),
vmax=data.max(),
extent=[1, nbins, max_sr, min_sr],
aspect=100)
plt.colorbar(im, ax=targ_ax)
targ_ax.set_xlabel('Day of Rotation')
targ_ax.set_ylabel('{0} Rotation number'.format(rtype.title()))
#targ_ax.pcolormesh(plot_view)
return
def test_toJSON_timeunaltered(self):
"""Test to check that stored datetimes aren't changed on write"""
data = dm.SpaceData()
data['Epoch'] = spt.tickrange('20200101', '20200102',
deltadays=datetime.timedelta(hours=1)).UTC
exptype = type(data['Epoch'][0]) #datetime.datetime
# save to file, then immediately clean up
fname = None
try:
with tempfile.NamedTemporaryFile(delete=False) as fp:
fname = fp.name
data.toJSONheadedASCII(fname)
finally:
if fname != None:
os.remove(fname)
restype = type(data['Epoch'][0])
self.assertEqual(exptype, restype)
def runQs():
#make figure
fig = plt.figure(figsize=(17, 21))
#set up test runs
nvals, rvals = 5, [2, 3, 4, 5, 6]
dates = spt.tickrange('2001-04-22T12:00:00', '2001-04-23T23:00:00', 1 / 24)
alpha = 45
#loop over radial positions, dates and quality flags
for rval in rvals:
Lstar = [[]] * nvals
for date, qual in itertools.product(dates.UTC, range(nvals)):
pos = dm.dmarray([-1 * rval, 0, 0], attrs={'sys': 'SM'})
data = lgmpy.get_Lstar(pos,
date,
alpha=alpha,
coord_system='SM',
Bfield='Lgm_B_cdip',
LstarThresh=20.0,
extended_out=True,
LstarQuality=qual)
try:
Lstar[qual].extend(data[alpha]['Lstar'])
except TypeError:
Lstar[qual].append(data[alpha]['Lstar'].tolist())
print('Did [-%d,0,0] for all qualities' % rval)
#make plots
fstr = '%d1%d' % (len(rvals), rval - rvals[0] + 1)
ax = fig.add_subplot(fstr)
ax.boxplot(Lstar)
#ax = plt.gca()
ax.set_title('LGM - Centred dipole [-%d, 0, 0]$_{SM}$; PA=%d$^{o}$' %
(rval, alpha))
ax.set_ylabel('L* (LanlGeoMag)')
ax.set_xlabel('Quality Flag')
ax.set_xticklabels([str(n) for n in range(5)])
tb.savepickle('lgm_cdip_lstar%d_alpha%d.pkl' % (rval, alpha),
{'Lstar': Lstar})
plt.ylim([rval - 0.04, rval + 0.04])
plt.savefig('lgm_cdip_verify%d.png' % alpha, dpi=300)
def runPA(quality):
#test for range of pitch angles
ticks = spt.tickrange('2002-04-18', '2002-04-19', 1)
loci = spc.Coords([[-4,0,0], [-5,0,0]], 'SM', 'car')
vals = []
for pp in range(1,91):
dum = lgmpy.get_Lstar(loci.data[1], ticks.UTC[1], alpha=pp, coord_system='SM',
Bfield='Lgm_B_cdip', extended_out=True, LstarQuality=quality)
try:
vals.extend(dum[pp]['Lstar'])
except (IndexError, TypeError): #get_Lstar returns a 0-D nan in some cases...
vals.extend([dum[pp]['Lstar'].tolist()])
fig = plt.figure()
expect = float(tb.hypot(loci.data[1]))
plt.plot(100*(dm.dmarray(vals)-expect)/expect, drawstyle='steps-mid')
plt.ylim([-0.01, 0.01])
plt.xlabel('Pitch Angle')
plt.ylabel('100*(L* - L*$_{exp}$)/L*$_{exp}$')
plt.title('Cent. dipole [-5,0,0]$_{SM}$'+'; Quality ({0})'.format(str(quality)))
figname = 'MagEphem_CDIP_test_q{0}'.format(str(quality)) #'LGM_L*_vs_PA_cdip_q{0}'.format(str(quality))
plt.savefig(''.join([figname,'.png']), dpi=300)
plt.savefig(''.join([figname,'.pdf']))
def solarRotationPlot(ticks, data, targ_ax=None, rtype='bartels', nbins=27):
'''Plots a 1-D time series as a Carrington or Bartels plot
'''
SRlength = {'bartels': 27.0, 'carrington':27.2753}
if not targ_ax:
#setup axes
fig, targ_ax = plt.subplots()
#data are 1-D, get range of SR numbers for Y-axis
min_sr = emp.getSolarRotation(ticks[0], rtype=rtype)
max_sr = emp.getSolarRotation(ticks[-1], rtype=rtype)
n_sr = 1+max_sr-min_sr
sr_values = np.linspace(min_sr, max_sr, n_sr)
#get bin size in time
start_time = emp.getSolarRotation(min_sr, rtype=rtype, reverse=True)[0] #convert back from Solar Rotation to date, so we get the start time of the rotation
end_time = emp.getSolarRotation(max_sr+1, rtype=rtype, reverse=True)[0] #add 1 to last solar rotation number then convert back to date
binned_times = spt.tickrange(start_time, end_time, SRlength[rtype.lower()]/nbins)
#now bin data on new time grid -- TODO: use windowMean from toolbox??
digital = np.digitize(ticks.RDT, binned_times.RDT)
bin_means = np.zeros(len(binned_times)-1)
bin_means.fill(np.nan)
for i in range(1, len(binned_times)):
try:
bin_means[i] = data[digital == i].mean()
except (ZeroDivisionError, IndexError): #ZDE happens when no data for that bin, IE is when bins extend beyond data
pass
plot_view = np.array(bin_means).reshape((n_sr, nbins))
im = targ_ax.imshow(plot_view, interpolation='nearest', vmin=data.min(), vmax=data.max(), extent=[1,nbins, max_sr,min_sr], aspect=100)
plt.colorbar(im, ax=targ_ax)
targ_ax.set_xlabel('Day of Rotation')
targ_ax.set_ylabel('{0} Rotation number'.format(rtype.title()))
#targ_ax.pcolormesh(plot_view)
return
def test_applySmartTimeTicks(self):
"""applySmartTimeTicks should have known behaviour"""
ticks = st.tickrange('2002-02-01T00:00:00',
'2002-02-07T00:00:00',
deltadays=1)
y = list(range(len(ticks)))
fig = plt.figure()
ax = fig.add_subplot(111)
line = ax.plot(ticks.UTC, y)
spacepy.plot.utils.applySmartTimeTicks(ax, ticks.UTC)
fp = sio()
fig.savefig(fp, format='png')
fp.close()
# should not have moved the ticks
real_ans = matplotlib.dates.date2num(
[datetime.datetime(2002, 2, i) for i in range(1, 8)])
numpy.testing.assert_almost_equal(real_ans, ax.get_xticks())
# should have named them 01 Feb, 02 Feb etc
try:
real_ans = ['{0:02d} Feb'.format(i + 1).decode() for i in range(7)]
except AttributeError: #Py3k
real_ans = ['{0:02d} Feb'.format(i + 1) for i in range(7)]
ans = [t.get_text() for t in ax.xaxis.get_majorticklabels()]
numpy.testing.assert_array_equal(real_ans, ans)
def plotSpectrogram(self, ecol=0, pvars=None, **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)
pvars : list
list of plotting variable names in order [Epoch-like (X axis), Flux-like (Z axis), Energy (Index var for Flux-like)]
ylim : list
2-element list with upper and lower bounds for y axis
'''
import spacepy.plot as splot
if 'Lm' not in self:
self.getLm()
sd = dm.SpaceData()
if pvars is None:
varname = self.attrs['varname']
enname = 'Energy'
epvar = 'Epoch'
else:
epvar = pvars[0]
varname = pvars[1]
enname = pvars[2]
if len(self['Lm']) != len(self[epvar]): #Lm needs interpolating to new timebase
import matplotlib.dates as mpd
tmptimetarg, tmptimesrc = mpd.date2num(self[epvar]), mpd.date2num(self['Epoch'])
sd['Lm'] = dm.dmarray(np.interp(tmptimetarg, tmptimesrc, self['Lm'], left=np.nan, right=np.nan))
else:
sd['Lm'] = self['Lm']
#filter any bad Lm
goodidx = sd['Lm']>1
sd['Lm'] = sd['Lm'][goodidx]
if 'ylim' in kwargs:
Lm_lim = kwargs['ylim']
del kwargs['ylim']
else:
Lm_lim = [2.0,8.0]
#TODO: allow user-definition of bins in time and Lm
sd['Epoch'] = dm.dmcopy(self[epvar])[goodidx] #TODO: assumes 1 pitch angle, generalize
try:
sd['1D_dataset'] = self[varname][goodidx,ecol] #TODO: assumes 1 pitch angle, generalize
except IndexError: #1-D
sd['1D_dataset'] = self[varname][goodidx]
bins = [[],[]]
if 'tbins' not in kwargs:
bins[0] = spt.tickrange(self[epvar][0], self[epvar][-1], 3/24.).UTC
else:
bins[0] = kwargs['tbins']
del kwargs['tbins']
if 'Lbins' not in kwargs:
bins[1] = np.arange(Lm_lim[0], Lm_lim[1], 1./4.)
else:
bins[1] = kwargs['Lbins']
del kwargs['Lbins']
spec = splot.spectrogram(sd, variables=['Epoch', 'Lm', '1D_dataset'], ylim=Lm_lim, bins=bins)
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:5.2f} {1}'.format(self[enname][ecol], self[enname].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
import numpy as np
import spacepy.LANLstar as LS
import spacepy.time as spt
import spacepy.omni as om
import datetime as datetime
st = datetime.datetime(2013, 3, 23)
en = datetime.datetime(2013, 3, 24)
delta = datetime.timedelta(minutes=5)
ticks = spt.tickrange(st, en, delta, 'UTC')
data = om.get_omni(ticks)
omni_t = data['ticks']
import numpy as np
from spacepy import pycdf
# cdf = pycdf.CDF('/Users/yangjian/Desktop/rbsp-a_magnetometer_4sec-gsm_emfisis-l3_20130323_v1.3.3.cdf')
# pos_data=cdf['coordinates'][...]
# epoch=cdf['Epoch'][...]
#
# pos_data_new=np.interp(omni_t,epoch,pos_data)
inputdict = {}
inputdict['Kp'] = data['Kp'] # Kp index
inputdict['Dst'] = data['Dst'] # Dst index (nT)
inputdict['dens'] = data['dens'] # solar wind density (/cc)
inputdict['velo'] = data['velo'] # solar wind velocity (km/s)
inputdict['Pdyn'] = data['Pdyn'] # solar wind dynamic pressure (nPa)
inputdict['ByIMF'] = data[
'ByIMF'] # GSM y component of IMF magnetic field (nT)
def setUp(self):
super(empFunctionTests, self).setUp()
self.ticks = spt.tickrange('2001-01-01T12:00:00','2001-01-04T00:00:00',.25)
self.omnivals = om.get_omni(self.ticks, dbase='Test')
from __future__ import division
import matplotlib.pyplot as plt
import itertools
import datetime as dt
import lgmpy, copy
import spacepy.time as spt
import spacepy.datamodel as dm
import spacepy.toolbox as tb
# make figure
fig = plt.figure(figsize=(17, 21))
# set up test runs
nvals, rvals = 5, [3, 4, 5, 6]
dates = spt.tickrange("2001-04-22T12:00:00", "2001-04-22T12:01:00", dt.timedelta(seconds=5))
alpha = 45
print("T89: alpha=%s, date range [%s, %s]" % (alpha, dates.UTC[0], dates.UTC[-1]))
# loop over radial positions, dates and quality flags
for rval in rvals:
print("Radial Distance: %s" % rval)
Lstar = [[]] * nvals
for date, qual in itertools.product(dates.UTC, range(nvals)):
print("%s, Quality=%d" % (date, qual))
pos = dm.dmarray([-1 * rval, 0, 0], attrs={"sys": "GSM"})
data = lgmpy.get_Lstar(
pos,
date,
alpha=alpha,
coord_system="GSM",
Bfield="Lgm_B_T89c",
if (igauss == 1):
# add Gaussian source term
rmod.add_source()
else:
# add PSD data
rmod.add_PSD_obs()
# Now, run the model over the enitre time range using the evolve method:
rmod.evolve()
if (igauss == 1):
# observation time delta
dt = datetime.timedelta(hours=6.0)
time = st.tickrange(start, end, dt, dtype='UTC')
# observation space delta
Lt = 2
Tgrid = rmod.ticks
nTAI = len(Tgrid)
# compute time delta
delta = Tgrid[1].UTC[0]-Tgrid[0].UTC[0]
delta = delta.seconds
# initialize arrays
Lstar = ['']*(len(time))
PSD = ['']*(len(time))
parser.print_help()
exit()
# TLEpath = '/n/space_data/TLE_DATABASE/26000_26999/26605' #ns41
partDir = '{0}000_{0}999'.format(int(options.SatNum) // 1000)
if os.path.isdir(os.path.join(options.TLEpath, partDir, options.SatNum)):
TLEpath = os.path.join(options.TLEpath, partDir, options.SatNum)
elif os.path.isdir(os.path.join(options.TLEpath, options.SatNum)):
TLEpath = os.path.join(options.TLEpath, options.SatNum)
else:
print('Specified TLE path not valid')
parser.print_help()
exit()
epochs = spt.tickrange(start_day,
end_day,
dt.timedelta(minutes=float(options.Delta)),
dtype='UTC')
epochs = epochs.UTC
print('Getting times between: {0} and {1}'.format(epochs[0], epochs[-1]))
# Set up number of cpus required for multiprocessing of different months
ncpus = cpu_count() - 1
nmonths = 1 + (epochs[-1].year - epochs[0].year) * \
12 + epochs[-1].month - epochs[0].month
if nmonths <= ncpus:
ncpus = nmonths
pool = Pool(ncpus)
if ncpus > 1:
print('Multiprocessing EphemFromTLE request')
# list of (lists of days per month to be processed)
mdateslist = splitMonths(epochs)
def test_tickrange3(self):
"""tickrange should return a known value for known input (test for bug #64 on SF tracker)"""
inval = ('2009-01-01', '2010-12-31 23:00', datetime.timedelta(hours=1))
real_ans = datetime.datetime(2010, 12, 31, 23)
ans = t.tickrange(*inval)
numpy.testing.assert_equal(real_ans, ans.UTC[-1])
# import spacepy.time as spt
# import spacepy.coordinates as spc
# import spacepy.irbempy as ib
# t = spt.Ticktock(['2002-02-02T12:00:00', '2002-02-02T12:10:00'], 'ISO')
# y = spc.Coords([[3,0,0],[2,0,0]], 'GEO', 'car')
# result=ib.get_Bfield(t,y)
# print result
# endregion
# region Description Fetch and Plot OMNI Data
import spacepy.omni as om
import spacepy.time as spt
import matplotlib.pyplot as plt
ticks = spt.tickrange('2007-03-07T00:00:00', \
'2007-05-16T00:00:00', 1./24.)
data = om.get_omni(ticks)
plt.Figure()
ax0 = plt.subplot(211)
plt.plot(data['UTC'], data['Dst'], 'r-')
ax1 = plt.subplot(212)
plt.plot(data['UTC'],data['velo'], 'k-')
ax0.set_ylabel('D$_{st}$ [nT]')
ax1.set_ylabel('V$_{sw}$ [km s$^{1}$]')
ax1.set_xlabel('UTC')
plt.show()
# endregion
from __future__ import division
import matplotlib.pyplot as plt
import itertools
import datetime as dt
import lgmpy, copy
import spacepy.time as spt
import spacepy.datamodel as dm
import spacepy.toolbox as tb
#make figure
fig = plt.figure(figsize=(17, 21))
#set up test runs
nvals, rvals = 5, [3, 4, 5, 6]
dates = spt.tickrange('2001-04-22T12:00:00', '2001-04-22T12:01:00',
dt.timedelta(seconds=5))
alpha = 45
print('T89: alpha=%s, date range [%s, %s]' %
(alpha, dates.UTC[0], dates.UTC[-1]))
#loop over radial positions, dates and quality flags
for rval in rvals:
print('Radial Distance: %s' % rval)
Lstar = [[]] * nvals
for date, qual in itertools.product(dates.UTC, range(nvals)):
print('%s, Quality=%d' % (date, qual))
pos = dm.dmarray([-1 * rval, 0, 0], attrs={'sys': 'GSM'})
data = lgmpy.get_Lstar(pos,
date,
alpha=alpha,
coord_system='GSM',
Bfield='Lgm_B_T89c',
def test_tickrange3(self):
"""tickrange should return a known value for known input (test for bug #64 on SF tracker)"""
inval = ("2009-01-01", "2010-12-31 23:00", datetime.timedelta(hours=1))
real_ans = datetime.datetime(2010, 12, 31, 23)
ans = t.tickrange(*inval)
numpy.testing.assert_equal(real_ans, ans.UTC[-1])
# TLEpath = '/n/space_data/TLE_DATABASE/26000_26999/26605' #ns41
partDir = '{0}000_{0}999'.format(int(options.SatNum) // 1000)
if os.path.isdir(os.path.join(options.TLEpath, partDir, options.SatNum)):
TLEpath = os.path.join(options.TLEpath, partDir, options.SatNum)
elif os.path.isdir(os.path.join(options.TLEpath, options.SatNum)):
TLEpath = os.path.join(options.TLEpath, options.SatNum)
else:
print('Specified TLE path not valid')
parser.print_help()
exit()
epochs = spt.tickrange(
start_day,
end_day,
dt.timedelta(
minutes=float(
options.Delta)),
dtype='UTC')
epochs = epochs.UTC
print('Getting times between: {0} and {1}'.format(epochs[0], epochs[-1]))
# Set up number of cpus required for multiprocessing of different months
ncpus = cpu_count() - 1
nmonths = 1 + (epochs[-1].year - epochs[0].year) * \
12 + epochs[-1].month - epochs[0].month
if nmonths <= ncpus:
ncpus = nmonths
pool = Pool(ncpus)
if ncpus > 1:
print('Multiprocessing EphemFromTLE request')
def plotSpectrogram(self, ecol=0, pvars=None, **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)
pvars : list
list of plotting variable names in order [Epoch-like (X axis), Flux-like (Z axis), Energy (Index var for Flux-like)]
ylim : list
2-element list with upper and lower bounds for y axis
'''
import spacepy.plot as splot
if 'Lm' not in self:
self.getLm()
sd = dm.SpaceData()
if pvars is None:
varname = self.attrs['varname']
enname = 'Energy'
epvar = 'Epoch'
else:
epvar = pvars[0]
varname = pvars[1]
enname = pvars[2]
if len(self['Lm']) != len(
self[epvar]): #Lm needs interpolating to new timebase
import matplotlib.dates as mpd
tmptimetarg, tmptimesrc = mpd.date2num(self[epvar]), mpd.date2num(
self['Epoch'])
sd['Lm'] = dm.dmarray(
np.interp(tmptimetarg,
tmptimesrc,
self['Lm'],
left=np.nan,
right=np.nan))
else:
sd['Lm'] = self['Lm']
#filter any bad Lm
goodidx = sd['Lm'] > 1
sd['Lm'] = sd['Lm'][goodidx]
if 'ylim' in kwargs:
Lm_lim = kwargs['ylim']
del kwargs['ylim']
else:
Lm_lim = [2.0, 8.0]
#TODO: allow user-definition of bins in time and Lm
sd['Epoch'] = dm.dmcopy(
self[epvar])[goodidx] #TODO: assumes 1 pitch angle, generalize
try:
sd['1D_dataset'] = self[varname][
goodidx, ecol] #TODO: assumes 1 pitch angle, generalize
except IndexError: #1-D
sd['1D_dataset'] = self[varname][goodidx]
bins = [[], []]
if 'tbins' not in kwargs:
bins[0] = spt.tickrange(self[epvar][0], self[epvar][-1],
3 / 24.).UTC
else:
bins[0] = kwargs['tbins']
del kwargs['tbins']
if 'Lbins' not in kwargs:
bins[1] = np.arange(Lm_lim[0], Lm_lim[1], 1. / 4.)
else:
bins[1] = kwargs['Lbins']
del kwargs['Lbins']
spec = splot.spectrogram(sd,
variables=['Epoch', 'Lm', '1D_dataset'],
ylim=Lm_lim,
bins=bins)
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:5.2f} {1}'.format(self[enname][ecol], self[enname].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 getPlasmaPause(ticks, model='M2002', LT='all', omnivals=None):
"""
Plasmapause location model(s)
CA1992 -- Carpenter, D. L., and R. R. Anderson, An ISEE/whistler
model of equatorial electron density in the magnetosphere,
J. Geophys. Res., 97, 1097, 1992.
M2002 -- Moldwin, M. B., L. Downward, H. K. Rassoul, R. Amin,
and R. R. Anderson, A new model of the location of the plasmapause:
CRRES results, J. Geophys. Res., 107(A11), 1339,
doi:10.1029/2001JA009211, 2002.
RT1970 -- Rycroft, M. J., and J. O. Thomas, The magnetospheric
plasmapause and the electron density trough at the alouette i
orbit, Planetary and Space Science, 18(1), 65-80, 1970
Parameters
==========
ticks : spacepy.time.Ticktock
TickTock object of desired times
Lpp_model : string, optional
'CA1992' or 'M2002' (default)
CA1992 returns the Carpenter and Anderson model,
M2002 returns the Moldwin et al. model
LT : int, float
requested local time sector, 'all' is valid option
omnivals : spacepy.datamodel.SpaceData, dict
dict-like containing UTC (datetimes) and Kp keys
Returns
=======
out : float
Plasmapause radius in Earth radii
Examples
========
>>> import spacepy.time as spt
>>> import spacepy.empiricals as emp
>>> ticks = spt.tickrange('2002-01-01T12:00:00','2002-01-04T00:00:00',.25)
>>> emp.getPlasmaPause(ticks)
array([ 6.42140002, 6.42140002, 6.42140002, 6.42140002, 6.42140002,
6.42140002, 6.42140002, 6.26859998, 5.772 , 5.6574 ,
5.6574 ])
"""
def calcLpp(Kpmax, A, B, power=1):
currLpp = A - B*Kpmax**power
return currLpp
model_list = ['CA1992', 'M2002', 'RT1970']
if model == 'CA1992':
if LT!='all':
print('No LT dependence currently supported for this model')
if model not in model_list:
raise ValueError("Please specify a valid model:\n{0}".format(' or '.join(model_list)))
if LT=='all':
parA = {'CA1992': 5.6, 'M2002': 5.39, 'RT1970': 5.64}
parB = {'CA1992': 0.46, 'M2002': 0.382, 'RT1970': 0.78}
priorvals = {'CA1992': datetime.timedelta(hours=24),
'M2002': datetime.timedelta(hours=12),
'RT1970': datetime.timedelta(0)}
A, B = parA[model], parB[model]
prior = priorvals[model]
else:
try:
float(LT)
except (ValueError, TypeError):
raise ValueError("Please specify a valid LT:\n'all' or a numeric type")
parA = {'CA1992': [5.6]*24,
'M2002': [5.7]*3+[6.05]*6+[5.2]*6+[4.45]*6+[5.7]*3,
'RT1970': [5.64]*24}
parB = {'CA1992': [0.46]*24,
'M2002': [0.42]*3+[0.573]*6+[0.425]*6+[0.167]*6+[0.42]*3,
'RT1970': [0.78]*24}
priorvals = {'CA1992': [datetime.timedelta(hours=24)]*24,
'M2002': [datetime.timedelta(hours=12)]*24,
'RT1970': [datetime.timedelta(0)]*24}
try:
LThr = long(LT)
except NameError:
LThr = int(LT)
prior = priorvals[model][LThr]
A, B = parA[model][LThr], parB[model][LThr]
st, en = ticks.UTC[0]-prior, ticks.UTC[-1]
if omnivals is None:
omdat = om.get_omni(spt.tickrange(st, en, 1.0/24.0), dbase='QDhourly')
else:
#now test for sanity of input
try:
assert isinstance(omnivals, dict)
except:
raise TypeError('Not a valid input type for omnivals, expected spacepy.datamodel.SpaceData')
try:
assert 'UTC' in omnivals
assert 'Kp' in omnivals
except:
raise KeyError('Required data not found in input dict-like (omnivals)')
omdat = omnivals
einds, oinds = tb.tOverlap([st, en], omdat['UTC'])
utc = np.array(omdat['UTC'])[oinds]
Kp = np.array(omdat['Kp'])[oinds]
Lpp = np.zeros(len(ticks))
if model == 'RT1970':
power = 0.5
else:
power = 1
for i, t1 in enumerate(ticks.UTC):
t0 = t1-prior
iprevday, dum = tb.tOverlap(utc, [t0, t1])
if iprevday:
Kpmax = max(Kp[iprevday])
Lpp[i] = calcLpp(Kpmax, A, B, power=power)
else:
Lpp[i] = np.nan
return Lpp
def getPlasmaPause(ticks, model='M2002', LT='all', omnivals=None):
"""
Plasmapause location model(s)
CA1992 -- Carpenter, D. L., and R. R. Anderson, An ISEE/whistler
model of equatorial electron density in the magnetosphere,
J. Geophys. Res., 97, 1097, 1992.
M2002 -- Moldwin, M. B., L. Downward, H. K. Rassoul, R. Amin,
and R. R. Anderson, A new model of the location of the plasmapause:
CRRES results, J. Geophys. Res., 107(A11), 1339,
doi:10.1029/2001JA009211, 2002.
RT1970 -- Rycroft, M. J., and J. O. Thomas, The magnetospheric
plasmapause and the electron density trough at the alouette i
orbit, Planetary and Space Science, 18(1), 65-80, 1970
Parameters
==========
ticks : spacepy.time.Ticktock
TickTock object of desired times
Lpp_model : string, optional
'CA1992' or 'M2002' (default)
CA1992 returns the Carpenter and Anderson model,
M2002 returns the Moldwin et al. model
LT : int, float
requested local time sector, 'all' is valid option
omnivals : spacepy.datamodel.SpaceData, dict
dict-like containing UTC (datetimes) and Kp keys
Returns
=======
out : float
Plasmapause radius in Earth radii
Examples
========
>>> import spacepy.time as spt
>>> import spacepy.empiricals as emp
>>> ticks = spt.tickrange('2002-01-01T12:00:00','2002-01-04T00:00:00',.25)
>>> emp.getPlasmaPause(ticks)
array([ 6.42140002, 6.42140002, 6.42140002, 6.42140002, 6.42140002,
6.42140002, 6.42140002, 6.26859998, 5.772 , 5.6574 ,
5.6574 ])
"""
def calcLpp(Kpmax, A, B, power=1):
currLpp = A - B * Kpmax**power
return currLpp
model_list = ['CA1992', 'M2002', 'RT1970']
if model == 'CA1992':
if LT != 'all':
print('No LT dependence currently supported for this model')
if model not in model_list:
raise ValueError("Please specify a valid model:\n{0}".format(
' or '.join(model_list)))
if LT == 'all':
parA = {'CA1992': 5.6, 'M2002': 5.39, 'RT1970': 5.64}
parB = {'CA1992': 0.46, 'M2002': 0.382, 'RT1970': 0.78}
priorvals = {
'CA1992': datetime.timedelta(hours=24),
'M2002': datetime.timedelta(hours=12),
'RT1970': datetime.timedelta(0)
}
A, B = parA[model], parB[model]
prior = priorvals[model]
else:
try:
float(LT)
except (ValueError, TypeError):
raise ValueError(
"Please specify a valid LT:\n'all' or a numeric type")
parA = {
'CA1992': [5.6] * 24,
'M2002':
[5.7] * 3 + [6.05] * 6 + [5.2] * 6 + [4.45] * 6 + [5.7] * 3,
'RT1970': [5.64] * 24
}
parB = {
'CA1992': [0.46] * 24,
'M2002':
[0.42] * 3 + [0.573] * 6 + [0.425] * 6 + [0.167] * 6 + [0.42] * 3,
'RT1970': [0.78] * 24
}
priorvals = {
'CA1992': [datetime.timedelta(hours=24)] * 24,
'M2002': [datetime.timedelta(hours=12)] * 24,
'RT1970': [datetime.timedelta(0)] * 24
}
try:
LThr = long(LT)
except NameError:
LThr = int(LT)
prior = priorvals[model][LThr]
A, B = parA[model][LThr], parB[model][LThr]
st, en = ticks.UTC[0] - prior, ticks.UTC[-1]
if omnivals is None:
omdat = om.get_omni(spt.tickrange(st, en, 1.0 / 24.0),
dbase='QDhourly')
else:
#now test for sanity of input
try:
assert isinstance(omnivals, dict)
except:
raise TypeError(
'Not a valid input type for omnivals, expected spacepy.datamodel.SpaceData'
)
try:
assert 'UTC' in omnivals
assert 'Kp' in omnivals
except:
raise KeyError(
'Required data not found in input dict-like (omnivals)')
omdat = omnivals
einds, oinds = tb.tOverlap([st, en], omdat['UTC'])
utc = np.array(omdat['UTC'])[oinds]
Kp = np.array(omdat['Kp'])[oinds]
Lpp = np.zeros(len(ticks))
if model == 'RT1970':
power = 0.5
else:
power = 1
for i, t1 in enumerate(ticks.UTC):
t0 = t1 - prior
iprevday, dum = tb.tOverlap(utc, [t0, t1])
if iprevday:
Kpmax = max(Kp[iprevday])
Lpp[i] = calcLpp(Kpmax, A, B, power=power)
else:
Lpp[i] = np.nan
return Lpp
