Code update

PiperOrigin-RevId: 677869063

Author: The swirl_lm Authors

swirl_lm/communication/send_recv.py

@@ -0,0 +1,143 @@
+# Copyright 2024 The swirl_lm Authors.
+#
+# 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.
+
+"""A library for communicating information across replicas.
+
+In an example of 4 replicas, with each replica has data with different sizes as:
+  replica 0: data = tf.constant([])
+  replica 1: data = tf.constant([1])
+  replica 2: data = tf.constant([2, 2])
+  replica 3: data = tf.constant([3, 3, 3])
+If data is shared in an order of 0 -> 1 -> 2 -> 3, the corresponding
+`source_dest_pairs` is [[0, 1], [1, 2], [2, 3]]. With a buffer size `n_max = 3`,
+calling `send_recv(data, source_dest_pairs, n_max)` provides the following:
+  replica 0: tf.constant([0, 0, 0])
+  replica 1: tf.constant([])
+  replica 2: tf.constant([1])
+  replica 3: tf.constant([2, 2]).
+
+Note that in the example above, the `source_dest_pairs` can be obtained by
+calling `source_dest_pairs_along_dim(np.array([[[0]], [[1]], [[2]], [[3]]]), 0,
+True, False)`,
+or `source_dest_pairs_along_dim(np.array([[[0]], [[1]], [[2]], [[3]]])
+*parse_dim('+x'))`.
+"""
+
+import re
+
+import numpy as np
+import tensorflow as tf
+
+
+def parse_dim(dim_info: str) -> tuple[int, bool, bool]:
+  """Parses a dimension string into a tuple (dim, forward, periodic).
+
+  Args:
+    dim_info: A string that has a structure '[-+][xyz]p?$'. The first character
+      is '-' or '+', which indicates the negative or positive direction,
+      respectively. The second character is one of 'x', 'y', and 'z', which
+      corresponds to dimension 0, 1, and 2, respectively. The optional last
+      character is 'p', which suggests the dimension is periodic if present.
+
+  Returns:
+    A 3-element tuple, with the first element being the dimension, the second
+    indicating whether the dimension is along the positive direction, and the
+    third indicating whether the dimension is periodic.
+
+  Raises:
+    ValueError if `dim_info` does not match '[-+][xyz]p?$'.
+  """
+  m = re.fullmatch(r'([-+])([xyz])(p?)', dim_info)
+  if m is None:
+    raise ValueError(
+        f'{dim_info} does not conform with the string structure for dimension'
+        ' info ("[-+][xyz]p?$").'
+    )
+
+  dim = 'xyz'.index(m.group(2))
+  forward = m.group(1) == '+'
+  periodic = m.group(3) == 'p'
+
+  return dim, forward, periodic
+
+
+def source_dest_pairs_along_dim(
+    replicas: np.ndarray, dim: int, forward: bool, periodic: bool
+) -> np.ndarray:
+  """Generates a 2-D array of source-target pairs along `dim` in the topology.
+
+  Args:
+    replicas: A 3-D tensor representing the topology of the partitions.
+    dim: The dimension of communication. Should be one of 0, 1, and 2.
+    forward: A boolean argument that indicates sending data from replicas with
+      lower indices to higher indices along the positive direction of the
+      topology. If it is `False`, communication in performed the opposite
+      direction, i.e. from the higher indices to lower indices.
+    periodic: An indicator of whether the topology is periodic. When using the
+      `source_dest_pairs` generated with this function, if `periodic` is
+      `True`, data from the last replica along `dim` will be send to the first
+      replica; otherwise the first replica returns all zeros with the same size
+      as the input. The first and last replica follows the direction specified
+      in `dim`.
+
+  Returns:
+    A 2-D array of size `[num_pairs, 2]`, with the columns being the
+    `replica_id` of the senders and the receivers, respectively.
+  """
+  rolled = np.roll(replicas, -1 if forward else 1, axis=dim)
+  trim = slice(
+      None if periodic or forward else 1,
+      None if periodic or not forward else -1,
+  )
+  stacked = np.moveaxis(np.stack([replicas, rolled]), dim + 1, 1)[:, trim]
+  return np.reshape(stacked, (2, -1)).T
+
+
+def send_recv(
+    data: tf.Tensor, source_dest_pairs: np.ndarray, n_max: int
+) -> tf.Tensor:
+  """Exchanges N-D `tf.Tensor`s across a list of (sender, receiver) pairs.
+
+  Args:
+    data: The n-dimensional tensor to be sent to a different replica. Dimension
+      0 of this tensor can have different sizes across replicas
+    source_dest_pairs: A 2-D numpy array of shape `[num_replicas, 2]`, with the
+      first column being the senders' `replica_id`, and the second one being the
+      receiver's `replica_id`.
+    n_max: The buffer size for the communication. It has to be greater or equal
+      to the maximum number of `data.shape[0]` across all replicas, otherwise a
+      runtime error will occur while padding the buffer for communication.
+
+  Returns:
+    An N-D tensor received from the sender replica specified in
+    `source_dest_pairs`.
+  """
+  # Because `CollectivePermute` permits transferring data that has the same
+  # shape across all replicas only, we need to pad the input data to satisfy
+  # this condition.
+  static_shape = data.get_shape()
+  u = tf.scatter_nd(
+      tf.range(tf.shape(data)[0])[:, tf.newaxis],
+      data,
+      (n_max, *static_shape[1:]),
+  )
+
+  n_received = tf.raw_ops.CollectivePermute(
+      input=tf.shape(data)[0], source_target_pairs=source_dest_pairs
+  )
+  w = tf.raw_ops.CollectivePermute(
+      input=u, source_target_pairs=source_dest_pairs
+  )
+  # Here we trim the padded data back to its original size.
+  return tf.gather_nd(w, tf.where(tf.range(n_max) < n_received))

