53 lines
1.8 KiB
Python
53 lines
1.8 KiB
Python
import numpy as np
|
|
|
|
def complete_field(G, G_0, f,
|
|
epsilon=1e-6,
|
|
max_iter=5000):
|
|
N = G.shape[0]
|
|
M = G.shape[1]
|
|
q_x_macro = 2*np.pi*np.fft.fftfreq(N, d=1/N)
|
|
q_y_macro = 2*np.pi*np.fft.fftfreq(M, d=1/M)
|
|
Qx, Qy = np.meshgrid(q_x_macro, q_y_macro)
|
|
|
|
|
|
|
|
# Step 0: Initial guess using the simpler Poisson solution
|
|
# ∇²u = -f / G_0 -> u = solve_poisson_fft(f) if we interpret as (∆u + f=0) or similar.
|
|
u_n = solve_poisson_fft(f)
|
|
u_n_tilde = np.fft.fft2(u_n)
|
|
|
|
# Pre-transform for efficiency
|
|
f_tilde = np.fft.fft2(f)
|
|
|
|
for _ in range(max_iter):
|
|
# Step 1: gradient of u_n in Fourier space
|
|
grad_u_n_tilde = (1j * Qx * u_n_tilde, 1j * Qy * u_n_tilde)
|
|
|
|
# Step 2: transform back to real space to get grad(u_n)
|
|
grad_u_n_x = np.fft.ifft2(grad_u_n_tilde[0]).real
|
|
grad_u_n_y = np.fft.ifft2(grad_u_n_tilde[1]).real
|
|
|
|
# Multiply (G - G_0) with grad(u_n) in real space, then transform to frequency space
|
|
temp_x = np.fft.fft2((G - G_0) * grad_u_n_x)
|
|
temp_y = np.fft.fft2((G - G_0) * grad_u_n_y)
|
|
|
|
# Step 3: compute divergence in Fourier space
|
|
# div(...) -> iQx * temp_x + iQy * temp_y
|
|
f_n_tilde = 1j * Qx * temp_x + 1j * Qy * temp_y
|
|
|
|
# Step 4: solve Poisson in Fourier space
|
|
# G_0 ∆ u = -(f_n + f) => ∆ u = -(f_n + f)/G_0
|
|
u_next_tilde = foh.solve_poisson_fft_tilde(-(f_n_tilde + f_tilde) / G_0, Qx, Qy)
|
|
u_next = np.fft.ifft2(u_next_tilde).real
|
|
|
|
# Step 5: check for convergence
|
|
if np.linalg.norm(u_next - np.fft.ifft2(u_n_tilde).real) < epsilon:
|
|
u_n_tilde = u_next_tilde
|
|
break
|
|
|
|
# Step 6: update solution
|
|
u_n_tilde = u_next_tilde
|
|
|
|
# Return final solution in real space
|
|
return np.fft.ifft2(u_n_tilde).real
|