DeepGPR 0.0.16__tar.gz → 0.0.18__tar.gz

{deepgpr-0.0.16 → deepgpr-0.0.18}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: DeepGPR
-Version: 0.0.16
+Version: 0.0.18
 Summary: PyTorch and CUDA for GPR FWI
 Author-email: Lei Liu <liulei990222@gmail.com>
 License-Expression: MIT

{deepgpr-0.0.16 → deepgpr-0.0.18}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "DeepGPR"
-version = "0.0.16"
+version = "0.0.18"
 authors = [
     { name = "Lei Liu", email = "liulei990222@gmail.com" }
 ]

{deepgpr-0.0.16 → deepgpr-0.0.18}/src/DeepGPR/common.py

@@ -80,6 +80,7 @@ def initialization(device, er,se,mr,source_amplitudes,source_location,receiver_l
         pmlthick: PML thickness as an int, list, or tensor.
         fdtd_order: Spatial finite-difference order used for the CFL check.
     """
+    device = torch.device("cpu" if device is None else device)
     dtype=torch.float32
     spacing = _normalize_grid_spacing(dx)
     try:
@@ -100,6 +101,8 @@ def initialization(device, er,se,mr,source_amplitudes,source_location,receiver_l
     elif len(se.shape) != 3:
         raise ValueError('The shape of sigma should be 2-d or 3-d.')
 
+    if mr is not None and not torch.is_tensor(mr):
+        raise TypeError("mr must be a PyTorch tensor or None.")
     if mr is not None:
         if len(mr.shape) == 2:
             mr = mr.reshape(*mr.shape, 1)
@@ -108,6 +111,8 @@ def initialization(device, er,se,mr,source_amplitudes,source_location,receiver_l
 
     if er.shape != se.shape:
         raise ValueError('The shape of epsilon and sigma should be the same.')
+    if any(size < 1 for size in er.shape):
+        raise ValueError("The material model dimensions must all be non-empty.")
     nx, ny, nz = er.shape
     mode = 2 if nz == 1 else 3
     er=er.to(device=device, dtype=dtype)
@@ -119,17 +124,36 @@ def initialization(device, er,se,mr,source_amplitudes,source_location,receiver_l
     else:
         raise ValueError('The shape of mr should be the same as epsilon and sigma.')
 
-    if source_location.shape[0] != receiver_location.shape[0]:
-        raise ValueError('The first dimension (nstep) of source_location and receiver_location should be the same.')
+    if not torch.is_tensor(source_location) or not torch.is_tensor(receiver_location):
+        raise TypeError("source_location and receiver_location must be PyTorch tensors.")
     if source_location.ndim != 3 or receiver_location.ndim != 3:
         raise ValueError("source_location and receiver_location must have shape (nstep, count, 3).")
     if source_location.shape[2] != 3 or receiver_location.shape[2] != 3:
         raise ValueError("The last dimension of source_location and receiver_location must be 3.")
+    if source_location.shape[0] != receiver_location.shape[0]:
+        raise ValueError('The first dimension (nstep) of source_location and receiver_location should be the same.')
     nstep=source_location.shape[0]
     nsr=source_location.shape[1]
     nrx=receiver_location.shape[1]
-    source_location=source_location.to(device=device, dtype=torch.int32).contiguous()
-    receiver_location=receiver_location.to(device=device, dtype=torch.int32).contiguous()
+    if nstep < 1:
+        raise ValueError("At least one shot is required.")
+    if nsr < 1:
+        raise ValueError("At least one source per shot is required.")
+    if nrx < 1:
+        raise ValueError("At least one receiver per shot is required.")
+    if device.type == "cuda" and nstep > 65535:
+        raise ValueError(
+            "CUDA supports at most 65535 shots in one DeepGPR compute call. "
+            "Split a larger acquisition batch into smaller calls."
+        )
+    for name, locations in (
+        ("source_location", source_location),
+        ("receiver_location", receiver_location),
+    ):
+        if locations.dtype == torch.bool or locations.is_complex():
+            raise TypeError(f"{name} must contain real integer-valued coordinates.")
+    source_location=source_location.to(device=device).contiguous()
+    receiver_location=receiver_location.to(device=device).contiguous()
 
     if not torch.is_tensor(source_amplitudes):
         raise TypeError("source_amplitudes must be a PyTorch tensor.")
@@ -165,6 +189,16 @@ def initialization(device, er,se,mr,source_amplitudes,source_location,receiver_l
             )
 
     shape_tensor = torch.tensor((nx, ny, nz), dtype=torch.int32, device=device)
+    source_integral = (
+        (torch.isfinite(source_location) & (source_location == source_location.trunc())).all()
+        if source_location.is_floating_point()
+        else torch.ones((), dtype=torch.bool, device=device)
+    )
+    receiver_integral = (
+        (torch.isfinite(receiver_location) & (receiver_location == receiver_location.trunc())).all()
+        if receiver_location.is_floating_point()
+        else torch.ones((), dtype=torch.bool, device=device)
+    )
     source_valid = ((source_location >= 0) & (source_location < shape_tensor)).all()
     receiver_valid = ((receiver_location >= 0) & (receiver_location < shape_tensor)).all()
     source_in_pml = _locations_in_pml(source_location, (nx, ny, nz), pml_values)
@@ -181,6 +215,8 @@ def initialization(device, er,se,mr,source_amplitudes,source_location,receiver_l
             (er.detach() * mr.detach()).amin(),
             source_valid.to(dtype),
             receiver_valid.to(dtype),
+            source_integral.to(dtype),
+            receiver_integral.to(dtype),
             source_in_pml.sum().to(dtype),
             receiver_in_pml.sum().to(dtype),
         )
@@ -188,7 +224,8 @@ def initialization(device, er,se,mr,source_amplitudes,source_location,receiver_l
     er_finite, se_finite, mr_finite, source_finite = (bool(value) for value in stats[:4])
     er_min, se_min, mr_min, min_er_mr = stats[4:8]
     source_valid, receiver_valid = (bool(value) for value in stats[8:10])
-    source_pml_count, receiver_pml_count = (int(value) for value in stats[10:12])
+    source_integral, receiver_integral = (bool(value) for value in stats[10:12])
+    source_pml_count, receiver_pml_count = (int(value) for value in stats[12:14])
 
     for name, finite in (
         ("er", er_finite), ("se", se_finite), ("mr", mr_finite),
@@ -202,6 +239,10 @@ def initialization(device, er,se,mr,source_amplitudes,source_location,receiver_l
         raise ValueError('The values of sigma is incorrect.(should be non-negative)')
     if mr_min <= 0:
         raise ValueError('The values of mr are incorrect (must be positive).')
+    if not source_integral:
+        raise ValueError("source_location must contain finite integer-valued coordinates.")
+    if not receiver_integral:
+        raise ValueError("receiver_location must contain finite integer-valued coordinates.")
     if not source_valid:
         raise ValueError(
             "Error: Source coordinates out of range! "
@@ -292,24 +333,46 @@ def pmlthick_revert(p, er):
         p: PML thickness as an int, list, or tensor.
         er: Relative permittivity tensor used to detect 2D or 3D mode.
     """
-    if isinstance(p, int):  
-        if er.shape[2] == 1:
-            return torch.tensor([p, p, p, p, 0, 0], dtype=torch.int32)
-        return torch.tensor([p]*6, dtype=torch.int32)
-    
-    elif isinstance(p, list):
-        if len(p) == 6:
-            return torch.tensor(p, dtype=torch.int32)
-        elif len(p) == 4:
-            return torch.tensor(p + [0, 0], dtype=torch.int32)
+    if isinstance(p, bool):
+        raise TypeError("PML thickness must contain integer values, not bool.")
+    if isinstance(p, int):
+        values = [p, p, p, p, 0, 0] if er.shape[2] == 1 else [p] * 6
+    elif isinstance(p, (list, tuple)):
+        if len(p) == 4:
+            values = [*p, 0, 0]
+        elif len(p) == 6:
+            values = list(p)
         else:
-            raise ValueError(f"Unsupported list length: {len(p)}. Must be 4 or 6.")
-    
+            raise ValueError(f"Unsupported PML length: {len(p)}. Must be 4 or 6.")
     elif isinstance(p, torch.Tensor):
-        return p.detach().to(device="cpu", dtype=torch.int32)
-    
+        if p.ndim != 1:
+            raise ValueError("PML thickness tensor must be one-dimensional.")
+        if p.numel() == 4:
+            values = [*p.detach().cpu().tolist(), 0, 0]
+        elif p.numel() == 6:
+            values = p.detach().cpu().tolist()
+        else:
+            raise ValueError(
+                f"Unsupported PML length: {p.numel()}. Must be 4 or 6."
+            )
     else:
-        raise TypeError(f"Unsupported type: {type(p)}")
+        raise TypeError(f"Unsupported PML thickness type: {type(p)}")
+
+    normalized = []
+    for value in values:
+        if isinstance(value, bool):
+            raise TypeError("PML thickness must contain integer values, not bool.")
+        try:
+            numeric = float(value)
+        except (TypeError, ValueError, OverflowError) as exc:
+            raise TypeError("PML thickness must contain finite integer values.") from exc
+        if not math.isfinite(numeric) or not numeric.is_integer():
+            raise ValueError("PML thickness must contain finite integer values.")
+        integer = int(numeric)
+        if integer < 0 or integer > torch.iinfo(torch.int32).max:
+            raise ValueError("PML thickness values must fit in non-negative int32.")
+        normalized.append(integer)
+    return torch.tensor(normalized, dtype=torch.int32)
 
 
 class TVRegularization(nn.Module):
@@ -734,8 +797,29 @@ def build_pml_phi(x0,xm,y0,ym,z0,zm,nstep,PML,device):
     z0EPhi1, z0EPhi2, z0HPhi1, z0HPhi2,
     zmEPhi1, zmEPhi2, zmHPhi1, zmHPhi2) = [torch.empty(0) for _ in range(24)]
 
