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

Allow subtracting the bkg simpler

package/FOC_reduction.py

@@ -36,12 +36,12 @@ def main(target=None, proposal_id=None, infiles=None, output_dir="./data", crop=
 
     # Background estimation
     error_sub_type = "freedman-diaconis"  # sqrt, sturges, rice, scott, freedman-diaconis (default) or shape (example (51, 51))
-    subtract_error = 2.0
+    subtract_error = 1.0
     display_bkg = False
 
     # Data binning
-    pxsize = 2
-    pxscale = "px"  # pixel, arcsec or full
+    pxsize = 0.05
+    pxscale = "arcsec"  # pixel, arcsec or full
     rebin_operation = "sum"  # sum or average
 
     # Alignement
@@ -54,8 +54,8 @@ def main(target=None, proposal_id=None, infiles=None, output_dir="./data", crop=
 
     # Smoothing
     smoothing_function = "combine"  # gaussian_after, weighted_gaussian_after, gaussian, weighted_gaussian or combine
-    smoothing_FWHM = 1.50       # If None, no smoothing is done
-    smoothing_scale = "px"      # pixel or arcsec
+    smoothing_FWHM = 0.10  # If None, no smoothing is done
+    smoothing_scale = "arcsec"  # pixel or arcsec
 
     # Rotation
     rotate_North = True
@@ -67,6 +67,16 @@ def main(target=None, proposal_id=None, infiles=None, output_dir="./data", crop=
     scale_vec = 5
     step_vec = 1  # plot all vectors in the array. if step_vec = 2, then every other vector will be plotted if step_vec = 0 then all vectors are displayed at full length
 
+    # Adaptive binning
+    # in order to perfrom optimal binning, there are several steps to follow:
+    # 1. Load the data again and preserve the full images
+    # 2. Skip the cropping step but use the same error and background estimation
+    # 3. Use the same alignment as the routine
+    # 4. Skip the rebinning step
+    # 5. Calulate the Stokes parameters without smoothing
+    optimal_binning = False
+    optimize = False
+
     # Pipeline start
 
     # Step 1:
@@ -111,297 +121,592 @@ def main(target=None, proposal_id=None, infiles=None, output_dir="./data", crop=
     if align_center is None:
         figtype = "_".join([figtype, "not_aligned"] if figtype != "" else ["not_aligned"])
 
-    #  Crop data to remove outside blank margins.
-    data_array, error_array, headers = proj_red.crop_array(
-        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)
+    if optimal_binning:
+        from lib.background import subtract_bkg
 
-    #  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)
+        options = {"optimize": optimize, "optimal_binning": True}
 
-    #  Estimate error from data background, estimated from sub-image of desired sub_shape.
-    background = None
-    data_array, error_array, headers, background = 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,
-    )
+        # Step 1: Load the data again and preserve the full images
+        _data_array, _headers = deepcopy(data_array), deepcopy(headers)  # Preserve full images
+        _data_mask = np.ones(_data_array[0].shape, dtype=bool)
 
-    # Rotate data to have same orientation
-    rotate_data = np.unique([np.round(float(head["ORIENTAT"]), 3) for head in headers]).size != 1
-    if rotate_data:
-        ang = np.mean([head["ORIENTAT"] for head in headers])
-        for head in headers:
-            head["ORIENTAT"] -= ang
-        data_array, error_array, data_mask, headers = proj_red.rotate_data(data_array, error_array, data_mask, headers)
-        if display_data:
+        # 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.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 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,
+        )
+        _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,
+            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()) ** 2
+                        / np.sum([h["filtnam1"] == head["filtnam1"] for h in _headers])
+                    )
+                ).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,
+            _header_stokes,
+            _data_mask,
+            figname,
+            data_folder=data_folder,
+            return_hdul=True,
+        )
+
+        # Step 6:
+        _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"],
+                *sci_not(
+                    _Stokes_hdul[0].data[_data_mask].sum() * _header_stokes["PHOTFLAM"],
+                    np.sqrt(_Stokes_hdul[3].data[0, 0][_data_mask].sum()) * _header_stokes["PHOTFLAM"],
+                    2,
+                    out=int,
+                ),
+            )
+        )
+        print("P_int = {0:.1f} ± {1:.1f} %".format(_header_stokes["p_int"] * 100.0, np.ceil(_header_stokes["sP_int"] * 1000.0) / 10.0))
+        print("PA_int = {0:.1f} ± {1:.1f} °".format(princ_angle(_header_stokes["pa_int"]), princ_angle(np.ceil(_header_stokes["sPA_int"] * 10.0) / 10.0)))
+        #  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),
+            )
+        )
+        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:
+            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]),
+                plots_folder=plots_folder,
+                **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"]),
+                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,
+                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}
+        #  Crop data to remove outside blank margins.
+        data_array, error_array, headers = proj_red.crop_array(
+            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
+            )
+
+        #  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,
+        )
+
+        #  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
+        )
+
+        if display_align:
+            print("Image shifts: {} \nShifts uncertainty: {}".format(shifts, error_shifts))
             proj_plots.plot_obs(
                 data_array,
                 headers,
-                savename="_".join([figname, "rotate_data"]),
+                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"]
                 ),
             )
 
