scalar_born batched tests passing

Author: ar4

tests/test_scalar_born.py

@@ -25,7 +25,7 @@ def test_born_scatter_v_batched_2d():
 
 def test_born_scatter_scatter_batched_2d():
     """Test Born propagation in a batched 2D scatter model."""
-    expected, actual = run_born_scatter_2d(dc=torch.tensor([[[50.0]], [[100.0]]]),
+    expected, actual = run_born_scatter_2d(dscatter=torch.tensor([[[50.0]], [[100.0]]]),
                                            propagator=scalarbornprop,
                                            dt=0.001,
                                            prop_kwargs={'pml_width': 30})
@@ -341,6 +341,7 @@ def scalarbornpropchained(model,
 
 def run_born_scatter(c,
                      dc,
+                     dscatter,
                      freq,
                      dx,
                      dt,
@@ -363,10 +364,10 @@ def run_born_scatter(c,
     else:
         min_c = c
         shot_c = [c] * num_shots
-    if isinstance(dc, torch.Tensor):
-        shot_dc = dc.flatten().tolist()
+    if isinstance(dscatter, torch.Tensor):
+        shot_dscatter = dscatter.flatten().tolist()
     else:
-        shot_dc = [dc] * num_shots
+        shot_dscatter = [dscatter] * num_shots
     if device is None:
         device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     model = torch.ones(*nx, device=device, dtype=dtype) * c
@@ -378,11 +379,11 @@ def run_born_scatter(c,
     x_s = _set_coords(num_shots, num_sources_per_shot, nx)
     x_r = _set_coords(num_shots, num_receivers_per_shot, nx)
     x_p = _set_coords(1, 1, nx, 'middle')[0, 0]
-    scatter = torch.zeros_like(model) * dc
-    if isinstance(dc, torch.Tensor):
-        scatter[..., x_p[0], x_p[1]] = dc.flatten()
+    scatter = torch.zeros_like(model) * dscatter
+    if isinstance(dscatter, torch.Tensor):
+        scatter[..., x_p[0], x_p[1]] = dscatter.flatten()
     else:
-        scatter[..., x_p[0], x_p[1]] = dc
+        scatter[..., x_p[0], x_p[1]] = dscatter
     #scatter[torch.split((x_p).long(), 1)] = dc
     sources = _set_sources(x_s, freq, dt, nt, dtype)
 
@@ -397,7 +398,7 @@ def run_born_scatter(c,
             for receiver in range(num_receivers_per_shot):
                 expected[shot, receiver, :] += \
                     scattered(x_r[shot, receiver], x_s[shot, source], x_p,
-                              dx, dt, shot_c[shot], shot_dc[shot],
+                              dx, dt, shot_c[shot], shot_dscatter[shot],
                               -sources['amplitude'][shot, source, :]).to(dtype)
 
     actual = propagator(model,
@@ -416,6 +417,7 @@ def run_born_scatter(c,
 
 def run_born_scatter_2d(c=1500,
                         dc=150,
+                        dscatter=150,
                         freq=25,
                         dx=(5, 5),
                         dt=0.0001,
@@ -430,7 +432,7 @@ def run_born_scatter_2d(c=1500,
                         **kwargs):
     """Runs run_born_scatter with default parameters for 2D."""
 
-    return run_born_scatter(c, dc, freq, dx, dt, nx, num_shots,
+    return run_born_scatter(c, dc, dscatter, freq, dx, dt, nx, num_shots,
                             num_sources_per_shot, num_receivers_per_shot,
                             propagator, prop_kwargs, device, dtype, **kwargs)
 
@@ -924,10 +926,11 @@ def test_born_gradcheck_v_batched():
 
 def test_born_gradcheck_scatter_batched():
     """Test gradcheck using a different scatter for each shot."""
-    run_born_gradcheck_2d(propagator=scalarbornprop, dc=torch.tensor([[[150.0]],[[1000.0]]]))
+    run_born_gradcheck_2d(propagator=scalarbornprop, dscatter=torch.tensor([[[150.0]],[[1000.0]]]))
 
 def run_born_gradcheck(c,
                        dc,
+                       dscatter,
                        freq,
                        dx,
                        dt,
@@ -965,15 +968,12 @@ def run_born_gradcheck(c,
                        nt_add=0):
     """Run PyTorch's gradcheck."""
     torch.manual_seed(1)
-    if isinstance(c, torch.Tensor):
-        min_c = c.abs().min().item()
-    else:
-        min_c = abs(c)
     if device is None:
         device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    model = (torch.ones(*nx, device=device, dtype=dtype) * c +
-             torch.rand(*nx, device=device, dtype=dtype) * dc)
-    scatter = torch.rand(*nx, device=device, dtype=dtype) * dc
+    model = torch.ones(*nx, device=device, dtype=dtype) * c
+    model += torch.rand(*model.shape, device=device, dtype=dtype) * dc
+    scatter = torch.rand(*nx, device=device, dtype=dtype) * dscatter
+    min_c = model.abs().min().item()
 
     nx = torch.Tensor(nx).long()
     dx = torch.Tensor(dx)
@@ -1065,6 +1065,7 @@ def run_born_gradcheck(c,
 
 def run_born_gradcheck_2d(c=1500,
                           dc=100,
+                          dscatter=100,
                           freq=25,
                           dx=(5, 5),
                           dt=0.001,
@@ -1084,6 +1085,7 @@ def run_born_gradcheck_2d(c=1500,
 
     return run_born_gradcheck(c,
                               dc,
+                              dscatter,
                               freq,
                               dx,
                               dt,