Capability index: GeoModels

Layers: GeoStack · GeoNumerics · GeoML · GeoModels.

Every public name in the libraries geostack builds on, each listed once under the shortest module path that exports it, with the first sentence of its docstring. Search this page before writing a loader, a model, a training step or a helper in a script or a notebook: if what you need is here, use it; if it is almost here, extend it where it lives. Regenerate with make capabilities after changing a public API (tests/test_capabilities.py checks it is current).

This page is the GeoModels layer. It may import GeoML, GeoNumerics, GeoStack.

Listed at plumax 0.1.2, somax 0.0.15, radtrax 0.0.0.

Upstream: `plumax`¶

Plume transport: Gaussian plume and puff, Lagrangian and LES.

`plumax.gauss_plume`¶

Name	Kind	What it does
`BRIGGS_DISPERSION_PARAMS`	constant	`{'A': Array([ 2.2e-01, 1.0e-04, -5.0e-01, 2.0e-01, 0.0...`
`MIN_WIND_SPEED`	constant	`0.5`
`STABILITY_CLASSES`	constant	`('A', 'B', 'C', 'D', 'E', 'F')`
`calculate_briggs_dispersion`	function	Evaluate σ_y(x), σ_z(x) from Briggs parameters.
`gaussian_plume_model`	function	NumPyro model for the steady-state Gaussian plume.
`get_dispersion_params`	function	Return the Briggs parameter vector for a stability class.
`infer_emission_rate`	function	Infer the plume emission rate via NUTS.
`plume_concentration`	function	Evaluate the steady-state Gaussian plume concentration.
`plume_concentration_vmap`	function	Vectorized version of plume_concentration.
`rotate_to_wind_frame`	function	Rotate fixed-frame coordinates into the wind-aligned frame.
`simulate_plume`	function	Simulate a steady-state Gaussian plume on a 3-D grid.

`plumax.gauss_puff`¶

Name	Kind	What it does
`DISPERSION_SCHEMES`	constant	`{'pg': ({'A': Array([-1.104 , 0.9878, -0.0076, 4.679 , ...`
`OUTurbulence`	class	Ornstein-Uhlenbeck turbulence for puff-release position disturbances.
`PG_DISPERSION_PARAMS`	constant	`{'A': Array([-1.104 , 0.9878, -0.0076, 4.679 , -1.7172,...`
`PuffState`	class	Container for the Lagrangian state of a puff ensemble at a single time.
`STABILITY_CLASSES`	constant	`('A', 'B', 'C', 'D', 'E', 'F')`
`WindSchedule`	class	Piecewise-linear time-varying wind field.
`calculate_briggs_dispersion_xyz`	function	Evaluate σ_x(s), σ_y(s), σ_z(s) from Briggs parameters.
`calculate_pg_dispersion`	function	Evaluate σ_x(s), σ_y(s), σ_z(s) from Pasquill-Gifford parameters.
`cumulative_wind_integrals`	function	Compute cumulative wind integrals via a single diffrax ODE solve.
`evolve_puffs`	function	Advance a puff ensemble to `current_time` under a time-varying wind.
`frequency_to_release_interval`	function	Convert a release frequency f [Hz] into a release interval Δt [s].
`gaussian_puff_model`	function	NumPyro model for the Gaussian puff forward with a constant Q prior.
`gaussian_puff_rw_model`	function	Random-walk state-space model for a time-varying emission rate Q_i.
`get_dispersion_scheme`	function	Return `(params_dict, calculator_fn)` for a named dispersion scheme.
`get_pg_params`	function	Return the Pasquill-Gifford parameter vector for a stability class.
`infer_emission_rate`	function	NUTS inference of a constant emission rate from puff observations.
`infer_emission_timeseries`	function	NUTS inference of a time-varying emission rate series Q_i.
`make_release_times`	function	Evenly spaced puff release times in `[t_start, t_end)`.
`puff_concentration`	function	Evaluate the concentration from a single Gaussian puff (JIT).
`puff_concentration_vmap`	function	Vectorized version of puff_concentration.
`release_durations`	function	Per-puff emission-interval durations [s] for a continuous source.
`release_interval_to_frequency`	function	Convert a release interval Δt [s] into a release frequency f [Hz].
`sample_ou_offsets`	function	Draw per-puff (Δx, Δy) offsets from a 2-D OU process at release times.
`simulate_puff`	function	Simulate a time-resolved Gaussian-puff dispersion field on a 3-D grid.
`simulate_puff_field`	function	Sum the contributions of all puffs at each receptor.

`plumax.lagrangian`¶

