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