upet 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,9 @@
1
+ @misc{PET-MAD-2025,
2
+ title={PET-MAD, a universal interatomic potential for advanced materials modeling},
3
+ author={Arslan Mazitov and Filippo Bigi and Matthias Kellner and Paolo Pegolo and Davide Tisi and Guillaume Fraux and Sergey Pozdnyakov and Philip Loche and Michele Ceriotti},
4
+ year={2025},
5
+ eprint={2503.14118},
6
+ archivePrefix={arXiv},
7
+ primaryClass={cond-mat.mtrl-sci},
8
+ url={https://arxiv.org/abs/2503.14118},
9
+ }
upet-0.1.0/LICENSE ADDED
@@ -0,0 +1,28 @@
1
+ BSD 3-Clause License
2
+
3
+ Copyright (c) 2023, Laboratory of Computational Science and Modeling
4
+
5
+ Redistribution and use in source and binary forms, with or without
6
+ modification, are permitted provided that the following conditions are met:
7
+
8
+ 1. Redistributions of source code must retain the above copyright notice, this
9
+ list of conditions and the following disclaimer.
10
+
11
+ 2. Redistributions in binary form must reproduce the above copyright notice,
12
+ this list of conditions and the following disclaimer in the documentation
13
+ and/or other materials provided with the distribution.
14
+
15
+ 3. Neither the name of the copyright holder nor the names of its
16
+ contributors may be used to endorse or promote products derived from
17
+ this software without specific prior written permission.
18
+
19
+ THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
20
+ AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
21
+ IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
22
+ DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
23
+ FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
24
+ DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
25
+ SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
26
+ CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
27
+ OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
28
+ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
upet-0.1.0/MANIFEST.in ADDED
@@ -0,0 +1,15 @@
1
+ graft src
2
+
3
+ include LICENSE
4
+ include README.md
5
+ include *.bib
6
+
7
+ prune docs
8
+ prune tests
9
+ prune .github
10
+ prune .tox
11
+
12
+ exclude .gitignore
13
+ exclude tox.ini
14
+
15
+ global-exclude *.py[cod] __pycache__/* *.so *.dylib
upet-0.1.0/PKG-INFO ADDED
@@ -0,0 +1,586 @@
1
+ Metadata-Version: 2.4
2
+ Name: upet
3
+ Version: 0.1.0
4
+ Summary: A universal interatomic potential for advanced materials modeling
5
+ Author: Filippo Bigi, Matthias Kellner, Paolo Pegolo, Davide Tisi, Guillaume Fraux, Sergey Pozdnyakov, Philip Loche
6
+ Author-email: Arslan Mazitov <arslan.mazitov@epfl.ch>, Michele Ceriotti <michele.ceriotti@epfl.ch>
7
+ Maintainer-email: Arslan Mazitov <arslan.mazitov@epfl.ch>, Filippo Bigi <filippo.bigi@epfl.ch>
8
+ License: BSD-3-Clause
9
+ Project-URL: repository, https://github.com/lab-cosmo/upet
10
+ Keywords: machine learning,molecular modeling
11
+ Classifier: Development Status :: 4 - Beta
12
+ Classifier: Intended Audience :: Science/Research
13
+ Classifier: Operating System :: POSIX
14
+ Classifier: Operating System :: MacOS :: MacOS X
15
+ Classifier: Operating System :: Microsoft :: Windows
16
+ Classifier: Programming Language :: Python
17
+ Classifier: Programming Language :: Python :: 3
18
+ Classifier: Topic :: Scientific/Engineering
19
+ Classifier: Topic :: Scientific/Engineering :: Bio-Informatics
20
+ Classifier: Topic :: Scientific/Engineering :: Chemistry
21
+ Classifier: Topic :: Scientific/Engineering :: Physics
22
+ Classifier: Topic :: Software Development :: Libraries
23
+ Classifier: Topic :: Software Development :: Libraries :: Python Modules
24
+ Requires-Python: >=3.10
25
+ Description-Content-Type: text/markdown
26
+ License-File: LICENSE
27
+ Requires-Dist: metatrain==2026.1
28
+ Requires-Dist: huggingface_hub
29
+ Requires-Dist: hf_xet
30
+ Requires-Dist: packaging
31
+ Requires-Dist: platformdirs
32
+ Requires-Dist: tqdm
33
+ Requires-Dist: scipy<1.17.0,>=1.15.0
34
+ Provides-Extra: deprecated
35
+ Requires-Dist: pet-neighbors-convert; extra == "deprecated"
36
+ Requires-Dist: torch-geometric; extra == "deprecated"
37
+ Requires-Dist: matscipy; extra == "deprecated"
38
+ Requires-Dist: scikit-learn; extra == "deprecated"
39
+ Requires-Dist: pathos; extra == "deprecated"
40
+ Requires-Dist: scipy; extra == "deprecated"
41
+ Dynamic: license-file
42
+
43
+ <div align="center" width="600">
44
+ <picture>
45
+ <source srcset="https://github.com/lab-cosmo/upet/raw/refs/heads/main/docs/static/upet-logo-with-text-dark.svg" media="(prefers-color-scheme: dark)">
46
+ <img src="https://github.com/lab-cosmo/upet/raw/refs/heads/main/docs/static/upet-logo-with-text.svg" alt="Figure">
47
+ </picture>
48
+ </div>
49
+
50
+ > [!WARNING]
51
+ > This repository is a successor of the PET-MAD repository, which is now deprecated.
52
+ > The package has been renamed to **UPET** to reflect the broader scope of the models
53
+ > and functionalities provided, going beyond the original PET-MAD model.
54
+ > Please use the version `1.4.4` of PET-MAD package if you want to use the old API.
55
+ > The older version of the README file with documentation is avaiable [here](docs/README_OLD.md).
56
+ > The migration guide from PET-MAD to UPET is available [here](docs/UPET_MIGRATION_GUIDE.md).
57
+
58
+ # UPET: Universal Models for Advanced Atomistic Simulations
59
+
60
+ This repository contains **UPET** models - universal interatomic potentials for
61
+ advanced materials modeling across the periodic table. This models are based on
62
+ the **Point Edge Transformer (PET)** architecture, trained on various popular
63
+ materials datasets, and are capable of predicting energies and forces in complex
64
+ atomistic simulations.
65
+
66
+ In addition, it contains **PET-MAD-DOS** - a universal model for predicting
67
+ the density of states (DOS) of materials, as well as their Fermi levels and bandgaps.
68
+ **PET-MAD-DOS** is using a slightly modified **PET** architecture, and the **MAD** dataset.
69
+
70
+ ## Key Features
71
+
72
+ - **Universality**: UPET models are generally-applicable, and can be used for
73
+ predicting energies and forces, as well as the density of states, Fermi levels,
74
+ and bandgaps for a wide range of materials and molecules.
75
+ - **Accuracy**: UPET models achieve high accuracy in various types of atomistic
76
+ simulations of organic and inorganic systems, comparable with system-specific
77
+ models, while being fast and efficient.
78
+ - **Efficiency**: UPET models are highly computationally efficient and have low
79
+ memory usage, what makes them suitable for large-scale simulations.
80
+ - **Infrastructure**: Various MD engines are available for diverse research and
81
+ application needs.
82
+ - **HPC Compatibility**: Efficient in HPC environments for extensive simulations.
83
+
84
+ ## Table of Contents
85
+ 1. [Installation](#installation)
86
+ 2. [Pre-trained Models](#pre-trained-models)
87
+ 3. [Interfaces for Atomistic Simulations](#interfaces-for-atomistic-simulations)
88
+ 4. [Usage](#usage)
89
+ - [ASE Interface](#ase-interface)
90
+ - [Basic usage](#basic-usage)
91
+ - [Non-conservative (direct) forces and stresses prediction](#non-conservative-direct-forces-and-stresses-prediction)
92
+ - [Evaluating UPET models on a dataset](#evaluating-upet-models-on-a-dataset)
93
+ - [Running UPET models with LAMMPS](#running-upet-models-with-lammps)
94
+ - [Uncertainty Quantification](#uncertainty-quantification)
95
+ - [Running UPET models with empirical dispersion corrections](#running-upet-models-with-empirical-dispersion-corrections)
96
+ - [Calculating the DOS, Fermi levels, and bandgaps](#calculating-the-dos-fermi-levels-and-bandgaps)
97
+ - [Dataset visualization with the PET-MAD featurizer](#dataset-visualization-with-the-pet-mad-featurizer)
98
+ 5. [Examples](#examples)
99
+ 6. [Fine-tuning](#fine-tuning)
100
+ 7. [Documentation](#documentation)
101
+ 8. [Citing PET-MAD](#citing-pet-mad)
102
+
103
+ ## Installation
104
+
105
+ You can install UPET using pip (not available yet):
106
+
107
+ ```bash
108
+ pip install upet
109
+ ```
110
+
111
+ Or directly from the GitHub repository:
112
+
113
+ ```bash
114
+ pip install git+https://github.com/lab-cosmo/upet.git
115
+ ```
116
+
117
+ ## Pre-trained Models
118
+
119
+ Currently, we provide the following pre-trained models:
120
+
121
+ | Name | Level of theory | Available sizes | To be used for | Training set |
122
+ |:------------|:-----------------------:|:----------------------:|:-----------------------:|:---------------------:|
123
+ | PET-MAD | PBEsol | S | materials and molecules | MAD |
124
+ | PET-OMAD | PBEsol | L | materials and molecules | OMat -> MAD |
125
+ | PET-OMATPES | r2SCAN | L | materials | OMat -> MATPES |
126
+ | PET-OMat | PBE | XS, S, M, L, XL | materials | OMat |
127
+ | PET-OAM | PBE (Materials Project) | L, XL | materials | OMat -> sAlex+MPtrj |
128
+ | PET-SPICE | ωB97M-D3 | S, L | molecules | SPICE |
129
+
130
+ All the checkpoints are available on the HuggingFace [repository](https://huggingface.co/lab-cosmo/upet).
131
+
132
+ ## Interfaces for Atomistic Simulations
133
+
134
+ UPET integrates with the following atomistic simulation engines:
135
+
136
+ - **Atomic Simulation Environment (ASE)**
137
+ - **LAMMPS** (including the KOKKOS support)
138
+ - **i-PI**
139
+ - **OpenMM** (coming soon)
140
+ - **GROMACS** (coming soon)
141
+
142
+ ## Usage
143
+
144
+ ### ASE Interface
145
+
146
+ #### Basic usage
147
+
148
+ You can use the UPET calculator, which is compatible with the Atomic
149
+ Simulation Environment (ASE). Model name can be obtained from the
150
+ by combining the model name and the size, e.g., `pet-mad-s`, `pet-omat-l`, etc.
151
+
152
+ ```python
153
+ from upet.calculator import UPETCalculator
154
+ from ase.build import bulk
155
+
156
+ atoms = bulk("Si", cubic=True, a=5.43, crystalstructure="diamond")
157
+ calculator = UPETCalculator(model="pet-mad-s", version="1.0.2", device="cpu")
158
+ atoms.calc = calculator
159
+ energy = atoms.get_potential_energy()
160
+ forces = atoms.get_forces()
161
+ ```
162
+
163
+ These ASE methods are ideal for single-structure evaluations, but they are
164
+ inefficient for the evaluation on a large number of pre-defined structures. To
165
+ perform efficient evaluation in that case, read [here](docs/README_BATCHED.md).
166
+
167
+ #### Non-conservative (direct) forces and stresses prediction
168
+
169
+ UPET models also supports the direct prediction of forces and stresses. In that case,
170
+ the forces and stresses are predicted as separate targets along with the energy
171
+ target, i.e. not computed as derivatives of the energy using the PyTorch
172
+ automatic differentiation. This approach typically leads to 2-3x speedup in the
173
+ evaluation time, since backward pass is disabled. However, as discussed in [this
174
+ preprint](https://arxiv.org/abs/2412.11569) it requires additional care to avoid
175
+ instabilities during the molecular dynamics simulations.
176
+
177
+ To use the non-conservative forces and stresses, you need to set the `non_conservative` parameter to `True` when initializing the `UPETCalculator` class.
178
+
179
+ ```python
180
+ from upet.calculator import UPETCalculator
181
+ from ase.build import bulk
182
+
183
+ atoms = bulk("Si", cubic=True, a=5.43, crystalstructure="diamond")
184
+ calculator = UPETCalculator(model="pet-mad-s", version="1.1.0", device="cpu", non_conservative=True)
185
+ atoms.calc = calculator
186
+ energy = atoms.get_potential_energy() # energy is computed as usual
187
+ forces = atoms.get_forces() # forces now are predicted as a separate target
188
+ stresses = atoms.get_stress() # stresses now are predicted as a separate target
189
+ ```
190
+
191
+ More details on how to make the direct forces MD simulations reliable are provided
192
+ in the [Atomistic Cookbook](https://atomistic-cookbook.org/examples/pet-mad-nc/pet-mad-nc.html).
193
+
194
+ ### Evaluating UPET models on a dataset
195
+
196
+ Efficient evaluation of UPET models on a desired dataset is also available from the
197
+ command line via [`metatrain`](https://github.com/metatensor/metatrain), which
198
+ is installed as a dependency of UPET. To evaluate the model, you first need
199
+ to fetch the UPET model from the HuggingFace repository:
200
+
201
+ ```bash
202
+ mtt export https://huggingface.co/lab-cosmo/upet/resolve/main/models/pet-mad-s-v1.0.2.ckpt -o model.pt
203
+ ```
204
+
205
+ Alternatively, you can fetch and save the model using the UPET Python API:
206
+
207
+ ```py
208
+ import upet
209
+
210
+ # Saving the latest version of UPET to a TorchScript file
211
+ upet.save_upet(
212
+ model="pet-mad",
213
+ size="s",
214
+ version="v1.0.2",
215
+ output="model.pt",
216
+ )
217
+ ```
218
+
219
+ Both these commands will download the model and convert it to TorchScript format. Next,
220
+ you need to create the `options.yaml` file and specify the dataset you want to
221
+ evaluate the model on (where the dataset is stored in `extxyz` format):
222
+
223
+ ```yaml
224
+ systems: your-test-dataset.xyz
225
+ targets:
226
+ energy:
227
+ key: "energy"
228
+ unit: "eV"
229
+ ```
230
+
231
+ Then, you can use the `mtt eval` command to evaluate the model on a dataset:
232
+
233
+ ```bash
234
+ mtt eval model.pt options.yaml --batch-size=16 --output=predictions.xyz
235
+ ```
236
+
237
+ This will create a file called `predictions.xyz` with the predicted energies and
238
+ forces for each structure in the dataset. More details on how to use `metatrain`
239
+ can be found in the [Metatrain documentation](https://metatensor.github.io/metatrain/latest/getting-started/usage.html#evaluation).
240
+
241
+ ### Uncertainty Quantification
242
+
243
+ UPET models can also be used to calculate the uncertainty of the energy prediction.
244
+ This feature is particularly important if you are interested in probing the model
245
+ on the data that is far away from the training data. Another important use case
246
+ is a propagation of the uncertainty of the energy prediction to other observables,
247
+ like phase transition temperatures, diffusion coefficients, etc.
248
+
249
+ To evaluate the uncertainty of the energy prediction, or to get an ensemble of energy
250
+ predictions, you can use the `get_energy_uncertainty` and `get_energy_ensemble` methods
251
+ of the `PETMADCalculator` class:
252
+
253
+ ```python
254
+ from upet.calculator import UPETCalculator
255
+ from ase.build import bulk
256
+
257
+ atoms = bulk("Si", cubic=True, a=5.43, crystalstructure="diamond")
258
+ calculator = UPETCalculator(model="pet-mad-s", version="1.0.2", device="cpu", calculate_uncertainty=True, calculate_ensemble=True)
259
+ atoms.calc = calculator
260
+ energy = atoms.get_potential_energy()
261
+
262
+ energy_uncertainty = calculator.get_energy_uncertainty(atoms, per_atom=False)
263
+ energy_ensemble = calculator.get_energy_ensemble(atoms, per_atom=False)
264
+ ```
265
+
266
+ Please note that the uncertainty quantification and ensemble prediction accepts the
267
+ `per_atom` flag, which indicates whether the uncertainty/ensemble should be computed
268
+ per atom or for the whole system. More details on the uncertainty quantification and shallow
269
+ ensemble method can be found in [this](https://doi.org/10.1088/2632-2153/ad594a) and
270
+ [this](https://doi.org/10.1088/2632-2153/ad805f) papers.
271
+
272
+
273
+ ## Running UPET models with LAMMPS
274
+
275
+ ### 1. Install LAMMPS with metatomic support
276
+
277
+ To use UPET with LAMMPS, follow the instructions
278
+ [here](https://docs.metatensor.org/metatomic/latest/engines/lammps.html#how-to-install-the-code)
279
+ to install lammps-metatomic. We recomend you also use conda to install prebuilt lammps binaries.
280
+
281
+ ### 2. Run LAMMPS with UPET
282
+
283
+ #### 2.1. CPU version
284
+
285
+ Fetch the UPET checkpoint from the HuggingFace repository:
286
+
287
+ ```bash
288
+ mtt export https://huggingface.co/lab-cosmo/upet/resolve/main/models/pet-mad-s-v1.0.2.ckpt -o model.pt
289
+ ```
290
+
291
+ This will download the model and convert it to TorchScript format compatible
292
+ with LAMMPS, using the `metatomic` and `metatrain` libraries, which UPET is
293
+ based on.
294
+
295
+ Other pre-trained UPET models can be prepared in the same way, e.g.,
296
+ ```bash
297
+ mtt export https://huggingface.co/lab-cosmo/upet/resolve/main/models/pet-omat-xs-v1.0.0.ckpt -o model.pt
298
+ mtt export https://huggingface.co/lab-cosmo/upet/resolve/main/models/pet-omatpes-l-v0.1.0.ckpt -o model.pt
299
+ ```
300
+
301
+ Prepare a lammps input file using `pair_style metatomic` and defining the
302
+ mapping from LAMMPS types in the data file to elements UPET can handle using
303
+ `pair_coeff` syntax. Here we indicate that lammps atom type 1 is Silicon (atomic
304
+ number 14).
305
+
306
+ ```
307
+ units metal
308
+ atom_style atomic
309
+
310
+ read_data silicon.data
311
+
312
+ pair_style metatomic model.pt device cpu # Change device to 'cuda' evaluate UPET on GPU
313
+ pair_coeff * * 14
314
+
315
+ neighbor 2.0 bin
316
+ timestep 0.001
317
+
318
+ dump myDump all xyz 10 trajectory.xyz
319
+ dump_modify myDump element Si
320
+
321
+ thermo_style multi
322
+ thermo 1
323
+
324
+ velocity all create 300 87287 mom yes rot yes
325
+
326
+ fix 1 all nvt temp 300 300 0.10
327
+
328
+ run 100
329
+ ```
330
+
331
+ Create the **`silicon.data`** data file for a silicon system.
332
+
333
+ ```
334
+ # LAMMPS data file for Silicon unit cell
335
+ 8 atoms
336
+ 1 atom types
337
+
338
+ 0.0 5.43 xlo xhi
339
+ 0.0 5.43 ylo yhi
340
+ 0.0 5.43 zlo zhi
341
+
342
+ Masses
343
+
344
+ 1 28.084999992775295 # Si
345
+
346
+ Atoms # atomic
347
+
348
+ 1 1 0 0 0
349
+ 2 1 1.3575 1.3575 1.3575
350
+ 3 1 0 2.715 2.715
351
+ 4 1 1.3575 4.0725 4.0725
352
+ 5 1 2.715 0 2.715
353
+ 6 1 4.0725 1.3575 4.0725
354
+ 7 1 2.715 2.715 0
355
+ 8 1 4.0725 4.0725 1.3575
356
+ ```
357
+
358
+ ```bash
359
+ lmp -in lammps.in # For serial version
360
+ mpirun -np 1 lmp -in lammps.in # For MPI version
361
+ ```
362
+
363
+ #### 2.2. KOKKOS-enabled GPU version
364
+
365
+ Running LAMMPS with KOKKOS and GPU support is similar to the CPU version, but
366
+ you need to change the `lammps.in` slightly and run `lmp` binary with a few
367
+ additional flags.
368
+
369
+ The updated `lammps.in` file looks like this:
370
+
371
+ ```
372
+ package kokkos newton on neigh half
373
+
374
+ units metal
375
+ atom_style atomic/kk
376
+
377
+ read_data silicon.data
378
+
379
+ pair_style metatomic/kk model.pt # This will use the same device as the kokkos simulation
380
+ pair_coeff * * 14
381
+
382
+ neighbor 2.0 bin
383
+ timestep 0.001
384
+
385
+ dump myDump all xyz 10 trajectory.xyz
386
+ dump_modify myDump element Si
387
+
388
+ thermo_style multi
389
+ thermo 1
390
+
391
+ velocity all create 300 87287 mom yes rot yes
392
+
393
+ fix 1 all nvt temp 300 300 0.10
394
+
395
+ run_style verlet/kk
396
+ run 100
397
+ ```
398
+
399
+ The **silicon.data** file remains the same.
400
+
401
+ To run the KOKKOS-enabled version of LAMMPS, you need to run
402
+
403
+ ```bash
404
+ lmp -in lammps.in -k on g 1 -sf kk # For serial version
405
+ mpirun -np 1 lmp -in lammps.in -k on g 1 -sf kk # For MPI version
406
+ ```
407
+
408
+ Here, the `-k on g 1 -sf kk` flags are used to activate the KOKKOS
409
+ subroutines. Specifically `g 1` is used to specify, how many GPUs are the
410
+ simulation is parallelized over, so if running the large systems on two or more
411
+ GPUs, this number should be adjusted accordingly.
412
+
413
+
414
+ ### 3. Important Notes
415
+
416
+ - For **CPU calculations**, use a single MPI task unless simulating large
417
+ systems (30+ Å box size). Multi-threading can be enabled via:
418
+
419
+ ```bash
420
+ export OMP_NUM_THREADS=4
421
+ ```
422
+
423
+ - For **GPU calculations**, use **one MPI task per GPU**.
424
+
425
+ ## Running UPET models with empirical dispersion corrections
426
+
427
+ ### In **ASE**:
428
+
429
+ You can combine the UPET calculator with the torch based implementation of
430
+ the D3 dispersion correction of `pfnet-research` - `torch-dftd`:
431
+
432
+ Within the UPET environment you can install `torch-dftd` via:
433
+
434
+ ```bash
435
+ pip install torch-dftd
436
+ ```
437
+
438
+ Then you can use the `D3Calculator` class to combine the two calculators:
439
+
440
+ ```python
441
+ from torch_dftd.torch_dftd3_calculator import TorchDFTD3Calculator
442
+ from upet.calculator import UPETCalculator
443
+ from ase.calculators.mixing import SumCalculator
444
+
445
+ device = "cuda" if torch.cuda.is_available() else "cpu"
446
+
447
+ calc_MAD = UPETCalculator(model="pet-mad-s", version="1.0.2", device=device)
448
+ dft_d3 = TorchDFTD3Calculator(device=device, xc="pbesol", damping="bj")
449
+
450
+ combined_calc = SumCalculator([calc_MAD, dft_d3])
451
+
452
+ # assign the calculator to the atoms object
453
+ atoms.calc = combined_calc
454
+
455
+ ```
456
+
457
+
458
+ ## Calculating the DOS, Fermi levels, and bandgaps
459
+
460
+ UPET package also allows the use of the **PET-MAD-DOS** model to predict
461
+ electronic density of states of materials, as well as their Fermi levels and
462
+ bandgaps. The **PET-MAD-DOS** model is also available in the **ASE** interface.
463
+
464
+ ```python
465
+ from upet.calculator import PETMADDOSCalculator
466
+
467
+ atoms = bulk("Si", cubic=True, a=5.43, crystalstructure="diamond")
468
+ pet_mad_dos_calculator = PETMADDOSCalculator(version="latest", device="cpu")
469
+
470
+ energies, dos = pet_mad_dos_calculator.calculate_dos(atoms)
471
+ ```
472
+
473
+ Predicting the densities of states for every atom in the crystal,
474
+ or a list of atoms, is also possible:
475
+
476
+ ```python
477
+ # Calculating the DOS for every atom in the crystal
478
+ energies, dos_per_atom = pet_mad_dos_calculator.calculate_dos(atoms, per_atom=True)
479
+
480
+ # Calculating the DOS for a list of atoms
481
+ atoms_1 = bulk("Si", cubic=True, a=5.43, crystalstructure="diamond")
482
+ atoms_2 = bulk("C", cubic=True, a=3.55, crystalstructure="diamond")
483
+
484
+ energies, dos = pet_mad_dos_calculator.calculate_dos([atoms_1, atoms_2], per_atom=False)
485
+ ```
486
+
487
+ Finally, you can use the `calculate_bandgap` and `calculate_efermi` methods to
488
+ predict the bandgap and Fermi level for the crystal:
489
+
490
+ ```python
491
+ bandgap = pet_mad_dos_calculator.calculate_bandgap(atoms)
492
+ fermi_level = pet_mad_dos_calculator.calculate_efermi(atoms)
493
+ ```
494
+
495
+ You can also re-use the DOS calculated earlier:
496
+
497
+ ```python
498
+ bandgap = pet_mad_dos_calculator.calculate_bandgap(atoms, dos=dos)
499
+ fermi_level = pet_mad_dos_calculator.calculate_efermi(atoms, dos=dos)
500
+ ```
501
+
502
+ This option is also available for a list of `ase.Atoms` objects:
503
+
504
+ ```python
505
+ bandgaps = pet_mad_dos_calculator.calculate_bandgap([atoms_1, atoms_2], dos=dos)
506
+ fermi_levels = pet_mad_dos_calculator.calculate_efermi([atoms_1, atoms_2], dos=dos)
507
+ ```
508
+
509
+
510
+ ## Dataset visualization with the PET-MAD featurizer
511
+
512
+ You can use PET-MAD last-layer features together with a pre-trained
513
+ sketch-map dimensionality reduction to obtain 2D and 3D representations
514
+ of a dataset, e.g. to identify structural or chemical motifs.
515
+ This can be used as a stand-alone feature builder, or combined with
516
+ the [chemiscope viewer](https://chemiscope.org) to generate an
517
+ interactive visualization.
518
+
519
+ ```python
520
+ import ase.io
521
+ import chemiscope
522
+ from upet.explore import PETMADFeaturizer
523
+
524
+ featurizer = PETMADFeaturizer(version="latest")
525
+
526
+ # Load structures
527
+ frames = ase.io.read("dataset.xyz", ":")
528
+
529
+ # You can just compute features
530
+ features = featurizer(frames, None)
531
+
532
+ # Or create an interactive visualization in a Jupyter notebook
533
+ chemiscope.explore(
534
+ frames,
535
+ featurize=featurizer
536
+ )
537
+ ```
538
+
539
+ ## Examples
540
+
541
+ More examples for **ASE, i-PI, and LAMMPS** are available in the [Atomistic
542
+ Cookbook](https://atomistic-cookbook.org/examples/pet-mad/pet-mad.html).
543
+
544
+ ## Fine-tuning
545
+
546
+ UPET models can be fine-tuned using the
547
+ [Metatrain](https://metatensor.github.io/metatrain/latest/advanced-concepts/fine-tuning.html)
548
+ library.
549
+
550
+ ## Documentation
551
+
552
+ Additional documentation can be found in the
553
+ [metatensor](https://docs.metatensor.org),
554
+ [metatomic](https://docs.metatensor.org/metatomic) and
555
+ [metatrain](https://metatensor.github.io/metatrain/) repositories.
556
+
557
+ - [Training a model](https://metatensor.github.io/metatrain/latest/getting-started/usage.html#training)
558
+ - [Fine-tuning](https://metatensor.github.io/metatrain/latest/advanced-concepts/fine-tuning.html)
559
+ - [LAMMPS interface](https://docs.metatensor.org/metatomic/latest/engines/lammps.html)
560
+ - [i-PI interface](https://docs.metatensor.org/metatomic/latest/engines/ipi.html)
561
+
562
+ ## Citing UPET Models
563
+
564
+ If you use any of the UPET models in your research, please cite the corresponding articles:
565
+
566
+ ```bibtex
567
+ @misc{PET-MAD-2025,
568
+ title={PET-MAD as a lightweight universal interatomic potential for advanced materials modeling},
569
+ author={Mazitov, Arslan and Bigi, Filippo and Kellner, Matthias and Pegolo, Paolo and Tisi, Davide and Fraux, Guillaume and Pozdnyakov, Sergey and Loche, Philip and Ceriotti, Michele},
570
+ journal={Nature Communications},
571
+ volume={16},
572
+ number={1},
573
+ pages={10653},
574
+ year={2025},
575
+ publisher={Nature Publishing Group UK London},
576
+ url={https://doi.org/10.1038/s41467-025-65662-7},
577
+ }
578
+ @misc{PET-MAD-DOS-2025,
579
+ title={A universal machine learning model for the electronic density of states},
580
+ author={Wei Bin How and Pol Febrer and Sanggyu Chong and Arslan Mazitov and Filippo Bigi and Matthias Kellner and Sergey Pozdnyakov and Michele Ceriotti},
581
+ year={2025},
582
+ eprint={2508.17418},
583
+ archivePrefix={arXiv},
584
+ primaryClass={physics.chem-ph},
585
+ url={https://arxiv.org/abs/2508.17418},
586
+ }