From 90a3b1d75c8be12e77f89b32d0d39056ffe9ade3 Mon Sep 17 00:00:00 2001
From: Ping He <phe@iastate.edu>
Date: Sun, 19 Nov 2023 08:39:59 -0600
Subject: [PATCH] Added options to use pre-computed PC mats and udpate part of
 PC mat values when solving the adj.

---
 dafoam/pyDAFoam.py                            | 122 ++++++-----
 .../DATurbulenceModel/DASpalartAllmarasFv3.C  |  41 ++++
 .../DATurbulenceModel/DASpalartAllmarasFv3.H  |   7 +
 .../DATurbulenceModel/DATurbulenceModel.C     |  12 ++
 .../DATurbulenceModel/DATurbulenceModel.H     |   7 +
 src/adjoint/DAResidual/DAResidual.C           |   7 +
 src/adjoint/DAResidual/DAResidual.H           |   3 +
 src/adjoint/DAResidual/DAResidualPimpleFoam.C | 193 ++++++++++++++++++
 src/adjoint/DAResidual/DAResidualPimpleFoam.H |   2 +
 src/adjoint/DASolver/DASolver.C               | 125 +++++++++++-
 src/adjoint/DASolver/DASolver.H               |   8 +
 src/pyDASolvers/DASolvers.H                   |  18 +-
 src/pyDASolvers/pyDASolvers.pyx               |   8 +
 13 files changed, 477 insertions(+), 76 deletions(-)

diff --git a/dafoam/pyDAFoam.py b/dafoam/pyDAFoam.py
index 7926179e..3ac42e57 100755
--- a/dafoam/pyDAFoam.py
+++ b/dafoam/pyDAFoam.py
@@ -465,6 +465,8 @@ def __init__(self):
             "periodicity": -1.0,
             "objFuncStartTime": -1.0,
             "objFuncEndTime": -1.0,
+            "PCMatPrecomputeInterval": 100,
+            "PCMatUpdateInterval": 1,
         }
 
         ## At which iteration should we start the averaging of objective functions.
@@ -861,8 +863,8 @@ def __init__(self, comm=None, options=None):
         # initialize the adjoint vector dict
         self.adjVectors = self._initializeAdjVectors()
 
-        # initialize the dRdWOldTPsi vectors
-        self._initializeTimeAccurateAdjointVectors()
+        # initialize the dRdWTPC dict for unsteady adjoint
+        self.dRdWTPCUnsteady = None
 
         if self.getOption("tensorflow")["active"]:
             TensorFlowHelper.options = self.getOption("tensorflow")
@@ -1016,62 +1018,6 @@ def _initializeAdjTotalDeriv(self):
 
         return adjTotalDeriv
 
