AtomVoxelizer 0.1.0__py3-none-any.whl

atomvoxelizer/__init__.py

@@ -0,0 +1,30 @@
+"""AtomVoxelizer public API."""
+
+from .voxelgrid import VoxelGrid, VoxelGridNumPy
+from .analysis import VoxelGridAnalysis, VoxelRegion
+
+__all__ = [
+    "VoxelGrid",
+    "VoxelGridAnalysis",
+    "VoxelGridCuPy",
+    "VoxelGridNumPy",
+    "VoxelGridNumba",
+    "VoxelGridTaichi",
+    "VoxelRegion",
+]
+
+
+def __getattr__(name):
+    if name == "VoxelGridCuPy":
+        from .cupy_backend import VoxelGridCuPy
+
+        return VoxelGridCuPy
+    if name == "VoxelGridNumba":
+        from .numba_backend import VoxelGridNumba
+
+        return VoxelGridNumba
+    if name == "VoxelGridTaichi":
+        from .taichi_backend import VoxelGridTaichi
+
+        return VoxelGridTaichi
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

atomvoxelizer/analysis.py

@@ -0,0 +1,207 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import numpy as np
+
+
+@dataclass(frozen=True)
+class VoxelRegion:
+    """Summary of one connected voxel region."""
+
+    label: int
+    voxel_count: int
+    volume: float
+    surface_area: float
+
+
+class VoxelGridAnalysis:
+    """Analyze connected voxel volumes and their surfaces."""
+
+    def __init__(self, voxel_grid):
+        self.voxel_grid = voxel_grid
+
+    @property
+    def grid(self):
+        return self.voxel_grid.to_numpy()
+
+    @property
+    def cell(self):
+        return self.voxel_grid.cell
+
+    @property
+    def gpts(self):
+        return self.voxel_grid.gpts
+
+    @property
+    def voxel_volume(self):
+        return abs(float(np.linalg.det(self.cell))) / float(np.prod(self.gpts))
+
+    def mask(self, min_value=None, max_value=None, threshold=None, above=True):
+        """Build a boolean mask from voxel values."""
+        if threshold is not None and (min_value is not None or max_value is not None):
+            raise ValueError("Specify either threshold or min_value/max_value, not both")
+
+        grid = self.grid
+        if threshold is not None:
+            return grid > threshold if above else grid < threshold
+
+        selected = np.ones(grid.shape, dtype=bool)
+        if min_value is not None:
+            selected &= grid >= min_value
+        if max_value is not None:
+            selected &= grid <= max_value
+        return selected
+
+    def connected_components(self, selected, connectivity=1, periodic=True):
+        """Label connected components in a boolean mask."""
+        try:
+            from skimage.measure import label
+        except ImportError as exc:  # pragma: no cover - depends on optional dependency
+            raise ImportError(
+                "VoxelGridAnalysis requires scikit-image. Install it with "
+                "`pip install AtomVoxelizer[analysis]`."
+            ) from exc
+
+        selected = np.asarray(selected, dtype=bool)
+        labels = label(selected, connectivity=connectivity)
+        if periodic:
+            labels = self._merge_periodic_labels(labels, selected)
+        return labels, int(labels.max())
+
+    def region_volume(self, selected):
+        """Return the volume represented by a boolean mask."""
+        return int(np.count_nonzero(selected)) * self.voxel_volume
+
+    def surface_area(self, selected, periodic=True):
+        """Estimate the surface area of a boolean region with marching cubes."""
+        try:
+            from skimage.measure import marching_cubes
+        except ImportError as exc:  # pragma: no cover - depends on optional dependency
+            raise ImportError(
+                "VoxelGridAnalysis requires scikit-image. Install it with "
+                "`pip install AtomVoxelizer[analysis]`."
+            ) from exc
+
+        selected = np.asarray(selected, dtype=bool)
+        if not np.any(selected):
+            return 0.0
+        if periodic and np.all(selected):
+            return 0.0
+
+        if periodic:
+            values = np.tile(selected.astype(np.float32), (3, 3, 3))
+            vertices, faces, _normals, _values = marching_cubes(values, level=0.5)
+            central_offset = self.gpts
+            centroids = vertices[faces].mean(axis=1)
+            in_central_cell = np.all((centroids >= central_offset) & (centroids < 2 * central_offset), axis=1)
+            faces = faces[in_central_cell]
+            vertices = vertices - central_offset
+        else:
+            values = np.pad(selected.astype(np.float32), 1, mode="constant", constant_values=0.0)
+            vertices, faces, _normals, _values = marching_cubes(values, level=0.5)
+            vertices = vertices - 1.0
+
+        if faces.size == 0:
+            return 0.0
+
+        real_vertices = self._index_vertices_to_real(vertices)
+        return self._mesh_surface_area(real_vertices, faces)
+
+    def analyze_regions(
+        self,
+        min_value=None,
+        max_value=None,
+        threshold=None,
+        above=True,
+        connectivity=1,
+        periodic=True,
+    ):
+        """Return volume and marching-cubes area for each connected region."""
+        selected = self.mask(min_value=min_value, max_value=max_value, threshold=threshold, above=above)
+        labels, label_count = self.connected_components(selected, connectivity=connectivity, periodic=periodic)
+
+        regions = []
+        for label_id in range(1, label_count + 1):
+            region_mask = labels == label_id
+            voxel_count = int(np.count_nonzero(region_mask))
+            regions.append(
+                VoxelRegion(
+                    label=label_id,
+                    voxel_count=voxel_count,
+                    volume=voxel_count * self.voxel_volume,
+                    surface_area=self.surface_area(region_mask, periodic=periodic),
+                )
+            )
+        return regions
+
+    @staticmethod
+    def volume_angstrom3_to_cm3_per_g(volume_angstrom3, mass_amu):
+        """Convert a cell/supercell volume from Angstrom^3 to cm^3/g."""
+        if mass_amu <= 0:
+            raise ValueError("mass_amu must be positive")
+        mass_g = float(mass_amu) * 1.66053906660e-24
+        return float(volume_angstrom3) * 1.0e-24 / mass_g
+
+    @staticmethod
+    def area_angstrom2_to_m2_per_g(area_angstrom2, mass_amu):
+        """Convert a cell/supercell area from Angstrom^2 to m^2/g."""
+        if mass_amu <= 0:
+            raise ValueError("mass_amu must be positive")
+        mass_g = float(mass_amu) * 1.66053906660e-24
+        return float(area_angstrom2) * 1.0e-20 / mass_g
+
+    def _index_vertices_to_real(self, vertices):
+        frac = vertices / self.gpts
+        return frac @ self.cell
+
+    @staticmethod
+    def _merge_periodic_labels(labels, selected):
+        label_count = int(labels.max())
+        if label_count == 0:
+            return labels
+
+        parent = np.arange(label_count + 1)
+
+        def find(label_id):
+            while parent[label_id] != label_id:
+                parent[label_id] = parent[parent[label_id]]
+                label_id = parent[label_id]
+            return label_id
+
+        def union(a, b):
+            if a == 0 or b == 0:
+                return
+            root_a = find(int(a))
+            root_b = find(int(b))
+            if root_a != root_b:
+                parent[root_b] = root_a
+
+        for axis in range(labels.ndim):
+            first_labels = np.take(labels, 0, axis=axis)
+            last_labels = np.take(labels, -1, axis=axis)
+            first_selected = np.take(selected, 0, axis=axis)
+            last_selected = np.take(selected, -1, axis=axis)
+            for a, b in zip(first_labels[first_selected & last_selected], last_labels[first_selected & last_selected]):
+                union(a, b)
+
+        root_to_new_label = {}
+        next_label = 1
+        merged = np.zeros_like(labels)
+        for label_id in range(1, label_count + 1):
+            root = find(label_id)
+            if root not in root_to_new_label:
+                root_to_new_label[root] = next_label
+                next_label += 1
+            merged[labels == label_id] = root_to_new_label[root]
+
+        return merged
+
+    @staticmethod
+    def _mesh_surface_area(vertices, faces):
+        triangles = vertices[faces]
+        cross = np.cross(triangles[:, 1] - triangles[:, 0], triangles[:, 2] - triangles[:, 0])
+        return float(0.5 * np.linalg.norm(cross, axis=1).sum())
+
+
+__all__ = ["VoxelGridAnalysis", "VoxelRegion"]

