OQuPaper performance test

examples/nt_corrs_example.py

@@ -0,0 +1,205 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Created on Thu Mar 28 09:24:44 2024
+
+@author: rmadw
+"""
+
+import sys
+sys.path.insert(0,'..')
+
+import oqupy
+#from oqupy.contractions import compute_correlations_nt
+import n_time_correlations as nt #with code implemented in oqupy, uncomment above and delete this import
+import numpy as np
+import matplotlib.pyplot as plt
+import matplotlib.colors
+from mpl_toolkits.axes_grid1 import make_axes_locatable
+plt.rc('font', size=8)          # controls default text sizes
+plt.rc('axes', titlesize=8)     # fontsize of the axes title
+plt.rc('axes', labelsize=8)    # fontsize of the x and y labels
+plt.rc('xtick', labelsize=8)    # fontsize of the tick labels
+plt.rc('ytick', labelsize=8)    # fontsize of the tick labels
+plt.rc('legend', fontsize=8)    # legend fontsize
+plt.rc('figure', titlesize=8)  # fontsize of the figure title
+from scipy.fft import fftfreq, fftshift, fft2, ifft
+import os
+
+os.chdir('/home/rmadw/Documents/OQuPy')
+PT_DIR_PATH = "./tests/data/process_tensors/"
+
+
+######################useful operators#########################################
+P_1 = np.array([[0., 0., 0.],
+                [0., 1., 0.],
+                [0., 0., 0.]], dtype=complex)
+
+P_2 = np.array([[0., 0., 0.],
+                [0., 0., 0.],
+                [0., 0., 1.]], dtype=complex)
+
+sigma_min = np.array([[0., 0., 0.],
+                      [0., 0., 1.],
+                      [0., 0., 0.]], dtype=complex)
+
+sigma_plus = np.array([[0., 0., 0.],
+                       [0., 0., 0.],
+                       [0., 1., 0.]], dtype=complex)
+
+######################compute the process tensor###############################
+omega_cutoff = 3.04
+alpha =0.1
+temperature = 13.09 #=100 K
+dt=0.1
+start_time = 0.
+end_time = dt*80
+dkmax=500
+epsrel=10**(-6)
+
+tempo_parameters = oqupy.TempoParameters(dt=dt, dkmax=dkmax, epsrel=epsrel)
+syst_int = P_1  - P_2
+correlations = oqupy.PowerLawSD(alpha=alpha,
+                                zeta=1,
+                                cutoff=omega_cutoff,
+                                cutoff_type='exponential',
+                                temperature=temperature)
+bath = oqupy.Bath(syst_int, correlations)
+
+pt_file_path = os.path.join(PT_DIR_PATH, "3ls_alpha0.1_zeta1.0_T13.09_cutoff3.04exp_tcut50.0_dt0.1_steps80_epsrel6.hdf5")
+process_tensor = oqupy.import_process_tensor(pt_file_path)
+
+######################define system + dipole operators#########################
+eps = 5.
+omeg= 2.
+reorg = 2.0*alpha*omega_cutoff
+system = oqupy.System((eps+reorg)*(P_1 + P_2)
+                      + omeg * (sigma_plus + sigma_min))
+
+dip_v = np.array([[0., 0., 1.],
+       [0., 0., 0.],[1., 0., 0.]], dtype=complex)
+
+initial_state = np.array([[1., 0., 0.],
+       [0., 0., 0.],[0., 0., 0.]], dtype=complex)
+
+##########################Calculate four-time correlations###################
+dipole_ops = [dip_v, dip_v, dip_v, dip_v]
+
+order_1 = ["left", "right", "right", "left"]
+order_2 = ["right", "left", "right", "left"]
+order_3 =  ["right", "right", "left", "left"]
+order_4 = ["left", "left", "left", "left"]
+
+ops_orders = [order_1, order_2, order_3, order_4]
+
+times_1 = (start_time, dt*40+dt)
+times_2 = dt*40
+times_3 = dt*40
+times_4 = (dt*40, end_time)
+
+ops_times = [times_1, times_2, times_3, times_4]
+
+cors=[]
+
+for i in range (len(ops_orders)):
+    cor = nt.compute_nt_correlations(system = system,
+                                      process_tensor=process_tensor,
+                                      dipole_ops = dipole_ops,
+                                      ops_times=ops_times,
+                                      ops_order=ops_orders[i],
+                                      dt = dt,
+                                      initial_state = initial_state,
+                                      start_time = start_time,
+                                      progress_type = "bar")
+    cors.append(cor)
+
+##########################Calculate two-time correlations###################
+order = ["left", "left"]
+
+times = [start_time, (start_time, end_time)]
+
+cor2 = nt.compute_nt_correlations(system = system,
+                                  process_tensor=process_tensor,
+                                  dipole_ops = [dip_v, dip_v],
+                                  ops_times=times,
+                                  ops_order=order,
+                                  dt = dt,
+                                  initial_state = initial_state,
+                                  start_time = start_time,
+                                  progress_type = "bar")
+
+pad=500
+
+t2 = cor2[0][1]
+
+f=fftshift(ifft(cor2[1][0], n=(t2.size+pad)))
+
+x=2*np.pi*fftshift(fftfreq(t2.size+pad,dt))
+
+
+######################Plot#########################################
+desired_folder = '/home/rmadw/Documents/OQuPy/examples'
+
+Rs = []
+for i in range (4):
+    R = cors[i][1][:,0,0,:]
+    R = R[::-1, :]
+    Rs.append(R)
+
+pad=500
+
+Rfs=[]
+for i in range (4):
+    Rpad = np.pad(Rs[i], ((0,pad),(0,pad+1)), 'constant')
+    Rf=fftshift((fft2(Rpad)))
+    Rfs.append(Rf)
+
+time = cors[0][0][0]
+f_time = 2*np.pi*fftshift(fftfreq(time.size+pad,dt))
+
+fig, ax = plt.subplots(1,1, figsize = (4/2.54, 3/2.54))
+yax = np.flip(Rfs[0].real) + np.flip(Rfs[1].real,1) + np.flip(Rfs[2].real,1) + np.flip(Rfs[3].real)
+cont1=ax.contour(f_time, f_time, yax, levels=15, linewidths=0.75)
+cbar = fig.colorbar(cont1)
+ax.set_xlim([-5, 15])
+ax.set_ylim([-5, 15])
+ax.set_aspect('equal', adjustable='box')
+ax.set_xlabel(r'$\omega_{detec}\,\, (ps^{-1})$')
+ax.set_ylabel(r'$\omega_{exc}\,\,(ps^{-1})$')
+ax.plot([0, 1], [0, 1], '--', color='gray', transform=ax.transAxes, linewidth=0.75)
+file_name = '2Dspectr_small_dt='+str(dt)+'_dkmax='+str(dkmax)+'_epsrel='+str(epsrel)+'_alph='+str(alpha)+'_temp='+str(np.round(temperature))+'omeg='+str(omeg)+'eps='+str(eps)+'wc='+str(omega_cutoff)+'t2='+str(0)+'.pdf'
+full_path = os.path.join(desired_folder, file_name)
+#plt.savefig(full_path, dpi=300, bbox_inches = "tight")
+
+fig2, ax2 = plt.subplots(1,1, figsize = (4/2.54, 3/2.54))
+ax2.plot(x, f, linewidth=1.5)
+ax2.set_xlim([-10.,25.])
+#ax2.set_ylim([-0.001,0.38])
+ax2.set(xlabel=r'$\omega\,(ps^{-1})$', ylabel=r'Linear absorption (arb. units)')
+file_name = '1Dspectr_small_dt='+str(dt)+'_dkmax='+str(dkmax)+'_epsrel='+str(epsrel)+'_alph='+str(alpha)+'_temp='+str(np.round(temperature))+'omeg='+str(omeg)+'eps='+str(eps)+'wc='+str(omega_cutoff)+'t2='+str(0)+'.pdf'
+full_path = os.path.join(desired_folder, file_name)
+#plt.savefig(full_path, dpi=300, bbox_inches = "tight")
+
+
+figs, axs = plt.subplots(2,1, figsize = (5/2.54, 7/2.54))
+axs[0].plot(x, f, linewidth=1.5)
+axs[0].set(xlabel=r'$\omega\,(ps^{-1})$', ylabel=r'Linear absorption (arb. units)')
+axs[0].set_xlim([-10, 25])
+cont1=axs[1].contour(f_time, f_time, yax, levels=15, linewidths = 0.5)
+axs[1].plot([0, 1], [0, 1], '--', color='gray', transform=axs[1].transAxes, linewidth=0.75)
+axs[1].set_aspect('equal', adjustable='box')
+axs[1].set_xlim([-10, 25])
+axs[1].set_ylim([-10, 25])
+axs[1].set(xlabel = r'$\omega_{detec}$ (ps$^{-1})$', ylabel = r'$\omega_{exc}\,\,(ps^{-1})$')
+cont1.monochrome = True
+for col, ls in zip(cont1.collections, cont1._process_linestyles()):
+    col.set_linestyle(ls)
+norm= matplotlib.colors.Normalize(vmin=cont1.cvalues.min(), vmax=cont1.cvalues.max())
+sm = plt.cm.ScalarMappable(norm=norm, cmap = cont1.cmap)
+sm.set_array([])
+divider = make_axes_locatable(axs[1])
+cax = divider.append_axes('right', size='5%', pad=0.05)
+fig.colorbar(sm, cax=cax, orientation='vertical')
+file_name = 'spectr_dt='+str(dt)+'_dkmax='+str(dkmax)+'_epsrel='+str(epsrel)+'_alph='+str(alpha)+'_temp='+str(np.round(temperature))+'omeg='+str(omeg)+'eps='+str(eps)+'wc='+str(omega_cutoff)+'t2='+str(0)+'.pdf'
+full_path = os.path.join(desired_folder, file_name)
+#plt.savefig(full_path, dpi=300, bbox_inches = "tight")

tests/performance/analysis/nt_corrs_plots.py

@@ -0,0 +1,67 @@
+# Copyright 2024 The TEMPO Collaboration
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Script to plot the multi-time correlations performance analysis results.
+"""
+
+import sys
+sys.path.insert(0,'.')
+
+import numpy as np
+import matplotlib.pyplot as plt
+plt.rc('font', size=8)          # controls default text sizes
+plt.rc('axes', titlesize=8)     # fontsize of the axes title
+plt.rc('axes', labelsize=8)    # fontsize of the x and y labels
+plt.rc('xtick', labelsize=8)    # fontsize of the tick labels
+plt.rc('ytick', labelsize=8)    # fontsize of the tick labels
+plt.rc('legend', fontsize=8)    # legend fontsize
+plt.rc('figure', titlesize=8)  # fontsize of the figure title
+
+import dill
+
+import oqupy
+import oqupy.operators as op
+
+import os
+
+
+os.chdir('/home/rmadw/Documents/OQuPy')
+
+with open("./tests/data/performance_results/nt_corrs.pkl", 'rb') as f:
+    all_results = dill.load(f)
+
+# -----------------------------------------------------------------------------
+pt4 = all_results[0][0][0]['corrs'][0,:]
+pt6 = all_results[0][0][3]['corrs'][0,:]
+pt8 = all_results[0][0][6]['corrs'][0,:]
+
+r1 = (np.abs(pt8)-np.abs(pt4)).max()
+r2 = (np.abs(pt8)-np.abs(pt6)).max()
+
+fig, ax = plt.subplots(figsize=(8/2.54,6/2.54), tight_layout=True)
+walltimes=[]
+times=[]
+for i in range(10):
+    result = all_results[0][1][i]['walltime']
+    N = all_results[0][1][i]['N']
+    walltimes.append(result)
+    times.append(N)
+ax.plot(times, walltimes, '.')
+ax.set_xlabel("Additional time steps")
+ax.set_ylabel("Walltime (s)")
+fig.savefig("./tests/data/plots/nt_corrs_time.pdf")
+
+# -----------------------------------------------------------------------------
+
+plt.show()

