you can now use "heatmap" from Plots to make a heatmap of a crss-sect…

…ion through either your original model setup:

- heatmap(model, field=:phase, y=0.1)
or through the result of a timestep of a simulation
heatmap(phase, field=:phase, y=0.1, timestep=1)

src/LaMEM.jl

@@ -40,8 +40,7 @@ export LaMEM_jll        # export LaMEM_jll as well & directories
 function __init__()
     #@require GLMakie = "e9467ef8-e4e7-5192-8a1a-b1aee30e663a" begin 
     @require Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80" begin
-
-        @eval include("MakieExt.jl")
+        @eval include("PlotsExt.jl")
     end
 end
 

src/LaMEM_ModelGeneration/Utils.jl

@@ -190,68 +190,72 @@ end
 
 
 """
-    x_vec, z_vec, data, axes_str = cross_section(model::LaMEM.Model, field=:phases; x=nothing, y=nothing, z=nothing)
+    data_tuple, axes_str = cross_section(model::LaMEM.Model, field=:phases; x=nothing, y=nothing, z=nothing)
 
 This creates a cross-section through the initial model setup & returns a 2D array
 """
-function cross_section(model::Model, field::Symbol =:phases; x=nothing, y=nothing, z=nothing)
-    Model3D = CartData(model.Grid.Grid, (Phases=model.Grid.Phases,Temp=model.Grid.Temp));
+function cross_section(model::Model, field::Symbol=:phase; x=nothing, y=nothing, z=nothing)
+    Model3D = CartData(model.Grid.Grid, (phase=model.Grid.Phases,temperature=model.Grid.Temp));
+
+    data_tuple, axes_str = cross_section(Model3D, field; x=x, y=y, z=z)
+
+    return data_tuple, axes_str
+end
+
+"""
+    Cross = cross_section(cart::CartData, field::Symbol =:phase; x=nothing, y=nothing, z=nothing)
+
+Creates a cross-section through the data and returns `x,z` coordinates
+"""
+function cross_section(cart::CartData, field::Symbol=:phase; x=nothing, y=nothing, z=nothing)
 
     if !isnothing(z); Depth_level = z*km; else Depth_level=nothing; end
     if !isnothing(x); Lon_level = x; else Lon_level=nothing; end
     if !isnothing(y); Lat_level = y; else Lat_level=nothing; end
 
-    Cross = CrossSectionVolume(Model3D, Depth_level=Depth_level, Lat_level=Lat_level, Lon_level=Lon_level)
+    Cross = CrossSectionVolume(cart, Depth_level=Depth_level, Lat_level=Lat_level, Lon_level=Lon_level)
 
+    data_tuple, axes_str = flatten(Cross, field,x,y,z)
+    return data_tuple, axes_str
+end
 
-    # Flatten arrays + coordinates
-    if      !isnothing(x)
-        Phase2D, Temp2D =   dropdims(Cross.fields.Phases,dims=1), dropdims(Cross.fields.Temp,dims=1)
-        X2D, Z2D = dropdims(Cross.y.val,dims=1), dropdims(Cross.z.val,dims=1)
-        x_vec, z_vec = X2D[:,1],  Z2D[1,:]
+"""
+Creates a 2D array out of a cross-section and a specified data field
+"""
+function flatten(cross::CartData, field::Symbol,x,y,z)
+    dim     =   findall(size(cross.x.val) .== 1)[1]
+    X       =   dropdims(cross.x.val, dims=dim)
+    Y       =   dropdims(cross.y.val, dims=dim)
+    Z       =   dropdims(cross.z.val, dims=dim)
+    data_or =   getfield(cross.fields,field)
+
+    # flatten data array
+    if isa(data_or, Array)
+        data = dropdims(data_or, dims=dim)
+    elseif isa(data_or,Tuple)
+        # vectors or tensors
+        data = ()
+        for d in data_or
+            data = (data..., dropdims(d, dims=dim) )
+        end
+    end
+
+    if  dim==1
         title_str = "x = $x"
         x_str,z_str = "y","z"
