Lorenz System¶

def lorenz(xyz, t, sigma, rho, beta):
    x, y, z = xyz  # parse variables
    dxdt = sigma * (y - x)
    dydt = x * (rho - z) - y
    dzdt = x * y - beta * z
    return [dxdt, dydt, dzdt]


sigma = 10.0
beta = 8 / 3
rho = 28.0

T = 50.0  # total time to run
dt = 0.01  # time step
time_points = np.arange(0.0, T, dt)

# Nontrivial steady state values for plotting
x2 = y2 = np.sqrt(beta * (rho - 1))  # nontrivial steady state
z2 = rho - 1

init = np.random.rand(3) * 2 - 1  # initial conditions between -1 and 1
sol2 = odeint(lorenz, init, time_points, args=(sigma, rho, beta))

plt.figure(figsize=(10, 6))
# Plot nontrivial steady states as horizontal lines
plt.axhline(x2, color="k", lw=1, linestyle="--", label=f"Steady state: x={x2:.2f}")
plt.axhline(-x2, color="k", lw=1, linestyle="--", label=f"Steady state: x={-x2:.2f}")

# Plot the single trajectory
plt.plot(time_points, sol2[:, 0], lw=2, label=f"Trajectory starting at {init}")

plt.xlabel(r"$t$")
plt.ylabel(r"$x$")
plt.title("Lorenz Attractor: Single Trajectory")
plt.legend()
plt.show()

plt.figure()
plt.plot(sol2[:, 0], sol2[:, 2])  # plot trajectories projected onto the (x,z) plane
plt.plot(
    [sol2[0, 0]], [sol2[0, 2]], "b^"
)  # blue trianble labels starting point of each trajectory
plt.plot([x2], [z2], "rx")  # steady state
plt.plot([-x2], [z2], "rx")  # steady state
plt.xlabel(r"$x$")
plt.ylabel(r"$z$")
plt.show()

# Let us now visualize all the solutions as trajectories in a 3-d plot

fig = plt.figure()
ax = fig.add_subplot(projection="3d")
ax.plot(sol2[:, 0], sol2[:, 1], sol2[:, 2])
ax.set_xlabel(r"$x$")
ax.set_ylabel(r"$y$")
ax.set_zlabel(r"$z$")
plt.show()

sigma = 10.0
beta = 8 / 3.0

t1 = 20.0  # time to add perturbation
T = 100.0
dt = 0.01
time_points = np.arange(0.0, 50.0, 0.01)

rho_list = np.geomspace(0.1, 1000.0, 401)  # list of rho values geometrically spaced
sol_list = []  # list of solutions of Lorenz equations

for rho in rho_list:
    init = np.random.rand(3) * 2 - 1  # random initial values
    sol = odeint(lorenz, init, time_points, args=(sigma, rho, beta))
    sol_list.append(sol)


rP = 1  # onset of instability at the origin, pitchfork bifurcation
rH = sigma * (sigma + beta + 3) / (sigma - beta - 1)  # onset of chaos, Hopf bifurcation

plt.figure(figsize=(8, 6))
plt.axvline(rP, lw=1, ls=":", label=r"$\rho = 1$")
plt.axvline(rH, lw=1, ls="--", label=r"$\rho = %.3f$" % rH)
plt.axvline(28, lw=1, ls="-", label=r"$\rho = 28$")
for i in range(len(rho_list)):
    rho = rho_list[i]
    sol = sol_list[i]
    y = sol[int(t1 / dt) :, 0]
    x = [rho] * len(y)
    plt.scatter(x, y, s=0.1, c="k", marker=".", edgecolor="none")

plt.xscale("log")
plt.xlabel(r"$\rho$")
plt.ylabel(r"$x$")
plt.legend()
plt.show()

plt.plot(sol2)

def prepare_data(data, y_col, x_cols, train_size=4000):
    """Split and reshape the data for each model."""
    train_data, test_data = data[:train_size], data[train_size:]

    # Extract y (target) and x (predictors)
    y_train, y_test = train_data[:, y_col].reshape(-1, 1), test_data[:, y_col].reshape(
        -1, 1
    )
    x_train, x_test = train_data[:, x_cols], test_data[:, x_cols]

    return x_train, y_train, x_test, y_test


