Hierarchical search spaces#

This notebook demonstrates Trieste’s HierarchicalSearchSpace, a search space that represents the conditional activation structure of a Generalized Disjunctive Program (GDP). Variables in such a space are gated by indicator variables (binary or categorical) that decide which other dimensions are meaningful.

The core primitives introduced:

BooleanSearchSpace: a one-dimensional discrete space restricted to {0, 1}.
CategoricalSearchSpace: a one-dimensional discrete space over K named or indexed categories.
HierarchyNode: a trieste tag-based primitive describing one disjunction via subspace_tags (the non-indicator subspaces it owns) and activity_condition_tags ({indicator_tag: required_value} gating the node). The owning HierarchicalSearchSpace resolves these tags to GPflow’s column-based representation internally (see space.to_gpflow_hierarchy()), so subspaces are supplied once, to the space, rather than to every node.
HierarchicalSearchSpace: a CollectionSearchSpace that wires the indicators and nodes together and exposes a tag-based query API for downstream consumers.

We walk through the API on two small examples: a Boolean indicator (the canonical four-variable disjunction) and a 3-ary categorical indicator. The notebook stops at the search-space layer.

[1]:

import numpy as np
import tensorflow as tf

from trieste.space import (
    BooleanSearchSpace,
    Box,
    CategoricalSearchSpace,
    HierarchicalSearchSpace,
    HierarchyNode,
)

np.random.seed(1793)
tf.random.set_seed(1793)

2026-06-17 10:47:28,791 INFO util.py:154 -- Missing packages: ['ipywidgets']. Run `pip install -U ipywidgets`, then restart the notebook server for rich notebook output.

Boolean-indicator example#

Consider a four-variable space with one Boolean indicator. Variable \(x_1\) is continuous and always active; \(y_1\) is the indicator; \(x_2\) is continuous and active only when \(y_1 = 1\); \(x_3\) is continuous and active only when \(y_1 = 0\).

We give each subspace a tag in the subspaces mapping (iteration order defines the flat-vector layout), then describe the hierarchy as three nodes — one unconditional (“shared”) and two gated by y1. The HierarchicalSearchSpace resolves the tag-based nodes to GPflow’s column-based representation internally.

[2]:

subspaces = {
    "x1": Box([0.0], [1.0]),  # unconditional
    "y1": BooleanSearchSpace(),  # Boolean indicator
    "x2": Box([0.0], [5.0]),  # active when y1 = 1
    "x3": Box([-1.0], [1.0]),  # active when y1 = 0
}
hierarchy = [
    HierarchyNode(
        "shared",
        subspace_tags=["x1"],
    ),
    HierarchyNode(
        "branch_A",
        subspace_tags=["x2"],
        activity_condition_tags={"y1": 1},
    ),
    HierarchyNode(
        "branch_B",
        subspace_tags=["x3"],
        activity_condition_tags={"y1": 0},
    ),
]
space = HierarchicalSearchSpace(subspaces, hierarchy)

print("dimension:", int(space.dimension))
print("indicator_tags:", space.indicator_tags)
print("non_indicator_tags:", space.non_indicator_tags)
print("indicator_value_sets:", space.indicator_value_sets)

dimension: 4
indicator_tags: ('y1',)
non_indicator_tags: ('x1', 'x2', 'x3')
indicator_value_sets: {'y1': (0, 1)}

Sampling and the flat-vector layout#

Points are represented as flat vectors [x1, y1, x2, x3] concatenated in tag order, the same convention as TaggedProductSearchSpace. For a sample with \(y_1 = 1\) the x3 column is inactive: its value is still present in the vector but has no semantic meaning. Downstream kernels (Arc, Wedge, Conditional) handle inactive columns according to their respective axioms.

[3]:

samples = space.sample(8)
print("samples shape:", samples.shape)
print(samples.numpy())

samples shape: (8, 4)
[[ 0.74731156  1.          4.37343108 -0.13034032]
 [ 0.45715611  0.          0.37282851 -0.62220801]
 [ 0.71670308  1.          0.63046947  0.03783136]
 [ 0.1830982   1.          2.42705078  0.06404234]
 [ 0.37016942  1.          3.12578231  0.84299597]
 [ 0.98365795  0.          0.12316771  0.86711156]
 [ 0.00584848  1.          1.70038628  0.97217024]
 [ 0.15572476  1.          2.27359396 -0.68564291]]

get_subspace_component slices the columns belonging to a given tag. It works on any batch shape because indexing is on the trailing axis.

[4]:

print("y1 column:", space.get_subspace_component("y1", samples).numpy().ravel())
print("x2 column:", space.get_subspace_component("x2", samples).numpy().ravel())

y1 column: [1. 0. 1. 1. 1. 0. 1. 1.]
x2 column: [4.37343108 0.37282851 0.63046947 2.42705078 3.12578231 0.12316771
 1.70038628 2.27359396]

Hierarchy queries#

Several methods expose the hierarchy programmatically. enumerate_tasks returns every indicator configuration; active_subspace_tags returns the non-indicator tags active for one configuration; is_active answers the per-tag question; and node_for_subspace returns the nodes that contain a given tag.

[5]:

