jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def plot_component(ax: plt.Axes, trajs: TimeTrace,
                   dname: str, fieldname: str, index: int,
                   label=None,
                   do_integrate=False,
                   do_schedule=False,
                   mult=1.0,
                   fix=False):
    """convenience method to plot one axis object
    :param ax: plt.Axes to plot
    :param trajs: csaf simulation output
    :param dname: component name to access
    :param fieldname: name of field to plot
    :param index: index in vector where field is located
    :param label: string containing "plot label (units)"
    :param do_integrate: integrate results (useful for transforming state variable)
    :param do_schedule: use a schedule plot (plot_schedule)
    :param mult: apply a multiplier (simple unit conversion)
    :param fix: center steady state at zero (easier to see features line overshoot, settle time, etc)
    """
    import numpy as np

    # times component
    t = trajs[dname].times
    # variable to plot
    x = np.array(getattr(trajs[dname], fieldname))[:, index]

    if do_schedule:
        plot_schedule(ax, t, x)
    else:
        if do_integrate:
            import scipy.integrate
            x = scipy.integrate.cumulative_trapezoid(x, t)
            t = t[:-1]
        if fix:
            x = x - x[-1]
        # pretty plot for zero order hold (ZOH)
        # times of plant
        t_plant = np.array(trajs["plant"].times)
        x_match = t_plant.copy()
        x_match[:] = x[1]
        idx_max = 0
        for idx, ti in enumerate(t):
            if idx < len(t) - 1:
                idxs = np.logical_and(t_plant >= ti, t_plant < t[idx+1])
                x_match[idxs] = x[idx+1]
                barr = np.where(idxs)[0]
                if len(barr) > 0:
                    idx_max = max(barr)
            else:
                x_match[idx_max:] = x[-1]

        # times now match up with t_plant
        t = t_plant
        ax.plot(t_plant, mult*np.array(x_match), label=label.split('(')[0])
        ax.grid(True)

    ax.set_xlim([min(t), max(t)])
    if label and not do_schedule:
        ax.set_ylabel(re.search(r'\((.*?)\)',label).group(1))
    ax.legend()
//  [python]
def ana_likelihood(rv, leff, fakedata, errors, model, prior, 
                   mylikelihood=mcmc.LogLikelihood, covariance_model_funct=None, otherp=None):
    ll = mylikelihood(xvals=leff, yvals=fakedata, errors=errors, 
                            model = model, flatprior=prior, covariance_model_funct=covariance_model_funct)  
    like = np.zeros_like(rv)
    for i in range(len(rv)):
        like[i] = np.exp(ll([rv[i]]))
    cumint = scipy.integrate.cumtrapz(like, x=rv)
    cumint = cumint / np.max(cumint)
    onesigma = np.interp(0.68, cumint, rv[1:])
    if otherp:
        other = np.interp(otherp, cumint, rv[1:])
        return like, cumint, onesigma, other
    else:
        return like, cumint, onesigma
//  [python]
def ana_likelihood(rv, leff, fakedata, errors, model, prior, 
                   mylikelihood=mcmc.LogLikelihood, covariance_model_funct=None, otherp=None):
    ll = mylikelihood(xvals=leff, yvals=fakedata, errors=errors, 
                            model = model, flatprior=prior, covariance_model_funct=covariance_model_funct)  
    like = np.zeros_like(rv)
    for i in range(len(rv)):
        like[i] = np.exp(ll([rv[i]]))
    cumint = scipy.integrate.cumulative_trapezoid(like, x=rv)
    cumint = cumint / np.max(cumint)
    onesigma = np.interp(0.68, cumint, rv[1:])
    if otherp:
        other = np.interp(otherp, cumint, rv[1:])
        return like, cumint, onesigma, other
    else:
        return like, cumint, onesigma
//  [python]
def measureConfidenceInterval(dist, bins=50, interval=0.68):
    """For single peaked quantities only. You need to figure that out first."""

    counts, edges = np.histogram(dist, bins=bins)
    normcounts = counts / float(len(dist))
    centers = (edges[:-1] + edges[1:]) / 2.
    densityorder = np.argsort(normcounts)[::-1]
    maxloc = centers[densityorder[0]]
    cumprob = integrate.cumtrapz(normcounts[densityorder])
    inInterval = cumprob <= interval
    leftEdge = np.min(centers[densityorder][inInterval])
    rightEdge = np.max(centers[densityorder][inInterval])
    binsize = edges[1:] - edges[:-1]
    density = counts / np.sum(binsize * counts)

    pdf = interpolate.interp1d(
        centers, density, bounds_error=False, fill_value=0., kind='cubic')

    def findBound(x, threshold):
        return pdf(x) - threshold

    def findThreshold(thresh):

        leftBound = optimize.fsolve(findBound, leftEdge, args=(thresh,))
        rightBound = optimize.fsolve(findBound, rightEdge, args=(thresh,))
        probability, err = integrate.quad(pdf, leftBound, rightBound)

        return probability - interval

    threshold = optimize.fsolve(findThreshold, pdf(leftEdge))
    leftEdge = optimize.fsolve(findBound, leftEdge, args=(threshold,))[0]
    rightEdge = optimize.fsolve(findBound, rightEdge, args=(threshold,))[0]
    err = np.array([maxloc - leftEdge, rightEdge - maxloc])

    return maxloc, err
//  [python]
def measureConfidenceInterval(dist, bins=50, interval=0.68):
    """For single peaked quantities only. You need to figure that out first."""

    counts, edges = np.histogram(dist, bins=bins)
    normcounts = counts / float(len(dist))
    centers = (edges[:-1] + edges[1:]) / 2.
    densityorder = np.argsort(normcounts)[::-1]
    maxloc = centers[densityorder[0]]
    cumprob = integrate.cumulative_trapezoid(normcounts[densityorder])
    inInterval = cumprob <= interval
    leftEdge = np.min(centers[densityorder][inInterval])
    rightEdge = np.max(centers[densityorder][inInterval])
    binsize = edges[1:] - edges[:-1]
    density = counts / np.sum(binsize * counts)

    pdf = interpolate.interp1d(
        centers, density, bounds_error=False, fill_value=0., kind='cubic')

    def findBound(x, threshold):
        return pdf(x) - threshold

    def findThreshold(thresh):

        leftBound = optimize.fsolve(findBound, leftEdge, args=(thresh,))
        rightBound = optimize.fsolve(findBound, rightEdge, args=(thresh,))
        probability, err = integrate.quad(pdf, leftBound, rightBound)

        return probability - interval

    threshold = optimize.fsolve(findThreshold, pdf(leftEdge))
    leftEdge = optimize.fsolve(findBound, leftEdge, args=(threshold,))[0]
    rightEdge = optimize.fsolve(findBound, rightEdge, args=(threshold,))[0]
    err = np.array([maxloc - leftEdge, rightEdge - maxloc])

    return maxloc, err
