def mass_streamfun(self):
from scipy import integrate
data = self._obj
# lonlen = len(data.lon)
if 'lon' in data.dims:
data = data.fillna(0).mean('lon')
levax = data.get_axis_num('lev')
stream = integrate.cumtrapz(data * np.cos(np.deg2rad(data.lat)), x=data.lev * 1e2, initial=0., axis=levax)
stream = stream * 2 * np.pi / cc.g * cc.rearth * 1e-9
stream = xr.DataArray(stream, coords=data.coords, dims=data.dims)
stream = stream.rename('ovt')
stream.attrs['long name'] = 'atmosphere overturning circulation'
stream.attrs['unit'] = 'Sv (1e9 kg/s)'
return stream
python类cumtrapz()的实例源码
def write_ETR(mtx, dt, i):
"""
Save the ETR in human readable format
"""
file_name = "electronTranferRates_fragment_{}.txt".format(i)
header = 'electron_Density electron_ETR(Nonadiabatic Adibatic) hole_density hole_ETR(Nonadiabatic Adibatic)'
# Density of Electron/hole
density_electron = mtx[1:, 0]
density_hole = mtx[1:, 3]
# Integrate the nonadiabatic/adiabatic components of the electron/hole ETR
int_elec_nonadia, int_elec_adia, int_hole_nonadia, int_hole_adia = [
integrate.cumtrapz(mtx[:, k], dx=dt) for k in [1, 2, 4, 5]]
# Join the data
data = np.stack(
(density_electron, int_elec_nonadia, int_elec_adia, density_hole,
int_hole_nonadia, int_hole_adia), axis=1)
# save the data in human readable format
np.savetxt(file_name, data, fmt='{:^3}'.format('%e'), header=header)
def another_integral(self,lst, time_lst=None):
acc = np.array(lst, dtype=np.float32)
if time_lst != None:
time = time_lst
else:
time = self.raw_data['time']
vel = integrate.cumtrapz(acc, time, initial=0)
dist = integrate.cumtrapz(vel, time, initial=0)
return acc, vel, dist
def integrate_array(x, y, axis=0):
y = np.array(y).swapaxes(axis, -1)
inty = np.zeros(y.shape)
myslice = [slice(0, i)
for i in aslist(y.shape)[:-1]]+[slice(1, y.shape[-1])]
inty[myslice] = scipyintegratecumtrapz(y, x)
return inty.swapaxes(axis, -1)
def compute_heat_transport(self, dsarray, method):
from scipy import integrate
'''
compute heat transport using surface(toa,sfc) flux
'''
lat_rad = np.deg2rad(dsarray.lat)
coslat = np.cos(lat_rad)
field = coslat * dsarray
if method is "Flux_adjusted":
field = field - field.mean("lat")
print("The heat transport is computed by Flux adjestment.")
elif method is "Flux":
print("The heat transport is computed by Flux.")
elif method is "Dynamic":
print("The heat transport is computed by dynamic method.")
raise ValueError("Dynamic method has not been implimented.")
else:
raise ValueError("Method is not supported.")
try:
latax = field.get_axis_num('lat')
except:
raise ValueError('No lat coordinate!')
integral = integrate.cumtrapz(field, x=lat_rad, initial=0., axis=latax)
transport = 1e-15 * 2 * np.math.pi * integral * cc.rearth **2 # unit in PW
if isinstance(field, xr.DataArray):
result = field.copy()
result.values = transport
return result.T
# heat transport
# @property
def ocn_heat_transport(self, dlat=1, grid='g16'):
from scipy import integrate
# check time dimension
if 'time' in self._obj.dims:
flux = self._obj.SHF.mean('time')
else:
flux = self._obj.SHF
area = dict(g16=utl.tarea_g16, g35=utl.tarea_g35, g37=utl.tarea_g37)
flux_area = flux * area[grid] * 1e-4 # convert to m2
#
lat_bins = np.arange(-90,91,dlat)
lat = np.arange(-89.5,90,dlat)
if 'TLAT' in flux_area.coords.keys():
flux_lat = flux_area.groupby_bins('TLAT', lat_bins, labels = lat).sum()
latax = flux_lat.get_axis_num('TLAT_bins')
elif 'ULAT' in flux_area.coords.keys():
flux_lat = flux_area.groupby_bins('ULAT', lat_bins, labels = lat).sum()
latax = flux_lat.get_axis_num('ULAT_bins')
flux_lat.values = flux_lat - flux_lat.mean() # remove bias
flux_lat.values = np.nan_to_num(flux_lat.values)
integral = integrate.cumtrapz(flux_lat, x=None, initial=0., axis=latax)
OHT = flux_lat.copy()
OHT.values = integral *1e-15
return OHT
def samples(lmin,lmax,samples,lstar,alpha,intsamples=100000,constant=False):
"""
this is the main function for lum_func_sample
lmin is the minimum luminosity to be considered
lmax is the maximum luminosity to be considered
samples of the number of galaxy luminosities to be returned
lstar is the lstar value for the schechter function
alpha is the alpha value for the schechter function
intsamples is the number of samples used internally for integration
constant is a boolean used for testing. If its true, it will output
samples number of galaxies all with lmin luminosity
"""
if constant:
gal = lmin*np.ones(samples)
return gal
loglmin = log10(lmin)
loglmax =log10(lmax)
#first, i will make an array of luminosities using logspace
l=np.logspace(loglmin,loglmax,intsamples)
lf=schechter(l,lstar,alpha)
#now, i make the cumulative integral of lf
cum = cumtrapz(lf,l)
#normalize cum so it goes from 0 to 1
cum = cum/np.max(cum)
#now i make the interpolating function
#note this is backwards to you plug in a number from 0 to 1
#and get back luminosities
finterp = interp1d(cum,l[:-1],bounds_error=False)
r=rand(samples)
gal = finterp(r)
return gal
def get_velocity_displacement(time_step, acceleration, units="cm/s/s",
velocity=None, displacement=None):
'''
Returns the velocity and displacement time series using simple integration.
By providing `velocity` or `displacement` as argument(s), you can speed up
this function by skipping either or both calculations.
:param time_step: float: Time-series time-step (s)
:param acceleration: numpy.ndarray: the acceleration
:param units: the acceleration units, either "m/s/s", "m/s**2", "m/s^2", "g", "cm/s/s",
"cm/s**2", "cm/s^2". The acceleration is supposed to
be in centimeters over seconds square: if 'units' is not one of the last three strings,
it will be conveted to cm/s^2 before calculation
:param velocity: numpt.ndarray or None: if None, the velocity will be computed. Otherwise it
is the already-computed vector of velocities and it will be returned by this function
:param displacement: numpt.ndarray or None: if None, the displacement will be computed.
Otherwise it is the already-computed vector of displacements and it will be returned by this
function
:returns:
velocity - Velocity Time series (cm/s)
displacement - Displacement Time series (cm)
'''
acceleration = ResponseSpectrum.convert_accel_units(acceleration, units)
if velocity is None:
velocity = time_step * cumtrapz(acceleration, initial=0.)
if displacement is None:
displacement = time_step * cumtrapz(velocity, initial=0.)
return velocity, displacement
def ocn_heat_transport(self, dlat=1, grid='g16', method='Flux_adjusted', lat_bd=90):
from scipy import integrate
flux = self._obj
area = dict(g16=utl.tarea_g16, g35=utl.tarea_g35, g37=utl.tarea_g37)
flux_area = flux.copy()
flux_area.values = flux * area[grid] * 1e-4 # convert to m2
lat_bins = np.arange(-90,91,dlat)
lat = np.arange(-89.5,90,dlat)
if 'TLAT' in flux_area.coords.keys():
flux_lat = flux_area.groupby_bins('TLAT', lat_bins, labels = lat).sum('stacked_nlat_nlon')
latax = flux_lat.get_axis_num('TLAT_bins')
elif 'ULAT' in flux_area.coords.keys():
flux_area = flux_area.rename({"ULAT":"TLAT"})
flux_lat = flux_area.groupby_bins('TLAT', lat_bins, labels = lat).sum('stacked_nlat_nlon')
latax = flux_lat.get_axis_num('TLAT_bins')
TLAT_bins = flux_lat.TLAT_bins
if method == "Flux_adjusted":
flux_lat = flux_lat.where(TLAT_bins < lat_bd) # north bound
flat_ave = flux_lat.mean('TLAT_bins')
flux_lat.values = flux_lat - flat_ave # remove bias
flux_lat = flux_lat.fillna(0)
print("The ocean heat trasnport is computed by Flux adjustment.")
elif method == "Flux":
flux_lat = flux_lat.fillna(0)
print("The ocean heat trasnport is computed by original flux.")
else:
raise ValueError("method is not suppoprted.")
flux_lat.values = -np.flip(flux_lat.values, latax) # integrate from north pole
integral = integrate.cumtrapz(flux_lat, x=None, initial=0., axis=latax)
OHT = flux_lat.copy()
OHT["TLAT_bins"] = np.flip(flux_lat.TLAT_bins.values, 0)
OHT.values = integral *1e-15
return OHT
