def lu(a): """ Performe LU decomposition on the matrix a so A = L*U """ dw = True if dw: def loop_body(l, u): c = get_iterator() l[c:, c - 1] = (u[c:, c - 1] / u[c - 1, c - 1:c]) u[c:, c - 1:] = u[c:, c - 1:] - l[c:, c - 1][:, None] * u[c - 1, c - 1:] u = a.copy() l = np.zeros_like(a) np.diagonal(l)[:] = 1.0 loop.do_while(loop_body, u.shape[0] - 2, l, u) return (l, u) else: u = a.copy() l = np.zeros_like(a) np.diagonal(l)[:] = 1.0 for c in range(1, u.shape[0]): l[c:, c - 1] = (u[c:, c - 1] / u[c - 1, c - 1:c]) u[c:, c - 1:] = u[c:, c - 1:] - l[c:, c - 1][:, None] * u[c - 1, c - 1:] np.flush() return (l, u)
def solve_tridiag(a, b, c, d, B): assert a.shape == b.shape and a.shape == c.shape and a.shape == d.shape #if not climate.is_bohrium: # return lapack.dgtsv(a.flatten()[1:],b.flatten(),c.flatten()[:-1],d.flatten())[3].reshape(a.shape) n = a.shape[-1] x, cp, dp = np.zeros_like(a), np.zeros_like(a), np.zeros_like(a) # initialize c-prime and d-prime cp[..., 0] = c[..., 0] / b[..., 0] dp[..., 0] = d[..., 0] / b[..., 0] # solve for vectors c-prime and d-prime def loop_body(cp, dp): i = get_iterator(1) m = b[..., i] - cp[..., i - 1] * a[..., i] fxa = 1.0 / m cp[..., i] = c[..., i] * fxa dp[..., i] = (d[..., i] - dp[..., i - 1] * a[..., i]) * fxa B.do_while(loop_body, n - 1, cp, dp) x[..., n - 1] = dp[..., n - 1] def loop_body(x): i = get_iterator() x[..., n - 2 - i] = dp[..., n - 2 - i] - cp[..., n - 2 - i] * x[..., n - 1 - i] B.do_while(loop_body, n - 1, x) return x
def adv_flux_superbee_wgrid(adv_fe, adv_fn, adv_ft, var, u_wgrid, v_wgrid, w_wgrid, maskW, dxt, dyt, dzw, cost, cosu, dt_tracer): """ Calculates advection of a tracer defined on Wgrid """ maskUtr = bh.zeros_like(maskW) maskUtr[:-1, :, :] = maskW[1:, :, :] * maskW[:-1, :, :] adv_fe[...] = 0. adv_fe[1:-2, 2:-2, :] = _adv_superbee(u_wgrid, var, maskUtr, dxt, 0, cost, cosu, dt_tracer) maskVtr = bh.zeros_like(maskW) maskVtr[:, :-1, :] = maskW[:, 1:, :] * maskW[:, :-1, :] adv_fn[...] = 0. adv_fn[2:-2, 1:-2, :] = _adv_superbee(v_wgrid, var, maskVtr, dyt, 1, cost, cosu, dt_tracer) maskWtr = bh.zeros_like(maskW) maskWtr[:, :, :-1] = maskW[:, :, 1:] * maskW[:, :, :-1] adv_ft[...] = 0. adv_ft[2:-2, 2:-2, :-1] = _adv_superbee(w_wgrid, var, maskWtr, dzw, 2, cost, cosu, dt_tracer)
def lu(a): """ Performe LU decomposition on the matrix a so A = L*U """ u = a.copy() l = np.zeros_like(a) np.diagonal(l)[:] = 1.0 for c in xrange(1,u.shape[0]): l[c:,c-1] = (u[c:,c-1]/u[c-1,c-1:c]) u[c:,c-1:] = u[c:,c-1:] - l[c:,c-1][:,None] * u[c-1,c-1:] np.flush(u) return (l,u)
def weighted_line(r0, c0, r1, c1, w, rmin=0, rmax=np.inf): # The algorithm below works fine if c1 >= c0 and c1-c0 >= abs(r1-r0). # If either of these cases are violated, do some switches. if abs(c1-c0) < abs(r1-r0): # Switch x and y, and switch again when returning. xx, yy, val = weighted_line(c0, r0, c1, r1, w, rmin=rmin, rmax=rmax) return (yy, xx, val) # At this point we know that the distance in columns (x) is greater # than that in rows (y). Possibly one more switch if c0 > c1. if c0 > c1: return weighted_line(r1, c1, r0, c0, w, rmin=rmin, rmax=rmax) # The following is now always < 1 in abs num=r1-r0 denom=c1-c0 slope = np.divide(num,denom,out=np.zeros_like(denom), where=denom!=0) # Adjust weight by the slope w *= np.sqrt(1+np.abs(slope)) / 2 # We write y as a function of x, because the slope is always <= 1 # (in absolute value) x = np.arange(c0, c1+1, dtype=float) y = x * slope + (c1*r0-c0*r1) / (c1-c0) # Now instead of 2 values for y, we have 2*np.ceil(w/2). # All values are 1 except the upmost and bottommost. thickness = np.ceil(w/2) yy = (np.floor(y).reshape(-1,1) + np.arange(-thickness-1,thickness+2).reshape(1,-1)) xx = np.repeat(x, yy.shape[1]) vals = trapez(yy, y.reshape(-1,1), w).flatten() yy = yy.flatten() # Exclude useless parts and those outside of the interval # to avoid parts outside of the picture mask = np.logical_and.reduce((yy >= rmin, yy < rmax, vals > 0)) return (yy[mask].astype(int), xx[mask].astype(int), vals[mask])
def integrate_tke(u, v, w, maskU, maskV, maskW, dxt, dxu, dyt, dyu, dzt, dzw, cost, cosu, kbot, kappaM, mxl, forc, forc_tke_surface, tke, dtke): tau = 0 taup1 = 1 taum1 = 2 dt_tracer = 1 dt_mom = 1 AB_eps = 0.1 alpha_tke = 1. c_eps = 0.7 K_h_tke = 2000. flux_east = bh.zeros_like(maskU) flux_north = bh.zeros_like(maskU) flux_top = bh.zeros_like(maskU) sqrttke = bh.sqrt(bh.maximum(0., tke[:, :, :, tau])) """ integrate Tke equation on W grid with surface flux boundary condition """ dt_tke = dt_mom # use momentum time step to prevent spurious oscillations """ vertical mixing and dissipation of TKE """ ks = kbot[2:-2, 2:-2] - 1 a_tri = bh.zeros_like(maskU[2:-2, 2:-2]) b_tri = bh.zeros_like(maskU[2:-2, 2:-2]) c_tri = bh.zeros_like(maskU[2:-2, 2:-2]) d_tri = bh.zeros_like(maskU[2:-2, 2:-2]) delta = bh.zeros_like(maskU[2:-2, 2:-2]) delta[:, :, :-1] = dt_tke / dzt[bh.newaxis, bh.newaxis, 1:] * alpha_tke * 0.5 \ * (kappaM[2:-2, 2:-2, :-1] + kappaM[2:-2, 2:-2, 1:]) a_tri[:, :, 1:-1] = -delta[:, :, :-2] / \ dzw[bh.newaxis, bh.newaxis, 1:-1] a_tri[:, :, -1] = -delta[:, :, -2] / (0.5 * dzw[-1]) b_tri[:, :, 1:-1] = 1 + (delta[:, :, 1:-1] + delta[:, :, :-2]) / dzw[bh.newaxis, bh.newaxis, 1:-1] \ + dt_tke * c_eps \ * sqrttke[2:-2, 2:-2, 1:-1] / mxl[2:-2, 2:-2, 1:-1] b_tri[:, :, -1] = 1 + delta[:, :, -2] / (0.5 * dzw[-1]) \ + dt_tke * c_eps / mxl[2:-2, 2:- 2, -1] * sqrttke[2:-2, 2:-2, -1] b_tri_edge = 1 + delta / dzw[bh.newaxis, bh.newaxis, :] \ + dt_tke * c_eps / mxl[2:-2, 2:-2, :] * sqrttke[2:-2, 2:-2, :] c_tri[:, :, :-1] = -delta[:, :, :-1] / dzw[bh.newaxis, bh.newaxis, :-1] d_tri[...] = tke[2:-2, 2:-2, :, tau] + dt_tke * forc[2:-2, 2:-2, :] d_tri[:, :, -1] += dt_tke * forc_tke_surface[2:-2, 2:-2] / (0.5 * dzw[-1]) sol, water_mask = solve_implicit(ks, a_tri, b_tri, c_tri, d_tri, b_edge=b_tri_edge) tke[2:-2, 2:-2, :, taup1] = where(water_mask, sol, tke[2:-2, 2:-2, :, taup1]) """ Add TKE if surface density flux drains TKE in uppermost box """ tke_surf_corr = bh.zeros(maskU.shape[:2]) mask = tke[2:-2, 2:-2, -1, taup1] < 0.0 tke_surf_corr[2:-2, 2:-2] = where( mask, -tke[2:-2, 2:-2, -1, taup1] * 0.5 * dzw[-1] / dt_tke, 0. ) tke[2:-2, 2:-2, -1, taup1] = bh.maximum(0., tke[2:-2, 2:-2, -1, taup1]) """ add tendency due to lateral diffusion """ flux_east[:-1, :, :] = K_h_tke * (tke[1:, :, :, tau] - tke[:-1, :, :, tau]) \ / (cost[bh.newaxis, :, bh.newaxis] * dxu[:-1, bh.newaxis, bh.newaxis]) * maskU[:-1, :, :] flux_east[-1, :, :] = 0. flux_north[:, :-1, :] = K_h_tke * (tke[:, 1:, :, tau] - tke[:, :-1, :, tau]) \ / dyu[bh.newaxis, :-1, bh.newaxis] * maskV[:, :-1, :] * cosu[bh.newaxis, :-1, bh.newaxis] flux_north[:, -1, :] = 0. tke[2:-2, 2:-2, :, taup1] += dt_tke * maskW[2:-2, 2:-2, :] * \ ((flux_east[2:-2, 2:-2, :] - flux_east[1:-3, 2:-2, :]) / (cost[bh.newaxis, 2:-2, bh.newaxis] * dxt[2:-2, bh.newaxis, bh.newaxis]) + (flux_north[2:-2, 2:-2, :] - flux_north[2:-2, 1:-3, :]) / (cost[bh.newaxis, 2:-2, bh.newaxis] * dyt[bh.newaxis, 2:-2, bh.newaxis])) """ add tendency due to advection """ adv_flux_superbee_wgrid( flux_east, flux_north, flux_top, tke[:, :, :, tau], u[..., tau], v[..., tau], w[..., tau], maskW, dxt, dyt, dzw, cost, cosu, dt_tracer ) dtke[2:-2, 2:-2, :, tau] = maskW[2:-2, 2:-2, :] * (-(flux_east[2:-2, 2:-2, :] - flux_east[1:-3, 2:-2, :]) / (cost[bh.newaxis, 2:-2, bh.newaxis] * dxt[2:-2, bh.newaxis, bh.newaxis]) - (flux_north[2:-2, 2:-2, :] - flux_north[2:-2, 1:-3, :]) / (cost[bh.newaxis, 2:-2, bh.newaxis] * dyt[bh.newaxis, 2:-2, bh.newaxis])) dtke[:, :, 0, tau] += -flux_top[:, :, 0] / dzw[0] dtke[:, :, 1:-1, tau] += - \ (flux_top[:, :, 1:-1] - flux_top[:, :, :-2]) / dzw[1:-1] dtke[:, :, -1, tau] += - \ (flux_top[:, :, -1] - flux_top[:, :, -2]) / \ (0.5 * dzw[-1]) """ Adam Bashforth time stepping """ tke[:, :, :, taup1] += dt_tracer * ((1.5 + AB_eps) * dtke[:, :, :, tau] - (0.5 + AB_eps) * dtke[:, :, :, taum1]) return tke, dtke, tke_surf_corr
def connected_components(a, connectivity=8): b = np.array([connectivity], dtype=a.dtype) c = np.zeros_like(a) ufuncs.extmethod("opencv_connected_components", c, a, b) return c
def isoneutral_diffusion_pre(maskT, maskU, maskV, maskW, dxt, dxu, dyt, dyu, dzt, dzw, cost, cosu, salt, temp, zt, K_iso, K_11, K_22, K_33, Ai_ez, Ai_nz, Ai_bx, Ai_by): """ Isopycnal diffusion for tracer following functional formulation by Griffies et al Code adopted from MOM2.1 """ epsln = 1e-20 iso_slopec = 1e-3 iso_dslope = 1e-3 K_iso_steep = 50. tau = 0 dTdx = bh.zeros_like(K_11) dSdx = bh.zeros_like(K_11) dTdy = bh.zeros_like(K_11) dSdy = bh.zeros_like(K_11) dTdz = bh.zeros_like(K_11) dSdz = bh.zeros_like(K_11) """ drho_dt and drho_ds at centers of T cells """ drdT = maskT * get_drhodT(salt[:, :, :, tau], temp[:, :, :, tau], bh.abs(zt)) drdS = maskT * get_drhodS(salt[:, :, :, tau], temp[:, :, :, tau], bh.abs(zt)) """ gradients at top face of T cells """ dTdz[:, :, :-1] = maskW[:, :, :-1] * \ (temp[:, :, 1:, tau] - temp[:, :, :-1, tau]) / \ dzw[bh.newaxis, bh.newaxis, :-1] dSdz[:, :, :-1] = maskW[:, :, :-1] * \ (salt[:, :, 1:, tau] - salt[:, :, :-1, tau]) / \ dzw[bh.newaxis, bh.newaxis, :-1] """ gradients at eastern face of T cells """ dTdx[:-1, :, :] = maskU[:-1, :, :] * (temp[1:, :, :, tau] - temp[:-1, :, :, tau]) \ / (dxu[:-1, bh.newaxis, bh.newaxis] * cost[bh.newaxis, :, bh.newaxis]) dSdx[:-1, :, :] = maskU[:-1, :, :] * (salt[1:, :, :, tau] - salt[:-1, :, :, tau]) \ / (dxu[:-1, bh.newaxis, bh.newaxis] * cost[bh.newaxis, :, bh.newaxis]) """ gradients at northern face of T cells """ dTdy[:, :-1, :] = maskV[:, :-1, :] * \ (temp[:, 1:, :, tau] - temp[:, :-1, :, tau]) \ / dyu[bh.newaxis, :-1, bh.newaxis] dSdy[:, :-1, :] = maskV[:, :-1, :] * \ (salt[:, 1:, :, tau] - salt[:, :-1, :, tau]) \ / dyu[bh.newaxis, :-1, bh.newaxis] def dm_taper(sx): """ tapering function for isopycnal slopes """ return 0.5 * (1. + bh.tanh((-bh.abs(sx) + iso_slopec) / iso_dslope)) """ Compute Ai_ez and K11 on center of east face of T cell. """ diffloc = bh.zeros_like(K_11) diffloc[1:-2, 2:-2, 1:] = 0.25 * (K_iso[1:-2, 2:-2, 1:] + K_iso[1:-2, 2:-2, :-1] + K_iso[2:-1, 2:-2, 1:] + K_iso[2:-1, 2:-2, :-1]) diffloc[1:-2, 2:-2, 0] = 0.5 * \ (K_iso[1:-2, 2:-2, 0] + K_iso[2:-1, 2:-2, 0]) sumz = bh.zeros_like(K_11)[1:-2, 2:-2] for kr in range(2): ki = 0 if kr == 1 else 1 for ip in range(2): drodxe = drdT[1 + ip:-2 + ip, 2:-2, ki:] * dTdx[1:-2, 2:-2, ki:] \ + drdS[1 + ip:-2 + ip, 2:-2, ki:] * dSdx[1:-2, 2:-2, ki:] drodze = drdT[1 + ip:-2 + ip, 2:-2, ki:] * dTdz[1 + ip:-2 + ip, 2:-2, :-1 + kr or None] \ + drdS[1 + ip:-2 + ip, 2:-2, ki:] * \ dSdz[1 + ip:-2 + ip, 2:-2, :-1 + kr or None] sxe = -drodxe / (bh.minimum(0., drodze) - epsln) taper = dm_taper(sxe) sumz[:, :, ki:] += dzw[bh.newaxis, bh.newaxis, :-1 + kr or None] * maskU[1:-2, 2:-2, ki:] \ * bh.maximum(K_iso_steep, diffloc[1:-2, 2:-2, ki:] * taper) Ai_ez[1:-2, 2:-2, ki:, ip, kr] = taper * \ sxe * maskU[1:-2, 2:-2, ki:] K_11[1:-2, 2:-2, :] = sumz / (4. * dzt[bh.newaxis, bh.newaxis, :]) """ Compute Ai_nz and K_22 on center of north face of T cell. """ diffloc[...] = 0 diffloc[2:-2, 1:-2, 1:] = 0.25 * (K_iso[2:-2, 1:-2, 1:] + K_iso[2:-2, 1:-2, :-1] + K_iso[2:-2, 2:-1, 1:] + K_iso[2:-2, 2:-1, :-1]) diffloc[2:-2, 1:-2, 0] = 0.5 * \ (K_iso[2:-2, 1:-2, 0] + K_iso[2:-2, 2:-1, 0]) sumz = bh.zeros_like(K_11)[2:-2, 1:-2] for kr in range(2): ki = 0 if kr == 1 else 1 for jp in range(2): drodyn = drdT[2:-2, 1 + jp:-2 + jp, ki:] * dTdy[2:-2, 1:-2, ki:] + \ drdS[2:-2, 1 + jp:-2 + jp, ki:] * dSdy[2:-2, 1:-2, ki:] drodzn = drdT[2:-2, 1 + jp:-2 + jp, ki:] * dTdz[2:-2, 1 + jp:-2 + jp, :-1 + kr or None] \ + drdS[2:-2, 1 + jp:-2 + jp, ki:] * \ dSdz[2:-2, 1 + jp:-2 + jp, :-1 + kr or None] syn = -drodyn / (bh.minimum(0., drodzn) - epsln) taper = dm_taper(syn) sumz[:, :, ki:] += dzw[bh.newaxis, bh.newaxis, :-1 + kr or None] \ * maskV[2:-2, 1:-2, ki:] * bh.maximum(K_iso_steep, diffloc[2:-2, 1:-2, ki:] * taper) Ai_nz[2:-2, 1:-2, ki:, jp, kr] = taper * \ syn * maskV[2:-2, 1:-2, ki:] K_22[2:-2, 1:-2, :] = sumz / (4. * dzt[bh.newaxis, bh.newaxis, :]) """ compute Ai_bx, Ai_by and K33 on top face of T cell. """ sumx = bh.zeros_like(K_11)[2:-2, 2:-2, :-1] sumy = bh.zeros_like(K_11)[2:-2, 2:-2, :-1] for kr in range(2): drodzb = drdT[2:-2, 2:-2, kr:-1 + kr or None] * dTdz[2:-2, 2:-2, :-1] \ + drdS[2:-2, 2:-2, kr:-1 + kr or None] * dSdz[2:-2, 2:-2, :-1] # eastward slopes at the top of T cells for ip in range(2): drodxb = drdT[2:-2, 2:-2, kr:-1 + kr or None] * dTdx[1 + ip:-3 + ip, 2:-2, kr:-1 + kr or None] \ + drdS[2:-2, 2:-2, kr:-1 + kr or None] * \ dSdx[1 + ip:-3 + ip, 2:-2, kr:-1 + kr or None] sxb = -drodxb / (bh.minimum(0., drodzb) - epsln) taper = dm_taper(sxb) sumx += dxu[1 + ip:-3 + ip, bh.newaxis, bh.newaxis] * \ K_iso[2:-2, 2:-2, :-1] * taper * \ sxb**2 * maskW[2:-2, 2:-2, :-1] Ai_bx[2:-2, 2:-2, :-1, ip, kr] = taper * \ sxb * maskW[2:-2, 2:-2, :-1] # northward slopes at the top of T cells for jp in range(2): facty = cosu[1 + jp:-3 + jp] * dyu[1 + jp:-3 + jp] drodyb = drdT[2:-2, 2:-2, kr:-1 + kr or None] * dTdy[2:-2, 1 + jp:-3 + jp, kr:-1 + kr or None] \ + drdS[2:-2, 2:-2, kr:-1 + kr or None] * \ dSdy[2:-2, 1 + jp:-3 + jp, kr:-1 + kr or None] syb = -drodyb / (bh.minimum(0., drodzb) - epsln) taper = dm_taper(syb) sumy += facty[bh.newaxis, :, bh.newaxis] * K_iso[2:-2, 2:-2, :-1] \ * taper * syb**2 * maskW[2:-2, 2:-2, :-1] Ai_by[2:-2, 2:-2, :-1, jp, kr] = taper * \ syb * maskW[2:-2, 2:-2, :-1] K_33[2:-2, 2:-2, :-1] = sumx / (4 * dxt[2:-2, bh.newaxis, bh.newaxis]) + \ sumy / (4 * dyt[bh.newaxis, 2:-2, bh.newaxis] * cost[bh.newaxis, 2:-2, bh.newaxis]) K_33[2:-2, 2:-2, -1] = 0.