{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}

-----------------------------------------------------------------------------
{-|
Module      : Math.Tensor.LinearAlgebra.Equations
Description : Linear tensor equations.
Copyright   : (c) Nils Alex, 2020
License     : MIT
Maintainer  : nils.alex@fau.de

Linear tensor equations.
-}
-----------------------------------------------------------------------------
module Math.Tensor.LinearAlgebra.Equations
  ( Equation
  , tensorToEquations
  , equationFromRational
  , equationsToSparseMat
  , equationsToMat
  , tensorsToSparseMat
  , tensorsToMat
  , systemRank
  , Solution
  , fromRref
  , fromRrefRev
  , fromRow
  , fromRowRev
  , applySolution
  , solveTensor
  , solveSystem
  , redefineIndets
  ) where

import Math.Tensor
  ( T
  , removeZerosT
  , toListT
  )
import Math.Tensor.LinearAlgebra.Scalar
  ( Poly (Const, Affine, NotSupported)
  , Lin (Lin)
  , polyMap
  , normalize
  )
import Math.Tensor.LinearAlgebra.Matrix
  ( rref
  , isrref
  , verify
  )

import qualified Numeric.LinearAlgebra.Data as HM
  ( Matrix
  , R
  , Z
  , fromLists
  , toLists
  )
import Numeric.LinearAlgebra (rank)

import Data.Maybe (mapMaybe)
import qualified Data.IntMap.Strict as IM
  ( IntMap
  , foldl'
  , map
  , assocs
  , null
  , findWithDefault
  , lookupMax
  , keys
  , fromList
  , (!)
  , difference
  , intersectionWith
  , mapKeys
  , empty
  )
import Data.List (nub, sort)
import Data.Ratio (numerator, denominator)

-- |A linear equation is a mapping from variable
-- indices to coefficients
type Equation a = IM.IntMap a

-- |Extract linear equations from tensor components.
-- The equations are normalized, sorted, and made unique.
tensorToEquations :: Integral a => T (Poly Rational) -> [Equation a]
tensorToEquations :: T (Poly Rational) -> [Equation a]
tensorToEquations = [Equation a] -> [Equation a]
forall a. Eq a => [a] -> [a]
nub ([Equation a] -> [Equation a])
-> (T (Poly Rational) -> [Equation a])
-> T (Poly Rational)
-> [Equation a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Equation a] -> [Equation a]
forall a. Ord a => [a] -> [a]
sort ([Equation a] -> [Equation a])
-> (T (Poly Rational) -> [Equation a])
-> T (Poly Rational)
-> [Equation a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Equation a -> Bool) -> [Equation a] -> [Equation a]
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not (Bool -> Bool) -> (Equation a -> Bool) -> Equation a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Equation a -> Bool
forall a. IntMap a -> Bool
IM.null) ([Equation a] -> [Equation a])
-> (T (Poly Rational) -> [Equation a])
-> T (Poly Rational)
-> [Equation a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (([Int], Poly Rational) -> Equation a)
-> [([Int], Poly Rational)] -> [Equation a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Poly Rational -> Equation a
forall a. Integral a => Poly Rational -> Equation a
equationFromRational (Poly Rational -> Equation a)
-> (([Int], Poly Rational) -> Poly Rational)
-> ([Int], Poly Rational)
-> Equation a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Poly Rational -> Poly Rational
forall a. (Fractional a, Eq a) => Poly a -> Poly a
normalize (Poly Rational -> Poly Rational)
-> (([Int], Poly Rational) -> Poly Rational)
-> ([Int], Poly Rational)
-> Poly Rational
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([Int], Poly Rational) -> Poly Rational
forall a b. (a, b) -> b
snd) ([([Int], Poly Rational)] -> [Equation a])
-> (T (Poly Rational) -> [([Int], Poly Rational)])
-> T (Poly Rational)
-> [Equation a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. T (Poly Rational) -> [([Int], Poly Rational)]
forall v. T v -> [([Int], v)]
toListT

-- |Extract linear equation with integral coefficients from polynomial
-- tensor component with rational coefficients.
-- Made made integral by multiplying with the @lcm@ of all denominators.
equationFromRational :: forall a.Integral a => Poly Rational -> Equation a
equationFromRational :: Poly Rational -> Equation a
equationFromRational (Affine Rational
x (Lin IntMap Rational
lin))
    | Rational
x Rational -> Rational -> Bool
forall a. Eq a => a -> a -> Bool
== Rational
0 = Equation a
lin'
    | Bool
otherwise = [Char] -> Equation a
forall a. HasCallStack => [Char] -> a
error [Char]
"affine equation not supported for the moment!"
  where
    fac :: a
    fac :: a
fac = (a -> Rational -> a) -> a -> IntMap Rational -> a
forall a b. (a -> b -> a) -> a -> IntMap b -> a
IM.foldl' (\a
acc Rational
v -> a -> a -> a
forall a. Integral a => a -> a -> a
lcm (Integer -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Rational -> Integer
forall a. Ratio a -> a
denominator Rational
v)) a
acc) a
1 IntMap Rational
lin
    lin' :: Equation a