atomvoxelizer/cupy_backend.py

@@ -0,0 +1,88 @@
+from __future__ import annotations
+
+import numpy as np
+
+from .voxelgrid import VoxelGrid, _cached_sphere_offsets
+
+try:
+    from .numba_backend import VoxelGridNumba as _BaseVoxelGrid
+except ImportError:  # pragma: no cover - depends on optional dependency
+    _BaseVoxelGrid = VoxelGrid
+
+try:
+    import cupy as cp
+except ImportError as exc:  # pragma: no cover - depends on optional dependency
+    raise ImportError(
+        "VoxelGridCuPy requires CuPy. Install it with `pip install AtomVoxelizer[cupy]`."
+    ) from exc
+
+
+class VoxelGridCuPy(_BaseVoxelGrid):
+    """CuPy-backed voxel grid.
+
+    The class keeps the same public API as :class:`atomvoxelizer.VoxelGrid`, uses
+    the Numba backend as its base when available, stores ``grid`` as a CuPy array,
+    and overrides the mutating sphere operations.
+    """
+
+    @property
+    def backend_name(self):
+        return "cupy"
+
+    def __init__(self, cell, resolution=None, gpts=None):
+        super().__init__(cell=cell, resolution=resolution, gpts=gpts)
+        self.grid = cp.asarray(self.grid)
+
+    def to_numpy(self):
+        """Return the voxel values as a NumPy array."""
+        return cp.asnumpy(self.grid)
+
+    def _sphere_indices(self, center, radius):
+        center_frac = np.asarray(center, dtype=np.float64) @ self.cell_inv % 1.0
+        center_idx = np.floor(center_frac * self.gpts).astype(np.int32)
+        offsets = _cached_sphere_offsets(float(radius), tuple(self.gpts), tuple(map(tuple, self.cell)))
+        indices = (offsets + center_idx) % self.gpts
+        return tuple(cp.asarray(indices[:, axis]) for axis in range(3))
+
+    def set_sphere(self, center, radius, value=1):
+        self.grid[self._sphere_indices(center, radius)] = value
+
+    def add_sphere(self, center, radius, value=1):
+        cp.add.at(self.grid, self._sphere_indices(center, radius), value)
+
+    def mul_sphere(self, center, radius, factor=2):
+        indices = self._sphere_indices(center, radius)
+        self.grid[indices] *= factor
+
+    def div_sphere(self, center, radius, factor=2):
+        indices = self._sphere_indices(center, radius)
+        self.grid[indices] /= factor
+
+    def add_spheres(self, centers, radii, value=1):
+        centers = np.asarray(centers, dtype=np.float64)
+        radii = np.asarray(radii, dtype=np.float64)
+        if centers.ndim != 2 or centers.shape[1] != 3:
+            raise ValueError("centers must have shape (N, 3)")
+        if radii.ndim != 1 or radii.shape[0] != centers.shape[0]:
+            raise ValueError("radii must have shape (N,)")
+
+        for center, radius in zip(centers, radii):
+            self.add_sphere(center, float(radius), value=value)
+
+    def set_spheres(self, centers, radii, value=1):
+        centers = np.asarray(centers, dtype=np.float64)
+        radii = np.asarray(radii, dtype=np.float64)
+        if centers.ndim != 2 or centers.shape[1] != 3:
+            raise ValueError("centers must have shape (N, 3)")
+        if radii.ndim != 1 or radii.shape[0] != centers.shape[0]:
+            raise ValueError("radii must have shape (N,)")
+
+        for center, radius in zip(centers, radii):
+            self.set_sphere(center, float(radius), value=value)
+
+    def clamp_grid(self, min_val=0.0, max_val=1.0):
+        self.grid = cp.clip(self.grid, min_val, max_val)
+
+    def synchronize(self):
+        """Synchronize the current CuPy device."""
+        cp.cuda.Stream.null.synchronize()

