Skip to content

abstract_dataloader.ext.graph

Composing transforms for data processing pipelines.

Programming Model

  • Data is represented as a dictionary with string keys and arbitrary values which are atomic from the perspective of transform composition.
  • Transforms are created from a directed acyclic graph (DAG) of nodes, where each node (Node) is a callable which takes a set of inputs and produces a set of outputs.

abstract_dataloader.ext.graph.Node dataclass

Node specification for a graph-based data processing transform.

Example Hydra Config

rsp:
  inputs:
    iq: radar
    aug: aug
  output: spectrum
  transform:
    _target_: grt.transforms.Spectrum
    ...

Attributes:

Name Type Description
transform Callable

callable to apply to the inputs.
output str | Sequence[str]

output data key (or output data keys for a node which returns multiple outputs).
inputs Mapping[str, str]

mapping of data keys to input argument names.
optional Mapping[str, str]

mapping of optional data keys to input argument names (i.e., they are only passed if present).
Source code in src/abstract_dataloader/ext/graph.py 
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
@dataclass
class Node:
    """Node specification for a graph-based data processing transform.

    !!! example "Example Hydra Config"

        ```yaml
        rsp:
          inputs:
            iq: radar
            aug: aug
          output: spectrum
          transform:
            _target_: grt.transforms.Spectrum
            ...
        ```

    Attributes:
        transform: callable to apply to the inputs.
        output: output data key (or output data keys for a node which returns
            multiple outputs).
        inputs: mapping of data keys to input argument names.
        optional: mapping of optional data keys to input argument names (i.e.,
            they are only passed if present).
    """

    transform: Callable
    output: str | Sequence[str]
    inputs: Mapping[str, str] = field(default_factory=dict)
    optional: Mapping[str, str] = field(default_factory=dict)

    def apply(self, data: dict[str, Any], name: str = "") -> dict[str, Any]:
        """Apply the node.

        Args:
            data: input data to process.
            name: node name (for error messages).

        Returns:
            Updated data, with any new keys added to the input data.
        """
        inputs = {k: data[v] for k, v in self.inputs.items()}
        for k, v in self.optional.items():
            if v in data:
                inputs[k] = data[v]

        output = self.transform(**inputs)

        if isinstance(self.output, str):
            data[self.output] = output
        else:  # Sequence[str]
            if not isinstance(output, Sequence):
                raise TypeError(
                    f"Node '{name}' output is expected to be a sequence due "
                    f"to output specification {self.output}: "
                    f"\n{wl.pformat(output)}\n")
            if len(self.output) != len(output):
                raise ValueError(
                    f"Node '{name}' output length mismatch: expected "
                    f"{len(self.output)} outputs ({self.output}), but got "
                    f"{len(output)} outputs:\n{wl.pformat(output)}\n")

            for o, v in zip(self.output, output):
                data[o] = v

        return data

apply 

apply(data: dict[str, Any], name: str = '') -> dict[str, Any]

Apply the node.

Parameters:

Name Type Description Default
data dict[str, Any]

input data to process.

 required
name str

node name (for error messages).

 ''

Returns:

Type Description
dict[str, Any]

Updated data, with any new keys added to the input data.
Source code in src/abstract_dataloader/ext/graph.py 
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
def apply(self, data: dict[str, Any], name: str = "") -> dict[str, Any]:
    """Apply the node.

    Args:
        data: input data to process.
        name: node name (for error messages).

    Returns:
        Updated data, with any new keys added to the input data.
    """
    inputs = {k: data[v] for k, v in self.inputs.items()}
    for k, v in self.optional.items():
        if v in data:
            inputs[k] = data[v]

    output = self.transform(**inputs)

    if isinstance(self.output, str):
        data[self.output] = output
    else:  # Sequence[str]
        if not isinstance(output, Sequence):
            raise TypeError(
                f"Node '{name}' output is expected to be a sequence due "
                f"to output specification {self.output}: "
                f"\n{wl.pformat(output)}\n")
        if len(self.output) != len(output):
            raise ValueError(
                f"Node '{name}' output length mismatch: expected "
                f"{len(self.output)} outputs ({self.output}), but got "
                f"{len(output)} outputs:\n{wl.pformat(output)}\n")

        for o, v in zip(self.output, output):
            data[o] = v

    return data

abstract_dataloader.ext.graph.Transform

Bases: Transform[dict[str, Any], dict[str, Any]]

Compose multiple callables forming a DAG into a transform.

Warning

Since the input data specifications are not provided at initialization, the graph execution order (or graph validity) is not statically determined, and result in runtime errors if invalid.

Parameters:

Name Type Description Default
outputs Mapping[str, str] | None

output data keys to produce as a mapping of output keys to graph data keys. If None, all values are returned.

 None
keep_all bool

keep references to all intermediate values and return them instead of decref-ing values which are no longer needed.

 False
nodes Node | dict[str, Any]

nodes in the graph, as keyword arguments where the key indicates a reference name for the node; any dict arguments are passed as key/value arguments to Node.

 {}
