Magnetopause script updates

analysator/pyCalculations/fieldtracer.py

@@ -51,7 +51,7 @@ def static_field_tracer( vlsvReader, x0, max_iterations, dx, direction='+', bvar
    ''' Field tracer in a static frame
 
        :param vlsvReader:         An open vlsv file
-       :param x:                  Starting point for the field trace
+       :param x0:                  Starting point for the field trace
        :param max_iterations:     The maximum amount of iteractions before the algorithm stops
        :param dx:                 One iteration step length
        :param direction:          '+' or '-' or '+-' Follow field in the plus direction or minus direction
@@ -61,6 +61,10 @@ def static_field_tracer( vlsvReader, x0, max_iterations, dx, direction='+', bvar
        :returns:                  List of coordinates
    '''
 
+   # Check the starting point
+   if (np.shape(x0) != (3,)): 
+      raise ValueError("Starting point should be a single [x,y,z] coordinate in 1-dimensional list or array")
+
    if(bvar != 'B'):
      warnings.warn("User defined tracing variable detected. fg, volumetric variable results may not work as intended, use face-values instead.")
 
@@ -117,7 +121,7 @@ def static_field_tracer( vlsvReader, x0, max_iterations, dx, direction='+', bvar
       x = np.arange(mins[0], maxs[0], dcell[0]) + 0.5*dcell[0]
       y = np.arange(mins[1], maxs[1], dcell[1]) + 0.5*dcell[1]
       z = np.arange(mins[2], maxs[2], dcell[2]) + 0.5*dcell[2]
-      coordinates = np.array([x,y,z])
+      coordinates = [x,y,z]
       # Debug:
       if( len(x) != sizes[0] ):
          logging.info("SIZE WRONG: " + str(len(x)) + " " + str(sizes[0]))
@@ -142,7 +146,7 @@ def static_field_tracer( vlsvReader, x0, max_iterations, dx, direction='+', bvar
       x = np.arange(mins[0], maxs[0], dcell[0]) + 0.5*dcell[0]
       y = np.arange(mins[1], maxs[1], dcell[1]) + 0.5*dcell[1]
       z = np.arange(mins[2], maxs[2], dcell[2]) + 0.5*dcell[2]
-      coordinates = np.array([x,y,z])
+      coordinates = [x,y,z]
       # Debug:
       if( len(x) != sizes[0] ):
          logging.info("SIZE WRONG: " + str(len(x)) + " " + str(sizes[0]))
@@ -359,6 +363,10 @@ def static_field_tracer_3d( vlsvReader, seed_coords, max_iterations, dx, directi
    # Standardize input: (N,3) np.array
    if type(seed_coords) != np.ndarray:
       raise TypeError("Please give a numpy array.")
+
+   # Transform a single (3,) coordinate to (1,3) if needed
+   if np.shape(seed_coords)==(3,):
+      seed_coords = np.array([seed_coords])
 
    # Cache and read variables:
    vg = None

scripts/magnetopause2d.py

@@ -1,17 +1,20 @@
 '''
-Finds the magnetopause position by tracing steamines of the plasma flow for two-dimensional Vlasiator runs. Needs the yt package.
+Finds the magnetopause position by tracing streamines of the plasma flow for two-dimensional Vlasiator runs. Needs the yt package.
 '''
 
 import numpy as np
 import analysator as pt
-import plot_colormap
 import yt
 from yt.visualization.api import Streamlines
 
 
-
-
 def interpolate(streamline, x_points):
+    """Interpolates a single streamline for make_magnetopause().
+
+        :param streamline: a single streamline to be interpolated
+        :param x_points: points in the x-axis to use for interpolation
+        :returns: the streamline as numpy array of x,z coordinate points where the x-axis coordinates are the points given to the function
+    """
 
     arr = np.array(streamline)
 
@@ -26,6 +29,11 @@ def interpolate(streamline, x_points):
 
 
 def make_streamlines(vlsvFileName):
+    """Traces streamlines of velocity field from outside the magnetosphere to magnetotail using the yt-package.
+
+        :param vlsvFileName: directory and file name of vlsv file, to be given to VlsvReader
+        :returns: streamlines as numpy array
+    """
     ## make streamlines
     boxre = [0,0]
 
@@ -81,6 +89,11 @@ def to_Re(m): #meters to Re
     return np.array(streamlines_pos.streamlines)
 
 def make_magnetopause(streams):
+    """Finds the mangetopause location based on streamlines.
+
+        :param streams: streamlines
+        :returns: magnetopause as coordinate points from tail on positive x-axis to tail on negative x-axis
+    """
 
     streampoints = np.reshape(streams, (streams.shape[0]*streams.shape[1], 3)) #all the points in one array
 
@@ -126,6 +139,12 @@ def make_magnetopause(streams):
 
 
 def polar_to_cartesian(r, phi):
+    """Converts polar coordinates to cartesian.
+
+        :param r: radius
+        :param phi: angular coordinate in degrees
+        :returns: y, z -coordinates in cartesian system
+    """
     phi = np.deg2rad(phi)
     y = r * np.cos(phi)
     z = r * np.sin(phi)
@@ -166,7 +185,7 @@ def external_plot(ax,XmeshXY=None, YmeshXY=None, pass_maps=None):
                 ax.plot(magnetopause[:,0], magnetopause[:,1], color='cyan', linewidth=1.0) 
 
 
-    plot_colormap.plot_colormap(
+    pt.plot.plot_colormap(
         filename=fileLocation+fileN,
         outputdir=outdir,
         nooverwrite=None,

scripts/magnetopause3d.py

@@ -1,24 +1,17 @@
 '''
-Finds the magnetopause position by tracing steamines of the plasma flow for three-dimensional Vlasiator runs. Needs the yt package.
+Finds the magnetopause position by tracing streamlines of the plasma flow for three-dimensional Vlasiator runs.
 '''
-from pyCalculations import ids3d
+
 import matplotlib.pyplot as plt
 import numpy as np
 import analysator as pt
-import plot_colormap3dslice
-import yt
-import math
-from mpl_toolkits import mplot3d
-from yt.visualization.api import Streamlines
-
-
-
-
-def to_Re(m): #meters to Re
-    return m/6371000
-
 
 def cartesian_to_polar(cartesian_coords): # for segments of plane
+    """Converts cartesian coordinates to polar (for the segments of the yz-planes).
+
+        :param cartesian_coords: (y,z) coordinates as list or array
+        :returns: the polar coordinates r, phi (angle in degrees)
+    """
     y,z = cartesian_coords[0], cartesian_coords[1]
     r = np.sqrt(z**2 + y**2)
     phi = np.arctan2(z, y)
@@ -27,13 +20,25 @@ def cartesian_to_polar(cartesian_coords): # for segments of plane
     return(r, phi)
 
 def polar_to_cartesian(r, phi):
+    """Converts polar coordinates of the yz-plane to cartesian coordinates.
+
+        :param r: radius of the segment (distance from the x-axis)
+        :param phi: the angle coordinate in degrees
+        :returns:  y, z -coordinates in cartesian system
+    """
     phi = np.deg2rad(phi)
     y = r * np.cos(phi)
     z = r * np.sin(phi)
     return(y, z)
 
 
 def interpolate(streamline, x_points):
+    """IInterpolates a single streamline for make_magnetopause(). 
+
+        :param streamline: a single streamline to be interpolated
+        :param x_points: points in the x-axis to use for interpolation
+        :returns: the streamline as numpy array of coordinate points where the x-axis coordinates are the points given to the function
+    """
     arr = np.array(streamline)
 
     # set arrays for interpolation
@@ -53,32 +58,31 @@ def interpolate(streamline, x_points):
 
 
 def make_surface(coords):
+    '''Defines a surface constructed of input coordinates as triangles.
 
-    '''
-    Defines surface constructed of input coordinates as triangles
-    Returns list of verts and vert indices of surface triangles
-
-    coordinates must be in form [...[c11, c21, c31, ... cn1],[c12, c22, c32, ... cn2],...
-    where cij = [xij, yij, zij], i marks slice, j marks yz-plane (x_coord) index 
-
-    How it works:
-    Three points make a triangle, triangles make the surface.
-    For every two planes next to each other:
-    - take every other point from plane1, every other from plane2 (in order!)
-    - from list of points: every three points closest to each other make a surface
-
-    Example:
-    plane 1: [v1, v2, v3, v4]
-    plane 2: [v5, v6, v7, v8]
-
-    -> list: [v1, v5, v2, v6, v3,...]
-    -> triangles:
-    v1 v5 v2
-    v5 v2 v6
-    v2 v6 v3
-    .
-    .
-    .
+        :param coords: points that make the surface
+        :returns: list of verts and vert indices of surface triangles
+
+        input coordinates must be in form [...[c11, c21, c31, ... cn1],[c12, c22, c32, ... cn2],...
+        where cij = [xij, yij, zij], i marks sector, j marks yz-plane (x_coord) index 
+
+        Three points make a triangle, triangles make the surface.
+        For every two planes next to each other:
+        - take every other point from plane1, every other from plane2 (in order!)
+        - from list of points: every three points closest to each other make a surface
+
+        Example:
+        plane 1: [v1, v2, v3, v4]
+        plane 2: [v5, v6, v7, v8]
+
+        -> list: [v1, v5, v2, v6, v3,...]
+        -> triangles:
+        v1 v5 v2
+        v5 v2 v6
+        v2 v6 v3
+        .
+        .
+        .
 
     '''
     verts = [] #points
@@ -121,60 +125,19 @@ def make_surface(coords):
     return verts, faces
 
 
-
 def make_streamlines(vlsvFileName):
-    ## make streamlines
-    boxre = [0]
-
-    # bulk file
-    f = pt.vlsvfile.VlsvReader(file_name=vlsvFileName)
+    """Traces streamlines of velocity field from outside the magnetosphere to magnetotail.
 
-    #get box coordinates from data
-    [xmin, ymin, zmin, xmax, ymax, zmax] = f.get_spatial_mesh_extent()
-    [xsize, ysize, zsize] = f.get_spatial_mesh_size()
-    [xsizefs, ysizefs, zsizefs] = f.get_fsgrid_mesh_size()
-    simext_m =[xmin,xmax,ymin,ymax,zmin,zmax]
-    sizes = np.array([xsize,ysize,zsize])
-    sizesfs = np.array([xsizefs,ysizefs,zsizefs])
+        :param vlsvFileName: directory and file name of .vlsv data file to use for VlsvReader
+        :returns: streamlines as numpy array
+    """
 
-    simext = list(map(to_Re, simext_m))
-    boxcoords=list(simext)
-
-    cellids = f.read_variable("CellID")
-    indexids = cellids.argsort()
-    cellids = cellids[indexids]
-
-    reflevel = ids3d.refinement_level(xsize, ysize, zsize, cellids[-1])
-
-
-    #Read the data from vlsv-file
-    V = f.read_variable("vg_v")
-
-    #from m to Re
-    V = np.array([[to_Re(i) for i in j ]for j in V])
-
-
-    V = V[indexids]
-
-    if np.ndim(V)==1:
-        shape = None
-    elif np.ndim(V)==2: # vector variable
-        shape = V.shape[1]
-    elif np.ndim(V)==3:  # tensor variable
-        shape = (V.shape[1], V.shape[2])
-
-
-    Vdpoints = ids3d.idmesh3d2(cellids, V, reflevel, xsize, ysize, zsize, shape)
-
-
-    Vxs = Vdpoints[:,:,:,0]
-    Vys = Vdpoints[:,:,:,1]
-    Vzs = Vdpoints[:,:,:,2]
+    f = pt.vlsvfile.VlsvReader(file_name=vlsvFileName)
 
-    data=dict(Vx=Vxs,Vy=Vys,Vz=Vzs)
+    RE = 6371000
 
     #Create streamline seeds (starting points for streamlines)
-    seedN = 50 #seeds per row, final seed count will be seedN*seedN !
+    seedN = 25 #seeds per row, final seed count will be seedN*seedN !
     streamline_seeds = np.zeros([seedN**2, 3])
 
     #range: np.arange(from, to, step)
@@ -185,27 +148,27 @@ def make_streamlines(vlsvFileName):
             streamline_seeds[k] = [20, i, j]
             k = k+1
 
+    dx = 0.4
+    iterations = int(50/(dx)) # iterations so that final step is at max -30 RE (in practice less)
 
-    #dataset in yt-form
-    yt_dataset = yt.load_uniform_grid(
-        data,
-        sizesfs,
-        bbox=np.array([[boxcoords[0], boxcoords[1]],
-                    [boxcoords[2],boxcoords[3]],
-                    [boxcoords[4],boxcoords[5]]]))
-
+    streams = pt.calculations.static_field_tracer_3d(
+        vlsvReader=f,
+        seed_coords=streamline_seeds*RE,
+        max_iterations=iterations,
+        dx=dx*RE,
+        direction='+',
+        grid_var='vg_v')
 
-    # data, seeds, dictionary positions, length of streamlines
-    streamlines_pos = yt.visualization.api.Streamlines(yt_dataset, streamline_seeds,
-                                                    "Vx", "Vy", "Vz", length=50, direction=1)
+    return streams*(1/RE)
 
-    # make the streamlines
-    streamlines_pos.integrate_through_volume()
 
-    return np.array(streamlines_pos.streamlines)
+def make_magnetopause(streams):
+    """Finds the mangetopause location based on streamlines.
 
+        :param streams: streamlines
+        :returns: the magnetopause position as coordinate points in numpy array
+    """
 
-def make_magnetopause(streams):
     streampoints = np.reshape(streams, (streams.shape[0]*streams.shape[1], 3)) #all the points in one array
 
     ## find the subsolar dayside point in the x-axis
@@ -322,7 +285,7 @@ def external_plot(ax,XmeshXY=None, YmeshXY=None, pass_maps=None, requestvariable
             ax.plot(xz_slice[:,0], xz_slice[:,2],  color='limegreen', linewidth=1.5)
 
 
-        plot_colormap3dslice.plot_colormap3dslice(
+        pt.plot.plot_colormap3dslice(
         filename=fileLocation+fileN,
         outputdir=outdir,
         run=run,
@@ -342,7 +305,7 @@ def external_plot(ax,XmeshXY=None, YmeshXY=None, pass_maps=None, requestvariable
                 return ['vg_v']
             ax.plot(xy_slice[:,0], xy_slice[:,1], color='limegreen', linewidth=1.5)
 
-        plot_colormap3dslice.plot_colormap3dslice(
+        pt.plot.plot_colormap3dslice(
         filename=fileLocation+fileN,
         outputdir=outdir,
         run=run,