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 "RandomTensor.h"
      10             : #include "SwiftUtils.h"
      11             : #include "TensorProblem.h"
      12             : #include <core/DeviceType.h>
      13             : 
      14             : registerMooseObject("SwiftApp", RandomTensor);
      15             : 
      16             : InputParameters
      17          60 : RandomTensor::validParams()
      18             : {
      19          60 :   InputParameters params = TensorOperator<>::validParams();
      20          60 :   params.addClassDescription("Uniform random IC with values between `min` and `max`.");
      21         120 :   params.addRequiredParam<Real>("min", "Minimum value.");
      22         120 :   params.addRequiredParam<Real>("max", "Maximum value.");
      23         120 :   params.addParam<int>("seed", "Random number seed.");
      24         120 :   params.addParam<bool>("generate_on_cpu",
      25         120 :                         true,
      26             :                         "To ensure reproducibility across devices it is recommended to generate "
      27             :                         "random tensors on the CPU.");
      28          60 :   return params;
      29           0 : }
      30             : 
      31          30 : RandomTensor::RandomTensor(const InputParameters & parameters)
      32          60 :   : TensorOperator<>(parameters), _generate_on_cpu(getParam<bool>("generate_on_cpu"))
      33             : {
      34          30 : }
      35             : 
      36             : void
      37          28 : RandomTensor::computeBuffer()
      38             : {
      39          56 :   const auto min = getParam<Real>("min");
      40          56 :   const auto max = getParam<Real>("max");
      41          56 :   if (isParamValid("seed"))
      42          84 :     torch::manual_seed(getParam<int>("seed"));
      43             : 
      44          28 :   if (_generate_on_cpu)
      45             :   {
      46          56 :     _u = torch::rand(_tensor_problem.getShape(),
      47          56 :                      MooseTensor::floatTensorOptions().device(torch::kCPU))
      48          84 :                  .to(MooseTensor::floatTensorOptions()) *
      49          56 :              (max - min) +
      50             :          min;
      51             :   }
      52             :   else
      53           0 :     _u = torch::rand(_tensor_problem.getShape(), MooseTensor::floatTensorOptions()) * (max - min) +
      54             :          min;
      55          28 : }