Add 2 example scripts, for FFT and Themis plotting

examples/fourier_transform_field.py

@@ -0,0 +1,43 @@
+#!/usr/bin/python
+#
+# Fourier transform a scalar variable, or one component of a vector, in a vlsv file
+# (entiere space)
+#
+import pytools as pt
+import numpy
+import pylab
+import sys
+
+if len(sys.argv) < 3:
+ sys.stderr.write("Syntax: fourier_transform_field.py <vlsvfile> <field> [<component>]\n")
+ sys.exit()
+
+filename = sys.argv[1]
+variable = sys.argv[2]
+
+# Open file
+f = pt.vlsvfile.VlsvReader(filename)
+
+# Read our field and cellIDs to sort it
+cellids = f.read_variable("CellID")
+a = f.read_variable(variable)
+
+if len(a.shape) > 1:
+ if len(sys.argv) > 3:
+ a = a[:,int(sys.argv[3])]
+ else:
+ a = numpy.sqrt( numpy.sum(a*a,1))
+
+xsize = f.read_parameter("xcells_ini")
+ysize = f.read_parameter("ycells_ini")
+
+# Sort the Field to be a proper 2D numpy-array
+a = a[cellids.argsort()].reshape([ysize,xsize])
+
+# 2D fourier-transform it.
+af = numpy.fft.fftshift(numpy.fft.fft2(a))
+
+
+# Plot it!
+pylab.plt.subplot().pcolormesh(numpy.log(abs(af)))
+pylab.plt.show()

examples/plot_themis_planes.py

@@ -0,0 +1,54 @@
+#!/usr/bin/python
+# -*- coding: utf-8 -*-
+#
+# Given input file and cellID, plot the velocity space as it would be seen by Themis ESA detector.
+# Creates output pngs in multiple cut planes:
+# x Direction | y Direction | filename
+# -----------------------+---------------------------+---------------------
+# parallel to B | perp. to B, in dir of v | themis_B_v.png
+# parallel to B | perp. to B, in dir of B×v | themis_B_Bxv.png
+# perp. to B, dir of v | perp. to B, in dir of Bxv | themis_B_perpperp.png
+# simulation x | simulation y | themis_x_y.png
+# simulation x | simulation z | themis_x_z.png
+# simulation y | simulation z | themis_y_z.png
+#
+# The plots where B does not directly specify an axis also have the magnetic
+# field vector drawn in as an arrow.
+
+import pytools as pt
+import numpy as np
+import pylab as pl
+import sys
+
+if len(sys.argv) < 2:
+ sys.stderr.write("Syntax: plot_themis_planes.py <vlsvfile> cellID\n")
+ sys.exit()
+
+filename = sys.argv[1]
+f = pt.vlsvfile.VlasiatorReader(filename)
+
+cellid = int(sys.argv[2])
+
+B = f.read_variable("B",cellid)
+v = f.read_variable("v",cellid)
+
+Bxv = np.cross(B,v)
+BxvxB = np.cross(Bxv,B)
+
+pt.calculations.themis_plot_phasespace_helistyle(f, cellid, B,v,smooth=True,xlabel=u"v∥B",ylabel=u"v∥V",vmin_min=1e-16)
+pl.savefig("themis_B_v.png")
+pt.calculations.themis_plot_phasespace_helistyle(f, cellid, B,Bxv,smooth=True,xlabel=u"v∥B",ylabel=u"v∥B×V", vmin_min=1e-17)
+pl.savefig("themis_B_Bxv.png")
+pt.calculations.themis_plot_phasespace_helistyle(f, cellid, Bxv,BxvxB,smooth=True,ylabel=u"v⟂B (in Plane of V)",xlabel=u"v⟂B (in Plane of B×V)", vmin_min=1e-17)
+pl.savefig("themis_B_perpperp.png")
+
+Bn = B / np.linalg.norm(B)
+ax = pt.calculations.themis_plot_phasespace_helistyle(f, cellid, np.array([1.,0,0]),np.array([0,1.,0.]),smooth=True,xlabel=u"vx",ylabel=u"vy", vmin_min=1e-17)
+ax.arrow(0,0,2000*Bn[0],2000*Bn[1],head_width=100,head_length=200, edgecolor='black', facecolor='black')
+pl.savefig("themis_x_y.png")
+ax = pt.calculations.themis_plot_phasespace_helistyle(f, cellid, np.array([1.,0,0]),np.array([0,0.,1.]),smooth=True,xlabel=u"vx",ylabel=u"vz", vmin_min=1e-17)
+ax.arrow(0,0,2000*Bn[0],2000*Bn[2],head_width=100,head_length=200, edgecolor='black', facecolor='black')
+pl.savefig("themis_x_z.png")
+ax = pt.calculations.themis_plot_phasespace_helistyle(f, cellid, np.array([0.,1.,0]),np.array([0,0.,1.]),smooth=True,xlabel=u"vy",ylabel=u"vz", vmin_min=1e-17)
+ax.arrow(0,0,2000*Bn[1],2000*Bn[2],head_width=100,head_length=200, edgecolor='black', facecolor='black')
+pl.savefig("themis_y_z.png")