-    #  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,
-        data_mask=data_mask,
-        background=background,
-        upsample_factor=10,
-        ref_center=align_center,
-        return_shifts=True,
-    )
+        #  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
+            )
 
-    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(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
-        )
-
-    # Plot array for checking output
-    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(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()) ** 2
-                    / np.sum([h["filtnam1"] == head["filtnam1"] for h in headers])
+        # Rotate data to have same orientation
+        rotate_data = np.unique([np.round(float(head["ORIENTAT"]), 3) for head in headers]).size != 1
+        if rotate_data:
+            ang = np.mean([head["ORIENTAT"] for head in headers])
+            for head in headers:
+                head["ORIENTAT"] -= ang
+            data_array, error_array, data_mask, headers = proj_red.rotate_data(data_array, error_array, data_mask, headers)
+            if display_data:
+                proj_plots.plot_obs(
+                    data_array,
+                    headers,
+                    savename="_".join([figname, "rotate_data"]),
+                    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"]
+                    ),
                 )
-            ).reshape(1, 1)
-            for bkg, head in zip(background, headers)
-        ]
-    )
 
-    # Step 2:
-    # Compute Stokes I, Q, U with smoothed polarized images
-    # SMOOTHING DISCUSSION :
-    # FWHM of FOC have been estimated at about 0.03" across 1500-5000 Angstrom band, which is about 2 detector pixels wide
-    # 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
-    )
-    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
-    )
+        # Plot array for checking output
+        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(
+                    vmin=data_array[data_array > 0.0].min() * headers[0]["photflam"], vmax=data_array[data_array > 0.0].max() * headers[0]["photflam"]
+                ),
+            )
 
-    # 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_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
+        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()) ** 2
+                        / np.sum([h["filtnam1"] == head["filtnam1"] for h in headers])
+                    )
+                ).reshape(1, 1)
+                for bkg, head in zip(background, headers)
+            ]
         )
 
-    # 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 2:
+        # Compute Stokes I, Q, U with smoothed polarized images
+        # SMOOTHING DISCUSSION :
+        # FWHM of FOC have been estimated at about 0.03" across 1500-5000 Angstrom band, which is about 2 detector pixels wide
+        # 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
+        )
+        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 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,
-        header_stokes,
-        data_mask,
-        figname,
-        data_folder=data_folder,
-        return_hdul=True,
-    )
-    outfiles.append("/".join([data_folder, Stokes_hdul[0].header["FILENAME"] + ".fits"]))
+        # 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_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
+            )
 
-    # Step 5:
-    # crop to desired region of interest (roi)
-    if crop:
-        figname += "_crop"
-        stokescrop = proj_plots.crop_Stokes(deepcopy(Stokes_hdul), norm=LogNorm())
-        stokescrop.crop()
-        stokescrop.write_to("/".join([data_folder, figname + ".fits"]))
-        Stokes_hdul, header_stokes = stokescrop.hdul_crop, stokescrop.hdul_crop[0].header
+        # 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 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,
+            header_stokes,
+            data_mask,
+            figname,
+            data_folder=data_folder,
+            return_hdul=True,
+        )
         outfiles.append("/".join([data_folder, Stokes_hdul[0].header["FILENAME"] + ".fits"]))
 
