def test_can_resize_with_particle_props():
ensemble = coldatoms.Ensemble(12)
ensemble.set_particle_property('mass', np.ones(12))
new_size = 15
ensemble.resize(new_size)
assert (ensemble.num_ptcls == new_size)
assert (ensemble.particle_properties['mass'].shape[0] == new_size)
def setup(self):
self.coulomb_force = coldatoms.CoulombForce()
self.ensemble = coldatoms.Ensemble(num_ptcls=3)
self.ensemble.x[1] = np.array([1.0, 1.0, 1.0])
self.ensemble.x[2] = np.array([2.0, 2.0, 2.0])
self.ensemble.ensemble_properties['charge'] = 1.0
self.f = np.zeros_like(self.ensemble.v)
def test_zero_velocities():
ensemble = coldatoms.Ensemble(num_ptcls=3)
ensemble.x = np.random.rand(3)
ensemble.v = np.zeros(3)
x_old = ensemble.x.copy()
coldatoms.drift_kick(1.0, ensemble)
assert (ensemble.x == x_old).all()
def create_ensemble(uE, omega_z, mass, charge):
num_ions = int(uE.size / 2)
x = uE[:num_ions]
y = uE[num_ions:]
r = np.sqrt(x**2 + y**2)
r_hat = np.transpose(np.array([x / r, y / r]))
phi_hat = np.transpose(np.array([-y / r, x / r]))
v = np.zeros([num_ions, 2], dtype=np.float64)
for i in range(num_ions):
v[i, 0] = omega_z * r[i] * phi_hat[i, 0]
v[i, 1] = omega_z * r[i] * phi_hat[i, 1]
ensemble = coldatoms.Ensemble(num_ions)
for i in range(num_ions):
ensemble.x[i, 0] = x[i]
ensemble.x[i, 1] = y[i]
ensemble.x[i, 2] = 0.0
ensemble.v[i, 0] = v[i, 0]
ensemble.v[i, 1] = v[i, 1]
ensemble.v[i, 2] = 0.0
ensemble.ensemble_properties['mass'] = mass
ensemble.ensemble_properties['charge'] = charge
return ensemble
def test_serialize_ensemble_rand_phase_space_with_ensemble_props():
num_ptcls = 2
ensemble = coldatoms.Ensemble(num_ptcls)
ensemble.x = np.random.random_sample(tuple(ensemble.x.shape))
ensemble.v = np.random.random_sample(tuple(ensemble.v.shape))
ensemble.ensemble_properties['mass'] = 100.0
check_json_roundtrip(ensemble)
def test_zero_accelerations():
ensemble = coldatoms.Ensemble(num_ptcls=3)
ensemble.x = np.random.rand(3)
ensemble.v = np.zeros(3)
v_old = ensemble.v.copy()
coldatoms.drift_kick(1.0, ensemble)
assert (ensemble.v == v_old).all()
def test_force_is_non_zero():
fluorescence = coldatoms.RadiationPressure(1.0e8,
hbar_k, intensity, detuning)
ensemble = coldatoms.Ensemble()
# In one millisecond we expect to scatter more than one photon
f = np.zeros_like(ensemble.v)
fluorescence.force(1.0e-3, ensemble, f)
assert(np.linalg.norm(f) > np.linalg.norm(hbar_k))
def test_can_use_multiple_sources():
ensemble = coldatoms.Ensemble(num_ptcls=17)
s = TrivialSource()
coldatoms.produce_ptcls(3.0, ensemble, [s, s])
assert (ensemble.num_ptcls == 17 + 2 * s.num_ptcls_produced(1.0))
for i in range(17, ensemble.num_ptcls):
assert (abs(ensemble.x[i, 0] - i) < 1.0e-6)
def test_have_to_provide_mass():
ensemble = coldatoms.Ensemble(num_ptcls=3)
ensemble.x = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]])
ensemble.v = np.zeros_like(ensemble.x)
k = [1.0, 2.0, 3.0]
harmonic = Harmonic(k)
coldatoms.drift_kick(0.1, ensemble, [harmonic])
def setup(self):
self.kx = 2.0
self.ky = 3.0
self.kz = -1.0
self.coulomb = coldatoms.CoulombForce()
self.harmonic_trap = coldatoms.HarmonicTrapPotential(self.kx, self.ky, self.kz)
self.ensemble = coldatoms.Ensemble(num_ptcls=3)
self.ensemble.x[1] = np.array([1.0, 1.0, 1.0])
self.ensemble.x[2] = np.array([2.0, 2.0, 2.0])
self.ensemble.ensemble_properties['charge'] = 1.0
def test_moving_parallel_to_face_wont_be_absorbed():
sink = coldatoms.SinkPlane(np.array([0.0, 0.0, 0.0]),
np.array([0.0, 0.0, 1.0]))
ensemble = coldatoms.Ensemble(num_ptcls=1)
ensemble.x[0] = np.array([0.0, 0.0, -0.1])
ensemble.v[0] = np.array([0.0, 1.0e5, 0.0])
dt = 1.0
times = sink.find_absorption_time(ensemble.x, ensemble.v, dt)
assert (times[0] < 0 or times[0] > dt)
def test_serialize_ensemble_rand_phase_space_with_particle_props():
num_ptcls = 120
ensemble = coldatoms.Ensemble(num_ptcls)
ensemble.x = np.random.random_sample(tuple(ensemble.x.shape))
ensemble.v = np.random.random_sample(tuple(ensemble.v.shape))
ensemble.particle_properties['mass'] = (np.random.random_sample(
ensemble.x.shape[0]))
ensemble.particle_properties['density_matrix'] = (np.random.random_sample(
(ensemble.x.shape[0], 4, 4)))
check_json_roundtrip(ensemble)
def test_copy_ensemble_is_deep():
num_ptcls = 12
ensemble = coldatoms.Ensemble(num_ptcls)
ensemble.x = np.random.random_sample(tuple(ensemble.x.shape))
orig_x = np.copy(ensemble.x)
cpy = ensemble.copy()
new_x = np.random.random_sample(tuple(ensemble.x.shape))
cpy.x = new_x
assert (arrays_close(ensemble.x, orig_x))
assert (arrays_close(cpy.x, new_x))
def test_serialize_ensemble():
num_ptcls = 12
ensemble = coldatoms.Ensemble(num_ptcls)
ensemble.x = np.random.random_sample(tuple(ensemble.x.shape))
ensemble.v = np.random.random_sample(tuple(ensemble.v.shape))
ensemble.particle_properties['mass'] = (np.random.random_sample(
ensemble.x.shape[0]))
ensemble.particle_properties['density_matrix'] = (np.random.random_sample(
(ensemble.x.shape[0], 4, 4)))
ensemble.ensemble_properties['charge'] = 1.7
check_json_roundtrip(ensemble)
def test_can_use_per_particle_masses():
ensemble = coldatoms.Ensemble(num_ptcls=3)
ensemble.x = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]])
ensemble.v = np.zeros_like(ensemble.x)
ensemble.set_particle_property('mass', np.ones(3))
dt = 0.02
k = [1.0, 2.0, 3.0]
harmonic = Harmonic(k)
coldatoms.drift_kick(dt, ensemble, [harmonic])
def test_particle_hittling_plane_will_be_absorbed():
sink = coldatoms.SinkPlane(np.array([0.0, 0.0, 0.0]),
np.array([0.0, 0.0, 1.0]))
ensemble = coldatoms.Ensemble(num_ptcls=1)
ensemble.x[0] = np.array([0.0, 0.0, -0.1])
ensemble.v[0] = np.array([0.0, 0.0, 1.0])
dt = 1.0
times = sink.find_absorption_time(ensemble.x, ensemble.v, dt)
assert (times[0] >= 0)
assert (times[0] <= dt)
def test_force_is_not_unreasonably_large():
fluorescence = coldatoms.RadiationPressure(1.0e8,
hbar_k, intensity, detuning)
ensemble = coldatoms.Ensemble()
# In one millisecond we expect to scatter no more than s * (gamma/2pi)* dt
# photons.
expected_number_of_recoils = (intensity.intensity *
(1.0e8 / 2.0 / np.pi) * 1.0e-3)
f = np.zeros_like(ensemble.v)
fluorescence.force(1.0e-3, ensemble, f)
assert(np.linalg.norm(f) <
3.0 * expected_number_of_recoils * np.linalg.norm(hbar_k))
def test_can_copy_ensemble():
num_ptcls = 12
ensemble = coldatoms.Ensemble(num_ptcls)
ensemble.x = np.random.random_sample(tuple(ensemble.x.shape))
ensemble.v = np.random.random_sample(tuple(ensemble.v.shape))
ensemble.particle_properties['blah'] = np.random.random_sample(
tuple(ensemble.v.shape))
cpy = ensemble.copy()
assert (arrays_close(ensemble.x, cpy.x))
assert (arrays_close(ensemble.v, cpy.v))
assert (arrays_close(ensemble.particle_properties['blah'],
cpy.particle_properties['blah']))
def test_recoil_force_is_consistent_with_random_walk():
fluorescence = coldatoms.RadiationPressure(1.0e8,
hbar_k, intensity, detuning)
ensemble = coldatoms.Ensemble()
# In one millisecond we expect to scatter no more than s * (gamma/2pi)* dt
# photons.
expected_number_of_recoils = (intensity.intensity *
(1.0e8 / 2.0 / np.pi) * 1.0e-3)
f = np.zeros_like(ensemble.v)
fluorescence.force(1.0e-3, ensemble, f)
assert(np.abs(f[0, 1]) <
3.0 * np.sqrt(expected_number_of_recoils) * np.linalg.norm(hbar_k))
def setup(self):
self.kx = 2.0
self.ky = 3.0
self.kz = -1.0
self.harmonic_trap_force = coldatoms.HarmonicTrapPotential(
self.kx, self.ky, self.kz)
self.num_ptcls = 5
self.ensemble = coldatoms.Ensemble(num_ptcls=self.num_ptcls)
self.ensemble.x[:, 0] = np.random.random(self.num_ptcls)
self.ensemble.x[:, 1] = np.random.random(self.num_ptcls)
self.ensemble.x[:, 2] = np.random.random(self.num_ptcls)
self.ensemble.ensemble_properties['charge'] = 1.0
self.f = np.zeros_like(self.ensemble.v)
def test_absorbed_particles_are_removed():
sink = coldatoms.SinkPlane(np.array([0.0, 0.0, 0.0]),
np.array([0.0, 0.0, 1.0]))
num_ptcls = 2
ensemble = coldatoms.Ensemble(num_ptcls=num_ptcls)
ensemble.x[0] = np.array([0.0, 0.0, -0.1])
ensemble.v[0] = np.array([0.0, 0.0, 1.0])
ensemble.x[1] = np.array([0.0, 0.0, -0.1])
ensemble.v[1] = np.array([0.0, 0.0, 1.0e-2])
dt = 1.0
coldatoms.process_sink(dt, ensemble, sink)
assert (ensemble.num_ptcls < num_ptcls)
def test_harmonic_potential_motion_is_bounded():
ensemble = coldatoms.Ensemble(num_ptcls=5)
ensemble.x = np.random.rand(5, 3)
ensemble.v = np.random.rand(5, 3)
initial_stddev_vel = np.linalg.norm(ensemble.v)
m = 3.4
ensemble.ensemble_properties['mass'] = m
harmonic = Harmonic([1.0, 2.0, 3.0])
for i in range(100):
coldatoms.drift_kick(1.0, ensemble, [harmonic])
stddev_vel = np.linalg.norm(ensemble.v)
assert (stddev_vel < 10.0 * initial_stddev_vel)
def test_recoil_force_works_for_multiple_atoms():
fluorescence = coldatoms.RadiationPressure(1.0e8,
hbar_k, intensity, detuning)
ensemble = coldatoms.Ensemble(10)
# In one millisecond we expect to scatter no more than s * (gamma/2pi)* dt
# photons.
expected_number_of_recoils = (intensity.intensity *
(1.0e8 / 2.0 / np.pi) * 1.0e-3)
f = np.zeros_like(ensemble.v)
fluorescence.force(1.0e-3, ensemble, f)
assert((np.linalg.norm(f, axis=1) > np.linalg.norm(hbar_k)).all())
assert((np.linalg.norm(f, axis=1) <
(3.0 * expected_number_of_recoils *
np.linalg.norm(hbar_k))).all())
assert((np.abs(f[:,1]) <
(3.0 * np.sqrt(expected_number_of_recoils) *
np.linalg.norm(hbar_k))).all())
def test_spot_check(self):
self.ensemble = coldatoms.Ensemble(2)
self.ensemble.x[0, :] = 0.0
self.ensemble.x[1, 0] = 1.3
self.ensemble.x[1, 1] = 1.5
self.ensemble.x[1, 2] = 1.7
q = 1.3
self.ensemble.ensemble_properties['charge'] = q
dt = 1.0e-2
r = self.ensemble.x[0, :] - self.ensemble.x[1, :]
dist = np.linalg.norm(r)
ke = 1.0 / (4.0 * np.pi * 8.854187817620e-12)
f_expected = ke * r * q * q / (dist**3)
f_expected *= dt
self.coulomb_force.force(dt, self.ensemble, self.f)
normalization = np.sqrt(
np.linalg.norm(self.f[0])**2 + np.linalg.norm(f_expected)**2)
assert(np.linalg.norm(self.f[0] - f_expected) / normalization < 1.0e-9)
def test_per_particle_charge_spot_check(self):
self.ensemble = coldatoms.Ensemble(2)
self.ensemble.x[0, :] = 0.0
self.ensemble.x[1, 0] = 1.3
self.ensemble.x[1, 1] = 1.5
self.ensemble.x[1, 2] = 1.7
q = 1.3 * np.array([1.6e-19, 2.0 * 1.6e-19])
self.ensemble.ensemble_properties = {}
self.ensemble.set_particle_property('charge', q)
dt = 1.0e-2
r = self.ensemble.x[0, :] - self.ensemble.x[1, :]
dist = np.linalg.norm(r)
ke = 1.0 / (4.0 * np.pi * 8.854187817620e-12)
f_expected = ke * r * q[0] * q[1] / (dist**3)
f_expected *= dt
self.coulomb_force.force(dt, self.ensemble, self.f)
normalization = np.sqrt(
np.linalg.norm(self.f[0])**2 + np.linalg.norm(f_expected)**2)
assert(np.linalg.norm(self.f[0] - f_expected) / normalization < 1.0e-9)
def measure_time(num_ptcls,
per_ptcl_charges=False,
num_iter=1,
use_reference_impl=False):
ensemble = coldatoms.Ensemble(num_ptcls=num_ptcls)
ensemble.x = np.random.random([num_ptcls, 3])
if per_ptcl_charges:
ensemble.set_particle_properties('charge', np.random.random(num_ptcls))
else:
ensemble.ensemble_properties['charge'] = 1.0
f = coldatoms.CoulombForce()
if use_reference_impl:
f.use_reference_implementations()
accumulated_force = np.zeros_like(ensemble.v)
t0 = time.time()
for i in range(num_iter):
f.force(1.0e-1, ensemble, accumulated_force)
t1 = time.time()
return t1 - t0
def test_harmonic_spot():
ensemble = coldatoms.Ensemble(num_ptcls=3)
ensemble.x = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]])
ensemble.v = np.zeros_like(ensemble.x)
t = 0
dt = 0.02
m = 4.0
ensemble.ensemble_properties['mass'] = m
k = [1.0, 2.0, 3.0]
harmonic = Harmonic(k)
for i in range(50):
coldatoms.drift_kick(dt, ensemble, [harmonic])
t += dt
omega_x = math.sqrt(k[0] / m)
assert (abs(ensemble.x[0, 0] - math.cos(t * omega_x)) < dt * dt)
omega_y = math.sqrt(k[1] / m)
assert (abs(ensemble.x[1, 1] - math.cos(t * omega_y)) < dt * dt)
omega_z = math.sqrt(k[2] / m)
assert (abs(ensemble.x[2, 2] - math.cos(t * omega_z)) < dt * dt)
def setup(self):
self.ensemble = coldatoms.Ensemble(num_ptcls=5)
self.ensemble.v[0] = np.array([1.0, 0.0, 0.0])
self.ensemble.ensemble_properties['charge'] = 3.3
self.ensemble.ensemble_properties['mass'] = 2.7
def test_no_sources_produce_no_particles():
ensemble = coldatoms.Ensemble(num_ptcls=17)
coldatoms.produce_ptcls(1.0, ensemble, [])
assert (ensemble.num_ptcls == 17)
def test_positions_get_generated():
ensemble = coldatoms.Ensemble(num_ptcls=17)
s = TrivialSource()
coldatoms.produce_ptcls(1.0, ensemble, [s])
for i in range(17, ensemble.num_ptcls):
assert (abs(ensemble.x[i, 0] - i) < 1.0e-6)