Name	Kind	What it does
`GaussianPosterior`	class	Closed-form Gaussian posterior over the source vector.
`HannaTurbulence`	class	Height-dependent Hanna (1982) turbulence as a `TurbulenceModel`.
`HomogeneousTurbulence`	class	Spatially uniform, well-mixed-exact turbulence statistics.
`LognormalPosterior`	class	Linearised lognormal posterior over a non-negative source vector.
`ParticleState`	class	State of an ensemble of Lagrangian particles.
`bin_positions`	function	Histogram particle `positions` `(n, 3)` onto a 3-D grid (counts).
`compute_footprint`	function	Compute a single receptor’s surface footprint by backward integration.
`hanna_profiles`	function	Hanna (1982) Monin–Obukhov-similarity turbulence profiles.
`integrate_particles`	function	Integrate the ensemble from `t0` to `t1` with fixed step `dt`.
`langevin_step`	function	Advance the ensemble one step `dt` with the Markov-1 update.
`linear_gaussian_inversion`	function	Closed-form Gaussian–Gaussian (BLUE) inversion of the source vector.
`lognormal_inversion`	function	Sign-constrained inversion via a lognormal prior linearised in log-space.
`matern32_covariance`	function	Build a Matérn-3/2 prior covariance over grid-cell coordinates.
`n_steps_for_horizon`	function	Number of fixed-`dt` steps covering `horizon` (ceil, float-robust).
`observation_covariance`	function	Assemble the diagonal observation-error covariance `R`.
`simulate_lagrangian`	function	Simulate a steady Lagrangian concentration field on a 3-D grid.
`step_durations`	function	Per-step durations summing to `horizon`; the final step takes the rest.
`uniform_wind`	function	Return a constant mean-wind field `t -> (u, v, w)`.
`wind_from_speed_direction`	function	Constant wind from speed + meteorological direction (degrees from North).

`plumax.les_fvm`¶

Name	Kind	What it does
`ColumnObservationOperator`	class	Sample the column-integrated field at fixed receptor cells.
`EddyDiffusivity`	class	Scalar or field-valued `(K_h, K_z)` eddy diffusivity.
`EulerianDispersionRHS`	class	Full tracer-transport tendency for diffrax.
`EulerianForward4DVar`	class	Differentiable map from an emission time series to column observations.
`FourDVarProblem`	class	A strong-constraint 4D-Var problem over a scalar emission time series.
`FourDVarResult`	class	Outcome of a 4D-Var minimisation.
`GaussianSource`	class	Gaussian point source with constant or time-varying emission rate.
`HorizontalBC`	class	Apply a :class:`BoundaryConditionSet` to every z-slice of a 3-D field.
`PlumeGrid3D`	class	Arakawa C-grid + cached interior coordinate arrays.
`PosteriorCovariance`	class	Laplace (Gauss-Newton) posterior covariance of a 4D-Var solution.
`PrescribedWindField`	class	Wind velocity evaluator on the full `[Nz, Ny, Nx]` field layout.
`VerticalBC`	class	Top/bottom ghost-slice update for a 3-D field.
`advection_tendency`	function	Flux-form advective tendency `-∇·(u C)` at interior T-points.
`apply_boundary_conditions`	function	Apply horizontal then vertical BCs to a 3-D tracer field.
`build_default_concentration_bc`	function	Build `(HorizontalBC, VerticalBC)` from user-facing BC specs.
`build_forward`	function	Assemble a differentiable :class:`EulerianForward4DVar`.
`build_problem`	function	Assemble a :class:`FourDVarProblem`, factorising the prior covariance.
`diffusion_tendency`	function	Full 3-D eddy diffusion tendency `∇·(K ∇C)` at interior T-points.
`laplace_sample`	function	Draw source-space samples from the Laplace posterior `N(S*, P_S)`.
`make_gaussian_source`	function	Build a :class:`GaussianSource` on `plume_grid`.
`make_grid`	function	Build a :class:`PlumeGrid3D` from `(x_min, x_max, nx_interior)` triples.
`pg_eddy_diffusivity`	function	Rough K-theory calibration from Pasquill-Gifford σ curves.
`posterior_covariance`	function	Gauss-Newton Laplace posterior covariance around a 4D-Var MAP.
`simulate_eulerian_dispersion`	function	Simulate 3-D Eulerian tracer transport under a prescribed wind.
`solve_4dvar`	function	Minimise the 4D-Var cost with L-BFGS in whitened space.
`stable_step_bound`	function	Maximum stable step for an explicit fixed-step transport solve.
`uniform_wind_field`	function	Space- and time-constant wind field.
`wind_field_from_callable`	function	General-purpose wind field from a user-supplied `(t, X, Y, Z)` function.
`wind_field_from_schedule`	function	Time-varying but spatially uniform wind field.

`plumax.met`¶

Name	Kind	What it does
`MetField`	class	Meteorological forcing on an analysis grid with a time axis.
`load_wrf`	function	Load a `wrfout` file onto an analysis grid.
`to_prescribed_wind`	function	Expose a :class:`MetField` as the wind object the Eulerian solver takes.

`plumax.population`¶

