def setUp(self):
self.mhd = pbs.Bats2d(
os.path.join(spacepy_testing.datadir, 'pybats_test',
'y0_binary.out'))
self.outs = pbs.Bats2d(
os.path.join(
spacepy_testing.datadir, 'pybats_test',
'y=0_mhd_1_e20140410-000000-000_20140410-000200-000.outs'))
class TestBats2d(unittest.TestCase):
'''
Test functionality of Bats2d objects.
'''
mhd = pbs.Bats2d('data/pybats_test/y0_binary.out')
def testCalc(self):
# Test all calculations:
self.mhd.calc_all()
def testMultispecies(self):
# Open file:
mhd = pbs.Bats2d('data/pybats_test/cut_multispecies.out',
format='ascii')
mspec_varnames = 'x Rho Ux Uy Uz Bx By Bz P OpRho OpUx OpUy OpUz OpP jx jy jz g rbody cuty cutz'.split(
)
mspec_units = 'km Mp/cc km/s km/s km/s nT nT nT nPa Mp/cc km/s km/s km/s nPa uA/m2 uA/m2 uA/m2'.split(
)
knownMultispecUnits = dict(zip(mspec_varnames, mspec_units))
# Test units are loaded correctly:
for v in mhd:
if v not in mspec_varnames: continue
self.assertEqual(knownMultispecUnits[v], mhd[v].attrs['units'])
mhd.calc_all(exclude=['calc_gradP', 'calc_vort'])
def makeplot(fname, ax):
data = bats.Bats2d(fn)
data.add_b_magsphere(target=ax, nOpen=10, nClosed=8)
ti_str = '{0}'.format(data.attrs['time'])
ax.set_xlim([-32, 16])
ax.set_ylim([-24, 24])
ax.set_xlabel('X$_{GSM}$ [R$_E$]')
ax.set_ylabel('Z$_{GSM}$ [R$_E$]')
ax.text(0.5, 1.01, ti_str, transform=ax.transAxes)
class TestBats2d(unittest.TestCase):
'''
Test functionality of Bats2d objects.
'''
mhd = pbs.Bats2d('data/pybats_test/y0_binary.out')
def testCalc(self):
# Test all calculations:
self.mhd.calc_all()
class TestExtraction(unittest.TestCase):
'''
Test Extraction class by opening a file with known solution.
'''
mhd = pbs.Bats2d('data/pybats_test/z0_sine.out')
def testExtract(self):
from numpy import pi, cos
analytic = lambda x: 1. + .5 * cos(x * pi / 10.)
extr = self.mhd.extract(range(-5, 6), [-8] * 11)
for x, rho in zip(extr['x'], extr['rho']):
self.assertAlmostEqual(rho, analytic(x), 2)
def testMultispecies(self):
# Open file:
mhd = pbs.Bats2d('data/pybats_test/cut_multispecies.out',
format='ascii')
mspec_varnames = 'x Rho Ux Uy Uz Bx By Bz P OpRho OpUx OpUy OpUz OpP jx jy jz g rbody cuty cutz'.split(
)
mspec_units = 'km Mp/cc km/s km/s km/s nT nT nT nPa Mp/cc km/s km/s km/s nPa uA/m2 uA/m2 uA/m2'.split(
)
knownMultispecUnits = dict(zip(mspec_varnames, mspec_units))
# Test units are loaded correctly:
for v in mhd:
if v not in mspec_varnames: continue
self.assertEqual(knownMultispecUnits[v], mhd[v].attrs['units'])
mhd.calc_all()
def testMultispecies(self):
# Open file:
mhd = pbs.Bats2d(
os.path.join(spacepy_testing.datadir, 'pybats_test',
'cut_multispecies.out'))
mspec_varnames = 'x Rho Ux Uy Uz Bx By Bz P OpRho OpUx OpUy OpUz OpP jx jy jz g rbody cuty cutz'.split(
)
mspec_units = 'km Mp/cc km/s km/s km/s nT nT nT nPa Mp/cc km/s km/s km/s nPa uA/m2 uA/m2 uA/m2'.split(
)
knownMultispecUnits = dict(zip(mspec_varnames, mspec_units))
# Test units are loaded correctly:
for v in mhd:
if v not in mspec_varnames: continue
self.assertEqual(knownMultispecUnits[v], mhd[v].attrs['units'])
mhd.calc_all(exclude=['calc_gradP', 'calc_vort'])
def convertOneFile(fname,
triangulate=True,
make_vectors=True,
keep_components=False):
'''Convert a single file from BATSRUS 2D to VTP
Parameters
----------
fname : str
Filename to convert to VTP
Other Parameters
----------------
triangulate : bool
If True, return a 2D Delaunay triangulation of the data. If False
just return the point mesh. Default is True.
make_vectors : bool
If True, combine vector state variables in VTK vectors. Default is True.
keep_components : bool
If True, keep components of vector state variables as scalars in output files.
Default is False. Warning: Setting this as False while make_vectors is False
will remove vector quantities from the output VTP files.
'''
# load BATSRUS 2D file
data = bts.Bats2d(fname)
outname = os.path.splitext(fname)[0]
# set output filename
fname = '.'.join((outname, 'vtp'))
# Initialize the VTK object
# Should work on y=0 and z=0 slices, probably also 3D IDL.
# Won't work on x=0 or arbitrary slice... (yet)
xs = data['x']
vec_dict = {'zeroy': False, 'zeroz': False}
if 'y' in data and 'z' not in data:
ys = data['y']
zs = np.zeros(len(xs))
vec_dict['zeroz'] = True
elif 'y' not in data and 'z' in data:
ys = np.zeros(len(xs))
zs = data['z']
vec_dict['zeroy'] = True
else: # Input data is probably 3D, but this should still work (untested)
ys = data['y']
zs = data['z']
points = np.c_[xs, ys, zs]
point_cloud = pv.PolyData(points)
if 'jx' in data:
if keep_components:
point_cloud["jx"] = data["jx"]
point_cloud["jy"] = data["jy"]
point_cloud["jz"] = data["jz"]
if make_vectors:
point_cloud['jvec_inplane'] = mergeComponents(
data["jx"], data["jy"], data["jz"], **vec_dict)
if 'bx' in data:
if keep_components:
point_cloud["bx"] = data["bx"]
point_cloud["by"] = data["by"]
point_cloud["bz"] = data["bz"]
if make_vectors:
point_cloud['bvec_inplane'] = mergeComponents(
data["bx"], data["by"], data["bz"], **vec_dict)
if 'ux' in data:
if keep_components:
point_cloud["ux"] = data["ux"]
point_cloud["uy"] = data["uy"]
point_cloud["uz"] = data["uz"]
if make_vectors:
point_cloud['uvec_inplane'] = mergeComponents(
data["ux"], data["uy"], data["uz"], **vec_dict)
if 'rho' in data:
point_cloud["rho"] = data["rho"]
if 'p' in data:
point_cloud["p"] = data["p"]
# Convert from points mesh to connected mesh using triangulation
if triangulate:
point_cloud.delaunay_2d(tol=1e-05,
alpha=0.0,
offset=1.0,
bound=False,
inplace=True,
edge_source=None,
progress_bar=False)
# write to binary XML output
point_cloud.save(fname)
def setUp(self):
self.pth = os.path.dirname(os.path.abspath(__file__))
self.mhd = pbs.Bats2d(
os.path.join(self.pth, 'data', 'pybats_test', 'z0_sine.out'))
def setUp(self):
self.mhd = pbs.Bats2d(
os.path.join(spacepy_testing.datadir, 'pybats_test',
'z0_sine.out'))
def plot_results(fileY, fileZ, iFile=0, nFiles=1):
'''
Create a nice plot of WTF is happening.
Kwargs iFile and nFiles are for reporting progress.
'''
# Get start times from file names:
time_y = pb.parse_filename_time(fileY)
time_z = pb.parse_filename_time(fileZ)
# Check that times match:
for t1, t2 in zip(time_y, time_z):
if t1 != t2:
print(
'File times do not match!\n\tY-File: {}\n\tZ-File: {}'.format(
t1, t2))
raise ValueError('File time mismatch')
# Use file time info get actual time:
if time_y[-1]:
tNow = time_y[-1] # Datetime in file name.
elif time_y[1]:
t_delta = dt.timedelta(seconds=time_y[1])
tNow = start + t_delta # Runtime in file name.
else:
tNow = start # Only iteration given.
# Check if files are outside time range.
if tNow < trange[0]: return 0
if tNow > trange[-1]: return 0
# Open files & calculate important values:
mY = pbs.Bats2d(fileY)
mZ = pbs.Bats2d(fileZ)
mY.calc_ndens()
mZ.calc_ndens()
mY.calc_utotal()
mZ.calc_utotal()
if mZ.attrs['runtime'] != mY.attrs['runtime']:
print('{}\n{}\n'.format(fileY, fileZ))
raise (ValueError('Times of files do not line up!'))
fig = plt.figure(figsize=[11, 8.5])
fig.subplots_adjust(left=0.08, right=.96, bottom=0.06, top=.93)
a1 = fig.add_subplot(221)
stuff = mY.add_contour('x',
'z',
v1,
91,
dolog=True,
xlim=xrng,
ylim=zrng,
add_cbar=True,
target=a1,
zlim=prng,
cmap=cmap1)
a1.set_title('{} ({})'.format(pb.mhdname_to_tex(v1),
mY[v1].attrs['units']))
stuff[-1].set_label('')
if args.mag: mY.add_b_magsphere(a1, DoOpen=True)
a2 = fig.add_subplot(222)
stuff = mZ.add_contour('x',
'y',
v2,
91,
dolog=True,
xlim=xrng,
ylim=yrng,
add_cbar=True,
target=a2,
zlim=drng,
cmap=cmap2)
a2.set_title('{} ({})'.format(pb.mhdname_to_tex(v2),
mZ[v2].attrs['units']))
stuff[-1].set_label('')
for s in sats:
s.add_sat_loc(tNow, a1, plane='XZ', dolabel=True, color='r', size=9)
s.add_sat_loc(tNow, a2, plane='XY', dolabel=True, color='r', size=9)
pos = stuff[-1].ax.get_position()
a3, a4 = fig.add_subplot(413), fig.add_subplot(414)
a5 = a4.twinx()
pos3, pos4 = a3.get_position(), a4.get_position()
a3.set_position([pos3.xmin, pos3.ymin, pos3.width * .95, pos3.height])
a4.set_position([pos4.xmin, pos4.ymin, pos4.width * .95, pos4.height])
a5.set_position([pos4.xmin, pos4.ymin, pos4.width * .95, pos4.height])
log.add_dst_quicklook(target=a3,
plot_obs=args.obsdst,
obs_kwargs={
'c': 'k',
'ls': '-'
},
color='r')
a3.legend(loc='best', fontsize=10)
a3.set_xlabel('')
applySmartTimeTicks(a3, trange, dolabel=False)
a3.grid(False, axis='y')
a3.set_ylabel('D$_{ST}$ ($nT$)')
a3.hlines(0, trange[0], trange[1], color='grey', linestyles='dashed')
ymin, ymax = a3.get_ylim()
a3.vlines(tNow, ymin, ymax, linestyles='solid', color='g', linewidths=2.)
a3.set_ylim([ymin, ymax])
a4.plot(imf['time'], imf['bz'], color=bzcolor, label='B$_Z$')
if args.by:
a4.plot(imf['time'], imf['by'], color=bycolor, label='B$_Y$')
a4.legend(loc='best', ncol=2)
a5.plot(imf['time'], imf['pram'], color=pdcolor, lw=1.5)
applySmartTimeTicks(a4, trange, dolabel=True)
a4.grid(False, axis='y')
a4.hlines(0, trange[0], trange[1], color=bzcolor, linestyles='dashed')
a4.set_ylabel('IMF ' + (not args.by) * 'B$_{Z}$ ' + '($nT$)',
color=bzcolor)
a5.set_ylabel('P$_{dyn}$ ($nPa$)', color=pdcolor)
a5.grid(False, axis='y')
ymin, ymax = a4.get_ylim()
a4.vlines(tNow, ymin, ymax, linestyles='solid', color='g', linewidths=2.)
a4.set_ylim([ymin, ymax])
fig.suptitle(args.title, size=18)
# Save the figure (normal operation) or show figure (debug mode).
if args.debug:
plt.show()
else:
fig.savefig(args.outdir + '/mhd{:%Y%m%d_%H%M%S}.png'.format(tNow))
# Clear and close figure.
if not args.debug:
fig.clf()
plt.close(fig)
percent_done = float(iFile + 1) / float(nFiles) * 100.0
return percent_done
if __name__ == '__main__':
# Find files, sort them.
yFiles = glob(prefix + 'y??_mhd*_[te]*.out')
zFiles = glob(prefix + 'z??_mhd*_[te]*.out')
yFiles.sort()
zFiles.sort()
if not (yFiles and zFiles):
raise (ValueError("Could not find any Y or Z slice files."))
if args.debug:
yFiles, zFiles = yFiles[:1], zFiles[:1]
if args.info:
mhd = pbs.Bats2d(yFiles[0])
mhd.calc_ndens()
mhd.calc_utotal()
mhd.tree()
exit()
#pool.map(plot_results, zip(yFiles, zFiles)[:3])
# Loop over files, making a plot for each set.
nFiles = len(yFiles)
print('\tFiles: working on {} files...'.format(nFiles))
stdout.write('\t{:5.1f}% complete...'.format(0))
if args.nthread > 1:
# Create pool of workers.
pool = mp.Pool(args.nthread)
