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 "LBMIsotropicGradient.h"
      10             : 
      11             : using namespace torch::indexing;
      12             : 
      13             : registerMooseObject("SwiftApp", LBMIsotropicGradient);
      14             : 
      15             : InputParameters
      16           0 : LBMIsotropicGradient::validParams()
      17             : {
      18           0 :   InputParameters params = LatticeBoltzmannOperator::validParams();
      19           0 :   params.addClassDescription("Compute isotropic gradient object.");
      20           0 :   params.addRequiredParam<TensorInputBufferName>("scalar_field",
      21             :                                                  "Scalar field to compute the gradient of");
      22             : 
      23           0 :   return params;
      24           0 : }
      25             : 
      26           0 : LBMIsotropicGradient::LBMIsotropicGradient(const InputParameters & parameters)
      27           0 :   : LatticeBoltzmannOperator(parameters), _scalar_field(getInputBuffer("scalar_field"))
      28             : {
      29           0 :   const unsigned int & dim = _domain.getDim();
      30             : 
      31             :   // Note: if D3Q19 stencil is used, isotropic gradient is NOT going to work,
      32             :   // because D3Q19 is NOT isotropic.
      33             : 
      34           0 :   if (_stencil._q == 19)
      35           0 :     mooseError("Isotropic gradient cannot be computed for D3Q19 stencil");
      36             : 
      37           0 :   _kernel = torch::zeros({3, 3, _domain.getDim()}, MooseTensor::floatTensorOptions());
      38             : 
      39           0 :   switch (dim)
      40             :   {
      41           0 :     case 3:
      42           0 :       mooseError("LBMIsotropicGradient is not implemented for 3D");
      43             :       break;
      44           0 :     case 2:
      45             :     {
      46             :       auto kernel_of_kernel =
      47           0 :           torch::index_select(_stencil._weights, 0, _stencil._reorder_indices).reshape({3, 3});
      48           0 :       auto ex3x3 = torch::index_select(_stencil._ex, 0, _stencil._reorder_indices).reshape({3, 3});
      49           0 :       auto ey3x3 = torch::index_select(_stencil._ey, 0, _stencil._reorder_indices).reshape({3, 3});
      50             : 
      51           0 :       _kernel.index_put_({Slice(), Slice(), 0}, kernel_of_kernel * ex3x3);
      52           0 :       _kernel.index_put_({Slice(), Slice(), 1}, kernel_of_kernel * ey3x3);
      53             : 
      54           0 :       _conv_options.bias(torch::Tensor()).stride({1, 1}).padding(0);
      55             :       break;
      56             :     }
      57             :   }
      58           0 : }
      59             : 
      60             : torch::Tensor
      61           0 : LBMIsotropicGradient::padScalarField()
      62             : {
      63             :   // because torch sucks at padding
      64             :   torch::Tensor right_pad_slice =
      65           0 :       _scalar_field.slice(1, _scalar_field.size(1) - _padding, _scalar_field.size(1));
      66           0 :   torch::Tensor left_pad_slice = _scalar_field.slice(1, 0, _padding);
      67             : 
      68           0 :   torch::Tensor padded_width = torch::cat({left_pad_slice, _scalar_field, right_pad_slice}, 1);
      69             : 
      70             :   torch::Tensor bottom_pad_slice =
      71           0 :       padded_width.slice(0, padded_width.size(0) - _padding, padded_width.size(0));
      72           0 :   torch::Tensor top_pad_slice = padded_width.slice(0, 0, _padding);
      73             : 
      74             :   torch::Tensor fully_padded_tensor =
      75           0 :       torch::cat({top_pad_slice, padded_width, bottom_pad_slice}, 0);
      76             : 
      77           0 :   return fully_padded_tensor;
      78           0 : }
      79             : 
      80             : void
      81           0 : LBMIsotropicGradient::computeBuffer()
      82             : {
      83             :   // check output buffer shape
      84           0 :   if ((unsigned int)_u.size(-1) != _domain.getDim())
      85           0 :     mooseError("Output buffer must have the same number of dimensions as the domain.");
      86             : 
      87             :   const unsigned int & dim = _domain.getDim();
      88           0 :   torch::Tensor kernel = _kernel.permute({2, 0, 1});
      89           0 :   kernel = kernel.unsqueeze(1);
      90             : 
      91           0 :   switch (dim)
      92             :   {
      93           0 :     case 3:
      94           0 :       mooseError("LBMIsotropicGradient is not implemented for 3D");
      95             :       break;
      96           0 :     case 2:
      97             :     {
      98           0 :       if (_scalar_field.dim() > 2)
      99           0 :         _scalar_field.squeeze_(-1);
     100             : 
     101           0 :       torch::Tensor input_field = padScalarField();
     102             : 
     103           0 :       input_field = input_field.unsqueeze(0).unsqueeze(0);
     104             : 
     105             :       torch::Tensor isotropic_gradient =
     106           0 :           torch::nn::functional::conv2d(input_field, kernel, _conv_options);
     107             : 
     108           0 :       isotropic_gradient = isotropic_gradient.squeeze(0);
     109             : 
     110           0 :       _u.index_put_({Slice(), Slice(), Slice(), 0},
     111           0 :                     isotropic_gradient.index({0, Slice(), Slice()}).unsqueeze(-1));
     112           0 :       _u.index_put_({Slice(), Slice(), Slice(), 1},
     113           0 :                     isotropic_gradient.index({1, Slice(), Slice()}).unsqueeze(-1));
     114           0 :       _u = _u / _lb_problem._cs2;
     115             :       break;
     116             :     }
     117             :   }
     118           0 :   _lb_problem.maskedFillSolids(_u, 0);
     119           0 : }