Name	Kind	What it does
`AdditiveMultiCovariatePOD`	class	Additive GLM with separate coefficients for each covariate.
`CloglogPOD`	class	Complementary log-log GLM for asymmetric detection tails.
`CoalMineVentilationIntensity`	class	Multi-scale operational intensity for underground coal mines.
`ConcentrationProxyPOD`	class	GLM on the log-concentration proxy ln(Q/(U·p)).
`ConstantIntensity`	class	Constant intensity for passive, continuously leaking infrastructure.
`DiurnalSeasonalIntensity`	class	Additive diurnal + seasonal intensity for thermally active biogenic sources.
`DiurnalSinusoidalIntensity`	class	Diurnal sinusoidal intensity for thermally driven infrastructure.
`EmissionCatalog`	class	A population of detected emission events.
`EmissionEvent`	class	A single detected emission event with its per-event posterior summary.
`FacilityConfig`	class	Immutable physical configuration of a single leaking facility.
`FullVaryingCoefficientPOD`	class	Full varying-coefficient GLM with spectral and radiometric controls.
`INTENSITY_REGISTRY`	constant	`{'constant': <class 'plumax.population.intensity.Constant...`
`InhomogeneousIntensityPosterior`	class	Posterior over a log-linear inhomogeneous intensity model.
`LandfillIntensity`	class	Multi-component intensity for municipal solid waste landfills.
`LivestockFeedlotIntensity`	class	Bimodal diurnal + seasonal intensity for concentrated livestock operations.
`LogLogisticPOD`	class	Logistic POD on log-flux: the lognormal CDF approximation.
`LogisticPOD`	class	Single-covariate logistic POD: the classic S-curve.
`OffshorePlatformIntensity`	class	Combined tidal and operational intensity for offshore oil & gas platforms.
`OperationalScheduleIntensity`	class	Step-function intensity for facilities with fixed shift schedules.
`POD_REGISTRY`	constant	`{'logistic': <class 'plumax.population.pod.LogisticPOD'>,...`
`ParadoxResult`	class	Container for a single Monte Carlo paradox simulation.
`PerEventPosterior`	class	Structural type for a scalar per-event `Q` posterior summary.
`PeriodicBatchIntensity`	class	Periodic narrow-pulse intensity for scheduled batch venting.
`PoissonRatePosterior`	class	Gamma posterior over a homogeneous Poisson rate `lambda`.
`ProbitPOD`	class	Probit GLM on the concentration proxy.
`SNRBasedPOD`	class	GLM on log-SNR: radiometric first-principles detection model.
`SeasonalSinusoidalIntensity`	class	Annual sinusoidal intensity for biogenic / climate-driven sources.
`SeasonallyModulatedDiurnalIntensity`	class	Diurnal intensity whose amplitude is seasonally modulated.
`SizeDistributionPosterior`	class	Posterior over the lognormal population size-distribution parameters.
`SpectralAwarePOD`	class	GLM with spectral band count and resolution as covariates.
`VaryingCoefficientPOD`	class	Varying coefficient GLM: sensitivity β₁ depends on environment.
`WeibullRenewalIntensity`	class	Weibull hazard function for pressure-driven, recharging sources.
`WetlandPermafrostIntensity`	class	Sigmoid-gated seasonal intensity for high-latitude wetland/permafrost.
`X_MAX_DEFAULT`	constant	`1000000.0`
`X_MIN_DEFAULT`	constant	`10.0`
`build_canonical_scenarios`	function	Six canonical test scenarios spanning the paradox regime.
`compute_E_Pd`	function	Compute E[P_d] = ∫₀^∞ P_d(Q) · f(Q) dQ by log-space quadrature.
`event_from_posterior`	function	Build an :class:`EmissionEvent` from any Tier II-IV posterior.
`fit_inhomogeneous_intensity`	function	Log-linear inhomogeneous Poisson-process intensity fit via NUTS.
`fit_lognormal_size_distribution`	function	Hierarchical-Bayes fit of the lognormal population size distribution.
`fit_poisson_rate`	function	Closed-form Gamma-Poisson posterior for a homogeneous event rate.
`logistic_pod`	function	Parametric probability of detection (logistic sigmoid).
`lognorm_cdf`	function	Lognormal CDF used as the probability-of-detection curve.
`lognormal_pdf`	function	Evaluate the Lognormal probability density function.
`pod_powerlaw_model`	function	NumPyro model for a POD-modified power law on observed fluxes.
`power_law`	function	Un-normalised power-law PDF: x⁻ᵅ.
`run_mcmc`	function	Run NumPyro NUTS on the POD-modified power-law model.
`run_scenario_grid`	function	Run the paradox simulation across multiple scenarios.
`simulate_paradox`	function	Run a single Monte Carlo realisation of the Missing Mass Paradox.

`plumax.coupled`¶

