junifer 0.0.7.dev332__py3-none-any.whl → 0.0.8__py3-none-any.whl

junifer/datagrabber/pattern.py

@@ -5,154 +5,158 @@
 #          Synchon Mandal <s.mandal@fz-juelich.de>
 # License: AGPL
 
+import glob
 import re
+from collections import defaultdict
 from copy import deepcopy
+from enum import Enum
 from pathlib import Path
-from typing import Optional, Union
+from typing import ClassVar
 
-import numpy as np
+from aenum import Enum as AEnum
+from aenum import extend_enum
+from pydantic import Field
 
 from ..api.decorators import register_datagrabber
 from ..typing import DataGrabberPatterns, Elements
-from ..utils import logger, raise_error
-from .base import BaseDataGrabber
+from ..utils import ensure_list, raise_error
+from .base import BaseDataGrabber, logger
 from .pattern_validation_mixin import PatternValidationMixin
 
 
-__all__ = ["PatternDataGrabber"]
+__all__ = [
+    "ConfoundsFormat",
+    "PatternDataGrabber",
+    "register_confounds_format",
+]
 
 
-# Accepted formats for confounds specification
-_CONFOUNDS_FORMATS = ("fmriprep", "adhoc")
+class ConfoundsFormat(str, AEnum):
+    """Accepted confounds format."""
 
+    FMRIPrep = "fmriprep"
+    AdHoc = "adhoc"
 
-@register_datagrabber
-class PatternDataGrabber(BaseDataGrabber, PatternValidationMixin):
-    """Concrete implementation for pattern-based data fetching.
 
-    Implements a DataGrabber that understands patterns to grab data.
+def register_confounds_format(name: str, alias: str) -> None:
+    """Register custom confounds format.
 
     Parameters
     ----------
-    types : list of str
-        The types of data to be grabbed.
-    patterns : dict
-        Data type patterns as a dictionary. It has the following schema:
-
-        * ``"T1w"`` :
-
-          .. code-block:: none
-
-            {
-              "mandatory": ["pattern", "space"],
-              "optional": {
-                  "mask": {
-                      "mandatory": ["pattern", "space"],
-                      "optional": []
-                  }
-              }
-            }
-
-        * ``"T2w"`` :
+    name : str
+        The confounds format name to be referred.
+    alias : str
+        The confounds format alias for string representation.
 
-          .. code-block:: none
+    """
+    extend_enum(ConfoundsFormat, name, alias)
 
-            {
-              "mandatory": ["pattern", "space"],
-              "optional": {
-                  "mask": {
-                      "mandatory": ["pattern", "space"],
-                      "optional": []
-                  }
-              }
-            }
 
-        * ``"BOLD"`` :
-
-          .. code-block:: none
-
-            {
-              "mandatory": ["pattern", "space"],
-              "optional": {
-                  "mask": {
-                      "mandatory": ["pattern", "space"],
-                      "optional": []
-                  }
-                  "confounds": {
-                      "mandatory": ["pattern", "format"],
-                      "optional": []
-                  }
-              }
-            }
+def _value(value: object) -> object:
+    """Get the value of an enum member, or the value itself.
 
-        * ``"Warp"`` :
+    Parameters
+    ----------
+    value : object
+        The value.
 
-          .. code-block:: none
+    Returns
+    -------
+    object
+        The value.
 
