Python scipy.integrate 模块，cumtrapz() 实例源码

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

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