10_pde.py

Download
python 437 lines 13.0 KB
  1"""
  2Partial Differential Equations (PDEs) - Numerical Solutions
  3
  4This script demonstrates:
  5- Heat equation (parabolic): ∂u/∂t = α ∂²u/∂x²
  6- Wave equation (hyperbolic): ∂²u/∂t² = c² ∂²u/∂x²
  7- Laplace equation (elliptic): ∂²u/∂x² + ∂²u/∂y² = 0
  8- Finite difference methods
  9- Relaxation method for Laplace equation
 10- Time evolution visualization
 11"""
 12
 13import numpy as np
 14
 15try:
 16    import matplotlib.pyplot as plt
 17    from matplotlib import animation
 18    HAS_MATPLOTLIB = True
 19except ImportError:
 20    HAS_MATPLOTLIB = False
 21    print("Warning: matplotlib not available, skipping visualizations")
 22
 23
 24def solve_heat_equation_1d(alpha, L, T, nx, nt, initial_condition, boundary_conditions):
 25    """
 26    Solve 1D heat equation: ∂u/∂t = α ∂²u/∂x²
 27    Using explicit finite difference (FTCS scheme)
 28
 29    boundary_conditions: tuple (left_type, left_val, right_type, right_val)
 30    where type is 'dirichlet' or 'neumann'
 31    """
 32    dx = L / (nx - 1)
 33    dt = T / nt
 34
 35    # Stability criterion: α*dt/dx² <= 0.5
 36    r = alpha * dt / dx**2
 37    if r > 0.5:
 38        print(f"Warning: Stability criterion violated! r={r:.3f} > 0.5")
 39
 40    x = np.linspace(0, L, nx)
 41    u = np.zeros((nt, nx))
 42
 43    # Initial condition
 44    u[0, :] = initial_condition(x)
 45
 46    # Time stepping
 47    for n in range(nt - 1):
 48        for i in range(1, nx - 1):
 49            u[n+1, i] = u[n, i] + r * (u[n, i+1] - 2*u[n, i] + u[n, i-1])
 50
 51        # Boundary conditions
 52        left_type, left_val, right_type, right_val = boundary_conditions
 53
 54        if left_type == 'dirichlet':
 55            u[n+1, 0] = left_val
 56        elif left_type == 'neumann':
 57            u[n+1, 0] = u[n+1, 1] - left_val * dx
 58
 59        if right_type == 'dirichlet':
 60            u[n+1, -1] = right_val
 61        elif right_type == 'neumann':
 62            u[n+1, -1] = u[n+1, -2] + right_val * dx
 63
 64    return x, u
 65
 66
 67def solve_wave_equation_1d(c, L, T, nx, nt, initial_u, initial_v, boundary_conditions):
 68    """
 69    Solve 1D wave equation: ∂²u/∂t² = c² ∂²u/∂x²
 70    Using explicit finite difference
 71
 72    initial_u: initial displacement
 73    initial_v: initial velocity
 74    """
 75    dx = L / (nx - 1)
 76    dt = T / nt
 77
 78    # Stability criterion: c*dt/dx <= 1 (CFL condition)
 79    cfl = c * dt / dx
 80    if cfl > 1:
 81        print(f"Warning: CFL condition violated! CFL={cfl:.3f} > 1")
 82
 83    x = np.linspace(0, L, nx)
 84    u = np.zeros((nt, nx))
 85
 86    # Initial conditions
 87    u[0, :] = initial_u(x)
 88
 89    # First time step using initial velocity
 90    r = (c * dt / dx)**2
 91    for i in range(1, nx - 1):
 92        u[1, i] = u[0, i] + dt * initial_v(x[i]) + 0.5 * r * (u[0, i+1] - 2*u[0, i] + u[0, i-1])
 93
 94    # Apply boundary conditions to first step
 95    left_type, left_val, right_type, right_val = boundary_conditions
 96    if left_type == 'dirichlet':
 97        u[1, 0] = left_val
 98    if right_type == 'dirichlet':
 99        u[1, -1] = right_val