-    elseif !isnothing(y)
-        Phase2D, Temp2D =   dropdims(Cross.fields.Phases,dims=2), dropdims(Cross.fields.Temp,dims=2)
-        X2D, Z2D = dropdims(Cross.x.val,dims=2), dropdims(Cross.z.val,dims=2)
-        x_vec, z_vec = X2D[:,1],  Z2D[1,:]
+        x, z = Y[:,1],  Z[1,:]
+    elseif  dim==2
         title_str = "y = $y"
         x_str,z_str = "x","z"
-    elseif !isnothing(z)
-        Phase2D, Temp2D =   dropdims(Cross.fields.Phases,dims=3), dropdims(Cross.fields.Temp,dims=3)
-        X2D, Z2D = dropdims(Cross.x.val,dims=3), dropdims(Cross.y.val,dims=3)
-
-        x_vec, z_vec = X2D[:,1],  Z2D[1,:]
+        x, z = X[:,1],  Z[1,:]
+    elseif dim==3
         title_str = "z = $z"
         x_str,z_str = "x","y"
-    end
-
-    # Create figure
-    if field == :phases
-        data = Phase2D
-        title_str = "Phases; "*title_str
-    else
-        data = Temp2D
-        title_str = "Temperature; "*title_str
+        x, z = X[:,1],  Y[1,:]
     end
     axes_str = (x_str=x_str, z_str=z_str, title_str=title_str)
-
-    return x_vec, z_vec, data, axes_str
+    data_tuple = (x=x,z=z,data=data)
+    return data_tuple, axes_str
 end
 
-
-"""
-    Cross = cross_section(cart::CartData, field::Symbol =:phases; x=nothing, y=nothing, z=nothing)
-
-Creates a cross-section through the data
-"""
-function cross_section(cart::CartData, field::Symbol =:phases; x=nothing, y=nothing, z=nothing)
-
-    if !isnothing(z); Depth_level = z*km; else Depth_level=nothing; end
-    if !isnothing(x); Lon_level = x; else Lon_level=nothing; end
-    if !isnothing(y); Lat_level = y; else Lat_level=nothing; end
-
-    Cross = CrossSectionVolume(cart, Depth_level=Depth_level, Lat_level=Lat_level, Lon_level=Lon_level)
-
-    return Cross
-end
+

src/PlotsExt.jl

@@ -0,0 +1,46 @@
+# Plotting extensions, that are only loaded when GLMakie is available
+println("adding Plots.jl plotting extensions for LaMEM")
+
+using LaMEM, GeophysicalModelGenerator
+using .Plots
+import .Plots: heatmap
+
+"""
+    Plots.heatmap(model::Union{Model,CartData}, args...; field::Symbol=:phase,  dim=1, x=nothing, y=nothing, z=nothing, aspect_ratio::Union{Real, Symbol}=:equal)
+   
+This plots a `Plots` heatmap of a cross-section through the LaMEM `model`, of the field  `field`. If that is a vector or tensor field, specify `dim` to indicate which of the fields you want to see.
+If the keyword `timestep` is specified, it will load a timestep 
+"""
+function Plots.heatmap(model::Union{Model,CartData}, args...; field::Symbol=:phase, timestep=nothing, dim=1, x=nothing, y=nothing, z=nothing, aspect_ratio::Union{Real, Symbol}=:equal)
+
+    if !isnothing(timestep)
+        # load a particular timestep
+        data_cart, time = Read_LaMEM_timestep(model,timestep)
+        model = data_cart
+    end
+
+    data_tuple, axes_str = cross_section(model, field; x=x, y=y, z=z)
+
+    if isa(data_tuple.data, Array)
+        data_field = data_tuple.data
+        cb_str = String(field)
+    elseif isa(data_tuple.data, Tuple)
+        data_field = data_tuple.data[dim]
+        cb_str = String(field)*"[$dim]"
+    end
+
+    title_str = axes_str.title_str
+    if !isnothing(timestep)
+        title_str=title_str*"; time=$time[1]"
+    end
+
+    hm = heatmap(data_tuple.x, data_tuple.z, data_field', 
+                aspect_ratio=aspect_ratio, 
+                xlabel=axes_str.x_str,
+                ylabel=axes_str.z_str,
+                title=title_str,
+                colorbar_title=cb_str,
+                args...)
+
+    return hm
+end