<commit_before># Copyright (C) 2016 Google Inc.
# Licensed under http://www.apache.org/licenses/LICENSE-2.0 <see LICENSE file>

"""
 GGRC email notification module hook to prepares email, email digest, notify
 email to recipients
"""

from collections import defaultdict
from datetime import date
from datetime import datetime
from logging import getLogger

from sqlalchemy import and_
from sqlalchemy.orm import joinedload
from werkzeug.exceptions import Forbidden
from google.appengine.api import mail

from ggrc import db
from ggrc import extensions
from ggrc import settings
from ggrc.models import Person
from ggrc.models import Notification
from ggrc.models import NotificationConfig
from ggrc.rbac import permissions
from ggrc.utils import merge_dict

# pylint: disable=invalid-name
logger = getLogger(__name__)

class Services(object):
  """Helper class for notification services.

  This class is a helper class for calling a notification service for a given
  object. The first call get_service_function must be done after all modules
  have been initialized.
  """

  services = []

  @classmethod
  def get_service_function(cls, name):
    """Get callback function for an object.

    This returns a service function which is a registered callback for an
    object that has a notification.

    Args:
      name: Name of an object for which we want to get a service function, such
        as "Request", "CycleTask", "Assessment", etc.

    Returns:
      callable: A function that takes a notification and returns a data dict
        with all the data for that notification.
    """
    if not cls.services:
      cls.services = extensions.get_module_contributions(
          "contributed_notifications")
    if name not in cls.services:
      raise ValueError("unknown service name: %s" % name)
    return cls.services[name]

  @classmethod
  def call_service(cls, notif):
    """Call data handler service for the object in the notification.

    Args:
      notif (Notification): Notification object for which we want to get the
        notification data dict.

    Returns:
      dict: Result of the data handler for the object in the notification.
    """
    service = cls.get_service_function(notif.object_type)
    return service(notif)

def get_filter_data(notification, people_cache):
  """Get filtered notification data.

  This function gets notification data for all users who should receive it. A
  single notification can be for multiple users (such as all assignees) but
  only some should receive it depending on if it's an instant notification or
  a daily digest and the specific users notification settings.

  Args:
    notification (Notification): Notification object for which we want to get
      data.

  Returns:
    dict: dictionary containing notification data for all users who should
      receive it, according to their notification settings.
  """
  result = {}
  data = Services.call_service(notification)

  for user, user_data in data.iteritems():
    if should_receive(notification, user_data, people_cache):
      result[user] = user_data
  return result

def get_notification_data(notifications):
  """Get notification data for all notifications.

  This function returns a filtered data for all notifications for the users
  that should receive it.

  Args:
    notifications (list of Notification): List of notification for which we
      want to get notification data.

  Returns:
    dict: Filtered dictionary containing all the data that should be sent for
      the given notification list.
  """
  if not notifications:
    return {}
  aggregate_data = {}
  people_cache = {}

  for notification in notifications:
    filtered_data = get_filter_data(notification, people_cache)
    aggregate_data = merge_dict(aggregate_data, filtered_data)

  # Remove notifications for objects without a contact (such as task groups)
  aggregate_data.pop("", None)

  return aggregate_data

def get_pending_notifications():
  """Get notification data for all future notifications.

  The data dict that get's returned here contains notification data grouped by
  dates on which the notifications should be received.

  Returns
    list of Notifications, data: a tuple of notifications that were handled
      and corresponding data for those notifications.
  """
  notifications = db.session.query(Notification).filter(
      Notification.sent_at.is_(None)).all()

  notif_by_day = defaultdict(list)
  for notification in notifications:
    notif_by_day[notification.send_on.date()].append(notification)

  data = defaultdict(dict)
  today = date.today()
  for day, notif in notif_by_day.iteritems():
    current_day = max(day, today)
    data[current_day] = merge_dict(data[current_day],
                                   get_notification_data(notif))

  return notifications, data

def get_daily_notifications():
  """Get notification data for all future notifications.

  Returns
    list of Notifications, data: a tuple of notifications that were handled
      and corresponding data for those notifications.
  """
  notifications = db.session.query(Notification).filter(
      and_(Notification.send_on <= datetime.today(),
           Notification.sent_at.is_(None)
           )).all()
  return notifications, get_notification_data(notifications)

def should_receive(notif, user_data, people_cache):
  """Check if a user should receive a notification or not.

  Args:
    notif (Notification): A notification entry that we are checking.
    user_data (dict): A dictionary containing data about user notifications.

  Returns:
    True if user should receive the given notification, or False otherwise.
  """
  force_notif = user_data.get("force_notifications", {}).get(notif.id, False)
  person_id = user_data["user"]["id"]
  if person_id in people_cache:
    person = people_cache[person_id]
  else:
    person = db.session.query(Person).options(
      joinedload('user_roles').joinedload('role'),
      joinedload('notification_configs')).filter(
        Person.id == person_id
      ).first()
    people_cache[person_id] = person

  # If the user has no access we should not send any emails
  if person.system_wide_role == "No Access":
    return False

  def is_enabled(notif_type):
    """Check user notification settings.

    Args:
      notif_type (string): can be "Email_Digest" or "Email_Now" based on what
        we wish to check.

    Returns:
      Boolean based on what settings users has stored or what the default
      setting is for the given notification.
    """
    result = NotificationConfig.query.filter(
        and_(NotificationConfig.person_id == person_id,
             NotificationConfig.notif_type == notif_type))
    if result.count() == 0:
      # If we have no results, we need to use the default value, which is
      # true for digest emails.
      return notif_type == "Email_Digest"
    return result.one().enable_flag

  has_digest = force_notif or is_enabled("Email_Digest")

  return has_digest

def send_daily_digest_notifications():
  """Send emails for today's or overdue notifications.

  Returns:
    str: String containing a simple list of who received the notification.
  """
  # pylint: disable=invalid-name
  notif_list, notif_data = get_daily_notifications()
  sent_emails = []
  subject = "gGRC daily digest for {}".format(date.today().strftime("%b %d"))
  for user_email, data in notif_data.iteritems():
    data = modify_data(data)
    email_body = settings.EMAIL_DIGEST.render(digest=data)
    send_email(user_email, subject, email_body)
    sent_emails.append(user_email)
  set_notification_sent_time(notif_list)
  return "emails sent to:
 {}".format("
".join(sent_emails))

def set_notification_sent_time(notif_list):
  """Set sent time to now for all notifications in the list.

  Args:
    notif_list (list of Notification): List of notification for which we want
      to modify sent_at field.
  """
  for notif in notif_list:
    notif.sent_at = datetime.now()
  db.session.commit()

def show_pending_notifications():
  """Get notification html for all future notifications.

  The data shown here is grouped by dates on which notifications should be
  sent.

  Note that the dates for all future notifications will be wrong since they are
  based on the current date which will be different when the notification is
  actually sent.

  Returns:
    str: Html containing all future notifications.
  """
  if not permissions.is_admin():
    raise Forbidden()
  _, notif_data = get_pending_notifications()

  for day_notif in notif_data.itervalues():
    for data in day_notif.itervalues():
      data = modify_data(data)
  return settings.EMAIL_PENDING.render(data=sorted(notif_data.iteritems()))

def show_daily_digest_notifications():
  """Get notification html for today's notifications.

  Returns:
    str: Html containing all today's notifications.
  """
  # pylint: disable=invalid-name
  if not permissions.is_admin():
    raise Forbidden()
  _, notif_data = get_daily_notifications()
  for data in notif_data.itervalues():
    data = modify_data(data)
  return settings.EMAIL_DAILY.render(data=notif_data)

def get_app_engine_email():
  """Get notification sender email.

  Return the email of the user that will be set as sender. This email should be
  authorized to send emails from the server. For more details, see Application
  Settings for email api authorized senders.

  Returns:
    Valid email address if it's set in the app settings.
  """
  email = getattr(settings, 'APPENGINE_EMAIL')
  return email if mail.is_email_valid(email) else None

def send_email(user_email, subject, body):
  """Helper function for sending emails.

  Args:
    user_email (string): Email of the recipient.
    subject (string): Email subject.
    body (basestring): Html body of the email. it can contain unicode
      characters and will be sent as a html mime type.
  """
  sender = get_app_engine_email()
  if not mail.is_email_valid(user_email):
    logger.error("Invalid email recipient: %s", user_email)
    return
  if not sender:
    logger.error("APPENGINE_EMAIL setting is invalid.")
    return

  message = mail.EmailMessage(sender=sender, subject=subject)

  message.to = user_email
  message.body = "TODO: add email in text mode."
  message.html = body

  message.send()

def modify_data(data):
  """Modify notification data dictionary.

  For easier use in templates, it joins the due_in and due today fields
  together.

  Args:
    data (dict): notification data.

  Returns:
    dict: the received dict with some additional fields for easier traversal
      in the notification template.
  """

  data["due_soon"] = {}
  if "due_in" in data:
    data["due_soon"].update(data["due_in"])
  if "due_today" in data:
    data["due_soon"].update(data["due_today"])

  # combine "my_tasks" from multiple cycles
  data["cycle_started_tasks"] = {}
  if "cycle_data" in data:
    for cycle in data["cycle_data"].values():
      if "my_tasks" in cycle:
        data["cycle_started_tasks"].update(cycle["my_tasks"])

  return data
<commit_msg>Optimize NotificationConfig fetching<commit_after># Copyright (C) 2016 Google Inc.
# Licensed under http://www.apache.org/licenses/LICENSE-2.0 <see LICENSE file>

"""
 GGRC email notification module hook to prepares email, email digest, notify
 email to recipients
"""

from collections import defaultdict
from datetime import date
from datetime import datetime
from logging import getLogger

from sqlalchemy import and_
from sqlalchemy.orm import joinedload
from werkzeug.exceptions import Forbidden
from google.appengine.api import mail

from ggrc import db
from ggrc import extensions
from ggrc import settings
from ggrc.models import Person
from ggrc.models import Notification
from ggrc.models import NotificationConfig
from ggrc.rbac import permissions
from ggrc.utils import merge_dict

# pylint: disable=invalid-name
logger = getLogger(__name__)

class Services(object):
  """Helper class for notification services.

  This class is a helper class for calling a notification service for a given
  object. The first call get_service_function must be done after all modules
  have been initialized.
  """

  services = []

  @classmethod
  def get_service_function(cls, name):
    """Get callback function for an object.

    This returns a service function which is a registered callback for an
    object that has a notification.

    Args:
      name: Name of an object for which we want to get a service function, such
        as "Request", "CycleTask", "Assessment", etc.

    Returns:
      callable: A function that takes a notification and returns a data dict
        with all the data for that notification.
    """
    if not cls.services:
      cls.services = extensions.get_module_contributions(
          "contributed_notifications")
    if name not in cls.services:
      raise ValueError("unknown service name: %s" % name)
    return cls.services[name]

  @classmethod
  def call_service(cls, notif):
    """Call data handler service for the object in the notification.

    Args:
      notif (Notification): Notification object for which we want to get the
        notification data dict.

    Returns:
      dict: Result of the data handler for the object in the notification.
    """
    service = cls.get_service_function(notif.object_type)
    return service(notif)

def get_filter_data(notification, people_cache):
  """Get filtered notification data.

  This function gets notification data for all users who should receive it. A
  single notification can be for multiple users (such as all assignees) but
  only some should receive it depending on if it's an instant notification or
  a daily digest and the specific users notification settings.

  Args:
    notification (Notification): Notification object for which we want to get
      data.

  Returns:
    dict: dictionary containing notification data for all users who should
      receive it, according to their notification settings.
  """
  result = {}
  data = Services.call_service(notification)

  for user, user_data in data.iteritems():
    if should_receive(notification, user_data, people_cache):
      result[user] = user_data
  return result

