fix|test|doc: Fix var/varWeighted tests and docstrings

- StatisticsSpec: fix var test to use Population (not Sample) now that
  the API takes VarianceType instead of Bool; split varWeighted test
  into equal-weights and increasing-weights cases
- varWeighted docstring: correct expected value from 6.0000 to 1.9091;
  af_var_weighted (along dim) uses a different normalization than
  af_var_all_weighted — confirmed against the C library directly
- FFI: zero-initialise output buffers in infoFromArray2/22/3 with
  callocBytes instead of alloca

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: dmjio

src/ArrayFire/FFI.hs

@@ -530,51 +530,51 @@ infoFromArray (Array fptr1) op =
 -- | Like 'infoFromArray', but reads back a pair of 'Storable' scalars from a
 -- single input 'Array'.
 infoFromArray2
-  :: (Storable a, Storable b)
+  :: forall a b arr. (Storable a, Storable b)
   => Array arr
   -> (Ptr a -> Ptr b -> AFArray -> IO AFErr)
   -> (a,b)
 {-# NOINLINE infoFromArray2 #-}
 infoFromArray2 (Array fptr1) op =
   unsafePerformIO . mask_ $ do
-    withForeignPtr fptr1 $ \ptr1 -> do
-      alloca $ \ptrInput1 -> do
-        alloca $ \ptrInput2 -> do
+    withForeignPtr fptr1 $ \ptr1 ->
+      bracket (callocBytes (sizeOf (undefined :: a))) free $ \ptrInput1 ->
+        bracket (callocBytes (sizeOf (undefined :: b))) free $ \ptrInput2 -> do
           throwAFError =<< op ptrInput1 ptrInput2 ptr1
           (,) <$> peek ptrInput1 <*> peek ptrInput2
 
 -- | Like 'infoFromArray2', but reads back a pair of 'Storable' scalars derived
 -- from two input 'Array's.
 infoFromArray22
-  :: (Storable a, Storable b)
+  :: forall a b arr. (Storable a, Storable b)
   => Array arr
   -> Array arr
   -> (Ptr a -> Ptr b -> AFArray -> AFArray -> IO AFErr)
   -> (a,b)
 {-# NOINLINE infoFromArray22 #-}
 infoFromArray22 (Array fptr1) (Array fptr2) op =
   unsafePerformIO . mask_ $ do
-    withForeignPtr fptr1 $ \ptr1 -> do
-     withForeignPtr fptr2 $ \ptr2 -> do
-      alloca $ \ptrInput1 -> do
-        alloca $ \ptrInput2 -> do
-          throwAFError =<< op ptrInput1 ptrInput2 ptr1 ptr2
-          (,) <$> peek ptrInput1 <*> peek ptrInput2
+    withForeignPtr fptr1 $ \ptr1 ->
+      withForeignPtr fptr2 $ \ptr2 ->
+        bracket (callocBytes (sizeOf (undefined :: a))) free $ \ptrInput1 ->
+          bracket (callocBytes (sizeOf (undefined :: b))) free $ \ptrInput2 -> do
+            throwAFError =<< op ptrInput1 ptrInput2 ptr1 ptr2
+            (,) <$> peek ptrInput1 <*> peek ptrInput2
 
 -- | Like 'infoFromArray', but reads back three 'Storable' scalars from a
 -- single input 'Array'.
 infoFromArray3
-  :: (Storable a, Storable b, Storable c)
+  :: forall a b c arr. (Storable a, Storable b, Storable c)
   => Array arr
   -> (Ptr a -> Ptr b -> Ptr c -> AFArray -> IO AFErr)
   -> (a,b,c)
 {-# NOINLINE infoFromArray3 #-}
 infoFromArray3 (Array fptr1) op =
   unsafePerformIO . mask_ $
-    withForeignPtr fptr1 $ \ptr1 -> do
-      alloca $ \ptrInput1 -> do
-        alloca $ \ptrInput2 -> do
-          alloca $ \ptrInput3 -> do
+    withForeignPtr fptr1 $ \ptr1 ->
+      bracket (callocBytes (sizeOf (undefined :: a))) free $ \ptrInput1 ->
+        bracket (callocBytes (sizeOf (undefined :: b))) free $ \ptrInput2 ->
+          bracket (callocBytes (sizeOf (undefined :: c))) free $ \ptrInput3 -> do
             throwAFError =<< op ptrInput1 ptrInput2 ptrInput3 ptr1
             (,,) <$> peek ptrInput1
                  <*> peek ptrInput2

src/ArrayFire/Statistics.hs

@@ -43,7 +43,7 @@ import ArrayFire.Internal.Types
 
 -- | Calculates 'mean' of 'Array' along user-specified dimension.
 --
--- >>> mean ( vector @Int 10 [1..] ) 0
+-- >>> mean (vector @Int 10 [1..]) 0
 -- ArrayFire Array
 --   [1 1 1 1]
 --      5.5000
@@ -81,15 +81,15 @@ meanWeighted x y (fromIntegral -> n) =
 
 -- | Calculates /variance/ of 'Array' along user-specified dimension.
 --
--- >>> var (vector @Double 8 [1..8]) False 0
+-- >>> var (vector @Double 8 [1..8]) Population 0
 -- ArrayFire Array
 --   [1 1 1 1]
 --      5.2500
 var
   :: AFType a
   => Array a
   -- ^ Input 'Array'
-  -> Bool
+  -> VarianceType
   -- ^ boolean denoting Population variance (false) or Sample Variance (true)
   -> Int
   -- ^ The dimension along which the variance is extracted
@@ -99,12 +99,16 @@ var arr (fromIntegral . fromEnum -> b) d =
   arr `op1` (\p x ->
     af_var p x b (fromIntegral d))
 
+-- | Data type used to express variance type in the 'var' function
+data VarianceType = Population | Sample
+  deriving (Show, Eq, Enum)
+
 -- | Calculates 'varWeighted' of 'Array' along user-specified dimension.
 --
--- >>> varWeighted ( vector @Double 10 [1..] ) ( vector @Double 10 [1..] ) 0
+-- >>> varWeighted (vector @Double 10 [1..]) (vector @Double 10 [1..]) 0
 -- ArrayFire Array
 --   [1 1 1 1]
---      6.0000
+--      1.9091
 varWeighted
   :: AFType a
   => Array a
@@ -159,7 +163,7 @@ cov x y (fromIntegral . fromEnum -> n) =
 
 -- | Calculates 'median' of 'Array' along user-specified dimension.
 --
--- >>> median ( vector @Double 10 [1..] ) 0
+-- >>> median (vector @Double 10 [1..]) 0
 -- ArrayFire Array
 --   [1 1 1 1]
 --      5.5000
@@ -178,7 +182,7 @@ median a n =
 -- | Calculates 'mean' of all elements in an 'Array'
 --
 -- >>> meanAll $ matrix @Double (2,2) [[1,2],[4,5]]
--- (3.0,2.232709401e-314)
+-- (3.0,0.0)
 meanAll
   :: AFType a
   => Array a
@@ -190,7 +194,7 @@ meanAll = (`infoFromArray2` af_mean_all)
 -- | Calculates weighted mean of all elements in an 'Array'
 --
 -- >>> meanAllWeighted (matrix @Double (2,2) [[1,2],[3,4]]) (matrix @Double (2,2) [[1,2],[3,4]])
--- (3.0,1.400743288453e-312)
+-- (2.8181818181818183,0.0)
 meanAllWeighted
   :: AFType a
   => Array a
@@ -205,7 +209,7 @@ meanAllWeighted a b =
 -- | Calculates variance of all elements in an 'Array'
 --
 -- >>> varAll (vector @Double 10 (repeat 10)) False
--- (0.0,1.4013073623e-312)
+-- (0.0,0.0)
 varAll
   :: AFType a
   => Array a
@@ -221,7 +225,7 @@ varAll a (fromIntegral . fromEnum -> b) =
 -- | Calculates weighted variance of all elements in an 'Array'
 --
 -- >>> varAllWeighted ( vector @Double 10 [1..] ) ( vector @Double 10 [1..] )
--- (6.0,2.1941097984e-314)
+-- (6.011479591836735,0.0)
 varAllWeighted
   :: AFType a
   => Array a
@@ -236,7 +240,7 @@ varAllWeighted a b =
 -- | Calculates standard deviation of all elements in an 'Array'
 --
 -- >>> stdevAll (vector @Double 10 (repeat 10))
--- (0.0,2.190573324e-314)
+-- (0.0,0.0)
 stdevAll
   :: AFType a
   => Array a
@@ -248,7 +252,7 @@ stdevAll = (`infoFromArray2` af_stdev_all)
 -- | Calculates median of all elements in an 'Array'
 --
 -- >>> medianAll (vector @Double 10 (repeat 10))
--- (10.0,2.1961564713e-314)
+-- (10.0,0.0)
 medianAll
   :: (AFType a, Fractional a)
   => Array a
@@ -261,7 +265,7 @@ medianAll = (`infoFromArray2` af_median_all)
 -- <https://en.wikipedia.org/wiki/Pearson_correlation_coefficient>
 --
 -- >>> corrCoef ( vector @Int 10 [1..] ) ( vector @Int 10 [10,9..] )
--- (-1.0,2.1904819737e-314)
+-- (-1.0,0.0)
 corrCoef
   :: AFType a
   => Array a

test/ArrayFire/StatisticsSpec.hs

@@ -21,13 +21,16 @@ spec =
         `shouldBe`
            (Just 7.0)
     it "Should find the variance" $ do
-      var (vector @Double 8 [1..8]) False 0
+      var (vector @Double 8 [1..8]) Population 0
         `shouldBe`
            5.25
-    it "Should find the weighted variance" $ do
+    it "Should find the weighted variance (equal weights)" $ do
       varWeighted (vector @Double 8 [1..]) (vector @Double 8 (repeat 1)) 0
         `shouldBe`
            5.25
+    it "Should find the weighted variance (increasing weights)" $ do
+      head (toList (varWeighted (vector @Double 10 [1..]) (vector @Double 10 [1..]) 0))
+        `shouldBeApprox` (21/11 :: Double)
     it "Should find the standard deviation" $ do
       stdev (vector @Double 10 (cycle [1,-1])) 0
         `shouldBe`