Skip to content

Repository files navigation

PySlice

A GPU-accelerated Python package for simulating vibrational electron energy loss spectroscopy (EELS) using the TACAW method (Time Autocorrelation of Auxiliary Wavefunctions). PySlice integrates molecular dynamics with multislice electron scattering calculations to predict momentum- and energy-resolved phonon spectra directly from atomic trajectories.

Features

  • TACAW Analysis: Convert time-domain electron scattering into frequency-domain phonon spectra
  • Integrated MD: Run molecular dynamics with universal ML potentials (ORB, MACE, CHGNet)
  • GPU Acceleration: PyTorch backend with automatic CUDA/MPS/CPU selection
  • Flexible Input: Load structures from CIF, LAMMPS, XYZ, ASE trajectories, or ASE Atoms objects
  • STEM Imaging: HAADF/ADF/BF imaging and 4D-STEM diffraction

Installation

# Clone the repository
git clone https://github.com/h-walk/PySlice.git
cd PySlice
# Install with pip. -e = editable mode. [fast] will install torch (technically optional, but provides extreme speed improvements).
pip install -e ".[fast]"# Install OVITO for trajectory loading
pip install ovito --find-links https://www.ovito.org/pip/
# Or using uv (recommended)
uv sync
# Optional ORB/MD support currently requires Python 3.12 because ORB's dm-tree# dependency does not publish Python 3.13 wheels.
uv python install 3.12
uv sync --python 3.12 --extra fast --extra md

Quick Start

Full TACAW Pipeline (MD → Multislice → Phonon Dispersion)

fromase.buildimportbulkfrompysliceimportORBMDCalculator,MultisliceCalculator,TACAWData# 1. Create structureatoms=bulk("Si", "diamond", a=5.431, cubic=True) * (10, 10, 2)
# 2. Run molecular dynamicsmd=ORBMDCalculator(model_name="orb-v3-direct-inf-omat")
md.setup(atoms, temperature=300, timestep=2.0, production_steps=500, save_interval=5)
trajectory=md.run()
# 3. Run multislice (parallel beam for TACAW)calc=MultisliceCalculator()
calc.setup(trajectory, aperture=0, voltage_eV=100e3, sampling=0.1, slice_thickness=0.5)
wf_data=calc.run()
# 4. Compute phonon spectrumtacaw=TACAWData(wf_data)
Z=tacaw.spectral_diffraction(15.0) # Diffraction at 15 THz

Load Existing Trajectory

frompyslice.io.loaderimportLoadertrajectory=Loader(
"hBN.lammpstrj",
timestep=0.005, # psatom_mapping={1: "B", 2: "N"}
).load()
# ASE trajectory or CIF/XYZ filetrajectory=Loader("silicon.cif").load()

HAADF-STEM Imaging

frompysliceimportLoader,MultisliceCalculator,HAADFDataimportnumpyasnp# Load your trajectorytrajectory=Loader(
"hBN.lammpstrj",
timestep=0.005, # psatom_mapping={1: "B", 2: "N"}
).load()
# Optional cropping in time and spacetrajectory=trajectory.get_random_timesteps(5).slice_positions([0,20],[0,20])
# Define probe scan gridxs=np.linspace(5,12,16) ; ys=np.linspace(5,12,16)
calc=MultisliceCalculator()
calc.setup(trajectory, aperture=30, voltage_eV=100e3, sampling=0.1, probe_xs=xs, probe_ys=ys)
wf_data=calc.run()
haadf=HAADFData(wf_data)
haadf.calculateADF(inner_mrad=60, outer_mrad=200)
haadf.plot()

TEM Diffraction

frompysliceimportLoader,MultisliceCalculatorimportnumpyasnp# Load your trajectorytrajectory=Loader(
"hBN.lammpstrj",
timestep=0.005, # psatom_mapping={1: "B", 2: "N"}
).load()
# Optional cropping in time and spacetrajectory=trajectory.get_random_timesteps(5).slice_positions([0,20],[0,20])
calc=MultisliceCalculator()
calc.setup(trajectory, aperture=0, voltage_eV=100e3, sampling=0.1)
wf_data=calc.run()
wf_data.plot(powerscaling=0.125) # Diffraction pattern

