DeepGPR 0.0.1__py3-none-any.whl

DeepGPR/compute2.py

@@ -0,0 +1,326 @@
+import torch
+import ctypes
+from . import c_lib
+from .common import initialization,build_pml_phi,create_or_separate,buildpmlcoeffs,check_tensors_for_nan_inf
+
+def compute(device, dx=None, dt=None, 
+            source_amplitudes=None,
+            source_location=None, 
+            receiver_location=None, 
+            er=None, se=None,mr=None, 
+            E=None,H=None,
+            PML=None,
+            pmlthick=10, source_direction=2, reciever_direction=2,
+            model_gradient_sampling_interval=1,
+            use_async_offload=False):
+            
+    er,se,nx,ny,nz,nt,nstep,nsr,nrx,ere,see,mr,mode,dtype,pmlthick,source_amplitudes=initialization(device,er,se,mr,source_amplitudes,source_location,receiver_location,dx,dt,pmlthick)
+
+    Ex,Ey,Ez=create_or_separate(E,nx,ny,nz,nstep,device,dtype)
+    Hx,Hy,Hz=create_or_separate(H,nx,ny,nz,nstep,device,dtype)
+
+    x0,xm,y0,ym,z0,zm,x01,x02,xm1,xm2,y01,y02,ym1,ym2,z01,z02,zm1,zm2=buildpmlcoeffs(er,mr,dt,dx,nx,ny,nz,pmlthick,device,dtype)
+
+    x0EPhi1,x0EPhi2,x0HPhi1,x0HPhi2,xmEPhi1,xmEPhi2,xmHPhi1,xmHPhi2,y0EPhi1,y0EPhi2,y0HPhi1,y0HPhi2,ymEPhi1,ymEPhi2,ymHPhi1,ymHPhi2,z0EPhi1,z0EPhi2,z0HPhi1,z0HPhi2,zmEPhi1,zmEPhi2,zmHPhi1,zmHPhi2=build_pml_phi(x0,xm,y0,ym,z0,zm,nstep,PML,device)
+
+    Ex,Ey,Ez,Hx,Hy,Hz,x0EPhi1,x0EPhi2,x0HPhi1,x0HPhi2,xmEPhi1,xmEPhi2,xmHPhi1,xmHPhi2,y0EPhi1,y0EPhi2,y0HPhi1,y0HPhi2,ymEPhi1,ymEPhi2,ymHPhi1,ymHPhi2,z0EPhi1,z0EPhi2,z0HPhi1,z0HPhi2,zmEPhi1,zmEPhi2,zmHPhi1,zmHPhi2,Eall,receiver_amplitudes = DeepGPR.apply(
+        er, se,Ex,Ey,Ez,Hx,Hy,Hz,x0EPhi1,x0EPhi2,x0HPhi1,x0HPhi2,xmEPhi1,xmEPhi2,xmHPhi1,xmHPhi2,y0EPhi1,y0EPhi2,y0HPhi1,y0HPhi2,ymEPhi1,ymEPhi2,ymHPhi1,ymHPhi2,z0EPhi1,z0EPhi2,z0HPhi1,z0HPhi2,zmEPhi1,zmEPhi2,zmHPhi1,zmHPhi2, mr,dx,nx,ny,nz,dt,nt,nstep,source_amplitudes,source_location,receiver_location,pmlthick,nsr,nrx,device,dtype,x0,xm,y0,ym,z0,zm,x01,x02,xm1,xm2,y01,y02,ym1,ym2,z01,z02,zm1,zm2,ere,see,source_direction, reciever_direction, model_gradient_sampling_interval, use_async_offload)
+
+    return Eall,(Ex,Ey,Ez),(Hx,Hy,Hz),(x0EPhi1,x0EPhi2,x0HPhi1,x0HPhi2,xmEPhi1,xmEPhi2,xmHPhi1,xmHPhi2,y0EPhi1,y0EPhi2,y0HPhi1,y0HPhi2,ymEPhi1,ymEPhi2,ymHPhi1,ymHPhi2,z0EPhi1,z0EPhi2,z0HPhi1,z0HPhi2,zmEPhi1,zmEPhi2,zmHPhi1,zmHPhi2),receiver_amplitudes
+
+
+class DeepGPR(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, er, se,Ex,Ey,Ez,Hx,Hy,Hz,x0EPhi1,x0EPhi2,x0HPhi1,x0HPhi2,xmEPhi1,xmEPhi2,xmHPhi1,xmHPhi2,y0EPhi1,y0EPhi2,y0HPhi1,y0HPhi2,ymEPhi1,ymEPhi2,ymHPhi1,ymHPhi2,z0EPhi1,z0EPhi2,z0HPhi1,z0HPhi2,zmEPhi1,zmEPhi2,zmHPhi1,zmHPhi2, mr,dx,nx,ny,nz,dt,
+                nt,nstep,source_amplitudes,source_location,receiver_location,
+                pmlthick,nsr,nrx,device,dtype,x0,xm,
+                y0,ym,z0,zm,x01,x02,xm1,xm2,
+                y01,y02,ym1,ym2,z01,z02,zm1,zm2,
+                ere,see,source_direction, reciever_direction, 
+                model_gradient_sampling_interval, use_async_offload):
+        
+        source_amplitudes = source_amplitudes.contiguous()
+        source_location=source_location.to(torch.int32).contiguous()
+        receiver_location=receiver_location.to(torch.int32).contiguous()
+        ctx.save_for_backward(er, se, mr,receiver_location,x0,xm,y0,ym,z0,zm,x01,x02,xm1,xm2,y01,y02,ym1,ym2,z01,z02,zm1,zm2,ere,see)
+        ctx.dx=dx
+        ctx.nx=nx
+        ctx.ny=ny
+        ctx.nz=nz
+        ctx.dt=dt
+        ctx.nt=nt
+        ctx.nrx=nrx
+        ctx.nsr=nsr
+        ctx.nstep=nstep
+        ctx.pmlthick=pmlthick
+        ctx.device=device
+        ctx.dtype=dtype
+        ctx.model_gradient_sampling_interval = model_gradient_sampling_interval
+        ctx.use_async_offload = use_async_offload
+
+        nt_saved = (nt + model_gradient_sampling_interval - 1) // model_gradient_sampling_interval
+        
+        # 仅仅控制显存分配策略，无需修改 C 端参数签名
+        if use_async_offload:
+            Eall = torch.zeros((nt_saved, nstep, nx, ny, nz), device='cpu', dtype=dtype).pin_memory()
+        else:
+            Eall = torch.zeros((nt_saved, nstep, nx, ny, nz), device=device, dtype=dtype).contiguous()
+
+        Eupdatecoffs0=torch.zeros((nx+1,ny+1,nz+1), device=device, dtype=dtype)
+        Eupdatecoffs1=torch.zeros((nx+1,ny+1,nz+1), device=device, dtype=dtype)
+        Eupdatecoffs4=torch.zeros((nx+1,ny+1,nz+1), device=device, dtype=dtype)
+        Hupdatecoffs0=torch.zeros((nx+1,ny+1,nz+1), device=device, dtype=dtype)
+        Hupdatecoffs1=torch.zeros((nx+1,ny+1,nz+1), device=device, dtype=dtype)
+        Hupdatecoffs4=torch.zeros((nx+1,ny+1,nz+1), device=device, dtype=dtype)
+
+        receiver_amplitudes = torch.zeros((nstep, 6, nt, nrx),device=device).contiguous()  
+
+        c_lib.forward(
+                ctypes.cast(ere.data_ptr(), ctypes.POINTER(ctypes.c_float)), 
+                ctypes.cast(see.data_ptr(), ctypes.POINTER(ctypes.c_float)), 
+                ctypes.cast(mr.data_ptr(), ctypes.POINTER(ctypes.c_float)),
+                ctypes.cast(Eall.data_ptr(), ctypes.POINTER(ctypes.c_float)),              
+                ctypes.cast(Ex.data_ptr(), ctypes.POINTER(ctypes.c_float)),
+                ctypes.cast(Ey.data_ptr(), ctypes.POINTER(ctypes.c_float)),
+                ctypes.cast(Ez.data_ptr(), ctypes.POINTER(ctypes.c_float)),
+                ctypes.cast(Hx.data_ptr(), ctypes.POINTER(ctypes.c_float)),
+                ctypes.cast(Hy.data_ptr(), ctypes.POINTER(ctypes.c_float)),
+                ctypes.cast(Hz.data_ptr(), ctypes.POINTER(ctypes.c_float)),
+
+                ctypes.cast(Eupdatecoffs0.data_ptr(), ctypes.POINTER(ctypes.c_float)), ctypes.cast(Eupdatecoffs1.data_ptr(), ctypes.POINTER(ctypes.c_float)),
+                ctypes.cast(Eupdatecoffs4.data_ptr(), ctypes.POINTER(ctypes.c_float)), ctypes.cast(Hupdatecoffs0.data_ptr(), ctypes.POINTER(ctypes.c_float)),
+                ctypes.cast(Hupdatecoffs1.data_ptr(), ctypes.POINTER(ctypes.c_float)), ctypes.cast(Hupdatecoffs4.data_ptr(), ctypes.POINTER(ctypes.c_float)),
+
+                ctypes.cast(x0EPhi1.data_ptr(), ctypes.POINTER(ctypes.c_float)), ctypes.cast(x0EPhi2.data_ptr(), ctypes.POINTER(ctypes.c_float)),
+                ctypes.cast(x0HPhi1.data_ptr(), ctypes.POINTER(ctypes.c_float)), ctypes.cast(x0HPhi2.data_ptr(), ctypes.POINTER(ctypes.c_float)),
+                ctypes.cast(xmEPhi1.data_ptr(), ctypes.POINTER(ctypes.c_float)), ctypes.cast(xmEPhi2.data_ptr(), ctypes.POINTER(ctypes.c_float)),
+                ctypes.cast(xmHPhi1.data_ptr(), ctypes.POINTER(ctypes.c_float)), ctypes.cast(xmHPhi2.data_ptr(), ctypes.POINTER(ctypes.c_float)),
+                ctypes.cast(y0EPhi1.data_ptr(), ctypes.POINTER(ctypes.c_float)), ctypes.cast(y0EPhi2.data_ptr(), ctypes.POINTER(ctypes.c_float)),
+                ctypes.cast(y0HPhi1.data_ptr(), ctypes.POINTER(ctypes.c_float)), ctypes.cast(y0HPhi2.data_ptr(), ctypes.POINTER(ctypes.c_float)),
+                ctypes.cast(ymEPhi1.data_ptr(), ctypes.POINTER(ctypes.c_float)), ctypes.cast(ymEPhi2.data_ptr(), ctypes.POINTER(ctypes.c_float)),
+                ctypes.cast(ymHPhi1.data_ptr(), ctypes.POINTER(ctypes.c_float)), ctypes.cast(ymHPhi2.data_ptr(), ctypes.POINTER(ctypes.c_float)),
+                ctypes.cast(z0EPhi1.data_ptr(), ctypes.POINTER(ctypes.c_float)), ctypes.cast(z0EPhi2.data_ptr(), ctypes.POINTER(ctypes.c_float)),
+                ctypes.cast(z0HPhi1.data_ptr(), ctypes.POINTER(ctypes.c_float)), ctypes.cast(z0HPhi2.data_ptr(), ctypes.POINTER(ctypes.c_float)),
+                ctypes.cast(zmEPhi1.data_ptr(), ctypes.POINTER(ctypes.c_float)), ctypes.cast(zmEPhi2.data_ptr(), ctypes.POINTER(ctypes.c_float)),
+                ctypes.cast(zmHPhi1.data_ptr(), ctypes.POINTER(ctypes.c_float)), ctypes.cast(zmHPhi2.data_ptr(), ctypes.POINTER(ctypes.c_float)),
+
+                pmlthick[0],pmlthick[1],pmlthick[2],
+                pmlthick[3],pmlthick[4],pmlthick[5],
+
+                ctypes.cast(x01.data_ptr(), ctypes.POINTER(ctypes.c_float)), ctypes.cast(xm1.data_ptr(), ctypes.POINTER(ctypes.c_float)),
+                ctypes.cast(y01.data_ptr(), ctypes.POINTER(ctypes.c_float)), ctypes.cast(ym1.data_ptr(), ctypes.POINTER(ctypes.c_float)),       
+                ctypes.cast(z01.data_ptr(), ctypes.POINTER(ctypes.c_float)), ctypes.cast(zm1.data_ptr(), ctypes.POINTER(ctypes.c_float)),                
+                ctypes.cast(x02.data_ptr(), ctypes.POINTER(ctypes.c_float)), ctypes.cast(xm2.data_ptr(), ctypes.POINTER(ctypes.c_float)),               
+                ctypes.cast(y02.data_ptr(), ctypes.POINTER(ctypes.c_float)), ctypes.cast(ym2.data_ptr(), ctypes.POINTER(ctypes.c_float)),         
+                ctypes.cast(z02.data_ptr(), ctypes.POINTER(ctypes.c_float)), ctypes.cast(zm2.data_ptr(), ctypes.POINTER(ctypes.c_float)),
+
+                dt, nt, nstep, nrx, dx,
+                ctypes.cast(receiver_location.data_ptr(), ctypes.POINTER(ctypes.c_int)), ctypes.cast(receiver_amplitudes.data_ptr(), ctypes.POINTER(ctypes.c_float)),
+                nx+1, ny+1, nz+1, nsr,
+                ctypes.cast(source_location.data_ptr(), ctypes.POINTER(ctypes.c_int)), ctypes.cast(source_amplitudes.data_ptr(), ctypes.POINTER(ctypes.c_float)),
+                source_direction,
+                model_gradient_sampling_interval)
+
+        check_tensors_for_nan_inf(d="forward",
+            Ex=Ex, Ey=Ey, Ez=Ez,
+            Hx=Hx, Hy=Hy, Hz=Hz,
+            x0EPhi1=x0EPhi1, x0EPhi2=x0EPhi2,
+            x0HPhi1=x0HPhi1, x0HPhi2=x0HPhi2,
+            xmEPhi1=xmEPhi1, xmEPhi2=xmEPhi2,
+            xmHPhi1=xmHPhi1, xmHPhi2=xmHPhi2,
+            y0EPhi1=y0EPhi1, y0EPhi2=y0EPhi2,
+            y0HPhi1=y0HPhi1, y0HPhi2=y0HPhi2,
+            ymEPhi1=ymEPhi1, ymEPhi2=ymEPhi2,
+            ymHPhi1=ymHPhi1, ymHPhi2=ymHPhi2,
+            z0EPhi1=z0EPhi1, z0EPhi2=z0EPhi2,
+            z0HPhi1=z0HPhi1, z0HPhi2=z0HPhi2,
+            zmEPhi1=zmEPhi1, zmEPhi2=zmEPhi2,
+            zmHPhi1=zmHPhi1, zmHPhi2=zmHPhi2
+        )
+
+        ctx.Eall = Eall
+        return (Ex,Ey,Ez,Hx,Hy,Hz,x0EPhi1,x0EPhi2,x0HPhi1,x0HPhi2,xmEPhi1,xmEPhi2,xmHPhi1,xmHPhi2,y0EPhi1,y0EPhi2,y0HPhi1,y0HPhi2,ymEPhi1,ymEPhi2,ymHPhi1,ymHPhi2,z0EPhi1,z0EPhi2,z0HPhi1,z0HPhi2,zmEPhi1,zmEPhi2,zmHPhi1,zmHPhi2,Eall,receiver_amplitudes[:,reciever_direction,:,:])
+
+    @staticmethod
+    def backward(ctx,gEx,gEy,gEz,gHx,gHy,gHz,gx0EPhi1,gx0EPhi2,gx0HPhi1,gx0HPhi2,gxmEPhi1,gxmEPhi2,gxmHPhi1,gxmHPhi2,gy0EPhi1,gy0EPhi2,gy0HPhi1,gy0HPhi2,gymEPhi1,gymEPhi2,gymHPhi1,gymHPhi2,gz0EPhi1,gz0EPhi2,gz0HPhi1,gz0HPhi2,gzmEPhi1,gzmEPhi2,gzmHPhi1,gzmHPhi2,gEall,gezreciver):
+        
+        sourceamp=gezreciver
+        er, se, mr,receiver_location,x0,xm,y0,ym,z0,zm,x01,x02,xm1,xm2,y01,y02,ym1,ym2,z01,z02,zm1,zm2,ere,see=ctx.saved_tensors
+        
+        ere=ere.contiguous()
+        see=see.contiguous()
+        er=er.contiguous()
+        se=se.contiguous()
+        mr=mr.contiguous()
+        receiver_location=receiver_location.contiguous()
+        x0=x0.contiguous()
+        xm=xm.contiguous()
+        y0=y0.contiguous()
+        ym=ym.contiguous()
+        z0=z0.contiguous()
+        zm=zm.contiguous()
+        x01=x01.contiguous()
+        x02=x02.contiguous()
+        xm1=xm1.contiguous()
+        xm2=xm2.contiguous()
+        y01=y01.contiguous()
+        y02=y02.contiguous()
+        ym1=ym1.contiguous()
+        ym2=ym2.contiguous()
+        z01=z01.contiguous()
+        z02=z02.contiguous()
+        zm1=zm1.contiguous()
+        zm2=zm2.contiguous()
+
+        dx=ctx.dx
+        nx=ctx.nx
+        ny=ctx.ny
+        nz=ctx.nz
+        dt=ctx.dt
+        nt=ctx.nt
+        nsr=ctx.nrx
+        nrx=ctx.nsr
+        dtype=ctx.dtype
+        nstep=ctx.nstep
+        pmlthick=ctx.pmlthick
+        device=ctx.device
+        Eall=ctx.Eall
+        model_gradient_sampling_interval = ctx.model_gradient_sampling_interval
+
+        Eall=Eall.contiguous()
+        gEx=gEx.contiguous()
+        gEy=gEy.contiguous()
+        gEz=gEz.contiguous()
+        gHx=gHx.contiguous()
+        gHy=gHy.contiguous()
+        gHz=gHz.contiguous()
+        gx0EPhi1=gx0EPhi1.contiguous()
+        gx0EPhi2=gx0EPhi2.contiguous()
+        gx0HPhi1=gx0HPhi1.contiguous()
+        gx0HPhi2=gx0HPhi2.contiguous()
+        gxmEPhi1=gxmEPhi1.contiguous()
+        gxmEPhi2=gxmEPhi2.contiguous()
+        gxmHPhi1=gxmHPhi1.contiguous()
+        gxmHPhi2=gxmHPhi2.contiguous()
+        gy0EPhi1=gy0EPhi1.contiguous()
+        gy0EPhi2=gy0EPhi2.contiguous()
+        gy0HPhi1=gy0HPhi1.contiguous()
+        gy0HPhi2=gy0HPhi2.contiguous()
+        gymEPhi1=gymEPhi1.contiguous()
+        gymEPhi2=gymEPhi2.contiguous()
+        gymHPhi1=gymHPhi1.contiguous()
+        gymHPhi2=gymHPhi2.contiguous()
+        gz0EPhi1=gz0EPhi1.contiguous()
+        gz0EPhi2=gz0EPhi2.contiguous()
+        gz0HPhi1=gz0HPhi1.contiguous()
+        gz0HPhi2=gz0HPhi2.contiguous()
+        gzmEPhi1=gzmEPhi1.contiguous()
+        gzmEPhi2=gzmEPhi2.contiguous()
+        gzmHPhi1=gzmHPhi1.contiguous()
+        gzmHPhi2=gzmHPhi2.contiguous()
+        sourceamp=sourceamp.contiguous()
+
+        Eupdatecoffs0=torch.zeros((nx+1,ny+1,nz+1), device=device, dtype=dtype)
+        Eupdatecoffs1=torch.zeros((nx+1,ny+1,nz+1), device=device, dtype=dtype)
+        Eupdatecoffs4=torch.zeros((nx+1,ny+1,nz+1), device=device, dtype=dtype)
+        Hupdatecoffs0=torch.zeros((nx+1,ny+1,nz+1), device=device, dtype=dtype)
+        Hupdatecoffs1=torch.zeros((nx+1,ny+1,nz+1), device=device, dtype=dtype)
+        Hupdatecoffs4=torch.zeros((nx+1,ny+1,nz+1), device=device, dtype=dtype)
+
+        if er.requires_grad:
+            grad_er=torch.zeros((nx,ny,nz),device=device,dtype=dtype).contiguous()
+            errequiregrad=1
+        else:
+            grad_er=torch.empty(0)
+            errequiregrad=0
+
+        if se.requires_grad: 
+            grad_se=torch.zeros((nx,ny,nz),device=device,dtype=dtype).contiguous()
+            serequiregrad=1
+        else:
+            grad_se=torch.empty(0)
+            serequiregrad=0
+
+        c_lib.backward(
+                ctypes.cast(ere.data_ptr(), ctypes.POINTER(ctypes.c_float)), 
+                ctypes.cast(see.data_ptr(), ctypes.POINTER(ctypes.c_float)), 
+                ctypes.cast(mr.data_ptr(), ctypes.POINTER(ctypes.c_float)),
+                ctypes.cast(Eall.data_ptr(), ctypes.POINTER(ctypes.c_float)),               
+                ctypes.cast(gEx.data_ptr(), ctypes.POINTER(ctypes.c_float)),
+                ctypes.cast(gEy.data_ptr(), ctypes.POINTER(ctypes.c_float)),
+                ctypes.cast(gEz.data_ptr(), ctypes.POINTER(ctypes.c_float)),
+                ctypes.cast(gHx.data_ptr(), ctypes.POINTER(ctypes.c_float)),
+                ctypes.cast(gHy.data_ptr(), ctypes.POINTER(ctypes.c_float)),
+                ctypes.cast(gHz.data_ptr(), ctypes.POINTER(ctypes.c_float)), 
+
+                ctypes.cast(Eupdatecoffs0.data_ptr(), ctypes.POINTER(ctypes.c_float)), ctypes.cast(Eupdatecoffs1.data_ptr(), ctypes.POINTER(ctypes.c_float)),
+                ctypes.cast(Eupdatecoffs4.data_ptr(), ctypes.POINTER(ctypes.c_float)), ctypes.cast(Hupdatecoffs0.data_ptr(), ctypes.POINTER(ctypes.c_float)),
+                ctypes.cast(Hupdatecoffs1.data_ptr(), ctypes.POINTER(ctypes.c_float)), ctypes.cast(Hupdatecoffs4.data_ptr(), ctypes.POINTER(ctypes.c_float)), 
+
+                ctypes.cast(gx0EPhi1.data_ptr(), ctypes.POINTER(ctypes.c_float)), ctypes.cast(gx0EPhi2.data_ptr(), ctypes.POINTER(ctypes.c_float)),
+                ctypes.cast(gx0HPhi1.data_ptr(), ctypes.POINTER(ctypes.c_float)), ctypes.cast(gx0HPhi2.data_ptr(), ctypes.POINTER(ctypes.c_float)),
+                ctypes.cast(gxmEPhi1.data_ptr(), ctypes.POINTER(ctypes.c_float)), ctypes.cast(gxmEPhi2.data_ptr(), ctypes.POINTER(ctypes.c_float)),
+                ctypes.cast(gxmHPhi1.data_ptr(), ctypes.POINTER(ctypes.c_float)), ctypes.cast(gxmHPhi2.data_ptr(), ctypes.POINTER(ctypes.c_float)),
+                ctypes.cast(gy0EPhi1.data_ptr(), ctypes.POINTER(ctypes.c_float)), ctypes.cast(gy0EPhi2.data_ptr(), ctypes.POINTER(ctypes.c_float)),
+                ctypes.cast(gy0HPhi1.data_ptr(), ctypes.POINTER(ctypes.c_float)), ctypes.cast(gy0HPhi2.data_ptr(), ctypes.POINTER(ctypes.c_float)),
+                ctypes.cast(gymEPhi1.data_ptr(), ctypes.POINTER(ctypes.c_float)), ctypes.cast(gymEPhi2.data_ptr(), ctypes.POINTER(ctypes.c_float)),
+                ctypes.cast(gymHPhi1.data_ptr(), ctypes.POINTER(ctypes.c_float)), ctypes.cast(gymHPhi2.data_ptr(), ctypes.POINTER(ctypes.c_float)),
+                ctypes.cast(gz0EPhi1.data_ptr(), ctypes.POINTER(ctypes.c_float)), ctypes.cast(gz0EPhi2.data_ptr(), ctypes.POINTER(ctypes.c_float)),
+                ctypes.cast(gz0HPhi1.data_ptr(), ctypes.POINTER(ctypes.c_float)), ctypes.cast(gz0HPhi2.data_ptr(), ctypes.POINTER(ctypes.c_float)),
+                ctypes.cast(gzmEPhi1.data_ptr(), ctypes.POINTER(ctypes.c_float)), ctypes.cast(gzmEPhi2.data_ptr(), ctypes.POINTER(ctypes.c_float)),
+                ctypes.cast(gzmHPhi1.data_ptr(), ctypes.POINTER(ctypes.c_float)), ctypes.cast(gzmHPhi2.data_ptr(), ctypes.POINTER(ctypes.c_float)), 
+
+                pmlthick[0],pmlthick[1],pmlthick[2],
+                pmlthick[3],pmlthick[4],pmlthick[5],
+
+                ctypes.cast(x01.data_ptr(), ctypes.POINTER(ctypes.c_float)), ctypes.cast(xm1.data_ptr(), ctypes.POINTER(ctypes.c_float)),
+                ctypes.cast(y01.data_ptr(), ctypes.POINTER(ctypes.c_float)), ctypes.cast(ym1.data_ptr(), ctypes.POINTER(ctypes.c_float)),       
+                ctypes.cast(z01.data_ptr(), ctypes.POINTER(ctypes.c_float)), ctypes.cast(zm1.data_ptr(), ctypes.POINTER(ctypes.c_float)),                
+                ctypes.cast(x02.data_ptr(), ctypes.POINTER(ctypes.c_float)), ctypes.cast(xm2.data_ptr(), ctypes.POINTER(ctypes.c_float)),               
+                ctypes.cast(y02.data_ptr(), ctypes.POINTER(ctypes.c_float)), ctypes.cast(ym2.data_ptr(), ctypes.POINTER(ctypes.c_float)),         
+                ctypes.cast(z02.data_ptr(), ctypes.POINTER(ctypes.c_float)), ctypes.cast(zm2.data_ptr(), ctypes.POINTER(ctypes.c_float)), 
+
+                dt, nt, nstep, nrx, dx,
+                nx+1, ny+1, nz+1, nsr,
+                ctypes.cast(receiver_location.data_ptr(), ctypes.POINTER(ctypes.c_int)), ctypes.cast(sourceamp.data_ptr(), ctypes.POINTER(ctypes.c_float)),
+                2,
+                ctypes.cast(grad_er.data_ptr(), ctypes.POINTER(ctypes.c_float)), ctypes.cast(grad_se.data_ptr(), ctypes.POINTER(ctypes.c_float)),errequiregrad,serequiregrad,
+                model_gradient_sampling_interval)  
+        
+        check_tensors_for_nan_inf(d="backward",
+            gEx=gEx, gEy=gEy, gEz=gEz,
+            gHx=gHx, gHy=gHy, gHz=gHz,
+            gx0EPhi1=gx0EPhi1, gx0EPhi2=gx0EPhi2,
+            gx0HPhi1=gx0HPhi1, gx0HPhi2=gx0HPhi2,
+            gxmEPhi1=gxmEPhi1, gxmEPhi2=gxmEPhi2,
+            gxmHPhi1=gxmHPhi1, gxmHPhi2=gxmHPhi2,
+            gy0EPhi1=gy0EPhi1, gy0EPhi2=gy0EPhi2,
+            gy0HPhi1=gy0HPhi1, gy0HPhi2=gy0HPhi2,
+            gymEPhi1=gymEPhi1, gymEPhi2=gymEPhi2,
+            gymHPhi1=gymHPhi1, gymHPhi2=gymHPhi2,
+            gz0EPhi1=gz0EPhi1, gz0EPhi2=gz0EPhi2,
+            gz0HPhi1=gz0HPhi1, gz0HPhi2=gz0HPhi2,
+            gzmEPhi1=gzmEPhi1, gzmEPhi2=gzmEPhi2,
+            gzmHPhi1=gzmHPhi1, gzmHPhi2=gzmHPhi2
+        )
+
+        ctx.Eall = None
+        del Eall,er, se, mr,receiver_location,x0,xm,y0,ym,z0,zm,x01,x02,xm1,xm2,y01,y02,ym1,ym2,z01,z02,zm1,zm2,ere,see, Eupdatecoffs0, Eupdatecoffs1, Eupdatecoffs4, Hupdatecoffs0, Hupdatecoffs1, Hupdatecoffs4
+
+        # 返回与前向传播参数一一对应（最后补齐 use_async_offload 占位）
+        return (
+                    grad_er, grad_se,         
+                    gEx,gEy,gEz, gHx,gHy,gHz, 
+                    gx0EPhi1,gx0EPhi2,gx0HPhi1,gx0HPhi2,
+                    gxmEPhi1,gxmEPhi2,gxmHPhi1,gxmHPhi2,
+                    gy0EPhi1,gy0EPhi2,gy0HPhi1,gy0HPhi2,
+                    gymEPhi1,gymEPhi2,gymHPhi1,gymHPhi2,
+                    gz0EPhi1,gz0EPhi2,gz0HPhi1,gz0HPhi2,
+                    gzmEPhi1,gzmEPhi2,gzmHPhi1,gzmHPhi2,   
+                    None, None, None, None, None, None, None, None,   
+                    None, None, None, None, None, None, None,
+                    None, None, None, None, None, None, None, None,
+                    None, None, None, None, None, None, None, None, 
+                    None, None, None, None, None, None, None, None, 
+                    None  
+                )