lin' = (Rational -> a) -> IntMap Rational -> Equation a
forall a b. (a -> b) -> IntMap a -> IntMap b
IM.map (\Rational
v -> Integer -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Rational -> Integer
forall a. Ratio a -> a
numerator Rational
v) a -> a -> a
forall a. Num a => a -> a -> a
* (a
fac a -> a -> a
forall a. Integral a => a -> a -> a
`div` Integer -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Rational -> Integer
forall a. Ratio a -> a
denominator Rational
v))) IntMap Rational
lin
equationFromRational (Const Rational
c)
    | Rational
c Rational -> Rational -> Bool
forall a. Eq a => a -> a -> Bool
== Rational
0    = Equation a
forall a. IntMap a
IM.empty
equationFromRational Poly Rational
_ = [Char] -> Equation a
forall a. HasCallStack => [Char] -> a
error [Char]
"equation can only be extracted from linear scalar!"

-- |Convert list of equations to sparse matrix representation of the
-- linear system.
equationsToSparseMat :: [Equation a] -> [((Int,Int), a)]
equationsToSparseMat :: [Equation a] -> [((Int, Int), a)]
equationsToSparseMat [Equation a]
xs = [[((Int, Int), a)]] -> [((Int, Int), a)]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[((Int, Int), a)]] -> [((Int, Int), a)])
-> [[((Int, Int), a)]] -> [((Int, Int), a)]
forall a b. (a -> b) -> a -> b
$ (Int -> Equation a -> [((Int, Int), a)])
-> [Int] -> [Equation a] -> [[((Int, Int), a)]]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Int
i Equation a
m -> ((Int, a) -> ((Int, Int), a)) -> [(Int, a)] -> [((Int, Int), a)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(Int
j,a
v) -> ((Int
i,Int
j),a
v)) (Equation a -> [(Int, a)]
forall a. IntMap a -> [(Int, a)]
IM.assocs Equation a
m)) [Int
1..] [Equation a]
xs

-- |Convert list of equations to dense matrix representation of the
-- linear system.
equationsToMat :: Integral a => [Equation a] -> [[a]]
equationsToMat :: [Equation a] -> [[a]]
equationsToMat [Equation a]
eqns = (Equation a -> Maybe [a]) -> [Equation a] -> [[a]]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (\Equation a
m -> if Equation a -> Bool
forall a. IntMap a -> Bool
IM.null Equation a
m
                                      then Maybe [a]
forall a. Maybe a
Nothing
                                      else [a] -> Maybe [a]
forall a. a -> Maybe a
Just ([a] -> Maybe [a]) -> [a] -> Maybe [a]
forall a b. (a -> b) -> a -> b
$ (Int -> a) -> [Int] -> [a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\Int
j -> a -> Int -> Equation a -> a
forall a. a -> Int -> IntMap a -> a
IM.findWithDefault a
0 Int
j Equation a
m) [Int
1..Int
maxVar]) [Equation a]
eqns
  where
    maxVar :: Int
