diff --git a/tests/test_stokes.py b/tests/test_stokes.py index 1792e5c2..7beeba98 100644 --- a/tests/test_stokes.py +++ b/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( diff --git a/waveorder/optim/losses.py b/waveorder/optim/losses.py index 57650996..14731cf0 100644 --- a/waveorder/optim/losses.py +++ b/waveorder/optim/losses.py @@ -243,8 +243,6 @@ 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: @@ -252,7 +250,7 @@ def _flatness_2d(img: Tensor) -> Tensor: 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], diff --git a/waveorder/stokes.py b/waveorder/stokes.py index bd276c18..ac57ff7f 100644 --- a/waveorder/stokes.py +++ b/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 diff --git a/waveorder/util.py b/waveorder/util.py index 42a22902..df97d7c4 100644 --- a/waveorder/util.py +++ b/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) diff --git a/waveorder/waveorder_reconstructor.py b/waveorder/waveorder_reconstructor.py index b5c89b33..3a18ea62 100644 --- a/waveorder/waveorder_reconstructor.py +++ b/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), diff --git a/waveorder/waveorder_simulator.py b/waveorder/waveorder_simulator.py index 5d99034c..712cedec 100644 --- a/waveorder/waveorder_simulator.py +++ b/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),