From 41529cf49cd29de27f642050e6c04f8e4a241c12 Mon Sep 17 00:00:00 2001
From: Matteo Rossi <teo.red90@gmail.com>
Date: Mon, 29 Jul 2019 13:43:29 +0300
Subject: [PATCH] Added target state to render

---
 gym_stirap/envs/stirap_env.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/gym_stirap/envs/stirap_env.py b/gym_stirap/envs/stirap_env.py
index fccd9a1..1be42b4 100644
--- a/gym_stirap/envs/stirap_env.py
+++ b/gym_stirap/envs/stirap_env.py
@@ -262,6 +262,8 @@ class StirapEnv(gym.Env):
 
             self.probability_line, = plt.plot(self.x, ys)
             self.potential_line, = plt.plot(self.x, 0.5*self.dx * self.potential(),'k')
+            plt.fill_between(self.x, 2 * np.abs(self.target_psi)**2 * self.dx, color='C2',alpha=.3)
+            
             plt.ylim([0, .2])
 
             if timeevolution:
-- 
GitLab