Factor out iterative refinement function

Author: matt-graham

s2fft/transforms/spherical.py

@@ -10,6 +10,7 @@
 from s2fft.transforms import otf_recursions as otf
 from s2fft.utils import healpix_ffts as hp
 from s2fft.utils import (
+    iterative_refinement,
     quadrature,
     quadrature_jax,
     resampling,
@@ -435,12 +436,12 @@ def forward(
             forward_kwargs = common_kwargs
             inverse_kwargs = {**common_kwargs, "method": "numpy"}
             forward_function = forward_numpy
-        flm = forward_function(f, **forward_kwargs)
-        for _ in range(iter):
-            f_recov = inverse(flm, **inverse_kwargs)
-            f_error = f - f_recov
-            flm += forward_function(f_error, **forward_kwargs)
-        return flm
+        return iterative_refinement.forward_with_iterative_refinement(
+            f=f,
+            n_iter=iter,
+            forward_function=partial(forward_function, **forward_kwargs),
+            backward_function=partial(inverse, **inverse_kwargs),
+        )
     elif method == "jax_ssht":
         if sampling.lower() == "healpix":
             raise ValueError("SSHT does not support healpix sampling.")

s2fft/utils/__init__.py

@@ -1,5 +1,6 @@
 from . import (
     healpix_ffts,
+    iterative_refinement,
     jax_primitive,
     quadrature,
     quadrature_jax,

s2fft/utils/iterative_refinement.py

@@ -0,0 +1,46 @@
+"""Iterative scheme for improving accuracy of linear transforms."""
+
+from collections.abc import Callable
+from typing import TypeVar
+
+T = TypeVar("T")
+
+
+def forward_with_iterative_refinement(
+    f: T,
+    n_iter: int,
+    forward_function: Callable[[T], T],
+    backward_function: Callable[[T], T],
+) -> T:
+    """
+    Apply forward transform with iterative refinement to improve accuracy.
+
+    `Iterative refinement <https://en.wikipedia.org/wiki/Iterative_refinement>`_ is a
+    general approach for improving the accuracy of numerial solutions to linear systems.
+    In the context of spherical harmonic transforms, given a forward transform which is
+    an _approximate_ inverse to a corresponding backward ('inverse') transform,
+    iterative refinement allows defining an iterative forward transform which is a more
+    accurate
+
+    Args:
+        f: Array argument to forward transform (signal on sphere) to compute iteratively
+           refined forward transform at.
+
+        n_iter: Number of refinement iterations to use, non-negative.
+
+        forward_function: Function computing forward transform (approximate inverse of
+            backward transform).
+
+        backward_function: Function computing backward ('inverse') transform.
+
+    Returns:
+        Array output from iteratively refined forward transform (spherical harmonic
+        coefficients).
+
+    """
+    flm = forward_function(f)
+    for _ in range(n_iter):
+        f_recov = backward_function(flm)
+        f_error = f - f_recov
+        flm += forward_function(f_error)
+    return flm