//  [python]
def initialize_points(self):
        """
        Set up the case where the user sent x and p(x) points instead
        of a function for p(x)
        """

        if self.method == "accum":
            # we need the cumulative probability distribution

            if self.cumulative:
                # we are done
                self.xvals = self.xinput
                self.norm = self.pofx[-1]
                self.pcum = self.pofx / self.norm

            else:
                # we must integrate and take a subset of x, since the
                # first value is not defined.

                import scipy.integrate

                pcum = scipy.integrate.cumtrapz(self.pofx, self.xinput)
                self.norm = pcum[-1]
                self.pcum = pcum / self.norm

                # interval is smaller, no integral in first point
                self.xvals = self.xinput[1:]
//  [python]
def initialize_points(self):
        """
        Set up the case where the user sent x and p(x) points instead
        of a function for p(x)
        """

        if self.method == "accum":
            # we need the cumulative probability distribution

            if self.cumulative:
                # we are done
                self.xvals = self.xinput
                self.norm = self.pofx[-1]
                self.pcum = self.pofx / self.norm

            else:
                # we must integrate and take a subset of x, since the
                # first value is not defined.

                import scipy.integrate

                pcum = scipy.integrate.cumulative_trapezoid(self.pofx, self.xinput)
                self.norm = pcum[-1]
                self.pcum = pcum / self.norm

                # interval is smaller, no integral in first point
                self.xvals = self.xinput[1:]
//  [python]
def initialize_func(self):
        """
        Set up the case where the user sent x and p(x) as a function
        """

        if self.method == "accum":
            # we need the cumulative probability distribution

            if self.cumulative:
                self.xvals = self.xinput
                self.norm = self.pofx(self.xinput[-1])
                self.pcum = self.pofx(self.xinput) / self.norm
            else:
                # we must integrate and take a subset of x, since the
                # first value is not defined.

                import scipy.integrate

                pofxvals = self.pofx(self.xinput)

                pcum = scipy.integrate.cumtrapz(pofxvals, self.xinput)
                self.norm = pcum[-1]
                self.pcum = pcum / self.norm

                # interval is smaller, no integral in first point
                self.xvals = self.xinput[1:]
//  [python]
def initialize_func(self):
        """
        Set up the case where the user sent x and p(x) as a function
        """

        if self.method == "accum":
            # we need the cumulative probability distribution

            if self.cumulative:
                self.xvals = self.xinput
                self.norm = self.pofx(self.xinput[-1])
                self.pcum = self.pofx(self.xinput) / self.norm
            else:
                # we must integrate and take a subset of x, since the
                # first value is not defined.

                import scipy.integrate

                pofxvals = self.pofx(self.xinput)

                pcum = scipy.integrate.cumulative_trapezoid(pofxvals, self.xinput)
                self.norm = pcum[-1]
                self.pcum = pcum / self.norm

                # interval is smaller, no integral in first point
                self.xvals = self.xinput[1:]
//  [python]
def RelaunchRays1(self,xp,eikonal,vg,a,L,material):
		'''Use wave data to create a new ray distribution.
		At present azimuthal phase variation is ignored.

		:param numpy.ndarray xp: phase space data with shape (Nb,4,8)
		:param numpy.ndarray eik: eikonal data with shape (Nb,4)
		:param numpy.ndarray vg: group velocity with shape (Nb,4,4)
		:param numpy.ndarray a: vector potential with shape (Nw,Nx,Ny)
		:param float L: length of the wave zone
		:param class material: dispersion object rays are emerging from'''
		# Assume vacuum for now
		wn,xn,yn,ext = self.GetGridInfo()
		a,xn = grid_tools.AddGhostCells(a,xn,1)
		# Work out the eikonal wavevectors indirectly
		# This is equivalent to k = grad(phase), but does not require unwrapping phase
		eps = np.max(np.abs(a))*1e-6
		if xn.shape[0]>1:
			dx = xn[1]-xn[0]
			adxastar = a*np.gradient(np.conj(a),dx,axis=1)
			kx = np.real(0.5j*(adxastar-np.conj(adxastar)))/(eps+np.abs(a))**2
		else:
			dx = 1.0
			kx = np.zeros(a.shape)
		phase = scipy.integrate.cumtrapz(kx,dx=dx,axis=1,initial=0.0)
		phase -= 0.25*kx[:,(0,),:]*dx
		# Rays keep their original frequency and transverse positions.
		# Frequency shifts are still accounted for because amplitude may change.
		rho,phi = self.GetCylCoords(xp)
		impact = np.where(np.logical_and(rho[:,0]<xn[-1],xp[:,0,4]<wn[-1]))[0]
		w = xp[impact,:,4]
		x = rho[impact,:]
		y = phi[impact,:]
		try:
			kr = grid_tools.DataFromGrid(w,x,y,wn,xn,yn,kx)
		except:
			kr = np.zeros(y.shape)
		kr[np.where(kr**2>=w**2)] = 0.0
		xp[impact,:,5] = kr*np.cos(y)
		xp[impact,:,6] = kr*np.sin(y)
		xp[impact,:,7] = np.sqrt(w**2 - xp[impact,:,5]**2 - xp[impact,:,6]**2)
		vg[impact,...] = xp[impact,...,4:8]/xp[impact,...,4:5]
		speed = np.sqrt(vg[impact,:,1]**2 + vg[impact,:,2]**2 + vg[impact,:,3]**2)
		xp[impact,:,0] += L/speed
		xp[impact,:,3] += L
		try:
			eikonal[impact,0] = grid_tools.DataFromGrid(w[:,0],x[:,0],y[:,0],wn,xn,yn,phase)
			eikonal[impact,1] = grid_tools.DataFromGrid(w[:,0],x[:,0],y[:,0],wn,xn,yn,np.abs(a))
		except:
			eikonal[impact,0] *= 1.0
			eikonal[impact,1] *= 1.0
		eikonal[impact,2] = 0.0
		eikonal[impact,3] = 0.0
