<commit_before># -*- coding: utf-8 -*-
"""
The following types are implemented with sanitize() and is_valid() functions:

 - Base64
 - ClassificationType
 - DNS
 - DateTime
 - FQDN
 - FeedName
 - GenericType
 - IPAddress
 - IPNetwork
 - MalwareName
 - String
 - URL
 - UUID
"""
import binascii
import datetime
import ipaddress
import json
import socket
from urllib.parse import urlparse

import dateutil.parser
import dns.resolver
import pytz

import intelmq.lib.utils as utils

__all__ = ['Base64', 'Boolean', 'ClassificationType', 'DateTime', 'FQDN',
           'Float', 'Accuracy', 'GenericType', 'IPAddress', 'IPNetwork',
           'Integer', 'JSON', 'MalwareName', 'String', 'URL', 'UUID',
           ]

class GenericType(object):

    @staticmethod
    def is_valid(value, sanitize=False):
        if sanitize:
            value = GenericType().sanitize(value)

        if not value:
            return False

        if not isinstance(value, str):
            return False

        if len(value) == 0:
            return False

        return True

    @staticmethod
    def sanitize(value):
        if not value:
            return None

        if isinstance(value, str):
            return value.strip()

        if isinstance(value, bytes):
            try:
                value = value.decode('utf-8')
            except UnicodeDecodeError:
                value = value.decode('utf-8', 'ignore')
            return value.strip()

        return None

class Base64(GenericType):
    """
    Base64 type. Always gives unicode strings.

    Sanitation encodes to base64 and accepts binary and unicode strings.
    """

    @staticmethod
    def is_valid(value, sanitize=False):
        if sanitize:
            value = GenericType().sanitize(value)
            value = Base64().sanitize(value)

        try:
            utils.base64_decode(value)
        except TypeError:
            return False

        if not GenericType().is_valid(value):
            return False

        return True

    @staticmethod
    def sanitize(value):
        value = utils.base64_encode(value)
        return value

class Boolean(GenericType):
    """
    Boolean type. Without sanitation only python bool is accepted.

    Sanitation accepts string 'true' and 'false' and integers 0 and 1.
    """

    @staticmethod
    def is_valid(value, sanitize=False):
        if isinstance(value, bool):
            return True
        else:
            if sanitize:
                value = Boolean().sanitize(value)
                if value is not None:
                    return True
            return False

    @staticmethod
    def sanitize(value):
        if isinstance(value, (str, bytes)):
            value = value.strip().lower()
            if value == 'true':
                return True
            elif value == 'false':
                return False
        elif isinstance(value, int):
            if value == 0:
                return False
            elif value == 1:
                return True
        return None

class ClassificationType(GenericType):

    allowed_values = ['spam',
                      'malware',
                      'botnet drone',
                      'ransomware',
                      'malware configuration',
                      'c&c',
                      'scanner',
                      'exploit',
                      'brute-force',
                      'ids alert',
                      'defacement',
                      'compromised',
                      'backdoor',
                      'ddos',
                      'dropzone',
                      'phishing',
                      'vulnerable service',
                      'blacklist',
                      'other',
                      'unknown'
                      ]

    @staticmethod
    def is_valid(value, sanitize=False):
        if sanitize:
            value = GenericType().sanitize(value)
            value = ClassificationType().sanitize(value)

        if not GenericType().is_valid(value):
            return False

        if not isinstance(value, str):
            return False

        if value not in ClassificationType().allowed_values:
            return False

        return True

class DateTime(GenericType):

    @staticmethod
    def is_valid(value, sanitize=False):
        if sanitize:
            value = GenericType().sanitize(value)
            value = DateTime().sanitize(value)

        if not GenericType().is_valid(value):
            return False

        if value != DateTime.__parse(value):
            return False

        return True

    @staticmethod
    def sanitize(value):
        value = DateTime.__parse(value)
        return GenericType().sanitize(value)

    @staticmethod
    def __parse(value):
        try:
            value = dateutil.parser.parse(value, fuzzy=True)
            value = value.astimezone(pytz.utc)
            value = value.isoformat()
        except ValueError:
            return None
        return utils.decode(value)

    @staticmethod
    def from_timestamp(tstamp, tzone='UTC'):
        """
        Returns ISO formated datetime from given timestamp.
        You can give timezone for given timestamp, UTC by default.
        """
        dtime = (datetime.datetime(1970, 1, 1, tzinfo=pytz.utc) +
                 datetime.timedelta(seconds=tstamp))
        localized = pytz.timezone(tzone).normalize(dtime)
        return str(localized.isoformat())

    @staticmethod
    def generate_datetime_now():
        value = datetime.datetime.now(pytz.timezone('UTC'))
        value = value.replace(microsecond=0)
        value = value.isoformat()
        # Is byte string in 2 and unicode string in 3, make unicode string
        return utils.decode(value)

class Float(GenericType):
    """
    Float type. Without sanitation only python float/integer/long is
    accepted. Boolean is excplicitly denied.

    Sanitation accepts strings and everything float() accepts.
    """

    @staticmethod
    def is_valid(value, sanitize=False):
        if sanitize:
            value = Float().sanitize(value)
            if value is not None:
                return True

        # Bool is subclass of int
        if isinstance(value, bool):
            return False
        if isinstance(value, (int, float)):
            return True

        return False

    @staticmethod
    def sanitize(value):
        try:
            return float(value)
        except (ValueError, TypeError):
            return None

class Accuracy(GenericType):
    """
    Accuracy type. A Float between 0 and 100.
    """

    @staticmethod
    def is_valid(value, sanitize=False):
        if sanitize:
            value = Accuracy.sanitize(value)
            if value is not None and value >= 0 and value <= 100:
                return True

        # Bool is subclass of int
        if isinstance(value, bool):
            return True
        if isinstance(value, (int, float)) and value >= 0 and value <= 100:
            return True

        return False

    @staticmethod
    def sanitize(value):
        try:
            if isinstance(value, bool):
                return float(value) * 100

            value = float(value)
            if value >= 0 or value <= 100:
                return value
        except (ValueError, TypeError):
            return None

class FQDN(GenericType):
    """
    Fully qualified domain name type.

    All valid domains are accepted, no IP addresses or URLs. Trailing dot is
    not allowed.
    """

    @staticmethod
    def is_valid(value, sanitize=False):
        if sanitize:
            value = GenericType().sanitize(value)
            value = FQDN().sanitize(value)

        if not GenericType().is_valid(value):
            return False

        if IPAddress().is_valid(value):
            return False

        if URL().is_valid(value):
            return False

        """
        # "localhost" is a valid hostname
        if len(value.split('.')) <= 1:
            return False
        """

        if value[-1] == '.':
            return False

        return True

    @staticmethod
    def sanitize(value):
        return value.rstrip('.')

    @staticmethod
    def to_ip(value):
        try:
            value = str(dns.resolver.query(value, 'A')[0])
        except dns.resolver.NXDOMAIN:  # domain not found
            value = None
        return value

class Integer(GenericType):
    """
    Integer type. Without sanitation only python integer/long is accepted.
    Bool is excplicitly denied.

    Sanitation accepts strings and everything int() accepts.
    """

    @staticmethod
    def is_valid(value, sanitize=False):
        if sanitize:
            value = Integer().sanitize(value)
            if value is not None:
                return True

        # Bool is subclass of int
        if isinstance(value, bool):
            return False
        if isinstance(value, int):
            return True

        return False

    @staticmethod
    def sanitize(value):
        try:
            return int(value)
        except (ValueError, TypeError):
            return None

class IPAddress(GenericType):

    @staticmethod
    def is_valid(value, sanitize=False):
        if sanitize:
            value = GenericType().sanitize(value)
            value = IPAddress().sanitize(value)

        if not GenericType().is_valid(value):
            return False

        try:
            address = ipaddress.ip_address(value)
        except ValueError:
            return False

        if address == ipaddress.ip_address('0.0.0.0'):
            return False

        return True

    @staticmethod
    def sanitize(value):

        try:
            network = ipaddress.ip_network(str(value))
        except ValueError:
            return None

        if network.num_addresses == 1:
            value = str(network.network_address)
        else:
            return None

        return GenericType().sanitize(value)

    @staticmethod
    def to_int(value):
        try:
            ip_integer = socket.inet_pton(socket.AF_INET, value)
        except socket.error:
            try:
                ip_integer = socket.inet_pton(socket.AF_INET6, value)
            except socket.error:
                return None

        ip_integer = int(binascii.hexlify(ip_integer), 16)
        return ip_integer

    @staticmethod
    def version(value):
        return ipaddress.ip_address(value).version

    @staticmethod
    def to_reverse(ip_addr):
        return str(dns.reversename.from_address(ip_addr))

class IPNetwork(GenericType):

    @staticmethod
    def is_valid(value, sanitize=False):
        if sanitize:
            value = GenericType().sanitize(value)
            value = IPNetwork().sanitize(value)

        if not GenericType().is_valid(value):
            return False

        try:
            ipaddress.ip_network(value)
        except ValueError:
            return False

        return True

    @staticmethod
    def sanitize(value):

        try:
            ipaddress.ip_network(str(value))
        except ValueError:
            return None

        return GenericType().sanitize(value)

    @staticmethod
    def version(value):
        return ipaddress.ip_network(str(value)).version

class JSON(GenericType):
    """
    JSON type.

    Sanitation accepts pythons dictionaries and JSON strings.

    Valid values are only unicode strings with JSON dictionaries.
    """

    @staticmethod
    def is_valid(value, sanitize=False):
        if sanitize:
            value = JSON().sanitize(value)

        if not isinstance(value, str):
            return False

        try:
            unpacked = json.loads(value)
        except ValueError:
            return False
        else:
            if isinstance(unpacked, dict) and unpacked != {}:
                return True

        return False

    @staticmethod
    def sanitize(value):
        if not value:
            return None
        if isinstance(value, (str, bytes)):
            sanitized = GenericType.sanitize(value)
            if JSON.is_valid(sanitized):
                return sanitized
        try:
            return GenericType().sanitize(json.dumps(value, sort_keys=True))
        except TypeError:
            return None

