perf: aggregator result-loading speedups (single-scan listing, cached summaries, faster samples parse)

autofit/aggregator/aggregator.py

@@ -152,6 +152,10 @@ def from_directory(
         The whole directory structure is traversed and a Phase object created for each directory that contains a
         metadata file.
 
+        Zipped search outputs are extracted as they are encountered, in the same traversal. A zip whose
+        extracted directory already exists is skipped, so repeat calls do not pay the extraction cost again;
+        delete the extracted directory to force re-extraction.
+
         Parameters
         ----------
         directory
@@ -164,12 +168,23 @@ def from_directory(
         """
         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 root, dirs, filenames in os.walk(directory, topdown=True):
+            for filename in filenames:
+                if filename.endswith(".zip"):
+                    extracted = Path(root) / filename[:-4]
+                    if extracted.exists():
+                        continue
+                    try:
+                        with zipfile.ZipFile(Path(root) / filename, "r") as f:
+                            f.extractall(extracted)
+                    except zipfile.BadZipFile:
+                        raise zipfile.BadZipFile(
+                            f"File is not a zip file: \n " f"{root} \n" f"{filename}"
+                        )
+                    dirs.append(filename[:-4])
 
             def should_add():
                 return not completed_only or ".completed" in filenames

autofit/aggregator/search_output.py

@@ -76,46 +76,53 @@ def parent_identifier(self) -> Optional[str]:
     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 files_path.rglob(f"*{suffix}"):
-            name = ".".join(file_path.relative_to(files_path).with_suffix("").parts)
-            outputs.append(FileOutput(name, file_path))
+    @cached_property
+    def _outputs_by_suffix(self) -> dict:
+        """
+        All output files in the files and image directories, grouped by suffix.
+
+        A single traversal of each directory serves every suffix — building the
+        json/pickle/csv/fits lists separately makes eight sweeps per search output,
+        which dominates load time when aggregating many results.
+        """
+        outputs = {".json": [], ".pickle": [], ".csv": [], ".fits": []}
+        for directory_name in ("files", "image"):
+            files_path = self.directory / directory_name
+            for file_path in files_path.rglob("*"):
+                if file_path.suffix in outputs:
+                    name = ".".join(
+                        file_path.relative_to(files_path).with_suffix("").parts
+                    )
+                    outputs[file_path.suffix].append(FileOutput(name, file_path))
         return outputs
 
     @cached_property
     def jsons(self) -> List[JSONOutput]:
         """
         The json files in the search output files directory
         """
-        return cast(List[JSONOutput], self._outputs(".json"))
+        return cast(List[JSONOutput], self._outputs_by_suffix[".json"])
 
     @cached_property
     def arrays(self):
         """
         The csv files in the search output files directory
         """
-        return self._outputs(".csv")
+        return self._outputs_by_suffix[".csv"]
 
     @cached_property
     def pickles(self):
         """
         The pickle files in the search output files directory
         """
-        return self._outputs(".pickle")
+        return self._outputs_by_suffix[".pickle"]
 
     @cached_property
     def fits(self):
         """
         The fits files in the search output files directory
         """
-        return self._outputs(".fits")
+        return self._outputs_by_suffix[".fits"]
 
     @property
     def max_log_likelihood(self) -> Optional[float]:
@@ -201,6 +208,8 @@ def __init__(self, directory: Path, reference: dict = None):
         self.__model = None
         self._samples = None
         self._latent_samples = None
+        self.__samples_summary = None
+        self.__latent_summary = None
 
         self.directory = directory
 
@@ -223,9 +232,11 @@ def samples_summary(self) -> SamplesSummary:
 
         This is loaded from a JSON file.
         """
-        summary = self.value("samples_summary")
-        summary.model = self.model
-        return summary
+        if self.__samples_summary is None:
+            summary = self.value("samples_summary")
+            summary.model = self.model
+            self.__samples_summary = summary
+        return self.__samples_summary
 
     @property
     def latent_summary(self) -> SamplesSummary:
@@ -234,9 +245,11 @@ def latent_summary(self) -> SamplesSummary:
 
         This is loaded from a JSON file.
         """
-        summary = self.value("latent.latent_summary")
-        summary.model = self.model
-        return summary
+        if self.__latent_summary is None:
+            summary = self.value("latent.latent_summary")
+            summary.model = self.model
+            self.__latent_summary = summary
+        return self.__latent_summary
 
     @property
     def instance(self):
@@ -251,7 +264,7 @@ def instance(self):
         except (AttributeError, NotImplementedError):
             return self.samples_summary.instance
 
-    @property
+    @cached_property
     def id(self) -> str:
         """
         The unique identifier of the search.

autofit/aggregator/summary/aggregate_csv/row.py

@@ -37,29 +37,33 @@ def _add_paths(self, kwargs):
             )
         }
 
-    @property
+    @cached_property
+    def _latent_summary(self):
+        return self.result.value("latent.latent_summary")
+
+    @cached_property
     def median_pdf_sample_kwargs(self) -> dict:
         """
         The median_pdf_sample arguments for the search from the samples_summary and latent_summary.
         """
-        samples_summary = self.result.value("samples_summary")
+        samples_summary = self.result.samples_summary
         kwargs = self._add_paths(samples_summary.median_pdf_sample.kwargs)
 
-        latent_summary = self.result.value("latent.latent_summary")
+        latent_summary = self._latent_summary
         if latent_summary is not None:
             kwargs.update(latent_summary.median_pdf_sample.kwargs)
 
         return kwargs
 