-    if PML==None:
-        PML=(None,None,None,None,None,None,None,None,None,None,None,None,None,None,None,None,None,None,None,None,None,None,None,None)
+    descriptors = (x0, xm, y0, ym, z0, zm)
+    if PML is None:
+        PML = (None,) * 24
+    elif not isinstance(PML, (list, tuple)) or len(PML) != 24:
+        raise ValueError("PML must contain exactly 24 CPML auxiliary tensors.")
+
+    for face, descriptor in enumerate(descriptors):
+        group = PML[4 * face:4 * face + 4]
+        supplied = tuple(value is not None for value in group)
+        if any(supplied) and not all(supplied):
+            raise ValueError(
+                f"CPML face {face} must provide all four auxiliary tensors or none."
+            )
+        if all(supplied) and not all(torch.is_tensor(value) for value in group):
+            raise TypeError(f"CPML face {face} state must contain PyTorch tensors.")
+        if (
+            all(supplied)
+            and descriptor.numel() == 0
+            and any(value.numel() != 0 for value in group)
+        ):
+            raise ValueError(
+                f"CPML face {face} state must be empty for a zero-thickness boundary."
+            )
 
     if x0.numel()!=0 and PML[0]==None:
         x0EPhi1=torch.zeros((nstep, int(x0[2]-x0[1]+1), int(x0[4]-x0[3]), int(x0[6]-x0[5]+1)),dtype=torch.float, device=device)
@@ -811,6 +895,50 @@ def build_pml_phi(x0,xm,y0,ym,z0,zm,nstep,PML,device):
         z0EPhi1,z0EPhi2,z0HPhi1,z0HPhi2,
         zmEPhi1,zmEPhi2,zmHPhi1,zmHPhi2,
     )
