feat: octhedral spherical harmonic transform - ectrans-based

[theissenhelen]

Description

This PR implements a spherical harmonic transform for octahedral reduced gaussian grids based on conventions used in ectrans.

What problem does this change solve?

What issue or task does this change relate to?

Part 1 of #599

Additional notes

As a contributor to the Anemoi framework, please ensure that your changes include unit tests, updates to any affected dependencies and documentation, and have been tested in a parallel setting (i.e., with multiple GPUs). As a reviewer, you are also responsible for verifying these aspects and requesting changes if they are not adequately addressed. For guidelines about those please refer to https://anemoi.readthedocs.io/en/latest/

By opening this pull request, I affirm that all authors agree to the Contributor License Agreement.

[OpheliaMiralles]

Merged in #678

[anaprietonem]

AS mentioned in #729 - So that's it's easier to work on the harmonisation of the Spectra transforms for Octahedral grids, I've reopened the two branches.

[samhatfield]

By the way, @sahahner and I are continuing to develop EcTransOctahedralSHT in a separate private repo. We're doing this because that repo has a very useful test suite which we haven't thought about how to integrate into anemoi-core yet. Whenever we reach a stable state in that repo we will update this one. Ideally we'd figure out how to get those tests into anemoi-core so all development can take place in this PR.

[samhatfield]

This PR now contains the state-of-the-art for the ecTrans-based octahedral SHT, and I've adapted it to PR #788's structure.

[samhatfield]

Dependency on ectrans4py removed. Now "Legendre assets" (Legendre polynomials, Gaussian weights, and "max wavenumber per latitude" arrays) must be precomputed and stored to disk, and the path to this location must be given when creating an EcTransOctahedralSHTModule object.

This script can be used to generate the assets:

generate_ectrans_assets.py

# (C) Copyright 2025 Anemoi contributors.
#
# This software is licensed under the terms of the Apache Licence Version 2.0
# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
#
# In applying this licence, ECMWF does not waive the privileges and immunities
# granted to it by virtue of its status as an intergovernmental organisation
# nor does it submit to any jurisdiction.

import argparse
import numpy as np
from pathlib import Path


def parse_grid(grid: str):
    # Octahedral grid
    if grid[0] == "o" or grid[0] == "O":
        n_lat_nh = int(grid[1:])
        truncation = n_lat_nh - 1
        lons_per_lat = [20 + 4 * i for i in range(n_lat_nh)]
        lons_per_lat += lons_per_lat[::-1]
        return n_lat_nh, truncation, lons_per_lat
    # TODO: support N grids too
    else:
        raise ValueError(f"Unsupported grid type: {grid}")


def generate_and_save(grid, num_grids, dir_path: str):
    """
    Fetch relevant arrays from ecTrans
    Note that all of these arrays (including the input points-per-latitude array) are
    specified across the full globe, pole to pole

    Parameters
    ----------
    grid : str
        grid specification as a str, e.g. "o16"
    num_grids : int
        total number of grids that will be computed in this script
        necessary because ectrans4py is a bit dumb and needs to know everything in advance
    dir_path : str
        directory to save generated assets to
    """

    print(f"Generating SHT assets for grid {grid}")

    try:
        import ectrans4py  # type: ignore
    except Exception as exc:  # pragma: no cover
        msg = (
            "ectrans4py is required to generate octahedral SHT assets. "
        )
        raise ModuleNotFoundError(msg) from exc

    file_path = dir_path + f"/sht_assets_{grid}.npz"

    n_lat_nh, truncation, lons_per_lat = parse_grid(grid)

    poly_size = sum(truncation + 2 - im for im in range(truncation + 1))

    (highest_zonal_wavenumber_per_lat, gaussian_weights, legendre_polynomials) = ectrans4py.get_legendre_assets(
        2 * n_lat_nh,
        truncation,
        2 * n_lat_nh,
        poly_size,
        np.array(lons_per_lat),
        num_grids
    )

    np.savez(
        file_path,
        legendre_polynomials=legendre_polynomials,
        gaussian_weights=gaussian_weights,
        highest_zonal_wavenumber_per_lat=highest_zonal_wavenumber_per_lat,
    )


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description=("Generate assets required to perform spherical harmonic transform with "
                     "EcTransOctahedralSHTModule")
    )
    parser.add_argument("output_dir", help="Directory to save generated assets to") 
    args = parser.parse_args()

    dir_path = args.output_dir

    if not Path(dir_path).exists():
        raise FileNotFoundError(f"Output directory {dir_path} does not exist")

    grids = ["o16", "o32", "o48", "o96", "o160", "o256", "o320", "o400", "o1280", "n320"]

    for grid in grids:
        generate_and_save(grid, len(grids), dir_path)