Radial Extension for SFNO

.gitignore

@@ -1,6 +1,7 @@
 *.DS_Store
 __pycache__
 *.so
+*.egg-info/
 checkpoints
 
 # Build artifacts

Changelog.md

@@ -4,6 +4,8 @@
 
 ### v0.9.3
 
+* Added a radial Poisson solver example (`RadialPoissonSolver` and `PoissonDataset` in `torch_harmonics.examples`), solving `lap u = f` on `(0, inf) x S2` and `[R, inf) x S2` subject to `u -> 0` at infinity. It combines an exact per-degree Green's operator in the radial direction with a spherical harmonic transform in the angular directions, and supports Dirichlet data on the inner sphere of the exterior domain. New tutorial notebook: `notebooks/poisson_equation.ipynb`.
+* Added `geometric_weights` to `torch_harmonics.quadrature`: geometrically spaced nodes, uniform in `log(x)`, together with the corresponding trapezoidal weights for the integral over `dx` on a positive interval. Intended for radial directions spanning several decades.
 * Fixed `trapezoidal_weights` returning float32 weights alongside float64 nodes, because the underlying `torch.ones(n)` inherited the default dtype. It was the only rule in `torch_harmonics.quadrature` doing so, and it capped the accuracy of everything derived from it at roughly 1e-7, including the latitude weights of the `"equiangular-trapezoidal"` grid. Weights are now float64 like the other rules, and the tolerances of the affected tests have been tightened to match the other grids.
 
 ### v0.9.2

docs/api/utilities.md

@@ -18,6 +18,11 @@ convolution layers:
 The longitudinal direction always uses equispaced nodes (see
 `precompute_longitudes`).
 
+`geometric_weights` is not a latitudinal rule: it returns nodes that are
+equispaced in $\log x$ on a positive interval, together with the corresponding
+trapezoidal weights for $\int f \, \mathrm{d}x$. It is intended for radial
+directions spanning several decades.
+
 ```{eval-rst}
 .. currentmodule:: torch_harmonics.quadrature
 
@@ -31,6 +36,7 @@ The longitudinal direction always uses equispaced nodes (see
    lobatto_weights
    clenshaw_curtiss_weights
    trapezoidal_weights
+   geometric_weights
 ```
 
 ## Plotting

docs/tutorials/index.md

@@ -23,6 +23,7 @@ maxdepth: 1
 caption: Applications
 ---
 helmholtz
+poisson_equation
 shallow_water_equations
 train_spherical_neural_operator
 stanford_2d3ds_dataset

docs/tutorials/poisson_equation.ipynb

@@ -0,0 +1 @@
+../../notebooks/poisson_equation.ipynb

tests/test_quadrature.py

@@ -37,7 +37,8 @@
 from testutils import compare_tensors, set_seed
 
 import torch_harmonics as th
-from torch_harmonics.quadrature import precompute_latitudes, precompute_longitudes, trapezoidal_weights
+from torch_harmonics.quadrature import geometric_weights, precompute_latitudes, precompute_longitudes, trapezoidal_weights
+
 
 _devices = [(torch.device("cpu"),)]
 if torch.cuda.is_available():
@@ -205,6 +206,76 @@ def test_normalization_consistency(self, nlat, nlon, batch_size, num_chan, grid,
         )
 
 
+class TestGeometricWeights(unittest.TestCase):
+    """Geometrically spaced quadrature nodes and weights on a positive interval."""
+
+    @parameterized.expand(
+        [
+            # n, a, b
+            [8, 1e-3, 1e3],
+            [64, 1e-3, 1e3],
+            [512, 1e-3, 1e3],
+            [33, 1.0, 2.0],
+            [65, 1e-6, 1.0],
+        ]
+    )
+    def test_inverse_integral(self, n, a, b, verbose=False):
+        """The rule is exact for f(x) = 1/x, whose integral over [a, b] is log(b / a).
+
+        The nodes are equispaced in t = log(x), so f dx/dt = 1 is constant and the
+        trapezoidal rule integrates it without discretization error at any n. This
+        pins down both the node placement and the dx/dt Jacobian carried by the
+        weights: dropping it turns the result into something n-dependent.
+        """
+
+        x, w = geometric_weights(n, a, b)
+
+        integral = (w / x).sum()
+        expected = torch.as_tensor(math.log(b / a), dtype=integral.dtype)
+
+        self.assertTrue(compare_tensors("inverse integral", integral, expected, atol=1e-6, rtol=1e-6, verbose=verbose))
+
+    @parameterized.expand(
+        [
+            # n, a, b
+            [16, 1e-3, 1e3],
+            [65, 1e-2, 1.0],
+        ]
+    )
+    def test_node_placement(self, n, a, b, verbose=False):
+        """Nodes span [a, b] and are geometrically spaced, i.e. a constant ratio apart."""
+
+        x, w = geometric_weights(n, a, b)
+
+        self.assertEqual(x.shape, (n,))
+        self.assertEqual(w.shape, (n,))
+        self.assertTrue(compare_tensors("endpoints", x[[0, -1]], torch.as_tensor([a, b], dtype=x.dtype), atol=1e-12, rtol=1e-12, verbose=verbose))
+
+        ratio = x[1:] / x[:-1]
+        expected = torch.full_like(ratio, (b / a) ** (1.0 / (n - 1)))
+        self.assertTrue(compare_tensors("node ratio", ratio, expected, atol=1e-12, rtol=1e-12, verbose=verbose))
+        self.assertTrue(torch.all(w > 0.0))
+
+    def test_convergence(self, verbose=False):
+        """For f(x) = 1, which is not exact, the error decays at second order in h = log(b / a) / (n - 1)."""
+
+        a, b = 1.0, 10.0
+        errors = [abs(geometric_weights(n, a, b)[1].sum().item() - (b - a)) for n in (64, 128, 256)]
+
+        for coarse, fine in zip(errors[:-1], errors[1:]):
+            self.assertGreater(coarse / fine, 3.5)
+
+    def test_invalid_bounds(self):
+        """A geometric grid is undefined for a non-positive lower bound.
+
+        Matched on the message rather than the type: math.log raises ValueError for
+        these inputs by itself, so a bare assertRaises would also pass if the explicit
+        bound check were removed.
+        """
+
+        for a in (0.0, -1.0):
+            with self.assertRaisesRegex(ValueError, "must be positive"):
+                geometric_weights(8, a, 10.0)
 class TestQuadratureWeightPrecision(unittest.TestCase):
     """Every quadrature rule must carry its weights in the same precision as its nodes."""
 

