Line data    Source code

       1             : /**********************************************************************/
       2             : /*                    DO NOT MODIFY THIS HEADER                       */
       3             : /*             Swift, a Fourier spectral solver for MOOSE             */
       4             : /*                                                                    */
       5             : /*            Copyright 2024 Battelle Energy Alliance, LLC            */
       6             : /*                        ALL RIGHTS RESERVED                         */
       7             : /**********************************************************************/
       8             : 
       9             : #include "LBMTensorBuffer.h"
      10             : #include "DomainAction.h"
      11             : #include "LatticeBoltzmannStencilBase.h"
      12             : #include "LatticeBoltzmannProblem.h"
      13             : 
      14             : #ifdef LIBMESH_HAVE_HDF5
      15             : #include "hdf5.h"
      16             : #endif
      17             : 
      18             : registerMooseObject("SwiftApp", LBMTensorBuffer);
      19             : 
      20             : InputParameters
      21           0 : LBMTensorBuffer::validParams()
      22             : {
      23             :   InputParameters params = TensorBuffer<torch::Tensor>::validParams();
      24           0 :   params.addRequiredParam<std::string>("buffer_type",
      25             :                                        "The buffer type can be either distribution function (df), "
      26             :                                        "macroscopic scalar (ms) or macroscopic vectorial (mv)");
      27             : 
      28           0 :   params.addParam<FileName>("file", "Optional path of the file to read tensor form.");
      29             : 
      30           0 :   params.addParam<bool>("is_integer", false, "Whether to specify integer dtype");
      31           0 :   params.addPrivateParam<TensorProblem *>("_tensor_problem", nullptr);
      32           0 :   params.addClassDescription("Tensor wrapper form LBM tensors");
      33             : 
      34           0 :   return params;
      35           0 : }
      36             : 
      37           0 : LBMTensorBuffer::LBMTensorBuffer(const InputParameters & parameters)
      38             :   : TensorBuffer<torch::Tensor>(parameters),
      39           0 :     _buffer_type(getParam<std::string>("buffer_type")),
      40           0 :     _lb_problem(dynamic_cast<LatticeBoltzmannProblem &>(
      41           0 :         *getCheckedPointerParam<TensorProblem *>("_tensor_problem"))),
      42           0 :     _stencil(_lb_problem.getStencil())
      43             : {
      44           0 : }
      45             : 
      46             : void
      47           0 : LBMTensorBuffer::init()
      48             : {
      49             :   int64_t dimension = 0;
      50           0 :   if (_buffer_type == "df")
      51           0 :     dimension = _stencil._q;
      52           0 :   else if (_buffer_type == "mv")
      53           0 :     dimension = _domain.getDim();
      54           0 :   else if (_buffer_type == "ms")
      55             :     dimension = 0;
      56             :   else
      57           0 :     mooseError("Buffer type ", _buffer_type, " is not recognized");
      58             : 
      59           0 :   std::vector<int64_t> shape(_domain.getShape().begin(), _domain.getShape().end());
      60             : 
      61           0 :   if (_domain.getDim() < 3)
      62           0 :     shape.push_back(1);
      63           0 :   if (dimension > 0)
      64           0 :     shape.push_back(static_cast<int64_t>(dimension));
      65             : 
      66           0 :   if (getParam<bool>("is_integer"))
      67           0 :     _u = torch::zeros(shape, MooseTensor::intTensorOptions());
      68             :   else
      69           0 :     _u = torch::zeros(shape, MooseTensor::floatTensorOptions());
      70             : 
      71           0 :   if (isParamValid("file"))
      72           0 :     readTensorFromHdf5();
      73           0 : }
      74             : 
      75             : void
      76           0 : LBMTensorBuffer::readTensorFromFile(const std::vector<int64_t> & shape)
      77             : {
      78           0 :   mooseDeprecated("readTensorFromFile is deprecated, use h5 reader readTensorFromHdf5 instead!");
      79             : 
      80           0 :   const FileName tensor_file = getParam<FileName>("file");
      81           0 :   mooseInfo("Loading tensor(s) from file \n" + tensor_file);
      82           0 :   std::ifstream file(tensor_file);
      83           0 :   if (!file.is_open())
      84           0 :     mooseError("Cannot open file " + tensor_file);
      85             : 
      86             :   // read file into standart vector
      87           0 :   std::vector<Real> fileData(shape[0] * shape[1] * shape[2]);
      88             : 
      89           0 :   for (unsigned int i = 0; i < fileData.size(); i++)
      90           0 :     if (!(file >> fileData[i]))
      91           0 :       mooseError("Insufficient data in the file");
      92             : 
      93           0 :   file.close();
      94             : 
      95             :   // reshape and write into torch tensor
      96           0 :   for (int64_t k = 0; k < shape[2]; k++)
      97           0 :     for (int64_t j = 0; j < shape[1]; j++)
      98           0 :       for (int64_t i = 0; i < shape[0]; i++)
      99             :       {
     100           0 :         if (getParam<bool>("is_integer"))
     101           0 :           _u.index_put_({i, j, k},
     102           0 :                         static_cast<int>(fileData[k * shape[1] * shape[0] + j * shape[0] + i]));
     103             :         else
     104           0 :           _u.index_put_({i, j, k}, fileData[k * shape[1] * shape[0] + j * shape[0] + i]);
     105             :       }
     106           0 : }
     107             : 
     108             : void
     109           0 : LBMTensorBuffer::readTensorFromHdf5()
     110             : {
     111             : #ifdef LIBMESH_HAVE_HDF5
     112           0 :   const FileName tensor_file_name = getParam<FileName>("file");
     113             : 
     114             :   auto tensor_file_char = tensor_file_name.c_str();
     115             : 
     116             :   // open file
     117           0 :   hid_t file_id = H5Fopen(tensor_file_char, H5F_ACC_RDONLY, H5P_DEFAULT);
     118           0 :   if (file_id < 0)
     119           0 :     mooseError("Failed to open h5 file");
     120             : 
     121           0 :   std::string dataset_name = tensor_file_name.substr(0, tensor_file_name.size() - 3);
     122           0 :   auto last_slash = dataset_name.find_last_of("/\\");
     123           0 :   if (last_slash != std::string::npos)
     124           0 :     dataset_name = dataset_name.substr(last_slash + 1);
     125             :   auto dataset_name_char = dataset_name.c_str();
     126             : 
     127             :   // open dataset
     128           0 :   hid_t dataset_id = H5Dopen2(file_id, dataset_name_char, H5P_DEFAULT);
     129           0 :   if (dataset_id < 0)
     130           0 :     mooseError("Failed to obtain dataset from h5 file");
     131             : 
     132             :   // get dataspace
     133           0 :   hid_t dataspace_id = H5Dget_space(dataset_id);
     134           0 :   if (dataspace_id < 0)
     135           0 :     mooseError("Failed to obtain dataspace from h5 dataset");
     136             : 
     137             :   // get the dimensions of the dataspace
     138           0 :   const hsize_t rank = H5Sget_simple_extent_ndims(dataspace_id);
     139           0 :   std::vector<hsize_t> dims(rank);
     140           0 :   H5Sget_simple_extent_dims(dataspace_id, dims.data(), NULL);
     141             : 
     142             :   // get memory type id
     143           0 :   hid_t datatype_id = H5Dget_type(dataset_id);
     144             : 
     145             :   // total number of elements in the buffer
     146             :   int64_t total_number_of_elements = 1;
     147           0 :   for (auto i : index_range(dims))
     148             :   {
     149           0 :     total_number_of_elements *= dims[i];
     150             :   }
     151             : 
     152             :   // make tensor
     153           0 :   std::vector<int64_t> torch_dims(dims.begin(), dims.end());
     154             : 
     155           0 :   if (getParam<bool>("is_integer"))
     156             :   {
     157             :     // create read buffer
     158           0 :     std::vector<int64_t> buffer(total_number_of_elements);
     159             :     // read data
     160           0 :     H5Dread(dataset_id, datatype_id, H5S_ALL, dataspace_id, H5P_DEFAULT, buffer.data());
     161             : 
     162           0 :     auto cpu_tensor = torch::from_blob(buffer.data(), torch_dims, torch::kInt64).clone();
     163           0 :     _u = cpu_tensor.to(MooseTensor::intTensorOptions());
     164           0 :   }
     165             :   else
     166             :   {
     167             :     // create read buffer
     168           0 :     std::vector<double> buffer(total_number_of_elements);
     169             :     // read data
     170           0 :     H5Dread(dataset_id, datatype_id, H5S_ALL, dataspace_id, H5P_DEFAULT, buffer.data());
     171             : 
     172           0 :     auto cpu_tensor = torch::from_blob(buffer.data(), torch_dims, torch::kFloat64).clone();
     173           0 :     _u = cpu_tensor.to(MooseTensor::floatTensorOptions());
     174           0 :   }
     175           0 :   while (_u.dim() < 3)
     176           0 :     _u.unsqueeze_(-1);
     177             : 
     178             :   // close everything
     179           0 :   H5Fclose(file_id);
     180           0 :   H5Dclose(dataset_id);
     181           0 :   H5Sclose(dataspace_id);
     182             : #else
     183             :   mooseError("MOOSE was built without HDF5 support.");
     184             : #endif
     185           0 : }
     186             : 
     187             : void
     188           0 : LBMTensorBuffer::makeCPUCopy()
     189             : {
     190           0 :   if (!_u.defined())
     191             :     return;
     192             : 
     193           0 :   if (_cpu_copy_requested)
     194             :   {
     195           0 :     if (_u.is_cpu())
     196           0 :       _u_cpu = _u.clone().contiguous();
     197             :     else
     198           0 :       _u_cpu = _u.cpu().contiguous();
     199             :   }
     200             : }