class MalwareName(GenericType):

    @staticmethod
    def is_valid(value, sanitize=False):
        if sanitize:
            value = GenericType().sanitize(value)
            value = MalwareName().sanitize(value)

        if not GenericType().is_valid(value):
            return False

        if value != value.lower():
            return False

        return True

    @staticmethod
    def sanitize(value):
        value = value.lower()
        return GenericType().sanitize(value)

class String(GenericType):

    @staticmethod
    def is_valid(value, sanitize=False):
        if sanitize:
            value = GenericType().sanitize(value)
            value = String().sanitize(value)

        if not GenericType().is_valid(value):
            return False

        if type(value) is not str:
            return False

        if len(value) == 0:
            return False

        return True

class URL(GenericType):

    @staticmethod
    def is_valid(value, sanitize=False):
        if sanitize:
            value = GenericType().sanitize(value)
            value = URL().sanitize(value)

        if not GenericType().is_valid(value):
            return False

        result = urlparse(value)
        if result.netloc == "":
            return False

        return True

    @staticmethod
    def sanitize(value):
        value = value.replace('hxxp://', 'http://')
        value = value.replace('hxxps://', 'https://')

        if urlparse(value).netloc != "":
            return GenericType().sanitize(value)

    @staticmethod
    def to_ip(url):
        value = urlparse(url)
        if value.netloc != "":
            return FQDN().to_ip(value.netloc)
        return None

    @staticmethod
    def to_domain_name(url):
        value = urlparse(url)
        if value.netloc != "" and not IPAddress.is_valid(value.netloc):
            return value.netloc
        return None

class UUID(String):
    pass
<commit_msg>Added piracy harmonization field<commit_after># -*- coding: utf-8 -*-
"""
The following types are implemented with sanitize() and is_valid() functions:

 - Base64
 - ClassificationType
 - DNS
 - DateTime
 - FQDN
 - FeedName
 - GenericType
 - IPAddress
 - IPNetwork
 - MalwareName
 - String
 - URL
 - UUID
"""
import binascii
import datetime
import ipaddress
import json
import socket
from urllib.parse import urlparse

import dateutil.parser
import dns.resolver
import pytz

import intelmq.lib.utils as utils

__all__ = ['Base64', 'Boolean', 'ClassificationType', 'DateTime', 'FQDN',
           'Float', 'Accuracy', 'GenericType', 'IPAddress', 'IPNetwork',
           'Integer', 'JSON', 'MalwareName', 'String', 'URL', 'UUID',
           ]

class GenericType(object):

    @staticmethod
    def is_valid(value, sanitize=False):
        if sanitize:
            value = GenericType().sanitize(value)

        if not value:
            return False

        if not isinstance(value, str):
            return False

        if len(value) == 0:
            return False

        return True

    @staticmethod
    def sanitize(value):
        if not value:
            return None

        if isinstance(value, str):
            return value.strip()

        if isinstance(value, bytes):
            try:
                value = value.decode('utf-8')
            except UnicodeDecodeError:
                value = value.decode('utf-8', 'ignore')
            return value.strip()

        return None

class Base64(GenericType):
    """
    Base64 type. Always gives unicode strings.

    Sanitation encodes to base64 and accepts binary and unicode strings.
    """

    @staticmethod
    def is_valid(value, sanitize=False):
        if sanitize:
            value = GenericType().sanitize(value)
            value = Base64().sanitize(value)

        try:
            utils.base64_decode(value)
        except TypeError:
            return False

        if not GenericType().is_valid(value):
            return False

        return True

    @staticmethod
    def sanitize(value):
        value = utils.base64_encode(value)
        return value

class Boolean(GenericType):
    """
    Boolean type. Without sanitation only python bool is accepted.

    Sanitation accepts string 'true' and 'false' and integers 0 and 1.
    """

    @staticmethod
    def is_valid(value, sanitize=False):
        if isinstance(value, bool):
            return True
        else:
            if sanitize:
                value = Boolean().sanitize(value)
                if value is not None:
                    return True
            return False

    @staticmethod
    def sanitize(value):
        if isinstance(value, (str, bytes)):
            value = value.strip().lower()
            if value == 'true':
                return True
            elif value == 'false':
                return False
        elif isinstance(value, int):
            if value == 0:
                return False
            elif value == 1:
                return True
        return None

class ClassificationType(GenericType):

    allowed_values = ['spam',
                      'malware',
                      'botnet drone',
                      'ransomware',
                      'malware configuration',
                      'c&c',
                      'scanner',
                      'piracy',
                      'exploit',
                      'brute-force',
                      'ids alert',
                      'defacement',
                      'compromised',
                      'backdoor',
                      'ddos',
                      'dropzone',
                      'phishing',
                      'vulnerable service',
                      'blacklist',
                      'other',
                      'unknown'
                      ]

    @staticmethod
    def is_valid(value, sanitize=False):
        if sanitize:
            value = GenericType().sanitize(value)
            value = ClassificationType().sanitize(value)

        if not GenericType().is_valid(value):
            return False

        if not isinstance(value, str):
            return False

        if value not in ClassificationType().allowed_values:
            return False

        return True

class DateTime(GenericType):

    @staticmethod
    def is_valid(value, sanitize=False):
        if sanitize:
            value = GenericType().sanitize(value)
            value = DateTime().sanitize(value)

        if not GenericType().is_valid(value):
            return False

        if value != DateTime.__parse(value):
            return False

        return True

    @staticmethod
    def sanitize(value):
        value = DateTime.__parse(value)
        return GenericType().sanitize(value)

    @staticmethod
    def __parse(value):
        try:
            value = dateutil.parser.parse(value, fuzzy=True)
            value = value.astimezone(pytz.utc)
            value = value.isoformat()
        except ValueError:
            return None
        return utils.decode(value)

    @staticmethod
    def from_timestamp(tstamp, tzone='UTC'):
        """
        Returns ISO formated datetime from given timestamp.
        You can give timezone for given timestamp, UTC by default.
        """
        dtime = (datetime.datetime(1970, 1, 1, tzinfo=pytz.utc) +
                 datetime.timedelta(seconds=tstamp))
        localized = pytz.timezone(tzone).normalize(dtime)
        return str(localized.isoformat())

    @staticmethod
    def generate_datetime_now():
        value = datetime.datetime.now(pytz.timezone('UTC'))
        value = value.replace(microsecond=0)
        value = value.isoformat()
        # Is byte string in 2 and unicode string in 3, make unicode string
        return utils.decode(value)

class Float(GenericType):
    """
    Float type. Without sanitation only python float/integer/long is
    accepted. Boolean is excplicitly denied.

    Sanitation accepts strings and everything float() accepts.
    """

    @staticmethod
    def is_valid(value, sanitize=False):
        if sanitize:
            value = Float().sanitize(value)
            if value is not None:
                return True

        # Bool is subclass of int
        if isinstance(value, bool):
            return False
        if isinstance(value, (int, float)):
            return True

        return False

    @staticmethod
    def sanitize(value):
        try:
            return float(value)
        except (ValueError, TypeError):
            return None

class Accuracy(GenericType):
    """
    Accuracy type. A Float between 0 and 100.
    """

    @staticmethod
    def is_valid(value, sanitize=False):
        if sanitize:
            value = Accuracy.sanitize(value)
            if value is not None and value >= 0 and value <= 100:
                return True

        # Bool is subclass of int
        if isinstance(value, bool):
            return True
        if isinstance(value, (int, float)) and value >= 0 and value <= 100:
            return True

        return False

    @staticmethod
    def sanitize(value):
        try:
            if isinstance(value, bool):
                return float(value) * 100

            value = float(value)
            if value >= 0 or value <= 100:
                return value
        except (ValueError, TypeError):
            return None

class FQDN(GenericType):
    """
    Fully qualified domain name type.

    All valid domains are accepted, no IP addresses or URLs. Trailing dot is
    not allowed.
    """

    @staticmethod
    def is_valid(value, sanitize=False):
        if sanitize:
            value = GenericType().sanitize(value)
            value = FQDN().sanitize(value)

        if not GenericType().is_valid(value):
            return False

        if IPAddress().is_valid(value):
            return False

        if URL().is_valid(value):
            return False

        """
        # "localhost" is a valid hostname
        if len(value.split('.')) <= 1:
            return False
        """

        if value[-1] == '.':
            return False

        return True

    @staticmethod
    def sanitize(value):
        return value.rstrip('.')

    @staticmethod
    def to_ip(value):
        try:
            value = str(dns.resolver.query(value, 'A')[0])
        except dns.resolver.NXDOMAIN:  # domain not found
            value = None
        return value

class Integer(GenericType):
    """
    Integer type. Without sanitation only python integer/long is accepted.
    Bool is excplicitly denied.

    Sanitation accepts strings and everything int() accepts.
    """

    @staticmethod
    def is_valid(value, sanitize=False):
        if sanitize:
            value = Integer().sanitize(value)
            if value is not None:
                return True

        # Bool is subclass of int
        if isinstance(value, bool):
            return False
        if isinstance(value, int):
            return True

        return False

    @staticmethod
    def sanitize(value):
        try:
            return int(value)
        except (ValueError, TypeError):
            return None

class IPAddress(GenericType):

    @staticmethod
    def is_valid(value, sanitize=False):
        if sanitize:
            value = GenericType().sanitize(value)
            value = IPAddress().sanitize(value)

        if not GenericType().is_valid(value):
            return False

        try:
            address = ipaddress.ip_address(value)
        except ValueError:
            return False

        if address == ipaddress.ip_address('0.0.0.0'):
            return False

        return True

    @staticmethod
    def sanitize(value):

        try:
            network = ipaddress.ip_network(str(value))
        except ValueError:
            return None

        if network.num_addresses == 1:
            value = str(network.network_address)
        else:
            return None

        return GenericType().sanitize(value)

    @staticmethod
    def to_int(value):
        try:
            ip_integer = socket.inet_pton(socket.AF_INET, value)
        except socket.error:
            try:
                ip_integer = socket.inet_pton(socket.AF_INET6, value)
            except socket.error:
                return None

        ip_integer = int(binascii.hexlify(ip_integer), 16)
        return ip_integer

    @staticmethod
    def version(value):
        return ipaddress.ip_address(value).version

    @staticmethod
    def to_reverse(ip_addr):
        return str(dns.reversename.from_address(ip_addr))

