ARM-DOE · zssherman · Nov 15, 2023 · Nov 14, 2023 · Nov 14, 2023
@@ -9,6 +9,7 @@
 import string
 import random
 import numpy as np
+from numpy.testing import assert_almost_equal
 import pandas as pd
 import pytest
 import pytz
@@ -761,3 +762,36 @@ def test_DatastreamParser():
     assert fn_obj.time is None
     assert fn_obj.ext is None
     del fn_obj
+
+
+def test_arm_site_location_search():
+    # Test for many facilities
+    test_dict_many = act.utils.arm_site_location_search(site_code='sgp', facility_code=None)
+    assert len(test_dict_many) > 30
+    assert list(test_dict_many)[0] == 'sgp A1'
+    assert list(test_dict_many)[2] == 'sgp A3'
+
+    assert_almost_equal(test_dict_many[list(test_dict_many)[0]]['latitude'], 37.843058)
+    assert_almost_equal(test_dict_many[list(test_dict_many)[0]]['longitude'], -97.020569)
+    assert_almost_equal(test_dict_many[list(test_dict_many)[2]]['latitude'], 37.626)
+    assert_almost_equal(test_dict_many[list(test_dict_many)[2]]['longitude'], -96.882)
+
+    # Test for one facility
+    test_dict_one = act.utils.arm_site_location_search(site_code='sgp', facility_code='I5')
+    assert len(test_dict_one) == 1
+    assert list(test_dict_one)[0] == 'sgp I5'
+    assert_almost_equal(test_dict_one[list(test_dict_one)[0]]['latitude'], 36.491178)
+    assert_almost_equal(test_dict_one[list(test_dict_one)[0]]['longitude'], -97.593936)
+
+    # Test for a facility with no latitude and longitude information
+    test_dict_no_coord = act.utils.arm_site_location_search(site_code='sgp', facility_code='A6')
+    assert list(test_dict_no_coord)[0] == 'sgp A6'
+    assert test_dict_no_coord[list(test_dict_no_coord)[0]]['latitude'] is None
+    assert test_dict_no_coord[list(test_dict_no_coord)[0]]['longitude'] is None
+
+    # Test for another site
+    test_dict_nsa = act.utils.arm_site_location_search(site_code='nsa', facility_code=None)
+    assert len(test_dict_nsa) > 5
+    assert list(test_dict_nsa)[0] == 'nsa C1'
+    assert test_dict_nsa[list(test_dict_nsa)[0]]['latitude'] == 71.323
+    assert test_dict_nsa[list(test_dict_nsa)[0]]['longitude'] == -156.615
@@ -21,6 +21,7 @@
             'height_adjusted_pressure',
             'height_adjusted_temperature',
             'convert_to_potential_temp',
+            'arm_site_location_search',
         ],
         'datetime_utils': [
             'dates_between',

@@ -6,13 +6,15 @@
 import importlib
 import warnings
 
+import json
 import metpy
 import numpy as np
 import pint
 import scipy.stats as stats
 import xarray as xr
 from pathlib import Path
 import re
+import requests
 
 spec = importlib.util.find_spec('pyart')
 if spec is not None:
@@ -1204,3 +1206,92 @@ def height_adjusted_pressure(
     adjusted_pressure = adjusted_pressure.to(press_var_units).magnitude
 
     return adjusted_pressure
+
+
+def arm_site_location_search(site_code='sgp', facility_code=None):
+    """
+    Parameters
+    ----------
+    site_code : str
+        ARM site code to retrieve facilities and coordinate information. Example and default
+        is 'sgp'.
+    facility_code : str or None
+        Facility code or codes for the ARM site provided. If None is provided, all facilities are returned.
+        Example string for multiple facilities is 'A4,I5'.
+
+    Returns
+    -------
+    coord_dict : dict
+        A dictionary containing the facility chosen coordinate information or all facilities
+        if None for facility_code and their respective coordinates.
+
+    """
+    headers = {
+        'Content-Type': 'application/json',
+    }
+    # Return all facilities if facility_code is None else set the query to include
+    # facility search
+    if facility_code is None:
+        query = "site_code:" + site_code
+    else:
+        query = "site_code:" + site_code + "AND facility_code:" + facility_code
+
+    # Search aggregation for elastic search
+    json_data = {
+        "aggs": {
+            "distinct_facility_code": {
+                "terms": {
+                    "field": "facility_code.keyword",
+                    "order": {
+                        "_key": "asc"
+                    },
+                    "size": 7000,
+                },
+                "aggs": {
+                    "hits": {
+                        "top_hits": {
+                            "_source": [
+                                "site_type",
+                                "site_code",
+                                "facility_code",
+                                "location",
+                            ],
+                            "size": 1
+                        },
+                    },
+                },
+            },
+        },
+        "size": 0,
+        "query": {
+            "query_string": {
+                "query": query,
+            },
+        },
+    }
+
+    # Uses requests to grab metadata from arm.gov.
+    response = requests.get('https://adc.arm.gov/elastic/metadata/_search', headers=headers, json=json_data)
+    # Loads the text to a dictionary
+    response_dict = json.loads(response.text)
+
+    # Searches dictionary for the site, facility and coordinate information.
+    coord_dict = {}
+    # Loop through each facility.
+    for i in range(len(response_dict['aggregations']['distinct_facility_code']['buckets'])):
+        site_info = response_dict['aggregations']['distinct_facility_code']['buckets'][i]['hits']['hits']['hits'][0]['_source']
+        site = site_info['site_code']
+        facility = site_info['facility_code']
+        # Some sites do not contain coordinate information, return None if that is the case.
+        if site_info['location'] is None:
+            coords = {'latitude': None,
+                      'longitude': None}
+        else:
+            lat, lon = site_info['location'].split(',')
+            lat = float(lat)
+            lon = float(lon)
+            coords = {'latitude': lat,
+                      'longitude': lon}
+        coord_dict.setdefault(site + ' ' + facility, coords)
+
+    return coord_dict