diff --git a/Setup.hs b/Setup.hs --- a/Setup.hs +++ b/Setup.hs @@ -1,2 +1,32 @@ +{- +Disable some errors and warnings during the haddock pass + (caused by compiler plugins and hs-boot) + -} +{-# OPTIONS_GHC -Wall #-} +module Main (main) where + import Distribution.Simple -main = defaultMain +import Distribution.Simple.Setup + +main :: IO () +main = defaultMainWithHooks simpleUserHooks + { confHook = \a -> confHook simpleUserHooks a . tweakFlags } + +tweakFlags :: ConfigFlags -> ConfigFlags +tweakFlags flags = flags + { configProgramArgs = addHaddockArgs (configProgramArgs flags) } + +addHaddockArgs :: [(String, [String])] -> [(String, [String])] +addHaddockArgs [] + = [("haddock", newHaddockGhcArgs)] +addHaddockArgs (("haddock", args):otherProgsArgs) + = ("haddock", args ++ newHaddockGhcArgs) : otherProgsArgs +addHaddockArgs (progArgs:otherProgsArgs) + = progArgs : addHaddockArgs otherProgsArgs + +newHaddockGhcArgs :: [String] +newHaddockGhcArgs = + [ "--optghc=-fdefer-type-errors" + , "--optghc=-fno-warn-deferred-type-errors" + , "--optghc=-fno-warn-missing-home-modules" + ] diff --git a/bench/subspacefolds.hs b/bench/subspacefolds.hs --- a/bench/subspacefolds.hs +++ b/bench/subspacefolds.hs @@ -23,47 +23,47 @@ seq t1 putStrLn $ "Created DataFrame, elapsed time is " ++ show (diffUTCTime t1 t0) putStrLn "\nRunning a ewfoldl on scalar elements..." - let rezEwf = ewfoldl @Float @DList @'[] (\a x -> return $! fromMaybe x a + fromMaybe 0 a / (x+1)) (Just 1) df + let rezEwf = ewfoldl' @Float @DList @'[] (\a x -> return $! fromMaybe x a + fromMaybe 0 a / (x+1)) (Just 1) df t2 <- rezEwf `seq` getCurrentTime seq t2 putStrLn $ "Done; elapsed time = " ++ show (diffUTCTime t2 t1) print rezEwf putStrLn "\nRunning a iwfoldl on scalar elements (not using idx)..." - let rezIwf = iwfoldl @Float @DList @'[] (\_ a x -> a + a / (x+1)) 1 df + let rezIwf = iwfoldl' @Float @DList @'[] (\_ a x -> a + a / (x+1)) 1 df t3 <- rezIwf `seq` getCurrentTime seq t3 putStrLn $ "Done; elapsed time = " ++ show (diffUTCTime t3 t2) print rezIwf putStrLn "\nRunning a iwfoldr on scalar elements (using fromEnum idx)..." - let rezIwf2 = iwfoldr @Float @DList @'[] (\i x a -> return $! fromMaybe 0 a + x / ((1+) . fromIntegral $ fromEnum i)) (Just 0) df + let rezIwf2 = iwfoldr' @Float @DList @'[] (\i x a -> return $! fromMaybe 0 a + x / ((1+) . fromIntegral $ fromEnum i)) (Just 0) df t4 <- rezIwf2 `seq` getCurrentTime seq t4 putStrLn $ "Done; elapsed time = " ++ show (diffUTCTime t4 t3) print rezIwf2 putStrLn "\nRunning a iwfoldl on scalar elements (enforcing idx)..." - let rezIwf3 = iwfoldl @Float @DList @'[] (\i a x -> i `seq` return $! fromMaybe 0 a + fromMaybe x a / (x+1)) (Just 1) df + let rezIwf3 = iwfoldl' @Float @DList @'[] (\i a x -> i `seq` return $! fromMaybe 0 a + fromMaybe x a / (x+1)) (Just 1) df t5 <- rezIwf3 `seq` getCurrentTime seq t5 putStrLn $ "Done; elapsed time = " ++ show (diffUTCTime t5 t4) print rezIwf3 putStrLn "\nRunning a ewfoldl on vector4 elements..." - let rezEwv1 = ewfoldl @Float @(Init DList) @'[Last DList] - (\a x -> return $! fromMaybe 2 a + fromMaybe 0 a / (1 + iwgen @_ @_ @'[] (\(Idx i:*U) -> x ! Idx (i+1) :* U )) ) + let rezEwv1 = ewfoldl' @Float @(Init DList) @'[Last DList] + (\a x -> return $! fromMaybe 2 a + fromMaybe 0 a / (1 + iwgen @_ @_ @'[] (\(Idx i:*U) -> x ! (i+1) )) ) (Just (3 :: DataFrame Float '[4])) df t6 <- rezEwv1 `seq` getCurrentTime seq t6 putStrLn $ "Done; elapsed time = " ++ show (diffUTCTime t6 t5) print rezEwv1 putStrLn "\nRunning a ewfoldr on vector3 elements..." - let rezEwv2 = ewfoldr @Float @(Init DList) @'[Last DList] - (\x a -> return $! fromMaybe 2 a + fromMaybe 1 a / (1 + iwgen @_ @_ @'[] (\(Idx i:*U) -> x ! Idx (i+1):* U ))) + let rezEwv2 = ewfoldr' @Float @(Init DList) @'[Last DList] + (\x a -> return $! fromMaybe 2 a + fromMaybe 1 a / (1 + iwgen @_ @_ @'[] (\(Idx i:*U) -> x ! (i+1) ))) (Just (3 :: DataFrame Float '[3])) df t7 <- rezEwv2 `seq` getCurrentTime seq t7 putStrLn $ "Done; elapsed time = " ++ show (diffUTCTime t7 t6) print rezEwv2 putStrLn "\nRunning a ewfoldr with matrix products..." - let rezEwm = ewfoldr @Float @(Take 3 DList) @(Drop 3 DList) + let rezEwm = ewfoldr' @Float @(Take 3 DList) @(Drop 3 DList) (\x a -> a + x %* (DF5 1 0.5 0.1 0.01 0.001) ) (1 :: DataFrame Float (Init (Drop 3 DList) +: 3)) df t8 <- rezEwm `seq` getCurrentTime @@ -74,4 +74,4 @@ putStrLn "Checking indexes" - print $ df ! (2:*1:*1:*3:*1:*U :: Idxs (Take 5 DList)) + print $ df.!2.!1.!1.!3.!1 diff --git a/easytensor.cabal b/easytensor.cabal --- a/easytensor.cabal +++ b/easytensor.cabal @@ -1,11 +1,13 @@ --- This file has been generated from package.yaml by hpack version 0.28.2. +cabal-version: 1.24 + +-- This file has been generated from package.yaml by hpack version 0.31.1. -- -- see: https://github.com/sol/hpack -- --- hash: c1b3d6ebe751438aebdd514a166566c79aa2a98083755f019015772fc95cb34e +-- hash: 75fa5546720cb775f5ce9a26382fcd60253bba851441e469156242ee52655f97 name: easytensor -version: 2.0.1.0 +version: 2.1.0.0 synopsis: Pure, type-indexed haskell vector, matrix, and tensor library. description: Pure, type-indexed haskell vector, matrix, and tensor library. Features dimensionality type-checking for all operations. Generic n-dimensional versions are implemented using low-level prim ops. Allows ad-hoc replacement with fixed low-dimensionality vectors and matrices without changing user interface. Please see the README on GitHub at category: math, geometry @@ -16,14 +18,18 @@ copyright: Copyright: (c) 2019 Artem Chirkin license: BSD3 license-file: LICENSE -build-type: Simple -cabal-version: >= 1.10 +build-type: Custom source-repository head type: git location: https://github.com/achirkin/easytensor subdir: easytensor +custom-setup + setup-depends: + Cabal + , base + flag unsafeindices description: Disable bound checks when accessing elements of a tensor, matrix, or vector. manual: True @@ -31,6 +37,7 @@ library exposed-modules: + Numeric.Basics Numeric.DataFrame Numeric.DataFrame.Contraction Numeric.DataFrame.Internal.Backend @@ -41,17 +48,24 @@ Numeric.DataFrame.SubSpace Numeric.DataFrame.Type Numeric.Matrix + Numeric.Matrix.Bidiagonal + Numeric.Matrix.LU + Numeric.Matrix.QR + Numeric.Matrix.SVD Numeric.PrimBytes Numeric.ProductOrd Numeric.ProductOrd.NonTransitive Numeric.ProductOrd.Partial Numeric.Quaternion Numeric.Scalar + Numeric.Subroutine.Householder + Numeric.Subroutine.SolveTriangular + Numeric.Subroutine.Sort Numeric.Vector other-modules: Numeric.Matrix.Internal - Numeric.Matrix.Internal.Mat44d - Numeric.Matrix.Internal.Mat44f + Numeric.Matrix.Internal.Double + Numeric.Matrix.Internal.Float Numeric.Scalar.Internal Numeric.Vector.Internal Numeric.Quaternion.Internal @@ -73,7 +87,7 @@ ghc-options: -Wall -Wcompat -Wtabs -Wmonomorphism-restriction -Wmissing-local-signatures -Wmissing-home-modules -Widentities -Wredundant-constraints build-depends: base >=4.11 && <5 - , constraints-deriving >=1.0.4.0 && <2 + , constraints-deriving >=1.1.1.0 && <2 , dimensions >=2 && <3 if flag(unsafeindices) cpp-options: -DUNSAFE_INDICES @@ -83,14 +97,19 @@ type: exitcode-stdio-1.0 main-is: Spec.hs other-modules: - Numeric.DataFrame.Arbitraries + Numeric.Arbitraries Numeric.DataFrame.BasicTest Numeric.DataFrame.SubSpaceTest + Numeric.Matrix.BidiagonalTest + Numeric.Matrix.LUTest + Numeric.Matrix.QRTest + Numeric.Matrix.SVDTest Numeric.MatrixDoubleTest Numeric.MatrixFloatTest Numeric.PrimBytesTest Numeric.QuaterDoubleTest Numeric.QuaterFloatTest + Numeric.Subroutine.SolveTriangularTest Spec.Util Paths_easytensor hs-source-dirs: @@ -100,7 +119,7 @@ Cabal , QuickCheck , base >=4.11 && <5 - , constraints-deriving >=1.0.4.0 && <2 + , constraints-deriving >=1.1.1.0 && <2 , dimensions >=2 && <3 , easytensor , template-haskell @@ -116,7 +135,7 @@ ghc-options: -Wall -Wcompat -Wtabs -Wmonomorphism-restriction -Wmissing-local-signatures -Wmissing-home-modules -Widentities -Wredundant-constraints build-depends: base >=4.11 && <5 - , constraints-deriving >=1.0.4.0 && <2 + , constraints-deriving >=1.1.1.0 && <2 , dimensions >=2 && <3 , easytensor default-language: Haskell2010 @@ -131,7 +150,7 @@ ghc-options: -Wall -Wcompat -Wtabs -Wmonomorphism-restriction -Wmissing-local-signatures -Wmissing-home-modules -Widentities -Wredundant-constraints build-depends: base >=4.11 && <5 - , constraints-deriving >=1.0.4.0 && <2 + , constraints-deriving >=1.1.1.0 && <2 , dimensions >=2 && <3 , easytensor , time diff --git a/src-base/Numeric/DataFrame/Internal/Backend/Family.hs b/src-base/Numeric/DataFrame/Internal/Backend/Family.hs --- a/src-base/Numeric/DataFrame/Internal/Backend/Family.hs +++ b/src-base/Numeric/DataFrame/Internal/Backend/Family.hs @@ -14,6 +14,8 @@ {-# LANGUAGE TypeFamilyDependencies #-} {-# LANGUAGE TypeInType #-} {-# LANGUAGE TypeOperators #-} +-- Combined GADT patterns are a bit of a trouble +{-# OPTIONS_GHC -fno-warn-overlapping-patterns #-} -- Constraints @b ~ BackendFamily t ds@ is not redundant; they allows to avoid AllowAmbigousTypes. {-# OPTIONS_GHC -fno-warn-redundant-constraints #-} {- | @@ -66,10 +68,6 @@ BackendFamily Double '[4] = DoubleX4 BackendFamily t ds = ArrayBase t ds --- | Promise that we are sure the backend is ArrayBase. -unsafeDefault :: forall t ds . Dict (BackendFamily t ds ~ ArrayBase t ds) -unsafeDefault = unsafeCoerce# (Dict @(ArrayBase t ds ~ ArrayBase t ds)) - -- | Singleton type used to determine the currently used DataFrame backend; -- establishes a good bijection @b <~> t ds@ (better than just using @BackendFamily@ everywhere). data BackendSing (t :: Type) (ds :: [Nat]) (backend :: Type) where @@ -99,33 +97,41 @@ #if defined(__HADDOCK__) || defined(__HADDOCK_VERSION__) instance KnownBackend t ds b where bSing = undefined #else -instance KnownBackend t '[] (ScalarBase t) where bSing = BSC -instance KnownBackend Float '[2] FloatX2 where bSing = BF2 -instance KnownBackend Float '[3] FloatX3 where bSing = BF3 -instance KnownBackend Float '[4] FloatX4 where bSing = BF4 -instance KnownBackend Double '[2] DoubleX2 where bSing = BD2 -instance KnownBackend Double '[3] DoubleX3 where bSing = BD3 -instance KnownBackend Double '[4] DoubleX4 where bSing = BD4 +instance KnownBackend t '[] (ScalarBase t) where bSing = BSC; {-# INLINE bSing #-} +instance KnownBackend Float '[2] FloatX2 where bSing = BF2; {-# INLINE bSing #-} +instance KnownBackend Float '[3] FloatX3 where bSing = BF3; {-# INLINE bSing #-} +instance KnownBackend Float '[4] FloatX4 where bSing = BF4; {-# INLINE bSing #-} +instance KnownBackend Double '[2] DoubleX2 where bSing = BD2; {-# INLINE bSing #-} +instance KnownBackend Double '[3] DoubleX3 where bSing = BD3; {-# INLINE bSing #-} +instance KnownBackend Double '[4] DoubleX4 where bSing = BD4; {-# INLINE bSing #-} instance PrimBytes t => KnownBackend t ds (ArrayBase t ds) where bSing = case unsafeDefault @t @ds of Dict -> BPB + {-# INLINE bSing #-} +instance {-# INCOHERENT #-} + (PrimBytes t, Dimensions ds, b ~ BackendFamily t ds) + => KnownBackend t ds b where + bSing = case (dims @ds, primTag @t undefined) of + (U, _) -> BSC + (D2 :* U, PTagFloat) -> BF2 + (D3 :* U, PTagFloat) -> BF3 + (D4 :* U, PTagFloat) -> BF4 + (D2 :* U, PTagDouble) -> BD2 + (D3 :* U, PTagDouble) -> BD3 + (D4 :* U, PTagDouble) -> BD4 + _ -> case unsafeDefault @t @ds of Dict -> BPB + {-# INLINE bSing #-} + +-- | Promise that we are sure the backend is ArrayBase. +unsafeDefault :: forall t ds . Dict (BackendFamily t ds ~ ArrayBase t ds) +unsafeDefault = unsafeCoerce# (Dict @(ArrayBase t ds ~ ArrayBase t ds)) #endif -- | Find an instance of `KnownBackend` class using `PrimBytes` and `Dimensions`. inferKnownBackend :: forall t ds b . (PrimBytes t, Dimensions ds, b ~ BackendFamily t ds) => Dict (KnownBackend t ds b) -inferKnownBackend = case (dims @ds, primTag @t undefined) of - (U, _) -> Dict - (d :* U, PTagFloat) - | Just Dict <- sameDim (D @2) d -> Dict - | Just Dict <- sameDim (D @3) d -> Dict - | Just Dict <- sameDim (D @4) d -> Dict - (d :* U, PTagDouble) - | Just Dict <- sameDim (D @2) d -> Dict - | Just Dict <- sameDim (D @3) d -> Dict - | Just Dict <- sameDim (D @4) d -> Dict - _ -> case unsafeDefault @t @ds of Dict -> Dict +inferKnownBackend = Dict {-# INLINE inferKnownBackend #-} diff --git a/src-base/Numeric/DataFrame/Internal/Backend/Family/ArrayBase.hs b/src-base/Numeric/DataFrame/Internal/Backend/Family/ArrayBase.hs --- a/src-base/Numeric/DataFrame/Internal/Backend/Family/ArrayBase.hs +++ b/src-base/Numeric/DataFrame/Internal/Backend/Family/ArrayBase.hs @@ -1,9 +1,9 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} {-# LANGUAGE ConstraintKinds #-} {-# LANGUAGE DataKinds #-} {-# LANGUAGE ExistentialQuantification #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE GADTs #-} {-# LANGUAGE MagicHash #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE PolyKinds #-} @@ -24,6 +24,7 @@ import Data.Int import Data.Word import GHC.Base hiding (foldr) +import GHC.Exts (TYPE) import Numeric.DataFrame.Internal.PrimArray import Numeric.Dimensions import Numeric.PrimBytes @@ -40,24 +41,23 @@ (# t -- Same value for each element; -- this is the cheapest way to initialize an array. -- It is also used for Num instances to avoid dependency on Dimensions. - | (# Int# -- Offset in Content array (measured in elements) + | (# CumulDims -- Steps array; [Word] + -- similar to steps in OpenCV::Mat, but counted in elements + -- rather than in bytes. + -- e.g. 0th element is the size of the content in elements + -- (data must be contiguous) + -- 1th element is the size of subspace data minus outer dim. + , Int# -- Offset in Content array (measured in elements) , ByteArray# -- Content; -- elements are stored row-by-row, plane-by-plane; -- (C-style, rather than Fortran-style); -- similar to C, OpenCV, etc. -- The data must be contiguous! -- (i.e. stride must be the same size as the element size). - , CumulDims -- Steps array; [Word] - -- similar to steps in OpenCV::Mat, but counted in elements - -- rather than in bytes. - -- e.g. 0th element is the size of the content in elements - -- (data must be contiguous) - -- 1th element is the size of subspace data minus outer dim. , Dict (PrimBytes t) #) #) - instance (PrimBytes t, Dimensions ds) => PrimBytes (ArrayBase t ds) where {-# SPECIALIZE instance Dimensions ds => PrimBytes (ArrayBase Float ds) #-} {-# SPECIALIZE instance Dimensions ds => PrimBytes (ArrayBase Double ds) #-} @@ -72,16 +72,13 @@ {-# SPECIALIZE instance Dimensions ds => PrimBytes (ArrayBase Word32 ds) #-} {-# SPECIALIZE instance Dimensions ds => PrimBytes (ArrayBase Word64 ds) #-} - getBytes (ArrayBase a ) = case a of - (# t | #) - | W# nw <- totalDim' @ds - , n <- word2Int# nw - , tbs <- byteSize t -> go tbs (tbs *# n) t - (# | (# _, arr, _, _ #) #) -> - -- very weird trick with touch# allows to workaround GHC bug - -- "internal error: ARR_WORDS object entered!" - -- TODO: report this - case runRW# (\s -> (# touch# arr s, arr #)) of (# _, ba #) -> ba + getBytes = withArrayContent' + (\t -> let tbs = byteSize t + in go tbs (tbs *# totalDim# @ds) t) + -- very weird trick with touch# allows to workaround GHC bug + -- "internal error: ARR_WORDS object entered!" + -- TODO: report this + (\_ _ arr -> case runRW# (\s -> (# touch# arr s, arr #)) of (# _, ba #) -> ba) where go :: Int# -> Int# -> t -> ByteArray# go tbs bsize t = case runRW# @@ -89,34 +86,29 @@ (# s1, mba #) -> unsafeFreezeByteArray# mba ( loop# 0# tbs bsize (\i -> writeBytes mba i t) s1 ) ) of (# _, ba #) -> ba - {-# NOINLINE go #-} {-# INLINE getBytes #-} - getBytesPinned (ArrayBase a) = case a of - (# t | #) - | W# nw <- totalDim' @ds - , n <- word2Int# nw - , tbs <- byteSize t -> go tbs (tbs *# n) t - (# | (# off, arr, _, _ #) #) -> - if isTrue# (isByteArrayPinned# arr) - then case runRW# (\s -> (# touch# arr s, arr #)) of (# _, ba #) -> ba - else case runRW# - ( \s0 -> case (# sizeofByteArray# arr - , byteAlign @t undefined - #) of - (# bsize, tba #) - | (# s1, mba #) <- newAlignedPinnedByteArray# bsize tba s0 - , s2 <- copyByteArray# arr off mba off bsize s1 - -> unsafeFreezeByteArray# mba s2 - ) of (# _, ba #) -> ba + getBytesPinned = withArrayContent' f + (const $ \off arr -> case runRW# (g off arr) of (# _, ba #) -> ba) where - go :: Int# -> Int# -> t -> ByteArray# - go tbs bsize t = case runRW# - ( \s0 -> case newAlignedPinnedByteArray# bsize (byteAlign t) s0 of - (# s1, mba #) -> unsafeFreezeByteArray# mba - ( loop# 0# tbs bsize (\i -> writeBytes mba i t) s1 ) - ) of (# _, ba #) -> ba - {-# NOINLINE go #-} + g :: Int# -> ByteArray# -> State# RealWorld -> (# State# RealWorld, ByteArray# #) + g off arr s0 + | isTrue# (isByteArrayPinned# arr) + = (# touch# arr s0, arr #) + | tba <- byteAlign @t undefined + , bsize <- sizeofByteArray# arr + , (# s1, mba #) <- newAlignedPinnedByteArray# bsize tba s0 + , s2 <- copyByteArray# arr off mba off bsize s1 + = unsafeFreezeByteArray# mba s2 + f :: t -> ByteArray# + f t + | tbs <- byteSize t + , bsize <- totalDim# @ds *# tbs + = case runRW# + ( \s0 -> case newAlignedPinnedByteArray# bsize (byteAlign t) s0 of + (# s1, mba #) -> unsafeFreezeByteArray# mba + ( loop# 0# tbs bsize (\i -> writeBytes mba i t) s1 ) + ) of (# _, ba #) -> ba {-# INLINE getBytesPinned #-} @@ -124,7 +116,7 @@ | tbs <- byteSize (undefined :: t) , (# offN, offRem #) <- quotRemInt# bOff tbs = case offRem of - 0# -> ArrayBase (# | (# offN, ba, steps, Dict #) #) + 0# -> fromElems' steps offN ba _ -> case cdTotalDim# steps of n -> go (tbs *# n) where steps = cumulDims $ dims @ds @@ -137,8 +129,7 @@ of (# s1, mba #) -> unsafeFreezeByteArray# mba (copyByteArray# ba bOff mba 0# bsize s1) - ) of (# _, r #) -> ArrayBase (# | (# 0#, r, steps, Dict #) #) - {-# NOINLINE go #-} + ) of (# _, r #) -> fromElems' steps 0# r {-# INLINE fromBytes #-} readBytes mba bOff s0 @@ -149,16 +140,17 @@ = case newByteArray# bsize s0 of (# s1, mba1 #) -> case unsafeFreezeByteArray# mba1 (copyMutableByteArray# mba bOff mba1 0# bsize s1) of - (# s2, ba #) -> (# s2, ArrayBase (# | (# 0#, ba, steps, Dict #) #) #) + (# s2, ba #) -> (# s2, fromElems' steps 0# ba #) {-# INLINE readBytes #-} - writeBytes mba bOff (ArrayBase c) - | tbs <- byteSize (undefined :: t) = case c of - (# t | #) | W# n <- totalDim' @ds -> - loop# bOff tbs (bOff +# word2Int# n *# tbs) (\i -> writeBytes mba i t) - (# | (# offContent, arrContent, steps, _ #) #) - | n <- cdTotalDim# steps-> - copyByteArray# arrContent (offContent *# tbs) mba bOff (n *# tbs) + writeBytes mba bOff + | tbs <- byteSize (undefined :: t) = withArrayContent' + (\t -> loop# bOff tbs + (bOff +# totalDim# @ds *# tbs) + (\i -> writeBytes mba i t)) + (\steps offContent arrContent -> + copyByteArray# arrContent (offContent *# tbs) + mba bOff (cdTotalDim# steps *# tbs)) {-# INLINE writeBytes #-} readAddr addr s0 @@ -169,29 +161,28 @@ = case newByteArray# bsize s0 of (# s1, mba1 #) -> case unsafeFreezeByteArray# mba1 (copyAddrToByteArray# addr mba1 0# bsize s1) of - (# s2, ba #) -> (# s2, ArrayBase (# | (# 0# , ba, steps, Dict #) #) #) + (# s2, ba #) -> (# s2, fromElems' steps 0# ba #) {-# INLINE readAddr #-} - writeAddr (ArrayBase c) addr - | tbs <- byteSize (undefined :: t) = case c of - (# t | #) | W# n <- totalDim' @ds -> - loop# 0# tbs (word2Int# n *# tbs) (\i -> writeAddr t (plusAddr# addr i)) - (# | (# offContent, arrContent, steps, _ #) #) - | n <- cdTotalDim# steps -> - copyByteArrayToAddr# arrContent (offContent *# tbs) addr (n *# tbs) + writeAddr a addr + | tbs <- byteSize (undefined :: t) = withArrayContent' + (\t -> loop# 0# tbs + (totalDim# @ds *# tbs) + (\i -> writeAddr t (plusAddr# addr i))) + (\steps offContent arrContent -> + copyByteArrayToAddr# arrContent (offContent *# tbs) + addr (cdTotalDim# steps *# tbs)) a {-# INLINE writeAddr #-} - byteSize _ = case totalDim' @ds of -- WARNING: slow! - W# n -> byteSize (undefined :: t) *# word2Int# n + byteSize _ = byteSize @t undefined *# totalDim# @ds -- WARNING: slow! {-# INLINE byteSize #-} - byteAlign _ = byteAlign (undefined :: t) + byteAlign _ = byteAlign @t undefined {-# INLINE byteAlign #-} - byteOffset (ArrayBase a) = case a of - (# _ | #) -> 0# - (# | (# off, _, _, _ #) #) -> off *# byteSize (undefined :: t) + byteOffset = withArrayContent' + (\_ -> 0#) (\_ off _ -> off *# byteSize @t undefined) {-# INLINE byteOffset #-} byteFieldOffset _ _ = negateInt# 1# @@ -199,8 +190,7 @@ indexArray ba off | steps <- cumulDims $ dims @ds - , n <- cdTotalDim# steps - = ArrayBase (# | (# off *# n, ba, steps, Dict #) #) + = fromElems' steps (off *# cdTotalDim# steps) ba {-# INLINE indexArray #-} @@ -222,17 +212,17 @@ (ArrayBase (# b | #)) = comb x (f a b) accumV2Idempotent x f comb - a@(ArrayBase (# | (# _, _, steps, Dict #) #)) + a@(ArrayBase (# | (# steps, _, _, Dict #) #)) b@(ArrayBase (# | _ #)) = foldr (comb . (\i -> f (ixOff i a) (ixOff i b))) x [0 .. fromIntegral (cdTotalDim steps) - 1] accumV2Idempotent x f comb (ArrayBase (# a | #)) - b@(ArrayBase (# | (# _, _, steps, Dict #) #)) + b@(ArrayBase (# | (# steps, _, _, Dict #) #)) = foldr (comb . (\i -> f a (ixOff i b))) x [0 .. fromIntegral (cdTotalDim steps) - 1] accumV2Idempotent x f comb - a@(ArrayBase (# | (# _, _, steps, Dict #) #)) + a@(ArrayBase (# | (# steps, _, _, Dict #) #)) (ArrayBase (# b | #)) = foldr (comb . (\i -> f (ixOff i a) b)) x [0 .. fromIntegral (cdTotalDim steps) - 1] @@ -242,7 +232,7 @@ mapV :: (t -> t) -> ArrayBase t ds -> ArrayBase t ds mapV f (ArrayBase (# t | #)) = ArrayBase (# f t | #) -mapV f x@(ArrayBase (# | (# offN, ba, steps, Dict #) #)) +mapV f x@(ArrayBase (# | (# steps, offN, ba, Dict #) #)) | tbs <- byteSize (undefEl x) , n <- cdTotalDim# steps = case runRW# @@ -251,7 +241,7 @@ ( loop1# n (\i -> writeArray mba i (f (indexArray ba (offN +# i)))) s1 ) - ) of (# _, r #) -> ArrayBase (# | (# 0#, r, steps, Dict #) #) + ) of (# _, r #) -> fromElems' steps 0# r {-# INLINE mapV #-} @@ -259,8 +249,8 @@ -> ArrayBase t ds -> ArrayBase t ds -> ArrayBase t ds zipV f (ArrayBase (# x | #)) b = mapV (f x) b zipV f a (ArrayBase (# y | #)) = mapV (`f` y) a -zipV f a@(ArrayBase (# | (# oa, ba, steps, Dict #) #)) - (ArrayBase (# | (# ob, bb, _, _ #) #)) +zipV f a@(ArrayBase (# | (# steps, oa, ba, Dict #) #)) + (ArrayBase (# | (# _, ob, bb, _ #) #)) | n <- cdTotalDim# steps = case runRW# ( \s0 -> case newByteArray# (byteSize (undefEl a) *# n) s0 of @@ -272,7 +262,7 @@ ) ) s1 ) - ) of (# _, r #) -> ArrayBase (# | (# 0#, r, steps, Dict #) #) + ) of (# _, r #) -> fromElems' steps 0# r {-# INLINE zipV #-} @@ -397,8 +387,8 @@ {-# SPECIALIZE instance Bounded (ArrayBase Word16 ds) #-} {-# SPECIALIZE instance Bounded (ArrayBase Word32 ds) #-} {-# SPECIALIZE instance Bounded (ArrayBase Word64 ds) #-} - maxBound = ArrayBase (# maxBound | #) - minBound = ArrayBase (# minBound | #) + maxBound = broadcast' maxBound + minBound = broadcast' minBound instance Num t => Num (ArrayBase t ds) where {-# SPECIALIZE instance Num (ArrayBase Float ds) #-} @@ -425,7 +415,7 @@ {-# INLINE abs #-} signum = mapV signum {-# INLINE signum #-} - fromInteger i = ArrayBase (# fromInteger i | #) + fromInteger = broadcast' . fromInteger {-# INLINE fromInteger #-} instance Fractional t => Fractional (ArrayBase t ds) where @@ -435,14 +425,14 @@ {-# INLINE (/) #-} recip = mapV recip {-# INLINE recip #-} - fromRational r = ArrayBase (# fromRational r | #) + fromRational = broadcast' . fromRational {-# INLINE fromRational #-} instance Floating t => Floating (ArrayBase t ds) where {-# SPECIALIZE instance Floating (ArrayBase Float ds) #-} {-# SPECIALIZE instance Floating (ArrayBase Double ds) #-} - pi = ArrayBase (# pi | #) + pi = broadcast' pi {-# INLINE pi #-} exp = mapV exp {-# INLINE exp #-} @@ -480,26 +470,9 @@ {-# INLINE atanh #-} instance PrimBytes t => PrimArray t (ArrayBase t ds) where - {-# SPECIALIZE instance PrimArray Float (ArrayBase Float ds) #-} - {-# SPECIALIZE instance PrimArray Double (ArrayBase Double ds) #-} - {-# SPECIALIZE instance PrimArray Int (ArrayBase Int ds) #-} - {-# SPECIALIZE instance PrimArray Word (ArrayBase Word ds) #-} - {-# SPECIALIZE instance PrimArray Int8 (ArrayBase Int8 ds) #-} - {-# SPECIALIZE instance PrimArray Int16 (ArrayBase Int16 ds) #-} - {-# SPECIALIZE instance PrimArray Int32 (ArrayBase Int32 ds) #-} - {-# SPECIALIZE instance PrimArray Int64 (ArrayBase Int64 ds) #-} - {-# SPECIALIZE instance PrimArray Word8 (ArrayBase Word8 ds) #-} - {-# SPECIALIZE instance PrimArray Word16 (ArrayBase Word16 ds) #-} - {-# SPECIALIZE instance PrimArray Word32 (ArrayBase Word32 ds) #-} - {-# SPECIALIZE instance PrimArray Word64 (ArrayBase Word64 ds) #-} - broadcast t = ArrayBase (# t | #) - {-# INLINE broadcast #-} - - ix# i (ArrayBase a) = case a of - (# t | #) -> t - (# | (# off, arr, _, _ #) #) -> indexArray arr (off +# i) - {-# INLINE ix# #-} + broadcast# t = ArrayBase (# t | #) + {-# INLINE broadcast# #-} gen# steps f z0 = go (byteSize @t undefined *# n) where @@ -510,57 +483,63 @@ (# s2, z1 #) -> case unsafeFreezeByteArray# mba s2 of (# s3, ba #) -> (# s3, (# z1, ba #) #) ) of (# _, (# z1, ba #) #) - -> (# z1, ArrayBase (# | (# 0# , ba, steps, Dict #) #) #) - {-# NOINLINE go #-} + -> (# z1, fromElems' steps 0# ba #) loop0 mba i z s | isTrue# (i ==# n) = (# s, z #) | otherwise = case f z of (# z', x #) -> loop0 mba (i +# 1#) z' (writeArray mba i x s) {-# INLINE gen# #-} - upd# steps i x (ArrayBase (# a | #)) = go (byteSize x) - where - n = cdTotalDim# steps - go tbs = case runRW# - ( \s0 -> case newByteArray# (tbs *# n) s0 of - (# s1, mba #) -> unsafeFreezeByteArray# mba - (writeArray mba i x - (loop1# n (\j -> writeArray mba j a) s1) - ) - ) of (# _, r #) -> ArrayBase (# | (# 0# , r, steps, Dict #) #) - {-# NOINLINE go #-} - upd# _ i x (ArrayBase (# | (# offN, ba, steps, Dict #) #)) - = go (byteSize x) + upd# steps i x = withArrayContent' f g where n = cdTotalDim# steps - go tbs = case runRW# - ( \s0 -> case newByteArray# (tbs *# n) s0 of - (# s1, mba #) -> unsafeFreezeByteArray# mba - (writeArray mba i x - (copyByteArray# ba (offN *# tbs) mba 0# (tbs *# n) s1) - ) - ) of (# _, r #) -> ArrayBase (# | (# 0#, r, steps, Dict #) #) - {-# NOINLINE go #-} + f a = go (byteSize x) + where + go tbs = case runRW# + ( \s0 -> case newByteArray# (tbs *# n) s0 of + (# s1, mba #) -> unsafeFreezeByteArray# mba + (writeArray mba i x + (loop1# n (\j -> writeArray mba j a) s1) + ) + ) of (# _, r #) -> fromElems' steps 0# r + g _ offN ba = go (byteSize x) + where + go tbs = case runRW# + ( \s0 -> case newByteArray# (tbs *# n) s0 of + (# s1, mba #) -> unsafeFreezeByteArray# mba + (writeArray mba i x + (copyByteArray# ba (offN *# tbs) mba 0# (tbs *# n) s1) + ) + ) of (# _, r #) -> fromElems' steps 0# r {-# INLINE upd# #-} - arrayContent# (ArrayBase a) = case a of - (# x | #) -> (# x | #) - (# | (# off, arr, cd, _ #) #) -> (# | (# cd, off, arr #) #) - {-# INLINE arrayContent# #-} + withArrayContent# = withArrayContent' + {-# INLINE withArrayContent# #-} - offsetElems (ArrayBase a) = case a of - (# _ | #) -> 0# - (# | (# off, _, _, _ #) #) -> off - {-# INLINE offsetElems #-} + fromElems# = fromElems' + {-# INLINE fromElems# #-} - uniqueOrCumulDims (ArrayBase a) = case a of - (# t | #) -> Left t - (# | (# _, _, steps, _ #) #) -> Right steps - {-# INLINE uniqueOrCumulDims #-} +withArrayContent' :: + forall (t :: Type) (ds :: [Nat]) (rep :: RuntimeRep) (r :: TYPE rep) + . (t -> r) + -> (CumulDims -> Int# -> ByteArray# -> r) + -> ArrayBase t ds -> r +withArrayContent' f _ (ArrayBase (# e | #)) = f e +withArrayContent' _ g (ArrayBase (# | (# steps, off, ba, _ #) #)) = g steps off ba +{-# INLINE withArrayContent' #-} - fromElems steps off ba = ArrayBase (# | (# off, ba, steps, Dict #) #) - {-# INLINE fromElems #-} +fromElems' :: forall (t :: Type) (ds :: [Nat]) + . PrimBytes t => CumulDims -> Int# -> ByteArray# -> ArrayBase t ds +fromElems' steps off ba = ArrayBase (# | (# steps, off, ba, Dict #) #) +{-# INLINE fromElems' #-} +broadcast' :: t -> ArrayBase t ds +broadcast' t = ArrayBase (# t | #) +{-# INLINE broadcast' #-} + +totalDim# :: forall (ds :: [Nat]) . Dimensions ds => Int# +totalDim# = case totalDim' @ds of W# n -> word2Int# n +{-# INLINE totalDim# #-} loop# :: Int# -- ^ initial value -> Int# -- ^ step diff --git a/src-base/Numeric/DataFrame/Internal/Backend/Family/DoubleX2.hs b/src-base/Numeric/DataFrame/Internal/Backend/Family/DoubleX2.hs --- a/src-base/Numeric/DataFrame/Internal/Backend/Family/DoubleX2.hs +++ b/src-base/Numeric/DataFrame/Internal/Backend/Family/DoubleX2.hs @@ -4,12 +4,10 @@ {-# LANGUAGE MagicHash #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE TypeApplications #-} -{-# LANGUAGE UnboxedSums #-} {-# LANGUAGE UnboxedTuples #-} {-# LANGUAGE UndecidableInstances #-} module Numeric.DataFrame.Internal.Backend.Family.DoubleX2 (DoubleX2 (..)) where - import GHC.Base import Numeric.DataFrame.Internal.PrimArray import Numeric.PrimBytes @@ -368,8 +366,8 @@ instance PrimArray Double DoubleX2 where - broadcast (D# x) = DoubleX2# x x - {-# INLINE broadcast #-} + broadcast# (D# x) = DoubleX2# x x + {-# INLINE broadcast# #-} ix# 0# (DoubleX2# a1 _) = D# a1 ix# 1# (DoubleX2# _ a2) = D# a2 @@ -386,8 +384,8 @@ upd# _ _ _ x = x {-# INLINE upd# #-} - arrayContent# x = (# | (# CumulDims [ELEM_N, 1], 0#, getBytes x #) #) - {-# INLINE arrayContent# #-} + withArrayContent# _ g x = g (CumulDims [ELEM_N, 1]) 0# (getBytes x) + {-# INLINE withArrayContent# #-} offsetElems _ = 0# {-# INLINE offsetElems #-} @@ -395,7 +393,7 @@ uniqueOrCumulDims _ = Right (CumulDims [ELEM_N, 1]) {-# INLINE uniqueOrCumulDims #-} - fromElems _ off ba = DoubleX2# + fromElems# _ off ba = DoubleX2# (indexDoubleArray# ba off) (indexDoubleArray# ba (off +# 1#)) - {-# INLINE fromElems #-} + {-# INLINE fromElems# #-} diff --git a/src-base/Numeric/DataFrame/Internal/Backend/Family/DoubleX3.hs b/src-base/Numeric/DataFrame/Internal/Backend/Family/DoubleX3.hs --- a/src-base/Numeric/DataFrame/Internal/Backend/Family/DoubleX3.hs +++ b/src-base/Numeric/DataFrame/Internal/Backend/Family/DoubleX3.hs @@ -4,7 +4,6 @@ {-# LANGUAGE MagicHash #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE TypeApplications #-} -{-# LANGUAGE UnboxedSums #-} {-# LANGUAGE UnboxedTuples #-} {-# LANGUAGE UndecidableInstances #-} module Numeric.DataFrame.Internal.Backend.Family.DoubleX3 (DoubleX3 (..)) where @@ -399,8 +398,8 @@ instance PrimArray Double DoubleX3 where - broadcast (D# x) = DoubleX3# x x x - {-# INLINE broadcast #-} + broadcast# (D# x) = DoubleX3# x x x + {-# INLINE broadcast# #-} ix# 0# (DoubleX3# a1 _ _) = D# a1 ix# 1# (DoubleX3# _ a2 _) = D# a2 @@ -420,8 +419,8 @@ upd# _ _ _ x = x {-# INLINE upd# #-} - arrayContent# x = (# | (# CumulDims [ELEM_N, 1], 0#, getBytes x #) #) - {-# INLINE arrayContent# #-} + withArrayContent# _ g x = g (CumulDims [ELEM_N, 1]) 0# (getBytes x) + {-# INLINE withArrayContent# #-} offsetElems _ = 0# {-# INLINE offsetElems #-} @@ -429,8 +428,8 @@ uniqueOrCumulDims _ = Right (CumulDims [ELEM_N, 1]) {-# INLINE uniqueOrCumulDims #-} - fromElems _ off ba = DoubleX3# + fromElems# _ off ba = DoubleX3# (indexDoubleArray# ba off) (indexDoubleArray# ba (off +# 1#)) (indexDoubleArray# ba (off +# 2#)) - {-# INLINE fromElems #-} + {-# INLINE fromElems# #-} diff --git a/src-base/Numeric/DataFrame/Internal/Backend/Family/DoubleX4.hs b/src-base/Numeric/DataFrame/Internal/Backend/Family/DoubleX4.hs --- a/src-base/Numeric/DataFrame/Internal/Backend/Family/DoubleX4.hs +++ b/src-base/Numeric/DataFrame/Internal/Backend/Family/DoubleX4.hs @@ -4,7 +4,6 @@ {-# LANGUAGE MagicHash #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE TypeApplications #-} -{-# LANGUAGE UnboxedSums #-} {-# LANGUAGE UnboxedTuples #-} {-# LANGUAGE UndecidableInstances #-} module Numeric.DataFrame.Internal.Backend.Family.DoubleX4 (DoubleX4 (..)) where @@ -426,8 +425,8 @@ instance PrimArray Double DoubleX4 where - broadcast (D# x) = DoubleX4# x x x x - {-# INLINE broadcast #-} + broadcast# (D# x) = DoubleX4# x x x x + {-# INLINE broadcast# #-} ix# 0# (DoubleX4# a1 _ _ _) = D# a1 ix# 1# (DoubleX4# _ a2 _ _) = D# a2 @@ -450,8 +449,8 @@ upd# _ _ _ x = x {-# INLINE upd# #-} - arrayContent# x = (# | (# CumulDims [ELEM_N, 1], 0#, getBytes x #) #) - {-# INLINE arrayContent# #-} + withArrayContent# _ g x = g (CumulDims [ELEM_N, 1]) 0# (getBytes x) + {-# INLINE withArrayContent# #-} offsetElems _ = 0# {-# INLINE offsetElems #-} @@ -459,9 +458,9 @@ uniqueOrCumulDims _ = Right (CumulDims [ELEM_N, 1]) {-# INLINE uniqueOrCumulDims #-} - fromElems _ off ba = DoubleX4# + fromElems# _ off ba = DoubleX4# (indexDoubleArray# ba off) (indexDoubleArray# ba (off +# 1#)) (indexDoubleArray# ba (off +# 2#)) (indexDoubleArray# ba (off +# 3#)) - {-# INLINE fromElems #-} + {-# INLINE fromElems# #-} diff --git a/src-base/Numeric/DataFrame/Internal/Backend/Family/FloatX2.hs b/src-base/Numeric/DataFrame/Internal/Backend/Family/FloatX2.hs --- a/src-base/Numeric/DataFrame/Internal/Backend/Family/FloatX2.hs +++ b/src-base/Numeric/DataFrame/Internal/Backend/Family/FloatX2.hs @@ -4,7 +4,6 @@ {-# LANGUAGE MagicHash #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE TypeApplications #-} -{-# LANGUAGE UnboxedSums #-} {-# LANGUAGE UnboxedTuples #-} {-# LANGUAGE UndecidableInstances #-} module Numeric.DataFrame.Internal.Backend.Family.FloatX2 (FloatX2 (..)) where @@ -368,8 +367,8 @@ instance PrimArray Float FloatX2 where - broadcast (F# x) = FloatX2# x x - {-# INLINE broadcast #-} + broadcast# (F# x) = FloatX2# x x + {-# INLINE broadcast# #-} ix# 0# (FloatX2# a1 _) = F# a1 ix# 1# (FloatX2# _ a2) = F# a2 @@ -386,8 +385,8 @@ upd# _ _ _ x = x {-# INLINE upd# #-} - arrayContent# x = (# | (# CumulDims [ELEM_N, 1], 0#, getBytes x #) #) - {-# INLINE arrayContent# #-} + withArrayContent# _ g x = g (CumulDims [ELEM_N, 1]) 0# (getBytes x) + {-# INLINE withArrayContent# #-} offsetElems _ = 0# {-# INLINE offsetElems #-} @@ -395,7 +394,7 @@ uniqueOrCumulDims _ = Right (CumulDims [ELEM_N, 1]) {-# INLINE uniqueOrCumulDims #-} - fromElems _ off ba = FloatX2# + fromElems# _ off ba = FloatX2# (indexFloatArray# ba off) (indexFloatArray# ba (off +# 1#)) - {-# INLINE fromElems #-} + {-# INLINE fromElems# #-} diff --git a/src-base/Numeric/DataFrame/Internal/Backend/Family/FloatX3.hs b/src-base/Numeric/DataFrame/Internal/Backend/Family/FloatX3.hs --- a/src-base/Numeric/DataFrame/Internal/Backend/Family/FloatX3.hs +++ b/src-base/Numeric/DataFrame/Internal/Backend/Family/FloatX3.hs @@ -4,7 +4,6 @@ {-# LANGUAGE MagicHash #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE TypeApplications #-} -{-# LANGUAGE UnboxedSums #-} {-# LANGUAGE UnboxedTuples #-} {-# LANGUAGE UndecidableInstances #-} module Numeric.DataFrame.Internal.Backend.Family.FloatX3 (FloatX3 (..)) where @@ -400,8 +399,8 @@ instance PrimArray Float FloatX3 where - broadcast (F# x) = FloatX3# x x x - {-# INLINE broadcast #-} + broadcast# (F# x) = FloatX3# x x x + {-# INLINE broadcast# #-} ix# 0# (FloatX3# a1 _ _) = F# a1 ix# 1# (FloatX3# _ a2 _) = F# a2 @@ -421,8 +420,8 @@ upd# _ _ _ x = x {-# INLINE upd# #-} - arrayContent# x = (# | (# CumulDims [ELEM_N, 1], 0#, getBytes x #) #) - {-# INLINE arrayContent# #-} + withArrayContent# _ g x = g (CumulDims [ELEM_N, 1]) 0# (getBytes x) + {-# INLINE withArrayContent# #-} offsetElems _ = 0# {-# INLINE offsetElems #-} @@ -430,8 +429,8 @@ uniqueOrCumulDims _ = Right (CumulDims [ELEM_N, 1]) {-# INLINE uniqueOrCumulDims #-} - fromElems _ off ba = FloatX3# + fromElems# _ off ba = FloatX3# (indexFloatArray# ba off) (indexFloatArray# ba (off +# 1#)) (indexFloatArray# ba (off +# 2#)) - {-# INLINE fromElems #-} + {-# INLINE fromElems# #-} diff --git a/src-base/Numeric/DataFrame/Internal/Backend/Family/FloatX4.hs b/src-base/Numeric/DataFrame/Internal/Backend/Family/FloatX4.hs --- a/src-base/Numeric/DataFrame/Internal/Backend/Family/FloatX4.hs +++ b/src-base/Numeric/DataFrame/Internal/Backend/Family/FloatX4.hs @@ -5,7 +5,6 @@ {-# LANGUAGE MagicHash #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE TypeApplications #-} -{-# LANGUAGE UnboxedSums #-} {-# LANGUAGE UnboxedTuples #-} {-# LANGUAGE UndecidableInstances #-} {-# OPTIONS_GHC -fno-warn-orphans #-} @@ -20,6 +19,7 @@ import Numeric.ProductOrd import qualified Numeric.ProductOrd.NonTransitive as NonTransitive import qualified Numeric.ProductOrd.Partial as Partial +import Unsafe.Coerce (unsafeCoerce) data FloatX4 = FloatX4# Float# Float# Float# Float# @@ -432,8 +432,8 @@ instance PrimArray Float FloatX4 where - broadcast (F# x) = FloatX4# x x x x - {-# INLINE broadcast #-} + broadcast# (F# x) = FloatX4# x x x x + {-# INLINE broadcast# #-} ix# 0# (FloatX4# a1 _ _ _) = F# a1 ix# 1# (FloatX4# _ a2 _ _) = F# a2 @@ -456,8 +456,8 @@ upd# _ _ _ x = x {-# INLINE upd# #-} - arrayContent# x = (# | (# CumulDims [ELEM_N, 1], 0#, getBytes x #) #) - {-# INLINE arrayContent# #-} + withArrayContent# _ g x = g (CumulDims [ELEM_N, 1]) 0# (getBytes x) + {-# INLINE withArrayContent# #-} offsetElems _ = 0# {-# INLINE offsetElems #-} @@ -465,13 +465,12 @@ uniqueOrCumulDims _ = Right (CumulDims [ELEM_N, 1]) {-# INLINE uniqueOrCumulDims #-} - fromElems _ off ba = FloatX4# + fromElems# _ off ba = FloatX4# (indexFloatArray# ba off) (indexFloatArray# ba (off +# 1#)) (indexFloatArray# ba (off +# 2#)) (indexFloatArray# ba (off +# 3#)) - {-# INLINE fromElems #-} - + {-# INLINE fromElems# #-} -------------------------------------------------------------------------------- -- Rewrite rules to improve efficiency of algorithms @@ -485,14 +484,13 @@ -------------------------------------------------------------------------------- getIdxOffset :: Idxs '[4] -> Int# -getIdxOffset is = case unsafeCoerce# is of - [W# i] -> word2Int# i - _ -> 0# +getIdxOffset is = case unsafeCoerce is of + ~[w] -> case w of W# i -> word2Int# i {-# INLINE getIdxOffset #-} {-# RULES -"index/FloatX4" forall i . index @Float @'[4] i - = unsafeCoerce# (ix# @Float @FloatX4 (getIdxOffset i)) +"index/FloatX4" forall i . index @Float @'[4] @'[] @'[4] i + = unsafeCoerce (ix# @Float @FloatX4 (getIdxOffset i)) #-} diff --git a/src-base/Numeric/DataFrame/Internal/Backend/Family/ScalarBase.hs b/src-base/Numeric/DataFrame/Internal/Backend/Family/ScalarBase.hs --- a/src-base/Numeric/DataFrame/Internal/Backend/Family/ScalarBase.hs +++ b/src-base/Numeric/DataFrame/Internal/Backend/Family/ScalarBase.hs @@ -9,7 +9,8 @@ module Numeric.DataFrame.Internal.Backend.Family.ScalarBase (ScalarBase (..)) where -import GHC.Base +import Data.Coerce +import Numeric.Basics import Numeric.DataFrame.Internal.PrimArray import Numeric.PrimBytes import Numeric.ProductOrd @@ -18,7 +19,7 @@ -- | Specialize ScalarBase type without any arrays newtype ScalarBase t = ScalarBase { _unScalarBase :: t } - deriving ( Enum, Eq, Bounded, Integral + deriving ( Enum, Eq, Bounded, Integral, Epsilon, RealExtras, RealFloatExtras , Num, Fractional, Floating, Ord, Real, RealFrac, RealFloat , PrimBytes) @@ -28,23 +29,23 @@ deriving instance Ord t => Ord (Partial.ProductOrd (ScalarBase t)) instance PrimBytes t => PrimArray t (ScalarBase t) where - broadcast = unsafeCoerce# - {-# INLINE broadcast #-} - ix# _ = unsafeCoerce# + broadcast# = coerce + {-# INLINE broadcast# #-} + ix# _ = coerce {-# INLINE ix# #-} - gen# _ = unsafeCoerce# + gen# _ = coerce {-# INLINE gen# #-} upd# _ 0# = const . ScalarBase upd# _ _ = const id {-# INLINE upd# #-} - arrayContent# x = (# _unScalarBase x | #) - {-# INLINE arrayContent# #-} + withArrayContent# f _ x = f (_unScalarBase x) + {-# INLINE withArrayContent# #-} offsetElems _ = 0# {-# INLINE offsetElems #-} uniqueOrCumulDims = Left . _unScalarBase {-# INLINE uniqueOrCumulDims #-} - fromElems _ off ba = indexArray ba off - {-# INLINE fromElems #-} + fromElems# _ off ba = indexArray ba off + {-# INLINE fromElems# #-} _suppressHlintUnboxedTuplesWarning :: () -> (# (), () #) _suppressHlintUnboxedTuplesWarning = undefined diff --git a/src/Numeric/Basics.hs b/src/Numeric/Basics.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/Basics.hs @@ -0,0 +1,238 @@ +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE ViewPatterns #-} +module Numeric.Basics + ( -- * Constants + pattern M_E, pattern M_LOG2E, pattern M_LOG10E, pattern M_LN2, pattern M_LN10 + , pattern M_SQRT2, pattern M_SQRT1_2 + , pattern M_PI, pattern M_PI_2, pattern M_PI_3, pattern M_PI_4 + , pattern M_1_PI, pattern M_2_PI, pattern M_2_SQRTPI + , Epsilon (..), pattern M_EPS + -- * Functions + , RealExtras (..), RealFloatExtras (..) + , negateUnless + -- * Exceptions + , IterativeMethod, TooManyIterations, tooManyIterations + ) where + +import Data.Int +import Data.Word +import GHC.Exception +import GHC.Stack +import Numeric.PrimBytes + +-- | \( e \) +pattern M_E :: (Eq a, Fractional a) => Fractional a => a +pattern M_E = 2.718281828459045235360287471352662498 + +-- | \( \log_2 e \) +pattern M_LOG2E :: (Eq a, Fractional a) => Fractional a => a +pattern M_LOG2E = 1.442695040888963407359924681001892137 + +-- | \( \log_{10} e \) +pattern M_LOG10E :: (Eq a, Fractional a) => Fractional a => a +pattern M_LOG10E = 0.434294481903251827651128918916605082 + +-- | \( \log_e 2 \) +pattern M_LN2 :: (Eq a, Fractional a) => Fractional a => a +pattern M_LN2 = 0.693147180559945309417232121458176568 + +-- | \( \log_e 10 \) +pattern M_LN10 :: (Eq a, Fractional a) => Fractional a => a +pattern M_LN10 = 2.302585092994045684017991454684364208 + +-- | \( \sqrt{2} \) +pattern M_SQRT2 :: (Eq a, Fractional a) => Fractional a => a +pattern M_SQRT2 = 1.414213562373095048801688724209698079 + +-- | \( \frac{1}{\sqrt{2}} \) +pattern M_SQRT1_2 :: (Eq a, Fractional a) => Fractional a => a +pattern M_SQRT1_2 = 0.707106781186547524400844362104849039 + +-- | \( \pi \) +pattern M_PI :: (Eq a, Fractional a) => Fractional a => a +pattern M_PI = 3.1415926535897932384626433832795028841971 +-- | \( \frac{\pi}{2} \) +pattern M_PI_2 :: (Eq a, Fractional a) => Fractional a => a +pattern M_PI_2 = 1.570796326794896619231321691639751442 + +-- | \( \frac{\pi}{3} \) +pattern M_PI_3 :: (Eq a, Fractional a) => Fractional a => a +pattern M_PI_3 = 1.047197551196597746154214461093167628 + +-- | \( \frac{\pi}{4} \) +pattern M_PI_4 :: (Eq a, Fractional a) => Fractional a => a +pattern M_PI_4 = 0.785398163397448309615660845819875721 + +-- | \( \frac{1}{\pi} \) +pattern M_1_PI :: (Eq a, Fractional a) => Fractional a => a +pattern M_1_PI = 0.318309886183790671537767526745028724 + +-- | \( \frac{2}{\pi} \) +pattern M_2_PI :: (Eq a, Fractional a) => Fractional a => a +pattern M_2_PI = 0.636619772367581343075535053490057448 + +-- | \( \frac{2}{\sqrt{\pi}} \) +pattern M_2_SQRTPI :: (Eq a, Fractional a) => Fractional a => a +pattern M_2_SQRTPI = 1.128379167095512573896158903121545172 + +-- | Define a meaningful precision for complex floating-point operations. +class (Eq a, Floating a) => Epsilon a where + -- | A small positive number that depends on the type of @a@. + -- Compare your values to @epsilon@ to tell if they are near zero. + epsilon :: a + +instance Epsilon Double where + epsilon = 1e-12 + +instance Epsilon Float where + epsilon = 1e-6 + +-- | A small positive number that depends on the type of @a@. +pattern M_EPS :: Epsilon a => a +pattern M_EPS <- ((epsilon ==) -> True) + where + M_EPS = epsilon + +-- | Negate if @False@. +-- This is a useful alternative to `signum`, +-- when @signum 0 == 0@ causing some troubles. +negateUnless :: Num t => Bool -> t -> t +negateUnless True = id +negateUnless False = negate +{-# INLINE negateUnless #-} + +-- | Extra functions for `Real` types. +class (Real a, PrimBytes a) => RealExtras a where + -- | @copysign x y@ returns a value with the magnitude of x and the sign of y. + -- + -- NB: in future, this function is to reimplemented using primops + -- and should behave the same way as its C analogue. + copysign :: a -> a -> a + copysign x y + | (x >= 0) == (y >= 0) = x + | otherwise = negate x + {-# INLINE copysign #-} + +instance RealExtras Int +instance RealExtras Int8 +instance RealExtras Int16 +instance RealExtras Int32 +instance RealExtras Int64 +instance RealExtras Word +instance RealExtras Word8 +instance RealExtras Word16 +instance RealExtras Word32 +instance RealExtras Word64 + +instance RealExtras Float where + copysign = c'copysignf + {-# INLINE copysign #-} + +instance RealExtras Double where + copysign = c'copysignd + {-# INLINE copysign #-} + +-- | Extra functions for `RealFrac` types. +class (Epsilon a, RealExtras a, RealFloat a) => RealFloatExtras a where + -- | \( \sqrt{ x^2 + y^2 } \). + -- + -- NB: in future, this function is to reimplemented using primops + -- and should behave the same way as its C analogue. + hypot :: a -> a -> a + hypot x y = scaleFloat ea (sqrt (an*an + bn*bn)) + where + x' = abs x + y' = abs y + (a,b) = if x' >= y' then (x', y') else (y', x') + (_, ea) = decodeFloat a + an = scaleFloat (negate ea) a + bn = scaleFloat (negate ea) b + {-# INLINE hypot #-} + -- | Maximum finite number representable by this FP type. + maxFinite :: a + +instance RealFloatExtras Float where + hypot = c'hypotf + {-# INLINE hypot #-} + maxFinite = 3.40282347e+38 + {-# INLINE maxFinite #-} + + +instance RealFloatExtras Double where + hypot = c'hypotd + {-# INLINE hypot #-} + maxFinite = 1.7976931348623157e+308 + {-# INLINE maxFinite #-} + +foreign import ccall unsafe "hypot" c'hypotd :: Double -> Double -> Double +foreign import ccall unsafe "hypotf" c'hypotf :: Float -> Float -> Float +foreign import ccall unsafe "copysign" c'copysignd :: Double -> Double -> Double +foreign import ccall unsafe "copysignf" c'copysignf :: Float -> Float -> Float + +{- | +This is just an alias for the `HasCallStack` constraint. +It servers two goals: + + 1. Document functions that use iterative methods and can fail + (supposedly in extremely rare cases). + 2. Provide a call stack in case of failure. + +Use `tooManyIterations` function where necessary if you implement an iterative +algorithm and annotate your implementation function with `IterativeMethod` constraint. + + -} +type IterativeMethod = HasCallStack + +-- | Throw a `TooManyIterations` exception. +tooManyIterations :: + IterativeMethod + => String -- ^ Label (e.g. function name) + -> a +tooManyIterations s + = throw TooManyIterations + { tmiUserInfo = s + , tmiCallStack = callStack + } + +{- | +Typically, this exception can occur in an iterative algorithm; +when the number of iterations exceeds a pre-defined limit, but the stop condition +is not fulfilled. + +Use the `IterativeMethod` constraint to annotate functions that throw this exception. + -} +data TooManyIterations + = TooManyIterations + { tmiUserInfo :: String + -- ^ Short description of the error, e.g. a function name. + , tmiCallStack :: CallStack + -- ^ Function call stack. + -- Note, this field is ignored in the `Eq` and `Ord` instances. + } + +-- | Note, this instance ignores `oodCallStack` +instance Eq TooManyIterations where + (==) a b = tmiUserInfo a == tmiUserInfo b + +-- | Note, this instance ignores `oodCallStack` +instance Ord TooManyIterations where + compare a b = compare (tmiUserInfo a) (tmiUserInfo b) + +instance Show TooManyIterations where + showsPrec p e = addLoc errStr + where + addLoc s + = let someE = errorCallWithCallStackException s (tmiCallStack e) + errc :: ErrorCall + errc = case fromException someE of + Nothing -> ErrorCall s + Just ec -> ec + in showsPrec p errc + errStr = unlines + [ "An iterative method has exceeded the limit for the iteration number." + , "User info: " ++ tmiUserInfo e + ] + +instance Exception TooManyIterations diff --git a/src/Numeric/DataFrame.hs b/src/Numeric/DataFrame.hs --- a/src/Numeric/DataFrame.hs +++ b/src/Numeric/DataFrame.hs @@ -7,6 +7,8 @@ -- * Functionality , module Numeric.DataFrame.SubSpace , module Numeric.DataFrame.Contraction + , module Numeric.Basics + , module Numeric.Subroutine.Sort ) where import Numeric.DataFrame.Internal.Backend () @@ -18,3 +20,7 @@ import Numeric.Matrix import Numeric.Scalar import Numeric.Vector + +import Numeric.Basics + +import Numeric.Subroutine.Sort (sort, sortBy) diff --git a/src/Numeric/DataFrame/Contraction.hs b/src/Numeric/DataFrame/Contraction.hs --- a/src/Numeric/DataFrame/Contraction.hs +++ b/src/Numeric/DataFrame/Contraction.hs @@ -1,16 +1,17 @@ -{-# LANGUAGE DataKinds #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE FunctionalDependencies #-} -{-# LANGUAGE InstanceSigs #-} -{-# LANGUAGE MagicHash #-} -{-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE TypeOperators #-} -{-# LANGUAGE UnboxedTuples #-} -{-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE FunctionalDependencies #-} +{-# LANGUAGE InstanceSigs #-} +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UnboxedTuples #-} +{-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE UndecidableSuperClasses #-} ----------------------------------------------------------------------------- -- | -- Module : Numeric.DataFrame.Contraction diff --git a/src/Numeric/DataFrame/IO.hs b/src/Numeric/DataFrame/IO.hs --- a/src/Numeric/DataFrame/IO.hs +++ b/src/Numeric/DataFrame/IO.hs @@ -2,14 +2,18 @@ {-# LANGUAGE ExistentialQuantification #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE InstanceSigs #-} {-# LANGUAGE MagicHash #-} {-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE PatternSynonyms #-} {-# LANGUAGE PolyKinds #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} {-# LANGUAGE TypeFamilies #-} {-# LANGUAGE TypeInType #-} {-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE ViewPatterns #-} ----------------------------------------------------------------------------- -- | -- Module : Numeric.DataFrame.IO @@ -22,208 +26,348 @@ module Numeric.DataFrame.IO ( IODataFrame (XIOFrame), SomeIODataFrame (..) - , newDataFrame, newPinnedDataFrame + , castDataFrame + , newDataFrame, newPinnedDataFrame, oneMoreDataFrame + , subDataFrameView, subDataFrameView' , copyDataFrame, copyMutableDataFrame , copyDataFrame', copyMutableDataFrame' + , copyDataFrameOff, copyMutableDataFrameOff , freezeDataFrame, unsafeFreezeDataFrame - , thawDataFrame, thawPinDataFrame, unsafeThawDataFrame + , thawDataFrame, thawPinDataFrame, unsafeThawDataFrame, withThawDataFrame , writeDataFrame, writeDataFrameOff , readDataFrame, readDataFrameOff - , withDataFramePtr, isDataFramePinned + , isDataFramePinned, getDataFrameSteps + , withDataFramePtr ) where -import GHC.Base -import GHC.IO (IO (..)) -import GHC.Ptr (Ptr (..)) +import GHC.IO (IO (..)) +import GHC.Ptr (Ptr (..)) +import Control.Monad.ST (RealWorld) +import Data.Coerce +import Data.Kind import Numeric.DataFrame.Internal.Mutable import Numeric.DataFrame.Internal.PrimArray import Numeric.DataFrame.Type import Numeric.Dimensions +import Unsafe.Coerce -- | Mutable DataFrame that lives in IO. -- Internal representation is always a MutableByteArray. -data family IODataFrame (t :: Type) (ns :: [k]) - --- | Pure wrapper on a mutable byte array -newtype instance IODataFrame t (ns :: [Nat]) = IODataFrame (MDataFrame RealWorld t (ns :: [Nat])) +newtype IODataFrame (t :: Type) (ns :: [k]) = IODataFrame (MDataFrame RealWorld t (ns :: [k])) -- | Data frame with some dimensions missing at compile time. -- Pattern-match against its constructor to get a Nat-indexed mutable data frame. -data instance IODataFrame t (xs :: [XNat]) - = forall (ns :: [Nat]) . Dimensions ns - => XIOFrame (IODataFrame t ns) +pattern XIOFrame :: forall (t :: Type) (xns :: [XNat]) . () + => forall (ns :: [Nat]) . (FixedDims xns ns, Dimensions ns) + => IODataFrame t ns -> IODataFrame t xns +pattern XIOFrame df <- (mkXIOFramePat -> XIOFramePat df) + where + XIOFrame = castDataFrame +{-# COMPLETE XIOFrame #-} +data XIOFramePat (t :: Type) (xns :: [XNat]) + = forall (ns :: [Nat]) . (FixedDims xns ns, Dimensions ns) + => XIOFramePat (IODataFrame t ns) + +mkXIOFramePat :: forall (t :: Type) (xns :: [XNat]) + . IODataFrame t xns -> XIOFramePat t xns +mkXIOFramePat df + | SomeDims ds <- fromSteps (getDataFrameSteps df) + , XDims (Dims :: Dims ns) <- (unsafeCoerce ds) :: Dims xns + = XIOFramePat @t @xns @ns (unsafeCoerce df) + | otherwise + = error "XIOFrame pattern: impossible args" + -- | Mutable DataFrame of unknown dimensionality data SomeIODataFrame (t :: Type) = forall (ns :: [Nat]) . Dimensions ns => SomeIODataFrame (IODataFrame t ns) +-- | Allow coercing between @XNat@-indexed and @Nat@-indexed Mutable DataFrames. +castDataFrame :: + forall (t :: Type) (xns :: [XNat]) (ns :: [Nat]) + . FixedDims xns ns + => IODataFrame t ns -> IODataFrame t xns +castDataFrame = coerce (castDataFrame# @t @xns @ns) +{-# INLINE castDataFrame #-} + -- | Create a new mutable DataFrame. -newDataFrame :: forall t (ns :: [Nat]) - . ( PrimBytes t, Dimensions ns) - => IO (IODataFrame t ns) -newDataFrame = IODataFrame <$> IO (newDataFrame# @t @ns) +newDataFrame :: + forall (t :: Type) ns + . (PrimBytes t, Dimensions ns) => IO (IODataFrame t ns) +newDataFrame = coerce (newDataFrame# @t @_ @ns) {-# INLINE newDataFrame #-} - -- | Create a new mutable DataFrame. -newPinnedDataFrame :: forall t (ns :: [Nat]) - . ( PrimBytes t, Dimensions ns) - => IO (IODataFrame t ns) -newPinnedDataFrame = IODataFrame <$> IO (newPinnedDataFrame# @t @ns) +newPinnedDataFrame :: + forall (t :: Type) ns + . (PrimBytes t, Dimensions ns) => IO (IODataFrame t ns) +newPinnedDataFrame = coerce (newPinnedDataFrame# @t @_ @ns) {-# INLINE newPinnedDataFrame #-} +-- | Create a new mutable DataFrame of the same size. +oneMoreDataFrame :: + forall (t :: Type) ns + . IODataFrame t ns -> IO (IODataFrame t ns) +oneMoreDataFrame = coerce (oneMoreDataFrame# @t @_ @ns) +{-# INLINE oneMoreDataFrame #-} +-- | View a part of a DataFrame. +-- +-- This function does not perform a copy. +-- All changes to a new DataFrame will be reflected in the original DataFrame as well. +-- +-- If any of the dims in @as@ or @b@ is unknown (@a ~ XN m@), +-- then this function is unsafe and can throw an `OutOfDimBounds` exception. +-- Otherwise, its safety is guaranteed by the type system. +subDataFrameView :: + forall (t :: Type) b bi bd as bs asbs + . (SubFrameIndexCtx b bi bd, KnownDim bd, ConcatList as (b :+ bs) asbs) + => Idxs (as +: bi) -> IODataFrame t asbs -> IODataFrame t (bd :+ bs) +subDataFrameView = coerce (subDataFrameView# @t @_ @b @bi @bd @as @bs @asbs) + +-- | View a part of a DataFrame. +-- +-- This function does not perform a copy. +-- All changes to a new DataFrame will be reflected in the original DataFrame as well. +-- +-- This is a simpler version of @subDataFrameView@ that allows +-- to view over one index at a time. +-- +-- If any of the dims in @as@ is unknown (@a ~ XN m@), +-- then this function is unsafe and can throw an `OutOfDimBounds` exception. +-- Otherwise, its safety is guaranteed by the type system. +subDataFrameView' :: + forall (t :: Type) as bs asbs + . ConcatList as bs asbs + => Idxs as -> IODataFrame t asbs -> IODataFrame t bs +subDataFrameView' = coerce (subDataFrameView'# @t @_ @as @bs @asbs) + -- | Copy one DataFrame into another mutable DataFrame at specified position. -- --- In contrast to @copyMDataFrame'@, this function allows to copy over a range of contiguous --- indices over a single dimension. +-- In contrast to @copyDataFrame'@, this function allows to copy over a range +-- of contiguous indices over a single dimension. -- For example, you can write a 3x4 matrix into a 7x4 matrix, starting at indices 0..3. -copyDataFrame :: forall (t :: Type) - (b :: Nat) (bi :: Nat) (bd :: Nat) - (as :: [Nat]) (bs :: [Nat]) (asbs :: [Nat]) - . ( b ~ (bi + bd - 1) - , PrimBytes t - , PrimBytes (DataFrame t (bd :+ bs)) - , ConcatList as (b :+ bs) asbs - ) - => Idxs (as +: bi) -> DataFrame t (bd :+ bs) - -> IODataFrame t asbs -> IO () -copyDataFrame ei df (IODataFrame mdf) = IO (copyDataFrame# ei df mdf) +-- +-- This function is safe (no `OutOfDimBounds` exception possible). +-- If any of the dims in @as@ is unknown (@a ~ XN m@), +-- you may happen to write data beyond dataframe bounds. +-- In this case, this function does nothing. +-- If (@b ~ XN m@) and (@Idx bi + Dim bd > Dim b@), this function copies only as +-- many elements as fits into the dataframe along this dimension (possibly none). +copyDataFrame :: + forall (t :: Type) b bi bd as bs asbs + . ( SubFrameIndexCtx b bi bd, KnownDim bd, ExactDims bs + , PrimArray t (DataFrame t (bd :+ bs)) + , ConcatList as (b :+ bs) asbs ) + => Idxs (as +: bi) -> DataFrame t (bd :+ bs) -> IODataFrame t asbs -> IO () +copyDataFrame = coerce (copyDataFrame# @t @_ @b @bi @bd @as @bs @asbs) {-# INLINE copyDataFrame #-} -- | Copy one mutable DataFrame into another mutable DataFrame at specified position. -- --- In contrast to @copyMDataFrame'@, this function allows to copy over a range of contiguous --- indices over a single dimension. +-- In contrast to @copyMutableDataFrame'@, this function allows to copy over a range +-- of contiguous indices over a single dimension. -- For example, you can write a 3x4 matrix into a 7x4 matrix, starting at indices 0..3. -copyMutableDataFrame :: forall (t :: Type) - (b :: Nat) (bi :: Nat) (bd :: Nat) - (as :: [Nat]) (bs :: [Nat]) (asbs :: [Nat]) - . ( b ~ (bi + bd - 1) - , PrimBytes t - , ConcatList as (b :+ bs) asbs - ) - => Idxs (as +: bi) -> IODataFrame t (bd :+ bs) - -> IODataFrame t asbs -> IO () -copyMutableDataFrame ei (IODataFrame mdfA) (IODataFrame mdfB) - = IO (copyMDataFrame# ei mdfA mdfB) +-- +-- This function is safe (no `OutOfDimBounds` exception possible). +-- If any of the dims in @as@ is unknown (@a ~ XN m@), +-- you may happen to write data beyond dataframe bounds. +-- In this case, this function does nothing. +-- If (@b ~ XN m@) and (@Idx bi + Dim bd > Dim b@), this function copies only as +-- many elements as fits into the dataframe along this dimension (possibly none). +copyMutableDataFrame :: + forall (t :: Type) b bi bd as bs asbs + . ( SubFrameIndexCtx b bi bd + , ExactDims bs + , PrimBytes t + , ConcatList as (b :+ bs) asbs ) + => Idxs (as +: bi) -> IODataFrame t (bd :+ bs) -> IODataFrame t asbs -> IO () +copyMutableDataFrame = coerce (copyMDataFrame# @t @_ @b @bi @bd @as @bs @asbs) {-# INLINE copyMutableDataFrame #-} -- | Copy one DataFrame into another mutable DataFrame at specified position. -- --- This is a simpler version of @copyDataFrame@ that allows to copy over one index at a time. -copyDataFrame' :: forall (t :: Type) - (as :: [Nat]) (bs :: [Nat]) (asbs :: [Nat]) - . ( PrimBytes t - , PrimBytes (DataFrame t bs) - , ConcatList as bs asbs - ) - => Idxs as -> DataFrame t bs - -> IODataFrame t asbs -> IO () -copyDataFrame' ei df (IODataFrame mdf) = IO (copyDataFrame'# ei df mdf) +-- This is a simpler version of @copyDataFrame@ that allows +-- to copy over one index at a time. +-- +-- This function is safe (no `OutOfDimBounds` exception possible). +-- If any of the dims in @as@ is unknown (@a ~ XN m@), +-- you may happen to write data beyond dataframe bounds. +-- In this case, this function does nothing. +copyDataFrame' :: + forall (t :: Type) as bs asbs + . ( ExactDims bs + , PrimArray t (DataFrame t bs) + , ConcatList as bs asbs ) + => Idxs as -> DataFrame t bs -> IODataFrame t asbs -> IO () +copyDataFrame' = coerce (copyDataFrame'# @t @_ @as @bs @asbs) {-# INLINE copyDataFrame' #-} -- | Copy one mutable DataFrame into another mutable DataFrame at specified position. -- --- This is a simpler version of @copyDataFrame@ that allows to copy over one index at a time. -copyMutableDataFrame' :: forall (t :: Type) - (as :: [Nat]) (bs :: [Nat]) (asbs :: [Nat]) - . ( PrimBytes t - , ConcatList as bs asbs - ) - => Idxs as -> IODataFrame t bs - -> IODataFrame t asbs -> IO () -copyMutableDataFrame' ei (IODataFrame mdfA) (IODataFrame mdfB) - = IO (copyMDataFrame'# ei mdfA mdfB) +-- This is a simpler version of @copyMutableDataFrame@ that allows +-- to copy over one index at a time. +-- +-- This function is safe (no `OutOfDimBounds` exception possible). +-- If any of the dims in @as@ is unknown (@a ~ XN m@), +-- you may happen to write data beyond dataframe bounds. +-- In this case, this function does nothing. +copyMutableDataFrame' :: + forall (t :: Type) as bs asbs + . (ExactDims bs, PrimBytes t, ConcatList as bs asbs) + => Idxs as -> IODataFrame t bs -> IODataFrame t asbs -> IO () +copyMutableDataFrame' = coerce (copyMDataFrame'# @t @_ @as @bs @asbs) {-# INLINE copyMutableDataFrame' #-} +-- | Copy one DataFrame into another mutable DataFrame by offset in +-- primitive elements. +-- +-- This is a low-level copy function; you have to keep in mind the row-major +-- layout of Mutable DataFrames. Offset bounds are not checked. +-- You will get an undefined behavior if you write beyond the DataFrame bounds. +copyDataFrameOff :: + forall (t :: Type) as bs asbs + . ( Dimensions bs + , PrimArray t (DataFrame t bs) + , ConcatList as bs asbs ) + => Int -> DataFrame t bs -> IODataFrame t asbs -> IO () +copyDataFrameOff = coerce (copyDataFrameOff# @t @_ @as @bs @asbs) +{-# INLINE copyDataFrameOff #-} + +-- | Copy one mutable DataFrame into another mutable DataFrame by offset in +-- primitive elements. +-- +-- This is a low-level copy function; you have to keep in mind the row-major +-- layout of Mutable DataFrames. Offset bounds are not checked. +-- You will get an undefined behavior if you write beyond the DataFrame bounds +copyMutableDataFrameOff :: + forall (t :: Type) as bs asbs + . (ExactDims bs, PrimBytes t, ConcatList as bs asbs) + => Int -> IODataFrame t bs -> IODataFrame t asbs -> IO () +copyMutableDataFrameOff = coerce (copyMDataFrameOff# @t @_ @as @bs @asbs) +{-# INLINE copyMutableDataFrameOff #-} + -- | Make a mutable DataFrame immutable, without copying. -unsafeFreezeDataFrame :: forall (t :: Type) (ns :: [Nat]) - . PrimArray t (DataFrame t ns) - => IODataFrame t ns -> IO (DataFrame t ns) -unsafeFreezeDataFrame (IODataFrame mdf) = IO (unsafeFreezeDataFrame# mdf) +unsafeFreezeDataFrame :: + forall (t :: Type) ns + . PrimArray t (DataFrame t ns) + => IODataFrame t ns -> IO (DataFrame t ns) +unsafeFreezeDataFrame = coerce (unsafeFreezeDataFrame# @t @_ @ns) {-# INLINE unsafeFreezeDataFrame #-} - -- | Copy content of a mutable DataFrame into a new immutable DataFrame. -freezeDataFrame :: forall (t :: Type) (ns :: [Nat]) - . PrimArray t (DataFrame t ns) - => IODataFrame t ns -> IO (DataFrame t ns) -freezeDataFrame (IODataFrame mdf) = IO (freezeDataFrame# mdf) +freezeDataFrame :: + forall (t :: Type) ns + . PrimArray t (DataFrame t ns) + => IODataFrame t ns -> IO (DataFrame t ns) +freezeDataFrame = coerce (freezeDataFrame# @t @_ @ns) {-# INLINE freezeDataFrame #-} -- | Create a new mutable DataFrame and copy content of immutable one in there. -thawDataFrame :: forall (t :: Type) (ns :: [Nat]) - . (Dimensions ns, PrimBytes (DataFrame t ns)) - => DataFrame t ns -> IO (IODataFrame t ns) -thawDataFrame df = IODataFrame <$> IO (thawDataFrame# df) +thawDataFrame :: + forall (t :: Type) ns + . (Dimensions ns, PrimArray t (DataFrame t ns)) + => DataFrame t ns -> IO (IODataFrame t ns) +thawDataFrame = coerce (thawDataFrame# @t @_ @ns) {-# INLINE thawDataFrame #-} -- | Create a new mutable DataFrame and copy content of immutable one in there. -- The result array is pinned and aligned. -thawPinDataFrame :: forall (t :: Type) (ns :: [Nat]) - . (Dimensions ns, PrimBytes (DataFrame t ns)) - => DataFrame t ns -> IO (IODataFrame t ns) -thawPinDataFrame df = IODataFrame <$> IO (thawPinDataFrame# df) +thawPinDataFrame :: + forall (t :: Type) ns + . (Dimensions ns, PrimArray t (DataFrame t ns)) + => DataFrame t ns -> IO (IODataFrame t ns) +thawPinDataFrame = coerce (thawPinDataFrame# @t @_ @ns) {-# INLINE thawPinDataFrame #-} -- | UnsafeCoerces an underlying byte array. -unsafeThawDataFrame :: forall (t :: Type) (ns :: [Nat]) - . ( Dimensions ns - , PrimBytes (DataFrame t ns), PrimBytes t) - => DataFrame t ns -> IO (IODataFrame t ns) -unsafeThawDataFrame df = IODataFrame <$> IO (unsafeThawDataFrame# df) +unsafeThawDataFrame :: + forall (t :: Type) ns + . (Dimensions ns, PrimArray t (DataFrame t ns)) + => DataFrame t ns + -> IO (IODataFrame t ns) +unsafeThawDataFrame = coerce (unsafeThawDataFrame# @t @_ @ns) {-# INLINE unsafeThawDataFrame #-} - --- | Write a single element at the specified index -writeDataFrame :: forall t (ns :: [Nat]) - . PrimBytes t - => IODataFrame t ns -> Idxs ns -> DataFrame t ('[] :: [Nat]) -> IO () -writeDataFrame (IODataFrame mdf) ei = IO . writeDataFrame# mdf ei . unsafeCoerce# -{-# INLINE writeDataFrame #-} - - --- | Read a single element at the specified index -readDataFrame :: forall (t :: Type) (ns :: [Nat]) - . PrimBytes t - => IODataFrame t ns -> Idxs ns -> IO (DataFrame t ('[] :: [Nat])) -readDataFrame (IODataFrame mdf) = unsafeCoerce# . IO . readDataFrame# mdf -{-# INLINE readDataFrame #-} - +-- | Given two continuations @f@ and @g@. +-- If the input DataFrame is a single broadcast value, use it in @f@. +-- Otherwise, create a new mutable DataFrame and copy content of immutable one +-- in there; then use it in @g@. +-- +-- This function is useful when @thawDataFrame@ cannot be used due to +-- @Dimensions ns@ constraint being not available. +withThawDataFrame :: + forall (t :: Type) ns r + . PrimArray t (DataFrame t ns) + => (t -> IO r) + -> (IODataFrame t ns -> IO r) + -> DataFrame t ns -> IO r +withThawDataFrame = coerce (withThawDataFrame# @t @_ @ns @r) --- | Write a single element at the specified element offset -writeDataFrameOff :: forall (t :: Type) (ns :: [Nat]) - . PrimBytes t - => IODataFrame t ns -> Int -> DataFrame t ('[] :: [Nat]) -> IO () -writeDataFrameOff (IODataFrame mdf) (I# i) - = IO . writeDataFrameOff# mdf i . unsafeCoerce# +-- | Write a single element at the specified element offset. +-- +-- This is a low-level write function; you have to keep in mind the row-major +-- layout of Mutable DataFrames. Offset bounds are not checked. +-- You will get an undefined behavior if you write beyond the DataFrame bounds. +writeDataFrameOff :: + forall (t :: Type) ns + . PrimBytes (DataFrame t ('[] :: KindOf ns)) + => IODataFrame t ns -> Int -> DataFrame t ('[] :: KindOf ns) -> IO () +writeDataFrameOff = coerce (writeDataFrameOff# @t @_ @ns) {-# INLINE writeDataFrameOff #-} +-- | Write a single element at the specified index. +-- +-- This function is safe (no `OutOfDimBounds` exception possible). +-- If any of the dims in @ns@ is unknown (@n ~ XN m@), +-- you may happen to write data beyond dataframe bounds. +-- In this case, this function does nothing. +writeDataFrame :: + forall (t :: Type) ns + . PrimBytes (DataFrame t ('[] :: KindOf ns)) + => IODataFrame t ns -> Idxs ns -> DataFrame t ('[] :: KindOf ns) -> IO () +writeDataFrame = coerce (writeDataFrame# @t @_ @ns) +{-# INLINE writeDataFrame #-} --- | Read a single element at the specified element offset -readDataFrameOff :: forall (t :: Type) (ns :: [Nat]) - . PrimBytes t - => IODataFrame t ns -> Int -> IO (DataFrame t ('[] :: [Nat])) -readDataFrameOff (IODataFrame mdf) (I# i) - = unsafeCoerce# (IO (readDataFrameOff# mdf i)) +-- | Read a single element at the specified element offset. +-- +-- This is a low-level read function; you have to keep in mind the row-major +-- layout of Mutable DataFrames. Offset bounds are not checked. +-- You will get an undefined behavior if you read beyond the DataFrame bounds. +readDataFrameOff :: + forall (t :: Type) ns + . PrimBytes (DataFrame t ('[] :: KindOf ns)) + => IODataFrame t ns -> Int -> IO (DataFrame t ('[] :: KindOf ns)) +readDataFrameOff = coerce (readDataFrameOff# @t @_ @ns) {-# INLINE readDataFrameOff #-} +-- | Read a single element at the specified index. +-- +-- If any of the dims in @ns@ is unknown (@n ~ XN m@), +-- then this function is unsafe and can throw an `OutOfDimBounds` exception. +-- Otherwise, its safety is guaranteed by the type system. +readDataFrame :: + forall (t :: Type) ns + . PrimBytes (DataFrame t ('[] :: KindOf ns)) + => IODataFrame t ns -> Idxs ns -> IO (DataFrame t ('[] :: KindOf ns)) +readDataFrame = coerce (readDataFrame# @t @_ @ns) +{-# INLINE readDataFrame #-} -- | Check if the byte array wrapped by this DataFrame is pinned, -- which means cannot be relocated by GC. -isDataFramePinned :: forall (k :: Type) (t :: Type) (ns :: [k]) - . KnownDimKind k - => IODataFrame t ns -> Bool -isDataFramePinned df = case dimKind @k of - DimNat -> case df of - IODataFrame x -> isDataFramePinned# x - DimXNat -> case df of - XIOFrame (IODataFrame x) -> isDataFramePinned# x +isDataFramePinned :: + forall (t :: Type) ns . IODataFrame t ns -> Bool +isDataFramePinned = coerce (isDataFramePinned# @t @_ @ns) +{-# INLINE isDataFramePinned #-} +-- | Get cumulative dimensions @ns@ of a @IODataFrame t ns@ +getDataFrameSteps :: + forall (t :: Type) ns . IODataFrame t ns -> CumulDims +getDataFrameSteps = coerce (getDataFrameSteps# @t @_ @ns) +{-# INLINE getDataFrameSteps #-} -- | Allow arbitrary IO operations on a pointer to the beginning of the data -- keeping the data from garbage collecting until the arg function returns. @@ -233,15 +377,9 @@ -- -- Warning: use this function on a pinned DataFrame only; -- otherwise, the data may be relocated before the arg fun finishes. -withDataFramePtr :: forall (k :: Type) (t :: Type) (ns :: [k]) (r :: Type) - . (PrimBytes t, KnownDimKind k) - => IODataFrame t ns - -> ( Ptr t -> IO r ) - -> IO r -withDataFramePtr df k = case dimKind @k of - DimNat -> case df of - IODataFrame x - -> IO $ withDataFramePtr# x (\p -> case k (Ptr p) of IO f -> f) - DimXNat -> case df of - XIOFrame (IODataFrame x) - -> IO $ withDataFramePtr# x (\p -> case k (Ptr p) of IO f -> f) +withDataFramePtr :: + forall (t :: Type) ns (r :: Type) + . PrimBytes t => IODataFrame t ns -> (Ptr t -> IO r) -> IO r +withDataFramePtr (IODataFrame x) k + = IO (withDataFramePtr# x (\p -> case k (Ptr p) of IO f -> f)) +{-# INLINE withDataFramePtr #-} diff --git a/src/Numeric/DataFrame/Internal/Backend.hs b/src/Numeric/DataFrame/Internal/Backend.hs --- a/src/Numeric/DataFrame/Internal/Backend.hs +++ b/src/Numeric/DataFrame/Internal/Backend.hs @@ -2,6 +2,7 @@ {-# LANGUAGE ConstraintKinds #-} {-# LANGUAGE DataKinds #-} {-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE PolyKinds #-} {-# LANGUAGE RoleAnnotations #-} @@ -9,15 +10,39 @@ {-# LANGUAGE TypeApplications #-} {-# LANGUAGE TypeFamilies #-} {-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UndecidableInstances #-} {-# OPTIONS_GHC -fno-warn-redundant-constraints #-} #if defined(__HADDOCK__) || defined(__HADDOCK_VERSION__) {-# OPTIONS_GHC -fno-warn-simplifiable-class-constraints #-} #else {-# OPTIONS_GHC -fplugin Data.Constraint.Deriving #-} #endif +{- +This module contains all generic and specific instances for all backend +(implementation) types. Its hs-boot counterpart provides only generic instances +(that do not depend on BackendFamily). +Import this module only at the root of the module hierarchy, where the specific +implementation-dependent modules do not need to depend on it. +But do not forget to import this module in the places that will be imported by +the user (such as Numeric.DataFrame). So that the user can benefit from the +compile-time knowledge of the specific backend implementation. +By design the import chain is as follows: + +NDI.Backend.Family.hs-boot NDI.Backend.Family + \ | + NDI.Backend.hs-boot NDI.Backend (instances only) + \ / + NDI.BackendI / + \ / + Numeric.DataFrame.Type / + \ / + \ / + Numeric.DataFrame + + -} module Numeric.DataFrame.Internal.Backend - ( DFBackend, Backend (..), BackendFamily, KnownBackend () + ( Backend (..), BackendFamily, KnownBackend () , inferKnownBackend, inferPrimElem -- * Auto-deriving instances , inferEq, inferOrd @@ -47,15 +72,15 @@ inferKnownBackend, inferPrimArray, inferPrimElem) -import Numeric.DataFrame.Internal.BackendI (I) --- | Implementation behind the DataFrame -type DFBackend (t :: Type) (ds :: [Nat]) = Backend I t ds (BackendFamily t ds) -- | Backend resolver: -- Use this constraint to find any class instances defined for all DataFrame implementations, -- e.g. @Num@, @PrimBytes@, etc. -type KnownBackend (t :: Type) (ds :: [Nat]) = Impl.KnownBackend t ds (BackendFamily t ds) +class Impl.KnownBackend t ds (BackendFamily t ds) + => KnownBackend (t :: Type) (ds :: [Nat]) +instance Impl.KnownBackend t ds (BackendFamily t ds) + => KnownBackend (t :: Type) (ds :: [Nat]) -- | A newtype wrapper for all DataFrame implementations. -- I need two layers of wrappers to provide default overlappable instances to @@ -65,12 +90,12 @@ newtype Backend (i :: Type) (t :: Type) (ds :: [Nat]) (backend :: Type) = Backend { _getBackend :: backend } type role Backend phantom phantom phantom representational -type instance DeriveContext (Backend I t ds b) = b ~ BackendFamily t ds --- When typechecker knows what the third parameter @b@ is, no instances overlap, +type instance DeriveContext (Backend i t ds b) = b ~ BackendFamily t ds +-- When typechecker knows what the fourth parameter @b@ is, no instances overlap, -- because @b@ is always always a concrete type. -- The "dynamic instances" derived via ToInstance have to be incoherent, -- because @BackendFamily t ds@ as a type family (before resolving it) --- is no more conrete than just @b@. +-- is no more concrete than just @b@. {-# ANN type Backend (DeriveAll' NoOverlap ["KnownBackend"]) #-} @@ -80,47 +105,48 @@ . (PrimBytes t, Dimensions ds) => Dict (KnownBackend t ds) inferKnownBackend - = Impl.inferKnownBackend @t @ds @(BackendFamily t ds) + = case Impl.inferKnownBackend @t @ds @(BackendFamily t ds) of + Dict -> Dict inferPrimElem - :: forall (t :: Type) (d :: Nat) (ds :: [Nat]) + :: forall (t :: Type) (d :: Nat) (ds :: [Nat]) (i :: Type) . KnownBackend t (d ': ds) - => DFBackend t (d ': ds) -> Dict (PrimBytes t) + => Backend i t (d ': ds) (BackendFamily t (d ': ds)) -> Dict (PrimBytes t) inferPrimElem = Impl.inferPrimElem @t @d @ds . _getBackend inferEq - :: forall (t :: Type) (ds :: [Nat]) (b :: Type) + :: forall (t :: Type) (ds :: [Nat]) (b :: Type) (i :: Type) . (Eq t, Impl.KnownBackend t ds b) - => Dict (Eq (Backend I t ds b)) + => Dict (Eq (Backend i t ds b)) inferEq = mapDict toBackend . mapDict (Sub (Impl.inferBackendInstance @t @ds)) $ inferDeriveContext @t @ds @b undefined inferOrd - :: forall (t :: Type) (ds :: [Nat]) (b :: Type) + :: forall (t :: Type) (ds :: [Nat]) (b :: Type) (i :: Type) . (Ord t, Impl.KnownBackend t ds b) - => Dict (Ord (Backend I t ds b)) + => Dict (Ord (Backend i t ds b)) inferOrd = mapDict toBackend . mapDict (Sub (Impl.inferBackendInstance @t @ds)) $ inferDeriveContext @t @ds @b undefined inferProductOrder - :: forall (t :: Type) (ds :: [Nat]) (b :: Type) + :: forall (t :: Type) (ds :: [Nat]) (b :: Type) (i :: Type) . (Ord t, Impl.KnownBackend t ds b) - => Dict (ProductOrder (Backend I t ds b)) + => Dict (ProductOrder (Backend i t ds b)) inferProductOrder = mapDict toBackend . mapDict (Sub (Impl.inferBackendInstance @t @ds)) $ inferDeriveContext @t @ds @b undefined inferPONonTransitive - :: forall (t :: Type) (ds :: [Nat]) (b :: Type) + :: forall (t :: Type) (ds :: [Nat]) (b :: Type) (i :: Type) . (Ord t, Impl.KnownBackend t ds b) - => Dict (Ord (NonTransitive.ProductOrd (Backend I t ds b))) + => Dict (Ord (NonTransitive.ProductOrd (Backend i t ds b))) inferPONonTransitive = mapDict (unsafeDerive NonTransitive.ProductOrd) . mapDict toBackend @@ -128,9 +154,9 @@ $ inferDeriveContext @t @ds @b undefined inferPOPartial - :: forall (t :: Type) (ds :: [Nat]) (b :: Type) + :: forall (t :: Type) (ds :: [Nat]) (b :: Type) (i :: Type) . (Ord t, Impl.KnownBackend t ds b) - => Dict (Ord (Partial.ProductOrd (Backend I t ds b))) + => Dict (Ord (Partial.ProductOrd (Backend i t ds b))) inferPOPartial = mapDict (unsafeDerive Partial.ProductOrd) . mapDict toBackend @@ -138,9 +164,9 @@ $ inferDeriveContext @t @ds @b undefined inferBounded - :: forall (t :: Type) (ds :: [Nat]) (b :: Type) + :: forall (t :: Type) (ds :: [Nat]) (b :: Type) (i :: Type) . (Bounded t, Impl.KnownBackend t ds b) - => Dict (Bounded (Backend I t ds b)) + => Dict (Bounded (Backend i t ds b)) inferBounded -- Here is an ugly trick: -- We don't have instance of @Bounded Float@ and @Bounded Double@. @@ -160,27 +186,27 @@ $ inferDeriveContext @t @ds @b undefined inferNum - :: forall (t :: Type) (ds :: [Nat]) (b :: Type) + :: forall (t :: Type) (ds :: [Nat]) (b :: Type) (i :: Type) . (Num t, Impl.KnownBackend t ds b) - => Dict (Num (Backend I t ds b)) + => Dict (Num (Backend i t ds b)) inferNum = mapDict toBackend . mapDict (Sub (Impl.inferBackendInstance @t @ds)) $ inferDeriveContext @t @ds @b undefined inferFractional - :: forall (t :: Type) (ds :: [Nat]) (b :: Type) + :: forall (t :: Type) (ds :: [Nat]) (b :: Type) (i :: Type) . (Fractional t, Impl.KnownBackend t ds b) - => Dict (Fractional (Backend I t ds b)) + => Dict (Fractional (Backend i t ds b)) inferFractional = mapDict toBackend . mapDict (Sub (Impl.inferBackendInstance @t @ds)) $ inferDeriveContext @t @ds @b undefined inferFloating - :: forall (t :: Type) (ds :: [Nat]) (b :: Type) + :: forall (t :: Type) (ds :: [Nat]) (b :: Type) (i :: Type) . (Floating t, Impl.KnownBackend t ds b) - => Dict (Floating (Backend I t ds b)) + => Dict (Floating (Backend i t ds b)) inferFloating = mapDict toBackend . mapDict (Sub (Impl.inferBackendInstance @t @ds)) @@ -188,18 +214,18 @@ inferPrimBytes - :: forall (t :: Type) (ds :: [Nat]) (b :: Type) + :: forall (t :: Type) (ds :: [Nat]) (b :: Type) (i :: Type) . (PrimBytes t, Dimensions ds, Impl.KnownBackend t ds b) - => Dict (PrimBytes (Backend I t ds b)) + => Dict (PrimBytes (Backend i t ds b)) inferPrimBytes = mapDict toBackend . mapDict (Sub (Impl.inferBackendInstance @t @ds)) $ inferDeriveContext @t @ds @b undefined inferPrimArray - :: forall (t :: Type) (ds :: [Nat]) (b :: Type) + :: forall (t :: Type) (ds :: [Nat]) (b :: Type) (i :: Type) . (PrimBytes t, Impl.KnownBackend t ds b) - => Dict (PrimArray t (Backend I t ds b)) + => Dict (PrimArray t (Backend i t ds b)) inferPrimArray = mapDict toBackend . mapDict (Sub (Impl.inferPrimArray @t @ds)) @@ -208,14 +234,14 @@ -- This is the rule that cannot be encoded in the type system, but enforced -- as an invariant: Backend t ds b implies DeriveContext t ds b -inferDeriveContext :: forall (t :: Type) (ds :: [Nat]) (b :: Type) - . Backend I t ds b - -> Dict (DeriveContext (Backend I t ds b)) +inferDeriveContext :: forall (t :: Type) (ds :: [Nat]) (b :: Type) (i :: Type) + . Backend i t ds b + -> Dict (DeriveContext (Backend i t ds b)) inferDeriveContext _ = unsafeCoerce (Dict :: Dict (b ~ b)) {-# INLINE inferDeriveContext #-} -- Backend is the newtype wrapper over b. -- It has the same represenation and I expect it to have the same instance behavior. -toBackend :: forall c t ds b . c b :- c (Backend I t ds b) +toBackend :: forall c t ds b i . c b :- c (Backend i t ds b) toBackend = unsafeDerive Backend {-# INLINE toBackend #-} diff --git a/src/Numeric/DataFrame/Internal/Backend.hs-boot b/src/Numeric/DataFrame/Internal/Backend.hs-boot --- a/src/Numeric/DataFrame/Internal/Backend.hs-boot +++ b/src/Numeric/DataFrame/Internal/Backend.hs-boot @@ -1,17 +1,22 @@ -{-# LANGUAGE CPP #-} -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE DataKinds #-} -{-# LANGUAGE ExplicitForAll #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE PolyKinds #-} -{-# LANGUAGE RoleAnnotations #-} -{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE ExplicitForAll #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE PolyKinds #-} +{-# LANGUAGE RoleAnnotations #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UndecidableInstances #-} #if defined(__HADDOCK__) || defined(__HADDOCK_VERSION__) {-# OPTIONS_GHC -fno-warn-simplifiable-class-constraints #-} #endif module Numeric.DataFrame.Internal.Backend - ( DFBackend, Backend (..), BackendFamily, KnownBackend () + ( Backend (..), BackendFamily, KnownBackend () , inferKnownBackend, inferPrimElem -- * Auto-deriving instances , inferEq, inferOrd @@ -32,21 +37,12 @@ import {-# SOURCE #-} Numeric.DataFrame.Internal.Backend.Family (BackendFamily) import {-# SOURCE #-} qualified Numeric.DataFrame.Internal.Backend.Family as Impl (KnownBackend) -import {-# SOURCE #-} Numeric.DataFrame.Internal.BackendI (I) --- | Implementation behind the DataFrame -type DFBackend (t :: Type) (ds :: [Nat]) = Backend I t ds (BackendFamily t ds) - --- | Backend resolver: --- Use this constraint to find any class instances defined for all DataFrame implementations, --- e.g. @Num@, @PrimBytes@, etc. -type KnownBackend (t :: Type) (ds :: [Nat]) = Impl.KnownBackend t ds (BackendFamily t ds) +class Impl.KnownBackend t ds (BackendFamily t ds) + => KnownBackend (t :: Type) (ds :: [Nat]) +instance Impl.KnownBackend t ds (BackendFamily t ds) + => KnownBackend (t :: Type) (ds :: [Nat]) --- | A newtype wrapper for all DataFrame implementations. --- I need two layers of wrappers to provide default overlappable instances to --- all type classes using KnownBackend mechanics. --- Type arguments are redundant here; --- nevertheless, they improve readability of error messages. newtype Backend (i :: Type) (t :: Type) (ds :: [Nat]) (backend :: Type) = Backend { _getBackend :: backend } type role Backend phantom phantom phantom representational @@ -57,61 +53,61 @@ => Dict (KnownBackend t ds) inferPrimElem - :: forall (t :: Type) (d :: Nat) (ds :: [Nat]) + :: forall (t :: Type) (d :: Nat) (ds :: [Nat]) (i :: Type) . KnownBackend t (d ': ds) - => DFBackend t (d ': ds) -> Dict (PrimBytes t) + => Backend i t (d ': ds) (BackendFamily t (d ': ds)) -> Dict (PrimBytes t) inferEq - :: forall (t :: Type) (ds :: [Nat]) (b :: Type) + :: forall (t :: Type) (ds :: [Nat]) (b :: Type) (i :: Type) . (Eq t, Impl.KnownBackend t ds b) - => Dict (Eq (Backend I t ds b)) + => Dict (Eq (Backend i t ds b)) inferOrd - :: forall (t :: Type) (ds :: [Nat]) (b :: Type) + :: forall (t :: Type) (ds :: [Nat]) (b :: Type) (i :: Type) . (Ord t, Impl.KnownBackend t ds b) - => Dict (Ord (Backend I t ds b)) + => Dict (Ord (Backend i t ds b)) inferProductOrder - :: forall (t :: Type) (ds :: [Nat]) (b :: Type) + :: forall (t :: Type) (ds :: [Nat]) (b :: Type) (i :: Type) . (Ord t, Impl.KnownBackend t ds b) - => Dict (ProductOrder (Backend I t ds b)) + => Dict (ProductOrder (Backend i t ds b)) inferPONonTransitive - :: forall (t :: Type) (ds :: [Nat]) (b :: Type) + :: forall (t :: Type) (ds :: [Nat]) (b :: Type) (i :: Type) . (Ord t, Impl.KnownBackend t ds b) - => Dict (Ord (NonTransitive.ProductOrd (Backend I t ds b))) + => Dict (Ord (NonTransitive.ProductOrd (Backend i t ds b))) inferPOPartial - :: forall (t :: Type) (ds :: [Nat]) (b :: Type) + :: forall (t :: Type) (ds :: [Nat]) (b :: Type) (i :: Type) . (Ord t, Impl.KnownBackend t ds b) - => Dict (Ord (Partial.ProductOrd (Backend I t ds b))) + => Dict (Ord (Partial.ProductOrd (Backend i t ds b))) inferBounded - :: forall (t :: Type) (ds :: [Nat]) (b :: Type) + :: forall (t :: Type) (ds :: [Nat]) (b :: Type) (i :: Type) . (Bounded t, Impl.KnownBackend t ds b) - => Dict (Bounded (Backend I t ds b)) + => Dict (Bounded (Backend i t ds b)) inferNum - :: forall (t :: Type) (ds :: [Nat]) (b :: Type) + :: forall (t :: Type) (ds :: [Nat]) (b :: Type) (i :: Type) . (Num t, Impl.KnownBackend t ds b) - => Dict (Num (Backend I t ds b)) + => Dict (Num (Backend i t ds b)) inferFractional - :: forall (t :: Type) (ds :: [Nat]) (b :: Type) + :: forall (t :: Type) (ds :: [Nat]) (b :: Type) (i :: Type) . (Fractional t, Impl.KnownBackend t ds b) - => Dict (Fractional (Backend I t ds b)) + => Dict (Fractional (Backend i t ds b)) inferFloating - :: forall (t :: Type) (ds :: [Nat]) (b :: Type) + :: forall (t :: Type) (ds :: [Nat]) (b :: Type) (i :: Type) . (Floating t, Impl.KnownBackend t ds b) - => Dict (Floating (Backend I t ds b)) + => Dict (Floating (Backend i t ds b)) inferPrimBytes - :: forall (t :: Type) (ds :: [Nat]) (b :: Type) + :: forall (t :: Type) (ds :: [Nat]) (b :: Type) (i :: Type) . (PrimBytes t, Dimensions ds, Impl.KnownBackend t ds b) - => Dict (PrimBytes (Backend I t ds b)) + => Dict (PrimBytes (Backend i t ds b)) inferPrimArray - :: forall (t :: Type) (ds :: [Nat]) (b :: Type) + :: forall (t :: Type) (ds :: [Nat]) (b :: Type) (i :: Type) . (PrimBytes t, Impl.KnownBackend t ds b) - => Dict (PrimArray t (Backend I t ds b)) + => Dict (PrimArray t (Backend i t ds b)) diff --git a/src/Numeric/DataFrame/Internal/BackendI.hs b/src/Numeric/DataFrame/Internal/BackendI.hs --- a/src/Numeric/DataFrame/Internal/BackendI.hs +++ b/src/Numeric/DataFrame/Internal/BackendI.hs @@ -7,31 +7,54 @@ {-# LANGUAGE TypeApplications #-} #if defined(__HADDOCK__) || defined(__HADDOCK_VERSION__) {-# LANGUAGE FlexibleInstances #-} -{-# OPTIONS_GHC -fdefer-type-errors #-} -{-# OPTIONS_GHC -fno-warn-deferred-type-errors #-} {-# OPTIONS_GHC -fno-warn-missing-methods #-} {-# OPTIONS_GHC -fno-warn-redundant-constraints #-} {-# OPTIONS_GHC -fno-warn-simplifiable-class-constraints #-} #else {-# OPTIONS_GHC -fplugin Data.Constraint.Deriving #-} #endif +{- +The purpose of this module is to provide incoherent instances of common type classes +for the type @Backend@. +It imports @Backend.hs-boot@ module and thus does not depend on the actual +implementation of @BackendFamily@. +Import this module instead of @Backend.hs-boot@ in the places that must not depend +on @BackendFamily@, such as @Numeric.DataFrame.Type@. +By design the import chain is as follows: + +NDI.Backend.Family.hs-boot NDI.Backend.Family + \ | + NDI.Backend.hs-boot NDI.Backend (instances only) + \ / + NDI.BackendI / + \ / + Numeric.DataFrame.Type / + \ / + \ / + Numeric.DataFrame + + -} module Numeric.DataFrame.Internal.BackendI - ( I + ( DFBackend, Backend (..), Backend.BackendFamily, Backend.KnownBackend () + , Backend.inferKnownBackend, Backend.inferPrimElem +#if !defined(__HADDOCK__) && !defined(__HADDOCK_VERSION__) , inferEq, inferOrd , inferProductOrder, inferPONonTransitive, inferPOPartial , inferBounded, inferNum , inferFractional, inferFloating , inferPrimBytes, inferPrimArray +#endif ) where - -import Data.Constraint.Deriving (OverlapMode (..), - ToInstance (..)) +#if !defined(__HADDOCK__) && !defined(__HADDOCK_VERSION__) +import Data.Constraint (Dict) +import Data.Constraint.Deriving (OverlapMode (..), ToInstance (..)) +#endif import Data.Kind (Type) import Numeric.DataFrame.Internal.PrimArray (PrimArray) -import Numeric.Dimensions (Dict, Dimensions, Nat) +import Numeric.Dimensions (Dimensions, Nat) import Numeric.PrimBytes (PrimBytes) import Numeric.ProductOrd (ProductOrder) import qualified Numeric.ProductOrd.NonTransitive as NonTransitive @@ -43,6 +66,9 @@ +-- | Implementation behind the DataFrame +type DFBackend (t :: Type) (ds :: [Nat]) = Backend I t ds (Impl.BackendFamily t ds) + {- | The instance keeper for the `Backend` type. Using this data as a tag to the `Backend` type allows to define `Backend` instances @@ -52,6 +78,10 @@ data I + +-- The instances below are bumb stubs needed only to compile haddock +#if !defined(__HADDOCK__) && !defined(__HADDOCK_VERSION__) + {-# ANN inferEq (ToInstance Incoherent) #-} inferEq :: forall (t :: Type) (ds :: [Nat]) (b :: Type) @@ -129,10 +159,7 @@ => Dict (PrimArray t (Backend I t ds b)) inferPrimArray = Backend.inferPrimArray @t @ds @b - - --- The instances below are bumb stubs needed only to compile haddock -#if defined(__HADDOCK__) || defined(__HADDOCK_VERSION__) +#else instance forall (t :: Type) (ds :: [Nat]) (b :: Type) diff --git a/src/Numeric/DataFrame/Internal/BackendI.hs-boot b/src/Numeric/DataFrame/Internal/BackendI.hs-boot deleted file mode 100644 --- a/src/Numeric/DataFrame/Internal/BackendI.hs-boot +++ /dev/null @@ -1,9 +0,0 @@ -module Numeric.DataFrame.Internal.BackendI (I) where - -{- | The instance keeper for the `Backend` type. - - Using this data as a tag to the `Backend` type allows to define `Backend` instances - in two different modules without any orphans. - - -} -data I diff --git a/src/Numeric/DataFrame/Internal/Mutable.hs b/src/Numeric/DataFrame/Internal/Mutable.hs --- a/src/Numeric/DataFrame/Internal/Mutable.hs +++ b/src/Numeric/DataFrame/Internal/Mutable.hs @@ -8,6 +8,7 @@ {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} {-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeInType #-} {-# LANGUAGE TypeOperators #-} {-# LANGUAGE UnboxedTuples #-} {-# OPTIONS_GHC -fno-warn-redundant-constraints #-} @@ -17,264 +18,511 @@ -- Copyright : (c) Artem Chirkin -- License : BSD3 -- --- Maintainer : chirkin@arch.ethz.ch -- --- Interfrace to perform primitive stateful operations on mutable frames. +-- Interface to perform primitive stateful operations on mutable frames. -- ----------------------------------------------------------------------------- module Numeric.DataFrame.Internal.Mutable ( MDataFrame () + , castDataFrame# , newDataFrame#, newPinnedDataFrame# + , oneMoreDataFrame# + , subDataFrameView#, subDataFrameView'# , copyDataFrame#, copyMDataFrame# , copyDataFrame'#, copyMDataFrame'# + , copyDataFrameOff#, copyMDataFrameOff# , freezeDataFrame#, unsafeFreezeDataFrame# - , thawDataFrame#, thawPinDataFrame#, unsafeThawDataFrame# + , thawDataFrame#, thawPinDataFrame#, unsafeThawDataFrame#, withThawDataFrame# , writeDataFrame#, writeDataFrameOff# , readDataFrame#, readDataFrameOff# , withDataFramePtr#, isDataFramePinned# + , getDataFrameSteps# ) where - import GHC.Base import Numeric.DataFrame.Internal.PrimArray import Numeric.DataFrame.Type import Numeric.Dimensions import Numeric.PrimBytes - -- | Mutable DataFrame type. -- Keeps element offset, number of elements, and a mutable byte storage -data MDataFrame s t (ns :: [Nat]) +data MDataFrame s t (ns :: [k]) = MDataFrame# Int# CumulDims (MutableByteArray# s) +-- | Allow coercing between @XNat@-indexed and @Nat@-indexed Mutable DataFrames. +castDataFrame# :: + forall (t :: Type) (xns :: [XNat]) (ns :: [Nat]) s + . FixedDims xns ns + => MDataFrame s t ns -> MDataFrame s t xns +castDataFrame# (MDataFrame# o c a) = MDataFrame# o c a +{-# INLINE castDataFrame# #-} -- | Create a new mutable DataFrame. -newDataFrame# :: forall t (ns :: [Nat]) s - . ( PrimBytes t, Dimensions ns) - => State# s -> (# State# s, MDataFrame s t ns #) -newDataFrame# s0 +newDataFrame# :: + forall (t :: Type) (k :: Type) (ns :: [k]) s + . (PrimBytes t, Dimensions ns) + => State# s -> (# State# s, MDataFrame s t ns #) +newDataFrame# | steps <- cumulDims $ dims @ns , n <- cdTotalDim# steps - , (# s1, mba #) <- newByteArray# (n *# byteSize @t undefined) s0 - = (# s1, MDataFrame# 0# steps mba #) + = if isTrue# (n ==# 0#) + then \s0 -> (# s0, error "Empty DataFrame (DF0)" #) + else \s0 -> case newByteArray# (n *# byteSize @t undefined) s0 of + (# s1, mba #) -> (# s1, MDataFrame# 0# steps mba #) {-# INLINE newDataFrame# #-} -- | Create a new mutable DataFrame. -newPinnedDataFrame# :: forall t (ns :: [Nat]) s - . ( PrimBytes t, Dimensions ns) - => State# s -> (# State# s, MDataFrame s t ns #) -newPinnedDataFrame# s0 +newPinnedDataFrame# :: + forall (t :: Type) (k :: Type) (ns :: [k]) s + . (PrimBytes t, Dimensions ns) + => State# s -> (# State# s, MDataFrame s t ns #) +newPinnedDataFrame# | steps <- cumulDims $ dims @ns , n <- cdTotalDim# steps - , (# s1, mba #) <- newAlignedPinnedByteArray# - (n *# byteSize @t undefined) - (byteAlign @t undefined) s0 - = (# s1, MDataFrame# 0# steps mba #) + = if isTrue# (n ==# 0#) + then \s0 -> (# s0, error "Empty DataFrame (DF0)" #) + else \s0 -> case newAlignedPinnedByteArray# + (n *# byteSize @t undefined) + (byteAlign @t undefined) s0 of + (# s1, mba #) -> (# s1, MDataFrame# 0# steps mba #) {-# INLINE newPinnedDataFrame# #-} +-- | Create a new mutable DataFrame of the same size. +oneMoreDataFrame# :: + forall (t :: Type) (k :: Type) (ns :: [k]) s + . MDataFrame s t ns -> State# s -> (# State# s, MDataFrame s t ns #) +oneMoreDataFrame# mdf@(MDataFrame# off steps mba) s0 + | 0# <- cdTotalDim# steps = (# s0, mdf #) + | (# s1, bs #) <- getSizeofMutableByteArray# mba s0 + , (# s2, mba' #) <- newByteArray# (bs -# off) s1 + = (# s2, MDataFrame# 0# steps mba' #) +{-# INLINE oneMoreDataFrame# #-} + +-- | View a part of a DataFrame. +-- +-- This function does not perform a copy. +-- All changes to a new DataFrame will be reflected in the original DataFrame as well. +-- +-- If any of the dims in @as@ or @b@ is unknown (@a ~ XN m@), +-- then this function is unsafe and can throw an `OutOfDimBounds` exception. +-- Otherwise, its safety is guaranteed by the type system. +subDataFrameView# :: + forall (t :: Type) (k :: Type) + (b :: k) (bi :: k) (bd :: k) + (as :: [k]) (bs :: [k]) (asbs :: [k]) s + . (SubFrameIndexCtx b bi bd, KnownDim bd, ConcatList as (b :+ bs) asbs) + => Idxs (as +: bi) -> MDataFrame s t asbs -> MDataFrame s t (bd :+ bs) +subDataFrameView# ei (MDataFrame# offM stepsM arr) + = MDataFrame# (case offA of I# i -> i) stepsA arr + where + (offA, stepsA) = getOffAndStepsSub (I# offM) stepsM ei (dim @bd) + +-- | View a part of a DataFrame. +-- +-- This function does not perform a copy. +-- All changes to a new DataFrame will be reflected in the original DataFrame as well. +-- +-- This is a simpler version of @subDataFrameView@ that allows +-- to view over one index at a time. +-- +-- If any of the dims in @as@ is unknown (@a ~ XN m@), +-- then this function is unsafe and can throw an `OutOfDimBounds` exception. +-- Otherwise, its safety is guaranteed by the type system. +subDataFrameView'# :: + forall (t :: Type) (k :: Type) (as :: [k]) (bs :: [k]) (asbs :: [k]) s + . ConcatList as bs asbs + => Idxs as -> MDataFrame s t asbs -> MDataFrame s t bs +subDataFrameView'# ei (MDataFrame# offM stepsM arr) + = MDataFrame# (case offA of I# i -> i) stepsA arr + where + (offA, stepsA) = getOffAndSteps (I# offM) stepsM ei + -- | Copy one DataFrame into another mutable DataFrame at specified position. -- --- In contrast to @copyDataFrame'@, this function allows to copy over a range of contiguous --- indices over a single dimension. +-- In contrast to @copyDataFrame'@, this function allows to copy over a range +-- of contiguous indices over a single dimension. -- For example, you can write a 3x4 matrix into a 7x4 matrix, starting at indices 0..3. -copyDataFrame# :: forall (t :: Type) - (b :: Nat) (bi :: Nat) (bd :: Nat) - (as :: [Nat]) (bs :: [Nat]) (asbs :: [Nat]) s - . ( b ~ (bi + bd - 1) - , PrimBytes t - , PrimBytes (DataFrame t (bd :+ bs)) - , ConcatList as (b :+ bs) asbs - ) - => Idxs (as +: bi) -> DataFrame t (bd :+ bs) -> MDataFrame s t asbs - -> State# s -> (# State# s, () #) -copyDataFrame# ei df (MDataFrame# off steps mba) s - | I# i <- cdIx steps ei - = (# writeBytes mba ((off +# i) *# byteSize @t undefined) df s, () #) +-- +-- This function is safe (no `OutOfDimBounds` exception possible). +-- If any of the dims in @as@ is unknown (@a ~ XN m@), +-- you may happen to write data beyond dataframe bounds. +-- In this case, this function does nothing. +-- If (@b ~ XN m@) and (@Idx bi + Dim bd > Dim b@), this function copies only as +-- many elements as fits into the dataframe along this dimension (possibly none). +copyDataFrame# :: + forall (t :: Type) (k :: Type) + (b :: k) (bi :: k) (bd :: k) + (as :: [k]) (bs :: [k]) (asbs :: [k]) s + . ( SubFrameIndexCtx b bi bd, KnownDim bd + , ExactDims bs + , PrimArray t (DataFrame t (bd :+ bs)) + , ConcatList as (b :+ bs) asbs ) + => Idxs (as +: bi) -> DataFrame t (bd :+ bs) -> MDataFrame s t asbs + -> State# s -> (# State# s, () #) +copyDataFrame# ei df (MDataFrame# offM stepsM arrDest) + | elS <- byteSize @t undefined + , Just (I# offDest, stepsB) + <- getOffAndStepsSubM (I# offM) stepsM ei (dim @bd) + , n <- cdTotalDim# stepsB + , isTrue# (n ># 0#) -- is there enough space to write anything? + = withArrayContent + (\e s -> (# fillArray arrDest offDest n e s, () #)) + (\_ offSrc arrSrc s -> + (# copyByteArray# arrSrc (offSrc *# elS) + arrDest (offDest *# elS) (n *# elS) s, () #)) df + | otherwise = \s -> (# s, () #) {-# INLINE copyDataFrame# #-} - {-# ANN copyMDataFrame# "HLint: ignore Use camelCase" #-} -- | Copy one mutable DataFrame into another mutable DataFrame at specified position. -- --- In contrast to @copyMDataFrame'@, this function allows to copy over a range of contiguous --- indices over a single dimension. +-- In contrast to @copyMutableDataFrame'@, this function allows to copy over a range +-- of contiguous indices over a single dimension. -- For example, you can write a 3x4 matrix into a 7x4 matrix, starting at indices 0..3. -copyMDataFrame# :: forall (t :: Type) - (b :: Nat) (bi :: Nat) (bd :: Nat) - (as :: [Nat]) (bs :: [Nat]) (asbs :: [Nat]) s - . ( b ~ (bi + bd - 1) - , PrimBytes t - , ConcatList as (b :+ bs) asbs - ) - => Idxs (as +: bi) -> MDataFrame s t (bd :+ bs) -> MDataFrame s t asbs - -> State# s -> (# State# s, () #) - -- all information for copying in this case is taken from the inside - -- of the data types at the term level; - -- steps zips with indices just fine independently of their types - -- and length. - -- Thus, it's fine here to just unsafeCoerce# an index. -copyMDataFrame# ei df = - copyMDataFrame'# - (unsafeCoerce# ei :: Idxs as) - (unsafeCoerce# df :: MDataFrame s t (b :+ bs)) +-- +-- This function is safe (no `OutOfDimBounds` exception possible). +-- If any of the dims in @as@ is unknown (@a ~ XN m@), +-- you may happen to write data beyond dataframe bounds. +-- In this case, this function does nothing. +-- If (@b ~ XN m@) and (@Idx bi + Dim bd > Dim b@), this function copies only as +-- many elements as fits into the dataframe along this dimension (possibly none). +copyMDataFrame# :: + forall (t :: Type) (k :: Type) + (b :: k) (bi :: k) (bd :: k) + (as :: [k]) (bs :: [k]) (asbs :: [k]) s + . ( SubFrameIndexCtx b bi bd + , ExactDims bs + , PrimBytes t + , ConcatList as (b :+ bs) asbs ) + => Idxs (as +: bi) -> MDataFrame s t (bd :+ bs) -> MDataFrame s t asbs + -> State# s -> (# State# s, () #) +copyMDataFrame# ei (MDataFrame# offA (CumulDims ~(bb:b:_)) arrA) + (MDataFrame# offM stepsM arrM) + | elS <- byteSize @t undefined + , Just (I# offDest, stepsB) + <- getOffAndStepsSubM (I# offM) stepsM ei (unsafeCoerce# (quot bb b)) + , n <- cdTotalDim# stepsB + , isTrue# (n ># 0#) -- is there enough space to write anything? + = \s -> (# copyMutableByteArray# arrA (offA *# elS) + arrM (offDest *# elS) (n *# elS) s + , () #) + | otherwise = \s -> (# s, () #) {-# INLINE copyMDataFrame# #-} {-# ANN copyDataFrame'# "HLint: ignore Use camelCase" #-} -- | Copy one DataFrame into another mutable DataFrame at specified position. -- --- This is a simpler version of @copyDataFrame@ that allows to copy over one index at a time. -copyDataFrame'# :: forall (t :: Type) (as :: [Nat]) (bs :: [Nat]) (asbs :: [Nat]) s - . ( PrimBytes t - , PrimBytes (DataFrame t bs) - , ConcatList as bs asbs - ) - => Idxs as -> DataFrame t bs -> MDataFrame s t asbs - -> State# s -> (# State# s, () #) -copyDataFrame'# ei df (MDataFrame# off steps mba) s - | I# i <- cdIx steps ei - = (# writeBytes mba ((off +# i) *# byteSize @t undefined) df s, () #) +-- This is a simpler version of @copyDataFrame@ that allows +-- to copy over one index at a time. +-- +-- This function is safe (no `OutOfDimBounds` exception possible). +-- If any of the dims in @as@ is unknown (@a ~ XN m@), +-- you may happen to write data beyond dataframe bounds. +-- In this case, this function does nothing. +copyDataFrame'# :: + forall (t :: Type) (k :: Type) (as :: [k]) (bs :: [k]) (asbs :: [k]) s + . ( ExactDims bs + , PrimArray t (DataFrame t bs) + , ConcatList as bs asbs ) + => Idxs as -> DataFrame t bs -> MDataFrame s t asbs + -> State# s -> (# State# s, () #) +copyDataFrame'# ei df (MDataFrame# offM stepsM arrDest) + | elS <- byteSize @t undefined + , Just (I# offDest, stepsA) <- getOffAndStepsM (I# offM) stepsM ei + , n <- cdTotalDim# stepsA + = withArrayContent + (\e s -> (# fillArray arrDest offDest n e s, () #)) + (\_ offSrc arrSrc s -> + (# copyByteArray# arrSrc (offSrc *# elS) + arrDest (offDest *# elS) (n *# elS) s, () #)) df + | otherwise = \s -> (# s, () #) {-# INLINE copyDataFrame'# #-} {-# ANN copyMDataFrame'# "HLint: ignore Use camelCase" #-} -- | Copy one mutable DataFrame into another mutable DataFrame at specified position. -- --- This is a simpler version of @copyMDataFrame@ that allows to copy over one index at a time. -copyMDataFrame'# :: forall (t :: Type) (as :: [Nat]) (bs :: [Nat]) (asbs :: [Nat]) s - . ( PrimBytes t - , ConcatList as bs asbs - ) - => Idxs as -> MDataFrame s t bs -> MDataFrame s t asbs - -> State# s -> (# State# s, () #) -copyMDataFrame'# ei (MDataFrame# offA stepsA arrA) (MDataFrame# offM stepsM arrM) s +-- This is a simpler version of @copyMutableDataFrame@ that allows +-- to copy over one index at a time. +-- +-- This function is safe (no `OutOfDimBounds` exception possible). +-- If any of the dims in @as@ is unknown (@a ~ XN m@), +-- you may happen to write data beyond dataframe bounds. +-- In this case, this function does nothing. +copyMDataFrame'# :: + forall (t :: Type) (k :: Type) (as :: [k]) (bs :: [k]) (asbs :: [k]) s + . (ExactDims bs, PrimBytes t, ConcatList as bs asbs) + => Idxs as -> MDataFrame s t bs -> MDataFrame s t asbs + -> State# s -> (# State# s, () #) +copyMDataFrame'# ei (MDataFrame# offA stepsA arrA) (MDataFrame# offM stepsM arrM) | elS <- byteSize @t undefined , lenA <- cdTotalDim# stepsA - , I# i <- cdIx stepsM ei - = (# copyMutableByteArray# arrA (offA *# elS) - arrM ((offM +# i) *# elS) (lenA *# elS) s - , () #) + , Just (I# i) <- cdIxM stepsM ei + = \s -> (# copyMutableByteArray# + arrA (offA *# elS) + arrM ((offM +# i) *# elS) (lenA *# elS) s, () #) + | otherwise = \s -> (# s, () #) {-# INLINE copyMDataFrame'# #-} +-- | Copy one DataFrame into another mutable DataFrame by offset in +-- primitive elements. +-- +-- This is a low-level copy function; you have to keep in mind the row-major +-- layout of Mutable DataFrames. Offset bounds are not checked. +-- You will get an undefined behavior if you write beyond the DataFrame bounds. +copyDataFrameOff# :: + forall (t :: Type) (k :: Type) (as :: [k]) (bs :: [k]) (asbs :: [k]) s + . ( Dimensions bs + , PrimArray t (DataFrame t bs) + , ConcatList as bs asbs ) + => Int -> DataFrame t bs -> MDataFrame s t asbs + -> State# s -> (# State# s, () #) +copyDataFrameOff# (I# off) df (MDataFrame# offM _ arrDest) + | elS <- byteSize @t undefined + , offDest <- offM +# off + = withArrayContent + (\e s -> + (# fillArray arrDest offDest + (case totalDim (dims @bs) of W# n -> word2Int# n) e s, () #)) + (\steps offSrc arrSrc s -> + (# copyByteArray# arrSrc (offSrc *# elS) + arrDest (offDest *# elS) + (cdTotalDim# steps *# elS) s, () #)) df +{-# INLINE copyDataFrameOff# #-} + +-- | Copy one mutable DataFrame into another mutable DataFrame by offset in +-- primitive elements. +-- +-- This is a low-level copy function; you have to keep in mind the row-major +-- layout of Mutable DataFrames. Offset bounds are not checked. +-- You will get an undefined behavior if you write beyond the DataFrame bounds. +copyMDataFrameOff# :: + forall (t :: Type) (k :: Type) (as :: [k]) (bs :: [k]) (asbs :: [k]) s + . (ExactDims bs, PrimBytes t, ConcatList as bs asbs) + => Int -> MDataFrame s t bs -> MDataFrame s t asbs + -> State# s -> (# State# s, () #) +copyMDataFrameOff# (I# off) (MDataFrame# offA stepsA arrA) + (MDataFrame# offM _ arrM) + | elS <- byteSize @t undefined + , lenA <- cdTotalDim# stepsA + = \s -> (# copyMutableByteArray# + arrA (offA *# elS) + arrM ((offM +# off) *# elS) (lenA *# elS) s, () #) +{-# INLINE copyMDataFrameOff# #-} + -- | Make a mutable DataFrame immutable, without copying. -unsafeFreezeDataFrame# :: forall (t :: Type) (ns :: [Nat]) s - . PrimArray t (DataFrame t ns) - => MDataFrame s t ns - -> State# s -> (# State# s, DataFrame t ns #) -unsafeFreezeDataFrame# (MDataFrame# offM steps arrM) s1 - | (# s2, arrA #) <- unsafeFreezeByteArray# arrM s1 - = (# s2, fromElems steps offM arrA #) +unsafeFreezeDataFrame# :: + forall (t :: Type) (k :: Type) (ns :: [k]) s + . PrimArray t (DataFrame t ns) + => MDataFrame s t ns + -> State# s -> (# State# s, DataFrame t ns #) +unsafeFreezeDataFrame# (MDataFrame# offM steps arrM) s0 + | 0# <- cdTotalDim# steps + = (# s0, error "Empty DataFrame (DF0)" #) + | (# s1, arrA #) <- unsafeFreezeByteArray# arrM s0 + = (# s1, fromElems steps offM arrA #) {-# INLINE unsafeFreezeDataFrame# #-} -- | Copy content of a mutable DataFrame into a new immutable DataFrame. -freezeDataFrame# :: forall (t :: Type) (ns :: [Nat]) s - . PrimArray t (DataFrame t ns) - => MDataFrame s t ns -> State# s -> (# State# s, DataFrame t ns #) +freezeDataFrame# :: + forall (t :: Type) (k :: Type) (ns :: [k]) s + . PrimArray t (DataFrame t ns) + => MDataFrame s t ns -> State# s -> (# State# s, DataFrame t ns #) freezeDataFrame# (MDataFrame# offM steps arrM) s0 + | 0# <- cdTotalDim# steps + = (# s0, error "Empty DataFrame (DF0)" #) | elS <- byteSize @t undefined , n <- cdTotalDim# steps , (# s1, mba #) <- newByteArray# (n *# elS) s0 , s2 <- copyMutableByteArray# arrM (offM *# elS) mba 0# (n *# elS) s1 , (# s3, arrA #) <- unsafeFreezeByteArray# mba s2 - = (# s3, fromElems steps 0# arrA #) + = (# s3, fromElems steps 0# arrA #) {-# INLINE freezeDataFrame# #-} -- | Create a new mutable DataFrame and copy content of immutable one in there. -thawDataFrame# :: forall (t :: Type) (ns :: [Nat]) s - . ( Dimensions ns - , PrimBytes (DataFrame t ns)) - => DataFrame t ns -> State# s -> (# State# s, MDataFrame s t ns #) +thawDataFrame# :: + forall (t :: Type) (k :: Type) (ns :: [k]) s + . (Dimensions ns, PrimArray t (DataFrame t ns)) + => DataFrame t ns -> State# s -> (# State# s, MDataFrame s t ns #) thawDataFrame# df s0 - | arrA <- getBytes df - , boff <- byteOffset df - , bsize <- byteSize df - , steps <- cumulDims $ dims @ns - , (# s1, arrM #) <- newByteArray# bsize s0 - , s2 <- copyByteArray# arrA boff arrM 0# bsize s1 - = (# s2, MDataFrame# 0# steps arrM #) + | nw == 0 + = (# s0, error "Empty DataFrame (DF0)" #) + | bSize <- case nw of W# w -> byteSize @t undefined *# word2Int# w + , (# s1, arrM #) <- newByteArray# bSize s0 + , r <- MDataFrame# 0# steps arrM + , (# s2, _ #) <- copyDataFrameOff# 0 df r s1 + = (# s2, r #) + where + nw = cdTotalDim steps + steps = getSteps (dims @ns) df {-# INLINE thawDataFrame# #-} -- | Create a new mutable DataFrame and copy content of immutable one in there. -- The result array is pinned and aligned. -thawPinDataFrame# :: forall (t :: Type) (ns :: [Nat]) s - . ( Dimensions ns - , PrimBytes (DataFrame t ns)) - => DataFrame t ns -> State# s -> (# State# s, MDataFrame s t ns #) +thawPinDataFrame# :: + forall (t :: Type) (k :: Type) (ns :: [k]) s + . (Dimensions ns, PrimArray t (DataFrame t ns)) + => DataFrame t ns -> State# s -> (# State# s, MDataFrame s t ns #) thawPinDataFrame# df s0 - | arrA <- getBytes df - , boff <- byteOffset df - , bsize <- byteSize df - , steps <- cumulDims $ dims @ns - , (# s1, arrM #) <- newAlignedPinnedByteArray# bsize (byteAlign df) s0 - , s2 <- copyByteArray# arrA boff arrM 0# bsize s1 - = (# s2, MDataFrame# 0# steps arrM #) + | nw == 0 + = (# s0, error "Empty DataFrame (DF0)" #) + | bSize <- case nw of W# w -> byteSize @t undefined *# word2Int# w + , (# s1, arrM #) <- newAlignedPinnedByteArray# bSize (byteAlign @t undefined) s0 + , r <- MDataFrame# 0# steps arrM + , (# s2, _ #) <- copyDataFrameOff# 0 df r s1 + = (# s2, r #) + where + nw = cdTotalDim steps + steps = getSteps (dims @ns) df {-# INLINE thawPinDataFrame# #-} -- | UnsafeCoerces an underlying byte array. -unsafeThawDataFrame# :: forall (t :: Type) (ns :: [Nat]) s - . ( Dimensions ns - , PrimBytes (DataFrame t ns), PrimBytes t) - => DataFrame t ns - -> State# s -> (# State# s, MDataFrame s t ns #) -unsafeThawDataFrame# df s0 - | elS <- byteSize @t undefined - , arrA <- getBytes df - , boff <- byteOffset df - , steps <- cumulDims $ dims @ns - = (# s0 - , MDataFrame# (quotInt# boff elS) steps (unsafeCoerce# arrA) - #) +unsafeThawDataFrame# :: + forall (t :: Type) (k :: Type) (ns :: [k]) s + . (Dimensions ns, PrimArray t (DataFrame t ns)) + => DataFrame t ns + -> State# s -> (# State# s, MDataFrame s t ns #) +unsafeThawDataFrame# = withArrayContent f g + where + f :: t -> State# s -> (# State# s, MDataFrame s t ns #) + f e s0 + | steps <- cumulDims (dims @ns) + , n <- cdTotalDim# steps + , (# s1, arrM #) <- newByteArray# (n *# byteSize @t undefined) s0 + = (# fillArray arrM 0# n e s1, MDataFrame# 0# steps arrM #) + g :: CumulDims -> Int# -> ByteArray# -> State# s -> (# State# s, MDataFrame s t ns #) + g steps off ba s0 + = (# s0, MDataFrame# off steps (unsafeCoerce# ba) #) {-# INLINE unsafeThawDataFrame# #-} +-- | Given two continuations @f@ and @g@. +-- If the input DataFrame is a single broadcast value, use it in @f@. +-- Otherwise, create a new mutable DataFrame and copy content of immutable one +-- in there; then use it in @g@. +-- +-- This function is useful when @thawDataFrame@ cannot be used due to +-- @Dimensions ns@ constraint being not available. +withThawDataFrame# :: + forall (t :: Type) (k :: Type) (ns :: [k]) (r :: Type) s + . PrimArray t (DataFrame t ns) + => (t -> State# s -> (# State# s, r #)) -- ^ f + -> (MDataFrame s t ns -> State# s -> (# State# s, r #)) -- ^ g + -> DataFrame t ns + -> State# s -> (# State# s, r #) +withThawDataFrame# f g = withArrayContent f g' + where + g' :: CumulDims -> Int# -> ByteArray# -> State# s -> (# State# s, r #) + g' steps eOff arrA s0 = case cdTotalDim# steps of + 0# -> g (error "Empty DataFrame (DF0)") s0 + elems + | elS <- byteSize @t undefined + , bsize <- elS *# elems + , (# s1, arrM #) <- newByteArray# bsize s0 + , s2 <- copyByteArray# arrA (eOff *# elS) arrM 0# bsize s1 + -> g (MDataFrame# 0# steps arrM) s2 +{-# INLINE withThawDataFrame# #-} --- | Write a single element at the specified element offset -writeDataFrameOff# :: forall (t :: Type) (ns :: [Nat]) s - . PrimBytes t - => MDataFrame s t ns -> Int# -> t -> State# s -> (# State# s, () #) -writeDataFrameOff# (MDataFrame# off _ mba) i x s +-- | Write a single element at the specified element offset. +-- +-- This is a low-level write function; you have to keep in mind the row-major +-- layout of Mutable DataFrames. Offset bounds are not checked. +-- You will get an undefined behavior if you write beyond the DataFrame bounds. +writeDataFrameOff# :: + forall (t :: Type) (k :: Type) (ns :: [k]) s + . PrimBytes (DataFrame t ('[] :: [k])) + => MDataFrame s t ns -> Int -> DataFrame t ('[] :: [k]) + -> State# s -> (# State# s, () #) +writeDataFrameOff# (MDataFrame# off _ mba) (I# i) x s = (# writeArray mba (off +# i) x s, () #) {-# INLINE writeDataFrameOff# #-} --- | Write a single element at the specified index -writeDataFrame# :: forall (t :: Type) (ns :: [Nat]) s - . PrimBytes t - => MDataFrame s t ns -> Idxs ns -> t -> State# s -> (# State# s, () #) +-- | Write a single element at the specified index. +-- +-- This function is safe (no `OutOfDimBounds` exception possible). +-- If any of the dims in @ns@ is unknown (@n ~ XN m@), +-- you may happen to write data beyond dataframe bounds. +-- In this case, this function does nothing. +writeDataFrame# :: + forall (t :: Type) (k :: Type) (ns :: [k]) s + . PrimBytes (DataFrame t ('[] :: [k])) + => MDataFrame s t ns -> Idxs ns -> DataFrame t ('[] :: [k]) + -> State# s -> (# State# s, () #) writeDataFrame# mdf@(MDataFrame# _ st _) ei - | I# i <- cdIx st ei = writeDataFrameOff# mdf i + | Just off <- (cdIxM st ei) + = writeDataFrameOff# mdf off + | otherwise = const (\s -> (# s, () #)) {-# INLINE writeDataFrame# #-} --- | Read a single element at the specified element offset -readDataFrameOff# :: forall (t :: Type) (ns :: [Nat]) s - . PrimBytes t - => MDataFrame s t ns -> Int# -> State# s -> (# State# s, t #) -readDataFrameOff# (MDataFrame# off _ mba) i = readArray mba (off +# i) +-- | Read a single element at the specified element offset. +-- +-- This is a low-level read function; you have to keep in mind the row-major +-- layout of Mutable DataFrames. Offset bounds are not checked. +-- You will get an undefined behavior if you read beyond the DataFrame bounds. +readDataFrameOff# :: + forall (t :: Type) (k :: Type) (ns :: [k]) s + . PrimBytes (DataFrame t ('[] :: [k])) + => MDataFrame s t ns -> Int + -> State# s -> (# State# s, DataFrame t ('[] :: [k]) #) +readDataFrameOff# (MDataFrame# off _ mba) (I# i) + = readArray @(DataFrame t ('[] :: [k])) mba (off +# i) {-# INLINE readDataFrameOff# #-} --- | Read a single element at the specified index -readDataFrame# :: forall (t :: Type) (ns :: [Nat]) s - . PrimBytes t - => MDataFrame s t ns -> Idxs ns -> State# s -> (# State# s, t #) +-- | Read a single element at the specified index. +-- +-- If any of the dims in @ns@ is unknown (@n ~ XN m@), +-- then this function is unsafe and can throw an `OutOfDimBounds` exception. +-- Otherwise, its safety is guaranteed by the type system. +readDataFrame# :: + forall (t :: Type) (k :: Type) (ns :: [k]) s + . PrimBytes (DataFrame t ('[] :: [k])) + => MDataFrame s t ns -> Idxs ns + -> State# s -> (# State# s, DataFrame t ('[] :: [k]) #) readDataFrame# mdf@(MDataFrame# _ st _) ei - | I# i <- cdIx st ei = readDataFrameOff# mdf i + = readDataFrameOff# mdf (cdIx st ei) {-# INLINE readDataFrame# #-} -- | Allow arbitrary operations on a pointer to the beginning of the data. -- Only possible with @RealWord@ state (thus, in @IO@) due to semantics of -- @touch#@ operation that keeps the data from being garbage collected. -withDataFramePtr# :: forall (t :: Type) (ns :: [Nat]) (r :: Type) - . PrimBytes t - => MDataFrame RealWorld t ns - -> ( Addr# -> State# RealWorld -> (# State# RealWorld, r #) ) - -> State# RealWorld -> (# State# RealWorld, r #) +withDataFramePtr# :: + forall (t :: Type) (k :: Type) (ns :: [k]) (r :: Type) + . PrimBytes t + => MDataFrame RealWorld t ns + -> (Addr# -> State# RealWorld -> (# State# RealWorld, r #)) + -> State# RealWorld -> (# State# RealWorld, r #) withDataFramePtr# (MDataFrame# off _ mba) k s0 | (# s1, a #) <- unsafeFreezeByteArray# mba s0 , (# s2, r #) <- k ( byteArrayContents# a `plusAddr#` (off *# byteSize @t undefined) ) s1 - = (# touch# mba s2, r #) + = (# touch# mba s2, r #) +{-# INLINE withDataFramePtr# #-} -- | Check if the byte array wrapped by this DataFrame is pinned, -- which means cannot be relocated by GC. -isDataFramePinned# :: forall (t :: Type) (ns :: [Nat]) s - . MDataFrame s t ns -> Bool +isDataFramePinned# :: + forall (t :: Type) (k :: Type) (ns :: [k]) s . MDataFrame s t ns -> Bool isDataFramePinned# (MDataFrame# _ _ mba) = isTrue# (isMutableByteArrayPinned# mba) + +-- | Get cumulative dimensions @ns@ of a @MDataFrame s t ns@ +getDataFrameSteps# :: + forall (t :: Type) (k :: Type) (ns :: [k]) s + . MDataFrame s t ns -> CumulDims +getDataFrameSteps# (MDataFrame# _ c _) = c +{-# INLINE getDataFrameSteps# #-} + +-- | Fill a mutable byte array with the same single element +fillArray :: PrimBytes t + => MutableByteArray# s + -> Int# -- ^ Offset in elements + -> Int# -- ^ Number of elements + -> t + -> State# s -> State# s +fillArray mba off n e + = let lim = off +# n + go i s | isTrue# (i >=# lim) = s + | otherwise = go (i +# 1#) (writeArray mba i e s) + in go off +{-# INLINE fillArray #-} diff --git a/src/Numeric/DataFrame/Internal/PrimArray.hs b/src/Numeric/DataFrame/Internal/PrimArray.hs --- a/src/Numeric/DataFrame/Internal/PrimArray.hs +++ b/src/Numeric/DataFrame/Internal/PrimArray.hs @@ -1,27 +1,40 @@ {-# LANGUAGE CPP #-} {-# LANGUAGE DataKinds #-} +{-# LANGUAGE ExplicitForAll #-} +{-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE FunctionalDependencies #-} +{-# LANGUAGE GADTs #-} {-# LANGUAGE MagicHash #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE PolyKinds #-} +{-# LANGUAGE TypeOperators #-} {-# LANGUAGE UnboxedTuples #-} module Numeric.DataFrame.Internal.PrimArray ( PrimArray (..), CumulDims (..) - , cumulDims, cdTotalDim, cdTotalDim#, cdIx + , cumulDims, cdTotalDim, cdTotalDim# + , cdIx, cdIxSub, getOffAndSteps, getOffAndStepsSub + , cdIxM, getOffAndStepsM, getOffAndStepsSubM , ixOff, unsafeFromFlatList, getSteps, fromSteps + , withArrayContent, fromElems, broadcast ) where -import Data.Monoid as Mon (Monoid (..)) -import Data.Semigroup as Sem (Semigroup (..)) -import GHC.Base (ByteArray#, Int (..), Int#, Word (..), word2Int#) +import Control.Arrow ((***)) +import Data.Monoid as Mon (Monoid (..)) +import Data.Semigroup as Sem (Semigroup (..)) +import GHC.Base hiding (foldr) + -- (ByteArray#, Int (..), Int#, RuntimeRep (..), Type, + -- Word (..), inline, runRW#, touch#, word2Int#, (+#)) +import GHC.Exts (TYPE) import Numeric.Dimensions import Numeric.PrimBytes + -- | Given @Dims ns@, @CumulativeDims@ is a list of length @Length ns + 1@; -- which cumulative @totalDim@ accumulated on the right. -- In particular, its first element is @totalDim ds@, -- its last element is always is always @1@. newtype CumulDims = CumulDims { unCumulDims :: [Word] } + deriving Show instance Sem.Semigroup CumulDims where CumulDims as <> CumulDims bs = CumulDims $ map (head bs *) (init as) ++ bs @@ -49,16 +62,128 @@ -- -- Note, you can take offset of subspace with CumulDims of larger space -- - very convenient! +-- +-- If any of the dims in @ns@ is unknown (@n ~ XN m@), +-- then this function is unsafe and can throw an `OutOfDimBounds` exception. +-- Otherwise, its safety is guaranteed by the type system. cdIx :: CumulDims -> Idxs ns -> Int -cdIx ~(CumulDims ~(_:steps)) - = fromIntegral . sum . zipWith (*) steps . listIdxs +cdIx steps + = fromIntegral . fst . getOffAndSteps' 0 0 (unCumulDims steps) . listIdxs +-- | Calculate offset of an Idxs. +-- +-- Also check if the last index plus dimVal of subN is not bigger than the +-- corresponding dim inside CumulDims; throw an `OutOfDimBounds` otherwise. +-- +-- If any of the dims in @ns@ is unknown (@n ~ XN m@), +-- then this function is unsafe and can throw an `OutOfDimBounds` exception. +-- Otherwise, its safety is guaranteed by the type system. +cdIxSub :: CumulDims -> Idxs (ns +: idxN) -> Dim subN -> Int +cdIxSub steps idxs d + = fromIntegral + . fst . getOffAndSteps' (dimVal d) 0 (unCumulDims steps) $ listIdxs idxs + +-- | Calculate offset of an Idxs and return remaining CumulDims. +-- +-- If any of the dims in @ns@ is unknown (@n ~ XN m@), +-- then this function is unsafe and can throw an `OutOfDimBounds` exception. +-- Otherwise, its safety is guaranteed by the type system. +getOffAndSteps :: Int -- ^ Initial offset + -> CumulDims -> Idxs ns -> (Int, CumulDims) +getOffAndSteps off0 steps + = (fromIntegral *** CumulDims) + . getOffAndSteps' 0 (fromIntegral off0) (unCumulDims steps) . listIdxs + +-- | Calculate offset of an Idxs and return remaining CumulDims. +-- +-- Also check if the last index plus dimVal of subN is not bigger than the +-- corresponding dim inside CumulDims; throw an `OutOfDimBounds` otherwise. +-- +-- If any of the dims in @ns@ is unknown (@n ~ XN m@), +-- then this function is unsafe and can throw an `OutOfDimBounds` exception. +-- Otherwise, its safety is guaranteed by the type system. +getOffAndStepsSub :: Int -- ^ Initial offset + -> CumulDims -> Idxs (ns +: idxN) + -> Dim subN -> (Int, CumulDims) +getOffAndStepsSub off0 steps idxs d + = (fromIntegral *** CumulDims) + . getOffAndSteps' (dimVal d) (fromIntegral off0) (unCumulDims steps) + $ listIdxs idxs + + +getOffAndSteps' :: Word -> Word -> [Word] -> [Word] -> (Word, [Word]) +getOffAndSteps' sub0 off0 steps0 is0 = go sub0 off0 steps0 is0 + where +#ifndef UNSAFE_INDICES + ds0 = case fromSteps (CumulDims steps0) of SomeDims x -> listDims x +#endif + go :: Word -> Word -> [Word] -> [Word] -> (Word, [Word]) + go 0 off steps [] = (off, steps) + go sub off ~(steps@(s:_)) [] = (off, sub*s : steps) +#ifndef UNSAFE_INDICES + go sub _ ~(bs:(s:_)) [i] + | b <- quot bs s + , sub > 0 && i + sub > b + = outOfDimBoundsNoCallStack + "{Calculating SubDataFrame offset}" + i b (Just sub) (Just (ds0, is0)) + go _ _ ~(bs:(s:_)) (i:_) + | b <- quot bs s + , i >= b + = outOfDimBoundsNoCallStack + "{Calculating SubDataFrame offset}" + i b Nothing (Just (ds0, is0)) +#endif + go sub off ~(_:steps@(s:_)) (i:ixs) + = go sub (off + i*s) steps ixs +{-# INLINE getOffAndSteps' #-} + + + +-- | Same as `cdIx`, but safe; returns @Nothing@ if out of bounds. +cdIxM :: CumulDims -> Idxs ns -> Maybe Int +cdIxM steps + = fmap (fromIntegral . fst) + . getOffAndStepsM' 0 0 (unCumulDims steps) . listIdxs + +-- | Same as `getOffAndSteps`, but safe; returns @Nothing@ if out of bounds. +-- Trims the first (slicing) dimension of the returned CumulDims to fit +-- the original dataframe if necessary. +getOffAndStepsM :: Int -- ^ Initial offset + -> CumulDims -> Idxs ns -> Maybe (Int, CumulDims) +getOffAndStepsM off0 steps + = fmap (fromIntegral *** CumulDims) + . getOffAndStepsM' 0 (fromIntegral off0) (unCumulDims steps) . listIdxs + +-- | Same as `getOffAndStepsSub`, but safe; returns @Nothing@ if out of bounds. +-- Trims the first (slicing) dimension of the returned CumulDims to fit +-- the original dataframe if necessary. +getOffAndStepsSubM :: Int -- ^ Initial offset + -> CumulDims -> Idxs (ns +: idxN) + -> Dim subN -> Maybe (Int, CumulDims) +getOffAndStepsSubM off0 steps idxs d + = fmap (fromIntegral *** CumulDims) + . getOffAndStepsM' (dimVal d) (fromIntegral off0) (unCumulDims steps) + $ listIdxs idxs + +getOffAndStepsM' :: Word -> Word -> [Word] -> [Word] -> Maybe (Word, [Word]) +getOffAndStepsM' = go + where + go :: Word -> Word -> [Word] -> [Word] -> Maybe (Word, [Word]) + go _ off steps [] = Just (off, steps) + go sub off ~(bs:steps@(s:_)) (i:ixs) + | is >= bs = Nothing + | not (null ixs) = go sub (off + is) steps ixs + | sub == 0 = Just (off + is, steps) + | otherwise = Just (off + is, min (bs - is) (sub*s) : steps) + where + is = i*s +{-# INLINE getOffAndStepsM' #-} + -- | Try to get @CumulDims@ from an array, -- and create it using @Dims@ if failed. getSteps :: PrimArray t a => Dims (ns :: [k]) -> a -> CumulDims -getSteps dds df = case uniqueOrCumulDims df of - Left _ -> cumulDims dds - Right ds -> ds +getSteps dds = withArrayContent (const $ cumulDims dds) (\cd _ _ -> cd) {-# INLINE getSteps #-} -- | Get @Dims@ by "de-accumulating" @CumulDims@. @@ -75,9 +200,14 @@ class PrimBytes t => PrimArray t a | a -> t where -- | Broadcast element into array - broadcast :: t -> a + -- + -- Warning: do not use this function at the call site; use `broadcast` + -- instead. Otherwise you will miss some rewrite rules. + broadcast# :: t -> a -- | Index an array given an offset ix# :: Int# -> a -> t + ix# i = withArrayContent id (\_ off arr -> indexArray arr (off +# i)) + {-# INLINE ix# #-} -- | Generate an array using an accumulator funtion gen# :: CumulDims -- ^ Dimensionality of the result array; @@ -92,17 +222,26 @@ -- types and is not checked at runtime. -> Int# -> t -> a -> a - -- | If the array represented as a single broadcasted value, - -- return this value. - -- Otherwise, return full array content: + -- | If the array is represented as a single broadcasted value, return this + -- this value. Otherwise, return the full array content: -- @CumulDims@, array offset (elements), byte array with the content. - arrayContent# :: a -> (# t | (# CumulDims, Int#, ByteArray# #) #) + -- + -- Warning: never use this function directly. Use `withArrayContent` instead. + -- There is a bug in GHC 8.6, such that certain optimizations + -- (probably, instance specialization/rewrite rules) break the code, + -- which is only observable at runtime. The effect is that the + -- content of a `ByteArray#` becomes a garbage. The workaround is + -- to use a non-inlinable wrapper to disable these optimizations. + -- In addition, the wrapper function has some rewrite rules, which + -- can potentially improve performance with other GHC versions. + withArrayContent# :: forall (rep :: RuntimeRep) (r :: TYPE rep) + . (t -> r) + -> (CumulDims -> Int# -> ByteArray# -> r) + -> a -> r -- | Offset of an array as a number of elements offsetElems :: a -> Int# - offsetElems a = case arrayContent# a of - (# _ | #) -> 0# - (# | (# _, o, _ #) #) -> o + offsetElems a = withArrayContent (\_ f -> f 0#) (const $ \o _ f -> f o) a (\i -> i) {-# INLINE offsetElems #-} -- | Normally, this returns a cumulative @totalDim@s. @@ -114,18 +253,25 @@ -- Note, this function returns the only unique element only if it is -- a such by construction (there is no equality checks involved). uniqueOrCumulDims :: a -> Either t CumulDims - uniqueOrCumulDims a = case arrayContent# a of - (# x | #) -> Left x - (# | (# cd, _, _ #) #) -> Right cd + uniqueOrCumulDims = withArrayContent Left (\cd _ _ -> Right cd) {-# INLINE uniqueOrCumulDims #-} - -- | Get array by its offset and cumulative dims in a ByteArray. + -- | Define an array by its offset and cumulative dims in a ByteArray. -- Both offset and dims are given in element number (not in bytes). -- - -- It is better to use this function instead of @fromBytes@ to avoid - -- recalculating @CumulDims@ for implementations that require it. - fromElems :: CumulDims -> Int# -> ByteArray# -> a + -- Warning: never use this function directly. Use `fromElems` instead. + -- There is a bug in GHC 8.6, such that certain optimizations + -- (probably, instance specialization/rewrite rules) break the code, + -- which is only observable at runtime. The effect is that the + -- content of a `ByteArray#` becomes a garbage. The workaround is + -- to use a non-inlinable wrapper to disable these optimizations. + -- In addition, the wrapper function has some rewrite rules, which + -- can potentially improve performance with other GHC versions. + fromElems# :: CumulDims -> Int# -> ByteArray# -> a +{-# WARNING fromElems# "Please, use fromElems instead." #-} +{-# WARNING withArrayContent# "Please, use withArrayContent instead." #-} + -- | Index array by an integer offset (starting from 0). ixOff :: PrimArray t a => Int -> a -> t ixOff (I# i) = ix# i @@ -138,3 +284,51 @@ where f [] = (# [], x0 #) f (x:xs) = (# xs, x #) + +-- | If the array is represented as a single broadcasted value, return this +-- this value. Otherwise, return the full array content: +-- @CumulDims@, array offset (elements), byte array with the content. +withArrayContent :: forall (t :: Type) (a :: Type) + (rep :: RuntimeRep) (r :: TYPE rep) + . PrimArray t a + => (t -> r) + -> (CumulDims -> Int# -> ByteArray# -> r) + -> a -> r +withArrayContent = withArrayContent# +#if __GLASGOW_HASKELL__ == 806 +{-# NOINLINE withArrayContent #-} +#else +{-# INLINE[1] withArrayContent #-} +#endif + +-- | Define an array by its offset and cumulative dims in a ByteArray. +-- Both offset and dims are given in element number (not in bytes). +-- +-- It is better to use this function instead of @fromBytes@ to avoid +-- recalculating @CumulDims@ for implementations that require it. +fromElems :: forall (t :: Type) (a :: Type) + . PrimArray t a => CumulDims -> Int# -> ByteArray# -> a +fromElems = fromElems# +#if __GLASGOW_HASKELL__ == 806 +{-# NOINLINE fromElems #-} +#else +{-# INLINE[1] fromElems #-} +#endif + +-- | Broadcast element into array +broadcast :: forall (t :: Type) (a :: Type) + . PrimArray t a => t -> a +broadcast = broadcast# +{-# INLINE[1] broadcast #-} + +{-# RULES + +"withArrayContent/id" + withArrayContent broadcast fromElems = id + +"withArrayContent+fromElems" forall f g cd off ba . + withArrayContent f g (fromElems cd off ba) = g cd off ba + +"withArrayContent+broadcast" forall f g e . + withArrayContent f g (broadcast e) = f e + #-} diff --git a/src/Numeric/DataFrame/ST.hs b/src/Numeric/DataFrame/ST.hs --- a/src/Numeric/DataFrame/ST.hs +++ b/src/Numeric/DataFrame/ST.hs @@ -2,14 +2,18 @@ {-# LANGUAGE ExistentialQuantification #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE InstanceSigs #-} {-# LANGUAGE MagicHash #-} {-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE PatternSynonyms #-} {-# LANGUAGE PolyKinds #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeApplications #-} {-# LANGUAGE TypeFamilies #-} {-# LANGUAGE TypeInType #-} {-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE ViewPatterns #-} ----------------------------------------------------------------------------- -- | -- Module : Numeric.DataFrame.ST @@ -22,203 +26,344 @@ module Numeric.DataFrame.ST ( STDataFrame (XSTFrame), SomeSTDataFrame (..) - , newDataFrame, newPinnedDataFrame + , castDataFrame + , newDataFrame, newPinnedDataFrame, oneMoreDataFrame + , subDataFrameView, subDataFrameView' , copyDataFrame, copyMutableDataFrame , copyDataFrame', copyMutableDataFrame' + , copyDataFrameOff, copyMutableDataFrameOff , freezeDataFrame, unsafeFreezeDataFrame - , thawDataFrame, thawPinDataFrame, unsafeThawDataFrame + , thawDataFrame, thawPinDataFrame, unsafeThawDataFrame, withThawDataFrame , writeDataFrame, writeDataFrameOff , readDataFrame, readDataFrameOff - , isDataFramePinned + , isDataFramePinned, getDataFrameSteps ) where - -import GHC.Base -import GHC.ST (ST (..)) +import GHC.ST (ST (..)) +import Data.Coerce +import Data.Kind import Numeric.DataFrame.Internal.Mutable import Numeric.DataFrame.Internal.PrimArray import Numeric.DataFrame.Type import Numeric.Dimensions +import Unsafe.Coerce -- | Mutable DataFrame that lives in ST. -- Internal representation is always a MutableByteArray. -data family STDataFrame s (t :: Type) (ns :: [k]) - --- | Pure wrapper on a mutable byte array -newtype instance STDataFrame s t (ns :: [Nat]) = STDataFrame (MDataFrame s t (ns :: [Nat])) +newtype STDataFrame s (t :: Type) (ns :: [k]) = STDataFrame (MDataFrame s t (ns :: [k])) -- | Data frame with some dimensions missing at compile time. -- Pattern-match against its constructor to get a Nat-indexed mutable data frame. -data instance STDataFrame s t (xs :: [XNat]) - = forall (ns :: [Nat]) . Dimensions ns - => XSTFrame (STDataFrame s t ns) +pattern XSTFrame :: forall s (t :: Type) (xns :: [XNat]) . () + => forall (ns :: [Nat]) . (FixedDims xns ns, Dimensions ns) + => STDataFrame s t ns -> STDataFrame s t xns +pattern XSTFrame df <- (mkXSTFramePat -> XSTFramePat df) + where + XSTFrame = castDataFrame +{-# COMPLETE XSTFrame #-} +data XSTFramePat s (t :: Type) (xns :: [XNat]) + = forall (ns :: [Nat]) . (FixedDims xns ns, Dimensions ns) + => XSTFramePat (STDataFrame s t ns) + +mkXSTFramePat :: forall s (t :: Type) (xns :: [XNat]) + . STDataFrame s t xns -> XSTFramePat s t xns +mkXSTFramePat df + | SomeDims ds <- fromSteps (getDataFrameSteps df) + , XDims (Dims :: Dims ns) <- (unsafeCoerce ds) :: Dims xns + = XSTFramePat @s @t @xns @ns (unsafeCoerce df) + | otherwise + = error "XSTFrame pattern: impossible args" + -- | Mutable DataFrame of unknown dimensionality data SomeSTDataFrame s (t :: Type) = forall (ns :: [Nat]) . Dimensions ns => SomeSTDataFrame (STDataFrame s t ns) +-- | Allow coercing between @XNat@-indexed and @Nat@-indexed Mutable DataFrames. +castDataFrame :: + forall (t :: Type) (xns :: [XNat]) (ns :: [Nat]) s + . FixedDims xns ns + => STDataFrame s t ns -> STDataFrame s t xns +castDataFrame = coerce (castDataFrame# @t @xns @ns) +{-# INLINE castDataFrame #-} + -- | Create a new mutable DataFrame. -newDataFrame :: forall t (ns :: [Nat]) s - . ( PrimBytes t, Dimensions ns) - => ST s (STDataFrame s t ns) -newDataFrame = STDataFrame <$> ST (newDataFrame# @t @ns) +newDataFrame :: + forall (t :: Type) ns s + . (PrimBytes t, Dimensions ns) => ST s (STDataFrame s t ns) +newDataFrame = coerce (newDataFrame# @t @_ @ns) {-# INLINE newDataFrame #-} - -- | Create a new mutable DataFrame. -newPinnedDataFrame :: forall t (ns :: [Nat]) s - . ( PrimBytes t, Dimensions ns) - => ST s (STDataFrame s t ns) -newPinnedDataFrame = STDataFrame <$> ST (newPinnedDataFrame# @t @ns) +newPinnedDataFrame :: + forall (t :: Type) ns s + . (PrimBytes t, Dimensions ns) => ST s (STDataFrame s t ns) +newPinnedDataFrame = coerce (newPinnedDataFrame# @t @_ @ns) {-# INLINE newPinnedDataFrame #-} +-- | Create a new mutable DataFrame of the same size. +oneMoreDataFrame :: + forall (t :: Type) ns s + . STDataFrame s t ns -> ST s (STDataFrame s t ns) +oneMoreDataFrame = coerce (oneMoreDataFrame# @t @_ @ns) +{-# INLINE oneMoreDataFrame #-} +-- | View a part of a DataFrame. +-- +-- This function does not perform a copy. +-- All changes to a new DataFrame will be reflected in the original DataFrame as well. +-- +-- If any of the dims in @as@ or @b@ is unknown (@a ~ XN m@), +-- then this function is unsafe and can throw an `OutOfDimBounds` exception. +-- Otherwise, its safety is guaranteed by the type system. +subDataFrameView :: + forall (t :: Type) b bi bd as bs asbs s + . (SubFrameIndexCtx b bi bd, KnownDim bd, ConcatList as (b :+ bs) asbs) + => Idxs (as +: bi) -> STDataFrame s t asbs -> STDataFrame s t (bd :+ bs) +subDataFrameView = coerce (subDataFrameView# @t @_ @b @bi @bd @as @bs @asbs) + +-- | View a part of a DataFrame. +-- +-- This function does not perform a copy. +-- All changes to a new DataFrame will be reflected in the original DataFrame as well. +-- +-- This is a simpler version of @subDataFrameView@ that allows +-- to view over one index at a time. +-- +-- If any of the dims in @as@ is unknown (@a ~ XN m@), +-- then this function is unsafe and can throw an `OutOfDimBounds` exception. +-- Otherwise, its safety is guaranteed by the type system. +subDataFrameView' :: + forall (t :: Type) as bs asbs s + . ConcatList as bs asbs + => Idxs as -> STDataFrame s t asbs -> STDataFrame s t bs +subDataFrameView' = coerce (subDataFrameView'# @t @_ @as @bs @asbs) + -- | Copy one DataFrame into another mutable DataFrame at specified position. -- --- In contrast to @copyMDataFrame'@, this function allows to copy over a range of contiguous --- indices over a single dimension. +-- In contrast to @copyDataFrame'@, this function allows to copy over a range +-- of contiguous indices over a single dimension. -- For example, you can write a 3x4 matrix into a 7x4 matrix, starting at indices 0..3. -copyDataFrame :: forall (t :: Type) - (b :: Nat) (bi :: Nat) (bd :: Nat) - (as :: [Nat]) (bs :: [Nat]) (asbs :: [Nat]) s - . ( b ~ (bi + bd - 1) - , PrimBytes t - , PrimBytes (DataFrame t (bd :+ bs)) - , ConcatList as (b :+ bs) asbs - ) - => Idxs (as +: bi) -> DataFrame t (bd :+ bs) - -> STDataFrame s t asbs -> ST s () -copyDataFrame ei df (STDataFrame mdf) = ST (copyDataFrame# ei df mdf) +-- +-- This function is safe (no `OutOfDimBounds` exception possible). +-- If any of the dims in @as@ is unknown (@a ~ XN m@), +-- you may happen to write data beyond dataframe bounds. +-- In this case, this function does nothing. +-- If (@b ~ XN m@) and (@Idx bi + Dim bd > Dim b@), this function copies only as +-- many elements as fits into the dataframe along this dimension (possibly none). +copyDataFrame :: + forall (t :: Type) b bi bd as bs asbs s + . ( SubFrameIndexCtx b bi bd, KnownDim bd, ExactDims bs + , PrimArray t (DataFrame t (bd :+ bs)) + , ConcatList as (b :+ bs) asbs ) + => Idxs (as +: bi) -> DataFrame t (bd :+ bs) -> STDataFrame s t asbs -> ST s () +copyDataFrame = coerce (copyDataFrame# @t @_ @b @bi @bd @as @bs @asbs) {-# INLINE copyDataFrame #-} -- | Copy one mutable DataFrame into another mutable DataFrame at specified position. -- --- In contrast to @copyMDataFrame'@, this function allows to copy over a range of contiguous --- indices over a single dimension. +-- In contrast to @copyMutableDataFrame'@, this function allows to copy over a range +-- of contiguous indices over a single dimension. -- For example, you can write a 3x4 matrix into a 7x4 matrix, starting at indices 0..3. -copyMutableDataFrame :: forall (t :: Type) - (b :: Nat) (bi :: Nat) (bd :: Nat) - (as :: [Nat]) (bs :: [Nat]) (asbs :: [Nat]) s - . ( b ~ (bi + bd - 1) - , PrimBytes t - , ConcatList as (b :+ bs) asbs - ) - => Idxs (as +: bi) -> STDataFrame s t (bd :+ bs) - -> STDataFrame s t asbs -> ST s () -copyMutableDataFrame ei (STDataFrame mdfA) (STDataFrame mdfB) - = ST (copyMDataFrame# ei mdfA mdfB) +-- +-- This function is safe (no `OutOfDimBounds` exception possible). +-- If any of the dims in @as@ is unknown (@a ~ XN m@), +-- you may happen to write data beyond dataframe bounds. +-- In this case, this function does nothing. +-- If (@b ~ XN m@) and (@Idx bi + Dim bd > Dim b@), this function copies only as +-- many elements as fits into the dataframe along this dimension (possibly none). +copyMutableDataFrame :: + forall (t :: Type) b bi bd as bs asbs s + . ( SubFrameIndexCtx b bi bd + , ExactDims bs + , PrimBytes t + , ConcatList as (b :+ bs) asbs ) + => Idxs (as +: bi) -> STDataFrame s t (bd :+ bs) -> STDataFrame s t asbs -> ST s () +copyMutableDataFrame = coerce (copyMDataFrame# @t @_ @b @bi @bd @as @bs @asbs) {-# INLINE copyMutableDataFrame #-} -- | Copy one DataFrame into another mutable DataFrame at specified position. -- --- This is a simpler version of @copyDataFrame@ that allows to copy over one index at a time. -copyDataFrame' :: forall (t :: Type) - (as :: [Nat]) (bs :: [Nat]) (asbs :: [Nat]) s - . ( PrimBytes t - , PrimBytes (DataFrame t bs) - , ConcatList as bs asbs - ) - => Idxs as -> DataFrame t bs - -> STDataFrame s t asbs -> ST s () -copyDataFrame' ei df (STDataFrame mdf) = ST (copyDataFrame'# ei df mdf) +-- This is a simpler version of @copyDataFrame@ that allows +-- to copy over one index at a time. +-- +-- This function is safe (no `OutOfDimBounds` exception possible). +-- If any of the dims in @as@ is unknown (@a ~ XN m@), +-- you may happen to write data beyond dataframe bounds. +-- In this case, this function does nothing. +copyDataFrame' :: + forall (t :: Type) as bs asbs s + . ( ExactDims bs + , PrimArray t (DataFrame t bs) + , ConcatList as bs asbs ) + => Idxs as -> DataFrame t bs -> STDataFrame s t asbs -> ST s () +copyDataFrame' = coerce (copyDataFrame'# @t @_ @as @bs @asbs) {-# INLINE copyDataFrame' #-} -- | Copy one mutable DataFrame into another mutable DataFrame at specified position. -- --- This is a simpler version of @copyDataFrame@ that allows to copy over one index at a time. -copyMutableDataFrame' :: forall (t :: Type) - (as :: [Nat]) (bs :: [Nat]) (asbs :: [Nat]) s - . ( PrimBytes t - , ConcatList as bs asbs - ) - => Idxs as -> STDataFrame s t bs - -> STDataFrame s t asbs -> ST s () -copyMutableDataFrame' ei (STDataFrame mdfA) (STDataFrame mdfB) - = ST (copyMDataFrame'# ei mdfA mdfB) +-- This is a simpler version of @copyMutableDataFrame@ that allows +-- to copy over one index at a time. +-- +-- This function is safe (no `OutOfDimBounds` exception possible). +-- If any of the dims in @as@ is unknown (@a ~ XN m@), +-- you may happen to write data beyond dataframe bounds. +-- In this case, this function does nothing. +copyMutableDataFrame' :: + forall (t :: Type) as bs asbs s + . (ExactDims bs, PrimBytes t, ConcatList as bs asbs) + => Idxs as -> STDataFrame s t bs -> STDataFrame s t asbs -> ST s () +copyMutableDataFrame' = coerce (copyMDataFrame'# @t @_ @as @bs @asbs) {-# INLINE copyMutableDataFrame' #-} +-- | Copy one DataFrame into another mutable DataFrame by offset in +-- primitive elements. +-- +-- This is a low-level copy function; you have to keep in mind the row-major +-- layout of Mutable DataFrames. Offset bounds are not checked. +-- You will get an undefined behavior if you write beyond the DataFrame bounds. +copyDataFrameOff :: + forall (t :: Type) as bs asbs s + . ( Dimensions bs + , PrimArray t (DataFrame t bs) + , ConcatList as bs asbs ) + => Int -> DataFrame t bs -> STDataFrame s t asbs -> ST s () +copyDataFrameOff = coerce (copyDataFrameOff# @t @_ @as @bs @asbs) +{-# INLINE copyDataFrameOff #-} + +-- | Copy one mutable DataFrame into another mutable DataFrame by offset in +-- primitive elements. +-- +-- This is a low-level copy function; you have to keep in mind the row-major +-- layout of Mutable DataFrames. Offset bounds are not checked. +-- You will get an undefined behavior if you write beyond the DataFrame bounds +copyMutableDataFrameOff :: + forall (t :: Type) as bs asbs s + . (ExactDims bs, PrimBytes t, ConcatList as bs asbs) + => Int -> STDataFrame s t bs -> STDataFrame s t asbs -> ST s () +copyMutableDataFrameOff = coerce (copyMDataFrameOff# @t @_ @as @bs @asbs) +{-# INLINE copyMutableDataFrameOff #-} + -- | Make a mutable DataFrame immutable, without copying. -unsafeFreezeDataFrame :: forall (t :: Type) (ns :: [Nat]) s - . PrimArray t (DataFrame t ns) - => STDataFrame s t ns -> ST s (DataFrame t ns) -unsafeFreezeDataFrame (STDataFrame mdf) = ST (unsafeFreezeDataFrame# mdf) +unsafeFreezeDataFrame :: + forall (t :: Type) ns s + . PrimArray t (DataFrame t ns) + => STDataFrame s t ns -> ST s (DataFrame t ns) +unsafeFreezeDataFrame = coerce (unsafeFreezeDataFrame# @t @_ @ns) {-# INLINE unsafeFreezeDataFrame #-} - -- | Copy content of a mutable DataFrame into a new immutable DataFrame. -freezeDataFrame :: forall (t :: Type) (ns :: [Nat]) s - . PrimArray t (DataFrame t ns) - => STDataFrame s t ns -> ST s (DataFrame t ns) -freezeDataFrame (STDataFrame mdf) = ST (freezeDataFrame# mdf) +freezeDataFrame :: + forall (t :: Type) ns s + . PrimArray t (DataFrame t ns) + => STDataFrame s t ns -> ST s (DataFrame t ns) +freezeDataFrame = coerce (freezeDataFrame# @t @_ @ns) {-# INLINE freezeDataFrame #-} -- | Create a new mutable DataFrame and copy content of immutable one in there. -thawDataFrame :: forall (t :: Type) (ns :: [Nat]) s - . (Dimensions ns, PrimBytes (DataFrame t ns)) - => DataFrame t ns -> ST s (STDataFrame s t ns) -thawDataFrame df = STDataFrame <$> ST (thawDataFrame# df) +thawDataFrame :: + forall (t :: Type) ns s + . (Dimensions ns, PrimArray t (DataFrame t ns)) + => DataFrame t ns -> ST s (STDataFrame s t ns) +thawDataFrame = coerce (thawDataFrame# @t @_ @ns) {-# INLINE thawDataFrame #-} -- | Create a new mutable DataFrame and copy content of immutable one in there. -- The result array is pinned and aligned. -thawPinDataFrame :: forall (t :: Type) (ns :: [Nat]) s - . (Dimensions ns, PrimBytes (DataFrame t ns)) - => DataFrame t ns -> ST s (STDataFrame s t ns) -thawPinDataFrame df = STDataFrame <$> ST (thawPinDataFrame# df) +thawPinDataFrame :: + forall (t :: Type) ns s + . (Dimensions ns, PrimArray t (DataFrame t ns)) + => DataFrame t ns -> ST s (STDataFrame s t ns) +thawPinDataFrame = coerce (thawPinDataFrame# @t @_ @ns) {-# INLINE thawPinDataFrame #-} -- | UnsafeCoerces an underlying byte array. -unsafeThawDataFrame :: forall (t :: Type) (ns :: [Nat]) s - . ( Dimensions ns - , PrimBytes (DataFrame t ns), PrimBytes t) - => DataFrame t ns -> ST s (STDataFrame s t ns) -unsafeThawDataFrame df = STDataFrame <$> ST (unsafeThawDataFrame# df) +unsafeThawDataFrame :: + forall (t :: Type) ns s + . (Dimensions ns, PrimArray t (DataFrame t ns)) + => DataFrame t ns + -> ST s (STDataFrame s t ns) +unsafeThawDataFrame = coerce (unsafeThawDataFrame# @t @_ @ns) {-# INLINE unsafeThawDataFrame #-} - --- | Write a single element at the specified index -writeDataFrame :: forall t (ns :: [Nat]) s - . PrimBytes t - => STDataFrame s t ns -> Idxs ns -> DataFrame t ('[] :: [Nat]) -> ST s () -writeDataFrame (STDataFrame mdf) ei = ST . writeDataFrame# mdf ei . unsafeCoerce# -{-# INLINE writeDataFrame #-} - - --- | Read a single element at the specified index -readDataFrame :: forall (t :: Type) (ns :: [Nat]) s - . PrimBytes t - => STDataFrame s t ns -> Idxs ns -> ST s (DataFrame t ('[] :: [Nat])) -readDataFrame (STDataFrame mdf) = unsafeCoerce# . ST . readDataFrame# mdf -{-# INLINE readDataFrame #-} - +-- | Given two continuations @f@ and @g@. +-- If the input DataFrame is a single broadcast value, use it in @f@. +-- Otherwise, create a new mutable DataFrame and copy content of immutable one +-- in there; then use it in @g@. +-- +-- This function is useful when @thawDataFrame@ cannot be used due to +-- @Dimensions ns@ constraint being not available. +withThawDataFrame :: + forall (t :: Type) ns r s + . PrimArray t (DataFrame t ns) + => (t -> ST s r) + -> (STDataFrame s t ns -> ST s r) + -> DataFrame t ns -> ST s r +withThawDataFrame = coerce (withThawDataFrame# @t @_ @ns @r) --- | Write a single element at the specified element offset -writeDataFrameOff :: forall (t :: Type) (ns :: [Nat]) s - . PrimBytes t - => STDataFrame s t ns -> Int -> DataFrame t ('[] :: [Nat]) -> ST s () -writeDataFrameOff (STDataFrame mdf) (I# i) - = ST . writeDataFrameOff# mdf i . unsafeCoerce# +-- | Write a single element at the specified element offset. +-- +-- This is a low-level write function; you have to keep in mind the row-major +-- layout of Mutable DataFrames. Offset bounds are not checked. +-- You will get an undefined behavior if you write beyond the DataFrame bounds. +writeDataFrameOff :: + forall (t :: Type) ns s + . PrimBytes (DataFrame t ('[] :: KindOf ns)) + => STDataFrame s t ns -> Int -> DataFrame t ('[] :: KindOf ns) -> ST s () +writeDataFrameOff = coerce (writeDataFrameOff# @t @_ @ns) {-# INLINE writeDataFrameOff #-} +-- | Write a single element at the specified index. +-- +-- This function is safe (no `OutOfDimBounds` exception possible). +-- If any of the dims in @ns@ is unknown (@n ~ XN m@), +-- you may happen to write data beyond dataframe bounds. +-- In this case, this function does nothing. +writeDataFrame :: + forall (t :: Type) ns s + . PrimBytes (DataFrame t ('[] :: KindOf ns)) + => STDataFrame s t ns -> Idxs ns -> DataFrame t ('[] :: KindOf ns) -> ST s () +writeDataFrame = coerce (writeDataFrame# @t @_ @ns) +{-# INLINE writeDataFrame #-} --- | Read a single element at the specified element offset -readDataFrameOff :: forall (t :: Type) (ns :: [Nat]) s - . PrimBytes t - => STDataFrame s t ns -> Int -> ST s (DataFrame t ('[] :: [Nat])) -readDataFrameOff (STDataFrame mdf) (I# i) - = unsafeCoerce# (ST (readDataFrameOff# mdf i)) +-- | Read a single element at the specified element offset. +-- +-- This is a low-level read function; you have to keep in mind the row-major +-- layout of Mutable DataFrames. Offset bounds are not checked. +-- You will get an undefined behavior if you read beyond the DataFrame bounds. +readDataFrameOff :: + forall (t :: Type) ns s + . PrimBytes (DataFrame t ('[] :: KindOf ns)) + => STDataFrame s t ns -> Int -> ST s (DataFrame t ('[] :: KindOf ns)) +readDataFrameOff = coerce (readDataFrameOff# @t @_ @ns) {-# INLINE readDataFrameOff #-} +-- | Read a single element at the specified index. +-- +-- If any of the dims in @ns@ is unknown (@n ~ XN m@), +-- then this function is unsafe and can throw an `OutOfDimBounds` exception. +-- Otherwise, its safety is guaranteed by the type system. +readDataFrame :: + forall (t :: Type) ns s + . PrimBytes (DataFrame t ('[] :: KindOf ns)) + => STDataFrame s t ns -> Idxs ns -> ST s (DataFrame t ('[] :: KindOf ns)) +readDataFrame = coerce (readDataFrame# @t @_ @ns) +{-# INLINE readDataFrame #-} -- | Check if the byte array wrapped by this DataFrame is pinned, -- which means cannot be relocated by GC. -isDataFramePinned :: forall (k :: Type) (t :: Type) (ns :: [k]) s - . KnownDimKind k - => STDataFrame s t ns -> Bool -isDataFramePinned df = case dimKind @k of - DimNat -> case df of - STDataFrame x -> isDataFramePinned# x - DimXNat -> case df of - XSTFrame (STDataFrame x) -> isDataFramePinned# x +isDataFramePinned :: + forall (t :: Type) ns s . STDataFrame s t ns -> Bool +isDataFramePinned = coerce (isDataFramePinned# @t @_ @ns) +{-# INLINE isDataFramePinned #-} + +-- | Get cumulative dimensions @ns@ of a @STDataFrame s t ns@ +getDataFrameSteps :: + forall (t :: Type) ns s . STDataFrame s t ns -> CumulDims +getDataFrameSteps = coerce (getDataFrameSteps# @t @_ @ns) +{-# INLINE getDataFrameSteps #-} + +_unusedTopBind :: STDataFrame s t ns +_unusedTopBind = STDataFrame undefined diff --git a/src/Numeric/DataFrame/SubSpace.hs b/src/Numeric/DataFrame/SubSpace.hs --- a/src/Numeric/DataFrame/SubSpace.hs +++ b/src/Numeric/DataFrame/SubSpace.hs @@ -1,503 +1,1265 @@ -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE DataKinds #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE FunctionalDependencies #-} -{-# LANGUAGE MagicHash #-} -{-# LANGUAGE MultiParamTypeClasses #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE TypeOperators #-} -{-# LANGUAGE UnboxedTuples #-} -{-# LANGUAGE UndecidableInstances #-} -{-# LANGUAGE UndecidableSuperClasses #-} ------------------------------------------------------------------------------ --- | --- Module : Numeric.DataFrame.SubSpace --- Copyright : (c) Artem Chirkin --- License : BSD3 --- --- ------------------------------------------------------------------------------ - -module Numeric.DataFrame.SubSpace - ( SubSpace (..), (!), element - , ewfoldMap, iwfoldMap - , ewzip, iwzip - , indexWise_, elementWise_ - ) where - -import Data.Kind -import GHC.Exts - -import Numeric.DataFrame.Internal.PrimArray -import Numeric.DataFrame.Type -import Numeric.Dimensions -import Numeric.PrimBytes - - --- | Operations on DataFrames --- --- @as@ is an indexing dimensionality --- --- @bs@ is an element dimensionality --- --- @t@ is an underlying data type (i.e. Float, Int, Double) --- -class ( ConcatList as bs asbs - , Dimensions as - , Dimensions bs - , Dimensions asbs - , PrimArray t (DataFrame t bs) - , PrimArray t (DataFrame t asbs) - , PrimBytes (DataFrame t bs) - , PrimBytes (DataFrame t asbs) - ) => SubSpace (t :: Type) (as :: [Nat]) (bs :: [Nat]) (asbs :: [Nat]) - | asbs as -> bs, asbs bs -> as, as bs -> asbs where - - -- | Unsafely get a sub-dataframe by its primitive element offset. - -- The offset is not checked to be aligned to the space structure or for bounds. - -- Arguments are zero-based primitive element offset and subset ("bs" element) - -- size (aka `totalDim` of sub dataframe) - -- - -- Normal indexing can be expressed in terms of `indexOffset#`: - -- - -- > index i = case (# dimVal (dim @bs), fromEnum i #) of (# I# n, I# j #) -> indexOffset# (n *# j) - indexOffset# :: Int# -- ^ Prim element offset - -> DataFrame t asbs -> DataFrame t bs - indexOffset# off df = case uniqueOrCumulDims df of - Left a -> broadcast a - Right steps -> fromElems (dropPref (dims @as) steps) (offsetElems df +# off) (getBytes df) - {-# INLINE [1] indexOffset# #-} - - -- | Unsafely update a sub-dataframe by its primitive element offset. - -- The offset is not checked to be aligned to the space structure or for bounds. - -- Arguments are zero-based primitive element offset and subset ("bs" element) - -- size (aka `totalDim` of sub dataframe) - -- - -- Normal updating can be expressed in terms of `indexOffset#`: - -- - -- > update i = case (# dimVal (dim @bs), fromEnum i #) of (# I# n, I# j #) -> updateOffset# (n *# j) - updateOffset# :: Int# -- ^ Prim element offset - -> DataFrame t bs -> DataFrame t asbs -> DataFrame t asbs - updateOffset# off x df - | steps <- getSteps (dims @asbs) df - , elemBS <- byteSize @t undefined = case runRW# - ( \s0 -> case newByteArray# (cdTotalDim# steps *# elemBS) s0 of - (# s1, mba #) -> unsafeFreezeByteArray# mba - ( writeBytes mba (off *# elemBS) x - ( writeBytes mba 0# df s1 ) - ) - ) of (# _, r #) -> fromElems steps 0# r - {-# INLINE [1] updateOffset# #-} - - -- | Get an element by its index in the dataframe - index :: Idxs as -> DataFrame t asbs -> DataFrame t bs - index i df = case uniqueOrCumulDims df of - Left a -> broadcast a - Right steps -> case cdIx steps i of - I# off -> fromElems (dropPref (dims @as) steps) (offsetElems df +# off) (getBytes df) - {-# INLINE [1] index #-} - - -- | Get a few contiguous elements. - -- - -- In a sense, this is just a more complicated version of `index`. - slice :: forall (bi :: Nat) (bd :: Nat) (b' :: Nat) (bs' :: [Nat]) - . ( b' ~ (bi + bd - 1), bs ~ (b' :+ bs'), KnownDim bd) - => Idxs (as +: bi) -> DataFrame t asbs -> DataFrame t (bd :+ bs') - slice i df - | _ :* Dims <- dims @bs - , Dict <- inferKnownBackend @t @(bd ': bs') - = case uniqueOrCumulDims df of - Left a -> broadcast a - Right steps -> case cdIx steps i of - I# off - | bsteps <- repHead (dimVal' @bd) (dropPref (dims @as) steps) - -> fromElems bsteps (offsetElems df +# off) (getBytes df) - | otherwise = error "SubSpace/slice: impossible arguments" - where - repHead y (CumulDims (_:x:xs)) = CumulDims (y*x:x:xs) - repHead _ steps = steps - - -- | Set a new value to an element - update :: Idxs as -> DataFrame t bs -> DataFrame t asbs -> DataFrame t asbs - update i x df - | steps <- getSteps (dims @asbs) df - , I# off <- cdIx steps i - , elemBS <- byteSize @t undefined = case runRW# - ( \s0 -> case newByteArray# (cdTotalDim# steps *# elemBS) s0 of - (# s1, mba #) -> unsafeFreezeByteArray# mba - ( writeBytes mba (off *# elemBS) x - ( writeBytes mba 0# df s1 ) - ) - ) of (# _, r #) -> fromElems steps 0# r - {-# INLINE [1] update #-} - - -- | Update a few contiguous elements - -- - -- In a sense, this is just a more complicated version of `update`. - updateSlice :: forall (bi :: Nat) (bd :: Nat) (b' :: Nat) (bs' :: [Nat]) - . ( b' ~ (bi + bd - 1), bs ~ (b' :+ bs'), KnownDim bd) - => Idxs (as +: bi) -> DataFrame t (bd :+ bs') -> DataFrame t asbs -> DataFrame t asbs - updateSlice i x df - | _ :* Dims <- dims @bs - , Dict <- inferKnownBackend @t @(bd ': bs') - , steps <- getSteps (dims @asbs) df - , I# off <- cdIx steps i - , elemBS <- byteSize @t undefined = case runRW# - ( \s0 -> case newByteArray# (cdTotalDim# steps *# elemBS) s0 of - (# s1, mba #) -> unsafeFreezeByteArray# mba - ( writeBytes mba (off *# elemBS) x - ( writeBytes mba 0# df s1 ) - ) - ) of (# _, r #) -> fromElems steps 0# r - | otherwise = error "SubSpace/updateSlice: impossible arguments" - {-# INLINE [1] updateSlice #-} - - -- | Map a function over each element of DataFrame - ewmap :: forall s (bs' :: [Nat]) (asbs' :: [Nat]) - . SubSpace s as bs' asbs' - => (DataFrame s bs' -> DataFrame t bs) - -> DataFrame s asbs' -> DataFrame t asbs - ewmap f df - | bsizeT <- byteSize @t undefined - , stepsAS <- cumulDims $ dims @as - , stepsBS <- cumulDims $ dims @bs - , stepsBS' <- cumulDims $ dims @bs' - , stepsASBS <- stepsAS <> stepsBS - , lenAS <- cdTotalDim# stepsAS - , lenBS <- cdTotalDim# stepsBS - , lenBS' <- cdTotalDim# stepsBS' - , lenBSB <- lenBS *# bsizeT - , lenASBSB <- lenAS *# lenBSB - = let go mba sourceOffE destOffB s - | isTrue# (destOffB >=# lenASBSB) - = s - | otherwise - = go mba (sourceOffE +# lenBS') (destOffB +# lenBSB) - (writeBytes mba destOffB (f (indexOffset# sourceOffE df)) s) - in case runRW# - ( \s0 -> case newByteArray# lenASBSB s0 of - (# s1, mba #) -> unsafeFreezeByteArray# mba ( go mba 0# 0# s1 ) - ) of (# _, r #) -> fromElems stepsASBS 0# r - {-# INLINE [1] ewmap #-} - - -- | Map a function over each element with its index of DataFrame - iwmap :: forall s (bs' :: [Nat]) (asbs' :: [Nat]) - . SubSpace s as bs' asbs' - => (Idxs as -> DataFrame s bs' -> DataFrame t bs) - -> DataFrame s asbs' -> DataFrame t asbs - iwmap f df - | bsizeT <- byteSize @t undefined - , stepsAS <- cumulDims $ dims @as - , stepsBS <- cumulDims $ dims @bs - , stepsBS' <- cumulDims $ dims @bs' - , stepsASBS <- stepsAS <> stepsBS - , lenAS <- cdTotalDim# stepsAS - , lenBS <- cdTotalDim# stepsBS - , lenBS' <- cdTotalDim# stepsBS' - , lenBSB <- lenBS *# bsizeT - , lenASBSB <- lenAS *# lenBSB - = let go _ [] _ _ s = s - go mba (i:is) sourceOffE destOffB s - = go mba is (sourceOffE +# lenBS') (destOffB +# lenBSB) - (writeBytes mba destOffB (f i (indexOffset# sourceOffE df)) s) - in case runRW# - ( \s0 -> case newByteArray# lenASBSB s0 of - (# s1, mba #) -> unsafeFreezeByteArray# mba ( go mba [minBound..maxBound] 0# 0# s1 ) - ) of (# _, r #) -> fromElems stepsASBS 0# r - {-# INLINE [1] iwmap #-} - - -- | Generate a DataFrame by repeating an element - ewgen :: DataFrame t bs -> DataFrame t asbs - ewgen df = case uniqueOrCumulDims df of - Left a -> broadcast a - Right stepsBS - | stepsAS <- cumulDims $ dims @as - , stepsASBS <- stepsAS <> stepsBS - , elS <- byteSize @t undefined - , lenBSB <- cdTotalDim# stepsBS *# elS - , lenASBSB <- cdTotalDim# stepsASBS *# elS - -> let go mba destOffB s - | isTrue# (destOffB >=# lenASBSB) - = s - | otherwise - = go mba (destOffB +# lenBSB) - (writeBytes mba destOffB df s) - in case runRW# - ( \s0 -> case newByteArray# lenASBSB s0 of - (# s1, mba #) -> unsafeFreezeByteArray# mba ( go mba 0# s1 ) - ) of (# _, r #) -> fromElems stepsASBS 0# r - {-# INLINE [1] ewgen #-} - - -- | Generate a DataFrame by iterating a function (index -> element) - iwgen :: (Idxs as -> DataFrame t bs) -> DataFrame t asbs - iwgen f - | stepsAS <- cumulDims $ dims @as - , stepsBS <- cumulDims $ dims @bs - , stepsASBS <- stepsAS <> stepsBS - , elS <- byteSize @t undefined - , lenBSB <- cdTotalDim# stepsBS *# elS - , lenASBSB <- cdTotalDim# stepsASBS *# elS - = let go _ [] _ s = s - go mba (i:is) destOffB s - = go mba is (destOffB +# lenBSB) (writeBytes mba destOffB (f i) s) - in case runRW# - ( \s0 -> case newByteArray# lenASBSB s0 of - (# s1, mba #) -> unsafeFreezeByteArray# mba ( go mba [minBound..maxBound] 0# s1 ) - ) of (# _, r #) -> fromElems stepsASBS 0# r - {-# INLINE [1] iwgen #-} - - -- | Left-associative fold of a DataFrame. - -- The fold is strict, so accumulater is evaluated to WHNF; - -- but you'd better make sure that the function is strict enough to not - -- produce memory leaks deeply inside the result data type. - ewfoldl :: (b -> DataFrame t bs -> b) -> b -> DataFrame t asbs -> b - ewfoldl f x0 df = case uniqueOrCumulDims df of - Left a -> - let b = broadcast a - go 0 x = x - go n x = go (n-1) $! f x b - in go (totalDim' @as) x0 - Right stepsASBS - | stepsBS <- dropPref (dims @as) stepsASBS - , lenBS <- cdTotalDim# stepsBS - , lenASBS <- cdTotalDim# stepsASBS - -> let go sourceOffE x - | isTrue# (sourceOffE >=# lenASBS) - = x - | otherwise - = go (sourceOffE +# lenBS) $! f x (indexOffset# sourceOffE df) - in go 0# x0 - - -- | Left-associative fold of a DataFrame with an index - -- The fold is strict, so accumulater is evaluated to WHNF; - -- but you'd better make sure that the function is strict enough to not - -- produce memory leaks deeply inside the result data type. - iwfoldl :: (Idxs as -> b -> DataFrame t bs -> b) -> b -> DataFrame t asbs -> b - iwfoldl f x0 df = case uniqueOrCumulDims df of - Left a -> - let b = broadcast a - go [] x = x - go (i:is) x = go is $! f i x b - in go [minBound..maxBound] x0 - Right stepsASBS - | stepsBS <- dropPref (dims @as) stepsASBS - , lenBS <- cdTotalDim# stepsBS - -> let go [] _ x = x - go (i:is) sourceOffE x - = go is (sourceOffE +# lenBS) $! f i x (indexOffset# sourceOffE df) - in go [minBound..maxBound] 0# x0 - - -- | Right-associative fold of a DataFrame - -- The fold is strict, so accumulater is evaluated to WHNF; - -- but you'd better make sure that the function is strict enough to not - -- produce memory leaks deeply inside the result data type. - ewfoldr :: (DataFrame t bs -> b -> b) -> b -> DataFrame t asbs -> b - ewfoldr f x0 df = case uniqueOrCumulDims df of - Left a -> - let b = broadcast a - go 0 x = x - go n x = go (n-1) $! f b x - in go (totalDim' @as) x0 - Right stepsASBS - | stepsBS <- dropPref (dims @as) stepsASBS - , lenBS <- cdTotalDim# stepsBS - , lenASBS <- cdTotalDim# stepsASBS - -> let go sourceOffE x - | isTrue# (sourceOffE <# 0#) - = x - | otherwise - = go (sourceOffE -# lenBS) $! f (indexOffset# sourceOffE df) x - in go (lenASBS -# lenBS) x0 - - -- | Right-associative fold of a DataFrame with an index - -- The fold is strict, so accumulater is evaluated to WHNF; - -- but you'd better make sure that the function is strict enough to not - -- produce memory leaks deeply inside the result data type. - iwfoldr :: (Idxs as -> DataFrame t bs -> b -> b) -> b -> DataFrame t asbs -> b - iwfoldr f x0 df = case uniqueOrCumulDims df of - Left a -> - let b = broadcast a - go [] x = x - go (i:is) x = go is $! f i b x - in go [maxBound, pred maxBound .. minBound] x0 - Right stepsASBS - | stepsBS <- dropPref (dims @as) stepsASBS - , lenBS <- cdTotalDim# stepsBS - , lenASBS <- cdTotalDim# stepsASBS - -> let go [] _ x = x - go (i:is) sourceOffE x - = go is (sourceOffE -# lenBS) $! f i (indexOffset# sourceOffE df) x - in go [maxBound, pred maxBound .. minBound] (lenASBS -# lenBS) x0 - - -- | Apply an applicative functor on each element (Lens-like traversal) - elementWise :: forall (s :: Type) (bs' :: [Nat]) (asbs' :: [Nat]) (f :: Type -> Type) - . ( Applicative f - , SubSpace s as bs' asbs' - ) - => (DataFrame s bs' -> f (DataFrame t bs)) - -> DataFrame s asbs' -> f (DataFrame t asbs) - elementWise = indexWise . const - {-# INLINE [1] elementWise #-} - - -- | Apply an applicative functor on each element with its index - -- (Lens-like indexed traversal) - indexWise :: forall (s :: Type) (bs' :: [Nat]) (asbs' :: [Nat]) (f :: Type -> Type) - . ( Applicative f - , SubSpace s as bs' asbs' - ) - => (Idxs as -> DataFrame s bs' -> f (DataFrame t bs)) - -> DataFrame s asbs' -> f (DataFrame t asbs) - indexWise f df = runWithState <$> iwfoldl applyF (pure initialState) df - where - steps = cumulDims $ dims @asbs - -- run a state-based continuation within RW - runWithState :: ( State# RealWorld -> (# State# RealWorld, (# MutableByteArray# RealWorld, Int# #) #)) - -> DataFrame t asbs - runWithState g = case runRW# - ( \s0 -> case g s0 of - (# s1, (# marr, _ #) #) -> unsafeFreezeByteArray# marr s1 - ) of (# _, arr #) -> fromElems steps 0# arr - - -- Prepare empty byte array for the result DataFrame and keep a current position counter - -- Input: state - -- Output: state + - -- ( current mutable byte array + current write position ) - initialState :: State# RealWorld -> (# State# RealWorld, (# MutableByteArray# RealWorld, Int# #) #) - initialState s0 = case newByteArray# (rezLength# *# rezElBSize#) s0 of - (# s1, marr #) -> (# s1, (# marr, 0# #) #) - - -- Given the result chunk, write it into a mutable array - updateChunk :: (State# RealWorld -> (# State# RealWorld, (# MutableByteArray# RealWorld, Int# #) #)) - -> DataFrame t bs - -> (State# RealWorld -> (# State# RealWorld, (# MutableByteArray# RealWorld, Int# #) #)) - updateChunk g dfChunk = case (# byteOffset dfChunk, getBytes dfChunk #) of - (# off#, arr# #) -> \s -> case g s of - (# s1, (# marr#, pos# #) #) -> case - copyByteArray# arr# (off# *# rezElBSize#) - marr# (pos# *# rezElBSize#) - (rezStepN# *# rezElBSize#) s1 of - s2 -> (# s2, (# marr#, pos# +# rezStepN# #) #) - - -- Apply applicative functor on each chunk and update a state. - applyF :: Idxs as - -> f (State# RealWorld -> (# State# RealWorld, (# MutableByteArray# RealWorld, Int# #) #)) - -> DataFrame s bs' - -> f (State# RealWorld -> (# State# RealWorld, (# MutableByteArray# RealWorld, Int# #) #)) - applyF idx s dfChunk = idx `seq` dfChunk `seq` updateChunk <$> s <*> f idx dfChunk - - -- Element byte size of the result DataFrame (byte size of s) - rezElBSize# = byteSize @t undefined - -- Number of primitive elements in the result DataFrame chunk - !(I# rezStepN#) = fromIntegral $ totalDim' @bs - -- Number of primitive elements in the result DataFrame - rezLength# = cdTotalDim# steps - - - --- | Apply an applicative functor on each element with its index --- (Lens-like indexed traversal) -indexWise_ :: forall t as bs asbs f b - . (SubSpace t as bs asbs, Applicative f) - => (Idxs as -> DataFrame t bs -> f b) - -> DataFrame t asbs -> f () -indexWise_ f = iwfoldr (\i -> (*>) . f i) (pure ()) - --- | Apply an applicative functor on each element (Lens-like traversal) -elementWise_ :: forall t as bs asbs f b - . (SubSpace t as bs asbs, Applicative f) - => (DataFrame t bs -> f b) - -> DataFrame t asbs -> f () -elementWise_ f = ewfoldr ((*>) . f) (pure ()) - - --- | Apply a functor over a single element (simple lens) -element :: forall t (as :: [Nat]) (bs :: [Nat]) (asbs :: [Nat]) f - . (SubSpace t as bs asbs, Applicative f) - => Idxs as - -> (DataFrame t bs -> f (DataFrame t bs)) - -> DataFrame t asbs -> f (DataFrame t asbs) -element i f df = flip (update i) df <$> f (index i df) -{-# INLINE element #-} - --- | Index an element (reverse arguments of `index`) -(!) :: SubSpace t (as :: [Nat]) (bs :: [Nat]) (asbs :: [Nat]) - => DataFrame t asbs -> Idxs as -> DataFrame t bs -(!) = flip index -{-# INLINE (!) #-} -infixl 4 ! - - - -ewfoldMap :: forall t (as :: [Nat]) (bs :: [Nat]) (asbs :: [Nat]) m - . (Monoid m, SubSpace t as bs asbs) - => (DataFrame t bs -> m) -> DataFrame t asbs -> m -ewfoldMap f = ewfoldl (\m b -> m `seq` (mappend m $! f b)) mempty -{-# INLINE ewfoldMap #-} - -iwfoldMap :: forall t (as :: [Nat]) (bs :: [Nat]) (asbs :: [Nat]) m - . ( Monoid m, SubSpace t as bs asbs) - => (Idxs as -> DataFrame t bs -> m) -> DataFrame t asbs -> m -iwfoldMap f = iwfoldl (\i m b -> m `seq` (mappend m $! f i b)) mempty -{-# INLINE iwfoldMap #-} - - - --- | Zip two spaces on a specified subspace index-wise (with index) -iwzip :: forall t (as :: [Nat]) (bs :: [Nat]) (asbs :: [Nat]) - s (bs' :: [Nat]) (asbs' :: [Nat]) - r (bs'' :: [Nat]) (asbs'' :: [Nat]) - . ( SubSpace t as bs asbs - , SubSpace s as bs' asbs' - , SubSpace r as bs'' asbs'' - ) - => (Idxs as -> DataFrame t bs -> DataFrame s bs' -> DataFrame r bs'') - -> DataFrame t asbs - -> DataFrame s asbs' - -> DataFrame r asbs'' -iwzip f dft dfs = iwmap g dft - where - g i dft' = f i dft' (index i dfs) -{-# INLINE iwzip #-} - --- | Zip two spaces on a specified subspace element-wise (without index) -ewzip :: forall t (as :: [Nat]) (bs :: [Nat]) (asbs :: [Nat]) - s (bs' :: [Nat]) (asbs' :: [Nat]) - r (bs'' :: [Nat]) (asbs'' :: [Nat]) - . ( SubSpace t as bs asbs - , SubSpace s as bs' asbs' - , SubSpace r as bs'' asbs'' - ) - => (DataFrame t bs -> DataFrame s bs' -> DataFrame r bs'') - -> DataFrame t asbs - -> DataFrame s asbs' - -> DataFrame r asbs'' -ewzip = iwzip . const -{-# INLINE ewzip #-} - - -instance ( ConcatList as bs asbs - , Dimensions as - , Dimensions bs - , Dimensions asbs - , PrimArray t (DataFrame t bs) - , PrimArray t (DataFrame t asbs) - , PrimBytes (DataFrame t bs) - , PrimBytes (DataFrame t asbs) - ) => SubSpace t (as :: [Nat]) (bs :: [Nat]) (asbs :: [Nat]) where - - - -dropPref :: Dims (ns :: [Nat]) -> CumulDims -> CumulDims -dropPref ds = CumulDims . drop (length $ listDims ds) . unCumulDims - - -unSc :: DataFrame (t :: Type) ('[] :: [Nat]) -> t -unSc = unsafeCoerce# - -{-# RULES -"ewgen/broadcast" ewgen = broadcast . unSc - - #-} +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE FunctionalDependencies #-} +{-# LANGUAGE InstanceSigs #-} +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE PolyKinds #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE UndecidableSuperClasses #-} +{-# OPTIONS_GHC -fno-warn-redundant-constraints #-} +{- | +Module : Numeric.DataFrame.SubSpace +Copyright : (c) Artem Chirkin +License : BSD3 + +This module provides a flexible interface to manipulate parts of a DataFrame. + +==== A note on indexing and slicing + +When you index or slice dataframes, the left part of the dimension list +(@as@ for indexing, plus @b@ for slicing) determines the mechanics of accessing +sub-dataframes. +If compiler knows all dimensions at compile time, it can guarantee that the +operation is safe (provided with valid indices). +Otherwise, you can get an `OutOfDimBounds` exception at runtime. + +When all dimensions in the indexing subspace satisfy @d :: Nat@ or @d ~ N n@, +slicing functions are safe to use, but you need some type-level proof for GHC +that the indices align. + +When any of the dimensions are unknown (@d ~ XN m@), these functions are unsafe +-- they can yield an `OutOfDimBounds` exception if you give a bad index. +But they are easy to use (no type-level proof needed). + + -} + +module Numeric.DataFrame.SubSpace + ( -- $indexingSafety + -- * Class definition + SubSpace (SubSpaceCtx), CanSlice + -- * Simple interface + -- + -- All functions in this section are named @sxxx@, where @s@ stands for + -- "simple". These allow using all functions of `SubSpace` with indexing + -- dimensionality @as@ fixed to a single dim @(as ~ '[a])@. + -- Use these functions if you are tired of @TypeApplications@ or find the + -- error messages too cryptic. + , sindexOffset, supdateOffset + , (.!), slookup, supdate, sslice, ssliceMaybe, supdateSlice + , sewgen, siwgen, sewmap, siwmap, sewzip, siwzip + , selement, selementWise, selementWise_, sindexWise, sindexWise_ + , sewfoldl, sewfoldl', sewfoldr, sewfoldr', sewfoldMap + , siwfoldl, siwfoldl', siwfoldr, siwfoldr', siwfoldMap + -- * Flexible interface + -- + -- Functions in this section allow you pick, fold, iterate, or do whatever + -- you want with arbitrary sub-dataframe dimensionalities: + -- e.g. a DataFrame of rank 3 can be processed as an 1D array of matrices + -- or a matrix of vectors, or a 3D array of scalars. + -- Often, you would need @TypeApplications@ to specify explicitly at least + -- the indexing subspace (parameter @as@). + , joinDataFrame, indexOffset, updateOffset + , index, Numeric.DataFrame.SubSpace.lookup + , update, slice, sliceMaybe, updateSlice + , ewgen, iwgen, ewmap, iwmap, ewzip, iwzip + , element, elementWise, elementWise_, indexWise, indexWise_ + , ewfoldl, ewfoldl', ewfoldr, ewfoldr', ewfoldMap + , iwfoldl, iwfoldl', iwfoldr, iwfoldr', iwfoldMap + ) where + +import Control.Arrow (first) +import Control.Monad +import Control.Monad.ST +import Data.Kind +import GHC.Base (Int (..), (+#)) +import Numeric.DataFrame.Internal.PrimArray +import Numeric.DataFrame.ST +import Numeric.DataFrame.Type +import Numeric.Dimensions +import qualified Numeric.TypedList as TL +import Unsafe.Coerce + +-- | Unsafely get a sub-dataframe by its primitive element offset. +-- The offset is not checked to be aligned to the space structure or for bounds. +-- +-- Warning: this function is utterly unsafe -- it does not even throw an exception if +-- the offset is too big; you just get an undefined behavior. +sindexOffset :: + forall t a bs + . SubSpace t '[a] bs (a :+ bs) + => Int -- ^ Prim element offset + -> DataFrame t (a :+ bs) -> DataFrame t bs +sindexOffset = indexOffset @t @'[a] @bs @(a :+ bs) +{-# INLINE sindexOffset #-} + +-- | Unsafely update a sub-dataframe by its primitive element offset. +-- The offset is not checked to be aligned to the space structure or for bounds. +-- +-- Warning: this function is utterly unsafe -- it does not even throw an exception if +-- the offset is too big; you just get an undefined behavior. +supdateOffset :: + forall t a bs + . SubSpace t '[a] bs (a :+ bs) + => Int -- ^ Prim element offset + -> DataFrame t bs -> DataFrame t (a :+ bs) -> DataFrame t (a :+ bs) +supdateOffset = updateOffset @t @'[a] @bs @(a :+ bs) +{-# INLINE supdateOffset #-} + +-- | Get an element by its index in the dataframe. +-- +-- If (@a ~ XN m@) then this function is unsafe and can throw +-- an `OutOfDimBounds` exception. +-- Otherwise, its safety is guaranteed by the type system. +(.!) :: + forall t a bs + . SubSpace t '[a] bs (a :+ bs) + => DataFrame t (a :+ bs) -> Idx a -> DataFrame t bs +(.!) x i = index @t @'[a] @bs @(a :+ bs) (i :* U) x +{-# INLINE (.!) #-} +infixl 4 .! + +-- | Set a new value to an element. +-- +-- If (@a ~ XN m@) and the index falls outside of the DataFrame dim, +-- then this function returns the original DataFrame. +supdate :: + forall t a bs + . SubSpace t '[a] bs (a :+ bs) + => Idx a -> DataFrame t bs -> DataFrame t (a :+ bs) -> DataFrame t (a :+ bs) +supdate = update @t @'[a] @bs @(a :+ bs) . (:*U) +{-# INLINE[1] supdate #-} + +-- | Map a function over each element of DataFrame. +sewmap :: + forall t a bs s bs' + . (SubSpace t '[a] bs (a :+ bs), SubSpace s '[a] bs' (a :+ bs')) + => (DataFrame s bs' -> DataFrame t bs) + -> DataFrame s (a :+ bs') -> DataFrame t (a :+ bs) +sewmap = ewmap @t @'[a] @bs @(a :+ bs) @s @bs' @(a :+ bs') +{-# INLINE sewmap #-} + +-- | Map a function over each element with its index of DataFrame. +siwmap :: + forall t a bs s bs' + . (SubSpace t '[a] bs (a :+ bs), SubSpace s '[a] bs' (a :+ bs')) + => (Idx a -> DataFrame s bs' -> DataFrame t bs) + -> DataFrame s (a :+ bs') -> DataFrame t (a :+ bs) +siwmap f = iwmap @t @'[a] @bs @(a :+ bs) @s @bs' @(a :+ bs') (\(i :* U) -> f i) +{-# INLINE[1] siwmap #-} + +-- | Generate a DataFrame by repeating an element. +sewgen :: + forall t a bs + . (SubSpace t '[a] bs (a :+ bs), Dimensions '[a]) + => DataFrame t bs -> DataFrame t (a :+ bs) +sewgen = ewgen @t @'[a] @bs @(a :+ bs) +{-# INLINE sewgen #-} + +-- | Generate a DataFrame by iterating a function (index -> element). +siwgen :: + forall t a bs + . (SubSpace t '[a] bs (a :+ bs), Dimensions '[a]) + => (Idx a -> DataFrame t bs) -> DataFrame t (a :+ bs) +siwgen f = iwgen @t @'[a] @bs @(a :+ bs) (\(i :* U) -> f i) +{-# INLINE[1] siwgen #-} + +-- | Left-associative lazy fold of a DataFrame. +-- Same rules apply as for `foldl`. +sewfoldl :: + forall t a bs b + . SubSpace t '[a] bs (a :+ bs) + => (b -> DataFrame t bs -> b) -> b -> DataFrame t (a :+ bs) -> b +sewfoldl = ewfoldl @t @'[a] @bs @(a :+ bs) @b +{-# INLINE sewfoldl #-} + +-- | Left-associative strict fold of a DataFrame. +-- Same rules apply as for `foldl'`. +sewfoldl' :: + forall t a bs b + . SubSpace t '[a] bs (a :+ bs) + => (b -> DataFrame t bs -> b) -> b -> DataFrame t (a :+ bs) -> b +sewfoldl' = ewfoldl' @t @'[a] @bs @(a :+ bs) @b +{-# INLINE sewfoldl' #-} + +-- | Left-associative lazy fold of a DataFrame with an index. +-- Same rules apply as for `foldl`. +siwfoldl :: + forall t a bs b + . SubSpace t '[a] bs (a :+ bs) + => (Idx a -> b -> DataFrame t bs -> b) -> b -> DataFrame t (a :+ bs) -> b +siwfoldl f = iwfoldl @t @'[a] @bs @(a :+ bs) @b (\(i :* U) -> f i) +{-# INLINE[1] siwfoldl #-} + +-- | Left-associative strict fold of a DataFrame with an index. +-- Same rules apply as for `foldl'`. +siwfoldl' :: + forall t a bs b + . SubSpace t '[a] bs (a :+ bs) + => (Idx a -> b -> DataFrame t bs -> b) -> b -> DataFrame t (a :+ bs) -> b +siwfoldl' f = iwfoldl' @t @'[a] @bs @(a :+ bs) @b (\(i :* U) -> f i) +{-# INLINE[1] siwfoldl' #-} + +-- | Right-associative lazy fold of a DataFrame. +-- Same rules apply as for `foldr`. +sewfoldr :: + forall t a bs b + . SubSpace t '[a] bs (a :+ bs) + => (DataFrame t bs -> b -> b) -> b -> DataFrame t (a :+ bs) -> b +sewfoldr = ewfoldr @t @'[a] @bs @(a :+ bs) @b +{-# INLINE sewfoldr #-} + +-- | Right-associative strict fold of a DataFrame. +-- Same rules apply as for `foldr'`. +sewfoldr' :: + forall t a bs b + . SubSpace t '[a] bs (a :+ bs) + => (DataFrame t bs -> b -> b) -> b -> DataFrame t (a :+ bs) -> b +sewfoldr' = ewfoldr' @t @'[a] @bs @(a :+ bs) @b +{-# INLINE sewfoldr' #-} + +-- | Right-associative lazy fold of a DataFrame with an index. +-- Same rules apply as for `foldr`. +siwfoldr :: + forall t a bs b + . SubSpace t '[a] bs (a :+ bs) + => (Idx a -> DataFrame t bs -> b -> b) -> b -> DataFrame t (a :+ bs) -> b +siwfoldr f = iwfoldr @t @'[a] @bs @(a :+ bs) @b (\(i :* U) -> f i) +{-# INLINE[1] siwfoldr #-} + +-- | Right-associative strict fold of a DataFrame with an index. +-- Same rules apply as for `foldr'`. +siwfoldr' :: + forall t a bs b + . SubSpace t '[a] bs (a :+ bs) + => (Idx a -> DataFrame t bs -> b -> b) -> b -> DataFrame t (a :+ bs) -> b +siwfoldr' f = iwfoldr' @t @'[a] @bs @(a :+ bs) @b (\(i :* U) -> f i) +{-# INLINE[1] siwfoldr' #-} + +-- | Apply an applicative functor on each element (Lens-like traversal). +selementWise :: + forall t a bs s bs' f + . (SubSpace t '[a] bs (a :+ bs), SubSpace s '[a] bs' (a :+ bs'), Applicative f) + => (DataFrame s bs' -> f (DataFrame t bs)) + -> DataFrame s (a :+ bs') -> f (DataFrame t (a :+ bs)) +selementWise = elementWise @t @'[a] @bs @(a :+ bs) @s @bs' @(a :+ bs') @f +{-# INLINE selementWise #-} + +-- | Apply an applicative functor on each element with its index +-- (Lens-like indexed traversal). +sindexWise :: + forall t a bs s bs' f + . (SubSpace t '[a] bs (a :+ bs), SubSpace s '[a] bs' (a :+ bs'), Applicative f) + => (Idx a -> DataFrame s bs' -> f (DataFrame t bs)) + -> DataFrame s (a :+ bs') -> f (DataFrame t (a :+ bs)) +sindexWise f = indexWise @t @'[a] @bs @(a :+ bs) @s @bs' @(a :+ bs') @f (\(i :* U) -> f i) +{-# INLINE[1] sindexWise #-} + +-- | Apply an applicative functor on each element (Lens-like traversal) +selementWise_ :: + forall t a bs f b + . (SubSpace t '[a] bs (a :+ bs), Applicative f) + => (DataFrame t bs -> f b) -> DataFrame t (a :+ bs) -> f () +selementWise_ = elementWise_ @t @'[a] @bs @(a :+ bs) @f @b +{-# INLINE selementWise_ #-} + +-- | Apply an applicative functor on each element with its index +-- (Lens-like indexed traversal) +sindexWise_ :: + forall t a bs f b + . (SubSpace t '[a] bs (a :+ bs), Applicative f) + => (Idx a -> DataFrame t bs -> f b) -> DataFrame t (a :+ bs) -> f () +sindexWise_ f = indexWise_ @t @'[a] @bs @(a :+ bs) @f @b (\(i :* U) -> f i) +{-# INLINE sindexWise_ #-} + +-- | Apply a functor over a single element (simple lens) +-- +-- If (@a ~ XN m@) and the index falls outside of the DataFrame Dim, the +-- argument Functor is not called and the result is @pure@ original DataFrame. +selement :: + forall t a bs f + . (SubSpace t '[a] bs (a :+ bs), Applicative f) + => Idx a + -> (DataFrame t bs -> f (DataFrame t bs)) + -> DataFrame t (a :+ bs) -> f (DataFrame t (a :+ bs)) +selement = element @t @'[a] @bs @(a :+ bs) @f . (:* U) +{-# INLINE selement #-} + +-- | Map each element of the DataFrame to a monoid, +-- and combine the results. +sewfoldMap :: + forall t a bs m + . (SubSpace t '[a] bs (a :+ bs), Monoid m) + => (DataFrame t bs -> m) -> DataFrame t (a :+ bs) -> m +sewfoldMap = ewfoldMap @t @'[a] @bs @(a :+ bs) @m +{-# INLINE sewfoldMap #-} + +-- | Map each element of the DataFrame and its index to a monoid, +-- and combine the results. +siwfoldMap :: + forall t a bs m + . (SubSpace t '[a] bs (a :+ bs), Monoid m) + => (Idx a -> DataFrame t bs -> m) -> DataFrame t (a :+ bs) -> m +siwfoldMap f = iwfoldMap @t @'[a] @bs @(a :+ bs) @m (\(i :* U) -> f i) +{-# INLINE[1] siwfoldMap #-} + +-- | Zip two spaces on a specified subspace element-wise (without index) +sewzip :: + forall t a bs l bsL r bsR + . (SubSpace t '[a] bs (a :+ bs), SubSpace l '[a] bsL (a :+ bsL), SubSpace r '[a] bsR (a :+ bsR)) + => (DataFrame l bsL -> DataFrame r bsR -> DataFrame t bs) + -> DataFrame l (a :+ bsL) -> DataFrame r (a :+ bsR) -> DataFrame t (a :+ bs) +sewzip = ewzip @t @'[a] @bs @(a :+ bs) @l @bsL @(a :+ bsL) @r @bsR @(a :+ bsR) +{-# INLINE sewzip #-} + +-- | Zip two spaces on a specified subspace index-wise (with index) +siwzip :: + forall t a bs l bsL r bsR + . (SubSpace t '[a] bs (a :+ bs), SubSpace l '[a] bsL (a :+ bsL), SubSpace r '[a] bsR (a :+ bsR)) + => (Idx a -> DataFrame l bsL -> DataFrame r bsR -> DataFrame t bs) + -> DataFrame l (a :+ bsL) -> DataFrame r (a :+ bsR) -> DataFrame t (a :+ bs) +siwzip f = iwzip @t @'[a] @bs @(a :+ bs) @l @bsL @(a :+ bsL) @r @bsR @(a :+ bsR) (\(i :* U) -> f i) +{-# INLINE siwzip #-} + +-- | Flatten a nested DataFrame, analogous to `Control.Monad.join`. +joinDataFrame :: + forall t as bs asbs + . (SubSpace t as bs asbs, PrimBytes (DataFrame t bs)) + => DataFrame (DataFrame t bs) as -> DataFrame t asbs +joinDataFrame = joinDataFrameI @_ @t @as @bs @asbs +{-# INLINE[1] joinDataFrame #-} + +-- | Unsafely get a sub-dataframe by its primitive element offset. +-- The offset is not checked to be aligned to the space structure or for bounds. +-- +-- Warning: this function is utterly unsafe -- it does not even throw an exception if +-- the offset is too big; you just get an undefined behavior. +indexOffset :: + forall t as bs asbs + . SubSpace t as bs asbs + => Int -- ^ Prim element offset + -> DataFrame t asbs -> DataFrame t bs +indexOffset = indexOffsetI @_ @t @as @bs @asbs +{-# INLINE[1] indexOffset #-} + +-- | Unsafely update a sub-dataframe by its primitive element offset. +-- The offset is not checked to be aligned to the space structure or for bounds. +-- +-- Warning: this function is utterly unsafe -- it does not even throw an exception if +-- the offset is too big; you just get an undefined behavior. +updateOffset :: + forall t as bs asbs + . SubSpace t as bs asbs + => Int -- ^ Prim element offset + -> DataFrame t bs -> DataFrame t asbs -> DataFrame t asbs +updateOffset = updateOffsetI @_ @t @as @bs @asbs +{-# INLINE[1] updateOffset #-} + +-- | Get an element by its index in the dataframe. +-- +-- If any of the dims in @as@ is unknown (@a ~ XN m@), +-- then this function is unsafe and can throw an `OutOfDimBounds` exception. +-- Otherwise, its safety is guaranteed by the type system. +index :: + forall t as bs asbs + . SubSpace t as bs asbs + => Idxs as -> DataFrame t asbs -> DataFrame t bs +index = indexI @_ @t @as @bs @asbs +{-# INLINE[1] index #-} + +-- | Set a new value to an element. +-- +-- If any of the dims in @as@ is unknown (@a ~ XN m@), +-- you may happen to update data beyond dataframe bounds. +-- In this case, the original DataFrame is returned. +update :: + forall t as bs asbs + . SubSpace t as bs asbs + => Idxs as -> DataFrame t bs -> DataFrame t asbs -> DataFrame t asbs +update = updateI @_ @t @as @bs @asbs +{-# INLINE[1] update #-} + +-- | Map a function over each element of DataFrame. +ewmap :: + forall t as bs asbs s bs' asbs' + . (SubSpace t as bs asbs, SubSpace s as bs' asbs') + => (DataFrame s bs' -> DataFrame t bs) + -> DataFrame s asbs' -> DataFrame t asbs +ewmap = ewmapI @_ @t @as @bs @asbs @s @bs' @asbs' +{-# INLINE[1] ewmap #-} + +-- | Map a function over each element with its index of DataFrame. +iwmap :: + forall t as bs asbs s bs' asbs' + . (SubSpace t as bs asbs, SubSpace s as bs' asbs') + => (Idxs as -> DataFrame s bs' -> DataFrame t bs) + -> DataFrame s asbs' -> DataFrame t asbs +iwmap = iwmapI @_ @t @as @bs @asbs @s @bs' @asbs' +{-# INLINE[1] iwmap #-} + +-- | Generate a DataFrame by repeating an element. +ewgen :: + forall t as bs asbs + . (SubSpace t as bs asbs, Dimensions as) + => DataFrame t bs -> DataFrame t asbs +ewgen = ewgenI @_ @t @as @bs @asbs +{-# INLINE[1] ewgen #-} + +-- | Generate a DataFrame by iterating a function (index -> element). +iwgen :: + forall t as bs asbs + . (SubSpace t as bs asbs, Dimensions as) + => (Idxs as -> DataFrame t bs) -> DataFrame t asbs +iwgen = iwgenI @_ @t @as @bs @asbs +{-# INLINE[1] iwgen #-} + +-- | Left-associative lazy fold of a DataFrame. +-- Same rules apply as for `foldl`. +ewfoldl :: + forall t as bs asbs b + . SubSpace t as bs asbs + => (b -> DataFrame t bs -> b) -> b -> DataFrame t asbs -> b +ewfoldl = ewfoldlI @_ @t @as @bs @asbs @b +{-# INLINE[1] ewfoldl #-} + +-- | Left-associative strict fold of a DataFrame. +-- Same rules apply as for `foldl'`. +ewfoldl' :: + forall t as bs asbs b + . SubSpace t as bs asbs + => (b -> DataFrame t bs -> b) -> b -> DataFrame t asbs -> b +ewfoldl' f z0 xs = ewfoldr f' id xs z0 + where + f' x k z = k $! f z x +{-# INLINE ewfoldl' #-} + +-- | Left-associative lazy fold of a DataFrame with an index. +-- Same rules apply as for `foldl`. +iwfoldl :: + forall t as bs asbs b + . SubSpace t as bs asbs + => (Idxs as -> b -> DataFrame t bs -> b) -> b -> DataFrame t asbs -> b +iwfoldl = iwfoldlI @_ @t @as @bs @asbs @b +{-# INLINE[1] iwfoldl #-} + +-- | Left-associative strict fold of a DataFrame with an index. +-- Same rules apply as for `foldl'`. +iwfoldl' :: + forall t as bs asbs b + . SubSpace t as bs asbs + => (Idxs as -> b -> DataFrame t bs -> b) -> b -> DataFrame t asbs -> b +iwfoldl' f z0 xs = iwfoldr f' id xs z0 + where + f' i x k z = k $! f i z x +{-# INLINE[1] iwfoldl' #-} + +-- | Right-associative lazy fold of a DataFrame. +-- Same rules apply as for `foldr`. +ewfoldr :: + forall t as bs asbs b + . SubSpace t as bs asbs + => (DataFrame t bs -> b -> b) -> b -> DataFrame t asbs -> b +ewfoldr = ewfoldrI @_ @t @as @bs @asbs @b +{-# INLINE[1] ewfoldr #-} + +-- | Right-associative strict fold of a DataFrame. +-- Same rules apply as for `foldr'`. +ewfoldr' :: + forall t as bs asbs b + . SubSpace t as bs asbs + => (DataFrame t bs -> b -> b) -> b -> DataFrame t asbs -> b +ewfoldr' f z0 xs = ewfoldl f' id xs z0 + where + f' k x z = k $! f x z +{-# INLINE ewfoldr' #-} + +-- | Right-associative lazy fold of a DataFrame with an index. +-- Same rules apply as for `foldr`. +iwfoldr :: + forall t as bs asbs b + . SubSpace t as bs asbs + => (Idxs as -> DataFrame t bs -> b -> b) -> b -> DataFrame t asbs -> b +iwfoldr = iwfoldrI @_ @t @as @bs @asbs @b +{-# INLINE[1] iwfoldr #-} + +-- | Right-associative strict fold of a DataFrame with an index. +-- Same rules apply as for `foldr'`. +iwfoldr' :: + forall t as bs asbs b + . SubSpace t as bs asbs + => (Idxs as -> DataFrame t bs -> b -> b) -> b -> DataFrame t asbs -> b +iwfoldr' f z0 xs = iwfoldl f' id xs z0 + where + f' i k x z = k $! f i x z +{-# INLINE[1] iwfoldr' #-} + +-- | Apply an applicative functor on each element (Lens-like traversal). +elementWise :: + forall t as bs asbs s bs' asbs' f + . (SubSpace t as bs asbs, SubSpace s as bs' asbs', Applicative f) + => (DataFrame s bs' -> f (DataFrame t bs)) + -> DataFrame s asbs' -> f (DataFrame t asbs) +elementWise = elementWiseI @_ @t @as @bs @asbs @s @bs' @asbs' @f +{-# INLINE[1] elementWise #-} + +-- | Apply an applicative functor on each element with its index +-- (Lens-like indexed traversal). +indexWise :: + forall t as bs asbs s bs' asbs' f + . (SubSpace t as bs asbs, SubSpace s as bs' asbs', Applicative f) + => (Idxs as -> DataFrame s bs' -> f (DataFrame t bs)) + -> DataFrame s asbs' -> f (DataFrame t asbs) +indexWise = indexWiseI @_ @t @as @bs @asbs @s @bs' @asbs' @f +{-# INLINE[1] indexWise #-} + +-- | Apply an applicative functor on each element (Lens-like traversal) +elementWise_ :: + forall t as bs asbs f b + . (SubSpace t as bs asbs, Applicative f) + => (DataFrame t bs -> f b) -> DataFrame t asbs -> f () +elementWise_ f = ewfoldr ((*>) . f) (pure ()) +{-# INLINE elementWise_ #-} + +-- | Apply an applicative functor on each element with its index +-- (Lens-like indexed traversal) +indexWise_ :: + forall t as bs asbs f b + . (SubSpace t as bs asbs, Applicative f) + => (Idxs as -> DataFrame t bs -> f b) -> DataFrame t asbs -> f () +indexWise_ f = iwfoldr (\i -> (*>) . f i) (pure ()) +{-# INLINE indexWise_ #-} + +-- | Apply a functor over a single element (simple lens) +-- +-- If any of the dims in @as@ is unknown (@a ~ XN m@) and any of the +-- corresponding indices fall outside of the DataFrame Dims, then the +-- argument Functor is not called and the result is @pure@ original DataFrame. +element :: + forall t as bs asbs f + . (SubSpace t as bs asbs, Applicative f) + => Idxs as + -> (DataFrame t bs -> f (DataFrame t bs)) + -> DataFrame t asbs -> f (DataFrame t asbs) +element i f df = case dimKind @(KindOfEl asbs) of + DimKNat -> flip (update i) df <$> f (index i df) + DimKXNat + | XFrame (dfN :: DataFrame t asbsN) <- df + , Just off <- cdIxM (getSteps (dims @asbsN) dfN) i + -> flip (updateOffset off) df <$> f (indexOffset off df) + | otherwise + -> pure df +{-# INLINE element #-} + + +-- | Map each element of the DataFrame to a monoid, +-- and combine the results. +ewfoldMap :: + forall t as bs asbs m + . (SubSpace t as bs asbs, Monoid m) + => (DataFrame t bs -> m) -> DataFrame t asbs -> m +ewfoldMap f = ewfoldr (mappend . f) mempty +{-# INLINE ewfoldMap #-} + +-- | Map each element of the DataFrame and its index to a monoid, +-- and combine the results. +iwfoldMap :: + forall t as bs asbs m + . (SubSpace t as bs asbs, Monoid m) + => (Idxs as -> DataFrame t bs -> m) -> DataFrame t asbs -> m +iwfoldMap f = iwfoldr (\i -> mappend . f i) mempty +{-# INLINE iwfoldMap #-} + +-- | Zip two spaces on a specified subspace element-wise (without index) +ewzip :: forall t as bs asbs l bsL asbsL r bsR asbsR + . (SubSpace t as bs asbs, SubSpace l as bsL asbsL, SubSpace r as bsR asbsR) + => (DataFrame l bsL -> DataFrame r bsR -> DataFrame t bs) + -> DataFrame l asbsL -> DataFrame r asbsR -> DataFrame t asbs +ewzip = iwzip . const +{-# INLINE ewzip #-} + +-- | Zip two spaces on a specified subspace index-wise (with index). +iwzip :: forall t as bs asbs l bsL asbsL r bsR asbsR + . (SubSpace t as bs asbs, SubSpace l as bsL asbsL, SubSpace r as bsR asbsR) + => (Idxs as -> DataFrame l bsL -> DataFrame r bsR -> DataFrame t bs) + -> DataFrame l asbsL -> DataFrame r asbsR -> DataFrame t asbs +iwzip f dfl dfr = case dimKind @(KindOfEl asbs) of + DimKNat -> iwmap (\i x -> f i x (index i dfr)) dfl + DimKXNat + | XFrame (_ :: DataFrame l asbsLN) <- dfl + , XFrame (_ :: DataFrame r asbsRN) <- dfr + , asbs <- unsafeCoerce -- minimum spans + $ zipWith min (listDims $ dims @asbsLN) (listDims $ dims @asbsRN) :: Dims asbs + , bs <- (dims :: Dims bs) + , as <- (dropSufDims bs asbs :: Dims as) + -> withLiftedConcatList @'Nothing @('Just (DimsBound bs)) @'Nothing as bs asbs $ + \(Dims :: Dims asn) (Dims :: Dims bsn) (Dims :: Dims asbsn) -> case () of + _ | Dict <- inferKnownBackend @t @asbsn + , Dict <- inferKnownBackend @t @bsn + -> XFrame (iwgen @t @asn @bsn @asbsn + (\iN -> let i = liftIdxs iN + in unliftXFrame (f i (index i dfl) (index i dfr)) + ) + ) +{-# INLINE iwzip #-} + +-- | DataFrames indexed by Nats and XNats require slightly different sets of +-- constraints to be sliced. +-- This family hides the difference, so that I could write one function for +-- both kinds. +type family CanSlice (t :: Type) (asbs :: [k]) :: Constraint where + CanSlice t (asbs :: [Nat]) + = ( PrimArray t (DataFrame t asbs), Dimensions asbs ) + CanSlice t (asbs :: [XNat]) + = ( ) + +-- | Get a few contiguous elements. +-- +-- In a sense, this is just a more complicated version of `sindex`. +-- +-- If (@b ~ XN m@) then this function is unsafe and can throw +-- an `OutOfDimBounds` exception. +-- Otherwise, its safety is guaranteed by the type system. +sslice :: + forall t b bi bd bs + . ( KnownDimKind (KindOfEl bs) + , CanSlice t (b :+ bs) + , SubFrameIndexCtx b bi bd, KnownDim bd + , PrimArray t (DataFrame t (bd :+ bs))) + => Idx bi -> DataFrame t (b :+ bs) -> DataFrame t (bd :+ bs) +sslice = slice @t @b @bi @bd @'[] @bs @(b :+ bs) . (:* U) +{-# INLINE[1] sslice #-} + +-- | Update a few contiguous elements. +-- +-- In a sense, this is just a more complicated version of `supdate`. +-- +-- If (@b ~ XN m@) and (@Idx bi + Dim bd > Dim b@), this function updates only as +-- many elements as fits into the dataframe along this dimension (possibly none). +supdateSlice :: + forall t b bi bd bs + . ( KnownDimKind (KindOfEl bs) + , CanSlice t (b :+ bs) + , SubFrameIndexCtx b bi bd + , KnownDim bd, ExactDims bs + , PrimArray t (DataFrame t (bd :+ bs))) + => Idx bi -> DataFrame t (bd :+ bs) + -> DataFrame t (b :+ bs) -> DataFrame t (b :+ bs) +supdateSlice = updateSlice @t @b @bi @bd @'[] @bs @(b :+ bs) . (:* U) +{-# INLINE[1] supdateSlice #-} + +-- | Get a few contiguous elements. +-- +-- In a sense, this is just a more complicated version of `index`. +-- +-- If any of the dims in @as@ or @b@ is unknown (@a ~ XN m@), +-- then this function is unsafe and can throw an `OutOfDimBounds` exception. +-- Otherwise, its safety is guaranteed by the type system. +slice :: + forall (t :: Type) b bi bd as bs asbs + . ( KnownDimKind (KindOfEl asbs) + , CanSlice t asbs + , SubFrameIndexCtx b bi bd, KnownDim bd + , ConcatList as (b :+ bs) asbs + , PrimArray t (DataFrame t (bd :+ bs))) + => Idxs (as +: bi) -> DataFrame t asbs -> DataFrame t (bd :+ bs) +slice i = case dimKind @(KindOfEl asbs) of + DimKNat + -> withArrayContent broadcast + (\steps off0 ba -> + let (off, bsteps) = getOffAndStepsSub (I# off0) steps i (dim @bd) + in fromElems bsteps (case off of I# o -> o) ba + ) + DimKXNat + -> \(XFrame (df :: DataFrame t asbsN)) -> withArrayContent broadcast + (\steps off0 ba -> + let (off, bsteps) = getOffAndStepsSub (I# off0) steps i (dim @bd) + in fromElems bsteps (case off of I# o -> o) ba + ) df +{-# INLINE[1] slice #-} + +-- | Update a few contiguous elements. +-- +-- In a sense, this is just a more complicated version of `update`. +-- +-- If any of the dims in @as@ is unknown (@a ~ XN m@), +-- you may happen to update data beyond dataframe bounds. +-- In this case, the original DataFrame is returned. +-- If (@b ~ XN m@) and (@Idx bi + Dim bd > Dim b@), this function updates only as +-- many elements as fits into the dataframe along this dimension (possibly none). +updateSlice :: + forall (t :: Type) b bi bd as bs asbs + . ( KnownDimKind (KindOfEl asbs) + , CanSlice t asbs + , SubFrameIndexCtx b bi bd + , KnownDim bd, ExactDims bs + , ConcatList as (b :+ bs) asbs + , PrimArray t (DataFrame t (bd :+ bs))) + => Idxs (as +: bi) -> DataFrame t (bd :+ bs) + -> DataFrame t asbs -> DataFrame t asbs +updateSlice i bdbsDf asbsDf = case dimKind @(KindOfEl asbs) of + DimKNat -> runST $ do + asbsM <- thawDataFrame asbsDf + copyDataFrame i bdbsDf asbsM + unsafeFreezeDataFrame asbsM + DimKXNat + | (XFrame (asbsDfN :: DataFrame t asbsN)) <- asbsDf -> runST $ do + asbsM <- thawDataFrame @t @asbsN asbsDfN + copyDataFrame i bdbsDf (castDataFrame @t @asbs asbsM) + XFrame <$> unsafeFreezeDataFrame asbsM +{-# INLINE[1] updateSlice #-} + + +-- | Get an element by its index in the dataframe. +-- This is a safe alternative to `(.!)` function when the index dimension +-- is not known at compile time (@a ~ XN m@). +slookup :: + forall t (a :: XNat) (bs :: [XNat]) + . (All KnownDimType bs, PrimBytes t) + => Idx a -> DataFrame t (a :+ bs) -> Maybe (DataFrame t bs) +slookup = Numeric.DataFrame.SubSpace.lookup @t @'[a] @bs @(a :+ bs) . (:*U) +{-# INLINE slookup #-} + +-- | Get an element by its index in the dataframe. +-- This is a safe alternative to `index` function when some of the dimensions +-- are not known at compile time (@d ~ XN m@). +lookup :: + forall t (as :: [XNat]) (bs :: [XNat]) (asbs :: [XNat]) + . (ConcatList as bs asbs, All KnownDimType bs, PrimBytes t) + => Idxs as -> DataFrame t asbs -> Maybe (DataFrame t bs) +lookup i (XFrame (df :: DataFrame t asbsN)) + | asbsN <- dims @asbsN + , Just off <- cdIxM (getSteps asbsN df) i + , asbs <- XDims asbsN :: Dims asbs + , (as, bs) <- prefSufDims i asbs + = withLiftedConcatList @'Nothing @'Nothing @('Just asbsN) as bs asbs $ + \(Dims :: Dims asn) (Dims :: Dims bsn) (_ :: Dims asbsn) + -> case inferKnownBackend @t @bsn of + Dict -> Just (XFrame (indexOffset @t @asn @bsn @asbsn off df)) + | otherwise = Nothing +{-# INLINE lookup #-} + +-- | Get a few contiguous elements. +-- +-- In a sense, this is just a more complicated version of `slookup`. +-- +-- This is a safe alternative to `sslice` function when the slice dimension +-- is not known at compile time (@b ~ XN m@). +ssliceMaybe :: + forall (t :: Type) (b :: XNat) (bi :: XNat) (bd :: XNat) (bs :: [XNat]) + . ( SubFrameIndexCtx b bi bd, KnownDim bd + , All KnownDimType bs, PrimBytes t) + => Idx bi -> DataFrame t (b :+ bs) -> Maybe (DataFrame t (bd :+ bs)) +ssliceMaybe = sliceMaybe @t @b @bi @bd @'[] @bs @(b :+ bs) . (:* U) +{-# INLINE ssliceMaybe #-} + + +-- | Get a few contiguous elements. +-- +-- In a sense, this is just a more complicated version of `lookup`. +-- +-- This is a safe alternative to `slice` function when some of the dimensions +-- are not known at compile time (@d ~ XN m@). +sliceMaybe :: + forall (t :: Type) (b :: XNat) bi bd as bs asbs + . ( SubFrameIndexCtx b bi bd, KnownDim bd + , ConcatList as (b :+ bs) asbs + , All KnownDimType bs, PrimBytes t) + => Idxs (as +: bi) -> DataFrame t asbs -> Maybe (DataFrame t (bd :+ bs)) +sliceMaybe i (XFrame (df :: DataFrame t asbsN)) + | asbsN <- dims @asbsN + , bd <- dim @bd + , Dict <- ( withKnownXDim @bd Dict + :: Dict (KnownDim (DimBound bd), KnownDimType bd, FixedDim bd (DimBound bd))) + , Just (I# off0, bsteps@(CumulDims (bbN:bN:_))) + <- getOffAndStepsSubM 0 (getSteps asbsN df) i bd + , bbN == bN * dimVal bd -- is there enough elements? + , asbs <- XDims asbsN :: Dims asbs + , XDims (Dims :: Dims bsN) <- + unsafeCoerce (dropSome (listIdxs i) (listDims asbs)) :: Dims bs + , Dict <- inferKnownBackend @t @(DimBound bd ': bsN) + = Just $ XFrame @_ @t @(bd :+ bs) @(DimBound bd ': bsN) + (withArrayContent broadcast + (\_ off -> fromElems bsteps (off0 +# off)) df) + | otherwise = Nothing +{-# INLINE sliceMaybe #-} + + +-- | Operations on DataFrames +-- +-- @as@ is an indexing dimensionality +-- +-- @bs@ is an element dimensionality +-- +-- @t@ is an underlying data type (i.e. Float, Int, Double) +-- +class ( ConcatList as bs asbs + , SubSpaceCtx t as bs asbs + , PrimBytes t + , KnownDimKind k + ) => SubSpace (t :: Type) (as :: [k]) (bs :: [k]) (asbs :: [k]) + | asbs as -> bs, asbs bs -> as, as bs -> asbs where + type SubSpaceCtx t as bs asbs :: Constraint + + joinDataFrameI :: PrimBytes (DataFrame t bs) + => DataFrame (DataFrame t bs) as -> DataFrame t asbs + indexOffsetI :: Int -> DataFrame t asbs -> DataFrame t bs + updateOffsetI :: Int -> DataFrame t bs -> DataFrame t asbs -> DataFrame t asbs + indexI :: Idxs as -> DataFrame t asbs -> DataFrame t bs + updateI :: Idxs as -> DataFrame t bs -> DataFrame t asbs -> DataFrame t asbs + ewmapI :: forall s (bs' :: [k]) (asbs' :: [k]) + . SubSpace s as bs' asbs' + => (DataFrame s bs' -> DataFrame t bs) + -> DataFrame s asbs' -> DataFrame t asbs + iwmapI :: forall s (bs' :: [k]) (asbs' :: [k]) + . SubSpace s as bs' asbs' + => (Idxs as -> DataFrame s bs' -> DataFrame t bs) + -> DataFrame s asbs' -> DataFrame t asbs + ewgenI :: Dimensions as => DataFrame t bs -> DataFrame t asbs + iwgenI :: Dimensions as => (Idxs as -> DataFrame t bs) -> DataFrame t asbs + ewfoldlI :: (b -> DataFrame t bs -> b) -> b -> DataFrame t asbs -> b + iwfoldlI :: (Idxs as -> b -> DataFrame t bs -> b) -> b -> DataFrame t asbs -> b + ewfoldrI :: (DataFrame t bs -> b -> b) -> b -> DataFrame t asbs -> b + iwfoldrI :: (Idxs as -> DataFrame t bs -> b -> b) -> b -> DataFrame t asbs -> b + elementWiseI :: forall (s :: Type) (bs' :: [k]) (asbs' :: [k]) (f :: Type -> Type) + . (Applicative f, SubSpace s as bs' asbs') + => (DataFrame s bs' -> f (DataFrame t bs)) + -> DataFrame s asbs' -> f (DataFrame t asbs) + elementWiseI = indexWiseI . const + {-# INLINE elementWiseI #-} + indexWiseI :: forall (s :: Type) (bs' :: [k]) (asbs' :: [k]) (f :: Type -> Type) + . (Applicative f, SubSpace s as bs' asbs') + => (Idxs as -> DataFrame s bs' -> f (DataFrame t bs)) + -> DataFrame s asbs' -> f (DataFrame t asbs) + +instance ( ConcatList as bs asbs + , SubSpaceCtx t as bs asbs + ) => SubSpace t (as :: [Nat]) (bs :: [Nat]) (asbs :: [Nat]) where + type SubSpaceCtx t as bs asbs + = ( Dimensions as, Dimensions bs, Dimensions asbs + , PrimArray t (DataFrame t asbs), PrimArray t (DataFrame t bs)) + + joinDataFrameI + | Dict <- inferKnownBackend @(DataFrame t bs) @as + = withArrayContent ewgenI + (\steps -> fromElems (steps `mappend` cumulDims (dims @bs))) + + indexOffsetI (I# off) = withArrayContent broadcast + (\steps off0 -> fromElems (dropPref (dims @as) steps) (off0 +# off)) + {-# INLINE indexOffsetI #-} + + updateOffsetI off bs asbs = runST $ do + asbsM <- thawDataFrame asbs + copyDataFrameOff off bs asbsM + unsafeFreezeDataFrame asbsM + {-# INLINE updateOffsetI #-} + + indexI i = withArrayContent broadcast + (\steps off0 -> fromElems (dropPref (dims @as) steps) + (case cdIx steps i of I# off -> off0 +# off)) + {-# INLINE indexI #-} + + updateI i bs asbs = runST $ do + asbsM <- thawDataFrame asbs + copyDataFrame' i bs asbsM + unsafeFreezeDataFrame asbsM + {-# INLINE updateI #-} + + ewmapI (f :: DataFrame s bs' -> DataFrame t bs) (asbs' :: DataFrame s asbs') + = case uniqueOrCumulDims asbs' of + Left e -> ewgen (f (broadcast e)) + Right _ -> runST $ do + let asTD = fromIntegral (totalDim' @as ) :: Int + bsTD = fromIntegral (totalDim' @bs ) :: Int + bs'TD = fromIntegral (totalDim' @bs') :: Int + asbsM <- newDataFrame + forM_ [0..asTD-1] $ \i -> + copyDataFrameOff (i*bsTD) (f (indexOffset (i*bs'TD) asbs')) asbsM + unsafeFreezeDataFrame asbsM + {-# INLINE ewmapI #-} + + iwmapI f (asbs' :: DataFrame s asbs') + = case uniqueOrCumulDims asbs' of + Left e -> iwgen (flip f (broadcast e)) + Right _ -> runST $ do + asbsM <- newDataFrame + forM_ [minBound..maxBound] $ \i -> + copyDataFrame' i (f i (index i asbs')) asbsM + unsafeFreezeDataFrame asbsM + {-# INLINE iwmapI #-} + + ewgenI bs = case uniqueOrCumulDims bs of + Left a -> broadcast a + Right bsSteps -> runST $ do + let asTD = fromIntegral (totalDim' @as) :: Int + bsTD = fromIntegral (cdTotalDim bsSteps) + asbsM <- newDataFrame + forM_ [0, bsTD .. asTD * bsTD - 1] $ \o -> copyDataFrameOff o bs asbsM + unsafeFreezeDataFrame asbsM + {-# INLINE ewgenI #-} + + iwgenI f = runST $ do + asbsM <- newDataFrame + forM_ [minBound..maxBound] $ \i -> copyDataFrame' i (f i) asbsM + unsafeFreezeDataFrame asbsM + {-# INLINE iwgenI #-} + + ewfoldlI f x0 asbs = case uniqueOrCumulDims asbs of + Left a -> + let b = broadcast a + go i | i < 0 = x0 + | otherwise = f (go (i-1)) b + in go (fromIntegral (totalDim' @as) - 1 :: Int) + Right asbsSteps -> + let bsSteps = dropPref (dims @as) asbsSteps + bsLen = max 1 (fromIntegral (cdTotalDim bsSteps)) :: Int + asbsLen = fromIntegral (cdTotalDim asbsSteps) :: Int + go i | i < 0 = x0 + | otherwise = f (go (i - bsLen)) (indexOffset i asbs) + in go (asbsLen - bsLen) + {-# INLINE ewfoldlI #-} + + iwfoldlI f x0 asbs + | not (nonVoidDims (dims @as)) = x0 + | otherwise = case uniqueOrCumulDims asbs of + Left a -> + let b = broadcast a + go i | i == minBound = f i x0 b + | otherwise = f i (go (pred i)) b + in go maxBound + Right _ -> + let go i | i == minBound = f i x0 (index i asbs) + | otherwise = f i (go (pred i)) (index i asbs) + in go maxBound + {-# INLINE iwfoldlI #-} + + ewfoldrI f x0 asbs = case uniqueOrCumulDims asbs of + Left a -> + let b = broadcast a + go i | i >= totalDim' @as = x0 + | otherwise = f b (go (i+1)) + in go 0 + Right asbsSteps -> + let bsSteps = dropPref (dims @as) asbsSteps + bsLen = max 1 (fromIntegral (cdTotalDim bsSteps)) :: Int + asbsLen = fromIntegral (cdTotalDim asbsSteps) :: Int + go i | i >= asbsLen = x0 + | otherwise = f (indexOffset i asbs) (go (i + bsLen)) + in go 0 + {-# INLINE ewfoldrI #-} + + iwfoldrI f x0 asbs + | not (nonVoidDims (dims @as)) = x0 + | otherwise = case uniqueOrCumulDims asbs of + Left a -> + let b = broadcast a + go i | i == maxBound = f i b x0 + | otherwise = f i b (go (succ i)) + in go minBound + Right _ -> + let go i | i == maxBound = f i (index i asbs) x0 + | otherwise = f i (index i asbs) (go (succ i)) + in go minBound + {-# INLINE iwfoldrI #-} + + indexWiseI (f :: Idxs as -> DataFrame s bs' -> f (DataFrame t bs)) + = fmap mkSTDF . iwfoldr fST (pure $ MakingDF (const (pure ()))) + where + mkSTDF :: MakingDF t asbs -> DataFrame t asbs + mkSTDF st = runST $ do + df <- newDataFrame + getMDF st df + unsafeFreezeDataFrame df + fST :: Idxs as -> DataFrame s bs' + -> f (MakingDF t asbs) -> f (MakingDF t asbs) + fST i bs' asbsSTF + = (\r st -> MakingDF (\x -> copyDataFrame' i r x *> getMDF st x)) + <$> f i bs' <*> asbsSTF + {-# INLINE indexWiseI #-} + +-- | Hide ST monad state with a DataFrame in works inside a plain type. +newtype MakingDF t asbs + = MakingDF { getMDF :: forall s . STDataFrame s t asbs -> ST s () } + +-- | Checks if all of the dimensions are non-zero. +nonVoidDims :: Dims ns -> Bool +nonVoidDims = all (0 <) . listDims + +dropPref :: Dims (ns :: [Nat]) -> CumulDims -> CumulDims +dropPref ds = CumulDims . dropSome (listDims ds) . unCumulDims + +dropSome :: [Word] -> [Word] -> [Word] +dropSome [] xs = xs +dropSome (_:as) xs = dropSome as (tail xs) + +instance ( ConcatList as bs asbs + , SubSpaceCtx t as bs asbs + ) => SubSpace t (as :: [XNat]) (bs :: [XNat]) (asbs :: [XNat]) where + type SubSpaceCtx t as bs asbs + = ( Dimensions bs, PrimBytes t, ExactDims bs + , All KnownDimType as, All KnownDimType bs, All KnownDimType asbs) + + joinDataFrameI (XFrame (df :: DataFrame (DataFrame t bs) asN)) = + let as = XDims (dims @asN) :: Dims as + bs = dims :: Dims bs + asbs = concatDims as bs :: Dims asbs + in withLiftedConcatList @('Just asN) @('Just (DimsBound bs)) @'Nothing + as bs asbs $ + \(Dims :: Dims asn) (Dims :: Dims bsn) (Dims :: Dims asbsn) -> case () of + _ | Dict <- inferKnownBackend @t @bsn + , Dict <- inferKnownBackend @t @asbsn + -> XFrame @_ @t @asbs @asbsn $ withArrayContent + -- know for sure its DataFrame t bsn + (ewgen @t @asn @bsn @asbsn . unliftXFrame) + (\steps -> fromElems (steps `mappend` cumulDims (dims @bs))) + df + {-# INLINE joinDataFrameI #-} + + indexOffsetI off (XFrame (df :: DataFrame t asbsN)) = + let bs = dims :: Dims bs + as = dropSufDims bs asbs :: Dims as + asbs = XDims (dims @asbsN) :: Dims asbs + in withLiftedConcatList @'Nothing @('Just (DimsBound bs)) @('Just asbsN) + as bs asbs $ + \(Dims :: Dims asn) (Dims :: Dims bsn) (Dims :: Dims asbsn) + -> case inferKnownBackend @t @bsn of + Dict -> XFrame (indexOffset @t @asn @bsn @asbsn off df) + {-# INLINE indexOffsetI #-} + + updateOffsetI off (XFrame (bsDf :: DataFrame t bsN)) + (XFrame (asbsDf :: DataFrame t asbsN)) = + let bs = dims :: Dims bs + as = dropSufDims bs asbs :: Dims as + asbs = XDims (dims @asbsN) :: Dims asbs + in withLiftedConcatList @'Nothing @('Just bsN) @('Just asbsN) + as bs asbs $ + \(Dims :: Dims asn) (Dims :: Dims bsn) (Dims :: Dims asbsn) + -> XFrame (updateOffset @t @asn @bsn @asbsn off bsDf asbsDf) + {-# INLINE updateOffsetI #-} + + indexI i (XFrame (df :: DataFrame t asbsN)) = + let bs = dims :: Dims bs + as = dropSufDims bs asbs :: Dims as + asbs = XDims (dims @asbsN) :: Dims asbs + in withLiftedConcatList @'Nothing @'Nothing @('Just asbsN) + as bs asbs $ + \(Dims :: Dims asn) (Dims :: Dims bsn) (Dims :: Dims asbsn) + -> case inferKnownBackend @t @bsn of + Dict -> XFrame (index @t @asn @bsn @asbsn (unsafeUnliftIdxs i) df) + {-# INLINE indexI #-} + + updateI i (XFrame (bsDf :: DataFrame t bsN)) + (XFrame (asbsDf :: DataFrame t asbsN)) = + let bs = dims :: Dims bs + as = dropSufDims bs asbs :: Dims as + asbs = XDims (dims @asbsN) :: Dims asbs + in withLiftedConcatList @'Nothing @('Just bsN) @('Just asbsN) + as bs asbs $ + \(Dims :: Dims asn) (Dims :: Dims bsn) (Dims :: Dims asbsn) + -> XFrame (update @t @asn @bsn @asbsn (unsafeUnliftIdxs i) bsDf asbsDf) + {-# INLINE updateI #-} + + ewmapI :: forall s (bs' :: [XNat]) (asbs' :: [XNat]) + . SubSpace s as bs' asbs' + => (DataFrame s bs' -> DataFrame t bs) + -> DataFrame s asbs' -> DataFrame t asbs + ewmapI f (XFrame (df :: DataFrame s asbsN')) = + let bs = dims :: Dims bs + bs' = dims :: Dims bs' + as = dropSufDims bs' asbs' :: Dims as + asbs' = XDims (dims @asbsN') :: Dims asbs' + asbs = concatDims as bs :: Dims asbs + in withLiftedConcatList @'Nothing @'Nothing @('Just asbsN') + as bs' asbs' $ + \(Dims :: Dims asn) (Dims :: Dims bsn') (Dims :: Dims asbsn') -> + withLiftedConcatList @('Just asn) @'Nothing @'Nothing + as bs asbs $ + \(Dims :: Dims asn) (Dims :: Dims bsn) (Dims :: Dims asbsn) -> + case inferKnownBackend @t @asbsn of + Dict + | Dict <- inferKnownBackend @t @bsn + , Dict <- inferKnownBackend @s @bsn' + -> XFrame (ewmap @t @asn @bsn @asbsn @s @bsn' @asbsn' + (unliftXFrame . f . XFrame) df) + {-# INLINE ewmapI #-} + + iwmapI :: forall s (bs' :: [XNat]) (asbs' :: [XNat]) + . SubSpace s as bs' asbs' + => (Idxs as -> DataFrame s bs' -> DataFrame t bs) + -> DataFrame s asbs' -> DataFrame t asbs + iwmapI f (XFrame (df :: DataFrame s asbsN')) = + let bs = dims :: Dims bs + bs' = dims :: Dims bs' + as = dropSufDims bs' asbs' :: Dims as + asbs' = XDims (dims @asbsN') :: Dims asbs' + asbs = concatDims as bs :: Dims asbs + in withLiftedConcatList @'Nothing @'Nothing @('Just asbsN') + as bs' asbs' $ + \(Dims :: Dims asn) (Dims :: Dims bsn') (Dims :: Dims asbsn') -> + withLiftedConcatList @('Just asn) @'Nothing @'Nothing + as bs asbs $ + \(Dims :: Dims asn) (Dims :: Dims bsn) (Dims :: Dims asbsn) -> case () of + _ | Dict <- inferKnownBackend @t @asbsn + , Dict <- inferKnownBackend @t @bsn + , Dict <- inferKnownBackend @s @bsn' + -> XFrame (iwmap @t @asn @bsn @asbsn @s @bsn' @asbsn' + (\i x -> unliftXFrame (f (liftIdxs i) (XFrame x))) df) + {-# INLINE iwmapI #-} + + + ewgenI (XFrame (df :: DataFrame t bsN)) = + let as = dims :: Dims as + bs = dims :: Dims bs + asbs = concatDims as bs :: Dims asbs + in withLiftedConcatList @'Nothing @('Just bsN) @'Nothing as bs asbs $ + \(Dims :: Dims asn) (Dims :: Dims bsn) (Dims :: Dims asbsn) -> + case inferKnownBackend @t @asbsn of + Dict -> XFrame (ewgen @t @asn @bsn @asbsn df) + {-# INLINE ewgenI #-} + + iwgenI f = + let as = dims :: Dims as + bs = dims :: Dims bs + asbs = concatDims as bs :: Dims asbs + in withLiftedConcatList @'Nothing @'Nothing @'Nothing as bs asbs $ + \(Dims :: Dims asn) (Dims :: Dims bsn) (Dims :: Dims asbsn) -> case () of + _ | Dict <- inferKnownBackend @t @asbsn + , Dict <- inferKnownBackend @t @bsn + -> XFrame (iwgen @t @asn @bsn @asbsn (unliftXFrame . f . liftIdxs)) + {-# INLINE iwgenI #-} + + ewfoldlI f x0 (XFrame (df :: DataFrame s asbsN)) = + let as = dropSufDims bs asbs :: Dims as + bs = dims :: Dims bs + asbs = XDims (dims @asbsN) :: Dims asbs + in withLiftedConcatList @'Nothing @'Nothing @('Just asbsN) as bs asbs $ + \(Dims :: Dims asn) (Dims :: Dims bsn) (Dims :: Dims asbsn) -> + case inferKnownBackend @t @bsn of + Dict -> ewfoldl @t @asn @bsn @asbsn (\b x -> f b (XFrame x)) x0 df + {-# INLINE ewfoldlI #-} + + iwfoldlI f x0 (XFrame (df :: DataFrame s asbsN)) = + let as = dropSufDims bs asbs :: Dims as + bs = dims :: Dims bs + asbs = XDims (dims @asbsN) :: Dims asbs + in withLiftedConcatList @'Nothing @'Nothing @('Just asbsN) as bs asbs $ + \(Dims :: Dims asn) (Dims :: Dims bsn) (Dims :: Dims asbsn) -> + case inferKnownBackend @t @bsn of + Dict -> iwfoldl @t @asn @bsn @asbsn + (\i b x -> f (liftIdxs i) b (XFrame x)) x0 df + {-# INLINE iwfoldlI #-} + + ewfoldrI f x0 (XFrame (df :: DataFrame s asbsN)) = + let as = dropSufDims bs asbs :: Dims as + bs = dims :: Dims bs + asbs = XDims (dims @asbsN) :: Dims asbs + in withLiftedConcatList @'Nothing @'Nothing @('Just asbsN) as bs asbs $ + \(Dims :: Dims asn) (Dims :: Dims bsn) (Dims :: Dims asbsn) -> + case inferKnownBackend @t @bsn of + Dict -> ewfoldr @t @asn @bsn @asbsn (\x -> f (XFrame x)) x0 df + {-# INLINE ewfoldrI #-} + + iwfoldrI f x0 (XFrame (df :: DataFrame s asbsN)) = + let as = dropSufDims bs asbs :: Dims as + bs = dims :: Dims bs + asbs = XDims (dims @asbsN) :: Dims asbs + in withLiftedConcatList @'Nothing @'Nothing @('Just asbsN) as bs asbs $ + \(Dims :: Dims asn) (Dims :: Dims bsn) (Dims :: Dims asbsn) -> + case inferKnownBackend @t @bsn of + Dict -> iwfoldr @t @asn @bsn @asbsn + (\i x -> f (liftIdxs i) (XFrame x)) x0 df + {-# INLINE iwfoldrI #-} + + indexWiseI :: forall s (bs' :: [XNat]) (asbs' :: [XNat]) f + . (Applicative f, SubSpace s as bs' asbs') + => (Idxs as -> DataFrame s bs' -> f (DataFrame t bs)) + -> DataFrame s asbs' -> f (DataFrame t asbs) + indexWiseI f (XFrame (df :: DataFrame s asbsN')) = + let bs = dims :: Dims bs + bs' = dims :: Dims bs' + as = dropSufDims bs' asbs' :: Dims as + asbs' = XDims (dims @asbsN') :: Dims asbs' + asbs = concatDims as bs :: Dims asbs + in withLiftedConcatList @'Nothing @'Nothing @('Just asbsN') + as bs' asbs' $ + \(Dims :: Dims asn) (Dims :: Dims bsn') (Dims :: Dims asbsn') -> + withLiftedConcatList @('Just asn) @'Nothing @'Nothing + as bs asbs $ + \(Dims :: Dims asn) (Dims :: Dims bsn) (Dims :: Dims asbsn) -> case () of + _ | Dict <- inferKnownBackend @t @asbsn + , Dict <- inferKnownBackend @t @bsn + , Dict <- inferKnownBackend @s @bsn' + -> XFrame <$> + (indexWise @t @asn @bsn @asbsn @s @bsn' @asbsn' + (\i x -> unliftXFrame <$> f (liftIdxs i) (XFrame x)) df + ) + {-# INLINE indexWiseI #-} + +dropSufDims :: ConcatList as bs asbs + => Dims bs -> Dims asbs -> Dims as +dropSufDims = unsafeCoerce dropSuf + where + dropSuf :: [Word] -> [Word] -> [Word] + dropSuf bs asbs = fst $ foldr f ([], bs) asbs + f :: Word -> ([Word], [Word]) -> ([Word], [Word]) + f x (as, []) = (x:as,[]) + f _ (_, _:rs) = ([], rs) + +prefSufDims :: ConcatList as bs asbs + => TypedList f as -> Dims asbs -> (Dims as, Dims bs) +prefSufDims = unsafeCoerce f + where + f :: [Word] -> [Word] -> ([Word], [Word]) + f [] asbs = ([], asbs) + f (_:is) ~(a:sbs) = first (a:) (f is sbs) + +concatDims :: ConcatList as bs asbs + => Dims as -> Dims bs -> Dims asbs +concatDims = TL.concat + +{- +A trick here is that both alternatives of 'MayEq' have the same representation +of a single equality constraint. +Thus, I can unsafeCoerce an equality later into MayEq to significantly +reduce amount of boilerplate. + -} +type family MayEq (mns :: (Maybe [Nat])) (ns :: [Nat]) :: Constraint where + MayEq 'Nothing bs = bs ~ bs + MayEq ('Just as) bs = as ~ bs + +-- Very unsafe function, because it does not check the content of XDims at all. +withLiftedConcatList :: + forall (asm :: Maybe [Nat]) (bsm :: Maybe [Nat]) (asbsm :: Maybe [Nat]) + (asx :: [XNat]) (bsx :: [XNat]) (asbsx :: [XNat]) r + . -- (ConcatList asx bsx asbsx) => + Dims asx -> Dims bsx -> Dims asbsx + -> ( forall (asn :: [Nat]) (bsn :: [Nat]) (asbsn :: [Nat]) + . ( ConcatList asn bsn asbsn + , FixedDims asx asn, FixedDims bsx bsn, FixedDims asbsx asbsn + , MayEq asm asn, MayEq bsm bsn, MayEq asbsm asbsn) + => Dims asn -> Dims bsn -> Dims asbsn -> r + ) -> r +withLiftedConcatList (XDims (asn :: Dims asn)) + (XDims (bsn :: Dims bsn)) + (XDims (asbsn :: Dims asbsn)) k + | Dict <- unsafeEqTypes @_ @asbsn @(Concat asn bsn) + , Dict <- unsafeCoerce (Dict @(asn ~ asn)) :: Dict (MayEq asm asn) + , Dict <- unsafeCoerce (Dict @(bsn ~ bsn)) :: Dict (MayEq bsm bsn) + , Dict <- unsafeCoerce (Dict @(asbsn ~ asbsn)) :: Dict (MayEq asbsm asbsn) + , Dict <- inferConcat @asn @bsn @asbsn = k asn bsn asbsn + +unliftXFrame :: forall (t :: Type) (xns :: [XNat]) (ns :: [Nat]) + . (ExactDims xns, FixedDims xns ns) + => DataFrame t xns -> DataFrame t ns +unliftXFrame (XFrame x) = unsafeCoerce x +{-# INLINE unliftXFrame #-} + +unsafeEqTypes :: forall k (a :: k) (b :: k) + . Dict (a ~ b) +unsafeEqTypes = unsafeCoerce (Dict :: Dict (a ~ a)) diff --git a/src/Numeric/DataFrame/Type.hs b/src/Numeric/DataFrame/Type.hs --- a/src/Numeric/DataFrame/Type.hs +++ b/src/Numeric/DataFrame/Type.hs @@ -12,6 +12,7 @@ {-# LANGUAGE MagicHash #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE PolyKinds #-} {-# LANGUAGE RankNTypes #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE StandaloneDeriving #-} @@ -22,7 +23,9 @@ {-# LANGUAGE UnboxedTuples #-} {-# LANGUAGE UndecidableInstances #-} {-# LANGUAGE ViewPatterns #-} - +#ifdef UNSAFE_INDICES +{-# OPTIONS_GHC -fno-warn-redundant-constraints #-} +#endif {- | The core @easytensor@ types. -} @@ -38,24 +41,29 @@ , DataFrame ( SingleFrame, MultiFrame, XFrame, (:*:), Z , S, DF2, DF3, DF4, DF5, DF6, DF7, DF8, DF9) #endif + , scalar, unScalar + , IndexFrame (..), SubFrameIndexCtx -- * Flexible assembling and disassembling , PackDF, packDF, unpackDF , appendDF, consDF, snocDF , fromFlatList, fromListWithDefault, fromList - , constrainDF + , constrainDF, asDiag -- * Infer type class instances , KnownBackend (), DFBackend, KnownBackends , InferKnownBackend (..), inferPrimElem -- * Re-exports - , Dim (..), Idx (), XNat (..), N, XN, Dims, Idxs, TypedList (..) + , Dim (..), Idx (), Nat, XNat (..), N, XN, Dims, Idxs, TypedList (..) , PrimBytes (), bSizeOf, bAlignOf, bFieldOffsetOf , PrimArray (), ixOff ) where +#if !(MIN_VERSION_base(4,13,0)) +import Data.Proxy (Proxy) +#endif +import Control.Arrow (second, (***)) import Data.Data -import Data.Proxy (Proxy) -import Data.Semigroup hiding (All) +import Data.Semigroup hiding (All, Min) import Data.Type.Lits import Data.Void import Foreign.Storable (Storable (..)) @@ -66,7 +74,9 @@ import qualified Text.ParserCombinators.ReadPrec as Read import qualified Text.Read as Read import qualified Text.Read.Lex as Read +import Unsafe.Coerce +import Numeric.Basics import Numeric.DataFrame.Internal.PrimArray import Numeric.Dimensions import Numeric.PrimBytes @@ -75,12 +85,11 @@ import qualified Numeric.ProductOrd.Partial as Partial import Numeric.TypedList (typeables) -import {-# SOURCE #-} Numeric.DataFrame.Internal.Backend (DFBackend, - KnownBackend) -import {-# SOURCE #-} qualified Numeric.DataFrame.Internal.Backend as Backend -import Numeric.DataFrame.Internal.BackendI () +import Numeric.DataFrame.Internal.BackendI (DFBackend, KnownBackend) +import qualified Numeric.DataFrame.Internal.BackendI as Backend + -- | Keep data in a primitive data frame -- and maintain information about Dimensions in the type system data family DataFrame (t :: l) (xs :: [k]) @@ -97,7 +106,7 @@ -- Pattern-match against its constructor to get a Nat-indexed data frame. data instance DataFrame (ts :: l) (xns :: [XNat]) = forall (ns :: [Nat]) - . (KnownXNatTypes xns, FixedDims xns ns, Dimensions ns, KnownBackends ts ns) + . (All KnownDimType xns, FixedDims xns ns, Dimensions ns, KnownBackends ts ns) => XFrame (DataFrame ts ns) -- | Data frame that has an unknown dimensionality at compile time. @@ -113,21 +122,36 @@ {-# COMPLETE Z, (:*:) #-} +#define PLEASE_STYLISH_HASKELL \ + forall (xs :: [Type]) (ns :: [Nat]) . () => \ + forall (y :: Type) (ys :: [Type]) . (xs ~ (y ': ys)) => \ + DataFrame y ns -> DataFrame ys ns -> DataFrame xs ns -- | Constructing a @MultiFrame@ using DataFrame columns -pattern (:*:) :: forall (xs :: [Type]) (ns :: [Nat]) . () - => forall (y :: Type) (ys :: [Type]) . (xs ~ (y ': ys)) - => DataFrame y ns -> DataFrame ys ns -> DataFrame xs ns +pattern (:*:) :: PLEASE_STYLISH_HASKELL pattern (:*:) x xs <- (MultiFrame (DataFrame' x :* (MultiFrame -> xs))) where (:*:) x (MultiFrame xs) = MultiFrame (DataFrame' x :* xs) infixr 6 :*: +#undef PLEASE_STYLISH_HASKELL -- | Empty MultiFrame pattern Z :: forall (xs :: [Type]) (ns :: [Nat]) . () => (xs ~ '[]) => DataFrame xs ns pattern Z = MultiFrame U +-- | Convert a scalar back to an ordinary type +unScalar :: DataFrame t ('[] :: [Nat]) -> t +-- rely on that Scalar is just two times newtype alias to t +unScalar = unsafeCoerce +{-# INLINE unScalar #-} + +-- | Convert any type to a scalar wrapper +scalar :: t -> DataFrame t ('[] :: [Nat]) +-- rely on that Scalar is just two times newtype alias to t +scalar = unsafeCoerce +{-# INLINE scalar #-} + -- | I use this kind-polymorphic constraint to generalize @XFrame@ and @SomeDataFrame@ -- over @SingleFrame@ and @MultiFrame@. type family KnownBackends (ts :: l) (ns :: [Nat]) :: Constraint where @@ -137,8 +161,8 @@ (KnownBackend t ns, KnownBackends ts ns) -- | Allow inferring @KnownBackends@ if you know the dimensions and the element types. -class InferKnownBackend (t :: k) ds where - -- Infer @KnownBackends@ if you know the dimensions and the element types. +class InferKnownBackend t (ds :: [Nat]) where + -- | Infer @KnownBackends@ if you know the dimensions and the element types. inferKnownBackend :: Dict (KnownBackends t ds) instance (PrimBytes t, Dimensions ds) => InferKnownBackend (t :: Type) ds where @@ -146,15 +170,17 @@ instance (RepresentableList ts, All PrimBytes ts, Dimensions ds) => InferKnownBackend (ts :: [Type]) ds where - inferKnownBackend = go (tList @_ @ts) + inferKnownBackend = go (tList @ts) where - go :: forall ss . All PrimBytes ss => TypeList ss -> Dict (KnownBackends ss ds) + go :: forall ss . All PrimBytes ss + => TypeList ss -> Dict (KnownBackends ss ds) go U = Dict go ((_ :: Proxy t) :* ts) = case Backend.inferKnownBackend @t @ds of Dict -> case go ts of Dict -> Dict + -- | All component data frames must satisfy a given constraint. type family AllFrames (f :: Type -> Constraint) (ts :: [Type]) (ds :: [Nat]) :: Constraint where @@ -162,6 +188,107 @@ AllFrames f (t ': ts) ds = (f (DataFrame t ds), AllFrames f ts ds) +-- | Index one dimension deep into a @DataFrame@. +class IndexFrame t d ds where + -- | Index one dimension deep into a @DataFrame@. + -- + -- Note, this function does not provide indexing safety at the type level; + -- it throws an `OutOfDimBounds` exception if an index is out of bounds + -- (unless @unsafeindices@ package flag is enabled, which is even more dangerous). + (!) :: DataFrame t (d ': ds) -> Word -> DataFrame t ds + +instance (PrimArray t (DataFrame t '[d]), KnownDim d) + => IndexFrame (t :: Type) (d :: Nat) '[] where + {-# INLINE (!) #-} + (!) df i +#ifndef UNSAFE_INDICES + | i >= dimVal' @d = outOfDimBoundsNoCallStack "IndexFrame.(!)" + i (dimVal' @d) Nothing Nothing + | otherwise +#endif + = S (ixOff (fromIntegral i) df) + +instance {-# INCOHERENT #-} + ( PrimArray t (DataFrame t (d ': ds)) + , KnownDim d, KnownBackend t ds) + => IndexFrame (t :: Type) (d :: Nat) ds where + {-# INLINE (!) #-} + (!) df i +#ifndef UNSAFE_INDICES + | i >= dimVal' @d = outOfDimBoundsNoCallStack "IndexFrame.(!)" + i (dimVal' @d) Nothing Nothing + | otherwise +#endif + = withArrayContent broadcast + (\(CumulDims ~(_:ss@(s:_))) off0 -> + fromElems (CumulDims ss) (case s*i of W# off -> off0 +# word2Int# off) + ) df + + +instance IndexFrame ('[] :: [Type]) (d :: Nat) ds where + {-# INLINE (!) #-} + (!) Z _ = Z + +instance (IndexFrame t d ds, IndexFrame ts d ds) + => IndexFrame ((t ': ts) :: [Type]) (d :: Nat) ds where + {-# INLINE (!) #-} + (!) (df :*: dfs) i = df ! i :*: dfs ! i + +instance PrimBytes t + => IndexFrame (t :: Type) (xd :: XNat) xds where + {-# INLINE (!) #-} + (!) (XFrame df) i + | D :* (Dims :: Dims ns) <- dims `inSpaceOf` df + , Dict <- inferKnownBackend @t @ns + = XFrame (df ! i) + (!) _ _ = error "IndexFrame: impossible pattern" + +instance (RepresentableList ts, All PrimBytes ts) + => IndexFrame (ts :: [Type]) (xd :: XNat) xds where + {-# INLINE (!) #-} + (!) (XFrame df) i + | (D :: Dim n) :* (Dims :: Dims ns) <- dims `inSpaceOf` df + , Dict <- getTsEvs @ts @n @ns tList + = XFrame (df ! i) + where + getTsEvs :: forall (as :: [Type]) (d :: Nat) (ds :: [Nat]) + . ( All PrimBytes as, KnownBackends as (d ': ds) + , Dimensions ds, KnownDim d) + => TypeList as -> Dict (IndexFrame as d ds, KnownBackends as ds) + getTsEvs U = Dict + getTsEvs ((_ :: Proxy a) :* as') + | Dict <- getTsEvs @_ @d @ds as' + , Dict <- inferKnownBackend @a @ds = Dict + (!) _ _ = error "IndexFrame: impossible pattern" + +{- | +This type family describes two strategies for dealing with dimensions when +slicing a @DataFrame@. + +If the original space dimension (@d@) is fixed at compile time +(@d :: Nat@ or @d ~ N n@), then we enforce the safety of a slicing operation +with types: + the index space is bounded by the size of the original space + minus the subframe size. + +If the original space dimension (@d@) is not fixed (@d ~ XN m@), then we give up. +Just let the user do the slicing as easy as possible, and throw an `OutOfDimBounds` +exception at runtime if the index plus the sliced size is greater than the +original size of the DataFrame. + -} +type family SubFrameIndexCtx (d :: k) (idxN :: k) (subN :: k) :: Constraint where + SubFrameIndexCtx (n :: Nat) idxN subN + = (n + 1) ~ (idxN + subN) + SubFrameIndexCtx (N n) idxN subN + = ( (n + 1) ~ (DimBound idxN + DimBound subN) + , idxN ~ N (DimBound idxN) + , subN ~ N (DimBound subN) + ) + SubFrameIndexCtx (XN m) idxN subN + = ( idxN ~ XN m + , subN ~ N (DimBound subN) + ) + deriving instance Typeable (DataFrame (t :: l) (xs :: [k])) deriving instance ( Data (DataFrame t xs) , Typeable t, Typeable xs, Typeable k, Typeable l) @@ -184,17 +311,24 @@ => Fractional (DataFrame t ds) deriving instance Floating (DFBackend t ds) => Floating (DataFrame t ds) +deriving instance Epsilon (DFBackend t ds) + => Epsilon (DataFrame t ds) deriving instance Real (DFBackend t ds) => Real (DataFrame t ds) +deriving instance RealExtras (DFBackend t ds) + => RealExtras (DataFrame t ds) deriving instance RealFrac (DFBackend t ds) => RealFrac (DataFrame t ds) deriving instance RealFloat (DFBackend t ds) => RealFloat (DataFrame t ds) +deriving instance RealFloatExtras (DFBackend t ds) + => RealFloatExtras (DataFrame t ds) deriving instance PrimBytes (DFBackend t ds) => PrimBytes (DataFrame t ds) deriving instance (PrimArray t (DFBackend t ds), PrimBytes t) => PrimArray t (DataFrame t ds) + instance Ord (NonTransitive.ProductOrd (DFBackend t ds)) => Ord (NonTransitive.ProductOrd (DataFrame t ds)) where (>) = coerce ((>) @(NonTransitive.ProductOrd (DFBackend t ds))) @@ -263,7 +397,7 @@ ) => Show (DataFrame (t :: Type) (ds :: [Nat])) where showsPrec p x = case dims @ds of U -> showParen (p >= 10) - $ showString "S " . showsPrec 10 (unsafeCoerce# x :: t) + $ showString "S " . showsPrec 10 (unScalar x) D0 :* _ -> showString "DF0" d :* _ -> showParen (p >= 10) $ unpackDF' ( \Dict f -> @@ -310,24 +444,24 @@ instance (All Read ts, All PrimBytes ts, RepresentableList ts, Dimensions ds) => Read (DataFrame (ts :: [Type]) (ds :: [Nat])) where - readPrec = readFixedMultiDF (tList @Type @ts) (dims @ds) + readPrec = readFixedMultiDF (tList @ts) (dims @ds) readList = Read.readListDefault readListPrec = Read.readListPrecDefault -instance (Read t, PrimBytes t, BoundedDims ds, All KnownXNatType ds) +instance (Read t, PrimBytes t, BoundedDims ds, All KnownDimType ds) => Read (DataFrame (t :: Type) (ds :: [XNat])) where readPrec = Read.parens . Read.prec 10 $ do Read.lift . Read.expect $ Read.Ident "XFrame" - Read.step $ readPrecBoundedDF (minDims @XNat @ds) + Read.step $ readPrecBoundedDF (minimalDims @ds) readList = Read.readListDefault readListPrec = Read.readListPrecDefault instance ( All Read ts, All PrimBytes ts, RepresentableList ts - , BoundedDims ds, All KnownXNatType ds) + , BoundedDims ds, All KnownDimType ds) => Read (DataFrame (ts :: [Type]) (ds :: [XNat])) where readPrec = Read.parens . Read.prec 10 $ do Read.lift . Read.expect $ Read.Ident "XFrame" - Read.step $ readBoundedMultiDF (tList @Type @ts) (minDims @XNat @ds) + Read.step $ readBoundedMultiDF (tList @ts) (minimalDims @ds) readList = Read.readListDefault readListPrec = Read.readListPrecDefault @@ -343,7 +477,7 @@ => Read (SomeDataFrame (ts :: [Type])) where readPrec = Read.parens . Read.prec 10 $ do Read.lift . Read.expect $ Read.Ident "SomeDataFrame" - Read.step $ readSomeMultiDF (tList @Type @ts) + Read.step $ readSomeMultiDF (tList @ts) readList = Read.readListDefault readListPrec = Read.readListPrecDefault @@ -381,7 +515,7 @@ -} readPrecBoundedDF :: forall (t :: Type) (ds :: [XNat]) - . (Read t, PrimBytes t, KnownXNatTypes ds) + . (Read t, PrimBytes t, All KnownDimType ds) => Dims ds -> Read.ReadPrec (DataFrame t ds) -- reading scalar readPrecBoundedDF U @@ -414,7 +548,7 @@ = Read.parens $ lookLex >>= \case Read.Ident ('D':'F':s) | Just (Dx (n :: Dim n)) <- (Read.readMaybe ('D':s) :: Maybe SomeDim) - , Just Dict <- indeedLE m n -- check if the DF dim is not less than m + , Just Dict <- lessOrEqDim m n -> case n of D0 -> do -- need to repack it under different constraint dims XFrame x <- readPrecBoundedDF @t (Dn D0 :* xns) @@ -428,12 +562,6 @@ (followedBy x) _ -> Read.pfail - where - indeedLE :: forall (a ::Nat) (b :: Nat) . Dim a -> Dim b -> Maybe (Dict (a <= b)) - indeedLE a b = case compareDim a b of - SLT -> Just Dict - SEQ -> Just Dict - SGT -> Nothing {- In this case we know Nothing about the dimensionality of a DataFrame. @@ -451,7 +579,7 @@ Read.+++ (lookLex >>= \case Read.Ident ('D':'F':s) - | Just (Dx (d :: Dim d)) <- (Read.readMaybe ('D':s) :: Maybe SomeDim) + | Just (Dx (d@D :: Dim d)) <- (Read.readMaybe ('D':s) :: Maybe SomeDim) -> case d of D0 | Dict <- inferKnownBackend @t @'[0] -> SomeDataFrame <$> readPrecFixedDF @t (D0 :* U) @@ -479,7 +607,7 @@ return (x :*: xs) readBoundedMultiDF :: forall (ts :: [Type]) (ds :: [XNat]) - . (All Read ts, All PrimBytes ts, KnownXNatTypes ds) + . (All Read ts, All PrimBytes ts, All KnownDimType ds) => TypeList ts -> Dims ds -> Read.ReadPrec (DataFrame ts ds) @@ -562,9 +690,9 @@ -- | A scalar DataFrame is just a newtype wrapper on a value. pattern S :: forall (t :: Type) . t -> DataFrame t ('[] :: [Nat]) -- rely on that Scalar is just two times newtype alias to t -pattern S x <- (unsafeCoerce# -> x) +pattern S x <- (unScalar -> x) where - S = unsafeCoerce# + S = scalar {-# COMPLETE S #-} @@ -800,25 +928,27 @@ -> Int -> r) -> DataFrame t (d ': ds) -> r -unpackDF' k df +unpackDF' k | d :* Dims <- dims @(d ': ds) , Dict <- inferKnownBackend @t @(d ': ds) , Dict <- inferKnownBackend @t @ds - = case arrayContent# df of - (# x | #) - | e <- broadcast x - -> let f :: forall (zr :: RuntimeRep) (z :: TYPE zr) - . (Int -> DataFrame t ds -> z) -> Int -> z - f consume = (`consume` e) - in k Dict f 0 - (# | (# cdims, off, arr #) #) - | cd <- CumulDims . tail $ unCumulDims cdims - , td <- cdTotalDim# cd - , n <- case dimVal d of W# w -> word2Int# w - -> let f :: forall (zr :: RuntimeRep) (z :: TYPE zr) - . (Int -> DataFrame t ds -> z) -> Int -> z - f consume (I# o) = consume (I# (o -# td)) (fromElems cd o arr) - in k Dict f (I# (off +# td *# (n -# 1#))) + = withArrayContent + ( \x -> + let e = broadcast x + f :: forall (zr :: RuntimeRep) (z :: TYPE zr) + . (Int -> DataFrame t ds -> z) -> Int -> z + f consume = (`consume` e) + in k Dict f 0 + ) + ( \cdims off arr -> + let cd = CumulDims . tail $ unCumulDims cdims + td = cdTotalDim# cd + n = case dimVal d of W# w -> word2Int# w + f :: forall (zr :: RuntimeRep) (z :: TYPE zr) + . (Int -> DataFrame t ds -> z) -> Int -> z + f consume (I# o) = consume (I# (o -# td)) (fromElems cd o arr) + in k Dict f (I# (off +# td *# (n -# 1#))) + ) | otherwise = error "Numeric.DataFrame.Type.unpackDF: impossible args" @@ -878,17 +1008,20 @@ -- | Unsafely copy two PrimBytes values into a third one. -unsafeAppendPB :: (PrimBytes x, PrimBytes y, PrimBytes z) +unsafeAppendPB :: forall x y z + . (PrimBytes x, PrimBytes y, PrimBytes z) => x -> y -> z -unsafeAppendPB x y - | sx <- byteSize x - = case runRW# - ( \s0 -> case newByteArray# (sx +# byteSize y) s0 of - (# s1, mba #) -> unsafeFreezeByteArray# mba - ( writeBytes mba sx y - ( writeBytes mba 0# x s1)) - ) of (# _, r #) -> fromBytes 0# r - +unsafeAppendPB x y = go (byteSize x) (byteSize y) + where + go :: Int# -> Int# -> z + go 0# _ = fromBytes (byteOffset y) (getBytes y) + go _ 0# = fromBytes (byteOffset x) (getBytes x) + go sx sy = case runRW# + ( \s0 -> case newByteArray# (sx +# sy) s0 of + (# s1, mba #) -> unsafeFreezeByteArray# mba + ( writeBytes mba sx y + ( writeBytes mba 0# x s1)) + ) of (# _, r #) -> fromBytes 0# r -- | Construct a DataFrame from a list of smaller DataFrames. -- @@ -911,35 +1044,76 @@ -- Pattern-match against the resulting @XFrame@ to find out its dimensionality. -- -- You must not provide an infinite list as an argument. -fromList :: forall (t :: Type) (ds :: [Nat]) - . (PrimBytes t, Dimensions ds) - => [DataFrame t ds] -> DataFrame t (XN 0 ': AsXDims ds) +fromList :: forall (t :: Type) (ds :: [Nat]) (xds :: [XNat]) + . (PrimBytes t, Dimensions ds, xds ~ Map N ds, ds ~ UnMap N xds) + => [DataFrame t ds] -> DataFrame t (XN 0 ': xds) fromList xs | Dx (D :: Dim n) <- someDimVal . fromIntegral $ length xs , Dict <- inferKnownBackend @t @(n ': ds) - , ds@(AsXDims (XDims ds')) <- dims @ds - , Just Dict <- sameDims ds ds' + , Dict <- unsafeEqTypes @_ @ds @(DimsBound xds) + , Dict <- inferExactFixedDims (dims @ds) = XFrame (fromListWithDefault @t @n @ds undefined xs) | otherwise - = error "Numeri.DataFrame.Type/fromList: impossible arguments" - + = case Dict @(ds ~ UnMap N xds) of + Dict -> -- just make GHC not complain about unused constraints. + error "Numeri.DataFrame.Type/fromList: impossible arguments" -- | Try to convert between @XNat@-indexed DataFrames. -- -- This is useful for imposing restrictions on unknown DataFrames, -- e.g. increasing the minimum number of elements. -constrainDF :: forall (ds :: [XNat]) (ys :: [XNat]) (l :: Type) (ts :: l) - . (BoundedDims ds, All KnownXNatType ds) - => DataFrame ts ys -> Maybe (DataFrame ts ds) -constrainDF (XFrame (df :: DataFrame ts ns)) +constrainDF :: forall (ds :: [XNat]) (ys :: [XNat]) t + . (BoundedDims ds, All KnownDimType ds) + => DataFrame t ys -> Maybe (DataFrame t ds) +constrainDF (XFrame (df :: DataFrame t ns)) | ns <- dims @ns - = case constrainDims @XNat @ds ns of + = case constrainDims @ds ns of Just (XDims (Dims :: Dims ms)) | Dict <- unsafeEqTypes @[Nat] @ns @ms -> Just $ XFrame df _ -> Nothing +-- | Place elements of a vector on the main diagonal of a matrix; +-- fill the rest of the matrix with zeroes. +-- +-- Note, this function is naturally generalized onto higher dimensions +-- (which can be seen from the type signature). +-- +-- Note, the argument of this function does not fully determine its result type. +-- This may cause some obscure type errors. +-- Specify type parameters @n@ and @m@ explicitly or make sure the result type is fixed. +asDiag :: forall (n :: Nat) (m :: Nat) (ds :: [Nat]) (t :: Type) + . ( Dimensions ds, KnownDim n, KnownDim m + , KnownBackend t (Min n m :+ ds) + , KnownBackend t (n :+ m :+ ds) + , PrimBytes t) + => DataFrame t (Min n m :+ ds) + -> DataFrame t (n :+ m :+ ds) +asDiag x + | elemSize <- byteSize @t undefined + , dMinNM@D <- minDim (D @n) (D @m) + , xba <- getBytes x + , ds <- dims @ds + , steps@(CumulDims (elemsNR:_:elemsNE:_)) <- cumulDims (D @n :* D @m :* ds) + , m <- dimVal' @m + , bsE <- case elemsNE of W# w -> word2Int# w *# elemSize + , bsR <- case elemsNR of W# w -> word2Int# w *# elemSize + , bsShift <- case m + 1 of W# w -> word2Int# w *# bsE + , bsLim <- case m * dimVal dMinNM of W# w -> word2Int# w *# bsE + = case runRW# + ( \s0 -> case newByteArray# bsR s0 of + (# s1, mba #) -> + let go offSrc offDst s + | isTrue# (offDst >=# bsLim) = s + | otherwise = go (offSrc +# bsE) (offDst +# bsShift) + (copyByteArray# xba offSrc mba offDst bsE s) + in unsafeFreezeByteArray# mba + (go (byteOffset x) 0# (setByteArray# mba 0# bsR 0# s1)) + ) of (# _, r #) -> fromElems steps 0# r + | otherwise + = error "Numeri.DataFrame.Type/asDiag: impossible arguments" +{-# INLINE asDiag #-} -- Need this for @packDF'@ to make @Int -> c z@ a proper second-order type -- parameterized by the result type. @@ -955,7 +1129,7 @@ U | S x <- v -> z S `k` x D :* (Dims :: Dims ns) - -> case inferTypeableCons @_ @ds of + -> case inferTypeableCons @ds of Dict -> -- PLZ don't ask me how does this work unpackDF' (\_ f -> runOff $ packDF' @@ -965,7 +1139,7 @@ gunfold k z _ = case typeableDims @ds of U -> k (z S) D :* (Dims :: Dims ns) - -> case inferTypeableCons @_ @ds of Dict -> packDF' k z + -> case inferTypeableCons @ds of Dict -> packDF' k z toConstr _ = case typeableDims @ds of U -> scalarFrameConstr d :* _ -> singleFrameConstr $ dimVal d @@ -980,11 +1154,11 @@ instance (AllFrames Data ts ds, Typeable ts, Typeable ds) => Data (DataFrame (ts :: [Type]) (ds :: [Nat])) where gfoldl _ z Z = z Z - gfoldl k z (x :*: xs) = case inferTypeableCons @Type @ts of + gfoldl k z (x :*: xs) = case inferTypeableCons @ts of Dict -> z (:*:) `k` x `k` xs - gunfold k z _ = case typeables @Type @ts of + gunfold k z _ = case typeables @ts of U -> z Z - _ :* _ -> case inferTypeableCons @_ @ts of Dict -> k (k (z (:*:))) + _ :* _ -> case inferTypeableCons @ts of Dict -> k (k (z (:*:))) toConstr Z = multiFrameZConstr toConstr (_ :*: _) = multiFrameConsConstr dataTypeOf _ = dataFrameDataType [multiFrameZConstr, multiFrameConsConstr] @@ -1043,13 +1217,14 @@ fromSingleFrame (dd@D :* (Dims :: Dims ds')) x | Dict <- inferKnownBackend @t @ds , Dict <- inferKnownBackend @t @ds' - = G.M1 $ case arrayContent# x of - (# e | #) -> fillRep @_ @ds' (const $ broadcast e) 0 dd - (# | (# cdims, off, arr #) #) - | cd <- CumulDims . tail $ unCumulDims cdims - , td <- cdTotalDim# cd - -> fillRep @_ @ds' - (\(W# i) -> fromElems cd (off +# td *# word2Int# i) arr) 0 dd + = G.M1 $ withArrayContent + (\e -> fillRep @_ @ds' (const $ broadcast e) 0 dd) + (\cdims off arr -> + let cd = CumulDims . tail $ unCumulDims cdims + td = cdTotalDim# cd + in fillRep @_ @ds' + (\(W# i) -> fromElems cd (off +# td *# word2Int# i) arr) 0 dd + ) x where fillRep :: forall (n :: Nat) (ns :: [Nat]) . (Word -> DataFrame t ns) @@ -1105,7 +1280,7 @@ where d2 = divDim d D2 d2' = d2 `plusDim` modDim d D2 - +{-# INLINE toSingleFrame #-} type family MultiFrameRepNil (ts :: [Type]) :: (Type -> Type) where MultiFrameRepNil '[] = G.C1 ('G.MetaCons "Z" 'G.PrefixI 'False) G.U1 @@ -1136,5 +1311,245 @@ +instance (AllFrames Eq ts ds, AllFrames Ord ts ds) + => Ord (DataFrame (ts :: [Type]) ds) where + compare Z Z = EQ + compare (a :*: as) (b :*: bs) = compare a b <> compare as bs + + +withFixedDF1 :: forall (l :: Type) (ts :: l) (xns :: [XNat]) + (rep :: RuntimeRep) (r :: TYPE rep) + . xns ~ Map 'N (DimsBound xns) + => ( forall (ns :: [Nat]) + . ( All KnownDimType xns, FixedDims xns ns + , Dimensions ns + , KnownBackends ts ns + , ns ~ DimsBound xns + , xns ~ Map 'N ns + ) => DataFrame ts ns -> r + ) -> DataFrame ts xns -> r +withFixedDF1 f (XFrame (df :: DataFrame ts ds)) + | Dict <- unsafeEqTypes @_ @ds @(DimsBound xns) + = f df +{-# INLINE withFixedDF1 #-} + +withFixedDF2 :: forall (l :: Type) (ts :: l) (xns :: [XNat]) + (rep :: RuntimeRep) (r :: TYPE rep) + . xns ~ Map 'N (DimsBound xns) + => ( forall (ns :: [Nat]) + . ( All KnownDimType xns, FixedDims xns ns + , Dimensions ns + , KnownBackends ts ns + , ns ~ DimsBound xns + , xns ~ Map 'N ns + ) => DataFrame ts ns -> DataFrame ts ns -> r + ) -> DataFrame ts xns -> DataFrame ts xns -> r +withFixedDF2 f (XFrame (a :: DataFrame ts as)) (XFrame (b :: DataFrame ts bs)) + | Dict <- unsafeEqTypes @_ @as @(DimsBound xns) + , Dict <- unsafeEqTypes @_ @bs @(DimsBound xns) + = f a b +{-# INLINE withFixedDF2 #-} + +instance ( xns ~ Map 'N (DimsBound xns) + , Eq (DataFrame ts xns) + , Ord (DataFrame ts (DimsBound xns)) + ) => Ord (DataFrame (ts :: l) (xns :: [XNat])) where + compare = withFixedDF2 compare + (>=) = withFixedDF2 (>=) + (<=) = withFixedDF2 (<=) + (>) = withFixedDF2 (>) + (<) = withFixedDF2 (<) + min = withFixedDF2 ((.) XFrame . min) + max = withFixedDF2 ((.) XFrame . max) + +instance ( Dimensions xns + , KnownBackends ts (DimsBound xns) + , Num (DataFrame ts (DimsBound xns)) + ) => Num (DataFrame (ts :: l) (xns :: [XNat])) where + (+) = withKnownXDims @xns $ withFixedDF2 ((.) XFrame . (+)) + (-) = withKnownXDims @xns $ withFixedDF2 ((.) XFrame . (-)) + (*) = withKnownXDims @xns $ withFixedDF2 ((.) XFrame . (*)) + negate = withKnownXDims @xns $ withFixedDF1 (XFrame . negate) + abs = withKnownXDims @xns $ withFixedDF1 (XFrame . abs) + signum = withKnownXDims @xns $ withFixedDF1 (XFrame . signum) + fromInteger = withKnownXDims @xns + (XFrame . fromInteger @(DataFrame ts (DimsBound xns))) + +instance ( Dimensions xns + , KnownBackends ts (DimsBound xns) + , Fractional (DataFrame ts (DimsBound xns)) + ) => Fractional (DataFrame (ts :: l) (xns :: [XNat])) where + (/) = withKnownXDims @xns $ withFixedDF2 ((.) XFrame . (/)) + recip = withKnownXDims @xns $ withFixedDF1 (XFrame . recip) + fromRational = withKnownXDims @xns + (XFrame . fromRational @(DataFrame ts (DimsBound xns))) + +instance ( Dimensions xns + , KnownBackends ts (DimsBound xns) + , Floating (DataFrame ts (DimsBound xns)) + ) => Floating (DataFrame (ts :: l) (xns :: [XNat])) where + pi = withKnownXDims @xns $ XFrame (pi :: DataFrame ts (DimsBound xns)) + exp = withKnownXDims @xns $ withFixedDF1 (XFrame . exp) + log = withKnownXDims @xns $ withFixedDF1 (XFrame . log) + sqrt = withKnownXDims @xns $ withFixedDF1 (XFrame . sqrt) + sin = withKnownXDims @xns $ withFixedDF1 (XFrame . sin) + cos = withKnownXDims @xns $ withFixedDF1 (XFrame . cos) + tan = withKnownXDims @xns $ withFixedDF1 (XFrame . tan) + asin = withKnownXDims @xns $ withFixedDF1 (XFrame . asin) + acos = withKnownXDims @xns $ withFixedDF1 (XFrame . acos) + atan = withKnownXDims @xns $ withFixedDF1 (XFrame . atan) + sinh = withKnownXDims @xns $ withFixedDF1 (XFrame . sinh) + cosh = withKnownXDims @xns $ withFixedDF1 (XFrame . cosh) + tanh = withKnownXDims @xns $ withFixedDF1 (XFrame . tanh) + asinh = withKnownXDims @xns $ withFixedDF1 (XFrame . asinh) + acosh = withKnownXDims @xns $ withFixedDF1 (XFrame . acosh) + atanh = withKnownXDims @xns $ withFixedDF1 (XFrame . atanh) + (**) = withKnownXDims @xns $ withFixedDF2 ((.) XFrame . (**)) + logBase = withKnownXDims @xns $ withFixedDF2 ((.) XFrame . logBase) + +instance ( xns ~ Map 'N (DimsBound xns) + , ProductOrder (DataFrame ts (DimsBound xns)) + ) => ProductOrder (DataFrame (ts :: l) (xns :: [XNat])) where + cmp = withFixedDF2 cmp + +instance ( Dimensions xns + , KnownBackends ts (DimsBound xns) + , Bounded (DataFrame ts (DimsBound xns)) + ) => Bounded (DataFrame (ts :: l) (xns :: [XNat])) where + minBound = withKnownXDims @xns (XFrame (minBound :: DataFrame ts (DimsBound xns))) + maxBound = withKnownXDims @xns (XFrame (maxBound :: DataFrame ts (DimsBound xns))) + +instance ( Dimensions xns + , KnownBackends ts (DimsBound xns) + , PrimBytes (DataFrame ts (DimsBound xns)) + ) => PrimBytes (DataFrame (ts :: l) (xns :: [XNat])) where + type PrimFields (DataFrame ts xns) + = PrimFields (DataFrame ts (DimsBound xns)) + getBytes = withKnownXDims @xns $ withFixedDF1 getBytes + getBytesPinned = withKnownXDims @xns $ withFixedDF1 getBytesPinned + fromBytes i ba + = withKnownXDims @xns (XFrame (fromBytes i ba :: DataFrame ts (DimsBound xns))) + readBytes mba i s0 + | (# s1, (a :: DataFrame ts (DimsBound xns)) #) <- readBytes mba i s0 + = (# s1, withKnownXDims @xns (XFrame a) #) + writeBytes mba i = withKnownXDims @xns $ withFixedDF1 (writeBytes mba i) + readAddr addr s0 + | (# s1, (a :: DataFrame ts (DimsBound xns)) #) <- readAddr addr s0 + = (# s1, withKnownXDims @xns (XFrame a) #) + writeAddr = withKnownXDims @xns $ withFixedDF1 writeAddr + byteSize _ = withKnownXDims @xns (byteSize @(DataFrame ts (DimsBound xns)) undefined) + byteAlign _ = withKnownXDims @xns (byteAlign @(DataFrame ts (DimsBound xns)) undefined) + byteOffset = withKnownXDims @xns (withFixedDF1 byteOffset) + byteFieldOffset n _ + = withKnownXDims @xns (byteFieldOffset @(DataFrame ts (DimsBound xns)) n undefined) + indexArray ba i + = withKnownXDims @xns (XFrame (indexArray ba i :: DataFrame ts (DimsBound xns))) + readArray mba i s0 + | (# s1, (a :: DataFrame ts (DimsBound xns)) #) <- readArray mba i s0 + = (# s1, withKnownXDims @xns (XFrame a) #) + writeArray mba i = withKnownXDims @xns (withFixedDF1 (writeArray mba i)) + +instance ( Dimensions xns + , KnownBackend t (DimsBound xns) + , PrimArray t (DataFrame t (DimsBound xns)) + , PrimBytes t + ) => PrimArray t (DataFrame t (xns :: [XNat])) where + broadcast# + = withKnownXDims @xns (XFrame . broadcast# @t @(DataFrame t (DimsBound xns))) + ix# i = withKnownXDims @xns (withFixedDF1 (ix# i)) + gen# cd f s0 + | (# s1, (a :: DataFrame t (DimsBound xns)) #) <- gen# cd f s0 + = (# s1, withKnownXDims @xns (XFrame a) #) + upd# cd i t = withKnownXDims @xns (withFixedDF1 (XFrame . upd# cd i t)) + withArrayContent# f g = withKnownXDims @xns (withFixedDF1 (withArrayContent f g)) + offsetElems = withKnownXDims @xns (withFixedDF1 offsetElems) + uniqueOrCumulDims = withKnownXDims @xns (withFixedDF1 uniqueOrCumulDims) + fromElems# cd i ba + = withKnownXDims @xns (XFrame (fromElems cd i ba :: DataFrame t (DimsBound xns))) + + +-- The following instance only make sense on scalars. + +instance ( Enum (DataFrame ts ('[] :: [Nat])) + , KnownBackends ts ('[] :: [Nat]) + ) => Enum (DataFrame (ts :: l) ('[] :: [XNat])) where + succ (XFrame x) = XFrame (succ x) + pred (XFrame x) = XFrame (pred x) + toEnum = XFrame . toEnum + fromEnum (XFrame x) = fromEnum x + enumFrom (XFrame x) = map XFrame $ enumFrom x + enumFromThen (XFrame x) (XFrame y) = map XFrame $ enumFromThen x y + enumFromTo (XFrame x) (XFrame y) = map XFrame $ enumFromTo x y + enumFromThenTo (XFrame x) (XFrame y) (XFrame z) + = map XFrame $ enumFromThenTo x y z + +instance ( Epsilon (DataFrame ts ('[] :: [Nat])) + , KnownBackends ts ('[] :: [Nat]) + , Eq (DataFrame ts ('[] :: [XNat])) + ) => Epsilon (DataFrame (ts :: l) ('[] :: [XNat])) where + epsilon = XFrame epsilon + +instance ( Real (DataFrame ts ('[] :: [Nat])) + , KnownBackends ts ('[] :: [Nat]) + , Eq (DataFrame ts ('[] :: [XNat])) + ) => Real (DataFrame (ts :: l) ('[] :: [XNat])) where + toRational (XFrame x) = toRational x + +instance ( Integral (DataFrame ts ('[] :: [Nat])) + , KnownBackends ts ('[] :: [Nat]) + , Eq (DataFrame ts ('[] :: [XNat])) + ) => Integral (DataFrame (ts :: l) ('[] :: [XNat])) where + quot (XFrame x) (XFrame y) = XFrame (quot x y) + rem (XFrame x) (XFrame y) = XFrame (rem x y) + div (XFrame x) (XFrame y) = XFrame (div x y) + mod (XFrame x) (XFrame y) = XFrame (mod x y) + quotRem (XFrame x) (XFrame y) = (XFrame *** XFrame) (quotRem x y) + divMod (XFrame x) (XFrame y) = (XFrame *** XFrame) (divMod x y) + toInteger (XFrame x) = toInteger x + +instance ( RealExtras (DataFrame t ('[] :: [Nat])) + , KnownBackend t ('[] :: [Nat]) + , Eq t + ) => RealExtras (DataFrame (t :: Type) ('[] :: [XNat])) where + copysign (XFrame x) (XFrame y) = XFrame (copysign x y) + +instance ( RealFrac (DataFrame t ('[] :: [Nat])) + , KnownBackend t ('[] :: [Nat]) + , Eq t + ) => RealFrac (DataFrame (t :: Type) ('[] :: [XNat])) where + properFraction (XFrame x) = second XFrame (properFraction x) + truncate (XFrame x) = truncate x + round (XFrame x) = round x + ceiling (XFrame x) = ceiling x + floor (XFrame x) = floor x + +instance ( RealFloat (DataFrame t ('[] :: [Nat])) + , KnownBackend t ('[] :: [Nat]) + , Eq t + ) => RealFloat (DataFrame (t :: Type) ('[] :: [XNat])) where + floatRadix = const (floatRadix @(DataFrame t ('[] :: [Nat])) undefined) + floatDigits = const (floatDigits @(DataFrame t ('[] :: [Nat])) undefined) + floatRange = const (floatRange @(DataFrame t ('[] :: [Nat])) undefined) + decodeFloat (XFrame x) = decodeFloat x + encodeFloat i = XFrame . encodeFloat i + exponent (XFrame x) = exponent x + significand (XFrame x) = XFrame (significand x) + scaleFloat i (XFrame x) = XFrame (scaleFloat i x) + isNaN (XFrame x) = isNaN x + isInfinite (XFrame x) = isInfinite x + isDenormalized (XFrame x) = isDenormalized x + isNegativeZero (XFrame x) = isNegativeZero x + isIEEE (XFrame x) = isIEEE x + atan2 (XFrame y) (XFrame x) = XFrame (atan2 y x) + +instance ( RealFloatExtras (DataFrame t ('[] :: [Nat])) + , KnownBackend t ('[] :: [Nat]) + , Eq t + ) => RealFloatExtras (DataFrame (t :: Type) ('[] :: [XNat])) where + hypot (XFrame x) (XFrame y) = XFrame (hypot x y) + maxFinite = XFrame maxFinite + + + unsafeEqTypes :: forall k (a :: k) (b :: k) . Dict (a ~ b) -unsafeEqTypes = unsafeCoerce# (Dict :: Dict (a ~ a)) +unsafeEqTypes = unsafeCoerce (Dict :: Dict (a ~ a)) diff --git a/src/Numeric/Matrix.hs b/src/Numeric/Matrix.hs --- a/src/Numeric/Matrix.hs +++ b/src/Numeric/Matrix.hs @@ -3,7 +3,6 @@ , SquareMatrix (..) , MatrixDeterminant (..) , MatrixInverse (..) - , MatrixLU (..), LUFact (..) , Matrix , HomTransform4 (..) , Mat22f, Mat23f, Mat24f @@ -14,10 +13,9 @@ , Mat42d, Mat43d, Mat44d , mat22, mat33, mat44 , (%*) - , pivotMat, luSolve ) where import Numeric.DataFrame.Internal.Backend () import Numeric.Matrix.Internal -import Numeric.Matrix.Internal.Mat44d () -import Numeric.Matrix.Internal.Mat44f () +import Numeric.Matrix.Internal.Double () +import Numeric.Matrix.Internal.Float () diff --git a/src/Numeric/Matrix/Bidiagonal.hs b/src/Numeric/Matrix/Bidiagonal.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/Matrix/Bidiagonal.hs @@ -0,0 +1,123 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE PolyKinds #-} +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE StandaloneDeriving #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UnboxedTuples #-} +{-# LANGUAGE UndecidableInstances #-} + +module Numeric.Matrix.Bidiagonal + ( BiDiag (..), biDiag, bidiagonalHouseholder + ) where + +import Control.Monad +import Control.Monad.ST +import Data.Kind +import Numeric.Basics +import Numeric.DataFrame.ST +import Numeric.DataFrame.SubSpace +import Numeric.DataFrame.Type +import Numeric.Dimensions +import Numeric.Matrix.Internal +import Numeric.Scalar.Internal +import Numeric.Subroutine.Householder +import Numeric.Vector.Internal + + +-- | Put two vectors on the main and first upper diagonal. +biDiag :: forall (t :: Type) (n :: Nat) (m :: Nat) + . (PrimBytes t, Num t) + => Dims '[n,m] + -> Vector t (Min n m) + -> Vector t (Min n m) + -> Matrix t n m +biDiag (dn@D :* dm@D :* U) a b = runST $ do + dnm@D <- pure $ minDim dn dm + rPtr <- thawDataFrame 0 + forM_ [0 .. dimVal dnm - 1] $ \i -> do + writeDataFrame rPtr (Idx i :* Idx i :* U) $ a ! i + when (i+1 < dimVal dm) $ + writeDataFrame rPtr (Idx i :* Idx (i+1) :* U) $ b ! i + unsafeFreezeDataFrame rPtr + +-- | Decomposition of a matrix \( A = U B V^\intercal \) such that +-- \( U \) and \( V \) are orthogonal and \( B \) is bidiagonal. +data BiDiag (t :: Type) (n :: Nat) (m :: Nat) + = BiDiag + { bdU :: Matrix t n n + -- ^ \( U \) left orthogonal matrix + , bdUDet :: Scalar t + -- ^ A shortcut for evaluating a determinant of \( |U| = \pm 1 \) + , bdAlpha :: Vector t (Min n m) + -- ^ Main diagonal of \( B \) + , bdBeta :: Vector t (Min n m) + -- ^ First upper diagonal of \( B \); + -- its last element equals zero if \( n \geq m \) + , bdV :: Matrix t m m + -- ^ \( B \) left orthogonal matrix + , bdVDet :: Scalar t + -- ^ A shortcut for evaluating a determinant of \( |V| = \pm 1 \) + } + +deriving instance ( Show t, PrimBytes t + , KnownDim n, KnownDim m, KnownDim (Min n m)) + => Show (BiDiag t n m) +deriving instance ( Eq t, PrimBytes t + , KnownDim n, KnownDim m, KnownDim (Min n m) + , KnownBackend t '[Min n m]) + => Eq (BiDiag t n m) + +{- | +Decompose a matrix \( A = U B V^\intercal \) such that +\( U \) and \( V \) are orthogonal and \( B \) is bidiagonal. + +The first returned number + -} +bidiagonalHouseholder :: + forall (t :: Type) (n :: Nat) (m :: Nat) + . (PrimBytes t, Ord t, Epsilon t, KnownDim n, KnownDim m) + => Matrix t n m + -> BiDiag t n m +bidiagonalHouseholder a = runST $ do + D <- pure $ minDim (dim @n) (dim @m) + tmpNPtr <- newDataFrame + tmpMPtr <- newDataFrame + uPtr <- thawDataFrame eye + bPtr <- thawDataFrame a + vPtr <- thawDataFrame eye + (ud, vd) <- + let f (ud, vd) i = do + ud' <- householderReflectionInplaceL tmpNPtr uPtr bPtr + (Idx (i - 1) :* Idx (i - 1) :* U) + vd' <- householderReflectionInplaceR tmpMPtr vPtr bPtr + (Idx (i - 1) :* Idx i :* U) + return (ud /= ud', vd /= vd') + in foldM f (False, False) [1 .. lim - 1] + + udn <- householderReflectionInplaceL tmpNPtr uPtr bPtr + (Idx (lim - 1) :* Idx (lim - 1) :* U) + vdn <- if (m > lim) + then householderReflectionInplaceR tmpMPtr vPtr bPtr + (Idx (lim - 1) :* Idx lim :* U) + else pure False + bdU <- unsafeFreezeDataFrame uPtr + bdV <- unsafeFreezeDataFrame vPtr + b <- unsafeFreezeDataFrame bPtr + let bdAlpha = iwgen @t @'[Min n m] + (\(Idx i :* U) -> index (Idx i :* Idx i :* U) b) + bdBeta = iwgen @t @'[Min n m] + (\(Idx i :* U) -> if i+1 < m then index (Idx i :* Idx (i+1) :* U) b else 0) + bdUDet = if ud /= udn then -1 else 1 + bdVDet = if vd /= vdn then -1 else 1 + return BiDiag {..} + where + n = dimVal' @n + m = dimVal' @m + lim = max 1 (min n m) diff --git a/src/Numeric/Matrix/Internal.hs b/src/Numeric/Matrix/Internal.hs --- a/src/Numeric/Matrix/Internal.hs +++ b/src/Numeric/Matrix/Internal.hs @@ -17,7 +17,6 @@ , SquareMatrix (..) , MatrixDeterminant (..) , MatrixInverse (..) - , MatrixLU (..), LUFact (..) , Matrix , HomTransform4 (..) , Mat22f, Mat23f, Mat24f @@ -28,21 +27,13 @@ , Mat42d, Mat43d, Mat44d , mat22, mat33, mat44 , (%*) - , pivotMat, luSolve ) where -import Control.Monad (foldM) -import Control.Monad.ST -import Data.Foldable (foldl', forM_) -import Data.List (delete) import GHC.Base import Numeric.DataFrame.Contraction ((%*)) import Numeric.DataFrame.Internal.PrimArray -import Numeric.DataFrame.ST -import Numeric.DataFrame.SubSpace import Numeric.DataFrame.Type import Numeric.Dimensions -import Numeric.PrimBytes import Numeric.Scalar.Internal import Numeric.Vector.Internal @@ -70,29 +61,6 @@ inverse :: Matrix t n n -> Matrix t n n --- | Result of LU factorization with Partial Pivoting --- @ PA = LU @. -data LUFact t n - = LUFact - { luLower :: Matrix t n n - -- ^ Lower triangular matrix @L@. - -- All elements on the diagonal of @L@ equal @1@. - , luUpper :: Matrix t n n - -- ^ Upper triangular matrix @U@ - , luPerm :: Matrix t n n - -- ^ Row permutation matrix @P@ - , luPermSign :: Scalar t - -- ^ Sign of permutation @luPermSign == det . luPerm@ - } - -deriving instance (Show t, PrimBytes t, KnownDim n) => Show (LUFact t n) -deriving instance (Eq (Matrix t n n), Eq t) => Eq (LUFact t n) - -class MatrixLU t (n :: Nat) where - -- | Compute LU factorization with Partial Pivoting - lu :: Matrix t n n -> LUFact t n - - -- | Operations on 4x4 transformation matrices and vectors in homogeneous coordinates. -- All angles are specified in radians. -- @@ -260,214 +228,3 @@ = let f 0# = ix# 0# df f k = ix# k df + f (k -# n1) in scalar $ f (n *# n -# 1#) - - -instance ( KnownDim n, Ord t, Fractional t - , PrimBytes t - , PrimArray t (Matrix t n n) - , PrimArray t (Vector t n) - , PrimBytes (Vector t n) - , PrimBytes (Matrix t n n) - ) - => MatrixInverse t n where - inverse m = runST $ do - m' <- newDataFrame - indexWise_ @t @'[n] @'[n] @'[n, n] @_ @() - (\is -> - indexWise_ @t @'[n] @'[] @'[n] @_ @() - (\(j:*U) -> writeDataFrame m' (j:*is)) . luSolve luM - ) eye - unsafeFreezeDataFrame m' - where - luM = lu m - - -- much nicer, but probably slower version: - -- (it is not particularly fast to do an extra transpose - -- before calculating the inverse) - -- ewmap (luSolve . lu $ transpose m) eye - - - -instance ( KnownDim n, Ord t, Fractional t - , PrimBytes t, PrimArray t (Matrix t n n)) - => MatrixDeterminant t n where - det m = prodF (luUpper f) * prodF (luLower f) * luPermSign f - where - f = lu m - n = case dimVal' @n of W# w -> word2Int# w - n1 = n +# 1# - nn1 = n *# n -# 1# - prodF a = scalar $ prodF' nn1 a - prodF' 0# a = ix# 0# a - prodF' k a = ix# k a * prodF' (k -# n1) a - - -instance ( KnownDim n, Ord t, Fractional t - , PrimBytes t, PrimArray t (Matrix t n n)) - => MatrixLU t n where - lu m' = case runRW# go of - (# _, (# bu, bl #) #) -> LUFact - { luLower = fromElems steps 0# bl - , luUpper = fromElems steps 0# bu - , luPerm = p - , luPermSign = si - } - where - dn = dimVal' @n - dnn = dn * dn - steps = CumulDims [dnn, dn, 1] - (m, p, si) = pivotMat m' - n = case dn of W# w -> word2Int# w - nn = case dnn of W# w -> word2Int# w - tbs = byteSize @t undefined - bsize = nn *# tbs - ixm i j = ix# (i *# n +# j) m - loop :: (Int# -> a -> State# s -> (# State# s, a #)) - -> Int# -> Int# -> a -> State# s -> (# State# s, a #) - loop f i k x s - | isTrue# (i ==# k) = (# s, x #) - | otherwise = case f i x s of - (# s', y #) -> loop f ( i +# 1# ) k y s' - go s0 - | (# s1, mbl #) <- newByteArray# bsize s0 - , (# s2, mbu #) <- newByteArray# bsize s1 - , s3 <- setByteArray# mbl 0# bsize 0# s2 - , s4 <- setByteArray# mbu 0# bsize 0# s3 - , readL <- \i j -> readArray @t mbl (i *# n +# j) - , readU <- \i j -> readArray @t mbu (i *# n +# j) - , writeL <- \i j -> writeArray @t mbl (i *# n +# j) - , writeU <- \i j -> writeArray @t mbu (i *# n +# j) - , computeU <- \i j -> - let f k x s - | (# s' , ukj #) <- readU k j s - , (# s'', lik #) <- readL i k s' - = (# s'', x - ukj * lik #) - in loop f 0# i (ixm i j) - , computeL' <- \i j -> - let f k x s - | (# s' , ukj #) <- readU k j s - , (# s'', lik #) <- readL i k s' - = (# s'', x - ukj * lik #) - in loop f 0# j (ixm i j) - , (# sr, () #) <- - loop - ( \j _ sj -> case sj of - sj0 - | sj1 <- writeL j j 1 sj0 - , (# sj2, () #) <- loop - ( \i _ sij0 -> case computeU i j sij0 of - (# sij1, uij #) -> (# writeU i j uij sij1, () #) - ) 0# j () sj1 - , (# sj3, ujj #) <- computeU j j sj2 - , sj4 <- writeU j j ujj sj3 - -> case ujj of - 0 -> loop - ( \i _ sij -> (# writeL i j 0 sij, () #) - ) (j +# 1#) n () sj4 - x -> loop - ( \i _ sij0 -> case computeL' i j sij0 of - (# sij1, lij #) -> (# writeL i j (lij / x) sij1, () #) - ) (j +# 1#) n () sj4 - ) 0# n () s4 - , (# sf0, bl #) <- unsafeFreezeByteArray# mbl sr - , (# sf1, bu #) <- unsafeFreezeByteArray# mbu sf0 - = (# sf1, (# bu, bl #) #) - - --- | Solve @Ax = b@ problem given LU decomposition of A. -luSolve :: forall (t :: Type) (n :: Nat) - . ( KnownDim n, Fractional t - , PrimArray t (Matrix t n n), PrimArray t (Vector t n)) - => LUFact t n -> Vector t n -> Vector t n -luSolve LUFact {..} b = x - where - -- Pb = LUx - pb = luPerm %* b - n# = case dimVal' @n of W# w -> word2Int# w - n = I# n# - -- Ly = Pb - y :: Vector t n - y = runST $ do - my <- newDataFrame - let ixA (I# i) (I# j) = scalar $ ix# (i *# n# +# j) luLower - ixB (I# i) = scalar $ ix# i pb - forM_ [0..n-1] $ \i -> do - v <- foldM ( \v j -> do - dj <- readDataFrameOff my j - return $ v - dj * ixA i j - ) (ixB i) [0..i-1] - writeDataFrameOff my i v - unsafeFreezeDataFrame my - -- Ux = y - x = runST $ do - mx <- newDataFrame - let ixA (I# i) (I# j) = scalar $ ix# (i *# n# +# j) luUpper - ixB (I# i) = scalar $ ix# i y - forM_ [n-1, n-2 .. 0] $ \i -> do - v <- foldM ( \v j -> do - dj <- readDataFrameOff mx j - return $ v - dj * ixA i j - ) (ixB i) [i+1..n-1] - writeDataFrameOff mx i (v / ixA i i) - unsafeFreezeDataFrame mx - - --- | Permute rows that the largest magnitude elements in columns are on diagonals. --- --- Invariants of result matrix: --- * forall j >= i: |M[i,i]| >= M[j,i] --- * if M[i,i] == 0 then forall j >= i: |M[i+1,i+1]| >= M[j,i+1] -pivotMat :: forall (t :: Type) (n :: Nat) - . (KnownDim n, PrimArray t (Matrix t n n), Ord t, Num t) - => Matrix t n n -> ( Matrix t n n -- permutated matrix - , Matrix t n n -- permutation matrix - , Scalar t -- sign of permutation matrix - ) -pivotMat mat - = ( let f ( _, [] ) = undefined - f ( j, i:is ) - | j == n = f (0, is) - | otherwise = (# (j+1, i:is), ix i j #) - in case gen# steps f (0,rowOrder) of - (# _, r #) -> r - , let f ( _, [] ) = undefined - f ( j, x:xs ) - | j == n = f (0, xs) - | otherwise = (# ( j + 1, x:xs), if j == x then 1 else 0 #) - in case gen# steps f (0,rowOrder) of - (# _, r #) -> r - , if countMisordered rowOrder `rem` 2 == 1 - then -1 else 1 - ) - where - dn = dimVal' @n - steps = CumulDims [dn * dn, dn, 1] - -- permuted row ordering - rowOrder = uncurry fillPass $ searchPass 0 [0..n-1] - -- matrix size - n = fromIntegral dn - n# = case n of I# x -> x - -- sign of permutations - countMisordered :: [Int] -> Int - countMisordered [] = 0 - countMisordered (i:is) = foldl' (\c j -> if i > j then succ c else c) 0 is - + countMisordered is - ix (I# i) (I# j) = ix# (i *# n# +# j) mat - findMax :: Int -> [Int] -> (t, Int) - findMax j = foldl' (\(ox, oi) i -> let x = abs (ix i j) - in if x > ox then (x, i) - else (ox, oi) - ) (0, 0) - -- search maximums, leaving Nothing where all rows are 0 - searchPass :: Int -> [Int] -> ([Int], [Maybe Int]) - searchPass j is - | j == n = (is, []) - | otherwise = case findMax j is of - (0, _) -> (Nothing:) <$> searchPass (j+1) is - (_, i) -> (Just i:) <$> searchPass (j+1) (delete i is) - -- replace Nothings with remaining row numbers - fillPass :: [Int] -> [Maybe Int] -> [Int] - fillPass _ [] = [] - fillPass js (Just i : is) = i : fillPass js is - fillPass (j:js) (Nothing : is) = j : fillPass js is - fillPass [] (Nothing : is) = 0 : fillPass [] is diff --git a/src/Numeric/Matrix/Internal/Double.hs b/src/Numeric/Matrix/Internal/Double.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/Matrix/Internal/Double.hs @@ -0,0 +1,199 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE UnboxedTuples #-} +{-# LANGUAGE ViewPatterns #-} +{-# OPTIONS_GHC -fno-warn-orphans #-} + +module Numeric.Matrix.Internal.Double () where + +import qualified Control.Monad.ST as ST +import qualified Numeric.DataFrame.ST as ST +import Numeric.Dimensions (KnownDim) +import Numeric.Matrix.Internal +import Numeric.Matrix.LU +import Numeric.Scalar.Internal +import Numeric.Vector.Internal + +{-# INLINE mkMat #-} +mkMat :: + Double -> Double -> Double -> Double -> + Double -> Double -> Double -> Double -> + Double -> Double -> Double -> Double -> + Double -> Double -> Double -> Double -> + Mat44d +mkMat + _11 _12 _13 _14 + _21 _22 _23 _24 + _31 _32 _33 _34 + _41 _42 _43 _44 + = ST.runST $ do + df <- ST.newDataFrame + ST.writeDataFrameOff df 0 $ scalar _11 + ST.writeDataFrameOff df 1 $ scalar _12 + ST.writeDataFrameOff df 2 $ scalar _13 + ST.writeDataFrameOff df 3 $ scalar _14 + ST.writeDataFrameOff df 4 $ scalar _21 + ST.writeDataFrameOff df 5 $ scalar _22 + ST.writeDataFrameOff df 6 $ scalar _23 + ST.writeDataFrameOff df 7 $ scalar _24 + ST.writeDataFrameOff df 8 $ scalar _31 + ST.writeDataFrameOff df 9 $ scalar _32 + ST.writeDataFrameOff df 10 $ scalar _33 + ST.writeDataFrameOff df 11 $ scalar _34 + ST.writeDataFrameOff df 12 $ scalar _41 + ST.writeDataFrameOff df 13 $ scalar _42 + ST.writeDataFrameOff df 14 $ scalar _43 + ST.writeDataFrameOff df 15 $ scalar _44 + ST.unsafeFreezeDataFrame df + +instance HomTransform4 Double where + {-# INLINE translate4 #-} + translate4 (unpackV4# -> (# x, y, z, _ #)) = mkMat + 1 0 0 0 + 0 1 0 0 + 0 0 1 0 + x y z 1 + + {-# INLINE translate3 #-} + translate3 (unpackV3# -> (# x, y, z #)) = mkMat + 1 0 0 0 + 0 1 0 0 + 0 0 1 0 + x y z 1 + + {-# INLINE rotateX #-} + rotateX a = mkMat + 1 0 0 0 + 0 c s 0 + 0 n c 0 + 0 0 0 1 + where + c = cos a + s = sin a + n = -s + + {-# INLINE rotateY #-} + rotateY a = mkMat + c 0 n 0 + 0 1 0 0 + s 0 c 0 + 0 0 0 1 + where + c = cos a + s = sin a + n = -s + + {-# INLINE rotateZ #-} + rotateZ a = mkMat + c s 0 0 + n c 0 0 + 0 0 1 0 + 0 0 0 1 + where + c = cos a + s = sin a + n = -s + + {-# INLINE rotate #-} + rotate (unpackV3# -> (# x, y, z #)) a = mkMat + (c+xxv) (yxv+zs) (zxv-ys) 0 + (xyv-zs) (c+yyv) (zyv+xs) 0 + (xzv+ys) (yzv-xs) (c+zzv) 0 + 0 0 0 1 + where + c = cos a + v = 1 - c -- v for versine + s = sin a + xxv = x * x * v + xyv = x * y * v + xzv = x * z * v + yxv = xyv + yyv = y * y * v + yzv = y * z * v + zxv = xzv + zyv = yzv + zzv = z * z * v + xs = x * s + ys = y * s + zs = z * s + + {-# INLINE rotateEuler #-} + rotateEuler x y z = mkMat + (cy*cz) (cx*sz+sx*sy*cz) (sx*sz-cx*sy*cz) 0 + (-cy*sz) (cx*cz-sx*sy*sz) (sx*cz+cx*sy*sz) 0 + sy (-sx*cy) (cx*cy) 0 + 0 0 0 1 + where + cx = cos x + sx = sin x + cy = cos y + sy = sin y + cz = cos z + sz = sin z + + {-# INLINE lookAt #-} + lookAt up cam foc = mkMat + xb1 yb1 zb1 0 + xb2 yb2 zb2 0 + xb3 yb3 zb3 0 + tx ty tz 1 + where + (# xb1, xb2, xb3 #) = unpackV3# xb + (# yb1, yb2, yb3 #) = unpackV3# yb + (# zb1, zb2, zb3 #) = unpackV3# zb + zb = normalized $ cam - foc -- Basis vector for "backward", since +Z is behind the camera + xb = normalized $ up `cross` zb -- Basis vector for "right" + yb = zb `cross` xb -- Basis vector for "up" + ncam = -cam + tx = unScalar $ xb `dot` ncam + ty = unScalar $ yb `dot` ncam + tz = unScalar $ zb `dot` ncam + + {-# INLINE perspective #-} + perspective n f fovy aspect = mkMat + dpw 0 0 0 + 0 dph 0 0 + 0 0 a (-1) + 0 0 b 0 + where + hpd = tan (fovy * 0.5) -- height/distance + wpd = aspect * hpd; -- width/distance + dph = recip hpd -- distance/height + dpw = recip wpd -- distance/width + nmf = n - f + a = (n + f) / nmf + b = 2 * n * f / nmf + + {-# INLINE orthogonal #-} + orthogonal n f w h = mkMat + iw 0 0 0 + 0 ih 0 0 + 0 0 a 0 + 0 0 b 1 + where + ih = 2 / h + iw = 2 / w + nmf = n - f + a = 2 / nmf + b = (n + f) / nmf + + {-# INLINE toHomPoint #-} + toHomPoint (unpackV3# -> (# x, y, z #)) = vec4 x y z 1 + + {-# INLINE toHomVector #-} + toHomVector (unpackV3# -> (# x, y, z #)) = vec4 x y z 0 + + {-# INLINE fromHom #-} + fromHom (unpackV4# -> (# x, y, z, w #)) + | w == 0 = vec3 x y z + | otherwise = vec3 (x/w) (y/w) (z/w) + +instance KnownDim n => MatrixInverse Double n where + inverse = inverseViaLU + +instance KnownDim n => MatrixDeterminant Double n where + det = detViaLU diff --git a/src/Numeric/Matrix/Internal/Float.hs b/src/Numeric/Matrix/Internal/Float.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/Matrix/Internal/Float.hs @@ -0,0 +1,199 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE UnboxedTuples #-} +{-# LANGUAGE ViewPatterns #-} +{-# OPTIONS_GHC -fno-warn-orphans #-} + +module Numeric.Matrix.Internal.Float () where + +import qualified Control.Monad.ST as ST +import qualified Numeric.DataFrame.ST as ST +import Numeric.Dimensions (KnownDim) +import Numeric.Matrix.Internal +import Numeric.Matrix.LU +import Numeric.Scalar.Internal +import Numeric.Vector.Internal + +{-# INLINE mkMat #-} +mkMat :: + Float -> Float -> Float -> Float -> + Float -> Float -> Float -> Float -> + Float -> Float -> Float -> Float -> + Float -> Float -> Float -> Float -> + Mat44f +mkMat + _11 _12 _13 _14 + _21 _22 _23 _24 + _31 _32 _33 _34 + _41 _42 _43 _44 + = ST.runST $ do + df <- ST.newDataFrame + ST.writeDataFrameOff df 0 $ scalar _11 + ST.writeDataFrameOff df 1 $ scalar _12 + ST.writeDataFrameOff df 2 $ scalar _13 + ST.writeDataFrameOff df 3 $ scalar _14 + ST.writeDataFrameOff df 4 $ scalar _21 + ST.writeDataFrameOff df 5 $ scalar _22 + ST.writeDataFrameOff df 6 $ scalar _23 + ST.writeDataFrameOff df 7 $ scalar _24 + ST.writeDataFrameOff df 8 $ scalar _31 + ST.writeDataFrameOff df 9 $ scalar _32 + ST.writeDataFrameOff df 10 $ scalar _33 + ST.writeDataFrameOff df 11 $ scalar _34 + ST.writeDataFrameOff df 12 $ scalar _41 + ST.writeDataFrameOff df 13 $ scalar _42 + ST.writeDataFrameOff df 14 $ scalar _43 + ST.writeDataFrameOff df 15 $ scalar _44 + ST.unsafeFreezeDataFrame df + +instance HomTransform4 Float where + {-# INLINE translate4 #-} + translate4 (unpackV4# -> (# x, y, z, _ #)) = mkMat + 1 0 0 0 + 0 1 0 0 + 0 0 1 0 + x y z 1 + + {-# INLINE translate3 #-} + translate3 (unpackV3# -> (# x, y, z #)) = mkMat + 1 0 0 0 + 0 1 0 0 + 0 0 1 0 + x y z 1 + + {-# INLINE rotateX #-} + rotateX a = mkMat + 1 0 0 0 + 0 c s 0 + 0 n c 0 + 0 0 0 1 + where + c = cos a + s = sin a + n = -s + + {-# INLINE rotateY #-} + rotateY a = mkMat + c 0 n 0 + 0 1 0 0 + s 0 c 0 + 0 0 0 1 + where + c = cos a + s = sin a + n = -s + + {-# INLINE rotateZ #-} + rotateZ a = mkMat + c s 0 0 + n c 0 0 + 0 0 1 0 + 0 0 0 1 + where + c = cos a + s = sin a + n = -s + + {-# INLINE rotate #-} + rotate (unpackV3# -> (# x, y, z #)) a = mkMat + (c+xxv) (yxv+zs) (zxv-ys) 0 + (xyv-zs) (c+yyv) (zyv+xs) 0 + (xzv+ys) (yzv-xs) (c+zzv) 0 + 0 0 0 1 + where + c = cos a + v = 1 - c -- v for versine + s = sin a + xxv = x * x * v + xyv = x * y * v + xzv = x * z * v + yxv = xyv + yyv = y * y * v + yzv = y * z * v + zxv = xzv + zyv = yzv + zzv = z * z * v + xs = x * s + ys = y * s + zs = z * s + + {-# INLINE rotateEuler #-} + rotateEuler x y z = mkMat + (cy*cz) (cx*sz+sx*sy*cz) (sx*sz-cx*sy*cz) 0 + (-cy*sz) (cx*cz-sx*sy*sz) (sx*cz+cx*sy*sz) 0 + sy (-sx*cy) (cx*cy) 0 + 0 0 0 1 + where + cx = cos x + sx = sin x + cy = cos y + sy = sin y + cz = cos z + sz = sin z + + {-# INLINE lookAt #-} + lookAt up cam foc = mkMat + xb1 yb1 zb1 0 + xb2 yb2 zb2 0 + xb3 yb3 zb3 0 + tx ty tz 1 + where + (# xb1, xb2, xb3 #) = unpackV3# xb + (# yb1, yb2, yb3 #) = unpackV3# yb + (# zb1, zb2, zb3 #) = unpackV3# zb + zb = normalized $ cam - foc -- Basis vector for "backward", since +Z is behind the camera + xb = normalized $ up `cross` zb -- Basis vector for "right" + yb = zb `cross` xb -- Basis vector for "up" + ncam = -cam + tx = unScalar $ xb `dot` ncam + ty = unScalar $ yb `dot` ncam + tz = unScalar $ zb `dot` ncam + + {-# INLINE perspective #-} + perspective n f fovy aspect = mkMat + dpw 0 0 0 + 0 dph 0 0 + 0 0 a (-1) + 0 0 b 0 + where + hpd = tan (fovy * 0.5) -- height/distance + wpd = aspect * hpd; -- width/distance + dph = recip hpd -- distance/height + dpw = recip wpd -- distance/width + nmf = n - f + a = (n + f) / nmf + b = 2 * n * f / nmf + + {-# INLINE orthogonal #-} + orthogonal n f w h = mkMat + iw 0 0 0 + 0 ih 0 0 + 0 0 a 0 + 0 0 b 1 + where + ih = 2 / h + iw = 2 / w + nmf = n - f + a = 2 / nmf + b = (n + f) / nmf + + {-# INLINE toHomPoint #-} + toHomPoint (unpackV3# -> (# x, y, z #)) = vec4 x y z 1 + + {-# INLINE toHomVector #-} + toHomVector (unpackV3# -> (# x, y, z #)) = vec4 x y z 0 + + {-# INLINE fromHom #-} + fromHom (unpackV4# -> (# x, y, z, w #)) + | w == 0 = vec3 x y z + | otherwise = vec3 (x/w) (y/w) (z/w) + +instance KnownDim n => MatrixInverse Float n where + inverse = inverseViaLU + +instance KnownDim n => MatrixDeterminant Float n where + det = detViaLU diff --git a/src/Numeric/Matrix/Internal/Mat44d.hs b/src/Numeric/Matrix/Internal/Mat44d.hs deleted file mode 100644 --- a/src/Numeric/Matrix/Internal/Mat44d.hs +++ /dev/null @@ -1,186 +0,0 @@ -{-# LANGUAGE MagicHash #-} -{-# LANGUAGE UnboxedTuples #-} -{-# LANGUAGE ViewPatterns #-} -{-# OPTIONS_GHC -fno-warn-orphans #-} - -module Numeric.Matrix.Internal.Mat44d () where - -import qualified Control.Monad.ST as ST -import qualified Numeric.DataFrame.ST as ST -import Numeric.Matrix.Internal -import Numeric.Scalar.Internal -import Numeric.Vector.Internal - -{-# INLINE mkMat #-} -mkMat :: - Double -> Double -> Double -> Double -> - Double -> Double -> Double -> Double -> - Double -> Double -> Double -> Double -> - Double -> Double -> Double -> Double -> - Mat44d -mkMat - _11 _12 _13 _14 - _21 _22 _23 _24 - _31 _32 _33 _34 - _41 _42 _43 _44 - = ST.runST $ do - df <- ST.newDataFrame - ST.writeDataFrameOff df 0 $ scalar _11 - ST.writeDataFrameOff df 1 $ scalar _12 - ST.writeDataFrameOff df 2 $ scalar _13 - ST.writeDataFrameOff df 3 $ scalar _14 - ST.writeDataFrameOff df 4 $ scalar _21 - ST.writeDataFrameOff df 5 $ scalar _22 - ST.writeDataFrameOff df 6 $ scalar _23 - ST.writeDataFrameOff df 7 $ scalar _24 - ST.writeDataFrameOff df 8 $ scalar _31 - ST.writeDataFrameOff df 9 $ scalar _32 - ST.writeDataFrameOff df 10 $ scalar _33 - ST.writeDataFrameOff df 11 $ scalar _34 - ST.writeDataFrameOff df 12 $ scalar _41 - ST.writeDataFrameOff df 13 $ scalar _42 - ST.writeDataFrameOff df 14 $ scalar _43 - ST.writeDataFrameOff df 15 $ scalar _44 - ST.unsafeFreezeDataFrame df - -instance HomTransform4 Double where - {-# INLINE translate4 #-} - translate4 (unpackV4# -> (# x, y, z, _ #)) = mkMat - 1 0 0 0 - 0 1 0 0 - 0 0 1 0 - x y z 1 - - {-# INLINE translate3 #-} - translate3 (unpackV3# -> (# x, y, z #)) = mkMat - 1 0 0 0 - 0 1 0 0 - 0 0 1 0 - x y z 1 - - {-# INLINE rotateX #-} - rotateX a = mkMat - 1 0 0 0 - 0 c s 0 - 0 n c 0 - 0 0 0 1 - where - c = cos a - s = sin a - n = -s - - {-# INLINE rotateY #-} - rotateY a = mkMat - c 0 n 0 - 0 1 0 0 - s 0 c 0 - 0 0 0 1 - where - c = cos a - s = sin a - n = -s - - {-# INLINE rotateZ #-} - rotateZ a = mkMat - c s 0 0 - n c 0 0 - 0 0 1 0 - 0 0 0 1 - where - c = cos a - s = sin a - n = -s - - {-# INLINE rotate #-} - rotate (unpackV3# -> (# x, y, z #)) a = mkMat - (c+xxv) (yxv+zs) (zxv-ys) 0 - (xyv-zs) (c+yyv) (zyv+xs) 0 - (xzv+ys) (yzv-xs) (c+zzv) 0 - 0 0 0 1 - where - c = cos a - v = 1 - c -- v for versine - s = sin a - xxv = x * x * v - xyv = x * y * v - xzv = x * z * v - yxv = xyv - yyv = y * y * v - yzv = y * z * v - zxv = xzv - zyv = yzv - zzv = z * z * v - xs = x * s - ys = y * s - zs = z * s - - {-# INLINE rotateEuler #-} - rotateEuler x y z = mkMat - (cy*cz) (cx*sz+sx*sy*cz) (sx*sz-cx*sy*cz) 0 - (-cy*sz) (cx*cz-sx*sy*sz) (sx*cz+cx*sy*sz) 0 - sy (-sx*cy) (cx*cy) 0 - 0 0 0 1 - where - cx = cos x - sx = sin x - cy = cos y - sy = sin y - cz = cos z - sz = sin z - - {-# INLINE lookAt #-} - lookAt up cam foc = mkMat - xb1 yb1 zb1 0 - xb2 yb2 zb2 0 - xb3 yb3 zb3 0 - tx ty tz 1 - where - (# xb1, xb2, xb3 #) = unpackV3# xb - (# yb1, yb2, yb3 #) = unpackV3# yb - (# zb1, zb2, zb3 #) = unpackV3# zb - zb = normalized $ cam - foc -- Basis vector for "backward", since +Z is behind the camera - xb = normalized $ up `cross` zb -- Basis vector for "right" - yb = zb `cross` xb -- Basis vector for "up" - ncam = -cam - tx = unScalar $ xb `dot` ncam - ty = unScalar $ yb `dot` ncam - tz = unScalar $ zb `dot` ncam - - {-# INLINE perspective #-} - perspective n f fovy aspect = mkMat - dpw 0 0 0 - 0 dph 0 0 - 0 0 a (-1) - 0 0 b 0 - where - hpd = tan (fovy * 0.5) -- height/distance - wpd = aspect * hpd; -- width/distance - dph = recip hpd -- distance/height - dpw = recip wpd -- distance/width - nmf = n - f - a = (n + f) / nmf - b = 2 * n * f / nmf - - {-# INLINE orthogonal #-} - orthogonal n f w h = mkMat - iw 0 0 0 - 0 ih 0 0 - 0 0 a 0 - 0 0 b 1 - where - ih = 2 / h - iw = 2 / w - nmf = n - f - a = 2 / nmf - b = (n + f) / nmf - - {-# INLINE toHomPoint #-} - toHomPoint (unpackV3# -> (# x, y, z #)) = vec4 x y z 1 - - {-# INLINE toHomVector #-} - toHomVector (unpackV3# -> (# x, y, z #)) = vec4 x y z 0 - - {-# INLINE fromHom #-} - fromHom (unpackV4# -> (# x, y, z, w #)) - | w == 0 = vec3 x y z - | otherwise = vec3 (x/w) (y/w) (z/w) diff --git a/src/Numeric/Matrix/Internal/Mat44f.hs b/src/Numeric/Matrix/Internal/Mat44f.hs deleted file mode 100644 --- a/src/Numeric/Matrix/Internal/Mat44f.hs +++ /dev/null @@ -1,186 +0,0 @@ -{-# LANGUAGE MagicHash #-} -{-# LANGUAGE UnboxedTuples #-} -{-# LANGUAGE ViewPatterns #-} -{-# OPTIONS_GHC -fno-warn-orphans #-} - -module Numeric.Matrix.Internal.Mat44f () where - -import qualified Control.Monad.ST as ST -import qualified Numeric.DataFrame.ST as ST -import Numeric.Matrix.Internal -import Numeric.Scalar.Internal -import Numeric.Vector.Internal - -{-# INLINE mkMat #-} -mkMat :: - Float -> Float -> Float -> Float -> - Float -> Float -> Float -> Float -> - Float -> Float -> Float -> Float -> - Float -> Float -> Float -> Float -> - Mat44f -mkMat - _11 _12 _13 _14 - _21 _22 _23 _24 - _31 _32 _33 _34 - _41 _42 _43 _44 - = ST.runST $ do - df <- ST.newDataFrame - ST.writeDataFrameOff df 0 $ scalar _11 - ST.writeDataFrameOff df 1 $ scalar _12 - ST.writeDataFrameOff df 2 $ scalar _13 - ST.writeDataFrameOff df 3 $ scalar _14 - ST.writeDataFrameOff df 4 $ scalar _21 - ST.writeDataFrameOff df 5 $ scalar _22 - ST.writeDataFrameOff df 6 $ scalar _23 - ST.writeDataFrameOff df 7 $ scalar _24 - ST.writeDataFrameOff df 8 $ scalar _31 - ST.writeDataFrameOff df 9 $ scalar _32 - ST.writeDataFrameOff df 10 $ scalar _33 - ST.writeDataFrameOff df 11 $ scalar _34 - ST.writeDataFrameOff df 12 $ scalar _41 - ST.writeDataFrameOff df 13 $ scalar _42 - ST.writeDataFrameOff df 14 $ scalar _43 - ST.writeDataFrameOff df 15 $ scalar _44 - ST.unsafeFreezeDataFrame df - -instance HomTransform4 Float where - {-# INLINE translate4 #-} - translate4 (unpackV4# -> (# x, y, z, _ #)) = mkMat - 1 0 0 0 - 0 1 0 0 - 0 0 1 0 - x y z 1 - - {-# INLINE translate3 #-} - translate3 (unpackV3# -> (# x, y, z #)) = mkMat - 1 0 0 0 - 0 1 0 0 - 0 0 1 0 - x y z 1 - - {-# INLINE rotateX #-} - rotateX a = mkMat - 1 0 0 0 - 0 c s 0 - 0 n c 0 - 0 0 0 1 - where - c = cos a - s = sin a - n = -s - - {-# INLINE rotateY #-} - rotateY a = mkMat - c 0 n 0 - 0 1 0 0 - s 0 c 0 - 0 0 0 1 - where - c = cos a - s = sin a - n = -s - - {-# INLINE rotateZ #-} - rotateZ a = mkMat - c s 0 0 - n c 0 0 - 0 0 1 0 - 0 0 0 1 - where - c = cos a - s = sin a - n = -s - - {-# INLINE rotate #-} - rotate (unpackV3# -> (# x, y, z #)) a = mkMat - (c+xxv) (yxv+zs) (zxv-ys) 0 - (xyv-zs) (c+yyv) (zyv+xs) 0 - (xzv+ys) (yzv-xs) (c+zzv) 0 - 0 0 0 1 - where - c = cos a - v = 1 - c -- v for versine - s = sin a - xxv = x * x * v - xyv = x * y * v - xzv = x * z * v - yxv = xyv - yyv = y * y * v - yzv = y * z * v - zxv = xzv - zyv = yzv - zzv = z * z * v - xs = x * s - ys = y * s - zs = z * s - - {-# INLINE rotateEuler #-} - rotateEuler x y z = mkMat - (cy*cz) (cx*sz+sx*sy*cz) (sx*sz-cx*sy*cz) 0 - (-cy*sz) (cx*cz-sx*sy*sz) (sx*cz+cx*sy*sz) 0 - sy (-sx*cy) (cx*cy) 0 - 0 0 0 1 - where - cx = cos x - sx = sin x - cy = cos y - sy = sin y - cz = cos z - sz = sin z - - {-# INLINE lookAt #-} - lookAt up cam foc = mkMat - xb1 yb1 zb1 0 - xb2 yb2 zb2 0 - xb3 yb3 zb3 0 - tx ty tz 1 - where - (# xb1, xb2, xb3 #) = unpackV3# xb - (# yb1, yb2, yb3 #) = unpackV3# yb - (# zb1, zb2, zb3 #) = unpackV3# zb - zb = normalized $ cam - foc -- Basis vector for "backward", since +Z is behind the camera - xb = normalized $ up `cross` zb -- Basis vector for "right" - yb = zb `cross` xb -- Basis vector for "up" - ncam = -cam - tx = unScalar $ xb `dot` ncam - ty = unScalar $ yb `dot` ncam - tz = unScalar $ zb `dot` ncam - - {-# INLINE perspective #-} - perspective n f fovy aspect = mkMat - dpw 0 0 0 - 0 dph 0 0 - 0 0 a (-1) - 0 0 b 0 - where - hpd = tan (fovy * 0.5) -- height/distance - wpd = aspect * hpd; -- width/distance - dph = recip hpd -- distance/height - dpw = recip wpd -- distance/width - nmf = n - f - a = (n + f) / nmf - b = 2 * n * f / nmf - - {-# INLINE orthogonal #-} - orthogonal n f w h = mkMat - iw 0 0 0 - 0 ih 0 0 - 0 0 a 0 - 0 0 b 1 - where - ih = 2 / h - iw = 2 / w - nmf = n - f - a = 2 / nmf - b = (n + f) / nmf - - {-# INLINE toHomPoint #-} - toHomPoint (unpackV3# -> (# x, y, z #)) = vec4 x y z 1 - - {-# INLINE toHomVector #-} - toHomVector (unpackV3# -> (# x, y, z #)) = vec4 x y z 0 - - {-# INLINE fromHom #-} - fromHom (unpackV4# -> (# x, y, z, w #)) - | w == 0 = vec3 x y z - | otherwise = vec3 (x/w) (y/w) (z/w) diff --git a/src/Numeric/Matrix/LU.hs b/src/Numeric/Matrix/LU.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/Matrix/LU.hs @@ -0,0 +1,201 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE PolyKinds #-} +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE StandaloneDeriving #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UndecidableInstances #-} +module Numeric.Matrix.LU + ( MatrixLU (..), LU (..) + , luSolveR, luSolveL + , detViaLU, inverseViaLU + ) where + +import Control.Monad +import Control.Monad.ST +import Data.Kind +import Numeric.DataFrame.Internal.PrimArray +import Numeric.DataFrame.ST +import Numeric.DataFrame.SubSpace +import Numeric.DataFrame.Type +import Numeric.Dimensions +import Numeric.Matrix.Internal +import Numeric.Scalar.Internal +import Numeric.Subroutine.SolveTriangular + + +-- | Result of LU factorization with Partial Pivoting +-- \( PA = LU \). +data LU (t :: Type) (n :: Nat) + = LU + { luLower :: Matrix t n n + -- ^ Unit lower triangular matrix \(L\). + -- All elements on the diagonal of @L@ equal @1@. + -- The rest of the elements satisfy \(|l_{ij}| \leq 1\). + , luUpper :: Matrix t n n + -- ^ Upper triangular matrix \(U\) + , luPerm :: Matrix t n n + -- ^ Row permutation matrix \(P\) + , luPermDet :: Scalar t + -- ^ Sign of permutation @luPermDet == det . luPerm@; \(|P| = \pm 1\). + } + +deriving instance (Show t, PrimBytes t, KnownDim n) => Show (LU t n) +deriving instance (Eq (Matrix t n n), Eq t) => Eq (LU t n) + +class (KnownDim n, Ord t, Fractional t, PrimBytes t, KnownBackend t '[n,n]) + => MatrixLU t (n :: Nat) where + -- | Compute LU factorization with Partial Pivoting + lu :: Matrix t n n -> LU t n + + +instance (KnownDim n, Ord t, Fractional t, PrimBytes t, KnownBackend t '[n,n]) + => MatrixLU t n where + lu a = runST $ do + pPtr <- unsafeThawDataFrame $ iwgen @_ @'[n] @'[] (\(Idx i :* U) -> S i) + uPtr <- thawDataFrame a + lPtr <- newDataFrame + temp <- newDataFrame + detPositive <- luInplace temp pPtr uPtr + p <- unsafeFreezeDataFrame pPtr + -- split U and L + forM_ [0..n-1] $ \i -> do + let ni = n*i + forM_ [0..n-1] $ \j -> case compare i j of + GT -> do + lij <- readDataFrameOff uPtr (ni + j) + writeDataFrameOff uPtr (ni + j) 0 + writeDataFrameOff lPtr (ni + j) lij + EQ -> writeDataFrameOff lPtr (ni + j) 1 + LT -> writeDataFrameOff lPtr (ni + j) 0 + luLower <- unsafeFreezeDataFrame lPtr + luUpper <- unsafeFreezeDataFrame uPtr + let luPermDet = if detPositive then 1 else -1 + luPerm = iwgen @_ @'[n,n] @'[] + (\(Idx i :* Idx j :* U) -> if S j == p ! i then 1 else 0) + return LU {..} + where + n = fromIntegral (dimVal' @n) :: Int + + +-- | Solve @Ax = b@ problem given LU decomposition of A. +luSolveR :: + forall t (n :: Nat) (ds :: [Nat]) + . (MatrixLU t n, Dimensions ds) + => LU t n -> DataFrame t (n :+ ds) -> DataFrame t (n :+ ds) +luSolveR LU {..} b = runST $ do + xPtr <- thawDataFrame (luPerm %* b) -- NB: wasting resources! + solveLowerTriangularR luLower xPtr + solveUpperTriangularR luUpper xPtr + unsafeFreezeDataFrame xPtr + +-- | Solve @xA = b@ problem given LU decomposition of A. +luSolveL :: + forall t (n :: Nat) (ds :: [Nat]) + . (MatrixLU t n, Dimensions ds) + => LU t n -> DataFrame t (ds +: n) -> DataFrame t (ds +: n) +luSolveL LU {..} b + | dn <- dim @n + , dds <- dims @ds + , Dims <- Snoc dds dn + , Dict <- Dict @(SnocList ds n _) + = runST $ do + xPtr <- thawDataFrame b + solveUpperTriangularL xPtr luUpper + solveLowerTriangularL xPtr luLower + (%* luPerm) <$> unsafeFreezeDataFrame xPtr +luSolveL _ _ = error "luSolveL: impossible pattern" + +-- | Calculate inverse of a matrix via LU decomposition +inverseViaLU :: forall (t :: Type) (n :: Nat) + . MatrixLU t n => Matrix t n n -> Matrix t n n +inverseViaLU a = runST $ do + xPtr <- unsafeThawDataFrame luPerm -- luPerm is only ever used once + solveLowerTriangularR luLower xPtr + solveUpperTriangularR luUpper xPtr + unsafeFreezeDataFrame xPtr + where + LU {..} = lu a +-- perfectly correct, but slightly slower versions: +-- inverseViaLU a = luSolveR (lu a) eye +-- inverseViaLU a = luSolveL (lu a) eye + +-- | Calculate determinant of a matrix via LU decomposition +detViaLU :: forall (t :: Type) (n :: Nat) + . MatrixLU t n => Matrix t n n -> Scalar t +detViaLU m = foldl (\x off -> scalar (ixOff off luUpper) * x) luPermDet [0,n+1..n*n] + where + n = fromIntegral (dimVal' @n) :: Int + LU {..} = lu m + +{- | +Run LU decomposition with partial pivoting inplace, such that +the upper triangular part of matrix \(A\) becomes \(U\) and +the lower triangular part (without diagonal) of matrix \(A\) becomes \(L\). + +\(U\) is upper triangular. +\(L\) is unit lower triangular; all diagonal elements of \(L\) are implicit and +equal to 1; the rest of the elements a smaller than one \(|l_{ij}| \leq 1\). + +Pivoting is represented as a permutation vector \(p\); +returned value is the sign of the permutation (positive if @True@, negative otherwise). + +NB: Initialize \(p\) with indices @0..n-1@. + +Reference: Algorithm 3.4.1 on p.128 + of "Matrix Computations" 4th edition by G. H. Golub and C. F. Van Loan. + -} +luInplace :: + forall (s :: Type) (t :: Type) (n :: Nat) + . (PrimBytes t, Fractional t, Ord t, KnownDim n) + => STDataFrame s t '[n] -- ^ Temporary buffer + -> STDataFrame s Word '[n] -- ^ Current state of permutation \(p\) + -> STDataFrame s t '[n,n] -- ^ Current state of \(A\) + -> ST s Bool +luInplace temp pPtr aPtr = foldM (\b -> fmap (b /=) . go) True [0..n-2] + where + n = fromIntegral (dimVal' @n) :: Int + + -- Runs an iteration of the algorithm; + --- returns whether there was a swap of rows. + go :: Int -> ST s Bool + go k = do + mu <- findPivot k + let swapped = k /= mu + when swapped $ swapRows k mu + akk <- readDataFrameOff aPtr (k*(n+1)) + when (akk /= 0) $ do + let rakk = recip akk + forM_ [k+1..n-1] $ \i -> do + let ni = n*i + aik <- (rakk *) <$> readDataFrameOff aPtr (ni + k) + writeDataFrameOff aPtr (ni + k) aik + forM_ [k+1..n-1] $ \j -> do + akj <- readDataFrameOff aPtr (n*k + j) + aij <- readDataFrameOff aPtr (ni + j) + writeDataFrameOff aPtr (ni + j) (aij - aik*akj) + return swapped + + findPivot :: Int -> ST s Int + findPivot k = snd <$> foldM findPivotF (0, k) [k..n-1] + where + findPivotF :: (Scalar t, Int) -> Int -> ST s (Scalar t, Int) + findPivotF aj@(a, _) i = do + x <- abs <$> readDataFrameOff aPtr (n*i + k) + return (if x > a then (x, i) else aj) + + swapRows :: Int -> Int -> ST s () + swapRows i j = do + let iPtr = subDataFrameView' (fromIntegral i :* U) aPtr + jPtr = subDataFrameView' (fromIntegral j :* U) aPtr + copyMutableDataFrame' U iPtr temp + copyMutableDataFrame' U jPtr iPtr + copyMutableDataFrame' U temp jPtr + t <- readDataFrameOff pPtr i + readDataFrameOff pPtr j >>= writeDataFrameOff pPtr i + writeDataFrameOff pPtr j t diff --git a/src/Numeric/Matrix/QR.hs b/src/Numeric/Matrix/QR.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/Matrix/QR.hs @@ -0,0 +1,231 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE PolyKinds #-} +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE StandaloneDeriving #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UndecidableInstances #-} +module Numeric.Matrix.QR + ( QR (..), LQ (..), MatrixQR (..) + , detViaQR, inverseViaQR + , qrSolveR, qrSolveL + ) where + + +import Control.Monad +import Control.Monad.ST +import Data.Kind +import Numeric.Basics +import Numeric.DataFrame.ST +import Numeric.DataFrame.SubSpace +import Numeric.DataFrame.Type +import Numeric.Dimensions +import Numeric.Matrix.Internal +import Numeric.Scalar.Internal +import Numeric.Subroutine.Householder +import Numeric.Subroutine.SolveTriangular + +import Unsafe.Coerce (unsafeCoerce) + +-- | Result of QR factorization +-- \( A = QR \). +data QR (t :: Type) (n :: Nat) (m :: Nat) + = QR + { qrQ :: Matrix t n n + -- ^ Orthogonal matrix \( Q \) + , qrQDet :: Scalar t + -- ^ A shortcut for evaluating a determinant of \( |Q| = \pm 1 \) + , qrR :: Matrix t n m + -- ^ Upper-triangular matrix \( R \) + } + +-- | Result of LQ factorization +-- \( A = LQ \). +data LQ (t :: Type) (n :: Nat) (m :: Nat) + = LQ + { lqL :: Matrix t n m + -- ^ Lower-triangular matrix \( L \) + , lqQ :: Matrix t m m + -- ^ Orthogonal matrix \( Q \) + , lqQDet :: Scalar t + -- ^ A shortcut for evaluating a determinant of \( |Q| = \pm 1 \) + } + +deriving instance ( Eq t, PrimBytes t + , KnownDim n, KnownDim m) => Eq (QR t n m) +deriving instance ( Show t, PrimBytes t + , KnownDim n, KnownDim m) => Show (QR t n m) +deriving instance ( Eq t, PrimBytes t + , KnownDim n, KnownDim m) => Eq (LQ t n m) +deriving instance ( Show t, PrimBytes t + , KnownDim n, KnownDim m) => Show (LQ t n m) + + +class (PrimBytes t, Ord t, Epsilon t, KnownDim n, KnownDim m) + => MatrixQR t (n :: Nat) (m :: Nat) where + -- | Compute QR factorization + qr :: Matrix t n m -> QR t n m + -- | Compute LQ factorization + lq :: Matrix t n m -> LQ t n m + +instance (PrimBytes t, Ord t, Epsilon t, KnownDim n, KnownDim m) + => MatrixQR t (n :: Nat) (m :: Nat) where + qr a + | lim == 0 = QR undefined undefined 1 + | otherwise = runST $ do + uPtr <- newDataFrame + pPtr <- unsafeThawDataFrame eye + rPtr <- thawDataFrame a + detNegative <- + let f x i = (x /=) <$> householderReflectionInplaceL + uPtr pPtr rPtr (Idx i :* Idx i :* U) + in foldM f False [0..lim - 1] + qrR <- unsafeFreezeDataFrame rPtr + qrQ <- unsafeFreezeDataFrame pPtr + let qrQDet = if detNegative then -1 else 1 + return QR {..} + where + n = dimVal' @n + m = dimVal' @m + lim = min n m + lq a + | lim == 0 = LQ undefined undefined 1 + | otherwise = runST $ do + uPtr <- newDataFrame + pPtr <- unsafeThawDataFrame eye + lPtr <- thawDataFrame a + detNegative <- + let f x i = (x /=) <$> householderReflectionInplaceR + uPtr pPtr lPtr (Idx i :* Idx i :* U) + in foldM f False [0..lim - 1] + lqL <- unsafeFreezeDataFrame lPtr + -- WARNING! todo: get rid of transpose. + lqQ <- transpose <$> unsafeFreezeDataFrame pPtr + let lqQDet = if detNegative then -1 else 1 + return LQ {..} + where + n = dimVal' @n + m = dimVal' @m + lim = min n m + +-- | Calculate determinant of a matrix via QR decomposition +detViaQR :: forall t n . MatrixQR t n n => Matrix t n n -> Scalar t +detViaQR m = foldl (\x off -> scalar (ixOff off qrR) * x) qrQDet [0,n+1..n*n] + where + n = fromIntegral (dimVal' @n) :: Int + QR {..} = qr m + +-- | Calculate inverse of a matrix via QR decomposition +inverseViaQR :: forall t n . MatrixQR t n n => Matrix t n n -> Matrix t n n +inverseViaQR = (`qrSolveR` eye) + +{- | +Compute a QR or LQ decomposition of matrix \( A : n \times m \), +and solve a system of linear equations \( Ax = b \). + +If \( n >= m \) QR decomposition is used; +if \( n > m \) this function solves linear least squares problem. +If \( n < m \) (underdetermined system) LQ decomposition is used + to yield a minimum norm solution. + -} +qrSolveR :: + forall t (n :: Nat) (m :: Nat) (ds :: [Nat]) + . (MatrixQR t n m, Dimensions ds) + => Matrix t n m -> DataFrame t (n :+ ds) -> DataFrame t (m :+ ds) +qrSolveR a b = case compareDim dn dm of + SEQ | Dict <- (unsafeCoerce (Dict @(m ~ m)) :: Dict (m ~ n)) + -> runST $ do + let QR {..} = qr a + xPtr <- thawDataFrame (transpose qrQ %* b) -- NB: make a stateful product for transposed mat + solveUpperTriangularR qrR xPtr + unsafeFreezeDataFrame xPtr + SGT | Dict <- unsafeCoerce (Dict @(m <= m)) :: Dict (m <= n) + , Dict <- unsafeCoerce (Dict @(n ~ n)) :: Dict ((((n - m) + 1) + m) ~ (n + 1) ) + , D <- minusDim dn dm `plusDim` D1 + -> runST $ do + let QR {..} = qr a + i0 :: Idxs '[n - m + 1] + i0 = Idx 0 :* U + xPtr <- thawDataFrame ( slice i0 $ transpose qrQ %* b + ) -- NB: make a stateful product for transposed mat + solveUpperTriangularR qrR xPtr + unsafeFreezeDataFrame xPtr + SLT | Dict <- unsafeCoerce (Dict @(m <= m)) :: Dict (n <= m) + , Dict <- unsafeCoerce (Dict @(n ~ n)) :: Dict ((((m - n) + 1) + n) ~ (m + 1)) + , Dict <- unsafeCoerce (Dict @(n ~ n)) :: Dict (((n + 1) + (m - n)) ~ (m + 1)) + , dd@D <- minusDim dm dn + , D <- dd `plusDim` D1 + , D <- dn `plusDim` D1 + -> runST $ do + let LQ {..} = lq a + i0 :: Idxs '[m - n + 1] + i0 = Idx 0 :* U + iz :: Idxs '[n + 1] + iz = maxBound :* U + xPtr <- newDataFrame + copyDataFrame i0 b xPtr + copyDataFrame iz (0 :: DataFrame t ((m - n) ': ds)) xPtr + solveLowerTriangularR lqL (subDataFrameView i0 xPtr) + (transpose lqQ %*) <$> unsafeFreezeDataFrame xPtr + -- NB: make a stateful product for transposed mat + _ -> error "qrSolveR: impossible pattern" + where + dn = dim @n + dm = dim @m + + +{- | +Compute a QR or LQ decomposition of matrix \( A : n \times m \), +and solve a system of linear equations \( xA = b \). + +If \( n <= m \) LQ decomposition is used; +if \( n < m \) this function solves linear least squares problem. +If \( n > m \) (underdetermined system) QR decomposition is used + to yield a minimum norm solution. + -} +qrSolveL :: + forall t (n :: Nat) (m :: Nat) (ds :: [Nat]) + . (MatrixQR t n m, Dimensions ds) + => Matrix t n m -> DataFrame t (ds +: m) -> DataFrame t (ds +: n) +qrSolveL a b + | Dims <- Snoc (dims @ds) dn + , Dims <- Snoc (dims @ds) dm + , Dict <- Dict @(SnocList ds n _) -- remind GHC that we have (ds +: n) + , Dict <- Dict @(SnocList ds m _) -- remind GHC that we have (ds +: m) + = case compareDim dn dm of + SEQ | Dict <- (unsafeCoerce (Dict @(m ~ m)) :: Dict (m ~ n)) + -> runST $ do + let LQ {..} = lq a + xPtr <- thawDataFrame (b %* transpose lqQ) -- NB: make a stateful product for transposed mat + solveLowerTriangularL xPtr lqL + unsafeFreezeDataFrame xPtr + SGT | Dict <- unsafeCoerce (Dict @(m <= m)) :: Dict (m <= n) + , Dict <- unsafeCoerce (Dict @(n ~ n)) :: Dict ((m + (n - m)) ~ n) + , D <- minusDim dn dm + -> runST $ do + let QR {..} = qr a + xPtr <- thawDataFrame b + solveUpperTriangularL xPtr qrR + ewmap @t @ds @'[n] (\x -> qrQ %* appendDF x (0 :: DataFrame t '[n - m])) + <$> unsafeFreezeDataFrame xPtr + SLT | Dict <- unsafeCoerce (Dict @(m <= m)) :: Dict (n <= m) + , Dict <- unsafeCoerce (Dict @(n ~ n)) :: Dict ((((m - n) + 1) + n) ~ (m + 1)) + , D <- minusDim dm dn `plusDim` D1 + -> runST $ do + let LQ {..} = lq a + i0 :: Idx (m - n + 1) + i0 = Idx 0 + xPtr <- thawDataFrame (b %* transpose lqQ) -- NB: make a stateful product for transposed mat + solveLowerTriangularL xPtr lqL + ewmap @t @ds @'[n] (sslice i0) + <$> unsafeFreezeDataFrame xPtr + _ -> error "qrSolveL/compareDim: impossible pattern" + | otherwise = error "qrSolveL: impossible pattern" + where + dn = dim @n + dm = dim @m diff --git a/src/Numeric/Matrix/SVD.hs b/src/Numeric/Matrix/SVD.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/Matrix/SVD.hs @@ -0,0 +1,799 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE PolyKinds #-} +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE StandaloneDeriving #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UndecidableInstances #-} +module Numeric.Matrix.SVD + ( MatrixSVD (..), SVD (..) + , svd1, svd2, svd3, svd3q + ) where + +import Control.Monad +import Control.Monad.ST +import Data.Kind +import Numeric.Basics +import Numeric.DataFrame.Internal.PrimArray +import Numeric.DataFrame.ST +import Numeric.DataFrame.SubSpace +import Numeric.DataFrame.Type +import Numeric.Dimensions +import Numeric.Matrix.Bidiagonal +import Numeric.Matrix.Internal +import Numeric.Quaternion.Internal +import Numeric.Scalar.Internal +import Numeric.Subroutine.Sort +import Numeric.Tuple +import Numeric.Vector.Internal + +-- | Result of SVD factorization +-- @ M = svdU %* asDiag svdS %* transpose svdV @. +-- +-- Invariants: +-- +-- * Singular values `svdS` are in non-increasing order and are non-negative. +-- * svdU and svdV are orthogonal matrices +-- * det svdU == 1 +-- +-- NB: +data SVD (t :: Type) (n :: Nat) (m :: Nat) + = SVD + { svdU :: Matrix t n n + -- ^ Left-singular basis matrix + , svdS :: Vector t (Min n m) + -- ^ Vector of singular values + , svdV :: Matrix t m m + -- ^ Right-singular basis matrix + } + +deriving instance ( Show t, PrimBytes t + , KnownDim n, KnownDim m, KnownDim (Min n m)) + => Show (SVD t n m) +deriving instance ( Eq (Matrix t n n) + , Eq (Matrix t m m) + , Eq (Vector t (Min n m))) + => Eq (SVD t n m) + +class RealFloatExtras t + => MatrixSVD (t :: Type) (n :: Nat) (m :: Nat) where + -- | Compute SVD factorization of a matrix + svd :: IterativeMethod => Matrix t n m -> SVD t n m + +-- | Obvious dummy implementation of SVD for 1x1 matrices +svd1 :: (PrimBytes t, Num t, Eq t) => Matrix t 1 1 -> SVD t 1 1 +svd1 m = SVD + { svdU = 1 + , svdS = broadcast $ abs x + , svdV = broadcast $ if x == 0 then 1 else signum x + } + where + x = ixOff 0 m + +-- | SVD of a 2x2 matrix can be computed analytically +-- +-- Related discussion: +-- +-- https://scicomp.stackexchange.com/questions/8899/robust-algorithm-for-2-times-2-svd/ +-- +-- https://ieeexplore.ieee.org/document/486688 +svd2 :: forall t . RealFloatExtras t => Matrix t 2 2 -> SVD t 2 2 +svd2 (DF2 (DF2 m00 m01) (DF2 m10 m11)) = + SVD + { svdU = DF2 (DF2 uc us) + (DF2 (negate us) uc) + , svdS = DF2 sigma1 (abs sigma2) + , svdV = DF2 (DF2 vc (sg2 (negate vs))) + (DF2 vs (sg2 vc )) + } + where + x1 = m00 - m11 -- 2F + x2 = m00 + m11 -- 2E + y1 = m01 + m10 -- 2G + y2 = m01 - m10 -- 2H + yy = y1*y2 + h1 = hypot x1 y1 -- h1 >= abs x1 + h2 = hypot x2 y2 -- h2 >= abs x2 + sigma1 = 0.5 * (h2 + h1) -- sigma1 >= abs sigma2 + sigma2 = 0.5 * (h2 - h1) -- can be negative, which is accounted by sg2 + sg2 = negateUnless (sigma2 >= 0) + hx1 = h1 + x1 + hx2 = h2 + x2 + hxhx = hx1*hx2 + hxy = hx1*y2 + yhx = y1*hx2 + (uc', us', vc', vs') = case (x1 > 0 || y1 /= 0, x2 > 0 || y2 /= 0) of + (True , True ) -> (hxhx + yy, hxy - yhx, hxhx - yy, hxy + yhx) + (True , False) -> (y1, hx1, -y1, hx1) + (False, True ) -> (y2, -hx2, -y2, hx2) + (False, False) -> (1, 0, -1, 0) + ru = recip $ hypot uc' us' + rv = recip $ hypot vc' vs' + uc = uc' * ru + us = us' * ru + vc = vc' * rv + vs = vs' * rv + +-- | Get SVD decomposition of a 3x3 matrix using `svd3q` function. +-- +-- This function reorders the singular components under the hood to make sure +-- @s1 >= s2 >= s3 >= 0@. +-- Thus, it has some overhead on top of `svd3q`. +svd3 :: forall t . (Quaternion t, RealFloatExtras t) => Matrix t 3 3 -> SVD t 3 3 +svd3 m = SVD + { svdU = toMatrix33 u + , svdS = DF3 s1 s2 s3' + , svdV = neg3If (s3 < 0) (toMatrix33 v) + } + where + (u, (DF3 s1 s2 s3), v) = svd3q m + s3' = abs s3 + neg3If :: Bool -> Matrix t 3 3 -> Matrix t 3 3 + neg3If False = id + neg3If True = ewmap @t @'[3] neg3 + neg3 :: Vector t 3 -> Vector t 3 + neg3 (DF3 a b c) = DF3 a b (negate c) + + +-- | Get SVD decomposition of a 3x3 matrix, with orthogonal matrices U and V +-- represented as quaternions. +-- Important: U and V are bound to be rotations at the expense of the last +-- singular value being possibly negative. +-- +-- This is an adoptation of a specialized 3x3 SVD algorithm described in +-- "Computing the Singular Value Decomposition of 3x3 matrices +-- with minimal branching and elementary floating point operations", +-- by A. McAdams, A. Selle, R. Tamstorf, J. Teran, E. Sifakis. +-- +-- http://pages.cs.wisc.edu/~sifakis/papers/SVD_TR1690.pdf +svd3q :: forall t . (Quaternion t, RealFloatExtras t) + => Matrix t 3 3 -> (Quater t, Vector t 3, Quater t) +svd3q m = (u, s, v) + where + v = jacobiEigenQ (transpose m %* m) + (s, u) = uncurry fixSigns $ qrDecomposition3 (m %* toMatrix33 v) + -- last bit: make sure s1 >= s2 >= 0 + fixSigns :: Vector t 3 -> Quater t -> (Vector t 3, Quater t) + fixSigns (DF3 s1 s2 s3) q@(Quater a b c d) = case (s1 >= 0, s2 >= 0) of + (True , True ) -> (mk3 s1 s2 s3, q) + (False, True ) -> (mk3 (negate s1) s2 (negate s3), Quater (-c) d a (-b)) + (True , False) -> (mk3 s1 (negate s2) (negate s3), Quater d c (-b) (-a)) + (False, False) -> (mk3 (negate s1) (negate s2) s3, Quater b (-a) d (-c)) + -- one more thing: + -- the singular values are ordered, but may have small errors; + -- as a result adjacent values may seem to be out of order by a very small number + mk3 :: Scalar t -> Scalar t -> Scalar t -> Vector t 3 + mk3 s1 s2 s3' = case (s1 >= s2, s1 >= abs s3, s2 >= abs s3) of + (True , True , True ) -> DF3 s1 s2 s3' -- s1 >= s2 >= s3 + (True , True , False) -> DF3 s1 s3 (cs s2) -- s1 >= s3 > s2 + (True , False, _ ) -> DF3 s3 s1 (cs s2) -- s3 > s1 >= s2 + (False, True , True ) -> DF3 s2 s1 s3' -- s2 > s1 >= s3 + (False, _ , False) -> DF3 s3 s2 (cs s1) -- s3 > s2 > s1 + (False, False, True ) -> DF3 s2 s3 (cs s1) -- s2 >= s3 > s1 + where + s3 = abs s3' + cs = negateUnless (s3' >= 0) + + +-- | Approximate values for cos (a/2) and sin (a/2) of a Givens rotation for +-- a 2x2 symmetric matrix. (Algorithm 2) +jacobiGivensQ :: forall t . RealFloatExtras t => t -> t -> t -> (t, t) +jacobiGivensQ aii aij ajj + | g*sh*sh < ch*ch = (w * ch, w * sh) + | otherwise = (c', s') + where + ch = 2 * (aii-ajj) + sh = aij + w = recip $ hypot ch sh + g = 5.82842712474619 :: t -- 3 + sqrt 8 + c' = 0.9238795325112867 :: t -- cos (pi/8) + s' = 0.3826834323650898 :: t -- sin (pi/8) + + +-- | A quaternion for a QR Givens iteration +qrGivensQ :: forall t . RealFloatExtras t => t -> t -> (t, t) +qrGivensQ a1 a2 + | a1 < 0 = (sh * w, ch * w) + | otherwise = (ch * w, sh * w) + where + rho2 = a1*a1 + a2*a2 + sh = if rho2 > M_EPS then a2 else 0 + ch = abs a1 + sqrt (max rho2 M_EPS) + w = recip $ hypot ch sh -- TODO: consider something like a hypot + + +-- | One iteration of the Jacobi algorithm on a symmetric 3x3 matrix +-- +-- The three words arguments are indices: +-- 0 <= i /= j /= k <= 2 +jacobiEigen3Iteration :: (Quaternion t, RealFloatExtras t) + => Int -> Int -> Int + -> STDataFrame s t '[3,3] + -> ST s (Quater t) +jacobiEigen3Iteration i j k sPtr = do + sii <- readDataFrameOff sPtr ii + sij <- readDataFrameOff sPtr ij + sjj <- readDataFrameOff sPtr jj + sik <- readDataFrameOff sPtr ik + sjk <- readDataFrameOff sPtr jk + -- Coefficients for a quaternion corresponding to a Givens rotation + let (ch, sh) = jacobiGivensQ sii sij sjj + a = ch*ch - sh*sh + b = 2 * sh*ch + aa = a * a + ab = a * b + bb = b * b + -- update the matrix + writeDataFrameOff sPtr ii $ + aa * sii + 2 * ab * sij + bb * sjj + writeDataFrameOff sPtr ij $ + ab * (sjj - sii) + (aa - bb) * sij + writeDataFrameOff sPtr jj $ + bb * sii - 2 * ab * sij + aa * sjj + writeDataFrameOff sPtr ik $ a * sik + b * sjk + writeDataFrameOff sPtr jk $ a * sjk - b * sik + + -- write the quaternion + qPtr <- unsafeThawDataFrame 0 + writeDataFrameOff qPtr k (negate sh) + writeDataFrameOff qPtr 3 ch + fromVec4 <$> unsafeFreezeDataFrame qPtr + where + ii = i*3 + i + ij = if i < j then i*3 + j else j*3 + i + jj = j*3 + j + ik = if i < k then i*3 + k else k*3 + i + jk = if j < k then j*3 + k else k*3 + j + + +-- | Total number of the Givens rotations during the Jacobi eigendecomposition +-- part of the 3x3 SVD equals eigenItersX3*3. +-- Value `eigenItersX3 = 6` corresponds to 18 iterations and gives a good precision. +eigenItersX3 :: Int +eigenItersX3 = 12 + +-- | Run a few iterations of the Jacobi algorithm on a real-valued 3x3 symmetric matrix. +-- The eigenvectors basis of such matrix is orthogonal, and can be represented as +-- a quaternion. +jacobiEigenQ :: forall t + . (Quaternion t, RealFloatExtras t) + => Matrix t 3 3 -> Quater t +jacobiEigenQ m = runST $ do + mPtr <- thawDataFrame m + q <- go eigenItersX3 mPtr 1 + s1 <- readDataFrameOff mPtr 0 + s2 <- readDataFrameOff mPtr 4 + s3 <- readDataFrameOff mPtr 8 + return $ sortQ s1 s2 s3 * q + where + go :: Int -> STDataFrame s t '[3,3] -> Quater t -> ST s (Quater t) + go 0 _ q = pure q + + -- -- primitive cyclic iteration; + -- -- fast, but the convergence is not perfect + -- -- + -- -- set eigenItersX3 = 6 for good precision + -- go n p q = do + -- q1 <- jacobiEigen3Iteration 0 1 2 p + -- q2 <- jacobiEigen3Iteration 1 2 0 p + -- q3 <- jacobiEigen3Iteration 2 0 1 p + -- go (n - 1) p (q3 * q2 * q1 * q) + + -- Pick the largest element on lower triangle; + -- slow because of branching, but has a better convergence + -- + -- set eigenItersX3 = 12 for good precision + -- (slightly faster than the cyclic version with -O0) + go n p q = do + a10 <- abs <$> readDataFrameOff p 1 + a20 <- abs <$> readDataFrameOff p 2 + a21 <- abs <$> readDataFrameOff p 5 + q' <- jiter n p a10 a20 a21 + go (n - 1) p (q' * q) + jiter :: Int -> STDataFrame s t '[3,3] + -> Scalar t -> Scalar t -> Scalar t -> ST s (Quater t) + jiter n p a10 a20 a21 + | gt2 a10 a20 a21 + = jacobiEigen3Iteration 0 1 2 p + | gt2 a20 a10 a21 + = jacobiEigen3Iteration 2 0 1 p + | gt2 a21 a10 a20 + = jacobiEigen3Iteration 1 2 0 p + | otherwise + = case mod n 3 of + 0 -> jacobiEigen3Iteration 0 1 2 p + 1 -> jacobiEigen3Iteration 2 0 1 p + _ -> jacobiEigen3Iteration 1 2 0 p + gt2 :: Scalar t -> Scalar t -> Scalar t -> Bool + gt2 a b c = case compare a b of + GT -> a >= c + EQ -> a > c + LT -> False + + -- Make such a quaternion that rotates the matrix so that: + -- abs s1 >= abs s2 >= abs s3 + -- Note, the corresponding singular values may be negative, which must be + -- taken into account later. + sortQ :: Scalar t -> Scalar t -> Scalar t -> Quater t + sortQ s1 s2 s3 = sortQ' (s1 >= s2) (s1 >= s3) (s2 >= s3) + sortQ' :: Bool -> Bool -> Bool -> Quater t + sortQ' True True True = Quater 0 0 0 1 -- s1 >= s2 >= s3 + sortQ' True True False = Quater M_SQRT1_2 0 0 (-M_SQRT1_2) -- s1 >= s3 > s2 + sortQ' True False _ = Quater 0.5 0.5 0.5 0.5 -- s3 > s1 >= s2 + sortQ' False True True = Quater 0 0 M_SQRT1_2 (-M_SQRT1_2) -- s2 > s1 >= s3 + sortQ' False _ False = Quater 0 M_SQRT1_2 0 (-M_SQRT1_2) -- s3 > s2 > s1 + sortQ' False False True = Quater 0.5 0.5 0.5 (-0.5) -- s2 >= s3 > s1 + + +-- | One Givens rotation for a QR algorithm on a 3x3 matrix +-- +-- The three words arguments are indices: +-- 0 <= i /= j /= k <= 2 +-- +-- if i < j then the eigen values are already sorted! +qrDecomp3Iteration :: (Quaternion t, RealFloatExtras t) + => Int -> Int -> Int + -> STDataFrame s t '[3,3] + -> ST s (Quater t) +qrDecomp3Iteration i j k sPtr = do + sii <- readDataFrameOff sPtr ii + sij <- readDataFrameOff sPtr ij + sji <- readDataFrameOff sPtr ji + sjj <- readDataFrameOff sPtr jj + sik <- readDataFrameOff sPtr ik + sjk <- readDataFrameOff sPtr jk + -- Coefficients for a quaternion corresponding to a Givens rotation + let (ch, sh) = qrGivensQ sii sji + a = ch*ch - sh*sh + b = 2 * sh*ch + -- update the matrix + writeDataFrameOff sPtr ii $ a * sii + b * sji + writeDataFrameOff sPtr ij $ a * sij + b * sjj + writeDataFrameOff sPtr ik $ a * sik + b * sjk + writeDataFrameOff sPtr ji 0 -- a * sji - b * sii + writeDataFrameOff sPtr jj $ a * sjj - b * sij + writeDataFrameOff sPtr jk $ a * sjk - b * sik + + -- write the quaternion + qPtr <- unsafeThawDataFrame 0 + writeDataFrameOff qPtr k (negateUnless leftTriple sh) + writeDataFrameOff qPtr 3 ch + fromVec4 <$> unsafeFreezeDataFrame qPtr + where + leftTriple = (j - i) /= 1 && (k - j) /= 1 + i3 = i*3 + j3 = j*3 + ii = i3 + i + ij = i3 + j + ik = i3 + k + ji = j3 + i + jj = j3 + j + jk = j3 + k + +-- | Run QR decomposition in context of 3x3 svd: AV = US = QR +-- The input here is matrix AV +-- The R upper-triangular matrix here is in fact a diagonal matrix Sigma; +-- The Q orthogonal matrix is matrix U in the svd decomposition, +-- represented here as a quaternion. +qrDecomposition3 :: (Quaternion t, RealFloatExtras t) + => Matrix t 3 3 -> (Vector t 3, Quater t) +qrDecomposition3 m = runST $ do + mPtr <- thawDataFrame m + q1 <- qrDecomp3Iteration 0 1 2 mPtr + q2 <- qrDecomp3Iteration 0 2 1 mPtr + q3 <- qrDecomp3Iteration 1 2 0 mPtr + sig0 <- readDataFrameOff mPtr 0 + sig1 <- readDataFrameOff mPtr 4 + sig2 <- readDataFrameOff mPtr 8 + return (DF3 sig0 sig1 sig2, q3 * q2 * q1) + + +instance RealFloatExtras t => MatrixSVD t 1 1 where + svd = svd1 + +instance RealFloatExtras t => MatrixSVD t 2 2 where + svd = svd2 + +instance (RealFloatExtras t, Quaternion t) => MatrixSVD t 3 3 where + svd = svd3 + +instance {-# INCOHERENT #-} + ( RealFloatExtras t, KnownDim n, KnownDim m) + => MatrixSVD t n m where + svd a = runST $ do + D <- pure dnm + Dict <- pure $ minIsSmaller dn dm -- GHC is not convinced :( + alphas <- unsafeThawDataFrame bdAlpha + betas <- unsafeThawDataFrame bdBeta + uPtr <- unsafeThawDataFrame bdU + vPtr <- unsafeThawDataFrame bdV + + -- remove last beta if m > n + bLast <- readDataFrameOff betas nm1 + when (abs bLast > M_EPS) $ + svdGolubKahanZeroCol alphas betas vPtr nm1 + + -- main routine for a bidiagonal matrix + let maxIter = 3*nm -- number of tries + withinIters <- svdBidiagonalInplace alphas betas uPtr vPtr nm maxIter + unless withinIters . tooManyIterations + $ "SVD - Givens rotation sweeps for a bidiagonal matrix (" + ++ show maxIter ++ " sweeps max)." + + -- sort singular values + sUnsorted <- unsafeFreezeDataFrame alphas + let sSorted :: Vector (Tuple '[t, Word]) (Min n m) + sSorted = sortBy (\(S (x :! _)) (S (y :! _)) -> compare y x) + $ iwmap @_ @_ @'[] (\(Idx i :* U) (S x) -> S (abs x :! i :! U) ) sUnsorted + svdS = ewmap @t @_ @'[] (\(S (x :! _)) -> S x) sSorted + perm = ewmap @Word @_ @'[] (\(S (_ :! i :! U)) -> S i) sSorted + pCount = + if nm < 2 + then 0 :: Word + else foldl (\s (i, j) -> if perm!i > perm!j then succ s else s) + 0 [(i, j) | i <- [0..nm2w], j <- [i+1..nm2w+1]] + pPositive = even pCount + + -- alphas and svdS are now out of sync, but that is not a problem + + -- make sure det U == 1 + when ((bdUDet < 0) == pPositive) $ do + readDataFrameOff alphas 0 >>= writeDataFrameOff alphas 0 . negate + forM_ [0..n - 1] $ \i -> + readDataFrameOff uPtr (i*n) >>= writeDataFrameOff uPtr (i*n) . negate + + -- negate negative singular values + forM_ [0..nm1] $ \i -> do + s <- readDataFrameOff alphas i + when (s < 0) $ do + writeDataFrameOff alphas i $ negate s + forM_ [0..m - 1] $ \j -> + readDataFrameOff vPtr (j*m + i) + >>= writeDataFrameOff vPtr (j*m + i) . negate + + -- apply permutations if necessary + if pCount == 0 + then do + svdU <- unsafeFreezeDataFrame uPtr + svdV <- unsafeFreezeDataFrame vPtr + return SVD {..} + else do + svdU' <- unsafeFreezeDataFrame uPtr + svdV' <- unsafeFreezeDataFrame vPtr + let svdU = iwgen @_ @_ @'[] $ \(i :* Idx j :* U) -> + if j >= dimVal dnm + then index (i :* Idx j :* U) svdU' + else index (i :* Idx (unScalar $ perm!j) :* U) svdU' + svdV = iwgen @_ @_ @'[] $ \(i :* Idx j :* U) -> + if j >= dimVal dnm + then index (i :* Idx j :* U) svdV' + else index (i :* Idx (unScalar $ perm!j) :* U) svdV' + return SVD {..} + where + n = fromIntegral $ dimVal dn :: Int + m = fromIntegral $ dimVal dm :: Int + dn = dim @n + dm = dim @m + dnm = minDim dn dm + nm1 = nm - 1 + nm = fromIntegral (dimVal dnm) :: Int + nm2w = fromIntegral (max (nm - 2) 0) :: Word + -- compute the bidiagonal form b first, solve svd for b. + BiDiag {..} = bidiagonalHouseholder a + + + +{- Compute svd for a square bidiagonal matrix inplace + \( B = U S V^\intercal \) + +@ + B = | a1 b1 0 ... 0 | + | 0 a2 b2 0 ... 0 | + | 0 0 a3 b3 ... 0 | + | ................. | + | 0 0 ... an1 bn1 | + | 0 0 ... 0 an | bn? (in case if n > m) +@ + -} +svdBidiagonalInplace :: + forall (s :: Type) (t :: Type) (n :: Nat) (m :: Nat) (nm :: Nat) + . ( IterativeMethod, RealFloatExtras t + , KnownDim n, KnownDim m, KnownDim nm, nm ~ Min n m) + => STDataFrame s t '[nm] -- ^ the main diagonal of B and then the singular values. + -> STDataFrame s t '[nm] -- ^ first upper diagonal of B + -> STDataFrame s t '[n,n] -- ^ U + -> STDataFrame s t '[m,m] -- ^ V + -> Int -- ^ 0 < q <= nm -- size of a reduced matrix, such that leftover is diagonal + -> Int -- iters + -> ST s Bool -- whether the algorithm succeeds within the maxIter +svdBidiagonalInplace _ _ _ _ 0 _ = pure True +svdBidiagonalInplace _ _ _ _ 1 _ = pure True +svdBidiagonalInplace _ _ _ _ _ 0 = pure False +svdBidiagonalInplace aPtr bPtr uPtr vPtr q' iter = do + Dict <- pure $ minIsSmaller (dim @n) (dim @m) + (p, q) <- findCounters q' + if (q /= 0) + then do + findZeroDiagonal p q >>= \case + Just k + | k == q-1 -> svdGolubKahanZeroCol aPtr bPtr vPtr (k-1) + | otherwise -> svdGolubKahanZeroRow aPtr bPtr uPtr k + Nothing -> svdGolubKahanStep aPtr bPtr uPtr vPtr p q + svdBidiagonalInplace aPtr bPtr uPtr vPtr q (iter - 1) + else return True + where + -- nm = fromIntegral $ dimVal' @nm :: Int + + -- Check if off-diagonal elements are close to zero and nullify them + -- if they are small along the way. + -- And find such indices p and q that satisfy condition in alg. 8.6.2 + -- on p. 492. of "Matrix Computations " (4-th edition). + -- Except these are inverted: + -- p -- is the starting index of B22 (last submatrix with non-zero superdiagonal) + -- q -- is the starting index of B33 (diagonal submatrix) + -- + -- that is, p and q determine the index and the size of next work piece. + findCounters :: Int -> ST s (Int, Int) + findCounters = goQ + where + checkEps :: Int -> ST s Bool + checkEps k = do + b <- abs <$> readDataFrameOff bPtr (k-1) + if b == 0 + then return True + else do + a1 <- abs <$> readDataFrameOff aPtr (k-1) + a2 <- abs <$> readDataFrameOff aPtr k + if b <= M_EPS * (max (a1 + a2) 1) + then True <$ writeDataFrameOff bPtr (k-1) 0 + else return False + goQ :: Int -> ST s (Int, Int) + goQ 0 = pure (0, 0) -- guard against calling with q == 0 + goQ 1 = pure (0, 0) -- 1x1 matrix is always diagonal + goQ k = checkEps (k-1) >>= \case + True -> goQ (k-1) + False -> flip (,) k <$> goP (k-2) + goP :: Int -> ST s Int + goP 0 = pure 0 + goP k = checkEps k >>= \case + True -> return k + False -> goP (k-1) + + -- For indices p and q (p < q), find the biggest index (< q) such that + -- a[p] == 0 + findZeroDiagonal :: Int -> Int -> ST s (Maybe Int) + findZeroDiagonal p q + | k < p = pure Nothing + | otherwise = do + ak <- readDataFrameOff aPtr k + if ak == 0 + then pure $ Just k + else if abs ak <= M_EPS + then Just k <$ writeDataFrameOff aPtr k 0 + else findZeroDiagonal p k + where + k = q - 1 + + +-- | Apply a series of column transformations to make b[k] (and whole column k+1) zero +-- (page 491, 1st paragraph) when a[k+1] == 0. +-- To make this element zero, I apply a series of Givens transforms on columns +-- (multiply on the right). +-- +-- Prerequisites: +-- * a[k+1] == 0 +-- * 0 <= k < min n (m-1) +-- Invariants: +-- * matrix \(B :: n \times m \) is bidiagonal, represented by two diagonals; +-- * matrix V is orthogonal +-- * \( B = A V^\intercal \), where \(A :: n \times m\) is an implicit original matrix +-- Results: +-- * Same bidiagonal matrix with b[k] == 0; i.e. (k+1)-th column is zero. +-- * matrix V is updated (multiplied on the right) +-- +-- NB: All changes are made inplace. +-- +svdGolubKahanZeroCol :: + forall (s :: Type) (t :: Type) (n :: Nat) (m :: Nat) + . (RealFloatExtras t, KnownDim n, KnownDim m, n <= m) + => STDataFrame s t '[n] -- ^ the main diagonal of \(B\) + -> STDataFrame s t '[n] -- ^ first upper diagonal of \(B\) + -> STDataFrame s t '[m,m] -- ^ \(V\) + -> Int -- ^ 0 <= k < min (n+1) m + -> ST s () +svdGolubKahanZeroCol aPtr bPtr vPtr k + | k < 0 || k >= lim = error $ unwords + [ "svdGolubKahanZeroCol: k =", show k + , "is outside of a valid range 0 <= k <", show lim] + -- this trick is to convince GHC that constraint (n <= m) is not redundant + | Dict <- Dict @(n <= m) = do + b <- readDataFrameOff bPtr k + writeDataFrameOff bPtr k 0 + foldM_ goGivens b [k, k-1 .. 0] + where + n = fromIntegral $ dimVal' @n :: Int + m = fromIntegral $ dimVal' @m :: Int + lim = min n (m-1) + goGivens :: Scalar t -> Int -> ST s (Scalar t) + goGivens 0 _ = return 0 -- non-diagonal element is nullified prematurely + goGivens b i = do + ai <- readDataFrameOff aPtr i + let rab = recip $ hypot b ai + c = ai*rab + s = b *rab + updateGivensMat vPtr i (k+1) c s + writeDataFrameOff aPtr i $ ai*c + b*s -- B[i,i] + if i == 0 + then return 0 + else do + bi1 <- readDataFrameOff bPtr (i - 1) -- B[i,i-1] + writeDataFrameOff bPtr (i - 1) $ bi1 * c + return $ negate (bi1 * s) + +-- | Apply a series of row transformations to make b[k] (and whole column k) zero +-- (page 490, last paragraph) when a[k] == 0. +-- To make this element zero, I apply a series of Givens transforms on rows +-- (multiply on the left). +-- +-- Prerequisites: +-- * a[k] == 0 +-- * 0 <= k < n - 1 +-- Invariants: +-- * matrix \(B :: m \times n \) is bidiagonal, represented by two diagonals; +-- * matrix U is orthogonal +-- * \( B = U A \), where \(A :: m \times n \) is an implicit original matrix +-- Results: +-- * Same bidiagonal matrix with b[k] == 0; i.e. k-th column is zero. +-- * matrix U is updated (multiplied on the right) +-- +-- NB: All changes are made inplace. +-- +svdGolubKahanZeroRow :: + forall (s :: Type) (t :: Type) (n :: Nat) (m :: Nat) + . (RealFloatExtras t, KnownDim n, KnownDim m, n <= m) + => STDataFrame s t '[n] -- ^ the main diagonal of B + -> STDataFrame s t '[n] -- ^ first upper diagonal of B + -> STDataFrame s t '[m,m] -- ^ U + -> Int -- ^ 0 <= k < n - 1 + -> ST s () +svdGolubKahanZeroRow aPtr bPtr uPtr k + | k < 0 || k >= n1 = error $ unwords + [ "svdGolubKahanZeroRow: k =", show k + , "is outside of a valid range 0 <= k <", show n1] + -- this trick is to convince GHC that constraint (n <= m) is not redundant + | Dict <- Dict @(n <= m) = do + b <- readDataFrameOff bPtr k + writeDataFrameOff bPtr k 0 + foldM_ goGivens b [k+1..n1] + where + n = fromIntegral $ dimVal' @n :: Int + n1 = n - 1 + goGivens :: Scalar t -> Int -> ST s (Scalar t) + goGivens 0 _ = return 0 -- non-diagonal element is nullified prematurely + goGivens b j = do + aj <- readDataFrameOff aPtr j + bj <- readDataFrameOff bPtr j + let rab = recip $ hypot b aj + c = aj*rab + s = b*rab + updateGivensMat uPtr k j c (negate s) + writeDataFrameOff aPtr j $ b*s + aj*c + writeDataFrameOff bPtr j $ bj*c + return $ negate (bj * s) + +-- | A Golub-Kahan bidiagonal SVD step on an unreduced matrix +svdGolubKahanStep :: + forall (s :: Type) (t :: Type) (n :: Nat) (m :: Nat) (nm :: Nat) + . ( RealFloatExtras t + , KnownDim n, KnownDim m, KnownDim nm, nm ~ Min n m) + => STDataFrame s t '[nm] -- ^ the main diagonal of B and then the singular values. + -> STDataFrame s t '[nm] -- ^ first upper diagonal of B + -> STDataFrame s t '[n,n] -- ^ U + -> STDataFrame s t '[m,m] -- ^ V + -> Int -- ^ p : 0 <= p < q <= nm; p <= q - 2 + -> Int -- ^ q : 0 <= p < q <= nm; p <= q - 2 + -> ST s () +svdGolubKahanStep aPtr bPtr uPtr vPtr p q + | p > q - 2 || p < 0 || q > nm + = error $ unwords + [ "svdGolubKahanStep: p =", show p, "and q =", show q + , "do not satisfy p <= q - 2 or 0 <= p < q <=", show nm] + | Dict <- Dict @(nm ~ Min n m) = do + (y,z) <- getWilkinsonShiftYZ + goGivens2 y z p + where + nm = fromIntegral $ dimVal' @nm :: Int + + -- get initial values for one recursion sweep. + -- Note, input must satisfy: q >= p+2 + getWilkinsonShiftYZ :: ST s (Scalar t, Scalar t) + getWilkinsonShiftYZ = do + a1 <- readDataFrameOff aPtr p + b1 <- readDataFrameOff bPtr p + am <- readDataFrameOff aPtr (q-2) + an <- readDataFrameOff aPtr (q-1) + bm <- if q >= p + 3 + then readDataFrameOff bPtr (q-3) + else pure 0 + bn <- readDataFrameOff bPtr (q-2) + let t11 = a1*a1 + t12 = a1*b1 + tmm = am*am + bm*bm + tnn = an*an + bn*bn + tnm = am*bn + d = 0.5*(tmm - tnn) + mu = tnn + d - negateUnless (d >= 0) (hypot d tnm) + return (t11 - mu, t12) + + -- yv = b[k-1]; zv = B[k-1,k+1] -- to be eliminated by 1st Givens r + -- yu = a[k]; zu = B[k+1,k-1] -- to be eliminated by 2nd Givens r + goGivens2 :: Scalar t -> Scalar t -> Int -> ST s () + goGivens2 yv zv k = do + a1 <- readDataFrameOff aPtr k -- B[k,k] + a2 <- readDataFrameOff aPtr (k+1) -- B[k+1,k+1] + b1 <- readDataFrameOff bPtr k -- B[k,k+1] + let a1' = a1*cv + b1*sv -- B[k,k] == yu + a2' = a2*cv -- B[k+1,k+1] + b0' = yv*cv + zv*sv -- B[k-1,k] + b1' = b1*cv - a1*sv -- B[k,k+1] + yu = a1' -- B[k,k] + zu = a2*sv -- B[k+1,k] + ryzu = recip $ hypot yu zu + cu = yu * ryzu + su = zu * ryzu + a1'' = yu *cu + zu *su + a2'' = a2'*cu - b1'*su + b1'' = b1'*cu + a2'*su + updateGivensMat vPtr k (k+1) cv sv + updateGivensMat uPtr k (k+1) cu su + + when (k > p) $ writeDataFrameOff bPtr (k-1) b0' + writeDataFrameOff bPtr k b1'' + writeDataFrameOff aPtr k a1'' + writeDataFrameOff aPtr (k+1) a2'' + when (k < q - 2) $ do + b2 <- readDataFrameOff bPtr (k+1) -- B[k+1,k+2] + let b2'' = b2*cu + zvn = b2*su + writeDataFrameOff bPtr (k+1) b2'' + goGivens2 b1'' zvn (k+1) + where + ryzv = recip $ hypot yv zv + cv = yv * ryzv + sv = zv * ryzv + +-- | Update a transformation matrix with a Givens transform (on the right) +updateGivensMat :: + forall (s :: Type) (t :: Type) (n :: Nat) + . (PrimBytes t, Num t, KnownDim n) + => STDataFrame s t '[n,n] + -> Int -> Int + -> Scalar t -> Scalar t -> ST s () +updateGivensMat p i j c s = forM_ [0..n-1] $ \k -> do + let nk = n*k + ioff = nk + i + joff = nk + j + uki <- readDataFrameOff p ioff + ukj <- readDataFrameOff p joff + writeDataFrameOff p ioff $ uki*c + ukj*s + writeDataFrameOff p joff $ ukj*c - uki*s + where + n = fromIntegral $ dimVal' @n :: Int + + +minIsSmaller :: forall (n :: Nat) (m :: Nat) + . Dim n -> Dim m -> Dict (Min n m <= n, Min n m <= m) +minIsSmaller dn dm + | Just Dict <- lessOrEqDim dnm dn + , Just Dict <- lessOrEqDim dnm dm + = Dict + | otherwise + = error "minIsSmaller: impossible type-level comparison" + where + dnm = minDim dn dm diff --git a/src/Numeric/PrimBytes.hs b/src/Numeric/PrimBytes.hs --- a/src/Numeric/PrimBytes.hs +++ b/src/Numeric/PrimBytes.hs @@ -90,7 +90,6 @@ import GHC.Generics import GHC.Int import GHC.IO (IO (..)) -import GHC.Ptr (Ptr (..)) import GHC.Word import Numeric.Dimensions import qualified Numeric.Tuple.Lazy as TL @@ -1328,7 +1327,7 @@ anyList :: forall (k :: Type) (xs :: [k]) . RepresentableList xs => [Any] -anyList = unsafeCoerce# (tList @_ @xs) +anyList = unsafeCoerce# (tList @xs) {-# INLINE anyList #-} instance RepresentableList xs => PrimBytes (Idxs (xs :: [k])) where @@ -1368,7 +1367,7 @@ go i (W# x :xs) s = go (plusAddr# i SIZEOF_HSWORD#) xs (writeWordOffAddr# i 0# x s) {-# INLINE writeAddr #-} - byteSize _ = case dimVal (order' @_ @xs) of + byteSize _ = case dimVal (order' @xs) of W# n -> byteSize (undefined :: Idx x) *# word2Int# n {-# INLINE byteSize #-} byteAlign _ = byteAlign (undefined :: Idx x) @@ -1376,14 +1375,14 @@ byteFieldOffset _ _ = negateInt# 1# {-# INLINE byteFieldOffset #-} indexArray ba off - | n@(W# n#) <- dimVal (order' @_ @xs) + | n@(W# n#) <- dimVal (order' @xs) = unsafeCoerce# (go (off *# word2Int# n#) n) where go _ 0 = [] go i n = W# (indexWordArray# ba i) : go (i +# 1#) (n-1) {-# INLINE indexArray #-} readArray mba off s - | n@(W# n#) <- dimVal (order' @_ @xs) + | n@(W# n#) <- dimVal (order' @xs) = unsafeCoerce# (go (off *# word2Int# n#) n s) where go _ 0 s0 = (# s0, [] #) @@ -1393,7 +1392,7 @@ = (# s2, W# w : xs #) {-# INLINE readArray #-} writeArray mba off is - | W# n# <- dimVal (order' @_ @xs) + | W# n# <- dimVal (order' @xs) = go (off *# word2Int# n#) (listIdxs is) where go _ [] s = s @@ -1437,7 +1436,7 @@ , L.All PrimBytes xs ) => PrimBytes (TS.Tuple xs) where type PrimFields (TS.Tuple xs) = TupleFields 1 xs - fromBytes off ba = go 0# (tList @_ @xs) + fromBytes off ba = go 0# (tList @xs) where go :: L.All PrimBytes ds => Int# -> TypeList ds -> TS.Tuple ds @@ -1447,7 +1446,7 @@ , n' <- roundUpInt n (byteAlign x) = TS.Id (fromBytes (off +# n') ba) :* go (n' +# byteSize x) ts {-# INLINE fromBytes #-} - readBytes mb off = go mb 0# (tList @_ @xs) + readBytes mb off = go mb 0# (tList @xs) where go :: L.All PrimBytes ds => MutableByteArray# s @@ -1469,7 +1468,7 @@ | n' <- roundUpInt n (byteAlign x) = go mb (n' +# byteSize x) xs ts (writeBytes mb (off +# n') x s) {-# INLINE writeBytes #-} - readAddr addr = go 0# (tList @_ @xs) + readAddr addr = go 0# (tList @xs) where go :: L.All PrimBytes ds => Int# -> TypeList ds -> State# s -> (# State# s, TS.Tuple ds #) @@ -1490,7 +1489,7 @@ | n' <- roundUpInt n (byteAlign x) = go (n' +# byteSize x) xs ts (writeAddr x (plusAddr# addr n') s) {-# INLINE writeAddr #-} - byteSize _ = go 0# 1# (tList @_ @xs) + byteSize _ = go 0# 1# (tList @xs) where go :: L.All PrimBytes ys => Int# -> Int# -> TypeList ys -> Int# go s a Empty = s `roundUpInt` a @@ -1499,7 +1498,7 @@ in go ( roundUpInt s xa +# byteSize x) ( maxInt a xa ) ps {-# INLINE byteSize #-} - byteAlign _ = go (tList @_ @xs) + byteAlign _ = go (tList @xs) where go :: L.All PrimBytes ys => TypeList ys -> Int# go Empty = 0# @@ -1507,7 +1506,7 @@ {-# INLINE byteAlign #-} byteFieldOffset name _ | Just n <- readMaybe $ symbolVal' name - = go (n-1) 0# (tList @_ @xs) + = go (n-1) 0# (tList @xs) | otherwise = negateInt# 1# where go :: L.All PrimBytes ys => Word -> Int# -> TypeList ys -> Int# diff --git a/src/Numeric/ProductOrd/NonTransitive.hs b/src/Numeric/ProductOrd/NonTransitive.hs --- a/src/Numeric/ProductOrd/NonTransitive.hs +++ b/src/Numeric/ProductOrd/NonTransitive.hs @@ -35,8 +35,6 @@ import Data.Data import Data.Foldable import Data.Kind (Type) -import Data.Monoid as Mon (Monoid (..)) -import Data.Semigroup as Sem (Semigroup (..)) import Foreign.Storable (Storable) import GHC.Generics diff --git a/src/Numeric/ProductOrd/Partial.hs b/src/Numeric/ProductOrd/Partial.hs --- a/src/Numeric/ProductOrd/Partial.hs +++ b/src/Numeric/ProductOrd/Partial.hs @@ -35,8 +35,6 @@ import Data.Data import Data.Foldable import Data.Kind (Type) -import Data.Monoid as Mon (Monoid (..)) -import Data.Semigroup as Sem (Semigroup (..)) import Foreign.Storable (Storable) import GHC.Generics diff --git a/src/Numeric/Quaternion/Internal.hs b/src/Numeric/Quaternion/Internal.hs --- a/src/Numeric/Quaternion/Internal.hs +++ b/src/Numeric/Quaternion/Internal.hs @@ -1,4 +1,5 @@ {-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE MagicHash #-} {-# LANGUAGE PatternSynonyms #-} {-# LANGUAGE ScopedTypeVariables #-} @@ -9,7 +10,9 @@ ( Quaternion (..), Quater(Quater) ) where +import Numeric.DataFrame.Type (KnownBackend) import Numeric.Matrix.Internal (Matrix) +import Numeric.PrimBytes (PrimBytes) import Numeric.Vector.Internal (Vector) import Text.Read @@ -21,7 +24,10 @@ {-# COMPLETE Quater #-} -- | Quaternion operations -class Quaternion t where +class ( Floating (Quater t), Floating t, Ord t, PrimBytes t + , KnownBackend t '[3], KnownBackend t '[4] + , KnownBackend t '[3,3], KnownBackend t '[4, 4] + ) => Quaternion t where -- | Quaternion data type. The ordering of coordinates is @(x,y,z,w)@, -- where @w@ is the argument, and @x y z@ are the components of a 3D vector data Quater t diff --git a/src/Numeric/Quaternion/Internal/QDouble.hs b/src/Numeric/Quaternion/Internal/QDouble.hs --- a/src/Numeric/Quaternion/Internal/QDouble.hs +++ b/src/Numeric/Quaternion/Internal/QDouble.hs @@ -10,19 +10,18 @@ {-# LANGUAGE UnboxedTuples #-} {-# LANGUAGE UndecidableInstances #-} {-# LANGUAGE ViewPatterns #-} -{-# OPTIONS_GHC -fno-warn-orphans #-} +{-# OPTIONS_GHC -fno-warn-orphans #-} module Numeric.Quaternion.Internal.QDouble ( QDouble, Quater (..) ) where import qualified Control.Monad.ST as ST import Data.Coerce (coerce) +import Numeric.Basics import Numeric.DataFrame.Internal.PrimArray import qualified Numeric.DataFrame.ST as ST import Numeric.DataFrame.Type -import Numeric.PrimBytes (PrimBytes) import Numeric.Quaternion.Internal -import Numeric.Scalar.Internal import Numeric.Vector.Internal type QDouble = Quater Double @@ -31,11 +30,11 @@ deriving instance PrimArray Double (Quater Double) instance Quaternion Double where - newtype Quater Double = QDouble Vec4d + newtype Quater Double = QDouble (Vector Double 4) {-# INLINE packQ #-} - packQ = coerce (vec4 :: Double -> Double -> Double -> Double -> Vec4d) + packQ = coerce (vec4 :: Double -> Double -> Double -> Double -> Vector Double 4) {-# INLINE unpackQ# #-} - unpackQ# = coerce (unpackV4# :: Vec4d -> (# Double, Double, Double, Double #)) + unpackQ# = coerce (unpackV4# :: Vector Double 4 -> (# Double, Double, Double, Double #)) {-# INLINE fromVecNum #-} fromVecNum (unpackV3# -> (# x, y, z #)) = packQ x y z {-# INLINE fromVec4 #-} @@ -47,7 +46,7 @@ {-# INLINE im #-} im (unpackQ# -> (# x, y, z, _ #)) = packQ x y z 0.0 {-# INLINE re #-} - re (unpackQ# -> (# _, _, _, w #)) = packQ 0.0 0.0 0.0 w + re (unpackQ# -> (# _, _, _, w #)) = packQ 0 0 0 w {-# INLINE imVec #-} imVec (unpackQ# -> (# x, y, z, _ #)) = vec3 x y z {-# INLINE taker #-} @@ -72,7 +71,7 @@ (l*z + d*k + t2 * (y*i - x*j)) {-# INLINE getRotScale #-} getRotScale a b = case (# unpackV3# a, unpackV3# b #) of - (# _, (# 0, 0, 0 #) #) -> packQ 0.0 0.0 0.0 0.0 + (# _, (# 0, 0, 0 #) #) -> packQ 0 0 0 0 (# (# 0, 0, 0 #), _ #) -> let x = (1 / 0 :: Double) in packQ x x x x (# (# a1, a2, a3 #), (# b1, b2, b3 #) #) -> let ma = sqrt (a1*a1 + a2*a2 + a3*a3) @@ -99,7 +98,7 @@ (# t1, t2, t3 #) -> packQ (t1 * r) (t2 * r) (t3 * r) (c / ma2) {-# INLINE axisRotation #-} axisRotation v a = case unpackV3# v of - (# 0, 0, 0 #) -> packQ 0.0 0.0 0.0 (if abs a < pi then 1 else -1) + (# 0, 0, 0 #) -> packQ 0 0 0 (negateUnless (abs a < M_PI) 1) (# x, y, z #) -> let c = cos (a * 0.5) s = sin (a * 0.5) @@ -261,39 +260,35 @@ = QDouble (a - b) {-# INLINE (-) #-} (unpackQ# -> (# a1, a2, a3, a4 #)) * (unpackQ# -> (# b1, b2, b3, b4 #)) - = packQ - ((a4 * b1) + - (a1 * b4) + - (a2 * b3) - - (a3 * b2) - ) - ((a4 * b2) - - (a1 * b3) + - (a2 * b4) + - (a3 * b1) - ) - ((a4 * b3) + - (a1 * b2) - - (a2 * b1) + - (a3 * b4) - ) - ((a4 * b4) - - (a1 * b1) - - (a2 * b2) - - (a3 * b3) - ) + = packQ ((a4 * b1) + (a1 * b4) + (a2 * b3) - (a3 * b2)) + ((a4 * b2) - (a1 * b3) + (a2 * b4) + (a3 * b1)) + ((a4 * b3) + (a1 * b2) - (a2 * b1) + (a3 * b4)) + ((a4 * b4) - (a1 * b1) - (a2 * b2) - (a3 * b3)) {-# INLINE (*) #-} negate (QDouble a) = QDouble (negate a) {-# INLINE negate #-} - abs = packQ 0.0 0.0 0.0 . sqrt . square + abs = packQ 0 0 0 . sqrt . square {-# INLINE abs #-} signum q@(unpackQ# -> (# x, y, z, w #)) - = case square q of - 0.0 -> packQ 0.0 0.0 0.0 0.0 - qd -> case recip (sqrt qd) of - s -> packQ (x * s) (y * s) (z * s) (w * s) + | qd == 0 = q + | otherwise = case ix + iy + iz + iw + nn of + 0 -> packQ (x * l) (y * l) (z * l) (w * l) + 1 -> packQ (copysign 1 x) 0 0 0 + 2 -> packQ 0 (copysign 1 y) 0 0 + 4 -> packQ 0 0 (copysign 1 z) 0 + 8 -> packQ 0 0 0 (copysign 1 w) + _ -> packQ n n n n + where + n = 0 / 0 :: Double + qd = x*x + y*y + z*z + w*w + ix = if isInfinite x then 1 else 0 :: Int + iy = if isInfinite y then 2 else 0 :: Int + iz = if isInfinite z then 4 else 0 :: Int + iw = if isInfinite w then 8 else 0 :: Int + nn = if isNaN x || isNaN y || isNaN z || isNaN w then 16 else 0 :: Int + l = recip (sqrt qd) {-# INLINE signum #-} - fromInteger = packQ 0.0 0.0 0.0 . fromInteger + fromInteger = packQ 0 0 0 . fromInteger {-# INLINE fromInteger #-} @@ -304,172 +299,253 @@ {-# INLINE (/) #-} a / b = a * recip b {-# INLINE fromRational #-} - fromRational = packQ 0.0 0.0 0.0 . fromRational + fromRational = packQ 0 0 0 . fromRational instance Floating QDouble where {-# INLINE pi #-} - pi = packQ 0.0 0.0 0.0 3.141592653589793 + pi = packQ 0 0 0 M_PI {-# INLINE exp #-} exp (unpackQ# -> (# x, y, z, w #)) - = case (# (x * x) + - (y * y) + - (z * z) - , exp w - #) of - (# 0.0, et #) -> packQ 0.0 0.0 0.0 et - (# mv2, et #) -> case sqrt mv2 of - mv -> case et * sin mv / mv of - l -> packQ (x * l) (y * l) (z * l) (et * cos mv) + | mv2 == 0 = packQ x y z ew + | otherwise = packQ (x * l) (y * l) (z * l) arg + where + mv2 = (x * x) + (y * y) + (z * z) + mv = sqrt mv2 + ew = exp w + l = ew * sin mv / mv + arg = ew * cos mv {-# INLINE log #-} log = log' (Vec3 1 0 0) {-# INLINE sqrt #-} sqrt = sqrt' (Vec3 1 0 0) {-# INLINE sin #-} sin (unpackQ# -> (# x, y, z, w #)) - = case (x * x) + (y * y) + (z * z) of - 0.0 -> packQ 0.0 0.0 0.0 (sin w) - mv2 -> case sqrt mv2 of - mv -> case cos w * sinh mv / mv of - l -> packQ (x * l) (y * l) (z * l) (sin w * cosh mv) + | mv2 == 0 = packQ x y z (sin w) + | otherwise = packQ (x * l) (y * l) (z * l) arg + where + mv2 = (x * x) + (y * y) + (z * z) + mv = sqrt mv2 + l = cos w * sinh mv / mv + arg = sin w * cosh mv {-# INLINE cos #-} cos (unpackQ# -> (# x, y, z, w #)) - = case (x * x) + (y * y) + (z * z) of - 0.0 -> packQ 0.0 0.0 0.0 (cos w) - mv2 -> case sqrt mv2 of - mv -> case sin w * sinh mv / negate mv of - l -> packQ (x * l) (y * l) (z * l) (cos w * cosh mv) + | mv2 == 0 = packQ x y z (cos w) + | otherwise = packQ (x * l) (y * l) (z * l) arg + where + mv2 = (x * x) + (y * y) + (z * z) + mv = sqrt mv2 + l = sin w * sinh mv / negate mv + arg = cos w * cosh mv {-# INLINE tan #-} tan (unpackQ# -> (# x, y, z, w #)) - = case (x * x) + (y * y) + (z * z) of - 0.0 -> packQ 0.0 0.0 0.0 (tan w) - mv2 -> - let mv = sqrt mv2 - chv = cosh mv - shv = sinh mv - ct = cos w - st = sin w - cq = recip ( (ct * ct * chv * chv) + (st * st * shv * shv) ) - l = chv * shv * cq / mv - in packQ (x * l) (y * l) (z * l) (ct * st * cq) + | mv2 == 0 = packQ x y z (tan w) + | isInfinite mv2 = signum (packQ x y z 0) + | otherwise = packQ (x * l) (y * l) (z * l) arg + where + mv2 = (x * x) + (y * y) + (z * z) + mv = sqrt mv2 + b = 2*mv + a = 2*w + sina = sin a + eb = exp (-b) + eb2 = eb*eb + d = 1 + eb2 + 2 * eb * cos a + rd = recip d + pa = M_PI - abs a + rd' = 2 / (b*b + pa*pa) + (l, arg) = + if d >= M_EPS + then ((1 - eb2) * rd / mv, 2 * sina * eb * rd) + else (2 * rd' , negate sina * rd') {-# INLINE sinh #-} sinh (unpackQ# -> (# x, y, z, w #)) - = case (x * x) + (y * y) + (z * z) of - 0.0 -> packQ 0.0 0.0 0.0 (sinh w) - mv2 -> case sqrt mv2 of - mv -> case cosh w * sin mv / mv of - l -> packQ (x * l) (y * l) (z * l) (sinh w * cos mv) + | mv2 == 0 = packQ x y z (sinh w) + | otherwise = packQ (x * l) (y * l) (z * l) arg + where + mv2 = (x * x) + (y * y) + (z * z) + mv = sqrt mv2 + l = cosh w * sin mv / mv + arg = sinh w * cos mv {-# INLINE cosh #-} cosh (unpackQ# -> (# x, y, z, w #)) - = case (x * x) + (y * y) + (z * z) of - 0.0 -> packQ 0.0 0.0 0.0 (cosh w) - mv2 -> case sqrt mv2 of - mv -> case sinh w * sin mv / mv of - l -> packQ (x * l) (y * l) (z * l) (cosh w * cos mv) + | mv2 == 0 = packQ x y z (cosh w) + | otherwise = packQ (x * l) (y * l) (z * l) arg + where + mv2 = (x * x) + (y * y) + (z * z) + mv = sqrt mv2 + l = sinh w * sin mv / mv + arg = cosh w * cos mv {-# INLINE tanh #-} tanh (unpackQ# -> (# x, y, z, w #)) - = case (x * x) + (y * y) + (z * z) of - 0.0 -> packQ 0.0 0.0 0.0 (tanh w) - mv2 -> - let mv = sqrt mv2 - cv = cos mv - sv = sin mv - cht = cosh w - sht = sinh w - cq = recip ( (cht * cht * cv * cv) + (sht * sht * sv * sv) ) - l = cv * sv * cq / mv - in packQ (x * l) (y * l) (z * l) (cht * sht * cq) + | mv2 == 0 = packQ x y z (tanh w) + | isInfinite mv2 = packQ 0 0 0 (signum w) + | otherwise = packQ (x * l) (y * l) (z * l) arg + where + mv2 = (x * x) + (y * y) + (z * z) + mv = sqrt mv2 + b = 2*w + a = 2*mv + eb = exp (- abs b) + eb2 = eb*eb + d = 1 + eb2 + 2 * eb * cos a + rd = recip d + pa = M_PI - a + rd' = 2 / (b*b + pa*pa) + (l, arg) = + if d >= M_EPS + then (2 * sin a * eb * rd / mv, copysign (1 - eb2) b * rd) + else (2 * rd' , b * rd') {-# INLINE asin #-} - asin q = -i * log' axis (i*q + sqrt' axis (1 - q*q)) - where - axis = sigVec q - i = fromVecNum axis 0 + -- The original formula: + -- asin q = -i * log (i*q + sqrt (1 - q*q)) + -- below is a more numerically stable version. + asin (unpackQ# -> (# x, y, z, w #)) + | v2 == 0 = if w2 <= 1 + then packQ x y z (asin w) + else packQ l 0 0 arg + | otherwise = packQ (x*c) (y*c) (z*c) arg + where + v2 = (x * x) + (y * y) + (z * z) + v = sqrt v2 + w2 = w*w + w1qq = 0.5 *(1 - w2 + v2) -- real part of (1 - q*q)/2 + l1qq = sqrt (w1qq*w1qq + w2*v2) -- length of (1 - q*q)/2 + sp2 = l1qq + w1qq + sn2 = l1qq - w1qq + sp = sqrt sp2 + sn = copysign (sqrt sn2) w + -- choose a more stable (symbolically equiv) version + dp = if 2 * v2 <= sp2 then sp - v else v2 / ((sp + v)*(sn2 + v2)) + dn = if 2 * w2 <= sn2 then w - sn else w2 / ((sn + w)*(sp2 + w2)) + (wD, vD) = case compare w1qq 0 of + GT -> (dp, w * dp / sp) + LT -> (v * dn / sn, dn) + EQ -> (-v, w) + l = -0.5 * log (wD*wD + vD*vD) + c = l / v + arg = atan2 vD wD {-# INLINE acos #-} - acos q = pi/2 - asin q + acos q = M_PI_2 - asin q {-# INLINE atan #-} -- atan q = i / 2 * log ( (i + q) / (i - q) ) atan (unpackQ# -> (# x, y, z, w #)) - = case (x * x) + (y * y) + (z * z) of - 0 -> packQ 0 0 0 (atan w) - v2 -> - let v = sqrt v2 - w2 = w*w - vw2 = w2 + v2 - t = 1 + vw2 - c = 0.25 * log ((t + 2*v) / (t - 2*v)) / v - in packQ (c*x) (c*y) (c*z) $ 0.5 * atan2 (2*w) (1 - vw2) + | v2 == 0 = packQ x y z (atan w) + | otherwise = packQ (x*c) (y*c) (z*c) arg + where + v2 = (x * x) + (y * y) + (z * z) + v = sqrt v2 + w2 = w*w + q2 = w2 + v2 + v' = v - 1 + mzero = w2 + v'*v' + (c, arg) = + if mzero == 0 + then ( sqrt maxFinite / v, 0) + else ( 0.25 * (log (1 + q2 + 2*v) - log mzero) / v + , 0.5 * atan2 (2*w) (1 - q2) ) {-# INLINE asinh #-} -- The original formula: -- asinh q = log (q + sqrt (q*q + 1)) -- below is a more numerically stable version. asinh (unpackQ# -> (# x, y, z, w #)) - | v2 == 0 - = packQ 0 0 0 (asinh w) - | w == 0 - = if v2 <= 1 - then let c = asin v / v - in packQ (c*x) (c*y) (c*z) 0 - else let c = 0.5 * pi / v - w' = 0.5 * log (2*v2 - 1 + 2 * v * sqrt ( v2 - 1 )) - in packQ (c*x) (c*y) (c*z) w' - | otherwise - = let c = atan2 v t / v - w' = 0.5 * log ( t2 + v2) - log t - + if w >= 0 - then log ( t + w ) - else log 2 + log w2 - log ( t - w ) - log ( r + w2 + v2 - 1 ) - in packQ (c*x) (c*y) (c*z) w' + | v2 == 0 = packQ x y z (asinh w) + | otherwise = packQ (x*c) (y*c) (z*c) arg where v2 = (x * x) + (y * y) + (z * z) - v = sqrt v2 + v = sqrt v2 w2 = w*w - wvp1 = w2 - v2 + 1 - r = sqrt $ wvp1*wvp1 + 4*w2*v2 - t = sqrt $ ( r + wvp1 ) * 0.5 - t2 = t*t + w1qq = 0.5 *(1 + w2 - v2) -- real part of (1 + q*q)/2 + l1qq = sqrt (w1qq*w1qq + w2*v2) -- length of (1 + q*q)/2 + sp2 = l1qq + w1qq + sn2 = l1qq - w1qq + sp = sqrt sp2 + sn = copysign (sqrt sn2) w + -- choose a more stable (symbolically equiv) version + dp = if 2 * w >= - sp then w + sp else w2 / ((sp - w)*(w2 + sn2)) + dn = if 2 * v <= - sn || sn >= 0 + then v + sn else v2 / ((v - sn)*(v2 + sp2)) + (wD, vD) = case compare w1qq 0 of + GT -> (dp, v * dp / sp) + LT -> (w * dn / sn, dn) + EQ -> (w, v) + c = atan2 vD wD / v + arg = 0.5 * log (wD*wD + vD*vD) {-# INLINE acosh #-} + -- The original formula: + -- asinh q = log (q + sqrt (q + 1) * sqrt (q - 1)) + -- below is a more numerically stable version. -- note, log (q + sqrt (q*q - 1)) would not work, because that would not -- be the principal value. - acosh q = log' axis (q + sqrt' axis (q + 1) * sqrt' axis (q - 1)) + acosh (unpackQ# -> (# x, y, z, w #)) + | v2 == 0 = packQ x y z (acosh w) + | otherwise = packQ (x*c) (y*c) (z*c) arg where - axis = sigVec q + v2 = (x * x) + (y * y) + (z * z) + v = sqrt v2 + w2 = w*w + w1qq = 0.5 *(w2 - v2 - 1) -- real part of (q*q - 1)/2 + l1qq = sqrt (w1qq*w1qq + w2*v2) -- length of (q*q - 1)/2 + sp2 = l1qq + w1qq + sn2 = l1qq - w1qq + sp = sqrt sp2 + sn = copysign (sqrt sn2) w + -- choose a more stable (symbolically equiv) version + dp = if 2 * w >= - sp then w + sp else w2 / ((w - sp)*(w2 + sn2)) + dn = if 2 * v <= - sn || sn >= 0 + then v + sn else v2 / ((sn - v)*(v2 + sp2)) + (wD, vD) = case compare w1qq 0 of + GT -> (dp, v * dp / sp) + LT -> (w * dn / sn, dn) + EQ -> (w, v) + c = atan2 vD wD / v + arg = 0.5 * log (wD*wD + vD*vD) {-# INLINE atanh #-} - atanh q = 0.5 * log' (sigVec q) ((1+q)/(1-q)) + -- atanh q = 0.5 * log ( (1 + q) / (1 - q) ) + atanh (unpackQ# -> (# x, y, z, w #)) + | v2 == 0 = packQ x y z (atanh w) + | otherwise = packQ (x*c) (y*c) (z*c) (copysign arg w) + where + v2 = (x * x) + (y * y) + (z * z) + v = sqrt v2 + w2 = w*w + q2 = w2 + v2 + w' = abs w - 1 + c = 0.5 * atan2 (2*v) (1 - q2) / v + arg = if w' == 0 + then (1/0) + else 0.25 * (log (1 + q2 + 2 * abs w) - log (v2 + w'*w')) + -- If q is negative real, provide a fallback axis to align log. -log' :: Vec3d -> QDouble -> QDouble +log' :: Vector Double 3 -> QDouble -> QDouble log' r (unpackQ# -> (# x, y, z, w #)) = case (x * x) + (y * y) + (z * z) of 0.0 | w >= 0 - -> packQ 0.0 0.0 0.0 (log w) + -> packQ 0 0 0 (log w) | Vec3 rx ry rz <- r - -> packQ (pi*rx) (pi*ry) (pi*rz) (log (negate w)) - mv2 -> case (# sqrt (mv2 + (w * w)), sqrt mv2 #) of - (# mq, mv #) -> case atan2 mv w / mv of - l -> packQ (x * l) (y * l) (z * l) (log mq) + -> packQ (M_PI*rx) (M_PI*ry) (M_PI*rz) (log (negate w)) + mv2 -> case (# mv2 + w * w, sqrt mv2 #) of + (# q2, mv #) -> case atan2 mv w / mv of + l -> packQ (x * l) (y * l) (z * l) (0.5 * log q2) -- If q is negative real, provide a fallback axis to align sqrt. -sqrt' :: Vec3d -> QDouble -> QDouble +sqrt' :: Vector Double 3 -> QDouble -> QDouble sqrt' r (unpackQ# -> (# x, y, z, w #)) - = case (x * x) + (y * y) + (z * z) of - 0.0 | w >= 0 - -> packQ 0.0 0.0 0.0 (sqrt w) - | Vec3 rx ry rz <- r - , sw <- sqrt (negate w) - -> packQ (sw*rx) (sw*ry) (sw*rz) 0 - mv2 -> - let mq = sqrt (mv2 + w * w) - l2 = sqrt mq - tq = w / (mq * 2.0) - sina = sqrt (0.5 - tq) * l2 / sqrt mv2 - in packQ (x * sina) (y * sina) (z * sina) (sqrt (0.5 + tq) * l2) - -sigVec :: QDouble -> Vec3d -sigVec (unpackQ# -> (# x, y, z, _ #)) - = case sqrt ((x * x) + (y * y) + (z * z)) of - 0 -> vec3 1 0 0 - l -> vec3 (x/l) (y/l) (z/l) + | v2 == 0 && w >= 0 + = packQ x y z (sqrt w) + | v2 == 0 + , Vec3 rx ry rz <- r + , sw <- sqrt (negate w) + = packQ (sw*rx) (sw*ry) (sw*rz) 0 + | otherwise + = packQ (x * c) (y * c) (z * c) arg + where + v2 = (x * x) + (y * y) + (z * z) + mq = sqrt (v2 + w * w) + arg = sqrt $ 0.5 * if w >= 0 then mq + w else v2 / (mq - w) + c = 0.5 / arg instance Eq QDouble where {-# INLINE (==) #-} diff --git a/src/Numeric/Quaternion/Internal/QFloat.hs b/src/Numeric/Quaternion/Internal/QFloat.hs --- a/src/Numeric/Quaternion/Internal/QFloat.hs +++ b/src/Numeric/Quaternion/Internal/QFloat.hs @@ -10,19 +10,18 @@ {-# LANGUAGE UnboxedTuples #-} {-# LANGUAGE UndecidableInstances #-} {-# LANGUAGE ViewPatterns #-} -{-# OPTIONS_GHC -fno-warn-orphans #-} +{-# OPTIONS_GHC -fno-warn-orphans #-} module Numeric.Quaternion.Internal.QFloat ( QFloat, Quater (..) ) where import qualified Control.Monad.ST as ST import Data.Coerce (coerce) +import Numeric.Basics import Numeric.DataFrame.Internal.PrimArray import qualified Numeric.DataFrame.ST as ST import Numeric.DataFrame.Type -import Numeric.PrimBytes (PrimBytes) import Numeric.Quaternion.Internal -import Numeric.Scalar.Internal import Numeric.Vector.Internal type QFloat = Quater Float @@ -31,11 +30,11 @@ deriving instance PrimArray Float (Quater Float) instance Quaternion Float where - newtype Quater Float = QFloat Vec4f + newtype Quater Float = QFloat (Vector Float 4) {-# INLINE packQ #-} - packQ = coerce (vec4 :: Float -> Float -> Float -> Float -> Vec4f) + packQ = coerce (vec4 :: Float -> Float -> Float -> Float -> Vector Float 4) {-# INLINE unpackQ# #-} - unpackQ# = coerce (unpackV4# :: Vec4f -> (# Float, Float, Float, Float #)) + unpackQ# = coerce (unpackV4# :: Vector Float 4 -> (# Float, Float, Float, Float #)) {-# INLINE fromVecNum #-} fromVecNum (unpackV3# -> (# x, y, z #)) = packQ x y z {-# INLINE fromVec4 #-} @@ -47,7 +46,7 @@ {-# INLINE im #-} im (unpackQ# -> (# x, y, z, _ #)) = packQ x y z 0.0 {-# INLINE re #-} - re (unpackQ# -> (# _, _, _, w #)) = packQ 0.0 0.0 0.0 w + re (unpackQ# -> (# _, _, _, w #)) = packQ 0 0 0 w {-# INLINE imVec #-} imVec (unpackQ# -> (# x, y, z, _ #)) = vec3 x y z {-# INLINE taker #-} @@ -72,7 +71,7 @@ (l*z + d*k + t2 * (y*i - x*j)) {-# INLINE getRotScale #-} getRotScale a b = case (# unpackV3# a, unpackV3# b #) of - (# _, (# 0, 0, 0 #) #) -> packQ 0.0 0.0 0.0 0.0 + (# _, (# 0, 0, 0 #) #) -> packQ 0 0 0 0 (# (# 0, 0, 0 #), _ #) -> let x = (1 / 0 :: Float) in packQ x x x x (# (# a1, a2, a3 #), (# b1, b2, b3 #) #) -> let ma = sqrt (a1*a1 + a2*a2 + a3*a3) @@ -99,7 +98,7 @@ (# t1, t2, t3 #) -> packQ (t1 * r) (t2 * r) (t3 * r) (c / ma2) {-# INLINE axisRotation #-} axisRotation v a = case unpackV3# v of - (# 0, 0, 0 #) -> packQ 0.0 0.0 0.0 (if abs a < pi then 1 else -1) + (# 0, 0, 0 #) -> packQ 0 0 0 (negateUnless (abs a < M_PI) 1) (# x, y, z #) -> let c = cos (a * 0.5) s = sin (a * 0.5) @@ -261,39 +260,35 @@ = QFloat (a - b) {-# INLINE (-) #-} (unpackQ# -> (# a1, a2, a3, a4 #)) * (unpackQ# -> (# b1, b2, b3, b4 #)) - = packQ - ((a4 * b1) + - (a1 * b4) + - (a2 * b3) - - (a3 * b2) - ) - ((a4 * b2) - - (a1 * b3) + - (a2 * b4) + - (a3 * b1) - ) - ((a4 * b3) + - (a1 * b2) - - (a2 * b1) + - (a3 * b4) - ) - ((a4 * b4) - - (a1 * b1) - - (a2 * b2) - - (a3 * b3) - ) + = packQ ((a4 * b1) + (a1 * b4) + (a2 * b3) - (a3 * b2)) + ((a4 * b2) - (a1 * b3) + (a2 * b4) + (a3 * b1)) + ((a4 * b3) + (a1 * b2) - (a2 * b1) + (a3 * b4)) + ((a4 * b4) - (a1 * b1) - (a2 * b2) - (a3 * b3)) {-# INLINE (*) #-} negate (QFloat a) = QFloat (negate a) {-# INLINE negate #-} - abs = packQ 0.0 0.0 0.0 . sqrt . square + abs = packQ 0 0 0 . sqrt . square {-# INLINE abs #-} signum q@(unpackQ# -> (# x, y, z, w #)) - = case square q of - 0.0 -> packQ 0.0 0.0 0.0 0.0 - qd -> case recip (sqrt qd) of - s -> packQ (x * s) (y * s) (z * s) (w * s) + | qd == 0 = q + | otherwise = case ix + iy + iz + iw + nn of + 0 -> packQ (x * l) (y * l) (z * l) (w * l) + 1 -> packQ (copysign 1 x) 0 0 0 + 2 -> packQ 0 (copysign 1 y) 0 0 + 4 -> packQ 0 0 (copysign 1 z) 0 + 8 -> packQ 0 0 0 (copysign 1 w) + _ -> packQ n n n n + where + n = 0 / 0 :: Float + qd = x*x + y*y + z*z + w*w + ix = if isInfinite x then 1 else 0 :: Int + iy = if isInfinite y then 2 else 0 :: Int + iz = if isInfinite z then 4 else 0 :: Int + iw = if isInfinite w then 8 else 0 :: Int + nn = if isNaN x || isNaN y || isNaN z || isNaN w then 16 else 0 :: Int + l = recip (sqrt qd) {-# INLINE signum #-} - fromInteger = packQ 0.0 0.0 0.0 . fromInteger + fromInteger = packQ 0 0 0 . fromInteger {-# INLINE fromInteger #-} @@ -304,172 +299,253 @@ {-# INLINE (/) #-} a / b = a * recip b {-# INLINE fromRational #-} - fromRational = packQ 0.0 0.0 0.0 . fromRational + fromRational = packQ 0 0 0 . fromRational instance Floating QFloat where {-# INLINE pi #-} - pi = packQ 0.0 0.0 0.0 3.141592653589793 + pi = packQ 0 0 0 M_PI {-# INLINE exp #-} exp (unpackQ# -> (# x, y, z, w #)) - = case (# (x * x) + - (y * y) + - (z * z) - , exp w - #) of - (# 0.0, et #) -> packQ 0.0 0.0 0.0 et - (# mv2, et #) -> case sqrt mv2 of - mv -> case et * sin mv / mv of - l -> packQ (x * l) (y * l) (z * l) (et * cos mv) + | mv2 == 0 = packQ x y z ew + | otherwise = packQ (x * l) (y * l) (z * l) arg + where + mv2 = (x * x) + (y * y) + (z * z) + mv = sqrt mv2 + ew = exp w + l = ew * sin mv / mv + arg = ew * cos mv {-# INLINE log #-} log = log' (Vec3 1 0 0) {-# INLINE sqrt #-} sqrt = sqrt' (Vec3 1 0 0) {-# INLINE sin #-} sin (unpackQ# -> (# x, y, z, w #)) - = case (x * x) + (y * y) + (z * z) of - 0.0 -> packQ 0.0 0.0 0.0 (sin w) - mv2 -> case sqrt mv2 of - mv -> case cos w * sinh mv / mv of - l -> packQ (x * l) (y * l) (z * l) (sin w * cosh mv) + | mv2 == 0 = packQ x y z (sin w) + | otherwise = packQ (x * l) (y * l) (z * l) arg + where + mv2 = (x * x) + (y * y) + (z * z) + mv = sqrt mv2 + l = cos w * sinh mv / mv + arg = sin w * cosh mv {-# INLINE cos #-} cos (unpackQ# -> (# x, y, z, w #)) - = case (x * x) + (y * y) + (z * z) of - 0.0 -> packQ 0.0 0.0 0.0 (cos w) - mv2 -> case sqrt mv2 of - mv -> case sin w * sinh mv / negate mv of - l -> packQ (x * l) (y * l) (z * l) (cos w * cosh mv) + | mv2 == 0 = packQ x y z (cos w) + | otherwise = packQ (x * l) (y * l) (z * l) arg + where + mv2 = (x * x) + (y * y) + (z * z) + mv = sqrt mv2 + l = sin w * sinh mv / negate mv + arg = cos w * cosh mv {-# INLINE tan #-} tan (unpackQ# -> (# x, y, z, w #)) - = case (x * x) + (y * y) + (z * z) of - 0.0 -> packQ 0.0 0.0 0.0 (tan w) - mv2 -> - let mv = sqrt mv2 - chv = cosh mv - shv = sinh mv - ct = cos w - st = sin w - cq = recip ( (ct * ct * chv * chv) + (st * st * shv * shv) ) - l = chv * shv * cq / mv - in packQ (x * l) (y * l) (z * l) (ct * st * cq) + | mv2 == 0 = packQ x y z (tan w) + | isInfinite mv2 = signum (packQ x y z 0) + | otherwise = packQ (x * l) (y * l) (z * l) arg + where + mv2 = (x * x) + (y * y) + (z * z) + mv = sqrt mv2 + b = 2*mv + a = 2*w + sina = sin a + eb = exp (-b) + eb2 = eb*eb + d = 1 + eb2 + 2 * eb * cos a + rd = recip d + pa = M_PI - abs a + rd' = 2 / (b*b + pa*pa) + (l, arg) = + if d >= M_EPS + then ((1 - eb2) * rd / mv, 2 * sina * eb * rd) + else (2 * rd' , negate sina * rd') {-# INLINE sinh #-} sinh (unpackQ# -> (# x, y, z, w #)) - = case (x * x) + (y * y) + (z * z) of - 0.0 -> packQ 0.0 0.0 0.0 (sinh w) - mv2 -> case sqrt mv2 of - mv -> case cosh w * sin mv / mv of - l -> packQ (x * l) (y * l) (z * l) (sinh w * cos mv) + | mv2 == 0 = packQ x y z (sinh w) + | otherwise = packQ (x * l) (y * l) (z * l) arg + where + mv2 = (x * x) + (y * y) + (z * z) + mv = sqrt mv2 + l = cosh w * sin mv / mv + arg = sinh w * cos mv {-# INLINE cosh #-} cosh (unpackQ# -> (# x, y, z, w #)) - = case (x * x) + (y * y) + (z * z) of - 0.0 -> packQ 0.0 0.0 0.0 (cosh w) - mv2 -> case sqrt mv2 of - mv -> case sinh w * sin mv / mv of - l -> packQ (x * l) (y * l) (z * l) (cosh w * cos mv) + | mv2 == 0 = packQ x y z (cosh w) + | otherwise = packQ (x * l) (y * l) (z * l) arg + where + mv2 = (x * x) + (y * y) + (z * z) + mv = sqrt mv2 + l = sinh w * sin mv / mv + arg = cosh w * cos mv {-# INLINE tanh #-} tanh (unpackQ# -> (# x, y, z, w #)) - = case (x * x) + (y * y) + (z * z) of - 0.0 -> packQ 0.0 0.0 0.0 (tanh w) - mv2 -> - let mv = sqrt mv2 - cv = cos mv - sv = sin mv - cht = cosh w - sht = sinh w - cq = recip ( (cht * cht * cv * cv) + (sht * sht * sv * sv) ) - l = cv * sv * cq / mv - in packQ (x * l) (y * l) (z * l) (cht * sht * cq) + | mv2 == 0 = packQ x y z (tanh w) + | isInfinite mv2 = packQ 0 0 0 (signum w) + | otherwise = packQ (x * l) (y * l) (z * l) arg + where + mv2 = (x * x) + (y * y) + (z * z) + mv = sqrt mv2 + b = 2*w + a = 2*mv + eb = exp (- abs b) + eb2 = eb*eb + d = 1 + eb2 + 2 * eb * cos a + rd = recip d + pa = M_PI - a + rd' = 2 / (b*b + pa*pa) + (l, arg) = + if d >= M_EPS + then (2 * sin a * eb * rd / mv, copysign (1 - eb2) b * rd) + else (2 * rd' , b * rd') {-# INLINE asin #-} - asin q = -i * log' axis (i*q + sqrt' axis (1 - q*q)) - where - axis = sigVec q - i = fromVecNum axis 0 + -- The original formula: + -- asin q = -i * log (i*q + sqrt (1 - q*q)) + -- below is a more numerically stable version. + asin (unpackQ# -> (# x, y, z, w #)) + | v2 == 0 = if w2 <= 1 + then packQ x y z (asin w) + else packQ l 0 0 arg + | otherwise = packQ (x*c) (y*c) (z*c) arg + where + v2 = (x * x) + (y * y) + (z * z) + v = sqrt v2 + w2 = w*w + w1qq = 0.5 *(1 - w2 + v2) -- real part of (1 - q*q)/2 + l1qq = sqrt (w1qq*w1qq + w2*v2) -- length of (1 - q*q)/2 + sp2 = l1qq + w1qq + sn2 = l1qq - w1qq + sp = sqrt sp2 + sn = copysign (sqrt sn2) w + -- choose a more stable (symbolically equiv) version + dp = if 2 * v2 <= sp2 then sp - v else v2 / ((sp + v)*(sn2 + v2)) + dn = if 2 * w2 <= sn2 then w - sn else w2 / ((sn + w)*(sp2 + w2)) + (wD, vD) = case compare w1qq 0 of + GT -> (dp, w * dp / sp) + LT -> (v * dn / sn, dn) + EQ -> (-v, w) + l = -0.5 * log (wD*wD + vD*vD) + c = l / v + arg = atan2 vD wD {-# INLINE acos #-} - acos q = pi/2 - asin q + acos q = M_PI_2 - asin q {-# INLINE atan #-} -- atan q = i / 2 * log ( (i + q) / (i - q) ) atan (unpackQ# -> (# x, y, z, w #)) - = case (x * x) + (y * y) + (z * z) of - 0 -> packQ 0 0 0 (atan w) - v2 -> - let v = sqrt v2 - w2 = w*w - vw2 = w2 + v2 - t = 1 + vw2 - c = 0.25 * log ((t + 2*v) / (t - 2*v)) / v - in packQ (c*x) (c*y) (c*z) $ 0.5 * atan2 (2*w) (1 - vw2) + | v2 == 0 = packQ x y z (atan w) + | otherwise = packQ (x*c) (y*c) (z*c) arg + where + v2 = (x * x) + (y * y) + (z * z) + v = sqrt v2 + w2 = w*w + q2 = w2 + v2 + v' = v - 1 + mzero = w2 + v'*v' + (c, arg) = + if mzero == 0 + then ( sqrt maxFinite / v, 0) + else ( 0.25 * (log (1 + q2 + 2*v) - log mzero) / v + , 0.5 * atan2 (2*w) (1 - q2) ) {-# INLINE asinh #-} -- The original formula: -- asinh q = log (q + sqrt (q*q + 1)) -- below is a more numerically stable version. asinh (unpackQ# -> (# x, y, z, w #)) - | v2 == 0 - = packQ 0 0 0 (asinh w) - | w == 0 - = if v2 <= 1 - then let c = asin v / v - in packQ (c*x) (c*y) (c*z) 0 - else let c = 0.5 * pi / v - w' = 0.5 * log (2*v2 - 1 + 2 * v * sqrt ( v2 - 1 )) - in packQ (c*x) (c*y) (c*z) w' - | otherwise - = let c = atan2 v t / v - w' = 0.5 * log ( t2 + v2) - log t - + if w >= 0 - then log ( t + w ) - else log 2 + log w2 - log ( t - w ) - log ( r + w2 + v2 - 1 ) - in packQ (c*x) (c*y) (c*z) w' + | v2 == 0 = packQ x y z (asinh w) + | otherwise = packQ (x*c) (y*c) (z*c) arg where v2 = (x * x) + (y * y) + (z * z) - v = sqrt v2 + v = sqrt v2 w2 = w*w - wvp1 = w2 - v2 + 1 - r = sqrt $ wvp1*wvp1 + 4*w2*v2 - t = sqrt $ ( r + wvp1 ) * 0.5 - t2 = t*t + w1qq = 0.5 *(1 + w2 - v2) -- real part of (1 + q*q)/2 + l1qq = sqrt (w1qq*w1qq + w2*v2) -- length of (1 + q*q)/2 + sp2 = l1qq + w1qq + sn2 = l1qq - w1qq + sp = sqrt sp2 + sn = copysign (sqrt sn2) w + -- choose a more stable (symbolically equiv) version + dp = if 2 * w >= - sp then w + sp else w2 / ((sp - w)*(w2 + sn2)) + dn = if 2 * v <= - sn || sn >= 0 + then v + sn else v2 / ((v - sn)*(v2 + sp2)) + (wD, vD) = case compare w1qq 0 of + GT -> (dp, v * dp / sp) + LT -> (w * dn / sn, dn) + EQ -> (w, v) + c = atan2 vD wD / v + arg = 0.5 * log (wD*wD + vD*vD) {-# INLINE acosh #-} + -- The original formula: + -- asinh q = log (q + sqrt (q + 1) * sqrt (q - 1)) + -- below is a more numerically stable version. -- note, log (q + sqrt (q*q - 1)) would not work, because that would not -- be the principal value. - acosh q = log' axis (q + sqrt' axis (q + 1) * sqrt' axis (q - 1)) + acosh (unpackQ# -> (# x, y, z, w #)) + | v2 == 0 = packQ x y z (acosh w) + | otherwise = packQ (x*c) (y*c) (z*c) arg where - axis = sigVec q + v2 = (x * x) + (y * y) + (z * z) + v = sqrt v2 + w2 = w*w + w1qq = 0.5 *(w2 - v2 - 1) -- real part of (q*q - 1)/2 + l1qq = sqrt (w1qq*w1qq + w2*v2) -- length of (q*q - 1)/2 + sp2 = l1qq + w1qq + sn2 = l1qq - w1qq + sp = sqrt sp2 + sn = copysign (sqrt sn2) w + -- choose a more stable (symbolically equiv) version + dp = if 2 * w >= - sp then w + sp else w2 / ((w - sp)*(w2 + sn2)) + dn = if 2 * v <= - sn || sn >= 0 + then v + sn else v2 / ((sn - v)*(v2 + sp2)) + (wD, vD) = case compare w1qq 0 of + GT -> (dp, v * dp / sp) + LT -> (w * dn / sn, dn) + EQ -> (w, v) + c = atan2 vD wD / v + arg = 0.5 * log (wD*wD + vD*vD) {-# INLINE atanh #-} - atanh q = 0.5 * log ((1+q)/(1-q)) + -- atanh q = 0.5 * log ( (1 + q) / (1 - q) ) + atanh (unpackQ# -> (# x, y, z, w #)) + | v2 == 0 = packQ x y z (atanh w) + | otherwise = packQ (x*c) (y*c) (z*c) (copysign arg w) + where + v2 = (x * x) + (y * y) + (z * z) + v = sqrt v2 + w2 = w*w + q2 = w2 + v2 + w' = abs w - 1 + c = 0.5 * atan2 (2*v) (1 - q2) / v + arg = if w' == 0 + then (1/0) + else 0.25 * (log (1 + q2 + 2 * abs w) - log (v2 + w'*w')) + -- If q is negative real, provide a fallback axis to align log. -log' :: Vec3f -> QFloat -> QFloat +log' :: Vector Float 3 -> QFloat -> QFloat log' r (unpackQ# -> (# x, y, z, w #)) = case (x * x) + (y * y) + (z * z) of 0.0 | w >= 0 - -> packQ 0.0 0.0 0.0 (log w) + -> packQ 0 0 0 (log w) | Vec3 rx ry rz <- r - -> packQ (pi*rx) (pi*ry) (pi*rz) (log (negate w)) - mv2 -> case (# sqrt (mv2 + (w * w)), sqrt mv2 #) of - (# mq, mv #) -> case atan2 mv w / mv of - l -> packQ (x * l) (y * l) (z * l) (log mq) + -> packQ (M_PI*rx) (M_PI*ry) (M_PI*rz) (log (negate w)) + mv2 -> case (# mv2 + w * w, sqrt mv2 #) of + (# q2, mv #) -> case atan2 mv w / mv of + l -> packQ (x * l) (y * l) (z * l) (0.5 * log q2) -- If q is negative real, provide a fallback axis to align sqrt. -sqrt' :: Vec3f -> QFloat -> QFloat +sqrt' :: Vector Float 3 -> QFloat -> QFloat sqrt' r (unpackQ# -> (# x, y, z, w #)) - = case (x * x) + (y * y) + (z * z) of - 0.0 | w >= 0 - -> packQ 0.0 0.0 0.0 (sqrt w) - | Vec3 rx ry rz <- r - , sw <- sqrt (negate w) - -> packQ (sw*rx) (sw*ry) (sw*rz) 0 - mv2 -> - let mq = sqrt (mv2 + w * w) - l2 = sqrt mq - tq = w / (mq * 2.0) - sina = sqrt (0.5 - tq) * l2 / sqrt mv2 - in packQ (x * sina) (y * sina) (z * sina) (sqrt (0.5 + tq) * l2) - -sigVec :: QFloat -> Vec3f -sigVec (unpackQ# -> (# x, y, z, _ #)) - = case sqrt ((x * x) + (y * y) + (z * z)) of - 0 -> vec3 1 0 0 - l -> vec3 (x/l) (y/l) (z/l) + | v2 == 0 && w >= 0 + = packQ x y z (sqrt w) + | v2 == 0 + , Vec3 rx ry rz <- r + , sw <- sqrt (negate w) + = packQ (sw*rx) (sw*ry) (sw*rz) 0 + | otherwise + = packQ (x * c) (y * c) (z * c) arg + where + v2 = (x * x) + (y * y) + (z * z) + mq = sqrt (v2 + w * w) + arg = sqrt $ 0.5 * if w >= 0 then mq + w else v2 / (mq - w) + c = 0.5 / arg instance Eq QFloat where {-# INLINE (==) #-} diff --git a/src/Numeric/Scalar/Internal.hs b/src/Numeric/Scalar/Internal.hs --- a/src/Numeric/Scalar/Internal.hs +++ b/src/Numeric/Scalar/Internal.hs @@ -11,11 +11,8 @@ , pattern S ) where - -import Data.Type.Lits (Nat) -import GHC.Base (unsafeCoerce#) -import Numeric.DataFrame.SubSpace (SubSpace (ewgen)) -import Numeric.DataFrame.Type (DataFrame, pattern S) +import Numeric.DataFrame.Internal.PrimArray (broadcast) +import Numeric.DataFrame.Type -- | Alias for zero-dimensional DataFrame type Scalar t = DataFrame t ('[] :: [Nat]) @@ -24,20 +21,8 @@ type Sci = Scalar Int type Scw = Scalar Word --- | Convert scalar back to ordinary type -unScalar :: Scalar t -> t --- rely on that Scalar is just two times newtype alias to t -unScalar = unsafeCoerce# -{-# INLINE unScalar #-} - --- | Convert any type to scalar wrapper -scalar :: t -> Scalar t --- rely on that Scalar is just two times newtype alias to t -scalar = unsafeCoerce# -{-# INLINE scalar #-} - -- | Broadcast scalar value onto a whole data frame -fromScalar :: SubSpace t ds '[] ds - => Scalar t -> DataFrame t ds -fromScalar = ewgen +fromScalar :: PrimArray t (DataFrame t ds) + => Scalar t -> DataFrame t (ds :: [Nat]) +fromScalar = broadcast . unScalar {-# INLINE fromScalar #-} diff --git a/src/Numeric/Subroutine/Householder.hs b/src/Numeric/Subroutine/Householder.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/Subroutine/Householder.hs @@ -0,0 +1,278 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE PolyKinds #-} +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE StandaloneDeriving #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UndecidableInstances #-} +module Numeric.Subroutine.Householder + ( householderReflectionInplaceR + , householderReflectionInplaceL + ) where + + +import Control.Monad +import Control.Monad.ST +import Data.Kind +import Numeric.Basics +import Numeric.DataFrame.ST +import Numeric.DataFrame.Type +import Numeric.Dimensions +import Numeric.Scalar.Internal + + +{- | Run a Householder transformation inplace. + + Given some orthongonal matrix \(P\), some matrix \(R\) and index \((k,l)\), + reflects \(R\) along some hyperplane, such that all elements of \(R\) + below index \( (k, l) \) become zeros, + then updates \(P\) with the inverse of the same transform as \(R\). + + Notes and invariants: + + 1. The transformation happens inplace for both matrices \(P\) and \(R\); + if \( R = P^\intercal A \), then \( R' = P^*PR = P'^\intercal A \), where + \( P' \) and \( R' \) are the updated versions of the input matrices, + \( P^* \) and \( A \) are implicit matrices. + + 2. All elements below and to the left of index \(k,l\) in \(R\) + are assumed (and not checked) to be zeros; + these are not touched by the subroutine to save flops. + + 3. A logical starting value for \(P\) is an identity matrix. + The subroutine can be used for a QR decomposition: + \( Q = P \). + + Returns @True@ if reflection has been performed, and @False@ if it was not needed. + This can be used to track the sign of @det P@. + -} +householderReflectionInplaceL :: + forall (s :: Type) (t :: Type) (n :: Nat) (m :: Nat) + . (PrimBytes t, Ord t, Epsilon t, KnownDim n, KnownDim m) + => STDataFrame s t '[n] -- ^ Temporary buffer for a Householder axis vector + -> STDataFrame s t '[n,n] -- ^ Current state of \(P^\intercal\) + -> STDataFrame s t '[n,m] -- ^ Current state of \(R\) + -> Idxs '[n,m] -- ^ Pivot element + -> ST s Bool +householderReflectionInplaceL u p r (Idx i :* Idx j :* U) + = householderReflectionInplaceL' u p r + (fromIntegral $ dimVal' @n) + (fromIntegral $ dimVal' @m) + (fromIntegral i) + (fromIntegral j) + +householderReflectionInplaceL' :: + forall (s :: Type) (t :: Type) (n :: Nat) (m :: Nat) + . (PrimBytes t, Epsilon t, Ord t) + => STDataFrame s t '[n] -- ^ Temporary buffer for a Householder axis vector + -> STDataFrame s t '[n,n] -- ^ \(P^\intercal\) + -> STDataFrame s t '[n,m] -- ^ \(R\) + -> Int -- ^ \(n\) + -> Int -- ^ \(m\) + -> Int -- ^ \( 0 \leq k < n \) + -> Int -- ^ \( 0 \leq l < m \) + -> ST s Bool +householderReflectionInplaceL' uPtr pPtr rPtr n m k l = do + -- pivot element (k,l) of new R + alpha <- getAlphaAndUpdateU + u2 <- getU2 + -- u2 == 0 means the column is already zeroed + if u2 > M_EPS + then do + let c = 2 / u2 -- a mult constant for updating matrices + -- update R + updateRl alpha + forM_ [l+1..m-1] $ updateRi c + -- update P + forM_ [0..n-1] $ updatePi c + return True + else return False + where + n' = n - k -- remaining rows + rOff0 = k*m + l -- offset of element (k,l) in matrix R + + -- u = Rk - alpha*ek + getAlphaAndUpdateU :: ST s (Scalar t) + getAlphaAndUpdateU = do + alpha' <- sqrt . fst <$> nTimesM n' + (\(r, off) -> do + x <- readDataFrameOff rPtr off + when (abs x <= M_EPS) $ writeDataFrameOff rPtr off 0 + return (r + x*x, off + m) + ) (0, rOff0) + x0 <- readDataFrameOff rPtr rOff0 + let alpha = if x0 >= 0 then negate alpha' else alpha' + -- update (lower part of) u + writeDataFrameOff uPtr k (x0 - alpha) + when (n' >= 1) $ void $ nTimesM (n' - 1) + (\(i, off) -> (i+1, off+m) <$ + (readDataFrameOff rPtr off >>= writeDataFrameOff uPtr i) + ) (k+1, rOff0+m) + return alpha + + -- l-th column of R zeroes below pivot + updateRl :: Scalar t -> ST s () + updateRl alpha = do + writeDataFrameOff rPtr rOff0 alpha + when (n' >= 1) $ void $ nTimesM (n' - 1) + (\off -> (off+m) <$ writeDataFrameOff rPtr off 0) (rOff0+m) + + -- update i-th column of R + updateRi :: Scalar t -> Int -> ST s () + updateRi c i = do + -- dot product of u and Ri + uRi <- fmap fst . flip (nTimesM n') (0, (k, k*m+i)) $ \(r, (j, off)) -> do + ju <- readDataFrameOff uPtr j + jiR <- readDataFrameOff rPtr off + return (r + ju * jiR, (j+1, off+m)) + let c' = c * uRi + -- update each element + void $ flip (nTimesM n') (k, k*m+i) $ \(j, off) -> do + ju <- readDataFrameOff uPtr j + jiR <- readDataFrameOff rPtr off + writeDataFrameOff rPtr off $ jiR - c'*ju + return (j+1, off+m) + + -- update i-th row of P + updatePi :: Scalar t -> Int -> ST s () + updatePi c i = do + let off0 = i*n + -- dot product of u and Pi + uPi <- fmap fst . flip (nTimesM n') (0, k) $ \(r, j) -> do + ju <- readDataFrameOff uPtr j + ijP <- readDataFrameOff pPtr (off0 + j) + return (r + ju * ijP, j+1) + let c' = c * uPi + -- update each element + forM_ [k..n-1] $ \j -> do + ju <- readDataFrameOff uPtr j + ijP <- readDataFrameOff pPtr (off0 + j) + writeDataFrameOff pPtr (off0 + j) $ ijP - c'*ju + + -- get module squared of u (for Q = I - 2 u*uT / |u|^2 ) + getU2 :: ST s (Scalar t) + getU2 = fst <$> nTimesM n' + (\(r, off) -> (\x -> (r + x*x, off + 1)) <$> readDataFrameOff uPtr off) (0, k) + +{- | Run a Householder transformation inplace. + + Similar to `householderReflectionInplaceR`, but works from right to left + - use to zero elements to the right from the pivot. + + Returns @True@ if reflection has been performed, and @False@ if it was not needed. + This can be used to track the sign of @det P@. + -} +householderReflectionInplaceR :: + forall (s :: Type) (t :: Type) (n :: Nat) (m :: Nat) + . (PrimBytes t, Ord t, Epsilon t, KnownDim n, KnownDim m) + => STDataFrame s t '[m] -- ^ Temporary buffer for a Householder axis vector + -> STDataFrame s t '[m,m] -- ^ Current state of \(P^\intercal\) + -> STDataFrame s t '[n,m] -- ^ Current state of \(R\) + -> Idxs '[n,m] -- ^ Pivot element + -> ST s Bool +householderReflectionInplaceR u p r (Idx i :* Idx j :* U) + = householderReflectionInplaceR' u p r + (fromIntegral $ dimVal' @n) + (fromIntegral $ dimVal' @m) + (fromIntegral i) + (fromIntegral j) + +householderReflectionInplaceR' :: + forall (s :: Type) (t :: Type) (n :: Nat) (m :: Nat) + . (PrimBytes t, Epsilon t, Ord t) + => STDataFrame s t '[m] -- ^ Temporary buffer for a Householder axis vector + -> STDataFrame s t '[m,m] -- ^ \(P^\intercal\) + -> STDataFrame s t '[n,m] -- ^ \(R\) + -> Int -- ^ \(n\) + -> Int -- ^ \(m\) + -> Int -- ^ \( 0 \leq k < n \) + -> Int -- ^ \( 0 \leq l < m \) + -> ST s Bool +householderReflectionInplaceR' uPtr pPtr rPtr n m k l = do + -- pivot element (k,l) of new R + alpha <- getAlphaAndUpdateU + u2 <- getU2 + -- u2 == 0 means the column is already zeroed + if u2 > M_EPS + then do + let c = 2 / u2 -- a mult constant for updating matrices + -- update R + updateRk alpha + forM_ [k+1..n-1] $ updateRi c + -- update P + forM_ [0..m-1] $ updatePi c + return True + else return False + where + m' = m - l -- remaining cols + rOff0 = k*m + l -- offset of element (k,l) in matrix R + + -- u = Rl - alpha*el + getAlphaAndUpdateU :: ST s (Scalar t) + getAlphaAndUpdateU = do + alpha' <- sqrt . fst <$> nTimesM m' + (\(r, off) -> do + x <- readDataFrameOff rPtr off + when (abs x <= M_EPS) $ writeDataFrameOff rPtr off 0 + return (r + x*x, off + 1) + ) (0, rOff0) + x0 <- readDataFrameOff rPtr rOff0 + let alpha = if x0 >= 0 then negate alpha' else alpha' + -- update (lower part of) u + writeDataFrameOff uPtr l (x0 - alpha) + forM_ [1..m'-1] $ \i -> + readDataFrameOff rPtr (rOff0 + i) >>= writeDataFrameOff uPtr (l + i) + return alpha + + -- k-th row of R zeroes below pivot + updateRk :: Scalar t -> ST s () + updateRk alpha = do + writeDataFrameOff rPtr rOff0 alpha + forM_ [rOff0+1..rOff0+m'-1] $ flip (writeDataFrameOff rPtr) 0 + + -- update i-th row of R + updateRi :: Scalar t -> Int -> ST s () + updateRi c i = do + let off0 = i*m + -- dot product of u and Ri + uRi <- fmap fst . flip (nTimesM m') (0, l) $ \(r, j) -> do + ju <- readDataFrameOff uPtr j + jiR <- readDataFrameOff rPtr (off0 + j) + return (r + ju * jiR, j+1) + let c' = c * uRi + -- update each element + forM_ [l..m-1] $ \j -> do + ju <- readDataFrameOff uPtr j + jiR <- readDataFrameOff rPtr (off0 + j) + writeDataFrameOff rPtr (off0 + j) $ jiR - c'*ju + + -- update i-th row of P + updatePi :: Scalar t -> Int -> ST s () + updatePi c i = do + let off0 = i*m + -- dot product of u and Pi + uPi <- fmap fst . flip (nTimesM m') (0, l) $ \(r, j) -> do + ju <- readDataFrameOff uPtr j + ijP <- readDataFrameOff pPtr (off0 + j) + return (r + ju * ijP, j+1) + let c' = c * uPi + -- update each element + forM_ [l..m-1] $ \j -> do + ju <- readDataFrameOff uPtr j + ijP <- readDataFrameOff pPtr (off0 + j) + writeDataFrameOff pPtr (off0 + j) $ ijP - c'*ju + + -- get module squared of u (for Q = I - 2 u*uT / |u|^2 ) + getU2 :: ST s (Scalar t) + getU2 = fst <$> nTimesM m' + (\(r, off) -> (\x -> (r + x*x, off + 1)) <$> readDataFrameOff uPtr off) (0, l) + + +nTimesM :: Monad m => Int -> (a -> m a) -> a -> m a +nTimesM 0 _ x = pure x +nTimesM n m x = m x >>= nTimesM (n - 1) m diff --git a/src/Numeric/Subroutine/SolveTriangular.hs b/src/Numeric/Subroutine/SolveTriangular.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/Subroutine/SolveTriangular.hs @@ -0,0 +1,173 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE PolyKinds #-} +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeOperators #-} +{- | +A few ways to solve a system of linear equations in ST monad. +The tesult is always computed inplace. + -} +module Numeric.Subroutine.SolveTriangular + ( solveUpperTriangularR + , solveUpperTriangularL + , solveLowerTriangularR + , solveLowerTriangularL + ) where + + +import Control.Monad +import Control.Monad.ST +import Data.Kind +import Numeric.DataFrame.Internal.PrimArray +import Numeric.DataFrame.ST +import Numeric.DataFrame.Type +import Numeric.Dimensions + + +{- | +Solve a system of linear equations \( Rx = b \) + or a linear least squares problem \( \min {|| Rx - b ||}^2 \), +where \( R \) is an upper-triangular matrix. + +DataFrame \( b \) is modified in-place; by the end of the process \( b_m = x \). + +NB: you can use `subDataFrameView` to truncate @b@ without performing a copy. + -} +solveUpperTriangularR :: + forall (s :: Type) (t :: Type) (n :: Nat) (m :: Nat) (ds :: [Nat]) + . (PrimBytes t, Fractional t, Eq t, KnownDim m, m <= n) + => DataFrame t '[n,m] -- ^ \(R\) + -> STDataFrame s t (m :+ ds) -- ^ Current state of \(b_m\) + -- (first @m@ rows of @b@) + -> ST s () +solveUpperTriangularR r bPtr | Dict <- Dict @(m <= n) = mapM_ go [m-1,m-2..0] + where + m = fromIntegral $ dimVal' @m :: Int + k = fromIntegral kw :: Int + CumulDims (_:kw:_) = getDataFrameSteps bPtr + go :: Int -> ST s () + go i | rii == 0 = forM_ [0..k-1] $ \j -> writeDataFrameOff bPtr (ki + j) 0 + | otherwise = forM_ [0..k-1] $ \j -> do + bij <- (rrii*) <$> readDataFrameOff bPtr (ki + j) + writeDataFrameOff bPtr (ki + j) bij + forM_ [0..i-1] $ \t -> do + let rti = scalar (ixOff (m*t + i) r) + ix = k*t + j + btj <- readDataFrameOff bPtr ix + writeDataFrameOff bPtr ix (btj - rti*bij) + where + mi = m*i + ki = k*i + rii = scalar (ixOff (mi + i) r) + rrii = recip rii + + +{- | +Solve a system of linear equations \( xR = b \), +where \( R \) is an upper-triangular matrix. + +DataFrame \( b \) is modified in-place; by the end of the process \( b = x_m \). +The \( (n - m) \) rows of \(R\) are not used. +Pad each dimension of \(x\) with \( (n - m) \) zeros if you want to get the full +solution. + -} +solveUpperTriangularL :: + forall (s :: Type) (t :: Type) (n :: Nat) (m :: Nat) (ds :: [Nat]) + . (PrimBytes t, Fractional t, Eq t, KnownDim m, m <= n) + => STDataFrame s t (ds +: m) -- ^ Current state of \(b\) + -- (first @m@ "columns" of x) + -> DataFrame t '[n,m] -- ^ \(R\) + -> ST s () +solveUpperTriangularL bPtr r | Dict <- Dict @(m <= n) = mapM_ go [0..m-1] + where + m = fromIntegral $ dimVal' @m :: Int + k = fromIntegral mkw `quot` m + CumulDims (mkw:_) = getDataFrameSteps bPtr + go :: Int -> ST s () + go i | rii == 0 = forM_ [0..k-1] $ \j -> writeDataFrameOff bPtr (m*j + i) 0 + | otherwise = forM_ [0..k-1] $ \j -> do + let mj = m*j + bji <- (rrii*) <$> readDataFrameOff bPtr (mj + i) + writeDataFrameOff bPtr (mj + i) bji + forM_ [i+1..m-1] $ \t -> do + let rit = scalar (ixOff (mi + t) r) + bjt <- readDataFrameOff bPtr (mj + t) + writeDataFrameOff bPtr (mj + t) (bjt - rit*bji) + where + mi = m*i + rii = scalar (ixOff (mi + i) r) + rrii = recip rii + +{- | +Solve a system of linear equations \( Lx = b \), +where \( L \) is a lower-triangular matrix. + +DataFrame \( b \) is modified in-place; by the end of the process \( b = x_n \). +The \( (m - n) \) columns of \(L\) are not used. +Pad \(x\) with \( (m - n) \) zero elements if you want to get the full solution. + -} +solveLowerTriangularR :: + forall (s :: Type) (t :: Type) (n :: Nat) (m :: Nat) (ds :: [Nat]) + . (PrimBytes t, Fractional t, Eq t, KnownDim n, KnownDim m, n <= m) + => DataFrame t '[n,m] -- ^ \(L\) + -> STDataFrame s t (n :+ ds) -- ^ Current state of \(b\) + -- (first @n@ elements of x) + -> ST s () +solveLowerTriangularR l bPtr | Dict <- Dict @(n <= m) = mapM_ go [0..n-1] + where + m = fromIntegral $ dimVal' @m :: Int + n = fromIntegral $ dimVal' @n :: Int + k = fromIntegral kw :: Int + CumulDims (_:kw:_) = getDataFrameSteps bPtr + go :: Int -> ST s () + go i | lii == 0 = forM_ [0..k-1] $ \j -> writeDataFrameOff bPtr (ki + j) 0 + | otherwise = forM_ [0..k-1] $ \j -> do + bij <- (rlii*) <$> readDataFrameOff bPtr (ki + j) + writeDataFrameOff bPtr (ki + j) bij + forM_ [i+1..n-1] $ \t -> do + let rti = scalar (ixOff (m*t + i) l) + ix = k*t + j + btj <- readDataFrameOff bPtr ix + writeDataFrameOff bPtr ix (btj - rti*bij) + where + mi = m*i + ki = k*i + lii = scalar (ixOff (mi + i) l) + rlii = recip lii + +{- | +Solve a system of linear equations \( xL = b \) + or a linear least squares problem \( \min {|| xL - b ||}^2 \), +where \( L \) is a lower-triangular matrix. + +DataFrame \( b \) is modified in-place; by the end of the process \( b_n = x \). +The last \( (m - n) \) columns of \(L\) and \(b\) and are not touched. + -} +solveLowerTriangularL :: + forall (s :: Type) (t :: Type) (n :: Nat) (m :: Nat) (ds :: [Nat]) + . (PrimBytes t, Fractional t, Eq t, KnownDim n, KnownDim m, n <= m) + => STDataFrame s t (ds +: m) -- ^ Current state of \(b\) + -> DataFrame t '[n,m] -- ^ \(L\) + -> ST s () +solveLowerTriangularL bPtr l | Dict <- Dict @(n <= m) = mapM_ go [n-1,n-2..0] + where + m = fromIntegral $ dimVal' @m :: Int + n = fromIntegral $ dimVal' @n :: Int + k = fromIntegral kmw `quot` m + CumulDims (kmw:_) = getDataFrameSteps bPtr + go :: Int -> ST s () + go i | lii == 0 = forM_ [0..k-1] $ \j -> writeDataFrameOff bPtr (m*j + i) 0 + | otherwise = forM_ [0..k-1] $ \j -> do + let mj = m*j + bji <- (rlii*) <$> readDataFrameOff bPtr (mj + i) + writeDataFrameOff bPtr (mj + i) bji + forM_ [0..i-1] $ \t -> do + let lit = scalar (ixOff (mi + t) l) + ix = m*j + t + bjt <- readDataFrameOff bPtr ix + writeDataFrameOff bPtr ix (bjt - lit*bji) + where + mi = m*i + lii = scalar (ixOff (mi + i) l) + rlii = recip lii diff --git a/src/Numeric/Subroutine/Sort.hs b/src/Numeric/Subroutine/Sort.hs new file mode 100644 --- /dev/null +++ b/src/Numeric/Subroutine/Sort.hs @@ -0,0 +1,250 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE PolyKinds #-} +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeInType #-} +{-# LANGUAGE TypeOperators #-} +module Numeric.Subroutine.Sort + ( SortBy (..), sortBy, sort + , SortableDataFrame + ) where + + +import Control.Monad +import Control.Monad.ST +import Control.Monad.ST.Unsafe +import Data.Kind +import Data.Type.Lits +import Numeric.DataFrame.Internal.PrimArray +import Numeric.DataFrame.ST +import Numeric.DataFrame.Type +import Numeric.Dimensions +import Unsafe.Coerce + +-- | Sort a @DataFrame@ along the first dimension. +-- +-- Note: the elements (which are of type @DataFrame t ns@) are compared +-- lexicographically. +sort :: forall (t :: Type) n ns + . ( SortableDataFrame t (n ': ns), Ord t, SortBy n) + => DataFrame t (n ': ns) + -> DataFrame t (n ': ns) +sort df = case dimKind @(KindOf n) of + DimKNat -> case uniqueOrCumulDims df of + Left _ -> df -- all equal, no need for sorting. + Right steps + | SomeDims (Dims :: Dims ms) <- fromSteps steps + , Dict <- (unsafeCoerce (Dict @(ns ~ ns)) :: Dict (ns ~ ms)) + -> sortBy compare df + | otherwise + -> error "sort/DimNat/uniqueOrCumulDims -- impossible pattern" + DimKXNat + | XFrame (df' :: DataFrame t ms) <- df + , D :* Dims <- dims @ms + -> XFrame (sortBy compare df') + | otherwise + -> error "sort/DimXNat -- impossible pattern" + +-- | Sort a @DataFrame@ along the first dimension using given comparison function. +sortBy :: forall (t :: Type) n ns + . ( SortableDataFrame t (n ': ns) + , SortBy n) + => (DataFrame t ns -> DataFrame t ns -> Ordering) + -> DataFrame t (n ': ns) + -> DataFrame t (n ': ns) +sortBy cmp df = case dimKind @(KindOf n) of + DimKNat -> runST $ + flip (withThawDataFrame (const $ pure df)) df $ \mdf -> do + sortByInplace + (\x y -> cmp <$> unsafeFreezeDataFrame x <*> unsafeFreezeDataFrame y) + mdf + unsafeFreezeDataFrame mdf + DimKXNat + | XFrame dfN <- df + , D :* Dims <- dims `inSpaceOf` dfN + -> XFrame (sortBy (\a b -> cmp (XFrame a) (XFrame b)) dfN) + | otherwise + -> error "sortBy/DimXNat -- impossible pattern" + +-- | The required context for sorting a DataFrame is slightly different +-- for @Nat@ and @XNat@ indexed arrays. +-- This type family abstracts away the difference. +type family SortableDataFrame (t :: Type) (ns :: [k]) :: Constraint where + SortableDataFrame t ((n ': ns) :: [Nat]) + = (PrimArray t (DataFrame t ns), PrimArray t (DataFrame t (n ': ns))) + SortableDataFrame t ((n ': ns) :: [XNat]) + = PrimBytes t + + +class BoundedDim n => SortBy n where + -- | Note, "Inplace" here means the input frame is modified. + -- It does not mean the algorithm does not use extra space (it does use). + sortByInplace :: PrimBytes t + => (STDataFrame s t ns -> STDataFrame s t ns -> ST s Ordering) + -- ^ must not modify state! + -> STDataFrame s t (n ': ns) + -> ST s () + + +instance SortBy 0 where + sortByInplace _ _ = pure () + +instance SortBy 1 where + sortByInplace _ _ = pure () + +instance SortBy 2 where + sortByInplace cmp xs = cmp a b >>= \case + GT -> do + tmp <- oneMoreDataFrame a + swapDF tmp a b + _ -> pure () + where + a = subDataFrameView' (Idx 0 :* U) xs + b = subDataFrameView' (Idx 1 :* U) xs + +instance SortBy 3 where + sortByInplace cmp xs = join $ + go <$> (unsafeDupableInterleaveST (oneMoreDataFrame a)) + <*> cmp a b <*> cmp b c <*> cmp a c + where + a = subDataFrameView' (Idx 0 :* U) xs + b = subDataFrameView' (Idx 1 :* U) xs + c = subDataFrameView' (Idx 2 :* U) xs + go tmp GT LT GT -- b < c < a + = swap3DF tmp a b c + go tmp LT GT GT -- c < a < b + = swap3DF tmp b a c + go tmp GT bc ac | bc /= GT && ac /= GT + = swapDF tmp a b + go tmp ab GT ac | ab /= GT && ac /= GT + = swapDF tmp b c + go tmp ab bc GT | ab /= LT && bc /= LT + = swapDF tmp a c + go _ _ _ _ = pure () + +instance SortBy 4 where + sortByInplace cmp xs = do + tmp <- unsafeDupableInterleaveST (oneMoreDataFrame a) + cmpSwap tmp a c + cmpSwap tmp b d + cmpSwap tmp a b + cmpSwap tmp c d + cmpSwap tmp b c + where + a = subDataFrameView' (Idx 0 :* U) xs + b = subDataFrameView' (Idx 1 :* U) xs + c = subDataFrameView' (Idx 2 :* U) xs + d = subDataFrameView' (Idx 3 :* U) xs + cmpSwap tmp x y = cmp x y >>= \case + GT -> swapDF tmp x y + _ -> pure () + +instance {-# INCOHERENT #-} + KnownDim n => SortBy (n :: Nat) where + sortByInplace cmp (xs :: STDataFrame s t (n ': ns)) = do + tmp <- oneMoreDataFrame xs + copyMutableDataFrame' U xs tmp + mergeSort D tmp xs + where + mergeSort :: Dim (d :: Nat) + -> STDataFrame s t (d ': ns) + -> STDataFrame s t (d ': ns) + -> ST s () + mergeSort D0 _ _ = pure () + mergeSort D1 _ _ = pure () + mergeSort (d@D :: Dim d) b a = do + d2l@D <- pure $ divDim d D2 + Just d2r@D <- pure $ minusDimM d d2l + d2li@D <- pure $ plusDim d2l D1 + d2ri@D <- pure $ plusDim d2r D1 + Just Dict <- pure $ sameDim (plusDim d D1) (plusDim d2li d2r) + Just Dict <- pure $ sameDim (plusDim d D1) (plusDim d2ri d2l) + let leA = subDataFrameView @t @d @(d - Div d 2 + 1) @(Div d 2) @'[] + (Idx 0 :* U) a + riA = subDataFrameView @t @d @(Div d 2 + 1) @(d - Div d 2) @'[] + (Idx (dimVal d2l) :* U) a + leB = subDataFrameView @t @d @(d - Div d 2 + 1) @(Div d 2) @'[] + (Idx 0 :* U) b + riB = subDataFrameView @t @d @(Div d 2 + 1) @(d - Div d 2) @'[] + (Idx (dimVal d2l) :* U) b + mergeSort d2l leA leB + mergeSort d2r riA riB + merge d2l d2r d leB riB a + merge :: forall (a :: Nat) (b :: Nat) (ab :: Nat) + . Dim a -> Dim b -> Dim ab + -> STDataFrame s t (a ': ns) + -> STDataFrame s t (b ': ns) + -> STDataFrame s t (ab ': ns) + -> ST s () + merge da@D db@D dab@D a b ab = foldM_ f (Just (0,0)) [0 .. dimVal dab - 1] + where + f Nothing _ = pure Nothing + f (Just (i,j)) k + | i >= dimVal da + , Dx dj@(D :: Dim j) <- someDimVal j + , D <- plusDim dj D1 + , Just bmj@D <- minusDimM db dj + , Just bmji@D <- minusDimM (plusDim dab D1) bmj + , Just Dict <- sameDim (plusDim dab D1) (plusDim bmji bmj) + , Just Dict <- sameDim (plusDim db D1) (dj `plusDim` D1 `plusDim` bmj) + = Nothing <$ copyMutableDataFrame @t @ab @(ab + 1 - (b - j)) + @(b - j) (Idx k :* U) + (subDataFrameView @t @b @(j + 1) @(b - j) (Idx j :* U) b) ab + | j >= dimVal db + , Dx di@(D :: Dim i) <- someDimVal i + , D <- plusDim di D1 + , Just bmi@D <- minusDimM da di + , Just bmii@D <- minusDimM (plusDim dab D1) bmi + , Just Dict <- sameDim (plusDim dab D1) (plusDim bmii bmi) + , Just Dict <- sameDim (plusDim da D1) (di `plusDim` D1 `plusDim` bmi) + = Nothing <$ copyMutableDataFrame (Idx k :* U) + (subDataFrameView @t @a @(i + 1) @(a - i) (Idx i :* U) a) ab + | otherwise + = cmp (subDataFrameView' (Idx i :* U) a) + (subDataFrameView' (Idx j :* U) b) >>= \case + GT -> Just (i, j + 1) + <$ copyMutableDataFrame' (Idx k :* U) + (subDataFrameView' (Idx j :* U) b) ab + _ -> Just (i + 1, j) + <$ copyMutableDataFrame' (Idx k :* U) + (subDataFrameView' (Idx i :* U) a) ab + + +instance BoundedDim xn => SortBy (xn :: XNat) where + sortByInplace cmp (XSTFrame xs) + | D :* _ <- dims `inSpaceOf` xs + = sortByInplace (\x y -> cmp (castDataFrame x) (castDataFrame y)) xs + | otherwise = error "sortByInplace: impossible pattern" + + +-- | Swap contents of two DataFrames +swapDF :: forall (s :: Type) (t :: Type) (ns :: [Nat]) + . PrimBytes t + => STDataFrame s t ns -- ^ Temporary buffer + -> STDataFrame s t ns + -> STDataFrame s t ns + -> ST s () +swapDF tmp a b = do + copyMutableDataFrame' U a tmp + copyMutableDataFrame' U b a + copyMutableDataFrame' U tmp b + +-- | Rotate left contents of three DataFrames +swap3DF :: forall (s :: Type) (t :: Type) (ns :: [Nat]) + . PrimBytes t + => STDataFrame s t ns -- ^ Temporary buffer + -> STDataFrame s t ns + -> STDataFrame s t ns + -> STDataFrame s t ns + -> ST s () +swap3DF tmp a b c = do + copyMutableDataFrame' U a tmp + copyMutableDataFrame' U b a + copyMutableDataFrame' U c b + copyMutableDataFrame' U tmp c diff --git a/src/Numeric/Vector/Internal.hs b/src/Numeric/Vector/Internal.hs --- a/src/Numeric/Vector/Internal.hs +++ b/src/Numeric/Vector/Internal.hs @@ -8,7 +8,9 @@ {-# LANGUAGE PatternSynonyms #-} {-# LANGUAGE PolyKinds #-} {-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} {-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeInType #-} {-# LANGUAGE TypeOperators #-} {-# LANGUAGE UnboxedTuples #-} {-# LANGUAGE UndecidableInstances #-} @@ -29,6 +31,7 @@ , det2, cross, (×) ) where +import Data.Kind import Numeric.DataFrame.SubSpace import Numeric.DataFrame.Type import Numeric.Scalar.Internal @@ -97,8 +100,8 @@ f _ = scalar b {-# INLINE vec2 #-} unpackV2# v = - (# unScalar (indexOffset# 0# v) - , unScalar (indexOffset# 1# v) #) + (# unScalar (indexOffset 0 v) + , unScalar (indexOffset 1 v) #) {-# INLINE unpackV2# #-} instance {-# OVERLAPPABLE #-} SubSpace t '[3] '[] '[3] => Vector3 t where @@ -109,9 +112,9 @@ f _ = scalar c {-# INLINE vec3 #-} unpackV3# v = - (# unScalar (indexOffset# 0# v) - , unScalar (indexOffset# 1# v) - , unScalar (indexOffset# 2# v) #) + (# unScalar (indexOffset 0 v) + , unScalar (indexOffset 1 v) + , unScalar (indexOffset 2 v) #) {-# INLINE unpackV3# #-} instance {-# OVERLAPPABLE #-} SubSpace t '[4] '[] '[4] => Vector4 t where @@ -123,10 +126,10 @@ f _ = scalar d {-# INLINE vec4 #-} unpackV4# v = - (# unScalar (indexOffset# 0# v) - , unScalar (indexOffset# 1# v) - , unScalar (indexOffset# 2# v) - , unScalar (indexOffset# 3# v) #) + (# unScalar (indexOffset 0 v) + , unScalar (indexOffset 1 v) + , unScalar (indexOffset 2 v) + , unScalar (indexOffset 3 v) #) {-# INLINE unpackV4# #-} @@ -137,7 +140,7 @@ , Num (Vector t n) , SubSpace t '[n] '[] '[n] ) - => Vector t n -> Vector t n -> Vector t n + => Vector t (n :: Nat) -> Vector t n -> Vector t n (.*.) a b = fromScalar . ewfoldl (+) 0 $ a * b infixl 7 .*. @@ -146,7 +149,7 @@ , Num (Vector t n) , SubSpace t '[n] '[] '[n] ) - => Vector t n -> Vector t n -> Scalar t + => Vector t (n :: Nat) -> Vector t n -> Scalar t dot a b = ewfoldl (+) 0 $ a * b -- | Dot product of two vectors @@ -155,41 +158,47 @@ , Num (Vector t n) , SubSpace t '[n] '[] '[n] ) - => Vector t n -> Vector t n -> Scalar t + => Vector t (n :: Nat) -> Vector t n -> Scalar t (·) = dot {-# INLINE (·) #-} -- | Sum of absolute values -normL1 :: ( Num t, SubSpace t '[n] '[] '[n] ) - => Vector t n -> Scalar t +normL1 :: ( Num (DataFrame t ('[] :: [k])), SubSpace t '[n] '[] '[n] ) + => Vector t (n :: k) -> DataFrame t ('[] :: [k]) normL1 = ewfoldr (\a -> (abs a +)) 0 -- | hypot function (square root of squares) -normL2 :: ( Floating t , SubSpace t '[n] '[] '[n] ) - => Vector t n -> Scalar t +normL2 :: ( Floating (DataFrame t ('[] :: [k])), SubSpace t '[n] '[] '[n] ) + => Vector t (n :: k) -> DataFrame t ('[] :: [k]) normL2 = sqrt . ewfoldr (\a -> (a*a +)) 0 -- | Normalize vector w.r.t. Euclidean metric (L2). -normalized :: ( Floating t , Fractional (Vector t n), SubSpace t '[n] '[] '[n] ) +normalized :: forall (k :: Type) (t :: Type) (n :: k) + . ( Floating (DataFrame t ('[] :: [k])) + , SubSpace t '[n] '[] '[n] ) => Vector t n -> Vector t n -normalized v = v / n +normalized v = ewmap @t @'[n] @'[] (*rn) v where - n = fromScalar . sqrt $ ewfoldr (\a -> (a*a +)) 0 v + rn = recip . sqrt $ ewfoldr (\a -> (a*a +)) 0 v -- | Maximum of absolute values -normLPInf :: ( Ord t, Num t , SubSpace t '[n] '[] '[n] ) - => Vector t n -> Scalar t +normLPInf :: ( Ord (DataFrame t ('[] :: [k])) + , Num (DataFrame t ('[] :: [k])) + , SubSpace t '[n] '[] '[n] ) + => Vector t (n :: k) -> DataFrame t ('[] :: [k]) normLPInf = ewfoldr (max . abs) 0 -- | Minimum of absolute values -normLNInf :: ( Ord t, Num t , SubSpace t '[n] '[] '[n] ) - => Vector t n -> Scalar t -normLNInf x = ewfoldr (min . abs) (scalar . abs $ ixOff 0 x) x +normLNInf :: ( Ord (DataFrame t ('[] :: [k])) + , Num (DataFrame t ('[] :: [k])) + , SubSpace t '[n] '[] '[n] ) + => Vector t (n :: k) -> DataFrame t ('[] :: [k]) +normLNInf x = ewfoldr (min . abs) (abs $ indexOffset 0 x) x -- | Norm in Lp space -normLP :: ( Floating t , SubSpace t '[n] '[] '[n] ) - => Int -> Vector t n -> Scalar t +normLP :: (Floating (DataFrame t ('[] :: [k])), SubSpace t '[n] '[] '[n]) + => Int -> Vector t (n :: k) -> DataFrame t ('[] :: [k]) normLP i' = (**ri) . ewfoldr (\a -> (a**i +)) 0 where i = fromIntegral i' @@ -202,23 +211,15 @@ -- | Take a determinant of a matrix composed from two 2D vectors. -- Like a cross product in 2D. -det2 :: ( Num t, SubSpace t '[2] '[] '[2] ) +det2 :: ( Num t, PrimBytes t ) => Vector t 2 -> Vector t 2 -> Scalar t -det2 a b = (a ! 0 :* U) * (b ! 1 :* U) - - (a ! 1 :* U) * (b ! 0 :* U) +det2 (DF2 a0 a1) (DF2 b0 b1) = a0*b1 - b0*a1 -- | Cross product -cross :: ( Num t, SubSpace t '[3] '[] '[3] ) +cross :: ( Num t, PrimBytes t ) => Vector t 3 -> Vector t 3 -> Vector t 3 -cross a b = vec3 ( unScalar - $ (a ! 1 :* U) * (b ! 2 :* U) - - (a ! 2 :* U) * (b ! 1 :* U) ) - ( unScalar - $ (a ! 2 :* U) * (b ! 0 :* U) - - (a ! 0 :* U) * (b ! 2 :* U) ) - ( unScalar - $ (a ! 0 :* U) * (b ! 1 :* U) - - (a ! 1 :* U) * (b ! 0 :* U) ) +cross (DF3 a0 a1 a2) (DF3 b0 b1 b2) + = DF3 (a1*b2 - a2*b1) (a2*b0 - a0*b2) (a0*b1 - a1*b0) -- | Cross product for two vectors in 3D diff --git a/test/Numeric/Arbitraries.hs b/test/Numeric/Arbitraries.hs new file mode 100644 --- /dev/null +++ b/test/Numeric/Arbitraries.hs @@ -0,0 +1,614 @@ +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE ExistentialQuantification #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE FunctionalDependencies #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE PartialTypeSignatures #-} +{-# LANGUAGE PolyKinds #-} +{-# LANGUAGE QuasiQuotes #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UndecidableInstances #-} +{-# OPTIONS_GHC -fno-warn-orphans #-} +-- | Provide instance of Arbitrary for all DataFrame types. +-- Also, this module is an example of fancy type inference and DataFrame +-- traversals with monadic actions. +module Numeric.Arbitraries where + +import Test.QuickCheck + +#if !(MIN_VERSION_base(4,13,0)) +import Control.Monad.Fail +#endif +import Data.Int +import Data.Kind (Type) +import Data.List (inits, tails) +import Data.Semigroup hiding (All) +import Data.Word +import Language.Haskell.TH (ExpQ, Q, TypeQ) +import qualified Language.Haskell.TH as TH +import Numeric.DataFrame +import Numeric.Dimensions +import Numeric.Quaternion +import qualified Numeric.Tuple.Lazy as LT +import qualified Numeric.Tuple.Strict as ST + +-- | Maximum number of elements in SomeDims lists +maxDims :: Word +maxDims = 5 + +-- | Some tests are rather slow when we have too many elements +maxTotalDim :: Word +maxTotalDim = 400 + +-- | Maximum value of @Dim (XN _)@ +maxDimSize :: Word +maxDimSize = 30 + +-- | Odds of using `fromScalar` constructor instead of element-by-element. +-- Interpret as "one to fromScalarChanceFactor" +fromScalarChanceFactor :: Int +fromScalarChanceFactor = 5 + +concatDimsQ :: TypeQ -> TypeQ -> TypeQ +concatDimsQ as bs = [t| Concat $(as) $(bs) |] + +aFewPropVariants :: ExpQ -> ExpQ +aFewPropVariants = TH.listE . replicate 100 + +someElemTypeFQ :: Q TypeQ +someElemTypeFQ = TH.runIO . generate $ elements elemTypesF + +someElemTypeIQ :: Q TypeQ +someElemTypeIQ = TH.runIO . generate $ elements elemTypesI + +someElemTypeQ :: Q TypeQ +someElemTypeQ = TH.runIO . generate . elements $ elemTypesI ++ elemTypesF + +elemTypesF :: [TypeQ] +elemTypesF = [ [t|Float|], [t|Double|] ] + +elemTypesI :: [TypeQ] +elemTypesI = + [ [t|Int|], [t|Int8|], [t|Int16|], [t|Int32|], [t|Int64|] + , [t|Word|], [t|Word8|], [t|Word16|], [t|Word32|], [t|Int64|] + ] + + +-- | Generate type (n :: Nat) at compile time +someDimNQ :: Q TypeQ +someDimNQ = TH.litT . TH.numTyLit . toInteger + <$> TH.runIO (generate (chooseDim 1)) + +-- | Generate type (n :: XNat) at compile time +someDimXQ :: Q TypeQ +someDimXQ = fmap pure $ do + d <- TH.litT . TH.numTyLit . toInteger + <$> TH.runIO (generate (chooseDim 1)) + TH.runIO (generate arbitrary) >>= + \isXN -> if isXN then [t|XN $(d)|] else [t|N $(d)|] + +-- | Generate type (N n) at compile time +someDimXfixedQ :: Q TypeQ +someDimXfixedQ = fmap pure $ do + d <- TH.litT . TH.numTyLit . toInteger + <$> TH.runIO (generate (chooseDim 1)) + [t| N $(d) |] + +-- | Generate type (ns :: [k]) at compile time +someDimsQ :: Q TypeQ +someDimsQ = do + isXNat <- TH.runIO (generate arbitrary) + if isXNat then someDimsXQ + else someDimsNQ + +-- | Generate type (ns :: [Nat]) at compile time +someDimsNQ :: Q TypeQ +someDimsNQ = fmap pure $ TH.runIO (removeDims <$> generate arbitrary) + >>= \(SomeDims ds) -> wordsToDimsNQ (listDims ds) + +-- | Generate type (Map 'N (ns :: [Nat])) at compile time +someDimsXQfixed :: Q TypeQ +someDimsXQfixed = fmap pure $ TH.runIO (removeDims <$> generate arbitrary) + >>= \(SomeDims ds) -> wordsToDimsXQfixed (listDims ds) + +-- | Generate type (Map 'N (ns :: [Nat])) at compile time +someDimsXQ :: Q TypeQ +someDimsXQ = fmap pure $ TH.runIO (removeDims <$> generate arbitrary) + >>= \(SomeDims ds) -> wordsToDimsXQ(listDims ds) + +-- | Generate type (ns :: [k]) at compile time +someLenAsBsQ :: Int -> Q (TypeQ, TypeQ) +someLenAsBsQ len = do + isXNat <- TH.runIO (generate arbitrary) + if isXNat then someLenAsBsXQ len + else someLenAsBsNQ len + +-- | Generate type (ns :: [k]) at compile time +someAsBsQ :: Q (TypeQ, TypeQ) +someAsBsQ = do + isXNat <- TH.runIO (generate arbitrary) + if isXNat then someAsBsXQ + else someAsBsNQ + +-- | Generate type (ns :: [k]) at compile time +someAsBsQfixed :: Q (TypeQ, TypeQ) +someAsBsQfixed = do + isXNat <- TH.runIO (generate arbitrary) + if isXNat then someAsBsXQfixed + else someAsBsNQ + +-- | Generate type (as :: [Nat], bs :: [Nat]) at compile time, +-- given a fixed length of the first dims (as). +someLenAsBsNQ :: Int -> Q (TypeQ, TypeQ) +someLenAsBsNQ len = do + (asw, bsw) <- TH.runIO (generate (genConcatDims len)) + as <- wordsToDimsNQ asw + bs <- wordsToDimsNQ bsw + return (pure as, pure bs) + +-- | Generate type (as :: [XNat], Map 'N bs :: [XNat]) at compile time, +-- given a fixed length of the first dims (as). +someLenAsBsXQ :: Int -> Q (TypeQ, TypeQ) +someLenAsBsXQ len = do + (asw, bsw) <- TH.runIO (generate (genConcatDims len)) + as <- wordsToDimsXQ asw + bs <- wordsToDimsXQfixed bsw + return (pure as, pure bs) + +-- | Generate type (Map 'N as :: [XNat], Map 'N bs :: [XNat]) at compile time, +-- given a fixed length of the first dims (as). +someLenAsBsXQfixed :: Int -> Q (TypeQ, TypeQ) +someLenAsBsXQfixed len = do + (asw, bsw) <- TH.runIO (generate (genConcatDims len)) + as <- wordsToDimsXQfixed asw + bs <- wordsToDimsXQfixed bsw + return (pure as, pure bs) + +-- | Generate type (as :: [Nat], bs :: [Nat]) at compile time. +someAsBsNQ :: Q (TypeQ, TypeQ) +someAsBsNQ = TH.runIO (generate (choose (0, fromIntegral maxDims - 1))) + >>= someLenAsBsNQ + +-- | Generate type (as :: [XNat], Map 'N bs :: [XNat]) at compile time. +someAsBsXQ :: Q (TypeQ, TypeQ) +someAsBsXQ = TH.runIO (generate (choose (0, fromIntegral maxDims - 1))) + >>= someLenAsBsXQ + +-- | Generate type (Map 'N as :: [XNat], Map 'N bs :: [XNat]) at compile time. +someAsBsXQfixed:: Q (TypeQ, TypeQ) +someAsBsXQfixed = TH.runIO (generate (choose (0, fromIntegral maxDims - 1))) + >>= someLenAsBsXQfixed + + +wordsToDimsNQ :: [Word] -> TypeQ +wordsToDimsNQ = (\ns -> [t| ($(ns) :: [Nat]) |]) . foldr f TH.promotedNilT + where + f :: Word -> TypeQ -> TypeQ + f w l = let d = TH.litT . TH.numTyLit $ toInteger w + in TH.promotedConsT `TH.appT` d `TH.appT` l + +wordsToDimsXQfixed :: [Word] -> TypeQ +wordsToDimsXQfixed = (\ns -> [t| ($(ns) :: [XNat]) |]) . foldr f TH.promotedNilT + where + f :: Word -> TypeQ -> TypeQ + f w l = let d = TH.litT . TH.numTyLit $ toInteger w + in TH.promotedConsT `TH.appT` [t| N $(d) |] `TH.appT` l + +wordsToDimsXQ :: [Word] -> TypeQ +wordsToDimsXQ = (\ns -> [t| ($(ns) :: [XNat]) |]) . foldr f TH.promotedNilT + where + f :: Word -> TypeQ -> TypeQ + f w l = let d = TH.litT . TH.numTyLit $ toInteger w + xd = TH.runIO (generate arbitrary) >>= + \isXN -> if isXN then [t|XN $(d)|] else [t|N $(d)|] + in TH.promotedConsT `TH.appT` xd `TH.appT` l + +-- | Generate two lists of Dims as Words, such that the combined list +-- is under the maxTotalDim limit, but the first list has exactly the given +-- number of dims. +genConcatDims :: Int -> Gen ([Word], [Word]) +genConcatDims len + = arbitrary >>= (\(SomeDims xs) -> checkLen (listDims xs)) . removeDims + where + checkLen :: [Word] -> Gen ([Word], [Word]) + checkLen xs = + let td = product xs + (as, bs) = splitAt len xs + l = length as + in if l == len then pure (as, bs) + else (,) <$> addDims (len - l) (maxTotalDim `quot` td) as + <*> pure bs + addDims :: Int -> Word -> [Word] -> Gen [Word] + addDims 0 _ xs = pure xs + addDims n lim xs = do + x <- choose (1, max 1 (min lim maxDimSize)) + addDims (n-1) (lim `quot` x) (x : xs) + + +-- | Remove dims from a dim list until its totalDim is less than maxTotalDim +removeDims :: SomeDims -> SomeDims +removeDims = removeDimsAbove maxTotalDim + +-- | Remove dims from a dim list until its totalDim is less than a given value +removeDimsAbove :: Word -> SomeDims -> SomeDims +removeDimsAbove _ (SomeDims U) = SomeDims U +removeDimsAbove z (SomeDims nns@(_ :* ns)) + | totalDim nns > z = removeDimsAbove z (SomeDims ns) + | otherwise = SomeDims nns + +-- | Reduce individual XN-dims untils its totalDim is less than maxTotalDim +reduceDims :: (All KnownDimType xns, BoundedDims xns) + => Dims (xns :: [XNat]) -> Dims (xns :: [XNat]) +reduceDims = reduceDims' 1 + +reduceDims' :: (All KnownDimType xns, BoundedDims xns) + => Word -> Dims (xns :: [XNat]) -> Dims (xns :: [XNat]) +reduceDims' _ U = U +reduceDims' l nns@(n :* ns) + | totalDim nns * l <= maxTotalDim = nns + | otherwise = case n of + Dn d -> n :* reduceDims' (l * dimVal d) ns + (Dx d :: Dim xn) -> case compareDim (dimBound @xn) D2 of + SLT -> Dx D2 :* reduceDims' (l * 2) ns + SEQ -> Dx D2 :* reduceDims' (l * 2) ns + SGT -> n :* reduceDims' (l * dimVal d) ns + +-- | Most of the time, we assume the error is proportional to the maginutude of +-- the biggest element. +maxElem :: (SubSpace t ds '[] ds, Ord t, Num t) + => DataFrame t (ds :: [Nat]) -> Scalar t +maxElem = ewfoldr' (max . abs) 0 + +rotateList :: [a] -> [[a]] +rotateList xs = init (zipWith (++) (tails xs) (inits xs)) + +class (RealFloatExtras t, Show t) => Approx t x | x -> t where + -- | Check if two values are approximately equal + approxEq :: t -- ^ Extra multiplier constant + -> x -> x -> Property + +-- | Check if two values are approximately equal. +(=~=) :: Approx t x => x -> x -> Property +(=~=) = approxEq 1 +infix 4 =~= + +instance Approx Double Double where + approxEq c a b = counterexample + (unlines + [ " Double approxEq failed:" + , " values: " ++ show (a, b) + , " error: " ++ show err + , " tolerance: " ++ show tol + ] + ) $ err <= tol + where + err = abs (a - b) + mel = abs a `max` abs b `max` 1 + tol = M_EPS*mel*c + +instance Approx Float Float where + approxEq c a b = counterexample + (unlines + [ " Double approxEq failed:" + , " values: " ++ show (a, b) + , " error: " ++ show err + , " tolerance: " ++ show tol + ] + ) $ err <= tol + where + err = abs (a - b) + mel = abs a `max` abs b `max` 1 + tol = M_EPS*mel*c + +instance MonadFail Gen where fail = error + +instance (Quaternion t, Arbitrary t, Num t) => Arbitrary (Quater t) where + arbitrary = sequence + [ Quater <$> arbitrary <*> arbitrary <*> arbitrary <*> arbitrary + , Quater <$> arbitrary <*> arbitrary <*> arbitrary <*> arbitrary + , Quater <$> arbitrary <*> arbitrary <*> arbitrary <*> arbitrary + , Quater <$> arbitrary <*> arbitrary <*> arbitrary <*> arbitrary + , Quater <$> arbitrary <*> pure 0 <*> pure 0 <*> pure 0 + , Quater <$> pure 0 <*> arbitrary <*> pure 0 <*> pure 0 + , Quater <$> pure 0 <*> pure 0 <*> arbitrary <*> pure 0 + , Quater <$> pure 0 <*> pure 0 <*> pure 0 <*> arbitrary + , Quater <$> arbitrary <*> arbitrary <*> pure 0 <*> pure 0 + , Quater <$> arbitrary <*> pure 0 <*> arbitrary <*> pure 0 + , Quater <$> arbitrary <*> pure 0 <*> pure 0 <*> arbitrary + , Quater <$> pure 0 <*> arbitrary <*> arbitrary <*> pure 0 + , Quater <$> pure 0 <*> arbitrary <*> pure 0 <*> arbitrary + , Quater <$> pure 0 <*> pure 0 <*> arbitrary <*> arbitrary + , Quater <$> arbitrary <*> arbitrary <*> arbitrary <*> pure 0 + , Quater <$> arbitrary <*> arbitrary <*> pure 0 <*> arbitrary + , Quater <$> arbitrary <*> pure 0 <*> arbitrary <*> arbitrary + , Quater <$> pure 0 <*> arbitrary <*> arbitrary <*> arbitrary + ] >>= elements + + shrink (Quater x y z t) + -- shrink either real or the whole imaginary part + = ($) <$> zipWith3 Quater (shrink x) (shrink y) (shrink z) <*> shrink t + +instance (RealFloatExtras t, Show t, Quaternion t) => Approx t (Quater t) where + approxEq c a b = counterexample + (unlines + [ " Quaternion approxEq failed:" + , " max elem: " ++ show mel + , " error: " ++ show err + , " tolerance: " ++ show tol + ] + ) $ err <= tol + where + v1 = toVec4 a + v2 = toVec4 b + err = unScalar $ maxElem (v1 - v2) + mel = unScalar $ 1 `max` maxElem v1 `max` maxElem v2 + tol = M_EPS*mel*c + +instance (Arbitrary t, PrimBytes t, Num t, Ord t, Dimensions ds) + => Arbitrary (DataFrame t (ds :: [Nat])) where + arbitrary = do + full <- (1 < ) <$> choose (1, fromScalarChanceFactor) + zeroChance <- choose (0, 8) + if full -- I want to check fromScalar code path sometimes + then arbitrary >>= elementWise @_ @ds @'[] (f zeroChance) . ewgen . scalar + else fromScalar . scalar <$> arbitrary + where + f :: (Arbitrary a, Num a) => Double -> Scalar a -> Gen (Scalar a) + f zeroChance _ = do + dice <- (zeroChance >=) <$> choose (0, 10) + if dice + then return 0 + else scalar <$> arbitrary + shrink df + | Just (Max ma) <- ewfoldMap @t @ds @'[] @ds + ((\x -> if x == 0 then Nothing else Just (Max x)) . abs) df + = [ ewmap (\x -> if abs x == ma then 0 else x) df + , ewmap (scalar . withAbs . unScalar) df + ] + | otherwise = [] + where + withAbs :: t -> t + withAbs x + | abs x <= 1 = 0 + | otherwise = signum x * closest2 (abs x) 1 + closest2 :: t -> t -> t + closest2 x b = if x <= b * 2 then b else closest2 x (b*2) + +instance (RealFloatExtras t, Show t, Dimensions ds) + => Approx (DataFrame t ('[] :: [Nat])) (DataFrame t (ds :: [Nat])) where + approxEq (S t) = approxEqDF t + +instance (RealFloatExtras t, Show t) + => Approx (DataFrame t ('[] :: [XNat])) (DataFrame t (ds :: [XNat])) where + approxEq (XFrame (S t)) (XFrame a) (XFrame b) = case sameDims da db of + Just Dict -> approxEqDF t a b + Nothing -> counterexample + (unlines + [ " DataFrame approxEq failed [XNat] due to different dims:" + , " dims a: " ++ show da + , " dims b: " ++ show db + ] + ) False + where + da = dims `inSpaceOf` a + db = dims `inSpaceOf` b + +approxEqDF :: + forall t (ds :: [Nat]) + . ( Dimensions ds + , Show t, RealFloatExtras t + , KnownBackend t ds) + => t -> DataFrame t ds -> DataFrame t (ds :: [Nat]) -> Property +approxEqDF c a b + | U <- dims @ds = counterexample + (unlines + [ " Scalar approxEq failed:" + , " values: " ++ show (unScalar a, unScalar b) + , " error: " ++ show err + , " tolerance: " ++ show tol + ] + ) $ err <= tol + | otherwise = counterexample + (unlines + [ " DataFrame approxEq failed:" + , " max elem: " ++ show mel + , " error: " ++ show err + , " tolerance: " ++ show tol + ] + ) $ err <= tol + where + err = unScalar $ maxElem (a - b) + mel = unScalar $ maxElem a `max` maxElem b `max` 1 + tol = M_EPS*mel*c + +instance ( All Arbitrary ts, All PrimBytes ts, All Num ts, All Ord ts + , RepresentableList ts, Dimensions ds) + => Arbitrary (DataFrame ts (ds :: [Nat])) where + -- We create arbitrary MultiFrame by combining several SingleFrames. + -- SingleFrames are "variables" or "columns" of a MultiFrame that are + -- independent byte arrays bounded by a common dimensions type signature. + arbitrary = -- Use RepresentableList to find out how many columns are there. + case tList @ts of + -- Zero columns, empty MultiFrame + U -> return Z + -- Cons-like construction. + -- Note, pattern matching TypeList brings RepresentableList evidence + -- for Tail ts. + _ :* (TypeList :: TypeList ts') -> do + at <- arbitrary + ats' <- arbitrary @(DataFrame ts' ds) + return (at :*: ats') + -- MultiFrame is a newtype wrapper on a TypedList. + -- Thus, we can always recover RepresentableList ts by using function @types@ + shrink (at :*: ats@(MultiFrame ats')) + | TypeList <- types ats' + = (:*:) <$> shrink at <*> shrink ats + shrink _ = [] + +-- | Generate a random word not smaller than the given number, and not bigger +-- than maxDimSize. Also tend to generate smaller values more often to have +-- a chance to pack more dims within maxTotalDim. +chooseDim :: Word -> Gen Word +chooseDim lowLim = f <$> choose (0, maxDimSize) + where + l = realToFrac lowLim :: Double + f :: Word -> Word + f w = let x = realToFrac w :: Double + n = realToFrac maxDimSize :: Double + in round $ l + max 0 (n - l) * (x / n) ** 3 + +instance KnownDim a => Arbitrary (Dim (N a)) where + arbitrary = return $ Dn (dim @a) + shrink _ = [] + +instance KnownDim m => Arbitrary (Dim (XN m)) where + arbitrary = do + dimN <- chooseDim $ dimVal' @m + case constrainDim @(XN m) (someDimVal dimN) of + Nothing -> error "impossible argument" + Just d -> return d + shrink _ = [] + +instance Arbitrary SomeDims where + arbitrary = do + dimN <- choose (0, maxDims) :: Gen Word + wdims <- mapM (\_ -> chooseDim 1) [1..dimN] + return $ someDimsVal wdims + shrink (SomeDims U) = [] + shrink (SomeDims (_ :* ds)) = [SomeDims ds] + +instance Arbitrary (Dims '[]) where + arbitrary = return U + shrink _ = [] + +instance (KnownDim n, Arbitrary (Dims xs)) => Arbitrary (Dims (N n ': xs)) where + arbitrary = (:*) <$> arbitrary <*> arbitrary + shrink _ = [] + +instance (KnownDim m, Arbitrary (Dims xs)) => Arbitrary (Dims (XN m ': xs)) where + arbitrary = (:*) <$> arbitrary <*> arbitrary + shrink _ = [] + +instance (Arbitrary t, PrimBytes t, Num t, Ord t) + => Arbitrary (SomeDataFrame t) where + arbitrary = do + SomeDims (Dims :: Dims ds) <- removeDims <$> arbitrary + SomeDataFrame <$> arbitrary @(DataFrame t ds) + shrink (SomeDataFrame df) = SomeDataFrame <$> shrink df + +-- All same as above, just change constraints a bit +instance ( All Arbitrary ts, All PrimBytes ts, All Num ts, All Ord ts + , RepresentableList ts) + => Arbitrary (SomeDataFrame ts) where + arbitrary = do + SomeDims ds <- removeDims <$> arbitrary + case ds of + (Dims :: Dims ds) -> case inferKnownBackend @ts @ds of + Dict -> SomeDataFrame <$> arbitrary @(DataFrame ts ds) + shrink (SomeDataFrame df) = SomeDataFrame <$> shrink df + + +instance ( Arbitrary t, PrimBytes t, Num t, Ord t + , Arbitrary (Dims xs), All KnownDimType xs, BoundedDims xs) + => Arbitrary (DataFrame t (xs :: [XNat])) where + arbitrary = do + XDims (Dims :: Dims ds) <- reduceDims <$> arbitrary @(Dims xs) + XFrame <$> arbitrary @(DataFrame t ds) + shrink (XFrame df) = XFrame <$> shrink df + +instance ( All Arbitrary ts, All PrimBytes ts, All Num ts, All Ord ts + , RepresentableList ts + , Arbitrary (Dims xs), All KnownDimType xs, BoundedDims xs) + => Arbitrary (DataFrame ts (xs :: [XNat])) where + arbitrary = do + ds <- reduceDims <$> arbitrary @(Dims xs) + case ds of + XDims (Dims :: Dims ds) -> case inferKnownBackend @ts @ds of + Dict -> XFrame <$> arbitrary @(DataFrame ts ds) + shrink (XFrame df) = XFrame <$> shrink df + + +instance KnownDim n => Arbitrary (Idx (n :: Nat)) where + arbitrary = Idx <$> choose (minBound, dimVal' @n - 1) + +instance KnownDim n => Arbitrary (Idx (N n)) where + arbitrary = Idx <$> choose (minBound, dimVal' @n - 1) + +instance KnownDim n => Arbitrary (Idx (XN n)) where + arbitrary = Idx <$> choose (minBound, dimVal' @n - 1) + +instance Dimensions ns => Arbitrary (Idxs (ns :: [Nat])) where + arbitrary = go (dims @ns) + where + go :: forall (bs :: [Nat]) . Dims bs -> Gen (Idxs bs) + go U = pure U + go (D :* bs) = (:*) <$> arbitrary <*> go bs + +instance (BoundedDims ns, All KnownDimType ns) => Arbitrary (Idxs (ns :: [XNat])) where + arbitrary = go (minimalDims @ns) + where + go :: forall (bs :: [XNat]) + . (BoundedDims bs, All KnownDimType bs) => Dims bs -> Gen (Idxs bs) + go U = pure U + go (Dn D :* bs) = (:*) <$> arbitrary <*> go bs + go (Dx D :* bs) = (:*) <$> arbitrary <*> go bs + + +instance (RepresentableList xs, All Arbitrary xs) => Arbitrary (ST.Tuple xs) where + arbitrary = go (tList @xs) + where + go :: forall (bs :: [Type]) + . All Arbitrary bs + => TypeList bs -> Gen (ST.Tuple bs) + go U = pure U + go (_ :* bs) = (ST.:$) <$> arbitrary <*> go bs + +instance (RepresentableList xs, All Arbitrary xs) => Arbitrary (LT.Tuple xs) where + arbitrary = go (tList @xs) + where + go :: forall (bs :: [Type]) + . All Arbitrary bs + => TypeList bs -> Gen (LT.Tuple bs) + go U = pure U + go (_ :* bs) = (LT.:$) <$> arbitrary <*> go bs + +data AnyMatrix +data NonSingular + +data SomeSquareMatrix prop t + = forall (n :: Nat) + . (KnownDim n, KnownBackend t '[n], KnownBackend t '[n, n]) + => SSM (DataFrame t '[n,n]) + +instance (Show t, PrimBytes t) => Show (SomeSquareMatrix prop t) where + show (SSM df) = show df + +instance (Arbitrary t, PrimBytes t, Num t, Ord t) + => Arbitrary (SomeSquareMatrix AnyMatrix t) where + arbitrary = do + Dx (D :: Dim n) <- arbitrary @(Dim (XN 2)) + SSM <$> arbitrary @(DataFrame t '[n,n]) + shrink (SSM df)= SSM <$> shrink df + +instance (Arbitrary t, PrimBytes t, Num t, Ord t) + => Arbitrary (SomeSquareMatrix NonSingular t) where + arbitrary = do + SSM (someMat :: DataFrame t '[n, n]) <- arbitrary @(SomeSquareMatrix AnyMatrix t) + -- https://en.wikipedia.org/wiki/Diagonally_dominant_matrix + return . SSM $ + iwmap @t @'[n] @'[n] @'[n,n] + @t @'[n] @'[n,n] + ( \i v -> + let s = ewfoldl (\a -> (a +) . abs) 1 v + in update i s v + ) someMat + shrink _ = [] diff --git a/test/Numeric/DataFrame/Arbitraries.hs b/test/Numeric/DataFrame/Arbitraries.hs deleted file mode 100644 --- a/test/Numeric/DataFrame/Arbitraries.hs +++ /dev/null @@ -1,274 +0,0 @@ -{-# LANGUAGE AllowAmbiguousTypes #-} -{-# LANGUAGE ConstraintKinds #-} -{-# LANGUAGE DataKinds #-} -{-# LANGUAGE ExistentialQuantification #-} -{-# LANGUAGE FlexibleContexts #-} -{-# LANGUAGE FlexibleInstances #-} -{-# LANGUAGE GADTs #-} -{-# LANGUAGE PartialTypeSignatures #-} -{-# LANGUAGE PolyKinds #-} -{-# LANGUAGE RankNTypes #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeApplications #-} -{-# LANGUAGE TypeOperators #-} -{-# LANGUAGE UndecidableInstances #-} -{-# OPTIONS_GHC -fno-warn-orphans #-} --- | Provide instance of Arbitrary for all DataFrame types. --- Also, this module is an example of fancy type inference and DataFrame --- traversals with monadic actions. -module Numeric.DataFrame.Arbitraries where - -import Test.QuickCheck - -import Data.Kind (Type) -import Data.Semigroup hiding (All) -import Numeric.DataFrame -import Numeric.Dimensions -import Numeric.Quaternion -import qualified Numeric.Tuple.Lazy as LT -import qualified Numeric.Tuple.Strict as ST - - -maxDims :: Word -maxDims = 5 - -maxDimSize :: Word -maxDimSize = 7 - -fromScalarChanceFactor :: Int -fromScalarChanceFactor = 5 - - -instance (Quaternion t, Arbitrary t, Num t) => Arbitrary (Quater t) where - arbitrary = sequence - [ Quater <$> arbitrary <*> arbitrary <*> arbitrary <*> arbitrary - , Quater <$> arbitrary <*> arbitrary <*> arbitrary <*> arbitrary - , Quater <$> arbitrary <*> arbitrary <*> arbitrary <*> arbitrary - , Quater <$> arbitrary <*> arbitrary <*> arbitrary <*> arbitrary - , Quater <$> arbitrary <*> pure 0 <*> pure 0 <*> pure 0 - , Quater <$> pure 0 <*> arbitrary <*> pure 0 <*> pure 0 - , Quater <$> pure 0 <*> pure 0 <*> arbitrary <*> pure 0 - , Quater <$> pure 0 <*> pure 0 <*> pure 0 <*> arbitrary - , Quater <$> arbitrary <*> arbitrary <*> pure 0 <*> pure 0 - , Quater <$> arbitrary <*> pure 0 <*> arbitrary <*> pure 0 - , Quater <$> arbitrary <*> pure 0 <*> pure 0 <*> arbitrary - , Quater <$> pure 0 <*> arbitrary <*> arbitrary <*> pure 0 - , Quater <$> pure 0 <*> arbitrary <*> pure 0 <*> arbitrary - , Quater <$> pure 0 <*> pure 0 <*> arbitrary <*> arbitrary - , Quater <$> arbitrary <*> arbitrary <*> arbitrary <*> pure 0 - , Quater <$> arbitrary <*> arbitrary <*> pure 0 <*> arbitrary - , Quater <$> arbitrary <*> pure 0 <*> arbitrary <*> arbitrary - , Quater <$> pure 0 <*> arbitrary <*> arbitrary <*> arbitrary - ] >>= elements - - shrink (Quater x y z t) - -- shrink either real or the whole imaginary part - = ($) <$> zipWith3 Quater (shrink x) (shrink y) (shrink z) <*> shrink t - - - -instance (Arbitrary t, PrimBytes t, Num t, Ord t, Dimensions ds) - => Arbitrary (DataFrame t (ds :: [Nat])) where - arbitrary - | -- First, we need to find out exact array implementation to use - -- inside this DataFrame. - -- We need to do that whenever exact value of ds is not known - Dict <- inferKnownBackend @t @ds - -- After that, GHC can infer all necessary fancy things like SubSpace - -- to do complex operations on sub-dimensions of a DataFrame. - -- - -- Note, we could put SubSpace into constraints of this instance as well. - -- That would render the above lines unnecessary, but would make - -- inference more difficult later. - = do - full <- (1 < ) <$> choose (1, fromScalarChanceFactor) - if full -- I want to check fromScalar code path sometimes - then arbitrary >>= elementWise @_ @ds @'[] f . ewgen . scalar - else fromScalar . scalar <$> arbitrary - where - f :: Arbitrary a => Scalar a -> Gen (Scalar a) - f _ = scalar <$> arbitrary - shrink df - | Dict <- inferKnownBackend @t @ds - , mma <- ewfoldMap @t @ds @'[] @ds - ((\x -> if x == 0 then Nothing else Just (Max x)) . abs) df - = case mma of - Nothing - -> [] -- all-zero is the most primitive DF possible - Just (Max ma) - -> [ ewmap (\x -> if abs x == ma then 0 else x) df - , ewmap (scalar . withAbs . unScalar) df - ] - where - withAbs :: t -> t - withAbs x - | abs x <= 1 = 0 - | otherwise = signum x * closest2 (abs x) 1 - closest2 :: t -> t -> t - closest2 x b = if x <= b * 2 then b else closest2 x (b*2) - -instance ( All Arbitrary ts, All PrimBytes ts, All Num ts, All Ord ts - , RepresentableList ts, Dimensions ds) - => Arbitrary (DataFrame ts (ds :: [Nat])) where - -- We create arbitrary MultiFrame by combining several SingleFrames. - -- SingleFrames are "variables" or "columns" of a MultiFrame that are - -- independent byte arrays bounded by a common dimensions type signature. - arbitrary = -- Use RepresentableList to find out how many columns are there. - case tList @_ @ts of - -- Zero columns, empty MultiFrame - U -> return Z - -- Cons-like construction. - -- Note, pattern matching TypeList brings RepresentableList evidence - -- for Tail ts. - _ :* (TypeList :: TypeList ts') -> do - at <- arbitrary - ats' <- arbitrary @(DataFrame ts' ds) - return (at :*: ats') - -- MultiFrame is a newtype wrapper on a TypedList. - -- Thus, we can always recover RepresentableList ts by using function @types@ - shrink (at :*: ats@(MultiFrame ats')) - | TypeList <- types ats' - = (:*:) <$> shrink at <*> shrink ats - shrink _ = [] - - -instance KnownDim a => Arbitrary (Dim (N a)) where - arbitrary = return $ Dn (dim @a) - shrink _ = [] - -instance KnownDim m => Arbitrary (Dim (XN m)) where - arbitrary = do - dimN <- choose (dimVal' @m, maxDims) - case constrainDim @XNat @(XN m) (someDimVal dimN) of - Nothing -> error "impossible argument" - Just d -> return d - shrink _ = [] - -instance Arbitrary SomeDims where - arbitrary = do - dimN <- choose (0, maxDims) :: Gen Word - wdims <- mapM (\_ -> choose (2, maxDimSize) :: Gen Word) [1..dimN] - return $ someDimsVal wdims - shrink (SomeDims U) = [] - shrink (SomeDims (_ :* ds)) = [SomeDims ds] - -instance Arbitrary (Dims '[]) where - arbitrary = return U - shrink _ = [] - -instance (KnownDim n, Arbitrary (Dims xs)) => Arbitrary (Dims (N n ': xs)) where - arbitrary = (:*) <$> arbitrary <*> arbitrary - shrink _ = [] - -instance (KnownDim m, Arbitrary (Dims xs)) => Arbitrary (Dims (XN m ': xs)) where - arbitrary = (:*) <$> arbitrary <*> arbitrary - shrink _ = [] - -instance (Arbitrary t, PrimBytes t, Num t, Ord t) - => Arbitrary (SomeDataFrame t) where - arbitrary = do - -- Generate random dimension list - SomeDims ds <- arbitrary - -- and pattern-match against it with Dims pattern. - -- This gives Dimensions ds evidence immediately. - case ds of - -- We also need to figure out an array implementation... - (Dims :: Dims ds) -> case inferKnownBackend @t @ds of - -- ... and generating a random DataFrame becomes a one-liner - Dict -> SomeDataFrame <$> arbitrary @(DataFrame t ds) - shrink (SomeDataFrame df) = SomeDataFrame <$> shrink df - --- All same as above, just change constraints a bit -instance ( All Arbitrary ts, All PrimBytes ts, All Num ts, All Ord ts - , RepresentableList ts) - => Arbitrary (SomeDataFrame ts) where - arbitrary = do - SomeDims ds <- arbitrary - case ds of - (Dims :: Dims ds) -> case inferKnownBackend @ts @ds of - Dict -> SomeDataFrame <$> arbitrary @(DataFrame ts ds) - shrink (SomeDataFrame df) = SomeDataFrame <$> shrink df - -instance ( Arbitrary t, PrimBytes t, Num t, Ord t - , Arbitrary (Dims xs), All KnownXNatType xs) - => Arbitrary (DataFrame t (xs :: [XNat])) where - arbitrary = do - ds <- arbitrary @(Dims xs) - case ds of - XDims (_ :: Dims ds) -> case inferKnownBackend @t @ds of - Dict -> XFrame <$> arbitrary @(DataFrame t ds) - shrink (XFrame df) = XFrame <$> shrink df - -instance ( All Arbitrary ts, All PrimBytes ts, All Num ts, All Ord ts - , RepresentableList ts - , Arbitrary (Dims xs), All KnownXNatType xs) - => Arbitrary (DataFrame ts (xs :: [XNat])) where - arbitrary = do - ds <- arbitrary @(Dims xs) - case ds of - XDims (_ :: Dims ds) -> case inferKnownBackend @ts @ds of - Dict -> XFrame <$> arbitrary @(DataFrame ts ds) - shrink (XFrame df) = XFrame <$> shrink df - - -instance KnownDim n => Arbitrary (Idx n) where - arbitrary = elements [0..] - -instance Dimensions ns => Arbitrary (Idxs ns) where - arbitrary = go (Dims @ns) - where - go :: forall (bs :: [Nat]) . Dims bs -> Gen (Idxs bs) - go U = pure U - go (D :* bs) = (:*) <$> arbitrary <*> go bs - -instance (RepresentableList xs, All Arbitrary xs) => Arbitrary (ST.Tuple xs) where - arbitrary = go (tList @Type @xs) - where - go :: forall (bs :: [Type]) - . All Arbitrary bs - => TypeList bs -> Gen (ST.Tuple bs) - go U = pure U - go (_ :* bs) = (ST.:$) <$> arbitrary <*> go bs - -instance (RepresentableList xs, All Arbitrary xs) => Arbitrary (LT.Tuple xs) where - arbitrary = go (tList @Type @xs) - where - go :: forall (bs :: [Type]) - . All Arbitrary bs - => TypeList bs -> Gen (LT.Tuple bs) - go U = pure U - go (_ :* bs) = (LT.:$) <$> arbitrary <*> go bs - -data AnyMatrix -data NonSingular - -data SomeSquareMatrix prop t - = forall (n :: Nat) - . (KnownDim n, KnownBackend t '[n], KnownBackend t '[n, n]) - => SSM (DataFrame t '[n,n]) - -instance (Show t, PrimBytes t) => Show (SomeSquareMatrix prop t) where - show (SSM df) = show df - -instance (Arbitrary t, PrimBytes t, Num t, Ord t) - => Arbitrary (SomeSquareMatrix AnyMatrix t) where - arbitrary = do - dx <- arbitrary @(Dim (XN 2)) - case dx of - Dx (D :: Dim n) -> case inferKnownBackend @t @'[n] of - Dict -> SSM <$> arbitrary @(DataFrame t '[n,n]) - shrink (SSM df)= SSM <$> shrink df - -instance (Arbitrary t, PrimBytes t, Num t, Ord t) - => Arbitrary (SomeSquareMatrix NonSingular t) where - arbitrary = do - SSM (someMat :: DataFrame t '[n, n]) <- arbitrary @(SomeSquareMatrix AnyMatrix t) - -- https://en.wikipedia.org/wiki/Diagonally_dominant_matrix - return . SSM $ - iwmap @t @'[n] @'[n] @'[n,n] - @t @'[n] @'[n,n] - ( \i v -> - let s = ewfoldl (\a -> (a +) . abs) 1 v - in update i s v - ) someMat - shrink _ = [] diff --git a/test/Numeric/DataFrame/BasicTest.hs b/test/Numeric/DataFrame/BasicTest.hs --- a/test/Numeric/DataFrame/BasicTest.hs +++ b/test/Numeric/DataFrame/BasicTest.hs @@ -21,12 +21,11 @@ {-# LANGUAGE TypeApplications #-} {-# LANGUAGE TypeOperators #-} {-# LANGUAGE UndecidableInstances #-} -{-# OPTIONS_GHC -fno-warn-overlapping-patterns #-} module Numeric.DataFrame.BasicTest (runTests) where +import Numeric.Arbitraries () import Numeric.DataFrame -import Numeric.DataFrame.Arbitraries () import qualified Numeric.ProductOrd.NonTransitive as NonTransitive import qualified Numeric.ProductOrd.Partial as Partial import Test.QuickCheck diff --git a/test/Numeric/DataFrame/SubSpaceTest.hs b/test/Numeric/DataFrame/SubSpaceTest.hs --- a/test/Numeric/DataFrame/SubSpaceTest.hs +++ b/test/Numeric/DataFrame/SubSpaceTest.hs @@ -1,55 +1,134 @@ {-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE GADTs #-} {-# LANGUAGE PolyKinds #-} +{-# LANGUAGE QuasiQuotes #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TemplateHaskell #-} {-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} {-# LANGUAGE TypeOperators #-} +-- overwise GHC dies trying to optimize 100 equal funs in each test +{-# OPTIONS_GHC -O0 #-} module Numeric.DataFrame.SubSpaceTest (runTests) where +import Data.Typeable (typeOf) +import Numeric.Arbitraries import Numeric.DataFrame -import Numeric.DataFrame.Arbitraries () import Numeric.Dimensions import Test.QuickCheck -type SFull = '[2,5,4,3,7] -type SPref = '[2,5,4] -type SSuff = '[3,7] +prop_IndexFrame0 :: Property +prop_IndexFrame0 = property $ oneof $( aFewPropVariants $ do + t <- someElemTypeQ + (_, ds) <- someLenAsBsQ 0 + [e| do + df <- arbitrary @(DataFrame $t $ds) + return $ conjoin + [ index U df === df + , indexOffset 0 df === df ] + |] + ) -prop_IndexDimMax :: DataFrame Int SPref -> DataFrame Int SSuff -> Bool -prop_IndexDimMax x y = - (z ! (maxBound `inSpaceOf` x)) == y - where - z = ewgen y :: DataFrame Int SFull +prop_IndexFrame1 :: Property +prop_IndexFrame1 = property $ oneof $( aFewPropVariants $ do + t <- someElemTypeQ + (as, bs) <- someLenAsBsQ 1 + [e| do + i@(Idx ii :* U) <- arbitrary @(Idxs $as) + df <- arbitrary @(DataFrame $t $(concatDimsQ as bs)) + return $ index i df === df ! ii + |] + ) -prop_IndexCustom1 :: DataFrame Word SSuff -> Bool -prop_IndexCustom1 x = (z ! 1:*3:*2) == x - where - z = ewgen x :: DataFrame Word SFull +prop_IndexFrame2 :: Property +prop_IndexFrame2 = property $ oneof $( aFewPropVariants $ do + t <- someElemTypeQ + (as, bs) <- someLenAsBsQ 2 + [e| do + i@((Idx ii1) :* (Idx ii2) :* U) <- arbitrary @(Idxs $as) + df <- arbitrary @(DataFrame $t $(concatDimsQ as bs)) + return $ index i df === df ! ii1 ! ii2 + |] + ) --- TODO:! --- prop_IndexCustom2 :: DataFrame Double SSuff -> Bool --- prop_IndexCustom2 x = (z ! 2:*2:*1) %* eye == x --- where --- z = ewgen x :: DataFrame Double SFull +prop_IndexOffset :: Property +prop_IndexOffset = property $ oneof $( aFewPropVariants $ do + t <- someElemTypeQ + (as, bs) <- someAsBsNQ + [e| do + i <- arbitrary @(Idxs $as) + df <- arbitrary @(DataFrame $t $(concatDimsQ as bs)) + let asw = listDims (dims @($as)) + bsN = totalDim (dims @($bs)) + isw = listIdxs i + off = fromIntegral . snd + . foldr (\(dn, di) (tn, to) -> (tn*dn, to + tn*di) ) (bsN, 0) + $ zip asw isw + return $ index i df === indexOffset off df + |] + ) -prop_Foldlr :: DataFrame Double SSuff -> Bool -prop_Foldlr x = - abs (ewfoldl (+) 10 z - ewfoldr @_ @SPref (+) 0 z - 10) - <= fromScalar (zmax * 0.0001) - where - z = ewgen x :: DataFrame Double SFull - zmax = ewfoldl @Double @SFull @'[] (max . abs) 0.001 z +prop_Foldlr :: Property +prop_Foldlr = property $ oneof $( aFewPropVariants $ do + t <- someElemTypeFQ + (as, bs) <- someAsBsQ + let asbs = concatDimsQ as bs + [e| do + df <- arbitrary @(DataFrame $t $asbs) + let lim = ewfoldr' @($t) @($asbs) (max.abs) 1 df + * ewfoldr' @($t) @($as) @($bs) (const (+1)) 0 df + l = ewfoldl' (+) 10 df + r = ewfoldr' (+) 0 df + 10 :: DataFrame $t $bs + return $ approxEq lim l r + |] + ) -prop_Ewmap :: DataFrame Double SFull -> Bool -prop_Ewmap x = x * 2 == ewmap @_ @_ @'[Last SFull] (*2) x +prop_LazyFoldlr :: Property +prop_LazyFoldlr = property $ oneof $( aFewPropVariants $ do + t <- someElemTypeQ + (as, bs) <- someAsBsQ + let asbs = concatDimsQ as bs + [e| do + n <- fromIntegral <$> choose (0, maxTotalDim `quot` 2) + df <- arbitrary @(DataFrame $t $asbs) + let S td = ewfoldr' @($t) @($as) @($bs) (const (+1)) 0 df + l = ewfoldl @($t) @($as) @($bs) (limitedList) (const []) df n + l' = ewfoldl' @($t) @($as) @($bs) (flip (:)) [] df + r = ewfoldr @($t) @($as) @($bs) (flip limitedList) (const []) df n + r' = ewfoldr' @($t) @($as) @($bs) (:) [] df -prop_ProdTranspose :: DataFrame Double '[2,6] -> DataFrame Double '[6,7] -> Bool -prop_ProdTranspose x y = transpose (x %* y) == transpose y %* transpose x + return $ + counterexample + ( unlines + [ "Failed with dims:" + , " n: " ++ show n + , " td: " ++ show td + , " @as: " ++ show (typeOf $ minimalDims @($as)) + , " @bs: " ++ show (typeOf $ minimalDims @($bs)) + ] + ) $ conjoin + [ if n >= td + then l === reverse r + else shouldErr (l === l') .&&. shouldErr (r === r') + , take n l === take n l' + , take n r === take n r' + ] + |] + ) where + shouldErr :: Property -> Property + shouldErr p = ioProperty $ checkErr + <$> quickCheckWithResult stdArgs { chatty = False } p + checkErr :: Result -> Property + checkErr Failure { theException = Just _} + = property True + checkErr _ + = counterexample "Should have failed with undefined." False + limitedList :: (Int -> [a]) -> a -> Int -> [a] + limitedList f a k | k <= 0 = undefined + | otherwise = a : f (k - 1) -prop_Eye :: DataFrame Double SFull -> Bool -prop_Eye x = eye %* x == x && x %* eye == x return [] runTests :: Int -> IO Bool diff --git a/test/Numeric/Matrix/BidiagonalTest.hs b/test/Numeric/Matrix/BidiagonalTest.hs new file mode 100644 --- /dev/null +++ b/test/Numeric/Matrix/BidiagonalTest.hs @@ -0,0 +1,144 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE PolyKinds #-} +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeOperators #-} + +module Numeric.Matrix.BidiagonalTest (runTests) where + + +import Control.Monad (join) +import Data.Kind +import Data.Monoid (All (..)) +import Numeric.Arbitraries +import Numeric.DataFrame +import Numeric.Dimensions +import Numeric.Matrix.Bidiagonal +import Test.QuickCheck + +validateBidiagonal :: forall (t :: Type) (n :: Nat) (m :: Nat) + . ( KnownDim n, KnownDim m + , RealFloatExtras t, Show t + , MatrixDeterminant t n + , MatrixDeterminant t m + ) + => Matrix t n m -> BiDiag t n m -> Property +validateBidiagonal a BiDiag {..} = + counterexample + (unlines + [ "failed a =~= u %* b %* transpose v:" + , "a: " ++ show a + , "a': " ++ show a' + , "u: " ++ show u + , "b: " ++ show b + , "v: " ++ show v + ] + ) (approxEq mag a a') + .&&. + counterexample + (unlines + [ "u is not quite orthogonal:" + , "a: " ++ show a + , "u: " ++ show u + , "b: " ++ show b + , "v: " ++ show v + ] + ) (approxEq mag eye (u %* transpose u)) + .&&. + counterexample + (unlines + [ "v is not quite orthogonal:" + , "a: " ++ show a + , "u: " ++ show u + , "b: " ++ show b + , "v: " ++ show v + ] + ) (approxEq mag eye (v %* transpose v)) + .&&. + counterexample + (unlines + [ "b is not upper-bidiagonal" + , "a: " ++ show a + , "u: " ++ show u + , "b: " ++ show b + , "v: " ++ show v + ] + ) (getAll $ iwfoldMap @t @'[n,m] + (\(Idx i :* Idx j :* U) x -> All (j == i || j == i + 1 || x == 0)) + b + ) + .&&. + counterexample + (unlines + [ "Incorrect determinant carried out for U" + , "a: " ++ show a + , "u: " ++ show u + , "bdUDet: " ++ show bdUDet + , "det u: " ++ show (det u) + ] + ) (approxEq mag bdUDet (det u)) + .&&. + counterexample + (unlines + [ "Incorrect determinant carried out for V" + , "a: " ++ show a + , "v: " ++ show v + , "bdVDet: " ++ show bdVDet + , "det v: " ++ show (det v) + ] + ) (approxEq mag bdVDet (det v)) + where + mag = fromIntegral (totalDim $ dims `inSpaceOf` a) ^ (2 :: Int) :: Scalar t + a' = u %* b %* transpose v + b = biDiag dims bdAlpha bdBeta + u = bdU + v = bdV + +prop_bidiagonalSimple :: Property +prop_bidiagonalSimple = once . conjoin $ map prop_bidiagonal $ xs + where + mkM :: Dims ([n,m]) -> [Double] -> DataFrame Double '[XN 1, XN 1] + mkM ds + | Just (XDims ds'@Dims) <- constrainDims ds :: Maybe (Dims '[XN 1, XN 1]) + = XFrame . fromFlatList ds' 0 + mkM _ = error "prop_qrSimple: bad dims" + variants :: Num a => [a] -> [[a]] + variants as = rotateList as ++ rotateList (map negate as) + xs :: [DataFrame Double '[XN 1, XN 1]] + xs = join + [ [mkM $ D2 :* D2 :* U, mkM $ D5 :* D2 :* U, mkM $ D3 :* D7 :* U] + <*> [repeat 0, repeat 1, repeat 2] + , mkM (D2 :* D2 :* U) <$> variants [3,2, 4,1] + , mkM (D3 :* D3 :* U) <$> variants [0,0,1, 3,2,0, 4,1,0] + , mkM (D4 :* D2 :* U) <$> variants [3,2, 4,1, 0,0, 0,2] + , mkM (D2 :* D3 :* U) <$> variants [3,2,0, 4,1,2] + , mkM (D2 :* D2 :* U) <$> variants [2, 0, 0, 0] + , mkM (D2 :* D2 :* U) <$> variants [4, 1, 0, 0] + , mkM (D2 :* D2 :* U) <$> variants [3, 1, -2, 0] + , [mkM $ D2 :* D3 :* U, mkM $ D3 :* D2 :* U] + <*> join + [ variants [2, 0, 0, 0, 0, 0] + , variants [4, 1, 0, 0, 0, 0] + , variants [3, 1, -2, 0, 0, 0] + , variants [3, 0, -2, 9, 0, 0] + ] + ] + +prop_bidiagonal :: DataFrame Double '[XN 1, XN 1] -> Property +prop_bidiagonal (XFrame x) + | n@D :* m@D :* U <- dims `inSpaceOf` x + , D <- minDim n m + = validateBidiagonal x (bidiagonalHouseholder x) +prop_bidiagonal _ = property False + + + + +return [] +runTests :: Int -> IO Bool +runTests n = $forAllProperties + $ quickCheckWithResult stdArgs { maxSuccess = n } diff --git a/test/Numeric/Matrix/LUTest.hs b/test/Numeric/Matrix/LUTest.hs new file mode 100644 --- /dev/null +++ b/test/Numeric/Matrix/LUTest.hs @@ -0,0 +1,104 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE PolyKinds #-} +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeOperators #-} + +module Numeric.Matrix.LUTest (runTests) where + + +import Control.Monad (join) +import Data.Kind +import Data.Monoid (All (..)) +import Numeric.Arbitraries +import Numeric.DataFrame +import Numeric.Dimensions +import Numeric.Matrix.LU +import Test.QuickCheck + + +validateLU :: forall (t :: Type) (n :: Nat) + . ( KnownDim n, RealFloatExtras t, Show t + , MatrixDeterminant t n) + => Matrix t n n -> LU t n -> Property +validateLU x lux@LU {..} = + counterexample + (unlines + [ "failed luPerm m =~= luLower %* luUpper:" + , "x: " ++ show x + , "px: " ++ show px + , "px': " ++ show px' + , "lu: " ++ show lux + ] + ) (px =~= px') + .&&. + counterexample + (unlines + [ "Failed an extra invariant det luPerm == luPermDet:" + , "lu:" ++ show lux + ] + ) (luPermDet =~= det luPerm) + .&&. + counterexample + (unlines + [ "luUpper is not upper-triangular" + , "lu:" ++ show lux + , "x:" ++ show x + ] + ) (getAll $ iwfoldMap @t @'[n,n] + (\(Idx i :* Idx j :* U) a -> All (j >= i || a == 0)) + luUpper + ) + .&&. + counterexample + (unlines + [ "luLower is not unit lower triangular," + ++ " or some of its elements are greater than one" + , "lu:" ++ show lux + , "x:" ++ show x + ] + ) (getAll $ iwfoldMap @t @'[n,n] + (\(Idx i :* Idx j :* U) a -> case compare i j of + GT -> All (abs a <= 1) + EQ -> All (a == 1) + LT -> All (a == 0) + ) luLower + ) + where + px = luPerm %* x + px' = luLower %* luUpper + + +manualMats :: forall t . (PrimBytes t, Num t) => [SomeSquareMatrix AnyMatrix t] +manualMats = join + [ [mkM D2, mkM D4, mkM D9] + <*> [repeat 0, repeat 1, repeat 2] + , mkM D2 <$> variants [3,2, 4,1] + , mkM D3 <$> variants [0,0,1, 3,2,0, 4,1,0] + , mkM D2 <$> variants [2, 0, 0, 0] + , mkM D2 <$> variants [4, 1, 0, 0] + , mkM D2 <$> variants [3, 1, -2, 0] + ] + where + mkM :: Dim (n :: Nat) -> [t] -> SomeSquareMatrix AnyMatrix t + mkM (d@D :: Dim n) = SSM . fromFlatList (d :* d :* U) 0 + mkM _ = error "manualMats: bad dims" + variants :: Num a => [a] -> [[a]] + variants as = rotateList as ++ rotateList (map negate as) + + +prop_luSimple :: Property +prop_luSimple = once . conjoin $ map prop_lu manualMats + +prop_lu :: SomeSquareMatrix AnyMatrix Double -> Property +prop_lu (SSM x) = validateLU x (lu x) + + +return [] +runTests :: Int -> IO Bool +runTests n = $forAllProperties + $ quickCheckWithResult stdArgs { maxSuccess = n } diff --git a/test/Numeric/Matrix/QRTest.hs b/test/Numeric/Matrix/QRTest.hs new file mode 100644 --- /dev/null +++ b/test/Numeric/Matrix/QRTest.hs @@ -0,0 +1,203 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE PolyKinds #-} +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeOperators #-} + +module Numeric.Matrix.QRTest (runTests) where + + +import Control.Monad (join) +import Data.Kind +import Data.Monoid (All (..)) +import Numeric.Arbitraries +import Numeric.DataFrame +import Numeric.Dimensions +import Numeric.Matrix.QR +import Test.QuickCheck + + +validateQR :: forall (t :: Type) (n :: Nat) (m :: Nat) + . ( KnownDim n, KnownDim m + , RealFloatExtras t, Show t + ) + => Matrix t n m -> QR t n m -> Property +validateQR x q@QR {..} = + counterexample + (unlines + [ "failed m =~= qrQ %* qrR:" + , "m: " ++ show x + , "m': " ++ show x' + , "qr:" ++ show q + ] + ) (x =~= x') + .&&. + counterexample + (unlines + [ "qrQ is not quite orthogonal:" + , "qr:" ++ show q + , "m:" ++ show x + ] + ) (eye =~= qrQ %* transpose qrQ) + .&&. + counterexample + (unlines + [ "qrR is not upper-triangular" + , "qr:" ++ show q + , "m:" ++ show x + ] + ) (getAll $ iwfoldMap @t @'[n,m] + (\(Idx i :* Idx j :* U) a -> All (j >= i || a <= tol)) + qrR + ) + where + tol = M_EPS * max 1 (maxElem x) + x' = qrQ %* qrR + +validateLQ :: forall (t :: Type) (n :: Nat) (m :: Nat) + . ( KnownDim n, KnownDim m + , RealFloatExtras t, Show t + ) + => Matrix t n m -> LQ t n m -> Property +validateLQ x q@LQ {..} = + counterexample + (unlines + [ "failed m =~= lqL %* lqQ:" + , "m: " ++ show x + , "m': " ++ show x' + , "lq:" ++ show q + ] + ) (x =~= x') + .&&. + counterexample + (unlines + [ "lqQ is not quite orthogonal:" + , "lq:" ++ show q + , "m:" ++ show x + ] + ) (eye =~= lqQ %* transpose lqQ) + .&&. + counterexample + (unlines + [ "lqL is not lower-triangular" + , "lq:" ++ show q + , "m:" ++ show x + ] + ) (getAll $ iwfoldMap @t @'[n,m] + (\(Idx i :* Idx j :* U) a -> All (j <= i || a <= tol)) + lqL + ) + where + tol = M_EPS * max 1 (maxElem x) + x' = lqL %* lqQ + +manualMats :: [DataFrame Double '[XN 1, XN 1]] +manualMats = join + [ [mkM $ D2 :* D2 :* U, mkM $ D5 :* D2 :* U, mkM $ D3 :* D7 :* U] + <*> [repeat 0, repeat 1, repeat 2] + , mkM (D2 :* D2 :* U) <$> variants [3,2, 4,1] + , mkM (D3 :* D3 :* U) <$> variants [0,0,1, 3,2,0, 4,1,0] + , mkM (D4 :* D2 :* U) <$> variants [3,2, 4,1, 0,0, 0,2] + , mkM (D2 :* D3 :* U) <$> variants [3,2,0, 4,1,2] + , mkM (D2 :* D2 :* U) <$> variants [2, 0, 0, 0] + , mkM (D2 :* D2 :* U) <$> variants [4, 1, 0, 0] + , mkM (D2 :* D2 :* U) <$> variants [3, 1, -2, 0] + , [mkM $ D2 :* D3 :* U, mkM $ D3 :* D2 :* U] + <*> join + [ variants [2, 0, 0, 0, 0, 0] + , variants [4, 1, 0, 0, 0, 0] + , variants [3, 1, -2, 0, 0, 0] + , variants [3, 0, -2, 9, 0, 0] + ] + ] + where + mkM :: Dims ([n,m]) -> [Double] -> DataFrame Double '[XN 1, XN 1] + mkM ds + | Just (XDims ds'@Dims) <- constrainDims ds :: Maybe (Dims '[XN 1, XN 1]) + = XFrame . fromFlatList ds' 0 + mkM _ = error "prop_qrSimple: bad dims" + variants :: Num a => [a] -> [[a]] + variants as = rotateList as ++ rotateList (map negate as) + +prop_lqSimple :: Property +prop_lqSimple = once . conjoin $ map prop_lq manualMats + +prop_lq :: DataFrame Double '[XN 1, XN 1] -> Property +prop_lq (XFrame x) + | n@D :* m@D :* U <- dims `inSpaceOf` x + , D <- minDim n m + = validateLQ x (lq x) +prop_lq _ = property False + +prop_qrSimple :: Property +prop_qrSimple = once . conjoin $ map prop_qr manualMats + +prop_qr :: DataFrame Double '[XN 1, XN 1] -> Property +prop_qr (XFrame x) + | n@D :* m@D :* U <- dims `inSpaceOf` x + , D <- minDim n m + = validateQR x (qr x) +prop_qr _ = property False + + +testQRSolve :: forall t n m + . (MatrixQR t n m, Show t, RealFloatExtras t) + => Matrix t n m -> Vector t n -> Property +testQRSolve a b = + counterexample + (unlines + [ "failed a' %* (a %* xr - b) =~= 0" + , "a: " ++ show a + , "b: " ++ show b + , "xr: " ++ show xr + , "zeroR: " ++ show zeroR + , "zeroL: " ++ show zeroL + ] + ) (nonDeficient ==> approxEq mag2 zeroR 0) + .&&. + counterexample + (unlines + [ "failed (xl %* a' - b) %* a =~= 0" + , "a': " ++ show a' + , "b: " ++ show b + , "xl: " ++ show xl + ] + ) (nonDeficient ==> approxEq mag2 zeroL 0) + .&&. + counterexample + (unlines + [ "failed xr =~= xl" + , "a': " ++ show a' + , "b: " ++ show b + , "xl: " ++ show xl + ] + ) (approxEq mag xr xl) + where + qrr = qrR (qr a) + nonDeficient = + foldl (\r i -> r && abs (qrr ! i ! i) > M_EPS * mag) True + [0 .. min (dimVal' @n) (dimVal' @m) - 1] + nm = totalDim (dims `inSpaceOf` a) + mag = max 1 (maxElem a `max` maxElem b) * fromIntegral nm + mag2 = mag*mag + a' = transpose a + xr = qrSolveR a b + xl = qrSolveL a' b + zeroR = a' %* (a %* xr - b) + zeroL = (xl %* a' - b) %* a + + +prop_qrSolve :: Dim (XN 1) -> Dim (XN 1) -> Property +prop_qrSolve (Dx (D :: Dim n)) (Dx (D :: Dim m)) = property $ do + a <- arbitrary @(DataFrame Double '[n,m]) + b <- arbitrary @(DataFrame Double '[n]) + return $ testQRSolve a b + +return [] +runTests :: Int -> IO Bool +runTests n = $forAllProperties + $ quickCheckWithResult stdArgs { maxSuccess = n } diff --git a/test/Numeric/Matrix/SVDTest.hs b/test/Numeric/Matrix/SVDTest.hs new file mode 100644 --- /dev/null +++ b/test/Numeric/Matrix/SVDTest.hs @@ -0,0 +1,182 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE PolyKinds #-} +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeOperators #-} + +module Numeric.Matrix.SVDTest (runTests) where + + +import Control.Monad +import Data.Kind +import Data.Maybe (isJust) +import Numeric.Arbitraries +import Numeric.DataFrame +import Numeric.Dimensions +import Numeric.Matrix.SVD +import Test.QuickCheck + +-- check invariants of SVD +validateSVD :: forall (t :: Type) (n :: Nat) (m :: Nat) + . ( KnownDim n, KnownDim m, KnownDim (Min n m) + , RealFloatExtras t, Show t + , MatrixDeterminant t n + ) + => t -> Matrix t n m -> SVD t n m -> Property +validateSVD extraTolerance x s@SVD {..} + | nm <- fromIntegral $ dimVal' @n * dimVal' @m = + counterexample + (unlines + [ "failed m =~= u %* asDiag s %* transpose v:" + , "m: " ++ show x + , "m': " ++ show x' + , "svd:" ++ show s + ] + ) (approxEq (S $ nm * nm * extraTolerance) x x') + .&&. + counterexample + (unlines + [ "SVD left-singular basis is not quite orthogonal:" + , "svd:" ++ show s + , "m:" ++ show x + ] + ) (eye =~= svdU %* transpose svdU) + .&&. + counterexample + (unlines + [ "Failed an extra invariant det svdU == 1:" + , "svd:" ++ show s + , "m:" ++ show x + ] + ) (1 =~= det svdU) + .&&. + counterexample + (unlines + [ "SVD right-singular basis is not quite orthogonal:" + , "svd:" ++ show s + , "m:" ++ show x + ] + ) (eye =~= svdV %* transpose svdV) + .&&. + counterexample + (unlines + [ "Bad singular values (should be decreasing, non-negative):" + , "svd:" ++ show s + , "m:" ++ show x + ] + ) + (isJust $ ewfoldr @_ @'[Min n m] @'[] checkIncreasing (Just 0) svdS) + where + checkIncreasing :: Scalar t -> Maybe (Scalar t) -> Maybe (Scalar t) + checkIncreasing _ Nothing = Nothing + checkIncreasing a (Just b) + | a >= b = Just a + | otherwise = Nothing + x' = svdU %* asDiag svdS %* transpose svdV + +prop_svd2x :: Property +prop_svd2x = once . conjoin $ map prop_svd2 $ xs ++ map negate xs + where + mkM :: [Double] -> Mat22d + mkM = fromFlatList dims 0 + xs :: [Mat22d] + xs = + [ 0, 1, 2, eye] + ++ map mkM (rotateList [4,0,0,0]) + ++ map mkM (rotateList [4,3,0,0]) + ++ map mkM (rotateList [4,-3,0,0]) + ++ map mkM (rotateList [3,0,-2,0]) + ++ map mkM (rotateList [3,1,-2,0]) + ++ map mkM (rotateList [1,1,3,3]) + ++ map ((eye + ). mkM) (rotateList [4,0,0,0]) + ++ map ((eye - ). mkM) (rotateList [4,0,0,0]) + +prop_svd2 :: Matrix Double 2 2 -> Property +prop_svd2 m = validateSVD 1 m (svd2 m) + +prop_svd3x :: Property +prop_svd3x = once . conjoin $ map prop_svd3 $ xs ++ map negate xs + where + mkM :: [Double] -> Mat33d + mkM = fromFlatList dims 0 + xs :: [Mat33d] + xs = + [ 0, 1, 2, eye] + ++ map mkM (rotateList [4,0,0,0,0,0,0,0,0]) + ++ map ((eye + ). mkM) (rotateList [4,0,0,0,0,0,0,0,0]) + ++ map ((eye - ). mkM) (rotateList [4,0,0,0,0,0,0,0,0]) + ++ map mkM (rotateList [4,1,0,0,0,0,0,0,0]) + ++ map mkM (rotateList [4,2,1,0,0,0,0,0,0]) + ++ map mkM (rotateList [4,2,1,0,0,7,0,0,0]) + ++ map mkM (rotateList [-5,0,1,0,0,7,1,0,10]) + ++ map mkM (rotateList [0,4.237014829207297,-2.0275007228244943 + ,0,-0.9998659636721602,0 + ,0,0,0.34652350292261475]) + ++ map mkM (rotateList [0,0,15.33245854346668 + ,0,-18.99787842921179,0 + ,0,29.953254288817046,0.78897034887446]) + ++ map mkM (rotateList [-306.5504258807584,0,0 + ,-58.78957738886663,0,95.51923744049921 + ,0,0,-292.92428690917154]) + ++ map mkM (rotateList [0,116.15251855804613,-94.95820664249142 + ,0,-69.551119690005,0 + ,0,0,46.77872074296825]) + + + +prop_svd3 :: Matrix Double 3 3 -> Property +prop_svd3 m = validateSVD extraTolerance m (svd3 m) + where + -- svd3 is fast, but not particularly precise, + -- though, 5-7 digits of precision are still there. + extraTolerance = 10e4 :: Double + + +prop_svdSimple :: Property +prop_svdSimple = once . conjoin $ prop_svd <$> dfs + +dfs :: [DataFrame Double '[XN 1, XN 1]] +dfs = xs + where + mkM :: Dims ([n,m]) -> [Double] -> DataFrame Double '[XN 1, XN 1] + mkM ds + | Just (XDims ds'@Dims) <- constrainDims ds :: Maybe (Dims '[XN 1, XN 1]) + = XFrame . fromFlatList ds' 0 + mkM _ = error "prop_qrSimple: bad dims" + variants :: Num a => [a] -> [[a]] + variants as = rotateList as ++ rotateList (map negate as) + xs :: [DataFrame Double '[XN 1, XN 1]] + xs = join + [ [mkM $ D2 :* D2 :* U, mkM $ D5 :* D2 :* U, mkM $ D3 :* D7 :* U] + <*> [repeat 0, repeat 1, repeat 2] + , mkM (D2 :* D2 :* U) <$> variants [3,2, 4,1] + , mkM (D3 :* D3 :* U) <$> variants [0,0,1, 3,2,0, 4,1,0] + , mkM (D4 :* D2 :* U) <$> variants [3,2, 4,1, 0,0, 0,2] + , mkM (D2 :* D3 :* U) <$> variants [3,2,0, 4,1,2] + , mkM (D2 :* D2 :* U) <$> variants [2, 0, 0, 0] + , mkM (D2 :* D2 :* U) <$> variants [4, 1, 0, 0] + , mkM (D2 :* D2 :* U) <$> variants [3, 1, -2, 0] + , [mkM $ D2 :* D3 :* U, mkM $ D3 :* D2 :* U] + <*> join + [ variants [2, 0, 0, 0, 0, 0] + , variants [4, 1, 0, 0, 0, 0] + , variants [3, 1, -2, 0, 0, 0] + , variants [3, 0, -2, 9, 0, 0] + ] + ] + +prop_svd :: DataFrame Double '[XN 1, XN 1] -> Property +prop_svd (XFrame x) + | n@D :* m@D :* U <- dims `inSpaceOf` x + , D <- minDim n m = validateSVD 1 x (svd x) +prop_svd _ = error "prop_svd: impossible pattern" + + +return [] +runTests :: Int -> IO Bool +runTests n = $forAllProperties + $ quickCheckWithResult stdArgs { maxSuccess = n } diff --git a/test/Numeric/MatrixDoubleTest.hs b/test/Numeric/MatrixDoubleTest.hs --- a/test/Numeric/MatrixDoubleTest.hs +++ b/test/Numeric/MatrixDoubleTest.hs @@ -11,77 +11,48 @@ import Data.Fixed -import Data.Semigroup +import Numeric.Arbitraries import Numeric.DataFrame -import Numeric.DataFrame.Arbitraries import Numeric.Dimensions import Test.QuickCheck type TestElem = Double -type TestDF = DataFrame TestElem -eps :: Scalar TestElem -eps = 0.000000001 - dropW :: (SubSpace t '[3] '[] '[3], SubSpace t '[4] '[] '[4]) => Vector t 4 -> Vector t 3 dropW (Vec4 x y z _) = Vec3 x y z --- | Most of the time, the error is proportional to the maginutude of the biggest element -maxElem :: (SubSpace TestElem ds '[] ds) - => TestDF (ds :: [Nat]) -> Scalar TestElem -maxElem = ewfoldl (\a -> max a . abs) 0 - --- | For operations like @det@, the error is proportional to maximum possible product --- over rows or columns -maxRows :: forall ds - . (SubSpace TestElem '[Head ds] (Tail ds) ds) - => TestDF (ds :: [Nat]) -> Scalar TestElem -maxRows = getProduct . ewfoldMap @_ @'[Head ds] @(Tail ds) @ds (Product . maxElem) - -approxEq :: - forall (ds :: [Nat]) . - ( - Dimensions ds, - Num (TestDF ds), - PrimBytes (TestDF ds), - PrimArray TestElem (TestDF ds) - ) => - TestDF ds -> TestDF ds -> Bool -approxEq a b = maxElem (a - b) <= eps * m - where - m = maxElem a `max` maxElem b -infix 4 `approxEq` - prop_detTranspose :: SomeSquareMatrix AnyMatrix TestElem -> Property -prop_detTranspose (SSM m) - = let a = det m - b = det $ transpose m - n = (case dims `inSpaceOf` m of dn :* _ -> dimVal dn) :: Word - mag = eps * maxRows m * fromIntegral (product [1..n]) - in counterexample - (unlines - [ "failed det/transpose:" - , "m: " ++ show m - , "mT: " ++ show (transpose m) - , show a ++ " /= " ++ show b - ++ " (tolerance: " ++ show mag ++ ")." - ] - ) $ abs (a - b) <= mag +prop_detTranspose (SSM m') = counterexample + (unlines + [ "Failed det/transpose:" + , "m: " ++ show m + , "mT: " ++ show (transpose m) + ] + ) $ approxEq mag a b + where + n2 = totalDim (dims `inSpaceOf` m) + mag = fromIntegral (n2*n2) + mm = recip $ maxElem m' `max` 1 + m = ewmap (*mm) m' + a = det m + b = det $ transpose m prop_inverse :: SomeSquareMatrix NonSingular TestElem -> Property -prop_inverse (SSM m) - = let mi = inverse m - aeq a b = maxElem (b - a) <= eps * maxRows m - in counterexample ("failed inverse:" ++ - show (m, mi, m %* mi, mi %* m)) $ - aeq eye (m %* mi) && aeq eye (mi %* m) - -prop_LU :: SomeSquareMatrix NonSingular TestElem -> Bool -prop_LU (SSM m) - = let f = lu m - aeq a b = maxElem (b - a) <= eps * maxRows m - in aeq (luPerm f %* m) (luLower f %* luUpper f) +prop_inverse (SSM m) = check (m %* mi) .&&. check (mi %* m) + where + n2 = totalDim (dims `inSpaceOf` m) + mag = maxElem m `max` maxElem mi `max` 1 + * fromIntegral (n2*n2) + mi = inverse m + check a = counterexample + ( unlines + [ "Failed inverse:" + , " m: " ++ show m + , " mi: " ++ show mi + , " eye': " ++ show a + ] + ) $ approxEq mag a eye prop_translate3vs4 :: Vector TestElem 4 -> Bool prop_translate3vs4 v = translate4 v == translate3 (dropW v) @@ -95,92 +66,96 @@ prop_rotateX :: Vector TestElem 4 -> Property prop_rotateX v@(Vec4 x y z w) = conjoin [ - v %* rotateX (-2 * pi) `approxEq` v, - v %* rotateX (-1.5 * pi) `approxEq` vec4 x (-z) y w, - v %* rotateX (-pi) `approxEq` vec4 x (-y) (-z) w, - v %* rotateX (-0.5 * pi) `approxEq` vec4 x z (-y) w, - v %* rotateX 0 `approxEq` v, - v %* rotateX (0.5 * pi) `approxEq` vec4 x (-z) y w, - v %* rotateX pi `approxEq` vec4 x (-y) (-z) w, - v %* rotateX (1.5 * pi) `approxEq` vec4 x z (-y) w, - v %* rotateX (2 * pi) `approxEq` v + v %* rotateX (-2 * pi) =~= v, + v %* rotateX (-1.5 * pi) =~= vec4 x (-z) y w, + v %* rotateX (-pi) =~= vec4 x (-y) (-z) w, + v %* rotateX (-0.5 * pi) =~= vec4 x z (-y) w, + v %* rotateX 0 =~= v, + v %* rotateX (0.5 * pi) =~= vec4 x (-z) y w, + v %* rotateX pi =~= vec4 x (-y) (-z) w, + v %* rotateX (1.5 * pi) =~= vec4 x z (-y) w, + v %* rotateX (2 * pi) =~= v ] prop_rotateY :: Vector TestElem 4 -> Property prop_rotateY v@(Vec4 x y z w) = conjoin [ - v %* rotateY (-2 * pi) `approxEq` v, - v %* rotateY (-1.5 * pi) `approxEq` vec4 z y (-x) w, - v %* rotateY (-pi) `approxEq` vec4 (-x) y (-z) w, - v %* rotateY (-0.5 * pi) `approxEq` vec4 (-z) y x w, - v %* rotateY 0 `approxEq` v, - v %* rotateY (0.5 * pi) `approxEq` vec4 z y (-x) w, - v %* rotateY pi `approxEq` vec4 (-x) y (-z) w, - v %* rotateY (1.5 * pi) `approxEq` vec4 (-z) y x w, - v %* rotateY (2 * pi) `approxEq` v + v %* rotateY (-2 * pi) =~= v, + v %* rotateY (-1.5 * pi) =~= vec4 z y (-x) w, + v %* rotateY (-pi) =~= vec4 (-x) y (-z) w, + v %* rotateY (-0.5 * pi) =~= vec4 (-z) y x w, + v %* rotateY 0 =~= v, + v %* rotateY (0.5 * pi) =~= vec4 z y (-x) w, + v %* rotateY pi =~= vec4 (-x) y (-z) w, + v %* rotateY (1.5 * pi) =~= vec4 (-z) y x w, + v %* rotateY (2 * pi) =~= v ] prop_rotateZ :: Vector TestElem 4 -> Property prop_rotateZ v@(Vec4 x y z w) = conjoin [ - v %* rotateZ (-2 * pi) `approxEq` v, - v %* rotateZ (-1.5 * pi) `approxEq` vec4 (-y) x z w, - v %* rotateZ (-pi) `approxEq` vec4 (-x) (-y) z w, - v %* rotateZ (-0.5 * pi) `approxEq` vec4 y (-x) z w, - v %* rotateZ 0 `approxEq` v, - v %* rotateZ (0.5 * pi) `approxEq` vec4 (-y) x z w, - v %* rotateZ pi `approxEq` vec4 (-x) (-y) z w, - v %* rotateZ (1.5 * pi) `approxEq` vec4 y (-x) z w, - v %* rotateZ (2 * pi) `approxEq` v + v %* rotateZ (-2 * pi) =~= v, + v %* rotateZ (-1.5 * pi) =~= vec4 (-y) x z w, + v %* rotateZ (-pi) =~= vec4 (-x) (-y) z w, + v %* rotateZ (-0.5 * pi) =~= vec4 y (-x) z w, + v %* rotateZ 0 =~= v, + v %* rotateZ (0.5 * pi) =~= vec4 (-y) x z w, + v %* rotateZ pi =~= vec4 (-x) (-y) z w, + v %* rotateZ (1.5 * pi) =~= vec4 y (-x) z w, + v %* rotateZ (2 * pi) =~= v ] prop_rotate :: TestElem -> Property prop_rotate a = conjoin [ - rotate (vec3 1 0 0) a `approxEq` rotateX a, - rotate (vec3 0 1 0) a `approxEq` rotateY a, - rotate (vec3 0 0 1) a `approxEq` rotateZ a + rotate (vec3 1 0 0) a =~= rotateX a, + rotate (vec3 0 1 0) a =~= rotateY a, + rotate (vec3 0 0 1) a =~= rotateZ a ] -prop_rotateEuler :: TestElem -> TestElem -> TestElem -> Bool +prop_rotateEuler :: TestElem -> TestElem -> TestElem -> Property prop_rotateEuler pitch yaw roll = - rotateEuler pitch yaw roll `approxEq` rotateZ roll %* rotateY yaw %* rotateX pitch + rotateEuler pitch yaw roll =~= rotateZ roll %* rotateY yaw %* rotateX pitch prop_lookAt :: Vector TestElem 3 -> Vector TestElem 3 -> Vector TestElem 3 -> Property prop_lookAt up cam foc = - (apart cam foc && apart up cam) ==> + (apart cam foc && nonColinear up zb && nonColinear zb xb) ==> conjoin [ - (normalized . fromHom $ toHomPoint foc %* m) `approxEq` vec3 0 0 (-1), - fromHom (toHomPoint cam %* m) `approxEq` 0, - fromHom (toHomVector xb %* m) `approxEq` vec3 1 0 0, - fromHom (toHomVector yb %* m) `approxEq` vec3 0 1 0, - fromHom (toHomVector zb %* m) `approxEq` vec3 0 0 1 + (normalized . fromHom $ toHomPoint foc %* m) `aeq` vec3 0 0 (-1), + fromHom (toHomPoint cam %* m) `aeq` 0, + fromHom (toHomVector xb %* m) `aeq` vec3 1 0 0, + fromHom (toHomVector yb %* m) `aeq` vec3 0 1 0, + fromHom (toHomVector zb %* m) `aeq` vec3 0 0 1 ] where + aeq :: Vector TestElem 3 -> Vector TestElem 3 -> Property + aeq = approxEq 10 apart :: Vector TestElem 3 -> Vector TestElem 3 -> Bool - apart a b = maxElem (a - b) > 0.01 * (maxElem a `max` maxElem b) + apart a b = normL2 (a - b) > max (2 * sqrt M_EPS) (0.1 * (normL2 a + normL2 b)) + nonColinear :: Vector TestElem 3 -> Vector TestElem 3 -> Bool + nonColinear a b = normL2 (cross a b) > 0.1 * abs (dot a b) m = lookAt up cam foc zb = normalized $ cam - foc xb = normalized $ up `cross` zb yb = zb `cross` xb -prop_perspective :: TestElem -> TestElem -> TestElem -> TestElem -> Bool +prop_perspective :: TestElem -> TestElem -> TestElem -> TestElem -> Property prop_perspective a b c d = - and [ - projectTo 0 0 n `approxEq` vec3 0 0 (-1), - projectTo 0 0 f `approxEq` vec3 0 0 1, - projectTo 1 1 n `approxEq` vec3 1 1 (-1), - projectTo 1 (-1) n `approxEq` vec3 1 (-1) (-1), - projectTo (-1) 1 n `approxEq` vec3 (-1) 1 (-1), - projectTo (-1) (-1) n `approxEq` vec3 (-1) (-1) (-1), - projectTo 1 1 f `approxEq` vec3 1 1 1, - projectTo 1 (-1) f `approxEq` vec3 1 (-1) 1, - projectTo (-1) 1 f `approxEq` vec3 (-1) 1 1, - projectTo (-1) (-1) f `approxEq` vec3 (-1) (-1) 1 + conjoin [ + projectTo 0 0 n =~= vec3 0 0 (-1), + projectTo 0 0 f =~= vec3 0 0 1, + projectTo 1 1 n =~= vec3 1 1 (-1), + projectTo 1 (-1) n =~= vec3 1 (-1) (-1), + projectTo (-1) 1 n =~= vec3 (-1) 1 (-1), + projectTo (-1) (-1) n =~= vec3 (-1) (-1) (-1), + projectTo 1 1 f =~= vec3 1 1 1, + projectTo 1 (-1) f =~= vec3 1 (-1) 1, + projectTo (-1) 1 f =~= vec3 (-1) 1 1, + projectTo (-1) (-1) f =~= vec3 (-1) (-1) 1 ] where n = 1.0 + mod' a 9.0 -- Near plane in range [1, 10) - f = n + 1.0 + mod' b 99.0 -- Far plane in range [n + 1, n + 100) + f = 2*n + mod' b 99.0 -- Far plane in range [2*n, 2*n + 100) fovy = (0.1 * pi) + mod' c (0.8 * pi) -- Y-axis field of view in range [0.1*pi, 0.9*pi) aspect = 0.25 + mod' d 4.0 -- Aspect ration in range [1/4, 4/1] hpd = tan (fovy * 0.5) -- height/distance @@ -188,23 +163,23 @@ m = perspective n f fovy aspect projectTo x' y' z = fromHom $ vec4 (x' * wpd * z) (y' * hpd * z) (-z) 1 %* m -prop_orthogonal :: TestElem -> TestElem -> TestElem -> TestElem -> Bool +prop_orthogonal :: TestElem -> TestElem -> TestElem -> TestElem -> Property prop_orthogonal a b c d = - and [ - projectTo 0 0 n `approxEq` vec3 0 0 (-1), - projectTo 0 0 f `approxEq` vec3 0 0 1, - projectTo 1 1 n `approxEq` vec3 1 1 (-1), - projectTo 1 (-1) n `approxEq` vec3 1 (-1) (-1), - projectTo (-1) 1 n `approxEq` vec3 (-1) 1 (-1), - projectTo (-1) (-1) n `approxEq` vec3 (-1) (-1) (-1), - projectTo 1 1 f `approxEq` vec3 1 1 1, - projectTo 1 (-1) f `approxEq` vec3 1 (-1) 1, - projectTo (-1) 1 f `approxEq` vec3 (-1) 1 1, - projectTo (-1) (-1) f `approxEq` vec3 (-1) (-1) 1 + conjoin [ + projectTo 0 0 n =~= vec3 0 0 (-1), + projectTo 0 0 f =~= vec3 0 0 1, + projectTo 1 1 n =~= vec3 1 1 (-1), + projectTo 1 (-1) n =~= vec3 1 (-1) (-1), + projectTo (-1) 1 n =~= vec3 (-1) 1 (-1), + projectTo (-1) (-1) n =~= vec3 (-1) (-1) (-1), + projectTo 1 1 f =~= vec3 1 1 1, + projectTo 1 (-1) f =~= vec3 1 (-1) 1, + projectTo (-1) 1 f =~= vec3 (-1) 1 1, + projectTo (-1) (-1) f =~= vec3 (-1) (-1) 1 ] where n = 1.0 + mod' a 9.0 -- Near plane in range [1, 10) - f = n + 1.0 + mod' b 99.0 -- Far plane in range [n + 1, n + 100) + f = 2*n + mod' b 99.0 -- Far plane in range [2*n, 2*n + 100) w = 1.0 + mod' c 9999.0 -- Width in range [1, 10000) h = 1.0 + mod' d 9999.0 -- Height in range [1, 10000) m = orthogonal n f w h @@ -216,11 +191,11 @@ prop_toHomVector :: Vector TestElem 3 -> Bool prop_toHomVector v@(Vec3 x y z) = toHomVector v == vec4 x y z 0 -prop_fromHom :: Vector TestElem 4 -> Bool +prop_fromHom :: Vector TestElem 4 -> Property prop_fromHom v@(Vec4 x y z w) = case w of - 0 -> fromHom v == vec3 x y z - _ -> fromHom v `approxEq` vec3 (x/w) (y/w) (z/w) + 0 -> fromHom v === vec3 x y z + _ -> fromHom v =~= vec3 (x/w) (y/w) (z/w) return [] runTests :: Int -> IO Bool diff --git a/test/Numeric/MatrixFloatTest.hs b/test/Numeric/MatrixFloatTest.hs --- a/test/Numeric/MatrixFloatTest.hs +++ b/test/Numeric/MatrixFloatTest.hs @@ -11,83 +11,48 @@ import Data.Fixed -import Data.Semigroup +import Numeric.Arbitraries import Numeric.DataFrame -import Numeric.DataFrame.Arbitraries import Numeric.Dimensions import Test.QuickCheck type TestElem = Float -type TestDF = DataFrame TestElem -eps :: Scalar TestElem -eps = 0.00001 - dropW :: (SubSpace t '[3] '[] '[3], SubSpace t '[4] '[] '[4]) => Vector t 4 -> Vector t 3 dropW (Vec4 x y z _) = Vec3 x y z --- | Most of the time, the error is proportional to the maginutude of the biggest element -maxElem :: (SubSpace TestElem ds '[] ds) - => TestDF (ds :: [Nat]) -> Scalar TestElem -maxElem = ewfoldl (\a -> max a . abs) 0 - --- | For operations like @det@, the error is proportional to maximum possible product --- over rows or columns -maxRows :: forall ds - . (SubSpace TestElem '[Head ds] (Tail ds) ds) - => TestDF (ds :: [Nat]) -> Scalar TestElem -maxRows = getProduct . ewfoldMap @_ @'[Head ds] @(Tail ds) @ds (Product . maxElem) - -approxEq :: - forall (ds :: [Nat]) . - ( - Dimensions ds, - Num (TestDF ds), - PrimBytes (TestDF ds), - PrimArray TestElem (TestDF ds) - ) => - TestDF ds -> TestDF ds -> Bool -approxEq a b = maxElem (a - b) <= eps * m - where - m = maxElem a `max` maxElem b -infix 4 `approxEq` - prop_detTranspose :: SomeSquareMatrix AnyMatrix TestElem -> Property -prop_detTranspose (SSM m) - = let a = det m - b = det $ transpose m - n = (case dims `inSpaceOf` m of dn :* _ -> dimVal dn) :: Word - mag = eps * maxRows m * fromIntegral (product [1..n]) - in counterexample - (unlines - [ "failed det/transpose:" - , "m: " ++ show m - , "mT: " ++ show (transpose m) - , show a ++ " /= " ++ show b - ++ " (tolerance: " ++ show mag ++ ")." - ] - ) $ abs (a - b) <= mag +prop_detTranspose (SSM m') = counterexample + (unlines + [ "Failed det/transpose:" + , "m: " ++ show m + , "mT: " ++ show (transpose m) + ] + ) $ approxEq mag a b + where + n2 = totalDim (dims `inSpaceOf` m) + mag = fromIntegral (n2*n2) + mm = recip $ maxElem m' `max` 1 + m = ewmap (*mm) m' + a = det m + b = det $ transpose m prop_inverse :: SomeSquareMatrix NonSingular TestElem -> Property -prop_inverse (SSM m) - = let mi = inverse m - aeq a b = maxElem (b - a) <= eps * maxRows m - in counterexample ("failed inverse:" ++ - show (m, mi, m %* mi, mi %* m)) $ - aeq eye (m %* mi) && aeq eye (mi %* m) - --- TODO: improve lu! --- lu factorization is very unstable in some cases --- --- >>> m = unsafeFromFlatList (Dims @'[4,4]) 0 [1,0,1,-2, 0,-1,1,2, 2,-4,2,-2, 4,-8,-2,2 ] :: Mat44f --- >>> det m == 0 --- >>> det (transpose m) == 60 -prop_LU :: SomeSquareMatrix NonSingular TestElem -> Bool -prop_LU (SSM m) - = let f = lu m - aeq a b = maxElem (b - a) <= eps * maxRows m - in aeq (luPerm f %* m) (luLower f %* luUpper f) +prop_inverse (SSM m) = check (m %* mi) .&&. check (mi %* m) + where + n2 = totalDim (dims `inSpaceOf` m) + mag = maxElem m `max` maxElem mi `max` 1 + * fromIntegral (n2*n2) + mi = inverse m + check a = counterexample + ( unlines + [ "Failed inverse:" + , " m: " ++ show m + , " mi: " ++ show mi + , " eye': " ++ show a + ] + ) $ approxEq mag a eye prop_translate3vs4 :: Vector TestElem 4 -> Bool prop_translate3vs4 v = translate4 v == translate3 (dropW v) @@ -101,92 +66,96 @@ prop_rotateX :: Vector TestElem 4 -> Property prop_rotateX v@(Vec4 x y z w) = conjoin [ - v %* rotateX (-2 * pi) `approxEq` v, - v %* rotateX (-1.5 * pi) `approxEq` vec4 x (-z) y w, - v %* rotateX (-pi) `approxEq` vec4 x (-y) (-z) w, - v %* rotateX (-0.5 * pi) `approxEq` vec4 x z (-y) w, - v %* rotateX 0 `approxEq` v, - v %* rotateX (0.5 * pi) `approxEq` vec4 x (-z) y w, - v %* rotateX pi `approxEq` vec4 x (-y) (-z) w, - v %* rotateX (1.5 * pi) `approxEq` vec4 x z (-y) w, - v %* rotateX (2 * pi) `approxEq` v + v %* rotateX (-2 * pi) =~= v, + v %* rotateX (-1.5 * pi) =~= vec4 x (-z) y w, + v %* rotateX (-pi) =~= vec4 x (-y) (-z) w, + v %* rotateX (-0.5 * pi) =~= vec4 x z (-y) w, + v %* rotateX 0 =~= v, + v %* rotateX (0.5 * pi) =~= vec4 x (-z) y w, + v %* rotateX pi =~= vec4 x (-y) (-z) w, + v %* rotateX (1.5 * pi) =~= vec4 x z (-y) w, + v %* rotateX (2 * pi) =~= v ] prop_rotateY :: Vector TestElem 4 -> Property prop_rotateY v@(Vec4 x y z w) = conjoin [ - v %* rotateY (-2 * pi) `approxEq` v, - v %* rotateY (-1.5 * pi) `approxEq` vec4 z y (-x) w, - v %* rotateY (-pi) `approxEq` vec4 (-x) y (-z) w, - v %* rotateY (-0.5 * pi) `approxEq` vec4 (-z) y x w, - v %* rotateY 0 `approxEq` v, - v %* rotateY (0.5 * pi) `approxEq` vec4 z y (-x) w, - v %* rotateY pi `approxEq` vec4 (-x) y (-z) w, - v %* rotateY (1.5 * pi) `approxEq` vec4 (-z) y x w, - v %* rotateY (2 * pi) `approxEq` v + v %* rotateY (-2 * pi) =~= v, + v %* rotateY (-1.5 * pi) =~= vec4 z y (-x) w, + v %* rotateY (-pi) =~= vec4 (-x) y (-z) w, + v %* rotateY (-0.5 * pi) =~= vec4 (-z) y x w, + v %* rotateY 0 =~= v, + v %* rotateY (0.5 * pi) =~= vec4 z y (-x) w, + v %* rotateY pi =~= vec4 (-x) y (-z) w, + v %* rotateY (1.5 * pi) =~= vec4 (-z) y x w, + v %* rotateY (2 * pi) =~= v ] prop_rotateZ :: Vector TestElem 4 -> Property prop_rotateZ v@(Vec4 x y z w) = conjoin [ - v %* rotateZ (-2 * pi) `approxEq` v, - v %* rotateZ (-1.5 * pi) `approxEq` vec4 (-y) x z w, - v %* rotateZ (-pi) `approxEq` vec4 (-x) (-y) z w, - v %* rotateZ (-0.5 * pi) `approxEq` vec4 y (-x) z w, - v %* rotateZ 0 `approxEq` v, - v %* rotateZ (0.5 * pi) `approxEq` vec4 (-y) x z w, - v %* rotateZ pi `approxEq` vec4 (-x) (-y) z w, - v %* rotateZ (1.5 * pi) `approxEq` vec4 y (-x) z w, - v %* rotateZ (2 * pi) `approxEq` v + v %* rotateZ (-2 * pi) =~= v, + v %* rotateZ (-1.5 * pi) =~= vec4 (-y) x z w, + v %* rotateZ (-pi) =~= vec4 (-x) (-y) z w, + v %* rotateZ (-0.5 * pi) =~= vec4 y (-x) z w, + v %* rotateZ 0 =~= v, + v %* rotateZ (0.5 * pi) =~= vec4 (-y) x z w, + v %* rotateZ pi =~= vec4 (-x) (-y) z w, + v %* rotateZ (1.5 * pi) =~= vec4 y (-x) z w, + v %* rotateZ (2 * pi) =~= v ] prop_rotate :: TestElem -> Property prop_rotate a = conjoin [ - rotate (vec3 1 0 0) a `approxEq` rotateX a, - rotate (vec3 0 1 0) a `approxEq` rotateY a, - rotate (vec3 0 0 1) a `approxEq` rotateZ a + rotate (vec3 1 0 0) a =~= rotateX a, + rotate (vec3 0 1 0) a =~= rotateY a, + rotate (vec3 0 0 1) a =~= rotateZ a ] -prop_rotateEuler :: TestElem -> TestElem -> TestElem -> Bool +prop_rotateEuler :: TestElem -> TestElem -> TestElem -> Property prop_rotateEuler pitch yaw roll = - rotateEuler pitch yaw roll `approxEq` rotateZ roll %* rotateY yaw %* rotateX pitch + rotateEuler pitch yaw roll =~= rotateZ roll %* rotateY yaw %* rotateX pitch prop_lookAt :: Vector TestElem 3 -> Vector TestElem 3 -> Vector TestElem 3 -> Property prop_lookAt up cam foc = - (apart cam foc && apart up cam) ==> + (apart cam foc && nonColinear up zb && nonColinear zb xb) ==> conjoin [ - (normalized . fromHom $ toHomPoint foc %* m) `approxEq` vec3 0 0 (-1), - fromHom (toHomPoint cam %* m) `approxEq` 0, - fromHom (toHomVector xb %* m) `approxEq` vec3 1 0 0, - fromHom (toHomVector yb %* m) `approxEq` vec3 0 1 0, - fromHom (toHomVector zb %* m) `approxEq` vec3 0 0 1 + (normalized . fromHom $ toHomPoint foc %* m) `aeq` vec3 0 0 (-1), + fromHom (toHomPoint cam %* m) `aeq` 0, + fromHom (toHomVector xb %* m) `aeq` vec3 1 0 0, + fromHom (toHomVector yb %* m) `aeq` vec3 0 1 0, + fromHom (toHomVector zb %* m) `aeq` vec3 0 0 1 ] where + aeq :: Vector TestElem 3 -> Vector TestElem 3 -> Property + aeq = approxEq 10 apart :: Vector TestElem 3 -> Vector TestElem 3 -> Bool - apart a b = maxElem (a - b) > 0.01 * (maxElem a `max` maxElem b) + apart a b = normL2 (a - b) > max (2 * sqrt M_EPS) (0.1 * (normL2 a + normL2 b)) + nonColinear :: Vector TestElem 3 -> Vector TestElem 3 -> Bool + nonColinear a b = normL2 (cross a b) > 0.1 * abs (dot a b) m = lookAt up cam foc zb = normalized $ cam - foc xb = normalized $ up `cross` zb yb = zb `cross` xb -prop_perspective :: TestElem -> TestElem -> TestElem -> TestElem -> Bool +prop_perspective :: TestElem -> TestElem -> TestElem -> TestElem -> Property prop_perspective a b c d = - and [ - projectTo 0 0 n `approxEq` vec3 0 0 (-1), - projectTo 0 0 f `approxEq` vec3 0 0 1, - projectTo 1 1 n `approxEq` vec3 1 1 (-1), - projectTo 1 (-1) n `approxEq` vec3 1 (-1) (-1), - projectTo (-1) 1 n `approxEq` vec3 (-1) 1 (-1), - projectTo (-1) (-1) n `approxEq` vec3 (-1) (-1) (-1), - projectTo 1 1 f `approxEq` vec3 1 1 1, - projectTo 1 (-1) f `approxEq` vec3 1 (-1) 1, - projectTo (-1) 1 f `approxEq` vec3 (-1) 1 1, - projectTo (-1) (-1) f `approxEq` vec3 (-1) (-1) 1 + conjoin [ + projectTo 0 0 n =~= vec3 0 0 (-1), + projectTo 0 0 f =~= vec3 0 0 1, + projectTo 1 1 n =~= vec3 1 1 (-1), + projectTo 1 (-1) n =~= vec3 1 (-1) (-1), + projectTo (-1) 1 n =~= vec3 (-1) 1 (-1), + projectTo (-1) (-1) n =~= vec3 (-1) (-1) (-1), + projectTo 1 1 f =~= vec3 1 1 1, + projectTo 1 (-1) f =~= vec3 1 (-1) 1, + projectTo (-1) 1 f =~= vec3 (-1) 1 1, + projectTo (-1) (-1) f =~= vec3 (-1) (-1) 1 ] where n = 1.0 + mod' a 9.0 -- Near plane in range [1, 10) - f = n + 1.0 + mod' b 99.0 -- Far plane in range [n + 1, n + 100) + f = 2*n + mod' b 99.0 -- Far plane in range [2*n, 2*n + 100) fovy = (0.1 * pi) + mod' c (0.8 * pi) -- Y-axis field of view in range [0.1*pi, 0.9*pi) aspect = 0.25 + mod' d 4.0 -- Aspect ration in range [1/4, 4/1] hpd = tan (fovy * 0.5) -- height/distance @@ -194,23 +163,23 @@ m = perspective n f fovy aspect projectTo x' y' z = fromHom $ vec4 (x' * wpd * z) (y' * hpd * z) (-z) 1 %* m -prop_orthogonal :: TestElem -> TestElem -> TestElem -> TestElem -> Bool +prop_orthogonal :: TestElem -> TestElem -> TestElem -> TestElem -> Property prop_orthogonal a b c d = - and [ - projectTo 0 0 n `approxEq` vec3 0 0 (-1), - projectTo 0 0 f `approxEq` vec3 0 0 1, - projectTo 1 1 n `approxEq` vec3 1 1 (-1), - projectTo 1 (-1) n `approxEq` vec3 1 (-1) (-1), - projectTo (-1) 1 n `approxEq` vec3 (-1) 1 (-1), - projectTo (-1) (-1) n `approxEq` vec3 (-1) (-1) (-1), - projectTo 1 1 f `approxEq` vec3 1 1 1, - projectTo 1 (-1) f `approxEq` vec3 1 (-1) 1, - projectTo (-1) 1 f `approxEq` vec3 (-1) 1 1, - projectTo (-1) (-1) f `approxEq` vec3 (-1) (-1) 1 + conjoin [ + projectTo 0 0 n =~= vec3 0 0 (-1), + projectTo 0 0 f =~= vec3 0 0 1, + projectTo 1 1 n =~= vec3 1 1 (-1), + projectTo 1 (-1) n =~= vec3 1 (-1) (-1), + projectTo (-1) 1 n =~= vec3 (-1) 1 (-1), + projectTo (-1) (-1) n =~= vec3 (-1) (-1) (-1), + projectTo 1 1 f =~= vec3 1 1 1, + projectTo 1 (-1) f =~= vec3 1 (-1) 1, + projectTo (-1) 1 f =~= vec3 (-1) 1 1, + projectTo (-1) (-1) f =~= vec3 (-1) (-1) 1 ] where n = 1.0 + mod' a 9.0 -- Near plane in range [1, 10) - f = n + 1.0 + mod' b 99.0 -- Far plane in range [n + 1, n + 100) + f = 2*n + mod' b 99.0 -- Far plane in range [2*n, 2*n + 100) w = 1.0 + mod' c 9999.0 -- Width in range [1, 10000) h = 1.0 + mod' d 9999.0 -- Height in range [1, 10000) m = orthogonal n f w h @@ -222,11 +191,11 @@ prop_toHomVector :: Vector TestElem 3 -> Bool prop_toHomVector v@(Vec3 x y z) = toHomVector v == vec4 x y z 0 -prop_fromHom :: Vector TestElem 4 -> Bool +prop_fromHom :: Vector TestElem 4 -> Property prop_fromHom v@(Vec4 x y z w) = case w of - 0 -> fromHom v == vec3 x y z - _ -> fromHom v `approxEq` vec3 (x/w) (y/w) (z/w) + 0 -> fromHom v === vec3 x y z + _ -> fromHom v =~= vec3 (x/w) (y/w) (z/w) return [] runTests :: Int -> IO Bool diff --git a/test/Numeric/PrimBytesTest.hs b/test/Numeric/PrimBytesTest.hs --- a/test/Numeric/PrimBytesTest.hs +++ b/test/Numeric/PrimBytesTest.hs @@ -28,11 +28,11 @@ import GHC.Exts import GHC.Generics import GHC.Word -import Numeric.DataFrame.Arbitraries () +import Numeric.Arbitraries () import Numeric.Dimensions import Numeric.PrimBytes -import qualified Numeric.Tuple.Lazy as LT -import qualified Numeric.Tuple.Strict as ST +import qualified Numeric.Tuple.Lazy as LT +import qualified Numeric.Tuple.Strict as ST import Spec.Util import Test.QuickCheck diff --git a/test/Numeric/QuaterDoubleTest.hs b/test/Numeric/QuaterDoubleTest.hs --- a/test/Numeric/QuaterDoubleTest.hs +++ b/test/Numeric/QuaterDoubleTest.hs @@ -2,42 +2,30 @@ {-# LANGUAGE TemplateHaskell #-} module Numeric.QuaterDoubleTest (runTests) where -import Numeric.DataFrame.Arbitraries () +import Numeric.Arbitraries +import Numeric.Basics import Numeric.Quaternion -import Numeric.Scalar import Numeric.Vector import Test.QuickCheck type T = Double -eps :: T -eps = 1e-7 - -- | Some non-linear function are very unstable; -- it would be a downting task to determine the uncertainty precisely. --- Instead, I just make sure the numbers are not so big to find at least +-- Instead, I just make sure the tolerance is small enough to find at least -- the most obvious bugs. -smaller :: Quater T -> Quater T -smaller q@(Quater x y z w) - = if square q >= 1 then Quater (f x) (f y) (f z) (f w) else q +-- This function increases the tolerance by the span of magnitudes in q. +qSpan :: Quater T -> T +qSpan (Quater a b c d) = asSpan . foldl mm (1,1) $ map (\x -> x*x) [a, b, c, d] where - f :: T -> T - f a | abs a > 2 = signum a * log (abs a) - | otherwise = a + mm :: (T,T) -> T -> (T,T) + mm (mi, ma) x + | x > M_EPS = (min mi x, max ma x) + | otherwise = (mi, ma) + asSpan :: (T,T) -> T + asSpan (mi, ma) = ma / mi -(=~=) :: Quater T -> Quater T -> Property -(=~=) = approxEq eps -infix 4 =~= --- Even lower precision - for naturally fragile operations -approxEq :: T -> Quater T -> Quater T -> Property -approxEq e a b = property $ err <= c * scalar e - where - v1 = toVec4 a - v2 = toVec4 b - c = 1 `max` normL1 v1 `max` normL1 v2 - err = normL1 $ v1 - v2 - prop_Eq :: Quater T -> Bool prop_Eq q = and [ q == q @@ -58,19 +46,20 @@ prop_GetRotScale :: Vector T 3 -> Vector T 3 -> Property prop_GetRotScale a b - = a /= 0 ==> approxEq eps' (fromVecNum (rotScale q a) 0) (fromVecNum b 0) + = normL2 a * ab > M_EPS * normL2 b + ==> approxEq (recip ab) b (rotScale q a) where q = getRotScale a b -- when a and b are almost opposite, precision of getRotScale suffers a lot - -- compensate it by increasing eps: - eps' = if cross a b == 0 - then eps - else eps / (max eps . min 1 . unScalar $ - 1 + normalized a `dot` normalized b) + -- compensate it by increasing tolerance: + ab = min 1 $ 1 + normalized a `dot` normalized b + prop_InverseRotScale :: Quater T -> Vector T 3 -> Property prop_InverseRotScale q v - = square q > eps ==> fromVecNum (rotScale (1/q) (rotScale q v)) 0 =~= fromVecNum v 0 + = min (recip s) s > M_EPS ==> v =~= rotScale (1/q) (rotScale q v) + where + s = square q prop_NegateToMatrix33 :: Quater T -> Bool prop_NegateToMatrix33 q = toMatrix33 q == toMatrix33 (negate q) @@ -95,53 +84,59 @@ prop_UnitQ :: Quater T -> Property prop_UnitQ q - = square q > eps ==> realToFrac (square (q / q)) =~= 1 + = square q > M_EPS ==> square (q / q) =~= 1 prop_ExpLog :: Quater T -> Property -prop_ExpLog q | square q < eps = log (exp q) =~= q - | otherwise = exp (log q) =~= q +prop_ExpLog q | square q < M_EPS = approxEq (qSpan q) q $ log (exp q) + | otherwise = approxEq (qSpan q) q $ exp (log q) prop_SinAsin :: Quater T -> Property -prop_SinAsin q = sin (asin q') =~= q' - where - q' = signum q +prop_SinAsin q = approxEq (qSpan q `max` qSpan q') q $ sin q' + where + q' = asin q prop_CosAcos :: Quater T -> Property -prop_CosAcos q = cos (acos q') =~= q' - where - q' = signum q +prop_CosAcos q = approxEq (qSpan q `max` qSpan q') q $ cos q' + where + q' = acos q prop_TanAtan :: Quater T -> Property -prop_TanAtan q = tan (atan q') =~= q' - where - q' = sqrt (smaller q) -- because this inverse fun is fragile +prop_TanAtan q = approxEq (qSpan q `max` qSpan q') q $ tan q' + where + q' = atan q prop_SinhAsinh :: Quater T -> Property -prop_SinhAsinh q = sinh (asinh q') =~= q' - where - q' = sqrt (smaller q) -- because this inverse fun is fragile +prop_SinhAsinh q = approxEq (qSpan q `max` qSpan q') q $ sinh q' + where + q' = asinh q prop_CoshAcosh :: Quater T -> Property -prop_CoshAcosh q = cosh (acosh q) =~= q +prop_CoshAcosh q = approxEq (qSpan q `max` qSpan q') q $ cosh q' + where + q' = acosh q prop_TanhAtanh :: Quater T -> Property -prop_TanhAtanh q = tanh (atanh q') =~= q' - where - q' = sqrt (smaller q) -- because this inverse fun is fragile +prop_TanhAtanh q = approxEq (qSpan q `max` qSpan q') q $ tanh q' + where + q' = atanh q prop_SqrtSqr :: Quater T -> Property -prop_SqrtSqr q = sqrt q * sqrt q =~= q +prop_SqrtSqr q = approxEq (qSpan q) q $ sqrt q * sqrt q prop_SinCos :: Quater T -> Property -prop_SinCos q = sin q' * sin q' + cos q' * cos q' =~= 1 - where - q' = sqrt (smaller q) -- because it involves exponents +prop_SinCos q' = approxEq (qSpan s `max` qSpan c) 1 $ c * c + s * s + where + q = signum q' -- avoid exploding exponents + s = sin q + c = cos q prop_SinhCosh :: Quater T -> Property -prop_SinhCosh q = cosh q' * cosh q' - sinh q' * sinh q' =~= 1 - where - q' = sqrt (smaller q) -- because it involves exponents +prop_SinhCosh q' = approxEq (qSpan s `max` qSpan c) 1 $ c * c - s * s + where + q = signum q' -- avoid exploding exponents + s = sinh q + c = cosh q prop_ReadShow :: Quater T -> Bool prop_ReadShow q = q == read (show q) diff --git a/test/Numeric/QuaterFloatTest.hs b/test/Numeric/QuaterFloatTest.hs --- a/test/Numeric/QuaterFloatTest.hs +++ b/test/Numeric/QuaterFloatTest.hs @@ -2,42 +2,30 @@ {-# LANGUAGE TemplateHaskell #-} module Numeric.QuaterFloatTest (runTests) where -import Numeric.DataFrame.Arbitraries () +import Numeric.Arbitraries +import Numeric.Basics import Numeric.Quaternion -import Numeric.Scalar import Numeric.Vector import Test.QuickCheck type T = Float -eps :: T -eps = 1e-3 - -- | Some non-linear function are very unstable; -- it would be a downting task to determine the uncertainty precisely. --- Instead, I just make sure the numbers are not so big to find at least +-- Instead, I just make sure the tolerance is small enough to find at least -- the most obvious bugs. -smaller :: Quater T -> Quater T -smaller q@(Quater x y z w) - = if square q >= 1 then Quater (f x) (f y) (f z) (f w) else q +-- This function increases the tolerance by the span of magnitudes in q. +qSpan :: Quater T -> T +qSpan (Quater a b c d) = asSpan . foldl mm (1,1) $ map (\x -> x*x) [a, b, c, d] where - f :: T -> T - f a | abs a > 2 = signum a * log (abs a) - | otherwise = a + mm :: (T,T) -> T -> (T,T) + mm (mi, ma) x + | x > M_EPS = (min mi x, max ma x) + | otherwise = (mi, ma) + asSpan :: (T,T) -> T + asSpan (mi, ma) = ma / mi -(=~=) :: Quater T -> Quater T -> Property -(=~=) = approxEq eps -infix 4 =~= --- Even lower precision - for naturally fragile operations -approxEq :: T -> Quater T -> Quater T -> Property -approxEq e a b = property $ err <= c * scalar e - where - v1 = toVec4 a - v2 = toVec4 b - c = 1 `max` normL1 v1 `max` normL1 v2 - err = normL1 $ v1 - v2 - prop_Eq :: Quater T -> Bool prop_Eq q = and [ q == q @@ -58,19 +46,20 @@ prop_GetRotScale :: Vector T 3 -> Vector T 3 -> Property prop_GetRotScale a b - = a /= 0 ==> approxEq eps' (fromVecNum (rotScale q a) 0) (fromVecNum b 0) + = normL2 a * ab > M_EPS * normL2 b + ==> approxEq (recip ab) b (rotScale q a) where q = getRotScale a b -- when a and b are almost opposite, precision of getRotScale suffers a lot - -- compensate it by increasing eps: - eps' = if cross a b == 0 - then eps - else eps / (max eps . min 1 . unScalar $ - 1 + normalized a `dot` normalized b) + -- compensate it by increasing tolerance: + ab = min 1 $ 1 + normalized a `dot` normalized b + prop_InverseRotScale :: Quater T -> Vector T 3 -> Property prop_InverseRotScale q v - = square q > eps ==> fromVecNum (rotScale (1/q) (rotScale q v)) 0 =~= fromVecNum v 0 + = min (recip s) s > M_EPS ==> v =~= rotScale (1/q) (rotScale q v) + where + s = square q prop_NegateToMatrix33 :: Quater T -> Bool prop_NegateToMatrix33 q = toMatrix33 q == toMatrix33 (negate q) @@ -95,53 +84,59 @@ prop_UnitQ :: Quater T -> Property prop_UnitQ q - = square q > eps ==> realToFrac (square (q / q)) =~= 1 + = square q > M_EPS ==> square (q / q) =~= 1 prop_ExpLog :: Quater T -> Property -prop_ExpLog q | square q < eps = log (exp q) =~= q - | otherwise = exp (log q) =~= q +prop_ExpLog q | square q < M_EPS = approxEq (qSpan q) q $ log (exp q) + | otherwise = approxEq (qSpan q) q $ exp (log q) prop_SinAsin :: Quater T -> Property -prop_SinAsin q = sin (asin q') =~= q' - where - q' = signum q +prop_SinAsin q = approxEq (qSpan q `max` qSpan q') q $ sin q' + where + q' = asin q prop_CosAcos :: Quater T -> Property -prop_CosAcos q = cos (acos q') =~= q' - where - q' = signum q +prop_CosAcos q = approxEq (qSpan q `max` qSpan q') q $ cos q' + where + q' = acos q prop_TanAtan :: Quater T -> Property -prop_TanAtan q = tan (atan q') =~= q' - where - q' = sqrt (smaller q) -- because this inverse fun is fragile +prop_TanAtan q = approxEq (qSpan q `max` qSpan q') q $ tan q' + where + q' = atan q prop_SinhAsinh :: Quater T -> Property -prop_SinhAsinh q = sinh (asinh q') =~= q' - where - q' = sqrt (smaller q) -- because this inverse fun is fragile +prop_SinhAsinh q = approxEq (qSpan q `max` qSpan q') q $ sinh q' + where + q' = asinh q prop_CoshAcosh :: Quater T -> Property -prop_CoshAcosh q = cosh (acosh q) =~= q +prop_CoshAcosh q = approxEq (qSpan q `max` qSpan q') q $ cosh q' + where + q' = acosh q prop_TanhAtanh :: Quater T -> Property -prop_TanhAtanh q = tanh (atanh q') =~= q' - where - q' = sqrt (smaller q) -- because this inverse fun is fragile +prop_TanhAtanh q = approxEq (qSpan q `max` qSpan q') q $ tanh q' + where + q' = atanh q prop_SqrtSqr :: Quater T -> Property -prop_SqrtSqr q = sqrt q * sqrt q =~= q +prop_SqrtSqr q = approxEq (qSpan q) q $ sqrt q * sqrt q prop_SinCos :: Quater T -> Property -prop_SinCos q = sin q' * sin q' + cos q' * cos q' =~= 1 - where - q' = sqrt (smaller q) -- because it involves exponents +prop_SinCos q' = approxEq (qSpan s `max` qSpan c) 1 $ c * c + s * s + where + q = signum q' -- avoid exploding exponents + s = sin q + c = cos q prop_SinhCosh :: Quater T -> Property -prop_SinhCosh q = cosh q' * cosh q' - sinh q' * sinh q' =~= 1 - where - q' = sqrt (smaller q) -- because it involves exponents +prop_SinhCosh q' = approxEq (qSpan s `max` qSpan c) 1 $ c * c - s * s + where + q = signum q' -- avoid exploding exponents + s = sinh q + c = cosh q prop_ReadShow :: Quater T -> Bool prop_ReadShow q = q == read (show q) diff --git a/test/Numeric/Subroutine/SolveTriangularTest.hs b/test/Numeric/Subroutine/SolveTriangularTest.hs new file mode 100644 --- /dev/null +++ b/test/Numeric/Subroutine/SolveTriangularTest.hs @@ -0,0 +1,249 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE PolyKinds #-} +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeOperators #-} + +module Numeric.Subroutine.SolveTriangularTest (runTests) where + + +import Control.Monad.ST +import Numeric.Arbitraries +import Numeric.DataFrame +import Numeric.DataFrame.ST +import Numeric.Dimensions +import Numeric.Subroutine.SolveTriangular +import Test.QuickCheck + + +arbitraryTriangular :: + forall (n :: Nat) (m :: Nat) + . (KnownDim n, KnownDim m, KnownBackend Double '[m]) + => Bool -> Gen (DataFrame Double '[n,m]) +arbitraryTriangular upper = iwmap f <$> arbitrary + where + f :: Idxs '[n] -> DataFrame Double '[m] -> DataFrame Double '[m] + f (Idx i :* U) x = iwmap g x + where + a = maxElem x + g :: Idxs '[m] -> Scalar Double -> Scalar Double + g (Idx j :* U) c + | if upper then j < i else j > i = 0 + | j == i = if c >= 0 then c + 1 + a else c - 1 - a + | otherwise = c + +-- | Test Rx = b +testSolveUpperTriangularR :: + forall (n :: Nat) (m :: Nat) (ds :: [Nat]) + . (KnownDim n, KnownDim m, m <= n, Dimensions ds) + => DataFrame Double '[n,m] -- ^ R + -> DataFrame Double (n :+ ds) -- ^ b + -> Property +testSolveUpperTriangularR r b + | dn <- dim @n + , dm <- dim @m + , di@D <- dn `minusDim` dm `plusDim` D1 + , i0 <- (Idx 0 :* U) `inSpaceOf` (di :* U) + , Just Dict <- sameDim (plusDim dn D1) (plusDim di dm) + = let bm :: DataFrame Double (m :+ ds) + bm = slice i0 b + x :: DataFrame Double (m :+ ds) + x = runST $ do + xPtr <- thawDataFrame bm + solveUpperTriangularR r xPtr + unsafeFreezeDataFrame xPtr + in counterexample + (unlines + [ "failed Rx = b:" + , "R: " ++ show r + , "x: " ++ show x + , "b: " ++ show b + , "Rx:" ++ show (r %* x) + ] + ) (approxEq (maxElem r) b (r %* x)) +testSolveUpperTriangularR _ _ = error "impossible pattern" + +prop_SolveUpperTriangularR :: Property +prop_SolveUpperTriangularR = property run + where + run :: Gen Property + run = do + Dx (m@D :: Dim m) <- arbitrary :: Gen (Dim (XN 1)) + Dx (D :: Dim n) <- arbitrary :: Gen (Dim (XN m)) + justVec <- arbitrary :: Gen Bool + SomeDims (Dims :: Dims ds) + <- if justVec + then pure (SomeDims U) + else removeDimsAbove 100 <$> arbitrary + r <- arbitraryTriangular @n @m True + b' <- arbitrary @(DataFrame Double (n :+ ds)) + let b = iwmap @_ @(n ': ds) @'[] + (\(Idx i :* _) x -> if i >= dimVal m then 0 else x) b' + return $ testSolveUpperTriangularR r b + +-- | Test xR = b +testSolveUpperTriangularL :: + forall (n :: Nat) (m :: Nat) (ds :: [Nat]) + . (KnownDim n, KnownDim m, m <= n, Dimensions ds) + => DataFrame Double (ds +: m) -- ^ b + -> DataFrame Double '[n,m] -- ^ r + -> Property +testSolveUpperTriangularL b r + | dn <- dim @n + , dm <- dim @m + , dnm@D <- minusDim dn dm + , Dims <- Snoc (dims @ds) dn + , Dims <- Snoc (dims @ds) dm + , Dict <- Dict @(SnocList ds n _) + , Dict <- Dict @(SnocList ds m _) + , Just Dict <- sameDim dn (plusDim dm dnm) + = let padZeroes :: DataFrame Double (ds +: m) -> DataFrame Double (ds +: n) + padZeroes = ewmap @_ @ds @'[n] $ \a -> appendDF a (0 `inSpaceOf` (dnm :* U)) + x :: DataFrame Double (ds +: n) + x = runST $ do + xPtr <- thawDataFrame b + solveUpperTriangularL xPtr r + padZeroes <$> unsafeFreezeDataFrame xPtr + in counterexample + (unlines + [ "failed xR = b:" + , "R: " ++ show r + , "x: " ++ show x + , "b: " ++ show b + , "xR:" ++ show (x %* r) + ] + ) (approxEq (maxElem r) b (x %* r)) +testSolveUpperTriangularL _ _ = error "impossible pattern" + +prop_SolveUpperTriangularL :: Property +prop_SolveUpperTriangularL = property run + where + run :: Gen Property + run = do + Dx (m@D :: Dim m) <- arbitrary :: Gen (Dim (XN 1)) + Dx (D :: Dim n) <- arbitrary :: Gen (Dim (XN m)) + justVec <- arbitrary :: Gen Bool + SomeDims (ds@Dims :: Dims ds) + <- if justVec + then pure (SomeDims U) + else removeDimsAbove 100 <$> arbitrary + Dims <- pure $ Snoc ds m + r <- arbitraryTriangular @n @m True + b <- arbitrary @(DataFrame Double (ds +: m)) + return $ testSolveUpperTriangularL b r + + +-- | Test Lx = b +testSolveLowerTriangularR :: + forall (n :: Nat) (m :: Nat) (ds :: [Nat]) + . (KnownDim n, KnownDim m, n <= m, Dimensions ds) + => DataFrame Double '[n,m] -- ^ L + -> DataFrame Double (n :+ ds) -- ^ b + -> Property +testSolveLowerTriangularR l b + | dn <- dim @n + , dm <- dim @m + , dmn@D <- dm `minusDim` dn + , Just Dict <- sameDim dm (dn `plusDim` dmn) + = let padZeroes :: DataFrame Double (n :+ ds) -> DataFrame Double (m :+ ds) + padZeroes z = appendDF z (0 :: DataFrame Double ((m-n) :+ ds)) + x :: DataFrame Double (m :+ ds) + x = runST $ do + xPtr <- thawDataFrame b + solveLowerTriangularR l xPtr + padZeroes <$> unsafeFreezeDataFrame xPtr + in counterexample + (unlines + [ "failed Lx = b:" + , "L: " ++ show l + , "x: " ++ show x + , "b: " ++ show b + , "Lx:" ++ show (l %* x) + ] + ) (approxEq (maxElem l) b (l %* x)) +testSolveLowerTriangularR _ _ = error "impossible pattern" + +prop_SolveLowerTriangularR :: Property +prop_SolveLowerTriangularR = property run + where + run :: Gen Property + run = do + Dx (D :: Dim n) <- arbitrary :: Gen (Dim (XN 1)) + Dx (D :: Dim m) <- arbitrary :: Gen (Dim (XN n)) + justVec <- arbitrary :: Gen Bool + SomeDims (Dims :: Dims ds) + <- if justVec + then pure (SomeDims U) + else removeDimsAbove 100 <$> arbitrary + l <- arbitraryTriangular @n @m False + b <- arbitrary @(DataFrame Double (n :+ ds)) + return $ testSolveLowerTriangularR l b + + +-- | Test xL = b +testSolveLowerTriangularL :: + forall (n :: Nat) (m :: Nat) (ds :: [Nat]) + . (KnownDim n, KnownDim m, n <= m, Dimensions ds) + => DataFrame Double (ds +: m) -- ^ b + -> DataFrame Double '[n,m] -- ^ L + -> Property +testSolveLowerTriangularL b l + | dn <- dim @n + , dm <- dim @m + , dmn@D <- minusDim dm dn + , Dims <- Snoc (dims @ds) dn + , Dims <- Snoc (dims @ds) dm + , Dict <- Dict @(SnocList ds n _) + , Dict <- Dict @(SnocList ds m _) + , Just Dict <- sameDim dm (plusDim dn dmn) + , di@D <- dm `minusDim` dn `plusDim` D1 + , i0 <- (Idx 0 :* U) `inSpaceOf` (di :* U) + , Just Dict <- sameDim (plusDim dm D1) (plusDim di dn) + = let dropLast :: DataFrame Double (ds +: m) -> DataFrame Double (ds +: n) + dropLast = ewmap @_ @ds @'[n] $ slice i0 + x :: DataFrame Double (ds +: n) + x = runST $ do + xPtr <- thawDataFrame b + solveLowerTriangularL xPtr l + dropLast <$> unsafeFreezeDataFrame xPtr + in counterexample + (unlines + [ "failed xL = b:" + , "R: " ++ show l + , "x: " ++ show x + , "b: " ++ show b + , "xL:" ++ show (x %* l) + ] + ) (approxEq (maxElem l) b (x %* l)) +testSolveLowerTriangularL _ _ = error "impossible pattern" + +prop_SolveLowerTriangularL :: Property +prop_SolveLowerTriangularL = property run + where + run :: Gen Property + run = do + Dx (n@D :: Dim n) <- arbitrary :: Gen (Dim (XN 1)) + Dx (m@D :: Dim m) <- arbitrary :: Gen (Dim (XN n)) + justVec <- arbitrary :: Gen Bool + SomeDims (ds@Dims :: Dims ds) + <- if justVec + then pure (SomeDims U) + else removeDimsAbove 100 <$> arbitrary + Dims <- pure $ Snoc ds m + Dict <- pure $ Dict @(SnocList ds m _) + l <- arbitraryTriangular @n @m False + b' <- arbitrary @(DataFrame Double (ds +: m)) + let b = ewmap @_ @ds @'[m] + (iwmap @_ @'[m] @'[] + (\(Idx i :* _) x -> if i >= dimVal n then 0 else x) + ) b' + return $ testSolveLowerTriangularL b l + +return [] +runTests :: Int -> IO Bool +runTests n = $forAllProperties + $ quickCheckWithResult stdArgs { maxSuccess = n } diff --git a/test/Spec.hs b/test/Spec.hs --- a/test/Spec.hs +++ b/test/Spec.hs @@ -5,11 +5,16 @@ import qualified Numeric.DataFrame.BasicTest import qualified Numeric.DataFrame.SubSpaceTest +import qualified Numeric.Matrix.BidiagonalTest +import qualified Numeric.Matrix.LUTest +import qualified Numeric.Matrix.QRTest +import qualified Numeric.Matrix.SVDTest import qualified Numeric.MatrixDoubleTest import qualified Numeric.MatrixFloatTest import qualified Numeric.PrimBytesTest import qualified Numeric.QuaterDoubleTest import qualified Numeric.QuaterFloatTest +import qualified Numeric.Subroutine.SolveTriangularTest -- | Collection of tests in detailed-0.9 format @@ -22,9 +27,15 @@ , test "QuaterDouble" $ Numeric.QuaterDoubleTest.runTests n , test "QuaterFloat" $ Numeric.QuaterFloatTest.runTests n , test "PrimBytes" $ Numeric.PrimBytesTest.runTests n + , test "Matrix.LU" $ Numeric.Matrix.LUTest.runTests n + , test "Matrix.QR" $ Numeric.Matrix.QRTest.runTests n + , test "Matrix.Bidiagonal" $ Numeric.Matrix.BidiagonalTest.runTests n + , test "Matrix.SVD" $ Numeric.Matrix.SVDTest.runTests n + , test "Subroutine.SolveTriangular" + $ Numeric.Subroutine.SolveTriangularTest.runTests n ] where - n = 1000 :: Int + n = 1000 :: Int