model/legacy/stochastic.py

# Copyright 2024 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Implementation of stochastic modules."""

import abc
import dataclasses
import enum
import logging
from typing import Any, Callable, Optional, Sequence, TypeVar, Union
import zlib

from dinosaur import coordinate_systems
from dinosaur import typing
import gin
import haiku as hk
import jax
import jax.numpy as jnp


tree_map = jax.tree_util.tree_map
tree_leaves = jax.tree_util.tree_leaves

Numeric = typing.Numeric
Quantity = typing.Quantity
_SOFTPLUS_INVERSE_1 = 0.5413248546129181

# CoreRandomState is advanced by a RandomField, and .to_*_values(core_state)
# produces the final (usable) random Array.
CoreRandomState = typing.Pytree
RandomnessState = typing.RandomnessState


def _validate_randomness_state(state: RandomnessState) -> None:
  """Validates that `state.core` is not `None`, raises an error otherwise."""
  if state.core is None:
    raise ValueError(
        f'Got {state.core=} when value is expected. '
        'Check how incoming randomness is initialized.'
    )


def make_positive_scalar(raw_parameter: typing.Array) -> jax.Array:
  """Positive [batch] scalar values, maps 0 --> 1 using a softplus(...)."""
  raw_parameter = jnp.asarray(raw_parameter)
  return jax.nn.softplus(raw_parameter + _SOFTPLUS_INVERSE_1)


# pylint: disable=logging-fstring-interpolation


################################################################################
# Single random fields that stand on their own.
################################################################################


class PreferredRepresentation(enum.Enum):
  """The preferred (for computational reasons) representation of a field."""

  NODAL = 'NODAL'
  MODAL = 'MODAL'


class RandomField(abc.ABC):
  """Base class for random fields."""

  def __init__(self, coords):
    self.coords = coords

  @property
  @abc.abstractmethod
  def preferred_representation(self) -> PreferredRepresentation | None:
    """The PreferredRepresentation for this field, or None if no preference."""

  @abc.abstractmethod
  def unconditional_sample(self, rng: typing.PRNGKeyArray) -> RandomnessState:
    """Sample the random field unconditionally."""

  @abc.abstractmethod
  def advance(self, state: RandomnessState) -> RandomnessState:
    """Updates the core state of a random field."""

  @abc.abstractmethod
  def to_modal_values(self, core_state: CoreRandomState) -> typing.Array | None:
    """Returns the modal rep. of the random field specified by this class."""

  @abc.abstractmethod
  def to_nodal_values(self, core_state: CoreRandomState) -> typing.Array | None:
    """Returns the nodal rep. of the random field specified by this class."""


RandomnessModule = Callable[..., RandomField]


_ADVANCE_SALT = zlib.crc32(b'advance')  # arbitrary uint32 value


T = TypeVar('T', typing.PRNGKeyArray, None)


def _prng_key_for_current_advance_step(
    randomness: typing.RandomnessState,
) -> typing.PRNGKeyArray | None:
  """Get a PRNG Key suitable for randomness in the current advance step."""
  if randomness.prng_key is None:
    return None
  salt = jnp.uint32(_ADVANCE_SALT) + jnp.uint32(randomness.prng_step)
  return jax.random.fold_in(randomness.prng_key, salt)


@gin.register
class NoRandomField(RandomField):
  """Module that disables randomness in a given module returning `None`."""

  def __init__(
      self,
      coords: coordinate_systems.CoordinateSystem,
      dt: float,
      physics_specs: Any,
      aux_features: Any,
      prefer_nodal: bool = True,
  ):
    """Constructs a ZerosRandomField.

    Args:
      coords: horizontal and vertical grid data.
      dt: nondimensionalized model time step.
      physics_specs: physical constants and definition of custom units.
      aux_features: additional static data.
      prefer_nodal: Whether this field should prefer a nodal representation.
    """
    super().__init__(coords)
    logging.info('[NGCM] Initializing NoRandomField')
    del dt, physics_specs, aux_features, prefer_nodal  # unused.

  @property
  def preferred_representation(self) -> PreferredRepresentation | None:
    return None

  def unconditional_sample(
      self, rng: typing.PRNGKeyArray | None
  ) -> RandomnessState:
    """Returns a zeros initialized state."""
    return RandomnessState(prng_key=rng, prng_step=0)  # pyrefly: ignore[unexpected-keyword]

  def advance(self, state: RandomnessState) -> RandomnessState:
    """Updates the state of a random gaussian field."""
    return RandomnessState(
        prng_key=state.prng_key, prng_step=state.prng_step + 1  # pyrefly: ignore[unexpected-keyword, unsupported-operation]
    )

  def to_nodal_values(self, core_state: CoreRandomState) -> typing.Array | None:
    del core_state  # unused.
    return None

  def to_modal_values(self, core_state: CoreRandomState) -> typing.Array | None:
    del core_state  # unused.
    return None


