"""
visualize.py - 3Dmol.js 3D visualization and iRASPA CIF export for PoreGCN HF Space.
Frontend contract (app.py imports):
from visualize import create_3d_visualization, export_iraspa_cif
create_3d_visualization():
Returns an HTML string (with embedded 3Dmol.js) ready for gr.HTML().
- Atoms colored by attribution via B-factor channel (bwr colormap):
blue = negative attribution, white = neutral, red = positive attribution
- Metal atoms rendered as larger spheres (radius 0.6) to visually mark SBU centers
- Non-metal atoms rendered as sticks (radius 0.15)
- Pores shown as a Mercury-style translucent yellow void isosurface
(rolling 1.2 A probe + marching cubes; CCDC convention)
- Per-cavity attribution carried by small opaque orange/blue beads
at cavity centres, sized by |attribution|
- Unit cell wireframe via addUnitCell()
export_iraspa_cif():
Writes a CIF with _atom_site_B_iso_or_equiv overwritten by normalised
attributions in [1, 99] (1=most negative, 50=neutral, 99=most positive).
Compatible with iRASPA's Temperature Factor colouring.
"""
from __future__ import annotations
import os
import re
from pathlib import Path
from typing import Optional
import numpy as np
# Bundled 3Dmol.js library, loaded once at module import.
# Avoids CDN dependency so the viewer works on networks that block cdnjs/jsdelivr.
_3DMOL_JS_PATH = Path(__file__).parent / "3Dmol-min.js"
_3DMOL_JS_INLINE: Optional[str] = None
if _3DMOL_JS_PATH.exists():
_3DMOL_JS_INLINE = _3DMOL_JS_PATH.read_text(encoding="utf-8")
def _inline_3dmol_library(viewer_html: str, width: int = 800, height: int = 600) -> str:
"""Patch py3Dmol HTML so 3Dmol.js runs in an isolated iframe.
Two reasons we need the iframe:
1. Gradio's gr.HTML component sanitises '
patched_viewer = re.sub(
r"loadScriptAsync\(['\"]https?://[^'\"]+['\"]\)",
"Promise.resolve()",
viewer_html,
count=1,
)
full_doc = (
""
f"{inline_script}"
f""
f"{patched_viewer}"
)
# HTML-escape the doc for the srcdoc attribute (escape quotes and ampersands)
srcdoc = full_doc.replace("&", "&").replace('"', """)
iframe_html = (
f''
)
return iframe_html
try:
import py3Dmol
PY3DMOL_OK = True
except ImportError:
PY3DMOL_OK = False
try:
from pymatgen.core import Structure
from pymatgen.io.cif import CifWriter
PYMATGEN_OK = True
except ImportError:
PYMATGEN_OK = False
# Metal elements that get the large-sphere treatment in the 3D viewer.
# Covers most MOF secondary building units (SBUs).
_METAL_ELEMENTS = {
'Cu', 'Zn', 'Co', 'Zr', 'Fe', 'Mn', 'Ni', 'Mg', 'Ca', 'V',
'Al', 'Cr', 'Ti', 'Hf', 'In', 'Ga', 'Y', 'Ce', 'Nd', 'Tb',
'Eu', 'Gd', 'La', 'Mo', 'W', 'Ru', 'Rh', 'Pd', 'Cd', 'Sn',
'Ba', 'Sr', 'Li', 'Na', 'K', 'Rb', 'Cs',
}
def _standard_cif_lattice(abc, angles_deg):
"""Build the standard CIF Cartesian lattice matrix from a/b/c/alpha/beta/gamma.
The standard convention (used by 3Dmol.js, Mercury, VESTA, iRASPA when
parsing a CIF that specifies only cell scalars):
a along +x
b in the +xy plane (positive y)
c chosen so c_z > 0
Pymatgen does NOT always store its `Lattice.matrix` in this orientation;
when a CIF is loaded, the matrix may be rotated to a different frame.
Atom positions written by pymatgen's CifWriter contain only fractional
coordinates plus the cell scalars, so the rendering viewer (3Dmol)
re-derives Cartesian positions in the standard convention. Any geometry
we compute alongside the structure must be transformed into the same
standard frame before being passed to the viewer.
Returns a 3x3 ndarray with row vectors a, b, c.
"""
a, b, c = abc
al, be, ga = [np.deg2rad(x) for x in angles_deg]
ax, ay, az = a, 0.0, 0.0
bx, by, bz = b * np.cos(ga), b * np.sin(ga), 0.0
cx = c * np.cos(be)
sin_ga = np.sin(ga) if abs(np.sin(ga)) > 1e-12 else 1e-12
cy = c * (np.cos(al) - np.cos(be) * np.cos(ga)) / sin_ga
cz = np.sqrt(max(c * c - cx * cx - cy * cy, 0.0))
return np.array([[ax, ay, az], [bx, by, bz], [cx, cy, cz]], dtype=float)
def _compute_void_mesh(
structure,
probe_radius: float = 1.2,
target_grid_spacing: float = 0.7,
n_per_axis_min: int = 16,
n_per_axis_max: int = 60,
max_vertices: int = 40000,
vdw_eff: float = 1.5,
):
"""Compute a Mercury-style contact-surface void mesh for the unit cell.
Rolls a `probe_radius` probe (default 1.2 A, helium-sized) through the
unit cell and runs marching cubes at the void boundary. Mirrors
Mercury's Display > Voids > Contact Surface (CCDC, see
https://www.ccdc.cam.ac.uk/discover/blog/how-to-search-visualize-and-analyse-mof-structures).
Atoms are replicated to a 3x3x3 super-image so distance queries respect
periodic boundaries. A constant effective vdW radius of 1.5 A is used
instead of per-element vdW; for typical MOF elements (C, N, O, H, Cu,
Zn etc.) the per-element value is within 0.2 A of this mean and the
visual is unchanged.
Robustness behaviour. Per-axis grid resolution is clamped to
[`n_per_axis_min`, `n_per_axis_max`] so very large unit cells do not
blow up compute or browser payload, and very small cells still get a
decent surface. Mesh vertex count is capped at `max_vertices`; if the
surface would exceed that, the function returns (None, None) and the
caller falls back to a no-mesh render (atom colouring plus attribution
beads). All anticipated failure modes (missing scikit-image, missing
scipy, degenerate lattice, all-void or all-solid grid, marching-cubes
failure, oversize mesh) return (None, None).
Returns
-------
(vertices_cart, faces) where
vertices_cart: np.ndarray [V, 3] in Cartesian Angstroms
faces: np.ndarray [F, 3] of vertex indices per triangle
or (None, None).
"""
# Optional dependencies; degrade silently if either is missing
try:
from skimage.measure import marching_cubes
except Exception:
return None, None
try:
from scipy.spatial import cKDTree
except Exception:
return None, None
try:
lattice = np.asarray(structure.lattice.matrix, dtype=float)
a, b, c = structure.lattice.abc
if min(a, b, c) <= 0.0:
return None, None
# Per-axis resolution clamped so any cell stays within the compute
# and payload budget. Tiny cells get more samples than nominal
# spacing would suggest; huge cells get coarser sampling.
nx = int(np.clip(np.ceil(a / target_grid_spacing),
n_per_axis_min, n_per_axis_max))
ny = int(np.clip(np.ceil(b / target_grid_spacing),
n_per_axis_min, n_per_axis_max))
nz = int(np.clip(np.ceil(c / target_grid_spacing),
n_per_axis_min, n_per_axis_max))
fx = np.linspace(0.0, 1.0, nx, endpoint=False)
fy = np.linspace(0.0, 1.0, ny, endpoint=False)
fz = np.linspace(0.0, 1.0, nz, endpoint=False)
grid_frac = np.stack(np.meshgrid(fx, fy, fz, indexing='ij'), axis=-1)
grid_cart = grid_frac.reshape(-1, 3) @ lattice
atom_cart = np.array([site.coords for site in structure], dtype=float)
if atom_cart.shape[0] == 0:
return None, None
shifts = []
for da in (-1, 0, 1):
for db in (-1, 0, 1):
for dc in (-1, 0, 1):
shifts.append(da * lattice[0] + db * lattice[1] + dc * lattice[2])
images = np.vstack([atom_cart + s for s in shifts])
tree = cKDTree(images)
cutoff = vdw_eff + probe_radius
dist, _ = tree.query(grid_cart, k=1)
void = (dist > cutoff).reshape(nx, ny, nz).astype(np.float32)
if void.max() < 0.5 or void.min() > 0.5:
return None, None
try:
verts_idx, faces, _, _ = marching_cubes(
void, level=0.5, allow_degenerate=False
)
except Exception:
return None, None
# Vertex cap so the JSON payload shipped to the iframe stays
# bounded even for unusual very-high-porosity structures
if verts_idx.shape[0] > max_vertices:
return None, None
# Convert to fractional coordinates (frame-agnostic), then to
# Cartesian using the STANDARD CIF convention so the mesh aligns
# with how 3Dmol.js positions the atoms it parses from the CIF.
# Using `lattice` (pymatgen's matrix) here would put the mesh in a
# different frame from the rendered framework when the source CIF
# uses a non-canonical lattice orientation.
verts_frac = verts_idx / np.array([nx, ny, nz], dtype=float)
# Defensive clip: drop any face whose vertices stray outside
# [0, 1) fractional. With endpoint=False linspace this should
# never trigger, but keeps the mesh strictly inside the cell
# wireframe even if marching_cubes produces an edge case.
eps = 1e-6
inside = np.all(
(verts_frac >= -eps) & (verts_frac <= 1.0 + eps), axis=1
)
if not inside.all():
face_inside = inside[faces].all(axis=1)
faces = faces[face_inside]
if faces.shape[0] == 0:
return None, None
std_lat = _standard_cif_lattice(structure.lattice.abc,
structure.lattice.angles)
verts_cart = verts_frac @ std_lat
return verts_cart.astype(float), faces.astype(int)
except Exception:
# Catch-all so a malformed structure never breaks the rest of the
# rendering pipeline. Caller falls back to no-mesh render.
return None, None
def _cluster_pore_vertices(
positions: np.ndarray,
radii: np.ndarray,
attrs: np.ndarray,
) -> list:
"""Cluster Voronoi pore vertices into cavity centres.
Greedy non-maximum-suppression style: walk vertices in descending order
of inscribed-sphere radius, take the largest as a cavity centre, absorb
all unassigned vertices within that radius. The largest inscribed sphere
in a cavity *is* the cavity centre by construction (Zeo++ semantics), and
nearby smaller-radius vertices are pinch points along channels feeding it.
Returns a list of dicts with center, radius, and aggregated attribution
statistics over cluster members. The cavity-scale sphere idiom follows
Lisensky and Yaghi, J. Chem. Educ. 2022, 99, 1998-2004 (Figures 1-8),
where one translucent sphere per crystallographic pore is drawn at the
cavity centre, sized to the cavity diameter.
"""
n = len(positions)
if n == 0:
return []
positions = np.asarray(positions, dtype=float)
radii = np.asarray(radii, dtype=float)
attrs = np.asarray(attrs, dtype=float)
order = np.argsort(radii)[::-1]
consumed = np.zeros(n, dtype=bool)
clusters: list = []
for i in order:
if consumed[i]:
continue
centre = positions[i]
cavity_radius = float(max(radii[i], 0.5)) # floor to avoid invisible spheres
d = np.linalg.norm(positions - centre, axis=1)
members = np.where((d <= cavity_radius) & (~consumed))[0]
if len(members) == 0:
members = np.array([i])
consumed[members] = True
member_attrs = attrs[members]
# Choose the cluster's representative attribution as the signed value
# of the member with the largest magnitude. This preserves the sign of
# the strongest contributor instead of cancelling positives and
# negatives in a mean.
peak_idx = int(np.argmax(np.abs(member_attrs)))
peak_attr = float(member_attrs[peak_idx])
clusters.append({
'center': centre,
'radius': cavity_radius,
'attr': peak_attr,
'members': members,
})
return clusters
# =============================================================================
# Private helper: CIF string with B-factors injected
# =============================================================================
def _cif_string_with_bfactors(structure, per_atom_attrs: np.ndarray) -> str:
"""
Build a CIF text string from a pymatgen Structure with per-atom B-factors
encoding XAI attributions on the iRASPA scale:
1 = most negative attribution
50 = neutral
99 = most positive attribution
Both create_3d_visualization and export_iraspa_cif delegate here so the
B-factor injection logic is defined exactly once.
Atom fractional coordinates are wrapped to [0, 1) before writing the CIF
so that no atoms appear outside the unit cell wireframe in the rendered
viewer. Source CIFs occasionally place atoms at fractional positions
slightly past 1.0 (numeric precision or symmetry-related sites), which
otherwise produce visually confusing atoms drifting outside the cell.
Args:
structure: pymatgen Structure
per_atom_attrs: np.ndarray [N_atoms], signed attributions
Returns:
CIF text string with _atom_site_B_iso_or_equiv column appended.
"""
abs_max = max(float(np.max(np.abs(per_atom_attrs))), 1e-8)
b_factors = 50.0 + 49.0 * (per_atom_attrs / abs_max)
# Wrap fractional coordinates to [0, 1) so 3Dmol does not render atoms
# outside the unit cell wireframe. Atom order is preserved so that
# per_atom_attrs indexing remains valid.
try:
wrapped_struct = Structure(
lattice=structure.lattice,
species=[site.species for site in structure],
coords=[(np.asarray(site.frac_coords, dtype=float) % 1.0)
for site in structure],
coords_are_cartesian=False,
)
except Exception:
wrapped_struct = structure
writer = CifWriter(wrapped_struct)
cif_text = str(writer)
lines = cif_text.split('\n')
new_lines = []
in_atom_block = False
in_atom_data = False
b_col_injected = False
atom_site_columns_seen: list = []
atom_counter = 0
for line in lines:
stripped = line.strip()
if stripped.startswith('_atom_site_'):
in_atom_block = True
atom_site_columns_seen.append(stripped)
new_lines.append(line)
continue
if in_atom_block and not stripped.startswith('_atom_site_'):
if '_atom_site_B_iso_or_equiv' not in '\n'.join(atom_site_columns_seen):
new_lines.append(' _atom_site_B_iso_or_equiv')
b_col_injected = True
in_atom_block = False
in_atom_data = True
if (
in_atom_data
and b_col_injected
and stripped
and not stripped.startswith('_')
and not stripped.startswith('loop_')
and not stripped.startswith('#')
and len(stripped.split()) >= 4
and atom_counter < len(b_factors)
):
bf = b_factors[atom_counter]
line = f'{line} {bf:.2f}'
atom_counter += 1
elif in_atom_data and not stripped:
in_atom_data = False
new_lines.append(line)
return '\n'.join(new_lines)
# =============================================================================
# create_3d_visualization
# =============================================================================
def create_3d_visualization(
structure,
per_atom_attrs: np.ndarray,
per_pore_attrs: np.ndarray,
pore_positions: np.ndarray,
pore_radii: np.ndarray,
property_name: str,
prediction_value: Optional[float] = None,
scenario: Optional[str] = None,
width: int = 800,
height: int = 600,
) -> str:
"""
Render the MOF structure in 3Dmol.js with chemistry-style aesthetics.
Color scheme (bwr = blue-white-red, low B-factor to high):
B = 1 (most negative attribution) -> blue
B = 50 (neutral) -> white
B = 99 (most positive attribution) -> red
Rendering:
Metal atoms: large indigo spheres marking SBU centres
Non-metal atoms: ball-and-stick with Jmol colours
Void mesh: translucent yellow isosurface at the contact-surface
boundary (probe radius 1.2 A, grid 0.7 A)
Attribution: small orange/blue beads at cavity centres,
opacity 0.95, radius 0.7-1.4 A scaled by |attr|,
top-12 by |peak attribution| above 10% threshold
Unit cell: subtle gray wireframe
Args:
structure: pymatgen Structure
per_atom_attrs: np.ndarray [N_atoms], signed attributions
per_pore_attrs: np.ndarray [N_pores], signed attributions
pore_positions: np.ndarray [N_pores, 3] Cartesian Angstroms
pore_radii: np.ndarray [N_pores] Angstroms
property_name: Property string for the info header
prediction_value: Optional float; shown in header if given
scenario: Optional scenario letter ('A'/'B'/'C'/'D')
width: Canvas width in pixels (default 800)
height: Canvas height in pixels (default 600)
Returns:
HTML string with embedded 3Dmol.js, ready for gr.HTML().
"""
if not PY3DMOL_OK:
return (
''
'py3Dmol not installed. Run: pip install py3Dmol'
'
'
)
if not PYMATGEN_OK:
return (
''
'pymatgen not installed. Run: pip install pymatgen'
'
'
)
per_atom_attrs = np.asarray(per_atom_attrs, dtype=float)
per_pore_attrs = np.asarray(per_pore_attrs, dtype=float)
pore_positions = np.asarray(pore_positions, dtype=float)
pore_radii = np.asarray(pore_radii, dtype=float)
# Build CIF string with B-factors
try:
cif_str = _cif_string_with_bfactors(structure, per_atom_attrs)
except Exception as exc:
return (
f''
f'CIF generation failed: {exc}'
f'
'
)
# Create the 3Dmol viewer
view = py3Dmol.view(width=width, height=height)
view.addModel(cif_str, 'cif')
# Polyhedra-style aesthetic: metals as bold blue spheres (the SBU centers),
# other atoms as smaller spheres in soft gray shades. Inspired by iRASPA /
# crystal-toolkit visualisation conventions where metal SBUs are blue
# polyhedra and the linker network sits behind in muted tones.
metal_style = {'color': '#60a5fa'} # cornflower blue, visible on black
nonmetal_style = {'colorscheme': 'Jmol'} # standard chemistry colors
bfactor_style = nonmetal_style # name retained for downstream code paths
# Determine which elements are present in the structure
elements_present = set(str(site.specie) for site in structure)
metals_present = sorted(elements_present & _METAL_ELEMENTS)
nonmetals_present = sorted(elements_present - _METAL_ELEMENTS)
# Default: ball-and-stick with Jmol coloring for non-metal framework atoms
view.setStyle(
{},
{
'sphere': {'radius': 0.28, **nonmetal_style},
'stick': {'radius': 0.12, **nonmetal_style},
},
)
# Metal SBU centers: large indigo spheres to visually mark the polyhedra
if metals_present:
view.setStyle(
{'elem': metals_present},
{
'sphere': {'radius': 0.85, **metal_style},
'stick': {'radius': 0.16, **metal_style},
},
)
# Hydrogens: tiny spheres so they do not dominate the view
if 'H' in nonmetals_present:
view.setStyle(
{'elem': 'H'},
{
'sphere': {'radius': 0.14, **nonmetal_style},
'stick': {'radius': 0.06, **nonmetal_style},
},
)
# Mercury-style void isosurface: rolls a 1.2 A probe through the unit
# cell and renders a translucent yellow mesh at the void boundary. This
# is the chemistry-paper convention for showing MOF porosity (CCDC blog,
# most MOF supplementary information documents). The mesh carries the
# geometric signal (where are the pores, what shape are they). Per-cavity
# attribution sign is carried by the small coloured beads added below;
# per-atom attribution is carried by the framework atom B-factor colour.
try:
verts_cart, faces = _compute_void_mesh(structure)
except Exception:
verts_cart, faces = None, None
if verts_cart is not None and faces is not None and len(faces) > 0:
try:
vertex_arr = [
{'x': float(v[0]), 'y': float(v[1]), 'z': float(v[2])}
for v in verts_cart
]
face_arr = faces.flatten().tolist()
view.addCustom({
'vertexArr': vertex_arr,
'faceArr': face_arr,
'color': '#fbbf24', # warm Mercury-style yellow
'opacity': 0.55,
})
except Exception:
# Mesh shape construction failed (e.g. py3Dmol API mismatch on
# some build). Skip the mesh and continue with the rest of the
# render so the user still sees atoms and attribution beads.
pass
# Per-cavity attribution glyphs: small opaque beads at cluster centres.
# Size scales with |attribution| so the strongest contributing cavity
# is visually largest. The mesh above carries cavity geometry; these
# beads carry the directional attribution signal.
n_attribution_beads = 0
n_pores = len(per_pore_attrs)
if n_pores > 0 and len(pore_positions) == n_pores:
clusters = _cluster_pore_vertices(pore_positions, pore_radii, per_pore_attrs)
if clusters:
cluster_abs = np.array([abs(c['attr']) for c in clusters])
cluster_abs_max = float(cluster_abs.max()) if cluster_abs.max() > 0 else 1e-8
threshold = 0.10 * cluster_abs_max
ranked = sorted(clusters, key=lambda c: abs(c['attr']), reverse=True)
shown = 0
for c in ranked:
if shown >= 12:
break
if abs(c['attr']) < threshold:
break
centre = c['center']
color = '#ea580c' if c['attr'] >= 0 else '#1d4ed8'
# Bead radius capped at 0.9 A so the markers read as small
# attribution glyphs rather than rivalling the SBU spheres
bead_radius = 0.5 + 0.4 * (abs(c['attr']) / cluster_abs_max)
view.addSphere({
'center': {'x': float(centre[0]),
'y': float(centre[1]),
'z': float(centre[2])},
'radius': float(bead_radius),
'color': color,
'opacity': 1.0,
})
# Persistent label next to the bead so the user can read the
# signed cavity attribution without hovering. Only labelled
# when the rounded value is meaningfully non-zero, so we do
# not litter the figure with "+0.00" noise on cavities whose
# contribution is below the readable display precision.
if abs(c['attr']) >= 0.005:
view.addLabel(
f'{c["attr"]:+.2f}',
{
'position': {
'x': float(centre[0]) + float(bead_radius) + 0.4,
'y': float(centre[1]),
'z': float(centre[2]),
},
'backgroundColor': '#0f172a',
'backgroundOpacity': 0.6,
'fontColor': 'white',
'fontSize': 10,
'borderThickness': 0,
'showBackground': True,
'inFront': True,
},
)
shown += 1
n_attribution_beads += 1
# Unit cell wireframe drawn manually as 12 line segments in white, plus
# x/y/z labels at the axis tips. We intentionally bypass 3Dmol's
# addUnitCell() because its `astyle.hidden` flag for the colored a/b/c
# axis arrows is not honoured reliably across versions, and the default
# red/green/blue arrows competed with the attribution colour signal.
cell_lat = _standard_cif_lattice(structure.lattice.abc,
structure.lattice.angles)
corners_frac = np.array([
[0, 0, 0], [1, 0, 0], [0, 1, 0], [0, 0, 1],
[1, 1, 0], [1, 0, 1], [0, 1, 1], [1, 1, 1],
], dtype=float)
corners_cart = corners_frac @ cell_lat
edge_pairs = []
for i_e in range(8):
for j_e in range(i_e + 1, 8):
if abs(corners_frac[j_e] - corners_frac[i_e]).sum() == 1.0:
edge_pairs.append((i_e, j_e))
for i_e, j_e in edge_pairs:
view.addLine({
'start': {'x': float(corners_cart[i_e, 0]),
'y': float(corners_cart[i_e, 1]),
'z': float(corners_cart[i_e, 2])},
'end': {'x': float(corners_cart[j_e, 0]),
'y': float(corners_cart[j_e, 1]),
'z': float(corners_cart[j_e, 2])},
'color': '#e2e8f0',
'linewidth': 2.0,
})
# Axis tip labels using the crystallographic a/b/c convention
# (standard for crystal-structure papers; was briefly x/y/z but reverted
# because a/b/c is the established convention for periodic frameworks).
for tip_idx, label_text in [(1, 'a'), (2, 'b'), (3, 'c')]:
tip = corners_cart[tip_idx]
view.addLabel(label_text, {
'position': {'x': float(tip[0]),
'y': float(tip[1]),
'z': float(tip[2])},
'fontColor': 'white',
'fontSize': 14,
'backgroundColor': 'black',
'backgroundOpacity': 0.55,
'borderThickness': 0,
'inFront': True,
})
# Compute attribution scale so the hover label can de-scale B-factor back to
# the original signed attribution value the user understands.
abs_max = float(np.abs(per_atom_attrs).max()) if len(per_atom_attrs) > 0 else 1.0
if abs_max < 1e-12:
abs_max = 1.0
# Hover tooltip on atoms: element + index + position + signed attribution
# Attribution is recovered from the B-factor channel (set in _cif_string_with_bfactors):
# B = 50 + 49 * (attr / abs_max) -> attr = (B - 50) / 49 * abs_max
hover_callback = (
'function(atom, viewer) {'
' if(atom.label) return;'
f' var absMax = {abs_max:.6f};'
' var attr = ((atom.b - 50.0) / 49.0) * absMax;'
' var sign = attr >= 0 ? "+" : "";'
' var line1 = atom.elem + " #" + atom.serial;'
' var line2 = "Attribution: " + sign + attr.toFixed(4);'
' var line3 = "Coords: (" + atom.x.toFixed(2) + ", " + atom.y.toFixed(2) + ", " + atom.z.toFixed(2) + ") A";'
' atom.label = viewer.addLabel(line1 + "\\n" + line2 + "\\n" + line3, {'
' position: atom, backgroundColor: "#0f172a", fontColor: "#ffffff",'
' fontSize: 11, padding: 5, borderRadius: 3, alignment: "topLeft"'
' });'
'}'
)
unhover_callback = (
'function(atom, viewer) {'
' if(atom.label) { viewer.removeLabel(atom.label); delete atom.label; }'
'}'
)
view.setHoverable({}, True, hover_callback, unhover_callback)
# Finalize
view.setBackgroundColor('black')
view.zoomTo()
viewer_html = _inline_3dmol_library(view._make_html(), width=width, height=height)
# Build a property/scenario header to display above the 3Dmol canvas
label_map = {'A': 'Trustworthy', 'B': 'Overconfident',
'C': 'Underconfident', 'D': 'Unreliable'}
header_parts = [f'XAI — {property_name}']
if prediction_value is not None:
header_parts.append(f'pred = {prediction_value:.3g}')
if scenario:
sc_label = label_map.get(scenario, '')
sc_colors = {
'A': ('#16a34a', '#dcfce7'),
'B': ('#d97706', '#fef3c7'),
'C': ('#2563eb', '#dbeafe'),
'D': ('#dc2626', '#fee2e2'),
}
fg, bg = sc_colors.get(scenario, ('#334155', '#f1f5f9'))
header_parts.append(
f'Scenario {scenario}: {sc_label}'
)
header_html = (
''
+ ' • '.join(header_parts)
+ ''
'Hover atoms for details • Drag to rotate • Scroll to zoom'
'
'
)
# Legend: only show entries that actually appear in the rendered scene.
# Metals are overridden to cornflower blue in the viewer (to be visible
# on the black background), so we override their swatch colour here too
# rather than using the Jmol palette default. Cavity attribution beads
# have no Jmol equivalent and get their own legend block when drawn.
JMOL_COLORS = {
'H': '#FFFFFF', 'C': '#909090', 'N': '#3050F8', 'O': '#FF0D0D',
'F': '#90E050', 'Cl': '#1FF01F', 'Br': '#A62929', 'I': '#940094',
'P': '#FF8000', 'S': '#FFFF30', 'B': '#FFB5B5', 'Si': '#F0C8A0',
'Cu': '#C88033', 'Zn': '#7D80B0', 'Co': '#F090A0', 'Fe': '#E06633',
'Ni': '#50D050', 'Mn': '#9C7AC7', 'Cr': '#8A99C7', 'V': '#A6A6AB',
'Ti': '#BFC2C7', 'Zr': '#94E0E0', 'Hf': '#4DC2FF', 'Mg': '#8AFF00',
'Ca': '#3DFF00', 'Al': '#BFA6A6', 'Cd': '#FFD98F', 'La': '#70D4FF',
'Ce': '#FFFFC7', 'Ag': '#C0C0C0', 'Au': '#FFD123', 'Pb': '#575961',
}
METAL_VIEWER_COLOR = '#60a5fa' # must match metal_style above
def _legend_chip(color: str, text: str) -> str:
border = 'border:1px solid #cbd5e1;' if color.upper() == '#FFFFFF' else ''
return (
f''
f'{text}'
f''
)
element_chips = []
for elem in sorted(elements_present):
if elem in _METAL_ELEMENTS:
color = METAL_VIEWER_COLOR
else:
color = JMOL_COLORS.get(elem, '#888888')
element_chips.append(_legend_chip(color, elem))
cavity_chips_html = ''
if n_attribution_beads > 0:
cavity_chips_html = (
'Cavity attribution:'
+ _legend_chip('#ea580c', '+ contribution')
+ _legend_chip('#1d4ed8', '− contribution')
)
legend_html = (
''
'Elements:'
+ ''.join(element_chips)
+ cavity_chips_html
+ '
'
)
return header_html + viewer_html + legend_html
# =============================================================================
# export_attribution_csv
# =============================================================================
def export_attribution_csv(
structure,
per_atom_attrs: np.ndarray,
pore_positions: np.ndarray,
pore_radii: np.ndarray,
per_pore_attrs: np.ndarray,
property_name: str,
output_path: str,
) -> str:
"""
Write per-atom and per-pore signed attribution to a CSV for downstream
analysis (e.g. ranking sites, identifying design hotspots, or feeding into
follow-up MD/DFT calculations on high-attribution motifs).
Columns:
kind 'atom' or 'pore'
idx sequential index within kind
element chemical symbol (atoms only; 'pore' for pore vertices)
x, y, z Cartesian coordinates (Angstroms)
radius inscribed-sphere radius (pores only; blank for atoms)
attribution signed attribution to
b_factor attribution mapped to 1-99 iRASPA convention
(1=most negative, 50=neutral, 99=most positive)
"""
abs_max = float(np.abs(per_atom_attrs).max()) if len(per_atom_attrs) > 0 else 1.0
if len(per_pore_attrs) > 0:
abs_max = max(abs_max, float(np.abs(per_pore_attrs).max()))
if abs_max < 1e-12:
abs_max = 1.0
rows = [
f'# PoreGCN per-atom and per-pore attribution to {property_name}',
f'# Generated by huggingface.co/spaces/catenate/PoreGCN',
f'# abs_max = {abs_max:.6f}',
'kind,idx,element,x,y,z,radius,attribution,b_factor',
]
# Atom rows
for i, site in enumerate(structure):
attr = float(per_atom_attrs[i]) if i < len(per_atom_attrs) else 0.0
b = 50.0 + 49.0 * (attr / abs_max)
b = max(1.0, min(99.0, b))
rows.append(
f'atom,{i},{site.specie},{site.coords[0]:.4f},'
f'{site.coords[1]:.4f},{site.coords[2]:.4f},,'
f'{attr:.6f},{b:.2f}'
)
# Pore rows
n_pores = len(per_pore_attrs)
if n_pores > 0 and len(pore_positions) == n_pores:
for j in range(n_pores):
attr = float(per_pore_attrs[j])
b = 50.0 + 49.0 * (attr / abs_max)
b = max(1.0, min(99.0, b))
r = float(pore_radii[j]) if j < len(pore_radii) else 0.0
rows.append(
f'pore,{j},pore,{pore_positions[j,0]:.4f},'
f'{pore_positions[j,1]:.4f},{pore_positions[j,2]:.4f},'
f'{r:.3f},{attr:.6f},{b:.2f}'
)
with open(output_path, 'w', encoding='utf-8') as f:
f.write('\n'.join(rows))
f.write('\n')
return output_path
# =============================================================================
# export_iraspa_cif
# =============================================================================
def export_iraspa_cif(
structure,
per_atom_attrs: np.ndarray,
output_path: str,
) -> str:
"""
Write a CIF file with _atom_site_B_iso_or_equiv encoding XAI attributions.
B-factor mapping (symmetric around 50):
1 = most negative attribution (atom drives property DOWN strongly)
50 = neutral (attribution ~ 0)
99 = most positive attribution (atom drives property UP strongly)
iRASPA usage:
1. File > Open >
2. Appearance > Atoms > Color by > Temperature Factor
3. Apply a diverging blue-white-red colormap
Delegates B-factor injection to _cif_string_with_bfactors() so the
encoding logic is shared with create_3d_visualization.
Args:
structure: pymatgen Structure
per_atom_attrs: np.ndarray [N_atoms] signed attributions
output_path: Destination .cif file path (string)
Returns:
output_path (string) for gr.File() download component
"""
if not PYMATGEN_OK:
raise ImportError('pymatgen is required: pip install pymatgen')
per_atom_attrs = np.asarray(per_atom_attrs, dtype=float)
cif_text = _cif_string_with_bfactors(structure, per_atom_attrs)
output_path = str(output_path)
os.makedirs(os.path.dirname(output_path) if os.path.dirname(output_path) else '.', exist_ok=True)
Path(output_path).write_text(cif_text, encoding='utf-8')
return output_path