Last active
June 22, 2025 16:23
-
-
Save nbelakovski/fea571de1f6e2b3010aa318c7f95e1ad to your computer and use it in GitHub Desktop.
Collocation with a monomial basis and multiple elements
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| import numpy as np | |
| from scipy.optimize import root | |
| from scipy.integrate import solve_ivp | |
| import matplotlib.pyplot as plt | |
| # Set up our approximation functions and a function to generate the system of equations we'll need | |
| def xtilde(t, a0_to_an): | |
| degree = len(a0_to_an) - 1 | |
| terms = [a0_to_an[i] * t**i for i in range(degree+1)] | |
| return sum(terms) | |
| def xtildedot(t, a0_to_an): | |
| degree = len(a0_to_an) - 1 | |
| terms = [a0_to_an[i] * i * t**(i-1) for i in range(1, degree+1)] | |
| return sum(terms) | |
| def create_system_of_eqns(dxdt, x0, t0, tf, degree, m=1): | |
| time_grid = np.linspace(t0, tf, m+1) | |
| elements = zip(time_grid[:-1], time_grid[1:]) | |
| elements = list(elements) | |
| collocation_points = [] | |
| for element in elements: | |
| collocation_points.append(np.linspace(element[0], element[1], degree)) | |
| def system(coeffs): | |
| coeffs = coeffs.reshape((m, degree+1)) | |
| coeffs1 = coeffs[0] | |
| eq_ic = [ | |
| xtilde(t0, coeffs1) - x0 | |
| ] | |
| eq_points = [ | |
| xtildedot(ti, coeffs1) - dxdt(ti, xtilde(ti, coeffs1)) | |
| for ti in collocation_points[0] | |
| ] | |
| eqs = eq_ic + eq_points | |
| for i in range(1, m): | |
| element = elements[i] | |
| coeffs_iminus1 = coeffs[i-1] | |
| coeffs_i = coeffs[i] | |
| eq_equality = [ | |
| xtilde(element[0], coeffs_iminus1) - xtilde(element[0], coeffs_i) | |
| ] | |
| eq_points = [ | |
| xtildedot(ti, coeffs_i) - dxdt(ti, xtilde(ti, coeffs_i)) | |
| for ti in collocation_points[i] | |
| ] | |
| eqs += eq_equality + eq_points | |
| return eqs | |
| def piecewise(t, coeffs): | |
| coeffs = coeffs.reshape((m, degree + 1)) | |
| if t <= elements[0][0]: | |
| return xtilde(t, coeffs[0]) | |
| elif t >= elements[-1][1]: | |
| return xtilde(t, coeffs[-1]) | |
| else: | |
| for i, element in enumerate(elements): | |
| if element[0] <= t <= element[1]: | |
| return xtilde(t, coeffs[i]) | |
| return system, piecewise | |
| # And now we can set up the problem | |
| def dxdt(t, x): | |
| return -(x+5)*np.cos(x) | |
| t0 = 0 | |
| x0 = 1 | |
| tf = 5 | |
| m = 5 | |
| degree = 7 | |
| system, piecewise = create_system_of_eqns(dxdt, x0, t0, tf, degree, m) | |
| initial_guess = np.zeros((m, (degree + 1))) | |
| solution = root(system, initial_guess) | |
| solution.x = solution.x.reshape((m, (degree + 1))) | |
| if solution.success: | |
| print(f'Solution converged in {solution.nfev} function evaluations') | |
| else: | |
| print(f"Solution failed to converge: {solution.message}") | |
| # Here we use standard numerical integration to solve the ODE and plot a comparison to the collocation results | |
| t_eval = np.linspace(t0, tf, num=100) | |
| sol = solve_ivp(dxdt, [t0, tf], [x0], t_eval=t_eval) | |
| plt.plot(sol.t, sol.y[0], label="$RK45$") | |
| plt.plot(sol.t, [piecewise(ti, solution.x) for ti in sol.t], label=r"$\tilde{x}$") | |
| plt.grid() | |
| plt.legend() | |
| plt.show() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment