test_cfv_dissdisp.py

Go to the documentation of this file.

import DNDSR
 
# print(DNDSR.__dict__)
 
from DNDSR import DNDS, Geom, CFV
from DNDSR.Geom.utils import *
import numpy as np
import json
 
# vfvSettings = json.loads(
#     "".join(
#         [
#             line if not line.strip().startswith("//") else ""
#             for line in """{
#     "maxOrder": 3,
#     "intOrder": 5,
#     "intOrderVR": 5,
#     "cacheDiffBase": true,
#     "jacobiRelax": 1,
#     "SORInstead": false,
#     "smoothThreshold": 1e-3,
#     "WBAP_nStd": 10.0,
#     "normWBAP": false,
#     "subs2ndOrder": 1,
#     "subs2ndOrderGGScheme": 0,
#     "baseSettings": {
#         "localOrientation": false,
#         "anisotropicLengths": false
#     },
#     "functionalSettings": {
#         // "scaleType": "MeanAACBB",
#         "dirWeightScheme": "HQM_OPT",
#         // "dirWeightScheme": "ManualDirWeight",
#         // "manualDirWeights": [
#         //     1.0,
#         //     1,
#         //     0,
#         //     0
#         // ],
#         "geomWeightScheme": "HQM_SD",
#         "geomWeightPower": 0.5,
#         "geomWeightBias": 1,
#         // "geomWeightScheme": "SD_Power",
#         // "geomWeightPower1": -0.5,
#         // "geomWeightPower2": 0.5,
#         // "useAnisotropicFunctional": true,
#         // // "anisotropicType": "InertiaCoordBB",
#         // "inertiaWeightPower": 0,
#         // "scaleMultiplier": 1,
#         "_tail": 0
#     }
# }
# """.splitlines()
#         ]
#     )
# )
# print(vfvSettings)
 
 
def default_VRSettings():
    return {
        "maxOrder": 3,
        "intOrder": 5,
        "intOrderVR": 5,
        "cacheDiffBase": True,
        "jacobiRelax": 1,
        "SORInstead": False,
        "smoothThreshold": 0.001,
        "WBAP_nStd": 10.0,
        "normWBAP": False,
        "subs2ndOrder": 1,
        "subs2ndOrderGGScheme": 0,
        "baseSettings": {"localOrientation": False, "anisotropicLengths": False},
        "functionalSettings": {
            "dirWeightScheme": "HQM_OPT",
            "geomWeightScheme": "HQM_SD",
            "geomWeightPower": 0.5,
            "geomWeightBias": 1,
            "_tail": 0,
        },
    }
 
 
