Skip to content

Commit

Permalink
use p_ref for theta calc instead of pmax.
Browse files Browse the repository at this point in the history
  • Loading branch information
@AaronDavidSchneider
AaronDavidSchneider committed Nov 17, 2023
1 parent c785de1 commit efe2715
Show file tree
Hide file tree
Showing 2 changed files with 66 additions and 41 deletions.
2 changes: 1 addition & 1 deletion .flake8
View file
Original file line number Diff line number Diff line change
@@ -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
105 changes: 65 additions & 40 deletions gcm_toolkit/utils/manipulations.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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"]]
)

Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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):
Expand All @@ -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)

Expand All @@ -549,24 +570,28 @@ 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
else:
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

0 comments on commit efe2715

Please sign in to comment.