Source code for schrodinger.application.matsci.qubec_utils

"""
Utilities for QUBEC.

Copyright Schrodinger, LLC. All rights reserved.
"""

import argparse
import json
import os

import requests

from schrodinger.application.matsci import enc
from schrodinger.application.matsci import jobutils
from schrodinger.application.matsci import msutils
from schrodinger.application.matsci import parserutils
from schrodinger.application.matsci import textlogger
from schrodinger.application.matsci.qb_sdk.api import enable_account
from schrodinger.application.matsci.qb_sdk.api import execute
from schrodinger.application.matsci.qb_sdk.parameters import QubecSdkError
from schrodinger.job import jobcontrol
from schrodinger.utils import fileutils

DEFAULT_JOB_NAME = 'qubec_job'

PROBLEM = 'problem'
SCF_PARAMETERS = 'scf_parameters'
QUANTUM_PARAMETERS = 'quantum_parameters'
GEOMETRY = 'geometry'
PROPERTIES = 'properties'
COUPLED_CLUSTER = 'Coupled Cluster'
CAS_MO_DELIMITER = ','

QUBEC_EUK = '.qubec_euk'
QUBEC_EU = f'.qubec_eu{enc.ENC_FILE_EXT}'
QUBEC_EPK = '.qubec_epk'
QUBEC_EP = f'.qubec_ep{enc.ENC_FILE_EXT}'
PAIRS = ((QUBEC_EUK, QUBEC_EU), (QUBEC_EPK, QUBEC_EP))

COMPLETED = 'completed'
FAILED = 'failed'
CANCELLED = 'cancelled'

HF_ENERGY_KEY = 'hf_energy'
HF_ENERGY_PROP_KEY = 'r_matsci_qubec_HF_energy_(au)'
CCSD_ENERGY_KEY = 'ccsd_energy'
CCSD_ENERGY_PROP_KEY = 'r_matsci_qubec_CCSD_energy_(au)'
FCI_ENERGY_KEY = 'fci_energy'
FCI_ENERGY_PROP_KEY = 'r_matsci_qubec_FCI_energy_(au)'
ENERGY_KEY = 'optimal_value'
ENERGY_PROP_KEY = 'r_matsci_qubec_energy_(au)'

ENERGY_KEY_PROP_DICT = {
    HF_ENERGY_KEY: HF_ENERGY_PROP_KEY,
    CCSD_ENERGY_KEY: CCSD_ENERGY_PROP_KEY,
    FCI_ENERGY_KEY: FCI_ENERGY_PROP_KEY,
    ENERGY_KEY: ENERGY_PROP_KEY
}

ENERGY_PER_ITER_KEY = 'history_values'
ENERGY_PER_ITER_PROP_KEY = 'r_matsci_qubec_energy_(au)_iter_{idx}'

# see MATSCI-10171 - many of the following QPERE properties can be very
# large integers so use strings instead

N_LOGI_QUBITS_KEY = 'Total logical qubits'
N_LOGI_QUBITS_PROP_KEY = 's_matsci_qubec_number_logical_qubits'

GS_N_PHYS_QUBITS_KEY = 'Ground state physical qubits'
GS_N_PHYS_QUBITS_PROP_KEY = 's_matsci_qubec_ground_state_physical_qubits'
GS_RUN_TIME_KEY = 'Ground state runtime'
GS_RUN_TIME_PROP_KEY = 's_matsci_qubec_ground_state_runtime'
GS_N_NON_CLIFFORD_KEY = 'Ground state non clifford gates'
GS_N_NON_CLIFFORD_PROP_KEY = 's_matsci_qubec_ground_state_non_clifford_gates'

ES_N_PHYS_QUBITS_KEY = 'Excited states physical qubits'
ES_N_PHYS_QUBITS_PROP_KEY = 's_matsci_qubec_excited_state_{idx}_physical_qubits'
ES_RUN_TIME_KEY = 'Excited states runtime'
ES_RUN_TIME_PROP_KEY = 's_matsci_qubec_excited_state_{idx}_runtime'
ES_N_NON_CLIFFORD_KEY = 'Excited states non clifford gates'
ES_N_NON_CLIFFORD_PROP_KEY = 's_matsci_qubec_excited_state_{idx}_non_clifford_gates'

