Merge branch 'master' of github.com:fmihpc/analysator

pyCalculations/calculations.py

@@ -39,8 +39,9 @@
 # List of functions and classes that should be imported into the interface
 from intpol_file import vlsv_intpol_file
 from intpol_points import vlsv_intpol_points
-from cutthrough import cut_through, cut_through_step
+from cutthrough import cut_through, cut_through_step, cut_through_curve, cut_through_swath
 from fourier import fourier
+from spectra import get_spectrum_energy, get_spectrum_alongaxis_vel
 from variable import VariableInfo
 from timeevolution import cell_time_evolution
 from pitchangle import pitch_angles

pyCalculations/cutthrough.py

@@ -32,7 +32,7 @@ def get_cellids_coordinates_distances( vlsvReader, xmax, xmin, xcells, ymax, ymi
        :type vlsvReader:        :class:`vlsvfile.VlsvReader`
        :param point1:           The starting point of a cut-through line
        :param point2:           The ending point of a cut-through line
-       :returns: Coordinates of for the cell cut-through as well as cell ids and distances
+       :returns: Coordinates of for the cell cut-through as well as cell ids and distances. NB: Last CID is a duplicate.
    '''
    point1 = np.array(point1, copy=False)
    point2 = np.array(point2, copy=False)
@@ -47,6 +47,13 @@ def get_cellids_coordinates_distances( vlsvReader, xmax, xmin, xcells, ymax, ymi
    epsilon = sys.float_info.epsilon*2
 
    iterator = point1
+
+   # Special case: both points in the same cell. 
+   #if(vlsvReader.get_cellid(point1) == vlsvReader.get_cellid(point2)):
+   #   cellids.append(vlsvReader.get_cellid(point2))
+
+
+
    unit_vector = (point2 - point1) / np.linalg.norm(point2 - point1 + epsilon)
    while True:
       # Get the cell id
@@ -58,9 +65,8 @@ def get_cellids_coordinates_distances( vlsvReader, xmax, xmin, xcells, ymax, ymi
       min_bounds = vlsvReader.get_cell_coordinates(cellid) - 0.5 * cell_lengths
       max_bounds = np.array([min_bounds[i] + cell_lengths[i] for i in range(0,3)])
       # Check which boundary we hit first when we move in the unit_vector direction:
-
-      coefficients_min = np.divide((min_bounds - iterator), unit_vector)
-      coefficients_max = np.divide((max_bounds - iterator) , unit_vector)
+      coefficients_min = np.divide((min_bounds - iterator), unit_vector, out=np.full(min_bounds.shape, sys.float_info.max), where=np.logical_not(unit_vector==0.0))
+      coefficients_max = np.divide((max_bounds - iterator), unit_vector, out=np.full(min_bounds.shape, sys.float_info.max), where=np.logical_not(unit_vector==0.0))
       # Remove negative coefficients:
       for i in range(3):
          if coefficients_min[i] <= 0:
@@ -71,10 +77,15 @@ def get_cellids_coordinates_distances( vlsvReader, xmax, xmin, xcells, ymax, ymi
 
 
       # Find the minimum coefficient for calculating the minimum distance from a boundary
-      coefficient = min([min(coefficients_min), min(coefficients_max)]) * 1.01
+      coefficient = min([min(coefficients_min), min(coefficients_max)]) * 1.000001
 
       # Get the cell id in the new coordinates
       newcoordinate = iterator + coefficient * unit_vector
+
+      # Check to make sure the iterator is not moving past point2 - if it is, use point2 instead:
+      if min((point2 - newcoordinate)* unit_vector) < 0:
+         newcoordinate = point2
+
       newcellid = vlsvReader.get_cellid( newcoordinate )
       # Check if valid cell id:
       if newcellid == 0:
@@ -107,7 +118,7 @@ def cut_through( vlsvReader, point1, point2 ):
        :type vlsvReader:        :class:`vlsvfile.VlsvReader`
        :param point1:           The starting point of a cut-through line
        :param point2:           The ending point of a cut-through line
-       :returns: an array containing cell ids, coordinates and distances in the following format: [cell ids, distances, coordinates]
+       :returns: an array containing cell ids, coordinates and distances in the following format: [cell ids, distances, coordinates]. NB. Last cellid is a duplicate.
 
        .. code-block:: python
 
@@ -151,11 +162,33 @@ def cut_through( vlsvReader, point1, point2 ):
    return get_cellids_coordinates_distances( vlsvReader, xmax, xmin, xcells, ymax, ymin, ycells, zmax, zmin, zcells, cell_lengths, point1, point2 )
 
 
+def cut_through_swath(vlsvReader, point1, point2, width, normal):
 
+   init_cut = cut_through(vlsvReader, point1, point2)
+   init_cids = init_cut[0].data
+
+   print('swath initial CellIds', init_cids)
+
+   #find the other vector spanning the swath
+   s = np.array(point2)-np.array(point1)
+   w = np.cross(s, np.array(normal))
+   w = width*w/np.linalg.norm(w)
+
+   out_cids = []
+   for cid in init_cids:
+      mid = vlsvReader.get_cell_coordinates(cid)
+      p1 = mid - w/2
+      p2 = mid + w/2
+      temp_cut = cut_through_step(vlsvReader, p1, p2)
+      out_cids.append(temp_cut[0].data)
+
+   init_cut[0] = np.array(out_cids)
+   print(init_cut[0])
+   return init_cut
 
 def cut_through_step( vlsvReader, point1, point2 ):
    ''' Returns cell ids and distances from point 1 to point 2, returning not every cell in a line
-       but rather the amount of cells which correspons with the largest axis-aligned component of the line.
+       but rather the amount of cells which corresponds with the largest axis-aligned component of the line.
 
        :param vlsvReader:       Some open VlsvReader
        :type vlsvReader:        :class:`vlsvfile.VlsvReader`
@@ -228,4 +261,100 @@ def cut_through_step( vlsvReader, point1, point2 ):
    # Return the coordinates, cellids and distances for processing
    from output import output_1d
    return output_1d( [np.array(cellids, copy=False), np.array(distances, copy=False), np.array(coordinates, copy=False)], ["CellID", "distances", "coordinates"], ["", "m", "m"] )
-
+
+
+# Take a curve (list of 3-coords), return cells and distances along curve as with cut_through.
+# NB: calls get_cells_coordinates_distances for each curve segment. Feel free to optimize your curves beforehand.
+def cut_through_curve(vlsvReader, curve):
+   ''' Returns cell ids and distances from point 1 for every cell in a line between given point1 and point2
+
+       :param vlsvReader:       Some open VlsvReader
+       :type vlsvReader:        :class:`vlsvfile.VlsvReader`
+       :param point1:           The starting point of a cut-through line
+       :returns: an array containing cell ids, edge distances, and the coordinates of edges in the following format: [cell ids, distances, coordinates]. NB. Last cellid is a duplicate.
+
+       .. code-block:: python
+
+          Example:
+          vlsvReader = VlsvReader(\"testfile.vlsv\")
+          cut_through_curve = cut_through(vlsvReader, [[0,0,0], [2,5e6,0]])
+          cellids = cut_through[0]
+          distances = cut_through[1]
+          print \"Cell ids: \" + str(cellids)
+          print \"Distance from point 1 for every cell: \" + str(distances)
+   '''
+   # init cut_through static values, then do what cut_through does for each segment
+   # Get parameters from the file to determine a good length between points (step length):
+   # Get xmax, xmin and xcells_ini
+   mesh_limits = vlsvReader.get_spatial_mesh_extent()
+   mesh_size = vlsvReader.get_spatial_mesh_size()
+   xmax = mesh_limits[3]
+   xmin = mesh_limits[0]
+   xcells = mesh_size[0]
+   # Do the same for y
+   ymax = mesh_limits[4]
+   ymin = mesh_limits[1]
+   ycells = mesh_size[1]
+   # And for z
+   zmax = mesh_limits[5]
+   zmin = mesh_limits[2]
+   zcells = mesh_size[2]
+
+
+   #Calculate cell lengths:
+   cell_lengths = np.array([(xmax - xmin)/(float)(xcells), (ymax - ymin)/(float)(ycells), (zmax - zmin)/(float)(zcells)])
+
+   if(len(curve) < 2):
+      print("ERROR, less than 2 points in a curve")
+      return None
+
+   cellIds = []
+   edges = []
+   coords = []
+   for i in range(len(curve)-1):
+     point1 = np.array(curve[i])
+     point2 = np.array(curve[i+1])
+     # Make sure point1 and point2 are inside bounds
+     if vlsvReader.get_cellid(point1) == 0:
+       print("ERROR, POINT1 IN CUT-THROUGH-CURVE OUT OF BOUNDS!")
+     if vlsvReader.get_cellid(point2) == 0:
+       print("ERROR, POINT2 IN CUT-THROUGH-CURVE OUT OF BOUNDS!")
+     cut = get_cellids_coordinates_distances( vlsvReader, xmax, xmin, xcells, ymax, ymin, ycells, zmax, zmin, zcells, cell_lengths, point1, point2)
+     ccid = cut[0].data
+     cedges = cut[1].data
+     try:
+       edgestart = edges[-1]
+     except:
+       edgestart = 0
+     for ci,citer in enumerate(ccid):
+       cellIds.append(citer)
+       edges.append(edgestart+cedges[ci])
+       coords.append(cut[2].data[ci])
+
+   #print('sum of diffs', np.sum(np.diff(edges)))
+   #reduce the cellIDs and edges
+   reduct = True
+   if reduct:
+     rCellIds = [cellIds[0]]
+     rEdges = [edges[0]]
+     rCoords = [coords[0]]
+     cid0 = cellIds[0]
+     for i in range(1, len(cellIds)-1):
+       c1 = cellIds[i]
+       if(cid0 == c1):
+         rEdges[-1] = edges[i]
+         rCoords[-1] = coords[i]
+         continue
+       rCellIds.append(c1)
+       rEdges.append(edges[i])
+       rCoords.append(coords[i])
+       cid0 = c1
+     rEdges.append(edges[-1])
+     rCoords.append(coords[-1])
+     #print('sum of r-diffs', np.sum(np.diff(rEdges)))
+   else:
+     rCellIds = cellIds
+     rEdges = edges
+     rCoords = coords
+   from output import output_1d
+   return output_1d( [np.array(rCellIds, copy=False), np.array(rEdges, copy=False), np.array(rCoords, copy=False)], ["CellID", "distances", "coordinates"], ["", "m", "m"] )

pyCalculations/fieldtracer.py

@@ -46,7 +46,7 @@ def dynamic_field_tracer( vlsvReader_list, x0, max_iterations, dx):
       x0 = x0 + v*dt
       iterations = iterations + 1
 
-def static_field_tracer( vlsvReader, x0, max_iterations, dx, direction='+', bvar='B' ):
+def static_field_tracer( vlsvReader, x0, max_iterations, dx, direction='+', bvar='B', centering='default', boundary_inner=-1):
    ''' Field tracer in a static frame
 
        :param vlsvReader:         An open vlsv file
@@ -55,10 +55,12 @@ def static_field_tracer( vlsvReader, x0, max_iterations, dx, direction='+', bvar
        :param dx:                 One iteration step length
        :param direction:          '+' or '-' or '+-' Follow field in the plus direction or minus direction
        :param bvar:               String, variable name to trace [default 'B']
+       :param centering:          String, variable centering: 'face', 'volume', 'node' [defaults to 'face']
+       :param boundary_inner:     Float, stop propagation if closer to origin than this value [default -1]
        :returns:                  List of coordinates
    '''
 
-   if(bvar is not 'B'):
+   if(bvar != 'B'):
      warnings.warn("User defined tracing variable detected. fg, volumetric variable results may not work as intended, use face-values instead.")
 
    if direction == '+-':
@@ -93,46 +95,85 @@ def static_field_tracer( vlsvReader, x0, max_iterations, dx, direction='+', bvar
    if zsize <= 1:
       indices = [1,0]
 
+   if 'vol' in bvar and centering == 'default':
+      warnings.warn("Found 'vol' in variable name, assuming volumetric variable and adjusting centering")
+      centering = 'volume'
+
    # Read face_B:
-   face_B = f.read_variable(bvar)
-   face_Bx = face_B[:,0]
-   face_By = face_B[:,1]
-   face_Bz = face_B[:,2]
-
-   face_Bx = face_Bx[cellids.argsort()].reshape(sizes[indices])
-   face_By = face_By[cellids.argsort()].reshape(sizes[indices])
-   face_Bz = face_Bz[cellids.argsort()].reshape(sizes[indices])
-
-   face_B = np.array([face_Bx, face_By, face_Bz])
-
-   # Create x, y, and z coordinates:
-   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])
-   # Debug:
-   if( len(x) != sizes[0] ):
-      print("SIZE WRONG: " + str(len(x)) + " " + str(sizes[0]))
-
-   # Create grid interpolation
-   interpolator_face_B_0 = interpolate.RectBivariateSpline(coordinates[indices[0]] - 0.5*dcell[indices[0]], coordinates[indices[1]], face_B[indices[0]], kx=2, ky=2, s=0)
-   interpolator_face_B_1 = interpolate.RectBivariateSpline(coordinates[indices[0]], coordinates[indices[1]] - 0.5*dcell[indices[1]], face_B[indices[1]], kx=2, ky=2, s=0)
-   interpolators = [interpolator_face_B_0, interpolator_face_B_1]#, interpolator_face_B_2]
+   if centering == 'face' or centering == 'default':
+      face_B = f.read_variable(bvar)
+      face_Bx = face_B[:,0]
+      face_By = face_B[:,1]
+      face_Bz = face_B[:,2]
+
+      face_Bx = face_Bx[cellids.argsort()].reshape(sizes[indices])
+      face_By = face_By[cellids.argsort()].reshape(sizes[indices])
+      face_Bz = face_Bz[cellids.argsort()].reshape(sizes[indices])
+
+      face_B = np.array([face_Bx, face_By, face_Bz])
+
+      # Create x, y, and z coordinates:
+      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])
+      # Debug:
+      if( len(x) != sizes[0] ):
+         print("SIZE WRONG: " + str(len(x)) + " " + str(sizes[0]))
+
+      # Create grid interpolation
+      interpolator_face_B_0 = interpolate.RectBivariateSpline(coordinates[indices[0]] - 0.5*dcell[indices[0]], coordinates[indices[1]], face_B[indices[0]], kx=2, ky=2, s=0)
+      interpolator_face_B_1 = interpolate.RectBivariateSpline(coordinates[indices[0]], coordinates[indices[1]] - 0.5*dcell[indices[1]], face_B[indices[1]], kx=2, ky=2, s=0)
+      interpolators = [interpolator_face_B_0, interpolator_face_B_1]#, interpolator_face_B_2]
+   elif centering == 'volume':
+      vol_B = f.read_variable(bvar)
+      vol_Bx = vol_B[:,0]
+      vol_By = vol_B[:,1]
+      vol_Bz = vol_B[:,2]
+
+      vol_Bx = vol_Bx[cellids.argsort()].reshape(sizes[indices])
+      vol_By = vol_By[cellids.argsort()].reshape(sizes[indices])
+      vol_Bz = vol_Bz[cellids.argsort()].reshape(sizes[indices])
+
+      vol_B = np.array([vol_Bx, vol_By, vol_Bz])
+
+      # Create x, y, and z coordinates:
+      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])
+      # Debug:
+      if( len(x) != sizes[0] ):
+         print("SIZE WRONG: " + str(len(x)) + " " + str(sizes[0]))
+
+      # Create grid interpolation
+      interpolator_vol_B_0 = interpolate.RectBivariateSpline(coordinates[indices[0]], coordinates[indices[1]], vol_B[indices[0]], kx=2, ky=2, s=0)
+      interpolator_vol_B_1 = interpolate.RectBivariateSpline(coordinates[indices[0]], coordinates[indices[1]], vol_B[indices[1]], kx=2, ky=2, s=0)
+      interpolators = [interpolator_vol_B_0, interpolator_vol_B_1]#, interpolator_face_B_2]
+   elif centering == 'node':
+      print("Nodal variables not implemented")
+      return
+   else:
+      print("Unrecognized centering:", centering)
+      return
 
    #######################################################
    if direction == '-':
       multiplier = -1
    else:
       multiplier = 1
 
-   points = [x0]
+   points = [np.array(x0)]
    for i in range(max_iterations):
-      previous_point = points[len(points)-2]
+      previous_point = points[-1]
       B_unit = np.zeros(3)
       B_unit[indices[0]] = interpolators[0](previous_point[indices[0]], previous_point[indices[1]])
       B_unit[indices[1]] = interpolators[1](previous_point[indices[0]], previous_point[indices[1]])
       B_unit = B_unit / float(np.linalg.norm(B_unit))
-      points.append( previous_point + multiplier*B_unit * dx )
+      next_point = previous_point + multiplier*B_unit * dx
+      if(np.linalg.norm(next_point) < boundary_inner):
+         break
+      points.append(next_point)
    #######################################################
 
    return points