-    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"],
-            *sci_not(
-                Stokes_hdul[0].data[data_mask].sum() * header_stokes["PHOTFLAM"],
-                np.sqrt(Stokes_hdul[3].data[0, 0][data_mask].sum()) * header_stokes["PHOTFLAM"],
-                2,
-                out=int,
-            ),
+        # Step 5:
+        # crop to desired region of interest (roi)
+        if crop:
+            figname += "_crop"
+            stokescrop = proj_plots.crop_Stokes(deepcopy(Stokes_hdul), norm=LogNorm())
+            stokescrop.crop()
+            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)
+        print(
+            "F_int({0:.0f} Angs) = ({1} ± {2})e{3} ergs.cm^-2.s^-1.Angs^-1".format(
+                header_stokes["PHOTPLAM"],
+                *sci_not(
+                    Stokes_hdul[0].data[data_mask].sum() * header_stokes["PHOTFLAM"],
+                    np.sqrt(Stokes_hdul[3].data[0, 0][data_mask].sum()) * header_stokes["PHOTFLAM"],
+                    2,
+                    out=int,
+                ),
+            )
         )
-    )
-    print("P_int = {0:.1f} ± {1:.1f} %".format(header_stokes["p_int"] * 100.0, np.ceil(header_stokes["sP_int"] * 1000.0) / 10.0))
-    print("PA_int = {0:.1f} ± {1:.1f} °".format(princ_angle(header_stokes["pa_int"]), princ_angle(np.ceil(header_stokes["sPA_int"] * 10.0) / 10.0)))
-    #  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["photflam"], np.sqrt(S_cov_bkg[0, 0][0, 0]) * header_stokes["photflam"], 2, out=int)
+        print("P_int = {0:.1f} ± {1:.1f} %".format(header_stokes["p_int"] * 100.0, np.ceil(header_stokes["sP_int"] * 1000.0) / 10.0))
+        print("PA_int = {0:.1f} ± {1:.1f} °".format(princ_angle(header_stokes["pa_int"]), princ_angle(np.ceil(header_stokes["sPA_int"] * 10.0) / 10.0)))
+        #  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),
+            )
         )
-    )
-    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:
-        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]),
-            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,
-            scale_vec=scale_vec,
-            savename="_".join([figname, "I"]),
-            plots_folder=plots_folder,
-            display="Intensity",
-        )
-        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_flux"]),
-            plots_folder=plots_folder,
-            display="Pol_Flux",
-        )
-        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",
-        )
-        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",
-        )
-        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",
-        )
-        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",
-        )
-        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",
-        )
-        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",
-        )
-    elif not interactive:
-        proj_plots.polarization_map(
-            deepcopy(Stokes_hdul), data_mask, SNRp_cut=SNRp_cut, SNRi_cut=SNRi_cut, savename=figname, plots_folder=plots_folder, display="integrate"
-        )
-    elif pxscale.lower() not in ["full", "integrate"]:
-        proj_plots.pol_map(Stokes_hdul, SNRp_cut=SNRp_cut, SNRi_cut=SNRi_cut, step_vec=step_vec, scale_vec=scale_vec, flux_lim=flux_lim)
+        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:
+            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]),
+                plots_folder=plots_folder,
+                **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"]),
+                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,
+                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
 

package/lib/background.py

