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 "NEML2TensorCompute.h"
      12             : #include "NEML2Utils.h"
      13             : 
      14             : #ifdef NEML2_ENABLED
      15             : #include "neml2/models/map_types_fwd.h"
      16             : #include "neml2/tensors/Scalar.h"
      17             : #include "neml2/tensors/Vec.h"
      18             : #include "neml2/tensors/R2.h"
      19             : #endif
      20             : 
      21             : registerMooseObject("SwiftApp", NEML2TensorCompute);
      22             : 
      23             : InputParameters
      24           0 : NEML2TensorCompute::validParams()
      25             : {
      26           0 :   InputParameters params = TensorOperatorBase::validParams();
      27           0 :   params.addClassDescription(NEML2Utils::docstring("Compute object wrapper for a NEML2 model"));
      28             : 
      29           0 :   params.addRequiredParam<DataFileName>(
      30             :       "neml2_input_file",
      31           0 :       NEML2Utils::docstring("Path to the NEML2 input file containing the NEML2 model(s)."));
      32           0 :   params.addParam<std::vector<std::string>>(
      33             :       "cli_args",
      34             :       {},
      35             :       "Additional command line arguments to use when parsing the NEML2 input file.");
      36             : 
      37           0 :   params.addRequiredParam<std::string>("neml2_model", "Model name");
      38             : 
      39           0 :   params.addParam<std::vector<TensorInputBufferName>>(
      40             :       "swift_inputs", "Swift buffer names that go into the NEML2 model.");
      41           0 :   params.addParam<std::vector<std::string>>(
      42             :       "neml2_inputs", "NEML2 variable names corresponding to the `swift_inputs`.");
      43             : 
      44           0 :   params.addParam<std::vector<std::string>>("neml2_outputs", "NEML2 model outputs.");
      45           0 :   params.addParam<std::vector<TensorInputBufferName>>(
      46             :       "swift_outputs", "Swift buffer name corresponding to the `neml2_outputs`.");
      47             : 
      48           0 :   return params;
      49           0 : }
      50             : 
      51           0 : NEML2TensorCompute::NEML2TensorCompute(const InputParameters & params)
      52           0 :   : TensorOperatorBase(params)
      53             : #ifdef NEML2_ENABLED
      54             :     ,
      55             :     _model(
      56             :         [this]()
      57             :         {
      58             :           const auto fname = getParam<DataFileName>("neml2_input_file");
      59             :           const auto cli_args = getParam<std::vector<std::string>>("cli_args");
      60             :           const auto model_name = getParam<std::string>("neml2_model");
      61             :           auto factory = neml2::load_input(std::string(fname), neml2::utils::join(cli_args, " "));
      62             :           return NEML2Utils::getModel(*factory, model_name);
      63             :         }())
      64             : #endif
      65             : {
      66           0 :   NEML2Utils::assertNEML2Enabled();
      67             : 
      68             : #ifdef NEML2_ENABLED
      69             :   for (const auto & [swift_input_name, neml2_input_name] :
      70             :        getParam<TensorInputBufferName, std::string>("swift_inputs", "neml2_inputs"))
      71             :   {
      72             :     const auto * input_buffer = &getInputBufferByName<>(swift_input_name);
      73             :     _input_mapping.emplace_back(
      74             :         input_buffer, neml2::LabeledAxisAccessor(NEML2Utils::parseVariableName(neml2_input_name)));
      75             :   }
      76             : 
      77             :   for (const auto & [neml2_output_name, swift_output_name] :
      78             :        getParam<std::string, TensorInputBufferName>("neml2_outputs", "swift_outputs"))
      79             :   {
      80             :     auto * output_buffer = &getOutputBufferByName<>(swift_output_name);
      81             :     _output_mapping.emplace_back(
      82             :         neml2::LabeledAxisAccessor(NEML2Utils::parseVariableName(neml2_output_name)),
      83             :         output_buffer);
      84             :   }
      85             : #endif
      86           0 : }
      87             : 
      88             : void
      89           0 : NEML2TensorCompute::init()
      90             : {
      91             : #ifdef NEML2_ENABLED
      92             :   neml2::diagnose(*_model);
      93             : #endif
      94           0 : }
      95             : 
      96             : void
      97           0 : NEML2TensorCompute::computeBuffer()
      98             : {
      99             : #ifdef NEML2_ENABLED
     100             :   neml2::ValueMap in;
     101             :   for (const auto & [tensor_ptr, label] : _input_mapping)
     102             :   {
     103             :     // convert tensors on the fly at runtime
     104             :     auto sizes = tensor_ptr->sizes();
     105             :     if (sizes.size() == _dim)
     106             :       in[label] = neml2::Scalar(*tensor_ptr);
     107             :     else if (sizes.size() == _dim + 1)
     108             :       in[label] = neml2::Vec(*tensor_ptr, _domain.getShape());
     109             :     else if (sizes.size() == _dim + 3)
     110             :       in[label] = neml2::R2(*tensor_ptr, _domain.getShape());
     111             :     else
     112             :       mooseError("Unsupported tensor dimension");
     113             :   }
     114             : 
     115             :   auto out = _model->value(in);
     116             : 
     117             :   for (const auto & [label, tensor_ptr] : _output_mapping)
     118             :     *tensor_ptr = out.at(label);
     119             : #endif
     120           0 : }