<commit_before>'''Tools for the lazy amongst us: automation of common HANDE analysis tasks.'''

import collections
import pandas as pd
import warnings

import pyblock
import pyhande.extract
import pyhande.analysis
import pyhande.weight
import math

def std_analysis(datafiles, start=0, select_function=None, extract_psips=False,
                reweight_history=0, mean_shift=0.0, arith_mean=False):
    '''Perform a 'standard' analysis of HANDE output files.

Parameters
----------
datafiles : list of strings
    names of files containing HANDE QMC calculation output.
start : int
    iteration from which the blocking analysis is performed.
select_function : function
    function which returns a boolean mask for the iterations to include in the
    analysis.  Not used if set to None (default).  Overrides ``start``.  See
    below for examples.
extract_psips : bool
    also extract the mean number of psips from the calculation.
reweight_history : integer
    reweight in an attempt to remove population control bias. According to
    [Umrigar93]_ this should be set to be a few correlation times.
mean_shift : float
    prevent the weights from beoming to large.

Returns
-------
info : list of :func:`collections.namedtuple`
    raw and analysed data, consisting of:

        metadata, data
            from :func:`pyhande.extract.extract_data_sets`.  If ``data``
            consists of several concatenated calculations, then the only
            ``metadata`` object is from the first calculation.
        data_len, reblock, covariance
            from :func:`pyblock.pd_utils.reblock`.  The projected energy
            estimator (evaluated by :func:`pyhande.analysis.projected_energy`)
            is included in ``reblock``.
        opt_block, no_opt_block
            from :func:`pyhande.analysis.qmc_summary`.  A 'pretty-printed'
            estimate string is included in ``opt_block``.

Examples
--------

The following are equivalent and will extract the data from the file called
hande.fciqmc.out, perform a blocking analysis from the 10000th iteration
onwards, calculated the projected energy estimator and find the optimal block
size from the blocking analysis:

>>> std_analysis(['hande.fciqmc.out'], 10000)
>>> std_analysis(['hande.fciqmc.out'],
...              select_function=lambda d: d['iterations'] > 10000)

References
----------
Umrigar93
    Umrigar et al., J. Chem. Phys. 99, 2865 (1993).
'''

    hande_out = pyhande.extract.extract_data_sets(datafiles)

    # Concat all QMC data (We did say 'lazy', so assumptions are being made...)
    data = []
    metadata = []
    for (md, df) in hande_out:
        if any(calc in md['calc_type'] for calc in ('FCIQMC', 'CCMC')):
            if reweight_history > 0:
                if 'mc_cycles' in md:
                    df = pyhande.weight.reweight(df, md['mc_cycles'],
                        md['tau'], reweight_history, mean_shift,
                        arith_mean=arith_mean)
                else:
                    df = pyhande.weight.reweight(df, md['qmc']['ncycles'],
                        md['qmc']['tau'], reweight_history, mean_shift,
                        arith_mean=arith_mean)
                df['W * \sum H_0j N_j'] = df['\sum H_0j N_j'] * df['Weight']
                df['W * N_0'] = df['N_0'] * df['Weight']
            data.append(df)
            metadata.append(md)
    if data:
        calcs_metadata, calcs = concat_calcs(metadata, data)
    else:
        raise ValueError('No data found in '+' '.join(datafiles))

    tuple_fields = ('metadata data data_len reblock covariance opt_block '
                   'no_opt_block'.split())
    info_tuple = collections.namedtuple('HandeInfo', tuple_fields)
    return_vals = []
    for calc,md in zip(calcs,calcs_metadata):
        # Reblock Monte Carlo data over desired window.
        if select_function is None:
            indx = calc['iterations'] > start
        else:
            indx = select_function(calc)
        to_block = []
        if extract_psips:
            to_block.append('# H psips')
        to_block.extend(['\sum H_0j N_j', 'N_0', 'Shift'])
        if reweight_history > 0:
            to_block.extend(['W * \sum H_0j N_j', 'W * N_0'])

        mc_data = calc.ix[indx, to_block]

        if mc_data['Shift'].iloc[0] == mc_data['Shift'].iloc[1]:
            warnings.warn('The blocking analysis starts from before the shift '
                         'begins to vary.')

        (data_len, reblock, covariance) = pyblock.pd_utils.reblock(mc_data)

        proje = pyhande.analysis.projected_energy(reblock, covariance, data_len)
        reblock = pd.concat([reblock, proje], axis=1)
        to_block.append('Proj. Energy')

        if reweight_history > 0:
            proje = pyhande.analysis.projected_energy(reblock, covariance,
                        data_len, sum_key='W * \sum H_0j N_j', ref_key='W * N_0',
                        col_name='Weighted Proj. E.')
            reblock = pd.concat([reblock, proje], axis=1)
            to_block.append('Weighted Proj. E.')

        # Summary (including pretty printing of estimates).
        (opt_block, no_opt_block) = pyhande.analysis.qmc_summary(reblock, to_block)

        # Now try to calculate an inefficiency according to
        # Understanding and improving the efficiency of full configuration interaction quantum Monte Carlo
        # W. A. Vigor, J. S. Spencer, M. J. Bearpark, and A. J. W. Thom
        # J. Chem. Phys. 144, 094110 (2016); doi: 10.1063/1.4943113

        #  We also try to provide some sort of error estimate for the inefficiency on the basis of the relative standard error error of the projected energy numerator..
        try:
            try:
                dtau = md['tau']
            except:
                dtau = md['qmc']['tau']
            err_proj_e = opt_block['standard error']['Proj. Energy']
            Np = opt_block['mean']['# H psips']
            N = calc['iterations'].iloc[-1] - start
            inefficiency = err_proj_e * math.sqrt(Np*N*dtau)
#    # Really we should care about the covariance etc. but it is really horrible for the
#    # projected energy
            err_err_proj_e = err_proj_e*math.sqrt( (opt_block['standard error error']['\sum H_0j N_j']/(opt_block['standard error']['\sum H_0j N_j']  ))**2  + ((opt_block['standard error error']['N_0'])/(opt_block['standard error']['N_0']))**2 )
            err_ineff = inefficiency*math.sqrt(((err_err_proj_e/err_proj_e)**2) + (opt_block['standard error']['# H psips']/(2*opt_block['mean']['# H psips']))**2)
            d = pd.DataFrame(data={'mean':inefficiency, 'standard error':err_ineff}, index = ['Inefficiency'])
            opt_block = opt_block.append(d)
        except:
            pass

        estimates = []
        for (name, row) in opt_block.iterrows():
            estimates.append(
                    pyblock.error.pretty_fmt_err(row['mean'], row['standard error'])
                           )
        opt_block['estimate'] = estimates

        return_vals.append(info_tuple(md, calc, data_len, reblock,
                                      covariance, opt_block, no_opt_block))

    return return_vals

def concat_calcs(metadata, data):
    '''Concatenate data from restarted calculations to analyse together.

Parameters
----------
metadata : list of dicts
    Extracted metadata for each calculation.
data : list of :class:`pandas.DataFrame`
    Output of each QMC calculation.

Returns
-------
calcs_metadata : list of dicts
    Metadata for each calculation, with duplicates from restarting dropped.
calcs : list of :class:`pandas.DataFrame`
    Output of each QMC calculation, with parts of a restarted calculation combined.
'''

    restart_uuids = [md['restart'].get('uuid_restart','') for md in metadata]
    uuids = [md['UUID'] for md in metadata]
    if any(restart_uuids) and all(uuids):
        data = list(data)
        metadata = list(metadata)
        calcs = []
        calcs_metadata = []
        while uuids:
            for indx in range(len(uuids)):
                if uuids[indx] not in restart_uuids:
                    # Found the end of a chain.
                    break
            uuid = uuids.pop(indx)
            restart = restart_uuids.pop(indx)
            calc = [data.pop(indx)]
            calcs_metadata.append(metadata.pop(indx))
            while restart and restart in uuids:
                indx = uuids.index(restart)
                uuid = uuids.pop(indx)
                restart = restart_uuids.pop(indx)
                calc.append(data.pop(indx))
                metadata.pop(indx)
            calcs.append(pd.concat(calc[::-1]))

    else:
        # Don't have UUID information.
        # Assume any restarted calculations are in the right order and contiguous
        # Check concatenating data is at least possibly sane.
        step = data[0]['iterations'].iloc[-1] - data[0]['iterations'].iloc[-2]
        prev_iteration = data[0]['iterations'].iloc[-1]
        calc_type = metadata[0]['calc_type']
        calcs = []
        calcs_metadata = [metadata[0]]
        xcalc = [data[0]]
        for i in range(1, len(data)):
            if metadata[i]['calc_type'] != calc_type or \
                    data[i]['iterations'].iloc[0] - step != prev_iteration or \
                    data[i]['iterations'].iloc[-1] - data[i]['iterations'].iloc[-2] != step:
                # Different calculation
                step = data[i]['iterations'].iloc[-1] - data[i]['iterations'].iloc[-2]
                calc_type = metadata[i]['calc_type']
                calcs.append(pd.concat(xcalc))
                xcalc = [data[i]]
                calcs_metadata.append(metadata[i])
            else:
                # Continuation of same calculation (probably)
                xcalc.append(data[i])
            prev_iteration = data[i]['iterations'].iloc[-1]
        calcs.append(pd.concat(xcalc, ignore_index=True))
    return calcs_metadata, calcs
<commit_msg>Pedantic comment<commit_after>'''Tools for the lazy amongst us: automation of common HANDE analysis tasks.'''

import collections
import pandas as pd
import warnings

import pyblock
import pyhande.extract
import pyhande.analysis
import pyhande.weight
import math

