{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}

-- |
-- Module      : Data.Massiv.Array.Numeric
-- Copyright   : (c) Alexey Kuleshevich 2018-2022
-- License     : BSD3
-- Maintainer  : Alexey Kuleshevich <lehins@yandex.ru>
-- Stability   : experimental
-- Portability : non-portable
module Data.Massiv.Array.Numeric (
  -- * Numeric
  Numeric,
  NumericFloat,
  liftNumArray2M,

  -- ** Pointwise addition
  (.+),
  (+.),
  (.+.),
  (!+!),
  sumArraysM,
  sumArrays',

  -- ** Pointwise subtraction
  (.-),
  (-.),
  (.-.),
  (!-!),

  -- ** Pointwise multiplication
  (.*),
  (*.),
  (.*.),
  (!*!),
  (.^),
  productArraysM,
  productArrays',

  -- ** Dot product
  (!.!),
  dotM,

  -- ** Matrix multiplication
  (.><),
  (!><),
  multiplyMatrixByVector,
  (><.),
  (><!),
  multiplyVectorByMatrix,
  (.><.),
  (!><!),
  multiplyMatrices,
  multiplyMatricesTransposed,

  -- * Norms
  normL2,

  -- * Simple matrices
  identityMatrix,
  lowerTriangular,
  upperTriangular,
  negateA,
  absA,
  signumA,

  -- * Integral
  quotA,
  remA,
  divA,
  modA,
  quotRemA,
  divModA,

  -- * Fractional
  (./),
  (/.),
  (./.),
  (!/!),
  (.^^),
  recipA,

  -- * Floating
  expA,
  logA,
  sqrtA,
  (.**),
  logBaseA,
  sinA,
  cosA,
  tanA,
  asinA,
  acosA,
  atanA,
  sinhA,
  coshA,
  tanhA,
  asinhA,
  acoshA,
  atanhA,

  -- * RealFrac
  truncateA,
  roundA,
  ceilingA,
  floorA,

  -- * RealFloat
  atan2A,
) where

import Control.Monad (when)
import Control.Scheduler
import qualified Data.Foldable as F
import Data.Function
import Data.Massiv.Array.Delayed.Pull
import Data.Massiv.Array.Delayed.Push
import Data.Massiv.Array.Manifest.Internal
import Data.Massiv.Array.Ops.Construct
import Data.Massiv.Array.Ops.Map as A
import Data.Massiv.Core
import Data.Massiv.Core.Common as A
import Data.Massiv.Core.Operations
import System.IO.Unsafe
import Prelude as P

infixr 8 .^, .^^
infixl 7 !*!, .*., .*, *., !/!, ./., ./, /., `quotA`, `remA`, `divA`, `modA`
infixl 6 !+!, .+., .+, +., !-!, .-., .-, -.

-- | Similar to `liftArray2M`, except it can be applied only to representations
-- with `Numeric` instance and result representation stays the same.
--
-- @since 1.0.0
liftNumArray2M
  :: (Index ix, Numeric r e, MonadThrow m)
  => (e -> e -> e)
  -> Array r ix e
  -> Array r ix e
  -> m (Array r ix e)
liftNumArray2M :: forall ix r e (m :: * -> *).
(Index ix, Numeric r e, MonadThrow m) =>
(e -> e -> e) -> Array r ix e -> Array r ix e -> m (Array r ix e)
liftNumArray2M e -> e -> e
f Array r ix e
a1 Array r ix e
a2
  | Array r ix e -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array r ix e -> Sz ix
size Array r ix e
a1 Sz ix -> Sz ix -> Bool
forall a. Eq a => a -> a -> Bool
== Array r ix e -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array r ix e -> Sz ix
size Array r ix e
a2 = Array r ix e -> m (Array r ix e)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Array r ix e -> m (Array r ix e))
-> Array r ix e -> m (Array r ix e)
forall a b. (a -> b) -> a -> b
$ (e -> e -> e) -> Array r ix e -> Array r ix e -> Array r ix e
forall ix.
Index ix =>
(e -> e -> e) -> Array r ix e -> Array r ix e -> Array r ix e
forall r e ix.
(Numeric r e, Index ix) =>
(e -> e -> e) -> Array r ix e -> Array r ix e -> Array r ix e
unsafeLiftArray2 e -> e -> e
f Array r ix e
a1 Array r ix e
a2
  | Sz ix -> Bool
