fix: CTIAgg on multi-dataset fits, and honour the input instance

autocti/aggregator/abstract.py

@@ -32,3 +32,27 @@ def __init__(
 
         self.use_dataset_full = use_dataset_full
         self.clocker_list = clocker_list
+
+
+def _cti_list_from(source, total_datasets: int):
+    """
+    Extract one CTI model per dataset from a model instance.
+
+    A single-analysis instance exposes ``instance.cti`` directly; a factor-graph
+    instance (multi-dataset fit) is an indexed collection with one child
+    instance per factor.
+    """
+    if hasattr(source, "cti"):
+        return [source.cti] * total_datasets
+
+    # A factor-graph instance also carries the FactorGraphModel itself as a
+    # trailing child, so only children with a CTI model are taken.
+    cti_list = [child.cti for child in source if hasattr(child, "cti")]
+
+    if len(cti_list) != total_datasets:
+        raise ValueError(
+            f"The instance contains {len(cti_list)} CTI models but the fit has "
+            f"{total_datasets} datasets."
+        )
+
+    return cti_list

autocti/aggregator/cti.py

@@ -1,57 +1,48 @@
 from __future__ import annotations
-import logging
-from typing import TYPE_CHECKING, Optional, Union
+from typing import TYPE_CHECKING, List, Optional, Union
 
 from autocti.aggregator.abstract import AggBase
+from autocti.aggregator.abstract import _cti_list_from
 
 if TYPE_CHECKING:
     from autocti.model.model_util import CTI1D
     from autocti.model.model_util import CTI2D
 
 import autofit as af
 
-logger = logging.getLogger(__name__)
 
-
-def _cti_from(fit: af.Fit) -> Union[CTI1D, CTI2D]:
+def _cti_from(
+    fit: af.Fit, instance: Optional[af.ModelInstance] = None
+) -> Union[CTI1D, CTI2D, List[Union[CTI1D, CTI2D]]]:
     """
-    Returns a list of `CTI` objects from a `PyAutoFit` sqlite database `Fit` object.
-
-    The results of a model-fit can be stored in a sqlite database, including the following attributes of the fit:
+    Returns the `CTI` model(s) of a `PyAutoFit` sqlite database `Fit` object.
 
-    - The model and its best fit parameters (e.g. `model.json`).
+    The results of a model-fit can be stored in a sqlite database, including the model and its best fit
+    parameters (e.g. `model.json`), which this function turns back into `CTI` objects for a given non-linear
+    search sample (e.g. the maximum likelihood model).
 
-    Each individual attribute can be loaded from the database via the `fit.value()` method.
-
-    This method combines all of these attributes and returns a `CTI` object for a given non-linear search sample
-    (e.g. the maximum likelihood model). This includes associating adapt images with their respective galaxies.
-
-    If multiple `CTI` objects were fitted simultaneously via analysis summing, the `fit.child_values()` method
-    is instead used to load lists of Tracers. This is necessary if each Tracer has different galaxies (e.g. certain
-    parameters vary across each dataset and `Analysis` object).
+    A fit to a single dataset returns one `CTI` object. A fit to multiple datasets, made with a
+    `FactorGraphModel`, has one child instance per dataset, and returns a list with one `CTI` object per dataset
+    (in the same order as the datasets were fitted), as the fit aggregators do.
 
     Parameters
     ----------
     fit
         A PyAutoFit database Fit object containing the generators of the results of model-fits.
+    instance
+        A manual instance that overwrites the max log likelihood instance in fit (e.g. for drawing the instance
+        randomly from the PDF).
 
     Returns
     -------
-    CTI
-        The list of CTI models computed via an instance of clocker.
+    The `CTI` model, or the list of `CTI` models of a multi-dataset fit.
     """
+    instance = instance if instance is not None else fit.instance
 
-    if len(fit.children) > 0:
-        logger.info(
-            """
-            Using database for a fit with multiple summed Analysis objects.
-
-            CTI objects do not fully support this yet (e.g. adapt images may not be set up correctly)
-            so proceed with caution!
-            """
-        )
+    if not fit.children:
+        return instance.cti
 
-    return fit.instance.cti
+    return _cti_list_from(source=instance, total_datasets=len(fit.children))
 
 
 class CTIAgg(AggBase):
@@ -73,9 +64,8 @@ class CTIAgg(AggBase):
     For example, if the `aggregator` contains 3 model-fits, this class can be used to create a generator which
     creates instances of the corresponding 3 `CTI` objects.
 
-    If multiple `CTI` objects were fitted simultaneously via analysis summing, the `fit.child_values()` method
-    is instead used to load lists of CTIs. This is necessary if each CTI has different galaxies (e.g. certain
-    parameters vary across each dataset and `Analysis` object).
+    A fit to a single dataset gives one `CTI` object per fit. A fit to multiple datasets at once, made with a
+    `FactorGraphModel`, gives a list with one `CTI` object per dataset (in the order the datasets were fitted).
 
     This can be done manually, but this object provides a more concise API.
 
@@ -87,9 +77,9 @@ class CTIAgg(AggBase):
 
     def object_via_gen_from(
         self, fit, instance: Optional[af.ModelInstance] = None
-    ) -> Union[CTI1D, CTI2D]:
+    ) -> Union[CTI1D, CTI2D, List[Union[CTI1D, CTI2D]]]:
         """
-        Returns a generator of `CTI` objects from an input aggregator.
+        Returns the `CTI` object(s) of a model-fit in the aggregator.
 
         See `__init__` for a description of how the `CTI` objects are created by this method.
 
@@ -103,7 +93,7 @@ def object_via_gen_from(
 
         Returns
         -------
-        CTI
-            A list of `CTI` objects whose parameters are a sample of the non-linear search.
+        The `CTI` object whose parameters are a sample of the non-linear search, or for a multi-dataset
+        (`FactorGraphModel`) fit a list with one `CTI` object per dataset.
         """
