def run_test(fld, seeds, plot2d=True, plot3d=True, add_title="", view_kwargs=None, show=False, scatter_mpl=False, mesh_mvi=True): interpolated_fld = viscid.interp_trilin(fld, seeds) seed_name = seeds.__class__.__name__ if add_title: seed_name += " " + add_title try: if not plot2d: raise ImportError from viscid.plot import vpyplot as vlt from matplotlib import pyplot as plt plt.clf() # plt.plot(seeds.get_points()[2, :], fld) mpl_plot_kwargs = dict() if interpolated_fld.is_spherical(): mpl_plot_kwargs['hemisphere'] = 'north' vlt.plot(interpolated_fld, **mpl_plot_kwargs) plt.title(seed_name) plt.savefig(next_plot_fname(__file__, series='2d')) if show: plt.show() if scatter_mpl: plt.clf() vlt.plot2d_line(seeds.get_points(), fld, symdir='z', marker='o') plt.savefig(next_plot_fname(__file__, series='2d')) if show: plt.show() except ImportError: pass try: if not plot3d: raise ImportError from viscid.plot import vlab _ = get_mvi_fig(offscreen=not show) try: if mesh_mvi: mesh = vlab.mesh_from_seeds(seeds, scalars=interpolated_fld) mesh.actor.property.backface_culling = True except RuntimeError: pass pts = seeds.get_points() p = vlab.points3d(pts[0], pts[1], pts[2], interpolated_fld.flat_data, scale_mode='none', scale_factor=0.02) vlab.axes(p) vlab.title(seed_name) if view_kwargs: vlab.view(**view_kwargs) vlab.savefig(next_plot_fname(__file__, series='3d')) if show: vlab.show(stop=True) except ImportError: pass
def minvar(B, p1, p2, n=40): """Find minimum variance eigenvectors of `B` between `p1` and `p2` Do minimum variance analysis (MVA) of vector field `B` between two points, i.e., `p1` and `p2` are like two spacecraft locations on either side of the boundary for the MVA. Args: B (:py:class:`viscid.field.VectorField`): Vector field for MVA p1 (sequence): xyz point on one side of the boundary p2 (sequence): xyz point on the other side of the boundary n (int): number of points to sample B for the MVA Returns: (evals, evecs) All are ordered toward increasing eigenvalue magnitude. `evecs` is a 2d array where the vectors are columns (2nd axis), i.e., the min variance eigenvector is evecs[0, :] """ p1 = np.array(p1).reshape((3,)) p2 = np.array(p2).reshape((3,)) line = viscid.Line(p1, p2, n) b = viscid.interp_trilin(B, line) return minvar_series(b.data)
def minvar(B, p1, p2, n=40): """Find minimum variance eigenvectors of `B` between `p1` and `p2` Do minimum variance analysis (MVA) of vector field `B` between two points, i.e., `p1` and `p2` are like two spacecraft locations on either side of the boundary for the MVA. Args: B (:py:class:`viscid.field.VectorField`): Vector field for MVA p1 (sequence): xyz point on one side of the boundary p2 (sequence): xyz point on the other side of the boundary n (int): number of points to sample B for the MVA Returns: (evals, evecs) All are ordered toward increasing eigenvalue magnitude. `evecs` is a 2d array where the vectors are columns (2nd axis), i.e., the min variance eigenvector is evecs[0, :] """ p1 = np.array(p1).reshape((3, )) p2 = np.array(p2).reshape((3, )) line = viscid.Line(p1, p2, n) b = viscid.interp_trilin(B, line) return minvar_series(b.data)
def run_test(fld, seeds, plot2d=True, plot3d=True, add_title="", view_kwargs=None, show=False, scatter_mpl=False, mesh_mvi=True): interpolated_fld = viscid.interp_trilin(fld, seeds) seed_name = seeds.__class__.__name__ if add_title: seed_name += " " + add_title try: if not plot2d: raise ImportError from viscid.plot import vpyplot as vlt from matplotlib import pyplot as plt plt.clf() # plt.plot(seeds.get_points()[2, :], fld) mpl_plot_kwargs = dict() if interpolated_fld.is_spherical(): mpl_plot_kwargs['hemisphere'] = 'north' vlt.plot(interpolated_fld, **mpl_plot_kwargs) plt.title(seed_name) plt.savefig(next_plot_fname(__file__, series='2d')) if show: plt.show() if scatter_mpl: plt.clf() vlt.plot2d_line(seeds.get_points(), fld, symdir='z', marker='o') plt.savefig(next_plot_fname(__file__, series='2d')) if show: plt.show() except ImportError: pass try: if not plot3d: raise ImportError vlab, _ = get_mvi_fig() try: if mesh_mvi: mesh = vlab.mesh_from_seeds(seeds, scalars=interpolated_fld) mesh.actor.property.backface_culling = True except RuntimeError: pass pts = seeds.get_points() p = vlab.points3d(pts[0], pts[1], pts[2], interpolated_fld.flat_data, scale_mode='none', scale_factor=0.02) vlab.axes(p) vlab.title(seed_name) if view_kwargs: vlab.view(**view_kwargs) vlab.savefig(next_plot_fname(__file__, series='3d')) if show: vlab.show(stop=True) except ImportError: pass
def run_test(fld, seeds, plot2d=True, plot3d=True, add_title="", view_kwargs=None, show=False): interpolated_fld = viscid.interp_trilin(fld, seeds) seed_name = seeds.__class__.__name__ if add_title: seed_name += " " + add_title try: if not plot2d: raise ImportError from matplotlib import pyplot as plt from viscid.plot import vpyplot as vlt plt.clf() # plt.plot(seeds.get_points()[2, :], fld) mpl_plot_kwargs = dict() if interpolated_fld.is_spherical(): mpl_plot_kwargs['hemisphere'] = 'north' vlt.plot(interpolated_fld, **mpl_plot_kwargs) plt.title(seed_name) plt.savefig(next_plot_fname(__file__, series='2d')) if show: plt.show() except ImportError: pass try: if not plot3d: raise ImportError from viscid.plot import vlab try: fig = _global_ns['figure'] vlab.clf() except KeyError: fig = vlab.figure(size=[1200, 800], offscreen=not show) _global_ns['figure'] = fig try: mesh = vlab.mesh_from_seeds(seeds, scalars=interpolated_fld) mesh.actor.property.backface_culling = True except RuntimeError: pass pts = seeds.get_points() p = vlab.points3d(pts[0], pts[1], pts[2], interpolated_fld.flat_data, scale_mode='none', scale_factor=0.02) vlab.axes(p) vlab.title(seed_name) if view_kwargs: vlab.view(**view_kwargs) vlab.savefig(next_plot_fname(__file__, series='3d')) if show: vlab.show() except ImportError: pass
def run_test(fld, seeds, plot2d=True, plot3d=True, add_title="", view_kwargs=None, show=False): interpolated_fld = viscid.interp_trilin(fld, seeds) seed_name = seeds.__class__.__name__ if add_title: seed_name += " " + add_title try: if not plot2d: raise ImportError from viscid.plot import mpl mpl.plt.clf() # mpl.plt.plot(seeds.get_points()[2, :], fld) mpl_plot_kwargs = dict() if interpolated_fld.is_spherical(): mpl_plot_kwargs['hemisphere'] = 'north' mpl.plot(interpolated_fld, **mpl_plot_kwargs) mpl.plt.title(seed_name) mpl.plt.savefig(next_plot_fname(__file__, series='2d')) if show: mpl.plt.show() except ImportError: pass try: if not plot3d: raise ImportError from viscid.plot import mvi try: fig = _global_ns['figure'] mvi.clf() except KeyError: fig = mvi.figure(size=[1200, 800], offscreen=not show) _global_ns['figure'] = fig try: mesh = mvi.mesh_from_seeds(seeds, scalars=interpolated_fld) mesh.actor.property.backface_culling = True except RuntimeError: pass pts = seeds.get_points() p = mvi.points3d(pts[0], pts[1], pts[2], interpolated_fld.flat_data, scale_mode='none', scale_factor=0.02) mvi.axes(p) mvi.title(seed_name) if view_kwargs: mvi.view(**view_kwargs) mvi.savefig(next_plot_fname(__file__, series='3d')) if show: mvi.show() except ImportError: pass
def project_along_line(line, fld): """Project a Vector Field Parallel to a streamline Args: line (ndarray): 3xN of points along the fld (VectorField): Field to interpolate and project onto the line """ fld_on_verts = viscid.interp_trilin(fld, line) dsvec = line[:, 2:] - line[:, :-2] dsvec = np.concatenate([dsvec[:, 0:1], dsvec, dsvec[:, -2:-1]], axis=1) dsvec = dsvec / np.linalg.norm(dsvec, axis=0) return np.sum(fld_on_verts * dsvec.T, axis=1)
def minvar(B, p1, p2, n=40): """Find minimum variance eigenvectors of `B` between `p1` and `p2` Do minimum variance analysis (MVA) of vector field `B` between two points, i.e., `p1` and `p2` are like two spacecraft locations on either side of the boundary for the MVA. Args: B (:py:class:`viscid.field.VectorField`): Vector field for MVA p1 (sequence): xyz point on one side of the boundary p2 (sequence): xyz point on the other side of the boundary n (int): number of points to sample B for the MVA Returns: (evals, evecs) All are ordered toward increasing eigenvalue magnitude. `evecs` is a 2d array where the vectors are columns (2nd axis), i.e., the min variance eigenvector is evecs[0, :] """ p1 = np.array(p1).reshape((3,)) p2 = np.array(p2).reshape((3,)) line = viscid.Line(p1, p2, n) b = viscid.interp_trilin(B, line) m = np.empty((3, 3)) for i in range(3): for j in range(i+1): bibj = np.average(b[:, i] * b[:, j]) m[i, j] = bibj - (np.average(b[:, i]) * np.average(b[:, j])) m[j, i] = bibj - (np.average(b[:, i]) * np.average(b[:, j])) evals, evecs = np.linalg.eig(m) eval_mags = np.abs(evals) sorted_inds = np.argsort(eval_mags) evec_min = evecs[:, sorted_inds[0]].flatten() evec_int = evecs[:, sorted_inds[1]].flatten() evec_max = evecs[:, sorted_inds[2]].flatten() eval_min = eval_mags[sorted_inds[0]] eval_int = eval_mags[sorted_inds[1]] warn_thresh = 0.05 if (eval_int - eval_min) / eval_min < warn_thresh: viscid.logger.warn("Minvar says minimum and intermediate eigenvalues " "are too close together: {0:g} - {1:g} < {2:g}%" "".format(eval_int, eval_min, warn_thresh)) return evals[sorted_inds], np.array([evec_min, evec_int, evec_max]).T
def integrate_along_lines(lines, fld, reduction="dot", mask_func=None): """Integrate the value of fld along a list of lines Args: lines (list): list of 3xN ndarrays, N needs not be the same for all lines fld (Field): Field to interpolate / integrate reduction (str): If fld is a vector field, what quantity to integrate. Can be "dot" to dot the vectors with ds along the line, or "mag" to integrate the magnitude. Returns: ndarray with shape (len(lines), ) """ arr = np.zeros((len(lines),), dtype=fld.dtype) cum_n = np.cumsum([0] + [line.shape[1] for line in lines]) all_verts = np.concatenate(lines, axis=1) fld_on_verts = viscid.interp_trilin(fld, all_verts).data for i, start, stop in izip(count(), cum_n[:-1], cum_n[1:]): ds = np.linalg.norm(lines[i][:, 1:] - lines[i][:, :-1], axis=0) if len(fld_on_verts.shape) > 1: reduction = reduction.strip().lower() if reduction == "dot": dsvec = lines[i][:, 1:] - lines[i][:, :-1] dsvec = dsvec / np.linalg.norm(dsvec, axis=0) values = 0.5 * (fld_on_verts[start:stop - 1, :] + fld_on_verts[start + 1:stop, :]) values = values * dsvec.T if mask_func is not None: values = np.ma.masked_where(mask_func(values), values) values = np.sum(values, axis=1) elif reduction in ["mag", "magnitude", "norm"]: mag = np.linalg.norm(fld_on_verts[start:stop], axis=1) values = 0.5 * (mag[start:stop - 1] + mag[start + 1:stop]) else: raise ValueError("Unknown reduction: {0}".format(reduction)) else: values = 0.5 * (fld_on_verts[start:stop - 1] + fld_on_verts[start + 1:stop]) arr[i] = np.sum(values * ds) return arr
def minvar(B, p1, p2, n=40): """Find minimum variance eigenvectors of `B` between `p1` and `p2` Do minimum variance analysis (MVA) of vector field `B` between two points, i.e., `p1` and `p2` are like two spacecraft locations on either side of the boundary for the MVA. Args: B (:py:class:`viscid.field.VectorField`): Vector field for MVA p1 (sequence): xyz point on one side of the boundary p2 (sequence): xyz point on the other side of the boundary n (int): number of points to sample B for the MVA Returns: (evals, evecs) All are ordered toward increasing eigenvalue magnitude. `evecs` is a 2d array where the vectors are columns (2nd axis), i.e., the min variance eigenvector is evecs[0, :] """ p1 = np.array(p1).reshape((3, )) p2 = np.array(p2).reshape((3, )) line = viscid.Line(p1, p2, n) b = viscid.interp_trilin(B, line) m = np.empty((3, 3)) for i in range(3): for j in range(i + 1): bibj = np.average(b[:, i] * b[:, j]) m[i, j] = bibj - (np.average(b[:, i]) * np.average(b[:, j])) m[j, i] = bibj - (np.average(b[:, i]) * np.average(b[:, j])) evals, evecs = np.linalg.eig(m) eval_mags = np.abs(evals) sorted_inds = np.argsort(eval_mags) evec_min = evecs[:, sorted_inds[0]].flatten() evec_int = evecs[:, sorted_inds[1]].flatten() evec_max = evecs[:, sorted_inds[2]].flatten() eval_min = eval_mags[sorted_inds[0]] eval_int = eval_mags[sorted_inds[1]] warn_thresh = 0.05 if (eval_int - eval_min) / eval_min < warn_thresh: viscid.logger.warn("Minvar says minimum and intermediate eigenvalues " "are too close together: {0:g} - {1:g} < {2:g}%" "".format(eval_int, eval_min, warn_thresh)) return evals[sorted_inds], np.array([evec_min, evec_int, evec_max]).T
def integrate_along_lines(lines, fld): """Integrate the value of fld along a list of lines Args: lines (list): list of 3xN ndarrays, N needs not be the same for all lines fld (Field): Field to interpolate / integrate Returns: ndarray with shape (len(lines), ) """ arr = np.zeros((len(lines),), dtype=fld.dtype) cum_n = np.cumsum([0] + [line.shape[1] for line in lines]) all_verts = np.concatenate(lines, axis=1) fld_on_verts = viscid.interp_trilin(fld, all_verts) for i, start, stop in izip(count(), cum_n[:-1], cum_n[1:]): ds = np.linalg.norm(lines[i][:, 1:] - lines[i][:, :-1], axis=0) values = 0.5 * (fld_on_verts[start:stop - 1] + fld_on_verts[start + 1:stop]) arr[i] = np.sum(values * ds) return arr
def prepare_lines(lines, scalars=None, do_connections=False, other=None): """Concatenate and standardize a list of lines Args: lines (list): Must be a list of 3xN or 4xN ndarrays of xyz(s) data for N points along the line. N need not be the same for all lines. Can alse be 6xN such that lines[:][3:, :] are interpreted as rgb colors scalars (sequence): can be one of:: - single hex color (ex, `#FD7400`) - sequence of Nlines hex colors - single rgb(a) tuple - sequence of Nlines rgb(a) sequences - sequence of N values that go with each vertex and will be mapped with a colormap - sequence of Nlines values that go each line and will be mapped with a colormap - Field. If `np.prod(fld.shape) in (N, nlines)` then the field is interpreted as a simple sequence of values (ie, the topology result from calc_streamlines for coloring each line). Otherwise, the field is interpolated onto each vertex. do_connections (bool): Whether or not to make connections array other (dict): a dictionary of other arrays that should be reshaped and the like the same way scalars is Returns: (vertices, scalars, connections, other) * vertices (ndarray): 3xN array of N xyz points. N is the sum of the lengths of all the lines * scalars (ndarray): N array of scalars, 3xN array of uint8 rgb values, 4xN array of uint8 rgba values, or None * connections (ndarray): Nx2 array of ints (indices along axis 1 of vertices) describing the forward and backward connectedness of the lines, or None * other (dict): a dict of N length arrays Raises: ValueError: If rgb data is not in a valid range or the shape of scalars is not understood """ nlines = len(lines) npts = [line.shape[1] for line in lines] N = np.sum(npts) first_idx = np.cumsum([0] + npts[:-1]) vertices = [np.asarray(line) for line in lines] vertices = np.concatenate(lines, axis=1) if vertices.dtype.kind not in 'fc': vertices = np.asarray(vertices, dtype='f') if vertices.shape[0] > 3: if scalars is not None: viscid.logger.warning("Overriding line scalars with scalars kwarg") else: scalars = vertices[3:, :] vertices = vertices[:3, :] if scalars is not None: scalars_are_strings = False if isinstance(scalars, viscid.field.Field): if np.prod(scalars.shape) in (nlines, N): scalars = np.asarray(scalars).reshape(-1) else: scalars = viscid.interp_trilin(scalars, vertices) if scalars.size != N: raise ValueError("Scalars was not a scalar field") elif isinstance(scalars, (list, tuple, viscid.string_types, np.ndarray)): # string types need some extra massaging if any(isinstance(s, viscid.string_types) for s in scalars): assert all(isinstance(s, viscid.string_types) for s in scalars) scalars_are_strings = True scalars = _string_colors_as_hex(scalars) elif isinstance(scalars, np.ndarray): scalars = scalars else: for i, si in enumerate(scalars): if not isinstance(si, np.ndarray): scalars[i] = np.asarray(si) scalars[i] = np.atleast_2d(scalars[i]) try: scalars = np.concatenate(scalars, axis=0) except ValueError: scalars = np.concatenate(scalars, axis=1) scalars = np.atleast_2d(scalars) if scalars.dtype == np.dtype('object'): raise RuntimeError("Scalars must be numbers, tuples of numbers " "that indicate rgb(a), or hex strings - they " "must not be python objects") if scalars.shape == (1, 1): scalars = scalars.repeat(N, axis=1) elif scalars.shape == (1, nlines) or scalars.shape == (nlines, 1): # one scalar for each line, so broadcast it scalars = scalars.reshape(nlines, 1) scalars = [scalars[i].repeat(ni) for i, ni in enumerate(npts)] scalars = np.concatenate(scalars, axis=0).reshape(1, N) elif scalars.shape == (N, 1) or scalars.shape == (1, N): # catch these so they're not interpreted as colors if # nlines == 1 and N == 3; ie. 1 line with 3 points scalars = scalars.reshape(1, N) elif scalars.shape in [(3, nlines), (nlines, 3), (4, nlines), (nlines, 4)]: # one rgb(a) color for each line, so broadcast it if (scalars.shape in [(3, nlines), (4, nlines)] and scalars.shape not in [(3, 3), (4, 4)]): # the guard against shapes (3, 3) and (4, 4) mean that # these square shapes are assumed Nlines x {3,4} scalars = scalars.T nccomps = scalars.shape[1] # 3 for rgb, 4 for rgba colors = [] for i, ni in enumerate(npts): c = scalars[i].reshape(nccomps, 1).repeat(ni, axis=1) colors.append(c) scalars = np.concatenate(colors, axis=1) elif scalars.shape in [(3, N), (N, 3), (4, N), (N, 4)]: # one rgb(a) color for each vertex if (scalars.shape in [(3, N), (4, N)] and scalars.shape not in [(3, 3), (4, 4)]): # the guard against shapes (3, 3) and (4, 4) mean that # these square shapes are assumed N x {3,4} scalars = scalars.T elif scalars.shape in [(1, 3), (3, 1), (1, 4), (4, 1)]: # interpret a single rgb(a) color, and repeat/broadcast it scalars = scalars.reshape([-1, 1]).repeat(N, axis=1) else: scalars = scalars.reshape(-1, N) # scalars now has shape (1, N), (3, N), or (4, N) if scalars_are_strings: # translate hex colors (#ff00ff) into rgb(a) values scalars = np.char.lstrip(scalars, '#') strlens = np.char.str_len(scalars) min_strlen, max_strlen = np.min(strlens), np.max(strlens) if min_strlen == max_strlen == 8: # 32-bit rgba (two hex chars per channel) scalars = _hexchar2int(scalars.astype('S8')).reshape(-1, 4).T elif min_strlen == max_strlen == 6: # 24-bit rgb (two hex chars per channel) scalars = _hexchar2int(scalars.astype('S6')).reshape(-1, 3).T else: raise NotImplementedError("This should never happen as " "scalars as colors should already " "be preprocessed appropriately") elif scalars.shape[0] == 1: # normal scalars, cast them down to a single dimension scalars = scalars.reshape(-1) elif scalars.shape[0] in (3, 4): # The scalars encode rgb data, standardize the result to a # 3xN or 4xN ndarray of 1 byte unsigned ints [0..255] if np.all(scalars >= 0) and np.all(scalars <= 1): scalars = (255 * scalars).round().astype('u1') elif np.all(scalars >= 0) and np.all(scalars < 256): scalars = scalars.round().astype('u1') else: raise ValueError("Rgb data should be in range [0, 1] or " "[0, 255], range given is [{0}, {1}]" "".format(np.min(scalars), np.max(scalars))) else: raise ValueError("Scalars should either be a number, or set of " "rgb values, shape is {0}".format(scalars.shape)) # scalars should now have shape (N, ) or be a uint8 array with shape # (3, N) or (4, N) encoding an rgb(a) color for each point [0..255] # ... done with scalars... # broadcast / reshape additional arrays given in other if other: for key, arr in other.items(): if arr is None: pass elif arr.shape == (1, nlines) or arr.shape == (nlines, 1): arr = arr.reshape(nlines, 1) arr = [arr[i].repeat(ni) for i, ni in enumerate(npts)] other[key] = np.concatenate(arr, axis=0).reshape(1, N) else: try: other[key] = arr.reshape(-1, N) except ValueError: viscid.logger.warning("Unknown dimension, dropping array {0}" "".format(key)) if do_connections: connections = [None] * nlines for i, ni in enumerate(npts): # i0 is the index of the first point of the i'th line in lines i0 = first_idx[i] connections[i] = np.vstack([np.arange(i0, i0 + ni - 1.5), np.arange(i0 + 1, i0 + ni - 0.5)]).T connections = np.concatenate(connections, axis=0).astype('i') else: connections = None return vertices, scalars, connections, other
def get_mp_info(pp, b, j, e, cache=True, cache_dir=None, slc="x=5.5f:11.0f, y=-4.0f:4.0f, z=-3.6f:3.6f", fit="mp_xloc", fit_p0=(9.0, 0.0, 0.0, 1.0, -1.0, -1.0)): """Get info about m-pause as flattened fields Notes: The first thing this function does is mask locations where the GSE-y current density < 1e-4. This masks out the bow shock and current free regions. This works for southward IMF, but it is not very general. Parameters: pp (ScalarcField): pressure b (VectorField): magnetic field j (VectorField): current density e (VectorField, None): electric field (same centering as b). If None, then the info that requires E will be filled with NaN cache (bool, str): Save to and load from cache, if "force", then don't load from cache if it exists, but do save a cache at the end cache_dir (str): Directory for cache, if None, same directory as that file to which the grid belongs slc (str): slice that gives a box that contains the m-pause fit (str): to which resulting field should the paraboloid be fit, defaults to mp_xloc, but pp_max_xloc might be useful in some circumstances fit_p0 (tuple): Initial guess vector for paraboloid fit Returns: dict: Unless otherwise noted, the entiries are 2D (y-z) fields - **mp_xloc** location of minimum abs(Bz), this works better than max of J^2 for FTEs - **mp_sheath_edge** location where Jy > 0.1 * Jy when coming in from the sheath side - **mp_sphere_edge** location where Jy > 0.1 * Jy when coming in from the sphere side - **mp_width** difference between m-sheath edge and msphere edge - **mp_shear** magnetic shear taken 6 grid points into the m-sheath / m-sphere - **pp_max** max pp - **pp_max_xloc** location of max pp - **epar_max** max e parallel - **epar_max_xloc** location of max e parallel - **paraboloid** numpy.recarray of paraboloid fit. The parameters are given in the 0th element, and the 1st element contains the 1-sigma values for the fit Raises: RuntimeError: if using MHD crds instead of GSE crds """ if not cache_dir: cache_dir = pp.find_info("_viscid_dirname", "./") run_name = pp.find_info("run", None) if cache and run_name: t = pp.time mp_fname = "{0}/{1}.mpause.{2:06.0f}".format(cache_dir, run_name, t) else: mp_fname = "" try: force = cache.strip().lower() == "force" except AttributeError: force = False try: if force or not mp_fname or not os.path.isfile(mp_fname + ".xdmf"): raise IOError() mp_info = {} with viscid.load_file(mp_fname + ".xdmf") as dat: fld_names = [ "mp_xloc", "mp_sheath_edge", "mp_sphere_edge", "mp_width", "mp_shear", "pp_max", "pp_max_xloc", "epar_max", "epar_max_xloc" ] for fld_name in fld_names: mp_info[fld_name] = dat[fld_name]["x=0"] except (IOError, KeyError): mp_info = {} crd_system = viscid.as_crd_system(b, None) if crd_system != 'gse': raise RuntimeError("get_mp_info can't work in MHD crds, " "switch to GSE please") if j.nr_patches == 1: pp_block = pp[slc] b_block = b[slc] j_block = j[slc] if e is None: e_block = np.nan * viscid.empty_like(j_block) else: e_block = e[slc] else: # interpolate an amr grid so we can proceed obnd = pp.get_slice_extent(slc) dx = np.min(pp.skeleton.L / pp.skeleton.n, axis=0) nx = np.ceil((obnd[1] - obnd[0]) / dx) vol = viscid.seed.Volume(obnd[0], obnd[1], nx, cache=True) pp_block = vol.wrap_field(viscid.interp_trilin(pp, vol), name="P").as_cell_centered() b_block = vol.wrap_field(viscid.interp_trilin(b, vol), name="B").as_cell_centered() j_block = vol.wrap_field(viscid.interp_trilin(j, vol), name="J").as_cell_centered() if e is None: e_block = np.nan * viscid.empty_like(j_block) else: e_block = vol.wrap_field(viscid.interp_trilin(e, vol), name="E").as_cell_centered() # jsq = viscid.dot(j_block, j_block) bsq = viscid.dot(b_block, b_block) # extract ndarrays and mask out bow shock / current free regions maskval = 1e-4 jy_mask = j_block['y'].data < maskval masked_bsq = 1.0 * bsq masked_bsq.data = np.ma.masked_where(jy_mask, bsq) xcc = j_block.get_crd_cc('x') nx = len(xcc) mp_xloc = np.argmin(masked_bsq, axis=0) # indices mp_xloc = mp_xloc.wrap(xcc[mp_xloc.data]) # location pp_max = np.max(pp_block, axis=0) pp_max_xloc = np.argmax(pp_block, axis=0) # indices pp_max_xloc = pp_max_xloc.wrap(xcc[pp_max_xloc.data]) # location epar = viscid.project(e_block, b_block) epar_max = np.max(epar, axis=0) epar_max_xloc = np.argmax(epar, axis=0) # indices epar_max_xloc = pp_max_xloc.wrap(xcc[epar_max_xloc.data]) # location _ret = find_mp_edges(j_block, 0.1, 0.1, maskval=maskval) sheath_edge, msphere_edge, mp_width, sheath_ind, sphere_ind = _ret # extract b and b**2 at sheath + 6 grid points and sphere - 6 grid pointns # clipping cases where things go outside the block. clipped ponints are # set to nan step = 6 # extract b if b_block.layout == "flat": comp_axis = 0 ic, _, iy, iz = np.ix_(*[np.arange(si) for si in b_block.shape]) ix = np.clip(sheath_ind + step, 0, nx - 1) b_sheath = b_block.data[ic, ix, iy, iz] ix = np.clip(sheath_ind - step, 0, nx - 1) b_sphere = b_block.data[ic, ix, iy, iz] elif b_block.layout == "interlaced": comp_axis = 3 _, iy, iz = np.ix_(*[np.arange(si) for si in b_block.shape[:-1]]) ix = np.clip(sheath_ind + step, 0, nx - 1) b_sheath = b_block.data[ix, iy, iz] ix = np.clip(sheath_ind - step, 0, nx - 1) b_sphere = b_block.data[ix, iy, iz] # extract b**2 bmag_sheath = np.sqrt(np.sum(b_sheath**2, axis=comp_axis)) bmag_sphere = np.sqrt(np.sum(b_sphere**2, axis=comp_axis)) costheta = (np.sum(b_sheath * b_sphere, axis=comp_axis) / (bmag_sphere * bmag_sheath)) costheta = np.where( (sheath_ind + step < nx) & (sphere_ind - step >= 0), costheta, np.nan) mp_shear = mp_width.wrap((180.0 / np.pi) * np.arccos(costheta)) # don't bother with pretty name since it's not written to file # plane_crds = b_block.crds.slice_keep('x=0', cc=True) # fld_kwargs = dict(center="Cell", time=b.time) mp_width.name = "mp_width" mp_xloc.name = "mp_xloc" sheath_edge.name = "mp_sheath_edge" msphere_edge.name = "mp_sphere_edge" mp_shear.name = "mp_shear" pp_max.name = "pp_max" pp_max_xloc.name = "pp_max_xloc" epar_max.name = "epar_max" epar_max_xloc.name = "epar_max_xloc" mp_info = {} mp_info["mp_width"] = mp_width mp_info["mp_xloc"] = mp_xloc mp_info["mp_sheath_edge"] = sheath_edge mp_info["mp_sphere_edge"] = msphere_edge mp_info["mp_shear"] = mp_shear mp_info["pp_max"] = pp_max mp_info["pp_max_xloc"] = pp_max_xloc mp_info["epar_max"] = epar_max mp_info["epar_max_xloc"] = epar_max_xloc # cache new fields to disk if mp_fname: viscid.save_fields(mp_fname + ".h5", mp_info.values()) try: _paraboloid_params = fit_paraboloid(mp_info[fit], p0=fit_p0) mp_info["paraboloid"] = _paraboloid_params except ImportError as _exception: try: msg = _exception.message except AttributeError: msg = _exception.msg mp_info["paraboloid"] = viscid.DeferredImportError(msg) mp_info["mp_width"].pretty_name = "Magnetopause Width" mp_info["mp_xloc"].pretty_name = "Magnetopause $X_{gse}$ Location" mp_info["mp_sheath_edge"].pretty_name = "Magnetosheath Edge" mp_info["mp_sphere_edge"].pretty_name = "Magnetosphere Edge" mp_info["mp_shear"].pretty_name = "Magnetic Shear" mp_info["pp_max"].pretty_name = "Max Pressure" mp_info["pp_max_xloc"].pretty_name = "Max Pressure Location" mp_info["epar_max"].pretty_name = "Max E Parallel" mp_info["epar_max_xloc"].pretty_name = "Max E Parallel Location" return mp_info
def prepare_lines(lines, scalars=None, do_connections=False, other=None): """Concatenate and standardize a list of lines Args: lines (list): Must be a list of 3xN or 4xN ndarrays of xyz(s) data for N points along the line. N need not be the same for all lines. Can alse be 6xN such that lines[:][3:, :] are interpreted as rgb colors scalars (sequence): can be one of:: - single hex color (ex, `#FD7400`) - sequence of Nlines hex colors - single rgb(a) tuple - sequence of Nlines rgb(a) sequences - sequence of N values that go with each vertex and will be mapped with a colormap - sequence of Nlines values that go each line and will be mapped with a colormap - Field. If `np.prod(fld.shape) in (N, nlines)` then the field is interpreted as a simple sequence of values (ie, the topology result from calc_streamlines for coloring each line). Otherwise, the field is interpolated onto each vertex. do_connections (bool): Whether or not to make connections array other (dict): a dictionary of other arrays that should be reshaped and the like the same way scalars is Returns: (vertices, scalars, connections, other) * vertices (ndarray): 3xN array of N xyz points. N is the sum of the lengths of all the lines * scalars (ndarray): N array of scalars, 3xN array of uint8 rgb values, 4xN array of uint8 rgba values, or None * connections (ndarray): Nx2 array of ints (indices along axis 1 of vertices) describing the forward and backward connectedness of the lines, or None * other (dict): a dict of N length arrays Raises: ValueError: If rgb data is not in a valid range or the shape of scalars is not understood """ nlines = len(lines) npts = [line.shape[1] for line in lines] N = np.sum(npts) first_idx = np.cumsum([0] + npts[:-1]) vertices = [np.asarray(line) for line in lines] vertices = np.concatenate(lines, axis=1) if vertices.dtype.kind not in 'fc': vertices = np.asarray(vertices, dtype='f') if vertices.shape[0] > 3: if scalars is not None: viscid.logger.warn("Overriding line scalars with scalars kwarg") else: scalars = vertices[3:, :] vertices = vertices[:3, :] if scalars is not None: scalars_are_strings = False if isinstance(scalars, viscid.field.Field): if np.prod(scalars.shape) in (nlines, N): scalars = np.asarray(scalars).reshape(-1) else: scalars = viscid.interp_trilin(scalars, vertices) if scalars.size != N: raise ValueError("Scalars was not a scalar field") elif isinstance(scalars, (list, tuple, viscid.string_types, np.ndarray)): # string types need some extra massaging if any(isinstance(s, viscid.string_types) for s in scalars): assert all(isinstance(s, viscid.string_types) for s in scalars) scalars_are_strings = True scalars = _string_colors_as_hex(scalars) elif isinstance(scalars, np.ndarray): scalars = scalars else: for i, si in enumerate(scalars): if not isinstance(si, np.ndarray): scalars[i] = np.asarray(si) scalars[i] = np.atleast_2d(scalars[i]) try: scalars = np.concatenate(scalars, axis=0) except ValueError: scalars = np.concatenate(scalars, axis=1) scalars = np.atleast_2d(scalars) if scalars.dtype == np.dtype('object'): raise RuntimeError("Scalars must be numbers, tuples of numbers " "that indicate rgb(a), or hex strings - they " "must not be python objects") if scalars.shape == (1, 1): scalars = scalars.repeat(N, axis=1) elif scalars.shape == (1, nlines) or scalars.shape == (nlines, 1): # one scalar for each line, so broadcast it scalars = scalars.reshape(nlines, 1) scalars = [scalars[i].repeat(ni) for i, ni in enumerate(npts)] scalars = np.concatenate(scalars, axis=0).reshape(1, N) elif scalars.shape == (N, 1) or scalars.shape == (1, N): # catch these so they're not interpreted as colors if # nlines == 1 and N == 3; ie. 1 line with 3 points scalars = scalars.reshape(1, N) elif scalars.shape in [(3, nlines), (nlines, 3), (4, nlines), (nlines, 4)]: # one rgb(a) color for each line, so broadcast it if (scalars.shape in [(3, nlines), (4, nlines)] and scalars.shape not in [(3, 3), (4, 4)]): # the guard against shapes (3, 3) and (4, 4) mean that # these square shapes are assumed Nlines x {3,4} scalars = scalars.T nccomps = scalars.shape[1] # 3 for rgb, 4 for rgba colors = [] for i, ni in enumerate(npts): c = scalars[i].reshape(nccomps, 1).repeat(ni, axis=1) colors.append(c) scalars = np.concatenate(colors, axis=1) elif scalars.shape in [(3, N), (N, 3), (4, N), (N, 4)]: # one rgb(a) color for each vertex if (scalars.shape in [(3, N), (4, N)] and scalars.shape not in [(3, 3), (4, 4)]): # the guard against shapes (3, 3) and (4, 4) mean that # these square shapes are assumed N x {3,4} scalars = scalars.T elif scalars.shape in [(1, 3), (3, 1), (1, 4), (4, 1)]: # interpret a single rgb(a) color, and repeat/broadcast it scalars = scalars.reshape([-1, 1]).repeat(N, axis=1) else: scalars = scalars.reshape(-1, N) # scalars now has shape (1, N), (3, N), or (4, N) if scalars_are_strings: # translate hex colors (#ff00ff) into rgb(a) values scalars = np.char.lstrip(scalars, '#') strlens = np.char.str_len(scalars) min_strlen, max_strlen = np.min(strlens), np.max(strlens) if min_strlen == max_strlen == 8: # 32-bit rgba (two hex chars per channel) scalars = _hexchar2int(scalars.astype('S8')).reshape(-1, 4).T elif min_strlen == max_strlen == 6: # 24-bit rgb (two hex chars per channel) scalars = _hexchar2int(scalars.astype('S6')).reshape(-1, 3).T else: raise NotImplementedError("This should never happen as " "scalars as colors should already " "be preprocessed appropriately") elif scalars.shape[0] == 1: # normal scalars, cast them down to a single dimension scalars = scalars.reshape(-1) elif scalars.shape[0] in (3, 4): # The scalars encode rgb data, standardize the result to a # 3xN or 4xN ndarray of 1 byte unsigned ints [0..255] if np.all(scalars >= 0) and np.all(scalars <= 1): scalars = (255 * scalars).round().astype('u1') elif np.all(scalars >= 0) and np.all(scalars < 256): scalars = scalars.round().astype('u1') else: raise ValueError("Rgb data should be in range [0, 1] or " "[0, 255], range given is [{0}, {1}]" "".format(np.min(scalars), np.max(scalars))) else: raise ValueError("Scalars should either be a number, or set of " "rgb values, shape is {0}".format(scalars.shape)) # scalars should now have shape (N, ) or be a uint8 array with shape # (3, N) or (4, N) encoding an rgb(a) color for each point [0..255] # ... done with scalars... # broadcast / reshape additional arrays given in other if other: for key, arr in other.items(): if arr is None: pass elif arr.shape == (1, nlines) or arr.shape == (nlines, 1): arr = arr.reshape(nlines, 1) arr = [arr[i].repeat(ni) for i, ni in enumerate(npts)] other[key] = np.concatenate(arr, axis=0).reshape(1, N) else: try: other[key] = arr.reshape(-1, N) except ValueError: viscid.logger.warn("Unknown dimension, dropping array {0}" "".format(key)) if do_connections: connections = [None] * nlines for i, ni in enumerate(npts): # i0 is the index of the first point of the i'th line in lines i0 = first_idx[i] connections[i] = np.vstack([ np.arange(i0, i0 + ni - 1.5), np.arange(i0 + 1, i0 + ni - 0.5) ]).T connections = np.concatenate(connections, axis=0).astype('i') else: connections = None return vertices, scalars, connections, other
def _main(): global offscreen_vlab parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--notwo", dest='notwo', action="store_true") parser.add_argument("--nothree", dest='nothree', action="store_true") parser.add_argument("--show", "--plot", action="store_true") args = viscid.vutil.common_argparse(parser, default_verb=0) plot2d = not args.notwo plot3d = not args.nothree # plot2d = True # plot3d = True # args.show = True offscreen_vlab = not args.show img = np.load(os.path.join(sample_dir, "logo.npy")) x = np.linspace(-1, 1, img.shape[0]) y = np.linspace(-1, 1, img.shape[1]) z = np.linspace(-1, 1, img.shape[2]) logo = viscid.arrays2field([x, y, z], img) if 1: viscid.logger.info('Testing Point with custom local coordinates...') pts = np.vstack([[-1, -0.5, 0, 0.5, 1], [-1, -0.5, 0, 0.5, 1], [ 0, 0.5, 1, 1.5, 2]]) local_crds = viscid.asarray_datetime64([0, 60, 120, 180, 240], conservative=True) seeds = viscid.Point(pts, local_crds=local_crds) run_test(logo, seeds, plot2d=plot2d, plot3d=plot3d, show=args.show) if 1: viscid.logger.info('Testing Line...') seeds = viscid.Line([-1, -1, 0], [1, 1, 2], n=5) run_test(logo, seeds, plot2d=plot2d, plot3d=plot3d, show=args.show) if 1: viscid.logger.info('Testing Plane...') seeds = viscid.Plane([0.0, 0.0, 0.0], [1, 1, 1], [1, 0, 0], 2, 2, nl=160, nm=170, NL_are_vectors=True) run_test(logo, seeds, plot2d=plot2d, plot3d=plot3d, show=args.show) if 1: viscid.logger.info('Testing Volume...') seeds = viscid.Volume([-0.8, -0.8, -0.8], [0.8, 0.8, 0.8], n=[64, 64, 3]) # note: can't make a 2d plot of the volume w/o a slice run_test(logo, seeds, plot2d=False, plot3d=plot3d, add_title="3d", show=args.show) if 1: viscid.logger.info('Testing Volume (with ignorable dim)...') seeds = viscid.Volume([-0.8, -0.8, 0.0], [0.8, 0.8, 0.0], n=[64, 64, 1]) run_test(logo, seeds, plot2d=plot2d, plot3d=plot3d, add_title="2d", show=args.show) if 1: viscid.logger.info('Testing Spherical Sphere (phi, theta)...') seeds = viscid.Sphere([0, 0, 0], r=1.0, ntheta=160, nphi=170, pole=[-1, -1, -1], theta_phi=False) run_test(logo, seeds, plot2d=plot2d, plot3d=plot3d, add_title="PT", show=args.show) if 1: viscid.logger.info('Testing Spherical Sphere (theta, phi)...') seeds = viscid.Sphere([0, 0, 0], r=1.0, ntheta=160, nphi=170, pole=[-1, -1, -1], theta_phi=True) run_test(logo, seeds, plot2d=plot2d, plot3d=plot3d, add_title="TP", show=args.show) if 1: viscid.logger.info('Testing Spherical Cap (phi, theta)...') seeds = viscid.SphericalCap(p0=[0, 0, 0], r=1.0, ntheta=64, nphi=80, pole=[-1, -1, -1], theta_phi=False) run_test(logo, seeds, plot2d=plot2d, plot3d=plot3d, add_title="PT", view_kwargs=dict(azimuth=180, elevation=180), show=args.show) if 1: viscid.logger.info('Testing Spherical Cap (theta, phi)...') seeds = viscid.SphericalCap(p0=[0, 0, 0], r=1.0, ntheta=64, nphi=80, pole=[-1, -1, -1], theta_phi=True) run_test(logo, seeds, plot2d=plot2d, plot3d=plot3d, add_title="TP", view_kwargs=dict(azimuth=180, elevation=180), show=args.show) if 1: viscid.logger.info('Testing Spherical Patch...') seeds = viscid.SphericalPatch(p0=[0, 0, 0], p1=[0, -0, -1], max_alpha=30.0, max_beta=59.9, nalpha=65, nbeta=80, r=0.5, roll=45.0) run_test(logo, seeds, plot2d=plot2d, plot3d=plot3d, show=args.show) if 1: # this spline test is very custom viscid.logger.info('Testing Spline...') try: import scipy.interpolate as interpolate except ImportError: msg = "XFail: ImportError (is scipy installed?)" if plot2d: try: from viscid.plot import vpyplot as vlt from matplotlib import pyplot as plt plt.clf() plt.annotate(msg, xy=(0.3, 0.4), xycoords='axes fraction') plt.savefig(next_plot_fname(__file__, series='2d')) plt.savefig(next_plot_fname(__file__, series='2d')) plt.savefig(next_plot_fname(__file__, series='3d')) if args.show: plt.show() except ImportError: pass else: knots = np.array([[ 0.2, 0.5, 0.0], [-0.2, 0.5, 0.2], [-0.2, 0.0, 0.4], [ 0.2, 0.0, 0.2], [ 0.2, -0.5, 0.0], [-0.2, -0.5, 0.2]]).T seed_name = "Spline" fld = logo seeds = viscid.Spline(knots) seed_pts = seeds.get_points() interp_fld = viscid.interp_trilin(fld, seeds) if plot2d: try: from viscid.plot import vpyplot as vlt from matplotlib import pyplot as plt plt.clf() vlt.plot(interp_fld) plt.title(seed_name) plt.savefig(next_plot_fname(__file__, series='2d')) if args.show: plt.show() plt.clf() from matplotlib import rcParams _ms = rcParams['lines.markersize'] plt.gca().scatter(knots[0, :], knots[1, :], s=(2 * _ms)**2, marker='^', color='y') plt.gca().scatter(seed_pts[0, :], seed_pts[1, :], s=(1.5 * _ms)**2, marker='o', color='k') vlt.plot2d_line(seed_pts, scalars=interp_fld.flat_data, symdir='z') plt.title(seed_name) plt.savefig(next_plot_fname(__file__, series='2d')) if args.show: plt.show() except ImportError: pass if plot3d: try: vlab, _ = get_mvi_fig() vlab.points3d(knots[0], knots[1], knots[2], color=(1.0, 1.0, 0), scale_mode='none', scale_factor=0.04) p = vlab.points3d(seed_pts[0], seed_pts[1], seed_pts[2], color=(0, 0, 0), scale_mode='none', scale_factor=0.03) vlab.plot_line(seed_pts, scalars=interp_fld.flat_data, tube_radius=0.01) vlab.axes(p) vlab.title(seed_name) vlab.mlab.roll(-90.0) vlab.savefig(next_plot_fname(__file__, series='3d')) if args.show: vlab.show(stop=True) except ImportError: pass if 1: viscid.logger.info('Testing RectilinearMeshPoints...') f = viscid.load_file(os.path.join(sample_dir, 'sample_xdmf.3d.[-1].xdmf')) slc = 'x=-40j:12j, y=-10j:10j, z=-10j:10j' b = f['b'][slc] z = b.get_crd('z') sheet_iz = np.argmin(b['x']**2, axis=2) sheet_pts = b['z=0:1'].get_points() sheet_pts[2, :] = z[sheet_iz].reshape(-1) isphere_mask = np.sum(sheet_pts[:2, :]**2, axis=0) < 5**2 day_mask = sheet_pts[0:1, :] > -1.0 sheet_pts[2, :] = np.choose(isphere_mask, [sheet_pts[2, :], 0]) sheet_pts[2, :] = np.choose(day_mask, [sheet_pts[2, :], 0]) nx, ny, _ = b.sshape sheet_seed = viscid.RectilinearMeshPoints(sheet_pts.reshape(3, nx, ny)) vx_sheet = viscid.interp_nearest(f['vx'], sheet_seed) try: if not plot2d: raise ImportError from viscid.plot import vpyplot as vlt from matplotlib import pyplot as plt vlt.clf() vlt.plot(vx_sheet, symmetric=True) plt.savefig(next_plot_fname(__file__, series='2d')) if args.show: vlt.show() except ImportError: pass try: if not plot3d: raise ImportError vlab, _ = get_mvi_fig() mesh = vlab.mesh_from_seeds(sheet_seed, scalars=vx_sheet, clim=(-400, 400)) vlab.plot_earth_3d(crd_system=b) vlab.view(azimuth=+90.0 + 45.0, elevation=90.0 - 25.0, distance=30.0, focalpoint=(-10.0, +1.0, +1.0)) vlab.title("RectilinearMeshPoints") vlab.savefig(next_plot_fname(__file__, series='3d')) if args.show: vlab.show(stop=True) except ImportError: pass # prevent weird xorg bad-instructions on tear down if 'figure' in _global_ns and _global_ns['figure'] is not None: from viscid.plot import vlab vlab.mlab.close(_global_ns['figure']) return 0
def main(): mhd_type = "C3" make_plots = 1 mhd_type = mhd_type.upper() if mhd_type.startswith("C"): if mhd_type in ("C",): f = viscid.load_file("$WORK/tmedium/*.3d.[-1].xdmf") elif mhd_type in ("C2", "C3"): f = viscid.load_file("$WORK/tmedium2/*.3d.[-1].xdmf") else: raise ValueError() catol = 1e-8 rtol = 2e-6 elif mhd_type in ("F", "FORTRAN"): f = viscid.load_file("$WORK/tmedium3/*.3df.[-1]") catol = 1e-8 rtol = 7e-2 else: raise ValueError() b = f['b_cc'] b1 = f['b_fc'] e_cc = f['e_cc'] e_ec = f['e_ec'] # divb = f['divB'] # viscid.interact() if True: bD = viscid.empty_like(b) bD.data = np.array(b.data) b1D = viscid.empty_like(b1) b1D.data = np.array(b1.data) mask5 = viscid.make_spherical_mask(bD, rmax=3.5) mask1_5 = viscid.make_spherical_mask(bD, rmax=1.5) viscid.fill_dipole(bD, mask=mask5) viscid.set_in_region(bD, bD, 0.0, 0.0, mask=mask1_5, out=bD) # compare_vectors(_b, bD, make_plots=True) mask5 = viscid.make_spherical_mask(b1D, rmax=3.5) mask1_5 = viscid.make_spherical_mask(b1D, rmax=1.5) viscid.fill_dipole(b1D, mask=mask5) viscid.set_in_region(b1D, b1D, 0.0, 0.0, mask=mask1_5, out=b1D) compare_vectors(bD["x=1:-1, y=1:-1, z=1:-1"], b1D.as_cell_centered(), make_plots=True) # plt.clf() # dkwargs = dict(symmetric=True, earth=True, clim=(-1e2, 1e2)) # ax1 = plt.subplot(311) # vlt.plot(viscid.div(b1)['y=0j'], **dkwargs) # plt.subplot(312, sharex=ax1, sharey=ax1) # vlt.plot(viscid.div(b)['y=0j'], **dkwargs) # plt.subplot(313, sharex=ax1, sharey=ax1) # vlt.plot(viscid.div(b1D)['y=0j'], **dkwargs) # vlt.show() bD = b1D = mask5 = mask1_5 = None # straight up interpolate b1 to cc crds and compare with b if True: b1_cc = viscid.interp_trilin(b1, b).as_flat() viscid.set_in_region(b, b, alpha=0.0, beta=0.0, out=b, mask=viscid.make_spherical_mask(b, rmax=5.0)) viscid.set_in_region(b1_cc, b1_cc, alpha=0.0, beta=0.0, out=b1_cc, mask=viscid.make_spherical_mask(b1_cc, rmax=5.0)) compare_vectors(b, b1_cc, make_plots=True) # make div? if True: # make seeds for 1.5x supersampling b1 n = 128 seeds = viscid.Volume((5.1, -0.02, -5.0), (12.0, 0.02, 5.0), (n, 3, n)) # do interpolation onto new seeds b2 = viscid.interp_trilin(b1, seeds) div_b = viscid.div(b) div_b1 = viscid.div(b1) div_b2 = viscid.div(b2) viscid.set_in_region(div_b, div_b, alpha=0.0, beta=0.0, out=div_b, mask=viscid.make_spherical_mask(div_b, rmax=5.0)) viscid.set_in_region(div_b1, div_b1, alpha=0.0, beta=0.0, out=div_b1, mask=viscid.make_spherical_mask(div_b1, rmax=5.0)) viscid.set_in_region(div_b2, div_b2, alpha=0.0, beta=0.0, out=div_b2, mask=viscid.make_spherical_mask(div_b2, rmax=5.0)) viscid.set_in_region(divb, divb, alpha=0.0, beta=0.0, out=divb, mask=viscid.make_spherical_mask(divb, rmax=5.0)) plt.clf() ax1 = vlt.subplot(311) vlt.plot(div_b['y=0j'], symmetric=True, earth=True) vlt.subplot(312, sharex=ax1, sharey=ax1) # vlt.plot(div_b1['y=0j'], symmetric=True, earth=True) vlt.plot(div_b2['y=0j'], symmetric=True, earth=True) vlt.subplot(313, sharex=ax1, sharey=ax1) vlt.plot(divb['y=0j'], symmetric=True, earth=True) vlt.show() return 0
def get_mp_info(pp, b, j, e, cache=True, cache_dir=None, slc="x=5.5j:11.0j, y=-4.0j:4.0j, z=-3.6j:3.6j", fit="mp_xloc", fit_p0=(9.0, 0.0, 0.0, 1.0, -1.0, -1.0)): """Get info about m-pause as flattened fields Notes: The first thing this function does is mask locations where the GSE-y current density < 1e-4. This masks out the bow shock and current free regions. This works for southward IMF, but it is not very general. Parameters: pp (ScalarcField): pressure b (VectorField): magnetic field j (VectorField): current density e (VectorField, None): electric field (same centering as b). If None, then the info that requires E will be filled with NaN cache (bool, str): Save to and load from cache, if "force", then don't load from cache if it exists, but do save a cache at the end cache_dir (str): Directory for cache, if None, same directory as that file to which the grid belongs slc (str): slice that gives a box that contains the m-pause fit (str): to which resulting field should the paraboloid be fit, defaults to mp_xloc, but pp_max_xloc might be useful in some circumstances fit_p0 (tuple): Initial guess vector for paraboloid fit Returns: dict: Unless otherwise noted, the entiries are 2D (y-z) fields - **mp_xloc** location of minimum abs(Bz), this works better than max of J^2 for FTEs - **mp_sheath_edge** location where Jy > 0.1 * Jy when coming in from the sheath side - **mp_sphere_edge** location where Jy > 0.1 * Jy when coming in from the sphere side - **mp_width** difference between m-sheath edge and msphere edge - **mp_shear** magnetic shear taken 6 grid points into the m-sheath / m-sphere - **pp_max** max pp - **pp_max_xloc** location of max pp - **epar_max** max e parallel - **epar_max_xloc** location of max e parallel - **paraboloid** numpy.recarray of paraboloid fit. The parameters are given in the 0th element, and the 1st element contains the 1-sigma values for the fit Raises: RuntimeError: if using MHD crds instead of GSE crds """ if not cache_dir: cache_dir = pp.find_info("_viscid_dirname", "./") run_name = pp.find_info("run", None) if cache and run_name: t = pp.time mp_fname = "{0}/{1}.mpause.{2:06.0f}".format(cache_dir, run_name, t) else: mp_fname = "" try: force = cache.strip().lower() == "force" except AttributeError: force = False try: if force or not mp_fname or not os.path.isfile(mp_fname + ".xdmf"): raise IOError() mp_info = {} with viscid.load_file(mp_fname + ".xdmf") as dat: fld_names = ["mp_xloc", "mp_sheath_edge", "mp_sphere_edge", "mp_width", "mp_shear", "pp_max", "pp_max_xloc", "epar_max", "epar_max_xloc"] for fld_name in fld_names: mp_info[fld_name] = dat[fld_name]["x=0"] except (IOError, KeyError): mp_info = {} crd_system = viscid.as_crd_system(b, None) if crd_system != 'gse': raise RuntimeError("get_mp_info can't work in MHD crds, " "switch to GSE please") if j.nr_patches == 1: pp_block = pp[slc] b_block = b[slc] j_block = j[slc] if e is None: e_block = np.nan * viscid.empty_like(j_block) else: e_block = e[slc] else: # interpolate an amr grid so we can proceed obnd = pp.get_slice_extent(slc) dx = np.min(pp.skeleton.L / pp.skeleton.n, axis=0) nx = np.ceil((obnd[1] - obnd[0]) / dx) vol = viscid.seed.Volume(obnd[0], obnd[1], nx, cache=True) pp_block = vol.wrap_field(viscid.interp_trilin(pp, vol), name="P").as_cell_centered() b_block = vol.wrap_field(viscid.interp_trilin(b, vol), name="B").as_cell_centered() j_block = vol.wrap_field(viscid.interp_trilin(j, vol), name="J").as_cell_centered() if e is None: e_block = np.nan * viscid.empty_like(j_block) else: e_block = vol.wrap_field(viscid.interp_trilin(e, vol), name="E").as_cell_centered() # jsq = viscid.dot(j_block, j_block) bsq = viscid.dot(b_block, b_block) # extract ndarrays and mask out bow shock / current free regions maskval = 1e-4 jy_mask = j_block['y'].data < maskval masked_bsq = 1.0 * bsq masked_bsq.data = np.ma.masked_where(jy_mask, bsq) xcc = j_block.get_crd_cc('x') nx = len(xcc) mp_xloc = np.argmin(masked_bsq, axis=0) # indices mp_xloc = mp_xloc.wrap(xcc[mp_xloc.data]) # location pp_max = np.max(pp_block, axis=0) pp_max_xloc = np.argmax(pp_block, axis=0) # indices pp_max_xloc = pp_max_xloc.wrap(xcc[pp_max_xloc.data]) # location epar = viscid.project(e_block, b_block) epar_max = np.max(epar, axis=0) epar_max_xloc = np.argmax(epar, axis=0) # indices epar_max_xloc = pp_max_xloc.wrap(xcc[epar_max_xloc.data]) # location _ret = find_mp_edges(j_block, 0.1, 0.1, maskval=maskval) sheath_edge, msphere_edge, mp_width, sheath_ind, sphere_ind = _ret # extract b and b**2 at sheath + 6 grid points and sphere - 6 grid pointns # clipping cases where things go outside the block. clipped ponints are # set to nan step = 6 # extract b if b_block.layout == "flat": comp_axis = 0 ic, _, iy, iz = np.ix_(*[np.arange(si) for si in b_block.shape]) ix = np.clip(sheath_ind + step, 0, nx - 1) b_sheath = b_block.data[ic, ix, iy, iz] ix = np.clip(sheath_ind - step, 0, nx - 1) b_sphere = b_block.data[ic, ix, iy, iz] elif b_block.layout == "interlaced": comp_axis = 3 _, iy, iz = np.ix_(*[np.arange(si) for si in b_block.shape[:-1]]) ix = np.clip(sheath_ind + step, 0, nx - 1) b_sheath = b_block.data[ix, iy, iz] ix = np.clip(sheath_ind - step, 0, nx - 1) b_sphere = b_block.data[ix, iy, iz] # extract b**2 bmag_sheath = np.sqrt(np.sum(b_sheath**2, axis=comp_axis)) bmag_sphere = np.sqrt(np.sum(b_sphere**2, axis=comp_axis)) costheta = (np.sum(b_sheath * b_sphere, axis=comp_axis) / (bmag_sphere * bmag_sheath)) costheta = np.where((sheath_ind + step < nx) & (sphere_ind - step >= 0), costheta, np.nan) mp_shear = mp_width.wrap((180.0 / np.pi) * np.arccos(costheta)) # don't bother with pretty name since it's not written to file # plane_crds = b_block.crds.slice_keep('x=0', cc=True) # fld_kwargs = dict(center="Cell", time=b.time) mp_width.name = "mp_width" mp_xloc.name = "mp_xloc" sheath_edge.name = "mp_sheath_edge" msphere_edge.name = "mp_sphere_edge" mp_shear.name = "mp_shear" pp_max.name = "pp_max" pp_max_xloc.name = "pp_max_xloc" epar_max.name = "epar_max" epar_max_xloc.name = "epar_max_xloc" mp_info = {} mp_info["mp_width"] = mp_width mp_info["mp_xloc"] = mp_xloc mp_info["mp_sheath_edge"] = sheath_edge mp_info["mp_sphere_edge"] = msphere_edge mp_info["mp_shear"] = mp_shear mp_info["pp_max"] = pp_max mp_info["pp_max_xloc"] = pp_max_xloc mp_info["epar_max"] = epar_max mp_info["epar_max_xloc"] = epar_max_xloc # cache new fields to disk if mp_fname: viscid.save_fields(mp_fname + ".h5", list(mp_info.values())) try: _paraboloid_params = fit_paraboloid(mp_info[fit], p0=fit_p0) mp_info["paraboloid"] = _paraboloid_params except ImportError as _exception: try: msg = _exception.message except AttributeError: msg = _exception.msg mp_info["paraboloid"] = viscid.DeferredImportError(msg) mp_info["mp_width"].pretty_name = "Magnetopause Width" mp_info["mp_xloc"].pretty_name = "Magnetopause $X_{gse}$ Location" mp_info["mp_sheath_edge"].pretty_name = "Magnetosheath Edge" mp_info["mp_sphere_edge"].pretty_name = "Magnetosphere Edge" mp_info["mp_shear"].pretty_name = "Magnetic Shear" mp_info["pp_max"].pretty_name = "Max Pressure" mp_info["pp_max_xloc"].pretty_name = "Max Pressure Location" mp_info["epar_max"].pretty_name = "Max E Parallel" mp_info["epar_max_xloc"].pretty_name = "Max E Parallel Location" return mp_info
def main(): mhd_type = "C3" make_plots = 1 mhd_type = mhd_type.upper() if mhd_type.startswith("C"): if mhd_type in ("C", ): f = viscid.load_file("$WORK/tmedium/*.3d.[-1].xdmf") elif mhd_type in ("C2", "C3"): f = viscid.load_file("$WORK/tmedium2/*.3d.[-1].xdmf") else: raise ValueError() catol = 1e-8 rtol = 2e-6 elif mhd_type in ("F", "FORTRAN"): f = viscid.load_file("$WORK/tmedium3/*.3df.[-1]") catol = 1e-8 rtol = 7e-2 else: raise ValueError() b = f['b_cc'] b1 = f['b_fc'] e_cc = f['e_cc'] e_ec = f['e_ec'] # divb = f['divB'] # viscid.interact() if True: bD = viscid.empty_like(b) bD.data = np.array(b.data) b1D = viscid.empty_like(b1) b1D.data = np.array(b1.data) mask5 = viscid.make_spherical_mask(bD, rmax=3.5) mask1_5 = viscid.make_spherical_mask(bD, rmax=1.5) viscid.fill_dipole(bD, mask=mask5) viscid.set_in_region(bD, bD, 0.0, 0.0, mask=mask1_5, out=bD) # compare_vectors(_b, bD, make_plots=True) mask5 = viscid.make_spherical_mask(b1D, rmax=3.5) mask1_5 = viscid.make_spherical_mask(b1D, rmax=1.5) viscid.fill_dipole(b1D, mask=mask5) viscid.set_in_region(b1D, b1D, 0.0, 0.0, mask=mask1_5, out=b1D) compare_vectors(bD["x=1:-1, y=1:-1, z=1:-1"], b1D.as_cell_centered(), make_plots=True) # plt.clf() # dkwargs = dict(symmetric=True, earth=True, clim=(-1e2, 1e2)) # ax1 = plt.subplot(311) # vlt.plot(viscid.div(b1)['y=0j'], **dkwargs) # plt.subplot(312, sharex=ax1, sharey=ax1) # vlt.plot(viscid.div(b)['y=0j'], **dkwargs) # plt.subplot(313, sharex=ax1, sharey=ax1) # vlt.plot(viscid.div(b1D)['y=0j'], **dkwargs) # vlt.show() bD = b1D = mask5 = mask1_5 = None # straight up interpolate b1 to cc crds and compare with b if True: b1_cc = viscid.interp_trilin(b1, b).as_flat() viscid.set_in_region(b, b, alpha=0.0, beta=0.0, out=b, mask=viscid.make_spherical_mask(b, rmax=5.0)) viscid.set_in_region(b1_cc, b1_cc, alpha=0.0, beta=0.0, out=b1_cc, mask=viscid.make_spherical_mask(b1_cc, rmax=5.0)) compare_vectors(b, b1_cc, make_plots=True) # make div? if True: # make seeds for 1.5x supersampling b1 n = 128 seeds = viscid.Volume((5.1, -0.02, -5.0), (12.0, 0.02, 5.0), (n, 3, n)) # do interpolation onto new seeds b2 = viscid.interp_trilin(b1, seeds) div_b = viscid.div(b) div_b1 = viscid.div(b1) div_b2 = viscid.div(b2) viscid.set_in_region(div_b, div_b, alpha=0.0, beta=0.0, out=div_b, mask=viscid.make_spherical_mask(div_b, rmax=5.0)) viscid.set_in_region(div_b1, div_b1, alpha=0.0, beta=0.0, out=div_b1, mask=viscid.make_spherical_mask(div_b1, rmax=5.0)) viscid.set_in_region(div_b2, div_b2, alpha=0.0, beta=0.0, out=div_b2, mask=viscid.make_spherical_mask(div_b2, rmax=5.0)) viscid.set_in_region(divb, divb, alpha=0.0, beta=0.0, out=divb, mask=viscid.make_spherical_mask(divb, rmax=5.0)) plt.clf() ax1 = vlt.subplot(311) vlt.plot(div_b['y=0j'], symmetric=True, earth=True) vlt.subplot(312, sharex=ax1, sharey=ax1) # vlt.plot(div_b1['y=0j'], symmetric=True, earth=True) vlt.plot(div_b2['y=0j'], symmetric=True, earth=True) vlt.subplot(313, sharex=ax1, sharey=ax1) vlt.plot(divb['y=0j'], symmetric=True, earth=True) vlt.show() return 0
def _main(): f = viscid.load_file('~/dev/work/xi_fte_001/*.3d.*.xdmf') time_slice = ':' times = np.array([grid.time for grid in f.iter_times(time_slice)]) # XYZ coordinates of virtual satelites in warped "plasma sheet coords" x_sat_psc = np.linspace(-30, 0, 31) # X (GSE == PSC) y_sat_psc = np.linspace(-10, 10, 21) # Y (GSE == PSC) z_sat_psc = np.linspace(-2, 2, 5) # Z in PSC (z=0 is the plasma sheet) # the GSE z location of the virtual satelites in the warped plasma sheet # coordinates, so sat_z_gse_ts['x=5j, y=1j, z=0j'] would give the # plasma sheet location at x=5.0, y=1.0 # These fields depend on time because the plasma sheet moves in time sat_z_gse_ts = viscid.zeros([times, x_sat_psc, y_sat_psc, z_sat_psc], crd_names='txyz', center='node', name='PlasmaSheetZ_GSE') vx_ts = viscid.zeros_like(sat_z_gse_ts) bz_ts = viscid.zeros_like(sat_z_gse_ts) for itime, grid in enumerate(f.iter_times(time_slice)): print("Processing time slice", itime, grid.time) gse_slice = 'x=-35j:0j, y=-15j:15j, z=-6j:6j' bx = grid['bx'][gse_slice] bx_argmin = np.argmin(bx**2, axis=2) z_gse = bx.get_crd('z') # ps_zloc_gse is the plasma sheet z location along the GGCM grid x/y ps_z_gse = viscid.zeros_like(bx[:, :, 0:1]) ps_z_gse[...] = z_gse[bx_argmin] # Note: Here you could apply a gaussian filter to # ps_z_gse[:, :, 0].data in order to smooth the surface # if desired. Scipy / Scikit-Image have some functions # that do this # ok, we found the plasma sheet z GSE location on the actual GGCM # grid, but we just want a subset of that grid for our virtual # satelites, so just interpolate the ps z location to our subset ps_z_gse_subset = viscid.interp_trilin(ps_z_gse, sat_z_gse_ts[itime, :, :, 0:1], wrap=True) # now we know the plasma sheet z location in GSE, and how far # apart we want the satelites in z, so put those two things together # to get a bunch of satelite locations sat_z_gse_ts[itime] = ps_z_gse_subset.data + z_sat_psc.reshape(1, 1, -1) # make a seed generator that we can use to fill the vx and bz # time series for this instant in time sat_loc_gse = sat_z_gse_ts[itime].get_points() sat_loc_gse[2, :] = sat_z_gse_ts[itime].data.reshape(-1) # slicing the field before doing the interpolation makes this # faster for hdf5 data, but probably for other data too vx_ts[itime] = viscid.interp_trilin(grid['vx'][gse_slice], sat_loc_gse, wrap=False ).reshape(vx_ts.shape[1:]) bz_ts[itime] = viscid.interp_trilin(grid['bz'][gse_slice], sat_loc_gse, wrap=False ).reshape(bz_ts.shape[1:]) # 2d plots of the plasma sheet z location to make sure we did the # interpolation correctly if False: # pylint: disable=using-constant-test from viscid.plot import vpyplot as vlt fig, (ax0, ax1) = vlt.subplots(2, 1) # pylint: disable=unused-variable vlt.plot(ps_z_gse, ax=ax0, clim=(-5, 5)) vlt.plot(ps_z_gse_subset, ax=ax1, clim=(-5, 5)) vlt.auto_adjust_subplots() vlt.show() # make a 3d plot of the plasma sheet surface to verify that it # makes sense if True: # pylint: disable=using-constant-test from viscid.plot import vlab fig = vlab.figure(size=(1280, 800), bgcolor=(1, 1, 1), fgcolor=(0, 0, 0)) vlab.clf() # plot the plasma sheet coloured by vx # Note: points closer to x = 0 are unsightly since the plasma # sheet criteria starts to fall apart on the flanks, so # just remove the first few rows ps_z_gse_tail = ps_z_gse['x=:-2.25j'] ps_mesh_shape = [3, ps_z_gse_tail.shape[0], ps_z_gse_tail.shape[1]] ps_pts = ps_z_gse_tail.get_points().reshape(ps_mesh_shape) ps_pts[2, :, :] = ps_z_gse_tail[:, :, 0] plasma_sheet = viscid.RectilinearMeshPoints(ps_pts) ps_vx = viscid.interp_trilin(grid['vx'][gse_slice], plasma_sheet) _ = vlab.mesh_from_seeds(plasma_sheet, scalars=ps_vx) vx_clim = (-1400, 1400) vx_cmap = 'viridis' vlab.colorbar(title='Vx', clim=vx_clim, cmap=vx_cmap, nb_labels=5) # plot satelite locations as dots colored by Vx with the same # limits and color as the plasma sheet mesh sat3d = vlab.points3d(sat_loc_gse[0], sat_loc_gse[1], sat_loc_gse[2], vx_ts[itime].data.reshape(-1), scale_mode='none', scale_factor=0.2) vlab.apply_cmap(sat3d, clim=vx_clim, cmap=vx_cmap) # plot Earth for reference cotr = viscid.Cotr(dip_tilt=0.0) # pylint: disable=not-callable vlab.plot_blue_marble(r=1.0, lines=False, ntheta=64, nphi=128, rotate=cotr, crd_system='mhd') vlab.plot_earth_3d(radius=1.01, night_only=True, opacity=0.5, crd_system='gse') vlab.view(azimuth=45, elevation=70, distance=35.0, focalpoint=[-9, 3, -1]) vlab.savefig('plasma_sheet_3d_{0:02d}.png'.format(itime)) vlab.show() try: vlab.mlab.close(fig) except TypeError: pass # this happens if the figure is already closed # now do what we will with the time series... this is not a good # presentation of this data, but you get the idea from viscid.plot import vpyplot as vlt fig, axes = vlt.subplots(4, 4, figsize=(12, 12)) for ax_row, yloc in zip(axes, np.linspace(-5, 5, len(axes))[::-1]): for ax, xloc in zip(ax_row, np.linspace(4, 7, len(ax_row))): vlt.plot(vx_ts['x={0}j, y={1}j, z=0j'.format(xloc, yloc)], ax=ax) ax.set_ylabel('') vlt.plt.title('x = {0:g}, y = {1:g}'.format(xloc, yloc)) vlt.plt.suptitle('Vx [km/s]') vlt.auto_adjust_subplots() vlt.show() return 0
def _main(): parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--notwo", dest='notwo', action="store_true") parser.add_argument("--nothree", dest='nothree', action="store_true") parser.add_argument("--show", "--plot", action="store_true") args = viscid.vutil.common_argparse(parser, default_verb=0) plot2d = not args.notwo plot3d = not args.nothree # plot2d = True # plot3d = True # args.show = True img = np.load(os.path.join(sample_dir, "logo.npy")) x = np.linspace(-1, 1, img.shape[0]) y = np.linspace(-1, 1, img.shape[1]) z = np.linspace(-1, 1, img.shape[2]) logo = viscid.arrays2field([x, y, z], img) if 1: viscid.logger.info('Testing Point with custom local coordinates...') pts = np.vstack([[-1, -0.5, 0, 0.5, 1], [-1, -0.5, 0, 0.5, 1], [ 0, 0.5, 1, 1.5, 2]]) local_crds = viscid.asarray_datetime64([0, 60, 120, 180, 240], conservative=True) seeds = viscid.Point(pts, local_crds=local_crds) run_test(logo, seeds, plot2d=plot2d, plot3d=plot3d, show=args.show) if 1: viscid.logger.info('Testing Line...') seeds = viscid.Line([-1, -1, 0], [1, 1, 2], n=5) run_test(logo, seeds, plot2d=plot2d, plot3d=plot3d, show=args.show) if 1: viscid.logger.info('Testing Plane...') seeds = viscid.Plane([0.0, 0.0, 0.0], [1, 1, 1], [1, 0, 0], 2, 2, nl=160, nm=170, NL_are_vectors=True) run_test(logo, seeds, plot2d=plot2d, plot3d=plot3d, show=args.show) if 1: viscid.logger.info('Testing Volume...') seeds = viscid.Volume([-0.8, -0.8, -0.8], [0.8, 0.8, 0.8], n=[64, 64, 3]) # note: can't make a 2d plot of the volume w/o a slice run_test(logo, seeds, plot2d=False, plot3d=plot3d, add_title="3d", show=args.show) if 1: viscid.logger.info('Testing Volume (with ignorable dim)...') seeds = viscid.Volume([-0.8, -0.8, 0.0], [0.8, 0.8, 0.0], n=[64, 64, 1]) run_test(logo, seeds, plot2d=plot2d, plot3d=plot3d, add_title="2d", show=args.show) if 1: viscid.logger.info('Testing Spherical Sphere (phi, theta)...') seeds = viscid.Sphere([0, 0, 0], r=1.0, ntheta=160, nphi=170, pole=[-1, -1, -1], theta_phi=False) run_test(logo, seeds, plot2d=plot2d, plot3d=plot3d, add_title="PT", show=args.show) if 1: viscid.logger.info('Testing Spherical Sphere (theta, phi)...') seeds = viscid.Sphere([0, 0, 0], r=1.0, ntheta=160, nphi=170, pole=[-1, -1, -1], theta_phi=True) run_test(logo, seeds, plot2d=plot2d, plot3d=plot3d, add_title="TP", show=args.show) if 1: viscid.logger.info('Testing Spherical Cap (phi, theta)...') seeds = viscid.SphericalCap(p0=[0, 0, 0], r=1.0, ntheta=64, nphi=80, pole=[-1, -1, -1], theta_phi=False) run_test(logo, seeds, plot2d=plot2d, plot3d=plot3d, add_title="PT", view_kwargs=dict(azimuth=180, elevation=180), show=args.show) if 1: viscid.logger.info('Testing Spherical Cap (theta, phi)...') seeds = viscid.SphericalCap(p0=[0, 0, 0], r=1.0, ntheta=64, nphi=80, pole=[-1, -1, -1], theta_phi=True) run_test(logo, seeds, plot2d=plot2d, plot3d=plot3d, add_title="TP", view_kwargs=dict(azimuth=180, elevation=180), show=args.show) if 1: viscid.logger.info('Testing Spherical Patch...') seeds = viscid.SphericalPatch(p0=[0, 0, 0], p1=[0, -0, -1], max_alpha=30.0, max_beta=59.9, nalpha=65, nbeta=80, r=0.5, roll=45.0) run_test(logo, seeds, plot2d=plot2d, plot3d=plot3d, show=args.show) if 1: # this spline test is very custom viscid.logger.info('Testing Spline...') try: import scipy.interpolate as interpolate except ImportError: msg = "XFail: ImportError (is scipy installed?)" if plot2d: try: from viscid.plot import vpyplot as vlt from matplotlib import pyplot as plt plt.clf() plt.annotate(msg, xy=(0.3, 0.4), xycoords='axes fraction') plt.savefig(next_plot_fname(__file__, series='2d')) plt.savefig(next_plot_fname(__file__, series='2d')) plt.savefig(next_plot_fname(__file__, series='3d')) if args.show: plt.show() except ImportError: pass else: knots = np.array([[ 0.2, 0.5, 0.0], [-0.2, 0.5, 0.2], [-0.2, 0.0, 0.4], [ 0.2, 0.0, 0.2], [ 0.2, -0.5, 0.0], [-0.2, -0.5, 0.2]]).T seed_name = "Spline" fld = logo seeds = viscid.Spline(knots) seed_pts = seeds.get_points() interp_fld = viscid.interp_trilin(fld, seeds) if plot2d: try: from viscid.plot import vpyplot as vlt from matplotlib import pyplot as plt plt.clf() vlt.plot(interp_fld) plt.title(seed_name) plt.savefig(next_plot_fname(__file__, series='2d')) if args.show: plt.show() plt.clf() from matplotlib import rcParams _ms = rcParams['lines.markersize'] plt.gca().scatter(knots[0, :], knots[1, :], s=(2 * _ms)**2, marker='^', color='y') plt.gca().scatter(seed_pts[0, :], seed_pts[1, :], s=(1.5 * _ms)**2, marker='o', color='k') vlt.plot2d_line(seed_pts, scalars=interp_fld.flat_data, symdir='z') plt.title(seed_name) plt.savefig(next_plot_fname(__file__, series='2d')) if args.show: plt.show() except ImportError: pass if plot3d: try: from viscid.plot import vlab _ = get_mvi_fig(offscreen=not args.show) vlab.points3d(knots[0], knots[1], knots[2], color=(1.0, 1.0, 0), scale_mode='none', scale_factor=0.04) p = vlab.points3d(seed_pts[0], seed_pts[1], seed_pts[2], color=(0, 0, 0), scale_mode='none', scale_factor=0.03) vlab.plot_line(seed_pts, scalars=interp_fld.flat_data, tube_radius=0.01) vlab.axes(p) vlab.title(seed_name) vlab.mlab.roll(-90.0) vlab.savefig(next_plot_fname(__file__, series='3d')) if args.show: vlab.show(stop=True) except ImportError: pass if 1: viscid.logger.info('Testing RectilinearMeshPoints...') f = viscid.load_file(os.path.join(sample_dir, 'sample_xdmf.3d.[-1].xdmf')) slc = 'x=-40j:12j, y=-10j:10j, z=-10j:10j' b = f['b'][slc] z = b.get_crd('z') sheet_iz = np.argmin(b['x']**2, axis=2) sheet_pts = b['z=0:1'].get_points() sheet_pts[2, :] = z[sheet_iz].reshape(-1) isphere_mask = np.sum(sheet_pts[:2, :]**2, axis=0) < 5**2 day_mask = sheet_pts[0:1, :] > -1.0 sheet_pts[2, :] = np.choose(isphere_mask, [sheet_pts[2, :], 0]) sheet_pts[2, :] = np.choose(day_mask, [sheet_pts[2, :], 0]) nx, ny, _ = b.sshape sheet_seed = viscid.RectilinearMeshPoints(sheet_pts.reshape(3, nx, ny)) vx_sheet = viscid.interp_nearest(f['vx'], sheet_seed) try: if not plot2d: raise ImportError from viscid.plot import vpyplot as vlt from matplotlib import pyplot as plt vlt.clf() vlt.plot(vx_sheet, symmetric=True) plt.savefig(next_plot_fname(__file__, series='2d')) if args.show: vlt.show() except ImportError: pass try: if not plot3d: raise ImportError from viscid.plot import vlab _ = get_mvi_fig(offscreen=not args.show) mesh = vlab.mesh_from_seeds(sheet_seed, scalars=vx_sheet, clim=(-400, 400)) vlab.plot_earth_3d(crd_system=b) vlab.view(azimuth=+90.0 + 45.0, elevation=90.0 - 25.0, distance=30.0, focalpoint=(-10.0, +1.0, +1.0)) vlab.title("RectilinearMeshPoints") vlab.savefig(next_plot_fname(__file__, series='3d')) if args.show: vlab.show(stop=True) except ImportError: pass # prevent weird xorg bad-instructions on tear down if 'figure' in _global_ns and _global_ns['figure'] is not None: from viscid.plot import vlab vlab.mlab.close(_global_ns['figure']) return 0
def prepare_lines(lines, scalars=None, do_connections=False, other=None): """Concatenate and standardize a list of lines Args: lines (list): Must be a list of 3xN or 4xN ndarrays of xyz(s) data for N points along the line. N need not be the same for all lines. Can alse be 6xN such that lines[:][3:, :] are interpreted as rgb colors scalars (ndarray, list): Can have shape 1xN for a single scalar or 3xN for an rgb color for each point. If the shape is 1xNlines, the scalar is broadcast so the whole line gets the same value, and likewise for 3xNlines and rgb colors. Can also be a list of hex color (#ffffff) strings. Otherwise, scalars is reshaped to -1xN. do_connections (bool): Whether or not to make connections array other (dict): a dictionary of other arrays that should be reshaped and the like the same way scalars is Returns: (vertices, scalars, connections, other) * vertices (ndarray): 3xN array of N xyz points. N is the sum of the lengths of all the lines * scalars (ndarray): N array of scalars, 3xN array of uint8 rgb values, or None * connections (ndarray): Nx2 array of ints (indices along axis 1 of vertices) describing the forward and backward connectedness of the lines, or None * other (dict): a dict of N length arrays Raises: ValueError: If rgb data is not in a valid range or the shape of scalars is not understood """ nlines = len(lines) npts = [line.shape[1] for line in lines] N = np.sum(npts) first_idx = np.cumsum([0] + npts[:-1]) vertices = [np.asarray(line) for line in lines] vertices = np.concatenate(lines, axis=1) if vertices.shape[0] > 3: if scalars is not None: viscid.logger.warn("Overriding line scalars with scalars kwarg") else: scalars = vertices[3:, :] vertices = vertices[:3, :] if scalars is not None: if isinstance(scalars, viscid.field.Field): scalars = viscid.interp_trilin(scalars, vertices) if scalars.size != N: raise ValueError("Scalars was not a scalar field") scalars = np.atleast_2d(scalars) if scalars.shape == (1, 1): scalars = scalars.repeat(N, axis=1) elif scalars.shape == (1, nlines) or scalars.shape == (nlines, 1): # one scalar for each line, so broadcast it scalars = scalars.reshape(nlines, 1) scalars = [scalars[i].repeat(ni) for i, ni in enumerate(npts)] scalars = np.concatenate(scalars, axis=0).reshape(1, N) elif scalars.shape == (N, 1) or scalars.shape == (1, N): # catch these so they're not interpreted as colors if # nlines == 1 and N == 3; ie. 1 line with 3 points scalars = scalars.reshape(1, N) elif scalars.shape == (3, nlines) or scalars.shape == (nlines, 3): # one rgb color for each line, so broadcast it if scalars.shape == (3, nlines): scalars = scalars.T colors = [] for i, ni in enumerate(npts): c = scalars[i].reshape(3, 1).repeat(ni, axis=1) colors.append(c) scalars = np.concatenate(colors, axis=1) else: scalars = scalars.reshape(-1, N) if scalars.dtype.kind == 'S': # translate hex colors (#ff00ff) into rgb values scalars = np.char.lstrip(scalars, '#').astype('S6') scalars = np.char.zfill(scalars, 6) scalars = np.frombuffer(np.char.decode(scalars, 'hex'), dtype='u1') scalars = scalars.reshape(-1, 3).T elif scalars.shape[0] == 1: # normal scalars scalars = scalars.reshape(-1) elif scalars.shape[0] == 3: # The scalars encode rgb data, standardize the result to a # 3xN ndarray of 1 byte unsigned ints (chars) if np.all(scalars >= 0) and np.all(scalars <= 1): scalars = (255 * scalars).round().astype('u1') elif np.all(scalars >= 0) and np.all(scalars < 256): scalars = scalars.round().astype('u1') else: raise ValueError("Rgb data should be in range [0, 1] or " "[0, 255], range given is [{0}, {1}]" "".format(np.min(scalars), np.max(scalars))) else: raise ValueError("Scalars should either be a number, or set of " "rgb values, shape is {0}".format(scalars.shape)) # broadcast / reshape additional arrays given in other if other: for key, arr in other.items(): if arr is None: pass elif arr.shape == (1, nlines) or arr.shape == (nlines, 1): arr = arr.reshape(nlines, 1) arr = [arr[i].repeat(ni) for i, ni in enumerate(npts)] other[key] = np.concatenate(arr, axis=0).reshape(1, N) else: try: other[key] = arr.reshape(-1, N) except ValueError: viscid.logger.warn("Unknown dimension, dropping array {0}" "".format(key)) if do_connections: connections = [None] * nlines for i, ni in enumerate(npts): # i0 is the index of the first point of the i'th line in lines i0 = first_idx[i] connections[i] = np.vstack([np.arange(i0, i0 + ni - 1.5), np.arange(i0 + 1, i0 + ni - 0.5)]).T connections = np.concatenate(connections, axis=0).astype('i') else: connections = None return vertices, scalars, connections, other