app.py

# Copyright 2024 D-Wave
#
#    Licensed under the Apache License, Version 2.0 (the "License");
#    you may not use this file except in compliance with the License.
#    You may obtain a copy of the License at
#
#        http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS,
#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#    See the License for the specific language governing permissions and
#    limitations under the License.

import dash
import dash_bootstrap_components as dbc
from dash import html, Input, Output, State
import datetime
import json
import numpy as np
import os

import dimod
from dwave.cloud import Client
from dwave.embedding import embed_bqm, is_valid_embedding
from dwave.system import DWaveSampler

from helpers.kz_calcs import *
from helpers.layouts_cards import *
from helpers.layouts_components import *
from helpers.plots import *
from helpers.qa import *
from helpers.tooltips import tool_tips

app = dash.Dash(__name__, external_stylesheets=[dbc.themes.BOOTSTRAP])

# Initialize: available QPUs, initial progress-bar status 
try:
    client = Client.from_config(client='qpu')
    qpus = {qpu.name: qpu for qpu in client.get_solvers(fast_anneal_time_range__covers=[0.005, 0.1])}
    if len(qpus) < 1:
        raise Exception    
    init_job_status = 'READY'
except Exception:
    qpus = {}
    client = None
    init_job_status = 'NO SOLVER'

# Dashboard-organization section
app.layout = dbc.Container([
    dbc.Row([                       # Top: logo
        dbc.Col([
            html.Img(
                src='assets/dwave_logo.png', 
                height='25px',
                style={'textAlign': 'left', 'margin': '10px 0px 15px 0px'}
            )
        ],
            width=3,
        )
    ]),
    dbc.Row([                        
        dbc.Col(                    # Left: control panel
            [
            control_card(
                solvers=qpus, 
                init_job_status=init_job_status
            ),
            *dbc_modal('modal_solver'),
            *[dbc.Tooltip(
            message, target=target, id=f'tooltip_{target}', style = dict())
            for target, message in tool_tips.items()]
            ],
            width=4,
            style={'minWidth': "30rem"},
        ),
        dbc.Col(                    # Right: display area
            graphs_card(),
            width=8,
            style={'minWidth': "60rem"},
        ),
    ]),
],
    fluid=True,
)

server = app.server
app.config['suppress_callback_exceptions'] = True

# Callbacks Section

@app.callback(
    Output('solver_modal', 'is_open'),
    Input('btn_simulate', 'n_clicks'),)
def alert_no_solver(dummy):
    """Notify if no quantum computer is accessible."""

    trigger_id = dash.callback_context.triggered[0]['prop_id'].split('.')[0]

    if trigger_id == 'btn_simulate':
        if not client:
            return True

    return False

@app.callback(
    Output('anneal_duration', 'disabled'),
    Output('coupling_strength', 'disabled'),
    Output('spins', 'options'),
    Output('qpu_selection', 'disabled'),
    Input('job_submit_state', 'children'),
    State('spins', 'options'))
def disable_buttons(job_submit_state, spins_options):        
    """Disable user input during job submissions."""

    trigger_id = dash.callback_context.triggered[0]['prop_id'].split('.')[0]

    if trigger_id !='job_submit_state':
        return dash.no_update, dash.no_update, dash.no_update, dash.no_update

    if job_submit_state in ['EMBEDDING', 'SUBMITTED', 'PENDING', 'IN_PROGRESS']:
        
        for inx, _ in enumerate(spins_options): 
            
            spins_options[inx]['disabled'] = True
        
        return  True, True, spins_options, True

    elif job_submit_state in ['COMPLETED', 'CANCELLED', 'FAILED']:

        for inx, _ in enumerate(spins_options): 
            spins_options[inx]['disabled'] = False
        
        return False, False, spins_options, False

    else:
        return dash.no_update, dash.no_update, dash.no_update, dash.no_update

