Merge pull request #127 from gaelforget/ERA5interp

from ERA5 grid to ECCO4 grid

Project.toml

@@ -1,7 +1,7 @@
 name = "MeshArrays"
 uuid = "cb8c808f-1acf-59a3-9d2b-6e38d009f683"
 authors = ["gaelforget <gforget@mit.edu>"]
-version = "0.3.8"
+version = "0.3.9"
 
 [deps]
 CatViews = "81a5f4ea-a946-549a-aa7e-2a7f63a27d31"

examples/ERA5/ERA5_module.jl

@@ -0,0 +1,167 @@
+
+module ERA5_to_ECCO4
+
+using MeshArrays, Interpolations, NCDatasets, DataFrames, Statistics, FortranFiles, Dates
+
+## list of variables etc
+
+list_in=["dlw","dsw","pres","rain","d2m","tmp2m_degC","u10m","ustr","v10m","vstr","wspeed"];
+list_ds=["msdwlwrf","msdwswrf","sp","tp","d2m","t2m","u10","metss","v10","mntss","..."]
+
+offset=zeros(12)
+offset[6]=-273.15
+factor=ones(12)
+factor[1]=-1.0
+factor[2]=-1.0
+factor[4]=1/3600
+factor[7]=1.0
+factor[8]=-1.0
+factor[9]=1.0
+factor[10]=-1.0
+
+## compute specific humidity
+
+Rdry=287.0597 ; Rvap=461.5250 ; a1=611.21 ; a3=17.502 ; a4=32.19 ; T0=273.16
+#Calculation of E saturation water vapour from Teten's formula
+E(dtas)=a1*exp(a3*(dtas-T0)/(dtas-a4))
+#Calculation of saturation specific humidity at 2m qsat  (equal to huss)
+qsat(ps,E)=(Rdry/Rvap)*E/(ps-((1-Rdry/Rvap)*E))
+
+## interpolation and bin average code
+
+function setup_arrays_ERA5(Γ)
+    lon=collect(0.0:0.25:359.75)
+    lat=collect(90.0:-0.25:-90.0)
+    ds_knn=calc_knn(lon,lat,Γ)
+    A,B=qt_arrays(lon,lat,Γ)
+    x=zeros(length(lon),length(lat))
+    (lon=lon,lat=lat,knn=ds_knn,A=A,B=B,x=x)
+end
+
+function calc_knn(q_lon,q_lat,Γ)
+    lo=[q_lon[i] for i in 1:length(q_lon), j in 1:length(q_lat)]
+    la=[q_lat[j] for i in 1:length(q_lon), j in 1:length(q_lat)]
+    (f,i,j,c)=knn(Γ.XC,Γ.YC,lo[:],la[:])
+    c
+end
+
+q_fx(lon)=mod(lon,360.0)*4+1
+q_fy(lat)=(-lat+90.0)*4+1
+
+function qt_arrays(q_lon,q_lat,Γ)
+  ni=length(q_lon)
+  nj=length(q_lat)
+  A=zeros(ni+1,nj)
+  B=MeshArray(Γ.XC.grid)
+  (A,B)
+end
+
+function calc_interp(a,A,B,Γ)
+    A[1:end-1,:].=a
+    A[end,:].=a[1,:]
+    itp = interpolate(A, BSpline(Linear()));
+    for i in eachindex(B)
+        B[i][:].=[itp(q_fx(Γ.XC[i][ij]),q_fy(Γ.YC[i][ij])) for ij in eachindex(B[i])]
+    end
+end
+
+storage_arrays(nx,ny,nt)=( y=NaN*zeros(nx,ny),n=NaN*zeros(nx,ny),z=NaN*zeros(nx,ny,nt),zi=NaN*zeros(nx,ny,nt) )
+
+storage_arrays(γ::gcmgrid,L::NamedTuple) = begin
+  (nx,ny)=γ.ioSize
+  (nnx,nny)=(length(L.lon),length(L.lat))
+
+  buffer2d=Array{Union{Missing, Float64}}(undef, nnx, nny)
+  buff2d  =Array{Union{Missing, Float64}}(undef, nnx, nny)
+  A=Array{Float32}(undef, nnx, nny)
+
+  ( y=NaN*zeros(nx,ny),yi=NaN*zeros(nx,ny),n=NaN*zeros(nx,ny) , 
+    buffer2d=buffer2d,buff2d=buff2d,A=A )
+end
+
+## average to 3h
+
+function ave3h(a,buffer2d,buff2d,y,d,r,rec0; path="ERA5/", variable=1) 
+    a.=0.0
+    for h in 1:3
+      rec=(d-1)*24+(r-1)*3+h
+      mo=findall(rec0.<rec)[end]
+      mo<10 ? mon="0$(mo)" : mon="$(mo)"
+
+      fil=joinpath(path,"$(y)","ERA5_$(y)_$(mon).nc")
+      ds=Dataset(fil)
+      jj=rec-rec0[mo]
+      ndims(ds["d2m"])==4 ? ii=(:,:,1,jj) : ii=(:,:,jj)
+      iii=(:,:,2,jj)
+      if variable==11
+        buffer2d.=ds["u10"][ii...]
+        ismissing(buffer2d[1]) ? buffer2d.=ds["u10"][iii...] : nothing
+        a.+= (buffer2d.^2)
+        buffer2d.=ds["v10"][ii...]
+        ismissing(buffer2d[1]) ? buffer2d.=ds["v10"][iii...] : nothing
+        a.+= (buffer2d.^2)
+      else
+        w=list_ds[variable]
+        buffer2d.=ds[w][ii...]
+        ismissing(buffer2d[1]) ? buffer2d.=ds[w][iii...] : nothing
+        a.+= factor[variable] .*(buffer2d .+offset[variable])
+      end
+      close(ds)
+    end
+    a.=a/3
+end
+
+function loop_over_years(variable=1,path="ERA5/")
+  y0=1941
+  ny=83
+  ndmax=366
+
+  γ=GridSpec("LatLonCap",MeshArrays.GRID_LLC90)
+  Γ=GridLoad(γ)
+
+  path_out=joinpath(path,"ERA5_llc90")
+
+  L=setup_arrays_ERA5(Γ)
+  df=DataFrame(:index=>L.knn[:],:x=>L.x[:])
+  S=storage_arrays(γ,L)
+
+  for y in y0.+(0:ny-1)
+    nd=min(Day(DateTime(y+1,1)-DateTime(y,1)).value,ndmax)
+    fil_out=joinpath(path_out,"llc90_ERA5_$(list_in[variable])_$y")
+    println(fil_out)
+    if !isfile(fil_out)
+      g=FortranFiles.FortranFile(fil_out,"w",access="direct",recl=90*1170*4,convert="big-endian")
+      dt=Dates.Day.([[DateTime(y,m+1,1)-DateTime(y,m,1) for m in 1:11];DateTime(y+1,1,1)-DateTime(y,12,1)])
+      rec0=[0;cumsum([a.value for a in dt])*24]
+
+      for d in 1:nd
+        mod(d,10)==0 ? println([y d]) : nothing
+        for r in 1:8
+          ave3h(S.A,S.buffer2d,S.buff2d,y,d,r,rec0; path=path, variable=variable)
+          variable==11 ? S.A.=sqrt.(S.A) : nothing        
+
+          calc_interp(S.A,L.A,L.B,Γ)
+          S.yi.=write(L.B)
+
+          df.x.=S.A[:]
+          gdf=groupby(df,:index)
+          df2=combine(gdf, :x => mean, nrow)
+
+          S.y.=NaN
+          [S.y[df2.index[i]]=df2.x_mean[i] for i in 1:size(df2,1)]
+          [S.y[i]=(isnan(S.y[i]) ? S.yi[i] : S.y[i]) for i in eachindex(S.y)]
+
+          write(g,rec=r+(d-1)*8,Float32.(write(L.B)))
+        end
+      end
+
+      close(g)
+    else
+      println("skipping : already done")
+    end
+  end
+  S
+end
+
+end #module ERA5interp
+

