diff --git a/analysator/calculations/magnetopause_sw_streamline_2d.py b/analysator/calculations/magnetopause_sw_streamline_2d.py index 5b37d10d..6a5df8af 100644 --- a/analysator/calculations/magnetopause_sw_streamline_2d.py +++ b/analysator/calculations/magnetopause_sw_streamline_2d.py @@ -1,10 +1,13 @@ ''' -Finds the magnetopause position by tracing streamlines of the plasma flow for two-dimensional Vlasiator runs. Needs the yt package. +Finds the magnetopause position by tracing streamlines of the plasma flow for two-dimensional Vlasiator runs. + +Streamlines are integrated with scipy using RegularGridInterpolator for the velocity field and solve_ivp for the ODE integration ''' import numpy as np +from scipy.interpolate import RegularGridInterpolator +from scipy.integrate import solve_ivp import analysator as pt -import yt def interpolate(streamline, x_points): @@ -17,75 +20,134 @@ def interpolate(streamline, x_points): arr = np.array(streamline) + # Drop points beyond which the streamline has no data + valid = np.isfinite(arr[:, 0]) & np.isfinite(arr[:, 2]) + arr = arr[valid] + + x_points = np.asarray(x_points, dtype=float) + if arr.shape[0] == 0: + return np.array([x_points, np.full_like(x_points, np.NaN)]) + # set arrays for interpolation xp = arr[:,0][::-1] zp = arr[:,2][::-1] - #interpolate z coordinates + # interpolate z coordinates z_points = np.interp(x_points, xp, zp, left=np.NaN, right=np.NaN) return np.array([x_points, z_points]) -def make_streamlines(vlsvfile, streamline_seeds=None,seeds_n=200, seeds_x0=20*6371000, seeds_range=[-5*6371000, 5*6371000], streamline_length=40*6371000): - """Traces streamlines of velocity field using the yt package. +def trace_streamline(x0, z0, vx_interp, vz_interp, bounds, length, max_step, n_points, rtol, atol, min_speed): + """Integrates one single streamline forward starting at (x0, z0). + + :param x0, z0: starting point, m + :param vx_interp, vz_interp: RegularGridInterpolator for the two velocity components (must return NaN outside their grid) + :param bounds: (xmin, xmax, zmin, zmax) -- valid range of the interpolators, m + :param length: maximum arc length to integrate, m + :param max_step: maximum solve_ivp step size, m + :param n_points: number of points to sample along the streamline + :param rtol, atol: solve_ivp error tolerances + :param min_speed: velocity magnitude (m/s) below which the local flow direction is treated as undefined -- guards against 0/0 in regions of exactly zero velocity (e.g. an inner boundary) + + :returns: (n_points, 3) array of x, y(=0), z points. Points beyond where the streamline left the domain are NaN. + """ + xmin, xmax, zmin, zmax = bounds + + def rhs(s, state): + x, z = state + vx = float(vx_interp((x, z))) + vz = float(vz_interp((x, z))) + speed = np.hypot(vx, vz) + if not np.isfinite(speed) or speed < min_speed: + return (0.0, 0.0) + return (vx / speed, vz / speed) + + def exit_domain(s, state): + x, z = state + return min(x-xmin, xmax-x, z-zmin, zmax-z) + exit_domain.terminal = True + exit_domain.direction = -1 + + sol = solve_ivp(rhs, [0.0, length], [x0, z0], events=exit_domain, dense_output=True, max_step=max_step, rtol=rtol, atol=atol) + + s_eval = np.linspace(0.0, length, n_points) + s_max = sol.t[-1] + + out = np.full((n_points, 3), np.nan) + valid = s_eval <= s_max + pts = sol.sol(s_eval[valid]) + out[valid, 0] = pts[0] + out[valid, 1] = 0.0 + out[valid, 2] = pts[1] + return out + + +def make_streamlines(vlsvfile, streamline_seeds=None, seeds_n=200, seeds_x0=20*6371000, seeds_range=(-5*6371000, 5*6371000), streamline_length=40*6371000, *, max_step=None, n_points=400, rtol=1e-6, atol=1.0, min_speed=1e-3): + """Traces streamlines of the velocity field. :param vlsvfile: directory and file name of .vlsv data file to use for VlsvReader :kwarg streamline_seeds: optional streamline starting points in numpy array (coordinates in meters including the y-coordinate 0.0) - :kwarg seeds_n: instead of streamline_seeds provide a number of streamlines to be traced - :kwarg seeds_x0: instead of streamline_seeds provide an x-coordinate for streamline starting points + :kwarg seeds_n: instead of streamline_seeds provide a number of streamlines to be traced + :kwarg seeds_x0: instead of streamline_seeds provide an x-coordinate for streamline starting points :kwarg seeds_range: instead of streamline_seeds provide [min, max] range to use for streamline starting point z-coordinates :kwarg streamline_length: streamline length + :kwarg max_step: maximum integration step, m (default: one grid cell of the run) + :kwarg n_points: number of points sampled along each returned streamline + :kwarg rtol, atol: scipy.integrate.solve_ivp error tolerances + :kwarg min_speed: velocity magnitude (m/s) treated as numerically zero - :returns: streamlines as numpy array + :returns: streamlines as numpy array, shape (n_seeds, n_points, 3). Points beyond which a streamline left the simulation domain are NaN. """ - + # bulk file f = pt.vlsvfile.VlsvReader(file_name=vlsvfile) # get box coordinates from data [xmin, ymin, zmin, xmax, ymax, zmax] = f.get_spatial_mesh_extent() - [xsize, ysize, zsize] = f.get_spatial_mesh_size() - simext =[xmin,xmax,ymin,ymax,zmin,zmax] - sizes = np.array([xsize,ysize,zsize]) - boxcoords=list(simext) + mesh_size = f.get_spatial_mesh_size() + [xsize, ysize, zsize] = mesh_size cellids = f.read_variable("CellID") - #Read the data from vlsv-file + # Read the data from vlsv-file Vx = f.read_variable("v", operator="x") Vz = f.read_variable("v", operator="z") + # Re-shape variable data + order = np.argsort(cellids) + Vxs = Vx[order].reshape(mesh_size, order="F") + Vzs = Vz[order].reshape(mesh_size, order="F") - #Re-shape variable data - Vxs=Vx[np.argsort(cellids)].reshape(f.get_spatial_mesh_size(), order="F") - Vys = np.zeros_like(Vxs) - Vzs=Vz[np.argsort(cellids)].reshape(f.get_spatial_mesh_size(), order="F") + # this routine is for 2D (x-z plane) runs: a single cell in y + Vxs2d = Vxs[:, 0, :] + Vzs2d = Vzs[:, 0, :] - data=dict(Vx=Vxs,Vy=Vys,Vz=Vzs) + # cell-centre coordinates of the (uniform) spatial mesh + dx = (xmax - xmin) / xsize + dz = (zmax - zmin) / zsize + x_coords = xmin + dx * (np.arange(xsize) + 0.5) + z_coords = zmin + dz * (np.arange(zsize) + 0.5) - #Create starting points for streamlines if they are not given - if streamline_seeds == None: - streamline_seeds = np.array([[seeds_x0, 0 ,i] for i in np.linspace(seeds_range[0], seeds_range[1], seeds_n)]) + vx_interp = RegularGridInterpolator((x_coords, z_coords), Vxs2d, method="linear", bounds_error=False, fill_value=np.nan) + vz_interp = RegularGridInterpolator((x_coords, z_coords), Vzs2d, method="linear", bounds_error=False, fill_value=np.nan) + bounds = (x_coords[0], x_coords[-1], z_coords[0], z_coords[-1]) - #streamline_seeds = np.array(streamline_seeds) - #dataset in yt-form - yt_dataset = yt.load_uniform_grid( - data, - sizes, - bbox=np.array([[boxcoords[0], boxcoords[1]], - [boxcoords[2],boxcoords[3]], - [boxcoords[4],boxcoords[5]]])) - + if max_step is None: + max_step = min(dx, dz) - #data, seeds, dictionary positions, step size - streamlines = yt.visualization.api.Streamlines(yt_dataset, streamline_seeds, - "Vx", "Vy", "Vz", length=streamline_length, direction=1) + # Create starting points for streamlines if they are not given + if streamline_seeds is None: + streamline_seeds = np.array([[seeds_x0, 0.0, i] for i in np.linspace(seeds_range[0], seeds_range[1], seeds_n)]) + else: + streamline_seeds = np.asarray(streamline_seeds) - #trace the streamlines with yt - streamlines.integrate_through_volume() - # return streamline positions - return np.array(streamlines.streamlines) + streamlines = np.array([ + trace_streamline(x0, z0, vx_interp, vz_interp, bounds, streamline_length, max_step, n_points, rtol, atol, min_speed) + for x0, _, z0 in streamline_seeds + ]) + + return streamlines def make_magnetopause(streamlines, end_x=-15*6371000, x_point_n=50): @@ -94,7 +156,7 @@ def make_magnetopause(streamlines, end_x=-15*6371000, x_point_n=50): :param streams: streamlines (coordinates in m) :kwarg end_x: tail end x-coordinate (how far along the negative x-axis the magnetopause is calculated) :kwarg x_point_n: integer, how many x-axis points the magnetopause will be divided in between the subsolar point and tail - + :returns: the magnetopause position as coordinate points in numpy array """ @@ -109,9 +171,10 @@ def make_magnetopause(streamlines, end_x=-15*6371000, x_point_n=50): ## define points in the x axis where to find magnetopause points on the yz-plane x_points = np.linspace(subsolar_x, end_x, x_point_n) - + ## interpolate more exact points for streamlines at exery x_point - new_streampoints = np.zeros((len(x_points), len(streamlines), 1)) # new array for keeping interpolated streamlines in form streamlines_new[x_point, streamline, z-coordinate] + # new array for keeping interpolated streamlines in form streamlines_new[x_point, streamline, z-coordinate] + new_streampoints = np.zeros((len(x_points), len(streamlines), 1)) for i,stream in enumerate(streamlines): interpolated_streamline = interpolate(stream, x_points) @@ -131,7 +194,7 @@ def make_magnetopause(streamlines, end_x=-15*6371000, x_point_n=50): if (pos.size == 0) or (neg.size == 0): raise ValueError('No streamlines found for x axis point, try adding streamlines or checking the x_points') - + # find points closest to x-axis and save found points pos_z_mpause[i] = [x_point, pos[pos.argmin()]] neg_z_mpause[i] = [x_point, neg[neg.argmax()]] @@ -141,22 +204,23 @@ def make_magnetopause(streamlines, end_x=-15*6371000, x_point_n=50): return magnetopause -def find_magnetopause_sw_streamline_2d(vlsvfile, streamline_seeds=None, seeds_n=200, seeds_x0=20*6371000, seeds_range=[-5*6371000, 5*6371000], streamline_length=45*6371000, end_x=-15*6371000, x_point_n=50): +def find_magnetopause_sw_streamline_2d(vlsvfile, streamline_seeds=None, seeds_n=200, seeds_x0=20*6371000, seeds_range=(-5*6371000, 5*6371000), streamline_length=45*6371000, end_x=-15*6371000, x_point_n=50, *, max_step=None, n_points=400, rtol=1e-6, atol=1.0, min_speed=1e-3): """Finds the magnetopause position by tracing streamlines of the velocity field for 2d runs. :param vlsvfile: directory and file name of .vlsv data file to use for VlsvReader :kwarg streamline_seeds: optional streamline starting points in numpy array (coordinates in meters including the y-coordinate 0.0) - :kwarg seeds_n: instead of streamline_seeds provide a number of streamlines to be traced - :kwarg seeds_x0: instead of streamline_seeds provide an x-coordinate for streamline starting points + :kwarg seeds_n: instead of streamline_seeds provide a number of streamlines to be traced + :kwarg seeds_x0: instead of streamline_seeds provide an x-coordinate for streamline starting points :kwarg seeds_range: instead of streamline_seeds provide [min, max] range to use for streamline starting point z-coordinates :kwarg streamline_length: streamline length for tracing :kwarg end_x: tail end x-coordinate (how far along the negative x-axis the magnetopause is calculated) :kwarg x_point_n: integer, how many x-axis points the magnetopause will be divided in between the subsolar point and tail + :kwarg max_step, n_points, rtol, atol, min_speed: passed through to make_streamlines(); see there :returns: the magnetopause position as coordinate points in numpy array """ - streamlines = make_streamlines(vlsvfile, streamline_seeds, seeds_n, seeds_x0, seeds_range, streamline_length) + streamlines = make_streamlines(vlsvfile, streamline_seeds, seeds_n, seeds_x0, seeds_range, streamline_length, max_step=max_step, n_points=n_points, rtol=rtol, atol=atol, min_speed=min_speed) magnetopause = make_magnetopause(streamlines, end_x, x_point_n) return magnetopause