fix: replace torch.tensor copy-constructs with as_tensor

tests/test_stokes.py

@@ -53,7 +53,7 @@ def test_stokes_recon(device):
                 s012 = stokes.stokes012_after_ar(*ar)
                 ar1 = stokes.estimate_ar_from_stokes012(*s012)
                 for i in range(3):
-                    tt.assert_close(torch.tensor(ar[i]), ar1[i])
+                    tt.assert_close(torch.as_tensor(ar[i]), ar1[i])
 
                 # Test attenuating depolarizing retarder (adr) functions
                 for depolarization in torch.arange(1e-3, 1, 0.1, device=device):
@@ -67,7 +67,7 @@ def test_stokes_recon(device):
                     adr1 = stokes.estimate_adr_from_stokes(*s0123)
 
                     for i in range(4):
-                        tt.assert_close(torch.tensor(adr[i]), adr1[i])
+                        tt.assert_close(torch.as_tensor(adr[i]), adr1[i])
 
 
 def test_stokes_after_adr_usage():
@@ -130,6 +130,21 @@ def test_copying(device):
     assert a[0] == 1
 
 
+def test_gradients_reach_copied_outputs():
+    """s0 and transmittance are copies of an input, so they must not detach."""
+    ones = torch.ones((2, 2))
+
+    transmittance = torch.ones((2, 2), requires_grad=True)
+    s0, _, _, _ = stokes.stokes_after_adr(ones, ones, transmittance, ones)
+    s0.sum().backward()
+    assert transmittance.grad is not None
+
+    s0_input = torch.ones((2, 2), requires_grad=True)
+    _, _, estimated, _ = stokes.estimate_adr_from_stokes(s0_input, ones, ones, ones)
+    estimated.sum().backward()
+    assert s0_input.grad is not None
+
+
 @pytest.mark.parametrize(*_DEVICE)
 def test_orientation_offset(device):
     ori = torch.tensor(

waveorder/optim/losses.py

@@ -243,16 +243,14 @@ def loss_fn(recon: Tensor) -> Tensor:
 
 
 def _make_spectral_flatness_loss(NA_det, wavelength, pixel_size, midband_fractions):
-    import numpy as np
-
     from waveorder import util
 
     def _flatness_2d(img: Tensor) -> Tensor:
         Y, X = img.shape
         device = img.device
 
         _, _, fxx, fyy = util.gen_coordinate((Y, X), pixel_size)
-        frr = torch.tensor(np.sqrt(fxx**2 + fyy**2), device=device)
+        frr = torch.sqrt(fxx**2 + fyy**2).to(device)
         cutoff = 2 * NA_det / wavelength
         mask = torch.logical_and(
             frr > cutoff * midband_fractions[0],

waveorder/stokes.py

@@ -170,7 +170,7 @@ def stokes_after_adr(retardance, orientation, transmittance, depolarization, inp
         raise NotImplementedError("input != cpl")
 
     # without copying transmittance, downstream changes to s0 will affect transmittance
-    s0 = torch.tensor(transmittance).clone()
+    s0 = torch.as_tensor(transmittance).clone()
     s1 = transmittance * depolarization * torch.sin(retardance) * torch.sin(2 * orientation)
     s2 = transmittance * depolarization * -torch.sin(retardance) * torch.cos(2 * orientation)
     s3 = transmittance * depolarization * torch.cos(retardance)
@@ -208,7 +208,7 @@ def stokes012_after_ar(retardance, orientation, transmittance, input="cpl"):
         raise NotImplementedError("input != cpl")
 
     # without copying transmittance, downstream changes to s0 will affect transmittance
-    s0 = torch.tensor(transmittance).clone()
+    s0 = torch.as_tensor(transmittance).clone()
     s1 = transmittance * torch.sin(retardance) * torch.sin(2 * orientation)
     s2 = transmittance * -torch.sin(retardance) * torch.cos(2 * orientation)
     return s0, s1, s2
@@ -273,7 +273,7 @@ def estimate_adr_from_stokes(s0, s1, s2, s3, input="cpl"):
     retardance = torch.arcsin(((s1**2 + s2**2) ** 0.5) / len_pol)
     orientation = _s12_to_orientation(s1, s2)
     # without copying s0, downstream changes to transmittance will affect s0
-    transmittance = torch.tensor(s0).clone()
+    transmittance = torch.as_tensor(s0).clone()
     depolarization = len_pol / s0
     return retardance, orientation, transmittance, depolarization
 
@@ -306,7 +306,7 @@ def estimate_ar_from_stokes012(s0, s1, s2, input="cpl"):
     retardance = torch.arcsin(((s1**2 + s2**2) ** 0.5) / s0)
     orientation = _s12_to_orientation(s1, s2)
     # without copying s0, downstream changes to transmittance will affect s0
-    transmittance = torch.tensor(s0).clone()
+    transmittance = torch.as_tensor(s0).clone()
     return retardance, orientation, transmittance
 
 
@@ -355,7 +355,7 @@ def mueller_from_stokes(
         raise NotImplementedError("direction must be `forward` or `inverse`")
 
     if direction == "forward":
-        M = torch.zeros((4, 4) + torch.tensor(s0).shape, device=s0.device)
+        M = torch.zeros((4, 4) + s0.shape, device=s0.device)
         denom = s1**2 + s2**2
         M[0, 0] = s0
         M[1, 1] = (s0 * s2**2 + s1**2 * s3) / denom

waveorder/util.py

@@ -114,7 +114,7 @@ def generate_star_target(yx_shape, blur_px=2, margin=60):
     x = np.arange(X) - X // 2
     y = np.arange(Y) - Y // 2
 
-    xx, yy = torch.tensor(np.meshgrid(x, y))
+    xx, yy = torch.as_tensor(np.stack(np.meshgrid(x, y)))
 
     rho = torch.sqrt(xx**2 + yy**2)
     theta = torch.arctan2(yy, xx)

waveorder/waveorder_reconstructor.py

@@ -540,7 +540,7 @@ def Hz_det_setup(self, phase_deconv, ph_deconv_layer, bire_in_plane_deconv, inc_
             # generate defocus kernel based on Pupil function and z_defocus
             self.Hz_det_2D = (
                 generate_propagation_kernel(
-                    torch.tensor(self.frr),
+                    torch.as_tensor(self.frr),
                     torch.tensor(self.Pupil_support),
                     self.lambda_illu,
                     torch.tensor(self.z_defocus),
@@ -580,7 +580,7 @@ def Hz_det_setup(self, phase_deconv, ph_deconv_layer, bire_in_plane_deconv, inc_
                 z = ifftshift((np.r_[0 : self.N_defocus_3D] - self.N_defocus_3D // 2) * self.psz)
             self.Hz_det_3D = (
                 generate_propagation_kernel(
-                    torch.tensor(self.frr),
+                    torch.as_tensor(self.frr),
                     torch.tensor(self.Pupil_support),
                     self.lambda_illu,
                     torch.tensor(z),
@@ -590,7 +590,7 @@ def Hz_det_setup(self, phase_deconv, ph_deconv_layer, bire_in_plane_deconv, inc_
             )
             self.G_fun_z_3D = (
                 generate_greens_function_z(
-                    torch.tensor(self.frr),
+                    torch.as_tensor(self.frr),
                     torch.tensor(self.Pupil_support),
                     self.lambda_illu,
                     torch.tensor(z),
@@ -772,7 +772,7 @@ def gen_WOTF(self):
         if self.N_pattern == 1:
             for i in range(self.N_defocus):
                 Hu_temp, Hp_temp = compute_weak_object_transfer_function_2d(
-                    torch.tensor(self.Source),
+                    torch.as_tensor(self.Source),
                     torch.tensor(self.Pupil_obj * self.Hz_det_2D[:, :, i]),
                 )
                 self.Hu[:, :, i] = Hu_temp.numpy()
@@ -781,7 +781,7 @@ def gen_WOTF(self):
             for i, j in itertools.product(range(self.N_defocus), range(self.N_pattern)):
                 idx = i * self.N_pattern + j
                 Hu_temp, Hp_temp = compute_weak_object_transfer_function_2d(
-                    torch.tensor(self.Source[j]),
+                    torch.as_tensor(self.Source[j]),
                     torch.tensor(self.Pupil_obj * self.Hz_det_2D[idx, :, :]),
                 )
                 self.Hu[:, :, idx] = Hu_temp.numpy()
@@ -903,7 +903,7 @@ def gen_2D_vec_WOTF(self, inc_option=False):
         # generate dyadic Green's tensor
         G_fun_z = (
             generate_greens_function_z(
-                torch.tensor(self.frr),
+                torch.as_tensor(self.frr),
                 torch.tensor(self.Pupil_support),
                 self.lambda_illu,
                 torch.tensor(self.z_defocus),
@@ -1241,7 +1241,7 @@ def gen_3D_vec_WOTF(self, inc_option):
             z = ifftshift((np.r_[0:N_defocus] - N_defocus // 2) * psz)
         G_fun_z = (
             generate_greens_function_z(
-                torch.tensor(self.frr),
+                torch.as_tensor(self.frr),
                 torch.tensor(self.Pupil_support),
                 self.lambda_illu,
                 torch.tensor(z),

waveorder/waveorder_simulator.py

@@ -109,7 +109,7 @@ def __init__(
 
         self.Hz_det = (
             generate_propagation_kernel(
-                torch.tensor(self.frr),
+                torch.as_tensor(self.frr),
                 torch.tensor(self.Pupil_support),
                 self.lambda_illu,
                 torch.tensor(self.z_defocus),