class IPNetwork(GenericType):

    @staticmethod
    def is_valid(value, sanitize=False):
        if sanitize:
            value = GenericType().sanitize(value)
            value = IPNetwork().sanitize(value)

        if not GenericType().is_valid(value):
            return False

        try:
            ipaddress.ip_network(value)
        except ValueError:
            return False

        return True

    @staticmethod
    def sanitize(value):

        try:
            ipaddress.ip_network(str(value))
        except ValueError:
            return None

        return GenericType().sanitize(value)

    @staticmethod
    def version(value):
        return ipaddress.ip_network(str(value)).version

class JSON(GenericType):
    """
    JSON type.

    Sanitation accepts pythons dictionaries and JSON strings.

    Valid values are only unicode strings with JSON dictionaries.
    """

    @staticmethod
    def is_valid(value, sanitize=False):
        if sanitize:
            value = JSON().sanitize(value)

        if not isinstance(value, str):
            return False

        try:
            unpacked = json.loads(value)
        except ValueError:
            return False
        else:
            if isinstance(unpacked, dict) and unpacked != {}:
                return True

        return False

    @staticmethod
    def sanitize(value):
        if not value:
            return None
        if isinstance(value, (str, bytes)):
            sanitized = GenericType.sanitize(value)
            if JSON.is_valid(sanitized):
                return sanitized
        try:
            return GenericType().sanitize(json.dumps(value, sort_keys=True))
        except TypeError:
            return None

class MalwareName(GenericType):

    @staticmethod
    def is_valid(value, sanitize=False):
        if sanitize:
            value = GenericType().sanitize(value)
            value = MalwareName().sanitize(value)

        if not GenericType().is_valid(value):
            return False

        if value != value.lower():
            return False

        return True

    @staticmethod
    def sanitize(value):
        value = value.lower()
        return GenericType().sanitize(value)

class String(GenericType):

    @staticmethod
    def is_valid(value, sanitize=False):
        if sanitize:
            value = GenericType().sanitize(value)
            value = String().sanitize(value)

        if not GenericType().is_valid(value):
            return False

        if type(value) is not str:
            return False

        if len(value) == 0:
            return False

        return True

class URL(GenericType):

    @staticmethod
    def is_valid(value, sanitize=False):
        if sanitize:
            value = GenericType().sanitize(value)
            value = URL().sanitize(value)

        if not GenericType().is_valid(value):
            return False

        result = urlparse(value)
        if result.netloc == "":
            return False

        return True

    @staticmethod
    def sanitize(value):
        value = value.replace('hxxp://', 'http://')
        value = value.replace('hxxps://', 'https://')

        if urlparse(value).netloc != "":
            return GenericType().sanitize(value)

    @staticmethod
    def to_ip(url):
        value = urlparse(url)
        if value.netloc != "":
            return FQDN().to_ip(value.netloc)
        return None

    @staticmethod
    def to_domain_name(url):
        value = urlparse(url)
        if value.netloc != "" and not IPAddress.is_valid(value.netloc):
            return value.netloc
        return None

class UUID(String):
    pass
<|endoftext|>


<commit_before>"""
Dishonest Casino Example
------------------------

We'll use the ubiquitous dishonest casino example to demonstrate
how to train a Hidden Markov Model (HMM) on somewhat realistic
test data (e.g. http://www.mcb111.org/w06/durbin_book.pdf
Chapter 3).

In this example, we suspect that a casino is trading out a fair
die (singular or dice) for a loaded die. We want to figure out
1) when the loaded die was used (i.e. the most likely path) 2) how
often the loaded die is used (i.e. the transition probabilities) and
3) the probabilities for each outcome of a roll for the loaded die
(i.e. the emission probabilities).
"""

# %%
# First, import necessary modules and functions.

import numpy as np
import matplotlib.pyplot as plt

from hmmlearn import hmm

# %%
# Now, let's act as the casino and exchange a fair die for a loaded one
# and generate a series of rolls that someone at the casino would
# observe.

# make our generative model with two components, a fair die and a
# loaded die
gen_model = hmm.MultinomialHMM(n_components=2, random_state=99)

# the first state is the fair die so let's start there so no one
# catches on rigth away
gen_model.startprob_ = np.array([1.0, 0.0])

# now let's say that we sneak the loaded die in:
# here, we have a 95% chance to continue using the fair die and a 5%
# chance to switch to the loaded die
# when we enter the loaded die state, we have a 90% chance of staying
# in that state and a 10% chance of leaving
gen_model.transmat_ = np.array([[0.95, 0.05],
                                [0.1, 0.9]])

# now let's set the emission means:
# the first state is a fair die with equal probabilities and the
# second is loaded by being biased toward rolling a six
gen_model.emissionprob_ = \
    np.array([[1 / 6, 1 / 6, 1 / 6, 1 / 6, 1 / 6, 1 / 6],
              [1 / 10, 1 / 10, 1 / 10, 1 / 10, 1 / 10, 1 / 2]])

# simulate the loaded dice rolls
rolls, gen_states = gen_model.sample(30000)

# plot states over time, let's just look at the first rolls for clarity
fig, ax = plt.subplots()
ax.plot(gen_states[:500])
ax.set_title('States over time')
ax.set_xlabel('Time (# of rolls)')
ax.set_ylabel('State')
fig.show()

# plot rolls for the fair and loaded states
fig, ax = plt.subplots()
ax.hist(rolls[gen_states == 0], label='fair', alpha=0.5,
        bins=np.arange(7) - 0.5, density=True)
ax.hist(rolls[gen_states == 1], label='loaded', alpha=0.5,
        bins=np.arange(7) - 0.5, density=True)
ax.set_title('Roll probabilities by state')
ax.set_xlabel('Count')
ax.set_ylabel('Roll')
ax.legend()
fig.show()

# %%
# Now, let's see if we can recover our hidden states, transmission matrix
# and emission probabilities.

