diff --git a/solarmach/pfss_utilities.py b/solarmach/pfss_utilities.py index 49494e7..24d7bfa 100644 --- a/solarmach/pfss_utilities.py +++ b/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