Repository: astro-informatics/s2fft (language: python)

Issue:
Forcing jax float64 mode
This is not really good practice, users can have legitimate reasons for not using float64 and if they import `s2fft` first, it will force the type without leaving them a choice:
https://github.com/astro-informatics/s2fft/blob/e9e6f78c685d2dd54f6344e29a35de65c09fb350/s2fft/transforms/spherical.py#L3

If you think float64 is really important for this library, you could instead send a warning back if the users have not enabled it. 


Discussion / hints:
The problem really starts to become apparent inside the recursions. With 64bit precision, at high L the error can grow from 1e-16 to around 1e-13 or so. With 32bit precision, without accounting for any other factors, this should correspond to an error from around 1e-6 to 1e-3 to even 1e-2, which isn't great.

I think a warning would be a good option around this though, but we will have to be clear that 64bit can be really quite important in this case. This hard coding of 64bit precision comes up throughout several functions, perhaps we should just include a warning at the high-level?

Fix patch:
diff --git a/notebooks/spherical_harmonic_transform.ipynb b/notebooks/spherical_harmonic_transform.ipynb
index df714e85..cfee00d0 100644
--- a/notebooks/spherical_harmonic_transform.ipynb
+++ b/notebooks/spherical_harmonic_transform.ipynb
@@ -28,6 +28,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "from jax.config import config\n",
+    "config.update(\"jax_enable_x64\", True)\n",
+    "\n",
     "import numpy as np\n",
     "import s2fft \n",
     "\n",
@@ -164,7 +167,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3.9.0 ('s2fft')",
+   "display_name": "Python 3.8.16 64-bit ('s2fft')",
    "language": "python",
    "name": "python3"
   },
@@ -178,12 +181,12 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.0"
+   "version": "3.8.16"
   },
   "orig_nbformat": 4,
   "vscode": {
    "interpreter": {
-    "hash": "3425e24474cbe920550266ea26b478634978cc419579f9dbcf479231067df6a3"
+    "hash": "d6019e21eb0d27eebd69283f1089b8b605b46cb058a452b887458f3af7017e46"
    }
   }
  },
diff --git a/notebooks/wigner_transform.ipynb b/notebooks/wigner_transform.ipynb
index bab81a2a..c1274705 100644
--- a/notebooks/wigner_transform.ipynb
+++ b/notebooks/wigner_transform.ipynb
@@ -28,6 +28,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "from jax.config import config\n",
+    "config.update(\"jax_enable_x64\", True) \n",
+    "\n",
     "import numpy as np\n",
     "import s2fft \n",
     "\n",
diff --git a/requirements/requirements-core.txt b/requirements/requirements-core.txt
index 1262799f..496be188 100644
--- a/requirements/requirements-core.txt
+++ b/requirements/requirements-core.txt
@@ -3,8 +3,7 @@ numpy>=1.20
 colorlog
 pyyaml
 pyssht
-jaxlib==0.3.10
-jax==0.3.13
+jax>=0.3.13
 
 # Packages required for plotting functionality
 pyvista
diff --git a/s2fft/__init__.py b/s2fft/__init__.py
index f8b4229d..37e0f1c2 100644
--- a/s2fft/__init__.py
+++ b/s2fft/__init__.py
@@ -7,3 +7,9 @@
     generate_precomputes_wigner,
     generate_precomputes_wigner_jax,
 )
+
+import logging
+from jax.config import config
+if config.read("jax_enable_x64") is False:
+    logger = logging.getLogger("s2fft")
+    logger.warning("JAX is not using 64-bit precision. This will dramatically affect numerical precision at even moderate L.")
\ No newline at end of file
diff --git a/s2fft/precompute_transforms/construct.py b/s2fft/precompute_transforms/construct.py
index aef08bf5..95d2b136 100644
--- a/s2fft/precompute_transforms/construct.py
+++ b/s2fft/precompute_transforms/construct.py
@@ -1,6 +1,3 @@
-from jax import config
-
-config.update("jax_enable_x64", True)
 import numpy as np
 from s2fft.sampling import s2_samples as samples
 from s2fft.utils import quadrature
diff --git a/s2fft/precompute_transforms/spherical.py b/s2fft/precompute_transforms/spherical.py
index 14a388ea..2a83c0e4 100644
--- a/s2fft/precompute_transforms/spherical.py
+++ b/s2fft/precompute_transforms/spherical.py
@@ -1,6 +1,5 @@
-from jax import jit, config
+from jax import jit 
 
