diff --git a/zoidberg/stencil_dagp_fv.py b/zoidberg/stencil_dagp_fv.py index 26a33e9..c8042da 100755 --- a/zoidberg/stencil_dagp_fv.py +++ b/zoidberg/stencil_dagp_fv.py @@ -3,7 +3,7 @@ import time -time_start = time.time() +time_start = time.time() # noqa from boututils.datafile import DataFile as DF import numpy as np @@ -41,61 +41,93 @@ def print(*args): my_print(*args) -def doit(fn): +def load(fn): with DF(fn) as f: RZ = [f[k] for k in "RZ"] - log("opened") - RZ = np.array(RZ) - log("read") + log("opened") + RZ = np.array(RZ) + _, a, b, c = RZ.shape - A = None - nx = 2 - spc = np.linspace(0, 0.5, nx, endpoint=True) - spc2 = np.linspace(0.5, 1, nx, endpoint=True) + coefsX = np.zeros((a - 1, b, c, 2)) + coefsZ = np.zeros((a - 2, b, c, 2)) + for d, x1 in zip((coefsX, coefsZ), "XZ"): + for i, x2 in enumerate("XZ"): + key = f"dagp_fv_{x1}{x2}" + fixup2(d[..., i], f[key]) + key = "dagp_fv_volume" + volume = f[key] + log("done reading") + return RZ, volume, coefsX, coefsZ + + +def doit(pols, plot=False): + # with DF(fn) as f: + # RZ = [f[k] for k in "RZ"] + # log("opened") + # RZ = np.array(RZ) + # log("read") + + RZs = [] ### Calculate Volume of the cell # ### Go in a line around the cell - # xf is the normalised coordinates in x - # yf is the normalised coordinates in z - # i and j are the offset in x and z - # i,j = 0 -> in x and z in lower direction - for xf, yf, i, j in ( - (spc2, 0.5, 1, 1), - (0.5, spc2[::-1], 1, 1), - (0.5, spc[::-1], 1, 0), - (spc2[::-1], 0.5, 1, 0), - (spc[::-1], 0.5, 0, 0), - (0.5, spc, 0, 0), - (spc, 0.5, 0, 1), - (0.5, spc2, 0, 1), - ): - log(f"doing offset ({i}, {j})") - dx = np.arange(2) + i - dy = np.arange(2) + j - - def slc(i): - i1 = -2 + i - return slice(i, i1 or None) - - xii = [ - [np.roll(RZ[:, slc(i)], -j + 1, axis=-1)[..., None] for i in dy] for j in dx - ] - xa, xb = xii - x00, x01 = xa - x10, x11 = xb - x = yf * (xf * x00 + (1 - xf) * x10) + (1 - yf) * (xf * x01 + (1 - xf) * x11) - y = x[1] - dy = (y[..., 1] - y[..., 0])[..., None] - xx = (x[0, ..., :-1] + x[0, ..., 1:]) / 2 - if A is None: - A = np.zeros(xx[..., 0].shape) - Ar = np.zeros(xx[..., 0].shape) - A -= np.sum(xx * dy, axis=-1) - Ar -= np.sum(0.5 * xx * xx * dy, axis=-1) - # plt.plot(x[0], x[1], "gx") + + A = np.empty((pols[0].nx, len(pols), pols[0].nz)) + Ar = np.empty_like(A) + + for gi, g in enumerate(pols): + n = 50 + x0 = np.arange(g.nx)[1:-1, None, None] + y0 = np.arange(g.nz)[None, :, None] + x1 = x0[:, :, 0] + y1 = y0[:, :, 0] + x = [] + y = [] + zero = np.zeros((g.nx - 2, g.nz, n)) + x += [np.linspace(x1 - 0.5, x1 + 0.5, n).transpose(1, 2, 0)] + y += [y0 + 0.5] + x += [x0 + 0.5] + y += [np.linspace(y1 + 0.5, y1 - 0.5, n).transpose(1, 2, 0)] + x += [np.linspace(x1 + 0.5, x1 - 0.5, n).transpose(1, 2, 0)] + y += [y0 - 0.5] + x += [x0 - 0.5] + y += [np.linspace(y1 - 0.5, y1 + 0.5, n).transpose(1, 2, 0)] + x = np.concatenate([k + zero for k in x], axis=-1) + y = np.concatenate([k + zero for k in y], axis=-1) + RZ = np.array(g.getCoordinate(x, y)) + + dy = RZ[1, ..., 1:] - RZ[1, ..., :-1] + xx = (RZ[0, ..., :-1] + RZ[0, ..., 1:]) / 2 + + A[1:-1, gi] = -np.sum(xx * dy, axis=-1) + Ar[1:-1, gi] = -np.sum(0.5 * xx * xx * dy, axis=-1) + + RZs += [RZ] + RZs = np.array(RZs) + print(RZs.shape) + # 3.160 s: (4, 2, 6, 32, 200) + RZs = RZs.transpose(1, 2, 0, 3, 4) + print(RZs.shape) + volume = Ar - A.shape, Ar.shape, RZ.shape + # f = + # V = [f[k] for k in ("dx", "dy", "dz", "J")] + + RZ = np.array( + [ + g.getCoordinate( + *np.meshgrid(np.arange(g.nx), np.arange(g.nz), indexing="ij") + ) + for g in pols + ] + ).transpose(1, 2, 0, 3) + + RZ = np.array([(g.R, g.Z) for g in pols]).transpose(1, 2, 0, 3) + + # volume = dx * dy * dz * J + # volume = V[0] * V[1] * V[2] * V[3] + # A.shape, Ar.shape, RZ.shape # AreaXplus # Z @@ -115,24 +147,20 @@ def slc(i): # |------------------------> x # # We are calculating the area of the surface denoted by X - # Note that we need to split it in two parts. Maybe? - nx = 2 - spc3 = np.linspace(0, 1, nx) - # pos goes in +z direction - # First segments starts at the centre of b,c,d,o - # and goes to centre of o,d - startpoint = toCent(RZ) - midpoint = (RZ[:, :-1] + RZ[:, 1:]) / 2 - endpoint = np.roll(startpoint, -1, -1) + # Note that we need to split it in two parts. log("calculating pos ...") - pos = np.concatenate( - [ - (1 - spc3[:-1]) * startpoint[..., None] + spc3[:-1] * midpoint[..., None], - (1 - spc3) * midpoint[..., None] + spc3 * endpoint[..., None], - ], - axis=-1, + # pos = RZs[..., :n] + tmp = np.meshgrid( + np.arange(g.nx - 1) + 0.5, + np.arange(g.nz) - 0.5, + np.linspace(0, 1, n), + indexing="ij", ) + tmp[1] += tmp[2] + print(tmp[0]) + print(tmp[1]) + pos = np.array([g.getCoordinate(*tmp[:2]) for g in pols]).transpose(1, 2, 0, 3, 4) log("done") # direction of edge, length ~ 1/nx dR = pos[..., :-1] - pos[..., 1:] @@ -143,7 +171,6 @@ def slc(i): assert dR.shape[0] == 2 dRr = np.array((-dR[1], dR[0])) dRr /= np.sqrt(np.sum(dRr**2, axis=0)) - print(dRr.shape, area.shape) dRr *= area dRr = np.sum(dRr, axis=-1) print(np.nanmean(dRr)) @@ -158,12 +185,9 @@ def slc(i): # dzR /= (np.sum(dzR**2, axis=0)) log("starting solve") dxzR = np.array((dxR, dzR)).transpose(2, 3, 4, 1, 0) - coefs = np.linalg.solve(dxzR, dRr.transpose(1, 2, 3, 0)) + coefsX = np.linalg.solve(dxzR, dRr.transpose(1, 2, 3, 0)) log("done") - areaX = area - coefsX = coefs - # In[154]: # AreaZplus @@ -182,20 +206,17 @@ def slc(i): # | g f e # | # |------------------------> x - nx = 3 - spc3 = np.linspace(0, 1, nx) - cent = np.roll(toCent(RZ), -1, -1) - startpoint = cent[:, :-1] # a - midpoint = (RZ[:, 1:-1] + np.roll(RZ[:, 1:-1], -1, -1)) / 2 # b - endpoint = cent[:, 1:] # c + log("concatenate") - pos = np.concatenate( - [ - (1 - spc3[:-1]) * startpoint[..., None] + spc3[:-1] * midpoint[..., None], - (1 - spc3) * midpoint[..., None] + spc3 * endpoint[..., None], - ], - axis=-1, + tmp = np.meshgrid( + np.arange(g.nx - 2) + 0.5, + np.arange(g.nz) + 0.5, + np.linspace(0, 1, n), + indexing="ij", ) + tmp[0] += tmp[2] + pos = np.array([g.getCoordinate(*tmp[:2]) for g in pols]).transpose(1, 2, 0, 3, 4) + dR = pos[..., :-1] - pos[..., 1:] dX = np.sqrt(np.sum(dR**2, axis=0)) area = dX * (pos[0, ..., 1:] + pos[0, ..., :-1]) * 0.5 @@ -211,15 +232,16 @@ def slc(i): dzR = (np.roll(RZ, -1, axis=-1) - RZ)[:, 1:-1] dxzR = np.array((dxR, dzR)).transpose(2, 3, 4, 1, 0) log("solving again") - coefs = np.linalg.solve(dxzR, dRr.transpose(1, 2, 3, 0)) + coefsZ = np.linalg.solve(dxzR, dRr.transpose(1, 2, 3, 0)) log("done") - coefs.shape - areaZ = area - coefsZ = coefs + test(RZ, volume, coefsX, coefsZ, plot=plot) + return write(RZ, volume, coefsX, coefsZ) # In[155]: + +def test(RZ, volume, coefsX, coefsZ, plot=False): inp = np.sin(RZ[1]) ana = -np.sin(RZ[1]) # + np.cos(RZ[0])/RZ[0] if 0: @@ -232,6 +254,10 @@ def slc(i): inp = RZ[0] ** 3 ana = 6 * np.sin(RZ[0]) + # print("Fuck it up!!!") + # coefsX[..., 1] *= -1 + # coefsZ[..., 0] *= -1 + # In[156]: def xp(ijk): @@ -285,26 +311,55 @@ def zm(i, j, k): r[0] = 0 r[-1] = 0 - print(fn, "error:", np.mean(l2(result[1:-1] / volume - ana[1:-1]))) - - def fixup(d): - c1 = np.zeros(RZ[0].shape) - if c1.shape[0] == d.shape[0] + 1: - c1[:-1] = d - else: - c1[1:-1] = d - # c1 = np.roll(c1, -1, -1) - return c1 - + if plot: + import matplotlib.pyplot as plt + + f, axs = plt.subplots(1, 3, sharex=True, sharey=True) + res = result / volume + for d, ax, t in zip( + (res, ana, res - ana), + axs, + ("computed", "analytical", "error"), + ): + slc = slice(1, -1), 1 + per = np.percentile(d[slc], [0, 2, 98, 100]) + print(per) + p = ax.pcolormesh(*[k[slc] for k in RZ], d[slc], vmin=per[1], vmax=per[2]) + ax.set_title(t) + plt.colorbar(p, ax=ax) + plt.show() + print("error:", np.mean(l2(result[1:-1] / volume[1:-1] - ana[1:-1]))) + + +def fixup(RZ, d): + c1 = np.zeros(RZ[0].shape) + if c1.shape[0] == d.shape[0] + 1: + c1[:-1] = d + else: + c1[1:-1] = d + # c1 = np.roll(c1, -1, -1) + return c1 + + +def fixup2(d, val): + if val.shape[0] == d.shape[0] + 1: + d[:] = val[:-1] + else: + d[:] = val[1:-1] + + +def write(RZ, volume, coefsX, coefsZ): + ret = {} log("writing") - with DF(fn, write=True) as f: - for d, x1 in zip((coefsX, coefsZ), "XZ"): - for i, x2 in enumerate("XZ"): - key = f"dagp_fv_{x1}{x2}" - f.write(key, fixup(d[..., i])) + for d, x1 in zip((coefsX, coefsZ), "XZ"): + for i, x2 in enumerate("XZ"): + key = f"dagp_fv_{x1}{x2}" + ret[key] = fixup(RZ, d[..., i]) key = "dagp_fv_volume" - f.write(key, fixup(volume)) + ret[key] = volume log("done") + print(ret.keys()) + return ret if __name__ == "__main__": @@ -312,7 +367,11 @@ def fixup(d): while flag in sys.argv: verbose += step sys.argv.remove(flag) + plot = "-p" in sys.argv + if plot: + sys.argv.remove("-p") for fn in sys.argv[1:]: print(fn) - doit(fn) + # doit(fn) + test(*load(fn), plot=plot)