torch_harmonics/examples/__init__.py

@@ -31,5 +31,7 @@
 
 from .pde_dataset import PdeDataset
 from .pde_sphere import SphereSolver
+from .poisson_dataset import PoissonDataset
+from .poisson_equation import RadialPoissonSolver
 from .shallow_water_equations import ShallowWaterSolver
 from .stanford_2d3ds_dataset import Stanford2D3DSDownloader, StanfordDatasetSubset, StanfordDepthDataset, StanfordSegmentationDataset, compute_stats_s2

torch_harmonics/examples/pde_sphere.py

@@ -77,7 +77,7 @@ def __init__(self, nlat, nlon, dt, lmax=None, mmax=None, grid="equiangular", rad
         self.nlon = nlon
         self.grid = grid
 
-        # physical sonstants
+        # physical constants
         self.register_buffer("radius", torch.as_tensor(radius, dtype=torch.float64))
         self.register_buffer("coeff", torch.as_tensor(coeff, dtype=torch.float64))
 

torch_harmonics/examples/poisson_dataset.py

@@ -0,0 +1,133 @@
+# coding=utf-8
+
+# SPDX-FileCopyrightText: Copyright (c) 2022 The torch-harmonics Authors. All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# 1. Redistributions of source code must retain the above copyright notice, this
+# list of conditions and the following disclaimer.
+#
+# 2. Redistributions in binary form must reproduce the above copyright notice,
+# this list of conditions and the following disclaimer in the documentation
+# and/or other materials provided with the distribution.
+#
+# 3. Neither the name of the copyright holder nor the names of its
+# contributors may be used to endorse or promote products derived from
+# this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+
+
+import torch
+
+from .poisson_equation import RadialPoissonSolver
+
+
+class PoissonDataset(torch.utils.data.Dataset):
+    """Custom Dataset class for Poisson training data
+
+    Parameters
+    ----------
+    dims : tuple, optional
+        Number of latitude, longitude and radial points, by default (64, 128, 256)
+    grid : str, optional
+        Angular grid type, by default "legendre-gauss"
+    domain : str, optional
+        Either "half-line" or "exterior", by default "half-line"
+    R : float, optional
+        Inner radius for exterior domain, by default None
+    nblobs : int or tuple of int, optional
+        Number of blobs in each source, by default (1, 8)
+    l_src : int, optional
+        Angular band limit of the source, by default 8
+    positive : bool, optional
+        Draw only positive sources, by default False
+    num_examples : int, optional
+        Number of examples, by default 32
+    device : torch.device, optional
+        Device to use, by default torch.device("cpu")
+    normalize : bool, optional
+        Whether to normalize the input and target, by default True
+
+    Returns
+    -------
+    inp : torch.Tensor
+        Source, shape (nr, nlat, nlon)
+    tar : torch.Tensor
+        Solution, shape (nr, nlat, nlon)
+    """
+
+    def __init__(
+        self,
+        dims=(64, 128, 256),
+        grid="legendre-gauss",
+        domain="half-line",
+        R=None,
+        nblobs=(1, 8),
+        l_src=8,
+        positive=False,
+        num_examples=32,
+        device=torch.device("cpu"),
+        normalize=True,
+    ):
+        self.num_examples = num_examples
+        self.device = device
+        self.normalize = normalize
+        self.nblobs = nblobs
+        self.l_src = l_src
+        self.positive = positive
+        self.nlat, self.nlon, self.nr = dims
+
+        self.solver = RadialPoissonSolver(
+            self.nlat,
+            self.nlon,
+            self.nr,
+            grid=grid,
+            domain=domain,
+            R=R,
+        ).to(self.device)
+
+    def __len__(self):
+        return self.num_examples
+
+    def _get_sample(self):
+        """Get one unscaled source + solution pair."""
+
+        f = self.solver.random_source(nblobs=self.nblobs, l_src=self.l_src, positive=self.positive)
+        u = self.solver.solve(f)
+
+        return f.float(), u.float()
+
+    def scale(self, f):
+        """
+        Scale factor for a pair: the L2 norm of the source, quadrature-weighted in the
+        radial direction by r**2 dr and approximated by an unweighted mean over the
+        angular directions.
+        """
+
+        w, r = self.solver.w, self.solver.r
+        return ((w * r**2) * (f**2).mean(dim=(-1, -2))).sum().sqrt()
+
+    def __getitem__(self, index):
+
+        with torch.inference_mode():
+            with torch.no_grad():
+                inp, tar = self._get_sample()
+
+                if self.normalize:
+                    s = self.scale(inp)
+                    inp, tar = inp / s, tar / s
+
+        return inp.clone(), tar.clone()