maxVar = [Int] -> Int
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum ([Int] -> Int) -> [Int] -> Int
forall a b. (a -> b) -> a -> b
$ (Equation a -> Maybe Int) -> [Equation a] -> [Int]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (((Int, a) -> Int) -> Maybe (Int, a) -> Maybe Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Int, a) -> Int
forall a b. (a, b) -> a
fst (Maybe (Int, a) -> Maybe Int)
-> (Equation a -> Maybe (Int, a)) -> Equation a -> Maybe Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Equation a -> Maybe (Int, a)
forall a. IntMap a -> Maybe (Int, a)
IM.lookupMax) [Equation a]
eqns

-- |Extract sparse matrix representation for the linear system given
-- by a list of existentially quantified tensors with polynomial values.
tensorsToSparseMat :: Integral a => [T (Poly Rational)] -> [((Int,Int), a)]
tensorsToSparseMat :: [T (Poly Rational)] -> [((Int, Int), a)]
tensorsToSparseMat = [Equation a] -> [((Int, Int), a)]
forall a. [Equation a] -> [((Int, Int), a)]
equationsToSparseMat ([Equation a] -> [((Int, Int), a)])
-> ([T (Poly Rational)] -> [Equation a])
-> [T (Poly Rational)]
-> [((Int, Int), a)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (T (Poly Rational) -> [Equation a])
-> [T (Poly Rational)] -> [Equation a]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap T (Poly Rational) -> [Equation a]
forall a. Integral a => T (Poly Rational) -> [Equation a]
tensorToEquations

-- |Extract dense matrix representation for the linear system given
-- by a list of existentially quantified tensors with polynomial values.
tensorsToMat :: Integral a => [T (Poly Rational)] -> [[a]]
tensorsToMat :: [T (Poly Rational)] -> [[a]]
tensorsToMat = [Equation a] -> [[a]]
forall a. Integral a => [Equation a] -> [[a]]
equationsToMat ([Equation a] -> [[a]])
-> ([T (Poly Rational)] -> [Equation a])
-> [T (Poly Rational)]
-> [[a]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (T (Poly Rational) -> [Equation a])
-> [T (Poly Rational)] -> [Equation a]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap T (Poly Rational) -> [Equation a]
forall a. Integral a => T (Poly Rational) -> [Equation a]
tensorToEquations

-- |Rank of a linear system. Uses dense svd provided by hmatrix.
matRank :: forall a.Integral a => [[a]] -> Int
matRank :: [[a]] -> Int
matRank []  = Int
0
matRank [[a]]
mat = Matrix R -> Int
forall t. Field t => Matrix t -> Int
rank (Matrix R
hmat :: HM.Matrix HM.R)
  where
    hmat :: Matrix R
hmat = [[R]] -> Matrix R
forall t. Element t => [[t]] -> Matrix t
HM.fromLists ([[R]] -> Matrix R) -> [[R]] -> Matrix R
forall a b. (a -> b) -> a -> b
$ ([a] -> [R]) -> [[a]] -> [[R]]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((a -> R) -> [a] -> [R]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Integral a, Num R) => a -> R
forall a b. (Integral a, Num b) => a -> b
fromIntegral @a @HM.R)) [[a]]
mat

-- |Rank of the linear system given by a list of existentially
-- quantified tensors with polynomial values.
systemRank :: [T (Poly Rational)] -> Int
systemRank :: [T (Poly Rational)] -> Int
systemRank = [[Int]] -> Int
forall a. Integral a => [[a]] -> Int
matRank ([[Int]] -> Int)
-> ([T (Poly Rational)] -> [[Int]]) -> [T (Poly Rational)] -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integral Int => [T (Poly Rational)] -> [[Int]]
forall a. Integral a => [T (Poly Rational)] -> [[a]]
tensorsToMat @Int

-- |The solution to a linear system is represented as a list of
-- substitution rules, stored as @'IM.IntMap' ('Poly' 'Rational')@.
type Solution = IM.IntMap (Poly Rational)

-- |Read substitution rules from reduced row echelon form
-- of a linear system.
fromRref :: HM.Matrix HM.Z -> Solution
fromRref :: Matrix Z -> Solution
fromRref Matrix Z
ref = [(Int, Poly Rational)] -> Solution
forall a. [(Int, a)] -> IntMap a
IM.fromList [(Int, Poly Rational)]
assocs
  where
    rows :: [[Z]]
