@@ -5,31 +5,31 @@ using CairoMakie
55
66function QuantumToolbox. plot_wigner (
77 library:: Val{:CairoMakie} ,
8- state:: QuantumObject{<:AbstractArray{T},OpType} ,
8+ state:: QuantumObject{<:AbstractArray{T},OpType} ;
99 xvec:: Union{Nothing,AbstractVector} = nothing ,
10- yvec:: Union{Nothing,AbstractVector} = nothing ;
10+ yvec:: Union{Nothing,AbstractVector} = nothing ,
1111 g:: Real = √ 2 ,
1212 method:: WignerSolver = WignerClenshaw (),
13- projection:: String = " 2d" ,
14- fig:: Union{Figure,Nothing} = nothing ,
15- ax:: Union{Axis,Nothing} = nothing ,
13+ projection:: Union{Val,Symbol} = Val (:two_dim ),
14+ location:: Union{GridPosition,Nothing} = nothing ,
1615 colorbar:: Bool = false ,
1716 kwargs... ,
1817) where {T,OpType<: Union{BraQuantumObject,KetQuantumObject,OperatorQuantumObject} }
19- projection == " 2d" || projection == " 3d" || throw (ArgumentError (" Unsupported projection: $projection " ))
18+ QuantumToolbox. getVal (projection) == :two_dim ||
19+ QuantumToolbox. getVal (projection) == :three_dim ||
20+ throw (ArgumentError (" Unsupported projection: $projection " ))
2021
2122 return _plot_wigner (
2223 library,
2324 state,
2425 xvec,
2526 yvec,
26- Val ( Symbol ( projection) ),
27+ QuantumToolbox . makeVal ( projection),
2728 g,
2829 method,
29- fig,
30- ax,
30+ location,
3131 colorbar;
32- kwargs...
32+ kwargs... ,
3333 )
3434end
3535
@@ -38,20 +38,17 @@ function _plot_wigner(
3838 state:: QuantumObject{<:AbstractArray{T},OpType} ,
3939 xvec:: AbstractVector ,
4040 yvec:: AbstractVector ,
41- projection:: Val{Symbol("2d") } ,
41+ projection:: Val{:two_dim } ,
4242 g:: Real ,
4343 method:: WignerSolver ,
44- fig:: Union{Figure,Nothing} ,
45- ax:: Union{Axis,Nothing} ,
44+ location:: Union{GridPosition,Nothing} ,
4645 colorbar:: Bool ;
4746 kwargs... ,
4847) where {T,OpType<: Union{BraQuantumObject,KetQuantumObject,OperatorQuantumObject} }
49- fig, ax = _getFigAx (fig, ax )
48+ fig, location = _getFigAndLocation (location )
5049
51- gridPos = _gridPosFromAx (ax)
52- CairoMakie. delete! (ax)
50+ lyt = GridLayout (location)
5351
54- lyt = GridLayout (gridPos)
5552 ax = Axis (lyt[1 , 1 ])
5653
5754 wig = wigner (state, xvec, yvec; g = g, method = method)
@@ -74,20 +71,17 @@ function _plot_wigner(
7471 state:: QuantumObject{<:AbstractArray{T},OpType} ,
7572 xvec:: AbstractVector ,
7673 yvec:: AbstractVector ,
77- projection:: Val{Symbol("3d") } ,
74+ projection:: Val{:three_dim } ,
7875 g:: Real ,
7976 method:: WignerSolver ,
80- fig:: Union{Figure,Nothing} ,
81- ax:: Union{Axis,Nothing} ,
77+ location:: Union{GridPosition,Nothing} ,
8278 colorbar:: Bool ;
8379 kwargs... ,
8480) where {T,OpType<: Union{BraQuantumObject,KetQuantumObject,OperatorQuantumObject} }
85- fig, ax = _getFigAx (fig, ax )
81+ fig, location = _getFigAndLocation (location )
8682
87- gridPos = _gridPosFromAx (ax)
88- CairoMakie. delete! (ax)
83+ lyt = GridLayout (location)
8984
90- lyt = GridLayout (gridPos)
9185 ax = Axis3 (lyt[1 , 1 ], azimuth = 1.775pi , elevation = pi / 16 , protrusions = (30 , 90 , 30 , 30 ), viewmode = :stretch )
9286
9387 wig = wigner (state, xvec, yvec; g = g, method = method)
@@ -106,22 +100,17 @@ function _plot_wigner(
106100 return fig, ax, surf
107101end
108102
109- _getFigAx (fig:: Figure , ax:: Axis ) = fig, ax
110- _getFigAx (fig:: Figure , :: Nothing ) = fig, Axis (fig[1 , 1 ])
111- _getFigAx (:: Nothing , ax:: Axis ) = _figFromChildren (ax), ax
112- function _getFigAx (:: Nothing , :: Nothing )
103+ function _getFigAndLocation (location:: Nothing )
113104 fig = Figure ()
114- ax = Axis (fig[1 , 1 ])
115- return fig, ax
105+ return fig, fig[1 , 1 ]
106+ end
107+ function _getFigAndLocation (location:: GridPosition )
108+ fig = _figFromChildren (location. layout)
109+ return fig, location
116110end
117111
118112_figFromChildren (children) = _figFromChildren (children. parent)
119113_figFromChildren (fig:: Figure ) = fig
120-
121- function _gridPosFromAx (ax:: Axis )
122- content = CairoMakie. Makie. GridLayoutBase. gridcontent (ax)
123- gl, sp, si = content. parent, content. span, content. side
124- return GridPosition (gl, sp, si)
125- end
114+ _figFromChildren (:: Nothing ) = throw (ArgumentError (" No Figure has been found at the top of the layout hierarchy." ))
126115
127116end
0 commit comments