Proposal: post-generation structural analysis, screening pipeline, and compositional guided sampling

[exopoiesis]

Context

We've been using CrystalFormer for generative design of layered mineral membranes (Fe-Ni-S sulfides for proton transport). Over the past few weeks we generated ~7000+ structures using both targeted and wild sampling, and built an ad-hoc Voronoi analysis on top to filter candidates by layeredness and void size.

This worked well — 91% of our targeted generations were layered, 87% had voids > 1 Å — but the analysis code is messy, single-purpose, and not reusable. We think there's value in building proper post-generation analysis and screening tools directly into CrystalFormer, and wanted to discuss the idea before investing time in a PR.

The gap

CrystalFormer does an excellent job at generation, but the workflow after generation is entirely on the user:

CrystalFormer generates structures → ??? → user somehow evaluates them

For many use cases (ionic conductors, membranes, porous materials, MOFs, zeolites), the key question isn't just "is this a valid crystal?" but "does this crystal have the structural features I need?" — channels, voids, layered morphology, specific compositions, etc.

Proposal: three modules

We'd like to contribute three chemistry-agnostic modules. None of these are tied to our specific Fe-S use case — they work for any crystal system.

1. Structural analysis (crystalformer/analysis/)

Voronoi-based characterization of generated structures:

Void analysis (voronoi.py):

• max_void_radius — largest inscribed sphere radius (Å)
• void_fraction — free volume fraction
• layeredness_score — anisotropy of Voronoi vertices (PCA-based, 0 = isotropic, 1 = perfectly layered)
• interlayer_spacing — layer separation for layered structures (Å)

Channel/percolation analysis (percolation.py):

• Build Voronoi network graph (nodes = Voronoi vertices, edges = faces between cells)
• Filter edges by bottleneck > r_probe (configurable, default 0.4 Å for H⁺)
• Check percolation through periodic boundaries in a/b/c directions via BFS
• Report: percolation_dimensionality (0D/1D/2D/3D), min_bottleneck along best channel path
• This answers the question "can an ion of radius r travel through this structure?" — useful for any ionic conductor screening

CLI:

python -m crystalformer.analysis structures.csv --r-probe 0.4 --check-percolation --output results.csv

Python API:

from crystalformer.analysis import VoronoiAnalyzer, PercolationAnalyzer

va = VoronoiAnalyzer()
result = va.analyze(structure)  # pymatgen Structure
# result.max_void_radius → 1.82 Å
# result.layeredness_score → 0.91

pa = PercolationAnalyzer(r_probe=0.4)
result = pa.analyze(structure)
# result.percolation_dimensionality → 2 (layered conductor)
# result.min_bottleneck → 0.62 Å
# result.percolates_c → True

2. Compositional guided sampling

Bias atom-type logits during autoregressive decoding to steer generation toward desired compositions — without retraining or fine-tuning.

Mechanism: At each atom-type sampling step, add a bias vector to the logits before softmax. Positive bias encourages an element, negative suppresses it. Zero = unchanged (default behavior preserved).

# "I want more structures with Li and O"
structures = sample_crystal(
    params, n_samples=1000,
    composition_bias={"Li": 2.0, "O": 1.5, "P": 1.0}
)

python -m crystalformer.sample --checkpoint model.pkl \
    --n-samples 1000 \
    --composition-bias "Li:2.0,O:1.5,P:1.0"

This is deliberately simple — not a hard constraint, just a soft nudge. It turns "generate random structures and hope for the right composition" into "generate structures enriched in desired elements." For us this increased Fe-S hit rate from ~12% to ~60% without any model changes.

Scope: Only modifies the sampling function — no changes to model architecture, training, or weights.

3. Screening pipeline (crystalformer/screen/)

A pluggable filter chain for batch screening of generated structures.

Built-in filters:

Filter	Criterion	Default
validity	Valid structure (atoms placed, no overlaps < 0.5 Å)	ON
voronoi	max_void_radius > threshold	r > 1.0 Å
percolation	percolation_dimensionality >= threshold	dim >= 1
composition	Contains specified elements	OFF
density	Density within range	OFF

Custom filters (user-extensible):

from crystalformer.screen import ScreeningPipeline, Filter

class StabilityFilter(Filter):
    """Example: user plugs in their own ML potential."""
    def __call__(self, structure) -> bool:
        energy = my_ml_model.predict(structure)
        return energy < self.threshold

pipeline = ScreeningPipeline([
    "validity",
    "voronoi:r_min=0.4",
    "percolation:dim=2",
    StabilityFilter(threshold=0.05),
])

results = pipeline.run("generated.csv", output="screened.csv")
# 5000 → 4200 valid → 1890 voronoi → 340 percolating → 52 stable

CLI:

python -m crystalformer.screen generated/ \
    --filters validity,voronoi,percolation \
    --r-probe 0.4 \
    --output screened.csv