rows   = Matrix Z -> [[Z]]
forall t. Element t => Matrix t -> [[t]]
HM.toLists Matrix Z
ref
    assocs :: [(Int, Poly Rational)]
assocs = ([Z] -> Maybe (Int, Poly Rational))
-> [[Z]] -> [(Int, Poly Rational)]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe [Z] -> Maybe (Int, Poly Rational)
forall a. Integral a => [a] -> Maybe (Int, Poly Rational)
fromRow [[Z]]
rows

fromRrefRev :: HM.Matrix HM.Z -> Solution
fromRrefRev :: Matrix Z -> Solution
fromRrefRev Matrix Z
ref = [(Int, Poly Rational)] -> Solution
forall a. [(Int, a)] -> IntMap a
IM.fromList [(Int, Poly Rational)]
assocs
  where
    rows :: [[Z]]
rows   = ([Z] -> [Z]) -> [[Z]] -> [[Z]]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [Z] -> [Z]
forall a. [a] -> [a]
reverse ([[Z]] -> [[Z]]) -> [[Z]] -> [[Z]]
forall a b. (a -> b) -> a -> b
$ Matrix Z -> [[Z]]
forall t. Element t => Matrix t -> [[t]]
HM.toLists Matrix Z
ref
    assocs :: [(Int, Poly Rational)]
assocs = ([Z] -> Maybe (Int, Poly Rational))
-> [[Z]] -> [(Int, Poly Rational)]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe [Z] -> Maybe (Int, Poly Rational)
forall a. Integral a => [a] -> Maybe (Int, Poly Rational)
fromRowRev [[Z]]
rows

-- |Read single substitution rule from single
-- row of reduced row echelon form.
fromRow :: forall a.Integral a => [a] -> Maybe (Int, Poly Rational)
fromRow :: [a] -> Maybe (Int, Poly Rational)
fromRow [a]
xs = case [(Int, a)]
assocs of
               []             -> Maybe (Int, Poly Rational)
forall a. Maybe a
Nothing
               [(Int
i,a
_)]        -> (Int, Poly Rational) -> Maybe (Int, Poly Rational)
forall a. a -> Maybe a
Just (Int
i, Rational -> Poly Rational
forall a. a -> Poly a
Const Rational
0)
               (Int
i, a
v):[(Int, a)]
assocs' -> let assocs'' :: [(Int, Rational)]
assocs'' = ((Int, a) -> (Int, Rational)) -> [(Int, a)] -> [(Int, Rational)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(Int
i',a
v') -> (Int
i', - a -> Rational
forall a b. (Integral a, Num b) => a -> b
fromIntegral @a @Rational a
v' Rational -> Rational -> Rational
forall a. Fractional a => a -> a -> a
/ a -> Rational
forall a b. (Integral a, Num b) => a -> b
fromIntegral @a @Rational a
v)) [(Int, a)]
assocs'
                                 in (Int, Poly Rational) -> Maybe (Int, Poly Rational)
forall a. a -> Maybe a
Just (Int
i, Rational -> Lin Rational -> Poly Rational
forall a. a -> Lin a -> Poly a
Affine Rational
0 (IntMap Rational -> Lin Rational
forall a. IntMap a -> Lin a
Lin ([(Int, Rational)] -> IntMap Rational
forall a. [(Int, a)] -> IntMap a
IM.fromList [(Int, Rational)]
assocs'')))
  where
    assocs :: [(Int, a)]
assocs = ((Int, a) -> Bool) -> [(Int, a)] -> [(Int, a)]
forall a. (a -> Bool) -> [a] -> [a]
filter ((a -> a -> Bool
forall a. Eq a => a -> a -> Bool
/=a
0)(a -> Bool) -> ((Int, a) -> a) -> (Int, a) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, a) -> a
forall a b. (a, b) -> b
snd) ([(Int, a)] -> [(Int, a)]) -> [(Int, a)] -> [(Int, a)]
forall a b. (a -> b) -> a -> b
$ [Int] -> [a] -> [(Int, a)]
forall a b. [a] -> [b] -> [(a, b)]
zip [(Int
1::Int)..] [a]
xs

