Adding support for jax and jaxoplanet

.github/workflows/pytest_testing.yml

@@ -26,7 +26,7 @@ jobs:
           sudo apt-get update
           sudo apt-get install -y libnetcdf-dev libhdf5-dev
           python -m pip install --timeout 60 --retries 5 --upgrade pip
-          pip install --timeout 60 --retries 5 ".[jwst,hst,test]"
+          pip install --timeout 60 --retries 5 ".[jwst,hst,test,jax]"
           mkdir $HOME/crds_cache
           echo "CRDS_PATH=$HOME/crds_cache" >> $GITHUB_ENV
           echo "CRDS_SERVER_URL=https://jwst-crds.stsci.edu" >> $GITHUB_ENV

environment.yml

@@ -27,23 +27,31 @@ dependencies:
     - pandas
     - photutils
     - pip # required for the pip installs below
-    - pytest[version='<8.0'] # Upper limit needed to avoid some conflict between pytest, an asdf package, and something in the pymc3 set of installs. Needed for testing
+    - pytest[version='<8.0'] # Upper limit needed to avoid some conflict between pytest and asdf. Needed for testing
     - pytest-cov # Needed for testing
     - pytest-doctestplus # Needed for testing
     - requests
     - scipy[version='>=1.8.0'] # Lower limit needed for scipy.spatial.QhullError
     - shapely
     - sphinx # Needed for documentation
     - sphinx-automodapi # Needed for documentation
+    - tinygp
     - tqdm
     - pip:
+        - arviz>=0.21  # Needed for jaxoplanet
+        - jax[cpu]>=0.5  # Needed for jaxoplanet
+        - jaxoplanet>=0.1.0  # Needed for jaxoplanet
+        - jaxopt>=0.8  # Needed for jaxoplanet
+        - numpyro>=0.19  # Needed for jaxoplanet
+        - numpyro_ext>=0.0.4  # Needed for jaxoplanet
+        - specutils>=1.20  # Needed for jaxoplanet
         - astraeus-io
         - batman-package
         - catwoman>=1.0.16 # Lower limit needed to avoid installation bug
         - celerite2 # Needed for GP
         - crds
         - exotic-ld==3.2.0 # Lower limit needed for updated JWST sensitivity files, upper limit needed for breaking changes
-        - fleck
+        - fleck@git+https://github.com/bmorris3/fleck.git@main
         - george # Needed for GP
         - image_registration@git+https://github.com/keflavich/image_registration@master # Need GitHub version to avoid np.float issue
         - jwst==1.20.2

setup.cfg

@@ -43,6 +43,7 @@ zip_safe = True
 python_requires= >=3.11,<3.14
 include_package_data = True
 install_requires =
+    arviz>=0.21
     astraeus-io
     astropy
     astroquery
@@ -56,24 +57,31 @@ install_requires =
     dynesty>1.0 # Lower limit needed for specific arguments
     emcee>3.0.0 # Lower limit needed for specific arguments
     exotic-ld==3.2.0 # Lower limit needed for updated JWST sensitivity files, upper limit needed for breaking changes
-    fleck
+    fleck@git+https://github.com/taylorbell57/fleck.git@main # Fork for fixed jax support and other bug fixes
     george # Needed for GP
     h5py>=3.7.0 # Lower limit needed to avoid issues with JWST I/O routines
+    jax[cpu]>=0.5 # Lower limit needed for compatibility with jaxoplanet
+    jaxoplanet>=0.1.0
+    jaxopt>=0.8 # Lower limit needed for compatibility with jax>=0.5
     lmfit
     matplotlib>=3.6 # Lower limit needed for set_layout_engine()
     mc3 # Needed for uncertainties in the RMS time-averaging plots in S5
     numpy>=1.20.0,<2.0 # numpy 2.0 introduces breaking changes for some packages
+    numpyro>=0.19 # Lower limit needed for compatibility with jax>=0.5
+    numpyro_ext>=0.0.4 # Lower limit needed for compatibility with numpyro>=0.19
     pandas
     pastasoss>=1.2.2
     photutils
     # planet-harmonica # Currently only supports for python 3.8 -- 3.10
