Blame for README.md

---
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

<!-- TODO: add the paper's BibTeX once it is available -->