Radial Extension for SFNO by tmarkmann · Pull Request #242 · NVIDIA/torch-harmonics

tmarkmann · 2026-08-15T16:17:55Z

Summary

This PR builds towards an extension of the current SFNO to 3D systems including a radial domain. Planned additions are new PDE examples, a Mellin Convolution and a new Radial Spherical Neural Operator Model. For now this PR only includes an addition of the spherical Poisson Equation for the half-line and exterior domain. More updates will be added to the PR in the future.

Roadmap

Poisson Equation
Mellin Convolution
Radial SFNO Model

Main Additions and Changes

torch_harmonics/examples/poisson_equation.py: Added Poisson Solver
torch_harmonics/examples/poisson_dataset.py: Added Dataset for training
torch_harmonics/quadrature.py: Added geometric weights
notebooks/poisson_equation.ipynb: Examples for Poisson

…rch-harmonics into tmarkmann/spherical-poisson

tmarkmann · 2026-08-15T16:22:49Z

@bonevbs

einops is not declared in pyproject.toml and is not used anywhere else in the repository. Since examples/__init__.py imports poisson_equation eagerly, the import broke every existing consumer of torch_harmonics.examples (the shallow water / depth / segmentation training scripts, the helmholtz and SFNO notebooks, Dockerfile.examples) for anyone without einops installed. The CI test suite does not import examples, so this would not have been caught there. All rearrange calls were pure axis insertions, so they are replaced by unsqueeze/reshape with no permutation or layout change. The duplicated logr broadcasts in _assemble are hoisted into locals and shape comments are added where the einops pattern strings previously carried that information. Output is bit-for-bit identical: the green and blift buffers, random_source, and solve all compare equal on both domains. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Replaces the hand-rolled if/elif over legendre_gauss_weights / lobatto_weights / clenshaw_curtiss_weights with the existing quadrature utility. The block was copied from pde_sphere.py and had no else branch; it was not reachable with an undefined `cost` because RealSHT is constructed first and already rejects unknown grids, but the duplication is unnecessary. precompute_latitudes returns colatitudes ordered from the north pole, so lats = pi/2 - colats reproduces the previous -arcsin(cost) values to within roundoff (<= 3.3e-16 on all three supported grids, both parities of nlat) with identical ordering. lats is only consumed by the plotting helpers, so solutions are unaffected. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Drops code paths that nothing reaches today. Each is small and self-contained, so it can come back with the work that needs it (Mellin convolution, radial SFNO, Newtonian form) rather than sitting unused in the meantime. - radial_grid: the "shell" branch. GreensOperator raises NotImplementedError for any domain but half-line/exterior, and RadialPoissonSolver passes its own domain straight through, so no caller could reach it. This also drops the now-unused trapezoidal_weights import. - RadialPoissonSolver.solve_density and the gravity buffer/ctor argument. Nothing in the repo or the notebook calls it; it is a constant prefactor on solve() that a caller can apply directly. - plot_sphere: the "orthographic" branch. The notebook only uses "mollweide" (and "log"/"polar" for plot_meridional), and this was the sole reference to cartopy, which is not a declared dependency of the package. Documentation, in the same pass: - Add geometric_weights to the quadrature autosummary in docs/api/utilities.md with a note that it is a radial rule, not a latitudinal one, so it does not belong in the grid-string table above it. - State the dense (lmax, nr, nr) buffer size and per-solve cost on RadialPoissonSolver; PoissonDataset calls solve() once per sample. - Correct PoissonDataset.scale, which uses an unweighted angular mean rather than the sphere measure the previous docstring claimed. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

torch_harmonics/ contains no bare assert statements since the v0.9.2 cleanup that converted them to torch._check and ValueError; this was the only one the branch reintroduced. Matches the style of plotting.py, which raises ValueError for an unknown projection, and formats the shape as a tuple so the message reads (64, 16, 32) rather than torch.Size([64, 16, 32]). Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Every other notebook in notebooks/ is symlinked into docs/tutorials/ and listed in its toctree; poisson_equation.ipynb was the only one left out. Filed under Applications next to helmholtz and shallow_water_equations, the other PDE solver tutorials. Verified with a local `sphinx -b html . build -W --keep-going`: the tutorial page and the geometric_weights autosummary stub both render, and the build is warning-free against this branch. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

