plot_theta_f.py

import odeg as od
import matplotlib.pyplot as plt
import numpy as np
from time import time

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-2, 2.01, 1e-2),
        'path': "../results",
        'mu_rel_target': 1 - 1e-10,
        }
    
    r0_grid = np.array([.5, 1.])
    rs_grid = np.array([2., 4., 8.])
    obs_count = 2
    
    fig_shape = (np.size(rs_grid), np.size(r0_grid))
    fig, axes = plt.subplots(*fig_shape, sharex = True, figsize = (7, 6))
    data = np.zeros((obs_count, *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[0, rs_ind, r0_ind, :] = rs * od.freeEnergy(**cfg) / cfg['pos_charge']
            data[1, rs_ind, r0_ind, :] = rs * od.fctExp(od.energyFct, **cfg) / cfg['pos_charge']
            
            line, = axes[rs_ind, r0_ind].plot(cfg['theta'],
                                              data[0, rs_ind, r0_ind, :])
            if (r0_ind, rs_ind) == (0, 0):
                line.set_label(r"$r_{\mathrm{s}}F /\cal{N}$")
            
            line, = axes[rs_ind, r0_ind].plot(cfg['theta'],
                                             data[1, rs_ind, r0_ind, :])
            if (r0_ind, rs_ind) == (0, 0):
                line.set_label(r"$r_{\mathrm{s}}\langle \hat H \rangle /\cal{N}$")
            
            print(f"rs:{rs}, r0:{r0}, time: {time()-start}")
            
            
    
    for ax in axes[-1]:
        ax.set_xlabel(r"$\Theta$")
        
    # hard-coded !!!!
    # axes[1,1].yaxis.set_ticks(np.arange(-1.09, -0.97, 0.03))
    
    axes[fig_shape[0]//2, 0].set_ylabel(r"(free) energy in units of $E_{\mathrm{h}}$")
        
    # 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=(-70, 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)