elastic batched tests pass

Author: ar4

tests/test_elastic.py

@@ -588,6 +588,20 @@ def test_gradcheck_only_mu_2d():
         source_requires_grad=False,
     )
 
+def test_gradcheck_batched_lamb_2d():
+    """Test gradcheck in a 2D model with batched lamb."""
+    run_gradcheck_2d(propagator=elasticprop,
+                     lamb=torch.tensor([[[DEFAULT_LAMB]],[[DEFAULT_LAMB*1.2]]]))
+
+def test_gradcheck_batched_mu_2d():
+    """Test gradcheck in a 2D model with batched mu."""
+    run_gradcheck_2d(propagator=elasticprop,
+                     mu=torch.tensor([[[DEFAULT_MU]],[[DEFAULT_MU*1.2]]]))
+
+def test_gradcheck_batched_buoyancy_2d():
+    """Test gradcheck in a 2D model with batched buoyancy."""
+    run_gradcheck_2d(propagator=elasticprop,
+                     buoyancy=torch.tensor([[[DEFAULT_BUOYANCY]],[[DEFAULT_BUOYANCY*1.2]]]))
 
 def _set_sources(x_s, freq, dt, nt, dtype=None, dpeak_time=0.3):
     """Create sources with amplitudes that have randomly shifted start times.
@@ -834,12 +848,12 @@ def run_gradcheck(mlamb,
     torch.manual_seed(1)
     if device is None:
         device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    lamb = (torch.ones(*nx, device=device, dtype=dtype) * mlamb +
-            torch.randn(*nx, dtype=dtype).to(device) * dlamb)
-    mu = (torch.ones(*nx, device=device, dtype=dtype) * mmu +
-          torch.randn(*nx, dtype=dtype).to(device) * dmu)
-    buoyancy = (torch.ones(*nx, device=device, dtype=dtype) * mbuoyancy +
-                torch.randn(*nx, dtype=dtype).to(device) * dbuoyancy)
+    lamb = torch.ones(*nx, device=device, dtype=dtype) * mlamb
+    lamb += torch.randn(*lamb.shape, dtype=dtype).to(device) * dlamb
+    mu = torch.ones(*nx, device=device, dtype=dtype) * mmu
+    mu += torch.randn(*mu.shape, dtype=dtype).to(device) * dmu
+    buoyancy = torch.ones(*nx, device=device, dtype=dtype) * mbuoyancy
+    buoyancy += torch.randn(*buoyancy.shape, dtype=dtype).to(device) * dbuoyancy
 
     nx = torch.Tensor(nx).long()
     dx = torch.Tensor(dx)
@@ -877,21 +891,34 @@ def run_gradcheck(mlamb,
     if isinstance(pml_width, int):
         pml_width = [pml_width for _ in range(4)]
 
-    if mlamb != 0:
+    if isinstance(mlamb, torch.Tensor):
+        min_mlamb = mlamb.abs().min()
+    else:
+        min_mlamb = mlamb
+    if isinstance(mmu, torch.Tensor):
+        min_mmu = mmu.abs().min()
+    else:
+        min_mmu = mmu
+    if isinstance(mbuoyancy, torch.Tensor):
+        min_mbuoyancy = mbuoyancy.abs().min()
+    else:
+        min_mbuoyancy = mbuoyancy
+
+    if min_mlamb != 0:
         lamb /= mlamb
-    if mmu != 0:
+    if min_mmu != 0:
         mu /= mmu
-    if mbuoyancy != 0:
+    if min_mbuoyancy != 0:
         buoyancy /= mbuoyancy
 
     lamb.requires_grad_(lamb_requires_grad)
     mu.requires_grad_(mu_requires_grad)
     buoyancy.requires_grad_(buoyancy_requires_grad)
 
     def wrap(lamb, mu, buoyancy, sources_y_amplitude, sources_x_amplitude):
-        if sources_y_amplitude is not None and mbuoyancy != 0:
+        if sources_y_amplitude is not None and min_mbuoyancy != 0:
             sources_y_amplitude = sources_y_amplitude / mbuoyancy / dt
-        if sources_x_amplitude is not None and mbuoyancy != 0:
+        if sources_x_amplitude is not None and min_mbuoyancy != 0:
             sources_x_amplitude = sources_x_amplitude / mbuoyancy / dt
         out = propagator(lamb * mlamb,
                          mu * mmu,