azrael417 · 2026-08-17T09:42:18Z

Hello Thorben — nice work, the solver is mathematically sound. Thanks for the MR. With the help of Claude, I pushed some changes to streamline it with the rest of the repo:

I validated the Green's operator against analytic solutions before touching anything: the monopole case u = exp(-ar²) converges cleanly at second order (rel L∞ 2.97e-4 → 7.39e-5 → 1.85e-5 for nr = 256/512/1024), the l=1 case lands at 3.2e-4, and the exterior harmonic lift is exact to 6.6e-15. The exterior image term makes the kernel vanish identically at r = R, so the homogeneous Dirichlet condition is exact rather than approximate.

I pushed five commits directly to the branch — happy to revert any of them if you disagree. Summary:

Removed the einops dependency (45a64bf)

This was the one blocker. einops isn't in pyproject.toml and isn't used anywhere else in the repo, and since examples/init.py imports poisson_equation eagerly, it broke every existing consumer of torch_harmonics.examples for anyone without it installed — the shallow water / depth / segmentation training scripts, the helmholtz and SFNO notebooks, Dockerfile.examples. CI wouldn't have caught it, because the test suite never imports examples.

All nine rearrange calls were pure axis insertions (no permutes), so they became unsqueeze/reshape. Verified bit-for-bit identical: the green and blift buffers, random_source, and solve all compare equal on both domains.

Reused precompute_latitudes (a19f51b)

The grid if/elif was copied from pde_sphere.py; the utility already does this and handles the unknown-grid case. lats = π/2 - colats reproduces the previous -arcsin(cost) values to ≤3.3e-16 on all three grids with identical ordering, and lats only feeds the plotting helpers, so solutions are untouched.

Removed currently-unreachable code (eb8e63e)

Since the Mellin convolution and radial SFNO are still ahead, these seemed better added alongside the work that needs them:

radial_grid's "shell" branch — GreensOperator rejects any domain but half-line/exterior, and the solver passes its own domain straight through, so nothing could reach it.
solve_density and the gravity buffer — no callers; it's a constant prefactor on solve().
plot_sphere's "orthographic" branch — the notebook only uses "mollweide", and this was the sole cartopy reference (also not a declared dependency).

Same commit documents geometric_weights in docs/api/utilities.md, notes the dense (lmax, nr, nr) buffer size (33 MB at nlat=64/nr=256, but 537 MB at nlat=256/nr=512, and PoissonDataset runs one solve per sample), and corrects the scale docstring — it uses an unweighted angular mean, not the sphere measure.

ValueError instead of assert in plot_sphere (12f9e1b)

torch_harmonics/ has had zero bare asserts since the v0.9.2 conversion, so this was the only one; matches plotting.py's style.

Wired the notebook into the docs (40e2275)

All 13 other notebooks are symlinked into docs/tutorials/ and in the toctree; this was the only one missing. Filed under Applications next to helmholtz. Confirmed with a local sphinx -b html . build -W --keep-going that the tutorial page and the geometric_weights stub both render warning-free.

"Bounds of the radial grid" is only true on the half-line. On the exterior domain radial_grid maps r = R + R * rho and r_min/r_max bound the reduced coordinate rho / R = (r - R) / R, so with R = 1 and the default r_min = 1e-2 the innermost node sits at r = 1.01, not 0.01. The defaults also differ by domain, (1e-1, 1e3) against (1e-2, 1e2), and were not stated at all. radial_grid already documents its vmin/vmax this way; the solver's docstring did not inherit it. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

- solve: document f and v0. v0 is the whole exterior-boundary feature and was undocumented; it is Dirichlet data on r = R, lifted harmonically as (R / r)**(l + 1) per degree, and homogeneous when omitted. Verified against the l = 0 and l = 1 lifts. - plot_meridional: name the supported projections ("log" and "polar") the way plot_sphere names "mollweide", and document ilon and rmax. Notes that "polar" draws the antipodal meridian too and wants a plain, non-polar axes. - PoissonDataset: list num_examples/device/normalize last, matching the constructor signature and the ordering PdeDataset already uses. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

