model/reference_code/stochastic_losses.py

# Copyright 2024 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Stochastic losses for NeuralGCM."""
import abc
from typing import Callable, Optional, Sequence
from dinosaur import typing
import gin
import jax
import jax.numpy as jnp
import model.reference_code.linear_transforms as linear_transforms
import model.reference_code.metrics_base as metrics_base
import model.reference_code.metrics_util as metrics_util
from model.legacy import model_utils
import numpy as np


Pytree = typing.Pytree
TrajectoryRepresentations = typing.TrajectoryRepresentations

AggregationTransformConstructor = metrics_util.AggregationTransformConstructor

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


def replicate(
    x: Pytree,
    axis_name: str = 'batch',
    times: Optional[int] = None,
) -> Pytree:
  """Replicated a pytree across devices."""
  if times is None:
    times = jax.local_device_count()

  def _replicate(_):
    return x

  return jax.pmap(_replicate, axis_name)(np.ones(times))


class EnergyLikeLoss(metrics_base.Loss, abc.ABC):
  """Energy-score like loss function.

  Both CRPS and EnergyScore take the form (with E expectation)
    E‖X - Y‖^β - ensemble_term_weight * E‖X - X'‖^β
  where for CRPS ‖⋅‖ is the L1 norm, and for EnergyScore it is the L2 norm.

  To create a general implementation, we decompose the norm as
    ‖Z‖ := _norm_reduction_fn(_norm_inner_fn(Z))

  For more see (21) and (22) in [1]; http://shortn/_Lyu0etEy1F

  References:
    [1]: Gneiting, T., & Raftery, A. E. (2007). Strictly proper scoring rules,
         prediction, and estimation. Journal of the American statistical
         Association, 102(477), 359-378.
  """

  def __init__(
      self,
      trajectory_spec: metrics_util.TrajectorySpec,
      components: Sequence[linear_transforms.LinearTransformConstructor],
      time_step: Optional[int | slice] = None,
      level: Optional[int] = None,
      getter: Callable[[Pytree], Pytree] = (
          metrics_util.filter_sim_time_and_diagnostics
      ),
      beta: float = 1.0,
      ensemble_term_weight: float = 0.5,
      is_nodal: bool = True,
      is_encoded: bool = False,
      coarsen_aggregation: AggregationTransformConstructor = (
          metrics_util.AggregateIdentity
      ),
      vector_norm_squared_aggregation: AggregationTransformConstructor = (
          metrics_util.AggregateIdentity
      ),
  ):
    """Constructs an instance of EnergyLikeLoss.

    Args:
      trajectory_spec: Specification of spatial and temporal trajectory sizes.
      components: Sequence of linear transformations to be applied to errors.
      time_step: Step or slice at which to compute loss, or None for all steps.
      level: Level to compute loss at, or None to use mean over all levels.
      getter: Function for extracting a sub-pytree on which errors are computed.
      beta: Power parameter of the loss. For energy score to be strictly proper
        beta must be belong to `(0, 2)`.
      ensemble_term_weight: Coefficient that specifcies how much weight is put
        on the terms that captures the spread of the 2-ensemble. For standard
        energy score this value should be set to `0.5`. It can be used to
        interpolate to other scoring rules that are not strictly proper. For
        example setting this value to `0.0` and setting `beta = 2.0` will result
        in a squared error loss.
      is_nodal: Indicator whether loss is computed in nodal space.
      is_encoded: Indicator whether loss is computed in encoded(model) space.
      coarsen_aggregation: Transform class that is used to aggregate errors
        before computing the loss elements. This enables defining losses on
        coarser representations that accentuate larger scale structure.
        Currently this argument should be used only by PatchEnergyLoss. Example
        coarsening operators include `RegriddingAggregation`, `TimeWindowSum`.
      vector_norm_squared_aggregation: Transform class that is used to aggregate
        components of the squared errors to form the distance for computing the
        energy score. Currently this argument should be used only by
        PatchEnergyLoss. Suitable aggregation methods include
        `RegriddingAggregation`, `TimeWindowSum`, `SumVariables`, which would
        correspond to vectors of (1) single level, time, variable, horizontal
        neighbors; (2) single level, variable, lon-lat, sequence of time values;
        (3) all variables at a single level, time, lon-lat.
    """
    self.coarsen_fn = coarsen_aggregation(
        trajectory_spec, is_nodal=is_nodal, is_encoded=is_encoded
    )
    self.vector_norm_squared_fn = vector_norm_squared_aggregation(
        self.coarsen_fn.out_trajectory_spec,
        is_nodal=is_nodal,
        is_encoded=is_encoded,
    )
    # parent class reductions are done on the final out_trajectory_spec.
    super().__init__(
        self.vector_norm_squared_fn.out_trajectory_spec,
        is_nodal=is_nodal,
        is_encoded=is_encoded,
    )
    self.components = components
    self.time_step = time_step
    self.level = level
    self.getter = getter
    # transform is applied to raw inputs which are aligned with trajectory_spec.
    self.transform = linear_transforms.ComposedTransformForLoss(
        trajectory_spec, self.components
    )
    self._beta = beta
    self._ensemble_term_weight = ensemble_term_weight

  def a_minus_cb(self, a: Pytree, c: float, b: Pytree) -> Pytree:
    """A - c * B."""
    return tree_map(lambda a_i, b_i: a_i - c * b_i, a, b)

  def ca_minus_b(self, c: float, a: Pytree, b: Pytree) -> Pytree:
    """c * A - B."""
    return tree_map(lambda a_i, b_i: c * a_i - b_i, a, b)

  def component_mean(self, tree: Pytree) -> jax.Array:
    """Mean over variable, time, pressure, lat, lon."""
    leaf_means = tree_leaves(self.mean_per_variable(tree))
    return sum(leaf_means) / len(leaf_means)

  def ensemble_mean(self, tree: Pytree) -> Pytree:
    return jax.lax.pmean(tree, 'ensemble')

  def _prepare(self, trajectory: TrajectoryRepresentations) -> Pytree:
    """Prepares target or predictions."""
    # Cannot consolidate with RMSE.prepare since this one
    # * does not take ensemble mean of trajectory.
    trajectory = metrics_util.extract_variable(
        trajectory,
        self.trajectory_spec,
        self.time_step,
        self.level,
        self.getter,
        self.is_nodal,
        self.is_encoded,
    )
    return trajectory

  def evaluate(
      self,
      prediction: TrajectoryRepresentations,
      target: TrajectoryRepresentations,
  ) -> Pytree:
    """Evaluates giving values of interest."""
    pv2ss = self._per_variable_spread_skill_errors(prediction, target)
    return self._spread_skill_and_loss(
        x_minus_y=pv2ss['x_minus_y'],
        x_minus_xprime=pv2ss['x_minus_xprime'],
    )['loss']

  def debug_loss_terms_instance(self) -> metrics_base.EvaluateFunctionWrapper:
    """Returns class that evaluates rel loss per variable and spread/skill."""

    def evaluate_fn(
        prediction: TrajectoryRepresentations,
        target: TrajectoryRepresentations,
    ) -> Pytree:
      # self.loss.evaluate takes ensemble mean (to evaluate on ensemble mean) if
      # needed.
      pv2ss = self._per_variable_spread_skill_errors(prediction, target)
      overall_spread_skill_loss = self._spread_skill_and_loss(
          x_minus_y=pv2ss['x_minus_y'],
          x_minus_xprime=pv2ss['x_minus_xprime'],
      )

      all_vars = pv2ss['x_minus_y'].keys()

      per_variable_terms = {
          var: self._spread_skill_and_loss(
              x_minus_y=pv2ss['x_minus_y'][var],
              x_minus_xprime=pv2ss['x_minus_xprime'][var],
          )
          for var in all_vars
      }
      # here we reduce terms by summation to expose relative contributions,
      # even though the actual total_loss might be different.
      per_variable_losses = {
          var: per_variable_terms[var]['loss'] for var in all_vars
      }
      sum_of_losses = sum(per_variable_losses.values())
      per_variable_relative_losses = tree_map(
          lambda x: x / sum_of_losses, per_variable_losses
      )
      return {
          'relative_loss': per_variable_relative_losses,
          'overall': overall_spread_skill_loss,
          'per_variable_spread': {
              var: per_variable_terms[var]['spread'] for var in all_vars
          },
          'per_variable_skill': {
              var: per_variable_terms[var]['skill'] for var in all_vars
          },
      }

    return metrics_base.EvaluateFunctionWrapper(evaluate_fn)

  def _per_variable_spread_skill_errors(
      self,
      prediction: TrajectoryRepresentations,
      target: TrajectoryRepresentations,
  ) -> Pytree:
    """Computes non-reduced loss terms (skill and spread) for each variable.

    Args:
      prediction: predicted 2-ensemble of trajectories with each component
        having shape [2, time_steps, vertical, lat_axis, lon_axis], with leading
        axis corresponding to different ensemble members and last two axes being
        either spherical harmonics numbers or lat, lon values.
      target: target trajectory replicated along the ensemble axis. The shape is
        expected to be exactly the same as `trajectory`.

    Returns:
      A dictionary with keys containing transformed variables.
        `x_minus_y` = prediction - target
        `x_minus_xprime` = difference of ensemble predictions
        `prediction` = prediction
    """
    ensemble_size = jax.lax.psum(1, 'ensemble')
    if ensemble_size != 2:
      raise ValueError(f'{ensemble_size=} is not 2')

    prediction = self.transform(self._prepare(prediction), target)
    target = self.transform(self._prepare(target), target)

    x_minus_y = tree_map(jnp.subtract, prediction, target)  # X_i - Y

    xprime = jax.lax.pshuffle(prediction, 'ensemble', (1, 0))
    x_minus_xprime = tree_map(jnp.subtract, prediction, xprime)  # X_i - X_j≠i

    return {
        'x_minus_y': x_minus_y,
        'x_minus_xprime': x_minus_xprime,
        'prediction': prediction,
    }

  @abc.abstractmethod
  def _spread_skill_and_loss(
      self,
      x_minus_y: Pytree,
      x_minus_xprime: Pytree,
  ) -> dict[str, jax.Array]:
    """Gets dictionary with 'spread', 'skill', and 'loss' entries."""


