Change plotting of alignment in rotational density example

examples/rot_dens/dens_sampling.py

@@ -91,6 +91,7 @@ def alignment(rotmat: List[np.ndarray],
         beta = fl['beta'][()]
         gamma = fl['gamma'][()]
         times = fl['times'][()]
+        field = fl['field'][()]
         cos2theta = fl['cos2theta'][()]
         costheta = fl['costheta'][()]
         cos2theta_wig = fl['cos2theta_wig'][()]
@@ -110,20 +111,32 @@ def alignment(rotmat: List[np.ndarray],
     # rotation matrix for Euler angle samples
     rotmat = [R.from_euler('zyz', pts).as_matrix() for pts in points]
 
+    # expectation values of alignment functions`
     costheta_, cos2theta_, costheta2d_, cos2theta2d_ = alignment(rotmat)
 
-    plt.plot(times, cos2theta, '-r', linewidth=2, label="$\cos^2\\theta$ analytic")
-    plt.plot(times, cos2theta_wig, '--g', linewidth=2, label="$\cos^2\\theta$ Wigner")
-    plt.plot(times, cos2theta_, ':b', linewidth=2, label="$\cos^2\\theta$ sampling")
+    # plot results
 
-    plt.plot(times, cos2theta2d_wig, '--g', linewidth=2, label="$\cos^2\\theta_{\\rm 2D}$ Wigner")
-    plt.plot(times, cos2theta2d_, ':b', linewidth=2, label="$\cos^2\\theta_{\\rm 2D}$ sampling")
+    fig, ax1 = plt.subplots()
 
-    plt.plot(times, costheta, '-r', linewidth=2, label="$\cos\\theta$ analytic")
-    plt.plot(times, costheta_, ':b', linewidth=2, label="$\cos\\theta$ sampling")
+    ax1.set_xlabel("Time in ps")
+    ax1.set_ylabel("Field in V/m")
+    ax1.plot(times, field, color='gray')
+    ax1.fill_between(times, field[:,-1], 0, color='gray', alpha=.1)
 
-    plt.xlabel("time in ps")
-    plt.legend()
+    ax2 = ax1.twinx()
+
+    ax2.set_ylabel("Alignment")
+    ax2.plot(times, cos2theta, '-r', linewidth=2, label="$\cos^2\\theta$ analytic")
+    ax2.plot(times, cos2theta_wig, '--g', linewidth=2, label="$\cos^2\\theta$ Wigner")
+    ax2.plot(times, cos2theta_, ':b', linewidth=2, label="$\cos^2\\theta$ Monte-Carlo")
+
+    ax2.plot(times, cos2theta2d_wig, '--b', linewidth=2, label="$\cos^2\\theta_{\\rm 2D}$ Wigner")
+    ax2.plot(times, cos2theta2d_, ':r', linewidth=2, label="$\cos^2\\theta_{\\rm 2D}$ Monte-Carlo")
+
+    ax2.plot(times, costheta, '-b', linewidth=2, label="$\cos\\theta$ analytic")
+    ax2.plot(times, costheta_, ':g', linewidth=2, label="$\cos\\theta$ Monte-Carlo")
+
+    plt.legend(loc='center left', bbox_to_anchor=(1.1, 0.5))
     plt.draw()
     plt.pause(0.0001)
-    plt.savefig(f"ocs_alignment_test.png")
+    plt.savefig(f"ocs_alignment_test.png", format = "png", dpi = 300, bbox_inches = "tight")

examples/rot_dens/ocs_alignment.py

@@ -60,10 +60,10 @@
 
 
     # set up 800nm Gaussian pulse
-    def field(t):
+    def field(t, omega_nm):
         nm = 1e-9
         fwhm = 10.0 # ps
-        omega = 800 * nm # in nm
+        omega = omega_nm * nm # in nm
         omega = 2 * np.pi * speed_of_light / omega * 1e-12  # in 1/ps
         t0 = 2.5 * fwhm / 2
         amp = 1e10 # in V/m
@@ -102,7 +102,7 @@ def field(t):
 
         # apply field to Hamiltonian
         thresh = 1e3 # thresh for considering field as zero
-        H.field(field(t), thresh=thresh)
+        H.field(field(t, 800), thresh=thresh)
 
         # update vector
         vecs, t_ = tdse.update(H, H0=h0, vecs=vecs, matvec_lib='scipy')
@@ -154,6 +154,7 @@ def field(t):
         fl.create_dataset("beta", data=beta)
         fl.create_dataset("gamma", data=gamma)
         fl.create_dataset("times", data=np.array(times))
+        fl.create_dataset("field", data=np.array([field(t, 1e10) for t in times]))
 
         # for comparison, also store the alignment functions
         fl.create_dataset("cos2theta", data=np.array(cos2th_ev))