Moist pressure on GPU

Examples/testNHSphere.jl

@@ -39,7 +39,7 @@ Microphysics = parsed_args["Microphysics"]
 RelCloud = parsed_args["RelCloud"]
 Rain = parsed_args["Rain"]
 Source = parsed_args["Source"]
-Force = parsed_args["Force"]
+Forcing = parsed_args["Forcing"]
 VerticalDiffusion = parsed_args["VerticalDiffusion"]
 JacVerticalDiffusion = parsed_args["JacVerticalDiffusion"]
 JacVerticalAdvection = parsed_args["JacVerticalAdvection"]
@@ -182,7 +182,7 @@ Model.VerticalDiffusion = VerticalDiffusion
 Model.JacVerticalDiffusion = JacVerticalDiffusion
 Model.JacVerticalAdvection = JacVerticalAdvection
 Model.Source = Source
-Model.Force = Force
+Model.Forcing = Forcing
 Model.Microphysics = Microphysics
 Model.RelCloud = RelCloud
 Model.Rain = Rain
@@ -212,22 +212,27 @@ if Problem == "Galewski"
 elseif Problem == "BaroWaveDrySphere"
   Profile = Examples.BaroWaveExample()(Param,Phys)
 elseif Problem == "HeldSuarezDrySphere"
-  Profile = Examples.HeldSuarezDryExample()(Param,Phys)
+  Model.InitialProfile, Model.Force = Examples.HeldSuarezDryExample()(Param,Phys)
 elseif Problem == "HeldSuarezMoistSphere"
-  Profile, Eddy = Examples.HeldSuarezMoistExample()(Param,Phys)
+  Profile, Force, Eddy = Examples.HeldSuarezMoistExample()(Param,Phys)
   Model.InitialProfile = Profile
+  Model.Force = Force
   Model.Eddy = Eddy
 end  
 
 @show "InitialConditions"
 U = GPU.InitialConditions(backend,FTB,CG,Metric,Phys,Global,Profile,Param)
 
-# Forcing
-if Problem == "HeldSuarezDrySphere" || Problem == "HeldSuarezMoistSphere"
-  Force = Examples.HeldSuarezForcing()(Param,Phys)
-else
-  Force =  Examples.NoForcing()(Param,Phys) 
-end
+# Pressure
+if Equation == "Compressible"
+  Pressure = Models.Compressible()(Phys)
+  Model.Pressure = Pressure
+elseif Equation == "CompressibleMoist"
+  Pressure = Models.CompressibleMoist()(Phys,Model.RhoPos,Model.ThPos,
+    Model.RhoVPos+NumV,Model.RhoCPos+NumV)
+  Model.Pressure = Pressure
+end  
+
 
 # Output
 Global.Output.vtkFileName = string(Problem*"_")
@@ -311,16 +316,16 @@ if Device == "CPU"  || Device == "GPU"
   nT = max(7 + NumTr, NumV + NumTr)
   @show "vor Timestepper"
   Integration.TimeStepper!(U,GPU.FcnGPU!,GPU.FcnPrepareGPU!,DyCore.JacSchurGPU!,
-    Grids.TransSphereX,CG,Metric,Phys,Exchange,Global,Param,Force,DiscType)
+    Grids.TransSphereX,CG,Metric,Phys,Exchange,Global,Param,DiscType)
 elseif Device == "CPU_P"  || Device == "GPU_P"
   Global.ParallelCom.NumberThreadGPU = NumberThreadGPU   
   nT = max(7 + NumTr, NumV + NumTr)
   Parallels.InitExchangeData3D(backend,FTB,nz,nT,Exchange)
   Integration.TimeStepper!(U,GPU.FcnGPU_P!,GPU.FcnPrepareGPU!,DyCore.JacSchurGPU!,
-    Grids.TransSphereX,CG,Metric,Phys,Exchange,Global,Param,Force,DiscType)
+    Grids.TransSphereX,CG,Metric,Phys,Exchange,Global,Param,DiscType)
 else
   nT = max(7 + NumTr, NumV + NumTr)
   Parallels.InitExchangeData3D(backend,FTB,nz,nT,Exchange)
   Integration.TimeStepper!(U,DyCore.Fcn!,DyCore.FcnPrepare!,DyCore.JacSchurGPU!,
-    Grids.TransSphereX,CG,Metric,Phys,Exchange,Global,Param,Force,DiscType)
+    Grids.TransSphereX,CG,Metric,Phys,Exchange,Global,Param,DiscType)
 end

