arrayfire-0.9.0.0: src/ArrayFire/Data.hs
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE FlexibleContexts #-}
--------------------------------------------------------------------------------
-- |
-- Module : ArrayFire.Data
-- Copyright : David Johnson (c) 2019-2026
-- License : BSD 3
-- Maintainer : David Johnson <code@dmj.io>
-- Stability : Experimental
-- Portability : GHC
--
-- Functions for populating 'Array' with Data.
--
-- @
-- >>> constant @Double [2,2] 2.0
-- ArrayFire Array
-- [2 2 1 1]
-- 2.0000 2.0000
-- 2.0000 2.0000
-- @
--
--------------------------------------------------------------------------------
module ArrayFire.Data where
import Control.Exception
import Control.Monad (when)
import Data.Complex
import Data.Int
import Data.Proxy
import Data.Word
import Foreign.C.Types
import Foreign.ForeignPtr
import Foreign.Marshal hiding (void)
import Foreign.Ptr (Ptr)
import Foreign.Storable
import System.IO.Unsafe
import Unsafe.Coerce
import Data.Bits
import ArrayFire.Exception
import ArrayFire.FFI
import ArrayFire.Internal.Array (af_get_dims)
import ArrayFire.Internal.Data
import ArrayFire.Internal.Defines
import ArrayFire.Internal.Types
import ArrayFire.Arith
-- | Bitwise complement of every element in an 'Array'
--
-- >>> A.bitNot (A.scalar @Int32 0)
-- ArrayFire Array
-- [1 1 1 1]
-- -1
bitNot
:: (AFType a, Bits a)
=> Array a
-> Array a
bitNot arr = arr `bitXor` ones
where
(d0, d1, d2, d3) = arr `infoFromArray4` af_get_dims
ones = constant
[ fromIntegral d0
, fromIntegral d1
, fromIntegral d2
, fromIntegral d3
]
(complement zeroBits)
-- | Creates an 'Array' from a scalar value from given dimensions
--
-- >>> constant @Double [2,2] 2.0
-- ArrayFire Array
-- [2 2 1 1]
-- 2.0000 2.0000
-- 2.0000 2.0000
constant
:: forall a . AFType a
=> [Int]
-- ^ Dimensions
-> a
-- ^ Scalar value
-> Array a
{-# NOINLINE constant #-}
constant dims val =
case dtyp of
x | x == c64 ->
cast $ constantComplex dims (unsafeCoerce val :: Complex Double)
| x == c32 ->
cast $ constantComplex dims (unsafeCoerce val :: Complex Float)
| x == s64 ->
cast $ constantLong dims (unsafeCoerce val :: Int)
| x == u64 ->
cast $ constantULong dims (unsafeCoerce val :: Word64)
| x == s32 ->
constant' dims (fromIntegral (unsafeCoerce val :: Int32) :: Double)
| x == s16 ->
constant' dims (fromIntegral (unsafeCoerce val :: Int16) :: Double)
| x == u32 ->
constant' dims (fromIntegral (unsafeCoerce val :: Word32) :: Double)
| x == u8 ->
constant' dims (fromIntegral (unsafeCoerce val :: Word8) :: Double)
| x == u16 ->
constant' dims (fromIntegral (unsafeCoerce val :: Word16) :: Double)
| x == f64 ->
constant' dims (unsafeCoerce val :: Double)
| x == b8 ->
constant' dims (fromIntegral (unsafeCoerce val :: CBool) :: Double)
| x == f32 ->
constant' dims (realToFrac (unsafeCoerce val :: Float))
| otherwise -> error "constant: Invalid array fire type"
where
dtyp = afType (Proxy @a)
-- Creates the array directly with the target dtype: @af_constant@ takes
-- the value as a C double for every non-complex, non-64-bit-integral
-- dtype. Routing through an f64 array and casting (as this used to do)
-- fails with AF_ERR_NO_DBL on OpenCL devices without fp64 support and
-- changes b8 semantics (the cast normalises non-zero values to 1).
constant'
:: [Int]
-- ^ Dimensions
-> Double
-- ^ Scalar value
-> Array a
constant' dims' val' =
unsafePerformIO . mask_ $ do
ptr <- calloca $ \ptrPtr -> do
withArray (fromIntegral <$> dims') $ \dimArray -> do
throwAFError =<< af_constant ptrPtr val' n dimArray dtyp
peek ptrPtr
Array <$>
newForeignPtr
af_release_array_finalizer
ptr
where
n = fromIntegral (length dims')
-- | Creates an 'Array (Complex Double)' from a scalar val'ue
--
-- @
-- >>> constantComplex [2,2] (2.0 :+ 2.0)
-- @
--
constantComplex
:: forall arr . (Real arr, AFType (Complex arr))
=> [Int]
-- ^ Dimensions
-> Complex arr
-- ^ Scalar val'ue
-> Array (Complex arr)
constantComplex dims' ((realToFrac -> x) :+ (realToFrac -> y)) = unsafePerformIO . mask_ $ do
ptr <- calloca $ \ptrPtr -> do
withArray (fromIntegral <$> dims') $ \dimArray -> do
throwAFError =<< af_constant_complex ptrPtr x y n dimArray typ
peek ptrPtr
Array <$>
newForeignPtr
af_release_array_finalizer
ptr
where
n = fromIntegral (length dims')
typ = afType (Proxy @(Complex arr))
-- | Creates an 'Array Int64' from a scalar val'ue
--
-- @
-- >>> constantLong [2,2] 2.0
-- @
--
constantLong
:: [Int]
-- ^ Dimensions
-> Int
-- ^ Scalar val'ue
-> Array Int
constantLong dims' val' = unsafePerformIO . mask_ $ do
ptr <- calloca $ \ptrPtr -> do
withArray (fromIntegral <$> dims') $ \dimArray -> do
throwAFError =<< af_constant_long ptrPtr (fromIntegral val') n dimArray
peek ptrPtr
Array <$>
newForeignPtr
af_release_array_finalizer
ptr
where
n = fromIntegral (length dims')
-- | Creates an 'Array Word64' from a scalar val'ue
--
-- @
-- >>> constantULong [2,2] 2.0
-- @
--
constantULong
:: [Int]
-> Word64
-> Array Word64
constantULong dims' val' = unsafePerformIO . mask_ $ do
ptr <- calloca $ \ptrPtr -> do
withArray (fromIntegral <$> dims') $ \dimArray -> do
throwAFError =<< af_constant_ulong ptrPtr (fromIntegral val') n dimArray
peek ptrPtr
Array <$>
newForeignPtr
af_release_array_finalizer
ptr
where
n = fromIntegral (length dims')
-- | Creates a range of values in an Array
--
-- >>> arange @Double [10] (-1)
-- ArrayFire Array
-- [10 1 1 1]
-- 0.0000
-- 1.0000
-- 2.0000
-- 3.0000
-- 4.0000
-- 5.0000
-- 6.0000
-- 7.0000
-- 8.0000
-- 9.0000
arange
:: forall a
. AFType a
=> [Int]
-> Int
-> Array a
{-# NOINLINE arange #-}
arange dims (fromIntegral -> k) = unsafePerformIO . mask_ $ do
ptr <- alloca $ \ptrPtr -> do
withArray (fromIntegral <$> dims) $ \dimArray -> do
throwAFError =<< af_range ptrPtr n dimArray k typ
peek ptrPtr
Array <$>
newForeignPtr
af_release_array_finalizer
ptr
where
n = fromIntegral (length dims)
typ = afType (Proxy @a)
-- | Create an sequence [0, dims.elements() - 1] and modify to specified dimensions dims and then tile it according to tile_dims.
--
-- <http://arrayfire.org/docs/group__data__func__iota.htm>
--
-- >>> iota @Double [5,3] []
-- ArrayFire Array
-- [5 3 1 1]
-- 0.0000 5.0000 10.0000
-- 1.0000 6.0000 11.0000
-- 2.0000 7.0000 12.0000
-- 3.0000 8.0000 13.0000
-- 4.0000 9.0000 14.0000
--
-- >>> iota @Double [5,3] [1,2]
-- ArrayFire Array
-- [5 6 1 1]
-- 0.0000 5.0000 10.0000 0.0000 5.0000 10.0000
-- 1.0000 6.0000 11.0000 1.0000 6.0000 11.0000
-- 2.0000 7.0000 12.0000 2.0000 7.0000 12.0000
-- 3.0000 8.0000 13.0000 3.0000 8.0000 13.0000
-- 4.0000 9.0000 14.0000 4.0000 9.0000 14.0000
iota
:: forall a . AFType a
=> [Int]
-- ^ is the array containing sizes of the dimension
-> [Int]
-- ^ is array containing the number of repetitions of the unit dimensions
-> Array a
-- ^ is the generated array
{-# NOINLINE iota #-}
iota dims tdims = unsafePerformIO . mask_ $ do
let dims' = take 4 (dims ++ repeat 1)
tdims' = take 4 (tdims ++ repeat 1)
ptr <- calloca $ \ptrPtr -> do
withArray (fromIntegral <$> dims') $ \dimArray ->
withArray (fromIntegral <$> tdims') $ \tdimArray -> do
throwAFError =<< af_iota ptrPtr 4 dimArray 4 tdimArray typ
peek ptrPtr
Array <$>
newForeignPtr
af_release_array_finalizer
ptr
where
typ = afType (Proxy @a)
-- | Creates the identity `Array` from given dimensions
--
-- >>> identity [2,2]
-- ArrayFire Array
-- [2 2 1 1]
-- 1.0000 0.0000
-- 0.0000 1.0000
identity
:: forall a . AFType a
=> [Int]
-- ^ Dimensions
-> Array a
{-# NOINLINE identity #-}
identity dims = unsafePerformIO . mask_ $ do
when (length dims > 4) $
throwIO AFException
{ afExceptionType = ArgError
, afExceptionCode = 202
, afExceptionMsg = "identity: ndims must be <= 4"
}
let dims' = take 4 (dims ++ repeat 1)
ptr <- calloca $ \ptrPtr -> mask_ $ do
withArray (fromIntegral <$> dims') $ \dimArray -> do
throwAFError =<< af_identity ptrPtr n dimArray typ
peek ptrPtr
Array <$>
newForeignPtr
af_release_array_finalizer
ptr
where
n = fromIntegral (length dims)
typ = afType (Proxy @a)
-- | Create a diagonal matrix from input array when extract is set to false
--
-- >>> diagCreate (vector @Double 2 [1..]) 0
-- ArrayFire Array
-- [2 2 1 1]
-- 1.0000 0.0000
-- 0.0000 2.0000
diagCreate
:: AFType a
=> Array a
-- ^ is the input array which is the diagonal
-> Int
-- ^ is the diagonal index
-> Array a
diagCreate x (fromIntegral -> n) =
x `op1` (\p a -> af_diag_create p a n)
-- | Create a diagonal matrix from input array when extract is set to false
--
-- >>> diagExtract (matrix @Double (2,2) [[1,2],[3,4]]) 0
-- ArrayFire Array
-- [2 1 1 1]
-- 1.0000
-- 4.0000
diagExtract
:: AFType a
=> Array a
-> Int
-> Array a
diagExtract x (fromIntegral -> n) =
x `op1` (\p a -> af_diag_extract p a n)
-- | Join two Arrays together along a specified dimension
--
-- >>> join 0 (matrix @Double (2,2) [[1,2],[3,4]]) (matrix @Double (2,2) [[5,6],[7,8]])
-- ArrayFire Array
-- [4 2 1 1]
-- 1.0000 3.0000
-- 2.0000 4.0000
-- 5.0000 7.0000
-- 6.0000 8.0000
--
join
:: Int
-> Array a
-> Array a
-> Array a
join (fromIntegral -> n) arr1 arr2 = op2 arr1 arr2 (\p a b -> af_join p n a b)
-- | Join many Arrays together along a specified dimension
--
-- >>> joinMany 0 [vector @Int 3 [1..], vector @Int 3 [1..]]
-- ArrayFire Array
-- [6 1 1 1]
-- 1
-- 2
-- 3
-- 1
-- 2
-- 3
joinMany
:: Int
-> [Array a]
-> Array a
{-# NOINLINE joinMany #-}
joinMany (fromIntegral -> n) (fmap (\(Array fp) -> fp) -> arrays) = unsafePerformIO . mask_ $ do
newPtr <- calloca $ \aPtr -> do
(throwAFError =<<) $
withManyForeignPtr arrays $ \(fromIntegral -> nArrays) fPtrsPtr ->
af_join_many aPtr n nArrays fPtrsPtr
peek aPtr
Array <$>
newForeignPtr af_release_array_finalizer newPtr
-- | Marshals a list of 'ForeignPtr' into a temporary, contiguous C array of
-- raw pointers, keeping every 'ForeignPtr' alive for the duration of the
-- action. The continuation receives the number of pointers and a pointer to
-- the array.
withManyForeignPtr :: [ForeignPtr a] -> (Int -> Ptr (Ptr a) -> IO b) -> IO b
withManyForeignPtr fptrs action = go [] fptrs
where
go ptrs [] = withArrayLen (reverse ptrs) action
go ptrs (fptr:others) = withForeignPtr fptr $ \ptr -> go (ptr : ptrs) others
-- | Tiles an Array according to specified dimensions
--
-- >>> tile @Double (scalar 22.0) [5,5]
-- ArrayFire Array
-- [5 5 1 1]
-- 22.0000 22.0000 22.0000 22.0000 22.0000
-- 22.0000 22.0000 22.0000 22.0000 22.0000
-- 22.0000 22.0000 22.0000 22.0000 22.0000
-- 22.0000 22.0000 22.0000 22.0000 22.0000
-- 22.0000 22.0000 22.0000 22.0000 22.0000
--
tile
:: Array a
-> [Int]
-> Array a
tile a (take 4 . (++repeat 1) -> [x,y,z,w]) =
a `op1` (\p k -> af_tile p k (fromIntegral x) (fromIntegral y) (fromIntegral z) (fromIntegral w))
tile _ _ = error "impossible"
-- | Reorders an Array according to newly specified dimensions
--
-- *FIX ME*
--
-- >>> reorder @Double (scalar 22.0) [5,5]
-- ArrayFire Array
-- [5 5 1 1]
-- 22.0000 22.0000 22.0000 22.0000 22.0000
-- 22.0000 22.0000 22.0000 22.0000 22.0000
-- 22.0000 22.0000 22.0000 22.0000 22.0000
-- 22.0000 22.0000 22.0000 22.0000 22.0000
-- 22.0000 22.0000 22.0000 22.0000 22.0000
--
reorder
:: Array a
-> [Int]
-> Array a
reorder a dims =
let base = take 4 dims
padding = filter (`notElem` base) [0..3]
in case take 4 (base ++ padding) of
[x,y,z,w] -> a `op1` (\p k -> af_reorder p k (fromIntegral x) (fromIntegral y) (fromIntegral z) (fromIntegral w))
_ -> error "impossible"
-- | Shift elements in an Array along a specified dimension (elements will wrap).
--
-- >>> shift (vector @Double 4 [1..]) 2 0 0 0
-- ArrayFire Array
-- [4 1 1 1]
-- 3.0000
-- 4.0000
-- 1.0000
-- 2.0000
--
shift
:: Array a
-> Int
-> Int
-> Int
-> Int
-> Array a
shift a (fromIntegral -> x) (fromIntegral -> y) (fromIntegral -> z) (fromIntegral -> w) =
a `op1` (\p k -> af_shift p k x y z w)
-- | Modify dimensions of array
--
-- >>> moddims (vector @Double 3 [1..]) [1,3]
-- ArrayFire Array
-- [1 3 1 1]
-- 1.0000 2.0000 3.0000
--
moddims
:: Array a
-> [Int]
-> Array a
{-# NOINLINE moddims #-}
moddims (Array fptr) dims =
unsafePerformIO . mask_ . withForeignPtr fptr $ \ptr -> do
newPtr <- calloca $ \aPtr -> do
withArray (fromIntegral <$> dims) $ \dimsPtr -> do
throwAFError =<< af_moddims aPtr ptr n dimsPtr
peek aPtr
Array <$> newForeignPtr af_release_array_finalizer newPtr
where
n = fromIntegral (length dims)
-- | Flatten an Array into a single dimension
--
-- >>> flat (matrix @Double (2,2) [[1..],[1..]])
-- ArrayFire Array
-- [4 1 1 1]
-- 1.0000
-- 2.0000
-- 1.0000
-- 2.0000
--
-- >>> flat $ cube @Int (2,2,2) [[[1,1],[1,1]],[[1,1],[1,1]]]
-- ArrayFire Array
-- [8 1 1 1]
-- 1
-- 1
-- 1
-- 1
-- 1
-- 1
-- 1
-- 1
--
flat
:: Array a
-> Array a
flat = (`op1` af_flat)
-- | Flip the values of an Array along a specified dimension
--
-- >>> matrix @Double (2,2) [[2,2],[3,3]]
-- ArrayFire Array
-- [2 2 1 1]
-- 2.0000 3.0000
-- 2.0000 3.0000
--
-- >>> A.flip (matrix @Double (2,2) [[2,2],[3,3]]) 1
-- ArrayFire Array
-- [2 2 1 1]
-- 3.0000 2.0000
-- 3.0000 2.0000
--
flip
:: Array a
-> Int
-> Array a
flip a (fromIntegral -> dim) =
a `op1` (\p k -> af_flip p k dim)
-- | Create a lower triangular matrix from input array.
--
-- >>> lower (constant [2,2] 10 :: Array Double) True
-- ArrayFire Array
-- [2 2 1 1]
-- 1.0000 0.0000
-- 10.0000 1.0000
--
lower
:: Array a
-- ^ is the input matrix
-> Bool
-- ^ boolean parameter specifying if the diagonal elements should be 1
-> Array a
lower a (fromIntegral . fromEnum -> b) =
a `op1` (\p k -> af_lower p k b)
-- | Create an upper triangular matrix from input array.
--
-- >>> upper (constant [2,2] 10 :: Array Double) True
-- ArrayFire Array
-- [2 2 1 1]
-- 1.0000 10.0000
-- 0.0000 1.0000
--
upper
:: Array a
-> Bool
-> Array a
upper a (fromIntegral . fromEnum -> b) =
a `op1` (\p k -> af_upper p k b)
-- | Selects elements from two arrays based on the values of a binary conditional array.
--
-- >>> cond = vector @CBool 5 [1,0,1,0,1]
-- >>> arr1 = vector @Double 5 (repeat 1)
-- >>> arr2 = vector @Double 5 (repeat 2)
-- >>> select cond arr1 arr2
-- ArrayFire Array
-- [5 1 1 1]
-- 1.0000
-- 2.0000
-- 1.0000
-- 2.0000
-- 1.0000
--
select
:: Array CBool
-- ^ is the conditional array
-> Array a
-- ^ is the array containing elements from the true part of the condition
-> Array a
-- ^ is the array containing elements from the false part of the condition
-> Array a
-- ^ is the output containing elements of a when cond is true else elements from b
select a b c = op3 a b c af_select
-- | Selects elements from two arrays based on the values of a binary conditional array.
--
-- <http://arrayfire.org/docs/group__data__func__select.htm#gab6886120d0bac4717276910e468bbe88>
--
-- >>> cond = vector @CBool 5 [1,0,1,0,1]
-- >>> arr1 = vector @Double 5 (repeat 1)
-- >>> x = 99
-- >>> selectScalarR cond arr1 x
-- ArrayFire Array
-- [5 1 1 1]
-- 1.0000
-- 99.0000
-- 1.0000
-- 99.0000
-- 1.0000
--
selectScalarR
:: Array CBool
-- ^ is the conditional array
-> Array a
-- ^ is the array containing elements from the true part of the condition
-> Double
-- ^ is a scalar assigned to out when cond is false
-> Array a
-- ^ the output containing elements of a when cond is true else elements from b
selectScalarR a b c = op2 a b (\p w x -> af_select_scalar_r p w x c)
-- | Selects elements from two arrays based on the values of a binary conditional array.
--
-- [ArrayFire Docs](http://arrayfire.org/docs/group__data__func__select.htm#ga0ccdc05779f88cab5095bce987c2da9d)
--
-- >>> cond = vector @CBool 5 [1,0,1,0,1]
-- >>> arr1 = vector @Double 5 (repeat 1)
-- >>> x = 99
-- >>> selectScalarL cond x arr1
-- ArrayFire Array
-- [5 1 1 1]
-- 99.0000
-- 1.0000
-- 99.0000
-- 1.0000
-- 99.0000
--
selectScalarL
:: Array CBool
-- ^ the conditional array
-> Double
-- ^ a scalar assigned to out when cond is true
-> Array a
-- ^ the array containing elements from the false part of the condition
-> Array a
-- ^ is the output containing elements of a when cond is true else elements from b
selectScalarL a n b = op2 a b (\p w x -> af_select_scalar_l p w n x)
-- af_err af_replace(af_array a, const af_array cond, const af_array b);
-- af_err af_replace_scalar(af_array a, const af_array cond, const double b);