merge

Author: Jammy2211

autoarray/mask/abstract_mask.py

@@ -7,8 +7,9 @@
 
 if use_jax:
     import jax
-from pathlib import Path
-from typing import Dict, Union
+
+import numpy as np
+from typing import Dict
 
 from autoconf.fitsable import output_to_fits
 

autoarray/mask/mask_1d.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from enum import Enum
 import logging
 import numpy as np
 from pathlib import Path
@@ -20,6 +21,10 @@
 logger = logging.getLogger(__name__)
 
 
+class Mask1DKeys(Enum):
+    PIXSCA = "PIXSCA"
+    ORIGIN = "ORIGIN"
+
 class Mask1D(Mask):
     def __init__(
         self,
@@ -172,7 +177,8 @@ def header_dict(self) -> Dict:
         -------
         A dictionary containing the pixel scale of the mask, which can be output to a .fits file.
         """
+
         return {
-            "PIXSCA": self.pixel_scales[0],
-            "ORIGIN": self.origin[0],
+            Mask1DKeys.PIXSCA: self.pixel_scales[0],
+            Mask1DKeys.ORIGIN: self.origin[0],
         }

autoarray/mask/mask_2d.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 import logging
+from enum import Enum
 import numpy as np
 from pathlib import Path
 from typing import TYPE_CHECKING, Dict, List, Tuple, Union
@@ -30,6 +31,13 @@
 logger = logging.getLogger(__name__)
 
 
+class Mask2DKeys(Enum):
+    PIXSCAY = "PIXSCAY"
+    PIXSCAX = "PIXSCAX"
+    ORIGINY = "ORIGINY"
+    ORIGINX = "ORIGINX"
+
+
 class Mask2D(Mask):
     # noinspection PyUnusedLocal
     def __init__(
@@ -657,11 +665,12 @@ def header_dict(self) -> Dict:
         -------
         A dictionary containing the pixel scale of the mask, which can be output to a .fits file.
         """
+
         return {
-            "PIXSCAY": self.pixel_scales[0],
-            "PIXSCAX": self.pixel_scales[1],
-            "ORIGINY": self.origin[0],
-            "ORIGINX": self.origin[1],
+            Mask2DKeys.PIXSCAY: self.pixel_scales[0],
+            Mask2DKeys.PIXSCAX: self.pixel_scales[1],
+            Mask2DKeys.ORIGINY: self.origin[0],
+            Mask2DKeys.ORIGINX: self.origin[1],
         }
 
     @property