fromRowRev :: forall a.Integral a => [a] -> Maybe (Int, Poly Rational)
fromRowRev :: [a] -> Maybe (Int, Poly Rational)
fromRowRev [a]
xs = case [(Int, a)]
assocs of
                  []             -> Maybe (Int, Poly Rational)
forall a. Maybe a
Nothing
                  [(Int
i,a
_)]        -> (Int, Poly Rational) -> Maybe (Int, Poly Rational)
forall a. a -> Maybe a
Just (Int
i, Rational -> Poly Rational
forall a. a -> Poly a
Const Rational
0)
                  (Int
i, a
v):[(Int, a)]
assocs' -> let assocs'' :: [(Int, Rational)]
assocs'' = ((Int, a) -> (Int, Rational)) -> [(Int, a)] -> [(Int, Rational)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(Int
i',a
v') -> (Int
i', - a -> Rational
forall a b. (Integral a, Num b) => a -> b
fromIntegral @a @Rational a
v' Rational -> Rational -> Rational
forall a. Fractional a => a -> a -> a
/ a -> Rational
forall a b. (Integral a, Num b) => a -> b
fromIntegral @a @Rational a
v)) [(Int, a)]
assocs'
                                    in (Int, Poly Rational) -> Maybe (Int, Poly Rational)
forall a. a -> Maybe a
Just (Int
i, Rational -> Lin Rational -> Poly Rational
forall a. a -> Lin a -> Poly a
Affine Rational
0 (IntMap Rational -> Lin Rational
forall a. IntMap a -> Lin a
Lin ([(Int, Rational)] -> IntMap Rational
forall a. [(Int, a)] -> IntMap a
IM.fromList [(Int, Rational)]
assocs'')))
  where
    assocs :: [(Int, a)]
assocs = [(Int, a)] -> [(Int, a)]
forall a. [a] -> [a]
reverse ([(Int, a)] -> [(Int, a)]) -> [(Int, a)] -> [(Int, a)]
forall a b. (a -> b) -> a -> b
$ ((Int, a) -> Bool) -> [(Int, a)] -> [(Int, a)]
forall a. (a -> Bool) -> [a] -> [a]
filter ((a -> a -> Bool
forall a. Eq a => a -> a -> Bool
/=a
0)(a -> Bool) -> ((Int, a) -> a) -> (Int, a) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, a) -> a
forall a b. (a, b) -> b
snd) ([(Int, a)] -> [(Int, a)]) -> [(Int, a)] -> [(Int, a)]
forall a b. (a -> b) -> a -> b
$ [Int] -> [a] -> [(Int, a)]
forall a b. [a] -> [b] -> [(a, b)]
zip [(Int
1::Int)..] [a]
xs

-- |Apply substitution rules to tensor component.
applySolution :: Solution -> Poly Rational -> Poly Rational
applySolution :: Solution -> Poly Rational -> Poly Rational
applySolution Solution
s (Affine Rational
x (Lin IntMap Rational
lin))
    | Rational
x Rational -> Rational -> Bool
forall a. Eq a => a -> a -> Bool
== Rational
0 = case Poly Rational
p of
                 Affine Rational
xFin (Lin IntMap Rational
linFin) -> if IntMap Rational -> Bool
forall a. IntMap a -> Bool
IM.null IntMap Rational
linFin
                                             then Rational -> Poly Rational
forall a. a -> Poly a
Const Rational
xFin
                                             else Poly Rational
p
                 Poly Rational
_ -> Poly Rational
p
    | Bool
otherwise = [Char] -> Poly Rational
forall a. HasCallStack => [Char] -> a
error [Char]
"affine equations not yet supported"
  where
    s' :: Solution