+
+    def expected_shapes(axis, descriptor):
+        if descriptor.numel() == 0:
+            return (None,) * 4
+        a = int(descriptor[2] - descriptor[1])
+        b = int(descriptor[4] - descriptor[3])
+        c = int(descriptor[6] - descriptor[5])
+        if axis == 0:
+            return (
+                (nstep, a + 1, b, c + 1),
+                (nstep, a + 1, b + 1, c),
+                (nstep, a, b + 1, c),
+                (nstep, a, b, c + 1),
+            )
+        if axis == 1:
+            return (
+                (nstep, a, b + 1, c + 1),
+                (nstep, a + 1, b + 1, c),
+                (nstep, a + 1, b, c),
+                (nstep, a, b, c + 1),
+            )
+        return (
+            (nstep, a, b + 1, c + 1),
+            (nstep, a + 1, b, c + 1),
+            (nstep, a + 1, b, c),
+            (nstep, a, b + 1, c),
+        )
+
+    for face, descriptor in enumerate(descriptors):
+        group = tensors[4 * face:4 * face + 4]
+        for component, (tensor, expected) in enumerate(
+            zip(group, expected_shapes(face // 2, descriptor))
+        ):
+            if expected is None:
+                if tensor.numel() != 0:
+                    raise ValueError(
+                        f"CPML face {face} component {component} must be empty."
+                    )
+            elif tuple(tensor.shape) != expected:
+                raise ValueError(
+                    f"CPML face {face} component {component} has shape "
+                    f"{tuple(tensor.shape)}; expected {expected}."
+                )
+
     return tuple(
         tensor.to(device=device, dtype=torch.float32).contiguous()
         for tensor in tensors

{deepgpr-0.0.16 → deepgpr-0.0.18}/src/DeepGPR/compute2.py

@@ -700,8 +700,12 @@ class DeepGPR(torch.autograd.Function):
         """
         
         source_amplitudes = source_amplitudes.contiguous()
-        source_location=source_location.to(torch.int32).contiguous()
-        receiver_location=receiver_location.to(torch.int32).contiguous()
+        source_location=source_location.to(
+            device=device, dtype=torch.int32
+        ).contiguous()
+        receiver_location=receiver_location.to(
+            device=device, dtype=torch.int32
+        ).contiguous()
         c_lib = get_deepgpr_lib(device)
         set_library_fdtd_order(c_lib, fdtd_order)
         ctx.save_for_backward(mu_r, source_location, receiver_location,

{deepgpr-0.0.16 → deepgpr-0.0.18}/src/DeepGPR/lib/deepgpr.cu

@@ -175,6 +175,146 @@ __device__ __forceinline__ float load_wavefield_value_device(
     } \
 } while (0)
 
+static int clamp_cpml_index(int value, int extent)
+{
+    if (value < 0) return 0;
+    if (value > extent) return extent;
+    return value;
+}
+
+/*
+ * Partition the CPML shell into six disjoint boxes. The x boxes own edges and
+ * corners, the y boxes exclude the x boxes, and the z boxes exclude both.
+ * Return the largest box so one grid row can be assigned to each face.
+ */
+static long long cpml_max_region_cells(
+    int NX, int NY, int NZ,
+    int pml0, int pml1, int pml2, int pml3, int pml4, int pml5,
+    int electric)
+{
+    int x_low_end = pml0 > 0
+        ? clamp_cpml_index(pml0 + (electric ? 1 : 0), NX) : 0;
+    int y_low_end = pml2 > 0
+        ? clamp_cpml_index(pml2 + (electric ? 1 : 0), NY) : 0;
+    int z_low_end = pml4 > 0
+        ? clamp_cpml_index(pml4 + (electric ? 1 : 0), NZ) : 0;
+    int x_high_begin = pml1 > 0
+        ? clamp_cpml_index(NX - 1 - pml1, NX) : NX;
+    int y_high_begin = pml3 > 0
+        ? clamp_cpml_index(NY - 1 - pml3, NY) : NY;
+    int z_high_begin = pml5 > 0
+        ? clamp_cpml_index(NZ - 1 - pml5, NZ) : NZ;
+
+    int x_mid_end = x_high_begin > x_low_end ? x_high_begin : x_low_end;
+    int y_mid_end = y_high_begin > y_low_end ? y_high_begin : y_low_end;
+    int z_mid_end = z_high_begin > z_low_end ? z_high_begin : z_low_end;
+    long long x_mid = x_mid_end - x_low_end;
+    long long y_mid = y_mid_end - y_low_end;
+
+    long long regions[6] = {
+        (long long)x_low_end * NY * NZ,
+        pml1 > 0 ? (long long)(NX - x_mid_end) * NY * NZ : 0,
+        x_mid * (long long)y_low_end * NZ,
+        pml3 > 0 ? x_mid * (long long)(NY - y_mid_end) * NZ : 0,
+        x_mid * y_mid * (long long)z_low_end,
+        pml5 > 0 ? x_mid * y_mid * (long long)(NZ - z_mid_end) : 0,
+    };
+    long long maximum = 0;
+    for (int region = 0; region < 6; ++region) {
+        if (regions[region] > maximum) maximum = regions[region];
+    }
+    return maximum;
+}
+
+template<int ELECTRIC>
+__device__ inline bool map_cpml_region_work(
+    int region, long long local,
+    int NX, int NY, int NZ,
+    int pml0, int pml1, int pml2, int pml3, int pml4, int pml5,
+    long long* i, long long* j, long long* k)
+{
+    int x_low_end = pml0 > 0 ? pml0 + ELECTRIC : 0;
+    int y_low_end = pml2 > 0 ? pml2 + ELECTRIC : 0;
+    int z_low_end = pml4 > 0 ? pml4 + ELECTRIC : 0;
+    if (x_low_end < 0) x_low_end = 0;
+    if (y_low_end < 0) y_low_end = 0;
+    if (z_low_end < 0) z_low_end = 0;
+    if (x_low_end > NX) x_low_end = NX;
+    if (y_low_end > NY) y_low_end = NY;
+    if (z_low_end > NZ) z_low_end = NZ;
+
+    int x_high_begin = pml1 > 0 ? NX - 1 - pml1 : NX;
+    int y_high_begin = pml3 > 0 ? NY - 1 - pml3 : NY;
+    int z_high_begin = pml5 > 0 ? NZ - 1 - pml5 : NZ;
+    if (x_high_begin < 0) x_high_begin = 0;
+    if (y_high_begin < 0) y_high_begin = 0;
+    if (z_high_begin < 0) z_high_begin = 0;
+
+    int x_mid_end = x_high_begin > x_low_end ? x_high_begin : x_low_end;
+    int y_mid_end = y_high_begin > y_low_end ? y_high_begin : y_low_end;
+    int z_mid_end = z_high_begin > z_low_end ? z_high_begin : z_low_end;
+
+    int i0 = 0, ni = 0;
+    int j0 = 0, nj = 0;
+    int k0 = 0, nk = 0;
+    switch (region) {
+        case 0:
+            ni = x_low_end;
+            nj = NY;
+            nk = NZ;
+            break;
+        case 1:
+            if (pml1 <= 0) return false;
+            i0 = x_mid_end;
+            ni = NX - x_mid_end;
+            nj = NY;
+            nk = NZ;
+            break;
+        case 2:
+            i0 = x_low_end;
+            ni = x_mid_end - x_low_end;
+            nj = y_low_end;
+            nk = NZ;
+            break;
+        case 3:
+            if (pml3 <= 0) return false;
+            i0 = x_low_end;
+            ni = x_mid_end - x_low_end;
+            j0 = y_mid_end;
+            nj = NY - y_mid_end;
+            nk = NZ;
+            break;
+        case 4:
+            i0 = x_low_end;
+            ni = x_mid_end - x_low_end;
+            j0 = y_low_end;
+            nj = y_mid_end - y_low_end;
+            nk = z_low_end;
+            break;
+        case 5:
+            if (pml5 <= 0) return false;
+            i0 = x_low_end;
+            ni = x_mid_end - x_low_end;
+            j0 = y_low_end;
+            nj = y_mid_end - y_low_end;
+            k0 = z_mid_end;
+            nk = NZ - z_mid_end;
+            break;
+        default:
+            return false;
+    }
+
+    if (ni <= 0 || nj <= 0 || nk <= 0) return false;
+    long long jk = (long long)nj * nk;
+    long long region_size = (long long)ni * jk;
+    if (local >= region_size) return false;
+    *i = i0 + local / jk;
+    long long remainder = local % jk;
+    *j = j0 + remainder / nk;
+    *k = k0 + remainder % nk;
+    return true;
+}
+
 static int g_fdtd_order = 2;
 
 DEEPGPR_API int deepgpr_abi_version(void)
@@ -439,7 +579,7 @@ __global__ void build_update_coeffs_gpu(const float* __restrict__ eps_r_pad, con
     float* __restrict__ ch_hist, float* __restrict__ ch_curl, float* __restrict__ ch_rhs,
     int NX_FIELDS, int NY_FIELDS, int NZ_FIELDS, float dt, float dx) 
 {
-    long long idx = blockIdx.x * blockDim.x + threadIdx.x;
+    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
     long long ny_nz = (long long)NY_FIELDS * NZ_FIELDS;
     if (idx >= (long long)NX_FIELDS * ny_nz) return;
 
@@ -489,7 +629,7 @@ __global__ void sample_receivers_gpu(
     int NX, int NY, int NZ, int N_ITER, int receiver_component)
 {
     long long field_stride = (long long)NX * NY * NZ;
-    long long work = blockIdx.x * blockDim.x + threadIdx.x;
+    long long work = (long long)blockIdx.x * blockDim.x + threadIdx.x;
     long long total = (long long)step * NRX;
     if (work >= total) return;
 
@@ -530,7 +670,7 @@ __global__ void inject_sources_gpu(
     int NX, int NY, int NZ, int nsrc, int polarisation, int nt) 
 {
     long long field_stride = (long long)NX * NY * NZ;
-    long long work = blockIdx.x * blockDim.x + threadIdx.x;
+    long long work = (long long)blockIdx.x * blockDim.x + threadIdx.x;
     long long total = (long long)step * nsrc;
     if (work >= total) return;
 
@@ -547,9 +687,10 @@ __global__ void inject_sources_gpu(
     long long id3 = i * NY * NZ + j * NZ + k;
     long long id4 = (long long)s * field_stride + id3;
 
-    if (polarisation == 0) Ex[id4] -= ce_rhs[id3] * scale;
-    else if (polarisation == 1) Ey[id4] -= ce_rhs[id3] * scale;
-    else if (polarisation == 2) Ez[id4] -= ce_rhs[id3] * scale;
+    float value = -ce_rhs[id3] * scale;
+    if (polarisation == 0) atomicAdd(&Ex[id4], value);
+    else if (polarisation == 1) atomicAdd(&Ey[id4], value);
+    else if (polarisation == 2) atomicAdd(&Ez[id4], value);
 }
 
 
@@ -577,9 +718,9 @@ __global__ void inject_sources_and_sample_gpu(
         float scale = srcwaveforms[src * nt + iteration] * dipole_length / (dx * dy * dz);
         float value = ce_rhs[material_idx] * scale;
 
-        if (source_component == 0) Ex[field_idx] -= value;
-        else if (source_component == 1) Ey[field_idx] -= value;
-        else Ez[field_idx] -= value;
+        if (source_component == 0) atomicAdd(&Ex[field_idx], -value);
+        else if (source_component == 1) atomicAdd(&Ey[field_idx], -value);
+        else atomicAdd(&Ez[field_idx], -value);
     }
 
     __syncthreads();
@@ -608,7 +749,7 @@ __global__ void update_e_gpu(
 {
     long long ny_nz = (long long)NY_FIELDS * NZ_FIELDS;
     long long field_stride = (long long)NX_FIELDS * ny_nz;
-    long long work = blockIdx.x * blockDim.x + threadIdx.x;
+    long long work = (long long)blockIdx.x * blockDim.x + threadIdx.x;
     if (work >= (long long)step * field_stride) return;
 
     long long idx = work % field_stride;
@@ -667,16 +808,16 @@ __global__ void cpml_e_gpu(
 {
     long long ny_nz = (long long)NY_FIELDS * NZ_FIELDS;
     long long field_stride = (long long)NX_FIELDS * ny_nz;
-    long long work = blockIdx.x * blockDim.x + threadIdx.x;
-    if (work >= (long long)step * field_stride) return;
-
-    int s = (int)(work / field_stride);
-    long long idx = work % field_stride;
-
-    long long i = idx / ny_nz;
-    long long rem = idx % ny_nz;
-    long long j = rem / NZ_FIELDS;
-    long long k = rem % NZ_FIELDS;
+    long long region_work = (long long)blockIdx.x * blockDim.x + threadIdx.x;
+    int region = (int)blockIdx.y;
+    int s = (int)blockIdx.z;
+    long long i, j, k;
+    if (s >= step ||
+        !map_cpml_region_work<1>(region, region_work,
+        NX_FIELDS, NY_FIELDS, NZ_FIELDS,
+        pml0, pml1, pml2, pml3, pml4, pml5, &i, &j, &k)) return;
+    long long idx = i * ny_nz + j * NZ_FIELDS + k;
+    long long work = (long long)s * field_stride + idx;
 
     bool in_x0 = (pml0 > 0 && i > 0 && i <= pml0 && j < NY_FIELDS && k < NZ_FIELDS);
     bool in_xm = (pml1 > 0 && i >= NX_FIELDS - 1 - pml1 && i < NX_FIELDS - 1 && j < NY_FIELDS && k < NZ_FIELDS);
@@ -818,7 +959,7 @@ __global__ void update_h_gpu(
 {
     long long ny_nz = (long long)NY_FIELDS * NZ_FIELDS;
     long long field_stride = (long long)NX_FIELDS * ny_nz;
-    long long work = blockIdx.x * blockDim.x + threadIdx.x;
+    long long work = (long long)blockIdx.x * blockDim.x + threadIdx.x;
     if (work >= (long long)step * field_stride) return;
 
     long long idx = work % field_stride;
@@ -862,7 +1003,7 @@ __global__ void adjoint_e_gpu(
 {
     long long ny_nz = (long long)NY * NZ;
     long long field_stride = (long long)NX * ny_nz;
-    long long work = blockIdx.x * blockDim.x + threadIdx.x;
+    long long work = (long long)blockIdx.x * blockDim.x + threadIdx.x;
     if (work >= (long long)step * field_stride) return;
     long long idx = work % field_stride;
     long long i = idx / ny_nz;
@@ -904,7 +1045,7 @@ __global__ void adjoint_h_gpu(
 {
     long long ny_nz = (long long)NY * NZ;
     long long field_stride = (long long)NX * ny_nz;
-    long long work = blockIdx.x * blockDim.x + threadIdx.x;
+    long long work = (long long)blockIdx.x * blockDim.x + threadIdx.x;
     if (work >= (long long)step * field_stride) return;
     long long idx = work % field_stride;
     long long i = idx / ny_nz;
@@ -961,16 +1102,16 @@ __global__ void cpml_h_gpu(
 {
     long long ny_nz = (long long)NY_FIELDS * NZ_FIELDS;
     long long field_stride = (long long)NX_FIELDS * ny_nz;
-    long long work = blockIdx.x * blockDim.x + threadIdx.x;
-    if (work >= (long long)step * field_stride) return;
-
-    int s = (int)(work / field_stride);
-    long long idx = work % field_stride;
-
-    long long i = idx / ny_nz;
-    long long rem = idx % ny_nz;
-    long long j = rem / NZ_FIELDS;
-    long long k = rem % NZ_FIELDS;
+    long long region_work = (long long)blockIdx.x * blockDim.x + threadIdx.x;
+    int region = (int)blockIdx.y;
+    int s = (int)blockIdx.z;
+    long long i, j, k;
+    if (s >= step ||
+        !map_cpml_region_work<0>(region, region_work,
+        NX_FIELDS, NY_FIELDS, NZ_FIELDS,
+        pml0, pml1, pml2, pml3, pml4, pml5, &i, &j, &k)) return;
+    long long idx = i * ny_nz + j * NZ_FIELDS + k;
+    long long work = (long long)s * field_stride + idx;
 
     bool in_x0 = (pml0 > 0 && i < pml0 && j < NY_FIELDS && k < NZ_FIELDS);
     bool in_xm = (pml1 > 0 && i >= NX_FIELDS - 1 - pml1 && i < NX_FIELDS - 1 && j < NY_FIELDS && k < NZ_FIELDS);
@@ -1115,14 +1256,16 @@ __global__ void adjoint_cpml_e_gpu(
 {
     long long ny_nz = (long long)NY * NZ;
     long long field_stride = (long long)NX * ny_nz;
-    long long work = blockIdx.x * blockDim.x + threadIdx.x;
-    if (work >= (long long)step * field_stride) return;
-    int s = (int)(work / field_stride);
-    long long idx = work % field_stride;
-    long long i = idx / ny_nz;
-    long long rem = idx % ny_nz;
-    long long j = rem / NZ;
-    long long k = rem % NZ;
+    long long region_work = (long long)blockIdx.x * blockDim.x + threadIdx.x;
+    int region = (int)blockIdx.y;
+    int s = (int)blockIdx.z;
+    long long i, j, k;
+    if (s >= step ||
+        !map_cpml_region_work<1>(region, region_work,
+        NX, NY, NZ, pml0, pml1, pml2, pml3, pml4, pml5,
+        &i, &j, &k)) return;
+    long long idx = i * ny_nz + j * NZ + k;
+    long long work = (long long)s * field_stride + idx;
     float upd = update[idx];
 
 #define APPLY_E_PML_GPU(R, P, p, q, stride, coord, n, spacing, field, source, sign) \
@@ -1215,14 +1358,16 @@ __global__ void adjoint_cpml_h_gpu(
 {
     long long ny_nz = (long long)NY * NZ;
     long long field_stride = (long long)NX * ny_nz;
-    long long work = blockIdx.x * blockDim.x + threadIdx.x;
-    if (work >= (long long)step * field_stride) return;
-    int s = (int)(work / field_stride);
-    long long idx = work % field_stride;
-    long long i = idx / ny_nz;
-    long long rem = idx % ny_nz;
-    long long j = rem / NZ;
-    long long k = rem % NZ;
+    long long region_work = (long long)blockIdx.x * blockDim.x + threadIdx.x;
+    int region = (int)blockIdx.y;
+    int s = (int)blockIdx.z;
+    long long i, j, k;
+    if (s >= step ||
+        !map_cpml_region_work<0>(region, region_work,
+        NX, NY, NZ, pml0, pml1, pml2, pml3, pml4, pml5,
+        &i, &j, &k)) return;
+    long long idx = i * ny_nz + j * NZ + k;
+    long long work = (long long)s * field_stride + idx;
     float upd = update[idx];
 
 #define APPLY_H_PML_GPU(R, P, p, q, stride, coord, n, spacing, field, source, sign) \
@@ -1324,7 +1469,7 @@ __global__ void adjoint_receivers_gpu(
     int NX, int NY, int NZ, int nsr, int polarisation, int iterations
 ){
     long long field_stride = (long long)NX * NY * NZ;
-    long long work = blockIdx.x * blockDim.x + threadIdx.x;
+    long long work = (long long)blockIdx.x * blockDim.x + threadIdx.x;
     long long total = (long long)step * nsr;
     if (work >= total) return;
 
@@ -1339,9 +1484,9 @@ __global__ void adjoint_receivers_gpu(
     float waveform_value = srcwaveforms[index];
     long long id4 = (long long)s * field_stride + i * NY * NZ + j * NZ + k;
 
-    if (polarisation == 0) lambda_ex[id4] += waveform_value;
-    else if (polarisation == 1) lambda_ey[id4] += waveform_value;
-    else if (polarisation == 2) lambda_ez[id4] += waveform_value;
+    if (polarisation == 0) atomicAdd(&lambda_ex[id4], waveform_value);
+    else if (polarisation == 1) atomicAdd(&lambda_ey[id4], waveform_value);
+    else if (polarisation == 2) atomicAdd(&lambda_ez[id4], waveform_value);
 }
 
 
@@ -1355,7 +1500,7 @@ __global__ void adjoint_source_injection_gpu(
     float* __restrict__ grad_source)
 {
     long long field_stride = (long long)NX * NY * NZ;
-    long long work = blockIdx.x * blockDim.x + threadIdx.x;
+    long long work = (long long)blockIdx.x * blockDim.x + threadIdx.x;
     long long total = (long long)step * nsrc;
     if (work >= total) return;
 
@@ -1392,7 +1537,7 @@ __global__ void save_e_snapshot_gpu(
 {
     long long nx1 = NX - 1, ny1 = NY - 1, nz1 = NZ - 1;
     long long total = nx1 * ny1 * nz1;
-    long long work = blockIdx.x * blockDim.x + threadIdx.x;
+    long long work = (long long)blockIdx.x * blockDim.x + threadIdx.x;
     if (work >= (long long)step * total) return;
 
     int s = (int)(work / total);
@@ -1425,7 +1570,7 @@ __global__ void save_rhs_snapshot_gpu(
 {
     long long nx1 = NX - 1, ny1 = NY - 1, nz1 = NZ - 1;
     long long total = nx1 * ny1 * nz1;
-    long long work = blockIdx.x * blockDim.x + threadIdx.x;
+    long long work = (long long)blockIdx.x * blockDim.x + threadIdx.x;
     if (work >= (long long)step * total) return;
 
     int s = (int)(work / total);
@@ -1491,7 +1636,7 @@ __global__ void accumulate_material_gradients_gpu(
     long long snap_stride = (long long)step * total_cells;
     long long component_stride = (long long)nt_saved * snap_stride;
     int components = (fwi_mode == 3) ? 3 : 1;
-    long long idx = blockIdx.x * blockDim.x + threadIdx.x;
+    long long idx = (long long)blockIdx.x * blockDim.x + threadIdx.x;
 
     if (idx >= total_cells) return;
 
@@ -1666,6 +1811,20 @@ DEEPGPR_API void forward(const float* __restrict__ eps_r_pad, const float* __res
     long long total_fields = (long long)NX_FIELDS * NY_FIELDS * NZ_FIELDS;
     dim3 grid_material(CEIL_DIV(total_fields, blockSize));
     dim3 grid_fields(CEIL_DIV((long long)step * total_fields, blockSize));
+    long long cpml_e_region_max = cpml_max_region_cells(
+        NX_FIELDS, NY_FIELDS, NZ_FIELDS,
+        pml0, pml1, pml2, pml3, pml4, pml5, 1);
+    long long cpml_h_region_max = cpml_max_region_cells(
+        NX_FIELDS, NY_FIELDS, NZ_FIELDS,
+        pml0, pml1, pml2, pml3, pml4, pml5, 0);
+    int has_cpml_e = has_cpml && cpml_e_region_max > 0;
+    int has_cpml_h = has_cpml && cpml_h_region_max > 0;
+    dim3 grid_cpml_e(
+        has_cpml_e ? CEIL_DIV(cpml_e_region_max, blockSize) : 1,
+        6, step);
+    dim3 grid_cpml_h(
+        has_cpml_h ? CEIL_DIV(cpml_h_region_max, blockSize) : 1,
+        6, step);
 
     build_update_coeffs_gpu<<<grid_material, blockSize, 0, stream_comp>>>(eps_r_pad, sigma_pad, mu_r_pad, ce_hist, ce_curl, ce_rhs, ch_hist, ch_curl, ch_rhs, NX_FIELDS, NY_FIELDS, NZ_FIELDS, dt, dx);
     CUDA_CHECK_LAST();
@@ -1717,8 +1876,8 @@ DEEPGPR_API void forward(const float* __restrict__ eps_r_pad, const float* __res
         if (fdtd_order == 8) {
             update_h_gpu<8><<<grid_fields, blockSize, 0, stream_comp>>>(ch_hist, ch_curl, Ex, Ey, Ez, Hx, Hy, Hz, step, NX_FIELDS, NY_FIELDS, NZ_FIELDS, dx, dy, dz);
             CUDA_CHECK_LAST();
-            if (has_cpml) {
-                cpml_h_gpu<8><<<grid_fields, blockSize, 0, stream_comp>>>(
+            if (has_cpml_h) {
+                cpml_h_gpu<8><<<grid_cpml_h, blockSize, 0, stream_comp>>>(
                     Ex, Ey, Ez, Hx, Hy, Hz, dx, dy, dz, step, NX_FIELDS, NY_FIELDS, NZ_FIELDS,
                     pml0, pml1, pml2, pml3, pml4, pml5, x0HR, xmHR, y0HR, ymHR, z0HR, zmHR, ch_rhs,
                     x0HPhi1, x0HPhi2, xmHPhi1, xmHPhi2, y0HPhi1, y0HPhi2, ymHPhi1, ymHPhi2, z0HPhi1, z0HPhi2, zmHPhi1, zmHPhi2);
@@ -1726,8 +1885,8 @@ DEEPGPR_API void forward(const float* __restrict__ eps_r_pad, const float* __res
             }
             update_e_gpu<8><<<grid_fields, blockSize, 0, stream_comp>>>(ce_hist, ce_curl, Ex, Ey, Ez, Hx, Hy, Hz, step, NX_FIELDS, NY_FIELDS, NZ_FIELDS, dx, dy, dz);
             CUDA_CHECK_LAST();
-            if (has_cpml) {
-                cpml_e_gpu<8><<<grid_fields, blockSize, 0, stream_comp>>>(
+            if (has_cpml_e) {
+                cpml_e_gpu<8><<<grid_cpml_e, blockSize, 0, stream_comp>>>(
                     Ex, Ey, Ez, Hx, Hy, Hz, dx, dy, dz, step, NX_FIELDS, NY_FIELDS, NZ_FIELDS,
                     pml0, pml1, pml2, pml3, pml4, pml5, x0ER, xmER, y0ER, ymER, z0ER, zmER, ce_rhs,
                     x0EPhi1, x0EPhi2, xmEPhi1, xmEPhi2, y0EPhi1, y0EPhi2, ymEPhi1, ymEPhi2, z0EPhi1, z0EPhi2, zmEPhi1, zmEPhi2);
@@ -1736,8 +1895,8 @@ DEEPGPR_API void forward(const float* __restrict__ eps_r_pad, const float* __res
         } else if (fdtd_order == 4) {
             update_h_gpu<4><<<grid_fields, blockSize, 0, stream_comp>>>(ch_hist, ch_curl, Ex, Ey, Ez, Hx, Hy, Hz, step, NX_FIELDS, NY_FIELDS, NZ_FIELDS, dx, dy, dz);
             CUDA_CHECK_LAST();
-            if (has_cpml) {
-                cpml_h_gpu<4><<<grid_fields, blockSize, 0, stream_comp>>>(
+            if (has_cpml_h) {
+                cpml_h_gpu<4><<<grid_cpml_h, blockSize, 0, stream_comp>>>(
                     Ex, Ey, Ez, Hx, Hy, Hz, dx, dy, dz, step, NX_FIELDS, NY_FIELDS, NZ_FIELDS,
                     pml0, pml1, pml2, pml3, pml4, pml5, x0HR, xmHR, y0HR, ymHR, z0HR, zmHR, ch_rhs,
                     x0HPhi1, x0HPhi2, xmHPhi1, xmHPhi2, y0HPhi1, y0HPhi2, ymHPhi1, ymHPhi2, z0HPhi1, z0HPhi2, zmHPhi1, zmHPhi2);
@@ -1745,8 +1904,8 @@ DEEPGPR_API void forward(const float* __restrict__ eps_r_pad, const float* __res
             }
             update_e_gpu<4><<<grid_fields, blockSize, 0, stream_comp>>>(ce_hist, ce_curl, Ex, Ey, Ez, Hx, Hy, Hz, step, NX_FIELDS, NY_FIELDS, NZ_FIELDS, dx, dy, dz);
             CUDA_CHECK_LAST();
-            if (has_cpml) {
-                cpml_e_gpu<4><<<grid_fields, blockSize, 0, stream_comp>>>(
+            if (has_cpml_e) {
+                cpml_e_gpu<4><<<grid_cpml_e, blockSize, 0, stream_comp>>>(
                     Ex, Ey, Ez, Hx, Hy, Hz, dx, dy, dz, step, NX_FIELDS, NY_FIELDS, NZ_FIELDS,
                     pml0, pml1, pml2, pml3, pml4, pml5, x0ER, xmER, y0ER, ymER, z0ER, zmER, ce_rhs,
                     x0EPhi1, x0EPhi2, xmEPhi1, xmEPhi2, y0EPhi1, y0EPhi2, ymEPhi1, ymEPhi2, z0EPhi1, z0EPhi2, zmEPhi1, zmEPhi2);
@@ -1755,8 +1914,8 @@ DEEPGPR_API void forward(const float* __restrict__ eps_r_pad, const float* __res
         } else {
             update_h_gpu<2><<<grid_fields, blockSize, 0, stream_comp>>>(ch_hist, ch_curl, Ex, Ey, Ez, Hx, Hy, Hz, step, NX_FIELDS, NY_FIELDS, NZ_FIELDS, dx, dy, dz);
             CUDA_CHECK_LAST();
-            if (has_cpml) {
-                cpml_h_gpu<2><<<grid_fields, blockSize, 0, stream_comp>>>(
+            if (has_cpml_h) {
+                cpml_h_gpu<2><<<grid_cpml_h, blockSize, 0, stream_comp>>>(
                     Ex, Ey, Ez, Hx, Hy, Hz, dx, dy, dz, step, NX_FIELDS, NY_FIELDS, NZ_FIELDS,
                     pml0, pml1, pml2, pml3, pml4, pml5, x0HR, xmHR, y0HR, ymHR, z0HR, zmHR, ch_rhs,
                     x0HPhi1, x0HPhi2, xmHPhi1, xmHPhi2, y0HPhi1, y0HPhi2, ymHPhi1, ymHPhi2, z0HPhi1, z0HPhi2, zmHPhi1, zmHPhi2);
@@ -1764,8 +1923,8 @@ DEEPGPR_API void forward(const float* __restrict__ eps_r_pad, const float* __res
             }
             update_e_gpu<2><<<grid_fields, blockSize, 0, stream_comp>>>(ce_hist, ce_curl, Ex, Ey, Ez, Hx, Hy, Hz, step, NX_FIELDS, NY_FIELDS, NZ_FIELDS, dx, dy, dz);
             CUDA_CHECK_LAST();
-            if (has_cpml) {
-                cpml_e_gpu<2><<<grid_fields, blockSize, 0, stream_comp>>>(
+            if (has_cpml_e) {
+                cpml_e_gpu<2><<<grid_cpml_e, blockSize, 0, stream_comp>>>(
                     Ex, Ey, Ez, Hx, Hy, Hz, dx, dy, dz, step, NX_FIELDS, NY_FIELDS, NZ_FIELDS,
                     pml0, pml1, pml2, pml3, pml4, pml5, x0ER, xmER, y0ER, ymER, z0ER, zmER, ce_rhs,
                     x0EPhi1, x0EPhi2, xmEPhi1, xmEPhi2, y0EPhi1, y0EPhi2, ymEPhi1, ymEPhi2, z0EPhi1, z0EPhi2, zmEPhi1, zmEPhi2);
@@ -1964,6 +2123,20 @@ DEEPGPR_API void backward(const float* __restrict__ eps_r_pad, const float* __re
     long long total_fields = (long long)NX_FIELDS * NY_FIELDS * NZ_FIELDS;
     dim3 grid_material(CEIL_DIV(total_fields, blockSize));
     dim3 grid_fields(CEIL_DIV((long long)step * total_fields, blockSize));
+    long long cpml_e_region_max = cpml_max_region_cells(
+        NX_FIELDS, NY_FIELDS, NZ_FIELDS,
+        pml0, pml1, pml2, pml3, pml4, pml5, 1);
+    long long cpml_h_region_max = cpml_max_region_cells(
+        NX_FIELDS, NY_FIELDS, NZ_FIELDS,
+        pml0, pml1, pml2, pml3, pml4, pml5, 0);
+    int has_cpml_e = has_cpml && cpml_e_region_max > 0;
+    int has_cpml_h = has_cpml && cpml_h_region_max > 0;
+    dim3 grid_cpml_e(
+        has_cpml_e ? CEIL_DIV(cpml_e_region_max, blockSize) : 1,
+        6, step);
+    dim3 grid_cpml_h(
+        has_cpml_h ? CEIL_DIV(cpml_h_region_max, blockSize) : 1,
+        6, step);
 
     build_update_coeffs_gpu<<<grid_material, blockSize, 0, stream_comp>>>(eps_r_pad, sigma_pad, mu_r_pad, ce_hist, ce_curl, ce_rhs, ch_hist, ch_curl, ch_rhs, NX_FIELDS, NY_FIELDS, NZ_FIELDS, dt, dx);
     CUDA_CHECK_LAST();
@@ -2032,8 +2205,8 @@ DEEPGPR_API void backward(const float* __restrict__ eps_r_pad, const float* __re
         }
 
         /* E CPML^T -> E update^T -> H CPML^T -> H update^T. */
-        if (has_cpml) {
-            LAUNCH_ORDER_KERNEL(adjoint_cpml_e_gpu, grid_fields, blockSize, stream_comp, fdtd_order,
+        if (has_cpml_e) {
+            LAUNCH_ORDER_KERNEL(adjoint_cpml_e_gpu, grid_cpml_e, blockSize, stream_comp, fdtd_order,
                 lambda_ex, lambda_ey, lambda_ez, lambda_hx, lambda_hy, lambda_hz, dx, dy, dz, step, NX_FIELDS, NY_FIELDS, NZ_FIELDS,
                 pml0, pml1, pml2, pml3, pml4, pml5, x0ER, xmER, y0ER, ymER, z0ER, zmER, ce_rhs,
                 x0EPhi1, x0EPhi2, xmEPhi1, xmEPhi2, y0EPhi1, y0EPhi2, ymEPhi1, ymEPhi2,
@@ -2044,8 +2217,8 @@ DEEPGPR_API void backward(const float* __restrict__ eps_r_pad, const float* __re
             ce_hist, ce_curl, lambda_ex, lambda_ey, lambda_ez, lambda_hx, lambda_hy, lambda_hz,
             step, NX_FIELDS, NY_FIELDS, NZ_FIELDS, dx, dy, dz);
         CUDA_CHECK_LAST();
-        if (has_cpml) {
-            LAUNCH_ORDER_KERNEL(adjoint_cpml_h_gpu, grid_fields, blockSize, stream_comp, fdtd_order,
+        if (has_cpml_h) {
+            LAUNCH_ORDER_KERNEL(adjoint_cpml_h_gpu, grid_cpml_h, blockSize, stream_comp, fdtd_order,
                 lambda_ex, lambda_ey, lambda_ez, lambda_hx, lambda_hy, lambda_hz, dx, dy, dz, step, NX_FIELDS, NY_FIELDS, NZ_FIELDS,
                 pml0, pml1, pml2, pml3, pml4, pml5, x0HR, xmHR, y0HR, ymHR, z0HR, zmHR, ch_rhs,
                 x0HPhi1, x0HPhi2, xmHPhi1, xmHPhi2, y0HPhi1, y0HPhi2, ymHPhi1, ymHPhi2,

{deepgpr-0.0.16 → deepgpr-0.0.18}/src/DeepGPR.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: DeepGPR
-Version: 0.0.16
+Version: 0.0.18
 Summary: PyTorch and CUDA for GPR FWI
 Author-email: Lei Liu <liulei990222@gmail.com>
 License-Expression: MIT

{deepgpr-0.0.16 → deepgpr-0.0.18}/tests/test_discrete_adjoint.py

@@ -67,23 +67,24 @@ def _locations(shape):
     return torch.tensor([[center]], dtype=torch.int32)
 
 
-def _state_problem(order, pml):
+def _state_problem(order, pml, device="cpu"):
     torch.manual_seed(1100 + order + pml)
+    device = torch.device(device)
     shape = (9, 10, 8)
     nx, ny, nz = shape
-    x = torch.linspace(0.0, 1.0, nx)[:, None, None]
-    y = torch.linspace(0.0, 1.0, ny)[None, :, None]
-    z = torch.linspace(0.0, 1.0, nz)[None, None, :]
+    x = torch.linspace(0.0, 1.0, nx, device=device)[:, None, None]
+    y = torch.linspace(0.0, 1.0, ny, device=device)[None, :, None]
+    z = torch.linspace(0.0, 1.0, nz, device=device)[None, None, :]
     eps_r = 3.5 + 0.7 * x + 0.2 * y + 0.1 * z
     sigma = 1.0e-4 + 1.0e-4 * (x + y + z) / 3.0
     mu_r = 1.0 + 0.15 * x + 0.0 * y + 0.05 * z
-    location = _locations(shape)
-    source = torch.zeros((1, 1, 1), dtype=torch.float32)
+    location = _locations(shape).to(device)
+    source = torch.zeros((1, 1, 1), dtype=torch.float32, device=device)
     spacing = (0.020, 0.017, 0.013)
     dt = 1.0e-11
 
     initial_e, initial_h, initial_cpml = DeepGPR.checkpoint_initial_field(
-        device="cpu",
+        device=device,
         dx=spacing,
         dt=dt,
         source_amplitudes=source,
@@ -102,7 +103,7 @@ def _state_problem(order, pml):
     ]
     state = [tensor.clone().requires_grad_(True) for tensor in base_state]
     result = DeepGPR.compute(
-        device="cpu",
+        device=device,
         dx=spacing,
         dt=dt,
         source_amplitudes=source,
@@ -269,6 +270,15 @@ class DiscreteAdjointTests(unittest.TestCase):
     def test_native_cpml_single_step_dot_product_all_faces(self):
         self.assertLess(_state_problem(order=4, pml=2), 3.0e-5)
 
+    @unittest.skipUnless(torch.cuda.is_available(), "CUDA is not available")
+    def test_cuda_cpml_single_step_dot_products_orders_2_4_8(self):
+        for order in (2, 4, 8):
+            with self.subTest(order=order):
+                self.assertLess(
+                    _state_problem(order=order, pml=2, device="cuda"),
+                    2.0e-4,
+                )
+
     def test_source_waveform_gradient_matches_finite_difference(self):
         nx, ny, nt = 12, 14, 60
         eps_r = torch.full((nx, ny), 4.0)
@@ -378,6 +388,143 @@ class DiscreteAdjointTests(unittest.TestCase):
         for reference, candidate in (*zip(cpu, cuda), *zip(cuda, cuda_offload)):
             torch.testing.assert_close(reference, candidate, rtol=5.0e-4, atol=1.0e-6)
 
+    @unittest.skipUnless(torch.cuda.is_available(), "CUDA is not available")
+    def test_cuda_coincident_sources_and_receivers_accumulate(self):
+        device = torch.device("cuda")
+        nt = 48
+        eps_r = torch.full((12, 14), 4.0, device=device)
+        sigma = torch.zeros_like(eps_r)
+        first = torch.linspace(-0.5, 0.8, nt, device=device)
+        second = torch.linspace(0.2, -0.4, nt, device=device)
+        source_pair = torch.stack((first, second), dim=0).unsqueeze(-1)
+        source_sum = (first + second).reshape(1, nt, 1)
+        one_source = torch.tensor([[[5, 5, 0]]], dtype=torch.int32, device=device)
+        two_sources = one_source.repeat(1, 2, 1)
+        one_receiver = torch.tensor([[[5, 8, 0]]], dtype=torch.int32, device=device)
+
+        common = dict(
+            device=device, dx=0.02, dt=3.0e-11,
+            receiver_location=one_receiver,
+            er=eps_r, se=sigma, pmlthick=0,
+            fdtd_order=2, mode=2,
+        )
+        pair_data = DeepGPR.compute(
+            source_amplitudes=source_pair,
+            source_location=two_sources,
+            **common,
+        )[-1]
+        sum_data = DeepGPR.compute(
+            source_amplitudes=source_sum,
+            source_location=one_source,
+            **common,
+        )[-1]
+        torch.cuda.synchronize()
+        torch.testing.assert_close(pair_data, sum_data, rtol=2.0e-5, atol=1.0e-6)
+
+        def receiver_gradient(receiver_location):
+            model = eps_r.detach().clone().requires_grad_(True)
+            data = DeepGPR.compute(
+                device=device, dx=0.02, dt=3.0e-11,
+                source_amplitudes=source_sum,
+                source_location=one_source,
+                receiver_location=receiver_location,
+                er=model, se=sigma, pmlthick=0,
+                fdtd_order=2, mode=2,
+                model_gradient_sampling_interval=1,
+                wavefield_storage_dtype=torch.float32,
+            )[-1]
+            data.sum().backward()
+            return model.grad
+
+        single_gradient = receiver_gradient(one_receiver)
+        duplicate_gradient = receiver_gradient(one_receiver.repeat(1, 2, 1))
+        torch.cuda.synchronize()
+        torch.testing.assert_close(
+            duplicate_gradient, 2.0 * single_gradient, rtol=2.0e-4, atol=1.0e-6
+        )
+
+    @unittest.skipUnless(torch.cuda.is_available(), "CUDA is not available")
+    def test_cuda_large_3d_all_face_cpml_async_bounds(self):
+        device = torch.device("cuda")
+        nx, ny, nz, nt = 120, 120, 100, 5
+        eps_r = torch.full(
+            (nx, ny, nz), 4.0, dtype=torch.float32, device=device,
+            requires_grad=True,
+        )
+        sigma = torch.full_like(eps_r, 1.0e-4)
+        source = torch.zeros((1, nt, 1), dtype=torch.float32, device=device)
+        source[0, 0, 0] = 1.0
+        source_location = torch.tensor(
+            [[[60, 60, 50]]], dtype=torch.int32
+        )
+        receiver_location = torch.tensor(
+            [[[60, 62, 50]]], dtype=torch.int32
+        )
+
+        result = DeepGPR.compute(
+            device=device,
+            dx=(0.020, 0.017, 0.014),
+            dt=1.0e-11,
+            source_amplitudes=source,
+            source_location=source_location,
+            receiver_location=receiver_location,
+            er=eps_r,
+            se=sigma,
+            pmlthick=10,
+            fdtd_order=2,
+            mode=3,
+            model_gradient_sampling_interval=3,
+            wavefield_storage_dtype=torch.float16,
+            use_async_offload=True,
+        )
+        final_fields = (*result[1], *result[2])
+        loss = result[-1].square().sum()
+        loss = loss + sum(field.square().mean() for field in final_fields)
+        loss.backward()
+        torch.cuda.synchronize()
+
+        self.assertTrue(torch.isfinite(result[-1]).all().item())
+        self.assertIsNotNone(eps_r.grad)
+        self.assertTrue(torch.isfinite(eps_r.grad).all().item())
+
+    @unittest.skipUnless(torch.cuda.is_available(), "CUDA is not available")
+    def test_cuda_asymmetric_cpml_and_incomplete_sampling_block(self):
+        device = torch.device("cuda")
+        shape = (31, 29, 27)
+        nt = 11
+        eps_r = torch.full(shape, 4.0, device=device, requires_grad=True)
+        sigma = torch.full_like(eps_r, 2.0e-4, requires_grad=True)
+        source = torch.linspace(-0.3, 0.7, nt, device=device).reshape(1, nt, 1)
+        source_location = torch.tensor(
+            [[[15, 14, 13]]], dtype=torch.int32, device=device
+        )
+        receiver_location = torch.tensor(
+            [[[16, 15, 14]]], dtype=torch.int32, device=device
+        )
+
+        receiver = DeepGPR.compute(
+            device=device,
+            dx=(0.020, 0.016, 0.013),
+            dt=1.0e-11,
+            source_amplitudes=source,
+            source_location=source_location,
+            receiver_location=receiver_location,
+            er=eps_r,
+            se=sigma,
+            pmlthick=[3, 5, 4, 2, 6, 3],
+            fdtd_order=8,
+            mode=3,
+            model_gradient_sampling_interval=4,
+            wavefield_storage_dtype=torch.bfloat16,
+            use_async_offload=True,
+        )[-1]
+        receiver.square().sum().backward()
+        torch.cuda.synchronize()
+
+        self.assertTrue(torch.isfinite(receiver).all().item())
+        self.assertTrue(torch.isfinite(eps_r.grad).all().item())
+        self.assertTrue(torch.isfinite(sigma.grad).all().item())
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

{deepgpr-0.0.16 → deepgpr-0.0.18}/tests/test_numerics.py

@@ -113,6 +113,70 @@ class NumericsValidationTests(unittest.TestCase):
                 location, location, 0.02, 3.0e-11, 2, 2,
             )
 
+    def test_initialization_rejects_empty_native_launch_dimensions(self):
+        source = torch.zeros((1, 4, 1))
+        location = torch.tensor([[[1, 1, 0]]], dtype=torch.int32)
+
+        with self.assertRaisesRegex(ValueError, "model dimensions"):
+            initialization(
+                "cpu", torch.empty((0, 4)), torch.empty((0, 4)), None,
+                source, location, location, 0.02, 3.0e-11, 0, 2,
+            )
+
+        empty_cases = (
+            (
+                torch.empty((0, 1, 3), dtype=torch.int32),
+                torch.empty((0, 1, 3), dtype=torch.int32),
+                "shot",
+            ),
+            (
+                torch.empty((1, 0, 3), dtype=torch.int32),
+                location,
+                "source",
+            ),
+            (
+                location,
+                torch.empty((1, 0, 3), dtype=torch.int32),
+                "receiver",
+            ),
+        )
+        er = torch.full((4, 4), 4.0)
+        se = torch.zeros_like(er)
+        for source_location, receiver_location, message in empty_cases:
+            with self.subTest(message=message), self.assertRaisesRegex(ValueError, message):
+                initialization(
+                    "cpu", er, se, None, source,
+                    source_location, receiver_location, 0.02, 3.0e-11, 0, 2,
+                )
+
+    def test_initialization_validates_coordinates_before_int32_conversion(self):
+        er = torch.full((4, 4), 4.0)
+        se = torch.zeros_like(er)
+        source = torch.zeros((1, 4, 1))
+        location = torch.tensor([[[1, 1, 0]]], dtype=torch.int64)
+
+        fractional = location.to(torch.float64)
+        fractional[0, 0, 1] = 1.5
+        with self.assertRaisesRegex(ValueError, "integer-valued"):
+            initialization(
+                "cpu", er, se, None, source,
+                fractional, location, 0.02, 3.0e-11, 0, 2,
+            )
+
+        wrapped = location.clone()
+        wrapped[0, 0, 0] = 2**32 + 1
+        with self.assertRaisesRegex(ValueError, "out of range"):
+            initialization(
+                "cpu", er, se, None, source,
+                wrapped, location, 0.02, 3.0e-11, 0, 2,
+            )
+
+        with self.assertRaisesRegex(ValueError, "finite integer"):
+            initialization(
+                "cpu", er, se, None, source,
+                location, location, 0.02, 3.0e-11, [0, 0, 0.5, 0], 2,
+            )
+
     def test_wavefield_storage_aliases(self):
         self.assertIs(_normalize_wavefield_storage_dtype("fp32"), torch.float32)
         self.assertIs(_normalize_wavefield_storage_dtype("fp16"), torch.float16)
@@ -228,6 +292,41 @@ class NumericsValidationTests(unittest.TestCase):
             _pml_phi_elements(nx, ny, nz, nstep, pml),
         )
 
+    def test_cpml_checkpoint_shapes_are_validated_before_native_call(self):
+        nx, ny, nt = 8, 10, 4
+        er = torch.full((nx, ny), 4.0)
+        se = torch.zeros_like(er)
+        source = torch.zeros((1, nt, 1))
+        location = torch.tensor([[[4, 5, 0]]], dtype=torch.int32)
+        _, _, pml_state = DeepGPR.checkpoint_initial_field(
+            device="cpu",
+            dx=0.02,
+            dt=3.0e-11,
+            source_amplitudes=source,
+            source_location=location,
+            receiver_location=location,
+            er=er,
+            se=se,
+            pmlthick=2,
+        )
+        bad_state = list(pml_state)
+        bad_state[0] = bad_state[0][..., :-1]
+
+        with self.assertRaisesRegex(ValueError, "CPML face 0 component 0"):
+            DeepGPR.compute(
+                device="cpu",
+                dx=0.02,
+                dt=3.0e-11,
+                source_amplitudes=source,
+                source_location=location,
+                receiver_location=location,
+                er=er,
+                se=se,
+                pmlthick=2,
+                PML=tuple(bad_state),
+                mode=2,
+            )
+
     def test_anisotropic_cpml_coefficients_use_axis_spacing(self):
         nx, ny, nz = 10, 12, 14
         er = torch.full((nx, ny, nz), 4.0)