Data Flow

Input Sources Processing Analysis Output
─────────────────────────────────────────────────────────────────────────────
CIF / XYZ / LAMMPS ─┬─→ Loader ─┬─→ ORBMDCalculator ─┐
ASE Atoms / .traj ─┘ │ (or FAIRChem) │
│ ↓
└───────────→ Trajectory
│
↓
MultisliceCalculator
(Probe → Potential → Propagate)
│
↓
WFData ψ(k,t)
│
┌─────────────────────────┼─────────────────────────┐
↓ ↓ ↓
TACAWData HAADFData WFData
FFT(t)→ω ∫|ψ|²dΩ (direct)
│ │ │
↓ ↓ ↓
Phonon Dispersion STEM Image Diffraction
Spectral Diffraction ADF/HAADF/BF CBED/LACBED
Spectrum Image 4D-STEM

Main Classes

Loader

Load atomic structures and trajectories from various formats.

frompyslice.io.loaderimportLoader# Supported: CIF, XYZ, LAMMPS dump, ASE .traj, ASE Atoms objectstraj=Loader("file.cif").load()
traj=Loader("dump.lammpstrj", timestep=0.01, atom_mapping={1: "B", 2: "N"}).load()

ORBMDCalculator / FAIRChemMDCalculator

Run molecular dynamics with universal ML potentials.

frompyslice.mdimportORBMDCalculatormd=ORBMDCalculator(model_name="orb-v3-direct-inf-omat", device="cuda")
md.setup(
atoms,
temperature=300, # Ktimestep=2.0, # fsproduction_steps=1000,
save_interval=5,
)
trajectory=md.run()

Trajectory

Container for atomic dynamics data.

trajectory.positions# (n_frames, n_atoms, 3)trajectory.velocities# (n_frames, n_atoms, 3)trajectory.atom_types# Atomic numberstrajectory.box_matrix# (3, 3) simulation celltrajectory.timestep# Frame spacing in ps

MultisliceCalculator

Compute exit wavefunctions via multislice algorithm.

frompyslice.multislice.calculatorsimportMultisliceCalculatorcalc=MultisliceCalculator()
calc.setup(
trajectory,
aperture=0, # mrad (0 = parallel beam)voltage_eV=100e3, # Accelerating voltagesampling=0.1, # Å/pixelslice_thickness=0.5, # Åprobe_positions=None, # Optional (N,2) array for STEM
)
wf_data=calc.run()

TACAWData

Frequency-domain phonon analysis.

frompyslice.postprocessing.tacaw_dataimportTACAWDatatacaw=TACAWData(wf_data)
# Analysis methodstacaw.frequencies# Available frequencies (THz)tacaw.spectral_diffraction(freq_THz) # k-space intensity at frequencytacaw.dispersion(kx_path, ky_path) # Phonon dispersion along k-pathtacaw.spectrum_image(freq_THz) # Real-space map at frequency (STEM)

HAADFData

STEM imaging analysis.

frompyslice.postprocessing.haadf_dataimportHAADFDatahaadf=HAADFData(wf_data)
haadf.calculateADF(inner_mrad=60, outer_mrad=200)
haadf.plot()

Examples

See the tests/ directory for detailed examples:

  • 00_probe.py - Probe wavefunction visualization
  • 01_potentials.py - Atomic potential calculations
  • 04_haadf.py - HAADF-STEM imaging
  • 05_tacaw.py - TACAW phonon spectroscopy
  • 06_loaders.py - Loading various file formats
  • 15_molecular_dynamics.py - MD with ORB potentials

Requirements

Core:

  • Python 3.10+
  • NumPy, SciPy, Matplotlib
  • ASE (Atomic Simulation Environment)
  • OVITO

Recommended:

  • PyTorch (GPU acceleration)

License

MIT License - see LICENSE file for details.

About

PySlice is a Python package for simulating and analyzing multslice simulations from molecular dynamics trajectories. In addition to standard multislice simulations such as diffraction and HAADF image generation, it implements the TACAW method to convert time-domain electron scattering data into frequency-domain spectra.

Topics

Resources

Stars

15 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages