def gaussian_measured(resolution, inverted=True):
x, y, dx, screen = make_lattice_2D(resolution, 5, 5)
wall = make_border_wall_2D(resolution, 5, 5, weight=100)
potential = wall
psi_ = exp(-50 * (x**2 + y**2)) + 0j
psi = psi_ * 0
psi[4:-4, 4:-4] = psi_[4:-4, 4:-4]
psi = normalize_2D(psi, dx)
episode_length = 0.1
mask = exp(-10000 * ((x - 0.2)**20 + (y + 0.3)**20))
if inverted:
mask = 1 - mask
return {
"dx": dx,
"screen": screen,
"psi": psi,
"potential": potential,
"episode_length": episode_length,
"measurements": {
0.02: {
"mask": mask,
"forced": True
},
},
}
def slit_base(resolution):
x, y, dx, screen = make_lattice_2D(resolution, 10, 10, extra_height=5)
potential = exp(-(8 * (x + 1))**4) * 2000
psi_ = exp(-10 * (x + 4)**2 - 10 * y**2 + 15j * x)
psi = psi_ * 0
psi[4:-4, 4:-4] = psi_[4:-4, 4:-4]
psi = normalize_2D(psi, dx)
episode_length = 0.5
return x, y, dx, screen, psi, potential, episode_length
def convex_mirror(resolution):
x, y, dx, screen = make_lattice_2D(resolution, 5, 4)
potential = exp(-(x - 3.3 + 0.15 * y * y)**4) * 7000
psi_ = exp(-30 * ((x + 1.5)**2 + y**2) + 35j * x)
psi = psi_ * 0
psi[4:-4, 4:-4] = psi_[4:-4, 4:-4]
psi = normalize_2D(psi, dx)
episode_length = 0.18
return dx, screen, psi, potential, episode_length
def moving_gaussian(resolution):
x, y, dx, screen = make_lattice_2D(resolution, 5, 2)
potential = 0 * x
psi_ = exp(-15 * ((x + 1.5)**2 + y**2) + 30j * x) + 0j
psi = psi_ * 0
psi[4:-4, 4:-4] = psi_[4:-4, 4:-4]
psi = normalize_2D(psi, dx)
episode_length = 0.1
return dx, screen, psi, potential, episode_length
def static_gaussian(resolution):
x, y, dx, screen = make_lattice_2D(resolution, 5, 5)
wall = make_border_wall_2D(resolution, 5, 5, weight=100)
potential = wall
psi_ = exp(-15 * (x**2 + y**2)) + 0j
psi = psi_ * 0
psi[4:-4, 4:-4] = psi_[4:-4, 4:-4]
psi = normalize_2D(psi, dx)
episode_length = 0.1
return dx, screen, psi, potential, episode_length
def gaussian_superposition(resolution):
x, y, dx, screen = make_lattice_2D(resolution, 5, 5)
potential = 0 * x
psi_ = exp(-60 * ((x + 1.5)**2 + (y + 0.2)**2)) + exp(-5 *
((x - 1.5)**2 +
(y - 0.2)**2)) + 0j
psi = psi_ * 0
psi[4:-4, 4:-4] = psi_[4:-4, 4:-4]
psi = normalize_2D(psi, dx)
episode_length = 0.15
return dx, screen, psi, potential, episode_length
def box_with_stuff(resolution):
x, y, dx, screen = make_lattice_2D(resolution, 10, 0.5)
wall = make_border_wall_2D(resolution, 10, 0.5, weight=1000)
potential = wall
potential += exp(-(x - 0)**2 - (y - 1.8)**2) * 200
potential += exp(-(x + y * 0.7)**8 - (x * 0.7 - y - 2)**8) * 1000
psi_ = exp(-2 * (x + 2.5)**2 - 2 * y**2 + 5j * x)
psi = psi_ * 0
psi[4:-4, 4:-4] = psi_[4:-4, 4:-4]
psi = normalize_2D(psi, dx)
episode_length = 3.0
return dx, screen, psi, potential, episode_length
def tunneling(resolution, weight=200, omega=14.75, extra_width=4):
x, y, dx, screen = make_lattice_2D(resolution,
7,
extra_width=extra_width,
extra_height=0)
potential = exp(-(2 * x)**6) * weight
potential = maximum(potential, (1 - exp(-(0.35 * y)**6)) * 1000)
psi_ = exp(-1.1 * ((x + 2.5)**2 + y**2) + omega * 1j * x)
psi = psi_ * 0
psi[4:-4, 4:-4] = psi_[4:-4, 4:-4]
psi = normalize_2D(psi, dx)
episode_length = 0.4
return dx, screen, psi, potential, episode_length
def harmonic_potential(resolution):
x, y, dx, screen = make_lattice_2D(resolution, 5, 1)
potential = (x * x + y * y) * 800
psi_ = exp(-50 * ((x + 1)**2 + (y - 0.2)**2) + 15j * y)
psi = psi_ * 0
psi[4:-4, 4:-4] = psi_[4:-4, 4:-4]
psi = normalize_2D(psi, dx)
episode_length = 0.4
return {
"dx": dx,
"screen": screen,
"psi": psi,
"potential": potential,
"episode_length": episode_length,
"dt": 0.001 * dx,
}
def colliding_gaussians(resolution):
x, y, dx, screen = make_lattice_2D(resolution, 5, 5)
potential = 0 * x
psi_ = exp(-15 * ((x + 1.5)**2 +
(y + 0.2)**2) + 31j * x) + exp(-17 *
((x - 1.5)**2 +
(y - 0.2)**2) - 28j * x)
psi = psi_ * 0
psi[4:-4, 4:-4] = psi_[4:-4, 4:-4]
psi = normalize_2D(psi, dx)
episode_length = 0.15
return {
"dx": dx,
"screen": screen,
"psi": psi,
"potential": potential,
"episode_length": episode_length,
}
type=float,
default=0.08)
args = parser.parse_args()
if args.folder and args.animate:
print('Animation is not supported while dumping to disk')
sys.exit()
if args.folder and not args.animate and not os.path.isdir(args.folder):
print("Target folder doesn't exist")
sys.exit()
width, height = RESOLUTIONS[args.resolution]
if args.episode == 'static_gaussian':
x, y, dx, screen = make_lattice_2D(args.resolution, 5, 5)
potential = 0 * x
psi = exp(-3 * (x**2 + y**2)) + 0j
psi = normalize_2D(psi[screen], dx)
total_samples = 16000
elif args.episode == 'gaussian_superposition':
dx, screen, psi, potential, episode_length = gaussian_superposition(
args.resolution)
i = 0
t = 0
dt = 0.004 * dx
while t < episode_length:
if i % 100 == 0:
print("Precalculating: {} % complete".format(
int(100 * t / episode_length)))
patch = advance_pde(t,