def get_notification_data(notifications):
  """Get notification data for all notifications.

  This function returns a filtered data for all notifications for the users
  that should receive it.

  Args:
    notifications (list of Notification): List of notification for which we
      want to get notification data.

  Returns:
    dict: Filtered dictionary containing all the data that should be sent for
      the given notification list.
  """
  if not notifications:
    return {}
  aggregate_data = {}
  people_cache = {}

  for notification in notifications:
    filtered_data = get_filter_data(notification, people_cache)
    aggregate_data = merge_dict(aggregate_data, filtered_data)

  # Remove notifications for objects without a contact (such as task groups)
  aggregate_data.pop("", None)

  return aggregate_data

def get_pending_notifications():
  """Get notification data for all future notifications.

  The data dict that get's returned here contains notification data grouped by
  dates on which the notifications should be received.

  Returns
    list of Notifications, data: a tuple of notifications that were handled
      and corresponding data for those notifications.
  """
  notifications = db.session.query(Notification).filter(
      Notification.sent_at.is_(None)).all()

  notif_by_day = defaultdict(list)
  for notification in notifications:
    notif_by_day[notification.send_on.date()].append(notification)

  data = defaultdict(dict)
  today = date.today()
  for day, notif in notif_by_day.iteritems():
    current_day = max(day, today)
    data[current_day] = merge_dict(data[current_day],
                                   get_notification_data(notif))

  return notifications, data

def get_daily_notifications():
  """Get notification data for all future notifications.

  Returns
    list of Notifications, data: a tuple of notifications that were handled
      and corresponding data for those notifications.
  """
  notifications = db.session.query(Notification).filter(
      and_(Notification.send_on <= datetime.today(),
           Notification.sent_at.is_(None)
           )).all()
  return notifications, get_notification_data(notifications)

def should_receive(notif, user_data, people_cache):
  """Check if a user should receive a notification or not.

  Args:
    notif (Notification): A notification entry that we are checking.
    user_data (dict): A dictionary containing data about user notifications.

  Returns:
    True if user should receive the given notification, or False otherwise.
  """
  force_notif = user_data.get("force_notifications", {}).get(notif.id, False)
  person_id = user_data["user"]["id"]
  if person_id in people_cache:
    person = people_cache[person_id]
  else:
    person = db.session.query(Person).options(
      joinedload('user_roles').joinedload('role'),
      joinedload('notification_configs')).filter(
        Person.id == person_id
      ).first()
    people_cache[person_id] = person

  # If the user has no access we should not send any emails
  if person.system_wide_role == "No Access":
    return False

  def is_enabled(notif_type, person):
    """Check user notification settings.

    Args:
      notif_type (string): can be "Email_Digest" or "Email_Now" based on what
        we wish to check.

    Returns:
      Boolean based on what settings users has stored or what the default
      setting is for the given notification.
    """

    notification_configs = [nc for nc in person.notification_configs
                               if nc.notif_type == notif_type]
    if len(notification_configs) == 0:
      # If we have no results, we need to use the default value, which is
      # true for digest emails.
      return notif_type == "Email_Digest"
    return notification_configs[0].enable_flag

  has_digest = force_notif or is_enabled("Email_Digest", person)

  return has_digest

def send_daily_digest_notifications():
  """Send emails for today's or overdue notifications.

  Returns:
    str: String containing a simple list of who received the notification.
  """
  # pylint: disable=invalid-name
  notif_list, notif_data = get_daily_notifications()
  sent_emails = []
  subject = "gGRC daily digest for {}".format(date.today().strftime("%b %d"))
  for user_email, data in notif_data.iteritems():
    data = modify_data(data)
    email_body = settings.EMAIL_DIGEST.render(digest=data)
    send_email(user_email, subject, email_body)
    sent_emails.append(user_email)
  set_notification_sent_time(notif_list)
  return "emails sent to:
 {}".format("
".join(sent_emails))

def set_notification_sent_time(notif_list):
  """Set sent time to now for all notifications in the list.

  Args:
    notif_list (list of Notification): List of notification for which we want
      to modify sent_at field.
  """
  for notif in notif_list:
    notif.sent_at = datetime.now()
  db.session.commit()

def show_pending_notifications():
  """Get notification html for all future notifications.

  The data shown here is grouped by dates on which notifications should be
  sent.

  Note that the dates for all future notifications will be wrong since they are
  based on the current date which will be different when the notification is
  actually sent.

  Returns:
    str: Html containing all future notifications.
  """
  if not permissions.is_admin():
    raise Forbidden()
  _, notif_data = get_pending_notifications()

  for day_notif in notif_data.itervalues():
    for data in day_notif.itervalues():
      data = modify_data(data)
  return settings.EMAIL_PENDING.render(data=sorted(notif_data.iteritems()))

def show_daily_digest_notifications():
  """Get notification html for today's notifications.

  Returns:
    str: Html containing all today's notifications.
  """
  # pylint: disable=invalid-name
  if not permissions.is_admin():
    raise Forbidden()
  _, notif_data = get_daily_notifications()
  for data in notif_data.itervalues():
    data = modify_data(data)
  return settings.EMAIL_DAILY.render(data=notif_data)

def get_app_engine_email():
  """Get notification sender email.

  Return the email of the user that will be set as sender. This email should be
  authorized to send emails from the server. For more details, see Application
  Settings for email api authorized senders.

  Returns:
    Valid email address if it's set in the app settings.
  """
  email = getattr(settings, 'APPENGINE_EMAIL')
  return email if mail.is_email_valid(email) else None

def send_email(user_email, subject, body):
  """Helper function for sending emails.

  Args:
    user_email (string): Email of the recipient.
    subject (string): Email subject.
    body (basestring): Html body of the email. it can contain unicode
      characters and will be sent as a html mime type.
  """
  sender = get_app_engine_email()
  if not mail.is_email_valid(user_email):
    logger.error("Invalid email recipient: %s", user_email)
    return
  if not sender:
    logger.error("APPENGINE_EMAIL setting is invalid.")
    return

  message = mail.EmailMessage(sender=sender, subject=subject)

  message.to = user_email
  message.body = "TODO: add email in text mode."
  message.html = body

  message.send()

def modify_data(data):
  """Modify notification data dictionary.

  For easier use in templates, it joins the due_in and due today fields
  together.

  Args:
    data (dict): notification data.

  Returns:
    dict: the received dict with some additional fields for easier traversal
      in the notification template.
  """

  data["due_soon"] = {}
  if "due_in" in data:
    data["due_soon"].update(data["due_in"])
  if "due_today" in data:
    data["due_soon"].update(data["due_today"])

  # combine "my_tasks" from multiple cycles
  data["cycle_started_tasks"] = {}
  if "cycle_data" in data:
    for cycle in data["cycle_data"].values():
      if "my_tasks" in cycle:
        data["cycle_started_tasks"].update(cycle["my_tasks"])

  return data
<|endoftext|>


<commit_before>"""
We have implemented RMSD calculation using both Kabsch algorithm and
via the quaterion-based characteristic polynomial.  Both are implemented
in pure python/numpy.

Notes
-----
This file works at a frame-by-frame level, not the trajectory level.

Reference
---------
Rapid calculation of RMSDs using a quaternion-based characteristic polynomial.
Acta Crystallogr A 61(4):478-480.
"""

import numpy as np
import scipy.optimize

class Transformation():
    def __init__(self, rotation, translation):
        self.rotation = rotation
        self.translation = translation

    def transform(self, xyz):
        """Apply transformation to xyz."""
        mu1 = xyz.mean(0)
        return xyz.dot(self.rotation) + self.translation - mu1.dot(self.rotation)

def compute_transformation(mobile, target):
    """Returns a Transformation object to align mobile onto target.

    Parameters
    ----------
    mobile : ndarray, shape = (n_atoms, 3)
        xyz coordinates of a SINGLE frame, to be aligned onto target.
    target : ndarray, shape = (n_atoms, 3)
        xyz coordinates of a SINGLE frame

    Returns
    -------
    T : Transformation aligning mobile to target.
    """
    translation, rotation = _compute_translation_and_rotation(mobile, target)
    return Transformation(rotation, translation)

def transform(mobile, target):
    """Align mobile onto target and return transformed coordinates.

    Parameters
    ----------
    mobile : ndarray, shape = (n_atoms, 3)
        xyz coordinates of a SINGLE frame, to be aligned onto target.
    target : ndarray, shape = (n_atoms, 3)
        xyz coordinates of a SINGLE frame

    Returns
    -------
    mobile_prime : transformed coordinates of mobile.
    """
    T = compute_transformation(mobile, target)
    mobile_prime = T.transform(mobile)
    return mobile_prime

def transform_inplace(mobile, target):
    """Align mobile onto target and adjust mobile's coordinates.

    Parameters
    ----------
    mobile : ndarray, shape = (n_atoms, 3)
        xyz coordinates of a SINGLE frame, to be aligned onto target.
    target : ndarray, shape = (n_atoms, 3)
        xyz coordinates of a SINGLE frame

    Notes
    -----
    This is an inplace operation on mobile.
    """
    mobile_prime = transform(mobile, target)
    mobile[:] = mobile_prime

def _compute_translation_and_rotation(mobile, target):
    """Returns the translation and rotation mapping mobile onto target.

    Parameters
    ----------
    mobile : ndarray, shape = (n_atoms, 3)
        xyz coordinates of a SINGLE frame, to be aligned onto target.
    target : ndarray, shape = (n_atoms, 3)
        xyz coordinates of a SINGLE frame

    Returns
    -------
    R : ndarray
        rotation
    """
    assert(mobile.shape[1] == 3)
    assert(mobile.shape == target.shape)

    mu1 = mobile.mean(0)
    mu2 = target.mean(0)

    translation = mu2

    mobile = mobile - mu1
    target = target - mu2

    correlation_matrix = np.dot(np.transpose(mobile), target)
    V, S, W_tr = np.linalg.svd(correlation_matrix)
    is_reflection = (np.linalg.det(V) * np.linalg.det(W_tr)) < 0.0
    if is_reflection:
        V[:, -1] = -V[:, -1]
    rotation = np.dot(V, W_tr)

    return translation, rotation

def rmsd_kabsch(xyz1, xyz2):
    """Returns the RMSD distance between conformations xyz1 and xyz2.

    Parameters
    ----------
    xyz1 : ndarray, shape = (n_atoms, 3)
        xyz1 coordinates of a SINGLE frame
    xyz2 : ndarray, shape = (n_atoms, 3)
        xyz2 coordinates of a SINGLE frame
    Returns
    -------
    rmsd : float
        RMSD between xyz1 and xyz2
    """
    xyz1_prime = transform(xyz1, xyz2)
    delta = xyz1_prime - xyz2
    rmsd = (delta ** 2.0).sum(1).mean() ** 0.5
    return rmsd

def _center(conformation):
    """Center and typecheck the conformation"""

    conformation = np.asarray(conformation)
    if not conformation.ndim == 2:
        raise ValueError('conformation must be two dimensional')
    _, three = conformation.shape
    if not three == 3:
        raise ValueError('conformation second dimension must be 3')

    centroid = np.mean(conformation, axis=0)
    centered = conformation - centroid
    return centered

