jax-cfd

jax_cfd/ml/models_configs/implicit_diffusion_dns_config.gin

+# Implicit diffusion direct numerical simulation (DNS) configuration file.
+#
+# Can be used as base configuration for other models.
+# For an example of LES model see `smagorinsky_config.gin`.
+
+# Imports of modules to get access to their configurables.
+import jax_cfd.ml.advections
+import jax_cfd.ml.decoders
+import jax_cfd.ml.diffusions
+import jax_cfd.ml.encoders
+import jax_cfd.ml.equations
+import jax_cfd.ml.forcings
+import jax_cfd.ml.interpolations
+import jax_cfd.ml.model_builder
+import jax_cfd.ml.model_utils
+import jax_cfd.ml.optimizer_modules
+import jax_cfd.ml.pressures
+import jax_cfd.ml.towers
+import jax_cfd.ml.viscosities
+
+
+# Interpolations  (lax_wendroff + TVD suitable for 2nd order advection scheme)
+C_INTERPOLATION_MODULE = @interpolations.transformed
+U_INTERPOLATION_MODULE = @interpolations.linear
+transformed.base_interpolation_module = @interpolations.lax_wendroff
+transformed.transformation = @interpolations.tvd_limiter_transformation
+
+# Advection (advection is solved explicitly)
+CONVECTION_MODULE = @advections.self_advection
+advections.self_advection.advection_module = @advections.modular_advection
+modular_advection.c_interpolation_module = %C_INTERPOLATION_MODULE
+modular_advection.u_interpolation_module = %U_INTERPOLATION_MODULE
+
+# Pressure
+PRESSURE_MODULE = @pressures.fast_diagonalization
+
+# Diffusion (diffusion is solved implicitly)
+DIFFUSION_MODULE = @diffusions.solve_fast_diag
+
+# Model (note the implicit diffusion)
+NS_MODULE = @equations.modular_navier_stokes_model
+modular_navier_stokes_model.convection_module = %CONVECTION_MODULE
+modular_navier_stokes_model.pressure_module = %PRESSURE_MODULE
+modular_navier_stokes_model.equation_solver = @equations.implicit_diffusion_navier_stokes
+implicit_diffusion_navier_stokes.diffusion_module = %DIFFUSION_MODULE
+
+# Model specifications
+ModularStepModel.advance_module = %NS_MODULE
+ModularStepModel.encoder_module = @encoders.aligned_array_encoder
+ModularStepModel.decoder_module = @decoders.aligned_array_decoder
+model_builder.get_model_cls.model_cls = @ModularStepModel
+

jax_cfd/ml/models_configs/smagorinsky_config.gin

+# Configuration file specifying Smagorinsky LES model.
+
+include 'third_party/py/jax_cfd/ml/models_configs/implicit_diffusion_dns_config.gin'
+
+# Smagorinsky model is equivalent to an implicit diffusion DNS solver with an
+# addition of a closure term. We solve it implicitly, together with diffusion.
+SMAGORINSKY_CONSTANT = 0.2
+
+DIFFUSION_MODULE = @diffusions.implicit_evm_solve_with_diffusion
+implicit_evm_solve_with_diffusion.viscosity_module = @viscosities.eddy_viscosity_model
+eddy_viscosity_model.viscosity_model = @viscosities.smagorinsky_viscosity
+smagorinsky_viscosity.cs = %SMAGORINSKY_CONSTANT
+
+# Model
+equations.implicit_diffusion_navier_stokes.diffusion_module = %DIFFUSION_MODULE

jax_cfd/ml/networks.py

