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 "ProjectTensorAux.h"
      10             : #include "DomainAction.h"
      11             : #include "TensorProblem.h"
      12             : #include "SwiftTypes.h"
      13             : 
      14             : registerMooseObject("SwiftApp", ProjectTensorAux);
      15             : 
      16             : InputParameters
      17          20 : ProjectTensorAux::validParams()
      18             : {
      19          20 :   InputParameters params = AuxKernel::validParams();
      20          20 :   params.addClassDescription("Project a Tensor buffer onto an auxiliary variable");
      21          40 :   params.addRequiredParam<TensorInputBufferName>("buffer", "The buffer to read from");
      22          20 :   return params;
      23           0 : }
      24             : 
      25          12 : ProjectTensorAux::ProjectTensorAux(const InputParameters & parameters)
      26             :   : AuxKernel(parameters),
      27             :     TensorProblemInterface(this),
      28             :     DomainInterface(this),
      29          12 :     _cpu_buffer(_tensor_problem.getRawCPUBuffer(getParam<TensorInputBufferName>("buffer"))),
      30          12 :     _dim(_domain.getDim()),
      31             :     _n(_domain.getGridSize()),
      32          12 :     _grid_spacing(_domain.getGridSpacing())
      33             : {
      34          12 : }
      35             : 
      36             : Real
      37       37004 : ProjectTensorAux::computeValue()
      38             : {
      39       37004 :   auto getElement = [this]()
      40             :   {
      41       37004 :     const Point shift(_grid_spacing(0) / 2.0, _grid_spacing(1) / 2.0, _grid_spacing(2) / 2.0);
      42       37004 :     Point p = isNodal() ? (*_current_node + shift) : _current_elem->vertex_average();
      43             : 
      44             :     using at::indexing::TensorIndex;
      45       37004 :     switch (_dim)
      46             :     {
      47           0 :       case 1:
      48           0 :         return _cpu_buffer.index({TensorIndex(int64_t(p(0) / _grid_spacing(0)) % _n[0])});
      49             : 
      50       37004 :       case 2:
      51       74008 :         return _cpu_buffer.index({TensorIndex(int64_t(p(0) / _grid_spacing(0)) % _n[0]),
      52      148016 :                                   TensorIndex(int64_t(p(1) / _grid_spacing(1)) % _n[1])});
      53             : 
      54           0 :       case 3:
      55           0 :         return _cpu_buffer.index({TensorIndex(int64_t(p(0) / _grid_spacing(0)) % _n[0]),
      56           0 :                                   TensorIndex(int64_t(p(1) / _grid_spacing(1)) % _n[1]),
      57           0 :                                   TensorIndex(int64_t(p(2) / _grid_spacing(2)) % _n[2])});
      58             :     }
      59             : 
      60           0 :     mooseError("Internal error (invalid dimension)");
      61       37004 :   };
      62             : 
      63       37004 :   const auto element = getElement();
      64             : 
      65       37004 :   if (_cpu_buffer.dtype() == torch::kFloat32)
      66           0 :     return element.item<float>();
      67       37004 :   else if (_cpu_buffer.dtype() == torch::kFloat64)
      68       37004 :     return element.item<double>();
      69             :   else
      70           0 :     mooseError("Unsupported output type");
      71             : }