atomvoxelizer/numba_backend.py

@@ -0,0 +1,135 @@
+from __future__ import annotations
+
+import numpy as np
+from numba import njit, prange
+
+from .voxelgrid import VoxelGrid, _cached_sphere_offsets
+
+
+@njit
+def _set_sphere_offsets(grid, center_idx, offsets, value):
+    nx, ny, nz = grid.shape
+    for n in range(offsets.shape[0]):
+        x = (center_idx[0] + offsets[n, 0]) % nx
+        y = (center_idx[1] + offsets[n, 1]) % ny
+        z = (center_idx[2] + offsets[n, 2]) % nz
+        grid[x, y, z] = value
+
+
+@njit
+def _add_sphere_offsets(grid, center_idx, offsets, value):
+    nx, ny, nz = grid.shape
+    for n in range(offsets.shape[0]):
+        x = (center_idx[0] + offsets[n, 0]) % nx
+        y = (center_idx[1] + offsets[n, 1]) % ny
+        z = (center_idx[2] + offsets[n, 2]) % nz
+        grid[x, y, z] += value
+
+
+@njit
+def _mul_sphere_offsets(grid, center_idx, offsets, factor):
+    nx, ny, nz = grid.shape
+    for n in range(offsets.shape[0]):
+        x = (center_idx[0] + offsets[n, 0]) % nx
+        y = (center_idx[1] + offsets[n, 1]) % ny
+        z = (center_idx[2] + offsets[n, 2]) % nz
+        grid[x, y, z] *= factor
+
+
+@njit
+def _div_sphere_offsets(grid, center_idx, offsets, divisor):
+    nx, ny, nz = grid.shape
+    for n in range(offsets.shape[0]):
+        x = (center_idx[0] + offsets[n, 0]) % nx
+        y = (center_idx[1] + offsets[n, 1]) % ny
+        z = (center_idx[2] + offsets[n, 2]) % nz
+        grid[x, y, z] /= divisor
+
+
+@njit
+def _set_many_sphere_offsets(grid, center_indices, offsets, value):
+    nx, ny, nz = grid.shape
+    for c in range(center_indices.shape[0]):
+        cx = center_indices[c, 0]
+        cy = center_indices[c, 1]
+        cz = center_indices[c, 2]
+        for n in range(offsets.shape[0]):
+            x = (cx + offsets[n, 0]) % nx
+            y = (cy + offsets[n, 1]) % ny
+            z = (cz + offsets[n, 2]) % nz
+            grid[x, y, z] = value
+
+
+@njit
+def _add_many_sphere_offsets(grid, center_indices, offsets, value):
+    nx, ny, nz = grid.shape
+    for c in range(center_indices.shape[0]):
+        cx = center_indices[c, 0]
+        cy = center_indices[c, 1]
+        cz = center_indices[c, 2]
+        for n in range(offsets.shape[0]):
+            x = (cx + offsets[n, 0]) % nx
+            y = (cy + offsets[n, 1]) % ny
+            z = (cz + offsets[n, 2]) % nz
+            grid[x, y, z] += value
+
+
+@njit(parallel=True)
+def _clamp_grid(grid, min_val, max_val):
+    nx, ny, nz = grid.shape
+    for i in prange(nx):
+        for j in range(ny):
+            for k in range(nz):
+                v = grid[i, j, k]
+                if v < min_val:
+                    grid[i, j, k] = min_val
+                elif v > max_val:
+                    grid[i, j, k] = max_val
+
+
+class VoxelGridNumba(VoxelGrid):
+    """Voxel grid with Numba-compiled mutation kernels."""
+
+    backend_name = "numba"
+
+    def set_sphere(self, center, radius, value=1):
+        center_idx = self._center_index(center)
+        offsets = self._sphere_offsets(radius)
+        _set_sphere_offsets(self.grid, center_idx, offsets, value)
+
+    def add_sphere(self, center, radius, value=1):
+        center_idx = self._center_index(center)
+        offsets = self._sphere_offsets(radius)
+        _add_sphere_offsets(self.grid, center_idx, offsets, value)
+
+    def mul_sphere(self, center, radius, factor=2):
+        center_idx = self._center_index(center)
+        offsets = self._sphere_offsets(radius)
+        _mul_sphere_offsets(self.grid, center_idx, offsets, factor)
+
+    def div_sphere(self, center, radius, factor=2):
+        center_idx = self._center_index(center)
+        offsets = self._sphere_offsets(radius)
+        _div_sphere_offsets(self.grid, center_idx, offsets, factor)
+
+    def add_spheres(self, centers, radii, value=1):
+        centers, radii = self._validate_spheres(centers, radii)
+        center_indices = self.positions_to_indices(centers)
+        for radius in np.unique(radii):
+            offsets = _cached_sphere_offsets(float(radius), tuple(self.gpts), tuple(map(tuple, self.cell)))
+            subset = center_indices[radii == radius]
+            _add_many_sphere_offsets(self.grid, subset, offsets, value)
+
+    def set_spheres(self, centers, radii, value=1):
+        centers, radii = self._validate_spheres(centers, radii)
+        center_indices = self.positions_to_indices(centers)
+        for radius in np.unique(radii):
+            offsets = _cached_sphere_offsets(float(radius), tuple(self.gpts), tuple(map(tuple, self.cell)))
+            subset = center_indices[radii == radius]
+            _set_many_sphere_offsets(self.grid, subset, offsets, value)
+
+    def clamp_grid(self, min_val=0.0, max_val=1.0):
+        _clamp_grid(self.grid, min_val, max_val)
+
+
+__all__ = ["VoxelGridNumba"]

