principal_invariants.h

#pragma once

#include <funcy/cmath/pow.h>
#include <funcy/concepts.h>
#include <funcy/linalg/cofactor.h>
#include <funcy/linalg/concepts.h>
#include <funcy/linalg/determinant.h>
#include <funcy/linalg/dimension.h>
#include <funcy/linalg/trace.h>
#include <funcy/util/chainer.h>
#include <funcy/util/type_traits.h>

#include <type_traits>

namespace funcy::linalg
{
    namespace detail
    {

        template < int row, Matrix Mat >
        [[nodiscard]] auto symmetric_cofactor_derivative( const Mat& A, const Mat& dA )
        {
            return compute_cofactor_directional_derivative< row, row >( A, dA ) +
                   compute_cofactor_directional_derivative< row, row >( dA, A );
        }

        template < int dim, Matrix Mat >
        [[nodiscard]] static auto sum_of_diagonal_cofactors( const Mat& A ) requires( dim == 2 )
        {
            return linalg::compute_cofactor< 0, 0 >( A ) + linalg::compute_cofactor< 1, 1 >( A );
        }

        template < int dim, Matrix Mat >
        [[nodiscard]] static auto
        sum_of_symmetric_cofactor_derivatives( const Mat& A, const Mat& B ) requires( dim == 2 )
        {
            return symmetric_cofactor_derivative< 0 >( A, B ) +
                   symmetric_cofactor_derivative< 1 >( A, B );
        }

        template < int dim, Matrix Mat >
        [[nodiscard]] static auto sum_of_diagonal_cofactors( const Mat& A ) requires( dim == 3 )
        {
            return linalg::compute_cofactor< 0, 0 >( A ) + linalg::compute_cofactor< 1, 1 >( A ) +
                   linalg::compute_cofactor< 2, 2 >( A );
        }

        template < int dim, Matrix Mat >
        [[nodiscard]] static auto
        sum_of_symmetric_cofactor_derivatives( const Mat& A, const Mat& B ) requires( dim == 3 )
        {
            return symmetric_cofactor_derivative< 0 >( A, B ) +
                   symmetric_cofactor_derivative< 1 >( A, B ) +
                   symmetric_cofactor_derivative< 2 >( A, B );
        }

        template < int dim, Matrix Mat >
        [[nodiscard]] static auto sum_of_diagonal_cofactors( const Mat& A ) requires( dim == -1 )
        {
            if ( rows( A ) == 2 )
                return sum_of_diagonal_cofactors< 2 >( A );
            /*if(rows(A) == 3)*/ return sum_of_diagonal_cofactors< 3 >( A );
        }

        template < int dim, Matrix Mat >
        [[nodiscard]] static auto
        sum_of_symmetric_cofactor_derivatives( const Mat& A, const Mat& B ) requires( dim == -1 )
        {
            if ( rows( A ) == 2 )
                return sum_of_symmetric_cofactor_derivatives< 2 >( A, B );
            /*if(rows(A) == 3)*/ return sum_of_symmetric_cofactor_derivatives< 3 >( A, B );
        }
    } // namespace detail

    template < Matrix Mat >
    class SecondPrincipalInvariant : public Chainer< SecondPrincipalInvariant< Mat > >
    {
    public:
        SecondPrincipalInvariant() = default;

        SecondPrincipalInvariant( const Mat& A )
        {
            update( A );
        }

        void update( const Mat& A )
        {
            if ( !initialized )
            {
                new ( &A_ ) Mat{ A };
                initialized = true;
            }
            else
                A_ = A;
            value = detail::sum_of_diagonal_cofactors< dim< Mat >() >( A );
        }

        [[nodiscard]] auto d0() const noexcept
        {
            return value;
        }

        [[nodiscard]] auto d1( const Mat& dA1 ) const
        {
            return detail::sum_of_symmetric_cofactor_derivatives< dim< Mat >() >( A_, dA1 );
        }

        [[nodiscard]] auto d2( const Mat& dA1, const Mat& dA2 ) const
        {
            return detail::sum_of_symmetric_cofactor_derivatives< dim< Mat >() >( dA1, dA2 );
        }

    private:
        Mat A_;
        std::decay_t< decltype( at( std::declval< Mat >(), 0, 0 ) ) > value = 0;
        bool initialized = false;
    };

    template < class Arg >
    [[nodiscard]] auto i1( Arg&& x )
    {
        return trace( std::forward< Arg >( x ) );
    }

    template < Matrix M >
    [[nodiscard]] auto i2( M&& A )
    {
        return SecondPrincipalInvariant( std::forward< M >( A ) );
    }

    template < Function F >
    [[nodiscard]] auto i2( const F& f )
    {
        return SecondPrincipalInvariant( f() )( f );
    }

    template < class Arg >
    [[nodiscard]] auto i3( const Arg& x )
    {
        return det( x );
    }

    template < class Arg, int n = dim< Arg >() >
    [[nodiscard]] auto mi1( const Arg& x )
    {
        return i1( x ) * pow< -1, n >( det( x ) );
    }

    template < class Arg, int n = dim< Arg >() >
    [[nodiscard]] auto mi2( const Arg& x )
    {
        return i2( x ) * pow< -2, n >( det( x ) );
    }
} // namespace funcy::linalg