-    def _initializeTimeAccurateAdjointVectors(self):
-        """
-        Initialize the dRdWTPsi vectors for time accurate adjoint.
-        Here we need to initialize current time step and two previous
-        time steps 0 and 00 for both state and residuals. This is
-        because the backward ddt scheme depends on U, U0, and U00
-        """
-        if self.getOption("unsteadyAdjoint")["mode"] == "timeAccurate":
-            objFuncDict = self.getOption("objFunc")
-            wSize = self.solver.getNLocalAdjointStates()
-            self.dRdW0TPsi = {}
-            self.dRdW00TPsi = {}
-            self.dR0dW0TPsi = {}
-            self.dR0dW00TPsi = {}
-            self.dR00dW0TPsi = {}
-            self.dR00dW00TPsi = {}
-            for objFuncName in objFuncDict:
-                if objFuncName in self.objFuncNames4Adj:
-                    vecA = PETSc.Vec().create(PETSc.COMM_WORLD)
-                    vecA.setSizes((wSize, PETSc.DECIDE), bsize=1)
-                    vecA.setFromOptions()
-                    vecA.zeroEntries()
-                    self.dRdW0TPsi[objFuncName] = vecA
-
-                    vecB = vecA.duplicate()
-                    vecB.zeroEntries()
-                    self.dRdW00TPsi[objFuncName] = vecB
-
-                    vecC = vecA.duplicate()
-                    vecC.zeroEntries()
-                    self.dR0dW0TPsi[objFuncName] = vecC
-
-                    vecD = vecA.duplicate()
-                    vecD.zeroEntries()
-                    self.dR0dW00TPsi[objFuncName] = vecD
-
-                    vecE = vecA.duplicate()
-                    vecE.zeroEntries()
-                    self.dR00dW0TPsi[objFuncName] = vecE
-
-                    vecF = vecA.duplicate()
-                    vecF.zeroEntries()
-                    self.dR00dW00TPsi[objFuncName] = vecF
-
-    def zeroTimeAccurateAdjointVectors(self):
-        if self.getOption("unsteadyAdjoint")["mode"] == "timeAccurate":
-            objFuncDict = self.getOption("objFunc")
-            for objFuncName in objFuncDict:
-                if objFuncName in self.objFuncNames4Adj:
-                    self.dRdW0TPsi[objFuncName].zeroEntries()
-                    self.dRdW00TPsi[objFuncName].zeroEntries()
-                    self.dR0dW0TPsi[objFuncName].zeroEntries()
-                    self.dR0dW00TPsi[objFuncName].zeroEntries()
-                    self.dR00dW0TPsi[objFuncName].zeroEntries()
-                    self.dR00dW00TPsi[objFuncName].zeroEntries()
-
     def _calcObjFuncNames4Adj(self):
         """
         Compute the objective function names for which we solve the adjoint equation
@@ -2337,6 +2283,9 @@ def solveAdjointUnsteady(self):
         self.setOption("runStatus", "solveAdjoint")
         self.updateDAOption()
 
+        PCMatPrecompute = self.getOption("unsteadyAdjoint")["PCMatPrecomputeInterval"]
+        PCMatUpdate = self.getOption("unsteadyAdjoint")["PCMatUpdateInterval"]
+
         self.adjointFail = 0
 
         # update the state to self.solver and self.solverAD
@@ -2370,13 +2319,48 @@ def solveAdjointUnsteady(self):
         # init dRdWTMF
         self.solverAD.initializedRdWTMatrixFree(self.xvVec, self.wVec)
 
-        # calc the preconditioner mat
-        self.dRdWTPC = PETSc.Mat().create(PETSc.COMM_WORLD)
-        self.solver.calcdRdWT(self.xvVec, self.wVec, 1, self.dRdWTPC)
-
-        # Initialize the KSP object
+        # precompute the KSP preconditioner Mat and save them to the self.dRdWTPCUnsteady dict
+        if self.dRdWTPCUnsteady is None:
+
+            self.dRdWTPCUnsteady = {}
+
+            # always calculate the PC mat for the endTime
+            self.solver.setTime(endTime, endTimeIndex)
+            self.solverAD.setTime(endTime, endTimeIndex)
+            # now we can read the variables
+            self.readStateVars(endTime, deltaT)
+            # calc the preconditioner mat for endTime
+            Info("Pre-Computing preconditiner mat for t = %f" % endTime)
+            dRdWTPC = PETSc.Mat().create(PETSc.COMM_WORLD)
+            self.solver.calcdRdWT(self.xvVec, self.wVec, 1, dRdWTPC)
+            # always update the PC mat values using OpenFOAM's fvMatrix
+            self.solver.calcPCMatWithFvMatrix(dRdWTPC)
+            self.dRdWTPCUnsteady[str(endTime)] = dRdWTPC
+
+            # if we define some extra PCMat in PCMatPrecomputeInterval, calculate them here
+            # and set them to the self.dRdWTPCUnsteady dict
+            if objFuncEndTimeIndex > PCMatPrecompute:
+                for timeIndex in range(objFuncEndTimeIndex - 1, 0, -1):
+                    if timeIndex % PCMatPrecompute == 0:
+                        t = timeIndex * deltaT
+                        Info("Pre-Computing preconditiner mat for t = %f" % t)
+                        # read the latest solution
+                        self.solver.setTime(t, timeIndex)
+                        self.solverAD.setTime(t, timeIndex)
+                        # now we can read the variables
+                        self.readStateVars(t, deltaT)
+
+                        # calc the preconditioner mat
+                        dRdWTPC = PETSc.Mat().create(PETSc.COMM_WORLD)
+                        self.solver.calcdRdWT(self.xvVec, self.wVec, 1, dRdWTPC)
+                        # always update the PC mat values using OpenFOAM's fvMatrix
+                        self.solver.calcPCMatWithFvMatrix(dRdWTPC)
+                        self.dRdWTPCUnsteady[str(t)] = dRdWTPC
+
+        # Initialize the KSP object using the PCMat from the endTime
+        PCMat = self.dRdWTPCUnsteady[str(endTime)]
         ksp = PETSc.KSP().create(PETSc.COMM_WORLD)
-        self.solverAD.createMLRKSPMatrixFree(self.dRdWTPC, ksp)
+        self.solverAD.createMLRKSPMatrixFree(PCMat, ksp)
 
         objFuncDict = self.getOption("objFunc")
         designVarDict = self.getOption("designVar")
@@ -2443,6 +2427,16 @@ def solveAdjointUnsteady(self):
                     else:
                         raise Error("ddtSchemeOrder not valid!" % ddtSchemeOrder)
 
+                    # check if we need to update the PC Mat vals or use the pre-computed PC matrix
+                    if str(timeVal) in list(self.dRdWTPCUnsteady.keys()):
+                        Info("Using pre-computed KSP PC mat for %f" % timeVal)
+                        PCMat = self.dRdWTPCUnsteady[str(timeVal)]
+                        self.solverAD.updateKSPPCMat(PCMat, ksp)
+                    if n % PCMatUpdate == 0 and n < endTimeIndex:
+                        # udpate part of the PC mat
+                        Info("Updating dRdWTPC mat value using OF fvMatrix")
+                        self.solver.calcPCMatWithFvMatrix(PCMat)
+
                     # now solve the adjoint eqn
                     self.adjointFail = self.solverAD.solveLinearEqn(ksp, dFdW, self.adjVectors[objFuncName])
 
diff --git a/src/adjoint/DAModel/DATurbulenceModel/DASpalartAllmarasFv3.C b/src/adjoint/DAModel/DATurbulenceModel/DASpalartAllmarasFv3.C
index 9c22bd47..84feda6b 100755
--- a/src/adjoint/DAModel/DATurbulenceModel/DASpalartAllmarasFv3.C
+++ b/src/adjoint/DAModel/DATurbulenceModel/DASpalartAllmarasFv3.C
@@ -695,6 +695,47 @@ void DASpalartAllmarasFv3::calcLduResidualTurb(volScalarField& nuTildaRes)
     // Below is not necessary, but it doesn't hurt
     nuTildaRes.correctBoundaryConditions();
 }
+
+void DASpalartAllmarasFv3::getFvMatrixFields(
+    const word varName,
+    scalarField& diag,
+    scalarField& upper,
+    scalarField& lower)
+{
+    /* 
+    Description:
+        return the diag(), upper(), and lower() scalarFields from the turbulence model's fvMatrix
+        this will be use to compute the preconditioner matrix
+    */
+
+    if (varName != "nuTilda")
+    {
+        FatalErrorIn(
+            "varName not valid. It has to be nuTilda")
+            << exit(FatalError);
+    }
+
+    const volScalarField chi(this->chi());
+    const volScalarField fv1(this->fv1(chi));
+
+    const volScalarField Stilda(
+        this->fv3(chi, fv1) * ::sqrt(2.0) * mag(skew(fvc::grad(U_)))
+        + this->fv2(chi, fv1) * nuTilda_ / sqr(kappa_ * y_));
+
+    fvScalarMatrix nuTildaEqn(
+        fvm::ddt(phase_, rho_, nuTilda_)
+            + fvm::div(phaseRhoPhi_, nuTilda_, "div(pc)")
+            - fvm::laplacian(phase_ * rho_ * DnuTildaEff(), nuTilda_)
+            - Cb2_ / sigmaNut_ * phase_ * rho_ * magSqr(fvc::grad(nuTilda_))
+        == Cb1_ * phase_ * rho_ * Stilda * nuTilda_ * betaFI_
+            - fvm::Sp(Cw1_ * phase_ * rho_ * fw(Stilda) * nuTilda_ / sqr(y_), nuTilda_));
+    
+    nuTildaEqn.relax();
+
+    diag = nuTildaEqn.D();
+    upper = nuTildaEqn.upper();
+    lower = nuTildaEqn.lower();
+}
 // * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * //
 
 } // End namespace Foam
diff --git a/src/adjoint/DAModel/DATurbulenceModel/DASpalartAllmarasFv3.H b/src/adjoint/DAModel/DATurbulenceModel/DASpalartAllmarasFv3.H
index 2cb76a87..fdba142b 100755
--- a/src/adjoint/DAModel/DATurbulenceModel/DASpalartAllmarasFv3.H
+++ b/src/adjoint/DAModel/DATurbulenceModel/DASpalartAllmarasFv3.H
@@ -161,6 +161,13 @@ public:
 
     /// calculate the turbulence residual using LDU matrix
     virtual void calcLduResidualTurb(volScalarField& nuTildaRes);
+
+    /// return the diag(), upper(), and lower() scalarFields from the turbulence model's fvMatrix
+    virtual void getFvMatrixFields(
+        const word varName,
+        scalarField& diag,
+        scalarField& upper,
+        scalarField& lower);
 };
 
 // * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * //
diff --git a/src/adjoint/DAModel/DATurbulenceModel/DATurbulenceModel.C b/src/adjoint/DAModel/DATurbulenceModel/DATurbulenceModel.C
index e229c043..2ed04794 100755
--- a/src/adjoint/DAModel/DATurbulenceModel/DATurbulenceModel.C
+++ b/src/adjoint/DAModel/DATurbulenceModel/DATurbulenceModel.C
@@ -542,6 +542,18 @@ void DATurbulenceModel::rhsSolvePseudoNuTildaEqn(const volScalarField& nuTildaSo
         << abort(FatalError);
 }
 
+/// return the diag(), upper(), and lower() scalarFields from the turbulence model's fvMatrix
+void DATurbulenceModel::getFvMatrixFields(
+    const word varName,
+    scalarField& diag,
+    scalarField& upper,
+    scalarField& lower)
+{
+    FatalErrorIn("DATurbulenceModel::getFvMatrixFields")
+        << "Child class not implemented!"
+        << abort(FatalError);
+}
+
 // * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * //
 
 } // End namespace Foam
diff --git a/src/adjoint/DAModel/DATurbulenceModel/DATurbulenceModel.H b/src/adjoint/DAModel/DATurbulenceModel/DATurbulenceModel.H
index 497da80a..dc936934 100755
--- a/src/adjoint/DAModel/DATurbulenceModel/DATurbulenceModel.H
+++ b/src/adjoint/DAModel/DATurbulenceModel/DATurbulenceModel.H
@@ -296,6 +296,13 @@ public:
 
     /// calculate the turbulence residual using LDU matrix
     virtual void calcLduResidualTurb(volScalarField& nuTildaRes);
+
+    /// return the diag(), upper(), and lower() scalarFields from the turbulence model's fvMatrix
+    virtual void getFvMatrixFields(
+        const word varName,
+        scalarField& diag,
+        scalarField& upper,
+        scalarField& lower);
 };
 
 // * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * //
diff --git a/src/adjoint/DAResidual/DAResidual.C b/src/adjoint/DAResidual/DAResidual.C
index e0d7b0ff..bd09a572 100755
--- a/src/adjoint/DAResidual/DAResidual.C
+++ b/src/adjoint/DAResidual/DAResidual.C
@@ -155,6 +155,13 @@ void DAResidual::masterFunction(
     }
 }
 
+void DAResidual::calcPCMatWithFvMatrix(Mat PCMat)
+{
+    FatalErrorIn("DAResidual::calcPCMatWithFvMatrix")
+        << "Child class not implemented!"
+        << abort(FatalError);
+}
+
 // * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * //
 
 } // End namespace Foam
diff --git a/src/adjoint/DAResidual/DAResidual.H b/src/adjoint/DAResidual/DAResidual.H
index 1d84b285..9246dbc7 100755
--- a/src/adjoint/DAResidual/DAResidual.H
+++ b/src/adjoint/DAResidual/DAResidual.H
@@ -120,6 +120,9 @@ public:
         const Vec xvVec,
         const Vec wVec,
         Vec resVec);
+    
+    /// calculating the adjoint preconditioner matrix using fvMatrix
+    virtual void calcPCMatWithFvMatrix(Mat PCMat);
 };
 
 // * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * //
diff --git a/src/adjoint/DAResidual/DAResidualPimpleFoam.C b/src/adjoint/DAResidual/DAResidualPimpleFoam.C
index 8a748c34..202c3d24 100755
--- a/src/adjoint/DAResidual/DAResidualPimpleFoam.C
+++ b/src/adjoint/DAResidual/DAResidualPimpleFoam.C
@@ -172,6 +172,199 @@ void DAResidualPimpleFoam::calcResiduals(const dictionary& options)
     normalizePhiResiduals(phiRes);
 }
 
+void DAResidualPimpleFoam::calcPCMatWithFvMatrix(Mat PCMat)
+{
+    /* 
+    Description:
+        Calculate the diagonal block of the preconditioner matrix dRdWTPC using the fvMatrix
+    */
+
+    const labelUList& owner = mesh_.owner();
+    const labelUList& neighbour = mesh_.neighbour();
+
+    PetscScalar val;
+
+    dictionary normStateDict = daOption_.getAllOptions().subDict("normalizeStates");
+    wordList normResDict = daOption_.getOption<wordList>("normalizeResiduals");
+
+    scalar UScaling = 1.0;
+    if (normStateDict.found("U"))
+    {
+        UScaling = normStateDict.getScalar("U");
+    }
+    scalar UResScaling = 1.0;
+
+    fvVectorMatrix UEqn(
+        fvm::ddt(U_)
+        + fvm::div(phi_, U_, "div(pc)")
+        + daTurb_.divDevReff(U_)
+        - fvSource_);
+
+    // set the val before relaxing UEqn!
+
+    // set diag
+    forAll(U_, cellI)
+    {
+        if (normResDict.found("URes"))
+        {
+            UResScaling = mesh_.V()[cellI];
+        }
+        for (label i = 0; i < 3; i++)
+        {
+            PetscInt rowI = daIndex_.getGlobalAdjointStateIndex("U", cellI, i);
+            PetscInt colI = rowI;
+            scalarField D = UEqn.D();
+            scalar val1 = D[cellI] * UScaling / UResScaling;
+            assignValueCheckAD(val, val1);
+            MatSetValues(PCMat, 1, &rowI, 1, &colI, &val, INSERT_VALUES);
+        }
+    }
+
+    // set lower/owner
+    for (label faceI = 0; faceI < daIndex_.nLocalInternalFaces; faceI++)
+    {
+        label ownerCellI = owner[faceI];
+        label neighbourCellI = neighbour[faceI];
+
+        if (normResDict.found("URes"))
+        {
+            UResScaling = mesh_.V()[neighbourCellI];
+        }
+
+        for (label i = 0; i < 3; i++)
+        {
+            PetscInt rowI = daIndex_.getGlobalAdjointStateIndex("U", neighbourCellI, i);
+            PetscInt colI = daIndex_.getGlobalAdjointStateIndex("U", ownerCellI, i);
+            scalar val1 = UEqn.lower()[faceI] * UScaling / UResScaling;
+            assignValueCheckAD(val, val1);
+            MatSetValues(PCMat, 1, &colI, 1, &rowI, &val, INSERT_VALUES);
+        }
+    }
+
+    // set upper/neighbour
+    for (label faceI = 0; faceI < daIndex_.nLocalInternalFaces; faceI++)
+    {
+        label ownerCellI = owner[faceI];
+        label neighbourCellI = neighbour[faceI];
+
+        if (normResDict.found("URes"))
+        {
+            UResScaling = mesh_.V()[ownerCellI];
+        }
+
+        for (label i = 0; i < 3; i++)
+        {
+            PetscInt rowI = daIndex_.getGlobalAdjointStateIndex("U", ownerCellI, i);
+            PetscInt colI = daIndex_.getGlobalAdjointStateIndex("U", neighbourCellI, i);
+            scalar val1 = UEqn.upper()[faceI] * UScaling / UResScaling;
+            assignValueCheckAD(val, val1);
+            MatSetValues(PCMat, 1, &colI, 1, &rowI, &val, INSERT_VALUES);
+        }
+    }
+
+    UEqn.relax();
+    label pRefCell = 0;
+    scalar pRefValue = 0.0;
+
+    volScalarField rAU(1.0 / UEqn.A());
+    autoPtr<volVectorField> HbyAPtr = nullptr;
+    label useConstrainHbyA = daOption_.getOption<label>("useConstrainHbyA");
+    if (useConstrainHbyA)
+    {
+        HbyAPtr.reset(new volVectorField(constrainHbyA(rAU * UEqn.H(), U_, p_)));
+    }
+    else
+    {
+        HbyAPtr.reset(new volVectorField("HbyA", U_));
+        HbyAPtr() = rAU * UEqn.H();
+    }
+    volVectorField& HbyA = HbyAPtr();
+
+    surfaceScalarField phiHbyA(
+        "phiHbyA",
+        fvc::flux(HbyA));
+
+    if (p_.needReference())
+    {
+        adjustPhi(phiHbyA, U_, p_);
+    }
+
+    tmp<volScalarField> rAtU(rAU);
+
+    if (pimple_.consistent())
+    {
+        rAtU = 1.0 / max(1.0 / rAU - UEqn.H1(), 0.1 / rAU);
+        phiHbyA +=
+            fvc::interpolate(rAtU() - rAU) * fvc::snGrad(p_) * mesh_.magSf();
+        HbyA -= (rAU - rAtU()) * fvc::grad(p_);
+    }
+
+    fvScalarMatrix pEqn(
+        fvm::laplacian(rAtU(), p_)
+        == fvc::div(phiHbyA));
+
+    pEqn.setReference(pRefCell, pRefValue);
+
+    // ********* p
+    scalar pScaling = 1.0;
+    if (normStateDict.found("p"))
+    {
+        pScaling = normStateDict.getScalar("p");
+    }
+    scalar pResScaling = 1.0;
+    // set diag
+    forAll(p_, cellI)
+    {
+        if (normResDict.found("pRes"))
+        {
+            pResScaling = mesh_.V()[cellI];
+        }
+
+        PetscInt rowI = daIndex_.getGlobalAdjointStateIndex("p", cellI);
+        PetscInt colI = rowI;
+        scalarField D = pEqn.D();
+        scalar val1 = D[cellI] * pScaling / pResScaling;
+        assignValueCheckAD(val, val1);
+        MatSetValues(PCMat, 1, &rowI, 1, &colI, &val, INSERT_VALUES);
+    }
+
+    // set lower/owner
+    for (label faceI = 0; faceI < daIndex_.nLocalInternalFaces; faceI++)
+    {
+        label ownerCellI = owner[faceI];
+        label neighbourCellI = neighbour[faceI];
+
+        if (normResDict.found("pRes"))
+        {
+            pResScaling = mesh_.V()[neighbourCellI];
+        }
+
+        PetscInt rowI = daIndex_.getGlobalAdjointStateIndex("p", neighbourCellI);
+        PetscInt colI = daIndex_.getGlobalAdjointStateIndex("p", ownerCellI);
+        scalar val1 = pEqn.lower()[faceI] * pScaling / pResScaling;
+        assignValueCheckAD(val, val1);
+        MatSetValues(PCMat, 1, &colI, 1, &rowI, &val, INSERT_VALUES);
+    }
+
+    // set upper/neighbour
+    for (label faceI = 0; faceI < daIndex_.nLocalInternalFaces; faceI++)
+    {
+        label ownerCellI = owner[faceI];
+        label neighbourCellI = neighbour[faceI];
+
+        if (normResDict.found("pRes"))
+        {
+            pResScaling = mesh_.V()[ownerCellI];
+        }
+
+        PetscInt rowI = daIndex_.getGlobalAdjointStateIndex("p", ownerCellI);
+        PetscInt colI = daIndex_.getGlobalAdjointStateIndex("p", neighbourCellI);
+        scalar val1 = pEqn.upper()[faceI] * pScaling / pResScaling;
+        assignValueCheckAD(val, val1);
+        MatSetValues(PCMat, 1, &colI, 1, &rowI, &val, INSERT_VALUES);
+    }
+}
+
 void DAResidualPimpleFoam::updateIntermediateVariables()
 {
     /* 
diff --git a/src/adjoint/DAResidual/DAResidualPimpleFoam.H b/src/adjoint/DAResidual/DAResidualPimpleFoam.H
index 5852f020..af6bafa5 100755
--- a/src/adjoint/DAResidual/DAResidualPimpleFoam.H
+++ b/src/adjoint/DAResidual/DAResidualPimpleFoam.H
@@ -86,6 +86,8 @@ public:
 
     /// update the boundary condition for all the states in the selected solver
     virtual void correctBoundaryConditions();
+
+    virtual void calcPCMatWithFvMatrix(Mat PCMat);
 };
 
 // * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * //
diff --git a/src/adjoint/DASolver/DASolver.C b/src/adjoint/DASolver/DASolver.C
index 9b71dc9d..659b9eeb 100644
--- a/src/adjoint/DASolver/DASolver.C
+++ b/src/adjoint/DASolver/DASolver.C
@@ -202,7 +202,7 @@ void DASolver::calcUnsteadyObjFuncs()
             {
                 label n = timeIndex - startTimeIndex + 1;
                 scalar objFuncVal = this->getObjFuncValue(objFuncName);
-                unsteadyObjFuncs_[objFuncName] = (unsteadyObjFuncs_[objFuncName] * (n - 1) + objFuncVal ) / n;
+                unsteadyObjFuncs_[objFuncName] = (unsteadyObjFuncs_[objFuncName] * (n - 1) + objFuncVal) / n;
             }
         }
         else
@@ -5079,6 +5079,17 @@ void DASolver::createMLRKSP(
     daLinearEqnPtr_->createMLRKSP(jacMat, jacPCMat, ksp);
 }
 
+void DASolver::updateKSPPCMat(
+    Mat PCMat,
+    KSP ksp)
+{
+    /*
+    Description:
+        Update the preconditioner matrix for the ksp object
+    */
+    KSPSetOperators(ksp, dRdWTMF_, PCMat);
+}
+
 void DASolver::createMLRKSPMatrixFree(
     const Mat jacPCMat,
     KSP ksp)
@@ -8552,6 +8563,106 @@ void DASolver::setTimeInstanceField(const label instanceI)
     }
 }
 
+
+void DASolver::calcPCMatWithFvMatrix(Mat PCMat)
+{
+    /*
+    Description:
+        calculate the PC mat using fvMatrix. Here we only calculate the block diagonal components, 
+        e.g., dR_U/dU, dR_p/dp, etc.
+    */
+
+    //DAUtility::writeMatrixASCII(PCMat, "MatOrig");
+
+    // MatZeroEntries(PCMat);
+
+    // non turbulence variables
+    daResidualPtr_->calcPCMatWithFvMatrix(PCMat);
+
+    // turbulence variables
+    DATurbulenceModel& daTurb = const_cast<DATurbulenceModel&>(daModelPtr_->getDATurbulenceModel());
+    const labelUList& owner = meshPtr_->owner();
+    const labelUList& neighbour = meshPtr_->neighbour();
+
+    PetscScalar val;
+
+    dictionary normStateDict = daOptionPtr_->getAllOptions().subDict("normalizeStates");
+    wordList normResDict = daOptionPtr_->getOption<wordList>("normalizeResiduals");
+    forAll(stateInfo_["modelStates"], idxI)
+    {
+        const word stateName = stateInfo_["modelStates"][idxI];
+        const word resName = stateName + "Res";
+        label nCells = meshPtr_->nCells();
+        label nInternalFaces = daIndexPtr_->nLocalInternalFaces;
+        scalarField D(nCells, 0.0);
+        scalarField upper(nInternalFaces, 0.0);
+        scalarField lower(nInternalFaces, 0.0);
+        daTurb.getFvMatrixFields(stateName, D, upper, lower);
+
+        scalar stateScaling = 1.0;
+        if (normStateDict.found(stateName))
+        {
+            stateScaling = normStateDict.getScalar(stateName);
+        }
+        scalar resScaling = 1.0;
+        // set diag
+        forAll(meshPtr_->cells(), cellI)
+        {
+            if (normResDict.found(resName))
+            {
+                resScaling = meshPtr_->V()[cellI];
+            }
+
+            PetscInt rowI = daIndexPtr_->getGlobalAdjointStateIndex(stateName, cellI);
+            PetscInt colI = rowI;
+            scalar val1 = D[cellI] * stateScaling / resScaling;
+            assignValueCheckAD(val, val1);
+            MatSetValues(PCMat, 1, &rowI, 1, &colI, &val, INSERT_VALUES);
+        }
+
+        // set lower/owner
+        for (label faceI = 0; faceI < daIndexPtr_->nLocalInternalFaces; faceI++)
+        {
+            label ownerCellI = owner[faceI];
+            label neighbourCellI = neighbour[faceI];
+
+            if (normResDict.found(resName))
+            {
+                resScaling = meshPtr_->V()[neighbourCellI];
+            }
+
+            PetscInt rowI = daIndexPtr_->getGlobalAdjointStateIndex(stateName, neighbourCellI);
+            PetscInt colI = daIndexPtr_->getGlobalAdjointStateIndex(stateName, ownerCellI);
+            scalar val1 = lower[faceI] * stateScaling / resScaling;
+            assignValueCheckAD(val, val1);
+            MatSetValues(PCMat, 1, &colI, 1, &rowI, &val, INSERT_VALUES);
+        }
+
+        // set upper/neighbour
+        for (label faceI = 0; faceI < daIndexPtr_->nLocalInternalFaces; faceI++)
+        {
+            label ownerCellI = owner[faceI];
+            label neighbourCellI = neighbour[faceI];
+
+            if (normResDict.found(resName))
+            {
+                resScaling = meshPtr_->V()[ownerCellI];
+            }
+
+            PetscInt rowI = daIndexPtr_->getGlobalAdjointStateIndex(stateName, ownerCellI);
+            PetscInt colI = daIndexPtr_->getGlobalAdjointStateIndex(stateName, neighbourCellI);
+            scalar val1 = upper[faceI] * stateScaling / resScaling;
+            assignValueCheckAD(val, val1);
+            MatSetValues(PCMat, 1, &colI, 1, &rowI, &val, INSERT_VALUES);
+        }
+    }
+
+    MatAssemblyBegin(PCMat, MAT_FINAL_ASSEMBLY);
+    MatAssemblyEnd(PCMat, MAT_FINAL_ASSEMBLY);
+
+    //DAUtility::writeMatrixASCII(PCMat, "MatNew");
+}
+
 void DASolver::readStateVars(
     scalar timeVal,
     label oldTimeLevel)
@@ -8602,12 +8713,10 @@ void DASolver::readStateVars(
             if (oldTimeLevel == 0)
             {
                 state = stateRead;
-                state.correctBoundaryConditions();
             }
             else if (oldTimeLevel == 1)
             {
                 state.oldTime() = stateRead;
-                state.oldTime().correctBoundaryConditions();
             }
             else if (oldTimeLevel == 2)
             {
@@ -8627,7 +8736,6 @@ void DASolver::readStateVars(
                 {
                     state.oldTime().oldTime() = stateRead;
                 }
-                state.oldTime().oldTime().correctBoundaryConditions();
             }
             else
             {
@@ -8659,12 +8767,10 @@ void DASolver::readStateVars(
             if (oldTimeLevel == 0)
             {
                 state = stateRead;
-                state.correctBoundaryConditions();
             }
             else if (oldTimeLevel == 1)
             {
                 state.oldTime() = stateRead;
-                state.oldTime().correctBoundaryConditions();
             }
             else if (oldTimeLevel == 2)
             {
@@ -8684,7 +8790,6 @@ void DASolver::readStateVars(
                 {
                     state.oldTime().oldTime() = stateRead;
                 }
-                state.oldTime().oldTime().correctBoundaryConditions();
             }
             else
             {
@@ -8716,12 +8821,10 @@ void DASolver::readStateVars(
             if (oldTimeLevel == 0)
             {
                 state = stateRead;
-                state.correctBoundaryConditions();
             }
             else if (oldTimeLevel == 1)
             {
                 state.oldTime() = stateRead;
-                state.oldTime().correctBoundaryConditions();
             }
             else if (oldTimeLevel == 2)
             {
@@ -8741,7 +8844,6 @@ void DASolver::readStateVars(
                 {
                     state.oldTime().oldTime() = stateRead;
                 }
-                state.oldTime().oldTime().correctBoundaryConditions();
             }
             else
             {
@@ -8802,6 +8904,9 @@ void DASolver::readStateVars(
             }
         }
     }
+
+    // update the BC and intermediate variables. This is important, e.g., for turbulent cases
+    this->updateStateBoundaryConditions();
 }
 
 void DASolver::setTimeInstanceVar(
diff --git a/src/adjoint/DASolver/DASolver.H b/src/adjoint/DASolver/DASolver.H
index 5f594de5..c9eeac42 100644
--- a/src/adjoint/DASolver/DASolver.H
+++ b/src/adjoint/DASolver/DASolver.H
@@ -343,6 +343,11 @@ public:
         const Mat jacPCMat,
         KSP ksp);
 
+    /// Update the preconditioner matrix for the ksp object
+    void updateKSPPCMat(
+        Mat PCMat,
+        KSP ksp);
+
     /// solve the linear equation given a ksp and right-hand-side vector
     label solveLinearEqn(
         const KSP ksp,
@@ -1112,6 +1117,9 @@ public:
         scalar timeVal,
         label oldTimeLevel = 0);
 
+    /// calculate the PC mat using fvMatrix
+    void calcPCMatWithFvMatrix(Mat PCMat);
+
     /// initialize tensorflow functions and interfaces for callback
     void initTensorFlowFuncs(
         pyComputeInterface computeInterface,
diff --git a/src/pyDASolvers/DASolvers.H b/src/pyDASolvers/DASolvers.H
index d4c59a6d..e6b627f4 100755
--- a/src/pyDASolvers/DASolvers.H
+++ b/src/pyDASolvers/DASolvers.H
@@ -194,6 +194,14 @@ public:
         DASolverPtr_->createMLRKSP(jacMat, jacPCMat, ksp);
     }
 
+    /// Update the preconditioner matrix for the ksp object
+    void updateKSPPCMat(
+        Mat PCMat,
+        KSP ksp)
+    {
+        DASolverPtr_->updateKSPPCMat(PCMat, ksp);
+    }
+
     /// create a multi-level, Richardson KSP object with matrix-free Jacobians
     void createMLRKSPMatrixFree(
         const Mat jacPCMat,
@@ -703,6 +711,12 @@ public:
         DASolverPtr_->readStateVars(timeVal, timeLevel);
     }
 
+    /// calculate the PC mat using fvMatrix
+    void calcPCMatWithFvMatrix(Mat PCMat)
+    {
+        DASolverPtr_->calcPCMatWithFvMatrix(PCMat);
+    }
+
     /// setTime for OF fields
     void setTime(scalar time, label timeIndex)
     {
@@ -742,7 +756,7 @@ public:
         return DASolverPtr_->getUnsteadyObjFuncStartTimeIndex();
 #endif
         FatalErrorIn("getUnsteadyObjFuncStartTimeIndex") << "should not call this func in AD mode!"
-                                                    << abort(FatalError);
+                                                         << abort(FatalError);
     }
 
     /// get the unsteady adjoint end time
@@ -752,7 +766,7 @@ public:
         return DASolverPtr_->getUnsteadyObjFuncEndTimeIndex();
 #endif
         FatalErrorIn("getUnsteadyObjFuncEndTimeIndex") << "should not call this func in AD mode!"
-                                                    << abort(FatalError);
+                                                       << abort(FatalError);
     }
 
     /// return the value of objective function at the given time instance and name
diff --git a/src/pyDASolvers/pyDASolvers.pyx b/src/pyDASolvers/pyDASolvers.pyx
index 37191e59..c01fbb9d 100755
--- a/src/pyDASolvers/pyDASolvers.pyx
+++ b/src/pyDASolvers/pyDASolvers.pyx
@@ -51,6 +51,7 @@ cdef extern from "DASolvers.H" namespace "Foam":
         void calcdFdWAD(PetscVec, PetscVec, char *, PetscVec)
         void createMLRKSP(PetscMat, PetscMat, PetscKSP)
         void createMLRKSPMatrixFree(PetscMat, PetscKSP)
+        void updateKSPPCMat(PetscMat, PetscKSP)
         void solveLinearEqn(PetscKSP, PetscVec, PetscVec)
         void calcdRdBC(PetscVec, PetscVec, char *, PetscMat)
         void calcdFdBC(PetscVec, PetscVec, char *, char *, PetscVec)
@@ -133,6 +134,7 @@ cdef extern from "DASolvers.H" namespace "Foam":
         int runFPAdj(PetscVec, PetscVec, PetscVec, PetscVec)
         void initTensorFlowFuncs(pyComputeInterface, void *, pyJacVecProdInterface, void *)
         void readStateVars(double, int)
+        void calcPCMatWithFvMatrix(PetscMat)
         double getEndTime()
         double getDeltaT()
         void setTime(double, int)
@@ -209,6 +211,9 @@ cdef class pyDASolvers:
     def createMLRKSPMatrixFree(self, Mat jacPCMat, KSP myKSP):
         self._thisptr.createMLRKSPMatrixFree(jacPCMat.mat, myKSP.ksp)
     
+    def updateKSPPCMat(self, Mat PCMat, KSP myKSP):
+        self._thisptr.updateKSPPCMat(PCMat.mat, myKSP.ksp)
+    
     def solveLinearEqn(self, KSP myKSP, Vec rhsVec, Vec solVec):
         self._thisptr.solveLinearEqn(myKSP.ksp, rhsVec.vec, solVec.vec)
 
@@ -521,6 +526,9 @@ cdef class pyDASolvers:
     def readStateVars(self, timeVal, timeLevel):
         self._thisptr.readStateVars(timeVal, timeLevel)
     
+    def calcPCMatWithFvMatrix(self, Mat PCMat):
+        self._thisptr.calcPCMatWithFvMatrix(PCMat.mat)
+    
     def setTime(self, time, timeIndex):
         self._thisptr.setTime(time, timeIndex)