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 "types.h"
#include <algorithm>
#include <array>
#include <basix/mdspan.hpp>
#include <cmath>
#include <string>
#include <type_traits>
 
namespace dolfinx::math
{
 
template <typename U, typename V>
std::array<typename U::value_type, 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 T>
auto det(const T* A, std::array<std::size_t, 2> shape)
{
  assert(shape[0] == shape[1]);
 
  // const int nrows = shape[0];
  switch (shape[0])
  {
  case 1:
    return *A;
  case 2:
    /* A(0, 0), A(0, 1), A(1, 0), A(1, 1) */
    return difference_of_products(A[0], A[1], A[2], A[3]);
  case 3:
  {
    // Leibniz formula combined with Kahan’s method for accurate
    // computation of 3 x 3 determinants
    T w0 = difference_of_products(A[3 + 1], A[3 + 2], A[3 * 2 + 1],
                                  A[2 * 3 + 2]);
    T w1 = difference_of_products(A[3], A[3 + 2], A[3 * 2], A[3 * 2 + 2]);
    T w2 = difference_of_products(A[3], A[3 + 1], A[3 * 2], A[3 * 2 + 1]);
    T w3 = difference_of_products(A[0], A[1], w1, w0);
    T w4 = std::fma(A[2], w2, w3);
    return w4;
  }
  default:
    throw std::runtime_error("math::det is not implemented for "
                             + std::to_string(A[0]) + "x" + std::to_string(A[1])
                             + " matrices.");
  }
}
 
template <typename Matrix>
auto det(Matrix A)
{
  static_assert(Matrix::rank() == 2, "Must be rank 2");
  assert(A.extent(0) == A.extent(1));
 
  using value_type = typename Matrix::value_type;
  const int nrows = A.extent(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.extent(0)) + "x"
                             + std::to_string(A.extent(1)) + " matrices.");
  }
}
 
template <typename U, typename V>
void inv(U A, V B)
{
  static_assert(U::rank() == 2, "Must be rank 2");
  static_assert(V::rank() == 2, "Must be rank 2");
 
  using value_type = typename U::value_type;
  const std::size_t nrows = A.extent(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.extent(0)) + "x"
                             + std::to_string(A.extent(1)) + " matrices.");
  }
}
 
template <typename U, typename V, typename P>
void dot(U A, V B, P C, bool transpose = false)
{
  static_assert(U::rank() == 2, "Must be rank 2");
  static_assert(V::rank() == 2, "Must be rank 2");
  static_assert(P::rank() == 2, "Must be rank 2");
 
  if (transpose)
  {
    assert(A.extent(0) == B.extent(1));
    for (std::size_t i = 0; i < A.extent(1); i++)
      for (std::size_t j = 0; j < B.extent(0); j++)
        for (std::size_t k = 0; k < A.extent(0); k++)
          C(i, j) += A(k, i) * B(j, k);
  }
  else
  {
    assert(A.extent(1) == B.extent(0));
    for (std::size_t i = 0; i < A.extent(0); i++)
      for (std::size_t j = 0; j < B.extent(1); j++)
        for (std::size_t k = 0; k < A.extent(1); k++)
          C(i, j) += A(i, k) * B(k, j);
  }
}
 
template <typename U, typename V>
void pinv(U A, V P)
{
  static_assert(U::rank() == 2, "Must be rank 2");
  static_assert(V::rank() == 2, "Must be rank 2");
 
  assert(A.extent(0) > A.extent(1));
  assert(P.extent(1) == A.extent(0));
  assert(P.extent(0) == A.extent(1));
  using T = typename U::value_type;
  if (A.extent(1) == 2)
  {
    std::array<T, 6> ATb;
    std::array<T, 4> ATAb, Invb;
    MDSPAN_IMPL_STANDARD_NAMESPACE::mdspan<
        T, MDSPAN_IMPL_STANDARD_NAMESPACE::extents<std::size_t, 2, 3>>
        AT(ATb.data(), 2, 3);
    MDSPAN_IMPL_STANDARD_NAMESPACE::mdspan<
        T, MDSPAN_IMPL_STANDARD_NAMESPACE::extents<std::size_t, 2, 2>>
        ATA(ATAb.data(), 2, 2);
    MDSPAN_IMPL_STANDARD_NAMESPACE::mdspan<
        T, MDSPAN_IMPL_STANDARD_NAMESPACE::extents<std::size_t, 2, 2>>
        Inv(Invb.data(), 2, 2);
 
    for (std::size_t i = 0; i < AT.extent(0); ++i)
      for (std::size_t j = 0; j < AT.extent(1); ++j)
        AT(i, j) = A(j, i);
 
    std::ranges::fill(ATAb, 0.0);
    for (std::size_t i = 0; i < P.extent(0); ++i)
      for (std::size_t j = 0; j < P.extent(1); ++j)
        P(i, j) = 0;
 
    // pinv(A) = (A^T * A)^-1 * A^T
    dot(AT, A, ATA);
    inv(ATA, Inv);
    dot(Inv, AT, P);
  }
  else if (A.extent(1) == 1)
  {
    T res = 0;
    for (std::size_t i = 0; i < A.extent(0); ++i)
      for (std::size_t j = 0; j < A.extent(1); ++j)
        res += A(i, j) * A(i, j);
 
    for (std::size_t i = 0; i < A.extent(0); ++i)
      for (std::size_t j = 0; j < A.extent(1); ++j)
        P(j, i) = (1 / res) * A(i, j);
  }
  else
  {
    throw std::runtime_error("math::pinv is not implemented for "
                             + std::to_string(A.extent(0)) + "x"
                             + std::to_string(A.extent(1)) + " matrices.");
  }
}
 
} // namespace dolfinx::math