@gin.register
class ZerosRandomField(RandomField):
  """Implements a constant random field identically equal to zero."""

  def __init__(
      self,
      coords: coordinate_systems.CoordinateSystem,
      dt: float,
      physics_specs: Any,
      aux_features: Any,
      prefer_nodal: bool = True,
  ):
    """Constructs a ZerosRandomField.

    Args:
      coords: horizontal and vertical grid data.
      dt: nondimensionalized model time step.
      physics_specs: physical constants and definition of custom units.
      aux_features: additional static data.
      prefer_nodal: Whether this field should prefer a nodal representation.
    """
    super().__init__(coords)
    logging.info('[NGCM] Initializing ZerosRandomField')
    del dt  # unused
    del physics_specs  # unused.
    del aux_features  # unused.
    self._prefer_nodal = prefer_nodal

  @property
  def preferred_representation(self) -> PreferredRepresentation | None:
    if self._prefer_nodal:
      return PreferredRepresentation.NODAL
    else:
      return PreferredRepresentation.MODAL

  def unconditional_sample(
      self, rng: typing.PRNGKeyArray | None
  ) -> RandomnessState:
    """Returns a zeros initialized state."""
    if self._prefer_nodal:
      core = jnp.zeros(self.coords.horizontal.nodal_shape)
    else:
      core = jnp.zeros(self.coords.horizontal.modal_shape)
    return RandomnessState(
        core=core,  # pyrefly: ignore[unexpected-keyword]
        nodal_value=jnp.zeros(self.coords.horizontal.nodal_shape),  # pyrefly: ignore[unexpected-keyword]
        modal_value=jnp.zeros(self.coords.horizontal.modal_shape),  # pyrefly: ignore[unexpected-keyword]
        prng_key=rng,  # pyrefly: ignore[unexpected-keyword]
        prng_step=0,  # pyrefly: ignore[unexpected-keyword]
    )

  def advance(self, state: RandomnessState) -> RandomnessState:
    """Updates the state of a random gaussian field."""
    _validate_randomness_state(state)
    return RandomnessState(
        core=jnp.zeros_like(state.core),  # pyrefly: ignore[bad-argument-type, unexpected-keyword]
        nodal_value=jnp.zeros(self.coords.horizontal.nodal_shape),  # pyrefly: ignore[unexpected-keyword]
        modal_value=jnp.zeros(self.coords.horizontal.modal_shape),  # pyrefly: ignore[unexpected-keyword]
        prng_key=state.prng_key,  # pyrefly: ignore[unexpected-keyword]
        prng_step=state.prng_step + 1,  # pyrefly: ignore[unexpected-keyword, unsupported-operation]
    )

  def to_nodal_values(self, core_state: CoreRandomState) -> typing.Array | None:
    """Returns the ready-for-use Zeros random field."""
    return jnp.zeros(self.coords.horizontal.nodal_shape)

  def to_modal_values(self, core_state: CoreRandomState) -> typing.Array | None:
    """Returns the ready-for-use Zeros random field."""
    return jnp.zeros(self.coords.horizontal.modal_shape)


