Implement a KinematicModel based on PrescribedVelocityFields

src/AtmosphereModels/atmosphere_model.jl

@@ -107,6 +107,7 @@ function AtmosphereModel(grid;
                          thermodynamic_constants = ThermodynamicConstants(eltype(grid)),
                          formulation = :LiquidIcePotentialTemperature,
                          dynamics = nothing,
+                         velocities = nothing,
                          moisture_density = DefaultValue(),
                          tracers = tuple(),
                          coriolis = nothing,
@@ -158,7 +159,17 @@ function AtmosphereModel(grid;
     dynamics = materialize_dynamics(dynamics, grid, boundary_conditions)
     formulation = materialize_formulation(formulation, dynamics, grid, boundary_conditions)
 
-    momentum, velocities = materialize_momentum_and_velocities(dynamics, grid, boundary_conditions)
+    # Materialize momentum and velocities
+    # If velocities is provided (e.g., PrescribedVelocityFields), use it
+    if isnothing(velocities)
+        momentum, velocities = materialize_momentum_and_velocities(dynamics, grid, boundary_conditions)
+    else
+        # Store velocity specification in dynamics for dispatch (e.g., PrescribedVelocityFields)
+        dynamics = update_dynamics_with_velocities(dynamics, velocities)
+        momentum, _ = materialize_momentum_and_velocities(dynamics, grid, boundary_conditions)
+        velocities = materialize_velocities(velocities, grid)
+    end
+
     microphysical_fields = materialize_microphysical_fields(microphysics, grid, boundary_conditions)
 
     tracers = NamedTuple(name => CenterField(grid, boundary_conditions=boundary_conditions[name]) for name in tracer_names)
@@ -272,10 +283,10 @@ Advection.cell_advection_timescale(model::AtmosphereModel) = cell_advection_time
 
 # Prognostic field names from dynamics + thermodynamic formulation + microphysics + tracers
 function prognostic_field_names(dynamics, formulation, microphysics, tracer_names)
-    default_names = (:ρu, :ρv, :ρw, :ρqᵗ)
+    momentum_names = prognostic_momentum_field_names(dynamics)
     formulation_names = prognostic_thermodynamic_field_names(formulation)
     microphysical_names = prognostic_field_names(microphysics)
-    return tuple(default_names..., formulation_names..., microphysical_names..., tracer_names...)
+    return tuple(momentum_names..., :ρqᵗ, formulation_names..., microphysical_names..., tracer_names...)
 end
 
 function field_names(dynamics, formulation, microphysics, tracer_names)

src/AtmosphereModels/dynamics_interface.jl

@@ -31,6 +31,21 @@ Create momentum and velocity fields for the given dynamics.
 """
 function materialize_momentum_and_velocities end
 
+"""
+    materialize_velocities(velocities, grid)
+
+Create velocity fields from a velocity specification (e.g., `PrescribedVelocityFields`).
+"""
+function materialize_velocities end
+
+"""
+    update_dynamics_with_velocities(dynamics, velocities)
+
+Update dynamics with velocity specification. Default is a no-op.
+For `PrescribedDynamics`, stores the `PrescribedVelocityFields` for dispatch.
+"""
+update_dynamics_with_velocities(dynamics, velocities) = dynamics
+
 """
     dynamics_pressure_solver(dynamics, grid)
 
@@ -140,6 +155,16 @@ initialize_model_thermodynamics!(model) = nothing  # default: do nothing
 ##### Prognostic fields interface
 #####
 
+"""
+    prognostic_momentum_field_names(dynamics)
+
+Return a tuple of prognostic momentum field names.
+
+For prognostic dynamics (anelastic, compressible), returns `(:ρu, :ρv, :ρw)`.
+For kinematic dynamics (prescribed velocities), returns an empty tuple.
+"""
+prognostic_momentum_field_names(::Any) = (:ρu, :ρv, :ρw)
+
 """
     prognostic_dynamics_field_names(dynamics)
 

src/AtmosphereModels/set_atmosphere_model.jl

@@ -28,6 +28,36 @@ end
 const settable_thermodynamic_variables = (:ρθ, :θ, :ρθˡⁱ, :θˡⁱ, :ρe, :e, :T)
 function set_thermodynamic_variable! end
 
+#####
+##### Velocity and momentum setting (extensible for kinematic models)
+#####
+
+"""
+    set_velocity!(model, name, value)
+
+Set the velocity component `name` (`:u`, `:v`, or `:w`) to `value`.
+Also updates the corresponding momentum field.
+"""
+function set_velocity!(model::AtmosphereModel, name::Symbol, value)
+    u = model.velocities[name]
+    set!(u, value)
+    ρ = dynamics_density(model.dynamics)
+    ϕ = model.momentum[Symbol(:ρ, name)]
+    set!(ϕ, ρ * u)
+    return nothing
+end
+
+"""
+    set_momentum!(model, name, value)
+
+Set the momentum component `name` (`:ρu`, `:ρv`, or `:ρw`) to `value`.
+"""
+function set_momentum!(model::AtmosphereModel, name::Symbol, value)
+    ρu = getproperty(model.momentum, name)
+    set!(ρu, value)
+    return nothing
+end
+
 """
 $(TYPEDSIGNATURES)
 
@@ -130,8 +160,7 @@ function Fields.set!(model::AtmosphereModel; time=nothing, enforce_mass_conserva
 
         # Prognostic variables
         if name ∈ propertynames(model.momentum)
-            ρu = getproperty(model.momentum, name)
-            set!(ρu, value)
+            set_momentum!(model, name, value)
 
         elseif name ∈ propertynames(model.tracers)
             c = getproperty(model.tracers, name)
@@ -163,13 +192,7 @@ function Fields.set!(model::AtmosphereModel; time=nothing, enforce_mass_conserva
             set!(ρqᵗ, ρ * qᵗ)
 
         elseif name ∈ (:u, :v, :w)
-            u = model.velocities[name]
-            set!(u, value)
-
-            ρ = dynamics_density(model.dynamics)
-            ϕ = model.momentum[Symbol(:ρ, name)]
-            value = ρ * u
-            set!(ϕ, value)
+            set_velocity!(model, name, value)
 
         elseif name ∈ settable_thermodynamic_variables
             set_thermodynamic_variable!(model, Val(name), value)

src/AtmosphereModels/update_atmosphere_model_state.jl

@@ -60,19 +60,16 @@ diagnostic_indices(::Bounded, N, H) = 1:N+1
 diagnostic_indices(::Periodic, N, H) = -H+1:N+H
 diagnostic_indices(::Flat, N, H) = 1:N
 
+#####
+##### Velocity and momentum computation
+#####
+
 """
 $(TYPEDSIGNATURES)
 
-Compute auxiliary model variables:
-
-- velocities from momentum and density (eg ``u = ρu / ρ``)
-
-- thermodynamic variables from the prognostic thermodynamic state,
-    * temperature ``T``, possibly involving saturation adjustment
-    * specific thermodynamic variable (``e = ρe / ρ`` or ``θ = ρθ / ρ``)
-    * moisture mass fraction ``qᵗ = ρqᵗ / ρ``
+Compute velocities from momentum: `u = ρu / ρ` for each velocity component.
 """
-function compute_auxiliary_variables!(model)
+function compute_velocities!(model::AtmosphereModel)
     grid = model.grid
     arch = grid.architecture
 
@@ -102,6 +99,66 @@ function compute_auxiliary_variables!(model)
     foreach(mask_immersed_field!, model.velocities)
     fill_halo_regions!(model.velocities)
 
+    return nothing
+end
+
+function compute_momentum_tendencies!(model::AtmosphereModel, model_fields)
+    grid = model.grid
+    arch = grid.architecture
+    Gρu = model.timestepper.Gⁿ.ρu
+    Gρv = model.timestepper.Gⁿ.ρv
+    Gρw = model.timestepper.Gⁿ.ρw
+
+    momentum_args = (
+        dynamics_density(model.dynamics),
+        model.advection.momentum,
+        model.velocities,
+        model.closure,
+        model.closure_fields,
+        model.momentum,
+        model.coriolis,
+        model.clock,
+        model_fields)
+
+    u_args = tuple(momentum_args..., model.forcing.ρu, model.dynamics)
+    v_args = tuple(momentum_args..., model.forcing.ρv, model.dynamics)
+
+    # Extra arguments for vertical velocity are required to compute buoyancy
+    w_args = tuple(momentum_args..., model.forcing.ρw,
+                   model.dynamics,
+                   model.formulation,
+                   model.temperature,
+                   model.specific_moisture,
+                   model.microphysics,
+                   model.microphysical_fields,
+                   model.thermodynamic_constants)
+
+    launch!(arch, grid, :xyz, compute_x_momentum_tendency!, Gρu, grid, u_args)
+    launch!(arch, grid, :xyz, compute_y_momentum_tendency!, Gρv, grid, v_args)
+    launch!(arch, grid, :xyz, compute_z_momentum_tendency!, Gρw, grid, w_args)
+
+    return nothing
+end
+
+"""
+$(TYPEDSIGNATURES)
+
+Compute auxiliary model variables:
+
+- velocities from momentum and density (eg ``u = ρu / ρ``)
+
+- thermodynamic variables from the prognostic thermodynamic state,
+    * temperature ``T``, possibly involving saturation adjustment
+    * specific thermodynamic variable (``e = ρe / ρ`` or ``θ = ρθ / ρ``)
+    * moisture mass fraction ``qᵗ = ρqᵗ / ρ``
+"""
+function compute_auxiliary_variables!(model)
+    grid = model.grid
+    arch = grid.architecture
+
+    # Compute velocities from momentum (skip for kinematic dynamics with prescribed velocities)
+    compute_velocities!(model)
+
     # Dispatch on thermodynamic formulation type
     compute_auxiliary_thermodynamic_variables!(model)
 
@@ -201,44 +258,14 @@ end
 function compute_tendencies!(model::AtmosphereModel)
     grid = model.grid
     arch = grid.architecture
-    Gρu = model.timestepper.Gⁿ.ρu
-    Gρv = model.timestepper.Gⁿ.ρv
-    Gρw = model.timestepper.Gⁿ.ρw
 
     model_fields = fields(model)
 
     #####
-    ##### Momentum tendencies
+    ##### Momentum tendencies (skip for kinematic dynamics)
     #####
 
-    momentum_args = (
-        dynamics_density(model.dynamics),
-        model.advection.momentum,
-        model.velocities,
-        model.closure,
-        model.closure_fields,
-        model.momentum,
-        model.coriolis,
-        model.clock,
-        model_fields)
-
-    u_args = tuple(momentum_args..., model.forcing.ρu, model.dynamics)
-    v_args = tuple(momentum_args..., model.forcing.ρv, model.dynamics)
-
-    # Extra arguments for vertical velocity are required to compute
-    # buoyancy:
-    w_args = tuple(momentum_args..., model.forcing.ρw,
-                   model.dynamics,
-                   model.formulation,
-                   model.temperature,
-                   model.specific_moisture,
-                   model.microphysics,
-                   model.microphysical_fields,
-                   model.thermodynamic_constants)
-
-    launch!(arch, grid, :xyz, compute_x_momentum_tendency!, Gρu, grid, u_args)
-    launch!(arch, grid, :xyz, compute_y_momentum_tendency!, Gρv, grid, v_args)
-    launch!(arch, grid, :xyz, compute_z_momentum_tendency!, Gρw, grid, w_args)
+    compute_momentum_tendencies!(model, model_fields)
 
     # Arguments common to energy density, moisture density, and tracer density tendencies:
     common_args = (

src/Breeze.jl

@@ -14,6 +14,8 @@ export
     AnelasticModel,
     CompressibleDynamics,
     CompressibleModel,
+    PrescribedDynamics,
+    KinematicModel,
     AtmosphereModel,
     StaticEnergyFormulation,
     LiquidIcePotentialTemperatureFormulation,
@@ -150,6 +152,9 @@ using .AnelasticEquations: AnelasticDynamics, AnelasticModel
 include("CompressibleEquations/CompressibleEquations.jl")
 using .CompressibleEquations: CompressibleDynamics, CompressibleModel
 
+include("KinematicDriver/KinematicDriver.jl")
+using .KinematicDriver: PrescribedDynamics, KinematicModel
+
 include("Microphysics/Microphysics.jl")
 using .Microphysics
 

src/KinematicDriver/KinematicDriver.jl

@@ -0,0 +1,42 @@
+"""
+    KinematicDriver
+
+Module implementing kinematic dynamics for atmosphere models.
+
+Kinematic dynamics prescribes the velocity field rather than solving for it,
+enabling isolated testing of microphysics, thermodynamics, and other physics
+without the complexity of solving the momentum equations.
+
+This is analogous to the `kin1d` driver in P3-microphysics.
+"""
+module KinematicDriver
+
+export
+    PrescribedDynamics,
+    KinematicModel
+
+using DocStringExtensions: TYPEDSIGNATURES
+using Adapt: Adapt, adapt
+
+using Oceananigans: Oceananigans, CenterField, XFaceField, YFaceField, ZFaceField, fields
+using Oceananigans.Architectures: architecture, on_architecture
+using Oceananigans.BoundaryConditions: fill_halo_regions!
+using Oceananigans.Fields: FunctionField, ZeroField, field
+using Oceananigans.Grids: Face, Center
+using Oceananigans.TimeSteppers: Clock, TimeSteppers
+using Oceananigans.Utils: prettysummary
+
+# Import PrescribedVelocityFields from Oceananigans
+using Oceananigans.Models.HydrostaticFreeSurfaceModels: PrescribedVelocityFields
+
+using Breeze.Thermodynamics: ReferenceState
+using Breeze.AtmosphereModels: AtmosphereModels, AtmosphereModel
+
+include("prescribed_dynamics.jl")
+
+# Type alias for kinematic models
+const KinematicModel = AtmosphereModel{<:PrescribedDynamics}
+
+include("kinematic_driver_time_stepping.jl")
+
+end # module

src/KinematicDriver/kinematic_driver_time_stepping.jl

@@ -0,0 +1,44 @@
+#####
+##### Time stepping for PrescribedDynamics (kinematic dynamics)
+#####
+
+using Oceananigans.Fields: set!, FunctionField
+
+#####
+##### Model initialization
+#####
+
+AtmosphereModels.initialize_model_thermodynamics!(m::KinematicModel) = set!(m, θ=m.dynamics.reference_state.potential_temperature)
+
+#####
+##### Velocity and momentum: no-ops (no momentum, velocities may be FunctionFields)
+#####
+
+AtmosphereModels.compute_velocities!(::KinematicModel) = nothing
+AtmosphereModels.compute_momentum_tendencies!(::KinematicModel, model_fields) = nothing
+
+#####
+##### Setting velocities for kinematic models
+#####
+
+# Dispatch on velocity specification stored in dynamics
+AtmosphereModels.set_velocity!(model::KinematicModel, name::Symbol, value) =
+    set_velocity!(model.dynamics.velocity_specification, model.velocities, name, value)
+
+# Regular velocity fields (velocity_specification is nothing): just set directly
+set_velocity!(::Nothing, velocities, name, value) = set!(velocities[name], value)
+
+# PrescribedVelocityFields: cannot be set
+set_velocity!(::PrescribedVelocityFields, velocities, name, value) = 
+    throw(ArgumentError("Cannot set velocity component '$name' of PrescribedVelocityFields."))
+
+# No momentum in kinematic models
+AtmosphereModels.set_momentum!(::KinematicModel, name::Symbol, value) = 
+    throw(ArgumentError("Cannot set momentum component '$name' of a KinematicModel."))
+
+#####
+##### Pressure correction: no-op for kinematic dynamics
+#####
+
+TimeSteppers.compute_pressure_correction!(::KinematicModel, Δt) = nothing
+TimeSteppers.make_pressure_correction!(::KinematicModel, Δt) = nothing