Use jit if openmp not available, and print Windows diagnostics

Author: ar4

.github/workflows/build.yml

@@ -9,7 +9,8 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        os: [ubuntu-latest, windows-2022, macos-13, macos-14]
+        #os: [ubuntu-latest, windows-2022, macos-13, macos-14]
+        os: [windows-2022, macos-13]
 
     steps:
       - uses: actions/checkout@v5
@@ -50,15 +51,15 @@ jobs:
           name: cibw-wheels-${{ matrix.os }}-${{ strategy.job-index }}
           path: ./wheelhouse/*.whl
 
-  build_sdist:
-    name: Build source distribution
-    runs-on: ubuntu-latest
-    steps:
-      - uses: actions/checkout@v5
-
-      - name: Build sdist
-        run: pipx run build --sdist
-
-      - uses: actions/upload-artifact@v4
-        with:
-          path: dist/*.tar.gz
+            #  build_sdist:
+            #    name: Build source distribution
+            #    runs-on: ubuntu-latest
+            #    steps:
+            #      - uses: actions/checkout@v5
+            #
+            #      - name: Build sdist
+            #        run: pipx run build --sdist
+            #
+            #      - uses: actions/upload-artifact@v4
+            #        with:
+            #          path: dist/*.tar.gz

pyproject.toml

@@ -51,16 +51,33 @@ 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"'
+# Activate VS environment and run extensive diagnostics
+before-test = '''
+    echo "--- Diagnostics before vcvars64.bat ---"
+    echo "PATH is:"
+    echo %PATH%
+    echo "Attempting to locate cl.exe:"
+    where cl
+    echo "Installed packages:"
+    pip freeze
+    echo "-----------------------------------------"
+    call "C:\Program Files\Microsoft Visual Studio\2022\Enterprise\VC\Auxiliary\Build\vcvars64.bat"
+    echo "--- Diagnostics after vcvars64.bat ---"
+    echo "PATH is:"
+    echo %PATH%
+    echo "INCLUDE is:"
+    echo %INCLUDE%
+    echo "LIB is:"
+    echo %LIB%
+    echo "Attempting to locate cl.exe:"
+    where cl
+    echo "Attempting to run cl.exe:"
+    cl /help
+    echo "--------------------------------------"
+'''
 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"
-# 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.py

@@ -2524,7 +2524,7 @@ def elastic_func(
 
     if python_backend:
         if python_backend is True:
-            mode = "compile"
+            mode = "compile" if deepwave.backend_utils.USE_OPENMP else "jit"
         elif isinstance(python_backend, str):
             mode = python_backend.lower()
         else:

src/deepwave/scalar.py

@@ -2348,7 +2348,7 @@ def scalar_func(
 
     if python_backend:
         if python_backend is True:
-            mode = "compile"
+            mode = "compile" if deepwave.backend_utils.USE_OPENMP else "jit"
         elif isinstance(python_backend, str):
             mode = python_backend.lower()
         else:

src/deepwave/scalar_born.py

@@ -1979,7 +1979,7 @@ def scalar_born_func(
 
     if python_backend:
         if python_backend is True:
-            mode = "compile"
+            mode = "compile" if deepwave.backend_utils.USE_OPENMP else "jit"
         elif isinstance(python_backend, str):
             mode = python_backend.lower()
         else:

tests/test_callbacks.py

@@ -413,8 +413,9 @@ def do_nothing(state: deepwave.common.CallbackState) -> None:
 
 
 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.
+    """Check that the backward callback is called the correct number of times.
+
+    Check it when no forward callback is provided.
     """
     v = torch.ones(10, 10) * 1500
     v.requires_grad_()

tests/test_callbacks_born.py

@@ -309,11 +309,13 @@ 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)
 
 
 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.
+    """Check that the backward callback is called the correct number of times.
+
+    Check it when no forward callback is provided.
     """
     v = torch.ones(10, 10) * 1500
     scatter = torch.ones(10, 10)

tests/test_callbacks_elastic.py

@@ -373,8 +373,9 @@ def do_nothing(state: deepwave.common.CallbackState) -> None:
 
 
 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.
+    """Check that the backward callback is called the correct number of times.
+
+    Check it when no forward callback is provided.
     """
     lamb = torch.ones(10, 10) * 2200
     mu = torch.ones(10, 10) * 1000

tests/test_scalar_born.py

@@ -623,7 +623,7 @@ def test_forward_cpu_gpu_match() -> None:
                 prop_kwargs={"python_backend": python},
             )
             for cpui, gpui in zip(actual_cpu, actual_gpu):
-                assert torch.allclose(cpui, gpui.cpu(), atol=1e-5)
+                assert torch.allclose(cpui, gpui.cpu(), atol=5e-5)
 
 
 def test_unused_source_receiver(