def plotDensity(psi1, psi2, n1, n2, name):
z = psi1._grid.z * 1e6
prj = ProjectionMeter.forPsi(psi1)
z_proj1 = prj.getZ(psi1) / 1e6 # cast to micrometers^-1
z_proj2 = prj.getZ(psi2) / 1e6 # cast to micrometers^-1
datas = [z_proj1]
plots = []
for comp in xrange(psi1.components):
plots += [
XYData(n1 + ", component " + str(comp + 1), z, z_proj1[comp],
xname="Z ($\\mu$m)", yname="Axial density ($\\mu$m$^{-1}$)", ymin=0),
XYData(n2 + ", component " + str(comp + 1), z, z_proj2[comp],
xname="Z ($\\mu$m)", yname="Axial density ($\\mu$m$^{-1}$)", ymin=0),
]
XYPlot(plots).save(name)
import numpy from beclab import * from beclab.meters import ProjectionMeter N = 20000 env = envs.cuda() constants = Constants(double=env.supportsDouble()) grid = UniformGrid.forN(env, constants, N, (128, 32, 16)) gs = SplitStepGroundState(env, constants, grid, dt=1e-5) psi = gs.create((N / 2, N / 2), precision=1e-8) prj = ProjectionMeter.forPsi(psi) # Create various density profiles prefix = "density_snapshots" # cast grid points to micrometers x = grid.x * 1e6 y = grid.y * 1e6 z = grid.z * 1e6 # Projection on Z axis z_proj = prj.getZ(psi) / 1e6 # cast to micrometers^-1 # Projection on XY and YZ planes xy_proj = prj.getXY(psi) / 1e12 # cast to micrometers^-2 yz_proj = prj.getYZ(psi) / 1e12 # cast to micrometers^-2
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)]
