Add precision arg for NUFFTs (#122)

AlecThomson · web-flow · commit c25249a77452 · 2024-03-27T09:05:58.000+11:00
* Add eps as param

* Add tests

* Fix api calls
diff --git a/RMtools_1D/do_RMsynth_1D.py b/RMtools_1D/do_RMsynth_1D.py
@@ -318,7 +318,6 @@ def run_rmsynth(
         phiArr_radm2=phiArr_radm2,
         weightArr=weightArr,
         nBits=nBits,
-        verbose=verbose,
         log=log,
         lam0Sq_m2=0 if super_resolution else None,
     )
diff --git a/RMtools_3D/do_RMsynth_3D.py b/RMtools_3D/do_RMsynth_3D.py
@@ -203,7 +203,6 @@ def run_rmsynth(
             phiArr_radm2=phiArr_radm2,
             weightArr=weightArr,
             nBits=32,
-            verbose=verbose,
             lam0Sq_m2=0 if super_resolution else None,
         )
     else:
diff --git a/RMutils/util_RM.py b/RMutils/util_RM.py
@@ -94,7 +94,7 @@ def do_rmsynth_planes(
     weightArr=None,
     lam0Sq_m2=None,
     nBits=32,
-    verbose=False,
+    eps=1e-6,
     log=print,
 ):
     """Perform RM-synthesis on Stokes Q and U cubes (1,2 or 3D). This version
@@ -109,15 +109,15 @@ def do_rmsynth_planes(
     lambdaSqArr_m2  ... vector of wavelength^2 values (assending freq order)
     phiArr_radm2    ... vector of trial Faraday depth values
     weightArr       ... vector of weights, default [None] is Uniform (all 1s)
+    lam0Sq_m2       ... force a reference lambda^2 value [None, calculated]
     nBits           ... precision of data arrays [32]
-    verbose         ... print feedback during calculation [False]
+    eps             ... NUFFT tolerance [1e-6] (see https://finufft.readthedocs.io/en/latest/python.html#module-finufft)
     log             ... function to be used to output messages [print]
 
     """
 
     # Default data types
     dtFloat = "float" + str(nBits)
-    dtComplex = "complex" + str(2 * nBits)
 
     # Set the weight array
     if weightArr is None:
@@ -195,7 +195,7 @@ def do_rmsynth_planes(
             x=a,
             c=np.ascontiguousarray(pCube.T),
             s=(phiArr_radm2[::-1] * 2).astype(a.dtype),
-            eps=1e-8,
+            eps=eps,
         )
         * KArr[..., None]
     ).T
@@ -228,6 +228,7 @@ def get_rmsf_planes(
     fitRMSF=False,
     fitRMSFreal=False,
     nBits=32,
+    eps=1e-6,
     verbose=False,
     log=print,
 ):
