Add plot_everything.jl

[skip ci][ci skip]

post_processing/plot/plot_everything.jl

@@ -0,0 +1,236 @@
+# plot_everything.jl: a simple Julia script to make a suite of plots for quick inspection as
+# part of CI.
+
+if length(ARGS) < 1
+    error("Run as plot_everything.jl <path of datadir>")
+end
+
+using Statistics # Used for `mean`
+using NCDatasets
+using CairoMakie
+
+# The NetCDF files contain the dataset we are interested in, but also all the
+# dimensions. We list the dimensions here so that we can isolate the actual
+# dataset.
+dimensions = ["x", "y", "z", "z_half", "time", "lon", "lat"]
+
+function contour_plot(
+    fig,
+    X,
+    Y,
+    Z,
+    ax_location;
+    title,
+    xlabel,
+    ylabel,
+    colorbar_label,
+)
+    # We put at most 4 plots on the same column, so we can the column and row index by
+    # modding times 4
+    yloc, xloc = divrem(ax_location[], 4)
+    # Indices start at 1
+    xloc += 1
+    yloc += 1
+
+    Axis(fig[xloc, yloc]; title, xlabel, ylabel)
+    # Custom_levels is a workaround for plotting constant fields with CairoMakie
+    custom_levels =
+        minimum(Z) ≈ maximum(Z) ? (minimum(Z):0.1:(minimum(Z) + 0.2)) : 25
+    cf = CairoMakie.contourf!(X, Y, Z, levels = custom_levels)
+    Colorbar(fig[xloc, yloc + 1], cf, label = colorbar_label)
+
+    ax_location[] += 1
+end
+
+calc_zonalave_timeave(x) =
+    dropdims(dropdims(mean(mean(x, dims = 1), dims = 4), dims = 4), dims = 1)
+calc_zonalave_variance(x) = calc_zonalave_timeave((x .- mean(x, dims = 1)) .^ 2)
+calc_zonalave_covariance(x, y) =
+    calc_zonalave_timeave((x .- mean(x, dims = 1)) .* (y .- mean(y, dims = 1)))
+
+function plot_one(file_path)
+
+    # Plot MAX_NUM_SNAPSHOTS time snapshots (unless not as many are available)
+    MAX_NUM_SNAPSHOTS = 3
+
+    name = join(split(basename(file_path), ".")[1:(end - 1)])
+
+    nc = NCDatasets.NCDataset(file_path)
+
+    # We split the interval, 1:length(nc["time"]) into MAX_NUM_SNAPSHOTS. When
+    # MAX_NUM_SNAPSHOTS is too large for the number of snapshots we have, we round
+    # everything down and remove the duplicates. At the end, we cast everything to Int
+    # because we will use this as indices.
+    selected_time_indices = Vector{Int}(
+        unique(floor.(range(1, length(nc["time"]), MAX_NUM_SNAPSHOTS))),
+    )
+
+    # keys(nc) contains all the datasets, which are all the dimensions + the one dataset we
+    # are interested in. So, to find the name of the variable we want to plot, we simply
+    # remove all the dimensions. NOTE: The assumption here is that there's only one dataset
+    # per file!
+    var_name = setdiff(keys(nc), dimensions)[]
+
+    var = nc[var_name]
+    var_long_name = nc[var_name].attrib["long_name"]
+    var_units = nc[var_name].attrib["units"]
+
+    # The last dimension is always time
+    var_dimnames = dimnames(var)
+    num_dims = length(var_dimnames)
+
+    # TODO: Find more elegant way to do this
+    if num_dims == 4
+        # If have 4 dimensions, we have x-y-z-t, or lat-long-z-t
+        xname, yname, zname, _ = var_dimnames
+
+        dim1 = nc[xname][:]
+        dim2 = nc[yname][:]
+        dimz = nc[zname][:]
+
+        xunits = nc[xname].attrib["units"]
+        yunits = nc[yname].attrib["units"]
+        zunits = nc[zname].attrib["units"]
+
+        zhalf_index = convert(Int, floor(length(nc[zname]) / 2))
+
+        zmin = nc[zname][1]
+        zmax = nc[zname][end]
+        zhalf = nc[zname][zhalf_index]
+
+        for time_index in selected_time_indices
+            time = nc["time"][time_index]
+
+            var_zmin = var[:, :, 1, time_index]
+            var_zmax = var[:, :, end, time_index]
+            var_zhalf = var[:, :, zhalf_index, time_index]
+            var_zonal =
+                mean(nc[var_name][:, :, :, time_index], dims = 1)[1, :, :]
+
+            fig = Figure(resolution = (800, 1600))
+            # We want to automatically increment ax_location, so we make it a reference
+            ax_location = Ref(0)
+
+            contour_plot(
+                fig,
+                dim1,
+                dim2,
+                var_zmin,
+                ax_location,
+                title = "$var_long_name, t = $time s, z = $zmin $zunits",
+                xlabel = "$xname [$xunits]",
+                ylabel = "$yname [$yunits]",
+                colorbar_label = "$var_name [$var_units]",
+            )
+            contour_plot(
+                fig,
+                dim1,
+                dim2,
+                var_zhalf,
+                ax_location,
+                title = "$var_long_name, t = $time s, z = $zhalf $zunits",
+                xlabel = "$xname [$xunits]",
+                ylabel = "$yname [$yunits]",
+                colorbar_label = "$var_name [$var_units]",
+            )
+            contour_plot(
+                fig,
+                dim1,
+                dim2,
+                var_zmax,
+                ax_location,
+                title = "$var_long_name, t = $time s, z = $zmax $zunits",
+                xlabel = "$xname [$xunits]",
+                ylabel = "$yname [$yunits]",
+                colorbar_label = "$var_name [$var_units]",
+            )
+            contour_plot(
+                fig,
+                dim2,
+                dimz,
+                var_zonal,
+                ax_location,
+                title = "$var_long_name, t = $time s, (averaged over $xname)",
+                xlabel = "$yname [$xunits]",
+                ylabel = "$zname [$zunits]",
+                colorbar_label = "$var_name [$var_units]",
+            )
+            output_path = joinpath(ARGS[1], name * "_$time_index.png")
+            save(output_path, fig)
+            @info "Plotted $output_path"
+        end
+
+        # Time averaged
+        fig = Figure(resolution = (800, 800))
+        ax_location = Ref(0)
+
+        contour_plot(
+            fig,
+            dim2,
+            dimz,
+            calc_zonalave_timeave(nc[var_name][:]),
+            ax_location,
+            title = "$var_long_name (averaged over all times and $xname)",
+            xlabel = "$yname [$xunits]",
+            ylabel = "$zname [$zunits]",
+            colorbar_label = "$var_name [$var_units]",
+        )
+
+        contour_plot(
+            fig,
+            dim2,
+            dimz,
+            calc_zonalave_variance(nc[var_name][:]),
+            ax_location,
+            title = "Variance $var_long_name (averaged over all times and $xname)",
+            xlabel = "$yname [$xunits]",
+            ylabel = "$zname [$zunits]",
+            colorbar_label = "$var_name [$var_units]",
+        )
+
+        output_path = joinpath(ARGS[1], name * "_time_averages.png")
+        save(output_path, fig)
+        @info "Plotted $output_path"
+
+    elseif num_dims == 3
+        # If we have 3 dimensions, it could be x-y-t, lat-long-t, or x-z-t
+        xname, yname, _ = var_dimnames
+
+        dim1 = nc[xname][:]
+        dim2 = nc[yname][:]
+
+        xunits = nc[xname].attrib["units"]
+        yunits = nc[yname].attrib["units"]
+
+        for time_index in selected_time_indices
+            time = nc["time"][time_index]
+            var_t_final = var[:, :, time_index]
+
+            fig = Figure(resolution = (800, 600))
+            # We want to automatically increment ax_location, so we make it a reference
+            ax_location = Ref(0)
+
+            contour_plot(
+                fig,
+                dim1,
+                dim2,
+                var_t_final,
+                ax_location,
+                title = "$var_long_name, t = $time s",
+                xlabel = "$xname [$xunits]",
+                ylabel = "$yname [$yunits]",
+                colorbar_label = "$var_name [$var_units]",
+            )
+
+            output_path = joinpath(ARGS[1], name * "_$time_index.png")
+            save(output_path, fig)
+            @info "Plotted $output_path"
+        end
+    end
+
+    NCDatasets.close(nc)
+
+end
+
+nc_files = filter(x -> x[(end - 2):end] == ".nc", readdir(ARGS[1]; join = true))
+map(plot_one, nc_files)