DeepGPR/multiscale.py

@@ -0,0 +1,63 @@
+
+import torch
+
+def design_fir_filter(cutoff: float, fs: float, numtaps: int) -> torch.Tensor:
+    """使用 Hamming 窗设计 FIR 滤波器"""
+    n = torch.arange(numtaps, dtype=torch.float32)
+    # 汉明窗
+    window = 0.54 - 0.46 * torch.cos(2 * torch.pi * n / (numtaps - 1))
+    # 正弦波
+    sinc = torch.sin(2 * torch.pi * (cutoff/fs) * (n - (numtaps-1)/2)) / (torch.pi * (n - (numtaps-1)/2))
+    # 处理中心点
+    center = (numtaps-1) // 2
+    sinc[center] = 2 * cutoff/fs
+    # 应用窗函数
+    h = window * sinc
+    # 归一化
+    return h / h.sum()
+
+def apply_filter(data: torch.Tensor, fs: float, cutoff: float) -> torch.Tensor:
+    """应用 FIR 滤波器到数据"""
+    numtaps = int(1 * (fs / cutoff))
+    fir_coeff = design_fir_filter(cutoff, fs, numtaps)
+    fir_coeff = fir_coeff.to(data.device)
+
+    if data.ndim == 1:
+        # 1D 数据处理 - 添加维度以支持反射填充
+        data_2d = data.view(1, 1, -1)
+        padded_data = torch.nn.functional.pad(data_2d, (numtaps-1, 0), mode='reflect')
+        filtered = torch.nn.functional.conv1d(
+            padded_data, 
+            fir_coeff.view(1, 1, -1), 
+            padding=0
+        )
+        return filtered.view(-1)
+
+    elif data.ndim == 3:
+        # 3D 数据处理
+        step, iterations, nrx = data.shape
+        # 重塑数据以使用批量处理
+        reshaped_data = data.permute(0, 2, 1).reshape(-1, 1, iterations)
+        # 填充数据
+        padded_data = torch.nn.functional.pad(reshaped_data, (numtaps-1, 0), mode='reflect')
+        # 应用卷积
+        filtered = torch.nn.functional.conv1d(
+            padded_data,
+            fir_coeff.view(1, 1, -1),
+            padding=0
+        )
+        # 重塑回原始形状
+        return filtered.view(step, nrx, iterations).permute(0, 2, 1)
+
+    else:
+        raise ValueError(f"不支持的数据维度: {data.ndim}。期望 1D 或 3D tensor。")
+    
+def hilbert_transform(data_in, p=1):
+  ns, nt, nr = data_in.shape
+  transforms = torch.fft.fftn(data_in,dim=1)
+  #print(transforms.shape)
+  transforms[:,1:nt//2,:]      *= 2.0
+  transforms[:,nt//2 + 1: nt,:]  = 0+0j
+  transforms[:,0,:] = 0;
+  data_out = torch.abs(torch.fft.ifftn(transforms,dim=1))**p
+  return data_out

DeepGPR/visual.py

@@ -0,0 +1,105 @@
+import torch
+import numpy as np
+import matplotlib.pyplot as plt
+
+
+def plot_survey_geometry(model, sources, receivers, dx=0.02, title="Survey Geometry"):
+
+    # 1. 数据转换与展平
+    if isinstance(model, torch.Tensor):
+        model = model.detach().cpu().numpy()
+    if isinstance(sources, torch.Tensor):
+        sources = sources.detach().cpu().numpy()
+    if isinstance(receivers, torch.Tensor):
+        receivers = receivers.detach().cpu().numpy()
+
+    # 处理多炮维度 [Shots, N, 3] -> [Total, 3]
+    if sources.ndim == 3: sources = sources.reshape(-1, 3)
+    if receivers.ndim == 3: receivers = receivers.reshape(-1, 3)
+
+    nx, ny, nz = model.shape
+    
+    # 2. 绘图
+    if nz == 1:
+        # ==========================================
+        # 2D 绘图模式 (垂直朝上，0在顶部)
+        # ==========================================
+        fig, ax = plt.subplots(figsize=(6, 4)) 
+        
+        slice_data = model[:, :, 0] # Shape: (nx, ny)
+        
+        # [修改 1] extent 乘以 dx，转换为真实物理距离
+        # extent=[左, 右, 下, 上] 
+        # 垂直轴从 nx*dx (底部) 到 0 (顶部)
+        extent_real = [0, ny * dx, nx * dx, 0]
+        
+        im = ax.imshow(slice_data, cmap='jet', origin='upper', 
+                       extent=extent_real, aspect='auto')
+        
+        plt.colorbar(im, ax=ax, label='Relative permittivity')
+
+        # [修改 2] 散点坐标乘以 dx
+        # Scatter x (水平) = sources[:, 1] * dx
+        # Scatter y (垂直) = sources[:, 0] * dx
+        ax.scatter(sources[:, 1] * dx, sources[:, 0] * dx, c='red', marker='*', s=150, 
+                   label='Sources', edgecolors='k', zorder=10)
+        
+        ax.scatter(receivers[:, 1] * dx, receivers[:, 0] * dx, c='white', marker='v', s=80, 
+                   label='Receivers', edgecolors='k', zorder=9, alpha=0.7)
+
+        # [修改 3] 标签改为物理单位
+        ax.set_xlabel('Distance (m)') 
+        ax.set_ylabel('Depth (m)') 
+        
+        # 保持 Y 轴方向控制 (0 在上)
+        if ax.get_ylim()[0] < ax.get_ylim()[1]: 
+             ax.invert_yaxis()
+
+        # ax.set_title(f"{title} ")
+        ax.legend(loc='lower right')
+        ax.grid(True, linestyle='--', alpha=0.3)
+
+    else:
+        # ==========================================
+        # 3D 绘图模式 (同步应用 dx 以保持一致)
+        # ==========================================
+        fig = plt.figure(figsize=(12, 10))
+        ax = fig.add_subplot(111, projection='3d')
+
+        # 散点坐标乘以 dx
+        ax.scatter(sources[:, 0] * dx, sources[:, 1] * dx, sources[:, 2] * dx, 
+                   c='red', marker='*', s=100, label='Sources')
+        ax.scatter(receivers[:, 0] * dx, receivers[:, 1] * dx, receivers[:, 2] * dx, 
+                   c='blue', marker='v', s=40, label='Receivers', alpha=0.6)
+
+        # 网格坐标乘以 dx
+        x_grid = np.arange(nx) * dx
+        y_grid = np.arange(ny) * dx
+        z_grid = np.arange(nz) * dx
+        
+        cx, cy, cz = nx // 2, ny // 2, nz // 2
+        
+        # 生成真实坐标的 Meshgrid
+        X, Y = np.meshgrid(x_grid, y_grid, indexing='ij') 
+        # 注意 offset 也要乘以 dx
+        ax.contourf(X, Y, model[:, :, cz], zdir='z', offset=cz * dx, cmap='viridis', alpha=0.5)
+        
+        X, Z = np.meshgrid(x_grid, z_grid, indexing='ij')
+        ax.contourf(X, model[:, cy, :], Z, zdir='y', offset=cy * dx, cmap='viridis', alpha=0.5)
+        
+        Y, Z = np.meshgrid(y_grid, z_grid, indexing='ij')
+        ax.contourf(model[cx, :, :], Y, Z, zdir='x', offset=cx * dx, cmap='viridis', alpha=0.5)
+
+        # 设置真实坐标范围
+        ax.set_xlim(0, nx * dx)
+        ax.set_ylim(0, ny * dx)
+        ax.set_zlim(0, nz * dx)
+        
+        ax.set_xlabel('X (m)')
+        ax.set_ylabel('Y (m)')
+        ax.set_zlabel('Z (m)')
+        ax.set_title(title)
+        ax.legend()
+
+    plt.tight_layout()
+    plt.show()

deepgpr-0.0.1.dist-info/METADATA

@@ -0,0 +1,86 @@
+Metadata-Version: 2.4
+Name: DeepGPR
+Version: 0.0.1
+Summary: PyTorch and CUDA for GPR FWI
+Author-email: Lei Liu <liulei990222@gmail.com>
+Classifier: Programming Language :: Python :: 3
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.7
+Description-Content-Type: text/markdown
+
+# DeepGPR
+
+DeepGPR provides a wave propagation module for PyTorch, designed for applications such as Ground Penetrating Radar (GPR) imaging and inversion. Its core concepts are derived from Deepwave. You can use it to perform both forward modeling and backpropagation—thereby enabling the simulation of wave propagation to generate synthetic data—as well as for Full Waveform Inversion (FWI). Furthermore, you can integrate this wave propagation functionality into a larger operational pipeline—incorporating various wavelets, loss functions, and other components—to achieve end-to-end forward and reverse propagation, powered by automatic differentiation and our high-performance operators.
+
+
+## Features
+
+Supports 2D and 3D forward modeling of Maxwell's equations—via the Finite-Difference Time-Domain (FDTD) method—for both single and multiple excitation scenarios.
+
+Gradients of the output receiver data can be computed with respect to model parameters (relative permittivity, conductivity), the initial wavefield, and source amplitudes.
+
+Utilizes CPML, allowing the width of the PML layer to be configured independently for each boundary.
+
+All operations are executed on the GPU.
+
+Supports techniques such as checkpointing, DDP, and the utilization of CPU memory to minimize GPU memory consumption, thereby enabling the execution of large-scale models.
+
+
+## Start
+
+Before use, you must ensure that you have an NVIDIA graphics card and have installed a CUDA-enabled version of PyTorch.
+
+DeepGPR can then be installed using
+
+```bash
+  pip install DeepGPR
+```
+
+A Small Forward Modeling Test
+
+```python
+import torch
+import DeepGPR
+import matplotlib.pyplot as plt
+
+# Set up the parameters and models
+device=torch.device("cuda")
+dx=0.02
+dt=3e-11
+nt=2000
+er = torch.ones(100, 100,1) * 2  
+er[50:,:]=5
+se = torch.zeros_like(er)  
+er.requires_grad_()
+source_location=torch.tensor([[[10,10,0]]],device=device,dtype=torch.int)
+receiver_location=torch.tensor([[[10,90,0]]],device=device,dtype=torch.int)
+freq=2e8
+peak_time = 1 / freq
+source_amplitudes = torch.zeros((1,nt,1),device=device)
+source_amplitudes[0,:,0]=DeepGPR.ricker(freq, nt, dt, peak_time).to(device)
+
+DeepGPR.plot_survey_geometry(er,source_location, receiver_location,dx)
+
+#forward modeling
+r = DeepGPR.compute(
+    device=device, dx=dx, dt=dt, 
+    source_amplitudes=source_amplitudes,
+    source_location=source_location, 
+    receiver_location=receiver_location, 
+    er=er, se=se
+)
+
+(r[-1]**2).sum().backward()
+
+_, ax = plt.subplots(1, 2, figsize=(10, 3))
+ax[0].plot(r[-1].detach().flatten().cpu().numpy())
+ax[0].set_title("Receiver data")
+ax[1].imshow(er.grad.detach())
+ax[1].set_title("Gradient")
+plt.show()
+```
+![result](./Fig/example1.png)
+
+![result](./Fig/example2.png)
+
+There are more examples in the ./examples.

deepgpr-0.0.1.dist-info/RECORD

@@ -0,0 +1,9 @@
+DeepGPR/__init__.py,sha256=mhXJhG4SQ2C1o8GAMygkKxR_annxX1Txku8rdQbmVi0,6663
+DeepGPR/common.py,sha256=If4uIzdLFSUdHTY3bHxB9Ws4LkB5cJiNlOAaY667-es,27915
+DeepGPR/compute2.py,sha256=4b1YWlzPhmD-_rFJ6yRB2iaKHlcOhJccgzbaOR9UJPQ,20230
+DeepGPR/multiscale.py,sha256=_tSFC3alHPDmMBPbRm1pJy5nZHQ-rtvgTjhupfhrPlc,2240
+DeepGPR/visual.py,sha256=Iynfiv_2DKVB7MCbgLvvmWn_m8jhrHoDLlHWyC9iHeY,3944
+deepgpr-0.0.1.dist-info/METADATA,sha256=wijleiqiX73S5EMxPNOWznbu1OnVC9JODpc7EOFsEvE,3089
+deepgpr-0.0.1.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
+deepgpr-0.0.1.dist-info/top_level.txt,sha256=Clp4TLG1b60qzO8D6ZYieWFKK39K1uU_Em9yRGwelTE,8
+deepgpr-0.0.1.dist-info/RECORD,,

deepgpr-0.0.1.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (82.0.1)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

deepgpr-0.0.1.dist-info/top_level.txt

@@ -0,0 +1 @@
+DeepGPR