Add parallelization option to calculate_pfss_solution using Ray

solarmach/pfss_utilities.py

@@ -14,6 +14,7 @@
 import glob
 import os
 import pickle
+import ray
 
 import astropy.constants as aconst
 import astropy.units as u
@@ -561,24 +562,40 @@ def sphere(radius, clr, dist=0):
     return trace
 
 
-def calculate_pfss_solution(gong_map, rss, coord_sys, nrho=35):
+@ray.remote  # (num_cpus=50)
+def pfss_ray(pfss_in):
+    """
+    Perform a Potential Field Source Surface (PFSS) extrapolation.
+
+    Parameters:
+    pfss_in (object): Input data for the PFSS extrapolation. This should be compatible with the pfsspy.pfss function.
+
+    Returns:
+    object: The result of the PFSS extrapolation.
+    """
+    return pfsspy.pfss(pfss_in)
+
+
+def calculate_pfss_solution(gong_map, rss, coord_sys, nrho=35, parallelize=True):
     """
     Calculates a Potential Field Source Surface (PFSS) solution based on a GONG map and parameters.
 
     Parameters:
     -----------
-    gong_map : {SunPy Map}
-        GONG map as SunPy Map object obtained with get_gong_map()
+    gong_map : {sunpy.map.Map}
+        GONG map in Carrington or Stonyhurst coordinates, obtained with get_gong_map()
     rss : {float}
         source surface height in solar radii
     coord_sys: {str}
         cordinate system used: either 'car' or 'Carrington', or 'sto' or 'Stonyhurst'
-    nrho : {float/int}
-        rho = ln(r) -> nrho is the amount of points in this logarithmic range
+    nrho : {float/int, optional}
+        rho = ln(r) -> nrho is the amount of points in this logarithmic range. Default is 35.
+    parallelize : {bool, optional}
+        Whether to parallelize the computation using Ray. Default is True.
 
     Returns:
     ----------
-    pfss_solution : {pfsspy solution object}
+    pfss_solution : {pfsspy.Output}
         The PFSS solution that can be used to plot magnetic field lines under the source surface
     """
     # GONG map is in Carrington coordinates
@@ -603,8 +620,17 @@ def calculate_pfss_solution(gong_map, rss, coord_sys, nrho=35):
     # The pfss input object, assembled from a gong map, resolution (nrho) and source surface height (rss)
     pfss_in = pfsspy.Input(gong_map, nrho, rss)
 
-    # This is the pfss solution, calculated from the input object
-    pfss_solution = pfsspy.pfss(pfss_in)
+    if parallelize:
+        # You can configure this init call to point to a separate machine too.
+        ray.init()
+
+        # wait for results
+        future = pfss_ray.remote(pfss_in)
+
+        pfss_solution = ray.get(future)
+        ray.shutdown()
+    else:
+        pfss_solution = pfsspy.pfss(pfss_in)
 
     return pfss_solution