# split our data into training and validation sets (50/50 split)
X_train = rolls[:rolls.shape[0] // 2]
X_validate = rolls[rolls.shape[0] // 2:]

# check optimal score
gen_score = gen_model.score(X_validate)

best_score = best_model = None
n_fits = 50
np.random.seed(13)
for idx in range(n_fits):
    model = hmm.MultinomialHMM(
        n_components=2, random_state=idx,
        init_params='se')  # don't init transition, set it below
    # we need to initialize with random transition matrix probabilities
    # because the default is an even likelihood transition
    # we know transitions are rare (otherwise the casino would get caught!)
    # so let's have an uniform random prior distribution on our transitions
    # between 0 and 0.25 to train our model faster (than uniform random 0 to 1)
    fair_to_loaded = np.random.random() / 4
    loaded_to_fair = np.random.random() / 4
    model.transmat_ = np.array([[1 - fair_to_loaded, fair_to_loaded],
                                [loaded_to_fair, 1 - loaded_to_fair]])
    model.fit(X_train)
    score = model.score(X_validate)
    print(f'Model #{idx}\tScore: {score}')
    idx += 1
    if best_score is None or score > best_score:
        best_model = model
        best_score = score

print(f'Generated score: {gen_score}\nBest score:      {best_score}')

# use the Viterbi algorithm to predict the most likely sequence of states
# given the model
states = best_model.predict(rolls)

# plot our recovered states compared to generated (aim 1)
fig, ax = plt.subplots()
ax.plot(gen_states[:500], label='generated')
ax.plot(states[:500] + 1.5, label='recovered')
ax.set_yticks([])
ax.set_title('States compared to generated')
ax.set_xlabel('Time (# rolls)')
ax.set_xlabel('State')
ax.legend()
fig.show()

# %%
# Let's check our learned transition probabilities and see if they match.

print(f'Transmission Matrix Generated:\n{gen_model.transmat_.round(3)}\n\n'
      f'Transmission Matrix Recovered:\n{best_model.transmat_.round(3)}\n\n')

# %%
# Finally, let's see if we can tell how the die is loaded.

print(f'Emission Matrix Generated:\n{gen_model.emissionprob_.round(3)}\n\n'
      f'Emission Matrix Recovered:\n{best_model.emissionprob_.round(3)}\n\n')

# %%
# In this case, we were able to get very good estimates of the transition and
# emission matrices, but decoding the states was imperfect. That's because
# the decoding algorithm is greedy and picks the most likely series of states
# which isn't necessarily what happens in real life. Even so, our model could
# tell us when to watch for the loaded die and we'd have a better chance at
# catching them red-handed.
<commit_msg>[ENH, MRG] Add Dirichlet random priors in example (#467)<commit_after>"""
Dishonest Casino Example
------------------------

We'll use the ubiquitous dishonest casino example to demonstrate
how to train a Hidden Markov Model (HMM) on somewhat realistic
test data (e.g. http://www.mcb111.org/w06/durbin_book.pdf
Chapter 3).

In this example, we suspect that a casino is trading out a fair
die (singular or dice) for a loaded die. We want to figure out
1) when the loaded die was used (i.e. the most likely path) 2) how
often the loaded die is used (i.e. the transition probabilities) and
3) the probabilities for each outcome of a roll for the loaded die
(i.e. the emission probabilities).
"""

# %%
# First, import necessary modules and functions.

import numpy as np
import matplotlib.pyplot as plt

from hmmlearn import hmm

# %%
# Now, let's act as the casino and exchange a fair die for a loaded one
# and generate a series of rolls that someone at the casino would
# observe.

# make our generative model with two components, a fair die and a
# loaded die
gen_model = hmm.MultinomialHMM(n_components=2, random_state=99)

# the first state is the fair die so let's start there so no one
# catches on right away
gen_model.startprob_ = np.array([1.0, 0.0])

# now let's say that we sneak the loaded die in:
# here, we have a 95% chance to continue using the fair die and a 5%
# chance to switch to the loaded die
# when we enter the loaded die state, we have a 90% chance of staying
# in that state and a 10% chance of leaving
gen_model.transmat_ = np.array([[0.95, 0.05],
                                [0.1, 0.9]])

# now let's set the emission means:
# the first state is a fair die with equal probabilities and the
# second is loaded by being biased toward rolling a six
gen_model.emissionprob_ = \
    np.array([[1 / 6, 1 / 6, 1 / 6, 1 / 6, 1 / 6, 1 / 6],
              [1 / 10, 1 / 10, 1 / 10, 1 / 10, 1 / 10, 1 / 2]])

# simulate the loaded dice rolls
rolls, gen_states = gen_model.sample(30000)

# plot states over time, let's just look at the first rolls for clarity
fig, ax = plt.subplots()
ax.plot(gen_states[:500])
ax.set_title('States over time')
ax.set_xlabel('Time (# of rolls)')
ax.set_ylabel('State')
fig.show()

# plot rolls for the fair and loaded states
fig, ax = plt.subplots()
ax.hist(rolls[gen_states == 0], label='fair', alpha=0.5,
        bins=np.arange(7) - 0.5, density=True)
ax.hist(rolls[gen_states == 1], label='loaded', alpha=0.5,
        bins=np.arange(7) - 0.5, density=True)
ax.set_title('Roll probabilities by state')
ax.set_xlabel('Count')
ax.set_ylabel('Roll')
ax.legend()
fig.show()

# %%
# Now, let's see if we can recover our hidden states, transmission matrix
# and emission probabilities.

# split our data into training and validation sets (50/50 split)
X_train = rolls[:rolls.shape[0] // 2]
X_validate = rolls[rolls.shape[0] // 2:]

# check optimal score
gen_score = gen_model.score(X_validate)

best_score = best_model = None
n_fits = 50
np.random.seed(13)
for idx in range(n_fits):
    model = hmm.MultinomialHMM(
        n_components=2, random_state=idx,
        init_params='se')  # don't init transition, set it below
    # we need to initialize with random transition matrix probabilities
    # because the default is an even likelihood transition
    # we know transitions are rare (otherwise the casino would get caught!)
    # so let's have an Dirichlet random prior with an alpha value of
    # (0.1, 0.9) to enforce our assumption transitions happen roughly 10%
    # of the time
    model.transmat_ = np.array([np.random.dirichlet([0.9, 0.1]),
                                np.random.dirichlet([0.1, 0.9])])
    model.fit(X_train)
    score = model.score(X_validate)
    print(f'Model #{idx}\tScore: {score}')
    if best_score is None or score > best_score:
        best_model = model
        best_score = score

print(f'Generated score: {gen_score}\nBest score:      {best_score}')

# use the Viterbi algorithm to predict the most likely sequence of states
# given the model
states = best_model.predict(rolls)

# plot our recovered states compared to generated (aim 1)
fig, ax = plt.subplots()
ax.plot(gen_states[:500], label='generated')
ax.plot(states[:500] + 1.5, label='recovered')
ax.set_yticks([])
ax.set_title('States compared to generated')
ax.set_xlabel('Time (# rolls)')
ax.set_xlabel('State')
ax.legend()
fig.show()

# %%
# Let's check our learned transition probabilities and see if they match.

print(f'Transmission Matrix Generated:\n{gen_model.transmat_.round(3)}\n\n'
      f'Transmission Matrix Recovered:\n{best_model.transmat_.round(3)}\n\n')

# %%
# Finally, let's see if we can tell how the die is loaded.

print(f'Emission Matrix Generated:\n{gen_model.emissionprob_.round(3)}\n\n'
      f'Emission Matrix Recovered:\n{best_model.emissionprob_.round(3)}\n\n')

# %%
# In this case, we were able to get very good estimates of the transition and
# emission matrices, but decoding the states was imperfect. That's because
# the decoding algorithm is greedy and picks the most likely series of states
# which isn't necessarily what happens in real life. Even so, our model could
# tell us when to watch for the loaded die and we'd have a better chance at
# catching them red-handed.
<|endoftext|>


<commit_before>
# Copyright (c) 2015, 2014 Computational Molecular Biology Group, Free University
# Berlin, 14195 Berlin, Germany.
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without modification,
# are permitted provided that the following conditions are met:
#
#  * Redistributions of source code must retain the above copyright notice, this
# list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation and/or
# other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS ``AS IS''
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR
# ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
# ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

r"""This module contains function for the Transition Path Theory (TPT)
analysis of Markov models.

__moduleauthor__ = "Benjamin Trendelkamp-Schroer, Frank Noe"

"""
import numpy as np
import api as tptapi

__all__ = ['ReactiveFlux']

class ReactiveFlux(object):
    r"""A->B reactive flux from transition path theory (TPT)

    This object describes a reactive flux, i.e. a network of fluxes from a set of source states A, to a set of
    sink states B, via a set of intermediate nodes. Every node has three properties: the stationary probability mu,
    the forward committor qplus and the backward committor qminus. Every pair of edges has the following properties:
    a flux, generally a net flux that has no unnecessary back-fluxes, and optionally a gross flux.

    Flux objects can be used to compute transition pathways (and their weights) from A to B, the total flux, the
    total transition rate or mean first passage time, and they can be coarse-grained onto a set discretization
    of the node set.

    Fluxes can be computed in EMMA using transition path theory - see :func:`pyemma.msm.tpt`

    Parameters
    ----------
    A : array_like
        List of integer state labels for set A
    B : array_like
        List of integer state labels for set B
    flux : (n,n) ndarray or scipy sparse matrix
        effective or net flux of A->B pathways
    mu : (n,) ndarray (optional)
        Stationary vector
    qminus : (n,) ndarray (optional)
        Backward committor for A->B reaction
    qplus : (n,) ndarray (optional)
        Forward committor for A-> B reaction
    gross_flux : (n,n) ndarray or scipy sparse matrix
        gross flux of A->B pathways, if available

    Notes
    -----
    Reactive flux contains a flux network from educt states (A) to product states (B).

    See also
    --------
    pyemma.msm.tpt

    """

    def __init__(self, A, B, flux,
                 mu=None, qminus=None, qplus=None, gross_flux=None):
        # set data
        self._A = A
        self._B = B
        self._flux = flux
        self._mu = mu
        self._qminus = qminus
        self._qplus = qplus
        self._gross_flux = gross_flux
        # compute derived quantities:
        self._totalflux = tptapi.total_flux(flux, A)
        self._kAB = tptapi.rate(self._totalflux, mu, qminus)

    @property
    def nstates(self):
        r"""Returns the number of states.

        """
        return np.shape(self._flux)[0]

    @property
    def A(self):
        r"""Returns the set of reactant (source) states.

        """
        return self._A

    @property
    def B(self):
        r"""Returns the set of product (target) states

        """
        return self._B

    @property
    def I(self):
        r"""Returns the set of intermediate states

        """
        return list(set(range(self.nstates)) - set(self._A) - set(self._B))

    @property
    def stationary_distribution(self):
        r"""Returns the stationary distribution
        """
        return self._mu

    @property
    def flux(self):
        r"""Returns the effective or net flux
        """
        return self._flux

    @property
    def net_flux(self):
        r"""Returns the effective or net flux
        """
        return self._flux

    @property
    def gross_flux(self):
        r"""Returns the gross A-->B flux
        """
        return self._gross_flux

    @property
    def committor(self):
        r"""Returns the forward committor probability
        """
        return self._qplus

    @property
    def forward_committor(self):
        r"""Returns the forward committor probability
        """
        return self._qplus

    @property
    def backward_committor(self):
        r"""Returns the backward committor probability
        """
        return self._qminus

    @property
    def total_flux(self):
        r"""Returns the total flux
        """
        return self._totalflux

    @property
    def rate(self):
        r"""Returns the rate (inverse mfpt) of A-->B transitions
        """
        return self._kAB

    @property
    def mfpt(self):
        r"""Returns the rate (inverse mfpt) of A-->B transitions
        """
        return 1.0 / self._kAB

    def pathways(self, fraction=1.0, maxiter=1000):
        r"""Decompose flux network into dominant reaction paths.

        Parameters
        ----------
        fraction : float, optional
            Fraction of total flux to assemble in pathway decomposition
        maxiter : int, optional
            Maximum number of pathways for decomposition

        Returns
        -------
        paths : list
            List of dominant reaction pathways
        capacities: list
            List of capacities corresponding to each reactions pathway in paths

        References
        ----------
        .. [1] P. Metzner, C. Schuette and E. Vanden-Eijnden.
            Transition Path Theory for Markov Jump Processes.
            Multiscale Model Simul 7: 1192-1219 (2009)

        """
        return tptapi.pathways(self.net_flux, self.A, self.B,
                               fraction=fraction, maxiter=maxiter)

    def _pathways_to_flux(self, paths, pathfluxes, n=None):
        r"""Sums up the flux from the pathways given

        Parameters
        -----------
        paths : list of int-arrays
        list of pathways

        pathfluxes : double-array
            array with path fluxes

        n : int
            number of states. If not set, will be automatically determined.

        Returns:
        --------
        flux : (n,n) ndarray of float
            the flux containing the summed path fluxes

        """
        if (n is None):
            n = 0
            for p in paths:
                n = max(n, np.max(p))
            n += 1

        # initialize flux
        F = np.zeros((n, n))
        for i in range(len(paths)):
            p = paths[i]
            for t in range(len(p) - 1):
                F[p[t], p[t + 1]] += pathfluxes[i]
        return F

    def major_flux(self, fraction=0.9):
        r"""Returns the main pathway part of the net flux comprising
        at most the requested fraction of the full flux.

        """
        (paths, pathfluxes) = self.pathways(fraction=fraction)
        return self._pathways_to_flux(paths, pathfluxes, n=self.nstates)

    # this will be a private function in tpt. only Parameter left will be the sets to be distinguished
    def _compute_coarse_sets(self, user_sets):
        r"""Computes the sets to coarse-grain the tpt flux to.

        Parameters
        ----------
            (tpt_sets, A, B) with
                tpt_sets : list of int-iterables
                sets of states that shall be distinguished in the coarse-grained flux.
            A : int-iterable
                set indexes in A
            B : int-iterable
                set indexes in B

        Returns
        -------
        sets : list of int-iterables
            sets to compute tpt on. These sets still respect the boundary between
            A, B and the intermediate tpt states.

        Notes
        -----
        Given the sets that the user wants to distinguish, the
        algorithm will create additional sets if necessary

           * If states are missing in user_sets, they will be put into a
            separate set
           * If sets in user_sets are crossing the boundary between A, B and the
             intermediates, they will be split at these boundaries. Thus each
             set in user_sets can remain intact or be split into two or three
             subsets

        """
        # set-ify everything
        setA = set(self.A)
        setB = set(self.B)
        setI = set(self.I)
        raw_sets = [set(user_set) for user_set in user_sets]

        # anything missing? Compute all listed states
        set_all = set(range(self.nstates))
        set_all_user = []
        for user_set in raw_sets:
            set_all_user += user_set
        set_all_user = set(set_all_user)
        # ... and add all the unlisted states in a separate set
        set_rest = set_all - set_all_user
        if len(set_rest) > 0:
            raw_sets.append(set_rest)

        # split sets
        Asets = []
        Isets = []
        Bsets = []
        for raw_set in raw_sets:
            s = raw_set.intersection(setA)
            if len(s) > 0:
                Asets.append(s)
            s = raw_set.intersection(setI)
            if len(s) > 0:
                Isets.append(s)
            s = raw_set.intersection(setB)
            if len(s) > 0:
                Bsets.append(s)
        tpt_sets = Asets + Isets + Bsets
        Aindexes = range(0, len(Asets))
        Bindexes = range(len(Asets) + len(Isets), len(tpt_sets))

        return (tpt_sets, Aindexes, Bindexes)

    def coarse_grain(self, user_sets):
        r"""Coarse-grains the flux onto user-defined sets.

        Parameters
        ----------
        user_sets : list of int-iterables
            sets of states that shall be distinguished in the coarse-grained flux.

        Returns
        -------
        (sets, tpt) : (list of int-iterables, tpt-object)
            sets contains the sets tpt is computed on. The tpt states of the new
            tpt object correspond to these sets of states in this order. Sets might
            be identical, if the user has already provided a complete partition that
            respects the boundary between A, B and the intermediates. If not, Sets
            will have more members than provided by the user, containing the
            "remainder" states and reflecting the splitting at the A and B
            boundaries.
            tpt contains a new tpt object for the coarse-grained flux. All its
            quantities (gross_flux, net_flux, A, B, committor, backward_committor)
            are coarse-grained to sets.

        Notes
        -----
        All user-specified sets will be split (if necessary) to
        preserve the boundary between A, B and the intermediate
        states.

        """
        # coarse-grain sets
        (tpt_sets, Aindexes, Bindexes) = self._compute_coarse_sets(user_sets)
        nnew = len(tpt_sets)

        # coarse-grain fluxHere we should branch between sparse and dense implementations, but currently there is only a
        F_coarse = tptapi.coarsegrain(self._gross_flux, tpt_sets)
        Fnet_coarse = tptapi.to_netflux(F_coarse)

        # coarse-grain stationary probability and committors - this can be done all dense
        pstat_coarse = np.zeros((nnew))
        forward_committor_coarse = np.zeros((nnew))
        backward_committor_coarse = np.zeros((nnew))
        for i in range(0, nnew):
            I = list(tpt_sets[i])
            muI = self._mu[I]
            pstat_coarse[i] = np.sum(muI)
            partialI = muI / pstat_coarse[i]  # normalized stationary probability over I
            forward_committor_coarse[i] = np.dot(partialI, self._qplus[I])
            backward_committor_coarse[i] = np.dot(partialI, self._qminus[I])

        res = ReactiveFlux(Aindexes, Bindexes, Fnet_coarse, mu=pstat_coarse,
                           qminus=backward_committor_coarse, qplus=forward_committor_coarse, gross_flux=F_coarse)
        return (tpt_sets, res)<commit_msg>[doc] fix wrong sections ("Returns:" -> "Returns")<commit_after>
# Copyright (c) 2015, 2014 Computational Molecular Biology Group, Free University
# Berlin, 14195 Berlin, Germany.
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without modification,
# are permitted provided that the following conditions are met:
#
#  * Redistributions of source code must retain the above copyright notice, this
# list of conditions and the following disclaimer.
#  * Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation and/or
# other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS ``AS IS''
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR
# ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
# ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

r"""This module contains function for the Transition Path Theory (TPT)
analysis of Markov models.

