{-# 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 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
  :: forall a. (AFType a, Bits a, Integral a)
  => Array a
  -> Array a
bitNot :: forall a. (AFType a, Bits a, Integral a) => Array a -> Array a
bitNot Array a
arr = Array a
arr Array a -> Array a -> Array a
forall a. AFType a => Array a -> Array a -> Array a
`bitXor` Array a
ones
  where
    (DimT
d0, DimT
d1, DimT
d2, DimT
d3) = Array a
arr Array a
-> (Ptr DimT
    -> Ptr DimT -> Ptr DimT -> Ptr DimT -> AFArray -> IO AFErr)
-> (DimT, DimT, DimT, DimT)
forall a b c d arr.
(Storable a, Storable b, Storable c, Storable d) =>
Array arr
-> (Ptr a -> Ptr b -> Ptr c -> Ptr d -> AFArray -> IO AFErr)
-> (a, b, c, d)
`infoFromArray4` Ptr DimT -> Ptr DimT -> Ptr DimT -> Ptr DimT -> AFArray -> IO AFErr
af_get_dims
    ones :: Array a
ones = [Int] -> Double -> Array a
forall a. AFType a => [Int] -> Double -> Array a
constant
      [ DimT -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral DimT
d0
      , DimT -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral DimT
d1
      , DimT -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral DimT
d2
      , DimT -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral DimT
d3
      ]
      (a -> Double
forall a b. (Integral a, Num b) => a -> b
fromIntegral (a -> a
forall a. Bits a => a -> a
complement (a
forall a. Bits a => a
zeroBits :: a)))

-- | Creates a constant 'Array' filled with a 'Double' scalar.
-- ArrayFire converts the value to the element type internally.
-- Use 'constantComplex' for complex arrays, 'constantLong' / 'constantULong'
-- for 64-bit integer arrays where the value exceeds 2^53.
constant
  :: forall a. AFType a
  => [Int]   -- ^ Dimensions
  -> Double  -- ^ Scalar value
  -> Array a
{-# NOINLINE constant #-}
constant :: forall a. AFType a => [Int] -> Double -> Array a
constant [Int]
dims Double
val =
  IO (Array a) -> Array a
forall a. IO a -> a
unsafePerformIO (IO (Array a) -> Array a)
-> (IO (Array a) -> IO (Array a)) -> IO (Array a) -> Array a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IO (Array a) -> IO (Array a)
forall a. IO a -> IO a
mask_ (IO (Array a) -> Array a) -> IO (Array a) -> Array a
forall a b. (a -> b) -> a -> b
$ do
    ptr <- (Ptr AFArray -> IO AFArray) -> IO AFArray
forall a b. Storable a => (Ptr a -> IO b) -> IO b
calloca ((Ptr AFArray -> IO AFArray) -> IO AFArray)
-> (Ptr AFArray -> IO AFArray) -> IO AFArray
forall a b. (a -> b) -> a -> b
$ \Ptr AFArray
ptrPtr -> do
      [DimT] -> (Ptr DimT -> IO AFArray) -> IO AFArray
forall a b. Storable a => [a] -> (Ptr a -> IO b) -> IO b
withArray (Int -> DimT
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> DimT) -> [Int] -> [DimT]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Int]
dims) ((Ptr DimT -> IO AFArray) -> IO AFArray)
-> (Ptr DimT -> IO AFArray) -> IO AFArray
forall a b. (a -> b) -> a -> b
$ \Ptr DimT
dimArray -> do
        AFErr -> IO ()
throwAFError (AFErr -> IO ()) -> IO AFErr -> IO ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Ptr AFArray -> Double -> CUInt -> Ptr DimT -> AFDtype -> IO AFErr
af_constant Ptr AFArray
ptrPtr Double
val CUInt
n Ptr DimT
dimArray AFDtype
dtyp
        Ptr AFArray -> IO AFArray
forall a. Storable a => Ptr a -> IO a
peek Ptr AFArray
ptrPtr
    Array <$> newForeignPtr af_release_array_finalizer ptr
  where
    n :: CUInt
n    = Int -> CUInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
dims)
    dtyp :: AFDtype
dtyp = Proxy a -> AFDtype
forall a. AFType a => Proxy a -> AFDtype
afType (forall t. Proxy t
forall {k} (t :: k). Proxy t
Proxy @a)

