"""
FMSteel3DBase Implementation: Grain detection, neighbor graphs, PAG clustering.
Phase 1 working implementation of the base grain structure class.
Entry point for the microstructure generation pipeline.
"""
import numpy as np
from typing import Dict, List, Tuple, Optional
from dataclasses import dataclass
import cc3d
[docs]
@dataclass
class PhysicalDimensions:
"""Container for physical domain dimensions (Lx, Ly, Lz)."""
Lx: float
Ly: float
Lz: float
[docs]
@classmethod
def from_tuple(cls, dims_tuple: Tuple[float, float, float]):
"""Create from (Lx, Ly, Lz) tuple."""
return cls(Lx=dims_tuple[0], Ly=dims_tuple[1], Lz=dims_tuple[2])
[docs]
def as_array(self) -> np.ndarray:
"""Return as (Lx, Ly, Lz) array."""
return np.array([self.Lx, self.Ly, self.Lz])
[docs]
class FMSteel3DBase:
"""
Foundation FM steel grain structure from labeled grain image (LFI).
Entry point to the pipeline. Detects grains, builds neighbor graph,
optionally cleans small grains, then enables PAG clustering.
"""
__slots__ = ('lgi', 'physical_dimensions', 'voxel_size', 'units',
'connectivity', 'grain_locs', 'n_grains', 'neigh_gid',
'min_grain_nvoxels', '_random_seed', '_verbosity', '_log_sink')
_VALID_UNITS = {'microns', 'mm', 'm'}
def __init__(self,
lgi: np.ndarray,
physical_dimensions: PhysicalDimensions,
voxel_size: float,
units: str = 'microns',
connectivity: int = 6,
grain_locs: Optional[Dict[int, np.ndarray]] = None,
n_grains: int = 0,
neigh_gid: Optional[Dict[int, List[int]]] = None,
min_grain_nvoxels: int = -1,
random_seed: Optional[int] = None,
verbosity: int = 0,
log_sink=None):
"""
Initialize FMSteel3DBase from component data.
Typically called internally by from_lfi classmethod. Direct instantiation
is allowed but not recommended.
Parameters
----------
lgi : np.ndarray
3D labeled grain image.
physical_dimensions : PhysicalDimensions
Physical domain dimensions (Lx, Ly, Lz).
voxel_size : float
Size of each voxel.
connectivity : int, optional
Grain connectivity (6, 18, 26). Default is 6 (face connectivity).
grain_locs : dict, optional
Pre-computed grain locations. If None, computed on demand.
n_grains : int, optional
Number of grains. If 0, computed from lgi.max().
neigh_gid : dict, optional
Pre-computed neighbor relationships. If None, computed on demand.
min_grain_nvoxels : int, optional
Minimum grain size. Default -1 (no cleanup).
random_seed : int, optional
Random seed for reproducibility.
"""
if units not in self._VALID_UNITS:
raise ValueError(
f"Invalid units '{units}'. Valid: {self._VALID_UNITS}"
)
self.lgi = lgi.astype(np.int32) if lgi.dtype != np.int32 else lgi
self.physical_dimensions = physical_dimensions
self.voxel_size = float(voxel_size)
self.units = units
self.connectivity = int(connectivity)
self.grain_locs = grain_locs or {}
self.n_grains = int(n_grains)
self.neigh_gid = neigh_gid or {}
self.min_grain_nvoxels = int(min_grain_nvoxels)
self._random_seed = random_seed
self._verbosity = int(verbosity)
self._log_sink = log_sink
def _emit(self, level: int, msg: str, component: str = 'BASE') -> None:
"""Emit a structured message if verbosity allows."""
if self._verbosity < int(level):
return
text = f"[{component}][L{int(level)}] {msg}"
if self._log_sink is not None:
self._log_sink(text)
else:
print(text)
[docs]
@classmethod
def from_lfi(cls,
lfi: np.ndarray,
physical_dimensions: Tuple[float, float, float],
voxel_size: float = 1.0,
units: str = 'microns',
connectivity: int = 6,
min_grain_nvoxels: int = -1,
grain_cleanup_max_passes: int = 100,
random_seed: Optional[int] = None,
verbosity: int = 0,
log_sink=None) -> 'FMSteel3DBase':
"""
Create FMSteel3DBase instance from a labeled grain image.
This is the primary factory method. It initializes grain locations,
computes neighbor relationships, and optionally performs cleanup
(dissolution of small grains) if min_grain_nvoxels >= 0.
Parameters
----------
lfi : np.ndarray
3D labeled grain image. Shape (nx, ny, nz). Values are grain IDs
(1, 2, ..., n_grains).
physical_dimensions : tuple[float, float, float]
Physical domain size as (Lx, Ly, Lz).
voxel_size : float, optional
Size of each voxel. Default 1.0. Must be positive.
connectivity : int, optional
Grain connectivity (6, 18, or 26). Default 6 (face connectivity only).
Values: 6 = face, 18 = face+edge, 26 = face+edge+corner.
min_grain_nvoxels : int, optional
Minimum grain size in voxels. If >= 0, triggers cleanup via
clean_gs_GMD_by_source_erosion_v1. Grains below threshold are
dissolved into neighboring larger grains. Default -1 (no cleanup).
random_seed : int, optional
Random seed for reproducibility. If provided, numpy random state
is set at initialization.
Returns
-------
FMSteel3DBase
Fully initialized grain structure instance.
Raises
------
ValueError
If lfi is not 3D, connectivity is invalid, or physical_dimensions are invalid.
Examples
--------
>>> lfi = np.random.randint(1, 100, size=(50, 50, 50))
>>> fm = FMSteel3DBase.from_lfi(lfi, physical_dimensions=(100.0, 100.0, 100.0))
>>> print(fm.n_grains)
"""
if random_seed is not None:
np.random.seed(random_seed)
lfi = np.asarray(lfi, dtype=np.int32)
phys_dims = (PhysicalDimensions.from_tuple(physical_dimensions)
if isinstance(physical_dimensions, (tuple, list))
else physical_dimensions)
instance = cls(lgi=lfi, physical_dimensions=phys_dims,
voxel_size=float(voxel_size), units=units,
connectivity=int(connectivity),
min_grain_nvoxels=int(min_grain_nvoxels), random_seed=random_seed,
verbosity=verbosity, log_sink=log_sink)
instance._emit(1, f"Initializing from LFI shape={tuple(lfi.shape)}, connectivity={connectivity}")
if min_grain_nvoxels >= 0:
instance._emit(2, f"Cleaning grains below {min_grain_nvoxels} voxels")
instance = instance.clean_small_grains(threshold=min_grain_nvoxels,
max_passes=grain_cleanup_max_passes)
instance.grain_locs = instance.compute_grain_locations()
instance.n_grains = len(instance.grain_locs)
instance.neigh_gid = instance.compute_neighbor_network(connectivity=connectivity)
instance._emit(1, f"Initialized grain structure with {instance.n_grains} grains")
return instance
[docs]
def compute_grain_locations(self) -> Dict[int, np.ndarray]:
"""
Compute voxel coordinates for each grain from LGI.
Single sort-and-split pass over all non-background voxels — O(N log N)
in voxel count, independent of grain count. Replaces an earlier
per-grain `np.argwhere(lgi == gid)` loop that rescanned the full array
once per grain (O(n_grains * N)); on a 100^3 domain with ~3000 grains
that loop took ~7s vs ~0.08s here (verified equivalent output on
randomized structures before replacing).
Returns
-------
dict[int, np.ndarray]
Maps grain_id → (n_voxels, 3) array of voxel coordinates.
Notes
-----
Results are cached in self.grain_locs after first call.
"""
coords = np.argwhere(self.lgi > 0)
if coords.shape[0] == 0:
return {}
labels = self.lgi[coords[:, 0], coords[:, 1], coords[:, 2]]
order = np.argsort(labels, kind='stable')
labels_sorted = labels[order]
coords_sorted = coords[order]
boundaries = np.flatnonzero(np.diff(labels_sorted)) + 1
groups = np.split(coords_sorted, boundaries)
starts = np.concatenate(([0], boundaries))
unique_labels = labels_sorted[starts]
return {int(lbl): grp for lbl, grp in zip(unique_labels, groups)}
@staticmethod
def _offset_slices(shape, dx: int, dy: int, dz: int):
"""Aligned (src, dst) slice tuples for a voxel shift by (dx, dy, dz)
with no wraparound: ``arr[dst][m]`` gives, for each True position in
``mask[src]``, the value of the neighbour at ``position + (dx,dy,dz)``.
"""
src_slices, dst_slices = [], []
for d, n in zip((dx, dy, dz), shape):
if d >= 0:
src_slices.append(slice(0, n - d))
dst_slices.append(slice(d, n))
else:
src_slices.append(slice(-d, n))
dst_slices.append(slice(0, n + d))
return tuple(src_slices), tuple(dst_slices)
[docs]
def compute_neighbor_network(self, connectivity: Optional[int] = None) -> Dict[int, List[int]]:
"""
Compute grain neighbor relationships using cc3d.region_graph.
Parameters
----------
connectivity : int, optional
Connectivity type (6, 18, 26). If None, uses self.connectivity.
Returns
-------
dict[int, list[int]]
Maps grain_id → list of neighbor grain_ids.
Notes
-----
Results are cached in self.neigh_gid after first call.
Uses cc3d.region_graph for efficient neighbor detection.
"""
if connectivity is None:
connectivity = self.connectivity
# cc3d.region_graph returns a set of frozensets of adjacent label pairs,
# not a NetworkX graph. Build the adjacency dict by iterating over edges.
edges = cc3d.region_graph(self.lgi, connectivity=connectivity)
neigh_dict: Dict[int, List[int]] = {}
for edge in edges:
pair = tuple(edge)
a, b = int(pair[0]), int(pair[1])
if a == 0 or b == 0:
continue
neigh_dict.setdefault(a, []).append(b)
neigh_dict.setdefault(b, []).append(a)
# Ensure every grain has an entry (even isolated ones with no touching neighbors)
for gid in np.unique(self.lgi):
if gid == 0:
continue
gid = int(gid)
if gid not in neigh_dict:
neigh_dict[gid] = []
else:
neigh_dict[gid] = sorted(list(set(neigh_dict[gid])))
return neigh_dict
[docs]
def clean_small_grains(self, threshold: int, parameter_metric: str = 'mean',
max_passes: int = 100) -> 'FMSteel3DBase':
"""
Dissolve grains smaller than threshold into larger neighbors.
Reimplements clean_gs_GMD_by_source_erosion_v1 from mcgs3_temporal_slice.py
to work independently. Iteratively merges small grains with their
largest neighboring grains until all remaining grains are >= threshold voxels.
Does NOT modify self -- returns a new FMSteel3DBase instance with the
cleaned grain structure. Callers must use the return value, e.g.
``fm = fm.clean_small_grains(threshold=10)``.
Parameters
----------
threshold : int
Minimum grain size (voxels). Grains with fewer voxels are dissolved.
parameter_metric : str, optional
Reserved for future sink-selection strategies ('mean', 'max', etc.,
matching the original clean_gs_GMD_by_source_erosion_v1 API).
Currently unused: this reimplementation always merges into the
largest neighboring grain regardless of the value passed.
Returns
-------
FMSteel3DBase
New instance with the cleaned grain structure.
Notes
-----
This is called automatically by from_lfi if min_grain_nvoxels >= 0.
Each pass is fully vectorized: grain sizes come from a single
`np.bincount`, and neighbour-label detection for ALL small grains in
the pass is done via whole-array shifted-slice comparisons (one pair
of slices per connectivity offset, not per grain/voxel), followed by
a single array-wide relabel. This replaced an earlier per-grain,
per-voxel, per-offset Python loop; verified to produce identical
grain-size distributions on randomized test structures (including a
deliberately chained small-into-small-into-large merge case), while
running roughly 500-650x faster on a 100^3 / ~3000-grain domain.
A pass's merge decisions are all computed from that pass's start-of-
pass snapshot and applied together, rather than the previous
implementation's incidental order-dependence (where a grain
processed later in the same pass could see an earlier grain's
already-applied merge) — final converged results are equivalent,
occasionally reaching convergence in a different number of passes.
"""
lgi_clean = self.lgi.copy()
n_before = int(np.sum(np.unique(lgi_clean) > 0))
self._emit(1, f"clean_small_grains: threshold={threshold} vox input={n_before} grains", 'CLEAN')
_conn = self.connectivity
_all = [(dx, dy, dz)
for dx in (-1, 0, 1)
for dy in (-1, 0, 1)
for dz in (-1, 0, 1)
if (dx, dy, dz) != (0, 0, 0)]
if _conn == 6:
_offsets = [(dx, dy, dz) for dx, dy, dz in _all if abs(dx)+abs(dy)+abs(dz) == 1]
elif _conn == 18:
_offsets = [(dx, dy, dz) for dx, dy, dz in _all if abs(dx)+abs(dy)+abs(dz) <= 2]
else:
_offsets = _all
for pass_i in range(max_passes):
counts = np.bincount(lgi_clean.ravel())
gids_all = np.arange(counts.size)
small_mask = (gids_all > 0) & (counts > 0) & (counts < threshold)
if not small_mask.any():
self._emit(2, f" converged after {pass_i} pass(es)", 'CLEAN')
break
self._emit(2, f" pass {pass_i + 1}: {int(small_mask.sum())} small grains to dissolve", 'CLEAN')
is_small = np.zeros(counts.size, dtype=bool)
is_small[gids_all[small_mask]] = True
is_small_voxel = is_small[lgi_clean]
pairs_own, pairs_neigh = [], []
for dx, dy, dz in _offsets:
src, dst = self._offset_slices(lgi_clean.shape, dx, dy, dz)
m_src = is_small_voxel[src]
if not m_src.any():
continue
own = lgi_clean[src][m_src]
neigh = lgi_clean[dst][m_src]
valid = (neigh > 0) & (neigh != own)
if valid.any():
pairs_own.append(own[valid])
pairs_neigh.append(neigh[valid])
if not pairs_own:
self._emit(2, " remaining small grains have no neighbours to merge into — stopping", 'CLEAN')
break
all_own = np.concatenate(pairs_own)
all_neigh = np.concatenate(pairs_neigh)
order = np.argsort(all_own, kind='stable')
own_sorted = all_own[order]
neigh_sorted = all_neigh[order]
neigh_sizes_sorted = counts[neigh_sorted]
boundaries = np.flatnonzero(np.diff(own_sorted)) + 1
starts = np.concatenate(([0], boundaries))
ends = np.concatenate((boundaries, [len(own_sorted)]))
best_map = {}
for s, e in zip(starts.tolist(), ends.tolist()):
grp_neigh = neigh_sorted[s:e]
grp_sizes = neigh_sizes_sorted[s:e]
best_map[int(own_sorted[s])] = int(grp_neigh[int(np.argmax(grp_sizes))])
# Resolve transitive chains within this pass (small->small->large
# collapses to small->large) so the single array-wide remap below
# stays correct even when a small grain's chosen sink is itself
# another small grain being merged in the same pass.
resolved = dict(best_map)
changed = True
while changed:
changed = False
for k in list(resolved.keys()):
v = resolved[k]
if v in resolved and resolved[v] != v:
resolved[k] = resolved[v]
changed = True
remap = np.arange(counts.size)
for k, v in resolved.items():
remap[k] = v
lgi_clean = remap[lgi_clean]
n_after = int(np.sum(np.unique(lgi_clean) > 0))
self._emit(1, f"clean_small_grains: done {n_before}→{n_after} grains", 'CLEAN')
new_inst = FMSteel3DBase(lgi=lgi_clean, physical_dimensions=self.physical_dimensions,
voxel_size=self.voxel_size, connectivity=self.connectivity,
min_grain_nvoxels=self.min_grain_nvoxels, random_seed=self._random_seed,
verbosity=self._verbosity, log_sink=self._log_sink)
new_inst.grain_locs = new_inst.compute_grain_locations()
new_inst.n_grains = len(new_inst.grain_locs)
new_inst.neigh_gid = new_inst.compute_neighbor_network()
return new_inst
[docs]
def remove_spikes(self):
"""
Remove spike voxels (face-connected same-grain neighbour count <= 1).
Each spike voxel is reassigned to whichever neighbouring grain has
the most face-touching voxels. Single vectorised detection pass,
small Python loop only over the (typically very few) spike voxels.
Adapted from StructureCleaner3D._fast_clean_spikes in the
twinned_simple_3d pipeline.
Returns
-------
(new_instance, spike_count) : tuple[FMSteel3DBase, int]
"""
lgi = self.lgi.copy()
c = np.zeros_like(lgi, dtype=np.int32)
for ax in range(3):
sl1 = [slice(None)] * 3
sl2 = [slice(None)] * 3
sl1[ax] = slice(1, None)
sl2[ax] = slice(None, -1)
eq = (lgi[tuple(sl1)] == lgi[tuple(sl2)]).astype(np.int32)
c[tuple(sl1)] += eq
c[tuple(sl2)] += eq
face_offsets = [(1,0,0),(-1,0,0),(0,1,0),(0,-1,0),(0,0,1),(0,0,-1)]
sx, sy, sz = lgi.shape
spike_coords = np.argwhere((lgi > 0) & (c <= 1))
self._emit(1, f"remove_spikes: detected {len(spike_coords)} candidate spike voxels", 'CLEAN')
spike_count = 0
for ix, iy, iz in spike_coords.tolist():
nb_counts: dict = {}
gid = int(lgi[ix, iy, iz])
for dx, dy, dz in face_offsets:
nx_, ny_, nz_ = ix+dx, iy+dy, iz+dz
if 0 <= nx_ < sx and 0 <= ny_ < sy and 0 <= nz_ < sz:
ng = int(lgi[nx_, ny_, nz_])
if ng > 0 and ng != gid:
nb_counts[ng] = nb_counts.get(ng, 0) + 1
if nb_counts:
lgi[ix, iy, iz] = max(nb_counts, key=nb_counts.get)
spike_count += 1
self._emit(1, f"remove_spikes: done {spike_count}/{len(spike_coords)} spikes repaired", 'CLEAN')
new_inst = FMSteel3DBase(lgi=lgi, physical_dimensions=self.physical_dimensions,
voxel_size=self.voxel_size, connectivity=self.connectivity,
min_grain_nvoxels=self.min_grain_nvoxels, random_seed=self._random_seed,
verbosity=self._verbosity, log_sink=self._log_sink)
new_inst.grain_locs = new_inst.compute_grain_locations()
new_inst.n_grains = len(new_inst.grain_locs)
new_inst.neigh_gid = new_inst.compute_neighbor_network()
return new_inst, spike_count
[docs]
def get_grain_statistics(self) -> Dict[str, any]:
"""
Compute basic statistics on grain structure.
Returns
-------
dict
Keys: 'n_grains', 'min_voxels', 'max_voxels', 'mean_voxels',
'median_voxels', 'total_voxels', 'domain_voxels'.
"""
if not self.grain_locs:
self.grain_locs = self.compute_grain_locations()
sizes = [len(coords) for coords in self.grain_locs.values()]
return {
'n_grains': self.n_grains,
'min_voxels': int(np.min(sizes)) if sizes else 0,
'max_voxels': int(np.max(sizes)) if sizes else 0,
'mean_voxels': float(np.mean(sizes)) if sizes else 0.0,
'median_voxels': float(np.median(sizes)) if sizes else 0.0,
'total_voxels': int(np.sum(sizes)),
'domain_voxels': int(np.prod(self.lgi.shape))
}
def _select_isolated_grains(
self,
grain_ids: list,
grain_sizes_map: dict,
target_iso_voxels: float,
strategy: str,
isolated_size_tol: float,
) -> set:
"""Greedy independent-set selection for isolated grains.
Builds an independent set in the grain adjacency graph (no two selected
grains are neighbours), stopping once the cumulative voxel count reaches
target_iso_voxels. Candidate ordering is controlled by *strategy*.
"""
if target_iso_voxels <= 0:
return set()
if strategy == 'smallest':
ordered = sorted(grain_ids, key=lambda g: grain_sizes_map[g])
elif strategy == 'largest':
ordered = sorted(grain_ids, key=lambda g: grain_sizes_map[g], reverse=True)
elif strategy == 'boundary':
lgi = self.lgi
bm = np.zeros(lgi.shape, dtype=bool)
bm[0, :, :] = bm[-1, :, :] = True
bm[:, 0, :] = bm[:, -1, :] = True
bm[:, :, 0] = bm[:, :, -1] = True
boundary_set = set(int(g) for g in np.unique(lgi[bm]) if g > 0)
b_list = [g for g in grain_ids if g in boundary_set]
i_list = [g for g in grain_ids if g not in boundary_set]
np.random.shuffle(b_list)
np.random.shuffle(i_list)
ordered = b_list + i_list
elif strategy == 'near_mean':
mean_sz = float(np.mean(list(grain_sizes_map.values())))
ordered_all = sorted(grain_ids, key=lambda g: abs(grain_sizes_map[g] - mean_sz))
tol_abs = isolated_size_tol * mean_sz
primary = [g for g in ordered_all
if abs(grain_sizes_map[g] - mean_sz) <= tol_abs]
secondary = [g for g in ordered_all if g not in set(primary)]
ordered = primary + secondary
else: # 'random' and unknown fallback
ordered = list(grain_ids)
np.random.shuffle(ordered)
isolated: set = set()
iso_voxels = 0.0
excluded: set = set() # neighbours of already-selected grains
for gid in ordered:
if gid in excluded:
continue
isolated.add(gid)
iso_voxels += grain_sizes_map[gid]
for ngid in self.neigh_gid.get(gid, []):
excluded.add(ngid)
if iso_voxels >= target_iso_voxels:
break
return isolated
[docs]
def generate_pag_clusters(self,
pag_size_distribution: Dict,
pag_grain_fraction: float = 1.0,
use_non_neigh_pag: bool = False,
isolated_grain_strategy: str = 'auto',
isolated_size_tol: float = 0.25,
random_seed: Optional[int] = None) -> 'FMSteel3DWithPAGs':
"""
Partition grains into PAGs (Prior Austenite Grains) via neighbour clustering.
Uses stochastic breadth-first-search on the grain neighbour graph to form
PAGs. Returns a new FMSteel3DWithPAGs instance.
Parameters
----------
pag_size_distribution : dict
Dict with keys 'sizes' (list of target cluster sizes) and
'probs' (list of corresponding probabilities, must sum to ~1.0).
Example: {'sizes': [3, 4, 5, 6, 7],
'probs': [0.10, 0.30, 0.40, 0.15, 0.05]}
pag_grain_fraction : float, optional
Volume fraction of the grain structure that participates in PAG clustering
(0.0–1.0, measured in voxels). The remaining volume becomes isolated grains.
Default 1.0 (all grains participate).
use_non_neigh_pag : bool, optional
If True, each new PAG seed is chosen from grains that do not yet
neighbour any formed PAG, keeping PAGs spatially separated.
As pag_grain_fraction increases this becomes impossible; once no
non-neighbour candidates remain the algorithm falls back to a
random seed from any remaining unclustered grain.
Default False (original random-seed behaviour).
isolated_grain_strategy : str, optional
Strategy for selecting which grains become isolated. Ignored when
pag_grain_fraction >= 1.0. Options:
'auto' — BFS runs until the volume target is met; remaining
grains become isolated (original behaviour).
'random' — random independent set (no two isolated grains touch).
'smallest' — smallest-first independent set.
'largest' — largest-first independent set.
'boundary' — domain-boundary grains preferred in the independent set.
'near_mean' — grains closest to mean grain size preferred (primary:
within isolated_size_tol of mean; fallback: by distance).
Default 'auto'.
isolated_size_tol : float, optional
Fractional tolerance around the mean grain size for the 'near_mean'
strategy. E.g. 0.25 accepts grains within ±25% of the mean voxel
count as primary candidates before falling back. Default 0.25.
random_seed : int, optional
Random seed for reproducibility.
Returns
-------
FMSteel3DWithPAGs
New instance with computed PAG hierarchy.
"""
if random_seed is not None:
np.random.seed(random_seed)
self._emit(
1,
f"Generating PAG clusters (vol_frac={pag_grain_fraction:.3f}, "
f"strategy={isolated_grain_strategy!r}, non_neigh={use_non_neigh_pag})",
component='PAG',
)
sizes = np.array(pag_size_distribution['sizes'])
probs = np.array(pag_size_distribution['probs'])
probs = probs / probs.sum()
grain_ids = [g for g in self.grain_locs.keys() if g != 0]
grain_sizes_map = {g: len(self.grain_locs[g]) for g in grain_ids}
total_voxels = float(sum(grain_sizes_map.values()))
target_iso_voxels = (1.0 - pag_grain_fraction) * total_voxels
if pag_grain_fraction >= 1.0 or target_iso_voxels <= 0 or isolated_grain_strategy == 'auto':
pre_isolated: set = set()
else:
pre_isolated = self._select_isolated_grains(
grain_ids=grain_ids,
grain_sizes_map=grain_sizes_map,
target_iso_voxels=target_iso_voxels,
strategy=isolated_grain_strategy,
isolated_size_tol=isolated_size_tol,
)
clusters_dict: dict = {}
clustered: set = set()
pag_id = 1
available = set(grain_ids) - pre_isolated
clustered_voxels = 0.0
target_clustered_voxels = pag_grain_fraction * total_voxels
pag_neighbor_grains: set = set()
while available:
if isolated_grain_strategy == 'auto' and clustered_voxels >= target_clustered_voxels:
break
# --- seed selection ---
if use_non_neigh_pag:
non_neigh = available - pag_neighbor_grains
pool = non_neigh if non_neigh else available
else:
pool = available
seed = int(np.random.choice(list(pool)))
target_sz = int(np.random.choice(sizes, p=probs))
# --- BFS cluster growth ---
cluster: set = set()
queue_bfs: list = [seed]
visited: set = set()
while queue_bfs and len(cluster) < target_sz:
cur = queue_bfs.pop(0)
if cur in visited or cur in clustered:
continue
visited.add(cur)
cluster.add(cur)
for ngid in self.neigh_gid.get(cur, []):
if (ngid not in visited and ngid not in clustered
and ngid in available and len(cluster) < target_sz):
queue_bfs.append(ngid)
if cluster:
clusters_dict[pag_id] = sorted(list(cluster))
clustered.update(cluster)
clustered_voxels += sum(grain_sizes_map[g] for g in cluster)
pag_id += 1
if use_non_neigh_pag:
for gid in cluster:
for ngid in self.neigh_gid.get(gid, []):
if ngid not in clustered:
pag_neighbor_grains.add(ngid)
pag_neighbor_grains -= clustered
available -= clustered
isolated = pre_isolated | (set(grain_ids) - clustered - pre_isolated)
neigh_clid = {}
for pid, glist in clusters_dict.items():
pag_neigh = set()
for gid in glist:
for ngid in self.neigh_gid.get(gid, []):
if ngid not in glist:
for opid, olist in clusters_dict.items():
if ngid in olist:
pag_neigh.add(opid)
break
neigh_clid[pid] = sorted(list(pag_neigh))
from .with_pags_3d import FMSteel3DWithPAGs
self._emit(
1,
f"Generated {len(clusters_dict)} PAGs; isolated grains={len(isolated)}",
component='PAG',
)
if clusters_dict:
grains_per_pag = [len(v) for v in clusters_dict.values()]
self._emit(
2,
f"grains/PAG min={min(grains_per_pag)}, max={max(grains_per_pag)}, mean={np.mean(grains_per_pag):.1f}",
component='PAG',
)
return FMSteel3DWithPAGs(parent=self, clusters_dict=clusters_dict, neigh_clid=neigh_clid,
pag_orientations={}, isolated_grains=isolated, random_seed=random_seed,
verbosity=self._verbosity, log_sink=self._log_sink)
__all__ = ['FMSteel3DBase', 'PhysicalDimensions']