add set_volume_mixing_ratio for well-mixed gases

NEWS.md

@@ -4,6 +4,21 @@ RRTMGP.jl Release Notes
 main
 ------
 
+v0.22.2
+-------
+- New **`set_volume_mixing_ratio!(solver, name, value)`**, the write counterpart
+  to `volume_mixing_ratio`. It updates a gas's volume mixing ratio in place,
+  correctly for **any** gas under either storage backend, without the caller
+  knowing the storage class: for spatially varying gases (and every gas under
+  per-layer `Vmr` storage) `value` broadcasts over the `(nlay, ncol)` field; for
+  a well-mixed gas under the default global-mean (`VmrGM`) storage it writes the
+  scalar slot. This fills a gap for well-mixed gases (e.g. prescribed,
+  time-varying CO₂): `volume_mixing_ratio` returns those as a **host scalar copy**,
+  so `volume_mixing_ratio(solver, "co2") .= value` did not write back. Spatially
+  varying gases (`"h2o"`, `"o3"`) were already writable through the getter's view
+  and remain so; use whichever reads more clearly (the setter needs no
+  storage-class knowledge; the getter view supports in-place read-modify-write).
+
 v0.22.1
 -------
 - Documented and added a regression test for the **copy-free flux-getter

Project.toml

@@ -1,7 +1,7 @@
 name = "RRTMGP"
 uuid = "a01a1ee8-cea4-48fc-987c-fc7878d79da1"
 authors = ["Climate Modeling Alliance"]
-version = "0.22.1"
+version = "0.22.2"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

docs/src/api.md

@@ -111,6 +111,7 @@ RRTMGP.aerosol_index
 
 ```@docs
 RRTMGP.volume_mixing_ratio
+RRTMGP.set_volume_mixing_ratio!
 VolumeMixingRatios.VolumeMixingRatioGlobalMean
 ```
 

docs/src/getters.md

@@ -31,7 +31,9 @@ the solver's own buffer**, following one invariant:
   have different behavior: `volume_mixing_ratio` for a well-mixed gas (with the
   global-mean `VmrGM` storage) returns a host scalar copied off the device, and
   [`heating_rate`](@ref RRTMGP.heating_rate) computes its result on demand into
-  a fresh array.
+  a fresh array. Because the well-mixed scalar is a copy, write it with
+  [`set_volume_mixing_ratio!`](@ref RRTMGP.set_volume_mixing_ratio!) rather than
+  broadcasting into the getter.
 
 Internal code that needs the full, boundary-extended buffer (rather than the
 domain view) uses the raw struct fields (e.g., `solver.net_flux_buffer`) or
@@ -160,7 +162,7 @@ mirror these: `clear_lw_flux_up`/`clear_lw_flux_dn`/`clear_lw_flux`,
 
 | Getter | Quantity |
 |---|---|
-| [`volume_mixing_ratio`](@ref RRTMGP.volume_mixing_ratio)`(s, name)` | gas volume mixing ratio: a `(nlay, ncol)` view for `"h2o"`/`"o3"`; a host scalar for well-mixed gases with the default global-mean (`VmrGM`) storage. With full per-layer `Vmr` storage, every gas is a `(nlay, ncol)` view |
+| [`volume_mixing_ratio`](@ref RRTMGP.volume_mixing_ratio)`(s, name)` | gas volume mixing ratio: a `(nlay, ncol)` view for `"h2o"`/`"o3"`; a host scalar for well-mixed gases with the default global-mean (`VmrGM`) storage. With full per-layer `Vmr` storage, every gas is a `(nlay, ncol)` view. Write with [`set_volume_mixing_ratio!`](@ref RRTMGP.set_volume_mixing_ratio!)`(s, name, value)` |
 | `center_z`, `face_z` | layer-center / level altitudes [m], or `nothing` if not provided |
 | `deep_atmosphere_inverse_scaling` | deep-atmosphere flux scaling `(nlev, ncol)`, or `nothing` |
 | [`radiation_method`](@ref RRTMGP.radiation_method) | the solver's radiation method |

src/api/getters.jl