def std_analysis(datafiles, start=0, select_function=None, extract_psips=False,
                reweight_history=0, mean_shift=0.0, arith_mean=False):
    '''Perform a 'standard' analysis of HANDE output files.

Parameters
----------
datafiles : list of strings
    names of files containing HANDE QMC calculation output.
start : int
    iteration from which the blocking analysis is performed.
select_function : function
    function which returns a boolean mask for the iterations to include in the
    analysis.  Not used if set to None (default).  Overrides ``start``.  See
    below for examples.
extract_psips : bool
    also extract the mean number of psips from the calculation.
reweight_history : integer
    reweight in an attempt to remove population control bias. According to
    [Umrigar93]_ this should be set to be a few correlation times.
mean_shift : float
    prevent the weights from beoming to large.

Returns
-------
info : list of :func:`collections.namedtuple`
    raw and analysed data, consisting of:

        metadata, data
            from :func:`pyhande.extract.extract_data_sets`.  If ``data``
            consists of several concatenated calculations, then the only
            ``metadata`` object is from the first calculation.
        data_len, reblock, covariance
            from :func:`pyblock.pd_utils.reblock`.  The projected energy
            estimator (evaluated by :func:`pyhande.analysis.projected_energy`)
            is included in ``reblock``.
        opt_block, no_opt_block
            from :func:`pyhande.analysis.qmc_summary`.  A 'pretty-printed'
            estimate string is included in ``opt_block``.

Examples
--------

The following are equivalent and will extract the data from the file called
hande.fciqmc.out, perform a blocking analysis from the 10000th iteration
onwards, calculated the projected energy estimator and find the optimal block
size from the blocking analysis:

>>> std_analysis(['hande.fciqmc.out'], 10000)
>>> std_analysis(['hande.fciqmc.out'],
...              select_function=lambda d: d['iterations'] > 10000)

References
----------
Umrigar93
    Umrigar et al., J. Chem. Phys. 99, 2865 (1993).
'''

    hande_out = pyhande.extract.extract_data_sets(datafiles)

    # Concat all QMC data (We did say 'lazy', so assumptions are being made...)
    data = []
    metadata = []
    for (md, df) in hande_out:
        if any(calc in md['calc_type'] for calc in ('FCIQMC', 'CCMC')):
            if reweight_history > 0:
                if 'mc_cycles' in md:
                    df = pyhande.weight.reweight(df, md['mc_cycles'],
                        md['tau'], reweight_history, mean_shift,
                        arith_mean=arith_mean)
                else:
                    df = pyhande.weight.reweight(df, md['qmc']['ncycles'],
                        md['qmc']['tau'], reweight_history, mean_shift,
                        arith_mean=arith_mean)
                df['W * \sum H_0j N_j'] = df['\sum H_0j N_j'] * df['Weight']
                df['W * N_0'] = df['N_0'] * df['Weight']
            data.append(df)
            metadata.append(md)
    if data:
        calcs_metadata, calcs = concat_calcs(metadata, data)
    else:
        raise ValueError('No data found in '+' '.join(datafiles))

    tuple_fields = ('metadata data data_len reblock covariance opt_block '
                   'no_opt_block'.split())
    info_tuple = collections.namedtuple('HandeInfo', tuple_fields)
    return_vals = []
    for calc,md in zip(calcs,calcs_metadata):
        # Reblock Monte Carlo data over desired window.
        if select_function is None:
            indx = calc['iterations'] > start
        else:
            indx = select_function(calc)
        to_block = []
        if extract_psips:
            to_block.append('# H psips')
        to_block.extend(['\sum H_0j N_j', 'N_0', 'Shift'])
        if reweight_history > 0:
            to_block.extend(['W * \sum H_0j N_j', 'W * N_0'])

        mc_data = calc.ix[indx, to_block]

        if mc_data['Shift'].iloc[0] == mc_data['Shift'].iloc[1]:
            warnings.warn('The blocking analysis starts from before the shift '
                         'begins to vary.')

        (data_len, reblock, covariance) = pyblock.pd_utils.reblock(mc_data)

        proje = pyhande.analysis.projected_energy(reblock, covariance, data_len)
        reblock = pd.concat([reblock, proje], axis=1)
        to_block.append('Proj. Energy')

        if reweight_history > 0:
            proje = pyhande.analysis.projected_energy(reblock, covariance,
                        data_len, sum_key='W * \sum H_0j N_j', ref_key='W * N_0',
                        col_name='Weighted Proj. E.')
            reblock = pd.concat([reblock, proje], axis=1)
            to_block.append('Weighted Proj. E.')

        # Summary (including pretty printing of estimates).
        (opt_block, no_opt_block) = pyhande.analysis.qmc_summary(reblock, to_block)

        # Now try to calculate an inefficiency according to
        # Understanding and improving the efficiency of full configuration interaction quantum Monte Carlo
        # W. A. Vigor, J. S. Spencer, M. J. Bearpark, and A. J. W. Thom
        # J. Chem. Phys. 144, 094110 (2016); doi: 10.1063/1.4943113

        #  We also try to provide some sort of error estimate for the inefficiency on the basis of the relative standard error error of the projected energy numerator..
        try:
            try:
                dtau = md['tau']
            except:
                dtau = md['qmc']['tau']
            err_proj_e = opt_block['standard error']['Proj. Energy']
            Np = opt_block['mean']['# H psips']
	#Review - [VAN]: are you not forgetting the start iteration itself for N. Also I would need to check where reblocking actually starts: (probably really irrelevant)
            N = calc['iterations'].iloc[-1] - start
            inefficiency = err_proj_e * math.sqrt(Np*N*dtau)
#    # Really we should care about the covariance etc. but it is really horrible for the
#    # projected energy
            err_err_proj_e = err_proj_e*math.sqrt( (opt_block['standard error error']['\sum H_0j N_j']/(opt_block['standard error']['\sum H_0j N_j']  ))**2  + ((opt_block['standard error error']['N_0'])/(opt_block['standard error']['N_0']))**2 )
            err_ineff = inefficiency*math.sqrt(((err_err_proj_e/err_proj_e)**2) + (opt_block['standard error']['# H psips']/(2*opt_block['mean']['# H psips']))**2)
            d = pd.DataFrame(data={'mean':inefficiency, 'standard error':err_ineff}, index = ['Inefficiency'])
            opt_block = opt_block.append(d)
        except:
            pass

        estimates = []
        for (name, row) in opt_block.iterrows():
            estimates.append(
                    pyblock.error.pretty_fmt_err(row['mean'], row['standard error'])
                           )
        opt_block['estimate'] = estimates

        return_vals.append(info_tuple(md, calc, data_len, reblock,
                                      covariance, opt_block, no_opt_block))

    return return_vals

def concat_calcs(metadata, data):
    '''Concatenate data from restarted calculations to analyse together.

Parameters
----------
metadata : list of dicts
    Extracted metadata for each calculation.
data : list of :class:`pandas.DataFrame`
    Output of each QMC calculation.

Returns
-------
calcs_metadata : list of dicts
    Metadata for each calculation, with duplicates from restarting dropped.
calcs : list of :class:`pandas.DataFrame`
    Output of each QMC calculation, with parts of a restarted calculation combined.
'''

    restart_uuids = [md['restart'].get('uuid_restart','') for md in metadata]
    uuids = [md['UUID'] for md in metadata]
    if any(restart_uuids) and all(uuids):
        data = list(data)
        metadata = list(metadata)
        calcs = []
        calcs_metadata = []
        while uuids:
            for indx in range(len(uuids)):
                if uuids[indx] not in restart_uuids:
                    # Found the end of a chain.
                    break
            uuid = uuids.pop(indx)
            restart = restart_uuids.pop(indx)
            calc = [data.pop(indx)]
            calcs_metadata.append(metadata.pop(indx))
            while restart and restart in uuids:
                indx = uuids.index(restart)
                uuid = uuids.pop(indx)
                restart = restart_uuids.pop(indx)
                calc.append(data.pop(indx))
                metadata.pop(indx)
            calcs.append(pd.concat(calc[::-1]))

    else:
        # Don't have UUID information.
        # Assume any restarted calculations are in the right order and contiguous
        # Check concatenating data is at least possibly sane.
        step = data[0]['iterations'].iloc[-1] - data[0]['iterations'].iloc[-2]
        prev_iteration = data[0]['iterations'].iloc[-1]
        calc_type = metadata[0]['calc_type']
        calcs = []
        calcs_metadata = [metadata[0]]
        xcalc = [data[0]]
        for i in range(1, len(data)):
            if metadata[i]['calc_type'] != calc_type or \
                    data[i]['iterations'].iloc[0] - step != prev_iteration or \
                    data[i]['iterations'].iloc[-1] - data[i]['iterations'].iloc[-2] != step:
                # Different calculation
                step = data[i]['iterations'].iloc[-1] - data[i]['iterations'].iloc[-2]
                calc_type = metadata[i]['calc_type']
                calcs.append(pd.concat(xcalc))
                xcalc = [data[i]]
                calcs_metadata.append(metadata[i])
            else:
                # Continuation of same calculation (probably)
                xcalc.append(data[i])
            prev_iteration = data[i]['iterations'].iloc[-1]
        calcs.append(pd.concat(xcalc, ignore_index=True))
    return calcs_metadata, calcs
<|endoftext|>


<commit_before>from __future__ import division

import numpy as np
import operator
import random

np.seterr(all='raise')

class CircleComputeException(Exception):
    pass

class UndefiendCircle():
    def __contains__(self, item): return False

class Circle:

    def __init__(self, x, y, r):
        self.x = x
        self.y = y
        self.r = r

    def center(self):
        return (self.x, self.y)

    def radius(self):
        return self.radius

    @staticmethod
    def singularity(p):
        x, y = p
        return Circle(x,y,0)

    @staticmethod
    def from_two_points(p1, p2):
        x1, y1 = p1
        x2, y2 = p2

        if x1 == x2 or y1 == y2:
            pass# raise Exception("slope == 0 or slope == oo")

        c_x = (x1 + x2) / 2
        c_y = (y1 + y2) / 2
        r = scipy.spatial.distance.euclidean(p1,p2) / 2
        return Circle(c_x, c_y, r)

    @staticmethod
    def from_three_points(p1, p2, p3):
        """ http://www.ehow.com/how_5899905_radius-three-points.html   """
        x1,y1 = p1
        x2,y2 = p2
        x3,y3 = p3

        x1,y1 = float(x1), float(y1)
        x2,y2 = float(x2), float(y2)
        x3,y3 = float(x3), float(y3)

        def central_point(p1, p2, p3, dim):
            foo = max(p1, p2, key=operator.itemgetter(dim))
            return min(p3, foo, key=operator.itemgetter(dim))

        # If the points are collinear, instead return a circle
        # which simply contains the central point and has the
        # others on its border

        if x1 == x2 == x3:
            top = max(p1, p2, p3, key=operator.itemgetter(1))
            bottom = min(p1, p2, p3, key=operator.itemgetter(1))
            return circle_from_two_points(top,bottom)

        if y1 == y2 == y3:
            left = max(p1, p2, p3, key=operator.itemgetter(0))
            right = min(p1, p2, p3, key=operator.itemgetter(0))
            return circle_from_two_points(left,right)

        x1s, y1s = x1 ** 2, y1 ** 2
        x2s, y2s = x2 ** 2, y2 ** 2
        x3s, y3s = x3 ** 2, y3 ** 2

        c_x = ((y1-y2) * (x3s+y3s) + (y2-y3) * (x1s + y1s) + (y3 - y1) * (x2s + y2s )) / \
              (2 * ((y1-y2)*x3  + (y3-y1)*x2 + (y2-y3)*x1))
        c_y = ((x1-x2) * (x3s+y3s) + (x2-x3) * (x1s + y1s) + (x3 - x1) * (x2s + y2s )) / \
              (2 * ((x1-x2)*y3  + (x3-x1)*y2 + (x2-x3)*y1))
        c_r = np.sqrt((c_x - x1)**2 + (c_y - y1)**2)

        return  Circle(c_x,c_y,c_r)

    def __contains__(self, p):
        return (p[0] - self.x)  ** 2 + (p[1] - self.y) ** 2 <= self.r ** 2

    def __str__(self):
        return "Center: ({}, {}), r: {}".format(self.x,self.y,self.r)

def minidisk(points):

    def empty(P):
        """ Checks whether a given set P is emtpy """
        return len(P) == 0

    def b_md(P, R):
        """ Let b_md(P, R) be the closed disk of smallest radius which
        contains all points in P with all points in R on its boundary.
        """
        if empty(R):
            return UndefiendCircle()
        elif len(R) == 1:
            [[x,y]] = R
            return Circle(x, y, 0)
        elif len(R) == 2:
            [p1, p2] = R
            return Circle.from_two_points(p1, p2)
        elif len(R) == 3:
            [p1, p2, p3] = R
            return Circle.from_three_points(p1, p2, p3)
        else:
            raise Exception("Should not happen!")

    def b_minidisk(P,R):
        """ Computes the smallest disk enclosing points P
        with points R on its boundary
        """
        if empty(P) or len(R) == 3:
            D = b_md(frozenset(), R)
        else:
            p = random.sample(P, 1)[0] # Take random element out of P
            p_as_set = frozenset((p,))
            D = b_minidisk(P - p_as_set, R)
            if p not in D:
                D = b_minidisk(P - p_as_set, R | p_as_set)
        return D

    return b_minidisk(points, frozenset() )
<commit_msg>Remove legacy files<commit_after><|endoftext|>


<commit_before>import functools
import logging
import multiprocessing
import signal
import socket
import sys

from oslo.config import cfg

from akanda.rug import event
from akanda.rug import health
from akanda.rug.openstack.common import log
from akanda.rug import metadata
from akanda.rug import notifications
from akanda.rug import scheduler
from akanda.rug import populate
from akanda.rug import worker
from akanda.rug.api import quantum as quantum_api

LOG = log.getLogger(__name__)

def shuffle_notifications(notification_queue, sched):
    """Copy messages from the notification queue into the scheduler.
    """
    while True:
        try:
            target, message = notification_queue.get()
            sched.handle_message(target, message)
        # FIXME(rods): if a signal arrive during an IO operation an
        # IOError is raised. We catch the exceptions in meantime
        # waiting for a better solution.
        except IOError:
            pass
        except KeyboardInterrupt:
            sched.stop()
            break

