xref: /honee/qfunctions/sgs_dd_model.h (revision 3ec34a3d2de566dd1c72d57ec6160791c84e75d8)
1 // Copyright (c) 2017-2023, Lawrence Livermore National Security, LLC and other CEED contributors.
2 // All Rights Reserved. See the top-level LICENSE and NOTICE files for details.
3 //
4 // SPDX-License-Identifier: BSD-2-Clause
5 //
6 // This file is part of CEED:  http://github.com/ceed
7 
8 /// @file
9 /// Structs and helper functions to evaluate data-driven subgrid-stress modeling
10 /// See 'Invariant data-driven subgrid stress modeling in the strain-rate eigenframe for large eddy simulation' 2022 and 'S-frame discrepancy
11 /// correction models for data-informed Reynolds stress closure' 2022
12 
13 #ifndef sgs_dd_model_h
14 #define sgs_dd_model_h
15 
16 #include <ceed.h>
17 
18 #include "newtonian_state.h"
19 #include "newtonian_types.h"
20 #include "sgs_dd_utils.h"
21 #include "utils.h"
22 #include "utils_eigensolver_jacobi.h"
23 
24 typedef struct SgsDDContext_ *SgsDDContext;
25 struct SgsDDContext_ {
26   CeedInt    num_inputs, num_outputs;
27   CeedInt    num_layers;
28   CeedInt    num_neurons;
29   CeedScalar alpha;
30 
31   struct NewtonianIdealGasContext_ gas;
32   struct {
33     size_t bias1, bias2;
34     size_t weight1, weight2;
35     size_t out_scaling;
36   } offsets;
37   size_t     total_bytes;
38   CeedScalar data[1];
39 };
40 
41 CEED_QFUNCTION_HELPER void LeakyReLU(CeedScalar *x, const CeedScalar alpha, const CeedInt N) {
42   for (CeedInt i = 0; i < N; i++) x[i] *= (x[i] < 0 ? alpha : 1.);
43 }
44 
45 CEED_QFUNCTION_HELPER void DataDrivenInference(const CeedScalar *inputs, CeedScalar *outputs, SgsDDContext sgsdd_ctx) {
46   const CeedInt     num_neurons = sgsdd_ctx->num_neurons;
47   const CeedInt     num_inputs  = sgsdd_ctx->num_inputs;
48   const CeedInt     num_outputs = sgsdd_ctx->num_outputs;
49   const CeedScalar  alpha       = sgsdd_ctx->alpha;
50   const CeedScalar *bias1       = &sgsdd_ctx->data[sgsdd_ctx->offsets.bias1];
51   const CeedScalar *bias2       = &sgsdd_ctx->data[sgsdd_ctx->offsets.bias2];
52   const CeedScalar *weight1     = &sgsdd_ctx->data[sgsdd_ctx->offsets.weight1];
53   const CeedScalar *weight2     = &sgsdd_ctx->data[sgsdd_ctx->offsets.weight2];
54   CeedScalar        V[20]       = {0.};
55 
56   CopyN(bias1, V, num_neurons);
57   MatVecNM(weight1, inputs, num_neurons, num_inputs, CEED_NOTRANSPOSE, V);
58   LeakyReLU(V, alpha, num_neurons);
59   CopyN(bias2, outputs, num_outputs);
60   MatVecNM(weight2, V, num_outputs, num_neurons, CEED_NOTRANSPOSE, outputs);
61 }
62 
63 CEED_QFUNCTION_HELPER void ComputeSgsDD_Fused(const CeedScalar grad_velo_aniso[3][3], const CeedScalar km_A_ij[6], const CeedScalar delta,
64                                               const CeedScalar viscosity, CeedScalar kmsgs_stress[6], SgsDDContext sgsdd_ctx) {
65   CeedScalar inputs[6], grad_velo_magnitude, eigenvectors[3][3], sgs_sframe_sym[6] = {0.}, new_bounds[6][2];
66   // Copying new_bounds because Sycl online compiler doesn't like direct casting the pointer
67   CopyN(&sgsdd_ctx->data[sgsdd_ctx->offsets.out_scaling], (CeedScalar *)new_bounds, 12);
68 
69   ComputeSgsDDInputs(grad_velo_aniso, km_A_ij, delta, viscosity, eigenvectors, inputs, &grad_velo_magnitude);
70   DataDrivenInference(inputs, sgs_sframe_sym, sgsdd_ctx);
71   ComputeSgsDDOutputs(sgs_sframe_sym, delta, eigenvectors, new_bounds, grad_velo_magnitude, kmsgs_stress);
72 }
73 
74 // @brief Calculate subgrid stress at nodes using anisotropic data-driven model
75 CEED_QFUNCTION_HELPER int ComputeSgsDDNodal_Fused(void *ctx, CeedInt Q, const CeedScalar *const *in, CeedScalar *const *out,
76                                                   StateVariable state_var) {
77   const CeedScalar(*q)[CEED_Q_VLA]            = (const CeedScalar(*)[CEED_Q_VLA])in[0];
78   const CeedScalar(*grad_velo)[3][CEED_Q_VLA] = (const CeedScalar(*)[3][CEED_Q_VLA])in[2];
79   const CeedScalar(*A_ij_delta)[CEED_Q_VLA]   = (const CeedScalar(*)[CEED_Q_VLA])in[3];
80   const CeedScalar(*inv_multiplicity)         = (const CeedScalar(*))in[4];
81   CeedScalar(*v)[CEED_Q_VLA]                  = (CeedScalar(*)[CEED_Q_VLA])out[0];
82 
83   const SgsDDContext             sgsdd_ctx = (SgsDDContext)ctx;
84   const NewtonianIdealGasContext gas       = &sgsdd_ctx->gas;
85 
86   CeedPragmaSIMD for (CeedInt i = 0; i < Q; i++) {
87     const CeedScalar qi[5]                 = {q[0][i], q[1][i], q[2][i], q[3][i], q[4][i]};
88     const CeedScalar grad_velo_aniso[3][3] = {
89         {grad_velo[0][0][i], grad_velo[0][1][i], grad_velo[0][2][i]},
90         {grad_velo[1][0][i], grad_velo[1][1][i], grad_velo[1][2][i]},
91         {grad_velo[2][0][i], grad_velo[2][1][i], grad_velo[2][2][i]}
92     };
93     const CeedScalar km_A_ij[6] = {A_ij_delta[0][i], A_ij_delta[1][i], A_ij_delta[2][i], A_ij_delta[3][i], A_ij_delta[4][i], A_ij_delta[5][i]};
94     const CeedScalar delta      = A_ij_delta[6][i];
95     const State      s          = StateFromQ(gas, qi, state_var);
96     CeedScalar       km_sgs[6];
97 
98     ComputeSgsDD_Fused(grad_velo_aniso, km_A_ij, delta, gas->mu / s.U.density, km_sgs, sgsdd_ctx);
99 
100     for (int j = 0; j < 6; j++) v[j][i] = inv_multiplicity[i] * km_sgs[j];
101   }
102   return 0;
103 }
104 
105 CEED_QFUNCTION(ComputeSgsDDNodal_Prim)(void *ctx, CeedInt Q, const CeedScalar *const *in, CeedScalar *const *out) {
106   return ComputeSgsDDNodal_Fused(ctx, Q, in, out, STATEVAR_PRIMITIVE);
107 }
108 
109 CEED_QFUNCTION(ComputeSgsDDNodal_Conserv)(void *ctx, CeedInt Q, const CeedScalar *const *in, CeedScalar *const *out) {
110   return ComputeSgsDDNodal_Fused(ctx, Q, in, out, STATEVAR_CONSERVATIVE);
111 }
112 
113 // @brief Calculate inputs to anisotropic data-driven model
114 CEED_QFUNCTION_HELPER int ComputeSgsDDNodal_Sequential_Inputs(void *ctx, CeedInt Q, const CeedScalar *const *in, CeedScalar *const *out,
115                                                               StateVariable state_var) {
116   const CeedScalar(*q)[CEED_Q_VLA]            = (const CeedScalar(*)[CEED_Q_VLA])in[0];
117   const CeedScalar(*grad_velo)[3][CEED_Q_VLA] = (const CeedScalar(*)[3][CEED_Q_VLA])in[1];
118   const CeedScalar(*A_ij_delta)[CEED_Q_VLA]   = (const CeedScalar(*)[CEED_Q_VLA])in[2];
119   const CeedScalar(*inv_multiplicity)         = (const CeedScalar(*))in[3];
120   CeedScalar(*eigenvectors_stored)            = out[0];
121   CeedScalar(*model_inputs)[CEED_Q_VLA]       = (CeedScalar(*)[CEED_Q_VLA])out[1];
122 
123   const SgsDDModelContext        sgsdd_ctx = (SgsDDModelContext)ctx;
124   const NewtonianIdealGasContext gas       = &sgsdd_ctx->gas;
125 
126   CeedPragmaSIMD for (CeedInt i = 0; i < Q; i++) {
127     const CeedScalar qi[5]                 = {q[0][i], q[1][i], q[2][i], q[3][i], q[4][i]};
128     const CeedScalar grad_velo_aniso[3][3] = {
129         {grad_velo[0][0][i], grad_velo[0][1][i], grad_velo[0][2][i]},
130         {grad_velo[1][0][i], grad_velo[1][1][i], grad_velo[1][2][i]},
131         {grad_velo[2][0][i], grad_velo[2][1][i], grad_velo[2][2][i]}
132     };
133     const CeedScalar km_A_ij[6] = {A_ij_delta[0][i], A_ij_delta[1][i], A_ij_delta[2][i], A_ij_delta[3][i], A_ij_delta[4][i], A_ij_delta[5][i]};
134     const CeedScalar delta      = A_ij_delta[6][i];
135     const State      s          = StateFromQ(gas, qi, state_var);
136 
137     CeedScalar model_inputs_i[6], grad_velo_magnitude, eigenvectors[3][3];
138     ComputeSgsDDInputs(grad_velo_aniso, km_A_ij, delta, gas->mu / s.U.density, eigenvectors, model_inputs_i, &grad_velo_magnitude);
139 
140     ScaleN(model_inputs_i, inv_multiplicity[i], 6);
141     StoredValuesPack(Q, i, 0, 6, model_inputs_i, (CeedScalar *)model_inputs);
142     StoredValuesPack(Q, i, 0, 9, (const CeedScalar *)eigenvectors, eigenvectors_stored);
143   }
144   return CEED_ERROR_SUCCESS;
145 }
146 
147 CEED_QFUNCTION(ComputeSgsDDNodal_Sequential_Inputs_Prim)(void *ctx, CeedInt Q, const CeedScalar *const *in, CeedScalar *const *out) {
148   return ComputeSgsDDNodal_Sequential_Inputs(ctx, Q, in, out, STATEVAR_PRIMITIVE);
149 }
150 
151 CEED_QFUNCTION(ComputeSgsDDNodal_Sequential_Inputs_Conserv)(void *ctx, CeedInt Q, const CeedScalar *const *in, CeedScalar *const *out) {
152   return ComputeSgsDDNodal_Sequential_Inputs(ctx, Q, in, out, STATEVAR_CONSERVATIVE);
153 }
154 
155 // @brief Calculates SGS from outputs of anisotropic data-driven model
156 CEED_QFUNCTION_HELPER int ComputeSgsDDNodal_Sequential_Outputs(void *ctx, CeedInt Q, const CeedScalar *const *in, CeedScalar *const *out,
157                                                                StateVariable state_var) {
158   const CeedScalar(*model_outputs)[CEED_Q_VLA] = (const CeedScalar(*)[CEED_Q_VLA])in[0];
159   const CeedScalar(*grad_velo)[3][CEED_Q_VLA]  = (const CeedScalar(*)[3][CEED_Q_VLA])in[1];
160   const CeedScalar(*A_ij_delta)[CEED_Q_VLA]    = (const CeedScalar(*)[CEED_Q_VLA])in[2];
161   const CeedScalar(*inv_multiplicity)          = (const CeedScalar(*))in[3];
162   const CeedScalar(*eigenvectors_stored)       = in[4];
163   CeedScalar(*kmsgs_stress)[CEED_Q_VLA]        = (CeedScalar(*)[CEED_Q_VLA])out[0];
164 
165   const SgsDDModelContext sgsdd_ctx = (SgsDDModelContext)ctx;
166   CeedScalar              new_bounds[6][2];
167   CopyN(&sgsdd_ctx->data[sgsdd_ctx->offsets.out_scaling], (CeedScalar *)new_bounds, 12);
168 
169   CeedPragmaSIMD for (CeedInt i = 0; i < Q; i++) {
170     CeedScalar       model_outputs_i[6]    = {model_outputs[0][i], model_outputs[1][i], model_outputs[2][i],
171                                               model_outputs[3][i], model_outputs[4][i], model_outputs[5][i]};
172     const CeedScalar grad_velo_aniso[3][3] = {
173         {grad_velo[0][0][i], grad_velo[0][1][i], grad_velo[0][2][i]},
174         {grad_velo[1][0][i], grad_velo[1][1][i], grad_velo[1][2][i]},
175         {grad_velo[2][0][i], grad_velo[2][1][i], grad_velo[2][2][i]}
176     };
177     const CeedScalar delta = A_ij_delta[6][i];
178 
179     StoredValuesUnpack(Q, i, 0, 6, model_outputs, model_outputs_i);
180     CeedScalar grad_velo_magnitude, eigenvectors[3][3], kmsgs_stress_i[6];
181     StoredValuesUnpack(Q, i, 0, 9, eigenvectors_stored, (CeedScalar *)eigenvectors);
182     grad_velo_magnitude = sqrt(DotN((CeedScalar *)grad_velo_aniso, (CeedScalar *)grad_velo_aniso, 9));
183     ComputeSgsDDOutputs(model_outputs_i, delta, eigenvectors, new_bounds, grad_velo_magnitude, kmsgs_stress_i);
184 
185     for (int j = 0; j < 6; j++) kmsgs_stress[j][i] = inv_multiplicity[i] * kmsgs_stress_i[j];
186   }
187   return CEED_ERROR_SUCCESS;
188 }
189 
190 // @brief Adds subgrid stress to residual (during IFunction evaluation)
191 CEED_QFUNCTION_HELPER int FluxSubgridStress(const StatePrimitive Y, const CeedScalar km_sgs[6], CeedScalar Flux[5][3]) {
192   CeedScalar sgs[3][3];
193 
194   KMUnpack(km_sgs, sgs);
195   for (CeedInt j = 0; j < 3; j++) {
196     Flux[0][j] = 0.;
197     for (CeedInt k = 0; k < 3; k++) Flux[k + 1][j] = sgs[k][j];
198     Flux[4][j] = Y.velocity[0] * sgs[0][j] + Y.velocity[1] * sgs[1][j] + Y.velocity[2] * sgs[2][j];
199   }
200   return 0;
201 }
202 
203 CEED_QFUNCTION_HELPER int IFunction_NodalSgs(void *ctx, CeedInt Q, const CeedScalar *const *in, CeedScalar *const *out, StateVariable state_var) {
204   const CeedScalar(*q)[CEED_Q_VLA]      = (const CeedScalar(*)[CEED_Q_VLA])in[0];
205   const CeedScalar(*q_data)             = in[1];
206   const CeedScalar(*km_sgs)[CEED_Q_VLA] = (const CeedScalar(*)[CEED_Q_VLA])in[2];
207   CeedScalar(*Grad_v)[5][CEED_Q_VLA]    = (CeedScalar(*)[5][CEED_Q_VLA])out[0];
208 
209   NewtonianIdealGasContext gas = (NewtonianIdealGasContext)ctx;
210 
211   CeedPragmaSIMD for (CeedInt i = 0; i < Q; i++) {
212     const CeedScalar qi[5] = {q[0][i], q[1][i], q[2][i], q[3][i], q[4][i]};
213     const State      s     = StateFromQ(gas, qi, state_var);
214 
215     CeedScalar wdetJ, dXdx[3][3];
216     QdataUnpack_3D(Q, i, q_data, &wdetJ, dXdx);
217 
218     CeedScalar       Flux[5][3];
219     const CeedScalar km_sgs_i[6] = {km_sgs[0][i], km_sgs[1][i], km_sgs[2][i], km_sgs[3][i], km_sgs[4][i], km_sgs[5][i]};
220     FluxSubgridStress(s.Y, km_sgs_i, Flux);
221 
222     for (CeedInt k = 0; k < 3; k++) {
223       for (CeedInt j = 0; j < 5; j++) {
224         Grad_v[k][j][i] = -wdetJ * (dXdx[k][0] * Flux[j][0] + dXdx[k][1] * Flux[j][1] + dXdx[k][2] * Flux[j][2]);
225       }
226     }
227   }
228   return 0;
229 }
230 
231 CEED_QFUNCTION(IFunction_NodalSgs_Conserv)(void *ctx, CeedInt Q, const CeedScalar *const *in, CeedScalar *const *out) {
232   return IFunction_NodalSgs(ctx, Q, in, out, STATEVAR_CONSERVATIVE);
233 }
234 
235 CEED_QFUNCTION(IFunction_NodalSgs_Prim)(void *ctx, CeedInt Q, const CeedScalar *const *in, CeedScalar *const *out) {
236   return IFunction_NodalSgs(ctx, Q, in, out, STATEVAR_PRIMITIVE);
237 }
238 
239 #endif  // sgs_dd_model_h
240