From a4518523076fbdb8b89a7f76bece4728f599a3e2 Mon Sep 17 00:00:00 2001
From: jgieseler <jan.gieseler@utu.fi>
Date: Tue, 15 Aug 2023 17:40:56 +0300
Subject: [PATCH] introduce function sc_distance()

---
 solarmach/__init__.py   | 53 +++++++++++++++++++++++++++++++++++++++++
 solarmach/tests/test.py | 12 +++++++++-
 2 files changed, 64 insertions(+), 1 deletion(-)

diff --git a/solarmach/__init__.py b/solarmach/__init__.py
index c9d6574..1134ede 100644
--- a/solarmach/__init__.py
+++ b/solarmach/__init__.py
@@ -23,6 +23,7 @@
 from matplotlib.legend_handler import HandlerPatch
 from sunpy import log
 from sunpy.coordinates import frames, get_horizons_coord
+from sunpy.time import parse_time
 
 from solarmach.pfss_utilities import calculate_pfss_solution, get_field_line_coords, get_gong_map, multicolorline, sphere, spheric2cartesian, vary_flines
 
@@ -1531,6 +1532,58 @@ def pfss_3d(self, active_area=(None, None, None, None), color_code='object'):
         return
 
 
+def sc_distance(sc1, sc2, dtime):
+    """
+    Obtain absolute distance between two bodies in 3d for a given datetime.
+
+    Parameters
+    ----------
+    sc1 : str
+        Name of body 1, e.g., planet or spacecraft
+    sc2 : str
+        Name of body 2, e.g., planet or spacecraft
+    dtime : datetime object or datetime-compatible str
+        Date (and time) of distance determination
+
+    Returns
+    -------
+    astropy units
+        Absolute distance between body 1 and 2 in AU.
+    """
+    # parse datetime:
+    if type(dtime) == str:
+        try:
+            obstime = parse_time(dtime)
+        except ValueError:
+            print(f"Unable to extract datetime from '{dtime}'. Please try a different format.")
+            return np.nan*u.AU
+
+    # standardize body names (e.g. 'PSP' => 'Parker Solar Probe')
+    try:
+        sc1 = body_dict[sc1][1]
+    except KeyError:
+        pass
+    #
+    try:
+        sc2 = body_dict[sc2][1]
+    except KeyError:
+        pass
+
+    try:
+        sc1_coord = get_horizons_coord(sc1, obstime, None)
+    except ValueError:
+        print(f"Unable to obtain position for '{sc1}' at {dtime}. Please try a different name or date.")
+        return np.nan*u.AU
+    #
+    try:
+        sc2_coord = get_horizons_coord(sc2, obstime, None)
+    except ValueError:
+        print(f"Unable to obtain position for '{sc2}' at {dtime}. Please try a different name or date.")
+        return np.nan*u.AU
+
+    return sc1_coord.separation_3d(sc2_coord)
+
+
 def _isstreamlit():
     """
     Function to check whether python code is run within streamlit
diff --git a/solarmach/tests/test.py b/solarmach/tests/test.py
index 9f01dd5..7ba9f81 100644
--- a/solarmach/tests/test.py
+++ b/solarmach/tests/test.py
@@ -6,7 +6,7 @@
 import numpy as np
 import pandas
 import pfsspy
-from solarmach import SolarMACH, print_body_list, get_gong_map, calculate_pfss_solution
+from solarmach import SolarMACH, print_body_list, get_gong_map, calculate_pfss_solution, sc_distance
 
 
 def test_print_body_list():
@@ -87,3 +87,13 @@ def test_solarmach_pfss():
                            numbered_markers=True, long_sector=[290, 328], long_sector_vsw=[400, 600],
                            long_sector_color='red', reference_vsw=400.0)
     assert isinstance(fig, matplotlib.figure.Figure)
+
+
+def test_sc_distance():
+    distance = sc_distance('SolO', 'PSP', "2020/12/12")
+    assert np.round(distance.value, 8) == 1.45237361
+    assert distance.unit == u.AU
+    #
+    distance = sc_distance('SolO', 'PSP', "2000/12/12")
+    assert np.isnan(distance.value)
+    assert distance.unit == u.AU