def rmsd_qcp(conformation1, conformation2):
    """Compute the RMSD with Theobald's quaterion-based characteristic
    polynomial

    Rapid calculation of RMSDs using a quaternion-based characteristic polynomial.
    Acta Crystallogr A 61(4):478-480.

    Parameters
    ----------
    conformation1 : np.ndarray, shape=(n_atoms, 3)
        The cartesian coordinates of the first conformation
    conformation2 : np.ndarray, shape=(n_atoms, 3)
        The cartesian coordinates of the second conformation

    Returns
    -------
    rmsd : float
        The root-mean square deviation after alignment between the two pointsets
    """
    # center and typecheck the conformations
    A = _center(conformation1)
    B = _center(conformation2)
    if not A.shape[0] == B.shape[0]:
        raise ValueError('conformation1 and conformation2 must have same number of atoms')
    n_atoms = len(A)

    # the inner product of the structures A and B
    G_A = np.einsum('ij,ij', A, A)
    G_B = np.einsum('ij,ij', B, B)
    # print 'GA', G_A, np.trace(np.dot(A.T, A))
    # print 'GB', G_B, np.trace(np.dot(B.T, B))

    # M is the inner product of the matrices A and B
    M = np.dot(B.T, A)

    # unpack the elements
    Sxx, Sxy, Sxz = M[0, :]
    Syx, Syy, Syz = M[1, :]
    Szx, Szy, Szz = M[2, :]

    # do some intermediate computations to assemble the characteristic
    # polynomial
    Sxx2 = Sxx * Sxx
    Syy2 = Syy * Syy
    Szz2 = Szz * Szz

    Sxy2 = Sxy * Sxy
    Syz2 = Syz * Syz
    Sxz2 = Sxz * Sxz

    Syx2 = Syx * Syx
    Szy2 = Szy * Szy
    Szx2 = Szx * Szx

    SyzSzymSyySzz2 = 2.0 * (Syz * Szy - Syy * Szz)
    Sxx2Syy2Szz2Syz2Szy2 = Syy2 + Szz2 - Sxx2 + Syz2 + Szy2

    # two of the coefficients
    C2 = -2.0 * (Sxx2 + Syy2 + Szz2 + Sxy2 + Syx2 + Sxz2 + Szx2 + Syz2 + Szy2)
    C1 = 8.0 * (Sxx * Syz * Szy + Syy * Szx * Sxz + Szz * Sxy * Syx - Sxx * Syy * Szz - Syz * Szx * Sxy - Szy * Syx * Sxz)

    SxzpSzx = Sxz + Szx
    SyzpSzy = Syz + Szy
    SxypSyx = Sxy + Syx
    SyzmSzy = Syz - Szy
    SxzmSzx = Sxz - Szx
    SxymSyx = Sxy - Syx
    SxxpSyy = Sxx + Syy
    SxxmSyy = Sxx - Syy
    Sxy2Sxz2Syx2Szx2 = Sxy2 + Sxz2 - Syx2 - Szx2

    # the other coefficient
    C0 = Sxy2Sxz2Syx2Szx2 * Sxy2Sxz2Syx2Szx2 \
        + (Sxx2Syy2Szz2Syz2Szy2 + SyzSzymSyySzz2) * (Sxx2Syy2Szz2Syz2Szy2 - SyzSzymSyySzz2) \
        + (-(SxzpSzx) * (SyzmSzy) + (SxymSyx) * (SxxmSyy - Szz)) * (-(SxzmSzx) * (SyzpSzy) + (SxymSyx) * (SxxmSyy + Szz)) \
        + (-(SxzpSzx) * (SyzpSzy) - (SxypSyx) * (SxxpSyy - Szz)) * (-(SxzmSzx) * (SyzmSzy) - (SxypSyx) * (SxxpSyy + Szz)) \
        + (+(SxypSyx) * (SyzpSzy) + (SxzpSzx) * (SxxmSyy + Szz)) * (-(SxymSyx) * (SyzmSzy) + (SxzpSzx) * (SxxpSyy + Szz)) \
        + (+(SxypSyx) * (SyzmSzy) + (SxzmSzx) * (SxxmSyy - Szz)) * (-(SxymSyx) * (SyzpSzy) + (SxzmSzx) * (SxxpSyy - Szz))

    E0 = (G_A + G_B) / 2.0
    f = lambda x: x ** 4.0 + C2 * x ** 2. + C1 * x + C0
    df = lambda x: 4 * x ** 3.0 + 2 * C2 * x + C1
    max_eigenvalue = scipy.optimize.newton(f, E0, df)
    rmsd = np.sqrt(np.abs(2.0 * (E0 - max_eigenvalue) / n_atoms))
    return rmsd
<commit_msg>using array ensure<commit_after>"""
We have implemented RMSD calculation using both Kabsch algorithm and
via the quaterion-based characteristic polynomial.  Both are implemented
in pure python/numpy.

Notes
-----
This file works at a frame-by-frame level, not the trajectory level.

Reference
---------
Rapid calculation of RMSDs using a quaternion-based characteristic polynomial.
Acta Crystallogr A 61(4):478-480.
"""

import numpy as np
import scipy.optimize
from mdtraj.utils.arrays import ensure_type

class Transformation():
    def __init__(self, rotation, translation):
        self.rotation = rotation
        self.translation = translation

    def transform(self, xyz):
        """Apply transformation to xyz."""
        mu1 = xyz.mean(0)
        return xyz.dot(self.rotation) + self.translation - mu1.dot(self.rotation)

def compute_transformation(mobile, target):
    """Returns a Transformation object to align mobile onto target.

    Parameters
    ----------
    mobile : ndarray, shape = (n_atoms, 3)
        xyz coordinates of a SINGLE frame, to be aligned onto target.
    target : ndarray, shape = (n_atoms, 3)
        xyz coordinates of a SINGLE frame

    Returns
    -------
    T : Transformation aligning mobile to target.
    """
    translation, rotation = compute_translation_and_rotation(mobile, target)
    return Transformation(rotation, translation)

def transform(mobile, target):
    """Align mobile onto target and return transformed coordinates.

    Parameters
    ----------
    mobile : ndarray, shape = (n_atoms, 3)
        xyz coordinates of a SINGLE frame, to be aligned onto target.
    target : ndarray, shape = (n_atoms, 3)
        xyz coordinates of a SINGLE frame

    Returns
    -------
    mobile_prime : transformed coordinates of mobile.
    """
    T = compute_transformation(mobile, target)
    mobile_prime = T.transform(mobile)
    return mobile_prime

def compute_translation_and_rotation(mobile, target):
    """Returns the translation and rotation mapping mobile onto target.

    Parameters
    ----------
    mobile : ndarray, shape = (n_atoms, 3)
        xyz coordinates of a SINGLE frame, to be aligned onto target.
    target : ndarray, shape = (n_atoms, 3)
        xyz coordinates of a SINGLE frame

    Returns
    -------
    R : ndarray
        rotation
    """

    ensure_type(mobile, 'float', 2, 'mobile', shape=(None, 3))
    ensure_type(target, 'float', 2, 'mobile', shape=(target.shape[0], 3))

    mu1 = mobile.mean(0)
    mu2 = target.mean(0)

    translation = mu2

    mobile = mobile - mu1
    target = target - mu2

    correlation_matrix = np.dot(np.transpose(mobile), target)
    V, S, W_tr = np.linalg.svd(correlation_matrix)
    is_reflection = (np.linalg.det(V) * np.linalg.det(W_tr)) < 0.0
    if is_reflection:
        V[:, -1] = -V[:, -1]
    rotation = np.dot(V, W_tr)

    return translation, rotation

def rmsd_kabsch(xyz1, xyz2):
    """Returns the RMSD distance between conformations xyz1 and xyz2.

    Parameters
    ----------
    xyz1 : ndarray, shape = (n_atoms, 3)
        xyz1 coordinates of a SINGLE frame
    xyz2 : ndarray, shape = (n_atoms, 3)
        xyz2 coordinates of a SINGLE frame
    Returns
    -------
    rmsd : float
        RMSD between xyz1 and xyz2
    """
    xyz1_prime = transform(xyz1, xyz2)
    delta = xyz1_prime - xyz2
    rmsd = (delta ** 2.0).sum(1).mean() ** 0.5
    return rmsd

def _center(conformation):
    """Center the conformation"""
    ensure_type(conformation, 'float', 2, 'conformation', shape=(None, 3))
    centroid = np.mean(conformation, axis=0)
    centered = conformation - centroid
    return centered

def rmsd_qcp(conformation1, conformation2):
    """Compute the RMSD with Theobald's quaterion-based characteristic
    polynomial

    Rapid calculation of RMSDs using a quaternion-based characteristic polynomial.
    Acta Crystallogr A 61(4):478-480.

    Parameters
    ----------
    conformation1 : np.ndarray, shape=(n_atoms, 3)
        The cartesian coordinates of the first conformation
    conformation2 : np.ndarray, shape=(n_atoms, 3)
        The cartesian coordinates of the second conformation

    Returns
    -------
    rmsd : float
        The root-mean square deviation after alignment between the two pointsets
    """
    ensure_type(conformation1, 'float', 2, 'conformation1', shape=(None, 3))
    ensure_type(conformation2, 'float', 2, 'conformation2', shape=(conformation1.shape[0], 3))

    A = _center(conformation1)
    B = _center(conformation2)
    if not A.shape[0] == B.shape[0]:
        raise ValueError('conformation1 and conformation2 must have same number of atoms')
    n_atoms = len(A)

    # the inner product of the structures A and B
    G_A = np.einsum('ij,ij', A, A)
    G_B = np.einsum('ij,ij', B, B)
    # print 'GA', G_A, np.trace(np.dot(A.T, A))
    # print 'GB', G_B, np.trace(np.dot(B.T, B))

    # M is the inner product of the matrices A and B
    M = np.dot(B.T, A)

    # unpack the elements
    Sxx, Sxy, Sxz = M[0, :]
    Syx, Syy, Syz = M[1, :]
    Szx, Szy, Szz = M[2, :]

    # do some intermediate computations to assemble the characteristic
    # polynomial
    Sxx2 = Sxx * Sxx
    Syy2 = Syy * Syy
    Szz2 = Szz * Szz

    Sxy2 = Sxy * Sxy
    Syz2 = Syz * Syz
    Sxz2 = Sxz * Sxz

    Syx2 = Syx * Syx
    Szy2 = Szy * Szy
    Szx2 = Szx * Szx

    SyzSzymSyySzz2 = 2.0 * (Syz * Szy - Syy * Szz)
    Sxx2Syy2Szz2Syz2Szy2 = Syy2 + Szz2 - Sxx2 + Syz2 + Szy2

    # two of the coefficients
    C2 = -2.0 * (Sxx2 + Syy2 + Szz2 + Sxy2 + Syx2 + Sxz2 + Szx2 + Syz2 + Szy2)
    C1 = 8.0 * (Sxx * Syz * Szy + Syy * Szx * Sxz + Szz * Sxy * Syx - Sxx * Syy * Szz - Syz * Szx * Sxy - Szy * Syx * Sxz)

    SxzpSzx = Sxz + Szx
    SyzpSzy = Syz + Szy
    SxypSyx = Sxy + Syx
    SyzmSzy = Syz - Szy
    SxzmSzx = Sxz - Szx
    SxymSyx = Sxy - Syx
    SxxpSyy = Sxx + Syy
    SxxmSyy = Sxx - Syy
    Sxy2Sxz2Syx2Szx2 = Sxy2 + Sxz2 - Syx2 - Szx2

    # the other coefficient
    C0 = Sxy2Sxz2Syx2Szx2 * Sxy2Sxz2Syx2Szx2 \
        + (Sxx2Syy2Szz2Syz2Szy2 + SyzSzymSyySzz2) * (Sxx2Syy2Szz2Syz2Szy2 - SyzSzymSyySzz2) \
        + (-(SxzpSzx) * (SyzmSzy) + (SxymSyx) * (SxxmSyy - Szz)) * (-(SxzmSzx) * (SyzpSzy) + (SxymSyx) * (SxxmSyy + Szz)) \
        + (-(SxzpSzx) * (SyzpSzy) - (SxypSyx) * (SxxpSyy - Szz)) * (-(SxzmSzx) * (SyzmSzy) - (SxypSyx) * (SxxpSyy + Szz)) \
        + (+(SxypSyx) * (SyzpSzy) + (SxzpSzx) * (SxxmSyy + Szz)) * (-(SxymSyx) * (SyzmSzy) + (SxzpSzx) * (SxxpSyy + Szz)) \
        + (+(SxypSyx) * (SyzmSzy) + (SxzmSzx) * (SxxmSyy - Szz)) * (-(SxymSyx) * (SyzpSzy) + (SxzmSzx) * (SxxpSyy - Szz))

    E0 = (G_A + G_B) / 2.0
    f = lambda x: x ** 4.0 + C2 * x ** 2. + C1 * x + C0
    df = lambda x: 4 * x ** 3.0 + 2 * C2 * x + C1
    max_eigenvalue = scipy.optimize.newton(f, E0, df)
    rmsd = np.sqrt(np.abs(2.0 * (E0 - max_eigenvalue) / n_atoms))
    return rmsd