//  [python]
def RelaunchRays1(self,xp,eikonal,vg,a,L,material):
		'''Use wave data to create a new ray distribution.
		At present azimuthal phase variation is ignored.

		:param numpy.ndarray xp: phase space data with shape (Nb,4,8)
		:param numpy.ndarray eik: eikonal data with shape (Nb,4)
		:param numpy.ndarray vg: group velocity with shape (Nb,4,4)
		:param numpy.ndarray a: vector potential with shape (Nw,Nx,Ny)
		:param float L: length of the wave zone
		:param class material: dispersion object rays are emerging from'''
		# Assume vacuum for now
		wn,xn,yn,ext = self.GetGridInfo()
		a,xn = grid_tools.AddGhostCells(a,xn,1)
		# Work out the eikonal wavevectors indirectly
		# This is equivalent to k = grad(phase), but does not require unwrapping phase
		eps = np.max(np.abs(a))*1e-6
		if xn.shape[0]>1:
			dx = xn[1]-xn[0]
			adxastar = a*np.gradient(np.conj(a),dx,axis=1)
			kx = np.real(0.5j*(adxastar-np.conj(adxastar)))/(eps+np.abs(a))**2
		else:
			dx = 1.0
			kx = np.zeros(a.shape)
		phase = scipy.integrate.cumulative_trapezoid(kx,dx=dx,axis=1,initial=0.0)
		phase -= 0.25*kx[:,(0,),:]*dx
		# Rays keep their original frequency and transverse positions.
		# Frequency shifts are still accounted for because amplitude may change.
		rho,phi = self.GetCylCoords(xp)
		impact = np.where(np.logical_and(rho[:,0]<xn[-1],xp[:,0,4]<wn[-1]))[0]
		w = xp[impact,:,4]
		x = rho[impact,:]
		y = phi[impact,:]
		try:
			kr = grid_tools.DataFromGrid(w,x,y,wn,xn,yn,kx)
		except:
			kr = np.zeros(y.shape)
		kr[np.where(kr**2>=w**2)] = 0.0
		xp[impact,:,5] = kr*np.cos(y)
		xp[impact,:,6] = kr*np.sin(y)
		xp[impact,:,7] = np.sqrt(w**2 - xp[impact,:,5]**2 - xp[impact,:,6]**2)
		vg[impact,...] = xp[impact,...,4:8]/xp[impact,...,4:5]
		speed = np.sqrt(vg[impact,:,1]**2 + vg[impact,:,2]**2 + vg[impact,:,3]**2)
		xp[impact,:,0] += L/speed
		xp[impact,:,3] += L
		try:
			eikonal[impact,0] = grid_tools.DataFromGrid(w[:,0],x[:,0],y[:,0],wn,xn,yn,phase)
			eikonal[impact,1] = grid_tools.DataFromGrid(w[:,0],x[:,0],y[:,0],wn,xn,yn,np.abs(a))
		except:
			eikonal[impact,0] *= 1.0
			eikonal[impact,1] *= 1.0
		eikonal[impact,2] = 0.0
		eikonal[impact,3] = 0.0
//  [python]
def GGP_cdf(GGPmodel, lat, degree):
    """
    Calculates the cumulative distribution functions (CDFs) of inclinations and declinations for a given GGP model and latitude.

    Parameters:
        GGPmodel (dict): Set of coefficients for a specific GGP model.
        lat (float): Latitude of interest in degrees.
        degree (int): Maximum degree of the GGP model.

    Returns:
        tuple: A tuple containing two PchipInterpolator objects representing the CDFs.
            - Icdf: PchipInterpolator function representing the CDF of inclinations for the given latitude.
            - Dcdf: PchipInterpolator function representing the CDF of declinations for the given latitude.
    """
 
    #set up grid for integration
    nI = 180
    nD = 361
    I0 = np.linspace(-90,90,nI)
    D0 = np.linspace(-180,180,nD)
    z0 = np.cos(np.deg2rad(I0))
    
    #integrate WRT to dec to get marginal distribution as function of inc
    pI = np.zeros(nI)
    for j in range(nI):
        di_block = []

        for i in range(nD):
            di_block.append([D0[i],I0[j]])
    
        xyz = np.asarray(pmag.dir2cart(di_block)).T
        p0 = GGP_pdf(GGPmodel,lat,degree,xyz)
        pI[j] = np.trapz(p0, x=np.deg2rad(D0)*z0[j])

    cI = cumtrapz(pI, x=np.deg2rad(I0),initial=0)
    Icdf = PchipInterpolator(np.deg2rad(I0),cI) #inclination marginal distribution
    Iinv = PchipInterpolator(cI,np.deg2rad(I0))
    
    #integrate WRT to inc to get marginal distribution as function of dec
    pD = np.zeros(nD)
    for j in range(nD):
        di_block = []

        for i in range(nI):
            di_block.append([D0[j],I0[i]])
    
        xyz = np.asarray(pmag.dir2cart(di_block)).T
        p0 = GGP_pdf(GGPmodel,lat,degree,xyz)
        pD[j] = np.trapz(p0*np.cos(np.deg2rad(I0)), x=np.deg2rad(I0))

    cD = cumtrapz(pD, x=np.deg2rad(D0),initial=0)
    Dcdf = PchipInterpolator(np.deg2rad(D0),cD) #declination marginal distribution
    Dinv = PchipInterpolator(cD,np.deg2rad(D0))
    
    return Icdf, Dcdf
//  [python]
def GGP_cdf(GGPmodel, lat, degree):
    """
    Calculates the cumulative distribution functions (CDFs) of inclinations and declinations for a given GGP model and latitude.

    Parameters:
        GGPmodel (dict): Set of coefficients for a specific GGP model.
        lat (float): Latitude of interest in degrees.
        degree (int): Maximum degree of the GGP model.

    Returns:
        tuple: A tuple containing two PchipInterpolator objects representing the CDFs.
            - Icdf: PchipInterpolator function representing the CDF of inclinations for the given latitude.
            - Dcdf: PchipInterpolator function representing the CDF of declinations for the given latitude.
    """
 
    #set up grid for integration
    nI = 180
    nD = 361
    I0 = np.linspace(-90,90,nI)
    D0 = np.linspace(-180,180,nD)
    z0 = np.cos(np.deg2rad(I0))
    
    #integrate WRT to dec to get marginal distribution as function of inc
    pI = np.zeros(nI)
    for j in range(nI):
        di_block = []

        for i in range(nD):
            di_block.append([D0[i],I0[j]])
    
        xyz = np.asarray(pmag.dir2cart(di_block)).T
        p0 = GGP_pdf(GGPmodel,lat,degree,xyz)
        pI[j] = np.trapz(p0, x=np.deg2rad(D0)*z0[j])

    cI = cumulative_trapezoid(pI, x=np.deg2rad(I0),initial=0)
    Icdf = PchipInterpolator(np.deg2rad(I0),cI) #inclination marginal distribution
    Iinv = PchipInterpolator(cI,np.deg2rad(I0))
    
    #integrate WRT to inc to get marginal distribution as function of dec
    pD = np.zeros(nD)
    for j in range(nD):
        di_block = []

        for i in range(nI):
            di_block.append([D0[j],I0[i]])
    
        xyz = np.asarray(pmag.dir2cart(di_block)).T
        p0 = GGP_pdf(GGPmodel,lat,degree,xyz)
        pD[j] = np.trapz(p0*np.cos(np.deg2rad(I0)), x=np.deg2rad(I0))

    cD = cumulative_trapezoid(pD, x=np.deg2rad(D0),initial=0)
    Dcdf = PchipInterpolator(np.deg2rad(D0),cD) #declination marginal distribution
    Dinv = PchipInterpolator(cD,np.deg2rad(D0))
    
    return Icdf, Dcdf
