From 62c7a49c596ea5bc90b75bce43c4a4b830bd2af1 Mon Sep 17 00:00:00 2001 From: Tom Date: Thu, 13 Feb 2025 17:32:14 +0000 Subject: [PATCH] Start filling out set operations --- docs/quickstart.md | 42 ++++++++++++++++++++++++++---- src/python/qubed/Qube.py | 26 +++++++++++++++--- src/python/qubed/set_operations.py | 21 +++++++++++++++ tests/test_basic_operations.py | 27 +++++++++++++++++++ 4 files changed, 107 insertions(+), 9 deletions(-) create mode 100644 src/python/qubed/set_operations.py create mode 100644 tests/test_basic_operations.py diff --git a/docs/quickstart.md b/docs/quickstart.md index ec4fbc3..7c624ad 100644 --- a/docs/quickstart.md +++ b/docs/quickstart.md @@ -9,7 +9,7 @@ jupytext: # Quickstart ## Installation -``` +```bash pip install qubed ``` @@ -29,17 +29,49 @@ q = Qube.from_dict({ "expver=0002": {"param=1":{}, "param=2":{}}, }, }) +print(f"{q.n_leaves = }, {q.n_nodes = }") q ``` -Compress the qube: +Compress it: ```{code-cell} python3 -q.compress() +cq = q.compress() +assert cq.n_leaves == q.n_leaves +print(f"{cq.n_leaves = }, {cq.n_nodes = }") +cq ``` -Load some example qubes: +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())) + +# 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 + +```{code-cell} python3 +A = Qube.from_dict({ + "a=1/2/3" : {"b=1/2/3" : {"c=1/2/3" : {}}}, + "a=5" : { "b=4" : { "c=4" : {}}} + }) + +B = Qube.from_dict({ + "a=1/2/3" : {"b=1/2/3" : {"c=1/2/3" : {}}}, + "a=5" : { "b=4" : { "c=4" : {}}} + }) + +A.print(name="A"), B.print(name="B"); + +A | B +``` + -### Set Operations \ No newline at end of file diff --git a/src/python/qubed/Qube.py b/src/python/qubed/Qube.py index 10a6057..93c2c79 100644 --- a/src/python/qubed/Qube.py +++ b/src/python/qubed/Qube.py @@ -6,6 +6,7 @@ from typing import Any, Callable, Hashable, Literal, Mapping from frozendict import frozendict +from . import set_operations from .tree_formatters import HTML, node_tree_to_html, node_tree_to_string from .value_types import DateRange, Enum, IntRange, TimeRange, Values @@ -31,6 +32,12 @@ class NodeData: def summary(self) -> str: return f"{self.key}={self.values.summary()}" if self.key != "root" else "root" + +@dataclass(frozen=True, eq=True, order=True) +class RootNodeData(NodeData): + "Helper class to print a custom root name" + def summary(self) -> str: + return self.key @dataclass(frozen=True, eq=True, order=True) class Qube: @@ -91,16 +98,21 @@ class Qube: return cls.make("root", Enum(("root",)), []) - def __str__(self, depth = None) -> str: - return "".join(node_tree_to_string(node=self, depth = depth)) + def __str__(self, depth = None, name = None) -> str: + node = dataclasses.replace(self, data = RootNodeData(key = name, values=self.values, metadata=self.metadata)) if name is not None else self + return "".join(node_tree_to_string(node=node, depth = depth)) - def print(self, depth = None): print(self.__str__(depth = depth)) + def print(self, depth = None, name: str | None = None): + print(self.__str__(depth = depth, name = name)) def html(self, depth = 2, collapse = True) -> HTML: return HTML(node_tree_to_html(self, depth = depth, collapse = collapse)) def _repr_html_(self) -> str: return node_tree_to_html(self, depth = 2, collapse = True) + + def __or__(self, other: "Qube") -> "Qube": + return set_operations.operation(self, other, set_operations.SetOperation.UNION) def __getitem__(self, args) -> 'Qube': @@ -111,7 +123,13 @@ class Qube: return dataclasses.replace(c, data = data) raise KeyError(f"Key {key} not found in children of {self.key}") - + @cached_property + def n_leaves(self) -> int: + return len(self.values) * (sum(c.n_leaves for c in self.children) if self.children else 1) + + @cached_property + def n_nodes(self) -> int: + return 1 + sum(c.n_nodes for c in self.children) def transform(self, func: 'Callable[[Qube], Qube | list[Qube]]') -> 'Qube': """ diff --git a/src/python/qubed/set_operations.py b/src/python/qubed/set_operations.py new file mode 100644 index 0000000..74d21fb --- /dev/null +++ b/src/python/qubed/set_operations.py @@ -0,0 +1,21 @@ +from enum import Enum +from collections import defaultdict + + +class SetOperation(Enum): + UNION = (1, 1, 1) + INTERSECTION = (0, 1, 0) + DIFFERENCE = (1, 0, 0) + SYMMETRIC_DIFFERENCE = (1, 0, 1) + + +def operation(A: "Qube", B : "Qube", type: SetOperation) -> "Qube": + # Sort nodes from both qubes by their keys + nodes_by_key = defaultdict(lambda : dict(A = [], B = [])) + for node in A.nodes: + nodes_by_key[node.key]["A"].append(node) + for key, ndoes + +# The root node is special so we need a helper method that we can recurse on +def _operation(A: list["Qube"], B : list["Qube"], type: SetOperation) -> "Qube": + pass \ No newline at end of file diff --git a/tests/test_basic_operations.py b/tests/test_basic_operations.py new file mode 100644 index 0000000..dd2883f --- /dev/null +++ b/tests/test_basic_operations.py @@ -0,0 +1,27 @@ +from qubed import Qube + + +def test_eq(): + d = { + "class=od" : { + "expver=0001": {"param=1":{}, "param=2":{}}, + "expver=0002": {"param=1":{}, "param=2":{}}, + }, + "class=rd" : { + "expver=0001": {"param=1":{}, "param=2":{}, "param=3":{}}, + "expver=0002": {"param=1":{}, "param=2":{}}, + }, + } + q = Qube.from_dict(d) + r = Qube.from_dict(d) + + assert q == r + +def test_n_leaves(): + q = Qube.from_dict({ + "a=1/2/3" : {"b=1/2/3" : {"c=1/2/3" : {}}}, + "a=5" : { "b=4" : { "c=4" : {}}} + }) + + # Size is 3*3*3 + 1*1*1 = 27 + 1 + assert q.n_leaves == 27 + 1 \ No newline at end of file