Enabled writing adjoint fields for post-processing

dafoam/mphys/mphys_dafoam.py

@@ -1044,6 +1044,13 @@ def solve_linear(self, d_outputs, d_residuals, mode):
 
                 # actually solving the adjoint linear equation using Petsc
                 fail = DASolver.solverAD.solveLinearEqn(DASolver.ksp, dFdW, self.psi)
+
+                # optionally write the adjoint vector as OpenFOAM field format for post-processing
+                # update the obj func name for solve_linear later
+                solveLinearObjFuncName = DASolver.getOption("solveLinearObjFuncName")
+                psi_array = DASolver.vec2Array(self.psi)
+                self.DASolver.writeAdjointFields(solveLinearObjFuncName, float(solutionTime), psi_array)
+
             elif adjEqnSolMethod == "fixedPoint":
                 solutionTime, renamed = DASolver.renameSolution(self.solution_counter)
                 if renamed:
@@ -1356,6 +1363,10 @@ def compute_jacvec_product(self, inputs, d_inputs, d_outputs, mode):
 
         if self.comm.rank == 0:
             print("Computing partials for ", list(funcsBar.keys()))
+
+        # update the obj func name for solve_linear later
+        DASolver.setOption("solveLinearObjFuncName", list(funcsBar.keys())[0])
+        DASolver.updateDAOption()
 
         # loop over all d_inputs keys and compute the partials accordingly
         for objFuncName in list(funcsBar.keys()):

dafoam/pyDAFoam.py

@@ -572,6 +572,10 @@ def __init__(self):
         ## used internally, so users should never change this option in the Python layer.
         self.runStatus = "None"
 
+        ## The current objective function name. This variable will be reset in DAFoamFunctions
+        ## Then, we know which objective function is used in solve_linear in DAFoamSolver
+        self.solveLinearObjFuncName = "None"
+
         ## Whether to print all options defined in pyDAFoam to screen before optimization.
         self.printPYDAFOAMOptions = False
 
@@ -706,6 +710,9 @@ def __init__(self):
         ## Whether to write deformed constraints to disk during optimization, i.e., DVCon.writeTecplot
         self.writeDeformedConstraints = False
 
+        ## Whether to write adjoint variables in OpenFOAM field format for post-processing
+        self.writeAdjointFields = False
+
         ## The max number of correctBoundaryConditions calls in the updateOFField function.
         self.maxCorrectBCCalls = 10
 
@@ -743,9 +750,9 @@ def __init__(self):
         ## tensorflow related functions
         self.tensorflow = {
             "active": False,
-            #"model1": {
+            # "model1": {
             #    "predictBatchSize": 1000
-            #}
+            # }
         }
 
 
@@ -927,9 +934,13 @@ def __init__(self, comm=None, options=None):
             TensorFlowHelper.options = self.getOption("tensorflow")
             TensorFlowHelper.initialize()
             # pass this helper function to the C++ layer
-            self.solver.initTensorFlowFuncs(TensorFlowHelper.predict, TensorFlowHelper.calcJacVecProd, TensorFlowHelper.setModelName)
+            self.solver.initTensorFlowFuncs(
+                TensorFlowHelper.predict, TensorFlowHelper.calcJacVecProd, TensorFlowHelper.setModelName
+            )
             if self.getOption("useAD")["mode"] in ["forward", "reverse"]:
-                self.solverAD.initTensorFlowFuncs(TensorFlowHelper.predict, TensorFlowHelper.calcJacVecProd, TensorFlowHelper.setModelName)
+                self.solverAD.initTensorFlowFuncs(
+                    TensorFlowHelper.predict, TensorFlowHelper.calcJacVecProd, TensorFlowHelper.setModelName
+                )
 
         Info("pyDAFoam initialization done!")
 
@@ -1329,6 +1340,14 @@ def writeDeformedFFDs(self, counter=None):
             else:
                 self.DVGeo.writeTecplot("deformedFFD_%d.dat" % counter)
 
