xref: /honee/problems/sgs_dd_model.c (revision 67263decd0c197aa72b4d5e710d60290ad1c37bc)
162b7942eSJames Wright // Copyright (c) 2017-2023, Lawrence Livermore National Security, LLC and other CEED contributors.
262b7942eSJames Wright // All Rights Reserved. See the top-level LICENSE and NOTICE files for details.
362b7942eSJames Wright //
462b7942eSJames Wright // SPDX-License-Identifier: BSD-2-Clause
562b7942eSJames Wright //
662b7942eSJames Wright // This file is part of CEED:  http://github.com/ceed
762b7942eSJames Wright 
862b7942eSJames Wright #include "../qfunctions/sgs_dd_model.h"
962b7942eSJames Wright 
1062b7942eSJames Wright #include <petscdmplex.h>
1162b7942eSJames Wright 
1262b7942eSJames Wright #include "../navierstokes.h"
1362b7942eSJames Wright 
14c38c977aSJames Wright typedef struct {
15c38c977aSJames Wright   CeedElemRestriction  elem_restr_grid_aniso, elem_restr_sgs;
16c38c977aSJames Wright   CeedVector           grid_aniso_ceed;
17ee1455b7SJames Wright   CeedQFunctionContext sgsdd_qfctx;
18c38c977aSJames Wright } *SGS_DD_ModelSetupData;
19c38c977aSJames Wright 
20c38c977aSJames Wright PetscErrorCode SGS_DD_ModelSetupDataDestroy(SGS_DD_ModelSetupData sgs_dd_setup_data) {
21c38c977aSJames Wright   PetscFunctionBeginUser;
22c38c977aSJames Wright   CeedElemRestrictionDestroy(&sgs_dd_setup_data->elem_restr_grid_aniso);
23c38c977aSJames Wright   CeedElemRestrictionDestroy(&sgs_dd_setup_data->elem_restr_sgs);
24c38c977aSJames Wright   CeedVectorDestroy(&sgs_dd_setup_data->grid_aniso_ceed);
25ee1455b7SJames Wright   CeedQFunctionContextDestroy(&sgs_dd_setup_data->sgsdd_qfctx);
26c38c977aSJames Wright 
27c38c977aSJames Wright   PetscCall(PetscFree(sgs_dd_setup_data));
28d949ddfcSJames Wright   PetscFunctionReturn(PETSC_SUCCESS);
29c38c977aSJames Wright }
30c38c977aSJames Wright 
31ee1455b7SJames Wright // @brief Create DM for storing subgrid stress at nodes
32*67263decSKenneth E. Jansen PetscErrorCode SGS_DD_ModelCreateDM(DM dm_source, DM *dm_sgs, PetscInt degree, PetscInt q_extra, PetscInt *num_components) {
33ee1455b7SJames Wright   PetscFE      fe;
34ee1455b7SJames Wright   PetscSection section;
35ee1455b7SJames Wright   PetscInt     dim;
36ee1455b7SJames Wright 
37ee1455b7SJames Wright   PetscFunctionBeginUser;
38ee1455b7SJames Wright   *num_components  = 6;
39*67263decSKenneth E. Jansen   PetscInt q_order = degree + q_extra;
40ee1455b7SJames Wright 
41ee1455b7SJames Wright   PetscCall(DMClone(dm_source, dm_sgs));
42ee1455b7SJames Wright   PetscCall(DMGetDimension(*dm_sgs, &dim));
43ee1455b7SJames Wright   PetscCall(PetscObjectSetName((PetscObject)*dm_sgs, "Subgrid Stress Projection"));
44ee1455b7SJames Wright 
45*67263decSKenneth E. Jansen   PetscCall(PetscFECreateLagrange(PETSC_COMM_SELF, dim, *num_components, PETSC_FALSE, degree, q_order, &fe));
46ee1455b7SJames Wright   PetscCall(PetscObjectSetName((PetscObject)fe, "Subgrid Stress Projection"));
47ee1455b7SJames Wright   PetscCall(DMAddField(*dm_sgs, NULL, (PetscObject)fe));
48ee1455b7SJames Wright   PetscCall(DMCreateDS(*dm_sgs));
49ee1455b7SJames Wright   PetscCall(DMPlexSetClosurePermutationTensor(*dm_sgs, PETSC_DETERMINE, NULL));
50ee1455b7SJames Wright 
51ee1455b7SJames Wright   PetscCall(DMGetLocalSection(*dm_sgs, &section));
52ee1455b7SJames Wright   PetscCall(PetscSectionSetFieldName(section, 0, ""));
53ee1455b7SJames Wright   PetscCall(PetscSectionSetComponentName(section, 0, 0, "KMSubgridStressXX"));
54ee1455b7SJames Wright   PetscCall(PetscSectionSetComponentName(section, 0, 1, "KMSubgridStressYY"));
55ee1455b7SJames Wright   PetscCall(PetscSectionSetComponentName(section, 0, 2, "KMSubgridStressZZ"));
56ee1455b7SJames Wright   PetscCall(PetscSectionSetComponentName(section, 0, 3, "KMSubgridStressYZ"));
57ee1455b7SJames Wright   PetscCall(PetscSectionSetComponentName(section, 0, 4, "KMSubgridStressXZ"));
58ee1455b7SJames Wright   PetscCall(PetscSectionSetComponentName(section, 0, 5, "KMSubgridStressXY"));
59ee1455b7SJames Wright 
60ee1455b7SJames Wright   PetscCall(PetscFEDestroy(&fe));
61ee1455b7SJames Wright 
62d949ddfcSJames Wright   PetscFunctionReturn(PETSC_SUCCESS);
63ee1455b7SJames Wright };
64ee1455b7SJames Wright 
65ee1455b7SJames Wright // @brief Create CeedOperator to calculate data-drive SGS at nodes
66ee1455b7SJames Wright PetscErrorCode SGS_DD_ModelSetupNodalEvaluation(Ceed ceed, User user, CeedData ceed_data, SGS_DD_ModelSetupData sgs_dd_setup_data) {
67ee1455b7SJames Wright   SGS_DD_Data         sgs_dd_data = user->sgs_dd_data;
68ee1455b7SJames Wright   CeedQFunction       qf_multiplicity, qf_sgs_dd_nodal;
69ee1455b7SJames Wright   CeedOperator        op_multiplicity, op_sgs_dd_nodal;
70*67263decSKenneth E. Jansen   CeedInt             num_elem, elem_size, num_comp_q, num_comp_grad_velo, num_comp_x, num_comp_grid_aniso;
71defe8520SJames Wright   PetscInt            dim;
724b0f6111SJames Wright   CeedVector          multiplicity, inv_multiplicity;
73ee1455b7SJames Wright   CeedElemRestriction elem_restr_inv_multiplicity, elem_restr_grad_velo, elem_restr_sgs;
74ee1455b7SJames Wright 
75ee1455b7SJames Wright   PetscFunctionBeginUser;
76ee1455b7SJames Wright   PetscCall(DMGetDimension(user->dm, &dim));
77ee1455b7SJames Wright   CeedElemRestrictionGetNumComponents(ceed_data->elem_restr_x, &num_comp_x);
78ee1455b7SJames Wright   CeedElemRestrictionGetNumComponents(ceed_data->elem_restr_q, &num_comp_q);
79ee1455b7SJames Wright   CeedElemRestrictionGetNumComponents(sgs_dd_setup_data->elem_restr_grid_aniso, &num_comp_grid_aniso);
80ee1455b7SJames Wright   CeedElemRestrictionGetNumElements(ceed_data->elem_restr_q, &num_elem);
81ee1455b7SJames Wright   CeedElemRestrictionGetElementSize(ceed_data->elem_restr_q, &elem_size);
82ee1455b7SJames Wright 
83ee1455b7SJames Wright   {  // Get velocity gradient information
84ee1455b7SJames Wright     CeedOperatorField op_field;
85ee1455b7SJames Wright     CeedOperatorGetFieldByName(user->grad_velo_proj->l2_rhs_ctx->op, "velocity gradient", &op_field);
86ee1455b7SJames Wright     CeedOperatorFieldGetElemRestriction(op_field, &elem_restr_grad_velo);
87ee1455b7SJames Wright     CeedElemRestrictionGetNumComponents(elem_restr_grad_velo, &num_comp_grad_velo);
88ee1455b7SJames Wright   }
89*67263decSKenneth E. Jansen   PetscCall(GetRestrictionForDomain(ceed, sgs_dd_data->dm_sgs, 0, 0, 0, 0, -1, 0, &elem_restr_sgs, NULL, NULL));
90ee1455b7SJames Wright   CeedElemRestrictionCreateVector(elem_restr_sgs, &sgs_dd_data->sgs_nodal_ceed, NULL);
91ee1455b7SJames Wright 
92ee1455b7SJames Wright   // -- Create inverse multiplicity for correcting nodal assembly
93ee1455b7SJames Wright   CeedElemRestrictionCreateVector(ceed_data->elem_restr_q, &multiplicity, NULL);
94ee1455b7SJames Wright   CeedElemRestrictionGetMultiplicity(ceed_data->elem_restr_q, multiplicity);
95ee1455b7SJames Wright   CeedElemRestrictionCreateStrided(ceed, num_elem, elem_size, 1, num_elem * elem_size, CEED_STRIDES_BACKEND, &elem_restr_inv_multiplicity);
96ee1455b7SJames Wright   CeedElemRestrictionCreateVector(elem_restr_inv_multiplicity, &inv_multiplicity, NULL);
97ee1455b7SJames Wright 
98ee1455b7SJames Wright   CeedQFunctionCreateInterior(ceed, 1, InverseMultiplicity, InverseMultiplicity_loc, &qf_multiplicity);
99ee1455b7SJames Wright   CeedQFunctionAddInput(qf_multiplicity, "multiplicity", num_comp_q, CEED_EVAL_NONE);
100ee1455b7SJames Wright   CeedQFunctionAddOutput(qf_multiplicity, "inverse multiplicity", 1, CEED_EVAL_NONE);
101ee1455b7SJames Wright 
102ee1455b7SJames Wright   CeedOperatorCreate(ceed, qf_multiplicity, NULL, NULL, &op_multiplicity);
103ee1455b7SJames Wright   CeedOperatorSetName(op_multiplicity, "SGS DD Model - Create Multiplicity Scaling");
104ee1455b7SJames Wright   CeedOperatorSetField(op_multiplicity, "multiplicity", ceed_data->elem_restr_q, CEED_BASIS_COLLOCATED, CEED_VECTOR_ACTIVE);
105ee1455b7SJames Wright   CeedOperatorSetField(op_multiplicity, "inverse multiplicity", elem_restr_inv_multiplicity, CEED_BASIS_COLLOCATED, CEED_VECTOR_ACTIVE);
106ee1455b7SJames Wright 
107ee1455b7SJames Wright   CeedOperatorApply(op_multiplicity, multiplicity, inv_multiplicity, CEED_REQUEST_IMMEDIATE);
108ee1455b7SJames Wright 
109ee1455b7SJames Wright   // -- Create operator for SGS DD model nodal evaluation
110ee1455b7SJames Wright   switch (user->phys->state_var) {
111ee1455b7SJames Wright     case STATEVAR_PRIMITIVE:
112ee1455b7SJames Wright       CeedQFunctionCreateInterior(ceed, 1, ComputeSGS_DDAnisotropicNodal_Prim, ComputeSGS_DDAnisotropicNodal_Prim_loc, &qf_sgs_dd_nodal);
113ee1455b7SJames Wright       break;
114ee1455b7SJames Wright     case STATEVAR_CONSERVATIVE:
115ee1455b7SJames Wright       CeedQFunctionCreateInterior(ceed, 1, ComputeSGS_DDAnisotropicNodal_Conserv, ComputeSGS_DDAnisotropicNodal_Conserv_loc, &qf_sgs_dd_nodal);
116ee1455b7SJames Wright       break;
117ee1455b7SJames Wright     default:
118ee1455b7SJames Wright       SETERRQ(PetscObjectComm((PetscObject)user->dm), PETSC_ERR_SUP,
119ee1455b7SJames Wright               "Anisotropic data-driven SGS nodal evaluation not available for chosen state variable");
120ee1455b7SJames Wright   }
121ee1455b7SJames Wright 
122ee1455b7SJames Wright   // Mesh/geometry order and solution basis order may differ, therefore must interpolate
123ee1455b7SJames Wright   CeedBasis basis_x_to_q;
124ee1455b7SJames Wright   PetscCall(CeedBasisCreateProjection(ceed_data->basis_x, ceed_data->basis_q, &basis_x_to_q));
125ee1455b7SJames Wright 
126ee1455b7SJames Wright   CeedQFunctionSetContext(qf_sgs_dd_nodal, sgs_dd_setup_data->sgsdd_qfctx);
127ee1455b7SJames Wright   CeedQFunctionAddInput(qf_sgs_dd_nodal, "q", num_comp_q, CEED_EVAL_NONE);
128ee1455b7SJames Wright   CeedQFunctionAddInput(qf_sgs_dd_nodal, "x", num_comp_x, CEED_EVAL_INTERP);
129ee1455b7SJames Wright   CeedQFunctionAddInput(qf_sgs_dd_nodal, "gradient velocity", num_comp_grad_velo, CEED_EVAL_NONE);
130ee1455b7SJames Wright   CeedQFunctionAddInput(qf_sgs_dd_nodal, "anisotropy tensor", num_comp_grid_aniso, CEED_EVAL_NONE);
131ee1455b7SJames Wright   CeedQFunctionAddInput(qf_sgs_dd_nodal, "inverse multiplicity", 1, CEED_EVAL_NONE);
132ee1455b7SJames Wright   CeedQFunctionAddOutput(qf_sgs_dd_nodal, "km_sgs", sgs_dd_data->num_comp_sgs, CEED_EVAL_NONE);
133ee1455b7SJames Wright 
134ee1455b7SJames Wright   CeedOperatorCreate(ceed, qf_sgs_dd_nodal, NULL, NULL, &op_sgs_dd_nodal);
135ee1455b7SJames Wright   CeedOperatorSetField(op_sgs_dd_nodal, "q", ceed_data->elem_restr_q, CEED_BASIS_COLLOCATED, user->q_ceed);
136ee1455b7SJames Wright   CeedOperatorSetField(op_sgs_dd_nodal, "x", ceed_data->elem_restr_x, basis_x_to_q, ceed_data->x_coord);
137ee1455b7SJames Wright   CeedOperatorSetField(op_sgs_dd_nodal, "gradient velocity", elem_restr_grad_velo, CEED_BASIS_COLLOCATED, CEED_VECTOR_ACTIVE);
138ee1455b7SJames Wright   CeedOperatorSetField(op_sgs_dd_nodal, "anisotropy tensor", sgs_dd_setup_data->elem_restr_grid_aniso, CEED_BASIS_COLLOCATED,
139ee1455b7SJames Wright                        sgs_dd_setup_data->grid_aniso_ceed);
140ee1455b7SJames Wright   CeedOperatorSetField(op_sgs_dd_nodal, "inverse multiplicity", elem_restr_inv_multiplicity, CEED_BASIS_COLLOCATED, inv_multiplicity);
141ee1455b7SJames Wright   CeedOperatorSetField(op_sgs_dd_nodal, "km_sgs", elem_restr_sgs, CEED_BASIS_COLLOCATED, CEED_VECTOR_ACTIVE);
142ee1455b7SJames Wright 
1434b0f6111SJames Wright   PetscCall(OperatorApplyContextCreate(user->grad_velo_proj->dm, sgs_dd_data->dm_sgs, ceed, op_sgs_dd_nodal, NULL, sgs_dd_data->sgs_nodal_ceed, NULL,
1444b0f6111SJames Wright                                        NULL, &sgs_dd_data->op_nodal_evaluation_ctx));
145ee1455b7SJames Wright 
146ee1455b7SJames Wright   sgs_dd_setup_data->elem_restr_sgs = elem_restr_sgs;
147ee1455b7SJames Wright 
148ee1455b7SJames Wright   CeedVectorDestroy(&multiplicity);
149ee1455b7SJames Wright   CeedVectorDestroy(&inv_multiplicity);
150ee1455b7SJames Wright   CeedBasisDestroy(&basis_x_to_q);
151ee1455b7SJames Wright   CeedElemRestrictionDestroy(&elem_restr_inv_multiplicity);
152ee1455b7SJames Wright   CeedQFunctionDestroy(&qf_multiplicity);
153ee1455b7SJames Wright   CeedQFunctionDestroy(&qf_sgs_dd_nodal);
154ee1455b7SJames Wright   CeedOperatorDestroy(&op_multiplicity);
155ee1455b7SJames Wright   CeedOperatorDestroy(&op_sgs_dd_nodal);
156d949ddfcSJames Wright   PetscFunctionReturn(PETSC_SUCCESS);
157ee1455b7SJames Wright }
158ee1455b7SJames Wright 
1599c678832SJames Wright // @brief Create CeedOperator to compute SGS contribution to the residual
1609c678832SJames Wright PetscErrorCode SGS_ModelSetupNodalIFunction(Ceed ceed, User user, CeedData ceed_data, SGS_DD_ModelSetupData sgs_dd_setup_data) {
1619c678832SJames Wright   SGS_DD_Data   sgs_dd_data = user->sgs_dd_data;
162*67263decSKenneth E. Jansen   CeedInt       num_comp_q, num_comp_qd, num_comp_x;
163defe8520SJames Wright   PetscInt      dim;
1649c678832SJames Wright   CeedQFunction qf_sgs_apply;
1659c678832SJames Wright   CeedOperator  op_sgs_apply;
1669c678832SJames Wright   CeedBasis     basis_sgs;
1679c678832SJames Wright 
1689c678832SJames Wright   PetscFunctionBeginUser;
1699c678832SJames Wright   PetscCall(DMGetDimension(user->dm, &dim));
1709c678832SJames Wright   CeedElemRestrictionGetNumComponents(ceed_data->elem_restr_q, &num_comp_q);
1719c678832SJames Wright   CeedElemRestrictionGetNumComponents(ceed_data->elem_restr_qd_i, &num_comp_qd);
1729c678832SJames Wright   CeedElemRestrictionGetNumComponents(ceed_data->elem_restr_x, &num_comp_x);
1739c678832SJames Wright 
174*67263decSKenneth E. Jansen   PetscCall(CreateBasisFromPlex(ceed, sgs_dd_data->dm_sgs, 0, 0, 0, 0, &basis_sgs));
1759c678832SJames Wright 
1769c678832SJames Wright   switch (user->phys->state_var) {
1779c678832SJames Wright     case STATEVAR_PRIMITIVE:
1789c678832SJames Wright       CeedQFunctionCreateInterior(ceed, 1, IFunction_NodalSubgridStress_Prim, IFunction_NodalSubgridStress_Prim_loc, &qf_sgs_apply);
1799c678832SJames Wright       break;
1809c678832SJames Wright     case STATEVAR_CONSERVATIVE:
1819c678832SJames Wright       CeedQFunctionCreateInterior(ceed, 1, IFunction_NodalSubgridStress_Conserv, IFunction_NodalSubgridStress_Conserv_loc, &qf_sgs_apply);
1829c678832SJames Wright       break;
1839c678832SJames Wright     default:
1849c678832SJames Wright       SETERRQ(PetscObjectComm((PetscObject)user->dm), PETSC_ERR_SUP, "Nodal SGS evaluation not available for chosen state variable");
1859c678832SJames Wright   }
1869c678832SJames Wright 
1879c678832SJames Wright   CeedQFunctionSetContext(qf_sgs_apply, sgs_dd_setup_data->sgsdd_qfctx);
1889c678832SJames Wright   CeedQFunctionAddInput(qf_sgs_apply, "q", num_comp_q, CEED_EVAL_INTERP);
1899c678832SJames Wright   CeedQFunctionAddInput(qf_sgs_apply, "qdata", num_comp_qd, CEED_EVAL_NONE);
1909c678832SJames Wright   CeedQFunctionAddInput(qf_sgs_apply, "x", num_comp_x, CEED_EVAL_INTERP);
1919c678832SJames Wright   CeedQFunctionAddInput(qf_sgs_apply, "km_sgs", sgs_dd_data->num_comp_sgs, CEED_EVAL_INTERP);
1929c678832SJames Wright   CeedQFunctionAddOutput(qf_sgs_apply, "Grad_v", num_comp_q * dim, CEED_EVAL_GRAD);
1939c678832SJames Wright 
1949c678832SJames Wright   CeedOperatorCreate(ceed, qf_sgs_apply, NULL, NULL, &op_sgs_apply);
1959c678832SJames Wright   CeedOperatorSetField(op_sgs_apply, "q", ceed_data->elem_restr_q, ceed_data->basis_q, CEED_VECTOR_ACTIVE);
1969c678832SJames Wright   CeedOperatorSetField(op_sgs_apply, "qdata", ceed_data->elem_restr_qd_i, CEED_BASIS_COLLOCATED, ceed_data->q_data);
1979c678832SJames Wright   CeedOperatorSetField(op_sgs_apply, "x", ceed_data->elem_restr_x, ceed_data->basis_x, ceed_data->x_coord);
1989c678832SJames Wright   CeedOperatorSetField(op_sgs_apply, "km_sgs", sgs_dd_setup_data->elem_restr_sgs, basis_sgs, sgs_dd_data->sgs_nodal_ceed);
1999c678832SJames Wright   CeedOperatorSetField(op_sgs_apply, "Grad_v", ceed_data->elem_restr_q, ceed_data->basis_q, CEED_VECTOR_ACTIVE);
2009c678832SJames Wright 
2019c678832SJames Wright   PetscCall(
2029c678832SJames Wright       OperatorApplyContextCreate(user->dm, user->dm, ceed, op_sgs_apply, user->q_ceed, user->g_ceed, NULL, NULL, &sgs_dd_data->op_sgs_apply_ctx));
2039c678832SJames Wright 
2049c678832SJames Wright   CeedOperatorDestroy(&op_sgs_apply);
2059c678832SJames Wright   CeedQFunctionDestroy(&qf_sgs_apply);
206d949ddfcSJames Wright   PetscFunctionReturn(PETSC_SUCCESS);
2079c678832SJames Wright }
2089c678832SJames Wright 
2099c678832SJames Wright // @brief Calculate and add data-driven SGS residual to the global residual
2109c678832SJames Wright PetscErrorCode SGS_DD_ModelApplyIFunction(User user, const Vec Q_loc, Vec G_loc) {
2119c678832SJames Wright   SGS_DD_Data  sgs_dd_data = user->sgs_dd_data;
2129c678832SJames Wright   Vec          VelocityGradient, SGSNodal_loc;
2139c678832SJames Wright   PetscMemType sgs_nodal_mem_type, q_mem_type;
2149c678832SJames Wright 
2159c678832SJames Wright   PetscFunctionBeginUser;
2169c678832SJames Wright   PetscCall(DMGetGlobalVector(user->grad_velo_proj->dm, &VelocityGradient));
2179c678832SJames Wright   PetscCall(VelocityGradientProjectionApply(user, Q_loc, VelocityGradient));
2189c678832SJames Wright 
2199c678832SJames Wright   // -- Compute Nodal SGS tensor
2209c678832SJames Wright   PetscCall(DMGetLocalVector(sgs_dd_data->dm_sgs, &SGSNodal_loc));
2219c678832SJames Wright   PetscCall(VecP2C(Q_loc, &q_mem_type, user->q_ceed));  // q_ceed is an implicit input
2229c678832SJames Wright 
2239c678832SJames Wright   PetscCall(ApplyCeedOperatorGlobalToLocal(VelocityGradient, SGSNodal_loc, sgs_dd_data->op_nodal_evaluation_ctx));
2249c678832SJames Wright 
2259c678832SJames Wright   PetscCall(VecC2P(user->q_ceed, q_mem_type, Q_loc));
2269c678832SJames Wright   PetscCall(VecP2C(SGSNodal_loc, &sgs_nodal_mem_type, sgs_dd_data->sgs_nodal_ceed));  // sgs_nodal_ceed is an implicit input
2279c678832SJames Wright 
2289c678832SJames Wright   // -- Compute contribution of the SGS stress
2299c678832SJames Wright   PetscCall(ApplyAddCeedOperatorLocalToLocal(Q_loc, G_loc, sgs_dd_data->op_sgs_apply_ctx));
2309c678832SJames Wright 
2319c678832SJames Wright   // -- Return local SGS vector
2329c678832SJames Wright   PetscCall(VecC2P(sgs_dd_data->sgs_nodal_ceed, sgs_nodal_mem_type, SGSNodal_loc));
2339c678832SJames Wright   PetscCall(DMRestoreLocalVector(sgs_dd_data->dm_sgs, &SGSNodal_loc));
2349c678832SJames Wright   PetscCall(DMRestoreGlobalVector(user->grad_velo_proj->dm, &VelocityGradient));
2359c678832SJames Wright 
236d949ddfcSJames Wright   PetscFunctionReturn(PETSC_SUCCESS);
2379c678832SJames Wright }
2389c678832SJames Wright 
23962b7942eSJames Wright // @brief B = A^T, A is NxM, B is MxN
24062b7942eSJames Wright PetscErrorCode TransposeMatrix(const PetscScalar *A, PetscScalar *B, const PetscInt N, const PetscInt M) {
24162b7942eSJames Wright   PetscFunctionBeginUser;
24262b7942eSJames Wright   for (PetscInt i = 0; i < N; i++) {
24362b7942eSJames Wright     for (PetscInt j = 0; j < M; j++) {
24462b7942eSJames Wright       B[j * N + i] = A[i * M + j];
24562b7942eSJames Wright     }
24662b7942eSJames Wright   }
247d949ddfcSJames Wright   PetscFunctionReturn(PETSC_SUCCESS);
24862b7942eSJames Wright }
24962b7942eSJames Wright 
25062b7942eSJames Wright // @brief Read neural network coefficients from file and put into context struct
25162b7942eSJames Wright PetscErrorCode SGS_DD_ModelContextFill(MPI_Comm comm, char data_dir[PETSC_MAX_PATH_LEN], SGS_DDModelContext *psgsdd_ctx) {
25262b7942eSJames Wright   SGS_DDModelContext sgsdd_ctx;
25362b7942eSJames Wright   PetscInt           num_inputs = (*psgsdd_ctx)->num_inputs, num_outputs = (*psgsdd_ctx)->num_outputs, num_neurons = (*psgsdd_ctx)->num_neurons;
25462b7942eSJames Wright   char               file_path[PETSC_MAX_PATH_LEN];
25562b7942eSJames Wright   PetscScalar       *temp;
25662b7942eSJames Wright 
25762b7942eSJames Wright   PetscFunctionBeginUser;
25862b7942eSJames Wright   {
25962b7942eSJames Wright     SGS_DDModelContext sgsdd_temp;
26062b7942eSJames Wright     PetscCall(PetscNew(&sgsdd_temp));
26162b7942eSJames Wright     *sgsdd_temp                     = **psgsdd_ctx;
26262b7942eSJames Wright     sgsdd_temp->offsets.bias1       = 0;
26362b7942eSJames Wright     sgsdd_temp->offsets.bias2       = sgsdd_temp->offsets.bias1 + num_neurons;
26462b7942eSJames Wright     sgsdd_temp->offsets.weight1     = sgsdd_temp->offsets.bias2 + num_neurons;
26562b7942eSJames Wright     sgsdd_temp->offsets.weight2     = sgsdd_temp->offsets.weight1 + num_neurons * num_inputs;
26662b7942eSJames Wright     sgsdd_temp->offsets.out_scaling = sgsdd_temp->offsets.weight2 + num_inputs * num_neurons;
26762b7942eSJames Wright     PetscInt total_num_scalars      = sgsdd_temp->offsets.out_scaling + 2 * num_outputs;
26862b7942eSJames Wright     sgsdd_temp->total_bytes         = sizeof(*sgsdd_ctx) + total_num_scalars * sizeof(sgsdd_ctx->data[0]);
26962b7942eSJames Wright     PetscCall(PetscMalloc(sgsdd_temp->total_bytes, &sgsdd_ctx));
27062b7942eSJames Wright     *sgsdd_ctx = *sgsdd_temp;
27162b7942eSJames Wright     PetscCall(PetscFree(sgsdd_temp));
27262b7942eSJames Wright   }
27362b7942eSJames Wright 
27462b7942eSJames Wright   PetscCall(PetscSNPrintf(file_path, sizeof file_path, "%s/%s", data_dir, "b1.dat"));
27562b7942eSJames Wright   PetscCall(PHASTADatFileReadToArrayReal(comm, file_path, &sgsdd_ctx->data[sgsdd_ctx->offsets.bias1]));
27662b7942eSJames Wright   PetscCall(PetscSNPrintf(file_path, sizeof file_path, "%s/%s", data_dir, "b2.dat"));
27762b7942eSJames Wright   PetscCall(PHASTADatFileReadToArrayReal(comm, file_path, &sgsdd_ctx->data[sgsdd_ctx->offsets.bias2]));
27862b7942eSJames Wright   PetscCall(PetscSNPrintf(file_path, sizeof file_path, "%s/%s", data_dir, "OutScaling.dat"));
27962b7942eSJames Wright   PetscCall(PHASTADatFileReadToArrayReal(comm, file_path, &sgsdd_ctx->data[sgsdd_ctx->offsets.out_scaling]));
28062b7942eSJames Wright 
28162b7942eSJames Wright   {
28262b7942eSJames Wright     PetscCall(PetscMalloc1(num_inputs * num_neurons, &temp));
28362b7942eSJames Wright     PetscCall(PetscSNPrintf(file_path, sizeof file_path, "%s/%s", data_dir, "w1.dat"));
28462b7942eSJames Wright     PetscCall(PHASTADatFileReadToArrayReal(comm, file_path, temp));
28562b7942eSJames Wright     PetscCall(TransposeMatrix(temp, &sgsdd_ctx->data[sgsdd_ctx->offsets.weight1], num_inputs, num_neurons));
28662b7942eSJames Wright     PetscCall(PetscFree(temp));
28762b7942eSJames Wright   }
28862b7942eSJames Wright   {
28962b7942eSJames Wright     PetscCall(PetscMalloc1(num_outputs * num_neurons, &temp));
29062b7942eSJames Wright     PetscCall(PetscSNPrintf(file_path, sizeof file_path, "%s/%s", data_dir, "w2.dat"));
29162b7942eSJames Wright     PetscCall(PHASTADatFileReadToArrayReal(comm, file_path, temp));
29262b7942eSJames Wright     PetscCall(TransposeMatrix(temp, &sgsdd_ctx->data[sgsdd_ctx->offsets.weight2], num_neurons, num_outputs));
29362b7942eSJames Wright     PetscCall(PetscFree(temp));
29462b7942eSJames Wright   }
29562b7942eSJames Wright 
29662b7942eSJames Wright   PetscCall(PetscFree(*psgsdd_ctx));
29762b7942eSJames Wright   *psgsdd_ctx = sgsdd_ctx;
298d949ddfcSJames Wright   PetscFunctionReturn(PETSC_SUCCESS);
29962b7942eSJames Wright }
30062b7942eSJames Wright 
30162b7942eSJames Wright PetscErrorCode SGS_DD_ModelSetup(Ceed ceed, User user, CeedData ceed_data, ProblemData *problem) {
302ee1455b7SJames Wright   PetscReal                alpha = 0;
30362b7942eSJames Wright   SGS_DDModelContext       sgsdd_ctx;
30462b7942eSJames Wright   MPI_Comm                 comm                           = user->comm;
305ee1455b7SJames Wright   char                     sgs_dd_dir[PETSC_MAX_PATH_LEN] = "./dd_sgs_parameters";
306c38c977aSJames Wright   SGS_DD_ModelSetupData    sgs_dd_setup_data;
307ee1455b7SJames Wright   NewtonianIdealGasContext gas;
30862b7942eSJames Wright   PetscFunctionBeginUser;
30962b7942eSJames Wright 
3109ab09d51SJames Wright   PetscCall(VelocityGradientProjectionSetup(ceed, user, ceed_data, problem));
3119ab09d51SJames Wright 
31262b7942eSJames Wright   PetscCall(PetscNew(&sgsdd_ctx));
31362b7942eSJames Wright 
3144b0f6111SJames Wright   PetscOptionsBegin(comm, NULL, "SGS Data-Driven Model Options", NULL);
31562b7942eSJames Wright   PetscCall(PetscOptionsReal("-sgs_model_dd_leakyrelu_alpha", "Slope parameter for Leaky ReLU activation function", NULL, alpha, &alpha, NULL));
31662b7942eSJames Wright   PetscCall(PetscOptionsString("-sgs_model_dd_parameter_dir", "Path to directory with model parameters (weights, biases, etc.)", NULL, sgs_dd_dir,
31762b7942eSJames Wright                                sgs_dd_dir, sizeof(sgs_dd_dir), NULL));
31862b7942eSJames Wright   PetscOptionsEnd();
31962b7942eSJames Wright 
3204b0f6111SJames Wright   sgsdd_ctx->num_layers  = 1;
32162b7942eSJames Wright   sgsdd_ctx->num_inputs  = 6;
32262b7942eSJames Wright   sgsdd_ctx->num_outputs = 6;
32362b7942eSJames Wright   sgsdd_ctx->num_neurons = 20;
32462b7942eSJames Wright   sgsdd_ctx->alpha       = alpha;
32562b7942eSJames Wright 
32601ab89c1SJames Wright   PetscCall(SGS_DD_ModelContextFill(comm, sgs_dd_dir, &sgsdd_ctx));
32762b7942eSJames Wright 
328ee1455b7SJames Wright   // -- Create DM for storing SGS tensor at nodes
329ee1455b7SJames Wright   PetscCall(PetscNew(&user->sgs_dd_data));
330*67263decSKenneth E. Jansen   PetscCall(
331*67263decSKenneth E. Jansen       SGS_DD_ModelCreateDM(user->dm, &user->sgs_dd_data->dm_sgs, user->app_ctx->degree, user->app_ctx->q_extra, &user->sgs_dd_data->num_comp_sgs));
332ee1455b7SJames Wright 
333c38c977aSJames Wright   PetscCall(PetscNew(&sgs_dd_setup_data));
334c38c977aSJames Wright 
335ee1455b7SJames Wright   CeedQFunctionContextGetDataRead(problem->apply_vol_ifunction.qfunction_context, CEED_MEM_HOST, &gas);
336ee1455b7SJames Wright   sgsdd_ctx->gas = *gas;
337ee1455b7SJames Wright   CeedQFunctionContextRestoreDataRead(problem->apply_vol_ifunction.qfunction_context, &gas);
338ee1455b7SJames Wright   CeedQFunctionContextCreate(user->ceed, &sgs_dd_setup_data->sgsdd_qfctx);
339ee1455b7SJames Wright   CeedQFunctionContextSetData(sgs_dd_setup_data->sgsdd_qfctx, CEED_MEM_HOST, CEED_USE_POINTER, sgsdd_ctx->total_bytes, sgsdd_ctx);
340ee1455b7SJames Wright   CeedQFunctionContextSetDataDestroy(sgs_dd_setup_data->sgsdd_qfctx, CEED_MEM_HOST, FreeContextPetsc);
341ee1455b7SJames Wright 
342c38c977aSJames Wright   // -- Compute and store anisotropy tensor
343c38c977aSJames Wright   PetscCall(GridAnisotropyTensorProjectionSetupApply(ceed, user, ceed_data, &sgs_dd_setup_data->elem_restr_grid_aniso,
344c38c977aSJames Wright                                                      &sgs_dd_setup_data->grid_aniso_ceed));
345c38c977aSJames Wright 
346ee1455b7SJames Wright   // -- Create Nodal Evaluation Operator
347ee1455b7SJames Wright   PetscCall(SGS_DD_ModelSetupNodalEvaluation(ceed, user, ceed_data, sgs_dd_setup_data));
348ee1455b7SJames Wright 
3499c678832SJames Wright   // -- Create Operator to evalutate residual of SGS stress
3509c678832SJames Wright   PetscCall(SGS_ModelSetupNodalIFunction(ceed, user, ceed_data, sgs_dd_setup_data));
3519c678832SJames Wright 
352c38c977aSJames Wright   PetscCall(SGS_DD_ModelSetupDataDestroy(sgs_dd_setup_data));
353d949ddfcSJames Wright   PetscFunctionReturn(PETSC_SUCCESS);
35462b7942eSJames Wright }
355ee1455b7SJames Wright 
356ee1455b7SJames Wright PetscErrorCode SGS_DD_DataDestroy(SGS_DD_Data sgs_dd_data) {
357ee1455b7SJames Wright   PetscFunctionBeginUser;
358d949ddfcSJames Wright   if (!sgs_dd_data) PetscFunctionReturn(PETSC_SUCCESS);
359ee1455b7SJames Wright 
360ee1455b7SJames Wright   CeedVectorDestroy(&sgs_dd_data->sgs_nodal_ceed);
361ee1455b7SJames Wright   PetscCall(OperatorApplyContextDestroy(sgs_dd_data->op_nodal_evaluation_ctx));
362ee1455b7SJames Wright   PetscCall(DMDestroy(&sgs_dd_data->dm_sgs));
363ee1455b7SJames Wright   PetscCall(PetscFree(sgs_dd_data));
364ee1455b7SJames Wright 
365d949ddfcSJames Wright   PetscFunctionReturn(PETSC_SUCCESS);
366ee1455b7SJames Wright }
367