xref: /libCEED/examples/fluids/src/setupts.c (revision 7d4c6defbc4b3e425caecd0a3042b4105f6c35a1)
13d8e8822SJeremy L Thompson // Copyright (c) 2017-2022, Lawrence Livermore National Security, LLC and other CEED contributors.
23d8e8822SJeremy L Thompson // All Rights Reserved. See the top-level LICENSE and NOTICE files for details.
377841947SLeila Ghaffari //
43d8e8822SJeremy L Thompson // SPDX-License-Identifier: BSD-2-Clause
577841947SLeila Ghaffari //
63d8e8822SJeremy L Thompson // This file is part of CEED:  http://github.com/ceed
777841947SLeila Ghaffari 
877841947SLeila Ghaffari /// @file
977841947SLeila Ghaffari /// Time-stepping functions for Navier-Stokes example using PETSc
1077841947SLeila Ghaffari 
1149aac155SJeremy L Thompson #include <ceed.h>
1249aac155SJeremy L Thompson #include <petscdmplex.h>
1349aac155SJeremy L Thompson #include <petscts.h>
1449aac155SJeremy L Thompson 
1577841947SLeila Ghaffari #include "../navierstokes.h"
16ca69d878SAdeleke O. Bankole #include "../qfunctions/newtonian_state.h"
1777841947SLeila Ghaffari 
1877841947SLeila Ghaffari // Compute mass matrix for explicit scheme
192b730f8bSJeremy L Thompson PetscErrorCode ComputeLumpedMassMatrix(Ceed ceed, DM dm, CeedData ceed_data, Vec M) {
2077841947SLeila Ghaffari   CeedQFunction        qf_mass;
2177841947SLeila Ghaffari   CeedOperator         op_mass;
223221f4d3SJames Wright   OperatorApplyContext op_mass_ctx;
233221f4d3SJames Wright   Vec                  Ones_loc;
2477841947SLeila Ghaffari   CeedInt              num_comp_q, q_data_size;
2577841947SLeila Ghaffari 
26f17d818dSJames Wright   PetscFunctionBeginUser;
27a424bcd0SJames Wright   PetscCallCeed(ceed, CeedElemRestrictionGetNumComponents(ceed_data->elem_restr_q, &num_comp_q));
28a424bcd0SJames Wright   PetscCallCeed(ceed, CeedElemRestrictionGetNumComponents(ceed_data->elem_restr_qd_i, &q_data_size));
2977841947SLeila Ghaffari 
30ef080ff9SJames Wright   PetscCall(CreateMassQFunction(ceed, num_comp_q, q_data_size, &qf_mass));
31a424bcd0SJames Wright   PetscCallCeed(ceed, CeedOperatorCreate(ceed, qf_mass, NULL, NULL, &op_mass));
32a424bcd0SJames Wright   PetscCallCeed(ceed, CeedOperatorSetField(op_mass, "u", ceed_data->elem_restr_q, ceed_data->basis_q, CEED_VECTOR_ACTIVE));
33356036faSJeremy L Thompson   PetscCallCeed(ceed, CeedOperatorSetField(op_mass, "qdata", ceed_data->elem_restr_qd_i, CEED_BASIS_NONE, ceed_data->q_data));
34a424bcd0SJames Wright   PetscCallCeed(ceed, CeedOperatorSetField(op_mass, "v", ceed_data->elem_restr_q, ceed_data->basis_q, CEED_VECTOR_ACTIVE));
3577841947SLeila Ghaffari 
363221f4d3SJames Wright   PetscCall(OperatorApplyContextCreate(NULL, dm, ceed, op_mass, NULL, NULL, NULL, NULL, &op_mass_ctx));
3777841947SLeila Ghaffari 
383221f4d3SJames Wright   PetscCall(DMGetLocalVector(dm, &Ones_loc));
393221f4d3SJames Wright   PetscCall(VecSet(Ones_loc, 1));
403221f4d3SJames Wright   PetscCall(ApplyCeedOperatorLocalToGlobal(Ones_loc, M, op_mass_ctx));
4177841947SLeila Ghaffari 
4277841947SLeila Ghaffari   // Invert diagonally lumped mass vector for RHS function
432b730f8bSJeremy L Thompson   PetscCall(VecReciprocal(M));
4477841947SLeila Ghaffari 
4577841947SLeila Ghaffari   // Cleanup
463221f4d3SJames Wright   PetscCall(OperatorApplyContextDestroy(op_mass_ctx));
473221f4d3SJames Wright   PetscCall(DMRestoreLocalVector(dm, &Ones_loc));
48a424bcd0SJames Wright   PetscCallCeed(ceed, CeedQFunctionDestroy(&qf_mass));
49a424bcd0SJames Wright   PetscCallCeed(ceed, CeedOperatorDestroy(&op_mass));
50ee4ca9cbSJames Wright   PetscFunctionReturn(PETSC_SUCCESS);
5177841947SLeila Ghaffari }
5277841947SLeila Ghaffari 
53a0b9a424SJames Wright // Insert Boundary values if it's a new time
54a0b9a424SJames Wright PetscErrorCode UpdateBoundaryValues(User user, Vec Q_loc, PetscReal t) {
55a0b9a424SJames Wright   PetscFunctionBeginUser;
56a0b9a424SJames Wright   if (user->time_bc_set != t) {
57a0b9a424SJames Wright     PetscCall(DMPlexInsertBoundaryValues(user->dm, PETSC_TRUE, Q_loc, t, NULL, NULL, NULL));
58a0b9a424SJames Wright     user->time_bc_set = t;
59a0b9a424SJames Wright   }
60ee4ca9cbSJames Wright   PetscFunctionReturn(PETSC_SUCCESS);
61a0b9a424SJames Wright }
62a0b9a424SJames Wright 
6377841947SLeila Ghaffari // RHS (Explicit time-stepper) function setup
6477841947SLeila Ghaffari //   This is the RHS of the ODE, given as u_t = G(t,u)
6577841947SLeila Ghaffari //   This function takes in a state vector Q and writes into G
6677841947SLeila Ghaffari PetscErrorCode RHS_NS(TS ts, PetscReal t, Vec Q, Vec G, void *user_data) {
6777841947SLeila Ghaffari   User        user = *(User *)user_data;
68*7d4c6defSJames Wright   Ceed        ceed = user->ceed;
69c798d55aSJames Wright   PetscScalar dt;
709ad5e8e4SJames Wright   Vec         Q_loc = user->Q_loc;
7177841947SLeila Ghaffari 
72f17d818dSJames Wright   PetscFunctionBeginUser;
735e82a6e1SJeremy L Thompson   // Update time dependent data
74a0b9a424SJames Wright   PetscCall(UpdateBoundaryValues(user, Q_loc, t));
75*7d4c6defSJames Wright   if (user->phys->solution_time_label) PetscCallCeed(ceed, CeedOperatorSetContextDouble(user->op_rhs_ctx->op, user->phys->solution_time_label, &t));
762b730f8bSJeremy L Thompson   PetscCall(TSGetTimeStep(ts, &dt));
77*7d4c6defSJames Wright   if (user->phys->timestep_size_label) PetscCallCeed(ceed, CeedOperatorSetContextDouble(user->op_rhs_ctx->op, user->phys->timestep_size_label, &dt));
7877841947SLeila Ghaffari 
799ad5e8e4SJames Wright   PetscCall(ApplyCeedOperatorGlobalToGlobal(Q, G, user->op_rhs_ctx));
8077841947SLeila Ghaffari 
81d6e67e47SJames Wright   // Inverse of the lumped mass matrix
82186595e6SJames Wright   PetscCall(VecPointwiseMult(G, G, user->M_inv));
83ee4ca9cbSJames Wright   PetscFunctionReturn(PETSC_SUCCESS);
8477841947SLeila Ghaffari }
8577841947SLeila Ghaffari 
86ca69d878SAdeleke O. Bankole // Surface forces function setup
87ca69d878SAdeleke O. Bankole static PetscErrorCode Surface_Forces_NS(DM dm, Vec G_loc, PetscInt num_walls, const PetscInt walls[], PetscScalar *reaction_force) {
88ca69d878SAdeleke O. Bankole   DMLabel            face_label;
89ca69d878SAdeleke O. Bankole   const PetscScalar *g;
90d6734f85SAdeleke O. Bankole   PetscInt           dof, dim = 3;
91ca69d878SAdeleke O. Bankole   MPI_Comm           comm;
92d6734f85SAdeleke O. Bankole   PetscSection       s;
93ca69d878SAdeleke O. Bankole 
94ca69d878SAdeleke O. Bankole   PetscFunctionBeginUser;
95ca69d878SAdeleke O. Bankole   PetscCall(PetscArrayzero(reaction_force, num_walls * dim));
96ca69d878SAdeleke O. Bankole   PetscCall(PetscObjectGetComm((PetscObject)dm, &comm));
97ca69d878SAdeleke O. Bankole   PetscCall(DMGetLabel(dm, "Face Sets", &face_label));
98ca69d878SAdeleke O. Bankole   PetscCall(VecGetArrayRead(G_loc, &g));
99ca69d878SAdeleke O. Bankole   for (PetscInt w = 0; w < num_walls; w++) {
100ca69d878SAdeleke O. Bankole     const PetscInt wall = walls[w];
101ca69d878SAdeleke O. Bankole     IS             wall_is;
102d6734f85SAdeleke O. Bankole     PetscCall(DMGetLocalSection(dm, &s));
103ca69d878SAdeleke O. Bankole     PetscCall(DMLabelGetStratumIS(face_label, wall, &wall_is));
104ca69d878SAdeleke O. Bankole     if (wall_is) {  // There exist such points on this process
105ca69d878SAdeleke O. Bankole       PetscInt        num_points;
106d6734f85SAdeleke O. Bankole       PetscInt        num_comp = 0;
107ca69d878SAdeleke O. Bankole       const PetscInt *points;
108d6734f85SAdeleke O. Bankole       PetscCall(PetscSectionGetFieldComponents(s, 0, &num_comp));
109ca69d878SAdeleke O. Bankole       PetscCall(ISGetSize(wall_is, &num_points));
110ca69d878SAdeleke O. Bankole       PetscCall(ISGetIndices(wall_is, &points));
111ca69d878SAdeleke O. Bankole       for (PetscInt i = 0; i < num_points; i++) {
112ca69d878SAdeleke O. Bankole         const PetscInt           p = points[i];
113ca69d878SAdeleke O. Bankole         const StateConservative *r;
114ca69d878SAdeleke O. Bankole         PetscCall(DMPlexPointLocalRead(dm, p, g, &r));
115d6734f85SAdeleke O. Bankole         PetscCall(PetscSectionGetDof(s, p, &dof));
116d6734f85SAdeleke O. Bankole         for (PetscInt node = 0; node < dof / num_comp; node++) {
117ca69d878SAdeleke O. Bankole           for (PetscInt j = 0; j < 3; j++) {
118d6734f85SAdeleke O. Bankole             reaction_force[w * dim + j] -= r[node].momentum[j];
119d6734f85SAdeleke O. Bankole           }
120ca69d878SAdeleke O. Bankole         }
121ca69d878SAdeleke O. Bankole       }
122ca69d878SAdeleke O. Bankole       PetscCall(ISRestoreIndices(wall_is, &points));
123ca69d878SAdeleke O. Bankole     }
124ca69d878SAdeleke O. Bankole     PetscCall(ISDestroy(&wall_is));
125ca69d878SAdeleke O. Bankole   }
126ca69d878SAdeleke O. Bankole   PetscCallMPI(MPI_Allreduce(MPI_IN_PLACE, reaction_force, dim * num_walls, MPIU_SCALAR, MPI_SUM, comm));
127ca69d878SAdeleke O. Bankole   //  Restore Vectors
128ca69d878SAdeleke O. Bankole   PetscCall(VecRestoreArrayRead(G_loc, &g));
129ee4ca9cbSJames Wright   PetscFunctionReturn(PETSC_SUCCESS);
130ca69d878SAdeleke O. Bankole }
131ca69d878SAdeleke O. Bankole 
13277841947SLeila Ghaffari // Implicit time-stepper function setup
1332b730f8bSJeremy L Thompson PetscErrorCode IFunction_NS(TS ts, PetscReal t, Vec Q, Vec Q_dot, Vec G, void *user_data) {
13477841947SLeila Ghaffari   User         user = *(User *)user_data;
135*7d4c6defSJames Wright   Ceed         ceed = user->ceed;
136c798d55aSJames Wright   PetscScalar  dt;
1375e82a6e1SJeremy L Thompson   Vec          Q_loc = user->Q_loc, Q_dot_loc = user->Q_dot_loc, G_loc;
13877841947SLeila Ghaffari   PetscMemType q_mem_type, q_dot_mem_type, g_mem_type;
13977841947SLeila Ghaffari 
140f17d818dSJames Wright   PetscFunctionBeginUser;
1415e82a6e1SJeremy L Thompson   // Get local vectors
142ca69d878SAdeleke O. Bankole   PetscCall(DMGetNamedLocalVector(user->dm, "ResidualLocal", &G_loc));
1435e82a6e1SJeremy L Thompson 
1445e82a6e1SJeremy L Thompson   // Update time dependent data
145a0b9a424SJames Wright   PetscCall(UpdateBoundaryValues(user, Q_loc, t));
146*7d4c6defSJames Wright   if (user->phys->solution_time_label) PetscCallCeed(ceed, CeedOperatorSetContextDouble(user->op_ifunction, user->phys->solution_time_label, &t));
1472b730f8bSJeremy L Thompson   PetscCall(TSGetTimeStep(ts, &dt));
148*7d4c6defSJames Wright   if (user->phys->timestep_size_label) PetscCallCeed(ceed, CeedOperatorSetContextDouble(user->op_ifunction, user->phys->timestep_size_label, &dt));
14977841947SLeila Ghaffari 
15077841947SLeila Ghaffari   // Global-to-local
151878eb0e7SJames Wright   PetscCall(DMGlobalToLocalBegin(user->dm, Q, INSERT_VALUES, Q_loc));
152878eb0e7SJames Wright   PetscCall(DMGlobalToLocalBegin(user->dm, Q_dot, INSERT_VALUES, Q_dot_loc));
153878eb0e7SJames Wright   PetscCall(DMGlobalToLocalEnd(user->dm, Q, INSERT_VALUES, Q_loc));
154878eb0e7SJames Wright   PetscCall(DMGlobalToLocalEnd(user->dm, Q_dot, INSERT_VALUES, Q_dot_loc));
15577841947SLeila Ghaffari 
15677841947SLeila Ghaffari   // Place PETSc vectors in CEED vectors
157c798d55aSJames Wright   PetscCall(VecReadP2C(Q_loc, &q_mem_type, user->q_ceed));
158c798d55aSJames Wright   PetscCall(VecReadP2C(Q_dot_loc, &q_dot_mem_type, user->q_dot_ceed));
159c798d55aSJames Wright   PetscCall(VecP2C(G_loc, &g_mem_type, user->g_ceed));
16077841947SLeila Ghaffari 
16177841947SLeila Ghaffari   // Apply CEED operator
16275d1798cSJames Wright   PetscCall(PetscLogEventBegin(FLUIDS_CeedOperatorApply, Q, G, 0, 0));
16375d1798cSJames Wright   PetscCall(PetscLogGpuTimeBegin());
164a424bcd0SJames Wright   PetscCallCeed(user->ceed, CeedOperatorApply(user->op_ifunction, user->q_ceed, user->g_ceed, CEED_REQUEST_IMMEDIATE));
16575d1798cSJames Wright   PetscCall(PetscLogGpuTimeEnd());
16675d1798cSJames Wright   PetscCall(PetscLogEventEnd(FLUIDS_CeedOperatorApply, Q, G, 0, 0));
16777841947SLeila Ghaffari 
16877841947SLeila Ghaffari   // Restore vectors
169c798d55aSJames Wright   PetscCall(VecReadC2P(user->q_ceed, q_mem_type, Q_loc));
170c798d55aSJames Wright   PetscCall(VecReadC2P(user->q_dot_ceed, q_dot_mem_type, Q_dot_loc));
171c798d55aSJames Wright   PetscCall(VecC2P(user->g_ceed, g_mem_type, G_loc));
17277841947SLeila Ghaffari 
17319ffbc25SJames Wright   if (user->app_ctx->sgs_model_type == SGS_MODEL_DATA_DRIVEN) {
174cbef7084SJames Wright     PetscCall(SgsDDModelApplyIFunction(user, Q_loc, G_loc));
17519ffbc25SJames Wright   }
176be34b3b7SJames Wright 
17777841947SLeila Ghaffari   // Local-to-Global
1782b730f8bSJeremy L Thompson   PetscCall(VecZeroEntries(G));
1792b730f8bSJeremy L Thompson   PetscCall(DMLocalToGlobal(user->dm, G_loc, ADD_VALUES, G));
18077841947SLeila Ghaffari 
18177841947SLeila Ghaffari   // Restore vectors
182ca69d878SAdeleke O. Bankole   PetscCall(DMRestoreNamedLocalVector(user->dm, "ResidualLocal", &G_loc));
183ee4ca9cbSJames Wright   PetscFunctionReturn(PETSC_SUCCESS);
18477841947SLeila Ghaffari }
18577841947SLeila Ghaffari 
1862b730f8bSJeremy L Thompson static PetscErrorCode FormPreallocation(User user, PetscBool pbdiagonal, Mat J, CeedVector *coo_values) {
187544be873SJed Brown   PetscCount ncoo;
188457e73b2SJames Wright   PetscInt  *rows_petsc, *cols_petsc;
18901f0e615SJames Wright   CeedInt   *rows_ceed, *cols_ceed;
190544be873SJed Brown 
191544be873SJed Brown   PetscFunctionBeginUser;
192544be873SJed Brown   if (pbdiagonal) {
19301f0e615SJames Wright     PetscCallCeed(user->ceed, CeedOperatorLinearAssemblePointBlockDiagonalSymbolic(user->op_ijacobian, &ncoo, &rows_ceed, &cols_ceed));
194544be873SJed Brown   } else {
195a424bcd0SJames Wright     PetscCallCeed(user->ceed, CeedOperatorLinearAssembleSymbolic(user->op_ijacobian, &ncoo, &rows_ceed, &cols_ceed));
19601f0e615SJames Wright   }
197457e73b2SJames Wright   PetscCall(IntArrayC2P(ncoo, &rows_ceed, &rows_petsc));
198457e73b2SJames Wright   PetscCall(IntArrayC2P(ncoo, &cols_ceed, &cols_petsc));
199457e73b2SJames Wright   PetscCall(MatSetPreallocationCOOLocal(J, ncoo, rows_petsc, cols_petsc));
200457e73b2SJames Wright   free(rows_petsc);
201457e73b2SJames Wright   free(cols_petsc);
202a424bcd0SJames Wright   PetscCallCeed(user->ceed, CeedVectorCreate(user->ceed, ncoo, coo_values));
203ee4ca9cbSJames Wright   PetscFunctionReturn(PETSC_SUCCESS);
204544be873SJed Brown }
205544be873SJed Brown 
2062b730f8bSJeremy L Thompson static PetscErrorCode FormSetValues(User user, PetscBool pbdiagonal, Mat J, CeedVector coo_values) {
207544be873SJed Brown   CeedMemType        mem_type = CEED_MEM_HOST;
208544be873SJed Brown   const PetscScalar *values;
209544be873SJed Brown   MatType            mat_type;
210544be873SJed Brown 
211544be873SJed Brown   PetscFunctionBeginUser;
212544be873SJed Brown   PetscCall(MatGetType(J, &mat_type));
2132b730f8bSJeremy L Thompson   if (strstr(mat_type, "kokkos") || strstr(mat_type, "cusparse")) mem_type = CEED_MEM_DEVICE;
214fbb546ddSJames Wright   if (pbdiagonal) {
21575d1798cSJames Wright     PetscCall(PetscLogEventBegin(FLUIDS_CeedOperatorAssemblePointBlockDiagonal, J, 0, 0, 0));
21675d1798cSJames Wright     PetscCall(PetscLogGpuTimeBegin());
217a424bcd0SJames Wright     PetscCallCeed(user->ceed, CeedOperatorLinearAssemblePointBlockDiagonal(user->op_ijacobian, coo_values, CEED_REQUEST_IMMEDIATE));
21875d1798cSJames Wright     PetscCall(PetscLogGpuTimeEnd());
21975d1798cSJames Wright     PetscCall(PetscLogEventEnd(FLUIDS_CeedOperatorAssemblePointBlockDiagonal, J, 0, 0, 0));
220544be873SJed Brown   } else {
22175d1798cSJames Wright     PetscCall(PetscLogEventBegin(FLUIDS_CeedOperatorAssemble, J, 0, 0, 0));
22275d1798cSJames Wright     PetscCall(PetscLogGpuTimeBegin());
223a424bcd0SJames Wright     PetscCallCeed(user->ceed, CeedOperatorLinearAssemble(user->op_ijacobian, coo_values));
22475d1798cSJames Wright     PetscCall(PetscLogGpuTimeEnd());
22575d1798cSJames Wright     PetscCall(PetscLogEventEnd(FLUIDS_CeedOperatorAssemble, J, 0, 0, 0));
226544be873SJed Brown   }
227a424bcd0SJames Wright   PetscCallCeed(user->ceed, CeedVectorGetArrayRead(coo_values, mem_type, &values));
228544be873SJed Brown   PetscCall(MatSetValuesCOO(J, values, INSERT_VALUES));
229a424bcd0SJames Wright   PetscCallCeed(user->ceed, CeedVectorRestoreArrayRead(coo_values, &values));
230ee4ca9cbSJames Wright   PetscFunctionReturn(PETSC_SUCCESS);
231544be873SJed Brown }
232544be873SJed Brown 
2332b730f8bSJeremy L Thompson PetscErrorCode FormIJacobian_NS(TS ts, PetscReal t, Vec Q, Vec Q_dot, PetscReal shift, Mat J, Mat J_pre, void *user_data) {
234e334ad8fSJed Brown   User      user = *(User *)user_data;
235a424bcd0SJames Wright   Ceed      ceed = user->ceed;
236d6bf345cSJed Brown   PetscBool J_is_shell, J_is_mffd, J_pre_is_shell;
237f17d818dSJames Wright 
238e334ad8fSJed Brown   PetscFunctionBeginUser;
239a424bcd0SJames Wright   if (user->phys->ijacobian_time_shift_label)
240a424bcd0SJames Wright     PetscCallCeed(ceed, CeedOperatorSetContextDouble(user->op_ijacobian, user->phys->ijacobian_time_shift_label, &shift));
241d6bf345cSJed Brown   PetscCall(PetscObjectTypeCompare((PetscObject)J, MATMFFD, &J_is_mffd));
242e334ad8fSJed Brown   PetscCall(PetscObjectTypeCompare((PetscObject)J, MATSHELL, &J_is_shell));
2432b730f8bSJeremy L Thompson   PetscCall(PetscObjectTypeCompare((PetscObject)J_pre, MATSHELL, &J_pre_is_shell));
244e334ad8fSJed Brown   if (!user->matrices_set_up) {
245e334ad8fSJed Brown     if (J_is_shell) {
246b07979f9SJames Wright       OperatorApplyContext op_ijacobian_ctx;
247b07979f9SJames Wright       OperatorApplyContextCreate(user->dm, user->dm, user->ceed, user->op_ijacobian, user->q_ceed, user->g_ceed, user->Q_dot_loc, NULL,
248b07979f9SJames Wright                                  &op_ijacobian_ctx);
249b07979f9SJames Wright       PetscCall(MatShellSetContext(J, op_ijacobian_ctx));
250b07979f9SJames Wright       PetscCall(MatShellSetContextDestroy(J, (PetscErrorCode(*)(void *))OperatorApplyContextDestroy));
251b07979f9SJames Wright       PetscCall(MatShellSetOperation(J, MATOP_MULT, (void (*)(void))MatMult_Ceed));
252b07979f9SJames Wright       PetscCall(MatShellSetOperation(J, MATOP_GET_DIAGONAL, (void (*)(void))MatGetDiag_Ceed));
253e334ad8fSJed Brown       PetscCall(MatSetUp(J));
254e334ad8fSJed Brown     }
255e334ad8fSJed Brown     if (!J_pre_is_shell) {
2562b730f8bSJeremy L Thompson       PetscCall(FormPreallocation(user, user->app_ctx->pmat_pbdiagonal, J_pre, &user->coo_values_pmat));
257544be873SJed Brown     }
258d6bf345cSJed Brown     if (J != J_pre && !J_is_shell && !J_is_mffd) {
259544be873SJed Brown       PetscCall(FormPreallocation(user, PETSC_FALSE, J, &user->coo_values_amat));
260544be873SJed Brown     }
261e334ad8fSJed Brown     user->matrices_set_up = true;
262e334ad8fSJed Brown   }
263e334ad8fSJed Brown   if (!J_pre_is_shell) {
2642b730f8bSJeremy L Thompson     PetscCall(FormSetValues(user, user->app_ctx->pmat_pbdiagonal, J_pre, user->coo_values_pmat));
265e334ad8fSJed Brown   }
266d6bf345cSJed Brown   if (user->coo_values_amat) {
267d6bf345cSJed Brown     PetscCall(FormSetValues(user, PETSC_FALSE, J, user->coo_values_amat));
268d6bf345cSJed Brown   } else if (J_is_mffd) {
269d6bf345cSJed Brown     PetscCall(MatAssemblyBegin(J, MAT_FINAL_ASSEMBLY));
270d6bf345cSJed Brown     PetscCall(MatAssemblyEnd(J, MAT_FINAL_ASSEMBLY));
271d6bf345cSJed Brown   }
272ee4ca9cbSJames Wright   PetscFunctionReturn(PETSC_SUCCESS);
273e334ad8fSJed Brown }
274e334ad8fSJed Brown 
2752b730f8bSJeremy L Thompson PetscErrorCode WriteOutput(User user, Vec Q, PetscInt step_no, PetscScalar time) {
27677841947SLeila Ghaffari   Vec         Q_loc;
27777841947SLeila Ghaffari   char        file_path[PETSC_MAX_PATH_LEN];
27877841947SLeila Ghaffari   PetscViewer viewer;
27977841947SLeila Ghaffari 
280f17d818dSJames Wright   PetscFunctionBeginUser;
28137cbb16aSJed Brown   if (user->app_ctx->checkpoint_vtk) {
28277841947SLeila Ghaffari     // Set up output
283f14660a4SJames Wright     PetscCall(DMGetLocalVector(user->dm, &Q_loc));
284f14660a4SJames Wright     PetscCall(PetscObjectSetName((PetscObject)Q_loc, "StateVec"));
285f14660a4SJames Wright     PetscCall(VecZeroEntries(Q_loc));
286f14660a4SJames Wright     PetscCall(DMGlobalToLocal(user->dm, Q, INSERT_VALUES, Q_loc));
28777841947SLeila Ghaffari 
28877841947SLeila Ghaffari     // Output
28937cbb16aSJed Brown     PetscCall(PetscSNPrintf(file_path, sizeof file_path, "%s/ns-%03" PetscInt_FMT ".vtu", user->app_ctx->output_dir, step_no));
290f14660a4SJames Wright 
2912b730f8bSJeremy L Thompson     PetscCall(PetscViewerVTKOpen(PetscObjectComm((PetscObject)Q), file_path, FILE_MODE_WRITE, &viewer));
292f14660a4SJames Wright     PetscCall(VecView(Q_loc, viewer));
293f14660a4SJames Wright     PetscCall(PetscViewerDestroy(&viewer));
29477841947SLeila Ghaffari     if (user->dm_viz) {
29577841947SLeila Ghaffari       Vec         Q_refined, Q_refined_loc;
29677841947SLeila Ghaffari       char        file_path_refined[PETSC_MAX_PATH_LEN];
29777841947SLeila Ghaffari       PetscViewer viewer_refined;
29877841947SLeila Ghaffari 
299f14660a4SJames Wright       PetscCall(DMGetGlobalVector(user->dm_viz, &Q_refined));
300f14660a4SJames Wright       PetscCall(DMGetLocalVector(user->dm_viz, &Q_refined_loc));
301f14660a4SJames Wright       PetscCall(PetscObjectSetName((PetscObject)Q_refined_loc, "Refined"));
302f14660a4SJames Wright 
303f14660a4SJames Wright       PetscCall(MatInterpolate(user->interp_viz, Q, Q_refined));
304f14660a4SJames Wright       PetscCall(VecZeroEntries(Q_refined_loc));
3052b730f8bSJeremy L Thompson       PetscCall(DMGlobalToLocal(user->dm_viz, Q_refined, INSERT_VALUES, Q_refined_loc));
306f14660a4SJames Wright 
30737cbb16aSJed Brown       PetscCall(
30837cbb16aSJed Brown           PetscSNPrintf(file_path_refined, sizeof file_path_refined, "%s/nsrefined-%03" PetscInt_FMT ".vtu", user->app_ctx->output_dir, step_no));
309f14660a4SJames Wright 
3102b730f8bSJeremy L Thompson       PetscCall(PetscViewerVTKOpen(PetscObjectComm((PetscObject)Q_refined), file_path_refined, FILE_MODE_WRITE, &viewer_refined));
311f14660a4SJames Wright       PetscCall(VecView(Q_refined_loc, viewer_refined));
312f14660a4SJames Wright       PetscCall(DMRestoreLocalVector(user->dm_viz, &Q_refined_loc));
313f14660a4SJames Wright       PetscCall(DMRestoreGlobalVector(user->dm_viz, &Q_refined));
314f14660a4SJames Wright       PetscCall(PetscViewerDestroy(&viewer_refined));
31577841947SLeila Ghaffari     }
316f14660a4SJames Wright     PetscCall(DMRestoreLocalVector(user->dm, &Q_loc));
31737cbb16aSJed Brown   }
31877841947SLeila Ghaffari 
31977841947SLeila Ghaffari   // Save data in a binary file for continuation of simulations
32069293791SJames Wright   if (user->app_ctx->add_stepnum2bin) {
32137cbb16aSJed Brown     PetscCall(PetscSNPrintf(file_path, sizeof file_path, "%s/ns-solution-%" PetscInt_FMT ".bin", user->app_ctx->output_dir, step_no));
32269293791SJames Wright   } else {
3232b730f8bSJeremy L Thompson     PetscCall(PetscSNPrintf(file_path, sizeof file_path, "%s/ns-solution.bin", user->app_ctx->output_dir));
32469293791SJames Wright   }
3252b730f8bSJeremy L Thompson   PetscCall(PetscViewerBinaryOpen(user->comm, file_path, FILE_MODE_WRITE, &viewer));
326f14660a4SJames Wright 
3270de6a49fSJames Wright   PetscInt32 token = PetscDefined(USE_64BIT_INDICES) ? FLUIDS_FILE_TOKEN_64 : FLUIDS_FILE_TOKEN_32;
3280de6a49fSJames Wright   PetscCall(PetscViewerBinaryWrite(viewer, &token, 1, PETSC_INT32));
3294de8550aSJed Brown   PetscCall(PetscViewerBinaryWrite(viewer, &step_no, 1, PETSC_INT));
3304de8550aSJed Brown   time /= user->units->second;  // Dimensionalize time back
3314de8550aSJed Brown   PetscCall(PetscViewerBinaryWrite(viewer, &time, 1, PETSC_REAL));
332f14660a4SJames Wright   PetscCall(VecView(Q, viewer));
333f14660a4SJames Wright   PetscCall(PetscViewerDestroy(&viewer));
334ee4ca9cbSJames Wright   PetscFunctionReturn(PETSC_SUCCESS);
335f14660a4SJames Wright }
336f14660a4SJames Wright 
337ca69d878SAdeleke O. Bankole // CSV Monitor
338ca69d878SAdeleke O. Bankole PetscErrorCode TSMonitor_WallForce(TS ts, PetscInt step_no, PetscReal time, Vec Q, void *ctx) {
339ca69d878SAdeleke O. Bankole   User              user = ctx;
340ca69d878SAdeleke O. Bankole   Vec               G_loc;
341ca69d878SAdeleke O. Bankole   PetscInt          num_wall = user->app_ctx->wall_forces.num_wall, dim = 3;
342ca69d878SAdeleke O. Bankole   const PetscInt   *walls  = user->app_ctx->wall_forces.walls;
343ca69d878SAdeleke O. Bankole   PetscViewer       viewer = user->app_ctx->wall_forces.viewer;
344ca69d878SAdeleke O. Bankole   PetscViewerFormat format = user->app_ctx->wall_forces.viewer_format;
345ca69d878SAdeleke O. Bankole   PetscScalar      *reaction_force;
346ca69d878SAdeleke O. Bankole   PetscBool         iascii;
347ca69d878SAdeleke O. Bankole 
348ca69d878SAdeleke O. Bankole   PetscFunctionBeginUser;
349ee4ca9cbSJames Wright   if (!viewer) PetscFunctionReturn(PETSC_SUCCESS);
350ca69d878SAdeleke O. Bankole   PetscCall(DMGetNamedLocalVector(user->dm, "ResidualLocal", &G_loc));
351ca69d878SAdeleke O. Bankole   PetscCall(PetscMalloc1(num_wall * dim, &reaction_force));
352ca69d878SAdeleke O. Bankole   PetscCall(Surface_Forces_NS(user->dm, G_loc, num_wall, walls, reaction_force));
353ca69d878SAdeleke O. Bankole   PetscCall(DMRestoreNamedLocalVector(user->dm, "ResidualLocal", &G_loc));
354ca69d878SAdeleke O. Bankole 
355ca69d878SAdeleke O. Bankole   PetscCall(PetscObjectTypeCompare((PetscObject)viewer, PETSCVIEWERASCII, &iascii));
356ca69d878SAdeleke O. Bankole 
357ca69d878SAdeleke O. Bankole   if (iascii) {
358ca69d878SAdeleke O. Bankole     if (format == PETSC_VIEWER_ASCII_CSV && !user->app_ctx->wall_forces.header_written) {
359ca69d878SAdeleke O. Bankole       PetscCall(PetscViewerASCIIPrintf(viewer, "Step,Time,Wall,ForceX,ForceY,ForceZ\n"));
360ca69d878SAdeleke O. Bankole       user->app_ctx->wall_forces.header_written = PETSC_TRUE;
361ca69d878SAdeleke O. Bankole     }
362ca69d878SAdeleke O. Bankole     for (PetscInt w = 0; w < num_wall; w++) {
363ca69d878SAdeleke O. Bankole       PetscInt wall = walls[w];
364ca69d878SAdeleke O. Bankole       if (format == PETSC_VIEWER_ASCII_CSV) {
365ca69d878SAdeleke O. Bankole         PetscCall(PetscViewerASCIIPrintf(viewer, "%" PetscInt_FMT ",%g,%" PetscInt_FMT ",%g,%g,%g\n", step_no, time, wall,
366ca69d878SAdeleke O. Bankole                                          reaction_force[w * dim + 0], reaction_force[w * dim + 1], reaction_force[w * dim + 2]));
367ca69d878SAdeleke O. Bankole 
368ca69d878SAdeleke O. Bankole       } else {
369ca69d878SAdeleke O. Bankole         PetscCall(PetscViewerASCIIPrintf(viewer, "Wall %" PetscInt_FMT " Forces: Force_x = %12g, Force_y = %12g, Force_z = %12g\n", wall,
370ca69d878SAdeleke O. Bankole                                          reaction_force[w * dim + 0], reaction_force[w * dim + 1], reaction_force[w * dim + 2]));
371ca69d878SAdeleke O. Bankole       }
372ca69d878SAdeleke O. Bankole     }
373ca69d878SAdeleke O. Bankole   }
374ca69d878SAdeleke O. Bankole   PetscCall(PetscFree(reaction_force));
375ee4ca9cbSJames Wright   PetscFunctionReturn(PETSC_SUCCESS);
376ca69d878SAdeleke O. Bankole }
377ca69d878SAdeleke O. Bankole 
378f14660a4SJames Wright // User provided TS Monitor
3792b730f8bSJeremy L Thompson PetscErrorCode TSMonitor_NS(TS ts, PetscInt step_no, PetscReal time, Vec Q, void *ctx) {
380f14660a4SJames Wright   User user = ctx;
381f14660a4SJames Wright 
382f17d818dSJames Wright   PetscFunctionBeginUser;
38337cbb16aSJed Brown   // Print every 'checkpoint_interval' steps
384894de27cSJames Wright   if (user->app_ctx->checkpoint_interval <= 0 || step_no % user->app_ctx->checkpoint_interval != 0 ||
38549aac155SJeremy L Thompson       (user->app_ctx->cont_steps == step_no && step_no != 0)) {
386ee4ca9cbSJames Wright     PetscFunctionReturn(PETSC_SUCCESS);
38749aac155SJeremy L Thompson   }
388f14660a4SJames Wright 
389f14660a4SJames Wright   PetscCall(WriteOutput(user, Q, step_no, time));
390ee4ca9cbSJames Wright   PetscFunctionReturn(PETSC_SUCCESS);
39177841947SLeila Ghaffari }
39277841947SLeila Ghaffari 
39377841947SLeila Ghaffari // TS: Create, setup, and solve
3942b730f8bSJeremy L Thompson PetscErrorCode TSSolve_NS(DM dm, User user, AppCtx app_ctx, Physics phys, Vec *Q, PetscScalar *f_time, TS *ts) {
39577841947SLeila Ghaffari   MPI_Comm    comm = user->comm;
39677841947SLeila Ghaffari   TSAdapt     adapt;
39777841947SLeila Ghaffari   PetscScalar final_time;
39877841947SLeila Ghaffari 
399f17d818dSJames Wright   PetscFunctionBeginUser;
4002b730f8bSJeremy L Thompson   PetscCall(TSCreate(comm, ts));
4012b730f8bSJeremy L Thompson   PetscCall(TSSetDM(*ts, dm));
40277841947SLeila Ghaffari   if (phys->implicit) {
4032b730f8bSJeremy L Thompson     PetscCall(TSSetType(*ts, TSBDF));
40477841947SLeila Ghaffari     if (user->op_ifunction) {
4052b730f8bSJeremy L Thompson       PetscCall(TSSetIFunction(*ts, NULL, IFunction_NS, &user));
40677841947SLeila Ghaffari     } else {  // Implicit integrators can fall back to using an RHSFunction
4072b730f8bSJeremy L Thompson       PetscCall(TSSetRHSFunction(*ts, NULL, RHS_NS, &user));
40877841947SLeila Ghaffari     }
409e334ad8fSJed Brown     if (user->op_ijacobian) {
4102b730f8bSJeremy L Thompson       PetscCall(DMTSSetIJacobian(dm, FormIJacobian_NS, &user));
411544be873SJed Brown       if (app_ctx->amat_type) {
412544be873SJed Brown         Mat Pmat, Amat;
4132b730f8bSJeremy L Thompson         PetscCall(DMCreateMatrix(dm, &Pmat));
4142b730f8bSJeremy L Thompson         PetscCall(DMSetMatType(dm, app_ctx->amat_type));
4152b730f8bSJeremy L Thompson         PetscCall(DMCreateMatrix(dm, &Amat));
4162b730f8bSJeremy L Thompson         PetscCall(TSSetIJacobian(*ts, Amat, Pmat, NULL, NULL));
4172b730f8bSJeremy L Thompson         PetscCall(MatDestroy(&Amat));
4182b730f8bSJeremy L Thompson         PetscCall(MatDestroy(&Pmat));
419544be873SJed Brown       }
420e334ad8fSJed Brown     }
42177841947SLeila Ghaffari   } else {
4229ad5e8e4SJames Wright     PetscCheck(user->op_rhs_ctx, comm, PETSC_ERR_ARG_NULL, "Problem does not provide RHSFunction");
4232b730f8bSJeremy L Thompson     PetscCall(TSSetType(*ts, TSRK));
4242b730f8bSJeremy L Thompson     PetscCall(TSRKSetType(*ts, TSRK5F));
4252b730f8bSJeremy L Thompson     PetscCall(TSSetRHSFunction(*ts, NULL, RHS_NS, &user));
42677841947SLeila Ghaffari   }
4272b730f8bSJeremy L Thompson   PetscCall(TSSetMaxTime(*ts, 500. * user->units->second));
4282b730f8bSJeremy L Thompson   PetscCall(TSSetExactFinalTime(*ts, TS_EXACTFINALTIME_STEPOVER));
42951b00d91SJames Wright   if (app_ctx->test_type == TESTTYPE_NONE) PetscCall(TSSetErrorIfStepFails(*ts, PETSC_FALSE));
4302b730f8bSJeremy L Thompson   PetscCall(TSSetTimeStep(*ts, 1.e-2 * user->units->second));
4318fb33541SJames Wright   if (app_ctx->test_type != TESTTYPE_NONE) {
4322b730f8bSJeremy L Thompson     PetscCall(TSSetMaxSteps(*ts, 10));
4332b730f8bSJeremy L Thompson   }
4342b730f8bSJeremy L Thompson   PetscCall(TSGetAdapt(*ts, &adapt));
4352b730f8bSJeremy L Thompson   PetscCall(TSAdaptSetStepLimits(adapt, 1.e-12 * user->units->second, 1.e2 * user->units->second));
4362b730f8bSJeremy L Thompson   PetscCall(TSSetFromOptions(*ts));
43715637395SJames Wright   user->time_bc_set = -1.0;   // require all BCs be updated
438d55646a4SJames Wright   if (app_ctx->cont_steps) {  // continue from previous timestep data
43977841947SLeila Ghaffari     PetscInt    count;
44077841947SLeila Ghaffari     PetscViewer viewer;
4412b730f8bSJeremy L Thompson 
4424de8550aSJed Brown     if (app_ctx->cont_time <= 0) {  // Legacy files did not include step number and time
4432b730f8bSJeremy L Thompson       PetscCall(PetscViewerBinaryOpen(comm, app_ctx->cont_time_file, FILE_MODE_READ, &viewer));
4444de8550aSJed Brown       PetscCall(PetscViewerBinaryRead(viewer, &app_ctx->cont_time, 1, &count, PETSC_REAL));
4452b730f8bSJeremy L Thompson       PetscCall(PetscViewerDestroy(&viewer));
4464de8550aSJed Brown       PetscCheck(app_ctx->cont_steps != -1, comm, PETSC_ERR_ARG_INCOMP,
4474de8550aSJed Brown                  "-continue step number not specified, but checkpoint file does not contain a step number (likely written by older code version)");
4484de8550aSJed Brown     }
4494de8550aSJed Brown     PetscCall(TSSetTime(*ts, app_ctx->cont_time * user->units->second));
450d8e0aecdSJed Brown     PetscCall(TSSetStepNumber(*ts, app_ctx->cont_steps));
45177841947SLeila Ghaffari   }
4528fb33541SJames Wright   if (app_ctx->test_type == TESTTYPE_NONE) {
4532b730f8bSJeremy L Thompson     PetscCall(TSMonitorSet(*ts, TSMonitor_NS, user, NULL));
4548fb33541SJames Wright   }
455ca69d878SAdeleke O. Bankole   if (app_ctx->wall_forces.viewer) {
456ca69d878SAdeleke O. Bankole     PetscCall(TSMonitorSet(*ts, TSMonitor_WallForce, user, NULL));
457ca69d878SAdeleke O. Bankole   }
458b7d66439SJames Wright   if (app_ctx->turb_spanstats_enable) {
459f5452247SJames Wright     PetscCall(TSMonitorSet(*ts, TSMonitor_TurbulenceStatistics, user, NULL));
460495b9769SJames Wright     CeedScalar previous_time = app_ctx->cont_time * user->units->second;
461a424bcd0SJames Wright     PetscCallCeed(user->ceed,
462a424bcd0SJames Wright                   CeedOperatorSetContextDouble(user->spanstats.op_stats_collect_ctx->op, user->spanstats.previous_time_label, &previous_time));
463a175e481SJames Wright   }
4644e9802d1SJames Wright   if (app_ctx->diff_filter_monitor) PetscCall(TSMonitorSet(*ts, TSMonitor_DifferentialFilter, user, NULL));
46577841947SLeila Ghaffari 
46677841947SLeila Ghaffari   // Solve
467d8e0aecdSJed Brown   PetscReal start_time;
468d8e0aecdSJed Brown   PetscInt  start_step;
4692b730f8bSJeremy L Thompson   PetscCall(TSGetTime(*ts, &start_time));
470d8e0aecdSJed Brown   PetscCall(TSGetStepNumber(*ts, &start_step));
4717b39487dSJeremy L Thompson 
472711a423aSJed Brown   PetscCall(PetscLogDefaultBegin());  // So we can use PetscLogStageGetPerfInfo without -log_view
4737b39487dSJeremy L Thompson   PetscPreLoadBegin(PETSC_FALSE, "Fluids Solve");
4747b39487dSJeremy L Thompson   PetscCall(TSSetTime(*ts, start_time));
475d8e0aecdSJed Brown   PetscCall(TSSetStepNumber(*ts, start_step));
4767b39487dSJeremy L Thompson   if (PetscPreLoadingOn) {
4777b39487dSJeremy L Thompson     // LCOV_EXCL_START
4787b39487dSJeremy L Thompson     SNES      snes;
4797b39487dSJeremy L Thompson     Vec       Q_preload;
4807b39487dSJeremy L Thompson     PetscReal rtol;
4817b39487dSJeremy L Thompson     PetscCall(VecDuplicate(*Q, &Q_preload));
4827b39487dSJeremy L Thompson     PetscCall(VecCopy(*Q, Q_preload));
4837b39487dSJeremy L Thompson     PetscCall(TSGetSNES(*ts, &snes));
4847b39487dSJeremy L Thompson     PetscCall(SNESGetTolerances(snes, NULL, &rtol, NULL, NULL, NULL));
4852b730f8bSJeremy L Thompson     PetscCall(SNESSetTolerances(snes, PETSC_DEFAULT, .99, PETSC_DEFAULT, PETSC_DEFAULT, PETSC_DEFAULT));
486fc4c3d69SJames Wright     PetscCall(TSSetSolution(*ts, Q_preload));
4877b39487dSJeremy L Thompson     PetscCall(TSStep(*ts));
4882b730f8bSJeremy L Thompson     PetscCall(SNESSetTolerances(snes, PETSC_DEFAULT, rtol, PETSC_DEFAULT, PETSC_DEFAULT, PETSC_DEFAULT));
4897b39487dSJeremy L Thompson     PetscCall(VecDestroy(&Q_preload));
4907b39487dSJeremy L Thompson     // LCOV_EXCL_STOP
4917b39487dSJeremy L Thompson   } else {
4922b730f8bSJeremy L Thompson     PetscCall(PetscBarrier((PetscObject)*ts));
4932b730f8bSJeremy L Thompson     PetscCall(TSSolve(*ts, *Q));
4947b39487dSJeremy L Thompson   }
4957b39487dSJeremy L Thompson   PetscPreLoadEnd();
4967b39487dSJeremy L Thompson 
4977b39487dSJeremy L Thompson   PetscCall(TSGetSolveTime(*ts, &final_time));
49877841947SLeila Ghaffari   *f_time = final_time;
4997b39487dSJeremy L Thompson 
5008fb33541SJames Wright   if (app_ctx->test_type == TESTTYPE_NONE) {
501f14660a4SJames Wright     PetscInt step_no;
502f14660a4SJames Wright     PetscCall(TSGetStepNumber(*ts, &step_no));
503a175e481SJames Wright     if (user->app_ctx->checkpoint_interval > 0 || user->app_ctx->checkpoint_interval == -1) {
504f14660a4SJames Wright       PetscCall(WriteOutput(user, *Q, step_no, final_time));
505f14660a4SJames Wright     }
506f14660a4SJames Wright 
50775d1798cSJames Wright     PetscLogStage      stage_id;
508711a423aSJed Brown     PetscEventPerfInfo stage_perf;
5097b39487dSJeremy L Thompson 
5107b39487dSJeremy L Thompson     PetscCall(PetscLogStageGetId("Fluids Solve", &stage_id));
511711a423aSJed Brown     PetscCall(PetscLogStageGetPerfInfo(stage_id, &stage_perf));
512711a423aSJed Brown     PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Time taken for solution (sec): %g\n", stage_perf.time));
51377841947SLeila Ghaffari   }
514ee4ca9cbSJames Wright   PetscFunctionReturn(PETSC_SUCCESS);
51577841947SLeila Ghaffari }
516