Initial upload: aesthetic-immich-starter
This commit is contained in:
Abderrahim-Mouhcine
@@ -0,0 +1,242 @@
|
||||
import numpy as np
|
||||
from numpy.testing import assert_array_almost_equal
|
||||
import pytest
|
||||
import matplotlib.pyplot as plt
|
||||
from matplotlib.testing.decorators import image_comparison
|
||||
import matplotlib.transforms as mtransforms
|
||||
|
||||
|
||||
def velocity_field():
|
||||
Y, X = np.mgrid[-3:3:100j, -3:3:200j]
|
||||
U = -1 - X**2 + Y
|
||||
V = 1 + X - Y**2
|
||||
return X, Y, U, V
|
||||
|
||||
|
||||
def swirl_velocity_field():
|
||||
x = np.linspace(-3., 3., 200)
|
||||
y = np.linspace(-3., 3., 100)
|
||||
X, Y = np.meshgrid(x, y)
|
||||
a = 0.1
|
||||
U = np.cos(a) * (-Y) - np.sin(a) * X
|
||||
V = np.sin(a) * (-Y) + np.cos(a) * X
|
||||
return x, y, U, V
|
||||
|
||||
|
||||
@image_comparison(['streamplot_startpoints.png'], remove_text=True, style='mpl20',
|
||||
tol=0.003)
|
||||
def test_startpoints():
|
||||
# Test varying startpoints. Also tests a non-default num_arrows argument.
|
||||
X, Y, U, V = velocity_field()
|
||||
start_x, start_y = np.meshgrid(np.linspace(X.min(), X.max(), 5),
|
||||
np.linspace(Y.min(), Y.max(), 5))
|
||||
start_points = np.column_stack([start_x.ravel(), start_y.ravel()])
|
||||
plt.streamplot(X, Y, U, V, start_points=start_points, num_arrows=4)
|
||||
plt.plot(start_x, start_y, 'ok')
|
||||
|
||||
|
||||
@image_comparison(['streamplot_colormap.png'], remove_text=True, style='mpl20',
|
||||
tol=0.022)
|
||||
def test_colormap():
|
||||
X, Y, U, V = velocity_field()
|
||||
plt.streamplot(X, Y, U, V, color=U, density=0.6, linewidth=2,
|
||||
cmap="autumn")
|
||||
plt.colorbar()
|
||||
|
||||
|
||||
@image_comparison(['streamplot_linewidth.png'], remove_text=True, style='mpl20',
|
||||
tol=0.03)
|
||||
def test_linewidth():
|
||||
X, Y, U, V = velocity_field()
|
||||
speed = np.hypot(U, V)
|
||||
lw = 5 * speed / speed.max()
|
||||
ax = plt.figure().subplots()
|
||||
ax.streamplot(X, Y, U, V, density=[0.5, 1], color='k', linewidth=lw, num_arrows=2)
|
||||
|
||||
|
||||
@image_comparison(['streamplot_masks_and_nans.png'],
|
||||
remove_text=True, style='mpl20')
|
||||
def test_masks_and_nans():
|
||||
X, Y, U, V = velocity_field()
|
||||
mask = np.zeros(U.shape, dtype=bool)
|
||||
mask[40:60, 80:120] = 1
|
||||
U[:20, :40] = np.nan
|
||||
U = np.ma.array(U, mask=mask)
|
||||
ax = plt.figure().subplots()
|
||||
with np.errstate(invalid='ignore'):
|
||||
ax.streamplot(X, Y, U, V, color=U, cmap="Blues")
|
||||
|
||||
|
||||
@image_comparison(['streamplot_maxlength.png'],
|
||||
remove_text=True, style='mpl20', tol=0.302)
|
||||
def test_maxlength():
|
||||
x, y, U, V = swirl_velocity_field()
|
||||
ax = plt.figure().subplots()
|
||||
ax.streamplot(x, y, U, V, maxlength=10., start_points=[[0., 1.5]],
|
||||
linewidth=2, density=2)
|
||||
assert ax.get_xlim()[-1] == ax.get_ylim()[-1] == 3
|
||||
# Compatibility for old test image
|
||||
ax.set(xlim=(None, 3.2555988021882305), ylim=(None, 3.078326760195413))
|
||||
|
||||
|
||||
@image_comparison(['streamplot_maxlength_no_broken.png'],
|
||||
remove_text=True, style='mpl20', tol=0.302)
|
||||
def test_maxlength_no_broken():
|
||||
x, y, U, V = swirl_velocity_field()
|
||||
ax = plt.figure().subplots()
|
||||
ax.streamplot(x, y, U, V, maxlength=10., start_points=[[0., 1.5]],
|
||||
linewidth=2, density=2, broken_streamlines=False)
|
||||
assert ax.get_xlim()[-1] == ax.get_ylim()[-1] == 3
|
||||
# Compatibility for old test image
|
||||
ax.set(xlim=(None, 3.2555988021882305), ylim=(None, 3.078326760195413))
|
||||
|
||||
|
||||
@image_comparison(['streamplot_direction.png'],
|
||||
remove_text=True, style='mpl20', tol=0.073)
|
||||
def test_direction():
|
||||
x, y, U, V = swirl_velocity_field()
|
||||
plt.streamplot(x, y, U, V, integration_direction='backward',
|
||||
maxlength=1.5, start_points=[[1.5, 0.]],
|
||||
linewidth=2, density=2)
|
||||
|
||||
|
||||
@image_comparison(['streamplot_integration.png'], style='mpl20', tol=0.05)
|
||||
def test_integration_options():
|
||||
# Linear potential flow over a lifting cylinder
|
||||
n = 50
|
||||
x, y = np.meshgrid(np.linspace(-2, 2, n), np.linspace(-3, 3, n))
|
||||
th = np.arctan2(y, x)
|
||||
r = np.sqrt(x**2 + y**2)
|
||||
vr = -np.cos(th) / r**2
|
||||
vt = -np.sin(th) / r**2 - 1 / r
|
||||
vx = vr * np.cos(th) - vt * np.sin(th) + 1.0
|
||||
vy = vr * np.sin(th) + vt * np.cos(th)
|
||||
|
||||
# Seed points
|
||||
n_seed = 50
|
||||
seed_pts = np.column_stack((np.full(n_seed, -1.75), np.linspace(-2, 2, n_seed)))
|
||||
|
||||
fig, axs = plt.subplots(3, 1, figsize=(6, 14))
|
||||
th_circ = np.linspace(0, 2 * np.pi, 100)
|
||||
for ax, max_val in zip(axs, [0.05, 1, 5]):
|
||||
ax_ins = ax.inset_axes([0.0, 0.7, 0.3, 0.35])
|
||||
for ax_curr, is_inset in zip([ax, ax_ins], [False, True]):
|
||||
ax_curr.streamplot(
|
||||
x,
|
||||
y,
|
||||
vx,
|
||||
vy,
|
||||
start_points=seed_pts,
|
||||
broken_streamlines=False,
|
||||
arrowsize=1e-10,
|
||||
linewidth=2 if is_inset else 0.6,
|
||||
color="k",
|
||||
integration_max_step_scale=max_val,
|
||||
integration_max_error_scale=max_val,
|
||||
)
|
||||
|
||||
# Draw the cylinder
|
||||
ax_curr.fill(
|
||||
np.cos(th_circ),
|
||||
np.sin(th_circ),
|
||||
color="w",
|
||||
ec="k",
|
||||
lw=6 if is_inset else 2,
|
||||
)
|
||||
|
||||
# Set axis properties
|
||||
ax_curr.set_aspect("equal")
|
||||
|
||||
# Set axis limits and show zoomed region
|
||||
ax_ins.set_xlim(-1.2, -0.7)
|
||||
ax_ins.set_ylim(-0.8, -0.4)
|
||||
ax_ins.set_yticks(())
|
||||
ax_ins.set_xticks(())
|
||||
|
||||
ax.set_ylim(-1.5, 1.5)
|
||||
ax.axis("off")
|
||||
ax.indicate_inset_zoom(ax_ins, ec="k")
|
||||
|
||||
fig.tight_layout()
|
||||
|
||||
|
||||
def test_streamplot_limits():
|
||||
ax = plt.axes()
|
||||
x = np.linspace(-5, 10, 20)
|
||||
y = np.linspace(-2, 4, 10)
|
||||
y, x = np.meshgrid(y, x)
|
||||
trans = mtransforms.Affine2D().translate(25, 32) + ax.transData
|
||||
plt.barbs(x, y, np.sin(x), np.cos(y), transform=trans)
|
||||
# The calculated bounds are approximately the bounds of the original data,
|
||||
# this is because the entire path is taken into account when updating the
|
||||
# datalim.
|
||||
assert_array_almost_equal(ax.dataLim.bounds, (20, 30, 15, 6),
|
||||
decimal=1)
|
||||
|
||||
|
||||
def test_streamplot_grid():
|
||||
u = np.ones((2, 2))
|
||||
v = np.zeros((2, 2))
|
||||
|
||||
# Test for same rows and columns
|
||||
x = np.array([[10, 20], [10, 30]])
|
||||
y = np.array([[10, 10], [20, 20]])
|
||||
|
||||
with pytest.raises(ValueError, match="The rows of 'x' must be equal"):
|
||||
plt.streamplot(x, y, u, v)
|
||||
|
||||
x = np.array([[10, 20], [10, 20]])
|
||||
y = np.array([[10, 10], [20, 30]])
|
||||
|
||||
with pytest.raises(ValueError, match="The columns of 'y' must be equal"):
|
||||
plt.streamplot(x, y, u, v)
|
||||
|
||||
x = np.array([[10, 20], [10, 20]])
|
||||
y = np.array([[10, 10], [20, 20]])
|
||||
plt.streamplot(x, y, u, v)
|
||||
|
||||
# Test for maximum dimensions
|
||||
x = np.array([0, 10])
|
||||
y = np.array([[[0, 10]]])
|
||||
|
||||
with pytest.raises(ValueError, match="'y' can have at maximum "
|
||||
"2 dimensions"):
|
||||
plt.streamplot(x, y, u, v)
|
||||
|
||||
# Test for equal spacing
|
||||
u = np.ones((3, 3))
|
||||
v = np.zeros((3, 3))
|
||||
x = np.array([0, 10, 20])
|
||||
y = np.array([0, 10, 30])
|
||||
|
||||
with pytest.raises(ValueError, match="'y' values must be equally spaced"):
|
||||
plt.streamplot(x, y, u, v)
|
||||
|
||||
# Test for strictly increasing
|
||||
x = np.array([0, 20, 40])
|
||||
y = np.array([0, 20, 10])
|
||||
|
||||
|
||||
def test_streamplot_integration_params():
|
||||
x = np.array([[10, 20], [10, 20]])
|
||||
y = np.array([[10, 10], [20, 20]])
|
||||
u = np.ones((2, 2))
|
||||
v = np.zeros((2, 2))
|
||||
|
||||
err_str = "The value of integration_max_step_scale must be > 0, got -0.5"
|
||||
with pytest.raises(ValueError, match=err_str):
|
||||
plt.streamplot(x, y, u, v, integration_max_step_scale=-0.5)
|
||||
|
||||
err_str = "The value of integration_max_error_scale must be > 0, got 0.0"
|
||||
with pytest.raises(ValueError, match=err_str):
|
||||
plt.streamplot(x, y, u, v, integration_max_error_scale=0.0)
|
||||
|
||||
|
||||
def test_streamplot_inputs(): # test no exception occurs.
|
||||
# fully-masked
|
||||
plt.streamplot(np.arange(3), np.arange(3),
|
||||
np.full((3, 3), np.nan), np.full((3, 3), np.nan),
|
||||
color=np.random.rand(3, 3))
|
||||
# array-likes
|
||||
plt.streamplot(range(3), range(3),
|
||||
np.random.rand(3, 3), np.random.rand(3, 3))
|
||||
Reference in New Issue
Block a user