fix pol_map.set_data_mask and some keyword saving

finish rebase display on main

package/FOC_reduction.py

@@ -20,6 +20,7 @@ import lib.reduction as proj_red  # Functions used in reduction pipeline
 import numpy as np
 from lib.utils import princ_angle, sci_not
 from matplotlib.colors import LogNorm
+from astropy.wcs import WCS
 
 
 def main(target=None, proposal_id=None, infiles=None, output_dir="./data", crop=False, interactive=False):
@@ -45,7 +46,7 @@ def main(target=None, proposal_id=None, infiles=None, output_dir="./data", crop=
     display_bkg = True
 
     # Data binning
-    pxsize = 0.10
+    pxsize = 0.05
     pxscale = "arcsec"  # pixel, arcsec or full
     rebin_operation = "sum"  # sum or average
 
@@ -59,10 +60,11 @@ 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 = 0.150  # If None, no smoothing is done
+    smoothing_FWHM = 0.075  # If None, no smoothing is done
     smoothing_scale = "arcsec"  # pixel or arcsec
 
     # Rotation
+    rotate_data = False
     rotate_North = True
 
     #  Polarization map output
@@ -142,7 +144,7 @@ def main(target=None, proposal_id=None, infiles=None, output_dir="./data", crop=
     )
 
     # Rotate data to have same orientation
-    rotate_data = np.unique([np.round(float(head["ORIENTAT"]), 3) for head in headers]).size != 1
+    rotate_data = rotate_data or 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:
@@ -158,7 +160,6 @@ def main(target=None, proposal_id=None, infiles=None, output_dir="./data", crop=
                     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,
@@ -182,6 +183,14 @@ def main(target=None, proposal_id=None, infiles=None, output_dir="./data", crop=
             norm=LogNorm(vmin=data_array[data_array > 0.0].min() * headers[0]["photflam"], vmax=data_array[data_array > 0.0].max() * headers[0]["photflam"]),
         )
 
+    flux_data, flux_error, flux_mask, flux_head = (
+        deepcopy(data_array.sum(axis=0)),
+        deepcopy(np.sqrt(np.sum(error_array**2, axis=0))),
+        deepcopy(data_mask),
+        deepcopy(headers[0]),
+    )
+    flux_head["EXPTIME"] = np.sum([head["EXPTIME"] for head in headers])
+
     #  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(
@@ -233,6 +242,10 @@ def main(target=None, proposal_id=None, infiles=None, output_dir="./data", crop=
         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
         )
+        flux_data, flux_error, flux_mask, flux_head = proj_red.rotate_data(np.array([flux_data]), np.array([flux_error]), flux_mask, [flux_head])
+        flux_data, flux_error, flux_head = flux_data[0], flux_error[0], flux_head[0]
+    elif not rotate_data:
+        figname += "_noROT"
 
     # 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)
@@ -240,7 +253,7 @@ def main(target=None, proposal_id=None, infiles=None, output_dir="./data", crop=
 
     # Step 4:
     # Save image to FITS.
-    figname = "_".join([figname, figtype]) if figtype != "" else figname
+    savename = "_".join([figname, figtype]) if figtype != "" else figname
     Stokes_hdul = proj_fits.save_Stokes(
         I_stokes,
         Q_stokes,
@@ -255,32 +268,29 @@ def main(target=None, proposal_id=None, infiles=None, output_dir="./data", crop=
         s_PA_P,
         header_stokes,
         data_mask,
-        figname,
+        savename,
         data_folder=data_folder,
         return_hdul=True,
+        flux_data=np.array([flux_data, flux_error, flux_mask]),
+        flux_head=flux_head,
     )
     outfiles.append("/".join([data_folder, Stokes_hdul[0].header["FILENAME"] + ".fits"]))
 
     # Step 5:
     # crop to desired region of interest (roi)
     if crop:
-        figname += "_crop"
+        savename += "_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
-        outfiles.append("/".join([data_folder, Stokes_hdul[0].header["FILENAME"] + ".fits"]))
+        outfiles.append("/".join([data_folder, Stokes_hdul["I_STOKES"].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"],
+            flux_head["PHOTPLAM"],
-            *sci_not(
+            *sci_not(flux_data[flux_mask].sum() * flux_head["PHOTFLAM"], np.sqrt(np.sum(flux_error[flux_mask] ** 2)) * flux_head["PHOTFLAM"], 2, out=int),
-                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))
@@ -303,12 +313,12 @@ def main(target=None, proposal_id=None, infiles=None, output_dir="./data", crop=
             flux_lim=flux_lim,
             step_vec=step_vec,
             scale_vec=scale_vec,
-            savename="_".join([figname]),
+            savename=figname,
             plots_folder=plots_folder,
         )
         for figtype, figsuffix in zip(
-            ["Intensity", "Pol_flux", "Pol_deg", "Pol_ang", "I_err", "P_err", "SNRi", "SNRp", "confp"],
+            ["FluxDensity", "Intensity", "Pol_flux", "Pol_deg", "Pol_ang", "I_err", "P_err", "SNRi", "SNRp", "confp"],
-            ["I", "P_flux", "P", "PA", "I_err", "P_err", "SNRi", "SNRp", "confP"],
+            ["F", "I", "P_flux", "P", "PA", "I_err", "P_err", "SNRi", "SNRp", "confP"],
         ):
             try:
                 proj_plots.polarization_map(
@@ -319,7 +329,7 @@ def main(target=None, proposal_id=None, infiles=None, output_dir="./data", crop=
                     flux_lim=flux_lim,
                     step_vec=step_vec,
                     scale_vec=scale_vec,
-                    savename="_".join([figname, figsuffix]),
+                    savename="_".join([savename, figsuffix]),
                     plots_folder=plots_folder,
                     display=figtype,
                 )
@@ -327,7 +337,7 @@ def main(target=None, proposal_id=None, infiles=None, output_dir="./data", crop=
                 pass
     elif not interactive:
         proj_plots.polarization_map(
-            deepcopy(Stokes_hdul), data_mask, P_cut=P_cut, SNRi_cut=SNRi_cut, savename=figname, plots_folder=plots_folder, display="integrate"
+            deepcopy(Stokes_hdul), data_mask, P_cut=P_cut, SNRi_cut=SNRi_cut, savename=savename, plots_folder=plots_folder, display="integrate"
         )
     elif pxscale.lower() not in ["full", "integrate"]:
         proj_plots.pol_map(Stokes_hdul, P_cut=P_cut, SNRi_cut=SNRi_cut, step_vec=step_vec, scale_vec=scale_vec, flux_lim=flux_lim)

package/__init__.py

@@ -1,2 +1,2 @@
 from .lib import *
-from . import FOC_reduction
+from .FOC_reduction import main

package/lib/background.py

@@ -18,7 +18,6 @@ import matplotlib.dates as mdates
 import matplotlib.pyplot as plt
 import numpy as np
 from astropy.time import Time
-from lib.plots import plot_obs
 from matplotlib.colors import LogNorm
 from matplotlib.patches import Rectangle
 from scipy.optimize import curve_fit
@@ -136,6 +135,8 @@ def display_bkg(data, background, std_bkg, headers, histograms=None, binning=Non
     fig2.subplots_adjust(hspace=0, wspace=0, right=1.0)
     fig2.colorbar(im2, ax=ax2, location="right", shrink=0.60, aspect=50, pad=0.025, label=r"Flux [$ergs \cdot cm^{-2} \cdot s^{-1} \cdot \AA^{-1}$]")
 
+    from .plots import plot_obs
+
     if savename is not None:
         this_savename = deepcopy(savename)
         if savename[-4:] not in [".png", ".jpg", ".pdf"]:
@@ -278,7 +279,7 @@ def bkg_fit(data, error, mask, headers, subtract_error=True, display=False, save
     return n_data_array, n_error_array, headers, background
 
 
-def bkg_hist(data, error, mask, headers, sub_type=None, subtract_error=True, display=False, savename=None, plots_folder=""):
+def bkg_hist(data, error, mask, headers, n_bins=None, subtract_error=True, display=False, savename=None, plots_folder=""):
     """
     ----------
     Inputs:
@@ -333,29 +334,15 @@ def bkg_hist(data, error, mask, headers, sub_type=None, subtract_error=True, dis
     for i, image in enumerate(data):
         # Compute the Count-rate histogram for the image
         n_mask = np.logical_and(mask, image > 0.0)
-        if sub_type is not None:
-            if isinstance(sub_type, int):
-                n_bins = sub_type
-            elif sub_type.lower() in ["square-root", "squareroot", "sqrt"]:
-                n_bins = np.fix(np.sqrt(image[n_mask].size)).astype(int)  # Square-root
-            elif sub_type.lower() in ["sturges"]:
-                n_bins = np.ceil(np.log2(image[n_mask].size)).astype(int) + 1  # Sturges
-            elif sub_type.lower() in ["rice"]:
-                n_bins = 2 * np.fix(np.power(image[n_mask].size, 1 / 3)).astype(int)  # Rice
-            elif sub_type.lower() in ["freedman-diaconis", "freedmandiaconis", "freedman", "diaconis"]:
-                n_bins = np.fix(
-                    (image[n_mask].max() - image[n_mask].min())
-                    / (2 * np.subtract(*np.percentile(image[n_mask], [75, 25])) / np.power(image[n_mask].size, 1 / 3))
-                ).astype(int)  # Freedman-Diaconis
-            else:  # Fallback
-                n_bins = np.fix((image[n_mask].max() - image[n_mask].min()) / (3.5 * image[n_mask].std() / np.power(image[n_mask].size, 1 / 3))).astype(
-                    int
-                )  # Scott
-        else:  # Default statistic
-            n_bins = np.fix((image[n_mask].max() - image[n_mask].min()) / (3.5 * image[n_mask].std() / np.power(image[n_mask].size, 1 / 3))).astype(
-                int
-            )  # Scott
 
+        if not isinstance(n_bins, int) and n_bins not in ["auto", "fd", "doane", "scott", "stone", "rice", "sturges", "sqrt"]:
+            match n_bins.lower():
+                case "square-root" | "squareroot":
+                    n_bins = "sqrt"
+                case "freedman-diaconis" | "freedmandiaconis":
+                    n_bins = "fd"
+                case _:
+                    n_bins = "scott"
         hist, bin_edges = np.histogram(np.log(image[n_mask]), bins=n_bins)
         histograms.append(hist)
         binning.append(np.exp(bin_centers(bin_edges)))

package/lib/fits.py

@@ -47,9 +47,9 @@ def get_obs_data(infiles, data_folder="", compute_flux=False):
             wcs_array.append(WCS(header=f[0].header, fobj=f).celestial)
             f.flush()
         # Save pixel area for flux density computation
-        if headers[i]["PXFORMT"] == "NORMAL":
+        if "PXFORMT" in headers[i].keys() and headers[i]["PXFORMT"] == "NORMAL":
             headers[i]["PXAREA"] = 1.96e-4  # 14x14 milliarcsec squared pixel area in arcsec^2
-        elif headers[i]["PXFORMT"] == "ZOOM":
+        elif "PXFORMT" in headers[i].keys() and headers[i]["PXFORMT"] == "ZOOM":
             headers[i]["PXAREA"] = 4.06e-4  # 29x14 milliarcsec squared pixel area in arcsec^2
         else:
             headers[i]["PXAREA"] = 1.0  # unknown default to 1 arcsec^2
@@ -90,10 +90,10 @@ def get_obs_data(infiles, data_folder="", compute_flux=False):
     # force WCS for POL60 to have same pixel size as POL0 and POL120
     is_pol60 = np.array([head["filtnam1"].lower() == "pol60" for head in headers], dtype=bool)
     cdelt = np.round(np.array([WCS(head).wcs.cdelt[:2] for head in headers]), 10)
-    if np.unique(cdelt[np.logical_not(is_pol60)], axis=0).size != 2:
+    if np.any(is_pol60) and np.unique(cdelt[np.logical_not(is_pol60)], axis=0).size != 2:
         print(np.unique(cdelt[np.logical_not(is_pol60)], axis=0))
         raise ValueError("Not all images have same pixel size")
-    else:
+    elif np.any(is_pol60):
         for i in np.arange(len(headers))[is_pol60]:
             headers[i]["cdelt1"], headers[i]["cdelt2"] = np.unique(cdelt[np.logical_not(is_pol60)], axis=0)[0]
 
@@ -106,7 +106,24 @@ def get_obs_data(infiles, data_folder="", compute_flux=False):
 
 
 def 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, filename, data_folder="", return_hdul=False
+    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,
+    filename,
+    data_folder="",
+    return_hdul=False,
+    flux_data=None,
+    flux_head=None,
 ):
     """
     Save computed polarimetry parameters to a single fits file,
@@ -162,7 +179,7 @@ def save_Stokes(
     header["PROPOSID"] = (header_stokes["PROPOSID"], "PEP proposal identifier for observation")
     header["TARGNAME"] = (header_stokes["TARGNAME"], "Target name")
     header["ORIENTAT"] = (header_stokes["ORIENTAT"], "Angle between North and the y-axis of the image")
-    header["FILENAME"] = (filename, "ORIGINAL FILENAME")
+    header["FILENAME"] = (filename, "Original filename")
     header["BKG_TYPE"] = (header_stokes["BKG_TYPE"], "Bkg estimation method used during reduction")
     header["BKG_SUB"] = (header_stokes["BKG_SUB"], "Amount of bkg subtracted from images")
     header["SMOOTH"] = (header_stokes["SMOOTH"] if "SMOOTH" in list(header_stokes.keys()) else "None", "Smoothing method used during reduction")
@@ -198,12 +215,30 @@ def save_Stokes(
     # Create HDUList object
     hdul = fits.HDUList([])
 
-    # Add I_stokes as PrimaryHDU
+    # Add Flux density as PrimaryHDU
-    header["datatype"] = ("I_stokes", "type of data stored in the HDU")
+    if flux_data is None:
+        header["DATATYPE"] = ("I_stokes", "type of data stored in the HDU")
         I_stokes[(1 - data_mask).astype(bool)] = 0.0
         primary_hdu = fits.PrimaryHDU(data=I_stokes, header=header)
         primary_hdu.name = "I_stokes"
         hdul.append(primary_hdu)
+    else:
+        flux_head["FILENAME"] = header["FILENAME"]
+        head = WCS(flux_head).deepcopy().to_header()
+        for key in [key for key in header.keys() if key not in ["SMOOTH", "SAMPLING"]]:
+            try:
+                head[key] = flux_head[key]
+            except KeyError:
+                head[key] = header[key]
+        header["DATATYPE"] = ("Flux_density", "type of data stored in the HDU")
+        primary_hdu = fits.PrimaryHDU(data=flux_data, header=head)
+        primary_hdu.name = "Flux_density"
+        hdul.append(primary_hdu)
+        header["DATATYPE"] = ("I_stokes", "type of data stored in the HDU")
+        I_stokes[(1 - data_mask).astype(bool)] = 0.0
+        image_hdu = fits.ImageHDU(data=I_stokes, header=header)
+        image_hdu.name = "I_stokes"
+        hdul.append(image_hdu)
 
     # Add Q, U, Stokes_cov, P, s_P, PA, s_PA to the HDUList
     for data, name in [
@@ -220,7 +255,7 @@ def save_Stokes(
         [data_mask, "Data_mask"],
     ]:
         hdu_header = header.copy()
-        hdu_header["datatype"] = name
+        hdu_header["DATATYPE"] = name
         if not name == "IQU_cov_matrix":
             data[(1 - data_mask).astype(bool)] = 0.0
         hdu = fits.ImageHDU(data=data, header=hdu_header)

package/lib/plots.py

@@ -133,6 +133,12 @@ def plot_obs(data_array, headers, rectangle=None, shifts=None, savename=None, pl
             ax_curr.arrow(x, y, dx, dy, length_includes_head=True, width=0.1, head_width=0.3, color="g")
             ax_curr.plot([x, x], [0, data.shape[0] - 1], "--", lw=2, color="g", alpha=0.85)
             ax_curr.plot([0, data.shape[1] - 1], [y, y], "--", lw=2, color="g", alpha=0.85)
+            # position of centroid
+            ax_curr.plot([data.shape[1] / 2, data.shape[1] / 2], [0, data.shape[0] - 1], "--", lw=2, color="b", alpha=0.85)
+            ax_curr.plot([0, data.shape[1] - 1], [data.shape[0] / 2, data.shape[0] / 2], "--", lw=2, color="b", alpha=0.85)
+        cr_x, cr_y = head["CRPIX1"], head["CRPIX2"]
+        # Plot WCS reference point
+        ax_curr.plot([cr_x], [cr_y], "+", lw=2, color="r", alpha=0.85)
         if rectangle is not None:
             x, y, width, height, angle, color = rectangle[i]
             ax_curr.add_patch(Rectangle((x, y), width, height, angle=angle, edgecolor=color, fill=False))
@@ -189,7 +195,7 @@ def plot_Stokes(Stokes, savename=None, plots_folder=""):
     for dataset in [stkI, stkQ, stkU]:
         dataset[np.logical_not(data_mask)] = np.nan
 
-    wcs = WCS(Stokes[0]).deepcopy()
+    wcs = WCS(Stokes["I_STOKES"]).deepcopy()
 
     # Plot figure
     plt.rcParams.update({"font.size": 14})
@@ -288,6 +294,9 @@ def polarization_map(
         The figure and ax created for interactive contour maps.
     """
     # Get data
+    flux = Stokes[0].data[0].copy() * Stokes[0].header["PHOTFLAM"]
+    flux_error = Stokes[0].data[1].copy() * Stokes[0].header["PHOTFLAM"]
+    flux_mask = Stokes[0].data[2].astype(bool).copy()
     stkI = Stokes["I_stokes"].data.copy()
     stkQ = Stokes["Q_stokes"].data.copy()
     stkU = Stokes["U_stokes"].data.copy()
@@ -302,6 +311,20 @@ def polarization_map(
             data_mask = np.zeros(stkI.shape).astype(bool)
             data_mask[stkI > 0.0] = True
 
+    wcs = WCS(Stokes["I_STOKES"]).deepcopy()
+    pivot_wav = Stokes["I_STOKES"].header["photplam"]
+    convert_flux = Stokes["I_STOKES"].header["photflam"]
+
+    # Get integrated flux values from sum
+    I_diluted = stkI[data_mask].sum() * convert_flux
+    I_diluted_err = np.sqrt(np.sum(stk_cov[0, 0][data_mask])) * convert_flux
+
+    # Get integrated polarization values from header
+    P_diluted = Stokes["I_STOKES"].header["P_int"]
+    P_diluted_err = Stokes["I_STOKES"].header["sP_int"]
+    PA_diluted = Stokes["I_STOKES"].header["PA_int"]
+    PA_diluted_err = Stokes["I_STOKES"].header["sPA_int"]
+
     # Compute confidence level map
     QN, UN, QN_ERR, UN_ERR = np.full((4, stkI.shape[0], stkI.shape[1]), np.nan)
     for nflux, sflux in zip([QN, UN, QN_ERR, UN_ERR], [stkQ, stkU, np.sqrt(stk_cov[1, 1]), np.sqrt(stk_cov[2, 2])]):
@@ -314,12 +337,8 @@ def polarization_map(
         for j in range(3):
             stk_cov[i][j][np.logical_not(data_mask)] = np.nan
 
-    wcs = WCS(Stokes[0]).deepcopy()
-    pivot_wav = Stokes[0].header["photplam"]
-    convert_flux = Stokes[0].header["photflam"]
-
     # Plot Stokes parameters map
-    if display is None or display.lower() in ["default"]:
+    if display is None or display.lower() in ["pol", "polarization", "polarisation", "pol_deg", "p"]:
         plot_Stokes(Stokes, savename=savename, plots_folder=plots_folder)
 
     # Compute SNR and apply cuts
@@ -363,7 +382,7 @@ def polarization_map(
         fig = plt.figure(figsize=(7 * ratiox, 7 * ratioy), layout="constrained")
     if ax is None:
         ax = fig.add_subplot(111, projection=wcs)
-        ax.set(aspect="equal", fc="k")  # , ylim=[-0.05 * stkI.shape[0], 1.05 * stkI.shape[0]])
+        ax.set(aspect="equal", fc="k", xlim=(0, stkI.shape[1]), ylim=(0, stkI.shape[0]))
     # fig.subplots_adjust(hspace=0, wspace=0, left=0.102, right=1.02)
 
     # ax.coords.grid(True, color='white', ls='dotted', alpha=0.5)
@@ -409,7 +428,25 @@ def polarization_map(
         ax.set_facecolor("white")
         font_color = "black"
 
-    if display.lower() in ["i", "intensity"]:
+    if display.lower() in ["f", "flux", "fluxdensity"]:
+        # If no display selected, show intensity map
+        display = "f"
+        if flux_lim is not None:
+            vmin, vmax = flux_lim
+        else:
+            vmin, vmax = np.max(flux[flux > 0.0]) / 2e3, np.max(flux[flux > 0.0])
+        imflux, cr = flux.copy(), WCS(Stokes[0].header).wcs.crpix.astype(int)
+        imflux[cr[1] - 1 : cr[1] + 2, cr[0] - 1 : cr[0] + 2] = np.nan
+        im = ax.imshow(
+            imflux, transform=ax.get_transform(WCS(Stokes[0].header).celestial), norm=LogNorm(vmin, vmax), aspect="equal", cmap=kwargs["cmap"], alpha=1.0
+        )
+        fig.colorbar(im, ax=ax, aspect=50, shrink=0.60, pad=0.015, fraction=0.03, label=r"$F_{\lambda}$ [$ergs \cdot cm^{-2} \cdot s^{-1} \cdot \AA$]")
+        levelsF = np.array([0.8, 2.0, 5.0, 10.0, 20.0, 50.0]) / 100.0 * vmax
+        print("Flux density contour levels : ", levelsF)
+        ax.contour(flux, levels=levelsF, transform=ax.get_transform(WCS(Stokes[0].header).celestial), colors="grey", linewidths=0.5)
+        ax.plot(*WCS(Stokes[1]).wcs.crpix, "g+")
+        I_diluted, I_diluted_err = np.sum(flux[flux_mask]), np.sqrt(np.sum(flux_error[flux_mask] ** 2))
+    elif display.lower() in ["i", "intensity"]:
         # If no display selected, show intensity map
         display = "i"
         if flux_lim is not None:
@@ -419,10 +456,13 @@ def polarization_map(
         else:
             vmin, vmax = 1.0 / 2.0 * np.median(np.sqrt(stk_cov[0, 0][stkI > 0.0]) * convert_flux), np.max(stkI[stkI > 0.0] * convert_flux)
         im = ax.imshow(stkI * convert_flux, norm=LogNorm(vmin, vmax), aspect="equal", cmap=kwargs["cmap"], alpha=1.0)
-        fig.colorbar(im, ax=ax, aspect=50, shrink=0.60, pad=0.025, label=r"$F_{\lambda}$ [$ergs \cdot cm^{-2} \cdot s^{-1} \cdot \AA^{-1}$]")
+        ax.plot(*WCS(Stokes[1]).wcs.crpix, "g+")
-        levelsI = np.array([0.8, 2.0, 5.0, 10.0, 20.0, 50.0]) / 100.0 * vmax
+        fig.colorbar(im, ax=ax, aspect=50, shrink=0.60, pad=0.015, fraction=0.03, label=r"$F_{\lambda}$ [$ergs \cdot cm^{-2} \cdot s^{-1} \cdot \AA^{-1}$]")
-        print("Flux density contour levels : ", levelsI)
+        # levelsI = np.array([0.8, 2.0, 5.0, 10.0, 20.0, 50.0]) / 100.0 * vmax
-        ax.contour(stkI * convert_flux, levels=levelsI, colors="grey", linewidths=0.5)
+        # print("Stokes I contour levels : ", levelsI)
+        # ax.contour(stkI * convert_flux, levels=levelsI, colors="grey", linewidths=0.5)
+        levelsF = np.array([0.8, 2.0, 5.0, 10.0, 20.0, 50.0]) / 100.0 * np.max(flux[flux > 0.0])
+        ax.contour(flux, levels=levelsF, transform=ax.get_transform(WCS(Stokes[0].header).celestial), colors="grey", linewidths=0.5)
     elif display.lower() in ["pf", "pol_flux"]:
         # Display polarization flux
         display = "pf"
@@ -436,7 +476,9 @@ def polarization_map(
             vmin, vmax = 1.0 / 2.0 * np.median(np.sqrt(stk_cov[0, 0][stkI > 0.0]) * convert_flux), np.max(stkI[stkI > 0.0] * convert_flux)
             pfmax = (stkI[stkI > 0.0] * pol[stkI > 0.0] * convert_flux).max()
         im = ax.imshow(stkI * convert_flux * pol, norm=LogNorm(vmin, vmax), aspect="equal", cmap=kwargs["cmap"], alpha=1.0)
-        fig.colorbar(im, ax=ax, aspect=50, shrink=0.60, pad=0.025, label=r"$F_{\lambda} \cdot P$ [$ergs \cdot cm^{-2} \cdot s^{-1} \cdot \AA^{-1}$]")
+        fig.colorbar(
+            im, ax=ax, aspect=50, shrink=0.60, pad=0.015, fraction=0.03, label=r"$F_{\lambda} \cdot P$ [$ergs \cdot cm^{-2} \cdot s^{-1} \cdot \AA^{-1}$]"
+        )
         # levelsPf = np.linspace(0.0.60, 0.50, 5) * pfmax
         levelsPf = np.array([1.73, 13.0, 33.0, 66.0]) / 100.0 * pfmax
         print("Polarized flux density contour levels : ", levelsPf)
@@ -446,13 +488,13 @@ def polarization_map(
         display = "p"
         vmin, vmax = 0.0, min(pol[np.isfinite(pol)].max(), 1.0) * 100.0
         im = ax.imshow(pol * 100.0, vmin=vmin, vmax=vmax, aspect="equal", cmap=kwargs["cmap"], alpha=1.0)
-        fig.colorbar(im, ax=ax, aspect=50, shrink=0.60, pad=0.025, label=r"$P$ [%]")
+        fig.colorbar(im, ax=ax, aspect=50, shrink=0.60, pad=0.015, fraction=0.03, label=r"$P$ [%]")
     elif display.lower() in ["pa", "pang", "pol_ang"]:
         # Display polarization degree map
         display = "pa"
         vmin, vmax = 0.0, 180.0
         im = ax.imshow(princ_angle(pang), vmin=vmin, vmax=vmax, aspect="equal", cmap=kwargs["cmap"], alpha=1.0)
-        fig.colorbar(im, ax=ax, aspect=50, shrink=0.60, pad=0.025, label=r"$\theta_P$ [°]")
+        fig.colorbar(im, ax=ax, aspect=50, shrink=0.60, pad=0.015, fraction=0.03, label=r"$\theta_P$ [°]")
     elif display.lower() in ["s_p", "pol_err", "pol_deg_err"]:
         # Display polarization degree error map
         display = "s_p"
@@ -462,7 +504,7 @@ def polarization_map(
         else:
             vmin, vmax = 0.0, 100.0
             im = ax.imshow(pol_err * 100.0, vmin=vmin, vmax=vmax, aspect="equal", cmap="inferno_r", alpha=1.0)
-        fig.colorbar(im, ax=ax, aspect=50, shrink=0.60, pad=0.025, label=r"$\sigma_P$ [%]")
+        fig.colorbar(im, ax=ax, aspect=50, shrink=0.60, pad=0.015, fraction=0.03, label=r"$\sigma_P$ [%]")
     elif display.lower() in ["s_i", "i_err"]:
         # Display intensity error map
         display = "s_i"
@@ -474,7 +516,7 @@ def polarization_map(
             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=kwargs["cmap"], alpha=1.0)
-        fig.colorbar(im, ax=ax, aspect=50, shrink=0.60, pad=0.025, label=r"$\sigma_I$ [$ergs \cdot cm^{-2} \cdot s^{-1} \cdot \AA^{-1}$]")
+        fig.colorbar(im, ax=ax, aspect=50, shrink=0.60, pad=0.015, fraction=0.03, label=r"$\sigma_I$ [$ergs \cdot cm^{-2} \cdot s^{-1} \cdot \AA^{-1}$]")
     elif display.lower() in ["snri"]:
         # Display I_stokes signal-to-noise map
         display = "snri"
@@ -486,7 +528,7 @@ def polarization_map(
             ax.contour(SNRi, levels=levelsSNRi, colors="grey", linewidths=0.5)
         else:
             im = ax.imshow(SNRi, aspect="equal", cmap=kwargs["cmap"], alpha=1.0)
-        fig.colorbar(im, ax=ax, aspect=50, shrink=0.60, pad=0.025, label=r"$I_{Stokes}/\sigma_{I}$")
+        fig.colorbar(im, ax=ax, aspect=50, shrink=0.60, pad=0.015, fraction=0.03, label=r"$I_{Stokes}/\sigma_{I}$")
     elif display.lower() in ["snr", "snrp"]:
         # Display polarization degree signal-to-noise map
         display = "snrp"
@@ -498,7 +540,7 @@ def polarization_map(
             ax.contour(SNRp, levels=levelsSNRp, colors="grey", linewidths=0.5)
         else:
             im = ax.imshow(SNRp, aspect="equal", cmap=kwargs["cmap"], alpha=1.0)
-        fig.colorbar(im, ax=ax, aspect=50, shrink=0.60, pad=0.025, label=r"$P/\sigma_{P}$")
+        fig.colorbar(im, ax=ax, aspect=50, shrink=0.60, pad=0.015, fraction=0.03, label=r"$P/\sigma_{P}$")
     elif display.lower() in ["conf", "confp"]:
         # Display polarization degree signal-to-noise map
         display = "confp"
@@ -507,27 +549,23 @@ def polarization_map(
         levelsconfp = np.array([0.500, 0.900, 0.990, 0.999])
         print("confp contour levels : ", levelsconfp)
         ax.contour(confP, levels=levelsconfp, colors="grey", linewidths=0.5)
-        fig.colorbar(im, ax=ax, aspect=50, shrink=0.60, pad=0.025, label=r"$Conf_{P}$")
+        fig.colorbar(im, ax=ax, aspect=50, shrink=0.60, pad=0.015, fraction=0.03, label=r"$Conf_{P}$")
     else:
         # Defaults to intensity map
         if flux_lim is not None:
             vmin, vmax = flux_lim
-        elif mask.sum() > 0.0:
-            vmin, vmax = 1.0 * np.mean(np.sqrt(stk_cov[0, 0][mask]) * convert_flux), np.max(stkI[stkI > 0.0] * convert_flux)
         else:
-            vmin, vmax = 1.0 * np.mean(np.sqrt(stk_cov[0, 0][stkI > 0.0]) * convert_flux), np.max(stkI[stkI > 0.0] * convert_flux)
+            vmin, vmax = np.max(flux[flux > 0.0] * convert_flux) / 2e3, np.max(flux[flux > 0.0] * convert_flux)
-        im = ax.imshow(stkI * convert_flux, norm=LogNorm(vmin, vmax), aspect="equal", cmap=kwargs["cmap"], alpha=1.0)
+        im = ax.imshow(
-        fig.colorbar(im, ax=ax, aspect=50, shrink=0.60, pad=0.025, label=r"$F_{\lambda}$ [$ergs \cdot cm^{-2} \cdot s^{-1} \cdot \AA$]")
+            flux * Stokes[0].header["PHOTFLAM"],
-
+            transform=ax.get_transform(WCS(Stokes[0].header).celestial),
-    # Get integrated flux values from sum
+            norm=LogNorm(vmin, vmax),
-    I_diluted = stkI[data_mask].sum()
+            aspect="equal",
-    I_diluted_err = np.sqrt(np.sum(stk_cov[0, 0][data_mask]))
+            cmap=kwargs["cmap"],
-
+            alpha=1.0,
-    # Get integrated polarization values from header
+        )
-    P_diluted = Stokes[0].header["P_int"]
+        fig.colorbar(im, ax=ax, aspect=50, shrink=0.60, pad=0.015, fraction=0.03, label=r"$F_{\lambda}$ [$ergs \cdot cm^{-2} \cdot s^{-1} \cdot \AA$]")
-    P_diluted_err = Stokes[0].header["sP_int"]
+        I_diluted, I_diluted_err = np.sum(flux[flux_mask]), np.sqrt(np.sum(flux_error[flux_mask] ** 2))
-    PA_diluted = Stokes[0].header["PA_int"]
-    PA_diluted_err = Stokes[0].header["sPA_int"]
 
     plt.rcParams.update({"font.size": 12})
     px_size = wcs.wcs.get_cdelt()[0] * 3600.0
@@ -545,12 +583,12 @@ def polarization_map(
         back_length=0.0,
         head_length=7.0,
         head_width=7.0,
-        angle=-Stokes[0].header["orientat"],
+        angle=-Stokes["I_STOKES"].header["orientat"],
         text_props={"ec": "k", "fc": font_color, "alpha": 1, "lw": 0.5},
         arrow_props={"ec": "k", "fc": font_color, "alpha": 1, "lw": 1},
     )
 
-    if display.lower() in ["i", "s_i", "snri", "pf", "p", "pa", "s_p", "snrp", "confp"] and step_vec != 0:
+    if display.lower() in ["f", "i", "s_i", "snri", "pf", "p", "pa", "s_p", "snrp", "confp"] and step_vec != 0:
         if scale_vec == -1:
             poldata[np.isfinite(poldata)] = 1.0 / 2.0
             step_vec = 1
@@ -1174,6 +1212,8 @@ class overplot_chandra(align_maps):
         other_data = deepcopy(self.other_data)
         other_wcs = self.other_wcs.deepcopy()
         if zoom != 1:
+            from scipy.ndimage import zoom as sc_zoom
+
             other_data = sc_zoom(other_data, zoom)
             other_wcs.wcs.crpix *= zoom
             other_wcs.wcs.cdelt /= zoom
@@ -2439,9 +2479,9 @@ class pol_map(object):
         self.conf = PCconf(self.QN, self.UN, self.QN_ERR, self.UN_ERR)
 
         # Get data
-        self.targ = self.Stokes[0].header["targname"]
+        self.targ = self.Stokes["I_STOKES"].header["targname"]
-        self.pivot_wav = self.Stokes[0].header["photplam"]
+        self.pivot_wav = self.Stokes["I_STOKES"].header["photplam"]
-        self.map_convert = self.Stokes[0].header["photflam"]
+        self.map_convert = self.Stokes["I_STOKES"].header["photflam"]
 
         # Create figure
         plt.rcParams.update({"font.size": 10})
@@ -2535,7 +2575,7 @@ class pol_map(object):
 
         def select_roi(event):
             if self.data is None:
-                self.data = self.Stokes[0].data
+                self.data = self.Stokes["I_STOKES"].data
             if self.selected:
                 self.selected = False
                 self.region = deepcopy(self.select_instance.mask.astype(bool))
@@ -2577,7 +2617,7 @@ class pol_map(object):
 
         def select_aperture(event):
             if self.data is None:
-                self.data = self.Stokes[0].data
+                self.data = self.Stokes["I_STOKES"].data
             if self.selected:
                 self.selected = False
                 self.select_instance.update_mask()
@@ -2634,7 +2674,7 @@ class pol_map(object):
 
         def select_slit(event):
             if self.data is None:
-                self.data = self.Stokes[0].data
+                self.data = self.Stokes["I_STOKES"].data
             if self.selected:
                 self.selected = False
                 self.select_instance.update_mask()
@@ -2911,7 +2951,31 @@ class pol_map(object):
 
     @property
     def wcs(self):
-        return WCS(self.Stokes[0].header).celestial.deepcopy()
+        return WCS(self.Stokes["I_STOKES"].header).celestial.deepcopy()
+
+    @property
+    def Flux(self):
+        return self.Stokes[0].data[0] * self.Stokes[0].header["PHOTFLAM"]
+
+    @property
+    def Flux_err(self):
+        return self.Stokes[0].data[1] * self.Stokes[0].header["PHOTFLAM"]
+
+    @property
+    def Flux_mask(self):
+        return self.Stokes[0].data[2].astype(bool)
+
+    @property
+    def Flux(self):
+        return self.Stokes[0].data[0] * self.Stokes[0].header["PHOTFLAM"]
+
+    @property
+    def Flux_err(self):
+        return self.Stokes[0].data[1] * self.Stokes[0].header["PHOTFLAM"]
+
+    @property
+    def Flux_mask(self):
+        return self.Stokes[0].data[2].astype(bool)
 
     @property
     def I(self):
@@ -2975,10 +3039,10 @@ class pol_map(object):
 
     @property
     def data_mask(self):
-        return self.Stokes["DATA_MASK"].data
+        return self.Stokes["DATA_MASK"].data.astype(bool)
 
     def set_data_mask(self, mask):
-        self.Stokes[np.argmax([self.Stokes[i].header["datatype"] == "Data_mask" for i in range(len(self.Stokes))])].data = mask.astype(float)
+        self.Stokes["DATA_MASK"].data = mask.astype(float)
 
     @property
     def cut(self):
@@ -3046,7 +3110,7 @@ class pol_map(object):
             back_length=0.0,
             head_length=7.5,
             head_width=7.5,
-            angle=-self.Stokes[0].header["orientat"],
+            angle=-self.Stokes["I_STOKES"].header["orientat"],
             color="white",
             text_props={"ec": None, "fc": "w", "alpha": 1, "lw": 0.4},
             arrow_props={"ec": None, "fc": "w", "alpha": 1, "lw": 1},
@@ -3054,18 +3118,20 @@ class pol_map(object):
         ax.add_artist(self.north_dir)
 
     def display(self, fig=None, ax=None, flux_lim=None):
-        norm = None
+        kwargs = dict([])
-        if self.display_selection is None:
-            self.display_selection = "total_flux"
         if flux_lim is None:
             flux_lim = self.flux_lim
-        if self.display_selection.lower() in ["total_flux"]:
+        if self.display_selection is None or self.display_selection.lower() in ["total_flux"]:
-            self.data = self.I * self.map_convert
+            self.data = self.Flux  # self.I * self.map_convert
             if flux_lim is None:
                 vmin, vmax = (1.0 / 2.0 * np.median(self.data[self.data > 0.0]), np.max(self.data[self.data > 0.0]))
             else:
                 vmin, vmax = flux_lim
-            norm = LogNorm(vmin, vmax)
+            kwargs["norm"] = LogNorm(vmin, vmax)
+            if ax is None:
+                kwargs["transform"] = self.ax.get_transform(WCS(self.Stokes[0].header).celestial)
+            else:
+                kwargs["transform"] = ax.get_transform(WCS(self.Stokes[0].header).celestial)
             label = r"$F_{\lambda}$ [$ergs \cdot cm^{-2} \cdot s^{-1} \cdot \AA^{-1}$]"
         elif self.display_selection.lower() in ["pol_flux"]:
             self.data = self.I * self.map_convert * self.P
@@ -3073,28 +3139,28 @@ class pol_map(object):
                 vmin, vmax = (1.0 / 2.0 * np.median(self.I[self.I > 0.0] * self.map_convert), np.max(self.I[self.I > 0.0] * self.map_convert))
             else:
                 vmin, vmax = flux_lim
-            norm = LogNorm(vmin, vmax)
+            kwargs["norm"] = LogNorm(vmin, vmax)
             label = r"$P \cdot F_{\lambda}$ [$ergs \cdot cm^{-2} \cdot s^{-1} \cdot \AA^{-1}$]"
         elif self.display_selection.lower() in ["pol_deg"]:
             self.data = self.P * 100.0
-            vmin, vmax = 0.0, np.max(self.data[self.P > self.s_P])
+            kwargs["vmin"], kwargs["vmax"] = 0.0, np.max(self.data[self.P > self.s_P])
             label = r"$P$ [%]"
         elif self.display_selection.lower() in ["pol_ang"]:
             self.data = princ_angle(self.PA)
-            vmin, vmax = 0, 180.0
+            kwargs["vmin"], kwargs["vmax"] = 0, 180.0
             label = r"$\theta_{P}$ [°]"
         elif self.display_selection.lower() in ["snri"]:
             s_I = np.sqrt(self.IQU_cov[0, 0])
             SNRi = np.zeros(self.I.shape)
             SNRi[s_I > 0.0] = self.I[s_I > 0.0] / s_I[s_I > 0.0]
             self.data = SNRi
-            vmin, vmax = 0.0, np.max(self.data[self.data > 0.0])
+            kwargs["vmin"], kwargs["vmax"] = 0.0, np.max(self.data[self.data > 0.0])
             label = r"$I_{Stokes}/\sigma_{I}$"
         elif self.display_selection.lower() in ["snrp"]:
             SNRp = np.zeros(self.P.shape)
             SNRp[self.s_P > 0.0] = self.P[self.s_P > 0.0] / self.s_P[self.s_P > 0.0]
             self.data = SNRp
-            vmin, vmax = 0.0, np.max(self.data[self.data > 0.0])
+            kwargs["vmin"], kwargs["vmax"] = 0.0, np.max(self.data[self.data > 0.0])
             label = r"$P/\sigma_{P}$"
 
         if fig is None:
@@ -3105,22 +3171,17 @@ class pol_map(object):
                 self.cbar.remove()
             if hasattr(self, "im"):
                 self.im.remove()
-            if norm is not None:
+            self.im = ax.imshow(self.data, aspect="equal", cmap="inferno", **kwargs)
-                self.im = ax.imshow(self.data, norm=norm, aspect="equal", cmap="inferno")
+            ax.set(xlim=(0, self.I.shape[1]), ylim=(0, self.I.shape[0]))
-            else:
-                self.im = ax.imshow(self.data, vmin=vmin, vmax=vmax, aspect="equal", cmap="inferno")
             plt.rcParams.update({"font.size": 14})
             self.cbar = fig.colorbar(self.im, ax=ax, aspect=50, shrink=0.75, pad=0.025, label=label)
             plt.rcParams.update({"font.size": 10})
             fig.canvas.draw_idle()
             return self.im
         else:
-            if norm is not None:
+            im = ax.imshow(self.data, aspect="equal", cmap="inferno", **kwargs)
-                im = ax.imshow(self.data, norm=norm, aspect="equal", cmap="inferno")
+            ax.set_xlim(0, self.I.shape[1])
-            else:
+            ax.set_ylim(0, self.I.shape[0])
-                im = ax.imshow(self.data, vmin=vmin, vmax=vmax, aspect="equal", cmap="inferno")
-            ax.set_xlim(0, self.data.shape[1])
-            ax.set_ylim(0, self.data.shape[0])
             plt.rcParams.update({"font.size": 14})
             fig.colorbar(im, ax=ax, aspect=50, shrink=0.75, pad=0.025, label=label)
             plt.rcParams.update({"font.size": 10})
@@ -3144,6 +3205,7 @@ class pol_map(object):
                 ax = self.ax
             if hasattr(self, "quiver"):
                 self.quiver.remove()
+            if self.pa_err:
                 self.quiver = ax.quiver(
                     X[:: self.step_vec, :: self.step_vec],
                     Y[:: self.step_vec, :: self.step_vec],
@@ -3156,12 +3218,11 @@ class pol_map(object):
                     headwidth=0.0,
                     headlength=0.0,
                     headaxislength=0.0,
-                width=0.3,
+                    width=0.1,
-                linewidth=0.6,
+                    # linewidth=0.6,
-                color="white",
+                    color="black",
                     edgecolor="black",
                 )
-            if self.pa_err:
                 XY_U_err1, XY_V_err1 = (
                     P_cut * np.cos(np.pi / 2.0 + (self.PA + 3.0 * self.s_PA) * np.pi / 180.0),
                     P_cut * np.sin(np.pi / 2.0 + (self.PA + 3.0 * self.s_PA) * np.pi / 180.0),
@@ -3210,6 +3271,25 @@ class pol_map(object):
                     edgecolor="black",
                     ls="dashed",
                 )
+            else:
+                self.quiver = ax.quiver(
+                    X[:: self.step_vec, :: self.step_vec],
+                    Y[:: self.step_vec, :: self.step_vec],
+                    XY_U[:: self.step_vec, :: self.step_vec],
+                    XY_V[:: self.step_vec, :: self.step_vec],
+                    units="xy",
+                    scale=1.0 / scale_vec,
+                    scale_units="xy",
+                    pivot="mid",
+                    headwidth=0.0,
+                    headlength=0.0,
+                    headaxislength=0.0,
+                    width=0.3,
+                    linewidth=0.6,
+                    color="white",
+                    edgecolor="black",
+                )
+
             fig.canvas.draw_idle()
             return self.quiver
         else:
@@ -3281,12 +3361,20 @@ class pol_map(object):
         str_conf = ""
         if self.region is None:
             s_I = np.sqrt(self.IQU_cov[0, 0])
-            I_reg = self.I.sum()
+            I_reg = (
-            I_reg_err = np.sqrt(np.sum(s_I**2))
+                np.sum(self.Flux[self.Flux_mask]) / self.map_convert
-            P_reg = self.Stokes[0].header["P_int"]
+                if self.display_selection is None or self.display_selection.lower() in ["total_flux"]
-            P_reg_err = self.Stokes[0].header["sP_int"]
+                else np.sum(self.I[self.data_mask])
-            PA_reg = self.Stokes[0].header["PA_int"]
+            )
-            PA_reg_err = self.Stokes[0].header["sPA_int"]
+            I_reg_err = (
+                np.sqrt(np.sum(self.Flux_err[self.Flux_mask] ** 2)) / self.map_convert
+                if self.display_selection is None or self.display_selection.lower() in ["total_flux"]
+                else np.sqrt(np.sum(s_I[self.data_mask] ** 2))
+            )
+            P_reg = self.Stokes["I_STOKES"].header["P_int"]
+            P_reg_err = self.Stokes["I_STOKES"].header["sP_int"]
+            PA_reg = self.Stokes["I_STOKES"].header["PA_int"]
+            PA_reg_err = self.Stokes["I_STOKES"].header["sPA_int"]
 
             s_I = np.sqrt(self.IQU_cov[0, 0])
             s_Q = np.sqrt(self.IQU_cov[1, 1])

package/lib/reduction.py

@@ -52,7 +52,6 @@ from scipy.ndimage import rotate as sc_rotate
 from scipy.ndimage import shift as sc_shift
 from scipy.signal import fftconvolve
 
-from .background import bkg_fit, bkg_hist, bkg_mini
 from .convex_hull import image_hull
 from .cross_correlation import phase_cross_correlation
 from .deconvolve import deconvolve_im, gaussian2d, gaussian_psf, zeropad
@@ -308,6 +307,8 @@ def crop_array(data_array, headers, error_array=None, data_mask=None, step=5, nu
         # Update CRPIX value in the associated header
         curr_wcs = WCS(crop_headers[i]).celestial.deepcopy()
         curr_wcs.wcs.crpix[:2] = curr_wcs.wcs.crpix[:2] - np.array([v_array[2], v_array[0]])
+        curr_wcs.array_shape = crop_array[i].shape
+        curr_wcs.wcs.set()
         crop_headers[i].update(curr_wcs.to_header())
         crop_headers[i]["naxis1"], crop_headers[i]["naxis2"] = crop_array[i].shape
 
@@ -523,20 +524,22 @@ def get_error(
     # estimated to less than 3%
     err_flat = data * 0.03
 
+    from .background import bkg_fit, bkg_hist, bkg_mini
+
     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
         )
         sub_type, subtract_error = "histogram ", str(int(subtract_error > 0.0))
     elif isinstance(sub_type, str):
-        if sub_type.lower() in ["auto"]:
+        if sub_type.lower() in ["fit"]:
             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
             )
             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
+                data, error, mask, headers, n_bins=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):
@@ -634,7 +637,8 @@ def rebin_array(data_array, error_array, headers, pxsize=2, scale="px", operatio
             pxsize, scale = "", "full"
         else:
             raise ValueError("'{0:s}' invalid scale for binning.".format(scale))
-    new_shape = np.ceil(min(image.shape / Dxy_arr, key=lambda x: x[0] + x[1])).astype(int)
+    new_shape_float = min(image.shape / Dxy_arr, key=lambda x: x[0] + x[1])
+    new_shape = np.ceil(new_shape_float).astype(int)
 
     for i, (image, error, header) in enumerate(list(zip(data_array, error_array, headers))):
         # Get current pixel size
@@ -663,8 +667,10 @@ def rebin_array(data_array, error_array, headers, pxsize=2, scale="px", operatio
         # Update header
         nw = w.deepcopy()
         nw.wcs.cdelt *= Dxy
+        # nw.wcs.crpix += np.abs(new_shape_float - new_shape) * np.array(new_shape) / Dxy
         nw.wcs.crpix /= Dxy
         nw.array_shape = new_shape
+        nw.wcs.set()
         new_header["NAXIS1"], new_header["NAXIS2"] = nw.array_shape
         new_header["PXAREA"] *= Dxy[0] * Dxy[1]
         for key, val in nw.to_header().items():
@@ -844,7 +850,10 @@ def align_data(
     new_crpix = np.array([wcs.wcs.crpix for wcs in headers_wcs]) + shifts[:, ::-1] + res_shift[::-1]
     for i in range(len(headers_wcs)):
         headers_wcs[i].wcs.crpix = new_crpix[0]
+        headers_wcs[i].array_shape = (res_shape, res_shape)
+        headers_wcs[i].wcs.set()
         headers[i].update(headers_wcs[i].to_header())
+        headers[i]["NAXIS1"], headers[i]["NAXIS2"] = res_shape, res_shape
 
     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)
@@ -1706,9 +1715,11 @@ def rotate_Stokes(I_stokes, Q_stokes, U_stokes, Stokes_cov, data_mask, header_st
 
     new_wcs.wcs.pc = np.dot(mrot, new_wcs.wcs.pc)
     new_wcs.wcs.crpix = np.dot(mrot, new_wcs.wcs.crpix - old_center[::-1]) + new_center[::-1]
+    new_wcs.array_shape = shape
     new_wcs.wcs.set()
     for key, val in new_wcs.to_header().items():
         new_header_stokes.set(key, val)
+    new_header_stokes["NAXIS1"], new_header_stokes["NAXIS2"] = new_wcs.array_shape
     new_header_stokes["ORIENTAT"] += ang
 
     # Nan handling :
@@ -1809,9 +1820,11 @@ def rotate_data(data_array, error_array, data_mask, headers):
         new_wcs = WCS(header).celestial.deepcopy()
         new_wcs.wcs.pc[:2, :2] = np.dot(mrot, new_wcs.wcs.pc[:2, :2])
         new_wcs.wcs.crpix[:2] = np.dot(mrot, new_wcs.wcs.crpix[:2] - old_center[::-1]) + new_center[::-1]
+        new_wcs.array_shape = shape
         new_wcs.wcs.set()
         for key, val in new_wcs.to_header().items():
             new_header[key] = val
+        new_header["NAXIS1"], new_header["NAXIS2"] = new_wcs.array_shape
         new_header["ORIENTAT"] = np.arccos(new_wcs.celestial.wcs.pc[0, 0]) * 180.0 / np.pi
         new_header["ROTATE"] = ang
         new_headers.append(new_header)