Skip to content

Commit

Permalink
Merge pull request #3 from fhaust/master
Browse files Browse the repository at this point in the history
added functions to deal with complex vectors
  • Loading branch information
mstksg committed Feb 4, 2020
2 parents cca61c8 + 5207c4e commit 91661e0
Showing 1 changed file with 45 additions and 0 deletions.
45 changes: 45 additions & 0 deletions src/Numeric/LinearAlgebra/Static/Backprop.hs
Original file line number Diff line number Diff line change
Expand Up @@ -123,6 +123,13 @@ module Numeric.LinearAlgebra.Static.Backprop (
, H.C
, H.M
, H.𝑖
, toComplex
, fromComplex
, complex
, real
, imag
, sqMagnitude
, magnitude
-- * Products
, (<>)
, (#>)
Expand Down Expand Up @@ -1305,3 +1312,41 @@ infixr 8 <·>
afSV :: Backprop a => BE.AddFunc (SV.Vector n a)
afSV = BE.AF (SV.zipWith add)
{-# INLINE afSV #-}



-- support for complex types
toComplex :: (KnownNat n, Reifies s W)
=> BVar s (H.R n) -> BVar s (H.R n) -> BVar s (H.C n)
toComplex = isoVar2 (curry H.toComplex) H.fromComplex
{-# INLINE toComplex #-}

fromComplex :: (KnownNat n, Reifies s W)
=> BVar s (H.C n) -> (BVar s (H.R n), BVar s (H.R n))
fromComplex c = let ri = isoVar H.fromComplex H.toComplex c in (ri ^^. _1, ri ^^. _2)
{-# INLINE fromComplex #-}

complex :: (KnownNat n, Reifies s W)
=> BVar s (H.R n) -> BVar s (H.C n)
complex = isoVar H.complex H.real
{-# INLINE complex #-}

real :: (KnownNat n, Reifies s W)
=> BVar s (H.C n) -> BVar s (H.R n)
real = isoVar H.real (\r -> H.toComplex (r, H.konst 0))
{-# INLINE real #-}

imag :: (KnownNat n, Reifies s W)
=> BVar s (H.C n) -> BVar s (H.R n)
imag = isoVar H.imag (\r -> H.toComplex (H.konst 0, r))
{-# INLINE imag #-}

sqMagnitude :: (KnownNat n, Reifies s W)
=> BVar s (H.C n) -> BVar s (H.R n)
sqMagnitude c = let (r,i) = fromComplex c in r**2 + i**2
{-# INLINE sqMagnitude #-}

magnitude :: (KnownNat n, Reifies s W)
=> BVar s (H.C n) -> BVar s (H.R n)
magnitude c = sqrt $ sqMagnitude c
{-# INLINE magnitude #-}

0 comments on commit 91661e0

Please sign in to comment.