From 7d6afad6f97dbbc9c521e27a64f619a9786d09d6 Mon Sep 17 00:00:00 2001 From: Hongyang Zhou Date: Wed, 29 Nov 2023 13:21:21 -0500 Subject: [PATCH] getdata -> getdata2d --- ext/BatsrusMakieExt/BatsrusMakieExt.jl | 2 +- ext/BatsrusMakieExt/typerecipe.jl | 2 +- src/Batsrus.jl | 6 +++- src/io.jl | 2 -- src/plot/plots.jl | 2 +- src/plot/pyplot.jl | 11 +++--- src/plot/utility.jl | 47 ++++++++++++++------------ src/select.jl | 3 +- src/vtk.jl | 3 -- test/runtests.jl | 3 ++ 10 files changed, 44 insertions(+), 37 deletions(-) diff --git a/ext/BatsrusMakieExt/BatsrusMakieExt.jl b/ext/BatsrusMakieExt/BatsrusMakieExt.jl index 17099f4..e3e9391 100644 --- a/ext/BatsrusMakieExt/BatsrusMakieExt.jl +++ b/ext/BatsrusMakieExt/BatsrusMakieExt.jl @@ -1,7 +1,7 @@ module BatsrusMakieExt using Batsrus, Printf -import Batsrus: findindex, hasunit, getunit, getdata +import Batsrus: findindex, hasunit, getunit, getdata2d import Batsrus.UnitfulBatsrus import Makie diff --git a/ext/BatsrusMakieExt/typerecipe.jl b/ext/BatsrusMakieExt/typerecipe.jl index 6c6b0c9..fa1c78e 100644 --- a/ext/BatsrusMakieExt/typerecipe.jl +++ b/ext/BatsrusMakieExt/typerecipe.jl @@ -19,7 +19,7 @@ end "Conversion for 2D plots." function Makie.convert_arguments(P::Makie.GridBased, bd::BATLData, var::String; plotrange=[-Inf,Inf,-Inf,Inf], plotinterval=0.1) - x, y, w = getdata(bd, var, plotrange, plotinterval) + x, y, w = getdata2d(bd, var, plotrange, plotinterval) unitx = getunit(bd, bd.head.variables[1]) unity = getunit(bd, bd.head.variables[2]) diff --git a/src/Batsrus.jl b/src/Batsrus.jl index 1758280..a17d1a2 100644 --- a/src/Batsrus.jl +++ b/src/Batsrus.jl @@ -5,7 +5,11 @@ module Batsrus using Printf, Requires -export BATLData +export BATLData, + load, readlogdata, readtecdata, showhead, # io + getvars, getvar, cutdata, subvolume, subsurface, # select + Batl, convertTECtoVTU, convertIDLtoVTK, readhead, readtree, getConnectivity, # vtk + getdata2d # plot/utility "Type for the file information." struct FileList diff --git a/src/io.jl b/src/io.jl index d12213a..f502c01 100644 --- a/src/io.jl +++ b/src/io.jl @@ -1,7 +1,5 @@ # All the IO related APIs. -export load, readlogdata, readtecdata, showhead - # Fortran binary format includes record end/start tags that needs to be skipped. # If there are continuous blocks, we usually skip 2*TAG between actual reading. const TAG = 4 # Fortran record tag size diff --git a/src/plot/plots.jl b/src/plot/plots.jl index 4933324..85907d3 100644 --- a/src/plot/plots.jl +++ b/src/plot/plots.jl @@ -27,7 +27,7 @@ using RecipesBase x, y end elseif ndim == 2 - x, y, w = getdata(bd, var, plotrange, plotinterval) + x, y, w = getdata2d(bd, var, plotrange, plotinterval) unitx = getunit(bd, bd.head.variables[1]) unity = getunit(bd, bd.head.variables[2]) diff --git a/src/plot/pyplot.jl b/src/plot/pyplot.jl index ba393ae..d0b4354 100644 --- a/src/plot/pyplot.jl +++ b/src/plot/pyplot.jl @@ -470,7 +470,7 @@ Wrapper over `contour` in matplotlib. function PyPlot.contour(bd::BATLData, var::AbstractString, levels::Int=0; ax=nothing, plotrange=[-Inf,Inf,-Inf,Inf], plotinterval=0.1, innermask=false, kwargs...) - Xi, Yi, Wi = getdata(bd, var, plotrange, plotinterval; innermask) + Xi, Yi, Wi = getdata2d(bd, var, plotrange, plotinterval; innermask) if isnothing(ax) ax = plt.gca() end @@ -490,7 +490,7 @@ Wrapper over `contourf` in matplotlib. function PyPlot.contourf(bd::BATLData, var::AbstractString, levels::Int=0; ax=nothing, plotrange=[-Inf,Inf,-Inf,Inf], plotinterval=0.1, innermask=false, kwargs...) - Xi, Yi, Wi = getdata(bd, var, plotrange, plotinterval; innermask) + Xi, Yi, Wi = getdata2d(bd, var, plotrange, plotinterval; innermask) if isnothing(ax) ax = plt.gca() end @@ -569,7 +569,8 @@ Wrapper over `plot_surface` in matplotlib. function PyPlot.plot_surface(bd::BATLData, var::AbstractString; plotrange=[-Inf,Inf,-Inf,Inf], plotinterval=0.1, innermask=false, kwargs...) - Xi, Yi, Wi = getdata(bd, var, plotrange, plotinterval; griddim=2, innermask) + xi, yi, Wi = getdata2d(bd, var, plotrange, plotinterval; innermask) + Xi, Yi = meshgrid(xi, yi) plot_surface(Xi, Yi, Wi; kwargs...) end @@ -584,11 +585,11 @@ Wrapper over `pcolormesh` in matplotlib. function PyPlot.pcolormesh(bd::BATLData, var::AbstractString, ax=nothing; plotrange=[-Inf,Inf,-Inf,Inf], plotinterval=0.1, innermask=false, kwargs...) - Xi, Yi, Wi = getdata(bd, var, plotrange, plotinterval; innermask) + xi, yi, Wi = getdata2d(bd, var, plotrange, plotinterval; innermask) if isnothing(ax) ax = plt.gca() end - c = ax.pcolormesh(Xi, Yi, Wi; kwargs...) + c = ax.pcolormesh(xi, yi, Wi; kwargs...) end """ diff --git a/src/plot/utility.jl b/src/plot/utility.jl index daa441f..b1a79c0 100644 --- a/src/plot/utility.jl +++ b/src/plot/utility.jl @@ -1,8 +1,15 @@ # Utility functions for plotting. -"Prepare 2D data arrays for passing to plotting functions." -function getdata(bd::BATLData, var::AbstractString, plotrange, plotinterval; griddim=1, - innermask::Bool=false) +""" + getdata2d(bd::BATLData, var::AbstractString, plotrange=[-Inf, Inf, -Inf, Inf], + plotinterval=Inf; innermask=false) + +Return 2D slices of data `var` from `bd`. If `plotrange` is not set, output data resolution +is the same as the original. If `innermask==true`, then the inner boundary cells are set to +NaN. +""" +function getdata2d(bd::BATLData, var::AbstractString, + plotrange::Vector=[-Inf, Inf, -Inf, Inf], plotinterval::Real=Inf; innermask::Bool=false) x, w, ndim = bd.x, bd.w, bd.head.ndim @assert ndim == 2 "data must be in 2D!" @@ -33,21 +40,23 @@ function getdata(bd::BATLData, var::AbstractString, plotrange, plotinterval; gri yrange = range(x[1,1,2], x[1,end,2], length=size(x,2)) if all(isinf.(plotrange)) xi, yi = xrange, yrange - Xi, Yi = meshgrid(xi, yi) Wi = w[:,:,varIndex_]' # Matplotlib does not accept view! else - if plotrange[1] == -Inf plotrange[1] = minimum(xrange) end - if plotrange[2] == Inf plotrange[2] = maximum(xrange) end - if plotrange[3] == -Inf plotrange[3] = minimum(yrange) end - if plotrange[4] == Inf plotrange[4] = maximum(yrange) end + if plotrange[1] == -Inf plotrange[1] = xrange[1] end + if plotrange[2] == Inf plotrange[2] = xrange[end] end + if plotrange[3] == -Inf plotrange[3] = yrange[1] end + if plotrange[4] == Inf plotrange[4] = yrange[end] end - xi = range(plotrange[1], stop=plotrange[2], step=plotinterval) - yi = range(plotrange[3], stop=plotrange[4], step=plotinterval) + if isinf(plotinterval) + xi = range(plotrange[1], stop=plotrange[2], step=xrange[2] - xrange[1]) + yi = range(plotrange[3], stop=plotrange[4], step=yrange[2] - yrange[1]) + else + xi = range(plotrange[1], stop=plotrange[2], step=plotinterval) + yi = range(plotrange[3], stop=plotrange[4], step=plotinterval) + end - spline = @views Spline2D(xrange, yrange, w[:,:,varIndex_]) - Xi, Yi = meshgrid(xi, yi) - wi = @views spline(Xi[:], Yi[:]) - Wi = reshape(wi, size(Xi)) + interp = @views cubic_spline_interpolation((xrange, yrange), w[:,:,varIndex_]) + Wi = [interp(i, j) for j in yi, i in xi] end end @@ -56,18 +65,14 @@ function getdata(bd::BATLData, var::AbstractString, plotrange, plotinterval; gri varIndex_ = findlast(x->x=="rbody", bd.head.variables) isnothing(varIndex_) && error("rbody not found in file header parameters!") ParamIndex_ = varIndex_ - ndim - bd.head.nw - @inbounds for i = eachindex(Xi, Yi) - if Xi[i]^2 + Yi[i]^2 < bd.head.eqpar[ParamIndex_]^2 + @inbounds for i in CartesianIndices(Wi) + if xi[i[1]]^2 + yi[i[2]]^2 < bd.head.eqpar[ParamIndex_]^2 Wi[i] = NaN end end end - if griddim == 1 - return xi, yi, Wi - else - return Xi, Yi, Wi - end + xi, yi, Wi end "Find variable index in the BATSRUS data." diff --git a/src/select.jl b/src/select.jl index 177d4d8..ae3621d 100644 --- a/src/select.jl +++ b/src/select.jl @@ -1,5 +1,4 @@ -export getvars, getvar, cutdata, subvolume, subsurface - +# Data manipulation. """ cutdata(data, var; plotrange=[-Inf,Inf,-Inf,Inf], dir="x", sequence=1) diff --git a/src/vtk.jl b/src/vtk.jl index f02b872..c13420e 100644 --- a/src/vtk.jl +++ b/src/vtk.jl @@ -2,9 +2,6 @@ using FortranFiles, WriteVTK, Glob, LightXML -export convertTECtoVTU, convertIDLtoVTK, readhead, readtree, getConnectivity -export Batl - # Named indexes of iTree_IA const status_ = Int8(1) # used, unused, to be refined, to be coarsened, etc. const level_ = Int8(2) # grid AMR level for this node diff --git a/test/runtests.jl b/test/runtests.jl index 266572c..438cd11 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -46,6 +46,9 @@ end @test bd.head.time == 25.6f0 @test extrema(bd.x) == (-127.5f0, 127.5f0) @test extrema(bd.w) == (-0.79985905f0, 1.9399388f0) + plotrange = [-10.0, 10.0, -Inf, Inf] + x, y, w = Batsrus.getdata2d(bd, "rho", plotrange) + @test w[1,end] == 0.6848978549242021 end @testset "Reading 2D unstructured binary" begin