+"""Network modules that interface with numerical methods."""
+
+import functools
+import itertools
+from typing import Callable, Optional, Tuple
+
+import gin
+import haiku as hk
+import jax
+import jax.numpy as jnp
+from jax_cfd.base import array_utils
+from jax_cfd.base import boundaries
+from jax_cfd.base import finite_differences
+from jax_cfd.base import grids
+from jax_cfd.base import interpolation
+from jax_cfd.ml import physics_specifications
+from jax_cfd.ml import towers
+import numpy as np
+
+
+def _identity(grid, dt, physics_specs):
+  del grid, dt, physics_specs  # unused.
+  return lambda x: x
+
+
+@gin.register
+def split_to_aligned_field(
+    grid: grids.Grid,
+    dt: float,
+    physics_specs: physics_specifications.BasePhysicsSpecs,
+    network_offsets: Optional[Tuple[Tuple[float, float], ...]] = None,
+):
+  """Returns module that splits inputs along last axis into GridArrayVector."""
+  del dt  # unused.
+  if hasattr(physics_specs, "combo_offsets"):
+    data_offsets = physics_specs.combo_offsets()
+  else:
+    data_offsets = grid.cell_faces
+
+  if hasattr(physics_specs, "combo_boundaries"):
+    boundary_conditions = physics_specs.combo_boundaries()
+  else:
+    boundary_conditions = tuple(
+        boundaries.periodic_boundary_conditions(grid.ndim)
+        for _ in range(grid.ndim))
+  network_offsets = network_offsets or data_offsets
+  def process(inputs):
+    split_inputs = array_utils.split_axis(inputs, -1)
+    output = tuple(
+        grids.GridVariable(grids.GridArray(x, offset, grid), bc) for x, offset,
+        bc in zip(split_inputs, network_offsets, boundary_conditions))
+    output = tuple(
+        interpolation.linear(x, offset)
+        for x, offset in zip(output, data_offsets))
+    return output
+  return hk.to_module(process)()
+
+
+@gin.configurable()
+def interpolate_gridvar(
+    grid: grids.Grid,
+    dt: float,
+    physics_specs: physics_specifications.BasePhysicsSpecs,
+    final_offsets: Optional[Tuple[Tuple[float, float], ...]] = None,
+    process_fn: Optional[Callable] = lambda x: x,  # pylint: disable=g-bare-generic
+):
+  """Returns module that splits inputs along last axis into GridArrayVector."""
+  del dt  # unused.
+  if hasattr(physics_specs, "combo_offsets"):
+    data_offsets = physics_specs.combo_offsets()
+  else:
+    data_offsets = grid.cell_faces
+  final_offsets = final_offsets or data_offsets
+
+  def process(inputs):
+    inputs = process_fn(inputs)
+    inputs = tuple(
+        interpolation.linear(x, offset)
+        for x, offset in zip(inputs, final_offsets))
+    return inputs
+
+  return hk.to_module(process)()
+
+
+@gin.register
+def aligned_field_from_split_divergence(
+    grid: grids.Grid,
+    dt: float,
+    physics_specs: physics_specifications.BasePhysicsSpecs,
+):
+  """Returns module that splits inputs along last axis into GridArrayVector."""
+  del dt, physics_specs  # unused.
+
+  def _shift_offset(offset, axis):
+    return tuple(o + 0.5 if i == axis else o for i, o in enumerate(offset))
+
+  flux_offsets = tuple(
+      _shift_offset(o, i) for i in range(grid.ndim)  # pylint: disable=g-complex-comprehension
+      for o in grid.cell_faces
+  )
+
+  def _to_grid_variables(grid_arrays):
+    # TODO(dkochkov) make boundary conditions configurable.
+    bc = boundaries.periodic_boundary_conditions(grid.ndim)
+    return tuple(grids.GridVariable(array, bc) for array in grid_arrays)
+
+  def process(inputs):
+    split_inputs = array_utils.split_axis(inputs, -1)
+    split_inputs = tuple(grids.GridArray(x, o, grid)
+                         for x, o in zip(split_inputs, flux_offsets))
+    # below we combine `grid.ndim`-sized sequences of arrays into a tuples.
+    # we do that by iterating over a `grid.ndim`-sized zip of the same iterator.
+    # For example:
+    # a = [1, 2, 3, 4]
+    # tuple(zip(*([iter(a)] * 2))) >>> ((1, 2), (3, 4))
+    split_inputs = tuple(zip(*[iter(split_inputs)] * grid.ndim))
+    tensor_inputs = grids.GridArrayTensor(split_inputs)
+    # to compute divergence we need to convert fluxes to GridVariable class.
+    grid_array_field = tuple(
+        -finite_differences.divergence(_to_grid_variables(tensor_inputs[i, :]))
+        for i in range(grid.ndim))
+    # since divergence removes the boundary conditions, we add them back.
+    return _to_grid_variables(grid_array_field)
+
+  return hk.to_module(process)()
+
+
+@gin.register
+def stack_aligned_field_with_neighbors(
+    grid: grids.Grid,
+    dt: float,
+    physics_specs: physics_specifications.BasePhysicsSpecs,
+    n_neighbors: int = 1,
+):
+  """Returns a module that stacks input field with neighbors along channels."""
+  del dt, physics_specs  # unused.
+  shifts = [i for i in np.arange(-n_neighbors, n_neighbors + 1) if i != 0]
+  shifts_and_axis = list(itertools.product(shifts, np.arange(grid.ndim)))
+  shifts_and_axis.append([0, 0])
+
+  def process(inputs):
+    inputs = tuple(jnp.expand_dims(x.data, axis=-1) for x in inputs)
+    array = array_utils.concat_along_axis(jax.tree_util.leaves(inputs), axis=-1)
+    arrays = tuple(
+        jnp.roll(array, *shift_and_axis) for shift_and_axis in shifts_and_axis)
+    return array_utils.concat_along_axis(arrays, axis=-1)
+
+  return hk.to_module(process)()
+
+
+@gin.register
+def stack_aligned_field(
+    grid: grids.Grid,
+    dt: float,
+    physics_specs: physics_specifications.BasePhysicsSpecs,
+):
+  """Returns a module that stacks GridArrayVector along the last axis."""
+  del grid, dt, physics_specs  # unused.
+
+  def process(inputs):
+    inputs = tuple(jnp.expand_dims(x.data, axis=-1) for x in inputs)
+    return array_utils.concat_along_axis(jax.tree_util.leaves(inputs), axis=-1)
+
+  return hk.to_module(process)()
+
+
+@gin.configurable
+def tower_module(
+    grid: grids.Grid,
+    dt: float,
+    physics_specs: physics_specifications.BasePhysicsSpecs,
+    tower_factory: towers.TowerFactory,
+    pre_process_module: Callable = _identity,  # pylint: disable=g-bare-generic
+    post_process_module: Callable = _identity,  # pylint: disable=g-bare-generic
+    num_output_channels: Optional[int] = None,
+    name: Optional[str] = None,
+):
+  """Constructs tower module with configured number of output channels."""
+  pre_process = pre_process_module(grid, dt, physics_specs)
+  post_process = post_process_module(grid, dt, physics_specs)
+
+  def forward_pass(x):
+    x = pre_process(x)
+    if num_output_channels is None:
+      network = tower_factory(x.shape[-1], grid.ndim)
+    else:
+      network = tower_factory(num_output_channels, grid.ndim)
+    return post_process(network(x))
+
+  return hk.to_module(forward_pass)(name=name)
+
+
+@gin.configurable
+def velocity_corrector_network_w_boundaries(
+    grid: grids.Grid,
+    dt: float,
+    physics_specs: physics_specifications.BasePhysicsSpecs,
+    tower_factory: towers.TowerFactory,
+    network_offsets: Tuple[Tuple[float, ...], ...],
+    num_output_channels: int,
+    name: Optional[str] = None,
+    process_fn: Optional[Callable] = _identity,  # pylint: disable=g-bare-generic
+):
+  """Returns a module that computes corrections to the velocity field."""
+  pre_process = functools.partial(
+      interpolate_gridvar, final_offsets=network_offsets, process_fn=process_fn)
+  post_process = interpolate_gridvar
+  return tower_module(
+      grid=grid, dt=dt, physics_specs=physics_specs,
+      tower_factory=tower_factory, pre_process_module=pre_process,
+      post_process_module=post_process, num_output_channels=num_output_channels,
+      name=name)
+
+
+@gin.register
+def velocity_corrector_network(
+    grid: grids.Grid,
+    dt: float,
+    physics_specs: physics_specifications.BasePhysicsSpecs,
+    tower_factory: towers.TowerFactory,
+    name: Optional[str] = None,
+):
+  """Returns a module that computes corrections to the velocity field."""
+  pre_process_module = stack_aligned_field
+  post_process_module = split_to_aligned_field
+  return tower_module(
+      grid=grid, dt=dt, physics_specs=physics_specs,
+      tower_factory=tower_factory, pre_process_module=pre_process_module,
+      post_process_module=post_process_module, num_output_channels=grid.ndim,
+      name=name)
+
+
+@gin.register
+def flux_corrector_network(
+    grid: grids.Grid,
+    dt: float,
+    physics_specs: physics_specifications.BasePhysicsSpecs,
+    tower_factory: towers.TowerFactory,
+    pre_process_module: Callable = stack_aligned_field,  # pylint: disable=g-bare-generic
+    name: Optional[str] = None,
+):
+  """Returns a module that computes corrections to the velocity fluxes."""
+  post_process_module = aligned_field_from_split_divergence
+  num_output_channels = grid.ndim ** 2
+  return tower_module(
+      grid=grid, dt=dt, physics_specs=physics_specs,
+      tower_factory=tower_factory, pre_process_module=pre_process_module,
+      post_process_module=post_process_module,
+      num_output_channels=num_output_channels, name=name)

jax_cfd/ml/nonlinearities.py

+"""Registry of nonlinearities that can be used in neural networks."""
+
+import gin
+import jax
+import jax.numpy as jnp
+
+
+relu = gin.external_configurable(jax.nn.relu)
+tanh = gin.external_configurable(jnp.tanh)
+softplus = gin.external_configurable(jax.nn.softplus)
+swish = gin.external_configurable(jax.nn.swish)
+elu = gin.external_configurable(jax.nn.elu)
+gelu = gin.external_configurable(jax.nn.gelu)

jax_cfd/ml/optimizer_modules.py

+"""Configurable optimizers from JAX."""
+import gin
+from jax.example_libraries import optimizers
+
+
+@gin.configurable
+def optimizer(value):
+  return value
+
+
+gin.external_configurable(optimizers.adam)
+gin.external_configurable(optimizers.momentum)
+gin.external_configurable(optimizers.nesterov)
+
+gin.external_configurable(optimizers.exponential_decay)
+gin.external_configurable(optimizers.inverse_time_decay)
+gin.external_configurable(optimizers.polynomial_decay)
+gin.external_configurable(optimizers.piecewise_constant)

jax_cfd/ml/physics_configs/kolmogorov_forcing.gin

+# Physics configuration file Navier-Stokes system with Kolmogorov forcing.
+import jax_cfd.ml.forcings
+import jax_cfd.ml.physics_specifications
+
+
+FORCING_MODULE = @forcings.kolmogorov_forcing
+
+forcings.kolmogorov_forcing.scale = 1.0
+forcings.kolmogorov_forcing.wavenumber = 4
+forcings.kolmogorov_forcing.linear_coefficient = -0.1
+
+DENSITY = 1.
+VISCOSITY = 0.001
+physics_specifications.NavierStokesPhysicsSpecs.density = %DENSITY
+physics_specifications.NavierStokesPhysicsSpecs.viscosity = %VISCOSITY
+physics_specifications.NavierStokesPhysicsSpecs.forcing_module = %FORCING_MODULE
+
+physics_specifications.get_physics_specs.physics_specs_cls = @NavierStokesPhysicsSpecs

jax_cfd/ml/physics_configs/kuramoto_sivashinsky.gin

+# Physics configuration file for kuramoto sivashinsky equation.
+import jax_cfd.ml.physics_specifications
+
+physics_specifications.KsPhysicsSpecs.forcing_module = None
+physics_specifications.get_physics_specs.physics_specs_cls = @KsPhysicsSpecs

jax_cfd/ml/physics_configs/linear_forcing.gin