Name	Kind	What it does
`CoupledForward`	class	Multi-instrument coupled forward: a list of per-instrument observations.
`FusionPosterior`	class	Closed-form joint posterior over `(Q, bias₁ … bias_M)`.
`Instrument`	class	Fixed observation geometry + error budget for one satellite instrument.
`LinearisedRadianceOperator`	class	Tangent-linear radiance operator — the linearised counterpart of the RTM.
`PlumeSource`	class	Static source + meteorology for the Tier I coupled forward.
`RadianceObservationOperator`	class	Map a per-receptor column-enhancement field to band radiances via the RTM.
`build_coupled_forward`	function	Assemble a :class:`CoupledForward`, precomputing unit-`Q` responses.
`column_mass_to_delta_vmr`	function	Convert a plume column mass enhancement [kg/m²] to a gas `ΔVMR`.
`column_response`	function	Per-receptor column enhancement at unit emission for one instrument.
`default_prior`	function	Build a diagonal prior `(θ_b, B)` for `(Q, bias₁ … bias_M)`.
`fuse_observations`	function	Closed-form joint posterior over `(Q, bias_inst)` from fused observations.
`linearise`	function	Build the :class:`LinearisedRadianceOperator` tangent to `operator` at 0.
`predict_observation`	function	Predict one instrument’s observation vector `y_inst`.
`radiance_response`	function	Per-receptor band radiances for one instrument under the RTM operator.
`radiance_response_linear`	function	Per-receptor linearised radiances for one instrument (differentiable).

`plumax.operators`¶

Name	Kind	What it does
`EulerianDispersion`	class	Eulerian finite-volume dispersion runner as a pipekit operator.
`GaussianPlume`	class	Steady-state Gaussian plume forward model as a pipekit operator.
`GaussianPuff`	class	Time-resolved Gaussian puff forward model as a pipekit operator.
`LagrangianDispersion`	class	Markov-1 Lagrangian dispersion forward model as a pipekit operator.

`plumax.adapters`¶

Name	Kind	What it does
`EulerianForwardModel`	class	Expose an Eulerian dispersion RHS as a `pipekit_cycle.ForwardModel`.
`RadtranObservationOperator`	class	Expose the normalised Beer–Lambert forward as an `ObservationOperator`.

`plumax.radtran`¶

Name	Kind	What it does
`ATM_TO_PA`	constant	`101325.0`
`BOLTZMANN_J_PER_K`	constant	`1.380649e-23`
`CM3_PER_M3`	constant	`1000000.0`
`ForwardResult`	class	Bundle of (radiance, Jacobian, transmittance) returned by the forward models.
`GroundSamplingDistance`	class	Block-mean down-sampling operator for a fixed integer factor.
`InstrumentSpec`	class	Multispectral instrument description.
`NBLookup`	class	Pre-computed normalised-brightness table, indexed by band and ΔX.
`ObservationGeometry`	class	Viewing geometry for a single pixel.
`PointSpreadFunction`	class	2D PSF blur applied independently per spectral channel.
`SpectralResponseFunction`	class	Band-integration operator for a multispectral / hyperspectral instrument.
`build_lowrank_covariance_operator`	function	Build a :class:`gaussx.LowRankUpdate` covariance operator from a scene.
`build_nb_lut`	function	Pre-tabulate band-integrated nB over a column-enhancement grid.
`forward_maclaurin`	function	Maclaurin-series forward model: expand `exp(-τ(VMR))` around VMR = 0.
`forward_maclaurin_normalized`	function	Maclaurin expansion of `exp(-Δτ)` around `ΔVMR = 0`.
`forward_nonlinear`	function	Exact Beer-Lambert forward model.
`forward_nonlinear_normalized`	function	Exact Beer-Lambert normalised by the background radiance.
`forward_taylor`	function	First-order Taylor expansion around `vmr_background`.
`forward_taylor_normalized`	function	Taylor expansion of `L_norm` around `ΔVMR = 0`.
`inject_plume`	function	Multiplicative plume injection: `scene_plume = scene · nB(ΔX)`.
`lookup_nb`	function	Look up per-band nB for every pixel of a 2-D column map.
`matched_filter_image`	function	Matched-filter abundance map, vectorised over spatial axes.
`matched_filter_image_op`	function	Vectorised matched filter using a structured covariance operator.
`matched_filter_pixel`	function	Matched-filter abundance for a single pixel.
`matched_filter_pixel_op`	function	Scalar matched filter using a structured covariance operator.
`matched_filter_snr`	function	Detection SNR for a matched-filter abundance map.
`matched_filter_snr_op`	function	Detection SNR when the covariance is a gaussx operator.
`number_density_cm3`	function	Ideal-gas total number density [molecules / cm³].
`robust_lowrank_covariance`	function	Robust low-rank covariance + inverse `(Σ, Σ⁻¹)`.
`target_bands`	function	Band-integrate a hi-res target spectrum through an SRF.
`target_spectrum_normalized_linear`	function	First-order linearised target `t(ν) = -σ · N · ΔVMR · L · AMF`.
`target_spectrum_normalized_nonlinear`	function	Exact normalised-radiance target `t(ν) = exp(-Δτ(ν)) − 1`.
`trimmed_mean_spectrum`	function	Per-channel trimmed mean `μ_b`.

`plumax.hapi_lut`¶

