Repository: rhayes777/PyAutoFit
Pull request #1099 (2025-01-31 17:20:30)

feature/aggregate csv
Resolves all issues in #1087 except for errors which still require a subtle refactor



Patch:
diff --git a/autofit/__init__.py b/autofit/__init__.py
index 152fa6bf9..a800104e5 100644
--- a/autofit/__init__.py
+++ b/autofit/__init__.py
@@ -25,6 +25,7 @@ from .non_linear.samples import Sample
 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 .database.aggregator import Query
 from autofit.aggregator.fit_interface import Fit
 from .aggregator.search_output import SearchOutput
diff --git a/autofit/aggregator/aggregate_csv.py b/autofit/aggregator/aggregate_csv.py
index a2f8a8a38..1af9b0054 100644
--- a/autofit/aggregator/aggregate_csv.py
+++ b/autofit/aggregator/aggregate_csv.py
@@ -1,4 +1,5 @@
-from typing import Optional, Callable, List, Union
+from abc import ABC, abstractmethod
+from typing import Optional, Callable, List, Union, Tuple
 
 from pathlib import Path
 
@@ -8,11 +9,29 @@ from autofit.aggregator.search_output import SearchOutput
 import csv
 
 
-class Column:
+class AbstractColumn(ABC):
+    def __init__(self, name: str):
+        """
+        A column in the summary table.
+
+        Parameters
+        ----------
+        name
+            An optional name for the column
+        """
+        self.name = name
+
+    @abstractmethod
+    def value(self, row: "Row"):
+        pass
+
+
+class Column(AbstractColumn):
     def __init__(
         self,
         argument: str,
         name: Optional[str] = None,
+        use_max_log_likelihood: Optional[bool] = False,
     ):
         """
         A column in the summary table.
@@ -24,11 +43,23 @@ class Column:
         name
             An optional name for the column
         """
