<jupyter_start><jupyter_text>Detailed Report<jupyter_code>from aiida import load_dbenv, is_dbenv_loaded
from aiida.backends import settings
if not is_dbenv_loaded():
    load_dbenv(profile=settings.AIIDADB_PROFILE)

from aiida.orm import load_node
from aiida.orm.querybuilder import QueryBuilder
from aiida.orm.calculation.work import WorkCalculation
from aiida.orm.calculation.job import JobCalculation

import numpy as np
import scipy.constants as const
import bqplot as bq
import ipywidgets as ipw
from IPython.display import display, clear_output, HTML
import re
import gzip
import matplotlib.pyplot as plt
from collections import OrderedDict
import urlparse
import io

from ase.data import covalent_radii, atomic_numbers
from ase.data.colors import cpk_colors
from ase.neighborlist import NeighborList

from IPython.display import FileLink
from base64 import b64encode
def get_calc_by_label(workcalc, label):
    qb = QueryBuilder()
    qb.append(WorkCalculation, filters={'uuid':workcalc.uuid})
    qb.append(JobCalculation, output_of=WorkCalculation, filters={'label':label})
    assert qb.count() == 1
    calc = qb.first()[0]
    assert(calc.get_state() == 'FINISHED')
    return calc
def get_calcs_by_label(workcalc, label):
    qb = QueryBuilder()
    qb.append(WorkCalculation, filters={'uuid':workcalc.uuid})
    qb.append(JobCalculation, output_of=WorkCalculation, filters={'label':label})
    calcs = [ n[0] for n in qb.all() ] 
    for calc in calcs: 
        assert(calc.get_state() == 'FINISHED')
    return calcs
url = urlparse.urlsplit(jupyter_notebook_url)
pk = urlparse.parse_qs(url.query)['pk'][0]
workcalc = load_node(pk=int(pk))

orbitals_calcs = get_calcs_by_label(workcalc, "export_orbitals")
bands_calc = get_calc_by_label(workcalc, "bands")
structure = bands_calc.inp.structure
ase_struct = structure.get_ase()

bands = bands_calc.out.output_band.get_bands()
if bands.ndim == 2:
    bands = bands[None,:,:]

nspins, nkpoints, nbands = bands.shape

vacuum_level = workcalc.get_extra('vacuum_level')
fermi_energy = workcalc.get_extra('fermi_energy')
total_energy = workcalc.get_extra('total_energy')
homo = workcalc.get_extra('homo')
lumo = workcalc.get_extra('lumo')
gap = workcalc.get_extra('gap')
abs_mag = workcalc.get_extra('absolute_magnetization_per_angstr')
tot_mag = workcalc.get_extra('total_magnetization_per_angstr')