-config.update("jax_enable_x64", True)
 import numpy as np
 import jax.numpy as jnp
 
diff --git a/s2fft/precompute_transforms/wigner.py b/s2fft/precompute_transforms/wigner.py
index 415ab7cd..5542a538 100644
--- a/s2fft/precompute_transforms/wigner.py
+++ b/s2fft/precompute_transforms/wigner.py
@@ -1,6 +1,4 @@
-from jax import jit, config
-
-config.update("jax_enable_x64", True)
+from jax import jit
 
 import numpy as np
 import jax.numpy as jnp
diff --git a/s2fft/recursions/price_mcewen.py b/s2fft/recursions/price_mcewen.py
index b0711aa7..27b99fec 100644
--- a/s2fft/recursions/price_mcewen.py
+++ b/s2fft/recursions/price_mcewen.py
@@ -1,6 +1,4 @@
-from jax import config, local_device_count
-
-config.update("jax_enable_x64", True)
+from jax import local_device_count
 
 import numpy as np
 import jax.numpy as jnp
diff --git a/s2fft/transforms/otf_recursions.py b/s2fft/transforms/otf_recursions.py
index 4d18b92d..b37e33f5 100644
--- a/s2fft/transforms/otf_recursions.py
+++ b/s2fft/transforms/otf_recursions.py
@@ -1,7 +1,3 @@
-from jax import config
-
-config.update("jax_enable_x64", True)
-
 import numpy as np
 import jax.numpy as jnp
 import jax.lax as lax
diff --git a/s2fft/transforms/spherical.py b/s2fft/transforms/spherical.py
index 5cedfcaa..eec6a2ae 100644
--- a/s2fft/transforms/spherical.py
+++ b/s2fft/transforms/spherical.py
@@ -1,6 +1,4 @@
-from jax import jit, config
-
-config.update("jax_enable_x64", True)
+from jax import jit
 
 import numpy as np
 import jax.numpy as jnp
diff --git a/s2fft/transforms/wigner.py b/s2fft/transforms/wigner.py
index f5349b4a..b48dee5a 100644
--- a/s2fft/transforms/wigner.py
+++ b/s2fft/transforms/wigner.py
@@ -1,6 +1,5 @@
-from jax import jit, config, pmap, local_device_count
+from jax import jit, pmap, local_device_count
 
-config.update("jax_enable_x64", True)
 
 import numpy as np
 import jax.numpy as jnp
diff --git a/s2fft/utils/healpix_ffts.py b/s2fft/utils/healpix_ffts.py
index 8e72550d..da2c1c66 100644
--- a/s2fft/utils/healpix_ffts.py
+++ b/s2fft/utils/healpix_ffts.py
@@ -1,6 +1,5 @@
-from jax import jit, config
+from jax import jit
 
-config.update("jax_enable_x64", True)
 import numpy as np
 import jax.numpy as jnp
 from s2fft.sampling import s2_samples as samples
diff --git a/s2fft/utils/quadrature_jax.py b/s2fft/utils/quadrature_jax.py
index 38fa82cb..492962bc 100644
--- a/s2fft/utils/quadrature_jax.py
+++ b/s2fft/utils/quadrature_jax.py
@@ -1,6 +1,5 @@
-from jax import jit, config
+from jax import jit
 
-config.update("jax_enable_x64", True)
 from functools import partial
 from s2fft.sampling import s2_samples as samples
 import jax.numpy as jnp
diff --git a/s2fft/utils/resampling_jax.py b/s2fft/utils/resampling_jax.py
index 0d134436..edcc87b5 100644
--- a/s2fft/utils/resampling_jax.py
+++ b/s2fft/utils/resampling_jax.py
@@ -1,6 +1,5 @@
-from jax import jit, config
+from jax import jit
 
-config.update("jax_enable_x64", True)
 import jax.numpy as jnp
 from s2fft.sampling import s2_samples as samples
 from functools import partial


Test patch:
diff --git a/requirements/requirements-tests.txt b/requirements/requirements-tests.txt
index 542b0894..a14631b6 100644
--- a/requirements/requirements-tests.txt
+++ b/requirements/requirements-tests.txt
@@ -2,6 +2,9 @@
 pytest 
 pytest-cov
 
+#jax dependencies
+jaxlib
+
 # Healpix Sampling
 healpy
 so3