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 "LBMApplyForces.h"
      10             : #include "LatticeBoltzmannProblem.h"
      11             : #include "LatticeBoltzmannStencilBase.h"
      12             : 
      13             : using namespace torch::indexing;
      14             : 
      15             : registerMooseObject("SwiftApp", LBMApplyForces);
      16             : 
      17             : InputParameters
      18           0 : LBMApplyForces::validParams()
      19             : {
      20           0 :   InputParameters params = LatticeBoltzmannOperator::validParams();
      21             :   // params.addRequiredParam<TensorInputBufferName>("f", "Distribution function");
      22           0 :   params.addParam<TensorInputBufferName>("velocity", "u", "Macroscopic velocity");
      23           0 :   params.addRequiredParam<TensorInputBufferName>("rho", "Macroscopic density");
      24           0 :   params.addRequiredParam<TensorInputBufferName>("forces", "LBM forces");
      25           0 :   params.addRequiredParam<std::string>("tau0", "Relaxation parameter");
      26           0 :   params.addClassDescription("Compute object for LB forces");
      27           0 :   return params;
      28           0 : }
      29             : 
      30           0 : LBMApplyForces::LBMApplyForces(const InputParameters & parameters)
      31             :   : LatticeBoltzmannOperator(parameters), /*_f(getInputBuffer("f"))*/
      32           0 :     _velocity(getInputBuffer("velocity")),
      33           0 :     _density(getInputBuffer("rho")),
      34           0 :     _forces(getInputBuffer("forces")),
      35           0 :     _tau(_lb_problem.getConstant<Real>(getParam<std::string>("tau0")))
      36             : {
      37           0 : }
      38             : 
      39             : void
      40           0 : LBMApplyForces::computeSourceTerm()
      41             : {
      42           0 :   const unsigned int & dim = _domain.getDim();
      43             : 
      44           0 :   if (_density.dim() < 3)
      45           0 :     _density.unsqueeze_(2);
      46             : 
      47           0 :   torch::Tensor rho_unsqueezed = _density.unsqueeze(-1);
      48           0 :   torch::Tensor Fx = _forces.select(3, 0).unsqueeze(3);
      49           0 :   torch::Tensor Fy = _forces.select(3, 1).unsqueeze(3);
      50             :   torch::Tensor Fz;
      51             : 
      52             :   // torch::Tensor ux = _velocity.select(3, 0).unsqueeze(3);
      53             :   // torch::Tensor uy = _velocity.select(3, 1).unsqueeze(3);
      54             :   // torch::Tensor uz;
      55             : 
      56           0 :   torch::Tensor e_xyz = torch::stack({_stencil._ex, _stencil._ey, _stencil._ez}, 0);
      57             : 
      58           0 :   switch (dim)
      59             :   {
      60           0 :     case 3:
      61             :     {
      62           0 :       Fz = _forces.select(3, 2).unsqueeze(3);
      63             :       // uz = _velocity.select(3, 2).unsqueeze(3);
      64           0 :       break;
      65             :     }
      66             :     case 2:
      67             :     {
      68           0 :       Fz = torch::zeros_like(rho_unsqueezed, MooseTensor::floatTensorOptions());
      69             :       // uz = torch::zeros_like(rho_unsqueezed, MooseTensor::floatTensorOptions());
      70           0 :       break;
      71             :     }
      72           0 :     default:
      73           0 :       mooseError("Unsupported dimensions for buffer _u");
      74             :   }
      75             : 
      76             :   // torch::Tensor Fxyz = torch::stack({Fx, Fy, Fz}, 3).squeeze(-1);
      77             :   // torch::Tensor Uxyz = torch::stack({ux, uy, uz}, 3).squeeze(-1);
      78             :   // torch::Tensor Fxyz_expanded = Fxyz.unsqueeze(-1);       // Shape: (Nx, Ny, Nz, 3, 1)
      79             :   // torch::Tensor Uxyz_expanded = Uxyz.unsqueeze(-2);       // Shape: (Nx, Ny, Nz, 1, 3)
      80             :   // torch::Tensor UF_outer = Fxyz_expanded * Uxyz_expanded; // Shape: (Nx, Ny, Nz, 3, 3)
      81             :   // torch::Tensor UF_outer_flat = UF_outer.flatten(-2, -1); // Shape: (Nx, Ny, Nz, 9)
      82             : 
      83           0 :   for (int64_t ic = 0; ic < _stencil._q; ic++)
      84             :   {
      85             :     // auto exyz_ic = e_xyz.index({Slice(), ic}).flatten(); // Shape (3)
      86             :     // torch::Tensor ccr = torch::outer(exyz_ic, exyz_ic) / _lb_problem._cs2 -
      87             :     //                     torch::eye(3, MooseTensor::floatTensorOptions()); // Shape (9)
      88             :     // auto ccr_flat = ccr.flatten();
      89             :     // torch::Tensor multiplied = UF_outer_flat * ccr_flat; // Shape: (Nx, Ny, Nz, 9)
      90             : 
      91             :     // // sum along the last dimension
      92             :     // torch::Tensor UFccr = multiplied.sum(/*dim=*/-1);
      93             : 
      94             :     // compute source
      95           0 :     _source_term.index_put_(
      96           0 :         {Slice(), Slice(), Slice(), ic},
      97           0 :         _stencil._weights[ic] * rho_unsqueezed.squeeze(-1) *
      98           0 :             ((_stencil._ex[ic] * Fx + _stencil._ey[ic] * Fy + _stencil._ez[ic] * Fz).squeeze(-1) /
      99           0 :              _lb_problem._cs2 /* + UFccr / 2.0 / _lb_problem._cs4 */));
     100             :   }
     101           0 : }
     102             : 
     103             : void
     104           0 : LBMApplyForces::computeBuffer()
     105             : {
     106           0 :   computeSourceTerm();
     107           0 :   _u += (1.0 - 1.0 / (2.0 * _tau)) * _source_term;
     108           0 :   _lb_problem.maskedFillSolids(_u, 0);
     109           0 : }