@gin.register
class GaussianRandomField(RandomField):
  """Implements gaussian random field with spatial and temporal correlations.

  This type of random fields is used in SPPT (stochastic physics
  parameterization tendencies) schemes, where each tendency due to physics
  parameterizations are multiplicatively perturbed by the value of such field.

  For implementation details see Appendix 8 in http://shortn/_56HCcQwmSS.

  With x ∈ EarthSurface, this field U is initialized at t=0 with
    U(0, x) = Σₖ Ψₖ(x) (1 - φ²)^(-0.5) σₖ γₖ σₖ ηₖ₀,
  where Ψₖ is the kth spherical harmonic basis function, φ² is the one timestep
  correlation, σₖ > 0 is a scaling factor, and ηₖ₀ are iid 1D unit Gaussians.

  With `variance` an init kwarg,
    E[U(0, x)] ≡ 0,
    1 / (4πR²) ∫ Var(U(0, x))dx = variance,
  regardless of coords (and the radius).

  Further states are generated with the recursion
    U(t + δ) = ϕ U(t) + σₖ ηₖₜ
  This ensures that U is stationary.

  In general,
    Cov(U(t, x), U(t + δ, y)) = ϕᵟ Σₖ Ψₖ(x) Ψₖ(y) (γₖ)².
  """

  def __init__(
      self,
      coords: coordinate_systems.CoordinateSystem,
      dt: float,
      physics_specs: Any,
      aux_features: Any,
      correlation_time: Union[jax.Array, Quantity, str] = gin.REQUIRED,
      correlation_length: Union[jax.Array, Quantity, str] = gin.REQUIRED,
      variance: Optional[Union[jax.Array, Quantity, str]] = gin.REQUIRED,
      clip: float = 6.0,
  ):
    """Constructs a GaussianRandomField.

    Args:
      coords: horizontal and vertical grid data.
      dt: nondimensionalized model time step.
      physics_specs: physical constants and definition of custom units.
      aux_features: additional static data.
      correlation_time: timescale with units over which autoregressive process
        decorrelates. Typical values in NWP range from hours to days.
      correlation_length: lengthscale with units over which random field is
        correlated. Typical values in NWP range from 500-2500 km.
      variance: The average (over EarthSurface) variance of the random field If
        None, this GRF always returns a zeros field and no RNGS are drawn.
      clip: number of standard deviations at which to clip randomness to ensure
        numerical stability.
    """
    del aux_features  # unused.
    super().__init__(coords)
    logging.info(
        '[NGCM] Initializing GaussianRandomField (possibly via'
        f' CenteredLognormalRandomField) with {variance=}, {correlation_time=},'
        f' {correlation_length=}'
    )

    tau = maybe_nondimensionalize(correlation_time, physics_specs)
    correlation_length = maybe_nondimensionalize(
        correlation_length, physics_specs
    )

    # In sampling, phi appears as 1 - phi**2 = 1 - exp(-2 dt / tau)
    self.one_minus_phi2 = -jnp.expm1(-2 * dt / tau)  # pyrefly: ignore[unsupported-operation]

    self.phi = jnp.exp(-dt / tau)  # pyrefly: ignore[unsupported-operation]

    self._variance = maybe_nondimensionalize(variance, physics_specs)  # σ²

    # [Palmer] states correlation_length = sqrt(2κT) / R, therefore
    self.kt = (correlation_length / self.coords.horizontal.radius) ** 2 / 2
    self.clip = clip

  @property
  def preferred_representation(self) -> PreferredRepresentation | None:
    return PreferredRepresentation.MODAL

  @property
  def _surf_area(self) -> jax.Array:
    """Surface area of sphere of radius self.coords.horizontal.radius."""
    return 4 * jnp.pi * self.coords.horizontal.radius**2  # pytype: disable=bad-return-type  # jnp-type

  def _sigma_array(self) -> jax.Array:
    """Array of σₙ from Appendix 8 in [Palmer] http://shortn/_56HCcQwmSS."""
    # n = [0, 1, ..., N]
    n = self.coords.horizontal.modal_axes[1]  # total wavenumbers.

    # Number of longitudinal wavenumbers at each total wavenumber n.
    #  L = 2n + 1, except for the last entry.
    n_longitudian_wavenumbers = self.coords.horizontal.mask.sum(axis=0)

    # sigmas_unnormed[n] is proportional to the standard deviation for each
    # longitudinal wavenumbers at each total wavenumber n.
    sigmas_unnormed = jnp.exp(-0.5 * self.kt * n * (n + 1))

    # The sum of unnormalized variance for all longitudinal wavenumbers at each
    # total wavenumber.
    sum_unnormed_vars = jnp.sum(n_longitudian_wavenumbers * sigmas_unnormed**2)

    # This is analogous to F₀ from [Palmer].
    # (normalization * sigmas_unnormed)² would sum to 1. The leading factor
    #   self._integrated_grf_variance * (1 - self.phi ** 2)
    # ensures that the AR(1) process has variance self._integrated_grf_variance.
    # We do not include the extra fator of 2 in the denominator. I do not know
    # why [Palmer] has this factor.
    normalization = jnp.sqrt(
        self._integrated_grf_variance()  # pyrefly: ignore[unsupported-operation]
        * self.one_minus_phi2
        / sum_unnormed_vars
    )

    # The factor of coords.horizontal.radius appears because our basis vectors
    # have L2 norm = radius.  See http://screen/9FYVXZ5cMHoGDZk
    return normalization * sigmas_unnormed / self.coords.horizontal.radius

  def unconditional_sample(self, rng: typing.PRNGKeyArray) -> RandomnessState:
    """Returns a randomly initialized state for the autoregressive process."""
    modal_shape = self.coords.horizontal.modal_shape
    rng, next_rng = jax.random.split(rng)
    if self.variance is None:
      return RandomnessState(
          core=jnp.zeros(modal_shape),  # pyrefly: ignore[unexpected-keyword]
          nodal_value=jnp.zeros(self.coords.horizontal.nodal_shape),  # pyrefly: ignore[unexpected-keyword]
          modal_value=jnp.zeros(modal_shape),  # pyrefly: ignore[unexpected-keyword]
          prng_key=next_rng,  # pyrefly: ignore[unexpected-keyword]
          prng_step=0,  # pyrefly: ignore[unexpected-keyword]
      )
    sigmas = self._sigma_array()
    weights = jnp.where(
        self.coords.horizontal.mask,
        jax.random.truncated_normal(rng, -self.clip, self.clip, modal_shape),
        jnp.zeros(modal_shape),
    )
    core = self.one_minus_phi2 ** (-0.5) * sigmas * weights
    return RandomnessState(
        core=core,  # pyrefly: ignore[unexpected-keyword]
        nodal_value=self.to_nodal_values(core),  # pyrefly: ignore[unexpected-keyword]
        modal_value=self.to_modal_values(core),  # pyrefly: ignore[unexpected-keyword]
        prng_key=next_rng,  # pyrefly: ignore[unexpected-keyword]
        prng_step=0,  # pyrefly: ignore[unexpected-keyword]
    )

  def advance(self, state: RandomnessState) -> RandomnessState:
    """Updates the CoreRandomState of a random gaussian field."""
    _validate_randomness_state(state)
    if self.variance is None:
      return RandomnessState(
          core=jnp.zeros_like(state.core),  # pyrefly: ignore[bad-argument-type, unexpected-keyword]
          nodal_value=jnp.zeros(self.coords.horizontal.nodal_shape),  # pyrefly: ignore[unexpected-keyword]
          modal_value=jnp.zeros(self.coords.horizontal.modal_shape),  # pyrefly: ignore[unexpected-keyword]
          prng_key=state.prng_key,  # pyrefly: ignore[unexpected-keyword]
          prng_step=state.prng_step + 1,  # pyrefly: ignore[unexpected-keyword, unsupported-operation]
      )
    modal_shape = self.coords.horizontal.modal_shape
    rng = _prng_key_for_current_advance_step(state)
    eta = jax.random.truncated_normal(rng, -self.clip, self.clip, modal_shape)  # pyrefly: ignore[bad-argument-type]
    next_core = state.core * self.phi + self._sigma_array() * jnp.where(  # pyrefly: ignore[unsupported-operation]
        self.coords.horizontal.mask, eta, jnp.zeros(modal_shape)
    )
    return RandomnessState(
        core=next_core,  # pyrefly: ignore[unexpected-keyword]
        nodal_value=self.to_nodal_values(next_core),  # pyrefly: ignore[unexpected-keyword]
        modal_value=self.to_modal_values(next_core),  # pyrefly: ignore[unexpected-keyword]
        prng_key=state.prng_key,  # pyrefly: ignore[unexpected-keyword]
        prng_step=state.prng_step + 1,  # pyrefly: ignore[unexpected-keyword, unsupported-operation]
    )

  @property
  def variance(self) -> Numeric | None:
    """An estimate of pointwise (in nodal space) variance of this random field.

    This random field is defined in spectral space, and has no precise
    pointwise variance quantity. However, it does have a precise integrated
    variance, which is used to define the field.

    If we assume the field is stationary (with higher spectral
    precision it is near stationary), then the average of this quantity is a
    good pointwise estimate. So define
      σ² := (1 / (4πR²)) ∫ Var(U(0, x))dx
          = (1 / (4πR²)) integrated_grf_variance

    Therefore the init parameter `variance` can be used to define
      `_integrated_grf_variance := variance * surf_area`
    and then `_integrated_grf_variance` is used to define this field. The result
    is a field with pointwise variance close to the init kwarg `variance`.

    Returns:
      Numeric estimate of pointwise variance.
    """
    return self._variance

  def _integrated_grf_variance(self) -> Numeric | None:
    """Integral of the GRF's variance over the earth's surface."""
    if self.variance is None:
      return self.variance
    return self.variance * self._surf_area

  def to_modal_values(self, core_state: CoreRandomState) -> typing.Array | None:
    """Returns the ready-for-use Gaussian random field."""
    return core_state

  def to_nodal_values(self, core_state: CoreRandomState) -> typing.Array | None:
    """Returns the ready-for-use Gaussian random field."""
    return self.coords.horizontal.to_nodal(core_state)


