"""
Export wrapper for asymptote (ASY)
For creating publication quality plots
Since ASY does not have the nested tree structure of threejs,
we just have to traverse the tree and draw each material as we see them.
TODO The code should also appends a set of special points at the end in case the user wants to add
more "hand drawn" features to the plot
"""
import logging
from itertools import chain
from jinja2 import Environment
from pymatgen.analysis.graphs import StructureGraph
from pymatgen.core.structure import Structure
from crystal_toolkit.helpers.utils import update_object_args
logger = logging.getLogger(__name__)
HEAD = """
import settings;
import solids;
size(300);
outformat="png";
defaultshininess = 0.8;
currentlight = light(0,0,400);
// Camera information
currentprojection=orthographic (
camera=(8,5,4),
up=(0,0,1),
target={{target}},
zoom=0.5
);
// Basic function for drawing spheres
void drawSpheres(triple[] C, real R, pen p=currentpen){
for(int i=0;i<C.length;++i){
draw(sphere(C[i],R).surface(
new pen(int i, real j){return p;}
)
);
}
}
// Draw a sphere without light
void drawSpheres_nolight(triple[] C, real R, pen p=currentpen){
material nlpen = material(diffusepen=opacity(1.0), emissivepen=p, shininess=0);
for(int i=0;i<C.length;++i){
revolution s_rev = sphere(C[i],R);
surface s_surf = surface(s_rev);
draw(s_surf, nlpen);
draw(s_rev.silhouette(100), black+linewidth(3));
}
}
// Draw a cylinder
void Draw(guide3 g,pen p=currentpen, real cylR=0.2){
draw(
cylinder(
point(g,0),cylR,arclength(g),point(g,1)-point(g,0)
).surface(
new pen(int i, real j){
return p;
}
)
);
}
// Draw a cylinder without light
void Draw_nolight(guide3 g,pen p=currentpen, real cylR=0.2){
material nlpen = material(diffusepen=opacity(1.0), emissivepen=p, shininess=0);
revolution s_rev = cylinder(point(g,0),cylR,arclength(g),point(g,1)-point(g,0));
surface s_surf = surface(s_rev);
draw(s_surf, nlpen);
draw(s_rev.silhouette(100), black+linewidth(3));
}
"""
TEMP_SPHERE = """
{% for val in positions %}
triple sphere{{loop.index}}={{val}};
{% endfor %}
triple[] spheres = {
{%- for val in positions -%}
sphere{{loop.index}}
{%- if not loop.last %},{% endif %}
{%- endfor -%}
};
drawSpheres(spheres, {{radius}}, rgb('{{color}}'));
"""
TEMP_SPHERE_NOLIGHT = """
{% for val in positions %}
triple sphere{{loop.index}}={{val}};
{% endfor %}
triple[] spheres = {
{%- for val in positions -%}
sphere{{loop.index}}
{%- if not loop.last %},{% endif %}
{%- endfor -%}
};
drawSpheres_nolight(spheres, {{radius}}, rgb('{{color}}'));
"""
TEMP_CYLINDER = """
pen connectPen=rgb('{{color}}');
{% for ipos, fpos in posPairs %}
triple IPOS = {{ipos}};
triple FPOS = {{fpos}};
Draw(IPOS--FPOS, connectPen, {{radius}});
{% endfor %}
"""
TEMP_LINE = """
pen connectPen=rgb('{{color}}');
{% for ipos, fpos in posPairs %}
triple IPOS = {{ipos}};
triple FPOS = {{fpos}};
draw(IPOS--FPOS, connectPen);
{% endfor %}
"""
TEMP_CYLINDER = """
pen connectPen=rgb('{{color}}');
{% for ipos, fpos in posPairs %}
triple IPOS = {{ipos}};
triple FPOS = {{fpos}};
Draw(IPOS--FPOS, connectPen, {{radius}});
{% endfor %}
"""
TEMP_SURF = """
real[][] A = {
{% for ipos in positions -%}
{{ipos}},
{% endfor %}
};
A = transpose(A);
material m = rgb('{{face_color}}') + opacity({{opac}});
triple a,b,c;
for(int i=0; i < A[0].length/3; ++i) {
a=(A[0][i*3],A[1][i*3],A[2][i*3]);
b=(A[0][i*3+1],A[1][i*3+1],A[2][i*3+1]);
c=(A[0][i*3+2],A[1][i*3+2],A[2][i*3+2]);
draw(surface(a--b--c--cycle),surfacepen = m); // draw i-th triangle
}
path3 no_show = path3(scale(0) * box((-1,-1),(1,1)));
draw(surface(no_show), surfacepen=m); // draw i-th triangle
"""
TEMP_SURF_NOLIGHT = """
real[][] A = {
{% for ipos in positions -%}
{{ipos}},
{% endfor %}
};
A = transpose(A);
material m = material(diffusepen=opacity({{opac}}), emissivepen=rgb("{{face_color}}"), shininess=0);
triple a,b,c;
for(int i=0; i < A[0].length/3; ++i) {
a=(A[0][i*3],A[1][i*3],A[2][i*3]);
b=(A[0][i*3+1],A[1][i*3+1],A[2][i*3+1]);
c=(A[0][i*3+2],A[1][i*3+2],A[2][i*3+2]);
draw(surface(a--b--c--cycle),surfacepen = m); // draw i-th triangle
}
path3 no_show = path3(scale(0) * box((-1,-1),(1,1)));
draw(surface(no_show), surfacepen=m); // draw i-th triangle
"""
# Functions for parsing specific shapes
# Each function should only take in the CTK object and the changes to the
# default display parameters
def _get_lines(ctk_scene, d_args=None):
"""return the ASY string output to draw cylinders
Arguments:
ctk_scene {Scene} -- CTK Scene Object with ctk_scene.type == 'cylinders'
Keyword Arguments:
d_args {dict} -- User defined defaults of the plot (default: {None})
"""
assert ctk_scene.type == "lines"
updated_defaults = update_object_args(
d_args, object_name="Lines", allowed_args=["linewidth", "color"]
)
ipos = map(tuple, ctk_scene.positions[0::2])
fpos = map(tuple, ctk_scene.positions[1::2])
posPairs = [*zip(ipos, fpos)]
linewidth = ctk_scene.linewidth or updated_defaults["linewidth"]
color = ctk_scene.color or updated_defaults["color"]
color = color.replace("#", "")
return (
Environment()
.from_string(TEMP_LINE)
.render(posPairs=posPairs, color=color, linewidth=linewidth)
)
def _get_spheres(ctk_scene, d_args=None):
"""return the ASY string output to draw the spheres
Arguments:
ctk_scene {Scene} -- CTK Scene Object with ctk_scene.type == 'spheres'
Keyword Arguments:
d_args {dict} -- User defined defaults of the plot (default: {None})
"""
assert ctk_scene.type == "spheres"
updated_defaults = update_object_args(
d_args, object_name="Spheres", allowed_args=["radius", "color"]
)
positions = [tuple(pos) for pos in ctk_scene.positions]
radius = ctk_scene.radius or updated_defaults["radius"]
color = ctk_scene.color or updated_defaults["color"]
color = color.replace("#", "")
if ctk_scene.phiStart or ctk_scene.phiEnd:
raise NotImplementedError
# phiStart = ctk_scene.phiStart or 0 # not yet implemented
# phiEnd = ctk_scene.phiEnd or 2*pi
return (
Environment()
.from_string(TEMP_SPHERE)
.render(positions=positions, radius=radius, color=color)
)
def _get_cylinders(ctk_scene, d_args=None):
"""return the ASY string output to draw cylinders
Arguments:
ctk_scene {Scene} -- CTK Scene Object with ctk_scene.type == 'cylinders'
Keyword Arguments:
d_args {dict} -- User defined defaults of the plot (default: {None})
"""
assert ctk_scene.type == "cylinders"
updated_defaults = update_object_args(
d_args, object_name="Cylinders", allowed_args=["radius", "color"]
)
posPairs = [[tuple(ipos), tuple(fpos)] for ipos, fpos in ctk_scene.positionPairs]
radius = ctk_scene.radius or updated_defaults["radius"]
color = ctk_scene.color or updated_defaults["color"]
color = color.replace("#", "")
return (
Environment()
.from_string(TEMP_CYLINDER)
.render(posPairs=posPairs, color=color, radius=radius)
)
def _get_surface(ctk_scene, d_args=None):
"""return the ASY string output to draw cylinders
Arguments:
ctk_scene {Scene} -- CTK Scene Object with ctk_scene.type == 'surface'
Keyword Arguments:
d_args {dict} -- User defined defaults of the plot (default: {None})
"""
assert ctk_scene.type == "surface"
if len(ctk_scene.positions) == 0:
return "" # print nothing
updated_defaults = update_object_args(
d_args, object_name="Surfaces", allowed_args=["opacity", "color", "edge_width"]
)
color = ctk_scene.color or updated_defaults["color"]
color = color.replace("#", "")
opacity = ctk_scene.opacity or updated_defaults["opacity"]
positions = tuple(map(lambda x: f"{{{x[0]}, {x[1]}, {x[2]}}}", ctk_scene.positions))
num_triangle = len(ctk_scene.positions) / 3.0
# sanity check the mesh must be triangles
assert num_triangle.is_integer()
# # make decision on transparency
# transparent = if obj_args['opacity'] < 0.99
#
# # asymptote just needs the xyz positions
num_triangle = int(num_triangle)
pos_xyz = tuple(
chain.from_iterable(
[
(positions[itr * 3], positions[itr * 3 + 1], positions[itr * 3 + 2])
for itr in range(num_triangle)
]
)
)
#
# # write the data array
data_array_asy = (
Environment()
.from_string(TEMP_SURF)
.render(positions=pos_xyz, face_color=color, opac=opacity)
)
return data_array_asy
# write the
[docs]def asy_write_data(input_scene_comp, fstream):
"""
parse a primitive display object in crystaltoolkit and print it to asymptote
input_scene_comp
fstream
"""
if input_scene_comp.type == "spheres":
asy_out = _get_spheres(input_scene_comp)
fstream.write(asy_out)
if input_scene_comp.type == "cylinders":
asy_out = _get_cylinders(input_scene_comp)
fstream.write(asy_out)
if input_scene_comp.type == "lines":
asy_out = _get_lines(input_scene_comp)
fstream.write(asy_out)
if input_scene_comp.type == "surface":
asy_out = _get_surface(input_scene_comp)
fstream.write(asy_out)
return
# TODO we can make the line solide for the foreground and dashed for the background
# This will require use to modify the way the line objects are generated
# at each vertex in the unit cell, we can evaluate the sum of all three lattice vectors from the point
# then the <vec_sum | vec_to_camera> for each vertex. The smallest
# normalized vertex contains the three lines that should be dashed
[docs]def filter_data(scene_data, fstream):
"""
Recursively traverse the scene_data dictionary to find objects to draw
"""
if "type" in scene_data:
asy_write_data(scene_data, fstream)
else:
for itr in scene_data["contents"]:
filter_data(itr, fstream)
[docs]def traverse_scene_object(scene_data, fstream):
"""
Traverse object
"""
for sub_object in scene_data.contents:
if isinstance(sub_object, list):
for iobj in sub_object:
traverse_scene_object(iobj)
continue
elif hasattr(sub_object, "type"):
asy_write_data(sub_object, fstream)
else:
traverse_scene_object(sub_object, fstream)
[docs]def write_ctk_scene_to_file(ctk_scene, file_name):
"""
ctk_scene : Scene object from crystaltoolkit
filename : Output asymptote file and location
"""
fstream = open(file_name, "w")
target = tuple(-ii for ii in ctk_scene.origin)
header = Environment().from_string(HEAD).render(target=target)
fstream.write(header)
traverse_scene_object(ctk_scene, fstream)
fstream.close()
[docs]def write_asy_file(renderable_object, file_name, **kwargs):
"""
Generate the scene object and write it to file
Args:
renderable_object: Object to be rendered
file_name: name of file
"""
if isinstance(renderable_object, Structure) or isinstance(
renderable_object, StructureGraph
):
kwargs["explicitly_calculate_polyhedra_hull"] = True
write_ctk_scene_to_file(renderable_object.get_scene(**kwargs), file_name)