Introduce visualization (#347)

* Add Wigner plotting via CairoMakie

* CairoMakie default lib

* docs

* formatting

* Fix typo

* Add error handling for unavailable plotting libraries in plot_wigner function

* Docs

* Update CairoMakie dependency and add CairoMakie extension tests

* Update documentation to reflect changes in Wigner function plotting and computation

* formatting

* Update changelog

* Improve docs

* Update CairoMakie imports and change axis labels

* fix CHANGELOG

* setup CI pipeline config for `CairoMakie` extension

* rebase accidently removed lines

* fix typo

* add `CairoMakie` extension doc page

* fix typo

* minor change in docs

* minor change in docs

---------

Co-authored-by: Lorenzo Fioroni <lorenzo.fioroni@epfl.ch>
Co-authored-by: Alberto Mercurio <alberto.mercurio96@gmail.com>

Author: 3 people

.github/workflows/CI.yml

@@ -10,6 +10,7 @@ on:
       - 'ext/**'
       - 'test/runtests.jl'
       - 'test/core-test/**'
+      - 'test/ext-test/cairomakie_ext.jl'
       - 'Project.toml'
   pull_request:
     branches:
@@ -20,6 +21,7 @@ on:
       - 'ext/**'
       - 'test/runtests.jl'
       - 'test/core-test/**'
+      - 'test/ext-test/cairomakie_ext.jl'
       - 'Project.toml'
     types:
       - opened
@@ -52,14 +54,21 @@ jobs:
         group:
           - 'Core'
 
-        # include:
+        include:
           # for core tests (intermediate versions)
           # - version: '1.x'
           #   node:
           #     os: 'ubuntu-latest'
           #     arch: 'x64'
           #   group: 'Core'
 
+          # for extension tests
+          - version: '1'
+            node:
+              os: 'ubuntu-latest'
+              arch: 'x64'
+            group: 'CairoMakie_Ext'
+
     steps:
       - uses: actions/checkout@v4
       - uses: julia-actions/setup-julia@v2

CHANGELOG.md

@@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Improve the construction of `QobjEvo`. ([#338], [#339])
 - Support `Base.zero` and `Base.one` for `AbstractQuantumObject`. ([#342], [#346])
+- Introduce visualization and function `plot_wigner` for easy plotting of Wigner functions. ([#86], [#292], [#347])
 
 ## [v0.23.1]
 Release date: 2024-12-06
@@ -45,7 +46,9 @@ Release date: 2024-11-13
 [v0.22.0]: https://github.com/qutip/QuantumToolbox.jl/releases/tag/v0.22.0
 [v0.23.0]: https://github.com/qutip/QuantumToolbox.jl/releases/tag/v0.23.0
 [v0.23.1]: https://github.com/qutip/QuantumToolbox.jl/releases/tag/v0.23.1
+[#86]: https://github.com/qutip/QuantumToolbox.jl/issues/86
 [#139]: https://github.com/qutip/QuantumToolbox.jl/issues/139
+[#292]: https://github.com/qutip/QuantumToolbox.jl/issues/292
 [#306]: https://github.com/qutip/QuantumToolbox.jl/issues/306
 [#309]: https://github.com/qutip/QuantumToolbox.jl/issues/309
 [#311]: https://github.com/qutip/QuantumToolbox.jl/issues/311
@@ -58,3 +61,4 @@ Release date: 2024-11-13
 [#339]: https://github.com/qutip/QuantumToolbox.jl/issues/339
 [#342]: https://github.com/qutip/QuantumToolbox.jl/issues/342
 [#346]: https://github.com/qutip/QuantumToolbox.jl/issues/346
+[#347]: https://github.com/qutip/QuantumToolbox.jl/issues/347

Project.toml

@@ -28,14 +28,17 @@ StochasticDiffEq = "789caeaf-c7a9-5a7d-9973-96adeb23e2a0"
 
 [weakdeps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
 
 [extensions]
 QuantumToolboxCUDAExt = "CUDA"
+QuantumToolboxCairoMakieExt = "CairoMakie"
 
 [compat]
 Aqua = "0.8"
 ArrayInterface = "6, 7"
 CUDA = "5"
+CairoMakie = "0.12"
 DiffEqBase = "6"
 DiffEqCallbacks = "4.2.1 - 4"
 DiffEqNoiseProcess = "5"
@@ -62,8 +65,9 @@ julia = "1.10"
 
 [extras]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
+CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Aqua", "JET", "Test"]
+test = ["Aqua", "CairoMakie", "JET", "Test"]

docs/make.jl

@@ -7,6 +7,9 @@ using DocumenterVitepress
 using DocumenterCitations
 using Changelog
 
+# Load of packages required to compile the extension documentation
+using CairoMakie
+
 DocMeta.setdocmeta!(QuantumToolbox, :DocTestSetup, :(using QuantumToolbox); recursive = true)
 
 # some options for `makedocs`
@@ -56,6 +59,7 @@ const PAGES = [
         "Two-time correlation functions" => "users_guide/two_time_corr_func.md",
         "Extensions" => [
             "users_guide/extensions/cuda.md",
+            "users_guide/extensions/cairomakie.md",
         ],
     ],
     "Tutorials" => [
@@ -76,7 +80,10 @@ const PAGES = [
 ]
 
 makedocs(;
-    modules = [QuantumToolbox],
+    modules = [
+        QuantumToolbox, 
+        Base.get_extension(QuantumToolbox, :QuantumToolboxCairoMakieExt),
+    ],
     authors = "Alberto Mercurio and Yi-Te Huang",
     repo = Remotes.GitHub("qutip", "QuantumToolbox.jl"),
     sitename = "QuantumToolbox.jl",

docs/src/resources/api.md

@@ -283,3 +283,9 @@ convert_unit
 row_major_reshape
 meshgrid
 ```
+
+## [Visualization](@id doc-API:Visualization)
+
+```@docs
+plot_wigner
+```

docs/src/tutorials/logo.md

@@ -67,25 +67,16 @@ Next, we construct the triangular cat state as a normalized superposition of thr
 normalize!(ψ)
 ```
 
-### Defining the Grid and calculating the Wigner function
+### Defining the Grid and plotting the Wigner function
 
-We define the grid for the Wigner function and calculate it using the [`wigner`](@ref) function. We shift the grid in the imaginary direction to ensure that the Wigner function is centered around the origin of the figure. The [`wigner`](@ref) function also supports the `g` scaling factor, which we put here equal to ``2``.
+We define the grid for the Wigner function and plot it using the [`plot_wigner`](@ref) function. This, internally calls the [`wigner`](@ref) function for the computation. We shift the grid in the imaginary direction to ensure that the Wigner function is centered around the origin of the figure. The [`wigner`](@ref) function also supports the `g` scaling factor, which we put here equal to ``2``.
 
 ```@example logo
 xvec = range(-ρ, ρ, 500) .* 1.5
 yvec = xvec .+ (abs(imag(α1)) - abs(imag(α2))) / 2
 
-wig = wigner(ψ, xvec, yvec, g = 2)
-```
-
-### Plotting the Wigner function
-
-Finally, we plot the Wigner function using the `heatmap` function from the `CairoMakie` package.
-
-```@example logo
 fig = Figure(size = (250, 250), figure_padding = 0)
-ax = Axis(fig[1, 1])
-heatmap!(ax, xvec, yvec, wig', colormap = :RdBu, interpolate = true, rasterize = 1)
+fig, ax, hm = plot_wigner(ψ, xvec = xvec, yvec = yvec, g = 2, library = Val(:CairoMakie), location = fig[1,1])
 hidespines!(ax)
 hidexdecorations!(ax)
 hideydecorations!(ax)
@@ -118,12 +109,8 @@ nothing # hide
 And the Wigner function becomes more uniform:
 
 ```@example logo
-wig = wigner(sol.states[end], xvec, yvec, g = 2)
-
 fig = Figure(size = (250, 250), figure_padding = 0)
-ax = Axis(fig[1, 1])
-
-img_wig = heatmap!(ax, xvec, yvec, wig', colormap = :RdBu, interpolate = true, rasterize = 1)
+fig, ax, hm = plot_wigner(sol.states[end], xvec = xvec, yvec = yvec, g = 2, library = Val(:CairoMakie), location = fig[1,1])
 hidespines!(ax)
 hidexdecorations!(ax)
 hideydecorations!(ax)
@@ -135,7 +122,7 @@ At this stage, we have finished to use the `QuantumToolbox` package. From now on
 
 ### Custom Colormap
 
-We define a custom colormap that changes depending on the Wigner function and spatial coordinates. Indeed, we want the three different colormaps, in the regions corresponding to the three coherent states, to match the colors of the Julia logo. We also want the colormap change to be smooth, so we use a Gaussian function to blend the colors. We introduce also a Wigner function dependent transparency to make the logo more appealing.
+We define a custom colormap that changes depending on the Wigner function and spatial coordinates. Indeed, we want the three different colormaps, in the regions corresponding to the three coherent states, to match the colors of the Julia logo. We also want the colormap change to be smooth, so we use a Gaussian function to blend the colors. We introduce also a Wigner function dependent transparency to make the logo more appealing. In order to do so, we are going to need the value of the wigner function at each point of the grid, rather than its plot. We will thus call the [`wigner`](@ref) function directly.
 
 ```@example logo
 function set_color_julia(x, y, wig::T, α1, α2, α3, cmap1, cmap2, cmap3, δ) where {T}
@@ -156,6 +143,7 @@ function set_color_julia(x, y, wig::T, α1, α2, α3, cmap1, cmap2, cmap3, δ) w
     return RGBAf(c_tot.r, c_tot.g, c_tot.b, alpha)
 end
 
+wig = wigner(sol.states[end], xvec, yvec, g = 2)
 X, Y = meshgrid(xvec, yvec)
 δ = 1.25 # Smoothing parameter for the Gaussian functions
 ```

docs/src/users_guide/extensions/cairomakie.md

@@ -0,0 +1,21 @@
+# [Extension for CairoMakie.jl](@id doc:CairoMakie)
+
+This is an extension to support visualization (plotting functions) using [`CairoMakie.jl`](https://github.com/MakieOrg/Makie.jl/tree/master/CairoMakie) library.
+
+This extension will be automatically loaded if user imports both `QuantumToolbox.jl` and [`CairoMakie.jl`](https://github.com/MakieOrg/Makie.jl/tree/master/CairoMakie):
+
+```julia
+using QuantumToolbox
+using CairoMakie
+```
+
+To plot with [`CairoMakie.jl`](https://github.com/MakieOrg/Makie.jl/tree/master/CairoMakie) library, specify the keyword argument `library = Val(:CairoMakie)` for the plotting functions.
+
+!!! warning "Beware of type-stability!"
+    If you want to keep type stability, it is recommended to use `Val(:CairoMakie)` instead of `: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.
+
+The supported plotting functions are listed as follows:
+
+| **Plotting Function** | **Description** |
+|:----------------------|:----------------|
+| [`plot_wigner`](@ref) | [Wigner quasipropability distribution](https://en.wikipedia.org/wiki/Wigner_quasiprobability_distribution) |