QPERE_KEY_PROP_DICT = {
    N_LOGI_QUBITS_KEY: N_LOGI_QUBITS_PROP_KEY,
    GS_N_PHYS_QUBITS_KEY: GS_N_PHYS_QUBITS_PROP_KEY,
    GS_RUN_TIME_KEY: GS_RUN_TIME_PROP_KEY,
    GS_N_NON_CLIFFORD_KEY: GS_N_NON_CLIFFORD_PROP_KEY,
    ES_N_PHYS_QUBITS_KEY: ES_N_PHYS_QUBITS_PROP_KEY,
    ES_RUN_TIME_KEY: ES_RUN_TIME_PROP_KEY,
    ES_N_NON_CLIFFORD_KEY: ES_N_NON_CLIFFORD_PROP_KEY
}


[docs]def qubec_error(msg):
    """
    Return an error message flagged with QUBEC.

    :type msg: str
    :param msg: the error message

    :rtype: str
    :return: the error message flagged with QUBEC
    """

    return f'QUBEC error: {msg}'


[docs]class QubecException(Exception):
    pass


def _write_enc_files(username, password):
    """
    Write encrypted files.

    :type username: str
    :param username: the username

    :type password: str
    :param password: the password
    """

    for (k_file, v_file), text in zip(PAIRS, (username, password)):
        enc.write_encrypted_text(text, v_file, k_file, key_file_path='')


def _get_username_password():
    """
    Return the username and password.

    :raise QubecException: if there is an issue

    :rtype: tuple
    :return: username and password
    """

    texts = []
    for k_file, v_file in PAIRS:
        if not (os.path.exists(k_file) and os.path.exists(v_file)):
            raise QubecException(
                f'Both {k_file} and {v_file} files must exist.')
        texts.append(enc.get_unencrypted_text(v_file, k_file, key_file_path=''))
    return tuple(texts)


def _delete_enc_files(path=''):
    """
    Delete the encrypted files.

    :type path: str
    :param path: the path
    """

    for files in PAIRS:
        for afile in files:
            file_path = os.path.join(path, afile)
            fileutils.force_remove(file_path)


[docs]def write_input_file(base_name, adict):
    """
    Write the input json file.

    :type base_name: str
    :param base_name: the base name of the input file to be written

    :type adict: dict
    :param adict: job parameters
    """

    file_name = base_name
    if not file_name.endswith('.json'):
        file_name += '.json'
    with open(file_name, 'w') as fh:
        json.dump(adict, fh, sort_keys=True, indent=4, separators=(',', ':'))


[docs]def read_input_file(base_name):
    """
    Read the input json file.

    :type base_name: str
    :param base_name: the base name of the input file to be read

    :raise QubecException: if there is an issue

    :rtype: dict
    :return: job parameters
    """

    file_name = base_name
    if not file_name.endswith('.json'):
        file_name += '.json'
    if not os.path.exists(file_name):
        raise QubecException(f'The file {file_name} does not exist.')
    with open(file_name, 'r') as fh:
        try:
            adict = json.load(fh)
        except json.decoder.JSONDecodeError:
            raise QubecException(
                f'The file {file_name} is not a valid json file.')
    return adict


[docs]def type_input(arg):
    """
    Validate the input.

    :type arg: str or unicode
    :param arg: the input file to validate

    :raise argparse.ArgumentTypeError: if the given input is invalid

    :rtype: str
    :return: the str-ed input file
    """

    arg = parserutils.type_json_file(arg)
    try:
        _get_username_password()
    except QubecException:
        raise argparse.ArgumentTypeError('Missing credentials.')
    return arg


[docs]class Job(jobcontrol.Job):
    """
    Job class to handle cancelling a QUBEC job.
    """

[docs]    def __init__(self, *args, **kwargs):
        """
        See parent class for documentation.
        """

        super().__init__(*args, **kwargs)
        self.qubec_job = None
        self.result = None

[docs]    def setQubecJob(self, qubec_job):
        """
        Set the QUBEC job.

        :type qubec_job: QubecJob
        :param qubec_job: the QUBEC job
        """

        self.qubec_job = qubec_job

[docs]    def cancel(self):
        """
        See parent class for documentation.
        """

        if self.qubec_job:
            with msutils.ignore_ssl_warnings():
                content = self.qubec_job.cancel()
            self.result = content.get('data', {}).get('result', {})
        super().cancel()


