clean up and rerun some old code, update dependencies, update python versions

Author: YigitElma

.github/workflows/jax_tests.yml

@@ -82,6 +82,8 @@ jobs:
             0.7.0,
             0.7.1,
             0.7.2,
+            0.8.0,
+            0.9.1,
           ]
         group: [1, 2]
     steps:
@@ -99,7 +101,6 @@ jobs:
           sed -i '1i\jax[cpu] == ${{ matrix.jax-version }}' ./requirements.txt
           cat ./requirements.txt
           pip install -r ./devtools/dev-requirements.txt
-          pip install matplotlib==3.9.2
       - name: Verify dependencies
         run: |
           python --version

.github/workflows/main_tests.yml

@@ -19,6 +19,7 @@ jobs:
           - "3.11"
           - "3.12"
           - "3.13"
+          - "3.14"
 
     steps:
       - uses: actions/checkout@v4

.github/workflows/notebook_tests.yml

@@ -15,7 +15,7 @@ jobs:
       fail-fast: false
       matrix:
         python-version:
-          - "3.10"
+          - "3.13"
 
     steps:
       - uses: actions/checkout@v4
@@ -29,7 +29,6 @@ jobs:
           python -m pip install --upgrade pip
           pip install -r dev-requirements.txt
 
-
       - name: Test notebooks with pytest and nbmake
         run: |
           pwd

.github/workflows/release.yml

@@ -12,7 +12,7 @@ jobs:
       - name: Set up Python
         uses: actions/setup-python@v5
         with:
-          python-version: "3.10"
+          python-version: "3.12"
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip

README.md

@@ -1,16 +1,8 @@
 # Zernipax
-A python library to calculate Zernike Polynomials fast and accurately using JAX.
+A python library to calculate Zernike Polynomials fast and accurately using JAX. Available on `PyPI`,
 
-To create conda environment:
 ```
-conda create --name zernipax-env 'python>=3.8, <=3.12'
-conda activate zernipax-env
-pip install -r requirements.txt
-```
-
-or, if you want to contribute
-```
-pip install -r dev-requirement.txt
+pip install zernipax
 ```
 
 ## GPU Support
@@ -24,16 +16,14 @@ pip install jax[cuda12]
 If you use this repository in your projects, please cite it as:
 ```
 @article{ELMACIOGLU2025129534,
-title = {ZERNIPAX: A fast and accurate Zernike polynomial calculator in Python},
-journal = {Applied Mathematics and Computation},
-volume = {505},
-pages = {129534},
-year = {2025},
-issn = {0096-3003},
-doi = {https://doi.org/10.1016/j.amc.2025.129534},
-url = {https://www.sciencedirect.com/science/article/pii/S0096300325002607},
-author = {Yigit Gunsur Elmacioglu and Rory Conlin and Daniel W. Dudt and Dario Panici and Egemen Kolemen},
-keywords = {Zernike polynomials, Optics, Astrophysics, Spectral simulations, Python, JAX, CPU/GPU computing},
-abstract = {Zernike polynomials serve as an orthogonal basis on the unit disc, and have proven to be effective in optics simulations, astrophysics, and more recently in plasma simulations. Unlike Bessel functions, Zernike polynomials are inherently finite and smooth at the disc center (r=0), ensuring continuous differentiability along the axis. This property makes them particularly suitable for simulations, requiring no additional handling at the origin. We developed ZERNIPAX, an open-source Python package capable of utilizing CPU/GPUs, leveraging Google's JAX package and available on GitHub as well as the Python software repository PyPI. Our implementation of the recursion relation between Jacobi polynomials significantly improves computation time compared to alternative methods by use of parallel computing while still performing more accurately for high-mode numbers.}
+    title = {ZERNIPAX: A fast and accurate Zernike polynomial calculator in Python},
+    journal = {Applied Mathematics and Computation},
+    volume = {505},
+    pages = {129534},
+    year = {2025},
+    issn = {0096-3003},
+    doi = {https://doi.org/10.1016/j.amc.2025.129534},
+    url = {https://www.sciencedirect.com/science/article/pii/S0096300325002607},
+    author = {Yigit Gunsur Elmacioglu, Rory Conlin, Daniel W. Dudt, Dario Panici and Egemen Kolemen},
 }
 ```

dev-requirements.txt

@@ -11,7 +11,7 @@ pre-commit
 
 # testing and benchmarking
 nbmake <= 1.5.5
