HDU -> FITS

.github/workflows/main.yml

@@ -22,7 +22,7 @@ jobs:
       uses: actions/setup-python@v2
       with:
         python-version: ${{ matrix.python-version }}
-    - uses: actions/cache@v2
+    - uses: actions/cache@v3
       id: cache-pip
       with:
         path: ~/.cache/pip

autofit/__init__.py

@@ -25,8 +25,9 @@
 from .non_linear.samples import load_from_table
 from .non_linear.samples import SamplesStored
 from .database.aggregator import Aggregator
-from .aggregator.aggregate_csv import AggregateCSV
-from .aggregator.aggregate_images import AggregateImages
+from .aggregator.summary.aggregate_csv import AggregateCSV
+from .aggregator.summary.aggregate_images import AggregateImages
+from .aggregator.summary.aggregate_fits import AggregateFITS, FitFITS
 from .database.aggregator import Query
 from autofit.aggregator.fit_interface import Fit
 from .aggregator.search_output import SearchOutput

autofit/aggregator/file_output.py

@@ -20,7 +20,7 @@ def __new__(cls, name, path: Path):
         elif suffix == ".csv":
             return super().__new__(ArrayOutput)
         elif suffix == ".fits":
-            return super().__new__(HDUOutput)
+            return super().__new__(FITSOutput)
         raise ValueError(f"File {path} is not a valid output file")
 
     def __init__(self, name: str, path: Path):
@@ -92,26 +92,12 @@ def value(self):
             return dill.load(f)
 
 
-class HDUOutput(FileOutput):
-    def __init__(self, name: str, path: Path):
-        super().__init__(name, path)
-        self._file = None
-
-    @property
-    def file(self):
-        if self._file is None:
-            self._file = open(self.path, "rb")
-        return self._file
-
+class FITSOutput(FileOutput):
     @property
     def value(self):
         """
         The contents of the fits file
         """
         from astropy.io import fits
 
-        return fits.PrimaryHDU.readfrom(self.file)
-
-    def __del__(self):
-        if self._file is not None:
-            self._file.close()
+        return fits.open(self.path)

autofit/aggregator/search_output.py

@@ -78,11 +78,15 @@ def files_path(self):
         return self.directory / "files"
 
     def _outputs(self, suffix):
+        return self._outputs_in_directory("files", suffix) + self._outputs_in_directory(
+            "image", suffix
+        )
+
+    def _outputs_in_directory(self, name: str, suffix: str):
+        files_path = self.directory / name
         outputs = []
-        for file_path in self.files_path.rglob(f"*{suffix}"):
-            name = ".".join(
-                file_path.relative_to(self.files_path).with_suffix("").parts
-            )
+        for file_path in files_path.rglob(f"*{suffix}"):
+            name = ".".join(file_path.relative_to(files_path).with_suffix("").parts)
             outputs.append(FileOutput(name, file_path))
         return outputs
 
@@ -108,7 +112,7 @@ def pickles(self):
         return self._outputs(".pickle")
 
     @cached_property
-    def hdus(self):
+    def fits(self):
         """
         The fits files in the search output files directory
         """
@@ -170,7 +174,7 @@ def value(self, name: str):
         for item in self.jsons:
             if item.name == name:
                 return item.value_using_reference(self._reference)
-        for item in self.pickles + self.arrays + self.hdus:
+        for item in self.pickles + self.arrays + self.fits:
             if item.name == name:
                 return item.value
 

autofit/aggregator/aggregate_csv.py → autofit/aggregator/summary/aggregate_csv.py

