diff --git a/package/lib/fits.py b/package/lib/fits.py
index 1c4ac2c..3604588 100755
--- a/package/lib/fits.py
+++ b/package/lib/fits.py
@@ -217,7 +217,7 @@ def save_Stokes(
 
     # Add Flux density as PrimaryHDU
     if flux_data is None:
-        header["datatype"] = ("I_stokes", "type of data stored in the HDU")
+        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"
@@ -230,11 +230,11 @@ def save_Stokes(
                 head[key] = flux_head[key]
             except KeyError:
                 head[key] = header[key]
-        header["datatype"] = ("Flux_density", "type of data stored in the HDU")
+        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")
+        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"
@@ -255,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)
diff --git a/package/lib/plots.py b/package/lib/plots.py
index e8b7588..061dcf7 100755
--- a/package/lib/plots.py
+++ b/package/lib/plots.py
@@ -2951,7 +2951,7 @@ 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):
@@ -3042,7 +3042,7 @@ class pol_map(object):
         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):