From 638f19e993e7239cd79efa7e30afb6777ed12074 Mon Sep 17 00:00:00 2001
From: zssherman <shermanx1991@gmail.com>
Date: Tue, 14 Nov 2023 17:04:12 -0600
Subject: [PATCH 1/2] ADD: Adding arm site and facility location coordinate
elastic search.
---
act/tests/test_utils.py | 34 +++++++++++++++
act/utils/__init__.py | 1 +
act/utils/data_utils.py | 91 +++++++++++++++++++++++++++++++++++++++++
3 files changed, 126 insertions(+)
diff --git a/act/tests/test_utils.py b/act/tests/test_utils.py
index 6bb385065c..e1068ba1c7 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'] == None
+ assert test_dict_no_coord[list(test_dict_no_coord)[0]]['longitude'] == 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..376c25f82d 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_code + ' ' + facility, coords)
+
+ return coord_dict
From b176b84912791ef8a5155353a8dc013a8e0ee5b4 Mon Sep 17 00:00:00 2001
From: zssherman <shermanx1991@gmail.com>
Date: Tue, 14 Nov 2023 17:11:47 -0600
Subject: [PATCH 2/2] STY: PEP8 Fixes.
---
act/tests/test_utils.py | 4 ++--
act/utils/data_utils.py | 18 +++++++++---------
2 files changed, 11 insertions(+), 11 deletions(-)
diff --git a/act/tests/test_utils.py b/act/tests/test_utils.py
index e1068ba1c7..4f4b140f31 100644
--- a/act/tests/test_utils.py
+++ b/act/tests/test_utils.py
@@ -786,8 +786,8 @@ def test_arm_site_location_search():
# 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'] == None
- assert test_dict_no_coord[list(test_dict_no_coord)[0]]['longitude'] == None
+ 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)
diff --git a/act/utils/data_utils.py b/act/utils/data_utils.py
index 376c25f82d..77fca5f58b 100644
--- a/act/utils/data_utils.py
+++ b/act/utils/data_utils.py
@@ -1261,13 +1261,13 @@ def arm_site_location_search(site_code='sgp', facility_code=None):
},
},
},
- },
- "size": 0,
- "query": {
- "query_string": {
- "query": query,
- },
- },
+ },
+ "size": 0,
+ "query": {
+ "query_string": {
+ "query": query,
+ },
+ },
}
# Uses requests to grab metadata from arm.gov.
@@ -1292,6 +1292,6 @@ def arm_site_location_search(site_code='sgp', facility_code=None):
lon = float(lon)
coords = {'latitude': lat,
'longitude': lon}
- coord_dict.setdefault(site_code + ' ' + facility, coords)
-
+ coord_dict.setdefault(site + ' ' + facility, coords)
+
return coord_dict