__moduleauthor__ = "Benjamin Trendelkamp-Schroer, Frank Noe"

"""
import numpy as np
import api as tptapi

__all__ = ['ReactiveFlux']

class ReactiveFlux(object):
    r"""A->B reactive flux from transition path theory (TPT)

    This object describes a reactive flux, i.e. a network of fluxes from a set of source states A, to a set of
    sink states B, via a set of intermediate nodes. Every node has three properties: the stationary probability mu,
    the forward committor qplus and the backward committor qminus. Every pair of edges has the following properties:
    a flux, generally a net flux that has no unnecessary back-fluxes, and optionally a gross flux.

    Flux objects can be used to compute transition pathways (and their weights) from A to B, the total flux, the
    total transition rate or mean first passage time, and they can be coarse-grained onto a set discretization
    of the node set.

    Fluxes can be computed in EMMA using transition path theory - see :func:`pyemma.msm.tpt`

    Parameters
    ----------
    A : array_like
        List of integer state labels for set A
    B : array_like
        List of integer state labels for set B
    flux : (n,n) ndarray or scipy sparse matrix
        effective or net flux of A->B pathways
    mu : (n,) ndarray (optional)
        Stationary vector
    qminus : (n,) ndarray (optional)
        Backward committor for A->B reaction
    qplus : (n,) ndarray (optional)
        Forward committor for A-> B reaction
    gross_flux : (n,n) ndarray or scipy sparse matrix
        gross flux of A->B pathways, if available

    Notes
    -----
    Reactive flux contains a flux network from educt states (A) to product states (B).

    See also
    --------
    pyemma.msm.tpt

    """

    def __init__(self, A, B, flux,
                 mu=None, qminus=None, qplus=None, gross_flux=None):
        # set data
        self._A = A
        self._B = B
        self._flux = flux
        self._mu = mu
        self._qminus = qminus
        self._qplus = qplus
        self._gross_flux = gross_flux
        # compute derived quantities:
        self._totalflux = tptapi.total_flux(flux, A)
        self._kAB = tptapi.rate(self._totalflux, mu, qminus)

    @property
    def nstates(self):
        r"""Returns the number of states.

        """
        return np.shape(self._flux)[0]

    @property
    def A(self):
        r"""Returns the set of reactant (source) states.

        """
        return self._A

    @property
    def B(self):
        r"""Returns the set of product (target) states

        """
        return self._B

    @property
    def I(self):
        r"""Returns the set of intermediate states

        """
        return list(set(range(self.nstates)) - set(self._A) - set(self._B))

    @property
    def stationary_distribution(self):
        r"""Returns the stationary distribution
        """
        return self._mu

    @property
    def flux(self):
        r"""Returns the effective or net flux
        """
        return self._flux

    @property
    def net_flux(self):
        r"""Returns the effective or net flux
        """
        return self._flux

    @property
    def gross_flux(self):
        r"""Returns the gross A-->B flux
        """
        return self._gross_flux

    @property
    def committor(self):
        r"""Returns the forward committor probability
        """
        return self._qplus

    @property
    def forward_committor(self):
        r"""Returns the forward committor probability
        """
        return self._qplus

    @property
    def backward_committor(self):
        r"""Returns the backward committor probability
        """
        return self._qminus

    @property
    def total_flux(self):
        r"""Returns the total flux
        """
        return self._totalflux

    @property
    def rate(self):
        r"""Returns the rate (inverse mfpt) of A-->B transitions
        """
        return self._kAB

    @property
    def mfpt(self):
        r"""Returns the rate (inverse mfpt) of A-->B transitions
        """
        return 1.0 / self._kAB

    def pathways(self, fraction=1.0, maxiter=1000):
        r"""Decompose flux network into dominant reaction paths.

        Parameters
        ----------
        fraction : float, optional
            Fraction of total flux to assemble in pathway decomposition
        maxiter : int, optional
            Maximum number of pathways for decomposition

        Returns
        -------
        paths : list
            List of dominant reaction pathways
        capacities: list
            List of capacities corresponding to each reactions pathway in paths

        References
        ----------
        .. [1] P. Metzner, C. Schuette and E. Vanden-Eijnden.
            Transition Path Theory for Markov Jump Processes.
            Multiscale Model Simul 7: 1192-1219 (2009)

        """
        return tptapi.pathways(self.net_flux, self.A, self.B,
                               fraction=fraction, maxiter=maxiter)

    def _pathways_to_flux(self, paths, pathfluxes, n=None):
        r"""Sums up the flux from the pathways given

        Parameters
        -----------
        paths : list of int-arrays
        list of pathways

        pathfluxes : double-array
            array with path fluxes

        n : int
            number of states. If not set, will be automatically determined.

        Returns
        -------
        flux : (n,n) ndarray of float
            the flux containing the summed path fluxes

        """
        if (n is None):
            n = 0
            for p in paths:
                n = max(n, np.max(p))
            n += 1

        # initialize flux
        F = np.zeros((n, n))
        for i in range(len(paths)):
            p = paths[i]
            for t in range(len(p) - 1):
                F[p[t], p[t + 1]] += pathfluxes[i]
        return F

    def major_flux(self, fraction=0.9):
        r"""Returns the main pathway part of the net flux comprising
        at most the requested fraction of the full flux.

        """
        (paths, pathfluxes) = self.pathways(fraction=fraction)
        return self._pathways_to_flux(paths, pathfluxes, n=self.nstates)

    # this will be a private function in tpt. only Parameter left will be the sets to be distinguished
    def _compute_coarse_sets(self, user_sets):
        r"""Computes the sets to coarse-grain the tpt flux to.

        Parameters
        ----------
            (tpt_sets, A, B) with
                tpt_sets : list of int-iterables
                sets of states that shall be distinguished in the coarse-grained flux.
            A : int-iterable
                set indexes in A
            B : int-iterable
                set indexes in B

        Returns
        -------
        sets : list of int-iterables
            sets to compute tpt on. These sets still respect the boundary between
            A, B and the intermediate tpt states.

        Notes
        -----
        Given the sets that the user wants to distinguish, the
        algorithm will create additional sets if necessary

           * If states are missing in user_sets, they will be put into a
            separate set
           * If sets in user_sets are crossing the boundary between A, B and the
             intermediates, they will be split at these boundaries. Thus each
             set in user_sets can remain intact or be split into two or three
             subsets

        """
        # set-ify everything
        setA = set(self.A)
        setB = set(self.B)
        setI = set(self.I)
        raw_sets = [set(user_set) for user_set in user_sets]

        # anything missing? Compute all listed states
        set_all = set(range(self.nstates))
        set_all_user = []
        for user_set in raw_sets:
            set_all_user += user_set
        set_all_user = set(set_all_user)
        # ... and add all the unlisted states in a separate set
        set_rest = set_all - set_all_user
        if len(set_rest) > 0:
            raw_sets.append(set_rest)

        # split sets
        Asets = []
        Isets = []
        Bsets = []
        for raw_set in raw_sets:
            s = raw_set.intersection(setA)
            if len(s) > 0:
                Asets.append(s)
            s = raw_set.intersection(setI)
            if len(s) > 0:
                Isets.append(s)
            s = raw_set.intersection(setB)
            if len(s) > 0:
                Bsets.append(s)
        tpt_sets = Asets + Isets + Bsets
        Aindexes = range(0, len(Asets))
        Bindexes = range(len(Asets) + len(Isets), len(tpt_sets))

        return (tpt_sets, Aindexes, Bindexes)

    def coarse_grain(self, user_sets):
        r"""Coarse-grains the flux onto user-defined sets.

        Parameters
        ----------
        user_sets : list of int-iterables
            sets of states that shall be distinguished in the coarse-grained flux.

        Returns
        -------
        (sets, tpt) : (list of int-iterables, tpt-object)
            sets contains the sets tpt is computed on. The tpt states of the new
            tpt object correspond to these sets of states in this order. Sets might
            be identical, if the user has already provided a complete partition that
            respects the boundary between A, B and the intermediates. If not, Sets
            will have more members than provided by the user, containing the
            "remainder" states and reflecting the splitting at the A and B
            boundaries.
            tpt contains a new tpt object for the coarse-grained flux. All its
            quantities (gross_flux, net_flux, A, B, committor, backward_committor)
            are coarse-grained to sets.

        Notes
        -----
        All user-specified sets will be split (if necessary) to
        preserve the boundary between A, B and the intermediate
        states.

        """
        # coarse-grain sets
        (tpt_sets, Aindexes, Bindexes) = self._compute_coarse_sets(user_sets)
        nnew = len(tpt_sets)

        # coarse-grain fluxHere we should branch between sparse and dense implementations, but currently there is only a
        F_coarse = tptapi.coarsegrain(self._gross_flux, tpt_sets)
        Fnet_coarse = tptapi.to_netflux(F_coarse)

        # coarse-grain stationary probability and committors - this can be done all dense
        pstat_coarse = np.zeros((nnew))
        forward_committor_coarse = np.zeros((nnew))
        backward_committor_coarse = np.zeros((nnew))
        for i in range(0, nnew):
            I = list(tpt_sets[i])
            muI = self._mu[I]
            pstat_coarse[i] = np.sum(muI)
            partialI = muI / pstat_coarse[i]  # normalized stationary probability over I
            forward_committor_coarse[i] = np.dot(partialI, self._qplus[I])
            backward_committor_coarse[i] = np.dot(partialI, self._qminus[I])

        res = ReactiveFlux(Aindexes, Bindexes, Fnet_coarse, mu=pstat_coarse,
                           qminus=backward_committor_coarse, qplus=forward_committor_coarse, gross_flux=F_coarse)
        return (tpt_sets, res)<|endoftext|>


<commit_before>import time
import sys
from IPython.html.widgets import FloatProgress, IntProgress
from IPython.display import display
from numpy import linspace

def progress_display(results, objectlist, sleep=10):
    while not results.ready():
        print("{:.1f} % done.".format(100*results.progress/len(objectlist)))
        sys.stdout.flush()
        time.sleep(sleep)

def int_progress(min, max):
    prog = IntProgress(min=min, max=max)
    display(prog)
    for i in linspace(min, max, 25):
        time.sleep(0.1)
        prog.value = i

def float_progress(min, max):
    prog = FloatProgress(min=min, max=max)
    display(prog)
    for i in linspace(min, max, 100):
        time.sleep(0.1)
        prog.value = i
<commit_msg>adding a notebook_progress_display using IPython widgets<commit_after>import time
import sys
from IPython.html.widgets import FloatProgress, IntProgress
from IPython.display import display
from numpy import linspace

def progress_display(results, objectlist, sleep=10):
    while not results.ready():
        print("{:.1f} % done.".format(100*results.progress/len(objectlist)))
        sys.stdout.flush()
        time.sleep(sleep)

def nb_progress_display(results, objectlist, sleep=1):
    prog = IntProgress(min=0, max=len(objectlist))
    display(prog)
    while not results.ready():
        prog.value = results.progress
        time.sleep(sleep)

def int_progress(min, max):
    prog = IntProgress(min=min, max=max)
    display(prog)
    for i in linspace(min, max, 25):
        time.sleep(0.1)
        prog.value = i

def float_progress(min, max):
    prog = FloatProgress(min=min, max=max)
    display(prog)
    for i in linspace(min, max, 100):
        time.sleep(0.1)
        prog.value = i
<|endoftext|>


<commit_before>"""
Django settings for japper project.

