def runPass():
env = envs.cuda(device_num=1)
constants = Constants(double=env.supportsDouble(),
a11=100.4, a12=97.7, a22=95.0, fx=f_ax, fy=f_ax, fz=f_rad, e_cut=e_cut)
# axial size of the N / 2 cloud ~ 8e-6 >> potentials_separation
# therefore we can safely use normal grid, provided that it has big enough border
box = constants.boxSizeForN(N, 3, border=1.2)
grid = UniformGrid(env, constants, shape,
(box[0] + potentials_separation * 2, box[1], box[2]))
#print constants.planeWaveModesForCutoff(constants.boxSizeForN(N, 3, border=2))
gs = SplitStepGroundState(env, constants, grid, dt=1e-6)
pulse = Pulse(env, constants, grid, f_rabi=f_rabi, f_detuning=f_detuning)
evolution = SplitStepEvolution(env, constants, grid,
potentials=split_potentials(constants, grid),
dt=1e-6)
n = ParticleNumberCollector(env, constants, grid, verbose=False)
v = VisibilityCollector(env, constants, grid)
# a = AxialProjectionCollector(env, constants, grid, pulse=pulse)
u = UncertaintyCollector(env, constants, grid)
s = SpinCloudCollector(env, constants, grid)
psi = gs.create((N, 0))
psi.toMSpace()
mode_data = numpy.abs(env.fromDevice(psi.data))[0, 0] # remember mode data
mask = buildProjectorMask(constants, grid)
psi.toXSpace()
psi.toWigner(128)
pulse.apply(psi, numpy.pi / 2)
evolution.run(psi, splitting_time, callbacks=[v, n, u, s],
callback_dt=splitting_time / 100)
env.synchronize()
env.release()
"""
times, heightmap = a.getData()
HeightmapPlot(
HeightmapData("test", heightmap,
xmin=0, xmax=splitting_time * 1e3,
ymin=grid.z[0] * 1e6,
ymax=grid.z[-1] * 1e6,
zmin=-1, zmax=1,
xname="T (ms)", yname="z ($\\mu$m)", zname="Spin projection")
).save('split_potentials_axial.pdf')
"""
times, n_stddev, xi_squared = u.getData()
XYData("Squeezing", times * 1000, numpy.log10(xi_squared),
xname="T (ms)", yname="log$_{10}$($\\xi^2$)",
description=parameters).save('split_potentials_xi.json')
times, vis = v.getData()
XYData('test', times * 1e3, vis,
xname="T (ms)", yname="$\\mathcal{V}$",
ymin=0, ymax=1,
description=parameters).save('split_potentials_vis.json')
times, phi, yps, Sx, Sy, Sz = s.getData()
return times, Sx, Sy, Sz
def runTest(env, dim, grid_type, prop_type, use_cutoff, use_big_grid):
# additional parameters
constants_kwds = {
'1d': dict(use_effective_area=True, fx=42e3, fy=42e3, fz=90),
'3d': {}
}[dim]
# total number of atoms in ground state
total_N = {
'1d': 60,
'3d': 50000
}[dim]
# number of lattice points (for no cutoff mode)
shape = {
('1d', 'uniform'): (64,),
('3d', 'uniform'): (128, 8, 8),
('1d', 'harmonic'): (50,),
('3d', 'harmonic'): (50, 10, 10)
}[(dim, grid_type)]
# time step for split-step propagation
ss_dt = {
'1d': 1e-6,
'3d': 1e-5
}[dim]
e_cut = {
'1d': 12000,
'3d': 3000
}[dim]
total_time = {
'1d': 0.1,
'3d': 1.0
}[dim]
# Prepare constants and grid
constants = Constants(double=env.supportsDouble(),
e_cut=(e_cut if use_cutoff else None),
**constants_kwds)
if grid_type == 'uniform':
if use_cutoff:
box_size = constants.boxSizeForN(total_N, len(shape))
min_shape = constants.planeWaveModesForCutoff(box_size)
# FFT supports 2**n dimensions only
shape = tuple([2 ** (log2(x - 1) + 1) for x in min_shape])
if use_big_grid:
shape = tuple([2 * x for x in shape])
grid = UniformGrid.forN(env, constants, total_N, shape)
elif grid_type == 'harmonic':
if use_cutoff:
shape = constants.harmonicModesForCutoff(len(shape))
if use_big_grid:
shape = tuple([x + 3 for x in shape])
grid = HarmonicGrid(env, constants, shape)
if grid_type == 'uniform':
gs = SplitStepGroundState(env, constants, grid, dt=ss_dt)
elif grid_type == 'harmonic':
gs = RK5HarmonicGroundState(env, constants, grid, eps=1e-7)
if grid_type == 'harmonic':
evolution = RK5HarmonicEvolution(env, constants, grid,
atol_coeff=1e-3, eps=1e-6, Nscale=total_N)
elif prop_type == 'split-step':
evolution = SplitStepEvolution(env, constants, grid, dt=ss_dt)
elif prop_type == 'rk5':
evolution = RK5IPEvolution(env, constants, grid, Nscale=total_N,
atol_coeff=1e-3, eps=1e-6)
pulse = Pulse(env, constants, grid, f_detuning=41, f_rabi=350)
a = AxialProjectionCollector(env, constants, grid, pulse=pulse)
p = ParticleNumberCollector(env, constants, grid, pulse=pulse)
v = VisibilityCollector(env, constants, grid)
# experiment
psi = gs.create((total_N, 0))
pulse.apply(psi, theta=0.5 * numpy.pi)
t1 = time.time()
evolution.run(psi, total_time, callbacks=[a, p, v], callback_dt=0.005)
env.synchronize()
t2 = time.time()
times, heightmap = a.getData()
# check that the final state is still projected
psi.toMSpace()
mode_data = env.fromDevice(psi.data)
mask = numpy.tile(buildProjectorMask(constants, grid),
(psi.components, 1) + (1,) * grid.dim)
masked_mode_data = mode_data * (1.0 - mask)
assert masked_mode_data.max() < 1e-6 * mode_data.max()
res = HeightmapPlot(
HeightmapData("test", heightmap,
xmin=0, xmax=total_time * 1e3,
ymin=grid.z[0] * 1e6,
ymax=grid.z[-1] * 1e6,
zmin=-1, zmax=1,
xname="T (ms)", yname="z ($\\mu$m)", zname="Spin projection")
)
times, Ns, Ntotals = p.getData()
times, vis = v.getData()
print (" Shape: {shape}, final N: {N} ({Ntotal}), V: {vis}\n" +
" Time spent: {t} s").format(
shape=shape, N=Ns[:,-1], Ntotal=Ntotals[-1], vis=vis[-1], t=t2-t1)
return res
def runTest(env, comp, grid_type, dim, gs_type, use_cutoff):
"""
Creates Thomas-Fermi ground state using different types of representations
"""
# additional parameters
constants_kwds = {
'1d': dict(use_effective_area=True, fx=42e3, fy=42e3, fz=90),
'3d': {}
}[dim]
# total number of atoms in ground state
total_N = {
'1d': 60,
'3d': 50000
}[dim]
# number of lattice points
shape = {
('1d', 'uniform'): (64,),
('3d', 'uniform'): (128, 8, 8),
('1d', 'harmonic'): (50,),
('3d', 'harmonic'): (50, 10, 10)
}[(dim, grid_type)]
# time step for split-step propagation
ss_dt = {
'1d': 1e-6,
'3d': 1e-5
}[dim]
# absolute precision for split-step algorithm
ss_precision = {
'1comp': 1e-6,
'2comp': 1e-8
}[comp]
# precision divided by time step for RK5 propagation
rk5_rprecision = {
('1d', '1comp'): 1,
('3d', '1comp'): 1,
('1d', '2comp'): 1e-2,
('3d', '2comp'): 1e-3
}[(dim, comp)]
# relative error tolerance for RK5 propagation
rk5_rtol = {
('1d', 'uniform'): 1e-9,
('1d', 'harmonic'): 1e-7,
('3d', 'uniform'): 1e-6,
('3d', 'harmonic'): 1e-6
}[(dim, grid_type)]
# absolute error tolerance divided by atom number for RK5 propagation
rk5_atol_coeff = 2e-3
target_N = {
'1comp': (total_N, 0),
'2comp': (int(total_N / 2 + 0.05 * total_N), int(total_N / 2 - 0.05 * total_N))
}[comp]
e_cut = {
'1d': 8000,
'3d': 7000
}[dim]
# Prepare constants and grid
constants = Constants(double=env.supportsDouble(),
e_cut=(e_cut if use_cutoff else None),
**constants_kwds)
if grid_type == 'uniform':
grid = UniformGrid.forN(env, constants, total_N, shape)
elif grid_type == 'harmonic':
grid = HarmonicGrid(env, constants, shape)
# Prepare 'apparatus'
args = (env, constants, grid)
if gs_type == "TF":
gs = TFGroundState(*args)
gs_kwds = {}
elif gs_type == "split-step":
gs = SplitStepGroundState(*args, dt=ss_dt)
gs_kwds = dict(precision=ss_precision)
elif gs_type == "rk5":
params = dict(eps=rk5_rtol, atol_coeff=rk5_atol_coeff)
gs_kwds = dict(relative_precision=rk5_rprecision)
if grid_type == 'uniform':
gs = RK5IPGroundState(*args, **params)
elif grid_type == 'harmonic':
gs = RK5HarmonicGroundState(*args, **params)
# Create ground state
t1 = time.time()
psi = gs.create(target_N, **gs_kwds)
t2 = time.time()
t_gs = t2 - t1
prj = ProjectionMeter.forPsi(psi)
obs = IntegralMeter.forPsi(psi)
# check that 2-component stats object works properly
N_xspace1 = psi.density_meter.getNTotal()
psi.toMSpace()
N_mspace = psi.density_meter.getNTotal()
mode_data = numpy.abs(env.fromDevice(psi.data)) # remember mode data
psi.toXSpace()
# population in x-space after double transformation (should not change)
N_xspace2 = psi.density_meter.getNTotal()
# calculate energy and chemical potential (per particle)
E = obs.getEPerParticle(psi) / constants.hbar / constants.wz
mu = obs.getMuPerParticle(psi) / constants.hbar / constants.wz
mu_tf = numpy.array(
[constants.muTF(N, dim=grid.dim) for N in target_N]
).sum() / constants.hbar / constants.wz
# Checks
norm = numpy.linalg.norm
target_norm = norm(numpy.array(target_N))
# check that number of particles right after GS creation is correct
assert N_xspace1.shape == (2,)
assert norm(N_xspace1 - numpy.array(target_N)) / target_norm < 1e-6
# check that double transform did not change N
assert norm(N_xspace1 - N_xspace2) / target_norm < 1e-6
# TODO: find out what causes difference even for uniform grid
assert norm(N_mspace - N_xspace2) / target_norm < 0.02
assert E.shape == (2,)
assert mu.shape == (2,)
if comp == '1comp':
assert E[1] == 0
else:
assert E[1] != 0
if comp == '1comp':
assert mu[1] == 0
else:
assert mu[1] != 0
# There should be some difference between analytical mu and numerical one,
# so it is more of a sanity check
assert abs(mu.sum() - mu_tf) / mu_tf < 0.35
# Check that GS is really restricted by mask
mask = numpy.tile(buildProjectorMask(constants, grid),
(psi.components, 1) + (1,) * grid.dim)
masked_mode_data = mode_data * (1.0 - mask)
assert masked_mode_data.max() < 1e-6 * mode_data.max()
E = E.sum()
mu = mu.sum()
Nx = N_xspace2.sum()
Nm = N_mspace.sum()
# Results
print (" Modes: {modes_num} out of {total_modes}\n" +
" N(x-space) = {Nx:.4f}, N(m-space) = {Nm:.4f}, " +
"E = {E:.4f} hbar wz, mu = {mu:.4f} hbar wz\n" +
" Time spent: {t_gs} s").format(
Nx=Nx, Nm=Nm, E=E, mu=mu, t_gs=t_gs, modes_num=int(mask.sum()), total_modes=mask.size)
z = grid.z * 1e6 # cast to micrometers
profile = prj.getZ(psi) / 1e6 # cast to micrometers^-1
plots = [
XYData(str(env) + ", " + grid_type + ", " + gs_type + " |" + str(i + 1) + ">",
z, profile[i],
xname="Z ($\\mu$m)", yname="Axial density ($\\mu$m$^{-1}$)", ymin=0,
linestyle=('-' if i == 0 else '--'))
for i in xrange(len(profile))
]
return plots[:(1 if comp == '1comp' else 2)]
