-- | Mobius inversion for the coarsening poset of partitions

{-# LANGUAGE BangPatterns, TypeSynonymInstances, FlexibleInstances #-}
module Math.RootLoci.Geometry.Mobius 
  ( Partition(..) 
  -- * The refinement poset of partitions
  , coarserThan , finerThan
  , (.==.) , (./=.) , (.<=.) , (.>=.) , (.<.) , (.>.) 
  -- * closures
  , fastClosure , fastAntiClosure
  , closureSet , closureSet'
  -- * Mobius function
  , zetaOf , mobiusOf
  -- * helpers
  , firstLevelDown , firstLevelUp  
  -- * set partitions
  , closureSetOfSetPartition
  , firstLevelDownSetP
  )
  where

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

import Data.List

import qualified Data.Map.Strict as Map ; import Data.Map.Strict (Map)
import qualified Data.Set        as Set ; import Data.Set        (Set)

import Math.Combinat.Partitions.Integer
import Math.Combinat.Partitions.Set
import Math.Combinat.Sets

import qualified Math.Algebra.Polynomial.FreeModule as ZMod

import Math.RootLoci.Algebra
import Math.RootLoci.Misc

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

{-
indicator :: Bool -> Integer
indicator b = if b then 1 else 0

kronecker' :: Partition -> ZMod Partition
kronecker' p = ZMod.singleton p 1

kronecker :: Partition -> Partition -> Integer
kronecker p q = indicator (p .==. q)

zeta :: Partition -> Partition -> Integer
zeta p q = indicator (p .<=. q)
-}

--------------------------------------------------------------------------------
-- * Mobius function

-- | Zeta function of the refinement poset
zetaOf :: Partition -> ZMod Partition
zetaOf :: Partition -> ZMod Partition
zetaOf = (Partition -> ZMod Partition) -> Partition -> ZMod Partition
forall a. (Partition -> a) -> Partition -> a
pcache Partition -> ZMod Partition
forall c. (Eq c, Num c) => Partition -> FreeMod c Partition
calc where
  calc :: Partition -> FreeMod c Partition
calc Partition
p = [(Partition, c)] -> FreeMod c Partition
forall c b. (Eq c, Num c, Ord b) => [(b, c)] -> FreeMod c b
ZMod.fromList ([(Partition, c)] -> FreeMod c Partition)
-> [(Partition, c)] -> FreeMod c Partition
forall a b. (a -> b) -> a -> b
$ (Partition -> (Partition, c)) -> [Partition] -> [(Partition, c)]
forall a b. (a -> b) -> [a] -> [b]
map (\Partition
p -> (Partition
p,c
1)) ([Partition] -> [(Partition, c)])
-> [Partition] -> [(Partition, c)]
forall a b. (a -> b) -> a -> b
$ Set Partition -> [Partition]
forall a. Set a -> [a]
Set.toList (Set Partition -> [Partition]) -> Set Partition -> [Partition]
forall a b. (a -> b) -> a -> b
$ Partition -> Set Partition
closureSet Partition
p

-- | Mobius function of the refinement poset
mobiusOf :: Partition -> ZMod Partition
mobiusOf :: Partition -> ZMod Partition
mobiusOf = (Partition -> ZMod Partition) -> Partition -> ZMod Partition
forall a. (Partition -> a) -> Partition -> a
pcache Partition -> ZMod Partition
calc where
  calc :: Partition -> ZMod Partition
calc    Partition
p = ZMod Partition -> ZMod Partition -> ZMod Partition
forall b c.
(Ord b, Eq c, Num c) =>
FreeMod c b -> FreeMod c b -> FreeMod c b
ZMod.sub (Partition -> Integer -> ZMod Partition
forall b c. (Ord b, Num c, Eq c) => b -> c -> FreeMod c b
ZMod.singleton Partition
p Integer
1) (Partition -> ZMod Partition
smaller Partition
p)
  smaller :: Partition -> ZMod Partition
smaller Partition
p = [ZMod Partition] -> ZMod Partition
forall b c. (Ord b, Eq c, Num c) => [FreeMod c b] -> FreeMod c b
ZMod.sum [ Partition -> ZMod Partition
mobiusOf Partition
q | Partition
q <- Set Partition -> [Partition]
forall a. Set a -> [a]
Set.toList (Partition -> Set Partition
closureSet' Partition
p) ]

--------------------------------------------------------------------------------
-- * The refinement poset of partitions

coarserThan :: Partition -> Partition -> Bool
coarserThan :: Partition -> Partition -> Bool
coarserThan Partition
p Partition
q = Partition -> Set Partition -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member Partition
p (Partition -> Set Partition
closureSet Partition
q)

finerThan :: Partition -> Partition -> Bool
finerThan :: Partition -> Partition -> Bool
finerThan Partition
q Partition
p = Partition -> Partition -> Bool
coarserThan Partition
p Partition
q

(.<=.) :: Partition -> Partition -> Bool
.<=. :: Partition -> Partition -> Bool
(.<=.) = Partition -> Partition -> Bool
coarserThan

(.>=.) :: Partition -> Partition -> Bool
.>=. :: Partition -> Partition -> Bool
(.>=.) = Partition -> Partition -> Bool
finerThan

(.==.) :: Partition -> Partition -> Bool
.==. :: Partition -> Partition -> Bool
(.==.) = Partition -> Partition -> Bool
forall a. Eq a => a -> a -> Bool
(==)

(./=.) :: Partition -> Partition -> Bool
./=. :: Partition -> Partition -> Bool
(./=.) = Partition -> Partition -> Bool
forall a. Eq a => a -> a -> Bool
(/=)

(.<.) :: Partition -> Partition -> Bool
.<. :: Partition -> Partition -> Bool
(.<.) Partition
p Partition
q = (Partition
p Partition -> Partition -> Bool
.<=. Partition
q) Bool -> Bool -> Bool
&& (Partition
p Partition -> Partition -> Bool
forall a. Eq a => a -> a -> Bool
/= Partition
q) 

(.>.) :: Partition -> Partition -> Bool
.>. :: Partition -> Partition -> Bool
(.>.) Partition
p Partition
q = (Partition
p Partition -> Partition -> Bool
.>=. Partition
q) Bool -> Bool -> Bool
&& (Partition
p Partition -> Partition -> Bool
forall a. Eq a => a -> a -> Bool
/= Partition
q) 

--------------------------------------------------------------------------------
-- | Efficient first level merge/split

insertRevSorted :: Int -> [Int] -> [Int]
insertRevSorted :: Int -> [Int] -> [Int]
insertRevSorted Int
x = [Int] -> [Int]
go where
  go :: [Int] -> [Int]
go yys :: [Int]
yys@(Int
y:[Int]
ys) = if Int
x Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
y then Int
x Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int]
yys else Int
y Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int] -> [Int]
go [Int]
ys
  go []         = [Int
x]