@@ -190,6 +190,9 @@ def __init__(self, aggregator: Aggregator):
         ----------
         aggregator
         """
+        if len(aggregator) == 0:
+            raise ValueError("The aggregator is empty.")
+
         self._aggregator = aggregator
         self._columns = []
 

autofit/aggregator/summary/aggregate_fits.py

@@ -0,0 +1,140 @@
+import re
+from enum import Enum
+from typing import List, Union
+
+from astropy.io import fits
+from pathlib import Path
+
+from autofit.aggregator.search_output import SearchOutput
+from autofit.aggregator import Aggregator
+
+
+def subplot_filename(subplot: Enum) -> str:
+    subplot_type = subplot.__class__
+    return (
+        re.sub(
+            r"([A-Z])",
+            r"_\1",
+            subplot_type.__name__.replace("FITS", ""),
+        )
+        .lower()
+        .lstrip("_")
+    )
+
+
+class FitFITS(Enum):
+    """
+    The HDUs that can be extracted from the fit.fits file.
+    """
+
+    ModelImage = "MODEL_IMAGE"
+    ResidualMap = "RESIDUAL_MAP"
+    NormalizedResidualMap = "NORMALIZED_RESIDUAL_MAP"
+    ChiSquaredMap = "CHI_SQUARED_MAP"
+
+
+class AggregateFITS:
+    def __init__(self, aggregator: Aggregator):
+        """
+        A class for extracting fits files from the aggregator.
+
+        Parameters
+        ----------
+        aggregator
+            The aggregator containing the fits files.
+        """
+        if len(aggregator) == 0:
+            raise ValueError("The aggregator is empty.")
+
+        self.aggregator = aggregator
+
+    @staticmethod
+    def _hdus(
+        result: SearchOutput,
+        *hdus: Enum,
+    ) -> List[fits.ImageHDU]:
+        """
+        Extract the HDUs from a given fits for a given search.
+
+        Parameters
+        ----------
+        result
+            The search output.
+        hdus
+            The HDUs to extract.
+
+        Returns
+        -------
+        The extracted HDUs.
+        """
+        row = []
+        for hdu in hdus:
+            source = result.value(subplot_filename(hdu))
+            source_hdu = source[source.index_of(hdu.value)]
+            row.append(
+                fits.ImageHDU(
+                    data=source_hdu.data,
+                    header=source_hdu.header,
+                )
+            )
+        return row
+
+    def extract_fits(self, *hdus: Enum) -> List[fits.HDUList]:
+        """
+        Extract the HDUs from the fits files for every search in the aggregator.
+
+        Return the result as a list of HDULists. The first HDU in each list is an empty PrimaryHDU.
+
+        Parameters
+        ----------
+        hdus
+            The HDUs to extract.
+
+        Returns
+        -------
+        The extracted HDUs.
+        """
+        output = [fits.PrimaryHDU()]
+        for result in self.aggregator:
+            output.extend(self._hdus(result, *hdus))
+
+        return fits.HDUList(output)
+
+    def output_to_folder(
+        self,
+        folder: Path,
+        *hdus: Enum,
+        name: Union[str, List[str]],
+    ):
+        """
+        Output the fits files for every search in the aggregator to a folder.
+
+        Only include HDUs specific in the hdus argument.
+
+        Parameters
+        ----------
+        folder
+            The folder to output the fits files to.
+        hdus
+            The HDUs to output.
+        name
+            The name of the fits file. This is the attribute of the search output that is used to name the file.
+            OR a list of names for each HDU.
+        """
+        folder.mkdir(parents=True, exist_ok=True)
+
+        for i, result in enumerate(self.aggregator):
+            if isinstance(name, str):
+                output_name = getattr(result, name)
+            else:
+                output_name = name[i]
+
+            hdu_list = fits.HDUList(
+                [fits.PrimaryHDU()]
+                + self._hdus(
+                    result,
+                    *hdus,
+                )
+            )
+            with open(folder / f"{output_name}.fits", "wb") as file:
+                hdu_list.writeto(file)

autofit/aggregator/aggregate_images.py → autofit/aggregator/summary/aggregate_images.py

@@ -1,6 +1,4 @@
-import re
 import sys
-from enum import Enum
 from typing import Optional, List, Union, Callable, Type
 from pathlib import Path
 
@@ -9,6 +7,22 @@
 from autofit.aggregator.search_output import SearchOutput
 from autofit.aggregator.aggregator import Aggregator
 
+import re
+from enum import Enum
+
+
+def subplot_filename(subplot: Enum) -> str:
+    subplot_type = subplot.__class__
+    return (
+        re.sub(
+            r"([A-Z])",
+            r"_\1",
+            subplot_type.__name__,
+        )
+        .lower()
+        .lstrip("_")
+    )
+
 
 class SubplotFit(Enum):
     """