<|endoftext|>


<commit_before>##############################################################################
# MDTraj: A Python Library for Loading, Saving, and Manipulating
#         Molecular Dynamics Trajectories.
# Copyright 2012-2014 Stanford University and the Authors
#
# Authors: Kyle A. Beauchamp
# Contributors: Robert McGibbon
#
# MDTraj is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as
# published by the Free Software Foundation, either version 2.1
# of the License, or (at your option) any later version.
#
# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with MDTraj. If not, see <http://www.gnu.org/licenses/>.
#############################################################################

##############################################################################
# Imports
##############################################################################

import os
from distutils.spawn import find_executable as _find_executable
import numpy as np

from mdtraj.utils import import_
from mdtraj.utils import enter_temp_directory

##############################################################################
# Globals
##############################################################################

# Possible names for the external commands -- these are expected
# to be found in the PATH.
SHIFTX2 = ['shiftx2.py']
SPARTA_PLUS = ['sparta+', 'SPARTA+', 'SPARTA+.linux']
PPM = ['ppm_linux_64.exe']

__all__ = ['chemical_shifts_shiftx2', 'chemical_shifts_ppm', 'chemical_shifts_spartaplus', "reindex_dataframe_by_atoms"]

def find_executable(names):
    for possible in names:
        result = _find_executable(possible)
        if result is not None:
            return result
    return None

##############################################################################
# Code
##############################################################################

def compute_chemical_shifts(trj, model="shiftx2", **kwargs):
    """Predict chemical shifts of a trajectory using ShiftX2.

    Parameters
    ----------
    trj : Trajectory
        Trajectory to predict shifts for.
    model : str, optional, default="shiftx2"
        The program to use for calculating chemical shifts.  Must be one
        of shiftx2, ppm, or sparta+

    Returns
    -------
    results : pandas DataFrame
        Dataframe containing results, with index consisting of
        (resSeq, atom_name) pairs and columns for each frame in trj.

    Notes
    -----
    You must have the appropriate chemical soft programs installed
    and in your executable path.

    Chemical shift prediction is for PROTEIN atoms; trajectory objects
    with ligands, solvent, ions, or other non-protein components may give
    UNKNOWN RESULTS.

    Please cite the appropriate reference, see docstrings for chemical_shifts_*
    for the various possible models.
    """
    if model == "shiftx2":
        return chemical_shifts_shiftx2(trj, **kwargs)
    elif model == "ppm":
        return chemical_shifts_ppm(trj, **kwargs)
    elif model == "sparta+":
        return chemical_shifts_spartaplus(trj, **kwargs)
    else:
        raise(ValueError("model must be one of shiftx2, ppm, or sparta+"))

def chemical_shifts_shiftx2(trj, pH=7.0, temperature=298.15):
    """Predict chemical shifts of a trajectory using ShiftX2.

    Parameters
    ----------
    trj : Trajectory
        Trajectory to predict shifts for.

    Returns
    -------
    results : pandas DataFrame
        Dataframe containing results, with index consisting of
        (resSeq, atom_name) pairs and columns for each frame in trj.

    Notes
    -----
    You must have ShiftX2 available on your path; see (http://www.shiftx2.ca/).

    Chemical shift prediction is for PROTEIN atoms; trajectory objects
    with ligands, solvent, ions, or other non-protein components may give
    UNKNOWN RESULTS.

    Please cite the appropriate reference below.

    References
    ----------
    .. [1] Beomsoo Han, Yifeng Liu, Simon Ginzinger, and David Wishart.
       "SHIFTX2: significantly improved protein chemical shift
       prediction." J. Biomol. NMR, 50, 1 43-57 (2011)
    """
    pd = import_('pandas')
    binary = find_executable(SHIFTX2)
    if binary is None:
        raise OSError('External command not found. Looked for %s in PATH. `chemical_shifts_shiftx2` requires the external program SHIFTX2, available at http://www.shiftx2.ca/' % ', '.join(SHIFTX2))

    results = []
    with enter_temp_directory():
        for i in range(trj.n_frames):
            trj[i].save("./trj%d.pdb" % i)
        cmd = "%s -b 'trj*.pdb' -p %.1f -t %.2f" % (binary, pH, temperature)

        return_flag = os.system(cmd)

        if return_flag != 0:
            raise(IOError("Could not successfully execute command '%s', check your ShiftX2 installation or your input trajectory." % cmd))

        for i in range(trj.n_frames):
            d = pd.read_csv("./trj%d.pdb.cs" % i)
            d.rename(columns={"NUM": "resSeq", "RES": "resName", "ATOMNAME": "name"}, inplace=True)
            d["frame"] = i
            results.append(d)

    results = pd.concat(results)
    results = results.pivot_table(rows=["resSeq", "name"], cols="frame", values="SHIFT")
    return results

def chemical_shifts_ppm(trj):
    """Predict chemical shifts of a trajectory using ppm.

    Parameters
    ----------
    trj : Trajectory
        Trajectory to predict shifts for.

    Returns
    -------
    results : pandas.DataFrame
        Dataframe containing results, with index consisting of
        (resSeq, atom_name) pairs and columns for each frame in trj.

    Notes
    -----
    You must have ppm available on your path; see
    (http://spin.ccic.ohio-state.edu/index.php/download/index).

    Chemical shift prediction is for PROTEIN atoms; trajectory objects
    with ligands, solvent, ions, or other non-protein components may give
    UNKNOWN RESULTS.

    Please cite the appropriate reference below.

    References
    ----------
    .. [1] Li, DW, and Bruschweiler, R. "PPM: a side-chain and backbone chemical
       shift predictor for the assessment of protein conformational ensembles."
       J Biomol NMR. 2012 Nov;54(3):257-65.
    """
    pd = import_('pandas')
    binary = find_executable(PPM)

    first_resSeq = trj.top.residue(0).resSeq

    if binary is None:
        raise OSError('External command not found. Looked for %s in PATH. `chemical_shifts_ppm` requires the external program PPM, available at http://spin.ccic.ohio-state.edu/index.php/download/index' % ', '.join(PPM))

    with enter_temp_directory():
        trj.save("./trj.pdb")
        cmd = "%s -pdb trj.pdb -mode detail" % binary

        return_flag = os.system(cmd)

        if return_flag != 0:
            raise(IOError("Could not successfully execute command '%s', check your PPM installation or your input trajectory." % cmd))

        d = pd.read_table("./bb_details.dat", index_col=False, header=None, sep="\s*").drop([3], axis=1)

        d = d.rename(columns={0: "resSeq", 1: "resName", 2: "name"})
        d["resSeq"] += first_resSeq - 1  # Fix bug in PPM that reindexes to 1
        d = d.drop("resName", axis=1)
        d = d.set_index(["resSeq", "name"])
        d.columns = np.arange(trj.n_frames)
        d.columns.name = "frame"

    return d

def _get_lines_to_skip(filename):
    """Determine the number of comment lines in a SPARTA+ output file."""
    format_string = """FORMAT %4d %4s %4s %9.3f %9.3f %9.3f %9.3f %9.3f %9.3f"""
    handle = open(filename)
    for i, line in enumerate(handle):
        if line.find(format_string) != -1:
            return i + 2

    raise(Exception("No format string found in SPARTA+ file!"))

def chemical_shifts_spartaplus(trj, rename_HN=True):
    """Predict chemical shifts of a trajectory using SPARTA+.

    Parameters
    ----------
    trj : Trajectory
        Trajectory to predict shifts for.
    rename_HN : bool, optional, default=True
        SPARTA+ calls the amide proton "HN" instead of the standard "H".
        When True, this option renames the output as "H" to match the PDB
        and BMRB nomenclature.

    Returns
    -------
    results : pandas.DataFrame
        Dataframe containing results, with index consisting of
        (resSeq, atom_name) pairs and columns for each frame in trj.

    Notes
    -----
    You must have SPARTA+ available on your path; see
    (http://spin.niddk.nih.gov/bax/software/SPARTA+/). Also, the SPARTAP_DIR
    environment variable must be set so that SPARTA+ knows where to find
    its database files.

    Chemical shift prediction is for PROTEIN atoms; trajectory objects
    with ligands, solvent, ions, or other non-protein components may give
    UNKNOWN RESULTS.

    Please cite the appropriate reference below.

    References
    ----------
    .. [1] Shen, Y., and Bax, Ad. "SPARTA+: a modest improvement in empirical
       NMR chemical shift prediction by means of an artificial neural network."
       J. Biomol. NMR, 48, 13-22 (2010)
    """
    pd = import_('pandas')
    binary = find_executable(SPARTA_PLUS)
    if binary is None:
        raise OSError('External command not found. Looked for %s in PATH. `chemical_shifts_spartaplus` requires the external program SPARTA+, available at http://spin.niddk.nih.gov/bax/software/SPARTA+/' % ', '.join(SPARTA_PLUS))

    names = ["resSeq", "resName", "name", "SS_SHIFT", "SHIFT", "RC_SHIFT", "HM_SHIFT", "EF_SHIFT", "SIGMA"]

    with enter_temp_directory():
        for i in range(trj.n_frames):
            trj[i].save("./trj%d.pdb" % i)

        cmd = "%s -in %s" % (binary, ' '.join("trj%d.pdb" % i for i in range(trj.n_frames)))

        return_flag = os.system(cmd)

        if return_flag != 0:
            raise(IOError("Could not successfully execute command '%s', check your SPARTA+ installation or your input trajectory." % cmd))

        lines_to_skip = _get_lines_to_skip("trj0_pred.tab")

        results = []
        for i in range(trj.n_frames):
            d = pd.read_table("./trj%d_pred.tab" % i, names=names, header=None, sep="\s*", skiprows=lines_to_skip)
            d["frame"] = i
            results.append(d)

    results = pd.concat(results)

    if rename_HN:
        results.name[results.name == "HN"] = "H"

    results = results.pivot_table(rows=["resSeq", "name"], cols="frame", values="SHIFT")

    return results

