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 "LBMComputeVelocity.h"
      10             : #include "LatticeBoltzmannProblem.h"
      11             : #include "LatticeBoltzmannStencilBase.h"
      12             : 
      13             : using namespace torch::indexing;
      14             : 
      15             : registerMooseObject("SwiftApp", LBMComputeVelocity);
      16             : 
      17             : InputParameters
      18           0 : LBMComputeVelocity::validParams()
      19             : {
      20           0 :   InputParameters params = LatticeBoltzmannOperator::validParams();
      21           0 :   params.addRequiredParam<TensorInputBufferName>("f", "Distribution function");
      22           0 :   params.addRequiredParam<TensorInputBufferName>("rho", "Density");
      23           0 :   params.addParam<TensorInputBufferName>("forces", "forces", "Force tensor");
      24           0 :   params.addParam<bool>("enable_forces", false, "Whether to enable forces or no");
      25           0 :   params.addParam<bool>("add_body_force", false, "Whether to enable forces or no");
      26           0 :   params.addParam<SwiftConstantName>("body_force_x", "0.0", "Body force to be added in x-dir");
      27           0 :   params.addParam<SwiftConstantName>("body_force_y", "0.0", "Body force to be added in y-dir");
      28           0 :   params.addParam<SwiftConstantName>("body_force_z", "0.0", "Body force to be added in z-dir");
      29           0 :   params.addClassDescription("Compute object for macroscopic velocity reconstruction.");
      30           0 :   return params;
      31           0 : }
      32             : 
      33           0 : LBMComputeVelocity::LBMComputeVelocity(const InputParameters & parameters)
      34             :   : LatticeBoltzmannOperator(parameters),
      35           0 :     _f(getInputBuffer("f")),
      36           0 :     _rho(getInputBuffer("rho")),
      37           0 :     _force_tensor(getInputBuffer("forces")),
      38           0 :     _body_force_constant_x(
      39           0 :         _lb_problem.getConstant<Real>(getParam<SwiftConstantName>("body_force_x"))),
      40           0 :     _body_force_constant_y(
      41           0 :         _lb_problem.getConstant<Real>(getParam<SwiftConstantName>("body_force_y"))),
      42           0 :     _body_force_constant_z(
      43           0 :         _lb_problem.getConstant<Real>(getParam<SwiftConstantName>("body_force_z")))
      44             : {
      45           0 :   if (getParam<bool>("add_body_force"))
      46             :   {
      47           0 :     std::vector<int64_t> shape = {_shape[0], _shape[1], _shape[2], _domain.getDim()};
      48           0 :     _body_forces = torch::zeros(shape, MooseTensor::floatTensorOptions());
      49             : 
      50             :     auto force_constants =
      51           0 :         torch::tensor({_body_force_constant_x, _body_force_constant_y, _body_force_constant_z},
      52           0 :                       MooseTensor::floatTensorOptions());
      53             : 
      54           0 :     for (int64_t d = 0; d < _domain.getDim(); d++)
      55             :     {
      56           0 :       auto t_index = torch::tensor({d}, MooseTensor::intTensorOptions());
      57           0 :       _body_forces.index_fill_(-1, t_index, force_constants[d]);
      58             :     }
      59           0 :   }
      60           0 : }
      61             : 
      62             : void
      63           0 : LBMComputeVelocity::computeBuffer()
      64             : {
      65           0 :   const unsigned int & dim = _domain.getDim();
      66             : 
      67           0 :   _u.index({Slice(), Slice(), Slice(), 0}) = torch::sum(_f * _stencil._ex, 3) / _rho;
      68             : 
      69           0 :   if (dim > 1)
      70           0 :     _u.index({Slice(), Slice(), Slice(), 1}) = torch::sum(_f * _stencil._ey, 3) / _rho;
      71           0 :   if (dim > 2)
      72           0 :     _u.index({Slice(), Slice(), Slice(), 2}) = torch::sum(_f * _stencil._ez, 3) / _rho;
      73             : 
      74             :   // include forces
      75           0 :   if (getParam<bool>("enable_forces"))
      76           0 :     _u += _force_tensor / (2.0 * _rho.unsqueeze(3));
      77             : 
      78           0 :   if (getParam<bool>("add_body_force"))
      79           0 :     _u += _body_forces / (2.0 * _rho.unsqueeze(3));
      80             : 
      81           0 :   _lb_problem.maskedFillSolids(_u, 0);
      82           0 : }