insertRevSorted2 :: Int -> Int -> [Int] -> [Int]
insertRevSorted2 :: Int -> Int -> [Int] -> [Int]
insertRevSorted2 Int
x Int
y = Int -> [Int] -> [Int]
insertRevSorted Int
x ([Int] -> [Int]) -> ([Int] -> [Int]) -> [Int] -> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> [Int] -> [Int]
insertRevSorted Int
y

-- | Example: 
-- 
-- > insertGroup [3,3] [[5,5,5],[4],[1,1,1,1]] == [5,5,5,4,3,3,1,1,1,1]
--
insertGroup_ :: [Int] -> [[Int]] -> [Int]
insertGroup_ :: [Int] -> [[Int]] -> [Int]
insertGroup_ zs :: [Int]
zs@(Int
z:[Int]
_) = [[Int]] -> [Int]
go where
  go :: [[Int]] -> [Int]
go (xs :: [Int]
xs@(Int
x:[Int]
_):[[Int]]
rest) = if Int
z Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
x then [Int]
zs [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ [Int]
xs [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ [[Int]] -> [Int]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Int]]
rest 
                                 else [Int]
xs [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ [[Int]] -> [Int]
go [[Int]]
rest
  go ([]      :[[Int]]
rest) = [[Int]] -> [Int]
go [[Int]]
rest
  go []              = [Int]
zs
insertGroup_ [] = [[Int]] -> [Int]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat

-- | These satisfy:
--
-- > concat . insertGroup what == insertGroup_ what
--
insertGroup :: [Int] -> [[Int]] -> [[Int]]
insertGroup :: [Int] -> [[Int]] -> [[Int]]
insertGroup zs :: [Int]
zs@(Int
z:[Int]
_) = [[Int]] -> [[Int]]
go where
  go :: [[Int]] -> [[Int]]
go (xs :: [Int]
xs@(Int
x:[Int]
_):[[Int]]
rest) = if Int
z Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
x then [Int]
zs [Int] -> [[Int]] -> [[Int]]
forall a. a -> [a] -> [a]
: [Int]
xs [Int] -> [[Int]] -> [[Int]]
forall a. a -> [a] -> [a]
: [[Int]]
rest 
                                 else [Int]
xs [Int] -> [[Int]] -> [[Int]]
forall a. a -> [a] -> [a]
: [[Int]] -> [[Int]]
go [[Int]]
rest
  go ([]      :[[Int]]
rest) = [[Int]] -> [[Int]]
go [[Int]]
rest
  go []              = [[Int]
zs]
insertGroup [] = [[Int]] -> [[Int]]
forall a. a -> a
id

insertGroup2_ :: [Int] -> [Int] -> [[Int]] -> [Int]
insertGroup2_ :: [Int] -> [Int] -> [[Int]] -> [Int]
insertGroup2_ [Int]
xs [Int]
ys = [Int] -> [[Int]] -> [Int]
insertGroup_ [Int]
xs ([[Int]] -> [Int]) -> ([[Int]] -> [[Int]]) -> [[Int]] -> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Int] -> [[Int]] -> [[Int]]
insertGroup [Int]
ys

