<jupyter_start><jupyter_text>Submit Nanoribbon Workchain<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.code import Code
from aiida.orm import load_node
from aiida.orm import Code, Computer
from aiida.orm.data.base import Int, Float, Str
from aiida.orm.querybuilder import QueryBuilder
from aiida.work.run import run, submit
from aiida.orm.data.structure import StructureData

import nglview
import numpy as np
import ipywidgets as ipw
from collections import OrderedDict
from IPython.display import display, clear_output, HTML
import nglview

from apps.surfaces.structure_browser import StructureBrowser
from apps.surfaces.nanoribbon.nanoribbonwork import NanoribbonWorkChain
def on_struct_change(c):
    refresh_structure_view()
    
struct_browser = StructureBrowser()
struct_browser.results.observe(on_struct_change, names='value')    
viewer = nglview.NGLWidget()

clear_output()
display(ipw.VBox([struct_browser, viewer]))
def refresh_structure_view():
    global viewer, atoms
    if hasattr(viewer, "component_0"):
        #viewer.clear_representations()
        viewer.component_0.remove_ball_and_stick()
        viewer.component_0.remove_ball_and_stick()
        viewer.component_0.remove_ball_and_stick()
        viewer.component_0.remove_unitcell()
        cid = viewer.component_0.id
        viewer.remove_component(cid)
    
    node = struct_browser.results.value
    if not node:
        return
    atoms = node.get_ase()
    viewer.add_component(nglview.ASEStructure(atoms)) # adds ball+stick
    viewer.add_unitcell()
    viewer.center_view()

    spin_ups = ",".join([str(i) for i, k in enumerate(node.get_site_kindnames()) if k.endswith("1")])
    spin_downs = ",".join([str(i) for i, k in enumerate(node.get_site_kindnames()) if k.endswith("2")])

    viewer.add_representation('ball+stick', selection="@"+spin_ups, color='red', aspectRatio=4.0, opacity=1.0)
    viewer.add_representation('ball+stick', selection="@"+spin_downs, color='green', aspectRatio=4.0, opacity=1.0)
def search_qe_codes(computer):

    qb = QueryBuilder()
    qb.append(Code)
    all_codes = qb.all()
    
    code_order = ['pw', 'pp', 'projwfc']
    
    qe_code_groups = OrderedDict()

    for code in all_codes:
        c = code[0]
        if c.get_computer().get_name() == computer.get_name():
            if any(x in c.label for x in code_order):

                # remove pw, pp or projwfc from name
                for substr in code_order:
                    if substr in c.label:
                        code_type = substr
                        group_label = (c.label).replace(substr, '')
                        break

                if group_label not in qe_code_groups:
                    qe_code_groups[group_label] = [None, None, None]

                qe_code_groups[group_label][code_order.index(code_type)] = c
    
    # update labels to list all codes
    old_keys = qe_code_groups.keys()
    for old_key in old_keys:
        if None in qe_code_groups[old_key]:
            qe_code_groups.pop(old_key)
            continue
        new_key = ""
        for code in qe_code_groups[old_key]:
            if new_key != "":
                new_key += ", "
            new_key += code.label
        qe_code_groups[new_key] = qe_code_groups.pop(old_key)
        
    return qe_code_groups
def update_code_selection(c):
    qe_code_groups.clear()
    qe_code_groups.update(search_qe_codes(all_computers[drop_computer.index][0]))
    drop_codes.options = qe_code_groups.keys()
    drop_codes.value = qe_code_groups.keys()[0]

qb = QueryBuilder()
qb.append(Computer, filters={'enabled': True})
all_computers = qb.all()

computer_selection = [comp[0].get_name() for comp in all_computers]
style = {'description_width': '120px'}
layout = {'width': '70%'}
drop_computer = ipw.Dropdown(description="Computer",
                             options=computer_selection,
                             style=style, layout=layout)

drop_computer.observe(update_code_selection)

qe_code_groups = search_qe_codes(all_computers[drop_computer.index][0])

drop_codes = ipw.Dropdown(description="Codes",
                             options = qe_code_groups.keys(),
                             style=style, layout=layout)

slider_precision = ipw.FloatSlider(description='Precision:',
                                   value=1.0, min=0.1, max=3.0, step=0.1,
                                   style=style, layout=layout)
text_calc_name = ipw.Text(description='Calculation Name: ',
                          placeholder='A great name.',
                          style=style, layout=layout)
display(ipw.VBox([drop_computer, drop_codes, slider_precision, text_calc_name]))

def on_submit(b):
    with submit_out:
        clear_output()
        if not struct_browser.results.value:
            print("Please select a structure.")
            return
        if not drop_computer.value:
            print("Please select a computer.")
            return
        if len(text_calc_name.value) < 3:
            print("Please enter a calculation name.")
            return

        pw_code, pp_code, projwfc_code = qe_code_groups[drop_codes.value]
                
        struct = struct_browser.results.value
        precision = Float(slider_precision.value)
        outputs = submit(NanoribbonWorkChain,
                         pw_code=pw_code,
                         pp_code=pp_code,
                         projwfc_code=projwfc_code,
                         structure=struct,
                         precision=precision)
        print(outputs)
        the_workcalc = load_node(outputs.pid)
        the_workcalc.description = text_calc_name.value

btn_submit = ipw.Button(description="Submit")
btn_submit.on_click(on_submit)
submit_out = ipw.Output()
display(btn_submit, submit_out)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.09774562584118439, 'fraction_numerical': 0.0058882907133243605, 'mean_word_length': 2.715625, 'pattern_counts': {'":': 1, '<': 6, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '395323', 'license': 'MIT', 'n_tokens_mistral': 1884, 'n_tokens_neox': 1818, 'n_words': 357, 'path': 'nanoribbon/submit.ipynb', 'repo_name': 'eimrek/mc-empa-surfaces', 'size': '5944'}