def register_and_load_opts(argv):

    # Set the logging format to include the process and thread, since
    # those aren't included in standard openstack logs but are useful
    # for the rug
    log_format = ':'.join('%(' + n + ')s'
                          for n in ['asctime',
                                    'levelname',
                                    'name',
                                    'processName',
                                    'threadName',
                                    'message'])
    log.set_defaults(log_format)

    # Configure the default log levels for some third-party packages
    # that are chatty
    cfg.set_defaults(log.log_opts,
                     default_log_levels=['amqplib=WARN',
                                         'qpid.messaging=INFO',
                                         'sqlalchemy=WARN',
                                         'keystoneclient=INFO',
                                         'stevedore=INFO',
                                         'eventlet.wsgi.server=WARN',
                                         'requests=WARN',
                                         ])

    cfg.CONF.register_opts([
        cfg.StrOpt('host',
                   default=socket.getfqdn(),
                   help="The hostname Akanda is running on"),

        # FIXME(dhellmann): Use a separate group for these auth params
        cfg.StrOpt('admin_user'),
        cfg.StrOpt('admin_password', secret=True),
        cfg.StrOpt('admin_tenant_name'),
        cfg.StrOpt('auth_url'),
        cfg.StrOpt('auth_strategy', default='keystone'),
        cfg.StrOpt('auth_region'),

        cfg.StrOpt('management_network_id'),
        cfg.StrOpt('external_network_id'),
        cfg.StrOpt('management_subnet_id'),
        cfg.StrOpt('router_image_uuid'),

        cfg.StrOpt('management_prefix', default='fdca:3ba5:a17a:acda::/64'),
        cfg.IntOpt('akanda_mgt_service_port', default=5000),
        cfg.IntOpt('router_instance_flavor', default=1),

        # needed for plugging locally into management network
        cfg.StrOpt('interface_driver'),
        cfg.StrOpt('ovs_integration_bridge', default='br-int'),
        cfg.BoolOpt('ovs_use_veth', default=False),
        cfg.IntOpt('network_device_mtu'),

        # needed for boot waiting
        cfg.IntOpt('boot_timeout', default=240),
        cfg.IntOpt('max_retries', default=3),
        cfg.IntOpt('retry_delay', default=1),

        cfg.StrOpt(
            'ignored_router_directory',
            default='/etc/akanda-rug/ignored',
            help='Directory to scan for routers to ignore for debugging',
        ),

    ])

    cfg.CONF.register_opts(metadata.metadata_opts)

    AGENT_OPTIONS = [
        cfg.StrOpt('root_helper', default='sudo'),
    ]

    cfg.CONF.register_opts(AGENT_OPTIONS, 'AGENT')

    # FIXME: Convert these to regular options, not command line options.
    cfg.CONF.register_cli_opts([
        cfg.IntOpt('health-check-period',
                   default=60,
                   help='seconds between health checks'),
        cfg.IntOpt('num-worker-processes',
                   short='w',
                   default=16,
                   help='the number of worker processes to run'),
        cfg.IntOpt('num-worker-threads',
                   short='t',
                   default=4,
                   help='the number of worker threads to run per process'),

        # FIXME(dhellmann): set up a group for these messaging params
        cfg.StrOpt('amqp-url',
                   default='amqp://guest:secrete@localhost:5672/',
                   help='connection for AMQP server'),
        cfg.StrOpt('incoming-notifications-exchange',
                   default='neutron',
                   help='name of the exchange where we receive notifications'),
        cfg.StrOpt('outgoing-notifications-exchange',
                   default='neutron',
                   help='name of the exchange where we send notifications'),
        cfg.StrOpt('rpc-exchange',
                   default='l3_agent_fanout',
                   help='name of the exchange where we receive RPC calls'),
    ])

    cfg.CONF(argv, project='akanda')

def main(argv=sys.argv[1:]):
    register_and_load_opts(argv)

    log.setup('akanda-rug')
    cfg.CONF.log_opt_values(LOG, logging.INFO)

    # Purge the mgt tap interface on startup
    quantum = quantum_api.Quantum(cfg.CONF)
    #TODO(mark): develop better way restore after machine reboot
    #quantum.purge_management_interface()

    # bring the mgt tap interface up
    quantum.ensure_local_service_port()

    # Set up the queue to move messages between the eventlet-based
    # listening process and the scheduler.
    notification_queue = multiprocessing.Queue()

    # Listen for notifications.
    notification_proc = multiprocessing.Process(
        target=notifications.listen,
        kwargs={
            'host_id': cfg.CONF.host,
            'amqp_url': cfg.CONF.amqp_url,
            'notifications_exchange_name':
            cfg.CONF.incoming_notifications_exchange,
            'rpc_exchange_name': cfg.CONF.rpc_exchange,
            'notification_queue': notification_queue,
        },
        name='notification-listener',
    )
    notification_proc.start()

    mgt_ip_address = quantum_api.get_local_service_ip(cfg.CONF).split('/')[0]
    metadata_proc = multiprocessing.Process(
        target=metadata.serve,
        args=(mgt_ip_address,),
        name='metadata-proxy'
    )
    metadata_proc.start()

    # Set up the notifications publisher
    publisher = notifications.Publisher(
        cfg.CONF.amqp_url,
        exchange_name=cfg.CONF.outgoing_notifications_exchange,
        topic='notifications.info',
    )

    # Set up a factory to make Workers that know how many threads to
    # run.
    worker_factory = functools.partial(
        worker.Worker,
        num_threads=cfg.CONF.num_worker_threads,
        notifier=publisher,
        ignore_directory=cfg.CONF.ignored_router_directory,
    )

    # Set up the scheduler that knows how to manage the routers and
    # dispatch messages.
    sched = scheduler.Scheduler(
        num_workers=cfg.CONF.num_worker_processes,
        worker_factory=worker_factory,
    )

    def debug_workers(signum, stack):
        sched.handle_message(
            'debug',
            event.Event('debug', '', event.POLL, {'verbose': 1})
        )

    signal.signal(signal.SIGUSR1, debug_workers)

    # Prepopulate the workers with existing routers on startup
    populate.pre_populate_workers(sched)

    # Set up the periodic health check
    health.start_inspector(cfg.CONF.health_check_period, sched)

    # Block the main process, copying messages from the notification
    # listener to the scheduler
    shuffle_notifications(notification_queue, sched)

    # Terminate the listening process
    notification_proc.terminate()
    metadata_proc.terminate()

    LOG.info('exiting')
<commit_msg>Parse CLI options after registering them<commit_after>import functools
import logging
import multiprocessing
import signal
import socket
import sys

from oslo.config import cfg

from akanda.rug import event
from akanda.rug import health
from akanda.rug.openstack.common import log
from akanda.rug import metadata
from akanda.rug import notifications
from akanda.rug import scheduler
from akanda.rug import populate
from akanda.rug import worker
from akanda.rug.api import quantum as quantum_api

LOG = log.getLogger(__name__)

def shuffle_notifications(notification_queue, sched):
    """Copy messages from the notification queue into the scheduler.
    """
    while True:
        try:
            target, message = notification_queue.get()
            sched.handle_message(target, message)
        # FIXME(rods): if a signal arrive during an IO operation an
        # IOError is raised. We catch the exceptions in meantime
        # waiting for a better solution.
        except IOError:
            pass
        except KeyboardInterrupt:
            sched.stop()
            break

def register_and_load_opts(argv):

    # Set the logging format to include the process and thread, since
    # those aren't included in standard openstack logs but are useful
    # for the rug
    log_format = ':'.join('%(' + n + ')s'
                          for n in ['asctime',
                                    'levelname',
                                    'name',
                                    'processName',
                                    'threadName',
                                    'message'])
    log.set_defaults(log_format)

    # Configure the default log levels for some third-party packages
    # that are chatty
    cfg.set_defaults(log.log_opts,
                     default_log_levels=['amqplib=WARN',
                                         'qpid.messaging=INFO',
                                         'sqlalchemy=WARN',
                                         'keystoneclient=INFO',
                                         'stevedore=INFO',
                                         'eventlet.wsgi.server=WARN',
                                         'requests=WARN',
                                         ])

    cfg.CONF.register_opts([
        cfg.StrOpt('host',
                   default=socket.getfqdn(),
                   help="The hostname Akanda is running on"),

        # FIXME(dhellmann): Use a separate group for these auth params
        cfg.StrOpt('admin_user'),
        cfg.StrOpt('admin_password', secret=True),
        cfg.StrOpt('admin_tenant_name'),
        cfg.StrOpt('auth_url'),
        cfg.StrOpt('auth_strategy', default='keystone'),
        cfg.StrOpt('auth_region'),

        cfg.StrOpt('management_network_id'),
        cfg.StrOpt('external_network_id'),
        cfg.StrOpt('management_subnet_id'),
        cfg.StrOpt('router_image_uuid'),

        cfg.StrOpt('management_prefix', default='fdca:3ba5:a17a:acda::/64'),
        cfg.IntOpt('akanda_mgt_service_port', default=5000),
        cfg.IntOpt('router_instance_flavor', default=1),

        # needed for plugging locally into management network
        cfg.StrOpt('interface_driver'),
        cfg.StrOpt('ovs_integration_bridge', default='br-int'),
        cfg.BoolOpt('ovs_use_veth', default=False),
        cfg.IntOpt('network_device_mtu'),

        # needed for boot waiting
        cfg.IntOpt('boot_timeout', default=240),
        cfg.IntOpt('max_retries', default=3),
        cfg.IntOpt('retry_delay', default=1),

        cfg.StrOpt(
            'ignored_router_directory',
            default='/etc/akanda-rug/ignored',
            help='Directory to scan for routers to ignore for debugging',
        ),

    ])

    cfg.CONF.register_opts(metadata.metadata_opts)

    AGENT_OPTIONS = [
        cfg.StrOpt('root_helper', default='sudo'),
    ]

    cfg.CONF.register_opts(AGENT_OPTIONS, 'AGENT')

    # FIXME: Convert these to regular options, not command line options.
    cfg.CONF.register_cli_opts([
        cfg.IntOpt('health-check-period',
                   default=60,
                   help='seconds between health checks'),
        cfg.IntOpt('num-worker-processes',
                   short='w',
                   default=16,
                   help='the number of worker processes to run'),
        cfg.IntOpt('num-worker-threads',
                   short='t',
                   default=4,
                   help='the number of worker threads to run per process'),

        # FIXME(dhellmann): set up a group for these messaging params
        cfg.StrOpt('amqp-url',
                   default='amqp://guest:secrete@localhost:5672/',
                   help='connection for AMQP server'),
        cfg.StrOpt('incoming-notifications-exchange',
                   default='neutron',
                   help='name of the exchange where we receive notifications'),
        cfg.StrOpt('outgoing-notifications-exchange',
                   default='neutron',
                   help='name of the exchange where we send notifications'),
        cfg.StrOpt('rpc-exchange',
                   default='l3_agent_fanout',
                   help='name of the exchange where we receive RPC calls'),
    ])