examples/plot_timestep.py

@@ -0,0 +1,77 @@
+#!/usr/bin/python
+# Plot a single vlasiator bulk file, can be used for continuous job monitoring.
+# 
+# Note: this requires python with matplotlib!
+# On hornet, do:
+# module load tools/python/2.7.8
+# before running this.
+
+import matplotlib;
+matplotlib.use('Agg')
+
+from matplotlib.colors import LogNorm
+import numpy as np;
+import pytools as pt;
+import pylab as pl;
+import numpy.ma as ma;
+import sys
+
+if len(sys.argv) < 2:
+ print("Syntax: python plot_timestep.py bulkfile.vlsv output.png")
+ sys.exit()
+
+# Open file
+filename = sys.argv[1] 
+outputname = sys.argv[2]
+f = pt.vlsvfile.VlasiatorReader(filename)
+
+# Themis colormap, as extracted from the themis tools' IDL file
+hot_desaturated_colors=[(71./255.,71./255.,219./255.),(0,0,91./255.),(0,1,1),(0,.5,0),(1,1,0),(1,96./255,0),(107./255,0,0),(224./255,76./255,76./255)]
+hot_desaturated_colormap = matplotlib.colors.LinearSegmentedColormap.from_list("hot_desaturated",hot_desaturated_colors)
+
+# Determine sizes
+xsize = f.read_parameter("xcells_ini")
+ysize = f.read_parameter("ycells_ini")
+zsize = f.read_parameter("zcells_ini")
+cellids = f.read_variable("CellID")
+
+# Juggle fields into correct order
+rho = f.read_variable("rho")
+rho = rho[cellids.argsort()].reshape([ysize,xsize])
+
+B = f.read_variable("B")
+#B = f.read_variable("perturbed_B") + f.read_variable("B")
+B_mag = np.array([np.linalg.norm(v) for v in B])
+B_mag = B_mag[cellids.argsort()].reshape([ysize,xsize])
+
+boundary_type = f.read_variable("Boundary_type");
+boundary_type = boundary_type[cellids.argsort()].reshape([ysize,xsize])
+boundary_type = (boundary_type != 1)
+
+maskedrho = ma.masked_array(rho,mask=boundary_type)
+maskedB = ma.masked_array(B_mag,mask=boundary_type)
+
+rhomin = np.min(maskedrho)
+rhomax = np.max(maskedrho)
+Bmin = np.min(maskedB)
+Bmax = np.max(maskedB)
+
+# Plot rho and B
+pl.rcParams['figure.figsize']= 30,38
+#pl.rcParams['figure.DPI']= xcells_ini/30
+
+ax1=pl.plt.subplot(2,1,1)
+ax1.set_aspect('equal','datalim')
+ax1.set_title("Rho (min = " + ("%.3e" % rhomin)+ ", max = " +("%.3e" % rhomax)+")")
+fig1=ax1.pcolormesh(maskedrho, norm=LogNorm(vmin=1e4, vmax=1e7),cmap=hot_desaturated_colormap);
+pl.plt.colorbar(fig1)
+
+ax2=pl.plt.subplot(2,1,2,sharex=ax1, sharey=ax1)
+ax2.set_title("B(min = " +("%.3e" % Bmin)+ ", max = " +("%.3e" % Bmax)+")")
+fig2=ax2.pcolormesh(maskedB, norm=LogNorm(vmin=1e-9, vmax=1e-7),cmap=hot_desaturated_colormap);
+pl.plt.colorbar(fig2)
+
+#fig1.tight_layout()
+#fig2.tight_layout()
+pl.plt.savefig(outputname,bbox_inches='tight',dpi=xsize/30)
+

pyCalculations/themis_observation.py