-- | Creates a constant complex 'Array' from a 'Complex' scalar.
constantComplex
  :: forall r. (Real r, AFType (Complex r))
  => [Int]      -- ^ Dimensions
  -> Complex r  -- ^ Scalar value
  -> Array (Complex r)
{-# NOINLINE constantComplex #-}
constantComplex :: forall r.
(Real r, AFType (Complex r)) =>
[Int] -> Complex r -> Array (Complex r)
constantComplex [Int]
dims ((r -> Double
forall a b. (Real a, Fractional b) => a -> b
realToFrac -> Double
x) :+ (r -> Double
forall a b. (Real a, Fractional b) => a -> b
realToFrac -> Double
y)) =
  IO (Array (Complex r)) -> Array (Complex r)
forall a. IO a -> a
unsafePerformIO (IO (Array (Complex r)) -> Array (Complex r))
-> (IO (Array (Complex r)) -> IO (Array (Complex r)))
-> IO (Array (Complex r))
-> Array (Complex r)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IO (Array (Complex r)) -> IO (Array (Complex r))
forall a. IO a -> IO a
mask_ (IO (Array (Complex r)) -> Array (Complex r))
-> IO (Array (Complex r)) -> Array (Complex r)
forall a b. (a -> b) -> a -> b
$ do
    ptr <- (Ptr AFArray -> IO AFArray) -> IO AFArray
forall a b. Storable a => (Ptr a -> IO b) -> IO b
calloca ((Ptr AFArray -> IO AFArray) -> IO AFArray)
-> (Ptr AFArray -> IO AFArray) -> IO AFArray
forall a b. (a -> b) -> a -> b
$ \Ptr AFArray
ptrPtr -> do
      [DimT] -> (Ptr DimT -> IO AFArray) -> IO AFArray
forall a b. Storable a => [a] -> (Ptr a -> IO b) -> IO b
withArray (Int -> DimT
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> DimT) -> [Int] -> [DimT]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Int]
dims) ((Ptr DimT -> IO AFArray) -> IO AFArray)
-> (Ptr DimT -> IO AFArray) -> IO AFArray
forall a b. (a -> b) -> a -> b
$ \Ptr DimT
dimArray -> do
        AFErr -> IO ()
throwAFError (AFErr -> IO ()) -> IO AFErr -> IO ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Ptr AFArray
-> Double -> Double -> CUInt -> Ptr DimT -> AFDtype -> IO AFErr
af_constant_complex Ptr AFArray
ptrPtr Double
x Double
y CUInt
n Ptr DimT
dimArray AFDtype
typ
        Ptr AFArray -> IO AFArray
forall a. Storable a => Ptr a -> IO a
peek Ptr AFArray
ptrPtr
    Array <$> newForeignPtr af_release_array_finalizer ptr
  where
    n :: CUInt
n   = Int -> CUInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
dims)
    typ :: AFDtype
typ = Proxy (Complex r) -> AFDtype
forall a. AFType a => Proxy a -> AFDtype
afType (forall t. Proxy t
forall {k} (t :: k). Proxy t
Proxy @(Complex r))

-- | Creates a constant 'Array' of 64-bit signed integers.
-- Preserves the full integer value without 'Double' rounding.
constantLong
  :: [Int]  -- ^ Dimensions
  -> Int    -- ^ Scalar value
  -> Array Int
{-# NOINLINE constantLong #-}
constantLong :: [Int] -> Int -> Array Int
constantLong [Int]
dims Int
val =
  IO (Array Int) -> Array Int
forall a. IO a -> a
unsafePerformIO (IO (Array Int) -> Array Int)
-> (IO (Array Int) -> IO (Array Int))
-> IO (Array Int)
-> Array Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IO (Array Int) -> IO (Array Int)
forall a. IO a -> IO a
mask_ (IO (Array Int) -> Array Int) -> IO (Array Int) -> Array Int
forall a b. (a -> b) -> a -> b
$ do
    ptr <- (Ptr AFArray -> IO AFArray) -> IO AFArray
forall a b. Storable a => (Ptr a -> IO b) -> IO b
calloca ((Ptr AFArray -> IO AFArray) -> IO AFArray)
-> (Ptr AFArray -> IO AFArray) -> IO AFArray
forall a b. (a -> b) -> a -> b
$ \Ptr AFArray
ptrPtr -> do
      [DimT] -> (Ptr DimT -> IO AFArray) -> IO AFArray
forall a b. Storable a => [a] -> (Ptr a -> IO b) -> IO b
withArray (Int -> DimT
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> DimT) -> [Int] -> [DimT]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Int]
dims) ((Ptr DimT -> IO AFArray) -> IO AFArray)
-> (Ptr DimT -> IO AFArray) -> IO AFArray
forall a b. (a -> b) -> a -> b
$ \Ptr DimT
dimArray -> do
        AFErr -> IO ()
