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

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

This module contains functions for displaying the
state of a spin-1/2 particle or other quantum two-level
system as a point on the Bloch Sphere.
-}

module Physics.Learn.BlochSphere
    ( VisObj
    , toPos
    , ketToPos
    , staticBlochSphere
    , displayStaticState
    , animatedBlochSphere
    , simulateBlochSphere
    , simulateBlochSphereK
    , stateProp
    , statePropK
    , evolutionBlochSphere
    , evolutionBlochSphereK
    , hamRabi
    )
    where

import qualified Physics.Learn.QuantumMat as M
import qualified Physics.Learn.Ket as K
import Physics.Learn.Ket
    ( Ket
    , Operator
    , (<>)
    , dagger
    )
import Numeric.LinearAlgebra
    ( Vector
    , Matrix
    , C
    , iC
--    , (<>)  -- matrix multiplication
--    , (|>)  -- vector definition
    , (!)   -- vector element access
    , (><)  -- matrix definition
    , scale
    , size
    )
import Data.Complex
    ( Complex(..)
    , conjugate
    , realPart
    , imagPart
    )
import Physics.Learn
    ( Position
    , v3FromPos
    , cart
    )
import Vis
    ( VisObject(..)
    , Flavour(..)
    , Options(..)
    , Camera0(..)
    , Euler(..)
    , defaultOpts
    , display
    , simulate
    , blue
    , red
    )
#if MIN_VERSION_base(4,11,0)
import Prelude hiding ((<>))
#endif

{-
3 ways to specify the state of a spin-1/2 particle:
Vector C
Ket
Position  (Bloch vector)

2 ways to specify a Hamiltonian:
Matrix C
Operator

3 choices for Vis' world:
(Float, Vector C)
(Float, Ket)
(Float, Position)
-}

-- | A Vis object.
type VisObj = VisObject Double

-- | Convert a 2x1 complex state vector for a qubit
--   into Bloch (x,y,z) coordinates.
toPos :: Vector C -> Position
toPos :: Vector C -> Position
toPos Vector C
v
    = if Vector C -> IndexOf Vector
forall (c :: * -> *) t. Container c t => c t -> IndexOf c
size Vector C
v Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
2
      then [Char] -> Position
forall a. HasCallStack => [Char] -> a
error [Char]
"toPos only for size 2 vectors"
      else let z1 :: C
z1 = Vector C
v Vector C -> Int -> C
forall c t. Indexable c t => c -> Int -> t
! Int
0
               z2 :: C
z2 = Vector C
v Vector C -> Int -> C
forall c t. Indexable c t => c -> Int -> t
! Int
1
           in Double -> Double -> Double -> Position
cart (Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* C -> Double
forall a. Complex a -> a
realPart (C -> C
forall a. Num a => Complex a -> Complex a
conjugate C
z1 C -> C -> C
forall a. Num a => a -> a -> a
* C
z2))
                   (Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* C -> Double
forall a. Complex a -> a
imagPart (C -> C
forall a. Num a => Complex a -> Complex a
conjugate C
z1 C -> C -> C
forall a. Num a => a -> a -> a
* C
z2))
                   (C -> Double
forall a. Complex a -> a
realPart (C -> C
forall a. Num a => Complex a -> Complex a
conjugate C
z1 C -> C -> C
forall a. Num a => a -> a -> a
* C
z1 C -> C -> C
forall a. Num a => a -> a -> a
- C -> C
forall a. Num a => Complex a -> Complex a
conjugate C
z2 C -> C -> C
forall a. Num a => a -> a -> a
* C
z2))

-- | Convert a qubit ket
--   into Bloch (x,y,z) coordinates.
ketToPos :: Ket -> Position
ketToPos :: Ket -> Position
ketToPos Ket
psi
    = if Ket -> Int
forall a b. Representable a b => a -> Int
K.dim Ket
psi Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
2
      then [Char] -> Position
forall a. HasCallStack => [Char] -> a
error [Char]
"ketToPos only for qubit kets"
      else let z1 :: C
z1 = Ket -> Bra
forall a b. Dagger a b => a -> b
dagger Ket
K.zp Bra -> Ket -> C
forall a b c. Mult a b c => a -> b -> c
<> Ket
psi
               z2 :: C
z2 = Ket -> Bra
forall a b. Dagger a b => a -> b
dagger Ket
K.zm Bra -> Ket -> C
forall a b c. Mult a b c => a -> b -> c
<> Ket
psi
           in Double -> Double -> Double -> Position