print("WorkCalculation PK: %i"%workcalc.pk)
print("total energy: %.3f eV"%total_energy)
print("gap: %.3f eV"%gap)
print("total magentization/A: %.3f"%abs_mag)
print("abs. magentization/A: %.3f"%tot_mag)
def plot_spin(ispin):
    
    center = (homo + lumo)/2.0
    x_sc = bq.LinearScale()
    y_sc = bq.LinearScale(min=center-3.0, max=center+3.0, )
    
    color_sc = bq.ColorScale(colors=['gray', 'red'], min=0.0, max=1.0)
    colors = np.zeros(nbands)
    
    Lx = structure.cell_lengths[0]
    x_max = np.pi / Lx
    ax_x = bq.Axis(label=u'kA^-1', scale=x_sc, grid_lines='solid', tick_format='.3f', tick_values=[0, x_max]) #, tick_values=[0.0, 0.5])
    ax_y = bq.Axis(label='eV', scale=y_sc, orientation='vertical', grid_lines='solid')
    
    x_data = np.linspace(0.0, x_max, nkpoints)
    y_datas = bands[ispin,:,:].transpose() - vacuum_level
    
    lines = bq.Lines(x=x_data, y=y_datas, color=colors, animate=True,
                     scales={'x': x_sc, 'y': y_sc, 'color': color_sc})

    homo_line = bq.Lines(x=[0, x_max], y=[homo, homo], line_style='dashed', colors=['red'], scales={'x': x_sc, 'y': y_sc})
    
    # Initialize the parabola as a random line and set visible to false
    # Later, when it is correctly set, show it.
    eff_mass_parabola = bq.Lines(x=[0, 0], y=[0, 0], visible=False, stroke_width=1.0,
                                 line_style='solid', colors=['blue'], scales={'x': x_sc, 'y': y_sc})
    
    ratio = 0.25
    layout = ipw.Layout(height="800px", width="200px")
    
    m_fig = dict(left=45, top=60, bottom=60, right=40)
    fig = bq.Figure(axes=[ax_x, ax_y], marks=[lines, homo_line, eff_mass_parabola], title='Spin %i'%ispin, 
                    layout=layout, fig_margin=m_fig,
                    min_aspect_ratio=ratio, max_aspect_ratio=ratio)

    def on_band_click(self, target):
        global selected_spin, selected_band 
        selected_spin = ispin
        selected_band = target['data']['index']
        on_band_change()

    lines.on_element_click(on_band_click)

    save_btn = ipw.Button(description="Download png")
    save_btn.on_click(lambda b: fig.save_png()) # save_png() does not work with unicode labels
    
    igor_link = mk_igor_link(ispin)
    
    layout = ipw.Layout(align_items="center", padding="5px", margin="0px")
    box = ipw.VBox([fig, save_btn, igor_link], layout=layout)
    return box, lines, eff_mass_parabola
def mk_igor_link(ispin):
    igorvalue = igor_bands(ispin)
    igorfile = b64encode(igorvalue)
    filename = ase_struct.get_chemical_formula() + "_bands_spin%d_pk%d.itx" % (ispin, structure.pk)

    html = '<a download="{}" href="'.format(filename)
    html += 'data:chemical/x-igor;name={};base64,{}"'.format(filename, igorfile)
    html += ' id="pdos_link"'
    html += ' target="_blank">Export itx-Bands</a>'

    return ipw.HTML(html)

def igor_bands(ispin):
    k_axis = np.linspace(0.0, np.pi / structure.cell_lengths[0], nkpoints)
    testio = io.StringIO()
    tosave = bands[ispin,:,:].transpose() -vacuum_level
    #with open('igor_bands-'+pk+'.itx', 'w') as f:
    with testio as f:
        f.write(u'IGOR\r')
        f.write(u'WAVES')
        f.write(u'\tx1'+(u'\ty{}'*nbands).format(*[x for x in range(nbands)])+u'\r')
        f.write(u'BEGIN\r')
        for i in range(nkpoints):
            f.write(u"\t{:.7f}".format(k_axis[i])) # first column k_axis
            f.write((u"\t{:.7f}"*nbands).format(*tosave[:,i])) # other columns the bands
            f.write(u"\r")
        f.write(u"END\r")
        f.write(u'X SetScale/P x {},{},"", x1; SetScale y 0,0,"", x1\r'.format(0, k_axis[1]-k_axis[0]))
        for idk in range(nbands):
            f.write((u'X SetScale/P x 0,1,"", y{0}; SetScale y 0,0,"", y{0}\r').format(str(idk)))
        return testio.getvalue()
def read_cube(fn):
    lines = gzip.open(fn).readlines()
    header = np.fromstring("".join(lines[2:6]), sep=' ').reshape(4,4)
    #print(header)
    natoms, nx, ny, nz = header[:,0].astype(int)
    cube = dict()
    cube['x0'] = header[0,1] # x origin
    cube['y0'] = header[0,2] # y origin
    cube['z0'] = header[0,3] # z origin
    cube['dx'] = header[1,1] # x step size
    cube['dy'] = header[2,2] # y step size
    cube['dz'] = header[3,3] # z step size
    cube['data'] = np.fromstring("".join(lines[natoms+6:]), sep=' ').reshape(nx, ny, nz)
    return cube
