From de41483be43cc0f08dcc916c9418270f0818a95f Mon Sep 17 00:00:00 2001 From: Jenny Wong Date: Thu, 9 Jul 2026 09:00:01 +0100 Subject: [PATCH 1/5] Fix selection --- src/earthkit/workflows/fluent.py | 41 ++++++++++++++----------- tests/earthkit_workflows/test_fluent.py | 4 +-- 2 files changed, 25 insertions(+), 20 deletions(-) diff --git a/src/earthkit/workflows/fluent.py b/src/earthkit/workflows/fluent.py index 84ab841..1d58106 100644 --- a/src/earthkit/workflows/fluent.py +++ b/src/earthkit/workflows/fluent.py @@ -751,11 +751,19 @@ def _validate_criteria(self, array: xr.DataArray, criteria: dict) -> tuple[bool, keys = list(criteria.keys()) new_criteria = criteria.copy() for key in keys: + if np.ndim(criteria[key]) == 1 and len(criteria[key]) == 1: + new_criteria[key] = criteria[key][0] if key not in array.dims: - if array.coords.get(key, None) == criteria[key]: - new_criteria.pop(key) - else: + coords = array.coords.get(key, None) + if coords is None: + return False, {} + coords = coords.data.tolist() + if not isinstance(coords, list): + coords = [coords] + if new_criteria[key] not in coords: return False, {} + if len(coords) == 1: + new_criteria.pop(key) return True, new_criteria def _select( @@ -777,7 +785,7 @@ def _select( for npath, narray in nodetree_arrays(nodes): # type: ignore[arg-type] if expand: keys = list(crit.keys()) - its = tuple(crit.values()) + its = tuple([x] if np.ndim(x) == 0 else x for x in crit.values()) crits = [dict(zip(keys, vals)) for vals in itertools.product(*its)] else: crits = [crit] @@ -1125,22 +1133,19 @@ def set_scalar_coords( self.nodes = nodetree_from_dict(nodetree) def _add_dimension(self, name: str, value: Any, axis: Optional[int] = None, path: Optional[str] = None): - new_tree = self.nodes.map_over_datasets(lambda ds: ds.expand_dims({name: [value]}, axis)) - assert isinstance(new_tree, xr.DataTree) - if path is not None: - self.nodes[path] = new_tree[path] - else: - self.nodes = new_tree + nodetree = {npath: narray for npath, narray in nodetree_arrays(self.nodes)} + selection = self.select(path=path) + for npath, narray in nodetree_arrays(selection.nodes): + nodetree[npath] = narray.expand_dims({name: [value]}, axis) + self.nodes = nodetree_from_dict(nodetree) def _squeeze_dimension(self, dim_name: str, drop: bool = False, path: Optional[str] = None): - new_tree = self.nodes.map_over_datasets( - lambda ds: ds.squeeze(dim_name, drop=drop) if dim_name in ds.dims and len(ds.coords[dim_name]) == 1 else ds - ) - assert isinstance(new_tree, xr.DataTree) - if path is not None: - self.nodes[path] = new_tree[path] - else: - self.nodes = new_tree + nodetree = {npath: narray for npath, narray in nodetree_arrays(self.nodes)} + selection = self.select(path=path) + for npath, narray in nodetree_arrays(selection.nodes): + if dim_name in narray.dims and len(narray.coords[dim_name]) == 1: + nodetree[npath] = narray.squeeze(dim_name, drop=drop) + self.nodes = nodetree_from_dict(nodetree) def __getattr__(self, attr): if attr in Action.REGISTRY: diff --git a/tests/earthkit_workflows/test_fluent.py b/tests/earthkit_workflows/test_fluent.py index d273a99..9f32fa1 100644 --- a/tests/earthkit_workflows/test_fluent.py +++ b/tests/earthkit_workflows/test_fluent.py @@ -388,7 +388,7 @@ def test_flatten_branches(): ({"dim_1": 4}, 1, [(2,)]), ({"path": "/branch1"}, 2, [(3, 4), (3, 4)]), ({"path": "/branch1", "dim_0": 1}, 2, [(4,), (4,)]), - ({"type": "A"}, 1, [(2, 5)]), + ({"type": ["AB"]}, 1, [(2, 5)]), ({"dim_1": 10}, 0, IndexError), ({"dim_0": [2], "dim_1": [0, 4]}, 0, IndexError), ({"dim_0": [2], "dim_1": [0, 4], "expand": True}, 2, [(1, 1), (1, 1)]), @@ -406,7 +406,7 @@ def test_select(selection, num_arrays, shapes_or_error): "/branch1/subbranch2": lambda data: np.where(data == 0, data, np.nan), } ) - subbranches.nodes["/branch2"].coords["type"] = "A" + subbranches.nodes["/branch2"].coords["type"] = "AB" if num_arrays > 0: select_dim = subbranches.sel(**selection) assert len(list(nodetree_arrays(select_dim.nodes))) == num_arrays From d97a5fb97a1ed784bd70dca85e5b801882749825 Mon Sep 17 00:00:00 2001 From: Jenny Wong Date: Thu, 9 Jul 2026 09:11:49 +0100 Subject: [PATCH 2/5] Allow override in adding dimension --- src/earthkit/workflows/fluent.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src/earthkit/workflows/fluent.py b/src/earthkit/workflows/fluent.py index 1d58106..ca7739f 100644 --- a/src/earthkit/workflows/fluent.py +++ b/src/earthkit/workflows/fluent.py @@ -371,8 +371,7 @@ def transform( for index, param in enumerate(params): new_res = func(self.select(path=path), *param) - if dim_name not in new_res.nodes.coords: - new_res._add_dimension(dim_name, dim_values[index], axis, path=path) + new_res._add_dimension(dim_name, dim_values[index], axis, path=path, override=True) if res is None: res = new_res else: @@ -1132,10 +1131,12 @@ def set_scalar_coords( nodetree[npath] = narray self.nodes = nodetree_from_dict(nodetree) - def _add_dimension(self, name: str, value: Any, axis: Optional[int] = None, path: Optional[str] = None): + def _add_dimension(self, name: str, value: Any, axis: Optional[int] = None, path: Optional[str] = None, override: bool = False): nodetree = {npath: narray for npath, narray in nodetree_arrays(self.nodes)} selection = self.select(path=path) for npath, narray in nodetree_arrays(selection.nodes): + if override and name in narray.dims and len(narray.coords[name]) == 1: + narray = narray.squeeze(name, drop=True) nodetree[npath] = narray.expand_dims({name: [value]}, axis) self.nodes = nodetree_from_dict(nodetree) From 36c79870bfb3b0f7a6baa9e340c08e785c6d7433 Mon Sep 17 00:00:00 2001 From: Jenny Wong Date: Fri, 10 Jul 2026 17:41:15 +0100 Subject: [PATCH 3/5] Tidy type annotations --- src/earthkit/workflows/fluent.py | 72 ++++++++++++++++---------------- 1 file changed, 36 insertions(+), 36 deletions(-) diff --git a/src/earthkit/workflows/fluent.py b/src/earthkit/workflows/fluent.py index ca7739f..46b8a67 100644 --- a/src/earthkit/workflows/fluent.py +++ b/src/earthkit/workflows/fluent.py @@ -314,10 +314,10 @@ def as_action(self, other) -> Action: def join( self, - other_action: "Action", + other_action: Action, dim: str | Coord, match_coord_values: bool = False, - ) -> "Action": + ) -> Action: node_arrays = {} for npath, narray in nodetree_arrays(self.nodes): oarray = nodetree_array(other_action.nodes, npath) @@ -336,12 +336,12 @@ def join( def transform( self, - func: Callable[..., "Action"], + func: Callable[..., Action], params: list, dim: str | Coord, axis: int = 0, path: Optional[str] = None, - ) -> "Action": + ) -> Action: """Create new nodes by applying function on action with different parameters. The result actions from applying function are joined along the specified dimension. @@ -388,10 +388,10 @@ def transform( def broadcast( self, - other_action: "Action", + other_action: Action, exclude: list[str] | None = None, path: Optional[str] = None, - ) -> "Action": + ) -> Action: """Broadcast nodes against nodes in other_action Parameters @@ -441,7 +441,7 @@ def expand( axis: int = 0, path: Optional[str] = None, backend_kwargs: dict = {}, - ) -> "Action": + ) -> Action: """Create new dimension in array of nodes of specified size by taking elements of internal data in each node. Indexing is taken along the specified axis dimension of internal data and graph execution will fail if @@ -482,7 +482,7 @@ def map( payload: PayloadFunc | Payload | np.ndarray[Any, Any] | list, yields: Coord | None = None, path: Optional[str] = None, - ) -> "Action": + ) -> Action: """Apply specified payload on all nodes. If argument is an array of payloads, this must be the same size as the array of nodes and each node gets a unique payload from the array @@ -545,7 +545,7 @@ def reduce( batch_size: int = 0, keep_dim: bool = False, path: Optional[str] = None, - ) -> "Action": + ) -> Action: """Reduction operation across the named dimension using the provided function in the payload. If batch_size > 1 and less than the size of the named dimension, the reduction will be computed first in @@ -643,7 +643,7 @@ def flatten( keep_dims: list[str] = [], path: Optional[str] = None, reset_coords: bool = False, - ) -> "Action": + ) -> Action: """Restructures node arrays by flattening arrays along all dims, except keep_dims, for node arrays along path. @@ -673,7 +673,7 @@ def flatten( node_arrays[npath] = node_arrays[npath].reset_coords(to_reset, drop=True) return type(self)(nodetree_from_dict(node_arrays)) - def set_path(self, path: str) -> "Action": + def set_path(self, path: str) -> Action: """Create path for current node array Parameters @@ -688,7 +688,7 @@ def set_path(self, path: str) -> "Action": raise NotImplementedError("Multiple node arrays present, can not set single path") return type(self)(nodetree_from_dict({path: nodetree_array(self.nodes)})) - def create_branches(self, expansion: dict[str, PayloadFunc | Payload]) -> "Action": + def create_branches(self, expansion: dict[str, PayloadFunc | Payload]) -> Action: """Create action containing new node arrays by splitting an existing node array by the specified functions in expansion @@ -713,7 +713,7 @@ def create_branches(self, expansion: dict[str, PayloadFunc | Payload]) -> "Actio node_arrays[path] = nodetree_array(action.map(func).nodes, parent) return type(self)(nodetree_from_dict(node_arrays)) - def combine_branches(self, dim: str, path: Optional[str] = None, force: bool = False) -> "Action": + def combine_branches(self, dim: str, path: Optional[str] = None, force: bool = False) -> Action: """Combine node arrays for leaves along path into a single node array Parameters @@ -773,7 +773,7 @@ def _select( expand: bool = False, backend_method: Union[Literal["sel"], Literal["isel"]] = "sel", **kwargs, - ) -> "Action": + ) -> Action: if backend_method not in ["sel", "isel"]: raise ValueError(f"backend_method must be 'sel' or 'isel', got {backend_method}") crit: dict = criteria or {} @@ -816,7 +816,7 @@ def select( path: Optional[str] = None, expand: bool = False, **kwargs, - ) -> "Action": + ) -> Action: """Create action contaning nodes match selection criteria Parameters @@ -841,7 +841,7 @@ def iselect( path: Optional[str] = None, expand: bool = False, **kwargs, - ) -> "Action": + ) -> Action: """Create action contaning nodes match index selection criteria Parameters @@ -868,7 +868,7 @@ def concatenate( path: Optional[str] = None, backend_kwargs: dict = {}, payload_metadata: dict | None = None, - ) -> "Action": + ) -> Action: return _combine_nodes(self, "concat", dim, batch_size, keep_dim, path, backend_kwargs) @capture_payload_metadata @@ -881,7 +881,7 @@ def stack( path: Optional[str] = None, backend_kwargs: dict = {}, payload_metadata: dict | None = None, - ) -> "Action": + ) -> Action: return _combine_nodes( self, "stack", @@ -901,7 +901,7 @@ def sum( path: Optional[str] = None, backend_kwargs: dict = {}, payload_metadata: dict | None = None, - ) -> "Action": + ) -> Action: return self.reduce( Payload(backends.sum, kwargs=backend_kwargs), dim=dim, @@ -919,7 +919,7 @@ def mean( path: Optional[str] = None, backend_kwargs: dict = {}, payload_metadata: dict | None = None, - ) -> "Action": + ) -> Action: action = self for npath, narray in nodetree_arrays(self.select(path=path).nodes): if len(dim) == 0: @@ -952,7 +952,7 @@ def std( path: Optional[str] = None, backend_kwargs: dict = {}, payload_metadata: dict | None = None, - ) -> "Action": + ) -> Action: action = self for npath, narray in nodetree_arrays(self.select(path=path).nodes): if len(dim) == 0: @@ -993,7 +993,7 @@ def max( path: Optional[str] = None, backend_kwargs: dict = {}, payload_metadata: dict | None = None, - ) -> "Action": + ) -> Action: return self.reduce( Payload(backends.max, kwargs=backend_kwargs), dim=dim, @@ -1011,7 +1011,7 @@ def min( path: Optional[str] = None, backend_kwargs: dict = {}, payload_metadata: dict | None = None, - ) -> "Action": + ) -> Action: return self.reduce( Payload(backends.min, kwargs=backend_kwargs), dim=dim, @@ -1029,7 +1029,7 @@ def prod( path: Optional[str] = None, backend_kwargs: dict = {}, payload_metadata: dict | None = None, - ) -> "Action": + ) -> Action: return self.reduce( Payload(backends.prod, kwargs=backend_kwargs), dim=dim, @@ -1041,10 +1041,10 @@ def prod( def __two_arg_method( self, method: Callable, - other: "Action | float", + other: Union[Action, float], path: Optional[str] = None, **kwargs, - ) -> "Action": + ) -> Action: if isinstance(other, Action): return self.join(other, "**datatype**", match_coord_values=True).reduce( Payload(method, kwargs=kwargs), dim="**datatype**", path=path @@ -1054,51 +1054,51 @@ def __two_arg_method( @capture_payload_metadata def subtract( self, - other: "Action | float", + other: Union[Action, float], path: Optional[str] = None, backend_kwargs: dict = {}, payload_metadata: dict | None = None, - ) -> "Action": + ) -> Action: return self.__two_arg_method(backends.subtract, other, path=path, **backend_kwargs) @capture_payload_metadata def divide( self, - other: "Action | float", + other: Union[Action, float], path: Optional[str] = None, backend_kwargs: dict = {}, payload_metadata: dict | None = None, - ) -> "Action": + ) -> Action: return self.__two_arg_method(backends.divide, other, path=path, **backend_kwargs) @capture_payload_metadata def add( self, - other: "Action | float", + other: Union[Action, float], path: Optional[str] = None, backend_kwargs: dict = {}, payload_metadata: dict | None = None, - ) -> "Action": + ) -> Action: return self.__two_arg_method(backends.add, other, path=path, **backend_kwargs) @capture_payload_metadata def multiply( self, - other: "Action | float", + other: Union[Action, float], path: Optional[str] = None, backend_kwargs: dict = {}, payload_metadata: dict | None = None, - ) -> "Action": + ) -> Action: return self.__two_arg_method(backends.multiply, other, path=path, **backend_kwargs) @capture_payload_metadata def power( self, - other: "Action | float", + other: Union[Action, float], path: Optional[str] = None, backend_kwargs: dict = {}, payload_metadata: dict | None = None, - ) -> "Action": + ) -> Action: return self.__two_arg_method(backends.pow, other, path=path, **backend_kwargs) def add_attributes(self, attrs: dict): From ab330743baeb16979bed8cfdd33235a029f44dd4 Mon Sep 17 00:00:00 2001 From: Jenny Wong Date: Wed, 22 Jul 2026 10:03:06 +0100 Subject: [PATCH 4/5] Revert change to transform --- src/earthkit/workflows/fluent.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/earthkit/workflows/fluent.py b/src/earthkit/workflows/fluent.py index 46b8a67..3588b79 100644 --- a/src/earthkit/workflows/fluent.py +++ b/src/earthkit/workflows/fluent.py @@ -23,7 +23,7 @@ from ._qubed import expand_as_qube from .graph import Graph, Output from .graph import Node as BaseNode -from .nodetree import combine_by_coords, nodetree_array, nodetree_arrays, nodetree_from_dict, nodetree_new_dimension +from .nodetree import combine_by_coords, nodetree_array, nodetree_arrays, nodetree_dimensions, nodetree_from_dict, nodetree_new_dimension PayloadFunc = Callable | str @@ -371,7 +371,8 @@ def transform( for index, param in enumerate(params): new_res = func(self.select(path=path), *param) - new_res._add_dimension(dim_name, dim_values[index], axis, path=path, override=True) + if dim_name not in nodetree_dimensions(new_res.nodes): + new_res._add_dimension(dim_name, dim_values[index], axis, path=path, override=True) if res is None: res = new_res else: From d6e3a9c6c7f26bb593d5d9d8b3a818f81560434a Mon Sep 17 00:00:00 2001 From: Jenny Wong Date: Wed, 22 Jul 2026 14:35:29 +0100 Subject: [PATCH 5/5] Remove redundant keys in selection criteria --- src/earthkit/workflows/fluent.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/src/earthkit/workflows/fluent.py b/src/earthkit/workflows/fluent.py index 3588b79..83b99d0 100644 --- a/src/earthkit/workflows/fluent.py +++ b/src/earthkit/workflows/fluent.py @@ -764,6 +764,12 @@ def _validate_criteria(self, array: xr.DataArray, criteria: dict) -> tuple[bool, return False, {} if len(coords) == 1: new_criteria.pop(key) + + # Remove from criteria coords that are dependent on other coords in criteria + dependencies = {name: [x for x in val.indexes.keys() if x != name] for name, val in array.coords.items()} + for key, dep in dependencies.items(): + if key in new_criteria and any(d in new_criteria for d in dep): + new_criteria.pop(key) return True, new_criteria def _select(