From 19bc860323aeb2c66907643441a57d4424609143 Mon Sep 17 00:00:00 2001 From: gaelforget Date: Thu, 19 Sep 2024 12:00:05 -0400 Subject: [PATCH] add Integration module --- Project.toml | 8 +- src/Integration.jl | 213 +++++++++++++++++++++++++++++++++++++++++++++ src/MeshArrays.jl | 4 +- src/Types.jl | 18 ++++ test/runtests.jl | 15 ++++ 5 files changed, 255 insertions(+), 3 deletions(-) create mode 100644 src/Integration.jl diff --git a/Project.toml b/Project.toml index 00383ad..926a140 100644 --- a/Project.toml +++ b/Project.toml @@ -6,21 +6,24 @@ version = "0.3.13" [deps] CatViews = "81a5f4ea-a946-549a-aa7e-2a7f63a27d31" Dates = "ade2ca70-3891-5945-98fb-dc099432e06a" +Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b" +Glob = "c27321d9-0574-5035-807b-f59d2c89b15c" LazyArtifacts = "4af54fe1-eca0-43a8-85a7-787d91b784e3" NearestNeighbors = "b8a86587-4115-5ab1-83bc-aa920d37bbce" Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" +SharedArrays = "1a1011a3-84de-559e-8e89-a11a2f7dc383" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d" [weakdeps] DataDeps = "124859b0-ceae-595e-8997-d05f6a7a8dfe" +GeoJSON = "61d90e0f-e114-555e-ac52-39dfb47a3ef9" +JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819" Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a" Proj = "c94c279d-25a6-4763-9509-64d165bea63e" -GeoJSON = "61d90e0f-e114-555e-ac52-39dfb47a3ef9" Shapefile = "8e980c4a-a4fe-5da2-b3a7-4b4b0353a2f4" -JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819" [extensions] MeshArraysDataDepsExt = ["DataDeps"] @@ -34,6 +37,7 @@ MeshArraysShapefileExt = ["Shapefile"] CatViews = "1.0" DataDeps = "0.7" GeoJSON = "0.6, 0.7, 0.8" +Glob = "1" JLD2 = "0.4, 0.5" Makie = "0.19, 0.20, 0.21" NearestNeighbors = "0.4" diff --git a/src/Integration.jl b/src/Integration.jl new file mode 100644 index 0000000..07a4350 --- /dev/null +++ b/src/Integration.jl @@ -0,0 +1,213 @@ + +module Integration + +using Distributed, SharedArrays, Glob +import MeshArrays: read_JLD2, write_JLD2 +import MeshArrays: GridLoad, GridLoadVar, GridSpec +import MeshArrays: demo, MeshArray, gridmask + +## + +example(;option=:hv,regions=:dlat_10,depths=[(0,7000)]) = begin + g=GridSpec(ID=:LLC90); Γ=GridLoad(g) + G=(hFacC=GridLoadVar("hFacC",g),RF=GridLoadVar("RF",g), + RC=GridLoadVar("RC",g),RAC=GridLoadVar("RAC",g), + DRF=GridLoadVar("DRF",g)) + G=merge(Γ,G) + + M=define_boxes(option=option, regions=regions, grid=G, depths=depths) + + diags=joinpath(pwd(),"diags") + files=glob("state_3d_set1*.data",diags) + + G,M,files +end + +## + +DEPTHS=[(0,100),(100,200),(200,300),(300,400),(400,500),(500,600),(600,700), + (700,800),(800,1000),(1000,1200),(1200,1400),(1400,1700),(1700,2000), + (2000,2500),(2500,3000),(3000,4000),(4000,5000),(5000,7000)] + +""" + define_regions(;option=:basins,grid::NamedTuple) + +Define regional integration mask (one value period region). +""" +function define_regions(;option=:basins,grid::NamedTuple) + if option==:basins + demo.ocean_basins() + elseif option==:global + mask=1.0*(grid.hFacC[:,1].>0) + (map=mask,name=["Global"]) + elseif option==:dlat_10 + mask=1.0*(grid.hFacC[:,1].>0) + la=grid.YC + 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)*grid.RAC) for b in 1:nb] + tmp3d=MeshArray(grid.XC.grid,Float32,nr) + + function func_v(X,d0,d1) + tmp2d.=0.0 + for k in 1:nr + tmp2d.+=zmsk(d0,d1,k)*X[:,k]*grid.DRF[k]*grid.hFacC[:,k]*grid.RAC + end + tmp2d + end + + 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 + + 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) + end + + 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) + end + + if option==:hv + gridmask(rgns.map,BXh.name,depths,BXh.hsum,BXv.vint,tmp2d,tmp3d) + else + gridmask(rgns.map,BX.name,depths,BX.volsum,[],tmp2d,tmp3d) + end +end + +#nonan(x)=[(isnan(y) ? 0.0 : y) for y in x] + +""" + loop(boxes::NamedTuple; files=String[], var=:THETA, rd=read) + +``` +begin +@everywhere using MeshArrays, MITgcm +@everywhere rd(F,var,tmp)=read(read_mdsio(F,var),tmp) +@everywhere G,M,files=Integration.example(option=:hv) +#,depths=Integration.DEPTHS) +end +#H_3d=Integration.loop_3d(M,files=files,rd=rd) +H_hv=Integration.loop_hv(M,files=files,rd=rd) +``` + +and to save results: + +``` +using JLD2; output_path="test.jld2" +jldsave(output_path; depths=M.depths, integral=H, volume=vol, name=M.names) +``` + +where vol is calculated as follows: + +``` +#option=:other +allones=1.0 .+0*G.hFacC +vol=[b(allones) for b in M.h_sum] +``` + +or + +``` +#option=:hv +M.tmp2d.=M.v_int[1](allones) +vol=[b(M.tmp2d) for b in M.h_sum] +``` +""" +function loop_3d(boxes::gridmask; files=String[], var=:THETA, rd=read) + nt=length(files) + nb=length(boxes.names) + BA=SharedArray{Float64}(nb,nt) + @sync @distributed for t in 1:nt + mod(t,10)==0 ? println(t) : nothing + F=files[t] + ext=split(F,".")[end] + boxes.tmp3d.=rd(F,var,boxes.tmp3d) + BA[:,t]=[b(boxes.tmp3d) for b in boxes.h_sum] + GC.gc() + end + BA +end + +## + +function loop_hv(boxes::gridmask; files=String[], var=:THETA, rd=read) + nt=length(files) + nh=length(boxes.names) + nv=length(boxes.depths) + BA=SharedArray{Float64}(nh,nv,nt) + @sync @distributed for t in 1:nt + mod(t,10)==0 ? println(t) : nothing + F=files[t] + ext=split(F,".")[end] + boxes.tmp3d.=rd(F,var,boxes.tmp3d) + for layer in 1:nv + boxes.tmp2d.=boxes.v_int[layer](boxes.tmp3d) + BA[:,layer,t]=[b(boxes.tmp2d) for b in boxes.h_sum] + #dH0[i]=nansum(write(dT*(G.b_i.==i))) + end + GC.gc() + end + BA +end + +end + diff --git a/src/MeshArrays.jl b/src/MeshArrays.jl index 178596c..34abaae 100644 --- a/src/MeshArrays.jl +++ b/src/MeshArrays.jl @@ -25,14 +25,16 @@ include("ReIndexing.jl") include("Interpolation.jl") include("VerticalDimension.jl") include("demo.jl") +include("Integration.jl") -export AbstractMeshArray, MeshArray, gcmgrid, varmeta, gridpath +export AbstractMeshArray, MeshArray, gcmgrid, varmeta, gridpath, gridmask export GridSpec, GridLoad, GridLoadVar, UnitGrid, simple_periodic_domain export exchange, Tiles, Tiles!, Interpolate, InterpolationFactors, knn, interpolation_setup #The following exch_UV differs from normal exchange; incl. exch_UV_N export exch_UV export nansum, nanmean, nanmax, nanmin export demo, land_mask +export Integration #export InnerArray, OuterArray #export smooth, mask diff --git a/src/Types.jl b/src/Types.jl index 078cb46..9919f50 100644 --- a/src/Types.jl +++ b/src/Types.jl @@ -236,3 +236,21 @@ Base.@kwdef struct gridpath W::Matrix S::Matrix end + +""" + gridmask + +gridmask data structure. +""" +Base.@kwdef struct gridmask + map::MeshArray + names::Array + depths::Array + h_sum::Array + v_int::Array + tmp2d::MeshArray + tmp3d::MeshArray +end + + + diff --git a/test/runtests.jl b/test/runtests.jl index 373f523..2738bbd 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -49,6 +49,21 @@ end @test isapprox(d[1][180,90],-2204.8384919029777) end +@testset "Regional Integration:" begin + G,M,files=Integration.example(option=:hv) + + allones=1.0 .+0*G.hFacC + vol0=sum(G.RAC*G.DRF*G.hFacC) + + M.tmp2d.=M.v_int[1](allones) + vol=[b(M.tmp2d) for b in M.h_sum] + @test isapprox(sum(vol),vol0) + + G,M,files=Integration.example(option=:whole) + vol=[b(allones) for b in M.h_sum] + @test isapprox(sum(vol),vol0) +end + @testset "Transport computations:" begin #Load grid and transport / vector field γ=GridSpec(ID=:LLC90)