def on_band_change():
    global selected_cube_files
    with info_out:
        clear_output()
        print("selected spin: %d"%selected_spin)
        print("selected band: %d"%selected_band)

        colors = np.zeros((nspins, nbands))
        colors[selected_spin, selected_band] = 1.0
        for ispin in range(nspins):
            band_plots[ispin].color = colors[ispin,:]
        
        # orbitals_calcs might use fewer nkpoints than bands_calc
        prev_calc = orbitals_calcs[0].inp.parent_calc_folder.inp.remote_folder
        nkpoints_lowres = prev_calc.res.number_of_k_points
        
        lower = nkpoints_lowres * selected_spin
        upper = lower + nkpoints_lowres
        selected_cube_files = []
        for fn in sorted([ fdr for orbitals_calc in orbitals_calcs for fdr in orbitals_calc.out.retrieved.get_folder_list() ]):
            m = re.match("aiida.filplot_K(\d\d\d)_B(\d\d\d)_orbital.cube.gz", fn)
            if not m:
                continue
            k, b = int(m.group(1)), int(m.group(2))
            if b != selected_band + 1:
                continue
            if lower < k and k <= upper:
                selected_cube_files.append(fn)

        n = len(selected_cube_files)
        kpoint_slider.max = max(n, 1)
        print("found %d cube files"%n)
        on_kpoint_change(None)
        
        ### -------------------------------------------
        ### Effective mass calculation and parabola plotting
        
        meff, parabola_fit, fit_kvals, fit_energies = calc_effective_mass(ispin=selected_spin)
        print("effective mass: %f"%meff)
        
        parab_k_arr = np.linspace(np.min(fit_kvals), np.max(fit_kvals), 20)
        parab_e_arr = parabola_fit[0]*parab_k_arr**2 + parabola_fit[1]*parab_k_arr + parabola_fit[2]
        eff_mass_parabolas[selected_spin].x = parab_k_arr
        eff_mass_parabolas[selected_spin].y = parab_e_arr
        eff_mass_parabolas[selected_spin].visible = True
        
        if nspins > 1:
            eff_mass_parabolas[(selected_spin+1)%2].visible = False
        
        ### -------------------------------------------
        
def on_kpoint_change(c):
    global selected_cube
    with kpnt_out:
        clear_output()
        i = kpoint_slider.value
        if i > len(selected_cube_files):
            print("Found no cube files")
            selected_cube = None
            height_slider.options = {"---":0}
            
        else:    
            fn = selected_cube_files[i-1]
            print(fn)
            for orbitals_calc in  orbitals_calcs:
                try:
                    absfn = orbitals_calc.out.retrieved.get_abs_path(fn)
                except:
                    continue
                selected_cube = read_cube(absfn)
                nz = selected_cube['data'].shape[2]
                z0 = selected_cube['z0']
                dz = selected_cube['dz']
        
                zmid = structure.cell_lengths[2] / 2.0
                options = OrderedDict()
                for i in range(nz):
                    z = (z0 + dz*i) * 0.529177 - zmid
                    options[u"%.3f Å"%z] = i
                height_slider.options = options
                break
        
            on_orb_plot_change(None)
def on_orb_plot_change(c):
    with orb_out:
        clear_output()
        if selected_cube is None:
            return

        fig, ax = plt.subplots()
        fig.dpi = 150.0
        vmin = 10 ** colormap_slider.value[0]
        vmax = 10 ** colormap_slider.value[1]
        
        cax = plot_cube(ax, selected_cube, height_slider.value, 'gray', vmin, vmax)
        fig.colorbar(cax, label='e/bohr^3', ticks=[vmin, vmax], format='%.0e', orientation='horizontal', shrink=0.3)
        
        plot_overlay_struct(ax, orb_alpha_slider.value)
        plt.show()