cart (Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* C -> Double
forall a. Complex a -> a
realPart (C -> C
forall a. Num a => Complex a -> Complex a
conjugate C
z1 C -> C -> C
forall a. Num a => a -> a -> a
* C
z2))
                   (Double
2 Double -> Double -> Double
forall a. Num a => a -> a -> a
* C -> Double
forall a. Complex a -> a
imagPart (C -> C
forall a. Num a => Complex a -> Complex a
conjugate C
z1 C -> C -> C
forall a. Num a => a -> a -> a
* C
z2))
                   (C -> Double
forall a. Complex a -> a
realPart (C -> C
forall a. Num a => Complex a -> Complex a
conjugate C
z1 C -> C -> C
forall a. Num a => a -> a -> a
* C
z1 C -> C -> C
forall a. Num a => a -> a -> a
- C -> C
forall a. Num a => Complex a -> Complex a
conjugate C
z2 C -> C -> C
forall a. Num a => a -> a -> a
* C
z2))

-- | A static 'VisObj' for the state of a qubit.
staticBlochSphere :: Position -> VisObj
staticBlochSphere :: Position -> VisObj
staticBlochSphere Position
r
    = Euler Double -> VisObj -> VisObj
forall a. Euler a -> VisObject a -> VisObject a
RotEulerDeg (Double -> Double -> Double -> Euler Double
forall a. a -> a -> a -> Euler a
Euler Double
270 Double
0 Double
0) (VisObj -> VisObj) -> VisObj -> VisObj
forall a b. (a -> b) -> a -> b
$ Euler Double -> VisObj -> VisObj
forall a. Euler a -> VisObject a -> VisObject a
RotEulerDeg (Double -> Double -> Double -> Euler Double
forall a. a -> a -> a -> Euler a
Euler Double
0 Double
180 Double
0) (VisObj -> VisObj) -> VisObj -> VisObj
forall a b. (a -> b) -> a -> b
$
      [VisObj] -> VisObj
forall a. [VisObject a] -> VisObject a
VisObjects [ Double -> Flavour -> Color -> VisObj
forall a. a -> Flavour -> Color -> VisObject a
Sphere Double
1 Flavour
Wireframe Color
blue
                 , V3 Double -> VisObj -> VisObj
forall a. V3 a -> VisObject a -> VisObject a
Trans (Position -> V3 Double
v3FromPos Position
r) (Double -> Flavour -> Color -> VisObj
forall a. a -> Flavour -> Color -> VisObject a
Sphere Double
0.05 Flavour
Solid Color
red)
                 ]

displayStaticBlochSphere :: Position -> IO ()
displayStaticBlochSphere :: Position -> IO ()
displayStaticBlochSphere Position
r
    = Options -> VisObj -> IO ()
forall b. Real b => Options -> VisObject b -> IO ()
display Options
myOptions (Position -> VisObj
staticBlochSphere Position
r)

-- | Display a qubit state vector as a point on the Bloch Sphere.
displayStaticState :: Vector C -> IO ()
displayStaticState :: Vector C -> IO ()
displayStaticState = Position -> IO ()
displayStaticBlochSphere (Position -> IO ()) -> (Vector C -> Position) -> Vector C -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector C -> Position
toPos

-- | Given a Bloch vector as a function of time,
--   return a 'VisObj' as a function of time.
animatedBlochSphere :: (Double -> Position) -> (Float -> VisObj)
animatedBlochSphere :: (Double -> Position) -> Float -> VisObj
animatedBlochSphere Double -> Position
f
    = Position -> VisObj
staticBlochSphere (Position -> VisObj) -> (Float -> Position) -> Float -> VisObj
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Double -> Position
f (Double -> Position) -> (Float -> Double) -> Float -> Position
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Float -> Double
forall a b. (Real a, Fractional b) => a -> b
realToFrac

-- | Given a sample rate, initial qubit state vector, and
--   state propagation function, produce a simulation.
--   The 'Float' in the state propagation function is the time
--   since the beginning of the simulation.
simulateBlochSphere :: Double -> Vector C -> (Float -> (Float,Vector C) -> (Float,Vector C)) -> IO ()
simulateBlochSphere :: Double
-> Vector C
-> (Float -> (Float, Vector C) -> (Float, Vector C))
-> IO ()
simulateBlochSphere Double
sampleRate Vector C
initial Float -> (Float, Vector C) -> (Float, Vector C)
statePropFunc
    = Options
