DeepGPR 0.0.17__tar.gz → 0.0.18__tar.gz

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

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

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

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

{deepgpr-0.0.17 → 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.17 → 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.17 → deepgpr-0.0.18}/src/DeepGPR/lib/deepgpr.cu

@@ -185,10 +185,9 @@ static int clamp_cpml_index(int value, int extent)
 /*
  * 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 number of compact boundary cells without passing a large layout
- * object to CUDA kernels, which is unreliable with CUDA 11.8's old frontend.
+ * Return the largest box so one grid row can be assigned to each face.
  */
-static long long cpml_cells_per_shot(
+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)
@@ -212,30 +211,34 @@ static long long cpml_cells_per_shot(
     long long x_mid = x_mid_end - x_low_end;
     long long y_mid = y_mid_end - y_low_end;
 
-    long long cells = (long long)x_low_end * NY * NZ;
-    if (pml1 > 0) cells += (long long)(NX - x_mid_end) * NY * NZ;
-    cells += x_mid * y_low_end * NZ;
-    if (pml3 > 0) cells += x_mid * (NY - y_mid_end) * NZ;
-    cells += x_mid * y_mid * z_low_end;
-    if (pml5 > 0) cells += x_mid * y_mid * (NZ - z_mid_end);
-    return cells;
+    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_work(
-    long long compact_work, int step, long long cells_per_shot,
+__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,
-    int* shot, long long* i, long long* j, long long* k)
+    long long* i, long long* j, long long* k)
 {
-    if (cells_per_shot == 0 ||
-        compact_work >= (long long)step * cells_per_shot) {
-        return false;
-    }
-
     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;
@@ -250,57 +253,61 @@ __device__ inline bool map_cpml_work(
     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 x_mid = x_mid_end - x_low_end;
-    int y_mid = y_mid_end - y_low_end;
-
-    long long local = compact_work % cells_per_shot;
-    long long region_size = (long long)x_low_end * NY * NZ;
-    int i0, j0, nj, k0, nk;
-    if (local < region_size) {
-        i0 = 0;
-        j0 = 0; nj = NY;
-        k0 = 0; nk = NZ;
-    } else {
-        local -= region_size;
-        region_size = pml1 > 0 ? (long long)(NX - x_mid_end) * NY * NZ : 0;
-        if (local < region_size) {
+
+    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;
-            j0 = 0; nj = NY;
-            k0 = 0; nk = NZ;
-        } else {
-            local -= region_size;
-            region_size = (long long)x_mid * y_low_end * NZ;
-            if (local < region_size) {
-                i0 = x_low_end;
-                j0 = 0; nj = y_low_end;
-                k0 = 0; nk = NZ;
-            } else {
-                local -= region_size;
-                region_size = pml3 > 0 ? (long long)x_mid * (NY - y_mid_end) * NZ : 0;
-                if (local < region_size) {
-                    i0 = x_low_end;
-                    j0 = y_mid_end; nj = NY - y_mid_end;
-                    k0 = 0; nk = NZ;
-                } else {
-                    local -= region_size;
-                    region_size = (long long)x_mid * y_mid * z_low_end;
-                    if (local < region_size) {
-                        i0 = x_low_end;
-                        j0 = y_low_end; nj = y_mid;
-                        k0 = 0; nk = z_low_end;
-                    } else {
-                        local -= region_size;
-                        i0 = x_low_end;
-                        j0 = y_low_end; nj = y_mid;
-                        k0 = z_mid_end; nk = NZ - z_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;
     }
 
-    *shot = (int)(compact_work / cells_per_shot);
+    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;
@@ -572,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;
 
@@ -622,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;
 
@@ -663,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;
 
@@ -680,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);
 }
 
 
@@ -710,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();
@@ -741,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;
@@ -796,17 +804,18 @@ __global__ void cpml_e_gpu(
     float* __restrict__ y0EPhi1, float* __restrict__ y0EPhi2,
     float* __restrict__ ymEPhi1, float* __restrict__ ymEPhi2,
     float* __restrict__ z0EPhi1, float* __restrict__ z0EPhi2,
-    float* __restrict__ zmEPhi1, float* __restrict__ zmEPhi2,
-    long long compact_cells_per_shot)
+    float* __restrict__ zmEPhi1, float* __restrict__ zmEPhi2)
 {
     long long ny_nz = (long long)NY_FIELDS * NZ_FIELDS;
     long long field_stride = (long long)NX_FIELDS * ny_nz;
-    long long compact_work = blockIdx.x * blockDim.x + threadIdx.x;
-    int s;
+    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 (!map_cpml_work<1>(compact_work, step, compact_cells_per_shot,
+    if (s >= step ||
+        !map_cpml_region_work<1>(region, region_work,
         NX_FIELDS, NY_FIELDS, NZ_FIELDS,
-        pml0, pml1, pml2, pml3, pml4, pml5, &s, &i, &j, &k)) return;
+        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;
 
@@ -950,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;
@@ -994,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;
@@ -1036,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;
@@ -1089,17 +1098,18 @@ __global__ void cpml_h_gpu(
     float* __restrict__ y0HPhi1, float* __restrict__ y0HPhi2,
     float* __restrict__ ymHPhi1, float* __restrict__ ymHPhi2,
     float* __restrict__ z0HPhi1, float* __restrict__ z0HPhi2,
-    float* __restrict__ zmHPhi1, float* __restrict__ zmHPhi2,
-    long long compact_cells_per_shot)
+    float* __restrict__ zmHPhi1, float* __restrict__ zmHPhi2)
 {
     long long ny_nz = (long long)NY_FIELDS * NZ_FIELDS;
     long long field_stride = (long long)NX_FIELDS * ny_nz;
-    long long compact_work = blockIdx.x * blockDim.x + threadIdx.x;
-    int s;
+    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 (!map_cpml_work<0>(compact_work, step, compact_cells_per_shot,
+    if (s >= step ||
+        !map_cpml_region_work<0>(region, region_work,
         NX_FIELDS, NY_FIELDS, NZ_FIELDS,
-        pml0, pml1, pml2, pml3, pml4, pml5, &s, &i, &j, &k)) return;
+        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;
 
@@ -1242,17 +1252,18 @@ __global__ void adjoint_cpml_e_gpu(
     const float* z0R, const float* zmR, const float* update,
     float* x0P1, float* x0P2, float* xmP1, float* xmP2,
     float* y0P1, float* y0P2, float* ymP1, float* ymP2,
-    float* z0P1, float* z0P2, float* zmP1, float* zmP2,
-    long long compact_cells_per_shot)
+    float* z0P1, float* z0P2, float* zmP1, float* zmP2)
 {
     long long ny_nz = (long long)NY * NZ;
     long long field_stride = (long long)NX * ny_nz;
-    long long compact_work = blockIdx.x * blockDim.x + threadIdx.x;
-    int s;
+    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 (!map_cpml_work<1>(compact_work, step, compact_cells_per_shot,
+    if (s >= step ||
+        !map_cpml_region_work<1>(region, region_work,
         NX, NY, NZ, pml0, pml1, pml2, pml3, pml4, pml5,
-        &s, &i, &j, &k)) return;
+        &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];
@@ -1343,17 +1354,18 @@ __global__ void adjoint_cpml_h_gpu(
     const float* z0R, const float* zmR, const float* update,
     float* x0P1, float* x0P2, float* xmP1, float* xmP2,
     float* y0P1, float* y0P2, float* ymP1, float* ymP2,
-    float* z0P1, float* z0P2, float* zmP1, float* zmP2,
-    long long compact_cells_per_shot)
+    float* z0P1, float* z0P2, float* zmP1, float* zmP2)
 {
     long long ny_nz = (long long)NY * NZ;
     long long field_stride = (long long)NX * ny_nz;
-    long long compact_work = blockIdx.x * blockDim.x + threadIdx.x;
-    int s;
+    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 (!map_cpml_work<0>(compact_work, step, compact_cells_per_shot,
+    if (s >= step ||
+        !map_cpml_region_work<0>(region, region_work,
         NX, NY, NZ, pml0, pml1, pml2, pml3, pml4, pml5,
-        &s, &i, &j, &k)) return;
+        &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];
@@ -1457,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;
 
@@ -1472,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);
 }
 
 
@@ -1488,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;
 
@@ -1525,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);
@@ -1558,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);
@@ -1624,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;
 
@@ -1799,16 +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_cells = cpml_cells_per_shot(
+    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_cells = cpml_cells_per_shot(
+    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_cells > 0;
-    int has_cpml_h = has_cpml && cpml_h_cells > 0;
-    dim3 grid_cpml_e(CEIL_DIV((long long)step * cpml_e_cells, blockSize));
-    dim3 grid_cpml_h(CEIL_DIV((long long)step * cpml_h_cells, blockSize));
+    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();
@@ -1864,8 +1880,7 @@ DEEPGPR_API void forward(const float* __restrict__ eps_r_pad, const float* __res
                 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,
-                    cpml_h_cells);
+                    x0HPhi1, x0HPhi2, xmHPhi1, xmHPhi2, y0HPhi1, y0HPhi2, ymHPhi1, ymHPhi2, z0HPhi1, z0HPhi2, zmHPhi1, zmHPhi2);
                 CUDA_CHECK_LAST();
             }
             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);
@@ -1874,8 +1889,7 @@ DEEPGPR_API void forward(const float* __restrict__ eps_r_pad, const float* __res
                 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,
-                    cpml_e_cells);
+                    x0EPhi1, x0EPhi2, xmEPhi1, xmEPhi2, y0EPhi1, y0EPhi2, ymEPhi1, ymEPhi2, z0EPhi1, z0EPhi2, zmEPhi1, zmEPhi2);
                 CUDA_CHECK_LAST();
             }
         } else if (fdtd_order == 4) {
@@ -1885,8 +1899,7 @@ DEEPGPR_API void forward(const float* __restrict__ eps_r_pad, const float* __res
                 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,
-                    cpml_h_cells);
+                    x0HPhi1, x0HPhi2, xmHPhi1, xmHPhi2, y0HPhi1, y0HPhi2, ymHPhi1, ymHPhi2, z0HPhi1, z0HPhi2, zmHPhi1, zmHPhi2);
                 CUDA_CHECK_LAST();
             }
             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);