-        self.argument = argument
-        self.name = name or argument.replace(
-            ".",
-            "_",
+        super().__init__(
+            name
+            or argument.replace(
+                ".",
+                "_",
+            )
         )
+        self.argument = argument
+        self.use_max_log_likelihood = use_max_log_likelihood
+
+    def value(self, row: "Row"):
+        if self.use_max_log_likelihood:
+            kwargs = row.max_likelihood_kwargs
+        else:
+            kwargs = row.median_pdf_sample_kwargs
+
+        return kwargs[self.path]
 
     @property
     def path(self) -> tuple:
@@ -38,12 +69,8 @@ class Column:
         return tuple(self.argument.split("."))
 
 
-class ComputedColumn:
-    def __init__(
-        self,
-        name: str,
-        compute: Callable,
-    ):
+class ComputedColumn(AbstractColumn):
+    def __init__(self, name: str, compute: Callable):
         """
         A column in the summary table that is computed from other columns.
 
@@ -54,16 +81,41 @@ class ComputedColumn:
         compute
             A function that takes the median_pdf_sample arguments and returns the computed value
         """
-        self.name = name
+        super().__init__(name)
         self.compute = compute
 
+    def value(self, row: "Row"):
+        try:
+            return self.compute(row.result.samples)
+        except AttributeError as e:
+            raise AssertionError(
+                "Cannot compute additional fields if no samples.json present"
+            ) from e
+
+
+class LabelColumn(AbstractColumn):
+    def __init__(self, name: str, values: list):
+        """
+        A column in the summary table that is a label.
+
+        Parameters
+        ----------
+        name
+            The name of the column
+        """
+        super().__init__(name)
+        self.values = values
+
+    def value(self, row: "Row"):
+        return self.values[row.number]
+
 
 class Row:
     def __init__(
         self,
         result: SearchOutput,
-        columns: List[Column],
-        computed_columns: List[ComputedColumn],
+        columns: list,
+        number: int,
     ):
         """
         A row in the summary table corresponding to one search.
@@ -74,44 +126,57 @@ class Row:
             The search output from the aggregator
         columns
             The columns to include in the summary. These are taken from samples_summary or latent_summary
-        computed_columns
-            Columns filled with values computed from the samples
         """
-        self._result = result
+        self.result = result
         self._columns = columns
-        self._computed_columns = computed_columns
+        self.number = number
+
+    def _all_paths(self, path: Tuple[str, ...]):
+        model = self.result.model
+        return model.all_paths_for_prior(model.object_for_path(path))
+
+    def _add_paths(self, kwargs):
+        return {
+            path: value
+            for key, value in kwargs.items()
+            for path in self._all_paths(key)
+        }
 
     @property
-    def kwargs(self) -> dict:
+    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")
-        kwargs = samples_summary.median_pdf_sample.kwargs
+        samples_summary = self.result.value("samples_summary")
+        kwargs = self._add_paths(samples_summary.median_pdf_sample.kwargs)
 
-        latent_summary = self._result.value("latent_summary")
+        latent_summary = self.result.value("latent.latent_summary")
         if latent_summary is not None:
             kwargs.update(latent_summary.median_pdf_sample.kwargs)
 
         return kwargs
 
+    @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.value("samples_summary")
+        kwargs = self._add_paths(samples_summary.median_pdf_sample.kwargs)
+
+        latent_summary = self.result.value("latent.latent_summary")
+        if latent_summary is not None:
+            kwargs.update(latent_summary.max_log_likelihood_sample.kwargs)
+
+        return kwargs
+
     def dict(self) -> dict:
         """
         The row as a dictionary including an id and one entry for each column.
         """
-        row = {"id": self._result.id}
+        row = {"id": self.result.id}
         for column in self._columns:
-            value = self.kwargs[column.path]
-            row[column.name] = value
-
-        for column in self._computed_columns:
-            try:
-                value = column.compute(self._result.samples)
-                row[column.name] = value
-            except AttributeError as e:
-                raise AssertionError(
-                    "Cannot compute additional fields if no samples.json present"
-                ) from e
+            row[column.name] = column.value(self)
 
         return row
 
@@ -127,12 +192,12 @@ class AggregateCSV:
         """
         self._aggregator = aggregator
         self._columns = []
-        self._computed_columns = []
 
     def add_column(
         self,
         argument: str,
         name: Optional[str] = None,
+        use_max_log_likelihood: Optional[bool] = False,
     ):
         """
         Add a column to the summary table.
@@ -146,11 +211,14 @@ class AggregateCSV:
             e.g. "galaxies.lens.bulge.centre.centre_0"
         name
             An optional name for the column. If not provided, the argument will be used.
+        use_max_log_likelihood
+            If True, the maximum likelihood value will be used instead of the median PDF value.
         """
         self._columns.append(
             Column(
                 argument,
                 name=name,
+                use_max_log_likelihood=use_max_log_likelihood,
             )
         )
 
@@ -169,23 +237,41 @@ class AggregateCSV:
         compute
             A function that takes the median_pdf_sample arguments and returns the computed value
         """
-        self._computed_columns.append(
+        self._columns.append(
             ComputedColumn(
                 name,
                 compute,
             )
         )
 
+    def add_label_column(
+        self,
+        name: str,
+        values: list,
+    ):
+        """
+        Add a column to the summary table that is a label.
+
+        Parameters
+        ----------
+        name
+            The name of the column
+        values
+            The values of the column
+        """
+        self._columns.append(
+            LabelColumn(
+                name,
+                values,
+            )
+        )
+
     @property
     def fieldnames(self) -> List[str]:
         """
         The fieldnames for the CSV file.
         """
-        return (
-            ["id"]
-            + [column.name for column in self._columns]
-            + [column.name for column in self._computed_columns]
-        )
+        return ["id"] + [column.name for column in self._columns]
 
     def save(self, path: Union[str, Path]):
         """
@@ -202,11 +288,11 @@ class AggregateCSV:
                 fieldnames=self.fieldnames,
             )
             writer.writeheader()
-            for result in self._aggregator:
+            for i, result in enumerate(self._aggregator):
                 row = Row(
                     result,
                     self._columns,
-                    self._computed_columns,
+                    i,
                 )
 
                 writer.writerow(row.dict())
diff --git a/autofit/non_linear/search/abstract_search.py b/autofit/non_linear/search/abstract_search.py
index 24454a1f1..5ad770ece 100644
--- a/autofit/non_linear/search/abstract_search.py
+++ b/autofit/non_linear/search/abstract_search.py
@@ -1036,7 +1036,7 @@ class NonLinearSearch(AbstractFactorOptimiser, ABC):
                     self.paths.save_latent_samples(latent_samples)
                     self.paths.save_samples_summary(
                         latent_samples.summary(),
-                        "latent_summary",
+                        "latent/latent_summary",
                     )
 
             self.perform_visualization(


Test patch:
diff --git a/test_autofit/aggregator/aggregate_summary/fit_1/files/latent_summary.json b/test_autofit/aggregator/aggregate_summary/fit_1/files/latent/latent_summary.json
similarity index 100%
rename from test_autofit/aggregator/aggregate_summary/fit_1/files/latent_summary.json
rename to test_autofit/aggregator/aggregate_summary/fit_1/files/latent/latent_summary.json
diff --git a/test_autofit/aggregator/aggregate_summary/fit_1/files/model.json b/test_autofit/aggregator/aggregate_summary/fit_1/files/model.json
new file mode 100644
index 000000000..ca459ceba
--- /dev/null
+++ b/test_autofit/aggregator/aggregate_summary/fit_1/files/model.json
@@ -0,0 +1,37 @@
+{
+  "type": "collection",
+  "arguments": {
+    "galaxies": {
+      "type": "collection",
+      "arguments": {
+        "lens": {
+          "type": "collection",
+          "arguments": {
+            "bulge": {
+              "type": "collection",
+              "arguments": {
+                "centre": {
+                  "type": "collection",
+                  "arguments": {
+                    "centre_0": {
+                      "type": "Gaussian",
+                      "id": 0,
+                      "mean": 0.0,
+                      "sigma": 0.1
+                    },
+                    "centre_1": {
+                      "type": "Gaussian",
+                      "id": 1,
+                      "mean": 0.0,
+                      "sigma": 0.1
+                    }
+                  }
+                }
+              }
+            }
+          }
+        }
+      }
+    }
+  }
+}
\ No newline at end of file
diff --git a/test_autofit/aggregator/aggregate_summary/fit_1/files/samples_summary.json b/test_autofit/aggregator/aggregate_summary/fit_1/files/samples_summary.json
index f7841950f..163ab48bc 100644
--- a/test_autofit/aggregator/aggregate_summary/fit_1/files/samples_summary.json
+++ b/test_autofit/aggregator/aggregate_summary/fit_1/files/samples_summary.json
@@ -13,9 +13,7 @@
           "type": "dict",
           "arguments": {
             "galaxies.lens.bulge.centre.centre_0": -1.0,
-            "galaxies.lens.bulge..centre.centre_1": -2.0,
-            "galaxies.lens.shear.gamma_1": 3.0,
-            "galaxies.lens.shear.gamma_2": 4.0
+            "galaxies.lens.bulge.centre.centre_1": -2.0
           }
         }
       }
@@ -31,9 +29,7 @@
           "type": "dict",
           "arguments": {
             "galaxies.lens.bulge.centre.centre_0": -1.0,
-            "galaxies.lens.bulge..centre.centre_1": -2.0,
-            "galaxies.lens.shear.gamma_1": 3.0,
-            "galaxies.lens.shear.gamma_2": 4.0
+            "galaxies.lens.bulge.centre.centre_1": -2.0
           }
         }
       }
diff --git a/test_autofit/aggregator/aggregate_summary/fit_2/files/latent_summary.json b/test_autofit/aggregator/aggregate_summary/fit_2/files/latent/latent_summary.json
similarity index 100%
rename from test_autofit/aggregator/aggregate_summary/fit_2/files/latent_summary.json
rename to test_autofit/aggregator/aggregate_summary/fit_2/files/latent/latent_summary.json
diff --git a/test_autofit/aggregator/aggregate_summary/fit_2/files/model.json b/test_autofit/aggregator/aggregate_summary/fit_2/files/model.json
new file mode 100644
index 000000000..ca459ceba
--- /dev/null
+++ b/test_autofit/aggregator/aggregate_summary/fit_2/files/model.json
@@ -0,0 +1,37 @@
+{
+  "type": "collection",
+  "arguments": {
+    "galaxies": {
+      "type": "collection",
+      "arguments": {
+        "lens": {
+          "type": "collection",
+          "arguments": {
+            "bulge": {
+              "type": "collection",
+              "arguments": {
+                "centre": {
+                  "type": "collection",
+                  "arguments": {
+                    "centre_0": {
+                      "type": "Gaussian",
+                      "id": 0,
+                      "mean": 0.0,
+                      "sigma": 0.1
+                    },
+                    "centre_1": {
+                      "type": "Gaussian",
+                      "id": 1,
+                      "mean": 0.0,
+                      "sigma": 0.1
+                    }
+                  }
+                }
+              }
+            }
+          }
+        }
+      }
+    }
+  }
+}
\ No newline at end of file
diff --git a/test_autofit/aggregator/aggregate_summary/fit_2/files/samples_summary.json b/test_autofit/aggregator/aggregate_summary/fit_2/files/samples_summary.json
index 99d92138a..32b1652ea 100644
--- a/test_autofit/aggregator/aggregate_summary/fit_2/files/samples_summary.json
+++ b/test_autofit/aggregator/aggregate_summary/fit_2/files/samples_summary.json
@@ -13,9 +13,7 @@
           "type": "dict",
           "arguments": {
             "galaxies.lens.bulge.centre.centre_0": -5.0,
-            "galaxies.lens.bulge..centre.centre_1": -6.0,
-            "galaxies.lens.shear.gamma_1": 7.0,
-            "galaxies.lens.shear.gamma_2": 8.0
+            "galaxies.lens.bulge.centre.centre_1": -6.0
           }
         }
       }
@@ -31,9 +29,7 @@
           "type": "dict",
           "arguments": {
             "galaxies.lens.bulge.centre.centre_0": -5.0,
-            "galaxies.lens.bulge..centre.centre_1": -6.0,
-            "galaxies.lens.shear.gamma_1": 7.0,
-            "galaxies.lens.shear.gamma_2": 8.0
+            "galaxies.lens.bulge.centre.centre_1": -6.0
           }
         }
       }
diff --git a/test_autofit/aggregator/test_aggregate_csv.py b/test_autofit/aggregator/test_aggregate_csv.py
index 6f9a88180..2469971d8 100644
--- a/test_autofit/aggregator/test_aggregate_csv.py
+++ b/test_autofit/aggregator/test_aggregate_csv.py
@@ -41,6 +41,20 @@ def test_writes(output_path, summary):
     assert dicts[1]["id"] is not None
 
 
+def test_add_label_colum(
+    output_path,
+    summary,
+    load_output,
+):
+    summary.add_label_column("label", ["a", "b"])
+    summary.save(output_path)
+
+    dicts = load_output()
+
+    assert dicts[0]["label"] == "a"
+    assert dicts[1]["label"] == "b"
+
+
 def test_add_column(
     output_path,
     summary,
@@ -55,6 +69,23 @@ def test_add_column(
     assert dicts[1]["galaxies_lens_bulge_centre_centre_0"] == "-5.0"
 
 
+def test_use_max_log_likelihood(
+    output_path,
+    summary,
+    load_output,
+):
+    summary.add_column(
+        "galaxies.lens.bulge.centre.centre_0",
+        use_max_log_likelihood=True,
+    )
+    summary.save(output_path)
+
+    dicts = load_output()
+
+    assert dicts[0]["galaxies_lens_bulge_centre_centre_0"] == "-1.0"
+    assert dicts[1]["galaxies_lens_bulge_centre_centre_0"] == "-5.0"
+
+
 def test_add_named_column(
     output_path,
     summary,