diff --git a/examples/04_advanced_examples/p04_mapping_fieldlines.py b/examples/04_advanced_examples/p04_mapping_fieldlines.py index 94ff234..5f4198e 100644 --- a/examples/04_advanced_examples/p04_mapping_fieldlines.py +++ b/examples/04_advanced_examples/p04_mapping_fieldlines.py @@ -35,7 +35,7 @@ # %% # Load in the magnetic field and scalar field files used for this example. # They are from a CORHEL-MAS thermodynamic MHD calculation for CR2282. -# Becuase they are not part of the standard datasets in mapflpy or psi-io +# Because they are not part of the standard datasets in mapflpy or psi-io # (yet), we fetch them manually and place them in the default cache location. diff --git a/examples/04_advanced_examples/p05_plotting_q.py b/examples/04_advanced_examples/p05_plotting_q.py new file mode 100644 index 0000000..f47a22a --- /dev/null +++ b/examples/04_advanced_examples/p05_plotting_q.py @@ -0,0 +1,103 @@ +""" +Plotting the Squashing Factor +========================================= + +This example demonstrates how to use and plot the :func:`~mapflpy.scripts.compute_q_on_surface` +to visualize key topology and morphology metrics of the magnetic field. +The :func:`~mapflpy.scripts.compute_q_on_surface` combines :func:`~mapflpy.scripts.expansion_factor` with + :func:`~mapflpy.utils.calc_jacobian` and :func:`~mapflpy.utils.calc_q` using + :func:`~mapflpy.scripts.map_pt_forward` or :func:`~mapflpy.scripts.map_pt_backward` + to calculate the squashing factor. + +For a more complete description of the squashing factor, see +`Titov et al. 2007 `_. +Generally, the squashing factor indicates how much a given flux tube distorts. +Other ways to imagine the squashing factor, Q, includes places where current sheets +are likely (but not guaranteed) to form as high Q lines indicate quasi-separatrix layers (QSLs). +Alternatively, high Q lines indicate different magnetic flux domains. + +Additionally, this plots :func:`~mapflpy.scripts.expansion_factor`, critical to models such as +WSA (see `Wang & Sheeley 1990 `_, +`Arge & Pizzo 2000 `_, and +`Arge et al. 2004 `_). +""" +import os +from psi_data import fetch_mas_data +import numpy as np +import matplotlib.pylab as plt +from mapflpy.scripts import map_pt_forward, expansion_factor, compute_q_on_surface + +# sphinx_gallery_start_ignore +if 'SPHINX_GALLERY_BUILD' not in os.environ: + import matplotlib + + matplotlib.use('TkAgg') +# sphinx_gallery_end_ignore + +# %% +# The squashing factor, Q, is a measure of the topology of the magnetic field. +# So, let's read in magnetic field files. We're loading in from +# a CORHEL-MAS thermodynamic MHD calculation for CR2282. These aren't +# currently standard datasets in mapflpy or psi-io, so we're fetching them +# manually and placing them in the default cache location. + +files = fetch_mas_data(domains="cor", variables="br,bt,bp,t") +magnetic_field_files = files.cor_br, files.cor_bt, files.cor_bp + +# %% +# As a quick look, we can immediately calculate Q +# specifying only the magnetic field files and visualize the output. + +# compute Q +squashing_factor_default = compute_q_on_surface(magnetic_field_files) +# plot Q +ax = plt.figure().add_subplot() +q_map = ax.pcolormesh(np.rad2deg(squashing_factor_default.p), 90 - np.rad2deg(squashing_factor_default.t), + np.log10(squashing_factor_default.q), + cmap='Grays') +ax.set_aspect("equal", adjustable="box") +ax.set_title('Log$_{10}$ Q') +plt.colorbar(q_map) +plt.show() + +# %% +# That should run fairly quickly, as the resolution is a little bit more than 1 degree. +# However, we can customize the map. We can, for example, +# change the direction of mapping to "bwd" and choose a trace_radius of 3. +# We can either specify a [start, end] for theta and phi with t_range and p_range, +# or specify our own array with t_arr and p_arr. +# We are intentionally picking a low resolution so this runs fast, you should use more points! + + +squashing_factor_3_bwd = compute_q_on_surface(magnetic_field_files, direction='bwd', nproc=4, trace_radius=3, + p_arr=np.linspace(0, 2 * np.pi, 80), t_arr=np.linspace(0, np.pi, 40)) +# and visualizing: +ax = plt.figure().add_subplot() +q_map = ax.pcolormesh(np.rad2deg(squashing_factor_3_bwd.p), 90 - np.rad2deg(squashing_factor_3_bwd.t), + np.log10(squashing_factor_3_bwd.q), + cmap='Grays') +ax.set_aspect("equal", adjustable="box") +ax.set_title('Log$_{10}$ Q') +plt.colorbar(q_map) +plt.show() + +# %% +# This wrapper makes it easy to get the squashing factor. If we're interested +# in just say, the expansion factor, we can plot that. + +# We first need to calculate a mapping on a set of given points +p_to_trace = np.linspace(0, 2 * np.pi, 100) +t_to_trace = np.linspace(0, np.pi, 50) + +# let's map and get the expansion factor +mapping = map_pt_forward(*magnetic_field_files, p_to_trace, t_to_trace) +ef, p_ef, t_ef = expansion_factor(magnetic_field_files, mapping, 3, p_to_trace, t_to_trace) + +# and now we can visualize +ax = plt.figure().add_subplot() +ef_map = ax.pcolormesh(np.rad2deg(p_to_trace), 90 - np.rad2deg(t_to_trace), np.log10(ef).T, + cmap='plasma') +ax.set_aspect("equal", adjustable="box") +ax.set_title('expansion factor') +plt.colorbar(ef_map) +plt.show() diff --git a/mapflpy/globals.py b/mapflpy/globals.py index 818e0d9..cde63fc 100644 --- a/mapflpy/globals.py +++ b/mapflpy/globals.py @@ -79,6 +79,26 @@ """ ) +# ------------------------------------------------------------------------------ +# Named tuple for storing squashing factor produced by the +# scripts in `mapflpy.scripts` +# ------------------------------------------------------------------------------ +SquashingFactor = namedtuple('Squashing_Factor', ['q', 'p', 't']) +SquashingFactor.__doc__ = ( + """Named tuple for storing q and the field on which it was calculated. + + Attributes + ---------- + q : ndarray + Array containing the squashing factor + p : ndarray + Array containing the phi (longitude) positions for q in radians. + t : ndarray + Array containing the theta (co-latitude) positions for q in radians. + + """ +) + # ------------------------------------------------------------------------------ # Type aliases for improved code readability. # ------------------------------------------------------------------------------ diff --git a/mapflpy/scripts.py b/mapflpy/scripts.py index 336d838..e8ea6c4 100644 --- a/mapflpy/scripts.py +++ b/mapflpy/scripts.py @@ -29,9 +29,10 @@ import numpy as np from numpy.typing import NDArray, ArrayLike -from mapflpy.globals import DEFAULT_BUFFER_SIZE, Traces, Mapping, PathType, DirectionType, ContextType +from mapflpy.globals import DEFAULT_BUFFER_SIZE, Traces, Mapping, PathType, DirectionType, ContextType, SquashingFactor from mapflpy.tracer import TracerMP -from mapflpy.utils import shift_phi_traces, shift_phi_lps, fetch_default_launch_points +from mapflpy.utils import shift_phi_traces, shift_phi_lps, fetch_default_launch_points, modulo_twopi, get_half_mesh, calc_jacobian, calc_q +from psi_io import interpolate_positions_from_hdf __all__ = [ "run_forward_tracing", @@ -414,6 +415,198 @@ def trace_function(launch_points): return mapping +def expansion_factor(b_files, mapping, trace_radius, pss, tss): + """ + Calculate the expansion factor of a given set of field line launch points and + their mapped end points. + + Parameters + ---------- + b_files : list of str + Specified the magnetic field (br, bt, bp ordered) files + mapping : :class:`~mapflpy.globals.Mapping` + A namedtuple containing mapping results (:class:`mapflpy.globals.Mapping`). + trace_radius : float + The radius from which to map. + pss : ndarray + phi points used to generate the mapping + tss : ndarray + theta points used to generate the mapping + + Returns + ------- + efl : ndarray + An array of the expansion factor + pss : ndarray + phi points used to generate the expansion factor + tss : ndarray + theta points used to generate the expansion factor + + + Notes + ----- + - This function manually handles periodic boundaries in phi + """ + + # for the sake of interp: make sure phi's periodicity is obeyed + if (pss[0] < 0) or (pss[-1] > 2 * np.pi): + pss = np.mod(pss, np.pi * 2) + # make 2d launch point meshes so we can get the initial magnetic field + tss2d, pss2d = np.meshgrid(tss, pss) + ones2d_ss = np.ones_like(pss2d) + + # we get [p, t] ordered arrays with the magnetic field values + # xfl0: radial launch point location + xfl0 = ones2d_ss * trace_radius + # xfl1: radial traced point location + xfl1 = np.transpose(mapping.r) + # bs0: radial magnetic field launch point + bs0 = interpolate_positions_from_hdf(b_files[0], xfl0, tss2d, pss2d) + # bs1: radial magnetic field traced point + bs1 = interpolate_positions_from_hdf(b_files[0], np.transpose(mapping.r), np.transpose(mapping.t), + np.transpose(mapping.p)) + + # logic checks: + # if we have a magnetic field of 0 in the denominator, set the expansion factor to 0: + efl_raw = np.where(bs0 == 0, 0, abs((bs0 * xfl0 ** 2) / (bs1 * xfl1 ** 2))) + # if we make it to the boundary, use our calculated expansion factor. otherwise, = 0: + efl = np.where(np.transpose(mapping.traced_to_boundary), efl_raw, 0) + + return efl, pss, tss + +def compute_q_on_surface(b_files, direction='fwd', nproc=4, trace_radius=1, p_arr=np.asarray([]), + t_arr=np.asarray([]), p_range=None, t_range=None, nppts=300, ntpts=150): + """ + This wrapper calculates the squashing factor for a specified slice on a specified grid. + + Parameters + ---------- + b_files : list of str + Specified the magnetic field (br, bt, bp ordered) files + direction : str + either "fwd," "bwd," or "fwdbwd" to specify what direction of tracing (last is averaged) + nproc: int, optional + The number of processes to spawn. This should be equal to or less than the number of threads that can be used on the machine. + trace_radius : float, optional + The radius from which to map. Defaults to 1. + p_arr : ndarray, optional + User-specified array for phi points used to generate the mapping. Must also specify t_arr. Either specify thse or t_range and p_range. + t_arr : ndarray, optional + User-specified array for theta points used to generate the mapping. Must also specify p_arr. Either specify thse or t_range and p_range. + p_range : list or ndarray of two float, optional + User-specified start and end point in theta. defaults to [0, 2*np.pi]. + t_range : list or ndarray of two float, optional + User-specified start and end point in theta. defaults to [0, np.pi]. + nppts : int, optional + Number of points desired in phi when p_range is in use + ntpts : int, optional + Number of points desired in theta when t_range is in use + + Returns + ------- + squashing_factor : :class:`~mapflpy.globals.SquashingFactor` + A namedtuple containing the q, p, t results (:class:`mapflpy.globals.SquashingFactor`). + + """ + # default ranges for t, p + if t_range is None: + t_range = [0, np.pi] + if p_range is None: + p_range = [0, 2 * np.pi] + + # creating theta and phi ranges for field lines from + # user-defined domains (or above default ranges) + if t_arr.shape[0] == 0: + tss_specified = np.linspace(t_range[0], t_range[1], ntpts) + pss_specified = np.linspace(p_range[0], p_range[1], nppts) + + pss = get_half_mesh(pss_specified) + + if (t_range[0] == 0) or (t_range[-1] == np.pi): + + th = get_half_mesh(tss_specified) + tss = np.copy(th) + + # clip the boundaries + tss[0] = 0.0 + tss[-1] = np.pi + clipped = True + else: + tss = get_half_mesh(tss_specified) + clipped = False + # using the user-defined mesh for t, p + else: + th = get_half_mesh(t_arr) + tss = np.copy(th) + pss = get_half_mesh(p_arr) + if (tss[0] < 0) or (tss[-1] > np.pi): + # clip the boundaries + tss[0] = 0.0 + tss[-1] = np.pi + clipped = True + else: + clipped = False + + # field-line tracing by direction. starting with forward + if direction == 'fwd': + print('fwd mapping') + mapping = map_pt_forward(*b_files, pss, tss, radius=trace_radius, nproc=nproc) + + # make the expansion factor. + ef_arr, p_ef, t_ef = expansion_factor(b_files, mapping, trace_radius, pss, tss) + + # make the components of the jacobian + dtdt, dtdp, dpdt, dpdp = calc_jacobian(mapping, pss, tss) + # put the expansion factor, jacobian, and field lines together to get q + q, p, t = calc_q(dtdt, dtdp, dpdt, dpdp, mapping, pss, tss, ef_arr, clipped) + + return SquashingFactor(q, p, t) + + # field-line tracing by direction. now backward + elif direction == 'bwd': + print('bwd mapping') + + mapping = map_pt_backward(*b_files, pss, tss, radius=trace_radius, nproc=nproc) + + # make the expansion factor. + ef_arr, p_ef, t_ef = expansion_factor(b_files, mapping, trace_radius, pss, tss) + + # make the components of the jacobian + dtdt, dtdp, dpdt, dpdp = calc_jacobian(mapping, pss, tss) + + # put the expansion factor, jacobian, and field lines together to get q + q, p, t = calc_q(dtdt, dtdp, dpdt, dpdp, mapping, pss, tss, ef_arr, clipped) + + return SquashingFactor(q, p, t) + + # field-line tracing by direction. now take the average of the forward/backward + elif direction == 'fwdbwd': + print('fwdbwd mapping') + + mapping_fwd = map_pt_forward(*b_files, pss, tss, radius=trace_radius, nproc=nproc) + mapping_bwd = map_pt_backward(*b_files, pss, tss, radius=trace_radius, nproc=nproc) + + # make the expansion factor. + ef_arr_fwd, p_ef_f, t_ef_f = expansion_factor(b_files, mapping_fwd, trace_radius, pss, tss) + ef_arr_bwd, p_ef_b, t_ef_b = expansion_factor(b_files, mapping_bwd, trace_radius, pss, tss) + + # make the components of the jacobian + dtdt_fwd, dtdp_fwd, dpdt_fwd, dpdp_fwd = calc_jacobian(mapping_fwd, pss, tss) + dtdt_bwd, dtdp_bwd, dpdt_bwd, dpdp_bwd = calc_jacobian(mapping_bwd, pss, tss) + + # put the expansion factor, jacobian, and field lines together to get q + q_fwd, p_fwd, t_fwd = calc_q(dtdt_fwd, dtdp_fwd, dpdt_fwd, dpdp_fwd, mapping_fwd, pss, tss, ef_arr_fwd, clipped) + q_bwd, p_bwd, t_bwd = calc_q(dtdt_bwd, dtdp_bwd, dpdt_bwd, dpdp_bwd, mapping_bwd, pss, tss, ef_arr_bwd, clipped) + + # reutrn the averaged quantity + p = 0.5 * (p_fwd + p_bwd) + t = 0.5 * (t_fwd + t_bwd) + q = 0.5 * (q_fwd + q_bwd) + + return SquashingFactor(q, p, t) + + else: + raise Exception("specify a valid direction: fwd, bwd, fwdbwd") def inter_domain_tracing(br_cor: PathType, bt_cor: PathType, diff --git a/mapflpy/utils.py b/mapflpy/utils.py index 319a7e6..3eb4864 100644 --- a/mapflpy/utils.py +++ b/mapflpy/utils.py @@ -503,9 +503,189 @@ def combine_and_pad_fieldlines(arrs: list | tuple, end_pos[:, i] = a[:, -1] return Traces(geometry, start_pos, end_pos, traced_to_boundary, integral) +def get_half_mesh(x): + """ + Return the FULL half mesh based on an array of 1D mesh locations. + + Assuming that the input x is the bounding exterior (main) mesh, this returns + the exterior (half) mesh locations. + + *** Unlike the "get_1d_mesh_properties" method in map_mesh, this routine + gives you the half mesh locations at the exterior, like in MAS (no clipping). + + Parameters + ---------- + x : 1D numpy array of n grid locations. + + Returns + ------- + xh: 1D numpy array of grid locations on the exterior half mesh (n+1). + """ + # Compute the interior (half) mesh positions + xh = 0.5*(x[1:] + x[0:-1]) + + # Compute the spacing centered around the interior half mesh + dxh = x[1:] - x[0:-1] + + # Add the boundary points to make the half mesh with exterior points + xh = np.concatenate([[xh[0] - dxh[0]], xh, [xh[-1] + dxh[-1]]]) + + return xh + +def modulo_twopi(x_arr): + """ + This returns the smallest value of x_arr, modulo 2*pi. + + Parameters + ---------- + x_arr : float or ndarray + Specified input value + + Returns + ------- + m_twopi : float or ndarray + The smallest value of each x_arr, modulo 2*pi + + Notes + ----- + - This uses :func:`~numpy.isclose` + """ + + # generate plus/minus 2pi + xm = np.abs(x_arr - (2 * np.pi)) + xp = np.abs(x_arr + (2 * np.pi)) + # check which value is the smallest + x_min = np.minimum(np.minimum(xm, np.abs(x_arr)), xp) + + # determine the appropriate value, modulo two pi, elementwise + m_twopi = np.where(np.isclose(xm, x_min), + x_arr - 2 * np.pi, + np.where(np.isclose(xp, x_min), x_arr + 2 * np.pi, + x_arr)) + return m_twopi + +def calc_jacobian(mapping, pss, tss): + """ + This calculates the Jacobian for a given field line mapping + and its launch points. + + Parameters + ---------- + mapping : :class:`~mapflpy.globals.Mapping` + A namedtuple containing mapping results (:class:`mapflpy.globals.Mapping`). + tss : ndarray + Theta points used to generate the mapping + pss : ndarray + Phi points used to generate the mapping + + Returns + ------- + dtdt, dtdp, dpdt, dpdp : ndarray + Each respective component of the jacobian as an ndarray + + """ + + tfl = mapping.t + pfl = mapping.p + + # get the spacing from launch points + dp = pss[1:] - pss[:-1] + dt = tss[1:] - tss[:-1] + + # and we need these as mesh grids since we're vectorized + dt_mg, dp_mg = np.meshgrid(dt, dp, indexing='ij') + + # get the components of the jacobian + # thetas + dtdt_m = (tfl[1:, :-1] - tfl[:-1, :-1]) / dt_mg + dtdt_p = (tfl[1:, 1:] - tfl[:-1, 1:]) / dt_mg + dtdp_m = (tfl[:-1, 1:] - tfl[:-1, :-1]) / dp_mg + dtdp_p = (tfl[1:, 1:] - tfl[1:, :-1]) / dp_mg + # phis + dpdt_m = modulo_twopi(pfl[1:, :-1] - pfl[:-1, :-1]) / dt_mg + dpdt_p = modulo_twopi(pfl[1:, 1:] - pfl[:-1, 1:]) / dt_mg + dpdp_m = modulo_twopi(pfl[:-1, 1:] - pfl[:-1, :-1]) / dp_mg + dpdp_p = modulo_twopi(pfl[1:, 1:] - pfl[1:, :-1]) / dp_mg + + # now combine + dtdt = 0.5 * (dtdt_m + dtdt_p) + dtdp = 0.5 * (dtdp_m + dtdp_p) + dpdt = 0.5 * (dpdt_m + dpdt_p) + dpdp = 0.5 * (dpdp_m + dpdp_p) + + return dtdt, dtdp, dpdt, dpdp + +def calc_q(dtdt, dtdp, dpdt, dpdp, mapping, pss, tss, ef_arr, clipped = False): + """ + Using components of a Jacobian, this calculates the squashing factor, q. + + Parameters + ---------- + dtdt : ndarray + dtdt component of a Jacobian from :func:`~calc_jacobian` + dtdp : ndarray + dtdp component of a Jacobian from :func:`~calc_jacobian` + dpdt : ndarray + dpdt component of a Jacobian from :func:`~calc_jacobian` + dpdp : ndarray + dpdp component of a Jacobian from :func:`~calc_jacobian` + mapping : :class:`~mapflpy.globals.Mapping` + A namedtuple containing mapping results (:class:`mapflpy.globals.Mapping`). + tss : ndarray + Theta points used to generate the mapping + pss : ndarray + Phi points used to generate the mapping + ef_arr : ndarray + An array of the expansion factor + clipped: logical + Specifies whether the input theta field is clipped + Returns + ------- + t_i : ndarray + theta values on which q is calculated + p_i : ndarray + phi values on which q is calculated + q : ndarray + squashing factor + """ + # get the averaged expansion factor: + efav = 0.25 * (ef_arr[:-1, :-1] + ef_arr[1:, :-1] + ef_arr[:-1, 1:] + ef_arr[1:, 1:]) + + # get the traced field lines in theta + tfl = mapping.t + # average the traced field lines in theta + tmav = 0.25 * (tfl[:-1, :-1] + tfl[1:, :-1] + tfl[:-1, 1:] + tfl[1:, 1:]) + + # create the mesh on which we compute Q + t_i = 0.5 * (tss[1:] + tss[0:-1]) + p_i = 0.5 * (pss[1:] + pss[0:-1]) + + # and we'll need meshgrid for vectorization + t_img, p_img = np.meshgrid(t_i, p_i, indexing='ij') + + # calculate the coefficients that'll make q + aa = np.sin(tmav) * dpdp / np.sin(t_img) + bb = np.sin(tmav) * dpdt + cc = dtdp / np.sin(t_img) + dd = dtdt + + # we did it! + q_raw = (aa ** 2 + bb ** 2 + cc ** 2 + dd ** 2) / np.transpose(efav) + + # final logic checks: expansion factor can't be 0 + q_checked = np.where(np.transpose(efav) == 0.0, 0, q_raw) + q = np.copy(q_checked) + if clipped == True: + # make sure the poles are now actually at the poles and not just close + q[0, :] = np.mean(q[0, :]) + q[-1, :] = np.mean(q[-1, :]) + # and finally change the theta ends to reflect the above change + t_i[-1] = np.pi + t_i[0] = 0.0 + return q, p_i, t_i def s2c(r, t, p): """ diff --git a/pyproject.toml b/pyproject.toml index c21855d..fb952b1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -16,7 +16,7 @@ requires = [ # ---------------- [project] name = "mapflpy" -version = "1.2.1" +version = "1.3.0" description = "Fast field line tracing for spherical vector fields" authors = [ {name = "Predictive Science Inc.", email = "webmaster@predsci.com"},