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)