+# Physics configuration file Navier-Stokes system with linear forcing.
+import jax_cfd.ml.forcings
+import jax_cfd.ml.physics_specifications
+
+FORCING_MODULE = @forcings.linear_forcing
+
+FORCING_SCALE = 0.05
+forcings.linear_forcing.scale = %FORCING_SCALE
+
+
+DENSITY = 1.
+VISCOSITY = 0.000665
+physics_specifications.NavierStokesPhysicsSpecs.density = %DENSITY
+physics_specifications.NavierStokesPhysicsSpecs.viscosity = %VISCOSITY
+physics_specifications.NavierStokesPhysicsSpecs.forcing_module = %FORCING_MODULE
+
+physics_specifications.get_physics_specs.physics_specs_cls = @NavierStokesPhysicsSpecs

jax_cfd/ml/physics_configs/no_forcing.gin

+# Physics configuration file Navier-Stokes system with Kolmogorov forcing.
+import jax_cfd.ml.forcings
+import jax_cfd.ml.physics_specifications
+
+FORCING_MODULE = @forcings.no_forcing
+DENSITY = 1.
+VISCOSITY = 0.001
+physics_specifications.NavierStokesPhysicsSpecs.density = %DENSITY
+physics_specifications.NavierStokesPhysicsSpecs.viscosity = %VISCOSITY
+physics_specifications.NavierStokesPhysicsSpecs.forcing_module = %FORCING_MODULE
+
+physics_specifications.get_physics_specs.physics_specs_cls = @NavierStokesPhysicsSpecs

jax_cfd/ml/physics_configs/taylor_green_forcing.gin

+# Physics configuration file Navier-Stokes system with Taylor-green forcing.
+import jax_cfd.ml.forcings
+import jax_cfd.ml.physics_specifications
+
+
+FORCING_MODULE = @forcings.taylor_green_forcing
+
+forcings.taylor_green_forcing.scale = 0.05
+forcings.taylor_green_forcing.wavenumber = 3
+forcings.taylor_green_forcing.linear_coefficient = -0.002
+
+DENSITY = 1.
+VISCOSITY = 0.001
+physics_specifications.NavierStokesPhysicsSpecs.density = %DENSITY
+physics_specifications.NavierStokesPhysicsSpecs.viscosity = %VISCOSITY
+physics_specifications.NavierStokesPhysicsSpecs.forcing_module = %FORCING_MODULE
+
+physics_specifications.get_physics_specs.physics_specs_cls = @NavierStokesPhysicsSpecs

jax_cfd/ml/physics_specifications.py

+"""Modules with PhysicsSpecifications for various equations.
+
+To ensure that all components of the pipeline obtain the expected PhysicsSpecs
+all modules (except specializing on a particular equation) must instantiate
+PhysicsSpecs objects using `get_physics_specs`, which should be
+configured appropriately.
+"""
+
+import dataclasses
+from typing import Optional
+
+import gin
+from jax_cfd.ml import forcings
+
+
+ForcingModule = forcings.ForcingModule
+
+
+@gin.configurable
+def get_physics_specs(physics_specs_cls=gin.REQUIRED):
+  """Returns an instance of `physics_specs_cls`, configured by gin."""
+  return physics_specs_cls()
+
+
+@gin.register
+@dataclasses.dataclass
+class BasePhysicsSpecs:
+  """Base class for keeping physical parameters and forcing module."""
+  forcing_module: Optional[ForcingModule]
+
+
+@gin.register
+@dataclasses.dataclass
+class KsPhysicsSpecs(BasePhysicsSpecs):
+  """Configurable physical parameters for Kuramoto-Sivashinsky models."""
+
+
+@gin.register
+@dataclasses.dataclass
+class NavierStokesPhysicsSpecs(BasePhysicsSpecs):
+  """Configurable physical parameters and modules for Navier-Stokes models."""
+  density: float
+  viscosity: float
+
+
+@gin.configurable
+@dataclasses.dataclass
+class SpectralNavierStokesPhysicsSpecs(BasePhysicsSpecs):
+  viscosity: float
+  drag: float
+  smooth: bool

jax_cfd/ml/pressures.py

+"""Models for pressure solvers.
+
+All modules are functions that return `pressure_solve` method that has the same
+signature as baseline methods e.g. `pressure.solve_fast_diag`.
+"""
+import functools
+from typing import Callable, Optional
+
+import gin
+
+from jax_cfd.base import grids
+from jax_cfd.base import pressure
+
+
+GridArray = grids.GridArray
+GridVariable = grids.GridVariable
+GridVariableVector = grids.GridVariableVector
+PressureSolveFn = Callable[
+    [GridVariableVector, Optional[GridVariable]], GridArray]
+PressureModule = Callable[..., PressureSolveFn]
+
+
+@gin.register
+def fast_diagonalization(grid, dt, physics_specs):
+  del grid, dt, physics_specs  # unused.
+  return pressure.solve_fast_diag
+
+
+@gin.register
+def conjugate_gradient(grid, dt, physics_specs, atol=1e-5, maxiter=32):
+  del grid, dt, physics_specs  # unused.
+  return functools.partial(pressure.solve_cg, atol=atol, maxiter=maxiter)

jax_cfd/ml/tiling.py

+"""Utilities for spatial tiling of periodic convolutions into batch dimensions.
+
+``layout`` tuple indicates how the corresponding spatial dimensions are layed
+out in space. In 2D:
+- `(1, 1)` indicates no tiling.
+- `(4, 2)` indicates 4 x-tiles and 2 y-tiles
+- `(16, 8)` indicates 16 x-tiles and 8 y-tiles
+
+Tiling is helpful for getting the highest performance convolutions on TPU. Per
+the TPU performance guide [1], batch dimensions on TPUs are tiled to multiples
+of 8 or 128. Thus the product of all elements in `layout` should typically be
+either 8 or 128.
+
+[1] https://cloud.google.com/tpu/docs/performance-guide.
+"""
+import functools
+from typing import Callable, Sequence, Tuple
+
+import einops
+import jax
+from jax import lax
+import jax.numpy as jnp
+
+
+Array = jnp.ndarray
+
+
+def _prod(xs):
+  # backport of math.prod() from Python 3.8+
+  result = 1
+  for x in xs:
+    result *= x
+  return result
+
+
+def _verify_layout(array, layout):
+  if array.ndim != len(layout) + 2 or array.shape[0] != _prod(layout):
+    raise ValueError(
+        f"array shape does not match layout: {array.shape} vs {layout}")
+
+
+def _layout_to_dict(layout):
+  return dict(zip(["bx", "by", "bz"], layout))
+
+
+def _tile_roll(array, layout, shift, axis):
+  """Roll along the "tiled" dimension."""
+  _verify_layout(array, layout)
+  sizes = _layout_to_dict(layout)
+  if len(layout) == 1:
+    array = jnp.roll(array, shift, axis=axis)
+  elif len(layout) == 2:
+    array = einops.rearrange(array, "(bx by) ... -> bx by ...", **sizes)
+    array = jnp.roll(array, shift, axis=axis)
+    array = einops.rearrange(array, "bx by ... -> (bx by) ...", **sizes)
+  elif len(layout) == 3:
+    array = einops.rearrange(array, "(bx by bz) ... -> bx by bz ...", **sizes)
+    array = jnp.roll(array, shift, axis=axis)
+    array = einops.rearrange(array, "bx by bz ... -> (bx by bz) ...", **sizes)
+  else:
+    raise NotImplementedError
+  return array
+
+
+def _halo_pad_1d(array, layout, axis, padding=(1, 1)):
+  """Pad for halo-exchange along a single array dimension."""
+  pad_left, pad_right = padding
+  spatial_axis = axis + 1
+  pieces = []
+
+  if pad_left:
+    # Note: importantly, dynamic_slice_in_dim raises an error for out of bounds
+    # access, which catches the edge case where a single array is insufficient
+    # padding.
+    start = array.shape[spatial_axis] - pad_left
+    input_right = lax.dynamic_slice_in_dim(array, start, pad_left, spatial_axis)
+    output_left = _tile_roll(input_right, layout, shift=+1, axis=axis)
+    pieces.append(output_left)
+
+  pieces.append(array)
+
+  if pad_right:
+    start = 0
+    input_left = lax.dynamic_slice_in_dim(array, start, pad_right, spatial_axis)
+    output_right = _tile_roll(input_left, layout, shift=-1, axis=axis)
+    pieces.append(output_right)
+
+  return jnp.concatenate(pieces, axis=spatial_axis)
+
+
+@functools.partial(jax.jit, static_argnums=(1, 2,))
+def _halo_exchange_pad(array: Array, layout: Tuple[int, ...],
+                       padding: Tuple[Tuple[int, int]]) -> Array:
+  """Pad with halo-exchange in N-dimensions."""
+  _verify_layout(array, layout)
+  if len(layout) != len(padding):
+    raise ValueError(f"invalid padding: {padding}")
+  out = array
+  for axis, pad in enumerate(padding):
+    out = _halo_pad_1d(out, layout, axis, pad)
+  return out
+
+
+def halo_exchange_pad(
+    array: Array,
+    layout: Tuple[int, ...],
+    padding: Sequence[Tuple[int, int]],
+) -> Array:
+  """Pad with halo-exchange in N-dimensions."""
+  return _halo_exchange_pad(
+      array, layout,
+      tuple(map(tuple, padding)))
+
+
+@functools.partial(jax.jit, static_argnums=(1,))
+def space_to_batch(array: Array, layout: Tuple[int, ...]) -> Array:
+  """Rearrange from space to batch dimensions."""
+  sizes = _layout_to_dict(layout)
+  if len(layout) == 1:
+    path = "(bx x) c -> (bx) x c"
+  elif len(layout) == 2:
+    path = "(bx x) (by y) c -> (bx by) x y c"
+  elif len(layout) == 3:
+    path = "(bx x) (by y) (bz z) c -> (bx by bz) x y z c"
+  else:
+    raise NotImplementedError
+  return einops.rearrange(array, path, **sizes)
+
+
+@functools.partial(jax.jit, static_argnums=(1,))
+def batch_to_space(array: Array, layout: Tuple[int, ...]) -> Array:
+  """Rearrange from batch to space dimensions."""
+  sizes = _layout_to_dict(layout)
+  if len(layout) == 1:
+    path = "(bx) x c -> (bx x) c"
+  elif len(layout) == 2:
+    path = "(bx by) x y c -> (bx x) (by y) c"
+  elif len(layout) == 3:
+    path = "(bx by bz) x y z c-> (bx x) (by y) (bz z) c"
+  else:
+    raise NotImplementedError
+  return einops.rearrange(array, path, **sizes)
+
+
+def apply_convolution(
+    conv: Callable[[Array], Array],
+    inputs: Array,
+    layout: Tuple[int, ...],
+    padding: Sequence[Tuple[int, ...]],
+) -> Array:
+  """Apply a valid convolution with tiling and periodic boundary conditions.
+
+  Args:
+    conv: function that calculates a convolution with valid boundary conditions
+      when applied to an array of shape [batch, [spatial dims], channel].
+    inputs: array of shape [[spatial dims], channel].
+    layout: tiling layout for implementing the operation.
+    padding: amount of periodic padding to add before and after each spatial
+      dimension.
+
+  Returns:
+    Convolved array.
+  """
+  if layout is None:
+    # TODO(shoyer): replace this with some sensible heuristic
+    layout = (1,) * len(padding)
+  tiled = space_to_batch(inputs, layout)
+  padded = halo_exchange_pad(tiled, layout, padding)
+  convolved = conv(padded)
+  output = batch_to_space(convolved, layout)
+  return output

