add polynomials

This commit is contained in:
rhiannon morris 2024-05-27 20:22:32 +02:00
parent c063107ecc
commit 1d8a6bb325
3 changed files with 313 additions and 0 deletions

189
lib/Quox/Polynomial.idr Normal file
View file

@ -0,0 +1,189 @@
module Quox.Polynomial
import public Quox.Syntax.Subst
import public Quox.Context
import public Quox.Semiring
import Data.Maybe
import Data.SortedMap
import Data.Singleton
import Quox.PrettyValExtra
%default total
%hide Prelude.toList
public export
Monom : Nat -> Type
Monom = Context' Nat
export %inline
mZero : {n : Nat} -> Monom n
mZero = replicate n 0
export %inline
mIsZero : Monom n -> Bool
mIsZero = all (== 0)
private
PolyBody : Type -> Nat -> Type
PolyBody coeff scope = SortedMap (Monom scope) coeff
private %inline
getScope : PolyBody coeff scope -> Maybe (Singleton scope)
getScope p = lengthPrf0 . fst <$> leftMost p
export
data Polynomial : (coeff : Type) -> (scope : Nat) -> Type where
Const : (k : coeff) -> Polynomial coeff scope
Many : (p : PolyBody coeff scope) -> Polynomial coeff scope
%name Polynomial p, q
-- the Const constructor exists so that `pconst` (and by extension, `one`, and
-- the TimesMonoid instance) doesn't need to know the scope size
export %inline
toConst' : List (Monom scope, coeff) -> Maybe coeff
toConst' [(m, k)] = k <$ guard (mIsZero m)
toConst' _ = Nothing
export %inline
toConst : PolyBody coeff scope -> Maybe coeff
toConst = toConst' . toList
export %inline
toTrimmedList : AddMonoid v => SortedMap k v -> List (k, v)
toTrimmedList = List.filter (not . isZero . snd) . toList
export %inline
toList : {scope : Nat} -> AddMonoid coeff =>
Polynomial coeff scope -> List (Monom scope, coeff)
toList (Const k) = [(mZero, k)]
toList (Many p) = toTrimmedList p
private %inline
addBody : AddMonoid coeff =>
PolyBody coeff scope -> PolyBody coeff scope -> PolyBody coeff scope
addBody = mergeWith (+.)
export %inline
fromList : AddMonoid coeff =>
List (Monom scope, coeff) -> Polynomial coeff scope
fromList xs = case toConst' xs of
Just k => Const k
Nothing => Many $ foldl add1 empty xs
where
add1 : PolyBody coeff scope -> (Monom scope, coeff) -> PolyBody coeff scope
add1 p (x, k) = p `addBody` singleton x k
export %inline
(AddMonoid coeff, Eq coeff) => Eq (Polynomial coeff scope) where
Const x == Const y = x == y
Const x == Many q = maybe False (== x) $ toConst q
Many p == Const y = maybe False (== y) $ toConst p
Many p == Many q = toTrimmedList p == toTrimmedList q
export %inline
(AddMonoid coeff, Ord coeff) => Ord (Polynomial coeff scope) where
compare (Const x) (Const y) = compare x y
compare (Const x) (Many q) = maybe LT (compare x) $ toConst q
compare (Many p) (Const y) = maybe GT (flip compare y) $ toConst p
compare (Many p) (Many q) = compare (toTrimmedList p) (toTrimmedList q)
export %inline
{scope : Nat} -> (AddMonoid coeff, Show coeff) =>
Show (Polynomial coeff scope) where
showPrec d p = showCon d "fromList" $ showArg (toList p)
export %inline
{scope : Nat} -> (AddMonoid coeff, PrettyVal coeff) =>
PrettyVal (Polynomial coeff scope) where
prettyVal p = Con "fromList" [prettyVal $ toList p]
export %inline
pconst : a -> Polynomial a n
pconst = Const
export %inline
(.at) : AddMonoid a => Polynomial a n -> Monom n -> a
(Const k).at m = if mIsZero m then k else zero
(Many p).at m = fromMaybe zero $ lookup m p
export %inline
(.at0) : AddMonoid a => Polynomial a n -> a
(Const k).at0 = k
(Many p).at0 = fromMaybe zero $ do
(m, k) <- leftMost p
guard $ mIsZero m
pure k
private %inline
addConstMany : AddMonoid coeff =>
coeff -> PolyBody coeff scope -> Polynomial coeff scope
addConstMany k p = case getScope p of
Nothing => Const k
Just (Val _) => Many $ p `addBody` singleton mZero k
export %inline
AddMonoid a => AddMonoid (Polynomial a n) where
zero = Const zero
isZero (Const k) = isZero k
isZero (Many p) = maybe False isZero $ toConst p
Const j +. Const k = Const $ j +. k
Const j +. Many q = addConstMany j q
Many p +. Const k = addConstMany k p
Many p +. Many q = Many $ p `addBody` q
export %inline
Semiring a => TimesMonoid (Polynomial a n) where
one = pconst one
Const j *. Const k = Const $ j *. k
Const j *. Many q = Many $ map (j *.) q
Many p *. Const k = Many $ map (*. k) p
Many p *. Many q = fromList $ do
(xs, i) <- toList p
(ys, j) <- toList q
let k = i *. j
guard $ not $ isZero k
pure (zipWith (+) xs ys, i *. j)
export %inline
Semiring a => VecSpace a (Polynomial a n) where
x *: xs = Const x *. xs
private
shiftMonom : Shift from to -> Monom from -> Monom to
shiftMonom SZ m = m
shiftMonom (SS by) m = shiftMonom by m :< 0
export
CanShift (Polynomial coeff) where
Const k // _ = Const k
Many p // by =
let xs = mapFst (shiftMonom by) <$> toList p in
Many $ fromList xs
export
TimesMonoid a => FromVarR (Polynomial a) where
fromVarR i l {n} =
let m = tabulateVar n $ \j => if i == j then 1 else 0 in
Many $ singleton m one
private
substMonom : Semiring a => {from, to : Nat} ->
Subst (Polynomial a) from to -> Monom from -> a -> Polynomial a to
substMonom th m k = sproduct {init = pconst k} $
zipWith (\i, p => getR th i noLoc ^. p) (allVars from) m
export
Semiring a => CanSubstSelfR (Polynomial a) where
Const k //? _ = Const k
Many p //? th = ssum $ map (uncurry $ substMonom th) $ toList p

