Release STAR-XPC model, inference code and model card

README.md

@@ -1,3 +1,190 @@
 ---
 license: mit
+library_name: pytorch
+pipeline_tag: tabular-regression
+tags:
+  - set-transformer
+  - chemberta
+  - radiation-shielding
+  - x-ray
+  - materials-science
+  - monte-carlo-surrogate
 ---
+
+# STAR-XPC
+
+STAR-XPC predicts how much of an X-ray beam passes through a polymer composite
+shield. You give it a polyethylene slab loaded with fillers, a slab thickness
+and a tube voltage. It returns the transmission ratio **I/I₀** in milliseconds. 
+The penEasy (PENELOPE) Monte Carlo simulations it was trained to replace take 
+hours per point.
+
+The composite is treated as a *set* of ingredients. Each ingredient is described
+by a ChemBERTa embedding of its SMILES and its weight fraction, and a Set
+Transformer maps the set plus four physics features to I/I₀. Because the
+chemistry enters through SMILES, you can also register materials the model was
+not trained on.
+
+## Quickstart
+
+```bash
+pip install torch numpy safetensors huggingface_hub
+```
+
+```python
+import sys
+from huggingface_hub import snapshot_download
+
+repo = snapshot_download("cmpe-netlab41/STAR-XPC")
+sys.path.insert(0, repo)
+from star_xpc import STARXPC
+
+model = STARXPC.from_pretrained(repo)
+
+# 40 wt% Bi + 30 wt% W in polyethylene (the remaining 30 wt%), 80 kVp beam
+y = model.predict({"Bi": 0.40, "W": 0.30}, thickness_mm=[0.25, 0.5, 1.0, 1.5], kvp=80)
+print(y)  # I/I0 at each thickness
+```
+
+`predict` works out the composite's density and mean excitation energy from the
+composition, so the composition, thickness and voltage are all you need.
+
+## Inputs
+
+| Argument | Meaning | Unit |
+|---|---|---|
+| `fillers` | Filler weight fractions of the **whole** composite, e.g. `{"Bi2O3": 0.4, "BaSO4": 0.4}`. Up to 10 fillers | fraction (0.4 = 40 wt%) |
+| `thickness_mm` | Slab thickness, scalar or array | mm |
+| `kvp` | X-ray tube voltage, scalar or array (broadcast with thickness) | kVp |
+| `polymer_fraction` | Polyethylene weight fraction. Defaults to `1 - sum(fillers)` | fraction |
+| `mass_den` | Composite density. Computed if omitted | g/cm³ |
+| `mean_excit` | Composite mean excitation energy. Computed if omitted | eV |
+
+## Materials
+
+The model ships with the materials it was trained on:
+
+- **Fillers:** B, B4C, BN, Ba, BaSO4, Bi, Bi2O3, Cu, Sb, Sn, W, WO3
+- **Matrix:** polyethylene (C2H4)
+
+### Registering a new material
+
+Any other filler can be registered from its SMILES. It is embedded with ChemBERTa
+(`seyonec/ChemBERTa-zinc-base-v1`, pinned revision) using exactly the recipe the
+training materials went through. For example, lead and gadolinium, which are not
+in the training data:
+
+```python
+# pip install transformers
+model.add_material("Pb", smiles="[Pb]", density=11.34)
+model.add_material("Gd", smiles="[Gd]", density=7.90)
+
+model.predict({"Pb": 0.70}, thickness_mm=[0.3, 0.6, 1.2], kvp=100)
+model.predict({"Gd": 0.40, "W": 0.30}, thickness_mm=1.0, kvp=120)
+```
+
+Registered this way, Pb and Gd reach NMSE % 0.14 and 0.16 against penEasy at
+100 kVp. That is in line with the held-out sets of trained materials (see
+[results](#held-out-test-results)).
+
+- **Formula:** read from the name when the name is a chemical formula (`"Pb"`,
+  `"Gd2O3"`). Otherwise pass `formula="..."`.
+- **Density:** in g/cm³, used to compute the composite density. Without a
+  formula or density, pass `mass_den=` and `mean_excit=` to `predict` instead.
+- **Style:** write SMILES the way the training materials were written. Use pure
+  elements as bracket atoms (`[Pb]`), oxides as covalent structures
+  (`O=[Gd]O[Gd]=O`, `O=[W](=O)=O`), and salts as ions
+  (`[Ba++].[O-]S(=O)(=O)[O-]`). Keep SMILES under 64 tokens, since longer ones
+  are truncated and you get a warning.
+- **Unregistered names:** using one raises an error that shows the
+  `add_material` call to make.
+- **Accuracy:** `predict` warns once for each material that is not in the
+  training data. Treat those predictions as estimates and confirm important
+  ones with Monte Carlo.
+
+## Physics features
+
+- **Mean excitation energy:** the Bragg additivity rule with PENELOPE's elemental
+  I-values for Z = 1–99, so it also works for registered materials. This is what
+  PENELOPE's `material` program does, and it matches the training data to within
+  10⁻⁴.
+- **Density:** the inverse rule of mixtures. It uses the training densities for
+  the bundled materials and the density you pass for registered ones.
+- **If you know better:** some boron-containing training samples used a slightly
+  different boron density. If you have the exact values from PENELOPE, pass
+  `mass_den=` and `mean_excit=`.
+
+## Batch prediction from a table
+
+`predict_frame` takes a pandas DataFrame in the training-data schema:
+`molecule_1..10`, `ratio_1..10` (filler wt%), `polymer`, `polymer_ratio`
+(fraction), `geo_number_mm` (**in units of 0.01 mm**, so 100 = 1 mm), `spc`
+(kVp), `mass_den` and `mean_excit`.
+
+```python
+import pandas as pd
+y = model.predict_frame(pd.read_csv("my_samples.csv"))
+```
+
+## Training domain and limitations
+
+- **Target:** the energy deposited in an air detector behind the slab, divided
+  by the same quantity with the slab replaced by air. It comes from penEasy
+  simulations of a fixed slab-and-detector geometry.
+- **Beams:** four spectra. 30 kVp is ISO N-30, 80 kVp is IEC RQR6, 100 kVp is
+  IEC RQR8 and 120 kVp is IEC RQR9. Other voltages are interpolation between
+  those spectra and are untested.
+- **Thickness:** 0.15–2.5 mm.
+- **Composition:** polyethylene at 20–40 wt%, with one to several fillers making
+  up the remaining 60–80 wt%.
+- **Output range:** the output layer is unbounded, so predictions near 0 or 1
+  can land slightly outside that range.
+- **Warnings:** `predict` warns when an input falls outside the training domain.
+
+## Held-out test results
+
+None of these test sets were used in training. Each has 10 thicknesses.
+
+| Test set | NMSE % | RMSE | MAPE (10 % trimmed) |
+|---|---:|---:|---:|
+| Bi–W, 80 kVp | 0.153 | 0.0119 | 7.59 % |
+| W–Cu–B, 100 kVp | 0.022 | 0.0074 | 1.04 % |
+| Bi–Sb, 120 kVp | 0.051 | 0.0111 | 2.85 % |
+| Bi₂O₃–BaSO₄, 80 kVp | 0.209 | 0.0147 | 5.15 % |
+| BaSO₄–WO₃, 100 kVp | 0.023 | 0.0076 | 1.98 % |
+| Bi₂O₃, 120 kVp | 0.013 | 0.0058 | 1.09 % |
+| **Pb** 70 wt% (registered), 100 kVp | 0.141 | 0.0143 | 4.93 % |
+| **Gd** 70 wt% (registered), 100 kVp | 0.163 | 0.0177 | 4.28 % |
+
+NMSE % = 100 · MSE / mean(y²).
+
+## How it works
+
+Each composite becomes a set of 11 slots (up to 10 fillers plus the polymer,
+zero-padded). Each slot is a 773-d vector:
+
+- the ingredient's 768-d ChemBERTa embedding, which is the mean of the last 4
+  hidden layers, masked-mean-pooled over tokens and L2-normalised;
+- its share of the composition;
+- four z-scored global features: tube voltage, composite density, thickness and
+  mean excitation energy.
+
+A Set Transformer (2 ISAB blocks with 64 inducing points, PMA with 8 seeds and a
+SAB; width 256, 32 heads) turns the set into a single output.
+
+In the training data, filler amounts were stored in wt% and the polymer amount
+as a fraction, and the two were renormalised together. `star_xpc.py`
+reproduces that encoding, so give `predict` ordinary weight fractions.
+
+## Files
+
+| File | Contents |
+|---|---|
+| `star_xpc.py` | Model definition and the `STARXPC` predictor |
+| `model.safetensors` | Set Transformer weights |
+| `embeddings.safetensors` | ChemBERTa embeddings of the 13 training materials |
+| `config.json` | Architecture, feature normalisation statistics, material table, training domain |
+
+## Citation
+
+<!-- TODO: add the paper's BibTeX once it is available -->

config.json

@@ -0,0 +1,121 @@
+{
+  "model_name": "STAR-XPC",
+  "architecture": {
+    "in_dimension": 773,
+    "out_dimension": 1,
+    "d": 256,
+    "m": 64,
+    "h": 32,
+    "k": 8,
+    "num_isabs": 2,
+    "num_rff_layers": 3
+  },
+  "max_slots": 11,
+  "embedding": {
+    "model_id": "seyonec/ChemBERTa-zinc-base-v1",
+    "revision": "761d6a18cf99db371e0b43baf3e2d21b3e865a20",
+    "dim": 768,
+    "max_length": 64
+  },
+  "polymer": "C2H4",
+  "thickness_unit_mm": 0.01,
+  "global_features": [
+    "spc",
+    "mass_den",
+    "geo_number_mm",
+    "mean_excit"
+  ],
+  "normalization": {
+    "spc_mean": 89.3144560357675,
+    "spc_std": 29.071914051102166,
+    "mass_den_mean": 2.2529541424398554,
+    "mass_den_std": 0.5317519959220683,
+    "thick_mean": 92.31850117096019,
+    "thick_std": 58.83530258224455,
+    "mean_excit_mean": 180.05545585669574,
+    "mean_excit_std": 75.43978864260697
+  },
+  "training_domain": {
+    "kvp": [
+      30,
+      80,
+      100,
+      120
+    ],
+    "thickness_mm": [
+      0.15,
+      2.5
+    ],
+    "polymer_fraction": [
+      0.2,
+      0.4
+    ]
+  },
+  "materials": {
+    "B": {
+      "smiles": "[B]",
+      "formula": "B",
+      "density": 2.37
+    },
+    "B4C": {
+      "smiles": "B12B3C14B2B43",
+      "formula": "B4C",
+      "density": 2.52
+    },
+    "BN": {
+      "smiles": "B#N",
+      "formula": "BN",
+      "density": 2.1
+    },
+    "Ba": {
+      "smiles": "[Ba]",
+      "formula": "Ba",
+      "density": 3.59
+    },
+    "BaSO4": {
+      "smiles": "[Ba++].[O-]S(=O)(=O)[O-]",
+      "formula": "BaSO4",
+      "density": 4.48
+    },
+    "Bi": {
+      "smiles": "[Bi]",
+      "formula": "Bi",
+      "density": 9.747
+    },
+    "Bi2O3": {
+      "smiles": "O=[Bi]O[Bi]=O",
+      "formula": "Bi2O3",
+      "density": 8.9018
+    },
+    "C2H4": {
+      "smiles": "C=C",
+      "formula": "C2H4",
+      "density": 0.95
+    },
+    "Cu": {
+      "smiles": "[Cu]",
+      "formula": "Cu",
+      "density": 8.96
+    },
+    "Sb": {
+      "smiles": "[Sb]",
+      "formula": "Sb",
+      "density": 6.697
+    },
+    "Sn": {
+      "smiles": "[Sn]",
+      "formula": "Sn",
+      "density": 7.265
+    },
+    "W": {
+      "smiles": "[W]",
+      "formula": "W",
+      "density": 19.254
+    },
+    "WO3": {
+      "smiles": "O=[W](=O)=O",
+      "formula": "WO3",
+      "density": 7.1596
+    }
+  }
+}

embeddings.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b1695dc7a6b6711814aa7f9d07b0f9a7a82c1c5dfaa3b02f1996f640929f3606
+size 40776

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cb8d7c4f6bbf706ec592a85bbf4aeb635ee8de6dade1ef82b468e588cb8a2e8b
+size 12029436

requirements.txt

@@ -0,0 +1,8 @@
+torch
+numpy
+safetensors
+huggingface_hub
+# optional: only for model.add_material() with a new SMILES
+transformers
+# optional: only for model.predict_frame()
+pandas

star_xpc.py

@@ -0,0 +1,527 @@
+"""STAR-XPC: Set-Transformer estimator of X-ray transmission through filled polymer composites.
+
+Self-contained inference code for the Hugging Face release. Needs ``torch``,
+``numpy`` and ``safetensors``; ``huggingface_hub`` to download from the Hub, and
+``transformers`` only to register a material that was not in training.
+
+    from star_xpc import STARXPC
+
+    model = STARXPC.from_pretrained("path/or/repo-id")
+    model.predict({"Bi": 0.40, "W": 0.30}, thickness_mm=[0.5, 1.0, 2.0], kvp=80)
+"""
+
+from __future__ import annotations
+
+import json
+import math
+import re
+import warnings
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn as nn
+from safetensors.torch import load_file
+
+__all__ = ["STARXPC", "SetTransformer", "ChemBERTaEmbedder", "UnknownMaterialError", "parse_formula"]
+
+
+# Set Transformer
+
+class Attention(nn.Module):
+    """Scaled dot-product attention."""
+
+    def __init__(self, temperature):
+        super().__init__()
+        self.temperature = temperature
+        self.softmax = nn.Softmax(dim=2)
+
+    def forward(self, queries, keys, values):
+        attention = self.softmax(torch.bmm(queries, keys.transpose(1, 2)) / self.temperature)
+        return torch.bmm(attention, values)
+
+
+class MultiheadAttention(nn.Module):
+    def __init__(self, d, h):
+        super().__init__()
+        assert d % h == 0
+        self.h = h
+        self.project_queries = nn.Linear(d, d)
+        self.project_keys    = nn.Linear(d, d)
+        self.project_values  = nn.Linear(d, d)
+        self.concatenation   = nn.Linear(d, d)
+        self.attention       = Attention(temperature=(d // h) ** 0.5)
+
+    def forward(self, queries, keys, values):
+        h = self.h
+        b, n, d = queries.size()
+        m = keys.size(1)
+        p = d // h
+
+        q = self.project_queries(queries).view(b, n, h, p).permute(2, 0, 1, 3).contiguous().view(h * b, n, p)
+        k = self.project_keys(keys).view(b, m, h, p).permute(2, 0, 1, 3).contiguous().view(h * b, m, p)
+        v = self.project_values(values).view(b, m, h, p).permute(2, 0, 1, 3).contiguous().view(h * b, m, p)
+
+        out = self.attention(q, k, v)                                              # [hb, n, p]
+        out = out.view(h, b, n, p).permute(1, 2, 0, 3).contiguous().view(b, n, d)  # [b, n, d]
+        return self.concatenation(out)
+
+
+class RFF(nn.Module):
+    """Row-wise feed-forward network."""
+
+    def __init__(self, d, num_layers=2):
+        super().__init__()
+        layers = []
+        for _ in range(num_layers):
+            layers.extend([nn.Linear(d, d), nn.ReLU(inplace=True)])
+        self.layers = nn.Sequential(*layers)
+
+    def forward(self, x):
+        return self.layers(x)
+
+
+class MultiheadAttentionBlock(nn.Module):
+    def __init__(self, d, h, rff):
+        super().__init__()
+        self.multihead = MultiheadAttention(d, h)
+        self.layer_norm1 = nn.LayerNorm(d)
+        self.layer_norm2 = nn.LayerNorm(d)
+        self.rff = rff
+
+    def forward(self, x, y):
+        h_ = self.layer_norm1(x + self.multihead(x, y, y))
+        return self.layer_norm2(h_ + self.rff(h_))
+
+
+class SetAttentionBlock(nn.Module):
+    def __init__(self, d, h, rff):
+        super().__init__()
+        self.mab = MultiheadAttentionBlock(d, h, rff)
+
+    def forward(self, x):
+        return self.mab(x, x)
+
+
+class InducedSetAttentionBlock(nn.Module):
+    def __init__(self, d, m, h, rff1, rff2):
+        super().__init__()
+        self.mab1 = MultiheadAttentionBlock(d, h, rff1)
+        self.mab2 = MultiheadAttentionBlock(d, h, rff2)
+        self.inducing_points = nn.Parameter(torch.zeros(1, m, d))
+
+    def forward(self, x):
+        p = self.inducing_points.repeat(x.size(0), 1, 1)
+        return self.mab2(x, self.mab1(p, x))
+
+
+class PoolingMultiheadAttention(nn.Module):
+    def __init__(self, d, k, h, rff):
+        super().__init__()
+        self.mab = MultiheadAttentionBlock(d, h, rff)
+        self.seed_vectors = nn.Parameter(torch.zeros(1, k, d))
+
+    def forward(self, z):
+        s = self.seed_vectors.repeat(z.size(0), 1, 1)
+        return self.mab(s, z)
+
+
+class SetTransformer(nn.Module):
+    def __init__(self, in_dimension, out_dimension, d=128, m=16, h=4, k=4,
+                 num_isabs=2, num_rff_layers=3):
+        super().__init__()
+        self.embed = nn.Sequential(nn.Linear(in_dimension, d), nn.ReLU(inplace=True))
+        self.encoder = nn.ModuleList([
+            InducedSetAttentionBlock(d, m, h, RFF(d, num_rff_layers), RFF(d, num_rff_layers))
+            for _ in range(num_isabs)
+        ])
+        self.pma = PoolingMultiheadAttention(d, k, h, RFF(d, num_rff_layers))
+        self.sab = SetAttentionBlock(d, h, RFF(d, num_rff_layers))
+        self.predictor = nn.Linear(k * d, out_dimension)
+
+    def forward(self, x):
+        """x: [b, n, in_dimension] -> [b, out_dimension]"""
+        x = self.embed(x)
+        for blk in self.encoder:
+            x = blk(x)
+        x = self.sab(self.pma(x))
+        b, k, d = x.shape
+        return self.predictor(x.view(b, k * d))
+
+
+# ChemBERTa embeddings
+
+class ChemBERTaEmbedder:
+    """768-d SMILES embedding: mean of the last 4 hidden layers, attention-masked
+    mean over tokens, then L2-normalised."""
+
+    def __init__(self, model_id="seyonec/ChemBERTa-zinc-base-v1", revision=None,
+                 max_length=64, device="cpu"):
+        from transformers import AutoModel, AutoTokenizer
+
+        self.tokenizer = AutoTokenizer.from_pretrained(model_id, revision=revision)
+        self.model = AutoModel.from_pretrained(model_id, revision=revision).to(device).eval()
+        self.max_length = max_length
+        self.device = device
+
+    @torch.no_grad()
+    def __call__(self, smiles: str) -> torch.Tensor:
+        n_tokens = len(self.tokenizer(smiles)["input_ids"])
+        if n_tokens > self.max_length:
+            warnings.warn(f"SMILES {smiles!r} is {n_tokens} tokens; only the first "
+                          f"{self.max_length} are embedded (as in training).", stacklevel=2)
+        enc = self.tokenizer(smiles, return_tensors="pt", padding=True, truncation=True,
+                             max_length=self.max_length).to(self.device)
+        hidden = self.model(**enc, output_hidden_states=True).hidden_states[-4:]
+        layers = torch.stack(hidden, dim=0).mean(dim=0)          # [1, seq, 768]
+        mask = enc["attention_mask"].unsqueeze(-1)               # [1, seq, 1]
+        emb = ((layers * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1e-9)).squeeze(0).cpu()
+        norm = emb.norm(p=2)
+        if norm.item() == 0:
+            raise ValueError(f"Zero-norm embedding for SMILES {smiles!r}")
+        return emb / norm
+
+
+# Composite physics features
+
+# symbol: (Z, atomic weight [g/mol], mean excitation energy I [eV]).
+_ELEMENTS = {
+    "H": (1, 1.00794, 19.2), "He": (2, 4.002602, 41.8), "Li": (3, 6.94, 40),
+    "Be": (4, 9.0121831, 63.7), "B": (5, 10.81, 76), "C": (6, 12.0107, 81),
+    "N": (7, 14.0067, 82), "O": (8, 15.999, 95), "F": (9, 18.998403163, 115),
+    "Ne": (10, 20.1797, 137), "Na": (11, 22.98976928, 149), "Mg": (12, 24.305, 156),
+    "Al": (13, 26.9815385, 166), "Si": (14, 28.085, 173), "P": (15, 30.973761998, 173),
+    "S": (16, 32.065, 180), "Cl": (17, 35.45, 174), "Ar": (18, 39.948, 188),
+    "K": (19, 39.0983, 190), "Ca": (20, 40.078, 191), "Sc": (21, 44.955908, 216),
+    "Ti": (22, 47.867, 233), "V": (23, 50.9415, 245), "Cr": (24, 51.9961, 257),
+    "Mn": (25, 54.938044, 272), "Fe": (26, 55.845, 286), "Co": (27, 58.933194, 297),
+    "Ni": (28, 58.6934, 311), "Cu": (29, 63.546, 322), "Zn": (30, 65.38, 330),
+    "Ga": (31, 69.723, 334), "Ge": (32, 72.63, 350), "As": (33, 74.921595, 347),
+    "Se": (34, 78.971, 348), "Br": (35, 79.904, 343), "Kr": (36, 83.798, 352),
+    "Rb": (37, 85.4678, 363), "Sr": (38, 87.62, 366), "Y": (39, 88.90584, 379),
+    "Zr": (40, 91.224, 393), "Nb": (41, 92.90637, 417), "Mo": (42, 95.95, 424),
+    "Tc": (43, 98.0, 428), "Ru": (44, 101.07, 441), "Rh": (45, 102.9055, 449),
+    "Pd": (46, 106.42, 470), "Ag": (47, 107.8682, 470), "Cd": (48, 112.414, 469),
+    "In": (49, 114.818, 488), "Sn": (50, 118.71, 488), "Sb": (51, 121.76, 487),
+    "Te": (52, 127.6, 485), "I": (53, 126.90447, 491), "Xe": (54, 131.293, 482),
+    "Cs": (55, 132.90545196, 488), "Ba": (56, 137.327, 491), "La": (57, 138.90547, 501),
+    "Ce": (58, 140.116, 523), "Pr": (59, 140.90766, 535), "Nd": (60, 144.242, 546),
+    "Pm": (61, 145.0, 560), "Sm": (62, 150.36, 574), "Eu": (63, 151.964, 580),
+    "Gd": (64, 157.25, 591), "Tb": (65, 158.92535, 614), "Dy": (66, 162.5, 628),
+    "Ho": (67, 164.93033, 650), "Er": (68, 167.259, 658), "Tm": (69, 168.93422, 674),
+    "Yb": (70, 173.045, 684), "Lu": (71, 174.9668, 694), "Hf": (72, 178.49, 705),
+    "Ta": (73, 180.94788, 718), "W": (74, 183.84, 727), "Re": (75, 186.207, 736),
+    "Os": (76, 190.23, 746), "Ir": (77, 192.217, 757), "Pt": (78, 195.084, 790),
+    "Au": (79, 196.966569, 790), "Hg": (80, 200.592, 800), "Tl": (81, 204.38, 810),
+    "Pb": (82, 207.231, 823), "Bi": (83, 208.9804, 823), "Po": (84, 209.0, 830),
+    "At": (85, 210.0, 825), "Rn": (86, 222.0, 794), "Fr": (87, 223.0, 827),
+    "Ra": (88, 226.0, 826), "Ac": (89, 227.0, 841), "Th": (90, 232.0377, 847),
+    "Pa": (91, 231.03588, 878), "U": (92, 238.02891, 890), "Np": (93, 237.0, 902),
+    "Pu": (94, 244.0, 921), "Am": (95, 243.0, 934), "Cm": (96, 247.0, 939),
+    "Bk": (97, 247.0, 952), "Cf": (98, 251.0, 966), "Es": (99, 252.0, 980),
+}
+
+_FORMULA_TOKEN = re.compile(r"([A-Z][a-z]?|\(|\)|\d+(?:\.\d+)?)")
+
+
+def parse_formula(formula: str) -> dict[str, float]:
+    """``"Gd2O3"`` -> ``{"Gd": 2, "O": 3}``. Understands parentheses
+    (``"Mg(OH)2"``) and fractional counts; raises ValueError otherwise."""
+    tokens = _FORMULA_TOKEN.findall(formula)
+    if not formula or "".join(tokens) != formula:
+        raise ValueError(f"Cannot parse chemical formula {formula!r}")
+    stack, i = [{}], 0
+    while i < len(tokens):
+        tok = tokens[i]
+        i += 1
+        if tok == "(":
+            stack.append({})
+            continue
+        if tok[0].isdigit():
+            raise ValueError(f"Unexpected number in formula {formula!r}")
+        mult = 1.0
+        if i < len(tokens) and tokens[i][0].isdigit():
+            mult = float(tokens[i])
+            i += 1
+        if tok == ")":
+            if len(stack) == 1:
+                raise ValueError(f"Unbalanced ')' in formula {formula!r}")
+            for el, c in stack.pop().items():
+                stack[-1][el] = stack[-1].get(el, 0.0) + c * mult
+        else:
+            if tok not in _ELEMENTS:
+                raise ValueError(f"Unknown element {tok!r} in formula {formula!r}")
+            stack[-1][tok] = stack[-1].get(tok, 0.0) + mult
+    if len(stack) != 1:
+        raise ValueError(f"Unbalanced '(' in formula {formula!r}")
+    return stack[0]
+
+
+_INVALID = {"-1", "-1.0", "", "nan", "None"}
+
+
+def _is_missing(v) -> bool:
+    return v is None or (isinstance(v, float) and math.isnan(v)) or str(v).strip() in _INVALID
+
+
+class UnknownMaterialError(KeyError):
+    def __str__(self):
+        return self.args[0]
+
+
+# Public predictor
+
+class STARXPC:
+    """Predicts the transmission ratio I/I0 of a filler + polyethylene composite slab."""
+
+    def __init__(self, model: SetTransformer, embeddings: dict, config: dict, device="cpu"):
+        self.device = torch.device(device)
+        self.model = model.to(self.device).eval()
+        self.config = config
+        self.embeddings = {k: v.float() for k, v in embeddings.items()}
+        self.materials = {k: dict(v) for k, v in config["materials"].items()}
+        self._embedder = None
+        self._warned = set()
+
+    @classmethod
+    def from_pretrained(cls, path_or_repo_id, revision=None, device="cpu") -> "STARXPC":
+        """Load from a local folder or a Hugging Face repo id."""
+        path = Path(path_or_repo_id)
+        if not path.is_dir():
+            from huggingface_hub import snapshot_download
+            path = Path(snapshot_download(str(path_or_repo_id), revision=revision))
+
+        config = json.loads((path / "config.json").read_text())
+        model = SetTransformer(**config["architecture"])
+        model.load_state_dict(load_file(path / "model.safetensors"), strict=True)
+        embeddings = load_file(path / "embeddings.safetensors")
+        return cls(model, embeddings, config, device=device)
+
+    # materials
+
+    @property
+    def available_materials(self) -> list[str]:
+        """Trained materials plus any registered with :meth:`add_material`."""
+        return sorted(self.embeddings)
+
+    def add_material(self, name: str, smiles: str, density: float | None = None,
+                     formula: str | dict | None = None, overwrite: bool = False) -> torch.Tensor:
+        """Register a new material from its SMILES (requires ``transformers``).
+
+        The SMILES is embedded with ChemBERTa exactly as the training materials
+        were. ``density`` (g/cm3) and ``formula`` (read from ``name`` when it is a
+        chemical formula, e.g. ``"Pb"`` or ``"Bi2O3"``) let :meth:`predict`
+        compute mass_den and mean_excit; without them, pass those to
+        :meth:`predict` yourself.
+        """
+        if name in self.embeddings and not overwrite:
+            raise ValueError(f"{name!r} already exists; pass overwrite=True to replace it.")
+        if formula is None:
+            try:
+                parse_formula(name)
+                formula = name
+            except ValueError:
+                formula = None
+        elif isinstance(formula, str):
+            parse_formula(formula)  # validate early
+        if self._embedder is None:
+            cfg = self.config["embedding"]
+            self._embedder = ChemBERTaEmbedder(cfg["model_id"], cfg.get("revision"),
+                                               cfg["max_length"], device=self.device)
+        self.embeddings[name] = self._embedder(smiles)
+        self.materials[name] = {"smiles": smiles, "formula": formula, "density": density}
+        return self.embeddings[name]
+
+    def _resolve(self, name) -> str:
+        name = str(name).strip()
+        if name in self.embeddings:
+            return name
+        matches = [k for k in self.embeddings if k.lower() == name.lower()]
+        if len(matches) == 1:
+            return matches[0]
+        raise UnknownMaterialError(
+            f"{name!r} is not a trained or registered material. Register it from its SMILES "
+            "(needs `pip install transformers`):\n\n"
+            f"    model.add_material({name!r}, smiles=\"...\", density=...)  # density in g/cm3\n\n"
+            "A pure element is a bracket atom, e.g. model.add_material(\"Pb\", smiles=\"[Pb]\", "
+            "density=11.34).\nTrained materials: " + ", ".join(sorted(self.config["materials"]))
+        )
+
+    def _counts(self, name: str) -> dict:
+        formula = self.materials[name].get("formula")
+        if not formula:
+            raise ValueError(f"{name!r} has no chemical formula: pass formula= to add_material(), "
+                             "or mean_excit= to predict().")
+        return dict(formula) if isinstance(formula, dict) else parse_formula(formula)
+
+    def _density(self, name: str) -> float:
+        rho = self.materials[name].get("density")
+        if not rho:
+            raise ValueError(f"{name!r} has no density: pass density= to add_material(), "
+                             "or mass_den= to predict().")
+        return float(rho)
+
+    def mass_density(self, mixture: dict) -> float:
+        """Composite density (g/cm3), inverse rule of mixtures over weight fractions."""
+        return 1.0 / sum(w / self._density(self._resolve(n)) for n, w in mixture.items())
+
+    def mean_excitation(self, mixture: dict) -> float:
+        """Composite mean excitation energy (eV), Bragg additivity rule as in PENELOPE."""
+        el_w = {}
+        for name, w in mixture.items():
+            counts = self._counts(self._resolve(name))
+            molar = sum(_ELEMENTS[e][1] * c for e, c in counts.items())
+            for e, c in counts.items():
+                el_w[e] = el_w.get(e, 0.0) + w * _ELEMENTS[e][1] * c / molar
+        z_over_a = {e: w * _ELEMENTS[e][0] / _ELEMENTS[e][1] for e, w in el_w.items()}
+        return math.exp(sum(za * math.log(_ELEMENTS[e][2]) for e, za in z_over_a.items())
+                        / sum(z_over_a.values()))
+
+    def composite_properties(self, fillers: dict, polymer_fraction: float) -> tuple[float, float]:
+        """(mass_den [g/cm3], mean_excit [eV]) of fillers + polyethylene."""
+        mixture = self._mixture(fillers, polymer_fraction)
+        return self.mass_density(mixture), self.mean_excitation(mixture)
+
+    def _mixture(self, fillers, polymer_fraction):
+        mixture = {self._resolve(k): float(v) for k, v in fillers.items()}
+        if polymer_fraction > 0:
+            poly = self.config["polymer"]
+            mixture[poly] = mixture.get(poly, 0.0) + polymer_fraction
+        return mixture
+
+    def _warn_untrained(self, names):
+        for name in names:
+            if name not in self.config["materials"] and name not in self._warned:
+                self._warned.add(name)
+                warnings.warn(
+                    f"{name!r} is not in the training data; its prediction relies on how well its "
+                    "ChemBERTa embedding generalises. Treat it as an estimate.", stacklevel=3)
+
+    # prediction
+
+    def predict(self, fillers: dict, thickness_mm, kvp, polymer_fraction: float | None = None,
+                mass_den: float | None = None, mean_excit: float | None = None):
+        """Transmission ratio I/I0 for one composite.
+
+        fillers:          weight fractions of the whole composite, e.g. {"Bi": 0.40, "W": 0.30}.
+        thickness_mm:     slab thickness in mm (scalar or array).
+        kvp:              X-ray tube voltage in kVp (scalar or array, broadcast with thickness).
+        polymer_fraction: polyethylene weight fraction; defaults to 1 - sum(fillers).
+        mass_den, mean_excit: g/cm3 and eV; computed from the composition when omitted.
+
+        Returns a float for scalar inputs, otherwise a numpy array.
+        """
+        resolved = {}
+        for k, v in fillers.items():
+            if float(v) > 0:
+                key = self._resolve(k)
+                resolved[key] = resolved.get(key, 0.0) + float(v)
+        fillers = resolved
+        if not fillers:
+            raise ValueError("At least one filler with a positive weight fraction is required.")
+        if len(fillers) > self.config["max_slots"] - 1:
+            raise ValueError(f"At most {self.config['max_slots'] - 1} fillers are supported.")
+        total = sum(fillers.values())
+        if total > 1.0 + 1e-6:
+            raise ValueError(f"Filler fractions sum to {total:.3f}; give weight fractions "
+                             "(0.40 for 40 wt%), not percentages.")
+        if polymer_fraction is None:
+            polymer_fraction = max(0.0, 1.0 - total)
+        if abs(total + polymer_fraction - 1.0) > 1e-3:
+            raise ValueError(f"Fillers ({total:.3f}) + polymer ({polymer_fraction:.3f}) must sum to 1.")
+
+        mixture = self._mixture(fillers, polymer_fraction)
+        mass_den = self.mass_density(mixture) if mass_den is None else mass_den
+        mean_excit = self.mean_excitation(mixture) if mean_excit is None else mean_excit
+
+        t, v = np.broadcast_arrays(np.asarray(thickness_mm, dtype=float), np.asarray(kvp, dtype=float))
+        shape = t.shape
+        t, v = t.ravel(), v.ravel()
+        self._check_domain(t, v, polymer_fraction)
+        self._warn_untrained(fillers)
+        
+        pairs = [(name, 100.0 * w) for name, w in fillers.items()]
+        if polymer_fraction > 0:
+            pairs.append((self.config["polymer"], polymer_fraction))
+
+        y = self._forward(
+            [pairs] * len(t),
+            spc=v,
+            mass_den=np.full(len(t), float(mass_den)),
+            geo=t / self.config["thickness_unit_mm"],
+            mean_excit=np.full(len(t), float(mean_excit)),
+        )
+        return float(y[0]) if shape == () else y.reshape(shape)
+
+    def predict_frame(self, df) -> np.ndarray:
+        """Predict for a pandas DataFrame in the training-CSV schema
+        (molecule_1..10, ratio_1..10, polymer, polymer_ratio, geo_number_mm,
+        spc, mass_den, mean_excit)."""
+        slots = self.config["max_slots"]
+        mol_cols = [f"molecule_{j}" for j in range(1, slots)] + ["polymer"]
+        rat_cols = [f"ratio_{j}" for j in range(1, slots)] + ["polymer_ratio"]
+
+        sets, used = [], set()
+        for row in df.to_dict("records"):
+            pairs = []
+            for mcol, rcol in zip(mol_cols, rat_cols):
+                m, r = row.get(mcol), row.get(rcol)
+                if _is_missing(m) or _is_missing(r):
+                    continue
+                try:
+                    r = float(r)
+                except (TypeError, ValueError):
+                    continue
+                if r > 0:
+                    name = self._resolve(m)
+                    pairs.append((name, r))
+                    used.add(name)
+            sets.append(pairs)
+        self._warn_untrained(sorted(used))
+
+        col = lambda c: df[c].to_numpy(dtype=float)
+        return self._forward(sets, col("spc"), col("mass_den"), col("geo_number_mm"), col("mean_excit"))
+
+    def _forward(self, sets, spc, mass_den, geo, mean_excit, batch_size=1024) -> np.ndarray:
+        n = self.config["normalization"]
+        slots = self.config["max_slots"]
+        dim = self.config["embedding"]["dim"]
+
+        gf = np.stack([
+            (spc        - n["spc_mean"])        / n["spc_std"],
+            (mass_den   - n["mass_den_mean"])   / n["mass_den_std"],
+            (geo        - n["thick_mean"])      / n["thick_std"],
+            (mean_excit - n["mean_excit_mean"]) / n["mean_excit_std"],
+        ], axis=1)
+
+        out = []
+        for start in range(0, len(sets), batch_size):
+            chunk = sets[start:start + batch_size]
+            emb = torch.zeros(len(chunk), slots, dim)
+            rat = torch.zeros(len(chunk), slots, 1)
+            for b, pairs in enumerate(chunk):
+                total = sum(r for _, r in pairs)
+                if not pairs or total <= 0:
+                    raise ValueError(f"Row {start + b} has no material with a positive ratio.")
+                for s, (name, r) in enumerate(pairs[:slots]):
+                    emb[b, s] = self.embeddings[name]
+                    rat[b, s, 0] = r / total
+            g = torch.tensor(gf[start:start + len(chunk)], dtype=torch.float32)
+            x = torch.cat([emb, rat, g.unsqueeze(1).expand(-1, slots, -1)], dim=-1).to(self.device)
+            with torch.no_grad():
+                out.append(self.model(x).squeeze(-1).cpu().numpy())
+        return np.concatenate(out) if out else np.empty(0, dtype=np.float32)
+
+    def _check_domain(self, thickness_mm, kvp, polymer_fraction):
+        dom = self.config.get("training_domain", {})
+        if "kvp" in dom and not np.isin(kvp, dom["kvp"]).all():
+            warnings.warn(f"Model was trained only at {dom['kvp']} kVp; other voltages are interpolation.",
+                          stacklevel=3)
+        if "thickness_mm" in dom:
+            lo, hi = dom["thickness_mm"]
+            if (thickness_mm < lo).any() or (thickness_mm > hi).any():
+                warnings.warn(f"Thickness outside the training range {lo}-{hi} mm.", stacklevel=3)
+        if "polymer_fraction" in dom:
+            lo, hi = dom["polymer_fraction"]
+            if not lo - 1e-9 <= polymer_fraction <= hi + 1e-9:
+                warnings.warn(f"Polymer fraction outside the training range {lo}-{hi}.", stacklevel=3)