atomvoxelizer/taichi_backend.py

@@ -0,0 +1,157 @@
+import os
+from pathlib import Path
+
+import numpy as np
+
+from .voxelgrid import VoxelGrid, _cached_sphere_offsets
+
+_default_cache = Path.home() / ".cache"
+if not os.access(_default_cache, os.W_OK):
+    os.environ.setdefault("XDG_CACHE_HOME", "/tmp/.cache")
+    os.environ.setdefault("TI_CACHE_HOME", "/tmp/taichi-cache")
+
+try:
+    import taichi as ti
+except ImportError as exc:  # pragma: no cover - depends on optional dependency
+    raise ImportError(
+        "VoxelGridTaichi requires Taichi. Install it with `pip install AtomVoxelizer[taichi]`."
+    ) from exc
+
+
+def _init_taichi_cpu_once():
+    try:
+        runtime = ti.lang.impl.get_runtime()
+        if runtime.prog is not None:
+            return
+    except Exception:
+        pass
+    ti.init(arch=ti.cpu, offline_cache=False)
+
+
+_init_taichi_cpu_once()
+
+
+@ti.kernel
+def _set_sphere_offsets(
+    grid: ti.template(),
+    center_idx: ti.types.ndarray(dtype=ti.i32, ndim=1),
+    offsets: ti.types.ndarray(dtype=ti.i32, ndim=2),
+    value: ti.f32,
+):
+    for n in range(offsets.shape[0]):
+        x = (center_idx[0] + offsets[n, 0]) % grid.shape[0]
+        y = (center_idx[1] + offsets[n, 1]) % grid.shape[1]
+        z = (center_idx[2] + offsets[n, 2]) % grid.shape[2]
+        grid[x, y, z] = value
+
+
+@ti.kernel
+def _add_sphere_offsets(
+    grid: ti.template(),
+    center_idx: ti.types.ndarray(dtype=ti.i32, ndim=1),
+    offsets: ti.types.ndarray(dtype=ti.i32, ndim=2),
+    value: ti.f32,
+):
+    for n in range(offsets.shape[0]):
+        x = (center_idx[0] + offsets[n, 0]) % grid.shape[0]
+        y = (center_idx[1] + offsets[n, 1]) % grid.shape[1]
+        z = (center_idx[2] + offsets[n, 2]) % grid.shape[2]
+        grid[x, y, z] += value
+
+
+@ti.kernel
+def _mul_sphere_offsets(
+    grid: ti.template(),
+    center_idx: ti.types.ndarray(dtype=ti.i32, ndim=1),
+    offsets: ti.types.ndarray(dtype=ti.i32, ndim=2),
+    factor: ti.f32,
+):
+    for n in range(offsets.shape[0]):
+        x = (center_idx[0] + offsets[n, 0]) % grid.shape[0]
+        y = (center_idx[1] + offsets[n, 1]) % grid.shape[1]
+        z = (center_idx[2] + offsets[n, 2]) % grid.shape[2]
+        grid[x, y, z] *= factor
+
+
+@ti.kernel
+def _div_sphere_offsets(
+    grid: ti.template(),
+    center_idx: ti.types.ndarray(dtype=ti.i32, ndim=1),
+    offsets: ti.types.ndarray(dtype=ti.i32, ndim=2),
+    divisor: ti.f32,
+):
+    for n in range(offsets.shape[0]):
+        x = (center_idx[0] + offsets[n, 0]) % grid.shape[0]
+        y = (center_idx[1] + offsets[n, 1]) % grid.shape[1]
+        z = (center_idx[2] + offsets[n, 2]) % grid.shape[2]
+        grid[x, y, z] /= divisor
+
+
+@ti.kernel
+def _clamp_grid(grid: ti.template(), min_val: ti.f32, max_val: ti.f32):
+    for i, j, k in ti.ndrange(grid.shape[0], grid.shape[1], grid.shape[2]):
+        value = grid[i, j, k]
+        if value < min_val:
+            grid[i, j, k] = min_val
+        elif value > max_val:
+            grid[i, j, k] = max_val
+
+
+class VoxelGridTaichi(VoxelGrid):
+    """Voxel grid with Taichi CPU kernels for mutating sphere operations."""
+
+    backend_name = "taichi-cpu"
+
+    def __init__(self, cell, resolution=None, gpts=None):
+        super().__init__(cell=cell, resolution=resolution, gpts=gpts)
+        self.grid = ti.field(dtype=ti.f32, shape=tuple(int(x) for x in self.gpts))
+
+    def to_numpy(self):
+        """Return the voxel values as a NumPy array."""
+        return self.grid.to_numpy()
+
+    def set_sphere(self, center, radius, value=1):
+        center_idx = self._center_index(center)
+        offsets = self._sphere_offsets(radius)
+        _set_sphere_offsets(self.grid, center_idx, offsets, float(value))
+
+    def add_sphere(self, center, radius, value=1):
+        center_idx = self._center_index(center)
+        offsets = self._sphere_offsets(radius)
+        _add_sphere_offsets(self.grid, center_idx, offsets, float(value))
+
+    def mul_sphere(self, center, radius, factor=2):
+        center_idx = self._center_index(center)
+        offsets = self._sphere_offsets(radius)
+        _mul_sphere_offsets(self.grid, center_idx, offsets, float(factor))
+
+    def div_sphere(self, center, radius, factor=2):
+        center_idx = self._center_index(center)
+        offsets = self._sphere_offsets(radius)
+        _div_sphere_offsets(self.grid, center_idx, offsets, float(factor))
+
+    def add_spheres(self, centers, radii, value=1):
+        centers, radii = self._validate_spheres(centers, radii)
+        center_indices = self.positions_to_indices(centers)
+        for radius in np.unique(radii):
+            offsets = _cached_sphere_offsets(float(radius), tuple(self.gpts), tuple(map(tuple, self.cell)))
+            for center_idx in center_indices[radii == radius]:
+                _add_sphere_offsets(self.grid, center_idx, offsets, float(value))
+
+    def set_spheres(self, centers, radii, value=1):
+        centers, radii = self._validate_spheres(centers, radii)
+        center_indices = self.positions_to_indices(centers)
+        for radius in np.unique(radii):
+            offsets = _cached_sphere_offsets(float(radius), tuple(self.gpts), tuple(map(tuple, self.cell)))
+            for center_idx in center_indices[radii == radius]:
+                _set_sphere_offsets(self.grid, center_idx, offsets, float(value))
+
+    def clamp_grid(self, min_val=0.0, max_val=1.0):
+        _clamp_grid(self.grid, float(min_val), float(max_val))
+
+    def synchronize(self):
+        """Synchronize Taichi kernels."""
+        ti.sync()
+
+
+__all__ = ["VoxelGridTaichi"]