100
101    # Time stepping
102    for n in range(1, nt - 1):
103        for i in range(1, nx - 1):
104            u[n+1, i] = 2*u[n, i] - u[n-1, i] + r * (u[n, i+1] - 2*u[n, i] + u[n, i-1])
105
106        # Boundary conditions
107        if left_type == 'dirichlet':
108            u[n+1, 0] = left_val
109        if right_type == 'dirichlet':
110            u[n+1, -1] = right_val
111
112    return x, u
113
114
115def solve_laplace_2d(nx, ny, max_iter=1000, tol=1e-5, boundary_func=None):
116    """
117    Solve 2D Laplace equation: ∂²u/∂x² + ∂²u/∂y² = 0
118    Using Jacobi relaxation method
119    """
120    u = np.zeros((nx, ny))
121
122    # Apply boundary conditions
123    if boundary_func is not None:
124        # Left and right boundaries
125        for j in range(ny):
126            y = j / (ny - 1)
127            u[0, j] = boundary_func('left', 0, y)
128            u[-1, j] = boundary_func('right', 1, y)
129
130        # Top and bottom boundaries
131        for i in range(nx):
132            x = i / (nx - 1)
133            u[i, 0] = boundary_func('bottom', x, 0)
134            u[i, -1] = boundary_func('top', x, 1)
135
136    # Jacobi iteration
137    for iteration in range(max_iter):
138        u_old = u.copy()
139
140        for i in range(1, nx - 1):
141            for j in range(1, ny - 1):
142                u[i, j] = 0.25 * (u_old[i+1, j] + u_old[i-1, j] +
143                                  u_old[i, j+1] + u_old[i, j-1])
144
145        # Check convergence
146        error = np.max(np.abs(u - u_old))
147        if error < tol:
148            print(f"Converged in {iteration} iterations, error={error:.2e}")
149            break
150
151    return u
152
153
154def solve_poisson_2d(nx, ny, source_func, max_iter=1000, tol=1e-5):
155    """
156    Solve 2D Poisson equation: ∂²u/∂x² + ∂²u/∂y² = f(x,y)
157    Using Jacobi relaxation
158    """
159    dx = 1.0 / (nx - 1)
160    dy = 1.0 / (ny - 1)
161
162    u = np.zeros((nx, ny))
163    f = np.zeros((nx, ny))
164
165    # Compute source term
166    for i in range(nx):
167        for j in range(ny):
168            x = i * dx
169            y = j * dy
170            f[i, j] = source_func(x, y)
171
172    # Jacobi iteration
173    for iteration in range(max_iter):
174        u_old = u.copy()
175
176        for i in range(1, nx - 1):
177            for j in range(1, ny - 1):
178                u[i, j] = 0.25 * (u_old[i+1, j] + u_old[i-1, j] +
179                                  u_old[i, j+1] + u_old[i, j-1] -
180                                  dx**2 * f[i, j])
181
182        # Check convergence
183        error = np.max(np.abs(u - u_old))
184        if error < tol:
185            print(f"Converged in {iteration} iterations")
186            break
187
188    return u
189
190
191# ============================================================================
192# MAIN DEMONSTRATIONS
193# ============================================================================
194
195print("=" * 70)
196print("PARTIAL DIFFERENTIAL EQUATIONS - NUMERICAL SOLUTIONS")
197print("=" * 70)
198
199# Test 1: Heat equation - temperature diffusion
200print("\n1. HEAT EQUATION - 1D DIFFUSION")
201print("-" * 70)
202alpha = 0.01  # Thermal diffusivity
203L = 1.0       # Length
204T = 10.0      # Total time
205nx = 50       # Spatial points
206nt = 500      # Time steps
207
208# Initial condition: Gaussian pulse
209initial_temp = lambda x: np.exp(-50 * (x - 0.5)**2)
210
211# Boundary conditions: Dirichlet (fixed temperature at ends)
212bc = ('dirichlet', 0.0, 'dirichlet', 0.0)
213
214x, u_heat = solve_heat_equation_1d(alpha, L, T, nx, nt, initial_temp, bc)
215
216print(f"Domain: [0, {L}], Time: [0, {T}]")
217print(f"Grid: {nx} × {nt}")
218print(f"Thermal diffusivity α = {alpha}")
219
220dt = T / nt
221dx = L / (nx - 1)
222r = alpha * dt / dx**2
223print(f"Stability parameter r = α*dt/dx² = {r:.4f} (must be ≤ 0.5)")
224
225# Check heat conservation (integral should decrease due to boundary)
226heat_initial = np.trapz(u_heat[0, :], x)
227heat_final = np.trapz(u_heat[-1, :], x)
228print(f"\nHeat content:")
229print(f"  Initial: {heat_initial:.6f}")
230print(f"  Final:   {heat_final:.6f}")
231
232# Test 2: Wave equation - vibrating string
233print("\n2. WAVE EQUATION - VIBRATING STRING")
234print("-" * 70)
235c = 1.0       # Wave speed
236L = 1.0
237T = 2.0
238nx = 100
239nt = 400
240
241# Initial displacement: plucked string
242initial_disp = lambda x: np.where(x < 0.5, 2*x, 2*(1-x))
243initial_vel = lambda x: 0.0  # Released from rest
244
245# Boundary conditions: fixed ends
246bc = ('dirichlet', 0.0, 'dirichlet', 0.0)
247
248x, u_wave = solve_wave_equation_1d(c, L, T, nx, nt, initial_disp, initial_vel, bc)
249
250dt = T / nt
251dx = L / (nx - 1)
252cfl = c * dt / dx
253print(f"Domain: [0, {L}], Time: [0, {T}]")
254print(f"Grid: {nx} × {nt}")
255print(f"Wave speed c = {c}")
256print(f"CFL number = c*dt/dx = {cfl:.4f} (must be ≤ 1)")
257
258# Energy conservation check
259def wave_energy(u_current, u_prev, dx, dt, c):
260    """Compute total energy: kinetic + potential"""
261    # Kinetic energy: (1/2) ∫ (∂u/∂t)² dx
262    u_t = (u_current - u_prev) / dt
263    kinetic = 0.5 * np.trapz(u_t**2, dx=dx)
264
265    # Potential energy: (1/2) c² ∫ (∂u/∂x)² dx
266    u_x = np.gradient(u_current, dx)
267    potential = 0.5 * c**2 * np.trapz(u_x**2, dx=dx)
268
269    return kinetic + potential
270
271energy_initial = wave_energy(u_wave[1, :], u_wave[0, :], dx, dt, c)
272energy_middle = wave_energy(u_wave[nt//2, :], u_wave[nt//2-1, :], dx, dt, c)
273energy_final = wave_energy(u_wave[-1, :], u_wave[-2, :], dx, dt, c)
274
275print(f"\nTotal energy (should be conserved):")
276print(f"  t={0:.2f}: E = {energy_initial:.6f}")
277print(f"  t={T/2:.2f}: E = {energy_middle:.6f}")
278print(f"  t={T:.2f}: E = {energy_final:.6f}")
279
280# Test 3: Laplace equation - steady-state temperature
281print("\n3. LAPLACE EQUATION - 2D STEADY-STATE")
282print("-" * 70)
283nx, ny = 50, 50
284
285# Boundary conditions: heated on top, cold on bottom, insulated sides
286def boundary_laplace(side, x, y):
287    if side == 'top':
288        return 100.0  # Hot
289    elif side == 'bottom':
290        return 0.0    # Cold
291    elif side == 'left':
292        return 50.0 * y  # Linear gradient
293    elif side == 'right':
294        return 50.0 * y
295    return 0.0
296
297u_laplace = solve_laplace_2d(nx, ny, max_iter=5000, tol=1e-6,
298                             boundary_func=boundary_laplace)
299
300print(f"Grid: {nx} × {ny}")
301print(f"Boundary conditions:")
302print(f"  Top:    T = 100°C")
303print(f"  Bottom: T = 0°C")
304print(f"  Sides:  T = 50y°C")
305print(f"\nSolution statistics:")
306print(f"  Min temperature: {np.min(u_laplace):.2f}°C")
307print(f"  Max temperature: {np.max(u_laplace):.2f}°C")
308print(f"  Center temperature: {u_laplace[nx//2, ny//2]:.2f}°C")
309
310# Test 4: Poisson equation with source
311print("\n4. POISSON EQUATION - WITH SOURCE TERM")
312print("-" * 70)
313nx, ny = 50, 50
314
315# Source term: point source at center
316def source_poisson(x, y):
317    r_squared = (x - 0.5)**2 + (y - 0.5)**2
318    return 1000 * np.exp(-100 * r_squared)
319
320u_poisson = solve_poisson_2d(nx, ny, source_poisson, max_iter=5000, tol=1e-6)
321
322print(f"Grid: {nx} × {ny}")
323print(f"Source: Gaussian centered at (0.5, 0.5)")
324print(f"\nSolution statistics:")
325print(f"  Min: {np.min(u_poisson):.4f}")
326print(f"  Max: {np.max(u_poisson):.4f}")
327print(f"  Center: {u_poisson[nx//2, ny//2]:.4f}")
328
329# Visualization
330if HAS_MATPLOTLIB:
331    print("\n" + "=" * 70)
332    print("VISUALIZATIONS")
333    print("=" * 70)
334
335    fig = plt.figure(figsize=(15, 10))
336
337    # Plot 1: Heat equation evolution
338    ax1 = plt.subplot(2, 3, 1)
339    time_indices = [0, nt//4, nt//2, -1]
340    times = [0, T/4, T/2, T]
341
342    for idx, t in zip(time_indices, times):
343        ax1.plot(x, u_heat[idx, :], linewidth=2, label=f't={t:.2f}')
344
345    ax1.set_xlabel('x')
346    ax1.set_ylabel('Temperature')
347    ax1.set_title('Heat Equation: 1D Diffusion')
348    ax1.legend()
349    ax1.grid(True, alpha=0.3)
350
351    # Plot 2: Heat equation space-time plot
352    ax2 = plt.subplot(2, 3, 2)
353    t_vals = np.linspace(0, T, nt)
354    X, T_mesh = np.meshgrid(x, t_vals)
355    im = ax2.contourf(X, T_mesh, u_heat, levels=20, cmap='hot')
356    ax2.set_xlabel('x')
357    ax2.set_ylabel('t')
358    ax2.set_title('Heat Equation: Space-Time')
359    plt.colorbar(im, ax=ax2)
360
361    # Plot 3: Wave equation snapshots
362    ax3 = plt.subplot(2, 3, 3)
363    time_indices = [0, nt//8, nt//4, nt//2]
364    times = [0, T/8, T/4, T/2]
365
366    for idx, t in zip(time_indices, times):
367        ax3.plot(x, u_wave[idx, :], linewidth=2, label=f't={t:.2f}')
368
369    ax3.set_xlabel('x')
370    ax3.set_ylabel('Displacement')
371    ax3.set_title('Wave Equation: Vibrating String')
372    ax3.legend()
373    ax3.grid(True, alpha=0.3)
374    ax3.set_ylim(-1.2, 1.2)
375
376    # Plot 4: Wave equation space-time
377    ax4 = plt.subplot(2, 3, 4)
378    t_vals = np.linspace(0, T, nt)
379    X, T_mesh = np.meshgrid(x, t_vals)
380    im = ax4.contourf(X, T_mesh, u_wave, levels=20, cmap='RdBu_r')
381    ax4.set_xlabel('x')
382    ax4.set_ylabel('t')
383    ax4.set_title('Wave Equation: Space-Time')
384    plt.colorbar(im, ax=ax4)
385
386    # Plot 5: Laplace equation solution
387    ax5 = plt.subplot(2, 3, 5)
388    x_grid = np.linspace(0, 1, nx)
389    y_grid = np.linspace(0, 1, ny)
390    X_grid, Y_grid = np.meshgrid(x_grid, y_grid)
391    im = ax5.contourf(X_grid, Y_grid, u_laplace.T, levels=20, cmap='coolwarm')
392    ax5.set_xlabel('x')
393    ax5.set_ylabel('y')
394    ax5.set_title('Laplace Equation: Steady-State')
395    plt.colorbar(im, ax=ax5)
396    ax5.set_aspect('equal')
397
398    # Plot 6: Poisson equation solution
399    ax6 = plt.subplot(2, 3, 6)
400    im = ax6.contourf(X_grid, Y_grid, u_poisson.T, levels=20, cmap='viridis')
401    ax6.set_xlabel('x')
402    ax6.set_ylabel('y')
403    ax6.set_title('Poisson Equation: With Source')
404    plt.colorbar(im, ax=ax6)
405    ax6.set_aspect('equal')
406
407    plt.tight_layout()
408    plt.savefig('/opt/projects/01_Personal/03_Study/examples/Mathematical_Methods/10_pde.png', dpi=150)
409    print("Saved visualization: 10_pde.png")
410    plt.close()
411
412    # Additional plot: 3D surface for Laplace equation
413    fig = plt.figure(figsize=(12, 5))
414
415    ax1 = fig.add_subplot(121, projection='3d')
416    ax1.plot_surface(X_grid, Y_grid, u_laplace.T, cmap='coolwarm', alpha=0.8)
417    ax1.set_xlabel('x')
418    ax1.set_ylabel('y')
419    ax1.set_zlabel('Temperature')
420    ax1.set_title('Laplace Equation: 3D View')
421
422    ax2 = fig.add_subplot(122, projection='3d')
423    ax2.plot_surface(X_grid, Y_grid, u_poisson.T, cmap='viridis', alpha=0.8)
424    ax2.set_xlabel('x')
425    ax2.set_ylabel('y')
426    ax2.set_zlabel('u')
427    ax2.set_title('Poisson Equation: 3D View')
428
429    plt.tight_layout()
430    plt.savefig('/opt/projects/01_Personal/03_Study/examples/Mathematical_Methods/10_pde_3d.png', dpi=150)
431    print("Saved visualization: 10_pde_3d.png")
432    plt.close()
433
434print("\n" + "=" * 70)
435print("DEMONSTRATION COMPLETE")
436print("=" * 70)