Merge pull request #6 from gaelforget/ImproveExchOpsPerformance

rewrite exchange methods and streamline operations to improve performance

Author: gaelforget

Project.toml

@@ -1,7 +1,7 @@
 name = "MeshArrays"
 uuid = "28126448-9d44-11e8-08f5-e59c47736fcb"
 authors = ["gaelforget <[email protected]>"]
-version = "0.1.3"
+version = "0.1.4"
 
 [deps]
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"

src/MeshArrays.jl

@@ -12,7 +12,7 @@ include("gcmfaces_convert.jl");
 include("gcmfaces_IO.jl");
 include("gcmfaces_demo.jl");
 
-export gcmfaces, exchange, gradient, smooth, mask, fsize
+export gcmfaces, exchange, gradient, convergence, smooth, mask, fsize
 export GCMGridSpec, GCMGridLoad, GCMGridOnes
 export demo1, demo2
 #The following functions rely on grid specs; currently via global vars.

src/gcmfaces_calc.jl

@@ -1,108 +1,170 @@
 
-function gradient(fld)
+## gradient methods
 
-  FLD=exchange(fld,1);
+function gradient(inFLD::gcmfaces)
+(dFLDdx, dFLDdy)=gradient(inFLD,true)
+return dFLDdx, dFLDdy
+end
+
+function gradient(inFLD::gcmfaces,doDIV::Bool)
+
+exFLD=exchange(inFLD,1)
+dFLDdx=gcmfaces(inFLD.nFaces,inFLD.grTopo)
+dFLDdy=gcmfaces(inFLD.nFaces,inFLD.grTopo)
+
+for a=1:inFLD.nFaces;
+  (s1,s2)=fsize(exFLD,a)
+  tmpA=view(exFLD.f[a],2:s1-1,2:s2-1)
+  tmpB=tmpA-view(exFLD.f[a],1:s1-2,2:s2-1)
+  tmpC=tmpA-view(exFLD.f[a],2:s1-1,1:s2-2)
+  if doDIV
+    dFLDdx.f[a]=tmpB./MeshArrays.DXC.f[a]
+    dFLDdy.f[a]=tmpC./MeshArrays.DYC.f[a]
+  else
+    dFLDdx.f[a]=tmpB
+    dFLDdy.f[a]=tmpC
+  end
+end
 
-  dFLDdx=gcmfaces(fld.nFaces,fld.grTopo);
-  dFLDdy=gcmfaces(fld.nFaces,fld.grTopo);
-  for iFace=1:FLD.nFaces;
-     tmpA=FLD.f[iFace][2:end-1,2:end-1];
-     tmpB=FLD.f[iFace][1:end-2,2:end-1];
-     dFLDdx.f[iFace]=(tmpA-tmpB)./MeshArrays.DXC.f[iFace];
-     tmpA=FLD.f[iFace][2:end-1,2:end-1];
-     tmpB=FLD.f[iFace][2:end-1,1:end-2];
-     dFLDdy.f[iFace]=(tmpA-tmpB)./MeshArrays.DYC.f[iFace];
-  end;
+return dFLDdx, dFLDdy
+end
+
+function gradient(inFLD::gcmfaces,iDXC::gcmfaces,iDYC::gcmfaces)
 
-  return dFLDdx, dFLDdy
+exFLD=exchange(inFLD,1)
+dFLDdx=gcmfaces(inFLD.nFaces,inFLD.grTopo)
+dFLDdy=gcmfaces(inFLD.nFaces,inFLD.grTopo)
 
+for a=1:inFLD.nFaces;
+  (s1,s2)=fsize(exFLD,a)
+  tmpA=view(exFLD.f[a],2:s1-1,2:s2-1)
+  tmpB=tmpA-view(exFLD.f[a],1:s1-2,2:s2-1)
+  tmpC=tmpA-view(exFLD.f[a],2:s1-1,1:s2-2)
+  dFLDdx.f[a]=tmpB.*iDXC.f[a]
+  dFLDdy.f[a]=tmpC.*iDYC.f[a]
 end
 
