Merge pull request #149 from gaelforget/v0p3p14a

add Integration module

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"

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.<lats[l+1]))].=l for l in 1:nl]
+  (map=mask,name=name)
+ else
+  error("unknown option")
+ end
+end
+
+layer_mask(dF,d0,d1)=begin
+  md=0*dF[1:end-1]
+  for k in 1:length(dF)-1
+    dF0=abs(dF[k])
+    dF1=abs(dF[k+1])
+    md[k]=max(0.0,min(d1,dF1)-max(d0,dF0))/(dF1-dF0)
+  end
+  md
+end
+
+"""
+    define_boxes(;option=:hv,grid::NamedTuple, regions=:basins, depths=[(0,7000)])
+
+Define regional integration function for each basin and depth range.
+"""
+function define_boxes(;option=:hv, grid::NamedTuple, regions=:basins, depths=[(0,7000)])
+  dep=(isa(depths,Tuple) ? [depths] : depths)
+  nd=length(dep)
+  rgns=define_regions(option=regions,grid=grid) 
+  nb=length(rgns.name)
+  allones=1.0 .+0*grid.hFacC
+  nr=length(grid.RC)
+
+  xymsk(b) = 1.0*(rgns.map.==b)
+  func_h(X,b)=sum(xymsk(b)*X)
+  tmp2d=MeshArray(grid.XC.grid,Float32)
+
+  zmsk(d0,d1,k) = layer_mask(grid.RF,d0,d1)[k] 
+  func(X,b,d0,d1)=sum([sum(xymsk(b)*zmsk(d0,d1,k)*X[:,k]*
+         grid.DRF[k]*grid.hFacC[:,k]*grid.RAC) for k in 1:nr])
+  ocn_surf=[sum(xymsk(b)*(grid.hFacC[:,1].>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
+

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

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
+
+
+

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)