The pipeline prints a summary table showing how many structures pass each stage — makes it easy to see where the funnel narrows.

Dependencies

Only pymatgen (already used for I/O) and scipy (for scipy.spatial.Voronoi), as optional:

[project.optional-dependencies]
analysis = ["pymatgen>=2024.1.1", "scipy>=1.10"]

No heavy ML dependencies. Core CrystalFormer stays lightweight.

File structure (all new files, minimal changes to existing code)

crystalformer/
├── analysis/
│   ├── __init__.py
│   ├── voronoi.py          # VoronoiAnalyzer
│   ├── percolation.py      # PercolationAnalyzer
│   └── __main__.py         # CLI
├── screen/
│   ├── __init__.py
│   ├── pipeline.py         # ScreeningPipeline, Filter base
│   ├── filters.py          # Built-in filters
│   └── __main__.py         # CLI
└── src/
    └── sample.py           # + optional composition_bias param

Questions before we start

1. Do you already have internal analysis/screening tools? We don't want to duplicate existing work. If you have something similar internally, we'd be happy to build on it instead.
2. Is this in scope for the main repo? If you'd prefer this to live as a separate package (e.g., crystalformer-analysis), that works too.
3. Single PR or incremental? We can do one PR with everything, or split into 2-3 (analysis → sampling → screening).
4. Any naming/API conventions you'd like us to follow?

Happy to discuss any aspect of this. We've been enjoying working with CrystalFormer and would like to contribute back to the project.

[zdcao121]

Thanks a lot for the detailed and thoughtful proposal — this is very valuable feedback, and we are glad to hear CrystalFormer has been useful in your workflow.

Your proposal is well motivated and technically clear. For the main CrystalFormer repo, however, we’d prefer to keep scope focused on the foundation model itself (training + inference). Integrating a large amount of post-generation analysis/screening code into the core would significantly increase the long-term maintenance burden.

At the moment, we already keep some evaluation/screening reference utilities under scripts/ (e.g., validity, novelty/uniqueness, relaxation, Ehull-based evaluation), mainly as examples rather than a complete framework. In practice, users across different application domains care about different metrics and screening funnels, and in today’s AI-coding environment, it is relatively straightforward for users to build task-specific post-processing pipelines on top of generated structures.

Given that, the most suitable path is:

1. Build this as a separate extension package/repo (e.g., crystalformer-analysis / crystalformer-screen), with CrystalFormer as an upstream dependency.
2. Keep interfaces lightweight and interoperable with CrystalFormer outputs.
3. Optionally, we can link/recommend the extension from our README/docs once it is stable.

[exopoiesis]

Thanks for the clear and thoughtful guidance — this makes a lot of sense.

We've split the proposal accordingly:

1. Compositional guided sampling — submitted as PR feat: compositional guided sampling (--composition_bias) #14. This is a small, self-contained feature (3 files touched, 17 tests) that lives entirely within src/ and adds a --composition_bias CLI flag. No new dependencies, no maintenance burden beyond what's already there. We think this belongs in the core since it's a natural extension of the sampling interface.

2. Structural analysis + screening pipeline — will live as a separate package with CrystalFormer as an upstream dependency, exactly as you suggested. We'll keep interfaces lightweight and interoperable with CrystalFormer outputs (CSV/CIF). Happy to share the link once it's stable.

Looking forward to your feedback on #14!

[zdcao121]

Why introduce a separate composition-bias parameter instead of using atom_mask?

My understanding is that atom_mask is already sufficient if the goal is to restrict allowed elements.

And if the goal is composition control, we can directly sample with a given composition vector (i.e., the crystal structure prediction setting): condition generation on the target composition and sample accordingly.

So practically:

• atom_mask handles element-set constraints.
• explicit composition-conditioned sampling handles stoichiometry/composition constraints.

Given this, could you clarify what additional benefit composition-bias provides beyond these two existing mechanisms?

[exopoiesis]

Great question — this is exactly the right thing to ask. Here's the distinction:

The gap between atom_mask and CSP

atom_mask is a hard binary constraint: an element is either allowed or blocked entirely. It answers "generate structures using ONLY these elements." This is useful but inflexible — it completely eliminates serendipitous discoveries of beneficial dopants or unexpected compositions.

CSP (formula-conditioned sampling) requires the exact stoichiometry upfront (e.g., Fe2NiS4). It answers "generate structures with THIS formula." This is powerful when you know what you want, but in exploratory materials design, the optimal composition is often unknown.

composition_bias fills the gap between these two. It answers: "I want structures enriched in Fe and S, but I don't know the exact formula — surprise me."

Concrete comparison

Mechanism	Constraint type	Requires formula?	Allows unexpected elements?
atom_mask	Hard block/allow	No	No (blocked = 0 probability)
CSP	Exact stoichiometry	Yes	No (fixed composition)
composition_bias	Soft nudge	No	Yes (just shifts probabilities)

Why this matters in practice

