From b73592c90d06c1f5bafdaab504d19d2ce9751cf9 Mon Sep 17 00:00:00 2001 From: David Bold Date: Fri, 8 Nov 2024 10:07:18 +0100 Subject: [PATCH] Change dagp_fv to use smooth volumes Use smooth volumes, based on the curvilinear grid. This should improve agreement of the volumes with the normal interpretation of the cells, and thus improve confinement in the normal interpretation. --- zoidberg/stencil_dagp_fv.py | 263 ++++++++++++++++++++++-------------- 1 file changed, 161 insertions(+), 102 deletions(-) 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)