xref: /libCEED/examples/fluids/src/setupts.c (revision 5e82a6e1725e39817d93729ca2d6f7dcbf054835)
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 
1177841947SLeila Ghaffari #include "../navierstokes.h"
1277841947SLeila Ghaffari #include "../qfunctions/mass.h"
1377841947SLeila Ghaffari 
1477841947SLeila Ghaffari // Compute mass matrix for explicit scheme
1577841947SLeila Ghaffari PetscErrorCode ComputeLumpedMassMatrix(Ceed ceed, DM dm, CeedData ceed_data,
1677841947SLeila Ghaffari                                        Vec M) {
1777841947SLeila Ghaffari   Vec            M_loc;
1877841947SLeila Ghaffari   CeedQFunction  qf_mass;
1977841947SLeila Ghaffari   CeedOperator   op_mass;
2077841947SLeila Ghaffari   CeedVector     m_ceed, ones_vec;
2177841947SLeila Ghaffari   CeedInt        num_comp_q, q_data_size;
2277841947SLeila Ghaffari   PetscErrorCode ierr;
2377841947SLeila Ghaffari   PetscFunctionBeginUser;
2477841947SLeila Ghaffari 
2577841947SLeila Ghaffari   // CEED Restriction
2677841947SLeila Ghaffari   CeedElemRestrictionGetNumComponents(ceed_data->elem_restr_q, &num_comp_q);
2777841947SLeila Ghaffari   CeedElemRestrictionGetNumComponents(ceed_data->elem_restr_qd_i, &q_data_size);
2877841947SLeila Ghaffari   CeedElemRestrictionCreateVector(ceed_data->elem_restr_q, &m_ceed, NULL);
2977841947SLeila Ghaffari   CeedElemRestrictionCreateVector(ceed_data->elem_restr_q, &ones_vec, NULL);
3077841947SLeila Ghaffari   CeedVectorSetValue(ones_vec, 1.0);
3177841947SLeila Ghaffari 
3277841947SLeila Ghaffari   // CEED QFunction
3377841947SLeila Ghaffari   CeedQFunctionCreateInterior(ceed, 1, Mass, Mass_loc, &qf_mass);
3477841947SLeila Ghaffari   CeedQFunctionAddInput(qf_mass, "q", num_comp_q, CEED_EVAL_INTERP);
35a61c78d6SJeremy L Thompson   CeedQFunctionAddInput(qf_mass, "qdata", q_data_size, CEED_EVAL_NONE);
3677841947SLeila Ghaffari   CeedQFunctionAddOutput(qf_mass, "v", num_comp_q, CEED_EVAL_INTERP);
3777841947SLeila Ghaffari 
3877841947SLeila Ghaffari   // CEED Operator
3977841947SLeila Ghaffari   CeedOperatorCreate(ceed, qf_mass, NULL, NULL, &op_mass);
4077841947SLeila Ghaffari   CeedOperatorSetField(op_mass, "q", ceed_data->elem_restr_q, ceed_data->basis_q,
4177841947SLeila Ghaffari                        CEED_VECTOR_ACTIVE);
42a61c78d6SJeremy L Thompson   CeedOperatorSetField(op_mass, "qdata", ceed_data->elem_restr_qd_i,
4377841947SLeila Ghaffari                        CEED_BASIS_COLLOCATED, ceed_data->q_data);
4477841947SLeila Ghaffari   CeedOperatorSetField(op_mass, "v", ceed_data->elem_restr_q, ceed_data->basis_q,
4577841947SLeila Ghaffari                        CEED_VECTOR_ACTIVE);
4677841947SLeila Ghaffari 
4777841947SLeila Ghaffari   // Place PETSc vector in CEED vector
4877841947SLeila Ghaffari   CeedScalar *m;
4977841947SLeila Ghaffari   PetscMemType m_mem_type;
5077841947SLeila Ghaffari   ierr = DMGetLocalVector(dm, &M_loc); CHKERRQ(ierr);
5177841947SLeila Ghaffari   ierr = VecGetArrayAndMemType(M_loc, (PetscScalar **)&m, &m_mem_type);
5277841947SLeila Ghaffari   CHKERRQ(ierr);
5377841947SLeila Ghaffari   CeedVectorSetArray(m_ceed, MemTypeP2C(m_mem_type), CEED_USE_POINTER, m);
5477841947SLeila Ghaffari 
5577841947SLeila Ghaffari   // Apply CEED Operator
5677841947SLeila Ghaffari   CeedOperatorApply(op_mass, ones_vec, m_ceed, CEED_REQUEST_IMMEDIATE);
5777841947SLeila Ghaffari 
5877841947SLeila Ghaffari   // Restore vectors
5977841947SLeila Ghaffari   CeedVectorTakeArray(m_ceed, MemTypeP2C(m_mem_type), NULL);
6077841947SLeila Ghaffari   ierr = VecRestoreArrayReadAndMemType(M_loc, (const PetscScalar **)&m);
6177841947SLeila Ghaffari   CHKERRQ(ierr);
6277841947SLeila Ghaffari 
6377841947SLeila Ghaffari   // Local-to-Global
6477841947SLeila Ghaffari   ierr = VecZeroEntries(M); CHKERRQ(ierr);
6577841947SLeila Ghaffari   ierr = DMLocalToGlobal(dm, M_loc, ADD_VALUES, M); CHKERRQ(ierr);
6677841947SLeila Ghaffari   ierr = DMRestoreLocalVector(dm, &M_loc); CHKERRQ(ierr);
6777841947SLeila Ghaffari 
6877841947SLeila Ghaffari   // Invert diagonally lumped mass vector for RHS function
6977841947SLeila Ghaffari   ierr = VecReciprocal(M); CHKERRQ(ierr);
7077841947SLeila Ghaffari 
7177841947SLeila Ghaffari   // Cleanup
7277841947SLeila Ghaffari   CeedVectorDestroy(&ones_vec);
7377841947SLeila Ghaffari   CeedVectorDestroy(&m_ceed);
7477841947SLeila Ghaffari   CeedQFunctionDestroy(&qf_mass);
7577841947SLeila Ghaffari   CeedOperatorDestroy(&op_mass);
7677841947SLeila Ghaffari 
7777841947SLeila Ghaffari   PetscFunctionReturn(0);
7877841947SLeila Ghaffari }
7977841947SLeila Ghaffari 
8077841947SLeila Ghaffari // RHS (Explicit time-stepper) function setup
8177841947SLeila Ghaffari //   This is the RHS of the ODE, given as u_t = G(t,u)
8277841947SLeila Ghaffari //   This function takes in a state vector Q and writes into G
8377841947SLeila Ghaffari PetscErrorCode RHS_NS(TS ts, PetscReal t, Vec Q, Vec G, void *user_data) {
8477841947SLeila Ghaffari   User           user = *(User *)user_data;
8577841947SLeila Ghaffari   PetscScalar    *q, *g;
86*5e82a6e1SJeremy L Thompson   Vec            Q_loc = user->Q_loc, G_loc;
8777841947SLeila Ghaffari   PetscMemType   q_mem_type, g_mem_type;
8877841947SLeila Ghaffari   PetscErrorCode ierr;
8977841947SLeila Ghaffari   PetscFunctionBeginUser;
9077841947SLeila Ghaffari 
91*5e82a6e1SJeremy L Thompson   // Get local vector
92*5e82a6e1SJeremy L Thompson   ierr = DMGetLocalVector(user->dm, &G_loc); CHKERRQ(ierr);
93*5e82a6e1SJeremy L Thompson 
94*5e82a6e1SJeremy L Thompson   // Update time dependent data
95*5e82a6e1SJeremy L Thompson   if (user->time != t) {
96*5e82a6e1SJeremy L Thompson     ierr = VecZeroEntries(Q_loc); CHKERRQ(ierr);
97*5e82a6e1SJeremy L Thompson     ierr = DMPlexInsertBoundaryValues(user->dm, PETSC_TRUE, Q_loc, t,
98*5e82a6e1SJeremy L Thompson                                       NULL, NULL, NULL); CHKERRQ(ierr);
99*5e82a6e1SJeremy L Thompson     if (user->phys->solution_time_label) {
10011436a05SJames Wright       CeedOperatorContextSetDouble(user->op_rhs, user->phys->solution_time_label, &t);
101*5e82a6e1SJeremy L Thompson     }
102*5e82a6e1SJeremy L Thompson     user->time = t;
103*5e82a6e1SJeremy L Thompson   }
10488626eedSJames Wright   if (user->phys->timestep_size_label) {
10588626eedSJames Wright     PetscScalar dt;
10688626eedSJames Wright     ierr = TSGetTimeStep(ts, &dt); CHKERRQ(ierr);
107*5e82a6e1SJeremy L Thompson     if (user->dt != dt) {
10888626eedSJames Wright       CeedOperatorContextSetDouble(user->op_rhs, user->phys->timestep_size_label,
10988626eedSJames Wright                                    &dt);
110*5e82a6e1SJeremy L Thompson       user->dt = dt;
111*5e82a6e1SJeremy L Thompson     }
11288626eedSJames Wright   }
11377841947SLeila Ghaffari 
11477841947SLeila Ghaffari   // Global-to-local
11577841947SLeila Ghaffari   ierr = DMGlobalToLocal(user->dm, Q, INSERT_VALUES, Q_loc); CHKERRQ(ierr);
11677841947SLeila Ghaffari 
11777841947SLeila Ghaffari   // Place PETSc vectors in CEED vectors
11877841947SLeila Ghaffari   ierr = VecGetArrayReadAndMemType(Q_loc, (const PetscScalar **)&q, &q_mem_type);
11977841947SLeila Ghaffari   CHKERRQ(ierr);
12077841947SLeila Ghaffari   ierr = VecGetArrayAndMemType(G_loc, &g, &g_mem_type); CHKERRQ(ierr);
12177841947SLeila Ghaffari   CeedVectorSetArray(user->q_ceed, MemTypeP2C(q_mem_type), CEED_USE_POINTER, q);
12277841947SLeila Ghaffari   CeedVectorSetArray(user->g_ceed, MemTypeP2C(g_mem_type), CEED_USE_POINTER, g);
12377841947SLeila Ghaffari 
12477841947SLeila Ghaffari   // Apply CEED operator
12577841947SLeila Ghaffari   CeedOperatorApply(user->op_rhs, user->q_ceed, user->g_ceed,
12677841947SLeila Ghaffari                     CEED_REQUEST_IMMEDIATE);
12777841947SLeila Ghaffari 
12877841947SLeila Ghaffari   // Restore vectors
12977841947SLeila Ghaffari   CeedVectorTakeArray(user->q_ceed, MemTypeP2C(q_mem_type), NULL);
13077841947SLeila Ghaffari   CeedVectorTakeArray(user->g_ceed, MemTypeP2C(g_mem_type), NULL);
13177841947SLeila Ghaffari   ierr = VecRestoreArrayReadAndMemType(Q_loc, (const PetscScalar **)&q);
13277841947SLeila Ghaffari   CHKERRQ(ierr);
13377841947SLeila Ghaffari   ierr = VecRestoreArrayAndMemType(G_loc, &g); CHKERRQ(ierr);
13477841947SLeila Ghaffari 
13577841947SLeila Ghaffari   // Local-to-Global
13677841947SLeila Ghaffari   ierr = VecZeroEntries(G); CHKERRQ(ierr);
13777841947SLeila Ghaffari   ierr = DMLocalToGlobal(user->dm, G_loc, ADD_VALUES, G); CHKERRQ(ierr);
13877841947SLeila Ghaffari 
13977841947SLeila Ghaffari   // Inverse of the lumped mass matrix (M is Minv)
14077841947SLeila Ghaffari   ierr = VecPointwiseMult(G, G, user->M); CHKERRQ(ierr);
14177841947SLeila Ghaffari 
14277841947SLeila Ghaffari   // Restore vectors
14377841947SLeila Ghaffari   ierr = DMRestoreLocalVector(user->dm, &G_loc); CHKERRQ(ierr);
14477841947SLeila Ghaffari 
14577841947SLeila Ghaffari   PetscFunctionReturn(0);
14677841947SLeila Ghaffari }
14777841947SLeila Ghaffari 
14877841947SLeila Ghaffari // Implicit time-stepper function setup
14977841947SLeila Ghaffari PetscErrorCode IFunction_NS(TS ts, PetscReal t, Vec Q, Vec Q_dot, Vec G,
15077841947SLeila Ghaffari                             void *user_data) {
15177841947SLeila Ghaffari   User              user = *(User *)user_data;
15277841947SLeila Ghaffari   const PetscScalar *q, *q_dot;
15377841947SLeila Ghaffari   PetscScalar       *g;
154*5e82a6e1SJeremy L Thompson   Vec               Q_loc = user->Q_loc, Q_dot_loc = user->Q_dot_loc, G_loc;
15577841947SLeila Ghaffari   PetscMemType      q_mem_type, q_dot_mem_type, g_mem_type;
15677841947SLeila Ghaffari   PetscErrorCode    ierr;
15777841947SLeila Ghaffari   PetscFunctionBeginUser;
15877841947SLeila Ghaffari 
159*5e82a6e1SJeremy L Thompson   // Get local vectors
160*5e82a6e1SJeremy L Thompson   ierr = DMGetLocalVector(user->dm, &G_loc); CHKERRQ(ierr);
161*5e82a6e1SJeremy L Thompson 
162*5e82a6e1SJeremy L Thompson   // Update time dependent data
163*5e82a6e1SJeremy L Thompson   if (user->time != t) {
164*5e82a6e1SJeremy L Thompson     ierr = VecZeroEntries(Q_loc); CHKERRQ(ierr);
165*5e82a6e1SJeremy L Thompson     ierr = DMPlexInsertBoundaryValues(user->dm, PETSC_TRUE, Q_loc, t,
166*5e82a6e1SJeremy L Thompson                                       NULL, NULL, NULL); CHKERRQ(ierr);
167*5e82a6e1SJeremy L Thompson     if (user->phys->solution_time_label) {
16811436a05SJames Wright       CeedOperatorContextSetDouble(user->op_ifunction,
16911436a05SJames Wright                                    user->phys->solution_time_label, &t);
170*5e82a6e1SJeremy L Thompson     }
171*5e82a6e1SJeremy L Thompson     user->time = t;
172*5e82a6e1SJeremy L Thompson   }
17388626eedSJames Wright   if (user->phys->timestep_size_label) {
17488626eedSJames Wright     PetscScalar dt;
17588626eedSJames Wright     ierr = TSGetTimeStep(ts, &dt); CHKERRQ(ierr);
176*5e82a6e1SJeremy L Thompson     if (user->dt != dt) {
17788626eedSJames Wright       CeedOperatorContextSetDouble(user->op_ifunction,
17888626eedSJames Wright                                    user->phys->timestep_size_label, &dt);
179*5e82a6e1SJeremy L Thompson       user->dt = dt;
180*5e82a6e1SJeremy L Thompson     }
18188626eedSJames Wright   }
18277841947SLeila Ghaffari 
18377841947SLeila Ghaffari   // Global-to-local
18477841947SLeila Ghaffari   ierr = DMGlobalToLocal(user->dm, Q, INSERT_VALUES, Q_loc); CHKERRQ(ierr);
18577841947SLeila Ghaffari   ierr = DMGlobalToLocal(user->dm, Q_dot, INSERT_VALUES, Q_dot_loc);
18677841947SLeila Ghaffari   CHKERRQ(ierr);
18777841947SLeila Ghaffari 
18877841947SLeila Ghaffari   // Place PETSc vectors in CEED vectors
18977841947SLeila Ghaffari   ierr = VecGetArrayReadAndMemType(Q_loc, &q, &q_mem_type); CHKERRQ(ierr);
19077841947SLeila Ghaffari   ierr = VecGetArrayReadAndMemType(Q_dot_loc, &q_dot, &q_dot_mem_type);
19177841947SLeila Ghaffari   CHKERRQ(ierr);
19277841947SLeila Ghaffari   ierr = VecGetArrayAndMemType(G_loc, &g, &g_mem_type); CHKERRQ(ierr);
19377841947SLeila Ghaffari   CeedVectorSetArray(user->q_ceed, MemTypeP2C(q_mem_type), CEED_USE_POINTER,
19477841947SLeila Ghaffari                      (PetscScalar *)q);
19577841947SLeila Ghaffari   CeedVectorSetArray(user->q_dot_ceed, MemTypeP2C(q_dot_mem_type),
196e6225c47SLeila Ghaffari                      CEED_USE_POINTER, (PetscScalar *)q_dot);
19777841947SLeila Ghaffari   CeedVectorSetArray(user->g_ceed, MemTypeP2C(g_mem_type), CEED_USE_POINTER, g);
19877841947SLeila Ghaffari 
19977841947SLeila Ghaffari   // Apply CEED operator
20077841947SLeila Ghaffari   CeedOperatorApply(user->op_ifunction, user->q_ceed, user->g_ceed,
20177841947SLeila Ghaffari                     CEED_REQUEST_IMMEDIATE);
20277841947SLeila Ghaffari 
20377841947SLeila Ghaffari   // Restore vectors
20477841947SLeila Ghaffari   CeedVectorTakeArray(user->q_ceed, MemTypeP2C(q_mem_type), NULL);
20577841947SLeila Ghaffari   CeedVectorTakeArray(user->q_dot_ceed, MemTypeP2C(q_dot_mem_type), NULL);
20677841947SLeila Ghaffari   CeedVectorTakeArray(user->g_ceed, MemTypeP2C(g_mem_type), NULL);
20777841947SLeila Ghaffari   ierr = VecRestoreArrayReadAndMemType(Q_loc, &q); CHKERRQ(ierr);
20877841947SLeila Ghaffari   ierr = VecRestoreArrayReadAndMemType(Q_dot_loc, &q_dot); CHKERRQ(ierr);
20977841947SLeila Ghaffari   ierr = VecRestoreArrayAndMemType(G_loc, &g); CHKERRQ(ierr);
21077841947SLeila Ghaffari 
21177841947SLeila Ghaffari   // Local-to-Global
21277841947SLeila Ghaffari   ierr = VecZeroEntries(G); CHKERRQ(ierr);
21377841947SLeila Ghaffari   ierr = DMLocalToGlobal(user->dm, G_loc, ADD_VALUES, G); CHKERRQ(ierr);
21477841947SLeila Ghaffari 
21577841947SLeila Ghaffari   // Restore vectors
21677841947SLeila Ghaffari   ierr = DMRestoreLocalVector(user->dm, &G_loc); CHKERRQ(ierr);
21777841947SLeila Ghaffari 
21877841947SLeila Ghaffari   PetscFunctionReturn(0);
21977841947SLeila Ghaffari }
22077841947SLeila Ghaffari 
221e334ad8fSJed Brown static PetscErrorCode MatMult_NS_IJacobian(Mat J, Vec Q, Vec G) {
222e334ad8fSJed Brown   User user;
223e334ad8fSJed Brown   const PetscScalar *q;
224e334ad8fSJed Brown   PetscScalar       *g;
225e334ad8fSJed Brown   PetscMemType      q_mem_type, g_mem_type;
226e334ad8fSJed Brown   PetscErrorCode    ierr;
227e334ad8fSJed Brown   PetscFunctionBeginUser;
228*5e82a6e1SJeremy L Thompson   ierr = MatShellGetContext(J, &user); CHKERRQ(ierr);
229*5e82a6e1SJeremy L Thompson   Vec               Q_loc = user->Q_dot_loc, // Note - Q_dot_loc has zero BCs
230*5e82a6e1SJeremy L Thompson                     G_loc;
231*5e82a6e1SJeremy L Thompson 
232e334ad8fSJed Brown   // Get local vectors
233e334ad8fSJed Brown   ierr = DMGetLocalVector(user->dm, &G_loc); CHKERRQ(ierr);
234e334ad8fSJed Brown 
235e334ad8fSJed Brown   // Global-to-local
236e334ad8fSJed Brown   ierr = DMGlobalToLocal(user->dm, Q, INSERT_VALUES, Q_loc); CHKERRQ(ierr);
237e334ad8fSJed Brown 
238e334ad8fSJed Brown   // Place PETSc vectors in CEED vectors
239e334ad8fSJed Brown   ierr = VecGetArrayReadAndMemType(Q_loc, &q, &q_mem_type); CHKERRQ(ierr);
240e334ad8fSJed Brown   ierr = VecGetArrayAndMemType(G_loc, &g, &g_mem_type); CHKERRQ(ierr);
241e334ad8fSJed Brown   CeedVectorSetArray(user->q_ceed, MemTypeP2C(q_mem_type), CEED_USE_POINTER,
242e334ad8fSJed Brown                      (PetscScalar *)q);
243e334ad8fSJed Brown   CeedVectorSetArray(user->g_ceed, MemTypeP2C(g_mem_type), CEED_USE_POINTER, g);
244e334ad8fSJed Brown 
245e334ad8fSJed Brown   // Apply CEED operator
246e334ad8fSJed Brown   CeedOperatorApply(user->op_ijacobian, user->q_ceed, user->g_ceed,
247e334ad8fSJed Brown                     CEED_REQUEST_IMMEDIATE);
248e334ad8fSJed Brown 
249e334ad8fSJed Brown   // Restore vectors
250e334ad8fSJed Brown   CeedVectorTakeArray(user->q_ceed, MemTypeP2C(q_mem_type), NULL);
251e334ad8fSJed Brown   CeedVectorTakeArray(user->g_ceed, MemTypeP2C(g_mem_type), NULL);
252e334ad8fSJed Brown   ierr = VecRestoreArrayReadAndMemType(Q_loc, &q); CHKERRQ(ierr);
253e334ad8fSJed Brown   ierr = VecRestoreArrayAndMemType(G_loc, &g); CHKERRQ(ierr);
254e334ad8fSJed Brown 
255e334ad8fSJed Brown   // Local-to-Global
256e334ad8fSJed Brown   ierr = VecZeroEntries(G); CHKERRQ(ierr);
257e334ad8fSJed Brown   ierr = DMLocalToGlobal(user->dm, G_loc, ADD_VALUES, G); CHKERRQ(ierr);
258e334ad8fSJed Brown 
259e334ad8fSJed Brown   // Restore vectors
260e334ad8fSJed Brown   ierr = DMRestoreLocalVector(user->dm, &G_loc); CHKERRQ(ierr);
261e334ad8fSJed Brown   PetscFunctionReturn(0);
262e334ad8fSJed Brown }
263e334ad8fSJed Brown 
264e334ad8fSJed Brown PetscErrorCode MatGetDiagonal_NS_IJacobian(Mat A, Vec D) {
265e334ad8fSJed Brown   User user;
266e334ad8fSJed Brown   Vec D_loc;
267e334ad8fSJed Brown   PetscScalar *d;
268e334ad8fSJed Brown   PetscMemType mem_type;
269e334ad8fSJed Brown 
270e334ad8fSJed Brown   PetscFunctionBeginUser;
271e334ad8fSJed Brown   MatShellGetContext(A, &user);
272e334ad8fSJed Brown   PetscCall(DMGetLocalVector(user->dm, &D_loc));
273e334ad8fSJed Brown   PetscCall(VecGetArrayAndMemType(D_loc, &d, &mem_type));
274e334ad8fSJed Brown   CeedVectorSetArray(user->g_ceed, MemTypeP2C(mem_type), CEED_USE_POINTER, d);
275e334ad8fSJed Brown   CeedOperatorLinearAssembleDiagonal(user->op_ijacobian, user->g_ceed,
276e334ad8fSJed Brown                                      CEED_REQUEST_IMMEDIATE);
277e334ad8fSJed Brown   CeedVectorTakeArray(user->g_ceed, MemTypeP2C(mem_type), NULL);
278e334ad8fSJed Brown   PetscCall(VecRestoreArrayAndMemType(D_loc, &d));
279e334ad8fSJed Brown   PetscCall(VecZeroEntries(D));
280e334ad8fSJed Brown   PetscCall(DMLocalToGlobal(user->dm, D_loc, ADD_VALUES, D));
281e334ad8fSJed Brown   PetscCall(DMRestoreLocalVector(user->dm, &D_loc));
282e334ad8fSJed Brown   VecViewFromOptions(D, NULL, "-diag_vec_view");
283e334ad8fSJed Brown   PetscFunctionReturn(0);
284e334ad8fSJed Brown }
285e334ad8fSJed Brown 
286e334ad8fSJed Brown PetscErrorCode FormIJacobian_NS(TS ts, PetscReal t, Vec Q, Vec Q_dot,
287e334ad8fSJed Brown                                 PetscReal shift, Mat J, Mat J_pre,
288e334ad8fSJed Brown                                 void *user_data) {
289e334ad8fSJed Brown   User user = *(User *)user_data;
290e334ad8fSJed Brown   PetscBool J_is_shell, J_pre_is_shell;
291e334ad8fSJed Brown   PetscFunctionBeginUser;
292e334ad8fSJed Brown   if (user->phys->ijacobian_time_shift_label)
293e334ad8fSJed Brown     CeedOperatorContextSetDouble(user->op_ijacobian,
294e334ad8fSJed Brown                                  user->phys->ijacobian_time_shift_label, &shift);
295e334ad8fSJed Brown   PetscCall(MatAssemblyBegin(J, MAT_FINAL_ASSEMBLY));
296e334ad8fSJed Brown   PetscCall(MatAssemblyEnd(J, MAT_FINAL_ASSEMBLY));
297e334ad8fSJed Brown   Vec coo_vec = NULL;
298e334ad8fSJed Brown   PetscCall(PetscObjectTypeCompare((PetscObject)J, MATSHELL, &J_is_shell));
299e334ad8fSJed Brown   PetscCall(PetscObjectTypeCompare((PetscObject)J_pre, MATSHELL,
300e334ad8fSJed Brown                                    &J_pre_is_shell));
301e334ad8fSJed Brown   if (!user->matrices_set_up) {
302e334ad8fSJed Brown     if (J_is_shell) {
303e334ad8fSJed Brown       PetscCall(MatShellSetContext(J, user));
304e334ad8fSJed Brown       PetscCall(MatShellSetOperation(J, MATOP_MULT,
305e334ad8fSJed Brown                                      (void (*)(void))MatMult_NS_IJacobian));
306e334ad8fSJed Brown       PetscCall(MatShellSetOperation(J, MATOP_GET_DIAGONAL,
307e334ad8fSJed Brown                                      (void (*)(void))MatGetDiagonal_NS_IJacobian));
308e334ad8fSJed Brown       PetscCall(MatSetUp(J));
309e334ad8fSJed Brown     }
310e334ad8fSJed Brown     if (!J_pre_is_shell) {
311e334ad8fSJed Brown       PetscCount ncoo;
312e334ad8fSJed Brown       PetscInt *rows, *cols;
313e334ad8fSJed Brown       PetscCall(CeedOperatorLinearAssembleSymbolic(user->op_ijacobian, &ncoo, &rows,
314e334ad8fSJed Brown                 &cols));
315e334ad8fSJed Brown       PetscCall(MatSetPreallocationCOOLocal(J_pre, ncoo, rows, cols));
316e334ad8fSJed Brown       free(rows);
317e334ad8fSJed Brown       free(cols);
318e334ad8fSJed Brown       CeedVectorCreate(user->ceed, ncoo, &user->coo_values);
319e334ad8fSJed Brown       user->matrices_set_up = true;
320e334ad8fSJed Brown       VecCreateSeq(PETSC_COMM_WORLD, ncoo, &coo_vec);
321e334ad8fSJed Brown     }
322e334ad8fSJed Brown   }
323e334ad8fSJed Brown   if (!J_pre_is_shell) {
324e334ad8fSJed Brown     CeedMemType mem_type = CEED_MEM_HOST;
325e334ad8fSJed Brown     const PetscScalar *values;
326e334ad8fSJed Brown     MatType mat_type;
327e334ad8fSJed Brown     PetscCall(MatGetType(J_pre, &mat_type));
328e334ad8fSJed Brown     //if (strstr(mat_type, "kokkos") || strstr(mat_type, "cusparse")) mem_type = CEED_MEM_DEVICE;
329e334ad8fSJed Brown     CeedOperatorLinearAssemble(user->op_ijacobian, user->coo_values);
330e334ad8fSJed Brown     CeedVectorGetArrayRead(user->coo_values, mem_type, &values);
331e334ad8fSJed Brown     if (coo_vec) {
332e334ad8fSJed Brown       VecPlaceArray(coo_vec, values);
333e334ad8fSJed Brown       VecViewFromOptions(coo_vec, NULL, "-coo_vec_view");
334e334ad8fSJed Brown       VecDestroy(&coo_vec);
335e334ad8fSJed Brown     }
336e334ad8fSJed Brown     PetscCall(MatSetValuesCOO(J_pre, values, INSERT_VALUES));
337e334ad8fSJed Brown     CeedVectorRestoreArrayRead(user->coo_values, &values);
338e334ad8fSJed Brown   }
339e334ad8fSJed Brown   PetscFunctionReturn(0);
340e334ad8fSJed Brown }
341e334ad8fSJed Brown 
34277841947SLeila Ghaffari // User provided TS Monitor
34377841947SLeila Ghaffari PetscErrorCode TSMonitor_NS(TS ts, PetscInt step_no, PetscReal time,
34477841947SLeila Ghaffari                             Vec Q, void *ctx) {
34577841947SLeila Ghaffari   User           user = ctx;
34677841947SLeila Ghaffari   Vec            Q_loc;
34777841947SLeila Ghaffari   char           file_path[PETSC_MAX_PATH_LEN];
34877841947SLeila Ghaffari   PetscViewer    viewer;
34977841947SLeila Ghaffari   PetscErrorCode ierr;
35077841947SLeila Ghaffari   PetscFunctionBeginUser;
35177841947SLeila Ghaffari 
35277841947SLeila Ghaffari   // Print every 'output_freq' steps
35377841947SLeila Ghaffari   if (step_no % user->app_ctx->output_freq != 0)
35477841947SLeila Ghaffari     PetscFunctionReturn(0);
35577841947SLeila Ghaffari 
35677841947SLeila Ghaffari   // Set up output
35777841947SLeila Ghaffari   ierr = DMGetLocalVector(user->dm, &Q_loc); CHKERRQ(ierr);
35877841947SLeila Ghaffari   ierr = PetscObjectSetName((PetscObject)Q_loc, "StateVec"); CHKERRQ(ierr);
35977841947SLeila Ghaffari   ierr = VecZeroEntries(Q_loc); CHKERRQ(ierr);
36077841947SLeila Ghaffari   ierr = DMGlobalToLocal(user->dm, Q, INSERT_VALUES, Q_loc); CHKERRQ(ierr);
36177841947SLeila Ghaffari 
36277841947SLeila Ghaffari   // Output
36308140895SJed Brown   ierr = PetscSNPrintf(file_path, sizeof file_path,
36408140895SJed Brown                        "%s/ns-%03" PetscInt_FMT ".vtu",
36577841947SLeila Ghaffari                        user->app_ctx->output_dir, step_no + user->app_ctx->cont_steps);
36677841947SLeila Ghaffari   CHKERRQ(ierr);
36777841947SLeila Ghaffari   ierr = PetscViewerVTKOpen(PetscObjectComm((PetscObject)Q), file_path,
36877841947SLeila Ghaffari                             FILE_MODE_WRITE, &viewer); CHKERRQ(ierr);
36977841947SLeila Ghaffari   ierr = VecView(Q_loc, viewer); CHKERRQ(ierr);
37077841947SLeila Ghaffari   ierr = PetscViewerDestroy(&viewer); CHKERRQ(ierr);
37177841947SLeila Ghaffari   if (user->dm_viz) {
37277841947SLeila Ghaffari     Vec         Q_refined, Q_refined_loc;
37377841947SLeila Ghaffari     char        file_path_refined[PETSC_MAX_PATH_LEN];
37477841947SLeila Ghaffari     PetscViewer viewer_refined;
37577841947SLeila Ghaffari 
37677841947SLeila Ghaffari     ierr = DMGetGlobalVector(user->dm_viz, &Q_refined); CHKERRQ(ierr);
37777841947SLeila Ghaffari     ierr = DMGetLocalVector(user->dm_viz, &Q_refined_loc); CHKERRQ(ierr);
37877841947SLeila Ghaffari     ierr = PetscObjectSetName((PetscObject)Q_refined_loc, "Refined");
37977841947SLeila Ghaffari     CHKERRQ(ierr);
38077841947SLeila Ghaffari     ierr = MatInterpolate(user->interp_viz, Q, Q_refined); CHKERRQ(ierr);
38177841947SLeila Ghaffari     ierr = VecZeroEntries(Q_refined_loc); CHKERRQ(ierr);
38277841947SLeila Ghaffari     ierr = DMGlobalToLocal(user->dm_viz, Q_refined, INSERT_VALUES, Q_refined_loc);
38377841947SLeila Ghaffari     CHKERRQ(ierr);
38477841947SLeila Ghaffari     ierr = PetscSNPrintf(file_path_refined, sizeof file_path_refined,
38508140895SJed Brown                          "%s/nsrefined-%03" PetscInt_FMT ".vtu", user->app_ctx->output_dir,
386e6225c47SLeila Ghaffari                          step_no + user->app_ctx->cont_steps);
38777841947SLeila Ghaffari     CHKERRQ(ierr);
38877841947SLeila Ghaffari     ierr = PetscViewerVTKOpen(PetscObjectComm((PetscObject)Q_refined),
389e6225c47SLeila Ghaffari                               file_path_refined, FILE_MODE_WRITE, &viewer_refined); CHKERRQ(ierr);
39077841947SLeila Ghaffari     ierr = VecView(Q_refined_loc, viewer_refined); CHKERRQ(ierr);
39177841947SLeila Ghaffari     ierr = DMRestoreLocalVector(user->dm_viz, &Q_refined_loc); CHKERRQ(ierr);
39277841947SLeila Ghaffari     ierr = DMRestoreGlobalVector(user->dm_viz, &Q_refined); CHKERRQ(ierr);
39377841947SLeila Ghaffari     ierr = PetscViewerDestroy(&viewer_refined); CHKERRQ(ierr);
39477841947SLeila Ghaffari   }
39577841947SLeila Ghaffari   ierr = DMRestoreLocalVector(user->dm, &Q_loc); CHKERRQ(ierr);
39677841947SLeila Ghaffari 
39777841947SLeila Ghaffari   // Save data in a binary file for continuation of simulations
39877841947SLeila Ghaffari   ierr = PetscSNPrintf(file_path, sizeof file_path, "%s/ns-solution.bin",
39977841947SLeila Ghaffari                        user->app_ctx->output_dir); CHKERRQ(ierr);
40077841947SLeila Ghaffari   ierr = PetscViewerBinaryOpen(user->comm, file_path, FILE_MODE_WRITE, &viewer);
40177841947SLeila Ghaffari   CHKERRQ(ierr);
40277841947SLeila Ghaffari   ierr = VecView(Q, viewer); CHKERRQ(ierr);
40377841947SLeila Ghaffari   ierr = PetscViewerDestroy(&viewer); CHKERRQ(ierr);
40477841947SLeila Ghaffari 
40577841947SLeila Ghaffari   // Save time stamp
40677841947SLeila Ghaffari   // Dimensionalize time back
40777841947SLeila Ghaffari   time /= user->units->second;
40877841947SLeila Ghaffari   ierr = PetscSNPrintf(file_path, sizeof file_path, "%s/ns-time.bin",
40977841947SLeila Ghaffari                        user->app_ctx->output_dir); CHKERRQ(ierr);
41077841947SLeila Ghaffari   ierr = PetscViewerBinaryOpen(user->comm, file_path, FILE_MODE_WRITE, &viewer);
41177841947SLeila Ghaffari   CHKERRQ(ierr);
41277841947SLeila Ghaffari   #if PETSC_VERSION_GE(3,13,0)
41377841947SLeila Ghaffari   ierr = PetscViewerBinaryWrite(viewer, &time, 1, PETSC_REAL);
41477841947SLeila Ghaffari   #else
41577841947SLeila Ghaffari   ierr = PetscViewerBinaryWrite(viewer, &time, 1, PETSC_REAL, true);
41677841947SLeila Ghaffari   #endif
41777841947SLeila Ghaffari   CHKERRQ(ierr);
41877841947SLeila Ghaffari   ierr = PetscViewerDestroy(&viewer); CHKERRQ(ierr);
41977841947SLeila Ghaffari 
42077841947SLeila Ghaffari   PetscFunctionReturn(0);
42177841947SLeila Ghaffari }
42277841947SLeila Ghaffari 
42377841947SLeila Ghaffari // TS: Create, setup, and solve
42477841947SLeila Ghaffari PetscErrorCode TSSolve_NS(DM dm, User user, AppCtx app_ctx, Physics phys,
42577841947SLeila Ghaffari                           Vec *Q, PetscScalar *f_time, TS *ts) {
42677841947SLeila Ghaffari   MPI_Comm       comm = user->comm;
42777841947SLeila Ghaffari   TSAdapt        adapt;
42877841947SLeila Ghaffari   PetscScalar    final_time;
42977841947SLeila Ghaffari   PetscErrorCode ierr;
43077841947SLeila Ghaffari   PetscFunctionBeginUser;
43177841947SLeila Ghaffari 
43277841947SLeila Ghaffari   ierr = TSCreate(comm, ts); CHKERRQ(ierr);
43377841947SLeila Ghaffari   ierr = TSSetDM(*ts, dm); CHKERRQ(ierr);
43477841947SLeila Ghaffari   if (phys->implicit) {
43577841947SLeila Ghaffari     ierr = TSSetType(*ts, TSBDF); CHKERRQ(ierr);
43677841947SLeila Ghaffari     if (user->op_ifunction) {
43777841947SLeila Ghaffari       ierr = TSSetIFunction(*ts, NULL, IFunction_NS, &user); CHKERRQ(ierr);
43877841947SLeila Ghaffari     } else { // Implicit integrators can fall back to using an RHSFunction
43977841947SLeila Ghaffari       ierr = TSSetRHSFunction(*ts, NULL, RHS_NS, &user); CHKERRQ(ierr);
44077841947SLeila Ghaffari     }
441e334ad8fSJed Brown     if (user->op_ijacobian) {
442e334ad8fSJed Brown       ierr = DMTSSetIJacobian(dm, FormIJacobian_NS, &user); CHKERRQ(ierr);
443e334ad8fSJed Brown     }
44477841947SLeila Ghaffari   } else {
44577841947SLeila Ghaffari     if (!user->op_rhs) SETERRQ(comm, PETSC_ERR_ARG_NULL,
44677841947SLeila Ghaffari                                  "Problem does not provide RHSFunction");
44777841947SLeila Ghaffari     ierr = TSSetType(*ts, TSRK); CHKERRQ(ierr);
44877841947SLeila Ghaffari     ierr = TSRKSetType(*ts, TSRK5F); CHKERRQ(ierr);
44977841947SLeila Ghaffari     ierr = TSSetRHSFunction(*ts, NULL, RHS_NS, &user); CHKERRQ(ierr);
45077841947SLeila Ghaffari   }
45177841947SLeila Ghaffari   ierr = TSSetMaxTime(*ts, 500. * user->units->second); CHKERRQ(ierr);
45277841947SLeila Ghaffari   ierr = TSSetExactFinalTime(*ts, TS_EXACTFINALTIME_STEPOVER); CHKERRQ(ierr);
45377841947SLeila Ghaffari   ierr = TSSetTimeStep(*ts, 1.e-2 * user->units->second); CHKERRQ(ierr);
45477841947SLeila Ghaffari   if (app_ctx->test_mode) {ierr = TSSetMaxSteps(*ts, 10); CHKERRQ(ierr);}
45577841947SLeila Ghaffari   ierr = TSGetAdapt(*ts, &adapt); CHKERRQ(ierr);
45677841947SLeila Ghaffari   ierr = TSAdaptSetStepLimits(adapt,
45777841947SLeila Ghaffari                               1.e-12 * user->units->second,
45877841947SLeila Ghaffari                               1.e2 * user->units->second); CHKERRQ(ierr);
45977841947SLeila Ghaffari   ierr = TSSetFromOptions(*ts); CHKERRQ(ierr);
460*5e82a6e1SJeremy L Thompson   user->time = -1.0; // require all BCs and ctx to be updated
461*5e82a6e1SJeremy L Thompson   user->dt   = -1.0;
46277841947SLeila Ghaffari   if (!app_ctx->cont_steps) { // print initial condition
46377841947SLeila Ghaffari     if (!app_ctx->test_mode) {
46477841947SLeila Ghaffari       ierr = TSMonitor_NS(*ts, 0, 0., *Q, user); CHKERRQ(ierr);
46577841947SLeila Ghaffari     }
46677841947SLeila Ghaffari   } else { // continue from time of last output
46777841947SLeila Ghaffari     PetscReal   time;
46877841947SLeila Ghaffari     PetscInt    count;
46977841947SLeila Ghaffari     PetscViewer viewer;
47077841947SLeila Ghaffari     char        file_path[PETSC_MAX_PATH_LEN];
47177841947SLeila Ghaffari     ierr = PetscSNPrintf(file_path, sizeof file_path, "%s/ns-time.bin",
47277841947SLeila Ghaffari                          app_ctx->output_dir); CHKERRQ(ierr);
47377841947SLeila Ghaffari     ierr = PetscViewerBinaryOpen(comm, file_path, FILE_MODE_READ, &viewer);
47477841947SLeila Ghaffari     CHKERRQ(ierr);
47577841947SLeila Ghaffari     ierr = PetscViewerBinaryRead(viewer, &time, 1, &count, PETSC_REAL);
47677841947SLeila Ghaffari     CHKERRQ(ierr);
47777841947SLeila Ghaffari     ierr = PetscViewerDestroy(&viewer); CHKERRQ(ierr);
47877841947SLeila Ghaffari     ierr = TSSetTime(*ts, time * user->units->second); CHKERRQ(ierr);
47977841947SLeila Ghaffari   }
48077841947SLeila Ghaffari   if (!app_ctx->test_mode) {
48177841947SLeila Ghaffari     ierr = TSMonitorSet(*ts, TSMonitor_NS, user, NULL); CHKERRQ(ierr);
48277841947SLeila Ghaffari   }
48377841947SLeila Ghaffari 
48477841947SLeila Ghaffari   // Solve
48577841947SLeila Ghaffari   double start, cpu_time_used;
48677841947SLeila Ghaffari   start = MPI_Wtime();
48777841947SLeila Ghaffari   ierr = PetscBarrier((PetscObject) *ts); CHKERRQ(ierr);
48877841947SLeila Ghaffari   ierr = TSSolve(*ts, *Q); CHKERRQ(ierr);
48977841947SLeila Ghaffari   cpu_time_used = MPI_Wtime() - start;
49077841947SLeila Ghaffari   ierr = TSGetSolveTime(*ts, &final_time); CHKERRQ(ierr);
49177841947SLeila Ghaffari   *f_time = final_time;
49277841947SLeila Ghaffari   ierr = MPI_Allreduce(MPI_IN_PLACE, &cpu_time_used, 1, MPI_DOUBLE, MPI_MIN,
49377841947SLeila Ghaffari                        comm); CHKERRQ(ierr);
49477841947SLeila Ghaffari   if (!app_ctx->test_mode) {
49577841947SLeila Ghaffari     ierr = PetscPrintf(PETSC_COMM_WORLD,
49677841947SLeila Ghaffari                        "Time taken for solution (sec): %g\n",
49777841947SLeila Ghaffari                        (double)cpu_time_used); CHKERRQ(ierr);
49877841947SLeila Ghaffari   }
49977841947SLeila Ghaffari   PetscFunctionReturn(0);
50077841947SLeila Ghaffari }
501