def main(argv=sys.argv[1:]):
    register_and_load_opts(argv)
    cfg.CONF(argv, project='akanda')

    log.setup('akanda-rug')
    cfg.CONF.log_opt_values(LOG, logging.INFO)

    # Purge the mgt tap interface on startup
    quantum = quantum_api.Quantum(cfg.CONF)
    #TODO(mark): develop better way restore after machine reboot
    #quantum.purge_management_interface()

    # bring the mgt tap interface up
    quantum.ensure_local_service_port()

    # Set up the queue to move messages between the eventlet-based
    # listening process and the scheduler.
    notification_queue = multiprocessing.Queue()

    # Listen for notifications.
    notification_proc = multiprocessing.Process(
        target=notifications.listen,
        kwargs={
            'host_id': cfg.CONF.host,
            'amqp_url': cfg.CONF.amqp_url,
            'notifications_exchange_name':
            cfg.CONF.incoming_notifications_exchange,
            'rpc_exchange_name': cfg.CONF.rpc_exchange,
            'notification_queue': notification_queue,
        },
        name='notification-listener',
    )
    notification_proc.start()

    mgt_ip_address = quantum_api.get_local_service_ip(cfg.CONF).split('/')[0]
    metadata_proc = multiprocessing.Process(
        target=metadata.serve,
        args=(mgt_ip_address,),
        name='metadata-proxy'
    )
    metadata_proc.start()

    # Set up the notifications publisher
    publisher = notifications.Publisher(
        cfg.CONF.amqp_url,
        exchange_name=cfg.CONF.outgoing_notifications_exchange,
        topic='notifications.info',
    )

    # Set up a factory to make Workers that know how many threads to
    # run.
    worker_factory = functools.partial(
        worker.Worker,
        num_threads=cfg.CONF.num_worker_threads,
        notifier=publisher,
        ignore_directory=cfg.CONF.ignored_router_directory,
    )

    # Set up the scheduler that knows how to manage the routers and
    # dispatch messages.
    sched = scheduler.Scheduler(
        num_workers=cfg.CONF.num_worker_processes,
        worker_factory=worker_factory,
    )

    def debug_workers(signum, stack):
        sched.handle_message(
            'debug',
            event.Event('debug', '', event.POLL, {'verbose': 1})
        )

    signal.signal(signal.SIGUSR1, debug_workers)

    # Prepopulate the workers with existing routers on startup
    populate.pre_populate_workers(sched)

    # Set up the periodic health check
    health.start_inspector(cfg.CONF.health_check_period, sched)

    # Block the main process, copying messages from the notification
    # listener to the scheduler
    shuffle_notifications(notification_queue, sched)

    # Terminate the listening process
    notification_proc.terminate()
    metadata_proc.terminate()

    LOG.info('exiting')
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

# Define your item pipelines here
#
# Don't forget to add your pipeline to the ITEM_PIPELINES setting
# See: http://doc.scrapy.org/en/latest/topics/item-pipeline.html

#python库
import sqlite3
import re
from scrapy import log
from urlparse import urlparse

from twisted.internet.threads import deferToThread
from scrapy.utils.serialize import ScrapyJSONEncoder

#自定义的库
from WebModel.items import PageItem, RulesetItem
# 数据库操作
from WebModel.database.databasehelper import getCliInstance
# [bloomfilter库,用以查重](https://github.com/jaybaird/python-bloomfilter/)
from pybloom import BloomFilter
# [域名解析库](https://pypi.python.org/pypi/publicsuffix/)
from WebModel.utils.publicsuffix import domain_getter, TYPE_DOMAIN, TYPE_IP
# scrapy-redis 的pipeline
import WebModel.utils.scrapy_redis.connection as connection
from WebModel.utils.rediskeys import domains_key, related_domains_key

class RedisPipeline(object):
	"""Pushes serialized item into a redis list/queue"""

	def __init__(self, server):
		# redis server
		self.server = server
		# 用来判断是否重复出现
		allowed = [
			"qq.com",
			"163.com",
			"people.com.cn",
			"xinhuanet.com",
			"cntv.cn",
			"ifeng.com",
			"hexun.com",
			"sina.com.cn",
			"sohu.com",
			"dbw.cn",
		]
		self.bloom_domain_filter = BloomFilter(capacity=32)
		for a in allowed:
			self.bloom_domain_filter.add(a)

		# 正则过滤, 一些博客
		self.qzone_filter = re.compile(r"^http://.*\.qzone\.qq\.com")
		self.wangyiblog_filter = re.compile(r"^http://.*\.blog\.163\.com")
		self.hexunblog_filter = re.compile(r"^http://.*\.blog\.hexun\.com")
		self.sohublog_filter = re.compile(r"http://.*\.blog\.sohu\.com")
		self.sohui_filter = re.compile(r"http://.*\.i\.sohu\.com")

		self.bloom_domain_vec = BloomFilter(capacity=1<<16, error_rate=0.001)
		self.bloom_netloc_vec = BloomFilter(capacity=1<<16, error_rate=0.001)

	@classmethod
	def from_crawler(cls, crawler):
		settings = crawler.settings
		server = connection.from_settings(settings)
		pipe = cls(server)
		pipe.settings = settings
		pipe.crawler = crawler
		return pipe

	def process_item(self, item, spider):
		return deferToThread(self._process_item, item, spider)

	def _process_item(self, item, spider):
		if isinstance(item, RulesetItem):
			self._process_ruleset(item, spider)
		elif isinstance(item, PageItem):
			self._process_page(item, spider)
		return item

	def _process_ruleset(self, item, spider):
		if item['update']:
			if item['domain'] in self.bloom_domain_vec:
				spider.log("BloomFilter 判断错误!域名%s未出现过"%item['domain'], level=log.WARNING)
			# 新域名,建立域名记录
			#redis数据库
			self.server.hset(domains_key%item['domain'], 'indegree', 1)
			self.server.hset(domains_key%item['domain'], 'outdegree', 0)
			#sqlite3数据库
			# self.dbcli.insertDomain(item['domain'])

			spider.log("爬取到新域名: "+item['domain'], level=log.INFO)

			# 把robots.txt内容存储进数据库
			if item['ruleset']:
				#redis数据库
				self.server.hset(domains_key%item['domain'], 'ruleset', item['ruleset'])
				#sqlite3数据库
				# self.dbcli.insertRuleset(item['ruleset'], item['domain'])

				spider.log("爬取到域名的robots规则: "+item['domain'], level=log.INFO)

	def _process_page(self, item, spider):
		# 记录爬取过此网址
		self.server.rpush(spider.URL_VISITED_KEY, item['url'])

		domain, ret_type = domain_getter.get_domain(item['url'])
		if not self.bloom_domain_vec.add(domain):
			self._initRelateDomain(domain, spider)
			if not self.server.exists(domains_key%domain):
				self._initGlobalDomain(domain)

		for link in item['links']:

			if self._isBlogLink(link):
				self.server.rpush(spider.BLOG_IGNORE_KEY, link)
				continue

			link_domain, ret_type = domain_getter.get_domain(link)
			# 判断域名是否和当前爬取域名D0相同
			if link_domain == spider.CRAWLING_DOMAIN:
				# D0的size+1
				self.server.hincrby(domains_key%link_domain, 'size', 1)
				self.server.hincrby(related_domains_key%(spider.CRAWLING_DOMAIN,domain), 'size', 1)
				# 判断其netloc是否出现过
				netloc = urlparse(link).netloc
				if not self.bloom_netloc_vec.add(netloc):
					# 未出现过,网页进入队列
					self.server.rpush(spider.URL_QUEUE_KEY, "http://"+netloc)
			else:# 不同,则判断域名是否已经出现过
				if self.bloom_domain_vec.add(link_domain):
					# 已出现过,对应的记录D1入度+1
					self.server.hincrby(domains_key%link_domain, 'indegree', 1)
					self.server.hincrby(related_domains_key%(spider.CRAWLING_DOMAIN,link_domain), 'indegree', 1)
				else:
					# 没出现过,则Domain中增加记录D1,D1记录入度初始化为1,出度初始化为0;D0出度+1
					self._initGlobalDomain(link_domain)
					self.server.hincrby(domains_key%domain, 'outdegree', 1)

					self._initRelateDomain(link_domain, spider)
					self.server.hincrby(related_domains_key%(spider.CRAWLING_DOMAIN,link_domain), 'outdegree', 1)

		for link in item['refused_links']:
			spider.log(link+" refuse by robots.txt", level=log.INFO)
			self.server.rpush(spider.ROBOT_REFUSED_KEY, link)

	def _initGlobalDomain(self, domain):
		key = domains_key%domain
		self._raw_init_domain(key)

	def _initRelateDomain(self, domain, spider):
		key = related_domains_key%(spider.CRAWLING_DOMAIN, domain)
		self._raw_init_domain(key)

	def _raw_init_domain(self, key):
		self.server.hset(key, 'indegree', 0)
		self.server.hset(key, 'outdegree', 0)
		self.server.hset(key, 'size', 1)

	def _isBlogLink(self,link):
		# qq 空间个性域名,过滤
		if self.qzone_filter.match(link):
			return True

		# 网易博客个性域名,过滤
		elif self.wangyiblog_filter.match(link):
			return True

		# 和讯blog,过滤
		elif self.hexunblog_filter.match(link):
			return True

		# 搜狐blog、个人展示页,过滤
		elif self.sohublog_filter.match(link) or self.sohui_filter.match(link):
			return True

		return False
<commit_msg>fix bug<commit_after># -*- coding: utf-8 -*-

# Define your item pipelines here
#
# Don't forget to add your pipeline to the ITEM_PIPELINES setting
# See: http://doc.scrapy.org/en/latest/topics/item-pipeline.html

#python库
import sqlite3
import re
from scrapy import log
from urlparse import urlparse

from twisted.internet.threads import deferToThread
from scrapy.utils.serialize import ScrapyJSONEncoder

#自定义的库
from WebModel.items import PageItem, RulesetItem
# 数据库操作
from WebModel.database.databasehelper import getCliInstance
# [bloomfilter库,用以查重](https://github.com/jaybaird/python-bloomfilter/)
from pybloom import BloomFilter
# [域名解析库](https://pypi.python.org/pypi/publicsuffix/)
from WebModel.utils.publicsuffix import domain_getter, TYPE_DOMAIN, TYPE_IP
# scrapy-redis 的pipeline
import WebModel.utils.scrapy_redis.connection as connection
from WebModel.utils.rediskeys import domains_key, related_domains_key

