plot_logfun.py

import numpy as np
from matplotlib import pyplot as plt

from logistic import iterate_f

def plot_trajectory(n, r, x0, fname="single_trajectory.png"):
    """
    Saves a plot of a single trajectory of the logistic function

    inputs
        n: int (number of iterations)
        r: float (r value for the logistic function)
        x0: float (between 0 and 1, starting point for the iteration)
        fname: str (filename to which to save the image)

    returns
        fig, ax (matplotlib objects)
    """
    l = iterate_f(n, x0, r)
    fig, ax = plt.subplots(figsize=(10,5))
    ax.plot(list(range(n)), l)
    fig.suptitle('Logistic Function')

    fig.savefig(fname)
    return fig, ax


def plot_bifurcation(start, end, step, fname="bifurcation.png", it=100000,
                     last=300):
    """
    Saves a plot of the bifurcation diagram of the logistic function. The
    `start`, `end`, and `step` parameters define for which r values to calculate
    the logistic function. If you space them too closely, it might take a very
    long time, if you dont plot enough, your bifurcation diagram won't be
    informative. Choose wisely!

    inputs
        start, end, step:  float (which r values to calculate the logistic
                                  function for)
        fname: str (filename to which to save the image)
        it: int (how many iterations to run for each r value)
        last: int (how many of the last iterates to plot)


    returns
        fig, ax (matplotlib objects)
    """
    r_range=np.arange(start, end, step)
    x=[]
    y=[]

    for r in r_range:
        l = iterate_f(it, 0.1, r)
        ll = l[len(l)-last::].copy()
        lll = np.unique(ll)
        y.extend(lll)
        x.extend(np.ones(len(lll))*r)

    fig, ax = plt.subplots(figsize=(20,10))
    ax.scatter(x, y, s=0.1, color='k')
    ax.set_xlabel("r")
    fig.savefig(fname)
    return fig, ax


if __name__=="__main__":
    plot_trajectory(100, 3.6, 0.1)
    plot_bifurcation(2.5, 4.2, 0.001)