Added method which reshapes a mask when loaded from a fits file.

Author: Jammy2211

autoarray/mask/mask.py

@@ -373,7 +373,7 @@ def elliptical_annular(
         )
 
     @classmethod
-    def from_fits(cls, file_path, pixel_scales, hdu=0, sub_size=1, origin=(0.0, 0.0)):
+    def from_fits(cls, file_path, pixel_scales, hdu=0, sub_size=1, origin=(0.0, 0.0), resized_mask_shape=None):
         """
         Loads the image from a .fits file.
 
@@ -391,13 +391,19 @@ def from_fits(cls, file_path, pixel_scales, hdu=0, sub_size=1, origin=(0.0, 0.0)
             if type(pixel_scales) is float or int:
                 pixel_scales = (float(pixel_scales), float(pixel_scales))
 
-        return cls(
+        mask = cls(
             array_util.numpy_array_2d_from_fits(file_path=file_path, hdu=hdu),
             pixel_scales=pixel_scales,
             sub_size=sub_size,
             origin=origin,
         )
 
+        if resized_mask_shape is not None:
+
+            mask = mask.mapping.resized_mask_from_new_shape(new_shape=resized_mask_shape)
+
+        return mask
+
     def output_to_fits(self, file_path, overwrite=False):
 
         array_util.numpy_array_2d_to_fits(

test_autoarray/unit/mask/test_mask.py

@@ -427,6 +427,8 @@ def test__mask_elliptical_annular_inverted__compare_to_array_util(self):
 
 
 class TestFromAndToFits:
+
+
     def test__load_and_output_mask_to_fits(self):
 
         mask = msk.Mask.from_fits(
@@ -459,6 +461,27 @@ def test__load_and_output_mask_to_fits(self):
         assert mask.pixel_scales == (1.0, 1.0)
         assert mask.origin == (2.0, 2.0)
 
+    def test__load_from_fits_with_resized_mask_shape(self):
+
+        mask = msk.Mask.from_fits(
+            file_path=test_data_dir + "3x3_ones.fits",
+            hdu=0,
+            sub_size=1,
+            pixel_scales=(1.0, 1.0),
+            resized_mask_shape=(1, 1)
+        )
+
+        assert mask.shape_2d == (1, 1)
+
+        mask = msk.Mask.from_fits(
+            file_path=test_data_dir + "3x3_ones.fits",
+            hdu=0,
+            sub_size=1,
+            pixel_scales=(1.0, 1.0),
+            resized_mask_shape=(5, 5)
+        )
+
+        assert mask.shape_2d == (5, 5)
 
 class TestSubQuantities:
     def test__sub_shape_is_shape_times_sub_size(self):