def reindex_dataframe_by_atoms(trj, frame):
    """Reindex chemical shift output to use atom number (serial) indexing.

    Parameters
    ----------
    trj : Trajectory
        Trajectory to predict shifts for.
    results : pandas.DataFrame
        Dataframe containing results, with index consisting of
        (resSeq, atom_name) pairs and columns for each frame in trj.

    Returns
    -------
    new_frame : pandas.DataFrame
        Dataframe containing results, with index consisting of atom
        indices (AKA the 'serial' entry in a PDB).  Columns correspond to
        each frame in trj.

    Notes
    -----
    Be aware that this function may DROP predictions if the atom naming
    is different between the input trajectory and the output of various
    chemical shift prediction tools.
    """

    top, bonds = trj.top.to_dataframe()
    top["serial"] = top.index
    top = top.set_index(["resSeq", "name"])

    new_frame = frame.copy()

    new_frame["serial"] = top.ix[new_frame.index].serial
    new_frame = new_frame.dropna().reset_index().set_index("serial").drop(["resSeq", "name"], axis=1)

    return new_frame
<commit_msg>pH and temperature arguments set to ShiftX2 defaults.<commit_after>##############################################################################
# MDTraj: A Python Library for Loading, Saving, and Manipulating
#         Molecular Dynamics Trajectories.
# Copyright 2012-2014 Stanford University and the Authors
#
# Authors: Kyle A. Beauchamp
# Contributors: Robert McGibbon
#
# MDTraj is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as
# published by the Free Software Foundation, either version 2.1
# of the License, or (at your option) any later version.
#
# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with MDTraj. If not, see <http://www.gnu.org/licenses/>.
#############################################################################

##############################################################################
# Imports
##############################################################################

import os
from distutils.spawn import find_executable as _find_executable
import numpy as np

from mdtraj.utils import import_
from mdtraj.utils import enter_temp_directory

##############################################################################
# Globals
##############################################################################

# Possible names for the external commands -- these are expected
# to be found in the PATH.
SHIFTX2 = ['shiftx2.py']
SPARTA_PLUS = ['sparta+', 'SPARTA+', 'SPARTA+.linux']
PPM = ['ppm_linux_64.exe']

__all__ = ['chemical_shifts_shiftx2', 'chemical_shifts_ppm', 'chemical_shifts_spartaplus', "reindex_dataframe_by_atoms"]

def find_executable(names):
    for possible in names:
        result = _find_executable(possible)
        if result is not None:
            return result
    return None

##############################################################################
# Code
##############################################################################

def compute_chemical_shifts(trj, model="shiftx2", **kwargs):
    """Predict chemical shifts of a trajectory using ShiftX2.

    Parameters
    ----------
    trj : Trajectory
        Trajectory to predict shifts for.
    model : str, optional, default="shiftx2"
        The program to use for calculating chemical shifts.  Must be one
        of shiftx2, ppm, or sparta+

    Returns
    -------
    results : pandas DataFrame
        Dataframe containing results, with index consisting of
        (resSeq, atom_name) pairs and columns for each frame in trj.

    Notes
    -----
    You must have the appropriate chemical soft programs installed
    and in your executable path.

    Chemical shift prediction is for PROTEIN atoms; trajectory objects
    with ligands, solvent, ions, or other non-protein components may give
    UNKNOWN RESULTS.

    Please cite the appropriate reference, see docstrings for chemical_shifts_*
    for the various possible models.
    """
    if model == "shiftx2":
        return chemical_shifts_shiftx2(trj, **kwargs)
    elif model == "ppm":
        return chemical_shifts_ppm(trj, **kwargs)
    elif model == "sparta+":
        return chemical_shifts_spartaplus(trj, **kwargs)
    else:
        raise(ValueError("model must be one of shiftx2, ppm, or sparta+"))

def chemical_shifts_shiftx2(trj, pH=5.0, temperature=298.00):
    """Predict chemical shifts of a trajectory using ShiftX2.

    Parameters
    ----------
    trj : Trajectory
        Trajectory to predict shifts for.

    Returns
    -------
    results : pandas DataFrame
        Dataframe containing results, with index consisting of
        (resSeq, atom_name) pairs and columns for each frame in trj.

    Notes
    -----
    You must have ShiftX2 available on your path; see (http://www.shiftx2.ca/).

    Chemical shift prediction is for PROTEIN atoms; trajectory objects
    with ligands, solvent, ions, or other non-protein components may give
    UNKNOWN RESULTS.

    Please cite the appropriate reference below.

    References
    ----------
    .. [1] Beomsoo Han, Yifeng Liu, Simon Ginzinger, and David Wishart.
       "SHIFTX2: significantly improved protein chemical shift
       prediction." J. Biomol. NMR, 50, 1 43-57 (2011)
    """
    pd = import_('pandas')
    binary = find_executable(SHIFTX2)
    if binary is None:
        raise OSError('External command not found. Looked for %s in PATH. `chemical_shifts_shiftx2` requires the external program SHIFTX2, available at http://www.shiftx2.ca/' % ', '.join(SHIFTX2))

    results = []
    with enter_temp_directory():
        for i in range(trj.n_frames):
            trj[i].save("./trj%d.pdb" % i)
        cmd = "%s -b 'trj*.pdb' -p %.1f -t %.2f" % (binary, pH, temperature)

        return_flag = os.system(cmd)

        if return_flag != 0:
            raise(IOError("Could not successfully execute command '%s', check your ShiftX2 installation or your input trajectory." % cmd))

        for i in range(trj.n_frames):
            d = pd.read_csv("./trj%d.pdb.cs" % i)
            d.rename(columns={"NUM": "resSeq", "RES": "resName", "ATOMNAME": "name"}, inplace=True)
            d["frame"] = i
            results.append(d)

    results = pd.concat(results)
    results = results.pivot_table(rows=["resSeq", "name"], cols="frame", values="SHIFT")
    return results

def chemical_shifts_ppm(trj):
    """Predict chemical shifts of a trajectory using ppm.

    Parameters
    ----------
    trj : Trajectory
        Trajectory to predict shifts for.

    Returns
    -------
    results : pandas.DataFrame
        Dataframe containing results, with index consisting of
        (resSeq, atom_name) pairs and columns for each frame in trj.

    Notes
    -----
    You must have ppm available on your path; see
    (http://spin.ccic.ohio-state.edu/index.php/download/index).

    Chemical shift prediction is for PROTEIN atoms; trajectory objects
    with ligands, solvent, ions, or other non-protein components may give
    UNKNOWN RESULTS.

    Please cite the appropriate reference below.

    References
    ----------
    .. [1] Li, DW, and Bruschweiler, R. "PPM: a side-chain and backbone chemical
       shift predictor for the assessment of protein conformational ensembles."
       J Biomol NMR. 2012 Nov;54(3):257-65.
    """
    pd = import_('pandas')
    binary = find_executable(PPM)

    first_resSeq = trj.top.residue(0).resSeq

    if binary is None:
        raise OSError('External command not found. Looked for %s in PATH. `chemical_shifts_ppm` requires the external program PPM, available at http://spin.ccic.ohio-state.edu/index.php/download/index' % ', '.join(PPM))

    with enter_temp_directory():
        trj.save("./trj.pdb")
        cmd = "%s -pdb trj.pdb -mode detail" % binary

        return_flag = os.system(cmd)

        if return_flag != 0:
            raise(IOError("Could not successfully execute command '%s', check your PPM installation or your input trajectory." % cmd))

        d = pd.read_table("./bb_details.dat", index_col=False, header=None, sep="\s*").drop([3], axis=1)

        d = d.rename(columns={0: "resSeq", 1: "resName", 2: "name"})
        d["resSeq"] += first_resSeq - 1  # Fix bug in PPM that reindexes to 1
        d = d.drop("resName", axis=1)
        d = d.set_index(["resSeq", "name"])
        d.columns = np.arange(trj.n_frames)
        d.columns.name = "frame"

    return d

def _get_lines_to_skip(filename):
    """Determine the number of comment lines in a SPARTA+ output file."""
    format_string = """FORMAT %4d %4s %4s %9.3f %9.3f %9.3f %9.3f %9.3f %9.3f"""
    handle = open(filename)
    for i, line in enumerate(handle):
        if line.find(format_string) != -1:
            return i + 2

    raise(Exception("No format string found in SPARTA+ file!"))

def chemical_shifts_spartaplus(trj, rename_HN=True):
    """Predict chemical shifts of a trajectory using SPARTA+.

    Parameters
    ----------
    trj : Trajectory
        Trajectory to predict shifts for.
    rename_HN : bool, optional, default=True
        SPARTA+ calls the amide proton "HN" instead of the standard "H".
        When True, this option renames the output as "H" to match the PDB
        and BMRB nomenclature.

    Returns
    -------
    results : pandas.DataFrame
        Dataframe containing results, with index consisting of
        (resSeq, atom_name) pairs and columns for each frame in trj.

    Notes
    -----
    You must have SPARTA+ available on your path; see
    (http://spin.niddk.nih.gov/bax/software/SPARTA+/). Also, the SPARTAP_DIR
    environment variable must be set so that SPARTA+ knows where to find
    its database files.

    Chemical shift prediction is for PROTEIN atoms; trajectory objects
    with ligands, solvent, ions, or other non-protein components may give
    UNKNOWN RESULTS.

    Please cite the appropriate reference below.

    References
    ----------
    .. [1] Shen, Y., and Bax, Ad. "SPARTA+: a modest improvement in empirical
       NMR chemical shift prediction by means of an artificial neural network."
       J. Biomol. NMR, 48, 13-22 (2010)
    """
    pd = import_('pandas')
    binary = find_executable(SPARTA_PLUS)
    if binary is None:
        raise OSError('External command not found. Looked for %s in PATH. `chemical_shifts_spartaplus` requires the external program SPARTA+, available at http://spin.niddk.nih.gov/bax/software/SPARTA+/' % ', '.join(SPARTA_PLUS))

    names = ["resSeq", "resName", "name", "SS_SHIFT", "SHIFT", "RC_SHIFT", "HM_SHIFT", "EF_SHIFT", "SIGMA"]

    with enter_temp_directory():
        for i in range(trj.n_frames):
            trj[i].save("./trj%d.pdb" % i)

        cmd = "%s -in %s" % (binary, ' '.join("trj%d.pdb" % i for i in range(trj.n_frames)))

        return_flag = os.system(cmd)

        if return_flag != 0:
            raise(IOError("Could not successfully execute command '%s', check your SPARTA+ installation or your input trajectory." % cmd))

        lines_to_skip = _get_lines_to_skip("trj0_pred.tab")

        results = []
        for i in range(trj.n_frames):
            d = pd.read_table("./trj%d_pred.tab" % i, names=names, header=None, sep="\s*", skiprows=lines_to_skip)
            d["frame"] = i
            results.append(d)

    results = pd.concat(results)

    if rename_HN:
        results.name[results.name == "HN"] = "H"

    results = results.pivot_table(rows=["resSeq", "name"], cols="frame", values="SHIFT")

    return results

def reindex_dataframe_by_atoms(trj, frame):
    """Reindex chemical shift output to use atom number (serial) indexing.

    Parameters
    ----------
    trj : Trajectory
        Trajectory to predict shifts for.
    results : pandas.DataFrame
        Dataframe containing results, with index consisting of
        (resSeq, atom_name) pairs and columns for each frame in trj.

    Returns
    -------
    new_frame : pandas.DataFrame
        Dataframe containing results, with index consisting of atom
        indices (AKA the 'serial' entry in a PDB).  Columns correspond to
        each frame in trj.

    Notes
    -----
    Be aware that this function may DROP predictions if the atom naming
    is different between the input trajectory and the output of various
    chemical shift prediction tools.
    """

    top, bonds = trj.top.to_dataframe()
    top["serial"] = top.index
    top = top.set_index(["resSeq", "name"])

    new_frame = frame.copy()

    new_frame["serial"] = top.ix[new_frame.index].serial
    new_frame = new_frame.dropna().reset_index().set_index("serial").drop(["resSeq", "name"], axis=1)

    return new_frame
