plotting streamlines with python - python

I am trying to plot streamlines and velocity potential for basic potential flows (uniform, source/sink, vortex, etc.)
I am just starting out with python and so I'm a little confused. I am following this guide..
I can plot the streamlines for the flow around a cylinder using this function
def cylinder_stream_function(U=1, R=1):
r = sympy.sqrt(x**2 + y**2)
theta = sympy.atan2(y, x)
return U * (r - R**2 / r) * sympy.sin(theta)
and it works. But when I change the return statement to
return U * r * sympy.cos(theta)
for uniform flow I get the following error
Traceback (most recent call last):
File "test.py", line 42, in
<module>
plot_streamlines(ax, u, v)
File "test.py", line 32, in plot_streamlines
ax.streamplot(X, Y, u(X, Y), v(X, Y), color='cornflowerblue')
File "/usr/local/lib/python3.6/site-packages/matplotlib/__init__.py",
line 1710, in inner
return func(ax, *args, **kwargs)
File "/usr/local/lib/python3.6/site-packages/matplotlib/axes/_axes.py",
line 4688, in streamplot
integration_direction=integration_direction)
File "/usr/local/lib/python3.6/site-packages/matplotlib/streamplot.py",
line 136, in streamplot
if (u.shape != grid.shape) or (v.shape != grid.shape):
AttributeError: 'int' object has no attribute 'shape'
I checked the type for the return object and it is <class 'sympy.core.mul.Mul'> with the first return statement and <class 'sympy.core.symbol.Symbol'> with the second one. Maybe this is related to why it doesn't work but I'm not sure how?
I plot the streamlines with the following
import numpy as np
import matplotlib.pyplot as plt
import sympy
from sympy.abc import x, y
def uniform_flow_stream_function(U=1):
r = sympy.sqrt(x**2 + y**2)
theta = sympy.atan2(y, x)
return U * r * sympy.sin(theta)
def velocity_field(psi):
u = sympy.lambdify((x, y), psi.diff(y), 'numpy')
v = sympy.lambdify((x, y), -psi.diff(x), 'numpy')
return u, v
def plot_streamlines(ax, u, v, xlim=(-4, 4), ylim=(-4, 4)):
x0, x1 = xlim
y0, y1 = ylim
# create a grid of values
Y, X = np.ogrid[y0:y1:100j, x0:x1:100j]
ax.streamplot(X, Y, u(X, Y), v(X, Y), color='cornflowerblue')
psi = uniform_flow_stream_function()
u, v = velocity_field(psi)
fig, ax = plt.subplots(figsize=(5, 5))
plot_streamlines(ax, u, v)
plt.show()
Can someone please help me understand why this doesn't work and how I can get it to work? Thank you!

The reason that this does not work is because of the class difference. Your function U * r * sympy.cos(theta)=y. This means that you are returning a function of only y. Therefore your -psi.diff(x)=0 and you get a integer for v.
It is impossible to plot streamlines with 1D data. So in order to plot streamlines in 2D you have to have both x and y in your uniform_flow_stream_function.

Related

Error, trying to plot a graph of solution to a Second Order ODE using Euler Method

