diff --git a/docs/quickstart.md b/docs/quickstart.md index ac0790c..a761776 100644 --- a/docs/quickstart.md +++ b/docs/quickstart.md @@ -45,18 +45,39 @@ cq Load a larger example qube (requires source checkout): ```{code-cell} python3 -from pathlib import Path -import json -data_path = Path("../tests/example_qubes/climate_dt.json") -with data_path.open("r") as f: - climate_dt = Qube.from_json(json.loads(f.read())) +import requests +qube_json = requests.get("https://github.com/ecmwf/qubed/raw/refs/heads/main/tests/example_qubes/climate_dt.json").json() +climate_dt = Qube.from_json(qube_json) # Using the html or print methods is optional but lets you specify things like the depth of the tree to display. print(f"{climate_dt.n_leaves = }, {climate_dt.n_nodes = }") climate_dt.html(depth=1) # Limit how much is open initially, click leave to see more. ``` -### Set Operations +Select a subset of the tree: + +```{code-cell} python3 +climate_dt.select({ + "activity": "scenariomip" +}).html(depth=1) +``` + +Use `.span("key")` to get the set of possibles values for a key, note this includes anywhere this key appears in the tree. + +```{code-cell} python3 +climate_dt.span("activity") +``` + +Use `.axes()` to get the span of every key in one go. + +```{code-cell} python3 +axes = climate_dt.axes() +for key, values in axes.items(): + print(f"{key} : {list(values)[:10]}") +``` + + + -### Command Line Usage + diff --git a/src/python/qubed/Qube.py b/src/python/qubed/Qube.py index 8020a1a..be0d951 100644 --- a/src/python/qubed/Qube.py +++ b/src/python/qubed/Qube.py @@ -142,10 +142,28 @@ class Qube: if not values: return None - return dataclasses.replace(node, values = values, children = not_none(select(c) for c in node.children)) + data = dataclasses.replace(node.data, values = values) + return dataclasses.replace(node, data=data, children = not_none(select(c) for c in node.children)) return dataclasses.replace(self, children = not_none(select(c) for c in self.children)) + def span(self, key: str) -> list[str]: + """ + Search the whole tree for any value that a given key takes anywhere. + """ + this = set(self.values) if self.key == key else set() + return sorted(this | set(v for c in self.children for v in c.span(key))) + + def axes(self) -> dict[str, set[str]]: + """ + Return a dictionary of all the spans of the keys in the qube. + """ + axes = defaultdict(set) + for c in self.children: + for k, v in c.axes().items(): + axes[k].update(v) + axes[self.key].update(self.values) + return dict(axes) @staticmethod def _insert(position: "Qube", identifier : list[tuple[str, list[str]]]):