class RedisPipeline(object):
	"""Pushes serialized item into a redis list/queue"""

	def __init__(self, server):
		# redis server
		self.server = server
		# 用来判断是否重复出现
		allowed = [
			"qq.com",
			"163.com",
			"people.com.cn",
			"xinhuanet.com",
			"cntv.cn",
			"ifeng.com",
			"hexun.com",
			"sina.com.cn",
			"sohu.com",
			"dbw.cn",
		]
		self.bloom_domain_filter = BloomFilter(capacity=32)
		for a in allowed:
			self.bloom_domain_filter.add(a)

		# 正则过滤, 一些博客
		self.qzone_filter = re.compile(r"^http://.*\.qzone\.qq\.com")
		self.wangyiblog_filter = re.compile(r"^http://.*\.blog\.163\.com")
		self.hexunblog_filter = re.compile(r"^http://.*\.blog\.hexun\.com")
		self.sohublog_filter = re.compile(r"http://.*\.blog\.sohu\.com")
		self.sohui_filter = re.compile(r"http://.*\.i\.sohu\.com")

		self.bloom_domain_vec = BloomFilter(capacity=1<<16, error_rate=0.001)
		self.bloom_netloc_vec = BloomFilter(capacity=1<<16, error_rate=0.001)

	@classmethod
	def from_crawler(cls, crawler):
		settings = crawler.settings
		server = connection.from_settings(settings)
		pipe = cls(server)
		pipe.settings = settings
		pipe.crawler = crawler
		return pipe

	def process_item(self, item, spider):
		return deferToThread(self._process_item, item, spider)

	def _process_item(self, item, spider):
		if isinstance(item, RulesetItem):
			self._process_ruleset(item, spider)
		elif isinstance(item, PageItem):
			self._process_page(item, spider)
		return item

	def _process_ruleset(self, item, spider):
		if item['update']:
			if item['domain'] in self.bloom_domain_vec:
				spider.log("BloomFilter 判断错误!域名%s未出现过"%item['domain'], level=log.WARNING)
			# 新域名,建立域名记录
			#redis数据库
			self.server.hset(domains_key%item['domain'], 'indegree', 1)
			self.server.hset(domains_key%item['domain'], 'outdegree', 0)
			#sqlite3数据库
			# self.dbcli.insertDomain(item['domain'])

			spider.log("爬取到新域名: "+item['domain'], level=log.INFO)

			# 把robots.txt内容存储进数据库
			if item['ruleset']:
				#redis数据库
				self.server.hset(domains_key%item['domain'], 'ruleset', item['ruleset'])
				#sqlite3数据库
				# self.dbcli.insertRuleset(item['ruleset'], item['domain'])

				spider.log("爬取到域名的robots规则: "+item['domain'], level=log.INFO)

	def _process_page(self, item, spider):
		# 记录爬取过此网址
		self.server.rpush(spider.URL_VISITED_KEY, item['url'])

		domain, ret_type = domain_getter.get_domain(item['url'])
		if not self.bloom_domain_vec.add(domain):
			self._initRelateDomain(domain, spider)
			if not self.server.exists(domains_key%domain):
				self._initGlobalDomain(domain)

		for link in item['links']:

			if self._isBlogLink(link):
				self.server.rpush(spider.BLOG_IGNORE_KEY, link)
				continue

			link_domain, ret_type = domain_getter.get_domain(link)
			# 判断域名是否和当前爬取域名D0相同
			if link_domain == spider.CRAWLING_DOMAIN:
				# D0的size+1
				self.server.hincrby(domains_key%link_domain, 'size', 1)
				self.server.hincrby(related_domains_key%(spider.CRAWLING_DOMAIN,domain), 'size', 1)
				# 判断其netloc是否出现过
				netloc = urlparse(link).netloc
				if not self.bloom_netloc_vec.add(netloc):
					# 未出现过,网页进入队列
					self.server.rpush(spider.URL_QUEUE_KEY, "http://"+netloc)
			else:# 不同,则判断域名是否已经出现过
				if self.bloom_domain_vec.add(link_domain):
					# 已出现过,对应的记录D1入度+1
					self.server.hincrby(domains_key%link_domain, 'indegree', 1)
					self.server.hincrby(related_domains_key%(spider.CRAWLING_DOMAIN,link_domain), 'indegree', 1)
				else:
					# 没出现过,则Domain中增加记录D1,D1记录入度初始化为1,出度初始化为0;D0出度+1
					if not self.server.exists(domains_key%domain):
						self._initGlobalDomain(link_domain)
					self.server.hincrby(domains_key%domain, 'outdegree', 1)

					self._initRelateDomain(link_domain, spider)
					self.server.hincrby(related_domains_key%(spider.CRAWLING_DOMAIN,link_domain), 'outdegree', 1)

		for link in item['refused_links']:
			spider.log(link+" refuse by robots.txt", level=log.INFO)
			self.server.rpush(spider.ROBOT_REFUSED_KEY, link)

	def _initGlobalDomain(self, domain):
		key = domains_key%domain
		self._raw_init_domain(key)

	def _initRelateDomain(self, domain, spider):
		key = related_domains_key%(spider.CRAWLING_DOMAIN, domain)
		self._raw_init_domain(key)

	def _raw_init_domain(self, key):
		self.server.hset(key, 'indegree', 0)
		self.server.hset(key, 'outdegree', 0)
		self.server.hset(key, 'size', 1)

	def _isBlogLink(self,link):
		# qq 空间个性域名,过滤
		if self.qzone_filter.match(link):
			return True

		# 网易博客个性域名,过滤
		elif self.wangyiblog_filter.match(link):
			return True

		# 和讯blog,过滤
		elif self.hexunblog_filter.match(link):
			return True

		# 搜狐blog、个人展示页,过滤
		elif self.sohublog_filter.match(link) or self.sohui_filter.match(link):
			return True

		return False
<|endoftext|>


<commit_before>from __future__ import print_function
import argparse
import os
import random
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.optim as optim
import torch.utils.data
import torchvision.datasets as dset
import torchvision.transforms as transforms
import torchvision.utils as vutils
from torch.autograd import Variable

parser = argparse.ArgumentParser()
parser.add_argument('--dataset', required=True, help='cifar10 | lsun | imagenet | folder | lfw ')
parser.add_argument('--dataroot', required=True, help='path to dataset')
parser.add_argument('--workers', type=int, help='number of data loading workers', default=2)
parser.add_argument('--batchSize', type=int, default=64, help='input batch size')
parser.add_argument('--imageSize', type=int, default=64, help='the height / width of the input image to network')
parser.add_argument('--nz', type=int, default=100, help='size of the latent z vector')
parser.add_argument('--ngf', type=int, default=64)
parser.add_argument('--ndf', type=int, default=64)
parser.add_argument('--niter', type=int, default=25, help='number of epochs to train for')
parser.add_argument('--lr', type=float, default=0.0002, help='learning rate, default=0.0002')
parser.add_argument('--beta1', type=float, default=0.5, help='beta1 for adam. default=0.5')
parser.add_argument('--cuda'  , action='store_true', help='enables cuda')
parser.add_argument('--ngpu'  , type=int, default=1, help='number of GPUs to use')
parser.add_argument('--netG', default='', help="path to netG (to continue training)")
parser.add_argument('--netD', default='', help="path to netD (to continue training)")
parser.add_argument('--outf', default='.', help='folder to output images and model checkpoints')

opt = parser.parse_args()
print(opt)

try:
    os.makedirs(opt.outf)
except OSError:
    pass
opt.manualSeed = random.randint(1, 10000) # fix seed
print("Random Seed: ", opt.manualSeed)
random.seed(opt.manualSeed)
torch.manual_seed(opt.manualSeed)

cudnn.benchmark = True

if torch.cuda.is_available() and not opt.cuda:
    print("WARNING: You have a CUDA device, so you should probably run with --cuda")

if opt.dataset in ['imagenet', 'folder', 'lfw']:
    # folder dataset
    dataset = dset.ImageFolder(root=opt.dataroot,
                               transform=transforms.Compose([
                                   transforms.Scale(opt.imageSize),
                                   transforms.CenterCrop(opt.imageSize),
                                   transforms.ToTensor(),
                                   transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
                               ]))
elif opt.dataset == 'lsun':
    dataset = dset.LSUN(db_path=opt.dataroot, classes=['bedroom_train'],
                        transform=transforms.Compose([
                            transforms.Scale(opt.imageSize),
                            transforms.CenterCrop(opt.imageSize),
                            transforms.ToTensor(),
                            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
                        ]))
elif opt.dataset == 'cifar10':
    dataset = dset.CIFAR10(root=opt.dataroot, download=True,
                           transform=transforms.Compose([
                               transforms.Scale(opt.imageSize),
                               transforms.ToTensor(),
                               transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
                           ])
    )
assert dataset
dataloader = torch.utils.data.DataLoader(dataset, batch_size=opt.batchSize,
                                         shuffle=True, num_workers=int(opt.workers))

ngpu = int(opt.ngpu)
nz = int(opt.nz)
ngf = int(opt.ngf)
ndf = int(opt.ndf)
nc = 3

# custom weights initialization called on netG and netD
def weights_init(m):
    classname = m.__class__.__name__
    if classname.find('Conv') != -1:
        m.weight.data.normal_(0.0, 0.02)
    elif classname.find('BatchNorm') != -1:
        m.weight.data.normal_(1.0, 0.02)
        m.bias.data.fill_(0)

class _netG(nn.Module):
    def __init__(self, ngpu):
        super(_netG, self).__init__()
        self.ngpu = ngpu
        self.main = nn.Sequential(
            # input is Z, going into a convolution
            nn.ConvTranspose2d(     nz, ngf * 8, 4, 1, 0, bias=False),
            nn.BatchNorm2d(ngf * 8),
            nn.ReLU(True),
            # state size. (ngf*8) x 4 x 4
            nn.ConvTranspose2d(ngf * 8, ngf * 4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf * 4),
            nn.ReLU(True),
            # state size. (ngf*4) x 8 x 8
            nn.ConvTranspose2d(ngf * 4, ngf * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf * 2),
            nn.ReLU(True),
            # state size. (ngf*2) x 16 x 16
            nn.ConvTranspose2d(ngf * 2,     ngf, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf),
            nn.ReLU(True),
            # state size. (ngf) x 32 x 32
            nn.ConvTranspose2d(    ngf,      nc, 4, 2, 1, bias=False),
            nn.Tanh()
            # state size. (nc) x 64 x 64
        )
    def forward(self, input):
        gpu_ids = None
        if isinstance(input.data, torch.cuda.FloatTensor) and self.ngpu > 1:
            gpu_ids = range(self.ngpu)
        return nn.parallel.data_parallel(self.main, input, gpu_ids)

netG = _netG(ngpu)
netG.apply(weights_init)
if opt.netG != '':
    netG.load_state_dict(torch.load(opt.netG))
print(netG)

class _netD(nn.Module):
    def __init__(self, ngpu):
        super(_netD, self).__init__()
        self.ngpu = ngpu
        self.main = nn.Sequential(
            # input is (nc) x 64 x 64
            nn.Conv2d(nc, ndf, 4, 2, 1, bias=False),
            nn.LeakyReLU(0.2, inplace=True),
            # state size. (ndf) x 32 x 32
            nn.Conv2d(ndf, ndf * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ndf * 2),
            nn.LeakyReLU(0.2, inplace=True),
            # state size. (ndf*2) x 16 x 16
            nn.Conv2d(ndf * 2, ndf * 4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ndf * 4),
            nn.LeakyReLU(0.2, inplace=True),
            # state size. (ndf*4) x 8 x 8
            nn.Conv2d(ndf * 4, ndf * 8, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ndf * 8),
            nn.LeakyReLU(0.2, inplace=True),
            # state size. (ndf*8) x 4 x 4
            nn.Conv2d(ndf * 8, 1, 4, 1, 0, bias=False),
            nn.Sigmoid()
        )
    def forward(self, input):
        gpu_ids = None
        if isinstance(input.data, torch.cuda.FloatTensor) and self.ngpu > 1:
            gpu_ids = range(self.ngpu)
        output = nn.parallel.data_parallel(self.main, input, gpu_ids)
        return output.view(-1, 1)

netD = _netD(ngpu)
netD.apply(weights_init)
if opt.netD != '':
    netD.load_state_dict(torch.load(opt.netD))
print(netD)

criterion = nn.BCELoss()

input = torch.FloatTensor(opt.batchSize, 3, opt.imageSize, opt.imageSize)
noise = torch.FloatTensor(opt.batchSize, nz, 1, 1)
fixed_noise = torch.FloatTensor(opt.batchSize, nz, 1, 1).normal_(0, 1)
label = torch.FloatTensor(opt.batchSize)
real_label = 1
fake_label = 0

if opt.cuda:
    netD.cuda()
    netG.cuda()
    criterion.cuda()
    input, label = input.cuda(), label.cuda()
    noise, fixed_noise = noise.cuda(), fixed_noise.cuda()

input = Variable(input)
label = Variable(label)
noise = Variable(noise)
fixed_noise = Variable(fixed_noise)

# setup optimizer
optimizerD = optim.Adam(netD.parameters(), lr = opt.lr, betas = (opt.beta1, 0.999))
optimizerG = optim.Adam(netG.parameters(), lr = opt.lr, betas = (opt.beta1, 0.999))

for epoch in range(opt.niter):
    for i, data in enumerate(dataloader, 0):
        ############################
        # (1) Update D network: maximize log(D(x)) + log(1 - D(G(z)))
        ###########################
        # train with real
        netD.zero_grad()
        real_cpu, _ = data
        batch_size = real_cpu.size(0)
        input.data.resize_(real_cpu.size()).copy_(real_cpu)
        label.data.resize_(batch_size).fill_(real_label)

        output = netD(input)
        errD_real = criterion(output, label)
        errD_real.backward()
        D_x = output.data.mean()

        # train with fake
        noise.data.resize_(batch_size, nz, 1, 1)
        noise.data.normal_(0, 1)
        fake = netG(noise)
        label.data.fill_(fake_label)
        output = netD(fake.detach())
        errD_fake = criterion(output, label)
        errD_fake.backward()
        D_G_z1 = output.data.mean()
        errD = errD_real + errD_fake
        optimizerD.step()

        ############################
        # (2) Update G network: maximize log(D(G(z)))
        ###########################
        netG.zero_grad()
        label.data.fill_(real_label) # fake labels are real for generator cost
        output = netD(fake)
        errG = criterion(output, label)
        errG.backward()
        D_G_z2 = output.data.mean()
        optimizerG.step()

        print('[%d/%d][%d/%d] Loss_D: %.4f Loss_G: %.4f D(x): %.4f D(G(z)): %.4f / %.4f'
              % (epoch, opt.niter, i, len(dataloader),
                 errD.data[0], errG.data[0], D_x, D_G_z1, D_G_z2))
        if i % 100 == 0:
            vutils.save_image(real_cpu,
                    '%s/real_samples.png' % opt.outf)
            fake = netG(fixed_noise)
            vutils.save_image(fake.data,
                    '%s/fake_samples_epoch_%03d.png' % (opt.outf, epoch))

    # do checkpointing
    torch.save(netG.state_dict(), '%s/netG_epoch_%d.pth' % (opt.outf, epoch))
    torch.save(netD.state_dict(), '%s/netD_epoch_%d.pth' % (opt.outf, epoch))
