<jupyter_start><jupyter_text>Select structure<jupyter_code>from apps.surfaces.structure_browser import StructureBrowser

atoms = None
slab_analyzed = None

def on_struct_change(c):
    global atoms, slab_analyzed
    structure = struct_browser.results.value
    if structure:
        atoms = structure.get_ase()
        atoms.pbc = [1, 1, 1]
        
        unobserve_cell_text()
        cell_x_text.value = atoms.cell[0,0]
        cell_y_text.value = atoms.cell[1,1]
        cell_z_text.value = atoms.cell[2,2]
        observe_cell_text()
        
        slab_analyzed = find_mol.analyze_slab(atoms)
        update_view(atoms, slab_analyzed)

def on_cell_change(c):
    global atoms
    atoms.cell = np.diag(np.array([
        cell_x_text.value,
        cell_y_text.value,
        cell_z_text.value
    ]))
    atoms.center()
    update_view(atoms, slab_analyzed)
    
def observe_cell_text():
    cell_x_text.observe(on_cell_change, names='value')
    cell_y_text.observe(on_cell_change, names='value')
    cell_z_text.observe(on_cell_change, names='value')
    
def unobserve_cell_text():
    cell_x_text.unobserve(on_cell_change, names='value')
    cell_y_text.unobserve(on_cell_change, names='value')
    cell_z_text.unobserve(on_cell_change, names='value')
    
struct_browser = StructureBrowser()
struct_browser.results.observe(on_struct_change, names='value')
viewer = nglview.NGLWidget()
clear_output()

cell_x_text = ipw.FloatText(description='cell x', value=0.0)
cell_y_text = ipw.FloatText(description='cell y', value=0.0)
cell_z_text = ipw.FloatText(description='cell z', value=0.0)

observe_cell_text()

display(ipw.VBox([struct_browser, viewer, cell_x_text, cell_y_text, cell_z_text]))

def update_view(atoms, slab_analyzed):
    
    details = slab_analyzed
    
    # delete all old components
    while hasattr(viewer, "component_0"):
        viewer.component_0.clear_representations()
        cid = viewer.component_0.id
        viewer.remove_component(cid)
    
    if atoms:
        
        mol_inds = [item for sublist in details['all_molecules'] for item in sublist]
        rest_inds = details['slabatoms']+details['bottom_H']+details['adatoms'] +details['unclassified']
        
        molecules_ase = atoms[mol_inds]
        rest_ase = atoms[rest_inds]
        
        # component 0: Molecule
        viewer.add_component(nglview.ASEStructure(molecules_ase), default_representation=False)
        #viewer.add_ball_and_stick(aspectRatio=1.8, radius=0.25, opacity=1.0, component=0)
        viewer.add_ball_and_stick(aspectRatio=1.5, scale=2.0, opacity=1.0, component=0)
        
        # component 1: Everything else
        viewer.add_component(nglview.ASEStructure(rest_ase), default_representation=False)
        viewer.add_ball_and_stick(aspectRatio=10.0, opacity=1.0, component=1)
        
        viewer.add_unitcell()
        viewer.center()
        
        # Orient camera to look from positive z
        cell_z = atoms.cell[2, 2]
        com = atoms.get_center_of_mass()
        def_orientation = viewer._camera_orientation
        top_z_orientation = [1.0, 0.0, 0.0, 0,
                             0.0, 1.0, 0.0, 0,
                             0.0, 0.0, -np.max([cell_z, 30.0]) , 0,
                             -com[0], -com[1], -com[2], 1]
        viewer._set_camera_orientation(top_z_orientation)
        #viewer.camera = 'orthographic'<jupyter_output><empty_output><jupyter_text>Select computer and codes<jupyter_code>qb = QueryBuilder()
qb.append(Computer, filters={'enabled': True}, project='name')
computer_names = [comp[0] for comp in qb.all()]

style = {'description_width': '120px'}
layout = {'width': '70%'}
drop_computer = ipw.Dropdown(description="Computer",
                             options=computer_names)

def comp_plugin_codes(computer_name, plugin_name):
    qb = QueryBuilder()
    qb.append(Computer, filters={'enabled': True}, project='name', tag='computer')
    qb.append(Code, project='*', has_computer='computer', filters={
        'attributes.input_plugin': plugin_name,
        'or': [{'extras': {'!has_key': 'hidden'}}, {'extras.hidden': False}]
    })
    qb.order_by({Code: {'id': 'desc'}})
    codes = qb.all()
    sel_codes = []
    for code in codes:
        if code[0] == computer_name:
            sel_codes.append(code[1])
    return sel_codes

def on_computer_change(c):
    global cp2k_codes, stm_codes
    cp2k_codes = comp_plugin_codes(drop_computer.value, 'cp2k')
    stm_codes = comp_plugin_codes(drop_computer.value, 'spm.stm')
    
    drop_cp2k.options = [c.label for c in cp2k_codes]
    drop_stm.options = [c.label for c in stm_codes]
    
    
drop_computer.observe(on_computer_change)

drop_cp2k = ipw.Dropdown(description="Cp2k code")

drop_stm = ipw.Dropdown(description="STM code")

on_computer_change(0)

elpa_check = ipw.Checkbox(
    value=True,
    description='use ELPA',
    disabled=False
)

display(drop_computer, drop_cp2k, drop_stm, elpa_check)<jupyter_output><empty_output><jupyter_text>Scanning tunnelling microscopy parameters<jupyter_code>style = {'description_width': '140px'}
layout = {'width': '50%'}
layout_small = {'width': '25%'}

