Merge pull request #90 from Bodigrim/moebius

Moebius and Mertens functions
Bodigrim · Jan 19, 2018 · 09dbf9c · 09dbf9c
2 parents 99c678a + 6ba2160
commit 09dbf9c
Show file tree

Hide file tree

Showing 23 changed files with 571 additions and 80 deletions.
diff --git a/Changes b/Changes
@@ -96,7 +96,7 @@
     Add new cabal flag check-bounds, which replaces all unsafe array functions with safe ones.
 
     Add basic functions on Gaussian integers.
-    Add Moebius mu-function.
+    Add Möbius mu-function.
 
     Forbid non-positive moduli in Math.NumberTheory.Moduli.
 
@@ -123,7 +123,7 @@
 0.4.0.1:
     Fixed Haddock bug
 0.4.0.0:
-    Added generalised Moebius inversion, to be continued
+    Added generalised Möbius inversion, to be continued
 0.3.0.0:
     Added modular square roots and Chinese remainder theorem
 0.2.0.6:

diff --git a/Math/NumberTheory/ArithmeticFunctions/Mertens.hs b/Math/NumberTheory/ArithmeticFunctions/Mertens.hs
@@ -0,0 +1,75 @@
+-- |
+-- Module:      Math.NumberTheory.ArithmeticFunctions.Mertens
+-- Copyright:   (c) 2018 Andrew Lelechenko
+-- Licence:     MIT
+-- Maintainer:  Andrew Lelechenko <[email protected]>
+-- Stability:   Provisional
+-- Portability: Non-portable (GHC extensions)
+--
+-- Values of <https://en.wikipedia.org/wiki/Mertens_function Mertens function>.
+--
+
+{-# LANGUAGE CPP        #-}
+{-# LANGUAGE LambdaCase #-}
+
+module Math.NumberTheory.ArithmeticFunctions.Mertens
+  ( mertens
+  ) where
+
+import qualified Data.Vector.Unboxed as U
+#if __GLASGOW_HASKELL__ < 709
+import Data.Word
+#endif
+
+import Math.NumberTheory.Powers.Cubes
+import Math.NumberTheory.Powers.Squares
+import Math.NumberTheory.ArithmeticFunctions.Moebius
+
+-- | Compute individual values of Mertens function in O(n^(2/3)) time and space.
+--
+-- The implementation follows Theorem 3.1 from <https://arxiv.org/pdf/1610.08551.pdf Computations of the Mertens function and improved bounds on the Mertens conjecture> by G. Hurst, excluding segmentation of sieves.
+mertens :: Word -> Int
+mertens 0 = 0
+mertens 1 = 1
+mertens x = sumMultMoebius lookupMus (\n -> sfunc (x `quot` n)) [1 .. x `quot` u]
+  where
+    u = (integerSquareRoot x + 1) `max` ((integerCubeRoot x) ^ (2 :: Word) `quot` 2)
+
+    sfunc :: Word -> Int
+    sfunc y
+      = 1
+      - sum [ U.unsafeIndex mes (fromIntegral $ y `quot` n) |  n <- [y `quot` u + 1 .. kappa] ]
+      + fromIntegral kappa * U.unsafeIndex mes (fromIntegral nu)
+      - sumMultMoebius lookupMus (\n -> fromIntegral $ y `quot` n) [1 .. nu]
+      where
+        nu = integerSquareRoot y
+        kappa = y `quot` (nu + 1)
+
+    -- cacheSize ~ u
+    cacheSize :: Word
+    cacheSize = u `max` (x `quot` u) `max` integerSquareRoot x
+
+    -- 1-based index
+    mus :: U.Vector Moebius
+    mus = sieveBlockMoebius 1 cacheSize
+
+    lookupMus :: Word -> Moebius
+    lookupMus i = U.unsafeIndex mus (fromIntegral (i - 1))
+
+    -- 0-based index
+    mes :: U.Vector Int
+    mes = U.scanl' go 0 mus
+      where
+        go acc = \case
+          MoebiusN -> acc - 1
+          MoebiusZ -> acc
+          MoebiusP -> acc + 1
+
+-- | Compute sum (map (\x -> runMoebius (mu x) * f x))
+sumMultMoebius :: (Word -> Moebius) -> (Word -> Int) -> [Word] -> Int
+sumMultMoebius mu f = foldl go 0
+  where
+    go acc i = case mu i of
+      MoebiusN -> acc - f i
+      MoebiusZ -> acc
+      MoebiusP -> acc + f i
diff --git a/Math/NumberTheory/ArithmeticFunctions/Moebius.hs b/Math/NumberTheory/ArithmeticFunctions/Moebius.hs
@@ -0,0 +1,173 @@
+-- |
+-- Module:      Math.NumberTheory.ArithmeticFunctions.Moebius
+-- Copyright:   (c) 2018 Andrew Lelechenko
+-- Licence:     MIT
+-- Maintainer:  Andrew Lelechenko <[email protected]>
+-- Stability:   Provisional
+-- Portability: Non-portable (GHC extensions)
+--
+-- Values of <https://en.wikipedia.org/wiki/Möbius_function Möbius function>.
+--
+
+{-# LANGUAGE BangPatterns          #-}
+{-# LANGUAGE MagicHash             #-}
+{-# LANGUAGE MultiParamTypeClasses #-}
+{-# LANGUAGE TypeFamilies          #-}
+
+module Math.NumberTheory.ArithmeticFunctions.Moebius
+  ( Moebius(..)
+  , runMoebius
+  , sieveBlockMoebius
+  ) where
+
+import Control.Monad (forM_, liftM)
+import Control.Monad.ST (runST)
+import Data.Bits
+import Data.Int
+import Data.Word
+import Data.Semigroup
+import qualified Data.Vector.Generic         as G
+import qualified Data.Vector.Generic.Mutable as M
+import qualified Data.Vector.Primitive as P
+import qualified Data.Vector.Unboxed         as U
+import qualified Data.Vector.Unboxed.Mutable as MU
+import GHC.Exts
+import GHC.Integer.GMP.Internals
+import Unsafe.Coerce
+
+import Math.NumberTheory.Primes (primes)
+import Math.NumberTheory.Powers.Squares (integerSquareRoot)
+import Math.NumberTheory.Utils.FromIntegral (wordToInt)
+
+import Math.NumberTheory.Logarithms
+
+-- | Represents three possible values of <https://en.wikipedia.org/wiki/Möbius_function Möbius function>.
+data Moebius
+  = MoebiusN -- ^ −1
+  | MoebiusZ -- ^  0
+  | MoebiusP -- ^  1
+  deriving (Eq, Ord, Show)
+
+-- | Convert to any numeric type.
+runMoebius :: Num a => Moebius -> a
+runMoebius m = fromInteger (S# (dataToTag# m -# 1#))
+
+fromMoebius :: Moebius -> Int8
+fromMoebius m = fromIntegral $ I# (dataToTag# m)
+{-# INLINE fromMoebius #-}
+
+toMoebius :: Int8 -> Moebius
+toMoebius i = let !(I# i#) = fromIntegral i in tagToEnum# i#
+{-# INLINE toMoebius #-}
+
+newtype instance U.MVector s Moebius = MV_Moebius (P.MVector s Int8)
+newtype instance U.Vector    Moebius = V_Moebius  (P.Vector    Int8)
+
+instance U.Unbox Moebius
+
+instance M.MVector U.MVector Moebius where
+  {-# INLINE basicLength #-}
+  {-# INLINE basicUnsafeSlice #-}
+  {-# INLINE basicOverlaps #-}
+  {-# INLINE basicUnsafeNew #-}
+  {-# INLINE basicInitialize #-}
+  {-# INLINE basicUnsafeReplicate #-}
+  {-# INLINE basicUnsafeRead #-}
+  {-# INLINE basicUnsafeWrite #-}
+  {-# INLINE basicClear #-}
+  {-# INLINE basicSet #-}
+  {-# INLINE basicUnsafeCopy #-}
+  {-# INLINE basicUnsafeGrow #-}
+  basicLength (MV_Moebius v) = M.basicLength v
+  basicUnsafeSlice i n (MV_Moebius v) = MV_Moebius $ M.basicUnsafeSlice i n v
+  basicOverlaps (MV_Moebius v1) (MV_Moebius v2) = M.basicOverlaps v1 v2
+  basicUnsafeNew n = MV_Moebius `liftM` M.basicUnsafeNew n
+  basicInitialize (MV_Moebius v) = M.basicInitialize v
+  basicUnsafeReplicate n x = MV_Moebius `liftM` M.basicUnsafeReplicate n (fromMoebius x)
+  basicUnsafeRead (MV_Moebius v) i = toMoebius `liftM` M.basicUnsafeRead v i
+  basicUnsafeWrite (MV_Moebius v) i x = M.basicUnsafeWrite v i (fromMoebius x)
+  basicClear (MV_Moebius v) = M.basicClear v
+  basicSet (MV_Moebius v) x = M.basicSet v (fromMoebius x)
+  basicUnsafeCopy (MV_Moebius v1) (MV_Moebius v2) = M.basicUnsafeCopy v1 v2
+  basicUnsafeMove (MV_Moebius v1) (MV_Moebius v2) = M.basicUnsafeMove v1 v2
+  basicUnsafeGrow (MV_Moebius v) n = MV_Moebius `liftM` M.basicUnsafeGrow v n
+
+instance G.Vector U.Vector Moebius where
+  {-# INLINE basicUnsafeFreeze #-}
+  {-# INLINE basicUnsafeThaw #-}
+  {-# INLINE basicLength #-}
+  {-# INLINE basicUnsafeSlice #-}
+  {-# INLINE basicUnsafeIndexM #-}
+  {-# INLINE elemseq #-}
+  basicUnsafeFreeze (MV_Moebius v) = V_Moebius `liftM` G.basicUnsafeFreeze v
+  basicUnsafeThaw (V_Moebius v) = MV_Moebius `liftM` G.basicUnsafeThaw v
+  basicLength (V_Moebius v) = G.basicLength v
+  basicUnsafeSlice i n (V_Moebius v) = V_Moebius $ G.basicUnsafeSlice i n v
+  basicUnsafeIndexM (V_Moebius v) i = toMoebius `liftM` G.basicUnsafeIndexM v i
+  basicUnsafeCopy (MV_Moebius mv) (V_Moebius v) = G.basicUnsafeCopy mv v
+  elemseq _ = seq
+
+instance Semigroup Moebius where
+  MoebiusZ <> _ = MoebiusZ
+  _ <> MoebiusZ = MoebiusZ
+  MoebiusP <> a = a
+  a <> MoebiusP = a
+  _ <> _ = MoebiusP
+
+instance Monoid Moebius where
+  mempty  = MoebiusP
+  mappend = (<>)
+
+-- | Evaluate the Möbius function over a block.
+-- Value of @f@ at 0, if zero falls into block, is undefined.
+--
+-- Based on the sieving algorithm from p. 3 of <https://arxiv.org/pdf/1610.08551.pdf Computations of the Mertens function and improved bounds on the Mertens conjecture> by G. Hurst. It is approximately 5x faster than 'Math.NumberTheory.ArithmeticFunctions.SieveBlock.sieveBlockUnboxed'.
+--
+-- > > sieveBlockMoebius 1 10
+-- > [MoebiusP, MoebiusN, MoebiusN, MoebiusZ, MoebiusN, MoebiusP, MoebiusN, MoebiusZ, MoebiusZ, MoebiusP]
+sieveBlockMoebius
+  :: Word
+  -> Word
+  -> U.Vector Moebius
+sieveBlockMoebius _ 0 = U.empty
+sieveBlockMoebius lowIndex' len'
+  = (unsafeCoerce :: U.Vector Word8 -> U.Vector Moebius) $ runST $ do
+    as <- MU.replicate len (0x80 :: Word8)
+    forM_ ps $ \p -> do
+      let offset  = negate lowIndex `mod` p
+          offset2 = negate lowIndex `mod` (p * p)
+          l :: Word8
+          l = fromIntegral $ intLog2 p .|. 1
+      forM_ [offset, offset + p .. len - 1] $ \ix -> do
+        MU.unsafeModify as (\y -> y + l) ix
+      forM_ [offset2, offset2 + p * p .. len - 1] $ \ix -> do
+        MU.unsafeWrite as ix 0
+    forM_ [0 .. len - 1] $ \ix -> do
+      MU.unsafeModify as (mapper ix) ix
+    U.unsafeFreeze as
+
+  where
+    lowIndex :: Int
+    lowIndex = wordToInt lowIndex'
+
+    len :: Int
+    len = wordToInt len'
+
+    highIndex :: Int
+    highIndex = lowIndex + len - 1
+
+    -- Bit fiddling in 'mapper' is correct only
+    -- if all sufficiently small (<= 191) primes has been sieved out.
+    ps :: [Int]
+    ps = takeWhile (<= (191 `max` integerSquareRoot highIndex)) $ map fromInteger primes
+
+    mapper :: Int -> Word8 -> Word8
+    mapper ix val
+      | val .&. 0x80 == 0x00
+      = 1
+      | fromIntegral (val .&. 0x7F) < intLog2 (ix + lowIndex) - 5
+        - (if ix + lowIndex >= 0x100000   then 2 else 0)
+        - (if ix + lowIndex >= 0x10000000 then 1 else 0)
+      = (val .&. 1) `shiftL` 1
+      | otherwise
+      = ((val .&. 1) `xor` 1) `shiftL` 1
diff --git a/Math/NumberTheory/ArithmeticFunctions/SieveBlock.hs b/Math/NumberTheory/ArithmeticFunctions/SieveBlock.hs
@@ -23,10 +23,11 @@ module Math.NumberTheory.ArithmeticFunctions.SieveBlock
   , additiveSieveBlockConfig
   , sieveBlock
   , sieveBlockUnboxed
+  , sieveBlockMoebius
   ) where
 
-import Control.Monad
-import Control.Monad.ST
+import Control.Monad (forM_)
+import Control.Monad.ST (runST)
 import Data.Coerce
 import qualified Data.Vector as V
 import qualified Data.Vector.Mutable as MV
@@ -36,13 +37,14 @@ import Data.Monoid
 #endif
 
 import Math.NumberTheory.ArithmeticFunctions.Class
+import Math.NumberTheory.ArithmeticFunctions.Moebius (sieveBlockMoebius)
 import Math.NumberTheory.ArithmeticFunctions.SieveBlock.Unboxed
 import Math.NumberTheory.Logarithms (integerLogBase')
-import Math.NumberTheory.Primes
+import Math.NumberTheory.Primes (primes)
 import Math.NumberTheory.Primes.Types
-import Math.NumberTheory.Powers.Squares
+import Math.NumberTheory.Powers.Squares (integerSquareRoot)
 import Math.NumberTheory.Utils (splitOff#)
-import Math.NumberTheory.Utils.FromIntegral
+import Math.NumberTheory.Utils.FromIntegral (wordToInt, intToWord)
 
 -- | 'runFunctionOverBlock' @f@ @x@ @l@ evaluates an arithmetic function
 -- for integers between @x@ and @x+l-1@ and returns a vector of length @l@.
@@ -87,6 +89,7 @@ sieveBlock
   -> Word
   -> Word
   -> V.Vector a
+sieveBlock _ _ 0 = V.empty
 sieveBlock (SieveBlockConfig empty f append) lowIndex' len' = runST $ do
 
     let lowIndex :: Int
@@ -122,7 +125,7 @@ sieveBlock (SieveBlockConfig empty f append) lowIndex' len' = runST $ do
         MV.unsafeWrite as ix (W# a'#)
         MV.unsafeModify bs (\y -> y `append` V.unsafeIndex fs (I# pow#)) ix
 
-    forM_ [0 .. MV.length as - 1] $ \k -> do
+    forM_ [0 .. len - 1] $ \k -> do
       a <- MV.unsafeRead as k
       MV.unsafeModify bs (\b -> if a /= 1 then b `append` f a 1 else b) k
 

diff --git a/Math/NumberTheory/ArithmeticFunctions/SieveBlock/Unboxed.hs b/Math/NumberTheory/ArithmeticFunctions/SieveBlock/Unboxed.hs
@@ -22,17 +22,18 @@ module Math.NumberTheory.ArithmeticFunctions.SieveBlock.Unboxed
   , sieveBlockUnboxed
   ) where
 
-import Control.Monad
-import Control.Monad.ST
+import Control.Monad (forM_)
+import Control.Monad.ST (runST)
 import qualified Data.Vector.Unboxed as V
 import qualified Data.Vector.Unboxed.Mutable as MV
 import GHC.Exts
 
+import Math.NumberTheory.ArithmeticFunctions.Moebius (Moebius)
 import Math.NumberTheory.Logarithms (integerLogBase')
-import Math.NumberTheory.Primes
-import Math.NumberTheory.Powers.Squares
+import Math.NumberTheory.Primes (primes)
+import Math.NumberTheory.Powers.Squares (integerSquareRoot)
 import Math.NumberTheory.Utils (splitOff#)
-import Math.NumberTheory.Utils.FromIntegral
+import Math.NumberTheory.Utils.FromIntegral (wordToInt, intToWord)
 
 -- | A record, which specifies a function to evaluate over a block.
 --
@@ -83,6 +84,7 @@ sieveBlockUnboxed
   -> Word
   -> Word
   -> V.Vector a
+sieveBlockUnboxed _ _ 0 = V.empty
 sieveBlockUnboxed (SieveBlockConfig empty f append) lowIndex' len' = runST $ do
 
     let lowIndex :: Int
@@ -118,7 +120,7 @@ sieveBlockUnboxed (SieveBlockConfig empty f append) lowIndex' len' = runST $ do
         MV.unsafeWrite as ix (W# a'#)
         MV.unsafeModify bs (\y -> y `append` V.unsafeIndex fs (I# pow#)) ix
 
-    forM_ [0 .. MV.length as - 1] $ \k -> do
+    forM_ [0 .. len - 1] $ \k -> do
       a <- MV.unsafeRead as k
       MV.unsafeModify bs (\b -> if a /= 1 then b `append` f a 1 else b) k
 
@@ -127,3 +129,4 @@ sieveBlockUnboxed (SieveBlockConfig empty f append) lowIndex' len' = runST $ do
 {-# SPECIALIZE sieveBlockUnboxed :: SieveBlockConfig Int  -> Word -> Word -> V.Vector Int  #-}
 {-# SPECIALIZE sieveBlockUnboxed :: SieveBlockConfig Word -> Word -> Word -> V.Vector Word #-}
 {-# SPECIALIZE sieveBlockUnboxed :: SieveBlockConfig Bool -> Word -> Word -> V.Vector Bool #-}
+{-# SPECIALIZE sieveBlockUnboxed :: SieveBlockConfig Moebius -> Word -> Word -> V.Vector Moebius #-}