def test_spot_check_per_particle(self):
Bz = 1.7
q = self.ensemble.ensemble_properties['charge']
m = self.ensemble.ensemble_properties['mass']
dt = 1.2e-7
self.ensemble.particle_properties['charge'] = q * np.ones(
self.ensemble.v.shape[0])
self.ensemble.particle_properties['mass'] = m * np.ones(
self.ensemble.v.shape[0])
self.ensemble.ensemble_properties = {}
initial_positions = np.random.random(self.ensemble.x.shape)
initial_velocities = np.random.random(self.ensemble.v.shape)
self.ensemble.x = np.copy(initial_positions)
self.ensemble.v = np.copy(initial_velocities)
coldatoms.bend_kick(dt, Bz, self.ensemble, [])
final_positions = np.copy(self.ensemble.x)
final_velocities = np.copy(self.ensemble.v)
# now compute the reference solution
self.ensemble.x = np.copy(initial_positions)
self.ensemble.v = np.copy(initial_velocities)
coldatoms.bend_kick(dt, Bz, self.ensemble, [], reference_impl=True)
final_positions_ref = np.copy(self.ensemble.x)
final_velocities_ref = np.copy(self.ensemble.v)
assert (np.linalg.norm(final_positions - final_positions_ref) < 1.0e-6)
assert (np.linalg.norm(final_velocities - final_velocities_ref) <
1.0e-6)
def test_multiple_steps_same_as_individual_steps(self):
Bz = 1.7
q = self.ensemble.ensemble_properties['charge']
m = self.ensemble.ensemble_properties['mass']
dt = 1.2e-7
num_steps = 4
initial_positions = np.random.random(self.ensemble.x.shape)
initial_velocities = np.random.random(self.ensemble.v.shape)
self.ensemble.x = np.copy(initial_positions)
self.ensemble.v = np.copy(initial_velocities)
coldatoms.bend_kick(dt, Bz, self.ensemble, [], num_steps=num_steps)
final_positions = np.copy(self.ensemble.x)
final_velocities = np.copy(self.ensemble.v)
# now compute the reference solution
self.ensemble.x = np.copy(initial_positions)
self.ensemble.v = np.copy(initial_velocities)
for i in range(num_steps):
coldatoms.bend_kick(dt, Bz, self.ensemble, [])
final_positions_ref = np.copy(self.ensemble.x)
final_velocities_ref = np.copy(self.ensemble.v)
assert (np.linalg.norm(final_positions - final_positions_ref) < 1.0e-6)
assert (np.linalg.norm(final_velocities - final_velocities_ref) <
1.0e-6)
def evolve_ensemble_with_damping(dt, t_max, ensemble, Bz, forces, dampings):
num_steps = int(np.floor(t_max / dt))
for i in range(num_steps):
coldatoms.bend_kick(dt, Bz, ensemble, forces, num_steps=1)
for d in dampings:
d.dampen(dt, ensemble)
fractional_dt = t_max - (num_steps * dt)
if (fractional_dt / dt > 1.0e-6):
coldatoms.bend_kick(fractional_dt, Bz, ensemble, forces, num_steps=1)
for d in dampings:
d.dampen(fractional_dt, ensemble)
def test_near_zero_field(self):
Bz = 1.0e-30
dt = 1.0
expected_position = self.ensemble.x[0] + dt * self.ensemble.v[0]
expected_velocity = self.ensemble.v[0]
coldatoms.bend_kick(dt, Bz, self.ensemble, [])
assert (np.linalg.norm(self.ensemble.x[0] - expected_position) <
1.0e-8)
assert (np.linalg.norm(self.ensemble.v[0] - expected_velocity) <
1.0e-8)
def test_force_gets_called(self):
Bz = 1.0e-30
dt = 1.0
class MockForce(object):
def force(self, dt, ensemble, f):
self.position = np.copy(ensemble.x)
expected_position = self.ensemble.x[0] + 0.5 * dt * self.ensemble.v[0]
force = MockForce()
coldatoms.bend_kick(dt, Bz, self.ensemble, [force])
assert (np.linalg.norm(force.position[0] - expected_position) < 1.0e-8)
def test_closed_circle(self):
Bz = 1.0
q = self.ensemble.ensemble_properties['charge']
m = self.ensemble.ensemble_properties['mass']
omegaB = Bz * q / m
dt = 2.0 * np.pi / omegaB
expected_position = self.ensemble.x[0]
expected_velocity = self.ensemble.v[0]
coldatoms.bend_kick(dt, Bz, self.ensemble, [])
assert (np.linalg.norm(self.ensemble.x[0] - expected_position) <
1.0e-8)
assert (np.linalg.norm(self.ensemble.v[0] - expected_velocity) <
1.0e-8)
def test_near_zero_field_ensemble_omegaB(self):
Bz = 1.0e-30
dt = 1.0
self.ensemble.particle_properties['charge'] = np.ones(
self.ensemble.v.shape[0])
self.ensemble.particle_properties['mass'] = np.ones(
self.ensemble.v.shape[0])
self.ensemble.ensemble_properties = {}
expected_position = self.ensemble.x[0] + dt * self.ensemble.v[0]
expected_velocity = self.ensemble.v[0]
coldatoms.bend_kick(dt, Bz, self.ensemble, [])
assert (np.linalg.norm(self.ensemble.x[0] - expected_position) <
1.0e-8)
assert (np.linalg.norm(self.ensemble.v[0] - expected_velocity) <
1.0e-8)
def test_must_provide_charge(self):
self.ensemble.ensemble_properties.pop('charge', None)
coldatoms.bend_kick(1.0, 1.0, self.ensemble, [])
def evolve_ensemble(dt, t_max, ensemble, Bz, forces):
num_steps = int(t_max / dt)
coldatoms.bend_kick(dt, Bz, ensemble, forces, num_steps=num_steps)
coldatoms.bend_kick(t_max - dt * num_steps, Bz, ensemble, forces)
def generate_seededPSD(storage_directory,
pMode,
sMode,
totalEnergy,
percentPrimaryE,
t_max=10.0e-3,
dt=1.0e-9,
num_dump=50):
"""
Params:
* storage_directory: directory to which to write results
* pMode: Mode to provide majority of energy
* sMode: Mode to seed with minimal energy
* totalEnergy: total energy of crystal
* percentPrimaryE: (Percentage of energy in the primary mode
* max_time: duration of evolution
* t_step: time step for each evolution
Returns:
Nothing. Instead, saves numpy arrays 'x_freq' and 'PSD_data', obtained from performing FFT^2 on the z position
time series.
MAKE SURE YOU HAVE CHOSEN AN EMPTY DIRECTORY (OR DONT CARE ABOUT OVERWRITING)
"""
os.makedirs(
storage_directory) #Should throw OS error if directory already exists
num_steps = int(np.ceil(t_max / dt))
modal_positions = [
] #instantiate array to hold z-positions as crystal evolves
# Create new ensemble from mode_analysis.
modal_ensemble = create_ensemble(mode_analysis.uE, mode_analysis.wrot,
mode_analysis.m_Be, mode_analysis.q)
x = mode_analysis.u[:mode_analysis.Nion]
y = mode_analysis.u[mode_analysis.Nion:]
dx = x.reshape((x.size, 1)) - x
dy = y.reshape((y.size, 1)) - y
rsep = np.sqrt(dx**2 + dy**2)
with np.errstate(divide='ignore'):
rsep3 = np.where(rsep != 0., rsep**(-3), 0)
K = np.diag((-1 + 0.5 * np.sum(rsep3, axis=0)))
K -= 0.5 * rsep3
K = np.mat(K)
Mmat = np.diag(mode_analysis.md)
Mmat = np.mat(Mmat)
# Based on desired mode number, establish ICs with z-components based on eigenmode
eigVect1a = mode_analysis.axialEvects[:, -2 * pMode]
eigVect1b = mode_analysis.axialEvects[:, -2 * pMode + 1]
eigVect2a = mode_analysis.axialEvects[:, -2 * sMode]
eigVect2b = mode_analysis.axialEvects[:, -2 * sMode + 1]
pCombo = eigVect1a + eigVect1b # Choose this combo to get a 'pure' position state
sCombo = eigVect2a + eigVect2b # Choose this combo to get a 'pure' position state
pri_pos = pCombo[:mode_analysis.Nion]
pri_vel = pCombo[mode_analysis.Nion:]
seed_pos = sCombo[:mode_analysis.Nion]
seed_vel = sCombo[mode_analysis.Nion:]
pri_x = np.mat(pri_pos.reshape((mode_analysis.Nion, 1)))
pri_v = np.mat(pri_vel.reshape((mode_analysis.Nion, 1)))
seed_x = np.mat(seed_pos.reshape((mode_analysis.Nion, 1)))
seed_v = np.mat(seed_vel.reshape((mode_analysis.Nion, 1)))
pri_energy = 0.5 * pri_v.H * Mmat * pri_v - 0.5 * pri_x.H * K * pri_x
seed_energy = 0.5 * seed_v.H * Mmat * seed_v - 0.5 * seed_x.H * K * seed_x
pri_E = totalEnergy * percentPrimaryE
seed_E = totalEnergy * (1 - percentPrimaryE)
pri_pos = np.sqrt(pri_E / pri_energy) * pri_pos
pri_vel = np.sqrt(pri_E / pri_energy) * pri_vel
seed_pos = np.sqrt(seed_E / seed_energy) * seed_pos
seed_vel = np.sqrt(seed_E / seed_energy) * seed_vel
modal_ensemble.x[:, 2] = (pri_pos + seed_pos) * mode_analysis.l0
modal_ensemble.v[:, 2] = (
pri_vel +
seed_vel) * mode_analysis.v0 # Should be within computer error of zero
# Establish positions and evolve, as per Dominic's Example
modal_positions.append(np.copy(modal_ensemble.x))
for i in range(num_steps // num_dump):
coldatoms.bend_kick(dt,
mode_analysis.B,
modal_ensemble, [coulomb_force, trap_potential],
num_steps=num_dump)
modal_positions.append(np.copy(modal_ensemble.x))
modal_positions = np.array(modal_positions)
# Convert time series to frequency data
delta_t = num_dump * dt
nu_nyquist = 0.5 / delta_t
nu_axis = np.linspace(0.0, 2.0 * nu_nyquist, modal_positions.shape[0])
freq_data = np.sum(np.abs(np.fft.fft(modal_positions[:, :, 2], axis=0))**2,
axis=1)
np.save(storage_directory + '/freqs', nu_axis)
np.save(storage_directory + '/PSD_data', freq_data)