From efe2715221b6650b7464d5ea2804cd9c11294770 Mon Sep 17 00:00:00 2001 From: Aaron David Schneider Date: Fri, 17 Nov 2023 10:21:14 +0100 Subject: [PATCH] use p_ref for theta calc instead of pmax. --- .flake8 | 2 +- gcm_toolkit/utils/manipulations.py | 105 ++++++++++++++++++----------- 2 files changed, 66 insertions(+), 41 deletions(-) diff --git a/.flake8 b/.flake8 index bc8713c..c97e15c 100644 --- a/.flake8 +++ b/.flake8 @@ -1,5 +1,5 @@ [flake8] -ignore = R0913, R0914, C0415, E0402, F401, E203, F811, W503 +ignore = R0913, R0914, C0415, E0402, F401, E203, F811, W503, C901 max-line-length = 88 max-complexity = 18 select = B,C,E,F,W,T4,B9 diff --git a/gcm_toolkit/utils/manipulations.py b/gcm_toolkit/utils/manipulations.py index cb529a0..e5159ec 100644 --- a/gcm_toolkit/utils/manipulations.py +++ b/gcm_toolkit/utils/manipulations.py @@ -304,7 +304,7 @@ def m_add_theta(dsi, var_key_out=None, temp_key="T"): theta : xarray.DataArray A dataArray with reduced dimensionality, containing the potential temperature """ - theta = dsi[c[temp_key]] * (dsi[c["Z"]].max() / dsi[c["Z"]]) ** ( + theta = dsi[c[temp_key]] * (dsi.attrs[c["p_ref"]] / dsi[c["Z"]]) ** ( dsi.attrs[c["R"]] / dsi.attrs[c["cp"]] ) @@ -437,8 +437,10 @@ def m_add_meridional_overturning(dsi, v_data="V", var_key_out=None): return psi -def m_extend_upward(temp_array, p_low, method=None, n_p=20, T_therm=None, - p_therm_high=None): + +def m_extend_upward( + temp_array, p_low, method=None, n_p=20, T_therm=None, p_therm_high=None +): """ Extend the given temperature DataArray upward to the given pressure value. Multiple methods can be used to extend the data upwards: @@ -474,58 +476,74 @@ def m_extend_upward(temp_array, p_low, method=None, n_p=20, T_therm=None, The extended temperature array. """ # Check if the input pressure is suitable - p_low_current = min(temp_array.Z.values) # lowest pressure in current data - if hasattr(p_low, "__len__"): # p_low is a list + p_low_current = min(temp_array.Z.values) # lowest pressure in current data + if hasattr(p_low, "__len__"): # p_low is a list p_low = [ip for ip in p_low if ip < p_low_current] p_low = sorted(p_low)[::-1] - elif isinstance(p_low, (int, float)): # p_low is a scalar - p_low = np.logspace(np.log10(p_low_current), np.log10(p_low), n_p+1) + elif isinstance(p_low, (int, float)): # p_low is a scalar + p_low = np.logspace(np.log10(p_low_current), np.log10(p_low), n_p + 1) p_low = p_low[1:] - else: # p_low is something else - wrt.write_status('ERROR', 'p_low should be a float, int, list, or array.') + else: # p_low is something else + wrt.write_status( + "ERROR", "p_low should be a float, int, list, or array." + ) # Now p_low consists of a sorted list of low pressures # Construct an extended dataArray (initially filled with zeros) - new_coords = dict(time= temp_array.time, Z=p_low, lat=temp_array.lat, - lon=temp_array.lon) + new_coords = dict( + time=temp_array.time, Z=p_low, lat=temp_array.lat, lon=temp_array.lon + ) # NB: in order to avoid problems with merging the extended and the original # temperature array, we need to copy over all coordinates. We just assign # them NaN values since we aren't interested in their values. - for c in temp_array.coords: - if c not in new_coords: - new_coords[c] = np.nan - nz = len(new_coords['Z']) - nlat = len(new_coords['lat']) - nlon = len(new_coords['lon']) - if 'time' in temp_array.dims: - nt = len(new_coords['time']) - time_coord = True # flag for an explicit time coordinate --> 4D array + for coord in temp_array.coords: + if coord not in new_coords: + new_coords[coord] = np.nan + + nz = len(new_coords["Z"]) + nlat = len(new_coords["lat"]) + nlon = len(new_coords["lon"]) + if "time" in temp_array.dims: + nt = len(new_coords["time"]) + time_coord = True # flag for an explicit time coordinate --> 4D array extended_values = np.zeros((nt, nz, nlat, nlon)) else: nt = 1 - time_coord = False # no explicit time coordinate --> 3D array + time_coord = False # no explicit time coordinate --> 3D array extended_values = np.zeros((nz, nlat, nlon)) # Fill the extended data according to the chosen method # ISOTHERMAL - if method == 'isothermal': - iso_slice = temp_array.isel(Z=-1).values # constant upper temperature + if method == "isothermal": + iso_slice = temp_array.isel(Z=-1).values # constant upper temperature for i in range(0, nz): if time_coord: - extended_values[:,i,:,:] = iso_slice # copy for each new pressure + extended_values[ + :, i, :, : + ] = iso_slice # copy for each new pressure else: - extended_values[i,:,:] = iso_slice # copy for each new pressure + extended_values[ + i, :, : + ] = iso_slice # copy for each new pressure # THERMOSPHERE - elif method == 'thermosphere': + elif method == "thermosphere": import math + # Check if all parameters are given if T_therm is None: - wrt.write_status('ERROR', 'Argument "T_therm" is required for the thermosphere method.') + wrt.write_status( + "ERROR", + 'Argument "T_therm" is required for the thermosphere method.', + ) if p_therm_high is None: - wrt.write_status('ERROR', 'Argument "p_therm_high" is required for the thermosphere method.') - p_therm_low = min(p_low) # minimum pressure for thermosphere extention + wrt.write_status( + "ERROR", + 'Argument "p_therm_high" is required for the thermosphere' + " method.", + ) + p_therm_low = min(p_low) # minimum pressure for thermosphere extention # For each time step... for it in range(0, nt): @@ -534,12 +552,15 @@ def m_extend_upward(temp_array, p_low, method=None, n_p=20, T_therm=None, for ilat in range(0, nlat): # ... save the top temperature as minimum if time_coord: - T_min = temp_array.isel(time=it, Z=-1, lat=ilat, lon=ilon) + T_min = temp_array.isel( + time=it, Z=-1, lat=ilat, lon=ilon + ) else: T_min = temp_array.isel(Z=-1, lat=ilat, lon=ilon) # ... calculate angle of incidence - mu = math.cos(math.radians(T_min.lon)) * \ - abs( math.cos(math.radians(T_min.lat)) ) + mu = math.cos(math.radians(T_min.lon)) * abs( + math.cos(math.radians(T_min.lat)) + ) # ... determine the 'thermospheric' upper temperature. T_max = max(T_min, T_therm * mu) @@ -549,8 +570,11 @@ def m_extend_upward(temp_array, p_low, method=None, n_p=20, T_therm=None, # ... if the pressure is below the thermosphere... if ip < p_therm_high and mu >= 0: # ... do the interpolation to T_max - T_new = T_min + (T_max-T_min)/(math.log10(p_therm_high)-math.log10(p_therm_low)) * (math.log10(p_therm_high)-math.log10(ip)) - else: # ... otherwise don't change anything. + T_new = T_min + (T_max - T_min) / ( + math.log10(p_therm_high) + - math.log10(p_therm_low) + ) * (math.log10(p_therm_high) - math.log10(ip)) + else: # ... otherwise don't change anything. T_new = T_min if time_coord: extended_values[it, i, ilat, ilon] = T_new @@ -558,15 +582,16 @@ def m_extend_upward(temp_array, p_low, method=None, n_p=20, T_therm=None, extended_values[i, ilat, ilon] = T_new else: msg = 'Specify either "isothermal" or "thermosphere" as method.' - wrt.write_status('ERROR', msg) + wrt.write_status("ERROR", msg) # Make the new extended data into a DataArray - extension = xr.DataArray(extended_values, - coords=new_coords, - dims=temp_array.dims) + extension = xr.DataArray( + extended_values, coords=new_coords, dims=temp_array.dims + ) # Combine the extended and original data - temp_array = xr.concat([temp_array, extension], dim='Z', - compat='no_conflicts') + temp_array = xr.concat( + [temp_array, extension], dim="Z", compat="no_conflicts" + ) return temp_array