# 这个python脚本基于达朗贝尔方程求解三维电磁场
# 达朗贝尔方程将势写为波动形式，从而从势能导出场
# 由于没有正确设置边界，在边界附近的物理不可靠
# 可能有bug，只测试了点电荷和z方向导线
# Gitee Repo

import numpy as np
import scipy
import torch

torch.set_default_device('cuda' if torch.cuda.is_available() else 'cpu')


L = 2
dx = 0.05
dt = 0.001
c = 1

sigma = (c*dt/dx)**2

x,y,z = torch.meshgrid(torch.arange(-L,L,dx),torch.arange(-L,L,dx),torch.arange(-L,L,dx))
n = x.shape[0]

Amu0 = torch.zeros((n,n,n,4))
Amu1 = torch.zeros((n,n,n,4))
Amu2 = torch.zeros((n,n,n,4))

TICK = 2000
for tick in range(TICK+2):
    diffi = Amu1[2:n,1:n-1,1:n-1,:] - 2*Amu1[1:n-1,1:n-1,1:n-1,:] + Amu1[0:n-2,1:n-1,1:n-1,:]
    diffj = Amu1[1:n-1,2:n,1:n-1,:] - 2*Amu1[1:n-1,1:n-1,1:n-1,:] + Amu1[1:n-1,0:n-2,1:n-1,:]
    diffk = Amu1[1:n-1,1:n-1,2:n,:] - 2*Amu1[1:n-1,1:n-1,1:n-1,:] + Amu1[1:n-1,1:n-1,0:n-2,:]
    
    Amu2[1:n-1,1:n-1,1:n-1,:] = 2*Amu1[1:n-1,1:n-1,1:n-1,:] - Amu0[1:n-1,1:n-1,1:n-1,:] + sigma*(diffi+diffj+diffk)

    jmu = 2*tick/TICK if tick < TICK/2 else 1
    
    # 点电荷
    #Amu2[int(n/2),int(n/2),int(n/2),0] += dt**2/dx**3*jmu
    
    # z轴导线
    #Amu2[int(n/2),int(n/2),1:n-1,3] += dt**2/dx**3*jmu

    Amu2[int(n/2),int(n/2),int(n/2),3] += dt**2/dx**3*np.sin(2*np.pi/1*tick*dt)
    
    print(tick)
    if tick == TICK:
        break
    
    Amu0,Amu1,Amu2=Amu1,Amu2,Amu0

print(Amu2.max())

E = torch.zeros((n,n,n,4))
B = torch.zeros((n,n,n,4))

E[1:n-1,1:n-1,1:n-1,1] = - (Amu2[1:n-1,1:n-1,1:n-1,1] - Amu0[1:n-1,1:n-1,1:n-1,1])/(2*dt) - (Amu1[2:n,1:n-1,1:n-1,0] - Amu1[0:n-2,1:n-1,1:n-1,0])/(2*dx)
E[1:n-1,1:n-1,1:n-1,2] = - (Amu2[1:n-1,1:n-1,1:n-1,2] - Amu0[1:n-1,1:n-1,1:n-1,2])/(2*dt) - (Amu1[1:n-1,2:n,1:n-1,0] - Amu1[1:n-1,0:n-2,1:n-1,0])/(2*dx)
E[1:n-1,1:n-1,1:n-1,3] = - (Amu2[1:n-1,1:n-1,1:n-1,3] - Amu0[1:n-1,1:n-1,1:n-1,3])/(2*dt) - (Amu1[1:n-1,1:n-1,2:n,0] - Amu1[1:n-1,1:n-1,0:n-2,0])/(2*dx)

B[1:n-1,1:n-1,1:n-1,1] =   (Amu1[1:n-1,2:n,1:n-1,3] - Amu1[1:n-1,0:n-2,1:n-1,3])/(2*dx) - (Amu1[1:n-1,1:n-1,2:n,2] - Amu1[1:n-1,1:n-1,0:n-2,2])/(2*dx)
B[1:n-1,1:n-1,1:n-1,2] =   (Amu1[1:n-1,1:n-1,2:n,1] - Amu1[1:n-1,1:n-1,0:n-2,1])/(2*dx) - (Amu1[2:n,1:n-1,1:n-1,3] - Amu1[0:n-2,1:n-1,1:n-1,3])/(2*dx)
B[1:n-1,1:n-1,1:n-1,3] =   (Amu1[2:n,1:n-1,1:n-1,2] - Amu1[0:n-2,1:n-1,1:n-1,2])/(2*dx) - (Amu1[1:n-1,2:n,1:n-1,1] - Amu1[1:n-1,0:n-2,1:n-1,1])/(2*dx)

scipy.io.savemat('EB.mat',{'L':L, 'x':x.cpu().numpy(),'y':y.cpu().numpy(),'z':z.cpu().numpy(),'E':E.cpu().numpy(), 'B':B.cpu().numpy(), 'Amu':Amu2.cpu().numpy()})