<commit_msg>Fix random seeding in DCGAN (#108)<commit_after>from __future__ import print_function
import argparse
import os
import random
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.optim as optim
import torch.utils.data
import torchvision.datasets as dset
import torchvision.transforms as transforms
import torchvision.utils as vutils
from torch.autograd import Variable

parser = argparse.ArgumentParser()
parser.add_argument('--dataset', required=True, help='cifar10 | lsun | imagenet | folder | lfw ')
parser.add_argument('--dataroot', required=True, help='path to dataset')
parser.add_argument('--workers', type=int, help='number of data loading workers', default=2)
parser.add_argument('--batchSize', type=int, default=64, help='input batch size')
parser.add_argument('--imageSize', type=int, default=64, help='the height / width of the input image to network')
parser.add_argument('--nz', type=int, default=100, help='size of the latent z vector')
parser.add_argument('--ngf', type=int, default=64)
parser.add_argument('--ndf', type=int, default=64)
parser.add_argument('--niter', type=int, default=25, help='number of epochs to train for')
parser.add_argument('--lr', type=float, default=0.0002, help='learning rate, default=0.0002')
parser.add_argument('--beta1', type=float, default=0.5, help='beta1 for adam. default=0.5')
parser.add_argument('--cuda', action='store_true', help='enables cuda')
parser.add_argument('--ngpu', type=int, default=1, help='number of GPUs to use')
parser.add_argument('--netG', default='', help="path to netG (to continue training)")
parser.add_argument('--netD', default='', help="path to netD (to continue training)")
parser.add_argument('--outf', default='.', help='folder to output images and model checkpoints')
parser.add_argument('--manualSeed', type=int, help='manual seed')

opt = parser.parse_args()
print(opt)

try:
    os.makedirs(opt.outf)
except OSError:
    pass

if opt.manualSeed is None:
    opt.manualSeed = random.randint(1, 10000)
print("Random Seed: ", opt.manualSeed)
random.seed(opt.manualSeed)
torch.manual_seed(opt.manualSeed)
if opt.cuda:
    torch.cuda.manual_seed_all(opt.manualSeed)

cudnn.benchmark = True

if torch.cuda.is_available() and not opt.cuda:
    print("WARNING: You have a CUDA device, so you should probably run with --cuda")

if opt.dataset in ['imagenet', 'folder', 'lfw']:
    # folder dataset
    dataset = dset.ImageFolder(root=opt.dataroot,
                               transform=transforms.Compose([
                                   transforms.Scale(opt.imageSize),
                                   transforms.CenterCrop(opt.imageSize),
                                   transforms.ToTensor(),
                                   transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
                               ]))
elif opt.dataset == 'lsun':
    dataset = dset.LSUN(db_path=opt.dataroot, classes=['bedroom_train'],
                        transform=transforms.Compose([
                            transforms.Scale(opt.imageSize),
                            transforms.CenterCrop(opt.imageSize),
                            transforms.ToTensor(),
                            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
                        ]))
elif opt.dataset == 'cifar10':
    dataset = dset.CIFAR10(root=opt.dataroot, download=True,
                           transform=transforms.Compose([
                               transforms.Scale(opt.imageSize),
                               transforms.ToTensor(),
                               transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
                           ])
    )
assert dataset
dataloader = torch.utils.data.DataLoader(dataset, batch_size=opt.batchSize,
                                         shuffle=True, num_workers=int(opt.workers))

ngpu = int(opt.ngpu)
nz = int(opt.nz)
ngf = int(opt.ngf)
ndf = int(opt.ndf)
nc = 3

# custom weights initialization called on netG and netD
def weights_init(m):
    classname = m.__class__.__name__
    if classname.find('Conv') != -1:
        m.weight.data.normal_(0.0, 0.02)
    elif classname.find('BatchNorm') != -1:
        m.weight.data.normal_(1.0, 0.02)
        m.bias.data.fill_(0)

class _netG(nn.Module):
    def __init__(self, ngpu):
        super(_netG, self).__init__()
        self.ngpu = ngpu
        self.main = nn.Sequential(
            # input is Z, going into a convolution
            nn.ConvTranspose2d(     nz, ngf * 8, 4, 1, 0, bias=False),
            nn.BatchNorm2d(ngf * 8),
            nn.ReLU(True),
            # state size. (ngf*8) x 4 x 4
            nn.ConvTranspose2d(ngf * 8, ngf * 4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf * 4),
            nn.ReLU(True),
            # state size. (ngf*4) x 8 x 8
            nn.ConvTranspose2d(ngf * 4, ngf * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf * 2),
            nn.ReLU(True),
            # state size. (ngf*2) x 16 x 16
            nn.ConvTranspose2d(ngf * 2,     ngf, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf),
            nn.ReLU(True),
            # state size. (ngf) x 32 x 32
            nn.ConvTranspose2d(    ngf,      nc, 4, 2, 1, bias=False),
            nn.Tanh()
            # state size. (nc) x 64 x 64
        )

    def forward(self, input):
        gpu_ids = None
        if isinstance(input.data, torch.cuda.FloatTensor) and self.ngpu > 1:
            gpu_ids = range(self.ngpu)
        return nn.parallel.data_parallel(self.main, input, gpu_ids)

netG = _netG(ngpu)
netG.apply(weights_init)
if opt.netG != '':
    netG.load_state_dict(torch.load(opt.netG))
print(netG)

class _netD(nn.Module):
    def __init__(self, ngpu):
        super(_netD, self).__init__()
        self.ngpu = ngpu
        self.main = nn.Sequential(
            # input is (nc) x 64 x 64
            nn.Conv2d(nc, ndf, 4, 2, 1, bias=False),
            nn.LeakyReLU(0.2, inplace=True),
            # state size. (ndf) x 32 x 32
            nn.Conv2d(ndf, ndf * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ndf * 2),
            nn.LeakyReLU(0.2, inplace=True),
            # state size. (ndf*2) x 16 x 16
            nn.Conv2d(ndf * 2, ndf * 4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ndf * 4),
            nn.LeakyReLU(0.2, inplace=True),
            # state size. (ndf*4) x 8 x 8
            nn.Conv2d(ndf * 4, ndf * 8, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ndf * 8),
            nn.LeakyReLU(0.2, inplace=True),
            # state size. (ndf*8) x 4 x 4
            nn.Conv2d(ndf * 8, 1, 4, 1, 0, bias=False),
            nn.Sigmoid()
        )

    def forward(self, input):
        gpu_ids = None
        if isinstance(input.data, torch.cuda.FloatTensor) and self.ngpu > 1:
            gpu_ids = range(self.ngpu)
        output = nn.parallel.data_parallel(self.main, input, gpu_ids)
        return output.view(-1, 1)

netD = _netD(ngpu)
netD.apply(weights_init)
if opt.netD != '':
    netD.load_state_dict(torch.load(opt.netD))
print(netD)

criterion = nn.BCELoss()

input = torch.FloatTensor(opt.batchSize, 3, opt.imageSize, opt.imageSize)
noise = torch.FloatTensor(opt.batchSize, nz, 1, 1)
fixed_noise = torch.FloatTensor(opt.batchSize, nz, 1, 1).normal_(0, 1)
label = torch.FloatTensor(opt.batchSize)
real_label = 1
fake_label = 0

if opt.cuda:
    netD.cuda()
    netG.cuda()
    criterion.cuda()
    input, label = input.cuda(), label.cuda()
    noise, fixed_noise = noise.cuda(), fixed_noise.cuda()

input = Variable(input)
label = Variable(label)
noise = Variable(noise)
fixed_noise = Variable(fixed_noise)

# setup optimizer
optimizerD = optim.Adam(netD.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))
optimizerG = optim.Adam(netG.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))

for epoch in range(opt.niter):
    for i, data in enumerate(dataloader, 0):
        ############################
        # (1) Update D network: maximize log(D(x)) + log(1 - D(G(z)))
        ###########################
        # train with real
        netD.zero_grad()
        real_cpu, _ = data
        batch_size = real_cpu.size(0)
        input.data.resize_(real_cpu.size()).copy_(real_cpu)
        label.data.resize_(batch_size).fill_(real_label)

        output = netD(input)
        errD_real = criterion(output, label)
        errD_real.backward()
        D_x = output.data.mean()

        # train with fake
        noise.data.resize_(batch_size, nz, 1, 1)
        noise.data.normal_(0, 1)
        fake = netG(noise)
        label.data.fill_(fake_label)
        output = netD(fake.detach())
        errD_fake = criterion(output, label)
        errD_fake.backward()
        D_G_z1 = output.data.mean()
        errD = errD_real + errD_fake
        optimizerD.step()

        ############################
        # (2) Update G network: maximize log(D(G(z)))
        ###########################
        netG.zero_grad()
        label.data.fill_(real_label)  # fake labels are real for generator cost
        output = netD(fake)
        errG = criterion(output, label)
        errG.backward()
        D_G_z2 = output.data.mean()
        optimizerG.step()

        print('[%d/%d][%d/%d] Loss_D: %.4f Loss_G: %.4f D(x): %.4f D(G(z)): %.4f / %.4f'
              % (epoch, opt.niter, i, len(dataloader),
                 errD.data[0], errG.data[0], D_x, D_G_z1, D_G_z2))
        if i % 100 == 0:
            vutils.save_image(real_cpu,
                    '%s/real_samples.png' % opt.outf)
            fake = netG(fixed_noise)
            vutils.save_image(fake.data,
                    '%s/fake_samples_epoch_%03d.png' % (opt.outf, epoch))

    # do checkpointing
    torch.save(netG.state_dict(), '%s/netG_epoch_%d.pth' % (opt.outf, epoch))
    torch.save(netD.state_dict(), '%s/netD_epoch_%d.pth' % (opt.outf, epoch))
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

"""
frequests
~~~~~~~~~

This module contains an asynchronous replica of ``requests.api``, powered
by futures. All API methods return a ``Request`` instance (as opposed to
``Response``). A list of requests can be sent with ``map()``.
"""

import sys

if sys.version_info[0] == 2:
    try:
        from futures import ThreadPoolExecutor
    except ImportError:
        raise RuntimeError('futures is required for frequests on python 2.X')
else:
    from concurrent.futures import ThreadPoolExecutor

try:
    from requests import api
except ImportError:
    raise RuntimeError('requests is required for frequests')

__all__ = (
    'map', 'imap',
    'get', 'options', 'head', 'post', 'put', 'patch', 'delete', 'request'
)

def patched(f):
    """Patches a given API function to not send."""

    def wrapped(*args, **kwargs):

        kwargs['return_response'] = False
        kwargs['prefetch'] = True

        config = kwargs.get('config', {})
        config.update(safe_mode=True)

        kwargs['config'] = config

        return f(*args, **kwargs)

    return wrapped

# Patched requests.api functions.
get = patched(api.get)
options = patched(api.options)
head = patched(api.head)
post = patched(api.post)
put = patched(api.put)
patch = patched(api.patch)
delete = patched(api.delete)
request = patched(api.request)

def send(r, prefetch=False):
    """Just sends the request using its send method and returns its response.  """
    r.send(prefetch=prefetch)
    return r.response

