# Copyright (c) 2022-2026, Yongchao Wu in Aalto University
# This file is from the mdapy project, released under the BSD 3-Clause License.
from typing import Tuple, Optional, Union, Iterable
import numpy as np
import polars as pl
from mdapy.voronoi import Container, Cell
from mdapy.system import System
from mdapy.build_lattice import build_crystal
from mdapy.box import Box
from mdapy import _polycrystal, _neighbor
from mdapy.parallel import get_num_threads
from mdapy.tool_function import _replicate_pos
from time import time
[docs]
class CreatePolycrystal:
"""
Create polycrystalline structures using Voronoi tessellation method.
This class generates polycrystalline atomic structures by:
1. Creating Voronoi cells from seed positions
2. Filling each cell with rotated unit cell atoms
3. Optionally decorating grain boundaries with graphene layers
4. Removing overlapping atoms based on distance criteria
Parameters
----------
unitcell : System
The unit cell structure to be replicated in each grain.
box : Union[int, float, Iterable[float], np.ndarray, Box]
Simulation box dimensions. Can be a single value (cubic),
array of 3 values, or Box object.
seed_number : int
Number of grains (Voronoi seeds) to generate.
seed_position : Optional[np.ndarray], default=None
Positions of Voronoi seeds with shape (seed_number, 3).
If None, randomly generated within the box.
theta_list : Optional[np.ndarray], default=None
Rotation angles (in degrees) for each grain with shape (seed_number, 3).
Represents rotations around x, y, z axes. If None, randomly generated.
randomseed : Optional[int], default=None
Random seed for reproducibility. If None, randomly chosen.
metal_overlap_dis : Optional[float], default=None
Minimum allowed distance between metal atoms (Angstrom).
Atoms closer than this will be removed.
add_graphene : bool, default=False
Whether to add graphene layers at grain boundaries.
metal_gra_overlap_dis : float, default=3.0
Minimum allowed distance between metal and carbon atoms (Angstrom).
face_threshold : float, default=0.0
Minimum face area (Angstrom²) for graphene decoration.
Faces smaller than this will be skipped.
need_rotation : bool, default=True
Whether to apply random rotations to grains.
Attributes
----------
unitcell : System
The unit cell structure.
box : Box
The simulation box.
seed_number : int
Number of grains.
seed_position : np.ndarray
Voronoi seed positions.
theta_list : np.ndarray
Rotation angles for each grain.
con : Container
Voronoi container with all cells.
randomseed : int
Random seed used.
Examples
--------
>>> from mdapy.build_lattice import build_crystal
>>> unit = build_crystal("Al", "fcc", 4.05)
>>> poly = CreatePolycrystal(unit, box=100, seed_number=10, metal_overlap_dis=2.0)
>>> system = poly.compute()
>>> system.write_xyz("polycrystal.xyz")
Notes
-----
- Triclinic boxes are not supported
- Free boundary conditions are not supported
"""
def __init__(
self,
unitcell: System,
box: Union[int, float, Iterable[float], np.ndarray, Box],
seed_number: int,
seed_position: Optional[np.ndarray] = None,
theta_list: Optional[np.ndarray] = None,
randomseed: Optional[int] = None,
metal_overlap_dis: Optional[float] = None,
add_graphene: bool = False,
metal_gra_overlap_dis: float = 3.0,
face_threshold: float = 0.0,
need_rotation: bool = True,
) -> None:
"""Initialize CreatePolycrystal with given parameters."""
self.unitcell = unitcell
self.box = Box(box)
if sum(self.box.boundary) != 3:
raise ValueError("Free boundary condition is not supported.")
if self.box.triclinic:
raise ValueError("Triclinic box is not supported")
self.seed_number = seed_number
self.metal_overlap_dis = metal_overlap_dis
self.add_graphene = add_graphene
self.metal_gra_overlap_dis = metal_gra_overlap_dis
self.need_rotation = need_rotation
self.face_threshold = face_threshold
# Initialize random number generator
if randomseed is None:
randomseed = np.random.randint(0, 1_000_000_000)
self.randomseed = int(randomseed)
self.rng = np.random.default_rng(self.randomseed)
# Initialize seed positions
if seed_position is None:
box_lengths = np.diag(self.box.box)
self.seed_position = self.rng.random((self.seed_number, 3)) * box_lengths
else:
if seed_position.shape != (self.seed_number, 3):
raise ValueError(
f"seed_position shape must be ({self.seed_number}, 3), "
f"got {seed_position.shape}"
)
self.seed_position = seed_position
# Initialize rotation angles
if theta_list is None:
self.theta_list = self.rng.uniform(-180, 180, (self.seed_number, 3))
else:
if theta_list.shape != (self.seed_number, 3):
raise ValueError(
f"theta_list shape must be ({self.seed_number}, 3), "
f"got {theta_list.shape}"
)
self.theta_list = theta_list
# Voronoi container (will be created in compute())
self.con: Optional[Container] = None
@staticmethod
def _get_rotation_matrix(
theta_deg: float, axis_tuple: Tuple[float, float, float]
) -> np.ndarray:
"""
Calculate rotation matrix using axis-angle representation.
Uses Rodrigues' rotation formula to compute the rotation matrix
that rotates vectors by theta_deg degrees around the given axis.
Parameters
----------
theta_deg : float
Rotation angle in degrees.
axis_tuple : Tuple[float, float, float]
Rotation axis as (x, y, z) tuple.
Returns
-------
np.ndarray
3x3 rotation matrix.
Raises
------
ValueError
If the rotation axis has zero length.
Notes
-----
The rotation matrix R is computed as:
R = I + sin(θ)K + (1-cos(θ))K²
where K is the skew-symmetric cross-product matrix of the normalized axis.
"""
theta = np.radians(theta_deg)
axis = np.array(axis_tuple, dtype=float)
norm = np.linalg.norm(axis)
if norm == 0:
raise ValueError("Rotation axis must be non-zero")
x, y, z = axis / norm
c = np.cos(theta)
s = np.sin(theta)
C = 1 - c
return np.array(
[
[c + C * x * x, C * x * y - s * z, C * x * z + s * y],
[C * y * x + s * z, c + C * y * y, C * y * z - s * x],
[C * z * x - s * y, C * z * y + s * x, c + C * z * z],
],
dtype=float,
)
def _get_plane_equation_coeffs_for_cell(self, cell: Cell) -> np.ndarray:
"""
Calculate plane equation coefficients for all faces of a Voronoi cell.
For each face, computes the plane equation: ax + by + cz + d = 0
where the normal vector (a, b, c) points inward to the cell.
Parameters
----------
cell : Cell
Voronoi cell object containing vertices and face information.
Returns
-------
np.ndarray
Array of shape (n_faces, 4) where each row is [a, b, c, d].
Raises
------
ValueError
If face vertices are degenerate (collinear).
Notes
-----
The normal vector is computed using cross product of two edge vectors,
then normalized. The direction is adjusted to point inward.
"""
n_faces = len(cell.face_vertices)
face_index = np.array([face[:3] for face in cell.face_vertices])
plane_pos = np.array([cell.vertices[face_index[i]] for i in range(n_faces)])
coeffs = np.zeros((n_faces, 4))
for i in range(n_faces):
p1, p2, p3 = plane_pos[i]
# Calculate normal vector using cross product
v1, v2 = p2 - p1, p3 - p1
n = np.cross(v1, v2)
norm_n = np.linalg.norm(n)
if norm_n < 1e-10:
raise ValueError(f"Degenerate face vertices at face {i}")
n = n / norm_n
d = -np.dot(n, p1)
# Ensure normal points inward to the cell
if np.dot(n, cell.pos) + d > 0:
n, d = -n, -d
coeffs[i, :3] = n
coeffs[i, 3] = d
return coeffs
def _generate_grain_atoms(
self, grain_idx: int, cell: Cell, data: pl.DataFrame, coeffs: np.ndarray
) -> np.ndarray:
"""
Generate atoms for a single grain by rotating and filtering unit cell.
Process:
1. Apply rotation matrices (Rx * Ry * Rz) to unit cell
2. Translate to grain center
3. Filter atoms inside Voronoi cell using plane equations
Parameters
----------
grain_idx : int
Index of the grain (0 to seed_number-1).
cell : Cell
Voronoi cell for this grain.
data : pl.DataFrame
Replicated unit cell data with columns ['x', 'y', 'z'].
coeffs : np.ndarray
Plane equation coefficients for cell faces.
Returns
-------
np.ndarray
Filtered atomic positions with shape (n_atoms, 3).
"""
if self.need_rotation:
R_x = self._get_rotation_matrix(
self.theta_list[grain_idx, 0], (1.0, 0.0, 0.0)
)
R_y = self._get_rotation_matrix(
self.theta_list[grain_idx, 1], (0.0, 1.0, 0.0)
)
R_z = self._get_rotation_matrix(
self.theta_list[grain_idx, 2], (0.0, 0.0, 1.0)
)
rotate_matrix = R_x @ R_y @ R_z
else:
rotate_matrix = self._get_rotation_matrix(0, (1.0, 0.0, 0.0))
pos_center = data.select("x", "y", "z").mean().to_numpy().ravel()
filtered_pos = _polycrystal.transform_and_filter(
data["x"].to_numpy(allow_copy=False),
data["y"].to_numpy(allow_copy=False),
data["z"].to_numpy(allow_copy=False),
rotate_matrix,
pos_center,
cell.pos,
coeffs,
get_num_threads(),
)
return filtered_pos
def _generate_gra_atoms(
self, cell: Cell, gra_data: pl.DataFrame, coeffs: np.ndarray
) -> np.ndarray:
"""
Generate graphene atoms on grain boundary faces.
Strategy:
1. For each face larger than face_threshold
2. Rotate graphene sheet to align with face normal
3. Translate to face center
4. Filter atoms inside face polygon
Parameters
----------
cell : Cell
Voronoi cell containing face information.
gra_data : pl.DataFrame
Graphene sheet data with columns ['x', 'y', 'z'].
coeffs : np.ndarray
Plane equation coefficients for cell faces.
Returns
-------
np.ndarray
Graphene atomic positions with shape (n_atoms, 3).
Raises
------
AssertionError
If no graphene atoms are generated for any face.
Notes
-----
- Initial graphene normal is [0, 0, 1] (xy-plane)
- Uses 2D projection and ray-casting for polygon filtering
- Z-tolerance of 0.5 Å is used for face proximity
"""
gra_normal = np.array([0.0, 0.0, 1.0]) # Initial graphene normal (xy-plane)
gra_positions_list = []
gra_pos = gra_data.select("x", "y", "z").to_numpy()
for face_idx in range(coeffs.shape[0]):
# Skip faces smaller than threshold
if cell.face_areas[face_idx] <= self.face_threshold:
continue
face_vertices = cell.vertices[cell.face_vertices[face_idx]]
face_normal = coeffs[face_idx, :3]
face_normal = face_normal / np.linalg.norm(face_normal)
face_center = np.mean(face_vertices, axis=0)
# Step 1: Compute rotation matrix to align graphene with face
rotation_matrix = self._compute_rotation_to_align_normals(
gra_normal, face_normal
)
# Step 2: Rotate graphene and translate to face center
rotated_gra_pos = gra_pos @ rotation_matrix.T
gra_centered = rotated_gra_pos - np.mean(rotated_gra_pos, axis=0)
gra_on_face = gra_centered + face_center
# Step 3: Filter atoms inside face polygon
atoms_in_face = self._filter_atoms_in_polygon(
gra_on_face, face_vertices, face_normal
)
if len(atoms_in_face) > 0:
gra_positions_list.append(atoms_in_face)
assert len(gra_positions_list) > 0, "No graphene atoms generated"
return np.vstack(gra_positions_list)
def _compute_rotation_to_align_normals(
self, source_normal: np.ndarray, target_normal: np.ndarray
) -> np.ndarray:
"""
Compute rotation matrix to align source normal with target normal.
Uses Rodrigues' rotation formula to find the rotation that maps
source_normal to target_normal.
Parameters
----------
source_normal : np.ndarray
Initial normal vector (will be normalized).
target_normal : np.ndarray
Target normal vector (will be normalized).
Returns
-------
np.ndarray
3x3 rotation matrix.
Notes
-----
Special cases:
- If normals are already aligned: returns identity matrix
- If normals are opposite: rotates 180° around a perpendicular axis
- General case: rotates around cross product of normals
"""
# Normalize
v1 = source_normal / np.linalg.norm(source_normal)
v2 = target_normal / np.linalg.norm(target_normal)
# Check if already aligned
dot_product = np.dot(v1, v2)
if np.isclose(dot_product, 1.0, atol=1e-6):
# Already aligned, return identity
return np.eye(3)
if np.isclose(dot_product, -1.0, atol=1e-6):
# Opposite directions, rotate 180° around perpendicular axis
if abs(v1[0]) < 0.9:
perp_axis = np.array([1.0, 0.0, 0.0])
else:
perp_axis = np.array([0.0, 1.0, 0.0])
rotation_axis = np.cross(v1, perp_axis)
rotation_axis = rotation_axis / np.linalg.norm(rotation_axis)
return self._get_rotation_matrix(180.0, tuple(rotation_axis))
# General case: compute rotation axis and angle
rotation_axis = np.cross(v1, v2)
rotation_axis = rotation_axis / np.linalg.norm(rotation_axis)
# Rotation angle
angle_rad = np.arccos(np.clip(dot_product, -1.0, 1.0))
angle_deg = np.degrees(angle_rad)
return self._get_rotation_matrix(angle_deg, tuple(rotation_axis))
def _filter_atoms_in_polygon(
self, points: np.ndarray, polygon_vertices: np.ndarray, face_normal: np.ndarray
) -> np.ndarray:
"""
Filter atoms that lie inside a 3D polygon face.
Method:
1. Build local coordinate system with face normal as z-axis
2. Project polygon and points to 2D plane
3. Use ray-casting algorithm to test if points are inside polygon
4. Check z-coordinate tolerance to ensure points are near face
Parameters
----------
points : np.ndarray
Points to test with shape (n_points, 3).
polygon_vertices : np.ndarray
Polygon vertices with shape (n_vertices, 3).
face_normal : np.ndarray
Normal vector of the face.
Returns
-------
np.ndarray
Filtered points inside polygon with shape (n_filtered, 3).
Notes
-----
- Z-tolerance is set to 0.5 Angstrom
- Uses vectorized ray-casting for efficiency
"""
# Build local coordinate system
# z-axis: face normal
local_z = face_normal / np.linalg.norm(face_normal)
# x-axis: from face center to first vertex (projected on face)
face_center = np.mean(polygon_vertices, axis=0)
temp_x = polygon_vertices[0] - face_center
temp_x = temp_x - np.dot(temp_x, local_z) * local_z # Project onto face
if np.linalg.norm(temp_x) < 1e-8:
# First vertex is at center, try another
temp_x = polygon_vertices[1] - face_center
temp_x = temp_x - np.dot(temp_x, local_z) * local_z
local_x = temp_x / np.linalg.norm(temp_x)
# y-axis: right-hand rule
local_y = np.cross(local_z, local_x)
# Transform matrix (world -> local)
transform_matrix = np.array([local_x, local_y, local_z])
# Project to 2D
vertices_local = (polygon_vertices - face_center) @ transform_matrix.T
polygon_2d = vertices_local[:, :2] # Take only x, y
points_local = (points - face_center) @ transform_matrix.T
points_2d = points_local[:, :2]
# Check z-coordinate (points must be close to face)
z_tolerance = 0.5 # Angstrom
close_to_face = np.abs(points_local[:, 2]) < z_tolerance
# Ray-casting algorithm
in_polygon = self._points_in_polygon_2d(polygon_2d, points_2d)
# Both conditions must be satisfied
valid_mask = close_to_face & in_polygon
return points[valid_mask]
def _points_in_polygon_2d(
self, polygon: np.ndarray, points: np.ndarray
) -> np.ndarray:
"""
Test if 2D points are inside a 2D polygon using ray-casting.
Casts a ray from each point to the right and counts intersections
with polygon edges. Odd number of intersections means inside.
Parameters
----------
polygon : np.ndarray
Polygon vertices with shape (n_vertices, 2).
points : np.ndarray
Points to test with shape (n_points, 2).
Returns
-------
np.ndarray
Boolean array indicating if each point is inside (n_points,).
Notes
-----
Vectorized implementation for efficiency. Handles edge cases:
- Points on vertices
- Points on edges
- Ray passing through vertices
"""
polygon = np.asarray(polygon, dtype=np.float32)
points = np.asarray(points, dtype=np.float32)
# Build polygon edges (current vertex -> next vertex)
p1 = polygon
p2 = np.roll(polygon, -1, axis=0)
# Vectorized computation
# Expand dimensions: points [N, 2] -> [N, 1, 2], polygon [V, 2] -> [1, V, 2]
pts = points[:, np.newaxis, :] # [N, 1, 2]
v1 = p1[np.newaxis, :, :] # [1, V, 2]
v2 = p2[np.newaxis, :, :] # [1, V, 2]
# Check if points are on vertices
on_vertex = np.any(np.all(np.isclose(pts, v1, atol=1e-6), axis=2), axis=1)
# Ray-casting: cast ray to the right, count intersections
# Condition 1: edge crosses point's y-coordinate
y_cross = (v1[:, :, 1] > pts[:, :, 1]) != (v2[:, :, 1] > pts[:, :, 1])
# Condition 2: intersection is to the right of point
x_intersect = (v2[:, :, 0] - v1[:, :, 0]) * (pts[:, :, 1] - v1[:, :, 1]) / (
v2[:, :, 1] - v1[:, :, 1] + 1e-10
) + v1[:, :, 0]
right_side = pts[:, :, 0] < x_intersect
# Count intersections
intersections = np.sum(y_cross & right_side, axis=1)
# Odd number of intersections -> inside polygon
inside = (intersections % 2 == 1) | on_vertex
return inside
def _get_pos(
self, verbose: bool = True
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
"""
Generate atomic positions for all grains.
Process:
1. Determine required unit cell replication
2. Generate metal atoms for each grain
3. Optionally generate graphene atoms for grain boundaries
4. Assign grain IDs and atom types
Returns
-------
Tuple[np.ndarray, np.ndarray, np.ndarray]
- all_pos: atomic positions (n_atoms, 3)
- all_grain_ids: grain ID for each atom (n_atoms,)
- all_types: atom type (1=metal, 2=carbon) (n_atoms,)
Notes
-----
- Unit cell is replicated based on maximum Voronoi cell radius
- Graphene lattice constant is 1.42 Angstrom
- Type 1 = metal, Type 2 = carbon
"""
# Calculate required unit cell replication
r_max = max(cell.cavity_radius for cell in self.con)
thickness = self.unitcell.box.get_thickness()
replicate_nums = np.ceil(r_max / thickness).astype(int)
# Replicate unit cell
data, _ = _replicate_pos(self.unitcell.data, self.unitcell.box, *replicate_nums)
# Prepare graphene if needed.
#
# Note: the new ``build_crystal("graphene", a=...)`` follows
# atomsk and treats ``a`` as the hexagonal in-plane lattice
# parameter (= C-C bond * sqrt(3) ≈ 2.46 Å). The old API used
# the C-C bond directly. We size the sheet so its (x, y)
# bounding box covers ≥ 2 * r_max, which is enough for any face
# of any Voronoi cell after centering at the face center.
if self.add_graphene:
cc_bond = 1.42 # Angstrom — graphene C-C bond
gra_lattice = cc_bond * 3**0.5 # hex in-plane parameter
target = 2.0 * r_max
x1 = int(np.ceil(target / gra_lattice))
# H2 contributes sqrt(3)/2 * a along y per cell.
y1 = int(np.ceil(target / (gra_lattice * 3**0.5 / 2.0)))
# `c` is just a vacuum spacer along z; atoms stay at z=0,
# only the box height uses it.
gra = build_crystal(
"C", "graphene", gra_lattice, nx=x1, ny=y1, nz=1, c=1.0,
)
# Preallocate lists
pos_list = []
grain_id_list = []
type_list = []
# Generate atoms for each grain
for n, cell in enumerate(self.con):
if verbose:
print(
f" Grain {n + 1:>3}/{len(self.con)}: "
f"Volume = {cell.volume:>10.2f} A^3 ",
end="",
)
coeffs = self._get_plane_equation_coeffs_for_cell(cell)
# Generate metal atoms
pos = self._generate_grain_atoms(n, cell, data, coeffs)
pos_list.append(pos)
N_metal = pos.shape[0]
type_list.extend([1] * N_metal)
# Generate graphene atoms if enabled
if self.add_graphene:
gra_pos = self._generate_gra_atoms(cell, gra.data, coeffs)
pos_list.append(gra_pos)
N_carbon = gra_pos.shape[0]
type_list.extend([2] * N_carbon)
N_total = N_metal + N_carbon
if verbose:
print(f"Metal = {N_metal:>6}, Carbon = {N_carbon:>6}")
else:
N_total = N_metal
if verbose:
print(f"Atoms = {N_metal:>6}")
# Assign grain IDs
grain_ids = np.full(N_total, n + 1, dtype=np.int32)
grain_id_list.append(grain_ids)
# Concatenate all atoms
all_pos = np.vstack(pos_list)
all_grain_ids = np.concatenate(grain_id_list)
all_types = np.array(type_list, dtype=np.int32)
return all_pos, all_grain_ids, all_types
[docs]
def compute(self, verbose: bool = True) -> System:
"""
Compute and generate the polycrystalline structure.
Main workflow:
1. Generate Voronoi tessellation from seeds
2. Fill each cell with rotated unit cell atoms
3. Optionally add graphene at grain boundaries
4. Remove overlapping atoms
5. Wrap atoms into simulation box
Parameters
----------
verbose : bool, default=True
If True, print detailed progress information.
If False, suppress all output.
Returns
-------
System
The generated polycrystalline system with attributes:
- data: polars DataFrame with columns ['element', 'x', 'y', 'z', 'grain_id', 'type']
- box: simulation box
Examples
--------
>>> poly = CreatePolycrystal(unitcell, box=100, seed_number=10)
>>> system = poly.compute(verbose=True)
>>> system.write_xyz("output.xyz")
"""
if verbose:
start = time()
print("=" * 70)
print(" " * 20 + "POLYCRYSTAL GENERATION")
print("=" * 70)
# Generate Voronoi tessellation
if verbose:
print("[1/5] Generating Voronoi tessellation...")
origin = self.box.origin.copy()
self.con = Container(self.seed_position, Box(self.box.box))
# Statistical information
volumes = np.array([cell.volume for cell in self.con])
ave_grain_volume = np.mean(volumes)
if verbose:
print(f" Number of grains: {self.seed_number}")
print(f" Average volume: {ave_grain_volume:>10.2f} A^3")
print(
f" Volume range: {volumes.min():>10.2f} - {volumes.max():<10.2f} A^3"
)
if self.add_graphene:
print(
f" Graphene enabled: Yes (threshold = {self.face_threshold} A^2)"
)
print(f" Random seed: {self.randomseed}")
# Generate atomic positions
if verbose:
print(f"[2/5] Generating atoms for {self.seed_number} grains...")
pos, grain_id, type_list = self._get_pos(verbose)
if verbose:
print(f" Total atoms generated: {len(pos):,}")
# Create dataframe
if verbose:
print("[3/5] Creating atomic structure...")
data = pl.DataFrame(
{
"x": pos[:, 0] + origin[0],
"y": pos[:, 1] + origin[1],
"z": pos[:, 2] + origin[2],
"grain_id": grain_id,
"type": type_list,
}
)
# Remove overlapping atoms
if verbose:
print("[4/5] Removing overlapping atoms...")
if self.add_graphene:
if verbose:
print(" Filtering: Metal-Metal, C-C, Metal-C overlaps")
print(
f" Metal-Metal distance: {self.metal_overlap_dis or 2.0:.2f} Å"
)
print(" C-C distance: 1.40 Å")
print(f" Metal-C distance: {self.metal_gra_overlap_dis:.2f} Å")
metal_metal_distance = (
self.metal_overlap_dis if self.metal_overlap_dis is not None else 2.0
)
cc_distance = 1.4 # C-C minimum distance (Angstrom)
metal_c_distance = self.metal_gra_overlap_dis
filter_mask = _neighbor.filter_overlap_atom_with_grain(
data["x"].to_numpy(allow_copy=False),
data["y"].to_numpy(allow_copy=False),
data["z"].to_numpy(allow_copy=False),
data["type"].to_numpy(allow_copy=False),
data["grain_id"].to_numpy(allow_copy=False),
self.box.box,
self.box.origin,
self.box.boundary,
metal_metal_distance,
cc_distance,
metal_c_distance,
get_num_threads(),
)
data = data.filter(filter_mask)
elif self.metal_overlap_dis is not None:
if verbose:
print(" Filtering: Metal-Metal overlaps")
print(f" Metal-Metal distance: {self.metal_overlap_dis:.2f} Å")
filter_mask = _neighbor.filter_overlap_atom(
data["x"].to_numpy(allow_copy=False),
data["y"].to_numpy(allow_copy=False),
data["z"].to_numpy(allow_copy=False),
self.box.box,
self.box.origin,
self.box.boundary,
self.metal_overlap_dis,
get_num_threads(),
)
data = data.filter(filter_mask)
if verbose:
removed = len(pos) - len(data)
print(f" Atoms removed: {removed:,} ({removed / len(pos) * 100:.2f}%)")
print(f" Atoms remaining: {len(data):,}")
# Add element information
if "element" in self.unitcell.data.columns:
element = self.unitcell.data["element"][0]
type2ele = {1: element, 2: "C"}
data = data.with_columns(
pl.col("type").replace_strict(type2ele).rechunk().alias("element")
).select("element", "x", "y", "z", "grain_id", "type")
# Create system and wrap atoms
if verbose:
print("[5/5] Finalizing structure...")
system = System(data=data, box=self.box)
system.wrap_pos()
# Final statistics
if verbose:
print("=" * 70)
print(" ✓ Polycrystal generation completed successfully!")
end = time()
print(f" ✓ Execution time: {end - start:.2f} seconds")
print("=" * 70 + "")
return system
if __name__ == "__main__":
pass