1 // SPDX-FileCopyrightText: Copyright (c) 2017-2025, HONEE contributors.
2 // SPDX-License-Identifier: Apache-2.0 OR BSD-2-Clause
3
4 #include <navierstokes.h>
5 #include <smartsim-impl.h>
6
7 typedef struct {
8 PetscInt write_data_interval;
9 size_t local_array_dims[2];
10 PetscBool overwrite_training_data;
11 } *SmartSimSolutionData;
12
TSMonitor_SmartSimSolutionSetup(TS ts,PetscViewerAndFormat * ctx)13 PetscErrorCode TSMonitor_SmartSimSolutionSetup(TS ts, PetscViewerAndFormat *ctx) {
14 SmartSimSolutionData smartsimsol;
15 Honee honee;
16 MPI_Comm comm = PetscObjectComm((PetscObject)ts);
17
18 PetscFunctionBeginUser;
19 PetscCall(TSGetApplicationContext(ts, &honee));
20 PetscCall(PetscNew(&smartsimsol));
21 smartsimsol->overwrite_training_data = PETSC_TRUE;
22
23 PetscOptionsBegin(comm, NULL, "SmartSim Solution Writing", NULL);
24 PetscCall(PetscOptionsBool("-ts_monitor_smartsim_solution_overwrite_data", "Overwrite old solution data in the database", NULL,
25 smartsimsol->overwrite_training_data, &smartsimsol->overwrite_training_data, NULL));
26 PetscOptionsEnd();
27
28 { // Get solution vector size
29 PetscSection output_section;
30 PetscInt num_dofs, num_comps;
31 DM dm, output_dm;
32
33 PetscCall(TSGetDM(ts, &dm));
34 PetscCall(DMGetOutputDM(dm, &output_dm));
35
36 PetscCall(DMGetGlobalVectorInfo(output_dm, &num_dofs, NULL, NULL));
37 PetscCall(DMGetGlobalSection(output_dm, &output_section));
38 PetscCall(PetscSectionGetFieldComponents(output_section, 0, &num_comps));
39 smartsimsol->local_array_dims[0] = num_dofs / num_comps;
40 smartsimsol->local_array_dims[1] = num_comps;
41 }
42 ctx->data = smartsimsol;
43 PetscFunctionReturn(PETSC_SUCCESS);
44 }
45
TSMonitor_SmartSimSolution(TS ts,PetscInt step_num,PetscReal solution_time,Vec Q,PetscViewerAndFormat * ctx)46 PetscErrorCode TSMonitor_SmartSimSolution(TS ts, PetscInt step_num, PetscReal solution_time, Vec Q, PetscViewerAndFormat *ctx) {
47 Honee honee;
48 SmartSimData smartsim;
49 Vec Q_output;
50 PetscMPIInt rank;
51 DM dm, output_dm;
52
53 PetscFunctionBeginUser;
54 SmartSimSolutionData smartsimsol = ctx->data;
55 PetscCall(TSGetApplicationContext(ts, &honee));
56 PetscCall(HoneeGetSmartSimData(honee, &smartsim));
57 PetscCallMPI(MPI_Comm_rank(honee->comm, &rank));
58
59 if (step_num % ctx->view_interval != 0) PetscFunctionReturn(PETSC_SUCCESS);
60
61 PetscCall(TSGetDM(ts, &dm));
62 PetscCall(DMGetOutputDM(dm, &output_dm));
63 PetscCall(DMGetGlobalVector(output_dm, &Q_output));
64
65 PetscCall(UpdateBoundaryValues(honee, honee->Q_loc, solution_time));
66 PetscCall(DMGlobalToLocal(honee->dm, Q, INSERT_VALUES, honee->Q_loc));
67 PetscCall(DMLocalToGlobal(output_dm, honee->Q_loc, INSERT_VALUES, Q_output));
68
69 { // -- Send solution data to SmartSim
70 char array_key[PETSC_MAX_PATH_LEN];
71 size_t array_key_len;
72 void *dataset;
73
74 if (smartsimsol->overwrite_training_data) {
75 PetscCall(PetscSNPrintf(array_key, sizeof array_key, "%s.flow_solution", smartsim->rank_id_name));
76 } else {
77 PetscCall(PetscSNPrintf(array_key, sizeof array_key, "%s.flow_solution.%" PetscInt_FMT, smartsim->rank_id_name, step_num));
78 }
79 PetscCall(PetscStrlen(array_key, &array_key_len));
80 PetscCallSmartRedis(CDataSet(array_key, array_key_len, &dataset));
81
82 {
83 const PetscScalar *Q_output_array;
84 PetscCall(VecGetArrayRead(Q_output, &Q_output_array));
85 PetscCall(PetscLogEventBegin(HONEE_SmartRedis_Write, 0, 0, 0, 0));
86 PetscCallSmartRedis(add_tensor(dataset, "solution", 8, (void *)Q_output_array, smartsimsol->local_array_dims, 2, SRTensorTypeDouble,
87 SRMemLayoutContiguous));
88 PetscCall(PetscLogEventEnd(HONEE_SmartRedis_Write, 0, 0, 0, 0));
89 PetscCall(VecRestoreArrayRead(Q_output, &Q_output_array));
90 }
91
92 PetscCallSmartRedis(add_meta_scalar(dataset, "step", 4, (void *)&step_num, SRMetadataTypePetscInt));
93 PetscCallSmartRedis(add_meta_scalar(dataset, "time", 4, (void *)&solution_time, SRMetadataTypeDouble));
94 PetscCall(PetscLogEventBegin(HONEE_SmartRedis_Write, 0, 0, 0, 0));
95 PetscCallSmartRedis(put_dataset(smartsim->client, dataset));
96 PetscCall(PetscLogEventEnd(HONEE_SmartRedis_Write, 0, 0, 0, 0));
97 PetscCallSmartRedis(DeallocateDataSet(&dataset));
98 }
99
100 PetscCall(DMRestoreGlobalVector(output_dm, &Q_output));
101 PetscFunctionReturn(PETSC_SUCCESS);
102 }
103