Skip to content

feat: add Log-Spectral Distance (LSD) metric and support spatial frequency analysis#508

Open
sohampatil01-svg wants to merge 2 commits intomllam:mainfrom
sohampatil01-svg:feat/log-spectral-distance
Open

feat: add Log-Spectral Distance (LSD) metric and support spatial frequency analysis#508
sohampatil01-svg wants to merge 2 commits intomllam:mainfrom
sohampatil01-svg:feat/log-spectral-distance

Conversation

@sohampatil01-svg
Copy link
Copy Markdown

@sohampatil01-svg sohampatil01-svg commented Mar 24, 2026

Describe your changes

This PR introduces a generalized Log-Spectral Distance (LSD) metric to neural-lam, providing a powerful tool for evaluating and potentially training models in the spatial frequency domain across all supported grid types.

Motivation

Standard point-wise metrics (MSE, MAE, NLL, CRPS) often fail to distinguish between physically realistic forecasts and overly smooth ("blurry") fields that result from "regression to the mean." LSD addresses this by comparing the power spectra of predictions and targets, directly penalizing discrepancies in energy distribution across different spatial scales.

Following a challenge to support irregular grids, this implementation now works for both:

  1. Regular Grids: Uses fast 2D RFFT (torch.fft.rfft2).
  2. Irregular/Unstructured Grids: Uses Graph Signal Processing (GSP) techniques. Specifically, it computes Laplacian Moments ($m_k = x^T L^k x$) using the Normalized Graph Laplacian ($L = I - D^{-1/2} A D^{-1/2}$). These moments act as spectral estimators that capture energy distribution across different "eigen-bands" of the graph, effectively providing a topology-aware spectral analysis.

Key Changes:

  • neural_lam/metrics.py:
    • Generalized log_spectral_distance to handle both regular and unstructured grids.
    • Added _compute_laplacian_moments helper for efficient GSP-based spectral analysis using power iterations.
  • neural_lam/models/ar_model.py:
    • Added self.edge_index to ARModel to store graph connectivity.
    • Updated training_step, validation_step, and test_step to pass edge_index to the loss function.
  • neural_lam/models/base_graph_model.py:
    • Automatically populates self.edge_index from the mesh graph (finest level for hierarchical models).
  • tests/test_lsd_training.py: Added tests for LSD training.
  • CHANGELOG.md: Added entry for generalized LSD support.

Dependencies

  • No new external dependencies (uses standard torch sparse operations and torch.fft available in modern PyTorch).

Issue Link

Solves #506

Type of change

  • 🐛 Bug fix (non-breaking change that fixes an issue)
  • ✨ New feature (non-breaking change that adds functionality)
  • 💥 Breaking change (fix or feature that would cause existing functionality to not work as expected)
  • 📖 Documentation (Addition or improvements to documentation)

Checklist before requesting a review

  • My branch is up-to-date with the target branch
  • I have performed a self-review of my code
  • For any new/modified functions/classes I have added docstrings that clearly describe its purpose, expected inputs and returned values
  • I have placed in-line comments to clarify the intent of any hard-to-understand passages of my code
  • I have updated the README (N/A - metric is internal to metrics.py and CLI)
  • I have added tests that prove my fix is effective or that my feature works
  • I have given the PR a name that clearly describes the change, written in imperative form
  • I have requested a reviewer and an assignee

Checklist for reviewers

Each PR comes with its own improvements and flaws. The reviewer should check the following:

  • the code is readable
  • the code is well tested
  • the code is documented (including return types and parameters)
  • the code is easy to maintain

Author checklist after completed review

  • I have added a line to the CHANGELOG describing this change, in a section reflecting type of change (add section where missing):
    • added: when you have added new functionality
    • changed: when default behaviour of the code has been changed
    • fixes: when your contribution fixes a bug
    • maintenance: when your contribution is relates to repo maintenance, e.g. CI/CD or documentation

Checklist for assignee

  • PR is up to date with the base branch
  • the tests pass
  • (if the PR is not just maintenance/bugfix) the PR is assigned to the next milestone. If it is not, propose it for a future milestone.
  • author has added an entry to the changelog (and designated the change as added, changed, fixed or maintenance)
  • Once the PR is ready to be merged, squash commits and merge the PR.

@sohampatil01-svg
Copy link
Copy Markdown
Author

@observingClouds Please have a Look at it !

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[ENHANCEMENT] Add Log-Spectral Distance (LSD) metric for spatial frequency analysis

1 participant