-            {
-              "mandatory": ["pattern", "src", "dst", "warper"],
-              "optional": []
-            }
-
-        * ``"VBM_GM"`` :
+    """
+    return value.value if isinstance(value, Enum) else value
+
+
+def _join_elements(
+    left: set[tuple[str, ...]],
+    left_keys: list[str],
+    right: set[tuple[str, ...]],
+    right_keys: list[str],
+) -> tuple[set[tuple[str, ...]], list[str]]:
+    """Join two sets of elements on their common keys.
+
+    Each data type (or pattern) has its own elements, with the values of the
+    replacements in its pattern. The elements of the datagrabber are the ones
+    available for all the data types, so the elements of two data types are
+    joined like database tables: an element of the first set is combined with
+    each element of the second set that has the same values for the common
+    keys, adding the values of the keys that only the second set has. An
+    element without a match in the other set is dropped.
+
+    For example, with BOLD elements ``(subject, session, task)``:
+    ``(sub-01, ses-1, rest)``, ``(sub-01, ses-1, movie)`` and
+    ``(sub-02, ses-1, rest)``:
+
+    * VBM elements ``(subject)``: ``(sub-01)`` and ``(sub-03)`` keep the BOLD
+      elements of ``sub-01``, as the VBM data only depends on the subject.
+    * DWI elements ``(subject, run)``: ``(sub-01, 1)`` and ``(sub-01, 2)``
+      give the elements ``(subject, session, task, run)`` for each BOLD
+      element of ``sub-01`` with each run, i.e., all the combinations.
+
+    The result is the same regardless of which set is first (except for the
+    order of the keys), so the order of the data types does not matter.
 
-          .. code-block:: none
+    Parameters
+    ----------
+    left : set of tuple of str
+        The first elements, as tuples of the values of ``left_keys``.
+    left_keys : list of str
+        The keys of the first elements.
+    right : set of tuple of str
+        The second elements, as tuples of the values of ``right_keys``.
+    right_keys : list of str
+        The keys of the second elements.
+
+    Returns
+    -------
+    set of tuple of str
+        The elements with the values of the first and the second elements that
+        have the same values for the common keys.
+    list of str
+        The keys of the elements: ``left_keys`` and the other ``right_keys``.
 
-            {
-              "mandatory": ["pattern", "space"],
-              "optional": []
-            }
+    """
+    common = [k for k in right_keys if k in left_keys]
+    other = [i for i, k in enumerate(right_keys) if k not in left_keys]
+    # Values of the other keys of the second elements, by the common values
+    right_by_common = defaultdict(list)
+    for element in right:
+        common_values = tuple(element[right_keys.index(k)] for k in common)
+        right_by_common[common_values].append(tuple(element[i] for i in other))
+    joined = {
+        element + other_values
+        for element in left
+        for other_values in right_by_common[
+            tuple(element[left_keys.index(k)] for k in common)
+        ]
+    }
+    return joined, left_keys + [right_keys[i] for i in other]
 
-        * ``"VBM_WM"`` :
 
-          .. code-block:: none
+@register_datagrabber
+class PatternDataGrabber(BaseDataGrabber, PatternValidationMixin):
+    """Concrete implementation for pattern-based data fetching.
 
-            {
-              "mandatory": ["pattern", "space"],
-              "optional": []
-            }
+    Implements a DataGrabber that understands patterns to grab data.
 
-        Basically, for each data type, one needs to provide ``mandatory`` keys
-        and can choose to also provide ``optional`` keys. The value for each
-        key is a string. So, one needs to provide necessary data types as a
-        dictionary, for example:
-
-        .. code-block:: none
-
-          {
-              "BOLD": {
-                "pattern": "...",
-                "space": "...",
-              },
-              "T1w": {
-                "pattern": "...",
-                "space": "...",
-              },
-          }
-
-        except ``Warp``, which needs to be a list of dictionaries as there can
-        be multiple spaces to warp (for example, with fMRIPrep):
-
-        .. code-block:: none
-
-          {
-              "Warp": [
-                {
-                  "pattern": "...",
-                  "src": "...",
-                  "dst": "...",
-                  "warper": "...",
-                },
-              ],
-          }
-
-        taken from :class:`.HCP1200`.
-    replacements : str or list of str
-        Replacements in the ``pattern`` key of each data type. The value needs
-        to be a list of all possible replacements.
-    datadir : str or pathlib.Path
-        The directory where the data is / will be stored.
-    confounds_format : {"fmriprep", "adhoc"} or None, optional
+    Parameters
+    ----------
+    types : :enum:`.DataType` or list of variants
+        The data type(s) to grab.
+    datadir : pathlib.Path
+        The path where the data is stored.
+    patterns : ``DataGrabberPatterns``
+        The datagrabber patterns. Check :class:`.DataTypeSchema` for the \
+        schema.
+    replacements : list of str or dict
+        All possible replacements in ``patterns.<data_type>.pattern``. As a
+        dictionary, the keys are the replacements and the values are the lists
+        of values to grab (or None to grab all of them), e.g.,
+        ``{"subject": None, "task": ["rest", "movie"]}`` to grab all the
+        subjects and only these tasks.
+    confounds_format : :enum:`.ConfoundsFormat` or None, optional
         The format of the confounds for the dataset (default None).
     partial_pattern_ok : bool, optional
         Whether to raise error if partial pattern for a data type is found.
