Merge pull request #65 from nichollsh/optim

Model optimisations

Author: nichollsh

CITATION.cff

@@ -5,7 +5,7 @@ authors:
   given-names: "Harrison"
   orcid: "https://orcid.org/0000-0002-8368-4641"
 title: "AGNI"
-version: 0.7.0
+version: 0.8.0
 doi: 10.xx/xx.xx
 date-released: 2024-09-11
 url: "https://github.com/nichollsh/AGNI"

Project.toml

@@ -1,7 +1,7 @@
 name = "AGNI"
 uuid = "ede838c1-9ec3-4ebe-8ae8-da4091b3f21c"
 authors = ["Harrison Nicholls <harrison.nicholls@physics.ox.ac.uk>"]
-version = "0.7.0"
+version = "0.8.0"
 
 [deps]
 ArgParse = "c7e460c6-2fb9-53a9-8c5b-16f535851c63"
@@ -14,6 +14,7 @@ LaTeXStrings = "b964fa9f-0449-5b57-a5c2-d3ea65f4040f"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
 LoggingExtras = "e6f89c97-d47a-5376-807f-9c37f3926c36"
+LoopVectorization = "bdcacae8-1622-11e9-2a5c-532679323890"
 NCDatasets = "85f8d34a-cbdd-5861-8df4-14fed0d494ab"
 OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
 PCHIPInterpolation = "afe20452-48d1-4729-9a8b-50fb251f06cd"

res/config/55cnce_chem.toml

@@ -7,7 +7,7 @@ title = "Roughly 55 Cancri e @ fO2=IW"
     albedo_b        = 0.0
     s0_fact         = 0.6652
     zenith_angle    = 60.0
-    surface_material= "res/surface_albedos/c9mb29.dat"
+    surface_material= "res/surface_albedos/basalt_tuff.dat"
     radius          = 1.1959e7
     gravity         = 22.304
     tmp_int         = 0.0

res/config/default.toml

@@ -46,7 +46,7 @@ title = "Default"                       # Name for this configuration file
     cloud           = false                     # Include water cloud radiative properties?
     aerosol         = false                     # Include aerosol radiative properties?
     overlap_method  = 4                         # Method for treating line overlap.
-    thermo_funct    = false                     # Use temperature-dependent thermodynamic properties?
+    thermo_funct    = true                     # Use temperature-dependent thermodynamic properties?
     sensible_heat   = false                     # Include sensible heat transport at the surface?
     latent_heat     = false                     # Include heat release from phase change
     convection      = true                      # Include heat transport by convection

src/atmosphere.jl

@@ -15,6 +15,7 @@ module atmosphere
     using LinearAlgebra
     using Statistics
     using Logging
+    using LoopVectorization
 
     # Local files
     include(joinpath(ENV["RAD_DIR"],"julia/src/SOCRATES.jl"))
@@ -73,10 +74,12 @@ module atmosphere
         grav_surf::Float64              # Surface gravity [m s-2]
         overlap_method::Int             # Absorber overlap method to be used
 
-        # Band edges
+        # Spectral bands
         nbands::Int
         bands_min::Array{Float64,1}    # Lower wavelength [m]
         bands_max::Array{Float64,1}    # Upper wavelength [m]
+        bands_cen::Array{Float64,1}    # Midpoint [m]
+        bands_wid::Array{Float64,1}    # Width [m]
 
         # Pressure-temperature grid (with i=1 at the top of the model)
         nlev_c::Int         # Cell centre (count)
@@ -151,6 +154,9 @@ module atmosphere
         band_u_sw::Array{Float64,2}         # up component, sw
         band_n_sw::Array{Float64,2}         # net upward, sw
 
+        # Surface planck emission (incl. emissivity)
+        surf_flux::Array{Float64, 1}
+
         # Contribution function (to outgoing flux) per-band
         contfunc_norm::Array{Float64,2}     # LW only, and normalised by maximum value
 
@@ -313,7 +319,7 @@ module atmosphere
         end
 
         # Code versions
-        atmos.AGNI_VERSION = "0.7.0"
+        atmos.AGNI_VERSION = "0.8.0"
         atmos.SOCRATES_VERSION = readchomp(joinpath(ENV["RAD_DIR"],"version"))
         @debug "AGNI VERSION = $(atmos.AGNI_VERSION)"
         @debug "Using SOCRATES at $(ENV["RAD_DIR"])"
@@ -572,7 +578,7 @@ module atmosphere
         # backup mixing ratios from current state
         for k in keys(atmos.gas_vmr)
             atmos.gas_ovmr[k] = zeros(Float64, atmos.nlev_c)
-            atmos.gas_ovmr[k][:] .= atmos.gas_vmr[k][:]
+            @turbo @. atmos.gas_ovmr[k] = atmos.gas_vmr[k]
         end
 
         # set condensation mask and yield values [kg]
