linting test_fancy_ptycho.py

tests/models/test_fancy_ptycho.py

@@ -1,13 +1,12 @@
 import pytest
-import cdtools
 import torch as t
 
 import cdtools
-from matplotlib import pyplot as plt
 
 # Force all reconstructions to use the same RNG seed
 t.manual_seed(0)
 
+
 def test_center_probe(lab_ptycho_cxi):
     dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(lab_ptycho_cxi)
     model = cdtools.models.FancyPtycho.from_dataset(
@@ -28,8 +27,8 @@ def test_center_probe(lab_ptycho_cxi):
     fourier_model.probe.data = cdtools.tools.propagators.far_field(
         base_probe
     )
-    
-    fourier_base_probe = fourier_model.probe.detach().clone()
+
+    fourier_model.probe.detach().clone()
     fourier_model.center_probes()
     fourier_centered_probe = fourier_model.probe.detach().clone()
     ifft_fourier_centered_probe = cdtools.tools.propagators.inverse_far_field(
@@ -45,26 +44,27 @@ def test_center_probe(lab_ptycho_cxi):
         atol=1e-4,
         rtol=1e-3
     )
-    
+
+
 @pytest.mark.slow
 def test_lab_ptycho(lab_ptycho_cxi, reconstruction_device, show_plot):
 
     print('\nTesting performance on the standard transmission ptycho dataset')
     dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(lab_ptycho_cxi)
-    
+
     model = cdtools.models.FancyPtycho.from_dataset(
         dataset,
-        n_modes=3, 
+        n_modes=3,
         oversampling=2,
         exponentiate_obj=True,
         dm_rank=2,
         probe_support_radius=120,
-        propagation_distance=5e-3, 
-        units='mm', 
+        propagation_distance=5e-3,
+        units='mm',
         obj_view_crop=-50,
-        use_qe_mask=True, # test this in the case where no qe mask is defined
+        use_qe_mask=True,  # test this in the case where no qe mask is defined
     )
-    
+
     print('Running reconstruction on provided reconstruction_device,',
           reconstruction_device)
     model.to(device=reconstruction_device)
@@ -75,11 +75,11 @@ def test_lab_ptycho(lab_ptycho_cxi, reconstruction_device, show_plot):
         if show_plot and model.epoch % 10 == 0:
             model.inspect(dataset)
 
-    for loss in model.Adam_optimize(50, dataset,  lr=0.005, batch_size=50):
+    for loss in model.Adam_optimize(50, dataset, lr=0.005, batch_size=50):
         print(model.report())
         if show_plot and model.epoch % 10 == 0:
             model.inspect(dataset)
-            
+
     model.tidy_probes()
 
     if show_plot:
@@ -94,7 +94,7 @@ def test_lab_ptycho(lab_ptycho_cxi, reconstruction_device, show_plot):
 def test_gold_balls(gold_ball_cxi, reconstruction_device, show_plot):
 
     print('\nTesting performance on the standard gold balls dataset')
-    
+
     dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(gold_ball_cxi)
 
     pad = 10
@@ -119,7 +119,7 @@ def test_gold_balls(gold_ball_cxi, reconstruction_device, show_plot):
           reconstruction_device)
     model.to(device=reconstruction_device)
     dataset.get_as(device=reconstruction_device)
-    
+
     for loss in model.Adam_optimize(20, dataset, lr=0.005, batch_size=50):
         print(model.report())
         if show_plot and model.epoch % 10 == 0:
@@ -135,7 +135,7 @@ def test_gold_balls(gold_ball_cxi, reconstruction_device, show_plot):
         print(model.report())
         if show_plot and model.epoch % 10 == 0:
             model.inspect(dataset)
-            
+
     model.tidy_probes()
 
     if show_plot:
@@ -146,5 +146,3 @@ def test_gold_balls(gold_ball_cxi, reconstruction_device, show_plot):
     # and choosing a rough value. If it triggers this assertion error,
     # something changed to make the final quality worse!
     assert model.loss_history[-1] < 0.0001
-
-