--- 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 training data were generated with [attenupy](https://github.com/cmpe-netlab41/attenupy), which designs, runs and collects these penEasy shielding studies. 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.19 and 0.20 against penEasy at 100 kVp. That is somewhat above the held-out sets of trained materials, which range from 0.01 to 0.13 (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. [attenupy](https://github.com/cmpe-netlab41/attenupy) sets up and runs those simulations. ## 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.129 | 0.0109 | 6.86 % | | W–Cu–B, 100 kVp | 0.00724 | 0.0042 | 0.65 % | | Bi–Sb, 120 kVp | 0.0469 | 0.0106 | 2.65 % | | Bi₂O₃–BaSO₄, 80 kVp | 0.0865 | 0.0095 | 4.89 % | | BaSO₄–WO₃, 100 kVp | 0.0167 | 0.0064 | 1.78 % | | Bi₂O₃, 120 kVp | 0.0113 | 0.0053 | 1.22 % | | **Pb** 70 wt% (registered), 100 kVp | 0.192 | 0.0167 | 5.24 % | | **Gd** 70 wt% (registered), 100 kVp | 0.198 | 0.0195 | 4.03 % | 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. The z-score statistics are those of the model's training split and are stored in `config.json`. 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