@app.callback(
    Output('quench_schedule_filename', 'children'),
    Output('quench_schedule_filename', 'style'),
    Input('qpu_selection', 'value'),)
def set_schedule(qpu_name):
    """Set the schedule for the selected QPU."""

    trigger_id = dash.callback_context.triggered[0]['prop_id'].split('.')[0]

    schedule_filename = 'FALLBACK_SCHEDULE.csv'  
    schedule_filename_style = {'color': 'red', 'fontSize': 12}
 
    if trigger_id == 'qpu_selection':

        for filename in [file for file in os.listdir('helpers') if 
                         'schedule.csv' in file.lower()]:
                        
            if qpu_name.split('.')[0] in filename:  # Accepts & reddens older versions
            
                schedule_filename = filename

                if qpu_name in filename:

                    schedule_filename_style = {'color': 'white', 'fontSize': 12} 

    return schedule_filename, schedule_filename_style

@app.callback(
    Output('embeddings_cached', 'data'),
    Output('embedding_is_cached', 'value'),
    Input('qpu_selection', 'value'),
    Input('embeddings_found', 'data'),
    State('embeddings_cached', 'data'),)
def cache_embeddings(qpu_name, embeddings_found, embeddings_cached):
    """Cache embeddings for the selected QPU."""

    trigger_id = dash.callback_context.triggered[0]['prop_id'].split('.')[0]
 
    if trigger_id == 'qpu_selection':

        embeddings_cached = {}  # Wipe out previous QPU's embeddings
          
        for filename in [file for file in os.listdir('helpers') if 
                         '.json' in file and 'emb_' in file]:

            if qpu_name.split('.')[0] in filename:

                with open(f'helpers/{filename}', 'r') as fp:
                    embeddings_cached = json.load(fp)

                embeddings_cached = json_to_dict(embeddings_cached)
                            
                # Validate that loaded embeddings' edges are still available on the selected QPU
                for length in list(embeddings_cached.keys()):
                    
                    source_graph = dimod.to_networkx_graph(create_bqm(num_spins=length)).edges 
                    target_graph = qpus[qpu_name].edges
                    emb = embeddings_cached[length]

                    if not is_valid_embedding(emb, source_graph, target_graph):

                        del embeddings_cached[length]

    if trigger_id == 'embeddings_found':

        if not isinstance(embeddings_found, str): # embeddings_found != 'needed' or 'not found'

            embeddings_cached = json_to_dict(embeddings_cached)
            embeddings_found = json_to_dict(embeddings_found)
            new_embedding = list(embeddings_found.keys())[0]
            embeddings_cached[new_embedding] = embeddings_found[new_embedding]

        else:
            return dash.no_update, dash.no_update

    return embeddings_cached, list(embeddings_cached.keys())

@app.callback(
    Output('sample_vs_theory', 'figure'),
    Input('kz_graph_display', 'value'),
    Input('coupling_strength', 'value'),
    Input('quench_schedule_filename', 'children'),
    Input('job_submit_state', 'children'),
    State('job_id', 'children'),
    State('anneal_duration', 'min'),
    State('anneal_duration', 'max'),
    State('anneal_duration', 'value'),
    State('spins', 'value'),
    State('embeddings_cached', 'data'),
    State('sample_vs_theory', 'figure'),)
def display_graphics_kink_density(kz_graph_display, J, schedule_filename, \
    job_submit_state, job_id, ta_min, ta_max, ta, \
    spins, embeddings_cached, figure):
    """Generate graphics for kink density based on theory and QPU samples."""

    trigger_id = dash.callback_context.triggered[0]['prop_id'].split('.')[0]

    if trigger_id in ['kz_graph_display', 'coupling_strength', 'quench_schedule_filename'] :
        
        fig = plot_kink_densities_bg(kz_graph_display, [ta_min, ta_max], J, schedule_filename)

        return fig
    
    if trigger_id == 'job_submit_state':

        if job_submit_state == 'COMPLETED':

            embeddings_cached = embeddings_cached = json_to_dict(embeddings_cached)

            sampleset_unembedded = get_samples(client, job_id, spins, J, embeddings_cached[spins])              
            _, kink_density = kink_stats(sampleset_unembedded, J)
            
            fig = plot_kink_density(kz_graph_display, figure, kink_density, ta)
            return fig
        
        else:
            return dash.no_update
        
    fig = plot_kink_densities_bg(kz_graph_display, [ta_min, ta_max], J, schedule_filename)
    return fig