examples/ERA5/ERA5_to_ECCO4.jl

@@ -0,0 +1,13 @@
+
+include("ERA5_module.jl")
+import Main.ERA5_to_ECCO4: loop_over_years
+
+v=1 #parse(Int64, ARGS[1])
+S=loop_over_years(v)
+
+if false
+    using MeshArrays, CairoMakie, DataDeps, JLD2
+    γ=GridSpec("LatLonCap",MeshArrays.GRID_LLC90)
+    λ=interpolation_setup()
+    heatmap(read(S.y-S.yi,γ),interpolation=λ)
+end

examples/ERA5/Project.toml

@@ -0,0 +1,9 @@
+[deps]
+DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
+Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
+FortranFiles = "c58ffaec-ab22-586d-bfc5-781a99fd0b10"
+Interpolations = "a98d9a8b-a2ab-59e6-89dd-64a1c18fca59"
+JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
+MeshArrays = "cb8c808f-1acf-59a3-9d2b-6e38d009f683"
+NCDatasets = "85f8d34a-cbdd-5861-8df4-14fed0d494ab"
+Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"

ext/MeshArraysMakieExt.jl

@@ -119,15 +119,29 @@ heatmap(MS,interpolation=λ,title="ocean depth") #same but w title
 heatmap(MS,x=lon,y=lat) #only for simple domains; will show MS[1]
 ```
 """
-function heatmap(MS::MeshArray;interpolation=nothing,globalmap=false,x=nothing,y=nothing,colorbar=true,title="",kwargs...)	
+heatmap(MS::MeshArray;interpolation=nothing,globalmap=false,x=nothing,y=nothing,colorbar=true,title="",kwargs...)=begin
+	if (!isnothing(interpolation))||globalmap||(!isnothing(x))
+		f=Figure()
+		ax=Axis(f[1,1],title=title)
+		hm1=heatmap!(ax,MS::MeshArray;
+				interpolation=interpolation,globalmap=globalmap,x=x,y=y,
+				kwargs...)
+		colorbar ? Colorbar(f[1,2], hm1, height = Relative(0.65)) : nothing
+		f
+	else
+		heatmap_tiled(MS;colorbar=colorbar,title=title,kwargs...)
+	end
+end
+
+function heatmap!(ax::Axis,MS::MeshArray;interpolation=nothing,globalmap=false,x=nothing,y=nothing,kwargs...)	
 	if !isnothing(interpolation)
-		heatmap_interpolation(MS,interpolation;colorbar=colorbar,title=title,kwargs...)
+		heatmap_interpolation!(ax,MS,interpolation;kwargs...)
 	elseif globalmap
-		heatmap_globalmap(MS;colorbar=colorbar,title=title,kwargs...)
+		heatmap_globalmap!(ax,MS;kwargs...)
 	elseif !isnothing(x)
-		heatmap_xy(MS,x,y;colorbar=colorbar,title=title,kwargs...)
+		heatmap_xy!(ax,MS,x,y;kwargs...)
 	else
-		heatmap_tiled(MS;colorbar=colorbar,title=title,kwargs...)
+		heatmap_tiled(MS;kwargs...)
 	end
 end
 
@@ -139,7 +153,6 @@ function heatmap_globalmap!(ax,MS::MeshArray;kwargs...)
 end
 
 function heatmap_globalmap(MS::MeshArray;title="",colorbar=true,kwargs...)
-
     fig = Figure(size = (900,900), backgroundcolor = :grey95)
     ax = Axis(fig[1,1],xlabel="i index",ylabel="j index",title=title)
 	hm1=heatmap_globalmap!(ax,MS;kwargs...)
@@ -165,7 +178,6 @@ function heatmap_interpolation!(ax,MS::MeshArray,λ::NamedTuple;kwargs...)
 end
 
 function heatmap_interpolation(MS::MeshArray,λ::NamedTuple;title="",colorbar=true,kwargs...)
-
     fig = Figure(size = (900,400), backgroundcolor = :grey95)
     ax = Axis(fig[1,1],xlabel="longitude",ylabel="latitude",title=title)
 	hm1=heatmap_interpolation!(ax,MS,λ;kwargs...)