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_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()
],
)