@@ -1895,8 +1908,7 @@ DEEPGPR_API void forward(const float* __restrict__ eps_r_pad, const float* __res
                 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,
-                    cpml_e_cells);
+                    x0EPhi1, x0EPhi2, xmEPhi1, xmEPhi2, y0EPhi1, y0EPhi2, ymEPhi1, ymEPhi2, z0EPhi1, z0EPhi2, zmEPhi1, zmEPhi2);
                 CUDA_CHECK_LAST();
             }
         } else {
@@ -1906,8 +1918,7 @@ DEEPGPR_API void forward(const float* __restrict__ eps_r_pad, const float* __res
                 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,
-                    cpml_h_cells);
+                    x0HPhi1, x0HPhi2, xmHPhi1, xmHPhi2, y0HPhi1, y0HPhi2, ymHPhi1, ymHPhi2, z0HPhi1, z0HPhi2, zmHPhi1, zmHPhi2);
                 CUDA_CHECK_LAST();
             }
             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);
@@ -1916,8 +1927,7 @@ DEEPGPR_API void forward(const float* __restrict__ eps_r_pad, const float* __res
                 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,
-                    cpml_e_cells);
+                    x0EPhi1, x0EPhi2, xmEPhi1, xmEPhi2, y0EPhi1, y0EPhi2, ymEPhi1, ymEPhi2, z0EPhi1, z0EPhi2, zmEPhi1, zmEPhi2);
                 CUDA_CHECK_LAST();
             }
         }