@@ -162,52 +166,57 @@ class PatternDataGrabber(BaseDataGrabber, PatternValidationMixin):
         powerful when used with :class:`.MultipleDataGrabber`
         (default True).
 
-    Raises
-    ------
-    ValueError
-        If ``confounds_format`` is invalid.
+    Attributes
+    ----------
+    skip_file_check
 
     """
 
-    def __init__(
-        self,
-        types: list[str],
-        patterns: DataGrabberPatterns,
-        replacements: Union[list[str], str],
-        datadir: Union[str, Path],
-        confounds_format: Optional[str] = None,
-        partial_pattern_ok: bool = False,
-    ) -> None:
-        # Convert replacements to list if not already
-        if not isinstance(replacements, list):
-            replacements = [replacements]
+    patterns: DataGrabberPatterns = Field(frozen=True)
+    replacements: list[str] | dict[str, list[str] | None] = Field(frozen=True)
+
+    # Fields with the values of the replacements to grab, by replacement, e.g.,
+    # ``{"task": "tasks"}``
+    _REPLACEMENT_FIELDS: ClassVar[dict[str, str]] = {}
+    confounds_format: ConfoundsFormat | None = Field(None, frozen=True)
+    partial_pattern_ok: bool = Field(False, frozen=True)
+
+    def validate_datagrabber_params(self) -> None:
+        """Run extra logical validation for datagrabber."""
         # Validate patterns
         self.validate_patterns(
-            types=types,
-            replacements=replacements,
-            patterns=patterns,
-            partial_pattern_ok=partial_pattern_ok,
+            types=self.types,
+            replacements=list(self.replacements),
+            patterns=self.patterns,
+            partial_pattern_ok=self.partial_pattern_ok,
         )
-        self.replacements = replacements
-        self.patterns = patterns
-        self.partial_pattern_ok = partial_pattern_ok
-
-        # Validate confounds format
-        if (
-            confounds_format is not None
-            and confounds_format not in _CONFOUNDS_FORMATS
-        ):
-            raise_error(
-                "Invalid value for `confounds_format`, should be one of "
-                f"{_CONFOUNDS_FORMATS}."
-            )
-        self.confounds_format = confounds_format
-
-        super().__init__(types=types, datadir=datadir)
+        # Validate the fields with the values of the replacements
+        for replacement, field in self._REPLACEMENT_FIELDS.items():
+            if replacement not in self.replacements:
+                raise_error(
+                    msg=(
+                        f"Replacement: `{replacement}` of the field `{field}` "
+                        f"is not one of the replacements: "
+                        f"{list(self.replacements)}"
+                    ),
+                    klass=ValueError,
+                )
+            if (
+                isinstance(self.replacements, dict)
+                and self.replacements[replacement] is not None
+            ):
+                raise_error(
+                    msg=(
+                        f"The values of the replacement `{replacement}` are "
+                        f"given by the field `{field}`, they cannot be given "
+                        "in `replacements`"
+                    ),
+                    klass=ValueError,
+                )
         logger.debug("Initializing PatternDataGrabber")
-        logger.debug(f"\tpatterns = {patterns}")
-        logger.debug(f"\treplacements = {replacements}")
-        logger.debug(f"\tconfounds_format = {confounds_format}")
+        logger.debug(f"\tpatterns = {self.patterns}")
+        logger.debug(f"\treplacements = {self.replacements}")
+        logger.debug(f"\tconfounds_format = {self.confounds_format}")
 
     @property
     def skip_file_check(self) -> bool:
@@ -242,18 +251,32 @@ class PatternDataGrabber(BaseDataGrabber, PatternValidationMixin):
             x for x in self.replacements if f"{{{x}}}" in pattern
         ]
         # Ops on re_pattern
-        # Remove negated unix glob pattern i.e., [!...] for re_pattern
+        # Remove negated unix glob pattern i.e., [!...] for re_pattern, as the
+        # character is part of the following replacement (e.g., [!f]{subject})
         re_pattern = re.sub(r"\[!.?\]", "", re_pattern)
         # Remove enclosing square brackets from unix glob pattern i.e., [...]
         # for re_pattern
         re_pattern = re.sub(r"\[|\]", "", re_pattern)
-        # Iteratively replace the first of each with a named group definition
-        for t_r in t_replacements:
-            re_pattern = re_pattern.replace(f"{{{t_r}}}", f"(?P<{t_r}>.*)", 1)
-        # Iteratively replace the second appearance of each with the named
-        # group back reference
-        for t_r in t_replacements:
-            re_pattern = re_pattern.replace(f"{{{t_r}}}", f"(?P={t_r})")
+        # Translate the replacements to named groups (the first appearance of
+        # each) or back references (the others), the unix glob wildcards to
+        # their regular expressions and escape the rest
+        parts = []
+        seen = set()
+        for token in re.split(r"(\{[^{}]*\}|\*|\?)", re_pattern):
+            name = token[1:-1]
+            if token.startswith("{") and name in t_replacements:
+                if name in seen:
+                    parts.append(f"(?P={name})")
+                else:
+                    parts.append(f"(?P<{name}>.*)")
+                    seen.add(name)
+            elif token == "*":
+                parts.append("[^/]*")
+            elif token == "?":
+                parts.append("[^/]")
+            else:
+                parts.append(re.escape(token))
+        re_pattern = "".join(parts)
         # Ops on glob_pattern
         # Iteratively replace replacements with wildcard i.e., *
         # for glob_pattern
@@ -283,9 +306,9 @@ class PatternDataGrabber(BaseDataGrabber, PatternValidationMixin):
             If element keys do not match with replacements.
 
         """