I am doing a project, where I want to use Euler's Method to show a solution to this Second Order Different equation
0=y''+y'+9.81y
So I started by changing the second-order into a system of first-order equations
y'=u, u'=f(t,y,u)
With initial condition
y(0)=180, u(0)=0
So I get two equation in the end
y[n + 1] = y[n] + u[n] * (t[n + 1] - t[n]), u[n + 1] = u[n] + f(u[0], y[n], t[0]) * (t[n + 1] - t[n])
This is my code
import numpy as np
import matplotlib.pyplot as plt
def odeEuler(f, y0, u0, t):
y = np.zeros(len(t))
u = np.zeros(len(t))
y[0] = y0
u[0] = u0
for n in range(0, len(t) - 1):
y[n + 1] = y[n] + u[n] * (t[n + 1] - t[n])
u[n + 1] = u[n] + f(u[0], y[n], t[0]) * (t[n + 1] - t[n])
return y, u
t = np.linspace(0, 100)
y0 = 180
u0 = 0
f = lambda u, y, t: -9.81 * y - u
y = odeEuler(f, y0, u0, t)
plt.plot(t, y, 'b.-')
plt.legend(['Euler'])
plt.axis([0, 100, 0, 200])
plt.grid(True)
plt.show()
However, when I run the code, it give me the error
Traceback (most recent call last):
File "/Users/huangy15/PycharmProjects/Draft/Damped Driven Pendulum/Praying this works.py", line 22, in <module>
plt.plot(t, y, 'b.-')
File "/Users/huangy15/PycharmProjects/Draft/Damped Driven Pendulum/venv/lib/python3.7/site-packages/matplotlib/pyplot.py", line 3021, in plot
**({"data": data} if data is not None else {}), **kwargs)
File "/Users/huangy15/PycharmProjects/Draft/Damped Driven Pendulum/venv/lib/python3.7/site-packages/matplotlib/axes/_axes.py", line 1605, in plot
lines = [*self._get_lines(*args, data=data, **kwargs)]
File "/Users/huangy15/PycharmProjects/Draft/Damped Driven Pendulum/venv/lib/python3.7/site-packages/matplotlib/axes/_base.py", line 315, in __call__
yield from self._plot_args(this, kwargs)
File "/Users/huangy15/PycharmProjects/Draft/Damped Driven Pendulum/venv/lib/python3.7/site-packages/matplotlib/axes/_base.py", line 501, in _plot_args
raise ValueError(f"x and y must have same first dimension, but "
ValueError: x and y must have same first dimension, but have shapes (50,) and (2, 50)
Can anyone help me check if my idea works, and if not, what other approaches I can take? Thanks!
Your odeEuler function is returning two variables: y and u, you need to select just one of them to plot.
Store the output of the function in two different variables will solve the problem:
y, u = odeEuler(f, y0, u0, t)

How to use interpolation with solve_ivp?

I am trying to make some code that calculates the trajectory of a particle, given its initial position and the velocity field. Here is an example of a working bit of code that nearly does what I want:
import numpy as np
from scipy.integrate import solve_ivp
import matplotlib.pyplot as plt
def velocity(t, coord):
# velocity gives the velocity of the particle
# coord gives the coordinate of the particle where
# coord[0] = x, coord[1] = y and coord[2] = z
return [coord[0], coord[1], -2 * coord[2]]
n_points = 100
lim = 5
coord0 = [2, 2, 5]
t = np.linspace(0, lim, n_points)
sol = solve_ivp(velocity, [0, lim], t_eval = t, y0 = coord0)
fig = plt.figure()
ax = fig.add_subplot(111, projection = '3d')
ln = ax.plot3D(sol.y[0], sol.y[1], sol.y[2])
ax.set_xlim((0, lim))
ax.set_ylim((0, lim))
ax.set_zlim((0, lim))
plt.savefig('particle_trajectory.png')
Here is the outputted figure:
I want the code to calculate the particle's trajectory if I only have an array describing the velocity field. My plan is to use interpolation so I can calculate the particle's trajectory using a similar method to that shown above. Here is my attempt to do this:
import numpy as np
from scipy.interpolate import RegularGridInterpolator
from scipy.integrate import solve_ivp
import matplotlib.pyplot as plt
def velocity(t, coord):
# velocity gives the velocity of the particle
# coord gives the coordinate of the particle where
# coord[0] = x, coord[1] = y and coord[2] = z
return [vx_interp([coord[0], coord[1], coord[2]]), \
vy_interp([coord[0], coord[1], coord[2]]), \
vz_interp([coord[0], coord[1], coord[2]])]
n_points = 100
lim = 5
# Create an example of a velocity array
x = np.linspace(0, lim, n_points)
y = np.linspace(0, lim, n_points)
z = np.linspace(0, lim, n_points)
X, Y, Z = np.meshgrid(x, y, z)
vx = X
vy = Y
vz = -2 * Z
vx_interp = RegularGridInterpolator((x, y, z), vx)
vy_interp = RegularGridInterpolator((x, y, z), vy)
vz_interp = RegularGridInterpolator((x, y, z), vz)
coord0 = [2, 2, 5]
t_range = np.linspace(0, 0.2 * lim, n_points)
sol = solve_ivp(velocity, [0, 2 * lim], t_eval = t_range, y0 = coord0)
But this gives the following error:
File "/home/lex/.local/lib/python3.8/site-packages/scipy/integrate/_ivp/ivp.py", line 542, in solve_ivp
solver = method(fun, t0, y0, tf, vectorized=vectorized, **options)
File "/home/lex/.local/lib/python3.8/site-packages/scipy/integrate/_ivp/rk.py", line 96, in __init__
self.h_abs = select_initial_step(
File "/home/lex/.local/lib/python3.8/site-packages/scipy/integrate/_ivp/common.py", line 111, in select_initial_step
f1 = fun(t0 + h0 * direction, y1)
File "/home/lex/.local/lib/python3.8/site-packages/scipy/integrate/_ivp/base.py", line 138, in fun
return self.fun_single(t, y)
File "/home/lex/.local/lib/python3.8/site-packages/scipy/integrate/_ivp/base.py", line 20, in fun_wrapped
return np.asarray(fun(t, y), dtype=dtype)
File "Python/low_and_lou/stack2.py", line 10, in velocity
return [vx_interp([coord[0], coord[1], coord[2]]), \
File "/home/lex/.local/lib/python3.8/site-packages/scipy/interpolate/interpolate.py", line 2509, in __call__
raise ValueError("One of the requested xi is out of bounds "
ValueError: One of the requested xi is out of bounds in dimension 2
I think the problem is that I am using a 'RegularGridInterpolator' inside of the solve_ivp routine. Do you know how to fix this problem?

Plotting vector field for first order differential equation

I'm trying to plot the direction fields for a simple velocity equation. I understand what I have to do when I'm working with two variables. I can understand the vector I have to create, but I don't understand how to do it for only one variable. My program is:
import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import odeint
def modelo2(y, t):
dydt = 32 - 0.16 * y
return dydt
t0 = 0 ; tf = 25 ; h = 0.1
t = np.arange(t0,tf+h,h)
for y0 in np.arange(0, 400, 25):
y = odeint(modelo2,y0,t )
plt.plot(t,y,'b')
x = np.arange(0, 400, 20)
z = np.arange(0, 400, 20)
X, Z = np.meshgrid(x, z)
U = modelo2(X,t)
V = modelo2 (Z, t)
plt.quiver(X, Z, U, V, scale = 70)
plt.quiver(X, Z, U, V, scale = 60)
plt.xlabel('time')
plt.ylabel('y(t)')
plt.axis([0,20,0, 500])
plt.show()
I get this
When I expect something like this
Can someone explain what I'm doing wrong?
Change this
U = modelo2(X,t)
V = modelo2 (Z, t)
to this
U = 1.0
V = modelo2(Z, None)
N = np.sqrt(U**2 + V**2)
U /= N
V /= N
As you can see you defined U wrong. Diving both U and V by N is necessary to normalise the magnitude of the vectors, otherwise their length in the plot will vary according to the strength of the field at each point. Just set U = np.ones(Z.shape) and don't divide either by N to see what I'm talking about.
Secondly, you need to set the following argument in plt.quiver()
plt.quiver(X, Z, U, V, angles='xy')
From the docs:
angles : {'uv', 'xy'} or array-like, optional, default: 'uv'
Method for determining the angle of the arrows.
- 'uv': The arrow axis aspect ratio is 1 so that
if *U* == *V* the orientation of the arrow on the plot is 45 degrees
counter-clockwise from the horizontal axis (positive to the right).
Use this if the arrows symbolize a quantity that is not based on
*X*, *Y* data coordinates.
- 'xy': Arrows point from (x, y) to (x+u, y+v).
Use this for plotting a gradient field, for example.
- Alternatively, arbitrary angles may be specified explicitly as an array
of values in degrees, counter-clockwise from the horizontal axis.
In this case *U*, *V* is only used to determine the length of the
arrows.
Note: inverting a data axis will correspondingly invert the
arrows only with ``angles='xy'``.
All in all, your code should look like this (with some minor variable name edits):
def modelo2(y, t):
dydt = 32 - 0.16 * y
return dydt
t0, tf, h = 0, 25, 0.1
t = np.arange(t0, tf+h, h)
ymin, ymax, ystep = 0, 400, 25
y = np.arange(ymin, ymax+ystep, ystep)
for y0 in y:
line = odeint(modelo2, y0, t)
plt.plot(t, line, 'b')
x = np.linspace(t0, tf, 20)
X, Y = np.meshgrid(x, y)
U = 1
V = modelo2(Y, None)
N = np.sqrt(U**2 + V**2)
U /= N
V /= N
plt.quiver(X, Y, U, V, angles='xy')
plt.xlabel('time')
plt.ylabel('y(t)')
plt.axis([t0, tf, ymin, ymax])
plt.show()
Result

How can I draw more lines?(python sympy)

when I type,
from sympy import *
from sympy.plotting import *
from sympy.plotting.plot import *
x, y = symbols('x y')
f = Function('f')
g = Function('g')
f = 1/((x+0.3)**2 + y**2) - 1/((x-0.3)**2 + y**2 )
g = (x+0.3)/sqrt((x+0.3)**2 + y**2) - (x-0.3)/sqrt((x-0.3)**2 + y**2)
p0 = Plot(ContourSeries(f,(x,-1.5,1.5),(y,-1.5,1.5)))
p1 = Plot(ContourSeries(g,(x,-1.5,1.5),(y,-1.5,1.5)))
p0.show()
p1.show()
p0 shows like first picture. The number of line is few.
I want to draw more line like second picture.
What is the solution?
The ContourSeries class doesn't expose the right info to do this but it is easy to extend. I called the parameter levels, is passed directly to matplotlib.
class MyContourSeries(ContourSeries):
def __init__(self, expr, var_start_end_x, var_start_end_y, **kwargs):
super(MyContourSeries, self).__init__(expr, var_start_end_x, var_start_end_y)
self.nb_of_points_x = kwargs.get('nb_of_points_x', 50)
self.nb_of_points_y = kwargs.get('nb_of_points_y', 50)
self.levels = kwargs.get('levels', 5)
def get_meshes(self):
mesh_x, mesh_y, f = super().get_meshes()
return (mesh_x, mesh_y, f, self.levels)
This should work with sympy 1.2 and 1.3 at least. Sympy plotting is using matplotlibs countour for this serie (https://github.com/sympy/sympy/blob/master/sympy/plotting/plot.py#L909)
elif s.is_contour:
self.ax.contour(*s.get_meshes())
which signature is matplotlib.pyplot.contour([X, Y,] Z, [levels], **kwargs)
As with the matplotlib contour function, you can used a fixed number of levels
p0 = Plot(MyContourSeries(f, (x, -1.5, 1.5), (y, -1.5, 1.5),
nb_of_points_x=50, nb_of_points_y=50, levels=100))
p0.show()
or pass the levels directly
p0 = Plot(MyContourSeries(f, (x, -1.5, 1.5), (y, -1.5, 1.5),
nb_of_points_x=50, nb_of_points_y=50, levels=np.arange(-50,50)))
p0.show()

Extracting 1D ellipse from 2D image

I've trying to simulate a 2D Sérsic profile and then testing an extraction routine on it. However, when I do a test by extracting all the points lying along an ellipse supposedly aligned with an image, I get a periodic function. It is meant to be a straight line since all points along the ellipse should have equal intensity, although there will be a small amount of deviation due to rounding errors in the rough coordinate estimation (get_I()).
from __future__ import division
import numpy as np
import matplotlib.pyplot as plt
from scipy.interpolate import NearestNDInterpolator
def rotate(x, y, angle):
x1 = x*np.cos(angle) + y*np.sin(angle)
y1 = y*np.cos(angle) - x*np.sin(angle)
return x1, y1
def sersic_1d(R, mu0, h, n, zp=0):
exponent = (R / h) ** (1 / n)
I0 = np.exp((zp - mu0) / 2.5)
return I0 * np.exp(-1.* exponent)
def sersic_2d(x, y, e, i, mu0, h, n, zp=0):
xp, yp = rotate(x, y, i)
alpha = np.arctan2(yp, xp * (1-e))
a = xp / np.cos(alpha)
b = a * (1 - e)
# R2 = (a*a) + ((1 - (e*e)) * yp*yp)
return sersic_1d(a, mu0, h, n, zp)
def ellipse(x0, y0, a, e, i, theta):
b = a * (1 - e)
x = a * np.cos(theta)
y = b * np.sin(theta)
x, y = rotate(x, y, i)
return x + x0, y + y0
def get_I(x, y, Z):
return Z[np.round(x).astype(int), np.round(y).astype(int)]
if __name__ == '__main__':
n = np.linspace(-100,100,1000)
nx, ny = np.meshgrid(n, n)
Z = sersic_2d(nx, ny, 0.5, 0., 0, 50, 1, 25)
theta = np.linspace(0, 2*np.pi, 1000.)
a = 100.
e = 0.5
i = np.pi / 4.
x, y = ellipse(0, 0, a, e, i, theta)
I = get_I(x, y, Z)
plt.plot(I)
# plt.imshow(Z)
plt.show()
However, What I actually get is a massive periodic function. I've checked the alignment and it's correct and the float-> int rounding errors can't account for this kind of shift?
Any ideas?
There are two things that strike me as odd, one of which for sure is not what you wanted, the other I'm not sure about because astronomy is not my field of expertise.
The first is in your function get_I:
def get_I(x, y, Z):
return Z[np.round(x).astype(int), np.round(y).astype(int)]
When you call that function, x an y outline an ellipse, with its center at the origin (0,0). That means x and y both become negative at some point. The indexing you perfom in that function will then take values from the array's last elements, because Z[0,0] is in fact the top left corner of the image (which you plotted, but commented), while Z[-1, -1] is the bottom right corner. What you want is to take the values of Z that are on the ellipse contour, but both have to have the same center. To do that, you would first make sure you use an uneven amount of samples for n (which ultimately defines the shape of Z) and second, you would add an indexing offset:
def get_I(x, y, Z):
offset = Z.shape[0]//2
return Z[np.round(y).astype(int) + offset, np.round(x).astype(int) + offset]
...
n = np.linspace(-100,100,1001) # changed from 1000 to 1001 to ensure a point of origin is present and that the image exhibits point symmetry
Also notice that I changed the order of y and x in get_I: that's because you first index along the rows (for which we usually take the y-coordinate) and only then along the columns (which map to the x-coordinate in most conventions).
The second item that struck me as unusual is that your ellipse has its axes at an angle of pi/4 with respect to the horizontal axis, whereas your sersic (which maps to the 2D array of Z) does not have a tilt at all.
Changing all that, I end up with this code:
from __future__ import division
import numpy as np
import matplotlib.pyplot as plt
def rotate(x, y, angle):
x1 = x*np.cos(angle) + y*np.sin(angle)
y1 = y*np.cos(angle) - x*np.sin(angle)
return x1, y1
def sersic_1d(R, mu0, h, n, zp=0):
exponent = (R / h) ** (1 / n)
I0 = np.exp((zp - mu0) / 2.5)
return I0 * np.exp(-1.* exponent)
def sersic_2d(x, y, e, ang, mu0, h, n, zp=0):
xp, yp = rotate(x, y, ang)
alpha = np.arctan2(yp, xp * (1-e))
a = xp / np.cos(alpha)
b = a * (1 - e)
return sersic_1d(a, mu0, h, n, zp)
def ellipse(x0, y0, a, e, i, theta):
b = a * (1 - e) # half of a
x = a * np.cos(theta)
y = b * np.sin(theta)
x, y = rotate(x, y, i) # rotated by 45deg
return x + x0, y + y0
def get_I(x, y, Z):
offset = Z.shape[0]//2
return Z[np.round(y).astype(int) + offset, np.round(x).astype(int) + offset]
#return Z[np.round(y).astype(int), np.round(x).astype(int)]
if __name__ == '__main__':
n = np.linspace(-100,100,1001) # changed
nx, ny = np.meshgrid(n, n)
ang = 0;#np.pi / 4.
Z = sersic_2d(nx, ny, 0.5, ang=0, mu0=0, h=50, n=1, zp=25)
f, ax = plt.subplots(1,2)
dn = n[1]-n[0]
ax[0].imshow(Z, cmap='gray', aspect='equal', extent=[-100-dn/2, 100+dn/2, -100-dn/2, 100+dn/2])
theta = np.linspace(0, 2*np.pi, 1000.)
a = 20. # decreased long axis of ellipse to see the intensity-map closer to the "center of the galaxy"
e = 0.5
x, y = ellipse(0,0, a, e, ang, theta)
I = get_I(x, y, Z)
ax[0].plot(x,y) # easier to see where you want the intensities
ax[1].plot(I)
plt.show()
and this image:
The intensity variations look like quantisation noise to me, with the exception of the peaks, which are due to the asymptote in sersic_1d.

Categories