math.h

// Copyright 2023, UChicago Argonne, LLC
// All Rights Reserved
// Software Name: NEML2 -- the New Engineering material Model Library, version 2
// By: Argonne National Laboratory
// OPEN SOURCE LICENSE (MIT)
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.
 
#pragma once
 
#include "neml2/tensors/BatchTensor.h"
 
namespace neml2
{
class SR2;
class WR2;
class SWR4;
class SSR4;
class WSR4;
namespace math
{
constexpr Real eps = std::numeric_limits<at::scalar_value_type<Real>::type>::epsilon();
 
constexpr Real sqrt2 = 1.4142135623730951;
constexpr Real invsqrt2 = 0.7071067811865475;
 
constexpr TorchSize mandel_reverse_index[3][3] = {{0, 5, 4}, {5, 1, 3}, {4, 3, 2}};
constexpr TorchSize mandel_index[6][2] = {{0, 0}, {1, 1}, {2, 2}, {1, 2}, {0, 2}, {0, 1}};
 
constexpr TorchSize skew_reverse_index[3][3] = {{0, 2, 1}, {2, 0, 0}, {1, 0, 0}};
constexpr Real skew_factor[3][3] = {{0.0, -1.0, 1.0}, {1.0, 0.0, -1.0}, {-1.0, 1.0, 0.0}};
 
inline constexpr Real
mandel_factor(TorchSize i)
{
  return i < 3 ? 1.0 : sqrt2;
}
 
struct ConstantTensors
{
  ConstantTensors();
 
  // Get the global constants
  static ConstantTensors & get();
 
  static const torch::Tensor & full_to_mandel_map();
  static const torch::Tensor & mandel_to_full_map();
  static const torch::Tensor & full_to_mandel_factor();
  static const torch::Tensor & mandel_to_full_factor();
  static const torch::Tensor & full_to_skew_map();
  static const torch::Tensor & skew_to_full_map();
  static const torch::Tensor & full_to_skew_factor();
  static const torch::Tensor & skew_to_full_factor();
 
private:
  torch::Tensor _full_to_mandel_map;
  torch::Tensor _mandel_to_full_map;
  torch::Tensor _full_to_mandel_factor;
  torch::Tensor _mandel_to_full_factor;
  torch::Tensor _full_to_skew_map;
  torch::Tensor _skew_to_full_map;
  torch::Tensor _full_to_skew_factor;
  torch::Tensor _skew_to_full_factor;
};
 
BatchTensor full_to_reduced(const BatchTensor & full,
                            const torch::Tensor & rmap,
                            const torch::Tensor & rfactors,
                            TorchSize dim = 0);
 
BatchTensor reduced_to_full(const BatchTensor & reduced,
                            const torch::Tensor & rmap,
                            const torch::Tensor & rfactors,
                            TorchSize dim = 0);
 
BatchTensor full_to_mandel(const BatchTensor & full, TorchSize dim = 0);
 
BatchTensor mandel_to_full(const BatchTensor & mandel, TorchSize dim = 0);
 
BatchTensor full_to_skew(const BatchTensor & full, TorchSize dim = 0);
 
BatchTensor skew_to_full(const BatchTensor & skew, TorchSize dim = 0);
 
BatchTensor jacrev(const BatchTensor & y, const BatchTensor & p);
 
BatchTensor
base_diag_embed(const BatchTensor & a, TorchSize offset = 0, TorchSize d1 = -2, TorchSize d2 = -1);
 
SR2 skew_and_sym_to_sym(const SR2 & e, const WR2 & w);
 
SSR4 d_skew_and_sym_to_sym_d_sym(const WR2 & w);
 
SWR4 d_skew_and_sym_to_sym_d_skew(const SR2 & e);
 
WR2 multiply_and_make_skew(const SR2 & a, const SR2 & b);
 
WSR4 d_multiply_and_make_skew_d_first(const SR2 & b);
 
WSR4 d_multiply_and_make_skew_d_second(const SR2 & a);
 
namespace linalg
{
BatchTensor vector_norm(const BatchTensor & v);
 
BatchTensor inv(const BatchTensor & m);
 
BatchTensor solve(const BatchTensor & A, const BatchTensor & B);
 
std::tuple<BatchTensor, BatchTensor> lu_factor(const BatchTensor & A, bool pivot = true);
 
BatchTensor lu_solve(const BatchTensor & LU,
                     const BatchTensor & pivots,
                     const BatchTensor & B,
                     bool left = true,
                     bool adjoint = false);
} // namespace linalg
} // namespace math
} // namespace neml2