import numpy as np
import exoplasim.filesupport
from exoplasim.filesupport import SUPPORTED
import os, glob
class _Constants:
def __init__(self):
'''Container for constants in SI units
Constants contained:
--------------------
h : Planck's constant
c : Speed of light
kb : Boltzmann constant
sb : Stefan-Boltzmann constant
R : ideal gas constant
pc : 1 parsec
ly : 1 light-year
AU : 1 astronomical unit
Rsun : solar radius
Rearth : Earth radius (by default the equatorial radius)
Rjup : Jupiter radius (by default the equatorial radius)
Rearth_* : Earth radii of various types. Polar (_po), equatorial (_eq),
arithmetic mean (_av), authalic (_auth), volumetric (_vol),
and rectifying (_rec) radii are all available.
Rjup_* : Jupiter radii of various types. Polar (_po), equatorial (_eq),
and arithmetic mean (_av) are available.
Msun : Solar mass
Mearth: Earth mass
Mjup: Jupiter mass
'''
self.h = 6.62607004e-34 #Planck's constant
self.c = 2.99792458e8 #Speed of light
self.kb = 1.38064852e-23 #Boltzmann constant
self.sb = 5.670374419e-8 #Stefan-Boltzmann constant
self.Rearth_eq = 6.3781e6 #Earth equatorial radius
self.Rsun = 6.957e8 #Radius of the sun (IAU)
self.Rearth_po = 6.3568e6 #Earth polar radius
self.Rearth_av = 6.3710088e6 #Earth arithmetic mean radius
self.Rearth_auth = 6.3710072e6 #Earth authalic radius (same area as reference ellipsoid)
self.Rearth_vol = 6.3710008e6 #Earth volumetric radius (same volume as reference ellipsoid)
self.Rearth_rec = 6.367445e6 #Earth rectifying radius (same polar cross section perimeter)
self.Rjup_eq = 7.1492e7 #Jupiter equatorial radius
self.Rjup_po = 6.6854e7 #Jupiter polar radius
self.Rjup_av = 6.9911e7 #Jupiter arithmetic mean radius
self.Rearth = self.Rearth_eq
self.Rjup = self.Rjup_eq
self.Msun = 1.98847e30 #Solar mass in kg
self.Mearth = 5.9722e24 #Earth mass in kg
self.Mjup = 1.8982e27 #Jupiter mass in kg
self.AU = 1.49597870691e11 #1 AU in meters
self.ly = 9.460730472580800e15 #1 light-year in meters
self.pc = 6.48000e5/np.pi*self.AU #1 parsec in meters
self.R = 8.31446261815324 #ideal gas constant
def __str__(self):
docstring = ("h (Planck's constant) : 6.62607004e-34 \n"+
"c (speed of light) : 2.99792458e8 \n"+
"kb (Boltzmann constant) : 1.38064852e-23 \n"+
"sb (Stefan-Boltzmann constant) : 5.670374419e-8 \n"+
"Rearth_eq (Earth equatorial radius) : 6.3781e6 \n"+
"Rsun (Solar radius) : 6.957e8 \n"+
"Rearth_po (Earth polar radius) : 6.3568e6 \n"+
"Rearth_av (Earth average radius) : 6.3710088e6 \n"+
"Rearth_auth (Earth authalic radius) : 6.3710072e6 \n"+
"Rearth_vol (Earth volumetric radius) : 6.3710008e6 \n"+
"Rearth_rec (Earth rectifying radius) : 6.367445e6 \n"+
"Rjup_eq (Jupiter equatorial radius) : 7.1492e7 \n"+
"Rjup_po (Jupiter polar radius) : 6.6854e7 \n"+
"Rjup_av (Jupiter average radius) : 6.9911e7 \n"+
"Rearth (Earth IAU radius) : Rearth_eq \n"+
"Rjup (Jupiter IAU radius) : Rjup_eq \n"+
"Msun (Solar mass) : 1.98847e30 \n"+
"Mearth (Earth mass) : 5.9722e24 \n"+
"Mjup (Jupiter mass) : 1.8982e27 \n"+
"AU (Astronomical Unit) : 1.49597870691e11 \n"+
"ly (light year) : 9.460730472580800e15\n"+
"pc (parsec) : 6.48000e5/pi*AU \n"+
"R (ideal gas constant) : 8.31446261815324")
return docstring
consts = _Constants()
[docs]def blackbody(wavelengths,temperature):
'''Compute the Planck function for a set of wavelengths and a given effective temperature.
Parameters
----------
wavelengths : array-like
Wavelengths in microns
temperature : float
Effective temperature in Kelvins
Returns
-------
array-like
Spectral radiance F_lambda(lambda,T) for the provided wavelengths assuming a perfect blackbody.
'''
y = wavelengths*1.0e-6
T = temperature
bb = 2*consts.h*consts.c**2/y**5 * 1/(np.exp(consts.h*consts.c/(y*consts.kb*T))-1) #W sr-1 m-3
bb *= 1e6 #W sr-1 m-2 um-1
return bb
def _loadnetcdf(filename):
import netCDF4 as nc
ncd = nc.Dataset(filename,"r")
return ncd,ncd.variables
def _loadnpsavez(filename):
npdata = np.load(filename)
return npdata
def _loadcsv(filename,buffersize=1):
import gzip
fileparts = filename.split('.')
rdataset = {}
metadata = {}
openfiles = []
if "tar" in fileparts[-2:]: #We're dealing with a tarball of some kind
import tarfile
with tarfile.open(filename,"r") as tarball:
members = tarball.getnames()
tarball.extractall()
else: #filename here should be the name of a directory containing only variable csv/txt/gz files
#Just a collection of CSV/TXT-type files in a subdirectory, which may be individually-compressed.
#These files can have .txt, .csv, or .gz file extensions.
members = glob.glob(filename+"/*")
dimensions = {}
for var in members:
#with open(var,"r") as csvf:
#data = np.loadtxt(csvf,delimiter=',')
if var[-3:]==".gz": #gzipped file
with gzip.open(var,"rt") as gzf:
header = gzf.readline()[1:].split(',')
else:
with open(var,"r") as txtf:
header = txtf.readline()[1:].split(',')
dims = []
k=0
while header[k]!="|||":
dims.append(int(header[k]))
k+=1
k+=1
meta = header[k:]
dims = tuple(dims)
dimensions[meta[0]] = dims
#rdataset[meta[0]] = np.reshape(data,dims)
metadata[meta[0]] = {}
try:
metadata[meta[0]]["standard_name"] = meta[1]
except:
metadata[meta[0]]["standard_name"] = meta[0]
try:
metadata[meta[0]]["long_name"] = meta[1]
except:
metadata[meta[0]]["long_name"] = meta[0]
try:
metadata[meta[0]]["units"] = meta[2]
except:
metadata[meta[0]]["units"] = "N/A"
try:
metadata[meta[0]]["code"] = int(meta[3])
except:
metadata[meta[0]]["code"] = -999
openfiles.append(var)
if "tar" in fileparts[-2:]:
for f in openfiles:
os.system("rm -rf %s"%f)
os.system("rm -rf %s"%(openfiles[0].split("/")[0]))
rdataset = _csvData(filename,shapes=dimensions,buffersize=buffersize)
return rdataset,metadata
def _loadhdf5(filename):
import h5py
hdfile = h5py.File(filename,"r")
return hdfile
class _csvData(dict):
'''An iterable dict-like object that supports all native dict methods, but accesses and manages a file archive or directory instead of a dictionary.
Parameters
----------
archive : str
Either a path to a tarball archive or a directory stem with file extension, e.g. MOST.002.csv
shapes : dict, optional
Dictionary of tuples giving the shapes of each variable in the archive
buffersize : int, optional
Number of data arrays to store in memory at a time
**kwargs : optional
Any additional keyword arguments to pass to the parent `dict` object. These will be accessible
via the usual dictionary methods.
Returns
-------
iterable
Supports all dictionary methods
'''
def __init__(self,archive,shapes={},buffersize=1,**kwargs):
self.archive = archive
self.tarball=False
self.buffersize=buffersize
self.dbuffer = {}
self.dbufferkeys = []
self.shapes = shapes
if "tar" in archive:
self.tarball=True
import tarfile
with tarfile.open(self.archive,"r") as tarball:
self.files = tarball.getnames()
else:
self.files = glob.glob(".".join(archive.split(".")[:-1])+"/*")
self.variables = []
self.filetree = {}
for f in self.files:
variable = f.split("_")[-1].split(".")[0]
self.filetree[variable] = f
self.variables.append(variable)
idx0 = list(self.filetree.keys())[0]
self.filestem = "_".join(self.filetree[idx0].split("_")[:-1])
self.extension = "."+self.filetree[idx0].split(".")[-1]
self.permanent = {}
dimkeys = ['lat','lon','lev','levp','time',"wvl"]
for key in dimkeys:
try:
self.permanent[key] = self.__getitem__(key,overridebuffer=True)
except:
pass
super(_csvData,self).__init__(**kwargs)
for key in self.filetree:
super(_csvData,self).__setitem__(key,self.filetree[key])
def __getitem__(self,key,overridebuffer=False):
'''Retrieve variable from archive
Parameters
----------
key : str
Variable name
overridebuffer : bool, optional
If True, don't store this variable in the buffer (useful for lat, lon, etc)
Returns
-------
numpy.ndarray
Variable data
'''
if key not in self.filetree: #This must have been passed via **kwargs
return super(_csvData,self).__getitem__(key)
if key in self.permanent:
return self.permanent[key]
if key not in self.dbuffer:
if self.tarball:
import tarfile
with tarfile.open(self.archive,"r") as tarball:
tarball.extract(self.filetree[key])
with open(self.filetree[key],"r") as csvf:
data = np.loadtxt(csvf,delimiter=',')
if key in self.shapes:
data = np.reshape(data,self.shapes[key])
os.system("rm -rf %s"%self.filetree[key])
os.system("rm -rf %s"%(self.filetree[key].split("/")[0]))
if not overridebuffer:
if len(self.dbufferkeys)==self.buffersize:
del self.dbuffer[self.dbufferkeys[0]]
self.dbufferkeys.remove(self.dbufferkeys[0])
self.dbuffer[key] = data
self.dbufferkeys.append(key)
else:
data = self.dbuffer[key]
if not overridebuffer:
self.dbufferkeys.remove(key)
self.dbufferkeys.append(key) #Move this key to the end so it's the last to be removed.
return data
def __setitem__(self,key,value):
'''Add array to archive
Unless the archive is an uncompressed tarfile that doesn't already have the variable `key`,
this will extract the archive, delete the original, write a new file, and re-tar
Parameters
----------
key : str
Variable name
value : numpy.ndarray
Variable data
'''
if key in self.filetree:
if self.tarball:
import tarfile
with tarfile.open(self.archive,"r") as tarball:
members = tarball.getnames()
tarball.extractall()
os.system("rm -rf "+self.archive)
print("Writing %8s to %s"%(key,fname))
np.savetxt(fname,value.astype("float32"),
header=(','.join(np.array(value.shape).astype(str))+',|||,'
+','.join([key,key,"user","-333"])),delimiter=',')
if self.tarball:
if "tar.gz" in self.archive or "tar.bz2" in self.archive or "tar.xz" in self.archive:
with tarfile.open(self.archive,"w:%s"%self.archive.split(".")[-1]) as tarball:
for var in members:
varname = var.split("_")[-1].split(".")[0]
tarball.add(var,arcname=varname)
else:
with tarfile.open(self.archive,"w") as tarball:
for var in members:
varname = var.split("_")[-1].split(".")[0]
tarball.add(var,arcname=varname)
for var in members:
os.system("rm -rf %s"%var)
else:
fname = lf.filestem+"_"+key+extension
self.variables.append(key)
self.filetree[key] = fname
print("Writing %8s to %s"%(key,fname))
np.savetxt(fname,value.astype("float32"),
header=(','.join(np.array(value.shape).astype(str))+',|||,'
+','.join([key,key,"user","-333"])),delimiter=',')
if self.tarball:
import tarfile
print("Packing %s in %s"%(fname,self.archive))
if "tar.gz" in self.archive or "tar.bz2" in self.archive or "tar.xz" in self.archive:
with tarfile.open(self.archive,"r") as tarball:
members = tarball.getnames()
tarball.extractall()
os.system("rm -rf "+self.archive)
with tarfile.open(self.archive,"w:%s"%self.archive.split(".")[-1]) as tarball:
for var in members:
varname = var.split("_")[-1].split(".")[0]
tarball.add(var,arcname=varname)
tarball.add(fname,arcname=key)
for var in members:
os.system("rm -rf %s"%var)
os.system("rm -rf %s"%fname)
os.system("rf -rf %s"%(members[0].split("/")[0]))
else:
with tarfile.open(self.archive,"a") as tarball:
tarball.add(fname,arcname=key)
os.system("rm -rf %s"%fname)
super(_csvData,self).__setitem__(key,fname)
class _Dataset:
def __init__(self,filename,csvbuffersize=1):
self.body=None
fileparts = filename.split('.')
if fileparts[-1] == "nc":
self.body,self.variables = _loadnetcdf(filename)
self.metadata = {}
for var in self.variables:
self.metadata[var] = {}
try:
self.metadata[var]["standard_name"]= self.variables[var].standard_name
except:
self.metadata[var]["standard_name"]= var
try:
self.metadata[var]["long_name"]= self.variables[var].long_name
except:
self.metadata[var]["long_name"] = var
try:
self.metadata[var]["units"] = self.variables[var].units
except:
self.metadata[var]["units"] = "N/A"
try:
self.metadata[var]["code"] = int(self.variables[var].code)
except:
self.metadata[var]["code"] = -999
elif fileparts[-1] == "npz" or fileparts[-1] == "npy":
self.variables=_loadnpsavez(filename)
meta = _loadnpsavez(filename[:-4]+"_metadata.npz")
self.metadata = {}
for var in self.variables:
self.metadata[var] = {}
try:
self.metadata[var]["standard_name"]= meta[var][1]
except:
self.metadata[var]["standard_name"]= var
try:
self.metadata[var]["long_name"]= meta[var][1]
except:
self.metadata[var]["long_name"] = var
try:
self.metadata[var]["units"] = meta[var][2]
except:
self.metadata[var]["units"] = "N/A"
try:
self.metadata[var]["code"] = int(meta[var][3])
except:
self.metadata[var]["code"] = -999
elif (fileparts[-1]=="tar" or \
fileparts[-2]+"."+fileparts[-1] in ("tar.gz","tar.bz2","tar.xz")):
self.variables,self.metadata =_loadcsv(filename,buffersize=csvbuffersize)
self.tarball = True
elif (fileparts[-1] in ("csv","txt","gz")):
self.variables,self.metadata =_loadcsv(filename,buffersize=csvbuffersize)
self.tarball = False
elif fileparts[-1] in ("hdf5","h5","he5"):
self.variables=_loadhdf5(filename)
self.metadata={}
for var in self.variables:
meta = self.variables.attrs[var]
self.metadata[var] = {}
try:
self.metadata[var]["standard_name"]= meta[1]
except:
self.metadata[var]["standard_name"]= var
try:
self.metadata[var]["long_name"]= meta[1]
except:
self.metadata[var]["long_name"] = var
try:
self.metadata[var]["units"] = meta[2]
except:
self.metadata[var]["units"] = "N/A"
try:
self.metadata[var]["code"] = int(meta[3])
except:
self.metadata[var]["code"] = -999
else:
raise DatafileError("Unsupported output format detected. Supported formats are:\n%s"%("\n\t".join(SUPPORTED)))
def close(self):
try:
if self.body is not None:
if type(self.body)==list:
for item in self.body:
if self.tarball:
os.system("rm %s"%item)
else:
pass
else:
self.body.close()
else:
self.variables.close()
except: #We have a standard python dictionary for variables, so ignore the error quietly
pass
[docs]def xcolorbar(mappable,fontsize=None,ticksize=None,**kwargs):
import matplotlib.pyplot as plt
cbar = plt.colorbar(mappable,**kwargs)
if fontsize is not None and "label" in kwargs:
cbar.set_label(label=kwargs["label"],fontsize=fontsize)
if ticksize is not None:
cbar.ax.tick_params(labelsize=ticksize)
return cbar
#import mpl_toolkits.basemap
#class _xBasemap(mpl_toolkits.basemap.Basemap):
#def __init__(self,**kwargs):
#super(_xBasemap,self).__init__(**kwargs)
#def xcolorbar(self,mappable,fontsize=None,ticksize=None,**kwargs):
#cbar = super(_xBasemap,self).colorbar(mappable,**kwargs)
#if fontsize is not None and "label" in kwargs:
#cbar.set_label(label=kwargs["label"],fontsize=fontsize)
#if ticksize is not None:
#cbar.ax.tick_params(labelsize=ticksize)
#return cbar
[docs]class DimensionError(Exception):
pass
[docs]class UnitError(Exception):
pass
[docs]class DatafileError(Exception):
pass
[docs]def parse(file,variable,lat=None,lon=None):
"""Retrieve a variable from a NetCDF file
Parameters
----------
file : str
Path to a NetCDF file
variable : str
Name of the variable to extract
lat,lon : str, optional
If the latitude and longitude arrays have non-standard names, specify them here.
Returns
-------
numpy.ndarray
Requested output field
"""
ncd=_Dataset(file)
variable = ncd.variables[variable][:]
if lat:
lt = ncd.variables[lat][:]
elif "lat" in ncd.variables:
lt = ncd.variables['lat'][:]
elif "lt" in ncd.variables:
lt = ncd.variables["lt"][:]
elif "lats" in ncd.variables:
lt = ncd.variables['lats'][:]
elif "latitude" in ncd.variables:
lt = ncd.variables['latitude'][:]
elif "latitudes" in ncd.variables:
lt = ncd.variables['latitudes'][:]
else:
raise DatafileError("Unknown datafile format; unsure how to extract latitude")
if lon:
ln = ncd.variables[lon][:]
elif "lat" in ncd.variables:
ln = ncd.variables['lon'][:]
elif "lt" in ncd.variables:
ln = ncd.variables["ln"][:]
elif "lats" in ncd.variables:
ln = ncd.variables['lons'][:]
elif "latitude" in ncd.variables:
ln = ncd.variables['longitude'][:]
elif "latitudes" in ncd.variables:
ln = ncd.variables['longitudes'][:]
else:
raise DatafileError("Unknown datafile format; unsure how to extract longitude")
ncd.close()
return ln,lt,variable
[docs]def make2d(variable,lat=None,lon=None,time=None,lev=None,ignoreNaNs=True,radius=6.371e6,
latitudes=None,longitudes=None):
"""Compress a variable in two dimensions by slicing or averaging.
Parameters
----------
variable : numpy.ndarray
The variable to operate on
lat,lon,lev : int, str, optional
Either an index on which to slice, or either of "sum" or "mean", indicating what
should be done along that axis.
time : int, optional
The time index on which to slice. If unspecified, a time average will be returned.
ignoreNaNs : bool, optional
If set, will use NaN-safe numpy operators.
radius : float, optional
Planet radius in meters (only used for summation)
latitudes: numpy.ndarray, optional
Latitude array--required if lat is "mean", or if either lat or lon is "sum"
longitudes: numpy.ndarray, optional
Longitude array--required if lon is "mean" or if either lat or lon is "sum"
Returns
-------
numpy.ndarray
A 2-D array
"""
if ignoreNaNs:
sumop = np.nansum
meanop = np.nanmean
else:
sumop = np.sum
meanop = np.mean
if len(variable.shape)==2:
return variable
if time:
try:
variable=variable[time,:]
except:
raise UnitError("You have probably passed a float time to a variable with no "+
"information about what that means. You should pass an integer "+
"time index instead")
elif time is None and len(variable.shape)>2:
variable=meanop(variable,axis=0)
elif time==0:
variable=variable[time,:]
if len(variable.shape)>2:
if lev is not None:
if type(lev)==int:
variable=variable[lev,:]
elif lev=="sum":
variable=sumop(variable,axis=0)
elif lev=="mean":
variable=meanop(variable,axis=0)
else:
raise UnitError("Unknown level specification")
elif lat is not None and lon is None:
if type(lat)==int:
variable=variable[:,lat,:]
elif lat=="sum":
if longitudes is None or latitudes is None:
raise Exception("You must provide latitudes and longitudes as 1D arrays.")
variable=latsum(variable,latitudes,dlon=longitudes[1]-longitudes[0],
radius=radius)
elif lat=="mean":
if latitudes is None:
raise Exception("You must provide latitudes as a 1D array.")
variable=latmean(variable,latitudes)
else:
raise UnitError("Unknown latitude specification")
elif lon is not None and lat is None:
if type(lon)==int:
variable=variable[:,:,lon]
elif lon=="sum":
if longitudes is None or latitudes is None:
raise Exception("You must provide latitudes and longitudes as 1D arrays.")
newvar = np.zeros(variable.shape[:-1])
gradlat = np.gradient(np.sin(latitudes*np.pi/180.))
for lt in range(len(latitudes)):
newvar[:,lt]=lonsum(variable,longitudes,dsinlat=gradlat[lt],radius=radius)
variable = newvar
elif lon=="mean":
if longitudes is None:
raise Exception("You must provide longitudes as a 1D array.")
variable=lonmean(variable,longitudes)
else:
raise UnitError("Unknown longitude specification")
else:
raise DimensionError("Inappropriate or insufficient dimensional constraints")
return variable
[docs]def spatialmath(variable,lat=None,lon=None,file=None,mean=True,time=None,
ignoreNaNs=True,lev=None,radius=6.371e6):
"""Compute spatial means or sums of data
Parameters
----------
variable : str, numpy.ndarray
The variable to operate on. Can either be a data array, or the name of a variable. If the latter, file must be specified.
lat,lon : numpy.ndarray, optional
Latitude and longitude arrays. If file is provided and lat and lon are not, they will be
extracted from the file.
file : str, optional
Path to a NetCDF output file to open and extract data from.
mean : bool, optional
If True, compute a global mean. If False, compute a global sum.
time : int, optional, or "all"
The time index on which to slice. If unspecified, a time average will be returned. If set to
"all", the time axis will be preserved.
ignoreNaNs : bool, optional
If True, use NaN-safe numpy operators.
lev : int, optional
If set, slice a 3D spatial array at the specified level.
radius : float, optional
Radius of the planet in meters. Only used if mean=False.
Returns
-------
float
"""
if ignoreNaNs:
sumop = np.nansum
meanop = np.nanmean
else:
sumop = np.sum
meanop = np.mean
if file:
ln,lt,variable = parse(file,variable,lat=lat,lon=lon)
else:
if type(lat)==type(None) or type(lon)==type(None):
raise DimensionError("Need to provide latitude and longitude data")
ln=lon
lt=lat
if time!="all":
variable = make2d(variable,time=time,lev=lev,ignoreNaNs=ignoreNaNs,longitudes=ln,latitudes=lt)
lt1 = np.zeros(len(lt)+1)
lt1[0] = 90
for n in range(0,len(lt)-1):
lt1[n+1] = 0.5*(lt[n]+lt[n+1])
lt1[-1] = -90
dln = np.diff(ln)[0]
ln1 = np.zeros(len(ln)+1)
ln1[0] = -dln
for n in range(0,len(ln)-1):
ln1[n+1] = 0.5*(ln[n]+ln[n+1])
ln1[-1] = 360.0-dln
lt1*=np.pi/180.0
ln1*=np.pi/180.0
darea = np.zeros((len(lt),len(ln)))
for jlat in range(0,len(lt)):
for jlon in range(0,len(ln)):
dln = ln1[jlon+1]-ln1[jlon]
darea[jlat,jlon] = abs(np.sin(lt1[jlat])-np.sin(lt1[jlat+1]))*abs(dln)
svar = variable*darea
if mean:
outvar = sumop(svar)/sumop(darea)
else:
outvar = sumop(svar) * radius**2
else:
ntimes = variable.shape[0]
variable_ = make2d(variable,time=0,lev=lev,ignoreNaNs=ignoreNaNs,longitudes=ln,latitudes=lt)
outvariable = np.zeros([ntimes,]+list(variable_.shape))
outvariable[0,...] = variable_[:]
for n in range(1,ntimes):
outvariable[n,...] = make2d(variable,time=n,lev=lev,ignoreNaNs=ignoreNaNs,
longitudes=ln,latitudes=lt)
lt1 = np.zeros(len(lt)+1)
lt1[0] = 90
for n in range(0,len(lt)-1):
lt1[n+1] = 0.5*(lt[n]+lt[n+1])
lt1[-1] = -90
dln = np.diff(ln)[0]
ln1 = np.zeros(len(ln)+1)
ln1[0] = -dln
for n in range(0,len(ln)-1):
ln1[n+1] = 0.5*(ln[n]+ln[n+1])
ln1[-1] = 360.0-dln
lt1*=np.pi/180.0
ln1*=np.pi/180.0
darea = np.zeros((len(lt),len(ln)))
for jlat in range(0,len(lt)):
for jlon in range(0,len(ln)):
dln = ln1[jlon+1]-ln1[jlon]
darea[jlat,jlon] = abs(np.sin(lt1[jlat])-np.sin(lt1[jlat+1]))*abs(dln)
svar = outvariable*darea[np.newaxis,:,:]
svar = np.reshape(svar,[ntimes,np.prod(variable_.shape)])
if mean:
outvar = sumop(svar,axis=1)/sumop(darea)
else:
outvar = sumop(svar,axis=1) * radius**2
return outvar
[docs]def latmean(variable,latitudes):
"""Compute meriodional mean (i.e. the variable that changes is latitude).
Compute the area-weighted mean of a latitude array :math:`x`\ , such that:
.. math::
\\bar{x} = \\frac{\sum_{i=1}^N |\\sin(\\phi_{i-1/2})-\\sin(\\phi_{i+1/2})|x_i}{\sum_{i=1}^N |\\sin(\\phi_{i-1/2})-\\sin(\\phi_{i+1/2})|}
Parameters
----------
variable : numpy.ndarray
Array to be averaged. Assumption is that if 2D, lat is the first dimension, if 3D, the second dimension, and if 4D. the 3rd dimension.
latitudes : array-like
Array or list of latitudes
Returns
-------
scalar or numpy.ndarray
Depending on the dimensionality of the input array, output may have 0, 1, or 2 dimensions.
"""
lt1 = np.zeros(len(latitudes)+1)
lt1[0] = 90
for n in range(0,len(latitudes)-1):
lt1[n+1] = 0.5*(latitudes[n]+latitudes[n+1])
lt1[-1] = -90
lt1*=np.pi/180.0
darea = np.zeros(latitudes.shape)
for jlat in range(0,len(latitudes)):
darea[jlat] = np.sin(lt1[jlat])-np.sin(lt1[jlat+1])
if len(variable.shape)==1:
return np.nansum(variable*darea)/np.nansum(darea)
elif len(variable.shape)==2:
return np.nansum(variable*darea[:,np.newaxis],axis=0)/np.nansum(darea[:,np.newaxis]*np.ones(variable.shape),axis=0)
elif len(variable.shape)==3:
return np.nansum(variable*darea[np.newaxis,:,np.newaxis],axis=1)/np.nansum(darea[np.newaxis,:,np.newaxis]*np.ones(variable.shape),axis=1)
elif len(variable.shape)==4:
return np.nansum(variable*darea[np.newaxis,np.newaxis,:,np.newaxis],axis=2)/np.nansum(darea[np.newaxis,np.newaxis,:,np.newaxis]*np.ones(variable.shape),axis=2)
else:
raise DimensionError("Variable must have 4 or fewer dimensions. Latitude should be the second-from the right-most dimension if there are 2 or more dimensions.")
[docs]def latsum(variable,latitudes,dlon=360.0,radius=6.371e6):
"""Compute meriodional sum (i.e. the variable that changes is latitude).
Compute the area-weighted sum of a latitude array :math:`x` given a longitude span :math:`\\Delta\\theta` and planet radius :math:`R`\ , such that:
.. math::
X = \sum_{i=1}^N |\\sin(\\phi_{i-1/2})-\\sin(\\phi_{i+1/2})|\\Delta\\theta R^2x_i
Parameters
----------
variable : numpy.ndarray
Array to be summed. Assumption is that if 2D, lat is the first dimension, if 3D, the second dimension, and if 4D. the 3rd dimension.
latitudes : array-like
Array or list of latitudes
dlon : float, optional
Longitude span in degrees.
radius : float, optional
Planet radius in meters.
Returns
-------
scalar or numpy.ndarray
Depending on the dimensionality of the input array, output may have 0, 1, or 2 dimensions.
"""
lt1 = np.zeros(len(latitudes)+1)
lt1[0] = 90
for n in range(0,len(latitudes)-1):
lt1[n+1] = 0.5*(latitudes[n]+latitudes[n+1])
lt1[-1] = -90
lt1*=np.pi/180.0
dlon *= np.pi/180.0
darea = np.zeros(latitudes.shape)
for jlat in range(0,len(latitudes)):
darea[jlat] = abs(np.sin(lt1[jlat])-np.sin(lt1[jlat+1]))*abs(dlon)*radius**2
if len(variable.shape)==1:
return np.nansum(variable*darea)
elif len(variable.shape)==2:
return np.nansum(variable*darea[:,np.newaxis],axis=0)
elif len(variable.shape)==3:
return np.nansum(variable*darea[np.newaxis,:,np.newaxis],axis=1)
elif len(variable.shape)==4:
return np.nansum(variable*darea[np.newaxis,np.newaxis,:,np.newaxis],axis=2)
else:
raise DimensionError("Variable must have 4 or fewer dimensions. Latitude should be the second-from the right-most dimension if there are 2 or more dimensions.")
[docs]def lonmean(variable,longitudes):
"""Compute zonal mean (i.e. the variable that changes is longitude).
Compute the area-weighted mean of a longitude array :math:`x`\ , such that:
.. math::
\\bar{x} = \\frac{\sum_{i=1}^N |\\theta_{i-1/2}-\\theta_{i+1/2}|x_i}{\sum_{i=1}^N |\\theta_{i-1/2}-\\theta_{i+1/2}|}
Parameters
----------
variable : numpy.ndarray
Array to be summed. Assumption is that longitude is always the last dimension.
Returns
-------
scalar or numpy.ndarray
Depending on the dimensionality of the input array, output may be a scalar or have N-1 dimensions.
"""
dlon = np.gradient(longitudes)
sumlon = np.nansum(dlon)
dlon = np.broadcast_to(dlon,variable.shape)
return np.nansum(variable*dlon,axis=-1)/sumlon
[docs]def lonsum(variable,longitudes,dsinlat=2.0,radius=6.371e6):
"""Compute zonal sum (i.e. the variable that changes is longitude).
Compute the area-weighted sum of a longitude array :math:`x` given a latitude span :math:`\\Delta\\sin\\phi` and planet radius :math:`R`\ , such that:
.. math::
X = \sum_{i=1}^N |\\theta_{i-1/2}-\\theta_{i+1/2}|\\Delta\\sin\\phi R^2x_i
Parameters
----------
variable : numpy.ndarray
Array to be summed. Assumption is that longitude is always the last dimension.
longitudes : array-like
Array or list of longitudes
dsinlat : float, optional
The sine-latitude span for the longitude span considered. The default is 2, corresponding to -90 degrees to 90 degrees.
radius : float, optional
Planet radius in meters.
Returns
-------
scalar or numpy.ndarray
Depending on the dimensionality of the input array, output may have 0, 1, or 2 dimensions.
"""
dlon = np.gradient(longitudes)*np.pi/180.0
darea = np.zeros(longitudes.shape)
for jlon in range(0,len(longitudes)):
darea[jlon] = abs(dsinlat)*abs(dlon[jlon])*radius**2
darea = np.broadcast_to(darea,variable.shape)
return np.nansum(variable*darea,axis=-1)
[docs]def cspatialmath(variable,lat=None,lon=None,file=None,mean=True,time=None,
ignoreNaNs=True,lev=None,radius=6.371e6,poles=False):
"""Compute spatial means or sums of data, but optionally don't go all the way to the poles.
Sometimes, saying that the latitudes covered go all the way to :math:`\pm90^\circ` results in
errors, and accurate accounting requires excluding the poles themselves. This function
is identical to spatialmath, except that it provides that option.
Parameters
----------
variable : str, numpy.ndarray
The variable to operate on. Can either be a data array, or the name of a variable. If the latter, file must be specified.
lat,lon : numpy.ndarray, optional
Latitude and longitude arrays. If file is provided and lat and lon are not, they will be
extracted from the file.
file : str, optional
Path to a NetCDF output file to open and extract data from.
mean : bool, optional
If True, compute a global mean. If False, compute a global sum.
time : int, optional
The time index on which to slice. If unspecified, a time average will be returned.
ignoreNaNs : bool, optional
If True, use NaN-safe numpy operators.
lev : int, optional
If set, slice a 3D spatial array at the specified level.
radius : float, optional
Radius of the planet in meters. Only used if mean=False.
poles : bool, optional
If False (default), exclude the poles.
Returns
-------
float
"""
if ignoreNaNs:
sumop = np.nansum
meanop = np.nanmean
else:
sumop = np.sum
meanop = np.mean
if file:
ln,lt,variable = parse(file,variable,lat=lat,lon=lon)
else:
if type(lat)==type(None) or type(lon)==type(None):
raise DimensionError("Need to provide latitude and longitude data")
ln=lon
lt=lat
variable = make2d(variable,time=time,lev=lev,ignoreNaNs=ignoreNaNs,
latitudes=lt,longitudes=ln)
lt1 = np.zeros(len(lt)+1)
dlt1 = abs(np.diff(lt)[0])
dlt2 = abs(np.diff(lt)[-1])
lt1[0] = lt[0]+0.5*dlt1
for n in range(0,len(lt)-1):
lt1[n+1] = 0.5*(lt[n]+lt[n+1])
lt1[-1] = lt[-1]-0.5*dlt2
if poles:
lt1[0] = 90.0
lt1[-1] = -90.0
dln = np.diff(ln)[0]
ln1 = np.zeros(len(ln)+1)
ln1[0] = ln[0]-dln*0.5
for n in range(0,len(ln)-1):
ln1[n+1] = 0.5*(ln[n]+ln[n+1])
ln1[-1] = ln[-1]+dln*0.5
lt1*=np.pi/180.0
ln1*=np.pi/180.0
darea = np.zeros((len(lt),len(ln)))
for jlat in range(0,len(lt)):
for jlon in range(0,len(ln)):
dln = ln1[jlon+1]-ln1[jlon]
darea[jlat,jlon] = abs(np.sin(lt1[jlat])-np.sin(lt1[jlat+1]))*abs(dln)
svar = variable*darea
if mean:
outvar = sumop(svar)/sumop(darea)
else:
outvar = sumop(svar) * radius**2
return outvar
[docs]def wrap2d(var):
'''Add one element to the longitude axis to allow for wrapping'''
newvar = np.zeros(np.array(var.shape)+np.array((0,1)))
newvar[:,:-1] = var[:,:]
newvar[:,-1] = var[:,0]
return newvar
[docs]def streamfxn(file,time=None):
'''Deprecated. Passes args to eqstream().'''
return eqstream(file,time=time)
[docs]def eqstream(file,radius=6.371e6,gravity=9.80665):
'''Compute the tidally-locked streamfunction
Parameters
----------
dataset : str or ExoPlaSim Dataset
Either path to ExoPlaSim Dataset of model output or an instance of the dataset.
plarad : float, optional
Planetary radius [m]
grav : float, optional
Surface gravity [m/s^2]
Returns
-------
numpy.ndarray(1D), numpy.ndarray(1D), numpy.ndarray(2D)
tidally-locked latitude, layer interface pressures, and TL streamfunction
'''
from scipy.integrate import cumtrapz
if type(dataset)==str:
dataset = _Dataset(dataset,"r")
lon = dataset.variables['lon'][:]
lat = dataset.variables['lat'][:]
#mva = np.nanmean(va_TL,axis=3) #tidally-locked meridional wind
#ps = spatialmath(dataset.variables['ps'][:],lon=lon,lat=lat)
lev = dataset.variables['lev'][:]
pa = dataset.variables['ps'][:,np.newaxis,:,:] * lev[np.newaxis,:,np.newaxis,np.newaxis] * 100.0
va = dataset.variables['va'][:]
nlon = len(lon)
nlat = len(lat)
ntime = pa.shape[0]
nlev = len(lev)
vadp = np.zeros(va.shape)
for nt in range(ntime):
for jlat in range(nlat):
for jlon in range(nlon):
vadp[nt,:jlat,jlon] = cumtrapz(va[nt,:,jlat,jlon],
x=pa[nt,:,jlat,jlon],
initial=0.0)
prefactor = 2*np.pi*radius/gravity*np.cos(lat*np.pi/180.0)
sign = 1 #-1 for synchronous, 1 for equatorial
stf = sign*prefactor[np.newaxis,np.newaxis,:,np.newaxis]*vadp
psurf = spatialmath(dataset.variables['ps'][:],lon=lon,lat=lat)
#mvadp = cumtrapz(mva,x=lev[:]*ps*100.0,axis=1) #integrate in Pa not hPa
#prefactor = 2*np.pi*plarad/grav #2piR/g
#stf = prefactor*np.cos(lat_TL[np.newaxis,np.newaxis,:]*np.pi/180.0)*mvadp
#pmid = 0.5*(lev[1:]+lev[:-1])*ps
return lat,psurf*lev,stf
[docs]def adist(lon1,lat1,lon2,lat2):
'''Return angular distance(s) in degrees between two points (or sets of points) on a sphere.
Parameters
----------
lon1 : float or numpy.ndarray
Longitudes of first point(s)
lat1 : float or numpy.ndarray
Latitudes of first point(s)
lon2 : float or numpy.ndarray
Longitudes of second point(s)
lat2 : float or numpy.ndarray
Latitudes of second point(s)
Returns
-------
float or numpy.ndarray
Angular distance(s) between given points
'''
rn1 = lon1*np.pi/180.0
rn2 = lon2*np.pi/180.0
rt1 = lat1*np.pi/180.0
rt2 = lat2*np.pi/180.0
dist = np.arccos(np.sin(rt1)*np.sin(rt2)+np.cos(rt1)*np.cos(rt2)*np.cos(rn1-rn2))
return dist*180.0/np.pi
[docs]def eq2tl_coords(lon,lat,substellar=0.0):
'''Compute tidally-locked coordinates of a set of equatorial lat-lon coordinates.
Transforms equatorial coordinates into a tidally-locked coordinate system where 0 degrees longitude is the substellar-south pole-antistellar meridian, and 90 degrees latitude is the substellar point, such that the evening hemisphere is 0-180 degrees longitude, the morning hemisphere is 180-360 degrees longitude, the north equatorial pole is at (0, 180), and easterly flow is counter-clockwise. Note that this differs from the coordinate system introduced in Koll & Abbot (2015) in that theirs is a left-handed coordinate system, with the south pole at (0, 180) and counter-clockwise easterly flow, which represents a south-facing observer inside the sphere, while ours is a right-handed coordinate system, representing a south-facing observer outside the sphere, which is the usual convention for spherical coordinate systems.
Parameters
----------
lon : numpy.ndarray
Longitudes in equatorial coordinates [degrees]
lat : numpy.ndarray
Latitudes in equatorial coordinates [degrees]
substellar : float, optional
Longitude of the substellar point. [degrees]
Returns
-------
numpy.ndarray, numpy.ndarray
Transformed longitudes and latitudes [degrees]
'''
lons,lats = np.meshgrid(lon,lat)
tlons = np.zeros(lons.shape)
tlats = np.zeros(lats.shape)
elons = lons*np.pi/180.0
elats = lats*np.pi/180.0
rss = substellar*np.pi/180.0
tlats = np.arcsin(np.cos(elats)*np.cos(elons-rss))*180.0/np.pi
tlons = np.arctan2(np.sin(elons-rss),-np.tan(elats))*180.0/np.pi
tlons[tlons<0] += 360.0
return tlons,tlats
[docs]def tl2eq_coords(lon,lat,substellar=0.0):
'''Compute equatorial coordinates of a set of tidally-locked lat-lon coordinates.
Transforms tidally-locked coordinates into the standard equatorial coordinate system. Note that in our tidally-locked coordinate system, 0 degrees longitude is the substellar-south pole-antistellar meridian, and 90 degrees latitude is the substellar point, such that the evening hemisphere is 0-180 degrees longitude, the morning hemisphere is 180-360 degrees longitude, the north equatorial pole is at (0, 180), and easterly flow is counter-clockwise. Note that this differs from the coordinate system introduced in Koll & Abbot (2015) in that theirs is a left-handed coordinate system, with the south pole at (0, 180) and counter-clockwise easterly flow, which represents a south-facing observer inside the sphere, while ours is a right-handed coordinate system, representing a south-facing observer outside the sphere, which is the usual convention for spherical coordinate systems.
Parameters
----------
lon : numpy.ndarray
Longitudes in tidally-locked coordinates [degrees]
lat : numpy.ndarray
Latitudes in tidally-locked coordinates [degrees]
substellar : float, optional
Longitude of the substellar point. [degrees]
Returns
-------
numpy.ndarray, numpy.ndarray
Transformed longitudes and latitudes [degrees]
'''
lons,lats = np.meshgrid(lon,lat)
qlons = np.zeros(lons.shape)
qlats = np.zeros(lats.shape)
tlons = lons*np.pi/180.0
tlats = lats*np.pi/180.0
rss = substellar*np.pi/180.0
qlats = np.arcsin(-np.cos(tlats)*np.cos(tlons))
qlons = np.arctan2(np.cos(tlats)*np.sin(tlons),np.sin(tlats)) + rss
qlats *= 180.0/np.pi
qlons *= 180.0/np.pi
qlons[qlons<0] += 360.0
return qlons,qlats
[docs]def eq2tl(variable,lon,lat,substellar=0.0, polemethod="interp"):
'''Transform a variable to tidally-locked coordinates
Note that in our tidally-locked coordinate system, 0 degrees longitude is the substellar-south pole-antistellar meridian, and 90 degrees latitude is the substellar point, such that the evening hemisphere is 0-180 degrees longitude, the morning hemisphere is 180-360 degrees longitude, the north equatorial pole is at (0, 180), and easterly flow is counter-clockwise. Note that this differs from the coordinate system introduced in Koll & Abbot (2015) in that theirs is a left-handed coordinate system, with the south pole at (0, 180) and counter-clockwise easterly flow, which represents a south-facing observer inside the sphere, while ours is a right-handed coordinate system, representing a south-facing observer outside the sphere, which is the usual convention for spherical coordinate systems.
Parameters
----------
variable : numpy.ndarray (2D, 3D, or 4D)
N-D data array to be transformed. Final two dimensions must be (lat,lon)
lon : numpy.ndarray
1D array of longitudes [deg]
lat : numpy.ndarray
1D array of latitudes [deg]
substellar : float, optional
Longitude of the substellar point (defaults to 0 degrees)
polemethod : str, optional
Interpolation method for polar latitudes. If "nearest", then instead of inverse-distance linear interpolation, will use nearest-neighbor. This is recommended for vector variables. For scalars, leave as "interp".
Returns
-------
numpy.ndarray, numpy.ndarray, numpy.ndarray
Transformed longitudes, latitudes, and data array.
'''
tlon = np.copy(lon)
tlat = np.copy(lat)
tlvariable = np.zeros(variable.shape)
nlon = len(lon)
nlat = len(lat)
elons,elats = tl2eq_coords(tlon,tlat,substellar=substellar)
elons *= np.pi/180.0
elats *= np.pi/180.0
for i in range(nlon):
for j in range(nlat):
rlon = elons[j,i]
rlat = elats[j,i]
dlon = rlon*180.0/np.pi
dlat = rlat*180.0/np.pi
if abs(dlat)>abs(lat).max():
if polemethod!="nearest":
jj = abs(lat-dlat).argmin()
distances = np.zeros(nlon)
for ii in range(nlon):
distances[ii] = 1.0/adist(lon[ii],dlon,lat[jj],dlat)
tlvariable[...,j,i] = np.average(variable[...,jj,:],weights=distances,axis=len(variable.shape)-2)
else:
ilat=np.argmin(abs(dlat-lat))
ilon=np.argmin(abs(dlon-lon))
tlvariable[...,j,i] = variable[...,ilat,ilon]
else:
latcomparison = abs(lat-dlat)
loncomparison = abs(lon-dlon)
colat = (latcomparison.min()==0.0) #We are colatitude
colon = (loncomparison.min()==0.0) #We are colongitude
if not colat and not colon:
jj = latcomparison.argmin()
ii = loncomparison.argmin()
if ii==0:
ln1 = abs(loncomparison[-1]-360.0-dlon) #if chosen, our indices will be -1 and 0
ln2 = loncomparison[ 1] #if chosen, our indices will be 0 and 1
elif ii==nlon-1:
ln1 = loncomparison[ii-1] #if chosen, our indices will be nlon-2 and nlon-1
ln2 = abs(360.0-dlon) #if chosen, our indices will be nlon-1 and 0
else:
ln1 = loncomparison[ii-1]
ln2 = loncomparison[ii+1]
if jj==0:
lt1 = abs(100.0-dlat) #shouldn't be chosen
lt2 = latcomparison[1]
elif jj==nlat-1:
lt1 = latcomparison[jj-1]
lt2 = abs(-100.0-dlat) #shouldn't be chosen
else:
lt1 = latcomparison[jj-1]
lt2 = latcomparison[jj+1]
ni = (ii + 2*np.argmin([ln1,ln2])-1)%nlon
nj = jj + 2*np.argmin([lt1,lt2])-1
if nj>=nlat or nj<0:
print(ii,jj,ni,nj,dlon,dlat)
raise
d11 = 1.0/adist(lon[ii],dlon,lat[jj],dlat)
d22 = 1.0/adist(lon[ni],dlon,lat[nj],dlat)
d21 = 1.0/adist(lon[ii],dlon,lat[nj],dlat)
d12 = 1.0/adist(lon[ni],dlon,lat[jj],dlat)
tlvariable[...,j,i] = np.average(np.array([variable[...,jj,ii],variable[...,jj,ni],
variable[...,nj,ii],variable[...,nj,ni]]),
weights=[d11,d12,d21,d22],axis=0)
elif colat: #only changing longitude
ii = loncomparison.argmin()
jj = latcomparison.argmin()
if ii==0:
ln1 = abs(loncomparison[-1]-360.0-dlon) #if chosen, our indices will be -1 and 0
ln2 = loncomparison[ 1] #if chosen, our indices will be 0 and 1
elif ii==nlon-1:
ln1 = loncomparison[ii-1] #if chosen, our indices will be nlon-2 and nlon-1
ln2 = abs(360.0-dlon) #if chosen, our indices will be nlon-1 and 0
else:
ln1 = loncomparison[ii-1]
ln2 = loncomparison[ii+1]
ni = (ii + 2*np.argmin([ln1,ln2])-1)%nlon
d1 = 1.0/adist(lon[ii],dlon,lat[jj],dlat)
d2 = 1.0/adist(lon[ni],dlon,lat[jj],dlat)
tlvariable[...,j,i] = np.average(np.array([variable[...,jj,ii],variable[...,jj,ni]]),
weights = [d1,d2],axis=0)
elif colon: #only changing latitude
ii = loncomparison.argmin()
jj = latcomparison.argmin()
if jj==0:
lt1 = abs(100.0-dlat) #shouldn't be chosen
lt2 = latcomparison[1]
elif jj==nlat-1:
lt1 = latcomparison[jj-1]
lt2 = abs(-100.0-dlat) #shouldn't be chosen
else:
lt1 = latcomparison[jj-1]
lt2 = latcomparison[jj+1]
nj = jj + 2*np.argmin([lt1,lt2])-1
d1 = 1.0/adist(lon[ii],dlon,lat[jj],dlat)
d2 = 1.0/adist(lon[ii],dlon,lat[nj],dlat)
tlvariable[...,j,i] = np.average(np.array([variable[...,jj,ii],variable[...,nj,ii]]),
weights = [d1,d2],axis=0)
else: #We coincide with a real point
ii = loncomparison.argmin()
jj = latcomparison.argmin()
tlvariable[...,j,i] = variable[...,jj,ii]
return tlon,tlat,tlvariable
[docs]def tl2eq(variable,lon,lat,substellar=0.0):
'''Transform a tidally-locked variable to standard equatorial coordinates
Note that in our tidally-locked coordinate system, 0 degrees longitude is the substellar-south pole-antistellar meridian, and 90 degrees latitude is the substellar point, such that the evening hemisphere is 0-180 degrees longitude, the morning hemisphere is 180-360 degrees longitude, the north equatorial pole is at (0, 180), and easterly flow is counter-clockwise. Note that this differs from the coordinate system introduced in Koll & Abbot (2015) in that theirs is a left-handed coordinate system, with the south pole at (0, 180) and counter-clockwise easterly flow, which represents a south-facing observer inside the sphere, while ours is a right-handed coordinate system, representing a south-facing observer outside the sphere, which is the usual convention for spherical coordinate systems.
Parameters
----------
variable : numpy.ndarray (2D, 3D, or 4D)
N-D data array to be transformed. Final two dimensions must be (lat,lon)
lon : numpy.ndarray
1D array of longitudes [deg]
lat : numpy.ndarray
1D array of latitudes [deg]
substellar : float, optional
Longitude of the substellar point (defaults to 0 degrees)
Returns
-------
numpy.ndarray, numpy.ndarray, numpy.ndarray
Transformed longitudes, latitudes, and data array.
'''
qlon = np.copy(lon)
qlat = np.copy(lat)
eqvariable = np.zeros(variable.shape)
nlon = len(lon)
nlat = len(lat)
qlons,qlats = eq2tl_coords(qlon,qlat,substellar=substellar)
qlons *= np.pi/180.0
qlats *= np.pi/180.0
for i in range(nlon):
for j in range(nlat):
rlon = qlons[j,i]
rlat = qlats[j,i]
dlon = rlon*180.0/np.pi
dlat = rlat*180.0/np.pi
if abs(dlat)>abs(lat).max():
jj = abs(lat-dlat).argmin()
distances = np.zeros(nlon)
for ii in range(nlon):
distances[ii] = 1.0/adist(lon[ii],dlon,lat[jj],dlat)
eqvariable[...,j,i] = np.average(variable[...,jj,:],weights=distances,axis=len(variable.shape)-2)
else:
latcomparison = abs(lat-dlat)
loncomparison = abs(lon-dlon)
colat = (latcomparison.min()==0.0) #We are colatitude
colon = (loncomparison.min()==0.0) #We are colongitude
if not colat and not colon:
jj = latcomparison.argmin()
ii = loncomparison.argmin()
if ii==0:
ln1 = abs(loncomparison[-1]-360.0-dlon) #if chosen, our indices will be -1 and 0
ln2 = loncomparison[ 1] #if chosen, our indices will be 0 and 1
elif ii==nlon-1:
ln1 = loncomparison[ii-1] #if chosen, our indices will be nlon-2 and nlon-1
ln2 = abs(360.0-dlon) #if chosen, our indices will be nlon-1 and 0
else:
ln1 = loncomparison[ii-1]
ln2 = loncomparison[ii+1]
if jj==0:
lt1 = abs(100.0-dlat) #shouldn't be chosen
lt2 = latcomparison[1]
elif jj==nlat-1:
lt1 = latcomparison[jj-1]
lt2 = abs(-100.0-dlat) #shouldn't be chosen
else:
lt1 = latcomparison[jj-1]
lt2 = latcomparison[jj+1]
ni = (ii + 2*np.argmin([ln1,ln2])-1)%nlon
nj = jj + 2*np.argmin([lt1,lt2])-1
if nj>=nlat or nj<0:
print(ii,jj,ni,nj,dlon,dlat)
raise
d11 = 1.0/adist(lon[ii],dlon,lat[jj],dlat)
d22 = 1.0/adist(lon[ni],dlon,lat[nj],dlat)
d21 = 1.0/adist(lon[ii],dlon,lat[nj],dlat)
d12 = 1.0/adist(lon[ni],dlon,lat[jj],dlat)
eqvariable[...,j,i] = np.average(np.array([variable[...,jj,ii],variable[...,jj,ni],
variable[...,nj,ii],variable[...,nj,ni]]),
weights=[d11,d12,d21,d22],axis=0)
elif colat: #only changing longitude
ii = loncomparison.argmin()
jj = latcomparison.argmin()
if ii==0:
ln1 = abs(loncomparison[-1]-360.0-dlon) #if chosen, our indices will be -1 and 0
ln2 = loncomparison[ 1] #if chosen, our indices will be 0 and 1
elif ii==nlon-1:
ln1 = loncomparison[ii-1] #if chosen, our indices will be nlon-2 and nlon-1
ln2 = abs(360.0-dlon) #if chosen, our indices will be nlon-1 and 0
else:
ln1 = loncomparison[ii-1]
ln2 = loncomparison[ii+1]
ni = (ii + 2*np.argmin([ln1,ln2])-1)%nlon
d1 = 1.0/adist(lon[ii],dlon,lat[jj],dlat)
d2 = 1.0/adist(lon[ni],dlon,lat[jj],dlat)
eqvariable[...,j,i] = np.average(np.array([variable[...,jj,ii],variable[...,jj,ni]]),
weights = [d1,d2],axis=0)
elif colon: #only changing latitude
ii = loncomparison.argmin()
jj = latcomparison.argmin()
if jj==0:
lt1 = abs(100.0-dlat) #shouldn't be chosen
lt2 = latcomparison[1]
elif jj==nlat-1:
lt1 = latcomparison[jj-1]
lt2 = abs(-100.0-dlat) #shouldn't be chosen
else:
lt1 = latcomparison[jj-1]
lt2 = latcomparison[jj+1]
nj = jj + 2*np.argmin([lt1,lt2])-1
d1 = 1.0/adist(lon[ii],dlon,lat[jj],dlat)
d2 = 1.0/adist(lon[ii],dlon,lat[nj],dlat)
eqvariable[...,j,i] = np.average(np.array([variable[...,jj,ii],variable[...,nj,ii]]),
weights = [d1,d2],axis=0)
else: #We coincide with a real point
ii = loncomparison.argmin()
jj = latcomparison.argmin()
eqvariable[...,j,i] = variable[...,jj,ii]
return qlon,qlat,eqvariable
[docs]def eq2tl_uv(u, v, lon, lat, substellar=0.0):
'''Transform velocity variables to tidally-locked coordinates
Parameters
----------
u : numpy.ndarray (2D, 3D, or 4D)
N-D data array of zonal velocities to be transformed. Final two dimensions must be (lat,lon)
v : numpy.ndarray (2D, 3D, or 4D)
N-D data array of meridional velocities to be transformed. Final two dimensions must be (lat,lon)
lon : numpy.ndarray
1D array of longitudes [deg]
lat : numpy.ndarray
1D array of latitudes [deg]
substellar : float, optional
Longitude of the substellar point (defaults to 0 degrees)
Returns
-------
numpy.ndarray, numpy.ndarray, numpy.ndarray, numpy.ndarray
Transformed longitudes, latitudes, and velocity data arrays.
'''
lon_tl,lat_tl,uq_tl = eq2tl(u,lon,lat,substellar=substellar,polemethod='nearest')
lon_tl,lat_tl,vq_tl = eq2tl(v,lon,lat,substellar=substellar,polemethod='nearest')
#lon_tl,lat_tl = eq2tl_coords(lon,lat,substellar=substellar)
lons,lats = np.meshgrid(lon,lat)
rlons = substellar*np.pi/180.0 - lons*np.pi/180.0
rlats = lats*np.pi/180.0
qlons,qlats = tl2eq_coords(lon,lat,substellar=substellar)
rqlons = substellar*np.pi/180.0 - qlons*np.pi/180.0
rqlats = qlats*np.pi/180.0
rtlons,rtlats = np.meshgrid(lon_tl,lat_tl)
rtlons *= np.pi/180.0
rtlats *= np.pi/180.0
if len(u.shape)==3:
rlons = rlons[np.newaxis,:,:]
rlats = rlats[np.newaxis,:,:]
rqlons = rqlons[np.newaxis,:,:]
rqlats = rqlats[np.newaxis,:,:]
rtlons = rtlons[np.newaxis,:,:]
rtlats = rtlats[np.newaxis,:,:]
elif len(u.shape)==4:
rqlons = rqlons[np.newaxis,np.newaxis,:,:]
rqlats = rqlats[np.newaxis,np.newaxis,:,:]
rlons = rlons[np.newaxis,np.newaxis,:,:]
rlats = rlats[np.newaxis,np.newaxis,:,:]
rtlons = rtlons[np.newaxis,np.newaxis,:,:]
rtlats = rtlats[np.newaxis,np.newaxis,:,:]
ufactor = -np.cos(rtlats)/(np.sin(rqlats)*(1+np.sin(rqlons)**2/np.tan(rqlats)**2))
vfactor = np.sin(rqlats)/np.sqrt(1-np.cos(rqlats)**2*np.cos(rqlons)**2)
u_tl = (vq_tl*np.sin(rqlons)/np.sin(rqlats) + uq_tl*np.cos(rqlons))*ufactor
v_tl = (uq_tl*np.sin(rqlons)/np.sin(rqlats) - vq_tl*np.cos(rqlons))*vfactor
return lon_tl,lat_tl,u_tl,v_tl
[docs]def tlstream(dataset,radius=6371.0e3,gravity=9.80665,substellar=0.0):
'''Compute the tidally-locked streamfunction
Parameters
----------
dataset : str or ExoPlaSim Dataset
Either path to ExoPlaSim Dataset of model output or an instance of the dataset.
radius : float, optional
Planetary radius [m]
gravity : float, optional
Surface gravity [m/s^2]
substellar : float, optional
Longitude of the substellar point in degrees.
Returns
-------
numpy.ndarray(1D), numpy.ndarray(1D), numpy.ndarray(2D)
tidally-locked latitude, layer interface pressures, and TL streamfunction
'''
from scipy.integrate import cumtrapz
if type(dataset)==str:
dataset = _Dataset(dataset,"r")
lon = dataset.variables['lon'][:]
lat = dataset.variables['lat'][:]
lon_TL,lat_TL,ua_TL,va_TL = eq2tl_uv(dataset.variables['ua'][:],
dataset.variables['va'][:],
lon,lat,substellar=substellar)
#mva = np.nanmean(va_TL,axis=3) #tidally-locked meridional wind
#ps = spatialmath(dataset.variables['ps'][:],lon=lon,lat=lat)
lev = dataset.variables['lev'][:]
lev = np.array(lev)
ps = dataset.variables['ps']
ps = np.array(ps)
pa = ps[:,np.newaxis,:,:] * lev[np.newaxis,:,np.newaxis,np.newaxis] * 100.0
nlon = len(lon)
nlat = len(lat)
ntime = pa.shape[0]
nlev = len(lev)
vadp = np.zeros(va_TL.shape)
for nt in range(ntime):
for jlat in range(nlat):
for jlon in range(nlon):
vadp[nt,:,jlat,jlon] = cumtrapz(va_TL[nt,:,jlat,jlon],x=pa[nt,:,jlat,jlon],initial=0.0)
prefactor = 2*np.pi*radius/gravity*np.cos(lat*np.pi/180.0)
sign = -1 #-1 for synchronous, 1 for equatorial
stf = sign*prefactor[np.newaxis,np.newaxis,:,np.newaxis]*vadp
psurf = spatialmath(dataset.variables['ps'][:],lon=lon,lat=lat)
#mvadp = cumtrapz(mva,x=lev[:]*ps*100.0,axis=1) #integrate in Pa not hPa
#prefactor = 2*np.pi*plarad/grav #2piR/g
#stf = prefactor*np.cos(lat_TL[np.newaxis,np.newaxis,:]*np.pi/180.0)*mvadp
#pmid = 0.5*(lev[1:]+lev[:-1])*ps
return lat_TL,psurf*lev,stf
[docs]def orthographic(lon,lat,imap,central_longitude=0,central_latitude=0,ny=200,nx=200,interp='bilinear'):
'''Perform an orthographic projection.
Parameters
----------
lon : numpy.ndarray (1D)
Longitude array [degrees]
lat : numpy.ndarray (1D)
Latitude array [degrees]
imap : numpy.ndarray
Data array to be projected. The first two dimensions should be (lat,lon)
central_longitude : float, optional
Longitude in degrees to be centered beneath the observer
central_latitude : float, optional
Latitude in degrees to be centered beneath the observer
ny : int, optional
Number of pixels in the Y direction for the output projection
nx : int, optional
Number of pixels in the X direction for the output projection
interp : str, optional
Interpolation to use. Currently only 'bilinear' is accepted; otherwise nearest-neighbor will be used.
Returns
-------
numpy.ndarray (ny,nx)
The projected output
'''
l0 = central_longitude
p0 = central_latitude
shape = np.array(imap.shape)
shape[0] = ny
shape[1] = nx
xymap = np.zeros(tuple(shape))
xymap[:] = 0.0
coords = np.zeros((ny,nx,2))
coords2 = np.zeros((ny,nx,2,3))
coords3 = np.zeros((ny,nx,2))
rad=0.5*(8*nx/18.0+8*ny/18.0)
if lon.min()<0: #we're in a -180 to 180 domain
if l0>180.0:
l0 -= 360.0
if lon.max()>180.0:
zlon[zlon>180]-=360.0
dl0 = l0 #lon-dl0 and l0-dl0 will rotate things so l0=0.
l0 -= dl0
imap = np.roll(imap,-int(dl0/np.diff(lon)[0]),axis=1) #l0 should now be centered in the data array
else:
if l0<0:
l0 += 360.0
dl0 = l0-180. #l0-dl0 = 180
l0 -= dl0
imap = np.roll(imap,-int(dl0/np.diff(lon)[0]),axis=1) #l0 should now be centered in the data array
#Add ghost cells
zlat = np.zeros(np.array(lat.shape)+2)
zlon = np.zeros(np.array(lon.shape)+2)
dlat = np.diff(lat)
dlon = np.diff(lon)
zlon[0] = lon[0]-dlon[0]
zlon[1:-1] = lon[:]
zlon[-1] = lon[-1]+dlon[-1]
zlat[0] = lat[0]-dlat[0]
zlat[-1] = lat[-1]+dlat[-1]
zlat[1:-1] = lat[:]
newshape = np.array(imap.shape)
newshape[0] += 2
newshape[1] += 2
zmap = np.zeros(tuple(newshape))
zmap[1:-1,1:-1] = imap[:]
#At the poles, ghost cells are just duplicates
zmap[0,...] = zmap[1,...]
zmap[-1,...] = zmap[-2,...]
#At the left and right edges, ghost cells are cyclical
zmap[:,0] = zmap[:,-2]
zmap[:,-1] = zmap[:,1]
p0 *= np.pi/180.0
l0 *= np.pi/180.0
xx = np.arange(nx)-nx/2
yy = np.arange(ny)-ny/2
for j in range(0,ny):
jy = yy[j]
for i in range(0,nx):
ix = xx[i]
if (ix**2+jy**2)<=rad**2:
rho = np.sqrt(ix**2+jy**2)
cc = np.arcsin(rho/rad)
if rho==0:
phi = p0
lamb = l0
else:
phi = np.arcsin(np.cos(cc)*np.sin(p0) + jy*np.sin(cc)*np.cos(p0)/rho)
lamb = l0 + np.arctan2(ix*np.sin(cc),(rho*np.cos(cc)*np.cos(p0)-jy*np.sin(cc)*np.sin(p0)))
phi *= 180.0/np.pi
lamb *= 180.0/np.pi
if lamb<0 and lon.min()>=0:
lamb += 360
if lamb>180 and lon.min()<0:
lamb -= 360
jlat = np.argmin(abs(phi-zlat))
jlon = np.argmin(abs(lamb-zlon))
if interp=="bilinear":
xlat1 = np.where(zlat<phi)[0]
xlat2 = np.where(zlat>=phi)[0]
xlon1 = np.where(zlon<lamb)[0]
xlon2 = np.where(zlon>=lamb)[0]
if len(xlat1)>0 and len(xlat2)>0 and len(xlon1)>0 and len(xlon2)>0:
jlat1 = np.where(zlat==zlat[xlat1].max())[0][0]
jlat2 = np.where(zlat==zlat[xlat2].min())[0][0]
jlon1 = np.where(zlon==zlon[xlon1].max())[0][0]
jlon2 = np.where(zlon==zlon[xlon2].min())[0][0]
p1 = zlat[jlat1]
p2 = zlat[jlat2]
l1 = zlon[jlon1]
l2 = zlon[jlon2]
#print "interpolating;",jlat1,jlat2,jlon1,jlon2,lamb#,zlon[np.where(zlon<lamb)]
dl = l2-l1
dp = p2-p1
try:
if dl==0 and dp>0:
fxy1 = zmap[jlat1,jlon]
fxy2 = zmap[jlat2,jlon]
xymap[j,i] = (p2-phi)/dp*fxy1 + (phi-p1)/dp*fxy2
elif dp==0 and dl>0:
fx1y = zmap[jlat,jlon1]
fx2y = zmap[jlat,jlon2]
xymap[j,i] = (l2-lamb)/dl*fx1y + (lamb-l1)/dl*fx2y
elif dp==0 and dl==0:
xymap[j,i] = zmap[jlat,jlon]
else:
fxy1 = (l2-lamb)/dl*zmap[jlat1,jlon1] + (lamb-l1)/dl*zmap[jlat1,jlon2]
fxy2 = (l2-lamb)/dl*zmap[jlat2,jlon1] + (lamb-l1)/dl*zmap[jlat2,jlon2]
xymap[j,i] = (p2-phi)/dp*fxy1 + (phi-p1)/dp*fxy2
except BaseException as err:
print(p1,p2,l1,l2,jlat1,jlat2,jlon1,jlon2)
raise err
else:
#print jlat1,jlat2,jlon1,jlon2
jlat1=-200
jlat2=-200
xymap[j,i] = zmap[jlat,jlon]
else:
#print "not bilinear"
xymap[j,i] = zmap[jlat,jlon]
return xymap
[docs]def load(filename,csvbuffersize=1):
'''Open a postprocessed ExoPlaSim output file.
Supported formats include netCDF, CSV/TXT (can be compressed), NumPy, and HDF5. If the data
archive is a group of files that are not tarballed, such as a directory of CSV/TXT or gzipped
files, then the filename should be the name of the directory with the final file extension.
For example, if the dataset is a group of CSV files in a folder called "MOST_output.002", then
`filename` ought to be "MOST_output.002.csv", even though no such file exists.
When accessing a file archive comprised of CSV/TXT files such as that described above, only part
of the archive will be extracted/read into memory at once, with the exception of the first read,
when the entire archive is extracted to read header information. Dimensional arrays, such as
latitude, longitude, etc will be ready into memory and stored as attributes of the returned
object (but are accessed with the usual dictionary pattern). Other data arrays however may need to
be extracted and read from the archive. A memory buffer exists to hold recently-accessed arrays
in memory, which will prioritize the most recently-accessed variables. The number of variables
that can be stored in memory can be set with the `csvbuffersize` keyword. The default is 1. This
means that the first time the variable is accessed, access times will be roughly the time it takes
to extract the file and read it into memory. Subsequent accesses, however, will use RAM speeds.
Once the variable has left the buffer, due to other variables being accessed, the next access will
return to file access speeds. This behavior is intended to mimic the npz, netcdf, and hdf5 protocols.
Parameters
----------
filename : str
Path to the file
csvbuffersize : int, optional
If the file (or group of files) is a file archive such as a directory, tarball, etc, this is
the number of variables to keep in a memory buffer when the archive is accessed.
Returns
-------
object
``gmct._Dataset`` object that can be queried like a netCDF file.
'''
#fileparts = filename.split('.')
#if fileparts[-1] == "nc":
if type(csvbuffersize)==str: #This is probably a legacy netcdf load mistake
try:
csvbuffersize = int(csvbuffersize)
except:
csvbuffersize = 1
output=_Dataset(filename,csvbuffersize=csvbuffersize) #Usually _Dataset calls load(), but _Dataset calls _loadnetcdf
#directly, so here we're going to defer to _Dataset and make use
#of the close() functionality
#elif fileparts[-1] == "npz" or fileparts[-1] == "npy":
#output=_loadnpsavez(filename)
#elif (fileparts[-1] in ("csv","txt","gz","tar") or \
#fileparts[-2]+"."+fileparts[-1] in ("tar.gz","tar.bz2","tar.xz")):
#output,meta,files=_loadcsv(filename)
#elif fileparts[-1] in ("hdf5","h5","he5"):
#output=_loadhdf5(filename)
#else:
#raise Exception("Unsupported output format detected. Supported formats are:\n%s"%("\n\t".join(SUPPORTED)))
return output
#def rhines(U,lat,lon,plarad=6371.0,daylen=15.0,beta=None):
#'''Return the nondimensional Rhines length scale L_R/a
#Parameters
#----------
#U : numpy.ndarray or float
#Characteristic velocity [m/s] of the atmospheric jets
#lat : numpy.ndarray
#1D array of latitudes [deg]
#lon : numpy.ndarray
#1D array of longitudes [deg]
#plarad : float, optional
#Planetary radius in km
#daylen : float, optional
#Planetary rotation periods in days
#Returns
#-------
#float
#'''
##if beta is None:
#Omega = 2*np.pi/(daylen*86400.0)
##beta = np.gradient(2*Omega*np.sin(lat*np.pi/180.0),
##plarad*1e3*(lat*np.pi/180.0))
##lons,lats = np.meshgrid(lon,lat)
##beta = np.ones(U.shape)*beta[np.newaxis,:,np.newaxis]
##beta = gt.spatialmath(beta,lat=lat,lon=lon)
#beta = 2*Omega/(plarad*1e3) #This is the equatorial approximation
#Lr = np.sqrt(U/beta)/(plarad*1e3)
#return Lr
#def getheight(T,P,P0,gascon,grav,lat=None,lon=None,lapse=None):
#ndims = len(T.shape)
#if lapse is None:
#altz = -gascon*T/grav*np.log(P/P0)
#else:
#if ndims==1:
#T0 = T[-1]
#elif ndims==2:
#T0 = np.nanmean(T[-1,:])
#elif ndims==3:
#if lat is None:
#lat = np.zeros(90,-90,num=T.shape[1])
#if lon is None:
#lon = np.zeros(0,360,num=T.shape[2])
#T0 = gt.spatialmath(T[-1,:,:],lat=lat,lon=lon)
#elif ndims==4:
#if lat is None:
#lat = np.zeros(90,-90,num=T.shape[1])
#if lon is None:
#lon = np.zeros(0,360,num=T.shape[2])
#T0 = gt.spatialmath(T[:,-1,:,:],lat=lat,lon=lon)
#altz = T0/lapse * ((P/P0)**(-lapse*gascon/grav)-1)
#return altz
#def Ngradient(var,x,axis=0):
#'''Assumes x has the same shape as var'''
#ndims = len(var.shape)
#if ndims==1: #why are you even using this function
#return np.gradient(var,x)
#elif ndims==2:
#grad = np.zeros(var.shape)
#if axis==0:
#for k in range(var.shape[1]):
#grad[:,k] = np.gradient(var[:,k],x[:,k])
#else:
#for k in range(var.shape[0]):
#grad[k,:] = np.gradient(var[k,:],x[k,:])
#elif ndims==3:
#grad = np.zeros(var.shape)
#if axis==0:
#for j in range(var.shape[1]):
#for k in range(var.shape[2]):
#grad[:,j,k] = np.gradient(var[:,j,k],x[:,j,k])
#elif axis==1:
#for j in range(var.shape[0]):
#for k in range(var.shape[2]):
#grad[j,:,k] = np.gradient(var[j,:,k],x[j,:,k])
#else:
#for j in range(var.shape[0]):
#for k in range(var.shape[1]):
#grad[j,k,:] = np.gradient(var[j,k,:],x[j,k,:])
#else:
#grad = np.zeros(var.shape)
#if axis==0:
#for j in range(var.shape[1]):
#for k in range(var.shape[2]):
#for l in range(var.shape[3]):
#grad[:,j,k,l] = np.gradient(var[:,j,k,l],
#x[:,j,k,l])
#elif axis==1:
#for j in range(var.shape[0]):
#for k in range(var.shape[2]):
#for l in range(var.shape[3]):
#grad[j,:,k,l] = np.gradient(var[j,:,k,l],
#x[j,:,k,l])
#elif axis==2:
#for j in range(var.shape[0]):
#for k in range(var.shape[1]):
#for l in range(var.shape[3]):
#grad[j,k,:,l] = np.gradient(var[j,k,:,l],
#x[j,k,:,l])
#else:
#for j in range(var.shape[0]):
#for k in range(var.shape[1]):
#for l in range(var.shape[2]):
#grad[j,k,l,:] = np.gradient(var[j,k,l,:],
#x[j,k,l,:])
#return grad
#def bruntvasaila(T,P,lat=None,lon=None,lapse=False,gascon=287.0,grav=9.80665,cp=1006.0):
#'''T and P must have the same shape'''
#P0 = np.nanmax(P)
#theta = T*(P0/P)**(gascon/cp)
#ndims = len(T.shape)
#if not lapse:
#altz = -gascon*T/grav*np.log(P/P0)
#else: #Iteratively find the lapse rates and altitudes
#altz = getheight(T,P,P0,gascon,grav,lapse=False,lat=lat,lon=lon)
#if ndims==4:
#lapse = Ngradient(T,altz,axis=1)
#else:
#lapse = Ngradient(T,altz,axis=0)
#for n in range(40):
#altz = getheight(T,P,P0,gascon,grav,lapse=lapse,lat=lat,lon=lon)
#if ndims==4:
#lapse = Ngradient(T,altz,axis=1)
#else:
#lapse = Ngradient(T,altz,axis=0)
#if ndims==4:
#N = np.sqrt(grav/theta * Ngradient(theta,altz,axis=1))
#else:
#N = np.sqrt(grav/theta * Ngradient(theta,altz,axis=0))
#return N
#def rossby(T,P,plarad=6371.0,daylen=15.0,gascon=287.0,grav=9.80665,
#lat=None,lon=None,lapse=False,cp=1006.0):
#'''Return the Rossby deformation number L_R/a
#Options
#-------
#T : Air temperature [K] -- numpy.ndarray
#P : Air pressure -- numpy.ndarray
#plarad : Planetary radius in km (optional)
#daylen : Planetary rotation period in days (optional)
#gascon : Specific gas constant (optional)
#grav : Surface gravity [m/s^2] (optional)
#cp : specific heat capacity of air at constant presure [J/kg]
#lat : Latitudes in degrees (optional)
#lon : Longitudes in degrees (optional)
#lapse : Whether to compute the lapse rate {True/False} (optional)
#Returns
#-------
#float
#'''
#N = bruntvasaila(T,P,lat=lat,lon=lon,lapse=lapse,gascon=gascon,
#grav=grav,cp=cp)
#Omega = 2*np.pi/(daylen*86400.0)
##beta = np.gradient(2*Omega*np.sin(lat*np.pi/180.0),
## plarad*1e3*(lat*np.pi/180.0))
##lons,lats = np.meshgrid(lon,lat)
##beta = np.ones(U.shape)*beta[np.newaxis,:,np.newaxis]
##beta = gt.spatialmath(beta,lat=lat,lon=lon)
#beta = 2*Omega/(plarad*1e3)
#scaleH = gascon*T/grav
#Ld = np.sqrt(N*scaleH/(2*beta))/(plarad*1e3)
#ndims = len(Ld.shape)
#return Ld
#if ndims<3:
#return np.nanmean(Ld)
#elif ndims==3:
#return gt.spatialmath(np.nanmean(Ld,axis=0),lon=lon,lat=lat)
#else:
#return gt.spatialmath(np.nanmean(Ld,axis=1),lon=lon,lat=lat)
#def deform(T,lat,lon,plarad=6371.0,daylen=15.0,gascon=287.0,grav=9.80665):
#'''Return the equatorial (Rossby) deformation length scale L_R/a
#Options
#-------
#T : characteristic temperature [K]
#plarad : Planetary radius in km (optional)
#daylen : Planetary rotation periods in days (optional)
#gascon : Specific gas constant (optional)
#grav : Surface gravity [m/s^2] (optional)
#Returns
#-------
#float
#'''
#Omega = 2*np.pi/(daylen*86400.0)
##beta = np.gradient(2*Omega*np.sin(lat*np.pi/180.0),
## plarad*1e3*(lat*np.pi/180.0))
##lons,lats = np.meshgrid(lon,lat)
##beta = np.ones(U.shape)*beta[np.newaxis,:,np.newaxis]
##beta = gt.spatialmath(beta,lat=lat,lon=lon)
#beta = 2*Omega/(plarad*1e3)
#scaleH = gascon*T/grav
#c = np.sqrt(grav*scaleH)
#Le = np.sqrt(c/beta)/(plarad*1e3)
#return Le