Line data    Source code

       1              : // Copyright 2024, UChicago Argonne, LLC
       2              : // All Rights Reserved
       3              : // Software Name: NEML2 -- the New Engineering material Model Library, version 2
       4              : // By: Argonne National Laboratory
       5              : // OPEN SOURCE LICENSE (MIT)
       6              : //
       7              : // Permission is hereby granted, free of charge, to any person obtaining a copy
       8              : // of this software and associated documentation files (the "Software"), to deal
       9              : // in the Software without restriction, including without limitation the rights
      10              : // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
      11              : // copies of the Software, and to permit persons to whom the Software is
      12              : // furnished to do so, subject to the following conditions:
      13              : //
      14              : // The above copyright notice and this permission notice shall be included in
      15              : // all copies or substantial portions of the Software.
      16              : //
      17              : // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
      18              : // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
      19              : // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
      20              : // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
      21              : // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
      22              : // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
      23              : // THE SOFTWARE.
      24              : 
      25              : #include "neml2/models/LinearInterpolation.h"
      26              : #include "neml2/tensors/functions/diff.h"
      27              : #include "neml2/tensors/Scalar.h"
      28              : #include "neml2/tensors/Vec.h"
      29              : #include "neml2/tensors/SR2.h"
      30              : #include "neml2/tensors/indexing.h"
      31              : 
      32              : namespace neml2
      33              : {
      34              : template <typename T>
      35              : OptionSet
      36            6 : LinearInterpolation<T>::expected_options()
      37              : {
      38            6 :   OptionSet options = Interpolation<T>::expected_options();
      39            6 :   options.doc() += " This object performs a _linear interpolation_.";
      40            6 :   options.set<bool>("define_second_derivatives") = true;
      41            6 :   return options;
      42            0 : }
      43              : 
      44              : template <typename T>
      45           10 : LinearInterpolation<T>::LinearInterpolation(const OptionSet & options)
      46           10 :   : Interpolation<T>(options)
      47              : {
      48           10 : }
      49              : 
      50              : template <typename T>
      51              : void
      52           29 : LinearInterpolation<T>::set_value(bool out, bool dout_din, bool d2out_din2)
      53              : {
      54           58 :   const auto slope =
      55           58 :       diff(this->_Y, 1, this->_Y.batch_dim() - 1) / diff(this->_X, 1, this->_X.batch_dim() - 1);
      56          116 :   const auto X0 = this->_X.batch_index({indexing::Ellipsis, indexing::Slice(indexing::None, -1)})
      57              :                       .batch_expand_as(slope);
      58          116 :   const auto X1 = this->_X.batch_index({indexing::Ellipsis, indexing::Slice(1, indexing::None)})
      59              :                       .batch_expand_as(slope);
      60          116 :   const auto Y0 = this->_Y.batch_index({indexing::Ellipsis, indexing::Slice(indexing::None, -1)})
      61              :                       .batch_expand_as(slope);
      62              : 
      63           29 :   const auto x = Scalar(this->_x);
      64           29 :   const auto loc =
      65              :       Scalar(at::logical_and(at::gt(x.batch_unsqueeze(-1), X0), at::le(x.batch_unsqueeze(-1), X1)));
      66           29 :   const auto si = mask<T>(slope, loc);
      67              : 
      68           29 :   if (out)
      69              :   {
      70           13 :     const auto X0i = mask<Scalar>(X0, loc);
      71           13 :     const auto Y0i = mask<T>(Y0, loc);
      72           13 :     this->_p = Y0i + si * (x - X0i);
      73           13 :   }
      74              : 
      75           29 :   if (dout_din)
      76           11 :     if (this->_x.is_dependent())
      77           11 :       this->_p.d(this->_x) = si;
      78              : 
      79              :   if (d2out_din2)
      80              :   {
      81              :     // zero
      82              :   }
      83          203 : }
      84              : 
      85              : #define REGISTER(T)                                                                                \
      86              :   using T##LinearInterpolation = LinearInterpolation<T>;                                           \
      87              :   register_NEML2_object(T##LinearInterpolation);                                                   \
      88              :   template class LinearInterpolation<T>
      89              : REGISTER(Scalar);
      90              : REGISTER(Vec);
      91              : REGISTER(SR2);
      92              : } // namespace neml2