elim_float_slider = ipw.FloatRangeSlider(
    value=[-2.0, 2.0],
    min=-4.0,
    max=4.0,
    step=0.1,
    description='Emin, Emax (eV):',
    disabled=False,
    continuous_update=False,
    orientation='horizontal',
    readout=True,
    readout_format='.1f',
    style=style, layout=layout)

de_floattext = ipw.BoundedFloatText(
                        description='dE (eV)',
                        min=0.01,
                        max=0.10,
                        step=0.01,
                        value=0.05,
                        style=style, layout=layout_small)

fwhms_text = ipw.Text(
                  description='FWHMs (eV)',
                  value='0.10',
                  style=style, layout=layout)

extrap_plane_floattext = ipw.BoundedFloatText(
                        description='Extrap plane (ang)',
                        min=1.0,
                        max=10.0,
                        step=0.1,
                        value=3.5,
                        style=style, layout=layout_small)

const_height_text = ipw.Text(description='Const. H (ang)',
                              value='3.0 5.0',
                              style=style, layout=layout)

const_current_text = ipw.Text(description='Const. cur. (isoval)',
                              value='1e-7',
                              style=style, layout=layout)

display(elim_float_slider, de_floattext, fwhms_text, extrap_plane_floattext, const_height_text, const_current_text)
def find_struct_wf(structure):
    qb = QueryBuilder()
    qb.append(Computer, filters={'enabled': True})
    selected_comp = qb.all()[drop_computer.index][0]
    
    # check stm
    extras = structure.get_extras()
    for ex_k in extras.keys():
        if ex_k.startswith(('stm', 'pdos', 'afm')):
            spm_workchain = load_node(extras[ex_k])
            cp2k_scf_calc = spm_workchain.get_outputs()[0]
            if cp2k_scf_calc.get_computer().hostname == selected_comp.hostname:
                wfn_path = cp2k_scf_calc.out.remote_folder.get_remote_path() + "/aiida-RESTART.wfn"
                # check if it exists
                file_exists = ! ssh {selected_comp.hostname} "if [ -f {wfn_path} ]; then echo 1 ; else echo 0 ; fi"
                if file_exists[0]:
                    print("Found .wfn from %s"%ex_k)
                    return wfn_path
    
    # check geo opt
    geo_opt_calc = structure.get_inputs()[0]
    geo_comp = geo_opt_calc.get_computer()
    if geo_comp.hostname == selected_comp.hostname:
        wfn_path = geo_opt_calc.out.remote_folder.get_remote_path() + "/aiida-RESTART.wfn"
        # check if it exists
        file_exists = ! ssh {selected_comp.hostname} "if [ -f {wfn_path} ]; then echo 1 ; else echo 0 ; fi"
        if file_exists[0]:
            print("Found .wfn from geo_opt")
            return wfn_path
    return ""

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
        
        extrap_plane = extrap_plane_floattext.value
        parent_dir = "parent_calc_folder/"
        
        energy_range_str = "%.2f %.2f %.3f" % (
            elim_float_slider.value[0], elim_float_slider.value[1], de_floattext.value
        )

        stm_params = ParameterData(dict={
            '--cp2k_input_file':    parent_dir+'aiida.inp',
            '--basis_set_file':     parent_dir+'BASIS_MOLOPT',
            '--xyz_file':           parent_dir+'geom.xyz',
            '--wfn_file':           parent_dir+'aiida-RESTART.wfn',
            '--hartree_file':       parent_dir+'aiida-HART-v_hartree-1_0.cube',
            '--output_file':        'stm.npz',
            '--eval_region':        ['G', 'G', 'G', 'G', 'n-2.0_C', 'p%.1f'%extrap_plane],
            '--dx':                 '0.15',
            '--eval_cutoff':        '14.0',
            '--extrap_extent':      '5.0',
            '--energy_range':       energy_range_str.split(),
            '--heights':            const_height_text.value.split(),
            '--isovalues':          const_current_text.value.split(),
            '--fwhms':              fwhms_text.value.split(),
        })
        
        cp2k_code = cp2k_codes[drop_cp2k.index]
        stm_code = stm_codes[drop_stm.index]
        
        struct = struct_browser.results.value
        
        cell = ArrayData()
        cell.set_array('cell', np.diag(atoms.cell))
        
       ## Try to access the restart-wfn file ##
        try:
            wfn_file_path = find_struct_wf(struct)
        except:
            wfn_file_path = ""
        if wfn_file_path == "":
            print("Didn't find any accessible .wfn file.")
    
        
        outputs = submit(STMWorkChain,
                 cp2k_code=cp2k_code,
                 structure=struct,
                 cell=cell,
                 wfn_file_path=Str(wfn_file_path),
                 elpa_switch=Bool(elpa_check.value),
                 stm_code=stm_code,
                 stm_params=stm_params
                )
                
        print(outputs)

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.11969552806850618, 'fraction_numerical': 0.016175071360608945, 'mean_word_length': 2.3283723875870805, 'pattern_counts': {'":': 1, '<': 13, '<?xml version=': 0, '>': 13, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '3', 'dir': 'jupyter-structured-clean-dedup', 'id': '609636', 'license': 'MIT', 'n_tokens_mistral': 3589, 'n_tokens_neox': 3469, 'n_words': 570, 'path': 'stm/submit_stm.ipynb', 'repo_name': 'hifabian/aiidalab-empa-scanning-probe', 'size': '10510'}