-        if list(element.keys()) != self.replacements:
+        if list(element.keys()) != self.get_element_keys():
             raise_error(
-                f"The element keys must be {self.replacements}, "
+                f"The element keys must be {self.get_element_keys()}, "
                 f"element has {list(element.keys())}."
             )
         # Remove negated unix glob pattern i.e., [!...]
@@ -323,9 +346,18 @@ class PatternDataGrabber(BaseDataGrabber, PatternValidationMixin):
         """
         # Replace element in the pattern for globbing
         resolved_pattern = self._replace_patterns_glob(element, pattern)
-        # Resolve path for wildcard
-        if "*" in resolved_pattern:
-            t_matches = list(self.datadir.absolute().glob(resolved_pattern))
+        # Resolve path for wildcards
+        if "*" in resolved_pattern or "?" in resolved_pattern:
+            # glob.glob (unlike pathlib before Python 3.12) also matches the
+            # broken symbolic links, e.g., the files of a DataLad dataset that
+            # are not downloaded yet.
+            # TODO: go back to Path.glob once Python 3.11 is dropped (end of
+            # life in October 2027)
+            fulldir = self.fulldir.absolute()
+            t_matches = [
+                fulldir / x
+                for x in glob.glob(resolved_pattern, root_dir=fulldir)
+            ]
             # Multiple matches
             if len(t_matches) > 1:
                 raise_error(
@@ -341,7 +373,8 @@ class PatternDataGrabber(BaseDataGrabber, PatternValidationMixin):
                 )
             path = t_matches[0]
         else:
-            path = self.datadir / resolved_pattern
+            # Absolute path, as for the patterns with wildcards
+            path = self.fulldir.absolute() / resolved_pattern
             if not self.skip_file_check:
                 if not path.exists() and not path.is_symlink():
                     raise_error(
@@ -352,23 +385,131 @@ class PatternDataGrabber(BaseDataGrabber, PatternValidationMixin):
 
         return path
 
+    def _get_type_patterns(self, data_type: str) -> list[str]:
+        """Get the patterns of a data type, including the nested ones.
+
+        Parameters
+        ----------
+        data_type : str
+            The data type.
+
+        Returns
+        -------
+        list of str
+            The patterns.
+
+        """
+        dtype_val = self.patterns[data_type]
+        # Conditional for list dtype vals like Warp
+        if not isinstance(dtype_val, list):
+            dtype_val = [dtype_val]
+        patterns = []
+        for val in dtype_val:
+            if "pattern" in val:
+                patterns.append(val["pattern"])
+            patterns.extend(
+                v["pattern"]
+                for v in val.values()
+                if isinstance(v, dict) and "pattern" in v
+            )
+        return patterns
+
+    def get_type_element_keys(self, data_type: str) -> list[str]:
+        """Get the element keys that the data of a data type depends on.
+
+        Parameters
+        ----------
+        data_type : str
+            The data type.
+
+        Returns
+        -------
+        list of str
+            The replacements in the patterns of the data type (including the
+            nested ones), in the order of ``replacements``.
+
+        """
+        patterns = self._get_type_patterns(data_type)
+        return [
+            x
+            for x in self.replacements
+            if any(f"{{{x}}}" in pattern for pattern in patterns)
+        ]
+
     def get_element_keys(self) -> list[str]:
         """Get element keys.
 
         For each item in the "element" tuple, this functions returns the
