Coverage for python/lsst/pipe/base/graph/_implDetails.py : 31%
Hot-keys on this page
r m x p toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1# This file is part of pipe_base.
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (http://www.lsst.org).
6# See the COPYRIGHT file at the top-level directory of this distribution
7# for details of code ownership.
8#
9# This program is free software: you can redistribute it and/or modify
10# it under the terms of the GNU General Public License as published by
11# the Free Software Foundation, either version 3 of the License, or
12# (at your option) any later version.
13#
14# This program is distributed in the hope that it will be useful,
15# but WITHOUT ANY WARRANTY; without even the implied warranty of
16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
17# GNU General Public License for more details.
18#
19# You should have received a copy of the GNU General Public License
20# along with this program. If not, see <http://www.gnu.org/licenses/>.
21from __future__ import annotations
22from collections import defaultdict
24__all__ = ("_DatasetTracker", "DatasetTypeName")
26import networkx as nx
27from typing import (DefaultDict, Generic, Optional, Set, TypeVar, NewType, Dict)
29from lsst.daf.butler import DatasetRef
31from .quantumNode import QuantumNode
32from ..pipeline import TaskDef
34# NewTypes
35DatasetTypeName = NewType("DatasetTypeName", str)
37# Generic type parameters
38_T = TypeVar("_T", DatasetTypeName, DatasetRef)
39_U = TypeVar("_U", TaskDef, QuantumNode)
42class _DatasetTracker(Generic[_T, _U]):
43 r"""This is a generic container for tracking keys which are produced or
44 consumed by some value. In the context of a QuantumGraph, keys may be
45 `~lsst.daf.butler.DatasetRef`\ s and the values would be Quanta that either
46 produce or consume those `~lsst.daf.butler.DatasetRef`\ s.
48 Prameters
49 ---------
50 createInverse : bool
51 When adding a key associated with a producer or consumer, also create
52 and inverse mapping that allows looking up all the keys associated with
53 some value. Defaults to False.
54 """
55 def __init__(self, createInverse: bool = False):
56 self._producers: Dict[_T, _U] = {}
57 self._consumers: DefaultDict[_T, Set[_U]] = defaultdict(set)
58 self._createInverse = createInverse
59 if self._createInverse:
60 self._itemsDict: DefaultDict[_U, Set[_T]] = defaultdict(set)
62 def addProducer(self, key: _T, value: _U):
63 """Add a key which is produced by some value.
65 Parameters
66 ----------
67 key : TypeVar
68 The type to track
69 value : TypeVar
70 The type associated with the production of the key
72 Raises
73 ------
74 ValueError
75 Raised if key is already declared to be produced by another value
76 """
77 if (existing := self._producers.get(key)) is not None and existing != value:
78 raise ValueError(f"Only one node is allowed to produce {key}, "
79 f"the current producer is {existing}")
80 self._producers[key] = value
81 if self._createInverse:
82 self._itemsDict[value].add(key)
84 def addConsumer(self, key: _T, value: _U):
85 """Add a key which is consumed by some value.
87 Parameters
88 ----------
89 key : TypeVar
90 The type to track
91 value : TypeVar
92 The type associated with the consumption of the key
93 """
94 self._consumers[key].add(value)
95 if self._createInverse:
96 self._itemsDict[value].add(key)
98 def getConsumers(self, key: _T) -> set[_U]:
99 """Return all values associated with the consumption of the supplied
100 key.
102 Parameters
103 ----------
104 key : TypeVar
105 The type which has been tracked in the _DatasetTracker
106 """
107 return self._consumers.get(key, set())
109 def getProducer(self, key: _T) -> Optional[_U]:
110 """Return the value associated with the consumption of the supplied
111 key.
113 Parameters
114 ----------
115 key : TypeVar
116 The type which has been tracked in the _DatasetTracker
117 """
118 return self._producers.get(key)
120 def getAll(self, key: _T) -> set[_U]:
121 """Return all consumers and the producer associated with the the
122 supplied key.
124 Parameters
125 ----------
126 key : TypeVar
127 The type which has been tracked in the _DatasetTracker
128 """
130 return self.getConsumers(key).union(x for x in (self.getProducer(key),) if x is not None)
132 @property
133 def inverse(self) -> Optional[DefaultDict[_U, Set[_T]]]:
134 """Return the inverse mapping if class was instantiated to create an
135 inverse, else return None.
136 """
137 return self._itemsDict if self._createInverse else None
139 def makeNetworkXGraph(self) -> nx.DiGraph:
140 """Create a NetworkX graph out of all the contained keys, using the
141 relations of producer and consumers to create the edges.
143 Returns:
144 graph : networkx.DiGraph
145 The graph created out of the supplied keys and their relations
146 """
147 graph = nx.DiGraph()
148 for entry in self._producers.keys() | self._consumers.keys():
149 producer = self.getProducer(entry)
150 consumers = self.getConsumers(entry)
151 # This block is for tasks that consume existing inputs
152 if producer is None and consumers:
153 for consumer in consumers:
154 graph.add_node(consumer)
155 # This block is for tasks that produce output that is not consumed
156 # in this graph
157 elif producer is not None and not consumers:
158 graph.add_node(producer)
159 # all other connections
160 else:
161 for consumer in consumers:
162 graph.add_edge(producer, consumer)
163 return graph
165 def keys(self) -> Set[_T]:
166 """Return all tracked keys.
167 """
168 return self._producers.keys() | self._consumers.keys()
170 def __contains__(self, key: _T) -> bool:
171 """Check if a key is in the _DatasetTracker
173 Parameters
174 ----------
175 key : TypeVar
176 The key to check
178 Returns
179 -------
180 contains : bool
181 Boolean of the presence of the supplied key
182 """
183 return key in self._producers or key in self._consumers