def _get_data():
# Make the tree
site_nodes = LengthNode(2)
line_nodes = [LengthNode(1), LengthNode(1)]
tank_nodes_site0 = [LengthNode(3)]
tank_nodes_site1 = [LengthNode(4)]
site_nodes.set_children(line_nodes)
line_nodes[0].set_children(tank_nodes_site0)
line_nodes[1].set_children(tank_nodes_site1)
nodes = [site_nodes]
nodes.extend(line_nodes)
nodes.extend(tank_nodes_site0)
nodes.extend(tank_nodes_site1)
level_names = "IDX_SITE", "IDX_LINE", "IDX_TANK"
flat_results = FlatData({
"sites[0].lines[0].tanks[0].Q_yield": 10.0,
"sites[0].lines[0].tanks[1].Q_yield": 6.0,
"sites[0].lines[0].tanks[2].Q_yield": 4.0,
"sites[1].lines[0].tanks[0].Q_yield": 20.0,
"sites[1].lines[0].tanks[1].Q_yield": 8.0,
"sites[1].lines[0].tanks[2].Q_yield": 3.0,
"sites[1].lines[0].tanks[3].Q_yield": 1.0,
})
tree = LengthTree(nodes, level_names)
print("tree", tree)
return tree, flat_results
def test_expand_path2(self):
"""Expand path"""
paths = (
HubitModelPath("segments[:@IDX_SEG].layers[:@IDX_LAY].test"),
HubitQueryPath("segments[:].layers[:].test"),
)
# path = "segments.:@IDX_SEG.layers.:@IDX_LAY.test"
seg_node = LengthNode(2)
lay_nodes = LengthNode(2), LengthNode(2)
seg_node.set_children(lay_nodes)
nodes = [seg_node]
nodes.extend(lay_nodes)
level_names = "IDX_SEG", "IDX_LAY"
tree = LengthTree(nodes, level_names)
# 2 values for segment 0 and 2 values for segment 1
expected_paths = [
[
"segments[0].layers[0].test",
"segments[0].layers[1].test",
],
[
"segments[1].layers[0].test",
"segments[1].layers[1].test",
],
]
for path in paths:
with self.subTest(path=path):
expanded_paths = tree.prune_from_path(path).expand_path(path)
self.assertSequenceEqual(expanded_paths, expected_paths)
def test_prune(self):
def print_test(tree):
nodes = tree.nodes_for_level[0]
for node in nodes:
print("node", node)
for child in node.children:
print("child", child)
idx_car_node = LengthNode(2)
idx_parts_nodes = LengthNode(5), LengthNode(4)
idx_car_node.set_children(idx_parts_nodes)
nodes = [idx_car_node]
nodes.extend(idx_parts_nodes)
level_names = "IDX_CAR", "IDX_CAR"
tree = LengthTree(nodes, level_names)
clipped_tree = tree.clip_at_level("IDX_CAR", inplace=False)
# nodes or bottom level
nodes = clipped_tree.nodes_for_level[-1]
# all children should be None at bottom level
children_are_leaves = [
all([isinstance(child, LeafNode) for child in node.children])
for node in nodes
]
with self.subTest():
self.assertTrue(all(children_are_leaves))
with self.subTest():
self.assertTrue(len(clipped_tree.level_names) == 1)
with self.subTest():
self.assertTrue(len(clipped_tree.nodes_for_level) == 1)
def test_expand_path_count_from_back(self):
"""Expand path with fixed index set to negative number"""
path = HubitQueryPath("segments[:].layers[-2].test")
seg_node = LengthNode(3)
lay_nodes = LengthNode(2), LengthNode(4), LengthNode(3)
seg_node.set_children(lay_nodes)
nodes = [seg_node]
nodes.extend(lay_nodes)
level_names = "IDX_SEG", "IDX_LAY"
tree = LengthTree(nodes, level_names)
expected_paths = [
"segments[0].layers[0].test",
"segments[1].layers[2].test",
"segments[2].layers[1].test",
]
# Since the tree is not pruned we must use flat=True
expanded_paths = tree.prune_from_path(path, inplace=False).expand_path(
path, flat=True)
self.assertSequenceEqual(expanded_paths, expected_paths)
# Index error. There are 2 layers on segment at index 0.
path = HubitQueryPath("segments[:].layers[-3].test")
with self.assertRaises(HubitIndexError) as context:
expanded_paths = tree.prune_from_path(
path, inplace=False).expand_path(path, flat=True)
tree, _ = _get_data()
qpath = HubitQueryPath("sites[1].lines[0].tanks[-1].Q_yield")
expanded_paths = tree.prune_from_path(
qpath, inplace=False).expand_path(qpath, flat=True)
assert expanded_paths == ["sites[1].lines[0].tanks[3].Q_yield"]
qpath = HubitQueryPath("sites[0].lines[-1].tanks[-1].Q_yield")
expanded_paths = tree.prune_from_path(
qpath, inplace=False).expand_path(qpath, flat=True)
assert expanded_paths == ["sites[0].lines[0].tanks[2].Q_yield"]
qpath = HubitQueryPath("sites[:].lines[-1].tanks[-1].Q_yield")
expanded_paths = tree.prune_from_path(
qpath, inplace=False).expand_path(qpath, flat=True)
assert expanded_paths == [
"sites[0].lines[0].tanks[2].Q_yield",
"sites[1].lines[0].tanks[3].Q_yield",
]
qpath = HubitQueryPath("sites[:].lines[:].tanks[:].Q_yield")
expanded_paths = tree.prune_from_path(
qpath, inplace=False).expand_path(qpath, flat=True)
def test_8(self):
"""
Test compression of indices. The query does not include any indices
"""
func = None
version = None
comp_yml = """
path: dummy,
func_name: dummy,
provides_results:
- name: mylist
path: factors
consumes_input:
- name: factors
path: list[:@IDX1].some_attr.factors
"""
component = HubitModelComponent.from_cfg(
yaml.load(comp_yml, Loader=yaml.FullLoader), 0
)
# Query something known to exist
querystr = "factors"
idx1_node = LengthNode(2)
nodes = [idx1_node]
level_names = ("IDX1",)
tree = LengthTree(nodes, level_names)
dummy_tree = DummyLengthTree()
_tree_for_idxcontext = {"": dummy_tree, tree.index_context: tree}
querystr = HubitQueryPath(querystr)
w = _Worker(
lambda x: x,
lambda x: x,
component,
querystr,
func,
version,
_tree_for_idxcontext,
dryrun=True,
)
def test_5(self):
"""
Initialize worker with ILOC locations in
query and ILOC wildcards in bindings
"""
func = None
version = None
cfg = """
path: dummy,
func_name: dummy,
provides_results:
- name: k_therm
path: segments[IDX_SEG].layers[:@IDX_LAY].k_therm
consumes_input:
- name: material
path: segments[IDX_SEG].layers[:@IDX_LAY].material
"""
component = HubitModelComponent.from_cfg(
yaml.load(cfg, Loader=yaml.FullLoader), 0
)
# Query something known to exist
querystr = HubitQueryPath("segments[0].layers[0].k_therm")
seg_node = LengthNode(2)
lay_nodes = LengthNode(2), LengthNode(2)
seg_node.set_children(lay_nodes)
nodes = [seg_node]
nodes.extend(lay_nodes)
level_names = "IDX_SEG", "IDX_LAY"
tree = LengthTree(nodes, level_names)
_tree_for_idxcontext = {tree.index_context: tree}
querystr = HubitQueryPath(querystr)
w = _Worker(
lambda x: x,
lambda x: x,
component,
querystr,
func,
version,
_tree_for_idxcontext,
dryrun=True,
)
def setUp(self):
# lengths = [['IDX_SEG', 2],
# ['IDX_LAY', [3, 4]],
# ['IDX_POS', [[1, 3, 2], [5, 1, 2, 4]]]]
seg_nodes = LengthNode(2)
lay_nodes = LengthNode(3), LengthNode(4)
pos_lay0_nodes = (
LengthNode(1),
LengthNode(3),
LengthNode(2),
)
pos_lay1_nodes = (
LengthNode(5),
LengthNode(1),
LengthNode(2),
LengthNode(4),
)
seg_nodes.set_children(lay_nodes)
lay_nodes[0].set_children(pos_lay0_nodes)
lay_nodes[1].set_children(pos_lay1_nodes)
nodes = [seg_nodes]
nodes.extend(lay_nodes)
nodes.extend(pos_lay0_nodes)
nodes.extend(pos_lay1_nodes)
level_names = "IDX_SEG", "IDX_LAY", "IDX_POS"
self.tree = LengthTree(nodes, level_names)
self.template_path = HubitModelPath(
"segments[:@IDX_SEG].layers[:@IDX_LAY].test.positions[:@IDX_POS]")
def test_7(self):
"""4-level tree"""
x1_nodes = LengthNode(2)
x2_nodes = LengthNode(1), LengthNode(3)
x3_0_nodes = [LengthNode(2)]
x3_1_nodes = LengthNode(1), LengthNode(2), LengthNode(4)
x4_0_0_nodes = LengthNode(1), LengthNode(3)
x4_1_0_nodes = [LengthNode(1)]
x4_1_1_nodes = LengthNode(2), LengthNode(2)
x4_1_2_nodes = (
LengthNode(1),
LengthNode(1),
LengthNode(1),
LengthNode(2),
)
x1_nodes.set_children(x2_nodes)
x2_nodes[0].set_children(x3_0_nodes)
x2_nodes[1].set_children(x3_1_nodes)
x3_0_nodes[0].set_children(x4_0_0_nodes)
x3_1_nodes[0].set_children(x4_1_0_nodes)
x3_1_nodes[1].set_children(x4_1_1_nodes)
x3_1_nodes[2].set_children(x4_1_2_nodes)
nodes = [x1_nodes]
nodes.extend(x2_nodes)
nodes.extend(x3_0_nodes)
nodes.extend(x3_1_nodes)
nodes.extend(x4_0_0_nodes)
nodes.extend(x4_1_0_nodes)
nodes.extend(x4_1_1_nodes)
nodes.extend(x4_1_2_nodes)
level_names = (
"IDX_X1",
"IDX_X2",
"IDX_X3",
"IDX_X4",
)
tree = LengthTree(nodes, level_names)
expected_lengths = [
2,
[1, 3],
[[2], [1, 2, 4]],
[[[1, 3]], [[1], [2, 2], [1, 1, 1, 2]]],
]
self.assertListEqual(tree.to_list(), expected_lengths)