A small routine to compute bifurcation diagrams

bifurcation.py

import numpy as np
import matplotlib.pyplot as plt

"""
Compute the trajectory of a function f: (R^n x R) -> R^n, where x_{t + 1} = f(x_t, beta). 
"""
def trajectory(f, beta, x0: np.ndarray, T: int) -> np.ndarray:

    n = x0.shape[0]
    X = np.zeros((T + 1, n), dtype = np.float64)
    X[0, :] = x0
    
    for t in range(1, T + 1):
        X[t, :] = f(X[t - 1, :], beta)

    return X

"""
Compute the orbits of a function f: (R^dim x R) -> R^dim, where x_{t + 1} = f(x_t, beta), for parameters beta in the parameters range.
"""
def orbits(f, parameter_range: np.ndarray, n = 1_000, tr = 100, dim = 2, idx = 1):
    x0 = np.random.normal(loc = 0, scale = 0.25, size = (n, dim))

    orbit = []
    for beta in parameter_range:
        beta_orbit = []
        for i in range(n):
            last = trajectory(f, beta, x0[i, :], tr)[-1, idx] # Save the last point
            beta_orbit.append(last)

        orbit.append(beta_orbit)


    return orbit

if __name__ == "__main__":
    d, s = 3/4, 1 # Default parameters
    c = 1 # 

    def phi(x: np.ndarray, beta):

        n = x.shape[0]
        x_prime = np.empty(n)

        intercept = 2 * d / s

        x_prime[0] = -x[0] * (1 - x[1]) / (intercept + 1 + x[1])

        profit_diff = s * np.square(x_prime[0] - x[0]) / 2  - c
        x_prime[1] = np.tanh(profit_diff * beta / 2)

        return x_prime

    # Computing a trajectory
    T = 80
    x0 = np.random.normal(size = 2)
    tr = trajectory(phi, 8, x0, T)

    fig, ax = plt.subplots()
    ax.plot(range(T + 1), tr[:, 0])
    fig.savefig("trajectory.png")
    
    # Computing and saving a bifurcation
    print("Computing...")
    parameters = np.arange(start = 0.5, stop = 10, step = 0.1)
    beta_orbits = orbits(phi, parameters, n = 2_000, tr = T)

    print("Plotting...")
    fig, ax = plt.subplots()

    for (x, y) in zip(parameters, beta_orbits):
        ax.scatter([x] * len(y), y, c = "black", s = 1)

    fig.savefig("bifurcation.png")