130 lines
3.0 KiB
Haskell
130 lines
3.0 KiB
Haskell
{-# OPTIONS_GHC -fno-warn-duplicate-constraints #-}
|
|
{-# LANGUAGE RankNTypes #-}
|
|
|
|
import Data.Array (Array, Ix)
|
|
import Data.Array.Base
|
|
|
|
-- | Element-wise combine the values of two arrays 'a' and 'b' with 'f'.
|
|
-- 'a' and 'b' must have the same bounds.
|
|
zipWithA :: (IArray arr a, IArray arr b, IArray arr c, Ix i) =>
|
|
(a -> b -> c) -> arr i a -> arr i b -> arr i c
|
|
zipWithA f a b =
|
|
case bounds a of
|
|
ba ->
|
|
if ba /= bounds b
|
|
then error "elemwise: bounds mismatch"
|
|
else
|
|
let n = numElements a
|
|
in unsafeArray ba [ (i, f (unsafeAt a i) (unsafeAt b i))
|
|
| i <- [0 .. n - 1]]
|
|
|
|
-- Convenient aliases for matrix-matrix element-wise operations.
|
|
type ElemOp a b c = (IArray arr a, IArray arr b, IArray arr c, Ix i) =>
|
|
arr i a -> arr i b -> arr i c
|
|
type ElemOp1 a = ElemOp a a a
|
|
|
|
infixl 6 +:, -:
|
|
infixl 7 *:, /:, `divE`
|
|
|
|
(+:), (-:), (*:) :: (Num a) => ElemOp1 a
|
|
(+:) = zipWithA (+)
|
|
(-:) = zipWithA (-)
|
|
(*:) = zipWithA (*)
|
|
|
|
divE :: (Integral a) => ElemOp1 a
|
|
divE = zipWithA div
|
|
|
|
(/:) :: (Fractional a) => ElemOp1 a
|
|
(/:) = zipWithA (/)
|
|
|
|
infixr 8 ^:, **:, ^^:
|
|
|
|
(^:) :: (Num a, Integral b) => ElemOp a b a
|
|
(^:) = zipWithA (^)
|
|
|
|
(**:) :: (Floating a) => ElemOp1 a
|
|
(**:) = zipWithA (**)
|
|
|
|
(^^:) :: (Fractional a, Integral b) => ElemOp a b a
|
|
(^^:) = zipWithA (^^)
|
|
|
|
-- Convenient aliases for matrix-scalar element-wise operations.
|
|
type ScalarOp a b c = (IArray arr a, IArray arr c, Ix i) =>
|
|
arr i a -> b -> arr i c
|
|
type ScalarOp1 a = ScalarOp a a a
|
|
|
|
samap :: (IArray arr a, IArray arr c, Ix i) =>
|
|
(a -> b -> c) -> arr i a -> b -> arr i c
|
|
samap f a s = amap (`f` s) a
|
|
|
|
infixl 6 +., -.
|
|
infixl 7 *., /., `divS`
|
|
|
|
(+.), (-.), (*.) :: (Num a) => ScalarOp1 a
|
|
(+.) = samap (+)
|
|
(-.) = samap (-)
|
|
(*.) = samap (*)
|
|
|
|
divS :: (Integral a) => ScalarOp1 a
|
|
divS = samap div
|
|
|
|
(/.) :: (Fractional a) => ScalarOp1 a
|
|
(/.) = samap (/)
|
|
|
|
infixr 8 ^., **., ^^.
|
|
|
|
(^.) :: (Num a, Integral b) => ScalarOp a b a
|
|
(^.) = samap (^)
|
|
|
|
(**.) :: (Floating a) => ScalarOp1 a
|
|
(**.) = samap (**)
|
|
|
|
(^^.) :: (Fractional a, Integral b) => ScalarOp a b a
|
|
(^^.) = samap (^^)
|
|
|
|
main :: IO ()
|
|
main = do
|
|
let m1, m2 :: (forall a. (Enum a, Num a) => Array (Int, Int) a)
|
|
m1 = listArray ((0, 0), (2, 3)) [1..]
|
|
m2 = listArray ((0, 0), (2, 3)) [10..]
|
|
s :: (forall a. Num a => a)
|
|
s = 99
|
|
putStrLn "m1"
|
|
print m1
|
|
putStrLn "m2"
|
|
print m2
|
|
putStrLn "s"
|
|
print s
|
|
putStrLn "m1 + m2"
|
|
print $ m1 +: m2
|
|
putStrLn "m1 - m2"
|
|
print $ m1 -: m2
|
|
putStrLn "m1 * m2"
|
|
print $ m1 *: m2
|
|
putStrLn "m1 `div` m2"
|
|
print $ m1 `divE` m2
|
|
putStrLn "m1 / m2"
|
|
print $ m1 /: m2
|
|
putStrLn "m1 ^ m2"
|
|
print $ m1 ^: m2
|
|
putStrLn "m1 ** m2"
|
|
print $ m1 **: m2
|
|
putStrLn "m1 ^^ m2"
|
|
print $ m1 ^^: m2
|
|
putStrLn "m1 + s"
|
|
print $ m1 +. s
|
|
putStrLn "m1 - s"
|
|
print $ m1 -. s
|
|
putStrLn "m1 * s"
|
|
print $ m1 *. s
|
|
putStrLn "m1 `div` s"
|
|
print $ m1 `divS` s
|
|
putStrLn "m1 / s"
|
|
print $ m1 /. s
|
|
putStrLn "m1 ^ s"
|
|
print $ m1 ^. s
|
|
putStrLn "m1 ** s"
|
|
print $ m1 **. s
|
|
putStrLn "m1 ^^ s"
|
|
print $ m1 ^^. s
|