Name	Kind	What it does
`ATMOSPHERIC_GASES`	constant	`{'CH4': GasConfig(CH4, M6, I1), 'CO2': GasConfig(CO2, M2,...`
`DEFAULT_VMR_NOMINAL`	constant	`{'CH4': 2e-06, 'CO2': 0.0004, 'H2O': 0.01, 'O2': 0.21, 'N...`
`GasConfig`	class	Configuration for a single HITRAN gas species.
`HAPI_CACHE_ENV`	constant	`'HAPI_CACHE_DIR'`
`LUTGridConfig`	class	Grid configuration shared across gases for a single LUT build.
`absorption_coefficient`	function	α(ν) = σ(ν) · N_gas [1/cm].
`air_mass_factor`	function	Plane-parallel two-way AMF for sun → surface → sensor.
`beers_law_from_lut`	function	End-to-end: σ-from-LUT → α → τ(ν) for a single homogeneous layer.
`build_lut_dataset`	function	Wrap the raw LUT arrays in a CF-conformant `xarray.Dataset`.
`compute_absorption_lut`	function	Evaluate σ(ν, T, P) on the full LUT grid.
`create_combined_lut`	function	Build a single NetCDF with one cross-section variable per gas.
`create_multi_gas_luts`	function	Build one separate `<gas>_absorption_lut.nc` per gas.
`default_cache_dir`	function	Return the HITRAN line-data cache directory.
`fetch_hitran_data`	function	Fetch HITRAN line parameters for `gas_config` into `cache_dir`.
`generate_single_gas_lut`	function	End-to-end single-gas pipeline: fetch → compute → wrap → (optionally) save.
`interpolate_cross_section`	function	Return σ(ν) at a given (T, p) via bilinear LUT interpolation.
`number_density`	function	Return total number density [molecules / cm³] from the ideal-gas law.
`plume_ratio_spectrum`	function	Differential-Beer-Lambert ratio τ_total(ν) / τ_bg(ν).
`save_lut`	function	Persist `ds` to `<output_dir>/<gas_name>_absorption_lut.nc` (zlib).
`transmittance`	function	τ(ν) = exp(-α · L_vert · AMF).

`plumax.assimilation`¶

Name	Kind	What it does
`Cost`	class	Bundle of (value, gradient, Hessian-vector product) for an optimiser.
`IdentityTransform`	class	Pass-through transform.
`RadianceObservationModel`	class	Bundle of precomputed pieces that defines the forward operator H.
`SolveResult`	class	Common result type for all three solvers.
`WhiteningTransform`	class	`δx = U ξ` where `B = U Uᵀ`.
`build_cost_x`	function	Build `J(δx) = ½ δxᵀ B⁻¹ δx + ½ (y − H(x_b + δx))ᵀ R⁻¹ (y − H(x_b + δx))`.
`build_cost_xi`	function	Build `J(ξ) = ½ ξᵀ ξ + ½ (y − H(x_b + U ξ))ᵀ R⁻¹ (…)`.
`build_diagonal_background`	function	Build `B = diag(σ²)` as a :class:`lineax.DiagonalLinearOperator`.
`build_kronecker_background`	function	Build `B = σ² · (ρ_y ⊗ ρ_x)` using :class:`gaussx.Kronecker`.
`build_lowrank_background`	function	Estimate `B` from a stack of state samples and wrap as `λI + U D Uᵀ`.
`degrees_of_freedom_for_signal`	function	Estimate `DFS = trace(I − Bₐ B⁻¹)` (Rodgers 2000, §2.5).
`finite_difference_grad`	function	Central finite differences — used only by tests to validate `jax.grad`.
`posterior_covariance_proxy`	function	Return a callable `v → Bₐ v` (model-space posterior covariance) via CG.
`reduced_chi_squared`	function	`χ²_red = (y - H(x̂))ᵀ R⁻¹ (y - H(x̂)) / dim(y)`.
`run_dual_psas`	function	Solve the dual 3D-Var system in observation space.
`run_gauss_newton`	function	Run :class:`optimistix.GaussNewton` for `min ½ ‖r(y)‖²`.
`run_lbfgs`	function	Minimise `cost.value` via :class:`optimistix.LBFGS`.

`plumax.matched_filter`¶

Name	Kind	What it does
`WelfordAccumulator`	class	Online mean and covariance for `(n_samples, n_bands)` streams.
`adaptive_window_background`	function	Per-pixel local mean and per-band variance over a square window.
`apply_image`	function	Matched-filter score map for a hyperspectral cube.
`apply_image_xarray`	function	Apply the matched filter to an xarray DataArray.
`apply_pixel`	function	Matched-filter score for a single pixel spectrum.
`detection_threshold`	function	Gaussian detection threshold on the MF score for a given FAR.
`estimate_cov_empirical`	function	Empirical covariance via :class:`sklearn.covariance.EmpiricalCovariance`.
`estimate_cov_lowrank`	function	Low-rank + Tikhonov covariance: `Σ = λI + V D Vᵀ`.
`estimate_cov_shrunk`	function	Shrinkage covariance estimator.
`estimate_mean`	function	Estimate the per-band background mean `μ` from a scene.
`gmm_cluster_background`	function	Fit a Gaussian mixture and build per-cluster background models.
`linear_target_from_obs`	function	Matched-filter target via the tangent-linear of the forward operator.
`matched_filter_snr`	function	Theoretical SNR of the matched filter at a given plume amplitude.
`nonlinear_target_from_obs`	function	Matched-filter target via a finite-amplitude difference.
`open_multi_scene`	function	Yield `(n_pixels, n_bands)` batches from a sequence of scene files.
`streaming_background`	function	Streamed `(μ, Σ)` from an iterable of batches.
`validate_mf_inputs`	function	Public fail-fast check for `(Σ, t)` compatibility.

