From 686f74a9296bb79fc0926efabc763f5cb995614b Mon Sep 17 00:00:00 2001 From: rasmusvt Date: Sun, 18 Sep 2022 15:50:54 +0200 Subject: [PATCH] Add initial script files for DFT processing --- nafuma/dft/__init__.py | 1 + nafuma/dft/electrons.py | 1109 +++++++++++++++++++++++++++++++++++++++ nafuma/dft/io.py | 545 +++++++++++++++++++ nafuma/dft/structure.py | 906 ++++++++++++++++++++++++++++++++ 4 files changed, 2561 insertions(+) create mode 100644 nafuma/dft/__init__.py create mode 100644 nafuma/dft/electrons.py create mode 100644 nafuma/dft/io.py create mode 100644 nafuma/dft/structure.py diff --git a/nafuma/dft/__init__.py b/nafuma/dft/__init__.py new file mode 100644 index 0000000..db87b5d --- /dev/null +++ b/nafuma/dft/__init__.py @@ -0,0 +1 @@ +from . import electrons, io, structure \ No newline at end of file diff --git a/nafuma/dft/electrons.py b/nafuma/dft/electrons.py new file mode 100644 index 0000000..d96d2ad --- /dev/null +++ b/nafuma/dft/electrons.py @@ -0,0 +1,1109 @@ +import re +import pandas as pd +import numpy as np +import matplotlib.pyplot as plt +import os +import linecache +import nafuma.dft as dft +import nafuma.auxillary as aux +import nafuma.plotting as btp + +from matplotlib.ticker import (MultipleLocator, FormatStrFormatter,AutoMinorLocator) +import importlib +import mpl_toolkits.axisartist as axisartist +from cycler import cycler +import itertools +import matplotlib.patches as mpatches + + + + +def count_electrons(pdos, orbital, interval=None, r=None, scale=None): + ''' Counts electrons the specified oribtals from a projected density of states DataFrame. Interval can be specified, as well as a scaling factor and whether the number should be rounded. + Inputs: + dos: either an individual DOS (as read from read_pdos()), or a list of DOSes. If a single DataFrame is passed, it will be appended to a list + orbital: list of which orbitals to count the electrons from + interval: a list specifying where to start counting from (lower limit) to where to stop counting (upper limit) in eV + r: Number of decimals points the number should be rounded to + scale: A scaling factor to scale the number of electrons to a desired size, e.g. if you have a set containing two atoms per unit cell and you want to know how many electrons per atom there is + + Output: + nelec: The total number of electrons given your choices + nelec_dos: A list where each element is the total number of electrons per DOS passed (e.g. you pass three PDOS from three individual atoms, then you will get total electron count per atom) + nelec_orbitals: A list of lists, where each list contains the number of electrons per orbital specified (e.g. you pass three PDOS from three individual atoms, you will get three lists each containing electron count per orbital specified) + ''' + + + + if not type(orbital) == list: + orbital = [orbital] + + if not type(pdos) == list: + pdos = [pdos] + + + nelec = 0 + nelec_per_dos = [] + nelec_per_orbital = [] + + for d in pdos: + + energy = d["Energy"] + + nelec_orbitals = [] + + for o in orbital: + orbital_dos = d[o] + dE = (energy.max()-energy.min()) / len(energy) + + + if not interval: + interval = [energy.min(), energy.max()] + + emin, emax = interval[0], interval[1] + + + nelec_orbital = 0 + for od, e in zip(orbital_dos, energy): + if e > emin and e < emax: + nelec_orbital += np.abs(od)*dE + #print(nelec_orbital) + + nelec += nelec_orbital + nelec_orbitals.append(nelec_orbital) + + + # Scale the values if specified + if scale: + + for ind, nelec_orbital in enumerate(nelec_orbitals): + nelec_orbitals[ind] = nelec_orbital * scale + + + # If rounding is specified, does so to the electron count per DOS and the electron count per orbital + if r: + # First sums the electron count per orbital, and then round this number + nelec_dos = np.round(sum(nelec_orbitals), r) + + # Then each individual orbital electron count + for ind, nelec_orbital in enumerate(nelec_orbitals): + nelec_orbitals[ind] = np.round(nelec_orbital, r) + + # If no rounding is specified, just adds the electron count per orbital together + else: + nelec_dos = sum(nelec_orbitals) + + # Appends the total electron count for this DOS to the list of all individual DOS electron count + nelec_per_dos.append(nelec_dos) + + # Appends the list of orbital electron counts to the list of all the individual DOS orbital electron count (phew...) + nelec_per_orbital.append(nelec_orbitals) + + + # The total electron count is then scaled in the end. At this point the other values will have been scaled already + if scale: + nelec = nelec * scale + + # And lastly round if this is specified. Again, the electron counts in the lists are already rounded so they don't have to be rounded again + if r: + nelec = np.round(nelec, r) + + return nelec, [nelec_per_dos, nelec_per_orbital] + + + +def integrate_coop(coopcar, interval=None, r=None, scale=None, interactions=None, kind='individual', up=True, down=True, collapse=False): + ''' As of now does not support not passing in interactions. Very much copy and paste from the plotting function - not every choice here might make sense for integration of COOP''' + + coopcar, coop_interactions = dft.io.read_coop(coopcar, collapse=collapse) + + # If interactions has been specified + if interactions: + + # Make interactions into a list of lists for correct looping below + if type(interactions[0]) != list: + interactions_list = [interactions] + else: + interactions_list = interactions + + for ind, interactions in enumerate(interactions_list): + + # Determine which columns to integrate if collapse is enabled + if collapse: + to_integrate = [2*(i-1)+3 for i in interactions] + + + # Make mean column for integration if mean mode is enabeld (is this sensible to include?) + if kind == 'avg' or kind == 'average' or kind == 'mean': + coopcar["mean"] = coopcar.iloc[:, to_integrate].mean(axis=1) + to_integrate = [coopcar.columns.get_loc('mean')] + + # Determine which columns to integrate if collapse is disabled and both up and down should be plotted + elif up and down: + to_integrate_up = [2*(i-1)+3 for i in interactions] + to_integrate_down = [2*(i-1)+5 +2*len(coop_interactions) for i in interactions] + to_integrate = to_integrate_up + to_integrate_down + + if kind == 'avg' or kind == 'average' or kind == 'mean': + coopcar["mean_up"] = coopcar.iloc[:, to_integrate_up].mean(axis=1) + coopcar["mean_down"] = coopcar.iloc[:, to_integrate_down].mean(axis=1) + to_integrate = [coopcar.columns.get_loc('mean_up'), coopcar.columns.get_loc('mean_down')] + + # Determine which columns to plot if collapse is disabled and only up should be plotted + elif up: + to_integrate = [2*(i-1)+3 for i in interactions] + + if kind == 'avg' or kind == 'average' or kind == 'mean': + coopcar["mean_up"] = coopcar.iloc[:, to_integrate].mean(axis=1) + to_integrate = [coopcar.columns.get_loc('mean_up')] + + + # Determine which columns to plot if collapse is disabled and only down should be plotted + elif down: + to_integrate = [2*(i-1)+5 +2*len(coop_interactions) for i in interactions] + + if kind == 'avg' or kind == 'average' or kind == 'mean': + coopcar["mean_down"] = coopcar.iloc[:, to_integrate].mean(axis=1) + to_integrate = [coopcar.columns.get_loc('mean_down')] + + + + bonding = 0 + antibonding = 0 + bonding_interactions = [] + antibonding_interactions = [] + difference_interactions = [] + percentage_bonding_interactions = [] + + + for integrate in to_integrate: + + bonding_interaction = 0 + antibonding_interaction = 0 + + coop = coopcar.iloc[:, integrate] + + energy = coopcar["Energy"] + dE = (energy.max()-energy.min()) / len(energy) + + # Sets interval to everything below the Fermi-level by default if not specified in function call + if not interval: + interval = [energy.min(), 0] + + emin, emax = interval[0], interval[1] + + + for c, e in zip(coop, energy): + if e > emin and e < emax: + if c > 0: + bonding_interaction += c*dE + elif c < 0: + antibonding_interaction += np.abs(c)*dE + + + bonding += bonding_interaction + antibonding += antibonding_interaction + + difference_interaction = bonding_interaction - antibonding_interaction + percentage_bonding_interaction = bonding_interaction / (bonding_interaction + antibonding_interaction) * 100 + + if scale: + bonding_interaction = bonding_interaction * scale + antibonding_interaction = antibonding_interaction * scale + difference_interaction = difference_interaction * scale + + if r: + bonding_interaction = np.round(bonding_interaction, r) + antibonding_interaction = np.round(antibonding_interaction, r) + difference_interaction = np.round(difference_interaction, r) + percentage_bonding_interaction = np.round(percentage_bonding_interaction, r) + + bonding_interactions.append(bonding_interaction) + antibonding_interactions.append(antibonding_interaction) + difference_interactions.append(difference_interaction) + percentage_bonding_interactions.append(percentage_bonding_interaction) + + difference = bonding - antibonding + percentage_bonding = (bonding/(bonding+antibonding)) * 100 + + if scale: + bonding = bonding * scale + antibonding = antibonding * scale + difference = difference * scale + + if r: + bonding = np.round(bonding, r) + antibonding = np.round(antibonding, r) + difference = np.round(difference, r) + percentage_bonding = np.round(percentage_bonding, r) + + return [bonding, antibonding, difference, percentage_bonding], [bonding_interactions, antibonding_interactions, difference_interactions, percentage_bonding_interactions] + + + + +def plot_partial_dos(data: dict, options={}): + + + required_options = ['atoms', 'orbitals', 'up', 'down', 'sum_atoms', 'collapse_spin', 'sum_orbitals', 'palettes', 'colours', 'fig', 'ax'] + + default_options = { + 'atoms': None, + 'orbitals': None, + 'up': True, + 'down': True, + 'sum_atoms': True, + 'collapse_spin': False, + 'sum_orbitals': False, + 'palettes': [('qualitative', 'Dark2_8')], + 'colours': None, + 'fig': None, + 'ax': None + + + } + + options = update_options(options=options, required_options=required_options, default_options=default_options) + + if not options['ax'] and not options['fig']: + fig, ax = btp.prepare_plot(options) + else: + fig, ax = options['fig'], options['ax'] + + species, *_ = dft.io.get_atoms(data['poscar']) + + pdos, options['dos_info'] = dft.io.read_pdos(data=data, options=options) # Extract projected DOS from DOSCAR, decomposed on individual atoms and orbitals Should yield list of N DataFrames where N is number of atoms in POSCAR + + + if not options['orbitals']: + options['orbitals'] = ['s', 'p1', 'p2', 'p3', 'd1', 'd2', 'd3', 'd4', 'd5'] if not options['sum_orbitals'] else ['s', 'p', 'd'] + + if not options['colours']: + colour_cycle = generate_colours(options=options) +# + #colours = [] + #for orbital in options['orbitals']: + # colours.append(next(colour_cycle)) +# + # else: + # colours = options['colours'] + + elif not isinstance(options['colours'], list): + new_colours = [] + for atom in options['atoms']: + new_colours.append([options['colours']]) + + options['colours'] = new_colours + + + print(options['colours']) + + + if not isinstance(options['orbitals'][0], list): + new_orbitals = [] + for atom in options['atoms']: + new_orbitals.append([options['orbitals']]) + + options['orbitals'] = new_orbitals + + + if options['atoms']: + for i, atom in enumerate(options['atoms']): + + if options['sum_atoms']: + for ind, specie in enumerate(species): + if specie == atom: + atom_index = ind + else: + atom_index = atom-1 + + + for j, orbital in enumerate(options['orbitals'][i]): + + colour = options['colours'][i][j] + + if options['dos_info']['spin_polarised']: + if options['up']: + pdos[atom_index].plot(x='Energy', y=orbital+'_u', ax=ax, c=colour) + + if options['down']: + pdos[atom_index].plot(x='Energy', y=orbital+'_d', ax=ax, c=colour) + else: + pdos[atom_index].plot(x='Energy', y=orbital, ax=ax, c=colour) + + + + btp.adjust_plot(fig=fig, ax=ax, options=options) + + return [pdos, ax, fig] + + + +def get_pdos_indices(poscar, atoms): + + species, atom_num, atoms_dict = dft.io.get_atoms(poscar) + + + + +def get_pdos(doscar='DOSCAR', nedos=301, headerlines=6, spin=True, adjust=True, manual_adjust=None): + + lines = dft.io.open_doscar(doscar) + + number_of_atoms = dft.io.get_number_of_atoms(doscar) + + if adjust: + e_fermi = dft.io.get_fermi_level(doscar) if not manual_adjust else manual_adjust + else: + e_fermi = 0 + + pdos = [] + + columns_non_spin = ["Energy", "s", "p_y", "p_z", "p_x", "d_xy", "d_yz", "d_z2-r2", "d_xz", "d_x2-y2"] + columns_spin = ["Energy", "s_up", "s_down", "p_y_up", "p_y_down", "p_z_up", "p_z_down", "p_x_up", "p_x_down", "d_xy_up", "d_xy_down", "d_yz_up", "d_yz_down", + "d_z2-r2_up", "d_z2-r2_down", "d_xz_up", "d_xz_down", "d_x2-y2_up", "d_x2-y2_down"] + + up = ['s_up', "p_y_up", "p_z_up", "p_x_up", "d_xy_up", "d_yz_up", "d_z2-r2_up", "d_xz_up", "d_x2-y2_up"] + down = ['s_down', "p_y_down", "p_z_down", "p_x_down", "d_xy_down", "d_yz_down", "d_z2-r2_down", "d_xz_down", "d_x2-y2_down"] + total = ["s", "p_y", "p_z", "p_x", "d_xy", "d_yz", "d_z2-r2", "d_xz", "d_x2-y2"] + + for i in range(1,number_of_atoms+1): + atom_dos = [] + + for j in range(headerlines+(nedos*i)+i,nedos+headerlines+(nedos*i)+i): + line = lines[j].strip() + values = line.split() + + for ind, value in enumerate(values): + values[ind] = float(value) + + values[0] = values[0] - e_fermi + atom_dos.append(values) + + + atom_df = pd.DataFrame(data=atom_dos, columns=columns_non_spin) if spin==False else pd.DataFrame(data=atom_dos, columns=columns_spin) + + if spin==True: + atom_df[["s_down"]] = -atom_df[["s_down"]] + atom_df[["p_y_down"]] = -atom_df[["p_y_down"]] + atom_df[["p_z_down"]] = -atom_df[["p_z_down"]] + atom_df[["p_x_down"]] = -atom_df[["p_x_down"]] + atom_df[["d_xy_down"]] = -atom_df[["d_xy_down"]] + atom_df[["d_yz_down"]] = -atom_df[["d_yz_down"]] + atom_df[["d_z2-r2_down"]] = -atom_df[["d_z2-r2_down"]] + atom_df[["d_xz_down"]] = -atom_df[["d_xz_down"]] + atom_df[["d_x2-y2_down"]] = -atom_df[["d_x2-y2_down"]] + + atom_df = atom_df.assign(total_up = atom_df[up].sum(axis=1)) + atom_df = atom_df.assign(total_down = atom_df[down].sum(axis=1)) + + elif spin==False: + atom_df = atom_df.assign(total = atom_df[total].sum(axis=1)) + + pdos.append(atom_df) + + return pdos + + + +def prepare_plot(options={}): + + rc_params = options['rc_params'] + format_params = options['format_params'] + + required_options = ['single_column_width', 'double_column_width', 'column_type', 'width_ratio', 'aspect_ratio', 'compress_width', 'compress_height', 'upscaling_factor', 'dpi'] + + default_options = { + 'single_column_width': 8.3, + 'double_column_width': 17.1, + 'column_type': 'single', + 'width_ratio': '1:1', + 'aspect_ratio': '1:1', + 'compress_width': 1, + 'compress_height': 1, + 'upscaling_factor': 1.0, + 'dpi': 600, + } + + options = update_options(format_params, required_options, default_options) + + + # Reset run commands + plt.rcdefaults() + + # Update run commands if any is passed (will pass an empty dictionary if not passed) + update_rc_params(rc_params) + + width = determine_width(options) + height = determine_height(options, width) + width, height = scale_figure(options=options, width=width, height=height) + + fig, ax = plt.subplots(figsize=(width, height), dpi=options['dpi']) + + return fig, ax + + +def prepare_dos_plot(width=None, height=None, square=True, dpi=None, colour_cycle=('qualitative', 'Dark2_8'), energyunit='eV', dosunit='arb. u.', scale=1, pdos=None): + + if not pdos: + linewidth = 3*scale + else: + linewidth = 3 + + axeswidth = 3*scale + + plt.rc('lines', linewidth=linewidth) + plt.rc('axes', linewidth=axeswidth) + + if square: + if not width: + width = 20 + + if not height: + height = width + + + fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(width, height), facecolor='w', dpi=dpi) + + + return fig, ax + +def prettify_dos_plot(fig, ax, options): + + + required_options = ['plot_kind', 'flip_xy', 'hide_x_labels', 'hide_y_labels', 'xlabel', 'ylabel', 'xunit', 'yunit', 'xlim', 'ylim', 'x_tick_locators', 'y_tick_locators', 'y_tick_format', 'x_tick_format', 'hide_x_ticks', 'hide_y_ticks', 'hide_x_ticklabels', 'hide_y_ticklabels', + 'colours', 'palettes', 'title', 'legend', 'labels', 'label_colours', 'legend_position', 'legend_ncol', 'subplots_adjust', 'text'] + + default_options = { + 'plot_kind': 'PDOS', # DOS/PDOS/COOP/COHP + 'flip_xy': False, + 'hide_x_labels': False, # Whether x labels should be hidden + 'hide_x_ticklabels': False, + 'hide_x_ticks': False, + 'hide_y_labels': False, # whether y labels should be hidden + 'hide_y_ticklabels': False, + 'hide_y_ticks': False, + 'xlabel': 'Energy', + 'ylabel': 'DOS', + 'xunit': r'eV', # The unit of the x-values in the curve plot + 'yunit': r'a.u.', # The unit of the y-values in the curve and bar plots + 'xlim': None, + 'ylim': None, + 'x_tick_locators': [1, .5], # Major and minor tick locators + 'y_tick_locators': [1, .5], + 'y_tick_format': None, + 'x_tick_format': None, + 'colours': None, + 'palettes': [('qualitative', 'Dark2_8'), ('qualitative', 'Paired_12')], + 'title': None, + 'legend': True, + 'labels': None, + 'label_colours': None, + 'legend_position': ['upper center', (0.20, 0.90)], # the position of the legend passed as arguments to loc and bbox_to_anchor respectively + 'legend_ncol': 1, + 'subplots_adjust': [0.1, 0.1, 0.9, 0.9], + 'text': None + } + + + if 'plot_kind' in options.keys(): + if 'ylabel' not in options.keys(): + if options['plot_kind'] == 'DOS': + options['ylabel'] = 'DOS' + elif options['plot_kind'] == 'PDOS': + options['ylabel'] = 'PDOS' + elif options['plot_kind'] == 'COOP': + options['ylabel'] = 'COOP' + elif options['plot_kind'] == 'COHP': + options['ylabel'] = 'COHP' + + + + options = update_options(options=options, required_options=required_options, default_options=default_options) + + + if options['flip_xy']: + + # Switch all the x- and y-specific values + options = swap_values(dict=options, key1='xlim', key2='ylim') + options = swap_values(dict=options, key1='xunit', key2='yunit') + options = swap_values(dict=options, key1='xlabel', key2='ylabel') + options = swap_values(dict=options, key1='x_tick_locators', key2='y_tick_locators') + options = swap_values(dict=options, key1='hide_x_labels', key2='hide_y_labels') + + # Set labels on x- and y-axes + if not options['hide_y_labels']: + ax.set_ylabel(f'{options["ylabel"]} [{options["yunit"]}]') + else: + ax.set_ylabel('') + + + + if not options['hide_x_labels']: + ax.set_xlabel(f'{options["xlabel"]} [{options["xunit"]}]') + else: + ax.set_xlabel('') + + + # Hide x- and y- ticklabels + if options['hide_y_ticklabels']: + ax.tick_params(axis='y', direction='in', which='both', labelleft=False, labelright=False) + if options['hide_x_ticklabels']: + ax.tick_params(axis='x', direction='in', which='both', labelbottom=False, labeltop=False) + + + # Hide x- and y-ticks: + if options['hide_y_ticks']: + ax.tick_params(axis='y', direction='in', which='both', left=False, right=False) + if options['hide_x_ticks']: + ax.tick_params(axis='x', direction='in', which='both', bottom=False, top=False) + + + + # Set multiple locators + ax.yaxis.set_major_locator(MultipleLocator(options['y_tick_locators'][0])) + ax.yaxis.set_minor_locator(MultipleLocator(options['y_tick_locators'][1])) + + ax.xaxis.set_major_locator(MultipleLocator(options['x_tick_locators'][0])) + ax.xaxis.set_minor_locator(MultipleLocator(options['x_tick_locators'][1])) + + # Change format of axis tick labels if specified: + + + + # Set title + if options['title']: + ax.set_title(options['title']) + + + if options['y_tick_format']: + ax.yaxis.set_major_formatter(FormatStrFormatter(options['y_tick_format'])) + if options['x_tick_format']: + ax.xaxis.set_major_formatter(FormatStrFormatter(options['x_tick_format'])) + + + # Create legend + + if ax.get_legend(): + ax.get_legend().remove() + + + if options['legend'] and options['labels']: + + + # Generate colours + if not options['colours']: + colour_cycle = generate_colours(palettes=options['palettes']) + + colours = [] + for label in options['labels']: + colours.append(next(colour_cycle)) + + + else: + colours = options['colours'] + + if options['label_colours']: + colours = options['label_colours'] + + # Create legend + patches = [] + for i, label in enumerate(options['labels']): + patches.append(mpatches.Patch(color=colours[i], label=label)) + + print(options['legend_ncol']) + + ax.legend(handles=patches, loc=options['legend_position'][0], bbox_to_anchor=options['legend_position'][1], frameon=False, ncol=options['legend_ncol']) + + + + # Adjust where the axes start within the figure. Default value is 10% in from the left and bottom edges. Used to make room for the plot within the figure size (to avoid using bbox_inches='tight' in the savefig-command, as this screws with plot dimensions) + plt.subplots_adjust(left=options['subplots_adjust'][0], bottom=options['subplots_adjust'][1], right=options['subplots_adjust'][2], top=options['subplots_adjust'][3]) + + + # If limits for x- and y-axes is passed, sets these. + if options['xlim']: + ax.set_xlim(options['xlim']) + + if options['ylim']: + ax.set_ylim(options['ylim']) + + + # Add custom text + if options['text']: + plt.text(x=options['text'][1][0], y=options['text'][1][1], s=options['text'][0]) + + + + if options['e_fermi']: + if options['flip_xy']: + ax.axhline(0, c='black', ls='dashed') + else: + ax.axvline(0, c='black', ls='dashed') + + if options['plot_kind'] == 'DOS' or options['plot_kind'] == 'PDOS': + if options['dos_info']['spin_polarised']: + if options['flip_xy']: + ax.axvline(0, c='black') + else: + ax.axhline(0, c='black') + elif options['plot_kind'] == 'COOP' or options['plot_kind'] == 'COHP': + if options['flip_xy']: + ax.axvline(0, c='black') + else: + ax.axhline(0, c='black') + + return fig, ax + + + +def plot_coop(plot_data, options): + ''' interactions = list with number of interaction (index + 1 of interactions list from read_coop)''' + + + required_options = ['plot_kind', 'mode', 'up', 'down', 'collapse', 'interactions', 'flip_xy', 'fill', 'colours', 'palettes'] + + default_options = { + 'plot_kind': 'COOP', + 'mode': 'individual', + 'fill': False, + 'up': True, + 'down': True, + 'collapse': False, + 'interactions': None, + 'palettes': [('qualitative', 'Dark2_8')], + 'colours': None, + 'flip_xy': False + + } + + options = update_options(options=options, required_options=required_options, default_options=default_options) + + + fig, ax = prepare_plot(options=options) + + coopcar, coop_interactions = dft.io.read_coop(plot_data=plot_data, options=options) + + + + if not options['colours']: + colour_cycle = generate_colours(palettes=options['palettes']) + + colours = [] + for interaction in range(len(coop_interactions)): + colours.append(next(colour_cycle)) + + else: + colours = options['colours'] + + + # If interactions has been specified + if options['interactions']: + + # Make interactions into a list of lists for correct looping below + if type(options['interactions'][0]) != list: + interactions_list = [options['interactions']] + else: + interactions_list = options['interactions'] + + for ind, interactions in enumerate(interactions_list): + + # Determine which columns to plot if collapse is enabled + if options['collapse']: + to_plot = [2*(i-1)+3 for i in interactions] + + # Make sum column for plotting if sum mode is enabled + if options['mode'] == 'sum': + coopcar["sum"] = coopcar.iloc[:, to_plot].sum(axis=1) + to_plot = ['sum'] + + + # Make mean column for plotting if mean mode is enabeld + elif options['mode'] == 'avg' or options['mode'] == 'average' or options['mode'] == 'mean': + coopcar["mean"] = coopcar.iloc[:, to_plot].mean(axis=1) + to_plot = ['mean'] + + # Determine which columns to plot if collapse is disabled and both up and down should be plotted + elif options['up'] and options['down']: + to_plot_up = [2*(i-1)+3 for i in interactions] + to_plot_down = [2*(i-1)+5 +2*len(coop_interactions) for i in interactions] + to_plot = to_plot_up + to_plot_down + + if options['mode'] == 'sum': + coopcar["sum_up"] = coopcar.iloc[:, to_plot_up].sum(axis=1) + coopcar["sum_down"] = coopcar.iloc[:, to_plot_down].sum(axis=1) + to_plot = ['sum_up', 'sum_down'] + + elif options['mode'] == 'avg' or options['mode'] == 'average' or options['mode'] == 'mean': + coopcar["mean_up"] = coopcar.iloc[:, to_plot_up].mean(axis=1) + coopcar["mean_down"] = coopcar.iloc[:, to_plot_down].mean(axis=1) + to_plot = ['mean_up', 'mean_down'] + + # Determine which columns to plot if collapse is disabled and only up should be plotted + elif options['up']: + to_plot = [2*(i-1)+3 for i in interactions] + + if options['mode'] == 'sum': + coopcar["sum_up"] = coopcar.iloc[:, to_plot].sum(axis=1) + to_plot = ['sum_up'] + + elif options['mode'] == 'avg' or options['mode'] == 'average' or options['mode'] == 'mean': + coopcar["mean_up"] = coopcar.iloc[:, to_plot].mean(axis=1) + to_plot = ['mean_up'] + + + # Determine which columns to plot if collapse is disabled and only down should be plotted + elif options['down']: + to_plot = [2*(i-1)+5 +2*len(coop_interactions) for i in interactions] + + if options['mode'] == 'sum': + coopcar["sum_down"] = coopcar.iloc[:, to_plot].sum(axis=1) + to_plot = ['sum_down'] + + elif options['mode'] == 'avg' or options['mode'] == 'average' or options['mode'] == 'mean': + coopcar["mean_down"] = coopcar.iloc[:, to_plot].mean(axis=1) + to_plot = ['mean_down'] + + + # Plot all columns as decided above + for j, column in enumerate(to_plot): + if options['fill']: + ax.fill_between(coopcar["Energy"], coopcar[column], 0, where=coopcar[column]>0, color=colours[ind]) + ax.fill_between(coopcar["Energy"], coopcar[column], 0, where=coopcar[column]<0, color=colours[ind+1]) + + else: + if options['mode'] == "individual": + colour = colours[j] + else: + colour = colours[ind] + + + if options['flip_xy']: + coopcar.plot(y='Energy', x=column, ax=ax, color=colour) + else: + coopcar.plot(x='Energy', y=column, ax=ax, color=colour) + + prettify_dos_plot(fig=fig, ax=ax, options=options) + + return coopcar + + + +def prettify_coop_plot(fig, ax, energyunit='eV', dosunit='arb. u.', xlim=None, ylim=None, title=None, hide_ylabels=False, hide_xlabels=False, hide_yvals=False, hide_xvals=False, flip_xy=False, pad_bottom=None, scale=1, colours=None, atoms=None, pdos=False, width=None, height=None, e_fermi=False, adjust=False, legend=False, labels=None, label_colours=None, xpad=0, ypad=0): + + # Set sizes of ticks, labes etc. + ticksize = 30*scale + labelsize = 30*scale + legendsize = 30*scale + titlesize = 30*scale + + linewidth = 3*scale + axeswidth = 3*scale + majorticklength = 20*scale + minorticklength = 10*scale + + plt.xticks(fontsize=ticksize) + plt.yticks(fontsize=ticksize) + + if flip_xy: + + # Set labels on x- and y-axes + if not hide_ylabels: + if ypad: + ax.set_ylabel('Energy [{}]'.format(energyunit), size=labelsize, labelpad=ypad) + else: + ax.set_ylabel('Energy [{}]'.format(energyunit), size=labelsize) + + if pdos: + if xpad: + ax.set_xlabel('COOP [{}]'.format(dosunit), size=labelsize, labelpad=xpad) + else: + ax.set_xlabel('COOP [{}]'.format(dosunit), size=labelsize) + + else: + if width >= 10: + if xpad: + ax.set_xlabel('COOP [{}]'.format(dosunit), size=labelsize, labelpad=xpad) + else: + ax.set_xlabel('COOP [{}]'.format(dosunit), size=labelsize) + + else: + if xpad: + ax.set_xlabel('COOP [{}]'.format(dosunit), size=labelsize, labelpad=xpad) + else: + ax.set_xlabel('COOP [{}]'.format(dosunit), size=labelsize) + + ax.tick_params(axis='y', direction='in', which='major', right=True, length=majorticklength, width=linewidth) + ax.tick_params(axis='y', direction='in', which='minor', right=True, length=minorticklength, width=linewidth) + + if hide_yvals: + ax.tick_params(axis='y', labelleft=False) + + ax.tick_params(axis='x', direction='in', which='major', bottom=False, labelbottom=False) + + ax.yaxis.set_major_locator(MultipleLocator(1)) + ax.yaxis.set_minor_locator(MultipleLocator(.5)) + + + + else: + # Set labels on x- and y-axes + if adjust: + if xpad: + ax.set_xlabel('E - E$_F$ [{}]'.format(energyunit), size=labelsize, labelpad=xpad) + else: + ax.set_xlabel('E - E$_F$ [{}]'.format(energyunit), size=labelsize) + + + else: + if xpad: + ax.set_xlabel('Energy [{}]'.format(energyunit), size=labelsize, labelpad=xpad) + else: + ax.set_xlabel('Energy [{}]'.format(energyunit), size=labelsize) + + + if height < 10: + if ypad: + ax.set_ylabel('COOP [{}]'.format(dosunit), size=labelsize, labelpad=ypad) + else: + ax.set_ylabel('COOP [{}]'.format(dosunit), size=labelsize) + + else: + if ypad: + ax.set_ylabel('Crystal orbital overlap population [{}]'.format(dosunit), size=labelsize, labelpad=ypad) + else: + ax.set_ylabel('Crystal orbital overlap population [{}]'.format(dosunit), size=labelsize) + + + ax.tick_params(axis='x', direction='in', which='major', bottom=True, top=True, length=majorticklength, width=linewidth) + ax.tick_params(axis='x', direction='in', which='minor', bottom=True, top=True, length=minorticklength, width=linewidth) + + + ax.tick_params(axis='y', which='major', direction='in', right=True, left=True, labelleft=True, length=majorticklength, width=linewidth) + ax.tick_params(axis='y', which='minor', direction='in', right=True, left=True, length=minorticklength, width=linewidth) + + if hide_ylabels: + ax.set_ylabel('') + if hide_xlabels: + ax.set_xlabel('') + if hide_yvals: + ax.tick_params(axis='y', which='both', labelleft=False) + if hide_xvals: + ax.tick_params(axis='x', which='both', labelbottom=False) + + + if ylim: + yspan = ylim[1] - ylim[0] + yloc = np.round(yspan / 4, 2) + + ax.yaxis.set_major_locator(MultipleLocator(yloc)) + ax.yaxis.set_minor_locator(MultipleLocator(yloc/2)) + + + + ax.xaxis.set_major_locator(MultipleLocator(1)) + ax.xaxis.set_minor_locator(MultipleLocator(.5)) + + + + plt.xlim(xlim) + plt.ylim(ylim) + + + if title: + ax.set_title(title, size=40) + + + + if legend: + patches = [] + + if label_colours: + colours=label_colours + + for ind, label in enumerate(labels): + patches.append(mpatches.Patch(color=colours[ind], label=label)) + + fig.legend(handles=patches, loc='upper right', ncol=len(labels), bbox_to_anchor=(0.8, 0.45), fontsize=legendsize/1.25, frameon=False) + + #bbox_to_anchor=(1.20, 0.91) + + + + if pad_bottom is not None: + bigax = fig.add_subplot(111) + bigax.set_facecolor([1,1,1,0]) + bigax.spines['top'].set_visible(False) + bigax.spines['bottom'].set_visible(True) + bigax.spines['left'].set_visible(False) + bigax.spines['right'].set_visible(False) + bigax.tick_params(labelcolor='w', color='w', direction='in', top=False, bottom=True, left=False, right=False, labelleft=False, pad=pad_bottom) + + + + if xpad: + ax.tick_params(axis='x', pad=xpad) + + if ypad: + ax.tick_params(axis='y', pad=ypad) + + if e_fermi: + if flip_xy: + plt.axhline(0, lw=linewidth, c='black', ls='dashed') + else: + plt.axvline(0, lw=linewidth, c='black', ls='dashed') + + + plt.axhline(0, lw=linewidth, c='black') + + return fig, ax + + + + +def get_unique_atoms(interactions): + ''' Get all the unique atoms involved in the interactions from the COOP-calculation + + Input: + interactions: list of interactions that comes as output from read_coop() + + Outut: + unique_atoms: list of unique atoms in the interactions list''' + + unique_atoms = [] + + for interaction in interactions: + + atoms = interaction.split('->') + + for atom in atoms: + if atom not in unique_atoms: + unique_atoms.append(atom) + + + unique_atoms.sort() + + return unique_atoms + + +def get_interactions_involving(interactions, targets): + ''' Get the indicies (+1) of all the interactions involving target. This list can be used as input to plot_coop(), as it is + then formatted the way that function accepts these interactions. + + Input: + interactions: list of interactions as output from read_coop() + target: the particular atom that should be involved in the interactions contained in the output list + + Output: + target_interactions: Indices (+1) of all the interactions involving target atom.''' + + target_interactions = [] + appended_interactions = [] + + + if type(targets) == list: + for target in targets: + for ind, interaction in enumerate(interactions): + if target in interaction.split('->') and interaction not in appended_interactions: + target_interactions.append(ind+1) + appended_interactions.append(interaction) + + else: + for ind, interaction in enumerate(interactions): + if targets in interaction.split('->'): + target_interactions.append(ind+1) + + + return target_interactions + + + + +def update_rc_params(rc_params): + ''' Update all passed run commands in matplotlib''' + + if rc_params: + for key in rc_params.keys(): + plt.rcParams.update({key: rc_params[key]}) + + + +def update_options(options, required_options, default_options): + ''' Update all passed options''' + + + for option in required_options: + if option not in options.keys(): + options[option] = default_options[option] + + + + return options + + + +def determine_width(options): + + conversion_cm_inch = 0.3937008 # cm to inch + + if options['column_type'] == 'single': + column_width = options['single_column_width'] + elif options['column_type'] == 'double': + column_width = options['double_column_width'] + + column_width *= conversion_cm_inch + + + width_ratio = [float(num) for num in options['width_ratio'].split(':')] + + + width = column_width * width_ratio[0]/width_ratio[1] + + + return width + + +def determine_height(options, width): + + aspect_ratio = [float(num) for num in options['aspect_ratio'].split(':')] + + height = width/(aspect_ratio[0] / aspect_ratio[1]) + + return height + + +def scale_figure(options, width, height): + width = width * options['upscaling_factor'] * options['compress_width'] + height = height * options['upscaling_factor'] * options['compress_height'] + + return width, height + + + +def swap_values(dict, key1, key2): + + key1_val = dict[key1] + dict[key1] = dict[key2] + dict[key2] = key1_val + + return dict + + + +def generate_colours(options: dict): + + if not isinstance(options['palettes'], list): + options['palettes'] = [options['palettes']] + + # Creates a list of all the colours that is passed in the colour_cycles argument. Then makes cyclic iterables of these. + colour_collection = [] + + for palette in options['palettes']: + mod = importlib.import_module("palettable.colorbrewer.%s" % palette[0]) + colour = getattr(mod, palette[1]).mpl_colors + colour_collection = colour_collection + colour + + colour_cycle = itertools.cycle(colour_collection) + + + return colour_cycle \ No newline at end of file diff --git a/nafuma/dft/io.py b/nafuma/dft/io.py new file mode 100644 index 0000000..0c32451 --- /dev/null +++ b/nafuma/dft/io.py @@ -0,0 +1,545 @@ +import pandas as pd +import matplotlib.pyplot as plt +import numpy as np + +import linecache + +import nafuma.auxillary as aux + + +def open_doscar(doscar='DOSCAR'): + with open(doscar, 'r') as dos: + lines = dos.readlines() + + return lines + +def open_poscar(poscar='POSCAR'): + with open(poscar, 'r') as pos: + lines = pos.readlines() + + return lines + +def open_outcar(outcar='OUTCAR'): + with open(outcar, 'r') as out: + lines = out.readlines() + + return lines + +def get_number_of_atoms(doscar='DOSCAR'): + lines = open_doscar(doscar) + + return int(lines[0].strip().split()[0]) + +def get_atoms(poscar): + + with open(poscar, 'r') as poscar: + lines = poscar.readlines() + + atoms = lines[5].split() + atom_num = lines[6].split() + + + atom_num = [int(num) for num in atom_num] + + atoms_dict = {} + + for ind, atom in enumerate(atoms): + atoms_dict[atom] = atom_num[ind] + + return atoms, atom_num, atoms_dict + + +def get_fermi_level(doscar='DOSCAR', headerlines=6): + lines = open_doscar(doscar) + + return float(lines[headerlines-1].strip().split()[3]) + + +def get_valence_electron_count(outcar='OUTCAR', poscar='POSCAR'): + lines = open_outcar(outcar) + + atoms, atoms_dict = get_atoms(poscar) + n_atoms = len(atoms) + + for line in lines: + line = line.strip() + if line[0:4] == "ZVAL": + valence_electrons = line.split()[-n_atoms:] + break + + return valence_electrons + + + + +def read_coop(data={}, options={}): + ''' Reads a COOPCAR.lobster file and prepares the DataFrame for further data handling. + + Input: + path: The path to the COOPCAR.lobster-file + + Output: + coopcar: A DataFrame containing the COOP data + interactions: A list of interactions''' + + + required_options = ['collapse'] + + default_options = { + 'collapse': False + } + + + options = aux.update_options(options=options, required_options=required_options, default_options=default_options) + + interactions = determine_interactions(data['coopcar']) + + coopcar = pd.read_csv(data['coopcar'], skiprows=3+len(interactions), header=None, delim_whitespace=True) + + spin_polarised = determine_spin_polarisation(coopcar, interactions) + + + # Create list of column names + # If the run is spin polarised, add all up states first and then all down states + if spin_polarised: + columns = ['Energy', 'avg_up', 'avg_int_up'] + for i in range(1, len(interactions)+1): + columns += ['interaction{}_up'.format(i), 'interaction{}_int_up'.format(i)] + columns += ['avg_down', 'avg_int_down'] + for i in range(1, len(interactions)+1): + columns += ['interaction{}_down'.format(i), 'interaction{}_int_down'.format(i)] + + # Otherwise just + else: + columns = ['Energy', 'avg', 'avg_int'] + for i in range(1, len(interactions)+1): + columns += ['interaction_{}'.format(i), 'interaction_{}_int'.format(i)] + + + + + coopcar.columns = columns + + + if options['collapse']: + + columns_collapsed = ['Energy'] + + for column in columns: + if column.split('_')[0] not in columns_collapsed: + columns_collapsed.append(''.join(column.split('_')[:-1])) + columns_collapsed.append(''.join(column.split('_')[:-1])+'_int') + + + for column in columns_collapsed: + if column != 'Energy': + coopcar[column] = coopcar[column+'_up'] + coopcar[column+'_down'] + + columns.remove('Energy') + coopcar.drop(columns, axis=1, inplace=True) + coopcar.columns = columns_collapsed + + + + + + return coopcar, interactions + + +def read_cohp(path, flip_sign=False): + ''' Reads a COHPCAR.lobster file and prepares the DataFrame for further data handling. + + Input: + path: The path to the COOPCAR.lobster-file + flip_sign: Boolean value to determine whether all COHP-values should be flipped (that is, return -COHP) + + Output: + coopcar: A DataFrame containing the COOP data + interactions: A list of interactions''' + + interactions = determine_interactions(path) + + cohpcar = pd.read_csv(path, skiprows=3+len(interactions), header=None, delim_whitespace=True) + + spin_polarised = determine_spin_polarisation(cohpcar, interactions) + + + # Create list of column names + # If the run is spin polarised, add all up states first and then all down states + if spin_polarised: + columns = ['Energy', 'avg_up', 'avg_up_int'] + for i in range(1, len(interactions)+1): + columns += ['interaction_{}_up'.format(i), 'interaction_{}_up_int'.format(i)] + columns += ['avg_down', 'avg_down_int'] + for i in range(1, len(interactions)+1): + columns += ['interaction_{}_down'.format(i), 'interaction_{}_up_int'.format(i)] + + # Otherwise just + else: + columns = ['Energy', 'avg', 'avg_int'] + for i in range(1, len(interactions)+1): + columns += ['interaction_{}'.format(i), 'interaction_{}_int'.format(i)] + + + cohpcar.columns = columns + + if flip_sign: + columns = columns[1:] + + for column in columns: + cohpcar[column] = -cohpcar[column] + + return cohpcar, interactions + + +def determine_interactions(path): + ''' Determines the number of interactions present in the COOPCAR.lobster file. + + Input: + path: The path to the COOPCAR.lobster-file + + Output: + interactions: A list of strings with the interactions in the COOPCAR.lobster file''' + + + with open(path, 'r') as coop: + lines = coop.readlines() + + + interactions = [] + for line in lines: + if line[0:2] == 'No': + interactions.append(line.split(':')[-1].split('(')[0]) + + + return interactions + + + +def determine_spin_polarisation(coopcar, interactions): + ''' Determines whether a COOPCAR.lobster file is spin polarised or not. + + Input: + coopcar: A DataFrame containing the COOP data + interactions: A list of interactions obtained from the determine_interactions function + + Output: + spin_polarised: Boolean value to indicate whether or not the COOP data is spin polarised''' + + number_of_interactions = len(interactions) + + spin_polarised = True if coopcar.shape[1] == 4*number_of_interactions + 5 else False + + return spin_polarised + + + + +def read_dos(path, flip_down=True): + + dos_info = get_doscar_information(path) + + with open(path, 'r') as doscar: + count = 0 + raw_data = [] + + while count < dos_info["NEDOS"] + 6: + if count >= 6: + data_line = [float(x) for x in doscar.readline().split()] + raw_data.append(data_line) + + else: + doscar.readline() + + count += 1 + + + dos = pd.DataFrame(raw_data) + + if dos_info["spin_polarised"]: + header_names = ['energy', 'total_up', 'total_down', 'total_integrated_up', 'total_integrated_down'] + else: + header_names = ['energy', 'total', 'total_integrated'] + + + dos.columns = header_names + + if dos_info["spin_polarised"] and flip_down: + dos["total_down"] = -dos["total_down"] + dos["total_integrated_down"] = -dos["total_integrated_down"] + + return dos + + +def read_pdos(data: dict, options={}): + + required_options = ['flip_down', 'sum_atoms', 'sum_orbitals', 'collapse_spin', 'adjust', 'manual_adjust', 'normalise', 'normalise_unit_cell', 'normalisation_factor'] + + default_options = { + 'flip_down': True, + 'sum_atoms': False, + 'sum_orbitals': False, + 'collapse_spin': False, + 'adjust': False, + 'manual_adjust': None, + 'normalise': False, + 'normalise_unit_cell': False, + 'normalisation_factor': None + } + + options = aux.update_options(options=options, required_options=required_options, default_options=default_options) + + # Get information about the DOSCAR-file (ENMIN, ENMAX, NEDOS, EFERMI and spin_polarised) + dos_info = get_doscar_information(data['doscar']) + + # Get information from the POSCAR-file (or CONTCAR) (which species, number of atoms per specie and a dictionary matching the two) + species, atom_num, atoms_dict = get_atoms(data['poscar']) + + + # Open DOSCAR and read all lines + with open(data['doscar'], 'r') as file: + doscar_data = file.readlines() + + # Make list of all individual atoms + atoms = [] + for specie in species: + for i in range(0,atoms_dict[specie]): + atoms.append(specie) + + + # Loop through all the atoms and make a DataFrame for each atom. + + pdos_full = [] + for ind, atom in enumerate(atoms): + + # Define line to start reading + start = int((1 + ind) * dos_info["NEDOS"] + 6 + (ind * 1)) + + + + add_d_orbitals = False + add_p_orbitals = False + # Check if d-orbitals are included (if DOSCAR comes from LOBSTER), otherwise they will have to be added to make the DataFrames the same size + if dos_info['kind'] == 'LOBSTER': + # Extract basis sets from each atom + basis_sets = doscar_data[start].split(';')[-1] + + if not 'd' in basis_sets: + add_d_orbitals = True + + if basis_sets.count('p') != 6: + add_p_orbitals = True + + + + # Start reading lines into a list and convert to DataFrame and cast all entries as float + pdos_atom = [] + for i in range(1,dos_info["NEDOS"]+1): + pdos_atom.append(doscar_data[start+i].split()) + + + pdos_atom = pd.DataFrame(pdos_atom) + pdos_atom = pdos_atom.astype(float) + + + + + # Give the columns names and add second set of p-orbitals, d-orbitals, or both if they don't exist in the data + if add_p_orbitals: + if dos_info['spin_polarised']: + pdos_atom.columns = ['Energy', 's_u', 's_d', 'p1_u', 'p1_d', 'p2_u', 'p2_d', 'p3_u', 'p3_d'] + pdos_atom[['2p1_u', '2p1_d', '2p2_u', '2p2_d', '2p3_u', '2p3_d']] = 0 + + if add_d_orbitals: + pdos_atom[['d1_u', 'd1_d', 'd2_u', 'd2_d', 'd3_u', 'd3_d', 'd4_u', 'd4_d', 'd5_u', 'd5_d']] = 0 + pdos_atom['tot_u'] = pdos_atom[['s_u', 'p1_u', 'p2_u', 'p3_u', '2p1_u', '2p2_u', '2p3_u', 'd1_u', 'd2_u', 'd3_u', 'd4_u', 'd5_u']].sum(axis=1) + pdos_atom['tot_d'] = pdos_atom[['s_d', 'p1_d', 'p2_d', 'p3_d', '2p1_d', '2p2_d', '2p3_d', 'd1_d', 'd2_d', 'd3_d', 'd4_d', 'd5_d']].sum(axis=1) + + else: + pdos_atom['tot_u'] = pdos_atom[['s_u', 'p1_u', 'p2_u', 'p3_u', '2p1_u', '2p2_u', '2p3_u', 'd1_u', 'd2_u', 'd3_u', 'd4_u', 'd5_u']].sum(axis=1) + pdos_atom['tot_d'] = pdos_atom[['s_d', 'p1_d', 'p2_d', 'p3_d', '2p1_d', '2p2_d', '2p3_d', 'd1_d', 'd2_d', 'd3_d', 'd4_d', 'd5_d']].sum(axis=1) + + + + elif add_d_orbitals: + if dos_info["spin_polarised"]: + pdos_atom.columns = ['Energy', 's_u', 's_d', 'p1_u', 'p1_d', 'p2_u', 'p2_d', 'p3_u', 'p3_d'] + pdos_atom[['d1_u', 'd1_d', 'd2_u', 'd2_d', 'd3_u', 'd3_d', 'd4_u', 'd4_d', 'd5_u', 'd5_d']] = 0 + pdos_atom['tot_u'] = pdos_atom[['s_u', 'p1_u', 'p2_u', 'p3_u', '2p1_u', '2p2_u', '2p3_u', 'd1_u', 'd2_u', 'd3_u', 'd4_u', 'd5_u']].sum(axis=1) + pdos_atom['tot_d'] = pdos_atom[['s_d', 'p1_d', 'p2_d', 'p3_d', '2p1_d', '2p2_d', '2p3_d', 'd1_d', 'd2_d', 'd3_d', 'd4_d', 'd5_d']].sum(axis=1) + + + else: + pdos_atom.columns = ['Energy', 's', 'p1', 'p2', 'p3'] + pdos_atom[['d1', 'd2', 'd3', 'd4', 'd5']] = 0 + pdos_atom['tot'] = pdos_atom[['s', 'p1', 'p2', 'p3', '2p1', '2p2', '2p3', 'd1', 'd2', 'd3', 'd4', 'd5']].sum(axis=1) + + else: + if dos_info["spin_polarised"]: + if dos_info['kind'] == 'LOBSTER': + pdos_atom.columns = ['Energy', 's_u', 's_d', 'p1_u', 'p1_d', 'p2_u', 'p2_d', 'p3_u', 'p3_d', '2p1_u', '2p1_d', '2p2_u', '2p2_d', '2p3_u', '2p3_d', 'd1_u', 'd1_d', 'd2_u', 'd2_d', 'd3_u', 'd3_d', 'd4_u', 'd4_d', 'd5_u', 'd5_d'] + pdos_atom['tot_u'] = pdos_atom[['s_u', 'p1_u', 'p2_u', 'p3_u', '2p1_u', '2p2_u', '2p3_u','d1_u', 'd2_u', 'd3_u', 'd4_u', 'd5_u']].sum(axis=1) + pdos_atom['tot_d'] = pdos_atom[['s_d', 'p1_d', 'p2_d', 'p3_d', '2p1_d', '2p2_d', '2p3_d', 'd1_d', 'd2_d', 'd3_d', 'd4_d', 'd5_d']].sum(axis=1) + + elif dos_info['kind'] == 'VASP': + pdos_atom.columns = ['Energy', 's_u', 's_d', 'p1_u', 'p1_d', 'p2_u', 'p2_d', 'p3_u', 'p3_d', 'd1_u', 'd1_d', 'd2_u', 'd2_d', 'd3_u', 'd3_d', 'd4_u', 'd4_d', 'd5_u', 'd5_d'] + pdos_atom['tot_u'] = pdos_atom[['s_u', 'p1_u', 'p2_u', 'p3_u', 'd1_u', 'd2_u', 'd3_u', 'd4_u', 'd5_u']].sum(axis=1) + pdos_atom['tot_d'] = pdos_atom[['s_d', 'p1_d', 'p2_d', 'p3_d', 'd1_d', 'd2_d', 'd3_d', 'd4_d', 'd5_d']].sum(axis=1) + + else: + pdos_atom.columns = ['Energy', 's', 'p1', 'p2', 'p3', '2p1', '2p2', '2p3', 'd1', 'd2', 'd3', 'd4', 'd5'] + pdos_atom['tot'] = pdos_atom[['s', 'p1', 'p2', 'p3', '2p1', '2p2', '2p3', 'd1', 'd2', 'd3', 'd4', 'd5']].sum(axis=1) + + + + # Change the sign of the spin down columns if flip_down is True + if options['flip_down'] and dos_info['spin_polarised'] and not options['collapse_spin']: + down = [orbital for orbital in pdos_atom.columns if '_d' in orbital] + + for orbital in down: + pdos_atom[orbital] = -pdos_atom[orbital] + + + if options['manual_adjust']: + pdos_atom["Energy"] = pdos_atom["Energy"] - options['manual_adjust'] + + + + pdos_full.append(pdos_atom) + + + # If sum_atoms is True, all DataFrames of a given specie will be added together + if options['sum_atoms']: + + # Initalise a new emtpy list that will become our pdos_full + pdos_full_sum_atoms = [] + + + start = 0 + # Loop through each specie so that each specie will get exactly one DataFrame + for specie in species: + # Initialise with first DataFrame of the specie + atom_pdos_summed = pdos_full[start] + + # Loop over DOSes and add if there's a match for specie + for ind, pdos in enumerate(pdos_full): + if atoms[ind] == specie and ind != start: + atom_pdos_summed = atom_pdos_summed + pdos + + + # Divide the Energy by the number of DataFrames added to get back to the original value of the Energy + atom_pdos_summed["Energy"] = atom_pdos_summed["Energy"] / atoms_dict[specie] + + + if options['normalise']: + + if dos_info["spin_polarised"]: + columns = atom_pdos_summed.columns + #columns = ['s_u', 's_d', 'p1_u', 'p1_d', 'p2_u', 'p2_d', 'p3_u', 'p3_d', '2p1_u', '2p1_d', '2p2_u', '2p2_d', '2p3_u', '2p3_d', 'd1_u', 'd1_d', 'd2_u', 'd2_d', 'd3_u', 'd3_d', 'd4_u', 'd4_d', 'd5_u', 'd5_d', 'tot_u', 'tot_d'] + for column in columns: + atom_pdos_summed[column] = atom_pdos_summed[column] / atoms_dict[specie] + + + # Append the new DataFrame for a given specie to the list + pdos_full_sum_atoms.append(atom_pdos_summed) + + start += atoms_dict[specie] + + + # Rename the list + pdos_full = pdos_full_sum_atoms + + + # If collapse_spin is True for a spin polarised DOSCAR, the up and down channels of each orbital will be added together. + if options['collapse_spin'] and dos_info['spin_polarised']: + + pdos_full_spin_collapsed = [] + for pdos in pdos_full: + temp_pdos = pd.DataFrame() + + temp_pdos["Energy"] = pdos["Energy"] + + temp_pdos['s'] = pdos[['s_u', 's_d']].sum(axis=1) + + temp_pdos['p1'] = pdos[['p1_u', 'p1_d']].sum(axis=1) + temp_pdos['p2'] = pdos[['p2_u', 'p2_d']].sum(axis=1) + temp_pdos['p3'] = pdos[['p3_u', 'p3_d']].sum(axis=1) + + temp_pdos['2p1'] = pdos[['2p1_u', '2p1_d']].sum(axis=1) + temp_pdos['2p2'] = pdos[['2p2_u', '2p2_d']].sum(axis=1) + temp_pdos['2p3'] = pdos[['2p3_u', '2p3_d']].sum(axis=1) + + temp_pdos['d1'] = pdos[['d1_u', 'd1_d']].sum(axis=1) + temp_pdos['d2'] = pdos[['d2_u', 'd2_d']].sum(axis=1) + temp_pdos['d3'] = pdos[['d3_u', 'd3_d']].sum(axis=1) + temp_pdos['d4'] = pdos[['d4_u', 'd4_d']].sum(axis=1) + temp_pdos['d5'] = pdos[['d5_u', 'd5_d']].sum(axis=1) + + temp_pdos['tot'] = pdos[['tot_u', 'tot_d']].sum(axis=1) + + pdos_full_spin_collapsed.append(temp_pdos) + + + pdos_full = pdos_full_spin_collapsed + dos_info['spin_polarised'] = False + + + + # If sum_orbitals is True, all columns belonging to a particular set of orbitals will be added together. + if options['sum_orbitals']: + pdos_full_sum_orbitals = [] + + for pdos in pdos_full: + temp_pdos = pd.DataFrame() + + temp_pdos["Energy"] = pdos["Energy"] + + if dos_info['spin_polarised']: + temp_pdos['s_u'] = pdos['s_u'] + temp_pdos['s_d'] = pdos['s_d'] + + temp_pdos['p_u'] = pdos[['p1_u', 'p2_u', 'p3_u']].sum(axis=1) + temp_pdos['p_d'] = pdos[['p1_d', 'p2_d', 'p3_d']].sum(axis=1) + + if len(pdos_full[0].columns) == 25: + temp_pdos['2p_u'] = pdos[['2p1_u', '2p2_u', '2p3_u']].sum(axis=1) + temp_pdos['2p_d'] = pdos[['2p1_d', '2p2_d', '2p3_d']].sum(axis=1) + + temp_pdos['d_u'] = pdos[['d1_u', 'd2_u', 'd3_u', 'd4_u', 'd5_u']].sum(axis=1) + temp_pdos['d_d'] = pdos[['d1_d', 'd2_d', 'd3_d', 'd4_d', 'd5_d']].sum(axis=1) + + temp_pdos['tot_u'] = pdos['tot_u'] + temp_pdos['tot_d'] = pdos['tot_d'] + + else: + temp_pdos['s'] = pdos['s'] + temp_pdos['p'] = pdos[['p1', 'p2', 'p3']].sum(axis=1) + temp_pdos['2p'] = pdos[['2p1', '2p2', '2p3']].sum(axis=1) + temp_pdos['d'] = pdos[['d1', 'd2', 'd3', 'd4', 'd5']].sum(axis=1) + temp_pdos['tot'] = pdos['tot'] + + pdos_full_sum_orbitals.append(temp_pdos) + + pdos_full = pdos_full_sum_orbitals + + + + + return pdos_full, dos_info + + +def get_doscar_information(path): + ''' Reads information from the DOSCAR''' + + kind = 'LOBSTER' if 'LOBSTER' in linecache.getline(path, 5) else 'VASP' + dos_info_raw = linecache.getline(path, 6).split() + spin_polarised = True if len(linecache.getline(path, 7).split()) == 5 else False + + + dos_info = {'ENMIN': float(dos_info_raw[0]), 'ENMAX': float(dos_info_raw[1]), 'NEDOS': int(dos_info_raw[2]), 'EFERMI': float(dos_info_raw[3]), 'spin_polarised': spin_polarised, 'kind': kind} + + return dos_info + + +#def get_bader_charges(poscar='POSCAR', acf='ACF.dat'): + diff --git a/nafuma/dft/structure.py b/nafuma/dft/structure.py new file mode 100644 index 0000000..04584cb --- /dev/null +++ b/nafuma/dft/structure.py @@ -0,0 +1,906 @@ +import math +import re +import pandas as pd +import numpy as np +from scipy.optimize import curve_fit + +import matplotlib.pyplot as plt +from matplotlib.ticker import (MultipleLocator, FormatStrFormatter,AutoMinorLocator) +import importlib +import matplotlib.patches as mpatches +import matplotlib.lines as mlines +from mpl_toolkits.axisartist.axislines import Subplot +from cycler import cycler +import itertools + +from ase import Atoms +from ase.io.trajectory import Trajectory +from ase import io +from ase.units import kJ +from ase.eos import EquationOfState +import os +import os.path + + +def read_eos_data(path, options): + ''' Reads volume and energy data from a energy-volume run and fits the data to an equation of state. Outputs a list with one pandas DataFrame containing the data points from the DFT-calculations, + one DataFrame containing the fitted curve data points and one dictionary with equilibrium volume, equilibrium energy and bulk modulus in GPa + + path: Path to the folder containing the energ.dat and POSCAR files. energ.dat must have two columns with volumes in the first, energy in the second separated by whitespace. + atoms_per_fu: Number of atoms per formula unit. Used to scale the values to be comparable with other calculations that may have a different sized unit cell. + eos: Type of equation of state to fit to. Same keywords as the ones used in ASE, as it simply calls ASE to fit the equation of state. + ''' + + required_options = ['atoms_per_fu', 'reference', 'eos'] + + default_options = { + 'atoms_per_fu': -1, # Scaling factor to output energy per f.u. + 'reference': 0, # Whether the energy should be relative to some reference energy (typically lowest energy) + 'eos': 'birchmurnaghan', # what type of EoS curve to fit the data to. Options: murnaghan, birch, birchmurnaghan, vinet, pouriertarantola + } + + + options = update_options(options=options, required_options=required_options, default_options=default_options) + + + # Make paths for the energ.dat and POSCAR files. + energ_path = os.path.join(path, 'energ.dat') + poscar_path = os.path.join(path, 'POSCAR') + + # Read POSCAR and calculate the scale factor to give values per formula unit + at = io.read(poscar_path) + + if options['atoms_per_fu'] == -1: + scale_factor = 1 + else: + scale_factor = options['atoms_per_fu'] / len(at) + + # Get the label + label = os.path.basename(path) + + # Reads the energ.dat file and structures the data into a pandas DataFrame. Then scales the values according to the scale factor. + dft_df = pd.read_csv(energ_path, delim_whitespace=True, header=None) + dft_df.columns = ['Configuration', 'Volume', 'Energy'] + dft_df['Energy'] = dft_df['Energy'] * scale_factor + dft_df['Volume'] = dft_df['Volume'] * scale_factor + + + dft_df["Energy"] = dft_df["Energy"] - options['reference'] # subtracts a reference energy if provided. THis value defaults to 0, so will not do anything if not provided. + + # Fit data to Equation of State using ASEs EquationOfState object. Makes a DataFrame out of the data points of the fitted curve. Also makes a ditionary of the equilibrium constants, + #then packages everything in a list which is returned by the function. + eos = EquationOfState(dft_df['Volume'].values, dft_df['Energy'].values, eos=options['eos']) + v0, e0, B = eos.fit() + + eos_df = pd.DataFrame(data={'Volume': eos.getplotdata()[4], 'Energy': eos.getplotdata()[5]}) + + equilibrium_constants = {'v0': v0, 'e0': e0,'B': B/kJ * 1.0e24} + + data = [dft_df, eos_df, equilibrium_constants, label] + + return data + + +def read_eos_datas(path, options): + + + required_options = ['subset', 'sort_by'] + + default_options = { + 'subset': None, # list with directory names of what you want to include + 'sort_by': 'e0', # whether the data should be sorted or not - relevant for bar plots, but also for the order of the entries in the legend in the EoScruve plot + } + + options = update_options(options=options, required_options=required_options, default_options=default_options) + + # If a subset of directories is not specified, will create a list of all directories in the path given. + if not options['subset']: + dirs = [dir for dir in os.listdir(path) if os.path.isdir(os.path.join(path, dir)) and dir[0] != '.'] + else: + dirs = options['subset'] + + + datas = [] + + + # Loop through all subdirectories and reads the data from these. Also appends the name of the directory to the list that is returned from the plot_eos_data() function + for dir in dirs: + subdir = os.path.join(path, dir) + data = read_eos_data(subdir, options) + datas.append(data) + + + # Sorts the data if sort is enabled. + if options['sort_by']: + datas = sort_data(datas, options['sort_by']) + + + return datas + + +def get_summarised_data(path, options): + + datas = read_eos_datas(path=path, options=options) + + summary = [] + for data in datas: + summary.append([data[3], data[2]['e0'], data[2]['v0'], data[2]['B']]) + + df = pd.DataFrame(summary) + df.columns = ['Label', 'E0', 'V0', 'B'] + + emin = df["E0"].min() + + df["dE0"] = df["E0"] - emin + + # Rearranging the columns + df = df[['Label', 'E0', 'dE0', 'V0', 'B']] + + return df + + +def plot_eos_data(path, options): + ''' Plots the data from the energy-volume curve runs. Allows plotting of just the energy-volume curves, a bar plot showing the equilibrium energies or both. + + path: path to where the data is located. It should point to a directory with subdirectories for each structure to be plotted. Inside each of these subdirectories there should be an energ.dat and a POSCAR file. + atoms_per_fu: Number of atoms per formula unit. Used to scale the values to be comparable with other calculations that may have a different sized unit cell. + dirs: List of directory names if only a subset of all available datasets is to be plotted. Defaults to None, and will thus get data from all subdirectories. + eos: Type of equation of state to fit to. Same keywords as the ones used in ASE, as it simply calls ASE to fit the equation of state. + width: Width of the total figure. Defaults to None, which will again default to width=20. + height: Height of the total figure. Defaults to None, which will again will default to height= width / phi where phi is the golden ratio. + dpi: Dots per inch of the figure. Defaults to pyplot's default + colour_cycles: List of tuples with sets of colours for the palettable colour collection. Defaults to two sets of in total 20 colours. Used for giving different colours to energy-volume curves. + energyunit: The energy unit. Defaults to eV per formula unit. Only used on the axis labels. + volumeunit: The volume unit. Defaults to Å^3. Only used on the axis labels. + xlim: Limits of the x-axes. List of min and max. If mode = both is used, has to contain two lists for each of the plots. As the x-limits for a bar plot is nonesense, should just contain a list with a NoneType. + ylim: Limits of the y-axes. List of min and max. If mode = both is used, has to contain two lists for each of the plots. + sort: Whether or not to sort the data from lowest to highest equilibrium energy. Defaults to True. + sort_by: What to sort by if sort is enabled. Defaults to e0. Other options: v0 = equilibrium volumes, B = bulk moduli. Alphabetical order sorting is not implemented. + mode: Determines what to plot. Defaults to energy-volume curves ('curves'). Other options: 'bars', bar-plot of equilibrium energies. 'both', both energy-volume curves and bar plots are plotted side-by-side. + highlight: Takes a list, either of booleans to highlight certain bars (must be the same length as the number of data sets). Alternatively can contain only names of the datasets to highlight. Defaults to None.''' + + + required_options = ['plot_kind', 'highlight', + 'reference', + 'eos', 'sort_by', + 'curve_colours', + 'bar_colours', + 'marker_cycle', + 'ylim', + 'legend_map', + 'rc_params', + 'legend'] + + + default_options = { + 'plot_kind': 'EoScurve', # EoScurve or EoSbars + 'highlight': None, # list with directory names (or Boolean array) of which bars to highlight. Only relevant to EoSbars + 'reference': 0, # Whether the energy should be relative to some reference energy (typically lowest energy) + 'eos': 'birchmurnaghan', # what type of EoS curve to fit the data to. Options: murnaghan, birch, birchmurnaghan, vinet, pouriertarantola + 'sort_by': 'e0', # whether the data should be sorted or not - relevant for bar plots, but also for the order of the entries in the legend in the EoScruve plot + 'curve_colours': [('qualitative', 'Dark2_8'), ('qualitative', 'Paired_12')], # a set of two colour cycles from the palettable package. Requires many colours for the EoScurve plot + 'bar_colours': [('qualitative', 'Dark2_3'), ('qualitative', 'Pastel2_3')], # set of two colour cycles from the palettable package. Extracts first element of each to make the highlighted and subdued colours respectively. Should be replaced in future by explicit passing of colours + 'marker_cycle': ('o', '*', '^', 'v', 'd', 'H', '8', '>', 'P', 'X'), # marker styles for the EoScurve plot + 'ylim': None, # y-limits (ist) + 'legend': True, + 'legend_map': None, # a dictionary with mappings between the folder names and what should appear in the legend + 'rc_params': None # dictionary of run commands to update plot style + } + + + options = update_options(options=options, required_options=required_options, default_options=default_options) + + # Create path to the data + datas = read_eos_datas(path=path, options=options) + + + ### PLOT THE ENERGY-VOLUME CURVES + if options['plot_kind'] == 'EoScurve': + + # Fetches a figure and axes object from the prepare_plot() function + fig, ax = prepare_plot(options=options) + + # Make an cyclic iterable of markers to be used for the calculated data points. + marker_cycle = itertools.cycle(options['marker_cycle']) + + + # Creates a list of all the colours that is passed in the colour_cycles argument. Then makes cyclic iterables of these. + colour_collection = [] + for cycle in options['curve_colours']: + mod = importlib.import_module("palettable.colorbrewer.%s" % cycle[0]) + colour = getattr(mod, cycle[1]).mpl_colors + colour_collection = colour_collection + colour + + colour_cycle = itertools.cycle(colour_collection) + + labels = [] + colours = [] + markers = [] + + + # For each of the data sets, extracts the data and plots them. + for data in datas: + dft_df, eos_df, label = data[0], data[1], data[3] + + + # If ylim is passed, only plot those that have a minimum energy below the max ylim parameter + if options['ylim']: + plot = True if dft_df["Energy"].min() < options['ylim'][1] else False + else: + plot = True + + if plot: + if options['label_map']: + labels.append(options['label_map'][label]) + + colours.append(next(colour_cycle)) + markers.append(next(marker_cycle)) + + dft_df.plot.scatter(x=1, y=2, ax=ax, marker=markers[-1], color=colours[-1], s=20) + eos_df.plot(x=0, y=1, ax=ax, color=colours[-1], label='_', ls='--') + + + if options['legend']: + options['legend_content'] = [labels, colours, markers] + + + ### PLOT THE BAR PLOTS + elif options['plot_kind'] == 'EoSbars': + + # Fetches a figure and axes object from the prepare_plot() function + fig, ax = prepare_plot(options=options) + + e0 = [] + labels = [] + colours = [] + + # Pick out colour for highlighting (NB! These colours are not passed as arguments, but could be in future) + + bar_colours = [] + for cycle in options['bar_colours']: + mod = importlib.import_module("palettable.colorbrewer.%s" % cycle[0]) + bar_colours.append(getattr(mod, cycle[1]).mpl_colors[0]) + + + # Loops through the datasets, picks out equilibrium volume and labels and sets colours according to the whether the highlight option is used or not. + for data in datas: + + if options['ylim']: + plot = True if data[2]['e0'] < options['ylim'][1] else False + else: + plot = True + + if plot: + + # Adds 100 if plotting in relative mode. The bases of the bar plots are sunk by 100 during plotting + adjustment = 100 if options['reference'] != 0 else 100 + print(adjustment) + + e0.append(data[2]['e0']+adjustment) + print(e0[-1]) + labels.append(options['label_map'][data[3]]) + + if options['highlight'] is not None: + if data[3] in options['highlight']: + colours.append(bar_colours[0]) + else: + colours.append(bar_colours[1]) + + elif options['highlight'] is not None and type(options['highlight'][0] == str): + if labels[-1] in options['highlight']: + colours.append(bar_colours[0]) + else: + colours.append(bar_colours[1]) + + else: + colours.append(bar_colours[0]) + + # Makes the bar plot. + bottom = -100 if options['reference'] != 0 else 0 + plt.bar(range(len(e0)), e0, color=colours, bottom=bottom) + plt.xticks(range(len(e0)), labels, rotation=90) + + + fig, ax = prettify_plot(fig=fig, ax=ax, options=options) + + + + +def sort_data(datas, sort_by='e0'): + ''' Bubble sort algorithm to sort the data sets''' + + l = len(datas) + + for i in range(0, l): + for j in range(0, l-i-1): + if datas[j][2]['{}'.format(sort_by)] > datas[j+1][2]['{}'.format(sort_by)]: + temp = datas[j] + datas[j] = datas[j+1] + datas[j+1] = temp + + return datas + + + +def prepare_plot(options={}): + + # Reset run commands + plt.rcdefaults() + + # Update run commands if any is passed + if 'rc_params' in options.keys(): + update_rc_params(options['rc_params']) + + + + required_options = ['single_column_width', 'double_column_width', 'column_type', 'width_ratio', 'aspect_ratio', 'compress_width', 'compress_height', 'upscaling_factor', 'dpi'] + default_options = { + 'single_column_width': 8.3, + 'double_column_width': 17.1, + 'column_type': 'single', + 'width_ratio': '1:1', + 'aspect_ratio': '1:1', + 'compress_width': 1, + 'compress_height': 1, + 'upscaling_factor': 1.0, + 'dpi': 600} + + options = update_options(options, required_options, default_options) + + width = determine_width(options) + height = determine_height(options, width) + width, height = scale_figure(options=options, width=width, height=height) + + fig, ax = plt.subplots(figsize=(width, height), dpi=options['dpi']) + + return fig, ax + + + + + + +def update_rc_params(rc_params): + ''' Update all passed run commands in matplotlib''' + + if rc_params: + for key in rc_params.keys(): + plt.rcParams.update({key: rc_params[key]}) + + + +def update_options(options, required_options, default_options): + ''' Update all passed options''' + + + for option in required_options: + if option not in options.keys(): + options[option] = default_options[option] + + + return options + + +def determine_width(options): + + conversion_cm_inch = 0.3937008 # cm to inch + + if options['column_type'] == 'single': + column_width = options['single_column_width'] + elif options['column_type'] == 'double': + column_width = options['double_column_width'] + + column_width *= conversion_cm_inch + + + width_ratio = [float(num) for num in options['width_ratio'].split(':')] + + + width = column_width * width_ratio[0]/width_ratio[1] + + + return width + + +def determine_height(options, width): + + aspect_ratio = [float(num) for num in options['aspect_ratio'].split(':')] + + height = width/(aspect_ratio[0] / aspect_ratio[1]) + + return height + +def scale_figure(options, width, height): + width = width * options['upscaling_factor'] * options['compress_width'] + height = height * options['upscaling_factor'] * options['compress_height'] + + return width, height + + + +def prepare_plot_old(width=None, height=None, dpi=None, energyunit='eV', volumeunit=r'Å$^3$', mode='curves', width_ratio=[1, 1], square=True, pad_bottom=None, scale=1, format_params=None): + '''Prepares pyplot figure and axes objects.''' + + + + linewidth = 3*scale + axeswidth = 3*scale + + plt.rc('lines', linewidth=linewidth) + plt.rc('axes', linewidth=axeswidth) + + + if square: + if not width: + width = 20 + + height = width + + + else: + if not width: + width = 20 + + + if not height: + golden_ratio = (math.sqrt(5) - 1) / 2 + height = width*golden_ratio + + + + if mode == 'curves': + + fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(width, height), facecolor='w', dpi=dpi) + + + if mode == 'bars': + + fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(width, height), facecolor='w', dpi=dpi) + + + if mode == 'both': + + fig, ax = plt.subplots(1, 2, figsize=(width, height), gridspec_kw={'width_ratios': width_ratio}) + + + return fig, ax + + +def prettify_plot(fig, ax, options): + '''Prepares pyplot figure and axes objects.''' + + required_options = ['plot_kind', 'hide_x_labels', 'hide_y_labels', 'xunit', 'yunit', 'legend_content', 'legend_position', 'x_tick_locators', 'y_tick_locators', 'tick_directions', 'subplots_adjust', 'xlim', 'ylim'] + + default_options = { + 'plot_kind': 'EoScurve', # EoScurve or EoSbars + 'hide_x_labels': False, # Whether x labels should be hidden + 'hide_y_labels': False, # whether y labels should be hidden + 'xunit': r'Å$^3$', # The unit of the x-values in the curve plot + 'yunit': r'eV f.u.$^{-1}$', # The unit of the y-values in the curve and bar plots + 'xlim': None, + 'ylim': None, + 'legend_content': None, + 'legend_position': ['upper center', (1.10, 0.90)], # the position of the legend passed as arguments to loc and bbox_to_anchor respectively + 'x_tick_locators': [10, 5], # Major and minor tick locators + 'y_tick_locators': [.1, .05], # Major and minor tick locators + 'tick_directions': 'in', # in or out + 'subplots_adjust': [0.1, 0.1, 0.9, 0.9] + } + + options = update_options(options=options, required_options=required_options, default_options=default_options) + + + if options['plot_kind'] == 'EoScurve': + + # Set labels on x- and y-axes + ax.set_xlabel('Volume [{}]'.format(options['xunit'])) + + if not options['hide_y_labels']: + ax.set_ylabel('Energy [{}]'.format(options['yunit'])) + else: + ax.set_ylabel('') + ax.tick_params(labelleft=False) + + + ax.xaxis.set_major_locator(MultipleLocator(options['x_tick_locators'][0])) + ax.xaxis.set_minor_locator(MultipleLocator(options['x_tick_locators'][1])) + + ax.yaxis.set_major_locator(MultipleLocator(options['y_tick_locators'][0])) + ax.yaxis.set_minor_locator(MultipleLocator(options['y_tick_locators'][1])) + + if ax.get_legend(): + ax.get_legend().remove() + + + if options['legend']: + labels = options['legend_content'][0] + colours = options['legend_content'][1] + markers = options['legend_content'][2] + + entries = [] + + for i in range(len(options['legend_content'][0])): + entries.append(mlines.Line2D([], [], label=labels[i], color=colours[i], marker=markers[i], linestyle='None')) + + + fig.legend(handles=entries, loc=options['legend_position'][0], bbox_to_anchor=options['legend_position'][1], frameon=False) + + + + + if options['plot_kind'] == 'EoSbars': + + if not options['hide_y_labels']: + ax.set_ylabel('Energy [{}]'.format(options['yunit'])) + + ax.yaxis.set_major_locator(MultipleLocator(options['y_tick_locators'][0])) + ax.yaxis.set_minor_locator(MultipleLocator(options['y_tick_locators'][1])) + + ax.tick_params(axis='x', which='minor', bottom=False, top=False) + + + + + # Adjust where the axes start within the figure. Default value is 10% in from the left and bottom edges. Used to make room for the plot within the figure size (to avoid using bbox_inches='tight' in the savefig-command, as this screws with plot dimensions) + plt.subplots_adjust(left=options['subplots_adjust'][0], bottom=options['subplots_adjust'][1], right=options['subplots_adjust'][2], top=options['subplots_adjust'][3]) + + + # If limits for x- and y-axes is passed, sets these. + if options['xlim'] is not None: + ax.set_xlim(options['xlim']) + + if options['ylim'] is not None: + ax.set_ylim(options['ylim']) + + + return fig, ax + + + +def prettify_plot_old(fig, ax, energyunit='eV', volumeunit=r'Å$^3$', mode='curves', legend_content=None, pad_bottom=None, scale=1, hide_ylabels=False, xpad=None, ypad=None): + '''Prepares pyplot figure and axes objects.''' + + # Set sizes of ticks, labes etc. + ticksize = 30*scale + labelsize = 30*scale + legendsize = 15*scale + titlesize = 30*scale + + linewidth = 3*scale + axeswidth = 3*scale + markersize = 15*scale + majorticklength = 20*scale + minorticklength = 10*scale + + xpad = 4 if not xpad else xpad + ypad = 4 if not ypad else ypad + + + if mode == 'curves': + + # Set labels on x- and y-axes + ax.set_xlabel('Volume [{}]'.format(volumeunit), size=labelsize, labelpad=xpad) + + if not hide_ylabels: + ax.set_ylabel('Energy [{}]'.format(energyunit), size=labelsize, labelpad=ypad) + else: + ax.set_ylabel('') + + # Set tick parameters + ax.tick_params(axis='both', direction='in', which='major', length=majorticklength, width=axeswidth, right=True, top=True, labelsize=ticksize) + ax.tick_params(axis='both', direction='in', which='minor', length=minorticklength, width=axeswidth, right=True, top=True, labelsize=ticksize) + + ax.tick_params(axis='x', pad=xpad) + ax.tick_params(axis='y', pad=ypad) + + if hide_ylabels: + ax.tick_params(labelleft=False) + + plt.xticks(fontsize=ticksize) + plt.yticks(fontsize=ticksize) + + + ax.xaxis.set_major_locator(MultipleLocator(10)) + ax.xaxis.set_minor_locator(MultipleLocator(5)) + + ax.yaxis.set_major_locator(MultipleLocator(.1)) + ax.yaxis.set_minor_locator(MultipleLocator(.05)) + + + ax.get_legend().remove() + if legend_content: + patches = [] + labels = legend_content[0] + colours = legend_content[1] + markers = legend_content[2] + + entries = [] + + for ind, label in enumerate(legend_content[0]): + entries.append(mlines.Line2D([], [], color=colours[ind], marker=markers[ind], linestyle='None', + markersize=markersize, label=labels[ind])) + + #patches.append(mpatches.Patch(color=colours[ind], label=labels[ind])) + + + fig.legend(handles=entries, loc='upper center', bbox_to_anchor=(1.10, 0.90), fontsize=legendsize, frameon=False) + + if pad_bottom is not None: + bigax = fig.add_subplot(111) + bigax.set_facecolor([1,1,1,0]) + bigax.spines['top'].set_visible(False) + bigax.spines['bottom'].set_visible(True) + bigax.spines['left'].set_visible(False) + bigax.spines['right'].set_visible(False) + bigax.tick_params(labelcolor='w', color='w', direction='in', top=False, bottom=True, left=False, right=False, labelleft=False, pad=pad_bottom) + + if mode == 'bars': + + + ax.tick_params(axis='both', direction='in', which='major', length=majorticklength, width=axeswidth, right=True, top=True) + ax.tick_params(axis='both', direction='in', which='minor', length=minorticklength, width=axeswidth, right=True, top=True) + + if not hide_ylabels: + ax.set_ylabel('Energy [{}]'.format(energyunit), size=labelsize, labelpad=ypad) + + ax.yaxis.set_major_locator(MultipleLocator(.1)) + ax.yaxis.set_minor_locator(MultipleLocator(.05)) + + ax.tick_params(axis='x', pad=xpad) + ax.tick_params(axis='y', pad=ypad) + + plt.xticks(fontsize=ticksize) + plt.yticks(fontsize=ticksize) + + if pad_bottom is not None: + bigax = fig.add_subplot(111) + bigax.set_facecolor([1,1,1,0]) + bigax.spines['top'].set_visible(False) + bigax.spines['bottom'].set_visible(True) + bigax.spines['left'].set_visible(False) + bigax.spines['right'].set_visible(False) + bigax.tick_params(labelcolor='w', color='w', direction='in', top=False, bottom=True, left=False, right=False, labelleft=False, pad=pad_bottom) + + if mode == 'both': + + # Set labels on x- and y-axes + ax[0].set_xlabel('Volume [{}]'.format(volumeunit), size=labelsize, labelpad=xpad) + ax[0].set_ylabel('Energy [{}]'.format(energyunit), size=labelsize, labelpad=ypad) + + # Set tick parameters + ax[0].tick_params(axis='both', direction='in', which='major', length=majorticklength, width=axeswidth, right=True, left=True, top=True, labelsize=ticksize) + ax[0].tick_params(axis='both', direction='in', which='minor', length=minorticklength, width=axeswidth, right=True, left=True, top=True, labelsize=ticksize) + + ax[0].tick_params(axis='x', pad=xpad) + ax[0].tick_params(axis='y', pad=ypad) + + ax[0].xaxis.set_major_locator(MultipleLocator(10)) + ax[0].xaxis.set_minor_locator(MultipleLocator(5)) + + ax[0].yaxis.set_major_locator(MultipleLocator(.1)) + ax[0].yaxis.set_minor_locator(MultipleLocator(.05)) + + plt.xticks(fontsize=ticksize) + plt.yticks(fontsize=ticksize) + + + ax[1].yaxis.set_major_locator(MultipleLocator(.2)) + ax[1].yaxis.set_minor_locator(MultipleLocator(.1)) + ax[1].yaxis.set_label_position('right') + ax[1].yaxis.tick_right() + ax[1].set_ylabel('Energy [{}]'.format(energyunit), size=labelsize, ypad=ypad) + ax[1].tick_params(axis='both', direction='in', which='major', length=majorticklength, width=axeswidth, left=True, right=True, top=True) + ax[1].tick_params(axis='both', direction='in', which='minor', length=minorticklength, width=axeswidth, left=True, right=True, top=True) + + ax[1].tick_params(axis='x', pad=xpad) + ax[1].tick_params(axis='y', pad=ypad) + + + plt.xticks(fontsize=ticksize) + plt.yticks(fontsize=ticksize) + + return fig, ax + + + +def parabola(V, a, b, c): + """parabola polynomial function + + this function is used to fit the data to get good guesses for + the equation of state fits + + a 4th order polynomial fit to get good guesses for + was not a good idea because for noisy data the fit is too wiggly + 2nd order seems to be sufficient, and guarantees a single minimum""" + + + E = (a * V**2) + (b * V) + c + + return E + + +def murnaghan(V, E0, V0, B0, BP): + 'From PRB 28,5480 (1983' + + E = E0 + ((B0 * V) / BP) * (((V0 / V)**BP) / (BP - 1) + 1) - ((V0 * B0) / (BP - 1)) + return E + + +def birch(V, E0, V0, B0, BP): + """ + From Intermetallic compounds: Principles and Practice, Vol. I: Principles + Chapter 9 pages 195-210 by M. Mehl. B. Klein, D. Papaconstantopoulos + paper downloaded from Web + + case where n=0 + """ + + E = (E0 + + 9 / 8 * B0 * V0 * ((V0 / V)**(2 / 3) - 1)**2 + + 9 / 16 * B0 * V0 * (BP - 4) * ((V0 / V)**(2 / 3) - 1)**3) + return E + + +def birchmurnaghan(V, E0, V0, B0, BP): + """ + BirchMurnaghan equation from PRB 70, 224107 + Eq. (3) in the paper. Note that there's a typo in the paper and it uses + inversed expression for eta. + """ + + eta = (V0 / V)**(1 / 3) + + E = E0 + 9 * B0 * V0 / 16 * (eta**2 - 1)**2 * (6 + BP * (eta**2 - 1) - 4 * eta**2) + + return E + + +def vinet(V, E0, V0, B0, BP): + 'Vinet equation from PRB 70, 224107' + + eta = (V / V0)**(1 / 3) + + E = (E0 + 2 * B0 * V0 / (BP - 1)**2 * + (2 - (5 + 3 * BP * (eta - 1) - 3 * eta) * + np.exp(-3 * (BP - 1) * (eta - 1) / 2))) + + return E + +def pouriertarantola(V, E0, V0, B0, BP): + 'Pourier-Tarantola equation from PRB 70, 224107' + + eta = (V / V0)**(1 / 3) + squiggle = -3 * np.log(eta) + + E = E0 + B0 * V0 * squiggle**2 / 6 * (3 + squiggle * (BP - 2)) + return E + + + +def get_initial_guesses(volume, energy): + + p = np.polyfit(volume, energy, deg=2) + + a, b, c = p[0], p[1], p[2] + + # Estimated from dE/dV = 2aV0 + b => V0 = -b / 2a + v0 = -b / (2*a) + + # Estimated by evaluating a parabola with a, b and c values at V = V0 + e0 = parabola(v0, a, b, c) + + # Estimated form B0 ~ V0 * d^2E / dV^2. d^2E / dV^2 = 2a. + b0 = 2 * a * v0 + + # Just a reasonable starting value + bp = 4 + + + return [e0, v0, b0, bp] + + + +def fit_eos_curve(volume, energy, p0, eos): + + eos_dict = {'murnaghan': murnaghan, 'birch': birch, 'birchmurnaghan': birchmurnaghan, 'vinet': vinet, 'pouriertarantola': pouriertarantola} + + func = eos_dict[eos] + + popt, pcov = curve_fit(func, volume, energy, p0) + + E0, V0, B0, BP = popt[0], popt[1], popt[2], popt[3] + + return [E0, V0, B0, BP] + + + + +def get_plotdata(volume, energy, equilibrium_values, eos): + + eos_dict = {'murnaghan': murnaghan, 'birch': birch, 'birchmurnaghan': birchmurnaghan, 'vinet': vinet, 'pouriertarantola': pouriertarantola} + + V = np.linspace(volume.min(), volume.max(), 100) + + E0, V0, B0, BP = equilibrium_values[0], equilibrium_values[1], equilibrium_values[2], equilibrium_values[3] + + print(E0, V0, B0, BP) + + func = eos_dict[eos] + + print(func) + + E = func(V, E0, V0, B0, BP) + + return E, V + + +def get_atoms(poscar): + + with open(poscar, 'r') as poscar: + lines = poscar.readlines() + + atoms = lines[5].split() + atom_num = lines[6].split() + + + atom_num = [int(num) for num in atom_num] + + atoms_dict = {} + + for ind, atom in enumerate(atoms): + atoms_dict[atom] = atom_num[ind] + + return atoms, atom_num, atoms_dict + + + +def get_equilibrium_data(path, atoms_per_formula_unit, eos=None): + + + if not eos: + eos = 'murnaghan' + + + dirs = [os.path.join(path, dir) for dir in os.listdir(path)] + + + + data = [] + + for dir in dirs: + atoms, atom_num, atoms_dict = get_atoms(os.path.join(dir, 'POSCAR')) + scaling_factor = sum(atom_num) / atoms_per_formula_unit + + label = dir.split('/')[-1] + + dft_df = pd.read_csv(os.path.join(dir, 'energ.dat'), header=None, delim_whitespace=True) + dft_df.columns = ['Volume', 'Energy'] + + volume = dft_df["Volume"].to_numpy() / scaling_factor + energy = dft_df["Energy"].to_numpy() / scaling_factor + + p0 = get_initial_guesses(volume, energy) + + equilibrium_constants = fit_eos_curve(volume, energy, p0, eos) + e0, v0, b0, bp = equilibrium_constants[0], equilibrium_constants[1], equilibrium_constants[2], equilibrium_constants[3] + + data.append([label, e0, v0, b0/kJ*1e24, bp]) + + + df = pd.DataFrame(data) + df.columns = ['Label', 'E0', 'V0', 'B0', 'Bp'] + df.sort_values(by='E0', ascending=True, inplace=True) + df.reset_index(inplace=True) + + E_min = df['E0'].min() + + df['dE'] = df['E0'] - E_min + + df = df[['Label', 'E0', 'dE', 'V0', 'B0', 'Bp']] + + + return df + + + +