atomvoxelizer/voxelgrid.py

@@ -0,0 +1,308 @@
+from __future__ import annotations
+
+from functools import lru_cache
+
+import numpy as np
+
+
+@lru_cache(maxsize=50)
+def _cached_sphere_mask(radius, gpts, cell):
+    nx, ny, nz = gpts
+    ix, iy, iz = np.meshgrid(
+        np.arange(nx),
+        np.arange(ny),
+        np.arange(nz),
+        indexing="ij",
+    )
+    frac_coords = (np.stack([ix, iy, iz], axis=-1) + 0.5) / gpts
+    center_frac = np.array([0.5, 0.5, 0.5])
+    disp_frac = frac_coords - center_frac
+    disp_frac -= np.round(disp_frac)
+    disp_mic = disp_frac @ cell
+    dist2 = np.sum(disp_mic**2, axis=-1)
+    return dist2 <= radius**2
+
+
+@lru_cache(maxsize=200)
+def _cached_sphere_offsets(radius, gpts, cell):
+    gpts_arr = np.array(gpts, dtype=np.int32)
+    cell_arr = np.array(cell, dtype=np.float64)
+    lengths = np.linalg.norm(cell_arr, axis=1)
+    max_offsets = np.ceil(radius / lengths * gpts_arr).astype(np.int32)
+
+    offsets = []
+    for dx in range(-max_offsets[0], max_offsets[0] + 1):
+        for dy in range(-max_offsets[1], max_offsets[1] + 1):
+            for dz in range(-max_offsets[2], max_offsets[2] + 1):
+                disp_frac = np.array([dx, dy, dz], dtype=np.float64) / gpts_arr
+                disp = disp_frac @ cell_arr
+                if np.dot(disp, disp) <= radius**2:
+                    offsets.append((dx, dy, dz))
+
+    return np.array(offsets, dtype=np.int32)
+
+
+class VoxelGrid:
+    """Periodic voxel grid implemented with NumPy only."""
+
+    backend_name = "numpy"
+
+    def __init__(self, cell, resolution=None, gpts=None):
+        self.cell = np.array(cell, dtype=np.float64)
+        self.cell_inv = np.linalg.inv(self.cell)
+
+        if resolution is None and gpts is None:
+            raise ValueError("Either resolution or gpts must be specified")
+        if resolution is not None and gpts is not None:
+            raise ValueError("Only one of resolution or gpts can be specified")
+
+        lengths = np.linalg.norm(self.cell, axis=1)
+        if resolution is not None:
+            self.gpts = np.ceil(lengths / resolution).astype(int)
+        else:
+            self.gpts = np.array(gpts, dtype=int)
+        self.resolution = lengths / self.gpts
+
+        self.grid = np.zeros(tuple(self.gpts), dtype=np.float32)
+
+    def to_numpy(self):
+        """Return the voxel values as a NumPy array."""
+        return self.grid
+
+    def position_to_index(self, r):
+        """Convert real-space position to voxel index using periodic wrapping."""
+        frac = np.asarray(r, dtype=np.float64) @ self.cell_inv
+        frac_wrapped = np.clip(frac % 1.0, 0.0, np.nextafter(1.0, 0.0))
+        idx = np.floor(frac_wrapped * self.gpts).astype(int)
+        return tuple(idx)
+
+    def index_to_position(self, i, j, k):
+        """Convert a grid index to the real-space voxel center."""
+        frac = (np.array([i, j, k]) + 0.5) / self.gpts
+        return frac @ self.cell
+
+    def _center_index(self, center):
+        center_frac = np.asarray(center, dtype=np.float64) @ self.cell_inv % 1.0
+        return np.floor(center_frac * self.gpts).astype(np.int32)
+
+    def _offset_indices(self, center_idx, offsets):
+        indices = (offsets + center_idx) % self.gpts
+        return tuple(indices[:, axis] for axis in range(3))
+
+    def _sphere_offsets(self, radius):
+        return _cached_sphere_offsets(float(radius), tuple(self.gpts), tuple(map(tuple, self.cell)))
+
+    def _sphere_indices(self, center, radius):
+        return self._offset_indices(self._center_index(center), self._sphere_offsets(radius))
+
+    def set_sphere(self, center, radius, value=1):
+        self.grid[self._sphere_indices(center, radius)] = value
+
+    def add_sphere(self, center, radius, value=1):
+        np.add.at(self.grid, self._sphere_indices(center, radius), value)
+
+    def mul_sphere(self, center, radius, factor=2):
+        self.grid[self._sphere_indices(center, radius)] *= factor
+
+    def div_sphere(self, center, radius, factor=2):
+        self.grid[self._sphere_indices(center, radius)] /= factor
+
+    def positions_to_indices(self, positions):
+        positions = np.asarray(positions, dtype=np.float64)
+        frac = positions @ self.cell_inv
+        frac_wrapped = np.clip(frac % 1.0, 0.0, np.nextafter(1.0, 0.0))
+        return np.floor(frac_wrapped * self.gpts).astype(np.int32)
+
+    def _validate_spheres(self, centers, radii):
+        centers = np.asarray(centers, dtype=np.float64)
+        radii = np.asarray(radii, dtype=np.float64)
+        if centers.ndim != 2 or centers.shape[1] != 3:
+            raise ValueError("centers must have shape (N, 3)")
+        if radii.ndim != 1 or radii.shape[0] != centers.shape[0]:
+            raise ValueError("radii must have shape (N,)")
+        return centers, radii
+
+    def add_spheres(self, centers, radii, value=1):
+        centers, radii = self._validate_spheres(centers, radii)
+        for center, radius in zip(centers, radii):
+            self.add_sphere(center, radius, value=value)
+
+    def set_spheres(self, centers, radii, value=1):
+        centers, radii = self._validate_spheres(centers, radii)
+        for center, radius in zip(centers, radii):
+            self.set_sphere(center, radius, value=value)
+
+    def clamp_grid(self, min_val=0.0, max_val=1.0):
+        np.clip(self.grid, min_val, max_val, out=self.grid)
+
+    def sample_voxels_in_range(self, min_val=0.0, max_val=1.0, min_dist=0.0, return_indices=False, seed=None):
+        """
+        Yield voxel positions or indices whose values lie in [min_val, max_val].
+
+        When returning real-space positions, ``min_dist`` enforces a minimum
+        Euclidean separation in Angstrom between yielded samples.
+        """
+        rng = np.random.default_rng(seed)
+        grid = self.to_numpy()
+        mask = (grid >= min_val) & (grid <= max_val)
+        candidates = np.argwhere(mask)
+
+        if candidates.shape[0] == 0:
+            raise ValueError("No voxels in specified value range.")
+        if return_indices and min_dist > 0:
+            raise ValueError("min_dist only supported when return_indices=False")
+
+        positions = candidates if return_indices else np.array([self.index_to_position(*idx) for idx in candidates])
+        selected = []
+        indices = rng.permutation(len(positions))
+        min_dist2 = min_dist**2
+
+        for i in indices:
+            pos = positions[i]
+            if min_dist > 0 and selected:
+                d2 = np.sum((np.array(selected) - pos) ** 2, axis=1)
+                if np.any(d2 < min_dist2):
+                    continue
+            selected.append(pos)
+            yield tuple(candidates[i]) if return_indices else pos
+
+    def plot_3D(self, threshold=0.1, s=5, draw_cell=True):
+        """Plot voxels with values above ``threshold`` in real space."""
+        import matplotlib.pyplot as plt
+
+        nx, ny, nz = self.gpts
+        ix, iy, iz = np.meshgrid(
+            np.arange(nx) + 0.5,
+            np.arange(ny) + 0.5,
+            np.arange(nz) + 0.5,
+            indexing="ij",
+        )
+
+        frac_coords = np.stack([ix / nx, iy / ny, iz / nz], axis=-1)
+        real_coords = frac_coords @ self.cell
+        grid = self.to_numpy()
+        mask = grid > threshold
+        xyz = real_coords[mask]
+        values = grid[mask]
+
+        fig = plt.figure()
+        ax = fig.add_subplot(projection="3d")
+        p = ax.scatter(xyz[:, 0], xyz[:, 1], xyz[:, 2], c=values, cmap="viridis", s=s)
+        fig.colorbar(p, ax=ax, label="Voxel value")
+
+        if draw_cell:
+            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],
+                ]
+            )
+            corners = corners_frac @ self.cell
+            edges = [
+                (0, 1),
+                (0, 2),
+                (0, 3),
+                (1, 4),
+                (1, 5),
+                (2, 4),
+                (2, 6),
+                (3, 5),
+                (3, 6),
+                (4, 7),
+                (5, 7),
+                (6, 7),
+            ]
+            for i, j in edges:
+                ax.plot(
+                    [corners[i, 0], corners[j, 0]],
+                    [corners[i, 1], corners[j, 1]],
+                    [corners[i, 2], corners[j, 2]],
+                    color="black",
+                )
+
+        all_coords = np.concatenate([xyz, corners]) if draw_cell else xyz
+        xlim = [all_coords[:, 0].min(), all_coords[:, 0].max()]
+        ylim = [all_coords[:, 1].min(), all_coords[:, 1].max()]
+        zlim = [all_coords[:, 2].min(), all_coords[:, 2].max()]
+        max_range = max(xlim[1] - xlim[0], ylim[1] - ylim[0], zlim[1] - zlim[0]) / 2.0
+        mid_x, mid_y, mid_z = np.mean(xlim), np.mean(ylim), np.mean(zlim)
+
+        ax.set_xlim(mid_x - max_range, mid_x + max_range)
+        ax.set_ylim(mid_y - max_range, mid_y + max_range)
+        ax.set_zlim(mid_z - max_range, mid_z + max_range)
+        ax.set_xlabel("x")
+        ax.set_ylabel("y")
+        ax.set_zlabel("z")
+        plt.tight_layout()
+        plt.show()
+
+    def plot_2D(self, axis="z", index=None, position=None, threshold=0.1, draw_cell=True, real_space=True):
+        """Plot a 2D slice of the voxel grid along ``axis``."""
+        import matplotlib.pyplot as plt
+
+        ax_map = {"x": 0, "y": 1, "z": 2}
+        if axis not in ax_map:
+            raise ValueError("Axis must be 'x', 'y', or 'z'")
+        ax_idx = ax_map[axis]
+
+        if index is not None and position is not None:
+            raise ValueError("Specify either `index` or `position`, not both")
+        if position is not None:
+            index = self.position_to_index(np.eye(3)[ax_idx] * position)[ax_idx]
+        if index is None:
+            index = self.gpts[ax_idx] // 2
+
+        shape = self.grid.shape
+        if not (0 <= index < shape[ax_idx]):
+            raise IndexError(f"{axis}-index {index} out of bounds (0 to {shape[ax_idx] - 1})")
+
+        axes = [0, 1, 2]
+        axes.remove(ax_idx)
+        ax1, ax2 = axes
+
+        slicers = [slice(None)] * 3
+        slicers[ax_idx] = index
+        slice_grid = self.to_numpy()[tuple(slicers)]
+
+        n1, n2 = self.gpts[ax1], self.gpts[ax2]
+        if real_space:
+            i1 = (np.arange(n1) + 0.5) / n1
+            i2 = (np.arange(n2) + 0.5) / n2
+            coords = np.meshgrid(i1, i2, indexing="ij")
+            frac_coords = np.stack(coords, axis=-1)
+            xy = frac_coords @ self.cell[[ax1, ax2], :]
+            xvals, yvals = xy[..., 0], xy[..., 1]
+        else:
+            xvals, yvals = np.meshgrid(np.arange(n1), np.arange(n2), indexing="ij")
+
+        mask = slice_grid > threshold
+        fig, ax = plt.subplots()
+        sc = ax.scatter(xvals[mask], yvals[mask], c=slice_grid[mask], cmap="viridis", s=10)
+        fig.colorbar(sc, ax=ax, label="Voxel value")
+
+        if draw_cell and real_space:
+            corners_frac = np.array([[0, 0], [1, 0], [1, 1], [0, 1], [0, 0]])
+            corners_real = corners_frac @ self.cell[[ax1, ax2], :]
+            ax.plot(corners_real[:, 0], corners_real[:, 1], "k--", lw=1)
+
+        ax.set_xlabel(f'{["x", "y", "z"][ax1]}' + (" [Angstrom]" if real_space else " (voxel)"))
+        ax.set_ylabel(f'{["x", "y", "z"][ax2]}' + (" [Angstrom]" if real_space else " (voxel)"))
+        ax.set_title(f"{axis.upper()} Slice at index {index}")
+        ax.set_aspect("equal")
+        plt.tight_layout()
+        plt.show()
+
+    def __repr__(self):
+        return f"VoxelGrid\n{self.cell} Cell\n{self.resolution} Resolution\n{self.gpts} gpts"
+
+
+VoxelGridNumPy = VoxelGrid
+
+
+__all__ = ["VoxelGrid", "VoxelGridNumPy", "_cached_sphere_mask", "_cached_sphere_offsets"]

