-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathodeint-2d-flow.py
170 lines (121 loc) · 5.58 KB
/
odeint-2d-flow.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
# -*- coding: utf-8 -*-
"""
Created on Sun Jul 26 15:47:51 2020
@author: artmenlope
Fun with scipy's odeint
"""
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from scipy.integrate import odeint
def acel_x(x, y):
"""
X-component of the acceleration vector.
"""
a_x = -np.cos(y)*np.sin(x)
return a_x
def acel_y(x, y):
"""
Y-component of the acceleration vector.
"""
a_y = -np.sin(y)*np.cos(x)
return a_y
def Mechanics(PositVeloc, t, n):
"""
Function defining the mechanics of the points. Created to be used with the
scipy.integrate.odeint function.
Letting xi and yi be the x and y coordinates of the point i and vxi and vyi
its velocity components on each axis:
PositVeloc must be a 1D Numpy array or a list of the form
[x1, vx1, y1, vy1, x2, vx2, y2, vy2, x3, vx3, y3, vy3, ...].
Calling axi and ayi to the components of the acceleration of the point i,
Mechanics returns an array ddt of the form
[vx1, ax1, vy1, ay1, vx2, ax2, vy2, ay2, vx3, ax3, vy3, ay3, ...]
"""
PositVeloc = np.asarray(PositVeloc) # Array of positions and velocities.
xs = PositVeloc[0::4] # Set of x coordinates.
vxs = PositVeloc[1::4] # Set of vx velocities.
ys = PositVeloc[2::4] # Set of y coordinates.
vys = PositVeloc[3::4] # Set of vy velocities.
axs = acel_x(xs, ys) # X-component of the accelerations. The acel_x function is defined externally.
ays = acel_y(xs, ys) # Y-component of the accelerations. The acel_y function is defined externally.
ddt_aux = np.stack((vxs, axs, vys, ays), axis=1) # Create a 2D Numpy array containing [vxi, axi, vyi, ayi] in each row.
ddt = np.reshape(ddt_aux, (4*n**2)) # Reshape ddt_aux to be of the form [vx1, ax1, vy1, ay1, vx2, ax2, vy2, ay2, ...]
return ddt
n = 50 # Number of points
lim = 1 # Limits of the point generation. The points will be generated on the square [-lim,lim]x[-lim,lim]
plotlim = lim # Plot visualization limits. Square of the form [-plotlim,plotlim]x[-plotlim,plotlim]
# Make a grid of points located at their initial positions.
x0 = np.linspace(-lim,lim,n)
y0 = np.linspace(-lim,lim,n)
X0, Y0 = np.meshgrid(x0, y0)
# Reshape the grid to obtain an array containing in each row the initial
# coordinates [x0i,y0i] of each point.
points = np.stack((X0, Y0), axis=-1).reshape((n**2,2))
x0s = points[:,0] # Set of initial x-coordinates. 1D array.
y0s = points[:,1] # Set of initial y-coordinates. 1D array.
# Initial velocity components of the points. 1D arrays.
vx0s = np.zeros(n**2)
vy0s = np.zeros(n**2)
tf = 50 # Final time value.
dt = 0.1 # Size of the time steps.
t = np.arange(0, tf+dt, dt) # 1D array of time values.
nt = len(t) # Number of time values.
# Create a 1D array of initial conditions for odeint.
initCond_aux = np.stack((x0s, vx0s, y0s, vy0s), axis=1) # Create a 2D Numpy array containing [xi, vxi, yi, vyi] in each row. See the Mechanics() function for more details.
init_cond = np.reshape(initCond_aux, (4*n**2)) # Reshape initCond_aux to be of the form [x1, vx1, y1, vy1, x2, vx2, y2, vy2, ...].
# We solve the position of all the dots at any time value of t.
abserr = 1.0e-6
relerr = 1.0e-4
solutions = odeint(Mechanics, init_cond, t, args=(n,), atol=abserr, rtol=relerr)
# Get the coordinates and velocities of the points at each time value.
xs = solutions.T[0::4]
vxs = solutions.T[1::4]
ys = solutions.T[2::4]
vys = solutions.T[3::4]
# Create the figure and the axes objects.
fig = plt.figure(figsize=(5,5))
ax = fig.add_subplot(111, xlim=(-plotlim,plotlim), ylim=(-plotlim,plotlim))
ax.set_aspect("equal")
ax.tick_params(axis='both',
which='both',
bottom=False,
left=False,
labelbottom=False,
labelleft=False) # Hide the ticks and tick labels of each axis.
def anim(i):
"""
Animation function. Used together with the FuncAnimation() function.
i is the frame's number (or the time step).
"""
# Clear the axes elements (delete the points of the previous frame).
# ax.collections = []
# ax.artists = []
ax.lines = []
# Get the point coordinates at the present time step i.
x = np.asarray(xs[:,i])
y = np.asarray(ys[:,i])
# 'Plot with markers' is sometimes faster than a 'scatter plot'.
# See https://pybonacci.org/2014/09/09/microentrada-rendimiento-de-scatterplots-en-matplotlib/
dots = ax.plot(x, y,
color="black",
alpha=1,
linewidth=0,
marker="o",
markersize=3 #,
# markerfacecolor="black"
)
# Scatter alternative
# dots = ax.scatter(x, y, c="black", alpha=0.8, s=5)
return dots
# Make the animation.
anim = FuncAnimation(fig,
func=anim,
frames=nt,
interval=0,
blit=True,
repeat=False)
# Snippet for saving the results.
# anim.save('<Path-where-the-video-will-be-stored>.mp4', writer="ffmpeg", fps=60)
# import os
# os.system("ffmpeg -i <Path-where-the-video-will-be-stored>.mp4 <Path-where-the-video-will-be-stored>.gif")