CONTRIBUTING.md asks for a Changelog.md entry on user-visible changes, and this PR adds public API in two places: RadialPoissonSolver / PoissonDataset exported from torch_harmonics.examples, and geometric_weights in the core quadrature module. Opens a v0.9.3 section, since v0.9.2 is already tagged. This mirrors how the v0.9.2 heading was opened right after v0.9.1 shipped, with entries then accumulating under it per PR. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

geometric_weights lands in the core quadrature module rather than in examples, so it belongs in tests/test_quadrature.py alongside the other rules. No test is added for the Poisson solver itself, since no example in the repo is covered by the suite today. The central case is f(x) = 1/x. The nodes are equispaced in t = log(x), so f dx/dt = 1 is constant and the trapezoidal rule integrates it exactly at any n: sum(w / x) must equal log(b / a) independent of resolution. That pins down both the node placement and the Jacobian carried by the weights. Supporting cases cover the endpoints and constant node ratio, second-order convergence for the non-exact f(x) = 1, and the non-positive lower bound. Tolerances are 1e-6 rather than machine epsilon because trapezoidal_weights builds its weights with torch.ones(n), which is float32 while the nodes are float64. That caps accuracy at ~1e-7; a float64 reference of the same rule hits 3.6e-15. Pre-existing on main and not addressed here. The bound check matches the message, not just the type: math.log raises ValueError for these inputs on its own, so a bare assertRaises would still pass if the explicit check were removed. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

azrael417

LGTM now, I applied some changes. Please have a look if they make sense. Then I recommend rebasing your code onto main once this MR lands.

tmarkmann · 2026-08-17T15:08:02Z

Hey Thorsten, thank you for your review! All of your points are fine for me and I will adapt future code to be more in line with the lib. Here are my thoughts on your points:

I think einops makes the code much more understandable but if it's not used throughout the library I will switch to reshapes from now to not break the library.
I copied that from examples/shallow_water_equations.py (and pde_sphere.py) so maybe it should be changed there as well.
I'm using some of the unused code in a downstream repository but I can readd those when they are actually used from torch-harmonics. One note about the cartopy plotting reference: I basically copied that from the examples/shallow_water_equations.py as well. Since it is a lazy import I think that should be fine and I can add an orthographic plot to the notebook (so it is called already).

One general question from my side is whether we wait until all the features are implemented and pushed to this PR or merge the smaller bits. I'm fine with both but @bonevbs might have a preference from earlier discussions.

azrael417 · 2026-08-18T04:37:59Z

Hello Thorben, thanks for the quick reply. We are trying to keep the numbers of package dependencies down as much as possible. Einops can be OK as long as it is contained in the examples, since those have their own set of dependences anyway, but most of the ops used were just permutes and reshapes, which I prefer doing in torch directly. Afair, I have removed all the einops from torch-harmonics otherwise, please let me know if that is not correct.

I have removed some of the unused code so that it can be added later in another targeted MR. So this MR adds the poisson equation and solver and some plumbing and a notebook. The next can then add the additional features like the GreensFunction and other classes. Does that make sense? Claude suggested to split it there. We will certainly add it back when you make a follow-up MR. This is part 1 of 3, right?

bonevbs · 2026-08-18T07:08:47Z

    return xlg, wlg


+def geometric_weights(n: int, a: Optional[float] = 1.0, b: Optional[float] = math.e) -> Tuple[torch.Tensor, torch.Tensor]:


can we rename this to radial weights? maybe even better to name them 'log_uniform_weights' ? given that they are uniform in logarithmic domain

also is there a good reason to name the parameters a and b? I would potentially name them inner_radius and log_base

I agree that 'log_uniform_weights' is a much better name. The parameters were called a and b to mirror 'trapezoidal_weights'. I think there are two options:

'log_uniform_weights' and stick to a and b

'radial_weights' and change to inner_radius and log_base

I would go for the first one if we ever add a shell as an option. I suppose that we could have uniform radial weights in that case and this could lead to confusion if we already reserve 'radial_weights'.

I am fine with 1, that sounds good. But make sure the doctoring describes what the parameters do. If we change the naming of the parameters here, we also should do that for the trapezoidal weights imo. Boris, what would be good names there? Then we can fix these issues in one go.

bonevbs

Hi @tmarkmann most of this already looks good, thank you! I left some comments. The bigger one is of this couldn't just reuse the PDEDataset as this would make it easier to just reuse existing training pipelines in the examples folder.

bonevbs · 2026-08-18T07:09:29Z

    return xlg, wlg


