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

@@ -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,34 +311,67 @@ 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__":
     for flag, step in (("-v", +1), ("-q", -1)):
         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)