def get_scene_and_legend(self) -> Tuple[Scene, Dict[str, str]]:
legend = Legend(self.graph.structure,
color_scheme="VESTA",
cmap_range=None)
legend.uniform_radius = 0.3
scene = get_structure_graph_scene(
self.graph,
draw_image_atoms=True,
bonded_sites_outside_unit_cell=False,
hide_incomplete_edges=True,
explicitly_calculate_polyhedra_hull=False,
group_by_symmetry=True,
draw_polyhedra=False,
legend=legend)
scene.name = "DefectStructureComponentScene"
lattice: Lattice = self.graph.structure.lattice
origin = -self.graph.structure.lattice.get_cartesian_coords(
[0.5, 0.5, 0.5])
scene_json = scene.to_json()
for idx, i in enumerate(self.interstitials, 1):
site = Site(species=DummySpecie(),
coords=lattice.get_cartesian_coords(i.frac_coords))
interstitial_scene = site.get_scene(origin=origin)
interstitial_scene.name = f"i{idx}"
interstitial_scene.contents[0].contents[0].tooltip = f"i{idx}"
scene_json["contents"].append(interstitial_scene.to_json())
return scene_json, legend.get_legend()
def test_msonable(self):
legend = Legend(self.struct)
legend_dict = legend.as_dict()
legend_from_dict = Legend.from_dict(legend_dict)
assert legend.get_legend() == legend_from_dict.get_legend()
def get_scene_and_legend(
graph: Optional[Union[StructureGraph, MoleculeGraph]],
name,
color_scheme=DEFAULTS["color_scheme"],
color_scale=None,
radius_strategy=DEFAULTS["radius_strategy"],
draw_image_atoms=DEFAULTS["draw_image_atoms"],
bonded_sites_outside_unit_cell=DEFAULTS[
"bonded_sites_outside_unit_cell"],
hide_incomplete_bonds=DEFAULTS["hide_incomplete_bonds"],
explicitly_calculate_polyhedra_hull=False,
scene_additions=None,
show_compass=DEFAULTS["show_compass"],
) -> Tuple[Scene, Dict[str, str]]:
# default scene name will be name of component, "_ct_..."
# strip leading _ since this will cause problems in JavaScript land
scene = Scene(name=name[1:])
if graph is None:
return scene, {}
struct_or_mol = StructureMoleculeComponent._get_struct_or_mol(graph)
# TODO: add radius_scale
legend = Legend(
struct_or_mol,
color_scheme=color_scheme,
radius_scheme=radius_strategy,
cmap_range=color_scale,
)
if isinstance(graph, StructureGraph):
scene = graph.get_scene(
draw_image_atoms=draw_image_atoms,
bonded_sites_outside_unit_cell=bonded_sites_outside_unit_cell,
hide_incomplete_edges=hide_incomplete_bonds,
explicitly_calculate_polyhedra_hull=
explicitly_calculate_polyhedra_hull,
legend=legend,
)
elif isinstance(graph, MoleculeGraph):
scene = graph.get_scene(legend=legend)
scene.name = name
if hasattr(struct_or_mol, "lattice"):
axes = struct_or_mol.lattice._axes_from_lattice()
# TODO: fix pop-in ?
axes.visible = show_compass
scene.contents.append(axes)
if scene_additions:
# TODO: need a Scene.from_json() to make this work
raise NotImplementedError
scene["contents"].append(scene_additions)
return scene.to_json(), legend.get_legend()
def get_scene_and_legend(
graph: Optional[Union[StructureGraph, MoleculeGraph]],
color_scheme=DEFAULTS["color_scheme"],
color_scale=None,
radius_strategy=DEFAULTS["radius_strategy"],
draw_image_atoms=DEFAULTS["draw_image_atoms"],
bonded_sites_outside_unit_cell=DEFAULTS[
"bonded_sites_outside_unit_cell"],
hide_incomplete_bonds=DEFAULTS["hide_incomplete_bonds"],
explicitly_calculate_polyhedra_hull=False,
scene_additions=None,
show_compass=DEFAULTS["show_compass"],
group_by_site_property=None,
) -> Tuple[Scene, Dict[str, str]]:
scene = Scene(name="StructureMoleculeComponentScene")
if graph is None:
return scene, {}
struct_or_mol = StructureMoleculeComponent._get_struct_or_mol(graph)
# TODO: add radius_scale
legend = Legend(
struct_or_mol,
color_scheme=color_scheme,
radius_scheme=radius_strategy,
cmap_range=color_scale,
)
if isinstance(graph, StructureGraph):
scene = graph.get_scene(
draw_image_atoms=draw_image_atoms,
bonded_sites_outside_unit_cell=bonded_sites_outside_unit_cell,
hide_incomplete_edges=hide_incomplete_bonds,
explicitly_calculate_polyhedra_hull=
explicitly_calculate_polyhedra_hull,
group_by_site_property=group_by_site_property,
legend=legend,
)
elif isinstance(graph, MoleculeGraph):
scene = graph.get_scene(legend=legend)
scene.name = "StructureMoleculeComponentScene"
if hasattr(struct_or_mol, "lattice"):
axes = struct_or_mol.lattice._axes_from_lattice()
axes.visible = show_compass
scene.contents.append(axes)
scene_json = scene.to_json()
if scene_additions:
# TODO: this might be cleaner if we had a Scene.from_json() method
scene_json["contents"].append(scene_additions)
return scene_json, legend.get_legend()
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 get_scene_and_legend(
graph: Optional[Union[StructureGraph, MoleculeGraph]],
color_scale=None,
radius_strategy=DEFAULTS["radius_strategy"],
draw_image_atoms=DEFAULTS["draw_image_atoms"],
bonded_sites_outside_unit_cell=DEFAULTS[
"bonded_sites_outside_unit_cell"],
hide_incomplete_bonds=DEFAULTS["hide_incomplete_bonds"],
explicitly_calculate_polyhedra_hull=False,
scene_additions=None,
show_compass=DEFAULTS["show_compass"],
) -> Tuple[Scene, Dict[str, str]]:
scene = Scene(name="AceStructureMoleculeComponentScene")
if graph is None:
return scene, {}
structure = StructureComponent._get_structure(graph)
# TODO: add radius_scale
legend = Legend(
structure,
color_scheme="VESTA",
radius_scheme=radius_strategy,
cmap_range=color_scale,
)
scene = graph.get_scene(
draw_image_atoms=draw_image_atoms,
bonded_sites_outside_unit_cell=bonded_sites_outside_unit_cell,
hide_incomplete_edges=hide_incomplete_bonds,
explicitly_calculate_polyhedra_hull=
explicitly_calculate_polyhedra_hull,
legend=legend,
)
scene.name = "StructureComponentScene"
if hasattr(structure, "lattice"):
axes = structure.lattice._axes_from_lattice()
axes.visible = show_compass
scene.contents.append(axes)
scene = scene.to_json()
if scene_additions:
# TODO: need a Scene.from_json() to make this work
# raise NotImplementedError
scene["contents"].append(scene_additions)
return scene, legend.get_legend()
def get_scene_from_molecule(self,
origin=None,
legend: Optional[Legend] = None):
"""
Create CTK objects for the lattice and sties
Args:
self: Structure object
origin: fractional coordinate of the origin
legend: Legend for the sites
Returns:
CTK scene object to be rendered
"""
origin = origin if origin else (0, 0, 0)
legend = legend or Legend(self)
primitives = defaultdict(list)
for idx, site in enumerate(self):
site_scene = site.get_scene(
origin=origin,
legend=legend,
)
for scene in site_scene.contents:
primitives[scene.name] += scene.contents
return Scene(
name=self.composition.reduced_formula,
contents=[Scene(name=k, contents=v) for k, v in primitives.items()],
origin=origin,
)
def get_molecule_graph_scene(
self,
origin=None,
explicitly_calculate_polyhedra_hull=False,
legend=None,
draw_polyhedra=False,
) -> Scene:
legend = legend or Legend(self.molecule)
primitives = defaultdict(list)
for idx, site in enumerate(self.molecule):
connected_sites = self.get_connected_sites(idx)
site_scene = site.get_scene(
connected_sites=connected_sites,
origin=origin,
explicitly_calculate_polyhedra_hull=
explicitly_calculate_polyhedra_hull,
legend=legend,
draw_polyhedra=draw_polyhedra,
)
for scene in site_scene.contents:
primitives[scene.name] += scene.contents
return Scene(
name=self.molecule.composition.reduced_formula,
contents=[Scene(name=k, contents=v) for k, v in primitives.items()],
origin=origin if origin else (0, 0, 0),
)
def update_scene_and_legend_and_colors(graph, display_options, scene_additions):
if not graph or not display_options:
raise PreventUpdate
display_options = self.from_data(display_options)
graph = self.from_data(graph)
scene, legend = self.get_scene_and_legend(
graph,
name=self.id(),
**display_options,
scene_additions=scene_additions,
)
color_options = [
{"label": "Jmol", "value": "Jmol"},
{"label": "VESTA", "value": "VESTA"},
{"label": "Accessible", "value": "accessible"},
]
struct_or_mol = self._get_struct_or_mol(graph)
site_props = Legend(struct_or_mol).analyze_site_props(struct_or_mol)
for site_prop_type in ("scalar", "categorical"):
if site_prop_type in site_props:
for prop in site_props[site_prop_type]:
color_options += [
{"label": f"Site property: {prop}", "value": prop}
]
return scene, legend, color_options
def get_structure_scene(
self,
draw_image_atoms=True,
legend: Optional[Legend] = None,
origin=None,
) -> Scene:
origin = origin or list(
-self.lattice.get_cartesian_coords([0.5, 0.5, 0.5]))
legend = legend or Legend(self)
primitives = defaultdict(list)
sites_to_draw = self._get_sites_to_draw(draw_image_atoms=draw_image_atoms)
for (idx, jimage) in sites_to_draw:
site_scene = self[idx].get_scene(origin=origin, legend=legend)
for scene in site_scene.contents:
primitives[scene.name] += scene.contents
primitives["unit_cell"].append(self.lattice.get_scene(origin=origin))
return Scene(
name=self.composition.reduced_formula,
contents=[Scene(name=k, contents=v) for k, v in primitives.items()],
origin=origin,
)
def test_get_radius(self):
legend = Legend(self.struct, radius_scheme="uniform")
assert legend.get_radius(sp=self.sp0) == 0.5
legend = Legend(self.struct, radius_scheme="covalent")
assert legend.get_radius(sp=self.sp1) == 0.66
legend = Legend(self.struct, radius_scheme="specified_or_average_ionic")
assert legend.get_radius(sp=self.sp2) == 0.94
def get_scene_and_legend(self,
scene_additions=None
) -> Tuple[Scene, Dict[str, str]]:
legend = Legend(self.graph.structure,
color_scheme="VESTA",
radius_scheme="uniform",
cmap_range=None)
legend.uniform_radius = 0.2
scene = get_structure_graph_scene(
self.graph,
draw_image_atoms=True,
bonded_sites_outside_unit_cell=False,
hide_incomplete_edges=True,
explicitly_calculate_polyhedra_hull=False,
group_by_symmetry=False,
legend=legend)
scene.name = "DefectStructureComponentScene"
# axes = graph.structure.lattice._axes_from_lattice()
# axes.visible = True
# scene.contents.append(axes)
scene = scene.to_json()
if scene_additions:
scene["contents"].append(scene_additions)
lattice = self.graph.structure.lattice
origin = -self.graph.structure.lattice.get_cartesian_coords(
[0.5, 0.5, 0.5])
for name, frac_coords in self.vacancy_sites:
site = Site(species=DummySpecie(name),
coords=lattice.get_cartesian_coords(frac_coords))
vac_scene = site.get_scene(origin=origin)
vac_scene.name = f"{name}_{frac_coords}"
vac_scene.contents[0].contents[0].tooltip = name
scene["contents"].append(vac_scene.to_json())
return scene, legend.get_legend()
def get_structure_scene(
self,
origin: List[float] = None,
legend: Optional[Legend] = None,
draw_image_atoms: bool = True,
) -> Scene:
"""
Create CTK objects for the lattice and sties
Args:
self: Structure object
origin: fractional coordinate of the origin
legend: Legend for the sites
draw_image_atoms: If true draw image atoms that are just outside the
periodic boundary
Returns:
CTK scene object to be rendered
"""
origin = origin or list(
-self.lattice.get_cartesian_coords([0.5, 0.5, 0.5]))
legend = legend or Legend(self)
primitives = defaultdict(list)
sites_to_draw = self._get_sites_to_draw(
draw_image_atoms=draw_image_atoms, )
for (idx, jimage) in sites_to_draw:
site = self[idx]
if jimage != (0, 0, 0):
site = PeriodicSite(
site.species,
np.add(site.frac_coords, jimage),
site.lattice,
properties=site.properties,
)
site_scene = site.get_scene(legend=legend, )
for scene in site_scene.contents:
primitives[scene.name] += scene.contents
primitives["unit_cell"].append(self.lattice.get_scene())
return Scene(
name="Structure",
origin=origin,
contents=[
Scene(name=k, contents=v, origin=origin)
for k, v in primitives.items()
],
)
def get_structure_scene(
self,
origin=None,
draw_image_atoms=True,
bonded_sites_outside_unit_cell=False,
legend: Optional[Legend] = None,
) -> Scene:
legend = legend or Legend(self.structure)
primitives = defaultdict(list)
sites_to_draw = self._get_sites_to_draw(
draw_image_atoms=draw_image_atoms,
bonded_sites_outside_unit_cell=bonded_sites_outside_unit_cell,
)
for (idx, jimage) in sites_to_draw:
site = self.structure[idx]
if jimage != (0, 0, 0):
connected_sites = self.get_connected_sites(idx, jimage=jimage)
site = PeriodicSite(
site.species,
np.add(site.frac_coords, jimage),
site.lattice,
properties=site.properties,
)
else:
connected_sites = self.get_connected_sites(idx)
site_scene = site.get_scene(origin=origin, legend=legend)
for scene in site_scene.contents:
primitives[scene.name] += scene.contents
primitives["unit_cell"].append(
self.structure.lattice.get_scene(origin=origin))
return Scene(
name=self.structure.composition.reduced_formula,
contents=[Scene(name=k, contents=v) for k, v in primitives.items()],
origin=origin,
)
def get_site_scene(
self,
connected_sites: List[ConnectedSite] = None,
# connected_site_metadata: None,
# connected_sites_to_draw,
connected_sites_not_drawn: List[ConnectedSite] = None,
hide_incomplete_edges: bool = False,
incomplete_edge_length_scale: Optional[float] = 1.0,
connected_sites_colors: Optional[List[str]] = None,
connected_sites_not_drawn_colors: Optional[List[str]] = None,
origin: Optional[List[float]] = None,
draw_polyhedra: bool = True,
explicitly_calculate_polyhedra_hull: bool = False,
bond_radius: float = 0.1,
legend: Optional[Legend] = None,
) -> Scene:
"""
Args:
connected_sites:
connected_sites_not_drawn:
hide_incomplete_edges:
incomplete_edge_length_scale:
connected_sites_colors:
connected_sites_not_drawn_colors:
origin:
explicitly_calculate_polyhedra_hull:
legend:
Returns:
"""
atoms = []
bonds = []
polyhedron = []
legend = legend or Legend(self)
# for disordered structures
is_ordered = self.is_ordered
phiStart, phiEnd = None, None
occu_start = 0.0
position = self.coords.tolist()
for idx, (sp, occu) in enumerate(self.species.items()):
if isinstance(sp, DummySpecie):
cube = Cubes(positions=[position],
color=legend.get_color(sp, site=self),
width=0.4)
atoms.append(cube)
else:
color = legend.get_color(sp, site=self)
radius = legend.get_radius(sp, site=self)
# TODO: make optional/default to None
# in disordered structures, we fractionally color-code spheres,
# drawing a sphere segment from phi_end to phi_start
# (think a sphere pie chart)
if not is_ordered:
phi_frac_end = occu_start + occu
phi_frac_start = occu_start
occu_start = phi_frac_end
phiStart = phi_frac_start * np.pi * 2
phiEnd = phi_frac_end * np.pi * 2
name = str(sp)
if occu != 1.0:
name += " ({}% occupancy)".format(occu)
name += f" ({position[0]:.3f}, {position[1]:.3f}, {position[2]:.3f})"
sphere = Spheres(
positions=[position],
color=color,
radius=radius,
phiStart=phiStart,
phiEnd=phiEnd,
clickable=True,
tooltip=name,
)
atoms.append(sphere)
if not is_ordered and not np.isclose(phiEnd, np.pi * 2):
# if site occupancy doesn't sum to 100%, cap sphere
sphere = Spheres(
positions=[position],
color="#ffffff",
radius=self.properties["display_radius"][0],
phiStart=phiEnd,
phiEnd=np.pi * 2,
)
atoms.append(sphere)
if connected_sites:
# TODO: more graceful solution here
# if ambiguous (disordered), re-use last color used
site_color = color
# TODO: can cause a bug if all vertices almost co-planar
# necessary to include center site in case it's outside polyhedra
all_positions = [self.coords]
for idx, connected_site in enumerate(connected_sites):
connected_position = connected_site.site.coords
bond_midpoint = np.add(position, connected_position) / 2
if connected_sites_colors:
color = connected_sites_colors[idx]
else:
color = site_color
cylinder = Cylinders(
positionPairs=[[position, bond_midpoint.tolist()]],
color=color,
radius=bond_radius,
)
bonds.append(cylinder)
all_positions.append(connected_position.tolist())
if connected_sites_not_drawn and not hide_incomplete_edges:
for idx, connected_site in enumerate(connected_sites_not_drawn):
connected_position = connected_site.site.coords
bond_midpoint = (incomplete_edge_length_scale *
np.add(position, connected_position) / 2)
if connected_sites_not_drawn_colors:
color = connected_sites_not_drawn_colors[idx]
else:
color = site_color
cylinder = Cylinders(
positionPairs=[[position, bond_midpoint.tolist()]],
color=color,
radius=bond_radius,
)
bonds.append(cylinder)
all_positions.append(connected_position.tolist())
# ensure intersecting polyhedra are not shown, defaults to choose by electronegativity
not_most_electro_negative = map(
lambda x:
(x.site.specie < self.specie) or (x.site.specie == self.specie),
connected_sites,
)
all_positions = [list(p) for p in all_positions]
if (draw_polyhedra and len(connected_sites) > 3
and not connected_sites_not_drawn
and not any(not_most_electro_negative)):
if explicitly_calculate_polyhedra_hull:
try:
# all_positions = [[0, 0, 0], [0, 0, 10], [0, 10, 0], [10, 0, 0]]
# gives...
# .convex_hull = [[2, 3, 0], [1, 3, 0], [1, 2, 0], [1, 2, 3]]
# .vertex_neighbor_vertices = [1, 2, 3, 2, 3, 0, 1, 3, 0, 1, 2, 0]
vertices_indices = Delaunay(all_positions).convex_hull
except Exception as e:
vertices_indices = []
vertices = [
all_positions[idx]
for idx in chain.from_iterable(vertices_indices)
]
polyhedron = [Surface(positions=vertices, color=site_color)]
else:
polyhedron = [
Convex(positions=all_positions, color=site_color)
]
return Scene(
self.species_string,
[
Scene("atoms", contents=atoms),
Scene("bonds", contents=bonds),
Scene("polyhedra", contents=polyhedron),
],
origin=origin,
)
def get_structure_graph_scene(
self,
origin=None,
draw_image_atoms=True,
bonded_sites_outside_unit_cell=True,
hide_incomplete_edges=False,
incomplete_edge_length_scale=0.3,
color_edges_by_edge_weight=True,
edge_weight_color_scale="coolwarm",
explicitly_calculate_polyhedra_hull=False,
legend: Optional[Legend] = None,
group_by_symmetry: bool = True,
) -> Scene:
origin = origin or list(
-self.structure.lattice.get_cartesian_coords([0.5, 0.5, 0.5]))
legend = legend or Legend(self.structure)
primitives = defaultdict(list)
sites_to_draw = self._get_sites_to_draw(
draw_image_atoms=draw_image_atoms,
bonded_sites_outside_unit_cell=bonded_sites_outside_unit_cell,
)
color_edges = False
if color_edges_by_edge_weight:
weights = [e[2].get("weight") for e in self.graph.edges(data=True)]
weights = np.array([w for w in weights if w])
if any(weights):
cmap = get_cmap(edge_weight_color_scale)
# try to keep color scheme symmetric around 0
weight_max = max([abs(min(weights)), max(weights)])
weight_min = -weight_max
def get_weight_color(weight):
if not weight:
weight = 0
x = (weight - weight_min) / (weight_max - weight_min)
return "#{:02x}{:02x}{:02x}".format(
*[int(c * 255) for c in cmap(x)[0:3]])
color_edges = True
idx_to_wyckoff = {}
if group_by_symmetry:
sga = SpacegroupAnalyzer(self.structure)
struct_sym = sga.get_symmetrized_structure()
for equiv_idxs, wyckoff in zip(struct_sym.equivalent_indices,
struct_sym.wyckoff_symbols):
for idx in equiv_idxs:
idx_to_wyckoff[idx] = wyckoff
for (idx, jimage) in sites_to_draw:
site = self.structure[idx]
if jimage != (0, 0, 0):
connected_sites = self.get_connected_sites(idx, jimage=jimage)
site = PeriodicSite(
site.species,
np.add(site.frac_coords, jimage),
site.lattice,
properties=site.properties,
)
else:
connected_sites = self.get_connected_sites(idx)
connected_sites = [
cs for cs in connected_sites
if (cs.index, cs.jimage) in sites_to_draw
]
connected_sites_not_drawn = [
cs for cs in connected_sites
if (cs.index, cs.jimage) not in sites_to_draw
]
if color_edges:
connected_sites_colors = [
get_weight_color(cs.weight) for cs in connected_sites
]
connected_sites_not_drawn_colors = [
get_weight_color(cs.weight) for cs in connected_sites_not_drawn
]
else:
connected_sites_colors = None
connected_sites_not_drawn_colors = None
site_scene = site.get_scene(
connected_sites=connected_sites,
connected_sites_not_drawn=connected_sites_not_drawn,
hide_incomplete_edges=hide_incomplete_edges,
incomplete_edge_length_scale=incomplete_edge_length_scale,
connected_sites_colors=connected_sites_colors,
connected_sites_not_drawn_colors=connected_sites_not_drawn_colors,
explicitly_calculate_polyhedra_hull=
explicitly_calculate_polyhedra_hull,
legend=legend,
)
for scene in site_scene.contents:
if group_by_symmetry and scene.name == "atoms" and idx in idx_to_wyckoff:
# will rename to e.g. atoms_N_4e
scene.name = f"atoms_{site_scene.name}_{idx_to_wyckoff[idx]}"
# this is a proof-of-concept to demonstrate hover labels, could create label
# automatically from site properties instead
scene.contents[
0].tooltip = f"{site_scene.name} ({idx_to_wyckoff[idx]})"
primitives[scene.name] += scene.contents
primitives["unit_cell"].append(self.structure.lattice.get_scene())
return Scene(
name="StructureGraph",
origin=origin,
contents=[
Scene(name=k, contents=v, origin=origin)
for k, v in primitives.items()
],
)
def get_structure_graph_scene(
self,
origin=None,
draw_image_atoms=True,
bonded_sites_outside_unit_cell=True,
hide_incomplete_edges=False,
incomplete_edge_length_scale=0.3,
color_edges_by_edge_weight=True,
edge_weight_color_scale="coolwarm",
explicitly_calculate_polyhedra_hull=False,
legend: Optional[Legend] = None,
) -> Scene:
origin = origin or list(
-self.structure.lattice.get_cartesian_coords([0.5, 0.5, 0.5]))
legend = legend or Legend(self.structure)
primitives = defaultdict(list)
sites_to_draw = self._get_sites_to_draw(
draw_image_atoms=draw_image_atoms,
bonded_sites_outside_unit_cell=bonded_sites_outside_unit_cell,
)
color_edges = False
if color_edges_by_edge_weight:
weights = [e[2].get("weight") for e in self.graph.edges(data=True)]
weights = np.array([w for w in weights if w])
if any(weights):
cmap = get_cmap(edge_weight_color_scale)
# try to keep color scheme symmetric around 0
weight_max = max([abs(min(weights)), max(weights)])
weight_min = -weight_max
def get_weight_color(weight):
if not weight:
weight = 0
x = (weight - weight_min) / (weight_max - weight_min)
return "#{:02x}{:02x}{:02x}".format(
*[int(c * 255) for c in cmap(x)[0:3]])
color_edges = True
for (idx, jimage) in sites_to_draw:
site = self.structure[idx]
if jimage != (0, 0, 0):
connected_sites = self.get_connected_sites(idx, jimage=jimage)
site = PeriodicSite(
site.species,
np.add(site.frac_coords, jimage),
site.lattice,
properties=site.properties,
)
else:
connected_sites = self.get_connected_sites(idx)
connected_sites = [
cs for cs in connected_sites
if (cs.index, cs.jimage) in sites_to_draw
]
connected_sites_not_drawn = [
cs for cs in connected_sites
if (cs.index, cs.jimage) not in sites_to_draw
]
if color_edges:
connected_sites_colors = [
get_weight_color(cs.weight) for cs in connected_sites
]
connected_sites_not_drawn_colors = [
get_weight_color(cs.weight) for cs in connected_sites_not_drawn
]
else:
connected_sites_colors = None
connected_sites_not_drawn_colors = None
site_scene = site.get_scene(
connected_sites=connected_sites,
connected_sites_not_drawn=connected_sites_not_drawn,
hide_incomplete_edges=hide_incomplete_edges,
incomplete_edge_length_scale=incomplete_edge_length_scale,
connected_sites_colors=connected_sites_colors,
connected_sites_not_drawn_colors=connected_sites_not_drawn_colors,
explicitly_calculate_polyhedra_hull=
explicitly_calculate_polyhedra_hull,
legend=legend,
)
for scene in site_scene.contents:
primitives[scene.name] += scene.contents
primitives["unit_cell"].append(self.structure.lattice.get_scene())
return Scene(
name=self.structure.composition.reduced_formula,
origin=origin,
contents=[
Scene(name=k, contents=v, origin=origin)
for k, v in primitives.items()
],
)
def get_site_scene(
self,
connected_sites: List[ConnectedSite] = None,
# connected_site_metadata: None,
# connected_sites_to_draw,
connected_sites_not_drawn: List[ConnectedSite] = None,
hide_incomplete_edges: bool = False,
incomplete_edge_length_scale: Optional[float] = 1.0,
connected_sites_colors: Optional[List[str]] = None,
connected_sites_not_drawn_colors: Optional[List[str]] = None,
origin: Optional[List[float]] = None,
draw_polyhedra: bool = True,
explicitly_calculate_polyhedra_hull: bool = False,
bond_radius: float = 0.1,
draw_magmoms: bool = True,
magmom_scale: float = 1.0,
legend: Optional[Legend] = None,
) -> Scene:
"""
Args:
connected_sites:
connected_sites_not_drawn:
hide_incomplete_edges:
incomplete_edge_length_scale:
connected_sites_colors:
connected_sites_not_drawn_colors:
origin:
explicitly_calculate_polyhedra_hull:
legend:
Returns:
"""
atoms = []
bonds = []
polyhedron = []
magmoms = []
legend = legend or Legend(self)
# for disordered structures
is_ordered = self.is_ordered
phiStart, phiEnd = None, None
occu_start = 0.0
position = self.coords.tolist()
radii = [legend.get_radius(sp, site=self) for sp in self.species.keys()]
max_radius = float(min(radii))
for idx, (sp, occu) in enumerate(self.species.items()):
if isinstance(sp, DummySpecie):
cube = Cubes(
positions=[position], color=legend.get_color(sp, site=self), width=0.4
)
atoms.append(cube)
else:
color = legend.get_color(sp, site=self)
radius = legend.get_radius(sp, site=self)
# TODO: make optional/default to None
# in disordered structures, we fractionally color-code spheres,
# drawing a sphere segment from phi_end to phi_start
# (think a sphere pie chart)
if not is_ordered:
phi_frac_end = occu_start + occu
phi_frac_start = occu_start
occu_start = phi_frac_end
phiStart = phi_frac_start * np.pi * 2
phiEnd = phi_frac_end * np.pi * 2
name = str(sp)
if occu != 1.0:
name += " ({}% occupancy)".format(occu)
name += f" ({position[0]:.3f}, {position[1]:.3f}, {position[2]:.3f})"
if self.properties:
for k, v in self.properties.items():
name += f" ({k} = {v})"
sphere = Spheres(
positions=[position],
color=color,
radius=radius,
phiStart=phiStart,
phiEnd=phiEnd,
clickable=True,
tooltip=name,
)
atoms.append(sphere)
# Add magmoms
if draw_magmoms:
if magmom := self.properties.get("magmom"):
# enforce type
magmom = np.array(Magmom(magmom).get_moment())
magmom = 2 * magmom_scale * max_radius * magmom
tail = np.array(position) - 0.5 * np.array(magmom)
head = np.array(position) + 0.5 * np.array(magmom)
arrow = Arrows(
positionPairs=[[tail, head]],
color="red",
radius=0.20,
headLength=0.5,
headWidth=0.4,
clickable=True,
)
magmoms.append(arrow)
def get_structure_graph_scene(
self,
origin=None,
draw_image_atoms=True,
bonded_sites_outside_unit_cell=True,
hide_incomplete_edges=False,
incomplete_edge_length_scale=0.3,
color_edges_by_edge_weight=False,
edge_weight_color_scale="coolwarm",
explicitly_calculate_polyhedra_hull=False,
legend: Optional[Legend] = None,
group_by_site_property: Optional[str] = None,
bond_radius: float = 0.1,
) -> Scene:
origin = origin or list(
-self.structure.lattice.get_cartesian_coords([0.5, 0.5, 0.5]))
legend = legend or Legend(self.structure)
# we get primitives from each site individually, then
# combine into one big Scene
primitives = defaultdict(list)
sites_to_draw = self._get_sites_to_draw(
draw_image_atoms=draw_image_atoms,
bonded_sites_outside_unit_cell=bonded_sites_outside_unit_cell,
)
color_edges = False
if color_edges_by_edge_weight:
weights = [e[2].get("weight") for e in self.graph.edges(data=True)]
weights = np.array([w for w in weights if w])
if any(weights):
cmap = get_cmap(edge_weight_color_scale)
# try to keep color scheme symmetric around 0
weight_max = max([abs(min(weights)), max(weights)])
weight_min = -weight_max
def get_weight_color(weight):
if not weight:
weight = 0
x = (weight - weight_min) / (weight_max - weight_min)
return "#{:02x}{:02x}{:02x}".format(
*[int(c * 255) for c in cmap(x)[0:3]])
color_edges = True
if group_by_site_property:
# we will create sub-scenes for each group of atoms
# for example, if the Structure has a "wyckoff" site property
# this might be used to allow grouping by Wyckoff position,
# this then changes mouseover/interaction behavior with this scene
grouped_atom_scene_contents = defaultdict(list)
for (idx, jimage) in sites_to_draw:
site = self.structure[idx]
if jimage != (0, 0, 0):
connected_sites = self.get_connected_sites(idx, jimage=jimage)
site = PeriodicSite(
site.species,
np.add(site.frac_coords, jimage),
site.lattice,
properties=site.properties,
)
else:
connected_sites = self.get_connected_sites(idx)
connected_sites = [
cs for cs in connected_sites
if (cs.index, cs.jimage) in sites_to_draw
]
connected_sites_not_drawn = [
cs for cs in connected_sites
if (cs.index, cs.jimage) not in sites_to_draw
]
if color_edges:
connected_sites_colors = [
get_weight_color(cs.weight) for cs in connected_sites
]
connected_sites_not_drawn_colors = [
get_weight_color(cs.weight) for cs in connected_sites_not_drawn
]
else:
connected_sites_colors = None
connected_sites_not_drawn_colors = None
site_scene = site.get_scene(
connected_sites=connected_sites,
connected_sites_not_drawn=connected_sites_not_drawn,
hide_incomplete_edges=hide_incomplete_edges,
incomplete_edge_length_scale=incomplete_edge_length_scale,
connected_sites_colors=connected_sites_colors,
connected_sites_not_drawn_colors=connected_sites_not_drawn_colors,
explicitly_calculate_polyhedra_hull=
explicitly_calculate_polyhedra_hull,
legend=legend,
bond_radius=bond_radius,
)
for scene in site_scene.contents:
if group_by_site_property and scene.name == "atoms":
group_name = f"{site.properties[group_by_site_property]}"
scene.contents[0].tooltip = group_name
grouped_atom_scene_contents[group_name] += scene.contents
else:
primitives[scene.name] += scene.contents
if group_by_site_property:
atoms_scenes: List[Scene] = []
for k, v in grouped_atom_scene_contents.items():
atoms_scenes.append(Scene(name=k, contents=v))
primitives["atoms"] = atoms_scenes
primitives["unit_cell"].append(self.structure.lattice.get_scene())
# why primitives comprehension? just make explicit! more readable
return Scene(
name="StructureGraph",
origin=origin,
contents=[
Scene(name=k, contents=v, origin=origin)
for k, v in primitives.items()
],
)
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",
}