xref: /honee/qfunctions/sgs_dd_model.h (revision 6cbc221e08f36ca553e7bfe766a8cbdb252ad96a)
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 for 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 "utils.h"
21 #include "utils_eigensolver_jacobi.h"
22 
23 typedef struct SGS_DD_ModelContext_ *SGS_DDModelContext;
24 struct SGS_DD_ModelContext_ {
25   CeedInt    num_inputs, num_outputs;
26   CeedInt    num_layers;
27   CeedInt    num_neurons;
28   CeedScalar alpha;
29 
30   struct NewtonianIdealGasContext_ gas;
31   struct {
32     size_t bias1, bias2;
33     size_t weight1, weight2;
34     size_t out_scaling;
35   } offsets;
36   size_t     total_bytes;
37   CeedScalar data[1];
38 };
39 
40 // @brief Calculate the inverse of the multiplicity, reducing to a single component
41 CEED_QFUNCTION(InverseMultiplicity)(void *ctx, CeedInt Q, const CeedScalar *const *in, CeedScalar *const *out) {
42   const CeedScalar(*multiplicity)[CEED_Q_VLA] = (const CeedScalar(*)[CEED_Q_VLA])in[0];
43   CeedScalar(*inv_multiplicity)               = (CeedScalar(*))out[0];
44 
45   CeedPragmaSIMD for (CeedInt i = 0; i < Q; i++) inv_multiplicity[i] = 1.0 / multiplicity[0][i];
46   return 0;
47 }
48 
49 // @brief Calculate Frobenius norm of velocity gradient from eigenframe quantities
50 CEED_QFUNCTION_HELPER CeedScalar VelocityGradientMagnitude(const CeedScalar strain_sframe[3], const CeedScalar vorticity_sframe[3]) {
51   return sqrt(Dot3(strain_sframe, strain_sframe) + 0.5 * Dot3(vorticity_sframe, vorticity_sframe));
52 };
53 
54 // @brief Denormalize outputs using min-max (de-)normalization
55 CEED_QFUNCTION_HELPER void DenormalizeDDOutputs(CeedScalar output[6], const CeedScalar new_bounds[6][2], const CeedScalar old_bounds[6][2]) {
56   CeedScalar bounds_ratio;
57   for (int i = 0; i < 6; i++) {
58     bounds_ratio = (new_bounds[i][1] - new_bounds[i][0]) / (old_bounds[i][1] - old_bounds[i][0]);
59     output[i]    = bounds_ratio * (output[i] - old_bounds[i][1]) + new_bounds[i][1];
60   }
61 }
62 
63 // @brief Change the order of basis vectors so that they align with vector and obey right-hand rule
64 // @details The e_1 and e_3 basis vectors are the closest aligned to the vector. The e_2 is set via  e_3 x e_1
65 // The basis vectors are assumed to form the rows of the basis matrix.
66 CEED_QFUNCTION_HELPER void OrientBasisWithVector(CeedScalar basis[3][3], const CeedScalar vector[3]) {
67   CeedScalar alignment[3] = {0.}, cross[3];
68 
69   MatVec3(basis, vector, CEED_NOTRANSPOSE, alignment);
70 
71   if (alignment[0] < 0) ScaleN(basis[0], -1, 3);
72   if (alignment[2] < 0) ScaleN(basis[2], -1, 3);
73 
74   Cross3(basis[2], basis[0], cross);
75   CeedScalar basis_1_orientation = Dot3(cross, basis[1]);
76   if (basis_1_orientation < 0) ScaleN(basis[1], -1, 3);
77 }
78 
79 CEED_QFUNCTION_HELPER void LeakyReLU(CeedScalar *x, const CeedScalar alpha, const CeedInt N) {
80   for (CeedInt i = 0; i < N; i++) x[i] *= (x[i] < 0 ? alpha : 1.);
81 }
82 
83 CEED_QFUNCTION_HELPER void DataDrivenInference(const CeedScalar *inputs, CeedScalar *outputs, SGS_DDModelContext sgsdd_ctx) {
84   const CeedInt     num_neurons = sgsdd_ctx->num_neurons;
85   const CeedInt     num_inputs  = sgsdd_ctx->num_inputs;
86   const CeedInt     num_outputs = sgsdd_ctx->num_outputs;
87   const CeedScalar  alpha       = sgsdd_ctx->alpha;
88   const CeedScalar *bias1       = &sgsdd_ctx->data[sgsdd_ctx->offsets.bias1];
89   const CeedScalar *bias2       = &sgsdd_ctx->data[sgsdd_ctx->offsets.bias2];
90   const CeedScalar *weight1     = &sgsdd_ctx->data[sgsdd_ctx->offsets.weight1];
91   const CeedScalar *weight2     = &sgsdd_ctx->data[sgsdd_ctx->offsets.weight2];
92   CeedScalar        V[20]       = {0.};
93 
94   CopyN(bias1, V, num_neurons);
95   MatVecNM(weight1, inputs, num_neurons, num_inputs, CEED_NOTRANSPOSE, V);
96   LeakyReLU(V, alpha, num_neurons);
97   CopyN(bias2, outputs, num_outputs);
98   MatVecNM(weight2, V, num_outputs, num_neurons, CEED_NOTRANSPOSE, outputs);
99 }
100 
101 CEED_QFUNCTION_HELPER void ComputeSGS_DDAnisotropic(const CeedScalar grad_velo_aniso[3][3], const CeedScalar km_A_ij[6], const CeedScalar delta,
102                                                     const CeedScalar viscosity, CeedScalar kmsgs_stress[6], SGS_DDModelContext sgsdd_ctx) {
103   CeedScalar strain_sframe[3] = {0.}, vorticity_sframe[3] = {0.}, eigenvectors[3][3];
104   CeedScalar A_ij[3][3] = {{0.}}, grad_velo_iso[3][3] = {{0.}};
105 
106   // -- Unpack anisotropy tensor
107   KMUnpack(km_A_ij, A_ij);
108 
109   // -- Transform physical, anisotropic velocity gradient to isotropic
110   MatMat3(grad_velo_aniso, A_ij, CEED_NOTRANSPOSE, CEED_NOTRANSPOSE, grad_velo_iso);
111 
112   {  // -- Get Eigenframe
113     CeedScalar kmstrain_iso[6], strain_iso[3][3];
114     CeedInt    work_vector[3] = {0};
115     KMStrainRate(grad_velo_iso, kmstrain_iso);
116     KMUnpack(kmstrain_iso, strain_iso);
117     Diagonalize3(strain_iso, strain_sframe, eigenvectors, work_vector, SORT_DECREASING_EVALS, true, 5);
118   }
119 
120   {  // -- Get vorticity in S-frame
121     CeedScalar rotation_iso[3][3];
122     RotationRate(grad_velo_iso, rotation_iso);
123     CeedScalar vorticity_iso[3] = {-2 * rotation_iso[1][2], 2 * rotation_iso[0][2], -2 * rotation_iso[0][1]};
124     OrientBasisWithVector(eigenvectors, vorticity_iso);
125     MatVec3(eigenvectors, vorticity_iso, CEED_NOTRANSPOSE, vorticity_sframe);
126   }
127 
128   // -- Setup DD model inputs
129   const CeedScalar grad_velo_magnitude = VelocityGradientMagnitude(strain_sframe, vorticity_sframe);
130   CeedScalar inputs[6] = {strain_sframe[0], strain_sframe[1], strain_sframe[2], vorticity_sframe[0], vorticity_sframe[1], viscosity / Square(delta)};
131   ScaleN(inputs, 1 / (grad_velo_magnitude + CEED_EPSILON), 6);
132 
133   CeedScalar sgs_sframe_sym[6] = {0.};
134   DataDrivenInference(inputs, sgs_sframe_sym, sgsdd_ctx);
135 
136   CeedScalar old_bounds[6][2] = {{0}};
137   for (int j = 0; j < 6; j++) old_bounds[j][1] = 1;
138   const CeedScalar(*new_bounds)[2] = (const CeedScalar(*)[2]) & sgsdd_ctx->data[sgsdd_ctx->offsets.out_scaling];
139   DenormalizeDDOutputs(sgs_sframe_sym, new_bounds, old_bounds);
140 
141   // Re-dimensionalize sgs_stress
142   ScaleN(sgs_sframe_sym, Square(delta) * Square(grad_velo_magnitude), 6);
143 
144   CeedScalar sgs_stress[3][3] = {{0.}};
145   {  // Rotate SGS Stress back to physical frame, SGS_physical = E^T SGS_sframe E
146     CeedScalar       Evec_sgs[3][3]   = {{0.}};
147     const CeedScalar sgs_sframe[3][3] = {
148         {sgs_sframe_sym[0], sgs_sframe_sym[3], sgs_sframe_sym[4]},
149         {sgs_sframe_sym[3], sgs_sframe_sym[1], sgs_sframe_sym[5]},
150         {sgs_sframe_sym[4], sgs_sframe_sym[5], sgs_sframe_sym[2]},
151     };
152     MatMat3(eigenvectors, sgs_sframe, CEED_TRANSPOSE, CEED_NOTRANSPOSE, Evec_sgs);
153     MatMat3(Evec_sgs, eigenvectors, CEED_NOTRANSPOSE, CEED_NOTRANSPOSE, sgs_stress);
154   }
155 
156   KMPack(sgs_stress, kmsgs_stress);
157 }
158 
159 // @brief Calculate subgrid stress at nodes using anisotropic data-driven model
160 CEED_QFUNCTION_HELPER int ComputeSGS_DDAnisotropicNodal(void *ctx, CeedInt Q, const CeedScalar *const *in, CeedScalar *const *out,
161                                                         StateFromQi_t StateFromQi) {
162   const CeedScalar(*q)[CEED_Q_VLA]            = (const CeedScalar(*)[CEED_Q_VLA])in[0];
163   const CeedScalar(*x)[CEED_Q_VLA]            = (const CeedScalar(*)[CEED_Q_VLA])in[1];
164   const CeedScalar(*grad_velo)[3][CEED_Q_VLA] = (const CeedScalar(*)[3][CEED_Q_VLA])in[2];
165   const CeedScalar(*A_ij_delta)[CEED_Q_VLA]   = (const CeedScalar(*)[CEED_Q_VLA])in[3];
166   const CeedScalar(*inv_multiplicity)         = (const CeedScalar(*))in[4];
167   CeedScalar(*v)[CEED_Q_VLA]                  = (CeedScalar(*)[CEED_Q_VLA])out[0];
168 
169   const SGS_DDModelContext       sgsdd_ctx = (SGS_DDModelContext)ctx;
170   const NewtonianIdealGasContext gas       = &sgsdd_ctx->gas;
171 
172   CeedPragmaSIMD for (CeedInt i = 0; i < Q; i++) {
173     const CeedScalar qi[5]                 = {q[0][i], q[1][i], q[2][i], q[3][i], q[4][i]};
174     const CeedScalar x_i[3]                = {x[0][i], x[1][i], x[2][i]};
175     const CeedScalar grad_velo_aniso[3][3] = {
176         {grad_velo[0][0][i], grad_velo[0][1][i], grad_velo[0][2][i]},
177         {grad_velo[1][0][i], grad_velo[1][1][i], grad_velo[1][2][i]},
178         {grad_velo[2][0][i], grad_velo[2][1][i], grad_velo[2][2][i]}
179     };
180     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]};
181     const CeedScalar delta      = A_ij_delta[6][i];
182     const State      s          = StateFromQi(gas, qi, x_i);
183     CeedScalar       km_sgs[6];
184 
185     ComputeSGS_DDAnisotropic(grad_velo_aniso, km_A_ij, delta, gas->mu / s.U.density, km_sgs, sgsdd_ctx);
186 
187     for (int j = 0; j < 6; j++) v[j][i] = inv_multiplicity[i] * km_sgs[j];
188   }
189   return 0;
190 }
191 
192 CEED_QFUNCTION(ComputeSGS_DDAnisotropicNodal_Prim)(void *ctx, CeedInt Q, const CeedScalar *const *in, CeedScalar *const *out) {
193   return ComputeSGS_DDAnisotropicNodal(ctx, Q, in, out, StateFromY);
194 }
195 
196 CEED_QFUNCTION(ComputeSGS_DDAnisotropicNodal_Conserv)(void *ctx, CeedInt Q, const CeedScalar *const *in, CeedScalar *const *out) {
197   return ComputeSGS_DDAnisotropicNodal(ctx, Q, in, out, StateFromU);
198 }
199 
200 // @brief Adds subgrid stress to residual (during IFunction evaluation)
201 CEED_QFUNCTION_HELPER int FluxSubgridStress(const StatePrimitive Y, const CeedScalar km_sgs[6], CeedScalar Flux[5][3]) {
202   CeedScalar sgs[3][3];
203 
204   KMUnpack(km_sgs, sgs);
205   for (CeedInt j = 0; j < 3; j++) {
206     Flux[0][j] = 0.;
207     for (CeedInt k = 0; k < 3; k++) Flux[k + 1][j] = sgs[k][j];
208     Flux[4][j] = Y.velocity[0] * sgs[0][j] + Y.velocity[1] * sgs[1][j] + Y.velocity[2] * sgs[2][j];
209   }
210   return 0;
211 }
212 
213 CEED_QFUNCTION_HELPER int IFunction_NodalSubgridStress(void *ctx, CeedInt Q, const CeedScalar *const *in, CeedScalar *const *out,
214                                                        StateFromQi_t StateFromQi, StateFromQi_fwd_t StateFromQi_fwd) {
215   const CeedScalar(*q)[CEED_Q_VLA]      = (const CeedScalar(*)[CEED_Q_VLA])in[0];
216   const CeedScalar(*q_data)[CEED_Q_VLA] = (const CeedScalar(*)[CEED_Q_VLA])in[1];
217   const CeedScalar(*x)[CEED_Q_VLA]      = (const CeedScalar(*)[CEED_Q_VLA])in[2];
218   const CeedScalar(*km_sgs)[CEED_Q_VLA] = (const CeedScalar(*)[CEED_Q_VLA])in[3];
219   CeedScalar(*Grad_v)[5][CEED_Q_VLA]    = (CeedScalar(*)[5][CEED_Q_VLA])out[0];
220 
221   SGS_DDModelContext       sgsdd_ctx = (SGS_DDModelContext)ctx;
222   NewtonianIdealGasContext gas       = &sgsdd_ctx->gas;
223 
224   CeedPragmaSIMD for (CeedInt i = 0; i < Q; i++) {
225     const CeedScalar qi[5]  = {q[0][i], q[1][i], q[2][i], q[3][i], q[4][i]};
226     const CeedScalar x_i[3] = {x[0][i], x[1][i], x[2][i]};
227     const State      s      = StateFromQi(gas, qi, x_i);
228 
229     const CeedScalar wdetJ      = q_data[0][i];
230     const CeedScalar dXdx[3][3] = {
231         {q_data[1][i], q_data[2][i], q_data[3][i]},
232         {q_data[4][i], q_data[5][i], q_data[6][i]},
233         {q_data[7][i], q_data[8][i], q_data[9][i]}
234     };
235 
236     CeedScalar       Flux[5][3];
237     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]};
238     FluxSubgridStress(s.Y, km_sgs_i, Flux);
239 
240     for (CeedInt j = 0; j < 3; j++) {
241       for (CeedInt k = 0; k < 5; k++) {
242         Grad_v[j][k][i] = -wdetJ * (dXdx[j][0] * Flux[k][0] + dXdx[j][1] * Flux[k][1] + dXdx[j][2] * Flux[k][2]);
243       }
244     }
245   }
246   return 0;
247 }
248 
249 CEED_QFUNCTION(IFunction_NodalSubgridStress_Conserv)(void *ctx, CeedInt Q, const CeedScalar *const *in, CeedScalar *const *out) {
250   return IFunction_NodalSubgridStress(ctx, Q, in, out, StateFromU, StateFromU_fwd);
251 }
252 
253 CEED_QFUNCTION(IFunction_NodalSubgridStress_Prim)(void *ctx, CeedInt Q, const CeedScalar *const *in, CeedScalar *const *out) {
254   return IFunction_NodalSubgridStress(ctx, Q, in, out, StateFromY, StateFromY_fwd);
255 }
256 
257 #endif  // sgs_dd_model_h
258