@app.callback(
    Output('spin_orientation', 'figure'),
    Input('spins', 'value'),
    Input('job_submit_state', 'children'),
    State('job_id', 'children'),
    State('coupling_strength', 'value'),
    State('embeddings_cached', 'data'),)
def display_graphics_spin_ring(spins, job_submit_state, job_id, J, embeddings_cached):
    """Generate graphics for spin-ring display."""

    trigger_id = dash.callback_context.triggered[0]['prop_id'].split('.')[0]
   
    if trigger_id == 'job_submit_state': 
    
        if job_submit_state == 'COMPLETED':

            embeddings_cached = embeddings_cached = json_to_dict(embeddings_cached)
            sampleset_unembedded = get_samples(client, job_id, spins, J, embeddings_cached[spins])
            kinks_per_sample, kink_density = kink_stats(sampleset_unembedded, J)
            best_indx = np.abs(kinks_per_sample - kink_density).argmin()
            best_sample = sampleset_unembedded.record.sample[best_indx]

            fig = plot_spin_orientation(num_spins=spins, sample=best_sample)
            return fig
        
        else:

            return dash.no_update

    fig = plot_spin_orientation(num_spins=spins, sample=None)
    return fig

@app.callback(
    Output('job_id', 'children'),
    Input('job_submit_time', 'children'),
    State('qpu_selection', 'value'),
    State('spins', 'value'),
    State('coupling_strength', 'value'),
    State('anneal_duration', 'value'),
    State('embeddings_cached', 'data'),)
def submit_job(job_submit_time, qpu_name, spins, J, ta_ns, embeddings_cached):
    """Submit job and provide job ID."""

    trigger_id = dash.callback_context.triggered[0]['prop_id'].split('.')[0]

    if trigger_id =='job_submit_time':

        solver = qpus[qpu_name]

        bqm = create_bqm(num_spins=spins, coupling_strength=J)

        embeddings_cached = json_to_dict(embeddings_cached)
        embedding = embeddings_cached[spins]

        bqm_embedded = embed_bqm(bqm, embedding, DWaveSampler(solver=solver.name).adjacency)

        computation = solver.sample_bqm(
            bqm=bqm_embedded,
            fast_anneal=True,
            annealing_time=0.001*ta_ns,     # SAPI anneal time units is microseconds
            auto_scale=False, 
            answer_mode='raw',              # Easier than accounting for num_occurrences
            num_reads=100, 
            label=f'Examples - Kibble-Zurek Simulation, submitted: {job_submit_time}',)   

        return computation.wait_id()

    return dash.no_update

@app.callback(
    Output('btn_simulate', 'disabled'),
    Output('wd_job', 'disabled'),
    Output('wd_job', 'interval'),
    Output('wd_job', 'n_intervals'),
    Output('job_submit_state', 'children'),
    Output('job_submit_time', 'children'),
    Output('embeddings_found', 'data'),
    Input('btn_simulate', 'n_clicks'),
    Input('wd_job', 'n_intervals'),
    State('job_id', 'children'),
    State('job_submit_state', 'children'),
    State('job_submit_time', 'children'),
    State('embedding_is_cached', 'value'),
    State('spins', 'value'),
    State('qpu_selection', 'value'),
    State('embeddings_found', 'data'),)
