feat(models): Add wrapper class for InverseReducedSHT

[samhatfield]

Description

PR #1106 added wrapper classes for inverse spherical harmonic transforms on a regular grid and octahedral grid. However no class was added for a "classic" reduced Gaussian grid, such as N320. This PR completes the set by adding InverseReducedSHT. This came up in discussions with @sahahner.

Note that there is no test suite for the wrapper classes, ReducedSHT, OctahedralSHT, InverseReducedSHT etc. Instead the test suite operates at the "helper class" level (i.e. the underlying grid-agnostic classes SphericalHarmonicTransform and InverseSphericalHarmonicTransform are tested). In future I am planning to refactor the test suite so it targets the wrapper classes, instead. Since the wrapper depends on the helper, both will be tested, rather than just the latter at the moment.

In the meantime, I tested this PR through these two tests which I didn't commit:

Test suite

@pytest.mark.parametrize("sht_setup", ["reduced"], indirect=True)
def test_idempotency_inverse_direct_reduced(sht_setup):
    """Test that the reduced SHT and its inverse are consistent."""

    import einops

    truncation = sht_setup["truncation"]
    dtype = sht_setup["dtype"]
    tolerance = sht_setup["tolerance"]

    device = "cuda" if torch.cuda.is_available() else "cpu"

    # Create reduced SHT and its inverse with the same truncation
    direct_reduced = ReducedSHT(grid="n320", truncation=truncation).to(device)
    inverse_reduced = InverseReducedSHT(grid="n320", truncation=truncation).to(device)

    before_spectral = random_spectral_array(truncation, dtype)

    gridded = einops.rearrange(inverse_reduced(before_spectral), 'b e p -> b 1 e p 1')
    after_spectral = einops.rearrange(direct_reduced(gridded), 'b 1 e l m 1 -> b e l m')

    assert torch.allclose(before_spectral, after_spectral, rtol=tolerance)


@pytest.mark.parametrize("sht_setup", ["reduced"], indirect=True)
def test_idempotency_direct_inverse_reduced(sht_setup):
    """Test that the reduced SHT and its inverse are consistent."""

    import einops

    truncation = sht_setup["truncation"]
    dtype = sht_setup["dtype"]
    tolerance = sht_setup["tolerance"]

    device = "cuda" if torch.cuda.is_available() else "cpu"

    # Create reduced SHT and its inverse with the same truncation
    direct_reduced = ReducedSHT(grid="n320", truncation=truncation).to(device)
    inverse_reduced = InverseReducedSHT(grid="n320", truncation=truncation).to(device)

    # Ensure the direct input is band-limited by constructing it via inverse.
    before_spectral = random_spectral_array(truncation, dtype)
    before_gridded = einops.rearrange(inverse_reduced(before_spectral), 'b e p -> b 1 e p 1')

    spectral = einops.rearrange(direct_reduced(before_gridded), 'b 1 e l m 1 -> b e l m')
    after_gridded = einops.rearrange(inverse_reduced(spectral), 'b e p -> b 1 e p 1')

    assert torch.allclose(before_gridded, after_gridded, rtol=tolerance)

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.