ttsim3d

Simulate 3D electrostatic potential maps from a PDB file in PyTorch. This package currently replicates theory laid out in Himes & Grigorieff (2021).

For a full list of changes, see the CHANGELOG.

Installation

ttsim3d is available on PyPi and can be installed via

pip install ttsim3d

From source

To create a source installation, first download/clone the repository, then run the install command

git clone https://github.com/teamtomo/ttsim3d.git
cd ttsim3d
pip install -e .

Optional development and testing dependencies can also be installed by running

pip install -e ".[dev,test]"

Running CLI program

Installation of the package creates the executable program ttsim3d-cli which takes in a PDB file along with other simulation options and outputs the simulated 3D scattering potential to a .mrc file. All options for the program can be printed by running:

ttsim3d-cli --help

Python objects

There are two user-facing classes in ttsim3d built upon Pydantic models for validating inputs and simulating a volume. The first class, ttsim3d.models.Simulator, holds reference to a PDB file and basic simulation parameters related to that structure. The second class, ttsim3d.models.SimulatorConfig is used to configure more advanced options, such as dose weighting and simulation upsampling. An basic use of these objects to run a simulation looks like

from ttsim3d.models import Simulator, SimulatorConfig

# Instantiate the configuration object 
sim_conf = SimulatorConfig(
    voltage=300.0,  # in keV
    apply_dose_weighting=True,
    dose_start=0.0,  # in e-/A^2
    dose_end=35.0,   # in e-/A^2
    upsampling=-1,   # auto
)

# Instantiate the simulator
sim = Simulator(
    pdb_filepath="some/path/to/structure.pdb",
    pixel_spacing=1.25,  # Angstroms
    volume_shape=(256, 256, 256),
    b_factor_scaling=1.0,
    additional_b_factor=15.0,  # Add to all atoms
    simulator_config=sim_conf,
)

# Run the simulation
volume = sim.run()
print(type(volume))  # torch.Tensor
print(volume.shape)  # (256, 256, 256)

# OR export the simulation to a mrc file
mrc_filepath = "some/path/to/simulated_structure/mrc"
sim.export_to_mrc(mrc_filepath)

Running on GPU or CPU

The ttsim3d package supports GPU accelerated simulations with PyTorch. Use the device argument to specify which device to run the simulation on.

# ...
# Assume same objects as above

# Run on the CPU
volume_cpu = sim.run(device="cpu")

# Run on the GPU (assumes CUDA is available)
# Both run on the zeroth GPU device
volume_gpu = sim.run(device="0")
volume_gpu = sim.run(device="cuda:0")

# Same argument can be passed to the `export_to_mrc` method
sim.export_to_mrc(mrc_filepath, device="cuda:0")

Working with configuration files

Simulation configurations can be saved to disk as either a YAML or JSON file by using the to_yaml or to_json methods of the Simulator class, respectively. Assuming the same sim object defined as above, you can export the configuration to a YAML file like this:

sim.to_yaml("some/path/to/simulation_config.yaml")

The contents of the YAML file will look something like this:

additional_b_factor: 15.0
b_factor_scaling: 1.0
center_atoms: true
pdb_filepath: some/path/to/structure.pdb
pixel_spacing: 1.25
remove_hydrogens: true
volume_shape:
- 256
- 256
- 256
simulator_config:
  apply_dose_weighting: true
  apply_dqe: true
  atom_batch_size: 16384
  crit_exposure_bfactor: -1
  dose_end: 35.0
  dose_filter_modify_signal: None
  dose_start: 0.0
  mtf_reference: k2_300kv
  store_volume: true
  upsampling: -1
  voltage: 300.0

Similarly, you can load in a configuration from a YAML file by using the from_yaml class method of the Simulator class:

from ttsim3d.models import Simulator

sim = Simulator.from_yaml("some/path/to/simulation_config.yaml")

Included MTF reference files

Users can specify MTF reference by a path to a star file under the SimulatorConfig.mtf_reference parameter. Applying a DQE filter during simulation can be turned on or off by the SimulatorConfig.apply_dqe parameter.

The ttsim3d package includes several common MTF reference files which are accessible by a string rather than a path. To see the list of included MTF reference files, run the following code snippet:

from ttsim3d import models

print(models.included_mtf_references())

The default MTF reference is "k2_300kV", which corresponds to the MTF of a Gatan K2 camera at 300 kV.