+    def writeAdjointFields(self, objFunc, writeTime, psi):
+        """
+        Write the adjoint variables in OpenFOAM field format for post-processing
+        """
+
+        if self.getOption("writeAdjointFields"):
+            self.solver.writeAdjointFields(objFunc, writeTime, psi)
+
     def writeTotalDeriv(self, fileName, sens, evalFuncs):
         """
         Write the total derivatives history to files in the json format

src/adjoint/DASolver/DASolver.C

@@ -10384,6 +10384,98 @@ void DASolver::writeSensMapField(
     }
 }
 
+void DASolver::writeAdjointFields(
+    const word objFunc,
+    const double writeTime,
+    const double* psi)
+{
+    /*
+    Description:
+        write the adjoint variables to the disk as OpenFOAM variables so they can be viewed
+        in ParaView
+    
+    Inputs:
+        writeTime: solution time the fields will be saved to
+        psi: the adjoint vector array, computed in the Python layer
+    */
+
+    runTimePtr_->setTime(writeTime, 0);
+
+    forAll(stateInfo_["volVectorStates"], idxI)
+    {
+        const word stateName = stateInfo_["volVectorStates"][idxI];
+        const volVectorField& state = meshPtr_->thisDb().lookupObject<volVectorField>(stateName);
+        word varName = "adjoint_" + objFunc + "_" + stateName;
+        volVectorField adjointVar(varName, state);
+        forAll(state, cellI)
+        {
+            for (label i = 0; i < 3; i++)
+            {
+                label localIdx = daIndexPtr_->getLocalAdjointStateIndex(stateName, cellI, i);
+                adjointVar[cellI][i] = psi[localIdx];
+            }
+        }
+        adjointVar.correctBoundaryConditions();
+        adjointVar.write();
+    }
+
+    forAll(stateInfo_["volScalarStates"], idxI)
+    {
+        const word stateName = stateInfo_["volScalarStates"][idxI];
+        const volScalarField& state = meshPtr_->thisDb().lookupObject<volScalarField>(stateName);
+        word varName = "adjoint_" + objFunc + "_" + stateName;
+        volScalarField adjointVar(varName, state);
+        forAll(state, cellI)
+        {
+            label localIdx = daIndexPtr_->getLocalAdjointStateIndex(stateName, cellI);
+            adjointVar[cellI] = psi[localIdx];
+        }
+        adjointVar.correctBoundaryConditions();
+        adjointVar.write();
+    }
+
+    forAll(stateInfo_["modelStates"], idxI)
+    {
+        const word stateName = stateInfo_["modelStates"][idxI];
+        const volScalarField& state = meshPtr_->thisDb().lookupObject<volScalarField>(stateName);
+        word varName = "adjoint_" + objFunc + "_" + stateName;
+        volScalarField adjointVar(varName, state);
+        forAll(state, cellI)
+        {
+            label localIdx = daIndexPtr_->getLocalAdjointStateIndex(stateName, cellI);
+            adjointVar[cellI] = psi[localIdx];
+        }
+        adjointVar.correctBoundaryConditions();
+        adjointVar.write();
+    }
+
+    forAll(stateInfo_["surfaceScalarStates"], idxI)
+    {
+        const word stateName = stateInfo_["surfaceScalarStates"][idxI];
+        const surfaceScalarField& state = meshPtr_->thisDb().lookupObject<surfaceScalarField>(stateName);
+        word varName = "adjoint_" + objFunc + "_" + stateName;
+        surfaceScalarField adjointVar(varName, state);
+
+        forAll(meshPtr_->faces(), faceI)
+        {
+            label localIdx = daIndexPtr_->getLocalAdjointStateIndex(stateName, faceI);
+
+            if (faceI < daIndexPtr_->nLocalInternalFaces)
+            {
+                adjointVar[faceI] = psi[localIdx];
+            }
+            else
+            {
+                label relIdx = faceI - daIndexPtr_->nLocalInternalFaces;
+                label patchIdx = daIndexPtr_->bFacePatchI[relIdx];
+                label faceIdx = daIndexPtr_->bFaceFaceI[relIdx];
+                adjointVar.boundaryFieldRef()[patchIdx][faceIdx] = psi[localIdx];
+            }
+        }
+        adjointVar.write();
+    }
+}
+
 // * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * //
 
 } // End namespace Foam

src/adjoint/DASolver/DASolver.H

@@ -236,7 +236,6 @@ protected:
     /// the tolerance of the primal objective std, used if primalObjStdTol->active is True
     scalar primalObjStdTol_;
 
-
 public:
     /// Runtime type information
     TypeName("DASolver");
@@ -1327,6 +1326,12 @@ public:
 
     /// calculate the objective function's std
     void calcObjStd(Time& runTime);
+
+    /// write the adjoint variables for all states
+    void writeAdjointFields(
+        const word objFunc,
+        const double writeTime,
+        const double* psi);
 };
 
 // * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * //

src/pyDASolvers/DASolvers.H

@@ -991,6 +991,15 @@ public:
             computeInterface, compute, jacVecProdInterface, jacVecProd, setModelNameInterface, setModelName);
     }
 
+    /// write the adjoint variables for all states
+    void writeAdjointFields(
+        const word objFunc,
+        const double writeTime,
+        const double* psi)
+    {
+        DASolverPtr_->writeAdjointFields(objFunc, writeTime, psi);
+    }
+
     /// write the sensitivity map for all wall surfaces
     void writeSensMapSurface(
         const word name,

src/pyDASolvers/pyDASolvers.pyx

@@ -157,6 +157,7 @@ cdef extern from "DASolvers.H" namespace "Foam":
         void writeSensMapSurface(char *, double *, double *, int, double)
         void writeSensMapField(char *, double *, char *, double)
         double getLatestTime()
+        void writeAdjointFields(char *, double, double *)
 
 # create python wrappers that call cpp functions
 cdef class pyDASolvers:
@@ -706,3 +707,12 @@ cdef class pyDASolvers:
 
     def getLatestTime(self):
         return self._thisptr.getLatestTime()
+
+    def writeAdjointFields(self, objFunc, writeTime, np.ndarray[double, ndim=1, mode="c"] psi):
+        nAdjStates = self.getNLocalAdjointStates()
+
+        assert len(psi) == nAdjStates, "invalid array size!"
+
+        cdef double *psi_data = <double*>psi.data
+
+        return self._thisptr.writeAdjointFields(objFunc.encode(), writeTime, psi_data)