Add Wigner plotting via CairoMakie #292
+349
−20
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,215 @@ | ||
| module QuantumToolboxCairoMakieExt | ||
|
|
||
| using QuantumToolbox | ||
| using CairoMakie: Axis, Axis3, Colorbar, Figure, GridLayout, heatmap!, surface!, GridPosition, @L_str, Reverse | ||
|
|
||
| @doc raw""" | ||
| plot_wigner( | ||
| library::Val{:CairoMakie}, | ||
| state::QuantumObject{DT,OpType}; | ||
| xvec::Union{Nothing,AbstractVector} = nothing, | ||
| yvec::Union{Nothing,AbstractVector} = nothing, | ||
| g::Real = √2, | ||
| method::WignerSolver = WignerClenshaw(), | ||
| projection::Union{Val,Symbol} = Val(:two_dim), | ||
| location::Union{GridPosition,Nothing} = nothing, | ||
| colorbar::Bool = false, | ||
| kwargs... | ||
| ) where {DT,OpType} | ||
|
|
||
| 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. | ||
|
|
||
| # Arguments | ||
| - `library::Val{:CairoMakie}`: The plotting library to use. | ||
| - `state::QuantumObject`: The quantum state for which the Wigner function is calculated. It can be either a [`Ket`](@ref), [`Bra`](@ref), or [`Operator`](@ref). | ||
| - `xvec::AbstractVector`: The x-coordinates of the phase space grid. Defaults to a linear range from -7.5 to 7.5 with 200 points. | ||
| - `yvec::AbstractVector`: The y-coordinates of the phase space grid. Defaults to a linear range from -7.5 to 7.5 with 200 points. | ||
| - `g::Real`: The scaling factor related to the value of ``\hbar`` in the commutation relation ``[x, y] = i \hbar`` via ``\hbar=2/g^2``. | ||
| - `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. | ||
| - `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. | ||
| - `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`. | ||
| - `colorbar::Bool`: Whether to include a colorbar in the plot. Default is `false`. | ||
| - `kwargs...`: Additional keyword arguments to pass to the plotting function. | ||
|
|
||
| # Returns | ||
| - `fig`: The figure object. | ||
| - `ax`: The axis object. | ||
| - `hm`: Either the heatmap or surface object, depending on the projection. | ||
|
|
||
| !!! note "Import library first" | ||
| [`CairoMakie`](https://github.com/MakieOrg/Makie.jl/tree/master/CairoMakie) must first be imported before using this function. | ||
|
|
||
| !!! warning "Beware of type-stability!" | ||
| 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. | ||
ytdHuang marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| """ | ||
| function QuantumToolbox.plot_wigner( | ||
| library::Val{:CairoMakie}, | ||
| state::QuantumObject{DT,OpType}; | ||
| xvec::Union{Nothing,AbstractVector} = LinRange(-7.5, 7.5, 200), | ||
| yvec::Union{Nothing,AbstractVector} = LinRange(-7.5, 7.5, 200), | ||
| g::Real = √2, | ||
| method::WignerSolver = WignerClenshaw(), | ||
| projection::Union{Val,Symbol} = Val(:two_dim), | ||
ytdHuang marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| location::Union{GridPosition,Nothing} = nothing, | ||
| colorbar::Bool = false, | ||
| kwargs..., | ||
| ) where {DT,OpType<:Union{BraQuantumObject,KetQuantumObject,OperatorQuantumObject}} | ||
| QuantumToolbox.getVal(projection) == :two_dim || | ||
| QuantumToolbox.getVal(projection) == :three_dim || | ||
ytdHuang marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| throw(ArgumentError("Unsupported projection: $projection")) | ||
|
|
||
| return _plot_wigner( | ||
| library, | ||
| state, | ||
| xvec, | ||
| yvec, | ||
| QuantumToolbox.makeVal(projection), | ||
| g, | ||
| method, | ||
| location, | ||
| colorbar; | ||
| kwargs..., | ||
| ) | ||
| end | ||
|
|
||
| function _plot_wigner( | ||
| ::Val{:CairoMakie}, | ||
| state::QuantumObject{DT,OpType}, | ||
| xvec::AbstractVector, | ||
| yvec::AbstractVector, | ||
| projection::Val{:two_dim}, | ||
| g::Real, | ||
| method::WignerSolver, | ||
| location::Union{GridPosition,Nothing}, | ||
| colorbar::Bool; | ||
| kwargs..., | ||
| ) where {DT,OpType<:Union{BraQuantumObject,KetQuantumObject,OperatorQuantumObject}} | ||
| fig, location = _getFigAndLocation(location) | ||
|
|
||
| lyt = GridLayout(location) | ||
|
|
||
| ax = Axis(lyt[1, 1]) | ||
|
|
||
| wig = wigner(state, xvec, yvec; g = g, method = method) | ||
| wlim = maximum(abs, wig) | ||
|
|
||
| kwargs = merge(Dict(:colormap => Reverse(:RdBu), :colorrange => (-wlim, wlim)), kwargs) | ||
| hm = heatmap!(ax, xvec, yvec, wig'; kwargs...) | ||
ytdHuang marked this conversation as resolved.
Show resolved
Hide resolved
|
||
|
|
||
| if colorbar | ||
| Colorbar(lyt[1, 2], hm) | ||
| end | ||
|
|
||
| ax.xlabel = L"\textrm{Re}(\alpha)" | ||
| ax.ylabel = L"\textrm{Im}(\alpha)" | ||
| return fig, ax, hm | ||
| end | ||
|
|
||
| function _plot_wigner( | ||
| ::Val{:CairoMakie}, | ||
| state::QuantumObject{DT,OpType}, | ||
| xvec::AbstractVector, | ||
| yvec::AbstractVector, | ||
| projection::Val{:three_dim}, | ||
| g::Real, | ||
| method::WignerSolver, | ||
| location::Union{GridPosition,Nothing}, | ||
| colorbar::Bool; | ||
| kwargs..., | ||
| ) where {DT,OpType<:Union{BraQuantumObject,KetQuantumObject,OperatorQuantumObject}} | ||
| fig, location = _getFigAndLocation(location) | ||
|
|
||
| lyt = GridLayout(location) | ||
|
|
||
| ax = Axis3(lyt[1, 1], azimuth = 1.775pi, elevation = pi / 16, protrusions = (30, 90, 30, 30), viewmode = :stretch) | ||
|
|
||
| wig = wigner(state, xvec, yvec; g = g, method = method) | ||
| wlim = maximum(abs, wig) | ||
|
|
||
| kwargs = merge(Dict(:colormap => :RdBu, :colorrange => (-wlim, wlim)), kwargs) | ||
| surf = surface!(ax, xvec, yvec, wig'; kwargs...) | ||
ytdHuang marked this conversation as resolved.
Show resolved
Hide resolved
|
||
|
|
||
| if colorbar | ||
| Colorbar(lyt[1, 2], surf) | ||
| end | ||
|
|
||
| ax.xlabel = L"\textrm{Re}(\alpha)" | ||
| ax.ylabel = L"\textrm{Im}(\alpha)" | ||
| ax.zlabel = "Wigner function" | ||
| return fig, ax, surf | ||
| end | ||
|
|
||
| raw""" | ||
| _getFigAndLocation(location::Nothing) | ||
|
|
||
| Create a new figure and return it, together with the GridPosition object pointing to the first cell. | ||
|
|
||
| # Arguments | ||
| - `location::Nothing` | ||
|
|
||
| # Returns | ||
| - `fig`: The figure object. | ||
| - `location`: The GridPosition object pointing to the first cell. | ||
| """ | ||
| function _getFigAndLocation(location::Nothing) | ||
| fig = Figure() | ||
| return fig, fig[1, 1] | ||
| end | ||
|
|
||
| raw""" | ||
| _getFigAndLocation(location::GridPosition) | ||
|
|
||
| Compute which figure does the location belong to and return it, together with the location itself. | ||
|
|
||
| # Arguments | ||
| - `location::GridPosition` | ||
|
|
||
| # Returns | ||
| - `fig`: The figure object. | ||
| - `location`: The GridPosition object. | ||
| """ | ||
| function _getFigAndLocation(location::GridPosition) | ||
| fig = _figFromChildren(location.layout) | ||
| return fig, location | ||
| end | ||
|
|
||
| raw""" | ||
| _figFromChildren(children::GridLayout) | ||
|
|
||
| Recursively find the figure object from the children layout. | ||
|
|
||
| # Arguments | ||
| - `children::GridLayout` | ||
|
|
||
| # Returns | ||
| - Union{Nothing, Figure, GridLayout}: The children's parent object. | ||
| """ | ||
| _figFromChildren(children) = _figFromChildren(children.parent) | ||
|
|
||
| raw""" | ||
| _figFromChildren(fig::Figure) | ||
|
|
||
| Return the figure object | ||
|
|
||
| # Arguments | ||
| - `fig::Figure` | ||
|
|
||
| # Returns | ||
| - `fig`: The figure object. | ||
| """ | ||
| _figFromChildren(fig::Figure) = fig | ||
|
|
||
| raw""" | ||
| _figFromChildren(::Nothing) | ||
|
|
||
| Throw an error if no figure has been found. | ||
|
|
||
| # Arguments | ||
| - `::Nothing` | ||
|
|
||
| # Throws | ||
| - `ArgumentError`: If no figure has been found. | ||
| """ | ||
| _figFromChildren(::Nothing) = throw(ArgumentError("No Figure has been found at the top of the layout hierarchy.")) | ||
|
|
||
| end | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.