Configfile (#13)

* added cache_dir option * rearrange cache_dir * Update changelog * remove print
lightkurve · Sep 24, 2024 · 5ca119c · 5ca119c
1 parent ae1e311
commit 5ca119c
Show file tree

Hide file tree

Showing 4 changed files with 26 additions and 21 deletions.
diff --git a/CHANGES.rst b/CHANGES.rst
@@ -3,7 +3,7 @@
 
 - Added the option to use the initial TESS commissioning PRF files [#9]
 - Modified tessprf.py and keplerprf.py to use only the closest PRF measurements to create the supersampled PRF [#8]
-
+- Added cache_dir flag to change default location for reading/writing engineer files [#13]
 1.0.3 (2024-07-30)
 ==================
 

diff --git a/src/lkprf/data.py b/src/lkprf/data.py
@@ -7,6 +7,7 @@
 from scipy.ndimage import label, uniform_filter
 
 from . import logger, PACKAGEDIR
+CACHEDIR = PACKAGEDIR + '/data/'
 
 __all__ = [
     "download_kepler_prf_file",
@@ -91,10 +92,10 @@ def _download_file(url, file_path):
                 raise http_err
 
 
-def download_kepler_prf_file(module: int, output: int):
+def download_kepler_prf_file(module: int, output: int, cache_dir: str = CACHEDIR):
     """Download a Kepler Module file"""
     filename = f"kplr{module:02}.{output}_2011265_prf.fits"
-    file_path = f"{PACKAGEDIR}/data/{filename}"
+    file_path = f"{cache_dir}{filename}"
     url = f"https://archive.stsci.edu/missions/kepler/fpc/prf/{filename}"
     logger.info(f"Downloading {module:02}.{output}")
     _download_file(url, file_path)
@@ -104,7 +105,7 @@ def download_kepler_prf_file(module: int, output: int):
     return
 
 
-def build_tess_prf_file(camera: int, ccd: int, sector: int):
+def build_tess_prf_file(camera: int, ccd: int, sector: int, cache_dir: str = CACHEDIR):
     """Download a set of TESS PRF files for a given camera/ccd"""
 
     def open_file(url: str):
@@ -138,7 +139,7 @@ def open_file(url: str):
         prefix = _tess_prefixes[camera][ccd]
         filename = f"tess-prf-cam{camera}-ccd{ccd}-sec4.fits"
 
-    file_path = f"{PACKAGEDIR}/data/{filename}"
+    file_path = f"{cache_dir}{filename}"
     # ensure the file_path exists
     if not os.path.exists(file_path):
         os.makedirs(os.path.dirname(file_path), exist_ok=True)
@@ -160,51 +161,50 @@ def open_file(url: str):
     return
 
 
-def get_kepler_prf_file(module: int, output: int):
+def get_kepler_prf_file(module: int, output: int, cache_dir: str = CACHEDIR):
     """Download a Kepler Module file"""
     filename = f"kplr{module:02}.{output}_2011265_prf.fits"
-    file_path = f"{PACKAGEDIR}/data/{filename}"
+    file_path = f"{cache_dir}{filename}"
     if not os.path.isfile(file_path):
         logger.info(
             f"No local files found, building Kepler PRF for Module {module}, output {output}."
         )
-        download_kepler_prf_file(module=module, output=output)
-    file_path = f"{PACKAGEDIR}/data/{filename}"
+        download_kepler_prf_file(module=module, output=output, cache_dir=cache_dir)
 
     hdulist = fitsio.FITS(file_path)
     return hdulist
 
 
-def get_tess_prf_file(camera: int, ccd: int, sector: int = 4):
+def get_tess_prf_file(camera: int, ccd: int, sector: int = 4, cache_dir: str = CACHEDIR):
     """Get a PRF file for a given camera/ccd/sector"""
     if sector <= 3:
         filename = f"tess-prf-cam{camera}-ccd{ccd}-sec1.fits"
     else:
         filename = f"tess-prf-cam{camera}-ccd{ccd}-sec4.fits"
-    file_path = f"{PACKAGEDIR}/data/{filename}"
+    file_path = f"{cache_dir}{filename}"
     if not os.path.isfile(file_path):
         logger.info(
             f"No local files found, building TESS PRF for Camera {camera}, CCD {ccd}."
         )
-        build_tess_prf_file(camera=camera, ccd=ccd, sector=sector)
-
+        build_tess_prf_file(camera=camera, ccd=ccd, sector=sector, cache_dir=cache_dir)
+    
     hdulist = fitsio.FITS(file_path)
     return hdulist
 
 
-def clear_kepler_cache():
+def clear_kepler_cache(cache_dir: str = CACHEDIR):
     for module in np.arange(26):
         for output in np.arange(1, 5):
             filename = f"kplr{module:02}.{output}_2011265_prf.fits"
-            file_path = f"{PACKAGEDIR}/data/{filename}"
+            file_path = f"{cache_dir}{filename}"
             if os.path.isfile(file_path):
                 os.remove(file_path)
 
 
-def clear_tess_cache():
+def clear_tess_cache(cache_dir: str = CACHEDIR):
     for camera in np.arange(1, 5):
         for ccd in np.arange(1, 5):
             filename = f"tess-prf-{camera}-{ccd}.fits"
-            file_path = f"{PACKAGEDIR}/data/{filename}"
+            file_path = f"{cache_dir}{filename}"
             if os.path.isfile(file_path):
                 os.remove(file_path)
diff --git a/src/lkprf/keplerprf.py b/src/lkprf/keplerprf.py
@@ -4,6 +4,7 @@
 import numpy as np
 from .utils import channel_to_module_output, LKPRFWarning
 from .data import get_kepler_prf_file
+from . import PACKAGEDIR
 import warnings
 
 from .prfmodel import PRF
@@ -21,18 +22,20 @@ class KeplerPRF(PRF):
     https://archive.stsci.edu/missions/kepler/commissioning_prfs/
     """
 
-    def __init__(self, channel: int):
+    def __init__(self, channel: int, cache_dir: str = PACKAGEDIR + '/data/'):
         super().__init__()
         self.channel = channel
         self.mission = "Kepler"
+        self.cache_dir = cache_dir
         self._prepare_prf()
+
 
     def __repr__(self):
         return f"KeplerPRF Object [Channel {self.channel}]"
 
     def _get_prf_data(self):
         module, output = channel_to_module_output(self.channel)
-        return get_kepler_prf_file(module=module, output=output)
+        return get_kepler_prf_file(module=module, output=output, cache_dir = self.cache_dir)
 
     def update_coordinates(self, targets: List[Tuple], shape: Tuple):
         row, column = self._unpack_targets(targets)

diff --git a/src/lkprf/tessprf.py b/src/lkprf/tessprf.py
@@ -4,6 +4,7 @@
 import numpy as np
 from .utils import LKPRFWarning
 from .data import get_tess_prf_file
+from . import PACKAGEDIR
 import warnings
 
 from .prfmodel import PRF
@@ -14,19 +15,20 @@ class TESSPRF(PRF):
     """A TESSPRF class. The TESS PRF measurements are supersampled by a factor of 9.
     Two PRF models were produced, one for sectors 1-3 and a second set for sectors 4+ """
 
-    def __init__(self, camera: int, ccd: int, sector: int = 4):
+    def __init__(self, camera: int, ccd: int, sector: int = 4, cache_dir: str = PACKAGEDIR + '/data/'):
         super().__init__()
         self.camera = camera
         self.ccd = ccd
         self.sector = sector
         self.mission = "TESS"
+        self.cache_dir = cache_dir
         self._prepare_prf()
 
     def __repr__(self):
         return f"TESSPRF Object [Camera {self.camera}, CCD {self.ccd}, Sector {self.sector}]"
 
     def _get_prf_data(self):
-        return get_tess_prf_file(camera=self.camera, ccd=self.ccd, sector=self.sector)
+        return get_tess_prf_file(camera=self.camera, ccd=self.ccd, sector=self.sector, cache_dir=self.cache_dir)
 
     def update_coordinates(self, targets: List[Tuple], shape: Tuple):
         row, column = self._unpack_targets(targets)