@gin.register(
    denylist=['coarsen_aggregation', 'vector_norm_squared_aggregation']
)
class CRPSLoss(EnergyLikeLoss):
  """CRPS loss on linearly transformed errors.

  CRPS takes the form (with E expectation)
    E‖X - Y‖^β - ensemble_term_weight * E‖X - X'‖^β
  where ‖⋅‖ is the L1 norm. It can be thought of as the sum of component-wise
  energy score losses.

  Based on formula 21 in [1]; http://shortn/_Lyu0etEy1F

  References:
    [1]: Gneiting, T., & Raftery, A. E. (2007). Strictly proper scoring rules,
         prediction, and estimation. Journal of the American statistical
         Association, 102(477), 359-378.
  """

  def _spread_skill_and_loss(
      self,
      x_minus_y: Pytree,
      x_minus_xprime: Pytree,
  ) -> dict[str, jax.Array]:
    """Gets dictionary with 'spread', 'skill', and 'loss' entries."""
    a_minus_cb = self.a_minus_cb
    ensemble_mean = self.ensemble_mean
    component_mean = self.component_mean

    def abs_beta(tree: Pytree) -> Pytree:
      return tree_map(lambda x: jnp.abs(x) ** self._beta, tree)

    # With X, X' two i.i.d. predictions,
    #   Skill  = (1/2)[ (1/N)Σₙ|Xₙ-Yₙ| + (1/N)Σₙ|Xₙ'-Yₙ| ]
    #   Spread = (1/N) Σₙ|Xₙ-Xₙ'|

    # Recall x_minus_y = X-Y on one device and X'-Y on another. So the ensemble
    # mean of this (which is all-reduced) is exactly Skill above.
    skill = component_mean(ensemble_mean(abs_beta(x_minus_y)))

    # One device has X-X' and the other has X'-X, so the ensemble mean is the
    # same on both devices.
    spread = component_mean(ensemble_mean(abs_beta(x_minus_xprime)))

    # Then CRPS = Skill - (1/2) Spread
    # However, this is unstable if Spread = 2Skill + ε, where |ε| << |Spread|.
    # In particular, up to numerical precision, CRPS will equal 0!
    # This can happen if Prob[Xₙ = 1] = p << 1, and Prob[Xₙ = 0] = 1 - p.
    # a stable estimate of CRPS is
    #   CRPS = C + C'  (an ensemble mean)
    # where
    #      C = (1/N) Σₙ[ |Xₙ-Yₙ| - (1/2) |Xₙ-Xₙ'| ]
    #      C'= (1/N) Σₙ[ |Xₙ'-Yₙ| - (1/2) |Xₙ'-Xₙ| ]
    # which should be re-written as
    #   CRPS = (1/(2N)) Σₙ[ |Xₙ-Yₙ| + |Xₙ'-Yₙ| - |Xₙ-Xₙ'| ]
    # The triangle inequality ensures the summands are non-negative.
    crps = component_mean(
        ensemble_mean(
            a_minus_cb(  # |Xₙ-Yₙ| - (1/2) |Xₙ-Xₙ'|
                abs_beta(x_minus_y),
                self._ensemble_term_weight,
                abs_beta(x_minus_xprime),
            )
        )
    )
    return {'spread': spread, 'skill': skill, 'loss': crps}


