Fix GPU support for redrock

README.md

@@ -2,7 +2,7 @@
 [![DOCS](https://readthedocs.org/projects/redmost/badge/?version=latest)](https://redmost.readthedocs.io/en/latest/) [![Coverage Status](https://coveralls.io/repos/github/mauritiusdadd/redmost/badge.svg?branch=main)](https://coveralls.io/github/mauritiusdadd/redmost?branch=main) [![Testing PyQt6](https://github.com/mauritiusdadd/redmost/actions/workflows/test_linux_pyqt6.yml/badge.svg)](https://github.com/mauritiusdadd/redmost/actions/workflows/test_linux_pyqt6.yml) [![Testing PySide6](https://github.com/mauritiusdadd/redmost/actions/workflows/test_linux_pyside6.yml/badge.svg)](https://github.com/mauritiusdadd/redmost/actions/workflows/test_linux_pyside6.yml)
 # REDMOST 
 
-<b>RED</b>shift <b>M</b>easurement <b>O</b>f <b>S</b>pec<b>T</b>ra is a Qt6 Graphical User Interface (GUI) to do redshift measurements on 1D spectra.
+<b>RED</b>shift <b>M</b>easurements <b>O</b>f <b>S</b>pec<b>T</b>ra is a Qt6 Graphical User Interface (GUI) to do redshift measurements on 1D spectra.
 
 ![Redmost main window](https://github.com/mauritiusdadd/redmost/blob/main/docs/pics/main_ui_spectrum.png?raw=true)
 

src/redmost/backends/rrock.py

@@ -2,6 +2,7 @@
 import uuid
 import pickle
 import argparse
+from packaging import version
 from typing import Tuple, List, Dict, Optional, Any, Callable
 
 import numpy as np
@@ -14,6 +15,8 @@
 
 from specutils import Spectrum1D  # type: ignore
 
+RR_GPU_MIN_VER = "0.16.0"
+
 try:
     import redrock
 except (ImportError, ModuleNotFoundError):
@@ -191,26 +194,31 @@ def run_redrock(targets: List[Target]) -> Tuple[Any, Any]:
         # Get the dictionary of wavelength grids
         dwave = dtargets.wavegrids()
 
+        print(f"Using redrock version {redrock.__version__}")
+
         _ = elapsed(
             start,
             "Distribution of {} targets".format(len(dtargets.all_target_ids)),
             comm=None
         )
 
+        opt_zfind_args = {}
+        opt_load_dist_templates_args = {}
+        use_gpu = False
+        if version.parse(redrock.__version__) >= version.parse(RR_GPU_MIN_VER):
+            print(f"GPU support: {use_gpu}")
+            opt_zfind_args['use_gpu'] = use_gpu
+            opt_load_dist_templates_args['use_gpu'] = use_gpu
+
         # Read the template data
         dtemplates = load_dist_templates(
             dwave,
             templates=None,
             comm=None,
             mp_procs=1,
-            # use_gpu=True,
-            # gpu_mode=True
+            **opt_load_dist_templates_args
         )
 
-        opt_zfind_args = {}
-        # TODO: do more tests before enabling this
-        # opt_zfind_args['use_gpu'] = False
-
         # Compute the redshifts, including both the coarse scan and the
         # refinement.  This function only returns data on the rank 0 process.
         start = elapsed(None, "", comm=None)