Source code for schrodinger.application.jaguar.neural_network_optimizer

"""
functions for energy minimization using a neural network potential energy surface
"""

import datetime
import json
import os
import time

import requests

from schrodinger import structure

HOST_VAR = "SCHRODINGER_ANI_HOSTNAME"
PROJ_VAR = "SCHRODINGER_ANI_PROJECT"
SUPPORTED_ELEMENTS = ["H", "C", "N", "O"]
MAE_STRING = 'mae_string'
FF_NO_PREOPT = 'b_ml_no_ff_min'


class UnsupportedElementError(ValueError):
    """
    Unsupported element type
    """


class AniServerError(RuntimeError):
    """
    Server is not responding
    """


class OptimizationFailure(AniServerError):
    """
    Optimization failed for an unknown reason
    """


# PaperPusherClient defined below is to avoid the following import:
#from paperpusher_client import PaperPusherClient


class _JobStatus(object):
    READY = 'READY'
    RUNNING = "RUNNING"
    COMPLETED = 'COMPLETED'
    FAILED = 'FAILED'


class _PaperPusherClient(object):

    def __init__(self, base_url, project):
        """
        :type base_url: str
        :param base_url: url of host
        :type project: str
        :param project: project name
        """
        self.base_url = base_url
        self.project = project

    def ping_server(self):
        """
        ping the server for health
        :return: dict of health status
        """
        url = "%s/v1/health" % self.base_url
        r = requests.get(url)
        return r.json()

    def is_running(self):
        """
        is there a server responding to requests at the current address
        """

        try:
            response = self.ping_server()
        except (requests.exceptions.ConnectionError,
                json.decoder.JSONDecodeError) as e:
            is_running = False
        else:
            is_running = response.get("is_up", False)

        return is_running

    def add_job(self, kwargs_string, block_for_results=False, timeout=120.0):
        """
        :type: kwargs_string, str
        :param kwargs_string: keyword arguments sent to host
        :type: block_for_results: bool
        :param block_for_results: if true poll until status is in [COMPLETED, FAILED]
                                  or until timeout seconds
        :type timeout: float
        :param timeout: wait this many seconds before returning to caller
                        if status does not change
        :return: job dictionary
        """
        url = "%s/v1/job" % self.base_url
        payload = {'project': self.project, 'kwargs': kwargs_string}
        payload = json.dumps(payload)
        r = requests.post(url, data=payload)
        r = r.json()

        if not block_for_results:
            return r

        start_time = datetime.datetime.now()
        end_time = datetime.datetime.now()
        job_id = r['id']
        details = r
        while (end_time - start_time).seconds < timeout:
            details = self.job_details(job_id)
            if details['status'] in [_JobStatus.COMPLETED, _JobStatus.FAILED]:
                return details
            end_time = datetime.datetime.now()
            time.sleep(2)

        return details

    def job_details(self, job_id):
        """
        :type job_id: str
        :param job_id: id of the job details are requested for
        :return: details of job
        """
        url = "%s/v1/job/%s" % (self.base_url, job_id)
        r = requests.get(url)
        return r.json()


def _get_client():
    """
    read env variables and return a _PaperPusherClient
    if server is not running a AniServerError is raised
    if env variables not set a ValueError is raised
    """

    try:
        hostname = os.environ[HOST_VAR]
        project = os.environ[PROJ_VAR]
    except KeyError:
        raise ValueError(
            "%s and %s env variables must be set to use ANI optimization" %
            (HOST_VAR, PROJ_VAR))
    client = _PaperPusherClient(hostname, project)

    if not client.is_running():
        raise AniServerError("ANI Server cannot be contacted")

    return client


def nn_optimize(st, ff_preopt=True):
    """
    Optimize a structure using the ANI potential energy function.
    If atom property s_ff_tors_atoms is present, that natural torsion is constrained

    use of this function requires two environment variables to be set
    %s and %s define the host and project
    name.

    If these variables are not defined a ValueError is raised
    If they are set but the server is not responding an AniServerError is raised
    If unsupported elements are present and the server is responding
        an UnsupportedElementError is raised
    If the optimization for any other reason we raise a OptimizationFailure

    :type st: Structure
    :param st: Structure to optimize
    """ % (HOST_VAR, PROJ_VAR)

    client = _get_client()

    # package structure into string, then string into json
    if not ff_preopt:
        st.property[FF_NO_PREOPT] = True

    mae_string = st.writeToString('maestro')

    if not ff_preopt:
        del st.property[FF_NO_PREOPT]

    json_to_send = {MAE_STRING: mae_string}
    send_data = json.dumps(json_to_send)

    # send data and get response
    response = client.add_job(send_data, block_for_results=True)

    # even if we don't support this system we send the data
    if any(at.element not in SUPPORTED_ELEMENTS for at in st.atom):
        raise UnsupportedElementError(
            "Cannot optimize structure with unsupported element types")

    if response["status"] == _JobStatus.COMPLETED:
        # load string from json, then structure from string
        json_result = json.loads(response['result'])
        mae_string = json_result[MAE_STRING]
        st = next(structure.StructureReader.fromString(mae_string))
    else:
        raise OptimizationFailure("optimization failed with status %s" %
                                  response["status"])

    return st