[Bodigrim#148] Generalize 'inverseSigma' and 'inverseTotient'

Math/NumberTheory/ArithmeticFunctions/Inverse.hs

@@ -14,7 +14,9 @@
 
 module Math.NumberTheory.ArithmeticFunctions.Inverse
   ( inverseTotient
+  , inverseJordan
   , inverseSigma
+  , inverseSigmaK
   , -- * Wrappers
     MinWord(..)
   , MaxWord(..)
@@ -264,12 +266,40 @@ inverseTotient
   => (a -> b)
   -> a
   -> b
-inverseTotient point = invertFunction point totientA invTotient
+inverseTotient = inverseJordan 1
 {-# SPECIALISE inverseTotient :: Semiring b => (Int -> b) -> Int -> b #-}
 {-# SPECIALISE inverseTotient :: Semiring b => (Word -> b) -> Word -> b #-}
 {-# SPECIALISE inverseTotient :: Semiring b => (Integer -> b) -> Integer -> b #-}
 {-# SPECIALISE inverseTotient :: Semiring b => (Natural -> b) -> Natural -> b #-}
 
+-- | The inverse for 'jordan_k' function, where 'k' is a nonnegative integer.
+--
+-- Generalizes the 'inverseTotient' function, which is 'inverseJordan' when 'k'
+-- is '1'.
+--
+-- The return value is parameterized by a 'Semiring', which allows
+-- various applications by providing different (multiplicative) embeddings.
+-- E. g., list all preimages (see a helper 'asSetOfPreimages'):
+--
+-- >>> import qualified Data.Set as S
+-- >>> import Data.Semigroup
+-- >>> S.mapMonotonic getProduct (inverseJordan (S.singleton . Product) 192)
+-- fromList [15,16]
+--
+-- Similarly to 'inverseTotient', it is possible to count and sum preimages, or
+-- get the maximum/minimum preimage.
+inverseJordan
+  :: (Semiring b, Euclidean a, UniqueFactorisation a, Ord a)
+  => Word
+  -> (a -> b)
+  -> a
+  -> b
+inverseJordan k point = invertFunction point (jordanA k) invTotient
+{-# SPECIALISE inverseJordan :: Semiring b => Word -> (Int -> b) -> Int -> b #-}
+{-# SPECIALISE inverseJordan :: Semiring b => Word -> (Word -> b) -> Word -> b #-}
+{-# SPECIALISE inverseJordan :: Semiring b => Word -> (Integer -> b) -> Integer -> b #-}
+{-# SPECIALISE inverseJordan :: Semiring b => Word -> (Natural -> b) -> Natural -> b #-}
+
 -- | The inverse for 'sigma' 1 function.
 --
 -- The return value is parameterized by a 'Semiring', which allows
@@ -302,12 +332,40 @@ inverseSigma
   => (a -> b)
   -> a
   -> b
-inverseSigma point = invertFunction point (sigmaA 1) invSigma
+inverseSigma = inverseSigmaK 1
 {-# SPECIALISE inverseSigma :: Semiring b => (Int -> b) -> Int -> b #-}
 {-# SPECIALISE inverseSigma :: Semiring b => (Word -> b) -> Word -> b #-}
 {-# SPECIALISE inverseSigma :: Semiring b => (Integer -> b) -> Integer -> b #-}
 {-# SPECIALISE inverseSigma :: Semiring b => (Natural -> b) -> Natural -> b #-}
 
+-- | The inverse for 'sigma_k' function, where 'k' is a nonnegative integer.
+--
+-- Generalizes the 'inverseSigma' function, which is 'inverseSigmaK' when 'k'
+-- is '1'.
+--
+-- The return value is parameterized by a 'Semiring', which allows
+-- various applications by providing different (multiplicative) embeddings.
+-- E. g., list all preimages (see a helper 'asSetOfPreimages'):
+--
+-- >>> import qualified Data.Set as S
+-- >>> import Data.Semigroup
+-- >>> S.mapMonotonic getProduct (inverseSigmaK 4 (S.singleton . Product) 6651267)
+-- fromList [50]
+--
+-- Similarly to 'inverseSigma', it is possible to count and sum preimages, or
+-- get the maximum/minimum preimage.
+inverseSigmaK
+  :: (Semiring b, Euclidean a, UniqueFactorisation a, Integral a)
+  => Word
+  -> (a -> b)
+  -> a
+  -> b
+inverseSigmaK k point = invertFunction point (sigmaA k) invSigma
+{-# SPECIALISE inverseSigmaK :: Semiring b => Word -> (Int -> b) -> Int -> b #-}
+{-# SPECIALISE inverseSigmaK :: Semiring b => Word -> (Word -> b) -> Word -> b #-}
+{-# SPECIALISE inverseSigmaK :: Semiring b => Word -> (Integer -> b) -> Integer -> b #-}
+{-# SPECIALISE inverseSigmaK :: Semiring b => Word -> (Natural -> b) -> Natural -> b #-}
+
 --------------------------------------------------------------------------------
 -- Wrappers
 

Math/NumberTheory/Quadratic/EisensteinIntegers.hs

@@ -292,8 +292,7 @@ divideByPrime p p' np k = go k 0
             where
                 (d1, z') = go1 c 0 z
                 d2 = c - d1
-                z'' = head $ drop (wordToInt d2)
-                    $ iterate (\g -> fromMaybe err $ (g * unPrime p) `quotEvenI` np) z'
+                z'' = iterate (\g -> fromMaybe err $ (g * unPrime p) `quotEvenI` np) z' !! max 0 (wordToInt d2)
 
         go1 :: Word -> Word -> EisensteinInteger -> (Word, EisensteinInteger)
         go1 0 d z = (d, z)

test-suite/Math/NumberTheory/ArithmeticFunctions/InverseTests.hs

@@ -36,9 +36,26 @@ import Math.NumberTheory.TestUtils
 totientProperty1 :: forall a. (Semiring a, Euclidean a, Integral a, UniqueFactorisation a, Ord a) => Positive a -> Bool
 totientProperty1 (Positive x) = x `S.member` asSetOfPreimages inverseTotient (totient x)
 
+jordanProperty1
+  :: forall a
+   . (Semiring a, Euclidean a, Integral a, UniqueFactorisation a, Ord a)
+  => Positive Word
+  -> Positive a
+  -> Bool
+jordanProperty1 (Positive k) (Positive x) =
+  x `S.member` asSetOfPreimages (inverseJordan k) (jordan k x)
+
 totientProperty2 :: (Semiring a, Euclidean a, Integral a, UniqueFactorisation a, Ord a) => Positive a -> Bool
 totientProperty2 (Positive x) = all (== x) (S.map totient (asSetOfPreimages inverseTotient x))
 
+jordanProperty2
+  :: (Semiring a, Euclidean a, Integral a, UniqueFactorisation a, Ord a)
+  => Positive Word
+  -> Positive a
+  -> Bool
+jordanProperty2 (Positive k) (Positive x) =
+  all (== x) (S.map (jordan k) (asSetOfPreimages (inverseJordan k) x))
+
 -- | http://oeis.org/A055506
 totientCountFactorial :: [Word]
 totientCountFactorial =
@@ -138,9 +155,29 @@ totientSpecialCases3 = zipWith mkAssert (tail factorial) totientMaxFactorial
 sigmaProperty1 :: forall a. (Semiring a, Euclidean a, UniqueFactorisation a, Integral a) => Positive a -> Bool
 sigmaProperty1 (Positive x) = x `S.member` asSetOfPreimages inverseSigma (sigma 1 x)
 
+sigmaKProperty1
+  :: forall a
+   . (Semiring a, Euclidean a, UniqueFactorisation a, Integral a)
+  => Positive Word
+  -> Positive a
+  -> Bool
+sigmaKProperty1 (Positive k') (Positive x) =
+  -- 'k' shouldn't be too large to avoid very slow tests.
+  let k = succ $ k' `Prelude.div` 50
+  in x `S.member` asSetOfPreimages (inverseSigmaK k) (sigma k x)
+
 sigmaProperty2 :: (Semiring a, Euclidean a, UniqueFactorisation a, Integral a) => Positive a -> Bool
 sigmaProperty2 (Positive x) = all (== x) (S.map (sigma 1) (asSetOfPreimages inverseSigma x))
 
+sigmaKProperty2
+  :: (Semiring a, Euclidean a, UniqueFactorisation a, Integral a)
+  => Positive Word
+  -> Positive a
+  -> Bool
+sigmaKProperty2 (Positive k') (Positive x) =
+  let k = succ $ k' `Prelude.div` 50
+  in all (== x) (S.map (sigma k) (asSetOfPreimages inverseSigma x))
+
 -- | http://oeis.org/A055486
 sigmaCountFactorial :: [Word]
 sigmaCountFactorial =
@@ -260,4 +297,14 @@ testSuite = testGroup "Inverse"
     , testGroup "max"
       (zipWith (\i a -> testCase ("factorial " ++ show i) a) [1..] sigmaSpecialCases3)
     ]
+
+  {-, testGroup "Jordan"
+    [ testIntegralPropertyNoLarge2 "forward"  jordanProperty1
+    , testIntegralPropertyNoLarge2 "backward" jordanProperty2
+    ]-}
+
+  , testGroup  "SigmaK"
+    [ testIntegralPropertyNoLarge2 "forward"  sigmaKProperty1
+    , testIntegralPropertyNoLarge2 "backward" sigmaKProperty2
+    ]
   ]

test-suite/Math/NumberTheory/TestUtils.hs

@@ -12,7 +12,6 @@
 {-# LANGUAGE DataKinds                  #-}
 {-# LANGUAGE FlexibleContexts           #-}
 {-# LANGUAGE FlexibleInstances          #-}
-{-# LANGUAGE KindSignatures             #-}
 {-# LANGUAGE MultiParamTypeClasses      #-}
 {-# LANGUAGE RankNTypes                 #-}
 {-# LANGUAGE ScopedTypeVariables        #-}
@@ -34,6 +33,7 @@ module Math.NumberTheory.TestUtils
   , NonZero(..)
   , testIntegralProperty
   , testIntegralPropertyNoLarge
+  , testIntegralPropertyNoLarge2
   , testSameIntegralProperty
   , testSameIntegralProperty3
   , testIntegral2Property
@@ -43,13 +43,12 @@ module Math.NumberTheory.TestUtils
   , assertEqualUpToEps
   ) where
 
-import Test.SmallCheck.Series (cons2)
 import Test.Tasty
 import Test.Tasty.HUnit       (Assertion, assertBool)
 import Test.Tasty.SmallCheck as SC
 import Test.Tasty.QuickCheck as QC hiding (Positive, getPositive, NonNegative, generate, getNonNegative)
 
-import Test.SmallCheck.Series (Positive(..), NonNegative(..), Serial(..), Series, generate, (\/))
+import Test.SmallCheck.Series (cons2, Positive(..), NonNegative(..), Serial(..), Series, generate, (\/))
 
 import Data.Bits
 import Data.Semiring (Semiring)
@@ -182,6 +181,25 @@ testIntegralPropertyNoLarge name f = testGroup name
   , QC.testProperty "quickcheck Natural" (f :: wrapper Natural -> bool)
   ]
 
+-- The same as 'testIntegralPropertyNoLarge', but with an additional
+-- 'forall wrapper. (TestableIntegral wrapper) => wrapper Word' argument.
+-- This is because the 'inverseSigmaK/inverseJordan' tests require an
+-- additional 'Word' argument for these functions' 'k' index.
+testIntegralPropertyNoLarge2
+  :: forall wrapper bool. (TestableIntegral wrapper, SC.Testable IO bool, QC.Testable bool)
+  => String -> (forall a. (Euclidean a, Semiring a, Integral a, Bits a, UniqueFactorisation a, Show a) => wrapper Word -> wrapper a -> bool) -> TestTree
+testIntegralPropertyNoLarge2 name f = testGroup
+  name
+  [ SC.testProperty "smallcheck Int" (f :: wrapper Word -> wrapper Int -> bool)
+  , SC.testProperty "smallcheck Word" (f :: wrapper Word -> wrapper Word -> bool)
+  , SC.testProperty "smallcheck Integer" (f :: wrapper Word -> wrapper Integer -> bool)
+  , SC.testProperty "smallcheck Natural" (f :: wrapper Word -> wrapper Natural -> bool)
+  , QC.testProperty "quickcheck Int" (f :: wrapper Word -> wrapper Int -> bool)
+  , QC.testProperty "quickcheck Word" (f :: wrapper Word -> wrapper Word -> bool)
+  , QC.testProperty "quickcheck Integer" (f :: wrapper Word -> wrapper Integer -> bool)
+  , QC.testProperty "quickcheck Natural" (f :: wrapper Word -> wrapper Natural -> bool)
+  ]
+
 testSameIntegralProperty
   :: forall wrapper1 wrapper2 bool. (TestableIntegral wrapper1, TestableIntegral wrapper2, SC.Testable IO bool, QC.Testable bool)
   => String -> (forall a. (Euclidean a, Integral a, Bits a, UniqueFactorisation a, Show a) => wrapper1 a -> wrapper2 a -> bool) -> TestTree