print("enumerate_tasks:", space.enumerate_tasks())
print("active for y1=1:", space.active_subspace_tags({"y1": 1}))
print("active for y1=0:", space.active_subspace_tags({"y1": 0}))
print("x2 is_active when y1=1:", space.is_active("x2", {"y1": 1}))
print("x2 is_active when y1=0:", space.is_active("x2", {"y1": 0}))
print("nodes containing 'x1':", [n.name for n in space.node_for_subspace("x1")])
print("nodes containing 'x2':", [n.name for n in space.node_for_subspace("x2")])

enumerate_tasks: [{'y1': 0}, {'y1': 1}]
active for y1=1: ['x1', 'x2']
active for y1=0: ['x1', 'x3']
x2 is_active when y1=1: True
x2 is_active when y1=0: False
nodes containing 'x1': ['shared']
nodes containing 'x2': ['branch_A']

Categorical-indicator example#

When a disjunction has more than two branches, a single \(K\)-ary categorical indicator is the natural representation. Suppose \(y_1\) now selects among three unit types \(\{0\colon\text{shared-only},\;1\colon\text{branch A},\;2\colon\text{branch B}\}\).

[6]:

subspaces_c = {
    "x1": Box([0.0], [1.0]),  # unconditional
    "y1": CategoricalSearchSpace(3),  # 3-ary indicator
    "x2": Box([0.0], [5.0]),  # active when y1 = 1
    "x3": Box([-1.0], [1.0]),  # active when y1 = 2
}
hierarchy_c = [
    HierarchyNode(
        "shared",
        subspace_tags=["x1"],
    ),
    HierarchyNode(
        "branch_A",
        subspace_tags=["x2"],
        activity_condition_tags={"y1": 1},
    ),
    HierarchyNode(
        "branch_B",
        subspace_tags=["x3"],
        activity_condition_tags={"y1": 2},
    ),
]
space_c = HierarchicalSearchSpace(subspaces_c, hierarchy_c)

print("indicator_value_sets:", space_c.indicator_value_sets)
print("enumerate_tasks:", space_c.enumerate_tasks())
print("active for y1=0:", space_c.active_subspace_tags({"y1": 0}))
print("active for y1=1:", space_c.active_subspace_tags({"y1": 1}))
print("active for y1=2:", space_c.active_subspace_tags({"y1": 2}))

indicator_value_sets: {'y1': (0, 1, 2)}
enumerate_tasks: [{'y1': 0}, {'y1': 1}, {'y1': 2}]
active for y1=0: ['x1']
active for y1=1: ['x1', 'x2']
active for y1=2: ['x1', 'x3']

Mixing indicator kinds#

Boolean and categorical indicators can be combined in a single HierarchicalSearchSpace. enumerate_tasks then returns the Cartesian product of every indicator’s permitted set: \(|\{0,1\}| \times |\{0,1,2\}| = 6\) tasks below.

[7]:

subspaces_mix = {
    "x1": Box([0.0], [1.0]),
    "y1": BooleanSearchSpace(),
    "y2": CategoricalSearchSpace(3),
    "x2": Box([0.0], [5.0]),
}
hierarchy_mix = [
    HierarchyNode(
        "shared",
        subspace_tags=["x1"],
    ),
    HierarchyNode(
        "branch",
        subspace_tags=["x2"],
        activity_condition_tags={"y1": 1, "y2": 2},
    ),
]
space_mix = HierarchicalSearchSpace(subspaces_mix, hierarchy_mix)
tasks = space_mix.enumerate_tasks()
print(f"number of tasks: {len(tasks)}")
for t in tasks:
    print(" ", t)

number of tasks: 6
  {'y1': 0, 'y2': 0}
  {'y1': 0, 'y2': 1}
  {'y1': 0, 'y2': 2}
  {'y1': 1, 'y2': 0}
  {'y1': 1, 'y2': 1}
  {'y1': 1, 'y2': 2}

Validation rules#

Construction fails fast with a ValueError when the inputs are inconsistent. A few representative rejections:

an indicator tag that does not point to a BooleanSearchSpace or a one-dimensional CategoricalSearchSpace;
an activity_condition required value that is outside the indicator’s permitted set (e.g. 5 for a 3-ary categorical, or any non-Boolean value for a BooleanSearchSpace);
an indicator that gates nothing (does not appear as a key in any node).

[8]:

def _expect_value_error(fn):
    try:
        fn()
    except ValueError as e:
        print("rejected:", str(e).split("\n")[0])
    else:
        raise AssertionError("expected ValueError but none was raised")


# Out-of-permitted-set required value: y1 is a 3-ary categorical, so 5 is invalid.
_subspaces_oor = {
    "x1": Box([0.0], [1.0]),
    "y1": CategoricalSearchSpace(3),
}
_expect_value_error(
    lambda: HierarchicalSearchSpace(
        _subspaces_oor,
        [
            HierarchyNode(
                "n",
                subspace_tags=["x1"],
                activity_condition_tags={"y1": 5},
            )
        ],
    )
)

# Indicator subspace must be dimension-1: a 2-D CategoricalSearchSpace is rejected.
_subspaces_2d = {
    "x1": Box([0.0], [1.0]),
    "y1": CategoricalSearchSpace([3, 2]),
}
_expect_value_error(
    lambda: HierarchicalSearchSpace(
        _subspaces_2d,
        [
            HierarchyNode(
                "n",
                subspace_tags=["x1"],
                activity_condition_tags={"y1": 1},
            )
        ],
    )
)

rejected: HierarchyNode 'n' requires value 5 for indicator 'y1', which is not in its permitted set [0, 1, 2].
rejected: Indicator 'y1' must be 1-dimensional, got 2 columns.