{-# LANGUAGE DataKinds #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UnboxedTuples #-}
{-# LANGUAGE UndecidableInstances #-}
{-# OPTIONS_GHC -fno-warn-overlapping-patterns #-}
module Numeric.DataFrame.Shape
( (<:>), (<::>), (<+:>)
, fromScalar
, DataFrameToList (..)
, fromListN, fromList
) where
import GHC.Base
import Numeric.DataFrame.Internal.Array.Class
import Numeric.DataFrame.Internal.Array.Family (inferASing, inferPrim,
inferPrimElem)
import Numeric.DataFrame.SubSpace
import Numeric.DataFrame.Type (DataFrame (..))
import Numeric.Dimensions
import Numeric.PrimBytes
import Numeric.Scalar as Scalar
import Numeric.TypedList (TypedList (..))
import qualified Numeric.TypedList as Dims
import Numeric.Vector
(<:>) :: forall (n :: Nat) (m :: Nat) (npm :: Nat) (ds :: [Nat])
(t :: Type)
. ( PrimBytes (DataFrame t (ds +: n))
, PrimBytes (DataFrame t (ds +: m))
, PrimBytes (DataFrame t (ds +: npm))
, npm ~ (n + m)
, n ~ (npm - m)
, m ~ (npm - n)
)
=> DataFrame t (ds +: n)
-> DataFrame t (ds +: m)
-> DataFrame t (ds +: npm)
(<:>) = appendDF
infixl 5 <:>
{-# INLINE [1] (<:>) #-}
(<::>) :: forall (ds :: [Nat]) (t :: Type)
. ( PrimBytes (DataFrame t ds)
, PrimBytes (DataFrame t ds)
, PrimBytes (DataFrame t (ds +: 2 :: [Nat]))
)
=> DataFrame t ds
-> DataFrame t ds
-> DataFrame t (ds +: 2 :: [Nat])
(<::>) = appendDF
infixl 5 <::>
{-# INLINE [1] (<::>) #-}
vec2t :: forall t . SubSpace t '[] '[2] '[2] => Scalar t -> Scalar t -> Vector t 2
vec2t = unsafeCoerce# (vec2 @t)
{-# INLINE vec2t #-}
{-# RULES
"<::>/vec2-Float" (<::>) = vec2t @Float
"<::>/vec2-Double" (<::>) = vec2t @Double
"<::>/vec2-Int" (<::>) = vec2t @Int
"<::>/vec2-Word" (<::>) = vec2t @Word
#-}
(<+:>) :: forall (ds :: [Nat]) (n :: Nat) (m :: Nat) (t :: Type)
. ( PrimBytes (DataFrame t (ds +: n))
, PrimBytes (DataFrame t ds)
, PrimBytes (DataFrame t (ds +: m))
, m ~ (n + 1)
)
=> DataFrame t (ds +: n)
-> DataFrame t ds
-> DataFrame t (ds +: m)
(<+:>) = appendDF
infixl 5 <+:>
{-# INLINE [1] (<+:>) #-}
fromList :: forall m ns t
. ( Dimensions ns
, PrimBytes t
)
=> Dim m
-> [DataFrame t ns]
-> Maybe (DataFrame t (AsXDims ns +: XN m))
fromList d xs = fromListN d (length xs) xs
class DataFrameToList (t :: Type) (ds :: [k]) (z :: k) where
toList :: DataFrame t (ds +: z) -> [DataFrame t ds]
instance ( Dimensions (ns +: z)
, PrimBytes t
)
=> DataFrameToList t (ns :: [Nat]) (z :: Nat) where
toList df
| Dims.Snoc (Dims :: Dims ns') dn <- dims @Nat @(ns +: z)
, E <- unsafeCoerce# (E @(ns ~ ns)) :: Evidence (ns ~ ns')
, E <- inferASing @t @ns
, E <- inferASing @t @(ns +: z)
, E <- inferPrim @t @(ns +: z)
, E <- inferPrim @t @ns
, n <- dimVal dn
, I# step <- fromIntegral $ totalDim' @ns
, off0 <- elemOffset df
, ba <- getBytes df
= let go 0 _ = []
go k off = fromElems off step ba : go (k-1) (off +# step)
in go n off0
| otherwise = []
instance DataFrameToList t (xns :: [XNat]) (xz :: XNat) where
toList (XFrame (df :: DataFrame t nsz))
| nsz <- dims @Nat @nsz
, xnsz <- xDims @(xns +: xz) nsz
, E <- unsafeCoerce# (E @(xns ~ xns)) :: Evidence (xns ~ Init (xns +: xz))
, xns <- Dims.init xnsz
, TypeList <- types xnsz
, EvList <- Dims.init (EvList @_ @KnownXNatType @(xns +: xz))
, Dims.Snoc (ns@Dims :: Dims ns) _ <- nsz
, Dims.Cons (_ :: Dim k) (_ :: Dims ks) <- nsz
, E <- inferPrimElem @t @k @ks
, XDims ns' <- xns
, Just E <- sameDims ns ns'
, E <- inferASing @t @ns
= map XFrame (toList df)
toList _ = []
fromListN :: forall (m :: Nat) (ns :: [Nat]) (t :: Type)
. ( Dimensions ns
, PrimBytes t
)
=> Dim m
-> Int
-> [DataFrame t ns]
-> Maybe (DataFrame t (AsXDims ns +: XN m))
fromListN Dim n@(I# n#) xs'
| n < 0 = Nothing
| Just dxn@(Dx dn@(D :: Dim n)) <- constrain @m (someDimVal (fromIntegral n))
, Just xs <- takeMaybe n xs'
, ns@(AsXDims xns) <- dims @Nat @ns
, nsn@Dims <- Dims.snoc ns dn
, xnsn <- Dims.snoc xns dxn
, EvList <- Dims.snoc (EvList @XNat @KnownXNatType @(AsXDims ns))
(toEvidence' (E @(KnownXNatType (XN m))))
, XDims nsn' <- xnsn
, Just E <- sameDims nsn nsn'
, E <- inferASing @t @ns
, E <- inferASing @t @(ns +: n)
, E <- inferPrim @t @ns
, E <- inferPrim @t @(ns +: n)
, I# partElN <- fromIntegral $ totalDim' @ns
, totalElN <- partElN *# n#
, elS <- byteSize @t undefined
, partBS <- partElN *# elS
= case runRW#
( \s0 -> case newByteArray# (totalElN *# elS) s0 of
(# s1, mba #) ->
let go _ [] s = s
go off (p : ps) s = go (off +# partBS) ps (writeBytes mba off p s)
in unsafeFreezeByteArray# mba (go 0# xs s1)
) of (# _, r #)
-> Just (XFrame (fromElems 0# totalElN r :: DataFrame t (ns +: n)))
fromListN _ _ _ = Nothing
takeMaybe :: Int -> [a] -> Maybe [a]
takeMaybe 0 _ = Just []
takeMaybe _ [] = Nothing
takeMaybe n (x : xs) = (x:) <$> takeMaybe (n-1) xs
appendDF :: (PrimBytes x, PrimBytes y, PrimBytes z)
=> x -> y -> z
appendDF 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
{-# INLINE appendDF #-}