jax_cfd/ml/time_integrators.py

+"""Methods for time integration of first order differential equations."""
+import gin
+import haiku as hk
+import jax
+
+
+# TODO(dkochkov) Include other integrators such as DP; RK methods;
+# TODO(dkochkov) Add option to include input as in funcutils.trajectory.
+
+
+@gin.register
+def euler_integrator(
+    derivative_module,
+    initial_state,
+    dt,
+    num_steps,
+):
+  """Integrates ode defined by `derivative_module` using euler method.
+
+  Args:
+    derivative_module: hk.Module that computes time derivative.
+    initial_state: initial state for time integration.
+    dt: time step.
+    num_steps: number time steps `dt` to integrate for.
+
+  Returns:
+   final state at time `t + num_steps * dt` and `dt` spaced trajectory.
+  """
+  def _single_step(state, _):
+    deriv = derivative_module(state)
+    next_state = jax.tree_util.tree_map(lambda x, dxdt: x + dt * dxdt, state, deriv)
+    return next_state, next_state
+
+  return hk.scan(_single_step, initial_state, None, num_steps)

jax_cfd/ml/towers.py

+"""Definitions of towers (neural networks based on multioke CNN layers)."""
+
+import functools
+from typing import Any, Callable, List, Optional, Tuple, Union
+import gin
+import jax
+import haiku as hk
+
+from jax_cfd.ml import layers
+from jax_cfd.ml import nonlinearities
+
+Array = layers.Array
+ConvModule = Callable[..., Any]
+ScaleFn = Callable[[Array, List[int]], Array]
+TowerFactory = Callable[..., Any]
+
+
+PERIODIC_CONV_MODULES = {
+    1: layers.PeriodicConv1D,
+    2: layers.PeriodicConv2D,
+    3: layers.PeriodicConv3D}
+
+PERIODIC_CONV_TRANSPOSE_MODULES = {
+    1: layers.PeriodicConvTranspose1D,
+    2: layers.PeriodicConvTranspose2D,
+    3: layers.PeriodicConvTranspose3D}
+
+
+@gin.register
+def periodic_convolution(
+    output_channels: int,
+    kernel_shape: Tuple[int, ...],
+    ndim: int,
+    **kwargs
+):
+  """Returns PeriodicConv module with specified parameters."""
+  return PERIODIC_CONV_MODULES[ndim](output_channels, kernel_shape, **kwargs)
+
+
+@gin.register
+def periodic_transpose_convolution(
+    output_channels: int,
+    kernel_shape: Tuple[int, ...],
+    ndim: int,
+    rate: Optional[int] = None,
+    **kwargs
+):
+  """Returns PeriodicConvTranspose module with specified parameters."""
+  if rate is not None and rate != 1:
+    raise ValueError('transpose convolutions do not support dilation rate')
+  return PERIODIC_CONV_TRANSPOSE_MODULES[ndim](
+      output_channels, kernel_shape, **kwargs)
+
+
+@gin.register
+def mirror_convolution(
+    output_channels: int,
+    kernel_shape: Tuple[int, ...],
+    ndim: int,
+    **kwargs
+):
+  """Returns MirrorConv2D module with specified parameters."""
+  del ndim
+  return layers.MirrorConv2D(
+      output_channels, kernel_shape, **kwargs)
+
+
+@gin.register
+def fixed_scale(inputs: Array,
+                axes: Tuple[int, ...],
+                rescaled_one: float = gin.REQUIRED) -> Array:
+  """Linearly scales `inputs` such that `1` maps to `rescaled_one`."""
+  del axes  # unused.
+  return inputs * rescaled_one
+
+
+@gin.register
+def fixed_scale_gridvar(
+    inputs: Array,
+    axes: Tuple[int, ...],
+    rescaled_one: float = gin.REQUIRED
+) ->Array:
+  """Linearly scales `inputs` such that `1` maps to `rescaled_one`."""
+  del axes  # unused.
+  return tuple(x.bc.impose_bc(x.array * rescaled_one) for x in inputs)  # pytype: disable=bad-return-type  # jax-devicearray
+
+
+@gin.register
+def scale_to_range(
+    inputs: Array,
+    axes: Tuple[int, ...],
+    min_value: float = gin.REQUIRED,
+    max_value: float = gin.REQUIRED,
+) -> Array:
+  """Dynamically scales `inputs` to be in `[min_value, max_value]` range.
+
+  This scaling function represents a shift and scale transform that forces every
+  `axes` slice of `inputs` to be exactly in range `[min_value, max_value]`.
+  For details see `layers.rescale_to_range`.
+
+  Args:
+    inputs: array values to be rescaled.
+    axes: tuple of ints representing axes over which the scaling is calculated.
+    min_value: minimum value to appear in the rescaled values.
+    max_value: maximum value to appear in the rescaled values.
+
+  Returns:
+    `inputs` scale to `[min_value, max_value]` range on every `axes` slice.
+  """
+  return layers.rescale_to_range(inputs, min_value, max_value, axes)
+
+
+@gin.register
+class MlpTowerFactory(hk.Module):
+  """Tower that applies shared MLP to inputs over spatial dimensions."""
+
+  def __init__(
+      self,
+      output_size: int,
+      ndim: int,
+      num_hidden_units: int,
+      num_hidden_layers: int,
+      nonlinearity: Callable[[Array], Array] = nonlinearities.relu,
+      inputs_scale_fn: ScaleFn = lambda x, axes: x,
+      output_scale_fn: ScaleFn = lambda x, axes: x,
+      name: Optional[str] = 'mlp_tower_factory',
+  ):
+    super().__init__(name=name)
+    output_sizes = [num_hidden_units] * num_hidden_layers + [output_size]
+    mlp_net = hk.nets.MLP(output_sizes, activation=nonlinearity)
+    for _ in range(ndim):
+      mlp_net = hk.vmap(mlp_net, split_rng=False)
+    ndim_axes = list(range(ndim))
+    self.inputs_scale_fn = functools.partial(inputs_scale_fn, axes=ndim_axes)
+    self.output_scale_fn = functools.partial(output_scale_fn, axes=ndim_axes)
+    self.mlp_tower = mlp_net
+
+  def __call__(self, inputs):
+    """Applied Mlp tower to `inputs`."""
+    return self.output_scale_fn(self.mlp_tower(self.inputs_scale_fn(inputs)))
+
+
+@gin.register
+def forward_tower_factory(
+    num_output_channels: int,
+    ndim: int,
+    num_hidden_channels: int = 16,
+    kernel_size: int = 3,
+    num_hidden_layers: int = 2,
+    rates: Union[int, Tuple[int, ...]] = 1,
+    strides: Union[int, Tuple[int, ...]] = 1,
+    output_kernel_size: int = 3,
+    output_dilation_rate: int = 1,
+    output_stride: int = 1,
+    conv_module: ConvModule = periodic_convolution,
+    nonlinearity: Callable[[Array], Array] = nonlinearities.relu,
+    inputs_scale_fn: ScaleFn = lambda x, axes: x,
+    output_scale_fn: ScaleFn = lambda x, axes: x,
+    name: Optional[str] = 'forward_cnn_tower',
+):
+  """Constructs parametrized feed-forward CNN tower.
+
+  Constructs CNN tower parametrized by fixed number of channels in hidden layers
+  and fixed square kernels.
+
+  Args:
+    num_output_channels: number of channels in the output layer.
+    ndim: number of spatial dimensions to expect in inputs to the network.
+    num_hidden_channels: number of channels to use in hidden layers.
+    kernel_size: size of the kernel to use along every dimension.
+    num_hidden_layers: number of hidden layers to construct in the tower.
+    rates: dilation rate(s) of the hidden layers.
+    strides: strides to use. Must be `int` or same a `num_hidden_layers`.
+    output_kernel_size: size of the output kernel to use along every dimension.
+    output_dilation_rate: dilation_rate of the output layer.
+    output_stride: stride of the final convolution.
+    conv_module: convolution module to use. Must accept
+      (output channels, kernel shape and ndim).
+    nonlinearity: nonlinearity function to apply between hidden layers.
+    inputs_scale_fn: scaling function to be applied to the inputs of the tower.
+      Must take inputs as argument and return an `Array` of the same shape.
+      Can expect an `axes` arguments specifying spatial axes in inputs.
+    output_scale_fn: similar to `inputs_scale_fn` but applied to outputs.
+    name: a name for this CNN tower. This name will appear in Xprof traces.
+
+  Returns:
+    CNN tower with specified configuration.
+  """
+  channels = (num_hidden_channels,) * num_hidden_layers
+  kernel_shapes = ((kernel_size,) * ndim,) * num_hidden_layers
+  output_kernel_shape = (output_kernel_size,) * ndim
+  return forward_flex_tower_factory(
+      num_output_channels=num_output_channels, ndim=ndim, channels=channels,
+      kernel_shapes=kernel_shapes, rates=rates, strides=strides,
+      output_kernel_shape=output_kernel_shape, output_rate=output_dilation_rate,
+      output_stride=output_stride, conv_module=conv_module,
+      nonlinearity=nonlinearity, inputs_scale_fn=inputs_scale_fn,
+      output_scale_fn=output_scale_fn, name=name)
+
+
+@gin.register
+def forward_flex_tower_factory(
+    num_output_channels: int,
+    ndim: int,
+    channels: Tuple[int, ...] = (16, 16),
+    kernel_shapes: Tuple[Tuple[int, ...], ...] = ((3, 3), (3, 3)),
+    rates: Tuple[int, ...] = (1, 1),
+    strides: Tuple[int, ...] = (1, 1),
+    output_kernel_shape: Tuple[int, ...] = (3, 3),
+    output_rate: int = 1,
+    output_stride: int = 1,
+    conv_module: ConvModule = periodic_convolution,
+    nonlinearity: Callable[[Array], Array] = nonlinearities.relu,
+    inputs_scale_fn: ScaleFn = lambda x, axes: x,
+    output_scale_fn: ScaleFn = lambda x, axes: x,
+    name: Optional[str] = 'forward_flex_cnn_tower',
+):
+  """Constructs CNN tower with specified architecture.
+
+  Args:
+    num_output_channels: number of channels in the output layer.
+    ndim: number of spatial dimensions to expect in inputs to the network.
+    channels: tuple specifying number of channels in hidden layers.
+    kernel_shapes: tuple specifying shapes of kernels in hidden layers.
+      Each entry must be a tuple that specifies a valid kernel_shape for the
+      provided `conv_module`. Must have the same length as `channels`.
+    rates: dilation rates of the convolutions.
+    strides: strides to use in convolutions.
+    output_kernel_shape: shape of the output kernel.
+    output_rate: dilation rate of the final convolution.
+    output_stride: stride of the final convolution.
+    conv_module: convolution module to use. Must accept
+      (output channels, kernel shape and ndim).
+    nonlinearity: nonlinearity function to apply between hidden layers.
+    inputs_scale_fn: scaling function to be applied to the inputs of the tower.
+      Must take `inputs`, `axes` arguments specifying input `Array` and
+      spatial dimensions and return an `Array` of the same shape as `inputs`.
+    output_scale_fn: similar to `inputs_scale_fn` but applied to outputs.
+    name: a name for this CNN tower. This name will appear in Xprof traces.
+
+  Returns:
+    CNN tower with specified architecture.
+  """
+  if isinstance(strides, int):
+    strides = (strides,) * len(channels)
+  if isinstance(rates, int):
+    rates = (rates,) * len(channels)
+
+  ndim_axes = list(range(ndim))
+  n_convs = len(channels)
+  if not all(len(arg) == n_convs for arg in [kernel_shapes, rates, strides]):
+    raise ValueError('conflicting lengths for channels/kernels/rates/strides: '
+                     f'{channels} / {kernel_shapes} / {rates} / {strides}')
+  def forward_pass(inputs):
+    components = [functools.partial(inputs_scale_fn, axes=ndim_axes)]
+    conv_args = zip(channels, kernel_shapes, rates, strides)
+    for num_channels, kernel_shape, rate, stride in conv_args:
+      components.append(conv_module(num_channels, kernel_shape, ndim, rate=rate,
+                                    stride=stride))
+      components.append(nonlinearity)
+    components.append(conv_module(num_output_channels, output_kernel_shape,
+                                  ndim, rate=output_rate, stride=output_stride))
+    components.append(functools.partial(output_scale_fn, axes=ndim_axes))
+    return hk.Sequential(components)(inputs)
+
+  module = hk.to_module(forward_pass)(name=name)
+  return jax.named_call(module, name=name)
+
+
+@gin.register
+def residual_block_tower_factory(
+    num_output_channels: int,
+    ndim: int,
+    num_blocks: int = 2,
+    block_factory: TowerFactory = forward_tower_factory,
+    skip_connection_fn: Callable[..., Array] = lambda x, block_num: x,
+    inputs_scale_fn: ScaleFn = lambda x, axes: x,
+    output_scale_fn: ScaleFn = lambda x, axes: x,
+    name: Optional[str] = 'residual_block_tower',
+):
+  """Constructs a tower with skip connections between blocks."""
+  def forward_pass(inputs):
+    inputs = inputs_scale_fn(inputs, list(range(ndim)))
+    for block_num in range(num_blocks - 1):
+      skip = skip_connection_fn(inputs, block_num)
+      block = block_factory(skip.shape[-1], ndim)
+      inputs = skip + block(inputs)
+    last_block = block_factory(num_output_channels, ndim)
+    return output_scale_fn(last_block(inputs), list(range(ndim)))
+
+  module = hk.to_module(forward_pass)(name=name)
+  return jax.named_call(module, name=name)
+
+
+@gin.register
+def residual_connection(*args, module_factory, **kwargs):
+  """Apply module_factory() as a residual correction to inputs."""
+  def forward_pass(inputs):
+    return inputs + module_factory(*args, **kwargs)(inputs)
+  return hk.to_module(forward_pass)()

