Repository: rhayes777/PyAutoFit
Pull request #862 (2023-10-30 10:32:53)

feature/scrape directory agg
Another step in aggregator unification. Moves functionality from scrape.py to the classic aggregator search output bringing its API into line with the database Fit object



Patch:
diff --git a/autofit/aggregator/aggregator.py b/autofit/aggregator/aggregator.py
index a4f25a93b..82f2d37d3 100755
--- a/autofit/aggregator/aggregator.py
+++ b/autofit/aggregator/aggregator.py
@@ -21,7 +21,7 @@ from shutil import rmtree
 from typing import List, Union, Iterator, Optional
 
 from .predicate import AttributePredicate
-from .search_output import SearchOutput
+from .search_output import SearchOutput, GridSearchOutput, GridSearch
 
 
 class AggregatorGroup:
@@ -76,10 +76,36 @@ class AggregatorGroup:
         return [group.values(name, parser=parser) for group in self.groups]
 
 
+def unzip_directory(directory: str):
+    """
+    Unzip all zip files in a directory recursively.
+    """
+    for root, _, filenames in os.walk(directory):
+        for filename in filenames:
+            if filename.endswith(".zip"):
+                try:
+                    with zipfile.ZipFile(path.join(root, filename), "r") as f:
+                        f.extractall(path.join(root, filename[:-4]))
+                except zipfile.BadZipFile:
+                    raise zipfile.BadZipFile(
+                        f"File is not a zip file: \n " f"{root} \n" f"{filename}"
+                    )
+
+
+def is_relative_to(path_a, path_b):
+    """Return True if the path is relative to another path or False."""
+    try:
+        path_a.relative_to(path_b)
+        return True
+    except ValueError:
+        return False
+
+
 class Aggregator:
     def __init__(
         self,
         search_outputs: List[SearchOutput],
+        grid_search_outputs: List[GridSearchOutput],
     ):
         """
         Class to aggregate phase results for all subdirectories in a given directory.
@@ -90,6 +116,23 @@ class Aggregator:
             A list of search_outputs
         """
         self.search_outputs = search_outputs
