diff --git a/PyAPD/__init__.py b/PyAPD/__init__.py index c0be987..f41b9bb 100644 --- a/PyAPD/__init__.py +++ b/PyAPD/__init__.py @@ -46,3 +46,36 @@ from .log_res_utils import ( reorder_variables as reorder_variables, ) +from .polycrystal_atomistic import ( + bulk_bcc_positions as bulk_bcc_positions, +) +from .polycrystal_atomistic import ( + bulk_lattice_positions as bulk_lattice_positions, +) +from .polycrystal_atomistic import ( + generate_bcc_polycrystal as generate_bcc_polycrystal, +) +from .polycrystal_atomistic import ( + generate_fcc_polycrystal as generate_fcc_polycrystal, +) +from .polycrystal_atomistic import ( + generate_hcp_polycrystal as generate_hcp_polycrystal, +) +from .polycrystal_atomistic import ( + generate_polycrystal as generate_polycrystal, +) +from .polycrystal_atomistic import ( + generate_sc_polycrystal as generate_sc_polycrystal, +) +from .polycrystal_atomistic import ( + min_interatomic_distance as min_interatomic_distance, +) +from .polycrystal_atomistic import ( + pbc_pair_cull as pbc_pair_cull, +) +from .polycrystal_atomistic import ( + shoemake_uniform_so3 as shoemake_uniform_so3, +) +from .polycrystal_atomistic import ( + write_lammps_data as write_lammps_data, +) diff --git a/PyAPD/apds.py b/PyAPD/apds.py index a0dcf79..f8635c8 100644 --- a/PyAPD/apds.py +++ b/PyAPD/apds.py @@ -42,6 +42,7 @@ def __init__( error_tolerance=0.01, pixel_size_prefactor=2, seed=-1, + periodic=False, ): """ Construct an anisotropic power diagram system. @@ -86,6 +87,11 @@ def __init__( Multiplier applied to the computed pixel count. Default: 2. seed : int Manual random seed (< 0 means no seed is set). Default: -1. + periodic : bool + If True, distances between seeds and points use the minimum-image + convention across the (rectilinear) domain, so the power diagram + is periodic and grains may wrap across the domain boundary. + Default: False (fully backward-compatible). """ self.N = N @@ -103,6 +109,15 @@ def __init__( self.set_domain(domain) + self.periodic = bool(periodic) + # Per-dimension edge lengths of the (rectilinear) domain. Kept on the + # same device/dtype as the domain so KeOps arithmetic stays homogeneous. + edge = (self.domain[:, 1] - self.domain[:, 0]).to( + device=self.device, dtype=self.dt + ) + self._L_tensor = edge # shape (D,), plain torch + self._L = LazyTensor(edge.view(1, 1, self.D)) # broadcast over (N, M, D) + self.set_X(X) self.set_As(As) @@ -387,6 +402,22 @@ def mask_pixels(self, mask): self.W = initial_guess_heuristic(self.As, self.target_masses, self.D) self.w = LazyTensor(self.W.view(self.N, 1, 1)) + def _displacement(self, y, x): + """Return ``y - x`` as a KeOps LazyTensor, wrapped to the minimum + image across the periodic domain when ``self.periodic`` is True. + + Uses the KeOps-compatible expression ``dy - L * round(dy / L)``, + which rounds ``dy / L`` to the nearest integer per component. + ``LazyTensor.floor()`` is unavailable on some KeOps builds and the + Python ``%`` operator is not defined between LazyTensors, so the + ``.round()`` form is used. + """ + dy = y - x + if self.periodic: + shift = (dy / self._L).round() + dy = dy - self._L * shift + return dy + def assemble_apd( self, record_time=False, verbose=False, color_by=None, backend="auto" ): @@ -396,7 +427,8 @@ def assemble_apd( if self.Y is None: self.assemble_pixels() start = time.time() - D_ij = ((self.y - self.x) | self.a.matvecmult(self.y - self.x)) - self.w + dy = self._displacement(self.y, self.x) + D_ij = (dy | self.a.matvecmult(dy)) - self.w # Find which grain each pixel belongs to grain_indices = D_ij.argmin(dim=0, backend=backend).ravel() time_taken = time.time() - start @@ -410,6 +442,34 @@ def assemble_apd( else: return grain_indices + def grain_of(self, points, backend="auto"): + """Return the grain index that owns each point in ``points``. + + Parameters + ---------- + points : torch.Tensor of shape (M, D) + Query points, in the same coordinate frame as ``self.domain``. + When ``self.periodic`` is True, query points are compared to each + seed using the minimum-image distance across the domain. + backend : str + KeOps reduction backend (forwarded to argmin). + + Returns + ------- + torch.Tensor of shape (M,), dtype int64 + For each row of ``points``, the grain index ``i`` minimising + ``(y - x_i) . A_i . (y - x_i) - w_i`` under the active metric. + """ + pts = points.to(device=self.device, dtype=self.dt) + if pts.ndim != 2 or pts.shape[1] != self.D: + raise ValueError( + f"grain_of: expected (M, {self.D}), got {tuple(pts.shape)}" + ) + y = LazyTensor(pts.view(1, pts.shape[0], self.D)) + dy = self._displacement(y, self.x) + D_ij = (dy | self.a.matvecmult(dy)) - self.w + return D_ij.argmin(dim=0, backend=backend).view(-1) + def plot_apd( self, color_by=None, mode="auto", alpha=None, ps_scale=False, marker_scale=20.0 ): @@ -603,7 +663,8 @@ def OT_dual_function(self, W, backend="auto"): self.W = W self.w = LazyTensor(self.W.view(self.N, 1, 1)) - D_ij = ((self.y - self.x) | self.a.matvecmult(self.y - self.x)) - self.w + dy = self._displacement(self.y, self.x) + D_ij = (dy | self.a.matvecmult(dy)) - self.w idx = D_ij.argmin(dim=0, backend=backend).view(-1) ind_select = torch.index_select(self.X, 0, idx) - self.Y @@ -623,7 +684,8 @@ def check_optimality( if self.Y is None: self.assemble_pixels() - D_ij = ((self.y - self.x) | self.a.matvecmult(self.y - self.x)) - self.w + dy = self._displacement(self.y, self.x) + D_ij = (dy | self.a.matvecmult(dy)) - self.w grain_indices = D_ij.argmin(dim=0, backend=backend).ravel() volumes = torch.bincount(grain_indices, self.PS, minlength=self.N) @@ -707,7 +769,8 @@ def adjust_X(self, backend="auto"): if not self.optimality: print("Find optimal W first!") else: - D_ij = ((self.y - self.x) | self.a.matvecmult(self.y - self.x)) - self.w + dy = self._displacement(self.y, self.x) + D_ij = (dy | self.a.matvecmult(dy)) - self.w grain_indices = D_ij.argmin(dim=0, backend=backend).ravel() normalisation = torch.bincount(grain_indices, self.PS, minlength=self.N) new_X0 = ( @@ -730,6 +793,12 @@ def adjust_X(self, backend="auto"): else: self.X = torch.stack([new_X0, new_X1], dim=1) + if self.periodic: + # Centroids may fall outside the box when a grain straddles the + # periodic seam; wrap them back into [domain[:, 0], domain[:, 1)). + origin = self.domain[:, 0] + self.X = origin + torch.remainder(self.X - origin, self._L_tensor) + self.optimality = False self.x = LazyTensor(self.X.view(self.N, 1, self.D)) diff --git a/PyAPD/polycrystal_atomistic.py b/PyAPD/polycrystal_atomistic.py new file mode 100644 index 0000000..82e3d17 --- /dev/null +++ b/PyAPD/polycrystal_atomistic.py @@ -0,0 +1,535 @@ +"""Periodic atomistic polycrystal generator built on PyAPD. + +The generator is crystal-agnostic: a grain map comes from a periodic +anisotropic power diagram, each grain is a randomly oriented copy of a +user-supplied lattice (an orthorhombic conventional cell plus a fractional +basis), and grains are filled in a space-filling way before grain-boundary +overlaps are culled. Cubic crystals (SC / BCC / FCC / diamond / B2 / ...) +use a scalar cell; orthorhombic, tetragonal and HCP (as an orthorhombic +supercell) use a per-axis cell. + +Public API +---------- +- generate_polycrystal(cell, basis, ...) : general driver (any lattice) +- generate_bcc_polycrystal(a, ...) : BCC convenience wrapper +- generate_fcc_polycrystal(a, ...) : FCC convenience wrapper +- generate_sc_polycrystal(a, ...) : simple-cubic convenience wrapper +- generate_hcp_polycrystal(a, c, ...) : HCP wrapper (orthorhombic cell) +- bulk_lattice_positions(cell, basis, counts): un-rotated reference lattice +- bulk_bcc_positions(a, m) : BCC reference lattice (compat) +- shoemake_uniform_so3(rng) : uniform-random SO(3) rotation +- pbc_pair_cull(positions, gids, ...) : cross-grain overlap removal +- min_interatomic_distance(cell, basis) : nearest-neighbour spacing +- write_lammps_data(path, ...) : LAMMPS atomic data-file writer + +Predefined fractional bases: SC_BASIS, BCC_BASIS, FCC_BASIS. HCP needs a +per-axis cell, so use generate_hcp_polycrystal / HCP_ORTHO_BASIS. +""" + +from __future__ import annotations + +import math +import warnings + +import numpy as np +import torch +from scipy.spatial import cKDTree + +from PyAPD.apds import apd_system + +# Fractional bases for cubic conventional cells (use with a scalar cell `a`). +SC_BASIS = np.array([[0.0, 0.0, 0.0]]) +BCC_BASIS = np.array([[0.0, 0.0, 0.0], [0.5, 0.5, 0.5]]) +FCC_BASIS = np.array( + [[0.0, 0.0, 0.0], [0.5, 0.5, 0.0], [0.5, 0.0, 0.5], [0.0, 0.5, 0.5]] +) +# HCP as a C-centred orthorhombic supercell: cell = (a, a*sqrt(3), c), 4 atoms. +HCP_ORTHO_BASIS = np.array( + [ + [0.0, 0.0, 0.0], + [0.5, 0.5, 0.0], + [0.5, 1.0 / 6.0, 0.5], + [0.0, 2.0 / 3.0, 0.5], + ] +) + + +def _as_cell(cell) -> np.ndarray: + """Coerce a scalar (cubic) or length-3 (orthorhombic) cell to shape (3,).""" + cell = np.asarray(cell, dtype=np.float64) + if cell.ndim == 0: + cell = np.repeat(cell, 3) + if cell.shape != (3,): + raise ValueError(f"cell must be a scalar or length-3; got shape {cell.shape}") + if np.any(cell <= 0): + raise ValueError(f"cell edge lengths must be positive; got {cell}") + return cell + + +def shoemake_uniform_so3(rng: np.random.Generator) -> np.ndarray: + """Uniform-random rotation matrix on SO(3) via Shoemake's quaternion method. + + Shoemake (1992), "Uniform Random Rotations", Graphics Gems III. + Returns a (3, 3) numpy array with det=+1 and R^T R = I. + """ + u1, u2, u3 = rng.uniform(0.0, 1.0, size=3) + q = np.array( + [ + np.sqrt(1.0 - u1) * np.sin(2.0 * np.pi * u2), + np.sqrt(1.0 - u1) * np.cos(2.0 * np.pi * u2), + np.sqrt(u1) * np.sin(2.0 * np.pi * u3), + np.sqrt(u1) * np.cos(2.0 * np.pi * u3), + ] + ) + x, y, z, w = q + # Standard quaternion-to-rotation-matrix. + return np.array( + [ + [1 - 2 * (y * y + z * z), 2 * (x * y - z * w), 2 * (x * z + y * w)], + [2 * (x * y + z * w), 1 - 2 * (x * x + z * z), 2 * (y * z - x * w)], + [2 * (x * z - y * w), 2 * (y * z + x * w), 1 - 2 * (x * x + y * y)], + ] + ) + + +def bulk_lattice_positions(cell, basis, counts) -> np.ndarray: + """Un-rotated reference lattice tiled `counts` times along each axis. + + Parameters + ---------- + cell : scalar or (3,) array + Orthorhombic conventional-cell edge lengths (scalar = cubic). + basis : (n_basis, 3) array + Fractional atom coordinates within the conventional cell, in [0, 1). + counts : int or (3,) int array + Number of cells along each axis. + + Returns + ------- + (n_basis * prod(counts), 3) float ndarray in [0, counts * cell). + """ + cell = _as_cell(cell) + basis = np.asarray(basis, dtype=np.float64).reshape(-1, 3) + counts = np.broadcast_to(np.asarray(counts, dtype=int), (3,)) + ii, jj, kk = np.meshgrid( + np.arange(counts[0]), + np.arange(counts[1]), + np.arange(counts[2]), + indexing="ij", + ) + idx = np.stack([ii.ravel(), jj.ravel(), kk.ravel()], axis=1).astype(np.float64) + # (n_cell, 1, 3) + (1, n_basis, 3) -> (n_cell, n_basis, 3), scaled by cell. + latt = (idx[:, None, :] + basis[None, :, :]) * cell + return latt.reshape(-1, 3) + + +def bulk_bcc_positions(a: float, m: int) -> np.ndarray: + """Reference BCC lattice tiled m x m x m in [0, m*a)^3 (compatibility helper). + + Two atoms per conventional cell: corner at (0,0,0) and body-centre at + (a/2, a/2, a/2). Returns 2*m**3 rows of shape (3,), dtype float64. + """ + return bulk_lattice_positions(a, BCC_BASIS, m) + + +def min_interatomic_distance(cell, basis) -> float: + """Nearest-neighbour spacing of the infinite lattice (cell, basis). + + Used to set the grain-boundary cull cutoff. Computed by tiling a 3x3x3 + block of cells and taking the smallest non-zero distance from any + central-cell atom; valid because the nearest neighbour of any site lies + within one cell along each axis. + """ + cell = _as_cell(cell) + basis = np.asarray(basis, dtype=np.float64).reshape(-1, 3) + offsets = np.array( + [[i, j, k] for i in (-1, 0, 1) for j in (-1, 0, 1) for k in (-1, 0, 1)], + dtype=np.float64, + ) + block = ((offsets[:, None, :] + basis[None, :, :]) * cell).reshape(-1, 3) + ref = basis * cell # central-cell atoms (a subset of `block`) + tree = cKDTree(block) + # k=2: nearest is the atom itself (distance 0), second is the true NN. + dists, _ = tree.query(ref, k=2) + return float(dists[:, 1].min()) + + +def pbc_pair_cull( + positions: np.ndarray, + grain_ids: np.ndarray, + cutoff: float, + box, + cull_same_grain: bool = False, +) -> tuple[np.ndarray, np.ndarray]: + """Remove one atom of every cross-grain pair within `cutoff` under PBC. + + Parameters + ---------- + positions : (N, 3) float ndarray in [0, box) + grain_ids : (N,) int ndarray, one grain label per atom + cutoff : float, max distance for an overlap (anything <= cutoff is a pair) + box : scalar or (3,) array, per-axis box edge length(s) (PBC per axis) + cull_same_grain : bool, default False + If True, all within-cutoff pairs are culled regardless of grain + membership. Use for rotate-and-wrap constructions where same-grain + near-overlaps are PBC-seam artifacts, not real lattice bonds. + + Returns + ------- + kept_positions, kept_grain_ids -- atoms surviving the cull. + + Notes + ----- + By default (cull_same_grain=False) same-grain pairs are NEVER dropped + (would corrupt a clean lattice). Set cull_same_grain=True when atoms + came from a rotate-and-wrap construction where same-grain near-overlaps + are geometric artifacts, not real lattice bonds. + """ + boxsize = _as_cell(box) + tree = cKDTree(positions, boxsize=boxsize) + pairs = tree.query_pairs(r=cutoff, output_type="ndarray") + if len(pairs) == 0: + return positions, grain_ids + if cull_same_grain: + # Cull every within-cutoff pair regardless of grain. + drop = np.unique(pairs[:, 1]) + else: + cross = grain_ids[pairs[:, 0]] != grain_ids[pairs[:, 1]] + drop = np.unique(pairs[cross, 1]) + keep_mask = np.ones(len(positions), dtype=bool) + keep_mask[drop] = False + return positions[keep_mask], grain_ids[keep_mask] + + +def _resolve_device(device: str) -> str: + if device != "auto": + return device + if torch.cuda.is_available(): + try: + _ = torch.zeros(1, device="cuda") + return "cuda" + except Exception as exc: # noqa: BLE001 + warnings.warn(f"CUDA unusable ({exc!r}); falling back to CPU.") + return "cpu" + + +def _box_dims_for_atom_count(cell, n_basis: int, n_atoms_target: int): + """Return (counts, box) for a roughly-cubic periodic box of ~n_atoms_target. + + counts[d] cells along axis d, box[d] = counts[d] * cell[d]. The physical + box is kept as close to cubic as the cell allows so grains stay equiaxed. + """ + cell = _as_cell(cell) + v_cell = float(np.prod(cell)) + n_cells = max(1.0, n_atoms_target / n_basis) + target_edge = (n_cells * v_cell) ** (1.0 / 3.0) + counts = np.maximum(1, np.round(target_edge / cell).astype(int)) + box = counts.astype(np.float64) * cell + return counts, box + + +def generate_polycrystal( + cell, + basis, + n_atoms_target: int, + n_grains: int, + ani_thres: float = 0.0, + lloyd_iters: int = 2, + seed: int = 0, + device: str = "auto", + cull_factor: float = 0.85, + nn_distance: float | None = None, +) -> dict: + """Build a periodic, space-filling atomistic polycrystal for any lattice. + + Parameters + ---------- + cell : float or (3,) array + Orthorhombic conventional-cell edge lengths in Angstrom. A scalar + means a cubic cell. For HCP use a C-centred orthorhombic supercell, + e.g. cell = (a, a*sqrt(3), c) with `HCP_ORTHO_BASIS`. + basis : (n_basis, 3) array + Fractional atom coordinates within the conventional cell, in [0, 1). + n_atoms_target : int + Approximate atom count. The box is snapped to an integer number of + cells per axis, kept as close to cubic (physically) as possible. + n_grains : int + Number of grains in the polycrystal. + ani_thres : float + PyAPD anisotropy threshold; 0 = equiaxed grains, larger = elongated. + lloyd_iters : int + Number of Lloyd iterations (each calls find_optimal_W + adjust_X). + 0 = pure find_optimal_W with sampled seeds. + seed : int + Single integer seeding (a) PyAPD seed-point sampling, (b) anisotropy + sampling, (c) per-grain SO(3) rotation sampling. + device : str + "auto" | "cpu" | "cuda". + cull_factor : float + Cross-grain pair-cull cutoff as a fraction of the lattice nearest- + neighbour distance. Default 0.85 follows the Atomsk convention. + nn_distance : float, optional + Nearest-neighbour distance for the cull cutoff. If None, it is + computed from (cell, basis) via `min_interatomic_distance`. + + Returns + ------- + dict with keys: + positions : (N, 3) float ndarray in [0, box) + grain_ids : (N,) int ndarray, 0..n_grains-1 + box : (3,) float ndarray, per-axis edge lengths + cells : (3,) int ndarray, cells per axis + cell : (3,) float ndarray, conventional-cell edges used + basis : (n_basis, 3) float ndarray, fractional basis used + rotations : (n_grains, 3, 3) ndarray, per-grain rotations + seeds : (n_grains, 3) ndarray, grain seed positions + apd : apd_system, the optimised system (handy for debugging) + """ + cell = _as_cell(cell) + basis = np.asarray(basis, dtype=np.float64).reshape(-1, 3) + n_basis = len(basis) + device = _resolve_device(device) + if nn_distance is None: + nn_distance = min_interatomic_distance(cell, basis) + + counts, box = _box_dims_for_atom_count(cell, n_basis, n_atoms_target) + Lx, Ly, Lz = box + + # 1. Periodic APD over the (orthorhombic) box ---------------------------- + domain = torch.tensor([[0.0, Lx], [0.0, Ly], [0.0, Lz]], dtype=torch.float32) + apd = apd_system( + N=n_grains, + D=3, + domain=domain, + periodic=True, + ani_thres=ani_thres, + device=device, + seed=seed, + ) + apd.find_optimal_W(verbose=False) + if lloyd_iters > 0: + apd.Lloyds_algorithm(K=lloyd_iters, verbosity_level=0) + seeds = apd.X.detach().cpu().numpy() + + # 2. Per-grain rotations ------------------------------------------------- + rng = np.random.default_rng(seed) + rotations = np.stack([shoemake_uniform_so3(rng) for _ in range(n_grains)]) + + # 3. Per-grain atom generation (space-filling) --------------------------- + # Each grain is filled by the lattice block that is the PREIMAGE of the box + # under that grain's rotation, rotated forward WITHOUT wrapping. Rotating a + # finite box-sized lattice about the seed and wrapping by the box is NOT + # space-filling for a generic rotation: it produces seam planes through each + # grain with paired overlaps (later culled) AND matching gaps, leaving a + # void equal to the cull fraction distributed through grain interiors. The + # preimage-block fill below is dense everywhere; only true cross-grain GB + # overlaps (and the unavoidable periodic-boundary seam) are then culled. + box_corners = np.array( + [[x, y, z] for x in (0.0, Lx) for y in (0.0, Ly) for z in (0.0, Lz)] + ) + chunks_pos = [] + chunks_gid = [] + for g in range(n_grains): + X_g = seeds[g] + R_g = rotations[g] + # Inverse-rotate the 8 box corners to bound the lattice block to build. + src = (box_corners - X_g) @ R_g + X_g + lo = np.floor(src.min(axis=0) / cell).astype(int) - 2 + hi = np.ceil(src.max(axis=0) / cell).astype(int) + 2 + ii, jj, kk = np.meshgrid( + *(np.arange(lo[d], hi[d] + 1) for d in range(3)), indexing="ij" + ) + idx = np.stack([ii.ravel(), jj.ravel(), kk.ravel()], axis=1).astype(np.float64) + latt = ((idx[:, None, :] + basis[None, :, :]) * cell).reshape(-1, 3) + rot = (latt - X_g) @ R_g.T + X_g # forward rotation about seed, no wrap + rot = rot[np.all((rot >= 0.0) & (rot < box), axis=1)] # keep in-box + rot_t = torch.from_numpy(rot).to(device=apd.device, dtype=apd.dt) + grain_id = apd.grain_of(rot_t).cpu().numpy() + mask = grain_id == g + chunks_pos.append(rot[mask]) + chunks_gid.append(np.full(int(mask.sum()), g, dtype=np.int64)) + + positions = np.concatenate(chunks_pos, axis=0) + grain_ids = np.concatenate(chunks_gid, axis=0) + + # 4. PBC pair-cull -------------------------------------------------------- + positions, grain_ids = pbc_pair_cull( + positions, + grain_ids, + cutoff=cull_factor * nn_distance, + box=box, + cull_same_grain=True, + ) + + return { + "positions": positions, + "grain_ids": grain_ids, + "box": box, + "cells": counts, + "cell": cell, + "basis": basis, + "rotations": rotations, + "seeds": seeds, + "apd": apd, + } + + +def generate_bcc_polycrystal(a: float, *args, **kwargs) -> dict: + """Convenience wrapper: BCC polycrystal with cubic cell `a`.""" + return generate_polycrystal(a, BCC_BASIS, *args, **kwargs) + + +def generate_fcc_polycrystal(a: float, *args, **kwargs) -> dict: + """Convenience wrapper: FCC polycrystal with cubic cell `a`.""" + return generate_polycrystal(a, FCC_BASIS, *args, **kwargs) + + +def generate_sc_polycrystal(a: float, *args, **kwargs) -> dict: + """Convenience wrapper: simple-cubic polycrystal with cubic cell `a`.""" + return generate_polycrystal(a, SC_BASIS, *args, **kwargs) + + +def generate_hcp_polycrystal(a: float, c: float, *args, **kwargs) -> dict: + """Convenience wrapper: HCP polycrystal via a C-centred orthorhombic cell. + + Uses cell = (a, a*sqrt(3), c) with the 4-atom orthorhombic HCP basis. + """ + cell = (a, a * math.sqrt(3.0), c) + return generate_polycrystal(cell, HCP_ORTHO_BASIS, *args, **kwargs) + + +def write_lammps_data( + path: str, + positions: np.ndarray, + box, + atom_type: int = 1, + mass: float = 55.845, + title: str = "Polycrystal generated by PyAPD", +) -> None: + """Write atoms to a LAMMPS 'atomic' style data file. + + `box` may be a scalar (cubic) or a length-3 array of per-axis edge lengths. + """ + box = _as_cell(box) + n = len(positions) + lines = [ + title, + "", + f"{n} atoms", + "1 atom types", + "", + f"0.0 {box[0]:.10f} xlo xhi", + f"0.0 {box[1]:.10f} ylo yhi", + f"0.0 {box[2]:.10f} zlo zhi", + "", + "Masses", + "", + f"1 {mass}", + "", + "Atoms # atomic", + "", + ] + for i, (x, y, z) in enumerate(positions, start=1): + lines.append(f"{i} {atom_type} {x:.10f} {y:.10f} {z:.10f}") + lines.append("") + with open(path, "w") as fh: + fh.write("\n".join(lines)) + + +# Crystal presets for the CLI: name -> (build cell from args, basis). +_CLI_CRYSTALS = { + "sc": SC_BASIS, + "bcc": BCC_BASIS, + "fcc": FCC_BASIS, + "hcp": HCP_ORTHO_BASIS, +} + + +def _main(argv: list[str] | None = None) -> int: + import argparse + import time + + parser = argparse.ArgumentParser( + prog="python -m PyAPD.polycrystal_atomistic", + description="Generate a periodic atomistic polycrystal LAMMPS data file via APDs.", + ) + parser.add_argument( + "--crystal", + default="bcc", + choices=sorted(_CLI_CRYSTALS), + help="crystal structure (default: bcc)", + ) + parser.add_argument( + "--a", type=float, default=2.856, help="lattice parameter a (Angstrom)" + ) + parser.add_argument( + "--c", + type=float, + default=None, + help="c lattice parameter (hcp only; default 1.633*a)", + ) + parser.add_argument( + "--n-atoms", type=int, default=350_000, help="approximate target atom count" + ) + parser.add_argument("--n-grains", type=int, default=20) + parser.add_argument( + "--ani-thres", + type=float, + default=0.0, + help="0 = equiaxed; >0 = elongated grains", + ) + parser.add_argument("--lloyd", type=int, default=2, help="Lloyd iterations") + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--device", default="auto", choices=["auto", "cpu", "cuda"]) + parser.add_argument("--cull-factor", type=float, default=0.85) + parser.add_argument("--output", required=True, help="output LAMMPS data file path") + parser.add_argument("--element", default="Fe", help="element symbol for the title") + parser.add_argument("--mass", type=float, default=55.845) + args = parser.parse_args(argv) + + common = dict( + n_atoms_target=args.n_atoms, + n_grains=args.n_grains, + ani_thres=args.ani_thres, + lloyd_iters=args.lloyd, + seed=args.seed, + device=args.device, + cull_factor=args.cull_factor, + ) + + t0 = time.time() + if args.crystal == "hcp": + c = args.c if args.c is not None else 1.633 * args.a + result = generate_hcp_polycrystal(args.a, c, **common) + else: + result = generate_polycrystal(args.a, _CLI_CRYSTALS[args.crystal], **common) + elapsed = time.time() - t0 + + write_lammps_data( + path=args.output, + positions=result["positions"], + box=result["box"], + atom_type=1, + mass=args.mass, + title=( + f"{args.crystal.upper()} {args.element} polycrystal | a={args.a} | " + f"N={args.n_grains} grains | seed={args.seed}" + ), + ) + + box = result["box"] + print(f"wrote {args.output}") + print(f" crystal : {args.crystal} ({len(result['basis'])} atoms/cell)") + print( + f" box edges : {box[0]:.3f} x {box[1]:.3f} x {box[2]:.3f} Angstrom " + f"(cells {tuple(int(v) for v in result['cells'])})" + ) + print(f" N atoms : {len(result['positions'])} (target {args.n_atoms})") + print(f" grains : {args.n_grains}") + print(f" ani_thres : {args.ani_thres}") + print(f" elapsed : {elapsed:.1f} s") + return 0 + + +if __name__ == "__main__": + raise SystemExit(_main()) diff --git a/notebooks/tutorials/periodic_polycrystal_bcc_fe.ipynb b/notebooks/tutorials/periodic_polycrystal_bcc_fe.ipynb new file mode 100644 index 0000000..a990911 --- /dev/null +++ b/notebooks/tutorials/periodic_polycrystal_bcc_fe.ipynb @@ -0,0 +1,168 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "7e3f3d8e", + "metadata": {}, + "source": [ + "# Periodic polycrystal generation via PyAPD\n", + "\n", + "End-to-end demo: build a ~350 000-atom periodic polycrystal with 20 curved,\n", + "anisotropy-aware grains and write it to a LAMMPS-readable data file.\n", + "\n", + "The generator is crystal-agnostic. We use BCC-Fe here; FCC / simple-cubic /\n", + "HCP and any custom (cell, basis) lattice work the same way (see the last cell)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "045a3872", + "metadata": {}, + "outputs": [], + "source": [ + "import time\n", + "import numpy as np\n", + "from PyAPD.polycrystal_atomistic import generate_bcc_polycrystal, write_lammps_data\n", + "\n", + "a = 2.856 # BCC Fe lattice parameter [Å]\n", + "N_TARGET = 350_000\n", + "N_GRAINS = 20\n", + "\n", + "t0 = time.time()\n", + "result = generate_bcc_polycrystal(\n", + " a=a,\n", + " n_atoms_target=N_TARGET,\n", + " n_grains=N_GRAINS,\n", + " ani_thres=0.0,\n", + " lloyd_iters=2,\n", + " seed=0,\n", + " device=\"auto\",\n", + ")\n", + "box = result[\"box\"] # (3,) per-axis edge lengths [Å]\n", + "print(f\"built {len(result['positions'])} atoms in {time.time()-t0:.1f}s \"\n", + " f\"(box = {box[0]:.1f} x {box[1]:.1f} x {box[2]:.1f} Å)\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6a08f9f7", + "metadata": {}, + "outputs": [], + "source": [ + "write_lammps_data(\n", + " path=\"bcc_fe_polycrystal.data\",\n", + " positions=result[\"positions\"],\n", + " box=result[\"box\"],\n", + " atom_type=1,\n", + " mass=55.845,\n", + " title=\"BCC Fe polycrystal | a=2.856 | 20 grains | curved GBs\",\n", + ")\n", + "print(\"wrote bcc_fe_polycrystal.data\")" + ] + }, + { + "cell_type": "markdown", + "id": "b52a6f3a", + "metadata": {}, + "source": [ + "## Visual check: grain-id slice\n", + "\n", + "A thin slab at z = Lz/2 to visualize the grain labels in 2D." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b3045082", + "metadata": {}, + "outputs": [], + "source": [ + "import matplotlib.pyplot as plt\n", + "\n", + "box = result[\"box\"]\n", + "pos = result[\"positions\"]\n", + "mask = np.abs(pos[:, 2] - box[2] / 2) < 5.0\n", + "plt.figure(figsize=(8, 8))\n", + "plt.scatter(\n", + " pos[mask, 0], pos[mask, 1], c=result[\"grain_ids\"][mask], s=1.5, cmap=\"tab20\"\n", + ")\n", + "plt.xlim(0, box[0]); plt.ylim(0, box[1])\n", + "plt.gca().set_aspect(\"equal\")\n", + "plt.title(\"Atoms in z=Lz/2 slab, coloured by grain\")\n", + "plt.savefig(\"polycrystal_slice.png\", dpi=120)" + ] + }, + { + "cell_type": "markdown", + "id": "5f7a31ec", + "metadata": {}, + "source": [ + "## LAMMPS relaxation (optional)\n", + "\n", + "Save as `in.relax` and run with an EAM Fe potential\n", + "(e.g. NIST `Fe_2.eam.fs`, which you supply):\n", + "\n", + "```text\n", + "units metal\n", + "atom_style atomic\n", + "boundary p p p\n", + "read_data bcc_fe_polycrystal.data\n", + "pair_style eam/fs\n", + "pair_coeff * * Fe_2.eam.fs Fe\n", + "thermo 10\n", + "minimize 1e-6 1e-8 200 1000\n", + "write_data bcc_fe_polycrystal.min.data\n", + "```\n", + "\n", + "Forces start finite and the energy decreases monotonically to ~0.05-0.1 eV/atom\n", + "above bulk BCC-Fe cohesive after relaxation (the grain-boundary excess energy)." + ] + }, + { + "cell_type": "markdown", + "id": "c1a2b3d4", + "metadata": {}, + "source": [ + "## Other crystal structures\n", + "\n", + "The same driver builds any lattice. Cubic crystals use a scalar cell; HCP /\n", + "tetragonal / orthorhombic use a per-axis cell. Convenience wrappers:\n", + "\n", + "```python\n", + "from PyAPD.polycrystal_atomistic import (\n", + " generate_fcc_polycrystal, generate_sc_polycrystal, generate_hcp_polycrystal,\n", + " generate_polycrystal,\n", + ")\n", + "\n", + "fcc = generate_fcc_polycrystal(a=3.615, n_atoms_target=200_000, n_grains=12) # e.g. Cu\n", + "hcp = generate_hcp_polycrystal(a=3.21, c=5.21, n_atoms_target=200_000, n_grains=12) # e.g. Mg\n", + "\n", + "# Fully custom lattice: orthorhombic cell + fractional basis.\n", + "import numpy as np\n", + "custom = generate_polycrystal(\n", + " cell=(3.0, 3.0, 5.0),\n", + " basis=np.array([[0.0, 0.0, 0.0], [0.5, 0.5, 0.5]]),\n", + " n_atoms_target=100_000, n_grains=8,\n", + ")\n", + "```\n", + "\n", + "The cull cutoff defaults to `0.85 x` the lattice nearest-neighbour distance\n", + "(computed automatically); pass `nn_distance=` to override." + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/tests/test_apds_periodic.py b/tests/test_apds_periodic.py new file mode 100644 index 0000000..231025b --- /dev/null +++ b/tests/test_apds_periodic.py @@ -0,0 +1,117 @@ +"""Tests for the periodic (minimum-image) distance support in apd_system.""" + +import pytest +import torch + +from PyAPD import apd_system + + +def test_periodic_flag_accepted(): + """apd_system must accept periodic=True without raising.""" + apd = apd_system(N=4, D=3, periodic=True, seed=0, device="cpu") + assert apd.periodic is True + + +def test_non_periodic_is_default(): + """periodic defaults to False (backward-compatible).""" + apd = apd_system(N=4, D=2, seed=0, device="cpu") + assert apd.periodic is False + + +def test_periodic_single_seed_uniform_labels(): + """A single seed in a periodic box must label every voxel as grain 0, + regardless of where the seed sits (corner, centre, edge).""" + L = 1.0 + for seed_pos in [[0.0, 0.0, 0.0], [0.5, 0.5, 0.5], [0.99, 0.01, 0.5]]: + X = torch.tensor([seed_pos]) + apd = apd_system( + X=X, + D=3, + domain=torch.tensor([[0.0, L], [0.0, L], [0.0, L]]), + periodic=True, + pixel_params=(16, 16, 16), + device="cpu", + ) + labels = apd.assemble_apd().cpu() + assert torch.all(labels == 0), ( + f"seed at {seed_pos} produced labels {labels.unique()}" + ) + + +def test_periodic_grain_wraps_across_seam(): + """Two seeds at x = 0.0 and x = 0.5 (y, z = 0.5). Under periodicity, + grain 0 (seed at the seam) owns x in [0, 0.25] u [0.75, 1.0] and grain 1 + owns x in [0.25, 0.75]. So the columns at x=0 and x=L (column 0 and column + M-1) must BOTH belong to grain 0 -- the grain wraps across the periodic + seam. Without periodicity, column M-1 would be closer to seed 1 (x=0.5) + than to seed 0 (x=0.0), so it would belong to grain 1.""" + L = 1.0 + M = 32 + X = torch.tensor([[0.0, 0.5, 0.5], [0.5, 0.5, 0.5]]) + apd = apd_system( + X=X, + D=3, + domain=torch.tensor([[0.0, L], [0.0, L], [0.0, L]]), + periodic=True, + pixel_params=(M, 16, 16), + device="cpu", + ) + labels = apd.assemble_apd().reshape(M, 16, 16).cpu() + # column 0 and column M-1 (both ~0.016 from the seam at x=0) must be in + # the same grain under periodicity, and that grain is grain 0. + assert torch.all(labels[0] == 0) + assert torch.all(labels[-1] == 0) + # and the middle column (x ~ 0.5) must be in grain 1 + assert torch.all(labels[M // 2] == 1) + + +def test_grain_of_matches_assemble_apd_on_voxel_centres(): + """grain_of() on the pixel centres must reproduce assemble_apd()'s output + exactly -- same metric, same data, same result.""" + L = 1.0 + apd = apd_system( + N=5, + D=3, + domain=torch.tensor([[0.0, L], [0.0, L], [0.0, L]]), + periodic=True, + pixel_params=(16, 16, 16), + seed=42, + device="cpu", + ) + apd.find_optimal_W(verbose=False) + voxel_labels = apd.assemble_apd().cpu() + # apd.Y holds the voxel centres (filled by assemble_pixels via assemble_apd). + point_labels = apd.grain_of(apd.Y).cpu() + assert torch.equal(voxel_labels, point_labels) + + +def test_grain_of_continuous_points(): + """grain_of must accept arbitrary (M, 3) torch tensors in the box.""" + L = 1.0 + apd = apd_system( + N=4, + D=3, + domain=torch.tensor([[0.0, L], [0.0, L], [0.0, L]]), + periodic=True, + seed=7, + device="cpu", + ) + apd.find_optimal_W(verbose=False) + pts = torch.tensor( + [ + [0.5, 0.5, 0.5], + [0.0, 0.0, 0.0], + [0.99, 0.01, 0.5], + [0.1, 0.9, 0.7], + ] + ) + labels = apd.grain_of(pts) + assert labels.shape == (4,) + assert int(labels.min()) >= 0 and int(labels.max()) < 4 + + +def test_grain_of_rejects_wrong_shape(): + """grain_of must validate the (M, D) shape of its input.""" + apd = apd_system(N=4, D=3, periodic=True, seed=1, device="cpu") + with pytest.raises(ValueError): + apd.grain_of(torch.zeros(5, 2)) # D=2 points into a D=3 system diff --git a/tests/test_polycrystal_atomistic.py b/tests/test_polycrystal_atomistic.py new file mode 100644 index 0000000..affe779 --- /dev/null +++ b/tests/test_polycrystal_atomistic.py @@ -0,0 +1,280 @@ +"""Unit tests for the polycrystal_atomistic module.""" + +import math + +import numpy as np +import pytest +from scipy.spatial import cKDTree + +from PyAPD.polycrystal_atomistic import ( + BCC_BASIS, + FCC_BASIS, + HCP_ORTHO_BASIS, + SC_BASIS, + bulk_bcc_positions, + bulk_lattice_positions, + min_interatomic_distance, + pbc_pair_cull, + shoemake_uniform_so3, +) + + +# -------------------------------------------------------------------------- +# Reference lattice builders +# -------------------------------------------------------------------------- +def test_bulk_bcc_positions_count_and_bounds(): + a = 2.856 + m = 3 + pos = bulk_bcc_positions(a, m) + assert pos.shape == (2 * m * m * m, 3), f"expected (54,3), got {pos.shape}" + L = m * a + assert np.all(pos >= 0.0) + assert np.all(pos < L) + + +def test_bulk_bcc_positions_both_sublattices(): + """BCC has the corner sublattice at (0,0,0) and the body-centre at + (a/2, a/2, a/2). Both must appear in m=2.""" + a = 2.856 + pos = bulk_bcc_positions(a, m=2) + assert np.any(np.all(np.isclose(pos, [0.0, 0.0, 0.0]), axis=1)) + assert np.any(np.all(np.isclose(pos, [a / 2, a / 2, a / 2]), axis=1)) + + +def test_bulk_lattice_positions_fcc_count(): + a = 3.6 + pos = bulk_lattice_positions(a, FCC_BASIS, 2) + assert pos.shape == (4 * 8, 3) + assert np.all(pos >= 0.0) and np.all(pos < 2 * a) + + +def test_bulk_lattice_positions_orthorhombic_counts_and_box(): + """Per-axis cell and per-axis counts: box is counts * cell on each axis.""" + cell = (2.0, 3.0, 4.0) + counts = (2, 3, 1) + pos = bulk_lattice_positions(cell, SC_BASIS, counts) + assert pos.shape == (1 * 2 * 3 * 1, 3) + box = np.asarray(counts) * np.asarray(cell) + assert np.all(pos[:, 0] < box[0]) + assert np.all(pos[:, 1] < box[1]) + assert np.all(pos[:, 2] < box[2]) + + +# -------------------------------------------------------------------------- +# Nearest-neighbour distance +# -------------------------------------------------------------------------- +def test_min_interatomic_distance_cubic(): + a = 2.856 + assert math.isclose(min_interatomic_distance(a, SC_BASIS), a, rel_tol=1e-9) + assert math.isclose( + min_interatomic_distance(a, BCC_BASIS), a * math.sqrt(3) / 2, rel_tol=1e-9 + ) + assert math.isclose( + min_interatomic_distance(a, FCC_BASIS), a / math.sqrt(2), rel_tol=1e-9 + ) + + +def test_min_interatomic_distance_orthorhombic_sc(): + """For simple-cubic on an orthorhombic cell the NN is the shortest edge.""" + cell = (2.0, 3.0, 5.0) + assert math.isclose(min_interatomic_distance(cell, SC_BASIS), 2.0, rel_tol=1e-9) + + +def test_min_interatomic_distance_hcp_is_a(): + """Ideal HCP nearest-neighbour distance is a (in-plane).""" + a, c = 3.2, 3.2 * 1.633 + cell = (a, a * math.sqrt(3), c) + assert math.isclose( + min_interatomic_distance(cell, HCP_ORTHO_BASIS), a, rel_tol=1e-6 + ) + + +# -------------------------------------------------------------------------- +# Rotations +# -------------------------------------------------------------------------- +def test_shoemake_uniform_so3_is_rotation(): + rng = np.random.default_rng(0) + R = shoemake_uniform_so3(rng) + assert R.shape == (3, 3) + np.testing.assert_allclose(R.T @ R, np.eye(3), atol=1e-6) + assert np.isclose(np.linalg.det(R), 1.0, atol=1e-6) + + +def test_shoemake_uniform_so3_distribution_haar(): + """1000 uniform samples -- mean rotation matrix should be ~ 0 (Haar mean).""" + rng = np.random.default_rng(1) + Rs = np.stack([shoemake_uniform_so3(rng) for _ in range(1000)]) + mean = Rs.mean(axis=0) + assert np.all(np.abs(mean) < 0.1), f"mean rotation has |entries| > 0.1: {mean}" + + +# -------------------------------------------------------------------------- +# PBC pair cull +# -------------------------------------------------------------------------- +def test_pbc_pair_cull_removes_only_cross_grain_overlaps(): + """Cross-grain overlaps are dropped, same-grain pairs kept; PBC respected.""" + L = 10.0 + pos = np.array( + [ + [1.0, 1.0, 1.0], # 0 + [1.4, 1.0, 1.0], # 1 (0.4 from 0, same grain -> keep) + [3.0, 1.0, 1.0], # 2 (1.6 from 1, far -> keep) + [3.4, 1.0, 1.0], # 3 (0.4 from 2, diff grain -> cull) + [0.1, 5.0, 5.0], # 4 (close to 5 across PBC seam, diff grain -> cull) + [9.9, 5.0, 5.0], # 5 + ] + ) + gids = np.array([0, 0, 0, 1, 0, 1]) + kept_pos, kept_gids = pbc_pair_cull(pos, gids, cutoff=0.5, box=L) + assert len(kept_pos) == 4, ( + f"expected 4 kept atoms, got {len(kept_pos)}: {kept_gids}" + ) + for idx in [0, 1, 2]: + assert np.any(np.all(np.isclose(kept_pos, pos[idx]), axis=1)) + + +def test_pbc_pair_cull_per_axis_box(): + """pbc_pair_cull accepts a per-axis box (orthorhombic PBC).""" + box = np.array([4.0, 6.0, 8.0]) + # two atoms close across the x-seam, different grains + pos = np.array([[0.1, 2.0, 2.0], [3.9, 2.0, 2.0], [2.0, 3.0, 4.0]]) + gids = np.array([0, 1, 0]) + kept_pos, _ = pbc_pair_cull(pos, gids, cutoff=0.5, box=box) + assert len(kept_pos) == 2 + + +# -------------------------------------------------------------------------- +# Box sizing +# -------------------------------------------------------------------------- +def test_box_dims_cubic_matches_reference(): + from PyAPD.polycrystal_atomistic import _box_dims_for_atom_count + + counts, box = _box_dims_for_atom_count(2.856, n_basis=2, n_atoms_target=250) + np.testing.assert_array_equal(counts, [5, 5, 5]) + np.testing.assert_allclose(box, 5 * 2.856) + + +def test_box_dims_orthorhombic_roughly_cubic(): + from PyAPD.polycrystal_atomistic import _box_dims_for_atom_count + + cell = (2.0, 4.0, 8.0) + counts, box = _box_dims_for_atom_count(cell, n_basis=1, n_atoms_target=4096) + # box should be as close to cubic (physically) as the cell allows + assert np.all(counts >= 1) + assert box.max() / box.min() < 1.6 # near-cubic physical box + # atom count is in the right ballpark + assert 0.5 * 4096 <= int(np.prod(counts)) <= 2.0 * 4096 + + +# -------------------------------------------------------------------------- +# End-to-end generation (needs KeOps; CPU-only here) +# -------------------------------------------------------------------------- +def _assert_valid_polycrystal(result, n_grains): + pos = result["positions"] + gids = result["grain_ids"] + box = np.asarray(result["box"]) + assert len(pos) == len(gids) and len(pos) > 0 + assert np.all(pos >= 0.0) and np.all(pos < box) + assert set(np.unique(gids)) == set(range(n_grains)) + # No residual overlaps below the cull cutoff (PBC). + cutoff = 0.85 * min_interatomic_distance(result["cell"], result["basis"]) + tree = cKDTree(pos, boxsize=box) + nn_dist, _ = tree.query(pos, k=2) + assert nn_dist[:, 1].min() >= cutoff - 1e-6 + + +@pytest.mark.parametrize( + "builder,kwargs,n_basis", + [ + ("generate_bcc_polycrystal", dict(a=2.856), 2), + ("generate_fcc_polycrystal", dict(a=3.6), 4), + ("generate_sc_polycrystal", dict(a=3.0), 1), + ], +) +def test_generate_cubic_polycrystals(builder, kwargs, n_basis): + import PyAPD.polycrystal_atomistic as pa + + n_grains = 4 + result = getattr(pa, builder)( + n_atoms_target=200, + n_grains=n_grains, + ani_thres=0.0, + lloyd_iters=0, + seed=0, + device="cpu", + **kwargs, + ) + _assert_valid_polycrystal(result, n_grains) + assert len(result["basis"]) == n_basis + + +def test_generate_hcp_polycrystal_orthorhombic(): + """HCP exercises the orthorhombic (non-cubic) cell path.""" + from PyAPD.polycrystal_atomistic import generate_hcp_polycrystal + + n_grains = 3 + result = generate_hcp_polycrystal( + a=3.2, + c=3.2 * 1.633, + n_atoms_target=200, + n_grains=n_grains, + lloyd_iters=0, + seed=1, + device="cpu", + ) + _assert_valid_polycrystal(result, n_grains) + # non-cubic box: the three edges are not all equal + box = np.asarray(result["box"]) + assert not np.allclose(box, box[0]) + + +def test_generate_polycrystal_deterministic(): + """Same seed -> identical positions and grain_ids.""" + from PyAPD.polycrystal_atomistic import generate_polycrystal + + kwargs = dict( + cell=2.856, + basis=BCC_BASIS, + n_atoms_target=120, + n_grains=3, + ani_thres=0.0, + lloyd_iters=0, + seed=7, + device="cpu", + ) + a_run = generate_polycrystal(**kwargs) + b_run = generate_polycrystal(**kwargs) + np.testing.assert_array_equal(a_run["positions"], b_run["positions"]) + np.testing.assert_array_equal(a_run["grain_ids"], b_run["grain_ids"]) + + +def test_write_lammps_data_roundtrip(tmp_path): + """Write a tiny polycrystal, read it back with ASE, compare (orthorhombic box).""" + pytest.importorskip("ase") + from ase.io import read as ase_read + + from PyAPD.polycrystal_atomistic import generate_hcp_polycrystal, write_lammps_data + + result = generate_hcp_polycrystal( + a=3.2, + c=3.2 * 1.633, + n_atoms_target=200, + n_grains=3, + lloyd_iters=0, + seed=0, + device="cpu", + ) + out = tmp_path / "tiny.data" + write_lammps_data( + path=str(out), + positions=result["positions"], + box=result["box"], + atom_type=1, + mass=47.867, + title="HCP tiny polycrystal test", + ) + atoms = ase_read(str(out), format="lammps-data", style="atomic") + assert len(atoms) == len(result["positions"]) + np.testing.assert_allclose(atoms.cell.lengths(), result["box"], atol=1e-6) + assert np.all(atoms.positions >= -1e-6) + assert np.all(atoms.positions < np.asarray(result["box"]) + 1e-6)