jax_cfd/ml/towers_test.py

+"""Tests for google3.research.simulation.whirl.models.towers."""
+
+import itertools
+from absl.testing import absltest
+from absl.testing import parameterized
+
+import gin
+import haiku as hk
+import jax
+from jax_cfd.base import test_util
+from jax_cfd.ml import towers  # pylint: disable=unused-import
+
+
+TOWERS = ['towers.forward_tower_factory', 'towers.residual_block_tower_factory']
+SCALE_FNS = ['towers.fixed_scale', 'towers.scale_to_range']
+NDIMS = [1, 2, 3]
+INPUT_CHANNELS = [1, 6]
+
+
+def test_parameters():
+  product = itertools.product(TOWERS, SCALE_FNS, NDIMS, INPUT_CHANNELS)
+  parameters = []
+  for tower, scale_fn, ndim, input_channels in product:
+    name = '_'.join([tower, scale_fn, f'{ndim}D', f'{input_channels}_channels'])
+    parameters.append(dict(
+        testcase_name=name,
+        tower_module=tower,
+        scale_fn_module=scale_fn,
+        ndim=ndim,
+        input_channels=input_channels))
+  return parameters
+
+
+@gin.configurable
+def forward_pass_module(
+    num_output_channels,
+    ndim,
+    tower_module=gin.REQUIRED
+):
+  """Constructs a function that initializes tower and applies it to inputs."""
+  def forward_pass(inputs):
+    return tower_module(num_output_channels, ndim)(inputs)
+
+  return forward_pass
+
+
+class TowersTest(test_util.TestCase):
+  """Tests towers construction, configuration and composition."""
+
+  def setUp(self):
+    """Configures all scale_fns that have gin.REQUIRED values."""
+    super().setUp()
+    gin.enter_interactive_mode()
+    config = '\n'.join([
+        'towers.fixed_scale.rescaled_one = 0.3',
+        'towers.scale_to_range.min_value = -1.23',
+        'towers.scale_to_range.max_value = 1.21'
+    ])
+    gin.parse_config(config)
+
+  @parameterized.named_parameters(*test_parameters())
+  def test_output_shapes(
+      self,
+      tower_module,
+      scale_fn_module,
+      ndim,
+      input_channels
+  ):
+    """Tests that towers produce outputs of expected shapes."""
+    gin.enter_interactive_mode()
+    config = '\n'.join([
+        f'forward_pass_module.tower_module = @{tower_module}',
+        f'{tower_module}.inputs_scale_fn = @{scale_fn_module}'
+    ])
+    gin.parse_config(config)
+
+    num_output_channels = 5
+    spatial_size = 17
+    rng = jax.random.PRNGKey(42)
+    inputs = jax.random.uniform(rng, (spatial_size,) * ndim + (input_channels,))
+
+    forward_pass = hk.without_apply_rng(
+        hk.transform(forward_pass_module(num_output_channels, ndim)))
+    params = forward_pass.init(rng, inputs)
+    output = forward_pass.apply(params, inputs)
+    expected_output_shape = inputs.shape[:-1] + (num_output_channels,)
+    actual_output_shape = output.shape
+    self.assertEqual(actual_output_shape, expected_output_shape)
+
+
+if __name__ == '__main__':
+  absltest.main()