@gin.register(
    denylist=['coarsen_aggregation', 'vector_norm_squared_aggregation']
)
class EnergyScoreLoss(EnergyLikeLoss):
  """Energy score loss on linearly transformed errors.

  EnergyScoreLoss takes the form (with E expectation)
    E‖X - Y‖^β - ensemble_term_weight * E‖X - X'‖^β
  where ‖⋅‖ is the L2 norm. It is a generalization of CRPS to
  multiple-dimensions.

  Based on formula 22 in [1]; http://shortn/_Lyu0etEy1F

  References:
    [1]: Gneiting, T., & Raftery, A. E. (2007). Strictly proper scoring rules,
         prediction, and estimation. Journal of the American Statistical
         Association, 102(477), 359-378.
  """

  def _spread_skill_and_loss(
      self,
      x_minus_y: Pytree,
      x_minus_xprime: Pytree,
  ) -> dict[str, jax.Array]:
    """Gets dictionary with 'spread', 'skill', and 'loss' entries."""
    a_minus_cb = self.a_minus_cb
    ensemble_mean = self.ensemble_mean
    component_mean = self.component_mean

    def sqrt_beta(x: jax.Array) -> jax.Array:
      return model_utils.safe_sqrt(x) ** self._beta

    def square(tree: Pytree) -> Pytree:
      return tree_map(jnp.square, tree)

    # With X, X' two i.i.d. predictions,
    #   Skill  = (1/2)[ ‖X-Y‖ + ‖X'-Y‖ ]
    #   Spread = ‖Xₙ-Xₙ'‖

    # Recall x_minus_y = X-Y on one device and X'-Y on another. So the ensemble
    # mean of this (which is all-reduced) is exactly Skill above.
    skill = ensemble_mean(sqrt_beta(component_mean(square(x_minus_y))))

    # One device has X-X' and the other has X'-X, so the ensemble mean is the
    # same on both devices. The call to ensemble_mean simply removes the
    # ensemble dim.
    spread = ensemble_mean(sqrt_beta(component_mean(square(x_minus_xprime))))

    # The straightforward implementation will lose resolution when the relative
    # difference between
    #   ‖X - X'‖  AND ‖X - Y‖ + ‖X' - Y‖,
    # is less than 1e-6. This is so unlikely that we do will not handle it.
    es_straightforward = a_minus_cb(skill, self._ensemble_term_weight, spread)
    es = es_straightforward

    # Unused demonstration of how to handle this co-linear case with lots of
    # extra complex operations.
    # if float(self._beta) != 1:
    #   es = es_straightforward
    # else:
    #   # If beta == 1, there is a high resolution fix.
    #   # See http://screen/BkvX57d9B9eqMrB
    #   #
    #   # alpha = ‖X - Y‖²
    #   alpha = component_mean(square(x_minus_y))
    #   # And if ensemble_term_weight == 1/2,
    #   # gamma_minus_alpha = ‖X - X'‖²/4 - ‖X - Y‖²
    #   #                  = (1/N) Σₙ[ (Xₙ-Xₙ')²/4 - (Xₙ-Yₙ)² ]
    #   gamma_minus_alpha = component_mean(
    #       self.ca_minus_b(
    #           self._ensemble_term_weight**2,
    #           square(x_minus_xprime),
    #           square(x_minus_y),
    #       )
    #   )
    #   # If gamma = ‖X - X'‖²/4 = 0 (e.g. at step=0), then (γ-α)/α = -1,
    #   # and then grad(sqrt1pm1) is NaN. However, in this case we can use the
    #   # straightforward version with no issues.
    #   gamma_minus_alpha_div_alpha = gamma_minus_alpha / alpha

    #   # Construct a "safe" input to use in the go/tf-where-nan trick.
    #   cutoff = -0.1
    #   safe_gamma_minus_alpha_div_alpha = jnp.maximum(
    #       gamma_minus_alpha / alpha, cutoff
    #   )

    #   # For γ ≈ α, safe_gamma_minus_alpha_div_alpha =
    #   # gamma_minus_alpha_div_alpha, and this code block will be used.
    #   es_for_small_diffs = ensemble_mean(
    #       # sqrt1pm1(z) = sqrt(z + 1) - 1, so
    #       # sqrt(α) * -1 * sqrt1pm1((γ-α)/α)
    #       # = sqrt(α) * (1 - sqrt((γ-α)/α) + 1)
    #       # = sqrt(α) * (1 - sqrt(γ/α))
    #       # = sqrt(α) - sqrt(γ)
    #       # = sqrt(‖X - Y‖²) - sqrt(‖X - X'‖²/4)
    #       jnp.sqrt(alpha)
    #       * -1
    #       * tfp.math.sqrt1pm1(safe_gamma_minus_alpha_div_alpha)
    #   )
    #   es = jnp.where(
    #       # Reminder that the triangle-inequality shows γ <= α always. So
    #       # (γ - α) / α < 0.1 is a "large diff" (despite being negative).
    #       gamma_minus_alpha_div_alpha < cutoff,
    #       es_straightforward,
    #       es_for_small_diffs,
    #   )
    return {'spread': spread, 'skill': skill, 'loss': es}