insertGroup2 :: [Int] -> [Int] -> [[Int]] -> [[Int]]
insertGroup2 :: [Int] -> [Int] -> [[Int]] -> [[Int]]
insertGroup2 [Int]
xs [Int]
ys = [Int] -> [[Int]] -> [[Int]]
insertGroup [Int]
xs ([[Int]] -> [[Int]]) -> ([[Int]] -> [[Int]]) -> [[Int]] -> [[Int]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Int] -> [[Int]] -> [[Int]]
insertGroup [Int]
ys

choose1 :: [a] -> [(a,[a])]
choose1 :: [a] -> [(a, [a])]
choose1 (a
x:[a]
xs) = (a
x,[a]
xs) (a, [a]) -> [(a, [a])] -> [(a, [a])]
forall a. a -> [a] -> [a]
: [ (a
y,a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
ys) | (a
y,[a]
ys) <- [a] -> [(a, [a])]
forall a. [a] -> [(a, [a])]
choose1 [a]
xs ]
choose1 []     = []

choose2 :: [a] -> [(a,a,[a])]
choose2 :: [a] -> [(a, a, [a])]
choose2 (a
x:[a]
xs) =  [ (a
x,a
y,[a]
ys  ) |   (a
y,[a]
ys) <- [a] -> [(a, [a])]
forall a. [a] -> [(a, [a])]
choose1 [a]
xs ]
               [(a, a, [a])] -> [(a, a, [a])] -> [(a, a, [a])]
forall a. [a] -> [a] -> [a]
++ [ (a
y,a
z,a
xa -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
zs) | (a
y,a
z,[a]
zs) <- [a] -> [(a, a, [a])]
forall a. [a] -> [(a, a, [a])]
choose2 [a]
xs ]
choose2 []     =  []

-- | Merging two parts
firstLevelDown :: Partition -> [Partition]
firstLevelDown :: Partition -> [Partition]
firstLevelDown (Partition [Int]
ps) = ([Partition]
one [Partition] -> [Partition] -> [Partition]
forall a. [a] -> [a] -> [a]
++ [Partition]
two) where
  gs :: [[Int]]
gs  = [Int] -> [[Int]]
forall a. Eq a => [a] -> [[a]]
group [Int]
ps
  one :: [Partition]
one = [ [Int] -> Partition
Partition ([Int] -> Partition) -> [Int] -> Partition
forall a b. (a -> b) -> a -> b
$ Int -> [Int] -> [Int]
insertRevSorted (Int
xInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
y) ([Int] -> [[Int]] -> [Int]
insertGroup_  [Int]
zs    [[Int]]
rest) | ((Int
x:Int
y:[Int]
zs)     ,[[Int]]
rest) <- [[Int]] -> [([Int], [[Int]])]
forall a. [a] -> [(a, [a])]
choose1 [[Int]]
gs ]
  two :: [Partition]
two = [ [Int] -> Partition
Partition ([Int] -> Partition) -> [Int] -> Partition
forall a b. (a -> b) -> a -> b
$ Int -> [Int] -> [Int]
insertRevSorted (Int
xInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
y) ([Int] -> [Int] -> [[Int]] -> [Int]
insertGroup2_ [Int]
xs [Int]
ys [[Int]]
rest) | ((Int
x:[Int]
xs),(Int
y:[Int]
ys),[[Int]]
rest) <- [[Int]] -> [([Int], [Int], [[Int]])]
forall a. [a] -> [(a, a, [a])]
choose2 [[Int]]
gs ]

-- | Splitting one part into two
firstLevelUp :: Partition -> [Partition]
firstLevelUp :: Partition -> [Partition]
firstLevelUp (Partition [Int]
ps) = [Partition]
one where
  gs :: [[Int]]