//  [python]
def ana_likelihood(self, rv, leff, fakedata, errors, model, prior,mylikelihood=mcmc.LogLikelihood, covariance_model_funct=None):

        ll = mylikelihood(xvals=leff, yvals=fakedata, errors=errors,
                                model = model, flatprior=prior, covariance_model_funct=covariance_model_funct, nbins=len(leff))

        like = np.zeros_like(rv)
        for i in range(len(rv)):
            like[i] = np.exp(ll([rv[i]]))
            #print(rv[i],ll([rv[i]]),like[i])
        maxL = rv[like == np.max(like)]
        cumint = scipy.integrate.cumtrapz(like, x=rv)
        cumint = cumint / np.max(cumint)
        onesigma = np.interp(0.68, cumint, rv[1:])

        return like, cumint, onesigma
//  [python]
def ana_likelihood(self, rv, leff, fakedata, errors, model, prior,mylikelihood=mcmc.LogLikelihood, covariance_model_funct=None):

        ll = mylikelihood(xvals=leff, yvals=fakedata, errors=errors,
                                model = model, flatprior=prior, covariance_model_funct=covariance_model_funct, nbins=len(leff))

        like = np.zeros_like(rv)
        for i in range(len(rv)):
            like[i] = np.exp(ll([rv[i]]))
            #print(rv[i],ll([rv[i]]),like[i])
        maxL = rv[like == np.max(like)]
        cumint = scipy.integrate.cumulative_trapezoid(like, x=rv)
        cumint = cumint / np.max(cumint)
        onesigma = np.interp(0.68, cumint, rv[1:])

        return like, cumint, onesigma
//  [python]
def set_filling(h,filling=0.5,nk=10,extrae=0.,
    mode="ED",**kwargs):
    """
    Set the filling of a Hamiltonian
    - nk = 10, number of kpoints in each direction
    - filling = 0.5, filling of the lattice
    - extrae = 0.0, number of extra electrons
    """
    if h.has_eh: # quick workaround
        ef = h.get_fermi4filling(filling,nk=nk) # fermi energy
        h.add_onsite(-ef)
        return
    fill = filling + extrae/h.intra.shape[0] # filling
    n = h.intra.shape[0]
    use_kpm = False
    if n>algebra.maxsize: # use the KPM method
        mode="KPM"
        print("Using KPM in set_filling")
    if mode=="KPM": # use KPM
        es,ds = h.get_dos(energies=np.linspace(-5.0,5.0,1000),
                mode="KPM",nk=nk,**kwargs)
        from scipy.integrate import cumtrapz
        di = cumtrapz(ds,es)
        ei = (es[0:len(es)-1] + es[1:len(es)])/2.
        di /= di[len(di)-1] # normalize
        from scipy.interpolate import interp1d
        f = interp1d(di,ei) # interpolating function
        efermi = f(fill) # get the fermi energy
    elif mode=="ED": # dense Hamiltonian, use ED
        es = eigenvalues(h,nk=nk,notime=True)
        efermi = get_fermi_energy(es,fill)
    else: raise
    h.shift_fermi(-efermi)
//  [python]
def set_filling(h,filling=0.5,nk=10,extrae=0.,
    mode="ED",**kwargs):
    """
    Set the filling of a Hamiltonian
    - nk = 10, number of kpoints in each direction
    - filling = 0.5, filling of the lattice
    - extrae = 0.0, number of extra electrons
    """
    if h.has_eh: # quick workaround
        ef = h.get_fermi4filling(filling,nk=nk) # fermi energy
        h.add_onsite(-ef)
        return
    fill = filling + extrae/h.intra.shape[0] # filling
    n = h.intra.shape[0]
    use_kpm = False
    if n>algebra.maxsize: # use the KPM method
        mode="KPM"
        print("Using KPM in set_filling")
    if mode=="KPM": # use KPM
        es,ds = h.get_dos(energies=np.linspace(-5.0,5.0,1000),
                mode="KPM",nk=nk,**kwargs)
        from scipy.integrate import cumulative_trapezoid
        di = cumulative_trapezoid(ds,es)
        ei = (es[0:len(es)-1] + es[1:len(es)])/2.
        di /= di[len(di)-1] # normalize
        from scipy.interpolate import interp1d
        f = interp1d(di,ei) # interpolating function
        efermi = f(fill) # get the fermi energy
    elif mode=="ED": # dense Hamiltonian, use ED
        es = eigenvalues(h,nk=nk,notime=True)
        efermi = get_fermi_energy(es,fill)
    else: raise
    h.shift_fermi(-efermi)
//  [python]
def GGP_cdf(GGPmodel, lat, degree):
    """
    Calculates the cumulative distribution functions (CDF) of inclinations and declinations for a given GGP model and latitude.

    Parameters:
        GGPmodel (dict): Set of coefficients for a specific GGP model.
        lat (float): Latitude of interest in degrees.
        degree (int): Maximum degree of the GGP model.

    Returns:
        tuple: A tuple of two PchipInterpolator objects representing the CDF of inclinations (Icdf) and declinations (Dcdf) for the given latitude.
    """

    
    #set up grid for integration
    nI = 180
    nD = 361
    I0 = np.linspace(-90,90,nI)
    D0 = np.linspace(-180,180,nD)
    z0 = np.cos(np.deg2rad(I0))
    
    #integrate WRT to dec to get marginal distribution as function of inc
    pI = np.zeros(nI)
    for j in range(nI):
        di_block = []

        for i in range(nD):
            di_block.append([D0[i],I0[j]])
    
        xyz = np.asarray(pmag.dir2cart(di_block)).T
        p0 = GGP_pdf(GGPmodel,lat,degree,xyz)
        pI[j] = np.trapz(p0, x=np.deg2rad(D0)*z0[j])

    cI = cumtrapz(pI, x=np.deg2rad(I0),initial=0)
    Icdf = PchipInterpolator(np.deg2rad(I0),cI) #inclination marginal distribution
    Iinv = PchipInterpolator(cI,np.deg2rad(I0))
    
    #integrate WRT to inc to get marginal distribution as function of dec
    pD = np.zeros(nD)
    for j in range(nD):
        di_block = []

        for i in range(nI):
            di_block.append([D0[j],I0[i]])
    
        xyz = np.asarray(pmag.dir2cart(di_block)).T
        p0 = GGP_pdf(GGPmodel,lat,degree,xyz)
        pD[j] = np.trapz(p0*np.cos(np.deg2rad(I0)), x=np.deg2rad(I0))

    cD = cumtrapz(pD, x=np.deg2rad(D0),initial=0)
    Dcdf = PchipInterpolator(np.deg2rad(D0),cD) #declination marginal distribution
    Dinv = PchipInterpolator(cD,np.deg2rad(D0))
    
    return Icdf, Dcdf
