def get_graph_data(graph, display_options):
color_scheme = display_options.get("color_scheme", "Jmol")
nodes = []
edges = []
struct_or_mol = StructureMoleculeComponent._get_struct_or_mol(graph)
legend = Legend(struct_or_mol, color_scheme=color_scheme)
for idx, node in enumerate(graph.graph.nodes()):
# TODO: fix for disordered
node_color = legend.get_color(
struct_or_mol[node].species.elements[0],
site=struct_or_mol[node])
nodes.append({
"id":
node,
"title":
f"{struct_or_mol[node].species_string} site "
f"({graph.get_coordination_of_site(idx)} neighbors)",
"color":
node_color,
})
for u, v, d in graph.graph.edges(data=True):
edge = {"from": u, "to": v, "arrows": ""}
to_jimage = d.get("to_jimage", (0, 0, 0))
# TODO: check these edge weights
if isinstance(struct_or_mol, Structure):
dist = struct_or_mol.get_distance(u, v, jimage=to_jimage)
else:
dist = struct_or_mol.get_distance(u, v)
edge["length"] = 50 * dist
if to_jimage != (0, 0, 0):
edge["arrows"] = "to"
label = f"{dist:.2f} Å to site at image vector {to_jimage}"
else:
label = f"{dist:.2f} Å between sites"
if label:
edge["title"] = label
# if 'weight' in d:
# label += f" {d['weight']}"
edges.append(edge)
return {"nodes": nodes, "edges": edges}
def test_get_color(self):
# test default
legend = Legend(self.struct, color_scheme="VESTA")
color = legend.get_color(self.sp0)
assert color == "#ffcccc"
# element-based color schemes shouldn't change if you supply a site
color = legend.get_color(self.sp0, site=self.site0)
assert color == "#ffcccc"
color = legend.get_color(self.sp2)
assert color == "#a67573"
assert legend.get_legend()["colors"] == {
"#a67573": "In",
"#fe0300": "O",
"#ffcccc": "H",
}
# test alternate
legend = Legend(self.struct, color_scheme="Jmol")
color = legend.get_color(self.sp0)
assert color == "#ffffff"
assert legend.get_legend()["colors"] == {
"#a67573": "In",
"#ff0d0d": "O",
"#ffffff": "H",
}
# test coloring by site properties
legend = Legend(self.struct, color_scheme="example_site_prop")
color = legend.get_color(self.sp0, site=self.site0)
assert color == "#b30326"
color = legend.get_color(self.sp1, site=self.site1)
assert color == "#dddcdb"
color = legend.get_color(self.sp2, site=self.site2)
assert color == "#7b9ef8"
assert legend.get_legend()["colors"] == {
"#7b9ef8": "-3.00",
"#b30326": "5.00",
"#dddcdb": "0.00",
}
# test accessible
legend = Legend(self.struct, color_scheme="accessible")
color = legend.get_color(self.sp0, site=self.site0)
assert color == "#ffffff"
color = legend.get_color(self.sp1, site=self.site1)
assert color == "#d55e00"
color = legend.get_color(self.sp2, site=self.site2)
assert color == "#cc79a7"
assert legend.get_legend()["colors"] == {
"#cc79a7": "In",
"#d55e00": "O",
"#ffffff": "H",
}
# test disordered
legend = Legend(self.struct_disordered)
color = legend.get_color(self.site_d_sp0, site=self.site_d)
assert color == "#a67573"
color = legend.get_color(self.site_d_sp1, site=self.site_d)
assert color == "#bfa6a6"
assert legend.get_legend()["colors"] == {
"#a67573": "In",
"#bfa6a6": "Al",
"#ff0d0d": "O",
"#ffffff": "H",
}
# test categorical
legend = Legend(self.struct,
color_scheme="example_categorical_site_prop")
assert legend.get_legend()["colors"] == {
"#377eb8": "8b",
"#e41a1c": "4a"
}
# test pre-defined
legend = Legend(self.struct_manual)
assert legend.get_legend()["colors"] == {
"#0000ff": "O2-",
"#00ff00": "In",
"#ff0000": "H",
}
