diff --git a/docs/_config.yml b/docs/_config.yml new file mode 100644 index 000000000..5033c032f --- /dev/null +++ b/docs/_config.yml @@ -0,0 +1,12 @@ +title: neural-lam Documentation +author: mllam contributors + +execute: + execute_notebooks: off + +repository: + url: https://github.com/mllam/neural-lam + +html: + use_repository_button: true + use_issues_button: true \ No newline at end of file diff --git a/docs/_toc.yml b/docs/_toc.yml new file mode 100644 index 000000000..db5054d0e --- /dev/null +++ b/docs/_toc.yml @@ -0,0 +1,6 @@ +format: jb-book +root: index + +chapters: + - file: quickstart + - file: create_reduced_meps_dataset \ No newline at end of file diff --git a/docs/index.md b/docs/index.md new file mode 100644 index 000000000..78cea6c29 --- /dev/null +++ b/docs/index.md @@ -0,0 +1,13 @@ +# neural-lam Documentation + +Welcome to the official documentation for neural-lam. + +This documentation provides: + +- A Quickstart guide for new users +- Dataset preparation notebooks +- API reference (auto-generated from docstrings) + +## Getting Started + +Start with the [Quickstart Guide](quickstart.md). \ No newline at end of file diff --git a/docs/quickstart.md b/docs/quickstart.md new file mode 100644 index 000000000..6829c9d46 --- /dev/null +++ b/docs/quickstart.md @@ -0,0 +1,50 @@ +# Quickstart Guide + +This guide demonstrates how to use neural-lam for computing evaluation metrics. + +## Installation + +Clone the repository and install in editable mode: + +```bash +git clone https://github.com/mllam/neural-lam.git +cd neural-lam +pip install -e . +``` + +## Minimal Example + +```python +import torch +from neural_lam.metrics import mse, get_metric + +# Dummy predictions +pred = torch.randn(2, 10, 3) +target = torch.randn(2, 10, 3) +pred_std = torch.ones_like(pred) + +# Compute MSE +loss = mse(pred, target, pred_std) +print("MSE:", loss) +``` + +## Using Registered Metrics + +You can also retrieve metrics dynamically: + +```python +metric_fn = get_metric("wmse") +value = metric_fn(pred, target, pred_std) +print("Weighted MSE:", value) +``` + +## Available Metrics + +Currently registered metrics: + +- mse +- wmse +- mae +- wmae +- nll +- crps_gauss \ No newline at end of file diff --git a/neural_lam/metrics.py b/neural_lam/metrics.py index 7db2cca6d..5f2b79b42 100644 --- a/neural_lam/metrics.py +++ b/neural_lam/metrics.py @@ -3,13 +3,16 @@ def get_metric(metric_name): - """ - Get a defined metric with given name + """Get a defined metric by name. - metric_name: str, name of the metric + Args: + metric_name (str): Name of the metric. Returns: - metric: function implementing the metric + Callable: Function implementing the selected metric. + + Raises: + AssertionError: If the metric name is not defined. """ metric_name_lower = metric_name.lower() assert ( @@ -19,62 +22,58 @@ def get_metric(metric_name): def mask_and_reduce_metric(metric_entry_vals, mask, average_grid, sum_vars): - """ - Masks and (optionally) reduces entry-wise metric values + """Mask and optionally reduce entry-wise metric values. + + (...,) represents any number of batch dimensions, potentially + different but broadcastable. - (...,) is any number of batch dimensions, potentially different - but broadcastable - metric_entry_vals: (..., N, d_state), prediction - mask: (N,), boolean mask describing which grid nodes to use in metric - average_grid: boolean, if grid dimension -2 should be reduced (mean over N) - sum_vars: boolean, if variable dimension -1 should be reduced (sum - over d_state) + Args: + metric_entry_vals (torch.Tensor): Tensor of shape (..., N, d_state). + mask (torch.Tensor or None): Boolean mask of shape (N,) describing + which grid nodes to include in the metric. + average_grid (bool): If True, reduce grid dimension (-2) + by taking the mean over N. + sum_vars (bool): If True, reduce variable dimension (-1) + by summing over d_state. Returns: - metric_val: One of (...,), (..., d_state), (..., N), (..., N, d_state), - depending on reduction arguments. + torch.Tensor: One of (...,), (..., d_state), (..., N), or + (..., N, d_state), depending on reduction arguments. """ - # Only keep grid nodes in mask if mask is not None: - metric_entry_vals = metric_entry_vals[ - ..., mask, : - ] # (..., N', d_state) - - # Optionally reduce last two dimensions - if average_grid: # Reduce grid first - metric_entry_vals = torch.mean( - metric_entry_vals, dim=-2 - ) # (..., d_state) - if sum_vars: # Reduce vars second - metric_entry_vals = torch.sum( - metric_entry_vals, dim=-1 - ) # (..., N) or (...,) + metric_entry_vals = metric_entry_vals[..., mask, :] + + if average_grid: + metric_entry_vals = torch.mean(metric_entry_vals, dim=-2) + + if sum_vars: + metric_entry_vals = torch.sum(metric_entry_vals, dim=-1) return metric_entry_vals def wmse(pred, target, pred_std, mask=None, average_grid=True, sum_vars=True): - """ - Weighted Mean Squared Error - - (...,) is any number of batch dimensions, potentially different - but broadcastable - pred: (..., N, d_state), prediction - target: (..., N, d_state), target - pred_std: (..., N, d_state) or (d_state,), predicted std.-dev. - mask: (N,), boolean mask describing which grid nodes to use in metric - average_grid: boolean, if grid dimension -2 should be reduced (mean over N) - sum_vars: boolean, if variable dimension -1 should be reduced (sum - over d_state) + """Weighted Mean Squared Error. + + (...,) represents any number of batch dimensions, potentially + different but broadcastable. + + Args: + pred (torch.Tensor): Predictions of shape (..., N, d_state). + target (torch.Tensor): Targets of shape (..., N, d_state). + pred_std (torch.Tensor): Predicted standard deviation of shape + (..., N, d_state) or (d_state,). + mask (torch.Tensor, optional): Boolean mask of shape (N,). + average_grid (bool): If True, average over grid dimension. + sum_vars (bool): If True, sum over variable dimension. Returns: - metric_val: One of (...,), (..., d_state), (..., N), (..., N, d_state), - depending on reduction arguments. + torch.Tensor: Metric value depending on reduction arguments. """ entry_mse = torch.nn.functional.mse_loss( pred, target, reduction="none" - ) # (..., N, d_state) - entry_mse_weighted = entry_mse / (pred_std**2) # (..., N, d_state) + ) + entry_mse_weighted = entry_mse / (pred_std**2) return mask_and_reduce_metric( entry_mse_weighted, @@ -85,51 +84,44 @@ def wmse(pred, target, pred_std, mask=None, average_grid=True, sum_vars=True): def mse(pred, target, pred_std, mask=None, average_grid=True, sum_vars=True): - """ - (Unweighted) Mean Squared Error - - (...,) is any number of batch dimensions, potentially different - but broadcastable - pred: (..., N, d_state), prediction - target: (..., N, d_state), target - pred_std: (..., N, d_state) or (d_state,), predicted std.-dev. - mask: (N,), boolean mask describing which grid nodes to use in metric - average_grid: boolean, if grid dimension -2 should be reduced (mean over N) - sum_vars: boolean, if variable dimension -1 should be reduced (sum - over d_state) + """Unweighted Mean Squared Error. + + Equivalent to weighted MSE with unit standard deviation. + + Args: + pred (torch.Tensor): Predictions of shape (..., N, d_state). + target (torch.Tensor): Targets of shape (..., N, d_state). + pred_std (torch.Tensor): Ignored; replaced with ones. + mask (torch.Tensor, optional): Boolean mask of shape (N,). + average_grid (bool): If True, average over grid dimension. + sum_vars (bool): If True, sum over variable dimension. Returns: - metric_val: One of (...,), (..., d_state), (..., N), (..., N, d_state), - depending on reduction arguments. + torch.Tensor: Metric value depending on reduction arguments. """ - # Replace pred_std with constant ones return wmse( pred, target, torch.ones_like(pred_std), mask, average_grid, sum_vars ) def wmae(pred, target, pred_std, mask=None, average_grid=True, sum_vars=True): - """ - Weighted Mean Absolute Error - - (...,) is any number of batch dimensions, potentially different - but broadcastable - pred: (..., N, d_state), prediction - target: (..., N, d_state), target - pred_std: (..., N, d_state) or (d_state,), predicted std.-dev. - mask: (N,), boolean mask describing which grid nodes to use in metric - average_grid: boolean, if grid dimension -2 should be reduced (mean over N) - sum_vars: boolean, if variable dimension -1 should be reduced (sum - over d_state) + """Weighted Mean Absolute Error. + + Args: + pred (torch.Tensor): Predictions of shape (..., N, d_state). + target (torch.Tensor): Targets of shape (..., N, d_state). + pred_std (torch.Tensor): Predicted standard deviation. + mask (torch.Tensor, optional): Boolean mask of shape (N,). + average_grid (bool): If True, average over grid dimension. + sum_vars (bool): If True, sum over variable dimension. Returns: - metric_val: One of (...,), (..., d_state), (..., N), (..., N, d_state), - depending on reduction arguments. + torch.Tensor: Metric value depending on reduction arguments. """ entry_mae = torch.nn.functional.l1_loss( pred, target, reduction="none" - ) # (..., N, d_state) - entry_mae_weighted = entry_mae / pred_std # (..., N, d_state) + ) + entry_mae_weighted = entry_mae / pred_std return mask_and_reduce_metric( entry_mae_weighted, @@ -140,50 +132,42 @@ def wmae(pred, target, pred_std, mask=None, average_grid=True, sum_vars=True): def mae(pred, target, pred_std, mask=None, average_grid=True, sum_vars=True): - """ - (Unweighted) Mean Absolute Error - - (...,) is any number of batch dimensions, potentially different - but broadcastable - pred: (..., N, d_state), prediction - target: (..., N, d_state), target - pred_std: (..., N, d_state) or (d_state,), predicted std.-dev. - mask: (N,), boolean mask describing which grid nodes to use in metric - average_grid: boolean, if grid dimension -2 should be reduced (mean over N) - sum_vars: boolean, if variable dimension -1 should be reduced (sum - over d_state) + """Unweighted Mean Absolute Error. + + Equivalent to weighted MAE with unit standard deviation. + + Args: + pred (torch.Tensor): Predictions of shape (..., N, d_state). + target (torch.Tensor): Targets of shape (..., N, d_state). + pred_std (torch.Tensor): Ignored; replaced with ones. + mask (torch.Tensor, optional): Boolean mask of shape (N,). + average_grid (bool): If True, average over grid dimension. + sum_vars (bool): If True, sum over variable dimension. Returns: - metric_val: One of (...,), (..., d_state), (..., N), (..., N, d_state), - depending on reduction arguments. + torch.Tensor: Metric value depending on reduction arguments. """ - # Replace pred_std with constant ones return wmae( pred, target, torch.ones_like(pred_std), mask, average_grid, sum_vars ) def nll(pred, target, pred_std, mask=None, average_grid=True, sum_vars=True): - """ - Negative Log Likelihood loss, for isotropic Gaussian likelihood - - (...,) is any number of batch dimensions, potentially different - but broadcastable - pred: (..., N, d_state), prediction - target: (..., N, d_state), target - pred_std: (..., N, d_state) or (d_state,), predicted std.-dev. - mask: (N,), boolean mask describing which grid nodes to use in metric - average_grid: boolean, if grid dimension -2 should be reduced (mean over N) - sum_vars: boolean, if variable dimension -1 should be reduced (sum - over d_state) + """Negative Log Likelihood for isotropic Gaussian likelihood. + + Args: + pred (torch.Tensor): Predictions of shape (..., N, d_state). + target (torch.Tensor): Targets of shape (..., N, d_state). + pred_std (torch.Tensor): Predicted standard deviation. + mask (torch.Tensor, optional): Boolean mask of shape (N,). + average_grid (bool): If True, average over grid dimension. + sum_vars (bool): If True, sum over variable dimension. Returns: - metric_val: One of (...,), (..., d_state), (..., N), (..., N, d_state), - depending on reduction arguments. + torch.Tensor: Metric value depending on reduction arguments. """ - # Broadcast pred_std if shaped (d_state,), done internally in Normal class - dist = torch.distributions.Normal(pred, pred_std) # (..., N, d_state) - entry_nll = -dist.log_prob(target) # (..., N, d_state) + dist = torch.distributions.Normal(pred, pred_std) + entry_nll = -dist.log_prob(target) return mask_and_reduce_metric( entry_nll, mask=mask, average_grid=average_grid, sum_vars=sum_vars @@ -193,34 +177,32 @@ def nll(pred, target, pred_std, mask=None, average_grid=True, sum_vars=True): def crps_gauss( pred, target, pred_std, mask=None, average_grid=True, sum_vars=True ): - """ - (Negative) Continuous Ranked Probability Score (CRPS) - Closed-form expression based on Gaussian predictive distribution - - (...,) is any number of batch dimensions, potentially different - but broadcastable - pred: (..., N, d_state), prediction - target: (..., N, d_state), target - pred_std: (..., N, d_state) or (d_state,), predicted std.-dev. - mask: (N,), boolean mask describing which grid nodes to use in metric - average_grid: boolean, if grid dimension -2 should be reduced (mean over N) - sum_vars: boolean, if variable dimension -1 should be reduced (sum - over d_state) + """Negative Continuous Ranked Probability Score (CRPS). + + Closed-form expression based on Gaussian predictive distribution. + + Args: + pred (torch.Tensor): Predictions of shape (..., N, d_state). + target (torch.Tensor): Targets of shape (..., N, d_state). + pred_std (torch.Tensor): Predicted standard deviation. + mask (torch.Tensor, optional): Boolean mask of shape (N,). + average_grid (bool): If True, average over grid dimension. + sum_vars (bool): If True, sum over variable dimension. Returns: - metric_val: One of (...,), (..., d_state), (..., N), (..., N, d_state), - depending on reduction arguments. + torch.Tensor: Metric value depending on reduction arguments. """ std_normal = torch.distributions.Normal( - torch.zeros((), device=pred.device), torch.ones((), device=pred.device) + torch.zeros((), device=pred.device), + torch.ones((), device=pred.device), ) - target_standard = (target - pred) / pred_std # (..., N, d_state) + target_standard = (target - pred) / pred_std entry_crps = -pred_std * ( torch.pi ** (-0.5) - 2 * torch.exp(std_normal.log_prob(target_standard)) - target_standard * (2 * std_normal.cdf(target_standard) - 1) - ) # (..., N, d_state) + ) return mask_and_reduce_metric( entry_crps, mask=mask, average_grid=average_grid, sum_vars=sum_vars