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             : #pragma once
      10             : 
      11             : #include "SwiftUtils.h"
      12             : #include "MooseObject.h"
      13             : #include "InputParameters.h"
      14             : #include "DomainInterface.h"
      15             : 
      16             : /**
      17             :  * Tensor wrapper arbitrary tensor value dimensions
      18             :  */
      19             : class TensorBufferBase : public MooseObject, public DomainInterface
      20             : {
      21             : public:
      22             :   static InputParameters validParams();
      23             : 
      24             :   TensorBufferBase(const InputParameters & parameters);
      25             : 
      26             :   /// assignment operator
      27             :   TensorBufferBase & operator=(const torch::Tensor & rhs);
      28             : 
      29             :   /// advance state, returns the new number of old states
      30             :   virtual std::size_t advanceState() = 0;
      31             : 
      32             :   /// clear old states
      33             :   virtual void clearStates() = 0;
      34             : 
      35             :   /// create a contiguous CPU copy of the current tensor
      36             :   virtual void makeCPUCopy() = 0;
      37             : 
      38             :   /// initialize the tensor
      39           0 :   virtual void init() {}
      40             : 
      41             :   /// get a raw torch tensor representation
      42             :   virtual const torch::Tensor & getRawTensor() const = 0;
      43             : 
      44             :   /// get a raw torch tensor representation
      45             :   virtual const torch::Tensor & getRawCPUTensor() = 0;
      46             : 
      47             :   /// expand the tensor to full dimensions
      48             :   void expand();
      49             : 
      50             :   const bool _reciprocal;
      51             : 
      52             :   const torch::IntArrayRef _domain_shape;
      53             : 
      54             :   const torch::TensorOptions _options;
      55             : };