//  [python]
def GGP_cdf(GGPmodel, lat, degree):
    """
    Calculates the cumulative distribution functions (CDF) of inclinations and declinations for a given GGP model and latitude.

    Parameters:
        GGPmodel (dict): Set of coefficients for a specific GGP model.
        lat (float): Latitude of interest in degrees.
        degree (int): Maximum degree of the GGP model.

    Returns:
        tuple: A tuple of two PchipInterpolator objects representing the CDF of inclinations (Icdf) and declinations (Dcdf) for the given latitude.
    """

    
    #set up grid for integration
    nI = 180
    nD = 361
    I0 = np.linspace(-90,90,nI)
    D0 = np.linspace(-180,180,nD)
    z0 = np.cos(np.deg2rad(I0))
    
    #integrate WRT to dec to get marginal distribution as function of inc
    pI = np.zeros(nI)
    for j in range(nI):
        di_block = []

        for i in range(nD):
            di_block.append([D0[i],I0[j]])
    
        xyz = np.asarray(pmag.dir2cart(di_block)).T
        p0 = GGP_pdf(GGPmodel,lat,degree,xyz)
        pI[j] = np.trapz(p0, x=np.deg2rad(D0)*z0[j])

    cI = scipy.integrate.cumulative_trapezoid(pI, x=np.deg2rad(I0),initial=0)
    Icdf = PchipInterpolator(np.deg2rad(I0),cI) #inclination marginal distribution
    Iinv = PchipInterpolator(cI,np.deg2rad(I0))
    
    #integrate WRT to inc to get marginal distribution as function of dec
    pD = np.zeros(nD)
    for j in range(nD):
        di_block = []

        for i in range(nI):
            di_block.append([D0[j],I0[i]])
    
        xyz = np.asarray(pmag.dir2cart(di_block)).T
        p0 = GGP_pdf(GGPmodel,lat,degree,xyz)
        pD[j] = np.trapz(p0*np.cos(np.deg2rad(I0)), x=np.deg2rad(I0))

    cD = scipy.integrate.cumulative_trapezoid(pD, x=np.deg2rad(D0),initial=0)
    Dcdf = PchipInterpolator(np.deg2rad(D0),cD) #declination marginal distribution
    Dinv = PchipInterpolator(cD,np.deg2rad(D0))
    
    return Icdf, Dcdf
//  [python]
def render(self, params, branchid, solution):
        try:
            os.makedirs('output/figures/%2.6f' % (params[0],))
        except:
            pass

        s = np.linspace(0, 1, 200)
        theta = [solution((s_,)) for s_ in s]
        x = scipy.integrate.cumtrapz(np.cos(theta), s, initial=0)
        y = scipy.integrate.cumtrapz(np.sin(theta), s, initial=0)
        plt.clf()
        plt.plot(x, y, '-b', linewidth=2)
        plt.grid()
        plt.xlabel(r'$s$')
        plt.ylabel(r'elastica')
        plt.title(r'$\lambda = %.3f$' % params[0])
        plt.savefig('output/figures/%2.6f/branchid-%d.png' % (params[0], branchid))
//  [python]
def render(self, params, branchid, solution):
        try:
            os.makedirs('output/figures/%2.6f' % (params[0],))
        except:
            pass

        s = np.linspace(0, 1, 200)
        theta = [solution((s_,)) for s_ in s]
        x = scipy.integrate.cumulative_trapezoid(np.cos(theta), s, initial=0)
        y = scipy.integrate.cumulative_trapezoid(np.sin(theta), s, initial=0)
        plt.clf()
        plt.plot(x, y, '-b', linewidth=2)
        plt.grid()
        plt.xlabel(r'$s$')
        plt.ylabel(r'elastica')
        plt.title(r'$\lambda = %.3f$' % params[0])
        plt.savefig('output/figures/%2.6f/branchid-%d.png' % (params[0], branchid))
//  [python]
def uintegrate(dudz, X, ubase=0.0):
    """ Integrate velocity shear from the first non-NaN value to the top. """
    U = -np.nan*np.empty_like(dudz)
    if isinstance(ubase, numbers.Number):
        ubase = ubase * np.ones(dudz.shape[1], dtype=np.float64)

    for jcol in range(dudz.shape[1]):
        # find the deepest non-NaN
        imax = np.max(np.arange(dudz.shape[0])[~np.isnan(dudz[:,jcol])])
        du = dudz[:imax+1,jcol]
        du[np.isnan(du)] = 0.0
        U[:imax+1,jcol] = scint.cumtrapz(du, x=X[:imax+1,jcol],
                                         initial=0.0)
        U[:imax+1,jcol] -= U[imax,jcol] - ubase[jcol]
    return U
//  [python]
def uintegrate(dudz, X, ubase=0.0):
    """ Integrate velocity shear from the first non-NaN value to the top. """
    U = -np.nan*np.empty_like(dudz)
    if isinstance(ubase, numbers.Number):
        ubase = ubase * np.ones(dudz.shape[1], dtype=np.float64)

    for jcol in range(dudz.shape[1]):
        # find the deepest non-NaN
        imax = np.max(np.arange(dudz.shape[0])[~np.isnan(dudz[:,jcol])])
        du = dudz[:imax+1,jcol]
        du[np.isnan(du)] = 0.0
        U[:imax+1,jcol] = scint.cumulative_trapezoid(du, x=X[:imax+1,jcol],
                                         initial=0.0)
        U[:imax+1,jcol] -= U[imax,jcol] - ubase[jcol]
    return U
//  [python]
def __cal_E_Jsc(self):
        fluxcumm = cumtrapz(AM15flux[::-1], self.Es[::-1], initial=0)
        fluxaboveE = fluxcumm[::-1] * -1 * self.intensity
        Jsc = fluxaboveE * q * 0.1  # mA/cm^2  (0.1: from A/m2 to mA/cm2)
        return Jsc
//  [python]
def __cal_E_Jsc(self):
        fluxcumm = cumulative_trapezoid(AM15flux[::-1], self.Es[::-1], initial=0)
        fluxaboveE = fluxcumm[::-1] * -1 * self.intensity
        Jsc = fluxaboveE * q * 0.1  # mA/cm^2  (0.1: from A/m2 to mA/cm2)
        return Jsc
//  [python]
def compute_emiss(I0, op, dx=1, axis=0, _sum=False):
    if axis==0:
        op = op[::-1]
    elif axis==1:
        op = op[:,::-1]
    #if mask is not None:
    #    # this part of simulation is opaque
    #    op[mask] = 1e3
   

    cop = cumtrapz(op, axis=axis, dx=dx)
    if axis==0:
        cop = np.vstack((np.zeros((1,cop.shape[1])), cop))
    elif axis==1:
        cop = np.hstack((np.zeros((cop.shape[0], 1)), cop))
    if axis==0:
        cop = cop[::-1]
    elif axis==1:
        cop = cop[:,::-1]
    em = I0*np.exp(-cop)
    if _sum:
        return np.sum(em, axis=axis)
    else:
        return em
