Upload folder using huggingface_hub

f4a39ee verified about 2 months ago

13.8 kB

# 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.
"""Defines `filtering` that aim to improve stability of integration."""

from typing import Any, Callable, Dict, Optional, Sequence, Union
from dinosaur import coordinate_systems
from dinosaur import filtering
from dinosaur import pytree_utils
from dinosaur import scales
from dinosaur import time_integration
from dinosaur import typing
import gin
import haiku as hk
import jax
import numpy as np


QuantityOrStr = Union[str, scales.Quantity]

StepFilterModule = Callable[..., typing.PyTreeStepFilterFn]
TransformModule = typing.TransformModule


#  =============================================================================
#  Step filters that attenuate modal components between time steps.
#  =============================================================================


@gin.register
class NoFilter(hk.Module):
  """Filter module that performs no filtering."""

  def __init__(
      self,
      coords: coordinate_systems.CoordinateSystem,
      dt: float,
      physics_specs: Any,
      aux_features: Dict[str, Any],
      name: Optional[str] = None,
  ):
    del coords, dt, physics_specs, aux_features  # unused.
    super().__init__(name=name)

  def __call__(
      self,
      u: typing.PyTreeState,
      u_next: typing.PyTreeState
  ) -> typing.PyTreeState:
    del u  # unused.
    return u_next


@gin.register
class ClipFilter(hk.Module):
  """Filter that clips highest total wavenumber in the next state."""

  def __init__(
      self,
      coords: coordinate_systems.CoordinateSystem,
      dt: float,
      physics_specs: Any,
      aux_features: Dict[str, Any],
      wavenumbers_to_clip: int = 1,
      name: Optional[str] = None,
  ):
    """See `time_integration.exponential_filter` for details."""
    del dt, physics_specs, aux_features  # unused.
    super().__init__(name=name)
    self.coords = coords
    self.wavenumbers_to_clip = wavenumbers_to_clip

  def __call__(
      self,
      u: typing.PyTreeState,
      u_next: typing.PyTreeState
  ) -> typing.PyTreeState:
    del u  # unused.
    return self.coords.horizontal.clip_wavenumbers(
        u_next, self.wavenumbers_to_clip
    )


@gin.register
class ExponentialLeapfrogFilter(hk.Module):
  """Filter that removes high frequency components from a spectral state."""

  def __init__(
      self,
      coords: coordinate_systems.CoordinateSystem,
      dt: float,
      physics_specs: Any,
      aux_features: Dict[str, Any],
      tau: QuantityOrStr = '0.010938',
      order: int = 18,
      cutoff: float = 0,
      name: Optional[str] = None,
  ):
    """See `time_integration.exponential_filter` for details."""
    del aux_features  # unused.
    super().__init__(name=name)
    tau = physics_specs.nondimensionalize(scales.Quantity(tau))
    self.filter_fn = time_integration.exponential_leapfrog_step_filter(
        coords.horizontal, dt, tau, order, cutoff)

  def __call__(
      self,
      u: typing.PyTreeState,
      u_next: typing.PyTreeState
  ) -> typing.PyTreeState:
    return self.filter_fn(u, u_next)


@gin.register
class ExponentialFilter(hk.Module):
  """Filter that removes high frequency components from a spectral state."""

  def __init__(
      self,
      coords: coordinate_systems.CoordinateSystem,
      dt: float,
      physics_specs: Any,
      aux_features: Dict[str, Any],
      tau: QuantityOrStr = '0.010938',
      order: int = 18,
      cutoff: float = 0,
      name: Optional[str] = None,
  ):
    """See `time_integration.exponential_step_filter` for details."""
    del aux_features  # unused.
    super().__init__(name=name)
    tau = physics_specs.nondimensionalize(scales.Quantity(tau))
    self.filter_fn = time_integration.exponential_step_filter(
        coords.horizontal, dt, tau, order, cutoff)

  def __call__(
      self,
      u: typing.PyTreeState,
      u_next: typing.PyTreeState
  ) -> typing.PyTreeState:
    return self.filter_fn(u, u_next)