For more information on this file, see
https://docs.djangoproject.com/en/1.7/topics/settings/

For the full list of settings and their values, see
https://docs.djangoproject.com/en/1.7/ref/settings/
"""

import datetime

import environ

from japper.monitoring.plugins import get_installed_apps

# Load environment
env = environ.Env(
    DEBUG=(bool, False),
    STATIC_ROOT=(str, None),
    ALLOWED_HOSTS=(list, []),
)

# Quick-start development settings - unsuitable for production
# See https://docs.djangoproject.com/en/1.7/howto/deployment/checklist/

# SECURITY WARNING: keep the secret key used in production secret!
SECRET_KEY = 'rk)5ko4=78lci42-c1$-*#tgdph!5qzq^qqrfn-gpjvtcs&vd_'

# SECURITY WARNING: don't run with debug turned on in production!
DEBUG = env('DEBUG')

TEMPLATE_DEBUG = True

ALLOWED_HOSTS = env('ALLOWED_HOSTS')

# Application definition
INSTALLED_APPS = get_installed_apps((
    'django.contrib.admin',
    'django.contrib.auth',
    'django.contrib.contenttypes',
    'django.contrib.sessions',
    'django.contrib.messages',
    'django.contrib.staticfiles',

    'crispy_forms',
    'menu',

    'japper',
    'japper.monitoring',
    'japper.users',
))

MIDDLEWARE_CLASSES = (
    'django.contrib.sessions.middleware.SessionMiddleware',
    'django.middleware.common.CommonMiddleware',
    'django.middleware.csrf.CsrfViewMiddleware',
    'django.contrib.auth.middleware.AuthenticationMiddleware',
    'django.contrib.auth.middleware.SessionAuthenticationMiddleware',
    'django.contrib.messages.middleware.MessageMiddleware',
    'django.middleware.clickjacking.XFrameOptionsMiddleware',
)

ROOT_URLCONF = 'japper.urls'

WSGI_APPLICATION = 'japper.wsgi.application'

TEMPLATE_CONTEXT_PROCESSORS = (
    'django.contrib.auth.context_processors.auth',
    'django.core.context_processors.debug',
    'django.core.context_processors.i18n',
    'django.core.context_processors.media',
    'django.core.context_processors.static',
    'django.core.context_processors.tz',
    'django.contrib.messages.context_processors.messages',
    'django.core.context_processors.request',
)

# Database
# https://docs.djangoproject.com/en/1.7/ref/settings/#databases

DATABASES = {
    'default': env.db(default='postgres://localhost/japper'),
}

# Internationalization
# https://docs.djangoproject.com/en/1.7/topics/i18n/

LANGUAGE_CODE = 'en-us'

TIME_ZONE = 'UTC'

USE_I18N = True

USE_L10N = True

USE_TZ = True

# Static files (CSS, JavaScript, Images)
# https://docs.djangoproject.com/en/1.7/howto/static-files/

STATIC_URL = '/static/'
STATIC_ROOT = env('STATIC_ROOT')

# Crispy forms settings

CRISPY_TEMPLATE_PACK = 'bootstrap3'

# Celery settings

CELERYBEAT_SCHEDULE = {
    'monitoring:fetch_check_results': {
        'task': 'japper.monitoring.tasks.fetch_check_results',
        'schedule': datetime.timedelta(minutes=1),
    },
    'monitoring:cleanup': {
        'task': 'japper.monitoring.tasks.cleanup',
        'schedule': datetime.timedelta(minutes=5),
    },
}

# Email settings

globals().update(env.email_url(default='smtp://localhost'))

# Auth settings

LOGIN_REDIRECT_URL = 'monitoring_problems'
LOGIN_URL = 'login'
LOGOUT_URL = 'logout'

# Proxy settings

SECURE_PROXY_SSL_HEADER = ('X-Forwarded-Proto', 'https')
<commit_msg>BROKER_URL can now be configured by an environment variable<commit_after>"""
Django settings for japper project.

For more information on this file, see
https://docs.djangoproject.com/en/1.7/topics/settings/

