Repository: rhayes777/PyAutoFit
Pull request #967 (2024-03-14 20:09:01)

Feature/samples filter
For large model fitting problems outputs may use up a lot of hard-disk space, therefore full customization of the 
outputs is supported. 

This is controlled by the `output.yaml` config file found in the `config` folder of the workspace. This file contains
a full description of all customization options.

This PR enables customization of the `samples.csv` file, where all samples with a weight below a certain value can be automatically removed.

In many use cases, the `samples.csv` takes up the significant majority of the hard-disk space, which for large-scale
model-fitting problems can exceed gigabytes and be prohibitive to the analysis. This PR enables a workflow where this is not an issue.


Patch:
diff --git a/autofit/config/output.yaml b/autofit/config/output.yaml
index 94ca0f7ae..a9182824f 100644
--- a/autofit/config/output.yaml
+++ b/autofit/config/output.yaml
@@ -21,6 +21,23 @@ default: true # If true then files which are not explicitly listed here are outp
 
 samples: true
 
+# The `samples.csv` file contains every accepted sampled value of every free parameter with its log likelihood and
+# weight. For certain searches, the majority of samples have a very low weight, which has no numerical impact on the
+# results of the model-fit. However, these samples are still output to the `samples.csv` file, taking up hard-disk space
+# and slowing down analysis of the samples (e.g. via the database).
+
+# The `samples_weight_threshold` below specifies the threshold value of the weight such that samples with a weight
+# below this value are not output to the `samples.csv` file. This can be used to reduce the size of the `samples.csv`
+# file and speed up analysis of the samples.
+
+# Note that for many searches (e.g. MCMC) all samples have equal weight, and thus this threshold has no impact and
+# there is no simple way to save hard-disk space. However, for nested sampling, the majority of samples have a very
+# low weight and this threshold can be used to save hard-disk space.
+
+# Set value to empty (e.g. delete 1.0e-10 below) to disable this feature.
+
+samples_weight_threshold: 1.0e-10
+
 ### Search Internal ###
 
 # The search internal folder which contains a saved state of the non-linear search, as a .pickle or .dill file.
diff --git a/autofit/example/analysis.py b/autofit/example/analysis.py
index 962cb2899..a86aa15a2 100644
--- a/autofit/example/analysis.py
+++ b/autofit/example/analysis.py
@@ -60,9 +60,12 @@ 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
-        )
+        try:
+            self.save_latent_variables(
+                fwmh=instance.fwhm
+            )
+        except AttributeError:
+            pass
 
         return log_likelihood
 
diff --git a/autofit/non_linear/plot/nest_plotters.py b/autofit/non_linear/plot/nest_plotters.py
index 21fe0969b..c03f3418e 100644
--- a/autofit/non_linear/plot/nest_plotters.py
+++ b/autofit/non_linear/plot/nest_plotters.py
@@ -2,6 +2,7 @@ from anesthetic.samples import NestedSamples
 from anesthetic import make_2d_axes
 from functools import wraps
 import numpy as np
+import warnings
 
 from autoconf import conf
 