def train_and_predict(x_train, y_train, x_test, y_test):
    """Train the model and make predictions."""
    basis_function = Polynomial(degree=2)
    model = FROLS(
        ylag=1,
        xlag=[[1], [1]],
        estimator=LeastSquares(),
        basis_function=basis_function,
    )
    model.fit(X=x_train, y=y_train)
    yhat = model.predict(
        X=x_test,
        y=y_test[: model.max_lag].reshape(-1, 1),
    )

    r = pd.DataFrame(
        results(
            model.final_model,
            model.theta,
            model.err,
            model.n_terms,
            err_precision=8,
            dtype="sci",
        ),
        columns=["Regressors", "Parameters", "ERR"],
    )
    print(r)
    return yhat


# First model: Using sol2[:, 0] for y and sol2[:, [1, 2]] for x
x_train, y_train, x_test, y_test = prepare_data(
    sol2, y_col=0, x_cols=[1, 2], train_size=4000
)
yhat_1 = train_and_predict(x_train, y_train, x_test, y_test)

# set 1 because is the maximum lag. Can set as model.max_lag as well
plot_results(y=y_test[1:], yhat=yhat_1[1:], n=1000)

# Second model: Using sol2[:, 1] for y and sol2[:, [0, 2]] for x
x_train, y_train, x_test, y_test = prepare_data(
    sol2, y_col=1, x_cols=[0, 2], train_size=4000
)
yhat_2 = train_and_predict(x_train, y_train, x_test, y_test)

plot_results(y=y_test[1:], yhat=yhat_2[1:], n=1000)

# Third model: Using sol2[:, 2] for y and sol2[:, [0, 1]] for x
x_train, y_train, x_test, y_test = prepare_data(
    sol2, y_col=2, x_cols=[0, 1], train_size=4000
)
yhat_3 = train_and_predict(x_train, y_train, x_test, y_test)

plot_results(y=y_test[1:], yhat=yhat_3[1:], n=1000)

# First model: Using predictions from yhat_1 (predicted x vs predicted z)
plt.figure()
plt.plot(yhat_1, yhat_3)  # plot predictions projected onto the (x, z) plane
plt.plot(
    [sol2[0, 0]], [sol2[0, 2]], "b^"
)  # blue triangle labels starting point of each trajectory
plt.plot([x2], [z2], "rx")  # steady state
plt.plot([-x2], [z2], "rx")  # steady state
plt.xlabel(r"$x$")
plt.ylabel(r"$z$")
plt.title("Butterfly Plot - Model 1 Predictions")
plt.show()

# Second model: Using predictions from yhat_2 (predicted y vs predicted z)
plt.figure()
plt.plot(yhat_2, yhat_3)  # plot predictions projected onto the (y, z) plane
plt.plot(
    [sol2[0, 1]], [sol2[0, 2]], "b^"
)  # blue triangle labels starting point of each trajectory
plt.plot([x2], [z2], "rx")  # steady state
plt.plot([-x2], [z2], "rx")  # steady state
plt.xlabel(r"$y$")
plt.ylabel(r"$z$")
plt.title("Butterfly Plot - Model 2 Predictions")
plt.show()

# Third model: Using predictions from yhat_3 (predicted z vs predicted x)
plt.figure()
plt.plot(yhat_3, yhat_1)  # plot predictions projected onto the (z, x) plane
plt.plot(
    [sol2[0, 2]], [sol2[0, 0]], "b^"
)  # blue triangle labels starting point of each trajectory
plt.plot([z2], [x2], "rx")  # steady state
plt.plot([z2], [-x2], "rx")  # steady state
plt.xlabel(r"$z$")
plt.ylabel(r"$x$")
plt.title("Butterfly Plot - Model 3 Predictions")
plt.show()

fig = plt.figure()
ax = fig.add_subplot(111, projection="3d")
ax.plot(yhat_1, yhat_2, yhat_3)
ax.set_xlabel(r"$x$")
ax.set_ylabel(r"$y$")
ax.set_zlabel(r"$z$")
ax.set_title("3D Butterfly Plot with Predicted Data")
plt.show()