[docs]class Qubec():
    """
    Manage a QUBEC job.
    """

[docs]    def __init__(self, st, params, logger=None):
        """
        Create an instance.

        :type st: schrodinger.structure.Structure
        :param st: the structure

        :type params: dict
        :param params: the QUBEC job parameters

        :type logger: logging.Logger or None
        :param logger: output logger or None if there isn't one
        """

        self.st = st
        self.params = params
        self.logger = logger

        self.username = None
        self.password = None
        self.qubec_job = None
        self.result = None

        jobbe = jobcontrol.get_backend()
        if jobbe:
            self.parent_job = Job(jobbe.job_id)
        else:
            self.parent_job = None

[docs]    def login(self):
        """
        Log in.

        :raise QubecException: if there is an issue
        """

        try:
            self.username, self.password = _get_username_password()
            with msutils.ignore_ssl_warnings():
                enable_account(self.username, password=self.password)
        except (QubecException, QubecSdkError) as err:
            if isinstance(err, QubecSdkError):
                err = qubec_error(str(err))
            raise QubecException(err)
        finally:
            _delete_enc_files()
        if self.logger:
            self.logger.info('Login successful.')
            self.logger.info('')

[docs]    def getParams(self):
        """
        Return the job parameters.

        :rtype: dict
        :return: the job parameters
        """

        params = self.params.copy()
        problem = params.get(PROBLEM, {})
        problem[GEOMETRY] = [
            (atom.element, tuple(atom.xyz)) for atom in self.st.atom
        ]
        return params

[docs]    def prepareOutput(self):
        """
        Prepare the output.
        """

        if self.parent_job:
            result = self.parent_job.result
        else:
            result = None
        self.result = result or self.qubec_job.result
        if not self.result:
            return

        if self.logger:
            self.logger.info('Results:')
            self.logger.info('')
            for key, value in self.result.items():
                self.logger.info(f'{key} = {value}')
                self.logger.info('')

        for key, prop in ENERGY_KEY_PROP_DICT.items():
            value = self.result.get(key)
            if value is not None:
                self.st.property[prop] = value

        energies = self.result.get(ENERGY_PER_ITER_KEY, [])
        for idx, energy in enumerate(energies):
            self.st.property[ENERGY_PER_ITER_PROP_KEY.format(idx=idx)] = energy

        for key, prop in QPERE_KEY_PROP_DICT.items():
            values = self.result.get(key)
            if isinstance(values, list):
                for idx, value in enumerate(values, 1):
                    self.st.property[prop.format(idx=idx)] = str(value)
            else:
                value = values
                if value is not None:
                    self.st.property[prop] = str(value)

        job_name = jobutils.get_jobname(DEFAULT_JOB_NAME)
        output_file = f'{job_name}_out.mae'
        self.st.write(output_file)
        jobutils.add_outfile_to_backend(output_file, set_structure_output=True)

[docs]    def run(self):
        """
        Run the job.

        :raise QubecException: if there is an issue
        """

        self.login()

        params = self.getParams()

        # see MATSCI-10077 for the json.decoder.JSONDecodeError
        # see MATSCI-10121 for the requests.exceptions.ConnectionError
        try:
            with msutils.ignore_ssl_warnings():
                self.qubec_job = execute(**params)
                if self.parent_job:
                    self.parent_job.setQubecJob(self.qubec_job)
                if self.logger:
                    self.logger.info(
                        f'Executing QUBEC job {self.qubec_job.job_id}.')
                    self.logger.info('')
                self.qubec_job.wait_completion()
        except (QubecSdkError, json.decoder.JSONDecodeError,
                requests.exceptions.ConnectionError) as err:
            err = qubec_error(str(err))
            raise QubecException(err)

        self.prepareOutput()

        if self.qubec_job.status == FAILED:
            raise QubecException('QUBEC job failed.')
        elif self.qubec_job.status == COMPLETED and not self.result:
            raise QubecException('QUBEC job completed but no results found.')
        elif self.qubec_job.status == CANCELLED:
            raise QubecException('QUBEC job cancelled.')
        elif self.qubec_job.error:
            raise QubecException(self.qubec_job.error)

        if self.logger:
            textlogger.log_msg('QUBEC job completed successfully',
                               timestamp=True,
                               logger=self.logger)