diff --git a/analysator/plot/__init__.py b/analysator/plot/__init__.py index e591a6e86..19a700f71 100644 --- a/analysator/plot/__init__.py +++ b/analysator/plot/__init__.py @@ -52,6 +52,7 @@ from .plot_themis_observation import themis_plot_detector,themis_plot_phasespace_contour,themis_plot_phasespace_helistyle,themis_observation,themis_observation_from_file from . import plot_helpers from .plot_colormap import plot_colormap +from .plot_cutthrough_timeseries import jplots from .plot_vdf import plot_vdf from .plot_vdfdiff import plot_vdfdiff from .plot_vdf_profiles import plot_vdf_profiles diff --git a/analysator/plot/plot_cutthrough_timeseries.py b/analysator/plot/plot_cutthrough_timeseries.py new file mode 100644 index 000000000..cdc4d01a8 --- /dev/null +++ b/analysator/plot/plot_cutthrough_timeseries.py @@ -0,0 +1,177 @@ + +import analysator as pt +import numpy as np +import matplotlib.pyplot as plt +from matplotlib.colors import BoundaryNorm,LogNorm +from matplotlib.ticker import MaxNLocator +from matplotlib.ticker import LogLocator +from scipy.ndimage import uniform_filter1d +import os +from analysator.calculations.lineout import lineout + +r_e = 6.371e6 #Maybe we should add an import that holds all the constants? + +def jplots( + var, + fnr1, + fnr2, + bulkpath, + bulkprefix, + point1, + point2, + outputname, + outputdir=None, + title=None, + lin=True, + vmin=None, + vmax=None, + filt=-1, + op="pass", + cmap="viridis", + re=False, + npoints=100, + interpolation_order=1, + draw=False, + nooverwrite=False +): + ''' + Function for plotting a cut-through timeseries, aka keogram, aka time-elongation map, etc. for a given variable, a given set of coordinates, and a range of times. + + :kwarg var: Variable to plot (e.g. "proton/vg_rho", "vg_b_vol") + :kwarg fnr1: First file number to plot + :kwarg fnr2: Last file number to plot + :kwarg bulkpath: Path to directory containing bulk files + :kwarg bulkprefix: Starting string of bulk file name (e.g. bulk, bulk1, bulk5) + :kwarg point1: First point on the line to plot + :kwarg point2: Last point on the line to plot + :kwarg lin: Whether to use linear or logarithmic scaling? (True: Use linear with 7 ticks, False: Logarithmic, int: Number of ticks for linear case ) + :kwarg vmin: Min value for the colormap scaling. + :kwarg vmax: Max value for the colormap scaling. + :kwarg title: Custom title, default None which then uses title=var. + :kwarg outputname: Name of output file + :kwarg outputdir: Path to output directory + :kwarg filt: Filter out temporally slowly changing signal? (<=0: no filtering, >0: filter with specified window size), default=-1 + :kwarg op: Variable operator, default="pass" + :kwarg cmap: Colormap, default="viridis" + :kwarg re: Input points given in Earth radii? default=False + :kwarg npoints: Number of points in line, default=100 + :kwarg interpolation_order: Order of interpolation (0 or 1), default=1 + :kward nooverwrite: Whether to overwrite if output file already exists, default=False + :kwarg draw: Whether draw on screen or output to file, default=False + + :returns: Outputs an image to a file or to the screen. + + .. code-block:: python + + # Example usage: + jplots(var="vg_b_vol",fnr1=600,fnr2=650,bulkpath='/wrk-vakka/group/spacephysics/vlasiator/3D/FIF/bulk1',bulkprefix='bulk1', + point1=[10,-20,0],point2=[20,-20,0],outputdir='./',outputname='test.png',op='x',re=True,npoints=50) + ''' + + + fnr_arr = np.arange(fnr1, fnr2 + 0.1, 1, dtype=int) + t_arr = np.zeros_like(fnr_arr).astype(float) + + fobj = pt.vlsvfile.VlsvReader( + os.path.join(bulkpath , bulkprefix + ".{}.vlsv".format(str(fnr1).zfill(7))) + ) + if re: + point1 = [point1[0] * r_e, point1[1] * r_e, point1[2] * r_e] + point2 = [point2[0] * r_e, point2[1] * r_e, point2[2] * r_e] + + lineout0 = lineout( + fobj, + point1, + point2, + variable=var, + operator=op, + points=npoints, + interpolation_order=interpolation_order, + ) + distances, _, _ = lineout0 + + fobj.optimize_clear_fileindex_for_cellid() + + distances = np.array(distances) / r_e + + data_arr = np.zeros((fnr_arr.size, distances.size), dtype=float) + + for idx in range(fnr_arr.size): + fnr = fnr_arr[idx] + vlsvobj = pt.vlsvfile.VlsvReader( + os.path.join(bulkpath , bulkprefix + ".{}.vlsv".format(str(fnr).zfill(7))) + ) + t_arr[idx] = vlsvobj.read_parameter("time") + + linecut = lineout( + vlsvobj, + point1, + point2, + variable=var, + operator=op, + points=npoints, + interpolation_order=interpolation_order, + ) + data_arr[idx, :] = linecut[2] + vlsvobj.optimize_clear_fileindex_for_cellid() + if filt > 0: + data_arr = data_arr - uniform_filter1d(data_arr, size=filt, axis=0) + + XmeshXY, YmeshXY = np.meshgrid(distances, t_arr) + + fig, ax = plt.subplots(1, 1, figsize=(8, 12), constrained_layout=True) + + if not lin and lin != 0: + lin = None + if lin is True and lin!=1: + lin = 7 + + if vmin is not None: + vminuse=vmin + else: + vminuse=np.ma.amin(data_arr) + if vmax is not None: + vmaxuse=vmax + else: + vmaxuse=np.ma.amax(data_arr) + + # If both values are zero, we have an empty array + if vmaxuse==vminuse==0: + raise ValueError("requested array is zero everywhere. Exiting.") + + cmapuse = pt.plot.get_cmap(cmap) + + if not lin: + norm = LogNorm(vmin=vminuse,vmax=vmaxuse) + ticks = LogLocator(base=10,subs=list(range(10))) # where to show labels + else: + levels = MaxNLocator(nbins=255).tick_values(vminuse,vmaxuse) + norm = BoundaryNorm(levels, ncolors=cmapuse.N, clip=True) + ticks = np.linspace(vminuse,vmaxuse,num=lin) + + im = ax.pcolormesh( + XmeshXY, + YmeshXY, + data_arr, + shading="gouraud", + cmap=cmap, + norm=norm, + rasterized=True, + ) + + ax.set_xlim(distances[0], distances[-1]) + ax.set_ylim(t_arr[0], t_arr[-1]) + ax.set_xlabel("Distance along cut [RE]", labelpad=10, fontsize=16) + ax.set_ylabel("Time [s]", labelpad=10, fontsize=16) + ax.set_title(title if title else var, pad=10, fontsize=16) + + cb = fig.colorbar(im, ax=ax,ticks=ticks) + + if not draw: + outputpath=pt.plot.output_path(outputname,None,outputdir,nooverwrite) + fig.savefig(outputpath, dpi=300) + plt.close(fig) + else: + return (fig,ax,XmeshXY,YmeshXY,data_arr) + + diff --git a/scripts/cutthrough_timeseries.py b/scripts/cutthrough_timeseries.py index ff9c98ae3..2fb278d35 100644 --- a/scripts/cutthrough_timeseries.py +++ b/scripts/cutthrough_timeseries.py @@ -58,11 +58,8 @@ def jplots( fnr_arr = np.arange(fnr1, fnr2 + 0.1, 1, dtype=int) t_arr = np.zeros_like(fnr_arr).astype(float) - if bulkpath[-1] != "/": - bulkpath += "/" - fobj = pt.vlsvfile.VlsvReader( - bulkpath + bulkprefix + ".{}.vlsv".format(str(fnr1).zfill(7)) + os.path.join(bulkpath , bulkprefix + ".{}.vlsv".format(str(fnr1).zfill(7))) ) if re: point1 = [point1[0] * r_e, point1[1] * r_e, point1[2] * r_e] @@ -88,7 +85,7 @@ def jplots( for idx in range(fnr_arr.size): fnr = fnr_arr[idx] vlsvobj = pt.vlsvfile.VlsvReader( - bulkpath + bulkprefix + ".{}.vlsv".format(str(fnr).zfill(7)) + os.path.join(bulkpath , bulkprefix + ".{}.vlsv".format(str(fnr).zfill(7))) ) t_arr[idx] = vlsvobj.read_parameter("time") @@ -134,10 +131,7 @@ def jplots( except OSError: pass - if outputdir[-1] != "/": - outputdir += "/" - - fig.savefig(outputdir + outputname, dpi=300) + fig.savefig(os.path.join(outputdir,outputname), dpi=300) plt.close(fig) else: return (fig,ax,XmeshXY,YmeshXY,data_arr) diff --git a/testpackage/testpackage_commons.py b/testpackage/testpackage_commons.py index 96babdb1b..0bd955fdf 100644 --- a/testpackage/testpackage_commons.py +++ b/testpackage/testpackage_commons.py @@ -6,6 +6,7 @@ def source_file_name(filename,fileLocation,time): if filename is None: + #This may not produce correct result but works currently if '2D' not in fileLocation: bulkname = "bulk1."+str(time).rjust(7,'0')+".vlsv" @@ -65,12 +66,27 @@ def source_file_name(filename,fileLocation,time): 'filename': None , 'cavitonparams': [6.6e6,2.64e6,4.e-9,10] }) +#For now this has just the jplots run, since that uses multiple files a method of forming the call arugments should be made +#since what is defined in the definition file only works for FIF (or runs with 501-531 times) +runs.append( { 'name': 'FIF', + 'verifydir': '/FIF/', + 'fileLocation': "/wrk-vakka/group/spacephysics/vlasiator/3D/FIF/bulk1", + 'fluxLocation': None, + 'funcs': [{'plot_cutthrough_timeseries':"jplots"}], + 'pops': ['avgs'], + 'time': 0, + 'skipped_args': None, + 'singletime': True, + 'filename': None, #restart file + 'nosubpops': True, # backstreaming / non-backstreaming + 'vlasiator5': True, + 'cavitonparams': [6.6e6,2.64e6,4.e-9,10] } ) runs.append( { 'name': 'FID', 'verifydir': '/FID/', 'fileLocation': datalocation+'/3D/FID/bulk1/', 'fluxLocation': None, - 'funcs': ['plot_colormap3dslice','plot_ionosphere','plot_isosurface'], + 'funcs': ['plot_colormap3dslice','plot_ionosphere','plot_isosurface'], 'pops': ['avgs'], 'time': 1000, 'skipped_args':{'plot_ionosphere':{"var":["ig_z","ig_p","ig_source","ig_residual"]}, @@ -215,7 +231,13 @@ def source_file_name(filename,fileLocation,time): functions = run['funcs'] else: - functions = sorted(list(set(run['funcs']) & set(funcs_to_use))) + functions=[] + for func in run['funcs']: + func_key=list(func.keys())[0] if type(func) is dict else func + if func_key in funcs_to_use: + functions.append(func) + + # functions = sorted(list(set(run['funcs']) & set(funcs_to_use))) if not functions: continue @@ -223,7 +245,11 @@ def source_file_name(filename,fileLocation,time): for j,func in enumerate(functions): #If one wants to test new calls to be added to the testpackage, they could be added in here after the testpackage_definitions/ lists have been parsed callindex = 0 - + func_dict=None + if type(func) is dict: + func_dict=func + func=list(func_dict.keys())[0] + #try to import the list of calls corresponding to the function to be tested. try: exec(f'import testpackage_definitions.testpackage_{func} as testpackage_{func}') @@ -240,7 +266,6 @@ def source_file_name(filename,fileLocation,time): skipped_args=run['skipped_args'] - if filename is not None: calls_in=v5restartcalls if vlasiator5 else restartcalls list_index=0 @@ -250,6 +275,8 @@ def source_file_name(filename,fileLocation,time): else: call = call.replace("var='V'","var='restart_V'") if skipped_args: + if func_dict: + func=func_dict[func] call=call_replace(call,func,skipped_args,required_args) if call is not None: callrunids.append(i) @@ -287,6 +314,8 @@ def source_file_name(filename,fileLocation,time): elif (("_backstream" in call) or ("_nonbackstream" in call)) and nosubpops: continue if skipped_args: + if func_dict: + func=func_dict[func] call=call_replace(call,func,skipped_args,required_args) if call is not None: list_inidices.append((1,list_index)) @@ -318,6 +347,8 @@ def source_file_name(filename,fileLocation,time): continue call = call.replace('REPLACEPOP',pop) if skipped_args: + if func_dict: + func=func_dict[func] call=call_replace(call,func,skipped_args,required_args) if call is not None: list_inidices.append((2,list_index)) @@ -361,13 +392,24 @@ def source_file_name(filename,fileLocation,time): list_index=list_inidices[j] funcid=funcids[j] - + #This could also be skipped since it's also done earlier but that would mean making new variable, checking it works blabla, im too lazy rn if not funcs_to_use: func = runs[runid]['funcs'][funcid] + if type(func) is dict: + func_dict=func + func=list(func_dict.keys())[0] else: - func = sorted(list(set(runs[runid]['funcs']) & set(funcs_to_use)))[funcid] + functions=[] + for func in runs[runid]['funcs']: + if type(func) is dict: + func=list(func.keys())[0] + if func in funcs_to_use: + functions.append(func) + if functions: + func=functions[funcid] + # func = sorted(list(set(runs[runid]['funcs']) & set(funcs_to_use)))[funcid] if not func: continue @@ -380,7 +422,7 @@ def source_file_name(filename,fileLocation,time): filename = runs[runid]['filename'] vlasiator5 = runs[runid]['vlasiator5'] singletime = runs[runid]['singletime'] - + #set custom expression variables for plot_colormap if 'plot_colormap' in func: exec(f"testpackage_{func}.cexp.level_bow_shock = runs[runid]['cavitonparams'][0]") @@ -445,9 +487,13 @@ def source_file_name(filename,fileLocation,time): # Many different plots + + if "jplots" in call: + bulkname="" + call=call.replace("DATALOCATION",fileLocation) + else: + f = pt.vlsvfile.VlsvReader(fileLocation+bulkname) print(j, runid, jrun, call,fileLocation+bulkname) - f = pt.vlsvfile.VlsvReader(fileLocation+bulkname) - try: exec(call) except Exception as e: diff --git a/testpackage/testpackage_definitions/testpackage_plot_cutthrough_timeseries.py b/testpackage/testpackage_definitions/testpackage_plot_cutthrough_timeseries.py new file mode 100644 index 000000000..1b1e1fc25 --- /dev/null +++ b/testpackage/testpackage_definitions/testpackage_plot_cutthrough_timeseries.py @@ -0,0 +1,5 @@ +v5nonrestartcalls=[ + + # 'pt.plot.jplots(var="vg_b_vol",fnr1=600,fnr2=650,bulkpath="/wrk-vakka/group/spacephysics/vlasiator/3D/FIF/bulk1",bulkprefix="bulk1",point1=[10,-20,0],point2=[20,-20,0],outputdir="verifydir+REPLACEINDEX",outputname="plot_jplots_REPLACEINDEX.png",op="x",re=True,npoints=50)', + "pt.plot.jplots(var='vg_b_vol',fnr1=501,fnr2=531,bulkpath='DATALOCATION',bulkprefix='bulk1',point1=[10,-20,0],point2=[20,-20,0],outputdir=outputLocation,outputname=REPLACEINDEX+'_plot_jplots.png',op='x',re=True,npoints=30)", +]