Adding the jplots from cutthrough_timeseries script to analysator plot module

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

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)
+
+

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)

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,15 +231,25 @@ 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
 
 
     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:

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)",
+]