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 #include "../qfunctions/sgs_dd_model.h" 9 10 #include <petscdmplex.h> 11 12 #include "../navierstokes.h" 13 14 typedef struct { 15 CeedElemRestriction elem_restr_grid_aniso, elem_restr_sgs; 16 CeedVector grid_aniso_ceed; 17 } *SGS_DD_ModelSetupData; 18 19 PetscErrorCode SGS_DD_ModelSetupDataDestroy(SGS_DD_ModelSetupData sgs_dd_setup_data) { 20 PetscFunctionBeginUser; 21 CeedElemRestrictionDestroy(&sgs_dd_setup_data->elem_restr_grid_aniso); 22 CeedElemRestrictionDestroy(&sgs_dd_setup_data->elem_restr_sgs); 23 CeedVectorDestroy(&sgs_dd_setup_data->grid_aniso_ceed); 24 25 PetscCall(PetscFree(sgs_dd_setup_data)); 26 PetscFunctionReturn(0); 27 } 28 29 // @brief B = A^T, A is NxM, B is MxN 30 PetscErrorCode TransposeMatrix(const PetscScalar *A, PetscScalar *B, const PetscInt N, const PetscInt M) { 31 PetscFunctionBeginUser; 32 for (PetscInt i = 0; i < N; i++) { 33 for (PetscInt j = 0; j < M; j++) { 34 B[j * N + i] = A[i * M + j]; 35 } 36 } 37 PetscFunctionReturn(0); 38 } 39 40 // @brief Read neural network coefficients from file and put into context struct 41 PetscErrorCode SGS_DD_ModelContextFill(MPI_Comm comm, char data_dir[PETSC_MAX_PATH_LEN], SGS_DDModelContext *psgsdd_ctx) { 42 SGS_DDModelContext sgsdd_ctx; 43 PetscInt num_inputs = (*psgsdd_ctx)->num_inputs, num_outputs = (*psgsdd_ctx)->num_outputs, num_neurons = (*psgsdd_ctx)->num_neurons; 44 char file_path[PETSC_MAX_PATH_LEN]; 45 PetscScalar *temp; 46 47 PetscFunctionBeginUser; 48 { 49 SGS_DDModelContext sgsdd_temp; 50 PetscCall(PetscNew(&sgsdd_temp)); 51 *sgsdd_temp = **psgsdd_ctx; 52 sgsdd_temp->offsets.bias1 = 0; 53 sgsdd_temp->offsets.bias2 = sgsdd_temp->offsets.bias1 + num_neurons; 54 sgsdd_temp->offsets.weight1 = sgsdd_temp->offsets.bias2 + num_neurons; 55 sgsdd_temp->offsets.weight2 = sgsdd_temp->offsets.weight1 + num_neurons * num_inputs; 56 sgsdd_temp->offsets.out_scaling = sgsdd_temp->offsets.weight2 + num_inputs * num_neurons; 57 PetscInt total_num_scalars = sgsdd_temp->offsets.out_scaling + 2 * num_outputs; 58 sgsdd_temp->total_bytes = sizeof(*sgsdd_ctx) + total_num_scalars * sizeof(sgsdd_ctx->data[0]); 59 PetscCall(PetscMalloc(sgsdd_temp->total_bytes, &sgsdd_ctx)); 60 *sgsdd_ctx = *sgsdd_temp; 61 PetscCall(PetscFree(sgsdd_temp)); 62 } 63 64 PetscCall(PetscSNPrintf(file_path, sizeof file_path, "%s/%s", data_dir, "b1.dat")); 65 PetscCall(PHASTADatFileReadToArrayReal(comm, file_path, &sgsdd_ctx->data[sgsdd_ctx->offsets.bias1])); 66 PetscCall(PetscSNPrintf(file_path, sizeof file_path, "%s/%s", data_dir, "b2.dat")); 67 PetscCall(PHASTADatFileReadToArrayReal(comm, file_path, &sgsdd_ctx->data[sgsdd_ctx->offsets.bias2])); 68 PetscCall(PetscSNPrintf(file_path, sizeof file_path, "%s/%s", data_dir, "OutScaling.dat")); 69 PetscCall(PHASTADatFileReadToArrayReal(comm, file_path, &sgsdd_ctx->data[sgsdd_ctx->offsets.out_scaling])); 70 71 { 72 PetscCall(PetscMalloc1(num_inputs * num_neurons, &temp)); 73 PetscCall(PetscSNPrintf(file_path, sizeof file_path, "%s/%s", data_dir, "w1.dat")); 74 PetscCall(PHASTADatFileReadToArrayReal(comm, file_path, temp)); 75 PetscCall(TransposeMatrix(temp, &sgsdd_ctx->data[sgsdd_ctx->offsets.weight1], num_inputs, num_neurons)); 76 PetscCall(PetscFree(temp)); 77 } 78 { 79 PetscCall(PetscMalloc1(num_outputs * num_neurons, &temp)); 80 PetscCall(PetscSNPrintf(file_path, sizeof file_path, "%s/%s", data_dir, "w2.dat")); 81 PetscCall(PHASTADatFileReadToArrayReal(comm, file_path, temp)); 82 PetscCall(TransposeMatrix(temp, &sgsdd_ctx->data[sgsdd_ctx->offsets.weight2], num_neurons, num_outputs)); 83 PetscCall(PetscFree(temp)); 84 } 85 86 PetscCall(PetscFree(*psgsdd_ctx)); 87 *psgsdd_ctx = sgsdd_ctx; 88 PetscFunctionReturn(0); 89 } 90 91 PetscErrorCode SGS_DD_ModelSetup(Ceed ceed, User user, CeedData ceed_data, ProblemData *problem) { 92 PetscReal alpha; 93 SGS_DDModelContext sgsdd_ctx; 94 MPI_Comm comm = user->comm; 95 char sgs_dd_dir[PETSC_MAX_PATH_LEN] = "./dd_sgs_data"; 96 SGS_DD_ModelSetupData sgs_dd_setup_data; 97 PetscFunctionBeginUser; 98 99 PetscCall(VelocityGradientProjectionSetup(ceed, user, ceed_data, problem)); 100 101 PetscCall(PetscNew(&sgsdd_ctx)); 102 103 PetscOptionsBegin(comm, NULL, "SGS Data-Drive Model Options", NULL); 104 PetscCall(PetscOptionsReal("-sgs_model_dd_leakyrelu_alpha", "Slope parameter for Leaky ReLU activation function", NULL, alpha, &alpha, NULL)); 105 PetscCall(PetscOptionsString("-sgs_model_dd_parameter_dir", "Path to directory with model parameters (weights, biases, etc.)", NULL, sgs_dd_dir, 106 sgs_dd_dir, sizeof(sgs_dd_dir), NULL)); 107 PetscOptionsEnd(); 108 109 sgsdd_ctx->num_layers = 2; 110 sgsdd_ctx->num_inputs = 6; 111 sgsdd_ctx->num_outputs = 6; 112 sgsdd_ctx->num_neurons = 20; 113 sgsdd_ctx->alpha = alpha; 114 115 // PetscCall(SGS_DD_ModelContextFill(comm, sgs_dd_dir, &sgsdd_ctx)); 116 117 PetscCall(PetscNew(&sgs_dd_setup_data)); 118 119 // -- Compute and store anisotropy tensor 120 PetscCall(GridAnisotropyTensorProjectionSetupApply(ceed, user, ceed_data, &sgs_dd_setup_data->elem_restr_grid_aniso, 121 &sgs_dd_setup_data->grid_aniso_ceed)); 122 123 PetscCall(SGS_DD_ModelSetupDataDestroy(sgs_dd_setup_data)); 124 PetscFunctionReturn(0); 125 } 126