Upstream: `somax`¶

Simple ocean and atmosphere models in JAX.

`somax.core`¶

Name	Kind	What it does
`BasisForcing`	class	Reduced-order forcing: a fixed space-time frame driven by a coefficient vector.
`Compose`	class	Operator composition: `Compose(outer, inner)(state) = outer(inner(state))`.
`ConstantForcing`	class	Time-independent forcing field.
`ConstantInTime`	class	Time-independent gate: `b(t) = ones(m)` (the `Phi_t = I` case).
`Diagnostics`	class	Base class for on-demand diagnostic quantities.
`DirichletHelmholtzCache`	class	Helmholtz cache for Dirichlet (DST-based) domains.
`ForcingProtocol`	class	Base class for forcing terms.
`ForcingTerm`	class	Lift a :class:`ForcingProtocol` field into the term algebra as a tendency.
`FourierInTime`	class	Cosine temporal gate `b_a(t) = cos(omega_a t + phase_a)`.
`GaussianWindowsInTime`	class	Localized temporal gate `b_a(t) = exp(-(t - tau_a)^2 / (2 T_a^2))`.
`HelmholtzCache`	class	Cached Helmholtz solver for repeated (lap - lambda) psi = f solves.
`InterpolatedForcing`	class	Data-driven forcing via diffrax path interpolation.
`ModalTransform`	class	Precomputed layer-to-mode and mode-to-layer transforms.
`MultimodalHelmholtzCache`	class	Batched Helmholtz cache for multilayer modal solves.
`NeumannHelmholtzCache`	class	Helmholtz cache for Neumann (DCT-based) domains.
`NoForcing`	class	Zero forcing (free evolution).
`Params`	class	Base class for differentiable model parameters.
`PeriodicHelmholtzCache`	class	Helmholtz cache for periodic (FFT-based) domains.
`PhysConsts`	class	Base class for frozen physical constants.
`Scaled`	class	A term multiplied by a scalar coefficient.
`ScaledModel`	class	Integrate `inner` in transformed coordinates.
`Scales`	class	Characteristic scales of a run.
`SeasonalWindForcing`	class	Sinusoidal wind forcing with learnable amplitude.
`SimulationCheckpointer`	class	Manages periodic checkpointing during long simulations.
`SomaxModel`	class	Abstract base class defining the somax model contract.
`SpatialBasis`	class	A precomputed spatial dictionary plus the per-mode prior std.
`State`	class	Base class for model state vectors.
`StateAffine`	class	Per-leaf affine map on a `State` pytree: `y = (x - loc) / scale`.
`StratificationProfile`	class	Discrete vertical stratification for a layered ocean model.
`Sum`	class	Additive composition of terms: `sum_i terms[i](t, state, args)`.
`TemporalBasis`	class	Maps a scalar time to per-atom temporal weights `b(t)`.
`Term`	class	A single additive contribution to a model right-hand side.
`TermFn`	class	Adapt a plain callable into a :class:`Term`.
`TermModel`	class	A :class:`SomaxModel` whose RHS is an assembled :class:`Term` tree.
`TransformedForcing`	class	Apply a pointwise transform to a base forcing (e.g. log-space synthesis).
`VectorBasisForcing`	class	Reduced-order vector forcing: a fixed vector frame driven by coeffs.
`VectorSpatialBasis`	class	A precomputed vector dictionary plus the per-mode prior std.
`add_to`	function	Build a placement that adds a field onto one named state component.
`add_vector_to`	function	Build a placement that adds a component-major field onto named components.
`as_parameter`	function	Coerce a factory argument into a `Params` leaf.
`build_diffrax_terms`	function	Build the diffrax term object for `term`, IMEX-aware.
`control_filter`	function	Boolean filter selecting only `coeffs` for gradient updates.
`explicit`	function	Tag `term` for the explicit stage of an IMEX integrator.
`frozen`	function	Hide a parameter from gradients while keeping it a runtime value.
`geostrophic_currents`	function	Vector preset: incompressible `(u, v)` current-error forcing.
`implicit`	function	Tag `term` for the implicit stage of an IMEX integrator.
`interval`	function	Constrain a parameter to the open interval `(lower, upper)`.
`matern_spectral_density`	function	Matérn power spectral density `S(omega)` at `omega = sqrt_lambda`.
`partition`	function	Split `term` into `(explicit, implicit)` sub-terms by kind.
`positive`	function	Constrain a parameter to be strictly positive.
`spatial_from_divfree`	function	Build a :class:`VectorSpatialBasis` of divergence-free velocity atoms.
`spatial_from_eof`	function	Build a :class:`SpatialBasis` from empirical orthogonal functions (PCA).
`spatial_from_fourier`	function	Build a :class:`SpatialBasis` from the box-Laplacian (HSGP) eigenbasis.
`spatial_from_gabor`	function	Build a :class:`SpatialBasis` from a geonnax dyadic radial-Gabor frame.
`spatial_from_graph_laplacian`	function	Build a :class:`SpatialBasis` from graph-Laplacian eigenvectors.
`spatial_from_rbf`	function	Build a :class:`SpatialBasis` from a geonnax placeable radial basis.
`spatial_from_spherical_rbf`	function	Build a :class:`SpatialBasis` from a geodesic (on-sphere) radial basis.
`spatial_from_wavelet`	function	Build a :class:`SpatialBasis` from the orthonormal 2D wavelet basis.
`ssh_geostrophic`	function	SSH geostrophic preset: a radial-Gabor frame with a wavenumber-law prior.
`sss_coastal`	function	Coastal SSS preset: placeable radial atoms with a prescribed prior.
`sst_frontal`	function	SST preset: a smooth Matérn (HSGP) field over the box-Laplacian eigenbasis.
`tile_in_time`	function	Lift a spatial dictionary into a separable space-time frame.
`trainable_mask`	function	Boolean pytree marking which leaves an optimiser may update.