122
lib/Quox/Semiring.idr Normal file
View file

@ -0,0 +1,122 @@
module Quox.Semiring
import Quox.Context
import public Quox.NatExtra
import Data.Bool.Decidable
import Quox.No
%hide Nat.isZero
export infixl 8 +.
export infixl 9 *., *:, :*
export infixr 10 ^.
public export
interface AddMonoid a where
zero : a
isZero : a -> Bool
(+.) : a -> a -> a
public export
interface TimesMonoid a where
one : a
(*.) : a -> a -> a
public export
0 Semiring : Type -> Type
Semiring a = (AddMonoid a, TimesMonoid a)
public export
interface (Semiring base, AddMonoid vec) => VecSpace base vec | vec where
(*:) : base -> vec -> vec
(:*) : vec -> base -> vec
x *: xs = xs :* x
xs :* x = x *: xs
namespace Foldable
public export %inline %tcinline
ssum : AddMonoid a => Foldable t => {default zero init : a} -> t a -> a
ssum = foldl (+.) init
public export %inline %tcinline
sproduct : TimesMonoid a => Foldable t => {default one init : a} -> t a -> a
sproduct = foldl (*.) init
namespace Context
public export %inline %tcinline
ssum : AddMonoid a => {default zero init : a} ->
Telescope' a from to -> a
ssum = foldl_ (+.) zero
public export %inline %tcinline
sproduct : TimesMonoid a => {default one init : a} ->
Telescope' a from to -> a
sproduct = foldl_ (*.) one
public export
(^.) : TimesMonoid a => a -> Nat -> a
x ^. 0 = one
x ^. S k = x *. x^.k
export %inline
[NumAdd] Eq a => Num a => AddMonoid a where
zero = 0; (+.) = (+); isZero = (== 0)
export %inline
[NumTimes] Num a => TimesMonoid a where
one = 1; (*.) = (*)
export %inline %hint
NatAM : AddMonoid Nat
NatAM = NumAdd
export %inline %hint
NatTM : TimesMonoid Nat
NatTM = NumTimes
export %inline %hint
IntegerAM : AddMonoid Integer
IntegerAM = NumAdd
export %inline %hint
IntegerTM : TimesMonoid Integer
IntegerTM = NumTimes
export %inline %hint
IntAM : AddMonoid Int
IntAM = NumAdd
export %inline %hint
IntTM : TimesMonoid Int
IntTM = NumTimes
export %inline
(from `LTE'` to, AddMonoid a) => AddMonoid (Telescope' a from to) where
zero = replicateLTE zero
(+.) = zipWith (+.)
isZero = all isZero
export %inline
(from `LTE'` to, Semiring a) => VecSpace a (Telescope' a from to) where
x *: xs = map (x *.) xs
public export
interface AddMonoid a => MonoAddMonoid a where
0 zeroPlus : (x, y : a) ->
No (isZero x) -> No (isZero y) -> No (isZero (x +. y))
export
MonoAddMonoid Nat where
zeroPlus (S _) (S _) Ah Ah = Ah
public export
0 MonoAddSemiring : Type -> Type
MonoAddSemiring a = (MonoAddMonoid a, TimesMonoid a)

View file

@ -21,6 +21,8 @@ modules =
Quox.No, Quox.No,
Quox.Log, Quox.Log,
Quox.Loc, Quox.Loc,
Quox.Semiring,
Quox.Polynomial,
Quox.Var, Quox.Var,
Quox.Scoped, Quox.Scoped,
Quox.OPE, Quox.OPE,