Repository: rhayes777/PyAutoFit
Pull request #964 (2024-03-11 11:46:09)

feature/custom quantities
Introduces latent variables.

Latent variables can be written by calling `save_latent_variables` on an instance of analysis.

```python
self.save_latent_variables(one=1.0, two=2.0)
```

This method must be called one time each analysis with the same arguments (or not called at all).

If the method is called then latent variables are either written to latent_variables.csv or to the database. They are included in scraping and aggregation as well as the result of a run.

Values can be accessed as a dictionary by iterating through or indexing latent variables.
```python
result.latent_variables[0]
> {"one": 1.0, "two": 2.0}

list(result.latent_variables)
> [{"one": 1.0, "two": 2.0}]
```

All values for some variables can be accessed by that variable name:
```python
result.latent_variables["one"]
> [1.0]




Patch:
diff --git a/autofit/__init__.py b/autofit/__init__.py
index ed2300b50..125adf193 100644
--- a/autofit/__init__.py
+++ b/autofit/__init__.py
@@ -66,6 +66,7 @@ from .mapper.prior_model.util import PriorModelNameValue
 from .non_linear.search.abstract_search import NonLinearSearch
 from .non_linear.analysis.analysis import Analysis
 from .non_linear.analysis.combined import CombinedAnalysis
+from .non_linear.analysis.latent_variables import LatentVariables
 from .non_linear.grid.grid_search import GridSearchResult
 from .non_linear.grid.sensitivity import Sensitivity
 from .non_linear.initializer import InitializerBall
diff --git a/autofit/aggregator/search_output.py b/autofit/aggregator/search_output.py
index 1d55a4452..b5883909c 100644
--- a/autofit/aggregator/search_output.py
+++ b/autofit/aggregator/search_output.py
@@ -9,6 +9,7 @@ from typing import Generator, Tuple, Optional, List, cast
 import dill
 
 from autoconf import cached_property
+
 from autofit.non_linear.samples.pdf import SamplesPDF
 from autofit.aggregator.file_output import (
     JSONOutput,
@@ -273,6 +274,18 @@ class SearchOutput(AbstractSearchOutput):
                 raise AttributeError("No samples found")
         return self._samples
 
+    @property
+    def latent_variables(self):
+        """
+        The latent variables of the search, parsed from a CSV file.
+        """
+        with open(self.files_path / "latent_variables.csv") as f:
+            reader = csv.reader(f)
+            headers = next(reader)
+            from autofit.non_linear.analysis.latent_variables import LatentVariables
+
+            return LatentVariables(headers, [list(map(float, row)) for row in reader])
+
     def names_and_paths(
         self,
         suffix: str,
diff --git a/autofit/database/aggregator/scrape.py b/autofit/database/aggregator/scrape.py
index 6b81b1535..c44bcf9b2 100644
--- a/autofit/database/aggregator/scrape.py
+++ b/autofit/database/aggregator/scrape.py
@@ -152,6 +152,11 @@ def _add_files_fit(fit: m.Fit, item: SearchOutput):
         fit.samples = item.samples
     except AttributeError:
         logger.warning(f"Failed to load samples for {fit.id}")
+    try:
+        fit.latent_variables = item.latent_variables
+    except FileNotFoundError:
+        logger.warning(f"Failed to load latent variables for {fit.id}")
+
     _add_files(fit, item)
 
 
@@ -171,6 +176,8 @@ def _add_files(fit: m.Fit, item: SearchOutput):
         fit.set_pickle(pickle_output.name, pickle_output.value)
 
     for array_output in item.arrays:
+        if array_output.name in ("samples", "latent_variables"):
+            continue
         try:
             fit.set_array(array_output.name, array_output.value)
         except ValueError:
diff --git a/autofit/database/migration/steps.py b/autofit/database/migration/steps.py
index 66636877e..7d2258426 100644
--- a/autofit/database/migration/steps.py
+++ b/autofit/database/migration/steps.py
@@ -17,6 +17,9 @@ steps = [
     Step(
         "CREATE TABLE hdu (id INTEGER NOT NULL, array_type VARCHAR, _header VARCHAR, PRIMARY KEY (id), FOREIGN KEY (id) REFERENCES array (id));"
     ),
+    Step(
+        "ALTER TABLE object ADD COLUMN latent_variables_for_id INTEGER;",
+    ),
 ]
 
 migrator = Migrator(*steps)
diff --git a/autofit/database/model/fit.py b/autofit/database/model/fit.py
index 7978eee0c..d808d136b 100644
--- a/autofit/database/model/fit.py
+++ b/autofit/database/model/fit.py
@@ -256,6 +256,9 @@ class Fit(Base):
     _samples = sa.orm.relationship(
         Object, uselist=False, foreign_keys=[Object.samples_for_id]
     )
+    _latent_variables = sa.orm.relationship(
+        Object, uselist=False, foreign_keys=[Object.latent_variables_for_id]
+    )
 
     @property
     @try_none
@@ -268,6 +271,15 @@ class Fit(Base):
             EfficientSamples(samples),
         )
 
+    @property
+    @try_none
+    def latent_variables(self) -> dict:
+        return self._latent_variables()
+
+    @latent_variables.setter
+    def latent_variables(self, latent_variables):
+        self._latent_variables = Object.from_object(latent_variables.efficient())
+
     @property
     def info(self):
         return {info.key: info.value for info in self._info}
diff --git a/autofit/database/model/model.py b/autofit/database/model/model.py
index 36cde5adb..79e40ebca 100644
--- a/autofit/database/model/model.py
+++ b/autofit/database/model/model.py
@@ -35,6 +35,11 @@ class Object(Base):
         "Fit", uselist=False, foreign_keys=[samples_for_id]
     )
 
+    latent_variables_for_id = sa.Column(sa.String, sa.ForeignKey("fit.id"))
+    latent_variables_for = sa.orm.relationship(
+        "Fit", uselist=False, foreign_keys=[latent_variables_for_id]
+    )
+
     children: List["Object"] = sa.orm.relationship(
         "Object",
         uselist=True,
diff --git a/autofit/example/analysis.py b/autofit/example/analysis.py
index 3659254f9..28d94e022 100644
--- a/autofit/example/analysis.py
+++ b/autofit/example/analysis.py
@@ -1,6 +1,6 @@
 import os
 import matplotlib.pyplot as plt
-from typing import List
+from typing import List, Optional
 
 from autofit.jax_wrapper import numpy as np
 
@@ -60,8 +60,13 @@ class Analysis(af.Analysis):
         chi_squared_map = (residual_map / self.noise_map) ** 2.0
         log_likelihood = -0.5 * sum(chi_squared_map)
 
+        self.save_latent_variables(
+            fwmh=instance.fwhm
+        )
+
         return log_likelihood
 
+
     def visualize(self, paths: af.DirectoryPaths, instance: af.ModelInstance, during_analysis : bool):
         """
         During a model-fit, the `visualize` method is called throughout the non-linear search and is used to output
diff --git a/autofit/non_linear/analysis/analysis.py b/autofit/non_linear/analysis/analysis.py
index 89841a881..2f2a64280 100644
--- a/autofit/non_linear/analysis/analysis.py
+++ b/autofit/non_linear/analysis/analysis.py
@@ -1,10 +1,12 @@
 import logging
 from abc import ABC
 import os
+from typing import Optional
 
 from autoconf import conf
 
 from autofit.mapper.prior_model.abstract import AbstractPriorModel
+from autofit.non_linear.analysis.latent_variables import LatentVariables
 from autofit.non_linear.paths.abstract import AbstractPaths
 from autofit.non_linear.paths.database import DatabasePaths
 from autofit.non_linear.paths.null import NullPaths
@@ -20,6 +22,39 @@ class Analysis(ABC):
     likelihood that some instance fits some data.
     """
 
+    @property
+    def latent_variables(self) -> Optional[LatentVariables]:
+        """
+        Custom quantities that are computed during the analysis.
+
+        If no latent variables have been saved, this will return None.
+        """
+        try:
+            return self._latent_variables
+        except AttributeError:
+            return None
+
+    def save_latent_variables(self, **kwargs: float):
+        """
+        Save latent variables that are computed during the analysis.
+
+        This should only be called once per a fit and must always be passed the same latent variables.
+
+        Parameters
+        ----------
+        kwargs
+            The latent variables to save.
+
+        Raises
+        ------
+        SamplesException
+            If the same latent variables are not passed to `add` each iteration.
+        """
+        if not hasattr(self, "_latent_variables"):
+            # noinspection PyAttributeOutsideInit
+            self._latent_variables = LatentVariables()
+        self._latent_variables.add(**kwargs)
+
     def with_model(self, model):
         """
         Associate an explicit model with this analysis. Instances of the model
@@ -139,6 +174,7 @@ class Analysis(ABC):
     def make_result(self, samples):
         return Result(
             samples=samples,
+            latent_variables=self.latent_variables,
         )
 
     def profile_log_likelihood_function(self, paths: AbstractPaths, instance):
diff --git a/autofit/non_linear/analysis/latent_variables.py b/autofit/non_linear/analysis/latent_variables.py
new file mode 100644
index 000000000..00519f87a
--- /dev/null
+++ b/autofit/non_linear/analysis/latent_variables.py
@@ -0,0 +1,116 @@
+from typing import List, Optional, Union
+
+import numpy as np
+
+from autofit import exc
+
+
+class LatentVariables:
+    def __init__(
+        self,
+        names: Optional[List[str]] = None,
+        values=None,
+    ):
+        """
+        A collection of latent variables computed during an optimisation.
+
+        Parameters
+        ----------
+        names
+            The names of the latent variables.
+        values
+            The values of the latent variables.
+        """
+        self.names = names or []
+        self.values = values or []
+
+        self._current_row = []
+
+    def _position(self, name: str):
+        """
+        Get the position of a custom quantity by its name.
+
+        Parameters
+        ----------
+        name
+            The name of the custom quantity.
+
+        Returns
+        -------
+        The column position of the custom quantity.
+        """
+        return self.names.index(name)
+
+    def add(self, **kwargs: float):
+        """
+        Add latent variables to the collection. This should be called once
+        per a fit and must always be passed the same latent variables.
+
+        Parameters
+        ----------
+        kwargs
+            The latent variables to add.
+
+        Raises
+        ------
+        SamplesException
+            If the same latent variables are not passed to `add` each iteration.
+        """
+        if self.names and set(kwargs.keys()) != set(self.names):
+            raise exc.SamplesException(
+                "The same latent variables must be passed to `add` each iteration."
+            )
+
+        for name in kwargs:
+            if name not in self.names:
+                self.names.append(name)
+
+        self.values.append([kwargs[name] for name in self.names])
+
+    def efficient(self) -> "LatentVariables":
+        """
+        Convert the values to a numpy array for efficient storage in the database.
+        """
+        return LatentVariables(
+            names=self.names,
+            values=np.array(self.values),
+        )
+
+    def _convert_values_to_dict(self, values):
+        return {name: value for name, value in zip(self.names, values)}
+
+    def __getitem__(self, item: Union[int, str]) -> Union[dict, List[float]]:
+        """
+        Retrieve a dictionary of latent variables given a row number or a list of
+        values for a specific custom quantity given its name.
+
+        Parameters
+        ----------
+        item
+            The row number or name of the custom quantity.
+
+        Returns
+        -------
+        The latent variables.
+        """
+        if isinstance(item, str):
+            return [value_list[self._position(item)] for value_list in self.values]
+        return self._convert_values_to_dict(self.values[item])
+
+    def __len__(self):
+        return len(self.values)
+
+    def __iter__(self):
+        for values in self.values:
+            yield self._convert_values_to_dict(values)
+
+    def minimise(self, index: int):
+        """
+        Minimise the latent variables to a single row.
+
+        Parameters
+        ----------
+        index
+            The row number to keep.
+        """
+        return LatentVariables(names=self.names, values=[self.values[index]])
diff --git a/autofit/non_linear/paths/abstract.py b/autofit/non_linear/paths/abstract.py
index 1c39bb8ac..1fcaeb7c9 100644
--- a/autofit/non_linear/paths/abstract.py
+++ b/autofit/non_linear/paths/abstract.py
@@ -13,6 +13,7 @@ import numpy as np
 
 from autoconf import conf
 from autofit.mapper.identifier import Identifier, IdentifierField
+
 from autofit.text import text_util
 from autofit.text.text_util import derived_info_from
 from autofit.tools.util import open_, zip_directory
@@ -421,6 +422,13 @@ class AbstractPaths(ABC):
         Save samples to the database
         """
 
+    @abstractmethod
+    def save_latent_variables(self, latent_variables, samples):
+        """
+        Save latent variables. These are values computed from an instance and output
+        during analysis.
+        """
+
     @abstractmethod
     def save_derived_quantities(self, samples):
         """
@@ -466,6 +474,10 @@ class AbstractPaths(ABC):
     def _samples_file(self) -> Path:
         return self._files_path / "samples.csv"
 
+    @property
+    def _latent_variables_file(self) -> Path:
+        return self._files_path / "latent_variables.csv"
+
     @property
     def _derived_quantities_file(self) -> Path:
         return self._files_path / "derived_quantities.csv"
diff --git a/autofit/non_linear/paths/database.py b/autofit/non_linear/paths/database.py
index b454b7ac0..fa309b2b0 100644
--- a/autofit/non_linear/paths/database.py
+++ b/autofit/non_linear/paths/database.py
@@ -1,5 +1,5 @@
 import shutil
-from typing import Dict, Optional, Union
+from typing import Optional, Union
 
 from autoconf.output import conditional_output
 from autofit.database.sqlalchemy_ import sa
@@ -10,6 +10,8 @@ from autofit.database.model import Fit
 from autoconf.dictable import to_dict
 from autofit.database.aggregator.info import Info
 
+# from autofit.non_linear.analysis.latent_variables import LatentVariables
+
 
 class DatabasePaths(AbstractPaths):
     def __init__(
@@ -268,6 +270,16 @@ class DatabasePaths(AbstractPaths):
         self.fit.samples = samples
         self.fit.set_json("samples_info", samples.samples_info)
 
+    def save_latent_variables(self, latent_variables, samples):
+        if not self.save_all_samples:
+            latent_variables = latent_variables.minimise(
+                samples.max_log_likelihood_index
+            )
+        self.fit.latent_variables = latent_variables
+
+    def load_latent_variables(self):
+        return self.fit.latent_variables
+
     def save_derived_quantities(self, samples):
         """
         Write out the derived quantities of the model.
diff --git a/autofit/non_linear/paths/directory.py b/autofit/non_linear/paths/directory.py
index 0d48502f6..c851eafdc 100644
--- a/autofit/non_linear/paths/directory.py
+++ b/autofit/non_linear/paths/directory.py
@@ -1,3 +1,4 @@
+import csv
 import shutil
 
 import dill
@@ -15,6 +16,8 @@ from autofit.text import formatter
 from autofit.tools.util import open_
 
 from .abstract import AbstractPaths
+
+# from ..analysis.latent_variables import LatentVariables
 from ..samples import load_from_table
 from autofit.non_linear.samples.pdf import SamplesPDF
 import numpy as np
@@ -245,6 +248,37 @@ class DirectoryPaths(AbstractPaths):
                         f"Could not save covariance matrix because of the following error:\n{e}"
                     )
 
+    def save_latent_variables(
+        self,
+        latent_variables,
+        samples,
+    ):
+        """
+        Write out the latent variables of the model to a file.
+
+        Parameters
+        ----------
+        latent_variables
+            The latent variables of the model
+        samples
+            The samples of the model
+        """
+        write_table(
+            filename=str(self._latent_variables_file),
+            headers=latent_variables.names,
+            rows=latent_variables.values,
+        )
+
+    def load_latent_variables(self):
+        with open(self._latent_variables_file, "r+", newline="") as f:
+            reader = csv.reader(f)
+            names = list(next(reader))
+            values = [list(map(float, row)) for row in reader]
+
+        from autofit.non_linear.analysis.latent_variables import LatentVariables
+
+        return LatentVariables(names=names, values=values)
+
     def save_derived_quantities(self, samples: Samples):
         """
         Write out the derived quantities of the model to a file.
diff --git a/autofit/non_linear/paths/null.py b/autofit/non_linear/paths/null.py
index 25cc64f94..da55dbc70 100644
--- a/autofit/non_linear/paths/null.py
+++ b/autofit/non_linear/paths/null.py
@@ -9,6 +9,9 @@ class NullPaths(AbstractPaths):
     Null version of paths object for avoiding writing of files to disk
     """
 
+    def save_latent_variables(self, latent_variables, samples):
+        pass
+
     def save_derived_quantities(self, samples):
         pass
 
diff --git a/autofit/non_linear/result.py b/autofit/non_linear/result.py
index 532456ba5..d59fa4524 100644
--- a/autofit/non_linear/result.py
+++ b/autofit/non_linear/result.py
@@ -165,8 +165,9 @@ class AbstractResult(ABC):
         """
         return self.samples.model_bounded(b)
 
+
 class Result(AbstractResult):
-    def __init__(self, samples: Samples):
+    def __init__(self, samples: Samples, latent_variables=None):
         """
         The result of a non-linear search, which includes:
 
@@ -184,6 +185,7 @@ class Result(AbstractResult):
             The samples of the non-linear search
         """
         self._samples = samples
+        self.latent_variables = latent_variables
 
         self.__model = None
 
@@ -220,9 +222,7 @@ class Result(AbstractResult):
 
     @property
     def model(self):
-
         if self.__model is None:
-
             self.__model = self.samples.model.mapper_from_prior_means(
                 means=self.samples.prior_means
             )
diff --git a/autofit/non_linear/samples/samples.py b/autofit/non_linear/samples/samples.py
index 1a5133e77..378bd0ee3 100644
--- a/autofit/non_linear/samples/samples.py
+++ b/autofit/non_linear/samples/samples.py
@@ -102,7 +102,7 @@ class Samples(SamplesInterface, ABC):
                 sample.parameter_lists_for_paths(
                     self.paths if sample.is_path_kwargs else self.names
                 ),
-                ignore_prior_limits=True
+                ignore_prior_limits=True,
             )
             for sample in self.sample_list
         ]
@@ -447,6 +447,13 @@ class Samples(SamplesInterface, ABC):
                 most_likely_sample = sample
         return most_likely_sample
 
+    @property
+    def max_log_likelihood_index(self) -> int:
+        """
+        The index of the sample with the highest log likelihood.
+        """
+        return int(np.argmax(self.log_likelihood_list))
+
     @to_instance
     def max_log_likelihood(self) -> List[float]:
         """
diff --git a/autofit/non_linear/search/abstract_search.py b/autofit/non_linear/search/abstract_search.py
index fd66bacbd..aa4bf2e17 100644
--- a/autofit/non_linear/search/abstract_search.py
+++ b/autofit/non_linear/search/abstract_search.py
@@ -818,9 +818,7 @@ class NonLinearSearch(AbstractFactorOptimiser, ABC):
         if os.environ.get("PYAUTOFIT_TEST_MODE") == "1":
             logger.warning(f"TEST MODE ON: SEARCH WILL SKIP SAMPLING\n\n")
 
-            config_dict = self.config_dict_test_mode_from(
-                config_dict=config_dict
-            )
+            config_dict = self.config_dict_test_mode_from(config_dict=config_dict)
 
         return config_dict
 
@@ -908,6 +906,13 @@ class NonLinearSearch(AbstractFactorOptimiser, ABC):
         if self.is_master:
             self.paths.save_samples(samples=samples)
 
+            latent_variables = analysis.latent_variables
+            if latent_variables:
+                self.paths.save_latent_variables(
+                    latent_variables,
+                    samples=samples,
+                )
+
             if not during_analysis:
                 try:
                     self.paths.save_derived_quantities(samples=samples)


Test patch:
diff --git a/test_autofit/analysis/test_latent_variables.py b/test_autofit/analysis/test_latent_variables.py
new file mode 100644
index 000000000..11cffbb79
--- /dev/null
+++ b/test_autofit/analysis/test_latent_variables.py
@@ -0,0 +1,102 @@
+import pytest
+import numpy as np
+
+import autofit as af
+from autofit import DirectoryPaths
+from autofit.exc import SamplesException
+from autofit.non_linear.analysis.latent_variables import LatentVariables
+
+
+class Analysis(af.Analysis):
+    def log_likelihood_function(self, instance):
+        self.save_latent_variables(centre=instance.centre)
+        return 1.0
+
+
+def test_latent_variables():
+    latent_variables = LatentVariables()
+    latent_variables.add(centre=1.0)
+
+    assert latent_variables.names == ["centre"]
+    assert latent_variables.values == [[1.0]]
+
+
+def test_multiple_quantities():
+    latent_variables = LatentVariables()
+    latent_variables.add(centre=1.0, intensity=2.0)
+
+    assert latent_variables.names == ["centre", "intensity"]
+    assert latent_variables.values == [[1.0, 2.0]]
+
+
+def test_multiple_iterations():
+    latent_variables = LatentVariables()
+    latent_variables.add(centre=1.0, intensity=2.0)
+    latent_variables.add(centre=3.0, intensity=4.0)
+
+    assert latent_variables.names == ["centre", "intensity"]
+    assert latent_variables.values == [[1.0, 2.0], [3.0, 4.0]]
+
+
+def test_split_addition():
+    latent_variables = LatentVariables()
+    latent_variables.add(centre=1.0)
+    with pytest.raises(SamplesException):
+        latent_variables.add(intensity=2.0)
+
+
+def test_analysis_latent_variables():
+    analysis = Analysis()
+    instance = af.Gaussian()
+    analysis.log_likelihood_function(instance=instance)
+
+    latent_variables = analysis.latent_variables
+    assert latent_variables.names == ["centre"]
+    assert latent_variables.values == [[instance.centre]]
+
+
+@pytest.fixture(name="latent_variables")
+def make_latent_variables():
+    return LatentVariables(names=["centre"], values=[[1.0]])
+
+
+def test_set_directory_paths(output_directory, latent_variables):
+    directory_paths = DirectoryPaths()
+    directory_paths.save_latent_variables(
+        latent_variables=latent_variables,
+        samples=None,
+    )
+    loaded = directory_paths.load_latent_variables()
+    assert loaded.names == ["centre"]
+    assert loaded.values == [[1.0]]
+
+
+def test_efficient(latent_variables):
+    assert latent_variables.efficient().values == np.array([[1.0]])
+
+
+class MockSamples:
+    @property
+    def max_log_likelihood_index(self):
+        return 0
+
+
+def test_set_database_paths(session, latent_variables):
+    database_paths = af.DatabasePaths(session)
+    database_paths.save_latent_variables(
+        latent_variables=latent_variables, samples=MockSamples()
+    )
+    loaded = database_paths.load_latent_variables()
+    assert loaded.names == ["centre"]
+    assert loaded.values == [[1.0]]
+
+
+def test_iter(latent_variables):
+    assert list(latent_variables) == [{"centre": 1.0}]
+    assert latent_variables[0] == {"centre": 1.0}
+    assert latent_variables["centre"] == [1.0]
+
+
+def test_minimise(latent_variables):
+    latent_variables.minimise(0)
+    assert latent_variables.values == [[1.0]]