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 "MooseObject.h"
      12             : #include <thread>
      13             : #include <torch/torch.h>
      14             : 
      15             : class TensorBufferBase;
      16             : class TensorProblem;
      17             : class DomainAction;
      18             : 
      19             : /**
      20             :  * Direct buffer output
      21             :  */
      22             : class TensorOutput : public MooseObject
      23             : {
      24             : public:
      25             :   static InputParameters validParams();
      26             : 
      27             :   TensorOutput(const InputParameters & parameters);
      28             : 
      29           0 :   virtual void init() {}
      30             : 
      31             :   virtual bool shouldRun(const ExecFlagType & execute_flag) const;
      32             : 
      33             :   void startOutput();
      34             :   void waitForCompletion();
      35             : 
      36             : protected:
      37             :   virtual void output() = 0;
      38             : 
      39             :   TensorProblem & _tensor_problem;
      40             :   const DomainAction & _domain;
      41             : 
      42             :   /// simulation time of the step that is being output
      43             :   const Real & _time;
      44             : 
      45             :   /// output file name prefix
      46             :   const std::string _file_base;
      47             : 
      48             :   std::thread _output_thread;
      49             : 
      50             :   /// The buffer this output object is outputting
      51             :   std::map<std::string, const torch::Tensor *> _out_buffers;
      52             : 
      53             :   const ExecFlagEnum _execute_on;
      54             : };