-        return _cti_from(fit=fit)
+        return _cti_from(fit=fit, instance=instance)

autocti/aggregator/fit_dataset_1d.py

@@ -2,6 +2,7 @@
 from typing import TYPE_CHECKING, List, Optional
 
 from autocti.aggregator.abstract import AggBase
+from autocti.aggregator.abstract import _cti_list_from
 
 if TYPE_CHECKING:
     from autocti.clocker.abstract import AbstractClocker
@@ -12,30 +13,6 @@
 from autocti.aggregator.dataset_1d import _dataset_1d_list_from
 
 
-def _cti_list_from(source, total_datasets: int):
-    """
-    Extract one CTI model per dataset from a model instance.
-
-    A single-analysis instance exposes ``instance.cti`` directly; a factor-graph
-    instance (multi-dataset fit) is an indexed collection with one child
-    instance per factor.
-    """
-    if hasattr(source, "cti"):
-        return [source.cti] * total_datasets
-
-    # A factor-graph instance also carries the FactorGraphModel itself as a
-    # trailing child, so only children with a CTI model are taken.
-    cti_list = [child.cti for child in source if hasattr(child, "cti")]
-
-    if len(cti_list) != total_datasets:
-        raise ValueError(
-            f"The instance contains {len(cti_list)} CTI models but the fit has "
-            f"{total_datasets} datasets."
-        )
-
-    return cti_list
-
-
 def _fit_dataset_1d_list_from(
     fit: af.Fit,
     instance: Optional[af.ModelInstance] = None,

autocti/aggregator/fit_imaging_ci.py

@@ -10,33 +10,10 @@
 import autofit as af
 
 from autocti.aggregator.abstract import AggBase
+from autocti.aggregator.abstract import _cti_list_from
 from autocti.aggregator.imaging_ci import _imaging_ci_list_from
 
 
-def _cti_list_from(source, total_datasets: int):
-    """
-    Extract one CTI model per dataset from a model instance.
-
-    A single-analysis instance exposes ``instance.cti`` directly; a factor-graph
-    instance (multi-dataset fit) is an indexed collection with one child
-    instance per factor.
-    """
-    if hasattr(source, "cti"):
-        return [source.cti] * total_datasets
-
-    # A factor-graph instance also carries the FactorGraphModel itself as a
-    # trailing child, so only children with a CTI model are taken.
-    cti_list = [child.cti for child in source if hasattr(child, "cti")]
-
-    if len(cti_list) != total_datasets:
-        raise ValueError(
-            f"The instance contains {len(cti_list)} CTI models but the fit has "
-            f"{total_datasets} datasets."
-        )
-
-    return cti_list
-
-
 def _fit_imaging_ci_list_from(
     fit: af.Fit,
     instance: Optional[af.ModelInstance] = None,

test_autocti/aggregator/test_aggregator_cti.py

@@ -59,12 +59,69 @@ def test__cti_all_above_weight_gen(dataset_1d_7, clocker_1d, samples_1d, model_1
         for cti in cti_gen:
             i += 1
 
+            # Each above-weight sample's own instance (samples_1d: all 1.0, then
+            # all 10.0), not the max likelihood instance repeated.
             if i == 1:
-                assert cti.trap_list[0].density == pytest.approx(10.0, 1.0e-4)
+                assert cti.trap_list[0].density == pytest.approx(1.0, 1.0e-4)
 
             if i == 2:
                 assert cti.trap_list[0].density == pytest.approx(10.0, 1.0e-4)
 
     assert i == 2
 
     clean(database_file=database_file)
+
+
+def test__cti_max_log_likelihood_gen__factor_graph_returns_one_cti_per_dataset(
+    dataset_1d_7, clocker_1d, samples_1d, model_1d
+):
+    # A multi-dataset fit is a FactorGraphModel: its instance has one child per
+    # dataset rather than a top-level `cti` attribute.
+    analysis = ac.AnalysisDataset1D(dataset=dataset_1d_7, clocker=clocker_1d)
+
+    agg = aggregator_from(
+        database_file=database_file,
+        analysis=[analysis, analysis],
+        model=model_1d,
+        samples=samples_1d,
+    )
+
+    cti_agg = ac.agg.CTIAgg(aggregator=agg)
+
+    i = 0
+
+    for cti_list in cti_agg.max_log_likelihood_gen_from():
+        i += 1
+
+        assert len(cti_list) == 2
+
+        for cti in cti_list:
+            assert cti.trap_list[0].density == pytest.approx(10.0, 1.0e-4)
+
+    assert i == 1
+
+    clean(database_file=database_file)
+
+
+def test__cti_object_via_gen_from__uses_input_instance(
+    dataset_1d_7, clocker_1d, samples_1d, model_1d
+):
+    analysis = ac.AnalysisDataset1D(dataset=dataset_1d_7, clocker=clocker_1d)
+
+    agg = aggregator_from(
+        database_file=database_file,
+        analysis=analysis,
+        model=model_1d,
+        samples=samples_1d,
+    )
+
+    instance = model_1d.instance_from_prior_medians()
+    instance.cti.trap_list[0].density = 123.0
+
+    cti_agg = ac.agg.CTIAgg(aggregator=agg)
+
+    cti = cti_agg.object_via_gen_from(fit=list(agg)[0], instance=instance)
+
+    assert cti.trap_list[0].density == pytest.approx(123.0, 1.0e-4)
+
+    clean(database_file=database_file)