@@ -104,8 +104,10 @@ def themis_plot_detector(vlsvReader, cellID, detector_axis=np.array([0,1,0])):
  print("Plotting...")
  cax = ax.pcolormesh(grid_theta,grid_r,values, norm=matplotlib.colors.LogNorm(vmin=vmin,vmax=vmax), cmap=themis_colormap)
  ax.grid(True)
- fig.colorbar(cax)
- pl.show()
+ cb = fig.colorbar(cax)
+ cb.set_label("Count rate")
+ #pl.show()
+ return ax
 
 def themis_plot_phasespace_contour(vlsvReader, cellID, plane_x=np.array([1.,0,0]), plane_y=np.array([0,0,1.]), smooth=False, xlabel="Vx", ylabel="Vy"):
  ''' Plots a contour view of phasespace, as seen by a themis detector, at the given cellID
@@ -148,9 +150,9 @@ def themis_plot_phasespace_contour(vlsvReader, cellID, plane_x=np.array([1.,0,0]
  cax = ax.contour(xi,yi,vi.T, levels=np.logspace(np.log10(vmin),np.log10(vmax),20), norm=matplotlib.colors.LogNorm(vmin=vmin,vmax=vmax))
  ax.grid(True)
  fig.colorbar(cax)
- pl.show()
+ #pl.show()
 
-def themis_plot_phasespace_helistyle(vlsvReader, cellID, plane_x=np.array([1.,0,0]), plane_y=np.array([0,0,1.]), smooth=True, xlabel="Vx", ylabel="Vy"):
+def themis_plot_phasespace_helistyle(vlsvReader, cellID, plane_x=np.array([1.,0,0]), plane_y=np.array([0,0,1.]), smooth=True, xlabel="Vx", ylabel="Vy", vmin_min=1e-16):
  ''' Plots a view of phasespace, as seen by a themis detector, at the given cellID, in the style that heli likes.
  :param vlsvReader: Some VlsvReader class with a file open
  :type vlsvReader: :class:`vlsvfile.VlsvReader`
