""" create simulation grids """ import typing as T import math import numpy as np from .readdata import readgrid pi = math.pi def makegrid_cart3d(p: T.Dict[str, T.Any]) -> T.Dict[str, T.Any]: """ make 3D grid and save to disk Parameters ----------- p: dict simulation parameters Returns ------- xg: dict simulation grid """ # %%create altitude grid # original Matlab params # p.alt_min = 80e3; # p.alt_max = 1000e3; # p.alt_scale = [10e3, 8e3, 500e3, 150e3]; if set(("alt_min", "alt_max", "alt_scale", "Bincl")) <= p.keys(): z = altitude_grid(p["alt_min"], p["alt_max"], p["Bincl"], p["alt_scale"]) elif "eqdir" in p and p["eqdir"].is_file(): print("makegrid_cart_3D: reusing grid from", p["eqdir"]) xeq = readgrid(p["eqdir"]) z = xeq["x1"] del xeq else: raise ValueError("must specify altitude grid parameters or grid file to reuse") # %% TRANSVERSE GRID (BASED ON SIZE OF CURRENT REGION SPECIFIED ABOVE) # EAST x = grid1d(p["xdist"], p["lxp"]) # NORTH y = grid1d(p["ydist"], p["lyp"]) # %% COMPUTE CELL WALL LOCATIONS lx2 = x.size xi = np.empty(lx2 + 1) xi[1:-1] = 1 / 2 * (x[1:] + x[:-1]) xi[0] = x[0] - 1 / 2 * (x[1] - x[0]) xi[-1] = x[-1] + 1 / 2 * (x[-1] - x[-2]) lx3 = y.size yi = np.empty(lx3 + 1) yi[1:-1] = 1 / 2 * (y[1:] + y[:-1]) yi[0] = y[0] - 1 / 2 * (y[1] - y[0]) yi[-1] = y[-1] + 1 / 2 * (y[-1] - y[-2]) lx1 = z.size zi = np.empty(lx1 + 1) zi[1:-1] = 1 / 2 * (z[1:] + z[:-1]) zi[0] = z[0] - 1 / 2 * (z[1] - z[0]) zi[-1] = z[-1] + 1 / 2 * (z[-1] - z[-2]) # %% GRAVITATIONAL FIELD COMPONENTS IN DIPOLE SYSTEM Re = 6370e3 G = 6.67428e-11 Me = 5.9722e24 r = z + Re g = G * Me / r ** 2 gz = np.tile(-g[:, None, None], (1, lx2, lx3)) assert gz.shape == (lx1, lx2, lx3) # DISTANCE EW AND NS (FROM ENU (or UEN in our case - cyclic permuted) COORD. SYSTEM) # #NEED TO BE CONVERTED TO DIPOLE SPHERICAL AND THEN # GLAT/GLONG - BASICALLY HERE WE ARE MAPPING THE CARTESIAN GRID ONTO THE # SURFACE OF A SPHERE THEN CONVERTING TO GEOGRAPHIC. # get the magnetic coordinates of the grid center, based on user input thetactr, phictr = geog2geomag(p["glat"], p["glon"]) # %% Center of earth distance r = Re + z r = np.tile(r[:, None, None], (1, lx2, lx3)) assert r.shape == (lx1, lx2, lx3) # %% Northward angular distance gamma2 = y / Re # must retain the sign of x3 theta = thetactr - gamma2 # minus because distance north is against theta's direction theta = np.tile(theta[None, None, :], (lx1, lx2, 1)) assert theta.shape == (lx1, lx2, lx3) # %% Eastward angular distance # gamma1=x/Re; %must retain the sign of x2 gamma1 = x / Re / np.sin(thetactr) # must retain the sign of x2, just use theta of center of grid phi = phictr + gamma1 phi = np.tile(phi[None, :, None], (lx1, 1, lx3)) assert phi.shape == (lx1, lx2, lx3) # %% COMPUTE THE GEOGRAPHIC COORDINATES OF EACH GRID POINT glatgrid, glongrid = geomag2geog(theta, phi) # %% COMPUTE ECEF CARTESIAN IN CASE THEY ARE NEEDED xECEF = r * np.sin(theta) * np.cos(phi) yECEF = r * np.sin(theta) * np.sin(phi) zECEF = r * np.cos(theta) # %% COMPUTE SPHERICAL ECEF UNIT VECTORS - CARTESIAN-ECEF COMPONENTS er = np.empty((lx1, lx2, lx3, 3)) etheta = np.empty_like(er) ephi = np.empty_like(er) er[:, :, :, 0] = np.sin(theta) * np.cos(phi) # xECEF-component of er er[:, :, :, 1] = np.sin(theta) * np.sin(phi) # yECEF er[:, :, :, 2] = np.cos(theta) # zECEF etheta[:, :, :, 0] = np.cos(theta) * np.cos(phi) etheta[:, :, :, 1] = np.cos(theta) * np.sin(phi) etheta[:, :, :, 2] = -np.sin(theta) ephi[:, :, :, 0] = -np.sin(phi) ephi[:, :, :, 1] = np.cos(phi) ephi[:, :, :, 2] = 0 # %% UEN UNIT VECTORS IN ECEF COMPONENTS e1 = er # up is the same direction as from ctr of earth e2 = ephi # e2 is same as ephi e3 = -etheta # etheta is positive south, e3 is pos. north # %% STORE RESULTS IN GRID DATA STRUCTURE xg = { "x1": z, "x2": x, "x3": y, "x1i": zi, "x2i": xi, "x3i": yi, } lx = np.array((xg["x1"].size, xg["x2"].size, xg["x3"].size)) xg["lx"] = lx xg["dx1f"] = np.append(xg["x1"][1:] - xg["x1"][:-1], xg["x1"][-1] - xg["x1"][-2]) # FWD DIFF xg["dx1b"] = np.insert(xg["x1"][1:] - xg["x1"][:-1], 0, xg["x1"][1] - xg["x1"][0]) # BACK DIFF xg["dx1h"] = xg["x1i"][1:-1] - xg["x1i"][:-2] # MIDPOINT DIFFS xg["dx2f"] = np.append(xg["x2"][1:] - xg["x2"][:-1], xg["x2"][-1] - xg["x2"][-2]) # FWD DIFF xg["dx2b"] = np.insert(xg["x2"][1:] - xg["x2"][:-1], 0, xg["x2"][1] - xg["x2"][0]) # BACK DIFF xg["dx2h"] = xg["x2i"][1:-1] - xg["x2i"][:-2] # MIDPOINT DIFFS xg["dx3f"] = np.append(xg["x3"][1:] - xg["x3"][:-1], xg["x3"][-1] - xg["x3"][-2]) # FWD DIFF xg["dx3b"] = np.insert(xg["x3"][1:] - xg["x3"][:-1], 0, xg["x3"][1] - xg["x3"][0]) # BACK DIFF xg["dx3h"] = xg["x3i"][1:-1] - xg["x3i"][:-2] # MIDPOINT DIFFS xg["h1"] = np.ones(xg["lx"]) xg["h2"] = np.ones(xg["lx"]) xg["h3"] = np.ones(xg["lx"]) xg["h1x1i"] = np.ones((lx[0] + 1, lx[1], lx[2])) xg["h2x1i"] = np.ones((lx[0] + 1, lx[1], lx[2])) xg["h3x1i"] = np.ones((lx[0] + 1, lx[1], lx[2])) xg["h1x2i"] = np.ones((lx[0], lx[1] + 1, lx[2])) xg["h2x2i"] = np.ones((lx[0], lx[1] + 1, lx[2])) xg["h3x2i"] = np.ones((lx[0], lx[1] + 1, lx[2])) xg["h1x3i"] = np.ones((lx[0], lx[1], lx[2] + 1)) xg["h2x3i"] = np.ones((lx[0], lx[1], lx[2] + 1)) xg["h3x3i"] = np.ones((lx[0], lx[1], lx[2] + 1)) # %% Cartesian, ECEF representation of curvilinar coordinates xg["e1"] = e1 xg["e2"] = e2 xg["e3"] = e3 # %% ECEF spherical coordinates xg["r"] = r xg["theta"] = theta xg["phi"] = phi # xg.rx1i=[]; xg.thetax1i=[]; # xg.rx2i=[]; xg.thetax2i=[]; # %% These are cartesian representations of the ECEF, spherical unit vectors xg["er"] = er xg["etheta"] = etheta xg["ephi"] = ephi xg["I"] = np.broadcast_to(p["Bincl"], (lx2, lx3)) # %% Cartesian ECEF coordinates xg["x"] = xECEF xg["z"] = zECEF xg["y"] = yECEF xg["alt"] = xg["r"] - Re # since we need a 3D array use xg.r here... xg["gx1"] = gz xg["gx2"] = np.zeros(xg["lx"]) xg["gx3"] = np.zeros(xg["lx"]) xg["Bmag"] = np.broadcast_to(-50000e-9, xg["lx"]) # minus for northern hemisphere... xg["glat"] = glatgrid xg["glon"] = glongrid # xg['xp']=x; xg['zp']=z; # xg['inull']=[]; xg["nullpts"] = np.zeros(xg["lx"]) # %% TRIM DATA STRUCTRE TO BE THE SIZE FORTRAN EXPECTS # note: xgf is xg == True xgf = xg # indices corresponding to non-ghost cells for 1 dimension i1 = slice(2, lx[0] - 2) i2 = slice(2, lx[1] - 2) i3 = slice(2, lx[2] - 2) # any dx variable will not need to first element (backward diff of two ghost cells) idx1 = slice(1, lx[0]) idx2 = slice(1, lx[1]) idx3 = slice(1, lx[2]) # x1-interface variables need only non-ghost cell values (left interface) plus one ix1i = slice(2, lx[0] - 1) ix2i = slice(2, lx[1] - 1) ix3i = slice(2, lx[2] - 1) # remove ghost cells # now that indices have been define we can go ahead and make this change xgf["lx"] = xgf["lx"] - 4 xgf["dx1b"] = xgf["dx1b"][idx1] xgf["dx2b"] = xgf["dx2b"][idx2] xgf["dx3b"] = xgf["dx3b"][idx3] xgf["x1i"] = xgf["x1i"][ix1i] xgf["x2i"] = xgf["x2i"][ix2i] xgf["x3i"] = xgf["x3i"][ix3i] xgf["dx1h"] = xgf["dx1h"][i1] xgf["dx2h"] = xgf["dx2h"][i2] xgf["dx3h"] = xgf["dx3h"][i3] xgf["h1x1i"] = xgf["h1x1i"][ix1i, i2, i3] xgf["h2x1i"] = xgf["h2x1i"][ix1i, i2, i3] xgf["h3x1i"] = xgf["h3x1i"][ix1i, i2, i3] xgf["h1x2i"] = xgf["h1x2i"][i1, ix2i, i3] xgf["h2x2i"] = xgf["h2x2i"][i1, ix2i, i3] xgf["h3x2i"] = xgf["h3x2i"][i1, ix2i, i3] xgf["h1x3i"] = xgf["h1x3i"][i1, i2, ix3i] xgf["h2x3i"] = xgf["h2x3i"][i1, i2, ix3i] xgf["h3x3i"] = xgf["h3x3i"][i1, i2, ix3i] xgf["gx1"] = xgf["gx1"][i1, i2, i3] xgf["gx2"] = xgf["gx2"][i1, i2, i3] xgf["gx3"] = xgf["gx3"][i1, i2, i3] xgf["glat"] = xgf["glat"][i1, i2, i3] xgf["glon"] = xgf["glon"][i1, i2, i3] xgf["alt"] = xgf["alt"][i1, i2, i3] xgf["Bmag"] = xgf["Bmag"][i1, i2, i3] xgf["I"] = xgf["I"][i2, i3] xgf["nullpts"] = xgf["nullpts"][i1, i2, i3] xgf["e1"] = xgf["e1"][i1, i2, i3, :] xgf["e2"] = xgf["e2"][i1, i2, i3, :] xgf["e3"] = xgf["e3"][i1, i2, i3, :] xgf["er"] = xgf["er"][i1, i2, i3, :] xgf["etheta"] = xgf["etheta"][i1, i2, i3, :] xgf["ephi"] = xgf["ephi"][i1, i2, i3, :] xgf["r"] = xgf["r"][i1, i2, i3] xgf["theta"] = xgf["theta"][i1, i2, i3] xgf["phi"] = xgf["phi"][i1, i2, i3] xgf["x"] = xgf["x"][i1, i2, i3] xgf["y"] = xgf["y"][i1, i2, i3] xgf["z"] = xgf["z"][i1, i2, i3] return xgf def geomag2geog(thetat: np.ndarray, phit: np.ndarray) -> T.Tuple[np.ndarray, np.ndarray]: """ geomagnetic to geographic """ # TODO: is OK to be hardcoded? thetan = math.radians(11) phin = math.radians(289) # enforce phit = [0,2pi] i = phit > 2 * pi phitcorrected = phit phitcorrected[i] = phit[i] - 2 * pi i = phit < 0 phitcorrected[i] = phit[i] + 2 * pi # thetag2p=acos(cos(thetat).*cos(thetan)-sin(thetat).*sin(thetan).*cos(phit)); thetag2p = np.arccos(np.cos(thetat) * np.cos(thetan) - np.sin(thetat) * np.sin(thetan) * np.cos(phitcorrected)) beta = np.arccos((np.cos(thetat) - np.cos(thetag2p) * np.cos(thetan)) / (np.sin(thetag2p) * np.sin(thetan))) phig2 = np.zeros_like(phitcorrected) i = phitcorrected > pi phig2[i] = phin - beta[i] i = phitcorrected <= pi phig2[i] = phin + beta[i] i = phig2 < 0 phig2[i] = phig2[i] + 2 * pi i = phig2 >= 2 * pi phig2[i] = phig2[i] - 2 * pi thetag2 = pi / 2 - thetag2p lat = np.degrees(thetag2) lon = np.degrees(phig2) return lat, lon def geog2geomag(lat: np.ndarray, lon: np.ndarray) -> T.Tuple[np.ndarray, np.ndarray]: """ geographic to geomagnetic """ # TODO: is this arbitrary or hardcoded OK? thetan = math.radians(11) phin = math.radians(289) # enforce [0,360] longitude lon = lon % 360 thetagp = pi / 2 - np.radians(lat) phig = np.radians(lon) thetat = np.arccos(np.cos(thetagp) * np.cos(thetan) + np.sin(thetagp) * np.sin(thetan) * np.cos(phig - phin)) argtmp = (np.cos(thetagp) - np.cos(thetat) * np.cos(thetan)) / (np.sin(thetat) * np.sin(thetan)) alpha = np.arccos(max(min(argtmp, 1), -1)) phit = np.empty_like(lat) i = ((phin > phig) & ((phin - phig) > pi)) | ((phin < phig) & ((phig - phin) < pi)) phit[i] = pi - alpha[i] phit[~i] = alpha[~i] + pi return thetat, phit def grid1d(dist: float, L: int) -> np.ndarray: """ generate 1D grid Parameters ---------- dist: float one-way distance from origin (meters) L: int number of cells Returns ------- x: np.ndarray 1D vector grid """ xmin = -dist / 2 xmax = dist / 2 if L == 1: # degenerate dimension # add 2 ghost cells on each side x = np.linspace(xmin, xmax, L + 4) else: # exclude the ghost cells when setting extents x = np.linspace(xmin, xmax, L) d0 = x[1] - x[0] d1 = x[-1] - x[-2] # now tack on ghost cells so they are outside user-specified region x = np.insert(x, 0, [x[0] - 2 * d0, x[0] - d0]) x = np.append(x, [x[-1] + d1, x[-1] + 2 * d1]) return x def altitude_grid(alt_min: float, alt_max: float, incl_deg: float, d: T.Tuple[float, float, float, float]) -> np.ndarray: if alt_min < 0 or alt_max < 0: raise ValueError("grid values must be positive") if alt_max <= alt_min: raise ValueError("grid_max must be greater than grid_min") alt = [alt_min] while alt[-1] < alt_max: # dalt=10+9.5*tanh((alt(i-1)-500)/150) dalt = d[0] + d[1] * math.tanh((alt[-1] - d[2]) / d[3]) alt.append(alt[-1] + dalt) alt = np.asarray(alt) if alt.size < 10: raise ValueError("grid too small") # %% tilt for magnetic inclination z = alt / math.sin(math.radians(incl_deg)) # %% add two ghost cells each to top and bottom dz1 = z[1] - z[0] dzn = z[-1] - z[-2] z = np.insert(z, 0, [z[0] - 2 * dz1, z[0] - dz1]) z = np.append(z, [z[-1] + dzn, z[-1] + 2 * dzn]) return z