def plot_earth_3d(figure=None, daycol=(1, 1, 1), nightcol=(0, 0, 0), radius=1.0, res=24, crd_system="gse", night_only=False, **kwargs): """Plot a black and white sphere (Earth) showing sunward direction Parameters: figure (mayavi.core.scene.Scene): specific figure, or None for :py:func:`mayavi.mlab.gcf` daycol (tuple, optional): color of dayside (RGB) nightcol (tuple, optional): color of nightside (RGB) res (optional): rosolution of teh sphere crd_system (str, other): One of ('mhd', 'gse'), or anything that returns from :py:func:`viscid.get_crd_system`. Returns: Tuple (day, night) as vtk sources """ if figure is None: figure = mlab.gcf() crd_system = viscid.get_crd_system(crd_system) if crd_system == "mhd": theta_dusk, theta_dawn = 270, 90 elif crd_system == "gse": theta_dusk, theta_dawn = 90, 270 else: # use GSE convention? theta_dusk, theta_dawn = 90, 270 night = BuiltinSurface(source='sphere', name='night') night.data_source.set(center=(0, 0, 0), radius=radius, start_theta=theta_dusk, end_theta=theta_dawn, theta_resolution=res, phi_resolution=res) mod = mlab.pipeline.surface(night, color=nightcol, figure=figure, **kwargs) mod.actor.property.backface_culling = True if not night_only: day = BuiltinSurface(source='sphere', name='day') day.data_source.set(center=(0, 0, 0), radius=radius, start_theta=theta_dawn, end_theta=theta_dusk, theta_resolution=res, phi_resolution=res) mod = mlab.pipeline.surface(day, color=daycol, figure=figure, **kwargs) mod.actor.property.backface_culling = True else: day = None return day, night
def make_dipole(m=(0, 0, -DEFAULT_STRENGTH), strength=None, l=None, h=None, n=None, twod=False, dtype='f8', nonuniform=False, crd_system='gse', name='b'): """Generate a dipole field with magnetic moment m [x, y, z]""" if l is None: l = [-5] * 3 if h is None: h = [5] * 3 if n is None: n = [256] * 3 x = np.array(np.linspace(l[0], h[0], n[0]), dtype=dtype) y = np.array(np.linspace(l[1], h[1], n[1]), dtype=dtype) z = np.array(np.linspace(l[2], h[2], n[2]), dtype=dtype) if twod: y = np.array(np.linspace(-0.1, 0.1, 2), dtype=dtype) if nonuniform: z += 0.01 * ((h[2] - l[2]) / n[2]) * np.sin(np.linspace(0, np.pi, n[2])) B = field.empty([x, y, z], nr_comps=3, name=name, center='cell', layout='interlaced', dtype=dtype) B.set_info('crd_system', viscid.get_crd_system(crd_system)) return fill_dipole(B, m=m, strength=strength)
def save_fields(cls, fname, flds, **kwargs): """ save some fields using the format given by the class """ # FIXME: this is only good for writing cartesian rectilnear flds # FIXME: axes are renamed if flds[0] is 1D or 2D assert len(flds) > 0 fname = os.path.expanduser(os.path.expandvars(fname)) # FIXME: all coordinates are saved as non-uniform, the proper # way to do this is to have let coordinate format its own # hdf5 / xdmf / numpy binary output clist = flds[0].crds.get_clist(full_arrays=True) crd_arrs = [np.array([0.0])] * 3 crd_names = ["x", "y", "z"] for i, c in enumerate(clist): crd_arrs[i] = c[1] crd_shape = [len(arr) for arr in crd_arrs] time = flds[0].time # write arrays to the hdf5 file with h5py.File(fname, 'w') as f: for axis_name, arr in zip(crd_names, crd_arrs): loc = cls._CRDS_GROUP + '/' + axis_name f[loc] = arr for fld in flds: loc = cls._FLD_GROUPS[fld.center.lower()] + '/' + fld.name # xdmf files use kji ordering f[loc] = fld.data.T # big bad openggcm time_str hack to put basetime into hdf5 file for fld in flds: try: tfmt = "%Y:%m:%d:%H:%M:%S.%f" sec_td = viscid.as_timedelta64(fld.time, 's') dtime = viscid.as_datetime(fld.basetime + sec_td).strftime(tfmt) epoch = viscid.readers.openggcm.GGCM_EPOCH ts = viscid.as_timedelta(fld.basetime - epoch).total_seconds() ts += fld.time timestr = "time= {0} {1:.16e} {2} 300c".format(fld.time, ts, dtime) f.create_group('openggcm') f['openggcm'].attrs['time_str'] = np.string_(timestr) break except viscid.NoBasetimeError: pass # now write an xdmf file xdmf_fname = os.path.splitext(fname)[0] + ".xdmf" relh5fname = "./" + os.path.basename(fname) with open(xdmf_fname, 'w') as f: xloc = cls._CRDS_GROUP + '/' + crd_names[0] yloc = cls._CRDS_GROUP + '/' + crd_names[1] zloc = cls._CRDS_GROUP + '/' + crd_names[2] dim_str = " ".join([str(l) for l in crd_shape][::-1]) f.write(cls._XDMF_TEMPLATE_BEGIN.format(time=time)) s = cls._XDMF_TEMPLATE_RECTILINEAR_GRID_BEGIN.format( grid_name="vgrid", crd_dims=dim_str, h5fname=relh5fname, xdim=crd_shape[0], ydim=crd_shape[1], zdim=crd_shape[2], xloc=xloc, yloc=yloc, zloc=zloc) f.write(s) for fld in flds: _crd_system = viscid.get_crd_system(fld, None) if _crd_system: f.write(cls._XDMF_INFO_TEMPLATE.format(name="crd_system", value=_crd_system)) break for fld in flds: fld = fld.as_flat().T dt = fld.dtype.name.rstrip("0123456789").title() precision = fld.dtype.itemsize fld_dim_str = " ".join([str(l) for l in fld.shape]) loc = cls._FLD_GROUPS[fld.center.lower()] + '/' + fld.name s = cls._XDMF_TEMPLATE_ATTRIBUTE.format( fld_name=fld.name, fld_type=fld.fldtype, center=fld.center.title(), dtype=dt, precision=precision, fld_dims=fld_dim_str, h5fname=relh5fname, fld_loc=loc) f.write(s) f.write(cls._XDMF_TEMPLATE_GRID_END) f.write(cls._XDMF_TEMPLATE_END)
def __init__(self, p0=(0, 0, 0), r=0.0, pole=(0, 0, 1), ntheta=20, nphi=20, thetalim=(0, 180.0), philim=(0, 360.0), roll=0.0, crd_system=None, theta_endpoint='auto', phi_endpoint='auto', pole_is_vector=True, theta_phi=False, cache=False, dtype=None): """Make seeds on the surface of a sphere Note: There is some funny business about the meaning of phi=0 and `crd_system`. By default, this seed generator is agnostic to coordinate systems and phi=0 always means the +x axis. If crd_system is 'gse', 'mhd', or an object whose find_info method returns a 'crd_system', then phi=0 means midnight. This is important when specifying a phi range or plotting on a matplotlib polar plot. Args: p0 (list, tuple, or ndarray): Origin of sphere as (x, y, z) r (float): Radius of sphere; or calculated from pole if 0 pole (list, tuple, or ndarray): Vector pointing in the direction of the north pole of the sphere. Defaults to (0, 0, 1). ntheta (int): Number of points in theta nphi (int): Number of points in phi thetalim (list): min and max theta (in degrees) philim (list): min and max phi (in degrees) roll (float): Roll the seeds around the pole by this angle in deg crd_system (str, Field): a crd system ('gse', 'mhd') or an object that has 'crd_system' info such that phi=0 means midnight instead of the +x axis. theta_endpoint (str): this is a bit of a hack to keep from having redundant seeds at poles. You probably just want auto here phi_endpoint (bool): if true, then let phi inclue upper value. This is false by default since 0 == 2pi. pole_is_vector (bool): Whether pole is a vector or a vector theta_phi (bool): If True, the uv and local representations are ordered (theta, phi), otherwise (phi, theta) Raises: ValueError: if thetalim or philim don't have 2 values each """ super(Sphere, self).__init__(cache=cache, dtype=dtype) self.p0 = np.asarray(p0, dtype=self.dtype) if pole_is_vector: self.pole = np.asarray(pole, dtype=self.dtype) else: if pole is None: pole = p0 + np.asarray([0, 0, 1], dtype=self.dtype) else: pole = np.asarray(pole, dtype=self.dtype) self.pole = pole - p0 if not r: r = np.linalg.norm(self.pole) self.pole = self.pole / np.linalg.norm(self.pole) if not len(thetalim) == len(philim) == 2: raise ValueError("thetalim and philim must have both min and max") try: roll = float(roll) except TypeError: pass # square away crd system if crd_system: if hasattr(crd_system, 'find_info'): crd_system = viscid.get_crd_system(crd_system, 'none') else: crd_system = 'none' if crd_system.strip().lower() == 'gse': crd_system_roll = -180.0 else: crd_system_roll = 0.0 self.r = r self.ntheta = ntheta self.nphi = nphi self.thetalim = np.deg2rad(thetalim) self.philim = np.deg2rad(philim) self.theta_endpoint = theta_endpoint self.phi_endpoint = phi_endpoint self.theta_phi = theta_phi self.roll = roll self.crd_system_roll = crd_system_roll
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.get_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