plot_three_particle_contribution.py

Go to the documentation of this file.

import h5py
from pylab import *
from glob import glob
from scipy.interpolate import griddata
from collections import OrderedDict
import os
from argparse import ArgumentParser

def find_nearest_index(array, value):
    idx = (np.abs(array-value)).argmin()
    return idx

parser = ArgumentParser()
parser.add_argument("--three_body_states", nargs="+")
parser.add_argument("--two_body_states", nargs="+")
parser.add_argument("--id", default="tmp")
args = parser.parse_args()
output_dir = os.path.abspath("tmp")

if args.id != "tmp":
    try:
        from sumatra.projects import load_project
        output_dir = os.path.join(os.path.abspath(load_project().data_store.root), args.id)
    except ImportError:
        pass

if not os.path.exists(output_dir):
    os.makedirs(output_dir)

states_files = args.two_body_states
if len(states_files) == 1:
    states_files = glob(states_files[0])

energy_min = inf
energy_max = -inf

energies = []
r12s = []

for statesFile in states_files:
    f = h5py.File(statesFile, "r")
    atomsMeta = f.get("atomMeta")
    energyOffset = atomsMeta.attrs["energyOffset"]
    if len(atomsMeta) != 2:
        raise Exception("Wrong number of atoms in atomsMeta. Found %d, should be 2." % len(atomsMeta))
    states = f.get("/states")
    for stateName in states:
        atoms = states.get(stateName)
        #r12 = atoms[1]["x"] - atoms[0]["x"]
        #r13 = sqrt((atoms[2]["x"] - atoms[0]["x"])**2 + (atoms[2]["y"] - atoms[0]["y"])**2)
        #angle = arctan2(atoms[2]["y"], atoms[2]["x"])
        r12 = atoms.attrs["r12"]
        energy = atoms.attrs["energy"] - energyOffset
        
        r12s.append(r12)
        energies.append(energy)
        
        energy_min = min(energy_min, energy)
        energy_max = max(energy_max, energy)
    f.close()

# Sort r12 and energies by r12
r12s, energies = [list(x) for x in zip(*sorted(zip(r12s, energies), key=lambda pair: pair[0]))]

two_body_r12s = array(r12s)
two_body_energies = array(energies)

states_files = args.three_body_states
if len(states_files) == 1:
    states_files = glob(states_files[0])

plots = {}

energy_min = inf
energy_max = -inf

energy_differences_min = inf
energy_differences_max = -inf

for statesFile in states_files:
    f = h5py.File(statesFile, "r")
    atomsMeta = f.get("atomMeta")
    states = f.get("/states")
    energyOffset = atomsMeta.attrs["energyOffset"]
    for stateName in states:
        atoms = states.get(stateName)
        #r12 = atoms[1]["x"] - atoms[0]["x"]
        #r13 = sqrt((atoms[2]["x"] - atoms[0]["x"])**2 + (atoms[2]["y"] - atoms[0]["y"])**2)
        #angle = arctan2(atoms[2]["y"], atoms[2]["x"])
        r12 = atoms.attrs["r12"]
        r13 = atoms.attrs["r13"]
        
        angle = atoms.attrs["angle"]
        plot_name = "%.4f" % angle
        energy = atoms.attrs["energy"] - energyOffset
        
        atom2_x = r12
#        
        atom3_x = cos(angle) * r13
        atom3_y = sin(angle) * r13
#        
        r23_vector = array([atom3_x - atom2_x, atom3_y - 0.0, 0.0 - 0.0])
        r23 = linalg.norm(r23_vector)
        
        index12 = find_nearest_index(two_body_r12s, r12)
        index13 = find_nearest_index(two_body_r12s, r13)
        index23 = find_nearest_index(two_body_r12s, r23)  
        
        energy_difference = energy - two_body_energies[index12] - two_body_energies[index13] - two_body_energies[index23]
        
        if r12 < 6.0 and r13 < 6.0:
            energy_difference = 0.0
        
        if not plots.has_key(plot_name):
            plots[plot_name] = {"r12s_r13s": [], "energies": [], "energy_differences": []}
        plots[plot_name]["r12s_r13s"].append([r12, r13])
        plots[plot_name]["energies"].append(energy)
        plots[plot_name]["energy_differences"].append(energy_difference)
        
        r12_min = atomsMeta.attrs["r12Min"]
        r12_max = atomsMeta.attrs["r12Max"]
        r13_min = atomsMeta.attrs["r13Min"]
        r13_max = atomsMeta.attrs["r13Max"]
        angle_min = atomsMeta.attrs["angleMin"]
        angle_max = atomsMeta.attrs["angleMax"]
        
        energy_min = min(energy_min, energy)
        energy_max = max(energy_max, energy)
        
        energy_differences_min = min(energy_differences_min, energy_difference)
        energy_differences_max = max(energy_differences_max, energy_difference)
    f.close()

print "Angle min:", angle_min, "max:", angle_max
print "r13 min:", r13_min, "max:", r13_max


r12s = linspace(r12_min, r12_max, 25)
r13s = linspace(r13_min, r13_max, 25)
grid_r12s, grid_r13s = meshgrid(r12s, r13s)
n_plots = len(plots)
n_plots_per_dim = int(sqrt(n_plots) + 1)
plot_counter = 1
plots = OrderedDict(sorted(plots.items()))

#vmin = energy_min
#vmax = energy_min + (energy_max - energy_min) * 0.5
vmin = energy_differences_min
vmax = energy_differences_max
#vmin = 0
#vmax = 0.5
#vmax = energy_differences_min + (energy_differences_max - energy_differences_min) * 0.1
print "vmin,vmax: ", vmin, vmax

fig = figure(figsize=(10,10))
for plot_name in plots:
    
    ax = fig.add_subplot(n_plots_per_dim, n_plots_per_dim, plot_counter)
    ax.set_aspect("equal")
    title(plot_name)
    
    values = plots[plot_name]
    
    grid_energies = griddata(array(values["r12s_r13s"]), array(values["energies"]), (grid_r12s, grid_r13s), method="nearest")
    grid_energy_differences = griddata(array(values["r12s_r13s"]), array(values["energy_differences"]), (grid_r12s, grid_r13s), method="nearest")
    #grid_energies = griddata(array(values["r12s_r13s"]), array(values["energy_differences"]), (grid_r12s, grid_r13s), method="nearest")
    
#    img = imshow(grid_energies, origin="lower", vmin=vmin, vmax=vmax,
#                 extent=[r12_min, r12_max, r13_min, r13_max], 
#                         interpolation="nearest", 
#                         aspect=(r12_max - r12_min) / (r13_max - r13_min))
    
    img = contour(r12s, r13s, grid_energy_differences, 50, vmin=vmin, vmax=vmax)
    colorbar()
    
    plot_counter += 1
    
fig.tight_layout()
#savefig(output_file + ".pdf")
#savefig(output_file + ".png")