# ToyTrack
[![Documentation Status](https://readthedocs.org/projects/toytrack/badge/?version=latest)](https://toytrack.readthedocs.io/en/latest/?badge=latest)
ToyTrack is a Python library for generating toy tracking events for particle physics.
**The goal**: To produce a "*good-enough*" event simulation, in as few lines as possible (currently 3 lines), as quickly as possible (currently 0.07 seconds for a 10,000-particle event).
## Installation
Use the package manager [pip](https://pip.pypa.io/en/stable/) to install ToyTrack.
```bash
pip install toytrack
```
Optionally, there are Pytorch dataloaders available for convenience. These require the `torch` package.
```bash
pip install toytrack[torch]
```
# Usage
## Vanilla Event
```python
from toytrack import ParticleGun, Detector, EventGenerator
# Initialize a particle gun which samples uniformly from pt between 2 and 20 GeV,
# initial direction phi between -pi and pi, and creation vertex vx and vy between -0.1 and 0.1 cm
# and which fires a normally-distributed number of particles with mean 20 and standard deviation 5
particle_gun = ParticleGun(num_particles=[20, 5, 'normal'], pt=[2, 20], pphi=[-np.pi, np.pi], vx=[-0.1, 0.1], vy=[-0.1, 0.1])
# Initialize a detector, a barrel-like detector with inner radius of 0.5 cm, and outer radius of 3 cm,
# with 10 layers
detector = Detector(dimension=2).add_from_template('barrel', min_radius=0.5, max_radius=3, number_of_layers=10)
# Initialize an event generator and generate an event
event = EventGenerator(particle_gun, detector).generate_event()
# Access the particles, hits and tracks as needed
particles = event.particles
hits = event.hits
tracks = event.tracks
# Plot the event
event.display()
```
![Example Event](https://raw.githubusercontent.com/murnanedaniel/ToyTrack/main/docs/imgs/example_event_vanilla.png)
## Event with Track Holes
```python
... # as above
# Initialize a detector that randomly drops 10% of hits. If an int N is given, then exactly
# N hits per track are dropped
detector = Detector(dimension=2, hole_inefficiency=0.1).add_from_template('barrel', min_radius=0.5, max_radius=3, number_of_layers=10)
... # as above
# Plot the event
event.display()
```
![Example Event with Holes](https://raw.githubusercontent.com/murnanedaniel/ToyTrack/main/docs/imgs/example_event_holes.png)
## Event with Noise Hits
```python
... # as above
# Generate event with between 30% and 70% noise hits per real hits. E.g. If the event has 100
# real hits, then between 30 and 70 noise hits will be generated. If an int N is given, then
# the absolute value of N noise hits are generated
event = EventGenerator(particle_gun, detector, noise=[0.3, 0.7]).generate_event()
... # as above
# Plot the event
event.display()
```
![Example Event with Noise](https://raw.githubusercontent.com/murnanedaniel/ToyTrack/main/docs/imgs/example_event_noise.png)
## Event with Multiple Particle Guns
```python
# Initialize one particle gun which samples uniformly from pt between 2 and 3 GeV,
# initial direction phi between -pi and pi, and creation vertex vx and vy between -0.1 and 0.1 cm
particle_gun_1 = ParticleGun(num_particles=[20, 5, 'normal'], pt=[2, 3], pphi=[-np.pi, np.pi], vx=[-0.1, 0.1], vy=[-0.1, 0.1])
# Initialize another particle gun which samples uniformly from pt between 100 and 200 GeV,
# initial direction phi between -pi and pi, and creation vertex vx and vy between -0.1 and 0.1 cm
particle_gun_2 = ParticleGun(num_particles=1, pt=[100, 200], pphi=[-np.pi, np.pi], vx=[-0.1, 0.1], vy=[-0.1, 0.1])
... # as above
# Initialize an event generator with a list of particle guns
event = EventGenerator([particle_gun_1, particle_gun_2], detector).generate_event()
... # as above
# Plot the event
event.display()
```
![Example Event with Multiple Particle Guns](https://raw.githubusercontent.com/murnanedaniel/ToyTrack/main/docs/imgs/example_event_multigun.png)
## Detector with Layer Safety Guarantee
We can ensure that each particle has exactly one hit per layer by setting the `layer_safety_guarantee` flag to `True`.
Consider an event with low pT particles:
```python
particle_gun = ParticleGun(num_particles=[20, 5, 'normal'], pt=[1, 3], pphi=[-np.pi, np.pi], vx=[-0.1, 0.1], vy=[-0.1, 0.1])
detector = Detector(dimension=2, layer_safety_guarantee=False).add_from_template('barrel', min_radius=0.5, max_radius=3, number_of_layers=10)
event = EventGenerator(particle_gun, detector).generate_event()
event.display()
```
![Example Event with Low pT Particles](https://raw.githubusercontent.com/murnanedaniel/ToyTrack/main/docs/imgs/example_event_lowpt.png)
```python
... # as above
# Initialize a detector WITH layer safety guarantee
detector = Detector(dimension=2, layer_safety_guarantee=True).add_from_template('barrel', min_radius=0.5, max_radius=3, number_of_layers=10)
... # as above
event.display()
```
![Example Event with Layer Safety Guarantee](https://raw.githubusercontent.com/murnanedaniel/ToyTrack/main/docs/imgs/example_event_lowpt_layer_safety.png)
Observe that particles with such low pT that they "curled" before the final layer are now removed.
# Performance
ToyTrack is designed to be fast. The following benchmarks were performed on a 64-core AMD EPYC 7763 (Milan) CPU.
![Scaling Study](https://raw.githubusercontent.com/murnanedaniel/ToyTrack/main/docs/imgs/time_scaling.png)
# Data Loading
## Pytorch Dataset
The `TracksDataset` class is a Pytorch dataset which can be used with a Pytorch dataloader. It can return either hitwise structure or trackwise structure.
```python
from toytrack.dataloaders import TracksDataset
config = {
"detector": {
"dimension": 2,
"hole_inefficiency": 0,
"min_radius": 0.5,
"max_radius": 3.,
"number_of_layers": 10,
},
"particle_guns": [
{
"num_particles": [20, 5, 'normal'],
"pt": [2, 20],
"pphi": [-3.14159, 3.14159],
"vx": [-0.1, 0.1],
"vy": [-0.1, 0.1]
}
],
"structure": "hitwise"
}
# initialize dataloader
dataset = TracksDataset(config)
# iterate over dataset
for sample in dataset:
x, mask, pids, event = sample["x"], sample["mask"], sample["pids"], sample["event"]
# x has shape (num_hits, num_features)
```
For trackwise structure:
```python
config = {
... # as above
"detector": {
... # as above
"layer_safety_guarantee": True # Must be True for trackwise structure
... # as above
}
"structure": "trackwise"
}
# initialize dataloader
dataset = TracksDataset(config)
# iterate over dataset
for sample in dataset:
x, mask, pids, event = sample["x"], sample["mask"], sample["pids"], sample["event"]
# x has shape (num_tracks, num_hits_per_track, num_features)
```
## Pytorch DataLoader
```python
from torch.utils.data import DataLoader
dataloader = DataLoader(dataset, batch_size=100, collate_fn=dataset.collate_fn)
# iterate over dataloader
for batch in dataloader:
x, mask, pids, event = batch["x"], batch["mask"], batch["pids"], batch["event"]
# Do something with the batch (which now has an extra batch dimension of size 100)
```
You can also get a straightforward parallelisation, with
```
dataloader = DataLoader(dataset, batch_size=100, collate_fn=dataset.collate_fn, num_workers=16)
```
## Transformations
You can apply transformations to the dataset by passing a list of transformations to the `TracksDataset` class. For example, to convert a trackwise dataset into patches of 3 hits each, you can do the following:
```python
from toytrack.transforms import TrackletPatchify
transform = TrackletPatchify(num_patches_per_track=3)
dataset = TracksDataset(config, transform=transform)
# iterate over dataset
for sample in dataset:
x, mask, pids, event = sample["x"], sample["mask"], sample["pids"], sample["event"]
# x has shape (num_tracks * num_patches_per_track, num_hits_per_patch, num_features)
```
![Example Event with Tracklet Patchify](https://raw.githubusercontent.com/murnanedaniel/ToyTrack/main/docs/imgs/example_event_tracklet_patches.png)
Raw data
{
"_id": null,
"home_page": "https://github.com/murnanedaniel/ToyTrack",
"name": "toytrack",
"maintainer": null,
"docs_url": null,
"requires_python": null,
"maintainer_email": null,
"keywords": null,
"author": "Daniel Murnane",
"author_email": "daniel.thomas.murnane@cern.ch",
"download_url": "https://files.pythonhosted.org/packages/03/52/372f247c114ab228d9c3900481c807e6e6b1d44f6a5a219b55e16948169e/toytrack-0.1.23.tar.gz",
"platform": null,
"description": "# ToyTrack\n\n[![Documentation Status](https://readthedocs.org/projects/toytrack/badge/?version=latest)](https://toytrack.readthedocs.io/en/latest/?badge=latest)\n\nToyTrack is a Python library for generating toy tracking events for particle physics. \n\n**The goal**: To produce a \"*good-enough*\" event simulation, in as few lines as possible (currently 3 lines), as quickly as possible (currently 0.07 seconds for a 10,000-particle event).\n\n## Installation\n\nUse the package manager [pip](https://pip.pypa.io/en/stable/) to install ToyTrack.\n\n```bash\npip install toytrack\n```\n\nOptionally, there are Pytorch dataloaders available for convenience. These require the `torch` package.\n\n```bash\npip install toytrack[torch]\n```\n\n# Usage\n\n## Vanilla Event\n\n```python\nfrom toytrack import ParticleGun, Detector, EventGenerator\n\n# Initialize a particle gun which samples uniformly from pt between 2 and 20 GeV, \n# initial direction phi between -pi and pi, and creation vertex vx and vy between -0.1 and 0.1 cm\n# and which fires a normally-distributed number of particles with mean 20 and standard deviation 5\nparticle_gun = ParticleGun(num_particles=[20, 5, 'normal'], pt=[2, 20], pphi=[-np.pi, np.pi], vx=[-0.1, 0.1], vy=[-0.1, 0.1])\n\n# Initialize a detector, a barrel-like detector with inner radius of 0.5 cm, and outer radius of 3 cm,\n# with 10 layers\ndetector = Detector(dimension=2).add_from_template('barrel', min_radius=0.5, max_radius=3, number_of_layers=10)\n\n# Initialize an event generator and generate an event\nevent = EventGenerator(particle_gun, detector).generate_event()\n\n# Access the particles, hits and tracks as needed\nparticles = event.particles\nhits = event.hits\ntracks = event.tracks\n\n# Plot the event\nevent.display()\n```\n\n![Example Event](https://raw.githubusercontent.com/murnanedaniel/ToyTrack/main/docs/imgs/example_event_vanilla.png)\n\n## Event with Track Holes\n\n```python\n... # as above\n\n# Initialize a detector that randomly drops 10% of hits. If an int N is given, then exactly\n# N hits per track are dropped\ndetector = Detector(dimension=2, hole_inefficiency=0.1).add_from_template('barrel', min_radius=0.5, max_radius=3, number_of_layers=10)\n\n... # as above\n\n# Plot the event\nevent.display()\n```\n\n![Example Event with Holes](https://raw.githubusercontent.com/murnanedaniel/ToyTrack/main/docs/imgs/example_event_holes.png)\n\n## Event with Noise Hits\n\n```python\n... # as above\n\n# Generate event with between 30% and 70% noise hits per real hits. E.g. If the event has 100 \n# real hits, then between 30 and 70 noise hits will be generated. If an int N is given, then\n# the absolute value of N noise hits are generated\nevent = EventGenerator(particle_gun, detector, noise=[0.3, 0.7]).generate_event()\n\n... # as above\n\n# Plot the event\nevent.display()\n```\n\n![Example Event with Noise](https://raw.githubusercontent.com/murnanedaniel/ToyTrack/main/docs/imgs/example_event_noise.png)\n\n## Event with Multiple Particle Guns\n\n```python\n# Initialize one particle gun which samples uniformly from pt between 2 and 3 GeV, \n# initial direction phi between -pi and pi, and creation vertex vx and vy between -0.1 and 0.1 cm\nparticle_gun_1 = ParticleGun(num_particles=[20, 5, 'normal'], pt=[2, 3], pphi=[-np.pi, np.pi], vx=[-0.1, 0.1], vy=[-0.1, 0.1])\n\n# Initialize another particle gun which samples uniformly from pt between 100 and 200 GeV, \n# initial direction phi between -pi and pi, and creation vertex vx and vy between -0.1 and 0.1 cm\nparticle_gun_2 = ParticleGun(num_particles=1, pt=[100, 200], pphi=[-np.pi, np.pi], vx=[-0.1, 0.1], vy=[-0.1, 0.1])\n\n... # as above\n\n# Initialize an event generator with a list of particle guns\nevent = EventGenerator([particle_gun_1, particle_gun_2], detector).generate_event()\n\n... # as above\n\n# Plot the event\nevent.display()\n```\n\n![Example Event with Multiple Particle Guns](https://raw.githubusercontent.com/murnanedaniel/ToyTrack/main/docs/imgs/example_event_multigun.png)\n\n## Detector with Layer Safety Guarantee\n\nWe can ensure that each particle has exactly one hit per layer by setting the `layer_safety_guarantee` flag to `True`.\n\nConsider an event with low pT particles:\n\n```python\nparticle_gun = ParticleGun(num_particles=[20, 5, 'normal'], pt=[1, 3], pphi=[-np.pi, np.pi], vx=[-0.1, 0.1], vy=[-0.1, 0.1])\n\ndetector = Detector(dimension=2, layer_safety_guarantee=False).add_from_template('barrel', min_radius=0.5, max_radius=3, number_of_layers=10)\n\nevent = EventGenerator(particle_gun, detector).generate_event()\n\nevent.display()\n```\n\n![Example Event with Low pT Particles](https://raw.githubusercontent.com/murnanedaniel/ToyTrack/main/docs/imgs/example_event_lowpt.png)\n\n```python\n... # as above\n\n# Initialize a detector WITH layer safety guarantee\ndetector = Detector(dimension=2, layer_safety_guarantee=True).add_from_template('barrel', min_radius=0.5, max_radius=3, number_of_layers=10)\n\n... # as above\n\nevent.display()\n```\n\n![Example Event with Layer Safety Guarantee](https://raw.githubusercontent.com/murnanedaniel/ToyTrack/main/docs/imgs/example_event_lowpt_layer_safety.png)\n\nObserve that particles with such low pT that they \"curled\" before the final layer are now removed.\n\n# Performance\n\nToyTrack is designed to be fast. The following benchmarks were performed on a 64-core AMD EPYC 7763 (Milan) CPU. \n\n![Scaling Study](https://raw.githubusercontent.com/murnanedaniel/ToyTrack/main/docs/imgs/time_scaling.png)\n\n# Data Loading\n\n## Pytorch Dataset\n\nThe `TracksDataset` class is a Pytorch dataset which can be used with a Pytorch dataloader. It can return either hitwise structure or trackwise structure.\n\n```python\nfrom toytrack.dataloaders import TracksDataset\n\nconfig = {\n \"detector\": {\n \"dimension\": 2,\n \"hole_inefficiency\": 0,\n \"min_radius\": 0.5,\n \"max_radius\": 3.,\n \"number_of_layers\": 10,\n },\n \"particle_guns\": [\n {\n \"num_particles\": [20, 5, 'normal'],\n \"pt\": [2, 20],\n \"pphi\": [-3.14159, 3.14159],\n \"vx\": [-0.1, 0.1],\n \"vy\": [-0.1, 0.1]\n }\n ],\n \"structure\": \"hitwise\"\n}\n\n# initialize dataloader\ndataset = TracksDataset(config)\n\n# iterate over dataset\nfor sample in dataset:\n x, mask, pids, event = sample[\"x\"], sample[\"mask\"], sample[\"pids\"], sample[\"event\"]\n # x has shape (num_hits, num_features)\n```\n\nFor trackwise structure:\n\n```python\nconfig = {\n ... # as above\n \"detector\": {\n ... # as above\n \"layer_safety_guarantee\": True # Must be True for trackwise structure\n ... # as above\n }\n \"structure\": \"trackwise\"\n}\n\n# initialize dataloader\ndataset = TracksDataset(config)\n\n# iterate over dataset\nfor sample in dataset:\n x, mask, pids, event = sample[\"x\"], sample[\"mask\"], sample[\"pids\"], sample[\"event\"]\n # x has shape (num_tracks, num_hits_per_track, num_features)\n```\n\n## Pytorch DataLoader\n\n```python\nfrom torch.utils.data import DataLoader\n\ndataloader = DataLoader(dataset, batch_size=100, collate_fn=dataset.collate_fn)\n\n# iterate over dataloader\nfor batch in dataloader:\n x, mask, pids, event = batch[\"x\"], batch[\"mask\"], batch[\"pids\"], batch[\"event\"]\n # Do something with the batch (which now has an extra batch dimension of size 100)\n```\n\nYou can also get a straightforward parallelisation, with \n```\ndataloader = DataLoader(dataset, batch_size=100, collate_fn=dataset.collate_fn, num_workers=16)\n```\n\n## Transformations\n\nYou can apply transformations to the dataset by passing a list of transformations to the `TracksDataset` class. For example, to convert a trackwise dataset into patches of 3 hits each, you can do the following:\n\n```python\nfrom toytrack.transforms import TrackletPatchify\n\ntransform = TrackletPatchify(num_patches_per_track=3)\ndataset = TracksDataset(config, transform=transform)\n\n# iterate over dataset\nfor sample in dataset:\n x, mask, pids, event = sample[\"x\"], sample[\"mask\"], sample[\"pids\"], sample[\"event\"]\n # x has shape (num_tracks * num_patches_per_track, num_hits_per_patch, num_features)\n```\n\n![Example Event with Tracklet Patchify](https://raw.githubusercontent.com/murnanedaniel/ToyTrack/main/docs/imgs/example_event_tracklet_patches.png)\n",
"bugtrack_url": null,
"license": null,
"summary": "A package for generating toy tracking data.",
"version": "0.1.23",
"project_urls": {
"Homepage": "https://github.com/murnanedaniel/ToyTrack"
},
"split_keywords": [],
"urls": [
{
"comment_text": "",
"digests": {
"blake2b_256": "4f6d777df521030a84ac586eced167a1c39bd6a7b7aa93b2980d4e0042a4249d",
"md5": "83edb683ffc0bd66bbff2d270f279581",
"sha256": "1c67e26f7610cebb19c7c390bbcff28aadded1ee1148244c89db1aa0c927bb59"
},
"downloads": -1,
"filename": "toytrack-0.1.23-py3-none-any.whl",
"has_sig": false,
"md5_digest": "83edb683ffc0bd66bbff2d270f279581",
"packagetype": "bdist_wheel",
"python_version": "py3",
"requires_python": null,
"size": 15791,
"upload_time": "2024-10-28T13:18:47",
"upload_time_iso_8601": "2024-10-28T13:18:47.034490Z",
"url": "https://files.pythonhosted.org/packages/4f/6d/777df521030a84ac586eced167a1c39bd6a7b7aa93b2980d4e0042a4249d/toytrack-0.1.23-py3-none-any.whl",
"yanked": false,
"yanked_reason": null
},
{
"comment_text": "",
"digests": {
"blake2b_256": "0352372f247c114ab228d9c3900481c807e6e6b1d44f6a5a219b55e16948169e",
"md5": "9a89552587196fcf73e925b122bb29ac",
"sha256": "84b270954ee0c64b0341c45403a78bedfc1700a08043f2a048ac0ac12fdca4af"
},
"downloads": -1,
"filename": "toytrack-0.1.23.tar.gz",
"has_sig": false,
"md5_digest": "9a89552587196fcf73e925b122bb29ac",
"packagetype": "sdist",
"python_version": "source",
"requires_python": null,
"size": 15785,
"upload_time": "2024-10-28T13:18:48",
"upload_time_iso_8601": "2024-10-28T13:18:48.081308Z",
"url": "https://files.pythonhosted.org/packages/03/52/372f247c114ab228d9c3900481c807e6e6b1d44f6a5a219b55e16948169e/toytrack-0.1.23.tar.gz",
"yanked": false,
"yanked_reason": null
}
],
"upload_time": "2024-10-28 13:18:48",
"github": true,
"gitlab": false,
"bitbucket": false,
"codeberg": false,
"github_user": "murnanedaniel",
"github_project": "ToyTrack",
"travis_ci": false,
"coveralls": false,
"github_actions": false,
"lcname": "toytrack"
}