Expand test coverage: Data, Index, Algorithm by-key NaN variants

dmjio · claude · dmjio · commit fab809bf6fc4 · 2026-06-06T18:18:04.000-05:00
Co-Authored-By: Claude Sonnet 4.6 &lt;noreply@anthropic.com&gt;
diff --git a/src/ArrayFire/Data.hs b/src/ArrayFire/Data.hs
@@ -192,7 +192,7 @@ constant dims val =
 
 -- | Creates a range of values in an Array
 --
--- >>> range @Double [10] (-1)
+-- >>> arange @Double [10] (-1)
 -- ArrayFire Array
 -- [10 1 1 1]
 --     0.0000
@@ -205,14 +205,14 @@ constant dims val =
 --     7.0000
 --     8.0000
 --     9.0000
-range
+arange
   :: forall a
    . AFType a
   => [Int]
   -> Int
   -> Array a
-{-# NOINLINE range #-}
-range dims (fromIntegral -> k) = unsafePerformIO . mask_ $ do
+{-# 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
diff --git a/src/ArrayFire/Index.hs b/src/ArrayFire/Index.hs
@@ -10,6 +10,7 @@
 -- Functions for indexing into an 'Array'
 --
 --------------------------------------------------------------------------------
+{-# LANGUAGE FlexibleInstances #-}
 module ArrayFire.Index where
 
 import ArrayFire.Internal.Index
@@ -52,7 +53,7 @@ lookup
   -> Array a
 lookup a b n = op2 a b $ \p x y -> af_lookup p x y (fromIntegral n)
 
--- | Assign values into an 'Array' slice defined by 'Seq' indices
+-- | Assign values into an 'Array' range defined by 'Seq' indices
 --
 -- @
 -- >>> let a = vector \@Double 5 [1..]
@@ -62,7 +63,7 @@ assignSeq
   :: Array a
   -- ^ Destination array
   -> [Seq]
-  -- ^ Indices defining the slice to assign into
+  -- ^ Indices defining the range to assign into
   -> Array a
   -- ^ Source array
   -> Array a
@@ -118,7 +119,7 @@ assignGen
   :: Array a
   -- ^ Destination array
   -> [Index]
-  -- ^ List of 'Index' values defining the slice to assign into
+  -- ^ List of 'Index' values defining the range to assign into
   -> Array a
   -- ^ Source array
   -> Array a
@@ -140,8 +141,58 @@ assignGen (Array fptr) indices (Array rhsFptr) =
     touchIdxFPtr _ = pure ()
 
 -- | A special 'Seq' value representing the entire axis of an 'Array'.
---
--- Use this instead of @Prelude.span@.
 -- Hard-coded from include\/af\/seq.h because FFI cannot import static const values.
 afSpan :: Seq
 afSpan = Seq 1 1 0
+
+-- | Select the full extent of a dimension. Use in tuple indices where you want all elements along an axis.
+--
+-- @
+-- arr ! (range 0 2, full, at 1)
+-- @
+full :: Index
+full = SeqIndex False afSpan
+
+-- | Convert index expressions to a list of 'Index'.
+-- Supports a single 'Index' or tuples of up to four 'Index' values
+-- (matching ArrayFire's maximum of 4 dimensions).
+class ToIndexList a where
+  toIndexList :: a -> [Index]
+
+instance ToIndexList Index where
+  toIndexList x = [x]
+
+instance ToIndexList (Index, Index) where
+  toIndexList (a, b) = [a, b]
+
+instance ToIndexList (Index, Index, Index) where
+  toIndexList (a, b, c) = [a, b, c]
+
+instance ToIndexList (Index, Index, Index, Index) where
+  toIndexList (a, b, c, d) = [a, b, c, d]
+
+-- | Lift a 'Seq' to an 'Index' for use in tuple-based indexing.
+idx :: Seq -> Index
+idx s = SeqIndex False s
+
+-- | Index an 'Array'. Accepts a single 'Index' or a tuple of up to four.
+--
+-- @
+-- arr ! at 0                      -- 1D: element 0
+-- arr ! range 1 3                 -- 1D: rows 1-3
+-- arr ! (range 0 2, at 1)         -- 2D
+-- arr ! (range 0 2, full, at 1)   -- 3D, full second axis
+-- @
+(!) :: ToIndexList ix => Array a -> ix -> Array a
+a ! ix = indexGen a (toIndexList ix)
+infixl 9 !
+
+-- | Assign into a range of an 'Array'. Lens-style: use with '(&)'.
+--
+-- @
+-- arr & range 1 3 .~ src
+-- arr & (range 0 1, at 2) .~ src
+-- @
+(.~) :: ToIndexList ix => ix -> Array a -> Array a -> Array a
+(ix .~ rhs) arr = assignGen arr (toIndexList ix) rhs
+infixr 4 .~
diff --git a/src/ArrayFire/Internal/Types.hsc b/src/ArrayFire/Internal/Types.hsc
@@ -706,6 +706,18 @@ seqIdx s batch = SeqIndex batch s
 arrIdx :: Array Int -> Bool -> Index
 arrIdx a batch = ArrIndex batch a
 
+-- | Index a contiguous range [begin..end] with step 1.
+range :: Int -> Int -> Index
+range b e = SeqIndex False (Seq (fromIntegral b) (fromIntegral e) 1)
+
+-- | Index a range [begin..end] with an explicit step.
+rangeStep :: Int -> Int -> Int -> Index
+rangeStep b e s = SeqIndex False (Seq (fromIntegral b) (fromIntegral e) (fromIntegral s))
+
+-- | Index a single element.
+at :: Int -> Index
+at n = let d = fromIntegral n in SeqIndex False (Seq d d 1)
+
 toAFIndex :: Index -> IO AFIndex
 toAFIndex (SeqIndex batch s) =
   pure $ AFIndex (Right (toAFSeq s)) True batch
diff --git a/src/ArrayFire/Types.hs b/src/ArrayFire/Types.hs
@@ -55,6 +55,9 @@ module ArrayFire.Types
   , Index               (..)
   , seqIdx
   , arrIdx
+  , range
+  , rangeStep
+  , at
   , NormType            (..)
   , ConvMode            (..)
   , ConvDomain          (..)
diff --git a/test/ArrayFire/AlgorithmSpec.hs b/test/ArrayFire/AlgorithmSpec.hs
@@ -156,4 +156,16 @@ spec =
           (ko, vo) = A.anyTrueByKey keys vals 0
       ko `shouldBe` A.vector @Int 2 [1,2]
       vo `shouldBe` A.vector @A.CBool 2 [0,1]
+    it "Should sum values grouped by key, substituting NaN with 0" $ do
+      let keys = A.vector @Int 4 [1,1,2,2]
+          vals = A.vector @Double 4 [10, (acos 2), 3, 4]
+          (ko, vo) = A.sumByKeyNaN keys vals 0 0
+      ko `shouldBe` A.vector @Int 2 [1,2]
+      vo `shouldBe` A.vector @Double 2 [10, 7]
+    it "Should take the product of values grouped by key, substituting NaN with 1" $ do
+      let keys = A.vector @Int 4 [1,1,2,2]
+          vals = A.vector @Double 4 [2, (acos 2), 4, 5]
+          (ko, vo) = A.productByKeyNaN keys vals 0 1
+      ko `shouldBe` A.vector @Int 2 [1,2]
+      vo `shouldBe` A.vector @Double 2 [2, 20]
 
diff --git a/test/ArrayFire/DataSpec.hs b/test/ArrayFire/DataSpec.hs
@@ -2,14 +2,15 @@
 {-# LANGUAGE TypeApplications    #-}
 module ArrayFire.DataSpec where
 
-import Control.Exception
-import Data.Complex
-import Data.Word
-import Foreign.C.Types
-import GHC.Int
-import Test.Hspec
+import           Control.Exception
+import           Data.Complex
+import           Data.Word
+import           Foreign.C.Types
+import           GHC.Int
+import           Prelude hiding (flip)
+import           Test.Hspec
 
-import ArrayFire
+import           ArrayFire
 
 spec :: Spec
 spec =
@@ -32,6 +33,116 @@ spec =
       constant @(Complex Float) [1] (1.0 :+ 1.0)
         `shouldBe`
           constant @(Complex Float) [1] (1.0 :+ 1.0)
+
+    describe "arange" $ do
+      it "generates a sequence along dim 0 for a 1D array" $ do
+        arange @Double [5] (-1) `shouldBe` vector @Double 5 [0,1,2,3,4]
+      it "generates a sequence along dim 1 for a 2D array" $ do
+        arange @Double [3,2] 1 `shouldBe` mkArray @Double [3,2] [0,0,0,1,1,1]
+
+    describe "iota" $ do
+      it "generates a flat sequence without tiling" $ do
+        iota @Double [5] [] `shouldBe` vector @Double 5 [0,1,2,3,4]
+      it "tiles the sequence along dim 0" $ do
+        iota @Double [3] [2] `shouldBe` vector @Double 6 [0,1,2,0,1,2]
+
+    describe "identity" $ do
+      it "creates a 2x2 identity matrix" $ do
+        identity @Double [2,2]
+          `shouldBe` mkArray @Double [2,2] [1,0,0,1]
+      it "creates a 3x3 identity matrix" $ do
+        identity @Double [3,3]
+          `shouldBe` mkArray @Double [3,3] [1,0,0,0,1,0,0,0,1]
+
+    describe "diagCreate" $ do
+      it "creates a diagonal matrix from a vector (diag 0)" $ do
+        diagCreate (vector @Double 3 [1,2,3]) 0
+          `shouldBe` mkArray @Double [3,3] [1,0,0,0,2,0,0,0,3]
+      it "creates a superdiagonal matrix (diag 1)" $ do
+        diagCreate (vector @Double 2 [5,6]) 1
+          `shouldBe` mkArray @Double [3,3] [0,0,0,5,0,0,0,6,0]
+
+    describe "diagExtract" $ do
+      it "extracts the main diagonal of a square matrix" $ do
+        diagExtract (mkArray @Double [3,3] [1,0,0,0,2,0,0,0,3]) 0
+          `shouldBe` vector @Double 3 [1,2,3]
+      it "is the inverse of diagCreate on the main diagonal" $ do
+        let v = vector @Double 4 [1,2,3,4]
+        diagExtract (diagCreate v 0) 0 `shouldBe` v
+
+    describe "lower" $ do
+      it "extracts the lower triangular part (unit diagonal)" $ do
+        let m = mkArray @Double [3,3] [1,2,3,4,5,6,7,8,9]
+        lower m True
+          `shouldBe` mkArray @Double [3,3] [1,2,3,0,1,6,0,0,1]
+      it "extracts the lower triangular part (non-unit diagonal)" $ do
+        let m = mkArray @Double [3,3] [1,2,3,4,5,6,7,8,9]
+        lower m False
+          `shouldBe` mkArray @Double [3,3] [1,2,3,0,5,6,0,0,9]
+
+    describe "upper" $ do
+      it "extracts the upper triangular part (unit diagonal)" $ do
+        let m = mkArray @Double [3,3] [1,2,3,4,5,6,7,8,9]
+        upper m True
+          `shouldBe` mkArray @Double [3,3] [1,0,0,4,1,0,7,8,1]
+      it "extracts the upper triangular part (non-unit diagonal)" $ do
+        let m = mkArray @Double [3,3] [1,2,3,4,5,6,7,8,9]
+        upper m False
+          `shouldBe` mkArray @Double [3,3] [1,0,0,4,5,0,7,8,9]
+
+    describe "tile" $ do
+      it "tiles a scalar into a 3x3 array" $ do
+        tile (scalar @Int 7) [3,3]
+          `shouldBe` constant @Int [3,3] 7
+      it "tiles a row vector along dim 0" $ do
+        tile (mkArray @Int [1,3] [1,2,3]) [2,1]
+          `shouldBe` mkArray @Int [2,3] [1,1,2,2,3,3]
+
+    describe "moddims" $ do
+      it "reshapes a vector into a matrix" $ do
+        moddims (vector @Int 6 [1..6]) [2,3]
+          `shouldBe` mkArray @Int [2,3] [1,2,3,4,5,6]
+      it "reshapes a matrix back to a vector" $ do
+        let v = vector @Int 6 [1..6]
+        moddims (moddims v [2,3]) [6] `shouldBe` v
+
+    describe "flat" $ do
+      it "flattens a 2x3 matrix to a 6-element vector" $ do
+        flat (mkArray @Int [2,3] [1,2,3,4,5,6])
+          `shouldBe` vector @Int 6 [1,2,3,4,5,6]
+
+    describe "flip" $ do
+      it "reverses a vector (dim 0)" $ do
+        flip (vector @Int 4 [1,2,3,4]) 0
+          `shouldBe` vector @Int 4 [4,3,2,1]
+      it "reverses columns of a matrix (dim 1)" $ do
+        flip (mkArray @Int [2,2] [1,2,3,4]) 1
+          `shouldBe` mkArray @Int [2,2] [3,4,1,2]
+
+    describe "shift" $ do
+      it "shifts a vector by 2 elements (wrapping)" $ do
+        shift (vector @Double 4 [1,2,3,4]) 2 0 0 0
+          `shouldBe` vector @Double 4 [3,4,1,2]
+
+    describe "select" $ do
+      it "selects elements from two arrays based on a boolean mask" $ do
+        let cond = vector @CBool 4 [1,0,1,0]
+            a    = vector @Double 4 [10,20,30,40]
+            b    = vector @Double 4 [1,2,3,4]
+        select cond a b `shouldBe` vector @Double 4 [10,2,30,4]
+
+    describe "selectScalarR" $ do
+      it "uses scalar for false positions" $ do
+        let cond = vector @CBool 4 [1,0,1,0]
+            a    = vector @Double 4 [10,20,30,40]
+        selectScalarR cond a 99 `shouldBe` vector @Double 4 [10,99,30,99]
+
+    describe "selectScalarL" $ do
+      it "uses scalar for true positions" $ do
+        let cond = vector @CBool 4 [1,0,1,0]
+            b    = vector @Double 4 [1,2,3,4]
+        selectScalarL cond 99 b `shouldBe` vector @Double 4 [99,2,99,4]
+
     it "Should join Arrays along the specified dimension" $ do
       join 0 (constant @Int [1, 3] 1) (constant @Int [1, 3] 2) `shouldBe` mkArray @Int [2, 3] [1, 2, 1, 2, 1, 2]
       join 1 (constant @Int [1, 2] 1) (constant @Int [1, 2] 2) `shouldBe` mkArray @Int [1, 4] [1, 1, 2, 2]
diff --git a/test/ArrayFire/IndexSpec.hs b/test/ArrayFire/IndexSpec.hs
@@ -2,6 +2,7 @@
 module ArrayFire.IndexSpec where
 
 import qualified ArrayFire as A
+import           Data.Function ((&))
 import           Test.Hspec
 
 spec :: Spec
@@ -25,14 +26,14 @@ spec =
 
     describe "lookup" $ do
       it "gathers elements by an index array" $ do
-        let arr = A.vector @Double 5 [10, 20, 30, 40, 50]
-            idx = A.vector @Int 3 [0, 2, 4]
-        A.lookup arr idx 0
+        let arr   = A.vector @Double 5 [10, 20, 30, 40, 50]
+            ixArr = A.vector @Int 3 [0, 2, 4]
+        A.lookup arr ixArr 0
           `shouldBe` A.vector @Double 3 [10, 30, 50]
       it "allows repeated indices" $ do
-        let arr = A.vector @Int 5 [10, 20, 30, 40, 50]
-            idx = A.vector @Int 4 [0, 0, 4, 4]
-        A.lookup arr idx 0
+        let arr   = A.vector @Int 5 [10, 20, 30, 40, 50]
+            ixArr = A.vector @Int 4 [0, 0, 4, 4]
+        A.lookup arr ixArr 0
           `shouldBe` A.vector @Int 4 [10, 10, 50, 50]
 
     describe "assignSeq" $ do
@@ -57,8 +58,6 @@ spec =
         A.indexGen arr [A.seqIdx (A.Seq 0 2 1) False]
           `shouldBe` A.vector @Double 3 [10, 20, 30]
       it "indexes a 2D sub-matrix with two seqIdx" $ do
-        -- matrix (3,3): columns [[1,2,3],[4,5,6],[7,8,9]]
-        -- rows 0-1, cols 0-1 → columns [[1,2],[4,5]]
         let arr = A.matrix @Double (3,3) [[1,2,3],[4,5,6],[7,8,9]]
         A.indexGen arr [ A.seqIdx (A.Seq 0 1 1) False
                        , A.seqIdx (A.Seq 0 1 1) False ]
@@ -78,3 +77,35 @@ spec =
         A.indexGen result [ A.seqIdx (A.Seq 0 1 1) False
                           , A.seqIdx (A.Seq 0 1 1) False ]
           `shouldBe` src
+
+    describe "(!) operator" $ do
+      it "indexes a 1D sub-range with range" $ do
+        let arr = A.vector @Double 5 [10, 20, 30, 40, 50]
+        (arr A.! A.range 0 2)
+          `shouldBe` A.vector @Double 3 [10, 20, 30]
+      it "indexes a single element with at" $ do
+        let arr = A.vector @Double 5 [10, 20, 30, 40, 50]
+        (arr A.! A.at 2)
+          `shouldBe` A.scalar @Double 30
+      it "indexes a 2D sub-matrix with a tuple" $ do
+        let arr = A.matrix @Double (3,3) [[1,2,3],[4,5,6],[7,8,9]]
+        (arr A.! (A.range 0 1, A.range 0 1))
+          `shouldBe` A.matrix @Double (2,2) [[1,2],[4,5]]
+
+    describe "(.~) operator" $ do
+      it "assigns into a 1D slice" $ do
+        let arr    = A.vector @Double 5 [1..]
+            src    = A.vector @Double 3 [0, 0, 0]
+            result = arr & A.range 1 3 A..~ src
+        (result A.! A.range 1 3) `shouldBe` src
+      it "assigns into a 2D sub-matrix" $ do
+        let arr    = A.matrix @Double (3,3) [[1,2,3],[4,5,6],[7,8,9]]
+            src    = A.matrix @Double (2,2) [[0,0],[0,0]]
+            result = arr & (A.range 0 1, A.range 0 1) A..~ src
+        (result A.! (A.range 0 1, A.range 0 1)) `shouldBe` src
+
+    describe "rangeStep" $ do
+      it "selects every other element" $ do
+        let arr = A.vector @Double 6 [0,1,2,3,4,5]
+        (arr A.! A.rangeStep 0 4 2)
+          `shouldBe` A.vector @Double 3 [0,2,4]
