plot_theta_sz2.py

import odeg as od
import matplotlib.pyplot as plt
import numpy as np
from time import time
from numba import njit
plt.rcParams['figure.dpi'] = 300
plt.rcParams['savefig.dpi'] = 300

# plots fxc and hxc over theta

if __name__ == '__main__':
    
    cfg = {
        'r0': 1,
        'rs': 2,
        'p_max': 5,
        'pos_charge': 5,
        'theta': np.arange(1e-4, 2.01, 1e-2),
        # 'mu_tol': 1e-6,
        'path': "../results",
        }
    
    r0_grid = np.array([.5, 1.])
    rs_grid = np.array([2., 4., 8.])
    
    fig_shape = (np.size(rs_grid), np.size(r0_grid))
    fig, axes = plt.subplots(*fig_shape, sharex = True, figsize = (7, 6))
    data = np.zeros((*fig_shape, cfg['theta'].size))
    
    for rs_ind, rs in enumerate(rs_grid):
        for r0_ind, r0 in enumerate(r0_grid):
            start = time()
            
            cfg['rs'] = rs
            cfg['r0'] = r0
            
            data[rs_ind, r0_ind, :] = od.fctExp(od.sz2Fct, **cfg)
            
            line, = axes[rs_ind, r0_ind].plot(cfg['theta'], data[rs_ind, r0_ind, :])
            
            print(f"rs:{rs}, r0:{r0}, time: {time()-start}")
            
            
    
    for ax in axes[-1]:
        ax.set_xlabel(r"$\Theta$")
    
    for ax in axes[:, 0]:
        ax.set_ylabel(r"$\langle \hat S_z ^2 \rangle$")
        
    # following 10 lines are copied from https://stackoverflow.com/a/25814386

    for ax, r0 in zip(axes[0], r0_grid):
        ax.annotate(r"$r_0 = $" + str(r0), xy=(0.5, 1), xytext=(0, 20),
                    xycoords='axes fraction', textcoords='offset points',
                    size='large', ha='center', va='baseline')
    
    for ax, rs in zip(axes[:,0], rs_grid):
        ax.annotate(r"$r_{\mathrm{s}} = $" + str(rs), xy=(0, 0.5), xytext=(-60, 0),
                    xycoords='axes fraction', textcoords='offset points',
                    size='large', ha='right', va='center')
            
    
    # fig.legend(loc = 'right')
    fig.tight_layout()
    fig.subplots_adjust(left=0.15, top=0.95)