Skip to content

Commit

Permalink
improved plotting
Browse files Browse the repository at this point in the history
  • Loading branch information
@orso82
orso82 committed Dec 16, 2024
1 parent e9257e3 commit 8065bd5
Showing 1 changed file with 84 additions and 58 deletions.
142 changes: 84 additions & 58 deletions src/plot.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -131,8 +131,9 @@ push!(document[:Plot], :help_plot!)
# ========= #
# pf_active #
# ========= #
@recipe function plot_pf_active_cx(pfa::pf_active, what::Symbol=:cx; time0=global_time(pfa), cname=:vik)
@recipe function plot_pf_active_cx(pfa::pf_active; what=:cx, time0=global_time(pfa), cname=:vik)
id = plot_help_id(pfa, what)
assert_type_and_record_argument(id, Symbol, "What to plot [:currents, :cx, :coils_flux]"; what)
assert_type_and_record_argument(id, Float64, "Time to plot"; time0)
assert_type_and_record_argument(id, Symbol, "Colormap name"; cname)

Expand Down Expand Up @@ -1001,17 +1002,56 @@ function layer_attrs(l::IMAS.build__layer)
return name, color
end

@recipe function plot_build_cx(bd::IMAS.build; cx=true, wireframe=false, equilibrium=true, pf_active=true, pf_passive=true, only=Symbol[], exclude_layers=Symbol[])
@recipe function plot_build(bd::IMAS.build; equilibrium=true, pf_active=true, pf_passive=true)
id = plot_help_id(bd)
assert_type_and_record_argument(id, Bool, "Plot cross section"; cx)
assert_type_and_record_argument(id, Bool, "Use wireframe"; wireframe)
assert_type_and_record_argument(id, Bool, "Include plot of equilibrium"; equilibrium)
assert_type_and_record_argument(id, Bool, "Include plot of pf_active"; pf_active)
assert_type_and_record_argument(id, Bool, "Include plot of pf_passive"; pf_passive)

dd = top_dd(bd)

legend_position --> :outerbottomright
aspect_ratio := :equal
grid --> :none

rmax = maximum(bd.layer[end].outline.r)
xlim --> [0, rmax]

if dd !== nothing && equilibrium
@series begin
cx := true
dd.equilibrium
end
end

@series begin
dd.build.layer
end

if dd !== nothing && pf_active
@series begin
colorbar --> :false
dd.pf_active
end
end

if dd !== nothing && pf_passive
@series begin
colorbar --> :false
color := :yellow
dd.pf_passive
end
end
end

@recipe function plot_build_layer(layers::IDSvector{<:IMAS.build__layer}; cx=true, wireframe=false, only=Symbol[], exclude_layers=Symbol[])
id = plot_help_id(layers)
assert_type_and_record_argument(id, Bool, "Plot cross section"; cx)
assert_type_and_record_argument(id, Bool, "Use wireframe"; wireframe)
assert_type_and_record_argument(id, AbstractVector{Symbol}, "Only include certain layers"; only)
assert_type_and_record_argument(id, AbstractVector{Symbol}, "Exclude certain layers"; exclude_layers)

dd = top_dd(bd)
bd = parent(layers)

base_linewidth = get(plotattributes, :linewidth, 1.0)

Expand All @@ -1021,19 +1061,11 @@ end

# cx
if cx

if dd !== nothing && equilibrium
@series begin
cx := true
dd.equilibrium
end
end

rmax = maximum(bd.layer[end].outline.r)
rmax = maximum(layers[end].outline.r)

# everything after first vacuum in _out_
if (isempty(only) || (:cryostat in only)) && :cryostat ∉ exclude_layers
for k in get_build_indexes(bd.layer; fs=_out_)[2:end]
for k in get_build_indexes(layers; fs=_out_)[2:end]
if !wireframe
@series begin
seriestype --> :shape
Expand All @@ -1042,10 +1074,10 @@ end
label --> (!wireframe ? "Cryostat" : "")
xlim --> [0, rmax]
join_outlines(
bd.layer[k].outline.r,
bd.layer[k].outline.z,
bd.layer[k-1].outline.r,
bd.layer[k-1].outline.z
layers[k].outline.r,
layers[k].outline.z,
layers[k-1].outline.r,
layers[k-1].outline.z
)
end
end
Expand All @@ -1055,14 +1087,14 @@ end
color --> :black
label --> ""
xlim --> [0, rmax]
bd.layer[k].outline.r[1:end-1], bd.layer[k].outline.z[1:end-1]
layers[k].outline.r[1:end-1], layers[k].outline.z[1:end-1]
end
end
end