throwAFError (AFErr -> IO ()) -> IO AFErr -> IO ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Ptr AFArray -> IntL -> CUInt -> Ptr DimT -> IO AFErr
af_constant_long Ptr AFArray
ptrPtr (Int -> IntL
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
val) CUInt
n Ptr DimT
dimArray
        Ptr AFArray -> IO AFArray
forall a. Storable a => Ptr a -> IO a
peek Ptr AFArray
ptrPtr
    Array <$> newForeignPtr af_release_array_finalizer ptr
  where n :: CUInt
n = Int -> CUInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
dims)

-- | Creates a constant 'Array' of 64-bit unsigned integers.
-- Preserves the full integer value without 'Double' rounding.
constantULong
  :: [Int]   -- ^ Dimensions
  -> Word64  -- ^ Scalar value
  -> Array Word64
{-# NOINLINE constantULong #-}
constantULong :: [Int] -> Word64 -> Array Word64
constantULong [Int]
dims Word64
val =
  IO (Array Word64) -> Array Word64
forall a. IO a -> a
unsafePerformIO (IO (Array Word64) -> Array Word64)
-> (IO (Array Word64) -> IO (Array Word64))
-> IO (Array Word64)
-> Array Word64
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IO (Array Word64) -> IO (Array Word64)
forall a. IO a -> IO a
mask_ (IO (Array Word64) -> Array Word64)
-> IO (Array Word64) -> Array Word64
forall a b. (a -> b) -> a -> b
$ do
    ptr <- (Ptr AFArray -> IO AFArray) -> IO AFArray
forall a b. Storable a => (Ptr a -> IO b) -> IO b
calloca ((Ptr AFArray -> IO AFArray) -> IO AFArray)
-> (Ptr AFArray -> IO AFArray) -> IO AFArray
forall a b. (a -> b) -> a -> b
$ \Ptr AFArray
ptrPtr -> do
      [DimT] -> (Ptr DimT -> IO AFArray) -> IO AFArray
forall a b. Storable a => [a] -> (Ptr a -> IO b) -> IO b
withArray (Int -> DimT
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> DimT) -> [Int] -> [DimT]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Int]
dims) ((Ptr DimT -> IO AFArray) -> IO AFArray)
-> (Ptr DimT -> IO AFArray) -> IO AFArray
forall a b. (a -> b) -> a -> b
$ \Ptr DimT
dimArray -> do
        AFErr -> IO ()
throwAFError (AFErr -> IO ()) -> IO AFErr -> IO ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Ptr AFArray -> UIntL -> CUInt -> Ptr DimT -> IO AFErr
af_constant_ulong Ptr AFArray
ptrPtr (Word64 -> UIntL
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
val) CUInt
n Ptr DimT
dimArray
        Ptr AFArray -> IO AFArray
forall a. Storable a => Ptr a -> IO a
peek Ptr AFArray
ptrPtr
    Array <$> newForeignPtr af_release_array_finalizer ptr
  where n :: CUInt
n = Int -> CUInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
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 :: forall a. AFType a => [Int] -> Int -> Array a
arange [Int]
dims (Int -> CInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral -> CInt
k) = IO (Array a) -> Array a
forall a. IO a -> a
unsafePerformIO (IO (Array a) -> Array a)
-> (IO (Array a) -> IO (Array a)) -> IO (Array a) -> Array a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IO (Array a) -> IO (Array a)
forall a. IO a -> IO a
mask_ (IO (Array a) -> Array a) -> IO (Array a) -> Array a
forall a b. (a -> b) -> a -> b
$ do
  ptr <- (Ptr AFArray -> IO AFArray) -> IO AFArray
forall a b. Storable a => (Ptr a -> IO b) -> IO b
alloca ((Ptr AFArray -> IO AFArray) -> IO AFArray)
-> (Ptr AFArray -> IO AFArray) -> IO AFArray
forall a b. (a -> b) -> a -> b
$ \Ptr AFArray
ptrPtr -> do
    [DimT] -> (Ptr DimT -> IO AFArray) -> IO AFArray
