xref: /honee/src/smartsim/solution.c (revision c4bd9656417ac488adba016fe5c2e786414c4baf)
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 
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 
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(
87           add_tensor(dataset, "solution", 8, (void *)Q_output_array, smartsimsol->local_array_dims, 2, SRTensorTypeDouble, 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