-    @property
+    @cached_property
     def max_likelihood_kwargs(self):
         """
         The median_pdf_sample arguments for the search from the samples_summary and latent_summary.
         """
         samples_summary = self.result.samples_summary
         kwargs = self._add_paths(samples_summary.median_pdf_sample.kwargs)
 
-        latent_summary = self.result.value("latent.latent_summary")
+        latent_summary = self._latent_summary
         if latent_summary is not None:
             kwargs.update(latent_summary.max_log_likelihood_sample.kwargs)
 
@@ -72,7 +76,7 @@ def values_at_sigma_1_kwargs(self) -> dict:
         """
         kwargs = self._add_paths(self.result.samples_summary.values_at_sigma_1)
 
-        latent_summary = self.result.value("latent.latent_summary")
+        latent_summary = self._latent_summary
         if latent_summary is not None:
             kwargs.update(
                 {
@@ -90,7 +94,7 @@ def values_at_sigma_3_kwargs(self) -> dict:
         """
         kwargs = self._add_paths(self.result.samples_summary.values_at_sigma_3)
 
-        latent_summary = self.result.value("latent.latent_summary")
+        latent_summary = self._latent_summary
         if latent_summary is not None:
             kwargs.update(
                 {
@@ -109,10 +113,10 @@ def dict(self) -> dict:
         for column in self._columns:
             value = column.value(self)
             if isinstance(value, dict):
-                for key, value in value.items():
-                    row[f"{column.name}_{key}" if key else column.name] = value
+                for key, item in value.items():
+                    row[f"{column.name}_{key}" if key else column.name] = item
             else:
-                row[column.name] = column.value(self)
+                row[column.name] = value
 
         return row
 

autofit/non_linear/samples/sample.py

@@ -271,19 +271,25 @@ def without_paths(self, paths: List[Tuple[str, ...]]) -> "Sample":
 def samples_from_iterator(iterator):
     samples = list()
 
-    headers = next(iterator)
-    headers = [header.strip() for header in headers]
-    for row in iterator:
-        d = {header: float(value) for header, value in zip(headers, row)}
+    headers = [header.strip() for header in next(iterator)]
+
+    # Classify and split the headers once — doing it per row (and letting Sample
+    # re-split every string key) dominates load time for large samples.csv files.
+    excluded = sample_args | {"log_posterior"}
+    arg_indices = {
+        header: index for index, header in enumerate(headers) if header in sample_args
+    }
+    kwarg_indices = [
+        (index, tuple(header.split(".")))
+        for index, header in enumerate(headers)
+        if header not in excluded
+    ]
 
+    for row in iterator:
         samples.append(
             Sample(
-                **{key: value for key, value in d.items() if key in sample_args},
-                kwargs={
-                    key: value
-                    for key, value in d.items()
-                    if key not in (sample_args | {"log_posterior"})
-                },
+                **{name: float(row[index]) for name, index in arg_indices.items()},
+                kwargs={key: float(row[index]) for index, key in kwarg_indices},
             )
         )
     return samples

test_autofit/aggregator/test_from_directory.py

@@ -0,0 +1,70 @@
+import shutil
+import zipfile
+from pathlib import Path
+
+import pytest
+
+from autofit.aggregator import Aggregator
+from autofit.aggregator.search_output import SearchOutput
+
+
+@pytest.fixture(name="scan_directory")
+def make_scan_directory(tmp_path):
+    source = Path(__file__).parent / "search_output"
+    destination = tmp_path / "search_output"
+    shutil.copytree(source, destination)
+    return tmp_path
+
+
+@pytest.fixture(name="zipped_directory")
+def make_zipped_directory(tmp_path):
+    source = Path(__file__).parent / "search_output"
+    with zipfile.ZipFile(tmp_path / "search_output.zip", "w") as f:
+        for path in source.rglob("*"):
+            if path.is_file():
+                f.write(path, path.relative_to(source))
+    return tmp_path
+
+
+def test_from_directory(scan_directory):
+    aggregator = Aggregator.from_directory(scan_directory)
+    assert len(aggregator) == 1
+
+
+def test_zip_extracted_and_loaded(zipped_directory):
+    aggregator = Aggregator.from_directory(zipped_directory)
+    assert len(aggregator) == 1
+    assert (zipped_directory / "search_output" / "metadata").exists()
+
+
+def test_zip_not_re_extracted(zipped_directory):
+    Aggregator.from_directory(zipped_directory)
+
+    (zipped_directory / "search_output" / ".completed").unlink()
+    aggregator = Aggregator.from_directory(zipped_directory)
+
+    assert len(aggregator) == 1
+    assert not (zipped_directory / "search_output" / ".completed").exists()
+
+
+def test_outputs_by_suffix(scan_directory):
+    search_output = SearchOutput(scan_directory / "search_output")
+
+    assert {output.name for output in search_output.jsons} == {
+        "directory.example",
+        "model",
+        "samples_info",
+        "search",
+    }
+    assert {output.name for output in search_output.pickles} == {"info"}
+    assert {output.name for output in search_output.fits} == {"psf"}
+    assert search_output.arrays == []
+
+
+def test_samples_summary_cached():
+    directory = (
+        Path(__file__).parent / "summary_files" / "aggregate_summary" / "fit_1"
+    )
+    search_output = SearchOutput(directory)
+
+    assert search_output.samples_summary is search_output.samples_summary