def map(requests, prefetch=True, size=1, **kwargs):
    """Concurrently converts a list of Requests to Responses.

    :param requests: a collection of Request objects.
    :param prefetch: If False, the content will not be downloaded immediately.
    :param size: Specifies the number of requests to make at a time. If 1, no throttling occurs.
    """

    requests = list(requests)

    with ThreadPoolExecutor(max_workers=size) as executor:
        responses = list(executor.map(send, requests, [prefetch]*len(requests), **kwargs))

    return responses

def imap(requests, prefetch=True, size=2, **kwargs):
    """Concurrently converts a generator object of Requests to
    a generator of Responses.

    :param requests: a generator of Request objects.
    :param prefetch: If False, the content will not be downloaded immediately.
    :param size: Specifies the number of requests to make at a time. default is 2
    """

    def prefetch_generator():
        while True:
            yield prefetch

    with ThreadPoolExecutor(max_workers=size) as executor:
        for response in executor.map(send, requests, prefetch_generator(), **kwargs):
            yield response
<commit_msg>requests 1.x compatibility changes<commit_after># -*- coding: utf-8 -*-

"""
frequests
~~~~~~~~~

This module contains an asynchronous replica of ``requests.api``, powered
by futures. All API methods return a ``Request`` instance (as opposed to
``Response``). A list of requests can be sent with ``map()``.
"""

import sys
from functools import partial

if sys.version_info[0] == 2:
    try:
        from futures import ThreadPoolExecutor
    except ImportError:
        raise RuntimeError('futures is required for frequests on python 2.X')
else:
    from concurrent.futures import ThreadPoolExecutor

try:
    from requests import Session
except ImportError:
    raise RuntimeError('requests is required for frequests')

__all__ = (
    'map', 'imap',
    'get', 'options', 'head', 'post', 'put', 'patch', 'delete', 'request'
)

class AsyncRequest(object):
    """ Asynchronous request.

    Accept same parameters as ``Session.request`` and some additional:

    :param session: Session which will do request
    :param callback: Callback called on response.
    Same as passing ``hooks={'response': callback}``
    """
    def __init__(self, method, url, **kwargs):
        #: Request method
        self.method = method
        #: URL to request
        self.url = url
        #: Associated ``Session``
        self.session = kwargs.pop('session', None)
        if self.session is None:
            self.session = Session()

        callback = kwargs.pop('callback', None)
        if callback:
            kwargs['hooks'] = {'response': callback}

        #: The rest arguments for ``Session.request``
        self.kwargs = kwargs
        #: Resulting ``Response``
        self.response = None

    def send(self, **kwargs):
        """
        Prepares request based on parameter passed to constructor and optional ``kwargs```.
        Then sends request and saves response to :attr:`response`

        :returns: ``Response``
        """
        merged_kwargs = {}
        merged_kwargs.update(self.kwargs)
        merged_kwargs.update(kwargs)
        self.response = self.session.request(self.method,
                                              self.url, **merged_kwargs)
        return self.response

def send(r, stream=False):
    """Just sends the request using its send method and returns its response.  """
    r.send(stream=stream)
    return r.response

# Shortcuts for creating AsyncRequest with appropriate HTTP method
get = partial(AsyncRequest, 'GET')
options = partial(AsyncRequest, 'OPTIONS')
head = partial(AsyncRequest, 'HEAD')
post = partial(AsyncRequest, 'POST')
put = partial(AsyncRequest, 'PUT')
patch = partial(AsyncRequest, 'PATCH')
delete = partial(AsyncRequest, 'DELETE')

# synonym
def request(method, url, **kwargs):
    return AsyncRequest(method, url, **kwargs)

def map(requests, stream=True, size=1, **kwargs):
    """Concurrently converts a list of Requests to Responses.

    :param requests: a collection of Request objects.
    :param stream: If False, the content will not be downloaded immediately.
    :param size: Specifies the number of requests to make at a time. If 1, no throttling occurs.
    """

    requests = list(requests)

    with ThreadPoolExecutor(max_workers=size) as executor:
        responses = list(executor.map(send, requests, [stream]*len(requests), **kwargs))

    return responses

def imap(requests, stream=True, size=2, **kwargs):
    """Concurrently converts a generator object of Requests to
    a generator of Responses.

    :param requests: a generator of Request objects.
    :param stream: If False, the content will not be downloaded immediately.
    :param size: Specifies the number of requests to make at a time. default is 2
    """

    def stream():
        while True:
            yield stream

    with ThreadPoolExecutor(max_workers=size) as executor:
        for response in executor.map(send, requests, stream(), **kwargs):
            yield response
<|endoftext|>


<commit_before>#! /usr/bin/env python2
from enum import Enum
import math, time

from nav_msgs.msg import Odometry
from sub_trajectory.msg import StabilityMode
from geometry_msgs.msg import Wrench

from dynamic_reconfigure.server import Server
from sub_trajectory.cfg import StabilityConfig

import tf
import rospy
import numpy as np

def rosToArray(msg): #Convert a ros message with 1-4 dimensions into a numpy array
    return np.array([getattr(msg, key) for key in ("x", "y", "z", "w") if hasattr(msg, key)])

def rotateVector(quat, vec):
	quat2 = list(vec)
	quat2.append(0.0)
	return tf.transformations.quaternion_multiply(
		tf.transformations.quaternion_multiply(quat, quat2),
		tf.transformations.quaternion_conjugate(quat)
	)[:3]

class ActiveStabilizer():
	def __init__(self):
		self.curDepthMode = StabilityMode.off
		self.curAngleMode = StabilityMode.off
		self.yawEnabled = False

		self.saveDepth = True
		self.saveOrientation = True

		self.targetDepth = None
		self.targetOrientation = None

		self.lastUpdateTime = None

		self.depthVelGain = 1
		self.depthPosGain = [1, 0, 1]
		self.depthLastPos = None
		self.depthIntegratedError = 0
		self.depthMaxForce = 5

		self.orientVelGain = 1
		self.orientPosGain = [1, 0, 1]
		self.orientLastError = None
		self.orientIntegratedError = np.array([0,0,0])
		self.orientMaxForce = 2

		self.reconfigureServer = Server(StabilityConfig, self.reconfigureCallback)

		self.angleModeSubscriber = rospy.Subscriber("/thrusters/angleMode", StabilityMode, self.angleModeCallback, queue_size=1)
		self.depthModeSubscriber = rospy.Subscriber("/thrusters/depthMode", StabilityMode, self.depthModeCallback, queue_size=1)
		self.odomSubcriber = rospy.Subscriber("/odometry/filtered",Odometry, self.callback, queue_size=1)

		self.stabilityPublisher = rospy.Publisher("/stabilityWrench", Wrench, queue_size=1)

		self.stabilityWrench = Wrench()

	def reconfigureCallback(self, config, level):
		self.depthVelGain = config["depth_vel_P"]
		self.depthPosGain = [config["depth_pos_P"],
							config["depth_pos_I"],
							config["depth_pos_D"]]
		self.depthMaxForce = config["depth_pos_Max"]

		self.orientVelGain = config["orientation_vel_P"]
		self.orientPosGain = [config["orientation_pos_P"],
							config["orientation_pos_I"],
							config["orientation_pos_D"]]
		self.orientMaxForce = config["orientation_pos_Max"]

		return config

	def depthModeCallback(self, msg):
		self.curDepthMode = msg.mode

		if self.curDepthMode == msg.position:
			self.depthLastPos = None

		if np.isfinite(msg.target.z):
			self.saveDepth = False
			self.targetDepth = msg.target.z
		elif self.curDepthMode != StabilityMode.position:
			self.saveDepth = True

	def angleModeCallback(self, msg):
		self.yawEnabled = msg.yawEnabled
		self.curAngleMode = msg.mode

		if self.curAngleMode == msg.position:
			self.orientLastError = None
		if np.isfinite(msg.target.w):
			self.saveOrientation = False
			self.targetOrientation = rosToArray(msg.target)
		if self.curAngleMode != StabilityMode.position:
			self.saveOrientation = True

	def callback(self, msg):
		timeNow = rospy.get_time()

		#TODO: How to reset the error integrations?
		if self.depthLastPos is None or self.curDepthMode is not StabilityMode.position:
			self.depthLastPos = msg.pose.pose.position.z

		if self.orientLastError is None or self.curAngleMode is not StabilityMode.position:
			self.orientLastError = np.array([0,0,0])
			self.orientIntegratedError = np.array([0,0,0])

		if self.curDepthMode == StabilityMode.off:
			self.stabilityWrench.force.x = 0
			self.stabilityWrench.force.y = 0
			self.stabilityWrench.force.z = 0

		elif self.curDepthMode == StabilityMode.velocity:
			quat = rosToArray(msg.pose.pose.orientation)
			counterVec = rotateVector(quat, [0,0,-self.depthVelGain*msg.twist.twist.linear.z])
			self.stabilityWrench.force.x = counterVec[0]
			self.stabilityWrench.force.y = counterVec[1]
			self.stabilityWrench.force.z = counterVec[2]

		elif self.curDepthMode == StabilityMode.position:
			if self.saveDepth:
				self.targetDepth = msg.pose.pose.position.z
				self.saveDepth = False
			error = msg.pose.pose.position.z - self.targetDepth

			proportionalCorrection = -self.depthPosGain[0]*error
			integralCorrection = 0
			derivativeCorrection = 0

			#TODO: Should we actually use the update time for this stuff?
			if self.lastUpdateTime is not None:
				#Depth derivative computation
				derivativeCorrection = -self.depthPosGain[2] * (msg.pose.pose.position.z - self.depthLastPos) / (timeNow - self.lastUpdateTime)
				self.depthLastPos = msg.pose.pose.position.z

				self.depthIntegratedError += error * (timeNow - self.lastUpdateTime)
				integralCorrection = -self.depthPosGain[1] * self.depthIntegratedError
				#Back-calculation integrator windup prevention
				if abs(proportionalCorrection + derivativeCorrection + integralCorrection) > self.depthMaxForce:
					self.depthIntegratedError = ((-np.sign(error)*self.depthMaxForce) - derivativeCorrection - proportionalCorrection)/-self.depthPosGain[1]
				integralCorrection = -self.depthPosGain[1] * self.depthIntegratedError

			#print([ proportionalCorrection, derivativeCorrection, integralCorrection])

			quat = rosToArray(msg.pose.pose.orientation)
			quat = tf.transformations.quaternion_conjugate(quat)
			counterVec = rotateVector(quat, [0,0,proportionalCorrection + derivativeCorrection + integralCorrection])
			self.stabilityWrench.force.x = counterVec[0]
			self.stabilityWrench.force.y = counterVec[1]
			self.stabilityWrench.force.z = counterVec[2]

		else:
			self.stabilityWrench.force.x = 0
			self.stabilityWrench.force.y = 0
			self.stabilityWrench.force.z = 0

		if self.curAngleMode == StabilityMode.off:
			self.stabilityWrench.torque.x = 0
			self.stabilityWrench.torque.y = 0
			self.stabilityWrench.torque.z = 0

		elif self.curAngleMode == StabilityMode.velocity:
			self.stabilityWrench.torque.x = -self.orientVelGain * msg.twist.twist.angular.x
			self.stabilityWrench.torque.y = -self.orientVelGain * msg.twist.twist.angular.y

			if self.yawEnabled:
				self.stabilityWrench.torque.z = -self.orientVelGain * msg.twist.twist.angular.z
			else:
				self.stabilityWrench.torque.z = 0.0

		elif self.curAngleMode == StabilityMode.position:
			if self.saveOrientation:
				self.targetOrientation = rosToArray(msg.pose.pose.orientation)
				self.saveOrientation = False

			error = tf.transformations.quaternion_multiply(
				self.targetOrientation,
				tf.transformations.quaternion_conjugate(rosToArray(msg.pose.pose.orientation))
			) #Find rotation from current to target position

			#TODO: Do we need to transform the error and how
			#subLocalError = tf.transformations.quaternion_multiply(
			#	error, rosToArray(msg.pose.orientation)
			#) #Rotate rotation from current to target by rotation from world to sub

			rpyError = tf.transformations.euler_from_quaternion(error)

			proportionalCorrection = -self.orientPosGain[0]*rpyError
			integralCorrection = np.array([0,0,0])
			derivativeCorrection = np.array([0,0,0])

			#TODO: Should we actually use the update time for this stuff?
			if self.lastUpdateTime is not None:
				#Depth derivative computation
				derivativeCorrection = -self.orientPosGain[2] * (rpyError - self.orientLastError) / (timeNow - self.lastUpdateTime)
				self.orientLastError = rpyError

				self.orientIntegratedError += rpyError * (timeNow - self.lastUpdateTime)
				integralCorrection = -self.orientPosGain[1] * self.orientIntegratedError
				#Back-calculation integrator windup prevention
				for i in range(3):
					if abs(proportionalCorrection[i] + derivativeCorrection[i] + integralCorrection[i]) > self.orientMaxForce:
						self.orientIntegratedError[i] = ((-np.sign(rpyError[i])*self.orientMaxForce) - derivativeCorrection[i] - proportionalCorrection[i])/-self.orientPosGain[1]

				integralCorrection = -self.orientPosGain[1] * self.orientIntegratedError

			self.stabilityWrench.torque.x = proportionalCorrection[0] + derivativeCorrection[0] + integralCorrection[0]
			self.stabilityWrench.torque.y = proportionalCorrection[1] + derivativeCorrection[1] + integralCorrection[1]
			if self.yawEnabled:
				self.stabilityWrench.torque.z = proportionalCorrection[2] + derivativeCorrection[2] + integralCorrection[2]
			else:
				self.stabilityWrench.torque.z = 0.0

		else:
			self.stabilityWrench.torque.x = 0
			self.stabilityWrench.torque.y = 0
			self.stabilityWrench.torque.z = 0

		self.stabilityPublisher.publish(self.stabilityWrench)
		self.lastUpdateTime = timeNow