# first vacuum in _out_
if !wireframe
k = get_build_indexes(bd.layer; fs=_out_)[1]
k = get_build_indexes(layers; fs=_out_)[1]
@series begin
seriestype --> :shape
linewidth := 0.0
Expand All @@ -1071,10 +1103,10 @@ end
label --> ""
xlim --> [0, rmax]
join_outlines(
bd.layer[k].outline.r,
bd.layer[k].outline.z,
get_build_layer(bd.layer; type=_plasma_).outline.r,
get_build_layer(bd.layer; type=_plasma_).outline.z
layers[k].outline.r,
layers[k].outline.z,
get_build_layer(layers; type=_plasma_).outline.r,
get_build_layer(layers; type=_plasma_).outline.z
)
end
if (isempty(only) || (:cryostat in only)) && :cryostat ∉ exclude_layers
Expand All @@ -1084,7 +1116,7 @@ end
color --> :black
label --> ""
xlim --> [0, rmax]
bd.layer[k].outline.r[1:end-1], bd.layer[k].outline.z[1:end-1]
layers[k].outline.r[1:end-1], layers[k].outline.z[1:end-1]
end
end
end
Expand All @@ -1096,13 +1128,13 @@ end
label --> ""
linestyle --> :dash
color --> :black
[0.0, 0.0], [minimum(bd.layer[end].outline.z), maximum(bd.layer[end].outline.z)]
[0.0, 0.0], [minimum(layers[end].outline.z), maximum(layers[end].outline.z)]
end

# all layers inside of the TF
if (isempty(only) || (:oh in only)) && (!(:oh in exclude_layers))
for k in get_build_indexes(bd.layer; fs=_in_)
layer = bd.layer[k]
for k in get_build_indexes(layers; fs=_in_)
layer = layers[k]
if getproperty(layer, :material, "not_vacuum") != "vacuum"
if !wireframe
@series begin
Expand All @@ -1127,9 +1159,9 @@ end
end

# all layers between the OH and the plasma
for k in get_build_indexes(bd.layer; fs=_hfs_)
l = bd.layer[k]
l1 = bd.layer[k+1]
for k in get_build_indexes(layers; fs=_hfs_)
l = layers[k]
l1 = layers[k+1]
poly = join_outlines(l.outline.r, l.outline.z, l1.outline.r, l1.outline.z)

# setup labels and colors
Expand Down Expand Up @@ -1159,7 +1191,7 @@ end

# plasma
if (isempty(only) || (:plasma in only)) && (!(:plasma in exclude_layers))
plasma_outline = outline(get_build_layer(bd.layer; type=_plasma_))
plasma_outline = outline(get_build_layer(layers; type=_plasma_))
@series begin
seriestype --> :path
linewidth := base_linewidth
Expand Down Expand Up @@ -1222,23 +1254,6 @@ end
end
end

if dd !== nothing && pf_active
@series begin
colorbar --> :false
xlim --> [0, rmax]
dd.pf_active
end
end

if dd !== nothing && pf_passive
@series begin
colorbar --> :false
xlim --> [0, rmax]
color := :yellow
dd.pf_passive
end
end

else # not-cx

@series begin
Expand All @@ -1251,7 +1266,7 @@ end
end

at = 0
for l in bd.layer
for l in layers
name, color = layer_attrs(l)
@series begin
seriestype --> :vspan
Expand All @@ -1274,19 +1289,31 @@ end
end
end

@recipe function plot_build_layer_outline(layer::IMAS.build__layer)
@recipe function plot_build_layer(layer::IMAS.build__layer)
@series begin
aspect_ratio := :equal
label --> layer.name
layer.outline.r, layer.outline.z
layer.outline
end
end

@recipe function plot_build_layer_outline(outline::IMAS.build__layer___outline)
@series begin
aspect_ratio := :equal
outline.r, outline.z
end
end

@recipe function plot_build_structure_outline(structure::IMAS.build__structure)
@recipe function plot_build_structure(structure::IMAS.build__structure)
@series begin
aspect_ratio := :equal
label --> structure.name
structure.outline.r, structure.outline.z
structure.outline
end
end

@recipe function plot_build_structure_outline(outline::IMAS.build__structure___outline)
@series begin
aspect_ratio := :equal
outline.r, outline.z
end
end

Expand Down Expand Up @@ -3024,7 +3051,6 @@ end
end
end


@recipe function plot(x::AbstractVector{<:Real}, y::AbstractVector{<:Measurement}, err::Symbol=:ribbon)
if err == :ribbon
ribbon := Measurements.uncertainty.(y)
Expand Down

0 comments on commit 8065bd5

Please sign in to comment.