-
Notifications
You must be signed in to change notification settings - Fork 621
/
qutip-structure.py
executable file
·213 lines (194 loc) · 7.23 KB
/
qutip-structure.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
#!/usr/bin/env python
import inspect
import pathlib
import warnings
import sys
import qutip
# This script currently relies on all packages being imported by the
# import qutip
# command. If in the future some packages are not imported, then you'll need
# to add more import lines below it to make sure they're all in. We do this
# rather than file-based discovery so we have more access to information
# included by the import system, such as which names are meant to be public.
# It also means that we can import Cythonised modules to investigate their
# internals as well.
root_directory = pathlib.Path(qutip.__file__).parent
# This list needs to populated manually at the moment. Each element of the
# list is a two-tuple (colour, modules), where the `colour` is the text colour
# in the output, and `modules` is a set of module names that will be that
# colour. You can also put package names into the set of modules---any
# submodules of that package will inherit the same colour. You don't need to
# include the "qutip." prefix to the modules. It's a list not a dictionary
# because the order is important to the output.
module_groups = [
# Solvers
("#0b5fa5", {
"mesolve", "mcsolve", "sesolve", "stochastic", "bloch_redfield",
"nonmarkov", "floquet", "essolve", "correlation", "steadystate",
"rhs_generate", "propagator", "eseries", "hsolve", "rcsolve",
"scattering", "piqs", "pdpsolve",
}),
# Options and settings
("#043c6b", {"settings", "configrc", "solver"}),
# Visualisation
("#3f8fd2", {
"bloch", "bloch3d", "sphereplot", "orbital", "visualization", "wigner",
"distributions", "tomography", "topology",
}),
# Operators
("#00ae68", {
"operators", "superoperator", "superop_reps", "subsystem_apply",
}),
# States
("#007143", {
"states", "continuous_variables", "qstate", "random_objects",
}),
# QIP
("#36d695", {"measurement"}),
# Metrics and distance measures
("#ff4500", {"entropy", "metrics", "countstat", "semidefinite"}),
# Core
("#692102", {
"qobj", "qobjevo", "expect", "tensor", "partial_transpose", "ptrace",
"cy", "fastsparse", "interpolate",
}),
# Utilities
("#bf5730", {
"fileio", "utilities", "ipynbtools", "sparse", "graph", "simdiag",
"permute", "demos", "about", "parallel", "version", "testing",
"hardware_info", "ui", "cite",
}),
]
# Set of modules that we don't want to include in the output. Any modules that
# are detected inside `qutip` but are not either in this set or the
# `module_groups` list will generate a warning when the script is run.
modules_ignored = {
"dimensions",
"logging_utils",
"matplotlib_utilities",
"legacy",
"qobjevo_codegen",
"_mkl",
"cy.pyxbuilder",
"cy.openmp",
"cy.graph_utils",
"cy.inter",
"cy.cqobjevo",
"cy.cqobjevo_factor",
"cy.codegen",
"cy.br_codegen",
"cy.ptrace",
}
def _our_tree(module, tree):
"""
Find the subtree corresponding to this module, creating any necessary
subtrees along the way.
"""
our_tree = tree
cur_name = ""
for part in module.__name__.split(".")[1:]:
cur_name = (cur_name + "." + part) if cur_name else part
if cur_name in modules_ignored:
return tree
try:
our_tree = our_tree[part]
except KeyError:
our_tree[part] = {}
our_tree = our_tree[part]
return our_tree
def _ignore(module, root):
if not module.__name__.startswith(root):
return True
name = module.__name__[len(root):]
if name in modules_ignored:
return True
while (idx := name.rfind(".")) > 0:
name = name[:idx]
if name in modules_ignored:
return True
return False
def python_object_tree(module, tree=None, seen=None, root=None, nobjects=0):
"""
Recursively access every accessible element of the given module, building
up a complete tree structure where the keys are the parts of the module
name, and the eventual leaves are public functions and classes defined in
that particular module (so ignoring any names that leak in from other
imports). For example,
>>> import qutip
>>> python_object_tree(qutip)
{
"mesolve" : {
"mesolve": <function qutip.mesolve.mesolve(...)>,
},
...
}
"""
tree = tree if tree is not None else {}
seen = seen if seen is not None else set()
root = root if root is not None else (module.__name__ + ".")
if module in seen:
return tree, nobjects
seen.add(module)
our_tree = _our_tree(module, tree)
for _, obj in inspect.getmembers(module):
if inspect.isclass(obj) or inspect.isroutine(obj):
object_module = inspect.getmodule(obj)
if object_module is module:
if not obj.__name__.startswith("_"):
our_tree[obj.__name__] = obj
nobjects += 1
continue
# Fall through, so we recursively comb through modules.
obj = object_module
if inspect.ismodule(obj) and not _ignore(obj, root):
if obj.__name__.startswith(root):
_, nobjects =\
python_object_tree(obj, tree, seen, root, nobjects)
# Also do our parent package, if we have one. In theory it's possible to
# get into a situation with packages and overzealous use of "del" in init
# scripts where a submodule may be accessible but its parent isn't.
parent = ".".join(module.__name__.split(".")[:-1])
if parent.startswith(root):
_, nobjects =\
python_object_tree(sys.modules[parent], tree, seen, root, nobjects)
return tree, nobjects
def _lookup_color(basename, index, color):
for i, (color_, modules) in enumerate(module_groups):
if basename in modules:
return i, color_
return index, color
def convert_to_d3_struct(in_tree, name, index=-1, color=None, basename=None):
out_struct = {}
children = []
color_default = "black"
index, color = _lookup_color(basename, index, color)
for key, value in in_tree.items():
nextname = (basename + "." + key) if basename else key
if isinstance(value, dict):
out = convert_to_d3_struct(value, key, index, color, nextname)
else:
out = {
"name": key,
"color": color or color_default,
"index": index,
}
children.append(out)
if name == "QuTiP" and basename is None:
# Don't warn for the base case.
color = color_default
if color is None:
modname = "qutip" + (("." + basename) if basename else "")
warnings.warn("handling unspecified module: " + modname)
out_struct["name"] = name
out_struct["color"] = color or color_default
out_struct["index"] = index
if children:
out_struct["children"] = sorted(children, key=lambda x: x["index"])
return out_struct
if __name__ == "__main__":
import json
tree, count = python_object_tree(qutip)
struct = convert_to_d3_struct(tree, "QuTiP")
with open("d3_data/qutip.json", "w") as f:
json.dump(struct, f)
print(count)