if __name__== "__main__":
	rospy.init_node('ActiveStability', anonymous=False)
	actStab = ActiveStabilizer()
	rospy.spin()
<commit_msg>Fix active stability divide by 0 and array type errors<commit_after>#! /usr/bin/env python2
from enum import Enum
import math, time

from nav_msgs.msg import Odometry
from sub_trajectory.msg import StabilityMode
from geometry_msgs.msg import Wrench
from std_msgs.msg import Float64
from dynamic_reconfigure.server import Server
from sub_trajectory.cfg import StabilityConfig

import tf
import rospy
import numpy as np

def rosToArray(msg): #Convert a ros message with 1-4 dimensions into a numpy array
    return np.array([getattr(msg, key) for key in ("x", "y", "z", "w") if hasattr(msg, key)])

def rotateVector(quat, vec):
	quat2 = list(vec)
	quat2.append(0.0)
	return tf.transformations.quaternion_multiply(
		tf.transformations.quaternion_multiply(quat, quat2),
		tf.transformations.quaternion_conjugate(quat)
	)[:3]

class ActiveStabilizer():
	def __init__(self):
		self.curDepthMode = StabilityMode.off
		self.curAngleMode = StabilityMode.off
		self.yawEnabled = False

		self.saveDepth = True
		self.saveOrientation = True

		self.targetDepth = None
		self.targetOrientation = None

		self.lastUpdateTime = None

		self.depthVelGain = 1
		self.depthPosGain = [1, 0, 1]
		self.depthLastPos = None
		self.depthIntegratedError = 0
		self.depthMaxForce = 5

		self.orientVelGain = 1
		self.orientPosGain = [1, 0, 1]
		self.orientLastError = None
		self.orientIntegratedError = np.array([0.0,0.0,0.0],dtype="float64")
		self.orientMaxForce = 2

		self.reconfigureServer = Server(StabilityConfig, self.reconfigureCallback)

		self.angleModeSubscriber = rospy.Subscriber("/thrusters/angleMode", StabilityMode, self.angleModeCallback, queue_size=1)
		self.depthModeSubscriber = rospy.Subscriber("/thrusters/depthMode", StabilityMode, self.depthModeCallback, queue_size=1)
		self.odomSubcriber = rospy.Subscriber("/odometry/filtered",Odometry, self.callback, queue_size=1)

		self.stabilityPublisher = rospy.Publisher("/stabilityWrench", Wrench, queue_size=1)

		self.stabilityWrench = Wrench()

		self.setPointPub = rospy.Publisher("/setpoint", Float64, queue_size=10)

	def reconfigureCallback(self, config, level):
		self.depthVelGain = config["depth_vel_P"]
		self.depthPosGain = [config["depth_pos_P"],
							config["depth_pos_I"],
							config["depth_pos_D"]]
		self.depthMaxForce = config["depth_pos_Max"]

		self.orientVelGain = config["orientation_vel_P"]
		self.orientPosGain = [config["orientation_pos_P"],
							config["orientation_pos_I"],
							config["orientation_pos_D"]]
		self.orientMaxForce = config["orientation_pos_Max"]

		return config

	def depthModeCallback(self, msg):
		self.curDepthMode = msg.mode

		if self.curDepthMode == msg.position:
			self.depthLastPos = None

		if np.isfinite(msg.target.z):
			self.saveDepth = False
			self.targetDepth = msg.target.z
		elif self.curDepthMode != StabilityMode.position:
			self.saveDepth = True

	def angleModeCallback(self, msg):
		self.yawEnabled = msg.yawEnabled
		self.curAngleMode = msg.mode

		if self.curAngleMode == msg.position:
			self.orientLastError = None
		if np.isfinite(msg.target.w):
			self.saveOrientation = False
			self.targetOrientation = rosToArray(msg.target)
		if self.curAngleMode != StabilityMode.position:
			self.saveOrientation = True

	def callback(self, msg):
		timeNow = rospy.get_time()

		#TODO: How to reset the error integrations?
		if self.depthLastPos is None or self.curDepthMode is not StabilityMode.position:
			self.depthLastPos = msg.pose.pose.position.z

		if self.orientLastError is None or self.curAngleMode is not StabilityMode.position:
			self.orientLastError = np.array([0,0,0])
			self.orientIntegratedError = np.array([0.0,0.0,0.0],dtype="float64")

		if self.curDepthMode == StabilityMode.off:
			self.stabilityWrench.force.x = 0
			self.stabilityWrench.force.y = 0
			self.stabilityWrench.force.z = 0

		elif self.curDepthMode == StabilityMode.velocity:
			quat = rosToArray(msg.pose.pose.orientation)
			counterVec = rotateVector(quat, [0,0,-self.depthVelGain*msg.twist.twist.linear.z])
			self.stabilityWrench.force.x = counterVec[0]
			self.stabilityWrench.force.y = counterVec[1]
			self.stabilityWrench.force.z = counterVec[2]

		elif self.curDepthMode == StabilityMode.position:
			if self.saveDepth:
				self.targetDepth = msg.pose.pose.position.z
				self.saveDepth = False
			self.setPointPub.publish(Float64(self.targetDepth))
			error = msg.pose.pose.position.z - self.targetDepth

			proportionalCorrection = -self.depthPosGain[0]*error
			integralCorrection = 0
			derivativeCorrection = 0

			#TODO: Should we actually use the update time for this stuff?
			if self.lastUpdateTime is not None:
				#Depth derivative computation
				derivativeCorrection = -self.depthPosGain[2] * (msg.pose.pose.position.z - self.depthLastPos) / (timeNow - self.lastUpdateTime)
				self.depthLastPos = msg.pose.pose.position.z

				if self.depthPosGain[1] > 0.0001:
					self.depthIntegratedError += error * (timeNow - self.lastUpdateTime)
				integralCorrection = -self.depthPosGain[1] * self.depthIntegratedError
				#Back-calculation integrator windup prevention
				if abs(proportionalCorrection + derivativeCorrection + integralCorrection) > self.depthMaxForce and self.depthPosGain[1] > 0.0001:
					self.depthIntegratedError = ((-np.sign(error)*self.depthMaxForce) - derivativeCorrection - proportionalCorrection)/-self.depthPosGain[1]
				integralCorrection = -self.depthPosGain[1] * self.depthIntegratedError

			#print([ proportionalCorrection, derivativeCorrection, integralCorrection])

			quat = rosToArray(msg.pose.pose.orientation)
			quat = tf.transformations.quaternion_conjugate(quat)
			counterVec = rotateVector(quat, [0,0,proportionalCorrection + derivativeCorrection + integralCorrection])
			self.stabilityWrench.force.x = counterVec[0]
			self.stabilityWrench.force.y = counterVec[1]
			self.stabilityWrench.force.z = counterVec[2]

		else:
			self.stabilityWrench.force.x = 0
			self.stabilityWrench.force.y = 0
			self.stabilityWrench.force.z = 0

		if self.curAngleMode == StabilityMode.off:
			self.stabilityWrench.torque.x = 0
			self.stabilityWrench.torque.y = 0
			self.stabilityWrench.torque.z = 0

		elif self.curAngleMode == StabilityMode.velocity:
			self.stabilityWrench.torque.x = -self.orientVelGain * msg.twist.twist.angular.x
			self.stabilityWrench.torque.y = -self.orientVelGain * msg.twist.twist.angular.y

			if self.yawEnabled:
				self.stabilityWrench.torque.z = -self.orientVelGain * msg.twist.twist.angular.z
			else:
				self.stabilityWrench.torque.z = 0.0

		elif self.curAngleMode == StabilityMode.position:
			if self.saveOrientation:
				self.targetOrientation = rosToArray(msg.pose.pose.orientation)
				self.saveOrientation = False

			error = tf.transformations.quaternion_multiply(
				self.targetOrientation,
				tf.transformations.quaternion_conjugate(rosToArray(msg.pose.pose.orientation))
			) #Find rotation from current to target position

			#TODO: Do we need to transform the error and how
			#subLocalError = tf.transformations.quaternion_multiply(
			#	error, rosToArray(msg.pose.orientation)
			#) #Rotate rotation from current to target by rotation from world to sub

			rpyError = np.array(tf.transformations.euler_from_quaternion(error))

			proportionalCorrection = -self.orientPosGain[0]*rpyError
			integralCorrection = np.array([0,0,0])
			derivativeCorrection = np.array([0,0,0])

			#TODO: Should we actually use the update time for this stuff?
			if self.lastUpdateTime is not None:
				#Depth derivative computation
				derivativeCorrection = -self.orientPosGain[2] * (rpyError - self.orientLastError) / (timeNow - self.lastUpdateTime)
				self.orientLastError = rpyError
				if self.orientPosGain[1] > 0.0001:
					self.orientIntegratedError += rpyError * (timeNow - self.lastUpdateTime)
				integralCorrection = -self.orientPosGain[1] * self.orientIntegratedError
				#Back-calculation integrator windup prevention
				if self.orientPosGain[1] > 0.0001:
					for i in range(3):
						if abs(proportionalCorrection[i] + derivativeCorrection[i] + integralCorrection[i]) > self.orientMaxForce:
							self.orientIntegratedError[i] = ((-np.sign(rpyError[i])*self.orientMaxForce) - derivativeCorrection[i] - proportionalCorrection[i])/-self.orientPosGain[1]

				integralCorrection = -self.orientPosGain[1] * self.orientIntegratedError

			self.stabilityWrench.torque.x = proportionalCorrection[0] + derivativeCorrection[0] + integralCorrection[0]
			self.stabilityWrench.torque.y = proportionalCorrection[1] + derivativeCorrection[1] + integralCorrection[1]
			if self.yawEnabled:
				self.stabilityWrench.torque.z = proportionalCorrection[2] + derivativeCorrection[2] + integralCorrection[2]
			else:
				self.stabilityWrench.torque.z = 0.0

		else:
			self.stabilityWrench.torque.x = 0
			self.stabilityWrench.torque.y = 0
			self.stabilityWrench.torque.z = 0

		self.stabilityPublisher.publish(self.stabilityWrench)
		self.lastUpdateTime = timeNow

if __name__== "__main__":
	rospy.init_node('ActiveStability', anonymous=False)
	actStab = ActiveStabilizer()
	rospy.spin()
<|endoftext|>