atomvoxelizer-0.1.0.dist-info/METADATA

@@ -0,0 +1,174 @@
+Metadata-Version: 2.4
+Name: AtomVoxelizer
+Version: 0.1.0
+Summary: Periodic atom-centered voxel grids for atomistic structures.
+Author: AtomVoxelizer contributors
+Project-URL: Homepage, https://gitlab.com/tgmaxson/atomvoxelizer
+Project-URL: Documentation, https://atomvoxelizer.readthedocs.io/
+Project-URL: Repository, https://gitlab.com/tgmaxson/atomvoxelizer
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3 :: Only
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Chemistry
+Classifier: Topic :: Scientific/Engineering :: Physics
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+Requires-Dist: matplotlib
+Requires-Dist: numpy
+Provides-Extra: numba
+Requires-Dist: numba; extra == "numba"
+Provides-Extra: cupy
+Requires-Dist: cupy; extra == "cupy"
+Requires-Dist: numba; extra == "cupy"
+Provides-Extra: taichi
+Requires-Dist: taichi; extra == "taichi"
+Provides-Extra: analysis
+Requires-Dist: scikit-image; extra == "analysis"
+Provides-Extra: examples
+Requires-Dist: ase; extra == "examples"
+Requires-Dist: requests; extra == "examples"
+Provides-Extra: docs
+Requires-Dist: sphinx; extra == "docs"
+Provides-Extra: dev
+Requires-Dist: ase; extra == "dev"
+Requires-Dist: numba; extra == "dev"
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: requests; extra == "dev"
+Requires-Dist: scikit-image; extra == "dev"
+Requires-Dist: sphinx; extra == "dev"
+Requires-Dist: taichi; extra == "dev"
+Provides-Extra: bench
+Requires-Dist: ase; extra == "bench"
+Requires-Dist: numba; extra == "bench"
+Requires-Dist: taichi; extra == "bench"
+Provides-Extra: publish
+Requires-Dist: build; extra == "publish"
+Requires-Dist: twine; extra == "publish"
+
+# AtomVoxelizer
+
+AtomVoxelizer builds periodic atom-centered voxel grids for atomistic structures.
+The core `VoxelGrid` class stores a 3D NumPy grid over a periodic cell and provides
+helpers for adding, setting, scaling, sampling, and plotting spherical regions.
+
+## Installation
+
+Install from this repository:
+
+```bash
+pip install .
+```
+
+Install optional acceleration backends with extras:
+
+```bash
+pip install ".[numba]"
+pip install ".[taichi]"
+pip install ".[cupy]"
+pip install ".[analysis]"
+```
+
+`VoxelGrid` is always the NumPy backend. Optional acceleration backends are
+explicit: `VoxelGridNumba`, `VoxelGridTaichi`, and `VoxelGridCuPy`.
+`VoxelGridAnalysis` provides connected-volume and marching-cubes surface-area
+analysis when the `analysis` extra is installed.
+
+For development, examples, tests, and documentation:
+
+```bash
+pip install -e ".[dev,examples]"
+```
+
+## Basic Usage
+
+```python
+import numpy as np
+
+from atomvoxelizer import VoxelGrid
+
+cell = np.eye(3) * 10.0
+grid = VoxelGrid(cell=cell, resolution=0.25)
+
+grid.add_sphere(center=np.array([5.0, 5.0, 5.0]), radius=1.0, value=1.0)
+grid.set_sphere(center=np.array([2.0, 2.0, 2.0]), radius=0.5, value=-1.0)
+grid.clamp_grid(min_val=-1.0, max_val=1.0)
+```
+
+## Zeolite Example
+
+The zeolite example and CIF files live in `examples/`.
+
+```bash
+pip install -e ".[examples]"
+python examples/zeolite_voxel.py BEA
+```
+
+The script reads a framework CIF, builds voxel grids at several resolutions, plots
+middle XZ slices, benchmarks supercell scaling, and opens a 3D scatter plot.
+
+The analysis example estimates pore volume and internal surface area:
+
+```bash
+pip install -e ".[examples,analysis]"
+python examples/zeolite_analysis.py BEA --resolution 0.25
+python examples/zeolite_analysis.py BEA --convergence 1.0 0.75 0.5 --plot bea_convergence.png
+```
+
+## Tests and Benchmarks
+
+Run the correctness tests with:
+
+```bash
+pytest
+```
+
+Run the backend benchmark with:
+
+```bash
+python benchmarks/benchmark_backends.py --backends numpy numba taichi cupy
+```
+
+Run the built-in structure benchmarks for a zeolite and a roughly 1000 atom Wulff
+construction with:
+
+```bash
+python benchmarks/benchmark_structures.py
+```
+
+Backends whose optional dependencies are not installed are reported as missing.
+
+## Documentation
+
+Documentation is scaffolded with Sphinx for Read the Docs.
+
+Build it locally with:
+
+```bash
+pip install -e ".[docs]"
+sphinx-build -b html docs/source docs/build/html
+```
+
+Read the Docs can use `.readthedocs.yaml` directly.
+
+## Publishing
+
+Build and check PyPI artifacts with:
+
+```bash
+pip install -e ".[publish]"
+python -m build
+twine check dist/*
+```
+
+Upload to TestPyPI first, then PyPI:
+
+```bash
+twine upload --repository testpypi dist/*
+twine upload dist/*
+```