<|endoftext|>


<commit_before># Michael A. Alcorn (malcorn@redhat.com)
# An implementation of the Deep Semantic Similarity Model (DSSM) found in [1].
# [1] Shen, Y., He, X., Gao, J., Deng, L., and Mesnil, G. 2014. A latent semantic model
#         with convolutional-pooling structure for information retrieval. In CIKM, pp. 101-110.
#         http://research.microsoft.com/pubs/226585/cikm2014_cdssm_final.pdf
# [2] http://research.microsoft.com/en-us/projects/dssm/
# [3] http://research.microsoft.com/pubs/238873/wsdm2015.v3.pdf

import numpy as np

from keras import backend
from keras.layers import Input, merge
from keras.layers.core import Dense, Lambda, Reshape
from keras.layers.convolutional import Convolution1D
from keras.models import Model

def R(vects):
    """
    Calculates the cosine similarity of two vectors.
    :param vects: a list of two vectors.
    :return: the cosine similarity of two vectors.
    """
    (x, y) = vects
    return backend.dot(x, backend.transpose(y)) / (x.norm(2) * y.norm(2)) # See equation (4)

LETTER_GRAM_SIZE = 3 # See section 3.2.
N_GRAM_SIZE = 3 # See section 3.2.
TOTAL_LETTER_GRAMS = int(3 * 1e4) # Determined from data. See section 3.2.
WORD_DEPTH = N_GRAM_SIZE * TOTAL_LETTER_GRAMS # See equation (1).
K = 300 # Dimensionality of the max-pooling layer. See section 3.4.
L = 128 # Dimensionality of latent semantic space. See section 3.5.
J = 4 # Number of random unclicked documents serving as negative examples for a query. See section 4.
FILTER_LENGTH = 1 # We only consider one time step for convolutions.

# Input tensors holding the query, positive (clicked) document, and negative (unclicked) documents.
# The first dimension is None because the queries and documents can vary in length.
query = Input(shape = (None, WORD_DEPTH))
pos_doc = Input(shape = (None, WORD_DEPTH))
neg_docs = [Input(shape = (None, WORD_DEPTH)) for j in range(J)]

# Query model. The paper uses separate neural nets for queries and documents (see section 5.2).

# In this step, we transform each word vector with WORD_DEPTH dimensions into its
# convolved representation with K dimensions. K is the number of kernels/filters
# being used in the operation. Essentially, the operation is taking the dot product
# of a single weight matrix (W_c) with each of the word vectors (l_t) from the
# query matrix (l_Q), adding a bias vector (b_c), and then applying the tanh function.
# That is, h_Q = tanh(W_c • l_Q + b_c). With that being said, that's not actually
# how the operation is being calculated here. To tie the weights of the weight
# matrix (W_c) together, we have to use a one-dimensional convolutional layer.
# Further, we have to transpose our query matrix (l_Q) so that time is the first
# dimension rather than the second (as described in the paper). That is, l_Q[0, :]
# represents our first word vector rather than l_Q[:, 0]. We can think of the weight
# matrix (W_c) as being similarly transposed such that each kernel is a column
# of W_c. Therefore, h_Q = tanh(l_Q • W_c + b_c) with l_Q, W_c, and b_c being
# the transposes of the matrices described in the paper.
query_conv = Convolution1D(K, FILTER_LENGTH, border_mode = "same", input_shape = (None, WORD_DEPTH), activation = "tanh")(query) # See equation (2).

# Next, we apply a max-pooling layer to the convolved query matrix. Keras provides
# its own max-pooling layers, but they cannot handle variable length input (as
# far as I can tell). As a result, I define my own max-pooling layer here. In the
# paper, the operation selects the maximum value for each row of h_Q, but, because
# we're using the transpose, we're selecting the maximum value for each column.
query_max = Lambda(lambda x: x.max(axis = 1), output_shape = (K,))(query_conv) # See section 3.4.

# In this step, we generate the semantic vector represenation of the query. This
# is a standard neural network dense layer, i.e., y = tanh(W_s • v + b_s).
query_sem = Dense(L, activation = "tanh", input_dim = K)(query_max) # See section 3.5.

# The document equivalent of the above query model.
doc_conv = Convolution1D(K, FILTER_LENGTH, border_mode = "same", input_shape = (None, WORD_DEPTH), activation = "tanh")
doc_max = Lambda(lambda x: x.max(axis = 1), output_shape = (K,))
doc_sem = Dense(L, activation = "tanh", input_dim = K)

pos_doc_conv = doc_conv(pos_doc)
neg_doc_convs = [doc_conv(neg_doc) for neg_doc in neg_docs]

pos_doc_max = doc_max(pos_doc_conv)
neg_doc_maxes = [doc_max(neg_doc_conv) for neg_doc_conv in neg_doc_convs]

pos_doc_sem = doc_sem(pos_doc_max)
neg_doc_sems = [doc_sem(neg_doc_max) for neg_doc_max in neg_doc_maxes]

# This layer calculates the cosine similarity between the semantic representations of
# a query and a document.
R_layer = Lambda(R, output_shape = (1,)) # See equation (4).

R_Q_D_p = R_layer([query_sem, pos_doc_sem]) # See equation (4).
R_Q_D_ns = [R_layer([query_sem, neg_doc_sem]) for neg_doc_sem in neg_doc_sems] # See equation (4).

concat_Rs = merge([R_Q_D_p] + R_Q_D_ns, mode = "concat")
concat_Rs = Reshape((J + 1, 1))(concat_Rs)

# In this step, we multiply each R(Q, D) value by gamma. In the paper, gamma is
# described as a smoothing factor for the softmax function, and it's set empirically
# on a held-out data set. We're going to learn gamma's value by pretending it's
# a single, 1 x 1 kernel.
with_gamma = Convolution1D(1, 1, border_mode = "same", input_shape = (J + 1, 1), activation = "linear")(concat_Rs) # See equation (5).

# Next, we exponentiate each of the gamma x R(Q, D) values.
exponentiated = Lambda(lambda x: backend.exp(x), output_shape = (J + 1,))(with_gamma) # See equation (5).
exponentiated = Reshape((J + 1,))(exponentiated)

# Finally, we use the softmax function to calculate the P(D+|Q).
prob = Lambda(lambda x: x[0][0] / backend.sum(x[0]), output_shape = (1,))(exponentiated) # See equation (5).

# We now have everything we need to define our model.
model = Model(input = [query, pos_doc] + neg_docs, output = prob)
model.compile(optimizer = "adadelta", loss = "binary_crossentropy")

# Build a random data set.
sample_size = 10
l_Qs = []
pos_l_Ds = []

for i in range(sample_size):
    query_len = np.random.randint(1, 10)
    l_Q = np.random.rand(1, query_len, WORD_DEPTH)
    l_Qs.append(l_Q)

    doc_len = np.random.randint(50, 500)
    l_D = np.random.rand(1, doc_len, WORD_DEPTH)
    pos_l_Ds.append(l_D)

neg_l_Ds = []
for i in range(sample_size):
    possibilities = list(range(sample_size))
    possibilities.remove(i)
    negatives = np.random.choice(possibilities, J)
    neg_l_Ds.append([pos_l_Ds[negative] for negative in negatives])

# Because we're using the "binary_crossentropy" loss function, we can pretend that
# we're solving a binary classification problem and that every sample is a member
# of the "1" class.
y = np.ones(1)

for i in range(sample_size):
    history = model.fit([l_Qs[i], pos_l_Ds[i]] + neg_l_Ds[i], y, nb_epoch = 1, verbose = 0)

# Here, I walk through an example of how to define a function for calculating output
# from the computational graph. Let's define a function that calculates R(Q, D+)
# for a given query and clicked document. The function depends on two inputs, query
# and pos_doc. That is, if you start at the point in the graph where R(Q, D+) is
# calculated and then backtrack as far as possible, you'll end up at two different
# starting points, query and pos_doc. As a result, we supply those inputs in a list
# to the function. This particular function only calculates a single output, but
# multiple outputs are possible (see the next example).
get_R_Q_D_p = backend.function([query, pos_doc], R_Q_D_p)
get_R_Q_D_p([l_Qs[0], pos_l_Ds[0]])

# A slightly more complex function. Notice that both neg_docs and our output are
# lists.
get_R_Q_D_ns = backend.function([query] + neg_docs, R_Q_D_ns)
get_R_Q_D_ns([l_Qs[0]] + neg_l_Ds[0])
<commit_msg>Change N_GRAM_SIZE variable name to more accurate WINDOW_SIZE.<commit_after># Michael A. Alcorn (malcorn@redhat.com)
# An implementation of the Deep Semantic Similarity Model (DSSM) found in [1].
# [1] Shen, Y., He, X., Gao, J., Deng, L., and Mesnil, G. 2014. A latent semantic model
#         with convolutional-pooling structure for information retrieval. In CIKM, pp. 101-110.
#         http://research.microsoft.com/pubs/226585/cikm2014_cdssm_final.pdf
# [2] http://research.microsoft.com/en-us/projects/dssm/
# [3] http://research.microsoft.com/pubs/238873/wsdm2015.v3.pdf

import numpy as np

from keras import backend
from keras.layers import Input, merge
from keras.layers.core import Dense, Lambda, Reshape
from keras.layers.convolutional import Convolution1D
from keras.models import Model

def R(vects):
    """
    Calculates the cosine similarity of two vectors.
    :param vects: a list of two vectors.
    :return: the cosine similarity of two vectors.
    """
    (x, y) = vects
    return backend.dot(x, backend.transpose(y)) / (x.norm(2) * y.norm(2)) # See equation (4)

LETTER_GRAM_SIZE = 3 # See section 3.2.
WINDOW_SIZE = 3 # See section 3.2.
TOTAL_LETTER_GRAMS = int(3 * 1e4) # Determined from data. See section 3.2.
WORD_DEPTH = WINDOW_SIZE * TOTAL_LETTER_GRAMS # See equation (1).
K = 300 # Dimensionality of the max-pooling layer. See section 3.4.
L = 128 # Dimensionality of latent semantic space. See section 3.5.
J = 4 # Number of random unclicked documents serving as negative examples for a query. See section 4.
FILTER_LENGTH = 1 # We only consider one time step for convolutions.

# Input tensors holding the query, positive (clicked) document, and negative (unclicked) documents.
# The first dimension is None because the queries and documents can vary in length.
query = Input(shape = (None, WORD_DEPTH))
pos_doc = Input(shape = (None, WORD_DEPTH))
neg_docs = [Input(shape = (None, WORD_DEPTH)) for j in range(J)]

# Query model. The paper uses separate neural nets for queries and documents (see section 5.2).

# In this step, we transform each word vector with WORD_DEPTH dimensions into its
# convolved representation with K dimensions. K is the number of kernels/filters
# being used in the operation. Essentially, the operation is taking the dot product
# of a single weight matrix (W_c) with each of the word vectors (l_t) from the
# query matrix (l_Q), adding a bias vector (b_c), and then applying the tanh function.
# That is, h_Q = tanh(W_c • l_Q + b_c). With that being said, that's not actually
# how the operation is being calculated here. To tie the weights of the weight
# matrix (W_c) together, we have to use a one-dimensional convolutional layer.
# Further, we have to transpose our query matrix (l_Q) so that time is the first
# dimension rather than the second (as described in the paper). That is, l_Q[0, :]
# represents our first word vector rather than l_Q[:, 0]. We can think of the weight
# matrix (W_c) as being similarly transposed such that each kernel is a column
# of W_c. Therefore, h_Q = tanh(l_Q • W_c + b_c) with l_Q, W_c, and b_c being
# the transposes of the matrices described in the paper.
query_conv = Convolution1D(K, FILTER_LENGTH, border_mode = "same", input_shape = (None, WORD_DEPTH), activation = "tanh")(query) # See equation (2).