Source code in src/abstract_dataloader/ext/graph.py 
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
class Transform(spec.Transform[dict[str, Any], dict[str, Any]]):
    """Compose multiple callables forming a DAG into a transform.

    !!! warning

        Since the input data specifications are not provided at initialization,
        the graph execution order (or graph validity) is not statically
        determined, and result in runtime errors if invalid.

    Args:
        outputs: output data keys to produce as a mapping of output keys to
            graph data keys. If `None`, all values are returned.
        keep_all: keep references to all intermediate values and return them
            instead of decref-ing values which are no longer needed.
        nodes: nodes in the graph, as keyword arguments where the key indicates
            a reference name for the node; any `dict` arguments are passed as
            key/value arguments to [`Node`][^.].
    """

    def __init__(
        self, outputs: Mapping[str, str] | None = None, keep_all: bool = False,
        **nodes: Node | dict[str, Any]
    ) -> None:
        self.nodes = {}
        for k, v in nodes.items():
            if not isinstance(v, Node):
                try:
                    v = Node(**v)
                except TypeError as e:
                    raise TypeError(
                        f"Node '{k}' is not a valid Node specification: "
                        f"{wl.pformat(v)}") from e
            self.nodes[k] = v

        self.outputs = outputs
        self.keep_all = keep_all

    def _err_disconnected(
        self, data: dict[str, Any], incomplete: dict[str, Node]
    ) -> ValueError:
        """Format error message for disconnected nodes."""
        remaining = {k: list(v.inputs.values()) for k, v in incomplete.items()}
        return ValueError(
            f"There are {len(incomplete)} nodes remaining, but "
            f"all remaining nodes have at least one missing input.\n"
            f"Current inputs: {list(data.keys())}\n"
            f"Remaining stage requirements: {wl.pformat(remaining)}")

    def _decref(
        self, data: dict[str, Any], incomplete: dict[str, Node]
    ) -> dict[str, Any]:
        """Decref unneeded data values."""
        if self.keep_all or self.outputs is None:
            return data
        else:
            keep = set(self.outputs.values())
            for node in self.nodes.values():
                keep |= set(node.inputs.values())
            return {k: v for k, v in data.items() if k in keep}

    def __call__(self, data: dict[str, Any]) -> dict[str, Any]:
        """Execute the transform graph on the input data.

        Args:
            data: input data to process.

        Returns:
            Processed data.
        """
        incomplete = self.nodes.copy()

        # Guaranteed to terminate:
        # Each loop removes one node, or raises an error.
        while len(incomplete) > 0:
            for name, node in incomplete.items():
                if all(k in data for k in node.inputs.values()):
                    data = node.apply(data, name=name)
                    incomplete.pop(name)
                    data = self._decref(data, incomplete)
                    break  # break back into the while
            else:
                raise self._err_disconnected(data, incomplete)

        if self.outputs is not None:
            if self.keep_all:
                for k, v in self.outputs.items():
                    data[k] = data[v]
                return data
            else:
                return {k: data[v] for k, v in self.outputs.items()}
        else:
            return data

    def children(self) -> Iterable[Any]:
        """Get all non-container child objects."""
        for node in self.nodes.values():
            yield node.transform

    def __repr__(self) -> str:
        """Friendly name."""
        return f"TransformGraph({', '.join(self.nodes.keys())})"

__call__ 

__call__(data: dict[str, Any]) -> dict[str, Any]

Execute the transform graph on the input data.

Parameters:

Name Type Description Default
data dict[str, Any]

input data to process.

 required

Returns:

Type Description
dict[str, Any]

Processed data.
Source code in src/abstract_dataloader/ext/graph.py 
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
def __call__(self, data: dict[str, Any]) -> dict[str, Any]:
    """Execute the transform graph on the input data.

    Args:
        data: input data to process.

    Returns:
        Processed data.
    """
    incomplete = self.nodes.copy()

    # Guaranteed to terminate:
    # Each loop removes one node, or raises an error.
    while len(incomplete) > 0:
        for name, node in incomplete.items():
            if all(k in data for k in node.inputs.values()):
                data = node.apply(data, name=name)
                incomplete.pop(name)
                data = self._decref(data, incomplete)
                break  # break back into the while
        else:
            raise self._err_disconnected(data, incomplete)

    if self.outputs is not None:
        if self.keep_all:
            for k, v in self.outputs.items():
                data[k] = data[v]
            return data
        else:
            return {k: data[v] for k, v in self.outputs.items()}
    else:
        return data

__repr__ 

__repr__() -> str

Friendly name.

Source code in src/abstract_dataloader/ext/graph.py 
187
188
189
def __repr__(self) -> str:
    """Friendly name."""
    return f"TransformGraph({', '.join(self.nodes.keys())})"

children 

children() -> Iterable[Any]

Get all non-container child objects.

Source code in src/abstract_dataloader/ext/graph.py 
182
183
184
185
def children(self) -> Iterable[Any]:
    """Get all non-container child objects."""
    for node in self.nodes.values():
        yield node.transform