-> Double
-> (Float, Vector C)
-> ((Float, Vector C) -> VisObj)
-> (Float -> (Float, Vector C) -> (Float, Vector C))
-> IO ()
forall b world.
Real b =>
Options
-> Double
-> world
-> (world -> VisObject b)
-> (Float -> world -> world)
-> IO ()
simulate Options
myOptions Double
sampleRate (Float
0,Vector C
initial) (Position -> VisObj
staticBlochSphere (Position -> VisObj)
-> ((Float, Vector C) -> Position) -> (Float, Vector C) -> VisObj
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector C -> Position
toPos (Vector C -> Position)
-> ((Float, Vector C) -> Vector C) -> (Float, Vector C) -> Position
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Float, Vector C) -> Vector C
forall a b. (a, b) -> b
snd) Float -> (Float, Vector C) -> (Float, Vector C)
statePropFunc

-- | Given a sample rate, initial qubit state ket, and
--   state propagation function, produce a simulation.
--   The 'Float' in the state propagation function is the time
--   since the beginning of the simulation.
simulateBlochSphereK :: Double -> Ket -> (Float -> (Float,Ket) -> (Float,Ket)) -> IO ()
simulateBlochSphereK :: Double -> Ket -> (Float -> (Float, Ket) -> (Float, Ket)) -> IO ()
simulateBlochSphereK Double
sampleRate Ket
initial Float -> (Float, Ket) -> (Float, Ket)
statePropFuncK
    = Options
-> Double
-> (Float, Ket)
-> ((Float, Ket) -> VisObj)
-> (Float -> (Float, Ket) -> (Float, Ket))
-> IO ()
forall b world.
Real b =>
Options
-> Double
-> world
-> (world -> VisObject b)
-> (Float -> world -> world)
-> IO ()
simulate Options
myOptions Double
sampleRate (Float
0,Ket
initial) (Position -> VisObj
staticBlochSphere (Position -> VisObj)
-> ((Float, Ket) -> Position) -> (Float, Ket) -> VisObj
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ket -> Position
ketToPos (Ket -> Position)
-> ((Float, Ket) -> Ket) -> (Float, Ket) -> Position
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Float, Ket) -> Ket
forall a b. (a, b) -> b
snd) Float -> (Float, Ket) -> (Float, Ket)
statePropFuncK

{-
-- | Given a sample rate, initial qubit state vector, and
--   state propagation function, produce a simulation.
--   The 'Float' in the state propagation function is the time
--   since the beginning of the simulation.
playBlochSphere :: Double -> Vector C -> (Float -> (Float,Vector C) -> (Float,Vector C)) -> IO ()
playBlochSphere sampleRate initial statePropFunc
    = play myOptions sampleRate (0,initial) (staticBlochSphere . toPos . snd) statePropFunc
-}

-- | Produce a state propagation function from a time-dependent Hamiltonian.
stateProp :: (Double -> Matrix C) -> Float -> (Float,Vector C) -> (Float,Vector C)
stateProp :: (Double -> Matrix C)
-> Float -> (Float, Vector C) -> (Float, Vector C)
stateProp Double -> Matrix C
ham Float
tNew (Float
tOld,Vector C
v)
    = (Float
tNew, Double -> Matrix C -> Vector C -> Vector C
M.timeEv (Float -> Double
forall a b. (Real a, Fractional b) => a -> b
realToFrac Float
dt) (Double -> Matrix C
ham Double
tMid) Vector C
v)
      where
        dt :: Float
dt = Float
tNew Float -> Float -> Float
forall a. Num a => a -> a -> a
- Float
tOld
        tMid :: Double
tMid = Float -> Double
forall a b. (Real a, Fractional b) => a -> b
realToFrac (Float -> Double) -> Float -> Double
forall a b. (a -> b) -> a -> b
$ (Float
tNew Float -> Float -> Float
forall a. Num a => a -> a -> a
+ Float
tOld) Float -> Float -> Float
forall a. Fractional a => a -> a -> a
/ Float
2

