xref: /libCEED/examples/fluids/src/setupts.c (revision e334ad8fca0b75b5d678f79cfd10f4127eb21e08)
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 
8577841947SLeila Ghaffari   User           user = *(User *)user_data;
8677841947SLeila Ghaffari   PetscScalar    *q, *g;
8777841947SLeila Ghaffari   Vec            Q_loc, G_loc;
8877841947SLeila Ghaffari   PetscMemType   q_mem_type, g_mem_type;
8977841947SLeila Ghaffari   PetscErrorCode ierr;
9077841947SLeila Ghaffari   PetscFunctionBeginUser;
9177841947SLeila Ghaffari 
9288626eedSJames Wright   // Update context field labels
9311436a05SJames Wright   if (user->phys->solution_time_label)
9411436a05SJames Wright     CeedOperatorContextSetDouble(user->op_rhs, user->phys->solution_time_label, &t);
9588626eedSJames Wright   if (user->phys->timestep_size_label) {
9688626eedSJames Wright     PetscScalar dt;
9788626eedSJames Wright     ierr = TSGetTimeStep(ts,&dt); CHKERRQ(ierr);
9888626eedSJames Wright     CeedOperatorContextSetDouble(user->op_rhs, user->phys->timestep_size_label,
9988626eedSJames Wright                                  &dt);
10088626eedSJames Wright   }
10177841947SLeila Ghaffari 
10277841947SLeila Ghaffari   // Get local vectors
10377841947SLeila Ghaffari   ierr = DMGetLocalVector(user->dm, &Q_loc); CHKERRQ(ierr);
10477841947SLeila Ghaffari   ierr = DMGetLocalVector(user->dm, &G_loc); CHKERRQ(ierr);
10577841947SLeila Ghaffari 
10677841947SLeila Ghaffari   // Global-to-local
10777841947SLeila Ghaffari   ierr = VecZeroEntries(Q_loc); CHKERRQ(ierr);
10877841947SLeila Ghaffari   ierr = DMGlobalToLocal(user->dm, Q, INSERT_VALUES, Q_loc); CHKERRQ(ierr);
10977841947SLeila Ghaffari   ierr = DMPlexInsertBoundaryValues(user->dm, PETSC_TRUE, Q_loc, 0.0,
11077841947SLeila Ghaffari                                     NULL, NULL, NULL); CHKERRQ(ierr);
11177841947SLeila Ghaffari   ierr = VecZeroEntries(G_loc); CHKERRQ(ierr);
11277841947SLeila Ghaffari 
11377841947SLeila Ghaffari   // Place PETSc vectors in CEED vectors
11477841947SLeila Ghaffari   ierr = VecGetArrayReadAndMemType(Q_loc, (const PetscScalar **)&q, &q_mem_type);
11577841947SLeila Ghaffari   CHKERRQ(ierr);
11677841947SLeila Ghaffari   ierr = VecGetArrayAndMemType(G_loc, &g, &g_mem_type); CHKERRQ(ierr);
11777841947SLeila Ghaffari   CeedVectorSetArray(user->q_ceed, MemTypeP2C(q_mem_type), CEED_USE_POINTER, q);
11877841947SLeila Ghaffari   CeedVectorSetArray(user->g_ceed, MemTypeP2C(g_mem_type), CEED_USE_POINTER, g);
11977841947SLeila Ghaffari 
12077841947SLeila Ghaffari   // Apply CEED operator
12177841947SLeila Ghaffari   CeedOperatorApply(user->op_rhs, user->q_ceed, user->g_ceed,
12277841947SLeila Ghaffari                     CEED_REQUEST_IMMEDIATE);
12377841947SLeila Ghaffari 
12477841947SLeila Ghaffari   // Restore vectors
12577841947SLeila Ghaffari   CeedVectorTakeArray(user->q_ceed, MemTypeP2C(q_mem_type), NULL);
12677841947SLeila Ghaffari   CeedVectorTakeArray(user->g_ceed, MemTypeP2C(g_mem_type), NULL);
12777841947SLeila Ghaffari   ierr = VecRestoreArrayReadAndMemType(Q_loc, (const PetscScalar **)&q);
12877841947SLeila Ghaffari   CHKERRQ(ierr);
12977841947SLeila Ghaffari   ierr = VecRestoreArrayAndMemType(G_loc, &g); CHKERRQ(ierr);
13077841947SLeila Ghaffari 
13177841947SLeila Ghaffari   // Local-to-Global
13277841947SLeila Ghaffari   ierr = VecZeroEntries(G); CHKERRQ(ierr);
13377841947SLeila Ghaffari   ierr = DMLocalToGlobal(user->dm, G_loc, ADD_VALUES, G); CHKERRQ(ierr);
13477841947SLeila Ghaffari 
13577841947SLeila Ghaffari   // Inverse of the lumped mass matrix (M is Minv)
13677841947SLeila Ghaffari   ierr = VecPointwiseMult(G, G, user->M); CHKERRQ(ierr);
13777841947SLeila Ghaffari 
13877841947SLeila Ghaffari   // Restore vectors
13977841947SLeila Ghaffari   ierr = DMRestoreLocalVector(user->dm, &Q_loc); CHKERRQ(ierr);
14077841947SLeila Ghaffari   ierr = DMRestoreLocalVector(user->dm, &G_loc); CHKERRQ(ierr);
14177841947SLeila Ghaffari 
14277841947SLeila Ghaffari   PetscFunctionReturn(0);
14377841947SLeila Ghaffari }
14477841947SLeila Ghaffari 
14577841947SLeila Ghaffari // Implicit time-stepper function setup
14677841947SLeila Ghaffari PetscErrorCode IFunction_NS(TS ts, PetscReal t, Vec Q, Vec Q_dot, Vec G,
14777841947SLeila Ghaffari                             void *user_data) {
14877841947SLeila Ghaffari   User              user = *(User *)user_data;
14977841947SLeila Ghaffari   const PetscScalar *q, *q_dot;
15077841947SLeila Ghaffari   PetscScalar       *g;
15177841947SLeila Ghaffari   Vec               Q_loc, Q_dot_loc, G_loc;
15277841947SLeila Ghaffari   PetscMemType      q_mem_type, q_dot_mem_type, g_mem_type;
15377841947SLeila Ghaffari   PetscErrorCode    ierr;
15477841947SLeila Ghaffari   PetscFunctionBeginUser;
15577841947SLeila Ghaffari 
15688626eedSJames Wright   // Update context field labels
15711436a05SJames Wright   if (user->phys->solution_time_label)
15811436a05SJames Wright     CeedOperatorContextSetDouble(user->op_ifunction,
15911436a05SJames Wright                                  user->phys->solution_time_label, &t);
16088626eedSJames Wright   if (user->phys->timestep_size_label) {
16188626eedSJames Wright     PetscScalar dt;
16288626eedSJames Wright     ierr = TSGetTimeStep(ts,&dt); CHKERRQ(ierr);
16388626eedSJames Wright     CeedOperatorContextSetDouble(user->op_ifunction,
16488626eedSJames Wright                                  user->phys->timestep_size_label, &dt);
16588626eedSJames Wright   }
16677841947SLeila Ghaffari 
16777841947SLeila Ghaffari   // Get local vectors
16877841947SLeila Ghaffari   ierr = DMGetLocalVector(user->dm, &Q_loc); CHKERRQ(ierr);
16977841947SLeila Ghaffari   ierr = DMGetLocalVector(user->dm, &Q_dot_loc); CHKERRQ(ierr);
17077841947SLeila Ghaffari   ierr = DMGetLocalVector(user->dm, &G_loc); CHKERRQ(ierr);
17177841947SLeila Ghaffari 
17277841947SLeila Ghaffari   // Global-to-local
17377841947SLeila Ghaffari   ierr = VecZeroEntries(Q_loc); CHKERRQ(ierr);
17477841947SLeila Ghaffari   ierr = DMGlobalToLocal(user->dm, Q, INSERT_VALUES, Q_loc); CHKERRQ(ierr);
17577841947SLeila Ghaffari   ierr = DMPlexInsertBoundaryValues(user->dm, PETSC_TRUE, Q_loc, 0.0,
17677841947SLeila Ghaffari                                     NULL, NULL, NULL); CHKERRQ(ierr);
17777841947SLeila Ghaffari   ierr = VecZeroEntries(Q_dot_loc); CHKERRQ(ierr);
17877841947SLeila Ghaffari   ierr = DMGlobalToLocal(user->dm, Q_dot, INSERT_VALUES, Q_dot_loc);
17977841947SLeila Ghaffari   CHKERRQ(ierr);
18077841947SLeila Ghaffari   ierr = VecZeroEntries(G_loc); CHKERRQ(ierr);
18177841947SLeila Ghaffari 
18277841947SLeila Ghaffari   // Place PETSc vectors in CEED vectors
18377841947SLeila Ghaffari   ierr = VecGetArrayReadAndMemType(Q_loc, &q, &q_mem_type); CHKERRQ(ierr);
18477841947SLeila Ghaffari   ierr = VecGetArrayReadAndMemType(Q_dot_loc, &q_dot, &q_dot_mem_type);
18577841947SLeila Ghaffari   CHKERRQ(ierr);
18677841947SLeila Ghaffari   ierr = VecGetArrayAndMemType(G_loc, &g, &g_mem_type); CHKERRQ(ierr);
18777841947SLeila Ghaffari   CeedVectorSetArray(user->q_ceed, MemTypeP2C(q_mem_type), CEED_USE_POINTER,
18877841947SLeila Ghaffari                      (PetscScalar *)q);
18977841947SLeila Ghaffari   CeedVectorSetArray(user->q_dot_ceed, MemTypeP2C(q_dot_mem_type),
190e6225c47SLeila Ghaffari                      CEED_USE_POINTER, (PetscScalar *)q_dot);
19177841947SLeila Ghaffari   CeedVectorSetArray(user->g_ceed, MemTypeP2C(g_mem_type), CEED_USE_POINTER, g);
19277841947SLeila Ghaffari 
19377841947SLeila Ghaffari   // Apply CEED operator
19477841947SLeila Ghaffari   CeedOperatorApply(user->op_ifunction, user->q_ceed, user->g_ceed,
19577841947SLeila Ghaffari                     CEED_REQUEST_IMMEDIATE);
19677841947SLeila Ghaffari 
19777841947SLeila Ghaffari   // Restore vectors
19877841947SLeila Ghaffari   CeedVectorTakeArray(user->q_ceed, MemTypeP2C(q_mem_type), NULL);
19977841947SLeila Ghaffari   CeedVectorTakeArray(user->q_dot_ceed, MemTypeP2C(q_dot_mem_type), NULL);
20077841947SLeila Ghaffari   CeedVectorTakeArray(user->g_ceed, MemTypeP2C(g_mem_type), NULL);
20177841947SLeila Ghaffari   ierr = VecRestoreArrayReadAndMemType(Q_loc, &q); CHKERRQ(ierr);
20277841947SLeila Ghaffari   ierr = VecRestoreArrayReadAndMemType(Q_dot_loc, &q_dot); CHKERRQ(ierr);
20377841947SLeila Ghaffari   ierr = VecRestoreArrayAndMemType(G_loc, &g); CHKERRQ(ierr);
20477841947SLeila Ghaffari 
20577841947SLeila Ghaffari   // Local-to-Global
20677841947SLeila Ghaffari   ierr = VecZeroEntries(G); CHKERRQ(ierr);
20777841947SLeila Ghaffari   ierr = DMLocalToGlobal(user->dm, G_loc, ADD_VALUES, G); CHKERRQ(ierr);
20877841947SLeila Ghaffari 
20977841947SLeila Ghaffari   // Restore vectors
21077841947SLeila Ghaffari   ierr = DMRestoreLocalVector(user->dm, &Q_loc); CHKERRQ(ierr);
21177841947SLeila Ghaffari   ierr = DMRestoreLocalVector(user->dm, &Q_dot_loc); CHKERRQ(ierr);
21277841947SLeila Ghaffari   ierr = DMRestoreLocalVector(user->dm, &G_loc); CHKERRQ(ierr);
21377841947SLeila Ghaffari 
21477841947SLeila Ghaffari   PetscFunctionReturn(0);
21577841947SLeila Ghaffari }
21677841947SLeila Ghaffari 
217*e334ad8fSJed Brown static PetscErrorCode MatMult_NS_IJacobian(Mat J, Vec Q, Vec G) {
218*e334ad8fSJed Brown   User user;
219*e334ad8fSJed Brown   const PetscScalar *q;
220*e334ad8fSJed Brown   PetscScalar       *g;
221*e334ad8fSJed Brown   Vec               Q_loc, G_loc;
222*e334ad8fSJed Brown   PetscMemType      q_mem_type, g_mem_type;
223*e334ad8fSJed Brown   PetscErrorCode    ierr;
224*e334ad8fSJed Brown   PetscFunctionBeginUser;
225*e334ad8fSJed Brown   MatShellGetContext(J, &user);
226*e334ad8fSJed Brown   // Get local vectors
227*e334ad8fSJed Brown   ierr = DMGetLocalVector(user->dm, &Q_loc); CHKERRQ(ierr);
228*e334ad8fSJed Brown   ierr = DMGetLocalVector(user->dm, &G_loc); CHKERRQ(ierr);
229*e334ad8fSJed Brown 
230*e334ad8fSJed Brown   // Global-to-local
231*e334ad8fSJed Brown   ierr = VecZeroEntries(Q_loc); CHKERRQ(ierr);
232*e334ad8fSJed Brown   ierr = DMGlobalToLocal(user->dm, Q, INSERT_VALUES, Q_loc); CHKERRQ(ierr);
233*e334ad8fSJed Brown   ierr = VecZeroEntries(G_loc); CHKERRQ(ierr);
234*e334ad8fSJed Brown 
235*e334ad8fSJed Brown   // Place PETSc vectors in CEED vectors
236*e334ad8fSJed Brown   ierr = VecGetArrayReadAndMemType(Q_loc, &q, &q_mem_type); CHKERRQ(ierr);
237*e334ad8fSJed Brown   ierr = VecGetArrayAndMemType(G_loc, &g, &g_mem_type); CHKERRQ(ierr);
238*e334ad8fSJed Brown   CeedVectorSetArray(user->q_ceed, MemTypeP2C(q_mem_type), CEED_USE_POINTER,
239*e334ad8fSJed Brown                      (PetscScalar *)q);
240*e334ad8fSJed Brown   CeedVectorSetArray(user->g_ceed, MemTypeP2C(g_mem_type), CEED_USE_POINTER, g);
241*e334ad8fSJed Brown 
242*e334ad8fSJed Brown   // Apply CEED operator
243*e334ad8fSJed Brown   CeedOperatorApply(user->op_ijacobian, user->q_ceed, user->g_ceed,
244*e334ad8fSJed Brown                     CEED_REQUEST_IMMEDIATE);
245*e334ad8fSJed Brown 
246*e334ad8fSJed Brown   // Restore vectors
247*e334ad8fSJed Brown   CeedVectorTakeArray(user->q_ceed, MemTypeP2C(q_mem_type), NULL);
248*e334ad8fSJed Brown   CeedVectorTakeArray(user->g_ceed, MemTypeP2C(g_mem_type), NULL);
249*e334ad8fSJed Brown   ierr = VecRestoreArrayReadAndMemType(Q_loc, &q); CHKERRQ(ierr);
250*e334ad8fSJed Brown   ierr = VecRestoreArrayAndMemType(G_loc, &g); CHKERRQ(ierr);
251*e334ad8fSJed Brown 
252*e334ad8fSJed Brown   // Local-to-Global
253*e334ad8fSJed Brown   ierr = VecZeroEntries(G); CHKERRQ(ierr);
254*e334ad8fSJed Brown   ierr = DMLocalToGlobal(user->dm, G_loc, ADD_VALUES, G); CHKERRQ(ierr);
255*e334ad8fSJed Brown 
256*e334ad8fSJed Brown   // Restore vectors
257*e334ad8fSJed Brown   ierr = DMRestoreLocalVector(user->dm, &Q_loc); CHKERRQ(ierr);
258*e334ad8fSJed Brown   ierr = DMRestoreLocalVector(user->dm, &G_loc); CHKERRQ(ierr);
259*e334ad8fSJed Brown   PetscFunctionReturn(0);
260*e334ad8fSJed Brown }
261*e334ad8fSJed Brown 
262*e334ad8fSJed Brown PetscErrorCode MatGetDiagonal_NS_IJacobian(Mat A, Vec D) {
263*e334ad8fSJed Brown   User user;
264*e334ad8fSJed Brown   Vec D_loc;
265*e334ad8fSJed Brown   PetscScalar *d;
266*e334ad8fSJed Brown   PetscMemType mem_type;
267*e334ad8fSJed Brown 
268*e334ad8fSJed Brown   PetscFunctionBeginUser;
269*e334ad8fSJed Brown   MatShellGetContext(A, &user);
270*e334ad8fSJed Brown   PetscCall(DMGetLocalVector(user->dm, &D_loc));
271*e334ad8fSJed Brown   PetscCall(VecGetArrayAndMemType(D_loc, &d, &mem_type));
272*e334ad8fSJed Brown   CeedVectorSetArray(user->g_ceed, MemTypeP2C(mem_type), CEED_USE_POINTER, d);
273*e334ad8fSJed Brown   CeedOperatorLinearAssembleDiagonal(user->op_ijacobian, user->g_ceed,
274*e334ad8fSJed Brown                                      CEED_REQUEST_IMMEDIATE);
275*e334ad8fSJed Brown   CeedVectorTakeArray(user->g_ceed, MemTypeP2C(mem_type), NULL);
276*e334ad8fSJed Brown   PetscCall(VecRestoreArrayAndMemType(D_loc, &d));
277*e334ad8fSJed Brown   PetscCall(VecZeroEntries(D));
278*e334ad8fSJed Brown   PetscCall(DMLocalToGlobal(user->dm, D_loc, ADD_VALUES, D));
279*e334ad8fSJed Brown   PetscCall(DMRestoreLocalVector(user->dm, &D_loc));
280*e334ad8fSJed Brown   VecViewFromOptions(D, NULL, "-diag_vec_view");
281*e334ad8fSJed Brown   PetscFunctionReturn(0);
282*e334ad8fSJed Brown }
283*e334ad8fSJed Brown 
284*e334ad8fSJed Brown PetscErrorCode FormIJacobian_NS(TS ts, PetscReal t, Vec Q, Vec Q_dot,
285*e334ad8fSJed Brown                                 PetscReal shift, Mat J, Mat J_pre,
286*e334ad8fSJed Brown                                 void *user_data) {
287*e334ad8fSJed Brown   User user = *(User *)user_data;
288*e334ad8fSJed Brown   PetscBool J_is_shell, J_pre_is_shell;
289*e334ad8fSJed Brown   PetscFunctionBeginUser;
290*e334ad8fSJed Brown   if (user->phys->ijacobian_time_shift_label)
291*e334ad8fSJed Brown     CeedOperatorContextSetDouble(user->op_ijacobian,
292*e334ad8fSJed Brown                                  user->phys->ijacobian_time_shift_label, &shift);
293*e334ad8fSJed Brown   PetscCall(MatAssemblyBegin(J, MAT_FINAL_ASSEMBLY));
294*e334ad8fSJed Brown   PetscCall(MatAssemblyEnd(J, MAT_FINAL_ASSEMBLY));
295*e334ad8fSJed Brown   Vec coo_vec = NULL;
296*e334ad8fSJed Brown   PetscCall(PetscObjectTypeCompare((PetscObject)J, MATSHELL, &J_is_shell));
297*e334ad8fSJed Brown   PetscCall(PetscObjectTypeCompare((PetscObject)J_pre, MATSHELL,
298*e334ad8fSJed Brown                                    &J_pre_is_shell));
299*e334ad8fSJed Brown   if (!user->matrices_set_up) {
300*e334ad8fSJed Brown     if (J_is_shell) {
301*e334ad8fSJed Brown       PetscCall(MatShellSetContext(J, user));
302*e334ad8fSJed Brown       PetscCall(MatShellSetOperation(J, MATOP_MULT,
303*e334ad8fSJed Brown                                      (void (*)(void))MatMult_NS_IJacobian));
304*e334ad8fSJed Brown       PetscCall(MatShellSetOperation(J, MATOP_GET_DIAGONAL,
305*e334ad8fSJed Brown                                      (void (*)(void))MatGetDiagonal_NS_IJacobian));
306*e334ad8fSJed Brown       PetscCall(MatSetUp(J));
307*e334ad8fSJed Brown     }
308*e334ad8fSJed Brown     if (!J_pre_is_shell) {
309*e334ad8fSJed Brown       PetscCount ncoo;
310*e334ad8fSJed Brown       PetscInt *rows, *cols;
311*e334ad8fSJed Brown       PetscCall(CeedOperatorLinearAssembleSymbolic(user->op_ijacobian, &ncoo, &rows,
312*e334ad8fSJed Brown                 &cols));
313*e334ad8fSJed Brown       PetscCall(MatSetPreallocationCOOLocal(J_pre, ncoo, rows, cols));
314*e334ad8fSJed Brown       free(rows);
315*e334ad8fSJed Brown       free(cols);
316*e334ad8fSJed Brown       CeedVectorCreate(user->ceed, ncoo, &user->coo_values);
317*e334ad8fSJed Brown       user->matrices_set_up = true;
318*e334ad8fSJed Brown       VecCreateSeq(PETSC_COMM_WORLD, ncoo, &coo_vec);
319*e334ad8fSJed Brown     }
320*e334ad8fSJed Brown   }
321*e334ad8fSJed Brown   if (!J_pre_is_shell) {
322*e334ad8fSJed Brown     CeedMemType mem_type = CEED_MEM_HOST;
323*e334ad8fSJed Brown     const PetscScalar *values;
324*e334ad8fSJed Brown     MatType mat_type;
325*e334ad8fSJed Brown     PetscCall(MatGetType(J_pre, &mat_type));
326*e334ad8fSJed Brown     //if (strstr(mat_type, "kokkos") || strstr(mat_type, "cusparse")) mem_type = CEED_MEM_DEVICE;
327*e334ad8fSJed Brown     CeedOperatorLinearAssemble(user->op_ijacobian, user->coo_values);
328*e334ad8fSJed Brown     CeedVectorGetArrayRead(user->coo_values, mem_type, &values);
329*e334ad8fSJed Brown     if (coo_vec) {
330*e334ad8fSJed Brown       VecPlaceArray(coo_vec, values);
331*e334ad8fSJed Brown       VecViewFromOptions(coo_vec, NULL, "-coo_vec_view");
332*e334ad8fSJed Brown       VecDestroy(&coo_vec);
333*e334ad8fSJed Brown     }
334*e334ad8fSJed Brown     PetscCall(MatSetValuesCOO(J_pre, values, INSERT_VALUES));
335*e334ad8fSJed Brown     CeedVectorRestoreArrayRead(user->coo_values, &values);
336*e334ad8fSJed Brown   }
337*e334ad8fSJed Brown   PetscFunctionReturn(0);
338*e334ad8fSJed Brown }
339*e334ad8fSJed Brown 
34077841947SLeila Ghaffari // User provided TS Monitor
34177841947SLeila Ghaffari PetscErrorCode TSMonitor_NS(TS ts, PetscInt step_no, PetscReal time,
34277841947SLeila Ghaffari                             Vec Q, void *ctx) {
34377841947SLeila Ghaffari   User           user = ctx;
34477841947SLeila Ghaffari   Vec            Q_loc;
34577841947SLeila Ghaffari   char           file_path[PETSC_MAX_PATH_LEN];
34677841947SLeila Ghaffari   PetscViewer    viewer;
34777841947SLeila Ghaffari   PetscErrorCode ierr;
34877841947SLeila Ghaffari   PetscFunctionBeginUser;
34977841947SLeila Ghaffari 
35077841947SLeila Ghaffari   // Print every 'output_freq' steps
35177841947SLeila Ghaffari   if (step_no % user->app_ctx->output_freq != 0)
35277841947SLeila Ghaffari     PetscFunctionReturn(0);
35377841947SLeila Ghaffari 
35477841947SLeila Ghaffari   // Set up output
35577841947SLeila Ghaffari   ierr = DMGetLocalVector(user->dm, &Q_loc); CHKERRQ(ierr);
35677841947SLeila Ghaffari   ierr = PetscObjectSetName((PetscObject)Q_loc, "StateVec"); CHKERRQ(ierr);
35777841947SLeila Ghaffari   ierr = VecZeroEntries(Q_loc); CHKERRQ(ierr);
35877841947SLeila Ghaffari   ierr = DMGlobalToLocal(user->dm, Q, INSERT_VALUES, Q_loc); CHKERRQ(ierr);
35977841947SLeila Ghaffari 
36077841947SLeila Ghaffari   // Output
36108140895SJed Brown   ierr = PetscSNPrintf(file_path, sizeof file_path,
36208140895SJed Brown                        "%s/ns-%03" PetscInt_FMT ".vtu",
36377841947SLeila Ghaffari                        user->app_ctx->output_dir, step_no + user->app_ctx->cont_steps);
36477841947SLeila Ghaffari   CHKERRQ(ierr);
36577841947SLeila Ghaffari   ierr = PetscViewerVTKOpen(PetscObjectComm((PetscObject)Q), file_path,
36677841947SLeila Ghaffari                             FILE_MODE_WRITE, &viewer); CHKERRQ(ierr);
36777841947SLeila Ghaffari   ierr = VecView(Q_loc, viewer); CHKERRQ(ierr);
36877841947SLeila Ghaffari   ierr = PetscViewerDestroy(&viewer); CHKERRQ(ierr);
36977841947SLeila Ghaffari   if (user->dm_viz) {
37077841947SLeila Ghaffari     Vec         Q_refined, Q_refined_loc;
37177841947SLeila Ghaffari     char        file_path_refined[PETSC_MAX_PATH_LEN];
37277841947SLeila Ghaffari     PetscViewer viewer_refined;
37377841947SLeila Ghaffari 
37477841947SLeila Ghaffari     ierr = DMGetGlobalVector(user->dm_viz, &Q_refined); CHKERRQ(ierr);
37577841947SLeila Ghaffari     ierr = DMGetLocalVector(user->dm_viz, &Q_refined_loc); CHKERRQ(ierr);
37677841947SLeila Ghaffari     ierr = PetscObjectSetName((PetscObject)Q_refined_loc, "Refined");
37777841947SLeila Ghaffari     CHKERRQ(ierr);
37877841947SLeila Ghaffari     ierr = MatInterpolate(user->interp_viz, Q, Q_refined); CHKERRQ(ierr);
37977841947SLeila Ghaffari     ierr = VecZeroEntries(Q_refined_loc); CHKERRQ(ierr);
38077841947SLeila Ghaffari     ierr = DMGlobalToLocal(user->dm_viz, Q_refined, INSERT_VALUES, Q_refined_loc);
38177841947SLeila Ghaffari     CHKERRQ(ierr);
38277841947SLeila Ghaffari     ierr = PetscSNPrintf(file_path_refined, sizeof file_path_refined,
38308140895SJed Brown                          "%s/nsrefined-%03" PetscInt_FMT ".vtu", user->app_ctx->output_dir,
384e6225c47SLeila Ghaffari                          step_no + user->app_ctx->cont_steps);
38577841947SLeila Ghaffari     CHKERRQ(ierr);
38677841947SLeila Ghaffari     ierr = PetscViewerVTKOpen(PetscObjectComm((PetscObject)Q_refined),
387e6225c47SLeila Ghaffari                               file_path_refined, FILE_MODE_WRITE, &viewer_refined); CHKERRQ(ierr);
38877841947SLeila Ghaffari     ierr = VecView(Q_refined_loc, viewer_refined); CHKERRQ(ierr);
38977841947SLeila Ghaffari     ierr = DMRestoreLocalVector(user->dm_viz, &Q_refined_loc); CHKERRQ(ierr);
39077841947SLeila Ghaffari     ierr = DMRestoreGlobalVector(user->dm_viz, &Q_refined); CHKERRQ(ierr);
39177841947SLeila Ghaffari     ierr = PetscViewerDestroy(&viewer_refined); CHKERRQ(ierr);
39277841947SLeila Ghaffari   }
39377841947SLeila Ghaffari   ierr = DMRestoreLocalVector(user->dm, &Q_loc); CHKERRQ(ierr);
39477841947SLeila Ghaffari 
39577841947SLeila Ghaffari   // Save data in a binary file for continuation of simulations
39677841947SLeila Ghaffari   ierr = PetscSNPrintf(file_path, sizeof file_path, "%s/ns-solution.bin",
39777841947SLeila Ghaffari                        user->app_ctx->output_dir); CHKERRQ(ierr);
39877841947SLeila Ghaffari   ierr = PetscViewerBinaryOpen(user->comm, file_path, FILE_MODE_WRITE, &viewer);
39977841947SLeila Ghaffari   CHKERRQ(ierr);
40077841947SLeila Ghaffari   ierr = VecView(Q, viewer); CHKERRQ(ierr);
40177841947SLeila Ghaffari   ierr = PetscViewerDestroy(&viewer); CHKERRQ(ierr);
40277841947SLeila Ghaffari 
40377841947SLeila Ghaffari   // Save time stamp
40477841947SLeila Ghaffari   // Dimensionalize time back
40577841947SLeila Ghaffari   time /= user->units->second;
40677841947SLeila Ghaffari   ierr = PetscSNPrintf(file_path, sizeof file_path, "%s/ns-time.bin",
40777841947SLeila Ghaffari                        user->app_ctx->output_dir); CHKERRQ(ierr);
40877841947SLeila Ghaffari   ierr = PetscViewerBinaryOpen(user->comm, file_path, FILE_MODE_WRITE, &viewer);
40977841947SLeila Ghaffari   CHKERRQ(ierr);
41077841947SLeila Ghaffari   #if PETSC_VERSION_GE(3,13,0)
41177841947SLeila Ghaffari   ierr = PetscViewerBinaryWrite(viewer, &time, 1, PETSC_REAL);
41277841947SLeila Ghaffari   #else
41377841947SLeila Ghaffari   ierr = PetscViewerBinaryWrite(viewer, &time, 1, PETSC_REAL, true);
41477841947SLeila Ghaffari   #endif
41577841947SLeila Ghaffari   CHKERRQ(ierr);
41677841947SLeila Ghaffari   ierr = PetscViewerDestroy(&viewer); CHKERRQ(ierr);
41777841947SLeila Ghaffari 
41877841947SLeila Ghaffari   PetscFunctionReturn(0);
41977841947SLeila Ghaffari }
42077841947SLeila Ghaffari 
42177841947SLeila Ghaffari // TS: Create, setup, and solve
42277841947SLeila Ghaffari PetscErrorCode TSSolve_NS(DM dm, User user, AppCtx app_ctx, Physics phys,
42377841947SLeila Ghaffari                           Vec *Q, PetscScalar *f_time, TS *ts) {
42477841947SLeila Ghaffari   MPI_Comm       comm = user->comm;
42577841947SLeila Ghaffari   TSAdapt        adapt;
42677841947SLeila Ghaffari   PetscScalar    final_time;
42777841947SLeila Ghaffari   PetscErrorCode ierr;
42877841947SLeila Ghaffari   PetscFunctionBeginUser;
42977841947SLeila Ghaffari 
43077841947SLeila Ghaffari   ierr = TSCreate(comm, ts); CHKERRQ(ierr);
43177841947SLeila Ghaffari   ierr = TSSetDM(*ts, dm); CHKERRQ(ierr);
43277841947SLeila Ghaffari   if (phys->implicit) {
43377841947SLeila Ghaffari     ierr = TSSetType(*ts, TSBDF); CHKERRQ(ierr);
43477841947SLeila Ghaffari     if (user->op_ifunction) {
43577841947SLeila Ghaffari       ierr = TSSetIFunction(*ts, NULL, IFunction_NS, &user); CHKERRQ(ierr);
43677841947SLeila Ghaffari     } else { // Implicit integrators can fall back to using an RHSFunction
43777841947SLeila Ghaffari       ierr = TSSetRHSFunction(*ts, NULL, RHS_NS, &user); CHKERRQ(ierr);
43877841947SLeila Ghaffari     }
439*e334ad8fSJed Brown     if (user->op_ijacobian) {
440*e334ad8fSJed Brown       ierr = DMTSSetIJacobian(dm, FormIJacobian_NS, &user); CHKERRQ(ierr);
441*e334ad8fSJed Brown     }
44277841947SLeila Ghaffari   } else {
44377841947SLeila Ghaffari     if (!user->op_rhs) SETERRQ(comm, PETSC_ERR_ARG_NULL,
44477841947SLeila Ghaffari                                  "Problem does not provide RHSFunction");
44577841947SLeila Ghaffari     ierr = TSSetType(*ts, TSRK); CHKERRQ(ierr);
44677841947SLeila Ghaffari     ierr = TSRKSetType(*ts, TSRK5F); CHKERRQ(ierr);
44777841947SLeila Ghaffari     ierr = TSSetRHSFunction(*ts, NULL, RHS_NS, &user); CHKERRQ(ierr);
44877841947SLeila Ghaffari   }
44977841947SLeila Ghaffari   ierr = TSSetMaxTime(*ts, 500. * user->units->second); CHKERRQ(ierr);
45077841947SLeila Ghaffari   ierr = TSSetExactFinalTime(*ts, TS_EXACTFINALTIME_STEPOVER); CHKERRQ(ierr);
45177841947SLeila Ghaffari   ierr = TSSetTimeStep(*ts, 1.e-2 * user->units->second); CHKERRQ(ierr);
45277841947SLeila Ghaffari   if (app_ctx->test_mode) {ierr = TSSetMaxSteps(*ts, 10); CHKERRQ(ierr);}
45377841947SLeila Ghaffari   ierr = TSGetAdapt(*ts, &adapt); CHKERRQ(ierr);
45477841947SLeila Ghaffari   ierr = TSAdaptSetStepLimits(adapt,
45577841947SLeila Ghaffari                               1.e-12 * user->units->second,
45677841947SLeila Ghaffari                               1.e2 * user->units->second); CHKERRQ(ierr);
45777841947SLeila Ghaffari   ierr = TSSetFromOptions(*ts); CHKERRQ(ierr);
45877841947SLeila Ghaffari   if (!app_ctx->cont_steps) { // print initial condition
45977841947SLeila Ghaffari     if (!app_ctx->test_mode) {
46077841947SLeila Ghaffari       ierr = TSMonitor_NS(*ts, 0, 0., *Q, user); CHKERRQ(ierr);
46177841947SLeila Ghaffari     }
46277841947SLeila Ghaffari   } else { // continue from time of last output
46377841947SLeila Ghaffari     PetscReal   time;
46477841947SLeila Ghaffari     PetscInt    count;
46577841947SLeila Ghaffari     PetscViewer viewer;
46677841947SLeila Ghaffari     char        file_path[PETSC_MAX_PATH_LEN];
46777841947SLeila Ghaffari     ierr = PetscSNPrintf(file_path, sizeof file_path, "%s/ns-time.bin",
46877841947SLeila Ghaffari                          app_ctx->output_dir); CHKERRQ(ierr);
46977841947SLeila Ghaffari     ierr = PetscViewerBinaryOpen(comm, file_path, FILE_MODE_READ, &viewer);
47077841947SLeila Ghaffari     CHKERRQ(ierr);
47177841947SLeila Ghaffari     ierr = PetscViewerBinaryRead(viewer, &time, 1, &count, PETSC_REAL);
47277841947SLeila Ghaffari     CHKERRQ(ierr);
47377841947SLeila Ghaffari     ierr = PetscViewerDestroy(&viewer); CHKERRQ(ierr);
47477841947SLeila Ghaffari     ierr = TSSetTime(*ts, time * user->units->second); CHKERRQ(ierr);
47577841947SLeila Ghaffari   }
47677841947SLeila Ghaffari   if (!app_ctx->test_mode) {
47777841947SLeila Ghaffari     ierr = TSMonitorSet(*ts, TSMonitor_NS, user, NULL); CHKERRQ(ierr);
47877841947SLeila Ghaffari   }
47977841947SLeila Ghaffari 
48077841947SLeila Ghaffari   // Solve
48177841947SLeila Ghaffari   double start, cpu_time_used;
48277841947SLeila Ghaffari   start = MPI_Wtime();
48377841947SLeila Ghaffari   ierr = PetscBarrier((PetscObject) *ts); CHKERRQ(ierr);
48477841947SLeila Ghaffari   ierr = TSSolve(*ts, *Q); CHKERRQ(ierr);
48577841947SLeila Ghaffari   cpu_time_used = MPI_Wtime() - start;
48677841947SLeila Ghaffari   ierr = TSGetSolveTime(*ts, &final_time); CHKERRQ(ierr);
48777841947SLeila Ghaffari   *f_time = final_time;
48877841947SLeila Ghaffari   ierr = MPI_Allreduce(MPI_IN_PLACE, &cpu_time_used, 1, MPI_DOUBLE, MPI_MIN,
48977841947SLeila Ghaffari                        comm); CHKERRQ(ierr);
49077841947SLeila Ghaffari   if (!app_ctx->test_mode) {
49177841947SLeila Ghaffari     ierr = PetscPrintf(PETSC_COMM_WORLD,
49277841947SLeila Ghaffari                        "Time taken for solution (sec): %g\n",
49377841947SLeila Ghaffari                        (double)cpu_time_used); CHKERRQ(ierr);
49477841947SLeila Ghaffari   }
49577841947SLeila Ghaffari   PetscFunctionReturn(0);
49677841947SLeila Ghaffari }
497