Acoustic refraction in wind shear — a differentiable example
This example does two things on top of the same compressible forward model. First, we simulate an acoustic pulse propagating through a wind shear layer using the fully compressible Euler equations, and observe how the shear refracts the wave: waves traveling with the wind bend downward (trapped near the surface), while waves traveling against the wind bend upward. The effective propagation speed for a wave in direction $\hat{\boldsymbol{n}}$ is
\[c^{ac} + \boldsymbol{u} \cdot \hat{\boldsymbol{n}}\]
where $cᵃᶜ$ is the acoustic sound speed and $\boldsymbol{u}$ is the wind velocity. Wavefronts tilt toward regions of lower effective propagation speed, "ducting" sound energy along the surface — which is why distant sounds are often heard more clearly downwind. For more on this topic, see
- Abdul-Razzak, H. and Ghan, S. J. (2000). A parameterization of aerosol activation: 2. Multiple aerosol types. Journal of Geophysical Research: Atmospheres 105, 6837–6844. ↩1 ↩2 ↩3
- Abdul-Razzak, H.; Ghan, S. J. and Rivera-Carpio, C. (1998). A parameterization of aerosol activation: 1. Single aerosol type. Journal of Geophysical Research: Atmospheres 103, 6123–6131. ↩1
- Baldauf, M. (2010). Linear stability analysis of Runge-Kutta-based partial time-splitting schemes for the Euler equations. Monthly Weather Review 138, 4475–4496. ↩1 ↩2 ↩3 ↩4
- Baldauf, M.; Seifert, A.; Förstner, J.; Majewski, D.; Raschendorfer, M. and Reinhardt, T. (2011). Operational Convective-Scale Numerical Weather Prediction with the COSMO Model: Description and Sensitivities. Monthly Weather Review 139, 3887–3905. ↩1
- Beare, R. J.; Macvean, M. K.; Holtslag, A. A.; Cuxart, J.; Esau, I.; Golaz, J.-C.; Jimenez, M. A.; Khairoutdinov, M.; Kosovic, B.; Lewellen, D.; Lund, T. S.; Lundquist, J. K.; Mccabe, A.; Moene, A. F.; Noh, Y.; Raasch, S. and Sullivan, P. (2006). An Intercomparison of Large-Eddy Simulations of the Stable Boundary Layer. Boundary-Layer Meteorology 118, 247–272. ↩1
- Berg, J.; Patton, E. G. and Sullivan, P. P. (2020). Large-eddy simulation of conditionally neutral boundary layers: A mesh resolution sensitivity study. Journal of the Atmospheric Sciences 77, 1969–1991. ↩1
- Bryan, G. H. and Fritsch, J. M. (2002). A benchmark simulation for moist nonhydrostatic numerical models. Monthly Weather Review 130, 2917–2928. ↩1 ↩2
- Cholette, M.; Morrison, H.; Milbrandt, J. A. and Thériault, J. M. (2019). Parameterization of the Bulk Liquid Fraction on Mixed-Phase Particles in the Predicted Particle Properties (P3) Scheme: Description and Idealized Simulations. Journal of the Atmospheric Sciences 76, 561–582. ↩1 ↩2
- Cober, S. G. and List, R. (1993). Measurements of the heat and mass transfer parameters characterizing conical graupel growth. Journal of the Atmospheric Sciences 50, 1591–1609. ↩1 ↩2
- Cronin, T. W. and Chavas, D. R. (2019). Dry and semidry tropical cyclones. Journal of the Atmospheric Sciences 76, 2193–2212. ↩1 ↩2 ↩3 ↩4
- Deardorff, J. W. (1980). Stratocumulus-capped mixed layers derived from a three-dimensional model. Boundary-Layer Meteorology 18, 495–527. ↩1 ↩2
- Durran, D. R. (2010). Numerical Methods for Fluid Dynamics: With Applications to Geophysics. 2nd Edition (Springer). ↩1
- Durran, D. R. and Klemp, J. B. (1982). On the effects of moisture on the Brunt-Väisälä frequency. Journal of the Atmospheric Sciences 39, 2152–2158. ↩1 ↩2 ↩3 ↩4
- Edson, J. B.; Jampana, V.; Weller, R. A.; Bigorre, S. P.; Plueddemann, A. J.; Fairall, C. W.; Miller, S. D.; Mahrt, L.; Vickers, D. and Hersbach, H. (2013). On the Exchange of Momentum over the Open Ocean. Journal of Physical Oceanography 43, 1589–1610. ↩1 ↩2
- Emanuel, K. A. (1986). An air-sea interaction theory for tropical cyclones. Part I: Steady-state maintenance. Journal of the Atmospheric Sciences 43, 585–605. ↩1
- Field, P. R.; Heymsfield, A. J. and Bansemer, A. (2007). Snow size distribution parameterization for midlatitude and tropical ice clouds. Journal of the Atmospheric Sciences 64, 4346–4365. ↩1
- Flatau, P. J.; Walko, R. L. and Cotton, W. R. (1992). Polynomial fits to saturation vapor pressure. Journal of Applied Meteorology 31, 1507–1513. ↩1 ↩2
- Gal-Chen, T. and Somerville, R. C. (1975). On the use of a coordinate transformation for the solution of the Navier-Stokes equations. Journal of Computational Physics 17, 209–228. ↩1
- Grabowski, W. W. and Morrison, H. (2008). Toward the mitigation of spurious cloud-edge supersaturation in cloud models. Monthly Weather Review 136, 1224–1234. ↩1
- Hall, W. D. and Pruppacher, H. R. (1976). The survival of ice particles falling from cirrus clouds in subsaturated air. Journal of the Atmospheric Sciences 33, 1995–2006. ↩1 ↩2 ↩3 ↩4
- Hallett, J. and Mossop, S. C. (1974). Production of secondary ice particles during the riming process. Nature 249, 26–28. ↩1
- Han, J. and Bretherton, C. S. (2019). TKE-Based Moist Eddy-Diffusivity Mass-Flux (EDMF) Parameterization for Vertical Turbulent Mixing. Weather and Forecasting 34, 869–886. ↩1
- Heymsfield, A. J. (2003). Properties of tropical and midlatitude ice cloud particle ensembles. Part I: Median mass diameters and terminal velocities. Journal of the Atmospheric Sciences 60, 2573–2591. ↩1 ↩2 ↩3
- Heymsfield, A. J.; Bansemer, A. and Twohy, C. H. (2007). Refinements to Ice Particle Mass Dimensional and Terminal Velocity Relationships for Ice Clouds. Part II: Evaluation and Parameterizations of Ensemble Ice Particle Sedimentation Velocities. Journal of the Atmospheric Sciences 64, 1068–1088. ↩1 ↩2 ↩3 ↩4
- Howard, L. N. (1961). Note on a paper of John W. Miles. Journal of Fluid Mechanics 10, 509–512. ↩1
- Jablonowski, C. and Williamson, D. L. (2006). A baroclinic instability test case for atmospheric model dynamical cores. Quarterly Journal of the Royal Meteorological Society 132, 2943–2975. ↩1
- Jordan, C. L. (1958). Mean soundings for the West Indies area. Journal of Meteorology 15, 91–97. ↩1 ↩2
- Kaul, C. M.; Teixeira, J. and Suzuki, K. (2015). Sensitivities in Large-Eddy Simulations of Mixed-Phase Arctic Stratocumulus Clouds Using a Simple Microphysics Approach. Monthly Weather Review 143, 4393–4421. ↩1
- Kessler, E. (1969). On the distribution and continuity of water substance in atmospheric circulations. Vol. 10 no. 32 of Meteorological Monographs (American Meteorological Society). ↩1 ↩2 ↩3
- Khairoutdinov, M. F.; Blossey, P. N. and Bretherton, C. S. (2022). Global system for atmospheric modeling: Model description and preliminary results. Journal of Advances in Modeling Earth Systems 14, e2021MS002968. ↩1
- Klemp, J. B. (2011). A terrain-following coordinate with smoothed coordinate surfaces. Monthly Weather Review 139, 2163–2169. ↩1 ↩2
- Klemp, J. B.; Skamarock, W. C. and Dudhia, J. (2007). Conservative split-explicit time integration methods for the compressible nonhydrostatic equations. Monthly Weather Review 135, 2897–2913. ↩1
- Klemp, J. B. and Wilhelmson, R. B. (1978). The simulation of three-dimensional convective storm dynamics. Journal of Atmospheric Sciences 35, 1070–1096. ↩1 ↩2 ↩3
- Knoth, O. and Wensch, J. (2014). Generalized split-explicit Runge-Kutta methods for the compressible Euler equations. Monthly Weather Review 142, 2067–2081. ↩1 ↩2
- Liu, L.; Gadde, S. N. and Stevens, R. J. (2021). Geostrophic drag law for conventionally neutral atmospheric boundary layers revisited. Quarterly Journal of the Royal Meteorological Society 147, 847–857. ↩1
- Milbrandt, J. A. and Morrison, H. (2016). Parameterization of cloud microphysics based on the prediction of bulk ice particle properties. Part III: Introduction of multiple free categories. Journal of the Atmospheric Sciences 73, 975–995. ↩1 ↩2 ↩3
- Milbrandt, J. A.; Morrison, H.; Dawson, D. T. and Paukert, M. (2021). A triple-moment representation of ice in the Predicted Particle Properties (P3) microphysics scheme. Journal of the Atmospheric Sciences 78, 439–458. ↩1 ↩2 ↩3
- Miles, J. W. (1961). On the stability of heterogeneous shear flows. Journal of Fluid Mechanics 10, 496-–508. ↩1
- Mitchell, D. L. (1996). Use of Mass- and Area-Dimensional Power Laws for Determining Precipitation Particle Terminal Velocities. Journal of Atmospheric Sciences 53, 1710–1723. ↩1 ↩2
- Mitchell, D. L. and Heymsfield, A. J. (2005). Refinements in the treatment of ice particle terminal velocities, highlighting aggregates. Journal of the Atmospheric Sciences 62, 1637–1644. ↩1 ↩2 ↩3 ↩4
- Moeng, C.-H. and Sullivan, P. P. (1994). A comparison of shear- and buoyancy-driven planetary boundary layer flows. Journal of the Atmospheric Sciences 51, 999–1022. ↩1 ↩2 ↩3
- Monteith, J. L. and Unsworth, M. H. (2014). Principles of Environmental Physics. 4th Edition (Academic Press). ↩1 ↩2
- Moon, Y. and Nolan, D. S. (2010). The dynamic response of the hurricane wind field to spiral rainband heating. Journal of the Atmospheric Sciences 67, 1779–1805. ↩1 ↩2 ↩3
- Morrison, H. and Grabowski, W. W. (2008). A novel approach for representing ice microphysics in models: Description and tests using a kinematic framework. Journal of the Atmospheric Sciences 65, 1528–1548. ↩1 ↩2 ↩3
- Morrison, H. and Milbrandt, J. A. (2015). Parameterization of cloud microphysics based on the prediction of bulk ice particle properties. Part I: Scheme description and idealized tests. Journal of the Atmospheric Sciences 72, 287–311. ↩1 ↩2 ↩3 ↩4 ↩5 ↩6 ↩7 ↩8 ↩9 ↩10 ↩11 ↩12 ↩13 ↩14 ↩15 ↩16 ↩17 ↩18 ↩19 ↩20 ↩21 ↩22 ↩23 ↩24 ↩25 ↩26 ↩27 ↩28 ↩29 ↩30 ↩31 ↩32 ↩33 ↩34 ↩35 ↩36 ↩37 ↩38 ↩39 ↩40 ↩41 ↩42 ↩43 ↩44 ↩45 ↩46 ↩47 ↩48 ↩49 ↩50 ↩51 ↩52 ↩53 ↩54 ↩55 ↩56 ↩57 ↩58 ↩59 ↩60 ↩61 ↩62 ↩63 ↩64 ↩65
- Morrison, H.; Milbrandt, J. A.; Bryan, G. H.; Ikeda, K.; Tessendorf, S. A. and Thompson, G. (2015). Parameterization of cloud microphysics based on the prediction of bulk ice particle properties. Part II: Case study comparisons with observations and other schemes. Journal of the Atmospheric Sciences 72, 312–339. ↩1 ↩2 ↩3 ↩4
- Murray, F. W. (1967). On the computation of saturation vapor pressure. Journal of Applied Meteorology 6, 203–204. ↩1 ↩2 ↩3
- Nakanishi, M. and Niino, H. (2009). Development of an Improved Turbulence Closure Model for the Atmospheric Boundary Layer. Journal of the Meteorological Society of Japan 87, 895–912. ↩1 ↩2
- Nishizawa, S. and Kitamura, Y. (2018). A Surface Flux Scheme Based on the Monin–Obukhov Similarity for Finite Volume Models. Journal of Advances in Modeling Earth Systems 10, 3159–3175. ↩1 ↩2 ↩3
- Nolan, D. S. (2001). The stabilizing effects of axial stretching on turbulent vortex dynamics. Physics of Fluids 13, 1724–1738. ↩1
- Ostashev, V. E. and Wilson, D. K. (2015). Acoustics in moving inhomogeneous media (CRC Press).
- Pedersen, J. G.; Gryning, S.-E. and Kelly, M. (2014). On the structure and adjustment of inversion-capped neutral atmospheric boundary-layer flows: Large-eddy simulation study. Boundary-Layer Meteorology 153, 43–62. ↩1
- Pierce, A. D. (2019). Acoustics: An introduction to its physical principles and applications (Springer Cham).
- Schneider, T. (2004). The tropopause and the thermal stratification in the extratropics of a dry atmosphere. Journal of the Atmospheric Sciences 61, 1317–1340. ↩1
- Shin, E. Y.; Yang, X. I. and Howland, M. F. (2025). Addressing Grid Convergence and Log-Layer Mismatch in Wall Modeled Large Eddy Simulations of Geophysical Flows Over Rough Surfaces and Canopies. Boundary-Layer Meteorology 191. ↩1 ↩2 ↩3 ↩4
- Shu, C.-W. (2009). High order weighted essentially nonoscillatory schemes for convection dominated problems. SIAM Review 51, 82–126. ↩1
- Shu, C.-W. and Osher, S. (1988). Efficient implementation of essentially non-oscillatory shock-capturing schemes. Journal of Computational Physics 77, 439–471. ↩1 ↩2 ↩3
- Siebesma, A. P.; Bretherton, C. S.; Brown, A.; Chlond, A.; Cuxart, J.; Duynkerke, P. G.; Jiang, H.; Khairoutdinov, M.; Lewellen, D.; Moeng, C.-H.; Sanchez, E.; Stevens, B. and Stevens, D. E. (2003). A large eddy simulation intercomparison study of shallow cumulus convection. Journal of the Atmospheric Sciences 60, 1201–1219. ↩1 ↩2 ↩3 ↩4 ↩5 ↩6 ↩7 ↩8 ↩9 ↩10 ↩11 ↩12 ↩13 ↩14 ↩15
- Skamarock, W. C. and Klemp, J. B. (1992). The Stability of Time-Split Numerical Methods for the Hydrostatic and the Nonhydrostatic Elastic Equations. Monthly Weather Review 120, 2109–2127. ↩1 ↩2 ↩3
- Skamarock, W. C. and Klemp, J. B. (2008). A time-split nonhydrostatic atmospheric model for weather research and forecasting applications. Journal of Computational Physics 227, 3465–3485. ↩1
- Skamarock, W. C.; Klemp, J. B.; Duda, M. G.; Fowler, L. D.; Park, S.-H. and Ringler, T. D. (2012). A Multiscale Nonhydrostatic Atmospheric Model Using Centroidal Voronoi Tesselations and C-Grid Staggering. Monthly Weather Review 140, 3090–3105. ↩1
- Stern, D. P. and Nolan, D. S. (2009). Reexamining the vertical structure of tangential winds in tropical cyclones: Observations and theory. Journal of the Atmospheric Sciences 66, 3579–3600. ↩1 ↩2
- Straka, J. M. (2009). Cloud and precipitation microphysics: Principles and parameterizations (Cambridge University Press). ↩1
- Sullivan, P. P.; McWilliams, J. C. and Moeng, C.-H. (1994). A subgrid-scale model for large-eddy simulation of planetary boundary-layer flows. Boundary-Layer Meteorology 71, 247–276. ↩1
- Ullrich, P. A.; Jablonowski, C.; Kent, J.; Lauritzen, P. H.; Nair, R.; Reed, K. A.; Zarzycki, C. M.; Hall, D. M.; Dazlich, D.; Heikes, R.; Konor, C.; Randall, D.; Dubos, T.; Meurdesoif, Y.; Chen, X.; Harris, L.; Kühnlein, C.; Lee, V.; Qaddouri, A.; Girard, C.; Giorgetta, M.; Reinert, D.; Klemp, J.; Park, S.-H.; Skamarock, W.; Miura, H.; Ohno, T.; Yoshida, R.; Walko, R.; Reinecke, A. and Viner, K. (2017). DCMIP2016: a review of non-hydrostatic dynamical core design and intercomparison of participating models. Geoscientific Model Development 10, 4477–4509. ↩1 ↩2
- Wagner, G. L.; Hillier, A.; Constantinou, N. C.; Silvestri, S.; Souza, A.; Burns, K.; Hill, C.; Campin, J.-M.; Marshall, J. and Ferrari, R. (2025). Formulation and calibration of CATKE, a one-equation parameterization for microscale ocean mixing. Journal of Advances in Modeling Earth Systems 17, e2024MS004522. ↩1 ↩2
- van Zanten, M. C.; Stevens, B.; Nuijens, L.; Siebesma, A. P.; Ackerman, A. S.; Burnet, F.; Cheng, A.; Couvreux, F.; Jiang, H.; Khairoutdinov, M.; Kogan, Y.; Lewellen, D. C.; Mechem, D.; Nakamura, K.; Noda, A.; Shipway, B. J.; Slawinska, J.; Wang, S. and Wyszogrodzki, A. (2011). Controls on precipitation and cloudiness in simulations of trade-wind cumulus as observed during RICO. Journal of Advances in Modeling Earth Systems 3, M06001. ↩1 ↩2 ↩3 ↩4 ↩5 ↩6 ↩7 ↩8 ↩9 ↩10 ↩11
- Zarzycki, C. M.; Jablonowski, C.; Kent, J.; Lauritzen, P. H.; Nair, R.; Reed, K. A.; Ullrich, P. A.; Hall, D. M.; Dazlich, D.; Heber, R.; Achatz, U.; Butter, T.; Galewsky, J.; Goodman, J.; Klein, R.; Lemarié, F.; Malardel, S.; Rauscher, S. A.; Schar, C.; Sprenger, M.; Taylor, M. A.; Vogl, C.; Wan, H. and Williamson, D. L. (2019). DCMIP2016: the splitting supercell test case. Geoscientific Model Development 12, 879–892. ↩1 ↩2 ↩3
Second, we use this setup as a minimal introduction to differentiable atmospheric simulation in Breeze. After running the forward problem, we take a gradient through the entire compressible time-stepping — asking which parts of the wind profile control how much acoustic energy ends up trapped near the surface? — using Enzyme.jl for reverse-mode automatic differentiation and Reactant.jl to compile the model down to XLA. The result is a 2D sensitivity field obtained in a single backward pass; the same answer via finite differences would cost one model rerun per grid cell.
We use stable stratification to suppress Kelvin-Helmholtz instability and a logarithmic wind profile consistent with the atmospheric surface layer.
using Breezeusing Oceananigans: Oceananigansusing Oceananigans.Unitsusing Printfusing CairoMakieGrid and model setup
Nx, Nz = 128, 64Lx, Lz = 1000, 200 # (m)grid = RectilinearGrid(size = (Nx, Nz), x = (-Lx/2, Lx/2), z = (0, Lz), topology = (Periodic, Flat, Bounded))128×1×64 RectilinearGrid{Float64, Periodic, Flat, Bounded} on CPU with 3×0×3 halo
├── Periodic x ∈ [-500.0, 500.0) regularly spaced with Δx=7.8125
├── Flat y
└── Bounded z ∈ [0.0, 200.0] regularly spaced with Δz=3.125This example is dry, so the θˡⁱ→T inversion has an exact closed form and needs no Newton steps. FixedIterations(2) selects the fixed-trip (unrolled) inversion with a tiny trip count: the differentiable pass at the end of this example takes a reverse-mode gradient through the compressible time step with Reactant/Enzyme, and the default tolerance-based NewtonSolver (a while loop) compiles to an XLA while op that does not differentiate cheaply, while a high iteration count needlessly inflates the traced graph (NumericalEarth/Breeze.jl#767). The forward and adjoint models use identical dynamics and formulation. We also use the full-pressure form (reference_state = nothing) so reinitializing the differentiated model does not rebuild a model-owned hydrostatic reference inside the traced loss.
formulation = LiquidIcePotentialTemperatureFormulation(temperature_solver = FixedIterations(2))model = AtmosphereModel(grid; formulation, dynamics = CompressibleDynamics(ExplicitTimeStepping(); reference_state = nothing))AtmosphereModel{CPU, RectilinearGrid}(time = 0 seconds, iteration = 0)
├── grid: 128×1×64 RectilinearGrid{Float64, Periodic, Flat, Bounded} on CPU with 3×0×3 halo
├── dynamics: CompressibleDynamics{ExplicitTimeStepping}
├── formulation: LiquidIcePotentialTemperatureFormulation
├── thermodynamic_constants: ThermodynamicConstants{Float64}
├── timestepper: SSPRungeKutta3
├── advection scheme:
│ ├── momentum: Centered(order=2)
│ ├── ρθ: Centered(order=2)
│ └── ρqᵛ: Centered(order=2)
├── forcing: @NamedTuple{ρᵈ::Returns{Float64}, ρu::Returns{Float64}, ρv::Returns{Float64}, ρw::Returns{Float64}, ρθ::Returns{Float64}, ρqᵛ::Returns{Float64}, ρE::Returns{Float64}}
├── tracers: ()
├── coriolis: Nothing
└── microphysics: NothingBackground state
We build a hydrostatically balanced reference state using ReferenceState. This provides the background density and pressure profiles.
constants = model.thermodynamic_constantsθ₀ = 300 # Reference potential temperature (K)p₀ = 101325 # Surface pressure (Pa)pˢᵗ = 1e5 # Standard pressure (Pa)reference = ReferenceState(grid, constants; base_pressure=p₀, potential_temperature=θ₀, standard_pressure=pˢᵗ)ReferenceState{Float64}(p₀=101325.0, θ₀=300.0, pˢᵗ=100000.0)The sound speed at the surface determines the acoustic wave propagation speed.
Rᵈ = constants.molar_gas_constant / constants.dry_air.molar_masscᵖᵈ = constants.dry_air.heat_capacityγ = cᵖᵈ / (cᵖᵈ - Rᵈ)cᵃᶜ = sqrt(γ * Rᵈ * θ₀)347.15629052210863The wind profile follows the classic log-law of the atmospheric surface layer.
U₀ = 20 # Surface velocity (m/s), u★/κℓ = 1 # Roughness length (m), like, shrubs and stuffUᵢ(z) = U₀ * log((z + ℓ) / ℓ)Uᵢ (generic function with 1 method)Initial conditions
We initialize a localized Gaussian density pulse representing an acoustic disturbance. For a rightward-propagating acoustic wave, the velocity perturbation is in phase with the density perturbation: $u' = (cᵃᶜ / ρ₀) ρ'$.
δρ = 0.01 # Density perturbation amplitude (kg/m³)σ = 20 # Pulse width (m)gaussian(x, z) = exp(-(x^2 + z^2) / 2σ^2)ρ₀ = interior(reference.density, 1, 1, 1)[]ρᵢ_func(x, z) = adiabatic_hydrostatic_density(z, p₀, θ₀, pˢᵗ, constants) + δρ * gaussian(x, z)uᵢ_func(x, z) = Uᵢ(z) # + (cᵃᶜ / ρ₀) * δρ * gaussian(x, z)set!(model, ρ=ρᵢ_func, θ=θ₀, u=uᵢ_func)Simulation setup
Acoustic waves travel fast ($cᵃᶜ ≈ 347$ m/s), so we need a small time step. The Courant–Friedrichs–Lewy (CFL) condition is based on the effective propagation speed $cᵃᶜ + \mathrm{max}(U)$.
Δx, Δz = Lx / Nx, Lz / NzΔt = 0.5 * min(Δx, Δz) / (cᵃᶜ + Uᵢ(Lz))stop_time = 0.5 # (s) — long enough for the wave to traverse the domain and for refraction to bend rays visiblysimulation = Simulation(model; Δt, stop_time)Oceananigans.Diagnostics.erroring_NaNChecker!(simulation)function progress(sim) u, v, w = sim.model.velocities msg = @sprintf("Iter: %d, t: %s, max|u|: %.2f m/s, max|w|: %.2f m/s", iteration(sim), prettytime(sim), maximum(abs, u), maximum(abs, w)) @info msgendadd_callback!(simulation, progress, IterationInterval(100))Output
We perturbation fields for density and x-velocity for visualization.
ρ = model.dynamics.dry_densityu, v, w = model.velocitiesρᵇᵍ = CenterField(grid)uᵇᵍ = XFaceField(grid)set!(ρᵇᵍ, (x, z) -> adiabatic_hydrostatic_density(z, p₀, θ₀, pˢᵗ, constants))set!(uᵇᵍ, (x, z) -> Uᵢ(z))ρ′ = Field(ρ - ρᵇᵍ)u′ = Field(u - uᵇᵍ)U = Average(u, dims = 1)R = Average(ρ, dims = 1)W² = Average(w^2, dims = 1)filename = "acoustic_wave.jld2"outputs = (; ρ′, u′, w, U, R, W²)simulation.output_writers[:jld2] = JLD2Writer(model, outputs; filename, schedule = TimeInterval(0.01), overwrite_files = true)run!(simulation)[ Info: Initializing simulation...
[ Info: Iter: 0, t: 0 seconds, max|u|: 105.91 m/s, max|w|: 0.00 m/s
[ Info: ... simulation initialization complete (27.460 seconds)
[ Info: Executing initial time step...
[ Info: ... initial time step complete (1.610 seconds).
[ Info: Iter: 100, t: 333.448 ms, max|u|: 105.91 m/s, max|w|: 0.48 m/s
[ Info: Simulation is stopping after running for 32.741 seconds.
[ Info: Simulation time 500 ms equals or exceeds stop time 500 ms.
Visualization
Load the saved perturbation fields and create a snapshot.
ρ′ts = FieldTimeSeries(filename, "ρ′")u′ts = FieldTimeSeries(filename, "u′")wts = FieldTimeSeries(filename, "w")Uts = FieldTimeSeries(filename, "U")Rts = FieldTimeSeries(filename, "R")W²ts = FieldTimeSeries(filename, "W²")times = ρ′ts.timesNt = length(times)fig = Figure(size = (900, 600), fontsize = 12)axρ = Axis(fig[1, 2]; aspect = 5, ylabel = "z (m)")axw = Axis(fig[2, 2]; aspect = 5, ylabel = "z (m)")axu = Axis(fig[3, 2]; aspect = 5, xlabel = "x (m)", ylabel = "z (m)")axR = Axis(fig[1, 1]; xlabel = "⟨ρ⟩ (kg/m³)")axW = Axis(fig[2, 1]; xlabel = "⟨w²⟩ (m²/s²)", limits = (extrema(W²ts), nothing))axU = Axis(fig[3, 1]; xlabel = "⟨u⟩ (m/s)")hidexdecorations!(axρ)hidexdecorations!(axw)colsize!(fig.layout, 1, Relative(0.2))n = Observable(Nt)ρ′n = @lift ρ′ts[$n]u′n = @lift u′ts[$n]wn = @lift wts[$n]Un = @lift Uts[$n]Rn = @lift Rts[$n]W²n = @lift W²ts[$n]ρlim = δρ / 4ulim = 1hmρ = heatmap!(axρ, ρ′n; colormap = :balance, colorrange = (-ρlim, ρlim))hmw = heatmap!(axw, wn; colormap = :balance, colorrange = (-ulim, ulim))hmu = heatmap!(axu, u′n; colormap = :balance, colorrange = (-ulim, ulim))lines!(axR, Rn)lines!(axW, W²n)lines!(axU, Un)Colorbar(fig[1, 3], hmρ; label = "ρ′ (kg/m³)")Colorbar(fig[2, 3], hmw; label = "w (m/s)")Colorbar(fig[3, 3], hmu; label = "u′ (m/s)")title = @lift "Acoustic wave in log-layer shear — t = $(prettytime(times[$n]))"fig[0, :] = Label(fig, title, fontsize = 16, tellwidth = false)CairoMakie.record(fig, "acoustic_wave.mp4", 1:Nt; framerate = 18, compression = 23) do nn n[] = nnendA differentiable workflow
The forward simulation above gives us the physics; the rest of the example treats that same forward model as a function and takes its gradient. This is the minimal pattern you'd reach for whenever you want to do data-assimilation, parameter calibration, or sensitivity analysis with a Breeze atmosphere — wrap the time-stepping in a scalar-valued loss, compile it with Reactant.@compile, and differentiate it with Enzyme.autodiff.
The pattern is:
- Rebuild the model on a
ReactantStategrid so all arrays are XLA buffers. - Choose a differentiated input (here, the initial wind field).
- Define a scalar
lossthat re-initializes the model from that input, runsnstepsoftime_step!inside a@traceloop, and reduces to a scalar diagnostic. - Wrap
lossin agrad_lossthat callsEnzyme.autodiff(...)with the input asDuplicatedand everything else asConst. - Compile once with
Reactant.@compile raise=true raise_first=true; run many times.
Why Reactant?
Reactant traces Julia code into an intermediate representation (StableHLO) that XLA can optimize and Enzyme can differentiate. The key requirement is that the model lives on ReactantState — Reactant's architecture in Oceananigans — so that all arrays are XLA buffers. We therefore rebuild the same physical setup on a new grid whose architecture is ReactantState().
using Reactant, CUDA # CUDA is required for loading the Reactant extensionusing Enzymeusing Statistics: meanusing Oceananigans.Architectures: ReactantStateusing Reactant: @traceReactant.set_default_backend("cpu")Rebuild the grid and model on ReactantState.
grid_ad = RectilinearGrid(ReactantState(); size = (Nx, Nz), x = (-Lx/2, Lx/2), z = (0, Lz), topology = (Periodic, Flat, Bounded))formulation_ad = LiquidIcePotentialTemperatureFormulation(temperature_solver = FixedIterations(2)) # fixed-trip, low-iteration EOS inversion so Enzyme can differentiate it cheaply (see forward model)model_ad = AtmosphereModel(grid_ad; formulation = formulation_ad, dynamics = CompressibleDynamics(ExplicitTimeStepping(); reference_state = nothing))AtmosphereModel{ReactantState, RectilinearGrid}(time = 0 seconds, iteration = Reactant.ConcretePJRTNumber{Int64, 1}(0))
├── grid: 128×1×64 RectilinearGrid{Float64, Periodic, Flat, Bounded} on ReactantState with 3×0×3 halo
├── dynamics: CompressibleDynamics{ExplicitTimeStepping}
├── formulation: LiquidIcePotentialTemperatureFormulation
├── thermodynamic_constants: ThermodynamicConstants{Float64}
├── timestepper: SSPRungeKutta3
├── advection scheme:
│ ├── momentum: Centered(order=2)
│ ├── ρθ: Centered(order=2)
│ └── ρqᵛ: Centered(order=2)
├── forcing: @NamedTuple{ρᵈ::Returns{Float64}, ρu::Returns{Float64}, ρv::Returns{Float64}, ρw::Returns{Float64}, ρθ::Returns{Float64}, ρqᵛ::Returns{Float64}, ρE::Returns{Float64}}
├── tracers: ()
├── coriolis: Nothing
└── microphysics: NothingFixed and varying fields
In this experiment the initial density pulse and hydrostatic background are held fixed; only the wind profile varies. We therefore precompute the total initial density once (background + Gaussian pulse) and use it as a Const. We also keep the standalone background $\bar\rho(z)$ around so the loss can subtract it from the model density to isolate the acoustic perturbation.
ρᵇᵍ = CenterField(grid_ad)set!(ρᵇᵍ, (x, z) -> adiabatic_hydrostatic_density(z, p₀, θ₀, pˢᵗ, constants))ρ_total = CenterField(grid_ad)set!(ρ_total, ρᵢ_func)128×1×64 Field{Oceananigans.Grids.Center, Oceananigans.Grids.Center, Oceananigans.Grids.Center} on Oceananigans.Grids.RectilinearGrid on ReactantState
├── grid: 128×1×64 RectilinearGrid{Float64, Periodic, Flat, Bounded} on ReactantState with 3×0×3 halo
├── boundary conditions: FieldBoundaryConditions
│ └── west: Periodic, east: Periodic, south: Nothing, north: Nothing, bottom: ZeroFlux, top: ZeroFlux, immersed: Nothing
└── data: 134×1×70 OffsetArray(::Reactant.ConcretePJRTArray{Float64,3}, -2:131, 1:1, -2:67) with eltype Float64 with indices -2:131×1:1×-2:67
└── max=1.18204, min=1.15363, mean=1.16299The initial wind field is the quantity we differentiate with respect to. Enzyme accumulates $∂J / ∂u_i$ into the shadow buffer $du_i$.
uᵢ = XFaceField(grid_ad)duᵢ = XFaceField(grid_ad)set!(uᵢ, (x, z) -> Uᵢ(z))set!(duᵢ, 0)128×1×64 Field{Oceananigans.Grids.Face, Oceananigans.Grids.Center, Oceananigans.Grids.Center} on Oceananigans.Grids.RectilinearGrid on ReactantState
├── grid: 128×1×64 RectilinearGrid{Float64, Periodic, Flat, Bounded} on ReactantState with 3×0×3 halo
├── boundary conditions: FieldBoundaryConditions
│ └── west: Periodic, east: Periodic, south: Nothing, north: Nothing, bottom: ZeroFlux, top: ZeroFlux, immersed: Nothing
└── data: 134×1×70 OffsetArray(::Reactant.ConcretePJRTArray{Float64,3}, -2:131, 1:1, -2:67) with eltype Float64 with indices -2:131×1:1×-2:67
└── max=0.0, min=0.0, mean=0.0The shadow model stores accumulated adjoints for every prognostic field.
dmodel_ad = Enzyme.make_zero(model_ad)AtmosphereModel{ReactantState, RectilinearGrid}(time = 0 seconds, iteration = Reactant.ConcretePJRTNumber{Int64, 1}(0))
├── grid: 128×1×64 RectilinearGrid{Float64, Periodic, Flat, Bounded} on ReactantState with 3×0×3 halo
├── dynamics: CompressibleDynamics{ExplicitTimeStepping}
├── formulation: LiquidIcePotentialTemperatureFormulation
├── thermodynamic_constants: ThermodynamicConstants{Float64}
├── timestepper: SSPRungeKutta3
├── advection scheme:
│ ├── momentum: Centered(order=2)
│ ├── ρθ: Centered(order=2)
│ └── ρqᵛ: Centered(order=2)
├── forcing: @NamedTuple{ρᵈ::Returns{Float64}, ρu::Returns{Float64}, ρv::Returns{Float64}, ρw::Returns{Float64}, ρθ::Returns{Float64}, ρqᵛ::Returns{Float64}, ρE::Returns{Float64}}
├── tracers: ()
├── coriolis: Nothing
└── microphysics: NothingTime step and integration length
We reuse the CFL-based time step and the exact number of iterations from the forward simulation above. The Simulation API is not used here because Reactant compiles a fixed-length traced loop instead. Gradient checkpointing requires a perfect-square step count, so we round up to the next perfect square.
Nt = simulation.model.clock.iterationNsteps = (isqrt(Nt - 1) + 1)^2169Defining the objective
We measure the mean squared acoustic density anomaly along the bottom of the domain:
\[J \;=\; \frac{1}{N_x}\sum_{i}\bigl[\rho(x_i, z_1) - \bar\rho(x_i, z_1)\bigr]^2\]
This is a global measure of how much acoustic energy ends up trapped near the surface. Averaging along the whole bottom row gives a sensitivity field that lights up wherever the wind affects any part of the surface response, making the ducting pattern visible across the entire domain.
The set! inside loss is what re-initializes the model from the current wind field on every backward evaluation. Without it, AD would differentiate a stale trajectory.
function loss(model, uᵢ, ρ_total, ρᵇᵍ, θ₀, Δt, nsteps) set!(model; ρ = ρ_total, θ = θ₀, u = uᵢ) @trace mincut=true checkpointing=true track_numbers=false for _ in 1:nsteps time_step!(model, Δt) end ρ₀ = interior(model.dynamics.dry_density, :, :, 1) ρᵇ₀ = interior(ρᵇᵍ, :, :, 1) return mean((ρ₀ .- ρᵇ₀).^2)endloss (generic function with 1 method)The gradient wrapper
grad_loss zeroes the adjoint buffer and calls Enzyme.autodiff in reverse mode. The model and the initial wind are Duplicated (primal + shadow); everything else is Const.
function grad_loss(model, dmodel, uᵢ, duᵢ, ρ_total, ρᵇᵍ, θ₀, Δt, nsteps) parent(duᵢ) .= 0 _, J = Enzyme.autodiff( Enzyme.set_strong_zero(Enzyme.ReverseWithPrimal), loss, Enzyme.Active, Enzyme.Duplicated(model, dmodel), Enzyme.Duplicated(uᵢ, duᵢ), Enzyme.Const(ρ_total), Enzyme.Const(ρᵇᵍ), Enzyme.Const(θ₀), Enzyme.Const(Δt), Enzyme.Const(nsteps)) return duᵢ, Jendgrad_loss (generic function with 1 method)Compilation and execution
Reactant.@compile traces the function once to build an XLA executable. The flags raise=true and raise_first=true ensure that every KernelAbstractions kernel is "raised" to StableHLO before Enzyme differentiates through it — a requirement for the backward pass.
@info "Compiling differentiated model — this may take a minute..."compiled_grad = Reactant.@compile raise=true raise_first=true sync=true grad_loss( model_ad, dmodel_ad, uᵢ, duᵢ, ρ_total, ρᵇᵍ, θ₀, Δt, Nsteps)@info "Running gradient..."du, J = compiled_grad(model_ad, dmodel_ad, uᵢ, duᵢ, ρ_total, ρᵇᵍ, θ₀, Δt, Nsteps)xs_u = xnodes(grid_ad, Face())zs = znodes(grid_ad, Center())@info @sprintf("Surface-mean (ρ - ρ̄)² = %.6e after %d steps", Float64(only(J)), Nsteps)[ Info: Compiling differentiated model — this may take a minute...
[ Info: Running gradient...
[ Info: Surface-mean (ρ - ρ̄)² = 1.094326e-07 after 169 steps
Sensitivity visualization
The heatmap shows $\partial J / \partial u_i(x, z)$: positive values are wind perturbations that would increase surface acoustic energy, negative values would decrease it. Because $J$ integrates along the entire bottom, the pattern reveals which parts of the wind profile feed energy into the surface duct from anywhere along it.
sensitivity = Array(interior(du, :, 1, :))abs_max = maximum(abs, sensitivity)fig_sens = Figure(size = (800, 350), fontsize = 12)Label(fig_sens[0, :], "∂J / ∂uᵢ (J = ⟨(ρ - ρ̄)²⟩ at surface, t=$(prettytime(Nsteps * Δt))", fontsize = 14, tellwidth = false)ax_sens = Axis(fig_sens[1, 1]; xlabel = "x (m)", ylabel = "z (m)")hm = heatmap!(ax_sens, xs_u, zs, sensitivity; colormap = :balance, colorrange = (-abs_max, abs_max))Colorbar(fig_sens[1, 2], hm; label = "∂J / ∂uᵢ")fig_sensJulia version and environment information
This example was executed with the following version of Julia:
using InteractiveUtils: versioninfoversioninfo()Julia Version 1.13.0
Commit d1c37793dd2 (2026-09-09 19:00 UTC)
Build Info:
Official https://julialang.org release
Platform Info:
OS: Linux (x86_64-linux-gnu)
CPU: 8 × AMD EPYC 7R13 Processor
WORD_SIZE: 64
LLVM: libLLVM-20.1.8 (ORCJIT, znver3)
GC: Built with stock GC
Threads: 1 default, 1 interactive, 5 GC (on 8 virtual cores)
Environment:
JULIA_NUM_GC_THREADS = 4,1
JULIA_GPG = 64B779A570972FFF7BFC2B54EAD471E1A1F2C10A
JULIA_LOAD_PATH = :@breeze
JULIA_PKG_SERVER_REGISTRY_PREFERENCE = eager
JULIA_VERSION = 1.13.0
JULIA_DEPOT_PATH = /usr/local/share/julia:
JULIA_PATH = /usr/local/julia
JULIA_PROJECT = @breeze
These were the top-level packages installed in the environment:
import PkgPkg.status()Status `/__w/Breeze.jl/Breeze.jl/docs/Project.toml`
[86bc3604] AtmosphericProfilesLibrary v0.1.10
[660aa2fb] Breeze v0.11.3 `.`
⌃ [052768ef] CUDA v6.2.2
[13f3f980] CairoMakie v0.15.15
[6a9e3e04] CloudMicrophysics v0.41.0
[e30172f5] Documenter v1.19.0
[daee34ce] DocumenterCitations v1.5.0
[b6400b83] DocumenterCodeBlocks v1.5.1
[1c52b33b] DocumenterLandingPage v0.2.2
[7da242da] Enzyme v0.13.206
⌅ [46192b85] GPUArraysCore v0.2.0
[63c18a36] KernelAbstractions v0.9.43
[98b081ad] Literate v2.21.0
[85f8d34a] NCDatasets v0.14.15
[9e8cae18] Oceananigans v0.113.4
[a01a1ee8] RRTMGP v1.0.1
[3c362404] Reactant v0.2.289
[276daf66] SpecialFunctions v2.9.0
[b77e0a4c] InteractiveUtils v1.11.0
[44cfe95a] Pkg v1.13.0
[9a3f8284] Random v1.11.0
Info Packages marked with ⌃ and ⌅ have new versions available. Those with ⌃ may be upgradable, but those with ⌅ are restricted by compatibility constraints from upgrading. To see why use `status --outdated`
This page was generated using Literate.jl.