xref: /honee/src/setupts.c (revision f0b653724bc7ff3458409f59c4d1e4b1ef161295)
1727da7e7SJeremy L Thompson // Copyright (c) 2017-2022, Lawrence Livermore National Security, LLC and other CEED contributors.
2727da7e7SJeremy L Thompson // All Rights Reserved. See the top-level LICENSE and NOTICE files for details.
3a515125bSLeila Ghaffari //
4727da7e7SJeremy L Thompson // SPDX-License-Identifier: BSD-2-Clause
5a515125bSLeila Ghaffari //
6727da7e7SJeremy L Thompson // This file is part of CEED:  http://github.com/ceed
7a515125bSLeila Ghaffari 
8a515125bSLeila Ghaffari /// @file
9a515125bSLeila Ghaffari /// Time-stepping functions for Navier-Stokes example using PETSc
10a515125bSLeila Ghaffari 
11a515125bSLeila Ghaffari #include "../navierstokes.h"
12a515125bSLeila Ghaffari #include "../qfunctions/mass.h"
13a515125bSLeila Ghaffari 
14a515125bSLeila Ghaffari // Compute mass matrix for explicit scheme
15a515125bSLeila Ghaffari PetscErrorCode ComputeLumpedMassMatrix(Ceed ceed, DM dm, CeedData ceed_data,
16a515125bSLeila Ghaffari                                        Vec M) {
17a515125bSLeila Ghaffari   Vec            M_loc;
18a515125bSLeila Ghaffari   CeedQFunction  qf_mass;
19a515125bSLeila Ghaffari   CeedOperator   op_mass;
20a515125bSLeila Ghaffari   CeedVector     m_ceed, ones_vec;
21a515125bSLeila Ghaffari   CeedInt        num_comp_q, q_data_size;
22a515125bSLeila Ghaffari   PetscErrorCode ierr;
23a515125bSLeila Ghaffari   PetscFunctionBeginUser;
24a515125bSLeila Ghaffari 
25a515125bSLeila Ghaffari   // CEED Restriction
26a515125bSLeila Ghaffari   CeedElemRestrictionGetNumComponents(ceed_data->elem_restr_q, &num_comp_q);
27a515125bSLeila Ghaffari   CeedElemRestrictionGetNumComponents(ceed_data->elem_restr_qd_i, &q_data_size);
28a515125bSLeila Ghaffari   CeedElemRestrictionCreateVector(ceed_data->elem_restr_q, &m_ceed, NULL);
29a515125bSLeila Ghaffari   CeedElemRestrictionCreateVector(ceed_data->elem_restr_q, &ones_vec, NULL);
30a515125bSLeila Ghaffari   CeedVectorSetValue(ones_vec, 1.0);
31a515125bSLeila Ghaffari 
32a515125bSLeila Ghaffari   // CEED QFunction
33a515125bSLeila Ghaffari   CeedQFunctionCreateInterior(ceed, 1, Mass, Mass_loc, &qf_mass);
34a515125bSLeila Ghaffari   CeedQFunctionAddInput(qf_mass, "q", num_comp_q, CEED_EVAL_INTERP);
353b1e209bSJeremy L Thompson   CeedQFunctionAddInput(qf_mass, "qdata", q_data_size, CEED_EVAL_NONE);
36a515125bSLeila Ghaffari   CeedQFunctionAddOutput(qf_mass, "v", num_comp_q, CEED_EVAL_INTERP);
37a515125bSLeila Ghaffari 
38a515125bSLeila Ghaffari   // CEED Operator
39a515125bSLeila Ghaffari   CeedOperatorCreate(ceed, qf_mass, NULL, NULL, &op_mass);
40a515125bSLeila Ghaffari   CeedOperatorSetField(op_mass, "q", ceed_data->elem_restr_q, ceed_data->basis_q,
41a515125bSLeila Ghaffari                        CEED_VECTOR_ACTIVE);
423b1e209bSJeremy L Thompson   CeedOperatorSetField(op_mass, "qdata", ceed_data->elem_restr_qd_i,
43a515125bSLeila Ghaffari                        CEED_BASIS_COLLOCATED, ceed_data->q_data);
44a515125bSLeila Ghaffari   CeedOperatorSetField(op_mass, "v", ceed_data->elem_restr_q, ceed_data->basis_q,
45a515125bSLeila Ghaffari                        CEED_VECTOR_ACTIVE);
46a515125bSLeila Ghaffari 
47a515125bSLeila Ghaffari   // Place PETSc vector in CEED vector
48a515125bSLeila Ghaffari   CeedScalar *m;
49a515125bSLeila Ghaffari   PetscMemType m_mem_type;
50a515125bSLeila Ghaffari   ierr = DMGetLocalVector(dm, &M_loc); CHKERRQ(ierr);
51a515125bSLeila Ghaffari   ierr = VecGetArrayAndMemType(M_loc, (PetscScalar **)&m, &m_mem_type);
52a515125bSLeila Ghaffari   CHKERRQ(ierr);
53a515125bSLeila Ghaffari   CeedVectorSetArray(m_ceed, MemTypeP2C(m_mem_type), CEED_USE_POINTER, m);
54a515125bSLeila Ghaffari 
55a515125bSLeila Ghaffari   // Apply CEED Operator
56a515125bSLeila Ghaffari   CeedOperatorApply(op_mass, ones_vec, m_ceed, CEED_REQUEST_IMMEDIATE);
57a515125bSLeila Ghaffari 
58a515125bSLeila Ghaffari   // Restore vectors
59a515125bSLeila Ghaffari   CeedVectorTakeArray(m_ceed, MemTypeP2C(m_mem_type), NULL);
60a515125bSLeila Ghaffari   ierr = VecRestoreArrayReadAndMemType(M_loc, (const PetscScalar **)&m);
61a515125bSLeila Ghaffari   CHKERRQ(ierr);
62a515125bSLeila Ghaffari 
63a515125bSLeila Ghaffari   // Local-to-Global
64a515125bSLeila Ghaffari   ierr = VecZeroEntries(M); CHKERRQ(ierr);
65a515125bSLeila Ghaffari   ierr = DMLocalToGlobal(dm, M_loc, ADD_VALUES, M); CHKERRQ(ierr);
66a515125bSLeila Ghaffari   ierr = DMRestoreLocalVector(dm, &M_loc); CHKERRQ(ierr);
67a515125bSLeila Ghaffari 
68a515125bSLeila Ghaffari   // Invert diagonally lumped mass vector for RHS function
69a515125bSLeila Ghaffari   ierr = VecReciprocal(M); CHKERRQ(ierr);
70a515125bSLeila Ghaffari 
71a515125bSLeila Ghaffari   // Cleanup
72a515125bSLeila Ghaffari   CeedVectorDestroy(&ones_vec);
73a515125bSLeila Ghaffari   CeedVectorDestroy(&m_ceed);
74a515125bSLeila Ghaffari   CeedQFunctionDestroy(&qf_mass);
75a515125bSLeila Ghaffari   CeedOperatorDestroy(&op_mass);
76a515125bSLeila Ghaffari 
77a515125bSLeila Ghaffari   PetscFunctionReturn(0);
78a515125bSLeila Ghaffari }
79a515125bSLeila Ghaffari 
80a515125bSLeila Ghaffari // RHS (Explicit time-stepper) function setup
81a515125bSLeila Ghaffari //   This is the RHS of the ODE, given as u_t = G(t,u)
82a515125bSLeila Ghaffari //   This function takes in a state vector Q and writes into G
83a515125bSLeila Ghaffari PetscErrorCode RHS_NS(TS ts, PetscReal t, Vec Q, Vec G, void *user_data) {
84a515125bSLeila Ghaffari 
85a515125bSLeila Ghaffari   User           user = *(User *)user_data;
86a515125bSLeila Ghaffari   PetscScalar    *q, *g;
87a515125bSLeila Ghaffari   Vec            Q_loc, G_loc;
88a515125bSLeila Ghaffari   PetscMemType   q_mem_type, g_mem_type;
89a515125bSLeila Ghaffari   PetscErrorCode ierr;
90a515125bSLeila Ghaffari   PetscFunctionBeginUser;
91a515125bSLeila Ghaffari 
92bb8a0c61SJames Wright   // Update context field labels
9392ada588SJames Wright   if (user->phys->solution_time_label)
9492ada588SJames Wright     CeedOperatorContextSetDouble(user->op_rhs, user->phys->solution_time_label, &t);
95bb8a0c61SJames Wright   if (user->phys->timestep_size_label) {
96bb8a0c61SJames Wright     PetscScalar dt;
97bb8a0c61SJames Wright     ierr = TSGetTimeStep(ts,&dt); CHKERRQ(ierr);
98bb8a0c61SJames Wright     CeedOperatorContextSetDouble(user->op_rhs, user->phys->timestep_size_label,
99bb8a0c61SJames Wright                                  &dt);
100bb8a0c61SJames Wright   }
101a515125bSLeila Ghaffari 
102a515125bSLeila Ghaffari   // Get local vectors
103a515125bSLeila Ghaffari   ierr = DMGetLocalVector(user->dm, &Q_loc); CHKERRQ(ierr);
104a515125bSLeila Ghaffari   ierr = DMGetLocalVector(user->dm, &G_loc); CHKERRQ(ierr);
105a515125bSLeila Ghaffari 
106a515125bSLeila Ghaffari   // Global-to-local
107a515125bSLeila Ghaffari   ierr = VecZeroEntries(Q_loc); CHKERRQ(ierr);
108a515125bSLeila Ghaffari   ierr = DMGlobalToLocal(user->dm, Q, INSERT_VALUES, Q_loc); CHKERRQ(ierr);
109a515125bSLeila Ghaffari   ierr = DMPlexInsertBoundaryValues(user->dm, PETSC_TRUE, Q_loc, 0.0,
110a515125bSLeila Ghaffari                                     NULL, NULL, NULL); CHKERRQ(ierr);
111a515125bSLeila Ghaffari   ierr = VecZeroEntries(G_loc); CHKERRQ(ierr);
112a515125bSLeila Ghaffari 
113a515125bSLeila Ghaffari   // Place PETSc vectors in CEED vectors
114a515125bSLeila Ghaffari   ierr = VecGetArrayReadAndMemType(Q_loc, (const PetscScalar **)&q, &q_mem_type);
115a515125bSLeila Ghaffari   CHKERRQ(ierr);
116a515125bSLeila Ghaffari   ierr = VecGetArrayAndMemType(G_loc, &g, &g_mem_type); CHKERRQ(ierr);
117a515125bSLeila Ghaffari   CeedVectorSetArray(user->q_ceed, MemTypeP2C(q_mem_type), CEED_USE_POINTER, q);
118a515125bSLeila Ghaffari   CeedVectorSetArray(user->g_ceed, MemTypeP2C(g_mem_type), CEED_USE_POINTER, g);
119a515125bSLeila Ghaffari 
120a515125bSLeila Ghaffari   // Apply CEED operator
121a515125bSLeila Ghaffari   CeedOperatorApply(user->op_rhs, user->q_ceed, user->g_ceed,
122a515125bSLeila Ghaffari                     CEED_REQUEST_IMMEDIATE);
123a515125bSLeila Ghaffari 
124a515125bSLeila Ghaffari   // Restore vectors
125a515125bSLeila Ghaffari   CeedVectorTakeArray(user->q_ceed, MemTypeP2C(q_mem_type), NULL);
126a515125bSLeila Ghaffari   CeedVectorTakeArray(user->g_ceed, MemTypeP2C(g_mem_type), NULL);
127a515125bSLeila Ghaffari   ierr = VecRestoreArrayReadAndMemType(Q_loc, (const PetscScalar **)&q);
128a515125bSLeila Ghaffari   CHKERRQ(ierr);
129a515125bSLeila Ghaffari   ierr = VecRestoreArrayAndMemType(G_loc, &g); CHKERRQ(ierr);
130a515125bSLeila Ghaffari 
131a515125bSLeila Ghaffari   // Local-to-Global
132a515125bSLeila Ghaffari   ierr = VecZeroEntries(G); CHKERRQ(ierr);
133a515125bSLeila Ghaffari   ierr = DMLocalToGlobal(user->dm, G_loc, ADD_VALUES, G); CHKERRQ(ierr);
134a515125bSLeila Ghaffari 
135a515125bSLeila Ghaffari   // Inverse of the lumped mass matrix (M is Minv)
136a515125bSLeila Ghaffari   ierr = VecPointwiseMult(G, G, user->M); CHKERRQ(ierr);
137a515125bSLeila Ghaffari 
138a515125bSLeila Ghaffari   // Restore vectors
139a515125bSLeila Ghaffari   ierr = DMRestoreLocalVector(user->dm, &Q_loc); CHKERRQ(ierr);
140a515125bSLeila Ghaffari   ierr = DMRestoreLocalVector(user->dm, &G_loc); CHKERRQ(ierr);
141a515125bSLeila Ghaffari 
142a515125bSLeila Ghaffari   PetscFunctionReturn(0);
143a515125bSLeila Ghaffari }
144a515125bSLeila Ghaffari 
145a515125bSLeila Ghaffari // Implicit time-stepper function setup
146a515125bSLeila Ghaffari PetscErrorCode IFunction_NS(TS ts, PetscReal t, Vec Q, Vec Q_dot, Vec G,
147a515125bSLeila Ghaffari                             void *user_data) {
148a515125bSLeila Ghaffari   User              user = *(User *)user_data;
149a515125bSLeila Ghaffari   const PetscScalar *q, *q_dot;
150a515125bSLeila Ghaffari   PetscScalar       *g;
151a515125bSLeila Ghaffari   Vec               Q_loc, Q_dot_loc, G_loc;
152a515125bSLeila Ghaffari   PetscMemType      q_mem_type, q_dot_mem_type, g_mem_type;
153a515125bSLeila Ghaffari   PetscErrorCode    ierr;
154a515125bSLeila Ghaffari   PetscFunctionBeginUser;
155a515125bSLeila Ghaffari 
156bb8a0c61SJames Wright   // Update context field labels
15792ada588SJames Wright   if (user->phys->solution_time_label)
15892ada588SJames Wright     CeedOperatorContextSetDouble(user->op_ifunction,
15992ada588SJames Wright                                  user->phys->solution_time_label, &t);
160bb8a0c61SJames Wright   if (user->phys->timestep_size_label) {
161bb8a0c61SJames Wright     PetscScalar dt;
162bb8a0c61SJames Wright     ierr = TSGetTimeStep(ts,&dt); CHKERRQ(ierr);
163bb8a0c61SJames Wright     CeedOperatorContextSetDouble(user->op_ifunction,
164bb8a0c61SJames Wright                                  user->phys->timestep_size_label, &dt);
165bb8a0c61SJames Wright   }
166a515125bSLeila Ghaffari 
167a515125bSLeila Ghaffari   // Get local vectors
168a515125bSLeila Ghaffari   ierr = DMGetLocalVector(user->dm, &Q_loc); CHKERRQ(ierr);
169a515125bSLeila Ghaffari   ierr = DMGetLocalVector(user->dm, &Q_dot_loc); CHKERRQ(ierr);
170a515125bSLeila Ghaffari   ierr = DMGetLocalVector(user->dm, &G_loc); CHKERRQ(ierr);
171a515125bSLeila Ghaffari 
172a515125bSLeila Ghaffari   // Global-to-local
173a515125bSLeila Ghaffari   ierr = VecZeroEntries(Q_loc); CHKERRQ(ierr);
174a515125bSLeila Ghaffari   ierr = DMGlobalToLocal(user->dm, Q, INSERT_VALUES, Q_loc); CHKERRQ(ierr);
175a515125bSLeila Ghaffari   ierr = DMPlexInsertBoundaryValues(user->dm, PETSC_TRUE, Q_loc, 0.0,
176a515125bSLeila Ghaffari                                     NULL, NULL, NULL); CHKERRQ(ierr);
177a515125bSLeila Ghaffari   ierr = VecZeroEntries(Q_dot_loc); CHKERRQ(ierr);
178a515125bSLeila Ghaffari   ierr = DMGlobalToLocal(user->dm, Q_dot, INSERT_VALUES, Q_dot_loc);
179a515125bSLeila Ghaffari   CHKERRQ(ierr);
180a515125bSLeila Ghaffari   ierr = VecZeroEntries(G_loc); CHKERRQ(ierr);
181a515125bSLeila Ghaffari 
182a515125bSLeila Ghaffari   // Place PETSc vectors in CEED vectors
183a515125bSLeila Ghaffari   ierr = VecGetArrayReadAndMemType(Q_loc, &q, &q_mem_type); CHKERRQ(ierr);
184a515125bSLeila Ghaffari   ierr = VecGetArrayReadAndMemType(Q_dot_loc, &q_dot, &q_dot_mem_type);
185a515125bSLeila Ghaffari   CHKERRQ(ierr);
186a515125bSLeila Ghaffari   ierr = VecGetArrayAndMemType(G_loc, &g, &g_mem_type); CHKERRQ(ierr);
187a515125bSLeila Ghaffari   CeedVectorSetArray(user->q_ceed, MemTypeP2C(q_mem_type), CEED_USE_POINTER,
188a515125bSLeila Ghaffari                      (PetscScalar *)q);
189a515125bSLeila Ghaffari   CeedVectorSetArray(user->q_dot_ceed, MemTypeP2C(q_dot_mem_type),
190139613f2SLeila Ghaffari                      CEED_USE_POINTER, (PetscScalar *)q_dot);
191a515125bSLeila Ghaffari   CeedVectorSetArray(user->g_ceed, MemTypeP2C(g_mem_type), CEED_USE_POINTER, g);
192a515125bSLeila Ghaffari 
193a515125bSLeila Ghaffari   // Apply CEED operator
194a515125bSLeila Ghaffari   CeedOperatorApply(user->op_ifunction, user->q_ceed, user->g_ceed,
195a515125bSLeila Ghaffari                     CEED_REQUEST_IMMEDIATE);
196a515125bSLeila Ghaffari 
197a515125bSLeila Ghaffari   // Restore vectors
198a515125bSLeila Ghaffari   CeedVectorTakeArray(user->q_ceed, MemTypeP2C(q_mem_type), NULL);
199a515125bSLeila Ghaffari   CeedVectorTakeArray(user->q_dot_ceed, MemTypeP2C(q_dot_mem_type), NULL);
200a515125bSLeila Ghaffari   CeedVectorTakeArray(user->g_ceed, MemTypeP2C(g_mem_type), NULL);
201a515125bSLeila Ghaffari   ierr = VecRestoreArrayReadAndMemType(Q_loc, &q); CHKERRQ(ierr);
202a515125bSLeila Ghaffari   ierr = VecRestoreArrayReadAndMemType(Q_dot_loc, &q_dot); CHKERRQ(ierr);
203a515125bSLeila Ghaffari   ierr = VecRestoreArrayAndMemType(G_loc, &g); CHKERRQ(ierr);
204a515125bSLeila Ghaffari 
205a515125bSLeila Ghaffari   // Local-to-Global
206a515125bSLeila Ghaffari   ierr = VecZeroEntries(G); CHKERRQ(ierr);
207a515125bSLeila Ghaffari   ierr = DMLocalToGlobal(user->dm, G_loc, ADD_VALUES, G); CHKERRQ(ierr);
208a515125bSLeila Ghaffari 
209a515125bSLeila Ghaffari   // Restore vectors
210a515125bSLeila Ghaffari   ierr = DMRestoreLocalVector(user->dm, &Q_loc); CHKERRQ(ierr);
211a515125bSLeila Ghaffari   ierr = DMRestoreLocalVector(user->dm, &Q_dot_loc); CHKERRQ(ierr);
212a515125bSLeila Ghaffari   ierr = DMRestoreLocalVector(user->dm, &G_loc); CHKERRQ(ierr);
213a515125bSLeila Ghaffari 
214a515125bSLeila Ghaffari   PetscFunctionReturn(0);
215a515125bSLeila Ghaffari }
216a515125bSLeila Ghaffari 
217*f0b65372SJed Brown static PetscErrorCode MatMult_NS_IJacobian(Mat J, Vec Q, Vec G) {
218*f0b65372SJed Brown   User user;
219*f0b65372SJed Brown   const PetscScalar *q;
220*f0b65372SJed Brown   PetscScalar       *g;
221*f0b65372SJed Brown   Vec               Q_loc, G_loc;
222*f0b65372SJed Brown   PetscMemType      q_mem_type, g_mem_type;
223*f0b65372SJed Brown   PetscErrorCode    ierr;
224*f0b65372SJed Brown   PetscFunctionBeginUser;
225*f0b65372SJed Brown   MatShellGetContext(J, &user);
226*f0b65372SJed Brown   // Get local vectors
227*f0b65372SJed Brown   ierr = DMGetLocalVector(user->dm, &Q_loc); CHKERRQ(ierr);
228*f0b65372SJed Brown   ierr = DMGetLocalVector(user->dm, &G_loc); CHKERRQ(ierr);
229*f0b65372SJed Brown 
230*f0b65372SJed Brown   // Global-to-local
231*f0b65372SJed Brown   ierr = VecZeroEntries(Q_loc); CHKERRQ(ierr);
232*f0b65372SJed Brown   ierr = DMGlobalToLocal(user->dm, Q, INSERT_VALUES, Q_loc); CHKERRQ(ierr);
233*f0b65372SJed Brown   ierr = VecZeroEntries(G_loc); CHKERRQ(ierr);
234*f0b65372SJed Brown 
235*f0b65372SJed Brown   // Place PETSc vectors in CEED vectors
236*f0b65372SJed Brown   ierr = VecGetArrayReadAndMemType(Q_loc, &q, &q_mem_type); CHKERRQ(ierr);
237*f0b65372SJed Brown   ierr = VecGetArrayAndMemType(G_loc, &g, &g_mem_type); CHKERRQ(ierr);
238*f0b65372SJed Brown   CeedVectorSetArray(user->q_ceed, MemTypeP2C(q_mem_type), CEED_USE_POINTER,
239*f0b65372SJed Brown                      (PetscScalar *)q);
240*f0b65372SJed Brown   CeedVectorSetArray(user->g_ceed, MemTypeP2C(g_mem_type), CEED_USE_POINTER, g);
241*f0b65372SJed Brown 
242*f0b65372SJed Brown   // Apply CEED operator
243*f0b65372SJed Brown   CeedOperatorApply(user->op_ijacobian, user->q_ceed, user->g_ceed,
244*f0b65372SJed Brown                     CEED_REQUEST_IMMEDIATE);
245*f0b65372SJed Brown 
246*f0b65372SJed Brown   // Restore vectors
247*f0b65372SJed Brown   CeedVectorTakeArray(user->q_ceed, MemTypeP2C(q_mem_type), NULL);
248*f0b65372SJed Brown   CeedVectorTakeArray(user->g_ceed, MemTypeP2C(g_mem_type), NULL);
249*f0b65372SJed Brown   ierr = VecRestoreArrayReadAndMemType(Q_loc, &q); CHKERRQ(ierr);
250*f0b65372SJed Brown   ierr = VecRestoreArrayAndMemType(G_loc, &g); CHKERRQ(ierr);
251*f0b65372SJed Brown 
252*f0b65372SJed Brown   // Local-to-Global
253*f0b65372SJed Brown   ierr = VecZeroEntries(G); CHKERRQ(ierr);
254*f0b65372SJed Brown   ierr = DMLocalToGlobal(user->dm, G_loc, ADD_VALUES, G); CHKERRQ(ierr);
255*f0b65372SJed Brown 
256*f0b65372SJed Brown   // Restore vectors
257*f0b65372SJed Brown   ierr = DMRestoreLocalVector(user->dm, &Q_loc); CHKERRQ(ierr);
258*f0b65372SJed Brown   ierr = DMRestoreLocalVector(user->dm, &G_loc); CHKERRQ(ierr);
259*f0b65372SJed Brown   PetscFunctionReturn(0);
260*f0b65372SJed Brown }
261*f0b65372SJed Brown 
262*f0b65372SJed Brown PetscErrorCode MatGetDiagonal_NS_IJacobian(Mat A, Vec D) {
263*f0b65372SJed Brown   User user;
264*f0b65372SJed Brown   Vec D_loc;
265*f0b65372SJed Brown   PetscScalar *d;
266*f0b65372SJed Brown   PetscMemType mem_type;
267*f0b65372SJed Brown 
268*f0b65372SJed Brown   PetscFunctionBeginUser;
269*f0b65372SJed Brown   MatShellGetContext(A, &user);
270*f0b65372SJed Brown   PetscCall(DMGetLocalVector(user->dm, &D_loc));
271*f0b65372SJed Brown   PetscCall(VecGetArrayAndMemType(D_loc, &d, &mem_type));
272*f0b65372SJed Brown   CeedVectorSetArray(user->g_ceed, MemTypeP2C(mem_type), CEED_USE_POINTER, d);
273*f0b65372SJed Brown   CeedOperatorLinearAssembleDiagonal(user->op_ijacobian, user->g_ceed,
274*f0b65372SJed Brown                                      CEED_REQUEST_IMMEDIATE);
275*f0b65372SJed Brown   CeedVectorTakeArray(user->g_ceed, MemTypeP2C(mem_type), NULL);
276*f0b65372SJed Brown   PetscCall(VecRestoreArrayAndMemType(D_loc, &d));
277*f0b65372SJed Brown   PetscCall(VecZeroEntries(D));
278*f0b65372SJed Brown   PetscCall(DMLocalToGlobal(user->dm, D_loc, ADD_VALUES, D));
279*f0b65372SJed Brown   PetscCall(DMRestoreLocalVector(user->dm, &D_loc));
280*f0b65372SJed Brown   VecViewFromOptions(D, NULL, "-diag_vec_view");
281*f0b65372SJed Brown   PetscFunctionReturn(0);
282*f0b65372SJed Brown }
283*f0b65372SJed Brown 
284*f0b65372SJed Brown PetscErrorCode FormIJacobian_NS(TS ts, PetscReal t, Vec Q, Vec Q_dot,
285*f0b65372SJed Brown                                 PetscReal shift, Mat J, Mat J_pre,
286*f0b65372SJed Brown                                 void *user_data) {
287*f0b65372SJed Brown   User user = *(User *)user_data;
288*f0b65372SJed Brown   PetscBool J_is_shell, J_pre_is_shell;
289*f0b65372SJed Brown   PetscFunctionBeginUser;
290*f0b65372SJed Brown   if (user->phys->ijacobian_time_shift_label)
291*f0b65372SJed Brown     CeedOperatorContextSetDouble(user->op_ijacobian,
292*f0b65372SJed Brown                                  user->phys->ijacobian_time_shift_label, &shift);
293*f0b65372SJed Brown   PetscCall(MatAssemblyBegin(J, MAT_FINAL_ASSEMBLY));
294*f0b65372SJed Brown   PetscCall(MatAssemblyEnd(J, MAT_FINAL_ASSEMBLY));
295*f0b65372SJed Brown   Vec coo_vec = NULL;
296*f0b65372SJed Brown   PetscCall(PetscObjectTypeCompare((PetscObject)J, MATSHELL, &J_is_shell));
297*f0b65372SJed Brown   PetscCall(PetscObjectTypeCompare((PetscObject)J_pre, MATSHELL,
298*f0b65372SJed Brown                                    &J_pre_is_shell));
299*f0b65372SJed Brown   if (!user->matrices_set_up) {
300*f0b65372SJed Brown     if (J_is_shell) {
301*f0b65372SJed Brown       PetscCall(MatShellSetContext(J, user));
302*f0b65372SJed Brown       PetscCall(MatShellSetOperation(J, MATOP_MULT,
303*f0b65372SJed Brown                                      (void (*)(void))MatMult_NS_IJacobian));
304*f0b65372SJed Brown       PetscCall(MatShellSetOperation(J, MATOP_GET_DIAGONAL,
305*f0b65372SJed Brown                                      (void (*)(void))MatGetDiagonal_NS_IJacobian));
306*f0b65372SJed Brown       PetscCall(MatSetUp(J));
307*f0b65372SJed Brown     }
308*f0b65372SJed Brown     if (!J_pre_is_shell) {
309*f0b65372SJed Brown       PetscCount ncoo;
310*f0b65372SJed Brown       PetscInt *rows, *cols;
311*f0b65372SJed Brown       PetscCall(CeedOperatorLinearAssembleSymbolic(user->op_ijacobian, &ncoo, &rows,
312*f0b65372SJed Brown                 &cols));
313*f0b65372SJed Brown       PetscCall(MatSetPreallocationCOOLocal(J_pre, ncoo, rows, cols));
314*f0b65372SJed Brown       free(rows);
315*f0b65372SJed Brown       free(cols);
316*f0b65372SJed Brown       CeedVectorCreate(user->ceed, ncoo, &user->coo_values);
317*f0b65372SJed Brown       user->matrices_set_up = true;
318*f0b65372SJed Brown       VecCreateSeq(PETSC_COMM_WORLD, ncoo, &coo_vec);
319*f0b65372SJed Brown     }
320*f0b65372SJed Brown   }
321*f0b65372SJed Brown   if (!J_pre_is_shell) {
322*f0b65372SJed Brown     CeedMemType mem_type = CEED_MEM_HOST;
323*f0b65372SJed Brown     const PetscScalar *values;
324*f0b65372SJed Brown     MatType mat_type;
325*f0b65372SJed Brown     PetscCall(MatGetType(J_pre, &mat_type));
326*f0b65372SJed Brown     //if (strstr(mat_type, "kokkos") || strstr(mat_type, "cusparse")) mem_type = CEED_MEM_DEVICE;
327*f0b65372SJed Brown     CeedOperatorLinearAssemble(user->op_ijacobian, user->coo_values);
328*f0b65372SJed Brown     CeedVectorGetArrayRead(user->coo_values, mem_type, &values);
329*f0b65372SJed Brown     if (coo_vec) {
330*f0b65372SJed Brown       VecPlaceArray(coo_vec, values);
331*f0b65372SJed Brown       VecViewFromOptions(coo_vec, NULL, "-coo_vec_view");
332*f0b65372SJed Brown       VecDestroy(&coo_vec);
333*f0b65372SJed Brown     }
334*f0b65372SJed Brown     PetscCall(MatSetValuesCOO(J_pre, values, INSERT_VALUES));
335*f0b65372SJed Brown     CeedVectorRestoreArrayRead(user->coo_values, &values);
336*f0b65372SJed Brown   }
337*f0b65372SJed Brown   PetscFunctionReturn(0);
338*f0b65372SJed Brown }
339*f0b65372SJed Brown 
340a515125bSLeila Ghaffari // User provided TS Monitor
341a515125bSLeila Ghaffari PetscErrorCode TSMonitor_NS(TS ts, PetscInt step_no, PetscReal time,
342a515125bSLeila Ghaffari                             Vec Q, void *ctx) {
343a515125bSLeila Ghaffari   User           user = ctx;
344a515125bSLeila Ghaffari   Vec            Q_loc;
345a515125bSLeila Ghaffari   char           file_path[PETSC_MAX_PATH_LEN];
346a515125bSLeila Ghaffari   PetscViewer    viewer;
347a515125bSLeila Ghaffari   PetscErrorCode ierr;
348a515125bSLeila Ghaffari   PetscFunctionBeginUser;
349a515125bSLeila Ghaffari 
350a515125bSLeila Ghaffari   // Print every 'output_freq' steps
351a515125bSLeila Ghaffari   if (step_no % user->app_ctx->output_freq != 0)
352a515125bSLeila Ghaffari     PetscFunctionReturn(0);
353a515125bSLeila Ghaffari 
354a515125bSLeila Ghaffari   // Set up output
355a515125bSLeila Ghaffari   ierr = DMGetLocalVector(user->dm, &Q_loc); CHKERRQ(ierr);
356a515125bSLeila Ghaffari   ierr = PetscObjectSetName((PetscObject)Q_loc, "StateVec"); CHKERRQ(ierr);
357a515125bSLeila Ghaffari   ierr = VecZeroEntries(Q_loc); CHKERRQ(ierr);
358a515125bSLeila Ghaffari   ierr = DMGlobalToLocal(user->dm, Q, INSERT_VALUES, Q_loc); CHKERRQ(ierr);
359a515125bSLeila Ghaffari 
360a515125bSLeila Ghaffari   // Output
361d940ca4eSJed Brown   ierr = PetscSNPrintf(file_path, sizeof file_path,
362d940ca4eSJed Brown                        "%s/ns-%03" PetscInt_FMT ".vtu",
363a515125bSLeila Ghaffari                        user->app_ctx->output_dir, step_no + user->app_ctx->cont_steps);
364a515125bSLeila Ghaffari   CHKERRQ(ierr);
365a515125bSLeila Ghaffari   ierr = PetscViewerVTKOpen(PetscObjectComm((PetscObject)Q), file_path,
366a515125bSLeila Ghaffari                             FILE_MODE_WRITE, &viewer); CHKERRQ(ierr);
367a515125bSLeila Ghaffari   ierr = VecView(Q_loc, viewer); CHKERRQ(ierr);
368a515125bSLeila Ghaffari   ierr = PetscViewerDestroy(&viewer); CHKERRQ(ierr);
369a515125bSLeila Ghaffari   if (user->dm_viz) {
370a515125bSLeila Ghaffari     Vec         Q_refined, Q_refined_loc;
371a515125bSLeila Ghaffari     char        file_path_refined[PETSC_MAX_PATH_LEN];
372a515125bSLeila Ghaffari     PetscViewer viewer_refined;
373a515125bSLeila Ghaffari 
374a515125bSLeila Ghaffari     ierr = DMGetGlobalVector(user->dm_viz, &Q_refined); CHKERRQ(ierr);
375a515125bSLeila Ghaffari     ierr = DMGetLocalVector(user->dm_viz, &Q_refined_loc); CHKERRQ(ierr);
376a515125bSLeila Ghaffari     ierr = PetscObjectSetName((PetscObject)Q_refined_loc, "Refined");
377a515125bSLeila Ghaffari     CHKERRQ(ierr);
378a515125bSLeila Ghaffari     ierr = MatInterpolate(user->interp_viz, Q, Q_refined); CHKERRQ(ierr);
379a515125bSLeila Ghaffari     ierr = VecZeroEntries(Q_refined_loc); CHKERRQ(ierr);
380a515125bSLeila Ghaffari     ierr = DMGlobalToLocal(user->dm_viz, Q_refined, INSERT_VALUES, Q_refined_loc);
381a515125bSLeila Ghaffari     CHKERRQ(ierr);
382a515125bSLeila Ghaffari     ierr = PetscSNPrintf(file_path_refined, sizeof file_path_refined,
383d940ca4eSJed Brown                          "%s/nsrefined-%03" PetscInt_FMT ".vtu", user->app_ctx->output_dir,
384139613f2SLeila Ghaffari                          step_no + user->app_ctx->cont_steps);
385a515125bSLeila Ghaffari     CHKERRQ(ierr);
386a515125bSLeila Ghaffari     ierr = PetscViewerVTKOpen(PetscObjectComm((PetscObject)Q_refined),
387139613f2SLeila Ghaffari                               file_path_refined, FILE_MODE_WRITE, &viewer_refined); CHKERRQ(ierr);
388a515125bSLeila Ghaffari     ierr = VecView(Q_refined_loc, viewer_refined); CHKERRQ(ierr);
389a515125bSLeila Ghaffari     ierr = DMRestoreLocalVector(user->dm_viz, &Q_refined_loc); CHKERRQ(ierr);
390a515125bSLeila Ghaffari     ierr = DMRestoreGlobalVector(user->dm_viz, &Q_refined); CHKERRQ(ierr);
391a515125bSLeila Ghaffari     ierr = PetscViewerDestroy(&viewer_refined); CHKERRQ(ierr);
392a515125bSLeila Ghaffari   }
393a515125bSLeila Ghaffari   ierr = DMRestoreLocalVector(user->dm, &Q_loc); CHKERRQ(ierr);
394a515125bSLeila Ghaffari 
395a515125bSLeila Ghaffari   // Save data in a binary file for continuation of simulations
396a515125bSLeila Ghaffari   ierr = PetscSNPrintf(file_path, sizeof file_path, "%s/ns-solution.bin",
397a515125bSLeila Ghaffari                        user->app_ctx->output_dir); CHKERRQ(ierr);
398a515125bSLeila Ghaffari   ierr = PetscViewerBinaryOpen(user->comm, file_path, FILE_MODE_WRITE, &viewer);
399a515125bSLeila Ghaffari   CHKERRQ(ierr);
400a515125bSLeila Ghaffari   ierr = VecView(Q, viewer); CHKERRQ(ierr);
401a515125bSLeila Ghaffari   ierr = PetscViewerDestroy(&viewer); CHKERRQ(ierr);
402a515125bSLeila Ghaffari 
403a515125bSLeila Ghaffari   // Save time stamp
404a515125bSLeila Ghaffari   // Dimensionalize time back
405a515125bSLeila Ghaffari   time /= user->units->second;
406a515125bSLeila Ghaffari   ierr = PetscSNPrintf(file_path, sizeof file_path, "%s/ns-time.bin",
407a515125bSLeila Ghaffari                        user->app_ctx->output_dir); CHKERRQ(ierr);
408a515125bSLeila Ghaffari   ierr = PetscViewerBinaryOpen(user->comm, file_path, FILE_MODE_WRITE, &viewer);
409a515125bSLeila Ghaffari   CHKERRQ(ierr);
410a515125bSLeila Ghaffari   #if PETSC_VERSION_GE(3,13,0)
411a515125bSLeila Ghaffari   ierr = PetscViewerBinaryWrite(viewer, &time, 1, PETSC_REAL);
412a515125bSLeila Ghaffari   #else
413a515125bSLeila Ghaffari   ierr = PetscViewerBinaryWrite(viewer, &time, 1, PETSC_REAL, true);
414a515125bSLeila Ghaffari   #endif
415a515125bSLeila Ghaffari   CHKERRQ(ierr);
416a515125bSLeila Ghaffari   ierr = PetscViewerDestroy(&viewer); CHKERRQ(ierr);
417a515125bSLeila Ghaffari 
418a515125bSLeila Ghaffari   PetscFunctionReturn(0);
419a515125bSLeila Ghaffari }
420a515125bSLeila Ghaffari 
421a515125bSLeila Ghaffari // TS: Create, setup, and solve
422a515125bSLeila Ghaffari PetscErrorCode TSSolve_NS(DM dm, User user, AppCtx app_ctx, Physics phys,
423a515125bSLeila Ghaffari                           Vec *Q, PetscScalar *f_time, TS *ts) {
424a515125bSLeila Ghaffari   MPI_Comm       comm = user->comm;
425a515125bSLeila Ghaffari   TSAdapt        adapt;
426a515125bSLeila Ghaffari   PetscScalar    final_time;
427a515125bSLeila Ghaffari   PetscErrorCode ierr;
428a515125bSLeila Ghaffari   PetscFunctionBeginUser;
429a515125bSLeila Ghaffari 
430a515125bSLeila Ghaffari   ierr = TSCreate(comm, ts); CHKERRQ(ierr);
431a515125bSLeila Ghaffari   ierr = TSSetDM(*ts, dm); CHKERRQ(ierr);
432a515125bSLeila Ghaffari   if (phys->implicit) {
433a515125bSLeila Ghaffari     ierr = TSSetType(*ts, TSBDF); CHKERRQ(ierr);
434a515125bSLeila Ghaffari     if (user->op_ifunction) {
435a515125bSLeila Ghaffari       ierr = TSSetIFunction(*ts, NULL, IFunction_NS, &user); CHKERRQ(ierr);
436a515125bSLeila Ghaffari     } else { // Implicit integrators can fall back to using an RHSFunction
437a515125bSLeila Ghaffari       ierr = TSSetRHSFunction(*ts, NULL, RHS_NS, &user); CHKERRQ(ierr);
438a515125bSLeila Ghaffari     }
439*f0b65372SJed Brown     if (user->op_ijacobian) {
440*f0b65372SJed Brown       ierr = DMTSSetIJacobian(dm, FormIJacobian_NS, &user); CHKERRQ(ierr);
441*f0b65372SJed Brown     }
442a515125bSLeila Ghaffari   } else {
443a515125bSLeila Ghaffari     if (!user->op_rhs) SETERRQ(comm, PETSC_ERR_ARG_NULL,
444a515125bSLeila Ghaffari                                  "Problem does not provide RHSFunction");
445a515125bSLeila Ghaffari     ierr = TSSetType(*ts, TSRK); CHKERRQ(ierr);
446a515125bSLeila Ghaffari     ierr = TSRKSetType(*ts, TSRK5F); CHKERRQ(ierr);
447a515125bSLeila Ghaffari     ierr = TSSetRHSFunction(*ts, NULL, RHS_NS, &user); CHKERRQ(ierr);
448a515125bSLeila Ghaffari   }
449a515125bSLeila Ghaffari   ierr = TSSetMaxTime(*ts, 500. * user->units->second); CHKERRQ(ierr);
450a515125bSLeila Ghaffari   ierr = TSSetExactFinalTime(*ts, TS_EXACTFINALTIME_STEPOVER); CHKERRQ(ierr);
451a515125bSLeila Ghaffari   ierr = TSSetTimeStep(*ts, 1.e-2 * user->units->second); CHKERRQ(ierr);
452a515125bSLeila Ghaffari   if (app_ctx->test_mode) {ierr = TSSetMaxSteps(*ts, 10); CHKERRQ(ierr);}
453a515125bSLeila Ghaffari   ierr = TSGetAdapt(*ts, &adapt); CHKERRQ(ierr);
454a515125bSLeila Ghaffari   ierr = TSAdaptSetStepLimits(adapt,
455a515125bSLeila Ghaffari                               1.e-12 * user->units->second,
456a515125bSLeila Ghaffari                               1.e2 * user->units->second); CHKERRQ(ierr);
457a515125bSLeila Ghaffari   ierr = TSSetFromOptions(*ts); CHKERRQ(ierr);
458a515125bSLeila Ghaffari   if (!app_ctx->cont_steps) { // print initial condition
459a515125bSLeila Ghaffari     if (!app_ctx->test_mode) {
460a515125bSLeila Ghaffari       ierr = TSMonitor_NS(*ts, 0, 0., *Q, user); CHKERRQ(ierr);
461a515125bSLeila Ghaffari     }
462a515125bSLeila Ghaffari   } else { // continue from time of last output
463a515125bSLeila Ghaffari     PetscReal   time;
464a515125bSLeila Ghaffari     PetscInt    count;
465a515125bSLeila Ghaffari     PetscViewer viewer;
466a515125bSLeila Ghaffari     char        file_path[PETSC_MAX_PATH_LEN];
467a515125bSLeila Ghaffari     ierr = PetscSNPrintf(file_path, sizeof file_path, "%s/ns-time.bin",
468a515125bSLeila Ghaffari                          app_ctx->output_dir); CHKERRQ(ierr);
469a515125bSLeila Ghaffari     ierr = PetscViewerBinaryOpen(comm, file_path, FILE_MODE_READ, &viewer);
470a515125bSLeila Ghaffari     CHKERRQ(ierr);
471a515125bSLeila Ghaffari     ierr = PetscViewerBinaryRead(viewer, &time, 1, &count, PETSC_REAL);
472a515125bSLeila Ghaffari     CHKERRQ(ierr);
473a515125bSLeila Ghaffari     ierr = PetscViewerDestroy(&viewer); CHKERRQ(ierr);
474a515125bSLeila Ghaffari     ierr = TSSetTime(*ts, time * user->units->second); CHKERRQ(ierr);
475a515125bSLeila Ghaffari   }
476a515125bSLeila Ghaffari   if (!app_ctx->test_mode) {
477a515125bSLeila Ghaffari     ierr = TSMonitorSet(*ts, TSMonitor_NS, user, NULL); CHKERRQ(ierr);
478a515125bSLeila Ghaffari   }
479a515125bSLeila Ghaffari 
480a515125bSLeila Ghaffari   // Solve
481a515125bSLeila Ghaffari   double start, cpu_time_used;
482a515125bSLeila Ghaffari   start = MPI_Wtime();
483a515125bSLeila Ghaffari   ierr = PetscBarrier((PetscObject) *ts); CHKERRQ(ierr);
484a515125bSLeila Ghaffari   ierr = TSSolve(*ts, *Q); CHKERRQ(ierr);
485a515125bSLeila Ghaffari   cpu_time_used = MPI_Wtime() - start;
486a515125bSLeila Ghaffari   ierr = TSGetSolveTime(*ts, &final_time); CHKERRQ(ierr);
487a515125bSLeila Ghaffari   *f_time = final_time;
488a515125bSLeila Ghaffari   ierr = MPI_Allreduce(MPI_IN_PLACE, &cpu_time_used, 1, MPI_DOUBLE, MPI_MIN,
489a515125bSLeila Ghaffari                        comm); CHKERRQ(ierr);
490a515125bSLeila Ghaffari   if (!app_ctx->test_mode) {
491a515125bSLeila Ghaffari     ierr = PetscPrintf(PETSC_COMM_WORLD,
492a515125bSLeila Ghaffari                        "Time taken for solution (sec): %g\n",
493a515125bSLeila Ghaffari                        (double)cpu_time_used); CHKERRQ(ierr);
494a515125bSLeila Ghaffari   }
495a515125bSLeila Ghaffari   PetscFunctionReturn(0);
496a515125bSLeila Ghaffari }
497