diff --git a/autofit/aggregator/aggregator.py b/autofit/aggregator/aggregator.py index 43ef1f23b..d949a5006 100755 --- a/autofit/aggregator/aggregator.py +++ b/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 diff --git a/autofit/aggregator/search_output.py b/autofit/aggregator/search_output.py index b1a4d0f93..5c573db39 100644 --- a/autofit/aggregator/search_output.py +++ b/autofit/aggregator/search_output.py @@ -76,17 +76,24 @@ 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 @@ -94,28 +101,28 @@ 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. diff --git a/autofit/aggregator/summary/aggregate_csv/row.py b/autofit/aggregator/summary/aggregate_csv/row.py index b73a31528..b1304b7ae 100644 --- a/autofit/aggregator/summary/aggregate_csv/row.py +++ b/autofit/aggregator/summary/aggregate_csv/row.py @@ -37,21 +37,25 @@ 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. @@ -59,7 +63,7 @@ def max_likelihood_kwargs(self): 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 diff --git a/autofit/non_linear/samples/sample.py b/autofit/non_linear/samples/sample.py index 6ed9bc7f3..1e40d991f 100644 --- a/autofit/non_linear/samples/sample.py +++ b/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 diff --git a/test_autofit/aggregator/test_from_directory.py b/test_autofit/aggregator/test_from_directory.py new file mode 100644 index 000000000..b8fbd3888 --- /dev/null +++ b/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