|
| 1 | +#!/usr/bin/env -S uv run --script |
| 2 | +# /// script |
| 3 | +# requires-python = ">=3.11" |
| 4 | +# dependencies = [ |
| 5 | +# "pymatgen>=2025.1", |
| 6 | +# "scikit-image>=0.24", |
| 7 | +# "trimesh>=4.5", |
| 8 | +# "numpy>=2.0", |
| 9 | +# "click>=8.1", |
| 10 | +# "boto3>=1.35", |
| 11 | +# ] |
| 12 | +# /// |
| 13 | +"""Pre-compute GLB-encoded isosurface meshes for a CHGCAR at several quantile |
| 14 | +iso-levels. Output meshes are tiny (~50-200 KB each) and load instantly in the |
| 15 | +browser via Three.js GLTFLoader, bypassing the download-and-march-cubes step. |
| 16 | +
|
| 17 | +Usage: |
| 18 | + generate-glb.py path/to/mp-XXX.CHGCAR [--s3 s3://bucket/prefix/] |
| 19 | +
|
| 20 | + Writes: |
| 21 | + <out_dir>/<material_id>/<quantile>.glb |
| 22 | + e.g. |
| 23 | + out/mp-1000020/0.50.glb (512 KB) |
| 24 | + out/mp-1000020/0.75.glb (300 KB) |
| 25 | + out/mp-1000020/0.90.glb (190 KB) |
| 26 | + out/mp-1000020/0.95.glb (140 KB) |
| 27 | + out/mp-1000020/0.99.glb (80 KB) |
| 28 | +
|
| 29 | +Iso-level quantiles default to [.50, .75, .90, .95, .99] — same range the |
| 30 | +ELvis iso-slider covers, matching the client-side quantile mapping. |
| 31 | +""" |
| 32 | +from __future__ import annotations |
| 33 | + |
| 34 | +import re |
| 35 | +import sys |
| 36 | +from functools import partial |
| 37 | +from pathlib import Path |
| 38 | + |
| 39 | +import click |
| 40 | +import numpy as np |
| 41 | +import trimesh |
| 42 | +from pymatgen.io.vasp import Chgcar |
| 43 | +from skimage import measure |
| 44 | + |
| 45 | +err = partial(print, file=sys.stderr) |
| 46 | + |
| 47 | + |
| 48 | +def iso_from_quantile(data: np.ndarray, q: float) -> float: |
| 49 | + flat = data.ravel() |
| 50 | + return float(np.quantile(flat, q)) |
| 51 | + |
| 52 | + |
| 53 | +def extract_mesh(grid: np.ndarray, iso: float, lattice: np.ndarray) -> trimesh.Trimesh | None: |
| 54 | + """Run marching cubes and convert fractional → Cartesian coordinates.""" |
| 55 | + try: |
| 56 | + verts, faces, normals, _ = measure.marching_cubes(grid, level=iso, allow_degenerate=False) |
| 57 | + except (ValueError, RuntimeError) as e: |
| 58 | + err(f' marching_cubes failed at iso={iso:.3f}: {e}') |
| 59 | + return None |
| 60 | + if len(verts) == 0: |
| 61 | + return None |
| 62 | + # `verts` is in voxel index space; normalize to fractional [0, 1] |
| 63 | + dims = np.array(grid.shape, dtype=np.float64) |
| 64 | + frac = verts / (dims - 1) |
| 65 | + # Fractional → Cartesian |
| 66 | + cart = frac @ lattice |
| 67 | + return trimesh.Trimesh(vertices=cart, faces=faces, vertex_normals=normals, process=False) |
| 68 | + |
| 69 | + |
| 70 | +@click.command() |
| 71 | +@click.option('-o', '--out-dir', type=click.Path(path_type=Path), default='out/glb', show_default=True) |
| 72 | +@click.option('-q', '--quantiles', default='0.50,0.75,0.90,0.95,0.99', help='Comma-separated iso quantiles') |
| 73 | +@click.option('-s', '--s3-prefix', default=None, help='Upload under this s3://bucket/prefix/ after generating') |
| 74 | +@click.option('-f', '--force', is_flag=True, help='Overwrite existing files') |
| 75 | +@click.argument('chgcar_path', type=click.Path(exists=True, path_type=Path)) |
| 76 | +def main(chgcar_path: Path, out_dir: Path, quantiles: str, s3_prefix: str | None, force: bool): |
| 77 | + mat_match = re.search(r'(mp-\d+)', chgcar_path.name) |
| 78 | + if not mat_match: |
| 79 | + raise click.ClickException(f'No mp-XXX ID in filename {chgcar_path.name!r}') |
| 80 | + mat_id = mat_match.group(1) |
| 81 | + |
| 82 | + qs = [float(x) for x in quantiles.split(',')] |
| 83 | + err(f'[{mat_id}] loading {chgcar_path}') |
| 84 | + chg = Chgcar.from_file(str(chgcar_path)) |
| 85 | + grid = chg.data['total'] |
| 86 | + lattice = chg.structure.lattice.matrix # 3x3, rows = a/b/c vectors |
| 87 | + err(f'[{mat_id}] grid {grid.shape}, min={grid.min():.3f}, max={grid.max():.3f}, lattice volume={chg.structure.lattice.volume:.2f} A^3') |
| 88 | + |
| 89 | + out_mat_dir = out_dir / mat_id |
| 90 | + out_mat_dir.mkdir(parents=True, exist_ok=True) |
| 91 | + |
| 92 | + uploaded = [] |
| 93 | + for q in qs: |
| 94 | + out_path = out_mat_dir / f'{q:.2f}.glb' |
| 95 | + if out_path.exists() and not force: |
| 96 | + err(f' q={q:.2f}: exists, skipping ({out_path})') |
| 97 | + continue |
| 98 | + iso = iso_from_quantile(grid, q) |
| 99 | + err(f' q={q:.2f} → iso={iso:.3f}, extracting mesh...') |
| 100 | + mesh = extract_mesh(grid, iso, lattice) |
| 101 | + if mesh is None: |
| 102 | + err(f' q={q:.2f}: empty mesh, skipping') |
| 103 | + continue |
| 104 | + mesh.export(str(out_path), file_type='glb') |
| 105 | + size_kb = out_path.stat().st_size / 1024 |
| 106 | + err(f' q={q:.2f}: wrote {out_path} ({size_kb:.1f} KB, {len(mesh.faces)} triangles)') |
| 107 | + uploaded.append(out_path) |
| 108 | + |
| 109 | + if s3_prefix and uploaded: |
| 110 | + import boto3 |
| 111 | + from urllib.parse import urlparse |
| 112 | + parsed = urlparse(s3_prefix) |
| 113 | + if parsed.scheme != 's3': |
| 114 | + raise click.ClickException(f'--s3 must be s3://bucket/prefix/: got {s3_prefix!r}') |
| 115 | + bucket = parsed.netloc |
| 116 | + key_prefix = parsed.path.lstrip('/') |
| 117 | + if not key_prefix.endswith('/'): |
| 118 | + key_prefix += '/' |
| 119 | + s3 = boto3.client('s3') |
| 120 | + for p in uploaded: |
| 121 | + key = f'{key_prefix}{mat_id}/{p.name}' |
| 122 | + s3.upload_file(str(p), bucket, key, ExtraArgs={'ContentType': 'model/gltf-binary'}) |
| 123 | + err(f' uploaded s3://{bucket}/{key}') |
| 124 | + |
| 125 | + |
| 126 | +if __name__ == '__main__': |
| 127 | + main() |
0 commit comments