|
| 1 | +module QuantumToolboxCairoMakieExt |
| 2 | + |
| 3 | +using QuantumToolbox |
| 4 | +using CairoMakie: Axis, Axis3, Colorbar, Figure, GridLayout, heatmap!, surface!, GridPosition, @L_str, Reverse |
| 5 | + |
| 6 | +@doc raw""" |
| 7 | + plot_wigner( |
| 8 | + library::Val{:CairoMakie}, |
| 9 | + state::QuantumObject{DT,OpType}; |
| 10 | + xvec::Union{Nothing,AbstractVector} = nothing, |
| 11 | + yvec::Union{Nothing,AbstractVector} = nothing, |
| 12 | + g::Real = √2, |
| 13 | + method::WignerSolver = WignerClenshaw(), |
| 14 | + projection::Union{Val,Symbol} = Val(:two_dim), |
| 15 | + location::Union{GridPosition,Nothing} = nothing, |
| 16 | + colorbar::Bool = false, |
| 17 | + kwargs... |
| 18 | + ) where {DT,OpType} |
| 19 | +
|
| 20 | +Plot the [Wigner quasipropability distribution](https://en.wikipedia.org/wiki/Wigner_quasiprobability_distribution) of `state` using the [`CairoMakie`](https://github.com/MakieOrg/Makie.jl/tree/master/CairoMakie) plotting library. |
| 21 | +
|
| 22 | +# Arguments |
| 23 | +- `library::Val{:CairoMakie}`: The plotting library to use. |
| 24 | +- `state::QuantumObject`: The quantum state for which the Wigner function is calculated. It can be either a [`Ket`](@ref), [`Bra`](@ref), or [`Operator`](@ref). |
| 25 | +- `xvec::AbstractVector`: The x-coordinates of the phase space grid. Defaults to a linear range from -7.5 to 7.5 with 200 points. |
| 26 | +- `yvec::AbstractVector`: The y-coordinates of the phase space grid. Defaults to a linear range from -7.5 to 7.5 with 200 points. |
| 27 | +- `g::Real`: The scaling factor related to the value of ``\hbar`` in the commutation relation ``[x, y] = i \hbar`` via ``\hbar=2/g^2``. |
| 28 | +- `method::WignerSolver`: The method used to calculate the Wigner function. It can be either `WignerLaguerre()` or `WignerClenshaw()`, with `WignerClenshaw()` as default. The `WignerLaguerre` method has the optional `parallel` and `tol` parameters, with default values `true` and `1e-14`, respectively. |
| 29 | +- `projection::Union{Val,Symbol}`: Whether to plot the Wigner function in 2D or 3D. It can be either `Val(:two_dim)` or `Val(:three_dim)`, with `Val(:two_dim)` as default. |
| 30 | +- `location::Union{GridPosition,Nothing}`: The location of the plot in the layout. If `nothing`, the plot is created in a new figure. Default is `nothing`. |
| 31 | +- `colorbar::Bool`: Whether to include a colorbar in the plot. Default is `false`. |
| 32 | +- `kwargs...`: Additional keyword arguments to pass to the plotting function. |
| 33 | +
|
| 34 | +# Returns |
| 35 | +- `fig`: The figure object. |
| 36 | +- `ax`: The axis object. |
| 37 | +- `hm`: Either the heatmap or surface object, depending on the projection. |
| 38 | +
|
| 39 | +!!! note "Import library first" |
| 40 | + [`CairoMakie`](https://github.com/MakieOrg/Makie.jl/tree/master/CairoMakie) must first be imported before using this function. |
| 41 | +
|
| 42 | +!!! warning "Beware of type-stability!" |
| 43 | + If you want to keep type stability, it is recommended to use `Val(:two_dim)` and `Val(:three_dim)` instead of `:two_dim` and `:three_dim`, respectively. Also, specify the library as `Val(:CairoMakie)` See [this link](https://docs.julialang.org/en/v1/manual/performance-tips/#man-performance-value-type) and the [related Section](@ref doc:Type-Stability) about type stability for more details. |
| 44 | +""" |
| 45 | +function QuantumToolbox.plot_wigner( |
| 46 | + library::Val{:CairoMakie}, |
| 47 | + state::QuantumObject{DT,OpType}; |
| 48 | + xvec::Union{Nothing,AbstractVector} = LinRange(-7.5, 7.5, 200), |
| 49 | + yvec::Union{Nothing,AbstractVector} = LinRange(-7.5, 7.5, 200), |
| 50 | + g::Real = √2, |
| 51 | + method::WignerSolver = WignerClenshaw(), |
| 52 | + projection::Union{Val,Symbol} = Val(:two_dim), |
| 53 | + location::Union{GridPosition,Nothing} = nothing, |
| 54 | + colorbar::Bool = false, |
| 55 | + kwargs..., |
| 56 | +) where {DT,OpType<:Union{BraQuantumObject,KetQuantumObject,OperatorQuantumObject}} |
| 57 | + QuantumToolbox.getVal(projection) == :two_dim || |
| 58 | + QuantumToolbox.getVal(projection) == :three_dim || |
| 59 | + throw(ArgumentError("Unsupported projection: $projection")) |
| 60 | + |
| 61 | + return _plot_wigner( |
| 62 | + library, |
| 63 | + state, |
| 64 | + xvec, |
| 65 | + yvec, |
| 66 | + QuantumToolbox.makeVal(projection), |
| 67 | + g, |
| 68 | + method, |
| 69 | + location, |
| 70 | + colorbar; |
| 71 | + kwargs..., |
| 72 | + ) |
| 73 | +end |
| 74 | + |
| 75 | +function _plot_wigner( |
| 76 | + ::Val{:CairoMakie}, |
| 77 | + state::QuantumObject{DT,OpType}, |
| 78 | + xvec::AbstractVector, |
| 79 | + yvec::AbstractVector, |
| 80 | + projection::Val{:two_dim}, |
| 81 | + g::Real, |
| 82 | + method::WignerSolver, |
| 83 | + location::Union{GridPosition,Nothing}, |
| 84 | + colorbar::Bool; |
| 85 | + kwargs..., |
| 86 | +) where {DT,OpType<:Union{BraQuantumObject,KetQuantumObject,OperatorQuantumObject}} |
| 87 | + fig, location = _getFigAndLocation(location) |
| 88 | + |
| 89 | + lyt = GridLayout(location) |
| 90 | + |
| 91 | + ax = Axis(lyt[1, 1]) |
| 92 | + |
| 93 | + wig = wigner(state, xvec, yvec; g = g, method = method) |
| 94 | + wlim = maximum(abs, wig) |
| 95 | + |
| 96 | + kwargs = merge(Dict(:colormap => Reverse(:RdBu), :colorrange => (-wlim, wlim)), kwargs) |
| 97 | + hm = heatmap!(ax, xvec, yvec, wig'; kwargs...) |
| 98 | + |
| 99 | + if colorbar |
| 100 | + Colorbar(lyt[1, 2], hm) |
| 101 | + end |
| 102 | + |
| 103 | + ax.xlabel = L"\textrm{Re}(\alpha)" |
| 104 | + ax.ylabel = L"\textrm{Im}(\alpha)" |
| 105 | + return fig, ax, hm |
| 106 | +end |
| 107 | + |
| 108 | +function _plot_wigner( |
| 109 | + ::Val{:CairoMakie}, |
| 110 | + state::QuantumObject{DT,OpType}, |
| 111 | + xvec::AbstractVector, |
| 112 | + yvec::AbstractVector, |
| 113 | + projection::Val{:three_dim}, |
| 114 | + g::Real, |
| 115 | + method::WignerSolver, |
| 116 | + location::Union{GridPosition,Nothing}, |
| 117 | + colorbar::Bool; |
| 118 | + kwargs..., |
| 119 | +) where {DT,OpType<:Union{BraQuantumObject,KetQuantumObject,OperatorQuantumObject}} |
| 120 | + fig, location = _getFigAndLocation(location) |
| 121 | + |
| 122 | + lyt = GridLayout(location) |
| 123 | + |
| 124 | + ax = Axis3(lyt[1, 1], azimuth = 1.775pi, elevation = pi / 16, protrusions = (30, 90, 30, 30), viewmode = :stretch) |
| 125 | + |
| 126 | + wig = wigner(state, xvec, yvec; g = g, method = method) |
| 127 | + wlim = maximum(abs, wig) |
| 128 | + |
| 129 | + kwargs = merge(Dict(:colormap => :RdBu, :colorrange => (-wlim, wlim)), kwargs) |
| 130 | + surf = surface!(ax, xvec, yvec, wig'; kwargs...) |
| 131 | + |
| 132 | + if colorbar |
| 133 | + Colorbar(lyt[1, 2], surf) |
| 134 | + end |
| 135 | + |
| 136 | + ax.xlabel = L"\textrm{Re}(\alpha)" |
| 137 | + ax.ylabel = L"\textrm{Im}(\alpha)" |
| 138 | + ax.zlabel = "Wigner function" |
| 139 | + return fig, ax, surf |
| 140 | +end |
| 141 | + |
| 142 | +raw""" |
| 143 | + _getFigAndLocation(location::Nothing) |
| 144 | + |
| 145 | + Create a new figure and return it, together with the GridPosition object pointing to the first cell. |
| 146 | +
|
| 147 | + # Arguments |
| 148 | + - `location::Nothing` |
| 149 | +
|
| 150 | + # Returns |
| 151 | + - `fig`: The figure object. |
| 152 | + - `location`: The GridPosition object pointing to the first cell. |
| 153 | +""" |
| 154 | +function _getFigAndLocation(location::Nothing) |
| 155 | + fig = Figure() |
| 156 | + return fig, fig[1, 1] |
| 157 | +end |
| 158 | + |
| 159 | +raw""" |
| 160 | + _getFigAndLocation(location::GridPosition) |
| 161 | + |
| 162 | + Compute which figure does the location belong to and return it, together with the location itself. |
| 163 | +
|
| 164 | + # Arguments |
| 165 | + - `location::GridPosition` |
| 166 | +
|
| 167 | + # Returns |
| 168 | + - `fig`: The figure object. |
| 169 | + - `location`: The GridPosition object. |
| 170 | +""" |
| 171 | +function _getFigAndLocation(location::GridPosition) |
| 172 | + fig = _figFromChildren(location.layout) |
| 173 | + return fig, location |
| 174 | +end |
| 175 | + |
| 176 | +raw""" |
| 177 | + _figFromChildren(children::GridLayout) |
| 178 | +
|
| 179 | + Recursively find the figure object from the children layout. |
| 180 | +
|
| 181 | + # Arguments |
| 182 | + - `children::GridLayout` |
| 183 | +
|
| 184 | + # Returns |
| 185 | + - Union{Nothing, Figure, GridLayout}: The children's parent object. |
| 186 | +""" |
| 187 | +_figFromChildren(children) = _figFromChildren(children.parent) |
| 188 | + |
| 189 | +raw""" |
| 190 | + _figFromChildren(fig::Figure) |
| 191 | +
|
| 192 | + Return the figure object |
| 193 | +
|
| 194 | + # Arguments |
| 195 | + - `fig::Figure` |
| 196 | +
|
| 197 | + # Returns |
| 198 | + - `fig`: The figure object. |
| 199 | +""" |
| 200 | +_figFromChildren(fig::Figure) = fig |
| 201 | + |
| 202 | +raw""" |
| 203 | + _figFromChildren(::Nothing) |
| 204 | +
|
| 205 | + Throw an error if no figure has been found. |
| 206 | +
|
| 207 | + # Arguments |
| 208 | + - `::Nothing` |
| 209 | +
|
| 210 | + # Throws |
| 211 | + - `ArgumentError`: If no figure has been found. |
| 212 | +""" |
| 213 | +_figFromChildren(::Nothing) = throw(ArgumentError("No Figure has been found at the top of the layout hierarchy.")) |
| 214 | + |
| 215 | +end |
0 commit comments