@@ -251,23 +251,18 @@ def bkg_fit(data, error, mask, headers, subtract_error=True, display=False, save
         weights = 1 / chi2**2
         weights /= weights.sum()
 
-        bkg = np.sum(weights * (coeff[:, 1] + np.abs(coeff[:, 2]) * subtract_error))
-
+        bkg = np.sum(weights*(coeff[:, 1]+np.abs(coeff[:, 2]) * subtract_error))
         error_bkg[i] *= bkg
 
-        n_error_array[i] = np.sqrt(n_error_array[i] ** 2 + error_bkg[i] ** 2)
-
-        # Substract background
-        if np.abs(subtract_error) > 0:
-            n_data_array[i][mask] = n_data_array[i][mask] - bkg
-            n_data_array[i][np.logical_and(mask, n_data_array[i] <= 1e-3 * bkg)] = 1e-3 * bkg
-
         std_bkg[i] = image[np.abs(image - bkg) / bkg < 1.0].std()
         background[i] = bkg
+    
+    if np.abs(subtract_error) > 0:
+        n_data_array, n_error_array, background, error_bkg = subtract_bkg(n_data_array, n_error_array, mask, background, error_bkg)
 
     if display:
         display_bkg(data, background, std_bkg, headers, histograms=histograms, binning=binning, coeff=coeff, savename=savename, plots_folder=plots_folder)
-    return n_data_array, n_error_array, headers, background
+    return n_data_array, n_error_array, headers, background, error_bkg
 
 
 def bkg_hist(data, error, mask, headers, sub_type=None, subtract_error=True, display=False, savename=None, plots_folder=""):
@@ -360,23 +355,19 @@ def bkg_hist(data, error, mask, headers, sub_type=None, subtract_error=True, dis
         # popt, pcov = curve_fit(gausspol, binning[-1], hist, p0=p0)
         popt, pcov = curve_fit(gauss, binning[-1], hist, p0=p0)
         coeff.append(popt)
+
         bkg = popt[1] + np.abs(popt[2]) * subtract_error
-
         error_bkg[i] *= bkg
 
-        n_error_array[i] = np.sqrt(n_error_array[i] ** 2 + error_bkg[i] ** 2)
-
-        # Substract background
-        if np.abs(subtract_error) > 0:
-            n_data_array[i][mask] = n_data_array[i][mask] - bkg
-            n_data_array[i][np.logical_and(mask, n_data_array[i] <= 1e-3 * bkg)] = 1e-3 * bkg
-
         std_bkg[i] = image[np.abs(image - bkg) / bkg < 1.0].std()
         background[i] = bkg
+    
+    if np.abs(subtract_error) > 0:
+        n_data_array, n_error_array, background, error_bkg = subtract_bkg(n_data_array, n_error_array, mask, background, error_bkg)
 
     if display:
         display_bkg(data, background, std_bkg, headers, histograms=histograms, binning=binning, coeff=coeff, savename=savename, plots_folder=plots_folder)
-    return n_data_array, n_error_array, headers, background
+    return n_data_array, n_error_array, headers, background, error_bkg
 
 
 def bkg_mini(data, error, mask, headers, sub_shape=(15, 15), subtract_error=True, display=False, savename=None, plots_folder=""):
@@ -458,19 +449,28 @@ def bkg_mini(data, error, mask, headers, sub_shape=(15, 15), subtract_error=True
         # Compute error : root mean square of the background
         sub_image = image[minima[0] : minima[0] + sub_shape[0], minima[1] : minima[1] + sub_shape[1]]
         # bkg =  np.std(sub_image)    # Previously computed using standard deviation over the background
-        bkg = np.sqrt(np.sum(sub_image**2) / sub_image.size) * subtract_error if subtract_error > 0 else np.sqrt(np.sum(sub_image**2) / sub_image.size)
+
+        bkg = np.sqrt(np.sum(sub_image**2)/sub_image.size)*subtract_error if subtract_error > 0 else np.sqrt(np.sum(sub_image**2)/sub_image.size)
         error_bkg[i] *= bkg
 
-        n_error_array[i] = np.sqrt(n_error_array[i] ** 2 + error_bkg[i] ** 2)
-
-        # Substract background
-        if np.abs(subtract_error) > 0:
-            n_data_array[i][mask] = n_data_array[i][mask] - bkg
-            n_data_array[i][np.logical_and(mask, n_data_array[i] <= 1e-3 * bkg)] = 1e-3 * bkg
-
         std_bkg[i] = image[np.abs(image - bkg) / bkg < 1.0].std()
         background[i] = bkg
 
+    if np.abs(subtract_error) > 0:
+        n_data_array, n_error_array, background, error_bkg = subtract_bkg(n_data_array, n_error_array, mask, background, error_bkg)
+    
     if display:
         display_bkg(data, background, std_bkg, headers, rectangle=rectangle, savename=savename, plots_folder=plots_folder)
-    return n_data_array, n_error_array, headers, background
+    return n_data_array, n_error_array, headers, background, error_bkg
+
+def subtract_bkg(data, error, mask, background, error_bkg):
+    assert data.ndim == 3, "Input data must have more than 1 image."
+    
+    n_data_array, n_error_array = deepcopy(data), deepcopy(error)
+    
+    for i in range(data.shape[0]):
+        n_data_array[i][mask] = n_data_array[i][mask] - background[i]
+        n_data_array[i][np.logical_and(mask, n_data_array[i] <= 1e-3 * background[i])] = 1e-3 * background[i]
+        n_error_array[i] = np.sqrt(n_error_array[i]**2 + error_bkg[i]**2)
+    
+    return n_data_array, n_error_array, background, error_bkg

package/lib/plots.py

@@ -43,6 +43,7 @@ prototypes :
 from copy import deepcopy
 from os.path import join as path_join
 
+
 import matplotlib.font_manager as fm
 import matplotlib.patheffects as pe
 import matplotlib.pyplot as plt
@@ -58,6 +59,7 @@ from mpl_toolkits.axes_grid1.anchored_artists import (
     AnchoredDirectionArrows,
     AnchoredSizeBar,
 )
+
 from scipy.ndimage import zoom as sc_zoom
 
 try:
@@ -65,6 +67,72 @@ try:
 except ImportError:
     from utils import princ_angle, rot2D, sci_not
 
+def adaptive_binning(I_stokes, Q_stokes, U_stokes, Stokes_cov):
+    shape = I_stokes.shape
+    
+    assert shape[0] == shape[1], "Only square images are supported"
+    assert shape[0] % 2 == 0, "Image size must be a power of 2"
+    
+    n = int(np.log2(shape[0]))
+    bin_map = np.zeros(shape)
+    bin_num = 0
+
+    for level in range(n):
+        grid_size = 2**level
+        temp_I = I_stokes.reshape(int(shape[0]/grid_size), grid_size, int(shape[1]/grid_size), grid_size).sum(1).sum(2)
+        temp_Q = Q_stokes.reshape(int(shape[0]/grid_size), grid_size, int(shape[1]/grid_size), grid_size).sum(1).sum(2)
+        temp_U = U_stokes.reshape(int(shape[0]/grid_size), grid_size, int(shape[1]/grid_size), grid_size).sum(1).sum(2)
+        temp_cov = Stokes_cov.reshape(3, 3, int(shape[0]/grid_size), grid_size, int(shape[1]/grid_size), grid_size).sum(3).sum(4)
+        temp_bin_map = bin_map.reshape(int(shape[0]/grid_size), grid_size, int(shape[1]/grid_size), grid_size).sum(1).sum(2)
+
+        temp_P = (temp_Q**2 + temp_U**2)**0.5 / temp_I
+        temp_P_err = (1 / temp_I) * np.sqrt((temp_Q**2 * temp_cov[1,1,:,:] + temp_U**2 * temp_cov[2,2,:,:] + 2. * temp_Q * temp_U * temp_cov[1,2,:,:]) / (temp_Q**2 + temp_U**2) + \
+                                            ((temp_Q / temp_I)**2 + (temp_U / temp_I)**2) * temp_cov[0,0,:,:] - \
+                                                2. * (temp_Q / temp_I) * temp_cov[0,1,:,:] - \
+                                                2. * (temp_U / temp_I) * temp_cov[0,2,:,:])
+
+        for i in range(int(shape[0]/grid_size)):
+            for j in range(int(shape[1]/grid_size)):
+                if  (temp_P[i,j] / temp_P_err[i,j] > 3) and (temp_bin_map[i,j] == 0): # the default criterion is 3 sigma in P
+                    bin_num += 1
+                    bin_map[i*grid_size:(i+1)*grid_size,j*grid_size:(j+1)*grid_size] = bin_num
+    
+    return bin_map, bin_num
+
+def plot_quiver(ax, stkI, stkQ, stkU, stk_cov, poldata, pangdata, step_vec=1., scale_vec=2., optimal_binning=False):
+    if optimal_binning:
+        bin_map, bin_num = adaptive_binning(stkI, stkQ, stkU, stk_cov)
+        shape = stkI.shape
+        
+        for i in range(1, bin_num+1):
+            bin = np.where(bin_map==i)
+            x_center, y_center = np.mean(bin, axis=1)
+            
+            if not (20 < x_center < shape[0]-20 and 20 < y_center < shape[1]-20): # avoid plotting vectors on the edges of the image
+                continue
+
+            bin_I = np.sum(stkI[bin])
+            bin_Q = np.sum(stkQ[bin])
+            bin_U = np.sum(stkU[bin])
+            bin_cov = np.zeros((3,3))
+            for i in range(3):
+                for j in range(3):
+                    bin_cov[i,j] = np.sum(stk_cov[i,j][bin])
+
+            poldata = np.sqrt(bin_Q**2 + bin_U**2) / bin_I
+            pangdata = 0.5 * np.arctan2(bin_U, bin_Q)
+            pangdata_err = (1 / (2. *(bin_Q**2 + bin_U**2))) * \
+                        np.sqrt(bin_U**2 * bin_cov[1,1] + bin_Q**2 * bin_cov[2,2] - 2. * bin_Q * bin_U * bin_cov[1,2])
+
+            ax.quiver(y_center, x_center, poldata * np.cos(np.pi/2.+pangdata), poldata * np.sin(np.pi/2.+pangdata), units='xy', angles='uv', scale=1./scale_vec, scale_units='xy', pivot='mid', headwidth=0., headlength=0., headaxislength=0., width=0.1, linewidth=0.5, color='white', edgecolor='white')
+            ax.quiver(y_center, x_center, poldata * np.cos(np.pi/2.+pangdata+pangdata_err), poldata * np.sin(np.pi/2.+pangdata+pangdata_err), units='xy', angles='uv', scale=1./scale_vec, scale_units='xy', pivot='mid', headwidth=0., headlength=0., headaxislength=0., width=0.1, linewidth=0.5, color='black', edgecolor='black', ls='dashed')
+            ax.quiver(y_center, x_center, poldata * np.cos(np.pi/2.+pangdata-pangdata_err), poldata * np.sin(np.pi/2.+pangdata-pangdata_err), units='xy', angles='uv', scale=1./scale_vec, scale_units='xy', pivot='mid', headwidth=0., headlength=0., headaxislength=0., width=0.1, linewidth=0.5, color='black', edgecolor='black', ls='dashed')
+
+    else:
+        X, Y = np.meshgrid(np.arange(stkI.shape[1]), np.arange(stkI.shape[0]))
+        U, V = poldata*np.cos(np.pi/2.+pangdata*np.pi/180.), poldata*np.sin(np.pi/2.+pangdata*np.pi/180.)
+        ax.quiver(X[::step_vec, ::step_vec], Y[::step_vec, ::step_vec], U[::step_vec, ::step_vec], V[::step_vec, ::step_vec], units='xy', angles='uv', scale=1./scale_vec, scale_units='xy', pivot='mid', headwidth=0., headlength=0., headaxislength=0., width=0.5, linewidth=0.75, color='w', edgecolor='k')
+
 
 def plot_obs(data_array, headers, rectangle=None, savename=None, plots_folder="", **kwargs):
     """
@@ -99,7 +167,10 @@ def plot_obs(data_array, headers, rectangle=None, savename=None, plots_folder=""
     plt.rcParams.update({"font.size": 10})
     nb_obs = np.max([np.sum([head["filtnam1"] == curr_pol for head in headers]) for curr_pol in ["POL0", "POL60", "POL120"]])
     shape = np.array((3, nb_obs))
-    fig, ax = plt.subplots(shape[0], shape[1], figsize=(3 * shape[1], 3 * shape[0]), dpi=200, layout="constrained", sharex=True, sharey=True)
+
+    fig, ax = plt.subplots(shape[0], shape[1], figsize=(3*shape[1], 3*shape[0]), dpi=200, layout='constrained',
+                            sharex=True, sharey=True)
+
     r_pol = dict(pol0=0, pol60=1, pol120=2)
     c_pol = dict(pol0=0, pol60=0, pol120=0)
     for i, (data, head) in enumerate(zip(data_array, headers)):
@@ -212,6 +283,7 @@ def plot_Stokes(Stokes, savename=None, plots_folder=""):
     return 0
 
 
+
 def polarization_map(
     Stokes,
     data_mask=None,
@@ -224,7 +296,9 @@ def polarization_map(
     savename=None,
     plots_folder="",
     display="default",
+    **kwargs
 ):
+
     """
     Plots polarization map from Stokes HDUList.
     ----------
@@ -275,11 +349,17 @@ def polarization_map(
         The figure and ax created for interactive contour maps.
     """
     # Get data
-    stkI = Stokes["I_stokes"].data.copy()
-    stk_cov = Stokes["IQU_cov_matrix"].data.copy()
-    pol = Stokes["Pol_deg_debiased"].data.copy()
-    pol_err = Stokes["Pol_deg_err"].data.copy()
-    pang = Stokes["Pol_ang"].data.copy()
+
+    optimal_binning = kwargs.get('optimal_binning', False)
+    
+    stkI = Stokes['I_stokes'].data.copy()
+    stkQ = Stokes['Q_stokes'].data.copy()
+    stkU = Stokes['U_stokes'].data.copy()
+    stk_cov = Stokes['IQU_cov_matrix'].data.copy()
+    pol = Stokes['Pol_deg_debiased'].data.copy()
+    pol_err = Stokes['Pol_deg_err'].data.copy()
+    pang = Stokes['Pol_ang'].data.copy()
+
     try:
         if data_mask is None:
             data_mask = Stokes["Data_mask"].data.astype(bool).copy()
@@ -392,11 +472,13 @@ def polarization_map(
         # Display intensity error map
         display = "s_i"
         if (SNRi > SNRi_cut).any():
+
             vmin, vmax = (
                 1.0 / 2.0 * np.median(np.sqrt(stk_cov[0, 0][stk_cov[0, 0] > 0.0]) * convert_flux),
                 np.max(np.sqrt(stk_cov[0, 0][stk_cov[0, 0] > 0.0]) * convert_flux),
             )
             im = ax.imshow(np.sqrt(stk_cov[0, 0]) * convert_flux, norm=LogNorm(vmin, vmax), aspect="equal", cmap="inferno_r", alpha=1.0)
+
         else:
             im = ax.imshow(np.sqrt(stk_cov[0, 0]) * convert_flux, aspect="equal", cmap="inferno", alpha=1.0)
         fig.colorbar(im, ax=ax, aspect=50, shrink=0.75, pad=0.025, label=r"$\sigma_I$ [$ergs \cdot cm^{-2} \cdot s^{-1} \cdot \AA^{-1}$]")
@@ -467,28 +549,10 @@ def polarization_map(
         if step_vec == 0:
             poldata[np.isfinite(poldata)] = 1.0 / 2.0
             step_vec = 1
-            scale_vec = 2.0
-        X, Y = np.meshgrid(np.arange(stkI.shape[1]), np.arange(stkI.shape[0]))
-        U, V = poldata * np.cos(np.pi / 2.0 + pangdata * np.pi / 180.0), poldata * np.sin(np.pi / 2.0 + pangdata * np.pi / 180.0)
-        ax.quiver(
-            X[::step_vec, ::step_vec],
-            Y[::step_vec, ::step_vec],
-            U[::step_vec, ::step_vec],
-            V[::step_vec, ::step_vec],
-            units="xy",
-            angles="uv",
-            scale=1.0 / scale_vec,
-            scale_units="xy",
-            pivot="mid",
-            headwidth=0.0,
-            headlength=0.0,
-            headaxislength=0.0,
-            width=0.5,
-            linewidth=0.75,
-            color="w",
-            edgecolor="k",
-        )
-        pol_sc = AnchoredSizeBar(ax.transData, scale_vec, r"$P$= 100 %", 4, pad=0.25, sep=5, borderpad=0.25, frameon=False, size_vertical=0.005, color="w")
+            scale_vec = 2.
+
+        plot_quiver(ax, stkI, stkQ, stkU, stk_cov, poldata, pangdata, step_vec=step_vec, scale_vec=scale_vec, optimal_binning=optimal_binning)
+        pol_sc = AnchoredSizeBar(ax.transData, scale_vec, r"$P$= 100 %", 4, pad=0.5, sep=5, borderpad=0.5, frameon=False, size_vertical=0.005, color='w')
 
         ax.add_artist(pol_sc)
         ax.add_artist(px_sc)
@@ -528,6 +592,7 @@ def polarization_map(
     # Display instrument FOV
     if rectangle is not None:
         x, y, width, height, angle, color = rectangle
+
         x, y = np.array([x, y]) - np.array(stkI.shape) / 2.0
         ax.add_patch(Rectangle((x, y), width, height, angle=angle, edgecolor=color, fill=False))
 
@@ -590,6 +655,7 @@ class align_maps(object):
         )
 
         plt.rcParams.update({"font.size": 10})
+
         fontprops = fm.FontProperties(size=16)
         self.fig_align = plt.figure(figsize=(20, 10))
         self.map_ax = self.fig_align.add_subplot(121, projection=self.map_wcs)

package/lib/reduction.py

@@ -528,27 +528,25 @@ def get_error(
     # estimated to less than 3%
     err_flat = data * 0.03
 
-    if sub_type is None:
-        n_data_array, c_error_bkg, headers, background = bkg_hist(
-            data, error, mask, headers, subtract_error=subtract_error, display=display, savename=savename, plots_folder=plots_folder
-        )
+
+    if (sub_type is None):
+        n_data_array, c_error_bkg, headers, background, error_bkg = bkg_hist(
+            data, error, mask, headers, subtract_error=subtract_error, display=display, savename=savename, plots_folder=plots_folder)
         sub_type, subtract_error = "histogram ", str(int(subtract_error > 0.0))
     elif isinstance(sub_type, str):
-        if sub_type.lower() in ["auto"]:
-            n_data_array, c_error_bkg, headers, background = bkg_fit(
-                data, error, mask, headers, subtract_error=subtract_error, display=display, savename=savename, plots_folder=plots_folder
-            )
+        if sub_type.lower() in ['auto']:
+            n_data_array, c_error_bkg, headers, background, error_bkg = bkg_fit(
+                data, error, mask, headers, subtract_error=subtract_error, display=display, savename=savename, plots_folder=plots_folder)
             sub_type, subtract_error = "histogram fit ", "mean+%.1fsigma" % subtract_error if subtract_error != 0.0 else 0.0
         else:
-            n_data_array, c_error_bkg, headers, background = bkg_hist(
-                data, error, mask, headers, sub_type=sub_type, subtract_error=subtract_error, display=display, savename=savename, plots_folder=plots_folder
-            )
+            n_data_array, c_error_bkg, headers, background, error_bkg = bkg_hist(
+                data, error, mask, headers, sub_type=sub_type, subtract_error=subtract_error, display=display, savename=savename, plots_folder=plots_folder)
             sub_type, subtract_error = "histogram ", "mean+%.1fsigma" % subtract_error if subtract_error != 0.0 else 0.0
     elif isinstance(sub_type, tuple):
-        n_data_array, c_error_bkg, headers, background = bkg_mini(
-            data, error, mask, headers, sub_shape=sub_type, subtract_error=subtract_error, display=display, savename=savename, plots_folder=plots_folder
-        )
+        n_data_array, c_error_bkg, headers, background, error_bkg = bkg_mini(
+            data, error, mask, headers, sub_shape=sub_type, subtract_error=subtract_error, display=display, savename=savename, plots_folder=plots_folder)
         sub_type, subtract_error = "minimal flux ", "mean+%.1fsigma" % subtract_error if subtract_error != 0.0 else 0.0
+
     else:
         print("Warning: Invalid subtype.")
 
@@ -560,7 +558,7 @@ def get_error(
     n_error_array = np.sqrt(err_wav**2 + err_psf**2 + err_flat**2 + c_error_bkg**2)
 
     if return_background:
-        return n_data_array, n_error_array, headers, background
+        return n_data_array, n_error_array, headers, background, error_bkg # return background error as well
     else:
         return n_data_array, n_error_array, headers
 
@@ -693,7 +691,7 @@ def rebin_array(data_array, error_array, headers, pxsize=2, scale="px", operatio
 
 
 def align_data(
-    data_array, headers, error_array=None, data_mask=None, background=None, upsample_factor=1.0, ref_data=None, ref_center=None, return_shifts=False
+    data_array, headers, error_array=None, data_mask=None, background=None, upsample_factor=1.0, ref_data=None, ref_center=None, return_shifts=False, optimal_binning=False
 ):
     """
     Align images in data_array using cross correlation, and rescale them to
@@ -772,12 +770,13 @@ def align_data(
     full_headers.append(headers[0])
     err_array = np.concatenate((error_array, [np.zeros(ref_data.shape)]), axis=0)
 
-    if data_mask is None:
-        full_array, err_array, full_headers = crop_array(full_array, full_headers, err_array, step=5, inside=False, null_val=0.0)
-    else:
-        full_array, err_array, data_mask, full_headers = crop_array(
-            full_array, full_headers, err_array, data_mask=data_mask, step=5, inside=False, null_val=0.0
-        )
+    if not optimal_binning:
+        if data_mask is None:
+            full_array, err_array, full_headers = crop_array(full_array, full_headers, err_array, step=5, inside=False, null_val=0.0)
+        else:
+            full_array, err_array, data_mask, full_headers = crop_array(
+                full_array, full_headers, err_array, data_mask=data_mask, step=5, inside=False, null_val=0.0
+            )
 
     data_array, ref_data, headers = full_array[:-1], full_array[-1], full_headers[:-1]
     error_array = err_array[:-1]
@@ -856,7 +855,9 @@ def align_data(
         headers[i].update(headers_wcs[i].to_header())
 
     data_mask = rescaled_mask.all(axis=0)
-    data_array, error_array, data_mask, headers = crop_array(rescaled_image, headers, rescaled_error, data_mask, null_val=0.01 * background)
+    
+    if not optimal_binning:
+        data_array, error_array, data_mask, headers = crop_array(rescaled_image, headers, rescaled_error, data_mask, null_val=0.01*background)
 
     if return_shifts:
         return data_array, error_array, headers, data_mask, shifts, errors
@@ -1847,4 +1848,4 @@ def rotate_data(data_array, error_array, data_mask, headers):
     for i in range(new_data_array.shape[0]):
         new_data_array[i][new_data_array[i] < 0.0] = 0.0
 
-    return new_data_array, new_error_array, new_data_mask, new_headers
+    return new_data_array, new_error_array, new_data_mask, new_headers