def simulate(dummy1, dummy2, job_id, job_submit_state, job_submit_time, \
             cached_embedding_lengths, spins, qpu_name, embeddings_found):
    """Manage simulation: embedding, job submission."""

    trigger_id = dash.callback_context.triggered[0]['prop_id'].split('.')[0]

    if not any(trigger_id == input for input in ['btn_simulate', 'wd_job']):
        return dash.no_update, dash.no_update, dash.no_update, \
            dash.no_update, dash.no_update, dash.no_update, dash.no_update

    if trigger_id == 'btn_simulate':

        if spins in cached_embedding_lengths:   

            submit_time = datetime.datetime.now().strftime('%c')
            job_submit_state = 'SUBMITTED'
            embedding = dash.no_update

        else:

            submit_time = dash.no_update
            job_submit_state = 'EMBEDDING'
            embedding = 'needed'

        disable_btn = True
        disable_watchdog = False

        return disable_btn, disable_watchdog, 0.5*1000, 0, job_submit_state, submit_time, embedding
    
    if job_submit_state == 'EMBEDDING':

        submit_time = dash.no_update
        embedding = dash.no_update

        if embeddings_found == 'needed':

            try:
                embedding = find_one_to_one_embedding(spins, qpus[qpu_name].edges)
                if embedding:
                    job_submit_state = 'EMBEDDING'  # Stay another WD to allow caching the embedding
                    embedding = {spins: embedding}
                else:
                    job_submit_state = 'FAILED'
                    embedding = 'not found'
            except Exception:
                job_submit_state = 'FAILED'
                embedding = 'not found'

        else:   # Found embedding last WD, so is cached, so now can submit job
            
            submit_time = datetime.datetime.now().strftime('%c')
            job_submit_state = 'SUBMITTED'

        return True, False, 0.2*1000, 0, job_submit_state, submit_time, embedding

    if any(job_submit_state == status for status in
        ['SUBMITTED', 'PENDING', 'IN_PROGRESS']):

        job_submit_state = get_job_status(client, job_id, job_submit_time)
        if not job_submit_state:
            job_submit_state = 'SUBMITTED'
            wd_time = 0.2*1000
        else:
            wd_time = 1*1000

        return True, False, wd_time, 0, job_submit_state, dash.no_update, dash.no_update

    if any(job_submit_state == status for status in ['COMPLETED', 'CANCELLED', 'FAILED']):

        disable_btn = False
        disable_watchdog = True

        return disable_btn, disable_watchdog, 0.1*1000, 0, dash.no_update, dash.no_update, dash.no_update

    else:   # Exception state: should only ever happen in testing
        return False, True, 0, 0, 'ERROR', dash.no_update, dash.no_update

@app.callback(
    Output('bar_job_status', 'value'),
    Output('bar_job_status', 'color'),
    Input('job_submit_state', 'children'),)
def set_progress_bar(job_submit_state):
    """Update progress bar for job submissions."""

    trigger_id = dash.callback_context.triggered[0]['prop_id'].split('.')[0]

    if trigger_id == 'job_submit_state':
 
        return job_bar_display[job_submit_state][0], job_bar_display[job_submit_state][1]
    
    return job_bar_display['READY'][0], job_bar_display['READY'][1]

@app.callback(
    *[Output(f'tooltip_{target}', component_property='style') for target in tool_tips.keys()],
    Input('tooltips_show', 'value'),)
def activate_tooltips(tooltips_show):
    """Activate or hide tooltips."""

    trigger = dash.callback_context.triggered
    trigger_id = trigger[0]['prop_id'].split('.')[0]

    if trigger_id == 'tooltips_show':
        if tooltips_show == 'off':
            return dict(display='none'), dict(display='none'), dict(display='none'), \
dict(display='none'), dict(display='none'), dict(display='none'), \
dict(display='none'), dict(display='none'), dict(display='none'), 

    return dict(), dict(), dict(), dict(), dict(), dict(), dict(), dict(), dict()


if __name__ == "__main__":
    app.run_server(debug=True)