@gin.register
class GaussianRandomFieldModule(GaussianRandomField, hk.Module):
  """Module wrapper of GaussianRandomField with trainable parameters."""

  def __init__(
      self,
      coords: coordinate_systems.CoordinateSystem,
      dt: float,
      physics_specs: Any,
      aux_features: Any,
      initial_correlation_time: Union[Quantity, str] = gin.REQUIRED,
      initial_correlation_length: Union[Quantity, str] = gin.REQUIRED,
      initial_variance: Optional[Union[Quantity, str]] = gin.REQUIRED,
      variance_bound: Optional[Union[Quantity, str]] = gin.REQUIRED,
      tune_variance: bool = True,
      clip: float = 6.0,
      name: Optional[str] = None,
  ):
    """Constructs a GaussianRandomFieldModule.

    Stochastic parameters are initialized at provided `initial_*` values.
    This hk.Module can then be used to tune values.

    Args:
      coords: horizontal and vertical grid data.
      dt: nondimensionalized model time step.
      physics_specs: physical constants and definition of custom units.
      aux_features: additional static data.
      initial_correlation_time: timescale with units over which autoregressive
        process decorrelates. Typical values in NWP range from hours to days.
      initial_correlation_length: lengthscale with units over which random field
        is correlated. Typical values in NWP range from 500-2500 km.
      initial_variance: The average (over EarthSurface) variance of the random
        field. If None, this GRF always returns a zeros field and no RNGS will
        be drawn
      variance_bound: If provided, an upper bound on tuned variance values.
      tune_variance: Whether variance should be a tunable hk.parameter, or fixed
      clip: number of standard deviations at which to clip randomness to ensure
        numerical stability.
      name: Something no one cares about and we just use None.
    """
    # You must call hk.Module.__init__ before initializing this class.
    hk.Module.__init__(self, name=name)

    correlation_time_raw = hk.get_parameter(
        'correlation_time_raw', shape=(), init=hk.initializers.Constant(0.0)
    )
    correlation_length_raw = hk.get_parameter(
        'correlation_length_raw', shape=(), init=hk.initializers.Constant(0.0)
    )

    if tune_variance:
      variance_raw = hk.get_parameter(
          'variance_raw', shape=(), init=hk.initializers.Constant(0.0)
      )
    else:
      variance_raw = 0.0

    initial_variance = maybe_nondimensionalize(initial_variance, physics_specs)
    _assert_positive_or_none(initial_variance, 'initial_variance')

    if initial_variance is None:
      variance = None
    elif variance_bound in {None, 'None'}:  # Allow strings for gin.
      variance = convert_hk_param_to_positive_scalar(
          variance_raw, initial_variance  # pyrefly: ignore[bad-argument-type]
      )
    else:
      variance_bound = maybe_nondimensionalize(variance_bound, physics_specs)
      _assert_positive_or_none(variance_bound, 'variance_bound')
      _assert_positive_or_none(
          variance_bound - initial_variance, 'variance_bound - initial_variance'  # pyrefly: ignore[unsupported-operation]
      )
      variance = convert_hk_param_to_bounded_scalar(
          variance_raw,  # pyrefly: ignore[bad-argument-type]
          initial_variance,
          low=0.0,
          high=variance_bound,  # pyrefly: ignore[bad-argument-type]
      )

    # We call GaussianRandomFieldModule.__init__ rather than super().__init__
    # since we don't want to call hk.Module.__init__ twice... although doing
    # that didn't hurt anything.
    GaussianRandomField.__init__(
        self,
        coords=coords,
        dt=dt,
        physics_specs=physics_specs,
        aux_features=aux_features,
        correlation_time=convert_hk_param_to_positive_scalar(
            correlation_time_raw,
            maybe_nondimensionalize(initial_correlation_time, physics_specs),  # pyrefly: ignore[bad-argument-type]
        ),
        correlation_length=convert_hk_param_to_positive_scalar(
            correlation_length_raw,
            maybe_nondimensionalize(initial_correlation_length, physics_specs),  # pyrefly: ignore[bad-argument-type]
        ),
        variance=variance,
        clip=clip,
    )


################################################################################
# Single random fields that are derived from "stand on their own" fields.
################################################################################


@gin.register
class CenteredLognormalRandomField(GaussianRandomField):
  """A lognormal random field shifted to have mean zero."""

  @property
  def preferred_representation(self) -> PreferredRepresentation | None:
    return PreferredRepresentation.NODAL

  def _integrated_grf_variance(self) -> jax.Array | None:
    """Integrated variance of the associated GRF (not this Lognormal field)."""
    if self.variance is None:
      return None
    # If Z ~ Normal(μ, σ²), then X ~ exp(Z) has
    #  variance = (exp(σ²) - 1) exp(2μ + σ²).
    # We have centered this field, which involved setting μ = -σ² / 2.
    # => variance = exp(σ²) - 1,
    # and thus
    #  σ² = log(1 + variance)
    return jnp.log1p(self.variance) * self._surf_area

  def to_nodal_values(self, core_state: CoreRandomState) -> jax.Array:
    """Returns the ready-for-use Lognormal random field."""
    if self.variance is None:
      grf_variance = 0.0
    else:
      grf_variance = self._integrated_grf_variance() / self._surf_area  # pyrefly: ignore[unsupported-operation]
    # If Z ~ Normal(μ, σ²), then X ~ exp(Z) has mean exp(μ + σ²/2).
    # To ensure E[X] = 1, we must set μ = -σ²/2.
    x = self.coords.horizontal.to_nodal(core_state)  # ~ Normal(0, σ²)
    return jnp.expm1(x - grf_variance / 2)  # ~ Exp(Normal(-σ²/2, σ²)) - 1

  def to_modal_values(self, core_state: CoreRandomState) -> jax.Array:
    """Returns the ready-for-use Lognormal random field."""
    return self.coords.horizontal.to_modal(self.to_nodal_values(core_state))


@gin.register
class CenteredLognormalRandomFieldModule(
    CenteredLognormalRandomField, GaussianRandomFieldModule
):
  """A lognormal random hk.Module field shifted to have mean zero."""


################################################################################
# Fields made from many different fields.
################################################################################


