Пример #1
0
def main():

    nstep = 500
    imode = 4
    v0 = np.array([1.0])
    fRHS = dydx_string
    fLOA = load_string
    fSCO = score_string
    fORD = rk4
    x, y = bvp_root(fRHS, fORD, fLOA, fSCO, v0, nstep, imode=imode)

    u = y[0, :]
    l = y[2, :]

    ftsz = 10
    plt.figure(num=1, figsize=(8, 8), dpi=100, facecolor='white')
    plt.subplot(211)
    plt.plot(x, u, linestyle='-', color='black', linewidth=1.0)
    plt.xlabel('x', fontsize=ftsz)
    plt.ylabel('u', fontsize=ftsz)
    util.rescaleplot(x, u, plt, 0.05)
    plt.tick_params(labelsize=ftsz)
    plt.subplot(212)
    plt.plot(x, l, linestyle='-', color='black', linewidth=1.0)
    plt.xlabel('x', fontsize=ftsz)
    plt.ylabel('$l$', fontsize=ftsz)
    util.rescaleplot(x, l, plt, 0.05)
    plt.tick_params(labelsize=ftsz)

    plt.show()
Пример #2
0
def main():
    parser = argparse.ArgumentParser(formatter_class=RawTextHelpFormatter)
    parser.add_argument("J",type=int,
                        help="number of support points")
    parser.add_argument("problem",type=str,
                        help="density field:\n"
                             "    point   : point charge")
    parser.add_argument("-j","--Jmax",type=int,default=None,help="maximum J")
    args        = parser.parse_args()
    J           = args.J
    problem     = args.problem
    Jmax        = args.Jmax

    

    if (Jmax):
        nx = int(np.log(Jmax//J)/np.log(2)+1)
    else:
        nx = 1
    x          = np.zeros(nx,dtype=object)
    rho        = np.zeros(nx,dtype=object)
    dx         = np.zeros(nx)
    phi        = np.zeros(nx,dtype=object)
    rhomin     = np.zeros(nx)
    rhomax     = np.zeros(nx)
    phimin     = np.zeros(nx)
    phimax     = np.zeros(nx)
    for i in range(nx):
        x[i],rho[i],dx[i] = init(problem,J*int(2**i)) 
        phi[i]            = poissonfft(rho[i],dx[i])
        rhomin[i]         = np.min(rho[i])
        rhomax[i]         = np.max(rho[i])
        phimin[i]         = np.min(phi[i])
        phimax[i]         = np.max(phi[i])
    rhomin     = np.min(rhomin)
    rhomax     = np.max(rhomax)
    phimin     = np.min(phimin)
    phimax     = np.max(phimax)

    ftsz       = 10
    plt.figure(num=1,figsize=(6,6),dpi=100,facecolor='white')
    plt.subplot(2,1,1)
    for i in range(nx):
        plt.plot(x[i],rho[i],linestyle='-',label=("J=%4i" % (J*2**i)))
    plt.xlabel('x')
    plt.ylabel(r"$\rho(x)$")
    plt.legend(fontsize=ftsz)
    plt.tick_params(labelsize=ftsz)
    util.rescaleplot(x[0],rho[0],plt,0.05,ymin=rhomin,ymax=rhomax)
    plt.subplot(2,1,2)
    for i in range(nx):
        plt.plot(x[i],phi[i],linestyle='-',label=("J=%4i" % (J*2**i)))
    plt.xlabel('x')
    plt.ylabel(r"$\Phi(x)$")
    plt.legend(fontsize=ftsz)
    plt.tick_params(labelsize=ftsz)
    util.rescaleplot(x[0],phi[0],plt,0.05,ymin=phimin,ymax=phimax)
    plt.tight_layout()
    plt.show()
Пример #3
0
def ode_check(x, y, it):

    n = x.size
    par = globalvar.get_odepar()
    z = y[0, :]  # altitude above ground
    vz = y[1, :]  # vertical velocity
    f = y[2, :]  # fuel mass
    thr = y[3, :]  # throttle

    accG = np.zeros(n)  # acceleration in units of g
    for k in range(n):
        accG[k] = (
            (dydx.lunarlanding(x[k], y[:, k], n /
                               (x[n - 1] - x[0])))[1]) / 9.81  # acceleration

    ftsz = 10
    plt.figure(num=1, figsize=(8, 8), dpi=100, facecolor='white')

    plt.subplot(321)
    plt.plot(x, z, linestyle='-', color='black', linewidth=1.0)
    plt.xlabel('t [s]', fontsize=ftsz)
    plt.ylabel('z(t) [m]', fontsize=ftsz)
    util.rescaleplot(x, z, plt, 0.05)
    plt.tick_params(labelsize=ftsz)

    plt.subplot(322)
    plt.plot(x, vz, linestyle='-', color='black', linewidth=1.0)
    plt.xlabel('t [s]', fontsize=ftsz)
    plt.ylabel('v$_z$ [m s$^{-1}$]', fontsize=ftsz)
    util.rescaleplot(x, vz, plt, 0.05)
    plt.tick_params(labelsize=ftsz)

    plt.subplot(323)
    plt.plot(x, f, linestyle='-', color='black', linewidth=1.0)
    plt.xlabel('t [s]', fontsize=ftsz)
    plt.ylabel('fuel [kg]', fontsize=ftsz)
    util.rescaleplot(x, f, plt, 0.05)
    plt.tick_params(labelsize=ftsz)

    plt.subplot(324)
    plt.plot(x, thr, linestyle='-', color='black', linewidth=1.0)
    plt.xlabel('t [s]', fontsize=ftsz)
    plt.ylabel('throttle', fontsize=ftsz)
    util.rescaleplot(x, thr, plt, 0.05)
    plt.tick_params(labelsize=ftsz)

    plt.subplot(325)
    plt.plot(x, accG, linestyle='-', color='black', linewidth=1.0)
    plt.xlabel('t [s]', fontsize=ftsz)
    plt.ylabel('acc/G', fontsize=ftsz)
    util.rescaleplot(x, accG, plt, 0.05)
    plt.tick_params(labelsize=ftsz)

    plt.tight_layout()

    plt.show()
Пример #4
0
def main():

    #   the function should accomplish the following:
    #   Test the 3 fixed stepsize integrators euler,rk2,rk4 by calculating their
    #   cumulative integration errors for increasing step sizes
    #   on the function y' = exp(y(x)-2*x)+2. This is given in the function "get_dydx" above.
    #   Use the integration interval x=[0,1], and the initial
    #   condition y[0] = -ln(2).
    #   (1) define an array containing the number of steps you want to test. Logarithmic spacing
    #       (in decades) might be useful.
    #   (2) loop through the integrators and the step numbers, calculate the
    #       integral and store the error. You'll need the analytical solution at x=1,
    #       see homework assignment.
    #   (3) Plot the errors against the stepsize as log-log plot, and print out the slopes.

    fINT = odeint.ode_ivp  # use the initial-value problem driver
    fORD = [step.euler, step.rk2,
            step.rk4]  # list of stepper functions to be run
    fRHS = get_dydx  # the RHS (derivative) for our test ODE
    fBVP = 0  # unused for this problem

    #????????????????? from here
    x0 = 0.0
    x1 = 1.0
    y0 = np.array([-np.log(2.0)])
    nstep = np.array([10, 30, 100, 300, 1000, 3000, 10000, 30000, 100000])
    err = np.zeros((nstep.size, len(fORD)))
    for j in range(len(fORD)):
        for i in range(nstep.size):
            x, y, it = fINT(fRHS, fORD[j], fBVP, x0, y0, x1, nstep[i])
            err[i, j] = np.abs(y[0, nstep[i]] - 2.0)
            print('[error_test]: j=%3d i=%3d f=%10s nstep=%7d err=%13.5e' %
                  (j, i, fORD[j].__name__, nstep[i], err[i, j]))

    lnstep = np.log10(nstep)
    lerr = np.log10(err)

    ftsz = 10
    plt.figure(num=1, figsize=(8, 8), dpi=100, facecolor='white')

    #   plot the error against step number, and print the slopes.
    plt.subplot(111)
    for j in range(len(fORD)):
        plt.plot(lnstep, lerr[:, j], linestyle='-', linewidth=1.0)
        print('[error_test]: j=%3d f=%10s slope=%10.2e' %
              (j, fORD[j].__name__,
               (lerr[1, j] - lerr[0, j]) / (lnstep[1] - lnstep[0])))
    plt.xlabel('log steps', fontsize=ftsz)
    plt.ylabel('log err', fontsize=ftsz)
    util.rescaleplot(lnstep, lerr, plt, 0.05)
    plt.tick_params(labelsize=ftsz)

    plt.show()
Пример #5
0
def main():
    parser = argparse.ArgumentParser(formatter_class=RawTextHelpFormatter)
    parser.add_argument("J", type=int, help="number of support points")
    parser.add_argument("problem",
                        type=str,
                        help="data field:\n"
                        "    sine   : well, sine")
    parser.add_argument("k", type=int, help="wave number")

    args = parser.parse_args()
    J = args.J
    problem = args.problem
    k = args.k

    x, f = init(problem, J, k)

    #???? here, you'll need to call your Fourier transforms
    # F = sft(f, -1)
    # ff = sft(F, 1)
    F = sft(f, 1)
    ff = sft(F, -1)
    F = abs(F)
    #?????

    ftsz = 10
    plt.figure(num=1, figsize=(6, 6), dpi=100, facecolor='white')
    plt.subplot(121)
    plt.plot(x, f, linestyle='-', linewidth=1, color='black', label='$f(x)$')
    plt.plot(x,
             ff,
             linestyle='-',
             linewidth=1,
             color='red',
             label='$F^{-1}[F[f]]$')
    plt.xlabel('x')
    plt.ylabel('f(x)')
    plt.legend(fontsize=ftsz)
    plt.tick_params(labelsize=ftsz)
    util.rescaleplot(x, f, plt, 0.05)
    plt.subplot(122)
    plt.plot(np.arange(J), F, 'o', color='black')
    plt.xlabel('k')
    plt.ylabel('F[f]')
    plt.tick_params(labelsize=ftsz)
    util.rescaleplot(np.arange(J), F, plt, 0.05)
    plt.show()
Пример #6
0
def lunarlander(tthron, **kwargs):
    nstep = 100
    x0 = 0.0
    x1 = 20.0
    y0 = np.array([5e2, -5.0, 1e2, 0.0])
    x, y, it = ode_bvp(dydx_lunarlanding,
                       step_rk4,
                       bvp_lunarlander,
                       x0,
                       y0,
                       x1,
                       nstep,
                       thr=tthron)
    fused = y[2, 0] - y[2, nstep - 1]
    vzend = y[1, nstep - 1]
    cst = np.sqrt(np.power(fused, 2) + np.power(vzend, 2))
    iplot = 0
    for key in kwargs:
        if (key == 'plot'):
            iplot = kwargs[key]
    if (iplot == 1):
        ftsz = 10
        plt.figure(num=2, figsize=(8, 8), dpi=100, facecolor='white')

        plt.subplot(221)
        plt.plot(x, y[0, :], linestyle='-', color='black', linewidth=1.0)
        plt.xlabel('t [s]', fontsize=ftsz)
        plt.ylabel('z(t) [m]', fontsize=ftsz)
        util.rescaleplot(x, y[0, :], plt, 0.05)
        plt.tick_params(labelsize=ftsz)

        plt.subplot(222)
        plt.plot(x, y[1, :], linestyle='-', color='black', linewidth=1.0)
        plt.xlabel('t [s]', fontsize=ftsz)
        plt.ylabel('v$_z$ [m s$^{-1}$]', fontsize=ftsz)
        util.rescaleplot(x, y[1, :], plt, 0.05)
        plt.tick_params(labelsize=ftsz)

        plt.subplot(223)
        plt.plot(x, y[2, :], linestyle='-', color='black', linewidth=1.0)
        plt.xlabel('t [s]', fontsize=ftsz)
        plt.ylabel('fuel [kg]', fontsize=ftsz)
        util.rescaleplot(x, y[2, :], plt, 0.05)
        plt.tick_params(labelsize=ftsz)

        plt.subplot(224)
        plt.plot(x, y[3, :], linestyle='-', color='black', linewidth=1.0)
        plt.xlabel('t [s]', fontsize=ftsz)
        plt.ylabel('throttle', fontsize=ftsz)
        util.rescaleplot(x, y[3, :], plt, 0.05)
        plt.tick_params(labelsize=ftsz)

    return cst
Пример #7
0
def ode_check(x, y, it):

    color = ['black', 'blue', 'red']
    label = [u"H", u"H\u2082", u"H\u207A"]
    tlabel = ['k1', 'k2', 'k3', 'k4']
    Myr = 1e6 * 3.6e3 * 3.65e2 * 2.4e1
    nspec = (y.shape)[0]
    n = x.size
    par = get_odepar()
    Ds = par[4]
    ys = get_h2eq(np.log10(Ds))
    print(
        'Equilibrium fractions: D=%13.5e A=%13.5e B=%13.5e C=%13.5e total=%13.5e'
        % (Ds, ys[0], ys[1], ys[2], np.sum(ys)))
    ys = ys * Ds
    th2 = np.zeros((4, n))
    for i in range(n):
        th2[:, i] = get_h2times(y[:, i])
        th2[2, :] = 0.0  # don't show CR ionization (really slow)
    t = x / Myr  # convert time in seconds to Myr
    tmin = np.min(t)
    tmax = np.max(t)
    nmin = np.min(np.array([np.nanmin(y), np.min(ys)]))
    nmax = np.max(np.array([np.nanmax(y), np.max(ys)]))
    ly = np.log10(y)
    lnmin = np.min(np.array([np.nanmin(ly), np.min(np.log10(ys))]))
    lnmax = np.max(np.array([np.nanmax(ly), np.max(np.log10(ys))]))

    # Note that the first call to subplot is slightly different, because we need the
    # axes object to reset the y-labels.
    ftsz = 10
    fig = plt.figure(num=1, figsize=(8, 8), dpi=100, facecolor='white')
    ax = fig.add_subplot(321)
    for i in range(nspec):
        plt.plot(t,
                 y[i, :],
                 linestyle='-',
                 color=color[i],
                 linewidth=1.0,
                 label=label[i])
        plt.plot([tmin, tmax], [ys[i], ys[i]],
                 linestyle='--',
                 color=color[i],
                 linewidth=1.0)
    plt.xlabel('t [Myr]', fontsize=ftsz)
    plt.ylabel('n$_\mathrm{\mathsf{H}}$', fontsize=ftsz)
    plt.legend(fontsize=10)
    util.rescaleplot(t, y, plt, 0.05)
    ylabels = ax.yaxis.get_ticklocs()
    ylabels1 = ylabels[1:ylabels.size - 1]
    ylabelstext = ['%4.2f' % (lb / np.max(ylabels1)) for lb in ylabels1]
    plt.yticks(ylabels1, ylabelstext)
    plt.tick_params(labelsize=ftsz)

    plt.subplot(322)
    for i in range(nspec):
        plt.plot(t,
                 ly[i, :],
                 linestyle='-',
                 color=color[i],
                 linewidth=1.0,
                 label=label[i])
        plt.plot(np.array([tmin, tmax]),
                 np.log10(np.array([ys[i], ys[i]])),
                 linestyle='--',
                 color=color[i],
                 linewidth=1.0)
    plt.xlabel('t [Myr]', fontsize=ftsz)
    plt.ylabel('log n [cm$^{-3}$]', fontsize=ftsz)
    util.rescaleplot(t, ly, plt, 0.05)
    plt.tick_params(labelsize=ftsz)

    plt.subplot(323)
    for i in range(nspec):
        plt.plot(t,
                 y[i, :] / ys[i],
                 linestyle='-',
                 color=color[i],
                 linewidth=1.0,
                 label=label[i])
    plt.xlabel('t [Myr]', fontsize=ftsz)
    plt.ylabel('n/n$_{eq}$', fontsize=ftsz)
    util.rescaleplot(t, np.array([0, 2]), plt, 0.05)
    plt.tick_params(labelsize=ftsz)

    plt.subplot(324)
    for i in range(nspec):
        plt.plot(t,
                 ly[i, :] - np.log10(ys[i]),
                 linestyle='-',
                 color=color[i],
                 linewidth=1.0,
                 label=label[i])
    plt.xlabel('t [Myr]', fontsize=ftsz)
    plt.ylabel('log n/n$_{eq}$', fontsize=ftsz)
    util.rescaleplot(t, np.array([-1, 1]), plt, 0.05)
    plt.tick_params(labelsize=ftsz)

    plt.subplot(325)
    wplot = np.array([0, 1, 3])
    for i in range(wplot.size):
        plt.plot(t,
                 np.log10(th2[wplot[i], :]),
                 linestyle='-',
                 linewidth=1.0,
                 label=tlabel[wplot[i]])
    plt.xlabel('t [Myr]', fontsize=ftsz)
    plt.ylabel('log reaction time [Myr]', fontsize=ftsz)
    plt.legend(fontsize=ftsz)
    util.rescaleplot(t, np.log10(th2[wplot, :]), plt, 0.05)
    plt.tick_params(labelsize=ftsz)

    plt.subplot(326)
    plt.plot(t[1:t.size],
             np.log10(it[1:t.size]),
             linestyle='-',
             color='black',
             linewidth=1.0,
             label='iterations')
    plt.xlabel('t [Myr]', fontsize=ftsz)
    plt.ylabel('log #', fontsize=ftsz)
    plt.legend(fontsize=ftsz)
    util.rescaleplot(t, np.log10(it[1:t.size]), plt, 0.05)
    plt.tick_params(labelsize=ftsz)
    plt.show()
Пример #8
0
def main():

    parser = argparse.ArgumentParser(formatter_class=RawTextHelpFormatter)
    parser.add_argument("T",
                        type=int,
                        help="length of Markov chain (depends on problem)")
    parser.add_argument("dorder",
                        type=int,
                        help="data model order: integer > 0")
    parser.add_argument("forder",
                        type=int,
                        help="polynomial fit order: integer > 0")
    parser.add_argument("sigamp",
                        type=float,
                        help="uncertainty amplitude in units of rms(y)")

    args = parser.parse_args()
    T = args.T
    dorder = args.dorder
    forder = args.forder
    sigamp = args.sigamp

    if (dorder == 0):
        raise Exception("[bayes_infer]: dorder = %i, but must be > 0\n" %
                        (dorder))
    if (forder == 0):
        raise Exception("[bayes_infer]: forder = %i, but must be > 0\n" %
                        (forder))

    dcoeff = np.array([1.0, 0.75, -1.5, 1.0])
    ndat = 30
    x, y, sy = make_data(ndat, dorder, dcoeff, sigamp)

    npar = forder + 2  # a parabola is of forder=2, but has 3 coefficients. We also fit the uncertainties.
    fGEN = gen_poly
    fPRI = pri_poly
    theta0 = np.zeros(npar) + 1.0
    dtheta = np.zeros((forder, npar))
    # note the separate dtheta for each parameter
    dtheta = np.array([
        [0.1, 0.1, 0.0, 0.0, 0.0, 0.0],  # constant and sigma
        [0.1, 0.1, 0.1, 0.0, 0.0, 0.0],  # linear and sigma
        [0.3, 0.3, 0.2, 0.1, 0.0, 0.0],  # quadratic and sigma
        [1.0, 1.5, 1.0, 0.1, 0.1, 0.0],  # cubic and sigma
        [1.0, 1.5, 1.0, 0.1, 0.1, 0.1]
    ])  # quartic and sigma
    T0 = np.array(
        T * (np.arange(forder) + 1.0)**2.5,
        dtype=int)  # need to increase MC length, but can't get too crazy
    nburn = np.array(0.1 * np.array(T0, dtype=float),
                     dtype=int)  # burn-in time. Problem-dependent
    nbin = 20
    pevidence = np.zeros(forder)
    ix = np.argsort(x)

    ftsz = 8
    fig1 = plt.figure(num=1,
                      figsize=(10, 2 * forder),
                      dpi=100,
                      facecolor='white')
    fig2 = plt.figure(num=2,
                      figsize=(10, 2 * forder),
                      dpi=100,
                      facecolor='white')

    for iford in range(0, forder):  # runs from 0 to forder-1.
        print('[bayes_infer]: iford=%1i T0=%7i nburn=%7i' %
              (iford + 1, T0[iford], nburn[iford]))
        if (T0[iford] <= nburn[iford]):
            print('[bayes_infer]: T0=%7i must be larger than nburn=%7i' %
                  (T0[iford], nburn[iford]))
            exit()
        inpar = iford + 2
        theta = (methast(T0[iford], theta0[0:inpar], dtheta[iford, 0:inpar], y,
                         x, fGEN, fPRI, iford))[:, nburn[iford]:T0[iford]]
        pevidence[iford] = marginalize_theta(theta[0:inpar], y, x, fGEN, fPRI,
                                             iford)
        hist = np.zeros((inpar, nbin, 2))
        mhist = np.zeros((inpar, 2))
        print('[bayes_infer]: iford=%1i evidence = %13.5e' %
              (iford, pevidence[iford]))
        print('[bayes_infer]: plotting histograms for iford=%1i' % (iford))
        for p in range(inpar):
            hist[p, :, :] = util.histogram(theta[p, :], nbin)
            mhist[p, 0] = np.sum(hist[p, :, 0] * hist[p, :, 1]) / np.sum(
                hist[p, :, 1])
            mhist[p, 1] = np.sqrt(
                np.sum(hist[p, :, 1] * (hist[p, :, 0] - mhist[p, 0])**2) /
                np.sum(hist[p, :, 1]))
            print('[bayes_infer]: iford=%1i   p=%1i theta=%13.5e+=%13.5e' %
                  (iford, p, mhist[p, 0], mhist[p, 1]))

            ax = fig1.add_subplot(forder, npar, p + (npar) * (iford) + 2)
            #ax.bar(hist[p,:,0],hist[p,:,1],width=(hist[p,1,0]-hist[p,0,0]),facecolor='green',align='center')
            ax.plot(hist[p, :, 0],
                    hist[p, :, 1],
                    color='black',
                    linewidth=1,
                    linestyle='-')
            ax.set_xlabel('p%1i' % (p), fontsize=ftsz)
            util.rescaleplot(hist[p, :, 0], hist[p, :, 1], ax, 0.05)
            ax.tick_params(labelsize=ftsz)

        print('[bayes_infer]: plotting data and fits for iford=%1i' % (iford))
        ax = fig1.add_subplot(forder, npar, (npar) * (iford) + 1)
        ax.plot(x[ix],
                fGEN(x[ix], mhist[:, 0], iford),
                linewidth=1,
                linestyle='-',
                color='red')
        ax.errorbar(x,
                    y,
                    yerr=sy,
                    fmt='o',
                    color='black',
                    mfc='none',
                    linewidth=1)
        ax.set_xlabel('x', fontsize=ftsz)
        ax.set_ylabel('y', fontsize=ftsz)
        util.rescaleplot(x, y, ax, 0.05)
        ax.tick_params(labelsize=ftsz)

        print('[bayes_infer]: plotting Markov chains for iford=%1i' % (iford))
        for p in range(iford + 2):
            ax = fig2.add_subplot(forder, forder + 1,
                                  p + (forder + 1) * (iford) + 1)
            ax.scatter(theta[p, :],
                       np.arange(max(theta.shape)),
                       color='black',
                       s=1)
            ax.set_xlabel('p%1i' % (p), fontsize=ftsz)
            ax.set_ylabel('t', fontsize=ftsz)
            util.rescaleplot(theta[p, :], np.arange(max(theta.shape)), ax,
                             0.05)
            ax.tick_params(labelsize=ftsz)

    pevidence = pevidence / np.sum(pevidence)
    fig3 = plt.figure(num=3, figsize=(6, 6), dpi=100, facecolor='white')
    ax = fig3.add_subplot(111)
    ax.plot(np.arange(forder),
            pevidence,
            linewidth=1,
            linestyle='-',
            color='black')
    ax.set_xlabel('order', fontsize=ftsz)
    ax.set_ylabel('p(D|M)', fontsize=ftsz)
    util.rescaleplot(np.arange(forder) + 1.0, pevidence, ax, 0.05)
    ax.tick_params(labelsize=ftsz)

    plt.show()
Пример #9
0
def ode_check(x, y, it, iprob):

    if (iprob == "lunarlander"):  # lunar lander problem
        n = x.size
        par = globalvar.get_odepar()
        z = y[0, :]  # altitude above ground
        vz = y[1, :]  # vertical velocity
        f = y[2, :]  # fuel mass
        thr = y[3, :]  # throttle

        accG = np.zeros(n)  # acceleration in units of g
        for k in range(n):
            accG[k] = ((dydx.lunarlanding(
                x[k], y[:, k], n /
                (x[n - 1] - x[0])))[1]) / 9.81  # acceleration

        ftsz = 10
        plt.figure(num=1, figsize=(8, 8), dpi=100, facecolor='white')

        plt.subplot(321)
        plt.plot(x, z, linestyle='-', color='black', linewidth=1.0)
        plt.xlabel('t [s]', fontsize=ftsz)
        plt.ylabel('z(t) [m]', fontsize=ftsz)
        util.rescaleplot(x, z, plt, 0.05)
        plt.tick_params(labelsize=ftsz)

        plt.subplot(322)
        plt.plot(x, vz, linestyle='-', color='black', linewidth=1.0)
        plt.xlabel('t [s]', fontsize=ftsz)
        plt.ylabel('v$_z$ [m s$^{-1}$]', fontsize=ftsz)
        util.rescaleplot(x, vz, plt, 0.05)
        plt.tick_params(labelsize=ftsz)

        plt.subplot(323)
        plt.plot(x, f, linestyle='-', color='black', linewidth=1.0)
        plt.xlabel('t [s]', fontsize=ftsz)
        plt.ylabel('fuel [kg]', fontsize=ftsz)
        util.rescaleplot(x, f, plt, 0.05)
        plt.tick_params(labelsize=ftsz)

        plt.subplot(324)
        plt.plot(x, thr, linestyle='-', color='black', linewidth=1.0)
        plt.xlabel('t [s]', fontsize=ftsz)
        plt.ylabel('throttle', fontsize=ftsz)
        util.rescaleplot(x, thr, plt, 0.05)
        plt.tick_params(labelsize=ftsz)

        plt.subplot(325)
        plt.plot(x, accG, linestyle='-', color='black', linewidth=1.0)
        plt.xlabel('t [s]', fontsize=ftsz)
        plt.ylabel('acc/G', fontsize=ftsz)
        util.rescaleplot(x, accG, plt, 0.05)
        plt.tick_params(labelsize=ftsz)

        plt.tight_layout()

        plt.show()

    elif (iprob == "kepler"):
        # for the direct Kepler problem, we check for energy and angular momentum conservation,
        # and for the center-of-mass position and velocity
        color = [
            'black', 'green', 'cyan', 'blue', 'red', 'black', 'black', 'black',
            'black'
        ]
        n = x.size
        par = globalvar.get_odepar()
        npar = par.size
        nbodies = par.size - 1
        gnewton = par[0]
        masses = par[1:npar]
        Egrav = np.zeros(n)
        indx = 2 * np.arange(nbodies)
        indy = 2 * np.arange(nbodies) + 1
        indvx = 2 * np.arange(nbodies) + 2 * nbodies
        indvy = 2 * np.arange(nbodies) + 2 * nbodies + 1
        E = np.zeros(n)  # total energy
        Lphi = np.zeros(n)  # angular momentum
        R = np.sqrt(
            np.power(y[indx[0], :] - y[indx[1], :], 2) +
            np.power(y[indy[0], :] - y[indy[1], :], 2))
        Rs = np.zeros(n)  # center of mass position
        vs = np.zeros(n)  # center of mass velocity
        for k in range(n):
            E[k] = 0.5 * np.sum(
                masses * (np.power(y[indvx, k], 2) + np.power(y[indvy, k], 2)))
            Lphi[k] = np.sum(
                masses * (y[indx, k] * y[indvy, k] - y[indy, k] * y[indvx, k]))
            Rsx = np.sum(masses * y[indx, k]) / np.sum(masses)
            Rsy = np.sum(masses * y[indy, k]) / np.sum(masses)
            vsx = np.sum(masses * y[indvx, k]) / np.sum(masses)
            vsy = np.sum(masses * y[indvy, k]) / np.sum(masses)
            Rs[k] = np.sqrt(Rsx * Rsx + Rsy * Rsy)
            vs[k] = np.sqrt(vsx * vsx + vsy * vsy)
        for j in range(nbodies):
            for i in range(
                    j):  # preventing double summation. Still O(N^2) though.
                dx = y[indx[j], :] - y[indx[i], :]
                dy = y[indy[j], :] - y[indy[i], :]
                Rt = np.sqrt(dx * dx + dy * dy)
                Egrav = Egrav - gnewton * masses[i] * masses[j] / Rt
        E = E + Egrav
        E = E / E[0]
        Lphi = Lphi / Lphi[0]
        for k in range(n):
            print(
                'k=%7i t=%13.5e E/E0=%20.12e L/L0=%20.12e Rs=%10.2e vs=%10.2e'
                % (k, x[k], E[k], Lphi[k], Rs[k], vs[k]))
        Eplot = E - 1.0
        Lplot = Lphi - 1.0
        # try Poincare cuts
        p1tot = np.zeros(n)
        p2tot = np.zeros(n)
        q1tot = np.zeros(n)
        q2tot = np.zeros(n)
        for k in range(n):
            p1tot[k] = np.sum(masses * y[indvx, k]) / np.sum(masses)
            p2tot[k] = np.sum(masses * y[indvy, k]) / np.sum(masses)
            q1tot[k] = np.sum(y[indx, k])
            q2tot[k] = np.sum(y[indy, k])
        thq = 1e-1
        thp = 1e-1
        print(
            '[ode_test]: min/max(q1), min/max(p1): %13.5e %13.5e %13.5e %13.5e'
            % (np.min(q1tot), np.max(q1tot), np.min(p1tot), np.max(p1tot)))
        tarr = (np.abs(p1tot) <= thp) & (np.abs(q1tot) <= thq)
        if (len(tarr) == 0):
            raise Exception('[ode_test]: no indices found at these thresholds')
        indp = np.where(tarr)[0]
        nind = indp.size
        print('[ode_test]: found %i elements for Poincare section\n' % (i))

        # now plot everything
        # (1) the orbits
        xmin = np.min(y[indx, :])
        xmax = np.max(y[indx, :])
        ymin = np.min(y[indy, :])
        ymax = np.max(y[indy, :])
        qmin = np.min(q2tot[indp])
        qmax = np.max(q2tot[indp])
        pmin = np.min(p2tot[indp])
        pmax = np.max(p2tot[indp])
        plt.figure(num=1, figsize=(8, 8), dpi=100, facecolor='white')
        plt.subplot(111)
        plt.xlim(1.05 * xmin, 1.05 * xmax)
        plt.ylim(1.05 * ymin, 1.05 * ymax)
        for k in range(nbodies):
            plt.plot(y[indx[k], :],
                     y[indy[k], :],
                     color=color[k],
                     linewidth=1.0,
                     linestyle='-')
        plt.axes().set_aspect('equal')
        plt.xlabel('x [AU]')
        plt.ylabel('y [AU]')
        # (2) the checks (total energy and angular momentum)
        plt.figure(num=2, dpi=100, facecolor='white')
        plt.subplot(311)
        plt.plot(x, Eplot, linestyle='-', color='black', linewidth=1.0)
        plt.xlabel('t [yr]')
        plt.ylabel('$\Delta$E/E')
        #plt.legend()
        plt.subplot(312)
        plt.plot(x, Lplot, linestyle='-', color='black', linewidth=1.0)
        plt.xlabel('t [yr]')
        plt.ylabel('$\Delta$L/L')
        #plt.legend()
        plt.subplot(313)
        plt.plot(q2tot[indp], p2tot[indp], '.', color='black')
        plt.xlabel('q$_2$')
        plt.ylabel('p$_2$')

        plt.tight_layout()

        plt.show()

    elif (iprob == "h2formation"):  # H2 formation
        color = ['black', 'blue', 'red']
        label = [u"H", u"H\u2082", u"H\u207A"]
        tlabel = ['k1', 'k2', 'k3', 'k4']
        Myr = 1e6 * 3.6e3 * 3.65e2 * 2.4e1
        nspec = (y.shape)[0]
        n = x.size
        par = globalvar.get_odepar()
        Ds = par[4]
        ys = get_h2eq(np.log10(Ds))
        print(
            'Equilibrium fractions: D=%13.5e A=%13.5e B=%13.5e C=%13.5e total=%13.5e'
            % (Ds, ys[0], ys[1], ys[2], np.sum(ys)))
        ys = ys * Ds
        th2 = np.zeros((4, n))
        for i in range(n):
            th2[:, i] = get_h2times(y[:, i])
        th2[2, :] = 0.0  # don't show CR ionization (really slow)
        t = x / Myr  # convert time in seconds to Myr
        tmin = np.min(t)
        tmax = np.max(t)
        nmin = np.min(np.array([np.nanmin(y), np.min(ys)]))
        nmax = np.max(np.array([np.nanmax(y), np.max(ys)]))
        ly = np.log10(y)
        lnmin = np.min(np.array([np.nanmin(ly), np.min(np.log10(ys))]))
        lnmax = np.max(np.array([np.nanmax(ly), np.max(np.log10(ys))]))

        # Note that the first call to subplot is slightly different, because we need the
        # axes object to reset the y-labels.
        ftsz = 10
        fig = plt.figure(num=1, figsize=(8, 8), dpi=100, facecolor='white')
        ax = fig.add_subplot(321)
        for i in range(nspec):
            plt.plot(t,
                     y[i, :],
                     linestyle='-',
                     color=color[i],
                     linewidth=1.0,
                     label=label[i])
            plt.plot([tmin, tmax], [ys[i], ys[i]],
                     linestyle='--',
                     color=color[i],
                     linewidth=1.0)
        plt.xlabel('t [Myr]', fontsize=ftsz)
        plt.ylabel('n$_\mathrm{\mathsf{H}}$', fontsize=ftsz)
        plt.legend(fontsize=10)
        util.rescaleplot(t, y, plt, 0.05)
        ylabels = ax.yaxis.get_ticklocs()
        ylabels1 = ylabels[1:ylabels.size - 1]
        ylabelstext = ['%4.2f' % (lb / np.max(ylabels1)) for lb in ylabels1]
        plt.yticks(ylabels1, ylabelstext)
        plt.tick_params(labelsize=ftsz)

        plt.subplot(322)
        for i in range(nspec):
            plt.plot(t,
                     ly[i, :],
                     linestyle='-',
                     color=color[i],
                     linewidth=1.0,
                     label=label[i])
            plt.plot(np.array([tmin, tmax]),
                     np.log10(np.array([ys[i], ys[i]])),
                     linestyle='--',
                     color=color[i],
                     linewidth=1.0)
        plt.xlabel('t [Myr]', fontsize=ftsz)
        plt.ylabel('log n [cm$^{-3}$]', fontsize=ftsz)
        util.rescaleplot(t, ly, plt, 0.05)
        plt.tick_params(labelsize=ftsz)

        plt.subplot(323)
        for i in range(nspec):
            plt.plot(t,
                     y[i, :] / ys[i],
                     linestyle='-',
                     color=color[i],
                     linewidth=1.0,
                     label=label[i])
        plt.xlabel('t [Myr]', fontsize=ftsz)
        plt.ylabel('n/n$_{eq}$', fontsize=ftsz)
        util.rescaleplot(t, np.array([0, 2]), plt, 0.05)
        plt.tick_params(labelsize=ftsz)

        plt.subplot(324)
        for i in range(nspec):
            plt.plot(t,
                     ly[i, :] - np.log10(ys[i]),
                     linestyle='-',
                     color=color[i],
                     linewidth=1.0,
                     label=label[i])
        plt.xlabel('t [Myr]', fontsize=ftsz)
        plt.ylabel('log n/n$_{eq}$', fontsize=ftsz)
        util.rescaleplot(t, np.array([-1, 1]), plt, 0.05)
        plt.tick_params(labelsize=ftsz)

        plt.subplot(325)
        wplot = np.array([0, 1, 3])
        for i in range(wplot.size):
            plt.plot(t,
                     np.log10(th2[wplot[i], :]),
                     linestyle='-',
                     linewidth=1.0,
                     label=tlabel[wplot[i]])
        plt.xlabel('t [Myr]', fontsize=ftsz)
        plt.ylabel('log reaction time [Myr]', fontsize=ftsz)
        plt.legend(fontsize=ftsz)
        util.rescaleplot(t, np.log10(th2[wplot, :]), plt, 0.05)
        plt.tick_params(labelsize=ftsz)

        plt.subplot(326)
        plt.plot(t[1:t.size],
                 np.log10(it[1:t.size]),
                 linestyle='-',
                 color='black',
                 linewidth=1.0,
                 label='iterations')
        plt.xlabel('t [Myr]', fontsize=ftsz)
        plt.ylabel('log #', fontsize=ftsz)
        plt.legend(fontsize=ftsz)
        util.rescaleplot(t, np.log10(it[1:t.size]), plt, 0.05)
        plt.tick_params(labelsize=ftsz)

        plt.tight_layout()

        plt.show()

    elif (iprob == "stiff1"):
        # Your code should show three plots in one window:
        # (1) the dependent variables y[0,:] and y[1,:] against x
        # (2) the residual y[0,:]-u and y[1,:]-v against x, see homework sheet
        # (3) the number of iterations  against x.
        xmin = np.min(x)
        xmax = np.max(x)
        ymin = np.min(y)
        ymax = np.max(y)

        ftsz = 10
        fig = plt.figure(num=1, figsize=(6, 8), dpi=100, facecolor='white')

        u = 2.0 * np.exp(-x) - np.exp(-1e3 * x)
        v = -np.exp(-x) + np.exp(-1e3 * x)

        diff = np.zeros((2, x.size))
        diff[0, :] = y[0] - u
        diff[1, :] = y[1] - v

        plt.subplot(311)
        plt.plot(x,
                 y[0, :],
                 linestyle='-',
                 color='blue',
                 linewidth=1.0,
                 label='u')
        plt.plot(x,
                 u,
                 linestyle='--',
                 color='blue',
                 linewidth=1.0,
                 label='u [analytic]')
        plt.plot(x,
                 y[1, :],
                 linestyle='-',
                 color='red',
                 linewidth=1.0,
                 label='v')
        plt.plot(x,
                 v,
                 linestyle='--',
                 color='red',
                 linewidth=1.0,
                 label='v [analytic]')
        plt.xlabel('x', fontsize=ftsz)
        plt.ylabel('y', fontsize=ftsz)
        plt.legend(fontsize=ftsz)
        util.rescaleplot(x, y, plt, 0.05)
        plt.tick_params(labelsize=ftsz)

        plt.subplot(312)
        plt.plot(x,
                 diff[0, :],
                 linestyle='-',
                 color='blue',
                 linewidth=1.0,
                 label='residual u')
        plt.plot(x,
                 diff[1, :],
                 linestyle='-',
                 color='red',
                 linewidth=1.0,
                 label='residual v')
        plt.xlabel('x', fontsize=ftsz)
        plt.ylabel('y-y[analytic]', fontsize=ftsz)
        plt.legend(fontsize=ftsz)
        util.rescaleplot(x, diff, plt, 0.05)
        plt.tick_params(labelsize=ftsz)

        plt.subplot(313)
        plt.plot(x[1:x.size],
                 np.log10(it[1:x.size]),
                 linestyle='-',
                 color='black',
                 linewidth=1.0,
                 label='iterations')
        plt.xlabel('x', fontsize=ftsz)
        plt.ylabel('log #', fontsize=ftsz)
        plt.legend(fontsize=ftsz)
        util.rescaleplot(x[1:x.size], np.log10(it[1:x.size]), plt, 0.05)
        plt.tick_params(labelsize=ftsz)

        plt.tight_layout()

        plt.show()

    elif (iprob == "stiff2"):

        xmin = np.min(x)
        xmax = np.max(x)
        ymin = np.min(y)
        ymax = np.max(y)

        ftsz = 10
        fig = plt.figure(num=1, figsize=(6, 8), dpi=100, facecolor='white')

        plt.subplot(211)
        for i in range(3):
            plt.plot(x,
                     y[i, :],
                     linestyle='-',
                     linewidth=1.0,
                     label='y[%i]' % (i))
        plt.xlabel('x', fontsize=ftsz)
        plt.ylabel('y', fontsize=ftsz)
        plt.legend(fontsize=ftsz)
        util.rescaleplot(x, y, plt, 0.05)
        plt.tick_params(labelsize=ftsz)

        plt.subplot(212)
        plt.plot(x[1:x.size],
                 np.log10(it[1:x.size]),
                 linestyle='-',
                 color='black',
                 linewidth=1.0,
                 label='iterations')
        plt.xlabel('x', fontsize=ftsz)
        plt.ylabel('log #', fontsize=ftsz)
        plt.legend(fontsize=ftsz)
        util.rescaleplot(x[1:x.size], np.log10(it[1:x.size]), plt, 0.05)
        plt.tick_params(labelsize=ftsz)

        plt.tight_layout()

        plt.show()

    elif ((iprob == "exponential") or (iprob == "gaussian")
          or (iprob == "tanh")):

        if (iprob == "gaussian"):
            f = dydx.gaussian(x, y, 1.0)
            sol = np.zeros(x.size)
            for i in range(x.size):
                sol[i] = 0.5 * math.erf(x[i] / np.sqrt(2.0))
        elif (iprob == "exponential"):
            f = dydx.exp(x, y, 1.0)
            sol = np.exp(x)
        elif (iprob == "tanh"):
            f = dydx.tanh(x, y, 1.0)
            sol = np.log(np.cosh(x))

        res = y[0, :] - sol

        ftsz = 10
        plt.figure(num=1, figsize=(8, 8), dpi=100, facecolor='white')

        plt.subplot(221)
        plt.plot(x, f, linestyle='-', color='black', linewidth=1.0)
        plt.xlabel('x', fontsize=ftsz)
        plt.ylabel("y'(x)", fontsize=ftsz)
        util.rescaleplot(x, f, plt, 0.05)
        plt.tick_params(labelsize=ftsz)

        plt.subplot(222)
        plt.plot(x,
                 y[0, :],
                 linestyle='-',
                 color='black',
                 linewidth=1.0,
                 label='y(x)')
        plt.plot(x,
                 sol,
                 linestyle='-',
                 color='red',
                 linewidth=1.0,
                 label='analytic')
        plt.xlabel('x', fontsize=ftsz)
        plt.ylabel("y(x)", fontsize=ftsz)
        plt.legend(fontsize=ftsz)
        util.rescaleplot(x, y, plt, 0.05)
        plt.tick_params(labelsize=ftsz)

        plt.subplot(223)
        plt.plot(x,
                 res,
                 linestyle='-',
                 color='black',
                 linewidth=1.0,
                 label='residual')
        plt.xlabel('x', fontsize=ftsz)
        plt.ylabel("y(x)", fontsize=ftsz)
        plt.legend(fontsize=ftsz)
        util.rescaleplot(x, res, plt, 0.05)
        plt.tick_params(labelsize=ftsz)

        plt.subplot(224)
        plt.plot(x[1:x.size],
                 it[1:x.size],
                 linestyle='-',
                 color='black',
                 linewidth=1.0,
                 label='iterations')
        plt.xlabel('x', fontsize=ftsz)
        plt.ylabel("#", fontsize=ftsz)
        plt.legend(fontsize=ftsz)
        util.rescaleplot(x, it, plt, 0.05)
        plt.tick_params(labelsize=ftsz)

        plt.tight_layout()

        plt.show()

    else:
        raise Exception('[ode_check]: invalid iprob %i\n' % (iprob))
Пример #10
0
def main():

    parser = argparse.ArgumentParser(formatter_class=RawTextHelpFormatter)
    parser.add_argument("stepper",
                        type=str,
                        default='euler',
                        help="stepping function:\n"
                        "   euler    : Euler step\n"
                        "   rk2      : Runge-Kutta 2nd order\n"
                        "   rk4      : Runge-Kutta 4th order\n"
                        "   backeuler: implicit Euler step\n")
    args = parser.parse_args()
    if (args.stepper == "euler"):
        fORD = step.euler
    elif (args.stepper == "rk2"):
        fORD = step.rk2
    elif (args.stepper == "rk4"):
        fORD = step.rk4
    elif (args.stepper == "backeuler"):
        fORD = step.backeuler
    else:
        raise Exception("invalid stepper %s" % (args.stepper))

    # initialization
    nstep = 100
    x0 = 0.0
    x1 = 1.0
    y0 = np.array([1.0, 0.0])
    fINT = odeint.ode_ivp
    fRHS = get_dydx
    fBVP = 0

    # solve
    x, y, it = fINT(fRHS, fORD, fBVP, x0, y0, x1, nstep)

    # Your code should show three plots in one window:
    # (1) the dependent variables y[0,:] and y[1,:] against x
    # (2) the residual y[0,:]-u and y[1,:]-v against x, see homework sheet
    # (3) the number of iterations  against x.
    # from here ??????
    xmin = np.min(x)
    xmax = np.max(x)
    ymin = np.min(y)
    ymax = np.max(y)
    ftsz = 10
    fig = plt.figure(num=1, figsize=(6, 8), dpi=100, facecolor='white')
    u = 2.0 * np.exp(-x) - np.exp(-1e3 * x)
    v = -np.exp(-x) + np.exp(-1e3 * x)
    diff = np.zeros((2, x.size))
    diff[0, :] = y[0] - u
    diff[1, :] = y[1] - v

    plt.subplot(311)
    plt.plot(x, y[0, :], linestyle='-', color='blue', linewidth=1.0, label='u')
    plt.plot(x,
             u,
             linestyle='--',
             color='blue',
             linewidth=1.0,
             label='u [analytic]')
    plt.plot(x, y[1, :], linestyle='-', color='red', linewidth=1.0, label='v')
    plt.plot(x,
             v,
             linestyle='--',
             color='red',
             linewidth=1.0,
             label='v [analytic]')
    plt.xlabel('x', fontsize=ftsz)
    plt.ylabel('y', fontsize=ftsz)
    plt.legend(fontsize=ftsz)
    util.rescaleplot(x, y, plt, 0.05)
    plt.tick_params(labelsize=ftsz)
    plt.subplot(312)
    plt.plot(x,
             diff[0, :],
             linestyle='-',
             color='blue',
             linewidth=1.0,
             label='residual u')
    plt.plot(x,
             diff[1, :],
             linestyle='-',
             color='red',
             linewidth=1.0,
             label='residual v')
    plt.xlabel('x', fontsize=ftsz)
    plt.ylabel('y-y[analytic]', fontsize=ftsz)
    plt.legend(fontsize=ftsz)
    util.rescaleplot(x, diff, plt, 0.05)
    plt.tick_params(labelsize=ftsz)
    plt.subplot(313)
    plt.plot(x[1:x.size],
             np.log10(it[1:x.size]),
             linestyle='-',
             color='black',
             linewidth=1.0,
             label='iterations')
    plt.xlabel('x', fontsize=ftsz)
    plt.ylabel('log #', fontsize=ftsz)
    plt.legend(fontsize=ftsz)
    util.rescaleplot(x[1:x.size], np.log10(it[1:x.size]), plt, 0.05)
    plt.tick_params(labelsize=ftsz)

    plt.show()
Пример #11
0
def check(X, E, s_prob, fPOT):
    s = X.shape
    N = s[0]  # number of time steps
    R = s[1]  # number of realizations

    xmin = np.min(X)
    xmax = np.max(X)

    if ((s_prob == 'free') or (s_prob == 'harmonic')
            or (s_prob == 'pertharmonic') or (s_prob == 'instanton')
            or (s_prob == 'box')):
        xmin = np.min(xmin)
        xmax = np.max(xmax)
        print("[check]: [xmin,xmax]=%13.5e,%13.5e" % (xmin, xmax))
        nbin = 200
        hist, edges = np.histogram(X[1:N - 1, :],
                                   nbin,
                                   range=(xmin, xmax),
                                   normed=False)
        xbin = 0.5 * (edges[0:edges.size - 1] + edges[1:edges.size])
        tothist = np.sum(hist.astype(float))
        hist = hist.astype(float) / tothist
        print('sum histogram: %13.5e' % (np.sum(hist)))
        psi2 = hist * float(nbin) / (xmax - xmin)
        psi2ana = get_psi2(xbin, s_prob)
        if (np.max(psi2ana) > 0.0):
            rmserr = np.sqrt(np.mean((psi2 - psi2ana)**2))
        else:
            rmserr = None
        pot = fPOT(xbin)
        print('int psi2     : %13.5e' % (np.sum(psi2) * (xbin[1] - xbin[0])))
        EX = np.mean(X)
        sX = np.std(X)
        EE = np.mean(np.mean(E, axis=0))
        sE = np.std(np.mean(E, axis=0))
        print("[check]: E[X] = %13.5e+-%13.5e  E[E]=%13.5e+-%13.5e" %
              (EX, sX, EE, sE))

        ftsz = 10
        # histogram of X - should be proportional to probability to find particle at given x.
        plt.figure(num=1, figsize=(8, 8), dpi=100, facecolor='white')
        plt.subplot(221)
        plt.bar(xbin,
                hist,
                width=(xbin[1] - xbin[0]),
                facecolor='green',
                align='center')
        plt.xlabel('$x$', fontsize=ftsz)
        plt.ylabel('$h(x)$', fontsize=ftsz)
        plt.tick_params(labelsize=ftsz)
        # the wave function - this needs to be appropriately normalized. See above.
        plt.subplot(222)
        plt.scatter(xbin, psi2, linewidth=1, color='red')
        if (rmserr != None):
            plt.plot(xbin,
                     psi2ana,
                     linewidth=1,
                     linestyle='-',
                     color='black',
                     label='analytic')
            plt.title('$<\Delta^2>^{1/2}$=%10.2e' % (rmserr), fontsize=ftsz)
        plt.xlabel('$x$', fontsize=ftsz)
        plt.ylabel('$|\psi(x)|^2$', fontsize=ftsz)
        plt.legend(fontsize=ftsz)
        util.rescaleplot(xbin, psi2, plt, 0.05)
        plt.tick_params(labelsize=ftsz)
        # the potential (just for informational purposes)
        plt.subplot(223)
        plt.plot(xbin,
                 pot,
                 linewidth=1,
                 linestyle='-',
                 color='black',
                 label='$V(x)$')
        plt.xlabel('$x$', fontsize=ftsz)
        plt.ylabel('$|\psi(x)|^2$', fontsize=ftsz)
        plt.legend(fontsize=ftsz)
        util.rescaleplot(xbin, pot, plt, 0.05)
        plt.tick_params(labelsize=ftsz)
        plt.show()
Пример #12
0
def ode_check(x, y, it, iprob):

    if (iprob == "h2formation"):  # H2 formation
        color = ['black', 'blue', 'red']
        label = [u"H", u"H\u2082", u"H\u207A"]
        tlabel = ['k1', 'k2', 'k3', 'k4']
        Myr = 1e6 * 3.6e3 * 3.65e2 * 2.4e1
        nspec = (y.shape)[0]
        n = x.size
        par = globalvar.get_odepar()
        Ds = par[4]
        ys = get_h2eq(np.log10(Ds))
        print(
            'Equilibrium fractions: D=%13.5e A=%13.5e B=%13.5e C=%13.5e total=%13.5e'
            % (Ds, ys[0], ys[1], ys[2], np.sum(ys)))
        ys = ys * Ds
        th2 = np.zeros((4, n))
        for i in range(n):
            th2[:, i] = get_h2times(y[:, i])
        th2[2, :] = 0.0  # don't show CR ionization (really slow)
        t = x / Myr  # convert time in seconds to Myr
        tmin = np.min(t)
        tmax = np.max(t)
        nmin = np.min(np.array([np.nanmin(y), np.min(ys)]))
        nmax = np.max(np.array([np.nanmax(y), np.max(ys)]))
        ly = np.log10(y)
        lnmin = np.min(np.array([np.nanmin(ly), np.min(np.log10(ys))]))
        lnmax = np.max(np.array([np.nanmax(ly), np.max(np.log10(ys))]))

        # Note that the first call to subplot is slightly different, because we need the
        # axes object to reset the y-labels.
        ftsz = 10
        fig = plt.figure(num=1, figsize=(8, 8), dpi=100, facecolor='white')
        ax = fig.add_subplot(321)
        for i in range(nspec):
            plt.plot(t,
                     y[i, :],
                     linestyle='-',
                     color=color[i],
                     linewidth=1.0,
                     label=label[i])
            plt.plot([tmin, tmax], [ys[i], ys[i]],
                     linestyle='--',
                     color=color[i],
                     linewidth=1.0)
        plt.xlabel('t [Myr]', fontsize=ftsz)
        plt.ylabel('n$_\mathrm{\mathsf{H}}$', fontsize=ftsz)
        plt.legend(fontsize=10)
        util.rescaleplot(t, y, plt, 0.05)
        ylabels = ax.yaxis.get_ticklocs()
        ylabels1 = ylabels[1:ylabels.size - 1]
        ylabelstext = ['%4.2f' % (lb / np.max(ylabels1)) for lb in ylabels1]
        plt.yticks(ylabels1, ylabelstext)
        plt.tick_params(labelsize=ftsz)

        plt.subplot(322)
        for i in range(nspec):
            plt.plot(t,
                     ly[i, :],
                     linestyle='-',
                     color=color[i],
                     linewidth=1.0,
                     label=label[i])
            plt.plot(np.array([tmin, tmax]),
                     np.log10(np.array([ys[i], ys[i]])),
                     linestyle='--',
                     color=color[i],
                     linewidth=1.0)
        plt.xlabel('t [Myr]', fontsize=ftsz)
        plt.ylabel('log n [cm$^{-3}$]', fontsize=ftsz)
        util.rescaleplot(t, ly, plt, 0.05)
        plt.tick_params(labelsize=ftsz)

        plt.subplot(323)
        for i in range(nspec):
            plt.plot(t,
                     y[i, :] / ys[i],
                     linestyle='-',
                     color=color[i],
                     linewidth=1.0,
                     label=label[i])
        plt.xlabel('t [Myr]', fontsize=ftsz)
        plt.ylabel('n/n$_{eq}$', fontsize=ftsz)
        util.rescaleplot(t, np.array([0, 2]), plt, 0.05)
        plt.tick_params(labelsize=ftsz)

        plt.subplot(324)
        for i in range(nspec):
            plt.plot(t,
                     ly[i, :] - np.log10(ys[i]),
                     linestyle='-',
                     color=color[i],
                     linewidth=1.0,
                     label=label[i])
        plt.xlabel('t [Myr]', fontsize=ftsz)
        plt.ylabel('log n/n$_{eq}$', fontsize=ftsz)
        util.rescaleplot(t, np.array([-1, 1]), plt, 0.05)
        plt.tick_params(labelsize=ftsz)

        plt.subplot(325)
        wplot = np.array([0, 1, 3])
        for i in range(wplot.size):
            plt.plot(t,
                     np.log10(th2[wplot[i], :]),
                     linestyle='-',
                     linewidth=1.0,
                     label=tlabel[wplot[i]])
        plt.xlabel('t [Myr]', fontsize=ftsz)
        plt.ylabel('log reaction time [Myr]', fontsize=ftsz)
        plt.legend(fontsize=ftsz)
        util.rescaleplot(t, np.log10(th2[wplot, :]), plt, 0.05)
        plt.tick_params(labelsize=ftsz)

        plt.subplot(326)
        plt.plot(t[1:t.size],
                 np.log10(it[1:t.size]),
                 linestyle='-',
                 color='black',
                 linewidth=1.0,
                 label='iterations')
        plt.xlabel('t [Myr]', fontsize=ftsz)
        plt.ylabel('log #', fontsize=ftsz)
        plt.legend(fontsize=ftsz)
        util.rescaleplot(t, np.log10(it[1:t.size]), plt, 0.05)
        plt.tick_params(labelsize=ftsz)

        print("[ode_test]: integration took %6i iterations" % (np.sum(it)))

        plt.tight_layout()
        plt.show()

    elif (iprob == "doubleexp"):
        xmin = np.min(x)
        xmax = np.max(x)
        ymin = np.min(y)
        ymax = np.max(y)

        ftsz = 10
        fig = plt.figure(num=1, figsize=(6, 8), dpi=100, facecolor='white')

        u = 2.0 * np.exp(-x) - np.exp(-1e3 * x)
        v = -np.exp(-x) + np.exp(-1e3 * x)

        diff = np.zeros((2, x.size))
        diff[0, :] = y[0] - u
        diff[1, :] = y[1] - v

        plt.subplot(311)
        plt.plot(x,
                 y[0, :],
                 linestyle='-',
                 color='blue',
                 linewidth=1.0,
                 label='u')
        plt.plot(x,
                 u,
                 linestyle='--',
                 color='blue',
                 linewidth=1.0,
                 label='u [analytic]')
        plt.plot(x,
                 y[1, :],
                 linestyle='-',
                 color='red',
                 linewidth=1.0,
                 label='v')
        plt.plot(x,
                 v,
                 linestyle='--',
                 color='red',
                 linewidth=1.0,
                 label='v [analytic]')
        plt.xlabel('x', fontsize=ftsz)
        plt.ylabel('y', fontsize=ftsz)
        plt.legend(fontsize=ftsz)
        util.rescaleplot(x, y, plt, 0.05)
        plt.tick_params(labelsize=ftsz)

        plt.subplot(312)
        plt.plot(x,
                 diff[0, :],
                 linestyle='-',
                 color='blue',
                 linewidth=1.0,
                 label='residual u')
        plt.plot(x,
                 diff[1, :],
                 linestyle='-',
                 color='red',
                 linewidth=1.0,
                 label='residual v')
        plt.xlabel('x', fontsize=ftsz)
        plt.ylabel('y-y[analytic]', fontsize=ftsz)
        plt.legend(fontsize=ftsz)
        util.rescaleplot(x, diff, plt, 0.05)
        plt.tick_params(labelsize=ftsz)

        plt.subplot(313)
        plt.plot(x[1:x.size],
                 np.log10(it[1:x.size]),
                 linestyle='-',
                 color='black',
                 linewidth=1.0,
                 label='iterations')
        plt.xlabel('x', fontsize=ftsz)
        plt.ylabel('log #', fontsize=ftsz)
        plt.legend(fontsize=ftsz)
        util.rescaleplot(x[1:x.size], np.log10(it[1:x.size]), plt, 0.05)
        plt.tick_params(labelsize=ftsz)

        print("[ode_test]: integration took %6i iterations" % (np.sum(it)))

        plt.tight_layout()
        plt.show()

    elif (iprob == "enrightpryce"):

        xmin = np.min(x)
        xmax = np.max(x)
        ymin = np.min(y)
        ymax = np.max(y)

        ftsz = 10
        fig = plt.figure(num=1, figsize=(6, 8), dpi=100, facecolor='white')

        plt.subplot(211)
        for i in range(3):
            plt.plot(x,
                     y[i, :],
                     linestyle='-',
                     linewidth=1.0,
                     label='y[%i]' % (i))
        plt.xlabel('x', fontsize=ftsz)
        plt.ylabel('y', fontsize=ftsz)
        plt.legend(fontsize=ftsz)
        util.rescaleplot(x, y, plt, 0.05)
        plt.tick_params(labelsize=ftsz)

        plt.subplot(212)
        plt.plot(x[1:x.size],
                 np.log10(it[1:x.size]),
                 linestyle='-',
                 color='black',
                 linewidth=1.0,
                 label='iterations')
        plt.xlabel('x', fontsize=ftsz)
        plt.ylabel('log #', fontsize=ftsz)
        plt.legend(fontsize=ftsz)
        util.rescaleplot(x[1:x.size], np.log10(it[1:x.size]), plt, 0.05)
        plt.tick_params(labelsize=ftsz)

        print("[ode_test]: integration took %6i iterations" % (np.sum(it)))

        plt.tight_layout()
        plt.show()

    else:
        raise Exception('[ode_check]: invalid iprob %i\n' % (iprob))
def main():

    parser = argparse.ArgumentParser(formatter_class=RawTextHelpFormatter)
    parser.add_argument("sinteg",
                        type=str,
                        help="function to be integrated:\n"
                        "   expabs   : double exponential\n"
                        "   gaussian : normal distribution\n"
                        "   lognormal: lognormal distribution")
    parser.add_argument("ipower",
                        type=int,
                        help="p-th order moment [p=0,1,2,...]")
    parser.add_argument("ssample",
                        type=str,
                        help="sampling function:\n"
                        "   uniform : uniform distribution\n"
                        "   cauchy  : Cauchy distribution\n"
                        "   normal  : normal distribution")
    parser.add_argument("N",
                        type=int,
                        help="number of support points in single integral")
    parser.add_argument("R",
                        type=int,
                        help="number of realizations (=experiments). R >0")
    args = parser.parse_args()
    sinteg = args.sinteg
    p = args.ipower
    ssample = args.ssample
    N = args.N
    R = args.R
    if (R <= 0):
        parser.error("R must be larger than zero")

    fINT, fRAN, fDEN, a, b, mu, sg = init(sinteg, ssample)

    if (R == 1):
        ev, sd = importsamp(fINT, fRAN, fDEN, N, a=a, b=b, mu=mu, sg=sg, p=p)
        x = a + (b - a) * np.arange(200) / 199.0
        fint = fINT(x, a=a, b=b, mu=mu, sg=sg, p=p)
        fden = fDEN(x, a=a, b=b, mu=mu, sg=sg, p=p)
        minf = np.min(np.array([np.min(fint), np.min(fden)]))
        maxf = np.max(np.array([np.max(fint), np.max(fden)]))
        print("[mci_importance]: I = %13.5e += %13.5e" % (ev, sd))
        ftsz = 10
        plt.figure(num=1, figsize=(8, 8), dpi=100, facecolor='white')
        plt.subplot(111)
        plt.plot(x,
                 fint,
                 linewidth=1,
                 linestyle='-',
                 color='black',
                 label='f(x)')
        plt.plot(x,
                 fden,
                 linewidth=1,
                 linestyle='-',
                 color='red',
                 label='p(x)')
        plt.xlabel('x', fontsize=ftsz)
        plt.ylabel('f(x),p(x)', fontsize=ftsz)
        plt.title('$\mu=$%11.3e $\sigma=$%11.3e' % (ev, sd))
        plt.legend()
        util.rescaleplot(x, np.array([minf, maxf]), plt, 0.05)
        plt.tick_params(labelsize=ftsz)

        plt.show()
    else:
        ev = np.zeros(R)
        sd = np.zeros(R)
        Ev = np.zeros(R)
        Sd = np.zeros(R)
        for i in range(R):
            for j in range(i):
                ev[j], sd[j] = importsamp(fINT,
                                          fRAN,
                                          fDEN,
                                          N,
                                          a=a,
                                          b=b,
                                          mu=mu,
                                          sg=sg,
                                          p=p)
            Ev[i] = np.mean(ev[0:i + 1])
            Sd[i] = np.std(ev[0:i + 1])
        ftsz = 10
        plt.figure(num=1, figsize=(6, 8), dpi=100, facecolor='white')
        plt.subplot(311)
        plt.plot(np.arange(R - 2) + 1,
                 Ev[2:R],
                 linestyle='-',
                 linewidth=1,
                 color='black')
        plt.xlabel('R', fontsize=ftsz)
        plt.ylabel('Ev', fontsize=ftsz)
        plt.tick_params(labelsize=ftsz)
        plt.subplot(312)
        plt.plot(np.log10(np.arange(R - 2) + 1),
                 np.log10(Sd[2:R]),
                 linestyle='-',
                 linewidth=1,
                 color='black')
        plt.xlabel('log R', fontsize=ftsz)
        plt.ylabel('log Var', fontsize=ftsz)
        plt.tick_params(labelsize=ftsz)
        plt.subplot(313)
        hist, edges = np.histogram(ev[0:R - 1],
                                   np.int(np.sqrt(R)),
                                   normed=False)
        x = 0.5 * (edges[0:edges.size - 1] + edges[1:edges.size])
        tothist = np.sum(hist.astype(float))
        hist = hist.astype(
            float
        ) / tothist  # it seems the "normed" keyword does not work in numpy/mathplotlib
        normd = mlab.normpdf(x, Ev[R - 1], Sd[R - 1]) * (x[1] - x[0])
        peak = mlab.normpdf(Ev[R - 1], Ev[R - 1], Sd[R - 1]) * (
            x[1] - x[0])  # need this for FWHM
        print('[mci_importance]: sum(hist)   = %13.5e %13.5e' %
              (np.sum(hist), tothist))
        print('[mci_importance]: expectation = %13.5e +-%12.5e' %
              (Ev[R - 1], Sd[R - 1]))
        plt.bar(x,
                hist,
                width=(x[1] - x[0]),
                facecolor='green',
                align='center')
        plt.plot(np.array([1.0, 1.0]) * Ev[R - 1],
                 np.array([0.0, 1.0]),
                 linestyle='--',
                 color='black',
                 linewidth=1.0)
        plt.plot(np.array([
            Ev[R - 1] - 0.5 * 2.36 * Sd[R - 1],
            Ev[R - 1] + 0.5 * 2.36 * Sd[R - 1]
        ]),
                 0.5 * np.array([peak, peak]),
                 linestyle='--',
                 color='black',
                 linewidth=1.0)
        plt.plot(x, normd, linestyle='--', color='red', linewidth=2.0)
        plt.xlabel('E[I]', fontsize=ftsz)
        plt.ylabel('frequency', fontsize=ftsz)
        plt.title('$\mu=$%11.3e $\sigma=$%11.3e' % (Ev[R - 1], Sd[R - 1]))
        plt.ylim(np.array([0.0, 1.05 * np.max(hist)]))
        plt.tick_params(labelsize=ftsz)

        plt.show()