+def geometric_weights(n: int, a: Optional[float] = 1.0, b: Optional[float] = math.e) -> Tuple[torch.Tensor, torch.Tensor]:


also is there a good reason to name the parameters a and b? I would potentially name them inner_radius and log_base

bonevbs · 2026-08-18T07:10:39Z

+from .poisson_equation import RadialPoissonSolver
+
+
+class PoissonDataset(torch.utils.data.Dataset):


torch-harmonics has a PDEDataset object. Couldn't we have reused that? or derived from that?

I did not extend the PDEDataset because I did not want to break existing code and some things are different for the poisson equation:

radial dimension: we need to either have a 2D or 3D dims object depending on the pde

domain: half-line or exterior which only is important for poisson

initial_condition: there is only random for poisson

the normalization for poisson needs to be separate for target and input (different to swe)

My first thought was that caring about all the divergences between SWE and Poisson is a bit hacky. But for the example poisson I can just fix all parameters like domain and rmin, rmax and then it should be ok. I will fix that!

Maybe we should rename the PDE dataset to SWEDataset? Since that is more descriptive? Boris, what do you think?

I do not mind having a separate class per example, in that way each example is self contained, up to the shared content from torch harmonics that is. But users can grab the example and their TH install and run it, no need to use classes or functions from an example helper folder or something like that.

bonevbs · 2026-08-18T07:15:03Z

For the einops, while I do agree it might make it more readable, we are usually trying to keep dependencies lightweight. So I prefer removing it unless it becomes absolutely essential.

tmarkmann · 2026-08-19T00:48:35Z

So this MR adds the poisson equation and solver and some plumbing and a notebook. The next can then add the additional features like the GreensFunction and other classes. Does that make sense? Claude suggested to split it there. We will certainly add it back when you make a follow-up MR. This is part 1 of 3, right?

Yes, this is part 1 of 3. Since the other two parts will have larger changes I agree to split it here. I will implement @bonevbs changes later today and then it is ready from my side.

tmarkmann added 10 commits July 26, 2026 23:48

Add drafts for poisson and heat equation

663790d

Add Poisson solver and notebook

10850a3

Remove heat equation

15be043

Add drafts for poisson and heat equation

0bc2c0f

Add Poisson solver and notebook

d4cfef6

Remove heat equation

cae38ea

Merge branch 'tmarkmann/spherical-poisson' of github.com:tmarkmann/to…

b5d8b71

…rch-harmonics into tmarkmann/spherical-poisson

Revert wrongfully changes

1e35c71

Cosmetic changes

314cfac

Run precommit

5dc37b9

tmarkmann marked this pull request as ready for review August 15, 2026 16:22

azrael417 and others added 5 commits August 17, 2026 11:02

azrael417 and others added 4 commits August 17, 2026 14:30

azrael417 previously approved these changes Aug 17, 2026

View reviewed changes

azrael417 mentioned this pull request Aug 17, 2026

Return float64 weights from trapezoidal_weights #243

Merged

Merge branch 'main' into tmarkmann/spherical-poisson

5544661

azrael417 dismissed their stale review via 5544661 August 18, 2026 07:07

bonevbs reviewed Aug 18, 2026

View reviewed changes

		return xlg, wlg


		def geometric_weights(n: int, a: Optional[float] = 1.0, b: Optional[float] = math.e) -> Tuple[torch.Tensor, torch.Tensor]:

		from .poisson_equation import RadialPoissonSolver


		class PoissonDataset(torch.utils.data.Dataset):

Uh oh!

Conversation

tmarkmann commented Aug 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Roadmap

Main Additions and Changes

Uh oh!

tmarkmann commented Aug 15, 2026

Uh oh!

azrael417 commented Aug 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

azrael417 left a comment

Choose a reason for hiding this comment

Uh oh!

tmarkmann commented Aug 17, 2026

Uh oh!

azrael417 commented Aug 18, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

azrael417 Aug 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

bonevbs left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

bonevbs commented Aug 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

tmarkmann commented Aug 19, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

tmarkmann commented Aug 15, 2026 •

edited

Loading

azrael417 commented Aug 17, 2026 •

edited

Loading

azrael417 Aug 19, 2026 •

edited

Loading

bonevbs commented Aug 18, 2026 •

edited

Loading