@gin.register
class BatchGaussianRandomFieldModule(hk.Module):
  """Batch of independent GaussianRandomFieldModules.

  These GRFs are meant to be fed into a neural network as generic "signals".

  The state arrays have leading batch dim indexing independent GRFs.
  """

  def __init__(
      self,
      coords: coordinate_systems.CoordinateSystem,
      dt: float,
      physics_specs: Any,
      aux_features: Any,
      initial_correlation_times: Sequence[Quantity | str] = gin.REQUIRED,  # pyrefly: ignore[bad-function-definition]
      initial_correlation_lengths: Sequence[Quantity | str] = gin.REQUIRED,  # pyrefly: ignore[bad-function-definition]
      variances: Sequence[Quantity | str] = gin.REQUIRED,  # pyrefly: ignore[bad-function-definition]
      field_subset: Optional[Sequence[int]] = None,
      n_fixed_fields: Optional[int] = None,
      clip: float = 6.0,
      name: Optional[str] = None,
  ):
    """Constructs a BatchGaussianRandomFieldModule.

    Correlation scales are initialized to `initial_*` args and will be tuned
    by Haiku optimizers. Variance will be fixed.

    Args:
      coords: horizontal and vertical grid data.
      dt: nondimensionalized model time step.
      physics_specs: physical constants and definition of custom units.
      aux_features: additional static data.
      initial_correlation_times: timescales with units over which autoregressive
        process decorrelates. Typical values in NWP range from hours to days.
      initial_correlation_lengths: lengthscale with units over which random
        field is correlated. Typical values in NWP range from 500-2500 km.
      variances: The average (over EarthSurface) variance of the random field.
        These are fixed arrays (not tunable hk.parameters).
      field_subset: Optional nonempty subset of indices into initial parameters.
        Specifies which fields to construct. If None, use all fields. E.g.,
        field_subset=[0, 5] means form 3 GRFs from the 0th and 5th parameter
        values.
      n_fixed_fields: Number of fields that use fixed parameters. These will
        be fixed at the trailing `n_fixed_fields` initial correlations. The
        total number of fields is unchanged, since these fixed fields replace
        learnable fields.
      clip: number of standard deviations at which to clip randomness to ensure
        numerical stability.
      name: Name to show in xprof.
    """
    ## You must call hk.Module.__init__ before initializing this class.
    hk.Module.__init__(self, name=name)

    lengths = [
        len(initial_correlation_times),
        len(initial_correlation_lengths),
        len(variances),
    ]
    if len(set(lengths)) != 1:
      raise ValueError(f'Argument lengths differed: {lengths=}')
    n_fixed_fields = n_fixed_fields or 0

    # Get subset of args using `field_subset`
    if field_subset is not None:
      if not field_subset:
        raise ValueError(
            '`field_subset` must be `None` or non-empty sequence. Found'
            f' {field_subset=}'
        )
      get_subset = lambda seq: [seq[i] for i in field_subset]
      initial_correlation_lengths = get_subset(initial_correlation_lengths)
      initial_correlation_times = get_subset(initial_correlation_times)
      variances = get_subset(variances)

    logging.info(
        '[NGCM] Initializing BatchGaussianRandomFieldModule with'
        f' {initial_correlation_times=}, and {initial_correlation_lengths=},'
        f' and {variances=}'
    )

    # Get Haiku parameters.
    self._n_fields = len(variances)
    self._variances = jnp.array(
        [nondimensionalize(v, physics_specs) for v in variances]
    )

    initial_correlation_lengths = jnp.array([  # pyrefly: ignore[bad-assignment]
        nondimensionalize(l, physics_specs) for l in initial_correlation_lengths
    ])
    correlation_lengths_raw = hk.get_parameter(
        'correlation_lengths_raw',
        shape=(self.n_fields - n_fixed_fields,),
        init=hk.initializers.Constant(0.0),
    )
    if n_fixed_fields:
      correlation_lengths_raw = jnp.concatenate([
          correlation_lengths_raw, jnp.zeros([n_fixed_fields])])
    self._correlation_lengths = convert_hk_param_to_positive_scalar(
        correlation_lengths_raw, initial_correlation_lengths  # pyrefly: ignore[bad-argument-type]
    )

    initial_correlation_times = jnp.array(  # pyrefly: ignore[bad-assignment]
        [nondimensionalize(t, physics_specs) for t in initial_correlation_times]
    )
    correlation_times_raw = hk.get_parameter(
        'correlation_times_raw',
        shape=(self.n_fields - n_fixed_fields,),
        init=hk.initializers.Constant(0.0),
    )
    if n_fixed_fields:
      correlation_times_raw = jnp.concatenate([
          correlation_times_raw, jnp.zeros([n_fixed_fields])])
    self._correlation_times = convert_hk_param_to_positive_scalar(
        correlation_times_raw, initial_correlation_times  # pyrefly: ignore[bad-argument-type]
    )

    def make_rf(correlation_time, correlation_length, variance):
      return GaussianRandomField(
          coords=coords,
          dt=dt,
          physics_specs=physics_specs,
          aux_features=aux_features,
          correlation_time=correlation_time,
          correlation_length=correlation_length,
          variance=variance,
          clip=clip,
      )

    self._make_rf = make_rf

  @property
  def n_fields(self) -> int:
    return self._n_fields

  def unconditional_sample(self, rng: typing.PRNGKeyArray) -> RandomnessState:
    """Sample the batch GRFs unconditionally."""
    logging.info(
        '[NGCM] Calling BatchGaussianRandomFieldModule.unconditional_sample'
    )

    def _unconditional_sample_one_rf(
        key, correlation_time, correlation_length, variance
    ):
      rf = self._make_rf(correlation_time, correlation_length, variance)
      return rf.unconditional_sample(key)

    rngs = jax.random.split(rng, self.n_fields + 1)
    rngs, next_rng = rngs[:-1], rngs[-1]
    sample = jax.vmap(_unconditional_sample_one_rf)(
        rngs,
        self._correlation_times,
        self._correlation_lengths,
        self._variances,
    )
    # We have RNG keys and steps associated with each field from vmap, but
    # RandomnessState should only have a single (scalar) RNG key/step.
    return dataclasses.replace(sample, prng_key=next_rng, prng_step=0)

  def advance(self, state: RandomnessState) -> RandomnessState:
    """Updates the state of the batch of GRFs."""
    logging.info('[NGCM] Calling BatchGaussianRandomFieldModule.advance')

    def _advance_one_rf(state, correlation_time, correlation_length, variance):
      rf = self._make_rf(correlation_time, correlation_length, variance)
      return rf.advance(state)

    rng = _prng_key_for_current_advance_step(state)
    rngs = jax.random.split(rng, self.n_fields)  # pyrefly: ignore[bad-argument-type]
    steps = jnp.ones(self.n_fields, int) * state.prng_step  # pyrefly: ignore[unsupported-operation]
    advanced = jax.vmap(_advance_one_rf)(
        dataclasses.replace(state, prng_key=rngs, prng_step=steps),  # pyrefly: ignore[bad-specialization]
        self._correlation_times,
        self._correlation_lengths,
        self._variances,
    )
    return dataclasses.replace(
        advanced, prng_key=state.prng_key, prng_step=state.prng_step + 1  # pyrefly: ignore[unsupported-operation]
    )


