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)]