+        self.grid_search_outputs = grid_search_outputs
+
+    def grid_searches(self):
+        """
+        A list of grid search outputs
+        """
+        return [
+            GridSearch(
+                output,
+                [
+                    search_output
+                    for search_output in self.search_outputs
+                    if is_relative_to(search_output.directory, output.directory)
+                ],
+            )
+            for output in self.grid_search_outputs
+        ]
 
     @classmethod
     def from_directory(
@@ -116,28 +159,31 @@ class Aggregator:
         """
         print("Aggregator loading search_outputs... could take some time.")
 
+        unzip_directory(directory)
+
         search_outputs = []
+        grid_search_outputs = []
 
         for root, _, filenames in os.walk(directory):
-            for filename in filenames:
-                if filename.endswith(".zip"):
-                    try:
-                        with zipfile.ZipFile(path.join(root, filename), "r") as f:
-                            f.extractall(path.join(root, filename[:-4]))
-                    except zipfile.BadZipFile:
-                        raise zipfile.BadZipFile(
-                            f"File is not a zip file: \n " f"{root} \n" f"{filename}"
-                        )
 
-        for root, _, filenames in os.walk(directory):
+            def should_add():
+                return not completed_only or ".completed" in filenames
+
             if "metadata" in filenames:
-                if not completed_only or ".completed" in filenames:
+                if should_add():
                     search_outputs.append(
                         SearchOutput(
                             Path(root),
                             reference=reference,
                         )
                     )
+            if ".is_grid_search" in filenames:
+                if should_add():
+                    grid_search_outputs.append(
+                        GridSearchOutput(
+                            Path(root),
+                        )
+                    )
 
         if len(search_outputs) == 0:
             print(f"\nNo search_outputs found in {directory}\n")
@@ -146,7 +192,7 @@ class Aggregator:
                 f"\n A total of {str(len(search_outputs))} search_outputs and results were found."
             )
 
-        return cls(search_outputs)
+        return cls(search_outputs, grid_search_outputs)
 
     def remove_unzipped(self):
         """
@@ -173,7 +219,10 @@ class Aggregator:
         An aggregator or phase
         """
         if isinstance(item, slice):
-            return Aggregator(self.search_outputs[item])
+            return Aggregator(
+                self.search_outputs[item],
+                self.grid_search_outputs,
+            )
         return self.search_outputs[item]
 
     def __len__(self):
@@ -207,7 +256,10 @@ class Aggregator:
             search_outputs = predicate.filter(search_outputs)
         search_outputs = list(search_outputs)
         print(f"Filter found a total of {str(len(search_outputs))} results")
-        return Aggregator(search_outputs=list(search_outputs))
+        return Aggregator(
+            search_outputs=list(search_outputs),
+            grid_search_outputs=self.grid_search_outputs,
+        )
 
     def values(self, name: str, parser=lambda o: o) -> Iterator:
         """
diff --git a/autofit/aggregator/file_output.py b/autofit/aggregator/file_output.py
new file mode 100644
index 000000000..38c25753e
--- /dev/null
+++ b/autofit/aggregator/file_output.py
@@ -0,0 +1,92 @@
+import json
+from abc import abstractmethod, ABC
+from pathlib import Path
+from typing import Union
+
+import numpy as np
+
+import dill
+
+from autoconf.dictable import from_dict
+
+
+class FileOutput(ABC):
+    def __new__(cls, name, path: Path):
+        suffix = path.suffix
+        if suffix == ".pickle":
+            return super().__new__(PickleOutput)
+        elif suffix == ".json":
+            return super().__new__(JSONOutput)
+        elif suffix == ".csv":
+            return super().__new__(ArrayOutput)
+        elif suffix == ".fits":
+            return super().__new__(HDUOutput)
+        raise ValueError(f"File {path} is not a valid output file")
+
+    def __init__(self, name: str, path: Path):
+        """
+        An output file from a fit, which is either a pickle, json, csv or fits file.
+
+        Parameters
+        ----------
+        name
+            The name of the output file. This is the path to the file from the root output directory
+            separated by '.' and without a suffix. e.g. output/directory/info.json -> directory.info
+        path
+            The path to the output file
+        """
+        self.name = name
+        self.path = path
+
+    @property
+    @abstractmethod
+    def value(self):
+        pass
+
+
+class ArrayOutput(FileOutput):
+    @property
+    def value(self) -> np.ndarray:
+        """
+        The array stored in the csv file
+        """
+        return np.loadtxt(self.path, delimiter=",")
+
+
+class JSONOutput(FileOutput):
+    @property
+    def dict(self) -> Union[dict, list]:
+        """
+        The dictionary stored in the json file
+        """
+        with open(self.path) as f:
+            return json.load(f)
+
+    @property
+    def value(self):
+        """
+        The object represented by the JSON
+        """
+        return from_dict(self.dict)
+
+
+class PickleOutput(FileOutput):
+    @property
+    def value(self):
+        """
+        The object stored in the pickle file
+        """
+        with open(self.path, "rb") as f:
+            return dill.load(f)
+
+
+class HDUOutput(FileOutput):
+    @property
+    def value(self):
+        """
+        The contents of the fits file
+        """
+        from astropy.io import fits
+
+        with open(self.path, "rb") as f:
+            return fits.PrimaryHDU.readfrom(f)
diff --git a/autofit/aggregator/search_output.py b/autofit/aggregator/search_output.py
index 30f59ea26..c18a17b4a 100644
--- a/autofit/aggregator/search_output.py
+++ b/autofit/aggregator/search_output.py
@@ -4,16 +4,18 @@ import logging
 import pickle
 from os import path
 from pathlib import Path
-from typing import Generator, Tuple, Optional
+from typing import Generator, Tuple, Optional, List
 
 import dill
 import numpy as np
 
-from autofit import SamplesPDF
+from autofit.non_linear.samples.pdf import SamplesPDF
+from autofit.aggregator.file_output import (
+    JSONOutput,
+    FileOutput,
+)
 from autofit.mapper.identifier import Identifier
 from autofit.non_linear.samples.sample import samples_from_iterator
-from autofit.non_linear.search import abstract_search
-from autofit.mapper.prior_model.abstract import AbstractPriorModel
 from autoconf.dictable import from_dict
 
 original_create_file_handle = dill._dill._create_filehandle
@@ -38,55 +40,77 @@ def _create_file_handle(*args, **kwargs):
 dill._dill._create_filehandle = _create_file_handle
 
 
-class SearchOutput:
-    """
-    @DynamicAttrs
-    """
+class AbstractSearchOutput:
+    def __init__(self, directory: Path):
+        self.directory = directory
 
-    def __new__(cls, directory: Path, reference: dict = None):
+    @property
+    def is_complete(self) -> bool:
         """
-        Represents the output of a single search. Comprises a metadata file and other dataset files.
-
-        Parameters
-        ----------
-        directory
-            The directory of the search
+        Whether the search has completed
         """
-        if (directory / ".is_grid_search").exists():
-            return super().__new__(GridSearchOutput)
-        return super().__new__(cls)
+        return (self.directory / ".completed").exists()
 
-    def __init__(self, directory: Path, reference: dict = None):
+    @property
+    def parent_identifier(self) -> Optional[str]:
         """
-        Represents the output of a single search. Comprises a metadata file and other dataset files.
+        Read the parent identifier for a fit in a directory.
 
-        Parameters
-        ----------
-        directory
-            The directory of the search
+        Defaults to None if no .parent_identifier file is found.
         """
-        self.__search = None
-        self.__model = None
-
-        self.directory = directory
-
-        self._reference = reference
-        self.file_path = directory / "metadata"
-
         try:
-            with open(self.file_path) as f:
-                self.text = f.read()
-                pairs = [
-                    line.split("=") for line in self.text.split("\n") if "=" in line
-                ]
-                self.__dict__.update({pair[0]: pair[1] for pair in pairs})
+            return (self.directory / ".parent_identifier").read_text()
         except FileNotFoundError:
-            pass
+            return None
 
     @property
     def files_path(self):
         return self.directory / "files"
 
+    def _outputs(self, suffix):
+        outputs = []
+        for file_path in self.files_path.rglob(f"*{suffix}"):
+            name = ".".join(
+                file_path.relative_to(self.files_path).with_suffix("").parts
+            )
+            outputs.append(FileOutput(name, file_path))
+        return outputs
+
+    @property
+    def jsons(self):
+        """
+        The json files in the search output files directory
+        """
+        return self._outputs(".json")
+
+    @property
+    def arrays(self):
+        """
+        The csv files in the search output files directory
+        """
+        return self._outputs(".csv")
+
+    @property
+    def pickles(self):
+        """
+        The pickle files in the search output files directory
+        """
+        return self._outputs(".pickle")
+
+    @property
+    def hdus(self):
+        """
+        The fits files in the search output files directory
+        """
+        return self._outputs(".fits")
+
+    @property
+    def log_likelihood(self) -> float:
+        """
+        The log likelihood of the maximum log likelihood sample
+        """
+        return self.samples.max_log_likelihood_sample.log_likelihood
+
     def __getattr__(self, item):
         """
         Attempt to load a pickle by the same name from the search output directory.
@@ -114,25 +138,61 @@ class SearchOutput:
         except (FileNotFoundError, ValueError):
             pass
 
+
+class SearchOutput(AbstractSearchOutput):
+    """
+    @DynamicAttrs
+    """
+
+    def __init__(self, directory: Path, reference: dict = None):
+        """
+        Represents the output of a single search. Comprises a metadata file and other dataset files.
+
+        Parameters
+        ----------
+        directory
+            The directory of the search
+        """
+        super().__init__(directory)
+        self.__search = None
+        self.__model = None
+
+        self.directory = directory
+
+        self._reference = reference
+        self.file_path = directory / "metadata"
+
+        try:
+            with open(self.file_path) as f:
+                self.text = f.read()
+                pairs = [
+                    line.split("=") for line in self.text.split("\n") if "=" in line
+                ]
+                self.__dict__.update({pair[0]: pair[1] for pair in pairs})
+        except FileNotFoundError:
+            pass
+
+    @property
+    def instance(self):
+        try:
+            return self.samples.max_log_likelihood()
+        except (AttributeError, NotImplementedError):
+            return None
+
     @property
     def id(self):
         return str(Identifier([self.search, self.model, self.search.unique_tag]))
 
     @property
-    def is_complete(self) -> bool:
+    def samples(self) -> SamplesPDF:
         """
-        Whether the search has completed
+        The samples of the search, parsed from a CSV containing individual samples
+        and a JSON containing metadata.
         """
-        return (self.directory / ".completed").exists()
-
-    @property
-    def samples(self):
         try:
-            info_path = self.files_path / "samples_info.json"
-            samples_path = self.files_path / "samples.csv"
-            with open(info_path) as f:
-                info_json = json.load(f)
-            with open(samples_path) as f:
+            info_json = JSONOutput("info", self.files_path / "samples_info.json").dict
+
+            with open(self.files_path / "samples.csv") as f:
                 sample_list = samples_from_iterator(csv.reader(f))
 
             return SamplesPDF.from_list_info_and_model(
@@ -197,7 +257,7 @@ class SearchOutput:
         return path.join(phase, dataset_name)
 
     @property
-    def search(self) -> abstract_search.NonLinearSearch:
+    def search(self):
         """
         The search object that was used in this phase
         """
@@ -226,15 +286,39 @@ class SearchOutput:
         return "<PhaseOutput {}>".format(self)
 
 
-class GridSearchOutput(SearchOutput):
+class GridSearchOutput(AbstractSearchOutput):
     @property
-    def parent_identifier(self) -> Optional[str]:
+    def unique_tag(self) -> str:
         """
-        Read the parent identifier for a fit in a directory.
+        The unique tag of the grid search.
+        """
+        with open(self.directory / ".is_grid_search") as f:
+            return f.read()
 
-        Defaults to None if no .parent_identifier file is found.
+    @property
+    def id(self):
+        return self.unique_tag
+
+
+class GridSearch:
+    def __init__(
+        self,
+        grid_search_output: GridSearchOutput,
+        search_outputs: List[SearchOutput],
+    ):
         """
-        try:
-            return (self.directory / ".parent_identifier").read_text()
-        except FileNotFoundError:
-            return None
+        Represents the output of a grid search. Comprises overall information from the grid search
+        and output from each individual search.
+
+        Parameters
+        ----------
+        grid_search_output
+            The output of the grid search
+        search_outputs
+            The outputs of each individual search performed as part of the grid search
+        """
+        self.grid_search_output = grid_search_output
+        self.search_outputs = search_outputs
+
+    def __getattr__(self, item):
+        return getattr(self.grid_search_output, item)
diff --git a/autofit/database/aggregator/scrape.py b/autofit/database/aggregator/scrape.py
index c0fc153a1..fd009bac2 100644
--- a/autofit/database/aggregator/scrape.py
+++ b/autofit/database/aggregator/scrape.py
@@ -1,34 +1,14 @@
-import csv
-import json
 import logging
-import os
-import pickle
 from pathlib import Path
-from typing import Optional, Union, Generator, Tuple
-
-import numpy as np
+from typing import Optional, Union
 
 from .. import model as m
 from ..sqlalchemy_ import sa
-from autofit.non_linear.samples.pdf import SamplesPDF
-from autofit.non_linear.samples.sample import samples_from_iterator
+from autofit.aggregator.search_output import SearchOutput
 
 logger = logging.getLogger(__name__)
 
 
-def _parent_identifier(directory: str) -> Optional[str]:
-    """
-    Read the parent identifier for a fit in a directory.
-
-    Defaults to None if no .parent_identifier file is found.
-    """
-    try:
-        with open(f"{directory}/.parent_identifier") as f:
-            return f.read()
-    except FileNotFoundError:
-        return None
-
-
 class Scraper:
     def __init__(
         self,
@@ -53,6 +33,13 @@ class Scraper:
         self.session = session
         self.reference = reference
 
+        from autofit.aggregator.aggregator import Aggregator as ClassicAggregator
+
+        self.aggregator = ClassicAggregator.from_directory(
+            self.directory,
+            reference=self.reference,
+        )
+
     def scrape(self):
         """
         Recursively scrape fits from the directory and
@@ -73,19 +60,8 @@ class Scraper:
         Generator yielding Fit database objects
         """
         logger.info(f"Scraping directory {self.directory}")
-        from autofit.aggregator.aggregator import Aggregator as ClassicAggregator
-
-        aggregator = ClassicAggregator.from_directory(
-            self.directory,
-            reference=self.reference,
-        )
-        logger.info(f"{len(aggregator)} searches found")
-        for item in aggregator:
-            parent_identifier = _parent_identifier(directory=item.directory)
-
-            model = item.model
-            samples = item.samples
-
+        logger.info(f"{len(self.aggregator)} searches found")
+        for item in self.aggregator:
             logger.info(
                 f"Creating fit for: "
                 f"{item.search.paths.path_prefix} "
@@ -94,37 +70,28 @@ class Scraper:
                 f"{item.id} "
             )
 
-            try:
-                instance = samples.max_log_likelihood()
-            except (AttributeError, NotImplementedError):
-                instance = None
-
             try:
                 fit = self._retrieve_model_fit(item)
                 logger.warning(f"Fit already existed with identifier {item.id}")
             except sa.orm.exc.NoResultFound:
-                try:
-                    log_likelihood = samples.max_log_likelihood_sample.log_likelihood
-                except AttributeError:
-                    log_likelihood = None
                 fit = m.Fit(
                     id=item.id,
                     name=item.search.name,
                     unique_tag=item.search.unique_tag,
-                    model=model,
-                    instance=instance,
+                    model=item.model,
+                    instance=item.instance,
                     is_complete=item.is_complete,
                     info=item.info,
-                    max_log_likelihood=log_likelihood,
-                    parent_id=parent_identifier,
+                    max_log_likelihood=item.log_likelihood,
+                    parent_id=item.parent_identifier,
                 )
 
-            _add_files(fit, Path(item.files_path))
+            _add_files(fit, item)
             for i, child_analysis in enumerate(item.child_analyses):
                 child_fit = m.Fit(
                     id=f"{item.id}_{i}",
                 )
-                _add_files(child_fit, child_analysis.files_path)
+                _add_files(child_fit, child_analysis)
                 fit.children.append(child_fit)
 
             yield fit
@@ -140,32 +107,21 @@ class Scraper:
         ------
         Fit objects representing grid searches with child fits associated
         """
-        from autofit.aggregator.aggregator import Aggregator as ClassicAggregator
-
-        for root, _, filenames in os.walk(self.directory):
-            if ".is_grid_search" in filenames:
-                path = Path(root)
-
-                is_complete = (path / ".completed").exists()
-
-                with open(path / ".is_grid_search") as f:
-                    unique_tag = f.read()
-
-                grid_search = m.Fit(
-                    id=path.name,
-                    unique_tag=unique_tag,
-                    is_grid_search=True,
-                    parent_id=_parent_identifier(root),
-                    is_complete=is_complete,
-                )
+        for item in self.aggregator.grid_searches():
+            grid_search = m.Fit(
+                id=item.id,
+                unique_tag=item.unique_tag,
+                is_grid_search=True,
+                parent_id=item.parent_identifier,
+                is_complete=item.is_complete,
+            )
 
-                _add_files(grid_search, path / "files")
+            _add_files(grid_search, item)
 
-                aggregator = ClassicAggregator.from_directory(root)
-                for item in aggregator:
-                    fit = self._retrieve_model_fit(item)
-                    grid_search.children.append(fit)
-                yield grid_search
+            for search in item.search_outputs:
+                fit = self._retrieve_model_fit(search)
+                grid_search.children.append(fit)
+            yield grid_search
 
     def _retrieve_model_fit(self, item) -> m.Fit:
         """
@@ -188,31 +144,7 @@ class Scraper:
         return self.session.query(m.Fit).filter(m.Fit.id == item.id).one()
 
 
-def names_and_paths(
-    files_path: Path,
-    suffix: str,
-) -> Generator[Tuple[str, Path], None, None]:
-    """
-    Get the names and paths of files with a given suffix in a directory.
-
-    Parameters
-    ----------
-    files_path
-        The path in which the files are stored
-    suffix
-        The suffix of the files to retrieve (e.g. ".json")
-
-    Returns
-    -------
-    A generator of tuples of the form (name, path) where name is the path to the file
-    joined by . without the suffix and path is the path to the file
-    """
-    for file in list(files_path.rglob(f"*{suffix}")):
-        name = ".".join(file.relative_to(files_path).with_suffix("").parts)
-        yield name, file
-
-
-def _add_files(fit: m.Fit, files_path: Path):
+def _add_files(fit: m.Fit, item: SearchOutput):
     """
     Load files from the path and add them to the database.
 
@@ -220,45 +152,23 @@ def _add_files(fit: m.Fit, files_path: Path):
     ----------
     fit
         A fit to which the pickles belong
-    files_path
-        The path in which the JSONs are stored
     """
-    info_path = files_path / "samples_info.json"
-    samples_path = files_path / "samples.csv"
-    if info_path.exists() and samples_path.exists():
-        with open(info_path) as f:
-            info_json = json.load(f)
-        with open(samples_path) as f:
-            sample_list = samples_from_iterator(csv.reader(f))
+    try:
+        fit.samples = item.samples
+    except AttributeError:
+        logger.warning(f"Failed to load samples for {fit.id}")
 
-        fit.samples = SamplesPDF.from_list_info_and_model(
-            sample_list=sample_list,
-            samples_info=info_json,
-            model=fit.model,
-        )
+    for json_output in item.jsons:
+        fit.set_json(json_output.name, json_output.dict)
 
-    for name, path in names_and_paths(files_path, ".json"):
-        with open(path) as f:
-            fit.set_json(name, json.load(f))
-    for name, path in names_and_paths(files_path, ".csv"):
-        try:
-            with open(path) as f:
-                fit.set_array(name, np.loadtxt(f, delimiter=","))
-        except ValueError:
-            logger.debug(f"Failed to load array from {path}")
+    for pickle_output in item.pickles:
+        fit.set_pickle(pickle_output.name, pickle_output.value)
 
-    for name, path in names_and_paths(files_path, ".pickle"):
+    for array_output in item.arrays:
         try:
-            with open(path, "r+b") as f:
-                fit[name] = pickle.load(f)
-        except (pickle.UnpicklingError, ModuleNotFoundError) as e:
-            if path.name == "dynesty.pickle":
-                continue
-
-            raise pickle.UnpicklingError(f"Failed to unpickle: {path}") from e
-
-    for name, path in names_and_paths(files_path, ".fits"):
-        from astropy.io import fits
+            fit.set_array(array_output.name, array_output.value)
+        except ValueError:
+            logger.debug(f"Failed to load array {array_output.name} for {fit.id}")
 
-        with open(path, "rb") as f:
-            fit.set_hdu(name, fits.PrimaryHDU.readfrom(f))
+    for hdu_output in item.hdus:
+        fit.set_hdu(hdu_output.name, hdu_output.value)
diff --git a/autofit/database/model/fit.py b/autofit/database/model/fit.py
index 53d247eb1..cf15a8035 100644
--- a/autofit/database/model/fit.py
+++ b/autofit/database/model/fit.py
@@ -368,6 +368,14 @@ class Fit(Base):
         new = Array(name=key, array=value)
         self.arrays = [p for p in self.arrays if p.name != key] + [new]
 
+    def set_pickle(self, key: str, value):
+        new = Pickle(name=key)
+        if isinstance(value, (str, bytes)):
+            new.string = value
+        else:
+            new.value = value
+        self.pickles = [p for p in self.pickles if p.name != key] + [new]
+
     def get_array(self, key: str) -> np.ndarray:
         """
         Retrieve an array from the database.
@@ -439,12 +447,7 @@ class Fit(Base):
         value
             A string, bytes or object
         """
-        new = Pickle(name=key)
-        if isinstance(value, (str, bytes)):
-            new.string = value
-        else:
-            new.value = value
-        self.pickles = [p for p in self.pickles if p.name != key] + [new]
+        self.set_pickle(key, value)
 
     def __delitem__(self, key):
         self.pickles = [p for p in self.pickles if p.name != key]


Test patch:
diff --git a/test_autofit/aggregator/test_scrape.py b/test_autofit/aggregator/test_scrape.py
index 63818dda4..178cd8510 100644
--- a/test_autofit/aggregator/test_scrape.py
+++ b/test_autofit/aggregator/test_scrape.py
@@ -1,5 +1,6 @@
 import pytest
 
+from autofit import SearchOutput
 from autofit.database.aggregator.scrape import _add_files
 
 
@@ -13,13 +14,22 @@ class MockFit:
     def __setitem__(self, key, value):
         pass
 
+    def set_pickle(self, key, value):
+        pass
+
+    def set_array(self, key, value):
+        pass
+
+    def set_hdu(self, key, value):
+        pass
+
 
 @pytest.fixture(name="fit")
 def make_fit(directory):
     fit = MockFit()
     _add_files(
         fit=fit,
-        files_path=directory / "search_output" / "files",
+        item=SearchOutput(directory / "search_output"),
     )
     return fit
 
diff --git a/test_autofit/aggregator/test_unification.py b/test_autofit/aggregator/test_unification.py
index 227e15f82..f4a238108 100644
--- a/test_autofit/aggregator/test_unification.py
+++ b/test_autofit/aggregator/test_unification.py
@@ -1,10 +1,6 @@
 import autofit as af
 
 
-def test_jsons(search_output):
-    assert len(list(search_output.names_and_paths(".json"))) == 3
-
-
 def test_id(search_output):
     assert isinstance(search_output.id, str)
 
diff --git a/test_autofit/non_linear/grid/conftest.py b/test_autofit/non_linear/grid/conftest.py
index e33597781..71a4148ed 100644
--- a/test_autofit/non_linear/grid/conftest.py
+++ b/test_autofit/non_linear/grid/conftest.py
@@ -2,21 +2,37 @@ import pytest
 
 import autofit as af
 
+
 @pytest.fixture(name="mapper")
 def make_mapper():
-    return af.Collection(
-        component=af.Model(
-            af.m.MockClassx2Tuple
-        )
-    )
+    return af.Collection(component=af.Model(af.m.MockClassx2Tuple))
 
 
 @pytest.fixture(name="grid_search")
 def make_grid_search(mapper):
     mock_search = af.m.MockSearch()
     mock_search.paths = af.DirectoryPaths(name="")
-    search = af.SearchGridSearch(
+    search = af.SearchGridSearch(number_of_steps=10, search=mock_search)
+    return search
+
+
+@pytest.fixture(name="sample_name_paths")
+def make_sample_name_paths():
+    return af.DirectoryPaths(name="sample_name")
+
+
+@pytest.fixture(name="grid_search_10_result")
+def make_grid_search_10_result(mapper, sample_name_paths):
+    grid_search = af.SearchGridSearch(
+        search=af.m.MockOptimizer(),
         number_of_steps=10,
-        search=mock_search
     )
-    return search
+    grid_search.search.paths = sample_name_paths
+    return grid_search.fit(
+        model=mapper,
+        analysis=af.m.MockAnalysis(),
+        grid_priors=[
+            mapper.component.one_tuple.one_tuple_0,
+            mapper.component.one_tuple.one_tuple_1,
+        ],
+    )
diff --git a/test_autofit/non_linear/grid/test_optimizer_grid_search.py b/test_autofit/non_linear/grid/test_optimizer_grid_search.py
index e4b21e4bc..922acc979 100644
--- a/test_autofit/non_linear/grid/test_optimizer_grid_search.py
+++ b/test_autofit/non_linear/grid/test_optimizer_grid_search.py
@@ -152,28 +152,6 @@ def empty_args():
     af.m.MockOptimizer.init_args = list()
 
 
-@pytest.fixture(name="sample_name_paths")
-def make_sample_name_paths():
-    return af.DirectoryPaths(name="sample_name")
-
-
-@pytest.fixture(name="grid_search_10_result")
-def make_grid_search_10_result(mapper, sample_name_paths):
-    grid_search = af.SearchGridSearch(
-        search=af.m.MockOptimizer(),
-        number_of_steps=10,
-    )
-    grid_search.search.paths = sample_name_paths
-    return grid_search.fit(
-        model=mapper,
-        analysis=af.m.MockAnalysis(),
-        grid_priors=[
-            mapper.component.one_tuple.one_tuple_0,
-            mapper.component.one_tuple.one_tuple_1,
-        ],
-    )
-
-
 def test_csv_headers(grid_search_10_result, sample_name_paths):
     with open(sample_name_paths.output_path / "results.csv") as f:
         reader = csv.reader(f)
diff --git a/test_autofit/non_linear/grid/test_paths/test_aggregator.py b/test_autofit/non_linear/grid/test_paths/test_aggregator.py
new file mode 100644
index 000000000..e89ebd7be
--- /dev/null
+++ b/test_autofit/non_linear/grid/test_paths/test_aggregator.py
@@ -0,0 +1,37 @@
+import pytest
+
+from autofit.aggregator.aggregator import Aggregator
+
+
+@pytest.fixture(name="aggregator")
+def make_aggregator(sample_name_paths, grid_search_10_result):
+    return Aggregator.from_directory(sample_name_paths.output_path)
+
+
+def test_aggregate(aggregator):
+    assert len(aggregator) == 100
+    assert len(aggregator.grid_search_outputs) == 1
+
+
+@pytest.fixture(name="grid_search")
+def make_grid_search(aggregator):
+    (grid_search,) = aggregator.grid_searches()
+    return grid_search
+
+
+def test_correspondence(grid_search, aggregator):
+    assert len(grid_search.search_outputs) == 100
+    assert grid_search.search_outputs[0] is list(aggregator)[0]
+
+
+def test_attributes(grid_search):
+    assert grid_search.is_complete
+    assert isinstance(grid_search.unique_tag, str)
+
+
+def test_files(grid_search):
+    assert grid_search.result
+
+
+def test_jsons(grid_search):
+    assert len(grid_search.jsons) == 1