Profiling-driven fixes (2026-07-25)

Project.toml

@@ -1,6 +1,6 @@
 name = "MeshArrays"
 uuid = "cb8c808f-1acf-59a3-9d2b-6e38d009f683"
-version = "0.5.8"
+version = "0.5.9"
 authors = ["gaelforget <gforget@mit.edu>"]
 
 [deps]

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.<lats[l+1]))].=l for l in 1:nl]
+
+  for l in 1:nl
+    for f in eachindex(mask.fIndex)
+        mf = mask.f[f]
+        laf = la.f[f]
+        for j in axes(mf,2), i in axes(mf,1)
+            if mf[i,j] > 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.<lats[i_a+1])
-	 *(lo.>=lons[i_o])*(lo.<lons[i_o+1]))].=length(name)
+	 *(lo.>=lons[i_o])*(lo.<lons[i_o+1]))]=length(name)
     end
    end
   end
@@ -90,75 +102,140 @@ layer_mask(dF,d0,d1)=begin
   md
 end
 
+"""
+    func_h_inner(X, msk)
+
+Zero-allocation horizontal sum: replaces sum(xymsk(b)*X).
+msk is pre-computed once per basin in define_sums.
+"""
+function func_h_inner(X::AbstractMeshArray, msk::AbstractMeshArray)
+    s = 0.0
+    nf = length(X.fIndex)
+    for f in 1:nf
+        Xf = X.f[f]
+        mf = msk.f[f]
+        @inbounds for j in axes(Xf,2), i in axes(Xf,1)
+            s += mf[i,j] * Xf[i,j]
+        end
+    end
+    s
+end
+
+"""
+    func_inner(X, msk, lmsk_d, grid, nr)
+
+Zero-allocation depth-integrated sum.
+- msk:    pre-computed horizontal mask, once per basin
+- lmsk_d: pre-computed layer_mask vector for one depth range (length nr)
+Replaces the loop over k with X[:,k] and hFacC[:,k] gcmarray slicing.
+"""
+function func_inner(X::AbstractMeshArray, msk::AbstractMeshArray,
+                    lmsk_d, grid::NamedTuple, nr::Int)
+    s = 0.0
+    nf = length(X.fIndex)
+    for k in 1:nr
+        zf = lmsk_d[k] * grid.DRF[k]
+        iszero(zf) && continue
+        for f in 1:nf
+            Xf = X.f[f, k]
+            hf = grid.hFacC.f[f, k]
+            mf = msk.f[f]
+            Rf = grid.RAC.f[f]
+            @inbounds for j in axes(Xf,2), i in axes(Xf,1)
+                s += zf * mf[i,j] * Xf[i,j] * hf[i,j] * Rf[i,j]
+            end
+        end
+    end
+    s
+end
+
 """
     define_sums(;option=:loops, grid::NamedTuple, regions=:global, depths=[(0,7000)])
 
 Define regional integration function for each basin and depth range.
 """
 function define_sums(;option=:loops, grid::NamedTuple, regions=:global, 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)
+  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.mask.==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])
-  tmp3d=MeshArray(grid.XC.grid,Float32,nr)
+  # pre-compute layer masks once per depth range — avoids re-allocation in hot loop
+  lmsk_cache = [layer_mask(grid.RF, d0, d1) for (d0,d1) in dep]
 
-  function func_v(X,d0,d1)
-    tmp2d.=0.0
+  tmp2d = MeshArray(grid.XC.grid, Float32)
+
+  function func_v_inner!(out::AbstractMeshArray, X::AbstractMeshArray,
+                       lmsk_d, grid::NamedTuple, nr::Int)
+    for f in 1:length(out.fIndex)
+        of = out.f[f]
+        @inbounds for j in axes(of,2), i in axes(of,1)
+            of[i,j] = 0.0
+        end
+    end
     for k in 1:nr
-      tmp2d.+=zmsk(d0,d1,k)*X[:,k]*grid.DRF[k]*grid.hFacC[:,k]*grid.RAC
+        zf = lmsk_d[k] * grid.DRF[k]
+        iszero(zf) && continue
+        for f in 1:length(out.fIndex)
+            of = out.f[f]
+            Xf = X.f[f,k]
+            hf = grid.hFacC.f[f,k]
+            Rf = grid.RAC.f[f]
+            @inbounds for j in axes(Xf,2), i in axes(Xf,1)
+                of[i,j] += zf * Xf[i,j] * hf[i,j] * Rf[i,j]
+            end
+        end
     end
-    tmp2d
-  end
+    out
+end
+
+  tmp3d = MeshArray(grid.XC.grid, Float32, nr)
 
   if option==:streamlined_loop
-  #ocn_surf=[sum(xymsk(b)*(grid.hFacC[:,1].>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

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).<mod(range[2].-range[1],360))
 
-function restrict_longitudes!(x::AbstractMeshArray,lon::AbstractMeshArray;range=(0.0,360.0))
+function restrict_longitudes!(x::AbstractMeshArray, lon::AbstractMeshArray;
+                               range=(0.0,360.0))
   for f in 1:x.grid.nFaces
-    x[f].=x[f].*is_in_lon_range.(lon[f],Ref(range))
+      xf = x.f[f]; lf = lon.f[f]
+      @inbounds for j in axes(xf,2), i in axes(xf,1)
+          xf[i,j] *= is_in_lon_range(lf[i,j], range)
+      end
   end
 end
 
 function MskToTab(msk::AbstractMeshArray)
-  n=Int(sum(msk .!= 0)); k=0
-  tab=Array{Int,2}(undef,n,4)
-  for i in eachindex(msk)
-    a=msk[i]
-    b=findall( a .!= 0)
-    for ii in eachindex(b)
-      k += 1
-      tab[k,:]=[i,b[ii][1],b[ii][2],a[b[ii]]]
-    end
+  # count without allocating
+  n = 0
+  for f in 1:length(msk.fIndex)
+      mf = msk.f[f]
+      @inbounds for j in axes(mf,2), i in axes(mf,1)
+          n += (mf[i,j] != 0)
+      end
+  end
+
+  k = 0
+  tab = Array{Int,2}(undef, n, 4)
+  for f in 1:length(msk.fIndex)
+      mf = msk.f[f]
+      @inbounds for j in axes(mf,2), i in axes(mf,1)
+          v = mf[i,j]
+          if v != 0
+              k += 1
+              tab[k,1] = f
+              tab[k,2] = i
+              tab[k,3] = j
+              tab[k,4] = Int(v)
+          end
+      end
   end
   return tab
 end
@@ -532,34 +558,29 @@ Compute edge mask (mskC,mskW,mskS) from domain interior mask (mskCint).
 This is used in `LatitudeCircles` and `Transect`.
 """
 function edge_mask(mskCint::AbstractMeshArray)
-  mskCint=1.0*mskCint
-
-  #treat the case of blank tiles:
-  #mskCint[findall(RAC.==0)].=NaN
-
-  mskC=similar(mskCint)
-  mskW=similar(mskCint)
-  mskS=similar(mskCint)
-
-  mskCint=exchange(mskCint).MA
-
-  for i in eachindex(mskCint)
-      tmp1=mskCint[i]
-      # tracer mask:
-      tmp2=tmp1[2:end-1,1:end-2]+tmp1[2:end-1,3:end]+
-      tmp1[1:end-2,2:end-1]+tmp1[3:end,2:end-1]
-      mskC[i]=1((tmp2.>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
 
 ##

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