# Next, we apply a max-pooling layer to the convolved query matrix. Keras provides
# its own max-pooling layers, but they cannot handle variable length input (as
# far as I can tell). As a result, I define my own max-pooling layer here. In the
# paper, the operation selects the maximum value for each row of h_Q, but, because
# we're using the transpose, we're selecting the maximum value for each column.
query_max = Lambda(lambda x: x.max(axis = 1), output_shape = (K,))(query_conv) # See section 3.4.

# In this step, we generate the semantic vector represenation of the query. This
# is a standard neural network dense layer, i.e., y = tanh(W_s • v + b_s).
query_sem = Dense(L, activation = "tanh", input_dim = K)(query_max) # See section 3.5.

# The document equivalent of the above query model.
doc_conv = Convolution1D(K, FILTER_LENGTH, border_mode = "same", input_shape = (None, WORD_DEPTH), activation = "tanh")
doc_max = Lambda(lambda x: x.max(axis = 1), output_shape = (K,))
doc_sem = Dense(L, activation = "tanh", input_dim = K)

pos_doc_conv = doc_conv(pos_doc)
neg_doc_convs = [doc_conv(neg_doc) for neg_doc in neg_docs]

pos_doc_max = doc_max(pos_doc_conv)
neg_doc_maxes = [doc_max(neg_doc_conv) for neg_doc_conv in neg_doc_convs]

pos_doc_sem = doc_sem(pos_doc_max)
neg_doc_sems = [doc_sem(neg_doc_max) for neg_doc_max in neg_doc_maxes]

# This layer calculates the cosine similarity between the semantic representations of
# a query and a document.
R_layer = Lambda(R, output_shape = (1,)) # See equation (4).

R_Q_D_p = R_layer([query_sem, pos_doc_sem]) # See equation (4).
R_Q_D_ns = [R_layer([query_sem, neg_doc_sem]) for neg_doc_sem in neg_doc_sems] # See equation (4).

concat_Rs = merge([R_Q_D_p] + R_Q_D_ns, mode = "concat")
concat_Rs = Reshape((J + 1, 1))(concat_Rs)

# In this step, we multiply each R(Q, D) value by gamma. In the paper, gamma is
# described as a smoothing factor for the softmax function, and it's set empirically
# on a held-out data set. We're going to learn gamma's value by pretending it's
# a single, 1 x 1 kernel.
with_gamma = Convolution1D(1, 1, border_mode = "same", input_shape = (J + 1, 1), activation = "linear")(concat_Rs) # See equation (5).

# Next, we exponentiate each of the gamma x R(Q, D) values.
exponentiated = Lambda(lambda x: backend.exp(x), output_shape = (J + 1,))(with_gamma) # See equation (5).
exponentiated = Reshape((J + 1,))(exponentiated)

# Finally, we use the softmax function to calculate the P(D+|Q).
prob = Lambda(lambda x: x[0][0] / backend.sum(x[0]), output_shape = (1,))(exponentiated) # See equation (5).

# We now have everything we need to define our model.
model = Model(input = [query, pos_doc] + neg_docs, output = prob)
model.compile(optimizer = "adadelta", loss = "binary_crossentropy")

# Build a random data set.
sample_size = 10
l_Qs = []
pos_l_Ds = []

for i in range(sample_size):
    query_len = np.random.randint(1, 10)
    l_Q = np.random.rand(1, query_len, WORD_DEPTH)
    l_Qs.append(l_Q)

    doc_len = np.random.randint(50, 500)
    l_D = np.random.rand(1, doc_len, WORD_DEPTH)
    pos_l_Ds.append(l_D)

neg_l_Ds = []
for i in range(sample_size):
    possibilities = list(range(sample_size))
    possibilities.remove(i)
    negatives = np.random.choice(possibilities, J)
    neg_l_Ds.append([pos_l_Ds[negative] for negative in negatives])

# Because we're using the "binary_crossentropy" loss function, we can pretend that
# we're solving a binary classification problem and that every sample is a member
# of the "1" class.
y = np.ones(1)

for i in range(sample_size):
    history = model.fit([l_Qs[i], pos_l_Ds[i]] + neg_l_Ds[i], y, nb_epoch = 1, verbose = 0)

# Here, I walk through an example of how to define a function for calculating output
# from the computational graph. Let's define a function that calculates R(Q, D+)
# for a given query and clicked document. The function depends on two inputs, query
# and pos_doc. That is, if you start at the point in the graph where R(Q, D+) is
# calculated and then backtrack as far as possible, you'll end up at two different
# starting points, query and pos_doc. As a result, we supply those inputs in a list
# to the function. This particular function only calculates a single output, but
# multiple outputs are possible (see the next example).
get_R_Q_D_p = backend.function([query, pos_doc], R_Q_D_p)
get_R_Q_D_p([l_Qs[0], pos_l_Ds[0]])

# A slightly more complex function. Notice that both neg_docs and our output are
# lists.
get_R_Q_D_ns = backend.function([query] + neg_docs, R_Q_D_ns)
get_R_Q_D_ns([l_Qs[0]] + neg_l_Ds[0])
<|endoftext|>


<commit_before>from selenium import selenium
import time
import sys
from ConfigParser import ConfigParser

MAX_TEST_LENGTH = 300

if len(sys.argv) > 2:
    filename = sys.argv[2]
else:
    filename = "config.cfg"

c = ConfigParser()
c.read(filename)

targets = {}

server = c.get('config', 'server')
url= c.get('config', 'url')

sections = c.sections()
for s in sections:
    if s == 'config':
       continue
    targets[s] = dict(c.items(s))
    targets[s]['name'] = s

if sys.argv[1] == "all":
    browsers = list(targets.values())
elif sys.argv[1] not in targets:
    print "Invalid target"
    sys.exit()
else:
    browsers = [targets[sys.argv[1]]]

if 1:
    for b in browsers:
        print "Running %s on %s" % (b['name'], b['host'])
        s = selenium(b['host'], 4444, "*%s" % b['browsercmd'], server)
        s.start()
        try:
            s.open(url)

            count = 0
            while count == 0:
                count = int(s.get_eval("window.document.getElementById('testtable').getElementsByTagName('tr').length"))
                time.sleep(5)

            ok = 0
            fail = 0
            last_change = time.time()
            while True:
                new_ok = int(s.get_eval('window.Test.AnotherWay._g_ok_pages'))
                new_fail = int(s.get_eval('window.Test.AnotherWay._g_fail_pages'))
                if new_ok != ok or new_fail != fail:
                    ok = new_ok
                    fail = new_fail
                    last_change = time.time()

                if (ok + fail) >= count:
                    break
                if time.time() - last_change > MAX_TEST_LENGTH:
                    raise Exception("Failed: with %s okay and %s failed, ran out of time: %s is more than %s" % (ok, fail, (time.time() - last_change), MAX_TEST_LENGTH))
                time.sleep(10)

            if fail:
                print "Failed: %s" % fail
                html = s.get_eval("window.document.getElementById('results').innerHTML").encode("utf-8")
                all_html = """

    <meta content="text/html; charset=utf-8" http-equiv="content-type" />

  %s""" % html

                f = open("fail.%s.%s.html" % (time.time(), b['name']), "w")
                f.write(all_html)
                f.close()
        except Exception, E:
            print "Error: ", E
        s.stop()

<commit_msg>Capture KeyboardInterrupt, and stop all tests. And configurability of max test timeout<commit_after>from selenium import selenium
import time
import sys
from ConfigParser import ConfigParser

MAX_TEST_LENGTH = 300
if len(sys.argv) > 2:
    filename = sys.argv[2]
else:
    filename = "config.cfg"

c = ConfigParser()
c.read(filename)

targets = {}

server = c.get('config', 'server')
url= c.get('config', 'url')
if c.has_option('config', 'timeout'):
    MAX_TEST_LENGTH = int(c.get('config', 'timeout'))

sections = c.sections()
for s in sections:
    if s == 'config':
       continue
    targets[s] = dict(c.items(s))
    targets[s]['name'] = s

if sys.argv[1] == "all":
    browsers = list(targets.values())
elif sys.argv[1] not in targets:
    print "Invalid target"
    sys.exit()
else:
    browsers = [targets[sys.argv[1]]]

keep_going = True

if 1:
    for b in browsers:
        if not keep_going:
            continue

        print "Running %s on %s" % (b['name'], b['host'])
        s = selenium(b['host'], 4444, "*%s" % b['browsercmd'], server)
        s.start()
        try:
            s.open_window(url, "test_running")
            time.sleep(2)
            s.select_window("test_running")
            time.sleep(2)
            s.refresh()

            count = 0
            while count == 0:
                count = int(s.get_eval("window.document.getElementById('testtable').getElementsByTagName('tr').length"))
                time.sleep(5)

            ok = 0
            fail = 0
            last_change = time.time()
            while True:
                new_ok = int(s.get_eval('window.Test.AnotherWay._g_ok_pages'))
                new_fail = int(s.get_eval('window.Test.AnotherWay._g_fail_pages'))
                if new_ok != ok or new_fail != fail:
                    ok = new_ok
                    fail = new_fail
                    last_change = time.time()

                if (ok + fail) >= count:
                    break
                if time.time() - last_change > MAX_TEST_LENGTH:
                    raise Exception("Failed: with %s okay and %s failed, ran out of time: %s is more than %s" % (ok, fail, (time.time() - last_change), MAX_TEST_LENGTH))
                time.sleep(10)

            if fail:
                print "Failed: %s" % fail
                html = s.get_eval("window.document.getElementById('results').innerHTML").encode("utf-8")
                all_html = """

    <meta content="text/html; charset=utf-8" http-equiv="content-type" />

  %s""" % html

                f = open("fail.%s.%s.html" % (time.time(), b['name']), "w")
                f.write(all_html)
                f.close()
        except KeyboardInterrupt, E:
            keep_going = False
            print "Stopped by keyboard interrupt"
        except Exception, E:
            print "Error: ", E
        s.stop()

<|endoftext|>


<commit_before>#! /usr/bin/env python

# ======================================================================
# matscipy - Python materials science tools
# https://github.com/libAtoms/matscipy
#
# Copyright (2014) James Kermode, King's College London
#                  Lars Pastewka, Karlsruhe Institute of Technology
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
# ======================================================================

import unittest
import numpy as np
from matscipy.neighbours import neighbour_list
from ase import Atoms
import matscipytest
from matscipy import spatial_correlation_functional