-    planet-harmonica@git+https://github.com/DavoGrant/harmonica@main
+    planet-harmonica@git+https://github.com/taylorbell57/harmonica@main # Fork for jax support
     requests
     scipy>=1.12.0 # Lower limit needed for stcal.ramp_fitting.likely_fit
     setuptools_scm>=6.2 # Needed to support pyproject.toml config
     shapely
+    specutils>=1.20 # Lower limit needed for compatibility with jaxoplanet
     stdatamodels
     svo_filters
+    tinygp
     tqdm
 
 
@@ -100,7 +108,7 @@ docs =
     sphinx-rtd-theme>=1.3.0 # Style updates and Sphinx compatibility
 test =
     flake8
-    pytest<8.0 # Upper limit due to asdf+pytest incompatibility
+    pytest<8.0 # Upper limit needed to avoid some conflict between pytest and asdf
     pytest-cov
     pytest-doctestplus
 jupyter =
@@ -142,6 +150,14 @@ per-file-ignores =
     ./demos/**/run_*.py:E402
     # Need to protect against crashes if jwst not installed
     **/test_*.py:E402
+    # Need to have unused imports in the AstroModel
+    **/models/AstroModel.py:F401
+    # Need to modify c compiler settings before importing starry
+    **/differentiable_models/*.py:E402
+    # Need to have unused imports in the init file
+    **/differentiable_models/__init__.py:F401
+    # Need to modify jax settings before importing celerite2.jax
+    **/jax_models/JaxModel.py:E402
     # Need to protect against crashes if hst packages not installed
     **/test_WFC3.py:F401,E402
     # Keep the unused numpy so users can pass in numpy values (e.g. np.inf)

src/eureka/S3_data_reduction/s3_reduce.py

@@ -389,9 +389,14 @@ def reduce(eventlabel, ecf_path=None, s2_meta=None, input_meta=None):
                     if not meta.ff_outlier and not meta.skip_bg:
                         data = inst.flag_bg(data, meta, log)
 
-                    # Do the background subtraction
-                    data = bg.BGsubtraction(data, meta, log,
-                                            m, meta.isplots_S3)
+                        # Do the background subtraction
+                        data = bg.BGsubtraction(data, meta, log,
+                                                m, meta.isplots_S3)
+                    else:
+                        data['bg'] = data.flux.copy()
+                        data.bg.values[:] = 0
+                        data.bg.attrs['bg_units'] = \
+                            data.flux.attrs['flux_units']
 
                     # Calulate and correct for 2D drift
                     if hasattr(inst, 'correct_drift2D'):

src/eureka/S4_generate_lightcurves/generate_LD.py

@@ -140,7 +140,7 @@ def exotic_ld(meta, spec, log, white=False):
                                     throughput_wavelengths < 30000)
             poly = np.polyfit(throughput_wavelengths[ind_use],
                               throughput[ind_use], deg=7)
-            wav_poly = np.linspace(2.733*1e4, throughput_wavelengths[0], 10000)
+            wav_poly = np.linspace(2.733e4, throughput_wavelengths[0], 10000)
             throughput_poly = np.polyval(poly, wav_poly) - 0.015
             # Make sure the throughput is always > 0
             throughput_poly[throughput_poly < 0] = 0
@@ -150,6 +150,26 @@ def exotic_ld(meta, spec, log, white=False):
             custom_throughput = np.append(throughput_poly, throughput)
             old_mode = mode
             mode = 'custom'
