bleykauf / aisim

Simulations for light-pulse atom interferometry
GNU General Public License v3.0
8 stars 5 forks source link

Cannot select a single atom from AtomicEnsemble with __getitem__ #49

Closed bleykauf closed 4 years ago

bleykauf commented 4 years ago

This fails:

pos_params = {
    'mean_x': 0.0,
    'std_x': 3.0e-3,  # cloud radius in m
    'mean_y': 0.0,
    'std_y': 3.0e-3,  # cloud radius in m
    'mean_z': 0.0,
    'std_z': 0,        # ignore z dimension, its not relevant here
}
vel_params = {
    'mean_vx': 0.0,
    # cloud velocity spread in m/s at tempearture of 3 uK
    'std_vx': ais.convert.vel_from_temp(3.0e-6),
    'mean_vy': 0.0,
    # cloud velocity spread in m/s at tempearture of 3 uK
    'std_vy': ais.convert.vel_from_temp(3.0e-6),
    'mean_vz': 0.0,
    'std_vz': ais.convert.vel_from_temp(0),
}

atoms = ais.create_random_ensemble_from_gaussian_distribution(
    pos_params, vel_params, int(1e4), state_kets=[1, 0, 0, 0])

atoms[0]

with

---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
<ipython-input-16-72aa0f35cda1> in <module>
----> 1 atoms[0]

c:\users\bastian\documents\programming\aisim\aisim\atoms.py in __getitem__(self, key)
     80             phase_space_vectors=self.phase_space_vectors[key][:],
     81             state_kets=self.state_kets[key][:],
---> 82             weights=self.weights[key])
     83         return new_instance
     84 

c:\users\bastian\documents\programming\aisim\aisim\atoms.py in __init__(self, phase_space_vectors, state_kets, time, weights)
     49     def __init__(self, phase_space_vectors, state_kets=[1, 0], time=0,
     50                  weights=None):
---> 51         assert phase_space_vectors.shape[1] == 6
     52         self.phase_space_vectors = phase_space_vectors
     53         self.state_kets = state_kets

IndexError: tuple index out of range

but

atoms[0:1]

works.