@@ -100,6 +114,9 @@ def __init__(
         aggregator
             The aggregator containing the fit results.
         """
+        if len(aggregator) == 0:
+            raise ValueError("The aggregator is empty.")
+
         self._aggregator = aggregator
         self._source_images = None
 
@@ -149,7 +166,7 @@ def output_to_folder(
         folder: Path,
         *subplots: Union[SubplotFit, List[Image.Image], Callable],
         subplot_width: Optional[int] = sys.maxsize,
-        name: str = "name",
+        name: Union[str, List[str]],
     ):
         """
         Output one subplot image for each fit in the aggregator.
@@ -171,8 +188,9 @@ def output_to_folder(
             images to wrap.
         name
             The attribute of each fit to use as the name of the output file.
+            OR a list of names, one for each fit.
         """
-        folder.mkdir(exist_ok=True)
+        folder.mkdir(exist_ok=True, parents=True)
 
         for i, result in enumerate(self._aggregator):
             image = self._matrix_to_image(
@@ -183,7 +201,13 @@ def output_to_folder(
                     subplot_width=subplot_width,
                 )
             )
-            image.save(folder / f"{getattr(result, name)}.png")
+
+            if isinstance(name, str):
+                output_name = getattr(result, name)
+            else:
+                output_name = name[i]
+
+            image.save(folder / f"{output_name}.png")
 
     @staticmethod
     def _matrix_for_result(
@@ -231,30 +255,30 @@ class name but using snake_case.
             The image for the subplot.
             """
             subplot_type = subplot_.__class__
-            name = (
-                re.sub(
-                    r"([A-Z])",
-                    r"_\1",
-                    subplot_type.__name__,
-                )
-                .lower()
-                .lstrip("_")
-            )
-
             if subplot_type not in _images:
-                _images[subplot_type] = SubplotFitImage(result.image(name))
+                _images[subplot_type] = SubplotFitImage(
+                    result.image(
+                        subplot_filename(subplot_),
+                    )
+                )
             return _images[subplot_type]
 
         matrix = []
         row = []
         for subplot in subplots:
-            if isinstance(subplot, SubplotFit):
+            if isinstance(subplot, Enum):
                 row.append(
                     get_image(subplot).image_at_coordinates(
                         *subplot.value,
                     )
                 )
             elif isinstance(subplot, list):
+                if not isinstance(subplot[i], Image.Image):
+                    raise TypeError(
+                        "The subplots must be of type Subplot or a list of "
+                        "images or a function that takes a SearchOutput as an "
+                        "argument."
+                    )
                 row.append(subplot[i])
             else:
                 try:

autofit/database/aggregator/scrape.py

@@ -183,5 +183,5 @@ def _add_files(fit: m.Fit, item: SearchOutput):
         except ValueError:
             logger.debug(f"Failed to load array {array_output.name} for {fit.id}")
 
-    for hdu_output in item.hdus:
-        fit.set_hdu(hdu_output.name, hdu_output.value)
+    for fits in item.fits:
+        fit.set_fits(fits.name, fits.value)

autofit/database/migration/steps.py

@@ -23,6 +23,15 @@
     Step(
         "ALTER TABLE object RENAME COLUMN latent_variables_for_id TO latent_samples_for_id;",
     ),
+    Step(
+        "CREATE TABLE fits (id INTEGER NOT NULL, name VARCHAR, fit_id VARCHAR, PRIMARY KEY (id), FOREIGN KEY (fit_id) REFERENCES fit (id));"
+    ),
+    Step(
+        "ALTER TABLE hdu ADD COLUMN fits_id INTEGER;",
+    ),
+    Step(
+        "ALTER TABLE hdu ADD COLUMN is_primary BOOLEAN;",
+    ),
 ]
 
 migrator = Migrator(*steps)