import subprocess
import tempfile
from pathlib import Path
import jinja2
import numpy as np
from pysisyphus.config import get_cmd
from pysisyphus.constants import BOHR2ANG
from pysisyphus.helpers_pure import interpolate_colors
CUBE_ISO_VAL = 0.001
TPL_BASE = """
{{ orient }}
function _setModelState() {
select;
Spacefill 0.0;
frank off;
font frank 16.0 SansSerif Plain;
select *;
set fontScaling false;
background white
frank off
set showhydrogens True;
}
_setModelState;
"""
CUBE_TPL = jinja2.Template(
"load {{ cube_fn }}\n"
+ TPL_BASE
+ """isosurface cutoff {{ isoval }} sign {{ colors }} "{{ cube_fn }}"
color isosurface translucent
{% if png_fn %}write image pngt "{{ png_fn }}"{% endif %}
"""
)
[docs]
def call_jmol(spt_str, show=False):
with tempfile.NamedTemporaryFile("w", suffix=".spt") as spt_handle:
spt_handle.write(spt_str)
spt_handle.flush()
jmol_cmd = [get_cmd("jmol"), "-n", spt_handle.name]
if show:
del jmol_cmd[1]
proc = subprocess.Popen(
jmol_cmd,
universal_newlines=True,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
)
proc.wait()
stdout = proc.stdout.read()
stderr = proc.stderr.read()
return stdout, stderr
[docs]
def view_cdd_cube(cdd_cube, isoval=CUBE_ISO_VAL, orient=""):
spt = CUBE_TPL.render(
orient=orient,
cube_fn=cdd_cube,
isoval=isoval,
colors="red blue",
)
with open("jmol.spt", "w") as handle:
handle.write(spt)
stdout, stderr = call_jmol(spt, show=True)
[docs]
def render_cdd_cube(cdd_cube, isoval=CUBE_ISO_VAL, orient=""):
png_fn = Path(cdd_cube).with_suffix(".png")
spt = CUBE_TPL.render(
orient=orient,
cube_fn=cdd_cube,
isoval=isoval,
colors="red blue",
png_fn=png_fn,
)
with open("jmol.spt", "w") as handle:
handle.write(spt)
stdout, stderr = call_jmol(spt)
return png_fn
[docs]
def render_geom_and_charges(geom, point_charges):
point_charges = point_charges.copy()
point_charges[:, :3] *= BOHR2ANG
charges = point_charges[:, -1]
cr = np.array((255, 0, 0)) # red
cw = np.array((255, 255, 255)) # white
cb = np.array((0, 0, 255)) # blue
chrg_min = charges.min()
chrg_max = charges.max()
print(
f"charges:\n{np.array2string(point_charges, precision=4)}\n"
f"min(charges): {chrg_min: .4f}\nmax(charges): {chrg_max: .4f}\n"
)
c1 = cb if chrg_min < 0.0 else cw
c2 = cr if chrg_max > 0.0 else cw
rgb_colors, _ = interpolate_colors(charges, c1, c2)
# Dump geometry to temporary file
with tempfile.NamedTemporaryFile("w", suffix=".xyz") as tmp_xyz:
tmp_xyz.write(geom.as_xyz())
tmp_xyz.flush()
spt = f"load {tmp_xyz.name};\n"
for i, ((x, y, z, pc), (r, g, b)) in enumerate(zip(point_charges, rgb_colors)):
id_ = f"chrg{i}"
spt += (
f"isosurface {id_} center {{{x} {y} {z}}} sphere 0.25;\n"
f"color ${id_} [{r} {g} {b}];\n"
)
print(f"SPT:\n\n{spt}")
call_jmol(spt, show=True)