-##
+return dFLDdx, dFLDdy
+end
+
+## mask methods
 
 function mask(fld::gcmfaces)
-  fldmsk=mask(fld,(NaN,Inf,0.),NaN)
-  return fldmsk
+fldmsk=mask(fld,NaN)
+return fldmsk
 end
 
 function mask(fld::gcmfaces, val::Number)
-  fldmsk=mask(fld,(NaN,Inf,0.),val)
+  fldmsk=gcmfaces(fld.nFaces,fld.grTopo)
+  for a=1:fld.nFaces
+    tmp1=copy(fld.f[a])
+    replace!(x -> !isfinite(x) ? val : x, tmp1 )
+    fldmsk.f[a]=tmp1
+  end
   return fldmsk
 end
 
-function mask(fld::gcmfaces, cond::Number, val::Number)
-  fldmsk=mask(fld,(cond,),val)
+function mask(fld::gcmfaces, val::Number, noval::Number)
+  fldmsk=gcmfaces(fld.nFaces,fld.grTopo)
+  for a=1:fld.nFaces
+    tmp1=copy(fld.f[a])
+    replace!(x -> x==noval ? val : x, tmp1  )
+    fldmsk.f[a]=tmp1
+  end
   return fldmsk
 end
 
-function mask(fld::gcmfaces,cond::Tuple,val::Number)
+## convergence methods
 
-  fldmsk=gcmfaces(fld.nFaces,fld.grTopo);
+function convergence(uFLD::gcmfaces,vFLD::gcmfaces);
 
-  for iFace=1:fld.nFaces;
-     tmp1=fld.f[iFace]
-     for i=1:length(cond)
-       isnan(cond[i]) ? tmp1[findall(isnan.(tmp1))] .= val : nothing
-       isinf(cond[i]) ? tmp1[findall(isinf.(tmp1))] .= val : nothing
-       isfinite(cond[i]) ? tmp1[findall(tmp1 .== cond[i])] .= val : nothing
-     end
-     fldmsk.f[iFace]=tmp1
-  end
+#important note:
+#  Normally uFLD, vFLD should not contain any NaN;
+#  if otherwise then something this may be needed:
+#  uFLD=mask(uFLD,0.0); vFLD=mask(vFLD,0.0);
 
-  return fldmsk
+CONV=gcmfaces(uFLD.nFaces,uFLD.grTopo)
 
+(tmpU,tmpV)=exch_UV(uFLD,vFLD)
+for a=1:tmpU.nFaces
+  (s1,s2)=fsize(uFLD,a)
+  tmpU1=view(tmpU.f[a],1:s1,1:s2)
+  tmpU2=view(tmpU.f[a],2:s1+1,1:s2)
+  tmpV1=view(tmpV.f[a],1:s1,1:s2)
+  tmpV2=view(tmpV.f[a],1:s1,2:s2+1)
+  CONV.f[a]=tmpU1-tmpU2+tmpV1-tmpV2
 end
 
-##
+return CONV
+end
+
+## smooth function
+
+function smooth(FLD::gcmfaces,DXCsm::gcmfaces,DYCsm::gcmfaces)
+
+#important note:
+#input FLD should be land masked (NaN/1) by caller if needed
 
-function smooth(fld::gcmfaces,DXCsm::gcmfaces,DYCsm::gcmfaces)
+#get land masks (NaN/1):
+mskC=fill(1.0,FLD) + 0.0 * mask(FLD)
+(mskW,mskS)=gradient(FLD,false)
+mskW=fill(1.0,FLD) + 0.0 * mask(mskW)
+mskS=fill(1.0,FLD) + 0.0 * mask(mskS)
 
-#get land mask:
-msk=fill(1.,fld) + 0. * mask(fld,NaN);
+#replace NaN with 0. in FLD and land masks:
+FLD=mask(FLD,0.0)
+mskC=mask(mskC,0.0)
+mskW=mask(mskW,0.0)
+mskS=mask(mskS,0.0)
+
+#get inverse grid spacing:
+iDXC=gcmfaces(FLD.nFaces,FLD.grTopo)
+iDYC=gcmfaces(FLD.nFaces,FLD.grTopo)
+for a=1:FLD.nFaces;
+  iDXC.f[a]=1.0./MeshArrays.DXC.f[a]
+  iDYC.f[a]=1.0./MeshArrays.DYC.f[a]
+end
 
 #Before scaling the diffusive operator ...
