Skip to content
Snippets Groups Projects
Commit efb5fe3b authored by Matteo Rossi's avatar Matteo Rossi
Browse files

Removed a test ssft environment which is now obsolete

parent 512021aa
Branches
No related tags found
No related merge requests found
......@@ -4,9 +4,3 @@ register(
id='stirap-v0',
entry_point='gym_stirap.envs:StirapEnv',
)
## We can add multiple environments
register(
id='stirap-ssft-v0',
entry_point='gym_stirap.envs:StirapEnvSSFT',
)
\ No newline at end of file
from gym_stirap.envs.stirap_env_ssft import StirapEnvSSFT
from gym_stirap.envs.stirap_env import StirapEnv
\ No newline at end of file
import gym
from gym import error, spaces, utils
from gym.utils import seeding
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1.inset_locator import inset_axes
from gym_stirap.envs.schrodinger import Schrodinger
class StirapEnvSSFT(gym.Env):
metadata = {'render.modes': ['human']}
def __init__(self):
# Simulation parameters
self.n = 512 # Number of points in space
self.xlim = 2
self.x = np.linspace(-self.xlim, self.xlim, self.n) # Positions vector
self.dx = self.x[1] - self.x[0]
self.timesteps = 301
self.totaltime = .5 # Total time of evolution
self.time = np.linspace(0, self.totaltime, self.timesteps)
self.dt = self.time[1] - self.time[0]
self.mass = 2.
self.hbar = 1.
# The SE needs to be evaluated for a higher number of timesteps
# in order to ensure convergence
self.internal_timesteps = 201
ti = np.linspace(0, self.dt, self.internal_timesteps)
self.internal_dt = ti[1] - ti[0]
# This gives the ratio between dt and dx^2
# (for the solution to be stable, it has to be ~< 0.1)
print("dt/dx^2 ratio ", self.internal_dt / self.dx**2)
# STIRAP parameters
self.g = 0 # Nonlinear interaction term
self.well_size = .5
# Position of the wells
self.left_center = -1.5 * self.well_size
self.right_center = 1.5 * self.well_size
# Strength of the potential
self.trap_strength = 2.e3
# Definition of the wells
# These vectors are one where the corresponding well is located
#self.left_well = np.where(np.abs(self.x - self.left_center) <= self.well_size/2, 1, 0)
#self.right_well = np.where(np.abs(self.x - self.right_center) <= self.well_size/2, 1, 0)
self.left_well = np.where(self.x < self.left_center + 2 * self.well_size / 3, 1, 0)
self.right_well = np.where(self.x > self.right_center - 2 * self.well_size /3, 1, 0)
# Initial state of the system
# self.initial_std = .1
# self.psi0 = (1+0j)*np.exp(-(self.x- self.left_center)**2/(2 * self.initial_std ** 2))
# self.psi0 = (1+0j)*np.exp(-(self.x- self.left_center)**2/(2 * self.initial_std ** 2)) / (np.pi**0.25 * np.sqrt(self.initial_std))
# Initial state is the ground state of the left trap
self.omega = np.sqrt(self.trap_strength/self.mass)
self.psi0 = (1+0j) * (self.mass * self.omega / (np.pi * self.hbar)) ** 0.25 * np.exp(-self.mass * self.omega * (self.x -self.left_center) ** 2/ (2*self.hbar))
print("Norm: ", self.norm(self.psi0, self.x))
# self.psi0 /= self.norm(self.psi0, self.x)
# Initialize solver
self.it = None
self.il = None
self.ir = None
# RL parameters
self.Delta = 2e2 * self.dx * self.dt # How much to change the well positions at each step
# Actions for a well: (- Delta, 0, + Delta)
tmp = self.Delta * np.array([-1, 0, 1])
# Carthesian product of actions on the two wells
self.actions = np.transpose([np.tile(tmp, len(tmp)), np.repeat(tmp, len(tmp))])
self.action_space = spaces.Discrete(len(self.actions))
self.observation_space = spaces.Box(low=np.array([0, 0, self.left_center - self.well_size/2, self.right_center - self.well_size/2]),
high=np.array([1, 1, self.left_center + self.well_size/2, self.right_center + self.well_size/2]),dtype=np.float32)
# For the rendering
self.viewer = None
# We compute the kinetic matrix once and for all
# self.K = self.kinematic_matrix(self.n) / (self.dx ** 2)
# Initializes the environment
self.reset()
self.solver = Schrodinger(self.x, self.psi0, self.potential(), m= self.mass, hbar=self.hbar)
def step(self, action):
assert self.action_space.contains(action), "%r (%s) invalid"%(action, type(action))
# Update the positions of the wells with the chosen action,
self.il += self.actions[action, 0]
self.ir += self.actions[action, 1]
# Evaluate the state at time t + dt
self.drive()
# This is the state that the agent sees (left/right population, position of the wells)
self.state = self.evaluate_populations() + (self.il, self.ir)
# Calculate the reward
reward = self.reward()
# Check if the simulation is done
self.it += 1
if self.it == self.timesteps:
done = True
else:
done = False
return np.array(self.state), reward, done, {}
def reset(self):
"""Reinitialize the system to the initial conditions"""
self.psi = self.psi0.copy() # We need a deep copy
self.il = self.left_center
self.ir = self.right_center
self.state = self.evaluate_populations() + (self.il, self.ir)
self.it = 0
return np.array(self.state)
def render(self, mode='human', close=True):
# Probability density
ys = np.abs(self.psi)**2 * self.dx
if self.viewer is None:
# Here we initialize the plot
self.viewer = plt.figure()
ax = self.viewer.gca()
self.probability_line, = plt.plot(self.x, ys)
self.potential_line, = plt.plot(self.x, self.dx * self.potential())
plt.ylim([0, .2])
plt.plot(self.x, self.left_well)
plt.plot(self.x, self.right_well)
self.axins = inset_axes(ax, width="30%", height="40%", loc='upper left')
self.axins.set_xlim([0, self.totaltime])
self.axins.set_ylim([0, 1])
self.rightpop_line, = self.axins.plot(self.time[self.it], self.state[1])
plt.pause(0.000001)
# Here we update the plot
self.probability_line.set_ydata(ys)
self.potential_line.set_ydata(self.dx * self.potential())
self.rightpop_line.set_xdata(np.append(self.rightpop_line.get_xdata(),self.time[self.it]))
self.rightpop_line.set_ydata(np.append(self.rightpop_line.get_ydata(),self.state[1]))
plt.pause(0.000001)
def reward(self):
""" Define the reward """
(pop_left, pop_right, _, _) = self.state
# We reward for the population in the last 90 % of time
#reward = 10 * pop_right * (self.it > .9*self.timesteps)
reward = pop_right
# We punish the agent for moving the traps too far away, or for crossing them
if self.il < (- self.xlim + 1 * self.well_size ) or self.il > 0:
reward -= 5
if self.ir < 0 or (self.ir > self.xlim - 1 * self.well_size):
reward -= 5
return reward
# These are auxiliary functions for the physical system
# Function that take as input the three parameters and returs
# the two figure of merits at each step
def drive(self):
v = self.potential()
self.solver.time_step(self.internal_dt, v, Nsteps = self.internal_timesteps)
self.psi = self.solver.psi_x
# We need to evaluate the SE for smaller dt for each simulation dt
#for i in range(self.internal_timesteps):
# self.psi = self.rk4(self.internal_dt, v, self.g, self.psi)
# Use Runge-Kutta scheme (more stable)
# Use Euler scheme (Slightly Faster)
# self.psi = self.eul(self.dt,v,self.g,self.psi)
def norm(self, psi, x):
return np.trapz(np.abs(psi)**2, x)
def evaluate_populations(self):
left = np.sum(np.abs(self.psi * self.left_well)**2) * self.dx
right = np.sum(np.abs(self.psi * self.right_well)**2) * self.dx
return (left, right)
def potential(self):
""" Returns the trapping potential """
#self.il = np.clip(self.il, self.left_center - self.well_size / 2, self.left_center + self.well_size / 2)
#self.ir = np.clip(self.ir, self.right_center - self.well_size / 2, self.right_center + self.well_size / 2)
#v = np.zeros_like(self.x)
v = 0.5 * self.trap_strength * ((self.x-self.il)**2) * (self.x ** 2) * ((self.x-self.ir)**2)
# + np.exp( (self.x-self.x[-1]-self.dx)** -2 ) + np.exp( (self.x-self.x[0]-self.dx)** -2)
return v
# Function for the kinetic energy matrix
def kinematic_matrix(self, n):
A = np.diagflat(-2*np.ones(n)) + np.diagflat(np.ones(n-1), 1) + np.diagflat(np.ones(n-1), -1)
A[0,:] = 0
A[-1,:] = 0
return A
#Function for the nonlinear term
def nonlin(self, psi):
return -1j * self.g * np.conjugate(psi) * psi * psi
#Function for single time step evolution
def der(self, v, g, psi):
# dpsi = 1j * np.dot(self.K, psi)
# dpsi = 1j * self.hbar ** 2 / (2 * self.mass) * np.gradient(np.gradient(psi, self.dx), self.dx)
# dpsi[0] = 0.j
# dpsi[-1] = 0.j
dpsi = 1j * self.hbar ** 2 / (2 * self.mass) * fastLaplace1d(psi) / self.dx**2 # Faster than gradient, probably less accurate
dpsi -= 1j * v * psi # Potential part
if self.g > 0:
dpsi += self.nonlin(psi)
return dpsi
#Euler scheme
def eul(self, h, v, g, psi):
der1 = h * self.der(v, g, psi)
psi += der1
norm = self.norm(psi, self.x)
if np.abs(1-norm) > 0.005:
print('Becoming unstable, norm= ', norm)
psi /= norm
return psi
# Runge-Kutta
def rk4(self, dt, v, g, psi):
k1 = dt * self.der(v,g, psi)
psi1 = self.psi + 0.5 * k1
k2 = dt * self.der(v,g,psi1)
psi2 = self.psi + 0.5 * k2
k3 = dt * self.der(v,g,psi2)
psi3 =self.psi + k3
k4 = dt * self.der(v,g,psi3)
psi += (k1 + 2*k2 + 2*k3 + k4)/6
norm = self.norm(psi, self.x)
if np.abs(1 - norm) > 0.005:
print('Becoming unstable, norm= ', norm)
psi /=norm
return psi
from scipy.ndimage import _nd_image
laplace_filter = np.asarray([1, -2, 1], dtype=np.float64)
def fastLaplace1d(arr):
"""Evaluates the Laplacian operator for a 1D complex vector"""
res = np.empty_like(arr)
_nd_image.correlate1d(np.real(arr), laplace_filter, 0, np.real(res), 2, 0.0, 0)
_nd_image.correlate1d(np.imag(arr), laplace_filter, 0, np.imag(res), 2, 0.0, 0)
return res
\ No newline at end of file
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment