Monomial.hs

{-# LANGUAGE AllowAmbiguousTypes   #-}
{-# LANGUAGE QuantifiedConstraints #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE ParallelListComp      #-}
{-# LANGUAGE TemplateHaskell       #-}
{-# LANGUAGE UndecidableInstances  #-}

module Math.Monomial
  ( Mon(..)
  , MonomialBasis(..), MonomialBasisQ(..), Deg, Vars
  , zeroMonomial, isZeroMonomial
  , totalDegree, isLinear, linearMonomial

  , split, mons

  , multiSubsetSumFaà, multiSubsetsSum, multiSubsetSum
  , partitionFaà, partitions

  , prodRuleQ

  ) where

-- base
import Data.Foldable
  ( toList )
import Data.Kind
  ( Type, Constraint )
import GHC.Exts
  ( proxy# )
import GHC.TypeNats
--  ( KnownNat, Nat, natVal' )
import Unsafe.Coerce
  ( unsafeCoerce )

-- template-haskell
import Language.Haskell.TH
  ( CodeQ )

-- brush-strokes
import Math.Linear
  ( Vec(..), Fin(..) )
import TH.Utils
  ( foldQ )

--------------------------------------------------------------------------------

-- | @Mon k n@ is the set of monomials in @n@ variables of degree less than or equal to @k@.
type Mon :: Nat -> Nat -> Type
newtype Mon k n = Mon { monDegs :: Vec n Word } -- sum <= k
  deriving stock ( Show, Eq, Ord )

type Deg  :: ( Type -> Type ) -> Nat
type Vars :: ( Type -> Type ) -> Nat
type family Deg  f
type family Vars f

zeroMonomial :: forall k n. KnownNat n => Mon k n
zeroMonomial = Mon $ Vec $ replicate ( fromIntegral $ word @n ) 0
{-# INLINE zeroMonomial #-}

isZeroMonomial :: Mon k n -> Bool
isZeroMonomial ( Mon ( Vec pows ) ) = all ( == 0 ) pows
{-# INLINE isZeroMonomial #-}

totalDegree :: Mon k n -> Word
totalDegree ( Mon ( Vec pows ) ) = sum pows

linearMonomial :: forall k n. ( KnownNat n, 1 <= k ) => Fin n -> Mon k n
linearMonomial ( Fin i' ) =
  Mon $ Vec $ replicate ( i - 1 ) 0 ++ [ 1 ] ++ replicate ( n - i ) 0
  where
    n, i :: Int
    n = fromIntegral $ word @n
    i = fromIntegral i'
{-# INLINE linearMonomial #-}

isLinear :: Mon k n -> Maybe ( Fin n )
isLinear ( Mon ( Vec pows ) ) = fmap Fin $ go 1 pows
  where
    go :: Word -> [ Word ] -> Maybe Word
    go _ [] = Nothing
    go j ( i : is )
      | i == 0
      = go ( j + 1 ) is
      | i == 1 && all ( == 0 ) is
      = Just j
      | otherwise
      = Nothing

-- | Co-multiplication of monomials.
split :: Mon k n -> [ ( Mon k n, Mon k n ) ]
split ( Mon ( Vec [] ) ) = [ ( Mon ( Vec [] ), Mon ( Vec [] ) ) ]
split ( Mon ( Vec ( d : ds ) ) ) =
      [ ( Mon ( Vec ( i : as ) ), Mon ( Vec ( ( ( d - i ) : bs ) ) ) )
      | i <- [ 0 .. d ]
      , ( Mon ( Vec as ), Mon ( Vec bs ) ) <- ( split ( Mon ( Vec ds ) ) )
      ]

-- | All monomials of degree less than or equal to @k@ in @n@ variables,
-- in lexicographic order.
mons :: forall k n. ( KnownNat n ) => Word -> [ Mon k n ]
mons k = unsafeCoerce ( mons' k ( word @n ) )

mons' :: Word -> Word -> [ [ Word ] ]
mons' k _ | k < 0 = []
mons' _ 0 = [ [] ]
mons' 0 n = [ replicate ( fromIntegral n ) 0 ]
mons' k n = [ i : is | i <- reverse [ 0 .. k ], is <- mons' ( k - i ) ( n - 1 ) ]

subs :: Mon k n -> [ ( Mon k n, Word ) ]
subs ( Mon ( Vec []) ) = [ ( Mon ( Vec [] ), maxBound ) ]
subs ( Mon ( Vec ( i : is ) ) )
  = [ ( Mon ( Vec ( j : js ) )
      , if j == 0 then mult else min ( i `quot` j ) mult )
    | j <- [ 0 .. i ]
    , ( Mon ( Vec js ), mult ) <- subs ( Mon ( Vec is ) )
    ]

word :: forall n. KnownNat n => Word
word = fromIntegral $ natVal' @n proxy#

-- | The factorial function \( n! = n \cdot (n-1) \cdot `ldots` \cdot 2 `cdot` 1 \).
factorial :: Word -> Word
factorial i = product [ 1 .. i ]

vecFactorial :: Vec n Word -> Word
vecFactorial ( Vec [] ) = 1
vecFactorial ( Vec ( i : is ) ) = factorial i * vecFactorial ( Vec is )

--------------------------------------------------------------------------------
-- Computations for the chain rule R^1 -> R^n -> R^m

-- | Faà di Bruno coefficient (naive implementation).
multiSubsetSumFaà :: Word -> Vec n [ ( Word, Word ) ] -> Word
multiSubsetSumFaà k multisubsets =
  factorial k `div`
    product [ factorial p * ( factorial i ) ^ p
            | multisubset <- toList multisubsets
            , ( i, p ) <- multisubset ]

-- | Computes the multisubsets of the given set which have the specified sum
-- and number of elements.
multiSubsetSum :: Word -- ^ size of multisubset
               -> Word -- ^ desired sum
               -> [ Word ] -- ^ set to pick from
               -> [ [ ( Word, Word ) ] ]
multiSubsetSum 0 0 _  = [ [] ]
multiSubsetSum 0 _ _  = []
multiSubsetSum _ _ [] = []
multiSubsetSum n s ( i : is ) =
  [ if j == 0 then js else ( i, j ) : js
  | j <- [ 0 .. n ]
  , js <- multiSubsetSum ( n - j ) ( s - i * j ) is
  ]

-- | @multiSubsetsSum is s ns@ computes all collection of multisubsets of @is@,
-- with sizes specified by @ns@, such that the total sum is @s@.
multiSubsetsSum :: forall n
                .  [ Word ]   -- ^ set to pick from
                -> Word       -- ^ desired total sum
                -> Vec n Word -- ^ sizes of each multisets
                -> [ Vec n [ ( Word, Word ) ] ]
multiSubsetsSum is = goMSS
  where
    goMSS :: forall i. Word -> Vec i Word -> [ Vec i [ ( Word, Word ) ] ]
    goMSS 0 ( Vec [] ) = [ Vec [] ]
    goMSS _ ( Vec [] ) = [ ]
    goMSS s ( Vec ( n : ns ) ) =
      [ Vec ( multi : rest )
      | s_i <- [ n * i_min .. s ]
      , multi <- multiSubsetSum n s_i is
      , Vec rest <- goMSS ( s - s_i ) ( Vec ns ) ]
    i_min = case is of
      [] -> 0
      _  -> max 0 $ minimum is

--------------------------------------------------------------------------------
-- Computations for the chain rule R^n -> R^1 -> R^1

partitionFaà :: Mon k n -> [ ( Mon k n, Word ) ] -> Word
partitionFaà ( Mon mon ) multiIndexes =
  vecFactorial mon `div`
    product [ factorial p * ( vecFactorial i ) ^ p
            | ( Mon i, p ) <- toList multiIndexes ]

-- | @partitions p mon@ computes all partitions of the monomial @mon@ into
-- @p@ (non-zero) parts, allowing repetition.
partitions :: forall k n
           .  Word    -- ^ number of parts
           -> Mon k n -- ^ monomial to sum to
           -> [ [ ( Mon k n, Word ) ] ]
partitions n_parts0 mon0 = go n_parts0 Nothing mon0
  where
    go :: Word -> Maybe ( Mon k n ) -> Mon k n -> [ [ ( Mon k n , Word ) ] ]
    go 0 _ mon
      | isZeroMonomial mon
      = [ [] ]
      | otherwise
      = []
    go 1 mbMaxMon mon
      |  isZeroMonomial mon
      || case mbMaxMon of { Just maxMon -> mon >= maxMon ; _ -> False }
      = []
      | otherwise
      = [ [ ( mon, 1 ) ] ]
    go n_parts mbMaxMon mon
      = [ ( part, l ) : parts
        | ( part, l_max ) <- subs mon
        , not ( isZeroMonomial part ) -- parts must be non-empty
        -- use a total ordering on monomials to ensure uniqueness
        , case mbMaxMon of
            Just maxMon -> part < maxMon
            Nothing     -> True
        , l <- [ 1 .. min n_parts l_max ]
        , parts <- go ( n_parts - l ) ( Just part ) ( subPart l mon part )
        ]

subPart :: Word -> Mon k n -> Mon k n -> Mon k n
subPart = unsafeCoerce subPart'

subPart' :: Word -> [ Word ] -> [ Word ] -> [ Word ]
subPart' m = zipWith ( \ i j -> i - m * j )

--------------------------------------------------------------------------------

-- | @'MonomialBasis' f@ exhibits @f u@ as a free @r@-module with basis the
-- monomials in @Vars u@ variables, of degree up to (and including) @Deg u@.
type MonomialBasis :: ( Type -> Type ) -> Constraint
class MonomialBasis f where
  monTabulate :: ( ( Mon ( Deg f ) ( Vars f ) ) -> u ) -> f u
  monIndex :: f u -> Mon ( Deg f ) ( Vars f ) -> u

-- | @'MonomialBasisQ' f@ exhibits @f u@ as a free @r@-module with basis the
-- monomials in @Vars u@ variables, of degree up to (and including) @Deg u@,
-- at compile-time (to be spliced in using TemplateHaskell).
type MonomialBasisQ :: ( Type -> Type ) -> Constraint
class MonomialBasisQ f where
  monTabulateQ :: ( ( Mon ( Deg f ) ( Vars f ) ) -> CodeQ u ) -> CodeQ ( f u )
  monIndexQ :: CodeQ ( f u ) -> Mon ( Deg f ) ( Vars f ) -> CodeQ u

--------------------------------------------------------------------------------

prodRuleQ :: forall f r. MonomialBasisQ f
          => CodeQ r -> CodeQ ( r -> r -> r ) -> CodeQ ( r -> r -> r )
            -- Ring r constraint (circumvent TH constraint problem)
          -> CodeQ ( f r ) -> CodeQ ( f r ) -> CodeQ ( f r )
prodRuleQ zero plus times df1 df2 =
  monTabulateQ @f \ mon ->
    [|| $$( foldQ plus zero
           [ [|| $$times $$( monIndexQ @f df1 m1 ) $$( monIndexQ @f df2 m2 ) ||]
           | (m1, m2) <- split mon
           ] )
    ||]