module Numeric.FFTW.Extra.NumberTheory where import qualified Data.Stream as Stream import qualified Data.NonEmpty as NonEmpty import qualified Data.List.HT as ListHT import Data.Maybe (catMaybes) import Data.Bits (Bits, shiftR, (.|.), (.&.)) import Data.Word (Word64) {- $setup >>> import Numeric.FFTW.Extra.NumberTheory >>> import qualified Test.QuickCheck as QC -} divideMaxPower :: (Integral n) => n -> n -> (Int,n) divideMaxPower fac = let go k n = case divMod n fac of (q,r) -> if r==0 then go (succ k) q else (k,n) in go 0 {- | >>> primeFactors (1::Int) [] >>> primeFactors (2::Int) [(2,1)] >>> primeFactors (3::Int) [(3,1)] >>> primeFactors (4::Int) [(2,2)] >>> primeFactors (44100::Int) [(2,2),(3,2),(5,2),(7,2)] >>> primeFactors (48000::Int) [(2,7),(3,1),(5,3)] prop> \(QC.Positive n) -> all (1<) $ map fst $ primeFactors (n::Int) prop> \(QC.Positive n) -> n == product (map (uncurry (^)) (primeFactors (n::Int))) -} primeFactors :: (Integral n) => n -> [(n, Int)] primeFactors 0 = error "Zero can not be factorized." primeFactors n = filter ((0<) . snd) $ let (k2,n2) = divideMaxPower 2 n in (2,k2) : let go di ni = if di*di <= ni then case divideMaxPower di ni of (ki,r) -> (di,ki) : go (di+2) r else if ni>1 then [(ni,1)] else [] in go 3 n2 {- | prop> \(QC.Positive n) -> elem 1 $ divisors (n::Int) prop> \(QC.Positive n) -> elem n $ divisors (n::Int) prop> \(QC.Positive n) -> all (\k -> mod n k == 0) $ divisors (n::Int) -} divisors :: (Integral n) => n -> [n] divisors n = map product $ traverse (\(p,k) -> take (succ k) $ iterate (p*) 1) $ primeFactors n {- | >>> nextDivisor 44100 (1000::Int) 1050 >>> nextDivisor 44100 (2000::Int) 2100 >>> nextDivisor 44100 (3000::Int) 3150 >>> nextDivisor 44100 (4000::Int) 4410 >>> nextDivisor 48000 (1000::Int) 1000 >>> nextDivisor 48000 (9000::Int) 9600 prop> :{ \(QC.Positive n) -> QC.forAll (QC.choose (1,n::Int)) $ \m -> let k = nextDivisor n m in m <= k && mod n k == 0 :} prop> :{ \(QC.Positive n) -> QC.forAll (QC.elements (divisors (n::Int))) $ \k -> nextDivisor n k == k :} -} nextDivisor :: (Integral n) => n -> n -> n nextDivisor n m = minimum $ filter (m<=) $ divisors n {- | >>> adjacentDivisors 44100 (1000::Int) (980,1050) >>> adjacentDivisors 44100 (2000::Int) (1764,2100) >>> adjacentDivisors 44100 (3000::Int) (2940,3150) >>> adjacentDivisors 44100 (4000::Int) (3675,4410) >>> adjacentDivisors 48000 (1000::Int) (1000,1000) >>> adjacentDivisors 48000 (9000::Int) (8000,9600) prop> :{ \(QC.Positive n) -> QC.forAll (QC.choose (1,n::Int)) $ \m -> let (lower,upper) = adjacentDivisors n m in lower <= m && m <= upper && mod n lower == 0 && mod n upper == 0 :} prop> :{ \(QC.Positive n) -> QC.forAll (QC.elements (divisors (n::Int))) $ \k -> adjacentDivisors n k == (k,k) :} -} adjacentDivisors :: (Integral n) => n -> n -> (n,n) adjacentDivisors n m = let ds = divisors n in (maximum $ filter (m>=) ds, minimum $ filter (m<=) ds) {- | prop> \(QC.NonNegative n) -> elem (ceilingPowerOfTwo (n::Int)) $ iterate (2*) 1 prop> \(QC.NonNegative n) -> n <= ceilingPowerOfTwo (n::Int) prop> \(QC.Positive n) -> ceilingPowerOfTwo (n::Int) < 2*n -} ceilingPowerOfTwo :: (Integral n, Bits n) => n -> n ceilingPowerOfTwo 0 = 1 ceilingPowerOfTwo n = (1+) $ fst $ head $ dropWhile (uncurry (/=)) $ ListHT.mapAdjacent (,) $ scanl (\m d -> shiftR m d .|. m) (n-1) $ iterate (2*) 1 {- | Compute square root and round to the next power of two. To put it more precisely: Compute logarithm to base four, round to the next integer with "round half up" semantics, use this as exponent for a power of two. >>> sqrtToPowerOfTwo (0::Integer) 0 >>> sqrtToPowerOfTwo (1::Integer) 1 >>> map sqrtToPowerOfTwo [7,8,9::Integer] [2,4,4] >>> map sqrtToPowerOfTwo [15,16,17::Integer] [4,4,4] prop> \(QC.Positive n) -> snd (divideMaxPower 2 (sqrtToPowerOfTwo (n::Integer))) == 1 prop> \(QC.NonNegative n) -> let m = sqrtToPowerOfTwo (n::Integer) in m*m <= 2*n && n <= 4*m*m -} sqrtToPowerOfTwo :: (Integral n) => n -> n sqrtToPowerOfTwo 0 = 0 sqrtToPowerOfTwo 1 = 1 sqrtToPowerOfTwo n = fst $ Stream.head $ Stream.dropWhile ((n>=) . snd) $ Stream.zip (Stream.iterate (2*) 1) (Stream.iterate (4*) 2) {- | Choose a 5-smooth number at least as the given one. It is not strictly the closest 5-smooth number because we allow at most two factors 3 and at most one factor 5. prop> \(QC.NonNegative n) -> n <= larger5Smooth (n::Int) prop> \(QC.Positive n) -> 4 * larger5Smooth (n::Int) < 5*n prop> :{ \(QC.Positive n) -> (1==) $ foldl (\k d -> snd $ divideMaxPower d k) (larger5Smooth (n::Int)) [2,3,5] :} -} larger5Smooth :: (Integral n, Bits n) => n -> n larger5Smooth n = let p2 = ceilingPowerOfTwo n rat numer denom = case divMod p2 denom of (q,0) -> Just (q*numer) _ -> Nothing in NonEmpty.minimum $ NonEmpty.cons p2 $ filter (n<=) $ catMaybes $ [rat 3 4, rat 5 8, rat 9 16, rat 15 16, rat 45 64] {- 128 16 144 18 160 20 168 21 (does not work well in Rank1.convolve) 180 22.5 192 24 200 25 240 30 256 32 -} efficientSizes :: Word64 efficientSizes = 0x002222278888CCEE nextEfficient :: (Integral n) => n -> n nextEfficient n = fromIntegral $ 32 - efficientSizes `shiftR` ((fromIntegral n-17)*4) .&. 0xF {- | Round to next higher number of the form @m*2^n@ with @m<32@ and @m@ known to be an efficient size. We found efficient sizes by measurements, though we only tested with FFTW's estimate planning. We tested all sizes from 1 to 256 (and did a batched FFT on 10000 data sets, in order to have something to measure) and marked every size as inefficient when there was a larger size that required absolutely less transform time. The remaining sizes are kept as efficient. >>> map ceilingEfficient [1..(40::Int)] [1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,18,18,20,20,24,24,24,24,25,30,30,30,30,30,32,32,36,36,36,36,40,40,40,40] prop> \(QC.NonNegative n) -> n <= ceilingEfficient (n::Int) prop> \(QC.NonNegative n) -> 5 * ceilingEfficient n <= (6*n :: Int) -} ceilingEfficient :: (Integral n) => n -> n ceilingEfficient n = (\(m,p) -> p * if m<=16 then m else nextEfficient m) $ Stream.head $ Stream.dropWhile ((>=32).fst) $ Stream.zip (Stream.iterate (\k -> div (k+1) 2) n) (Stream.iterate (2*) 1)