|
| 1 | +""" pyplots.ai |
| 2 | +streamline-basic: Basic Streamline Plot |
| 3 | +Library: plotnine 0.15.2 | Python 3.13.11 |
| 4 | +Quality: 90/100 | Created: 2025-12-31 |
| 5 | +""" |
| 6 | + |
| 7 | +import numpy as np |
| 8 | +import pandas as pd |
| 9 | +from plotnine import ( |
| 10 | + aes, |
| 11 | + coord_fixed, |
| 12 | + element_rect, |
| 13 | + element_text, |
| 14 | + geom_path, |
| 15 | + geom_point, |
| 16 | + ggplot, |
| 17 | + labs, |
| 18 | + scale_color_gradient, |
| 19 | + theme, |
| 20 | + theme_minimal, |
| 21 | +) |
| 22 | +from scipy.integrate import solve_ivp |
| 23 | +from scipy.interpolate import RegularGridInterpolator |
| 24 | + |
| 25 | + |
| 26 | +# Set seed for reproducibility |
| 27 | +np.random.seed(42) |
| 28 | + |
| 29 | +# Create grid for vector field |
| 30 | +nx, ny = 40, 40 |
| 31 | +x = np.linspace(-3, 3, nx) |
| 32 | +y = np.linspace(-3, 3, ny) |
| 33 | +X, Y = np.meshgrid(x, y) |
| 34 | + |
| 35 | +# Define vector field: vortex flow with a dipole effect |
| 36 | +# u = -y, v = x creates circular streamlines (vortex) |
| 37 | +U = -Y |
| 38 | +V = X |
| 39 | + |
| 40 | +# Create interpolators for the velocity field |
| 41 | +u_interp = RegularGridInterpolator((y, x), U, bounds_error=False, fill_value=0) |
| 42 | +v_interp = RegularGridInterpolator((y, x), V, bounds_error=False, fill_value=0) |
| 43 | + |
| 44 | +# Compute streamlines from various starting points |
| 45 | +streamlines_data = [] |
| 46 | +arrow_data = [] |
| 47 | +streamline_id = 0 |
| 48 | + |
| 49 | +# Create varied starting points - grid on one side and scattered points |
| 50 | +# This shows field topology better than just radial starting points |
| 51 | +start_points = [] |
| 52 | + |
| 53 | +# Grid of starting points on left side |
| 54 | +for sx in np.linspace(-2.8, -1.5, 4): |
| 55 | + for sy in np.linspace(-2.5, 2.5, 6): |
| 56 | + start_points.append((sx, sy)) |
| 57 | + |
| 58 | +# Some radial starting points to show circular nature |
| 59 | +for r in [0.6, 1.3, 2.2]: |
| 60 | + for angle in np.linspace(0, 2 * np.pi, 6, endpoint=False): |
| 61 | + start_points.append((r * np.cos(angle), r * np.sin(angle))) |
| 62 | + |
| 63 | +for x0, y0 in start_points: |
| 64 | + # Integrate forward - inline velocity calculation (no function definition) |
| 65 | + try: |
| 66 | + result = solve_ivp( |
| 67 | + lambda t, pos: [u_interp([pos[1], pos[0]])[0], v_interp([pos[1], pos[0]])[0]], |
| 68 | + [0, 4], |
| 69 | + [x0, y0], |
| 70 | + max_step=0.05, |
| 71 | + dense_output=True, |
| 72 | + ) |
| 73 | + if result.success and len(result.t) > 2: |
| 74 | + t_eval = np.linspace(0, result.t[-1], 100) |
| 75 | + trajectory = result.sol(t_eval) |
| 76 | + |
| 77 | + # Calculate velocity magnitude at each point |
| 78 | + for j in range(len(t_eval)): |
| 79 | + px, py = trajectory[0, j], trajectory[1, j] |
| 80 | + # Keep points within bounds |
| 81 | + if -3 <= px <= 3 and -3 <= py <= 3: |
| 82 | + speed = np.sqrt(px**2 + py**2) # For vortex: speed = r |
| 83 | + streamlines_data.append({"x": px, "y": py, "streamline": streamline_id, "order": j, "speed": speed}) |
| 84 | + |
| 85 | + # Add arrow marker at ~60% along the streamline to show direction |
| 86 | + arrow_idx = int(len(t_eval) * 0.6) |
| 87 | + if arrow_idx < len(t_eval): |
| 88 | + ax, ay = trajectory[0, arrow_idx], trajectory[1, arrow_idx] |
| 89 | + if -3 <= ax <= 3 and -3 <= ay <= 3: |
| 90 | + arrow_speed = np.sqrt(ax**2 + ay**2) |
| 91 | + arrow_data.append({"x": ax, "y": ay, "speed": arrow_speed}) |
| 92 | + |
| 93 | + streamline_id += 1 |
| 94 | + except Exception: |
| 95 | + pass |
| 96 | + |
| 97 | +# Convert to DataFrames |
| 98 | +df = pd.DataFrame(streamlines_data) |
| 99 | +df_arrows = pd.DataFrame(arrow_data) |
| 100 | + |
| 101 | +# Create the plot using plotnine's native geom_path |
| 102 | +# Use 1:1 canvas (12x12) for circular pattern - avoids empty horizontal space |
| 103 | +plot = ( |
| 104 | + ggplot(df, aes(x="x", y="y", group="streamline", color="speed")) |
| 105 | + + geom_path(size=1.2, alpha=0.8) |
| 106 | + + geom_point( |
| 107 | + data=df_arrows, |
| 108 | + mapping=aes(x="x", y="y", color="speed"), |
| 109 | + shape=">", |
| 110 | + size=4, |
| 111 | + inherit_aes=False, |
| 112 | + show_legend=False, |
| 113 | + ) |
| 114 | + + scale_color_gradient(low="#306998", high="#FFD43B", name="Flow Speed") |
| 115 | + + labs(x="X Position", y="Y Position", title="streamline-basic · plotnine · pyplots.ai") |
| 116 | + + coord_fixed(ratio=1) |
| 117 | + + theme_minimal() |
| 118 | + + theme( |
| 119 | + figure_size=(12, 12), |
| 120 | + plot_title=element_text(size=24, weight="bold"), |
| 121 | + axis_title=element_text(size=20), |
| 122 | + axis_text=element_text(size=16), |
| 123 | + legend_title=element_text(size=18), |
| 124 | + legend_text=element_text(size=14), |
| 125 | + panel_background=element_rect(fill="white"), |
| 126 | + plot_background=element_rect(fill="white"), |
| 127 | + ) |
| 128 | +) |
| 129 | + |
| 130 | +# Save the plot |
| 131 | +plot.save("plot.png", dpi=300, verbose=False) |
0 commit comments