diff --git a/Project.toml b/Project.toml index 18f588d..31a7e62 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "MeshArrays" uuid = "cb8c808f-1acf-59a3-9d2b-6e38d009f683" -version = "0.5.8" +version = "0.5.9" authors = ["gaelforget "] [deps] diff --git a/src/Integration.jl b/src/Integration.jl index 363b733..bec0864 100644 --- a/src/Integration.jl +++ b/src/Integration.jl @@ -48,7 +48,19 @@ function define_regions(;option=:global,grid::NamedTuple) lats=[-90 ; -75:10:75 ; 90] nl=length(lats)-1 name=[Symbol("lat_$(lats[l])_to_$(lats[l+1])") for l in 1:nl] - [mask[findall((mask.>0)*(la.>=lats[l])*(la. 0 && laf[i,j] >= lats[l] && laf[i,j] < lats[l+1] + mf[i,j] = Float64(l) + end + end + end + end + (mask=mask,name=name) elseif isa(option,Tuple) dlo=option[1]; dla=option[2] @@ -70,7 +82,7 @@ function define_regions(;option=:global,grid::NamedTuple) t_o="$(lons[i_o])Eto$(lons[i_o+1])E" push!(name,Symbol(t_a*"_"*t_o)) mask[findall((mask.>0)*(la.>=lats[i_a])*(la.=lons[i_o])*(lo.=lons[i_o])*(lo.0)*grid.RAC) for b in 1:nb] - BX=(name=String[],volsum=Function[],volume=Float64[], - ocn_surf=Float64[],tmp2d=tmp2d,tmp3d=tmp3d) - for b in 1:nb - for d in 1:nd - (d0,d1)=dep[d] - n=string(rgns.name[b])*"_dep_$(d0)_to_$(d1)" - @inline f=X->func(X,b,d0,d1) - v=f(allones) - push!(BX.name,n) - push!(BX.volsum,f) - push!(BX.volume,v) - #push!(BX.ocn_surf,ocn_surf[b]) - end - end + BX = (name=String[], volsum=Function[], volume=Float64[], + ocn_surf=Float64[], tmp2d=tmp2d, tmp3d=tmp3d) + for b in 1:nb + msk_b = xymsk(b) # once per basin + for d in 1:nd + (d0,d1) = dep[d] + lmsk_d = lmsk_cache[d] + n = string(rgns.name[b])*"_dep_$(d0)_to_$(d1)" + @inline f = X -> func_inner(X, msk_b, lmsk_d, grid, nr) + v = f(allones) + push!(BX.name, n) + push!(BX.volsum, f) + push!(BX.volume, v) + end + end end - BXh=(name=String[],hsum=Function[],tmp2d=tmp2d,tmp3d=tmp3d) + BXh = (name=String[], hsum=Function[], tmp2d=tmp2d, tmp3d=tmp3d) for b in 1:nb - n=string(rgns.name[b]) - @inline f=X->func_h(X,b) - push!(BXh.name,n) - push!(BXh.hsum,f) + msk_b = xymsk(b) # once per basin + n = string(rgns.name[b]) + @inline f = X -> func_h_inner(X, msk_b) + push!(BXh.name, n) + push!(BXh.hsum, f) end - BXv=(name=String[],vint=Function[],tmp2d=tmp2d,tmp3d=tmp3d) + BXv = (name=String[], vint=Function[], tmp2d=tmp2d, tmp3d=tmp3d) for d in 1:nd - (d0,d1)=dep[d] - n="$(d0)-$(d1)m" - @inline f=X->func_v(X,d0,d1) - push!(BXv.name,n) - push!(BXv.vint,f) + (d0,d1) = dep[d] + lmsk_d = lmsk_cache[d] # already computed above + n = "$(d0)-$(d1)m" + @inline f = X -> func_v_inner!(tmp2d, X, lmsk_d, grid, nr) + push!(BXv.name, n) + push!(BXv.vint, f) end if option==:loops - gridmask(rgns.mask,BXh.name,depths,BXh.hsum,BXv.vint,tmp2d,tmp3d) + gridmask(rgns.mask, BXh.name, depths, BXh.hsum, BXv.vint, tmp2d, tmp3d) elseif option==:streamlined_loop - gridmask(rgns.mask,BX.name,depths,BX.volsum,[],tmp2d,tmp3d) + gridmask(rgns.mask, BX.name, depths, BX.volsum, [], tmp2d, tmp3d) else error("unknown option") end diff --git a/src/Operations.jl b/src/Operations.jl index fb10b6d..5cadbbc 100644 --- a/src/Operations.jl +++ b/src/Operations.jl @@ -277,8 +277,8 @@ function land_mask(m::AbstractMeshArray) for I in eachindex(m.f) μ.f[I]=copy(m.f[I]) end - μ[findall(μ.>0.0)].=1.0 - μ[findall(μ.==0.0)].=NaN + μ[findall(μ.>0.0)]=1.0 + μ[findall(μ.==0.0)]=NaN μ end @@ -488,8 +488,16 @@ Returns a `gridpath` by default, or a `NamedTuple` with fields `lat`, `name`, """ function LatitudeCircle(lat,Γ::NamedTuple; format=:gridpath, range=(0.0,360.0)) - mskCint=1*(Γ.YC .>= lat) - mskC,mskW,mskS=edge_mask(mskCint) + + mskCint = similar(Γ.YC) + for f in 1:mskCint.grid.nFaces + yf = Γ.YC.f[f]; mf = mskCint.f[f] + @inbounds for j in axes(mf,2), i in axes(mf,1) + mf[i,j] = Float64(yf[i,j] >= lat) + end + end + + mskC,mskW,mskS=edge_mask(mskCint) restrict_longitudes!(mskC,Γ.XC,range=range) restrict_longitudes!(mskS,Γ.XS,range=range) restrict_longitudes!(mskW,Γ.XW,range=range) @@ -505,22 +513,40 @@ end is_in_lon_range(x,range)=(range[2].-range[1]>=360)|| (mod(x-range[1],360).0).&(tmp1[2:end-1,2:end-1].==0)) - # velocity masks: - mskW[i]=tmp1[2:end-1,2:end-1] - tmp1[1:end-2,2:end-1] - mskS[i]=tmp1[2:end-1,2:end-1] - tmp1[2:end-1,1:end-2] + mskC = similar(mskCint) + mskW = similar(mskCint) + mskS = similar(mskCint) + + exFLD = exchange(mskCint).MA # still needed for halo + + for i in eachindex(mskC.fIndex) + tmp1 = exFLD.f[i] # halo array, size (s1+2, s2+2) + mCf = mskC.f[i] + mWf = mskW.f[i] + mSf = mskS.f[i] + s1, s2 = size(mCf) + @inbounds for j in 1:s2, ii in 1:s1 + c = tmp1[ii+1, j+1] + neigh = tmp1[ii+1, j ] + tmp1[ii+1, j+2] + + tmp1[ii, j+1] + tmp1[ii+2, j+1] + mCf[ii,j] = Float64((neigh > 0) & (c == 0)) + mWf[ii,j] = c - tmp1[ii, j+1] + mSf[ii,j] = c - tmp1[ii+1, j ] + end end - #treat the case of blank tiles: - #mskC[findall(isnan.(mskC))].=0.0 - #mskW[findall(isnan.(mskW))].=0.0 - #mskS[findall(isnan.(mskS))].=0.0 - - return mskC,mskW,mskS + return mskC, mskW, mskS end ## diff --git a/src/Polygons.jl b/src/Polygons.jl index e0c7a2b..f8e56a2 100644 --- a/src/Polygons.jl +++ b/src/Polygons.jl @@ -48,9 +48,9 @@ function to_Polygons(XG,YG,ff=1) IJ1=1:n1; IJ2=1:n2 for i in IJ1 for j in IJ2 - x=[XG[ff][i,j] XG[ff][i+1,j] XG[ff][i+1,j+1] XG[ff][i,j+1] XG[ff][i,j]] + x=Float64[XG[ff][i,j] XG[ff][i+1,j] XG[ff][i+1,j+1] XG[ff][i,j+1] XG[ff][i,j]] treat_180lon!(x) - y=[YG[ff][i,j] YG[ff][i+1,j] YG[ff][i+1,j+1] YG[ff][i,j+1] YG[ff][i,j]] + y=Float64[YG[ff][i,j] YG[ff][i+1,j] YG[ff][i+1,j+1] YG[ff][i,j+1] YG[ff][i,j]] arr2[i,j]=GI.Polygon([ GI.LinearRing(GI.Point.(zip(vec(x),vec(y)))) ]) # arr2[i,j]=GI.LineString(GI.Point.(zip(vec(x),vec(y)))) end @@ -64,9 +64,9 @@ function to_LineStrings2D(XG,YG,ff=1;do_sphere=true) IJ1=1:n1; IJ2=1:n2 for i in IJ1 for j in IJ2 - x=[XG[ff][i,j] XG[ff][i+1,j] XG[ff][i+1,j+1] XG[ff][i,j+1] XG[ff][i,j]] + x=Float64[XG[ff][i,j] XG[ff][i+1,j] XG[ff][i+1,j+1] XG[ff][i,j+1] XG[ff][i,j]] treat_180lon!(x) - y=[YG[ff][i,j] YG[ff][i+1,j] YG[ff][i+1,j+1] YG[ff][i,j+1] YG[ff][i,j]] + y=Float64[YG[ff][i,j] YG[ff][i+1,j] YG[ff][i+1,j+1] YG[ff][i,j+1] YG[ff][i,j]] arr2[i,j]=GI.LineString(GI.Point.(zip(vec(x),vec(y)))) end end diff --git a/src/ReadWrite.jl b/src/ReadWrite.jl index c279b10..e090218 100644 --- a/src/ReadWrite.jl +++ b/src/ReadWrite.jl @@ -19,7 +19,8 @@ function _load_binary(fil::String, y::AbstractMeshArray) fid = open(fil) xx = Array{eltype(y),1}(undef, n1*n2*n3*n4) read!(fid, xx) - xx = reshape(hton.(xx), (n1,n2,n3,n4)) + map!(hton, xx, xx) + xx = reshape(xx, (n1,n2,n3,n4)) close(fid) xx end @@ -236,10 +237,11 @@ function read_tiles(fil::String,x::AbstractMeshArray) fid = open(fil) xx = Array{eltype(x),2}(undef,(n1*n2,n3)) read!(fid,xx) - xx = reshape(hton.(xx),(n1,n2,n3)) + map!(hton, xx, xx) + xx = reshape(xx,(n1,n2,n3)) close(fid) - return x.grid.read(xx,x) + return read_tiles(xx,x) end diff --git a/src/types/gcmarray.jl b/src/types/gcmarray.jl index 180254a..9cdb7a3 100644 --- a/src/types/gcmarray.jl +++ b/src/types/gcmarray.jl @@ -241,11 +241,11 @@ end import Base: display; display(X::AbstractMeshArray)=show(X) -function Base.similar(A::gcmarray;m::varmeta=defaultmeta) +function Base.similar(A::gcmarray; m::varmeta=defaultmeta) if ndims(A)==1 - B=gcmarray(similar(A.grid),eltype(A),copy(A.fSize),copy(A.fIndex); meta=m) + B = gcmarray(A.grid, eltype(A), A.fSize, A.fIndex; meta=m) else - B=gcmarray(similar(A.grid),eltype(A),copy(A.fSize),copy(A.fIndex),size(A)[2:end]...; meta=m) + B = gcmarray(A.grid, eltype(A), A.fSize, A.fIndex, size(A)[2:end]...; meta=m) end return B end @@ -255,15 +255,13 @@ end Base.BroadcastStyle(::Type{<:AbstractMeshArray}) = Broadcast.ArrayStyle{AbstractMeshArray}() function Base.similar(bc::Broadcast.Broadcasted{Broadcast.ArrayStyle{AbstractMeshArray}}, ::Type{ElType}) where ElType - # Scan the inputs for the gcmarray: - A = find_gcmarray(bc) - # Create the gcmarray output: - if ndims(A)==1 - B=gcmarray(similar(A.grid),ElType,copy(A.fSize),copy(A.fIndex)) - else - B=gcmarray(similar(A.grid),ElType,copy(A.fSize),copy(A.fIndex),size(A)[2:end]...) - end - return B + A = find_gcmarray(bc) + if ndims(A)==1 + B = gcmarray(A.grid, ElType, A.fSize, A.fIndex) + else + B = gcmarray(A.grid, ElType, A.fSize, A.fIndex, size(A)[2:end]...) + end + return B end find_gcmarray(bc::Base.Broadcast.Broadcasted) = find_gcmarray(bc.args) @@ -276,20 +274,28 @@ find_gcmarray(::Any, rest) = find_gcmarray(rest) import Base: copyto! -# Specialize this method if all you want to do is specialize on typeof(dest) +# Rewrite a Broadcasted tree: replace AbstractMeshArray args with their face-I inner arrays. +# Scalars and plain arrays pass through unchanged. +function _bc_face(bc::Broadcast.Broadcasted, I) + new_args = map(a -> _bc_face(a, I), bc.args) + Broadcast.Broadcasted(bc.f, new_args) # axes=nothing → recomputed from face arrays +end +_bc_face(a::AbstractMeshArray, I) = ndims(a.f) == 1 ? a.f[I[1]] : a.f[I] +_bc_face(a, I) = a # scalar, plain array, etc. + @inline function copyto!(dest::AbstractMeshArray, bc::Broadcast.Broadcasted{Nothing}) axes(dest) == axes(bc) || throwdm(axes(dest), axes(bc)) - # Performance optimization: broadcast!(identity, dest, A) is equivalent to copyto!(dest, A) if indices match - if bc.f === identity && bc.args isa Tuple{AbstractArray} # only a single input argument to broadcast! - A = bc.args[1] - if axes(dest) == axes(A) - return copyto!(dest, A) + for I in CartesianIndices(dest.f) + face_args = map(a -> _bc_face(a, I), bc.args) + array_args = filter(a -> isa(a, AbstractArray), face_args) + if !isempty(array_args) + result_axes = Base.Broadcast.combine_axes(array_args...) + result_size = map(length, result_axes) + if size(dest.f[I]) != result_size + dest.f[I] = similar(dest.f[I], result_size) + end end - end - bc′ = Broadcast.preprocess(dest, bc) - @simd for I in eachindex(bc′) - #@inbounds dest[I] = bc′[I] - @inbounds dest[I] = gcmarray_getindex_evalf(bc′,I) + broadcast!(bc.f, dest.f[I], face_args...) end return dest end diff --git a/src/types/gcmvector.jl b/src/types/gcmvector.jl index b4b6c96..06068bd 100644 --- a/src/types/gcmvector.jl +++ b/src/types/gcmvector.jl @@ -62,10 +62,12 @@ end Base.BroadcastStyle(::Type{<:gcmvector}) = Broadcast.ArrayStyle{gcmvector}() function Base.similar(bc::Broadcast.Broadcasted{Broadcast.ArrayStyle{gcmvector}}, ::Type{ElType}) where ElType - # Scan the inputs for the gcmarray: A = find_gcmvector(bc) - # Create the gcmvector output: - return gcmvector{ElType,ndims(A)}(A.grid,similar(A.f),copy(A.fSize),copy(A.fIndex)) + f = similar(A.f, Array{ElType,1}) + for I in eachindex(A.f) + f[I] = similar(A.f[I], ElType) + end + return gcmvector{ElType,ndims(A)}(A.grid,f,copy(A.fSize),copy(A.fIndex)) end find_gcmvector(bc::Base.Broadcast.Broadcasted) = find_gcmvector(bc.args) @@ -77,17 +79,17 @@ find_gcmvector(::Any, rest) = find_gcmvector(rest) # Specialize this method if all you want to do is specialize on typeof(dest) @inline function copyto!(dest::gcmvector, bc::Broadcast.Broadcasted{Nothing}) axes(dest) == axes(bc) || throwdm(axes(dest), axes(bc)) - # Performance optimization: broadcast!(identity, dest, A) is equivalent to copyto!(dest, A) if indices match - if bc.f === identity && bc.args isa Tuple{AbstractArray} # only a single input argument to broadcast! - A = bc.args[1] - if axes(dest) == axes(A) - return copyto!(dest, A) + for I in CartesianIndices(dest.f) + face_args = map(a -> _bc_face(a, I), bc.args) + array_args = filter(a -> isa(a, AbstractArray), face_args) + if !isempty(array_args) + result_axes = Base.Broadcast.combine_axes(array_args...) + result_size = map(length, result_axes) + if size(dest.f[I]) != result_size + dest.f[I] = similar(dest.f[I], result_size) + end end - end - bc′ = Broadcast.preprocess(dest, bc) - @simd for I in eachindex(bc′) - #@inbounds dest[I] = bc′[I] - @inbounds dest[I] = gcmarray_getindex_evalf(bc′,I) + broadcast!(bc.f, dest.f[I], face_args...) end return dest end @@ -126,11 +128,8 @@ function Base.getindex(A::gcmarray{T,N}, B::gcmvector{CartesianIndex{2},N}) wher end function Base.view(A::gcmarray{T,N}, B::gcmvector{CartesianIndex{2},N}) where {T,N} - tmpOut=missing - for a in eachindex(A) - ismissing(tmpOut) ? tmpOut=view(A[a],B[a]) : tmpOut=CatView(tmpOut,view(A[a],B[a])) - end - return tmpOut + views = [view(A[a], B[a]) for a in eachindex(A)] + return CatView(views...) end function Base.setindex!(A::gcmarray{T,N}, B::gcmvector{T,N}, C::gcmvector{CartesianIndex{2},N}) where {T,N} diff --git a/test/testsets/transport.jl b/test/testsets/transport.jl index d348fb6..eeded10 100644 --- a/test/testsets/transport.jl +++ b/test/testsets/transport.jl @@ -36,7 +36,7 @@ uv=Dict("U"=>Tx,"V"=>Ty,"dimensions"=>["x","y"]) L=-85.0:5.0:85.0; LC=LatitudeCircles(L,Γ,format=:gridpath) T=Array{Float64,1}(undef,length(LC)) - [T[i]=1e-6*ThroughFlow(uv,LC[i],Γ) for i=1:length(LC)] + for i=1:length(LC); T[i]=1e-6*ThroughFlow(uv,LC[i],Γ); end plot(LC) plot(LC[1])