forall a b. Storable a => [a] -> (Ptr a -> IO b) -> IO b
withArray (Int -> DimT
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> DimT) -> [Int] -> [DimT]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Int]
dims) ((Ptr DimT -> IO AFArray) -> IO AFArray)
-> (Ptr DimT -> IO AFArray) -> IO AFArray
forall a b. (a -> b) -> a -> b
$ \Ptr DimT
dimArray -> do
      AFErr -> IO ()
throwAFError (AFErr -> IO ()) -> IO AFErr -> IO ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Ptr AFArray -> CUInt -> Ptr DimT -> CInt -> AFDtype -> IO AFErr
af_range Ptr AFArray
ptrPtr CUInt
n Ptr DimT
dimArray CInt
k AFDtype
typ
      Ptr AFArray -> IO AFArray
forall a. Storable a => Ptr a -> IO a
peek Ptr AFArray
ptrPtr
  Array <$>
    newForeignPtr
      af_release_array_finalizer
        ptr
      where
        n :: CUInt
n = Int -> CUInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
dims)
        typ :: AFDtype
typ = Proxy a -> AFDtype
forall a. AFType a => Proxy a -> AFDtype
afType (forall t. Proxy t
forall {k} (t :: k). Proxy t
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 :: forall a. AFType a => [Int] -> [Int] -> Array a
iota [Int]
dims [Int]
tdims = IO (Array a) -> Array a
forall a. IO a -> a
unsafePerformIO (IO (Array a) -> Array a)
-> (IO (Array a) -> IO (Array a)) -> IO (Array a) -> Array a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IO (Array a) -> IO (Array a)
forall a. IO a -> IO a
mask_ (IO (Array a) -> Array a) -> IO (Array a) -> Array a
forall a b. (a -> b) -> a -> b
$ do
  let dims' :: [Int]
dims' = Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
take Int
4 ([Int]
dims [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ Int -> [Int]
forall a. a -> [a]
repeat Int
1)
      tdims' :: [Int]
tdims' =  Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
take Int
4 ([Int]
tdims [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ Int -> [Int]
forall a. a -> [a]
repeat Int
1)
  ptr <- (Ptr AFArray -> IO AFArray) -> IO AFArray
forall a b. Storable a => (Ptr a -> IO b) -> IO b
calloca ((Ptr AFArray -> IO AFArray) -> IO AFArray)
-> (Ptr AFArray -> IO AFArray) -> IO AFArray
forall a b. (a -> b) -> a -> b
$ \Ptr AFArray
ptrPtr -> do
    [DimT] -> (Ptr DimT -> IO AFArray) -> IO AFArray
forall a b. Storable a => [a] -> (Ptr a -> IO b) -> IO b
withArray (Int -> DimT
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> DimT) -> [Int] -> [DimT]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Int]
dims') ((Ptr DimT -> IO AFArray) -> IO AFArray)
-> (Ptr DimT -> IO AFArray) -> IO AFArray
forall a b. (a -> b) -> a -> b
$ \Ptr DimT
dimArray ->
      [DimT] -> (Ptr DimT -> IO AFArray) -> IO AFArray
forall a b. Storable a => [a] -> (Ptr a -> IO b) -> IO b
withArray (Int -> DimT
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> DimT) -> [Int] -> [DimT]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Int]
tdims') ((Ptr DimT -> IO AFArray) -> IO AFArray)
-> (Ptr DimT -> IO AFArray) -> IO AFArray
forall a b. (a -> b) -> a -> b
$ \Ptr DimT
tdimArray -> do
        AFErr -> IO ()
throwAFError (AFErr -> IO ()) -> IO AFErr -> IO ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Ptr AFArray
-> CUInt -> Ptr DimT -> CUInt -> Ptr DimT -> AFDtype -> IO AFErr
af_iota Ptr AFArray
ptrPtr CUInt
4 Ptr DimT
dimArray CUInt
4 Ptr DimT
tdimArray AFDtype
typ
        Ptr AFArray -> IO AFArray
forall a. Storable a => Ptr a -> IO a
peek Ptr AFArray
ptrPtr
  Array <$>
    newForeignPtr
      af_release_array_finalizer
        ptr
      where
        typ :: AFDtype
typ = Proxy a -> AFDtype
forall a. AFType a => Proxy a -> AFDtype
afType (forall t. Proxy t
forall {k} (t :: k). Proxy t
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 :: forall a. AFType a => [Int] -> Array a
identity [Int]
dims = IO (Array a) -> Array a
forall a. IO a -> a
unsafePerformIO (IO (Array a) -> Array a)
-> (IO (Array a) -> IO (Array a)) -> IO (Array a) -> Array a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IO (Array a) -> IO (Array a)
forall a. IO a -> IO a
mask_ (IO (Array a) -> Array a) -> IO (Array a) -> Array a
forall a b. (a -> b) -> a -> b
$ do
  Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when ([Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
dims Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
4) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
    AFException -> IO ()
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO AFException
      { afExceptionType :: AFExceptionType
afExceptionType = AFExceptionType
ArgError
      , afExceptionCode :: Int
afExceptionCode = Int
202
      , afExceptionMsg :: String
afExceptionMsg  = String
"identity: ndims must be <= 4"
      }
  let dims' :: [Int]
dims' = Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
take Int
4 ([Int]
dims [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ Int -> [Int]
forall a. a -> [a]
repeat Int
1)
  ptr <- (Ptr AFArray -> IO AFArray) -> IO AFArray
forall a b. Storable a => (Ptr a -> IO b) -> IO b
calloca ((Ptr AFArray -> IO AFArray) -> IO AFArray)
-> (Ptr AFArray -> IO AFArray) -> IO AFArray
forall a b. (a -> b) -> a -> b
$ \Ptr AFArray
ptrPtr -> IO AFArray -> IO AFArray
forall a. IO a -> IO a
mask_ (IO AFArray -> IO AFArray) -> IO AFArray -> IO AFArray
forall a b. (a -> b) -> a -> b
$ do
    [DimT] -> (Ptr DimT -> IO AFArray) -> IO AFArray
forall a b. Storable a => [a] -> (Ptr a -> IO b) -> IO b
withArray (Int -> DimT
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> DimT) -> [Int] -> [DimT]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Int]
dims') ((Ptr DimT -> IO AFArray) -> IO AFArray)
-> (Ptr DimT -> IO AFArray) -> IO AFArray
forall a b. (a -> b) -> a -> b
$ \Ptr DimT
dimArray -> do
      AFErr -> IO ()
throwAFError (AFErr -> IO ()) -> IO AFErr -> IO ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Ptr AFArray -> CUInt -> Ptr DimT -> AFDtype -> IO AFErr
af_identity Ptr AFArray
ptrPtr CUInt
n Ptr DimT
dimArray AFDtype
typ
      Ptr AFArray -> IO AFArray
forall a. Storable a => Ptr a -> IO a
peek Ptr AFArray
ptrPtr
  Array <$>
    newForeignPtr
      af_release_array_finalizer
        ptr
      where
        n :: CUInt
n = Int -> CUInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
dims)
        typ :: AFDtype
typ = Proxy a -> AFDtype
forall a. AFType a => Proxy a -> AFDtype
afType (forall t. Proxy t
forall {k} (t :: k). Proxy t
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 :: forall a. AFType a => Array a -> Int -> Array a
diagCreate Array a
x (Int -> CInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral -> CInt
n) =
  Array a
x Array a -> (Ptr AFArray -> AFArray -> IO AFErr) -> Array a
forall a b.
Array a -> (Ptr AFArray -> AFArray -> IO AFErr) -> Array b
`op1` (\Ptr AFArray
p AFArray
a -> Ptr AFArray -> AFArray -> CInt -> IO AFErr
af_diag_create Ptr AFArray
p AFArray
a CInt
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 :: forall a. AFType a => Array a -> Int -> Array a
diagExtract Array a
x (Int -> CInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral -> CInt
n) =
  Array a
x Array a -> (Ptr AFArray -> AFArray -> IO AFErr) -> Array a
forall a b.
Array a -> (Ptr AFArray -> AFArray -> IO AFErr) -> Array b
`op1` (\Ptr AFArray
p AFArray
a -> Ptr AFArray -> AFArray -> CInt -> IO AFErr
af_diag_extract Ptr AFArray
p AFArray
a CInt
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 :: forall a. Int -> Array a -> Array a -> Array a
join (Int -> CInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral -> CInt
n) Array a
arr1 Array a
arr2 = Array a
-> Array a
-> (Ptr AFArray -> AFArray -> AFArray -> IO AFErr)
-> Array a
forall b a c.
Array b
-> Array a
-> (Ptr AFArray -> AFArray -> AFArray -> IO AFErr)
-> Array c
op2 Array a
arr1 Array a
arr2 (\Ptr AFArray
p AFArray
a AFArray
b -> Ptr AFArray -> CInt -> AFArray -> AFArray -> IO AFErr
af_join Ptr AFArray
p CInt
n AFArray
a AFArray
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 :: forall a. Int -> [Array a] -> Array a
joinMany (Int -> CInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral -> CInt
n) ((Array a -> ForeignPtr ()) -> [Array a] -> [ForeignPtr ()]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(Array ForeignPtr ()
fp) -> ForeignPtr ()
fp) -> [ForeignPtr ()]
arrays) = IO (Array a) -> Array a
forall a. IO a -> a
unsafePerformIO (IO (Array a) -> Array a)
-> (IO (Array a) -> IO (Array a)) -> IO (Array a) -> Array a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IO (Array a) -> IO (Array a)
forall a. IO a -> IO a
mask_ (IO (Array a) -> Array a) -> IO (Array a) -> Array a
forall a b. (a -> b) -> a -> b
$ do
  newPtr <- (Ptr AFArray -> IO AFArray) -> IO AFArray
forall a b. Storable a => (Ptr a -> IO b) -> IO b
calloca ((Ptr AFArray -> IO AFArray) -> IO AFArray)
-> (Ptr AFArray -> IO AFArray) -> IO AFArray
forall a b. (a -> b) -> a -> b
$ \Ptr AFArray
aPtr -> do
    (AFErr -> IO ()
throwAFError (AFErr -> IO ()) -> IO AFErr -> IO ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<<) (IO AFErr -> IO ()) -> IO AFErr -> IO ()
forall a b. (a -> b) -> a -> b
$
      [ForeignPtr ()] -> (Int -> Ptr AFArray -> IO AFErr) -> IO AFErr
forall a b. [ForeignPtr a] -> (Int -> Ptr (Ptr a) -> IO b) -> IO b
withManyForeignPtr [ForeignPtr ()]
arrays ((Int -> Ptr AFArray -> IO AFErr) -> IO AFErr)
-> (Int -> Ptr AFArray -> IO AFErr) -> IO AFErr
forall a b. (a -> b) -> a -> b
$ \(Int -> CUInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral -> CUInt
nArrays) Ptr AFArray
fPtrsPtr ->
        Ptr AFArray -> CInt -> CUInt -> Ptr AFArray -> IO AFErr
af_join_many Ptr AFArray
aPtr CInt
n CUInt
nArrays Ptr AFArray
fPtrsPtr
    Ptr AFArray -> IO AFArray
forall a. Storable a => Ptr a -> IO a
peek Ptr AFArray
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 :: forall a b. [ForeignPtr a] -> (Int -> Ptr (Ptr a) -> IO b) -> IO b
withManyForeignPtr [ForeignPtr a]
fptrs Int -> Ptr (Ptr a) -> IO b
action = [Ptr a] -> [ForeignPtr a] -> IO b
go [] [ForeignPtr a]
fptrs
  where
    go :: [Ptr a] -> [ForeignPtr a] -> IO b
go [Ptr a]
ptrs [] = [Ptr a] -> (Int -> Ptr (Ptr a) -> IO b) -> IO b
forall a b. Storable a => [a] -> (Int -> Ptr a -> IO b) -> IO b
withArrayLen ([Ptr a] -> [Ptr a]
forall a. [a] -> [a]
reverse [Ptr a]
ptrs) Int -> Ptr (Ptr a) -> IO b
action
    go [Ptr a]
ptrs (ForeignPtr a
fptr:[ForeignPtr a]
others) = ForeignPtr a -> (Ptr a -> IO b) -> IO b
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr a
fptr ((Ptr a -> IO b) -> IO b) -> (Ptr a -> IO b) -> IO b
forall a b. (a -> b) -> a -> b
$ \Ptr a
ptr -> [Ptr a] -> [ForeignPtr a] -> IO b
go (Ptr a
ptr Ptr a -> [Ptr a] -> [Ptr a]
forall a. a -> [a] -> [a]
: [Ptr a]
ptrs) [ForeignPtr a]
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 :: forall a. Array a -> [Int] -> Array a
tile Array a
a (Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
take Int
4 ([Int] -> [Int]) -> ([Int] -> [Int]) -> [Int] -> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++Int -> [Int]
forall a. a -> [a]
repeat Int
1) -> [Int
x,Int
y,Int
z,Int
w]) =
  Array a
a Array a -> (Ptr AFArray -> AFArray -> IO AFErr) -> Array a
forall a b.
Array a -> (Ptr AFArray -> AFArray -> IO AFErr) -> Array b
`op1` (\Ptr AFArray
p AFArray
k -> Ptr AFArray
-> AFArray -> CUInt -> CUInt -> CUInt -> CUInt -> IO AFErr
af_tile Ptr AFArray
p AFArray
k (Int -> CUInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
x) (Int -> CUInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
y) (Int -> CUInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
z) (Int -> CUInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
w))
tile Array a
_ [Int]
_ = String -> Array a
forall a. HasCallStack => String -> a
error String
"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 :: forall a. Array a -> [Int] -> Array a
reorder Array a
a [Int]
dims =
  let base :: [Int]
base    = Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
take Int
4 [Int]
dims
      padding :: [Int]
padding = (Int -> Bool) -> [Int] -> [Int]
forall a. (a -> Bool) -> [a] -> [a]
filter (Int -> [Int] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`notElem` [Int]
base) [Int
0..Int
3]
  in case Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
take Int
4 ([Int]
base [Int] -> [Int] -> [Int]
forall a. [a] -> [a] -> [a]
++ [Int]
padding) of
    [Int
x,Int
y,Int
z,Int
w] -> Array a
a Array a -> (Ptr AFArray -> AFArray -> IO AFErr) -> Array a
forall a b.
Array a -> (Ptr AFArray -> AFArray -> IO AFErr) -> Array b
`op1` (\Ptr AFArray
p AFArray
k -> Ptr AFArray
-> AFArray -> CUInt -> CUInt -> CUInt -> CUInt -> IO AFErr
af_reorder Ptr AFArray
p AFArray
k (Int -> CUInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
x) (Int -> CUInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
y) (Int -> CUInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
z) (Int -> CUInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
w))
    [Int]
_          -> String -> Array a
forall a. HasCallStack => String -> a
error String
"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 :: forall a. Array a -> Int -> Int -> Int -> Int -> Array a
shift Array a
a (Int -> CInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral -> CInt
x) (Int -> CInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral -> CInt
y) (Int -> CInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral -> CInt
z) (Int -> CInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral -> CInt
w) =
  Array a
a Array a -> (Ptr AFArray -> AFArray -> IO AFErr) -> Array a
forall a b.
Array a -> (Ptr AFArray -> AFArray -> IO AFErr) -> Array b
`op1` (\Ptr AFArray
p AFArray
k -> Ptr AFArray -> AFArray -> CInt -> CInt -> CInt -> CInt -> IO AFErr
af_shift Ptr AFArray
p AFArray
k CInt
x CInt
y CInt
z CInt
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 :: forall a. Array a -> [Int] -> Array a
moddims (Array ForeignPtr ()
fptr) [Int]
dims =
  IO (Array a) -> Array a
forall a. IO a -> a
unsafePerformIO (IO (Array a) -> Array a)
-> ((AFArray -> IO (Array a)) -> IO (Array a))
-> (AFArray -> IO (Array a))
-> Array a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IO (Array a) -> IO (Array a)
forall a. IO a -> IO a
mask_ (IO (Array a) -> IO (Array a))
-> ((AFArray -> IO (Array a)) -> IO (Array a))
-> (AFArray -> IO (Array a))
-> IO (Array a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ForeignPtr () -> (AFArray -> IO (Array a)) -> IO (Array a)
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr ()
fptr ((AFArray -> IO (Array a)) -> Array a)
-> (AFArray -> IO (Array a)) -> Array a
forall a b. (a -> b) -> a -> b
$ \AFArray
ptr -> do
    newPtr <- (Ptr AFArray -> IO AFArray) -> IO AFArray
forall a b. Storable a => (Ptr a -> IO b) -> IO b
calloca ((Ptr AFArray -> IO AFArray) -> IO AFArray)
-> (Ptr AFArray -> IO AFArray) -> IO AFArray
forall a b. (a -> b) -> a -> b
$ \Ptr AFArray
aPtr -> do
      [DimT] -> (Ptr DimT -> IO AFArray) -> IO AFArray
forall a b. Storable a => [a] -> (Ptr a -> IO b) -> IO b
withArray (Int -> DimT
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> DimT) -> [Int] -> [DimT]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Int]
dims) ((Ptr DimT -> IO AFArray) -> IO AFArray)
-> (Ptr DimT -> IO AFArray) -> IO AFArray
forall a b. (a -> b) -> a -> b
$ \Ptr DimT
dimsPtr -> do
        AFErr -> IO ()
throwAFError (AFErr -> IO ()) -> IO AFErr -> IO ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Ptr AFArray -> AFArray -> CUInt -> Ptr DimT -> IO AFErr
af_moddims Ptr AFArray
aPtr AFArray
ptr CUInt
n Ptr DimT
dimsPtr
        Ptr AFArray -> IO AFArray
forall a. Storable a => Ptr a -> IO a
peek Ptr AFArray
aPtr
    Array <$> newForeignPtr af_release_array_finalizer newPtr
  where
    n :: CUInt
n = Int -> CUInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
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 :: forall a. Array a -> Array a
flat = (Array a -> (Ptr AFArray -> AFArray -> IO AFErr) -> Array a
forall a b.
Array a -> (Ptr AFArray -> AFArray -> IO AFErr) -> Array b
`op1` Ptr AFArray -> AFArray -> IO AFErr
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 :: forall a. Array a -> Int -> Array a
flip Array a
a (Int -> CUInt
forall a b. (Integral a, Num b) => a -> b
fromIntegral -> CUInt
dim) =
  Array a
a Array a -> (Ptr AFArray -> AFArray -> IO AFErr) -> Array a
forall a b.
Array a -> (Ptr AFArray -> AFArray -> IO AFErr) -> Array b
`op1` (\Ptr AFArray
p AFArray
k -> Ptr AFArray -> AFArray -> CUInt -> IO AFErr
af_flip Ptr AFArray
p AFArray
k CUInt
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 :: forall a. Array a -> Bool -> Array a
lower Array a
a (Int -> CBool
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> CBool) -> (Bool -> Int) -> Bool -> CBool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bool -> Int
forall a. Enum a => a -> Int
fromEnum -> CBool
b) =
  Array a
a Array a -> (Ptr AFArray -> AFArray -> IO AFErr) -> Array a
forall a b.
Array a -> (Ptr AFArray -> AFArray -> IO AFErr) -> Array b
`op1` (\Ptr AFArray
p AFArray
k -> Ptr AFArray -> AFArray -> CBool -> IO AFErr
af_lower Ptr AFArray
p AFArray
k CBool
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 :: forall a. Array a -> Bool -> Array a
upper Array a
a (Int -> CBool
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> CBool) -> (Bool -> Int) -> Bool -> CBool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bool -> Int
forall a. Enum a => a -> Int
fromEnum -> CBool
b) =
  Array a
a Array a -> (Ptr AFArray -> AFArray -> IO AFErr) -> Array a
forall a b.
Array a -> (Ptr AFArray -> AFArray -> IO AFErr) -> Array b
`op1` (\Ptr AFArray
p AFArray
k -> Ptr AFArray -> AFArray -> CBool -> IO AFErr
af_upper Ptr AFArray
p AFArray
k CBool
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 :: forall a. Array CBool -> Array a -> Array a -> Array a
select Array CBool
a Array a
b Array a
c = Array CBool
-> Array a
-> Array a
-> (Ptr AFArray -> AFArray -> AFArray -> AFArray -> IO AFErr)
-> Array a
forall b a.
Array b
-> Array a
-> Array a
-> (Ptr AFArray -> AFArray -> AFArray -> AFArray -> IO AFErr)
-> Array a
op3 Array CBool
a Array a
b Array a
c Ptr AFArray -> AFArray -> AFArray -> AFArray -> IO AFErr
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 :: forall a. Array CBool -> Array a -> Double -> Array a
selectScalarR Array CBool
a Array a
b Double
c = Array CBool
-> Array a
-> (Ptr AFArray -> AFArray -> AFArray -> IO AFErr)
-> Array a
forall b a c.
Array b
-> Array a
-> (Ptr AFArray -> AFArray -> AFArray -> IO AFErr)
-> Array c
op2 Array CBool
a Array a
b (\Ptr AFArray
p AFArray
w AFArray
x -> Ptr AFArray -> AFArray -> AFArray -> Double -> IO AFErr
af_select_scalar_r Ptr AFArray
p AFArray
w AFArray
x Double
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 :: forall a. Array CBool -> Double -> Array a -> Array a
selectScalarL Array CBool
a Double
n Array a
b = Array CBool
-> Array a
-> (Ptr AFArray -> AFArray -> AFArray -> IO AFErr)
-> Array a
forall b a c.
Array b
-> Array a
-> (Ptr AFArray -> AFArray -> AFArray -> IO AFErr)
-> Array c
op2 Array CBool
a Array a
b (\Ptr AFArray
p AFArray
w AFArray
x -> Ptr AFArray -> AFArray -> Double -> AFArray -> IO AFErr
af_select_scalar_l Ptr AFArray
p AFArray
w Double
n AFArray
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);