Fix bug in elastic.cu and continue effort to fix MacOS and Windows builds

Author: ar4

pyproject.toml

@@ -51,11 +51,16 @@ repair-wheel-command = "auditwheel repair -w {dest_dir} {wheel} --exclude libcud
 before-test = "pip install torch --index-url https://download.pytorch.org/whl/cpu"
 
 [tool.cibuildwheel.windows]
+# Activate VS environment so torch.compile can find the compiler.
+before-test = 'call "C:\Program Files\Microsoft Visual Studio\2022\Enterprise\VC\Auxiliary\Build\vcvars64.bat"'
 test-command = "pytest tests/test_scalar.py"
 
 [tool.cibuildwheel.macos]
+# Install libomp for linking during build and for runtime in the test env.
+before-build = "brew install libomp"
 before-test = "brew install libomp"
-environment = { OMP_PREFIX = "{brew_prefix}/opt/libomp" }
+# Set flags to find and enable OpenMP during the build.
+environment = { CFLAGS="-Xclang -fopenmp -I$(brew --prefix libomp)/include", LDFLAGS="-L$(brew --prefix libomp)/lib" }
 test-command = "pytest tests/test_scalar.py"
 
 [tool.scikit-build.sdist]

src/deepwave/elastic.cu

@@ -531,10 +531,10 @@ __global__ void backward_kernel_v(
     if (mu_requires_grad) {
       grad_mu_yx[i] += sigmaxy[i] * dvydxdvxdy_store[i] * (DW_DTYPE)step_ratio;
     }
-    if (pml_y == 0 || pml_y == 2) {
+    if (pml_y) {
       m_vxy[i] = mu_yx_shot[j] * dt * ayh[y] * sigmaxy[i] + ayh[y] * m_vxy[i];
     }
-    if (pml_x == 0 || pml_x == 2) {
+    if (pml_x) {
       m_vyx[i] = mu_yx_shot[j] * dt * axh[x] * sigmaxy[i] + axh[x] * m_vyx[i];
     }
     sigmaxy[i] += ((pml_y ? -DIFFYH1(SIGMAXY_Y_PML) : -DIFFYH1(SIGMAXY_Y)) +

tests/test_callbacks.py

@@ -410,3 +410,63 @@ def do_nothing(state: deepwave.common.CallbackState) -> None:
     )
     for i in range(len(out1)):
         assert torch.allclose(out1[i], out2[i])
+
+
+def test_scalar_backward_callback_only_call_count() -> None:
+    """Check that the backward callback is called the correct number of times
+    when no forward callback is provided.
+    """
+    v = torch.ones(10, 10) * 1500
+    v.requires_grad_()
+    dx = 5.0
+    dt = 0.004
+    nt = 20
+    source_amplitudes = torch.zeros(1, 1, nt)
+    source_amplitudes[0, 0, 5] = 1
+    source_locations = torch.zeros(1, 1, 2, dtype=torch.long)
+    source_locations[0, 0, 0] = 5
+    source_locations[0, 0, 1] = 5
+    receiver_locations = torch.zeros(1, 1, 2, dtype=torch.long)
+    receiver_locations[0, 0, 0] = 5
+    receiver_locations[0, 0, 1] = 5
+
+    class Counter:
+        """A simple counter class for callbacks."""
+
+        def __init__(self) -> None:
+            self.count = 0
+
+        def __call__(self, state: deepwave.common.CallbackState) -> None:
+            """Increments the counter."""
+            self.count += 1
+
+    # Test with a frequency that divides nt evenly
+    backward_counter = Counter()
+    out = deepwave.scalar(
+        v,
+        dx,
+        dt,
+        source_amplitudes=source_amplitudes,
+        source_locations=source_locations,
+        receiver_locations=receiver_locations,
+        backward_callback=backward_counter,
+        callback_frequency=2,
+    )
+    out[-1].sum().backward()
+    assert backward_counter.count == nt / 2
+
+    # Test with a frequency that does not divide nt evenly
+    v.grad.zero_()
+    backward_counter = Counter()
+    out = deepwave.scalar(
+        v,
+        dx,
+        dt,
+        source_amplitudes=source_amplitudes,
+        source_locations=source_locations,
+        receiver_locations=receiver_locations,
+        backward_callback=backward_counter,
+        callback_frequency=3,
+    )
+    out[-1].sum().backward()
+    assert backward_counter.count == (nt + 2) // 3

tests/test_callbacks_born.py

@@ -309,30 +309,74 @@ def do_nothing(state: deepwave.common.CallbackState) -> None:
     for i in range(len(out1)):
         assert torch.allclose(out1[i], out2[i])
     assert torch.allclose(grad1_v, grad2_v)
-    assert torch.allclose(grad1_scatter, grad2_scatter)
-    v.grad.zero_()
-    scatter.grad.zero_()
+
+
+def test_scalar_born_backward_callback_only_call_count() -> None:
+    """Check that the backward callback is called the correct number of times
+    when no forward callback is provided.
+    """
+    v = torch.ones(10, 10) * 1500
+    scatter = torch.ones(10, 10)
+    v.requires_grad_()
+    scatter.requires_grad_()
+    dx = 5.0
+    dt = 0.004
+    nt = 20
+    source_amplitudes = torch.zeros(1, 1, nt)
+    source_amplitudes[0, 0, 5] = 1
+    source_locations = torch.zeros(1, 1, 2, dtype=torch.long)
+    source_locations[0, 0, 0] = 5
+    source_locations[0, 0, 1] = 5
+    receiver_locations = torch.zeros(1, 1, 2, dtype=torch.long)
+    receiver_locations[0, 0, 0] = 5
+    receiver_locations[0, 0, 1] = 5
+
+    class Counter:
+        """A simple counter class for callbacks."""
+
+        def __init__(self) -> None:
+            self.count = 0
+
+        def __call__(self, state: deepwave.common.CallbackState) -> None:
+            """Increments the counter."""
+            self.count += 1
+
+    # Test with a frequency that divides nt evenly
+    backward_counter = Counter()
+    out = deepwave.scalar_born(
+        v,
+        scatter,
+        dx,
+        dt,
+        source_amplitudes=source_amplitudes,
+        source_locations=source_locations,
+        receiver_locations=receiver_locations,
+        backward_callback=backward_counter,
+        callback_frequency=2,
+    )
+    out[-1].sum().backward()
+    grad1_scatter = scatter.grad.detach()
+    assert backward_counter.count == nt / 2
 
     # Test with a frequency that does not divide nt evenly
-    out3 = deepwave.scalar_born(
+    v.grad.zero_()
+    scatter.grad.zero_()
+    backward_counter = Counter()
+    out = deepwave.scalar_born(
         v,
         scatter,
         dx,
         dt,
         source_amplitudes=source_amplitudes,
         source_locations=source_locations,
         receiver_locations=receiver_locations,
-        forward_callback=do_nothing,
-        backward_callback=do_nothing,
+        backward_callback=backward_counter,
         callback_frequency=3,
     )
-    out3[-1].sum().backward()
-    grad3_v = v.grad.clone()
-    grad3_scatter = scatter.grad.clone()
-    for i in range(len(out1)):
-        assert torch.allclose(out1[i], out3[i])
-    assert torch.allclose(grad1_v, grad3_v)
-    assert torch.allclose(grad1_scatter, grad3_scatter)
+    out[-1].sum().backward()
+    grad2_scatter = scatter.grad.detach()
+    assert backward_counter.count == (nt + 2) // 3
+    assert torch.allclose(grad1_scatter, grad2_scatter)
 
 
 def test_scalar_born_multishot_equivalence() -> None:

tests/test_callbacks_elastic.py

@@ -370,3 +370,73 @@ def do_nothing(state: deepwave.common.CallbackState) -> None:
     assert torch.allclose(grad1_lamb, grad3_lamb)
     assert torch.allclose(grad1_mu, grad3_mu)
     assert torch.allclose(grad1_buoyancy, grad3_buoyancy)
+
+
+def test_elastic_backward_callback_only_call_count() -> None:
+    """Check that the backward callback is called the correct number of times
+    when no forward callback is provided.
+    """
+    lamb = torch.ones(10, 10) * 2200
+    mu = torch.ones(10, 10) * 1000
+    buoyancy = torch.ones(10, 10) * 1 / 2200
+    lamb.requires_grad_()
+    mu.requires_grad_()
+    buoyancy.requires_grad_()
+    dx = 5.0
+    dt = 0.004
+    nt = 20
+    source_amplitudes_y = torch.zeros(1, 1, nt)
+    source_amplitudes_y[0, 0, 5] = 1
+    source_locations_y = torch.zeros(1, 1, 2, dtype=torch.long)
+    source_locations_y[0, 0, 0] = 5
+    source_locations_y[0, 0, 1] = 5
+    receiver_locations_y = torch.zeros(1, 1, 2, dtype=torch.long)
+    receiver_locations_y[0, 0, 0] = 5
+    receiver_locations_y[0, 0, 1] = 5
+
+    class Counter:
+        """A simple counter class for callbacks."""
+
+        def __init__(self) -> None:
+            self.count = 0
+
+        def __call__(self, state: deepwave.common.CallbackState) -> None:
+            """Increments the counter."""
+            self.count += 1
+
+    # Test with a frequency that divides nt evenly
+    backward_counter = Counter()
+    out = deepwave.elastic(
+        lamb,
+        mu,
+        buoyancy,
+        dx,
+        dt,
+        source_amplitudes_y=source_amplitudes_y,
+        source_locations_y=source_locations_y,
+        receiver_locations_y=receiver_locations_y,
+        backward_callback=backward_counter,
+        callback_frequency=2,
+    )
+    out[-1].sum().backward()
+    assert backward_counter.count == nt / 2
+
+    # Test with a frequency that does not divide nt evenly
+    lamb.grad.zero_()
+    mu.grad.zero_()
+    buoyancy.grad.zero_()
+    backward_counter = Counter()
+    out = deepwave.elastic(
+        lamb,
+        mu,
+        buoyancy,
+        dx,
+        dt,
+        source_amplitudes_y=source_amplitudes_y,
+        source_locations_y=source_locations_y,
+        receiver_locations_y=receiver_locations_y,
+        backward_callback=backward_counter,
+        callback_frequency=3,
+    )
+    out[-1].sum().backward()
+    assert backward_counter.count == (nt + 2) // 3