//  [python]
def compute_emiss(I0, op, dx=1, axis=0, _sum=False):
    if axis==0:
        op = op[::-1]
    elif axis==1:
        op = op[:,::-1]
    #if mask is not None:
    #    # this part of simulation is opaque
    #    op[mask] = 1e3
   

    cop = cumulative_trapezoid(op, axis=axis, dx=dx)
    if axis==0:
        cop = np.vstack((np.zeros((1,cop.shape[1])), cop))
    elif axis==1:
        cop = np.hstack((np.zeros((cop.shape[0], 1)), cop))
    if axis==0:
        cop = cop[::-1]
    elif axis==1:
        cop = cop[:,::-1]
    em = I0*np.exp(-cop)
    if _sum:
        return np.sum(em, axis=axis)
    else:
        return em
//  [python]
def calc_source_trans(self, Esrc, Isrc):
        
        I0 = np.interp(self.source_energy, Esrc, Isrc, left=0, right=0)
         
        arg = spi.cumtrapz(np.array(self.mesh_sam_abscoeff), self.source_ysrc, axis=0, initial=0)
        
        # beer-lambert law: I = I0*exp(-rho.mu.length)
        self.source_I = np.array([I0])*np.exp(-arg)
//  [python]
def calc_source_trans(self, Esrc, Isrc):
        
        I0 = np.interp(self.source_energy, Esrc, Isrc, left=0, right=0)
         
        arg = spi.cumulative_trapezoid(np.array(self.mesh_sam_abscoeff), self.source_ysrc, axis=0, initial=0)
        
        # beer-lambert law: I = I0*exp(-rho.mu.length)
        self.source_I = np.array([I0])*np.exp(-arg)
//  [python]
def __cal_E_J0(self):
        '''
        Calculate and return E vs J0, the dark saturation current
        J0 = q * (integrate(phi dE) from E to infinity)  / EQE_EL
        phi is the black body radiation at T (flux vs energy)

        '''
        phi = 2 * np.pi * (((self.Es*eV)**2) * eV / ((h**3) * (c**2)) / (
                           np.exp(self.Es*eV / (k*self.T)) - 1))

        # fluxcumm = sp.integrate.cumtrapz(phi[::-1], self.Es[::-1], initial=0)
        fluxcumm = cumtrapz(phi[::-1], self.Es[::-1], initial=0)
        fluxaboveE = fluxcumm[::-1] * -1
        J0 = fluxaboveE * q * 0.1 / self.EQE_EL  # (0.1: from A/m2 to mA/cm2)
        J0[-1] = np.nan  # avoid divide by zero when calculating Voc
        return J0
//  [python]
def __cal_E_J0(self):
        '''
        Calculate and return E vs J0, the dark saturation current
        J0 = q * (integrate(phi dE) from E to infinity)  / EQE_EL
        phi is the black body radiation at T (flux vs energy)

        '''
        phi = 2 * np.pi * (((self.Es*eV)**2) * eV / ((h**3) * (c**2)) / (
                           np.exp(self.Es*eV / (k*self.T)) - 1))

        # fluxcumm = sp.integrate.cumtrapz(phi[::-1], self.Es[::-1], initial=0)
        fluxcumm = cumulative_trapezoid(phi[::-1], self.Es[::-1], initial=0)
        fluxaboveE = fluxcumm[::-1] * -1
        J0 = fluxaboveE * q * 0.1 / self.EQE_EL  # (0.1: from A/m2 to mA/cm2)
        J0[-1] = np.nan  # avoid divide by zero when calculating Voc
        return J0
//  [python]
def _compute_gravity(self):
        # Calculate the gravity of the planet, based on a density profile.
        # There is a check of the surface value is within reason for this
        # model, otherwise values for PREM are used.

        density = self.table_density[::-1]
        radii = self.table_radius[::-1]
        g = scipy.integrate.cumtrapz(constants.G*4.*np.pi*density*radii*radii, x=radii, initial=0)
        g[1:] = g[1:] / radii[1:] / radii[1:]

        self.table_gravity = g[::-1]
//  [python]
def _compute_gravity(self):
        # Calculate the gravity of the planet, based on a density profile.
        # There is a check of the surface value is within reason for this
        # model, otherwise values for PREM are used.

        density = self.table_density[::-1]
        radii = self.table_radius[::-1]
        g = scipy.integrate.cumulative_trapezoid(constants.G*4.*np.pi*density*radii*radii, x=radii, initial=0)
        g[1:] = g[1:] / radii[1:] / radii[1:]

        self.table_gravity = g[::-1]
//  [python]
def _compute_pressure(self):
        # Calculate the pressure profile based on density and gravity.  This integrates
        # the equation for hydrostatic equilibrium  P = rho g z.
        radii = self.table_radius
        density = self.table_density
        gravity = self.gravity(self.earth_radius-radii)

        # convert radii to depths
        depth = self.earth_radius - radii
        pressure = scipy.integrate.cumtrapz(gravity*density, x=depth, initial=0.)

        self.table_pressure = pressure
//  [python]
def _compute_pressure(self):
        # Calculate the pressure profile based on density and gravity.  This integrates
        # the equation for hydrostatic equilibrium  P = rho g z.
        radii = self.table_radius
        density = self.table_density
        gravity = self.gravity(self.earth_radius-radii)

        # convert radii to depths
        depth = self.earth_radius - radii
        pressure = scipy.integrate.cumulative_trapezoid(gravity*density, x=depth, initial=0.)

        self.table_pressure = pressure
//  [python]
def integrate(self, **kwargs):
        """
        DESCRIPTION:
            Method to integrate selected columns respect to time.
            -- Using scipy.integrate.cumtrapz
        VARIABLES:
            optional:
            keys: (list - default ['x','y','z','f'] provide limited key-list
        """


        logger.info('--- Integrating started at %s ' % str(datetime.now()))

        keys = kwargs.get('keys')
        if not keys:
            keys = ['x','y','z']

        array = [[] for key in KEYLIST]
        ndtype = False
        if len(self.ndarray[0])>0:
            ndtype = True
            t = self.ndarray[0]
            array[0] = t
        else:
            t = self._get_column('time')
        for key in keys:
            if ndtype:
                ind = KEYLIST.index(key)
                val = self.ndarray[ind]
                array[ind] = np.asarray(val)
            else:
                val = self._get_column(key)
            dval = sp.integrate.cumtrapz(np.asarray(val),t)
            dval = np.insert(dval, 0, 0) # Prepend 0 to maintain original length
            if ndtype:
                ind = KEYLIST.index('d'+key)
                array[ind] = np.asarray(dval)
            else:
                self._put_column(dval, 'd'+key)

        self.ndarray = np.asarray(array)
        logger.info('--- integration finished at %s ' % str(datetime.now()))
        return self
