hpcflow 0.1.15__py3-none-any.whl → 0.2.0a271__py3-none-any.whl

hpcflow/sdk/core/loop.py

@@ -0,0 +1,1320 @@
+"""
+A looping construct for a workflow.
+There are multiple types of loop,
+notably looping over a set of values or until a condition holds.
+"""
+
+from __future__ import annotations
+
+from collections import defaultdict
+import copy
+from pprint import pp
+import pprint
+from typing import Dict, List, Optional, Tuple, Union, Any
+from warnings import warn
+from collections import defaultdict
+from itertools import chain
+from typing import cast, TYPE_CHECKING
+from typing_extensions import override
+
+from hpcflow.sdk.core.app_aware import AppAware
+from hpcflow.sdk.core.actions import EARStatus
+from hpcflow.sdk.core.skip_reason import SkipReason
+from hpcflow.sdk.core.errors import LoopTaskSubsetError
+from hpcflow.sdk.core.json_like import ChildObjectSpec, JSONLike
+from hpcflow.sdk.core.loop_cache import LoopCache, LoopIndex
+from hpcflow.sdk.core.enums import InputSourceType, TaskSourceType
+from hpcflow.sdk.core.utils import check_valid_py_identifier, nth_key, nth_value
+from hpcflow.sdk.utils.strings import shorten_list_str
+from hpcflow.sdk.log import TimeIt
+
+if TYPE_CHECKING:
+    from collections.abc import Iterable, Iterator, Mapping, Sequence
+    from typing import Any, ClassVar
+    from typing_extensions import Self, TypeIs
+    from rich.status import Status
+    from ..typing import DataIndex, ParamSource
+    from .parameters import SchemaInput, InputSource
+    from .rule import Rule
+    from .task import WorkflowTask
+    from .types import IterableParam
+    from .workflow import Workflow, WorkflowTemplate
+
+
+# @dataclass
+# class StoppingCriterion:
+#     parameter: Parameter
+#     condition: ConditionLike
+
+
+# @dataclass
+# class Loop:
+#     parameter: Parameter
+#     stopping_criteria: StoppingCriterion
+#     # TODO: should be a logical combination of these (maybe provide a superclass in valida to re-use some logic there?)
+#     maximum_iterations: int
+
+
+class Loop(JSONLike):
+    """
+    A loop in a workflow template.
+
+    Parameters
+    ----------
+    tasks: list[int | ~hpcflow.app.WorkflowTask]
+        List of task insert IDs or workflow tasks.
+    num_iterations:
+        Number of iterations to perform.
+    name: str
+        Loop name.
+    non_iterable_parameters: list[str]
+        Specify input parameters that should not iterate.
+    termination: v~hpcflow.app.Rule
+        Stopping criterion, expressed as a rule.
+    termination_task: int | ~hpcflow.app.WorkflowTask
+        Task at which to evaluate the termination condition.
+    """
+
+    _child_objects: ClassVar[tuple[ChildObjectSpec, ...]] = (
+        ChildObjectSpec(name="termination", class_name="Rule"),
+    )
+
+    @classmethod
+    def __is_WorkflowTask(cls, value) -> TypeIs[WorkflowTask]:
+        return isinstance(value, cls._app.WorkflowTask)
+
+    def __init__(
+        self,
+        tasks: Iterable[int | str | WorkflowTask],
+        num_iterations: int,
+        name: str | None = None,
+        non_iterable_parameters: list[str] | None = None,
+        termination: Rule | None = None,
+        termination_task: int | str | WorkflowTask | None = None,
+    ) -> None:
+        _task_refs: list[int | str] = [
+            ref_i.insert_ID if self.__is_WorkflowTask(ref_i) else ref_i for ref_i in tasks
+        ]
+
+        if termination_task is None:
+            _term_task_ref = _task_refs[-1]  # terminate on final task by default
+        elif self.__is_WorkflowTask(termination_task):
+            _term_task_ref = termination_task.insert_ID
+        else:
+            _term_task_ref = termination_task
+
+        if _term_task_ref not in _task_refs:
+            raise ValueError(
+                f"If specified, `termination_task` (provided: {termination_task!r}) must "
+                f"refer to a task that is part of the loop. Available task references "
+                f"are: {_task_refs!r}."
+            )
+
+        self._task_refs = _task_refs
+        self._task_insert_IDs = (
+            cast("list[int]", _task_refs)
+            if all(isinstance(ref_i, int) for ref_i in _task_refs)
+            else None
+        )
+
+        self._num_iterations = num_iterations
+        self._name = check_valid_py_identifier(name) if name else name
+        self._non_iterable_parameters = non_iterable_parameters or []
+        self._termination = termination
+        self._termination_task_ref = _term_task_ref
+        self._termination_task_insert_ID = (
+            _term_task_ref if isinstance(_term_task_ref, int) else None
+        )
+
+        self._workflow_template: WorkflowTemplate | None = (
+            None  # assigned by parent WorkflowTemplate
+        )
+
+    @override
+    def _postprocess_to_dict(self, d: dict[str, Any]) -> dict[str, Any]:
+        out = super()._postprocess_to_dict(d)
+        return {k.lstrip("_"): v for k, v in out.items()}
+
+    @classmethod
+    def _json_like_constructor(cls, json_like: dict) -> Self:
+        """Invoked by `JSONLike.from_json_like` instead of `__init__`."""
+        if "task_insert_IDs" in json_like:
+            insert_IDs = json_like.pop("task_insert_IDs")
+        else:
+            insert_IDs = json_like.pop("tasks")  # e.g. from YAML
+
+        if "termination_task_insert_ID" in json_like:
+            tt_iID = json_like.pop("termination_task_insert_ID")
+        elif "termination_task" in json_like:
+            tt_iID = json_like.pop("termination_task")  # e.g. from YAML
+        else:
+            tt_iID = None
+
+        task_refs = None
+        term_task_ref = None
+        if "task_refs" in json_like:
+            task_refs = json_like.pop("task_refs")
+        if "termination_task_ref" in json_like:
+            term_task_ref = json_like.pop("termination_task_ref")
+
+        obj = cls(tasks=insert_IDs, termination_task=tt_iID, **json_like)
+        if task_refs:
+            obj._task_refs = task_refs
+        if term_task_ref is not None:
+            obj._termination_task_ref = term_task_ref
+
+        return obj
+
+    @property
+    def task_refs(self) -> tuple[int | str, ...]:
+        """Get the list of task references (insert IDs or task unique names) that define
+        the extent of the loop."""
+        return tuple(self._task_refs)
+
+    @property
+    def task_insert_IDs(self) -> tuple[int, ...]:
+        """Get the list of task insert_IDs that define the extent of the loop."""
+        assert self._task_insert_IDs
+        return tuple(self._task_insert_IDs)
+
+    @property
+    def name(self) -> str | None:
+        """
+        The name of the loop, if one was provided.
+        """
+        return self._name
+
+    @property
+    def num_iterations(self) -> int:
+        """
+        The number of loop iterations to do.
+        """
+        return self._num_iterations
+
+    @property
+    def non_iterable_parameters(self) -> Sequence[str]:
+        """
+        Which parameters are not iterable.
+        """
+        return self._non_iterable_parameters
+
+    @property
+    def termination(self) -> Rule | None:
+        """
+        A termination rule for the loop, if one is provided.
+        """
+        return self._termination
+
+    @property
+    def termination_task_ref(self) -> int | str:
+        """
+        The unique name of the task at which the loop will terminate.
+        """
+        return self._termination_task_ref
+
+    @property
+    def termination_task_insert_ID(self) -> int:
+        """
+        The insert ID of the task at which the loop will terminate.
+        """
+        assert self._termination_task_insert_ID is not None
+        return self._termination_task_insert_ID
+
+    @property
+    def termination_task(self) -> WorkflowTask:
+        """
+        The task at which the loop will terminate.
+        """
+        if (wt := self.workflow_template) is None:
+            raise RuntimeError(
+                "Workflow template must be assigned to retrieve task objects of the loop."
+            )
+        assert wt.workflow
+        return wt.workflow.tasks.get(insert_ID=self.termination_task_insert_ID)
+
+    @property
+    def workflow_template(self) -> WorkflowTemplate | None:
+        """
+        The workflow template that contains this loop.
+        """
+        return self._workflow_template
+
+    @workflow_template.setter
+    def workflow_template(self, template: WorkflowTemplate):
+        self._workflow_template = template
+
+    def __workflow(self) -> None | Workflow:
+        if (wt := self.workflow_template) is None:
+            return None
+        return wt.workflow
+
+    @property
+    def task_objects(self) -> tuple[WorkflowTask, ...]:
+        """
+        The tasks in the loop.
+        """
+        if not (wf := self.__workflow()):
+            raise RuntimeError(
+                "Workflow template must be assigned to retrieve task objects of the loop."
+            )
+        return tuple(wf.tasks.get(insert_ID=t_id) for t_id in self.task_insert_IDs)
+
+    def _validate_against_workflow(self, workflow: Workflow) -> None:
+        """Validate the loop parameters against the associated workflow."""
+
+        names = workflow.get_task_unique_names(map_to_insert_ID=True)
+        if self._task_insert_IDs is None:
+            err = lambda ref_i: (
+                f"Loop {self.name + ' ' if self.name else ''}has an invalid "
+                f"task reference {ref_i!r}. Such a task does not exist in "
+                f"the associated workflow, which has task names/insert IDs: "
+                f"{names!r}."
+            )
+            self._task_insert_IDs = []
+            for ref_i in self._task_refs:
+                if isinstance(ref_i, str):
+                    # reference is task unique name, so retrieve the insert ID:
+                    try:
+                        iID = names[ref_i]
+                    except KeyError:
+                        raise ValueError(err(ref_i)) from None
+                else:
+                    try:
+                        workflow.tasks.get(insert_ID=ref_i)
+                    except ValueError:
+                        raise ValueError(err(ref_i)) from None
+                    iID = ref_i
+                self._task_insert_IDs.append(iID)
+
+        if self._termination_task_insert_ID is None:
+            tt_ref = self._termination_task_ref
+            if isinstance(tt_ref, str):
+                tt_iID = names[tt_ref]
+            else:
+                tt_iID = tt_ref
+            self._termination_task_insert_ID = tt_iID
+
+    def __repr__(self) -> str:
+        num_iterations_str = ""
+        if self.num_iterations is not None:
+            num_iterations_str = f", num_iterations={self.num_iterations!r}"
+
+        name_str = ""
+        if self.name:
+            name_str = f", name={self.name!r}"
+
+        return (
+            f"{self.__class__.__name__}("
+            f"task_insert_IDs={self.task_insert_IDs!r}{num_iterations_str}{name_str}"
+            f")"
+        )
+
+    def __deepcopy__(self, memo: dict[int, Any]) -> Self:
+        kwargs = self.to_dict()
+        kwargs["tasks"] = kwargs.pop("task_insert_IDs")
+        kwargs["termination_task"] = kwargs.pop("termination_task_insert_ID")
+        task_refs = kwargs.pop("task_refs", None)
+        term_task_ref = kwargs.pop("termination_task_ref", None)
+        obj = self.__class__(**copy.deepcopy(kwargs, memo))
+        obj._workflow_template = self._workflow_template
+        obj._task_refs = task_refs
+        obj._termination_task_ref = term_task_ref
+        return obj
+
+
+class WorkflowLoop(AppAware):
+    """
+    Class to represent a :py:class:`.Loop` that is bound to a
+    :py:class:`~hpcflow.app.Workflow`.
+
+    Parameters
+    ----------
+    index: int
+        The index of this loop in the workflow.
+    workflow: ~hpcflow.app.Workflow
+        The workflow containing this loop.
+    template: Loop
+        The loop that this was generated from.
+    num_added_iterations:
+        Description of what iterations have been added.
+    iterable_parameters:
+        Description of what parameters are being iterated over.
+    output_parameters:
+        Description of what parameter are output from this loop, and the final task insert
+        ID from which they are output.
+    parents: list[str]
+        The paths to the parent entities of this loop.
+    """
+
+    def __init__(
+        self,
+        index: int,
+        workflow: Workflow,
+        template: Loop,
+        num_added_iterations: dict[tuple[int, ...], int],
+        iterable_parameters: dict[str, IterableParam],
+        output_parameters: dict[str, int],
+        parents: list[str],
+    ) -> None:
+        self._index = index
+        self._workflow = workflow
+        self._template = template
+        self._num_added_iterations = num_added_iterations
+        self._iterable_parameters = iterable_parameters
+        self._output_parameters = output_parameters
+        self._parents = parents
+
+        # appended to when adding an empty loop to the workflow that is a parent of this
+        # loop; reset and added to `self._parents` on dump to disk:
+        self._pending_parents: list[str] = []
+
+        # used for `num_added_iterations` when a new loop iteration is added, or when
+        # parents are append to; reset to None on dump to disk. Each key is a tuple of
+        # parent loop indices and each value is the number of pending new iterations:
+        self._pending_num_added_iterations: dict[tuple[int, ...], int] | None = None
+
+        self._validate()
+
+    @TimeIt.decorator
+    def _validate(self) -> None:
+        # task subset must be a contiguous range of task indices:
+        task_indices = self.task_indices
+        task_min, task_max = task_indices[0], task_indices[-1]
+        if task_indices != tuple(range(task_min, task_max + 1)):
+            raise LoopTaskSubsetError(self.name, self.task_indices)
+
+    def __repr__(self) -> str:
+        return (
+            f"{self.__class__.__name__}(template={self.template!r}, "
+            f"num_added_iterations={self.num_added_iterations!r})"
+        )
+
+    @property
+    def num_added_iterations(self) -> Mapping[tuple[int, ...], int]:
+        """
+        The number of added iterations.
+        """
+        if self._pending_num_added_iterations:
+            return self._pending_num_added_iterations
+        else:
+            return self._num_added_iterations
+
+    @property
+    def __pending(self) -> dict[tuple[int, ...], int]:
+        if not self._pending_num_added_iterations:
+            self._pending_num_added_iterations = dict(self._num_added_iterations)
+        return self._pending_num_added_iterations
+
+    def _initialise_pending_added_iters(self, added_iters: Iterable[int]):
+        if not self._pending_num_added_iterations:
+            self._pending_num_added_iterations = dict(self._num_added_iterations)
+        if (added_iters_key := tuple(added_iters)) not in (pending := self.__pending):
+            pending[added_iters_key] = 1
+
+    def _increment_pending_added_iters(self, added_iters_key: Iterable[int]):
+        self.__pending[tuple(added_iters_key)] += 1
+
+    def _update_parents(self, parent: WorkflowLoop):
+        assert parent.name
+        self._pending_parents.append(parent.name)
+
+        self._pending_num_added_iterations = {
+            (*k, 0): v
+            for k, v in (
+                self._pending_num_added_iterations or self._num_added_iterations
+            ).items()
+        }
+
+        self.workflow._store.update_loop_parents(
+            index=self.index,
+            num_added_iters=self.num_added_iterations,
+            parents=self.parents,
+        )
+
+    def _reset_pending_num_added_iters(self) -> None:
+        self._pending_num_added_iterations = None
+
+    def _accept_pending_num_added_iters(self) -> None:
+        if self._pending_num_added_iterations:
+            self._num_added_iterations = dict(self._pending_num_added_iterations)
+            self._reset_pending_num_added_iters()
+
+    def _reset_pending_parents(self) -> None:
+        self._pending_parents = []
+
+    def _accept_pending_parents(self) -> None:
+        self._parents += self._pending_parents
+        self._reset_pending_parents()
+
+    @property
+    def index(self) -> int:
+        """
+        The index of this loop within its workflow.
+        """
+        return self._index
+
+    @property
+    def task_refs(self) -> tuple[int | str, ...]:
+        """
+        The list of task references (insert IDs or task unique names) that define the
+        extent of the loop."""
+        return self.template.task_refs
+
+    @property
+    def task_insert_IDs(self) -> tuple[int, ...]:
+        """
+        The insertion IDs of the tasks inside this loop.
+        """
+        return self.template.task_insert_IDs
+
+    @property
+    def task_objects(self) -> tuple[WorkflowTask, ...]:
+        """
+        The tasks in this loop.
+        """
+        return self.template.task_objects
+
+    @property
+    def task_indices(self) -> tuple[int, ...]:
+        """
+        The list of task indices that define the extent of the loop.
+        """
+        return tuple(task.index for task in self.task_objects)
+
+    @property
+    def workflow(self) -> Workflow:
+        """
+        The workflow containing this loop.
+        """
+        return self._workflow
+
+    @property
+    def template(self) -> Loop:
+        """
+        The loop template for this loop.
+        """
+        return self._template
+
+    @property
+    def parents(self) -> Sequence[str]:
+        """
+        The parents of this loop.
+        """
+        return self._parents + self._pending_parents
+
+    @property
+    def name(self) -> str:
+        """
+        The name of this loop, if one is defined.
+        """
+        assert self.template.name
+        return self.template.name
+
+    @property
+    def iterable_parameters(self) -> dict[str, IterableParam]:
+        """
+        The parameters that are being iterated over.
+        """
+        return self._iterable_parameters
+
+    @property
+    def output_parameters(self) -> dict[str, int]:
+        """
+        The parameters that are outputs of this loop, and the final task insert ID from
+        which each parameter is output.
+        """
+        return self._output_parameters
+
+    @property
+    def num_iterations(self) -> int:
+        """
+        The number of iterations.
+        """
+        return self.template.num_iterations
+
+    @property
+    def downstream_tasks(self) -> Iterator[WorkflowTask]:
+        """Tasks that are not part of the loop, and downstream from this loop."""
+        tasks = self.workflow.tasks
+        for idx in range(self.task_objects[-1].index + 1, len(tasks)):
+            yield tasks[idx]
+
+    @property
+    def upstream_tasks(self) -> Iterator[WorkflowTask]:
+        """Tasks that are not part of the loop, and upstream from this loop."""
+        tasks = self.workflow.tasks
+        for idx in range(0, self.task_objects[0].index):
+            yield tasks[idx]
+
+    @staticmethod
+    @TimeIt.decorator
+    def _find_iterable_and_output_parameters(
+        loop_template: Loop,
+    ) -> tuple[dict[str, IterableParam], dict[str, int]]:
+        all_inputs_first_idx: dict[str, int] = {}
+        all_outputs_idx: dict[str, list[int]] = defaultdict(list)
+        for task in loop_template.task_objects:
+            for typ in task.template.all_schema_input_types:
+                all_inputs_first_idx.setdefault(typ, task.insert_ID)
+            for typ in task.template.all_schema_output_types:
+                all_outputs_idx[typ].append(task.insert_ID)
+
+        # find input parameters that are also output parameters at a later/same task:
+        iterable_params: dict[str, IterableParam] = {}
+        for typ, first_idx in all_inputs_first_idx.items():
+            if typ in all_outputs_idx and first_idx <= all_outputs_idx[typ][0]:
+                iterable_params[typ] = {
+                    "input_task": first_idx,
+                    "output_tasks": all_outputs_idx[typ],
+                }
+
+        for non_iter in loop_template.non_iterable_parameters:
+            iterable_params.pop(non_iter, None)
+
+        final_out_tasks = {k: v[-1] for k, v in all_outputs_idx.items()}
+
+        return iterable_params, final_out_tasks
+
+    @classmethod
+    @TimeIt.decorator
+    def new_empty_loop(
+        cls,
+        index: int,
+        workflow: Workflow,
+        template: Loop,
+        iter_loop_idx: Sequence[Mapping[str, int]],
+    ) -> WorkflowLoop:
+        """
+        Make a new empty loop.
+
+        Parameters
+        ----------
+        index: int
+            The index of the loop to create.
+        workflow: ~hpcflow.app.Workflow
+            The workflow that will contain the loop.
+        template: Loop
+            The template for the loop.
+        iter_loop_idx: list[dict]
+            Iteration information from parent loops.
+        """
+        parent_loops = cls._get_parent_loops(index, workflow, template)
+        parent_names = [i.name for i in parent_loops]
+        num_added_iters: dict[tuple[int, ...], int] = {}
+        for i in iter_loop_idx:
+            num_added_iters[tuple([i[j] for j in parent_names])] = 1
+
+        iter_params, out_params = cls._find_iterable_and_output_parameters(template)
+        return cls(
+            index=index,
+            workflow=workflow,
+            template=template,
+            num_added_iterations=num_added_iters,
+            iterable_parameters=iter_params,
+            output_parameters=out_params,
+            parents=parent_names,
+        )
+
+    @classmethod
+    @TimeIt.decorator
+    def _get_parent_loops(
+        cls,
+        index: int,
+        workflow: Workflow,
+        template: Loop,
+    ) -> list[WorkflowLoop]:
+        parents: list[WorkflowLoop] = []
+        passed_self = False
+        self_tasks = set(template.task_insert_IDs)
+        for loop_i in workflow.loops:
+            if loop_i.index == index:
+                passed_self = True
+                continue
+            other_tasks = set(loop_i.task_insert_IDs)
+            if self_tasks.issubset(other_tasks):
+                if (self_tasks == other_tasks) and not passed_self:
+                    continue
+                parents.append(loop_i)
+        return parents
+
+    @TimeIt.decorator
+    def get_parent_loops(self) -> list[WorkflowLoop]:
+        """Get loops whose task subset is a superset of this loop's task subset. If two
+        loops have identical task subsets, the first loop in the workflow loop list is
+        considered the child."""
+        return self._get_parent_loops(self.index, self.workflow, self.template)
+
+    @TimeIt.decorator
+    def get_child_loops(self) -> list[WorkflowLoop]:
+        """Get loops whose task subset is a subset of this loop's task subset. If two
+        loops have identical task subsets, the first loop in the workflow loop list is
+        considered the child."""
+        children: list[WorkflowLoop] = []
+        passed_self = False
+        self_tasks = set(self.task_insert_IDs)
+        for loop_i in self.workflow.loops:
+            if loop_i.index == self.index:
+                passed_self = True
+                continue
+            other_tasks = set(loop_i.task_insert_IDs)
+            if self_tasks.issuperset(other_tasks):
+                if (self_tasks == other_tasks) and passed_self:
+                    continue
+                children.append(loop_i)
+
+        # order by depth, so direct child is first:
+        return sorted(children, key=lambda x: len(next(iter(x.num_added_iterations))))
+
+    @TimeIt.decorator
+    def add_iteration(
+        self,
+        parent_loop_indices: Mapping[str, int] | None = None,
+        cache: LoopCache | None = None,
+        status: Status | None = None,
+    ) -> None:
+        """
+        Add an iteration to this loop.
+
+        Parameters
+        ----------
+        parent_loop_indices:
+            Where have any parent loops got up to?
+        cache:
+            A cache used to make adding the iteration more efficient.
+            One will be created if it is not supplied.
+        """
+        if not cache:
+            cache = LoopCache.build(self.workflow)
+        assert cache is not None
+        parent_loops = self.get_parent_loops()
+        child_loops = self.get_child_loops()
+        parent_loop_indices_ = parent_loop_indices or {
+            loop.name: 0 for loop in parent_loops
+        }
+
+        iters_key = tuple(parent_loop_indices_[p_nm] for p_nm in self.parents)
+        cur_loop_idx = self.num_added_iterations[iters_key] - 1
+
+        # keys are (task.insert_ID and element.index)
+        all_new_data_idx: dict[tuple[int, int], DataIndex] = {}
+
+        # initialise a new `num_added_iterations` key on each child loop:
+        iters_key_dct = {
+            **parent_loop_indices_,
+            self.name: cur_loop_idx + 1,
+        }
+        for child in child_loops:
+            child._initialise_pending_added_iters(
+                iters_key_dct.get(j, 0) for j in child.parents
+            )
+
+            # needed for the case where an inner loop has only one iteration, meaning
+            # `add_iteration` will not be called recursively on it:
+            self.workflow._store.update_loop_num_iters(
+                index=child.index,
+                num_added_iters=child.num_added_iterations,
+            )
+
+        for task in self.task_objects:
+            new_loop_idx = LoopIndex(iters_key_dct) + {
+                child.name: 0
+                for child in child_loops
+                if task.insert_ID in child.task_insert_IDs
+            }
+            added_iter_IDs: list[int] = []
+            for elem_idx in range(task.num_elements):
+                elem_ID = task.element_IDs[elem_idx]
+
+                new_data_idx: DataIndex = {}
+
+                # copy resources from zeroth iteration:
+                zeroth_iter_ID, zi_iter_data_idx = cache.zeroth_iters[elem_ID]
+                zi_elem_ID, zi_idx = cache.iterations[zeroth_iter_ID]
+                zi_data_idx = nth_value(cache.data_idx[zi_elem_ID], zi_idx)
+
+                for key, val in zi_data_idx.items():
+                    if key.startswith("resources."):
+                        new_data_idx[key] = val
+
+                for inp in task.template.all_schema_inputs:
+                    is_inp_task = False
+                    if iter_dat := self.iterable_parameters.get(inp.typ):
+                        is_inp_task = task.insert_ID == iter_dat["input_task"]
+
+                    inp_key = f"inputs.{inp.typ}"
+
+                    if is_inp_task:
+                        assert iter_dat is not None
+                        inp_dat_idx = self.__get_looped_index(
+                            task,
+                            elem_ID,
+                            cache,
+                            iter_dat,
+                            inp,
+                            parent_loops,
+                            parent_loop_indices_,
+                            child_loops,
+                            cur_loop_idx,
+                        )
+                        new_data_idx[inp_key] = inp_dat_idx
+                    else:
+                        orig_inp_src = cache.elements[elem_ID]["input_sources"][inp_key]
+                        inp_dat_idx = None
+
+                        if orig_inp_src.source_type is InputSourceType.LOCAL:
+                            # keep locally defined inputs from original element
+                            inp_dat_idx = zi_data_idx[inp_key]
+
+                        elif orig_inp_src.source_type is InputSourceType.DEFAULT:
+                            # keep default value from original element
+                            try:
+                                inp_dat_idx = zi_data_idx[inp_key]
+                            except KeyError:
+                                # if this input is required by a conditional action, and
+                                # that condition is not met, then this input will not
+                                # exist in the action-run data index, so use the initial
+                                # iteration data index:
+                                inp_dat_idx = zi_iter_data_idx[inp_key]
+
+                        elif orig_inp_src.source_type is InputSourceType.TASK:
+                            inp_dat_idx = self.__get_task_index(
+                                task,
+                                orig_inp_src,
+                                cache,
+                                elem_ID,
+                                inp,
+                                inp_key,
+                                parent_loop_indices_,
+                                all_new_data_idx,
+                            )
+
+                        if inp_dat_idx is None:
+                            raise RuntimeError(
+                                f"Could not find a source for parameter {inp.typ} "
+                                f"when adding a new iteration for task {task!r}."
+                            )
+
+                        new_data_idx[inp_key] = inp_dat_idx
+
+                # add any locally defined sub-parameters:
+                inp_statuses = cache.elements[elem_ID]["input_statuses"]
+                inp_status_inps = set(f"inputs.{inp}" for inp in inp_statuses)
+                for sub_param_i in inp_status_inps.difference(new_data_idx):
+                    sub_param_data_idx_iter_0 = zi_data_idx
+                    try:
+                        sub_param_data_idx = sub_param_data_idx_iter_0[sub_param_i]
+                    except KeyError:
+                        # as before, if this input is required by a conditional action,
+                        # and that condition is not met, then this input will not exist in
+                        # the action-run data index, so use the initial iteration data
+                        # index:
+                        sub_param_data_idx = zi_data_idx[sub_param_i]
+
+                    new_data_idx[sub_param_i] = sub_param_data_idx
+
+                for out in task.template.all_schema_outputs:
+                    path_i = f"outputs.{out.typ}"
+                    p_src: ParamSource = {"type": "EAR_output"}
+                    new_data_idx[path_i] = self.workflow._add_unset_parameter_data(p_src)
+
+                schema_params = set(i for i in new_data_idx if len(i.split(".")) == 2)
+                all_new_data_idx[task.insert_ID, elem_idx] = new_data_idx
+
+                iter_ID_i = self.workflow._store.add_element_iteration(
+                    element_ID=elem_ID,
+                    data_idx=new_data_idx,
+                    schema_parameters=list(schema_params),
+                    loop_idx=new_loop_idx,
+                )
+                if cache:
+                    cache.add_iteration(
+                        iter_ID=iter_ID_i,
+                        task_insert_ID=task.insert_ID,
+                        element_ID=elem_ID,
+                        loop_idx=new_loop_idx,
+                        data_idx=new_data_idx,
+                    )
+
+                added_iter_IDs.append(iter_ID_i)
+
+            task.initialise_EARs(iter_IDs=added_iter_IDs)
+
+        self._increment_pending_added_iters(
+            parent_loop_indices_[p_nm] for p_nm in self.parents
+        )
+        self.workflow._store.update_loop_num_iters(
+            index=self.index,
+            num_added_iters=self.num_added_iterations,
+        )
+
+        # add iterations to fixed-number-iteration children only:
+        for child in child_loops[::-1]:
+            if child.num_iterations is not None:
+                if status:
+                    status_prev = str(status.status).rstrip(".")
+                for iter_idx in range(child.num_iterations - 1):
+                    if status:
+                        status.update(
+                            f"{status_prev} --> ({child.name!r}): iteration "
+                            f"{iter_idx + 2}/{child.num_iterations}."
+                        )
+                    par_idx = {parent_name: 0 for parent_name in child.parents}
+                    if parent_loop_indices:
+                        par_idx.update(parent_loop_indices)
+                    par_idx[self.name] = cur_loop_idx + 1
+                    child.add_iteration(parent_loop_indices=par_idx, cache=cache)
+
+        self.__update_loop_downstream_data_idx(parent_loop_indices_)
+
+    def __get_src_ID_and_groups(
+        self,
+        elem_ID: int,
+        iter_dat: IterableParam,
+        inp: SchemaInput,
+        cache: LoopCache,
+        task: WorkflowTask,
+    ) -> tuple[int, Sequence[int]]:
+        # `cache.elements` contains only elements that are part of the
+        # loop, so indexing a dependent element may raise:
+        src_elem_IDs = {}
+        for k, v in cache.element_dependents[elem_ID].items():
+            try:
+                if cache.elements[k]["task_insert_ID"] == iter_dat["output_tasks"][-1]:
+                    src_elem_IDs[k] = v
+            except KeyError:
+                continue
+
+        # consider groups
+        single_data = inp.single_labelled_data
+        assert single_data is not None
+        inp_group_name = single_data.get("group")
+        grouped_elems = [
+            src_elem_j_ID
+            for src_elem_j_ID, src_elem_j_dat in src_elem_IDs.items()
+            if any(nm == inp_group_name for nm in src_elem_j_dat["group_names"])
+        ]
+
+        if not grouped_elems and len(src_elem_IDs) > 1:
+            raise NotImplementedError(
+                f"Multiple elements found in the iterable parameter "
+                f"{inp!r}'s latest output task (insert ID: "
+                f"{iter_dat['output_tasks'][-1]}) that can be used "
+                f"to parametrise the next iteration of task "
+                f"{task.unique_name!r}: "
+                f"{list(src_elem_IDs)!r}."
+            )
+
+        elif not src_elem_IDs:
+            # TODO: maybe OK?
+            raise NotImplementedError(
+                f"No elements found in the iterable parameter "
+                f"{inp!r}'s latest output task (insert ID: "
+                f"{iter_dat['output_tasks'][-1]}) that can be used "
+                f"to parametrise the next iteration."
+            )
+
+        return nth_key(src_elem_IDs, 0), grouped_elems
+
+    def __get_looped_index(
+        self,
+        task: WorkflowTask,
+        elem_ID: int,
+        cache: LoopCache,
+        iter_dat: IterableParam,
+        inp: SchemaInput,
+        parent_loops: list[WorkflowLoop],
+        parent_loop_indices: Mapping[str, int],
+        child_loops: list[WorkflowLoop],
+        cur_loop_idx: int,
+    ):
+        # source from final output task of previous iteration, with all parent
+        # loop indices the same as previous iteration, and all child loop indices
+        # maximised:
+
+        # identify element(s) from which this iterable input should be
+        # parametrised:
+        if task.insert_ID == iter_dat["output_tasks"][-1]:
+            # single-task loop
+            src_elem_ID = elem_ID
+            grouped_elems: Sequence[int] = []
+        else:
+            # multi-task loop
+            src_elem_ID, grouped_elems = self.__get_src_ID_and_groups(
+                elem_ID, iter_dat, inp, cache, task
+            )
+
+        child_loop_max_iters: dict[str, int] = {}
+        parent_loop_same_iters = {
+            loop.name: parent_loop_indices[loop.name] for loop in parent_loops
+        }
+        child_iter_parents = {
+            **parent_loop_same_iters,
+            self.name: cur_loop_idx,
+        }
+        for loop in child_loops:
+            if iter_dat["output_tasks"][-1] in loop.task_insert_IDs:
+                i_num_iters = loop.num_added_iterations[
+                    tuple(child_iter_parents[j] for j in loop.parents)
+                ]
+                i_max = i_num_iters - 1
+                child_iter_parents[loop.name] = i_max
+                child_loop_max_iters[loop.name] = i_max
+
+        loop_idx_key = LoopIndex(child_loop_max_iters)
+        loop_idx_key.update(parent_loop_same_iters)
+        loop_idx_key[self.name] = cur_loop_idx
+
+        # identify the ElementIteration from which this input should be
+        # parametrised:
+        if grouped_elems:
+            src_data_idx = [
+                cache.data_idx[src_elem_ID][loop_idx_key] for src_elem_ID in grouped_elems
+            ]
+            if not src_data_idx:
+                raise RuntimeError(
+                    f"Could not find a source iteration with loop_idx: "
+                    f"{loop_idx_key!r}."
+                )
+            return [i[f"outputs.{inp.typ}"] for i in src_data_idx]
+        else:
+            return cache.data_idx[src_elem_ID][loop_idx_key][f"outputs.{inp.typ}"]
+
+    def __get_task_index(
+        self,
+        task: WorkflowTask,
+        orig_inp_src: InputSource,
+        cache: LoopCache,
+        elem_ID: int,
+        inp: SchemaInput,
+        inp_key: str,
+        parent_loop_indices: Mapping[str, int],
+        all_new_data_idx: Mapping[tuple[int, int], DataIndex],
+    ) -> int | list[int]:
+        if orig_inp_src.task_ref not in self.task_insert_IDs:
+            # source the data_idx from the iteration with same parent
+            # loop indices as the new iteration to add:
+            src_data_idx = next(
+                di_k
+                for li_k, di_k in cache.data_idx[elem_ID].items()
+                if all(li_k.get(p_k) == p_v for p_k, p_v in parent_loop_indices.items())
+            )
+
+            # could be multiple, but they should all have the same
+            # data index for this parameter:
+            return src_data_idx[inp_key]
+
+        is_group = (
+            inp.single_labelled_data is not None
+            and "group" in inp.single_labelled_data
+            # this input is a group, assume for now all elements
+        )
+
+        # same task/element, but update iteration to the just-added
+        # iteration:
+        assert orig_inp_src.task_source_type is not None
+        key_prefix = orig_inp_src.task_source_type.name.lower()
+        prev_dat_idx_key = f"{key_prefix}s.{inp.typ}"
+        new_sources: list[tuple[int, int]] = []
+        for (tiID, e_idx), _ in all_new_data_idx.items():
+            if tiID == orig_inp_src.task_ref:
+                # find which element in that task `element`
+                # depends on:
+                src_elem_IDs = cache.element_dependents[
+                    self.workflow.tasks.get(insert_ID=tiID).element_IDs[e_idx]
+                ]
+                # `cache.elements` contains only elements that are part of the loop, so
+                # indexing a dependent element may raise:
+                src_elem_IDs_i = []
+                for k, _v in src_elem_IDs.items():
+                    try:
+                        if (
+                            cache.elements[k]["task_insert_ID"] == task.insert_ID
+                            and k == elem_ID
+                            # filter src_elem_IDs_i for matching element IDs
+                        ):
+
+                            src_elem_IDs_i.append(k)
+                    except KeyError:
+                        continue
+
+                if len(src_elem_IDs_i) == 1:
+                    new_sources.append((tiID, e_idx))
+
+        if is_group:
+            # Convert into simple list of indices
+            return list(
+                chain.from_iterable(
+                    self.__as_sequence(all_new_data_idx[src][prev_dat_idx_key])
+                    for src in new_sources
+                )
+            )
+        else:
+            assert len(new_sources) == 1
+            return all_new_data_idx[new_sources[0]][prev_dat_idx_key]
+
+    @staticmethod
+    def __as_sequence(seq: int | Iterable[int]) -> Iterable[int]:
+        if isinstance(seq, int):
+            yield seq
+        else:
+            yield from seq
+
+    def __update_loop_downstream_data_idx(
+        self,
+        parent_loop_indices: Mapping[str, int],
+    ):
+        # update data indices of loop-downstream tasks that depend on task outputs from
+        # this loop:
+
+        # keys: iter or run ID, values: dict of param type and new parameter index
+        iter_new_data_idx: dict[int, DataIndex] = defaultdict(dict)
+        run_new_data_idx: dict[int, DataIndex] = defaultdict(dict)
+
+        param_sources = self.workflow.get_all_parameter_sources()
+
+        # keys are parameter type, then task insert ID, then data index keys mapping to
+        # their updated values:
+        all_updates: dict[str, dict[int, dict[int, int]]] = defaultdict(
+            lambda: defaultdict(dict)
+        )
+
+        for task in self.downstream_tasks:
+            for elem in task.elements:
+                for param_typ, param_out_task_iID in self.output_parameters.items():
+                    if param_typ in task.template.all_schema_input_types:
+                        # this element's input *might* need updating, only if it has a
+                        # task input source type that is this loop's output task for this
+                        # parameter:
+                        elem_src = elem.input_sources[f"inputs.{param_typ}"]
+                        if (
+                            elem_src.source_type is InputSourceType.TASK
+                            and elem_src.task_source_type is TaskSourceType.OUTPUT
+                            and elem_src.task_ref == param_out_task_iID
+                        ):
+                            for iter_i in elem.iterations:
+
+                                # do not modify element-iterations of previous iterations
+                                # of the current loop:
+                                skip_iter = False
+                                for k, v in parent_loop_indices.items():
+                                    if iter_i.loop_idx.get(k) != v:
+                                        skip_iter = True
+                                        break
+
+                                if skip_iter:
+                                    continue
+
+                                # update the iteration data index and any pending runs:
+                                iter_old_di = iter_i.data_idx[f"inputs.{param_typ}"]
+
+                                is_group = True
+                                if not isinstance(iter_old_di, list):
+                                    is_group = False
+                                    iter_old_di = [iter_old_di]
+
+                                iter_old_run_source = [
+                                    param_sources[i]["EAR_ID"] for i in iter_old_di
+                                ]
+                                iter_old_run_objs = self.workflow.get_EARs_from_IDs(
+                                    iter_old_run_source
+                                )  # TODO: use cache
+
+                                # need to check the run source is actually from the loop
+                                # output task (it could be from a previous iteration of a
+                                # separate loop in this task):
+                                if any(
+                                    i.task.insert_ID != param_out_task_iID
+                                    for i in iter_old_run_objs
+                                ):
+                                    continue
+
+                                iter_new_iters = [
+                                    i.element.iterations[-1] for i in iter_old_run_objs
+                                ]
+
+                                # note: we can cast to int, because output keys never
+                                # have multiple data indices (unlike input keys):
+                                iter_new_dis = [
+                                    cast("int", i.get_data_idx()[f"outputs.{param_typ}"])
+                                    for i in iter_new_iters
+                                ]
+
+                                # keep track of updates so we can also update task-input
+                                # type sources:
+                                all_updates[param_typ][task.insert_ID].update(
+                                    dict(zip(iter_old_di, iter_new_dis))
+                                )
+
+                                iter_new_data_idx[iter_i.id_][f"inputs.{param_typ}"] = (
+                                    iter_new_dis if is_group else iter_new_dis[0]
+                                )
+
+                                for run_j in iter_i.action_runs:
+                                    if run_j.status is EARStatus.pending:
+                                        try:
+                                            old_di = run_j.data_idx[f"inputs.{param_typ}"]
+                                        except KeyError:
+                                            # not all actions will include this input
+                                            continue
+
+                                        is_group = True
+                                        if not isinstance(old_di, list):
+                                            is_group = False
+                                            old_di = [old_di]
+
+                                        old_run_source = [
+                                            param_sources[i]["EAR_ID"] for i in old_di
+                                        ]
+                                        old_run_objs = self.workflow.get_EARs_from_IDs(
+                                            old_run_source
+                                        )  # TODO: use cache
+
+                                        # need to check the run source is actually from the loop
+                                        # output task (it could be from a previous action in this
+                                        # element-iteration):
+                                        if any(
+                                            i.task.insert_ID != param_out_task_iID
+                                            for i in old_run_objs
+                                        ):
+                                            continue
+
+                                        new_iters = [
+                                            i.element.iterations[-1] for i in old_run_objs
+                                        ]
+
+                                        # note: we can cast to int, because output keys
+                                        # never have multiple data indices (unlike input
+                                        # keys):
+                                        new_dis = [
+                                            cast(
+                                                "int",
+                                                i.get_data_idx()[f"outputs.{param_typ}"],
+                                            )
+                                            for i in new_iters
+                                        ]
+
+                                        run_new_data_idx[run_j.id_][
+                                            f"inputs.{param_typ}"
+                                        ] = (new_dis if is_group else new_dis[0])
+
+                        elif (
+                            elem_src.source_type is InputSourceType.TASK
+                            and elem_src.task_source_type is TaskSourceType.INPUT
+                        ):
+                            # parameters are that sourced from inputs of other tasks,
+                            # might need to be updated if those other tasks have
+                            # themselves had their data indices updated:
+                            assert elem_src.task_ref
+                            ups_i = all_updates.get(param_typ, {}).get(elem_src.task_ref)
+
+                            if ups_i:
+                                # if a further-downstream task has a task-input source
+                                # that points to this task, this will also need updating:
+                                all_updates[param_typ][task.insert_ID].update(ups_i)
+
+                            else:
+                                continue
+
+                            for iter_i in elem.iterations:
+
+                                # update the iteration data index and any pending runs:
+                                iter_old_di = iter_i.data_idx[f"inputs.{param_typ}"]
+
+                                is_group = True
+                                if not isinstance(iter_old_di, list):
+                                    is_group = False
+                                    iter_old_di = [iter_old_di]
+
+                                iter_new_dis = [ups_i.get(i, i) for i in iter_old_di]
+
+                                if iter_new_dis != iter_old_di:
+                                    iter_new_data_idx[iter_i.id_][
+                                        f"inputs.{param_typ}"
+                                    ] = (iter_new_dis if is_group else iter_new_dis[0])
+
+                                for run_j in iter_i.action_runs:
+                                    if run_j.status is EARStatus.pending:
+                                        try:
+                                            old_di = run_j.data_idx[f"inputs.{param_typ}"]
+                                        except KeyError:
+                                            # not all actions will include this input
+                                            continue
+
+                                        is_group = True
+                                        if not isinstance(old_di, list):
+                                            is_group = False
+                                            old_di = [old_di]
+
+                                        new_dis = [ups_i.get(i, i) for i in old_di]
+
+                                        if new_dis != old_di:
+                                            run_new_data_idx[run_j.id_][
+                                                f"inputs.{param_typ}"
+                                            ] = (new_dis if is_group else new_dis[0])
+
+        # now update data indices (TODO: including in cache!)
+        if iter_new_data_idx:
+            self.workflow._store.update_iter_data_indices(iter_new_data_idx)
+
+        if run_new_data_idx:
+            self.workflow._store.update_run_data_indices(run_new_data_idx)
+
+    def test_termination(self, element_iter) -> bool:
+        """Check if a loop should terminate, given the specified completed element
+        iteration."""
+        if self.template.termination:
+            return self.template.termination.test(element_iter)
+        return False
+
+    @TimeIt.decorator
+    def get_element_IDs(self):
+        elem_IDs = [
+            j
+            for i in self.task_insert_IDs
+            for j in self.workflow.tasks.get(insert_ID=i).element_IDs
+        ]
+        return elem_IDs
+
+    @TimeIt.decorator
+    def get_elements(self):
+        return self.workflow.get_elements_from_IDs(self.get_element_IDs())
+
+    @TimeIt.decorator
+    def skip_downstream_iterations(self, elem_iter) -> list[int]:
+        """
+        Parameters
+        ----------
+        elem_iter
+            The element iteration whose subsequent iterations should be skipped.
+        dep_element_IDs
+            List of elements that are dependent (recursively) on the element
+            of `elem_iter`.
+        """
+        current_iter_idx = elem_iter.loop_idx[self.name]
+        current_task_iID = elem_iter.task.insert_ID
+        self._app.logger.info(
+            f"setting loop {self.name!r} iterations downstream of current iteration "
+            f"index {current_iter_idx} to skip"
+        )
+        elements = self.get_elements()
+
+        # TODO: fix for multiple loop cycles
+        warn(
+            "skip downstream iterations does not work correctly for multiple loop cycles!"
+        )
+
+        to_skip = []
+        for elem in elements:
+            for iter_i in elem.iterations:
+                if iter_i.loop_idx[self.name] > current_iter_idx or (
+                    iter_i.loop_idx[self.name] == current_iter_idx
+                    and iter_i.task.insert_ID > current_task_iID
+                ):
+                    to_skip.extend(iter_i.EAR_IDs_flat)
+
+        self._app.logger.info(
+            f"{len(to_skip)} runs will be set to skip: {shorten_list_str(to_skip)}"
+        )
+        self.workflow.set_EAR_skip({k: SkipReason.LOOP_TERMINATION for k in to_skip})
+
+        return to_skip