In our use case (designing layered iron sulfide membranes), we know we want Fe and S, but:

• We don't know if Fe₂S₃, FeS, Fe₃NiS₄, or something else is optimal
• Ni might be beneficial (it is — pentlandite Fe₄.₅Ni₄.₅S₈ turned out to be our best candidate), but we wouldn't have known to include it in a formula
• Cu substitution might help (chalcopyrite CuFeS₂ is interesting too)

With atom_mask=Fe,S,Ni: we'd miss Cu-containing structures entirely.
With CSP --formula FeNiS2: we'd miss Fe₃NiS₄, Fe₅Ni₄S₈, FeS, etc.
With --composition_bias "Fe:3.0,S:3.0": the model explores freely but generates ~50% Fe-containing and ~40% S-containing structures instead of the baseline ~6% and ~12%.

Empirical validation (just ran this today)

We tested on the pretrained checkpoint (epoch 44000), SG=129 (P4/nmm), 16 structures per run:

Run	Structures with Fe	%
--composition_bias "Fe:3.0,S:3.0,Ni:2.0"	8/16	50%
No bias (baseline)	1/16	6%

The biased run still generates diverse compositions (Ti, Mg, As, In appear alongside Fe), which is exactly what you want for discovery — enriched in target elements, but not locked to a specific formula.

Composability

composition_bias composes cleanly with existing mechanisms:

# Bias + mask: "prefer Fe and S, but NEVER allow radioactive elements"
--composition_bias "Fe:3.0,S:3.0" --remove_radioactive

# Bias + CSP: "generate FeNiS2, but if the model hesitates on Ni, nudge it"
--formula FeNiS2 --composition_bias "Ni:1.0"

# Bias alone: "explore iron sulfide space freely"
--composition_bias "Fe:3.0,S:3.0"

The implementation is minimal (1 line added to the sampling loop: a_logit = a_logit + ... + _composition_bias), zero overhead when unused (default is a zero vector), and the atom_mask hard block always overrides the soft bias (as tested — see PR #14 test test_atom_mask_overrides_positive_bias).

In short: atom_mask and CSP are precise tools for when you know what you want. composition_bias is a discovery tool for when you know the direction but not the destination.

[zdcao121]

Thanks for the detailed explanation — this distinction is very clear and helpful.

I still wonder whether composition_bias is necessary in the most common practical workflow.

In many projects, we already know the element space (or a target chemical system), but not the best stoichiometry yet. In that setting, one reasonable path is:

1. Use atom_mask to remove clearly irrelevant elements and keep sampling inside the target composition space.
2. Once promising compositions are identified, switch to crystal structure prediction (formula-conditioned sampling/CSP) to test specific stoichiometries and evaluate structural stability.

From this perspective, composition_bias may be less critical, because we can first constrain the search space with atom_mask and then do focused CSP-based stability checks.

A second concern is controllability: the logits bias is manually specified by users. In practice, it can be hard to choose robust bias magnitudes (Fe:3.0 vs Fe:1.0, etc.), and tuning these values may introduce extra trial-and-error or user-dependent behavior.

So my current view is:

• atom_mask + CSP already covers a large part of the "known direction, unknown ratio" workflow.
• composition_bias is interesting for exploratory discovery, but may not be essential as a core interface unless we also provide clear guidance or automatic calibration for bias strength.

Happy to hear your thoughts.

[exopoiesis]

Thanks for the thoughtful pushback — these are fair points, especially on controllability.

One thing I'd note: the atom_mask + CSP workflow works well when the element space is small and well-defined. But there's an implicit assumption that the user can enumerate the "clearly irrelevant" elements upfront. In practice, this boundary is often fuzzy — especially in exploratory settings where beneficial dopants or substitutions aren't known in advance.

For example, in our iron sulfide work, Ni turned out to be critical (pentlandite Fe4.5Ni4.5S8), but we had no reason to include it in the mask initially. With atom_mask=Fe,S we'd have missed it entirely. With atom_mask=Fe,S,Ni,Co,Cu,Mn,... we're essentially guessing the answer before asking the question — and the mask gives equal weight to all included elements, which may not reflect domain knowledge (e.g., "Fe and S are essential, Ni is plausible, Mn is a long shot").

That said, I agree that raw logit values are not a user-friendly interface, and adding auto-calibration would increase complexity without a clear win. So rather than pushing this into the core, we'll move composition_bias into our extension package (crystalformer-x) and continue developing it there — including experimenting with more intuitive parameterizations.

We'll close PR #14 accordingly. Thanks again for the open and constructive discussion — it's been genuinely helpful in shaping our thinking.

[zdcao121]

Moving composition_bias to crystalformer-x sounds good. That keeps the core repo focused while giving space to iterate on this interface externally.

Please share updates as it matures; we can revisit the README links later.

[exopoiesis]

All resolved — composition_bias lives in crystalformer-x, will share updates as it matures. Thanks!