def __getitem__(self, key):
key = Key(key)
try:
return self.output_dict[key]
except KeyError:
assert self.parent is not None
return self.parent.collections[self.cat][key]
def make_key_subset_of(key_set):
key_set = Key(key_set)
def key_subset_of(key):
return set(key_set.key).issubset(set(Key(key).key))
return key_subset_of
def source_selector(parent, modality="all", location="all"):
locations_set = set(get_ancestral_metadata(parent, "locations"))
assert location in locations_set or location == "all", f"Location {location} not in {locations_set}"
modality_set = set(get_ancestral_metadata(parent, "modalities"))
assert modality in modality_set or modality == "all", f"Modality {modality} not in {modality_set}"
loc, mod = location, modality
root = parent / f"{loc=}-{mod=}"
# Prepare a set of viable outputs
valid_locations = set()
for pair in parent.meta["sources"]:
loc, mod = pair["loc"], pair["mod"]
good_location = location == "all" or pair["loc"] == location
good_modality = modality == "all" or pair["mod"] == modality
if good_location and good_modality:
valid_locations.update({Key(f"{loc=}-{mod=}")})
# Aggregate all relevant sources
selected = 0
for key, node in parent.outputs.items():
if key in valid_locations:
selected += 1
root.acquire_node(key=key, node=node)
if not selected:
logger.exception(
f"No wearable keys found in {sorted(parent.outputs.keys())}")
raise KeyError
return root
def acquire_node(self,
node: NodeWrapper,
key: Optional[Union[Key, str]] = None) -> None:
if key is None:
raise NotImplementedError
key = Key(key)
if key in self.nodes:
raise KeyError(
f"Cannot acquire {key} since a node of this name's already in {self.nodes.keys()} of {self}"
)
self.nodes[key] = node
def __getitem__(self, key: Union[Key, str]) -> NodeWrapper:
assert self.validate_key(key)
key = Key(key)
if key in self.graph.nodes:
return self.graph.nodes[key]
if self.graph.parent is not None:
return self.parent_group[key]
logger.exception(
f"Unable to find key '{key}' in graph - reached root.")
raise KeyError
def get_or_create(
self,
key: Union[Key, str],
func: Callable,
args: Optional[Union[Any, Tuple[Any]]] = None,
kwargs: Optional[Dict[str, Any]] = None,
backend: Optional[str] = None,
) -> NodeWrapper:
key = Key(key)
if key in self.nodes:
return self.nodes[key]
return self.instantiate_node(key=key,
func=func,
backend=backend,
args=args,
kwargs=kwargs)
def make_output(self, key, func, backend=None, kwargs=None):
assert callable(func)
key = Key(key)
if kwargs is None:
kwargs = dict()
assert key not in self.output_dict
node = self.graph.node(func=func,
name=self.graph.build_path(key),
backend=backend,
kwargs=dict(
key=key,
**kwargs,
))
return node
def instantiate_node(
self,
key: Union[Key, str],
func: Callable,
args: Optional[Union[Any, List[Any], Tuple[Any]]] = None,
kwargs: Optional[Dict[str, Any]] = None,
backend: Optional[str] = None,
force_add: bool = False,
) -> NodeWrapper:
key = Key(key)
if not force_add:
assert key not in self.nodes
name = absolute_node_name(identifier=self.identifier, key=key)
backend = validate_backend(backend, key)
return self.make_node(name=name,
key=key,
func=func,
backend=backend,
args=args,
kwargs=kwargs)
def validate_key_set_membership(key, key_set):
return Key(key) in set(map(Key, key_set))
def append_output(self, key, node):
key = Key(key)
assert key not in self.output_dict
self.output_dict[key] = node
def key_subset_of(key):
return set(key_set.key).issubset(set(Key(key).key))