-- | Produce a state propagation function from a time-dependent Hamiltonian.
statePropK :: (Double -> Operator) -> Float -> (Float,Ket) -> (Float,Ket)
statePropK :: (Double -> Operator) -> Float -> (Float, Ket) -> (Float, Ket)
statePropK Double -> Operator
ham Float
tNew (Float
tOld,Ket
psi)
    = (Float
tNew, Double -> Operator -> Ket -> Ket
K.timeEv (Float -> Double
forall a b. (Real a, Fractional b) => a -> b
realToFrac Float
dt) (Double -> Operator
ham Double
tMid) Ket
psi)
      where
        dt :: Float
dt = Float
tNew Float -> Float -> Float
forall a. Num a => a -> a -> a
- Float
tOld
        tMid :: Double
tMid = Float -> Double
forall a b. (Real a, Fractional b) => a -> b
realToFrac (Float -> Double) -> Float -> Double
forall a b. (a -> b) -> a -> b
$ (Float
tNew Float -> Float -> Float
forall a. Num a => a -> a -> a
+ Float
tOld) Float -> Float -> Float
forall a. Fractional a => a -> a -> a
/ Float
2

-- | Given an initial qubit state and a time-dependent Hamiltonian,
--   produce a visualization.
evolutionBlochSphere :: Vector C -> (Double -> Matrix C) -> IO ()
evolutionBlochSphere :: Vector C -> (Double -> Matrix C) -> IO ()
evolutionBlochSphere Vector C
psi0 Double -> Matrix C
ham
    = Double
-> Vector C
-> (Float -> (Float, Vector C) -> (Float, Vector C))
-> IO ()
simulateBlochSphere Double
0.01 Vector C
psi0 ((Double -> Matrix C)
-> Float -> (Float, Vector C) -> (Float, Vector C)
stateProp Double -> Matrix C
ham)

-- | Given an initial qubit ket and a time-dependent Hamiltonian,
--   produce a visualization.
evolutionBlochSphereK :: Ket -> (Double -> Operator) -> IO ()
evolutionBlochSphereK :: Ket -> (Double -> Operator) -> IO ()
evolutionBlochSphereK Ket
psi0 Double -> Operator
ham
    = Double -> Ket -> (Float -> (Float, Ket) -> (Float, Ket)) -> IO ()
simulateBlochSphereK Double
0.01 Ket
psi0 ((Double -> Operator) -> Float -> (Float, Ket) -> (Float, Ket)
statePropK Double -> Operator
ham)

myOptions :: Options
myOptions :: Options
myOptions = Options
defaultOpts {optWindowName = "Bloch Sphere"
                        ,optInitialCamera = Just (Camera0 75 20 4)}

{-
staticBz1 :: IO ()
staticBz1 = evolutionBlochSphere M.xp (const (scale 0.9 M.sz))

staticBz2 :: IO ()
staticBz2 = evolutionBlochSphere ((2|>) [(cos (pi / 8)), (sin (pi / 8))]) (const M.sz)

staticBy1 :: IO ()
staticBy1 = evolutionBlochSphere M.xp (const M.sy)
-}

-- | Hamiltonian for nuclear magnetic resonance.
--   Explain omega0, omegaR, omega.
hamRabi :: Double ->  Double ->  Double ->  Double -> Matrix C
hamRabi :: Double -> Double -> Double -> Double -> Matrix C
hamRabi Double
omega0 Double
omegaR Double
omega Double
t
    = let h11 :: C
h11 = Double
omega0 Double -> Double -> C
forall a. a -> a -> Complex a
:+ Double
0
          h12 :: C
h12 = (Double
omegaR Double -> Double -> C
forall a. a -> a -> Complex a
:+ Double
0) C -> C -> C
forall a. Num a => a -> a -> a
* C -> C
forall a. Floating a => a -> a
exp (-C
iC C -> C -> C
forall a. Num a => a -> a -> a
* ((Double
omega Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
t) Double -> Double -> C
forall a. a -> a -> Complex a
:+ Double
0))
      in C -> Matrix C -> Matrix C
forall t (c :: * -> *). Linear t c => t -> c t -> c t
scale (C
1C -> C -> C
forall a. Fractional a => a -> a -> a
/C
2) (Matrix C -> Matrix C) -> Matrix C -> Matrix C
forall a b. (a -> b) -> a -> b
$ (Int
2Int -> Int -> [C] -> Matrix C
forall a. Storable a => Int -> Int -> [a] -> Matrix a
><Int
2) [C
h11, C
h12, (C -> C
forall a. Num a => Complex a -> Complex a
conjugate C
h12), (-C
h11)]

-- need to scale time

-- a pi pulse