-pytest ~= 8.0.0
+pytest ~= 9.0
 pytest-benchmark
 pytest-cov >= 2.6.0
 pytest-monitor

notebooks/benchmarks.ipynb

@@ -44,12 +44,11 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "res = 12\n",
-    "# basis = ZernikePolynomial(L=res, M=res, spectral_indexing=\"ansi\", sym=\"cos\")\n",
     "basis = FourierZernikeBasis(L=res, M=res, N=res)\n",
     "r = np.linspace(0, 1, 1000)"
    ]
@@ -63,15 +62,15 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "WARNING:2025-10-19 18:04:12,507:jax._src.xla_bridge:794: An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.\n"
+      "WARNING:2026-02-24 16:12:18,619:jax._src.xla_bridge:794: An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.\n"
      ]
     },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "9.89 ms ± 827 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n",
-      "46.7 ms ± 582 μs per loop (mean ± std. dev. of 7 runs, 10 loops each)\n"
+      "10.9 ms ± 367 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n",
+      "46.4 ms ± 664 μs per loop (mean ± std. dev. of 7 runs, 10 loops each)\n"
      ]
     }
    ],
@@ -96,13 +95,13 @@
      "output_type": "stream",
      "text": [
       "zernike_radial_poly, 0th derivative\n",
-      "5.61 ms ± 318 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n",
+      "4.81 ms ± 273 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n",
       "zernike_radial_poly, 1st derivative\n",
-      "5.82 ms ± 79.8 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n",
+      "4.81 ms ± 133 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n",
       "zernike_radial_poly, 2nd derivative\n",
-      "5.83 ms ± 148 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n",
+      "4.6 ms ± 273 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n",
       "zernike_radial_poly, 3rd derivative\n",
-      "6.46 ms ± 467 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
+      "4.75 ms ± 115 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
      ]
     }
    ],
@@ -126,40 +125,16 @@
     "# Exact computation\n",
     "mpmath.mp.dps = 100\n",
     "c = zernike_radial_coeffs(basis.modes[:, 0], basis.modes[:, 1], exact=True)\n",
-    "zt0 = np.array([np.asarray(mpmath.polyval(list(ci), r), dtype=float) for ci in c]).T\n",
-    "zt1 = np.array([np.asarray(mpmath.polyval(list(ci), r), dtype=float) for ci in polyder_vec(c, 1, exact=True)]).T\n",
-    "zt2 = np.array([np.asarray(mpmath.polyval(list(ci), r), dtype=float) for ci in polyder_vec(c, 2, exact=True)]).T\n",
     "zt3 = np.array([np.asarray(mpmath.polyval(list(ci), r), dtype=float) for ci in polyder_vec(c, 3, exact=True)]).T\n",
     "# CPU version\n",
-    "zr0 = zernike_radial_cpu(r, basis.modes[:, 0], basis.modes[:, 1], 0)\n",
-    "zr1 = zernike_radial_cpu(r, basis.modes[:, 0], basis.modes[:, 1], 1)\n",
-    "zr2 = zernike_radial_cpu(r, basis.modes[:, 0], basis.modes[:, 1], 2)\n",
     "zr3 = zernike_radial_cpu(r, basis.modes[:, 0], basis.modes[:, 1], 3)\n",
     "# GPU version\n",
-    "zd0 = zernike_radial_gpu(\n",
-    "    r[:, np.newaxis], basis.modes[:, 0], basis.modes[:, 1], 0\n",
-    ")\n",
-    "zd1 = zernike_radial_gpu(\n",
-    "    r[:, np.newaxis], basis.modes[:, 0], basis.modes[:, 1], 1\n",
-    ")\n",
-    "zd2 = zernike_radial_gpu(\n",
-    "    r[:, np.newaxis], basis.modes[:, 0], basis.modes[:, 1], 2\n",
-    ")\n",
     "zd3 = zernike_radial_gpu(\n",
     "    r[:, np.newaxis], basis.modes[:, 0], basis.modes[:, 1], 3\n",
     ")\n",
     "\n",
     "mpmath.mp.dps = 10\n",
     "# Polynomial computation\n",