def plot_cube(ax, cube, z, cmap, vmin=-1, vmax=+1):
    assert cube['x0'] == 0.0 and cube['y0'] == 0.0
    
    a = np.flip(cube['data'][:,:,z].transpose(), axis=0)
    aa = np.tile(a, (1, 2))
    x2 = cube['dx'] * aa.shape[1] * 0.529177
    y2 = cube['dy'] * aa.shape[0] * 0.529177
    
    ax.set_xlabel(u'Å')
    ax.set_ylabel(u'Å')
    ax.set_xlim(0, x2)
    ax.set_ylim(0, y2)
    
    cax = ax.imshow(aa, extent=[0,x2,0,y2], cmap=cmap, vmin=vmin, vmax=vmax)
    return cax
def plot_overlay_struct(ax, alpha):
    if alpha == 0:
        return
    
    # plot overlayed structure
    s = ase_struct.repeat((2,1,1))
    cov_radii = [covalent_radii[a.number] for a in s]
    nl = NeighborList(cov_radii, bothways = True, self_interaction = False)
    nl.update(s)

    for at in s:
        #circles
        x,y,z = at.position
        n = atomic_numbers[at.symbol]
        ax.add_artist(plt.Circle((x,y), covalent_radii[n]*0.5, color=cpk_colors[n], fill=True, clip_on=True, alpha=alpha))
        #bonds
        nlist = nl.get_neighbors(at.index)[0]
        for theneig in nlist:
            x,y,z = (s[theneig].position +  at.position)/2
            x0,y0,z0 = at.position
            if (x-x0)**2 + (y-y0)**2 < 2 :
                ax.plot([x0,x],[y0,y],color=cpk_colors[n],linewidth=2,linestyle='-', alpha=alpha)
band_plots = []
boxes = []
eff_mass_parabolas = []
for ispin in range(nspins):
    box, plot, eff_mass_parabola = plot_spin(ispin)
    boxes.append(box)
    band_plots.append(plot)
    eff_mass_parabolas.append(eff_mass_parabola)

layout = ipw.Layout(padding="5px", margin="0px")
info_out = ipw.Output(layout=layout)
kpnt_out = ipw.Output(layout=layout)
orb_out = ipw.Output(layout=layout)

layout = ipw.Layout(width="400px")

### -----------------------------
### Slider to control how many points of the band to use for parabolic fit

# Odd values of fit have better accuracy, so it's worth it to disable even values
efm_fit_slider = ipw.IntSlider(description="eff. mass fit", min=3, max=15, step=2, continuous_update=False, layout=layout)
# Only if a band is selected, selecting a new effective mass fit will update the plot and infopanel
on_efm_fit_change = lambda c: on_band_change() if 'selected_spin' in globals() else None
efm_fit_slider.observe(on_efm_fit_change, names='value')
### -----------------------------

kpoint_slider = ipw.IntSlider(description="k-point", min=1, max=1, continuous_update=False, layout=layout)
kpoint_slider.observe(on_kpoint_change, names='value')

height_slider = ipw.SelectionSlider(description="height", options={"---":0}, continuous_update=False, layout=layout)
height_slider.observe(on_orb_plot_change, names='value')

orb_alpha_slider = ipw.FloatSlider(description="opacity", value=0.5, max=1.0, continuous_update=False, layout=layout)
orb_alpha_slider.observe(on_orb_plot_change, names='value')

colormap_slider = ipw.IntRangeSlider(description='colormap', min=-10, max=-1, value=[-6, -3], continuous_update=False, layout=layout)
colormap_slider.observe(on_orb_plot_change, names='value')

layout = ipw.Layout(align_items="center")
side_box = ipw.VBox([info_out, efm_fit_slider, kpoint_slider, height_slider, orb_alpha_slider, colormap_slider, 
                     kpnt_out, orb_out], layout=layout)
