Core v0.1.0

# Source code for modulus.metrics.climate.reduction

# Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# 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.

import torch
from torch import Tensor
from abc import ABC
from modulus.metrics.general.reduction import WeightedMean, WeightedVariance

def _compute_lat_weights(lat: Tensor) -> Tensor:
"""Computes weighting for latitude reduction

Parameters
----------
lat : Tensor
A one-dimension tensor [H] representing the latitudes at which the function will
return weights for

Returns
-------
Tensor
Latitude weight tensor [H]
"""

nlat = len(lat)
lat_weight = torch.abs(torch.cos(torch.pi * (lat / 180)))

lat_weight = lat_weight / lat_weight.sum()
return lat_weight

[docs]def zonal_mean(x: Tensor, lat: Tensor, dim: int = -2, keepdims: bool = False) -> Tensor:
"""Computes zonal mean, weighting over the latitude direction that is specified by dim

Parameters
----------
x : Tensor
The tensor [..., H, W] over which the mean will be computed
lat : Tensor
A one-dimension tensor representing the latitudes at which the function will
return weights for
dim : int, optional
The int specifying which dimension of x the reduction will occur, by default -2
keepdims : bool, optional
Keep aggregated dimension, by default False

Returns
-------
Tensor
Zonal mean tensor of x over the latitude dimension
"""
weights = _compute_lat_weights(lat)
wm = WeightedMean(weights)
return wm(x, dim=dim, keepdims=keepdims)

[docs]def zonal_var(
x: Tensor,
lat: Tensor,
std: bool = False,
dim: int = -2,
keepdims: bool = False,
) -> Tensor:
"""Computes zonal variance, weighting over the latitude direction

Parameters
----------
x : Tensor
The tensor [..., H, W] over which the variance will be computed
lat : Tensor
A one-dimension tensor [H] representing the latitudes at which the function will
return weights for
std : bool, optional
Return zonal standard deviation, by default False
dim : int, optional
The int specifying which dimension of x the reduction will occur, by default -2
keepdims : bool, optional
Keep aggregated dimension, by default False

Returns
-------
Tensor
The variance (or standard deviation) of x over the latitude dimension
"""
weights = _compute_lat_weights(lat)
ws = WeightedVariance(weights)
var = ws(x, dim=dim, keepdims=keepdims)
if std:
return torch.sqrt(var)
else:
return var

[docs]def global_mean(x: Tensor, lat: Tensor, keepdims: bool = False) -> Tensor:
"""Computes global mean

This function computs the global mean of a lat/lon grid by weighting over the
latitude direction and then averaging over longitude

Parameters
----------
x : Tensor
The lat/lon tensor [..., H, W] over which the mean will be computed
lat : Tensor
A one-dimension tensor [H] representing the latitudes at which the function will
return weights for
keepdims : bool, optional
Keep aggregated dimension, by default False

Returns
-------
Tensor
Global mean tensor
"""
assert (
x.ndim > 2
), "Expected x to have at least two dimensions, with the last two dimensions representing lat and lon respectively"

# Mean out the latitudes
lat_reduced = zonal_mean(x, lat, dim=-2, keepdims=keepdims)

# Return after reduction across longitudes
return torch.mean(lat_reduced, dim=-1, keepdims=keepdims)

[docs]def global_var(
x: Tensor,
lat: Tensor,
std: bool = False,
keepdims: bool = False,
) -> Tensor:
"""Computes global variance

This function computs the global variance of a lat/lon grid by weighting over the
latitude direction and then averaging over longitude

Parameters
----------
x : Tensor
The lat/lon tensor [..., H, W] over which the variance will be computed
lat : Tensor
A one-dimension tensor [H] representing the latitudes at which the function will
return weights for
std : bool, optional
Return global standard deviation, by default False
keepdims : bool, optional
Keep aggregated dimension, by default False

Returns
-------
Tensor
Global variance tensor
"""
assert (
x.ndim > 2
), "Expected x to have at least two dimensions, with the last two dimensions representing lat and lon respectively"

# Take global mean, incorporated weights
gm = global_mean(x, lat, keepdims=True)

# Take var of lat
lat_reduced = zonal_mean((x - gm) ** 2, lat, dim=-2, keepdims=keepdims)

# Take var over longitude
long_reduce = torch.mean(lat_reduced, dim=-1, keepdims=keepdims)

if std:
return torch.sqrt(long_reduce)
else:
return long_reduce

© Copyright 2023, NVIDIA Modulus Team. Last updated on Aug 8, 2023.