@gin.register
class DictOfGaussianRandomFieldModules(hk.Module):
  """Dictionary of independent GaussianRandomFieldModules.

  These GRFs are meant to be fed into a neural network as generic "signals".
  """

  def __init__(
      self,
      coords: coordinate_systems.CoordinateSystem,
      dt: float,
      physics_specs: Any,
      aux_features: Any,
      initial_correlation_times: Sequence[Quantity | str] = gin.REQUIRED,  # pyrefly: ignore[bad-function-definition]
      initial_correlation_lengths: Sequence[Quantity | str] = gin.REQUIRED,  # pyrefly: ignore[bad-function-definition]
      variances: Sequence[Quantity | str] = gin.REQUIRED,  # pyrefly: ignore[bad-function-definition]
      field_names: Optional[Sequence[str]] = None,
      field_subset: Optional[Sequence[int]] = None,
      clip: float = 6.0,
      name: Optional[str] = None,
  ):
    """Constructs a DictOfGaussianRandomFieldModules.

    Correlation scales are initialized to `initial_*` args and will be tuned
    by Haiku optimizers. Variance will be fixed.

    Args:
      coords: horizontal and vertical grid data.
      dt: nondimensionalized model time step.
      physics_specs: physical constants and definition of custom units.
      aux_features: additional static data.
      initial_correlation_times: timescales with units over which autoregressive
        process decorrelates. Typical values in NWP range from hours to days.
      initial_correlation_lengths: lengthscale with units over which random
        field is correlated. Typical values in NWP range from 500-2500 km.
      variances: The average (over EarthSurface) variance of the random field.
        These are fixed arrays (not tunable hk.parameters).
      field_names: Optional names to give the fields. If None, the fields are
        named like "GRF0", "GRF1",...
      field_subset: Optional nonempty subset of indices into initial parameters.
        Specifies which fields to construct. If None, use all fields. E.g.,
        field_subset=[0, 5] means form 3 GRFs from the 0th and 5th parameter
        values.
      clip: number of standard deviations at which to clip randomness to ensure
        numerical stability.
      name: Name to show in xprof.
    """
    ## You must call hk.Module.__init__ before initializing this class.
    hk.Module.__init__(self, name=name)
    logging.info(
        '[NGCM] Initializing DictOfGaussianRandomFieldModules with'
        f' {initial_correlation_times=}, and {initial_correlation_lengths=},'
        f' and {variances=}'
    )

    field_names = field_names or [
        f'GRF{i}' for i in range(len(initial_correlation_times))
    ]

    lengths = [
        len(initial_correlation_times),
        len(initial_correlation_lengths),
        len(variances),
        len(field_names),
    ]
    if len(set(lengths)) != 1:
      raise ValueError(f'Argument lengths differed: {lengths=}')

    if field_subset is not None:
      if not field_subset:
        raise ValueError(
            '`field_subset` must be `None` or non-empty sequence. Found'
            f' {field_subset=}'
        )
      subset = lambda seq: [seq[i] for i in field_subset]
      field_names = subset(field_names)
      initial_correlation_lengths = subset(initial_correlation_lengths)
      initial_correlation_times = subset(initial_correlation_times)
      variances = subset(variances)

    self._field_names = tuple(field_names)

    self._random_fields = {}
    for tau, lam, var, field_name in zip(
        initial_correlation_times,
        initial_correlation_lengths,
        variances,
        self.field_names,
        strict=True,
    ):
      self._random_fields[field_name] = GaussianRandomFieldModule(
          coords,
          dt,
          physics_specs,
          aux_features,
          initial_correlation_time=tau,
          initial_correlation_length=lam,
          initial_variance=var,
          tune_variance=False,
          variance_bound=None,
          clip=clip,
          name=field_name,
      )

  @property
  def n_fields(self) -> int:
    return len(self._random_fields)

  @property
  def field_names(self) -> tuple[str, ...]:
    return self._field_names

  def unconditional_sample(self, rng: typing.PRNGKeyArray) -> RandomnessState:
    """Sample the random field unconditionally."""
    core = {}
    nodal_values = {}
    modal_values = {}
    *rngs, next_rng = jax.random.split(rng, self.n_fields + 1)
    for (name, rf), sample_key in zip(self._random_fields.items(), rngs):
      rvs = rf.unconditional_sample(sample_key)
      core[name] = rvs.core
      nodal_values[name] = rvs.nodal_value
      modal_values[name] = rvs.modal_value
    return RandomnessState(
        core=core,  # pyrefly: ignore[unexpected-keyword]
        nodal_value=nodal_values,  # pyrefly: ignore[unexpected-keyword]
        modal_value=modal_values,  # pyrefly: ignore[unexpected-keyword]
        prng_key=next_rng,  # pyrefly: ignore[unexpected-keyword]
        prng_step=0,  # pyrefly: ignore[unexpected-keyword]
    )

  def advance(self, state: RandomnessState) -> RandomnessState:
    """Updates the core state of a random field."""
    core = {}
    nodal_values = {}
    modal_values = {}
    rng = _prng_key_for_current_advance_step(state)
    rngs = jax.random.split(rng, self.n_fields)  # pyrefly: ignore[bad-argument-type]
    for (name, rf), sample_key in zip(self._random_fields.items(), rngs):
      # rvs is a RandomnessState.
      rvs = rf.advance(
          RandomnessState(state.core[name], prng_key=sample_key, prng_step=0)  # pyrefly: ignore[bad-argument-count, unexpected-keyword, unsupported-operation]
      )
      core[name] = rvs.core
      nodal_values[name] = rvs.nodal_value
      modal_values[name] = rvs.modal_value
    return RandomnessState(
        core=core,  # pyrefly: ignore[unexpected-keyword]
        nodal_value=nodal_values,  # pyrefly: ignore[unexpected-keyword]
        modal_value=modal_values,  # pyrefly: ignore[unexpected-keyword]
        prng_key=state.prng_key,  # pyrefly: ignore[unexpected-keyword]
        prng_step=state.prng_step + 1,  # pyrefly: ignore[unexpected-keyword, unsupported-operation]
    )


