jwt625/BPM

1import numpy as np2 3# Global factors; these might be computed more dynamically in a full implementation.4laplacian_factor = None5index_factor = None6 7def compute_dE_dz(E_slice, n_r2_slice, dx, n0, sigma_x, k0):8    """9    Compute the derivative dE/dz using the BPM equation:10    11      ∂E/∂z = (i/(2 k0 n0)) (∂^2 E/∂x^2) + i (k0/(2 n0)) [n_r^2 - n0^2] E - sigma(x) E12    """13    # Finite-difference Laplacian in x14    laplacian_E = (np.roll(E_slice, 1, axis=0) - 2 * E_slice + np.roll(E_slice, -1, axis=0)) / dx**215    laplacian_term = (1j / (2 * k0 * n0)) * laplacian_E16    index_term = 1j * (k0 / (2 * n0)) * (n_r2_slice - n0**2) * E_slice17    damping_term = - sigma_x * E_slice18    return laplacian_term + index_term + damping_term19 20def run_bpm(E, n_r2, x, z, dx, dz, n0, sigma_x, wavelength):21    """22    Run the BPM propagation using an RK4 integrator.23    24    Parameters:25      E: initial field (2D array, shape (len(x), len(z)); only E[:,0] is used)26      n_r2: refractive index squared distribution (2D array, shape (len(x), len(z)))27      x, z: transverse and propagation coordinates28      dx, dz: grid spacings in x and z29      n0: background refractive index30      sigma_x: 1D array for PML damping in x31      wavelength: wavelength in um32    33    Returns:34      E: propagated field (2D array)35    """36    k0 = 2 * np.pi / wavelength37    Nz = len(z)38    for zi in range(1, Nz):39        E_prev = E[:, zi-1]40        n_r2_slice = n_r2[:, zi-1]41        k1 = dz * compute_dE_dz(E_prev, n_r2_slice, dx, n0, sigma_x, k0)42        k2 = dz * compute_dE_dz(E_prev + k1/2, n_r2_slice, dx, n0, sigma_x, k0)43        k3 = dz * compute_dE_dz(E_prev + k2/2, n_r2_slice, dx, n0, sigma_x, k0)44        k4 = dz * compute_dE_dz(E_prev + k3, n_r2_slice, dx, n0, sigma_x, k0)45        E[:, zi] = E_prev + (k1 + 2*k2 + 2*k3 + k4) / 646    return E47