s' = (Poly Rational -> Rational -> Poly Rational)
-> Solution -> IntMap Rational -> Solution
forall a b c. (a -> b -> c) -> IntMap a -> IntMap b -> IntMap c
IM.intersectionWith (\Poly Rational
row Rational
v -> if Rational
v Rational -> Rational -> Bool
forall a. Eq a => a -> a -> Bool
== Rational
0
                                        then [Char] -> Poly Rational
forall a. HasCallStack => [Char] -> a
error [Char]
"value 0 encountered in linear scalar"
                                        else (Rational -> Rational) -> Poly Rational -> Poly Rational
forall a b. (a -> b) -> Poly a -> Poly b
polyMap (Rational
vRational -> Rational -> Rational
forall a. Num a => a -> a -> a
*) Poly Rational
row) Solution
s IntMap Rational
lin
    lin' :: IntMap Rational
lin' = IntMap Rational -> Solution -> IntMap Rational
forall a b. IntMap a -> IntMap b -> IntMap a
IM.difference IntMap Rational
lin Solution
s
    p0 :: Poly Rational
p0 = if IntMap Rational -> Bool
forall a. IntMap a -> Bool
IM.null IntMap Rational
lin'
         then Rational -> Poly Rational
forall a. a -> Poly a
Const Rational
0
         else Rational -> Lin Rational -> Poly Rational
forall a. a -> Lin a -> Poly a
Affine Rational
0 (IntMap Rational -> Lin Rational
forall a. IntMap a -> Lin a
Lin IntMap Rational
lin')

    p :: Poly Rational
p = (Poly Rational -> Poly Rational -> Poly Rational)
-> Poly Rational -> Solution -> Poly Rational
forall a b. (a -> b -> a) -> a -> IntMap b -> a
IM.foldl' Poly Rational -> Poly Rational -> Poly Rational
forall a. Num a => a -> a -> a
(+) Poly Rational
p0 Solution
s'

    {-
    lin' = IM.foldlWithKey'
             (\lin' i sub ->
                 case Affine 0 (Lin lin') + sub of
                   Affine 0 (Lin lin'') -> IM.delete i lin''
                   Const 0              -> IM.empty
                   _                    -> error "affine equations not yet supported")
             lin s'
    -}
applySolution Solution
_ Poly Rational
_ = [Char] -> Poly Rational
forall a. HasCallStack => [Char] -> a
error [Char]
"only linear equations supported"

-- |Apply substitution rules to all components of a tensor.
solveTensor :: Solution -> T (Poly Rational) -> T (Poly Rational)
solveTensor :: Solution -> T (Poly Rational) -> T (Poly Rational)
solveTensor Solution
sol = T (Poly Rational) -> T (Poly Rational)
forall v. (Eq v, Num v) => T v -> T v
removeZerosT (T (Poly Rational) -> T (Poly Rational))
-> (T (Poly Rational) -> T (Poly Rational))
-> T (Poly Rational)
-> T (Poly Rational)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Poly Rational -> Poly Rational)
-> T (Poly Rational) -> T (Poly Rational)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Solution -> Poly Rational -> Poly Rational
applySolution Solution
sol)

-- |Solve a linear system and apply solution to the tensorial
-- indeterminants.
solveSystem ::
     [T (Poly Rational)] -- ^ Tensorial linear system
  -> [T (Poly Rational)] -- ^ List of indeterminant tensors
  -> [T (Poly Rational)] -- ^ Solved indeterminant tensors
solveSystem :: [T (Poly Rational)] -> [T (Poly Rational)] -> [T (Poly Rational)]
solveSystem [T (Poly Rational)]
system [T (Poly Rational)]
indets
    | Bool
wrongSolution = [Char] -> [T (Poly Rational)]
forall a. HasCallStack => [Char] -> a
error [Char]
"Wrong solution found. May be an Int64 overflow."
    | Bool
otherwise     = [T (Poly Rational)]
indets'
  where
    mat :: Matrix Z
mat = [[Z]] -> Matrix Z
forall t. Element t => [[t]] -> Matrix t
HM.fromLists ([[Z]] -> Matrix Z) -> [[Z]] -> Matrix Z
forall a b. (a -> b) -> a -> b
$ [T (Poly Rational)] -> [[Z]]
forall a. Integral a => [T (Poly Rational)] -> [[a]]
tensorsToMat @HM.Z [T (Poly Rational)]
system
    ref :: Matrix Z