-tmp0=DXCsm/MeshArrays.DXC;
-tmp0=mask(msk*tmp0,0.);
+tmp0=DXCsm*iDXC*mskW;
 tmp00=maximum(tmp0);
-tmp0=DYCsm/MeshArrays.DYC;
-tmp0=mask(msk*tmp0,0.);
+tmp0=DYCsm*iDYC*mskS;
 tmp00=max(tmp00,maximum(tmp0));
 
 #... determine a suitable time period:
 nbt=ceil(1.1*2*tmp00^2);
 dt=1.;
 T=nbt*dt;
+#println("nbt="*"$nbt")
+
+#diffusion operator times DYG / DXG:
+KuxFac=mskW*DXCsm*DXCsm/T/2.0*MeshArrays.DYG;
+KvyFac=mskS*DYCsm*DYCsm/T/2.0*MeshArrays.DXG;
 
-#diffusion operator:
-Kux=DXCsm*DXCsm/T/2;
-Kvy=DYCsm*DYCsm/T/2;
+#time steping factor:
+dtFac=dt*mskC/MeshArrays.RAC;
 
 #loop:
 for it=1:nbt
-  (dTdxAtU,dTdyAtV)=gradient(fld);
-  tmpU=dTdxAtU*Kux*MeshArrays.DYG;
-  tmpV=dTdyAtV*Kvy*MeshArrays.DXG;
+  (dTdxAtU,dTdyAtV)=gradient(FLD,iDXC,iDYC);
+  tmpU=gcmfaces(FLD.nFaces,FLD.grTopo)
+  tmpV=gcmfaces(FLD.nFaces,FLD.grTopo)
+  for a=1:FLD.nFaces
+      tmpU.f[a]=dTdxAtU.f[a].*KuxFac.f[a];
+      tmpV.f[a]=dTdyAtV.f[a].*KvyFac.f[a];
+  end
   tmpC=convergence(tmpU,tmpV);
-  fld=fld-dt*msk*tmpC/MeshArrays.RAC;
+  for a=1:FLD.nFaces
+      FLD.f[a]=FLD.f[a]-dtFac.f[a].*tmpC.f[a];
+  end
 end
 
-return fld
+#Apply land mask (NaN/1) to end result:
+mskC=mask(mskC,NaN,0.0)
+FLD=mskC*FLD
+
+return FLD
 
 end
 
 ##
-
-function convergence(fldU::gcmfaces,fldV::gcmfaces);
-
-  (tmpU,tmpV)=exch_UV(fldU,fldV);
-  tmpU=mask(tmpU,NaN,0);
-  tmpV=mask(tmpV,NaN,0);
-  tmpC=gcmfaces(tmpU.nFaces,tmpU.grTopo);
-  for iFace=1:tmpU.nFaces;
-    tmp1=tmpU.f[iFace][1:end-1,:,:,:]-tmpU.f[iFace][2:end,:,:,:];
-    tmp2=tmpV.f[iFace][:,1:end-1,:,:]-tmpV.f[iFace][:,2:end,:,:];
-    tmpC.f[iFace]=dropdims(tmp1+tmp2,dims=(3,4));
-  end;
-  return tmpC
-
-end

src/gcmfaces_demo.jl

@@ -63,9 +63,9 @@ Rini=convert2gcmfaces(tmp1);
 
 #apply land mask
 if ndims(MeshArrays.hFacC.f[1])>2
-    tmp1=mask(view(MeshArrays.hFacC,:,:,1),NaN);
+    tmp1=mask(view(MeshArrays.hFacC,:,:,1),NaN,0);
 else
-    tmp1=mask(MeshArrays.hFacC,NaN);
+    tmp1=mask(MeshArrays.hFacC,NaN,0);
 end
 msk=fill(1.,tmp1) + 0. *tmp1;
 Rini=msk*Rini;