@gin.register
class HorizontalDiffusionFilter(hk.Module):
  """Filter that applies implicit diffusion operator to a spectral state."""

  def __init__(
      self,
      coords: coordinate_systems.CoordinateSystem,
      dt: float,
      physics_specs: Any,
      aux_features: Dict[str, Any],
      tau: QuantityOrStr = '',
      order: int = 1,
      name: Optional[str] = None,
  ):
    """See `time_integration.horizontal_diffusion_filter` for details."""
    del aux_features  # unused.
    super().__init__(name=name)
    tau = physics_specs.nondimensionalize(scales.Quantity(tau))
    self.filter_fn = time_integration.horizontal_diffusion_step_filter(
        coords.horizontal, dt, tau, order)

  def __call__(
      self,
      u: typing.PyTreeState,
      u_next: typing.PyTreeState
  ) -> typing.PyTreeState:
    del u  # unused
    return self.filter_fn(u_next)  # pytype: disable=wrong-arg-count  # always-use-return-annotations


@gin.register
class RobertAsselinLeapfrogFilter(hk.Module):
  """Time smoothing filter."""

  def __init__(
      self,
      coords: coordinate_systems.CoordinateSystem,
      dt: float,
      physics_specs: Any,
      aux_features: Dict[str, Any],
      strength: float = 0.05,
      name: Optional[str] = None,
  ):
    """See `time_integration.robert_asselin_leapfrog_filter` for details."""
    del dt, coords, physics_specs, aux_features  # unused.
    super().__init__(name=name)
    self.filter_fn = time_integration.robert_asselin_leapfrog_filter(strength)

  def __call__(
      self,
      u: typing.PyTreeState,
      u_next: typing.PyTreeState
  ) -> typing.PyTreeState:
    return self.filter_fn(u, u_next)


@gin.register
class LearnedExponentialFilter(hk.Module):
  """Low pass filter with learned parameters."""

  def __init__(
      self,
      coords: coordinate_systems.CoordinateSystem,
      dt: float,
      physics_specs: Any,
      aux_features: Dict[str, Any],
      name: Optional[str] = None,
  ):
    del dt, physics_specs, aux_features  # unused.
    self.coords = coords
    self.a_init = hk.initializers.Constant(16)
    self.p_init = hk.initializers.Constant(18)
    self.c_init = hk.initializers.Constant(0)
    super().__init__(name=name)

  def __call__(
      self,
      u: typing.PyTreeState,
      u_next: typing.PyTreeState
  ) -> typing.PyTreeState:
    del u  # unused.
    a_logit = hk.get_parameter('attenuation_logit', shape=(), init=self.a_init)
    p_logit = hk.get_parameter('order_logit', shape=(), init=self.p_init)
    c_logit = hk.get_parameter('threshold_logit', shape=(), init=self.c_init)
    a = jax.nn.softplus(a_logit)
    p = jax.nn.softplus(p_logit)
    c = jax.nn.sigmoid(c_logit)
    filter_fn = filtering.exponential_filter(self.coords.horizontal, a, p, c)  # pytype: disable=wrong-arg-types  # jax-nn-types
    return filter_fn(u_next)


@gin.register
class SequentialStepFilter(hk.Module):
  """Filter module that combines multiple step filters applied sequentially."""

  def __init__(
      self,
      coords: coordinate_systems.CoordinateSystem,
      dt: float,
      physics_specs: Any,
      aux_features: Dict[str, Any],
      filter_modules: Sequence[StepFilterModule],
      name: Optional[str] = None,
  ):
    super().__init__(name=name)
    self.filter_fns = [module(coords, dt, physics_specs, aux_features)
                       for module in filter_modules]

  def __call__(
      self,
      u: typing.PyTreeState,
      u_next: typing.PyTreeState
  ) -> typing.PyTreeState:
    for filter_fn in self.filter_fns:
      u_next = filter_fn(u, u_next)
    return u_next