jax_cfd/ml/train_utils.py

+# Copyright 2021 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
+#
+#      http://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.
+
+"""Often used training utility funcutils."""
+
+import collections
+
+from typing import Any, Callable, Iterable, Iterator, Mapping, Optional, Tuple, Union
+
+import gin
+import jax
+from jax_cfd.base import array_utils
+from jax_cfd.base import grids
+
+
+# TODO(dkochkov): make utility functions agnostic of jax_cfd/other models;
+Array = grids.Array
+Field = Tuple[Array, ...]
+GridArray = grids.GridArray
+GridArrayVector = grids.GridArrayVector
+
+# TODO(jamieas): Replace `Any` with well-defined types.
+IntOrArray = Union[int, Array]
+NavierStokesState = Tuple[GridArrayVector, GridArray, Optional[GridArray]]
+Velocity = Field
+OptimizerState = Any
+ModelParams = Any
+ModelGradients = ModelParams
+EMAParams = ModelParams
+StepAndOptimizerState = Tuple[IntOrArray, OptimizerState]
+StepOptAndEMAState = Tuple[IntOrArray, OptimizerState, ModelParams]
+LossValue = Array
+LossFunction = Callable[[GridArrayVector, GridArrayVector], LossValue]
+LossAndGradFunction = Callable[[ModelParams, Array],
+                               Tuple[LossValue, ModelGradients]]
+MetricFunction = Callable[[Velocity, Velocity], Array]
+TrainStepFunction = Callable[[StepAndOptimizerState, Array],
+                             Tuple[OptimizerState, Array]]
+EvalStepFunction = Callable[[OptimizerState, Field],
+                            Mapping[str, Array]]
+TrajectoryFunction = Callable[[ModelParams, Velocity],
+                              Tuple[Velocity, Velocity]]
+
+
+#
+#  Note that all functions below deal with *batched* inputs.
+#
+
+
+def loss_and_gradient(
+    trajectory_fn: TrajectoryFunction,
+    loss_fn: LossFunction
+) -> LossAndGradFunction:
+  """Returns a function that computes loss and the gradient of the loss.
+
+  Args:
+    trajectory_fn: a function that accepts `params` and `initial_velocity`
+      and returns a trajectory of velocities.
+    loss_fn: a function that accepts a predicted trajectory and a ground truth
+      trajectory, returning a scalar loss value.
+
+  Returns:
+    A function that accepts `params, initial_velocity, target_trajectory` and
+    returns the loss and the gradient of the loss.
+  """
+  def _loss(params: ModelParams, target_trajectory: Velocity) -> LossValue:
+    """Returns loss value and gradient with respect to model parameters."""
+    _, predicted_trajectory = trajectory_fn(params, target_trajectory)
+    loss = loss_fn(predicted_trajectory, target_trajectory)  # type: ignore
+    return loss
+  return jax.value_and_grad(_loss)
+
+
+def train_step(
+    loss_and_grad_fn: LossAndGradFunction,
+    update_fn: Callable[[int, ModelGradients, OptimizerState], OptimizerState],
+    get_params_fn: Callable[[OptimizerState], ModelParams]
+) -> TrainStepFunction:
+  """Returns a function that performs a single training step.
+
+  Args:
+    loss_and_grad_fn: a function that accepts `params, initial_velocity,
+      target_trajectory` and returns the loss and the gradient of the loss.
+    update_fn: a function that accepts `step_num, gradients, optimizer_state`
+      and returns an updated optimizer state.
+    get_params_fn: a function that accepts `optimizer_state` and returns model
+      params. If the state is encoded by params, this should be the identity.
+
+  Returns:
+    A function that performs a single training step.
+  """
+  def _train_step(
+      step_and_state: StepAndOptimizerState,
+      target_trajectory: Array,
+  ) -> Tuple[StepAndOptimizerState, LossValue]:
+    """A function that performs a single training step."""
+    step, optimizer_state = step_and_state
+    params = get_params_fn(optimizer_state)
+    loss, grad = loss_and_grad_fn(params, target_trajectory)
+    optimizer_state = update_fn(step, grad, optimizer_state)
+    return (step + 1, optimizer_state), loss
+  return _train_step
+
+
+def eval_batch(
+    trajectory_fn: TrajectoryFunction,
+    metric_funcs: Mapping[str, MetricFunction],
+) -> EvalStepFunction:
+  """Returns a function that performs a single evaluation step.
+
+  Args:
+    trajectory_fn: a function that accepts `params` and `initial_velocity`
+      and returns a trajectory of velocities.
+    metric_funcs: a dictionary mapping strings to metric funcutils.
+
+  Returns:
+    A function that performs a single evaluation step.
+  """
+  def _eval_batch(
+      params: ModelParams,
+      target_trajectory: Velocity,
+  ) -> Mapping[str, Array]:
+    """A function that performs a single training step."""
+    _, predicted_trajectory = trajectory_fn(params, target_trajectory)
+    metric_values = {k: metric(predicted_trajectory, target_trajectory)
+                     for k, metric in metric_funcs.items()}
+    return metric_values
+  return _eval_batch
+
+
+def streaming_mean(
+    batches: Iterable[Velocity],
+    eval_fn: Callable[[Field], Mapping[str, Array]],
+) -> Mapping[str, Array]:
+  """Runs evaluation on `eval_data`.
+
+  Args:
+    batches: an iterable of batched velocity trajectories.
+    eval_fn: a function that performs a single evaluation step.
+
+  Returns:
+    A dict mapping strings to metric values.
+
+  Raises:
+    RuntimeError: if there are no batches to iterate over.
+  """
+  # TODO(jamieas): update to accommodate non-scalar metrics.
+  eval_metrics = collections.defaultdict(float)
+  count = 0
+  for batch in batches:
+    batch_metrics = eval_fn(batch)
+    for k, v in batch_metrics.items():
+      eval_metrics[k] += v
+    count += 1
+  if not count:
+    raise RuntimeError("no batches to iterate over")
+  return {k: v / count for k, v in eval_metrics.items()}
+
+
+@gin.register
+def identity(batch: Tuple[Array, ...], rng: Array = None) -> Tuple[Array, ...]:
+  """Identity preprocessing function that does not modify the `batch`."""
+  del rng  # unused.
+  return batch
+
+
+# TODO(dkochkov) consider adding an option to perform pressure projection step.
+@gin.configurable
+def add_noise_to_input_frame(
+    batch: Tuple[Array, ...],
+    rng: Array,
+    scale: float = 1e-2,
+    **kwargs
+) -> Tuple[Array, ...]:
+  """Adds noise to the 0th time frame in the `batch`.
+
+  Args:
+    batch: original batch to which the noise will be added.
+    rng: random number key to be used to generate noise.
+    scale: scale of the normal noise to be added.
+    **kwargs: other keyword arguments. Not used.
+
+  Returns:
+    batch with noise added along the 0th time slice.
+  """
+  del kwargs  # unused.
+  time_zero_slice = array_utils.slice_along_axis(batch, 1, 0)
+  shapes = jax.tree_util.tree_map(lambda x: x.shape, time_zero_slice)
+  rngs = jax.random.split(rng, len(jax.tree_util.leaves(time_zero_slice)))
+  # TODO(dkochkov) add `split_like` method to `array_utils.py`.
+  rngs = jax.tree_util.unflatten(jax.tree_util.structure(time_zero_slice), rngs)
+  noise_fn = lambda key, s: scale * jax.random.truncated_normal(key, -2., 2., s)
+  noise = jax.tree_util.tree_map(noise_fn, rngs, shapes)
+  add_noise_fn = lambda x, n: x.at[:, 0, ...].add(n)
+  return jax.tree_util.tree_map(add_noise_fn, batch, noise)
+
+
+def preprocess(
+    data_iterator: Iterator[Tuple[Array, ...]],
+    rng_stream: Iterator[Array],
+    preprocess_fn: Callable[..., Tuple[Array, ...]]
+):
+  """Generator that applies `preprocess_fn` to entries of the `data_iterator`.
+
+  Args:
+    data_iterator: numpy iterator holding the data.
+    rng_stream: stream of random numbers to be used by `preprocess_fn`.
+    preprocess_fn: preprocessing function to be applied to each batch of data.
+
+  Yields:
+    Batch of data from `data_iterator` preprocessed with `preprocess_fn`.
+  """
+  preprocess_fn = jax.jit(preprocess_fn)
+  while True:
+    rng = next(rng_stream)
+    yield preprocess_fn(next(data_iterator), rng)

