From 433ebce6af72e59fc81a73dd93c4b7b6f084ee87 Mon Sep 17 00:00:00 2001 From: Hananeh Oliaei Date: Tue, 19 May 2026 14:55:11 -0400 Subject: [PATCH 1/5] Add PADS generation script and config Adds create_pads.py to reconstruct charge density grids from atomic radial profiles using a Numba-accelerated kernel, and config_pads.yaml to src/electrai/configs/. Also adds numba as a project dependency. Co-Authored-By: Claude Sonnet 4.6 --- pyproject.toml | 1 + src/electrai/configs/config_pads.yaml | 9 + src/electrai/pads_generation/create_pads.py | 257 ++++++++++++++++++++ uv.lock | 97 ++++++++ 4 files changed, 364 insertions(+) create mode 100644 src/electrai/configs/config_pads.yaml create mode 100644 src/electrai/pads_generation/create_pads.py diff --git a/pyproject.toml b/pyproject.toml index d8920e88..fca495e5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -36,6 +36,7 @@ dependencies = [ "pymatgen>=2025.10.7", "zarr>=3.1.3", "hydra-core>=1.3.2", + "numba>=0.65.1", ] [project.optional-dependencies] diff --git a/src/electrai/configs/config_pads.yaml b/src/electrai/configs/config_pads.yaml new file mode 100644 index 00000000..568bdd90 --- /dev/null +++ b/src/electrai/configs/config_pads.yaml @@ -0,0 +1,9 @@ +paths: + input_file: /scratch/gpfs/ROSENGROUP/ho0950/electrai/data/MP/chgcars/label/mp-1775579.CHGCAR # or /scratch/gpfs/ROSENGROUP/ho0950/electrai/data/MP/jsongz/label/mp-1775579.CHGCAR.json.gz + output_dir: ./data + stats_dir: ./stats + radial_dir: /scratch/gpfs/ROSENGROUP/common/globus_share_OA/mp/chg_datasets/atomic_radial_profiles + zval_file: /scratch/gpfs/ROSENGROUP/common/globus_share_OA/mp/chg_datasets/zval.json + +pads: + extra_shells: 1 diff --git a/src/electrai/pads_generation/create_pads.py b/src/electrai/pads_generation/create_pads.py new file mode 100644 index 00000000..9e551d86 --- /dev/null +++ b/src/electrai/pads_generation/create_pads.py @@ -0,0 +1,257 @@ +from __future__ import annotations + +import argparse +import json +import shutil +from itertools import product as iproduct +from pathlib import Path + +import numpy as np +import yaml +from numba import njit, prange + +from electrai.zarr_conversion.convert_to_zarr import load_chgcar +from electrai.zarr_conversion.zarr_writer import write_chgcar_to_zarr + + +# -------------------- ARGS -------------------- +def _derive_task_id(path: Path) -> str: + name = path.name + if name.endswith(".json.gz"): + return name[: -len(".json.gz")] + if name.lower().endswith(".chgcar"): + return name[: -len(".chgcar")] + return path.stem + + +def parse_args(): + p = argparse.ArgumentParser() + p.add_argument( + "--config", + type=Path, + default=Path(__file__).parent.parent / "configs" / "config_pads.yaml", + ) + p.add_argument( + "--input", type=Path, help="Path to .json.gz or .CHGCAR file (overrides config)" + ) + return p.parse_args() + + +# -------------------- NUMBA CORE -------------------- +@njit(parallel=True, fastmath=True, cache=True) +def _accumulate_images( + grid_cart, atom_images_cart, r_valid, d_valid, r_max, r_start, dr +): + nvox = grid_cart.shape[0] + nimg = atom_images_cart.shape[0] + r_max2 = r_max * r_max + M = r_valid.shape[0] + + total = np.zeros(nvox, dtype=np.float32) + + for vi in prange(nvox): + gx = grid_cart[vi, 0] + gy = grid_cart[vi, 1] + gz = grid_cart[vi, 2] + acc = np.float32(0.0) + + for ii in range(nimg): + dx = gx - atom_images_cart[ii, 0] + dy = gy - atom_images_cart[ii, 1] + dz = gz - atom_images_cart[ii, 2] + d2 = dx * dx + dy * dy + dz * dz + if d2 > r_max2: + continue + + r = np.sqrt(d2) + + idx = int((r - r_start) / dr) + if idx < 0: + acc += d_valid[0] + elif idx < M - 1: + t = (r - (r_start + idx * dr)) / dr + acc += d_valid[idx] + t * (d_valid[idx + 1] - d_valid[idx]) + + total[vi] = acc + + return total + + +# -------------------- GRID -------------------- +def build_voxel_grid(cell, grid_shape): + Nx, Ny, Nz = grid_shape + cell = np.asarray(cell, dtype=np.float32) + + fx = (np.arange(Nx, dtype=np.float32) + 0.5) / Nx + fy = (np.arange(Ny, dtype=np.float32) + 0.5) / Ny + fz = (np.arange(Nz, dtype=np.float32) + 0.5) / Nz + + gx, gy, gz = np.meshgrid(fx, fy, fz, indexing="ij") + grid_frac = np.stack([gx.ravel(), gy.ravel(), gz.ravel()], axis=1) + + return grid_frac @ cell + + +# -------------------- RECONSTRUCTION -------------------- +def reconstruct_from_radial_function( + radial_r, + radial_density, + ref_point, + cell, + grid_shape, + extra_shells=1, + prebuilt_grid=None, +): + r_valid = np.asarray(radial_r, dtype=np.float32) + d_valid = np.asarray(radial_density, dtype=np.float32) + cell = np.asarray(cell, dtype=np.float32) + + r_max = float(r_valid[-1]) + r_start = float(r_valid[0]) + dr = float(r_valid[1] - r_valid[0]) + + grid_cart = ( + prebuilt_grid + if prebuilt_grid is not None + else build_voxel_grid(cell, grid_shape) + ) + + lat_lengths = np.linalg.norm(cell, axis=1) + nmax = int(np.ceil(r_max / lat_lengths.min())) + extra_shells + + shifts = np.array( + list(iproduct(range(-nmax, nmax + 1), repeat=3)), dtype=np.float32 + ) + atom0_frac = np.asarray(ref_point, dtype=np.float32) + atom_images_cart = (atom0_frac[None, :] + shifts) @ cell + + corners = np.array( + [[i, j, k] for i in [0, 1] for j in [0, 1] for k in [0, 1]], dtype=np.float32 + ) + corners_cart = corners @ cell + + lo = corners_cart.min(axis=0) + hi = corners_cart.max(axis=0) + + clamped = np.clip(atom_images_cart, lo, hi) + dist_to_box = np.linalg.norm(atom_images_cart - clamped, axis=1) + atom_images_cart = atom_images_cart[dist_to_box <= r_max] + + total_flat = _accumulate_images( + grid_cart, atom_images_cart, r_valid, d_valid, r_max, r_start, dr + ) + return total_flat.reshape(grid_shape) + + +# -------------------- TASK -------------------- +def process_task( + input_path: Path, + task_id: str, + output_dir: Path, + radial_cache: dict, + zval_dict: dict, + extra_shells: int, +): + out_zarr = output_dir / f"{task_id}.zarr" + + cd = load_chgcar(input_path) + + structure = cd.structure + grid_size = cd.data["total"].shape + cell_matrix = structure.lattice.matrix + + chg_label = np.array(cd.data["total"], dtype=np.float64) + + if out_zarr.exists(): + shutil.rmtree(out_zarr) + + new_grid = np.zeros(grid_size, dtype=np.float64) + grid_cart = build_voxel_grid(cell_matrix, grid_size) + voxel_volume = structure.lattice.volume / np.prod(grid_size) + + for site in structure: + specie = str(site.specie) + radial_r, radial_density = radial_cache[specie] + + atom_contribution = reconstruct_from_radial_function( + radial_r, + radial_density, + ref_point=site.frac_coords, + cell=cell_matrix, + grid_shape=grid_size, + extra_shells=extra_shells, + prebuilt_grid=grid_cart, + ) + nelecs = zval_dict[specie] + t = (nelecs / voxel_volume) / np.sum(atom_contribution) + atom_contribution *= t + new_grid += atom_contribution + + new_grid *= structure.lattice.volume + cd.data["total"] = new_grid + + write_chgcar_to_zarr(cd, out_zarr, write_diff=False, chunks=False) + + n_vox = chg_label.size + den = np.sum(chg_label) + nmae = float(np.sum(np.abs(chg_label - new_grid)) / den) if den != 0 else np.nan + nelec_diff = float((np.sum(new_grid) - np.sum(chg_label)) / n_vox) + + return nmae, nelec_diff + + +# -------------------- MAIN -------------------- +def main(): + args = parse_args() + + with args.config.open() as f: + cfg = yaml.safe_load(f) + + paths = cfg["paths"] + input_path = ( + Path(args.input) if args.input is not None else Path(paths["input_file"]) + ) + output_dir = Path(paths["output_dir"]) + stats_dir = Path(paths["stats_dir"]) + radial_dir = Path(paths["radial_dir"]) + zval_path = Path(paths["zval_file"]) + extra_shells = cfg["pads"]["extra_shells"] + + task_id = _derive_task_id(input_path) + + output_dir.mkdir(parents=True, exist_ok=True) + stats_dir.mkdir(parents=True, exist_ok=True) + + radial_cache: dict[str, tuple[np.ndarray, np.ndarray]] = {} + for fp in radial_dir.glob("*.npy"): + specie = fp.stem + arr = np.load(fp) + r, d = arr[:, 0].astype(np.float32), arr[:, 1].astype(np.float32) + valid = np.isfinite(r) & np.isfinite(d) + r, d = r[valid], d[valid] + order = np.argsort(r) + r, d = r[order], d[order] + # NaN removal leaves non-uniform spacing; resample to uniform grid + # so the O(1) index in _accumulate_images stays correct. + dr_raw = np.float32(0.0025) + n_uniform = round((r[-1] - r[0]) / dr_raw) + 1 + r_uniform = np.linspace(r[0], r[-1], n_uniform, dtype=np.float32) + d_uniform = np.interp(r_uniform, r, d).astype(np.float32) + radial_cache[specie] = (r_uniform, d_uniform) + + with zval_path.open() as f: + zval_dict: dict[str, float] = json.load(f) + + nmae, nelec_diff = process_task( + input_path, task_id, output_dir, radial_cache, zval_dict, extra_shells + ) + + print(f"{task_id} nmae={nmae:.6e} nelec_diff={nelec_diff:.6e}", flush=True) # noqa: T201 + + result_path = stats_dir / f"{task_id}.json" + with result_path.open("w") as f: + json.dump({"nmae": nmae, "nelec_diff": nelec_diff}, f) + + +if __name__ == "__main__": + main() diff --git a/uv.lock b/uv.lock index fbe3c234..9984001e 100644 --- a/uv.lock +++ b/uv.lock @@ -170,6 +170,12 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/78/b6/6307fbef88d9b5ee7421e68d78a9f162e0da4900bc5f5793f6d3d0e34fb8/annotated_types-0.7.0-py3-none-any.whl", hash = "sha256:1f02e8b43a8fbbc3f3e0d4f0f4bfc8131bcb4eebe8849b8e5c773f3a1c582a53", size = 13643, upload-time = "2024-05-20T21:33:24.1Z" }, ] +[[package]] +name = "antlr4-python3-runtime" +version = "4.9.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/3e/38/7859ff46355f76f8d19459005ca000b6e7012f2f1ca597746cbcd1fbfe5e/antlr4-python3-runtime-4.9.3.tar.gz", hash = "sha256:f224469b4168294902bb1efa80a8bf7855f24c99aef99cbefc1bcd3cce77881b", size = 117034, upload-time = "2021-11-06T17:52:23.524Z" } + [[package]] name = "attrs" version = "25.4.0" @@ -715,7 +721,9 @@ name = "electrai" version = "0.0.1" source = { editable = "." } dependencies = [ + { name = "hydra-core" }, { name = "lightning" }, + { name = "numba" }, { name = "numpy" }, { name = "pymatgen" }, { name = "pyyaml" }, @@ -759,11 +767,13 @@ dev = [ requires-dist = [ { name = "cairosvg", marker = "extra == 'docs'", specifier = ">=2.7.1" }, { name = "fire", marker = "extra == 'zarr-conversion'", specifier = ">=0.5.0" }, + { name = "hydra-core", specifier = ">=1.3.2" }, { name = "lightning", specifier = "~=2.5.6" }, { name = "mkdocs-gen-files", marker = "extra == 'docs'", specifier = ">=0.5.0" }, { name = "mkdocs-literate-nav", marker = "extra == 'docs'", specifier = ">=0.6.0" }, { name = "mkdocs-material", marker = "extra == 'docs'", specifier = ">=9.4.0" }, { name = "mkdocstrings", extras = ["python"], marker = "extra == 'docs'", specifier = ">=0.22.0" }, + { name = "numba", specifier = ">=0.65.1" }, { name = "numcodecs", marker = "extra == 'zarr-conversion'", specifier = ">=0.16.3" }, { name = "numpy", specifier = "~=2.3.3" }, { name = "pillow", marker = "extra == 'docs'", specifier = ">=10.0.0" }, @@ -1027,6 +1037,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/2a/b1/9ff6578d789a89812ff21e4e0f80ffae20a65d5dd84e7a17873fe3b365be/griffe-1.14.0-py3-none-any.whl", hash = "sha256:0e9d52832cccf0f7188cfe585ba962d2674b241c01916d780925df34873bceb0", size = 144439, upload-time = "2025-09-05T15:02:27.511Z" }, ] +[[package]] +name = "hydra-core" +version = "1.3.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "antlr4-python3-runtime" }, + { name = "omegaconf" }, + { name = "packaging" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/6d/8e/07e42bc434a847154083b315779b0a81d567154504624e181caf2c71cd98/hydra-core-1.3.2.tar.gz", hash = "sha256:8a878ed67216997c3e9d88a8e72e7b4767e81af37afb4ea3334b269a4390a824", size = 3263494, upload-time = "2023-02-23T18:33:43.03Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c6/50/e0edd38dcd63fb26a8547f13d28f7a008bc4a3fd4eb4ff030673f22ad41a/hydra_core-1.3.2-py3-none-any.whl", hash = "sha256:fa0238a9e31df3373b35b0bfb672c34cc92718d21f81311d8996a16de1141d8b", size = 154547, upload-time = "2023-02-23T18:33:40.801Z" }, +] + [[package]] name = "idna" version = "3.11" @@ -1199,6 +1223,34 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/de/73/3d757cb3fc16f0f9794dd289bcd0c4a031d9cf54d8137d6b984b2d02edf3/lightning_utilities-0.15.2-py3-none-any.whl", hash = "sha256:ad3ab1703775044bbf880dbf7ddaaac899396c96315f3aa1779cec9d618a9841", size = 29431, upload-time = "2025-08-06T13:57:38.046Z" }, ] +[[package]] +name = "llvmlite" +version = "0.47.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/01/88/a8952b6d5c21e74cbf158515b779666f692846502623e9e3c39d8e8ba25f/llvmlite-0.47.0.tar.gz", hash = "sha256:62031ce968ec74e95092184d4b0e857e444f8fdff0b8f9213707699570c33ccc", size = 193614, upload-time = "2026-03-31T18:29:53.497Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/34/0b/b9d1911cfefa61399821dfb37f486d83e0f42630a8d12f7194270c417002/llvmlite-0.47.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:74090f0dcfd6f24ebbef3f21f11e38111c4d7e6919b54c4416e1e357c3446b07", size = 37232770, upload-time = "2026-03-31T18:28:26.765Z" }, + { url = "https://files.pythonhosted.org/packages/46/27/5799b020e4cdfb25a7c951c06a96397c135efcdc21b78d853bbd9c814c7d/llvmlite-0.47.0-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:ca14f02e29134e837982497959a8e2193d6035235de1cb41a9cb2bd6da4eedbb", size = 56275177, upload-time = "2026-03-31T18:28:31.01Z" }, + { url = "https://files.pythonhosted.org/packages/7e/51/48a53fedf01cb1f3f43ef200be17ebf83c8d9a04018d3783c1a226c342c2/llvmlite-0.47.0-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:12a69d4bb05f402f30477e21eeabe81911e7c251cecb192bed82cd83c9db10d8", size = 55128631, upload-time = "2026-03-31T18:28:36.046Z" }, + { url = "https://files.pythonhosted.org/packages/a2/50/59227d06bdc96e23322713c381af4e77420949d8cd8a042c79e0043096cc/llvmlite-0.47.0-cp311-cp311-win_amd64.whl", hash = "sha256:c37d6eb7aaabfa83ab9c2ff5b5cdb95a5e6830403937b2c588b7490724e05327", size = 38138400, upload-time = "2026-03-31T18:28:40.076Z" }, + { url = "https://files.pythonhosted.org/packages/fa/48/4b7fe0e34c169fa2f12532916133e0b219d2823b540733651b34fdac509a/llvmlite-0.47.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:306a265f408c259067257a732c8e159284334018b4083a9e35f67d19792b164f", size = 37232769, upload-time = "2026-03-31T18:28:43.735Z" }, + { url = "https://files.pythonhosted.org/packages/e6/4b/e3f2cd17822cf772a4a51a0a8080b0032e6d37b2dbe8cfb724eac4e31c52/llvmlite-0.47.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:5853bf26160857c0c2573415ff4efe01c4c651e59e2c55c2a088740acfee51cd", size = 56275178, upload-time = "2026-03-31T18:28:48.342Z" }, + { url = "https://files.pythonhosted.org/packages/b6/55/a3b4a543185305a9bdf3d9759d53646ed96e55e7dfd43f53e7a421b8fbae/llvmlite-0.47.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:003bcf7fa579e14db59c1a1e113f93ab8a06b56a4be31c7f08264d1d4072d077", size = 55128632, upload-time = "2026-03-31T18:28:52.901Z" }, + { url = "https://files.pythonhosted.org/packages/2f/f5/d281ae0f79378a5a91f308ea9fdb9f9cc068fddd09629edc0725a5a8fde1/llvmlite-0.47.0-cp312-cp312-win_amd64.whl", hash = "sha256:f3079f25bdc24cd9d27c4b2b5e68f5f60c4fdb7e8ad5ee2b9b006007558f9df7", size = 38138692, upload-time = "2026-03-31T18:28:57.147Z" }, + { url = "https://files.pythonhosted.org/packages/77/6f/4615353e016799f80fa52ccb270a843c413b22361fadda2589b2922fb9b0/llvmlite-0.47.0-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:a3c6a735d4e1041808434f9d440faa3d78d9b4af2ee64d05a66f351883b6ceec", size = 37232771, upload-time = "2026-03-31T18:29:01.324Z" }, + { url = "https://files.pythonhosted.org/packages/31/b8/69f5565f1a280d032525878a86511eebed0645818492feeb169dfb20ae8e/llvmlite-0.47.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:2699a74321189e812d476a43d6d7f652f51811e7b5aad9d9bba842a1c7927acb", size = 56275178, upload-time = "2026-03-31T18:29:05.748Z" }, + { url = "https://files.pythonhosted.org/packages/d6/da/b32cafcb926fb0ce2aa25553bf32cb8764af31438f40e2481df08884c947/llvmlite-0.47.0-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6c6951e2b29930227963e53ee152441f0e14be92e9d4231852102d986c761e40", size = 55128632, upload-time = "2026-03-31T18:29:11.235Z" }, + { url = "https://files.pythonhosted.org/packages/46/9f/4898b44e4042c60fafcb1162dfb7014f6f15b1ec19bf29cfea6bf26df90d/llvmlite-0.47.0-cp313-cp313-win_amd64.whl", hash = "sha256:c2e9adf8698d813a9a5efb2d4370caf344dbc1e145019851fee6a6f319ba760e", size = 38138695, upload-time = "2026-03-31T18:29:15.43Z" }, + { url = "https://files.pythonhosted.org/packages/1c/d4/33c8af00f0bf6f552d74f3a054f648af2c5bc6bece97972f3bfadce4f5ec/llvmlite-0.47.0-cp314-cp314-macosx_12_0_arm64.whl", hash = "sha256:de966c626c35c9dff5ae7bf12db25637738d0df83fc370cf793bc94d43d92d14", size = 37232773, upload-time = "2026-03-31T18:29:19.453Z" }, + { url = "https://files.pythonhosted.org/packages/64/1d/a760e993e0c0ba6db38d46b9f48f6c7dceb8ac838824997fb9e25f97bc04/llvmlite-0.47.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:ddbccff2aeaff8670368340a158abefc032fe9b3ccf7d9c496639263d00151aa", size = 56275176, upload-time = "2026-03-31T18:29:24.149Z" }, + { url = "https://files.pythonhosted.org/packages/84/3b/e679bc3b29127182a7f4aa2d2e9e5bea42adb93fb840484147d59c236299/llvmlite-0.47.0-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d4a7b778a2e144fc64468fb9bf509ac1226c9813a00b4d7afea5d988c4e22fca", size = 55128631, upload-time = "2026-03-31T18:29:29.536Z" }, + { url = "https://files.pythonhosted.org/packages/be/f7/19e2a09c62809c9e63bbd14ce71fb92c6ff7b7b3045741bb00c781efc3c9/llvmlite-0.47.0-cp314-cp314-win_amd64.whl", hash = "sha256:694e3c2cdc472ed2bd8bd4555ca002eec4310961dd58ef791d508f57b5cc4c94", size = 39153826, upload-time = "2026-03-31T18:29:33.681Z" }, + { url = "https://files.pythonhosted.org/packages/40/a1/581a8c707b5e80efdbbe1dd94527404d33fe50bceb71f39d5a7e11bd57b7/llvmlite-0.47.0-cp314-cp314t-macosx_12_0_arm64.whl", hash = "sha256:92ec8a169a20b473c1c54d4695e371bde36489fc1efa3688e11e99beba0abf9c", size = 37232772, upload-time = "2026-03-31T18:29:37.952Z" }, + { url = "https://files.pythonhosted.org/packages/11/03/16090dd6f74ba2b8b922276047f15962fbeea0a75d5601607edb301ba945/llvmlite-0.47.0-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:fa1cbd800edd3b20bc141521f7fd45a6185a5b84109aa6855134e81397ffe72b", size = 56275178, upload-time = "2026-03-31T18:29:42.58Z" }, + { url = "https://files.pythonhosted.org/packages/f5/cb/0abf1dd4c5286a95ffe0c1d8c67aec06b515894a0dd2ac97f5e27b82ab0b/llvmlite-0.47.0-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f6725179b89f03b17dabe236ff3422cb8291b4c1bf40af152826dfd34e350ae8", size = 55128632, upload-time = "2026-03-31T18:29:46.939Z" }, + { url = "https://files.pythonhosted.org/packages/4f/79/d3bbab197e86e0ff4f9c07122895b66a3e0d024247fcff7f12c473cb36d9/llvmlite-0.47.0-cp314-cp314t-win_amd64.whl", hash = "sha256:6842cf6f707ec4be3d985a385ad03f72b2d724439e118fcbe99b2929964f0453", size = 39153839, upload-time = "2026-03-31T18:29:51.004Z" }, +] + [[package]] name = "markdown" version = "3.9" @@ -1655,6 +1707,38 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/eb/8d/776adee7bbf76365fdd7f2552710282c79a4ead5d2a46408c9043a2b70ba/networkx-3.5-py3-none-any.whl", hash = "sha256:0030d386a9a06dee3565298b4a734b68589749a544acbb6c412dc9e2489ec6ec", size = 2034406, upload-time = "2025-05-29T11:35:04.961Z" }, ] +[[package]] +name = "numba" +version = "0.65.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "llvmlite" }, + { name = "numpy" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/f6/c5/db2ac3685833d626c0dcae6bd2330cd68433e1fd248d15f70998160d3ad7/numba-0.65.1.tar.gz", hash = "sha256:19357146c32fe9ed25059ab915e8465fb13951cf6b0aace3826b76886373ab23", size = 2765600, upload-time = "2026-04-24T02:02:56.551Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/96/b3/650500c2eab4534d98e9166f4298e0f3c69c742afdf24e6eabccd1f16ad8/numba-0.65.1-cp311-cp311-macosx_12_0_arm64.whl", hash = "sha256:7020d74b19cdb8cff16506542fdd510756e28c5e7f3bd0b7f574f0f42272fcd9", size = 2680563, upload-time = "2026-04-24T02:02:18.414Z" }, + { url = "https://files.pythonhosted.org/packages/44/0b/0615dbedb98f5b32a35a53290fbdc6e22306968109278d7e58df82d7a9f6/numba-0.65.1-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:f80ed83774b5173abd6581cd8d2165d1d38e13d2e5c8155c0c0b421784745420", size = 3745018, upload-time = "2026-04-24T02:02:20.252Z" }, + { url = "https://files.pythonhosted.org/packages/49/aa/4361698f35bf63bff67dfe6c90493731177f48ede954f77b0588731537bc/numba-0.65.1-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:7ed425a43b0a5f9772f2f4e2dd0bbd12eabecae1af0b24efcfd4e053f012aac6", size = 3450962, upload-time = "2026-04-24T02:02:22.449Z" }, + { url = "https://files.pythonhosted.org/packages/bd/9a/af61ec03b3116c161fd7a06b9e8a265729a8718458333e8ffbb06d9a3978/numba-0.65.1-cp311-cp311-win_amd64.whl", hash = "sha256:df40a5028a975b9ea66f6a2a3f7abbdbd541a863070e34ed367aff21141248e4", size = 2747417, upload-time = "2026-04-24T02:02:24.43Z" }, + { url = "https://files.pythonhosted.org/packages/57/bc/76f8f8c5cf9adee47fdb7bbb03be8900f76f902d451d7477cf12b845e1de/numba-0.65.1-cp312-cp312-macosx_12_0_arm64.whl", hash = "sha256:ac3f1e77c352dd0ea9712732c2d8f9ca507717435eec5b5013bf138ac33c4a08", size = 2681371, upload-time = "2026-04-24T02:02:26.105Z" }, + { url = "https://files.pythonhosted.org/packages/69/47/a415af0283e4db0398104c6d1c11c9861a98dc67a7aa442a7769ed5d6196/numba-0.65.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:52bc6f3ceb8fcaff9b2ae26b4c6b1e9fee39db8d355534c0fe4f39a901246b84", size = 3802467, upload-time = "2026-04-24T02:02:27.712Z" }, + { url = "https://files.pythonhosted.org/packages/46/36/246f73ec99cfeab2f2cb2ce7d4218766cc36a2da418901223f4f4da9c813/numba-0.65.1-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:90ca10b3463bae0bd70589726fe3c77d01d6b5fc86bee54bcdf9fb6b47c28977", size = 3502628, upload-time = "2026-04-24T02:02:29.763Z" }, + { url = "https://files.pythonhosted.org/packages/db/9e/3c679b2ee078425b9e99a91e44f8d132a6830d8ccce5227bc5e9181aeed8/numba-0.65.1-cp312-cp312-win_amd64.whl", hash = "sha256:5971c632be2a2351500431f46213821dba8d02b18a9f7d02fd36bd2743e41a6a", size = 2750611, upload-time = "2026-04-24T02:02:31.477Z" }, + { url = "https://files.pythonhosted.org/packages/79/37/14a4579049c1eb673afd0de0cb4842982acd55b9ce2643e763db858bcea0/numba-0.65.1-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:1735c15c1134a5108b4d6a5c77fc0947924ea066a738dc09a52008c13df9cad3", size = 2681344, upload-time = "2026-04-24T02:02:33.65Z" }, + { url = "https://files.pythonhosted.org/packages/a0/22/b8d873f6466b20aa563fc9b33acd48dec89a07803ddaa2f1c8ca1cd33126/numba-0.65.1-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:c09f49117ef255e1f1c6dad0c7a1ed39868243862a73be5706793241a3755f1b", size = 3810619, upload-time = "2026-04-24T02:02:36.041Z" }, + { url = "https://files.pythonhosted.org/packages/62/08/e16a8b5d9a018962ebb5c66be662317cde32b9f5dab08441f90bed5522fb/numba-0.65.1-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:594a8680b3fadac99e97e489b1fd89007177e5336713745c3b769528c635a464", size = 3509783, upload-time = "2026-04-24T02:02:38.245Z" }, + { url = "https://files.pythonhosted.org/packages/fd/a5/03c970d57f4c1741354837353ce39fb5206952ae1dba8922d29c86f64805/numba-0.65.1-cp313-cp313-win_amd64.whl", hash = "sha256:85be74c0d036842699a30058f82fb88fc5ffdc59f7615cab5792ea92914c9b62", size = 2750534, upload-time = "2026-04-24T02:02:39.903Z" }, + { url = "https://files.pythonhosted.org/packages/4f/2e/8aed9b726d9ba5f11ad287645fd479e88278db3060a25cb1225d730eb2b7/numba-0.65.1-cp314-cp314-macosx_12_0_arm64.whl", hash = "sha256:33f5eb68eb1c843511615d14663ce60258525d6a4c65ab040e2c2b0c4cf17450", size = 2681554, upload-time = "2026-04-24T02:02:41.812Z" }, + { url = "https://files.pythonhosted.org/packages/87/96/f3eb235fafa82a34e2ab5dd7dc9ffff998ebf5f0bbc23fa56a96aeb44da6/numba-0.65.1-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:71e73029bf53a62cc6afcf96be4bd942290d8b4c55f0a454fb536158115790f7", size = 3779602, upload-time = "2026-04-24T02:02:43.726Z" }, + { url = "https://files.pythonhosted.org/packages/09/90/b0f09b48752d23640b8284f22aa597737e8adaddc7fbfacc4708b7f73a4c/numba-0.65.1-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3a07635e0be926b9bdbffb09137c230fb13f6ec0e564914ba937cee12ce3eb35", size = 3479532, upload-time = "2026-04-24T02:02:45.427Z" }, + { url = "https://files.pythonhosted.org/packages/56/46/3f7fc04fb853559e74b210e0b62c19974ec844cefec611f9e535f4da3761/numba-0.65.1-cp314-cp314-win_amd64.whl", hash = "sha256:2a20fcdabdefbdacf88d85caf70c3b18c4bcb7ebb8f82e6a19486383dd26ab63", size = 2752637, upload-time = "2026-04-24T02:02:47.664Z" }, + { url = "https://files.pythonhosted.org/packages/81/7b/c1a341a9067367778f4152a5f01061cf281fb09582c92c510ec4918cabf6/numba-0.65.1-cp314-cp314t-macosx_12_0_arm64.whl", hash = "sha256:548dd4b3a4508d5062768d1514b2cd7b015f9a25ec7af651c50dee243965e652", size = 2684600, upload-time = "2026-04-24T02:02:49.653Z" }, + { url = "https://files.pythonhosted.org/packages/03/36/98ddbcf3e4f04a6dd07e1c67249955920579ba4af6bb6868e3088f4ed282/numba-0.65.1-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:78abc28feff2c2ff8307fff3975b6438352759c9acb797ecd6b1fb6e7e39e31d", size = 3817198, upload-time = "2026-04-24T02:02:51.266Z" }, + { url = "https://files.pythonhosted.org/packages/a3/83/0dad21057ece5a835599f5d24099b091703995e23dbbf894f259e91c010b/numba-0.65.1-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ee7676cb389555805f9b9a1840cbcd1ea6c8bd5376ab6918e3a29c5ea1dbda20", size = 3533862, upload-time = "2026-04-24T02:02:52.987Z" }, + { url = "https://files.pythonhosted.org/packages/32/36/8be7118ffd4c8440881046eac3d0982cc5ab42909508cf5d67024d62a2e4/numba-0.65.1-cp314-cp314t-win_amd64.whl", hash = "sha256:20609346e3bd75204950dcbbfe383a8d7dbf4902f442aedbf00f97fef4aa8f38", size = 2758237, upload-time = "2026-04-24T02:02:54.612Z" }, +] + [[package]] name = "numcodecs" version = "0.16.3" @@ -1902,6 +1986,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a2/eb/86626c1bbc2edb86323022371c39aa48df6fd8b0a1647bc274577f72e90b/nvidia_nvtx_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:5b17e2001cc0d751a5bc2c6ec6d26ad95913324a4adb86788c944f8ce9ba441f", size = 89954, upload-time = "2025-03-07T01:42:44.131Z" }, ] +[[package]] +name = "omegaconf" +version = "2.3.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "antlr4-python3-runtime" }, + { name = "pyyaml" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/09/48/6388f1bb9da707110532cb70ec4d2822858ddfb44f1cdf1233c20a80ea4b/omegaconf-2.3.0.tar.gz", hash = "sha256:d5d4b6d29955cc50ad50c46dc269bcd92c6e00f5f90d23ab5fee7bfca4ba4cc7", size = 3298120, upload-time = "2022-12-08T20:59:22.753Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e3/94/1843518e420fa3ed6919835845df698c7e27e183cb997394e4a670973a65/omegaconf-2.3.0-py3-none-any.whl", hash = "sha256:7b4df175cdb08ba400f45cae3bdcae7ba8365db4d165fc65fd04b050ab63b46b", size = 79500, upload-time = "2022-12-08T20:59:19.686Z" }, +] + [[package]] name = "orjson" version = "3.11.4" From c06eb5c690655730a774a6c034102bcb00d7320a Mon Sep 17 00:00:00 2001 From: Hananeh Oliaei Date: Tue, 19 May 2026 16:11:13 -0400 Subject: [PATCH 2/5] Fix zarr writer passing chunks=False to zarr v3 zarr v3 raises an error when chunks=False is passed to root.create(). Resolve falsy chunks to the full array shape (single chunk) instead. Co-Authored-By: Claude Sonnet 4.6 --- src/electrai/zarr_conversion/zarr_writer.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/src/electrai/zarr_conversion/zarr_writer.py b/src/electrai/zarr_conversion/zarr_writer.py index d16e32c3..7faa3f9f 100644 --- a/src/electrai/zarr_conversion/zarr_writer.py +++ b/src/electrai/zarr_conversion/zarr_writer.py @@ -97,14 +97,19 @@ def write_chgcar_to_zarr( # Store total charge density total_density = np.array(charge_data["total"], dtype=np.float32) - root.create(name="charge_density_total", data=total_density, chunks=chunks) + resolved_chunks = chunks if chunks else total_density.shape + root.create( + name="charge_density_total", data=total_density, chunks=resolved_chunks + ) logger.debug(f"Stored total charge density with shape {total_density.shape}") # Store diff charge density (if present and write_diff is True) diff_density_raw = charge_data.get("diff") if write_diff and diff_density_raw is not None: diff_density = np.array(diff_density_raw, dtype=np.float32) - root.create(name="charge_density_diff", data=diff_density, chunks=chunks) + root.create( + name="charge_density_diff", data=diff_density, chunks=resolved_chunks + ) logger.debug(f"Stored diff charge density with shape {diff_density.shape}") # Store structure information as JSON From 1a5ba67607f8155115732608e9fd99d43a5f3b64 Mon Sep 17 00:00:00 2001 From: Betsy Cannon Date: Wed, 20 May 2026 11:21:42 -0400 Subject: [PATCH 3/5] Add __init__.py to pads_generation package Mirrors the empty __init__.py in zarr_conversion/ so the package structure is consistent. Addresses PR #136 review feedback. Co-Authored-By: Claude Opus 4.7 (1M context) --- src/electrai/pads_generation/__init__.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) create mode 100644 src/electrai/pads_generation/__init__.py diff --git a/src/electrai/pads_generation/__init__.py b/src/electrai/pads_generation/__init__.py new file mode 100644 index 00000000..e69de29b From c98295d18f385fd9ca8bd8e7f5fb9b4946a51bb4 Mon Sep 17 00:00:00 2001 From: Betsy Cannon Date: Wed, 20 May 2026 11:39:19 -0400 Subject: [PATCH 4/5] Use chunks=None instead of chunks=False in create_pads write_chgcar_to_zarr declares chunks: tuple[int, int, int] | None, so passing False was a type mismatch even though it worked at runtime. Addresses PR #136 review feedback. Co-Authored-By: Claude Opus 4.7 (1M context) --- src/electrai/pads_generation/create_pads.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/electrai/pads_generation/create_pads.py b/src/electrai/pads_generation/create_pads.py index 9e551d86..f77712a0 100644 --- a/src/electrai/pads_generation/create_pads.py +++ b/src/electrai/pads_generation/create_pads.py @@ -190,7 +190,7 @@ def process_task( new_grid *= structure.lattice.volume cd.data["total"] = new_grid - write_chgcar_to_zarr(cd, out_zarr, write_diff=False, chunks=False) + write_chgcar_to_zarr(cd, out_zarr, write_diff=False, chunks=None) n_vox = chg_label.size den = np.sum(chg_label) From ca1f6dff34240920fa07c6a3acfdba80b2b7c981 Mon Sep 17 00:00:00 2001 From: Hananeh Oliaei Date: Fri, 22 May 2026 09:55:22 -0400 Subject: [PATCH 5/5] Use site.species_string instead of deprecated site.specie Co-Authored-By: Claude Sonnet 4.6 --- src/electrai/pads_generation/create_pads.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/electrai/pads_generation/create_pads.py b/src/electrai/pads_generation/create_pads.py index f77712a0..ea5fb463 100644 --- a/src/electrai/pads_generation/create_pads.py +++ b/src/electrai/pads_generation/create_pads.py @@ -170,7 +170,7 @@ def process_task( voxel_volume = structure.lattice.volume / np.prod(grid_size) for site in structure: - specie = str(site.specie) + specie = site.species_string radial_r, radial_density = radial_cache[specie] atom_contribution = reconstruct_from_radial_function(