-        corresponding key, that is, the ``replacements`` of patterns defined
-        in the constructor.
+        corresponding key: the replacements in the patterns of the data
+        types to grab (see :meth:`.get_type_element_keys`), e.g., without
+        ``task`` if only anatomical data types are grabbed.
 
         Returns
         -------
         list of str
-            The element keys.
+            The element keys, in the order of ``replacements``.
+
+        """
+        keys = {
+            k for t in self.get_types() for k in self.get_type_element_keys(t)
+        }
+        return [x for x in self.replacements if x in keys]
+
+    def get_replacement_values(self) -> dict[str, list[str]]:
+        """Get the values of the replacements to grab.
+
+        The values are given by ``replacements`` (as a dictionary) or by the
+        fields of the replacements (see ``_REPLACEMENT_FIELDS``).
+
+        Returns
+        -------
+        dict
+            The lists of values of the replacements to grab, by replacement.
+            The replacements without values (to grab all of them) are not
+            included.
 
         """
-        return self.replacements
+        values = {}
+        if isinstance(self.replacements, dict):
+            values = {
+                k: list(v)
+                for k, v in self.replacements.items()
+                if v is not None
+            }
+        for replacement, field in self._REPLACEMENT_FIELDS.items():
+            field_values = getattr(self, field)
+            if field_values is not None:
+                values[replacement] = [
+                    _value(x) for x in ensure_list(field_values)
+                ]
+        return values
+
+    def _check_element(self, element: dict) -> None:
+        """Check the element has values of the replacements to grab.
+
+        Parameters
+        ----------
+        element : dict
+            The element, as given to :meth:`.get_item`.
+
+        Raises
+        ------
+        ValueError
+            If a value of the element is not one of the values to grab.
+
+        """
+        for key, values in self.get_replacement_values().items():
+            if key in element and _value(element[key]) not in values:
+                raise_error(
+                    msg=(
+                        f"The value `{element[key]}` of `{key}` is not one of "
+                        f"the values to grab: {values}"
+                    ),
+                    klass=ValueError,
+                )
 
     def get_item(self, **element: dict) -> dict[str, dict]:
-        """Implement single element indexing for the datagrabber.
+        """Get the specified item from the dataset.
 
         This method constructs a real path to the requested item's data, by
         replacing the ``patterns`` with actual values passed via ``**element``.
@@ -454,8 +595,8 @@ class PatternDataGrabber(BaseDataGrabber, PatternValidationMixin):
         """Implement fetching list of elements in the dataset.
 
         It will use regex to search for "replacements" in the "patterns" and
