math.h

// Copyright (C) 2021 Igor Baratta
//
// This file is part of DOLFINx (https://www.fenicsproject.org)
//
// SPDX-License-Identifier:    LGPL-3.0-or-later
 
#pragma once
 
#include <array>
#include <cmath>
#include <type_traits>
#include <xtensor/xfixed.hpp>
#include <xtensor/xtensor.hpp>
 
namespace dolfinx::math
{
 
template <typename U, typename V>
xt::xtensor_fixed<typename U::value_type, xt::xshape<3>> cross(const U& u,
                                                               const V& v)
{
  assert(u.size() == 3);
  assert(v.size() == 3);
  return {u[1] * v[2] - u[2] * v[1], u[2] * v[0] - u[0] * v[2],
          u[0] * v[1] - u[1] * v[0]};
}
 
template <typename T>
T difference_of_products(T a, T b, T c, T d) noexcept
{
  T w = b * c;
  T err = std::fma(-b, c, w);
  T diff = std::fma(a, d, -w);
  return diff + err;
}
 
template <typename Matrix>
auto det(const Matrix& A)
{
  using value_type = typename Matrix::value_type;
  assert(A.shape(0) == A.shape(1));
  assert(A.dimension() == 2);
 
  const int nrows = A.shape(0);
  switch (nrows)
  {
  case 1:
    return A(0, 0);
  case 2:
    return difference_of_products(A(0, 0), A(0, 1), A(1, 0), A(1, 1));
  case 3:
  {
    // Leibniz formula combined with Kahan’s method for accurate
    // computation of 3 x 3 determinants
    value_type w0 = difference_of_products(A(1, 1), A(1, 2), A(2, 1), A(2, 2));
    value_type w1 = difference_of_products(A(1, 0), A(1, 2), A(2, 0), A(2, 2));
    value_type w2 = difference_of_products(A(1, 0), A(1, 1), A(2, 0), A(2, 1));
    value_type w3 = difference_of_products(A(0, 0), A(0, 1), w1, w0);
    value_type w4 = std::fma(A(0, 2), w2, w3);
    return w4;
  }
  default:
    throw std::runtime_error("math::det is not implemented for "
                             + std::to_string(A.shape(0)) + "x"
                             + std::to_string(A.shape(1)) + " matrices.");
  }
}
 
template <typename U, typename V>
void inv(const U& A, V&& B)
{
  using value_type = typename U::value_type;
  const std::size_t nrows = A.shape(0);
  switch (nrows)
  {
  case 1:
    B(0, 0) = 1 / A(0, 0);
    break;
  case 2:
  {
    value_type idet = 1. / det(A);
    B(0, 0) = idet * A(1, 1);
    B(0, 1) = -idet * A(0, 1);
    B(1, 0) = -idet * A(1, 0);
    B(1, 1) = idet * A(0, 0);
    break;
  }
  case 3:
  {
    value_type w0 = difference_of_products(A(1, 1), A(1, 2), A(2, 1), A(2, 2));
    value_type w1 = difference_of_products(A(1, 0), A(1, 2), A(2, 0), A(2, 2));
    value_type w2 = difference_of_products(A(1, 0), A(1, 1), A(2, 0), A(2, 1));
    value_type w3 = difference_of_products(A(0, 0), A(0, 1), w1, w0);
    value_type det = std::fma(A(0, 2), w2, w3);
    assert(det != 0.);
    value_type idet = 1 / det;
 
    B(0, 0) = w0 * idet;
    B(1, 0) = -w1 * idet;
    B(2, 0) = w2 * idet;
    B(0, 1) = difference_of_products(A(0, 2), A(0, 1), A(2, 2), A(2, 1)) * idet;
    B(0, 2) = difference_of_products(A(0, 1), A(0, 2), A(1, 1), A(1, 2)) * idet;
    B(1, 1) = difference_of_products(A(0, 0), A(0, 2), A(2, 0), A(2, 2)) * idet;
    B(1, 2) = difference_of_products(A(1, 0), A(0, 0), A(1, 2), A(0, 2)) * idet;
    B(2, 1) = difference_of_products(A(2, 0), A(0, 0), A(2, 1), A(0, 1)) * idet;
    B(2, 2) = difference_of_products(A(0, 0), A(1, 0), A(0, 1), A(1, 1)) * idet;
    break;
  }
  default:
    throw std::runtime_error("math::inv is not implemented for "
                             + std::to_string(A.shape(0)) + "x"
                             + std::to_string(A.shape(1)) + " matrices.");
  }
}
 
template <typename U, typename V, typename P>
void dot(const U& A, const V& B, P&& C, bool transpose = false)
{
  if (transpose)
  {
    assert(A.shape(0) == B.shape(1));
    for (std::size_t i = 0; i < A.shape(1); i++)
      for (std::size_t j = 0; j < B.shape(0); j++)
        for (std::size_t k = 0; k < A.shape(0); k++)
          C(i, j) += A(k, i) * B(j, k);
  }
  else
  {
    assert(A.shape(1) == B.shape(0));
    for (std::size_t i = 0; i < A.shape(0); i++)
      for (std::size_t j = 0; j < B.shape(1); j++)
        for (std::size_t k = 0; k < A.shape(1); k++)
          C(i, j) += A(i, k) * B(k, j);
  }
}
 
template <typename U, typename V>
void pinv(const U& A, V&& P)
{
  auto shape = A.shape();
  assert(shape[0] > shape[1]);
  assert(P.shape(1) == shape[0]);
  assert(P.shape(0) == shape[1]);
  using T = typename U::value_type;
  if (shape[1] == 2)
  {
    xt::xtensor_fixed<T, xt::xshape<2, 3>> AT;
    xt::xtensor_fixed<T, xt::xshape<2, 2>> ATA;
    xt::xtensor_fixed<T, xt::xshape<2, 2>> Inv;
    AT = xt::transpose(A);
    std::fill(ATA.begin(), ATA.end(), 0.0);
    std::fill(P.begin(), P.end(), 0.0);
 
    // pinv(A) = (A^T * A)^-1 * A^T
    dolfinx::math::dot(AT, A, ATA);
    dolfinx::math::inv(ATA, Inv);
    dolfinx::math::dot(Inv, AT, P);
  }
  else if (shape[1] == 1)
  {
    auto res = std::transform_reduce(A.begin(), A.end(), 0., std::plus<T>(),
                                     [](const auto v) { return v * v; });
    P = (1 / res) * xt::transpose(A);
  }
  else
  {
    throw std::runtime_error("math::pinv is not implemented for "
                             + std::to_string(A.shape(0)) + "x"
                             + std::to_string(A.shape(1)) + " matrices.");
  }
}
 
} // namespace dolfinx::math