@@ -161,9 +163,9 @@ def themis_plot_phasespace_helistyle(vlsvReader, cellID, plane_x=np.array([1.,0,
 
  matrix = simulation_to_observation_frame(plane_x,plane_y)
 
- angles, energies, vmin, vmax, values = themis_observation_from_file( vlsvReader=vlsvReader, cellid=cellID, matrix=matrix, countrates=False)
+ angles, energies, vmin, vmax, values = themis_observation_from_file( vlsvReader=vlsvReader, cellid=cellID, matrix=matrix, countrates=False, noise=False)
  if vmin == 0:
- vmin = 1e-15
+ vmin = vmin_min
  if vmax < vmin:
  vmax = vmin*10
 
@@ -188,13 +190,27 @@ def themis_plot_phasespace_helistyle(vlsvReader, cellID, plane_x=np.array([1.,0,
  ax.set_title("Phasespace at cell " + str(cellID))
  ax.set_xlabel(xlabel+" (km/s)")
  ax.set_ylabel(ylabel+" (km/s)")
+ vi *= 10.**-3 # Go from s^3 m^-6 to s^3 km^-3 cm^-3
+ vmin *= 10.**-3
+ vmax *= 5*10.**-3
  cax = ax.pcolormesh(xi,yi,vi.T, norm=matplotlib.colors.LogNorm(vmin=vmin,vmax=vmax), vmin=vmin, vmax=vmax, cmap=pl.get_cmap("Blues"), shading='flat')
  cax2 = ax.contourf(xi,yi,vi.T, levels=np.logspace(np.log10(vmin),np.log10(vmax),20), norm=matplotlib.colors.LogNorm(vmin=vmin,vmax=vmax), vmin=vmin, vmax=vmax, cmap=pl.get_cmap("Blues"))
  #cax3 = ax.contour(xi,yi,vi.T, levels=np.logspace(np.log10(vmin),np.log10(vmax),20), norm=matplotlib.colors.LogNorm(vmin=vmin,vmax=vmax), cmap=pl.get_cmap("binary"))
  ax.grid(True)
- fig.colorbar(cax)
- pl.show()
-def themis_observation_from_file( vlsvReader, cellid, matrix=np.array([[1,0,0],[0,1,0],[0,0,1]]), countrates=True, interpolate=True,binOffset=[0.,0.]):
+
+ # Draw a circle at themis max extents
+ angles = np.linspace(0,2*np.pi,180)
+ x = 2188*np.cos(angles)
+ y = 2188*np.sin(angles)
+ ax.plot(x,y)
+
+ cb = fig.colorbar(cax)
+ cb.set_label("Phasespace density (s^3 km^-3 cm^3)")
+ #pl.show()
+
+ return ax
+
+def themis_observation_from_file( vlsvReader, cellid, matrix=np.array([[1,0,0],[0,1,0],[0,0,1]]), countrates=True, interpolate=True, noise=False, binOffset=[0.,0.]):
  ''' Calculates artificial THEMIS EMS observation from the given cell
  :param vlsvReader: Some VlsvReader class with a file open
  :type vlsvReader: :class:`vlsvfile.VlsvReader`
@@ -230,7 +246,7 @@ def themis_observation_from_file( vlsvReader, cellid, matrix=np.array([[1,0,0],[
  vcellids = velocity_cell_data.keys()
  # Get a list of velocity coordinates:
  velocity_coordinates = vlsvReader.get_velocity_cell_coordinates(vcellids)
- angles,energies,detector_values = themis_observation(velocity_cell_data, velocity_coordinates,matrix,dvx,countrates=countrates, interpolate=interpolate,binOffset=binOffset)
+ angles,energies,detector_values = themis_observation(velocity_cell_data, velocity_coordinates,matrix,dvx,countrates=countrates, interpolate=interpolate, noise=noise, binOffset=binOffset)
 
  # Calc min and max
  val_min = np.min(detector_values)
@@ -239,7 +255,7 @@ def themis_observation_from_file( vlsvReader, cellid, matrix=np.array([[1,0,0],[
 
 
 
-def themis_observation(velocity_cell_data, velocity_coordinates, matrix, dv=30e3, countrates=True, interpolate=True, binOffset=[0.,0.]):
+def themis_observation(velocity_cell_data, velocity_coordinates, matrix, dv=30e3, countrates=True, interpolate=True, binOffset=[0.,0.], noise=False):
  ''' Calculates artificial THEMIS EMS observation from the given velocity space data
 
  :param velocity_cell_data: velocity cell information as obtained from vlsvReader
@@ -294,7 +310,7 @@ def themis_observation(velocity_cell_data, velocity_coordinates, matrix, dv=30e3
  continue
 
  target_val = avgs[i]
- if countrates:
+ if countrates or noise:
  # Calculate actual countrate from phasespace density
  # (=rho*g/E*v*dt)
  #deltaOmega = dv*dv/v_abs[i] # solid angle covered by the phase space cell (approx)
@@ -306,6 +322,14 @@ def themis_observation(velocity_cell_data, velocity_coordinates, matrix, dv=30e3
 
  target_val = avgs[i] * velcell_volume * effectiveArea * v_abs[i] * detector_timestep
 
+ if noise:
+ # Draw a poisson-distributed random number based on the countrate + some background
+ target_val = np.random.poisson(target_val) + np.random.random() * 2
+ if not countrates:
+ # convert back to phasespace density
+ target_val /= velcell_volume * effectiveArea * v_abs[i] * detector_timestep
+
+
  if interpolate:
  # Linearly interpolate
  angle_frac = angle_bins[i] - int(angle_bins[i])

pyMayaVi/mayavigrid.py

@@ -435,6 +435,7 @@ def __picker_callback( self, picker ):
  # Plot pitch angle distribution:
  from themis_observation import themis_plot_detector
  themis_plot_detector(self.vlsvReader,cellid,detectoraxis)
+ pl.show()
  elif (self.picker == "Themis_contour"):
  # Find the nearest cell id with distribution:
  # Read cell ids with velocity distribution in:
@@ -482,6 +483,7 @@ def __picker_callback( self, picker ):
  self.__add_label( cellid )
  from themis_observation import themis_plot_phasespace_contour
  themis_plot_phasespace_contour(self.vlsvReader, cellid, plane[0], plane[1],xlabel=labels[0],ylabel=labels[1])
+ pl.show()
  elif (self.picker == "Themis_helistyle"):
 
  # Find the nearest cell id with distribution:
@@ -530,6 +532,7 @@ def __picker_callback( self, picker ):
  self.__add_label( cellid )
  from themis_observation import themis_plot_phasespace_helistyle
  themis_plot_phasespace_helistyle(self.vlsvReader, cellid, plane[0], plane[1],xlabel=labels[0],ylabel=labels[1])
+ pl.show()
 
  elif self.picker in self.picker_functions:
  # Call the corresponding pickerfunction