@@ -729,59 +735,45 @@ module atmosphere
     """
     function calc_layer_props!(atmos::atmosphere.Atmos_t; ignore_errors::Bool=false)::Bool
         if !atmos.is_param
-            error("atmosphere parameters have not been set")
+            error("Atmosphere parameters have not been set")
         end
 
-        ok::Bool = true
-        dz_max::Float64 = 1e9
-
-        # Set pressure arrays in SOCRATES
-        atmos.atm.p[1, :] .= atmos.p[:]
-        atmos.atm.p_level[1, 0:end] .= atmos.pl[:]
-
-        # Set MMW at each level
-        fill!(atmos.layer_mmw, 0.0)
-        for i in 1:atmos.nlev_c
-            for gas in atmos.gas_names
-                atmos.layer_mmw[i] += atmos.gas_vmr[gas][i] * atmos.gas_dat[gas].mmw
-            end
-        end
-
-        # Temporary values
-        mmr::Float64 = 0.0
-        g1::Float64 = 0.0; p1::Float64 = 0.0; t1::Float64 = 0.0
-        g2::Float64 = 0.0; p2::Float64 = 0.0; t2::Float64 = 0.0
-        dzc::Float64= 0.0; dzl::Float64 = 0.0
-        GMpl::Float64 = atmos.grav_surf * (atmos.rp^2.0)
-
         # Reset arrays
         fill!(atmos.z         ,  0.0)
         fill!(atmos.zl        ,  0.0)
         fill!(atmos.layer_grav,  0.0)
         fill!(atmos.layer_thick, 0.0)
         fill!(atmos.layer_density,0.0)
-        fill!(atmos.layer_cp     ,0.0)
-        fill!(atmos.layer_kc     ,0.0)
         fill!(atmos.layer_mass   ,0.0)
+        fill!(atmos.layer_cp, 0.0)
+        fill!(atmos.layer_kc, 0.0)
+        fill!(atmos.layer_mmw, 0.0)
 
-        # Integrate from bottom upwards
-        for i in range(start=atmos.nlev_c, stop=1, step=-1)
+        # Temporary values
+        g1::Float64 = 0.0; p1::Float64 = 0.0; t1::Float64 = 0.0
+        g2::Float64 = 0.0; p2::Float64 = 0.0; t2::Float64 = 0.0
+        dzc::Float64= 0.0; dzl::Float64 = 0.0
+        GMpl::Float64 = atmos.grav_surf * (atmos.rp^2.0)
+        mmr::Float64 = 0.0
+        ok::Bool = true
+        dz_max::Float64 = 1e8
 
-            # Set cp, kc at this level
-            atmos.layer_cp[i] = 0.0
-            atmos.layer_kc[i] = 0.0
-            for gas in atmos.gas_names
+        # Set MMW at each level
+        for gas in atmos.gas_names
+            @turbo @. atmos.layer_mmw += atmos.gas_vmr[gas] * atmos.gas_dat[gas].mmw
+        end
+
+        # Set cp, kc at each level
+        for gas in atmos.gas_names
+            @inbounds for i in 1:atmos.nlev_c
                 mmr = atmos.gas_vmr[gas][i] * atmos.gas_dat[gas].mmw/atmos.layer_mmw[i]
                 atmos.layer_cp[i] += mmr * phys.get_Cp(atmos.gas_dat[gas], atmos.tmp[i])
                 atmos.layer_kc[i] += mmr * phys.get_Kc(atmos.gas_dat[gas], atmos.tmp[i])
             end
+        end
 
-            # Temporarily copy this cp, kc to the level above
-            #     since they are needed for doing the hydrostatic integration
-            if i > 1
-                atmos.layer_cp[i-1] = atmos.layer_cp[i]
-                atmos.layer_kc[i-1] = atmos.layer_kc[i]
-            end
+        # Integrate from bottom upwards
+        for i in range(start=atmos.nlev_c, stop=1, step=-1)
 
             # Technically, g and z should be integrated as coupled equations,
             # but here they are not. This loose integration has been found to
@@ -832,7 +824,7 @@ module atmosphere
 
         # Mass (per unit area, kg m-2) and density (kg m-3)
         # Loop from top downards
-        for i in 1:atmos.nlev_c
+        @inbounds for i in 1:atmos.nlev_c
             atmos.layer_mass[i] = (atmos.pl[i+1] - atmos.pl[i])/atmos.layer_grav[i]
             atmos.atm.mass[1, i] = atmos.layer_mass[i]          # pass to SOCRATES
 
@@ -885,8 +877,8 @@ module atmosphere
         atmos.p[1:end] .= 0.5 .* (atmos.pl[1:end-1] .+ atmos.pl[2:end])
 
         # Finally, convert arrays to 'real' pressure units
-        atmos.p[:]  .= 10.0 .^ atmos.p[:]
-        atmos.pl[:] .= 10.0 .^ atmos.pl[:]
+        @turbo @. atmos.p  = 10.0 ^ atmos.p
+        @turbo @. atmos.pl = 10.0 ^ atmos.pl
 
         return nothing
     end
@@ -990,13 +982,17 @@ module atmosphere
         end
 
         atmos.nbands = atmos.spectrum.Basic.n_band
-        atmos.bands_max = zeros(Float64, atmos.spectrum.Basic.n_band)
-        atmos.bands_min = zeros(Float64, atmos.spectrum.Basic.n_band)
+        atmos.bands_max = zeros(Float64, atmos.nbands)
+        atmos.bands_min = zeros(Float64, atmos.nbands)
+        atmos.bands_cen = zeros(Float64, atmos.nbands)
+        atmos.bands_wid = zeros(Float64, atmos.nbands)
 
         for i in 1:atmos.nbands
             atmos.bands_min[i] = atmos.spectrum.Basic.wavelength_short[i]
             atmos.bands_max[i] = atmos.spectrum.Basic.wavelength_long[i]
         end
+        @turbo @. atmos.bands_cen = 0.5 * (atmos.bands_max + atmos.bands_min)
+        @turbo @. atmos.bands_wid = abs(atmos.bands_max - atmos.bands_min)
 
         # modules_gen/dimensions_field_cdf_ucf.f90
         npd_direction = 1                   # Maximum number of directions for radiances
@@ -1065,7 +1061,7 @@ module atmosphere
         SOCRATES.allocate_bound(atmos.bound, atmos.dimen, atmos.spectrum)
 
         # Fill with zeros - will be set inside of radtrans function at call time
-        atmos.bound.flux_ground[1, :] .= 0.0
+        fill!(atmos.bound.flux_ground, 0.0)
 
 
         ###########################################
@@ -1103,7 +1099,7 @@ module atmosphere
         end
 
         # Calculate the weighting for the bands.
-        atmos.control.weight_band .= 1.0
+        fill!(atmos.control.weight_band, 1.0)
 
         # 'Entre treatment of optical depth for direct solar flux (0/1/2)'
         # '0: no scaling; 1: delta-scaling; 2: circumsolar scaling'
@@ -1318,17 +1314,17 @@ module atmosphere
             push!(_alb_s, _alb_s[end])
 
             # convert ss albedo to gamma values (eq 14.3 from Hapke 2012)
-            _alb_s[:] .= sqrt.(1 .- _alb_s[:]) # operating in place
+            @turbo @. _alb_s = sqrt(1.0 - _alb_s) # operating in place
 
             # create interpolator on gamma
             _gamma::Interpolator = Interpolator(_alb_w, _alb_s)
 
             # use interpolator to fill band values
             ga::Float64 = 0.0
             mu::Float64 = cosd(atmos.zenith_degrees)
-            for i in 1:atmos.nbands
+            @inbounds for i in 1:atmos.nbands
                 # evaluate gamma at band centre, converting from m to nm
-                ga = _gamma(0.5 * 1.0e9 * (atmos.bands_min[i]+atmos.bands_max[i]))
+                ga = _gamma(1.0e9 * atmos.bands_cen[i])
 
                 # calculate dh reflectance (eq 3 from Hammond24)
                 atmos.surf_r_arr[i] = (1-ga) / (1+2*ga*mu)
@@ -1361,6 +1357,8 @@ module atmosphere
         atmos.band_u_sw =         zeros(Float64, (atmos.nlev_l,atmos.nbands))
         atmos.band_n_sw =         zeros(Float64, (atmos.nlev_l,atmos.nbands))
 
+        atmos.surf_flux =         zeros(Float64, atmos.nbands)
+
         atmos.contfunc_norm =     zeros(Float64, (atmos.nlev_c,atmos.nbands))
 
         atmos.flux_sens =         0.0
@@ -1614,7 +1612,7 @@ module atmosphere
             # matched?
             if match
                 N_g = data[i,:]  # number densities for this gas
-                atmos.gas_vmr[g_in][:] .+= N_g[:] ./ N_t[:]    # VMR for this gas
+                @. atmos.gas_vmr[g_in] += N_g / N_t    # VMR for this gas
             end
         end
 
@@ -1724,9 +1722,9 @@ module atmosphere
         # Reset phase change flags
         # Reset condensation yield values
         for g in atmos.gas_names
-           atmos.gas_vmr[g][1:end-1] .= atmos.gas_vmr[g][end]
-           atmos.gas_sat[g][:] .= false
-           atmos.gas_yield[g][:] .= 0.0
+           fill!(atmos.gas_vmr[g][1:end-1], atmos.gas_vmr[g][end])
+           fill!(atmos.gas_sat[g],          false)
+           fill!(atmos.gas_yield[g],        0.0)
         end
 
         # Reset water cloud
@@ -1990,7 +1988,7 @@ module atmosphere
         if back_interp
             clamp!(atmos.tmpl, atmos.tmp_floor, atmos.tmp_ceiling)
             itp = Interpolator(log.(atmos.pl), atmos.tmpl)
-            atmos.tmp[:] .= itp.(log.(atmos.p[:]))
+            @. atmos.tmp = itp(log(atmos.p))
         end
 
         return nothing

src/dump.jl

@@ -1,4 +1,4 @@
-# Contains save/load module
+# Write atmosphere to NetCDF or CSV file
 
 # Not for direct execution
 if (abspath(PROGRAM_FILE) == @__FILE__)