//  [python]
def integrate(self, **kwargs):
        """
        DESCRIPTION:
            Method to integrate selected columns respect to time.
            -- Using scipy.integrate.cumtrapz
        VARIABLES:
            optional:
            keys: (list - default ['x','y','z','f'] provide limited key-list
        """


        logger.info('--- Integrating started at %s ' % str(datetime.now()))

        keys = kwargs.get('keys')
        if not keys:
            keys = ['x','y','z']

        array = [[] for key in KEYLIST]
        ndtype = False
        if len(self.ndarray[0])>0:
            ndtype = True
            t = self.ndarray[0]
            array[0] = t
        else:
            t = self._get_column('time')
        for key in keys:
            if ndtype:
                ind = KEYLIST.index(key)
                val = self.ndarray[ind]
                array[ind] = np.asarray(val)
            else:
                val = self._get_column(key)
            dval = sp.integrate.cumulative_trapezoid(np.asarray(val),t)
            dval = np.insert(dval, 0, 0) # Prepend 0 to maintain original length
            if ndtype:
                ind = KEYLIST.index('d'+key)
                array[ind] = np.asarray(dval)
            else:
                self._put_column(dval, 'd'+key)

        self.ndarray = np.asarray(array)
        logger.info('--- integration finished at %s ' % str(datetime.now()))
        return self
//  [python]
def get_confidence_HDI(post, grid, kde=None, interval=0.683, ax=None, color="k", lab="auto", median=False, ls='--'):

    if kde is not None:
        N_samples = 10**6
        post /= np.sum(post)
        samples   = np.random.choice(grid, size=N_samples, replace=True, p=post)
        med    = np.median(samples)
        lo, hi = np.percentile(samples, 16), np.percentile(samples, 84)
        post      = gaussian_kde(samples,bw_method=kde).evaluate(grid)
    
    if kde is None and median:
        raise ValueError("kde must be provided to compute the median. Too lazy to write other code.")
    
    # Higher resolution grid
    xx       = np.linspace(grid[0], grid[-1], 10**6)
    post_int = interp1d(grid, post, kind="cubic", bounds_error=False)
    yy       = post_int(xx)

    # Renormalize
    norm     = np.trapz(yy, xx)
    yy       = yy/norm

    # Compute CDF and inverse CDF
    cdf = cumtrapz(yy, xx)
    ppf = interp1d(cdf[cdf > 0.], xx[1:][cdf > 0.], fill_value=0., bounds_error=False)

    def width(low):
        w = ppf(interval + low) - ppf(low)
        return w if w > 0. else np.inf

    # Find the interval
    HDI_low = fmin(width, 1. - interval, disp=False)[0]

    if not median:
        xmax = xx[np.argmax(yy)]
        xlow = ppf(HDI_low)
        xup  = ppf(HDI_low + interval)
    else:
        xmax = med
        xlow = lo
        xup  = hi 
    
    
    if ax is not None:
        if lab == "auto":
            lab = r"$H_0={:.0f}^{{+{:.0f}}}_{{-{:.0f}}}$".format(xmax, xup-xmax, xmax-xlow)

        ax.fill_between(xx, 0, yy, where=(xx >= xlow) & (xx <= xup), color=color, alpha=0.2, linewidth=0.0)
        ax.plot([xmax,xmax], [0, float(post_int(xmax))/norm], color=color, linestyle=ls, label=lab)
        ax.plot(xx, yy, lw=1, color=color, ls="-")
    
    return xmax, xlow, xup, post_int
//  [python]
def get_confidence_HDI(post, grid, kde=None, interval=0.683, ax=None, color="k", lab="auto", median=False, ls='--'):

    if kde is not None:
        N_samples = 10**6
        post /= np.sum(post)
        samples   = np.random.choice(grid, size=N_samples, replace=True, p=post)
        med    = np.median(samples)
        lo, hi = np.percentile(samples, 16), np.percentile(samples, 84)
        post      = gaussian_kde(samples,bw_method=kde).evaluate(grid)
    
    if kde is None and median:
        raise ValueError("kde must be provided to compute the median. Too lazy to write other code.")
    
    # Higher resolution grid
    xx       = np.linspace(grid[0], grid[-1], 10**6)
    post_int = interp1d(grid, post, kind="cubic", bounds_error=False)
    yy       = post_int(xx)

    # Renormalize
    norm     = np.trapz(yy, xx)
    yy       = yy/norm

    # Compute CDF and inverse CDF
    cdf = cumulative_trapezoid(yy, xx)
    ppf = interp1d(cdf[cdf > 0.], xx[1:][cdf > 0.], fill_value=0., bounds_error=False)

    def width(low):
        w = ppf(interval + low) - ppf(low)
        return w if w > 0. else np.inf

    # Find the interval
    HDI_low = fmin(width, 1. - interval, disp=False)[0]

    if not median:
        xmax = xx[np.argmax(yy)]
        xlow = ppf(HDI_low)
        xup  = ppf(HDI_low + interval)
    else:
        xmax = med
        xlow = lo
        xup  = hi 
    
    
    if ax is not None:
        if lab == "auto":
            lab = r"$H_0={:.0f}^{{+{:.0f}}}_{{-{:.0f}}}$".format(xmax, xup-xmax, xmax-xlow)

        ax.fill_between(xx, 0, yy, where=(xx >= xlow) & (xx <= xup), color=color, alpha=0.2, linewidth=0.0)
        ax.plot([xmax,xmax], [0, float(post_int(xmax))/norm], color=color, linestyle=ls, label=lab)
        ax.plot(xx, yy, lw=1, color=color, ls="-")
    
    return xmax, xlow, xup, post_int
