Source code for schrodinger.application.jaguar.scan

"""
Support functions and classes for Jaguar scans.

Copyright Schrodinger, LLC. All rights reserved.

"""

import math
import re
from past.utils import old_div

__cvs_version__ = "$Revision: 1.2 $"
_version = __cvs_version__


class Scan(object):
    """
    A class for storing information about geometry scan variables.

    """

    precision = 1.0e-3

    def __init__(self,
                 var_name,
                 initial=None,
                 final=None,
                 steps=None,
                 step_size=None,
                 at_values=None):
        """
        This constructor can be used for two different types of scans -
        regular and irregular. For regular scans, the 'initial' argument
        must be specified along with two out of the three arguments 'final,'
        'steps,' or 'step_size.' For irregular scans, the only argument that
        should be specified is 'at_values.'

        Attributes

        var_name (str)
            The name of the variable the scan is based on.

        initial (float)
            The starting point for a scan.

        final (float)
            The final point for a scan.

        steps (float)
            The number of steps in the scan

        step_size (float)
            The step size of a regular scan.

        at_values (sequence of floats)
            Specific values at which a

        """
        self.var_name = var_name
        self._initial = initial
        self._final = final
        self._steps = steps
        self._step_size = step_size
        self._at_values = at_values

    def _getFinal(self):
        if self._final is not None:
            return self._final
        elif self._at_values:
            return self._at_values[-1]
        else:
            return self._initial + (self._steps - 1) * self._step_size

    final = property(_getFinal, doc="final scan value")

    def _getInitial(self):
        if self._initial is not None:
            return self._initial
        else:
            return self._at_values[0]

    initial = property(_getInitial, doc="initial scan value")

    def _getSteps(self):
        if self._steps is not None:
            return self._steps
        elif self._at_values:
            return len(self._at_values)
        else:
            return 1 + int(
                math.floor(
                    old_div((self._final - self._initial), self._step_size)))

    steps = property(_getSteps, doc="number of steps in the scan")

    def _getStepSize(self):
        if self._step_size is not None:
            return self._step_size
        elif self._at_values:
            return None
        else:
            return old_div((self._final - self._initial), (self._steps - 1))

    step_size = property(_getStepSize,
                         doc="step size (None for 'at values' scans)")

    def _getAtValues(self):
        if self._at_values:
            return self._at_values
        else:
            return []

    at_values = property(
        _getAtValues, doc="an array of the scan values for 'at values' scans")

    def __eq__(self, other):
        """
        Equality comparison only looks at numeric quantities, not the
        variable name.

        """
        if self.steps != other.steps:
            return False
        if abs(self.initial - other.initial) > self.precision:
            return False
        if abs(self.final - other.final) > self.precision:
            return False
        if abs(self.step_size - other.step_size) > self.precision:
            return False
        return True

    def __ne__(self, other):
        return not self.__eq__(other)

    @staticmethod
    def parse(zvarstring):
        """
        Parse a scan definition string and return a Scan object.

        """
        # This is defined as a static method (as opposed to a module level
        # function) to allow it to be used when 'from scan import Scan' is
        # the import method.

        zvarsplit = zvarstring.split("=")
        if len(zvarsplit) == 2:
            zvarname, assignment = zvarsplit
        else:
            # The output files print 'ohlen at 1.11 1.12' without the equals
            # sign, so assume that a missing '=' is for this type of scan.
            zvarname, assignment = zvarstring.split(" ", 1)

        zvarname = zvarname.strip()

        scan = None

        # FROM ... TO ... BY
        match = re.match(r"\s*(from\s+)?(\S+)\s+to\s+(\S+)\s+by\s+(\S+)\s*",
                         assignment, re.IGNORECASE)
        if match:
            scan = Scan(zvarname,
                        initial=float(match.group(2)),
                        final=float(match.group(3)),
                        step_size=float(match.group(4)))

        # FROM ... TO ... IN
        if not scan:
            match = re.match(r"\s*(from\s+)?(\S+)\s+to\s+(\S+)\s+in\s+(\S+)\s*",
                             assignment, re.IGNORECASE)
            if match:
                scan = Scan(zvarname,
                            initial=float(match.group(2)),
                            final=float(match.group(3)),
                            steps=int(match.group(4)))

        # FROM ... BY ... IN
        if not scan:
            match = re.match(r"\s*(from\s+)?(\S+)\s+by\s+(\S+)\s+in\s+(\S+)\s*",
                             assignment, re.IGNORECASE)
            if match:
                scan = Scan(zvarname,
                            initial=float(match.group(2)),
                            step_size=float(match.group(3)),
                            steps=int(match.group(4)))

        # AT ...
        if not scan:
            match = re.match(r"\s*at\s+", assignment, re.IGNORECASE)
            if match:
                items = assignment.split()
                at_values = [float(item) for item in items[1:]]
                scan = Scan(zvarname, at_values=at_values)

        return scan