@@ -2113,16 +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_cells = cpml_cells_per_shot(
+    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_cells = cpml_cells_per_shot(
+    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_cells > 0;
-    int has_cpml_h = has_cpml && cpml_h_cells > 0;
-    dim3 grid_cpml_e(CEIL_DIV((long long)step * cpml_e_cells, blockSize));
-    dim3 grid_cpml_h(CEIL_DIV((long long)step * cpml_h_cells, blockSize));
+    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();
@@ -2196,7 +2210,7 @@ DEEPGPR_API void backward(const float* __restrict__ eps_r_pad, const float* __re
                 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,
-                z0EPhi1, z0EPhi2, zmEPhi1, zmEPhi2, cpml_e_cells);
+                z0EPhi1, z0EPhi2, zmEPhi1, zmEPhi2);
             CUDA_CHECK_LAST();
         }
         LAUNCH_ORDER_KERNEL(adjoint_e_gpu, grid_fields, blockSize, stream_comp, fdtd_order,
@@ -2208,7 +2222,7 @@ DEEPGPR_API void backward(const float* __restrict__ eps_r_pad, const float* __re
                 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,
-                z0HPhi1, z0HPhi2, zmHPhi1, zmHPhi2, cpml_h_cells);
+                z0HPhi1, z0HPhi2, zmHPhi1, zmHPhi2);
             CUDA_CHECK_LAST();
         }
         LAUNCH_ORDER_KERNEL(adjoint_h_gpu, grid_fields, blockSize, stream_comp, fdtd_order,

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

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

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

@@ -388,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.17 → 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)