{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE ViewPatterns #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}
module Data.Poly.Internal.Multi
( MultiPoly(..)
, VMultiPoly
, UMultiPoly
, toMultiPoly
, toMultiPoly'
, leading
, monomial
, monomial'
, scale
, scale'
, pattern X
, pattern Y
, pattern Z
, pattern X'
, pattern Y'
, pattern Z'
, eval
, eval'
, subst
, subst'
, substitute
, substitute'
, deriv
, deriv'
, integral
, integral'
, Poly
, VPoly
, UPoly
, unPoly
, segregate
, unsegregate
) where
import Prelude hiding (quot, gcd)
import Control.Arrow
import Control.DeepSeq
import Data.Euclidean (Field, quot)
import Data.Finite
import Data.Kind
import Data.List (intersperse)
import Data.Semiring (Semiring(..), Ring())
import qualified Data.Semiring as Semiring
import qualified Data.Vector as V
import qualified Data.Vector.Generic as G
import qualified Data.Vector.Generic.Sized as SG
import qualified Data.Vector.Unboxed as U
import qualified Data.Vector.Unboxed.Sized as SU
import qualified Data.Vector.Sized as SV
import GHC.Exts (IsList(..))
import GHC.TypeNats (KnownNat, Nat, type (+), type (<=))
import Data.Poly.Internal.Multi.Core
newtype MultiPoly (v :: Type -> Type) (n :: Nat) (a :: Type) = MultiPoly
{ unMultiPoly :: v (SU.Vector n Word, a)
}
deriving instance Eq (v (SU.Vector n Word, a)) => Eq (MultiPoly v n a)
deriving instance Ord (v (SU.Vector n Word, a)) => Ord (MultiPoly v n a)
deriving instance NFData (v (SU.Vector n Word, a)) => NFData (MultiPoly v n a)
type VMultiPoly (n :: Nat) (a :: Type) = MultiPoly V.Vector n a
type UMultiPoly (n :: Nat) (a :: Type) = MultiPoly U.Vector n a
type Poly (v :: Type -> Type) (a :: Type) = MultiPoly v 1 a
type VPoly (a :: Type) = Poly V.Vector a
type UPoly (a :: Type) = Poly U.Vector a
unPoly
:: (G.Vector v (Word, a), G.Vector v (SU.Vector 1 Word, a))
=> Poly v a
-> v (Word, a)
unPoly = G.map (first SU.head) . unMultiPoly
instance (Eq a, Semiring a, G.Vector v (SU.Vector n Word, a)) => IsList (MultiPoly v n a) where
type Item (MultiPoly v n a) = (SU.Vector n Word, a)
fromList = toMultiPoly' . G.fromList
fromListN = (toMultiPoly' .) . G.fromListN
toList = G.toList . unMultiPoly
instance (Show a, KnownNat n, G.Vector v (SU.Vector n Word, a)) => Show (MultiPoly v n a) where
showsPrec d (MultiPoly xs)
| G.null xs
= showString "0"
| otherwise
= showParen (d > 0)
$ foldl (.) id
$ intersperse (showString " + ")
$ G.foldl (\acc (is, c) -> showCoeff is c : acc) [] xs
where
showCoeff is c
= showsPrec 7 c . foldl (.) id
( map ((showString " * " .) . uncurry showPower)
$ filter ((/= 0) . fst)
$ zip (SU.toList is) (finites :: [Finite n]))
showPower :: Word -> Finite n -> String -> String
showPower 1 n = showString (showVar n)
showPower i n = showString (showVar n) . showString ("^" ++ show i)
showVar :: Finite n -> String
showVar = \case
0 -> "X"
1 -> "Y"
2 -> "Z"
k -> "X" ++ show k
toMultiPoly
:: (Eq a, Num a, G.Vector v (SU.Vector n Word, a))
=> v (SU.Vector n Word, a)
-> MultiPoly v n a
toMultiPoly = MultiPoly . normalize (/= 0) (+)
toMultiPoly'
:: (Eq a, Semiring a, G.Vector v (SU.Vector n Word, a))
=> v (SU.Vector n Word, a)
-> MultiPoly v n a
toMultiPoly' = MultiPoly . normalize (/= zero) plus
instance (Eq a, Num a, KnownNat n, G.Vector v (SU.Vector n Word, a)) => Num (MultiPoly v n a) where
MultiPoly xs + MultiPoly ys = MultiPoly $ plusPoly (/= 0) (+) xs ys
MultiPoly xs - MultiPoly ys = MultiPoly $ minusPoly (/= 0) negate (-) xs ys
negate (MultiPoly xs) = MultiPoly $ G.map (fmap negate) xs
abs = id
signum = const 1
fromInteger n = case fromInteger n of
0 -> MultiPoly G.empty
m -> MultiPoly $ G.singleton (0, m)
MultiPoly xs * MultiPoly ys = MultiPoly $ convolution (/= 0) (+) (*) xs ys
{-# INLINE (+) #-}
{-# INLINE (-) #-}
{-# INLINE negate #-}
{-# INLINE fromInteger #-}
{-# INLINE (*) #-}
instance (Eq a, Semiring a, KnownNat n, G.Vector v (SU.Vector n Word, a)) => Semiring (MultiPoly v n a) where
zero = MultiPoly G.empty
one
| (one :: a) == zero = zero
| otherwise = MultiPoly $ G.singleton (0, one)
plus (MultiPoly xs) (MultiPoly ys) = MultiPoly $ plusPoly (/= zero) plus xs ys
times (MultiPoly xs) (MultiPoly ys) = MultiPoly $ convolution (/= zero) plus times xs ys
{-# INLINE zero #-}
{-# INLINE one #-}
{-# INLINE plus #-}
{-# INLINE times #-}
fromNatural n = if n' == zero then zero else MultiPoly $ G.singleton (0, n')
where
n' :: a
n' = fromNatural n
{-# INLINE fromNatural #-}
instance (Eq a, Ring a, KnownNat n, G.Vector v (SU.Vector n Word, a)) => Ring (MultiPoly v n a) where
negate (MultiPoly xs) = MultiPoly $ G.map (fmap Semiring.negate) xs
leading :: G.Vector v (SU.Vector 1 Word, a) => Poly v a -> Maybe (Word, a)
leading (MultiPoly v)
| G.null v = Nothing
| otherwise = Just $ first SU.head $ G.last v
scale
:: (Eq a, Num a, KnownNat n, G.Vector v (SU.Vector n Word, a))
=> SU.Vector n Word
-> a
-> MultiPoly v n a
-> MultiPoly v n a
scale yp yc = MultiPoly . scaleInternal (/= 0) (*) yp yc . unMultiPoly
scale'
:: (Eq a, Semiring a, KnownNat n, G.Vector v (SU.Vector n Word, a))
=> SU.Vector n Word
-> a
-> MultiPoly v n a
-> MultiPoly v n a
scale' yp yc = MultiPoly . scaleInternal (/= zero) times yp yc . unMultiPoly
monomial
:: (Eq a, Num a, G.Vector v (SU.Vector n Word, a))
=> SU.Vector n Word
-> a
-> MultiPoly v n a
monomial p c
| c == 0 = MultiPoly G.empty
| otherwise = MultiPoly $ G.singleton (p, c)
monomial'
:: (Eq a, Semiring a, G.Vector v (SU.Vector n Word, a))
=> SU.Vector n Word
-> a
-> MultiPoly v n a
monomial' p c
| c == zero = MultiPoly G.empty
| otherwise = MultiPoly $ G.singleton (p, c)
eval
:: (Num a, G.Vector v (SU.Vector n Word, a), G.Vector u a)
=> MultiPoly v n a
-> SG.Vector u n a
-> a
eval = substitute (*)
{-# INLINE eval #-}
eval'
:: (Semiring a, G.Vector v (SU.Vector n Word, a), G.Vector u a)
=> MultiPoly v n a
-> SG.Vector u n a
-> a
eval' = substitute' times
{-# INLINE eval' #-}
subst
:: (Eq a, Num a, KnownNat m, G.Vector v (SU.Vector n Word, a), G.Vector w (SU.Vector m Word, a))
=> MultiPoly v n a
-> SV.Vector n (MultiPoly w m a)
-> MultiPoly w m a
subst = substitute (scale 0)
{-# INLINE subst #-}
subst'
:: (Eq a, Semiring a, KnownNat m, G.Vector v (SU.Vector n Word, a), G.Vector w (SU.Vector m Word, a))
=> MultiPoly v n a
-> SV.Vector n (MultiPoly w m a)
-> MultiPoly w m a
subst' = substitute' (scale' 0)
{-# INLINE subst' #-}
substitute
:: forall v u n a b.
(G.Vector v (SU.Vector n Word, a), G.Vector u b, Num b)
=> (a -> b -> b)
-> MultiPoly v n a
-> SG.Vector u n b
-> b
substitute f (MultiPoly cs) xs = G.foldl' go 0 cs
where
go :: b -> (SU.Vector n Word, a) -> b
go acc (ps, c) = acc + f c (doMonom ps)
doMonom :: SU.Vector n Word -> b
doMonom = SU.ifoldl' (\acc i p -> acc * ((xs `SG.index` i) ^ p)) 1
{-# INLINE substitute #-}
substitute'
:: forall v u n a b.
(G.Vector v (SU.Vector n Word, a), G.Vector u b, Semiring b)
=> (a -> b -> b)
-> MultiPoly v n a
-> SG.Vector u n b
-> b
substitute' f (MultiPoly cs) xs = G.foldl' go zero cs
where
go :: b -> (SU.Vector n Word, a) -> b
go acc (ps, c) = acc `plus` f c (doMonom ps)
doMonom :: SU.Vector n Word -> b
doMonom = SU.ifoldl' (\acc i p -> acc `times` ((xs `SG.index` i) Semiring.^ p)) one
{-# INLINE substitute' #-}
deriv
:: (Eq a, Num a, G.Vector v (SU.Vector n Word, a))
=> Finite n
-> MultiPoly v n a
-> MultiPoly v n a
deriv i (MultiPoly xs) = MultiPoly $ derivPoly
(/= 0)
(\ps -> ps SU.// [(i, ps `SU.index` i - 1)])
(\ps c -> fromIntegral (ps `SU.index` i) * c)
xs
{-# INLINE deriv #-}
deriv'
:: (Eq a, Semiring a, G.Vector v (SU.Vector n Word, a))
=> Finite n
-> MultiPoly v n a
-> MultiPoly v n a
deriv' i (MultiPoly xs) = MultiPoly $ derivPoly
(/= zero)
(\ps -> ps SU.// [(i, ps `SU.index` i - 1)])
(\ps c -> fromNatural (fromIntegral (ps `SU.index` i)) `times` c)
xs
{-# INLINE deriv' #-}
integral
:: (Fractional a, G.Vector v (SU.Vector n Word, a))
=> Finite n
-> MultiPoly v n a
-> MultiPoly v n a
integral i (MultiPoly xs)
= MultiPoly
$ G.map (\(ps, c) -> let p = ps `SU.index` i in
(ps SU.// [(i, p + 1)], c / fromIntegral (p + 1))) xs
{-# INLINE integral #-}
integral'
:: (Field a, G.Vector v (SU.Vector n Word, a))
=> Finite n
-> MultiPoly v n a
-> MultiPoly v n a
integral' i (MultiPoly xs)
= MultiPoly
$ G.map (\(ps, c) -> let p = ps `SU.index` i in
(ps SU.// [(i, p + 1)], c `quot` Semiring.fromIntegral (p + 1))) xs
{-# INLINE integral' #-}
pattern X
:: (Eq a, Num a, KnownNat n, 1 <= n, G.Vector v (SU.Vector n Word, a))
=> MultiPoly v n a
pattern X <- (isVar 0 -> True)
where X = var 0
pattern X'
:: (Eq a, Semiring a, KnownNat n, 1 <= n, G.Vector v (SU.Vector n Word, a))
=> MultiPoly v n a
pattern X' <- (isVar' 0 -> True)
where X' = var' 0
pattern Y
:: (Eq a, Num a, KnownNat n, 2 <= n, G.Vector v (SU.Vector n Word, a))
=> MultiPoly v n a
pattern Y <- (isVar 1 -> True)
where Y = var 1
pattern Y'
:: (Eq a, Semiring a, KnownNat n, 2 <= n, G.Vector v (SU.Vector n Word, a))
=> MultiPoly v n a
pattern Y' <- (isVar' 1 -> True)
where Y' = var' 1
pattern Z
:: (Eq a, Num a, KnownNat n, 3 <= n, G.Vector v (SU.Vector n Word, a))
=> MultiPoly v n a
pattern Z <- (isVar 2 -> True)
where Z = var 2
pattern Z'
:: (Eq a, Semiring a, KnownNat n, 3 <= n, G.Vector v (SU.Vector n Word, a))
=> MultiPoly v n a
pattern Z' <- (isVar' 2 -> True)
where Z' = var' 2
var
:: forall v n a.
(Eq a, Num a, KnownNat n, G.Vector v (SU.Vector n Word, a))
=> Finite n
-> MultiPoly v n a
var i
| (1 :: a) == 0 = MultiPoly G.empty
| otherwise = MultiPoly $ G.singleton
(SU.generate (\j -> if i == j then 1 else 0), 1)
{-# INLINE var #-}
var'
:: forall v n a.
(Eq a, Semiring a, KnownNat n, G.Vector v (SU.Vector n Word, a))
=> Finite n
-> MultiPoly v n a
var' i
| (one :: a) == zero = MultiPoly G.empty
| otherwise = MultiPoly $ G.singleton
(SU.generate (\j -> if i == j then 1 else 0), one)
{-# INLINE var' #-}
isVar
:: forall v n a.
(Eq a, Num a, KnownNat n, G.Vector v (SU.Vector n Word, a))
=> Finite n
-> MultiPoly v n a
-> Bool
isVar i (MultiPoly xs)
| (1 :: a) == 0 = G.null xs
| otherwise = G.length xs == 1 && G.unsafeHead xs == (SU.generate (\j -> if i == j then 1 else 0), 1)
{-# INLINE isVar #-}
isVar'
:: forall v n a.
(Eq a, Semiring a, KnownNat n, G.Vector v (SU.Vector n Word, a))
=> Finite n
-> MultiPoly v n a
-> Bool
isVar' i (MultiPoly xs)
| (one :: a) == zero = G.null xs
| otherwise = G.length xs == 1 && G.unsafeHead xs == (SU.generate (\j -> if i == j then 1 else 0), one)
{-# INLINE isVar' #-}
groupOn :: (G.Vector v a, Eq b) => (a -> b) -> v a -> [v a]
groupOn f = go
where
go xs
| G.null xs = []
| otherwise = case mk of
Nothing -> [xs]
Just k -> G.unsafeTake (k + 1) xs : go (G.unsafeDrop (k + 1) xs)
where
fy = f (G.unsafeHead xs)
mk = G.findIndex ((/= fy) . f) (G.unsafeTail xs)
segregate
:: (G.Vector v (SU.Vector (1 + m) Word, a), G.Vector v (SU.Vector m Word, a))
=> MultiPoly v (1 + m) a
-> VPoly (MultiPoly v m a)
segregate
= MultiPoly
. G.fromList
. map (\vs -> (SU.take (fst (G.unsafeHead vs)), MultiPoly $ G.map (first SU.tail) vs))
. groupOn (SU.head . fst)
. unMultiPoly
unsegregate
:: (G.Vector v (SU.Vector (1 + m) Word, a), G.Vector v (SU.Vector m Word, a))
=> VPoly (MultiPoly v m a)
-> MultiPoly v (1 + m) a
unsegregate
= MultiPoly
. G.concat
. G.toList
. G.map (\(v, MultiPoly vs) -> G.map (first (v SU.++)) vs)
. unMultiPoly