diff --git a/REQUIREMENTS.txt b/REQUIREMENTS.txt new file mode 100644 index 00000000..ba365124 --- /dev/null +++ b/REQUIREMENTS.txt @@ -0,0 +1,7 @@ +arm_pyart +numpy +scipy +cartopy +pandas +dask +six diff --git a/pydda/constraints/__init__.py b/pydda/constraints/__init__.py index 62639719..6e9bbe35 100644 --- a/pydda/constraints/__init__.py +++ b/pydda/constraints/__init__.py @@ -20,6 +20,9 @@ make_constraint_from_wrf add_hrrr_constraint_to_grid + make_constraint_from_era_interim + download_needed_era_data + get_iem_obs """ @@ -27,3 +30,4 @@ from .model_data import add_hrrr_constraint_to_grid from .model_data import make_constraint_from_era_interim from .model_data import download_needed_era_data +from .station_data import get_iem_obs diff --git a/pydda/constraints/station_data.py b/pydda/constraints/station_data.py new file mode 100644 index 00000000..e389b607 --- /dev/null +++ b/pydda/constraints/station_data.py @@ -0,0 +1,149 @@ +import json +import pandas as pd +import warnings +import numpy as np +import pyart +import time + +from datetime import datetime, timedelta +from six import StringIO + +try: + from urllib.request import urlopen +except ImportError: + from urllib2 import urlopen + + +def get_iem_obs(Grid, window=60.): + """ + Returns all of the station observations from the Iowa Mesonet for a given Grid in the format + needed for PyDDA. + + Parameters + ---------- + Grid: pyART Grid + The Grid to retrieve the station data for. + window: float + The window (in minutes) to look for the nearest observation in time. + + Returns + ------- + station_data: list of dicts + A list of dictionaries containing the following entries as keys: + + *lat* - Latitude of the site (float) + + *lon* - Longitude of the site (float) + + *u* - u wind at the site (float) + + *v* - v wind at the site (float) + + *w* - w wind at the site (assumed to be 0) (float) + + *site_id* - Station ID (string) + + *x*, *y*, *z* - The (x, y, z) coordinates of the site in the Grid. (floats) + """ + + # First query the database for all of the JSON info for every station + # Only add stations whose lat/lon are within the Grid's boundaries + regions = """AF AL_ AI_ AQ_ AG_ AR_ AK AL AM_ AO_ AS_ AR AW_ AU_ AT_ + AZ_ BA_ BE_ BB_ BG_ BO_ BR_ BF_ BT_ BS_ BI_ BM_ BB_ BY_ BZ_ BJ_ BW_ AZ CA CA_AB + CA_BC CD_ CK_ CF_ CG_ CL_ CM_ CO CO_ CN_ CR_ CT CU_ CV_ CY_ CZ_ DE DK_ DJ_ DM_ DO_ + DZ EE_ ET_ FK_ FM_ FJ_ FI_ FR_ GF_ PF_ GA_ GM_ GE_ DE_ GH_ GI_ KY_ GB_ GR_ GL_ GD_ + GU_ GT_ GN_ GW_ GY_ HT_ HN_ HK_ HU_ IS_ IN_ ID_ IR_ IQ_ IE_ IL_ IT_ CI_ JM_ JP_ + JO_ KZ_ KE_ KI_ KW_ LA_ LV_ LB_ LS_ LR_ LY_ LT_ LU_ MK_ MG_ MW_ MY_ MV_ ML_ CA_MB + MH_ MR_ MU_ YT_ MX_ MD_ MC_ MA_ MZ_ MM_ NA_ NP_ AN_ NL_ CA_NB NC_ CA_NF NF_ NI_ + NE_ NG_ MP_ KP_ CA_NT NO_ CA_NS CA_NU OM_ CA_ON PK_ PA_ PG_ PY_ PE_ PH_ PN_ PL_ + PT_ CA_PE PR_ QA_ CA_QC RO_ RU_RW_ SH_ KN_ LC_ VC_ WS_ ST_ CA_SK SA_ SN_ RS_ SC_ + SL_ SG_ SK_ SI_ SB_ SO_ ZA_ KR_ ES_ LK_ SD_ SR_ SZ_ SE_ CH_ SY_ TW_ TJ_ TZ_ TH_ + TG_ TO_ TT_ TU TN_ TR_ TM_ UG_ UA_ AE_ UN_ UY_ UZ_ VU_ VE_ VN_ VI_ YE_ CA_YT ZM_ ZW_ + EC_ EG_ FL GA GQ_ HI HR_ IA ID IL IO_ IN KS KH_ KY KM_ LA MA MD ME + MI MN MO MS MT NC ND NE NH NJ NM NV NY OH OK OR PA RI SC SV_ SD TD_ TN TX UT VA VT VG_ + WA WI WV WY""" + + networks = ["AWOS"] + grid_lon_min = Grid.point_longitude["data"].min() + grid_lon_max = Grid.point_longitude["data"].max() + grid_lat_min = Grid.point_latitude["data"].min() + grid_lat_max = Grid.point_latitude["data"].max() + for region in regions.split(): + networks.append("%s_ASOS" % (region,)) + + site_list = [] + elevations = [] + for network in networks: + # Get metadata + uri = ("https://mesonet.agron.iastate.edu/" "geojson/network/%s.geojson" + ) % (network,) + data = urlopen(uri) + jdict = json.load(data) + for site in jdict["features"]: + lat = site["geometry"]["coordinates"][1] + lon = site["geometry"]["coordinates"][0] + if lat >= grid_lat_min and lat <= grid_lat_max and lon >= grid_lon_min and lon <= grid_lon_max: + site_list.append((site["properties"]["sid"], site["properties"]["elevation"])) + + + # Get the timestamp for each request + grid_time = datetime.strptime(Grid.time["units"], + "seconds since %Y-%m-%dT%H:%M:%SZ") + start_time = grid_time - timedelta(minutes=window / 2.) + end_time = grid_time + timedelta(minutes=window / 2.) + + SERVICE = "http://mesonet.agron.iastate.edu/cgi-bin/request/asos.py?" + service = SERVICE + "data=all&tz=Etc/UTC&format=comma&latlon=yes&" + + service += start_time.strftime("year1=%Y&month1=%m&day1=%d&") + service += end_time.strftime("year2=%Y&month2=%m&day2=%d&") + station_obs = [] + for stations, elevations in site_list: + uri = "%s&station=%s" % (service, stations) + print("Downloading: %s" % (stations,)) + data = _download_data(uri) + buf = StringIO() + buf.write(data) + buf.seek(0) + + my_df = pd.read_csv(buf, skiprows=5) + stat_dict = {} + if len(my_df['lat'].values) == 0: + warnings.warn( + "No data available at station %s between time %s and %s" % + (stations, start_time.strftime('%Y-%m-%d %H:%M:%S'), + end_time.strftime('%Y-%m-%d %H:%M:%S'))) + else: + stat_dict['lat'] = my_df['lat'].values[0] + stat_dict['lon'] = my_df['lon'].values[0] + stat_dict['x'], stat_dict['y'] = pyart.core.geographic_to_cartesian( + stat_dict['lon'], stat_dict['lat'], Grid.get_projparams()) + stat_dict['x'] = stat_dict['x'][0] + stat_dict['y'] = stat_dict['y'][0] + stat_dict['z'] = elevations - Grid.origin_altitude["data"][0] + if my_df['drct'].values[0] == 'M': + continue + drct = float(my_df['drct'].values[0]) + s_ms = float(my_df['sknt'].values[0]) * 0.514444 + stat_dict['u'] = -np.sin(np.deg2rad(drct)) * s_ms + stat_dict['v'] = -np.cos(np.deg2rad(drct)) * s_ms + stat_dict['site_id'] = stations + station_obs.append(stat_dict) + buf.close() + + return station_obs + +def _download_data(uri): + attempt = 0 + while attempt < 6: + try: + data = urlopen(uri, timeout=300).read().decode("utf-8") + if data is not None and not data.startswith("ERROR"): + return data + except Exception as exp: + print("download_data(%s) failed with %s" % (uri, exp)) + time.sleep(5) + attempt += 1 + + print("Exhausted attempts to download, returning empty data") + return "" diff --git a/pydda/cost_functions/__init__.py b/pydda/cost_functions/__init__.py index 093c238f..50ed67f8 100644 --- a/pydda/cost_functions/__init__.py +++ b/pydda/cost_functions/__init__.py @@ -72,6 +72,8 @@ calculate_model_cost calculate_model_gradient calculate_fall_speed + calculate_point_cost + calculate_point_gradient """ @@ -88,4 +90,5 @@ from .cost_functions import calculate_vertical_vorticity_gradient from .cost_functions import calculate_model_cost from .cost_functions import calculate_model_gradient +from .cost_functions import calculate_point_cost, calculate_point_gradient from .cost_functions import J_function, grad_J diff --git a/pydda/cost_functions/cost_functions.py b/pydda/cost_functions/cost_functions.py index e5029158..7dd8af75 100644 --- a/pydda/cost_functions/cost_functions.py +++ b/pydda/cost_functions/cost_functions.py @@ -3,99 +3,22 @@ import scipy.ndimage.filters -def J_function(winds, vrs, azs, els, wts, u_back, v_back, u_model, - v_model, w_model, Co, Cm, Cx, Cy, Cz, Cb, Cv, Cmod, - Ut, Vt, grid_shape, dx, dy, dz, z, rmsVr, weights, - bg_weights, model_weights, upper_bc, - print_out=False): +def J_function(winds, parameters): """ Calculates the total cost function. This typically does not need to be called directly as get_dd_wind_field is a wrapper around this function and - grad_J. In order to add more terms to the cost function, modify this - function and grad_J. + :py:func:`pydda.cost_functions.grad_J`. + In order to add more terms to the cost function, modify this + function and :py:func:`pydda.cost_functions.grad_J`. Parameters ---------- winds: 1-D float array The wind field, flattened to 1-D for f_min. The total size of the array will be a 1D array of 3*nx*ny*nz elements. - vrs: List of 3D float arrays - List of radial velocities from each radar. All arrays in list - must have the same dimensions. - azs: List of 3D float arrays - List of azimuths from each radar. All arrays in list - must have the same dimensions. - els: List of 3D float arrays - List of elevations from each radar. All arrays in list - must have the same dimensions. - wts: List of 3D float arrays - Float array containing fall speeds from radar. All arrays in list - must have the same dimensions. - u_back: 1D float array (number of vertical levels): - Background u wind. This takes in a 1D float array of length nz, - with each element corresponding to the u component of the wind - at a given vertical level from the sounding. - v_back: 1D float array (number of vertical levels): - Background v wind. This takes in a 1D float array of length nz, - with each element corresponding to the v component of the wind - at a given vertical level from the sounding. - u_model: list of 3D float arrays - U from each model integrated into the retrieval. The U from - each model is given as a list of array of the same dimensions, - with the U from the model interpolated on to the radar analysis grid. - v_model: list of 3D float arrays - V from each model integrated into the retrieval. The V from each model - is given as a list of array of the same dimensions, with the V from - the model interpolated on to the radar analysis grid. - w_model: - W from each model integrated into the retrieval. The W from each model - is given as a list of array of the same dimensions, with the W from - the model interpolated on to the radar analysis grid. - Co: float - Weighting coefficient for data constraint. - Cm: float - Weighting coefficient for mass continuity constraint. - Cx: float - Smoothing coefficient for x-direction - Cy: float - Smoothing coefficient for y-direction - Cz: float - Smoothing coefficient for z-direction - Cb: float - Coefficient for sounding constraint - Cv: float - Weight for cost function related to vertical vorticity equation. - Cmod: float - Coefficient for model constraint - Ut: float - Prescribed storm motion. This is only needed if Cv is not zero. - Vt: float - Prescribed storm motion. This is only needed if Cv is not zero. - grid_shape: - Shape of wind grid - dx: - Spacing of grid in x direction - dy: - Spacing of grid in y direction - dz: - Spacing of grid in z direction - z: - Grid vertical levels in m - rmsVr: float - The sum of squares of velocity/num_points. Use for normalization - of data weighting coefficient - weights: n_radars by z_bins by y_bins by x_bins float array - Data weights for each pair of radars. This is usually automatically - determined by get_dd_wind_field. - bg_weights: z_bins by y_bins by x_bins float array - Data weights for sounding constraint. - model_weights: n_models by z_bins by y_bins by x_bins float array - Data weights for each model. - upper_bc: bool - True to enforce w=0 at top of domain (impermeability condition), - False to not enforce impermeability at top of domain - print_out: bool - Set to True to print out the value of the cost function. + parameters: DDParameters + The parameters for the cost function evaluation as specified by the + :py:func:`pydda.retrieval.DDParameters` class. Returns ------- @@ -103,132 +26,90 @@ def J_function(winds, vrs, azs, els, wts, u_back, v_back, u_model, The value of the cost function """ winds = np.reshape(winds, - (3, grid_shape[0], grid_shape[1], grid_shape[2])) + (3, parameters.grid_shape[0], parameters.grid_shape[1], + parameters.grid_shape[2])) Jvel = calculate_radial_vel_cost_function( - vrs, azs, els, winds[0], winds[1], winds[2], wts, rmsVr=rmsVr, - weights=weights, coeff=Co) + parameters.vrs, parameters.azs, parameters.els, + winds[0], winds[1], winds[2], parameters.wts, rmsVr=parameters.rmsVr, + weights=parameters.weights, coeff=parameters.Co) - if(Cm > 0): + if(parameters.Cm > 0): Jmass = calculate_mass_continuity( - winds[0], winds[1], winds[2], z, dx, dy, dz, coeff=Cm) + winds[0], winds[1], winds[2], parameters.z, + parameters.dx, parameters.dy, parameters.dz, + coeff=parameters.Cm) else: Jmass = 0 - if(Cx > 0 or Cy > 0 or Cz > 0): + if(parameters.Cx > 0 or parameters.Cy > 0 or parameters.Cz > 0): Jsmooth = calculate_smoothness_cost( - winds[0], winds[1], winds[2], Cx=Cx, Cy=Cy, Cz=Cz) + winds[0], winds[1], winds[2], Cx=parameters.Cx, + Cy=parameters.Cy, Cz=parameters.Cz) else: Jsmooth = 0 - if(Cb > 0): + if(parameters.Cb > 0): Jbackground = calculate_background_cost( - winds[0], winds[1], winds[2], bg_weights, u_back, v_back, Cb) + winds[0], winds[1], winds[2], parameters.bg_weights, + parameters.u_back, parameters.v_back, parameters.Cb) else: Jbackground = 0 - if(Cv > 0): + if(parameters.Cv > 0): Jvorticity = calculate_vertical_vorticity_cost( - winds[0], winds[1], winds[2], dx, dy, dz, Ut, Vt, coeff=Cv) + winds[0], winds[1], winds[2], parameters.dx, + parameters.dy, parameters.dz, parameters.Ut, + parameters.Vt, coeff=parameters.Cv) else: Jvorticity = 0 - if(Cmod > 0): + if(parameters.Cmod > 0): Jmod = calculate_model_cost( - winds[0], winds[1], winds[2], model_weights, u_model, v_model, - w_model, coeff=Cmod) + winds[0], winds[1], winds[2], + parameters.model_weights, parameters.u_model, + parameters.v_model, + parameters.w_model, coeff=parameters.Cmod) else: Jmod = 0 - if(print_out is True): - print(('| Jvel | Jmass | Jsmooth | Jbg | Jvort | Jmodel ' + - '| Max w ')) + if parameters.Cpoint > 0: + Jpoint = calculate_point_cost( + winds[0], winds[1], parameters.x, parameters.y, parameters.z, + parameters.point_list, Cp=parameters.Cpoint, roi=parameters.roi) + else: + Jpoint = 0 + + if(parameters.print_out is True): + print(('| Jvel | Jmass | Jsmooth | Jbg | Jvort | Jmodel | Jpoint |' + + ' Max w ')) print(('|' + "{:9.4f}".format(Jvel) + '|' + "{:9.4f}".format(Jmass) + '|' + "{:9.4f}".format(Jsmooth) + '|' + "{:9.4f}".format(Jbackground) + '|' + "{:9.4f}".format(Jvorticity) + '|' + "{:9.4f}".format(Jmod) + '|' + - "{:9.4f}".format(np.abs(winds[2]).max()))) + "{:9.4f}".format(Jpoint)) + '|' + + "{:9.4f}".format(np.ma.max(np.ma.abs(winds[2])))) - return Jvel + Jmass + Jsmooth + Jbackground + Jvorticity + Jmod + return Jvel + Jmass + Jsmooth + Jbackground + Jvorticity + Jmod + Jpoint -def grad_J(winds, vrs, azs, els, wts, u_back, v_back, u_model, - v_model, w_model, Co, Cm, Cx, Cy, Cz, Cb, Cv, Cmod, - Ut, Vt, grid_shape, dx, dy, dz, z, rmsVr, - weights, bg_weights, model_weights, upper_bc, - print_out=False): +def grad_J(winds, parameters): """ Calculates the gradient of the cost function. This typically does not need to be called directly as get_dd_wind_field is a wrapper around this - function and J_function. In order to add more terms to the cost function, - modify this function and grad_J. + function and :py:func:`pydda.cost_functions.J_function`. + In order to add more terms to the cost function, + modify this function and :py:func:`pydda.cost_functions.grad_J`. Parameters ---------- winds: 1-D float array The wind field, flattened to 1-D for f_min - vrs: List of float arrays - List of radial velocities from each radar - azs: List of float arrays - List of azimuths from each radar - els: List of float arrays - List of elevations from each radar - wts: List of float arrays - Float array containing fall speed from radar. - u_back: 1D float array (number of vertical levels): - Background u wind - v_back: 1D float array (number of vertical levels): - Background u wind - u_model: list of 3D float arrays - U from each model integrated into the retrieval - v_model: list of 3D float arrays - V from each model integrated into the retrieval - w_model: - W from each model integrated into the retrieval - Co: float - Weighting coefficient for data constraint. - Cm: float - Weighting coefficient for mass continuity constraint. - Cx: float - Smoothing coefficient for x-direction - Cy: float - Smoothing coefficient for y-direction - Cz: float - Smoothing coefficient for z-direction - Cb: float - Coefficient for sounding constraint - Cv: float - Weight for cost function related to vertical vorticity equation. - Cmod: float - Coefficient for model constraint - Ut: float - Prescribed storm motion. This is only needed if Cv is not zero. - Vt: float - Prescribed storm motion. This is only needed if Cv is not zero. - grid_shape: - Shape of wind grid - dx: - Spacing of grid in x direction - dy: - Spacing of grid in y direction - dz: - Spacing of grid in z direction - z: - Grid vertical levels in m - rmsVr: float - The sum of squares of velocity/num_points. Use for normalization - of data weighting coefficient - weights: n_radars by z_bins by y_bins x x_bins float array - Data weights for each pair of radars - bg_weights: z_bins by y_bins x x_bins float array - Data weights for sounding constraint - model_weights: n_models by z_bins by y_bins by x_bins float array - Data weights for each model. - upper_bc: bool - True to enforce w=0 at top of domain (impermeability condition), - False to not enforce impermeability at top of domain + parameters: DDParameters + The parameters for the cost function evaluation as specified by the + :py:func:`pydda.retrieve.DDParameters` class. Returns ------- @@ -236,36 +117,48 @@ def grad_J(winds, vrs, azs, els, wts, u_back, v_back, u_model, Gradient vector of cost function """ winds = np.reshape(winds, - (3, grid_shape[0], grid_shape[1], grid_shape[2])) + (3, parameters.grid_shape[0], + parameters.grid_shape[1], parameters.grid_shape[2])) grad = calculate_grad_radial_vel( - vrs, els, azs, winds[0], winds[1], winds[2], wts, weights, - rmsVr, coeff=Co, upper_bc=upper_bc) + parameters.vrs, parameters.els, parameters.azs, + winds[0], winds[1], winds[2], parameters.wts, parameters.weights, + parameters.rmsVr, coeff=parameters.Co, upper_bc=parameters.upper_bc) - if(Cm > 0): + if(parameters.Cm > 0): grad += calculate_mass_continuity_gradient( - winds[0], winds[1], winds[2], z, dx, dy, dz, coeff=Cm, - upper_bc=upper_bc) + winds[0], winds[1], winds[2], parameters.z, + parameters.dx, parameters.dy, parameters.dz, + coeff=parameters.Cm, upper_bc=parameters.upper_bc) - if(Cx > 0 or Cy > 0 or Cz > 0): + if(parameters.Cx > 0 or parameters.Cy > 0 or parameters.Cz > 0): grad += calculate_smoothness_gradient( - winds[0], winds[1], winds[2], Cx=Cx, Cy=Cy, Cz=Cz, - upper_bc=upper_bc) + winds[0], winds[1], winds[2], Cx=parameters.Cx, + Cy=parameters.Cy, Cz=parameters.Cz, upper_bc=parameters.upper_bc) - if(Cb > 0): + if(parameters.Cb > 0): grad += calculate_background_gradient( - winds[0], winds[1], winds[2], bg_weights, u_back, v_back, Cb, - upper_bc=upper_bc) + winds[0], winds[1], winds[2], parameters.bg_weights, + parameters.u_back, parameters.v_back, parameters.Cb, + upper_bc=parameters.upper_bc) - if(Cv > 0): + if(parameters.Cv > 0): grad += calculate_vertical_vorticity_gradient( - winds[0], winds[1], winds[2], dx, dy, dz, Ut, Vt, coeff=Cv) + winds[0], winds[1], winds[2], parameters.dx, + parameters.dy, parameters.dz, parameters.Ut, + parameters.Vt, coeff=parameters.Cv) - if(Cmod > 0): + if(parameters.Cmod > 0): grad += calculate_model_gradient( - winds[0], winds[1], winds[2], model_weights, u_model, v_model, - w_model, coeff=Cmod) + winds[0], winds[1], winds[2], + parameters.model_weights, parameters.u_model, parameters.v_model, + parameters.w_model, coeff=parameters.Cmod) - if(print_out is True): + if parameters.Cpoint > 0: + grad += calculate_point_gradient( + winds[0], winds[1], parameters.x, parameters.y, parameters.z, + parameters.point_list, Cp=parameters.Cpoint, roi=parameters.roi) + + if(parameters.print_out is True): print('Norm of gradient: ' + str(np.linalg.norm(grad, np.inf))) return grad @@ -522,6 +415,105 @@ def calculate_smoothness_gradient(u, v, w, Cx=1e-5, Cy=1e-5, Cz=1e-5, return y.flatten() +def calculate_point_cost(u, v, x, y, z, point_list, Cp=1e-3, roi=500.0): + """ + Calculates the cost function related to point observations. A mean square error cost + function term is applied to points that are within the sphere of influence + whose radius is determined by *roi*. + + Parameters + ---------- + u: Float array + Float array with u component of wind field + v: Float array + Float array with v component of wind field + x: Float array + X coordinates of grid centers + y: Float array + Y coordinates of grid centers + z: Float array + Z coordinated of grid centers + point_list: list of dicts + List of point constraints. + Each member is a dict with keys of "u", "v", to correspond + to each component of the wind field and "x", "y", "z" + to correspond to the location of the point observation. + + In addition, "site_id" gives the METAR code (or name) to the station. + Cp: float + The weighting coefficient of the point cost function. + roi: float + Radius of influence of observations + + Returns + ------- + J: float + The cost function related to the difference between wind field and points. + + """ + J = 0.0 + for the_point in point_list: + # Instead of worrying about whole domain, just find points in radius of influence + # Since we know that the weight will be zero outside the sphere of influence anyways + the_box = np.where(np.logical_and.reduce( + (np.abs(x - the_point["x"]) < roi, np.abs(y - the_point["y"]) < roi, + np.abs(z - the_point["z"]) < roi))) + J += np.sum(((u[the_box] - the_point["u"])**2 + (v[the_box] - the_point["v"])**2)) + + return J * Cp + + +def calculate_point_gradient(u, v, x, y, z, point_list, Cp=1e-3, roi=500.0): + """ + Calculates the gradient of the cost function related to point observations. + A mean square error cost function term is applied to points that are within the sphere of influence + whose radius is determined by *roi*. + + Parameters + ---------- + u: Float array + Float array with u component of wind field + v: Float array + Float array with v component of wind field + x: Float array + X coordinates of grid centers + y: Float array + Y coordinates of grid centers + z: Float array + Z coordinated of grid centers + point_list: list of dicts + List of point constraints. Each member is a dict with keys of "u", "v", + to correspond to each component of the wind field and "x", "y", "z" + to correspond to the location of the point observation. + + In addition, "site_id" gives the METAR code (or name) to the station. + Cp: float + The weighting coefficient of the point cost function. + roi: float + Radius of influence of observations + + Returns + ------- + gradJ: float array + The gradient of the cost function related to the difference between wind field and points. + + """ + + gradJ_u = np.zeros_like(u) + gradJ_v = np.zeros_like(v) + gradJ_w = np.zeros_like(u) + + for the_point in point_list: + the_box = np.where(np.logical_and.reduce( + (np.abs(x - the_point["x"]) < roi, np.abs(y - the_point["y"]) < roi, + np.abs(z - the_point["z"]) < roi))) + gradJ_u[the_box] += 2 * (u[the_box] - the_point["u"]) + gradJ_v[the_box] += 2 * (v[the_box] - the_point["v"]) + + gradJ = np.stack([gradJ_u, gradJ_v, gradJ_w], axis=0).flatten() + return gradJ * Cp + + def calculate_mass_continuity(u, v, w, z, dx, dy, dz, coeff=1500.0, anel=1): """ Calculates the mass continuity cost function by taking the divergence @@ -826,7 +818,7 @@ def calculate_vertical_vorticity_gradient(u, v, w, dx, dy, dz, Ut, Vt, coeff=1e-5): """ Calculates the gradient of the cost function due to deviance from vertical - vorticity equation. This is done by taking the functional deriviative of + vorticity equation. This is done by taking the functional derivative of the vertical vorticity cost function. Parameters @@ -993,7 +985,7 @@ def calculate_model_gradient(u, v, w, weights, u_model, v_model: list of 3D float arrays Meridional wind field from model w_model: list of 3D float arrays - Vetical wind field from model + Vertical wind field from model coeff: float Weight of background constraint to total cost function diff --git a/pydda/retrieval/__init__.py b/pydda/retrieval/__init__.py index 3e75115e..72193c73 100644 --- a/pydda/retrieval/__init__.py +++ b/pydda/retrieval/__init__.py @@ -18,9 +18,11 @@ get_dd_wind_field get_dd_wind_field_nested get_bca + DDParameters """ from .wind_retrieve import get_dd_wind_field from .wind_retrieve import get_bca +from .wind_retrieve import DDParameters from .nesting import get_dd_wind_field_nested diff --git a/pydda/retrieval/wind_retrieve.py b/pydda/retrieval/wind_retrieve.py index b509fb30..c1d4d194 100644 --- a/pydda/retrieval/wind_retrieve.py +++ b/pydda/retrieval/wind_retrieve.py @@ -20,17 +20,145 @@ from .angles import add_azimuth_as_field, add_elevation_as_field -def get_dd_wind_field(Grids, u_init, v_init, w_init, vel_name=None, +class DDParameters(object): + """ + This is a helper class for inserting more arguments into the :func:`pydda.cost_functions.J_function` and + :func:`pydda.cost_functions.grad_J` function. Since these cost functions take numerous parameters, this class + will store the needed parameters as one positional argument for easier readability of the code. + + In addition, class members can be added here so that those contributing more constraints to the variational + framework can add any parameters they may need. + + Attributes + ---------- + vrs: List of float arrays + List of radial velocities from each radar + azs: List of float arrays + List of azimuths from each radar + els: List of float arrays + List of elevations from each radar + wts: List of float arrays + Float array containing fall speed from radar. + u_back: 1D float array (number of vertical levels) + Background u wind + v_back: 1D float array (number of vertical levels) + Background u wind + u_model: list of 3D float arrays + U from each model integrated into the retrieval + v_model: list of 3D float arrays + V from each model integrated into the retrieval + w_model: + W from each model integrated into the retrieval + Co: float + Weighting coefficient for data constraint. + Cm: float + Weighting coefficient for mass continuity constraint. + Cx: float + Smoothing coefficient for x-direction + Cy: float + Smoothing coefficient for y-direction + Cz: float + Smoothing coefficient for z-direction + Cb: float + Coefficient for sounding constraint + Cv: float + Weight for cost function related to vertical vorticity equation. + Cmod: float + Coefficient for model constraint + Cpoint: float + Coefficient for point constraint + Ut: float + Prescribed storm motion. This is only needed if Cv is not zero. + Vt: float + Prescribed storm motion. This is only needed if Cv is not zero. + grid_shape: + Shape of wind grid + dx: + Spacing of grid in x direction + dy: + Spacing of grid in y direction + dz: + Spacing of grid in z direction + x: + E-W grid levels in m + y: + N-S grid levels in m + z: + Grid vertical levels in m + rmsVr: float + The sum of squares of velocity/num_points. Use for normalization + of data weighting coefficient + weights: n_radars by z_bins by y_bins x x_bins float array + Data weights for each pair of radars + bg_weights: z_bins by y_bins x x_bins float array + Data weights for sounding constraint + model_weights: n_models by z_bins by y_bins by x_bins float array + Data weights for each model. + point_list: list or None + point_list: list of dicts + List of point constraints. Each member is a dict with keys of "u", "v", + to correspond to each component of the wind field and "x", "y", "z" + to correspond to the location of the point observation in the Grid's + Cartesian coordinates. + roi: float + The radius of influence of each point observation in m. + upper_bc: bool + True to enforce w=0 at top of domain (impermeability condition), + False to not enforce impermeability at top of domain + """ + def __init__(self): + self.Ut = np.nan + self.Vt = np.nan + self.rmsVr = np.nan + self.grid_shape = None + self.Cmod = np.nan + self.Cpoint = np.nan + self.u_back = None + self.v_back = None + self.wts = [] + self.vrs = [] + self.azs = [] + self.els = [] + self.weights = [] + self.bg_weights = [] + self.model_weights = [] + self.u_model = [] + self.v_model = [] + self.w_model = [] + self.x = None + self.y = None + self.z = None + self.dx = np.nan + self.dy = np.nan + self.dz = np.nan + self.Co = 1.0 + self.Cm = 1500.0 + self.Cx = 0.0 + self.Cy = 0.0 + self.Cz = 0.0 + self.Cb = 0.0 + self.Cv = 0.0 + self.Cmod = 0.0 + self.Cpoint = 0.0 + self.Ut = 0.0 + self.Vt = 0.0 + self.upper_bc = True + self.roi = 1000.0 + self.frz = 4500.0 + self.point_list = [] + + +def get_dd_wind_field(Grids, u_init, v_init, w_init, points=None, vel_name=None, refl_field=None, u_back=None, v_back=None, z_back=None, frz=4500.0, Co=1.0, Cm=1500.0, Cx=0.0, - Cy=0.0, Cz=0.0, Cb=0.0, Cv=0.0, Cmod=0.0, + Cy=0.0, Cz=0.0, Cb=0.0, Cv=0.0, Cmod=0.0, Cpoint=0.0, Ut=None, Vt=None, filt_iterations=2, mask_outside_opt=False, weights_obs=None, weights_model=None, weights_bg=None, max_iterations=200, mask_w_outside_opt=True, - filter_window=9, filter_order=4, min_bca=30.0, + filter_window=9, filter_order=3, min_bca=30.0, max_bca=150.0, upper_bc=True, model_fields=None, - output_cost_functions=True): + output_cost_functions=True, roi=1000.0): """ This function takes in a list of Py-ART Grid objects and derives a wind field. Every Py-ART Grid in Grids must have the same grid @@ -62,6 +190,9 @@ def get_dd_wind_field(Grids, u_init, v_init, w_init, vel_name=None, w_init: 3D ndarray The intial guess for the vertical wind field, input as a 3D array with the same shape as the fields in Grids. + points: None or list of dicts + Point observations as returned by :func:`pydda.constraints.get_iem_obs`. Set + to None to disable. vel_name: string Name of radial velocity field. Setting to None will have PyDDA attempt to automatically detect the velocity field name. @@ -93,6 +224,8 @@ def get_dd_wind_field(Grids, u_init, v_init, w_init, vel_name=None, Weight for cost function related to vertical vorticity equation. Cmod: float Weight for cost function related to custom constraints. + Cpoint: float + Weight for cost function related to point observations. weights_obs: list of floating point arrays or None List of weights for each point in grid from each radar in Grids. Set to None to let PyDDA determine this automatically. @@ -112,13 +245,13 @@ def get_dd_wind_field(Grids, u_init, v_init, w_init, vel_name=None, If this is greater than 0, PyDDA will run a low pass filter on the retrieved wind field and then do the optimization step for filt_iterations iterations. Set to 0 to disable the low pass filter. - max_outside_opt: bool + mask_outside_opt: bool If set to true, wind values outside the multiple doppler lobes will be masked, i.e. if less than 2 radars provide coverage for a given point. max_iterations: int The maximum number of iterations to run the optimization loop for. - max_w_outside_opt: bool + mask_w_outside_opt: bool If set to true, vertical winds outside the multiple doppler lobes will be masked, i.e. if less than 2 radars provide coverage for a given point. @@ -150,6 +283,9 @@ def get_dd_wind_field(Grids, u_init, v_init, w_init, vel_name=None, output_cost_functions: bool Set to True to output the value of each cost function every 10 iterations. + roi: float + Radius of influence for the point observations. The point observation will + not hold any weight outside this radius. Returns ======= @@ -158,8 +294,6 @@ def get_dd_wind_field(Grids, u_init, v_init, w_init, vel_name=None, are displayable by the visualization module. """ - num_evaluations = 0 - # We have to have a prescribed storm motion for vorticity constraint if(Ut is None or Vt is None): if(Cv != 0.0): @@ -169,6 +303,9 @@ def get_dd_wind_field(Grids, u_init, v_init, w_init, vel_name=None, if not isinstance(Grids, list): raise ValueError('Grids has to be a list!') + parameters = DDParameters() + parameters.Ut = Ut + parameters.Vt = Vt # Ensure that all Grids are on the same coordinate system prev_grid = Grids[0] @@ -192,20 +329,19 @@ def get_dd_wind_field(Grids, u_init, v_init, w_init, vel_name=None, # Disable background constraint if none provided if(u_back is None or v_back is None): - u_back2 = np.zeros(u_init.shape[0]) - v_back2 = np.zeros(v_init.shape[0]) - C8 = 0.0 + parameters.u_back = np.zeros(u_init.shape[0]) + parameters.v_back = np.zeros(v_init.shape[0]) else: # Interpolate sounding to radar grid print('Interpolating sounding to radar grid') u_interp = interp1d(z_back, u_back, bounds_error=False) v_interp = interp1d(z_back, v_back, bounds_error=False) - u_back2 = u_interp(Grids[0].z['data']) - v_back2 = v_interp(Grids[0].z['data']) + parameters.u_back = u_interp(Grids[0].z['data']) + parameters.v_back = v_interp(Grids[0].z['data']) print('Interpolated U field:') - print(u_back2) + print(parameters["u_back"]) print('Interpolated V field:') - print(v_back2) + print(parameters["v_back"]) print('Grid levels:') print(Grids[0].z['data']) @@ -217,22 +353,19 @@ def get_dd_wind_field(Grids, u_init, v_init, w_init, vel_name=None, if vel_name is None: vel_name = pyart.config.get_field_name('corrected_velocity') winds = np.stack([u_init, v_init, w_init]) - wts = [] - vrs = [] - azs = [] - els = [] + # Set up wind fields and weights from each radar - weights = np.zeros( + parameters.weights = np.zeros( (len(Grids), u_init.shape[0], u_init.shape[1], u_init.shape[2])) - bg_weights = np.zeros(v_init.shape) + parameters.bg_weights = np.zeros(v_init.shape) if(model_fields is not None): - mod_weights = np.ones( + parameters.model_weights = np.ones( (len(model_fields), u_init.shape[0], u_init.shape[1], u_init.shape[2])) else: - mod_weights = np.zeros( + parameters.model_weights = np.zeros( (1, u_init.shape[0], u_init.shape[1], u_init.shape[2])) if(model_fields is None): @@ -242,17 +375,16 @@ def get_dd_wind_field(Grids, u_init, v_init, w_init, vel_name=None, bca = np.zeros( (len(Grids), len(Grids), u_init.shape[1], u_init.shape[2])) - M = np.zeros(len(Grids)) sum_Vr = np.zeros(len(Grids)) for i in range(len(Grids)): - wts.append(cost_functions.calculate_fall_speed(Grids[i], - refl_field=refl_field)) + parameters.wts.append(cost_functions.calculate_fall_speed(Grids[i], + refl_field=refl_field, frz=frz)) add_azimuth_as_field(Grids[i], dz_name=refl_field) add_elevation_as_field(Grids[i], dz_name=refl_field) - vrs.append(Grids[i].fields[vel_name]['data']) - azs.append(Grids[i].fields['AZ']['data']*np.pi/180) - els.append(Grids[i].fields['EL']['data']*np.pi/180) + parameters.vrs.append(Grids[i].fields[vel_name]['data']) + parameters.azs.append(Grids[i].fields['AZ']['data']*np.pi/180) + parameters.els.append(Grids[i].fields['EL']['data']*np.pi/180) if(len(Grids) > 1): for i in range(len(Grids)): @@ -267,41 +399,41 @@ def get_dd_wind_field(Grids, u_init, v_init, w_init, vel_name=None, Grids[i].point_y['data'][0], Grids[i].get_projparams()) - for k in range(vrs[i].shape[0]): + for k in range(parameters.vrs[i].shape[0]): if(weights_obs is None): - cur_array = weights[i, k] + cur_array = parameters.weights[i, k] cur_array[np.logical_and( - ~vrs[i][k].mask, + ~parameters.vrs[i][k].mask, np.logical_and( bca[i, j] >= math.radians(min_bca), bca[i, j] <= math.radians(max_bca)))] += 1 - weights[i, k] = cur_array + parameters.weights[i, k] = cur_array else: - weights[i, k] = weights_obs[i][k, :, :] + parameters.weights[i, k] = weights_obs[i][k, :, :] if(weights_obs is None): - cur_array = weights[j, k] + cur_array = parameters.weights[j, k] cur_array[np.logical_and( - ~vrs[j][k].mask, + ~parameters.vrs[j][k].mask, np.logical_and( bca[i, j] >= math.radians(min_bca), bca[i, j] <= math.radians(max_bca)))] += 1 - weights[j, k] = cur_array + parameters.weights[j, k] = cur_array else: - weights[j, k] = weights_obs[j][k, :, :] + parameters.weights[j, k] = weights_obs[j][k, :, :] if(weights_bg is None): - cur_array = bg_weights[k] + cur_array = parameters.bg_weights[k] cur_array[np.logical_or( bca[i, j] >= math.radians(min_bca), bca[i, j] <= math.radians(max_bca))] = 1 - cur_array[vrs[i][k].mask] = 0 - bg_weights[i] = cur_array + cur_array[parameters.vrs[i][k].mask] = 0 + parameters.bg_weights[i] = cur_array else: - bg_weights[i] = weights_bg[i] + parameters.bg_weights[i] = weights_bg[i] print("Calculating weights for models...") - coverage_grade = weights.sum(axis=0) + coverage_grade = parameters.weights.sum(axis=0) coverage_grade = coverage_grade/coverage_grade.max() # Weigh in model input more when we have no coverage @@ -310,165 +442,142 @@ def get_dd_wind_field(Grids, u_init, v_init, w_init, vel_name=None, if(model_fields is not None): if(weights_model is None): for i in range(len(model_fields)): - mod_weights[i] = 1 - (coverage_grade/(len(Grids)+1)) + parameters.model_weights[i] = 1 - (coverage_grade/(len(Grids)+1)) else: for i in range(len(model_fields)): - mod_weights[i] = weights_model[i] + parameters.model_weights[i] = weights_model[i] else: if weights_obs is None: - weights[0] = np.where(~vrs[0].mask, 1, 0) + parameters.weights[0] = np.where(~parameters.vrs[0].mask, 1, 0) else: - weights[0] = weights_obs[0] + parameters.weights[0] = weights_obs[0] if weights_bg is None: - bg_weights = np.where(~vrs[0].mask, 0, 1) + parameters.bg_weights = np.where(~parameters.vrs[0].mask, 0, 1) else: - bg_weights = weights_bg + parameters.bg_weights = weights_bg - weights[weights > 0] = 1 - bg_weights[bg_weights > 0] = 1 - sum_Vr = np.nansum(np.square(vrs*weights)) - rmsVr = np.sqrt(np.nansum(sum_Vr)/np.nansum(weights)) + parameters.weights[parameters.weights > 0] = 1 + parameters.bg_weights[parameters.bg_weights > 0] = 1 + sum_Vr = np.nansum(np.square(parameters.vrs * parameters.weights)) + parameters.rmsVr = np.sqrt(np.nansum(sum_Vr) / np.nansum(parameters.weights)) del bca - grid_shape = u_init.shape + parameters.grid_shape = u_init.shape # Parse names of velocity field winds = winds.flatten() - ndims = len(winds) - - print(("Starting solver ")) - dx = np.diff(Grids[0].x['data'], axis=0)[0] - dy = np.diff(Grids[0].y['data'], axis=0)[0] - dz = np.diff(Grids[0].z['data'], axis=0)[0] - print('rmsVR = ' + str(rmsVr)) - print('Total points:' + str(weights.sum())) - z = Grids[0].point_z['data'] - the_time = time.time() + print("Starting solver ") + parameters.dx = np.diff(Grids[0].x['data'], axis=0)[0] + parameters.dy = np.diff(Grids[0].y['data'], axis=0)[0] + parameters.dz = np.diff(Grids[0].z['data'], axis=0)[0] + print('rmsVR = ' + str(parameters.rmsVr)) + print('Total points: %d' % parameters.weights.sum()) + parameters.z = Grids[0].point_z['data'] + parameters.x = Grids[0].point_x['data'] + parameters.y = Grids[0].point_y['data'] bt = time.time() # First pass - no filter - wcurr = w_init - wprev = 100*np.ones(w_init.shape) wprevmax = 99 wcurrmax = w_init.max() iterations = 0 - warnflag = 99999 - coeff_max = np.max([Co, Cb, Cm, Cx, Cy, Cz, Cb]) bounds = [(-x, x) for x in 100*np.ones(winds.shape)] - u_model = [] - v_model = [] - w_model = [] if(model_fields is not None): for the_field in model_fields: u_field = ("U_" + the_field) v_field = ("V_" + the_field) w_field = ("W_" + the_field) - u_model.append(Grids[0].fields[u_field]["data"]) - v_model.append(Grids[0].fields[v_field]["data"]) - w_model.append(Grids[0].fields[w_field]["data"]) + parameters.u_model.append(Grids[0].fields[u_field]["data"]) + parameters.v_model.append(Grids[0].fields[v_field]["data"]) + parameters.w_model.append(Grids[0].fields[w_field]["data"]) + + parameters.Co = Co + parameters.Cm = Cm + parameters.Cx = Cx + parameters.Cy = Cy + parameters.Cz = Cz + parameters.Cb = Cb + parameters.Cv = Cv + parameters.Cmod = Cmod + parameters.Cpoint = Cpoint + parameters.roi = roi + parameters.upper_bc = upper_bc + parameters.points = points + parameters.point_list = points while(iterations < max_iterations and (abs(wprevmax-wcurrmax) > 0.02)): wprevmax = wcurrmax - winds = fmin_l_bfgs_b(J_function, winds, args=(vrs, azs, els, - wts, u_back, v_back, - u_model, v_model, - w_model, - Co, Cm, Cx, Cy, Cz, Cb, - Cv, Cmod, Ut, Vt, - grid_shape, - dx, dy, dz, z, rmsVr, - weights, bg_weights, - mod_weights, - upper_bc, - False), + parameters.print_out = False + winds = fmin_l_bfgs_b(J_function, winds, args=(parameters,), maxiter=10, pgtol=1e-3, bounds=bounds, fprime=grad_J, disp=0, iprint=-1) - if(output_cost_functions is True): - J_function(winds[0], vrs, azs, els, wts, u_back, v_back, - u_model, v_model, w_model, - Co, Cm, Cx, Cy, Cz, Cb, Cv, Cmod, Ut, Vt, - grid_shape, dx, dy, dz, z, rmsVr, - weights, bg_weights, mod_weights, - upper_bc, True) - grad_J(winds[0], vrs, azs, els, wts, u_back, v_back, - u_model, v_model, w_model, - Co, Cm, Cx, Cy, Cz, Cb, Cv, Cmod, Ut, Vt, - grid_shape, dx, dy, dz, z, rmsVr, - weights, bg_weights, mod_weights, - upper_bc, True) - warnflag = winds[2]['warnflag'] - winds = np.reshape(winds[0], (3, grid_shape[0], grid_shape[1], - grid_shape[2])) + parameters.print_out = True + if output_cost_functions is True: + J_function(winds[0], parameters) + grad_J(winds[0], parameters) + winds = np.reshape( + winds[0], (3, parameters.grid_shape[0], parameters.grid_shape[1], parameters.grid_shape[2])) iterations = iterations+10 print('Iterations before filter: ' + str(iterations)) wcurrmax = winds[2].max() winds = np.stack([winds[0], winds[1], winds[2]]) winds = winds.flatten() - if(filt_iterations > 0): + if filt_iterations > 0: print('Applying low pass filter to wind field...') - winds = np.reshape(winds, (3, grid_shape[0], grid_shape[1], - grid_shape[2])) - winds[0] = savgol_filter(winds[0], 9, 3, axis=0) - winds[0] = savgol_filter(winds[0], 9, 3, axis=1) - winds[0] = savgol_filter(winds[0], 9, 3, axis=2) - winds[1] = savgol_filter(winds[1], 9, 3, axis=0) - winds[1] = savgol_filter(winds[1], 9, 3, axis=1) - winds[1] = savgol_filter(winds[1], 9, 3, axis=2) - winds[2] = savgol_filter(winds[2], 9, 3, axis=0) - winds[2] = savgol_filter(winds[2], 9, 3, axis=1) - winds[2] = savgol_filter(winds[2], 9, 3, axis=2) + winds = np.reshape(winds, (3, parameters.grid_shape[0], parameters.grid_shape[1], + parameters.grid_shape[2])) + winds[0] = savgol_filter(winds[0], filter_window, filter_order, axis=0) + winds[0] = savgol_filter(winds[0], filter_window, filter_order, axis=1) + winds[0] = savgol_filter(winds[0], filter_window, filter_order, axis=2) + winds[1] = savgol_filter(winds[1], filter_window, filter_order, axis=0) + winds[1] = savgol_filter(winds[1], filter_window, filter_order, axis=1) + winds[1] = savgol_filter(winds[1], filter_window, filter_order, axis=2) + winds[2] = savgol_filter(winds[2], filter_window, filter_order, axis=0) + winds[2] = savgol_filter(winds[2], filter_window, filter_order, axis=1) + winds[2] = savgol_filter(winds[2], filter_window, filter_order, axis=2) winds = np.stack([winds[0], winds[1], winds[2]]) winds = winds.flatten() iterations = 0 while(iterations < filt_iterations): winds = fmin_l_bfgs_b( - J_function, winds, args=(vrs, azs, els, - wts, u_back, v_back, - u_model, v_model, w_model, - Co, Cm, Cx, Cy, Cz, Cb, - Cv, Cmod, Ut, Vt, - grid_shape, - dx, dy, dz, z, rmsVr, - weights, bg_weights, - mod_weights, - upper_bc, - False), + J_function, winds, args=(parameters,), maxiter=10, pgtol=1e-3, bounds=bounds, fprime=grad_J, disp=0, iprint=-1) - - warnflag = winds[2]['warnflag'] - winds = np.reshape(winds[0], (3, grid_shape[0], grid_shape[1], - grid_shape[2])) + parameters.print_out = False + winds = np.reshape( + winds[0], (3, parameters.grid_shape[0], parameters.grid_shape[1], parameters.grid_shape[2])) iterations = iterations+1 print('Iterations after filter: ' + str(iterations)) winds = np.stack([winds[0], winds[1], winds[2]]) winds = winds.flatten() + print("Done! Time = " + "{:2.1f}".format(time.time() - bt)) # First pass - no filter - the_winds = np.reshape(winds, (3, grid_shape[0], grid_shape[1], - grid_shape[2])) + the_winds = np.reshape( + winds, (3, parameters.grid_shape[0], parameters.grid_shape[1], parameters.grid_shape[2])) u = the_winds[0] v = the_winds[1] w = the_winds[2] - where_mask = np.sum(weights, axis=0) + np.sum(mod_weights, axis=0) + where_mask = np.sum(parameters.weights, axis=0) + \ + np.sum(parameters.model_weights, axis=0) u = np.ma.array(u) w = np.ma.array(w) v = np.ma.array(v) - if(mask_outside_opt is True): + if mask_outside_opt is True: u = np.ma.masked_where(where_mask < 1, u) v = np.ma.masked_where(where_mask < 1, v) w = np.ma.masked_where(where_mask < 1, w) - if(mask_w_outside_opt is True): + if mask_w_outside_opt is True: w = np.ma.masked_where(where_mask < 1, w) u_field = deepcopy(Grids[0].fields[vel_name]) diff --git a/pydda/tests/test_cost_functions.py b/pydda/tests/test_cost_functions.py index fb1e52df..cc8440bb 100644 --- a/pydda/tests/test_cost_functions.py +++ b/pydda/tests/test_cost_functions.py @@ -158,6 +158,44 @@ def test_vert_vorticity(): assert cost > 0 +def test_point_cost(): + u = 1 * np.ones((10, 10, 10)) + v = 1 * np.ones((10, 10, 10)) + w = 0 * np.ones((10, 10, 10)) + + x = np.linspace(-10, 10, 10) + y = np.linspace(-10, 10, 10) + z = np.linspace(-10, 10, 10) + x, y, z = np.meshgrid(x, y, z) + + my_point1 = {'x': 0, 'y': 0, 'z': 0, 'u': 2., 'v': 2., 'w': 0.} + cost = pydda.cost_functions.calculate_point_cost(u, v, x, y, z, [my_point1], roi=2.0) + grad = pydda.cost_functions.calculate_point_gradient(u, v, x, y, z, [my_point1], roi=2.0) + + assert cost > 0 + assert np.all(grad <= 0) + + my_point1 = {'x': 0, 'y': 0, 'z': 0, 'u': -2., 'v': -2., 'w': 0.} + my_point2 = {'x': 3, 'y': 3, 'z': 0, 'u': 2., 'v': 2., 'w': 0.} + + cost = pydda.cost_functions.calculate_point_cost(u, v, x, y, z, [my_point1], roi=2.0) + grad = pydda.cost_functions.calculate_point_gradient(u, v, x, y, z, [my_point1], roi=2.0) + assert cost > 0 + assert np.all(grad >= 0) + + cost = pydda.cost_functions.calculate_point_cost(u, v, x, y, z, [my_point1, my_point2], roi=2.0) + grad = pydda.cost_functions.calculate_point_gradient(u, v, x, y, z, [my_point1, my_point2], roi=2.0) + assert cost > 0 + assert ~np.all(grad >= 0) + + my_point1 = {'x': 0, 'y': 0, 'z': 0, 'u': 1., 'v': 1., 'w': 0.} + my_point2 = {'x': 3, 'y': 3, 'z': 0, 'u': 1., 'v': 1., 'w': 0.} + cost = pydda.cost_functions.calculate_point_cost(u, v, x, y, z, [my_point1, my_point2], roi=2.0) + grad = pydda.cost_functions.calculate_point_gradient(u, v, x, y, z, [my_point1, my_point2], roi=2.0) + assert cost == 0 + assert np.all(grad == 0) + + def test_model_cost(): u = 10*np.ones((10, 10, 10)) v = 10*np.ones((10, 10, 10)) diff --git a/pydda/tests/test_initialization.py b/pydda/tests/test_initialization.py index b10ffa5a..bcf79ffe 100644 --- a/pydda/tests/test_initialization.py +++ b/pydda/tests/test_initialization.py @@ -52,3 +52,12 @@ def test_make_wind_field_from_profile(): assert np.all(np.round(u) == 1) assert np.all(np.round(v) == 1) assert np.all(w == 0.0) + + +def test_get_iem_data(): + Grid = pyart.testing.make_empty_grid( + (20, 20, 20), ((0, 100000.), (-100000., 100000.), (-100000., 100000.))) + station_obs = pydda.constraints.get_iem_obs(Grid) + names = [x['site_id'] for x in station_obs] + assert names == ['P28', 'WLD', 'WDG', 'SWO', 'END'] +