ref = Matrix Z -> Matrix Z
rref Matrix Z
mat
    wrongSolution :: Bool
wrongSolution = Bool -> Bool
not (Matrix Z -> Bool
forall a. (Num a, Ord a, Container Vector a) => Matrix a -> Bool
isrref Matrix Z
ref Bool -> Bool -> Bool
&& Matrix Z -> Matrix Z -> Bool
verify Matrix Z
mat Matrix Z
ref)
    sol :: Solution
sol = Matrix Z -> Solution
fromRref Matrix Z
ref
    indets' :: [T (Poly Rational)]
indets' = (T (Poly Rational) -> T (Poly Rational))
-> [T (Poly Rational)] -> [T (Poly Rational)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Solution -> T (Poly Rational) -> T (Poly Rational)
solveTensor Solution
sol) [T (Poly Rational)]
indets

-- |Relabelling of the indeterminants present in a list of tensors.
-- Redefines the labels of @n@ indeterminants as @[1..n]@, preserving
-- the previous order.
redefineIndets :: [T (Poly v)] -> [T (Poly v)]
redefineIndets :: [T (Poly v)] -> [T (Poly v)]
redefineIndets [T (Poly v)]
indets = (T (Poly v) -> T (Poly v)) -> [T (Poly v)] -> [T (Poly v)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Poly v -> Poly v) -> T (Poly v) -> T (Poly v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\case
                                       Const v
c -> v -> Poly v
forall a. a -> Poly a
Const v
c
                                       Poly v
NotSupported -> Poly v
forall a. Poly a
NotSupported
                                       Affine v
a (Lin IntMap v
lin) ->
                                         v -> Lin v -> Poly v
forall a. a -> Lin a -> Poly a
Affine v
a (IntMap v -> Lin v
forall a. IntMap a -> Lin a
Lin ((Int -> Int) -> IntMap v -> IntMap v
forall a. (Int -> Int) -> IntMap a -> IntMap a
IM.mapKeys (IntMap Int
varMap IntMap Int -> Int -> Int
forall a. IntMap a -> Int -> a
IM.!) IntMap v
lin)))) [T (Poly v)]
indets
  where
    comps :: [Poly v]
comps = ([Int], Poly v) -> Poly v
forall a b. (a, b) -> b
snd (([Int], Poly v) -> Poly v) -> [([Int], Poly v)] -> [Poly v]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (T (Poly v) -> [([Int], Poly v)])
-> [T (Poly v)] -> [([Int], Poly v)]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap T (Poly v) -> [([Int], Poly v)]
forall v. T v -> [([Int], v)]
toListT [T (Poly v)]
indets
    vars :: [Int]
vars = [Int] -> [Int]
forall a. Eq a => [a] -> [a]
nub ([Int] -> [Int]) -> [Int] -> [Int]
forall a b. (a -> b) -> a -> b
$ [[Int]] -> [Int]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[Int]] -> [Int]) -> [[Int]] -> [Int]
forall a b. (a -> b) -> a -> b
$ (Poly v -> Maybe [Int]) -> [Poly v] -> [[Int]]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (\case
                                       Affine v
_ (Lin IntMap v
lin) -> [Int] -> Maybe [Int]
forall a. a -> Maybe a
Just ([Int] -> Maybe [Int]) -> [Int] -> Maybe [Int]
forall a b. (a -> b) -> a -> b
$ IntMap v -> [Int]
forall a. IntMap a -> [Int]
IM.keys IntMap v
lin
                                       Poly v
_                  -> Maybe [Int]
forall a. Maybe a
Nothing) [Poly v]
comps
    varMap :: IntMap Int
varMap = [(Int, Int)] -> IntMap Int
forall a. [(Int, a)] -> IntMap a
IM.fromList ([(Int, Int)] -> IntMap Int) -> [(Int, Int)] -> IntMap Int
forall a b. (a -> b) -> a -> b
$ [Int] -> [Int] -> [(Int, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int]
vars [(Int
1::Int)..]