atomvoxelizer-0.1.0.dist-info/RECORD

@@ -0,0 +1,10 @@
+atomvoxelizer/__init__.py,sha256=SA3jmOk191wtX3E_Zs3TbH9A56xzJ-VVStKgBE33eb4,741
+atomvoxelizer/analysis.py,sha256=3NNXJvzPiZ8K52TZ89IPLahC3_PPb7KBMtuFnvPr2DE,7436
+atomvoxelizer/cupy_backend.py,sha256=t95ZZmtIcsGT_A-2x9hlxdk_gX9ziU_t5bfwxzvhGl4,3404
+atomvoxelizer/numba_backend.py,sha256=7jSrwcZOpzR6TnmdmLmeRC83KutlH_y_RZ4CoFhNLpQ,4680
+atomvoxelizer/taichi_backend.py,sha256=NgnGkFAEcMdSx--EHkpo8QBd2ZxeF6ll_oD_FUSvleU,5449
+atomvoxelizer/voxelgrid.py,sha256=HSSUAVDWLkXQ2rEhwgbZEFI4w1tdA8Zw9H6iy-Krb3c,11721
+atomvoxelizer-0.1.0.dist-info/METADATA,sha256=M00nvGjj5ml1_R-o50AUbVJurArZb9szdZEeEwwYzs4,4760
+atomvoxelizer-0.1.0.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
+atomvoxelizer-0.1.0.dist-info/top_level.txt,sha256=KiH4w3N7wUa8hwDdDLNwfgHHj8BK1uE5LSzb63tzA7E,14
+atomvoxelizer-0.1.0.dist-info/RECORD,,

atomvoxelizer-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (82.0.1)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

atomvoxelizer-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+atomvoxelizer