For the full list of settings and their values, see
https://docs.djangoproject.com/en/1.7/ref/settings/
"""

import datetime

import environ

from japper.monitoring.plugins import get_installed_apps

# Load environment
env = environ.Env(
    DEBUG=(bool, False),
    STATIC_ROOT=(str, None),
    ALLOWED_HOSTS=(list, []),
    BROKER_URL=(str, 'amqp://guest:guest@localhost:5672//')
)

# Quick-start development settings - unsuitable for production
# See https://docs.djangoproject.com/en/1.7/howto/deployment/checklist/

# SECURITY WARNING: keep the secret key used in production secret!
SECRET_KEY = 'rk)5ko4=78lci42-c1$-*#tgdph!5qzq^qqrfn-gpjvtcs&vd_'

# SECURITY WARNING: don't run with debug turned on in production!
DEBUG = env('DEBUG')

TEMPLATE_DEBUG = True

ALLOWED_HOSTS = env('ALLOWED_HOSTS')

# Application definition
INSTALLED_APPS = get_installed_apps((
    'django.contrib.admin',
    'django.contrib.auth',
    'django.contrib.contenttypes',
    'django.contrib.sessions',
    'django.contrib.messages',
    'django.contrib.staticfiles',

    'crispy_forms',
    'menu',

    'japper',
    'japper.monitoring',
    'japper.users',
))

MIDDLEWARE_CLASSES = (
    'django.contrib.sessions.middleware.SessionMiddleware',
    'django.middleware.common.CommonMiddleware',
    'django.middleware.csrf.CsrfViewMiddleware',
    'django.contrib.auth.middleware.AuthenticationMiddleware',
    'django.contrib.auth.middleware.SessionAuthenticationMiddleware',
    'django.contrib.messages.middleware.MessageMiddleware',
    'django.middleware.clickjacking.XFrameOptionsMiddleware',
)

ROOT_URLCONF = 'japper.urls'

WSGI_APPLICATION = 'japper.wsgi.application'

TEMPLATE_CONTEXT_PROCESSORS = (
    'django.contrib.auth.context_processors.auth',
    'django.core.context_processors.debug',
    'django.core.context_processors.i18n',
    'django.core.context_processors.media',
    'django.core.context_processors.static',
    'django.core.context_processors.tz',
    'django.contrib.messages.context_processors.messages',
    'django.core.context_processors.request',
)

# Database
# https://docs.djangoproject.com/en/1.7/ref/settings/#databases

DATABASES = {
    'default': env.db(default='postgres://localhost/japper'),
}

# Internationalization
# https://docs.djangoproject.com/en/1.7/topics/i18n/

LANGUAGE_CODE = 'en-us'

TIME_ZONE = 'UTC'

USE_I18N = True

USE_L10N = True

USE_TZ = True

# Static files (CSS, JavaScript, Images)
# https://docs.djangoproject.com/en/1.7/howto/static-files/

STATIC_URL = '/static/'
STATIC_ROOT = env('STATIC_ROOT')

# Crispy forms settings

CRISPY_TEMPLATE_PACK = 'bootstrap3'

# Celery settings

CELERYBEAT_SCHEDULE = {
    'monitoring:fetch_check_results': {
        'task': 'japper.monitoring.tasks.fetch_check_results',
        'schedule': datetime.timedelta(minutes=1),
    },
    'monitoring:cleanup': {
        'task': 'japper.monitoring.tasks.cleanup',
        'schedule': datetime.timedelta(minutes=5),
    },
}

# Email settings

globals().update(env.email_url(default='smtp://localhost'))

# Auth settings

LOGIN_REDIRECT_URL = 'monitoring_problems'
LOGIN_URL = 'login'
LOGOUT_URL = 'logout'

# Proxy settings

SECURE_PROXY_SSL_HEADER = ('X-Forwarded-Proto', 'https')

# Celery settings

BROKER_URL = env('BROKER_URL')
<|endoftext|>


<commit_before>from darch import searchers as se, surrogates as su, core as co
import benchmarks.datasets as datasets
import numpy as np
import torch
import matplotlib.pyplot as pyplot
import seaborn as sns; sns.set()

# TODO: Change the below imports when using PyTorch
import darch.contrib.search_spaces.tensorflow.dnn as search_dnn
from darch.contrib.evaluators.tensorflow.classification import SimpleClassifierEvaluator
from darch.contrib.search_spaces.tensorflow.common import D
from darch.contrib.datasets.dataset import InMemoryDataset
from sklearn.model_selection import train_test_split

# Python >= 3.5

# Some requirements of surrogate models:
# 1. They should be fast to evaluate (justification)
# 2. They should be accurate (replacement)

class CLSTMSurrogateModel(torch.nn.Module):
    """ The actual CLSTM model
    """
    def __init__(self, character_list, hidden_size, embedding_size):
        super(CLSTMSurrogateModel, self).__init__()
        # One embedding feeds 4 LSTMs
        self.embeddings = torch.nn.Embedding(len(character_list), embedding_size)
        self.lstm_in = []
        self.h0 = []
        self.c0 = []
        for i in range(4):
            # One LSTM for each input feature
            self.lstm_in.append(torch.nn.LSTM(embedding_size, hidden_size, 1))
            # Learnable initial hidden states and cell states
            self.h0.append(torch.nn.Parameter(torch.randn(1, 1, hidden_size)))
            self.c0.append(torch.nn.Parameter(torch.randn(1, 1, hidden_size)))
        # LSTM to read the concatenated outputs of the input LSTMs
        self.lstm_out = torch.nn.LSTM(4 * hidden_size, hidden_size, 1)
        self.h_out = torch.nn.Parameter(torch.randn(1, 1, hidden_size))
        self.c_out = torch.nn.Parameter(torch.randn(1, 1, hidden_size))
        self.fc_out = torch.nn.Linear(hidden_size, 1)
    def forward(self, feats):
        out_0 = self.embeddings(feats[0]).unsqueeze(1)
        out_1 = self.embeddings(feats[1]).unsqueeze(1)
        out_2 = self.embeddings(feats[2]).unsqueeze(1)
        out_3 = self.embeddings(feats[3]).unsqueeze(1)
        _, (out_0, _) = self.lstm_in[0](out_0, (self.h0[0], self.c0[0]))
        _, (out_1, _) = self.lstm_in[1](out_1, (self.h0[1], self.c0[1]))
        _, (out_2, _) = self.lstm_in[2](out_2, (self.h0[2], self.c0[2]))
        _, (out_3, _) = self.lstm_in[3](out_3, (self.h0[3], self.c0[3]))
        print(out_0.size())
        out = torch.cat((out_0, out_1, out_2, out_3), dim=2)
        print(out.size())
        _, (out, _) = self.lstm_out(out, (self.h_out, self.c_out))
        out = self.fc_out(out)
        print(out)
        return out

class CLSTMSurrogate(su.SurrogateModel):
    """ The CLSTM Surrogate Function
    """
    # Character LSTM: One LSTM for each feature vector,
    # For a total of 3, then concat the outputs and learn
    character_list = [chr(i) for i in range(ord('A'), ord('Z'))] + [
        chr(i) for i in range(ord('a'), ord('z') + 1)] + [
        chr(i) for i in range(ord('0'), ord('9') + 1)] + [
        '.', ':', '-', '_', '<', '>', '/', '=', '*', ' ', '|']
    char_to_index = {ch : idx for (idx, ch) in enumerate(character_list)}
    # Each input LSTM has an embedding size and hidden state of 128
    embedding_size = 128
    hidden_size = 128

    def __init__(self):
        self.model = CLSTMSurrogateModel(self.character_list, self.embedding_size, self.hidden_size)
        self.optimizer = torch.optim.Adam(self.model.parameters())
        self.loss_fn = torch.nn.MSELoss()

    def preprocess(self, feats):
        # Feats is a 4-tuple of a list of strings. Need to convert this into a tensor
        # To make this into a tensor, we pad * symbols to the right of all shorter string lengths
        # Convert strings to tensor by mapping the characters to indicies in the character_list
        # Final output is a list of 4 Char Tensors
        # max_len = max([max(len(s) for s in feat) for feat in feats]) + 2
        output = []
        for i, feat in enumerate(feats):
            # vec_feat = [self.char_to_index['*']]
            vec_feat = []
            for obj in feat:
                for char in obj:
                    vec_feat.append(self.char_to_index[char])
            # print(len(vec_feat), max_len)
            # while len(vec_feat) < max_len:
            #     vec_feat.append(self.char_to_index['*'])
            output.append(torch.autograd.Variable(torch.LongTensor(vec_feat)))
        return output

    def eval(self, feats):
        val = self.model(self.preprocess(feats))
        # print(val)
        return float(val)

    def update(self, val, feats):
        self.optimizer.zero_grad()               # Zero out the gradient buffer
        # TODO: refactor api to pass this in as an arg?
        out = self.model(self.preprocess(feats)) # Need to copmute network output
        loss = self.loss_fn(out, val)
        loss.backward()
        self.optimizer.step()

def test_clstm_surrogate():
    ## TODO: Remove Tensorflow dependency so TF doesn't eat up all GPU memory
    ## TODO: Only use PyTorch in the benchmark?

    ## Choose our dataset and split train/valid/test data
    datasets.IRIS('.temp/data')
    # For tensorflow usage, hardcoding in the data, once I change the contrib to PyTorch, can use the dataloaders
    def iris_labels(label):
            if label == b'Iris-setosa':
                return 0
            elif label == b'Iris-versicolor':
                return 1
            elif label == b'Iris-virginica':
                return 2
    data = np.loadtxt('.temp/data/iris.data', delimiter=',', usecols=[0,1,2,3])
    labels = np.loadtxt('.temp/data/iris.data', delimiter=',', usecols=[4], converters={4: iris_labels})
    # Split data
    X_train, X_test, y_train, y_test = train_test_split(data, labels, test_size=0.33)
    X_valid, X_test, y_valid, y_test = train_test_split(X_test, y_test, test_size=0.50)
    # Define datasets for contrib tensorflow evaluator
    train_dataset = InMemoryDataset(X_train, y_train, True)
    val_dataset = InMemoryDataset(X_valid, y_valid, False)
    test_dataset = InMemoryDataset(X_test, y_test, False)
    num_classes = 3

    ## Declare the model evaluator
    evaluator = SimpleClassifierEvaluator(train_dataset, val_dataset, num_classes,
        './temp', max_eval_time_in_minutes=1.0, log_output_to_terminal=True,
        batch_size=256)

    ## Define our search space
    def search_space_fn():
        co.Scope.reset_default_scope()
        inputs, outputs = search_dnn.dnn_net(2)
        return inputs, outputs, {'learning_rate_init' : D([1e-2, 1e-3, 1e-4, 1e-5])}

    ## Define our surrogate models
    clstm_sur = CLSTMSurrogate()
    baseline_sur = su.HashingSurrogate(1024, 100000)

    ## Choose our searching algorithm
    searcher_clstm = se.SMBOSearcher(search_space_fn, clstm_sur, 32, 0.2)
    searcher_baseline = se.SMBOSearcher(search_space_fn, baseline_sur, 32, 0.2)

    ## Search the model space
    clstm_pred_accs = []
    baseline_pred_accs = []
    clstm_true_accs = []
    baseline_true_accs = []

    for _ in range(128):
        for model, searcher, pred_accs, true_accs in zip([clstm_sur, baseline_sur], [searcher_clstm, searcher_baseline],
            [clstm_pred_accs, baseline_pred_accs], [clstm_true_accs, baseline_true_accs]):
            # Sample from our searcher
            (inputs, outputs, hs, hyperp_value_hist, searcher_eval_token) = searcher.sample()
            # Since Searcher doesn't return the score, we recompute it again, should be cheap since that's a goal
            feats = su.extract_features(inputs, outputs, hs)
            score = model.eval(feats)
            pred_accs.append(score)
            # Get the true score
            val_acc = evaluator.eval(inputs, outputs, hs)['val_acc']
            true_accs.append(val_acc)
            searcher.update(val_acc, searcher_eval_token)

if __name__ == '__main__':
    test_clstm_surrogate()<commit_msg>Proper return in eval for CLSTM surrogate<commit_after>from darch import searchers as se, surrogates as su, core as co
import benchmarks.datasets as datasets
import numpy as np
import torch
import matplotlib.pyplot as pyplot
import seaborn as sns; sns.set()

# TODO: Change the below imports when using PyTorch
import darch.contrib.search_spaces.tensorflow.dnn as search_dnn
from darch.contrib.evaluators.tensorflow.classification import SimpleClassifierEvaluator
from darch.contrib.search_spaces.tensorflow.common import D
from darch.contrib.datasets.dataset import InMemoryDataset
from sklearn.model_selection import train_test_split

# Python >= 3.5

# Some requirements of surrogate models:
# 1. They should be fast to evaluate (justification)
# 2. They should be accurate (replacement)

class CLSTMSurrogateModel(torch.nn.Module):
    """ The actual CLSTM model
    """
    def __init__(self, character_list, hidden_size, embedding_size):
        super(CLSTMSurrogateModel, self).__init__()
        # One embedding feeds 4 LSTMs
        self.embeddings = torch.nn.Embedding(len(character_list), embedding_size)
        self.lstm_in = []
        self.h0 = []
        self.c0 = []
        for i in range(4):
            # One LSTM for each input feature
            self.lstm_in.append(torch.nn.LSTM(embedding_size, hidden_size, 1))
            # Learnable initial hidden states and cell states
            self.h0.append(torch.nn.Parameter(torch.randn(1, 1, hidden_size)))
            self.c0.append(torch.nn.Parameter(torch.randn(1, 1, hidden_size)))
        # LSTM to read the concatenated outputs of the input LSTMs
        self.lstm_out = torch.nn.LSTM(4 * hidden_size, hidden_size, 1)
        self.h_out = torch.nn.Parameter(torch.randn(1, 1, hidden_size))
        self.c_out = torch.nn.Parameter(torch.randn(1, 1, hidden_size))
        self.fc_out = torch.nn.Linear(hidden_size, 1)
    def forward(self, feats):
        out_0 = self.embeddings(feats[0]).unsqueeze(1)
        out_1 = self.embeddings(feats[1]).unsqueeze(1)
        out_2 = self.embeddings(feats[2]).unsqueeze(1)
        out_3 = self.embeddings(feats[3]).unsqueeze(1)
        _, (out_0, _) = self.lstm_in[0](out_0, (self.h0[0], self.c0[0]))
        _, (out_1, _) = self.lstm_in[1](out_1, (self.h0[1], self.c0[1]))
        _, (out_2, _) = self.lstm_in[2](out_2, (self.h0[2], self.c0[2]))
        _, (out_3, _) = self.lstm_in[3](out_3, (self.h0[3], self.c0[3]))
        out = torch.cat((out_0, out_1, out_2, out_3), dim=2)
        _, (out, _) = self.lstm_out(out, (self.h_out, self.c_out))
        out = self.fc_out(out)
        return out

class CLSTMSurrogate(su.SurrogateModel):
    """ The CLSTM Surrogate Function
    """
    # Character LSTM: One LSTM for each feature vector,
    # For a total of 3, then concat the outputs and learn
    character_list = [chr(i) for i in range(ord('A'), ord('Z'))] + [
        chr(i) for i in range(ord('a'), ord('z') + 1)] + [
        chr(i) for i in range(ord('0'), ord('9') + 1)] + [
        '.', ':', '-', '_', '<', '>', '/', '=', '*', ' ', '|']
    char_to_index = {ch : idx for (idx, ch) in enumerate(character_list)}
    # Each input LSTM has an embedding size and hidden state of 128
    embedding_size = 128
    hidden_size = 128

    def __init__(self):
        self.model = CLSTMSurrogateModel(self.character_list, self.embedding_size, self.hidden_size)
        self.optimizer = torch.optim.Adam(self.model.parameters())
        self.loss_fn = torch.nn.MSELoss()

    def preprocess(self, feats):
        # Feats is a 4-tuple of a list of strings. Need to convert this into a tensor
        # To make this into a tensor, we pad * symbols to the right of all shorter string lengths
        # Convert strings to tensor by mapping the characters to indicies in the character_list
        # Final output is a list of 4 Char Tensors
        # max_len = max([max(len(s) for s in feat) for feat in feats]) + 2
        output = []
        for i, feat in enumerate(feats):
            # vec_feat = [self.char_to_index['*']]
            vec_feat = []
            for obj in feat:
                for char in obj:
                    vec_feat.append(self.char_to_index[char])
            # print(len(vec_feat), max_len)
            # while len(vec_feat) < max_len:
            #     vec_feat.append(self.char_to_index['*'])
            output.append(torch.autograd.Variable(torch.LongTensor(vec_feat)))
        return output

    def eval(self, feats):
        return self.model(self.preprocess(feats)).data[0, 0, 0]

    def update(self, val, feats):
        self.optimizer.zero_grad()               # Zero out the gradient buffer
        # TODO: refactor api to pass this in as an arg?
        out = self.model(self.preprocess(feats)) # Need to copmute network output
        loss = self.loss_fn(out, val)
        loss.backward()
        self.optimizer.step()

def test_clstm_surrogate():
    ## TODO: Remove Tensorflow dependency so TF doesn't eat up all GPU memory
    ## TODO: Only use PyTorch in the benchmark?

    ## Choose our dataset and split train/valid/test data
    datasets.IRIS('.temp/data')
    # For tensorflow usage, hardcoding in the data, once I change the contrib to PyTorch, can use the dataloaders
    def iris_labels(label):
            if label == b'Iris-setosa':
                return 0
            elif label == b'Iris-versicolor':
                return 1
            elif label == b'Iris-virginica':
                return 2
    data = np.loadtxt('.temp/data/iris.data', delimiter=',', usecols=[0,1,2,3])
    labels = np.loadtxt('.temp/data/iris.data', delimiter=',', usecols=[4], converters={4: iris_labels})
    # Split data
    X_train, X_test, y_train, y_test = train_test_split(data, labels, test_size=0.33)
    X_valid, X_test, y_valid, y_test = train_test_split(X_test, y_test, test_size=0.50)
    # Define datasets for contrib tensorflow evaluator
    train_dataset = InMemoryDataset(X_train, y_train, True)
    val_dataset = InMemoryDataset(X_valid, y_valid, False)
    test_dataset = InMemoryDataset(X_test, y_test, False)
    num_classes = 3

    ## Declare the model evaluator
    evaluator = SimpleClassifierEvaluator(train_dataset, val_dataset, num_classes,
        './temp', max_eval_time_in_minutes=1.0, log_output_to_terminal=True,
        batch_size=256)

    ## Define our search space
    def search_space_fn():
        co.Scope.reset_default_scope()
        inputs, outputs = search_dnn.dnn_net(2)
        return inputs, outputs, {'learning_rate_init' : D([1e-2, 1e-3, 1e-4, 1e-5])}

    ## Define our surrogate models
    clstm_sur = CLSTMSurrogate()
    baseline_sur = su.HashingSurrogate(1024, 100000)

    ## Choose our searching algorithm
    searcher_clstm = se.SMBOSearcher(search_space_fn, clstm_sur, 32, 0.2)
    searcher_baseline = se.SMBOSearcher(search_space_fn, baseline_sur, 32, 0.2)

    ## Search the model space
    clstm_pred_accs = []
    baseline_pred_accs = []
    clstm_true_accs = []
    baseline_true_accs = []

    for _ in range(128):
        for model, searcher, pred_accs, true_accs in zip([clstm_sur, baseline_sur], [searcher_clstm, searcher_baseline],
            [clstm_pred_accs, baseline_pred_accs], [clstm_true_accs, baseline_true_accs]):
            # Sample from our searcher
            (inputs, outputs, hs, hyperp_value_hist, searcher_eval_token) = searcher.sample()
            # Since Searcher doesn't return the score, we recompute it again, should be cheap since that's a goal
            feats = su.extract_features(inputs, outputs, hs)
            score = model.eval(feats)
            pred_accs.append(score)
            # Get the true score
            val_acc = evaluator.eval(inputs, outputs, hs)['val_acc']
            true_accs.append(val_acc)
            searcher.update(val_acc, searcher_eval_token)

if __name__ == '__main__':
    test_clstm_surrogate()<|endoftext|>