//  [python]
def PTV2del(Presi,Tempi,Vpri,hgt,cdict,verbose=False):
    '''Computes the delay function given Pressure, Temperature and Vapor pressure.

    Args:
        * Presi (np.ndarray) : Pressure at height levels.
        * Tempi (np.ndarray) : Temperature at height levels.
        * Vpri  (np.ndarray) : Vapor pressure at height levels.
        * hgt   (np.ndarray) : Height levels.
        * cdict (np.ndarray) : Dictionary of constants.

    Returns:
        * DDry2 (np.ndarray) : Dry component of atmospheric delay.
        * DWet2 (np.ndarray) : Wet component of atmospheric delay.

    .. note::
        Computes refractive index at each altitude and integrates the delay using cumtrapz.'''

    if verbose:
        print('PROGRESS: COMPUTING DELAY FUNCTIONS')
    nhgt = len(hgt)              #Number of height points
    nlat = Presi.shape[0]        #Number of stations
    nlon = Presi.shape[1]
    WonT = Vpri/Tempi
    WonT2 = WonT/Tempi

    k1 = cdict['k1']
    Rd = cdict['Rd']
    Rv = cdict['Rv']
    k2 = cdict['k2']
    k3 = cdict['k3']
    g = cdict['g']

    #Dry delay
    DDry2 = np.zeros((nlat,nlon,nhgt))
    DDry2[:,:,:] = k1*Rd*(Presi[:,:,:] - Presi[:,:,-1][:,:,np.newaxis])*1.0e-6/g

    #Wet delay
    S1 = intg.cumtrapz(WonT,x=hgt,axis=-1)
    val = 2*S1[:,:,-1]-S1[:,:,-2]
    val = val[:,:,None]
    S1 = np.concatenate((S1,val),axis=-1)
    del WonT

    S2 = intg.cumtrapz(WonT2,x=hgt,axis=-1)
    val = 2*S2[:,:,-1]-S2[:,:,-2]
    val = val[:,:,None]
    S2 = np.concatenate((S2,val),axis=-1)
    DWet2 = -1.0e-6*((k2-k1*Rd/Rv)*S1+k3*S2)
    
    for i in range(nlat):
        for j in range(nlon):
            DWet2[i,j,:]  = DWet2[i,j,:] - DWet2[i,j,-1]

    return DDry2,DWet2
//  [python]
def PTV2del(Presi,Tempi,Vpri,hgt,cdict,verbose=False):
    '''Computes the delay function given Pressure, Temperature and Vapor pressure.

    Args:
        * Presi (np.ndarray) : Pressure at height levels.
        * Tempi (np.ndarray) : Temperature at height levels.
        * Vpri  (np.ndarray) : Vapor pressure at height levels.
        * hgt   (np.ndarray) : Height levels.
        * cdict (np.ndarray) : Dictionary of constants.

    Returns:
        * DDry2 (np.ndarray) : Dry component of atmospheric delay.
        * DWet2 (np.ndarray) : Wet component of atmospheric delay.

    .. note::
        Computes refractive index at each altitude and integrates the delay using cumtrapz.'''

    if verbose:
        print('PROGRESS: COMPUTING DELAY FUNCTIONS')
    nhgt = len(hgt)              #Number of height points
    nlat = Presi.shape[0]        #Number of stations
    nlon = Presi.shape[1]
    WonT = Vpri/Tempi
    WonT2 = WonT/Tempi

    k1 = cdict['k1']
    Rd = cdict['Rd']
    Rv = cdict['Rv']
    k2 = cdict['k2']
    k3 = cdict['k3']
    g = cdict['g']

    #Dry delay
    DDry2 = np.zeros((nlat,nlon,nhgt))
    DDry2[:,:,:] = k1*Rd*(Presi[:,:,:] - Presi[:,:,-1][:,:,np.newaxis])*1.0e-6/g

    #Wet delay
    S1 = intg.cumulative_trapezoid(WonT,x=hgt,axis=-1)
    val = 2*S1[:,:,-1]-S1[:,:,-2]
    val = val[:,:,None]
    S1 = np.concatenate((S1,val),axis=-1)
    del WonT

    S2 = intg.cumulative_trapezoid(WonT2,x=hgt,axis=-1)
    val = 2*S2[:,:,-1]-S2[:,:,-2]
    val = val[:,:,None]
    S2 = np.concatenate((S2,val),axis=-1)
    DWet2 = -1.0e-6*((k2-k1*Rd/Rv)*S1+k3*S2)
    
    for i in range(nlat):
        for j in range(nlon):
            DWet2[i,j,:]  = DWet2[i,j,:] - DWet2[i,j,-1]

    return DDry2,DWet2
//  [python]
def integrate_ps(simfolder, f0, natoms, pi, pf, nsims=1, 
    return_values=False):
    """
    Carry out the reversible scaling integration
    
    Parameters
    ----------
    simfolder : string
        main simulation folder
    f0 : float
        initial free energy for integration
    nsims : int, optional
        number of independent switching
    
    Returns
    -------
    None
    
    Notes
    -----
    Writes the output in a file pressure_sweep.dat
    
    """

    ws = []

    for i in range(1, nsims+1):
        _, fp, fvol, _ = np.loadtxt(os.path.join(simfolder, "ps.forward_%d.dat"%i), unpack=True, comments="#")
        _, bp, bvol, _ = np.loadtxt(os.path.join(simfolder, "ps.backward_%d.dat"%i), unpack=True, comments="#")
        
        fvol = fvol/natoms
        bvol = bvol/natoms
        
        fp = fp/(10000*160.21766208)
        bp = bp/(10000*160.21766208)

        wf = cumtrapz(fvol, fp, initial=0)
        wb = cumtrapz(bvol[::-1], bp[::-1], initial=0)

        w = (wf + wb)/2
        ws.append(w)

    wmean = np.mean(ws, axis=0)
    werr = np.std(ws, axis=0)
    
    press = np.linspace(pi, pf, len(wmean))

    f = f0 + wmean

    if not return_values:
        outfile = os.path.join(simfolder, "pressure_sweep.dat")
        np.savetxt(outfile, np.column_stack((press, f, werr)))
    else:
        return (press, f, werr)
//  [python]
def integrate_ps(simfolder, f0, natoms, pi, pf, nsims=1, 
    return_values=False):
    """
    Carry out the reversible scaling integration
    
    Parameters
    ----------
    simfolder : string
        main simulation folder
    f0 : float
        initial free energy for integration
    nsims : int, optional
        number of independent switching
    
    Returns
    -------
    None
    
    Notes
    -----
    Writes the output in a file pressure_sweep.dat
    
    """

    ws = []

    for i in range(1, nsims+1):
        _, fp, fvol, _ = np.loadtxt(os.path.join(simfolder, "ps.forward_%d.dat"%i), unpack=True, comments="#")
        _, bp, bvol, _ = np.loadtxt(os.path.join(simfolder, "ps.backward_%d.dat"%i), unpack=True, comments="#")
        
        fvol = fvol/natoms
        bvol = bvol/natoms
        
        fp = fp/(10000*160.21766208)
        bp = bp/(10000*160.21766208)

        wf = scipy.integrate.cumulative_trapezoid(fvol, fp, initial=0)
        wb = scipy.integrate.cumulative_trapezoid(bvol[::-1], bp[::-1], initial=0)

        w = (wf + wb)/2
        ws.append(w)

    wmean = np.mean(ws, axis=0)
    werr = np.std(ws, axis=0)
    
    press = np.linspace(pi, pf, len(wmean))

    f = f0 + wmean

    if not return_values:
        outfile = os.path.join(simfolder, "pressure_sweep.dat")
        np.savetxt(outfile, np.column_stack((press, f, werr)))
    else:
        return (press, f, werr)