Add utilities to automatically download external data

.github/workflows/ci.yml

@@ -25,24 +25,16 @@ jobs:
     with:
       envs: |
         - linux: check-style
-  latest_data_cache:
-    uses: ./.github/workflows/retrieve_cache.yml
-    with:
-      cache_path: /tmp/data
   test:
-    needs: [ latest_data_cache ]
     uses: OpenAstronomy/github-actions-workflows/.github/workflows/tox.yml@924441154cf3053034c6513d5e06c69d262fb9a6  # v1.13.0
     with:
       libraries: |
         brew:
           - eigen
           - fftw
       setenv: |
-        WEBBPSF_PATH: ${{ needs.latest_data_cache.outputs.cache_path }}/webbpsf-data/
-        GALSIM_CAT_PATH: ${{ needs.latest_data_cache.outputs.cache_path }}/galsim_data/real_galaxy_catalog_23.5_example.fits
         FFTW_DIR: /opt/homebrew/opt/fftw/lib/
-      cache-path: ${{ needs.latest_data_cache.outputs.cache_path }}
-      cache-key: ${{ needs.latest_data_cache.outputs.cache_key }}
+
       envs: |
         - linux: py310-oldestdeps-cov-xdist
         - linux: py310-xdist

.gitignore

@@ -71,6 +71,7 @@ instance/
 
 # Sphinx documentation
 docs/_build/
+docs/api/
 
 # PyBuilder
 .pybuilder/

romanisim/psf.py

@@ -18,8 +18,15 @@
 import galsim
 from galsim import roman
 from .bandpass import galsim2roman_bandpass, roman2galsim_bandpass
+from .util import import_webbpsf
 from romanisim import log
 
+__all__ = [
+    "make_one_psf",
+    "make_psf",
+    "VariablePSF",
+]
+
 
 def make_one_psf(sca, filter_name, wcs=None, webbpsf=True, pix=None,
                  chromatic=False, oversample=4, extra_convolution=None, **kw):
@@ -75,7 +82,7 @@ def make_one_psf(sca, filter_name, wcs=None, webbpsf=True, pix=None,
     if chromatic:
         log.warning('romanisim does not yet support chromatic PSFs '
                     'with webbpsf')
-    import webbpsf as wpsf
+    wpsf = import_webbpsf()
     filter_name = galsim2roman_bandpass[filter_name]
     wfi = wpsf.WFI()
     wfi.detector = f'SCA{sca:02d}'

romanisim/tests/test_catalog.py

@@ -6,6 +6,7 @@
 import numpy as np
 import galsim
 from romanisim import catalog
+from romanisim.util import get_galsim_data_path
 from astropy.coordinates import SkyCoord
 from astropy import units as u
 from romanisim import log
@@ -16,7 +17,7 @@ def test_make_dummy_catalog():
     cen = SkyCoord(ra=5 * u.deg, dec=-10 * u.deg)
     radius = 0.2
     nobj = 100
-    fn = os.environ.get('GALSIM_CAT_PATH', None)
+    fn = get_galsim_data_path()
     if fn is not None:
         fn = str(fn)
     cat = catalog.make_dummy_catalog(

romanisim/tests/test_image.py

@@ -24,13 +24,15 @@
 from astropy.time import Time
 from astropy import table
 import asdf
-import webbpsf
 from astropy.modeling.functional_models import Sersic2D
 import pytest
 from romanisim import log
+from romanisim.util import import_webbpsf, get_galsim_data_path
 from roman_datamodels.stnode import WfiScienceRaw, WfiImage
 import romanisim.bandpass
 
+webbpsf = import_webbpsf()
+
 
 def test_in_bounds():
     bounds = galsim.BoundsI(0, 1000, 0, 1000)
@@ -573,7 +575,7 @@ def test_make_test_catalog_and_images():
     # public interface, and may be removed.  We'll settle for just
     # testing that it runs.
     roman.n_pix = 100
-    fn = os.environ.get('GALSIM_CAT_PATH', None)
+    fn = get_galsim_data_path()
     if fn is not None:
         fn = str(fn)
     res = image.make_test_catalog_and_images(usecrds=False,

romanisim/util.py

@@ -1,5 +1,11 @@
 """Miscellaneous utility routines.
 """
+import os
+import tarfile
+import warnings
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from urllib.request import urlretrieve
 
 import numpy as np
 from astropy.coordinates import SkyCoord
@@ -12,6 +18,26 @@
 from scipy import integrate
 
 
+__all__ = [
+    "skycoord",
+    "celestialcoord",
+    "scalergb",
+    "random_points_in_cap",
+    "random_points_in_king",
+    "random_points_at_radii",
+    "add_more_metadata",
+    "update_pointing_and_wcsinfo_metadata",
+    "king_profile",
+    "sample_king_distances",
+    "decode_context_times",
+    "default_image_meta",
+    "update_photom_keywords",
+    "merge_dicts",
+    "import_webbpsf",
+    "get_galsim_data_path",
+]
+
+
 def skycoord(celestial):
     """Turn a CelestialCoord into a SkyCoord.
 
@@ -539,3 +565,68 @@
         else:
             a[key] = b[key]
     return a
+
+
+def import_webbpsf():
+    """
+    Import webbpsf and download the latest data files if necessary.
+
+    Returns
+    -------
+    webbpsf : module
+        webbpsf module
+    """
+    # Grab the path to the WEBBPSF data files from the environment
+    # or use the default location in the user's home directory.
+    path = os.getenv("WEBBPSF_PATH") or Path.home() / "data" / "webbpsf-data"
+    path = Path(path)
+
+    # Make sure the directory exists
+    path.mkdir(parents=True, exist_ok=True)
+
+    # Set the environment variable so that webbpsf can find the data files
+    os.environ["WEBBPSF_PATH"] = str(path)
+
+    # download the data files if necessary
+    if not any(path.iterdir()):
+        # Download the data files to a temporary directory
+        warnings.warn("Downloading WEBBPSF data files. This may take a few minutes...")
+
+        with TemporaryDirectory() as tmpdir:
+            # This is the URL for the latest data files according to the WEBBPSF documentation
+            url = "https://stsci.box.com/shared/static/qxpiaxsjwo15ml6m4pkhtk36c9jgj70k.gz"
+            filename = Path(tmpdir) / "webbpsf-data-LATEST.tar.gz"
+            urlretrieve(url, filename)  # download the tarball
+
+            # Extract the tarball
+            with tarfile.open(filename) as tar:
+                tar.extractall(path.parent, filter="fully_trusted")
+
+            if not any(path.iterdir()):
+                raise ImportError("Failed to download WEBBPSF data files")
+
+    # Import webbpsf and return it
+    import webbpsf
+
+    return webbpsf
+
+
+def get_galsim_data_path():
+    path = os.environ.get("GALSIM_CAT_PATH") or Path.home() / "data" / "galsim-data"
+    path.mkdir(parents=True, exist_ok=True)
+    os.environ["GALSIM_CAT_PATH"] = str(path)
+
+    url = "https://github.com/GalSim-developers/GalSim/raw/releases/2.4/examples/data/"
+
+    data_files = (
+        "real_galaxy_catalog_23.5_example.fits",
+        "real_galaxy_catalog_23.5_example_selection.fits",
+        "real_galaxy_catalog_23.5_example_fits.fits",
+    )
+
+    for filename in data_files:
+        if not (path / filename).exists():
+            warnings.warn(f"Downloading {filename}...")
+            urlretrieve(url + filename, path / filename)
+
+    return path / data_files[0]