@@ -314,3 +314,48 @@ function _volume_mixing_ratio(
 end
 _volume_mixing_ratio(s::RRTMGPSolver, vmr::VolumeMixingRatios.Vmr, name::AbstractString) =
     _domain_view(s, view(vmr.vmr, _lookup_tables(s).idx_gases_sw[name], :, :))
+
+"""
+    set_volume_mixing_ratio!(s::RRTMGPSolver, name::AbstractString, value)
+
+Set the volume mixing ratio for gas `name` to `value`, returning `value`.
+
+This is the write counterpart to [`volume_mixing_ratio`](@ref). For `"h2o"`/`"o3"`
+(and, with per-layer `Vmr` storage, every gas), `value` broadcasts over the
+`(nlay, ncol)` field, so it may be a scalar or a domain-sized array. For a
+well-mixed gas with the default global-mean (`VmrGM`) storage, the mixing ratio
+is a single scalar; pass a scalar `value`.
+
+Unlike `volume_mixing_ratio`, this is the supported way to update a well-mixed
+gas: that getter returns a read-only host copy, so `volume_mixing_ratio(s, name) .= value`
+does not write back. Use this to update time-varying trace gases (e.g. prescribed CO₂).
+
+Available names are:
+$(gas_names_sw_docs())
+"""
+set_volume_mixing_ratio!(s::RRTMGPSolver, name::AbstractString, value) =
+    _set_volume_mixing_ratio!(s, _atmospheric_state(s).vmr, name, value)
+
+function _set_volume_mixing_ratio!(
+    s::RRTMGPSolver,
+    vmr::VolumeMixingRatios.VmrGM,
+    name::AbstractString,
+    value,
+)
+    name == "h2o" && return (_domain_view(s, vmr.vmr_h2o) .= value; value)
+    name == "o3" && return (_domain_view(s, vmr.vmr_o3) .= value; value)
+    idx = _lookup_tables(s).idx_gases_sw[name]
+    # The water-vapor continuum pseudo-gases share the h2o slot (see the getter);
+    # route them to `vmr_h2o` rather than the unused well-mixed slot 1.
+    idx == 1 && return (_domain_view(s, vmr.vmr_h2o) .= value; value)
+    # A well-mixed gas is a single value. `fill!` over a 1-element view writes the
+    # scalar without scalar-indexing the GPU (a bare `vmr.vmr[idx] = value` would).
+    fill!(view(vmr.vmr, idx:idx), value)
+    return value
+end
+_set_volume_mixing_ratio!(
+    s::RRTMGPSolver,
+    vmr::VolumeMixingRatios.Vmr,
+    name::AbstractString,
+    value,
+) = (_domain_view(s, view(vmr.vmr, _lookup_tables(s).idx_gases_sw[name], :, :)) .= value; value)

test/standalone_spectral.jl

@@ -76,6 +76,17 @@ using NCDatasets
     @test co2 isa Number
     @test co2 ≈ 420e-6
 
+    # set_volume_mixing_ratio! writes a well-mixed gas back (the getter is a copy,
+    # so broadcasting into it would not persist); read it back through the getter
+    RRTMGP.set_volume_mixing_ratio!(solver, "co2", FT(500e-6))
+    @test RRTMGP.volume_mixing_ratio(solver, "co2") ≈ 500e-6
+    # ... and writes a spatially varying gas in place (returns the value it set)
+    @test RRTMGP.set_volume_mixing_ratio!(solver, "h2o", FT(0.02)) == FT(0.02)
+    @test all(==(FT(0.02)), Array(RRTMGP.volume_mixing_ratio(solver, "h2o")))
+    # ... the continuum pseudo-gases route to the shared h2o slot
+    RRTMGP.set_volume_mixing_ratio!(solver, "h2o_self", FT(0.03))
+    @test all(==(FT(0.03)), Array(RRTMGP.volume_mixing_ratio(solver, "h2o")))
+
     # per-gas Vmr storage returns a (nlay, ncol) view for every gas
     ngas = solver.lookups.ngas_sw
     vmr_all = RRTMGP.VolumeMixingRatios.Vmr(zeros(FT, ngas, nlay, 1))
@@ -84,4 +95,7 @@ using NCDatasets
     co2_view = RRTMGP._volume_mixing_ratio(solver, vmr_all, "co2")
     @test size(co2_view) == (nlay, 1)
     @test all(==(FT(400e-6)), Array(co2_view))
+    # the Vmr setter branch broadcasts a scalar over the (nlay, ncol) field
+    RRTMGP._set_volume_mixing_ratio!(solver, vmr_all, "co2", FT(410e-6))
+    @test all(==(FT(410e-6)), Array(co2_view))
 end