class SumOfRandomFields(RandomField):
  """RandomField that is the sum of multiple fields."""

  def __init__(self, random_fields: Sequence[RandomField]):
    self._random_fields = list(random_fields)  # Shallow copy
    coords = self._random_fields[0].coords
    if any(rf.coords != coords for rf in self._random_fields):
      raise ValueError(f'All fields must have the same coords. Found {coords=}')
    super().__init__(coords)

  @property
  def preferred_representation(self) -> PreferredRepresentation | None:
    n_nodal = sum(
        rf.preferred_representation == PreferredRepresentation.NODAL
        for rf in self._random_fields
    )
    n_modal = sum(
        rf.preferred_representation == PreferredRepresentation.MODAL
        for rf in self._random_fields
    )
    if n_nodal > n_modal:
      return PreferredRepresentation.NODAL
    elif n_nodal < n_modal:
      return PreferredRepresentation.MODAL
    return None

  def unconditional_sample(self, rng: typing.PRNGKeyArray) -> RandomnessState:
    """Sample the random field unconditionally."""
    rvs = []
    *rngs, next_rng = jax.random.split(rng, len(self._random_fields) + 1)
    for rf, sample_key in zip(self._random_fields, rngs, strict=True):
      rvs.append(rf.unconditional_sample(sample_key).core)
    return RandomnessState(
        core=rvs,  # pyrefly: ignore[unexpected-keyword]
        nodal_value=self.to_nodal_values(rvs),  # pyrefly: ignore[unexpected-keyword]
        modal_value=self.to_modal_values(rvs),  # pyrefly: ignore[unexpected-keyword]
        prng_key=next_rng,  # pyrefly: ignore[unexpected-keyword]
        prng_step=0,  # pyrefly: ignore[unexpected-keyword]
    )

  def advance(self, state: RandomnessState) -> RandomnessState:
    """Updates the core state of a random field."""
    rvs = []
    rng = _prng_key_for_current_advance_step(state)
    rngs = jax.random.split(rng, len(self._random_fields))  # pyrefly: ignore[bad-argument-type]
    for rf, s, k in zip(
        self._random_fields, state.core, rngs, strict=True  # pyrefly: ignore[bad-argument-type]
    ):
      rs = RandomnessState(s, prng_key=k, prng_step=state.prng_step)  # pyrefly: ignore[bad-argument-count, unexpected-keyword]
      rvs.append(rf.advance(rs).core)
    return RandomnessState(
        core=rvs,  # pyrefly: ignore[unexpected-keyword]
        nodal_value=self.to_nodal_values(rvs),  # pyrefly: ignore[unexpected-keyword]
        modal_value=self.to_modal_values(rvs),  # pyrefly: ignore[unexpected-keyword]
        prng_key=state.prng_key,  # pyrefly: ignore[unexpected-keyword]
        prng_step=state.prng_step + 1,  # pyrefly: ignore[unexpected-keyword, unsupported-operation]
    )

  def to_modal_values(self, core_state: CoreRandomState) -> typing.Array | None:
    """Finishes `core_state` by summing components."""
    modal_sum = 0.0
    nodal_sum = 0.0
    for rf, s in zip(self._random_fields, core_state, strict=True):
      if rf.preferred_representation == PreferredRepresentation.NODAL:
        nodal_sum += rf.to_nodal_values(s)  # pyrefly: ignore[unsupported-operation]
      elif rf.preferred_representation in [PreferredRepresentation.MODAL, None]:
        modal_sum += rf.to_modal_values(s)  # pyrefly: ignore[unsupported-operation]
    return modal_sum + self.coords.horizontal.to_modal(nodal_sum)

  def to_nodal_values(self, core_state: CoreRandomState) -> typing.Array | None:
    """Finishes `core_state` by summing components."""
    modal_sum = 0.0
    nodal_sum = 0.0
    for rf, s in zip(self._random_fields, core_state, strict=True):
      if rf.preferred_representation == PreferredRepresentation.MODAL:
        modal_sum += rf.to_modal_values(s)  # pyrefly: ignore[unsupported-operation]
      elif rf.preferred_representation in [PreferredRepresentation.NODAL, None]:
        nodal_sum += rf.to_nodal_values(s)  # pyrefly: ignore[unsupported-operation]
    return nodal_sum + self.coords.horizontal.to_nodal(modal_sum)


class SumOfGaussianLikeRandomFields(SumOfRandomFields, abc.ABC):
  """Base class for sum of independent Gaussian-like random fields."""

  def __init__(
      self,
      coords: coordinate_systems.CoordinateSystem,
      dt: float,
      physics_specs: Any,
      aux_features: Any,
      correlation_times: Sequence[
          Union[jax.Array, Quantity, str]
      ] = gin.REQUIRED,  # pyrefly: ignore[bad-function-definition]
      correlation_lengths: Sequence[
          Union[jax.Array, Quantity, str]
      ] = gin.REQUIRED,  # pyrefly: ignore[bad-function-definition]
      variances: Sequence[Union[jax.Array, Quantity, str]] = gin.REQUIRED,  # pyrefly: ignore[bad-function-definition]
      clip: float = 6.0,
  ):
    """Constructs a SumOfGaussianLikeRandomFields."""
    n_fields = len(correlation_times)
    variances = variances or [None] * n_fields
    random_fields = []
    logging.info(
        '[NGCM] Initializing SumOfGaussianLikeRandomFields with '
        f'{variances=}, {correlation_times=}, {correlation_lengths=}'
    )
    for tau, lam, var in zip(
        correlation_times, correlation_lengths, variances, strict=True
    ):
      random_fields.append(
          self.get_cls_constructor()(
              coords,
              dt,
              physics_specs,
              aux_features,
              correlation_time=tau,
              correlation_length=lam,
              variance=var,
              clip=clip,
          )
      )

    super().__init__(random_fields)

  @abc.abstractmethod
  def get_cls_constructor(self) -> type[GaussianRandomField]:
    """Gets class constructor that is initialized with Gaussian-like kwargs."""