tests/performance/analysis/nt_corrs_run.py

@@ -0,0 +1,110 @@
+# Copyright 2024 The TEMPO Collaboration
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Script to run the multi-time correlations performance analysis.
+"""
+
+import sys
+sys.path.insert(0,'.')
+sys.path.append('/home/rmadw/Documents/OQuPy/tests/')
+
+import oqupy
+import dill
+import os
+import numpy as np
+import re
+
+from performance.run_all import run_all
+from performance.nt_corrs import ALL_TESTS
+from performance.nt_corrs import REQUIRED_PTS
+
+os.chdir('/home/rmadw/Documents/OQuPy')
+PT_DIR_PATH = "./tests/data/process_tensors/"
+
+
+def pt_3ls_exists(name, pt_dir = PT_DIR_PATH):
+      """Check the existence of a precomputed process tensor in the process tensor directory."""
+      return os.path.isfile(os.path.join(pt_dir,f"{name}.hdf5"))
+
+def generate_3ls_pt(name, pt_dir = PT_DIR_PATH):
+      """
+      Generate a process tensor from 'name' using PT-TEMPO and write it into the
+      process tensor directory 'pt_dir'.
+
+      """
+      P_1 = np.array([[0., 0., 0.],
+                      [0., 1., 0.],
+                      [0., 0., 0.]], dtype=complex)
+
+      P_2 = np.array([[0., 0., 0.],
+                      [0., 0., 0.],
+                      [0., 0., 1.]], dtype=complex)
+
+
+      res = re.match('3ls_alpha0.1_zeta1.0_T13.09_cutoff3.04exp_tcut50.0_dt0.1_steps80_epsrel([\d]+)', name)
+      if res is None:
+          raise ValueError(f"Invalid PT name '{name}'")
+      p = {
+              'epsrelExp':int(res[1]),
+      }
+
+      file_path = os.path.join(PT_DIR_PATH,f"{name}.hdf5")
+
+      dt = 0.1
+      steps = 80
+      epsrel = 10**(-p['epsrelExp'])
+      end_time = dt*steps
+      dkmax = 500
+
+      correlations = oqupy.PowerLawSD(alpha=0.1,
+                                      zeta=1.0,
+                                      cutoff=3.04,
+                                      cutoff_type='exponential',
+                                      temperature=13.09)
+      bath = oqupy.Bath(P_1 - P_2, correlations)
+      pt_tempo_parameters = oqupy.TempoParameters(dt=dt, dkmax=dkmax, epsrel=epsrel)
+
+      pt = oqupy.pt_tempo_compute(
+            bath=bath,
+            start_time=0.0,
+            end_time=end_time,
+            parameters=pt_tempo_parameters,
+            process_tensor_file=file_path,
+            overwrite=True,
+      )
+
+      pt.close()
+
+
+# -- preparation --------------------------------------------------------------
+
+print("# Prepare computations:")
+
+
+for pt_name in REQUIRED_PTS:
+    if pt_3ls_exists(pt_name):
+        print(f"Process tensor '{pt_name}' already exists. ")
+    else:
+        print(f"Generating process tensor '{pt_name}' ... ")
+        generate_3ls_pt(pt_name)
+
+print("... preparation done.")
+
+# -- computation --------------------------------------------------------------
+
+all_results = run_all(ALL_TESTS)
+with open('./tests/data/performance_results/nt_corrs.pkl', 'wb') as f:
+    dill.dump(all_results, f)
+
+print("... all done.")