Jobs/JobNHHeldSuarezMoistSphere

@@ -11,13 +11,14 @@ mpirun -n 1 julia --project Examples/testNHSphere.jl \
   --ProfpBGrd="" \
   --ProfRhoBGrd="" \
   --Source=true \
-  --Force=true \
+  --Forcing=true \
   --Curl=false \
   --ModelType="VectorInvariant" \
   --VerticalDiffusion=true \
   --Coriolis=true \
   --Upwind=true \
   --HorLimit=false \
+  --Equation=CompressibleMoist \
   --Buoyancy=true \
   --Decomp="EqualArea" \
   --SimDays=300 \

src/DyCore/Fcn.jl

@@ -164,7 +164,7 @@ function FcnPrepare!(U,CG,Metric,Phys,Cache,Exchange,Global,
   Parallels.ExchangeData3DRecv!(Aux2DG,Exchange)
 end
 
-function Fcn!(F,U,CG,Metric,Phys,Cache,Exchange,Global,Param,Force,
+function Fcn!(F,U,CG,Metric,Phys,Cache,Exchange,Global,Param,
   DiscType::Val{:VectorInvariant})
 
 (;  RhoPos,

src/DyCore/GlobalVariables.jl

@@ -283,7 +283,7 @@ Base.@kwdef mutable struct ModelStruct{FT}
   Thermo::String
   ModelType::String
   Source::Bool
-  Force::Bool
+  Forcing::Bool
   Damping::Bool
   Geos::Bool
   Relax::FT
@@ -314,6 +314,8 @@ Base.@kwdef mutable struct ModelStruct{FT}
   StretchType::String
   InitialProfile::Any
   Eddy::Any
+  Force::Any
+  Pressure::Any
 end
 
 function ModelStruct{FT}() where FT <:AbstractFloat
@@ -344,7 +346,7 @@ function ModelStruct{FT}() where FT <:AbstractFloat
   Thermo = ""
   ModelType = "VectorInvariant"
   Source = false
-  Force = false
+  Forcing = false
   Damping = false
   Geos = false
   Relax = 0.0
@@ -375,6 +377,8 @@ function ModelStruct{FT}() where FT <:AbstractFloat
   StretchType = ""
   InitialProfile = ""
   Eddy = ""
+  Force = ""
+  Pressure = ""
   return ModelStruct{FT}(
    Problem,
    Profile,
@@ -403,7 +407,7 @@ function ModelStruct{FT}() where FT <:AbstractFloat
    Thermo,
    ModelType,
    Source,
-   Force,
+   Forcing,
    Damping,
    Geos,
    Relax,
@@ -434,6 +438,8 @@ function ModelStruct{FT}() where FT <:AbstractFloat
    StretchType,
    InitialProfile,
    Eddy,
+   Force,
+   Pressure,
    )
 end  
 

src/DyCore/parse_commandline.jl

@@ -167,8 +167,8 @@ function parse_commandline()
     arg_type = Bool
     default = false
 
-    "--Force"
-    help = "Force"
+    "--Forcing"
+    help = "Forcing"
     arg_type = Bool
     default = false
 

src/Examples/initial.jl

@@ -220,8 +220,8 @@ end
 
 Base.@kwdef struct HeldSuarezDryExample <: Example end
 
-function (profile::HeldSuarezDryExample)(Param,Phys)
-  function local_profile(x,time)
+function (::HeldSuarezDryExample)(Param,Phys)
+  function profile(x,time)
     FT = eltype(x)
     (Lon,Lat,R)= Grids.cart2sphere(x[1],x[2],x[3])
     z=max(R-Phys.RadEarth,FT(0));
@@ -234,7 +234,25 @@ function (profile::HeldSuarezDryExample)(Param,Phys)
     w = FT(0)
     return (Rho,uS,vS,w,Th)
   end
-  return local_profile
+  function Force(U,p,lat)
+    FT = eltype(U)
+    Sigma = p / Phys.p0
+    height_factor = max(FT(0), (Sigma - Param.sigma_b) / (FT(1) - Param.sigma_b))
+    Fu = -(Param.k_f * height_factor) * U[2]
+    Fv = -(Param.k_f * height_factor) * U[3]
+    if Sigma < FT(0.7)
+      kT = Param.k_a + (Param.k_s - Param.k_a) * height_factor * cos(lat) * cos(lat) * cos(lat) * cos(lat)
+    else
+      kT = FT(0)
+    end
+    Teq = (Param.T_equator - Param.DeltaT_y * sin(lat) * sin(lat) -
+      Param.DeltaTh_z * log(Sigma) * cos(lat) * cos(lat)) * Sigma^Phys.kappa
+    Teq = max(Param.T_min, Teq)
+    DeltaT =  kT * (Phys.p0 * Sigma / (U[1] * Phys.Rd) - Teq)
+    FRhoTh  = -U[1] * DeltaT / Sigma^Phys.kappa
+    return FT(0),Fu,Fv,FT(0),FRhoTh
+  end
+  return profile,Force
 end
 
 Base.@kwdef struct HeldSuarezMoistExample <: Example end
@@ -255,6 +273,24 @@ function (profile::HeldSuarezMoistExample)(Param,Phys)
     qc = FT(0)
     return (Rho,uS,vS,w,Th,qv,qc)
   end
+  function Force(U,p,lat)
+    FT = eltype(U)
+    Sigma = p / Phys.p0
+    height_factor = max(FT(0), (Sigma - Param.sigma_b) / (FT(1) - Param.sigma_b))
+    Fu = -(Param.k_f * height_factor) * U[2]
+    Fv = -(Param.k_f * height_factor) * U[3]
+    if Sigma < FT(0.7)
+      kT = Param.k_a + (Param.k_s - Param.k_a) * height_factor * cos(lat) * cos(lat) * cos(lat) * cos(lat)
+    else
+      kT = FT(0)
+    end
+    Teq = (Param.T_equator - Param.DeltaT_y * sin(lat) * sin(lat) -
+      Param.DeltaTh_z * log(Sigma) * cos(lat) * cos(lat)) * Sigma^Phys.kappa
+    Teq = max(Param.T_min, Teq)
+    DeltaT =  kT * (Phys.p0 * Sigma / (U[1] * Phys.Rd) - Teq)
+    FRhoTh  = -U[1] * DeltaT / Sigma^Phys.kappa
+    return FT(0),Fu,Fv,FT(0),FRhoTh
+  end
   function Eddy(uStar,p,dz)
     K = Param.CE * uStar * dz / 2
     if p < Param.p_pbl
@@ -263,5 +299,5 @@ function (profile::HeldSuarezMoistExample)(Param,Phys)
     end
     return K
   end
-  return local_profile,Eddy
+  return local_profile,Force,Eddy
 end

src/GPU/DiagnosticKernel.jl

@@ -1,20 +1,17 @@
-@kernel function PressureKernel!(p,@Const(U),Phys)
+@kernel function PressureKernel!(Pressure,p,@Const(U))
   Iz,IC = @index(Global, NTuple)
 
   NumG = @uniform @ndrange()[2]
 
   if IC <= NumG
-    @inbounds p[Iz,IC] = PressureGPU(view(U,Iz,IC,:),Phys)
+    @inbounds p[Iz,IC] = Pressure(view(U,Iz,IC,:))
   end
 end
 
-@inline function PressureGPU(U,Phys)
-  @inbounds Phys.p0 * fast_powGPU(Phys.Rd * U[5] / Phys.p0, 1 / (1 - Phys.kappa))
-end
+#@inline function PressureGPU(U,Phys)
+#  @inbounds Phys.p0 * fast_powGPU(Phys.Rd * U[5] / Phys.p0, 1 / (1 - Phys.kappa))
+#end
 
-# we may be hitting a slow path:
-# https://stackoverflow.com/questions/14687665/very-slow-stdpow-for-bases-very-close-to-1
-fast_powGPU(x::FT, y::FT) where {FT <: AbstractFloat} = exp(y * log(x))
 
 @kernel function uStarCoefficientKernel!(uStar,@Const(U),@Const(dXdxI),@Const(nS),@Const(Glob))
   ID,IF = @index(Global, NTuple)
@@ -23,7 +20,7 @@ fast_powGPU(x::FT, y::FT) where {FT <: AbstractFloat} = exp(y * log(x))
 
   if IF <= NumF
     ind = Glob[ID,IF]
-    @inbounds @views uStar[ID,IF] = uStarCoefficientGPU(U[ind,2],U[ind,3],U[ind,4],dXdxI[:,ID,IF],nS[:,ID,IF])
+    @inbounds uStar[ID,IF] = uStarCoefficientGPU(U[1,ind,2],U[1,ind,3],U[1,ind,4],view(dXdxI,3,:,1,ID,1,IF),view(nS,:,ID,IF))
   end
 end
 
@@ -84,15 +81,15 @@ function FcnPrepareGPU!(U,FE,Metric,Phys,Cache,Exchange,Global,Param,DiscType)
   KV = Cache.KV
   dz = Metric.dz
   Eddy = Global.Model.Eddy
+  Pressure = Global.Model.Pressure
 
   KPressureKernel! = PressureKernel!(backend,group)
-  KPressureKernel!(p,U,Phys,ndrange=ndrange)
+  KPressureKernel!(Pressure,p,U,ndrange=ndrange)
 
   if Global.Model.VerticalDiffusion
-    @show "uStar Eddy"  
     KuStarCoefficientKernel! = uStarCoefficientKernel!(backend,groupS)
     KEddyCoefficientKernel! = EddyCoefficientKernel!(backend,groupK)
-    @views KuStarCoefficientKernel!(uStar,U[1,:,:],dXdxI[3,:,1,:,1,:],nS,Glob,ndrange=ndrangeS)
+    KuStarCoefficientKernel!(uStar,U,dXdxI,nS,Glob,ndrange=ndrangeS)
     KEddyCoefficientKernel!(Eddy,KV,uStar,p,dz,Glob,ndrange=ndrangeK)
   end   
 

src/GPU/FcnGPU.jl

@@ -44,7 +44,7 @@ function FcnAdvectionGPU!(F,U,time,FE,Metric,Phys,Cache,Global,Param,Profile)
   KernelAbstractions.synchronize(backend)
 end
 
-function FcnGPU!(F,U,FE,Metric,Phys,Cache,Exchange,Global,Param,Force,DiscType)
+function FcnGPU!(F,U,FE,Metric,Phys,Cache,Exchange,Global,Param,DiscType)
 
   backend = get_backend(F)
   FT = eltype(F)
@@ -67,6 +67,7 @@ function FcnGPU!(F,U,FE,Metric,Phys,Cache,Exchange,Global,Param,Force,DiscType)
   Koeff = Global.Model.HyperDDiv
   Temp1 = Cache.Temp1
   NumberThreadGPU = Global.ParallelCom.NumberThreadGPU
+  Force = Global.Model.Force
 
 
   KoeffCurl = Global.Model.HyperDCurl
@@ -148,18 +149,18 @@ function FcnGPU!(F,U,FE,Metric,Phys,Cache,Exchange,Global,Param,Force,DiscType)
   end    
 
 
-  if Global.Model.Force
+  if Global.Model.Forcing
     NDoFG = min(div(NumberThreadGPU,Nz),NDoF)
     groupG = (Nz, NDoFG)  
     ndrangeG = (Nz, NDoF)  
     KForceKernel! = ForceKernel!(backend, groupG)
-    KForceKernel!(F,U,p,lat,Force,ndrange=ndrangeG)  
+    KForceKernel!(Force,F,U,p,lat,ndrange=ndrangeG)  
     KernelAbstractions.synchronize(backend)
   end  
 
 end
 
-function FcnGPU_P!(F,U,FE,Metric,Phys,Cache,Exchange,Global,Param,Force,DiscType)
+function FcnGPU_P!(F,U,FE,Metric,Phys,Cache,Exchange,Global,Param,DiscType)
 
   backend = get_backend(F)
   FT = eltype(F)
@@ -258,13 +259,13 @@ function FcnGPU_P!(F,U,FE,Metric,Phys,Cache,Exchange,Global,Param,Force,DiscType
   Parallels.ExchangeData3DRecv!(F,Exchange)
   KernelAbstractions.synchronize(backend)
 
-  if Global.Model.Force
+  if Global.Model.Forcingc
     lat = Metric.lat  
     NDoFG = min(div(NumberThreadGPU,Nz),NDoF)
     groupG = (Nz, NDoFG)  
     ndrangeG = (Nz, NDoF)  
     KForceKernel! = ForceKernel!(backend, groupG)
-    KForceKernel!(F,U,p,lat,Force,ndrange=ndrangeG)  
+    KForceKernel!(Force,F,U,p,lat,ndrange=ndrangeG)  
     KernelAbstractions.synchronize(backend)
   end  
 end

src/GPU/OperatorKernel.jl

@@ -1269,7 +1269,7 @@ end
 end
 
 
-@kernel function ForceKernel!(F,U,p,lat,Force)
+@kernel function ForceKernel!(Force,F,U,p,lat)
   Iz,IG = @index(Global, NTuple)
   NG = @uniform @ndrange()[2]
 

src/Integration/RosenbrockSchur.jl

@@ -1,5 +1,5 @@
 function RosenbrockSchur!(V,dt,Fcn!,FcnPrepare!,Jac,CG,Metric,Phys,Cache,JCache,Exchange,
-  Global,Param,Force,DiscType)
+  Global,Param,DiscType)
 
   ROS = Global.TimeStepper.ROS
   nStage = ROS.nStage
@@ -15,7 +15,7 @@ function RosenbrockSchur!(V,dt,Fcn!,FcnPrepare!,Jac,CG,Metric,Phys,Cache,JCache,
       @views @. V = V + ROS.a[iStage,jStage] * k[:,:,:,jStage]
     end
     FcnPrepare!(V,CG,Metric,Phys,Cache,Exchange,Global,Param,DiscType)
-    Fcn!(fV,V,CG,Metric,Phys,Cache,Exchange,Global,Param,Force,DiscType)
+    Fcn!(fV,V,CG,Metric,Phys,Cache,Exchange,Global,Param,DiscType)
     if iStage == 1
       Jac(JCache,V,CG,Metric,Phys,Cache,Global,Param,DiscType)
     end  

src/Integration/RungeKuttaExplicit.jl

@@ -10,7 +10,7 @@ function RungeKuttaExplicit!(V,dt,Fcn!,FcnPrepare!,CG,Metric,Phys,Cache,Global,P
       @views @. V = V + dt * RK.ARKE[iStage,jStage] * f[:,:,:,jStage]
     end
     @views FcnPrepare!(V,CG,Metric,Phys,Cache,Global,Param,DiscType)
-    @views Fcn!(f[:,:,:,iStage],V,CG,Metric,Phys,Cache,Global,Param,Force,DiscType)
+    @views Fcn!(f[:,:,:,iStage],V,CG,Metric,Phys,Cache,Global,Param,DiscType)
   end
   @. V = Vn
   @inbounds for iStage=1:RK.nStage