gs  = [Int] -> [[Int]]
forall a. Eq a => [a] -> [[a]]
group [Int]
ps
  one :: [Partition]
one = [ [Int] -> Partition
Partition ([Int] -> Partition) -> [Int] -> Partition
forall a b. (a -> b) -> a -> b
$ Int -> Int -> [Int] -> [Int]
insertRevSorted2 Int
x (Int
zInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
x) ([Int] -> [[Int]] -> [Int]
insertGroup_  [Int]
zs [[Int]]
rest) | ((Int
z:[Int]
zs),[[Int]]
rest) <- [[Int]] -> [([Int], [[Int]])]
forall a. [a] -> [(a, [a])]
choose1 [[Int]]
gs , Int
x<-[Int
1..Int -> Int -> Int
forall a. Integral a => a -> a -> a
div Int
z Int
2] ]

-- | Sanity check
firstLevelDownNaive :: Partition -> [Partition]
firstLevelDownNaive :: Partition -> [Partition]
firstLevelDownNaive (Partition [Int]
ps) = [Partition] -> [Partition]
forall a. Ord a => [a] -> [a]
unique [ [Int] -> Partition
mkPartition ( Int
xInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
y Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int]
zs ) | ([Int
x,Int
y],[Int]
zs) <- Int -> [Int] -> [([Int], [Int])]
forall a. Int -> [a] -> [([a], [a])]
choose' Int
2 [Int]
ps ]

firstLevelUpNaive :: Partition -> [Partition]
firstLevelUpNaive :: Partition -> [Partition]
firstLevelUpNaive (Partition [Int]
ps) = [Partition] -> [Partition]
forall a. Ord a => [a] -> [a]
unique [ [Int] -> Partition
mkPartition ( Int
x Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: Int
zInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
x Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int]
zs ) | ([Int
z],[Int]
zs) <- Int -> [Int] -> [([Int], [Int])]
forall a. Int -> [a] -> [([a], [a])]
choose' Int
1 [Int]
ps , Int
x<-[Int
1..Int
zInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1] ]

checkDown :: Partition -> Bool
checkDown :: Partition -> Bool
checkDown Partition
p = ([Partition] -> [Partition]
forall a. Ord a => [a] -> [a]
sort (Partition -> [Partition]
firstLevelDown Partition
p) [Partition] -> [Partition] -> Bool
forall a. Eq a => a -> a -> Bool
== Partition -> [Partition]
firstLevelDownNaive Partition
p)

checkUp :: Partition -> Bool
checkUp :: Partition -> Bool
checkUp Partition
p = ([Partition] -> [Partition]
forall a. Ord a => [a] -> [a]
sort (Partition -> [Partition]
firstLevelUp Partition
p) [Partition] -> [Partition] -> Bool
forall a. Eq a => a -> a -> Bool
== Partition -> [Partition]
firstLevelUpNaive Partition
p)

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

-- | Fast computation of a single closure
fastClosure :: Partition -> Set Partition
fastClosure :: Partition -> Set Partition
fastClosure Partition
p = Set Partition -> [Partition] -> Set Partition
go Set Partition
forall a. Set a
Set.empty [Partition
p] where
  go :: Set Partition -> [Partition] -> Set Partition
go !Set Partition
acc (Partition
p:[Partition]
ps) = case Partition -> Set Partition -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member Partition
p Set Partition
acc of
    Bool
True  -> Set Partition -> [Partition] -> Set Partition
go Set Partition
acc [Partition]
ps
    Bool
False -> Set Partition -> [Partition] -> Set Partition
go (Partition -> Set Partition -> Set Partition
forall a. Ord a => a -> Set a -> Set a
Set.insert Partition
p Set Partition
acc) (Partition -> [Partition]
firstLevelDown Partition
p [Partition] -> [Partition] -> [Partition]
forall a. [a] -> [a] -> [a]
++ [Partition]
ps)
  go !Set Partition
acc []     = Set Partition
acc

-- | Fast computation of a single \"anticlosure\" (opposite poset)
fastAntiClosure :: Partition -> Set Partition
fastAntiClosure :: Partition -> Set Partition
fastAntiClosure Partition
p = Set Partition -> [Partition] -> Set Partition
go Set Partition
forall a. Set a
Set.empty [Partition
p] where
  go :: Set Partition -> [Partition] -> Set Partition
go !Set Partition
acc (Partition
p:[Partition]
ps) = case Partition -> Set Partition -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member Partition
p Set Partition
acc of
    Bool
True  -> Set Partition -> [Partition] -> Set Partition
go Set Partition
acc [Partition]
ps
    Bool