-        return the intersection of the results for each type i.e., build a
-        list of elements that have all the required types.
+        join the results for each type on their common replacements i.e.,
+        build a list of elements that have all the required types.
 
         Returns
         -------
@@ -463,22 +604,10 @@ class PatternDataGrabber(BaseDataGrabber, PatternValidationMixin):
             The list of elements that can be grabbed in the dataset.
 
         """
+        # The elements found so far, as tuples of the values of `keys`
         elements = None
-
-        # Iterate by number of replacements. At least one pattern must have
-        # all of them.
-        replacement_in_patterns = [
-            np.sum([x in self.patterns[t_type] for x in self.replacements])
-            for t_type in self.types
-        ]
-        order = np.argsort(replacement_in_patterns)
-
-        # Start from the one with the most replacements and then
-        # filter out elements based on the ones with less replacements.
-        for t_idx in reversed(order):
-            t_type = self.types[t_idx]
-            types_element = set()
-
+        keys = []
+        for t_type in self.get_types():
             # Data type dictionary
             patterns = self.patterns[t_type]
             # Conditional for list dtype vals like Warp
@@ -496,46 +625,48 @@ class PatternDataGrabber(BaseDataGrabber, PatternValidationMixin):
                             pattern = v["pattern"]
                             break
 
-                # Replace the pattern
+                # Find the elements of the pattern
                 (
                     re_pattern,
                     glob_pattern,
-                    t_replacements,
+                    t_keys,
                 ) = self._replace_patterns_regex(pattern)
-                for fname in self.datadir.glob(glob_pattern):
-                    suffix = fname.relative_to(self.datadir).as_posix()
+                t_elements = set()
+                # glob.glob (unlike pathlib before Python 3.12) also matches
+                # the broken symbolic links, e.g., the files of a DataLad
+                # dataset that are not downloaded yet.
+                # TODO: go back to Path.glob once Python 3.11 is dropped (end
+                # of life in October 2027)
+                for fname in glob.glob(glob_pattern, root_dir=self.fulldir):
+                    suffix = Path(fname).as_posix()
                     m = re.match(re_pattern, suffix)
                     if m is not None:
-                        # Find the groups of replacements present in the
-                        # pattern. If one replacement is not present, set it
-                        # to None. We will take care of this in the
-                        # intersection.
-                        t_element = tuple([m.group(k) for k in t_replacements])
-                        if len(self.replacements) == 1:
-                            t_element = t_element[0]
-                        types_element.add(t_element)
-                # TODO: does this make sense as elements is always None
+                        t_elements.add(tuple(m.group(k) for k in t_keys))
+                # Keep the elements with the values to grab
+                replacement_values = self.get_replacement_values()
+                restricted = [
+                    (i, set(replacement_values[k]))
+                    for i, k in enumerate(t_keys)
+                    if k in replacement_values
+                ]
+                t_elements = {
+                    element
+                    for element in t_elements
+                    if all(element[i] in values for i, values in restricted)
+                }
+                # Keep the elements available for all the data types
                 if elements is None:
-                    elements = types_element
+                    elements, keys = t_elements, t_keys
                 else:
-                    # Do the intersection by filtering out elements in which
-                    # the replacements are not None
-                    if t_replacements == self.replacements:
-                        elements.intersection(types_element)
-                    else:
-                        t_repl_idx = [
-                            i
-                            for i, v in enumerate(self.replacements)
-                            if v in t_replacements
-                        ]
-                        new_elements = set()
-                        for t_element in elements:
-                            if (
-                                tuple(np.array(t_element)[t_repl_idx])
-                                in types_element
-                            ):
-                                new_elements.add(t_element)
-                        elements = new_elements
+                    elements, keys = _join_elements(
+                        elements, keys, t_elements, t_keys
+                    )
         if elements is None:
-            elements = set()
-        return list(elements)
+            return []
+        # Sort the values as the element keys
+        element_keys = self.get_element_keys()
+        idx = [keys.index(k) for k in element_keys]
+        out = sorted(tuple(element[i] for i in idx) for element in elements)
+        if len(element_keys) == 1:
+            out = [element[0] for element in out]
+        return out