class SumOfGaussianLikeRandomFieldsModule(
    SumOfRandomFields, hk.Module, abc.ABC
):
  """Base class for sums of independent Gaussian-like RandomFieldModules."""

  def __init__(
      self,
      coords: coordinate_systems.CoordinateSystem,
      dt: float,
      physics_specs: Any,
      aux_features: Any,
      initial_correlation_times: Sequence[Quantity | str] = gin.REQUIRED,  # pyrefly: ignore[bad-function-definition]
      initial_correlation_lengths: Sequence[Quantity | str] = gin.REQUIRED,  # pyrefly: ignore[bad-function-definition]
      initial_variances: Optional[Sequence[Quantity | str]] = gin.REQUIRED,  # pyrefly: ignore[bad-function-definition]
      variance_bounds: Optional[Sequence[Quantity | str]] = gin.REQUIRED,  # pyrefly: ignore[bad-function-definition]
      clip: float = 6.0,
      name: Optional[str] = None,
  ):
    """Constructs a SumOfGaussianLikeRandomFieldsModule."""
    # You must call hk.Module.__init__ before initializing this class.
    hk.Module.__init__(self, name=name)

    n_fields = len(initial_correlation_times)
    initial_variances = initial_variances or [None] * n_fields
    variance_bounds = variance_bounds or [None] * n_fields
    random_fields = []
    for tau, lam, var, bound in zip(
        initial_correlation_times,
        initial_correlation_lengths,
        initial_variances,
        variance_bounds,
        strict=True,
    ):
      random_fields.append(
          self.get_cls_constructor()(
              coords,
              dt,
              physics_specs,
              aux_features,
              initial_correlation_time=tau,
              initial_correlation_length=lam,
              initial_variance=var,
              variance_bound=bound,
              clip=clip,
              name=name,
          )
      )
    # We call SumOfRandomFields.__init__ rather than super().__init__
    # since we don't want to call hk.Module.__init__ twice... although doing
    # that didn't hurt anything.
    SumOfRandomFields.__init__(self, random_fields)

  @abc.abstractmethod
  def get_cls_constructor(self) -> type[GaussianRandomFieldModule]:
    """Gets class constructor that is initialized with Gaussian-like kwargs."""


@gin.register
class SumOfGaussianRandomFields(SumOfGaussianLikeRandomFields):

  def get_cls_constructor(self) -> type[GaussianRandomField]:
    return GaussianRandomField


@gin.register
class SumOfGaussianRandomFieldsModule(SumOfGaussianLikeRandomFieldsModule):
  """A sum of independent GaussianRandomFieldModules."""

  def get_cls_constructor(self) -> type[GaussianRandomFieldModule]:
    return GaussianRandomFieldModule


@gin.register
class SumOfCenteredLognormalRandomFields(SumOfGaussianLikeRandomFields):

  def get_cls_constructor(self) -> type[CenteredLognormalRandomField]:
    return CenteredLognormalRandomField


@gin.register
class SumOfCenteredLognormalRandomFieldsModule(
    SumOfGaussianLikeRandomFieldsModule
):
  """A sum of independent CenteredLognormalRandomFieldModules."""

  def get_cls_constructor(self) -> type[CenteredLognormalRandomFieldModule]:
    return CenteredLognormalRandomFieldModule


################################################################################
# Helper functions for creating fields.
################################################################################


def convert_hk_param_to_positive_scalar(
    param: jax.Array,
    initial_value: Numeric,
) -> jax.Array:
  """Converts [batch] scalar parameter to scalar value using Softplus."""
  return initial_value * make_positive_scalar(param)  # pyrefly: ignore[bad-return]


def _sigmoid(low: Numeric, high: Numeric, x: jax.Array) -> jax.Array:
  """Numerically stable sigmoid, adapted from tfp.bijectors.Sigmoid."""
  diff = high - low
  left = low + diff * jax.nn.sigmoid(x)
  right = high - diff * jax.nn.sigmoid(-x)
  return jnp.where(x < 0, left, right)


def _inv_sigmoid(low: Numeric, high: Numeric, x: jax.Array) -> jax.Array:
  """Inverse sigmoid, adapted from tfp.bijectors.Sigmoid."""
  return jnp.log(x - low) - jnp.log(high - x)


def convert_hk_param_to_bounded_scalar(
    param: jax.Array,
    initial_value: Numeric,
    low: Numeric,
    high: Numeric,
) -> jax.Array:
  """Converts a [batch] scalar parameter to scalar value using Sigmoid."""
  offset = _inv_sigmoid(low, high, initial_value)  # pyrefly: ignore[bad-argument-type]
  return _sigmoid(low, high, offset + param)


def nondimensionalize(
    x: Union[typing.Numeric, Quantity, str],
    physics_specs: Any,
) -> typing.Numeric:
  if isinstance(x, (Quantity, str)):  # pyrefly: ignore[invalid-argument]
    return physics_specs.nondimensionalize(Quantity(x))
  else:
    return x


def maybe_nondimensionalize(
    x: Optional[Union[typing.Numeric, Quantity, str]],
    physics_specs: Any,
) -> None | typing.Numeric:
  """Calls nondimensionalize on Quantity or str, otherwise passthrough."""
  if x == 'None':  # Allow strings for gin
    return None
  return nondimensionalize(x, physics_specs)


def _assert_positive_or_none(x: typing.Numeric | None, name: str) -> None:
  if x is None:
    return
  if x <= 0:
    raise ValueError(f'{name}={x} but should have been positive or None')