`somax.domain`¶

Name	Kind	What it does
`Domain`	class	Domain class for a rectangular domain
`TimeDomain`	class	TimeDomain(tmin, tmax, dt)

`somax.operators`¶

Name	Kind	What it does
`Burgers2DOp`	class	Serializable pipekit Operator for the term-based 2D Burgers model.
`SomaxModelOp`	class	A pipekit Operator wrapping any built somax forward model.

`somax.solvers`¶

Name	Kind	What it does
`imex_solver`	function	Build a `KenCarp3` IMEX solver with a matrix-free implicit stage.
`imex_stepsize_controller`	function	Build an adaptive PID controller for an IMEX solve.

`somax.da`¶

Name	Kind	What it does
`SomaxDynamics`	class	Adapt a somax forward model to filterax’s flat-vector dynamics.
`SomaxForwardModel`	class	Adapt a somax model to vardax’s flat-vector `ForwardModel`.
`SubsampleObs`	class	Linear sub-sampling observation operator `H(x) = x[indices]`.
`make_ensemble`	function	Build a Gaussian-perturbed flat ensemble around a base state.
`state_to_vector`	function	Flatten a state pytree to a 1-D vector plus its inverse.

`somax.flows`¶

Name	Kind	What it does
`StateAffineBijection`	class	A :class:`StateAffine` as a flat flowjax bijection.
`learned_prior`	function	Compose a fixed physical scaling with a learned flow into a prior.
`to_flowjax`	function	Adapt a :class:`StateAffine` to the flowjax bijection interface.

`somax.eval`¶

Name	Kind	What it does
`compute_eval_metrics`	function	Compute every applicable reference-free metric for `(model, state)`.
`geostrophic_imbalance`	function	Dimensionless ageostrophic fraction of a shallow-water-type state.
`kinetic_energy`	function	Domain-integrated kinetic energy `0.5 ∫ (u² + v²) dA`.
`qg_balance_residual`	function	Dimensionless PV-inversion residual for a barotropic QG model.
`rms_divergence`	function	Root-mean-square horizontal divergence `sqrt(<(∇·u)²>)`.
`total_enstrophy`	function	Domain-integrated enstrophy `0.5 ∫ ζ² dA` with `ζ = ∂v/∂x - ∂u/∂y`.

`somax.io`¶

Name	Kind	What it does
`append_to_dataset`	function	Append a Dataset to an existing zarr store along a dimension.
`apply_scale_metadata`	function	Attach scale / transform metadata to an exported Dataset.
`dataset_to_state`	function	Reconstruct a single State from an xarray Dataset.
`load_dataset`	function	Open a zarr store as an xarray Dataset.
`save_dataset`	function	Persist a Dataset to a zarr store.
`scales_attrs`	function	Flat Dataset attrs describing the scale set a run used.
`snapshots_to_dataset`	function	Convert a time-stacked PyTree of states to an xarray Dataset.
`state_to_dataset`	function	Convert a single model State to an xarray Dataset.
`transform_attrs`	function	Per-variable `loc` / `scale` attrs for an affine transform.

`somax.monitor`¶