@gin.register
class LayeredStepFilter(hk.Module):
  """Filter decorator that uses varying time-scales at different levels."""

  def __init__(
      self,
      coords: coordinate_systems.CoordinateSystem,
      dt: float,
      physics_specs: Any,
      aux_features: Dict[str, Any],
      filter_module: Union[HorizontalDiffusionFilter, ExponentialFilter],
      tau_vals: Union[Sequence[float], np.ndarray],
      tau_units: QuantityOrStr,
      name: Optional[str] = None,
  ):
    super().__init__(name=name)
    tau = (np.asarray(tau_vals) * scales.Quantity(tau_units))
    tau = tau[:, np.newaxis, np.newaxis]  # add spatial axes.
    self.filter_fn = filter_module(
        coords, dt, physics_specs, aux_features, tau=tau)  # pyrefly: ignore[bad-argument-count, unexpected-keyword]

  def __call__(
      self,
      u: typing.PyTreeState,
      u_next: typing.PyTreeState
  ) -> typing.PyTreeState:
    return self.filter_fn(u, u_next)  # pyrefly: ignore[not-callable]


@gin.register
class MaskedFilter(hk.Module):
  """Filter that is only applied to a part of the state."""

  def __init__(
      self,
      coords: coordinate_systems.CoordinateSystem,
      dt: float,
      physics_specs: Any,
      aux_features: Dict[str, Any],
      filter_module: StepFilterModule,
      mask: typing.Pytree,
      name: Optional[str] = None,
  ):
    super().__init__(name=name)
    self.filter_fn = filter_module(coords, dt, physics_specs, aux_features)
    self.mask = mask

  def __call__(
      self,
      u: typing.PyTreeState,
      u_next: typing.PyTreeState
  ) -> typing.PyTreeState:
    mask = type(u_next)(**self.mask)  # convert to same structure.
    return jax.tree_util.tree_map(
        lambda x, y, b: self.filter_fn(x, y) if b else y,
        u, u_next, mask)


@gin.register
class FilterFromTransform(hk.Module):
  """Filter module that wraps a transform module."""

  def __init__(
      self,
      coords: coordinate_systems.CoordinateSystem,
      dt: float,
      physics_specs: Any,
      aux_features: Dict[str, Any],
      transform_module: TransformModule,
      name: Optional[str] = None,
  ):
    super().__init__(name=name)
    self.transform_fn = transform_module(
        coords, dt, physics_specs, aux_features)

  def __call__(
      self,
      u: typing.PyTreeState,
      u_next: typing.PyTreeState
  ) -> typing.PyTreeState:
    return self.transform_fn(u_next)


@gin.register
class FixGlobalMeanFilter(hk.Module):
  """Filter that removes the change in the global mean of certain keys."""

  def __init__(
      self,
      coords: coordinate_systems.CoordinateSystem,
      dt: float,
      physics_specs: Any,
      aux_features: Dict[str, Any],
      keys: tuple[str, ...] = ('log_surface_pressure',),
      name: Optional[str] = None,
  ):
    del aux_features  # unused.
    super().__init__(name=name)
    self.keys = keys

  def __call__(
      self, u: typing.PyTreeState, u_next: typing.PyTreeState
  ) -> typing.PyTreeState:
    u_dict, _ = pytree_utils.as_dict(u)
    u_dict, _ = pytree_utils.flatten_dict(u_dict)
    u_next_dict, from_dict_fn = pytree_utils.as_dict(u_next)
    u_next_dict, _ = pytree_utils.flatten_dict(u_next_dict)
    for key in self.keys:
      global_mean = u_dict[key][..., 0]
      u_next_dict[key] = u_next_dict[key].at[..., 0].set(global_mean)
    u_next_dict = pytree_utils.unflatten_dict(u_next_dict)
    return from_dict_fn(u_next_dict)


#  =============================================================================
#  Filters that act on modal variables without time-step context.
#  =============================================================================