@@ -39,7 +40,7 @@ def log_value_error(func):
         """
         try:
             return func(self, *args, **kwargs)
-        except (ValueError, KeyError, AssertionError, IndexError, np.linalg.LinAlgError):
+        except (ValueError, KeyError, AssertionError, IndexError, TypeError, np.linalg.LinAlgError):
             self.log_plot_exception(func.__name__)
 
     return wrapper
@@ -74,12 +75,17 @@ class NestPlotter(SamplesPlotter):
             columns=self.model.parameter_labels_with_superscripts_latex
         )
 
+        from pandas.errors import SettingWithCopyWarning
+        warnings.filterwarnings("ignore", category=SettingWithCopyWarning)
+
         fig, axes = make_2d_axes(
             self.model.parameter_labels_with_superscripts_latex,
             figsize=figsize,
             facecolor=config_dict["facecolor"],
         )
 
+        warnings.filterwarnings("default", category=SettingWithCopyWarning)
+
         # prior = samples.prior()
         # prior.plot_2d(axes, alpha=0.9, label="prior")
         samples.plot_2d(
diff --git a/autofit/non_linear/samples/nest.py b/autofit/non_linear/samples/nest.py
index ea2c72aae..b0cdebc14 100644
--- a/autofit/non_linear/samples/nest.py
+++ b/autofit/non_linear/samples/nest.py
@@ -89,7 +89,7 @@ class SamplesNest(SamplesPDF):
         """
         The total number of accepted samples performed by the nested sampler.
         """
-        return len(self.log_likelihood_list)
+        return self.samples_info["total_accepted_samples"]
 
     @property
     def acceptance_ratio(self) -> float:
diff --git a/autofit/non_linear/samples/samples.py b/autofit/non_linear/samples/samples.py
index d4afe8181..8d692af50 100644
--- a/autofit/non_linear/samples/samples.py
+++ b/autofit/non_linear/samples/samples.py
@@ -382,6 +382,37 @@ class Samples(SamplesInterface, ABC):
         """
         return self.parameter_lists[sample_index]
 
