-
Notifications
You must be signed in to change notification settings - Fork 0
/
streamtube_examples_1.py
59 lines (37 loc) · 1.36 KB
/
streamtube_examples_1.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
# -*- coding: utf-8 -*-
"""
Created on Thu Sep 10 23:07:52 2020
@author: usuario
"""
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from scipy.integrate import odeint
from mpl_toolkits.mplot3d import Axes3D
import streamtubes as st # Import the defined module located in the working directory.
def f(x, y, z):
"""
Arbitrary function defining the radius
of the tube at a given point (x, y, z).
"""
return np.cos(y)**2 + 0.2
# Parameters.
lim = 5 # The plot will have the limits [-lim, lim] on each axis.
num_polygons = 100 # Number of sections that compose each streamtube.
num_sides = 10 # Number of sides of the regular polygon defining the sections.
# Generate the path of the tube.
x = np.linspace(-lim,lim,num_polygons)
y = np.sin(x)
z = np.cos(x)
# Obtain the radius of the tube at each point in its trajectory.
rs = f(x, y, z)
# Plot the results.
plt.close("all")
fig = plt.figure(figsize=(6,6))
ax = fig.add_subplot(111, projection="3d", xlim=(-lim,lim), ylim=(-lim,lim), zlim=(-lim,lim))
ax.set_xlabel("x")
ax.set_ylabel("y")
ax.set_zlabel("z")
st.plot_streamtube(ax, x, y, z, rs, num_sides=num_sides, color="C0", alpha=0.4, linewidths=0.5, cmap_name="RdYlBu_r", vmin=np.min(rs), vmax=np.max(rs))
plt.tight_layout()
plt.show()