swirl_lm/equations/scalars.proto

@@ -126,7 +126,7 @@ message TotalEnergy {
 }
 
 // Defines configurations for total humidity.
-// Next id: 6
+// Next id: 7
 message Humidity {
   // An option of whether the effect of subsidence velocity is included in the
   // total humidity equation.
@@ -137,6 +137,8 @@ message Humidity {
   // An option of whether condensation is included in the equations for
   // humidity.
   optional bool include_condensation = 3;
+  // An option for including sedimentation in the convective flux.
+  optional bool include_sedimentation = 6;
   // Specifies the option of microphysics, which will be used if
   // `include_precipitation` is true.
   oneof microphysics {

swirl_lm/equations/source_function/humidity.py

@@ -70,6 +70,10 @@ def __init__(
         self._scalar_params.HasField('humidity') and
         self._scalar_params.humidity.include_condensation)
 
+    self._include_sedimentation = (
+        self._scalar_params.HasField('humidity') and
+        self._scalar_params.humidity.include_sedimentation)
+
     if scalar_name in ('q_r', 'q_s'):
       assert self._include_precipitation, (
           'Calculating q_r without setting include_precipitation to True is not'
@@ -80,6 +84,7 @@ def __init__(
     if (
         self._include_precipitation
         or self._include_condensation
+        or (self._include_sedimentation and scalar_name == 'q_t')
         or scalar_name in ('q_r', 'q_s')
     ):
       assert self._scalar_params.humidity.HasField('microphysics'), (
@@ -332,18 +337,82 @@ def source_fn(
       )
 
     # Compute source terms
-    if self._scalar_name == 'q_t' and self._include_subsidence:
-      subsidence_source = eq_utils.source_by_subsidence_velocity(
-          self._deriv_lib,
-          states[common.KEY_RHO],
-          thermo_states['zz'],
-          thermo_states['q_c'],
-          self._g_dim,
-          additional_states,
-      )
+    if self._scalar_name == 'q_t':
+      if self._include_subsidence:
+        subsidence_source = eq_utils.source_by_subsidence_velocity(
+            self._deriv_lib,
+            states[common.KEY_RHO],
+            thermo_states['zz'],
+            thermo_states['q_c'],
+            self._g_dim,
+            additional_states,
+        )
+        source = tf.nest.map_structure(tf.math.add, source, subsidence_source)
 
-      # Add external source, e.g. sponge forcing and subsidence.
-      source = tf.nest.map_structure(tf.math.add, source, subsidence_source)
+      if self._include_sedimentation:
+        assert self._microphysics is not None, (
+            'Terminal velocity for q_t requires a microphysics model but is'
+            ' undefined.'
+        )
+        w_l = self._microphysics.terminal_velocity(
+            'q_r',
+            {'rho_thermal': states['rho_thermal'], 'q_r': thermo_states['q_l']},
+            additional_states,
+        )
+        if 'q_r' in states:
+          # The sedimentation velocity of liquid-phase cloud cannot be larger
+          # than that of rain.
+          w_r = self._microphysics.terminal_velocity(
+              'q_r',
+              {
+                  'rho_thermal': states['rho_thermal'],
+                  'q_r': states['q_r'],
+              },
+              additional_states,
+          )
+          w_l = tf.nest.map_structure(
+              lambda w_l, w_r: tf.where(tf.greater(w_l, w_r), w_r, w_l),
+              w_l,
+              w_r,
+          )
+        w_i = self._microphysics.terminal_velocity(
+            'q_s',
+            {'rho_thermal': states['rho_thermal'], 'q_s': thermo_states['q_i']},
+            additional_states,
+        )
+        if 'q_s' in states:
+          # The sedimentation velocity of ice-phase cloud cannot be larger than
+          # that of snow.
+          w_s = self._microphysics.terminal_velocity(
+              'q_s',
+              {
+                  'rho_thermal': states['rho_thermal'],
+                  'q_s': states['q_s'],
+              },
+              additional_states,
+          )
+          w_i = tf.nest.map_structure(
+              lambda w_i, w_s: tf.where(tf.greater(w_i, w_s), w_s, w_i),
+              w_i,
+              w_s,
+          )
+        sedimentation_flux_fn = lambda rho, q_l, q_i, w_l, w_i: rho * (
+            q_l * w_l + q_i * w_i
+        )
+        sedimentation_flux = tf.nest.map_structure(
+            sedimentation_flux_fn,
+            states['rho'],
+            thermo_states['q_l'],
+            thermo_states['q_i'],
+            w_l,
+            w_i,
+        )
+        sedimentation_source = self._deriv_lib.deriv_centered(
+            sedimentation_flux, self._g_dim, additional_states
+        )
+        source = tf.nest.map_structure(
+            tf.math.add, source, sedimentation_source
+        )
 
     if self._include_condensation:
       assert isinstance(self._thermodynamics.model, water.Water), (

swirl_lm/numerics/interpolation.py

@@ -16,8 +16,10 @@
 
 import enum
 import functools
+import itertools
 from typing import Sequence, Tuple, TypeAlias
 
+from swirl_lm.utility import common_ops
 from swirl_lm.utility import get_kernel_fn
 from swirl_lm.utility import types
 import tensorflow as tf
@@ -440,3 +442,64 @@ def muscl(r: FlowFieldVal) -> FlowFieldVal:
       diff1,
   )
   return v_neg, v_pos
+
+
+def trilinear_interpolation(
+    field_data: tf.Tensor,
+    points: tf.Tensor,
+    grid_spacing: tf.Tensor,
+    domain_min_pt: tuple[float, float, float] = (0.0, 0.0, 0.0),
+) -> tf.Tensor:
+  """Linear interpolation on a 3-D orthogonal, uniform grid.
+
+  Performs trilinear interpolation by calculating the local point coordinate
+  indices and then interpolating the surrounding field data at the point.
+  Points provided outside of the core domain give invalid solutions.
+
+  Note that coordinates in this function follow the order of the dimensions of
+  `field_data` as a 3D tensor instead of the physical-coordinates orientation in
+  Swirl-LM. For instance, the first element in `grid_spacing` and
+  `domain_min_pt`, as well as the first column in `points`, are associated with
+  the 0th dimensions of `field_data`, instead of the 'x' axis in Swirl-LM.
+
+  Args:
+    field_data: A 3D tensor of field scalars without halos.
+    points: An 2D tensor (n, 3) of n coordinate points in 3D space to
+      interpolate at. Points must fall within the range of coordinates
+      associated with the core. If points are outside of this domain, the
+      function will return invalid results.
+    grid_spacing: A three element tensor defining the grid spacing along the
+      three dimensions.
+    domain_min_pt: A three element tuple defining the minimum coordinate
+      location within the entire physical domain encompassing all cores, default
+      is (0, 0, 0).
+
+  Returns:
+    An n element tensor containing interpolated data values at the n supplied
+      points.
+  """
+  core_spacing = grid_spacing * (
+      tf.cast(field_data.shape, dtype=tf.float32) - 1
+  )
+
+  # Normalizes points to be within the range (0, 0, 0,) to (nx, ny, nz) for nx,
+  # ny, and nz nodes in the core.
+  points_norm = (
+      (points - tf.constant(domain_min_pt)) % core_spacing
+  ) / grid_spacing
+  ijk = tf.floor(points_norm)
+  points_norm -= ijk
+  ijk = tf.cast(ijk, dtype=tf.int32)
+  i, j, k = ijk[:, 0], ijk[:, 1], ijk[:, 2]
+  x0, x1, x2 = points_norm[:, 0], points_norm[:, 1], points_norm[:, 2]
+
+  values = tf.zeros(points.shape[0], dtype=field_data.dtype)
+  for p, q, l in itertools.product(range(2), range(2), range(2)):
+    v = common_ops.gather(field_data, tf.stack([i + p, j + q, k + l], axis=-1))
+    values += v * (
+        ((1 - p) + (2 * p - 1) * x0)
+        * ((1 - q) + (2 * q - 1) * x1)
+        * ((1 - l) + (2 * l - 1) * x2)
+    )
+
+  return values

swirl_lm/physics/combustion/wood.py

@@ -227,9 +227,11 @@ def _radiative_emission(
       emissivity (< 1).
 
   Returns:
-    The radiation source term due to emission.
+    The radiation source term due to emission. If `t` is less than `t_ambient`,
+    the radiation term is 0, i.e., radiation energy can only be lost to
+    ambient conditions.
   """
-  return _SIGMA * k / l * (t**4 - t_ambient**4)
+  return tf.maximum(_SIGMA * k / l * (t**4 - t_ambient**4), 0.0)
 
 
 def _evaporation(