@gin.register
class DataNoFilter(hk.Module):
  """Filter module that performs no filtering."""

  def __init__(
      self,
      coords: coordinate_systems.CoordinateSystem,
      dt: float,
      physics_specs: Any,
      aux_features: Dict[str, Any],
      name: Optional[str] = None,
  ):
    del coords, dt, physics_specs, aux_features  # unused.
    super().__init__(name=name)

  def __call__(self, inputs: typing.PyTreeState) -> typing.PyTreeState:
    return inputs


@gin.register
class DataExponentialFilter(hk.Module):
  """Filter that removes high frequency components from a modal data."""

  def __init__(
      self,
      coords: coordinate_systems.CoordinateSystem,
      dt: float,
      physics_specs: Any,
      aux_features: Dict[str, Any],
      attenuation: float = 16,
      order: int = 18,
      cutoff: float = 0,
      name: Optional[str] = None,
  ):
    """See `filtering.exponential_filter` for details."""
    del dt, physics_specs, aux_features  # unused.
    super().__init__(name=name)
    self.filter_fn = filtering.exponential_filter(
        coords.horizontal, attenuation, order, cutoff)

  def __call__(self, inputs: typing.PyTreeState) -> typing.PyTreeState:
    return self.filter_fn(inputs)


@gin.register
class PerVariableDataFilter(hk.Module):
  """Filter module that applies different filters for each variable."""

  def __init__(
      self,
      coords: coordinate_systems.CoordinateSystem,
      dt: float,
      physics_specs: Any,
      aux_features: Dict[str, Any],
      per_variable_filters: Dict[str, Any],
      name: Optional[str] = None,
  ):
    """See `filtering.exponential_filter` for details."""
    super().__init__(name=name)
    self.filter_fns = jax.tree_util.tree_map(
        lambda m: m(coords, dt, physics_specs, aux_features),
        per_variable_filters)

  def __call__(self, inputs: typing.PyTreeState) -> typing.PyTreeState:
    inputs_dict, from_dict_fn = pytree_utils.as_dict(inputs)
    return from_dict_fn(jax.tree_util.tree_map(
        lambda x, fn: fn(x), inputs_dict, self.filter_fns))

	# 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.
	"""Defines `filtering` that aim to improve stability of integration."""

	from typing import Any, Callable, Dict, Optional, Sequence, Union
	from dinosaur import coordinate_systems
	from dinosaur import filtering
	from dinosaur import pytree_utils
	from dinosaur import scales
	from dinosaur import time_integration
	from dinosaur import typing
	import gin
	import haiku as hk
	import jax
	import numpy as np


	QuantityOrStr = Union[str, scales.Quantity]

	StepFilterModule = Callable[..., typing.PyTreeStepFilterFn]
	TransformModule = typing.TransformModule


	# =============================================================================
	# Step filters that attenuate modal components between time steps.
	# =============================================================================


	@gin.register
	class NoFilter(hk.Module):
	"""Filter module that performs no filtering."""

	def __init__(
	self,
	coords: coordinate_systems.CoordinateSystem,
	dt: float,
	physics_specs: Any,
	aux_features: Dict[str, Any],
	name: Optional[str] = None,
	):
	del coords, dt, physics_specs, aux_features # unused.
	super().__init__(name=name)

	def __call__(
	self,
	u: typing.PyTreeState,
	u_next: typing.PyTreeState
	) -> typing.PyTreeState:
	del u # unused.
	return u_next


	@gin.register
	class ClipFilter(hk.Module):
	"""Filter that clips highest total wavenumber in the next state."""

	def __init__(
	self,
	coords: coordinate_systems.CoordinateSystem,
	dt: float,
	physics_specs: Any,
	aux_features: Dict[str, Any],
	wavenumbers_to_clip: int = 1,
	name: Optional[str] = None,
	):
	"""See `time_integration.exponential_filter` for details."""
	del dt, physics_specs, aux_features # unused.
	super().__init__(name=name)
	self.coords = coords
	self.wavenumbers_to_clip = wavenumbers_to_clip

	def __call__(
	self,
	u: typing.PyTreeState,
	u_next: typing.PyTreeState
	) -> typing.PyTreeState:
	del u # unused.
	return self.coords.horizontal.clip_wavenumbers(
	u_next, self.wavenumbers_to_clip
	)


	@gin.register
	class ExponentialLeapfrogFilter(hk.Module):
	"""Filter that removes high frequency components from a spectral state."""

	def __init__(
	self,
	coords: coordinate_systems.CoordinateSystem,
	dt: float,
	physics_specs: Any,
	aux_features: Dict[str, Any],
	tau: QuantityOrStr = '0.010938',
	order: int = 18,
	cutoff: float = 0,
	name: Optional[str] = None,
	):
	"""See `time_integration.exponential_filter` for details."""
	del aux_features # unused.
	super().__init__(name=name)
	tau = physics_specs.nondimensionalize(scales.Quantity(tau))
	self.filter_fn = time_integration.exponential_leapfrog_step_filter(
	coords.horizontal, dt, tau, order, cutoff)

	def __call__(
	self,
	u: typing.PyTreeState,
	u_next: typing.PyTreeState
	) -> typing.PyTreeState:
	return self.filter_fn(u, u_next)


	@gin.register
	class ExponentialFilter(hk.Module):
	"""Filter that removes high frequency components from a spectral state."""

	def __init__(
	self,
	coords: coordinate_systems.CoordinateSystem,
	dt: float,
	physics_specs: Any,
	aux_features: Dict[str, Any],
	tau: QuantityOrStr = '0.010938',
	order: int = 18,
	cutoff: float = 0,
	name: Optional[str] = None,
	):
	"""See `time_integration.exponential_step_filter` for details."""
	del aux_features # unused.
	super().__init__(name=name)
	tau = physics_specs.nondimensionalize(scales.Quantity(tau))
	self.filter_fn = time_integration.exponential_step_filter(
	coords.horizontal, dt, tau, order, cutoff)

	def __call__(
	self,
	u: typing.PyTreeState,
	u_next: typing.PyTreeState
	) -> typing.PyTreeState:
	return self.filter_fn(u, u_next)


	@gin.register
	class HorizontalDiffusionFilter(hk.Module):
	"""Filter that applies implicit diffusion operator to a spectral state."""

	def __init__(
	self,
	coords: coordinate_systems.CoordinateSystem,
	dt: float,
	physics_specs: Any,
	aux_features: Dict[str, Any],
	tau: QuantityOrStr = '',
	order: int = 1,
	name: Optional[str] = None,
	):
	"""See `time_integration.horizontal_diffusion_filter` for details."""
	del aux_features # unused.
	super().__init__(name=name)
	tau = physics_specs.nondimensionalize(scales.Quantity(tau))
	self.filter_fn = time_integration.horizontal_diffusion_step_filter(
	coords.horizontal, dt, tau, order)

	def __call__(
	self,
	u: typing.PyTreeState,
	u_next: typing.PyTreeState
	) -> typing.PyTreeState:
	del u # unused
	return self.filter_fn(u_next) # pytype: disable=wrong-arg-count # always-use-return-annotations


	@gin.register
	class RobertAsselinLeapfrogFilter(hk.Module):
	"""Time smoothing filter."""

	def __init__(
	self,
	coords: coordinate_systems.CoordinateSystem,
	dt: float,
	physics_specs: Any,
	aux_features: Dict[str, Any],
	strength: float = 0.05,
	name: Optional[str] = None,
	):
	"""See `time_integration.robert_asselin_leapfrog_filter` for details."""
	del dt, coords, physics_specs, aux_features # unused.
	super().__init__(name=name)
	self.filter_fn = time_integration.robert_asselin_leapfrog_filter(strength)

	def __call__(
	self,
	u: typing.PyTreeState,
	u_next: typing.PyTreeState
	) -> typing.PyTreeState:
	return self.filter_fn(u, u_next)


	@gin.register
	class LearnedExponentialFilter(hk.Module):
	"""Low pass filter with learned parameters."""

	def __init__(
	self,
	coords: coordinate_systems.CoordinateSystem,
	dt: float,
	physics_specs: Any,
	aux_features: Dict[str, Any],
	name: Optional[str] = None,
	):
	del dt, physics_specs, aux_features # unused.
	self.coords = coords
	self.a_init = hk.initializers.Constant(16)
	self.p_init = hk.initializers.Constant(18)
	self.c_init = hk.initializers.Constant(0)
	super().__init__(name=name)

	def __call__(
	self,
	u: typing.PyTreeState,
	u_next: typing.PyTreeState
	) -> typing.PyTreeState:
	del u # unused.
	a_logit = hk.get_parameter('attenuation_logit', shape=(), init=self.a_init)
	p_logit = hk.get_parameter('order_logit', shape=(), init=self.p_init)
	c_logit = hk.get_parameter('threshold_logit', shape=(), init=self.c_init)
	a = jax.nn.softplus(a_logit)
	p = jax.nn.softplus(p_logit)
	c = jax.nn.sigmoid(c_logit)
	filter_fn = filtering.exponential_filter(self.coords.horizontal, a, p, c) # pytype: disable=wrong-arg-types # jax-nn-types
	return filter_fn(u_next)


	@gin.register
	class SequentialStepFilter(hk.Module):
	"""Filter module that combines multiple step filters applied sequentially."""

	def __init__(
	self,
	coords: coordinate_systems.CoordinateSystem,
	dt: float,
	physics_specs: Any,
	aux_features: Dict[str, Any],
	filter_modules: Sequence[StepFilterModule],
	name: Optional[str] = None,
	):
	super().__init__(name=name)
	self.filter_fns = [module(coords, dt, physics_specs, aux_features)
	for module in filter_modules]

	def __call__(
	self,
	u: typing.PyTreeState,
	u_next: typing.PyTreeState
	) -> typing.PyTreeState:
	for filter_fn in self.filter_fns:
	u_next = filter_fn(u, u_next)
	return u_next


	@gin.register
	class LayeredStepFilter(hk.Module):
	"""Filter decorator that uses varying time-scales at different levels."""

	def __init__(
	self,
	coords: coordinate_systems.CoordinateSystem,
	dt: float,
	physics_specs: Any,
	aux_features: Dict[str, Any],
	filter_module: Union[HorizontalDiffusionFilter, ExponentialFilter],
	tau_vals: Union[Sequence[float], np.ndarray],
	tau_units: QuantityOrStr,
	name: Optional[str] = None,
	):
	super().__init__(name=name)
	tau = (np.asarray(tau_vals) * scales.Quantity(tau_units))
	tau = tau[:, np.newaxis, np.newaxis] # add spatial axes.
	self.filter_fn = filter_module(
	coords, dt, physics_specs, aux_features, tau=tau) # pyrefly: ignore[bad-argument-count, unexpected-keyword]

	def __call__(
	self,
	u: typing.PyTreeState,
	u_next: typing.PyTreeState
	) -> typing.PyTreeState:
	return self.filter_fn(u, u_next) # pyrefly: ignore[not-callable]


	@gin.register
	class MaskedFilter(hk.Module):
	"""Filter that is only applied to a part of the state."""

	def __init__(
	self,
	coords: coordinate_systems.CoordinateSystem,
	dt: float,
	physics_specs: Any,
	aux_features: Dict[str, Any],
	filter_module: StepFilterModule,
	mask: typing.Pytree,
	name: Optional[str] = None,
	):
	super().__init__(name=name)
	self.filter_fn = filter_module(coords, dt, physics_specs, aux_features)
	self.mask = mask

	def __call__(
	self,
	u: typing.PyTreeState,
	u_next: typing.PyTreeState
	) -> typing.PyTreeState:
	mask = type(u_next)(**self.mask) # convert to same structure.
	return jax.tree_util.tree_map(
	lambda x, y, b: self.filter_fn(x, y) if b else y,
	u, u_next, mask)


	@gin.register
	class FilterFromTransform(hk.Module):
	"""Filter module that wraps a transform module."""

	def __init__(
	self,
	coords: coordinate_systems.CoordinateSystem,
	dt: float,
	physics_specs: Any,
	aux_features: Dict[str, Any],
	transform_module: TransformModule,
	name: Optional[str] = None,
	):
	super().__init__(name=name)
	self.transform_fn = transform_module(
	coords, dt, physics_specs, aux_features)

	def __call__(
	self,
	u: typing.PyTreeState,
	u_next: typing.PyTreeState
	) -> typing.PyTreeState:
	return self.transform_fn(u_next)


	@gin.register
	class FixGlobalMeanFilter(hk.Module):
	"""Filter that removes the change in the global mean of certain keys."""

	def __init__(
	self,
	coords: coordinate_systems.CoordinateSystem,
	dt: float,
	physics_specs: Any,
	aux_features: Dict[str, Any],
	keys: tuple[str, ...] = ('log_surface_pressure',),
	name: Optional[str] = None,
	):
	del aux_features # unused.
	super().__init__(name=name)
	self.keys = keys

	def __call__(
	self, u: typing.PyTreeState, u_next: typing.PyTreeState
	) -> typing.PyTreeState:
	u_dict, _ = pytree_utils.as_dict(u)
	u_dict, _ = pytree_utils.flatten_dict(u_dict)
	u_next_dict, from_dict_fn = pytree_utils.as_dict(u_next)
	u_next_dict, _ = pytree_utils.flatten_dict(u_next_dict)
	for key in self.keys:
	global_mean = u_dict[key][..., 0]
	u_next_dict[key] = u_next_dict[key].at[..., 0].set(global_mean)
	u_next_dict = pytree_utils.unflatten_dict(u_next_dict)
	return from_dict_fn(u_next_dict)


	# =============================================================================
	# Filters that act on modal variables without time-step context.
	# =============================================================================


	@gin.register
	class DataNoFilter(hk.Module):
	"""Filter module that performs no filtering."""

	def __init__(
	self,
	coords: coordinate_systems.CoordinateSystem,
	dt: float,
	physics_specs: Any,
	aux_features: Dict[str, Any],
	name: Optional[str] = None,
	):
	del coords, dt, physics_specs, aux_features # unused.
	super().__init__(name=name)

	def __call__(self, inputs: typing.PyTreeState) -> typing.PyTreeState:
	return inputs


	@gin.register
	class DataExponentialFilter(hk.Module):
	"""Filter that removes high frequency components from a modal data."""

	def __init__(
	self,
	coords: coordinate_systems.CoordinateSystem,
	dt: float,
	physics_specs: Any,
	aux_features: Dict[str, Any],
	attenuation: float = 16,
	order: int = 18,
	cutoff: float = 0,
	name: Optional[str] = None,
	):
	"""See `filtering.exponential_filter` for details."""
	del dt, physics_specs, aux_features # unused.
	super().__init__(name=name)
	self.filter_fn = filtering.exponential_filter(
	coords.horizontal, attenuation, order, cutoff)

	def __call__(self, inputs: typing.PyTreeState) -> typing.PyTreeState:
	return self.filter_fn(inputs)


	@gin.register
	class PerVariableDataFilter(hk.Module):
	"""Filter module that applies different filters for each variable."""

	def __init__(
	self,
	coords: coordinate_systems.CoordinateSystem,
	dt: float,
	physics_specs: Any,
	aux_features: Dict[str, Any],
	per_variable_filters: Dict[str, Any],
	name: Optional[str] = None,
	):
	"""See `filtering.exponential_filter` for details."""
	super().__init__(name=name)
	self.filter_fns = jax.tree_util.tree_map(
	lambda m: m(coords, dt, physics_specs, aux_features),
	per_variable_filters)

	def __call__(self, inputs: typing.PyTreeState) -> typing.PyTreeState:
	inputs_dict, from_dict_fn = pytree_utils.as_dict(inputs)
	return from_dict_fn(jax.tree_util.tree_map(
	lambda x, fn: fn(x), inputs_dict, self.filter_fns))