fix: if multiple generators, use coupled trial type generator

Author: arjunsridhar12345

src/dynamic_foraging_processing/processing/_trial_table.py

@@ -10,6 +10,8 @@
 
 import numpy as np
 import pandas as pd
+from aind_behavior_dynamic_foraging.task_logic import AindDynamicForagingTaskLogic
+from aind_behavior_dynamic_foraging.task_logic.trial_generators import TrialGeneratorSpec
 from aind_behavior_dynamic_foraging.task_logic.trial_models import Trial, TrialOutcome
 from aind_behavior_services.task.distributions import Distribution, DistributionFamily
 from contraqctor.contract import Dataset
@@ -274,7 +276,42 @@ def _auto_water(trial: t.Optional[Trial], *, side: bool) -> t.Optional[int]:
     # ------------------------------------------------------------------ #
     # Session-level (constant across trials) columns
     # ------------------------------------------------------------------ #
-    def _session_columns(self, task_logic: t.Optional[t.Any]) -> t.Dict[str, t.Any]:
+    @staticmethod
+    def _coupled_generator(
+        generator: TrialGeneratorSpec,
+    ) -> t.Optional[TrialGeneratorSpec]:
+        """Resolve the ``CoupledTrialGenerator`` carrying the session parameters.
+
+        ``task_parameters.trial_generator`` may be a single generator spec or a
+        ``TrialGeneratorCompositeSpec`` wrapping several sub-generators in
+        ``.generators``. The block-length, ITI, delay, and reward-probability
+        summaries come from the coupled generator; this unwraps the composite
+        and returns it, or returns ``generator`` unchanged when it is not a
+        composite.
+
+        Parameters
+        ----------
+        generator : TrialGeneratorSpec
+            The ``trial_generator`` from the task parameters.
+
+        Returns
+        -------
+        TrialGeneratorSpec or None
+            The coupled generator, or ``None`` if a composite contains none.
+        """
+        sub_generators = getattr(generator, "generators", None)
+        if sub_generators is None:
+            return generator
+        for sub in sub_generators:
+            if getattr(sub, "type", None) == "CoupledTrialGenerator":
+                return sub
+        logger.warning(
+            "No CoupledTrialGenerator among composite trial generators; "
+            "session columns will be null."
+        )
+        return None
+
+    def _session_columns(self, task_logic: AindDynamicForagingTaskLogic) -> t.Dict[str, t.Any]:
         """Return the per-session trial columns derived from the task logic.
 
         These (block/ITI/delay distribution summaries and reward structure) are
@@ -284,7 +321,9 @@ def _session_columns(self, task_logic: t.Optional[t.Any]) -> t.Dict[str, t.Any]:
         if task_logic is None:
             return columns
 
-        generator = task_logic.task_parameters.trial_generator
+        generator = self._coupled_generator(task_logic.task_parameters.trial_generator)
+        if generator is None:
+            return columns
 
         block_beta, block_min, block_max = self._distribution_stats(generator.block_length)
         iti_beta, iti_min, iti_max = self._distribution_stats(

tests/test_processing/test_trial_table.py

@@ -9,6 +9,8 @@
 )
 from aind_behavior_dynamic_foraging.task_logic.trial_generators import (
     CoupledTrialGeneratorSpec,
+    TrialGeneratorCompositeSpec,
+    UncoupledTrialGeneratorSpec,
 )
 from aind_behavior_dynamic_foraging.task_logic.trial_models import TrialOutcome
 from aind_behavior_services.task.distributions import (
@@ -224,6 +226,41 @@ def test_build_exponential_quiescent_sets_delay_beta():
     assert table.iloc[0]["delay_max"] == 1.0
 
 
+# --------------------------------------------------------------------------- #
+# _coupled_generator — composite trial generators
+# --------------------------------------------------------------------------- #
+def test_coupled_generator_returns_single_generator_unchanged():
+    """A non-composite generator (no ``.generators``) is returned as-is."""
+    spec = _task_logic().task_parameters.trial_generator
+    assert TrialTableBuilder._coupled_generator(spec) is spec
+
+
+def test_coupled_generator_unwraps_composite():
+    """The coupled sub-generator is selected from a composite."""
+    coupled = _task_logic().task_parameters.trial_generator
+    composite = TrialGeneratorCompositeSpec(generators=[UncoupledTrialGeneratorSpec(), coupled])
+    resolved = TrialTableBuilder._coupled_generator(composite)
+    assert resolved.type == "CoupledTrialGenerator"
+    assert resolved.min_block_reward == coupled.min_block_reward
+
+
+def test_coupled_generator_none_when_composite_has_no_coupled(caplog):
+    """A composite without a coupled generator yields ``None`` and warns."""
+    composite = TrialGeneratorCompositeSpec(generators=[UncoupledTrialGeneratorSpec()])
+    assert TrialTableBuilder._coupled_generator(composite) is None
+    assert "No CoupledTrialGenerator" in caplog.text
+
+
+def test_session_columns_empty_when_no_coupled_generator():
+    """A composite task logic with no coupled generator yields no session columns."""
+    task_logic = AindDynamicForagingTaskLogic(
+        task_parameters=AindDynamicForagingTaskParameters(
+            trial_generator=TrialGeneratorCompositeSpec(generators=[UncoupledTrialGeneratorSpec()])
+        )
+    )
+    assert TrialTableBuilder(_Node({}))._session_columns(task_logic) == {}
+
+
 # --------------------------------------------------------------------------- #
 # _load — error handling
 # --------------------------------------------------------------------------- #