Name	Kind	What it does
`BaseMonitor`	class	Inert base monitor — override only the hooks you care about.
`ChunkInfo`	class	Context handed to a monitor at a diagnostic-chunk boundary.
`ConservationDriftMonitor`	class	MONITOR (optionally FAIL-HARD): track drift of conserved invariants.
`EnergyGrowthMonitor`	class	MONITOR (optionally FAIL-HARD): flag large energy jumps between chunks.
`Monitor`	class	Structural protocol for a simulation monitor.
`MonitorVerdict`	class	A monitor’s response to a chunk.
`NonFiniteMonitor`	class	FAIL-HARD: terminate the run if any state field holds NaN/Inf.
`SolverHealthMonitor`	class	MONITOR: surface diffrax per-chunk solver statistics.
`ThroughputMonitor`	class	MONITOR: report simulated seconds per wallclock second per chunk.
`WatchdogMonitor`	class	FAIL-HARD: terminate if cumulative wallclock exceeds a ceiling.
`default_monitors`	function	The runner’s default monitor set — preserves legacy behavior.

Upstream: `radtrax`¶

Differentiable radiative transfer for methane remote sensing.

`radtrax`¶

Name	Kind	What it does
`RadtraxError`	class	Base class for every exception raised by `radtrax`.

`radtrax.atmosphere`¶

Name	Kind	What it does
`ATM_TO_PA`	constant	`101325.0`
`BOLTZMANN_J_PER_K`	constant	`1.380649e-23`
`CM3_PER_M3`	constant	`1000000.0`
`air_mass_factor`	function	Plane-parallel two-way air-mass factor, sun → surface → sensor.
`number_density`	function	Total number density from the ideal-gas law [molecules / cm³].

`radtrax.spectroscopy`¶

Name	Kind	What it does
`absorption_coefficient`	function	α(ν) = σ(ν) · N_gas [1/cm].
`beers_law_from_lut`	function	End-to-end: σ-from-LUT → α → τ(ν) for a single homogeneous layer.
`interpolate_cross_section`	function	Return σ(ν) at a given (T, p) via bilinear LUT interpolation.
`plume_ratio_spectrum`	function	Differential-Beer-Lambert ratio τ_total(ν) / τ_bg(ν).
`transmittance`	function	τ(ν) = exp(-α · L_vert · AMF).

`radtrax.instrument`¶

Name	Kind	What it does
`GroundSamplingDistance`	class	Block-mean down-sampling operator for a fixed integer factor.
`PointSpreadFunction`	class	2D PSF blur applied independently per spectral channel.
`SpectralResponseFunction`	class	Band-integration operator for a multispectral / hyperspectral instrument.

`radtrax.slab`¶

Name	Kind	What it does
`ATM_TO_PA`	constant	`101325.0`
`BOLTZMANN_J_PER_K`	constant	`1.380649e-23`
`CM3_PER_M3`	constant	`1000000.0`
`ForwardResult`	class	Bundle of (radiance, Jacobian, transmittance) returned by the forward models.
`InstrumentSpec`	class	Multispectral instrument description.
`NBLookup`	class	Pre-computed normalised-brightness table, indexed by band and ΔX.
`ObservationGeometry`	class	Viewing geometry for a single pixel.
`build_nb_lut`	function	Pre-tabulate band-integrated nB over a column-enhancement grid.
`forward_maclaurin`	function	Maclaurin-series forward model: expand `exp(-τ(VMR))` around VMR = 0.
`forward_maclaurin_normalized`	function	Maclaurin expansion of `exp(-Δτ)` around `ΔVMR = 0`.
`forward_nonlinear`	function	Exact Beer-Lambert forward model.
`forward_nonlinear_normalized`	function	Exact Beer-Lambert normalised by the background radiance.
`forward_taylor`	function	First-order Taylor expansion around `vmr_background`.
`forward_taylor_normalized`	function	Taylor expansion of `L_norm` around `ΔVMR = 0`.
`inject_plume`	function	Multiplicative plume injection: `scene_plume = scene · nB(ΔX)`.
`lookup_nb`	function	Look up per-band nB for every pixel of a 2-D column map.
`number_density_cm3`	function	Ideal-gas total number density [molecules / cm³].

`radtrax.geo`¶

Name	Kind	What it does
`sentinel2_srf`	function	Build a Sentinel-2 SRF from the agency table georeader bundles.

Capability index: GeoModels

Upstream: plumax¶

plumax.gauss_plume¶

plumax.gauss_puff¶

plumax.lagrangian¶

plumax.les_fvm¶

plumax.met¶

plumax.population¶

plumax.coupled¶

plumax.operators¶

plumax.adapters¶

plumax.radtran¶

plumax.hapi_lut¶

plumax.assimilation¶

plumax.matched_filter¶

Upstream: somax¶

somax.core¶

somax.domain¶

somax.operators¶

somax.solvers¶

somax.da¶

somax.flows¶

somax.eval¶

somax.io¶

somax.monitor¶

Upstream: radtrax¶

radtrax¶

radtrax.atmosphere¶

radtrax.spectroscopy¶

radtrax.instrument¶

radtrax.slab¶

radtrax.geo¶