{-# OPTIONS_GHC -Wall #-}
{-# LANGUAGE Trustworthy #-}
{-# LANGUAGE CPP #-}

{- | 
Module      :  Physics.Learn.BeamStack
Copyright   :  (c) Scott N. Walck 2016-2018
License     :  BSD3 (see LICENSE)
Maintainer  :  Scott N. Walck <walck@lvc.edu>
Stability   :  experimental

Splitters, recombiners, and detectors for Stern-Gerlach
experiments.
-}

-- Spin-1/2 mixed states.

module Physics.Learn.BeamStack
    (
    -- * Core laboratory components
      BeamStack()
    , randomBeam
    , split
    , recombine
    , applyBField
    , dropBeam
    , flipBeams
    , numBeams
    , detect
    -- * Standard splitters
    , splitX
    , splitY
    , splitZ
    -- * Standard magnetic fields
    , applyBFieldX
    , applyBFieldY
    , applyBFieldZ
    -- * Standard combiners
    , recombineX
    , recombineY
    , recombineZ
    -- * Filters
    , xpFilter
    , xmFilter
    , zpFilter
    , zmFilter
    )
    where

import Physics.Learn.QuantumMat
    ( zp
    , zm
    , nm
    , np
    , couter
    , oneQubitMixed
    )
import Numeric.LinearAlgebra
    ( C
    , Vector
    , Matrix
    , iC
    , (<>)
    , kronecker
    , fromLists
    , toList
    , toLists
    , scale
    , size
    , takeDiag
    , ident
    , tr
    )
import Data.Complex
    ( Complex(..)
    , realPart
    , imagPart
    )
import Data.List
    ( intercalate
    )
#if MIN_VERSION_base(4,11,0)
import Prelude hiding ((<>))
#endif

data BeamStack = BeamStack (Matrix C)

showOneBeam :: Double -> String
showOneBeam :: Double -> String
showOneBeam Double
r = String
"Beam of intensity " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Double
r

instance Show BeamStack where
    show :: BeamStack -> String
show BeamStack
b = forall a. [a] -> [[a]] -> [a]
intercalate String
"\n" forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map Double -> String
showOneBeam (BeamStack -> [Double]
detect BeamStack
b)

{-
unBeamStack :: BeamStack -> Matrix C
unBeamStack (BeamStack m) = m
-}

--------------------
-- Core functions --
--------------------

-- | A beam of randomly oriented spin-1/2 particles.
randomBeam :: BeamStack
randomBeam :: BeamStack
randomBeam = Matrix C -> BeamStack
BeamStack Matrix C
oneQubitMixed

extendWithZeros :: Matrix C -> Matrix C
extendWithZeros :: Matrix C -> Matrix C
extendWithZeros Matrix C
m
    = let (Int
_,Int
q) = forall (c :: * -> *) t. Container c t => c t -> IndexOf c
size Matrix C
m
          ml :: [[C]]
ml = forall t. Element t => Matrix t -> [[t]]
toLists Matrix C
m
      in forall t. Element t => [[t]] -> Matrix t
fromLists forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (forall a. [a] -> [a] -> [a]
++ [C
0,C
0]) [[C]]
ml
             forall a. [a] -> [a] -> [a]
++ [forall a. Int -> a -> [a]
replicate (Int
qforall a. Num a => a -> a -> a
+Int
2) C
0, forall a. Int -> a -> [a]
replicate (Int
qforall a. Num a => a -> a -> a
+Int
2) C
0]

-- reduce row and column size by 2
reduceMat :: Matrix C -> Matrix C
reduceMat :: Matrix C -> Matrix C
reduceMat Matrix C
m
    = let (Int
p,Int
q) = forall (c :: * -> *) t. Container c t => c t -> IndexOf c
size Matrix C
m
          ml :: [[C]]
ml = forall t. Element t => Matrix t -> [[t]]
toLists Matrix C
m
      in forall t. Element t => [[t]] -> Matrix t
fromLists forall a b. (a -> b) -> a -> b
$ forall a. Int -> [a] -> [a]
take (Int
pforall a. Num a => a -> a -> a
-Int
2) forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (forall a. Int -> [a] -> [a]
take (Int
qforall a. Num a => a -> a -> a
-Int
2)) [[C]]
ml

checkedRealPart :: C -> Double
checkedRealPart :: C -> Double
checkedRealPart C
c
    = let eps :: Double
eps = Double
1e-14
      in if forall a. Complex a -> a
imagPart C
c forall a. Ord a => a -> a -> Bool
< Double
eps
         then forall a. Complex a -> a
realPart C
c
         else forall a. HasCallStack => String -> a
error forall a b. (a -> b) -> a -> b
$ String
"checkRealPart: imagPart = " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show (forall a. Complex a -> a
imagPart C
c)

-- | Return the intensities of a stack of beams.
detect :: BeamStack -> [Double]
detect :: BeamStack -> [Double]
detect (BeamStack Matrix C
m)
    = [C] -> [Double]
addAlternate forall a b. (a -> b) -> a -> b
$ forall a. Storable a => Vector a -> [a]
toList forall a b. (a -> b) -> a -> b
$ forall t. Element t => Matrix t -> Vector t
takeDiag Matrix C
m

addAlternate :: [C] -> [Double]
addAlternate :: [C] -> [Double]
addAlternate [] = []
addAlternate [C
_] = forall a. HasCallStack => String -> a
error String
"addAlternate needs even number of elements"
addAlternate (C
x:C
y:[C]
xs) = C -> Double
checkedRealPart (C
xforall a. Num a => a -> a -> a
+C
y) forall a. a -> [a] -> [a]
: [C] -> [Double]
addAlternate [C]
xs

-- | Remove the most recent beam from the stack.
dropBeam :: BeamStack -> BeamStack
dropBeam :: BeamStack -> BeamStack
dropBeam (BeamStack Matrix C
m) = Matrix C -> BeamStack
BeamStack (Matrix C -> Matrix C
reduceMat Matrix C
m)

-- | Return the number of beams in a 'BeamStack'.
numBeams :: BeamStack -> Int
numBeams :: BeamStack -> Int
numBeams (BeamStack Matrix C
m)
    = let (Int
p,Int
_) = forall (c :: * -> *) t. Container c t => c t -> IndexOf c
size Matrix C
m
      in Int
p forall a. Integral a => a -> a -> a
`div` Int
2

-- | Interchange the two most recent beams on the stack.
flipBeams :: BeamStack -> BeamStack
flipBeams :: BeamStack -> BeamStack
flipBeams (BeamStack Matrix C
m)
    = let (Int
d,Int
_) = forall (c :: * -> *) t. Container c t => c t -> IndexOf c
size Matrix C
m
          fl :: Matrix C
fl = Int -> Matrix C
flipMat Int
d
      in Matrix C -> BeamStack
BeamStack forall a b. (a -> b) -> a -> b
$ Matrix C
fl forall t. Numeric t => Matrix t -> Matrix t -> Matrix t
<> Matrix C
m forall t. Numeric t => Matrix t -> Matrix t -> Matrix t
<> forall m mt. Transposable m mt => m -> mt
tr Matrix C
fl

flipMat :: Int -> Matrix C
flipMat :: Int -> Matrix C
flipMat Int
d = Int -> Matrix C -> Matrix C
bigM Int
d (forall t. Element t => [[t]] -> Matrix t
fromLists [[C
0,C
0,C
1,C
0]
                              ,[C
0,C
0,C
0,C
1]
                              ,[C
1,C
0,C
0,C
0]
                              ,[C
0,C
1,C
0,C
0]])

-- Turn a 2x2 into a dxd.
bigM2 :: Int -> Matrix C -> Matrix C
bigM2 :: Int -> Matrix C -> Matrix C
bigM2 Int
d Matrix C
m
    | Int
d forall a. Ord a => a -> a -> Bool
< Int
2      = forall a. HasCallStack => String -> a
error String
"bigM2 requires d >= 2"
    | forall a. Integral a => a -> Bool
odd Int
d      = forall a. HasCallStack => String -> a
error String
"bigM2 requires even d"
    | Bool
otherwise  = forall t. Element t => [[t]] -> Matrix t
fromLists forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (forall a. [a] -> [a] -> [a]
++ [C
0,C
0]) (forall t. Element t => Matrix t -> [[t]]
toLists (forall a. (Num a, Element a) => Int -> Matrix a
ident (Int
dforall a. Num a => a -> a -> a
-Int
2)))
                   forall a. [a] -> [a] -> [a]
++ forall a b. (a -> b) -> [a] -> [b]
map (forall a. Int -> a -> [a]
replicate (Int
dforall a. Num a => a -> a -> a
-Int
2) C
0 forall a. [a] -> [a] -> [a]
++) (forall t. Element t => Matrix t -> [[t]]
toLists Matrix C
m)

-- Turn a 4x4 into a dxd.
bigM :: Int -> Matrix C -> Matrix C
bigM :: Int -> Matrix C -> Matrix C
bigM Int
d Matrix C
m
    | Int
d forall a. Ord a => a -> a -> Bool
< Int
4      = forall a. HasCallStack => String -> a
error String
"bigM requires d >= 4"
    | forall a. Integral a => a -> Bool
odd Int
d      = forall a. HasCallStack => String -> a
error String
"bigM requires even d"
    | Bool
otherwise  = forall t. Element t => [[t]] -> Matrix t
fromLists forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map (forall a. [a] -> [a] -> [a]
++ [C
0,C
0,C
0,C
0]) (forall t. Element t => Matrix t -> [[t]]
toLists (forall a. (Num a, Element a) => Int -> Matrix a
ident (Int
dforall a. Num a => a -> a -> a
-Int
4)))
                   forall a. [a] -> [a] -> [a]
++ forall a b. (a -> b) -> [a] -> [b]
map (forall a. Int -> a -> [a]
replicate (Int
dforall a. Num a => a -> a -> a
-Int
4) C
0 forall a. [a] -> [a] -> [a]
++) (forall t. Element t => Matrix t -> [[t]]
toLists Matrix C
m)

s :: Double -> Double -> Matrix C
s :: Double -> Double -> Matrix C
s Double
theta Double
phi = forall t. Product t => Matrix t -> Matrix t -> Matrix t
kronecker (Vector C
u Vector C -> Vector C -> Matrix C
`couter` Vector C
u) (Double -> Double -> Vector C
np Double
theta Double
phi Vector C -> Vector C -> Matrix C
`couter` Double -> Double -> Vector C
np Double
theta Double
phi)
            forall a. Num a => a -> a -> a
+ forall t. Product t => Matrix t -> Matrix t -> Matrix t
kronecker (Vector C
l Vector C -> Vector C -> Matrix C
`couter` Vector C
u) (Double -> Double -> Vector C
nm Double
theta Double
phi Vector C -> Vector C -> Matrix C
`couter` Double -> Double -> Vector C
nm Double
theta Double
phi)
            forall a. Num a => a -> a -> a
+ forall t. Product t => Matrix t -> Matrix t -> Matrix t
kronecker (Vector C
u Vector C -> Vector C -> Matrix C
`couter` Vector C
l) (Double -> Double -> Vector C
nm Double
theta Double
phi Vector C -> Vector C -> Matrix C
`couter` Double -> Double -> Vector C
nm Double
theta Double
phi)
            forall a. Num a => a -> a -> a
+ forall t. Product t => Matrix t -> Matrix t -> Matrix t
kronecker (Vector C
l Vector C -> Vector C -> Matrix C
`couter` Vector C
l) (Double -> Double -> Vector C
np Double
theta Double
phi Vector C -> Vector C -> Matrix C
`couter` Double -> Double -> Vector C
np Double
theta Double
phi)

u :: Vector C
u :: Vector C
u = Vector C
zp

l :: Vector C
l :: Vector C
l = Vector C
zm

-- | Given angles describing the orientation of the splitter,
--   removes an incoming beam from the stack and replaces
--   it with two beams, a spin-up and a spin-down beam.
--   The spin-down beam is the most recent beam on the stack.
split :: Double -> Double -> BeamStack -> BeamStack
split :: Double -> Double -> BeamStack -> BeamStack
split Double
theta Double
phi (BeamStack Matrix C
m)
    = let m' :: Matrix C
m' = Matrix C -> Matrix C
extendWithZeros Matrix C
m
          (Int
p,Int
_) = forall (c :: * -> *) t. Container c t => c t -> IndexOf c
size Matrix C
m'
          ss :: Matrix C
ss = Int -> Matrix C -> Matrix C
bigM Int
p (Double -> Double -> Matrix C
s Double
theta Double
phi)
      in Matrix C -> BeamStack
BeamStack forall a b. (a -> b) -> a -> b
$ Matrix C
ss forall t. Numeric t => Matrix t -> Matrix t -> Matrix t
<> Matrix C
m' forall t. Numeric t => Matrix t -> Matrix t -> Matrix t
<> forall m mt. Transposable m mt => m -> mt
tr Matrix C
ss

-- | Given angles describing the orientation of the recombiner,
--   returns a single beam from an incoming pair of beams.
recombine :: Double -> Double -> BeamStack -> BeamStack
recombine :: Double -> Double -> BeamStack -> BeamStack
recombine Double
theta Double
phi (BeamStack Matrix C
m)
    = let (Int
d,Int
_) = forall (c :: * -> *) t. Container c t => c t -> IndexOf c
size Matrix C
m
          ss :: Matrix C
ss = Int -> Matrix C -> Matrix C
bigM Int
d (Double -> Double -> Matrix C
s Double
theta Double
phi)
      in BeamStack -> BeamStack
dropBeam forall a b. (a -> b) -> a -> b
$ Matrix C -> BeamStack
BeamStack forall a b. (a -> b) -> a -> b
$ Matrix C
ss forall t. Numeric t => Matrix t -> Matrix t -> Matrix t
<> Matrix C
m forall t. Numeric t => Matrix t -> Matrix t -> Matrix t
<> forall m mt. Transposable m mt => m -> mt
tr Matrix C
ss

mag2x2 :: Double -> Double -> Double -> Matrix C
mag2x2 :: Double -> Double -> Double -> Matrix C
mag2x2 Double
theta Double
phi Double
omegaT
    = let z :: C
z = C
iC forall a. Num a => a -> a -> a
* (Double
omegaT forall a. a -> a -> Complex a
:+ Double
0) forall a. Fractional a => a -> a -> a
/ C
2
          np' :: Vector C
np' = Double -> Double -> Vector C
np Double
theta Double
phi
          nm' :: Vector C
nm' = Double -> Double -> Vector C
nm Double
theta Double
phi
      in forall t (c :: * -> *). Linear t c => t -> c t -> c t
scale (forall a. Floating a => a -> a
exp   C
z ) (Vector C
np' Vector C -> Vector C -> Matrix C
`couter` Vector C
np')
       forall a. Num a => a -> a -> a
+ forall t (c :: * -> *). Linear t c => t -> c t -> c t
scale (forall a. Floating a => a -> a
exp (-C
z)) (Vector C
nm' Vector C -> Vector C -> Matrix C
`couter` Vector C
nm')

-- | Given angles describing the direction of a
--   uniform magnetic field, and given an angle
--   describing the product of the Larmor frequency
--   and the time, return an output beam from an
--   input beam.
applyBField :: Double -> Double -> Double -> BeamStack -> BeamStack
applyBField :: Double -> Double -> Double -> BeamStack -> BeamStack
applyBField Double
theta Double
phi Double
omegaT (BeamStack Matrix C
m)
    = let (Int
d,Int
_) = forall (c :: * -> *) t. Container c t => c t -> IndexOf c
size Matrix C
m
          uu :: Matrix C
uu = Int -> Matrix C -> Matrix C
bigM2 Int
d (Double -> Double -> Double -> Matrix C
mag2x2 Double
theta Double
phi Double
omegaT)
      in Matrix C -> BeamStack
BeamStack forall a b. (a -> b) -> a -> b
$ Matrix C
uu forall t. Numeric t => Matrix t -> Matrix t -> Matrix t
<> Matrix C
m forall t. Numeric t => Matrix t -> Matrix t -> Matrix t
<> forall m mt. Transposable m mt => m -> mt
tr Matrix C
uu

-----------------------
-- Derived functions --
-----------------------

-- | A Stern-Gerlach splitter in the x direction.
splitX :: BeamStack -> BeamStack
splitX :: BeamStack -> BeamStack
splitX = Double -> Double -> BeamStack -> BeamStack
split (forall a. Floating a => a
piforall a. Fractional a => a -> a -> a
/Double
2) Double
0

-- | A Stern-Gerlach splitter in the y direction.
splitY :: BeamStack -> BeamStack
splitY :: BeamStack -> BeamStack
splitY = Double -> Double -> BeamStack -> BeamStack
split (forall a. Floating a => a
piforall a. Fractional a => a -> a -> a
/Double
2) (forall a. Floating a => a
piforall a. Fractional a => a -> a -> a
/Double
2)

-- | A Stern-Gerlach splitter in the z direction.
splitZ :: BeamStack -> BeamStack
splitZ :: BeamStack -> BeamStack
splitZ = Double -> Double -> BeamStack -> BeamStack
split Double
0 Double
0

-- | Given an angle in radians
--   describing the product of the Larmor frequency
--   and the time, apply a magnetic in the x direction
--   to the most recent beam on the stack.
applyBFieldX :: Double -> BeamStack -> BeamStack
applyBFieldX :: Double -> BeamStack -> BeamStack
applyBFieldX = Double -> Double -> Double -> BeamStack -> BeamStack
applyBField (forall a. Floating a => a
piforall a. Fractional a => a -> a -> a
/Double
2) Double
0

-- | Given an angle in radians
--   describing the product of the Larmor frequency
--   and the time, apply a magnetic in the y direction
--   to the most recent beam on the stack.
applyBFieldY :: Double -> BeamStack -> BeamStack
applyBFieldY :: Double -> BeamStack -> BeamStack
applyBFieldY = Double -> Double -> Double -> BeamStack -> BeamStack
applyBField (forall a. Floating a => a
piforall a. Fractional a => a -> a -> a
/Double
2) (forall a. Floating a => a
piforall a. Fractional a => a -> a -> a
/Double
2)

-- | Given an angle in radians
--   describing the product of the Larmor frequency
--   and the time, apply a magnetic in the z direction
--   to the most recent beam on the stack.
applyBFieldZ :: Double -> BeamStack -> BeamStack
applyBFieldZ :: Double -> BeamStack -> BeamStack
applyBFieldZ = Double -> Double -> Double -> BeamStack -> BeamStack
applyBField Double
0 Double
0

-- | A Stern-Gerlach recombiner in the x direction.
recombineX :: BeamStack -> BeamStack
recombineX :: BeamStack -> BeamStack
recombineX = Double -> Double -> BeamStack -> BeamStack
recombine (forall a. Floating a => a
piforall a. Fractional a => a -> a -> a
/Double
2) Double
0

-- | A Stern-Gerlach recombiner in the y direction.
recombineY :: BeamStack -> BeamStack
recombineY :: BeamStack -> BeamStack
recombineY = Double -> Double -> BeamStack -> BeamStack
recombine (forall a. Floating a => a
piforall a. Fractional a => a -> a -> a
/Double
2) (forall a. Floating a => a
piforall a. Fractional a => a -> a -> a
/Double
2)

-- | A Stern-Gerlach recombiner in the z direction.
recombineZ :: BeamStack -> BeamStack
recombineZ :: BeamStack -> BeamStack
recombineZ = Double -> Double -> BeamStack -> BeamStack
recombine Double
0 Double
0

-- | Filter for spin-up particles in the x direction.
xpFilter :: BeamStack -> BeamStack
xpFilter :: BeamStack -> BeamStack
xpFilter = BeamStack -> BeamStack
dropBeam forall b c a. (b -> c) -> (a -> b) -> a -> c
. BeamStack -> BeamStack
splitX

-- | Filter for spin-down particles in the x direction.
xmFilter :: BeamStack -> BeamStack
xmFilter :: BeamStack -> BeamStack
xmFilter = BeamStack -> BeamStack
dropBeam forall b c a. (b -> c) -> (a -> b) -> a -> c
. BeamStack -> BeamStack
flipBeams forall b c a. (b -> c) -> (a -> b) -> a -> c
. BeamStack -> BeamStack
splitX

-- | Filter for spin-up particles in the z direction.
zpFilter :: BeamStack -> BeamStack
zpFilter :: BeamStack -> BeamStack
zpFilter = BeamStack -> BeamStack
dropBeam forall b c a. (b -> c) -> (a -> b) -> a -> c
. BeamStack -> BeamStack
splitZ

-- | Filter for spin-down particles in the z direction.
zmFilter :: BeamStack -> BeamStack
zmFilter :: BeamStack -> BeamStack
zmFilter = BeamStack -> BeamStack
dropBeam forall b c a. (b -> c) -> (a -> b) -> a -> c
. BeamStack -> BeamStack
flipBeams forall b c a. (b -> c) -> (a -> b) -> a -> c
. BeamStack -> BeamStack
splitZ