+    def samples_above_weight_threshold_from(self, weight_threshold: float) -> "Samples":
+        """
+        Returns a new `Samples` object containing only the samples with a weight above the input threshold.
+
+        This function can be used after a non-linear search is complete, to reduce the samples to only the high weight
+        values. The benefit of this is that the corresponding `samples.csv` file will be reduced in hard-disk size.
+
+        For large libraries of results can significantly reduce the overall hard-disk space used and speed up the
+        time taken to load the samples from a .csv file and perform analysis on them.
+
+        For a sufficiently low threshold, this has a neglible impact on the numerical accuracy of the results, and
+        even higher values can be used for aggresive use cases where hard-disk space is at a premium.
+
+        Parameters
+        ----------
+        weight_threshold
+            The threshold of weight at which a sample is included in the new `Samples` object.
+        """
+        sample_list = []
+
+        for sample in self.sample_list:
+            if sample.weight > weight_threshold:
+                sample_list.append(sample)
+
+        return self.__class__(
+            model=self.model,
+            sample_list=sample_list,
+            samples_info=self.samples_info,
+        )
+
+
     def minimise(self) -> "Samples":
         """
         A copy of this object with only important samples retained
diff --git a/autofit/non_linear/search/abstract_search.py b/autofit/non_linear/search/abstract_search.py
index ff0931559..52da5f786 100644
--- a/autofit/non_linear/search/abstract_search.py
+++ b/autofit/non_linear/search/abstract_search.py
@@ -891,13 +891,13 @@ class NonLinearSearch(AbstractFactorOptimiser, ABC):
         self.iterations += self.iterations_per_update
         if during_analysis:
             self.logger.info(
-                f"Fit Still Running: Updating results after {self.iterations} iterations (see "
-                f"output folder for latest visualization, samples, etc.)"
+                f"""Fit Still Running: Updating results after {self.iterations} iterations (see
+                output folder for latest visualization, samples, etc.)"""
             )
         else:
             self.logger.info(
-                f"Fit Complete: Updating final results (see "
-                f"output folder for final visualization, samples, etc.)"
+                "Fit Complete: Updating final results (see "
+                "output folder for final visualization, samples, etc.)"
             )
 
         if not isinstance(self.paths, DatabasePaths) and not isinstance(
@@ -913,7 +913,22 @@ class NonLinearSearch(AbstractFactorOptimiser, ABC):
             return samples
 
         if self.is_master:
-            self.paths.save_samples(samples=samples)
+
+            samples_for_csv = samples
+
+            samples_weight_threshold = conf.instance["output"]["samples_weight_threshold"]
+
+            if samples_weight_threshold is not None:
+                samples_for_csv = samples_for_csv.samples_above_weight_threshold_from(
+                    weight_threshold=samples_weight_threshold
+                )
+
+                logger.info(
+                    f"Samples with weight less than {samples_weight_threshold} removed from samples.csv and not used to "
+                    f"compute parameter estimates errors, etc."
+                )
+
+            self.paths.save_samples(samples=samples_for_csv)
 
             latent_variables = analysis.latent_variables
             if latent_variables:
diff --git a/autofit/non_linear/search/nest/dynesty/search/abstract.py b/autofit/non_linear/search/nest/dynesty/search/abstract.py
index d301a1da9..28b65b6fb 100644
--- a/autofit/non_linear/search/nest/dynesty/search/abstract.py
+++ b/autofit/non_linear/search/nest/dynesty/search/abstract.py
@@ -211,6 +211,7 @@ class AbstractDynesty(AbstractNest, ABC):
         return {
             "log_evidence": np.max(search_internal.results.logz),
             "total_samples": int(np.sum(search_internal.results.ncall)),
+            "total_accepted_samples": len(search_internal.results.logl),
             "time": self.timer.time if self.timer else None,
             "number_live_points": self.number_live_points,
         }
diff --git a/autofit/non_linear/search/nest/nautilus/search.py b/autofit/non_linear/search/nest/nautilus/search.py
index 65e7d5fa1..467f685fc 100644
--- a/autofit/non_linear/search/nest/nautilus/search.py
+++ b/autofit/non_linear/search/nest/nautilus/search.py
@@ -459,6 +459,7 @@ class Nautilus(abstract_nest.AbstractNest):
         return {
             "log_evidence": search_internal.evidence(),
             "total_samples": int(search_internal.n_like),
+            "total_accepted_samples": int(search_internal.n_like),
             "time": self.timer.time if self.timer else None,
             "number_live_points": int(search_internal.n_live),
         }
diff --git a/autofit/non_linear/search/nest/ultranest/search.py b/autofit/non_linear/search/nest/ultranest/search.py
index 07f208fb0..934b16bd0 100644
--- a/autofit/non_linear/search/nest/ultranest/search.py
+++ b/autofit/non_linear/search/nest/ultranest/search.py
@@ -214,6 +214,7 @@ class UltraNest(abstract_nest.AbstractNest):
         return {
             "log_evidence": search_internal["logz"],
             "total_samples": search_internal["ncall"],
+            "total_accepted_samples": len(search_internal["weighted_samples"]["logl"]),
             "time": self.timer.time if self.timer else None,
             "number_live_points": self.config_dict_run["min_num_live_points"]
         }
diff --git a/docs/cookbooks/search.rst b/docs/cookbooks/search.rst
index 0a232eb4c..88b9b1984 100644
--- a/docs/cookbooks/search.rst
+++ b/docs/cookbooks/search.rst
@@ -11,6 +11,7 @@ It first covers standard options available for all non-linear searches:
 
 - **Example Fit**: A simple example of a non-linear search to remind us how it works.
 - **Output To Hard-Disk**: Output results to hard-disk so they can be inspected and used to restart a crashed search.
+- **Output Customization**: Customize the output of a non-linear search to hard-disk.
 - **Unique Identifier**: Ensure results are output in unique folders, so tthey do not overwrite each other.
 - **Iterations Per Update**: Control how often non-linear searches output results to hard-disk.
 - **Parallelization**: Use parallel processing to speed up the sampling of parameter space.
@@ -78,6 +79,32 @@ These outputs are fully described in the scientific workflow example.
 
     search = af.Emcee(path_prefix=path.join("folder_0", "folder_1"), name="example_mcmc")
 
+Output Customization
+--------------------
+
+For large model fitting problems outputs may use up a lot of hard-disk space, therefore full customization of the 
+outputs is supported. 
+
+This is controlled by the ``output.yaml`` config file found in the ``config`` folder of the workspace. This file contains
+a full description of all customization options.
+
+A few examples of the options available include:
+
+- Control over every file which is output to the ``files`` folder (e.g. ``model.json``, ``samples.csv``, etc.).
+
+- For the ``samples.csv`` file, all samples with a weight below a certain value can be automatically removed.
+
+- Customization of the ``samples_summary.json`` file, which summarizes the results of the model-fit  (e.g. the maximum 
+  log likelihood model, the median PDF model and 3 sigma error). These results are computed using the full set of
+  samples, ensuring samples removal via a weight cut does not impact the results.
+
+In many use cases, the ``samples.csv`` takes up the significant majority of the hard-disk space, which for large-scale
+model-fitting problems can exceed gigabytes and be prohibitive to the analysis. 
+
+Careful customization of the ``output.yaml`` file enables a workflow where the ``samples.csv`` file is never output, 
+but all important information is output in the ``samples_summary.json`` file using the full samples to compute all 
+results to high numerical accuracy.
+
 Unique Identifier
 -----------------
 


Test patch:
diff --git a/test_autofit/non_linear/samples/test_nest.py b/test_autofit/non_linear/samples/test_nest.py
index e691cfab1..58a0f2381 100644
--- a/test_autofit/non_linear/samples/test_nest.py
+++ b/test_autofit/non_linear/samples/test_nest.py
@@ -72,6 +72,7 @@ def test__acceptance_ratio_is_correct():
             "total_samples" : 10,
             "log_evidence" : 0.0,
             "number_live_points" : 5,
+            "total_accepted_samples": 5
             }
     )
 
diff --git a/test_autofit/non_linear/samples/test_samples.py b/test_autofit/non_linear/samples/test_samples.py
index 35436b6f4..40c5f45a5 100644
--- a/test_autofit/non_linear/samples/test_samples.py
+++ b/test_autofit/non_linear/samples/test_samples.py
@@ -134,6 +134,36 @@ def test__instance_from_sample_index():
     assert instance.mock_class.four == 4.0
 
 
+def test__samples_above_weight_threshold_from():
+
+    model = af.Collection(mock_class=af.m.MockClassx4)
+
+    parameters = [
+        [1.0, 2.0, 3.0, 4.0],
+        [5.0, 6.0, 7.0, 8.0],
+        [1.0, 2.0, 3.0, 4.0],
+        [1.0, 2.0, 3.0, 4.0],
+        [1.1, 2.1, 3.1, 4.1],
+    ]
+
+    samples_x5 = af.m.MockSamples(
+        model=model,
+        sample_list=af.Sample.from_lists(
+            model=model,
+            parameter_lists=parameters,
+            log_likelihood_list=[0.0, 0.0, 0.0, 0.0, 0.0],
+            log_prior_list=[0.0, 0.0, 0.0, 0.0, 0.0],
+            weight_list=[0.2, 0.2, 1.0, 1.0, 1.0],
+        ),
+    )
+
+    samples_above_weight_threshold = samples_x5.samples_above_weight_threshold_from(
+        weight_threshold=0.5
+    )
+
+    assert len(samples_above_weight_threshold) == 3
+    assert samples_above_weight_threshold.sample_list[0].weight == 1.0
+
 def test__addition_of_samples(samples_x5):
     samples = samples_x5 + samples_x5
 
diff --git a/test_autofit/text/test_text_util.py b/test_autofit/text/test_text_util.py
index 465b48796..532c5f51e 100644
--- a/test_autofit/text/test_text_util.py
+++ b/test_autofit/text/test_text_util.py
@@ -57,6 +57,7 @@ def test__search_summary_to_file(model):
         ),
         samples_info={
             "time": "1",
+            "total_accepted_samples" : 2
         }
     )
 
@@ -82,6 +83,7 @@ def test__search_summary_to_file(model):
         ),
         samples_info={
             "total_samples": 10,
+            "total_accepted_samples": 2,
             "time": "1",
             "number_live_points": 1,
             "log_evidence": 1.0