False -> Set Partition -> [Partition] -> Set Partition
go (Partition -> Set Partition -> Set Partition
forall a. Ord a => a -> Set a -> Set a
Set.insert Partition
p Set Partition
acc) (Partition -> [Partition]
firstLevelUp Partition
p [Partition] -> [Partition] -> [Partition]
forall a. [a] -> [a] -> [a]
++ [Partition]
ps)
  go !Set Partition
acc []     = Set Partition
acc

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

-- | Caches and reuses all closures (lazily), this is the fastest version
closureSet :: Partition -> Set Partition 
closureSet :: Partition -> Set Partition
closureSet = Partition -> Set Partition
cached where
  cached :: Partition -> Set Partition
cached = (Partition -> Set Partition) -> Partition -> Set Partition
forall key a. CacheKey key => (key -> a) -> key -> a
monoCache Partition -> Set Partition
calc 
  calc :: Partition -> Set Partition
calc Partition
p = Set Partition -> [Partition] -> Set Partition
go (Partition -> Set Partition
forall a. a -> Set a
Set.singleton Partition
p) (Partition -> [Partition]
firstLevelDown Partition
p) where
    go :: Set Partition -> [Partition] -> Set Partition
go !Set Partition
acc (Partition
p:[Partition]
ps) = case Partition -> Set Partition -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member Partition
p Set Partition
acc of
      Bool
True  -> Set Partition -> [Partition] -> Set Partition
go Set Partition
acc [Partition]
ps
      Bool
False -> Set Partition -> [Partition] -> Set Partition
go (Set Partition -> Set Partition -> Set Partition
forall a. Ord a => Set a -> Set a -> Set a
Set.union Set Partition
acc (Partition -> Set Partition
cached Partition
p)) [Partition]
ps
    go !Set Partition
acc []     = Set Partition
acc

-- | The closure without the stratum itself
closureSet' :: Partition -> Set Partition
closureSet' :: Partition -> Set Partition
closureSet' Partition
p = Partition -> Set Partition -> Set Partition
forall a. Ord a => a -> Set a -> Set a
Set.delete Partition
p (Partition -> Set Partition
closureSet Partition
p)

--------------------------------------------------------------------------------
-- * set partitions

firstLevelDownSetP :: SetPartition -> [SetPartition]
firstLevelDownSetP :: SetPartition -> [SetPartition]
firstLevelDownSetP (SetPartition [[Int]]
ps) =
  [ [[Int]] -> SetPartition
toSetPartition ( ([Int]
x[Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++[Int]
y) [Int] -> [[Int]] -> [[Int]]
forall a. a -> [a] -> [a]
: [[Int]]
zs ) | ([[Int]
x,[Int]
y],[[Int]]
zs) <- Int -> [[Int]] -> [([[Int]], [[Int]])]
forall a. Int -> [a] -> [([a], [a])]
choose' Int
2 [[Int]]
ps ]
  
closureSetOfSetPartition :: SetPartition -> Set SetPartition  
closureSetOfSetPartition :: SetPartition -> Set SetPartition
closureSetOfSetPartition = SetPartition -> Set SetPartition
cached where
  cached :: SetPartition -> Set SetPartition
cached = (SetPartition -> Set SetPartition)
-> SetPartition -> Set SetPartition
forall key a. CacheKey key => (key -> a) -> key -> a
monoCache SetPartition -> Set SetPartition
calc
  calc :: SetPartition -> Set SetPartition
calc SetPartition
p = Set SetPartition -> [SetPartition] -> Set SetPartition
go (SetPartition -> Set SetPartition
forall a. a -> Set a
Set.singleton SetPartition
p) (SetPartition -> [SetPartition]
firstLevelDownSetP SetPartition
p) where
    go :: Set SetPartition -> [SetPartition] -> Set SetPartition
go !Set SetPartition
acc (SetPartition
p:[SetPartition]
ps) = case SetPartition -> Set SetPartition -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member SetPartition
p Set SetPartition
acc of
      Bool
True  -> Set SetPartition -> [SetPartition] -> Set SetPartition
go Set SetPartition
acc [SetPartition]
ps
      Bool
False -> Set SetPartition -> [SetPartition] -> Set SetPartition
go (Set SetPartition -> Set SetPartition -> Set SetPartition
forall a. Ord a => Set a -> Set a -> Set a
Set.union Set SetPartition
acc (SetPartition -> Set SetPartition
cached SetPartition
p)) [SetPartition]
ps
    go !Set SetPartition
acc []     = Set SetPartition
acc
 
--------------------------------------------------------------------------------