From 8d5c2bf62ab078f7748fda673c6e467a9a91d7cb Mon Sep 17 00:00:00 2001
From: Michael Hearne <mhearne@usgs.gov>
Date: Fri, 17 Jul 2020 07:15:25 -0600
Subject: [PATCH 1/2] Adding distance/az to magnitude data frames

---
 libcomcat/dataframes.py            | 56 +++++++++++++++++++++++++++++-
 tests/libcomcat/dataframes_test.py |  8 ++---
 2 files changed, 59 insertions(+), 5 deletions(-)

diff --git a/libcomcat/dataframes.py b/libcomcat/dataframes.py
index 1996469..3b85d2c 100644
--- a/libcomcat/dataframes.py
+++ b/libcomcat/dataframes.py
@@ -223,6 +223,46 @@ def get_arrival(event, pickid):
         return None
 
 
+def get_pick(event, waveid):
+    """Find the pick object in a Catalog Event corresponding to input waveform id.
+
+    Args:
+        event (Event): Obspy Catalog Event object.
+        ampid (str): Amplitude ID string.
+
+    Returns:
+      Amplitude: Obspy Catalog amplitude object.
+    """
+    if waveid is None:
+        return None
+    idlist = [pick.waveform_id for pick in event.picks]
+    if waveid not in idlist:
+        return None
+    idx = idlist.index(waveid)
+    pick = event.picks[idx]
+    return pick
+
+
+def get_amplitude(catevent, ampid):
+    """Find the amplitude object in a Catalog Event corresponding to input amplitude id.
+
+    Args:
+        event (Event): Obspy Catalog Event object.
+        ampid (str): Amplitude ID string.
+
+    Returns:
+      Amplitude: Obspy Catalog amplitude object.
+    """
+    if ampid is None:
+        return None
+    idlist = [amp.resource_id for amp in catevent.amplitudes]
+    if ampid not in idlist:
+        return None
+    idx = idlist.index(ampid)
+    amplitude = catevent.amplitudes[idx]
+    return amplitude
+
+
 def get_magnitude_data_frame(detail, catalog, magtype):
     """Return a Pandas DataFrame consisting of magnitude data.
 
@@ -247,7 +287,7 @@ def get_magnitude_data_frame(detail, catalog, magtype):
     """
     columns = columns = ['Channel', 'Type', 'Amplitude',
                          'Period', 'Status', 'Magnitude',
-                         'Weight']
+                         'Weight', 'Distance', 'Azimuth']
     df = pd.DataFrame(columns=columns)
     phasedata = detail.getProducts('phase-data', source=catalog)[0]
     quakeurl = phasedata.getContentURL('quakeml.xml')
@@ -298,6 +338,18 @@ def get_magnitude_data_frame(detail, catalog, magtype):
 
                 row['Channel'] = fmt % tpl
                 row['Type'] = smag.station_magnitude_type
+
+                distance = np.nan
+                azimuth = np.nan
+                amplitude = get_amplitude(catevent, smag.amplitude_id)
+                if amplitude is not None:
+                    waveform_id = amplitude.waveform_id
+                    pick = get_pick(catevent, waveform_id)
+                    if pick is not None:
+                        arrival = get_arrival(catevent, pick.resource_id)
+                        if arrival is not None:
+                            distance = arrival.distance
+                            azimuth = arrival.azimuth
                 if amp is not None:
                     row['Amplitude'] = amp.generic_amplitude
                     row['Period'] = amp.period
@@ -308,6 +360,8 @@ def get_magnitude_data_frame(detail, catalog, magtype):
                     row['Status'] = 'automatic'
                 row['Magnitude'] = smag.mag
                 row['Weight'] = contribution.weight
+                row['Distance'] = distance
+                row['Azimuth'] = azimuth
                 df = df.append(row, ignore_index=True)
     df = df[columns]
     return df
diff --git a/tests/libcomcat/dataframes_test.py b/tests/libcomcat/dataframes_test.py
index 464f70e..510cfce 100755
--- a/tests/libcomcat/dataframes_test.py
+++ b/tests/libcomcat/dataframes_test.py
@@ -26,9 +26,9 @@ def get_datadir():
     return cassettes, datadir
 
 
-def test_phase_dataframe():
+def test_magnitude_dataframe():
     cassettes, datadir = get_datadir()
-    tape_file = os.path.join(cassettes, 'dataframes_phase.yaml')
+    tape_file = os.path.join(cassettes, 'dataframes_magnitude.yaml')
     with vcr.use_cassette(tape_file):
         detail = get_event_by_id('us1000778i')  # 2016 NZ event
         df = get_magnitude_data_frame(detail, 'us', 'mb')
@@ -36,9 +36,9 @@ def test_phase_dataframe():
             df['Magnitude'].sum(), 756.8100000000001)
 
 
-def test_magnitude_dataframe():
+def test_phase_dataframe():
     cassettes, datadir = get_datadir()
-    tape_file = os.path.join(cassettes, 'dataframes_magnitude.yaml')
+    tape_file = os.path.join(cassettes, 'dataframes_phase.yaml')
     with vcr.use_cassette(tape_file):
         detail = get_event_by_id('us1000778i')  # 2016 NZ event
         df = get_phase_dataframe(detail, catalog='us')

From 80fa48b5bd696b516b9593e93b121279c850fb25 Mon Sep 17 00:00:00 2001
From: Michael Hearne <mhearne@usgs.gov>
Date: Fri, 17 Jul 2020 07:18:09 -0600
Subject: [PATCH 2/2] Updated docs to reflect new changes to mag dataframe

---
 docs/api.md             | 2 ++
 libcomcat/dataframes.py | 2 ++
 2 files changed, 4 insertions(+)

diff --git a/docs/api.md b/docs/api.md
index 5e15e3f..85ad9ae 100644
--- a/docs/api.md
+++ b/docs/api.md
@@ -391,6 +391,8 @@ The output is a dataframe with the following columns.
 - Status: Whether it is manual or automatic.
 - Magnitude: The magnitude derived locally for the station.
 - Weight: The weight of the magnitude.
+- Distance: The distance from the earthquake.
+- Azimuth: Azimuth (degrees) from epicenter to station.
 
 
 ### Phase Dataframe
diff --git a/libcomcat/dataframes.py b/libcomcat/dataframes.py
index 3b85d2c..056233a 100644
--- a/libcomcat/dataframes.py
+++ b/libcomcat/dataframes.py
@@ -281,6 +281,8 @@ def get_magnitude_data_frame(detail, catalog, magtype):
             - Status: "manual" or "automatic".
             - Magnitude: Locally determined magnitude.
             - Weight: Magnitude weight.
+            - Distance: Distance from epicenter to station.
+            - Azimuth (degrees) from epicenter to station.
     Raises:
         AttributeError if input DetailEvent does not have a phase-data product
             for the input catalog.