class TestSpatialCorrelationFunctional(unittest.TestCase):

    def test_peak_count(self):
        n=50

        xyz=np.zeros((n,3))
        xyz[:,0]=np.arange(n)
        values=np.random.rand(len(xyz))

        atoms=Atoms(positions=xyz)
        cell=np.array([[n,0,0],[0,n,0],[0,0,n]])

        length_cutoff= n/2.
        output_gridsize= 0.1
        FFT_cutoff= 7.5
        approx_FFT_gridsize= 1.
        SCF=spatial_correlation_functional.spatial_correlation_functional(atoms, values, cell, length_cutoff, output_gridsize, FFT_cutoff, approx_FFT_gridsize,dim=None,delta='simple',norm=True)
        SCF2=spatial_correlation_functional.spatial_correlation_functional(atoms, values, cell, length_cutoff, output_gridsize, FFT_cutoff, approx_FFT_gridsize/50.*n,dim=None,delta='simple',norm=True)
        SCF3=spatial_correlation_functional.spatial_correlation_functional(atoms, values, cell, length_cutoff, output_gridsize, FFT_cutoff, approx_FFT_gridsize/20.*n,dim=None,delta='simple',norm=True)

        SCF-=SCF.min()
        SCF2-=SCF2.min()
        SCF3-=SCF3.min()

        self.assertTrue((np.isfinite(SCF/SCF)).sum()==np.floor(length_cutoff))
        self.assertTrue((np.isfinite(SCF2/SCF2)).sum()==int(np.floor(FFT_cutoff)+np.ceil(np.ceil(length_cutoff-FFT_cutoff)*50./n)))
        self.assertTrue((np.isfinite(SCF3/SCF3)).sum()==int(np.floor(FFT_cutoff)+np.ceil(np.ceil(length_cutoff-FFT_cutoff)*20./n)))

    def test_directional_spacing(self):
        n=50

        xyz=np.zeros((n**3,3))
        m=np.meshgrid(np.arange(0,n,1),np.arange(0,2*n,2),np.arange(0,3*n,3))
        xyz[:,0]=m[0].reshape((-1))
        xyz[:,1]=m[0].reshape((-1))
        xyz[:,2]=m[2].reshape((-1))

        values=np.random.rand(len(xyz))
        atoms=Atoms(positions=xyz)
        cell=np.array([[n,0,0],[0,2*n,0],[0,0,3*n]])

        length_cutoff= n
        output_gridsize= 0.1
        FFT_cutoff= 0.
        approx_FFT_gridsize= 1.

        SCF0=spatial_correlation_functional.spatial_correlation_functional(atoms, values, cell, length_cutoff, output_gridsize, FFT_cutoff, approx_FFT_gridsize, dim=0, delta='simple', norm=True)
        SCF1=spatial_correlation_functional.spatial_correlation_functional(atoms, values, cell, length_cutoff, output_gridsize, FFT_cutoff, approx_FFT_gridsize, dim=1, delta='simple', norm=True)
        SCF2=spatial_correlation_functional.spatial_correlation_functional(atoms, values, cell, length_cutoff, output_gridsize, FFT_cutoff, approx_FFT_gridsize, dim=2, delta='simple', norm=True)

        SCF0-=SCF0.min()
        SCF1-=SCF1.min()
        SCF2-=SCF2.min()

        n_peaks0=np.isfinite(SCF0/SCF0).sum()
        n_peaks1=np.isfinite(SCF1/SCF1).sum()
        n_peaks2=np.isfinite(SCF2/SCF2).sum()

        self.assertTrue(n_peaks0/2.-n_peaks1 < 2)
        self.assertTrue(n_peaks0/3.-n_peaks2 < 2)

###
if __name__ == '__main__':
    unittest.main()
<commit_msg>test too slow -> speedup<commit_after>#! /usr/bin/env python

# ======================================================================
# matscipy - Python materials science tools
# https://github.com/libAtoms/matscipy
#
# Copyright (2014) James Kermode, King's College London
#                  Lars Pastewka, Karlsruhe Institute of Technology
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
# ======================================================================

import unittest
import numpy as np
from matscipy.neighbours import neighbour_list
from ase import Atoms
import matscipytest
from matscipy import spatial_correlation_functional

class TestSpatialCorrelationFunctional(unittest.TestCase):

    def test_peak_count(self):
        n=30

        xyz=np.zeros((n,3))
        xyz[:,0]=np.arange(n)
        values=np.random.rand(len(xyz))

        atoms=Atoms(positions=xyz)
        cell=np.array([[n,0,0],[0,n,0],[0,0,n]])

        length_cutoff= n/2.
        output_gridsize= 0.1
        FFT_cutoff= 7.5
        approx_FFT_gridsize= 1.
        SCF=spatial_correlation_functional.spatial_correlation_functional(atoms, values, cell, length_cutoff, output_gridsize, FFT_cutoff, approx_FFT_gridsize,dim=None,delta='simple',norm=True)
        SCF2=spatial_correlation_functional.spatial_correlation_functional(atoms, values, cell, length_cutoff, output_gridsize, FFT_cutoff, approx_FFT_gridsize/50.*n,dim=None,delta='simple',norm=True)
        SCF3=spatial_correlation_functional.spatial_correlation_functional(atoms, values, cell, length_cutoff, output_gridsize, FFT_cutoff, approx_FFT_gridsize/20.*n,dim=None,delta='simple',norm=True)

        SCF-=SCF.min()
        SCF2-=SCF2.min()
        SCF3-=SCF3.min()

        self.assertTrue((np.isfinite(SCF/SCF)).sum()==np.floor(length_cutoff))
        self.assertTrue((np.isfinite(SCF2/SCF2)).sum()==int(np.floor(FFT_cutoff)+np.ceil(np.ceil(length_cutoff-FFT_cutoff)*50./n)))
        self.assertTrue((np.isfinite(SCF3/SCF3)).sum()==int(np.floor(FFT_cutoff)+np.ceil(np.ceil(length_cutoff-FFT_cutoff)*20./n)))

    def test_directional_spacing(self):
        n=30

        xyz=np.zeros((n**3,3))
        m=np.meshgrid(np.arange(0,n,1),np.arange(0,2*n,2),np.arange(0,3*n,3))
        xyz[:,0]=m[0].reshape((-1))
        xyz[:,1]=m[0].reshape((-1))
        xyz[:,2]=m[2].reshape((-1))

        values=np.random.rand(len(xyz))
        atoms=Atoms(positions=xyz)
        cell=np.array([[n,0,0],[0,2*n,0],[0,0,3*n]])

        length_cutoff= n
        output_gridsize= 0.1
        FFT_cutoff= 0.
        approx_FFT_gridsize= 1.

        SCF0=spatial_correlation_functional.spatial_correlation_functional(atoms, values, cell, length_cutoff, output_gridsize, FFT_cutoff, approx_FFT_gridsize, dim=0, delta='simple', norm=True)
        SCF1=spatial_correlation_functional.spatial_correlation_functional(atoms, values, cell, length_cutoff, output_gridsize, FFT_cutoff, approx_FFT_gridsize, dim=1, delta='simple', norm=True)
        SCF2=spatial_correlation_functional.spatial_correlation_functional(atoms, values, cell, length_cutoff, output_gridsize, FFT_cutoff, approx_FFT_gridsize, dim=2, delta='simple', norm=True)

        SCF0-=SCF0.min()
        SCF1-=SCF1.min()
        SCF2-=SCF2.min()

        n_peaks0=np.isfinite(SCF0/SCF0).sum()
        n_peaks1=np.isfinite(SCF1/SCF1).sum()
        n_peaks2=np.isfinite(SCF2/SCF2).sum()

        self.assertTrue(n_peaks0/2.-n_peaks1 < 2)
        self.assertTrue(n_peaks0/3.-n_peaks2 < 2)

###
if __name__ == '__main__':
    unittest.main()
<|endoftext|>


<commit_before>#! /usr/bin/env python

# ======================================================================
# matscipy - Python materials science tools
# https://github.com/libAtoms/matscipy
#
# Copyright (2014) James Kermode, King's College London
#                  Lars Pastewka, Karlsruhe Institute of Technology
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
# ======================================================================

import unittest
import numpy as np
from matscipy.neighbours import neighbour_list
from ase import Atoms
import matscipytest
from matscipy import spatial_correlation_functional

class TestSpatialCorrelationFunctional(unittest.TestCase):

    def test_peak_count(self):
        n=100

        xyz=np.zeros((n,3))
        xyz[:,0]=np.arange(n)
        values=np.random.rand(len(xyz))

        atoms=Atoms(positions=xyz)
        cell=np.array([[n,0,0],[0,n,0],[0,0,n]])

        length_cutoff= n/2.
        output_gridsize= 0.1
        FFT_cutoff= 7.5
        approx_FFT_gridsize= 1.
        SCF=spatial_correlation_functional.spatial_correlation_functional(atoms, values, cell, length_cutoff, output_gridsize, FFT_cutoff, approx_FFT_gridsize)
        SCF2=spatial_correlation_functional.spatial_correlation_functional(atoms, values, cell, length_cutoff, output_gridsize, FFT_cutoff, approx_FFT_gridsize/50.*n)
        SCF3=spatial_correlation_functional.spatial_correlation_functional(atoms, values, cell, length_cutoff, output_gridsize, FFT_cutoff, approx_FFT_gridsize/20.*n)

        SCF-=SCF.min()
        SCF2-=SCF2.min()
        SCF3-=SCF3.min()

        self.assertTrue((np.isfinite(SCF/SCF)).sum()==np.floor(length_cutoff))
        self.assertTrue((np.isfinite(SCF2/SCF2)).sum()==int(np.floor(FFT_cutoff)+np.ceil(np.ceil(length_cutoff-FFT_cutoff)*50./n)))
        self.assertTrue((np.isfinite(SCF3/SCF3)).sum()==int(np.floor(FFT_cutoff)+np.ceil(np.ceil(length_cutoff-FFT_cutoff)*20./n)))

###
if __name__ == '__main__':
    unittest.main()
<commit_msg>syntax to fit new version of function<commit_after>#! /usr/bin/env python

# ======================================================================
# matscipy - Python materials science tools
# https://github.com/libAtoms/matscipy
#
# Copyright (2014) James Kermode, King's College London
#                  Lars Pastewka, Karlsruhe Institute of Technology
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
# ======================================================================

import unittest
import numpy as np
from matscipy.neighbours import neighbour_list
from ase import Atoms
import matscipytest
from matscipy import spatial_correlation_functional

class TestSpatialCorrelationFunctional(unittest.TestCase):

    def test_peak_count(self):
        n=100

        xyz=np.zeros((n,3))
        xyz[:,0]=np.arange(n)
        values=np.random.rand(len(xyz))

        atoms=Atoms(positions=xyz)
        cell=np.array([[n,0,0],[0,n,0],[0,0,n]])

        length_cutoff= n/2.
        output_gridsize= 0.1
        FFT_cutoff= 7.5
        approx_FFT_gridsize= 1.
        SCF=spatial_correlation_functional.spatial_correlation_functional(atoms, values, cell, length_cutoff, output_gridsize, FFT_cutoff, approx_FFT_gridsize,dim=None,delta='simple',norm=True)
        SCF2=spatial_correlation_functional.spatial_correlation_functional(atoms, values, cell, length_cutoff, output_gridsize, FFT_cutoff, approx_FFT_gridsize/50.*n,dim=None,delta='simple',norm=True)
        SCF3=spatial_correlation_functional.spatial_correlation_functional(atoms, values, cell, length_cutoff, output_gridsize, FFT_cutoff, approx_FFT_gridsize/20.*n,dim=None,delta='simple',norm=True)

        SCF-=SCF.min()
        SCF2-=SCF2.min()
        SCF3-=SCF3.min()

        self.assertTrue((np.isfinite(SCF/SCF)).sum()==np.floor(length_cutoff))
        self.assertTrue((np.isfinite(SCF2/SCF2)).sum()==int(np.floor(FFT_cutoff)+np.ceil(np.ceil(length_cutoff-FFT_cutoff)*50./n)))
        self.assertTrue((np.isfinite(SCF3/SCF3)).sum()==int(np.floor(FFT_cutoff)+np.ceil(np.ceil(length_cutoff-FFT_cutoff)*20./n)))

###
if __name__ == '__main__':
    unittest.main()
<|endoftext|>