boxes.append(side_box)
display(ipw.HBox(boxes))<jupyter_output><empty_output><jupyter_text>Spin Density<jupyter_code>try:
    spinden_calc = get_calc_by_label(workcalc, "export_spinden")
except:
    spinden_calc = None
    print("Could not find spin density")
    
if spinden_calc:
    fn = spinden_calc.out.retrieved.get_abs_path("_spin.cube.gz")
    spinden_cube = read_cube(fn)
    spinden_cube['data'] *= 2000 # normalize scale
    def on_spinden_plot_change(c):
        with spinden_out:
            clear_output()
            fig, ax = plt.subplots()
            fig.dpi = 150.0
            cax = plot_cube(ax, spinden_cube, 1, 'seismic')
            fig.colorbar(cax,  label='arbitrary unit')
            plot_overlay_struct(ax, spinden_alpha_slider.value)
            plt.show()
        
    spinden_alpha_slider = ipw.FloatSlider(description="opacity", value=0.5, max=1.0, continuous_update=False)
    spinden_alpha_slider.observe(on_spinden_plot_change, names='value')
    spinden_out = ipw.Output()
    display(spinden_out, spinden_alpha_slider)

    on_spinden_plot_change(None)
## m*
def calc_effective_mass(ispin):
    # m* = hbar^2*[d^2E/dk^2]^-1
    hbar = const.value('Planck constant over 2 pi in eV s')
    el_mass = const.m_e*1e-20/const.eV # in eV*s^2/ang^2
    
    band = bands[ispin].transpose()[selected_band] - vacuum_level
    k_axis = np.linspace(0.0, np.pi / structure.cell_lengths[0], nkpoints)
    
    num_fit_points = efm_fit_slider.value
    
    if np.amax(band)>=lumo:
        # conduction band, let's search for effective electron mass (lowest point in energy)
        parabola_ind = np.argmin(band)
    else:
        # valence band, effective hole mass (highest point in energy)
        parabola_ind = np.argmax(band)
    
    # extend band and k values to neighbouring regions
    band_ext = np.concatenate([np.flip(band, 0)[:-1], band, np.flip(band, 0)[1:]])
    k_vals_ext = np.concatenate([-np.flip(k_axis, 0)[:-1], k_axis, k_axis[-1] + k_axis[1:]])
    
    # define fitting region
    i_min = parabola_ind - int(np.ceil(num_fit_points/2.0)) + len(band)
    i_max = parabola_ind + int(np.floor(num_fit_points/2.0)) + len(band)
    
    fit_energies = band_ext[i_min:i_max]
    fit_kvals = k_vals_ext[i_min:i_max]
    
    #print(k_axis[parabola_ind], band[parabola_ind])
    #print(fit_kvals)
    #print(fit_energies)
    
    parabola_fit = np.polyfit(fit_kvals, fit_energies, 2)
    
    meff = hbar**2/(2*parabola_fit[0])/el_mass
    
    # restrict fitting values to "main region"
    main_region_mask = (fit_kvals >= k_axis[0]) & (fit_kvals <= k_axis[-1])
    fit_energies = fit_energies[main_region_mask]
    fit_kvals = fit_kvals[main_region_mask]
    
    return meff, parabola_fit, fit_kvals, fit_energies<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.14310719459198454, 'fraction_numerical': 0.01635683244809271, 'mean_word_length': 3.093132411067194, 'pattern_counts': {'":': 2, '<': 15, '<?xml version=': 0, '>': 15, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 3, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '2', 'dir': 'jupyter-structured-clean-dedup', 'id': '495543', 'license': 'MIT', 'n_tokens_mistral': 6187, 'n_tokens_neox': 6018, 'n_words': 1133, 'path': 'nanoribbon/show.ipynb', 'repo_name': 'eimrek/mc-empa-surfaces', 'size': '16568'}