class WaveTester:
 
    def __init__(
        self,
        mpi,
        vfvSettings=default_VRSettings(),
        ax=1.0,
        ay=0.0,
        sigma=0.0,
        AR=1.0,
        ob=0.0,
        batch_size=1,
    ):
        self.batch_size = batch_size
        self.mpi = mpi
        self.ax = ax
        self.ay = ay
        self.sigma = sigma
 
        meshFile = os.path.join(
            os.path.dirname(__file__),
            "..",
            "..",
            "data",
            "mesh",
            "Uniform_5x5_wave.cgns",
        )
 
        transform = np.array(
            [
                [1, 0, 0],
                [0, 1.0 / AR, 0],
                [0, 0, 1],
            ],
            dtype=np.float64,
        )
        transform = (
            np.array(
                [
                    [1, ob, 0],
                    [0, 1, 0],
                    [0, 0, 1],
                ],
                dtype=np.float64,
            )
            @ transform
        )
 
        mesh, reader, name2Id = create_mesh_from_CGNS(
            meshFile,
            mpi,
            2,
            periodic_geometry={
                "translation1": [5, 0, 0],
                "translation2": [0, 5, 0],
            },
        )
        mesh.SetPeriodicGeometry(
            translation1=transform @ np.array([5, 0, 0]),
            translation2=transform @ np.array([0, 5, 0]),
        )
        coord_father = np.array(mesh.coords.father.data(), copy=False).reshape(-1, 3)
        coord_son = np.array(mesh.coords.son.data(), copy=False).reshape(-1, 3)
 
        coord_father[:] = coord_father @ transform.transpose()
        coord_son[:] = coord_son @ transform.transpose()
 
        meshBnd, readerBnd = create_bnd_mesh(mesh)
 
        self.reader = reader
        self.name2Id = name2Id
        self.meshBnd = meshBnd
        self.readerBnd = readerBnd
 
        vfv = CFV.VariationalReconstruction_2(mpi, mesh)
        self.vfvSettings = vfvSettings
        vfv.ParseSettings(vfvSettings)
        vfv.SetPeriodicTransformationsNoOp()
 
        bcid_2_bcweight_map = {}
        for name, id in name2Id.n2id_map.items():
            # if name == "WALL":
            bcid_2_bcweight_map[(id, 0)] = 1.0
            if name.startswith("PERIODIC"):
                bcid_2_bcweight_map[(id, 0)] = 1.0
                bcid_2_bcweight_map[(id, 1)] = 1.0
                bcid_2_bcweight_map[(id, 2)] = 1.0
                bcid_2_bcweight_map[(id, 3)] = 1.0
        vfv.ConstructMetrics()
        vfv.ConstructBaseAndWeight_map(bcid_2_bcweight_map)
        vfv.ConstructRecCoeff()
 
        self.mesh = mesh
        self.vfv = vfv
        self.eval = CFV.ModelEvaluator(
            mesh, vfv, {"ax": ax, "ay": ay, "sigma": sigma}, batch_size
        )
 
        u_real, rhs_real = [CFV.tUDof_D() for _ in range(2)]
        uRec_real, uRecNew_real = [CFV.tURec_D() for _ in range(2)]
        u_imag, rhs_imag = [CFV.tUDof_D() for _ in range(2)]
        uRec_imag, uRecNew_imag = [CFV.tURec_D() for _ in range(2)]
        self.u_list = [u_real, rhs_real, u_imag, rhs_imag]
        self.uRec_list = [uRec_real, uRecNew_real, uRec_imag, uRecNew_imag]
 
        for u_ in self.u_list:
            vfv.BuildUDof_D(u_, batch_size)
        for uRec_ in self.uRec_list:
            vfv.BuildURec_D(uRec_, batch_size)
 
        nFree = 0
        for iCell in range(mesh.NumCell()):
            if (
                np.linalg.norm(
                    vfv.GetCellBary(iCell).flatten()
                    - transform @ np.array([0.5, 0.5, 0])
                )
                < 1e-10
            ):
                self.iCellFree = iCell
                nFree += 1
        if not nFree == 1:
            raise ValueError("nFree not 1")
 
    def update_vfv_settings(self, vfvSettings):
        mpi = self.mpi
        vfv = self.vfv
        mesh = self.mesh
        eval = self.eval
        name2Id = self.name2Id
        vfv.ParseSettings(vfvSettings)
        vfv.SetPeriodicTransformationsNoOp()
 
        bcid_2_bcweight_map = {}
        for name, id in name2Id.n2id_map.items():
            # if name == "WALL":
            bcid_2_bcweight_map[(id, 0)] = 1.0
            if name.startswith("PERIODIC"):
                bcid_2_bcweight_map[(id, 0)] = 1.0
                bcid_2_bcweight_map[(id, 1)] = 1.0
                bcid_2_bcweight_map[(id, 2)] = 1.0
                bcid_2_bcweight_map[(id, 3)] = 1.0
        vfv.ConstructMetrics()
        vfv.ConstructBaseAndWeight_map(bcid_2_bcweight_map)
        vfv.ConstructRecCoeff()
 
    def update_eval_settings(self, ax=1.0, ay=0, sigma=0.0):
        self.ax = ax
        self.ay = ay
        self.sigma = sigma
        self.eval.ParseSettings({"ax": ax, "ay": ay, "sigma": sigma})
 
    def get_rhsFreeComplex(self):
        u_real, rhs_real, u_imag, rhs_imag = self.u_list
        return np.array(rhs_real[self.iCellFree]) + 1j * np.array(
            rhs_imag[self.iCellFree]
        )
 
    def get_uFreeComplex(self):
        u_real, rhs_real, u_imag, rhs_imag = self.u_list
        return np.array(u_real[self.iCellFree]) + 1j * np.array(u_imag[self.iCellFree])
 
    def set_uFree(self, r: np.ndarray, i: np.ndarray):
        u_real, rhs_real, u_imag, rhs_imag = self.u_list
        np.array(u_real[self.iCellFree], copy=False)[:] = r
        np.array(u_imag[self.iCellFree], copy=False)[:] = i
 
    def test_one_wave(
        self,
        kx: np.ndarray,
        ky: np.ndarray,
        n_iter=10000,
        tol=1e-15,
        n_print=0,
        n_iter_min=1,
        u_free=1 + 0j,
        rhsOptions=CFV.ModelEvaluator.EvaluateRHSOptions(),
    ):
        if kx.size != self.batch_size or ky.size != self.batch_size:
            raise ValueError("input k size not right")
 
        kx = kx.flatten()
        ky = ky.flatten()
        vfv = self.vfv
        mesh = self.mesh
        eval = self.eval
 
        u_real, rhs_real, u_imag, rhs_imag = self.u_list
        uRec_real, uRecNew_real, uRec_imag, uRecNew_imag = self.uRec_list
 
        self.set_uFree(u_free.real, u_free.imag)
        self.uSync(kx, ky)
 
        self.DoReconstruction(kx, ky, n_iter, tol, n_print, n_iter_min=n_iter_min)
 
        eval.EvaluateRHS(rhs_real, u_real, uRec_real, 0.0, options=rhsOptions)
        eval.EvaluateRHS(rhs_imag, u_imag, uRec_imag, 0.0, options=rhsOptions)
        return (
            np.array(rhs_real[self.iCellFree]) + 1j * np.array(rhs_imag[self.iCellFree])
        ).flatten()
 
    def test_conv_rate(
        self,
        dTau=10,
        dT=1e100,
        kx: np.ndarray = np.array([0]),
        ky: np.ndarray = np.array([0]),
        n_iter=10000,
        tol=1e-15,
        n_print=0,
        n_iter_min=1,
        singlegrid_niter=1,
        multigrid_niters=(0, 0),
        use_diff_Jacobi=False,
        top_override=-1,
        multigrid_res_fact=(1.0, 1.0),
        multigrid_dtau_fact=(1.0, 1.0),
    ):
        if kx.size != self.batch_size or ky.size != self.batch_size:
            raise ValueError("input k size not right")
 
        kx = kx.flatten()
        ky = ky.flatten()
        
        vfv = self.vfv
        mesh = self.mesh
        eval = self.eval
 
        u_real, rhs_real, u_imag, rhs_imag = self.u_list
        uRec_real, uRecNew_real, uRec_imag, uRecNew_imag = self.uRec_list
 
        rhsParamsGeneral = {
            "n_iter": n_iter,
            "tol": tol,
            "n_print": n_print,
            "n_iter_min": n_iter_min,
        }
        rhsOptionsTop = eval.EvaluateRHSOptions()
        if top_override == 1:
            rhsOptionsTop.direct2ndRec = True
            rhsOptionsTop.direct2ndRec1stConv = False
        if top_override == 0:
            rhsOptionsTop.direct2ndRec = True
            rhsOptionsTop.direct2ndRec1stConv = True
 
        # build the fv jacobian
        J = np.complex128(0)
        c2f = np.array(mesh.cell2face[self.iCellFree])
        xcC = vfv.GetCellBary(self.iCellFree)
 
        for ic2f, iFace in enumerate(c2f):
            iCellOther = mesh.CellFaceOther(self.iCellFree, iFace)
            assert iCellOther != DNDS.UnInitIndex
            normOut = vfv.GetFaceNormFromCell(iFace, self.iCellFree, -1, -1)
            if not mesh.CellIsFaceBack(self.iCellFree, iFace):
                normOut *= -1.0
            a_out = normOut[0] * self.ax + normOut[1] * self.ay
 
            a_vis = (
                self.sigma
                * vfv.GetFaceArea(iFace)
                * (
                    1.0 / vfv.GetCellVol(self.iCellFree)
                    + 1.0 / vfv.GetCellVol(iCellOther)
                )
            )
 
            xcr = vfv.GetCellBary(iCellOther) - xcC
            wave_val = np.exp(1j * (kx * xcr[0] + ky * xcr[1])).reshape(-1, 1)
            dFdu = (1 + wave_val) * 0.5 * a_out - 0.5 * np.abs(a_out) * (wave_val - 1)
            dFdu += -0.5 * a_vis * (wave_val - 1)
 
            J -= dFdu * vfv.GetFaceArea(iFace) / vfv.GetCellVol(self.iCellFree)
 
        if use_diff_Jacobi:
            eps = 1e-6
            J_top = (
                self.test_one_wave(
                    kx,
                    ky,
                    rhsOptions=rhsOptionsTop,
                    **rhsParamsGeneral,
                )
                - self.test_one_wave(
                    kx,
                    ky,
                    rhsOptions=rhsOptionsTop,
                    u_free=0j + 1 - eps,
                    **rhsParamsGeneral,
                )
            ) / eps
        else:
            J_top = J
 
        if use_diff_Jacobi:
            eps = 1e-6
            options = eval.EvaluateRHSOptions()
            options.direct2ndRec = True
            options.direct2ndRec1stConv = False
            J_p1 = (
                self.test_one_wave(kx, ky, rhsOptions=options, **rhsParamsGeneral)
                - self.test_one_wave(
                    kx, ky, u_free=0j + 1 - eps, rhsOptions=options, **rhsParamsGeneral
                )
            ) / eps
        else:
            J_p1 = J
 
        np.array(u_real[self.iCellFree], copy=False)[:] = 1
        np.array(u_imag[self.iCellFree], copy=False)[:] = 0
        self.uSync(kx, ky)
 
        ################# OTop
 
        for iter in range(singlegrid_niter):
            self.DoReconstruction(kx, ky, n_iter, tol, n_print, n_iter_min=n_iter_min)
            eval.EvaluateRHS(rhs_real, u_real, uRec_real, 0.0, options=rhsOptionsTop)
            eval.EvaluateRHS(rhs_imag, u_imag, uRec_imag, 0.0, options=rhsOptionsTop)
 
            uFreeNew = self.get_uFreeComplex() + (
                self.get_rhsFreeComplex() - self.get_uFreeComplex() / dT
            ) / (-J_top + 1.0 / (dTau * 1.0) + 1.0 / dT)
            self.set_uFree(np.real(uFreeNew), np.imag(uFreeNew))
            self.uSync(kx, ky)
 
        self.DoReconstruction(kx, ky, n_iter, tol, n_print, n_iter_min=n_iter_min)
        eval.EvaluateRHS(rhs_real, u_real, uRec_real, 0.0, options=rhsOptionsTop)
        eval.EvaluateRHS(rhs_imag, u_imag, uRec_imag, 0.0, options=rhsOptionsTop)
        rhs_top = self.get_rhsFreeComplex() - self.get_uFreeComplex() / dT
 
        ################# O1
 
        options1 = eval.EvaluateRHSOptions()
        options1.direct2ndRec = True
        options1.direct2ndRec1stConv = False
        rhs1_init = None
        # not need reconstruction
        # self.DoReconstruction(kx, ky, n_iter, tol, n_print, n_iter_min=n_iter_min)
        eval.EvaluateRHS(rhs_real, u_real, uRec_real, 0.0, options=options1)
        eval.EvaluateRHS(rhs_imag, u_imag, uRec_imag, 0.0, options=options1)
        rhs1_init = self.get_rhsFreeComplex() - self.get_uFreeComplex() / dT
        for iter in range(multigrid_niters[0]):
            uFreeNew = self.get_uFreeComplex() + (
                self.get_rhsFreeComplex()
                - self.get_uFreeComplex() / dT
                - rhs1_init
                + rhs_top * multigrid_res_fact[0]
            ) / (-J_p1 + 1.0 / (dTau * multigrid_dtau_fact[0]) + 1.0 / dT)
            self.set_uFree(np.real(uFreeNew), np.imag(uFreeNew))
            self.uSync(kx, ky)
            eval.EvaluateRHS(rhs_real, u_real, uRec_real, 0.0, options=options1)
            eval.EvaluateRHS(rhs_imag, u_imag, uRec_imag, 0.0, options=options1)
        rhs_1 = (
            self.get_rhsFreeComplex()
            - self.get_uFreeComplex() / dT
            - rhs1_init
            + rhs_top * multigrid_res_fact[0]
        )
 
        ################# O0
 
        options0 = eval.EvaluateRHSOptions()
        options0.direct2ndRec = True
        options0.direct2ndRec1stConv = True
        rhs0_init = None
        # not need reconstruction
        # self.DoReconstruction(kx, ky, n_iter, tol, n_print, n_iter_min=n_iter_min)
        eval.EvaluateRHS(rhs_real, u_real, uRec_real, 0.0, options=options0)
        eval.EvaluateRHS(rhs_imag, u_imag, uRec_imag, 0.0, options=options0)
        rhs0_init = self.get_rhsFreeComplex() - self.get_uFreeComplex() / dT
        for iter in range(multigrid_niters[1]):
            uFreeNew = self.get_uFreeComplex() + (
                self.get_rhsFreeComplex()
                - self.get_uFreeComplex() / dT
                - rhs0_init
                + rhs_1 * multigrid_res_fact[1]
            ) / (-J + 1.0 / (dTau * multigrid_dtau_fact[1]) + 1.0 / dT)
            self.set_uFree(np.real(uFreeNew), np.imag(uFreeNew))
            self.uSync(kx, ky)
            eval.EvaluateRHS(rhs_real, u_real, uRec_real, 0.0, options=options0)
            eval.EvaluateRHS(rhs_imag, u_imag, uRec_imag, 0.0, options=options0)
 
        return self.get_uFreeComplex().flatten()
 
    def uRecSync(self, kx: np.ndarray, ky: np.ndarray):
        vfv = self.vfv
        uRec_real, uRecNew_real, uRec_imag, uRecNew_imag = self.uRec_list
        for iCell in range(self.mesh.NumCell()):
            if iCell == self.iCellFree:
                continue
            xc = vfv.GetCellBary(iCell)
            xcr = xc - vfv.GetCellBary(self.iCellFree)
            wave_val = np.exp(1j * (kx * xcr[0] + ky * xcr[1]))
            wave_val = wave_val.reshape(1, -1)
            # fmt: off
            np.array(uRec_real[iCell], copy=False)[:] = \
                np.array(uRec_real[self.iCellFree]) * np.real(wave_val) - \
                np.array(uRec_imag[self.iCellFree]) * np.imag(wave_val)
            np.array(uRec_imag[iCell], copy=False)[:] = \
                np.array(uRec_imag[self.iCellFree]) * np.real(wave_val) + \
                np.array(uRec_real[self.iCellFree]) * np.imag(wave_val)
            # fmt: on
 
    def uSync(self, kx: np.ndarray, ky: np.ndarray):
        vfv = self.vfv
        u_real, rhs_real, u_imag, rhs_imag = self.u_list
        for iCell in range(self.mesh.NumCell()):
            if iCell == self.iCellFree:
                continue
            xc = vfv.GetCellBary(iCell)
            xcr = xc - vfv.GetCellBary(self.iCellFree)
            wave_val = np.exp(1j * (kx * xcr[0] + ky * xcr[1])).reshape(-1, 1)
            # fmt: off
            np.array(u_real[iCell], copy=False)[:] = \
                np.array(u_real[self.iCellFree]) * np.real(wave_val) - \
                np.array(u_imag[self.iCellFree]) * np.imag(wave_val)
            np.array(u_imag[iCell], copy=False)[:] = \
                np.array(u_imag[self.iCellFree]) * np.real(wave_val) + \
                np.array(u_real[self.iCellFree]) * np.imag(wave_val)
            # fmt: on
 
    def DoReconstruction(
        self, kx, ky, n_iter=10000, tol=1e-15, n_print=0, n_iter_min=1
    ):
        vfv = self.vfv
        mesh = self.mesh
        eval = self.eval
        u_real, rhs_real, u_imag, rhs_imag = self.u_list
        uRec_real, uRecNew_real, uRec_imag, uRecNew_imag = self.uRec_list
 
        # print(vfv)
 
        if self.vfvSettings["subs2ndOrder"] == 0 or self.vfvSettings["maxOrder"] > 1:
            # use matrices:
            c2f = np.array(mesh.cell2face[self.iCellFree])
            matrixAAInvB = vfv.matrixAAInvB
            vectorAInvB = vfv.vectorAInvB
            AInv = np.array(matrixAAInvB[self.iCellFree, 0])
            # print(AInv)
            M, N = AInv.shape
            assert M == N
            MatC = np.eye(M, M, dtype=np.complex128)
            MatC = np.repeat(MatC[np.newaxis, :, :], self.batch_size, axis=0)
            RhsC = np.zeros((self.batch_size, M, 1), dtype=np.complex128)
            xcC = vfv.GetCellBary(self.iCellFree)
            uC = np.array(u_real[self.iCellFree]) + 1j * np.array(
                u_imag[self.iCellFree]
            )
 
            for ic2f, iFace in enumerate(c2f):
                iCellOther = mesh.CellFaceOther(self.iCellFree, iFace)
                xcr = vfv.GetCellBary(iCellOther) - xcC
                wave_val = np.exp(1j * (kx * xcr[0] + ky * xcr[1]))
                # print(np.array(matrixAAInvB[self.iCellFree, ic2f + 2], copy=False))
                # fmt: off
                MatC -= (
                    wave_val[:, np.newaxis, np.newaxis]
                    * np.array(matrixAAInvB[self.iCellFree, ic2f + 1], copy=False)[np.newaxis, :, :]
                )
                # fmt: on
                uOther = np.array(u_real[iCellOther]) + 1j * np.array(
                    u_imag[iCellOther]
                )
                RhsC += (uOther - uC)[:, :, np.newaxis] * np.array(
                    vectorAInvB[self.iCellFree, ic2f], copy=False
                )[np.newaxis, :, :]
            # print(np.linalg.eig(MatC))
            uRecFree = np.linalg.solve(MatC, RhsC)
            # fmt: off
            np.array(uRec_real[self.iCellFree], copy=False)[:] = np.real(uRecFree).reshape(-1, M).T
            np.array(uRec_imag[self.iCellFree], copy=False)[:] = np.imag(uRecFree).reshape(-1, M).T
            # fmt: on
            self.uRecSync(kx, ky)
            return
 
        for iCell in range(self.mesh.NumCell()):
            np.array(uRec_real[iCell], copy=False)[:] = 0.0
            np.array(uRec_imag[iCell], copy=False)[:] = 0.0
            # print(f"=== Cell {iCell}")
            # print(xcr)
            # print(wave_val)
            # print(u_real[iCell].tolist())
            # print(u_imag[iCell].tolist())
 
        for iter in range(1, 1 + n_iter):
            uRecPrevFree = np.array(
                (
                    np.array(uRec_real[self.iCellFree]),
                    np.array(uRec_imag[self.iCellFree]),
                )
            )
            eval.DoReconstructionIter(uRec_real, uRecNew_real, u_real, 0.0, False)
            eval.DoReconstructionIter(uRec_imag, uRecNew_imag, u_imag, 0.0, False)
            self.uRecSync(kx, ky)
            uRecNewFree = np.array(
                (
                    np.array(uRec_real[self.iCellFree]),
                    np.array(uRec_imag[self.iCellFree]),
                )
            )
            # print((uRecPrevFree - uRecNewFree).reshape((2, 9)))
            incNorm = np.linalg.norm((uRecPrevFree - uRecNewFree).flatten(), ord=1)
 
            if iter >= n_iter_min and n_print > 0 and iter % n_print == 0:
                print(f"iter {iter} [{incNorm}]")
            if incNorm < tol:
                break
 
        # print(np.array(uRec_real[self.iCellFree]).flatten())
        # print(np.real(uRecFree).flatten())
 
 
def test():
    print(CFV.tUDof_1)
    # CFV.VariationalReconstruction_2()
 
    mpi = DNDS.MPIInfo()
    mpi.setWorld()
    tester = WaveTester(mpi, batch_size=3)
    # print(
    #     tester.test_one_wave(
    #         kx=np.pi * np.array([0, 0.1, 0.1]),
    #         ky=np.pi * np.array([0, 0, 0]),
    #     )
    # )
 
    # kxs = np.linspace(0, 1, 101) * np.pi
    # kappaNum = np.zeros_like(kxs, dtype=np.complex128)
    # for ikx, kx in enumerate(kxs):
    #     kappaNum[ikx] = 1j * tester.test_one_wave(kx, 0.0)
 
    print(
        tester.test_conv_rate(
            10,
            dT=0.1,
            kx=np.pi * np.array([0, 0.1, 0.1]),
            ky=np.pi * np.array([0, 0, 0]),
            multigrid_niters=(0, 0),
        )
    )
 
 
if __name__ == "__main__":
    test()