easytensor-0.3.0.0: src-base/Numeric/Array/Family/ArrayD.hs
{-# LANGUAGE CPP #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UnboxedTuples #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE BangPatterns #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}
-----------------------------------------------------------------------------
-- |
-- Module : Numeric.Array.Family.ArrayD
-- Copyright : (c) Artem Chirkin
-- License : BSD3
--
-- Maintainer : chirkin@arch.ethz.ch
--
--
-----------------------------------------------------------------------------
module Numeric.Array.Family.ArrayD () where
import GHC.Base (runRW#)
import GHC.Prim
import GHC.Types (Double (..), Int (..),
RuntimeRep (..), isTrue#)
import Numeric.Array.ElementWise
import Numeric.Array.Family
import Numeric.Commons
import Numeric.DataFrame.Type
import Numeric.Dimensions
import Numeric.Dimensions.Traverse
import Numeric.TypeLits
import Numeric.Matrix.Type
#include "MachDeps.h"
#define ARR_TYPE ArrayD
#define ARR_FROMSCALAR FromScalarD#
#define ARR_CONSTR ArrayD#
#define EL_TYPE_BOXED Double
#define EL_TYPE_PRIM Double#
#define EL_RUNTIME_REP 'DoubleRep
#define EL_CONSTR D#
#define EL_SIZE SIZEOF_HSDOUBLE#
#define EL_ALIGNMENT ALIGNMENT_HSDOUBLE#
#define EL_ZERO 0.0##
#define EL_ONE 1.0##
#define EL_MINUS_ONE -1.0##
#define INDEX_ARRAY indexDoubleArray#
#define WRITE_ARRAY writeDoubleArray#
#define OP_EQ (==##)
#define OP_NE (/=##)
#define OP_GT (>##)
#define OP_GE (>=##)
#define OP_LT (<##)
#define OP_LE (<=##)
#define OP_PLUS (+##)
#define OP_MINUS (-##)
#define OP_TIMES (*##)
#define OP_NEGATE negateDouble#
#include "Array.h"
instance Num (ArrayD ds) where
(+) = zipV (+##)
{-# INLINE (+) #-}
(-) = zipV (-##)
{-# INLINE (-) #-}
(*) = zipV (*##)
{-# INLINE (*) #-}
negate = mapV negateDouble#
{-# INLINE negate #-}
abs = mapV (\x -> if isTrue# (x >=## 0.0##)
then x
else negateDouble# x
)
{-# INLINE abs #-}
signum = mapV (\x -> if isTrue# (x >## 0.0##)
then 1.0##
else if isTrue# (x <## 0.0##)
then -1.0##
else 0.0##
)
{-# INLINE signum #-}
fromInteger = broadcastArray . fromInteger
{-# INLINE fromInteger #-}
instance Fractional (ArrayD ds) where
(/) = zipV (/##)
{-# INLINE (/) #-}
recip = mapV (1.0## /##)
{-# INLINE recip #-}
fromRational = broadcastArray . fromRational
{-# INLINE fromRational #-}
instance Floating (ArrayD ds) where
pi = broadcastArray pi
{-# INLINE pi #-}
exp = mapV expDouble#
{-# INLINE exp #-}
log = mapV logDouble#
{-# INLINE log #-}
sqrt = mapV sqrtDouble#
{-# INLINE sqrt #-}
sin = mapV sinDouble#
{-# INLINE sin #-}
cos = mapV cosDouble#
{-# INLINE cos #-}
tan = mapV tanDouble#
{-# INLINE tan #-}
asin = mapV asinDouble#
{-# INLINE asin #-}
acos = mapV acosDouble#
{-# INLINE acos #-}
atan = mapV atanDouble#
{-# INLINE atan #-}
sinh = mapV sinDouble#
{-# INLINE sinh #-}
cosh = mapV coshDouble#
{-# INLINE cosh #-}
tanh = mapV tanhDouble#
{-# INLINE tanh #-}
(**) = zipV (**##)
{-# INLINE (**) #-}
logBase = zipV (\x y -> logDouble# y /## logDouble# x)
{-# INLINE logBase #-}
asinh = mapV (\x -> logDouble# (x +##
sqrtDouble# (1.0## +## x *## x)))
{-# INLINE asinh #-}
acosh = mapV (\x -> case x +## 1.0## of
y -> logDouble# ( x +## y *##
sqrtDouble# ((x -## 1.0##) /## y)
)
)
{-# INLINE acosh #-}
atanh = mapV (\x -> 0.5## *##
logDouble# ((1.0## +## x) /## (1.0## -## x)))
{-# INLINE atanh #-}
instance (KnownNat n, KnownNat m, ArrayD '[n,m] ~ Array Double '[n,m], 2 <= n, 2 <= m)
=> MatrixCalculus Double n m where
transpose (KnownDataFrame (ArrayD# offs nm arr)) = case runRW#
( \s0 -> case newByteArray# bs s0 of
(# s1, marr #) -> case loop2# n m
(\i j s' -> writeDoubleArray# marr (j +# m *# i)
(indexDoubleArray# arr (offs +# j *# n +# i)) s'
) s1 of
s2 -> unsafeFreezeByteArray# marr s2
) of (# _, r #) -> fromBytes (# 0#, nm, r #)
where
n = case fromInteger $ natVal (Proxy @n) of I# np -> np
m = case fromInteger $ natVal (Proxy @m) of I# mp -> mp
bs = n *# m *# EL_SIZE
transpose (KnownDataFrame (FromScalarD# x)) = unsafeCoerce# $ FromScalarD# x
instance ( KnownDim n, ArrayD '[n,n] ~ Array Double '[n,n] )
=> SquareMatrixCalculus Double n where
eye = case runRW#
( \s0 -> case newByteArray# bs s0 of
(# s1, marr #) -> case loop1# n
(\j s' -> writeDoubleArray# marr (j *# n1) 1.0## s'
) (setByteArray# marr 0# bs 0# s1) of
s2 -> unsafeFreezeByteArray# marr s2
) of (# _, r #) -> fromBytes (# 0#, n *# n, r #)
where
n1 = n +# 1#
n = case dimVal' @n of I# np -> np
bs = n *# n *# EL_SIZE
{-# INLINE eye #-}
diag (KnownDataFrame (Scalar (D# v))) = case runRW#
( \s0 -> case newByteArray# bs s0 of
(# s1, marr #) -> case loop1# n
(\j s' -> writeDoubleArray# marr (j *# n1) v s'
) (setByteArray# marr 0# bs 0# s1) of
s2 -> unsafeFreezeByteArray# marr s2
) of (# _, r #) -> fromBytes (# 0#, n *# n, r #)
where
n1 = n +# 1#
n = case dimVal' @n of I# np -> np
bs = n *# n *# EL_SIZE
{-# INLINE diag #-}
det (KnownDataFrame (ArrayD# off nsqr arr)) = case runRW#
( \s0 -> case newByteArray# bs s0 of
(# s1, mat #) -> case newByteArray#
(n *# EL_SIZE)
(copyByteArray# arr offb mat 0# bs s1) of
(# s2, vec #) ->
let f i x s | isTrue# (i >=# n) = (# s, x #)
| otherwise =
let !(# s' , j #) = maxInRowRem# n n i mat s
!(# s'', x' #) = if isTrue# (i /=# j)
then (# swapCols# n i j vec mat s'
, negateDouble# x #)
else (# s', x #)
!(# s''', y #) = clearRowEnd# n n i mat s''
in if isTrue# (0.0## ==## y)
then (# s''', 0.0## #)
else f (i +# 1#) (x' *## y) s'''
in f 0# 1.0## s2
) of (# _, r #) -> KnownDataFrame (Scalar (D# r))
where
n = case dimVal' @n of I# np -> np
offb = off *# EL_SIZE
bs = nsqr *# EL_SIZE
det (KnownDataFrame (FromScalarD# _)) = 0
{-# INLINE det #-}
trace (KnownDataFrame (ArrayD# off nsqr a)) = KnownDataFrame (Scalar (D# (loop' 0# 0.0##)))
where
n1 = n +# 1#
n = case dimVal' @n of I# np -> np
loop' i acc | isTrue# (i ># nsqr) = acc
| otherwise = loop' (i +# n1)
(indexDoubleArray# a (off +# i) +## acc)
trace (KnownDataFrame (FromScalarD# x)) = KnownDataFrame (Scalar (D# (x *## n)))
where
n = case fromIntegral (dimVal' @n) of D# np -> np
{-# INLINE trace #-}
instance (KnownNat n, ArrayD '[n,n] ~ Array Double '[n,n], 2 <= n) => MatrixInverse Double n where
inverse (KnownDataFrame (ArrayD# offs nsqr arr)) = case runRW#
( \s0 -> case newByteArray# (bs *# 2#) s0 of
(# s1, mat #) -> case newByteArray# (vs *# 2#)
-- copy original matrix to the top of an augmented matrix
(loop1# n (\i s -> writeDoubleArray# mat
(i *# nn +# i +# n) 1.0##
(copyByteArray# arr (offb +# i *# vs)
mat (2# *# i *# vs) vs s))
(setByteArray# mat 0# (bs *# 2#) 0# s1)
) of
(# s2, vec #) ->
let f i s | isTrue# (i >=# n) = s
| otherwise =
let !(# s' , j #) = maxInRowRem# nn n i mat s
s'' = if isTrue# (i /=# j) then swapCols# nn i j vec mat s'
else s'
!(# s''', _ #) = clearRowAll# nn n i mat s''
in f (i +# 1#) s'''
in unsafeFreezeByteArray# mat
( shrinkMutableByteArray# mat bs
(-- copy inverse matrix from the augmented part
loop1# n (\i s ->
copyMutableByteArray# mat
(2# *# i *# vs +# vs)
mat (i *# vs) vs s)
(f 0# s2)
)
)
) of (# _, r #) -> KnownDataFrame (ArrayD# 0# nsqr r)
where
nn = 2# *# n
n = case fromInteger $ natVal (Proxy @n) of I# np -> np
vs = n *# EL_SIZE
bs = n *# n *# EL_SIZE
offb = offs *# EL_SIZE
inverse (KnownDataFrame (FromScalarD# _)) = error "Cannot take inverse of a degenerate matrix"
-----------------------------------------------------------------------------
-- Helpers
-----------------------------------------------------------------------------
-- #ifndef UNSAFE_INDICES
-- | isTrue# ( (i ># dim# _x)
-- `orI#` (i <=# 0#)
-- ) = error $ "Bad index " ++
-- show (I# i) ++ " for " ++ show (dim _x) ++ "D vector"
-- | otherwise
-- #endif
-- | Swap columns i and j. Does not check if i or j is larger than matrix width m
swapCols# :: Int# -- n
-> Int# -- ith column to swap
-> Int# -- jth column to swap
-> MutableByteArray# s -- buffer byte array of length of n elems
-> MutableByteArray# s -- byte array of matrix
-> State# s -- previous state
-> State# s -- next state
swapCols# n i j vec mat s0 =
-- copy ith column to bugger vec
case copyMutableByteArray# mat (i *# bs) vec 0# bs s0 of
s1 -> case copyMutableByteArray# mat (j *# bs) mat (i *# bs) bs s1 of
s2 -> copyMutableByteArray# vec 0# mat (j *# bs) bs s2
where
bs = n *# EL_SIZE
-- | Starting from i-th row and i+1-th column, substract a multiple of i-th column from i+1 .. m columns,
-- such that there are only zeroes in i-th row and i+1..m columns elements.
clearRowEnd# :: Int# -- n
-> Int# -- m
-> Int# -- ith column to remove from all others
-> MutableByteArray# s -- byte array of matrix
-> State# s -- previous state
-> (# State# s, Double# #) -- next state and a diagonal element
clearRowEnd# n m i mat s0 = (# loop' (i +# 1#) s1, y' #)
where
y0 = (n +# 1#) *# i +# 1# -- first element in source column
!(# s1, y' #) = readDoubleArray# mat ((n +# 1#) *# i) s0 -- diagonal element, must be non-zero
yrc = 1.0## /## y'
n' = n -# i -# 1#
loop' k s | isTrue# (k >=# m) = s
| otherwise = loop' (k +# 1#)
( let x0 = k *# n +# i
!(# s', a' #) = readDoubleArray# mat x0 s
s'' = writeDoubleArray# mat x0 0.0## s'
a = a' *## yrc
in multNRem# n' (x0 +# 1#) y0 a mat s''
)
-- | Substract a multiple of i-th column from 0 .. i-1 and i+1 .. m columns,
-- such that there are only zeroes in i-th row everywhere except i-th column
-- Assuming that elements in 0..i-1 columnts and in i-th row are zeroes, so they do not affect other columns.
-- After all columns updated, divide i-th row by its diagonal element, so (i,i) element has 1.
clearRowAll# :: Int# -- n
-> Int# -- m
-> Int# -- ith column to remove from all others
-> MutableByteArray# s -- byte array of matrix
-> State# s -- previous state
-> (# State# s, Double# #) -- next state and a diagonal element
clearRowAll# n m i mat s0 = (# divLoop (i +# 1#)
(writeDoubleArray# mat ((n +# 1#) *# i) 1.0##
(loop' 0# i (loop' (i +# 1#) m s1))), y' #)
where
y0 = (n +# 1#) *# i +# 1# -- first element in source column
!(# s1, y' #) = readDoubleArray# mat ((n +# 1#) *# i) s0 -- diagonal element, must be non-zero
yrc = 1.0## /## y'
n' = n -# i -# 1#
loop' k km s | isTrue# (k >=# km) = s
| otherwise = loop' (k +# 1#) km
( let x0 = k *# n +# i
!(# s', a' #) = readDoubleArray# mat x0 s
s'' = writeDoubleArray# mat x0 0.0## s'
a = a' *## yrc
in multNRem# n' (x0 +# 1#) y0 a mat s''
)
divLoop k s | isTrue# (k >=# n) = s
| otherwise = divLoop (k +# 1#)
( let x0 = n *# i +# k
!(# s', x #) = readDoubleArray# mat x0 s
in writeDoubleArray# mat x0 (x *## yrc) s'
)
-- | Remove a multiple of one row from another one.
-- do: xi = xi - yi*a
multNRem# :: Int# -- n - nr of elements to go through
-> Int# -- start idx of x (update)
-> Int# -- start idx of y (read)
-> Double# -- multiplier a
-> MutableByteArray# s -- byte array of matrix
-> State# s -- previous state
-> State# s -- next state
multNRem# 0# _ _ _ _ s = s
multNRem# n x0 y0 a mat s = multNRem# (n -# 1#) (x0 +# 1#) (y0 +# 1#) a mat
( case readDoubleArray# mat y0 s of
(# s1, y #) -> case readDoubleArray# mat x0 s1 of
(# s2, x #) -> writeDoubleArray# mat x0 (x -## y *## a) s2
)
-- | Gives index of maximum (absolute) element in i-th row, starting from i-th element only.
-- If i >= m then returns i.
maxInRowRem# :: Int# -- n
-> Int# -- m
-> Int# -- ith column to start to search for and a row to look in
-> MutableByteArray# s -- byte array of matrix
-> State# s -- previous state
-> (# State# s, Int# #) -- next state
maxInRowRem# n m i mat s0 = loop' i (abs# v) i s1
where
!(# s1, v #) = readDoubleArray# mat ((n +# 1#) *# i) s0
abs# x = if isTrue# (x >=## 0.0##) then x else negateDouble# x
loop' ok ov k s | isTrue# (k >=# m) = (# s, ok #)
| otherwise = case readDoubleArray# mat (n *# k +# i) s of
(# s', v' #) -> if isTrue# (abs# v' >## ov)
then loop' k (abs# v') (k +# 1#) s'
else loop' ok ov (k +# 1#) s'
-- | Do something in a loop for int i from 0 to n-1 and j from 0 to m-1
loop2# :: Int# -> Int# -> (Int# -> Int#-> State# s -> State# s)
-> State# s -> State# s
loop2# n m f = loop0 0# 0#
where
loop0 i j s | isTrue# (j ==# m) = s
| isTrue# (i ==# n) = loop0 0# (j +# 1#) s
| otherwise = case f i j s of s1 -> loop0 (i +# 1#) j s1
{-# INLINE loop2# #-}