diff --git a/docs/src/python/py_qubed.md b/docs/src/python/py_qubed.md index 8582923..5fa7551 100644 --- a/docs/src/python/py_qubed.md +++ b/docs/src/python/py_qubed.md @@ -229,9 +229,139 @@ print(q) --- -### Query +### Structural utilities + +These methods implement structural helpers to work with complex Qubes. + +#### `axes() -> dict[str, list[str]]` + +Return all dimension names and the union of their coordinate values across the +entire Qube. This is an alias for `all_unique_dim_coords()` with a name that +matches the terminology used in the qubed-utils helper library. + +```python +q = Qube.from_datacube({"param": "2t/tp", "time": "0/1/2"}, ["param", "time"]) +ax = q.axes() +# {'param': ['2t', 'tp'], 'time': ['0', '1', '2']} +assert set(ax["param"]) == {"2t", "tp"} +``` + +#### `dimensions() -> set[str]` + +Return the set of all dimension names present anywhere in the Qube. + +```python +q = Qube.from_datacube({"param": "2t/tp", "time": "0/1/2"}, ["param", "time"]) +assert q.dimensions() == {"param", "time"} +``` + +#### `common_dimensions() -> set[str]` + +Return the set of dimension names that appear in **every** leaf path +(datacube). For a Qube with uniform depth this equals `dimensions()`. For +an irregular Qube where some branches are shallower, only the dimensions +shared by all branches are returned. + +```python +q1 = Qube.from_datacube({"param": "2t", "time": "0/1"}, ["param", "time"]) +q2 = Qube.from_datacube({"param": "msl"}, ["param"]) +q1.append(q2) +assert q1.common_dimensions() == {"param"} # "time" absent in branch 2 +``` + +#### `expand(dimension: dict[str, list]) -> None` + +Wrap the entire Qube tree under one or more new outer dimensions. Each key +in `dimension` becomes a new dimension name; the associated list supplies its +coordinate values. + +Dimensions are applied in dict insertion order. The **last** entry in the +dict becomes the outermost dimension of the resulting tree. The operation +mutates the Qube in place. + +```python +q = Qube.from_datacube({"param": "2t/tp", "time": "0/1/2"}, ["param", "time"]) +q.expand({"ensemble": ["ens1", "ens2"]}) +assert "ensemble" in q.dimensions() +assert "param" in q.dimensions() +# The new dimension wraps the original tree: +# root +# └── ensemble=ens1/ens2 +# └── param=2t/tp +# └── time=0/1/2 + +# Multiple dimensions at once: +q.expand({"member": ["m1"], "batch": ["b1", "b2"]}) +# "batch" ends up outermost since it was last in the dict. +``` + +#### `collapse(axis: str | list[str]) -> None` + +Remove one or more dimensions from the Qube. `axis` may be a single +dimension name or a list of names. Children of removed nodes are re-parented +to their grandparent, preserving the rest of the structure. The result is +automatically compressed. + +Raises `ValueError` if any of the specified dimensions do not exist. + +```python +q = Qube.from_datacube( + {"param": "2t/tp", "time": "0/1/2", "level": "1000/850"}, + ["param", "time", "level"], +) +q.collapse("level") +assert "level" not in q.dimensions() +assert "param" in q.dimensions() + +# Remove multiple dimensions at once: +q.collapse(["time", "param"]) +assert q.dimensions() == set() +``` + +#### `coxpand(axis: str | list[str], dimension: dict[str, list]) -> None` + +Collapse one or more dimensions and then expand with new ones in a single +call. Equivalent to `collapse(axis)` followed by `expand(dimension)`. + +Useful for replacing a dimension with a different one while preserving the +rest of the structure. + +```python +q = Qube.from_datacube({"param": "2t/tp", "time": "0/1/2"}, ["param", "time"]) +q.coxpand("time", {"step": ["s1", "s2"]}) +assert "time" not in q.dimensions() +assert "step" in q.dimensions() +assert "param" in q.dimensions() +``` + +#### `contains(item: str | dict | Qube) -> bool` + +Check whether the Qube contains a given dimension or set of values. + +| `item` type | Meaning | +|---|---| +| `str` | Returns `True` if the named dimension exists anywhere in the Qube. | +| `dict[str, list]` | Returns `True` if every key exists as a dimension **and** every listed value is present in that dimension's coordinate set. | +| `Qube` | Returns `True` if every dimension+value from the other Qube is also present here (subset check). | + +```python +q = Qube.from_datacube({"param": "2t/tp", "time": "0/1/2"}, ["param", "time"]) + +assert q.contains("param") # True +assert q.contains("level") # False + +assert q.contains({"param": ["2t"]}) # True +assert q.contains({"param": ["xyz"]}) # False +assert q.contains({"param": ["2t"], "time": ["0"]}) # True +assert q.contains({"time": ["0", "999"]}) # False – 999 absent + +subset = Qube.from_datacube({"param": "2t", "time": "0"}, ["param", "time"]) +assert q.contains(subset) # True +``` + +--- + -#### `all_unique_dim_coords() -> dict[str, list[str]]` Return a dictionary mapping each dimension name to a list of all coordinate values that appear anywhere in the Qube. diff --git a/docs/src/rust/qubed.md b/docs/src/rust/qubed.md index 20eaa5c..2399e8b 100644 --- a/docs/src/rust/qubed.md +++ b/docs/src/rust/qubed.md @@ -66,6 +66,7 @@ let q = Qube::from_json(json!({ | `append_datacube` | `fn append_datacube(&mut self, dc: Datacube, order: Option<&[String]>, accept_existing_order: bool)` | Append a single Datacube | | `drop` | `fn drop(&mut self, to_drop: I) -> Result<(), String>` | Remove one or more dimensions, re-parenting their children, then compress | | `squeeze` | `fn squeeze(&mut self) -> Result<(), String>` | Drop every dimension whose union of values has length 1 | +| `expand` | `fn expand(&mut self, key: &str, values: Coordinates) -> Result<(), String>` | Wrap the entire tree under a new outer dimension | **Example — building programmatically:** ```rust @@ -116,6 +117,44 @@ q.squeeze().unwrap(); // class=1 is the only value for that dimension, so it is dropped ``` +**Example — expand:** +```rust +use qubed::{Qube, Coordinates}; + +let mut q = Qube::from_ascii("root\n└── param=2t/tp\n └── time=0/1/2").unwrap(); +q.expand("ensemble", Coordinates::from_string("ens1/ens2")).unwrap(); + +// Tree is now: +// root +// └── ensemble=ens1/ens2 +// └── param=2t/tp +// └── time=0/1/2 + +let dims = q.dimensions(); +assert!(dims.contains("ensemble")); +assert!(dims.contains("param")); +assert!(dims.contains("time")); +``` + +**Example — common_dimensions:** +```rust +use qubed::{Qube, Datacube, Coordinates}; + +let mut dc1 = Datacube::new(); +dc1.add_coordinate("param", Coordinates::from_string("2t/tp")); +dc1.add_coordinate("time", Coordinates::from_string("0/1/2")); +let mut q = Qube::from_datacube(&dc1, Some(&["param".to_string(), "time".to_string()])); + +let mut dc2 = Datacube::new(); +dc2.add_coordinate("param", Coordinates::from_string("msl")); +let mut other = Qube::from_datacube(&dc2, None); +q.append(&mut other); + +let common = q.common_dimensions(); +assert!(common.contains("param")); +assert!(!common.contains("time")); // "time" absent in the second branch +``` + ### Compression ```rust @@ -169,7 +208,9 @@ Remove branches that don't contain **all** of the specified dimensions. | `to_datacubes` | `fn to_datacubes(&self) -> Vec` | Decompose into leaf-path datacubes | | `datacube_count` | `fn datacube_count(&self) -> usize` | Count leaf identifiers without expansion | | `is_empty` | `fn is_empty(&self) -> bool` | True if root has no children and no coordinates | -| `all_unique_dim_coords` | `fn all_unique_dim_coords(&mut self) -> BTreeMap` | Union of all coordinates per dimension | +| `all_unique_dim_coords` | `fn all_unique_dim_coords(&self) -> BTreeMap` | Union of all coordinates per dimension | +| `dimensions` | `fn dimensions(&self) -> HashSet` | Set of all dimension names present in the Qube | +| `common_dimensions` | `fn common_dimensions(&self) -> HashSet` | Dimension names present in **every** leaf path | | `root` | `fn root(&self) -> NodeIdx` | Root node index | | `node` | `fn node(&self, id: NodeIdx) -> Option` | Read-only reference to a node | | `dimension` | `fn dimension(&self, s: &str) -> Option` | Look up dimension by name | diff --git a/py_qubed/src/helpers.rs b/py_qubed/src/helpers.rs new file mode 100644 index 0000000..6534ca4 --- /dev/null +++ b/py_qubed/src/helpers.rs @@ -0,0 +1,43 @@ +//! Internal helpers used by the PyO3 bindings. +//! +//! Nothing in this module is exposed to Python; the functions here exist solely +//! to keep `lib.rs` readable by factoring out auxiliary logic. + +use pyo3::exceptions::PyTypeError; +use pyo3::prelude::*; +use pyo3::types::{PyDict, PyList}; + +/// Check that every key→[values] entry in `dict` is satisfied by `axes`. +/// +/// Returns `Ok(true)` when every dimension key in `dict` is present in `axes` +/// and every queried value for that dimension is found in `axes`'s value list. +/// Returns `Ok(false)` on the first mismatch. +pub(crate) fn check_dict_against_axes( + dict: &Bound<'_, PyDict>, + axes: &std::collections::BTreeMap>, + _py: Python<'_>, +) -> PyResult { + for (k, v) in dict.iter() { + let key: String = + k.extract().map_err(|_| PyTypeError::new_err("contains: dict keys must be strings"))?; + + let query_vals: Vec = if v.is_instance_of::() { + let lst = v.downcast::().map_err(|e| PyTypeError::new_err(e.to_string()))?; + lst.iter().map(|it| Ok(it.str()?.extract::()?)).collect::>()? + } else { + vec![v.str()?.extract::()?] + }; + + match axes.get(&key) { + None => return Ok(false), + Some(cur_vals) => { + for qval in &query_vals { + if !cur_vals.contains(qval) { + return Ok(false); + } + } + } + } + } + Ok(true) +} diff --git a/py_qubed/src/lib.rs b/py_qubed/src/lib.rs index e45f4bb..8ca3150 100644 --- a/py_qubed/src/lib.rs +++ b/py_qubed/src/lib.rs @@ -3,10 +3,13 @@ use ::qubed::Datacube; use ::qubed::Qube; use ::qubed::select::SelectMode; use pyo3::exceptions::PyTypeError; +use pyo3::exceptions::PyValueError; use pyo3::prelude::*; -use pyo3::types::{PyDict, PyList, PyModule}; +use pyo3::types::{PyDict, PyList, PyModule, PySet}; use serde_json::Value as JsonValue; +mod helpers; + #[pyclass(name = "Qube", unsendable)] pub struct PyQube { inner: Qube, @@ -213,6 +216,279 @@ impl PyQube { pub fn __repr__(&self) -> PyResult { Ok(format!("PyQube(root_id={:?})", self.inner.root())) } + + // ----------------------------------------------------------------------- + // New utility methods (mirrors of the @support_qubed_output functions + // in forecast-in-a-box/fiab_plugin_ecmwf/qubed_utils.py) + // ----------------------------------------------------------------------- + + /// Return all dimension names and their union of coordinate values. + /// + /// This is an alias for `all_unique_dim_coords()` and returns the same + /// ``dict[str, list[str]]`` mapping. The name `axes` matches the + /// terminology used in the qubed-utils helper library. + /// + /// ```python + /// q = Qube.from_datacube({"param": "2t/tp", "time": "0/1/2"}, ["param", "time"]) + /// ax = q.axes() + /// assert set(ax["param"]) == {"2t", "tp"} + /// assert set(ax["time"]) == {"0", "1", "2"} + /// ``` + pub fn axes(&self, py: Python<'_>) -> PyResult> { + let dim_coords = self.inner.all_unique_dim_coords(); + let py_dict = PyDict::new(py); + + for (dimension, coordinates) in dim_coords { + let coord_str = coordinates.to_string(); + let values: Vec<&str> = if coord_str.is_empty() { + vec![] + } else if coord_str.contains('/') { + coord_str.split('/').collect() + } else { + vec![&coord_str] + }; + + let py_list = PyList::empty(py); + for value in values { + py_list.append(value)?; + } + py_dict.set_item(dimension, py_list)?; + } + + Ok(py_dict.into_any().unbind()) + } + + /// Return the set of all dimension names present anywhere in the Qube. + /// + /// ```python + /// q = Qube.from_datacube({"param": "2t/tp", "time": "0/1/2"}, ["param", "time"]) + /// assert q.dimensions() == {"param", "time"} + /// ``` + pub fn dimensions(&self, py: Python<'_>) -> PyResult> { + let dims = self.inner.dimensions(); + let py_set = PySet::new(py, dims.iter())?; + Ok(py_set.into_any().unbind()) + } + + /// Return the set of dimension names present in **every** leaf path. + /// + /// For a Qube with uniform depth this equals `dimensions()`. When + /// branches differ in depth, only dimensions that appear in *all* + /// branches are returned. + /// + /// ```python + /// q1 = Qube.from_datacube({"param": "2t", "time": "0/1"}, ["param", "time"]) + /// q2 = Qube.from_datacube({"param": "msl"}, ["param"]) + /// q1.append(q2) + /// assert q1.common_dimensions() == {"param"} + /// ``` + pub fn common_dimensions(&self, py: Python<'_>) -> PyResult> { + let common = self.inner.common_dimensions(); + let py_set = PySet::new(py, common.iter())?; + Ok(py_set.into_any().unbind()) + } + + /// Wrap the entire Qube under one or more new outer dimensions. + /// + /// `dimension` is a ``dict[str, list]`` where each key becomes a new + /// dimension name and the corresponding list supplies its coordinate + /// values. Dimensions are applied in insertion order: the *last* entry + /// in the dict becomes the outermost dimension. + /// + /// The operation mutates the Qube in place and does **not** return a new + /// object (unlike the Python-reference implementation which follows an + /// immutable style). + /// + /// ```python + /// q = Qube.from_datacube({"param": "2t/tp", "time": "0/1/2"}, ["param", "time"]) + /// q.expand({"ensemble": ["ens1", "ens2"]}) + /// assert "ensemble" in q.dimensions() + /// assert "param" in q.dimensions() + /// assert "time" in q.dimensions() + /// ``` + pub fn expand(&mut self, dimension: Bound<'_, PyDict>) -> PyResult<()> { + for (k, v) in dimension.iter() { + let key: String = k + .extract() + .map_err(|_| PyTypeError::new_err("expand: dimension keys must be strings"))?; + + let coords = if v.is_instance_of::() { + let lst = + v.downcast::().map_err(|e| PyTypeError::new_err(e.to_string()))?; + let joined = join_pylist_as_path(lst)?; + Coordinates::from_string(&joined) + } else { + let py_str = v.str()?; + let s: String = py_str.extract()?; + Coordinates::from_string(&s) + }; + + self.inner.expand(&key, coords).map_err(PyTypeError::new_err)?; + } + Ok(()) + } + + /// Remove one or more dimensions from the Qube. + /// + /// `axis` may be a single dimension name (``str``) or a list of names. + /// All children of the removed nodes are re-parented to their + /// grandparent, preserving the rest of the structure. The result is + /// automatically compressed. + /// + /// Raises ``ValueError`` if any of the specified dimensions do not exist + /// in the Qube. + /// + /// ```python + /// q = Qube.from_datacube( + /// {"param": "2t/tp", "time": "0/1/2", "level": "1000/850"}, + /// ["param", "time", "level"], + /// ) + /// q.collapse("level") + /// assert "level" not in q.dimensions() + /// assert "param" in q.dimensions() + /// ``` + pub fn collapse(&mut self, axis: Bound<'_, PyAny>) -> PyResult<()> { + // Accept either a single string or a list of strings. + let axes: Vec = if axis.is_instance_of::() { + let lst = axis.downcast::().map_err(|e| PyTypeError::new_err(e.to_string()))?; + lst.iter() + .map(|item| { + item.str() + .and_then(|s| s.extract::()) + .map_err(|_| PyTypeError::new_err("collapse: axis names must be strings")) + }) + .collect::>()? + } else { + let s: String = axis + .extract() + .map_err(|_| PyTypeError::new_err("collapse: axis must be a str or list[str]"))?; + vec![s] + }; + + // Validate that every requested dimension actually exists. + let existing_dims = self.inner.dimensions(); + let missing: Vec<&str> = axes + .iter() + .filter(|a| !existing_dims.contains(a.as_str())) + .map(String::as_str) + .collect(); + + if !missing.is_empty() { + return Err(PyValueError::new_err(format!( + "Dimension(s) '{}' not found in qube (existing: {:?})", + missing.join(", "), + { + let mut sorted: Vec<&str> = existing_dims.iter().map(String::as_str).collect(); + sorted.sort(); + sorted + } + ))); + } + + self.inner.drop(axes).map_err(PyTypeError::new_err) + } + + /// Collapse then expand in a single operation. + /// + /// Equivalent to calling `collapse(axis)` followed by `expand(dimension)` + /// on the same Qube. Useful for replacing one dimension with another + /// (e.g. replacing "level" with "step"). + /// + /// ```python + /// q = Qube.from_datacube( + /// {"param": "2t/tp", "time": "0/1/2"}, + /// ["param", "time"], + /// ) + /// q.coxpand("time", {"step": ["s1", "s2"]}) + /// assert "time" not in q.dimensions() + /// assert "step" in q.dimensions() + /// ``` + pub fn coxpand( + &mut self, + axis: Bound<'_, PyAny>, + dimension: Bound<'_, PyDict>, + ) -> PyResult<()> { + self.collapse(axis)?; + self.expand(dimension) + } + + /// Check whether the Qube contains a given dimension or set of values. + /// + /// * **`str`** — returns ``True`` if the named dimension exists anywhere + /// in the Qube. + /// * **`dict[str, list]`** — returns ``True`` if every key in the dict is + /// a dimension that exists *and* every value listed for that key is + /// present in that dimension's coordinate set. + /// * **`Qube`** — checks that every dimension/value present in the other + /// Qube also exists in this one (subset check). + /// + /// ```python + /// q = Qube.from_datacube({"param": "2t/tp", "time": "0/1/2"}, ["param", "time"]) + /// + /// assert q.contains("param") is True + /// assert q.contains("level") is False + /// assert q.contains({"param": ["2t"]}) is True + /// assert q.contains({"param": ["xyz"]}) is False + /// assert q.contains({"time": ["0", "1"]}) is True + /// ``` + pub fn contains(&self, py: Python<'_>, item: Bound<'_, PyAny>) -> PyResult { + // Helper: convert a Coordinates value to a Vec of individual values. + let coords_to_strings = |coords: &::qubed::Coordinates| -> Vec { + let s = coords.to_string(); + if s.is_empty() { + vec![] + } else if s.contains('/') { + s.split('/').map(|v| v.to_owned()).collect() + } else { + vec![s] + } + }; + + // Pre-compute current axes as Vec per dimension. + let current_axes: std::collections::BTreeMap> = self + .inner + .all_unique_dim_coords() + .into_iter() + .map(|(k, v)| (k, coords_to_strings(&v))) + .collect(); + + // --- branch 1: item is a plain string (dimension name check) --- + if let Ok(dim_str) = item.extract::() { + return Ok(current_axes.contains_key(&dim_str)); + } + + // --- branch 2: item is a dict --- + if let Ok(dict) = item.downcast::() { + return helpers::check_dict_against_axes(dict, ¤t_axes, py); + } + + // --- branch 3: item is a Qube --- + if let Ok(other_cell) = item.downcast::() { + let other_ref = other_cell.borrow(); + let other_axes: std::collections::BTreeMap> = other_ref + .inner + .all_unique_dim_coords() + .into_iter() + .map(|(k, v)| (k, coords_to_strings(&v))) + .collect(); + + for (key, other_vals) in &other_axes { + match current_axes.get(key) { + None => return Ok(false), + Some(cur_vals) => { + for oval in other_vals { + if !cur_vals.contains(oval) { + return Ok(false); + } + } + } + } + } + return Ok(true); + } + + Err(PyTypeError::new_err("contains: item must be a str, dict[str, list], or Qube")) + } } fn pydict_to_datacube(datacube: Bound<'_, PyDict>) -> PyResult { diff --git a/py_qubed/tests/test_qubed_utils.py b/py_qubed/tests/test_qubed_utils.py new file mode 100644 index 0000000..1b0edb5 --- /dev/null +++ b/py_qubed/tests/test_qubed_utils.py @@ -0,0 +1,397 @@ +"""Tests for the utility methods added to mirror the @support_qubed_output +helpers from forecast-in-a-box/fiab_plugin_ecmwf/qubed_utils.py. + +Covered methods +--------------- +axes() – all dimension names → list of coordinate values +dimensions() – set of all dimension names +common_dimensions() – dimensions present in every leaf path +expand() – wrap tree under new outer dimension(s) +collapse() – remove dimension(s) (alias for drop with validation) +coxpand() – collapse then expand in one call +contains() – check whether a dimension / dict of values exists +""" +from qubed import Qube +import pytest + + +# --------------------------------------------------------------------------- +# Shared fixtures +# --------------------------------------------------------------------------- + +def make_param_time_qube() -> Qube: + """Qube with two dimensions: param={2t, tp} × time={0, 1, 2}.""" + return Qube.from_datacube({"param": "2t/tp", "time": "0/1/2"}, ["param", "time"]) + + +def make_three_dim_qube() -> Qube: + """Qube with three dimensions: class, param, time.""" + return Qube.from_datacube( + {"class": "od", "param": "2t/tp", "time": "0/1/2"}, + ["class", "param", "time"], + ) + + +# --------------------------------------------------------------------------- +# axes() +# --------------------------------------------------------------------------- + +class TestAxes: + def test_returns_dict_of_all_dimension_values(self): + q = make_param_time_qube() + ax = q.axes() + assert isinstance(ax, dict) + assert set(ax.keys()) == {"param", "time"} + + def test_param_values_are_correct(self): + q = make_param_time_qube() + assert set(q.axes()["param"]) == {"2t", "tp"} + + def test_integer_time_values_as_strings(self): + q = make_param_time_qube() + # Stored internally as integers; returned as strings by the binding. + assert set(q.axes()["time"]) == {"0", "1", "2"} + + def test_leading_zeros_preserved(self): + q = Qube.from_datacube({"expver": "0001/0002"}, None) + ax = q.axes() + assert "0001" in ax["expver"] + assert "0002" in ax["expver"] + + def test_does_not_include_root(self): + q = make_param_time_qube() + assert "root" not in q.axes() + + def test_axes_equals_all_unique_dim_coords(self): + q = make_param_time_qube() + assert q.axes() == q.all_unique_dim_coords() + + +# --------------------------------------------------------------------------- +# dimensions() +# --------------------------------------------------------------------------- + +class TestDimensions: + def test_returns_a_set(self): + q = make_param_time_qube() + assert isinstance(q.dimensions(), set) + + def test_correct_names(self): + q = make_param_time_qube() + assert q.dimensions() == {"param", "time"} + + def test_three_dimensions(self): + q = make_three_dim_qube() + assert q.dimensions() == {"class", "param", "time"} + + def test_excludes_root(self): + q = make_param_time_qube() + assert "root" not in q.dimensions() + + def test_empty_qube_has_no_dimensions(self): + assert Qube().dimensions() == set() + + def test_matches_axes_keys(self): + q = make_three_dim_qube() + assert q.dimensions() == set(q.axes().keys()) + + +# --------------------------------------------------------------------------- +# common_dimensions() +# --------------------------------------------------------------------------- + +class TestCommonDimensions: + def test_uniform_depth_returns_all_dims(self): + q = make_param_time_qube() + assert q.common_dimensions() == {"param", "time"} + + def test_irregular_tree_returns_intersection(self): + # Branch 1 has param + time; branch 2 has only param. + q1 = Qube.from_datacube({"param": "2t", "time": "0/1"}, ["param", "time"]) + q2 = Qube.from_datacube({"param": "msl"}, ["param"]) + q1.append(q2) + common = q1.common_dimensions() + assert "param" in common + assert "time" not in common + + def test_disjoint_dims_gives_empty_intersection(self): + q1 = Qube.from_datacube({"dim_a": "v1"}, None) + q2 = Qube.from_datacube({"dim_b": "v2"}, None) + q1.append(q2) + assert q1.common_dimensions() == set() + + def test_empty_qube(self): + assert Qube().common_dimensions() == set() + + def test_single_branch(self): + q = Qube.from_datacube({"a": "1", "b": "2", "c": "3"}, ["a", "b", "c"]) + assert q.common_dimensions() == {"a", "b", "c"} + + +# --------------------------------------------------------------------------- +# expand() +# --------------------------------------------------------------------------- + +class TestExpand: + def test_adds_new_outer_dimension(self): + q = make_param_time_qube() + q.expand({"ensemble": ["ens1", "ens2"]}) + assert "ensemble" in q.dimensions() + + def test_original_dimensions_preserved(self): + q = make_param_time_qube() + q.expand({"ensemble": ["ens1", "ens2"]}) + assert "param" in q.dimensions() + assert "time" in q.dimensions() + + def test_new_dimension_values_correct(self): + q = make_param_time_qube() + q.expand({"ensemble": ["ens1", "ens2"]}) + ax = q.axes() + assert set(ax["ensemble"]) == {"ens1", "ens2"} + + def test_original_values_unchanged(self): + q = make_param_time_qube() + original_params = set(q.axes()["param"]) + q.expand({"ensemble": ["ens1", "ens2"]}) + assert set(q.axes()["param"]) == original_params + + def test_expand_multiple_dimensions_in_one_call(self): + q = make_param_time_qube() + q.expand({"ensemble": ["ens1", "ens2"], "member": ["m1", "m2"]}) + dims = q.dimensions() + assert "ensemble" in dims + assert "member" in dims + + def test_expand_on_empty_qube(self): + q = Qube() + q.expand({"ensemble": ["ens1"]}) + assert "ensemble" in q.dimensions() + + def test_expand_twice_nests_last_outermost(self): + q = make_param_time_qube() + q.expand({"inner_dim": ["i1", "i2"]}) + q.expand({"outer_dim": ["o1", "o2"]}) + ascii_repr = q.to_ascii() + # outer_dim must appear before inner_dim in the ASCII tree + assert ascii_repr.index("outer_dim") < ascii_repr.index("inner_dim") + + def test_expand_integer_values(self): + q = make_param_time_qube() + q.expand({"step": [1, 2, 3]}) + ax = q.axes() + assert set(ax["step"]) == {"1", "2", "3"} + + def test_expand_is_reflected_in_datacubes(self): + q = make_param_time_qube() + q.expand({"ensemble": ["ens1", "ens2"]}) + # Every datacube should now include the ensemble dimension. + for dc in q.to_datacubes(): + assert "ensemble" in dc + + +# --------------------------------------------------------------------------- +# collapse() +# --------------------------------------------------------------------------- + +class TestCollapse: + def test_single_string_removes_dimension(self): + q = make_param_time_qube() + q.collapse("time") + assert "time" not in q.dimensions() + + def test_original_dims_preserved_after_collapse(self): + q = make_param_time_qube() + q.collapse("time") + assert "param" in q.dimensions() + + def test_list_of_dimensions(self): + q = make_three_dim_qube() + q.collapse(["class", "time"]) + dims = q.dimensions() + assert "class" not in dims + assert "time" not in dims + assert "param" in dims + + def test_nonexistent_dimension_raises_value_error(self): + q = make_param_time_qube() + with pytest.raises(ValueError, match="nonexistent"): + q.collapse("nonexistent") + + def test_nonexistent_in_list_raises_value_error(self): + q = make_param_time_qube() + with pytest.raises(ValueError): + q.collapse(["param", "does_not_exist"]) + + def test_collapse_and_check_via_contains(self): + q = make_param_time_qube() + q.collapse("time") + assert q.contains("time") is False + assert q.contains("param") is True + + def test_collapse_then_dimensions_updated(self): + q = make_three_dim_qube() + q.collapse("class") + assert q.dimensions() == {"param", "time"} + + def test_wrong_type_raises(self): + q = make_param_time_qube() + with pytest.raises((TypeError, ValueError)): + q.collapse(42) # type: ignore[arg-type] + + +# --------------------------------------------------------------------------- +# coxpand() +# --------------------------------------------------------------------------- + +class TestCoxpand: + def test_removes_collapsed_dimension(self): + q = make_param_time_qube() + q.coxpand("time", {"step": ["s1", "s2"]}) + assert "time" not in q.dimensions() + + def test_adds_expanded_dimension(self): + q = make_param_time_qube() + q.coxpand("time", {"step": ["s1", "s2"]}) + assert "step" in q.dimensions() + + def test_preserves_non_touched_dimension(self): + q = make_param_time_qube() + q.coxpand("time", {"step": ["s1", "s2"]}) + assert "param" in q.dimensions() + + def test_new_dimension_values(self): + q = make_param_time_qube() + q.coxpand("time", {"step": ["s1", "s2"]}) + assert set(q.axes()["step"]) == {"s1", "s2"} + + def test_list_axis_collapse(self): + q = make_three_dim_qube() + q.coxpand(["class", "time"], {"batch": ["b1"]}) + dims = q.dimensions() + assert "class" not in dims + assert "time" not in dims + assert "batch" in dims + assert "param" in dims + + def test_coxpand_on_nonexistent_raises(self): + q = make_param_time_qube() + with pytest.raises(ValueError): + q.coxpand("nonexistent", {"new": ["v1"]}) + + +# --------------------------------------------------------------------------- +# contains() +# --------------------------------------------------------------------------- + +class TestContains: + # --- string (dimension name) checks --- + def test_existing_dimension_returns_true(self): + q = make_param_time_qube() + assert q.contains("param") is True + + def test_missing_dimension_returns_false(self): + q = make_param_time_qube() + assert q.contains("level") is False + + def test_root_is_not_a_dimension(self): + q = make_param_time_qube() + assert q.contains("root") is False + + # --- dict checks --- + def test_dict_existing_key_and_value_returns_true(self): + q = make_param_time_qube() + assert q.contains({"param": ["2t"]}) is True + + def test_dict_missing_value_returns_false(self): + q = make_param_time_qube() + assert q.contains({"param": ["xyz"]}) is False + + def test_dict_missing_dimension_returns_false(self): + q = make_param_time_qube() + assert q.contains({"level": ["1000"]}) is False + + def test_dict_multiple_dims_all_present(self): + q = make_param_time_qube() + assert q.contains({"param": ["2t", "tp"], "time": ["0", "1"]}) is True + + def test_dict_one_value_absent_returns_false(self): + q = make_param_time_qube() + # "0" exists but "999" does not + assert q.contains({"time": ["0", "999"]}) is False + + def test_dict_single_value_string(self): + q = make_param_time_qube() + # scalar value (not a list) should still work + assert q.contains({"param": "2t"}) is True + + def test_dict_integer_value_as_string(self): + q = make_param_time_qube() + # integers are stored as integers internally but returned as strings + assert q.contains({"time": ["0"]}) is True + + def test_empty_dict_always_true(self): + q = make_param_time_qube() + assert q.contains({}) is True + + # --- Qube checks --- + def test_qube_subset_returns_true(self): + q = make_param_time_qube() + # A Qube that only covers "2t" and time "0" is a subset. + subset = Qube.from_datacube({"param": "2t", "time": "0"}, ["param", "time"]) + assert q.contains(subset) is True + + def test_qube_with_extra_value_returns_false(self): + q = make_param_time_qube() + other = Qube.from_datacube({"param": "999"}, None) + assert q.contains(other) is False + + def test_qube_with_extra_dimension_returns_false(self): + q = make_param_time_qube() + other = Qube.from_datacube({"level": "1000"}, None) + assert q.contains(other) is False + + def test_invalid_type_raises(self): + q = make_param_time_qube() + with pytest.raises(TypeError): + q.contains(42) # type: ignore[arg-type] + + +# --------------------------------------------------------------------------- +# Integration: expand → collapse → contains round-trip +# --------------------------------------------------------------------------- + +class TestRoundTrips: + def test_expand_then_collapse_restores_original_dims(self): + q = make_param_time_qube() + original_dims = q.dimensions() + q.expand({"ensemble": ["e1", "e2"]}) + q.collapse("ensemble") + # After round-trip the ensemble dimension should be gone. + assert "ensemble" not in q.dimensions() + # Original dimensions should still be present. + for d in original_dims: + assert d in q.dimensions() + + def test_expanded_qube_contains_new_dimension(self): + q = make_param_time_qube() + q.expand({"ensemble": ["ens1", "ens2"]}) + assert q.contains("ensemble") is True + assert q.contains({"ensemble": ["ens1"]}) is True + assert q.contains({"ensemble": ["unknown"]}) is False + + def test_coxpand_then_contains(self): + q = make_three_dim_qube() + q.coxpand("class", {"run_type": ["fc", "an"]}) + assert q.contains("class") is False + assert q.contains("run_type") is True + assert q.contains({"run_type": ["fc", "an"]}) is True + + def test_common_dimensions_after_expand_is_consistent(self): + q = make_param_time_qube() + q.expand({"ensemble": ["ens1", "ens2"]}) + # Since expand wraps the whole tree uniformly, all leaf paths have ensemble. + common = q.common_dimensions() + assert "ensemble" in common + assert "param" in common + assert "time" in common diff --git a/qubed/src/qube.rs b/qubed/src/qube.rs index c909495..e5c4d74 100644 --- a/qubed/src/qube.rs +++ b/qubed/src/qube.rs @@ -194,15 +194,14 @@ impl Qube { Ok(node_id) } - pub fn all_unique_dim_coords(&mut self) -> BTreeMap { - // TODO + pub fn all_unique_dim_coords(&self) -> BTreeMap { let mut map: BTreeMap = BTreeMap::new(); for (_id, node) in self.nodes.iter() { if let Some(dim_str) = self.dimension_str(&node.dim) { let coords = node.coords.clone(); if coords.is_empty() { - continue; // Skip empty coordinates + continue; // Skip empty coordinates (incl. the virtual root node) } // if there is no entry in map for this dimension, just fill it in with coords, otherwise extend the current entry with coords map.entry(dim_str.to_string()) @@ -213,6 +212,135 @@ impl Qube { map } + /// Returns the set of all dimension names present anywhere in the Qube. + /// + /// This is the set of keys from [`all_unique_dim_coords`]. + /// + /// # Examples + /// ``` + /// use qubed::Qube; + /// let q = Qube::from_ascii("root\n└── class=od\n └── param=1/2").unwrap(); + /// let dims = q.dimensions(); + /// assert!(dims.contains("class")); + /// assert!(dims.contains("param")); + /// assert!(!dims.contains("root")); + /// ``` + pub fn dimensions(&self) -> HashSet { + self.all_unique_dim_coords().into_keys().collect() + } + + /// Returns the set of dimension names present in **every** leaf path (datacube). + /// + /// For a Qube with uniform depth this equals [`dimensions`]. For an + /// irregular Qube some branches may be missing a dimension; only those + /// that appear in *all* branches are returned. + /// + /// # Examples + /// ``` + /// use qubed::Qube; + /// use qubed::Datacube; + /// use qubed::Coordinates; + /// + /// // Both datacubes share "param"; only one has "time". + /// let mut dc1 = Datacube::new(); + /// dc1.add_coordinate("param", Coordinates::from_string("2t/tp")); + /// dc1.add_coordinate("time", Coordinates::from_string("0/1/2")); + /// let mut qube = Qube::from_datacube(&dc1, Some(&["param".to_string(), "time".to_string()])); + /// + /// let mut dc2 = Datacube::new(); + /// dc2.add_coordinate("param", Coordinates::from_string("msl")); + /// let mut other = Qube::from_datacube(&dc2, None); + /// + /// qube.append(&mut other); + /// + /// let common = qube.common_dimensions(); + /// assert!(common.contains("param")); + /// assert!(!common.contains("time")); + /// ``` + pub fn common_dimensions(&self) -> HashSet { + let datacubes = self.to_datacubes(); + if datacubes.is_empty() { + return HashSet::new(); + } + + let mut iter = datacubes.iter().map(|dc| { + dc.coordinates() + .iter() + .filter(|(_, v)| !v.is_empty()) // exclude the virtual root node + .map(|(k, _)| k.clone()) + .collect::>() + }); + + let first = match iter.next() { + Some(s) => s, + None => return HashSet::new(), + }; + + iter.fold(first, |acc, keys| acc.intersection(&keys).cloned().collect()) + } + + /// Wraps the entire existing Qube tree under a new outer dimension. + /// + /// All current children of the root are re-parented to a new node that + /// has the given `key` and `values` as its dimension and coordinates. + /// The result is that `key` becomes the outermost dimension of the Qube. + /// + /// Calling `expand` multiple times nests the dimensions from the inside + /// out: each call wraps the *current* tree, so the last call produces + /// the outermost dimension. + /// + /// # Examples + /// ``` + /// use qubed::{Qube, Coordinates}; + /// + /// let mut q = Qube::from_ascii("root\n└── param=2t/tp\n └── time=0/1/2").unwrap(); + /// q.expand("ensemble", Coordinates::from_string("ens1/ens2")).unwrap(); + /// + /// let dims = q.dimensions(); + /// assert!(dims.contains("ensemble")); + /// assert!(dims.contains("param")); + /// assert!(dims.contains("time")); + /// ``` + pub fn expand(&mut self, key: &str, values: Coordinates) -> Result<(), String> { + let root_id = self.root_id; + + // 1. Clone root's current children before any mutation. + let old_root_children: BTreeMap> = self + .nodes + .get(root_id) + .ok_or_else(|| "Root node not found".to_string())? + .children + .clone(); + + // 2. Clear root's children so get_or_create_child starts with a clean slate. + if let Some(root) = self.nodes.get_mut(root_id) { + root.children.clear(); + root.structural_hash.store(0, Ordering::Release); + } + + // 3. Create the new dimension node as the sole child of root. + let new_node_id = self.get_or_create_child(key, root_id, Some(values))?; + + // 4. Move the saved children into the new node. + if let Some(new_node) = self.nodes.get_mut(new_node_id) { + new_node.children = old_root_children.clone(); + } + + // 5. Fix parent pointers for the moved subtree roots. + for child_ids in old_root_children.values() { + for &child_id in child_ids.iter() { + if let Some(child) = self.nodes.get_mut(child_id) { + child.parent = Some(new_node_id); + } + } + } + + // 6. Invalidate cached structural hashes up to (and including) root. + self.invalidate_ancestors(new_node_id); + + Ok(()) + } + pub fn remove_node(&mut self, id: NodeIdx) -> Result<(), String> { let node = self.nodes.remove(id).ok_or_else(|| format!("Node {:?} not found", id))?; @@ -799,6 +927,99 @@ mod tests { assert!(class1_node.children(qube.dimension("expver").unwrap()).is_some()); } + #[test] + fn test_dimensions_returns_dim_names() { + let q = Qube::from_ascii( + "root\n└── class=od\n ├── expver=0001\n │ └── param=1\n └── expver=0002\n └── param=2", + ) + .unwrap(); + let dims = q.dimensions(); + assert!(dims.contains("class")); + assert!(dims.contains("expver")); + assert!(dims.contains("param")); + assert!(!dims.contains("root"), "root should not appear as a dimension"); + } + + #[test] + fn test_common_dimensions_uniform_depth() { + let q = Qube::from_ascii("root\n└── class=od\n └── param=1/2").unwrap(); + let common = q.common_dimensions(); + assert!(common.contains("class")); + assert!(common.contains("param")); + } + + #[test] + fn test_common_dimensions_irregular_depth() { + use crate::Datacube; + // Branch 1: param + time + let mut dc1 = Datacube::new(); + dc1.add_coordinate("param", Coordinates::from_string("2t/tp")); + dc1.add_coordinate("time", Coordinates::from_string("0/1/2")); + let mut qube = Qube::from_datacube(&dc1, Some(&["param".to_string(), "time".to_string()])); + + // Branch 2: only param + let mut dc2 = Datacube::new(); + dc2.add_coordinate("param", Coordinates::from_string("msl")); + let mut other = Qube::from_datacube(&dc2, None); + + qube.append(&mut other); + + let common = qube.common_dimensions(); + assert!(common.contains("param"), "'param' should be common"); + assert!(!common.contains("time"), "'time' is absent in one branch"); + } + + #[test] + fn test_expand_wraps_tree_under_new_outer_dimension() { + let mut q = Qube::from_ascii("root\n└── param=2t/tp\n └── time=0/1/2").unwrap(); + q.expand("ensemble", Coordinates::from_string("ens1/ens2")).unwrap(); + + let dims = q.dimensions(); + assert!(dims.contains("ensemble")); + assert!(dims.contains("param")); + assert!(dims.contains("time")); + + let ascii = q.to_ascii(); + assert!(ascii.contains("ensemble=ens1/ens2"), "new dimension should appear in ascii"); + } + + #[test] + fn test_expand_on_empty_qube() { + let mut q = Qube::new(); + q.expand("ensemble", Coordinates::from_string("ens1/ens2")).unwrap(); + + let dims = q.dimensions(); + assert!(dims.contains("ensemble")); + } + + #[test] + fn test_expand_twice_nests_outermost_last() { + let mut q = Qube::from_ascii("root\n└── param=2t").unwrap(); + q.expand("ensemble", Coordinates::from_string("ens1")).unwrap(); + q.expand("member", Coordinates::from_string("m1/m2")).unwrap(); + + let ascii = q.to_ascii(); + // "member" was added last so it must appear higher (earlier) in the tree + let member_pos = ascii.find("member").expect("member not found"); + let ensemble_pos = ascii.find("ensemble").expect("ensemble not found"); + let param_pos = ascii.find("param").expect("param not found"); + assert!(member_pos < ensemble_pos, "member should be outer of ensemble"); + assert!(ensemble_pos < param_pos, "ensemble should be outer of param"); + } + + #[test] + fn test_expand_preserves_original_coords() { + let mut q = Qube::from_ascii("root\n└── param=2t/tp\n └── time=0/1/2").unwrap(); + q.expand("ensemble", Coordinates::from_string("ens1/ens2")).unwrap(); + + let all = q.all_unique_dim_coords(); + // Original coords must still be present + let param_str = all.get("param").unwrap().to_string(); + assert!(param_str.contains("2t") && param_str.contains("tp")); + let ens_str = all.get("ensemble").unwrap().to_string(); + assert!(ens_str.contains("ens1") && ens_str.contains("ens2")); + } + #[test] fn test_squeeze() -> Result<(), String> { let input = r#"root