-    "zp0 = zernike_radial_poly(\n",
-    "    r[:, np.newaxis], basis.modes[:, 0], basis.modes[:, 1], dr=0, exact=False\n",
-    ")\n",
-    "zp1 = zernike_radial_poly(\n",
-    "    r[:, np.newaxis], basis.modes[:, 0], basis.modes[:, 1], dr=1, exact=False\n",
-    ")\n",
-    "zp2 = zernike_radial_poly(\n",
-    "    r[:, np.newaxis], basis.modes[:, 0], basis.modes[:, 1], dr=2, exact=False\n",
-    ")\n",
     "zp3 = zernike_radial_poly(\n",
     "    r[:, np.newaxis], basis.modes[:, 0], basis.modes[:, 1], dr=3, exact=False\n",
     ")"
@@ -187,7 +162,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 10,
    "metadata": {},
    "outputs": [
     {
@@ -196,11 +171,11 @@
      "text": [
       "zernike_radial, derivative order: 0\n",
       "# With no duplicate modes (might have lacking modes)\n",
-      "929 μs ± 327 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n",
-      "# With all the checks necessary but no reverse mode AutoDiff capable\n",
-      "11.3 ms ± 480 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n",
+      "1.24 ms ± 308 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n",
+      "# With all the checks necessary\n",
+      "12.9 ms ± 274 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n",
       "# With all the checks necessary but less efficient on CPU\n",
-      "41.1 ms ± 1.19 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)\n"
+      "44.6 ms ± 418 μs per loop (mean ± std. dev. of 7 runs, 10 loops each)\n"
      ]
     }
    ],
@@ -211,7 +186,7 @@
     "print(\"# With no duplicate modes (might have lacking modes)\")\n",
     "%timeit _ = zernike_radial_unique(r, basis.modes[:,0], basis.modes[:,1], dr).block_until_ready()\n",
     "\n",
-    "print(\"# With all the checks necessary but no reverse mode AutoDiff capable\")\n",
+    "print(\"# With all the checks necessary\")\n",
     "%timeit _ = zernike_radial_cpu(r, basis.modes[:,0], basis.modes[:,1], dr).block_until_ready()\n",
     "\n",
     "print(\"# With all the checks necessary but less efficient on CPU\")\n",
@@ -220,7 +195,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": 11,
    "metadata": {},
    "outputs": [
     {
@@ -229,11 +204,11 @@
      "text": [
       "zernike_radial, derivative order: 1\n",
       "# With no duplicate modes (might have lacking modes)\n",
-      "1.42 ms ± 303 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n",
-      "# With all the checks necessary but no reverse mode AutoDiff capable\n",
-      "11.1 ms ± 303 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n",
+      "1.93 ms ± 6.67 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n",
+      "# With all the checks necessary\n",
+      "13.2 ms ± 303 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n",
       "# With all the checks necessary but less efficient on CPU\n",
-      "68.9 ms ± 460 μs per loop (mean ± std. dev. of 7 runs, 10 loops each)\n"
+      "75.4 ms ± 1.11 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)\n"
      ]
     }
    ],
@@ -244,7 +219,7 @@
     "print(\"# With no duplicate modes (might have lacking modes)\")\n",
     "%timeit _ = zernike_radial_unique(r, basis.modes[:,0], basis.modes[:,1], dr).block_until_ready()\n",
     "\n",
-    "print(\"# With all the checks necessary but no reverse mode AutoDiff capable\")\n",
+    "print(\"# With all the checks necessary\")\n",
     "%timeit _ = zernike_radial_cpu(r, basis.modes[:,0], basis.modes[:,1], dr).block_until_ready()\n",
     "\n",
     "print(\"# With all the checks necessary but less efficient on CPU\")\n",

notebooks/tutorial-basis.ipynb

@@ -1,4 +1,4 @@
-jax[cpu] >= 0.3.2, <= 0.6.2
-matplotlib >= 3.5.0, <= 4.0.0
+jax[cpu] >= 0.3.2
+matplotlib >= 3.5.0
 mpmath >= 1.0.0
-numpy >= 1.20.0, < 2.0.0
+numpy >= 1.20.0

requirements.txt

@@ -33,7 +33,7 @@
     keywords="zernike polynomials, optics, astrophysics, spectral "
     + "simulation, basis, orthogonal polynomials, parallel computing, "
     + "JAX",
-    packages=find_packages(exclude=["tests", "notebooks", "figures"]),
+    packages=find_packages(include=["zernipax"]),
     include_package_data=True,
     install_requires=requirements,
     python_requires=">=3.10",