jax_cfd/ml/viscosities.py

+"""Models for closure terms and effective viscosities."""
+
+import functools
+from typing import Any, Callable, Optional, Tuple, Union
+
+import gin
+import haiku as hk
+import jax
+import jax.numpy as jnp
+from jax_cfd.base import boundaries
+from jax_cfd.base import grids
+from jax_cfd.base import subgrid_models
+from jax_cfd.ml import interpolations
+from jax_cfd.ml import physics_specifications
+from jax_cfd.ml import towers
+import numpy as np
+
+Array = Union[np.ndarray, jax.Array]
+GridArray = grids.GridArray
+GridArrayVector = grids.GridArrayVector
+GridVariableVector = grids.GridVariableVector
+InterpolationModule = interpolations.InterpolationModule
+ViscosityFn = Callable[[grids.GridArrayTensor, GridArrayVector, grids.Grid],
+                       grids.GridArrayTensor]
+ViscosityModule = Callable[..., ViscosityFn]
+
+
+@gin.register
+def smagorinsky_viscosity(
+    grid: grids.Grid,
+    dt: float,
+    physics_specs: physics_specifications.NavierStokesPhysicsSpecs,
+    viscosity_scale: Optional[float] = None,
+    cs: float = 0.2,
+    interpolation_module: InterpolationModule = interpolations.linear,
+) -> ViscosityFn:
+  """Constructs a Smagorinsky viscosity model."""
+  del viscosity_scale  # unused.
+  interpolate = interpolation_module(grid, dt, physics_specs)
+  viscosity_fn = functools.partial(
+      subgrid_models.smagorinsky_viscosity, dt=dt, cs=cs,
+      interpolate_fn=interpolate)
+  return hk.to_module(viscosity_fn)(name='smagorinsky_viscosity')
+
+
+@gin.register
+def learned_scalar_viscosity(
+    grid: grids.Grid,
+    dt: float,
+    physics_specs: physics_specifications.NavierStokesPhysicsSpecs,
+    viscosity_scale: float,
+    interpolate_module: InterpolationModule = interpolations.linear,
+    tower_factory: Callable[..., Any] = towers.forward_tower_factory,
+) -> ViscosityFn:
+  """Constructs an learned, scalar-valued viscosity model."""
+  def interpolate(
+      c: GridArray,
+      offset: Tuple[float, ...],
+      v: Optional[GridVariableVector] = None,
+      dt: Optional[float] = None,
+  ) -> grids.GridArray:
+    """Interpolation method wrapped for GridArray using periodic BC."""
+    bc = boundaries.periodic_boundary_conditions(grid.ndim)
+    c_bc = grids.GridVariable(c, bc)
+    interp_var = interpolate_module(grid, dt, physics_specs)(
+        c_bc, offset, v, dt)
+    return interp_var.array
+
+  def viscosity_fn(
+      s_ij: grids.GridArrayTensor,
+      v: GridArrayVector,
+  ) -> grids.GridArrayTensor:
+    """Computes effective eddy viscosity using learned components.
+
+    This viscosity model computes parametric scalar viscosity that is
+    interpolated to the offsets of the strain rate tensor.
+
+    Args:
+      s_ij: strain rate tensor that is equal to the forward finite difference
+        derivatives of the velocity field `(d(u_i)/d(x_j) + d(u_j)/d(x_i)) / 2`.
+      v: velocity field.
+
+    Returns:
+      tensor containing values of the eddy viscosity at the same grid offsets
+      as the strain tensor `s_ij`.
+    """
+    s_ij_offsets = [array.offset for array in s_ij.ravel()]
+    unique_offsets = list(set(s_ij_offsets))
+    viscosity_net = tower_factory(1, grid.ndim)
+    inputs = jnp.stack([u.data for u in v], axis=-1)
+    predicted_viscosity = (viscosity_scale + 1e-6) * viscosity_net(inputs)
+    predicted_viscosity = grids.GridArray(
+        data=jnp.squeeze(predicted_viscosity, -1),
+        offset=grid.cell_center, grid=grid)
+    interpolated_viscosities = {
+        offset: interpolate(predicted_viscosity, offset, v, dt)  # pytype: disable=wrong-arg-types  # always-use-return-annotations
+        for offset in unique_offsets}
+    viscosities = [interpolated_viscosities[offset] for offset in s_ij_offsets]
+    tree_def = jax.tree_util.tree_structure(s_ij)
+    return jax.tree_util.unflatten(tree_def, [x.data for x in viscosities])
+
+  return hk.to_module(viscosity_fn)()
+
+
+@gin.register
+def learned_scalar_viscosity_from_gradients(
+    grid: grids.Grid,
+    dt: float,
+    physics_specs: physics_specifications.NavierStokesPhysicsSpecs,
+    viscosity_scale: float,
+    interpolate_module: InterpolationModule = interpolations.linear,
+    tower_factory: Callable[..., Any] = towers.forward_tower_factory,
+) -> ViscosityFn:
+  """Constructs a scalar viscosity model predicted from velocity gradients."""
+  def interpolate(
+      c: GridArray,
+      offset: Tuple[float, ...],
+      v: Optional[GridVariableVector] = None,
+      dt: Optional[float] = None,
+  ) -> grids.GridArray:
+    """Interpolation method wrapped for GridArray using periodic BC."""
+    bc = boundaries.periodic_boundary_conditions(grid.ndim)
+    c_bc = grids.GridVariable(c, bc)
+    interp_var = interpolate_module(grid, dt, physics_specs)(
+        c_bc, offset, v, dt)
+    return interp_var.array
+
+  def viscosity_fn(
+      s_ij: grids.GridArrayTensor,
+      v: GridArrayVector,
+  ) -> grids.GridArrayTensor:
+    """Computes effective eddy viscosity using learned components.
+
+    This viscosity model computes parametric scalar viscosity that is
+    interpolated to the offsets of the strain rate tensor.
+
+    Args:
+      s_ij: strain rate tensor that is equal to the forward finite difference
+        derivatives of the velocity field `(d(u_i)/d(x_j) + d(u_j)/d(x_i)) / 2`.
+      v: velocity field.
+
+    Returns:
+      tensor containing values of the eddy viscosity at the same grid offsets
+      as the strain tensor `s_ij`.
+    """
+    s_ij_offsets = [array.offset for array in s_ij.ravel()]
+    unique_offsets = list(set(s_ij_offsets))
+    viscosity_net = tower_factory(1, grid.ndim)
+    cell_center = grid.cell_center
+    interpolate_to_center = lambda x: interpolate(x, cell_center, v, dt)  # pytype: disable=wrong-arg-types
+    centered_s_ij = np.vectorize(interpolate_to_center)(s_ij)
+    inputs = jnp.stack([array.data for array in centered_s_ij.ravel()], axis=-1)
+    predicted_viscosity = (viscosity_scale + 1e-6) * viscosity_net(inputs)
+    predicted_viscosity = grids.GridArray(
+        data=jnp.squeeze(predicted_viscosity, -1),
+        offset=grid.cell_center, grid=grid)
+    interpolated_viscosities = {
+        offset: interpolate(predicted_viscosity, offset, v, dt)  # pytype: disable=wrong-arg-types  # always-use-return-annotations
+        for offset in unique_offsets}
+    viscosities = [interpolated_viscosities[offset] for offset in s_ij_offsets]
+    tree_def = jax.tree_util.tree_structure(s_ij)
+    return jax.tree_util.unflatten(tree_def, [x.data for x in viscosities])
+
+  return hk.to_module(viscosity_fn)()
+
+
+@gin.register
+def learned_tensor_viscosity(
+    grid: grids.Grid,
+    dt: float,
+    physics_specs: physics_specifications.NavierStokesPhysicsSpecs,
+    viscosity_scale: float,
+    tower_factory: Callable[..., Any] = towers.forward_tower_factory,
+) -> ViscosityFn:
+  """Constructs an learned, tensor-valued viscosity model."""
+  del dt, physics_specs
+
+  def viscosity_fn(
+      s_ij: grids.GridArrayTensor,
+      v: GridArrayVector,
+  ) -> grids.GridArrayTensor:
+    """Computes effective eddy viscosity using learned components.
+
+    This viscosity model computes parametric tensor viscosity that predicts
+    independent values at all offsets of the strain rate tensor.
+
+    Args:
+      s_ij: strain rate tensor that is equal to the forward finite difference
+        derivatives of the velocity field `(d(u_i)/d(x_j) + d(u_j)/d(x_i)) / 2`.
+      v: velocity field.
+
+    Returns:
+      tensor containing values of the eddy viscosity at the same grid offsets
+      as the strain tensor `s_ij`.
+    """
+    s_ij_offsets = [array.offset for array in s_ij.ravel()]
+    unique_offsets = list(set(s_ij_offsets))
+    num_offsets = len(unique_offsets)
+    viscosity_net = tower_factory(num_offsets, grid.ndim)
+    inputs = jnp.stack([u.data for u in v], axis=-1)
+    viscosities = (viscosity_scale + 1e-6) * viscosity_net(inputs)
+    viscosities = jnp.split(viscosities, np.arange(1, num_offsets), axis=-1)
+    viscosities_dict = {
+        offset: jnp.squeeze(visc, axis=-1)  # remove channel dimension.
+        for offset, visc in zip(unique_offsets, viscosities)}
+    viscosities = [viscosities_dict[offset] for offset in s_ij_offsets]
+    tree_def = jax.tree_util.tree_structure(s_ij)
+    return jax.tree_util.unflatten(tree_def, viscosities)
+
+  return hk.to_module(viscosity_fn)()
+
+
+@gin.register
+def eddy_viscosity_model(
+    grid: grids.Grid,
+    dt: float,
+    physics_specs: physics_specifications.NavierStokesPhysicsSpecs,
+    v: Optional[GridArrayVector] = None,
+    viscosity_scale: Optional[float] = None,
+    viscosity_model: ViscosityModule = smagorinsky_viscosity,
+):
+  """Constructs eddy viscosity model that computes accelerations.
+
+  Eddy viscosity models compute a turbulence closure term as a divergence of the
+  subgrid-scale stress tensor, which is expressed as velocity dependent
+  viscosity times the rate of strain tensor. This module delegates computation
+  of the eddy-viscosity to `viscosity_model` function. Note that if outputs of
+  the `viscosity_model` are not unrestricted to interpolations of a scalar
+  field, this model can represent almost arbitrary stress tensor.
+  For details see go/whirl-evm.
+
+  Args:
+    grid: grid on which the Navier-Stokes equation is discretized.
+    dt: time step to use for time evolution.
+    physics_specs: physical parameters of the simulation module.
+    v: optional velocity field that is used to precompute values in models.
+    viscosity_scale: the kinematic viscosity of the fluid.
+    viscosity_model: function that generates a `viscosity_fn`.
+
+  Returns:
+    Function that computes accelerations due to eddy viscosity model.
+  """
+  del v  # unused.
+  if viscosity_scale is None:
+    viscosity_scale = physics_specs.viscosity
+  viscosity = viscosity_model(
+      grid, dt, physics_specs, viscosity_scale=viscosity_scale)
+  evm_fn = functools.partial(subgrid_models.evm_model, viscosity_fn=viscosity)
+  return hk.to_module(evm_fn)(name='eddy_viscosity_model')

jax_cfd/spectral/__init__.py

+# Copyright 2021 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
+#
+#      http://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.
+
+"""Pseudospectral codes (no ML)."""
+
+import jax_cfd.spectral.equations
+import jax_cfd.spectral.time_stepping
+import jax_cfd.spectral.types
+import jax_cfd.spectral.utils