@@ -250,6 +251,7 @@ def get_rmsf_planes(
     fitRMSF         ... fit the main lobe of the RMSF with a Gaussian [False]
     fitRMSFreal     ... fit RMSF.real, rather than abs(RMSF) [False]
     nBits           ... precision of data arrays [32]
+    eps             ... NUFFT tolerance [1e-6] (see https://finufft.readthedocs.io/en/latest/python.html#module-finufft)
     verbose         ... print feedback during calculation [False]
     log             ... function to be used to output messages [print]
 
@@ -379,7 +381,7 @@ def get_rmsf_planes(
                 x=a,
                 c=np.ascontiguousarray(weightCube.T),
                 s=(phiArr_radm2[::-1] * 2).astype(a.dtype),
-                eps=1e-8,
+                eps=eps,
             )
             * KArr[..., None]
         ).T
diff --git a/tests/nufft_test.py b/tests/nufft_test.py
@@ -385,59 +385,72 @@ def make_fake_data() -> FakeData:
 def test_rmsynth() -> None:
     """Test the NUFFT RM-synthesis routine agaist DFT."""
     fake_data = make_fake_data()
-    tick = time()
-    FDFcube, lam0Sq_m2 = do_rmsynth_planes(
-        dataQ=fake_data.stokes_q,
-        dataU=fake_data.stokes_u,
-        lambdaSqArr_m2=fake_data.lsq,
-        phiArr_radm2=fake_data.phis,
-        weightArr=fake_data.weights,
-        lam0Sq_m2=fake_data.lsq_0,
-    )
-    tock = time()
-    logger.info(f"Time taken for NUFFT: {(tock - tick)*1000:0.2f} ms")
-
-    tick = time()
-    FDFcube_old, lam0Sq_m2_old = do_rmsynth_planes_old(
-        dataQ=fake_data.stokes_q,
-        dataU=fake_data.stokes_u,
-        lambdaSqArr_m2=fake_data.lsq,
-        phiArr_radm2=fake_data.phis,
-        weightArr=fake_data.weights,
-        lam0Sq_m2=fake_data.lsq_0,
-    )
-    tock = time()
-    logger.info(f"Time taken for DFT: {(tock - tick)*1000:0.2f} ms")
+    for eps in [1e-4, 1e-5, 1e-6, 1e-8]:
+        tick = time()
+        FDFcube, lam0Sq_m2 = do_rmsynth_planes(
+            dataQ=fake_data.stokes_q,
+            dataU=fake_data.stokes_u,
+            lambdaSqArr_m2=fake_data.lsq,
+            phiArr_radm2=fake_data.phis,
+            weightArr=fake_data.weights,
+            lam0Sq_m2=fake_data.lsq_0,
+            eps=eps,
+        )
+        tock = time()
+        logger.info(f"Time taken for NUFFT: {(tock - tick)*1000:0.2f} ms")
+
+        tick = time()
+        FDFcube_old, lam0Sq_m2_old = do_rmsynth_planes_old(
+            dataQ=fake_data.stokes_q,
+            dataU=fake_data.stokes_u,
+            lambdaSqArr_m2=fake_data.lsq,
+            phiArr_radm2=fake_data.phis,
+            weightArr=fake_data.weights,
+            lam0Sq_m2=fake_data.lsq_0,
+        )
+        tock = time()
+        logger.info(f"Time taken for DFT: {(tock - tick)*1000:0.2f} ms")
 
-    assert np.allclose(FDFcube, FDFcube_old, rtol=1e-6, atol=1e-6)
-    assert np.allclose(lam0Sq_m2, lam0Sq_m2_old, rtol=1e-6, atol=1e-6)
+        # fiNUFFT can't go below 1e-8 in precision!
+        if eps == 1e-8:
+            assert np.allclose(FDFcube, FDFcube_old, rtol=eps * 10, atol=eps * 10)
+            assert np.allclose(lam0Sq_m2, lam0Sq_m2_old, rtol=eps * 10, atol=eps * 10)
+        else:
+            assert np.allclose(FDFcube, FDFcube_old, rtol=eps, atol=eps)
+            assert np.allclose(lam0Sq_m2, lam0Sq_m2_old, rtol=eps, atol=eps)
 
 
 def test_rmsf():
     """Test the NUFFT RMSF routine agaist DFT."""
     fake_data = make_fake_data()
-    tick = time()
-    RMSFcube, phi2Arr, fwhmRMSFArr, statArr, _ = get_rmsf_planes(
-        lambdaSqArr_m2=fake_data.lsq,
-        phiArr_radm2=fake_data.phis,
-        weightArr=fake_data.weights,
-        lam0Sq_m2=fake_data.lsq_0,
-    )
-    tock = time()
-    logger.info(f"Time taken for NUFFT: {(tock - tick)*1000:0.2f} ms")
-
-    tick = time()
-    RMSFcube_old, phi2Arr_old, fwhmRMSFArr_old, statArr_old = get_rmsf_planes_old(
-        lambdaSqArr_m2=fake_data.lsq,
-        phiArr_radm2=fake_data.phis,
-        weightArr=fake_data.weights,
-        lam0Sq_m2=fake_data.lsq_0,
-    )
-    tock = time()
-    logger.info(f"Time taken for DFT: {(tock - tick)*1000:0.2f} ms")
-
-    for new, old in zip(
-        [RMSFcube, phi2Arr, fwhmRMSFArr, statArr],
-        [RMSFcube_old, phi2Arr_old, fwhmRMSFArr_old, statArr_old],
-    ):
-        assert np.allclose(new, old, rtol=1e-6, atol=1e-6)
+    for eps in [1e-4, 1e-5, 1e-6, 1e-8]:
+        tick = time()
+        RMSFcube, phi2Arr, fwhmRMSFArr, statArr, _ = get_rmsf_planes(
+            lambdaSqArr_m2=fake_data.lsq,
+            phiArr_radm2=fake_data.phis,
+            weightArr=fake_data.weights,
+            lam0Sq_m2=fake_data.lsq_0,
+            eps=eps,
+        )
+        tock = time()
+        logger.info(f"Time taken for NUFFT: {(tock - tick)*1000:0.2f} ms")
+
+        tick = time()
+        RMSFcube_old, phi2Arr_old, fwhmRMSFArr_old, statArr_old = get_rmsf_planes_old(
+            lambdaSqArr_m2=fake_data.lsq,
+            phiArr_radm2=fake_data.phis,
+            weightArr=fake_data.weights,
+            lam0Sq_m2=fake_data.lsq_0,
+        )
+        tock = time()
+        logger.info(f"Time taken for DFT: {(tock - tick)*1000:0.2f} ms")
+
+        # fiNUFFT can't go below 1e-8 in precision!
+        for new, old in zip(
+            [RMSFcube, phi2Arr, fwhmRMSFArr, statArr],
+            [RMSFcube_old, phi2Arr_old, fwhmRMSFArr_old, statArr_old],
+        ):
+            if eps == 1e-8:
+                assert np.allclose(new, old, rtol=eps * 10, atol=eps * 10)
+            else:
+                assert np.allclose(new, old, rtol=eps, atol=eps)

Original file line number	Diff line number	Diff line change
`@@ -318,7 +318,6 @@ def run_rmsynth(`
`318`	`318`	`phiArr_radm2=phiArr_radm2,`
`319`	`319`	`weightArr=weightArr,`
`320`	`320`	`nBits=nBits,`
`321`		`- verbose=verbose,`
`322`	`321`	`log=log,`
`323`	`322`	`lam0Sq_m2=0 if super_resolution else None,`
`324`	`323`	`)`
Original file line number	Diff line number	Diff line change
`@@ -203,7 +203,6 @@ def run_rmsynth(`
`203`	`203`	`phiArr_radm2=phiArr_radm2,`
`204`	`204`	`weightArr=weightArr,`
`205`	`205`	`nBits=32,`
`206`		`- verbose=verbose,`
`207`	`206`	`lam0Sq_m2=0 if super_resolution else None,`
`208`	`207`	`)`
`209`	`208`	`else:`