forall ix. Index ix => Sz ix -> Bool
isZeroSz Sz ix
sz1 Bool -> Bool -> Bool
&& Sz ix -> Bool
forall ix. Index ix => Sz ix -> Bool
isZeroSz Sz ix
sz2 = Array r ix e -> m (Array r ix e)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Array r ix e -> m (Array r ix e))
-> Array r ix e -> m (Array r ix e)
forall a b. (a -> b) -> a -> b
$ Sz ix -> Array r ix e -> Array r ix e
forall r ix ix' e.
(Size r, Index ix, Index ix') =>
Sz ix' -> Array r ix e -> Array r ix' e
forall ix ix' e.
(Index ix, Index ix') =>
Sz ix' -> Array r ix e -> Array r ix' e
unsafeResize Sz ix
forall ix. Index ix => Sz ix
zeroSz Array r ix e
a1
  | Bool
otherwise = SizeException -> m (Array r ix e)
forall e a. (HasCallStack, Exception e) => e -> m a
forall (m :: * -> *) e a.
(MonadThrow m, HasCallStack, Exception e) =>
e -> m a
throwM (SizeException -> m (Array r ix e))
-> SizeException -> m (Array r ix e)
forall a b. (a -> b) -> a -> b
$ Sz ix -> Sz ix -> SizeException
forall ix. Index ix => Sz ix -> Sz ix -> SizeException
SizeMismatchException Sz ix
sz1 Sz ix
sz2
  where
    !sz1 :: Sz ix
sz1 = Array r ix e -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array r ix e -> Sz ix
size Array r ix e
a1
    !sz2 :: Sz ix
sz2 = Array r ix e -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array r ix e -> Sz ix
size Array r ix e
a2
{-# INLINE liftNumArray2M #-}

applyExactSize2M
  :: (Index ix, Size r, MonadThrow m)
  => (Array r ix e -> Array r ix e -> Array r ix e)
  -> Array r ix e
  -> Array r ix e
  -> m (Array r ix e)
applyExactSize2M :: forall ix r (m :: * -> *) e.
(Index ix, Size r, MonadThrow m) =>
(Array r ix e -> Array r ix e -> Array r ix e)
-> Array r ix e -> Array r ix e -> m (Array r ix e)
applyExactSize2M Array r ix e -> Array r ix e -> Array r ix e
f Array r ix e
a1 Array r ix e
a2
  | Array r ix e -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array r ix e -> Sz ix
size Array r ix e
a1 Sz ix -> Sz ix -> Bool
forall a. Eq a => a -> a -> Bool
== Array r ix e -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array r ix e -> Sz ix
size Array r ix e
a2 = Array r ix e -> m (Array r ix e)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Array r ix e -> m (Array r ix e))
-> Array r ix e -> m (Array r ix e)
forall a b. (a -> b) -> a -> b
$! Array r ix e -> Array r ix e -> Array r ix e
f Array r ix e
a1 Array r ix e
a2
  | Sz ix -> Bool
forall ix. Index ix => Sz ix -> Bool
isZeroSz Sz ix
sz1 Bool -> Bool -> Bool
&& Sz ix -> Bool
forall ix. Index ix => Sz ix -> Bool
isZeroSz Sz ix
sz2 = Array r ix e -> m (Array r ix e)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Array r ix e -> m (Array r ix e))
-> Array r ix e -> m (Array r ix e)
forall a b. (a -> b) -> a -> b
$! Sz ix -> Array r ix e -> Array r ix e
forall r ix ix' e.
(Size r, Index ix, Index ix') =>
Sz ix' -> Array r ix e -> Array r ix' e
forall ix ix' e.
(Index ix, Index ix') =>
Sz ix' -> Array r ix e -> Array r ix' e
unsafeResize Sz ix
forall ix. Index ix => Sz ix
zeroSz Array r ix e
a1
  | Bool
otherwise = SizeException -> m (Array r ix e)
forall e a. (HasCallStack, Exception e) => e -> m a
forall (m :: * -> *) e a.
(MonadThrow m, HasCallStack, Exception e) =>
e -> m a
throwM (SizeException -> m (Array r ix e))
-> SizeException -> m (Array r ix e)
forall a b. (a -> b) -> a -> b
$! Sz ix -> Sz ix -> SizeException
forall ix. Index ix => Sz ix -> Sz ix -> SizeException
SizeMismatchException Sz ix
sz1 Sz ix
sz2
  where
    !sz1 :: Sz ix
sz1 = Array r ix e -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array r ix e -> Sz ix
size Array r ix e
a1
    !sz2 :: Sz ix
sz2 = Array r ix e -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array r ix e -> Sz ix
size Array r ix e
a2
{-# INLINE applyExactSize2M #-}

-- | Add two arrays together pointwise. Same as `!+!` but produces monadic computation
-- that allows for handling failure.
--
-- /__Throws Exception__/: `SizeMismatchException` when array sizes do not match.
--
-- @since 0.4.0
(.+.) :: (Index ix, Numeric r e, MonadThrow m) => Array r ix e -> Array r ix e -> m (Array r ix e)
.+. :: forall ix r e (m :: * -> *).
(Index ix, Numeric r e, MonadThrow m) =>
Array r ix e -> Array r ix e -> m (Array r ix e)
(.+.) = (Array r ix e -> Array r ix e -> Array r ix e)
-> Array r ix e -> Array r ix e -> m (Array r ix e)
forall ix r (m :: * -> *) e.
(Index ix, Size r, MonadThrow m) =>
(Array r ix e -> Array r ix e -> Array r ix e)
-> Array r ix e -> Array r ix e -> m (Array r ix e)
applyExactSize2M Array r ix e -> Array r ix e -> Array r ix e
forall ix. Index ix => Array r ix e -> Array r ix e -> Array r ix e
forall r e ix.
(Numeric r e, Index ix) =>
Array r ix e -> Array r ix e -> Array r ix e
additionPointwise
{-# INLINE (.+.) #-}

-- | Add two arrays together pointwise. Prefer to use monadic version of this function
-- `.+.` whenever possible, because it is better to avoid partial functions.
--
-- [Partial] Mismatched array sizes will result in an impure exception being thrown.
--
-- ====__Example__
--
-- >>> let a1 = Ix1 0 ... 10
-- >>> let a2 = Ix1 20 ... 30
-- >>> a1 !+! a2
-- Array D Seq (Sz1 11)
--   [ 20, 22, 24, 26, 28, 30, 32, 34, 36, 38, 40 ]
--
-- @since 0.5.6
(!+!) :: (HasCallStack, Index ix, Numeric r e) => Array r ix e -> Array r ix e -> Array r ix e
!+! :: forall ix r e.
(HasCallStack, Index ix, Numeric r e) =>
Array r ix e -> Array r ix e -> Array r ix e
(!+!) Array r ix e
a1 Array r ix e
a2 = Either SomeException (Array r ix e) -> Array r ix e
forall a. HasCallStack => Either SomeException a -> a
throwEither (Array r ix e
a1 Array r ix e -> Array r ix e -> Either SomeException (Array r ix e)
forall ix r e (m :: * -> *).
(Index ix, Numeric r e, MonadThrow m) =>
Array r ix e -> Array r ix e -> m (Array r ix e)
.+. Array r ix e
a2)
{-# INLINE (!+!) #-}

-- | Add a scalar to each element of the array. Array is on the left.
--
-- @since 0.1.0
(.+) :: (Index ix, Numeric r e) => Array r ix e -> e -> Array r ix e
.+ :: forall ix r e.
(Index ix, Numeric r e) =>
Array r ix e -> e -> Array r ix e
(.+) = Array r ix e -> e -> Array r ix e
forall ix. Index ix => Array r ix e -> e -> Array r ix e
forall r e ix.
(Numeric r e, Index ix) =>
Array r ix e -> e -> Array r ix e
plusScalar
{-# INLINE (.+) #-}

-- | Add a scalar to each element of the array. Array is on the right.
--
-- @since 0.4.0
(+.) :: (Index ix, Numeric r e) => e -> Array r ix e -> Array r ix e
+. :: forall ix r e.
(Index ix, Numeric r e) =>
e -> Array r ix e -> Array r ix e
(+.) = (Array r ix e -> e -> Array r ix e)
-> e -> Array r ix e -> Array r ix e
forall a b c. (a -> b -> c) -> b -> a -> c
flip Array r ix e -> e -> Array r ix e
forall ix. Index ix => Array r ix e -> e -> Array r ix e
forall r e ix.
(Numeric r e, Index ix) =>
Array r ix e -> e -> Array r ix e
plusScalar
{-# INLINE (+.) #-}

-- | Subtract two arrays pointwise. Same as `!-!` but produces monadic computation that
-- allows for handling failure.
--
-- /__Throws Exception__/: `SizeMismatchException` when array sizes do not match.
--
-- @since 0.4.0
(.-.)
  :: (Index ix, Numeric r e, MonadThrow m) => Array r ix e -> Array r ix e -> m (Array r ix e)
.-. :: forall ix r e (m :: * -> *).
(Index ix, Numeric r e, MonadThrow m) =>
Array r ix e -> Array r ix e -> m (Array r ix e)
(.-.) = (Array r ix e -> Array r ix e -> Array r ix e)
-> Array r ix e -> Array r ix e -> m (Array r ix e)
forall ix r (m :: * -> *) e.
(Index ix, Size r, MonadThrow m) =>
(Array r ix e -> Array r ix e -> Array r ix e)
-> Array r ix e -> Array r ix e -> m (Array r ix e)
applyExactSize2M Array r ix e -> Array r ix e -> Array r ix e
forall ix. Index ix => Array r ix e -> Array r ix e -> Array r ix e
forall r e ix.
(Numeric r e, Index ix) =>
Array r ix e -> Array r ix e -> Array r ix e
subtractionPointwise
{-# INLINE (.-.) #-}

-- | Subtract one array from another pointwise. Prefer to use monadic version of this
-- function `.-.` whenever possible, because it is better to avoid partial functions.
--
-- [Partial] Mismatched array sizes will result in an impure exception being thrown.
--
-- ====__Example__
--
-- >>> let a1 = Ix1 0 ... 10
-- >>> let a2 = Ix1 20 ... 30
-- >>> a1 !-! a2
-- Array D Seq (Sz1 11)
--   [ -20, -20, -20, -20, -20, -20, -20, -20, -20, -20, -20 ]
--
-- @since 0.5.6
(!-!) :: (Index ix, Numeric r e) => Array r ix e -> Array r ix e -> Array r ix e
!-! :: forall ix r e.
(Index ix, Numeric r e) =>
Array r ix e -> Array r ix e -> Array r ix e
(!-!) Array r ix e
a1 Array r ix e
a2 = Either SomeException (Array r ix e) -> Array r ix e
forall a. HasCallStack => Either SomeException a -> a
throwEither (Array r ix e
a1 Array r ix e -> Array r ix e -> Either SomeException (Array r ix e)
forall ix r e (m :: * -> *).
(Index ix, Numeric r e, MonadThrow m) =>
Array r ix e -> Array r ix e -> m (Array r ix e)
.-. Array r ix e
a2)
{-# INLINE (!-!) #-}

-- | Subtract a scalar from each element of the array. Array is on the left.
--
-- @since 0.4.0
(.-) :: (Index ix, Numeric r e) => Array r ix e -> e -> Array r ix e
.- :: forall ix r e.
(Index ix, Numeric r e) =>
Array r ix e -> e -> Array r ix e
(.-) = Array r ix e -> e -> Array r ix e
forall ix. Index ix => Array r ix e -> e -> Array r ix e
forall r e ix.
(Numeric r e, Index ix) =>
Array r ix e -> e -> Array r ix e
minusScalar
{-# INLINE (.-) #-}

-- | Subtract each element of the array from a scalar. Array is on the right.
--
-- @since 0.5.6
(-.) :: (Index ix, Numeric r e) => e -> Array r ix e -> Array r ix e
-. :: forall ix r e.
(Index ix, Numeric r e) =>
e -> Array r ix e -> Array r ix e
(-.) = e -> Array r ix e -> Array r ix e
forall ix. Index ix => e -> Array r ix e -> Array r ix e
forall r e ix.
(Numeric r e, Index ix) =>
e -> Array r ix e -> Array r ix e
scalarMinus
{-# INLINE (-.) #-}

-- | Multiply two arrays together pointwise. Same as `!*!` but produces monadic
-- computation that allows for handling failure.
--
-- /__Throws Exception__/: `SizeMismatchException` when array sizes do not match.
--
-- @since 0.4.0
(.*.)
  :: (Index ix, Numeric r e, MonadThrow m) => Array r ix e -> Array r ix e -> m (Array r ix e)
.*. :: forall ix r e (m :: * -> *).
(Index ix, Numeric r e, MonadThrow m) =>
Array r ix e -> Array r ix e -> m (Array r ix e)
(.*.) = (Array r ix e -> Array r ix e -> Array r ix e)
-> Array r ix e -> Array r ix e -> m (Array r ix e)
forall ix r (m :: * -> *) e.
(Index ix, Size r, MonadThrow m) =>
(Array r ix e -> Array r ix e -> Array r ix e)
-> Array r ix e -> Array r ix e -> m (Array r ix e)
applyExactSize2M Array r ix e -> Array r ix e -> Array r ix e
forall ix. Index ix => Array r ix e -> Array r ix e -> Array r ix e
forall r e ix.
(Numeric r e, Index ix) =>
Array r ix e -> Array r ix e -> Array r ix e
multiplicationPointwise
{-# INLINE (.*.) #-}

-- | Multiplication of two arrays pointwise,
-- i.e. <https://en.wikipedia.org/wiki/Hadamard_product_(matrices) Hadamard product>.
-- Prefer to use monadic version of this function `.*.` whenever possible,
-- because it is better to avoid partial functions.
--
-- [Partial] Mismatched array sizes will result in an impure exception being thrown.
--
-- ====__Example__
--
-- >>> let a1 = Ix1 0 ... 10
-- >>> let a2 = Ix1 20 ... 30
-- >>> a1 !*! a2
-- Array D Seq (Sz1 11)
--   [ 0, 21, 44, 69, 96, 125, 156, 189, 224, 261, 300 ]
--
-- @since 0.5.6
(!*!) :: (Index ix, Numeric r e) => Array r ix e -> Array r ix e -> Array r ix e
!*! :: forall ix r e.
(Index ix, Numeric r e) =>
Array r ix e -> Array r ix e -> Array r ix e
(!*!) Array r ix e
a1 Array r ix e
a2 = Either SomeException (Array r ix e) -> Array r ix e
forall a. HasCallStack => Either SomeException a -> a
throwEither (Array r ix e
a1 Array r ix e -> Array r ix e -> Either SomeException (Array r ix e)
forall ix r e (m :: * -> *).
(Index ix, Numeric r e, MonadThrow m) =>
Array r ix e -> Array r ix e -> m (Array r ix e)
.*. Array r ix e
a2)
{-# INLINE (!*!) #-}

-- | Multiply each element of the array by a scalar value. Scalar is on the right.
--
-- ====__Example__
--
-- >>> let arr = Ix1 20 ..: 25
-- >>> arr
-- Array D Seq (Sz1 5)
--   [ 20, 21, 22, 23, 24 ]
-- >>> arr .* 10
-- Array D Seq (Sz1 5)
--   [ 200, 210, 220, 230, 240 ]
--
-- @since 0.4.0
(.*) :: (Index ix, Numeric r e) => Array r ix e -> e -> Array r ix e
.* :: forall ix r e.
(Index ix, Numeric r e) =>
Array r ix e -> e -> Array r ix e
(.*) = Array r ix e -> e -> Array r ix e
forall ix. Index ix => Array r ix e -> e -> Array r ix e
forall r e ix.
(Numeric r e, Index ix) =>
Array r ix e -> e -> Array r ix e
multiplyScalar
{-# INLINE (.*) #-}

-- | Multiply each element of the array by a scalar value. Scalar is on the left.
--
-- ====__Example__
--
-- >>> let arr = Ix1 20 ..: 25
-- >>> arr
-- Array D Seq (Sz1 5)
--   [ 20, 21, 22, 23, 24 ]
-- >>> 10 *. arr
-- Array D Seq (Sz1 5)
--   [ 200, 210, 220, 230, 240 ]
--
-- @since 0.4.0
(*.) :: (Index ix, Numeric r e) => e -> Array r ix e -> Array r ix e
*. :: forall ix r e.
(Index ix, Numeric r e) =>
e -> Array r ix e -> Array r ix e
(*.) = (Array r ix e -> e -> Array r ix e)
-> e -> Array r ix e -> Array r ix e
forall a b c. (a -> b -> c) -> b -> a -> c
flip Array r ix e -> e -> Array r ix e
forall ix. Index ix => Array r ix e -> e -> Array r ix e
forall r e ix.
(Numeric r e, Index ix) =>
Array r ix e -> e -> Array r ix e
multiplyScalar
{-# INLINE (*.) #-}

-- | Raise each element of the array to a power.
--
-- ====__Example__
--
-- >>> let arr = Ix1 20 ..: 25
-- >>> arr
-- Array D Seq (Sz1 5)
--   [ 20, 21, 22, 23, 24 ]
-- >>> arr .^ 3
-- Array D Seq (Sz1 5)
--   [ 8000, 9261, 10648, 12167, 13824 ]
--
-- @since 0.4.0
(.^) :: (Index ix, Numeric r e) => Array r ix e -> Int -> Array r ix e
.^ :: forall ix r e.
(Index ix, Numeric r e) =>
Array r ix e -> Int -> Array r ix e
(.^) = Array r ix e -> Int -> Array r ix e
forall ix. Index ix => Array r ix e -> Int -> Array r ix e
forall r e ix.
(Numeric r e, Index ix) =>
Array r ix e -> Int -> Array r ix e
powerPointwise
{-# INLINE (.^) #-}

-- | Dot product of two vectors.
--
-- [Partial] Throws an impure exception when lengths of vectors do not match
--
-- @since 0.5.6
(!.!) :: (Numeric r e, Source r e) => Vector r e -> Vector r e -> e
!.! :: forall r e.
(Numeric r e, Source r e) =>
Vector r e -> Vector r e -> e
(!.!) Vector r e
v1 Vector r e
v2 = Either SomeException e -> e
forall a. HasCallStack => Either SomeException a -> a
throwEither (Either SomeException e -> e) -> Either SomeException e -> e
forall a b. (a -> b) -> a -> b
$ Vector r e -> Vector r e -> Either SomeException e
forall r e (m :: * -> *).
(FoldNumeric r e, Source r e, MonadThrow m) =>
Vector r e -> Vector r e -> m e
dotM Vector r e
v1 Vector r e
v2
{-# INLINE (!.!) #-}

-- | Dot product of two vectors.
--
-- /__Throws Exception__/: `SizeMismatchException` when lengths of vectors do not match
--
-- @since 0.5.6
dotM :: (FoldNumeric r e, Source r e, MonadThrow m) => Vector r e -> Vector r e -> m e
dotM :: forall r e (m :: * -> *).
(FoldNumeric r e, Source r e, MonadThrow m) =>
Vector r e -> Vector r e -> m e
dotM Vector r e
v1 Vector r e
v2
  | Vector r e -> Sz Int
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array r ix e -> Sz ix
size Vector r e
v1 Sz Int -> Sz Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Vector r e -> Sz Int
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array r ix e -> Sz ix
size Vector r e
v2 = SizeException -> m e
forall e a. (HasCallStack, Exception e) => e -> m a
forall (m :: * -> *) e a.
(MonadThrow m, HasCallStack, Exception e) =>
e -> m a
throwM (SizeException -> m e) -> SizeException -> m e
forall a b. (a -> b) -> a -> b
$ Sz Int -> Sz Int -> SizeException
forall ix. Index ix => Sz ix -> Sz ix -> SizeException
SizeMismatchException (Vector r e -> Sz Int
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array r ix e -> Sz ix
size Vector r e
v1) (Vector r e -> Sz Int
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array r ix e -> Sz ix
size Vector r e
v2)
  | Comp
comp Comp -> Comp -> Bool
forall a. Eq a => a -> a -> Bool
== Comp
Seq = e -> m e
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (e -> m e) -> e -> m e
forall a b. (a -> b) -> a -> b
$! Vector r e -> Vector r e -> e
forall ix. Index ix => Array r ix e -> Array r ix e -> e
forall r e ix.
(FoldNumeric r e, Index ix) =>
Array r ix e -> Array r ix e -> e
unsafeDotProduct Vector r e
v1 Vector r e
v2
  | Bool
otherwise = e -> m e
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (e -> m e) -> e -> m e
forall a b. (a -> b) -> a -> b
$! IO e -> e
forall a. IO a -> a
unsafePerformIO (IO e -> e) -> IO e -> e
forall a b. (a -> b) -> a -> b
$ Vector r e -> Vector r e -> IO e
forall (m :: * -> *) ix r b.
(MonadUnliftIO m, Index ix, FoldNumeric r b, Source r b) =>
Array r ix b -> Array r ix b -> m b
unsafeDotProductIO Vector r e
v1 Vector r e
v2
  where
    comp :: Comp
comp = Vector r e -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
forall ix e. Array r ix e -> Comp
getComp Vector r e
v1 Comp -> Comp -> Comp
forall a. Semigroup a => a -> a -> a
<> Vector r e -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
forall ix e. Array r ix e -> Comp
getComp Vector r e
v2
{-# INLINE dotM #-}

unsafeDotProductIO
  :: (MonadUnliftIO m, Index ix, FoldNumeric r b, Source r b)
  => Array r ix b
  -> Array r ix b
  -> m b
unsafeDotProductIO :: forall (m :: * -> *) ix r b.
(MonadUnliftIO m, Index ix, FoldNumeric r b, Source r b) =>
Array r ix b -> Array r ix b -> m b
unsafeDotProductIO Array r ix b
v1 Array r ix b
v2 = do
  [b]
results <-
    Comp -> (Scheduler RealWorld b -> m ()) -> m [b]
forall (m :: * -> *) a b.
MonadUnliftIO m =>
Comp -> (Scheduler RealWorld a -> m b) -> m [a]
withScheduler Comp
comp ((Scheduler RealWorld b -> m ()) -> m [b])
-> (Scheduler RealWorld b -> m ()) -> m [b]
forall a b. (a -> b) -> a -> b
$ \Scheduler RealWorld b
scheduler ->
      Int -> Int -> (Int -> Int -> m ()) -> m ()
forall a. Int -> Int -> (Int -> Int -> a) -> a
splitLinearly (Scheduler RealWorld b -> Int
forall s a. Scheduler s a -> Int
numWorkers Scheduler RealWorld b
scheduler) Int
totalLength ((Int -> Int -> m ()) -> m ()) -> (Int -> Int -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \Int
chunkLength Int
slackStart -> IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ do
        let n :: Sz Int
n = Int -> Sz Int
forall ix. ix -> Sz ix
SafeSz Int
chunkLength
        Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> IO ()) -> IO ()
forall (f :: * -> *) a.
Applicative f =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> f a) -> f ()
loopA_ Int
0 (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
slackStart) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
chunkLength) ((Int -> IO ()) -> IO ()) -> (Int -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \ !Int
start ->
          Scheduler RealWorld b -> IO b -> IO ()
forall s (m :: * -> *) a.
MonadPrimBase s m =>
Scheduler s a -> m a -> m ()
scheduleWork Scheduler RealWorld b
scheduler (IO b -> IO ()) -> IO b -> IO ()
forall a b. (a -> b) -> a -> b
$
            b -> IO b
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (b -> IO b) -> b -> IO b
forall a b. (a -> b) -> a -> b
$!
              Array r Int b -> Array r Int b -> b
forall ix. Index ix => Array r ix b -> Array r ix b -> b
forall r e ix.
(FoldNumeric r e, Index ix) =>
Array r ix e -> Array r ix e -> e
unsafeDotProduct (Int -> Sz Int -> Array r ix b -> Array r Int b
forall ix.
Index ix =>
Int -> Sz Int -> Array r ix b -> Array r Int b
forall r e ix.
(Source r e, Index ix) =>
Int -> Sz Int -> Array r ix e -> Array r Int e
unsafeLinearSlice Int
start Sz Int
n Array r ix b
v1) (Int -> Sz Int -> Array r ix b -> Array r Int b
forall ix.
Index ix =>
Int -> Sz Int -> Array r ix b -> Array r Int b
forall r e ix.
(Source r e, Index ix) =>
Int -> Sz Int -> Array r ix e -> Array r Int e
unsafeLinearSlice Int
start Sz Int
n Array r ix b
v2)
        Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
slackStart Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
totalLength) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
          let k :: Sz Int
k = Int -> Sz Int
forall ix. ix -> Sz ix
SafeSz (Int
totalLength Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
slackStart)
          Scheduler RealWorld b -> IO b -> IO ()
forall s (m :: * -> *) a.
MonadPrimBase s m =>
Scheduler s a -> m a -> m ()
scheduleWork Scheduler RealWorld b
scheduler (IO b -> IO ()) -> IO b -> IO ()
forall a b. (a -> b) -> a -> b
$
            b -> IO b
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (b -> IO b) -> b -> IO b
forall a b. (a -> b) -> a -> b
$!
              Array r Int b -> Array r Int b -> b
forall ix. Index ix => Array r ix b -> Array r ix b -> b
forall r e ix.
(FoldNumeric r e, Index ix) =>
Array r ix e -> Array r ix e -> e
unsafeDotProduct (Int -> Sz Int -> Array r ix b -> Array r Int b
forall ix.
Index ix =>
Int -> Sz Int -> Array r ix b -> Array r Int b
forall r e ix.
(Source r e, Index ix) =>
Int -> Sz Int -> Array r ix e -> Array r Int e
unsafeLinearSlice Int
slackStart Sz Int
k Array r ix b
v1) (Int -> Sz Int -> Array r ix b -> Array r Int b
forall ix.
Index ix =>
Int -> Sz Int -> Array r ix b -> Array r Int b
forall r e ix.
(Source r e, Index ix) =>
Int -> Sz Int -> Array r ix e -> Array r Int e
unsafeLinearSlice Int
slackStart Sz Int
k Array r ix b
v2)
  b -> m b
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (b -> m b) -> b -> m b
forall a b. (a -> b) -> a -> b
$! (b -> b -> b) -> b -> [b] -> b
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
F.foldl' b -> b -> b
forall a. Num a => a -> a -> a
(+) b
0 [b]
results
  where
    totalLength :: Int
totalLength = Sz ix -> Int
forall ix. Index ix => Sz ix -> Int
totalElem (Array r ix b -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array r ix e -> Sz ix
size Array r ix b
v1)
    comp :: Comp
comp = Array r ix b -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
forall ix e. Array r ix e -> Comp
getComp Array r ix b
v1 Comp -> Comp -> Comp
forall a. Semigroup a => a -> a -> a
<> Array r ix b -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
forall ix e. Array r ix e -> Comp
getComp Array r ix b
v2
{-# INLINE unsafeDotProductIO #-}

-- | Compute L2 norm of an array.
--
-- @since 0.5.6
normL2 :: (FoldNumeric r e, Source r e, Index ix, Floating e) => Array r ix e -> e
normL2 :: forall r e ix.
(FoldNumeric r e, Source r e, Index ix, Floating e) =>
Array r ix e -> e
normL2 Array r ix e
v
  | Array r ix e -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
forall ix e. Array r ix e -> Comp
getComp Array r ix e
v Comp -> Comp -> Bool
forall a. Eq a => a -> a -> Bool
== Comp
Seq = e -> e
forall a. Floating a => a -> a
sqrt (e -> e) -> e -> e
forall a b. (a -> b) -> a -> b
$! Array r ix e -> Int -> e
forall ix. Index ix => Array r ix e -> Int -> e
forall r e ix.
(FoldNumeric r e, Index ix) =>
Array r ix e -> Int -> e
powerSumArray Array r ix e
v Int
2
  | Bool
otherwise = e -> e
forall a. Floating a => a -> a
sqrt (e -> e) -> e -> e
forall a b. (a -> b) -> a -> b
$! IO e -> e
forall a. IO a -> a
unsafePerformIO (IO e -> e) -> IO e -> e
forall a b. (a -> b) -> a -> b
$ Array r ix e -> Int -> IO e
forall (m :: * -> *) ix r b.
(MonadUnliftIO m, Index ix, FoldNumeric r b, Source r b) =>
Array r ix b -> Int -> m b
powerSumArrayIO Array r ix e
v Int
2
{-# INLINE normL2 #-}

powerSumArrayIO
  :: (MonadUnliftIO m, Index ix, FoldNumeric r b, Source r b)
  => Array r ix b
  -> Int
  -> m b
powerSumArrayIO :: forall (m :: * -> *) ix r b.
(MonadUnliftIO m, Index ix, FoldNumeric r b, Source r b) =>
Array r ix b -> Int -> m b
powerSumArrayIO Array r ix b
v Int
p = do
  [b]
results <-
    Comp -> (Scheduler RealWorld b -> m ()) -> m [b]
forall (m :: * -> *) a b.
MonadUnliftIO m =>
Comp -> (Scheduler RealWorld a -> m b) -> m [a]
withScheduler (Array r ix b -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
forall ix e. Array r ix e -> Comp
getComp Array r ix b
v) ((Scheduler RealWorld b -> m ()) -> m [b])
-> (Scheduler RealWorld b -> m ()) -> m [b]
forall a b. (a -> b) -> a -> b
$ \Scheduler RealWorld b
scheduler ->
      Int -> Int -> (Int -> Int -> m ()) -> m ()
forall a. Int -> Int -> (Int -> Int -> a) -> a
splitLinearly (Scheduler RealWorld b -> Int
forall s a. Scheduler s a -> Int
numWorkers Scheduler RealWorld b
scheduler) Int
totalLength ((Int -> Int -> m ()) -> m ()) -> (Int -> Int -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \Int
chunkLength Int
slackStart -> IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ do
        let n :: Sz Int
n = Int -> Sz Int
forall ix. ix -> Sz ix
SafeSz Int
chunkLength
        Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> IO ()) -> IO ()
forall (f :: * -> *) a.
Applicative f =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> f a) -> f ()
loopA_ Int
0 (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
slackStart) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
chunkLength) ((Int -> IO ()) -> IO ()) -> (Int -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \ !Int
start ->
          Scheduler RealWorld b -> IO b -> IO ()
forall s (m :: * -> *) a.
MonadPrimBase s m =>
Scheduler s a -> m a -> m ()
scheduleWork Scheduler RealWorld b
scheduler (IO b -> IO ()) -> IO b -> IO ()
forall a b. (a -> b) -> a -> b
$ b -> IO b
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (b -> IO b) -> b -> IO b
forall a b. (a -> b) -> a -> b
$! Array r Int b -> Int -> b
forall ix. Index ix => Array r ix b -> Int -> b
forall r e ix.
(FoldNumeric r e, Index ix) =>
Array r ix e -> Int -> e
powerSumArray (Int -> Sz Int -> Array r ix b -> Array r Int b
forall ix.
Index ix =>
Int -> Sz Int -> Array r ix b -> Array r Int b
forall r e ix.
(Source r e, Index ix) =>
Int -> Sz Int -> Array r ix e -> Array r Int e
unsafeLinearSlice Int
start Sz Int
n Array r ix b
v) Int
p
        Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
slackStart Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
totalLength) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
          let k :: Sz Int
k = Int -> Sz Int
forall ix. ix -> Sz ix
SafeSz (Int
totalLength Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
slackStart)
          Scheduler RealWorld b -> IO b -> IO ()
forall s (m :: * -> *) a.
MonadPrimBase s m =>
Scheduler s a -> m a -> m ()
scheduleWork Scheduler RealWorld b
scheduler (IO b -> IO ()) -> IO b -> IO ()
forall a b. (a -> b) -> a -> b
$ b -> IO b
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (b -> IO b) -> b -> IO b
forall a b. (a -> b) -> a -> b
$! Array r Int b -> Int -> b
forall ix. Index ix => Array r ix b -> Int -> b
forall r e ix.
(FoldNumeric r e, Index ix) =>
Array r ix e -> Int -> e
powerSumArray (Int -> Sz Int -> Array r ix b -> Array r Int b
forall ix.
Index ix =>
Int -> Sz Int -> Array r ix b -> Array r Int b
forall r e ix.
(Source r e, Index ix) =>
Int -> Sz Int -> Array r ix e -> Array r Int e
unsafeLinearSlice Int
slackStart Sz Int
k Array r ix b
v) Int
p
  b -> m b
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (b -> m b) -> b -> m b
forall a b. (a -> b) -> a -> b
$! (b -> b -> b) -> b -> [b] -> b
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
F.foldl' b -> b -> b
forall a. Num a => a -> a -> a
(+) b
0 [b]
results
  where
    totalLength :: Int
totalLength = Sz ix -> Int
forall ix. Index ix => Sz ix -> Int
totalElem (Array r ix b -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array r ix e -> Sz ix
size Array r ix b
v)
{-# INLINE powerSumArrayIO #-}

-- | Multiply a matrix by a column vector. Same as `!><` but produces monadic
-- computation that allows for handling failure.
--
-- /__Throws Exception__/: `SizeMismatchException` when inner dimensions of arrays do not match.
--
-- @since 0.5.6
(.><)
  :: (MonadThrow m, FoldNumeric r e, Source r e)
  => Matrix r e
  -- ^ Matrix
  -> Vector r e
  -- ^ Column vector (Used many times, so make sure it is computed)
  -> m (Vector D e)
.>< :: forall (m :: * -> *) r e.
(MonadThrow m, FoldNumeric r e, Source r e) =>
Matrix r e -> Vector r e -> m (Vector D e)
(.><) Matrix r e
mm Vector r e
v
  | Int
mCols Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
n = SizeException -> m (Vector D e)
forall e a. (HasCallStack, Exception e) => e -> m a
forall (m :: * -> *) e a.
(MonadThrow m, HasCallStack, Exception e) =>
e -> m a
throwM (SizeException -> m (Vector D e))
-> SizeException -> m (Vector D e)
forall a b. (a -> b) -> a -> b
$ Sz Ix2 -> Sz Ix2 -> SizeException
forall ix. Index ix => Sz ix -> Sz ix -> SizeException
SizeMismatchException (Matrix r e -> Sz Ix2
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array r ix e -> Sz ix
size Matrix r e
mm) (Int -> Int -> Sz Ix2
Sz2 Int
n Int
1)
  | Int
mRows Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 Bool -> Bool -> Bool
|| Int
mCols Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = Vector D e -> m (Vector D e)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Vector D e -> m (Vector D e)) -> Vector D e -> m (Vector D e)
forall a b. (a -> b) -> a -> b
$ Comp -> Vector D e -> Vector D e
forall r ix e. Strategy r => Comp -> Array r ix e -> Array r ix e
forall ix e. Comp -> Array D ix e -> Array D ix e
setComp Comp
comp Vector D e
forall r ix e. Load r ix e => Array r ix e
empty
  | Bool
otherwise = Vector D e -> m (Vector D e)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Vector D e -> m (Vector D e)) -> Vector D e -> m (Vector D e)
forall a b. (a -> b) -> a -> b
$ Comp -> Sz Int -> (Int -> e) -> Vector D e
forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray Comp
comp (Int -> Sz Int
Sz1 Int
mRows) ((Int -> e) -> Vector D e) -> (Int -> e) -> Vector D e
forall a b. (a -> b) -> a -> b
$ \Int
i ->
      Vector r e -> Vector r e -> e
forall ix. Index ix => Array r ix e -> Array r ix e -> e
forall r e ix.
(FoldNumeric r e, Index ix) =>
Array r ix e -> Array r ix e -> e
unsafeDotProduct (Int -> Sz Int -> Matrix r e -> Vector r e
forall ix. Index ix => Int -> Sz Int -> Array r ix e -> Vector r e
forall r e ix.
(Source r e, Index ix) =>
Int -> Sz Int -> Array r ix e -> Array r Int e
unsafeLinearSlice (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
n) Sz Int
sz Matrix r e
mm) Vector r e
v
  where
    comp :: Comp
comp = Matrix r e -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
forall ix e. Array r ix e -> Comp
getComp Matrix r e
mm Comp -> Comp -> Comp
forall a. Semigroup a => a -> a -> a
<> Vector r e -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
forall ix e. Array r ix e -> Comp
getComp Vector r e
v
    Sz2 Int
mRows Int
mCols = Matrix r e -> Sz Ix2
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array r ix e -> Sz ix
size Matrix r e
mm
    sz :: Sz Int
sz@(Sz1 Int
n) = Vector r e -> Sz Int
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array r ix e -> Sz ix
size Vector r e
v
{-# INLINE (.><) #-}

-- | Multiply matrix by a column vector. Same as `.><` but returns computed version of a vector
--
-- /__Throws Exception__/: `SizeMismatchException` when inner dimensions of arrays do not match.
--
-- @since 0.5.7
multiplyMatrixByVector
  :: (MonadThrow m, Numeric r e, Manifest r e)
  => Matrix r e
  -- ^ Matrix
  -> Vector r e
  -- ^ Column vector (Used many times, so make sure it is computed)
  -> m (Vector r e)
multiplyMatrixByVector :: forall (m :: * -> *) r e.
(MonadThrow m, Numeric r e, Manifest r e) =>
Matrix r e -> Vector r e -> m (Vector r e)
multiplyMatrixByVector Matrix r e
mm Vector r e
v = Array D Int e -> Vector r e
forall r ix e r'.
(Manifest r e, Load r' ix e) =>
Array r' ix e -> Array r ix e
compute (Array D Int e -> Vector r e)
-> m (Array D Int e) -> m (Vector r e)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Matrix r e
mm Matrix r e -> Vector r e -> m (Array D Int e)
forall (m :: * -> *) r e.
(MonadThrow m, FoldNumeric r e, Source r e) =>
Matrix r e -> Vector r e -> m (Vector D e)
.>< Vector r e
v
{-# INLINE multiplyMatrixByVector #-}

-- | Multiply a matrix by a column vector
--
-- [Partial] Throws impure exception when inner dimensions do not agree
--
-- @since 0.5.6
(!><)
  :: (Numeric r e, Source r e)
  => Matrix r e
  -- ^ Matrix
  -> Vector r e
  -- ^ Column vector (Used many times, so make sure it is computed)
  -> Vector D e
!>< :: forall r e.
(Numeric r e, Source r e) =>
Matrix r e -> Vector r e -> Vector D e
(!><) Matrix r e
mm Vector r e
v = Either SomeException (Vector D e) -> Vector D e
forall a. HasCallStack => Either SomeException a -> a
throwEither (Matrix r e
mm Matrix r e -> Vector r e -> Either SomeException (Vector D e)
forall (m :: * -> *) r e.
(MonadThrow m, FoldNumeric r e, Source r e) =>
Matrix r e -> Vector r e -> m (Vector D e)
.>< Vector r e
v)
{-# INLINE (!><) #-}

-- | Multiply a row vector by a matrix. Same as `><!` but produces monadic computation
-- that allows for handling failure.
--
-- /__Throws Exception__/: `SizeMismatchException` when inner dimensions of arrays do not match.
--
-- @since 0.5.6
(><.)
  :: (MonadThrow m, Numeric r e, Manifest r e)
  => Vector r e
  -- ^ Row vector
  -> Matrix r e
  -- ^ Matrix
  -> m (Vector r e)
><. :: forall (m :: * -> *) r e.
(MonadThrow m, Numeric r e, Manifest r e) =>
Vector r e -> Matrix r e -> m (Vector r e)
(><.) = Vector r e -> Matrix r e -> m (Vector r e)
forall (m :: * -> *) r e.
(MonadThrow m, Numeric r e, Manifest r e) =>
Vector r e -> Matrix r e -> m (Vector r e)
multiplyVectorByMatrix
{-# INLINE (><.) #-}

-- | Multiply a row vector by a matrix. Same as `><.` but returns computed vector instead of
-- a delayed one.
--
-- /__Throws Exception__/: `SizeMismatchException` when inner dimensions of arrays do not match.
--
-- @since 0.5.7
multiplyVectorByMatrix
  :: (MonadThrow m, Numeric r e, Manifest r e)
  => Vector r e
  -- ^ Row vector
  -> Matrix r e
  -- ^ Matrix
  -> m (Vector r e)
multiplyVectorByMatrix :: forall (m :: * -> *) r e.
(MonadThrow m, Numeric r e, Manifest r e) =>
Vector r e -> Matrix r e -> m (Vector r e)
multiplyVectorByMatrix Vector r e
v Matrix r e
mm
  | Int
mRows Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
n = SizeException -> m (Vector r e)
forall e a. (HasCallStack, Exception e) => e -> m a
forall (m :: * -> *) e a.
(MonadThrow m, HasCallStack, Exception e) =>
e -> m a
throwM (SizeException -> m (Vector r e))
-> SizeException -> m (Vector r e)
forall a b. (a -> b) -> a -> b
$ Sz Ix2 -> Sz Ix2 -> SizeException
forall ix. Index ix => Sz ix -> Sz ix -> SizeException
SizeMismatchException (Int -> Int -> Sz Ix2
Sz2 Int
1 Int
n) (Matrix r e -> Sz Ix2
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array r ix e -> Sz ix
size Matrix r e
mm)
  | Int
mRows Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 Bool -> Bool -> Bool
|| Int
mCols Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = Vector r e -> m (Vector r e)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Vector r e -> m (Vector r e)) -> Vector r e -> m (Vector r e)
forall a b. (a -> b) -> a -> b
$ (forall s. ST s (Vector r e)) -> Vector r e
forall a. (forall s. ST s a) -> a
runST (Comp -> MArray (PrimState (ST s)) r Int e -> ST s (Vector r e)
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Comp -> MArray (PrimState m) r ix e -> m (Array r ix e)
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
Comp -> MArray (PrimState m) r ix e -> m (Array r ix e)
unsafeFreeze Comp
comp (MArray s r Int e -> ST s (Vector r e))
-> ST s (MArray s r Int e) -> ST s (Vector r e)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Sz Int -> ST s (MArray (PrimState (ST s)) r Int e)
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Sz ix -> m (MArray (PrimState m) r ix e)
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
Sz ix -> m (MArray (PrimState m) r ix e)
unsafeNew Sz Int
forall ix. Index ix => Sz ix
zeroSz)
  | Bool
otherwise =
      Vector r e -> m (Vector r e)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Vector r e -> m (Vector r e)) -> Vector r e -> m (Vector r e)
forall a b. (a -> b) -> a -> b
$!
        IO (Vector r e) -> Vector r e
forall a. IO a -> a
unsafePerformIO (IO (Vector r e) -> Vector r e) -> IO (Vector r e) -> Vector r e
forall a b. (a -> b) -> a -> b
$ do
          MArray RealWorld r Int e
mv <- Sz Int -> e -> IO (MArray (PrimState IO) r Int e)
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Sz ix -> e -> m (MArray (PrimState m) r ix e)
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
Sz ix -> e -> m (MArray (PrimState m) r ix e)
newMArray (Int -> Sz Int
forall ix. Index ix => ix -> Sz ix
Sz Int
mCols) e
0
          Comp -> (Scheduler RealWorld () -> IO ()) -> IO ()
withMassivScheduler_ Comp
comp ((Scheduler RealWorld () -> IO ()) -> IO ())
-> (Scheduler RealWorld () -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Scheduler RealWorld ()
scheduler -> do
            let loopCols :: e -> Int -> Int -> Int -> IO ()
loopCols e
x Int
ivto =
                  ((Int -> Int -> IO ()) -> Int -> Int -> IO ())
-> Int -> Int -> IO ()
forall a. (a -> a) -> a
fix (((Int -> Int -> IO ()) -> Int -> Int -> IO ())
 -> Int -> Int -> IO ())
-> ((Int -> Int -> IO ()) -> Int -> Int -> IO ())
-> Int
-> Int
-> IO ()
forall a b. (a -> b) -> a -> b
$ \Int -> Int -> IO ()
go Int
im Int
iv ->
                    Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
iv Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
ivto) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
                      e
_ <- MArray (PrimState IO) r Int e -> (e -> IO e) -> Int -> IO e
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> (e -> m e) -> Int -> m e
unsafeLinearModify MArray RealWorld r Int e
MArray (PrimState IO) r Int e
mv (\e
a -> e -> IO e
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (e -> IO e) -> e -> IO e
forall a b. (a -> b) -> a -> b
$ e
a e -> e -> e
forall a. Num a => a -> a -> a
+ Matrix r e -> Int -> e
forall ix. Index ix => Array r ix e -> Int -> e
forall r e ix. (Source r e, Index ix) => Array r ix e -> Int -> e
unsafeLinearIndex Matrix r e
mm Int
im e -> e -> e
forall a. Num a => a -> a -> a
* e
x) Int
iv
                      Int -> Int -> IO ()
go (Int
im Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int
iv Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
                loopRows :: Int -> Int -> Int -> IO ()
loopRows Int
i0 Int
from Int
to =
                  (((Int -> IO ()) -> Int -> IO ()) -> Int -> IO ())
-> Int -> ((Int -> IO ()) -> Int -> IO ()) -> IO ()
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((Int -> IO ()) -> Int -> IO ()) -> Int -> IO ()
forall a. (a -> a) -> a
fix Int
i0 (((Int -> IO ()) -> Int -> IO ()) -> IO ())
-> ((Int -> IO ()) -> Int -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int -> IO ()
go Int
i ->
                    Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
mRows) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
                      e -> Int -> Int -> Int -> IO ()
loopCols (Vector r e -> Int -> e
forall ix. Index ix => Array r ix e -> Int -> e
forall r e ix. (Source r e, Index ix) => Array r ix e -> Int -> e
unsafeLinearIndex Vector r e
v Int
i) Int
to (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
mCols Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
from) Int
from
                      Int -> IO ()
go (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
            Scheduler RealWorld () -> Int -> (Int -> Int -> IO ()) -> IO ()
forall s (m :: * -> *).
MonadPrimBase s m =>
Scheduler s () -> Int -> (Int -> Int -> m ()) -> m ()
splitLinearlyM_ Scheduler RealWorld ()
scheduler Int
mCols (Int -> Int -> Int -> IO ()
loopRows Int
0)
          Comp -> MArray (PrimState IO) r Int e -> IO (Vector r e)
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Comp -> MArray (PrimState m) r ix e -> m (Array r ix e)
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
Comp -> MArray (PrimState m) r ix e -> m (Array r ix e)
unsafeFreeze Comp
comp MArray RealWorld r Int e
MArray (PrimState IO) r Int e
mv
  where
    comp :: Comp
comp = Matrix r e -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
forall ix e. Array r ix e -> Comp
getComp Matrix r e
mm Comp -> Comp -> Comp
forall a. Semigroup a => a -> a -> a
<> Vector r e -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
forall ix e. Array r ix e -> Comp
getComp Vector r e
v
    Sz2 Int
mRows Int
mCols = Matrix r e -> Sz Ix2
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array r ix e -> Sz ix
size Matrix r e
mm
    Sz1 Int
n = Vector r e -> Sz Int
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array r ix e -> Sz ix
size Vector r e
v
{-# INLINE multiplyVectorByMatrix #-}

-- | Multiply a row vector by a matrix.
--
-- [Partial] Throws impure exception when inner dimensions do not agree
--
-- @since 0.5.6
(><!)
  :: (Numeric r e, Manifest r e)
  => Vector r e
  -- ^ Row vector (Used many times, so make sure it is computed)
  -> Matrix r e
  -- ^ Matrix
  -> Vector r e
><! :: forall r e.
(Numeric r e, Manifest r e) =>
Vector r e -> Matrix r e -> Vector r e
(><!) Vector r e
v Matrix r e
mm = Either SomeException (Vector r e) -> Vector r e
forall a. HasCallStack => Either SomeException a -> a
throwEither (Vector r e
v Vector r e -> Matrix r e -> Either SomeException (Vector r e)
forall (m :: * -> *) r e.
(MonadThrow m, Numeric r e, Manifest r e) =>
Vector r e -> Matrix r e -> m (Vector r e)
><. Matrix r e
mm)
{-# INLINE (><!) #-}

-- | Multiply two matrices together.
--
-- [Partial] Inner dimension must agree
--
-- ====__Examples__
--
-- >>> import Data.Massiv.Array
-- >>> a1 = makeArrayR P Seq (Sz2 5 6) $ \(i :. j) -> i + j
-- >>> a2 = makeArrayR P Seq (Sz2 6 5) $ \(i :. j) -> i - j
-- >>> a1 !><! a2
-- Array P Seq (Sz (5 :. 5))
--   [ [ 55, 40, 25, 10, -5 ]
--   , [ 70, 49, 28, 7, -14 ]
--   , [ 85, 58, 31, 4, -23 ]
--   , [ 100, 67, 34, 1, -32 ]
--   , [ 115, 76, 37, -2, -41 ]
--   ]
--
-- @since 0.5.6
(!><!) :: (Numeric r e, Manifest r e) => Matrix r e -> Matrix r e -> Matrix r e
!><! :: forall r e.
(Numeric r e, Manifest r e) =>
Matrix r e -> Matrix r e -> Matrix r e
(!><!) Matrix r e
a1 Matrix r e
a2 = Either SomeException (Matrix r e) -> Matrix r e
forall a. HasCallStack => Either SomeException a -> a
throwEither (Matrix r e
a1 Matrix r e -> Matrix r e -> Either SomeException (Matrix r e)
forall r e (m :: * -> *).
(Numeric r e, Manifest r e, MonadThrow m) =>
Matrix r e -> Matrix r e -> m (Matrix r e)
`multiplyMatrices` Matrix r e
a2)
{-# INLINE (!><!) #-}

-- | Matrix multiplication. Same as `!><!` but produces monadic computation that allows
-- for handling failure.
--
-- /__Throws Exception__/: `SizeMismatchException` when inner dimensions of arrays do not match.
--
-- @since 0.5.6
(.><.) :: (Numeric r e, Manifest r e, MonadThrow m) => Matrix r e -> Matrix r e -> m (Matrix r e)
.><. :: forall r e (m :: * -> *).
(Numeric r e, Manifest r e, MonadThrow m) =>
Matrix r e -> Matrix r e -> m (Matrix r e)
(.><.) = Matrix r e -> Matrix r e -> m (Matrix r e)
forall r e (m :: * -> *).
(Numeric r e, Manifest r e, MonadThrow m) =>
Matrix r e -> Matrix r e -> m (Matrix r e)
multiplyMatrices
{-# INLINE (.><.) #-}

-- | Synonym for `.><.`
--
-- @since 0.5.6
multiplyMatrices
  :: (Numeric r e, Manifest r e, MonadThrow m) => Matrix r e -> Matrix r e -> m (Matrix r e)
multiplyMatrices :: forall r e (m :: * -> *).
(Numeric r e, Manifest r e, MonadThrow m) =>
Matrix r e -> Matrix r e -> m (Matrix r e)
multiplyMatrices Matrix r e
arrA Matrix r e
arrB
  -- mA == 1 = -- TODO: call multiplyVectorByMatrix
  -- nA == 1 = -- TODO: call multiplyMatrixByVector
  | Int
nA Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
mB = SizeException -> m (Matrix r e)
forall e a. (HasCallStack, Exception e) => e -> m a
forall (m :: * -> *) e a.
(MonadThrow m, HasCallStack, Exception e) =>
e -> m a
throwM (SizeException -> m (Matrix r e))
-> SizeException -> m (Matrix r e)
forall a b. (a -> b) -> a -> b
$ Sz Ix2 -> Sz Ix2 -> SizeException
forall ix. Index ix => Sz ix -> Sz ix -> SizeException
SizeMismatchException (Matrix r e -> Sz Ix2
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array r ix e -> Sz ix
size Matrix r e
arrA) (Matrix r e -> Sz Ix2
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array r ix e -> Sz ix
size Matrix r e
arrB)
  | Matrix r e -> Bool
forall ix r e. (Index ix, Size r) => Array r ix e -> Bool
isEmpty Matrix r e
arrA Bool -> Bool -> Bool
|| Matrix r e -> Bool
forall ix r e. (Index ix, Size r) => Array r ix e -> Bool
isEmpty Matrix r e
arrB = Matrix r e -> m (Matrix r e)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Matrix r e -> m (Matrix r e)) -> Matrix r e -> m (Matrix r e)
forall a b. (a -> b) -> a -> b
$ (forall s. ST s (Matrix r e)) -> Matrix r e
forall a. (forall s. ST s a) -> a
runST (Comp -> MArray (PrimState (ST s)) r Ix2 e -> ST s (Matrix r e)
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Comp -> MArray (PrimState m) r ix e -> m (Array r ix e)
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
Comp -> MArray (PrimState m) r ix e -> m (Array r ix e)
unsafeFreeze Comp
comp (MArray s r Ix2 e -> ST s (Matrix r e))
-> ST s (MArray s r Ix2 e) -> ST s (Matrix r e)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< Sz Ix2 -> ST s (MArray (PrimState (ST s)) r Ix2 e)
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Sz ix -> m (MArray (PrimState m) r ix e)
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
Sz ix -> m (MArray (PrimState m) r ix e)
unsafeNew Sz Ix2
forall ix. Index ix => Sz ix
zeroSz)
  | Bool
otherwise = Matrix r e -> m (Matrix r e)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Matrix r e -> m (Matrix r e)) -> Matrix r e -> m (Matrix r e)
forall a b. (a -> b) -> a -> b
$! IO (Matrix r e) -> Matrix r e
forall a. IO a -> a
unsafePerformIO (IO (Matrix r e) -> Matrix r e) -> IO (Matrix r e) -> Matrix r e
forall a b. (a -> b) -> a -> b
$ do
      MArray RealWorld r Ix2 e
marrC <- Sz Ix2 -> e -> IO (MArray (PrimState IO) r Ix2 e)
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Sz ix -> e -> m (MArray (PrimState m) r ix e)
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
Sz ix -> e -> m (MArray (PrimState m) r ix e)
newMArray (Ix2 -> Sz Ix2
forall ix. ix -> Sz ix
SafeSz (Int
mA Int -> Int -> Ix2
:. Int
nB)) e
0
      Comp -> (Scheduler RealWorld () -> IO ()) -> IO ()
forall (m :: * -> *) a b.
MonadUnliftIO m =>
Comp -> (Scheduler RealWorld a -> m b) -> m ()
withScheduler_ Comp
comp ((Scheduler RealWorld () -> IO ()) -> IO ())
-> (Scheduler RealWorld () -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Scheduler RealWorld ()
scheduler -> do
        let withC00 :: Int -> Int -> (Int -> e -> IO ()) -> IO ()
withC00 Int
iA Int
jB Int -> e -> IO ()
f =
              let !ixC00 :: Int
ixC00 = Int
iA Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
nB Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
jB
               in Int -> e -> IO ()
f Int
ixC00 (e -> IO ()) -> IO e -> IO ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< MArray (PrimState IO) r Ix2 e -> Int -> IO e
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> m e
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> m e
unsafeLinearRead MArray RealWorld r Ix2 e
MArray (PrimState IO) r Ix2 e
marrC Int
ixC00
            withC01 :: Int -> (Int -> e -> IO ()) -> IO ()
withC01 Int
ixC00 Int -> e -> IO ()
f =
              let !ixC01 :: Int
ixC01 = Int
ixC00 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
               in Int -> e -> IO ()
f Int
ixC01 (e -> IO ()) -> IO e -> IO ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< MArray (PrimState IO) r Ix2 e -> Int -> IO e
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> m e
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> m e
unsafeLinearRead MArray RealWorld r Ix2 e
MArray (PrimState IO) r Ix2 e
marrC Int
ixC01
            withC10 :: Int -> (Int -> e -> IO ()) -> IO ()
withC10 Int
ixC00 Int -> e -> IO ()
f =
              let !ixC10 :: Int
ixC10 = Int
ixC00 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
nB
               in Int -> e -> IO ()
f Int
ixC10 (e -> IO ()) -> IO e -> IO ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< MArray (PrimState IO) r Ix2 e -> Int -> IO e
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> m e
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> m e
unsafeLinearRead MArray RealWorld r Ix2 e
MArray (PrimState IO) r Ix2 e
marrC Int
ixC10
            withC11 :: Int -> (Int -> e -> IO ()) -> IO ()
withC11 Int
ixC01 Int -> e -> IO ()
f =
              let !ixC11 :: Int
ixC11 = Int
ixC01 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
nB
               in Int -> e -> IO ()
f Int
ixC11 (e -> IO ()) -> IO e -> IO ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< MArray (PrimState IO) r Ix2 e -> Int -> IO e
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> m e
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> m e
unsafeLinearRead MArray RealWorld r Ix2 e
MArray (PrimState IO) r Ix2 e
marrC Int
ixC11
            withB00 :: Int -> Int -> (Int -> e -> IO ()) -> IO ()
withB00 Int
iB Int
jB Int -> e -> IO ()
f =
              let !ixB00 :: Int
ixB00 = Int
iB Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
nB Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
jB
               in Int -> e -> IO ()
f Int
ixB00 (e -> IO ()) -> e -> IO ()
forall a b. (a -> b) -> a -> b
$! Matrix r e -> Int -> e
forall ix. Index ix => Array r ix e -> Int -> e
forall r e ix. (Source r e, Index ix) => Array r ix e -> Int -> e
unsafeLinearIndex Matrix r e
arrB Int
ixB00
            withB00B10 :: Int -> Int -> (Int -> e -> Int -> e -> IO ()) -> IO ()
withB00B10 Int
iB Int
jB Int -> e -> Int -> e -> IO ()
f =
              Int -> Int -> (Int -> e -> IO ()) -> IO ()
withB00 Int
iB Int
jB ((Int -> e -> IO ()) -> IO ()) -> (Int -> e -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
ixB00 e
b00 ->
                let !ixB10 :: Int
ixB10 = Int
ixB00 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
nB
                 in Int -> e -> Int -> e -> IO ()
f Int
ixB00 e
b00 Int
ixB10 (e -> IO ()) -> e -> IO ()
forall a b. (a -> b) -> a -> b
$! Matrix r e -> Int -> e
forall ix. Index ix => Array r ix e -> Int -> e
forall r e ix. (Source r e, Index ix) => Array r ix e -> Int -> e
unsafeLinearIndex Matrix r e
arrB Int
ixB10
            withA00 :: Int -> Int -> (Int -> e -> IO ()) -> IO ()
withA00 Int
iA Int
jA Int -> e -> IO ()
f =
              let !ixA00 :: Int
ixA00 = Int
iA Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
nA Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
jA
               in Int -> e -> IO ()
f Int
ixA00 (e -> IO ()) -> e -> IO ()
forall a b. (a -> b) -> a -> b
$! Matrix r e -> Int -> e
forall ix. Index ix => Array r ix e -> Int -> e
forall r e ix. (Source r e, Index ix) => Array r ix e -> Int -> e
unsafeLinearIndex Matrix r e
arrA Int
ixA00
            withA00A10 :: Int -> Int -> (Int -> e -> Int -> e -> IO ()) -> IO ()
withA00A10 Int
iA Int
jA Int -> e -> Int -> e -> IO ()
f =
              Int -> Int -> (Int -> e -> IO ()) -> IO ()
withA00 Int
iA Int
jA ((Int -> e -> IO ()) -> IO ()) -> (Int -> e -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
ixA00 e
a00 ->
                let !ixA10 :: Int
ixA10 = Int
ixA00 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
nA
                 in Int -> e -> Int -> e -> IO ()
f Int
ixA00 e
a00 Int
ixA10 (e -> IO ()) -> e -> IO ()
forall a b. (a -> b) -> a -> b
$! Matrix r e -> Int -> e
forall ix. Index ix => Array r ix e -> Int -> e
forall r e ix. (Source r e, Index ix) => Array r ix e -> Int -> e
unsafeLinearIndex Matrix r e
arrA Int
ixA10
        let loopColsB_UnRowBColA_UnRowA :: e -> e -> e -> e -> Int -> Int -> Int -> IO ()
loopColsB_UnRowBColA_UnRowA e
a00 e
a01 e
a10 e
a11 Int
iA Int
iB Int
jB
              | Int
jB Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n2B = do
                  Int -> Int -> (Int -> e -> Int -> e -> IO ()) -> IO ()
withB00B10 Int
iB Int
jB ((Int -> e -> Int -> e -> IO ()) -> IO ())
-> (Int -> e -> Int -> e -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
ixB00 e
b00 Int
ixB10 e
b10 -> do
                    let !b01 :: e
b01 = Matrix r e -> Int -> e
forall ix. Index ix => Array r ix e -> Int -> e
forall r e ix. (Source r e, Index ix) => Array r ix e -> Int -> e
unsafeLinearIndex Matrix r e
arrB (Int
ixB00 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
                        !b11 :: e
b11 = Matrix r e -> Int -> e
forall ix. Index ix => Array r ix e -> Int -> e
forall r e ix. (Source r e, Index ix) => Array r ix e -> Int -> e
unsafeLinearIndex Matrix r e
arrB (Int
ixB10 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
                    Int -> Int -> (Int -> e -> IO ()) -> IO ()
withC00 Int
iA Int
jB ((Int -> e -> IO ()) -> IO ()) -> (Int -> e -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
ixC00 e
c00 -> do
                      MArray (PrimState IO) r Ix2 e -> Int -> e -> IO ()
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
unsafeLinearWrite MArray RealWorld r Ix2 e
MArray (PrimState IO) r Ix2 e
marrC Int
ixC00 (e
c00 e -> e -> e
forall a. Num a => a -> a -> a
+ e
a00 e -> e -> e
forall a. Num a => a -> a -> a
* e
b00 e -> e -> e
forall a. Num a => a -> a -> a
+ e
a01 e -> e -> e
forall a. Num a => a -> a -> a
* e
b10)
                      Int -> (Int -> e -> IO ()) -> IO ()
withC01 Int
ixC00 ((Int -> e -> IO ()) -> IO ()) -> (Int -> e -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
ixC01 e
c01 -> do
                        MArray (PrimState IO) r Ix2 e -> Int -> e -> IO ()
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
unsafeLinearWrite MArray RealWorld r Ix2 e
MArray (PrimState IO) r Ix2 e
marrC Int
ixC01 (e
c01 e -> e -> e
forall a. Num a => a -> a -> a
+ e
a00 e -> e -> e
forall a. Num a => a -> a -> a
* e
b01 e -> e -> e
forall a. Num a => a -> a -> a
+ e
a01 e -> e -> e
forall a. Num a => a -> a -> a
* e
b11)
                        Int -> (Int -> e -> IO ()) -> IO ()
withC10 Int
ixC00 ((Int -> e -> IO ()) -> IO ()) -> (Int -> e -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
ixC10 e
c10 ->
                          MArray (PrimState IO) r Ix2 e -> Int -> e -> IO ()
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
unsafeLinearWrite MArray RealWorld r Ix2 e
MArray (PrimState IO) r Ix2 e
marrC Int
ixC10 (e
c10 e -> e -> e
forall a. Num a => a -> a -> a
+ e
a10 e -> e -> e
forall a. Num a => a -> a -> a
* e
b00 e -> e -> e
forall a. Num a => a -> a -> a
+ e
a11 e -> e -> e
forall a. Num a => a -> a -> a
* e
b10)
                        Int -> (Int -> e -> IO ()) -> IO ()
withC11 Int
ixC01 ((Int -> e -> IO ()) -> IO ()) -> (Int -> e -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
ixC11 e
c11 ->
                          MArray (PrimState IO) r Ix2 e -> Int -> e -> IO ()
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
unsafeLinearWrite MArray RealWorld r Ix2 e
MArray (PrimState IO) r Ix2 e
marrC Int
ixC11 (e
c11 e -> e -> e
forall a. Num a => a -> a -> a
+ e
a10 e -> e -> e
forall a. Num a => a -> a -> a
* e
b01 e -> e -> e
forall a. Num a => a -> a -> a
+ e
a11 e -> e -> e
forall a. Num a => a -> a -> a
* e
b11)
                  e -> e -> e -> e -> Int -> Int -> Int -> IO ()
loopColsB_UnRowBColA_UnRowA e
a00 e
a01 e
a10 e
a11 Int
iA Int
iB (Int
jB Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
2)
              | Int
jB Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
nB = Int -> Int -> (Int -> e -> Int -> e -> IO ()) -> IO ()
withB00B10 Int
iB Int
jB ((Int -> e -> Int -> e -> IO ()) -> IO ())
-> (Int -> e -> Int -> e -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
_ e
b00 Int
_ e
b10 ->
                  Int -> Int -> (Int -> e -> IO ()) -> IO ()
withC00 Int
iA Int
jB ((Int -> e -> IO ()) -> IO ()) -> (Int -> e -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
ixC00 e
c00 -> do
                    MArray (PrimState IO) r Ix2 e -> Int -> e -> IO ()
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
unsafeLinearWrite MArray RealWorld r Ix2 e
MArray (PrimState IO) r Ix2 e
marrC Int
ixC00 (e
c00 e -> e -> e
forall a. Num a => a -> a -> a
+ e
a00 e -> e -> e
forall a. Num a => a -> a -> a
* e
b00 e -> e -> e
forall a. Num a => a -> a -> a
+ e
a01 e -> e -> e
forall a. Num a => a -> a -> a
* e
b10)
                    Int -> (Int -> e -> IO ()) -> IO ()
withC10 Int
ixC00 ((Int -> e -> IO ()) -> IO ()) -> (Int -> e -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
ixC10 e
c10 ->
                      MArray (PrimState IO) r Ix2 e -> Int -> e -> IO ()
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
unsafeLinearWrite MArray RealWorld r Ix2 e
MArray (PrimState IO) r Ix2 e
marrC Int
ixC10 (e
c10 e -> e -> e
forall a. Num a => a -> a -> a
+ e
a10 e -> e -> e
forall a. Num a => a -> a -> a
* e
b00 e -> e -> e
forall a. Num a => a -> a -> a
+ e
a11 e -> e -> e
forall a. Num a => a -> a -> a
* e
b10)
              | Bool
otherwise = () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

            loopColsB_UnRowBColA_RowA :: e -> e -> Int -> Int -> Int -> IO ()
loopColsB_UnRowBColA_RowA e
a00 e
a01 Int
iA Int
iB Int
jB
              | Int
jB Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n2B = do
                  Int -> Int -> (Int -> e -> Int -> e -> IO ()) -> IO ()
withB00B10 Int
iB Int
jB ((Int -> e -> Int -> e -> IO ()) -> IO ())
-> (Int -> e -> Int -> e -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
ixB00 e
b00 Int
ixB10 e
b10 -> do
                    let !b01 :: e
b01 = Matrix r e -> Int -> e
forall ix. Index ix => Array r ix e -> Int -> e
forall r e ix. (Source r e, Index ix) => Array r ix e -> Int -> e
unsafeLinearIndex Matrix r e
arrB (Int
ixB00 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
                        !b11 :: e
b11 = Matrix r e -> Int -> e
forall ix. Index ix => Array r ix e -> Int -> e
forall r e ix. (Source r e, Index ix) => Array r ix e -> Int -> e
unsafeLinearIndex Matrix r e
arrB (Int
ixB10 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
                    Int -> Int -> (Int -> e -> IO ()) -> IO ()
withC00 Int
iA Int
jB ((Int -> e -> IO ()) -> IO ()) -> (Int -> e -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
ixC00 e
c00 -> do
                      MArray (PrimState IO) r Ix2 e -> Int -> e -> IO ()
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
unsafeLinearWrite MArray RealWorld r Ix2 e
MArray (PrimState IO) r Ix2 e
marrC Int
ixC00 (e
c00 e -> e -> e
forall a. Num a => a -> a -> a
+ e
a00 e -> e -> e
forall a. Num a => a -> a -> a
* e
b00 e -> e -> e
forall a. Num a => a -> a -> a
+ e
a01 e -> e -> e
forall a. Num a => a -> a -> a
* e
b10)
                      Int -> (Int -> e -> IO ()) -> IO ()
withC01 Int
ixC00 ((Int -> e -> IO ()) -> IO ()) -> (Int -> e -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
ixC01 e
c01 ->
                        MArray (PrimState IO) r Ix2 e -> Int -> e -> IO ()
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
unsafeLinearWrite MArray RealWorld r Ix2 e
MArray (PrimState IO) r Ix2 e
marrC Int
ixC01 (e
c01 e -> e -> e
forall a. Num a => a -> a -> a
+ e
a00 e -> e -> e
forall a. Num a => a -> a -> a
* e
b01 e -> e -> e
forall a. Num a => a -> a -> a
+ e
a01 e -> e -> e
forall a. Num a => a -> a -> a
* e
b11)
                  e -> e -> Int -> Int -> Int -> IO ()
loopColsB_UnRowBColA_RowA e
a00 e
a01 Int
iA Int
iB (Int
jB Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
2)
              | Int
jB Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
nB = Int -> Int -> (Int -> e -> Int -> e -> IO ()) -> IO ()
withB00B10 Int
iB Int
jB ((Int -> e -> Int -> e -> IO ()) -> IO ())
-> (Int -> e -> Int -> e -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
_ e
b00 Int
_ e
b10 ->
                  Int -> Int -> (Int -> e -> IO ()) -> IO ()
withC00 Int
iA Int
jB ((Int -> e -> IO ()) -> IO ()) -> (Int -> e -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
ixC00 e
c00 ->
                    MArray (PrimState IO) r Ix2 e -> Int -> e -> IO ()
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
unsafeLinearWrite MArray RealWorld r Ix2 e
MArray (PrimState IO) r Ix2 e
marrC Int
ixC00 (e
c00 e -> e -> e
forall a. Num a => a -> a -> a
+ e
a00 e -> e -> e
forall a. Num a => a -> a -> a
* e
b00 e -> e -> e
forall a. Num a => a -> a -> a
+ e
a01 e -> e -> e
forall a. Num a => a -> a -> a
* e
b10)
              | Bool
otherwise = () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

            loopColsB_RowBColA_UnRowA :: e -> e -> Int -> Int -> Int -> IO ()
loopColsB_RowBColA_UnRowA e
a00 e
a10 Int
iA Int
iB Int
jB
              | Int
jB Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n2B = do
                  Int -> Int -> (Int -> e -> IO ()) -> IO ()
withB00 Int
iB Int
jB ((Int -> e -> IO ()) -> IO ()) -> (Int -> e -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
ixB00 e
b00 -> do
                    let !b01 :: e
b01 = Matrix r e -> Int -> e
forall ix. Index ix => Array r ix e -> Int -> e
forall r e ix. (Source r e, Index ix) => Array r ix e -> Int -> e
unsafeLinearIndex Matrix r e
arrB (Int
ixB00 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
                    Int -> Int -> (Int -> e -> IO ()) -> IO ()
withC00 Int
iA Int
jB ((Int -> e -> IO ()) -> IO ()) -> (Int -> e -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
ixC00 e
c00 -> do
                      MArray (PrimState IO) r Ix2 e -> Int -> e -> IO ()
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
unsafeLinearWrite MArray RealWorld r Ix2 e
MArray (PrimState IO) r Ix2 e
marrC Int
ixC00 (e
c00 e -> e -> e
forall a. Num a => a -> a -> a
+ e
a00 e -> e -> e
forall a. Num a => a -> a -> a
* e
b00)
                      Int -> (Int -> e -> IO ()) -> IO ()
withC01 Int
ixC00 ((Int -> e -> IO ()) -> IO ()) -> (Int -> e -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
ixC01 e
c01 -> do
                        MArray (PrimState IO) r Ix2 e -> Int -> e -> IO ()
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
unsafeLinearWrite MArray RealWorld r Ix2 e
MArray (PrimState IO) r Ix2 e
marrC Int
ixC01 (e
c01 e -> e -> e
forall a. Num a => a -> a -> a
+ e
a00 e -> e -> e
forall a. Num a => a -> a -> a
* e
b01)
                        Int -> (Int -> e -> IO ()) -> IO ()
withC10 Int
ixC00 ((Int -> e -> IO ()) -> IO ()) -> (Int -> e -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
ixC10 e
c10 ->
                          MArray (PrimState IO) r Ix2 e -> Int -> e -> IO ()
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
unsafeLinearWrite MArray RealWorld r Ix2 e
MArray (PrimState IO) r Ix2 e
marrC Int
ixC10 (e
c10 e -> e -> e
forall a. Num a => a -> a -> a
+ e
a10 e -> e -> e
forall a. Num a => a -> a -> a
* e
b00)
                        Int -> (Int -> e -> IO ()) -> IO ()
withC11 Int
ixC01 ((Int -> e -> IO ()) -> IO ()) -> (Int -> e -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
ixC11 e
c11 ->
                          MArray (PrimState IO) r Ix2 e -> Int -> e -> IO ()
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
unsafeLinearWrite MArray RealWorld r Ix2 e
MArray (PrimState IO) r Ix2 e
marrC Int
ixC11 (e
c11 e -> e -> e
forall a. Num a => a -> a -> a
+ e
a10 e -> e -> e
forall a. Num a => a -> a -> a
* e
b01)
                  e -> e -> Int -> Int -> Int -> IO ()
loopColsB_RowBColA_UnRowA e
a00 e
a10 Int
iA Int
iB (Int
jB Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
2)
              | Int
jB Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
nB = Int -> Int -> (Int -> e -> IO ()) -> IO ()
withB00 Int
iB Int
jB ((Int -> e -> IO ()) -> IO ()) -> (Int -> e -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
_ e
b00 ->
                  Int -> Int -> (Int -> e -> IO ()) -> IO ()
withC00 Int
iA Int
jB ((Int -> e -> IO ()) -> IO ()) -> (Int -> e -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
ixC00 e
c00 -> do
                    MArray (PrimState IO) r Ix2 e -> Int -> e -> IO ()
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
unsafeLinearWrite MArray RealWorld r Ix2 e
MArray (PrimState IO) r Ix2 e
marrC Int
ixC00 (e
c00 e -> e -> e
forall a. Num a => a -> a -> a
+ e
a00 e -> e -> e
forall a. Num a => a -> a -> a
* e
b00)
                    Int -> (Int -> e -> IO ()) -> IO ()
withC10 Int
ixC00 ((Int -> e -> IO ()) -> IO ()) -> (Int -> e -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
ixC10 e
c10 ->
                      MArray (PrimState IO) r Ix2 e -> Int -> e -> IO ()
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
unsafeLinearWrite MArray RealWorld r Ix2 e
MArray (PrimState IO) r Ix2 e
marrC Int
ixC10 (e
c10 e -> e -> e
forall a. Num a => a -> a -> a
+ e
a10 e -> e -> e
forall a. Num a => a -> a -> a
* e
b00)
              | Bool
otherwise = () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

            loopColsB_RowBColA_RowA :: e -> Int -> Int -> Int -> IO ()
loopColsB_RowBColA_RowA e
a00 Int
iA Int
iB Int
jB
              | Int
jB Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n2B = do
                  Int -> Int -> (Int -> e -> IO ()) -> IO ()
withB00 Int
iB Int
jB ((Int -> e -> IO ()) -> IO ()) -> (Int -> e -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
ixB00 e
b00 -> do
                    let !b01 :: e
b01 = Matrix r e -> Int -> e
forall ix. Index ix => Array r ix e -> Int -> e
forall r e ix. (Source r e, Index ix) => Array r ix e -> Int -> e
unsafeLinearIndex Matrix r e
arrB (Int
ixB00 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
                    Int -> Int -> (Int -> e -> IO ()) -> IO ()
withC00 Int
iA Int
jB ((Int -> e -> IO ()) -> IO ()) -> (Int -> e -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
ixC00 e
c00 -> do
                      MArray (PrimState IO) r Ix2 e -> Int -> e -> IO ()
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
unsafeLinearWrite MArray RealWorld r Ix2 e
MArray (PrimState IO) r Ix2 e
marrC Int
ixC00 (e
c00 e -> e -> e
forall a. Num a => a -> a -> a
+ e
a00 e -> e -> e
forall a. Num a => a -> a -> a
* e
b00)
                      Int -> (Int -> e -> IO ()) -> IO ()
withC01 Int
ixC00 ((Int -> e -> IO ()) -> IO ()) -> (Int -> e -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
ixC01 e
c01 -> do
                        MArray (PrimState IO) r Ix2 e -> Int -> e -> IO ()
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
unsafeLinearWrite MArray RealWorld r Ix2 e
MArray (PrimState IO) r Ix2 e
marrC Int
ixC01 (e
c01 e -> e -> e
forall a. Num a => a -> a -> a
+ e
a00 e -> e -> e
forall a. Num a => a -> a -> a
* e
b01)
                  e -> Int -> Int -> Int -> IO ()
loopColsB_RowBColA_RowA e
a00 Int
iA Int
iB (Int
jB Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
2)
              | Int
jB Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
nB = Int -> Int -> (Int -> e -> IO ()) -> IO ()
withB00 Int
iB Int
jB ((Int -> e -> IO ()) -> IO ()) -> (Int -> e -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
_ e
b00 ->
                  Int -> Int -> (Int -> e -> IO ()) -> IO ()
withC00 Int
iA Int
jB ((Int -> e -> IO ()) -> IO ()) -> (Int -> e -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
ixC00 e
c00 ->
                    MArray (PrimState IO) r Ix2 e -> Int -> e -> IO ()
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
MArray (PrimState m) r ix e -> Int -> e -> m ()
unsafeLinearWrite MArray RealWorld r Ix2 e
MArray (PrimState IO) r Ix2 e
marrC Int
ixC00 (e
c00 e -> e -> e
forall a. Num a => a -> a -> a
+ e
a00 e -> e -> e
forall a. Num a => a -> a -> a
* e
b00)
              | Bool
otherwise = () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

            loopRowsB_UnRowA :: Int -> Int -> IO ()
loopRowsB_UnRowA Int
iA Int
iB
              | Int
iB Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
m2B = do
                  Int -> Int -> (Int -> e -> Int -> e -> IO ()) -> IO ()
withA00A10 Int
iA Int
iB ((Int -> e -> Int -> e -> IO ()) -> IO ())
-> (Int -> e -> Int -> e -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
ixA00 e
a00 Int
ixA10 e
a10 -> do
                    let !a01 :: e
a01 = Matrix r e -> Int -> e
forall ix. Index ix => Array r ix e -> Int -> e
forall r e ix. (Source r e, Index ix) => Array r ix e -> Int -> e
unsafeLinearIndex Matrix r e
arrA (Int
ixA00 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
                        !a11 :: e
a11 = Matrix r e -> Int -> e
forall ix. Index ix => Array r ix e -> Int -> e
forall r e ix. (Source r e, Index ix) => Array r ix e -> Int -> e
unsafeLinearIndex Matrix r e
arrA (Int
ixA10 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
                    e -> e -> e -> e -> Int -> Int -> Int -> IO ()
loopColsB_UnRowBColA_UnRowA e
a00 e
a01 e
a10 e
a11 Int
iA Int
iB Int
0
                  Int -> Int -> IO ()
loopRowsB_UnRowA Int
iA (Int
iB Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
2)
              | Int
iB Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
mB =
                  Int -> Int -> (Int -> e -> Int -> e -> IO ()) -> IO ()
withA00A10 Int
iA Int
iB ((Int -> e -> Int -> e -> IO ()) -> IO ())
-> (Int -> e -> Int -> e -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
_ e
a00 Int
_ e
a10 -> e -> e -> Int -> Int -> Int -> IO ()
loopColsB_RowBColA_UnRowA e
a00 e
a10 Int
iA Int
iB Int
0
              | Bool
otherwise = () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

            loopRowsB_RowA :: Int -> Int -> IO ()
loopRowsB_RowA Int
iA Int
iB
              | Int
iB Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
m2B = do
                  Int -> Int -> (Int -> e -> IO ()) -> IO ()
withA00 Int
iA Int
iB ((Int -> e -> IO ()) -> IO ()) -> (Int -> e -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
ixA00 e
a00 -> do
                    let !a01 :: e
a01 = Matrix r e -> Int -> e
forall ix. Index ix => Array r ix e -> Int -> e
forall r e ix. (Source r e, Index ix) => Array r ix e -> Int -> e
unsafeLinearIndex Matrix r e
arrA (Int
ixA00 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
                    e -> e -> Int -> Int -> Int -> IO ()
loopColsB_UnRowBColA_RowA e
a00 e
a01 Int
iA Int
iB Int
0
                  Int -> Int -> IO ()
loopRowsB_RowA Int
iA (Int
iB Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
2)
              | Int
iB Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
mB = Int -> Int -> (Int -> e -> IO ()) -> IO ()
withA00 Int
iA Int
iB ((Int -> e -> IO ()) -> IO ()) -> (Int -> e -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Int
_ e
a00 -> e -> Int -> Int -> Int -> IO ()
loopColsB_RowBColA_RowA e
a00 Int
iA Int
iB Int
0
              | Bool
otherwise = () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

            loopRowsA :: Int -> IO ()
loopRowsA Int
iA
              | Int
iA Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
m2A = do
                  Scheduler RealWorld () -> IO () -> IO ()
forall s (m :: * -> *).
MonadPrimBase s m =>
Scheduler s () -> m () -> m ()
scheduleWork_ Scheduler RealWorld ()
scheduler (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ Int -> Int -> IO ()
loopRowsB_UnRowA Int
iA Int
0
                  Int -> IO ()
loopRowsA (Int
iA Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
2)
              | Int
iA Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
mA = Scheduler RealWorld () -> IO () -> IO ()
forall s (m :: * -> *).
MonadPrimBase s m =>
Scheduler s () -> m () -> m ()
scheduleWork_ Scheduler RealWorld ()
scheduler (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ Int -> Int -> IO ()
loopRowsB_RowA Int
iA Int
0
              | Bool
otherwise = () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
        Int -> IO ()
loopRowsA Int
0

      Comp -> MArray (PrimState IO) r Ix2 e -> IO (Matrix r e)
forall r e ix (m :: * -> *).
(Manifest r e, Index ix, PrimMonad m) =>
Comp -> MArray (PrimState m) r ix e -> m (Array r ix e)
forall ix (m :: * -> *).
(Index ix, PrimMonad m) =>
Comp -> MArray (PrimState m) r ix e -> m (Array r ix e)
unsafeFreeze Comp
comp MArray RealWorld r Ix2 e
MArray (PrimState IO) r Ix2 e
marrC
  where
    comp :: Comp
comp = Matrix r e -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
forall ix e. Array r ix e -> Comp
getComp Matrix r e
arrA Comp -> Comp -> Comp
forall a. Semigroup a => a -> a -> a
<> Matrix r e -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
forall ix e. Array r ix e -> Comp
getComp Matrix r e
arrB
    m2A :: Int
m2A = Int
mA Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
mA Int -> Int -> Int
forall a. Integral a => a -> a -> a
`rem` Int
2
    m2B :: Int
m2B = Int
mB Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
mB Int -> Int -> Int
forall a. Integral a => a -> a -> a
`rem` Int
2
    n2B :: Int
n2B = Int
nB Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
nB Int -> Int -> Int
forall a. Integral a => a -> a -> a
`rem` Int
2
    Sz (Int
mA :. Int
nA) = Matrix r e -> Sz Ix2
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array r ix e -> Sz ix
size Matrix r e
arrA
    Sz (Int
mB :. Int
nB) = Matrix r e -> Sz Ix2
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array r ix e -> Sz ix
size Matrix r e
arrB
{-# INLINEABLE multiplyMatrices #-}

-- | Computes the matrix-matrix multiplication where second matrix is transposed (i.e. M
-- x N')
--
-- > m1 .><. transpose m2 == multiplyMatricesTransposed m1 m2
--
-- @since 0.5.6
multiplyMatricesTransposed
  :: (Numeric r e, Manifest r e, MonadThrow m)
  => Matrix r e
  -> Matrix r e
  -> m (Matrix D e)
multiplyMatricesTransposed :: forall r e (m :: * -> *).
(Numeric r e, Manifest r e, MonadThrow m) =>
Matrix r e -> Matrix r e -> m (Matrix D e)
multiplyMatricesTransposed Matrix r e
arr1 Matrix r e
arr2
  | Int
n1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
m2 = SizeException -> m (Matrix D e)
forall e a. (HasCallStack, Exception e) => e -> m a
forall (m :: * -> *) e a.
(MonadThrow m, HasCallStack, Exception e) =>
e -> m a
throwM (SizeException -> m (Matrix D e))
-> SizeException -> m (Matrix D e)
forall a b. (a -> b) -> a -> b
$ Sz Ix2 -> Sz Ix2 -> SizeException
forall ix. Index ix => Sz ix -> Sz ix -> SizeException
SizeMismatchException (Matrix r e -> Sz Ix2
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array r ix e -> Sz ix
size Matrix r e
arr1) (Int -> Int -> Sz Ix2
Sz2 Int
m2 Int
n2)
  | Matrix r e -> Bool
forall ix r e. (Index ix, Size r) => Array r ix e -> Bool
isEmpty Matrix r e
arr1 Bool -> Bool -> Bool
|| Matrix r e -> Bool
forall ix r e. (Index ix, Size r) => Array r ix e -> Bool
isEmpty Matrix r e
arr2 = Matrix D e -> m (Matrix D e)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Matrix D e -> m (Matrix D e)) -> Matrix D e -> m (Matrix D e)
forall a b. (a -> b) -> a -> b
$ Comp -> Matrix D e -> Matrix D e
forall r ix e. Strategy r => Comp -> Array r ix e -> Array r ix e
forall ix e. Comp -> Array D ix e -> Array D ix e
setComp Comp
comp Matrix D e
forall r ix e. Load r ix e => Array r ix e
empty
  | Bool
otherwise =
      Matrix D e -> m (Matrix D e)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Matrix D e -> m (Matrix D e)) -> Matrix D e -> m (Matrix D e)
forall a b. (a -> b) -> a -> b
$
        Comp -> Sz Ix2 -> (Ix2 -> e) -> Matrix D e
forall r ix e.
Load r ix e =>
Comp -> Sz ix -> (ix -> e) -> Array r ix e
makeArray Comp
comp (Ix2 -> Sz Ix2
forall ix. ix -> Sz ix
SafeSz (Int
m1 Int -> Int -> Ix2
:. Int
n2)) ((Ix2 -> e) -> Matrix D e) -> (Ix2 -> e) -> Matrix D e
forall a b. (a -> b) -> a -> b
$ \(Int
i :. Int
j) ->
          Array r Int e -> Array r Int e -> e
forall ix. Index ix => Array r ix e -> Array r ix e -> e
forall r e ix.
(FoldNumeric r e, Index ix) =>
Array r ix e -> Array r ix e -> e
unsafeDotProduct (Int -> Sz Int -> Matrix r e -> Array r Int e
forall ix.
Index ix =>
Int -> Sz Int -> Array r ix e -> Array r Int e
forall r e ix.
(Source r e, Index ix) =>
Int -> Sz Int -> Array r ix e -> Array r Int e
unsafeLinearSlice (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
n1) Sz Int
n Matrix r e
arr1) (Int -> Sz Int -> Matrix r e -> Array r Int e
forall ix.
Index ix =>
Int -> Sz Int -> Array r ix e -> Array r Int e
forall r e ix.
(Source r e, Index ix) =>
Int -> Sz Int -> Array r ix e -> Array r Int e
unsafeLinearSlice (Int
j Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
n1) Sz Int
n Matrix r e
arr2)
  where
    comp :: Comp
comp = Matrix r e -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
forall ix e. Array r ix e -> Comp
getComp Matrix r e
arr1 Comp -> Comp -> Comp
forall a. Semigroup a => a -> a -> a
<> Matrix r e -> Comp
forall r ix e. Strategy r => Array r ix e -> Comp
forall ix e. Array r ix e -> Comp
getComp Matrix r e
arr2
    n :: Sz Int
n = Int -> Sz Int
forall ix. ix -> Sz ix
SafeSz Int
n1
    SafeSz (Int
m1 :. Int
n1) = Matrix r e -> Sz Ix2
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array r ix e -> Sz ix
size Matrix r e
arr1
    SafeSz (Int
n2 :. Int
m2) = Matrix r e -> Sz Ix2
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array r ix e -> Sz ix
size Matrix r e
arr2
{-# INLINE multiplyMatricesTransposed #-}

-- | Create an indentity matrix.
--
-- ==== __Example__
--
-- >>> import Data.Massiv.Array
-- >>> identityMatrix 5
-- Array DL Seq (Sz (5 :. 5))
--   [ [ 1, 0, 0, 0, 0 ]
--   , [ 0, 1, 0, 0, 0 ]
--   , [ 0, 0, 1, 0, 0 ]
--   , [ 0, 0, 0, 1, 0 ]
--   , [ 0, 0, 0, 0, 1 ]
--   ]
--
-- @since 0.3.6
identityMatrix :: Num e => Sz1 -> Matrix DL e
identityMatrix :: forall e. Num e => Sz Int -> Matrix DL e
identityMatrix (Sz Int
n) =
  Sz Ix2
-> e
-> (forall {m :: * -> *}. Monad m => (Ix2 -> e -> m Bool) -> m ())
-> Array DL Ix2 e
forall ix e.
Index ix =>
Sz ix
-> e
-> (forall (m :: * -> *). Monad m => (ix -> e -> m Bool) -> m ())
-> Array DL ix e
makeLoadArrayS (Int -> Int -> Sz Ix2
Sz2 Int
n Int
n) e
0 ((forall {m :: * -> *}. Monad m => (Ix2 -> e -> m Bool) -> m ())
 -> Array DL Ix2 e)
-> (forall {m :: * -> *}. Monad m => (Ix2 -> e -> m Bool) -> m ())
-> Array DL Ix2 e
forall a b. (a -> b) -> a -> b
$ \Ix2 -> e -> m Bool
w -> Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> m Bool) -> m ()
forall (f :: * -> *) a.
Applicative f =>
Int -> (Int -> Bool) -> (Int -> Int) -> (Int -> f a) -> f ()
loopA_ Int
0 (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
n) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) ((Int -> m Bool) -> m ()) -> (Int -> m Bool) -> m ()
forall a b. (a -> b) -> a -> b
$ \Int
i -> Ix2 -> e -> m Bool
w (Int
i Int -> Int -> Ix2
:. Int
i) e
1
{-# INLINE identityMatrix #-}

-- | Create a lower triangular (L in LU decomposition) matrix of size @NxN@
--
-- ==== __Example__
--
-- >>> import Data.Massiv.Array
-- >>> lowerTriangular Seq 5 (\(i :. j) -> i + j)
-- Array DL Seq (Sz (5 :. 5))
--   [ [ 0, 0, 0, 0, 0 ]
--   , [ 1, 2, 0, 0, 0 ]
--   , [ 2, 3, 4, 0, 0 ]
--   , [ 3, 4, 5, 6, 0 ]
--   , [ 4, 5, 6, 7, 8 ]
--   ]
--
-- @since 0.5.2
lowerTriangular :: forall e. Num e => Comp -> Sz1 -> (Ix2 -> e) -> Matrix DL e
lowerTriangular :: forall e. Num e => Comp -> Sz Int -> (Ix2 -> e) -> Matrix DL e
lowerTriangular Comp
comp (Sz1 Int
n) Ix2 -> e
f = Comp -> Sz Ix2 -> Loader e -> Array DL Ix2 e
forall ix e. Comp -> Sz ix -> Loader e -> Array DL ix e
DLArray Comp
comp (Ix2 -> Sz Ix2
forall ix. ix -> Sz ix
SafeSz (Int
n Int -> Int -> Ix2
:. Int
n)) Scheduler s ()
-> Int
-> (Int -> e -> ST s ())
-> (Int -> Sz Int -> e -> ST s ())
-> ST s ()
Loader e
load
  where
    load :: Loader e
    load :: Loader e
load Scheduler s ()
scheduler Int
startAt Int -> e -> ST s ()
uWrite Int -> Sz Int -> e -> ST s ()
uSet = do
      Array D Int Int -> (Int -> ST s ()) -> ST s ()
forall r a ix (m :: * -> *) b.
(Source r a, Index ix, Monad m) =>
Array r ix a -> (a -> m b) -> m ()
forM_ (Int
0 Int -> Int -> Array D Int Int
forall ix. Index ix => ix -> ix -> Array D ix ix
..: Int
n) ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
i -> do
        let !k :: Int
k = Int
startAt Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
n
        Scheduler s () -> ST s () -> ST s ()
forall s (m :: * -> *).
MonadPrimBase s m =>
Scheduler s () -> m () -> m ()
scheduleWork_ Scheduler s ()
scheduler (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
          Array D Int Int -> (Int -> ST s ()) -> ST s ()
forall r a ix (m :: * -> *) b.
(Source r a, Index ix, Monad m) =>
Array r ix a -> (a -> m b) -> m ()
forM_ (Int
0 Int -> Int -> Array D Int Int
forall ix. Index ix => ix -> ix -> Array D ix ix
... Int
i) ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
j -> Int -> e -> ST s ()
uWrite (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
j) (Ix2 -> e
f (Int
i Int -> Int -> Ix2
:. Int
j))
          Int -> Sz Int -> e -> ST s ()
uSet (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Int -> Sz Int
forall ix. Index ix => ix -> Sz ix
Sz (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)) e
0
{-# INLINE lowerTriangular #-}

-- | Create an upper triangular (U in LU decomposition) matrix of size @NxN@
--
-- ==== __Example__
--
-- >>> import Data.Massiv.Array
-- >>> upperTriangular Par 5 (\(i :. j) -> i + j)
-- Array DL Par (Sz (5 :. 5))
--   [ [ 0, 1, 2, 3, 4 ]
--   , [ 0, 2, 3, 4, 5 ]
--   , [ 0, 0, 4, 5, 6 ]
--   , [ 0, 0, 0, 6, 7 ]
--   , [ 0, 0, 0, 0, 8 ]
--   ]
--
-- @since 0.5.2
upperTriangular :: forall e. Num e => Comp -> Sz1 -> (Ix2 -> e) -> Matrix DL e
upperTriangular :: forall e. Num e => Comp -> Sz Int -> (Ix2 -> e) -> Matrix DL e
upperTriangular Comp
comp (Sz1 Int
n) Ix2 -> e
f = Comp -> Sz Ix2 -> Loader e -> Array DL Ix2 e
forall ix e. Comp -> Sz ix -> Loader e -> Array DL ix e
DLArray Comp
comp (Ix2 -> Sz Ix2
forall ix. ix -> Sz ix
SafeSz (Int
n Int -> Int -> Ix2
:. Int
n)) Scheduler s ()
-> Int
-> (Int -> e -> ST s ())
-> (Int -> Sz Int -> e -> ST s ())
-> ST s ()
Loader e
load
  where
    load :: Loader e
    load :: Loader e
load Scheduler s ()
scheduler Int
startAt Int -> e -> ST s ()
uWrite Int -> Sz Int -> e -> ST s ()
uSet = do
      Array D Int Int -> (Int -> ST s ()) -> ST s ()
forall r a ix (m :: * -> *) b.
(Source r a, Index ix, Monad m) =>
Array r ix a -> (a -> m b) -> m ()
forM_ (Int
0 Int -> Int -> Array D Int Int
forall ix. Index ix => ix -> ix -> Array D ix ix
..: Int
n) ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
i -> do
        let !k :: Int
k = Int
startAt Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
n
        Scheduler s () -> ST s () -> ST s ()
forall s (m :: * -> *).
MonadPrimBase s m =>
Scheduler s () -> m () -> m ()
scheduleWork_ Scheduler s ()
scheduler (ST s () -> ST s ()) -> ST s () -> ST s ()
forall a b. (a -> b) -> a -> b
$ do
          Int -> Sz Int -> e -> ST s ()
uSet Int
k (Int -> Sz Int
forall ix. ix -> Sz ix
SafeSz Int
i) e
0
          Array D Int Int -> (Int -> ST s ()) -> ST s ()
forall r a ix (m :: * -> *) b.
(Source r a, Index ix, Monad m) =>
Array r ix a -> (a -> m b) -> m ()
forM_ (Int
i Int -> Int -> Array D Int Int
forall ix. Index ix => ix -> ix -> Array D ix ix
..: Int
n) ((Int -> ST s ()) -> ST s ()) -> (Int -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Int
j -> Int -> e -> ST s ()
uWrite (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
j) (Ix2 -> e
f (Int
i Int -> Int -> Ix2
:. Int
j))
{-# INLINE upperTriangular #-}

-- | Negate each element of the array
--
-- @since 0.4.0
negateA :: (Index ix, Numeric r e) => Array r ix e -> Array r ix e
negateA :: forall ix r e.
(Index ix, Numeric r e) =>
Array r ix e -> Array r ix e
negateA = (e -> e) -> Array r ix e -> Array r ix e
forall ix. Index ix => (e -> e) -> Array r ix e -> Array r ix e
forall r e ix.
(Numeric r e, Index ix) =>
(e -> e) -> Array r ix e -> Array r ix e
unsafeLiftArray e -> e
forall a. Num a => a -> a
negate
{-# INLINE negateA #-}

-- | Apply `abs` to each element of the array
--
-- @since 0.4.0
absA :: (Index ix, Numeric r e) => Array r ix e -> Array r ix e
absA :: forall ix r e.
(Index ix, Numeric r e) =>
Array r ix e -> Array r ix e
absA = Array r ix e -> Array r ix e
forall ix. Index ix => Array r ix e -> Array r ix e
forall r e ix.
(Numeric r e, Index ix) =>
Array r ix e -> Array r ix e
absPointwise
{-# INLINE absA #-}

-- | Apply `signum` to each element of the array
--
-- @since 0.4.0
signumA :: (Index ix, Numeric r e) => Array r ix e -> Array r ix e
signumA :: forall ix r e.
(Index ix, Numeric r e) =>
Array r ix e -> Array r ix e
signumA = (e -> e) -> Array r ix e -> Array r ix e
forall ix. Index ix => (e -> e) -> Array r ix e -> Array r ix e
forall r e ix.
(Numeric r e, Index ix) =>
(e -> e) -> Array r ix e -> Array r ix e
unsafeLiftArray e -> e
forall a. Num a => a -> a
signum
{-# INLINE signumA #-}

-- | Divide each element of one array by another pointwise. Same as `!/!` but produces
-- monadic computation that allows for handling failure.
--
-- /__Throws Exception__/: `SizeMismatchException` when array sizes do not match.
--
-- @since 0.4.0
(./.)
  :: (Index ix, NumericFloat r e, MonadThrow m)
  => Array r ix e
  -> Array r ix e
  -> m (Array r ix e)
./. :: forall ix r e (m :: * -> *).
(Index ix, NumericFloat r e, MonadThrow m) =>
Array r ix e -> Array r ix e -> m (Array r ix e)
(./.) = (Array r ix e -> Array r ix e -> Array r ix e)
-> Array r ix e -> Array r ix e -> m (Array r ix e)
forall ix r (m :: * -> *) e.
(Index ix, Size r, MonadThrow m) =>
(Array r ix e -> Array r ix e -> Array r ix e)
-> Array r ix e -> Array r ix e -> m (Array r ix e)
applyExactSize2M Array r ix e -> Array r ix e -> Array r ix e
forall ix. Index ix => Array r ix e -> Array r ix e -> Array r ix e
forall r e ix.
(NumericFloat r e, Index ix) =>
Array r ix e -> Array r ix e -> Array r ix e
divisionPointwise
{-# INLINE (./.) #-}

-- | Divide two arrays pointwise. Prefer to use monadic version of this function `./.`
-- whenever possible, because it is better to avoid partial functions.
--
-- [Partial] Mismatched array sizes will result in an impure exception being thrown.
--
-- ====__Example__
--
-- >>> let arr1 = fromIntegral <$> (Ix1 20 ..: 25) :: Array D Ix1 Float
-- >>> let arr2 = fromIntegral <$> (Ix1 100 ..: 105) :: Array D Ix1 Float
-- >>> arr1 !/! arr2
-- Array D Seq (Sz1 5)
--   [ 0.2, 0.20792079, 0.21568628, 0.22330096, 0.23076923 ]
--
-- @since 0.5.6
(!/!) :: (Index ix, NumericFloat r e) => Array r ix e -> Array r ix e -> Array r ix e
!/! :: forall ix r e.
(Index ix, NumericFloat r e) =>
Array r ix e -> Array r ix e -> Array r ix e
(!/!) Array r ix e
a1 Array r ix e
a2 = Either SomeException (Array r ix e) -> Array r ix e
forall a. HasCallStack => Either SomeException a -> a
throwEither (Array r ix e
a1 Array r ix e -> Array r ix e -> Either SomeException (Array r ix e)
forall ix r e (m :: * -> *).
(Index ix, NumericFloat r e, MonadThrow m) =>
Array r ix e -> Array r ix e -> m (Array r ix e)
./. Array r ix e
a2)
{-# INLINE (!/!) #-}

-- | Divide a scalar value by each element of the array.
--
-- > e /. arr == e *. recipA arr
--
-- ====__Example__
--
-- >>> let arr = fromIntegral <$> (Ix1 20 ..: 25) :: Array D Ix1 Float
-- >>> arr
-- Array D Seq (Sz1 5)
--   [ 20.0, 21.0, 22.0, 23.0, 24.0 ]
-- >>> 100 /. arr
-- Array D Seq (Sz1 5)
--   [ 5.0, 4.7619047, 4.5454545, 4.347826, 4.1666665 ]
--
-- @since 0.5.6
(/.) :: (Index ix, NumericFloat r e) => e -> Array r ix e -> Array r ix e
/. :: forall ix r e.
(Index ix, NumericFloat r e) =>
e -> Array r ix e -> Array r ix e
(/.) = e -> Array r ix e -> Array r ix e
forall ix. Index ix => e -> Array r ix e -> Array r ix e
forall r e ix.
(NumericFloat r e, Index ix) =>
e -> Array r ix e -> Array r ix e
scalarDivide
{-# INLINE (/.) #-}

-- | Divide each element of the array by a scalar value.
--
-- ====__Example__
--
-- >>> let arr = fromIntegral <$> (Ix1 20 ..: 25) :: Array D Ix1 Float
-- >>> arr
-- Array D Seq (Sz1 5)
--   [ 20.0, 21.0, 22.0, 23.0, 24.0 ]
-- >>> arr ./ 100
-- Array D Seq (Sz1 5)
--   [ 0.2, 0.21, 0.22, 0.23, 0.24 ]
--
-- @since 0.4.0
(./) :: (Index ix, NumericFloat r e) => Array r ix e -> e -> Array r ix e
./ :: forall ix r e.
(Index ix, NumericFloat r e) =>
Array r ix e -> e -> Array r ix e
(./) = Array r ix e -> e -> Array r ix e
forall ix. Index ix => Array r ix e -> e -> Array r ix e
forall r e ix.
(NumericFloat r e, Index ix) =>
Array r ix e -> e -> Array r ix e
divideScalar
{-# INLINE (./) #-}

(.^^)
  :: (Index ix, Numeric r e, Fractional e, Integral b)
  => Array r ix e
  -> b
  -> Array r ix e
.^^ :: forall ix r e b.
(Index ix, Numeric r e, Fractional e, Integral b) =>
Array r ix e -> b -> Array r ix e
(.^^) Array r ix e
arr b
n = (e -> e) -> Array r ix e -> Array r ix e
forall ix. Index ix => (e -> e) -> Array r ix e -> Array r ix e
forall r e ix.
(Numeric r e, Index ix) =>
(e -> e) -> Array r ix e -> Array r ix e
unsafeLiftArray (e -> b -> e
forall a b. (Fractional a, Integral b) => a -> b -> a
^^ b
n) Array r ix e
arr
{-# INLINE (.^^) #-}

-- | Apply reciprocal to each element of the array.
--
-- > recipA arr == 1 /. arr
--
-- @since 0.4.0
recipA :: (Index ix, NumericFloat r e) => Array r ix e -> Array r ix e
recipA :: forall ix r e.
(Index ix, NumericFloat r e) =>
Array r ix e -> Array r ix e
recipA = Array r ix e -> Array r ix e
forall ix. Index ix => Array r ix e -> Array r ix e
forall r e ix.
(NumericFloat r e, Index ix) =>
Array r ix e -> Array r ix e
recipPointwise
{-# INLINE recipA #-}

-- | Apply exponent to each element of the array.
--
-- > expA arr == map exp arr
--
-- @since 0.4.0
expA :: (Index ix, NumericFloat r e) => Array r ix e -> Array r ix e
expA :: forall ix r e.
(Index ix, NumericFloat r e) =>
Array r ix e -> Array r ix e
expA = (e -> e) -> Array r ix e -> Array r ix e
forall ix. Index ix => (e -> e) -> Array r ix e -> Array r ix e
forall r e ix.
(Numeric r e, Index ix) =>
(e -> e) -> Array r ix e -> Array r ix e
unsafeLiftArray e -> e
forall a. Floating a => a -> a
exp
{-# INLINE expA #-}

-- | Apply square root to each element of the array.
--
-- > sqrtA arr == map sqrt arr
--
-- @since 0.4.0
sqrtA :: (Index ix, NumericFloat r e) => Array r ix e -> Array r ix e
sqrtA :: forall ix r e.
(Index ix, NumericFloat r e) =>
Array r ix e -> Array r ix e
sqrtA = (e -> e) -> Array r ix e -> Array r ix e
forall ix. Index ix => (e -> e) -> Array r ix e -> Array r ix e
forall r e ix.
(Numeric r e, Index ix) =>
(e -> e) -> Array r ix e -> Array r ix e
unsafeLiftArray e -> e
forall a. Floating a => a -> a
sqrt
{-# INLINE sqrtA #-}

-- | Apply logarithm to each element of the array.
--
-- > logA arr == map log arr
--
-- @since 0.4.0
logA :: (Index ix, NumericFloat r e) => Array r ix e -> Array r ix e
logA :: forall ix r e.
(Index ix, NumericFloat r e) =>
Array r ix e -> Array r ix e
logA = (e -> e) -> Array r ix e -> Array r ix e
forall ix. Index ix => (e -> e) -> Array r ix e -> Array r ix e
forall r e ix.
(Numeric r e, Index ix) =>
(e -> e) -> Array r ix e -> Array r ix e
unsafeLiftArray e -> e
forall a. Floating a => a -> a
log
{-# INLINE logA #-}

-- | Apply logarithm to each element of the array where the base is in the same cell in
-- the second array.
--
-- > logBaseA arr1 arr2 == zipWith logBase arr1 arr2
--
-- [Partial] Throws an error when arrays do not have matching sizes
--
-- @since 0.4.0
logBaseA
  :: (Index ix, Source r1 e, Source r2 e, Floating e)
  => Array r1 ix e
  -> Array r2 ix e
  -> Array D ix e
logBaseA :: forall ix r1 e r2.
(Index ix, Source r1 e, Source r2 e, Floating e) =>
Array r1 ix e -> Array r2 ix e -> Array D ix e
logBaseA = (e -> e -> e) -> Array r1 ix e -> Array r2 ix e -> Array D ix e
forall ix r1 a r2 b e.
(HasCallStack, Index ix, Source r1 a, Source r2 b) =>
(a -> b -> e) -> Array r1 ix a -> Array r2 ix b -> Array D ix e
liftArray2' e -> e -> e
forall a. Floating a => a -> a -> a
logBase
{-# INLINE logBaseA #-}

-- TODO: siwtch to
-- (breaking) logBaseA :: Array r ix e -> e -> Array D ix e
-- logBasesM :: Array r ix e -> Array r ix e -> m (Array D ix e)

-- | Apply power to each element of the array where the power value is in the same cell
-- in the second array.
--
-- > arr1 .** arr2 == zipWith (**) arr1 arr2
--
-- [Partial] Throws an error when arrays do not have matching sizes
--
-- @since 0.4.0
(.**)
  :: (Index ix, Source r1 e, Source r2 e, Floating e)
  => Array r1 ix e
  -> Array r2 ix e
  -> Array D ix e
.** :: forall ix r1 e r2.
(Index ix, Source r1 e, Source r2 e, Floating e) =>
Array r1 ix e -> Array r2 ix e -> Array D ix e
(.**) = (e -> e -> e) -> Array r1 ix e -> Array r2 ix e -> Array D ix e
forall ix r1 a r2 b e.
(HasCallStack, Index ix, Source r1 a, Source r2 b) =>
(a -> b -> e) -> Array r1 ix a -> Array r2 ix b -> Array D ix e
liftArray2' e -> e -> e
forall a. Floating a => a -> a -> a
(**)
{-# INLINE (.**) #-}

-- TODO:
-- !**! :: Array r1 ix e -> Array r2 ix e -> Array D ix e
-- .**. :: Array r1 ix e -> Array r2 ix e -> m (Array D ix e)
-- (breaking) .** :: Array r1 ix e -> e -> Array D ix e

-- | Apply sine function to each element of the array.
--
-- > sinA arr == map sin arr
--
-- @since 0.4.0
sinA :: (Index ix, NumericFloat r e) => Array r ix e -> Array r ix e
sinA :: forall ix r e.
(Index ix, NumericFloat r e) =>
Array r ix e -> Array r ix e
sinA = (e -> e) -> Array r ix e -> Array r ix e
forall ix. Index ix => (e -> e) -> Array r ix e -> Array r ix e
forall r e ix.
(Numeric r e, Index ix) =>
(e -> e) -> Array r ix e -> Array r ix e
unsafeLiftArray e -> e
forall a. Floating a => a -> a
sin
{-# INLINE sinA #-}

-- | Apply cosine function to each element of the array.
--
-- > cosA arr == map cos arr
--
-- @since 0.4.0
cosA :: (Index ix, NumericFloat r e) => Array r ix e -> Array r ix e
cosA :: forall ix r e.
(Index ix, NumericFloat r e) =>
Array r ix e -> Array r ix e
cosA = (e -> e) -> Array r ix e -> Array r ix e
forall ix. Index ix => (e -> e) -> Array r ix e -> Array r ix e
forall r e ix.
(Numeric r e, Index ix) =>
(e -> e) -> Array r ix e -> Array r ix e
unsafeLiftArray e -> e
forall a. Floating a => a -> a
cos
{-# INLINE cosA #-}

-- | Apply tangent function to each element of the array.
--
-- > tanA arr == map tan arr
--
-- @since 0.4.0
tanA :: (Index ix, NumericFloat r e) => Array r ix e -> Array r ix e
tanA :: forall ix r e.
(Index ix, NumericFloat r e) =>
Array r ix e -> Array r ix e
tanA = (e -> e) -> Array r ix e -> Array r ix e
forall ix. Index ix => (e -> e) -> Array r ix e -> Array r ix e
forall r e ix.
(Numeric r e, Index ix) =>
(e -> e) -> Array r ix e -> Array r ix e
unsafeLiftArray e -> e
forall a. Floating a => a -> a
tan
{-# INLINE tanA #-}

-- | Apply arcsine function to each element of the array.
--
-- > asinA arr == map asin arr
--
-- @since 0.4.0
asinA :: (Index ix, NumericFloat r e) => Array r ix e -> Array r ix e
asinA :: forall ix r e.
(Index ix, NumericFloat r e) =>
Array r ix e -> Array r ix e
asinA = (e -> e) -> Array r ix e -> Array r ix e
forall ix. Index ix => (e -> e) -> Array r ix e -> Array r ix e
forall r e ix.
(Numeric r e, Index ix) =>
(e -> e) -> Array r ix e -> Array r ix e
unsafeLiftArray e -> e
forall a. Floating a => a -> a
asin
{-# INLINE asinA #-}

-- | Apply arctangent function to each element of the array.
--
-- > atanA arr == map atan arr
--
-- @since 0.4.0
atanA :: (Index ix, NumericFloat r e) => Array r ix e -> Array r ix e
atanA :: forall ix r e.
(Index ix, NumericFloat r e) =>
Array r ix e -> Array r ix e
atanA = (e -> e) -> Array r ix e -> Array r ix e
forall ix. Index ix => (e -> e) -> Array r ix e -> Array r ix e
forall r e ix.
(Numeric r e, Index ix) =>
(e -> e) -> Array r ix e -> Array r ix e
unsafeLiftArray e -> e
forall a. Floating a => a -> a
atan
{-# INLINE atanA #-}

-- | Apply arccosine function to each element of the array.
--
-- > acosA arr == map acos arr
--
-- @since 0.4.0
acosA :: (Index ix, NumericFloat r e) => Array r ix e -> Array r ix e
acosA :: forall ix r e.
(Index ix, NumericFloat r e) =>
Array r ix e -> Array r ix e
acosA = (e -> e) -> Array r ix e -> Array r ix e
forall ix. Index ix => (e -> e) -> Array r ix e -> Array r ix e
forall r e ix.
(Numeric r e, Index ix) =>
(e -> e) -> Array r ix e -> Array r ix e
unsafeLiftArray e -> e
forall a. Floating a => a -> a
acos
{-# INLINE acosA #-}

-- | Apply hyperbolic sine function to each element of the array.
--
-- > sinhA arr == map sinh arr
--
-- @since 0.4.0
sinhA :: (Index ix, NumericFloat r e) => Array r ix e -> Array r ix e
sinhA :: forall ix r e.
(Index ix, NumericFloat r e) =>
Array r ix e -> Array r ix e
sinhA = (e -> e) -> Array r ix e -> Array r ix e
forall ix. Index ix => (e -> e) -> Array r ix e -> Array r ix e
forall r e ix.
(Numeric r e, Index ix) =>
(e -> e) -> Array r ix e -> Array r ix e
unsafeLiftArray e -> e
forall a. Floating a => a -> a
sinh
{-# INLINE sinhA #-}

-- | Apply hyperbolic tangent function to each element of the array.
--
-- > tanhA arr == map tanh arr
--
-- @since 0.4.0
tanhA :: (Index ix, NumericFloat r e) => Array r ix e -> Array r ix e
tanhA :: forall ix r e.
(Index ix, NumericFloat r e) =>
Array r ix e -> Array r ix e
tanhA = (e -> e) -> Array r ix e -> Array r ix e
forall ix. Index ix => (e -> e) -> Array r ix e -> Array r ix e
forall r e ix.
(Numeric r e, Index ix) =>
(e -> e) -> Array r ix e -> Array r ix e
unsafeLiftArray e -> e
forall a. Floating a => a -> a
tanh
{-# INLINE tanhA #-}

-- | Apply hyperbolic cosine function to each element of the array.
--
-- > coshA arr == map cosh arr
--
-- @since 0.4.0
coshA :: (Index ix, NumericFloat r e) => Array r ix e -> Array r ix e
coshA :: forall ix r e.
(Index ix, NumericFloat r e) =>
Array r ix e -> Array r ix e
coshA = (e -> e) -> Array r ix e -> Array r ix e
forall ix. Index ix => (e -> e) -> Array r ix e -> Array r ix e
forall r e ix.
(Numeric r e, Index ix) =>
(e -> e) -> Array r ix e -> Array r ix e
unsafeLiftArray e -> e
forall a. Floating a => a -> a
cosh
{-# INLINE coshA #-}

-- | Apply inverse hyperbolic sine function to each element of the array.
--
-- > asinhA arr == map asinh arr
--
-- @since 0.4.0
asinhA :: (Index ix, NumericFloat r e) => Array r ix e -> Array r ix e
asinhA :: forall ix r e.
(Index ix, NumericFloat r e) =>
Array r ix e -> Array r ix e
asinhA = (e -> e) -> Array r ix e -> Array r ix e
forall ix. Index ix => (e -> e) -> Array r ix e -> Array r ix e
forall r e ix.
(Numeric r e, Index ix) =>
(e -> e) -> Array r ix e -> Array r ix e
unsafeLiftArray e -> e
forall a. Floating a => a -> a
asinh
{-# INLINE asinhA #-}

-- | Apply inverse hyperbolic cosine function to each element of the array.
--
-- > acoshA arr == map acosh arr
--
-- @since 0.4.0
acoshA :: (Index ix, NumericFloat r e) => Array r ix e -> Array r ix e
acoshA :: forall ix r e.
(Index ix, NumericFloat r e) =>
Array r ix e -> Array r ix e
acoshA = (e -> e) -> Array r ix e -> Array r ix e
forall ix. Index ix => (e -> e) -> Array r ix e -> Array r ix e
forall r e ix.
(Numeric r e, Index ix) =>
(e -> e) -> Array r ix e -> Array r ix e
unsafeLiftArray e -> e
forall a. Floating a => a -> a
acosh
{-# INLINE acoshA #-}

-- | Apply inverse hyperbolic tangent function to each element of the array.
--
-- > atanhA arr == map atanh arr
--
-- @since 0.4.0
atanhA :: (Index ix, NumericFloat r e) => Array r ix e -> Array r ix e
atanhA :: forall ix r e.
(Index ix, NumericFloat r e) =>
Array r ix e -> Array r ix e
atanhA = (e -> e) -> Array r ix e -> Array r ix e
forall ix. Index ix => (e -> e) -> Array r ix e -> Array r ix e
forall r e ix.
(Numeric r e, Index ix) =>
(e -> e) -> Array r ix e -> Array r ix e
unsafeLiftArray e -> e
forall a. Floating a => a -> a
atanh
{-# INLINE atanhA #-}

-- | Perform a pointwise quotient where first array contains numerators and the second
-- one denominators
--
-- > quotA arr1 arr2 == zipWith quot arr1 arr2
--
-- [Partial] Mismatched array sizes will result in an impure exception being thrown.
--
-- @since 0.1.0
quotA
  :: (HasCallStack, Index ix, Source r1 e, Source r2 e, Integral e)
  => Array r1 ix e
  -> Array r2 ix e
  -> Array D ix e
quotA :: forall ix r1 e r2.
(HasCallStack, Index ix, Source r1 e, Source r2 e, Integral e) =>
Array r1 ix e -> Array r2 ix e -> Array D ix e
quotA = (e -> e -> e) -> Array r1 ix e -> Array r2 ix e -> Array D ix e
forall ix r1 a r2 b e.
(HasCallStack, Index ix, Source r1 a, Source r2 b) =>
(a -> b -> e) -> Array r1 ix a -> Array r2 ix b -> Array D ix e
liftArray2' e -> e -> e
forall a. Integral a => a -> a -> a
quot
{-# INLINE quotA #-}

-- | Perform a pointwise remainder computation
--
-- > remA arr1 arr2 == zipWith rem arr1 arr2
--
-- [Partial] Mismatched array sizes will result in an impure exception being thrown.
--
-- @since 0.1.0
remA
  :: (HasCallStack, Index ix, Source r1 e, Source r2 e, Integral e)
  => Array r1 ix e
  -> Array r2 ix e
  -> Array D ix e
remA :: forall ix r1 e r2.
(HasCallStack, Index ix, Source r1 e, Source r2 e, Integral e) =>
Array r1 ix e -> Array r2 ix e -> Array D ix e
remA = (e -> e -> e) -> Array r1 ix e -> Array r2 ix e -> Array D ix e
forall ix r1 a r2 b e.
(HasCallStack, Index ix, Source r1 a, Source r2 b) =>
(a -> b -> e) -> Array r1 ix a -> Array r2 ix b -> Array D ix e
liftArray2' e -> e -> e
forall a. Integral a => a -> a -> a
rem
{-# INLINE remA #-}

-- | Perform a pointwise integer division where first array contains numerators and the
-- second one denominators
--
-- > divA arr1 arr2 == zipWith div arr1 arr2
--
-- [Partial] Mismatched array sizes will result in an impure exception being thrown.
--
-- @since 0.1.0
divA
  :: (HasCallStack, Index ix, Source r1 e, Source r2 e, Integral e)
  => Array r1 ix e
  -> Array r2 ix e
  -> Array D ix e
divA :: forall ix r1 e r2.
(HasCallStack, Index ix, Source r1 e, Source r2 e, Integral e) =>
Array r1 ix e -> Array r2 ix e -> Array D ix e
divA = (e -> e -> e) -> Array r1 ix e -> Array r2 ix e -> Array D ix e
forall ix r1 a r2 b e.
(HasCallStack, Index ix, Source r1 a, Source r2 b) =>
(a -> b -> e) -> Array r1 ix a -> Array r2 ix b -> Array D ix e
liftArray2' e -> e -> e
forall a. Integral a => a -> a -> a
div
{-# INLINE divA #-}

-- TODO:
--  * Array r ix e -> Array r ix e -> m (Array r ix e)
--  * Array r ix e -> e -> Array r ix e
--  * e -> Array r ix e -> Array r ix e

-- | Perform a pointwise modulo computation
--
-- > modA arr1 arr2 == zipWith mod arr1 arr2
--
-- [Partial] Mismatched array sizes will result in an impure exception being thrown.
--
-- @since 0.1.0
modA
  :: (HasCallStack, Index ix, Source r1 e, Source r2 e, Integral e)
  => Array r1 ix e
  -> Array r2 ix e
  -> Array D ix e
modA :: forall ix r1 e r2.
(HasCallStack, Index ix, Source r1 e, Source r2 e, Integral e) =>
Array r1 ix e -> Array r2 ix e -> Array D ix e
modA = (e -> e -> e) -> Array r1 ix e -> Array r2 ix e -> Array D ix e
forall ix r1 a r2 b e.
(HasCallStack, Index ix, Source r1 a, Source r2 b) =>
(a -> b -> e) -> Array r1 ix a -> Array r2 ix b -> Array D ix e
liftArray2' e -> e -> e
forall a. Integral a => a -> a -> a
mod
{-# INLINE modA #-}

-- | Perform a pointwise quotient with remainder where first array contains numerators
-- and the second one denominators
--
-- > quotRemA arr1 arr2 == zipWith quotRem arr1 arr2
--
-- [Partial] Mismatched array sizes will result in an impure exception being thrown.
--
-- @since 0.1.0
quotRemA
  :: (HasCallStack, Index ix, Source r1 e, Source r2 e, Integral e)
  => Array r1 ix e
  -> Array r2 ix e
  -> (Array D ix e, Array D ix e)
quotRemA :: forall ix r1 e r2.
(HasCallStack, Index ix, Source r1 e, Source r2 e, Integral e) =>
Array r1 ix e -> Array r2 ix e -> (Array D ix e, Array D ix e)
quotRemA Array r1 ix e
arr1 = Array D ix (e, e) -> (Array D ix e, Array D ix e)
forall ix r e1 e2.
(Index ix, Source r (e1, e2)) =>
Array r ix (e1, e2) -> (Array D ix e1, Array D ix e2)
A.unzip (Array D ix (e, e) -> (Array D ix e, Array D ix e))
-> (Array r2 ix e -> Array D ix (e, e))
-> Array r2 ix e
-> (Array D ix e, Array D ix e)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (e -> e -> (e, e))
-> Array r1 ix e -> Array r2 ix e -> Array D ix (e, e)
forall ix r1 a r2 b e.
(HasCallStack, Index ix, Source r1 a, Source r2 b) =>
(a -> b -> e) -> Array r1 ix a -> Array r2 ix b -> Array D ix e
liftArray2' e -> e -> (e, e)
forall a. Integral a => a -> a -> (a, a)
quotRem Array r1 ix e
arr1
{-# INLINE quotRemA #-}

-- | Perform a pointwise integer division with modulo where first array contains
-- numerators and the second one denominators
--
-- > divModA arr1 arr2 == zipWith divMod arr1 arr2
--
-- [Partial] Mismatched array sizes will result in an impure exception being thrown.
--
-- @since 0.1.0
divModA
  :: (HasCallStack, Index ix, Source r1 e, Source r2 e, Integral e)
  => Array r1 ix e
  -> Array r2 ix e
  -> (Array D ix e, Array D ix e)
divModA :: forall ix r1 e r2.
(HasCallStack, Index ix, Source r1 e, Source r2 e, Integral e) =>
Array r1 ix e -> Array r2 ix e -> (Array D ix e, Array D ix e)
divModA Array r1 ix e
arr1 = Array D ix (e, e) -> (Array D ix e, Array D ix e)
forall ix r e1 e2.
(Index ix, Source r (e1, e2)) =>
Array r ix (e1, e2) -> (Array D ix e1, Array D ix e2)
A.unzip (Array D ix (e, e) -> (Array D ix e, Array D ix e))
-> (Array r2 ix e -> Array D ix (e, e))
-> Array r2 ix e
-> (Array D ix e, Array D ix e)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (e -> e -> (e, e))
-> Array r1 ix e -> Array r2 ix e -> Array D ix (e, e)
forall ix r1 a r2 b e.
(HasCallStack, Index ix, Source r1 a, Source r2 b) =>
(a -> b -> e) -> Array r1 ix a -> Array r2 ix b -> Array D ix e
liftArray2' e -> e -> (e, e)
forall a. Integral a => a -> a -> (a, a)
divMod Array r1 ix e
arr1
{-# INLINE divModA #-}

-- | Truncate each element of the array.
--
-- > truncateA arr == map truncate arr
--
-- @since 0.1.0
truncateA :: (Index ix, Source r a, RealFrac a, Integral e) => Array r ix a -> Array D ix e
truncateA :: forall ix r a e.
(Index ix, Source r a, RealFrac a, Integral e) =>
Array r ix a -> Array D ix e
truncateA = (a -> e) -> Array r ix a -> Array D ix e
forall ix r e' e.
(Index ix, Source r e') =>
(e' -> e) -> Array r ix e' -> Array D ix e
A.map a -> e
forall b. Integral b => a -> b
forall a b. (RealFrac a, Integral b) => a -> b
truncate
{-# INLINE truncateA #-}

-- | Round each element of the array.
--
-- > truncateA arr == map truncate arr
--
-- @since 0.1.0
roundA :: (Index ix, Source r a, RealFrac a, Integral e) => Array r ix a -> Array D ix e
roundA :: forall ix r a e.
(Index ix, Source r a, RealFrac a, Integral e) =>
Array r ix a -> Array D ix e
roundA = (a -> e) -> Array r ix a -> Array D ix e
forall ix r e' e.
(Index ix, Source r e') =>
(e' -> e) -> Array r ix e' -> Array D ix e
A.map a -> e
forall b. Integral b => a -> b
forall a b. (RealFrac a, Integral b) => a -> b
round
{-# INLINE roundA #-}

-- | Ceiling of each element of the array.
--
-- > truncateA arr == map truncate arr
--
-- @since 0.1.0
ceilingA :: (Index ix, Source r a, RealFrac a, Integral e) => Array r ix a -> Array D ix e
ceilingA :: forall ix r a e.
(Index ix, Source r a, RealFrac a, Integral e) =>
Array r ix a -> Array D ix e
ceilingA = (a -> e) -> Array r ix a -> Array D ix e
forall ix r e' e.
(Index ix, Source r e') =>
(e' -> e) -> Array r ix e' -> Array D ix e
A.map a -> e
forall b. Integral b => a -> b
forall a b. (RealFrac a, Integral b) => a -> b
ceiling
{-# INLINE ceilingA #-}

-- | Floor each element of the array.
--
-- > truncateA arr == map truncate arr
--
-- @since 0.1.0
floorA :: (Index ix, Source r a, RealFrac a, Integral e) => Array r ix a -> Array D ix e
floorA :: forall ix r a e.
(Index ix, Source r a, RealFrac a, Integral e) =>
Array r ix a -> Array D ix e
floorA = (a -> e) -> Array r ix a -> Array D ix e
forall ix r e' e.
(Index ix, Source r e') =>
(e' -> e) -> Array r ix e' -> Array D ix e
A.map a -> e
forall b. Integral b => a -> b
forall a b. (RealFrac a, Integral b) => a -> b
floor
{-# INLINE floorA #-}

-- | Perform atan2 pointwise
--
-- > atan2A arr1 arr2 == zipWith atan2 arr1 arr2
--
-- /__Throws Exception__/: `SizeMismatchException` when array sizes do not match.
--
-- @since 0.1.0
atan2A
  :: (Index ix, Numeric r e, RealFloat e, MonadThrow m)
  => Array r ix e
  -> Array r ix e
  -> m (Array r ix e)
atan2A :: forall ix r e (m :: * -> *).
(Index ix, Numeric r e, RealFloat e, MonadThrow m) =>
Array r ix e -> Array r ix e -> m (Array r ix e)
atan2A = (e -> e -> e) -> Array r ix e -> Array r ix e -> m (Array r ix e)
forall ix r e (m :: * -> *).
(Index ix, Numeric r e, MonadThrow m) =>
(e -> e -> e) -> Array r ix e -> Array r ix e -> m (Array r ix e)
liftNumArray2M e -> e -> e
forall a. RealFloat a => a -> a -> a
atan2
{-# INLINE atan2A #-}

-- | Same as `sumArraysM`, compute sum of arrays pointwise. All arrays must have the same
-- size, otherwise it will result in an error.
--
-- @since 1.0.0
sumArrays' :: (HasCallStack, Foldable t, Load r ix e, Numeric r e) => t (Array r ix e) -> Array r ix e
sumArrays' :: forall (t :: * -> *) r ix e.
(HasCallStack, Foldable t, Load r ix e, Numeric r e) =>
t (Array r ix e) -> Array r ix e
sumArrays' = Either SomeException (Array r ix e) -> Array r ix e
forall a. HasCallStack => Either SomeException a -> a
throwEither (Either SomeException (Array r ix e) -> Array r ix e)
-> (t (Array r ix e) -> Either SomeException (Array r ix e))
-> t (Array r ix e)
-> Array r ix e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. t (Array r ix e) -> Either SomeException (Array r ix e)
forall (t :: * -> *) r ix e (m :: * -> *).
(Foldable t, Load r ix e, Numeric r e, MonadThrow m) =>
t (Array r ix e) -> m (Array r ix e)
sumArraysM
{-# INLINE sumArrays' #-}

-- | Compute sum of arrays pointwise. All arrays must have the same size.
--
-- ====__Examples__
--
-- >>> import Data.Massiv.Array as A
-- >>> sumArraysM [] :: IO (Array P Ix3 Int)
-- Array P Seq (Sz (0 :> 0 :. 0))
--   [  ]
-- >>> arr = A.makeArrayR P Seq (Sz3 4 5 6) $ \(i :> j :. k) -> i + j * k
-- >>> arr
-- Array P Seq (Sz (4 :> 5 :. 6))
--   [ [ [ 0, 0, 0, 0, 0, 0 ]
--     , [ 0, 1, 2, 3, 4, 5 ]
--     , [ 0, 2, 4, 6, 8, 10 ]
--     , [ 0, 3, 6, 9, 12, 15 ]
--     , [ 0, 4, 8, 12, 16, 20 ]
--     ]
--   , [ [ 1, 1, 1, 1, 1, 1 ]
--     , [ 1, 2, 3, 4, 5, 6 ]
--     , [ 1, 3, 5, 7, 9, 11 ]
--     , [ 1, 4, 7, 10, 13, 16 ]
--     , [ 1, 5, 9, 13, 17, 21 ]
--     ]
--   , [ [ 2, 2, 2, 2, 2, 2 ]
--     , [ 2, 3, 4, 5, 6, 7 ]
--     , [ 2, 4, 6, 8, 10, 12 ]
--     , [ 2, 5, 8, 11, 14, 17 ]
--     , [ 2, 6, 10, 14, 18, 22 ]
--     ]
--   , [ [ 3, 3, 3, 3, 3, 3 ]
--     , [ 3, 4, 5, 6, 7, 8 ]
--     , [ 3, 5, 7, 9, 11, 13 ]
--     , [ 3, 6, 9, 12, 15, 18 ]
--     , [ 3, 7, 11, 15, 19, 23 ]
--     ]
--   ]
-- >>> sumArraysM $ outerSlices arr
-- Array P Seq (Sz (5 :. 6))
--   [ [ 6, 6, 6, 6, 6, 6 ]
--   , [ 6, 10, 14, 18, 22, 26 ]
--   , [ 6, 14, 22, 30, 38, 46 ]
--   , [ 6, 18, 30, 42, 54, 66 ]
--   , [ 6, 22, 38, 54, 70, 86 ]
--   ]
-- >>> sumArraysM $ innerSlices arr
-- Array D Seq (Sz (4 :. 5))
--   [ [ 0, 15, 30, 45, 60 ]
--   , [ 6, 21, 36, 51, 66 ]
--   , [ 12, 27, 42, 57, 72 ]
--   , [ 18, 33, 48, 63, 78 ]
--   ]
--
-- @since 1.0.0
sumArraysM
  :: (Foldable t, Load r ix e, Numeric r e, MonadThrow m) => t (Array r ix e) -> m (Array r ix e)
sumArraysM :: forall (t :: * -> *) r ix e (m :: * -> *).
(Foldable t, Load r ix e, Numeric r e, MonadThrow m) =>
t (Array r ix e) -> m (Array r ix e)
sumArraysM t (Array r ix e)
as =
  case t (Array r ix e) -> [Array r ix e]
forall a. t a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
F.toList t (Array r ix e)
as of
    [] -> Array r ix e -> m (Array r ix e)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Array r ix e
forall r ix e. Load r ix e => Array r ix e
empty
    (Array r ix e
x : [Array r ix e]
xs) -> (Array r ix e -> Array r ix e -> m (Array r ix e))
-> Array r ix e -> [Array r ix e] -> m (Array r ix e)
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
F.foldlM Array r ix e -> Array r ix e -> m (Array r ix e)
forall ix r e (m :: * -> *).
(Index ix, Numeric r e, MonadThrow m) =>
Array r ix e -> Array r ix e -> m (Array r ix e)
(.+.) Array r ix e
x [Array r ix e]
xs
{-# INLINE sumArraysM #-}

-- OPTIMIZE: Allocate a single result array and write sums into it incrementally.

-- | Same as `productArraysM`. Compute product of arrays pointwise. All arrays must have
-- the same size, otherwise it
-- will result in an error.
--
-- @since 1.0.0
productArrays'
  :: (HasCallStack, Foldable t, Load r ix e, Numeric r e) => t (Array r ix e) -> Array r ix e
productArrays' :: forall (t :: * -> *) r ix e.
(HasCallStack, Foldable t, Load r ix e, Numeric r e) =>
t (Array r ix e) -> Array r ix e
productArrays' = Either SomeException (Array r ix e) -> Array r ix e
forall a. HasCallStack => Either SomeException a -> a
throwEither (Either SomeException (Array r ix e) -> Array r ix e)
-> (t (Array r ix e) -> Either SomeException (Array r ix e))
-> t (Array r ix e)
-> Array r ix e
forall b c a. (b -> c) -> (a -> b) -> a -> c
. t (Array r ix e) -> Either SomeException (Array r ix e)
forall (t :: * -> *) r ix e (m :: * -> *).
(Foldable t, Load r ix e, Numeric r e, MonadThrow m) =>
t (Array r ix e) -> m (Array r ix e)
productArraysM
{-# INLINE productArrays' #-}

-- | Compute product of arrays pointwise. All arrays must have the same size.
--
-- ====__Examples__
--
-- >>> import Data.Massiv.Array as A
-- >>> productArraysM [] :: IO (Array P Ix3 Int)
-- Array P Seq (Sz (0 :> 0 :. 0))
--   [  ]
-- >>> arr = A.makeArrayR P Seq (Sz3 4 5 6) $ \(i :> j :. k) -> i + j * k
-- >>> arr
-- Array P Seq (Sz (4 :> 5 :. 6))
--   [ [ [ 0, 0, 0, 0, 0, 0 ]
--     , [ 0, 1, 2, 3, 4, 5 ]
--     , [ 0, 2, 4, 6, 8, 10 ]
--     , [ 0, 3, 6, 9, 12, 15 ]
--     , [ 0, 4, 8, 12, 16, 20 ]
--     ]
--   , [ [ 1, 1, 1, 1, 1, 1 ]
--     , [ 1, 2, 3, 4, 5, 6 ]
--     , [ 1, 3, 5, 7, 9, 11 ]
--     , [ 1, 4, 7, 10, 13, 16 ]
--     , [ 1, 5, 9, 13, 17, 21 ]
--     ]
--   , [ [ 2, 2, 2, 2, 2, 2 ]
--     , [ 2, 3, 4, 5, 6, 7 ]
--     , [ 2, 4, 6, 8, 10, 12 ]
--     , [ 2, 5, 8, 11, 14, 17 ]
--     , [ 2, 6, 10, 14, 18, 22 ]
--     ]
--   , [ [ 3, 3, 3, 3, 3, 3 ]
--     , [ 3, 4, 5, 6, 7, 8 ]
--     , [ 3, 5, 7, 9, 11, 13 ]
--     , [ 3, 6, 9, 12, 15, 18 ]
--     , [ 3, 7, 11, 15, 19, 23 ]
--     ]
--   ]
-- >>> productArraysM $ outerSlices arr
-- Array P Seq (Sz (5 :. 6))
--   [ [ 0, 0, 0, 0, 0, 0 ]
--   , [ 0, 24, 120, 360, 840, 1680 ]
--   , [ 0, 120, 840, 3024, 7920, 17160 ]
--   , [ 0, 360, 3024, 11880, 32760, 73440 ]
--   , [ 0, 840, 7920, 32760, 93024, 212520 ]
--   ]
-- >>> productArraysM $ innerSlices arr
-- Array D Seq (Sz (4 :. 5))
--   [ [ 0, 0, 0, 0, 0 ]
--   , [ 1, 720, 10395, 58240, 208845 ]
--   , [ 64, 5040, 46080, 209440, 665280 ]
--   , [ 729, 20160, 135135, 524880, 1514205 ]
--   ]
--
-- @since 1.0.0
productArraysM
  :: (Foldable t, Load r ix e, Numeric r e, MonadThrow m) => t (Array r ix e) -> m (Array r ix e)
productArraysM :: forall (t :: * -> *) r ix e (m :: * -> *).
(Foldable t, Load r ix e, Numeric r e, MonadThrow m) =>
t (Array r ix e) -> m (Array r ix e)
productArraysM t (Array r ix e)
as =
  case t (Array r ix e) -> [Array r ix e]
forall a. t a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
F.toList t (Array r ix e)
as of
    [] -> Array r ix e -> m (Array r ix e)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Array r ix e
forall r ix e. Load r ix e => Array r ix e
empty
    (Array r ix e
x : [Array r ix e]
xs) -> (Array r ix e -> Array r ix e -> m (Array r ix e))
-> Array r ix e -> [Array r ix e] -> m (Array r ix e)
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
F.foldlM Array r ix e -> Array r ix e -> m (Array r ix e)
forall ix r e (m :: * -> *).
(Index ix, Numeric r e, MonadThrow m) =>
Array r ix e -> Array r ix e -> m (Array r ix e)
(.*.) Array r ix e
x [Array r ix e]
xs
{-# INLINE productArraysM #-}