+        elif (mode == 'JWST_MIRI_LRS' and
+                wavelength_range[0][0] > throughput_edges[0]/1e4):
+            # Extrapolate throughput to the blue edge of the filter if needed
+            log.writelog("WARNING: Extrapolating ExoTiC-LD throughput file to "
+                         "get closer to the blue edge of the filter.")
+
+            # The following polynomial was estimated by TJB on Dec 19, 2025
+            ind_use = throughput_wavelengths < 100000
+            poly = np.polyfit(throughput_wavelengths[ind_use],
+                              throughput[ind_use], deg=3)
+            wav_poly = np.linspace(4e4, throughput_wavelengths[0], 1000)
+            throughput_poly = np.polyval(poly, wav_poly)
+            # Make sure the throughput is always > 0
+            throughput_poly[throughput_poly < 0] = 0
+            # Prepend extrapolated throughput and then switch to custom
+            # throughput mode
+            custom_wavelengths = np.append(wav_poly, throughput_wavelengths)
+            custom_throughput = np.append(throughput_poly, throughput)
+            old_mode = mode
+            mode = 'custom'
 
         if mode == 'custom' and meta.isplots_S4 >= 3:
             plots_s4.plot_extrapolated_throughput(meta, throughput_wavelengths,

src/eureka/S4_generate_lightcurves/outliers.py

@@ -30,8 +30,8 @@ def get_outliers(meta, spec):
     optmask = spec.optmask.values[:, iwmin:iwmax]
     norm_lcdata, norm_lcerr = util.normalize_spectrum(meta, optspec, opterr,
                                                       optmask=optmask)
-    norm_lcdata = norm_lcdata.filled(np.nan)
-    norm_lcerr = norm_lcerr.filled(np.nan)
+    norm_lcdata = np.ma.filled(norm_lcdata, np.nan)
+    norm_lcerr = np.ma.filled(norm_lcerr, np.nan)
 
     # Compute unbinned LC MAD values, then scale
     numx = norm_lcdata.shape[1]

src/eureka/S5_lightcurve_fitting/__init__.py

@@ -1,7 +1,26 @@
 """
 Package to fit models to light curve data
 """
+# Do some checks here to avoid spamming these warnings
+try:
+    import os
+    import multiprocessing
+
+    # Must come before any import that touches JAX
+    ncpu = multiprocessing.cpu_count()
+    os.environ["XLA_FLAGS"] = f"--xla_force_host_platform_device_count={ncpu}"
+
+    import jax
+    jax.config.update("jax_enable_x64", True)
+    from . import jax_models
+    from . import jax_lightcurve
+except ModuleNotFoundError:
+    # Don't require that the jax and jaxoplanet packages be installed
+    # but also don't raise a warning here to avoid excessive spam
+    pass
+
 from . import fitters
+# from . import gradient_fitters
 from . import lightcurve
 from . import likelihood
 from . import limb_darkening_fit

src/eureka/S5_lightcurve_fitting/fitters.py

@@ -89,6 +89,16 @@ def lsqfitter(lc, model, meta, log, calling_function='lsq', **kwargs):
             'your initial parameter values are outside of the bounds of its '
             'prior.')
 
+    # Plot Harmonica string starting point
+    if 'harmonica_tr' in meta.run_myfuncs and meta.isplots_S5 >= 3:
+        plots.plot_harmonica_string(lc, model, meta,
+                                    fitter=calling_function+'StartingPoint')
+
+    if not np.isfinite(start_lnprob):
+        raise AssertionError(
+            'The starting lnprob value must be finite. Most likely, one of '
+            'your priors is outside of its bounds.')
+
     def neg_lnprob(theta, lc, model, prior1, prior2, priortype, freenames):
         return -lnprob(theta, lc, model, prior1, prior2, priortype, freenames)
     global lsq_t0
@@ -315,6 +325,11 @@ def emceefitter(lc, model, meta, log, **kwargs):
         plots.plot_harmonica_string(lc, model, meta,
                                     fitter='emceeStartingPoint')
 
+    # Plot Harmonica string starting point
+    if 'harmonica_tr' in meta.run_myfuncs and meta.isplots_S5 >= 3:
+        plots.plot_harmonica_string(lc, model, meta,
+                                    fitter='emceeStartingPoint')
+
     # Initialize tread pool
     if meta.ncpu > 1:
         pool = Pool(meta.ncpu)