merge CZ fork to testing, prepare pipeline for clenup and fix

package/FOC_reduction.py

@@ -5,26 +5,19 @@ Main script where are progressively added the steps for the FOC pipeline reducti
 """
 
 # Project libraries
-
 from copy import deepcopy
-import os
 from os import system
 from os.path import exists as path_exists
 
-from matplotlib.colors import LogNorm
+import lib.fits as proj_fits  # Functions to handle fits files
+import lib.plots as proj_plots  # Functions for plotting data
+import lib.reduction as proj_red  # Functions used in reduction pipeline
 import numpy as np
-
+from lib.utils import princ_angle, sci_not
-from lib.background import subtract_bkg
+from matplotlib.colors import LogNorm
-import lib.fits as proj_fits        # Functions to handle fits files
-import lib.reduction as proj_red    # Functions used in reduction pipeline
-import lib.plots as proj_plots      # Functions for plotting data
-from lib.utils import sci_not, princ_angle
 
 
-
+def main(target=None, proposal_id=None, infiles=None, output_dir="./data", crop=False, interactive=False):
-
-
-def main(target=None, proposal_id=None, data_dir=None, infiles=None, output_dir="./data", crop=False, interactive=False):
     # Reduction parameters
     # Deconvolution
     deconvolve = False
@@ -42,10 +35,8 @@ def main(target=None, proposal_id=None, data_dir=None, infiles=None, output_dir=
     display_crop = False
 
     # Background estimation
-
     error_sub_type = "freedman-diaconis"  # sqrt, sturges, rice, scott, freedman-diaconis (default) or shape (example (51, 51))
     subtract_error = 1.0
-
     display_bkg = False
 
     # Data binning
@@ -55,7 +46,6 @@ def main(target=None, proposal_id=None, data_dir=None, infiles=None, output_dir=
 
     # Alignement
     align_center = "center"  # If None will not align the images
-
     display_align = False
     display_data = False
 
@@ -64,7 +54,7 @@ def main(target=None, proposal_id=None, data_dir=None, infiles=None, output_dir=
 
     # Smoothing
     smoothing_function = "combine"  # gaussian_after, weighted_gaussian_after, gaussian, weighted_gaussian or combine
-    smoothing_FWHM = 0.1    # If None, no smoothing is done
+    smoothing_FWHM = 0.10  # If None, no smoothing is done
     smoothing_scale = "arcsec"  # pixel or arcsec
 
     # Rotation
@@ -84,47 +74,37 @@ def main(target=None, proposal_id=None, data_dir=None, infiles=None, output_dir=
     # 3. Use the same alignment as the routine
     # 4. Skip the rebinning step
     # 5. Calulate the Stokes parameters without smoothing
-    optimal_binning = True
+    optimal_binning = False
     optimize = False
-    
+
     # Pipeline start
 
     # Step 1:
     #  Get data from fits files and translate to flux in erg/cm²/s/Angstrom.
     outfiles = []
-    if data_dir is None:
+    if infiles is not None:
-        if infiles is not None:
+        prod = np.array([["/".join(filepath.split("/")[:-1]), filepath.split("/")[-1]] for filepath in infiles], dtype=str)
-            prod = np.array([["/".join(filepath.split('/')[:-1]), filepath.split('/')[-1]] for filepath in infiles], dtype=str)
+        obs_dir = "/".join(infiles[0].split("/")[:-1])
-            obs_dir = "/".join(infiles[0].split("/")[:-1])
+        if not path_exists(obs_dir):
-            if not path_exists(obs_dir):
+            system("mkdir -p {0:s} {1:s}".format(obs_dir, obs_dir.replace("data", "plots")))
-                system("mkdir -p {0:s} {1:s}".format(obs_dir, obs_dir.replace("data", "plots")))
-            if target is None:
-                target = input("Target name:\n>")
-        else:
-            from lib.query import retrieve_products
-            target, products = retrieve_products(target, proposal_id, output_dir=output_dir)
-            prod = products.pop()
-            for prods in products:
-                outfiles.append(main(target=target, infiles=["/".join(pr) for pr in prods], output_dir=output_dir, crop=crop, interactive=interactive))
-        data_folder = prod[0][0]
-        
-        infiles = [p[1] for p in prod]
-        data_array, headers = proj_fits.get_obs_data(infiles, data_folder=data_folder, compute_flux=True)
-    
-    else:
-        infiles = [f for f in os.listdir(data_dir) if f.endswith('.fits') and f.startswith('x')]
-        data_folder = data_dir
         if target is None:
-                target = input("Target name:\n>")
+            target = input("Target name:\n>")
-        
+    else:
-        data_array, headers = proj_fits.get_obs_data(infiles, data_folder=data_folder, compute_flux=True)
+        from lib.query import retrieve_products
 
+        target, products = retrieve_products(target, proposal_id, output_dir=output_dir)
+        prod = products.pop()
+        for prods in products:
+            outfiles.append(main(target=target, infiles=["/".join(pr) for pr in prods], output_dir=output_dir, crop=crop, interactive=interactive)[0])
+    data_folder = prod[0][0]
     try:
         plots_folder = data_folder.replace("data", "plots")
     except ValueError:
         plots_folder = "."
     if not path_exists(plots_folder):
         system("mkdir -p {0:s} ".format(plots_folder))
+    infiles = [p[1] for p in prod]
+    data_array, headers = proj_fits.get_obs_data(infiles, data_folder=data_folder, compute_flux=True)
 
     figname = "_".join([target, "FOC"])
     figtype = ""
@@ -133,65 +113,129 @@ def main(target=None, proposal_id=None, data_dir=None, infiles=None, output_dir=
             figtype = "".join(["b", "{0:.2f}".format(pxsize), pxscale])  # additionnal informations
         else:
             figtype = "full"
-
     if smoothing_FWHM is not None and smoothing_scale is not None:
         smoothstr = "".join([*[s[0] for s in smoothing_function.split("_")], "{0:.2f}".format(smoothing_FWHM), smoothing_scale])
         figtype = "_".join([figtype, smoothstr] if figtype != "" else [smoothstr])
-
     if deconvolve:
         figtype = "_".join([figtype, "deconv"] if figtype != "" else ["deconv"])
-
     if align_center is None:
         figtype = "_".join([figtype, "not_aligned"] if figtype != "" else ["not_aligned"])
-    
+
     if optimal_binning:
-        options = {'optimize': optimize, 'optimal_binning': True}
+        from lib.background import subtract_bkg
-        
+
+        options = {"optimize": optimize, "optimal_binning": True}
+
         # Step 1: Load the data again and preserve the full images
-        _data_array, _headers = deepcopy(data_array), deepcopy(headers) # Preserve full images
+        _data_array, _headers = deepcopy(data_array), deepcopy(headers)  # Preserve full images
         _data_mask = np.ones(_data_array[0].shape, dtype=bool)
-        
+
         # Step 2: Skip the cropping step but use the same error and background estimation (I don't understand why this is wrong)
-        data_array, error_array, headers = proj_red.crop_array(data_array, headers, step=5, null_val=0., inside=True,
+        data_array, error_array, headers = proj_red.crop_array(
-                                                                display=display_crop, savename=figname, plots_folder=plots_folder)
+            data_array, headers, step=5, null_val=0.0, inside=True, display=display_crop, savename=figname, plots_folder=plots_folder
+        )
         data_mask = np.ones(data_array[0].shape, dtype=bool)
-        
+
         background = None
-        _, _, _, background, error_bkg = proj_red.get_error(data_array, headers, error_array, data_mask=data_mask, sub_type=error_sub_type, subtract_error=subtract_error, display=display_bkg, savename="_".join([figname, "errors"]), plots_folder=plots_folder, return_background=True)
+        _, _, _, background, error_bkg = proj_red.get_error(
-        
+            data_array,
+            headers,
+            error_array,
+            data_mask=data_mask,
+            sub_type=error_sub_type,
+            subtract_error=subtract_error,
+            display=display_bkg,
+            savename="_".join([figname, "errors"]),
+            plots_folder=plots_folder,
+            return_background=True,
+        )
+
         # _background is the same as background, but for the optimal binning
         _background = None
-        _data_array, _error_array, _, = proj_red.get_error(_data_array, _headers, error_array=None, data_mask=_data_mask, sub_type=error_sub_type, subtract_error=False, display=display_bkg, savename="_".join([figname, "errors"]), plots_folder=plots_folder, return_background=False)
+        _data_array, _error_array, _ = proj_red.get_error(
+            _data_array,
+            _headers,
+            error_array=None,
+            data_mask=_data_mask,
+            sub_type=error_sub_type,
+            subtract_error=False,
+            display=display_bkg,
+            savename="_".join([figname, "errors"]),
+            plots_folder=plots_folder,
+            return_background=False,
+        )
         _error_bkg = np.ones_like(_data_array) * error_bkg[:, 0, 0, np.newaxis, np.newaxis]
         _data_array, _error_array, _background, _ = subtract_bkg(_data_array, _error_array, _data_mask, background, _error_bkg)
 
         # Step 3: Align and rescale images with oversampling. (has to disable croping in align_data function)
-        _data_array, _error_array, _headers, _, shifts, error_shifts = proj_red.align_data(_data_array, _headers, error_array=_error_array, background=_background,
+        _data_array, _error_array, _headers, _, shifts, error_shifts = proj_red.align_data(
-                                                                                                upsample_factor=10, ref_center=align_center, return_shifts=True, optimal_binning=True)
+            _data_array,
+            _headers,
+            error_array=_error_array,
+            background=_background,
+            upsample_factor=10,
+            ref_center=align_center,
+            return_shifts=True,
+            optimal_binning=True,
+        )
         print("Image shifts: {} \nShifts uncertainty: {}".format(shifts, error_shifts))
         _data_mask = np.ones(_data_array[0].shape, dtype=bool)
-        
+
         # Step 4: Compute Stokes I, Q, U
         _background = np.array([np.array(bkg).reshape(1, 1) for bkg in _background])
-        _background_error = np.array([np.array(np.sqrt((bkg-_background[np.array([h['filtnam1'] == head['filtnam1'] for h in _headers], dtype=bool)].mean())
+        _background_error = np.array(
-                                    ** 2/np.sum([h['filtnam1'] == head['filtnam1'] for h in _headers]))).reshape(1, 1) for bkg, head in zip(_background, _headers)])
+            [
-        
+                np.array(
-        _I_stokes, _Q_stokes, _U_stokes, _Stokes_cov, _header_stokes = proj_red.compute_Stokes(_data_array, _error_array, _data_mask, _headers, 
+                    np.sqrt(
-                                                                                FWHM=None, scale=smoothing_scale, smoothing=smoothing_function, transmitcorr=transmitcorr)
+                        (bkg - _background[np.array([h["filtnam1"] == head["filtnam1"] for h in _headers], dtype=bool)].mean()) ** 2
-        _I_bkg, _Q_bkg, _U_bkg, _S_cov_bkg, _header_bkg = proj_red.compute_Stokes(_background, _background_error, np.array(True).reshape(1, 1), _headers, 
+                        / np.sum([h["filtnam1"] == head["filtnam1"] for h in _headers])
-                                                                    FWHM=None, scale=smoothing_scale, smoothing=smoothing_function, transmitcorr=False)
+                    )
-        
+                ).reshape(1, 1)
+                for bkg, head in zip(_background, _headers)
+            ]
+        )
+
+        _I_stokes, _Q_stokes, _U_stokes, _Stokes_cov, _header_stokes = proj_red.compute_Stokes(
+            _data_array, _error_array, _data_mask, _headers, FWHM=None, scale=smoothing_scale, smoothing=smoothing_function, transmitcorr=transmitcorr
+        )
+        _I_bkg, _Q_bkg, _U_bkg, _S_cov_bkg, _header_bkg = proj_red.compute_Stokes(
+            _background,
+            _background_error,
+            np.array(True).reshape(1, 1),
+            _headers,
+            FWHM=None,
+            scale=smoothing_scale,
+            smoothing=smoothing_function,
+            transmitcorr=False,
+        )
+
         # Step 5: Compute polarimetric parameters (polarization degree and angle).
         _P, _debiased_P, _s_P, _s_P_P, _PA, _s_PA, _s_PA_P = proj_red.compute_pol(_I_stokes, _Q_stokes, _U_stokes, _Stokes_cov, _header_stokes)
         _P_bkg, _debiased_P_bkg, _s_P_bkg, _s_P_P_bkg, _PA_bkg, _s_PA_bkg, _s_PA_P_bkg = proj_red.compute_pol(_I_bkg, _Q_bkg, _U_bkg, _S_cov_bkg, _header_bkg)
-        
+
         # Step 6: Save image to FITS.
         figname = "_".join([figname, figtype]) if figtype != "" else figname
-        _Stokes_hdul = proj_fits.save_Stokes(_I_stokes, _Q_stokes, _U_stokes, _Stokes_cov, _P, _debiased_P, _s_P, _s_P_P, _PA, _s_PA, _s_PA_P,
+        _Stokes_hdul = proj_fits.save_Stokes(
-                                    _header_stokes, _data_mask, figname, data_folder=data_folder, return_hdul=True)
+            _I_stokes,
-        
+            _Q_stokes,
+            _U_stokes,
+            _Stokes_cov,
+            _P,
+            _debiased_P,
+            _s_P,
+            _s_P_P,
+            _PA,
+            _s_PA,
+            _s_PA_P,
+            _header_stokes,
+            _data_mask,
+            figname,
+            data_folder=data_folder,
+            return_hdul=True,
+        )
+
         # Step 6:
-        _data_mask = _Stokes_hdul['data_mask'].data.astype(bool)
+        _data_mask = _Stokes_hdul["data_mask"].data.astype(bool)
         print(
             "F_int({0:.0f} Angs) = ({1} ± {2})e{3} ergs.cm^-2.s^-1.Angs^-1".format(
                 _header_stokes["PHOTPLAM"],
@@ -208,66 +252,196 @@ def main(target=None, proposal_id=None, data_dir=None, infiles=None, output_dir=
         #  Background values
         print(
             "F_bkg({0:.0f} Angs) = ({1} ± {2})e{3} ergs.cm^-2.s^-1.Angs^-1".format(
-                _header_stokes["PHOTFLAM"], *sci_not(_I_bkg[0, 0] * _header_stokes["PHOTFLAM"], np.sqrt(_S_cov_bkg[0, 0][0, 0]) * _header_stokes["PHOTFLAM"], 2, out=int)
+                _header_stokes["PHOTFLAM"],
+                *sci_not(_I_bkg[0, 0] * _header_stokes["PHOTFLAM"], np.sqrt(_S_cov_bkg[0, 0][0, 0]) * _header_stokes["PHOTFLAM"], 2, out=int),
             )
         )
         print("P_bkg = {0:.1f} ± {1:.1f} %".format(_debiased_P_bkg[0, 0] * 100.0, np.ceil(_s_P_bkg[0, 0] * 1000.0) / 10.0))
         print("PA_bkg = {0:.1f} ± {1:.1f} °".format(princ_angle(_PA_bkg[0, 0]), princ_angle(np.ceil(_s_PA_bkg[0, 0] * 10.0) / 10.0)))
-        
+
         #  Plot polarization map (Background is either total Flux, Polarization degree or Polarization degree error).
-        if pxscale.lower() not in ['full', 'integrate'] and not interactive:
+        if pxscale.lower() not in ["full", "integrate"] and not interactive:
-            proj_plots.polarization_map(deepcopy(_Stokes_hdul), _data_mask, SNRp_cut=SNRp_cut, SNRi_cut=SNRi_cut, flux_lim=flux_lim,
+            proj_plots.polarization_map(
-                                        step_vec=step_vec, vec_scale=scale_vec, savename="_".join([figname]), plots_folder=plots_folder, **options)
+                deepcopy(_Stokes_hdul),
-            proj_plots.polarization_map(deepcopy(_Stokes_hdul), _data_mask, SNRp_cut=SNRp_cut, SNRi_cut=SNRi_cut, flux_lim=flux_lim, step_vec=step_vec,
+                _data_mask,
-                                        vec_scale=scale_vec, savename="_".join([figname, "I"]), plots_folder=plots_folder, display='Intensity', **options)
+                SNRp_cut=SNRp_cut,
-            proj_plots.polarization_map(deepcopy(_Stokes_hdul), _data_mask, SNRp_cut=SNRp_cut, SNRi_cut=SNRi_cut, flux_lim=flux_lim, step_vec=step_vec,
+                SNRi_cut=SNRi_cut,
-                                        vec_scale=scale_vec, savename="_".join([figname, "P_flux"]), plots_folder=plots_folder, display='Pol_Flux', **options)
+                flux_lim=flux_lim,
-            proj_plots.polarization_map(deepcopy(_Stokes_hdul), _data_mask, SNRp_cut=SNRp_cut, SNRi_cut=SNRi_cut, flux_lim=flux_lim, step_vec=step_vec,
+                step_vec=step_vec,
-                                        vec_scale=scale_vec, savename="_".join([figname, "P"]), plots_folder=plots_folder, display='Pol_deg', **options)
+                vec_scale=scale_vec,
-            proj_plots.polarization_map(deepcopy(_Stokes_hdul), _data_mask, SNRp_cut=SNRp_cut, SNRi_cut=SNRi_cut, flux_lim=flux_lim, step_vec=step_vec,
+                savename="_".join([figname]),
-                                        vec_scale=scale_vec, savename="_".join([figname, "PA"]), plots_folder=plots_folder, display='Pol_ang', **options)
+                plots_folder=plots_folder,
-            proj_plots.polarization_map(deepcopy(_Stokes_hdul), _data_mask, SNRp_cut=SNRp_cut, SNRi_cut=SNRi_cut, flux_lim=flux_lim, step_vec=step_vec,
+                **options,
-                                        vec_scale=scale_vec, savename="_".join([figname, "I_err"]), plots_folder=plots_folder, display='I_err', **options)
+            )
-            proj_plots.polarization_map(deepcopy(_Stokes_hdul), _data_mask, SNRp_cut=SNRp_cut, SNRi_cut=SNRi_cut, flux_lim=flux_lim, step_vec=step_vec,
+            proj_plots.polarization_map(
-                                        vec_scale=scale_vec, savename="_".join([figname, "P_err"]), plots_folder=plots_folder, display='Pol_deg_err', **options)
+                deepcopy(_Stokes_hdul),
-            proj_plots.polarization_map(deepcopy(_Stokes_hdul), _data_mask, SNRp_cut=SNRp_cut, SNRi_cut=SNRi_cut, flux_lim=flux_lim, step_vec=step_vec,
+                _data_mask,
-                                        vec_scale=scale_vec, savename="_".join([figname, "SNRi"]), plots_folder=plots_folder, display='SNRi', **options)
+                SNRp_cut=SNRp_cut,
-            proj_plots.polarization_map(deepcopy(_Stokes_hdul), _data_mask, SNRp_cut=SNRp_cut, SNRi_cut=SNRi_cut, flux_lim=flux_lim, step_vec=step_vec,
+                SNRi_cut=SNRi_cut,
-                                        vec_scale=scale_vec, savename="_".join([figname, "SNRp"]), plots_folder=plots_folder, display='SNRp', **options)
+                flux_lim=flux_lim,
+                step_vec=step_vec,
+                vec_scale=scale_vec,
+                savename="_".join([figname, "I"]),
+                plots_folder=plots_folder,
+                display="Intensity",
+                **options,
+            )
+            proj_plots.polarization_map(
+                deepcopy(_Stokes_hdul),
+                _data_mask,
+                SNRp_cut=SNRp_cut,
+                SNRi_cut=SNRi_cut,
+                flux_lim=flux_lim,
+                step_vec=step_vec,
+                vec_scale=scale_vec,
+                savename="_".join([figname, "P_flux"]),
+                plots_folder=plots_folder,
+                display="Pol_Flux",
+                **options,
+            )
+            proj_plots.polarization_map(
+                deepcopy(_Stokes_hdul),
+                _data_mask,
+                SNRp_cut=SNRp_cut,
+                SNRi_cut=SNRi_cut,
+                flux_lim=flux_lim,
+                step_vec=step_vec,
+                vec_scale=scale_vec,
+                savename="_".join([figname, "P"]),
+                plots_folder=plots_folder,
+                display="Pol_deg",
+                **options,
+            )
+            proj_plots.polarization_map(
+                deepcopy(_Stokes_hdul),
+                _data_mask,
+                SNRp_cut=SNRp_cut,
+                SNRi_cut=SNRi_cut,
+                flux_lim=flux_lim,
+                step_vec=step_vec,
+                vec_scale=scale_vec,
+                savename="_".join([figname, "PA"]),
+                plots_folder=plots_folder,
+                display="Pol_ang",
+                **options,
+            )
+            proj_plots.polarization_map(
+                deepcopy(_Stokes_hdul),
+                _data_mask,
+                SNRp_cut=SNRp_cut,
+                SNRi_cut=SNRi_cut,
+                flux_lim=flux_lim,
+                step_vec=step_vec,
+                vec_scale=scale_vec,
+                savename="_".join([figname, "I_err"]),
+                plots_folder=plots_folder,
+                display="I_err",
+                **options,
+            )
+            proj_plots.polarization_map(
+                deepcopy(_Stokes_hdul),
+                _data_mask,
+                SNRp_cut=SNRp_cut,
+                SNRi_cut=SNRi_cut,
+                flux_lim=flux_lim,
+                step_vec=step_vec,
+                vec_scale=scale_vec,
+                savename="_".join([figname, "P_err"]),
+                plots_folder=plots_folder,
+                display="Pol_deg_err",
+                **options,
+            )
+            proj_plots.polarization_map(
+                deepcopy(_Stokes_hdul),
+                _data_mask,
+                SNRp_cut=SNRp_cut,
+                SNRi_cut=SNRi_cut,
+                flux_lim=flux_lim,
+                step_vec=step_vec,
+                vec_scale=scale_vec,
+                savename="_".join([figname, "SNRi"]),
+                plots_folder=plots_folder,
+                display="SNRi",
+                **options,
+            )
+            proj_plots.polarization_map(
+                deepcopy(_Stokes_hdul),
+                _data_mask,
+                SNRp_cut=SNRp_cut,
+                SNRi_cut=SNRi_cut,
+                flux_lim=flux_lim,
+                step_vec=step_vec,
+                vec_scale=scale_vec,
+                savename="_".join([figname, "SNRp"]),
+                plots_folder=plots_folder,
+                display="SNRp",
+                **options,
+            )
         elif not interactive:
-            proj_plots.polarization_map(deepcopy(_Stokes_hdul), _data_mask, SNRp_cut=SNRp_cut, SNRi_cut=SNRi_cut,
+            proj_plots.polarization_map(
-                                        savename=figname, plots_folder=plots_folder, display='integrate', **options)
+                deepcopy(_Stokes_hdul),
-        elif pxscale.lower() not in ['full', 'integrate']:
+                _data_mask,
+                SNRp_cut=SNRp_cut,
+                SNRi_cut=SNRi_cut,
+                savename=figname,
+                plots_folder=plots_folder,
+                display="integrate",
+                **options,
+            )
+        elif pxscale.lower() not in ["full", "integrate"]:
             proj_plots.pol_map(_Stokes_hdul, SNRp_cut=SNRp_cut, SNRi_cut=SNRi_cut, flux_lim=flux_lim)
-    
+
     else:
-        options = {'optimize': optimize, 'optimal_binning': False}
+        options = {"optimize": optimize, "optimal_binning": False}
         #  Crop data to remove outside blank margins.
-        data_array, error_array, headers = proj_red.crop_array(data_array, headers, step=5, null_val=0., inside=True,
+        data_array, error_array, headers = proj_red.crop_array(
-                                                                display=display_crop, savename=figname, plots_folder=plots_folder)
+            data_array, headers, step=5, null_val=0.0, inside=True, display=display_crop, savename=figname, plots_folder=plots_folder
+        )
         data_mask = np.ones(data_array[0].shape, dtype=bool)
 
         #  Deconvolve data using Richardson-Lucy iterative algorithm with a gaussian PSF of given FWHM.
         if deconvolve:
-            data_array = proj_red.deconvolve_array(data_array, headers, psf=psf, FWHM=psf_FWHM, scale=psf_scale, shape=psf_shape, iterations=iterations, algo=algo)
+            data_array = proj_red.deconvolve_array(
+                data_array, headers, psf=psf, FWHM=psf_FWHM, scale=psf_scale, shape=psf_shape, iterations=iterations, algo=algo
+            )
 
         #  Estimate error from data background, estimated from sub-image of desired sub_shape.
         background = None
-        data_array, error_array, headers, background, error_bkg = proj_red.get_error(data_array, headers, error_array, data_mask=data_mask, sub_type=error_sub_type, subtract_error=subtract_error, display=display_bkg, savename="_".join([figname, "errors"]), plots_folder=plots_folder, return_background=True)
+        data_array, error_array, headers, background, error_bkg = proj_red.get_error(
+            data_array,
+            headers,
+            error_array,
+            data_mask=data_mask,
+            sub_type=error_sub_type,
+            subtract_error=subtract_error,
+            display=display_bkg,
+            savename="_".join([figname, "errors"]),
+            plots_folder=plots_folder,
+            return_background=True,
+        )
 
         #  Align and rescale images with oversampling.
         data_array, error_array, headers, data_mask, shifts, error_shifts = proj_red.align_data(
-            data_array, headers, error_array=error_array, background=background, upsample_factor=10, ref_center=align_center, return_shifts=True)
+            data_array, headers, error_array=error_array, background=background, upsample_factor=10, ref_center=align_center, return_shifts=True
-            
+        )
+
         if display_align:
             print("Image shifts: {} \nShifts uncertainty: {}".format(shifts, error_shifts))
-            proj_plots.plot_obs(data_array, headers, savename="_".join([figname, str(align_center)]), plots_folder=plots_folder, norm=LogNorm(
+            proj_plots.plot_obs(
-                vmin=data_array[data_array > 0.].min()*headers[0]['photflam'], vmax=data_array[data_array > 0.].max()*headers[0]['photflam']))
+                data_array,
+                headers,
+                savename="_".join([figname, str(align_center)]),
+                plots_folder=plots_folder,
+                norm=LogNorm(
+                    vmin=data_array[data_array > 0.0].min() * headers[0]["photflam"], vmax=data_array[data_array > 0.0].max() * headers[0]["photflam"]
+                ),
+            )
 
         #  Rebin data to desired pixel size.
         if (pxsize is not None) and not (pxsize == 1 and pxscale.lower() in ["px", "pixel", "pixels"]):
             data_array, error_array, headers, Dxy, data_mask = proj_red.rebin_array(
-                data_array, error_array, headers, pxsize=pxsize, scale=pxscale, operation=rebin_operation, data_mask=data_mask)
+                data_array, error_array, headers, pxsize=pxsize, scale=pxscale, operation=rebin_operation, data_mask=data_mask
+            )
 
         # Rotate data to have same orientation
         rotate_data = np.unique([np.round(float(head["ORIENTAT"]), 3) for head in headers]).size != 1
@@ -288,13 +462,29 @@ def main(target=None, proposal_id=None, data_dir=None, infiles=None, output_dir=
                 )
 
         # Plot array for checking output
-        if display_data and pxscale.lower() not in ['full', 'integrate']:
+        if display_data and pxscale.lower() not in ["full", "integrate"]:
-            proj_plots.plot_obs(data_array, headers, savename="_".join([figname, "rebin"]), plots_folder=plots_folder, norm=LogNorm(
+            proj_plots.plot_obs(
-                vmin=data_array[data_array > 0.].min()*headers[0]['photflam'], vmax=data_array[data_array > 0.].max()*headers[0]['photflam']))
+                data_array,
+                headers,
+                savename="_".join([figname, "rebin"]),
+                plots_folder=plots_folder,
+                norm=LogNorm(
+                    vmin=data_array[data_array > 0.0].min() * headers[0]["photflam"], vmax=data_array[data_array > 0.0].max() * headers[0]["photflam"]
+                ),
+            )
 
         background = np.array([np.array(bkg).reshape(1, 1) for bkg in background])
-        background_error = np.array([np.array(np.sqrt((bkg-background[np.array([h['filtnam1'] == head['filtnam1'] for h in headers], dtype=bool)].mean())
+        background_error = np.array(
-                                    ** 2/np.sum([h['filtnam1'] == head['filtnam1'] for h in headers]))).reshape(1, 1) for bkg, head in zip(background, headers)])
+            [
+                np.array(
+                    np.sqrt(
+                        (bkg - background[np.array([h["filtnam1"] == head["filtnam1"] for h in headers], dtype=bool)].mean()) ** 2
+                        / np.sum([h["filtnam1"] == head["filtnam1"] for h in headers])
+                    )
+                ).reshape(1, 1)
+                for bkg, head in zip(background, headers)
+            ]
+        )
 
         # Step 2:
         # Compute Stokes I, Q, U with smoothed polarized images
@@ -303,16 +493,28 @@ def main(target=None, proposal_id=None, data_dir=None, infiles=None, output_dir=
         # see Jedrzejewski, R.; Nota, A.; Hack, W. J., A Comparison Between FOC and WFPC2
         # Bibcode : 1995chst.conf...10J
         I_stokes, Q_stokes, U_stokes, Stokes_cov, header_stokes = proj_red.compute_Stokes(
-            data_array, error_array, data_mask, headers, FWHM=smoothing_FWHM, scale=smoothing_scale, smoothing=smoothing_function, transmitcorr=transmitcorr)
+            data_array, error_array, data_mask, headers, FWHM=smoothing_FWHM, scale=smoothing_scale, smoothing=smoothing_function, transmitcorr=transmitcorr
-        I_bkg, Q_bkg, U_bkg, S_cov_bkg, header_bkg = proj_red.compute_Stokes(background, background_error, np.array(True).reshape(
+        )
-            1, 1), headers, FWHM=None, scale=smoothing_scale, smoothing=smoothing_function, transmitcorr=False)
+        I_bkg, Q_bkg, U_bkg, S_cov_bkg, header_bkg = proj_red.compute_Stokes(
+            background,
+            background_error,
+            np.array(True).reshape(1, 1),
+            headers,
+            FWHM=None,
+            scale=smoothing_scale,
+            smoothing=smoothing_function,
+            transmitcorr=False,
+        )
 
         # Step 3:
         # Rotate images to have North up
         if rotate_North:
             I_stokes, Q_stokes, U_stokes, Stokes_cov, data_mask, header_stokes = proj_red.rotate_Stokes(
-                I_stokes, Q_stokes, U_stokes, Stokes_cov, data_mask, header_stokes, SNRi_cut=None)
+                I_stokes, Q_stokes, U_stokes, Stokes_cov, data_mask, header_stokes, SNRi_cut=None
-            I_bkg, Q_bkg, U_bkg, S_cov_bkg, data_mask_bkg, header_bkg = proj_red.rotate_Stokes(I_bkg, Q_bkg, U_bkg, S_cov_bkg, np.array(True).reshape(1, 1), header_bkg, SNRi_cut=None)
+            )
+            I_bkg, Q_bkg, U_bkg, S_cov_bkg, data_mask_bkg, header_bkg = proj_red.rotate_Stokes(
+                I_bkg, Q_bkg, U_bkg, S_cov_bkg, np.array(True).reshape(1, 1), header_bkg, SNRi_cut=None
+            )
 
         # Compute polarimetric parameters (polarization degree and angle).
         P, debiased_P, s_P, s_P_P, PA, s_PA, s_PA_P = proj_red.compute_pol(I_stokes, Q_stokes, U_stokes, Stokes_cov, header_stokes)
@@ -321,8 +523,24 @@ def main(target=None, proposal_id=None, data_dir=None, infiles=None, output_dir=
         # Step 4:
         # Save image to FITS.
         figname = "_".join([figname, figtype]) if figtype != "" else figname
-        Stokes_hdul = proj_fits.save_Stokes(I_stokes, Q_stokes, U_stokes, Stokes_cov, P, debiased_P, s_P, s_P_P, PA, s_PA, s_PA_P,
+        Stokes_hdul = proj_fits.save_Stokes(
-                                            header_stokes, data_mask, figname, data_folder=data_folder, return_hdul=True)
+            I_stokes,
+            Q_stokes,
+            U_stokes,
+            Stokes_cov,
+            P,
+            debiased_P,
+            s_P,
+            s_P_P,
+            PA,
+            s_PA,
+            s_PA_P,
+            header_stokes,
+            data_mask,
+            figname,
+            data_folder=data_folder,
+            return_hdul=True,
+        )
         outfiles.append("/".join([data_folder, Stokes_hdul[0].header["FILENAME"] + ".fits"]))
 
         # Step 5:
@@ -331,11 +549,11 @@ def main(target=None, proposal_id=None, data_dir=None, infiles=None, output_dir=
             figname += "_crop"
             stokescrop = proj_plots.crop_Stokes(deepcopy(Stokes_hdul), norm=LogNorm())
             stokescrop.crop()
-            stokescrop.write_to("/".join([data_folder, figname+".fits"]))
+            stokescrop.write_to("/".join([data_folder, figname + ".fits"]))
             Stokes_hdul, header_stokes = stokescrop.hdul_crop, [dataset.header for dataset in stokescrop.hdul_crop]
             outfiles.append("/".join([data_folder, Stokes_hdul[0].header["FILENAME"] + ".fits"]))
 
-        data_mask = Stokes_hdul['data_mask'].data.astype(bool)
+        data_mask = Stokes_hdul["data_mask"].data.astype(bool)
         print(
             "F_int({0:.0f} Angs) = ({1} ± {2})e{3} ergs.cm^-2.s^-1.Angs^-1".format(
                 header_stokes["PHOTPLAM"],
@@ -352,55 +570,161 @@ def main(target=None, proposal_id=None, data_dir=None, infiles=None, output_dir=
         #  Background values
         print(
             "F_bkg({0:.0f} Angs) = ({1} ± {2})e{3} ergs.cm^-2.s^-1.Angs^-1".format(
-                header_stokes["PHOTPLAM"], *sci_not(I_bkg[0, 0] * header_stokes["PHOTPLAM"], np.sqrt(S_cov_bkg[0, 0][0, 0]) * header_stokes["PHOTPLAM"], 2, out=int)
+                header_stokes["PHOTPLAM"],
+                *sci_not(I_bkg[0, 0] * header_stokes["PHOTPLAM"], np.sqrt(S_cov_bkg[0, 0][0, 0]) * header_stokes["PHOTPLAM"], 2, out=int),
             )
         )
         print("P_bkg = {0:.1f} ± {1:.1f} %".format(debiased_P_bkg[0, 0] * 100.0, np.ceil(s_P_bkg[0, 0] * 1000.0) / 10.0))
         print("PA_bkg = {0:.1f} ± {1:.1f} °".format(princ_angle(PA_bkg[0, 0]), princ_angle(np.ceil(s_PA_bkg[0, 0] * 10.0) / 10.0)))
         #  Plot polarization map (Background is either total Flux, Polarization degree or Polarization degree error).
-        if pxscale.lower() not in ['full', 'integrate'] and not interactive:
+        if pxscale.lower() not in ["full", "integrate"] and not interactive:
-            proj_plots.polarization_map(deepcopy(Stokes_hdul), data_mask, SNRp_cut=SNRp_cut, SNRi_cut=SNRi_cut, flux_lim=flux_lim,
+            proj_plots.polarization_map(
-                                        step_vec=step_vec, scale_vec=scale_vec, savename="_".join([figname]), plots_folder=plots_folder, **options)
+                deepcopy(Stokes_hdul),
-            proj_plots.polarization_map(deepcopy(Stokes_hdul), data_mask, SNRp_cut=SNRp_cut, SNRi_cut=SNRi_cut, flux_lim=flux_lim, step_vec=step_vec,
+                data_mask,
-                                        scale_vec=scale_vec, savename="_".join([figname, "I"]), plots_folder=plots_folder, display='Intensity', **options)
+                SNRp_cut=SNRp_cut,
-            proj_plots.polarization_map(deepcopy(Stokes_hdul), data_mask, SNRp_cut=SNRp_cut, SNRi_cut=SNRi_cut, flux_lim=flux_lim, step_vec=step_vec,
+                SNRi_cut=SNRi_cut,
-                                        scale_vece=scale_vec, savename="_".join([figname, "P_flux"]), plots_folder=plots_folder, display='Pol_Flux', **options)
+                flux_lim=flux_lim,
-            proj_plots.polarization_map(deepcopy(Stokes_hdul), data_mask, SNRp_cut=SNRp_cut, SNRi_cut=SNRi_cut, flux_lim=flux_lim, step_vec=step_vec,
+                step_vec=step_vec,
-                                        scale_vec=scale_vec, savename="_".join([figname, "P"]), plots_folder=plots_folder, display='Pol_deg', **options)
+                scale_vec=scale_vec,
-            proj_plots.polarization_map(deepcopy(Stokes_hdul), data_mask, SNRp_cut=SNRp_cut, SNRi_cut=SNRi_cut, flux_lim=flux_lim, step_vec=step_vec,
+                savename="_".join([figname]),
-                                        scale_vec=scale_vec, savename="_".join([figname, "PA"]), plots_folder=plots_folder, display='Pol_ang', **options)
+                plots_folder=plots_folder,
-            proj_plots.polarization_map(deepcopy(Stokes_hdul), data_mask, SNRp_cut=SNRp_cut, SNRi_cut=SNRi_cut, flux_lim=flux_lim, step_vec=step_vec,
+                **options,
-                                        scale_vec=scale_vec, savename="_".join([figname, "I_err"]), plots_folder=plots_folder, display='I_err', **options)
+            )
-            proj_plots.polarization_map(deepcopy(Stokes_hdul), data_mask, SNRp_cut=SNRp_cut, SNRi_cut=SNRi_cut, flux_lim=flux_lim, step_vec=step_vec,
+            proj_plots.polarization_map(
-                                        scale_vec=scale_vec, savename="_".join([figname, "P_err"]), plots_folder=plots_folder, display='Pol_deg_err', **options)
+                deepcopy(Stokes_hdul),
-            proj_plots.polarization_map(deepcopy(Stokes_hdul), data_mask, SNRp_cut=SNRp_cut, SNRi_cut=SNRi_cut, flux_lim=flux_lim, step_vec=step_vec,
+                data_mask,
-                                        scale_vec=scale_vec, savename="_".join([figname, "SNRi"]), plots_folder=plots_folder, display='SNRi', **options)
+                SNRp_cut=SNRp_cut,
-            proj_plots.polarization_map(deepcopy(Stokes_hdul), data_mask, SNRp_cut=SNRp_cut, SNRi_cut=SNRi_cut, flux_lim=flux_lim, step_vec=step_vec,
+                SNRi_cut=SNRi_cut,
-                                        scale_vec=scale_vec, savename="_".join([figname, "SNRp"]), plots_folder=plots_folder, display='SNRp', **options)
+                flux_lim=flux_lim,
+                step_vec=step_vec,
+                scale_vec=scale_vec,
+                savename="_".join([figname, "I"]),
+                plots_folder=plots_folder,
+                display="Intensity",
+                **options,
+            )
+            proj_plots.polarization_map(
+                deepcopy(Stokes_hdul),
+                data_mask,
+                SNRp_cut=SNRp_cut,
+                SNRi_cut=SNRi_cut,
+                flux_lim=flux_lim,
+                step_vec=step_vec,
+                scale_vece=scale_vec,
+                savename="_".join([figname, "P_flux"]),
+                plots_folder=plots_folder,
+                display="Pol_Flux",
+                **options,
+            )
+            proj_plots.polarization_map(
+                deepcopy(Stokes_hdul),
+                data_mask,
+                SNRp_cut=SNRp_cut,
+                SNRi_cut=SNRi_cut,
+                flux_lim=flux_lim,
+                step_vec=step_vec,
+                scale_vec=scale_vec,
+                savename="_".join([figname, "P"]),
+                plots_folder=plots_folder,
+                display="Pol_deg",
+                **options,
+            )
+            proj_plots.polarization_map(
+                deepcopy(Stokes_hdul),
+                data_mask,
+                SNRp_cut=SNRp_cut,
+                SNRi_cut=SNRi_cut,
+                flux_lim=flux_lim,
+                step_vec=step_vec,
+                scale_vec=scale_vec,
+                savename="_".join([figname, "PA"]),
+                plots_folder=plots_folder,
+                display="Pol_ang",
+                **options,
+            )
+            proj_plots.polarization_map(
+                deepcopy(Stokes_hdul),
+                data_mask,
+                SNRp_cut=SNRp_cut,
+                SNRi_cut=SNRi_cut,
+                flux_lim=flux_lim,
+                step_vec=step_vec,
+                scale_vec=scale_vec,
+                savename="_".join([figname, "I_err"]),
+                plots_folder=plots_folder,
+                display="I_err",
+                **options,
+            )
+            proj_plots.polarization_map(
+                deepcopy(Stokes_hdul),
+                data_mask,
+                SNRp_cut=SNRp_cut,
+                SNRi_cut=SNRi_cut,
+                flux_lim=flux_lim,
+                step_vec=step_vec,
+                scale_vec=scale_vec,
+                savename="_".join([figname, "P_err"]),
+                plots_folder=plots_folder,
+                display="Pol_deg_err",
+                **options,
+            )
+            proj_plots.polarization_map(
+                deepcopy(Stokes_hdul),
+                data_mask,
+                SNRp_cut=SNRp_cut,
+                SNRi_cut=SNRi_cut,
+                flux_lim=flux_lim,
+                step_vec=step_vec,
+                scale_vec=scale_vec,
+                savename="_".join([figname, "SNRi"]),
+                plots_folder=plots_folder,
+                display="SNRi",
+                **options,
+            )
+            proj_plots.polarization_map(
+                deepcopy(Stokes_hdul),
+                data_mask,
+                SNRp_cut=SNRp_cut,
+                SNRi_cut=SNRi_cut,
+                flux_lim=flux_lim,
+                step_vec=step_vec,
+                scale_vec=scale_vec,
+                savename="_".join([figname, "SNRp"]),
+                plots_folder=plots_folder,
+                display="SNRp",
+                **options,
+            )
         elif not interactive:
-            proj_plots.polarization_map(deepcopy(Stokes_hdul), data_mask, SNRp_cut=SNRp_cut, SNRi_cut=SNRi_cut,
+            proj_plots.polarization_map(
-                                        savename=figname, plots_folder=plots_folder, display='integrate', **options)
+                deepcopy(Stokes_hdul),
-        elif pxscale.lower() not in ['full', 'integrate']:
+                data_mask,
+                SNRp_cut=SNRp_cut,
+                SNRi_cut=SNRi_cut,
+                savename=figname,
+                plots_folder=plots_folder,
+                display="integrate",
+                **options,
+            )
+        elif pxscale.lower() not in ["full", "integrate"]:
             proj_plots.pol_map(Stokes_hdul, SNRp_cut=SNRp_cut, SNRi_cut=SNRi_cut, flux_lim=flux_lim)
 
-
     return outfiles
 
 
 if __name__ == "__main__":
     import argparse
 
-    parser = argparse.ArgumentParser(description='Query MAST for target products')
+    parser = argparse.ArgumentParser(description="Query MAST for target products")
-    parser.add_argument('-t', '--target', metavar='targetname', required=False, help='the name of the target', type=str, default=None)
+    parser.add_argument("-t", "--target", metavar="targetname", required=False, help="the name of the target", type=str, default=None)
-    parser.add_argument('-p', '--proposal_id', metavar='proposal_id', required=False, help='the proposal id of the data products', type=int, default=None)
+    parser.add_argument("-p", "--proposal_id", metavar="proposal_id", required=False, help="the proposal id of the data products", type=int, default=None)
-    parser.add_argument('-d', '--data_dir', metavar='directory_path', required=False, help='directory path to the data products', type=str, default=None)
+    parser.add_argument("-f", "--files", metavar="path", required=False, nargs="*", help="the full or relative path to the data products", default=None)
-    parser.add_argument('-f', '--files', metavar='path', required=False, nargs='*', help='the full or relative path to the data products', default=None)
+    parser.add_argument(
-    parser.add_argument('-o', '--output_dir', metavar='directory_path', required=False,
+        "-o", "--output_dir", metavar="directory_path", required=False, help="output directory path for the data products", type=str, default="./data"
-                        help='output directory path for the data products', type=str, default="./data")
+    )
-    parser.add_argument('-c', '--crop', action='store_true', required=False, help='whether to crop the analysis region')
+    parser.add_argument("-c", "--crop", action="store_true", required=False, help="whether to crop the analysis region")
-    parser.add_argument('-i', '--interactive', action='store_true', required=False, help='whether to output to the interactive analysis tool')
+    parser.add_argument("-i", "--interactive", action="store_true", required=False, help="whether to output to the interactive analysis tool")
-    
     args = parser.parse_args()
-    exitcode = main(target=args.target, proposal_id=args.proposal_id, data_dir=args.data_dir, infiles=args.files,
+    exitcode = main(
-                    output_dir=args.output_dir, crop=args.crop, interactive=args.interactive)
+        target=args.target, proposal_id=args.proposal_id, infiles=args.files, output_dir=args.output_dir, crop=args.crop, interactive=args.interactive
-    print("Finished with ExitCode: ", exitcode)
+    )
+    print("Written to: ", exitcode)

package/lib/plots.py

@@ -406,7 +406,7 @@ def polarization_map(
     plt.rcdefaults()
     ratiox = max(int(stkI.shape[1]/(stkI.shape[0])),1)
     ratioy = max(int((stkI.shape[0])/stkI.shape[1]),1)
-    fig, ax = plt.subplots(figsize=(6*ratiox, 6*ratioy), layout="compressed", subplot_kw=dict(projection=wcs))
+    fig, ax = plt.subplots(figsize=(7*ratiox, 7*ratioy), layout="compressed", subplot_kw=dict(projection=wcs))
     ax.set(aspect="equal", fc="k", ylim=[-stkI.shape[0]*0.10,stkI.shape[0]*1.15])
     # fig.subplots_adjust(hspace=0, wspace=0, left=0.102, right=1.02)
 
@@ -531,8 +531,8 @@ def polarization_map(
         ax.transAxes,
         "E",
         "N",
-        length=-0.05,
+        length=-0.07,
-        fontsize=0.02,
+        fontsize=0.03,
         loc=1,
         aspect_ratio=-(stkI.shape[1]/(stkI.shape[0]*1.25)),
         sep_y=0.01,
@@ -736,7 +736,7 @@ class align_maps(object):
                 length=-0.08,
                 fontsize=0.03,
                 loc=1,
-                aspect_ratio=-(self.map_data.shape[1]/self.map_data.shape[0]),
+                aspect_ratio=-(self.map_ax.get_xbound()[1]-self.map_ax.get_xbound()[0])/(self.map_ax.get_ybound()[1]-self.map_ax.get_ybound()[0]),
                 sep_y=0.01,
                 sep_x=0.01,
                 angle=-self.map_header["orientat"],
@@ -788,13 +788,13 @@ class align_maps(object):
             )
         if "ORIENTAT" in list(self.other_header.keys()):
             north_dir2 = AnchoredDirectionArrows(
-                self.map_ax.transAxes,
+                self.other_ax.transAxes,
                 "E",
                 "N",
                 length=-0.08,
                 fontsize=0.03,
                 loc=1,
-                aspect_ratio=-(self.other_data.shape[1]/self.other_data.shape[0]),
+                aspect_ratio=-(self.other_ax.get_xbound()[1]-self.other_ax.get_xbound()[0])/(self.other_ax.get_ybound()[1]-self.other_ax.get_ybound()[0]),
                 sep_y=0.01,
                 sep_x=0.01,
                 angle=-self.other_header["orientat"],
@@ -1338,7 +1338,9 @@ class overplot_pol(align_maps):
         pol[SNRi < SNRi_cut] = np.nan
 
         plt.rcParams.update({"font.size": 16})
-        self.fig_overplot, self.ax_overplot = plt.subplots(figsize=(11, 10), subplot_kw=dict(projection=self.other_wcs))
+        ratiox = max(int(stkI.shape[1]/stkI.shape[0]),1)
+        ratioy = max(int(stkI.shape[0]/stkI.shape[1]),1)
+        self.fig_overplot, self.ax_overplot = plt.subplots(figsize=(10*ratiox, 10*ratioy), subplot_kw=dict(projection=self.other_wcs))
         self.fig_overplot.subplots_adjust(hspace=0, wspace=0, bottom=0.1, left=0.1, top=0.80, right=1.02)
 
         self.ax_overplot.set_xlabel(label="Right Ascension (J2000)")
@@ -1393,11 +1395,12 @@ class overplot_pol(align_maps):
         )
 
         # Display full size polarization vectors
+        px_scale = self.wcs_UV.wcs.get_cdelt()[0]/self.other_wcs.wcs.get_cdelt()[0]
         if scale_vec is None:
-            self.scale_vec = 2.0
+            self.scale_vec = 2.0*px_scale
             pol[np.isfinite(pol)] = 1.0 / 2.0
         else:
-            self.scale_vec = scale_vec
+            self.scale_vec = scale_vec*px_scale
         step_vec = 1
         self.X, self.Y = np.meshgrid(np.arange(stkI.shape[1]), np.arange(stkI.shape[0]))
         self.U, self.V = pol * np.cos(np.pi / 2.0 + pang * np.pi / 180.0), pol * np.sin(np.pi / 2.0 + pang * np.pi / 180.0)
@@ -1414,8 +1417,8 @@ class overplot_pol(align_maps):
             headwidth=0.0,
             headlength=0.0,
             headaxislength=0.0,
-            width=0.5,
+            width=0.5*px_scale,
-            linewidth=0.75,
+            linewidth=0.3*px_scale,
             color="white",
             edgecolor="black",
             transform=self.ax_overplot.get_transform(self.wcs_UV),
@@ -1454,7 +1457,7 @@ class overplot_pol(align_maps):
             length=-0.08,
             fontsize=0.03,
             loc=1,
-            aspect_ratio=-(stkI.shape[1]/stkI.shape[0]),
+            aspect_ratio=-(self.ax_overplot.get_xbound()[1]-self.ax_overplot.get_xbound()[0])/(self.ax_overplot.get_ybound()[1]-self.ax_overplot.get_ybound()[0]),
             sep_y=0.01,
             sep_x=0.01,
             angle=-self.Stokes_UV[0].header["orientat"],

package/lib/reduction.py

@@ -217,9 +217,9 @@ def bin_ndarray(ndarray, new_shape, operation="sum"):
             elif operation.lower() in ["mean", "average", "avg"]:
                 ndarray = ndarray.mean(-1 * (i + 1))
     else:
-        row_comp = np.mat(get_row_compressor(ndarray.shape[0], new_shape[0], operation))
+        row_comp = np.asmatrix(get_row_compressor(ndarray.shape[0], new_shape[0], operation))
-        col_comp = np.mat(get_column_compressor(ndarray.shape[1], new_shape[1], operation))
+        col_comp = np.asmatrix(get_column_compressor(ndarray.shape[1], new_shape[1], operation))
-        ndarray = np.array(row_comp * np.mat(ndarray) * col_comp)
+        ndarray = np.array(row_comp * np.asmatrix(ndarray) * col_comp)
 
     return ndarray