xref: /honee/src/setupts.c (revision e2f84137d0f7db9b4666917d058148e3dc471d7d)
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   User           user = *(User *)user_data;
85a515125bSLeila Ghaffari   PetscScalar    *q, *g;
86*e2f84137SJeremy L Thompson   Vec            Q_loc = user->Q_loc, G_loc;
87a515125bSLeila Ghaffari   PetscMemType   q_mem_type, g_mem_type;
88a515125bSLeila Ghaffari   PetscErrorCode ierr;
89a515125bSLeila Ghaffari   PetscFunctionBeginUser;
90a515125bSLeila Ghaffari 
91*e2f84137SJeremy L Thompson   // Get local vector
92*e2f84137SJeremy L Thompson   ierr = DMGetLocalVector(user->dm, &G_loc); CHKERRQ(ierr);
93*e2f84137SJeremy L Thompson 
94*e2f84137SJeremy L Thompson   // Update time dependent data
95*e2f84137SJeremy L Thompson   if (user->time != t) {
96*e2f84137SJeremy L Thompson     ierr = VecZeroEntries(Q_loc); CHKERRQ(ierr);
97*e2f84137SJeremy L Thompson     ierr = DMPlexInsertBoundaryValues(user->dm, PETSC_TRUE, Q_loc, t,
98*e2f84137SJeremy L Thompson                                       NULL, NULL, NULL); CHKERRQ(ierr);
99*e2f84137SJeremy L Thompson     if (user->phys->solution_time_label) {
10092ada588SJames Wright       CeedOperatorContextSetDouble(user->op_rhs, user->phys->solution_time_label, &t);
101*e2f84137SJeremy L Thompson     }
102*e2f84137SJeremy L Thompson     user->time = t;
103*e2f84137SJeremy L Thompson   }
104bb8a0c61SJames Wright   if (user->phys->timestep_size_label) {
105bb8a0c61SJames Wright     PetscScalar dt;
106bb8a0c61SJames Wright     ierr = TSGetTimeStep(ts, &dt); CHKERRQ(ierr);
107*e2f84137SJeremy L Thompson     if (user->dt != dt) {
108bb8a0c61SJames Wright       CeedOperatorContextSetDouble(user->op_rhs, user->phys->timestep_size_label,
109bb8a0c61SJames Wright                                    &dt);
110*e2f84137SJeremy L Thompson       user->dt = dt;
111*e2f84137SJeremy L Thompson     }
112bb8a0c61SJames Wright   }
113a515125bSLeila Ghaffari 
114a515125bSLeila Ghaffari   // Global-to-local
115a515125bSLeila Ghaffari   ierr = DMGlobalToLocal(user->dm, Q, INSERT_VALUES, Q_loc); CHKERRQ(ierr);
116a515125bSLeila Ghaffari 
117a515125bSLeila Ghaffari   // Place PETSc vectors in CEED vectors
118a515125bSLeila Ghaffari   ierr = VecGetArrayReadAndMemType(Q_loc, (const PetscScalar **)&q, &q_mem_type);
119a515125bSLeila Ghaffari   CHKERRQ(ierr);
120a515125bSLeila Ghaffari   ierr = VecGetArrayAndMemType(G_loc, &g, &g_mem_type); CHKERRQ(ierr);
121a515125bSLeila Ghaffari   CeedVectorSetArray(user->q_ceed, MemTypeP2C(q_mem_type), CEED_USE_POINTER, q);
122a515125bSLeila Ghaffari   CeedVectorSetArray(user->g_ceed, MemTypeP2C(g_mem_type), CEED_USE_POINTER, g);
123a515125bSLeila Ghaffari 
124a515125bSLeila Ghaffari   // Apply CEED operator
125a515125bSLeila Ghaffari   CeedOperatorApply(user->op_rhs, user->q_ceed, user->g_ceed,
126a515125bSLeila Ghaffari                     CEED_REQUEST_IMMEDIATE);
127a515125bSLeila Ghaffari 
128a515125bSLeila Ghaffari   // Restore vectors
129a515125bSLeila Ghaffari   CeedVectorTakeArray(user->q_ceed, MemTypeP2C(q_mem_type), NULL);
130a515125bSLeila Ghaffari   CeedVectorTakeArray(user->g_ceed, MemTypeP2C(g_mem_type), NULL);
131a515125bSLeila Ghaffari   ierr = VecRestoreArrayReadAndMemType(Q_loc, (const PetscScalar **)&q);
132a515125bSLeila Ghaffari   CHKERRQ(ierr);
133a515125bSLeila Ghaffari   ierr = VecRestoreArrayAndMemType(G_loc, &g); CHKERRQ(ierr);
134a515125bSLeila Ghaffari 
135a515125bSLeila Ghaffari   // Local-to-Global
136a515125bSLeila Ghaffari   ierr = VecZeroEntries(G); CHKERRQ(ierr);
137a515125bSLeila Ghaffari   ierr = DMLocalToGlobal(user->dm, G_loc, ADD_VALUES, G); CHKERRQ(ierr);
138a515125bSLeila Ghaffari 
139a515125bSLeila Ghaffari   // Inverse of the lumped mass matrix (M is Minv)
140a515125bSLeila Ghaffari   ierr = VecPointwiseMult(G, G, user->M); CHKERRQ(ierr);
141a515125bSLeila Ghaffari 
142a515125bSLeila Ghaffari   // Restore vectors
143a515125bSLeila Ghaffari   ierr = DMRestoreLocalVector(user->dm, &G_loc); CHKERRQ(ierr);
144a515125bSLeila Ghaffari 
145a515125bSLeila Ghaffari   PetscFunctionReturn(0);
146a515125bSLeila Ghaffari }
147a515125bSLeila Ghaffari 
148a515125bSLeila Ghaffari // Implicit time-stepper function setup
149a515125bSLeila Ghaffari PetscErrorCode IFunction_NS(TS ts, PetscReal t, Vec Q, Vec Q_dot, Vec G,
150a515125bSLeila Ghaffari                             void *user_data) {
151a515125bSLeila Ghaffari   User              user = *(User *)user_data;
152a515125bSLeila Ghaffari   const PetscScalar *q, *q_dot;
153a515125bSLeila Ghaffari   PetscScalar       *g;
154*e2f84137SJeremy L Thompson   Vec               Q_loc = user->Q_loc, Q_dot_loc = user->Q_dot_loc, G_loc;
155a515125bSLeila Ghaffari   PetscMemType      q_mem_type, q_dot_mem_type, g_mem_type;
156a515125bSLeila Ghaffari   PetscErrorCode    ierr;
157a515125bSLeila Ghaffari   PetscFunctionBeginUser;
158a515125bSLeila Ghaffari 
159*e2f84137SJeremy L Thompson   // Get local vectors
160*e2f84137SJeremy L Thompson   ierr = DMGetLocalVector(user->dm, &G_loc); CHKERRQ(ierr);
161*e2f84137SJeremy L Thompson 
162*e2f84137SJeremy L Thompson   // Update time dependent data
163*e2f84137SJeremy L Thompson   if (user->time != t) {
164*e2f84137SJeremy L Thompson     ierr = VecZeroEntries(Q_loc); CHKERRQ(ierr);
165*e2f84137SJeremy L Thompson     ierr = DMPlexInsertBoundaryValues(user->dm, PETSC_TRUE, Q_loc, t,
166*e2f84137SJeremy L Thompson                                       NULL, NULL, NULL); CHKERRQ(ierr);
167*e2f84137SJeremy L Thompson     if (user->phys->solution_time_label) {
16892ada588SJames Wright       CeedOperatorContextSetDouble(user->op_ifunction,
16992ada588SJames Wright                                    user->phys->solution_time_label, &t);
170*e2f84137SJeremy L Thompson     }
171*e2f84137SJeremy L Thompson     user->time = t;
172*e2f84137SJeremy L Thompson   }
173bb8a0c61SJames Wright   if (user->phys->timestep_size_label) {
174bb8a0c61SJames Wright     PetscScalar dt;
175bb8a0c61SJames Wright     ierr = TSGetTimeStep(ts, &dt); CHKERRQ(ierr);
176*e2f84137SJeremy L Thompson     if (user->dt != dt) {
177bb8a0c61SJames Wright       CeedOperatorContextSetDouble(user->op_ifunction,
178bb8a0c61SJames Wright                                    user->phys->timestep_size_label, &dt);
179*e2f84137SJeremy L Thompson       user->dt = dt;
180*e2f84137SJeremy L Thompson     }
181bb8a0c61SJames Wright   }
182a515125bSLeila Ghaffari 
183a515125bSLeila Ghaffari   // Global-to-local
184a515125bSLeila Ghaffari   ierr = DMGlobalToLocal(user->dm, Q, INSERT_VALUES, Q_loc); CHKERRQ(ierr);
185a515125bSLeila Ghaffari   ierr = DMGlobalToLocal(user->dm, Q_dot, INSERT_VALUES, Q_dot_loc);
186a515125bSLeila Ghaffari   CHKERRQ(ierr);
187a515125bSLeila Ghaffari 
188a515125bSLeila Ghaffari   // Place PETSc vectors in CEED vectors
189a515125bSLeila Ghaffari   ierr = VecGetArrayReadAndMemType(Q_loc, &q, &q_mem_type); CHKERRQ(ierr);
190a515125bSLeila Ghaffari   ierr = VecGetArrayReadAndMemType(Q_dot_loc, &q_dot, &q_dot_mem_type);
191a515125bSLeila Ghaffari   CHKERRQ(ierr);
192a515125bSLeila Ghaffari   ierr = VecGetArrayAndMemType(G_loc, &g, &g_mem_type); CHKERRQ(ierr);
193a515125bSLeila Ghaffari   CeedVectorSetArray(user->q_ceed, MemTypeP2C(q_mem_type), CEED_USE_POINTER,
194a515125bSLeila Ghaffari                      (PetscScalar *)q);
195a515125bSLeila Ghaffari   CeedVectorSetArray(user->q_dot_ceed, MemTypeP2C(q_dot_mem_type),
196139613f2SLeila Ghaffari                      CEED_USE_POINTER, (PetscScalar *)q_dot);
197a515125bSLeila Ghaffari   CeedVectorSetArray(user->g_ceed, MemTypeP2C(g_mem_type), CEED_USE_POINTER, g);
198a515125bSLeila Ghaffari 
199a515125bSLeila Ghaffari   // Apply CEED operator
200a515125bSLeila Ghaffari   CeedOperatorApply(user->op_ifunction, user->q_ceed, user->g_ceed,
201a515125bSLeila Ghaffari                     CEED_REQUEST_IMMEDIATE);
202a515125bSLeila Ghaffari 
203a515125bSLeila Ghaffari   // Restore vectors
204a515125bSLeila Ghaffari   CeedVectorTakeArray(user->q_ceed, MemTypeP2C(q_mem_type), NULL);
205a515125bSLeila Ghaffari   CeedVectorTakeArray(user->q_dot_ceed, MemTypeP2C(q_dot_mem_type), NULL);
206a515125bSLeila Ghaffari   CeedVectorTakeArray(user->g_ceed, MemTypeP2C(g_mem_type), NULL);
207a515125bSLeila Ghaffari   ierr = VecRestoreArrayReadAndMemType(Q_loc, &q); CHKERRQ(ierr);
208a515125bSLeila Ghaffari   ierr = VecRestoreArrayReadAndMemType(Q_dot_loc, &q_dot); CHKERRQ(ierr);
209a515125bSLeila Ghaffari   ierr = VecRestoreArrayAndMemType(G_loc, &g); CHKERRQ(ierr);
210a515125bSLeila Ghaffari 
211a515125bSLeila Ghaffari   // Local-to-Global
212a515125bSLeila Ghaffari   ierr = VecZeroEntries(G); CHKERRQ(ierr);
213a515125bSLeila Ghaffari   ierr = DMLocalToGlobal(user->dm, G_loc, ADD_VALUES, G); CHKERRQ(ierr);
214a515125bSLeila Ghaffari 
215a515125bSLeila Ghaffari   // Restore vectors
216a515125bSLeila Ghaffari   ierr = DMRestoreLocalVector(user->dm, &G_loc); CHKERRQ(ierr);
217a515125bSLeila Ghaffari 
218a515125bSLeila Ghaffari   PetscFunctionReturn(0);
219a515125bSLeila Ghaffari }
220a515125bSLeila Ghaffari 
221f0b65372SJed Brown static PetscErrorCode MatMult_NS_IJacobian(Mat J, Vec Q, Vec G) {
222f0b65372SJed Brown   User user;
223f0b65372SJed Brown   const PetscScalar *q;
224f0b65372SJed Brown   PetscScalar       *g;
225f0b65372SJed Brown   PetscMemType      q_mem_type, g_mem_type;
226f0b65372SJed Brown   PetscErrorCode    ierr;
227f0b65372SJed Brown   PetscFunctionBeginUser;
228*e2f84137SJeremy L Thompson   ierr = MatShellGetContext(J, &user); CHKERRQ(ierr);
229*e2f84137SJeremy L Thompson   Vec               Q_loc = user->Q_dot_loc, // Note - Q_dot_loc has zero BCs
230*e2f84137SJeremy L Thompson                     G_loc;
231*e2f84137SJeremy L Thompson 
232f0b65372SJed Brown   // Get local vectors
233f0b65372SJed Brown   ierr = DMGetLocalVector(user->dm, &G_loc); CHKERRQ(ierr);
234f0b65372SJed Brown 
235f0b65372SJed Brown   // Global-to-local
236f0b65372SJed Brown   ierr = DMGlobalToLocal(user->dm, Q, INSERT_VALUES, Q_loc); CHKERRQ(ierr);
237f0b65372SJed Brown 
238f0b65372SJed Brown   // Place PETSc vectors in CEED vectors
239f0b65372SJed Brown   ierr = VecGetArrayReadAndMemType(Q_loc, &q, &q_mem_type); CHKERRQ(ierr);
240f0b65372SJed Brown   ierr = VecGetArrayAndMemType(G_loc, &g, &g_mem_type); CHKERRQ(ierr);
241f0b65372SJed Brown   CeedVectorSetArray(user->q_ceed, MemTypeP2C(q_mem_type), CEED_USE_POINTER,
242f0b65372SJed Brown                      (PetscScalar *)q);
243f0b65372SJed Brown   CeedVectorSetArray(user->g_ceed, MemTypeP2C(g_mem_type), CEED_USE_POINTER, g);
244f0b65372SJed Brown 
245f0b65372SJed Brown   // Apply CEED operator
246f0b65372SJed Brown   CeedOperatorApply(user->op_ijacobian, user->q_ceed, user->g_ceed,
247f0b65372SJed Brown                     CEED_REQUEST_IMMEDIATE);
248f0b65372SJed Brown 
249f0b65372SJed Brown   // Restore vectors
250f0b65372SJed Brown   CeedVectorTakeArray(user->q_ceed, MemTypeP2C(q_mem_type), NULL);
251f0b65372SJed Brown   CeedVectorTakeArray(user->g_ceed, MemTypeP2C(g_mem_type), NULL);
252f0b65372SJed Brown   ierr = VecRestoreArrayReadAndMemType(Q_loc, &q); CHKERRQ(ierr);
253f0b65372SJed Brown   ierr = VecRestoreArrayAndMemType(G_loc, &g); CHKERRQ(ierr);
254f0b65372SJed Brown 
255f0b65372SJed Brown   // Local-to-Global
256f0b65372SJed Brown   ierr = VecZeroEntries(G); CHKERRQ(ierr);
257f0b65372SJed Brown   ierr = DMLocalToGlobal(user->dm, G_loc, ADD_VALUES, G); CHKERRQ(ierr);
258f0b65372SJed Brown 
259f0b65372SJed Brown   // Restore vectors
260f0b65372SJed Brown   ierr = DMRestoreLocalVector(user->dm, &G_loc); CHKERRQ(ierr);
261f0b65372SJed Brown   PetscFunctionReturn(0);
262f0b65372SJed Brown }
263f0b65372SJed Brown 
264f0b65372SJed Brown PetscErrorCode MatGetDiagonal_NS_IJacobian(Mat A, Vec D) {
265f0b65372SJed Brown   User user;
266f0b65372SJed Brown   Vec D_loc;
267f0b65372SJed Brown   PetscScalar *d;
268f0b65372SJed Brown   PetscMemType mem_type;
269f0b65372SJed Brown 
270f0b65372SJed Brown   PetscFunctionBeginUser;
271f0b65372SJed Brown   MatShellGetContext(A, &user);
272f0b65372SJed Brown   PetscCall(DMGetLocalVector(user->dm, &D_loc));
273f0b65372SJed Brown   PetscCall(VecGetArrayAndMemType(D_loc, &d, &mem_type));
274f0b65372SJed Brown   CeedVectorSetArray(user->g_ceed, MemTypeP2C(mem_type), CEED_USE_POINTER, d);
275f0b65372SJed Brown   CeedOperatorLinearAssembleDiagonal(user->op_ijacobian, user->g_ceed,
276f0b65372SJed Brown                                      CEED_REQUEST_IMMEDIATE);
277f0b65372SJed Brown   CeedVectorTakeArray(user->g_ceed, MemTypeP2C(mem_type), NULL);
278f0b65372SJed Brown   PetscCall(VecRestoreArrayAndMemType(D_loc, &d));
279f0b65372SJed Brown   PetscCall(VecZeroEntries(D));
280f0b65372SJed Brown   PetscCall(DMLocalToGlobal(user->dm, D_loc, ADD_VALUES, D));
281f0b65372SJed Brown   PetscCall(DMRestoreLocalVector(user->dm, &D_loc));
282f0b65372SJed Brown   VecViewFromOptions(D, NULL, "-diag_vec_view");
283f0b65372SJed Brown   PetscFunctionReturn(0);
284f0b65372SJed Brown }
285f0b65372SJed Brown 
286f0b65372SJed Brown PetscErrorCode FormIJacobian_NS(TS ts, PetscReal t, Vec Q, Vec Q_dot,
287f0b65372SJed Brown                                 PetscReal shift, Mat J, Mat J_pre,
288f0b65372SJed Brown                                 void *user_data) {
289f0b65372SJed Brown   User user = *(User *)user_data;
290f0b65372SJed Brown   PetscBool J_is_shell, J_pre_is_shell;
291f0b65372SJed Brown   PetscFunctionBeginUser;
292f0b65372SJed Brown   if (user->phys->ijacobian_time_shift_label)
293f0b65372SJed Brown     CeedOperatorContextSetDouble(user->op_ijacobian,
294f0b65372SJed Brown                                  user->phys->ijacobian_time_shift_label, &shift);
295f0b65372SJed Brown   PetscCall(MatAssemblyBegin(J, MAT_FINAL_ASSEMBLY));
296f0b65372SJed Brown   PetscCall(MatAssemblyEnd(J, MAT_FINAL_ASSEMBLY));
297f0b65372SJed Brown   Vec coo_vec = NULL;
298f0b65372SJed Brown   PetscCall(PetscObjectTypeCompare((PetscObject)J, MATSHELL, &J_is_shell));
299f0b65372SJed Brown   PetscCall(PetscObjectTypeCompare((PetscObject)J_pre, MATSHELL,
300f0b65372SJed Brown                                    &J_pre_is_shell));
301f0b65372SJed Brown   if (!user->matrices_set_up) {
302f0b65372SJed Brown     if (J_is_shell) {
303f0b65372SJed Brown       PetscCall(MatShellSetContext(J, user));
304f0b65372SJed Brown       PetscCall(MatShellSetOperation(J, MATOP_MULT,
305f0b65372SJed Brown                                      (void (*)(void))MatMult_NS_IJacobian));
306f0b65372SJed Brown       PetscCall(MatShellSetOperation(J, MATOP_GET_DIAGONAL,
307f0b65372SJed Brown                                      (void (*)(void))MatGetDiagonal_NS_IJacobian));
308f0b65372SJed Brown       PetscCall(MatSetUp(J));
309f0b65372SJed Brown     }
310f0b65372SJed Brown     if (!J_pre_is_shell) {
311f0b65372SJed Brown       PetscCount ncoo;
312f0b65372SJed Brown       PetscInt *rows, *cols;
313f0b65372SJed Brown       PetscCall(CeedOperatorLinearAssembleSymbolic(user->op_ijacobian, &ncoo, &rows,
314f0b65372SJed Brown                 &cols));
315f0b65372SJed Brown       PetscCall(MatSetPreallocationCOOLocal(J_pre, ncoo, rows, cols));
316f0b65372SJed Brown       free(rows);
317f0b65372SJed Brown       free(cols);
318f0b65372SJed Brown       CeedVectorCreate(user->ceed, ncoo, &user->coo_values);
319f0b65372SJed Brown       user->matrices_set_up = true;
320f0b65372SJed Brown       VecCreateSeq(PETSC_COMM_WORLD, ncoo, &coo_vec);
321f0b65372SJed Brown     }
322f0b65372SJed Brown   }
323f0b65372SJed Brown   if (!J_pre_is_shell) {
324f0b65372SJed Brown     CeedMemType mem_type = CEED_MEM_HOST;
325f0b65372SJed Brown     const PetscScalar *values;
326f0b65372SJed Brown     MatType mat_type;
327f0b65372SJed Brown     PetscCall(MatGetType(J_pre, &mat_type));
328f0b65372SJed Brown     //if (strstr(mat_type, "kokkos") || strstr(mat_type, "cusparse")) mem_type = CEED_MEM_DEVICE;
329f0b65372SJed Brown     CeedOperatorLinearAssemble(user->op_ijacobian, user->coo_values);
330f0b65372SJed Brown     CeedVectorGetArrayRead(user->coo_values, mem_type, &values);
331f0b65372SJed Brown     if (coo_vec) {
332f0b65372SJed Brown       VecPlaceArray(coo_vec, values);
333f0b65372SJed Brown       VecViewFromOptions(coo_vec, NULL, "-coo_vec_view");
334f0b65372SJed Brown       VecDestroy(&coo_vec);
335f0b65372SJed Brown     }
336f0b65372SJed Brown     PetscCall(MatSetValuesCOO(J_pre, values, INSERT_VALUES));
337f0b65372SJed Brown     CeedVectorRestoreArrayRead(user->coo_values, &values);
338f0b65372SJed Brown   }
339f0b65372SJed Brown   PetscFunctionReturn(0);
340f0b65372SJed Brown }
341f0b65372SJed Brown 
342a515125bSLeila Ghaffari // User provided TS Monitor
343a515125bSLeila Ghaffari PetscErrorCode TSMonitor_NS(TS ts, PetscInt step_no, PetscReal time,
344a515125bSLeila Ghaffari                             Vec Q, void *ctx) {
345a515125bSLeila Ghaffari   User           user = ctx;
346a515125bSLeila Ghaffari   Vec            Q_loc;
347a515125bSLeila Ghaffari   char           file_path[PETSC_MAX_PATH_LEN];
348a515125bSLeila Ghaffari   PetscViewer    viewer;
349a515125bSLeila Ghaffari   PetscErrorCode ierr;
350a515125bSLeila Ghaffari   PetscFunctionBeginUser;
351a515125bSLeila Ghaffari 
352a515125bSLeila Ghaffari   // Print every 'output_freq' steps
353a515125bSLeila Ghaffari   if (step_no % user->app_ctx->output_freq != 0)
354a515125bSLeila Ghaffari     PetscFunctionReturn(0);
355a515125bSLeila Ghaffari 
356a515125bSLeila Ghaffari   // Set up output
357a515125bSLeila Ghaffari   ierr = DMGetLocalVector(user->dm, &Q_loc); CHKERRQ(ierr);
358a515125bSLeila Ghaffari   ierr = PetscObjectSetName((PetscObject)Q_loc, "StateVec"); CHKERRQ(ierr);
359a515125bSLeila Ghaffari   ierr = VecZeroEntries(Q_loc); CHKERRQ(ierr);
360a515125bSLeila Ghaffari   ierr = DMGlobalToLocal(user->dm, Q, INSERT_VALUES, Q_loc); CHKERRQ(ierr);
361a515125bSLeila Ghaffari 
362a515125bSLeila Ghaffari   // Output
363d940ca4eSJed Brown   ierr = PetscSNPrintf(file_path, sizeof file_path,
364d940ca4eSJed Brown                        "%s/ns-%03" PetscInt_FMT ".vtu",
365a515125bSLeila Ghaffari                        user->app_ctx->output_dir, step_no + user->app_ctx->cont_steps);
366a515125bSLeila Ghaffari   CHKERRQ(ierr);
367a515125bSLeila Ghaffari   ierr = PetscViewerVTKOpen(PetscObjectComm((PetscObject)Q), file_path,
368a515125bSLeila Ghaffari                             FILE_MODE_WRITE, &viewer); CHKERRQ(ierr);
369a515125bSLeila Ghaffari   ierr = VecView(Q_loc, viewer); CHKERRQ(ierr);
370a515125bSLeila Ghaffari   ierr = PetscViewerDestroy(&viewer); CHKERRQ(ierr);
371a515125bSLeila Ghaffari   if (user->dm_viz) {
372a515125bSLeila Ghaffari     Vec         Q_refined, Q_refined_loc;
373a515125bSLeila Ghaffari     char        file_path_refined[PETSC_MAX_PATH_LEN];
374a515125bSLeila Ghaffari     PetscViewer viewer_refined;
375a515125bSLeila Ghaffari 
376a515125bSLeila Ghaffari     ierr = DMGetGlobalVector(user->dm_viz, &Q_refined); CHKERRQ(ierr);
377a515125bSLeila Ghaffari     ierr = DMGetLocalVector(user->dm_viz, &Q_refined_loc); CHKERRQ(ierr);
378a515125bSLeila Ghaffari     ierr = PetscObjectSetName((PetscObject)Q_refined_loc, "Refined");
379a515125bSLeila Ghaffari     CHKERRQ(ierr);
380a515125bSLeila Ghaffari     ierr = MatInterpolate(user->interp_viz, Q, Q_refined); CHKERRQ(ierr);
381a515125bSLeila Ghaffari     ierr = VecZeroEntries(Q_refined_loc); CHKERRQ(ierr);
382a515125bSLeila Ghaffari     ierr = DMGlobalToLocal(user->dm_viz, Q_refined, INSERT_VALUES, Q_refined_loc);
383a515125bSLeila Ghaffari     CHKERRQ(ierr);
384a515125bSLeila Ghaffari     ierr = PetscSNPrintf(file_path_refined, sizeof file_path_refined,
385d940ca4eSJed Brown                          "%s/nsrefined-%03" PetscInt_FMT ".vtu", user->app_ctx->output_dir,
386139613f2SLeila Ghaffari                          step_no + user->app_ctx->cont_steps);
387a515125bSLeila Ghaffari     CHKERRQ(ierr);
388a515125bSLeila Ghaffari     ierr = PetscViewerVTKOpen(PetscObjectComm((PetscObject)Q_refined),
389139613f2SLeila Ghaffari                               file_path_refined, FILE_MODE_WRITE, &viewer_refined); CHKERRQ(ierr);
390a515125bSLeila Ghaffari     ierr = VecView(Q_refined_loc, viewer_refined); CHKERRQ(ierr);
391a515125bSLeila Ghaffari     ierr = DMRestoreLocalVector(user->dm_viz, &Q_refined_loc); CHKERRQ(ierr);
392a515125bSLeila Ghaffari     ierr = DMRestoreGlobalVector(user->dm_viz, &Q_refined); CHKERRQ(ierr);
393a515125bSLeila Ghaffari     ierr = PetscViewerDestroy(&viewer_refined); CHKERRQ(ierr);
394a515125bSLeila Ghaffari   }
395a515125bSLeila Ghaffari   ierr = DMRestoreLocalVector(user->dm, &Q_loc); CHKERRQ(ierr);
396a515125bSLeila Ghaffari 
397a515125bSLeila Ghaffari   // Save data in a binary file for continuation of simulations
398a515125bSLeila Ghaffari   ierr = PetscSNPrintf(file_path, sizeof file_path, "%s/ns-solution.bin",
399a515125bSLeila Ghaffari                        user->app_ctx->output_dir); CHKERRQ(ierr);
400a515125bSLeila Ghaffari   ierr = PetscViewerBinaryOpen(user->comm, file_path, FILE_MODE_WRITE, &viewer);
401a515125bSLeila Ghaffari   CHKERRQ(ierr);
402a515125bSLeila Ghaffari   ierr = VecView(Q, viewer); CHKERRQ(ierr);
403a515125bSLeila Ghaffari   ierr = PetscViewerDestroy(&viewer); CHKERRQ(ierr);
404a515125bSLeila Ghaffari 
405a515125bSLeila Ghaffari   // Save time stamp
406a515125bSLeila Ghaffari   // Dimensionalize time back
407a515125bSLeila Ghaffari   time /= user->units->second;
408a515125bSLeila Ghaffari   ierr = PetscSNPrintf(file_path, sizeof file_path, "%s/ns-time.bin",
409a515125bSLeila Ghaffari                        user->app_ctx->output_dir); CHKERRQ(ierr);
410a515125bSLeila Ghaffari   ierr = PetscViewerBinaryOpen(user->comm, file_path, FILE_MODE_WRITE, &viewer);
411a515125bSLeila Ghaffari   CHKERRQ(ierr);
412a515125bSLeila Ghaffari   #if PETSC_VERSION_GE(3,13,0)
413a515125bSLeila Ghaffari   ierr = PetscViewerBinaryWrite(viewer, &time, 1, PETSC_REAL);
414a515125bSLeila Ghaffari   #else
415a515125bSLeila Ghaffari   ierr = PetscViewerBinaryWrite(viewer, &time, 1, PETSC_REAL, true);
416a515125bSLeila Ghaffari   #endif
417a515125bSLeila Ghaffari   CHKERRQ(ierr);
418a515125bSLeila Ghaffari   ierr = PetscViewerDestroy(&viewer); CHKERRQ(ierr);
419a515125bSLeila Ghaffari 
420a515125bSLeila Ghaffari   PetscFunctionReturn(0);
421a515125bSLeila Ghaffari }
422a515125bSLeila Ghaffari 
423a515125bSLeila Ghaffari // TS: Create, setup, and solve
424a515125bSLeila Ghaffari PetscErrorCode TSSolve_NS(DM dm, User user, AppCtx app_ctx, Physics phys,
425a515125bSLeila Ghaffari                           Vec *Q, PetscScalar *f_time, TS *ts) {
426a515125bSLeila Ghaffari   MPI_Comm       comm = user->comm;
427a515125bSLeila Ghaffari   TSAdapt        adapt;
428a515125bSLeila Ghaffari   PetscScalar    final_time;
429a515125bSLeila Ghaffari   PetscErrorCode ierr;
430a515125bSLeila Ghaffari   PetscFunctionBeginUser;
431a515125bSLeila Ghaffari 
432a515125bSLeila Ghaffari   ierr = TSCreate(comm, ts); CHKERRQ(ierr);
433a515125bSLeila Ghaffari   ierr = TSSetDM(*ts, dm); CHKERRQ(ierr);
434a515125bSLeila Ghaffari   if (phys->implicit) {
435a515125bSLeila Ghaffari     ierr = TSSetType(*ts, TSBDF); CHKERRQ(ierr);
436a515125bSLeila Ghaffari     if (user->op_ifunction) {
437a515125bSLeila Ghaffari       ierr = TSSetIFunction(*ts, NULL, IFunction_NS, &user); CHKERRQ(ierr);
438a515125bSLeila Ghaffari     } else { // Implicit integrators can fall back to using an RHSFunction
439a515125bSLeila Ghaffari       ierr = TSSetRHSFunction(*ts, NULL, RHS_NS, &user); CHKERRQ(ierr);
440a515125bSLeila Ghaffari     }
441f0b65372SJed Brown     if (user->op_ijacobian) {
442f0b65372SJed Brown       ierr = DMTSSetIJacobian(dm, FormIJacobian_NS, &user); CHKERRQ(ierr);
443f0b65372SJed Brown     }
444a515125bSLeila Ghaffari   } else {
445a515125bSLeila Ghaffari     if (!user->op_rhs) SETERRQ(comm, PETSC_ERR_ARG_NULL,
446a515125bSLeila Ghaffari                                  "Problem does not provide RHSFunction");
447a515125bSLeila Ghaffari     ierr = TSSetType(*ts, TSRK); CHKERRQ(ierr);
448a515125bSLeila Ghaffari     ierr = TSRKSetType(*ts, TSRK5F); CHKERRQ(ierr);
449a515125bSLeila Ghaffari     ierr = TSSetRHSFunction(*ts, NULL, RHS_NS, &user); CHKERRQ(ierr);
450a515125bSLeila Ghaffari   }
451a515125bSLeila Ghaffari   ierr = TSSetMaxTime(*ts, 500. * user->units->second); CHKERRQ(ierr);
452a515125bSLeila Ghaffari   ierr = TSSetExactFinalTime(*ts, TS_EXACTFINALTIME_STEPOVER); CHKERRQ(ierr);
453a515125bSLeila Ghaffari   ierr = TSSetTimeStep(*ts, 1.e-2 * user->units->second); CHKERRQ(ierr);
454a515125bSLeila Ghaffari   if (app_ctx->test_mode) {ierr = TSSetMaxSteps(*ts, 10); CHKERRQ(ierr);}
455a515125bSLeila Ghaffari   ierr = TSGetAdapt(*ts, &adapt); CHKERRQ(ierr);
456a515125bSLeila Ghaffari   ierr = TSAdaptSetStepLimits(adapt,
457a515125bSLeila Ghaffari                               1.e-12 * user->units->second,
458a515125bSLeila Ghaffari                               1.e2 * user->units->second); CHKERRQ(ierr);
459a515125bSLeila Ghaffari   ierr = TSSetFromOptions(*ts); CHKERRQ(ierr);
460*e2f84137SJeremy L Thompson   user->time = -1.0; // require all BCs and ctx to be updated
461*e2f84137SJeremy L Thompson   user->dt   = -1.0;
462a515125bSLeila Ghaffari   if (!app_ctx->cont_steps) { // print initial condition
463a515125bSLeila Ghaffari     if (!app_ctx->test_mode) {
464a515125bSLeila Ghaffari       ierr = TSMonitor_NS(*ts, 0, 0., *Q, user); CHKERRQ(ierr);
465a515125bSLeila Ghaffari     }
466a515125bSLeila Ghaffari   } else { // continue from time of last output
467a515125bSLeila Ghaffari     PetscReal   time;
468a515125bSLeila Ghaffari     PetscInt    count;
469a515125bSLeila Ghaffari     PetscViewer viewer;
470a515125bSLeila Ghaffari     char        file_path[PETSC_MAX_PATH_LEN];
471a515125bSLeila Ghaffari     ierr = PetscSNPrintf(file_path, sizeof file_path, "%s/ns-time.bin",
472a515125bSLeila Ghaffari                          app_ctx->output_dir); CHKERRQ(ierr);
473a515125bSLeila Ghaffari     ierr = PetscViewerBinaryOpen(comm, file_path, FILE_MODE_READ, &viewer);
474a515125bSLeila Ghaffari     CHKERRQ(ierr);
475a515125bSLeila Ghaffari     ierr = PetscViewerBinaryRead(viewer, &time, 1, &count, PETSC_REAL);
476a515125bSLeila Ghaffari     CHKERRQ(ierr);
477a515125bSLeila Ghaffari     ierr = PetscViewerDestroy(&viewer); CHKERRQ(ierr);
478a515125bSLeila Ghaffari     ierr = TSSetTime(*ts, time * user->units->second); CHKERRQ(ierr);
479a515125bSLeila Ghaffari   }
480a515125bSLeila Ghaffari   if (!app_ctx->test_mode) {
481a515125bSLeila Ghaffari     ierr = TSMonitorSet(*ts, TSMonitor_NS, user, NULL); CHKERRQ(ierr);
482a515125bSLeila Ghaffari   }
483a515125bSLeila Ghaffari 
484a515125bSLeila Ghaffari   // Solve
485a515125bSLeila Ghaffari   double start, cpu_time_used;
486a515125bSLeila Ghaffari   start = MPI_Wtime();
487a515125bSLeila Ghaffari   ierr = PetscBarrier((PetscObject) *ts); CHKERRQ(ierr);
488a515125bSLeila Ghaffari   ierr = TSSolve(*ts, *Q); CHKERRQ(ierr);
489a515125bSLeila Ghaffari   cpu_time_used = MPI_Wtime() - start;
490a515125bSLeila Ghaffari   ierr = TSGetSolveTime(*ts, &final_time); CHKERRQ(ierr);
491a515125bSLeila Ghaffari   *f_time = final_time;
492a515125bSLeila Ghaffari   ierr = MPI_Allreduce(MPI_IN_PLACE, &cpu_time_used, 1, MPI_DOUBLE, MPI_MIN,
493a515125bSLeila Ghaffari                        comm); CHKERRQ(ierr);
494a515125bSLeila Ghaffari   if (!app_ctx->test_mode) {
495a515125bSLeila Ghaffari     ierr = PetscPrintf(PETSC_COMM_WORLD,
496a515125bSLeila Ghaffari                        "Time taken for solution (sec): %g\n",
497a515125bSLeila Ghaffari                        (double)cpu_time_used); CHKERRQ(ierr);
498a515125bSLeila Ghaffari   }
499a515125bSLeila Ghaffari   PetscFunctionReturn(0);
500a515125bSLeila Ghaffari }
501