Fixed codestyle issues.

specreduce/tests/test_align_along_trace.py

@@ -18,19 +18,19 @@ def mk_test_image():
 def test_align_spectrum_along_trace_bad_input():
     image, trace = mk_test_image()
     with pytest.raises(ValueError, match='Unre'):
-        im = align_spectrum_along_trace(image, None)
+        im = align_spectrum_along_trace(image, None)   # noqa
 
     with pytest.raises(ValueError, match='method must be'):
-        im = align_spectrum_along_trace(image, trace, method='int')
+        im = align_spectrum_along_trace(image, trace, method='int')   # noqa
 
     with pytest.raises(ValueError, match='Spectral axis length'):
-        im = align_spectrum_along_trace(image.T, trace, method='interpolate', disp_axis=0)
+        im = align_spectrum_along_trace(image.T, trace, method='interpolate', disp_axis=0)   # noqa
 
     with pytest.raises(ValueError, match='Displacement axis must be'):
-        im = align_spectrum_along_trace(image, trace, disp_axis=2)
+        im = align_spectrum_along_trace(image, trace, disp_axis=2)  # noqa
 
     with pytest.raises(ValueError, match='The number of image dimensions must be'):
-        im = align_spectrum_along_trace(np.zeros((3,6,9)), trace)
+        im = align_spectrum_along_trace(np.zeros((3, 6, 9)), trace)  # noqa
 
 
 @pytest.mark.parametrize("method, truth_data, truth_mask, truth_ucty",
@@ -39,18 +39,18 @@ def test_align_spectrum_along_trace_bad_input():
                                      [0, 0, 0, 0.25, 0.5, 0.25, 0, 0, 0]]).T,
                            np.array([[0, 0, 0, 0, 0, 0, 0, 1, 1],
                                      [1, 1, 0, 0, 0, 0, 0, 0, 0]]).astype(bool).T,
-                           np.array([[1. , 1. , 1. , 1. , 1. , 1. , 1. , 1. , 1. ],
+                           np.array([[1., 1., 1., 1., 1., 1., 1., 1., 1.],
                                      [0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5]]).T),
                           ('shift',
-                           np.array([[0. , 0. , 0. , 0. , 1. , 0. , 0. , 0. , 0. ],
-                                     [0. , 0. , 0. , 0.5, 0.5, 0. , 0. , 0. , 0. ]]).T,
+                           np.array([[0., 0., 0., 0., 1., 0., 0., 0., 0.],
+                                     [0., 0., 0., 0.5, 0.5, 0., 0., 0., 0.]]).T,
                            np.array([[0, 0, 0, 0, 0, 0, 0, 0, 1],
                                      [1, 0, 0, 0, 0, 0, 0, 0, 0]]).astype(bool).T,
-                           np.ones((9,2)))],
+                           np.ones((9, 2)))],
                          ids=('method=interpolate', 'method=shift'))
 def test_align_spectrum_along_trace(method, truth_data, truth_mask, truth_ucty):
     image, trace = mk_test_image()
-    im =  align_spectrum_along_trace(image, trace, method=method)
+    im = align_spectrum_along_trace(image, trace, method=method)
     assert im.shape == image.shape
     assert im.unit == u.DN
     assert im.uncertainty.uncertainty_type == 'var'

specreduce/utils/utils.py

@@ -9,16 +9,17 @@
 from specreduce.tracing import Trace, FlatTrace
 from specreduce.extract import _ap_weight_image, _align_along_trace
 
-__all__ = ['measure_cross_dispersion_profile', '_align_along_trace', 'align_2d_spectrum_along_trace']
+__all__ = ['measure_cross_dispersion_profile', '_align_along_trace',
+           'align_2d_spectrum_along_trace']
 
 
 def _get_image_ndim(image):
     if isinstance(image, np.ndarray):
         return image.ndim
     elif isinstance(image, NDData):
         return image.data.ndim
     else:
         raise ValueError('Unrecognized image data format.')
 
 
 def align_2d_spectrum_along_trace(image: NDData | np.ndarray,
@@ -61,74 +62,74 @@
       require additional pre-processing (e.g., geometric rectification) before using
       this function.
     """
     if _get_image_ndim(image) > 2:
         raise ValueError('The number of image dimensions must be 2.')
     if not (0 <= disp_axis <= 1):
         raise ValueError('Displacement axis must be either 0 or 1.')
 
     if isinstance(trace, Trace):
         trace = trace.trace.data
     elif isinstance(trace, (np.ndarray, Number)):
         pass
     else:
         raise ValueError('Unrecognized trace format.')
 
     image = _ImageParser()._parse_image(image, disp_axis=disp_axis)
     data = image.data
     mask = image.mask | ~np.isfinite(data)
     ucty = image.uncertainty.represent_as(VarianceUncertainty).array
 
     if disp_axis == 0:
         data = data.T
         mask = mask.T
         ucty = ucty.T
 
     n_rows = data.shape[0]
     n_cols = data.shape[1]
 
     rows = np.broadcast_to(np.arange(n_rows)[:, None], data.shape)
     cols = np.broadcast_to(np.arange(n_cols), data.shape)
 
     if method == 'interpolate':
         # Calculate the real and integer shifts
         # and the interpolation weights.
         shifts = trace - n_rows / 2.0
         k = np.floor(shifts).astype(int)
         a = shifts - k
 
         # Calculate the shifted indices and mask the
         # edge pixels without information.
         ix1 = rows + k
         ix2 = ix1 + 1
         shift_mask = (ix1 < 0) | (ix1 > n_rows - 2)
         ix1 = np.clip(ix1, 0, n_rows - 1)
         ix2 = np.clip(ix2, 0, n_rows - 1)
 
         # Shift the data, uncertainties, and the mask using linear
         # interpolation.
         data_r = (1.0-a)*data[ix1, cols] + a*data[ix2, cols]
         ucty_r = (1.0-a)**2*ucty[ix1, cols] + a**2*ucty[ix2, cols]
         mask_r = mask[ix1, cols] | mask[ix2, cols] | shift_mask
 
     elif method == 'shift':
         shifts = trace.astype(int) - n_rows // 2
         ix = rows + shifts
         shift_mask = (ix < 0) | (ix > n_rows - 1)
         ix = np.clip(ix, 0, n_rows - 1)
 
         data_r = data[ix, cols]
         ucty_r = ucty[ix, cols]
         mask_r = mask[ix, cols] | shift_mask
 
     else:
         raise ValueError("method must be either 'interpolate' or 'shift'.")
 
     if disp_axis == 0:
         data_r = data_r.T
         mask_r = mask_r.T
         ucty_r = ucty_r.T
 
     return Spectrum1D(data_r * image.unit, mask=mask_r, meta=image.meta,
                       uncertainty=VarianceUncertainty(ucty_r).represent_as(image.uncertainty))
 
 
@@ -307,7 +308,7 @@
 
     # now that we have figured out the mask for the window in cross-disp. axis,
     # select only the pixel(s) we want to include in measuring the avg. profile
     pixel_mask = np.ones(image.shape)
     pixel_mask[:, pixels] = 0
 
     # combine these masks to isolate the rows and cols used to measure profile