diff --git a/act/tests/test_utils.py b/act/tests/test_utils.py index 6bb385065c..4f4b140f31 100644 --- a/act/tests/test_utils.py +++ b/act/tests/test_utils.py @@ -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 diff --git a/act/utils/__init__.py b/act/utils/__init__.py index e7cae20635..cd885b0cb1 100644 --- a/act/utils/__init__.py +++ b/act/utils/__init__.py @@ -21,6 +21,7 @@ 'height_adjusted_pressure', 'height_adjusted_temperature', 'convert_to_potential_temp', + 'arm_site_location_search', ], 'datetime_utils': [ 'dates_between', diff --git a/act/utils/data_utils.py b/act/utils/data_utils.py index 66242897f7..77fca5f58b 100644 --- a/act/utils/data_utils.py +++ b/act/utils/data_utils.py @@ -6,6 +6,7 @@ import importlib import warnings +import json import metpy import numpy as np import pint @@ -13,6 +14,7 @@ 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