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

"""
This module defines filters for Transmuter object.
"""

from __future__ import division

__author__ = "Will Richards, Shyue Ping Ong, Stephen Dacek"
__copyright__ = "Copyright 2011, The Materials Project"
__version__ = "1.0"
__maintainer__ = "Will Richards"
__email__ = "wrichards@mit.edu"
__date__ = "Sep 25, 2012"

from pymatgen.core.periodic_table import smart_element_or_specie
from pymatgen.serializers.json_coders import MSONable
from pymatgen.analysis.structure_matcher import StructureMatcher,\
    ElementComparator
import abc

class AbstractStructureFilter(MSONable):
    """
    Abstract structure filter class.
    """
    __metaclass__ = abc.ABCMeta

    @abc.abstractmethod
    def test(self, structure):
        """
        Returns a boolean for any structure. Structures that return true are
        kept in the Transmuter object during filtering.
        """
        return

    @staticmethod
    def from_dict(d):
        for trans_modules in ['filters']:
            mod = __import__('pymatgen.alchemy.' + trans_modules,
                             globals(), locals(), [d['@class']], -1)
            if hasattr(mod, d['@class']):
                trans = getattr(mod, d['@class'])
                return trans(**d['init_args'])
        raise ValueError("Invalid filter dict")

class ContainsSpecieFilter(AbstractStructureFilter):
    """
    Filter for structures containing certain elements or species.
    By default compares by atomic number
    """

    def __init__(self, species, strict_compare=False, AND=True, exclude=False):
        """
        Args:
            species:
                list of species to look for
            AND:
                whether all species must be present to pass (or fail)
                filter.
            strict_compare:
                if true, compares objects by specie or element object
                if false, compares atomic number
            exclude:
                if true, returns false for any structures with the specie
                (excludes them from the Transmuter)
        """
        self._species = map(smart_element_or_specie, species)
        self._strict = strict_compare
        self._AND = AND
        self._exclude = exclude

    def test(self, structure):
        #set up lists to compare
        if not self._strict:
            #compare by atomic number
            atomic_number = lambda x: x.Z
            filter_set = set(map(atomic_number, self._species))
            structure_set = set(map(atomic_number,
                                structure.composition.elements))
        else:
            #compare by specie or element object
            filter_set = set(self._species)
            structure_set = set(structure.composition.elements)

        if self._AND and filter_set <= structure_set:
            #return true if we aren't excluding since all are in structure
            return not self._exclude
        elif (not self._AND) and filter_set & structure_set:
            #return true if we aren't excluding since one is in structure
            return not self._exclude
        else:
            #return false if we aren't excluding otherwise
            return self._exclude

    def __repr__(self):
        return "\n".join(["ContainsSpecieFilter with parameters:",
                          "species = {}".format(self._species),
                          "strict_compare = {}".format(self._strict),
                          "AND = {}".format(self._AND),
                          "exclude = {}".format(self._exclude)])

    @property
    def to_dict(self):
        return {"version": __version__, "@module": self.__class__.__module__,
                "@class": self.__class__.__name__,
                "init_args": {"species": [str(sp) for sp in self._species],
                              "strict_compare": self._strict,
                              "AND": self._AND,
                              "exclude": self._exclude}}

class SpecieProximityFilter(AbstractStructureFilter):
    """
    This filter removes structures that have certain species that are too close
    together.
    """

    def __init__(self, specie_and_min_dist_dict):
        """
        Args:
            specie_and_min_dist_dict:
                A species string to float mapping. For example, {"Na+": 1}
                means that all Na+ ions must be at least 1 Angstrom away from
                each other. Multiple species criteria can be applied. Note that
                the testing is done based on the actual object. If you have a
                structure with Element, you must use {"Na":1} instead to filter
                based on Element and not Specie.
        """
        self.specie_and_min_dist = {smart_element_or_specie(k): v
                                    for k, v
                                    in specie_and_min_dist_dict.items()}

    def test(self, structure):
        all_species = set(self.specie_and_min_dist.keys())
        for site in structure:
            species = site.species_and_occu.keys()
            sp_to_test = set(species).intersection(all_species)
            if sp_to_test:
                max_r = max([self.specie_and_min_dist[sp]
                             for sp in sp_to_test])
                nn = structure.get_neighbors(site, max_r)
                for sp in sp_to_test:
                    for (nnsite, dist) in nn:
                        if sp in nnsite.species_and_occu.keys():
                            if dist < self.specie_and_min_dist[sp]:
                                return False
        return True

    @property
    def to_dict(self):
        return {"version": __version__, "@module": self.__class__.__module__,
                "@class": self.__class__.__name__,
                "init_args": {"specie_and_min_dist_dict":
                              {str(sp): v
                              for sp, v in self.specie_and_min_dist.items()}}}

class RemoveDuplicatesFilter(AbstractStructureFilter):
    """
    This filter removes exact duplicate structures from the transmuter.
    """

    def __init__(self,
                 structure_matcher=StructureMatcher(
                     comparator=ElementComparator())):
        """
        Args:
            comparator:
                The comparator to be used in the matching
        """
        self._structure_list = []
        if isinstance(structure_matcher, dict):
            self._sm = StructureMatcher.from_dict(structure_matcher)
        else:
            self._sm = structure_matcher

    def test(self, structure):
        if not self._structure_list:
            self._structure_list.append(structure)
            return True

        for s in self._structure_list:
            if self._sm._comparator.get_structure_hash(structure) ==\
                    self._sm._comparator.get_structure_hash(s):
                if self._sm.fit(s, structure):
                    return False

        self._structure_list.append(structure)
        return True

    @property
    def to_dict(self):
        return {"version": __version__, "@module": self.__class__.__module__,
                "@class": self.__class__.__name__,
                "init_args": {"structure_matcher": self._sm.to_dict}}
<commit_msg>added filter to check charge balance for structures in the transmuter. Makes doing substitutions easier<commit_after>#!/usr/bin/env python

"""
This module defines filters for Transmuter object.
"""

from __future__ import division

__author__ = "Will Richards, Shyue Ping Ong, Stephen Dacek"
__copyright__ = "Copyright 2011, The Materials Project"
__version__ = "1.0"
__maintainer__ = "Will Richards"
__email__ = "wrichards@mit.edu"
__date__ = "Sep 25, 2012"

from pymatgen.core.periodic_table import smart_element_or_specie
from pymatgen.serializers.json_coders import MSONable
from pymatgen.analysis.structure_matcher import StructureMatcher,\
    ElementComparator
import abc

class AbstractStructureFilter(MSONable):
    """
    Abstract structure filter class.
    """
    __metaclass__ = abc.ABCMeta

    @abc.abstractmethod
    def test(self, structure):
        """
        Returns a boolean for any structure. Structures that return true are
        kept in the Transmuter object during filtering.
        """
        return

    @staticmethod
    def from_dict(d):
        for trans_modules in ['filters']:
            mod = __import__('pymatgen.alchemy.' + trans_modules,
                             globals(), locals(), [d['@class']], -1)
            if hasattr(mod, d['@class']):
                trans = getattr(mod, d['@class'])
                return trans(**d['init_args'])
        raise ValueError("Invalid filter dict")

class ContainsSpecieFilter(AbstractStructureFilter):
    """
    Filter for structures containing certain elements or species.
    By default compares by atomic number
    """

    def __init__(self, species, strict_compare=False, AND=True, exclude=False):
        """
        Args:
            species:
                list of species to look for
            AND:
                whether all species must be present to pass (or fail)
                filter.
            strict_compare:
                if true, compares objects by specie or element object
                if false, compares atomic number
            exclude:
                if true, returns false for any structures with the specie
                (excludes them from the Transmuter)
        """
        self._species = map(smart_element_or_specie, species)
        self._strict = strict_compare
        self._AND = AND
        self._exclude = exclude

    def test(self, structure):
        #set up lists to compare
        if not self._strict:
            #compare by atomic number
            atomic_number = lambda x: x.Z
            filter_set = set(map(atomic_number, self._species))
            structure_set = set(map(atomic_number,
                                structure.composition.elements))
        else:
            #compare by specie or element object
            filter_set = set(self._species)
            structure_set = set(structure.composition.elements)

        if self._AND and filter_set <= structure_set:
            #return true if we aren't excluding since all are in structure
            return not self._exclude
        elif (not self._AND) and filter_set & structure_set:
            #return true if we aren't excluding since one is in structure
            return not self._exclude
        else:
            #return false if we aren't excluding otherwise
            return self._exclude

    def __repr__(self):
        return "\n".join(["ContainsSpecieFilter with parameters:",
                          "species = {}".format(self._species),
                          "strict_compare = {}".format(self._strict),
                          "AND = {}".format(self._AND),
                          "exclude = {}".format(self._exclude)])

    @property
    def to_dict(self):
        return {"version": __version__, "@module": self.__class__.__module__,
                "@class": self.__class__.__name__,
                "init_args": {"species": [str(sp) for sp in self._species],
                              "strict_compare": self._strict,
                              "AND": self._AND,
                              "exclude": self._exclude}}

class SpecieProximityFilter(AbstractStructureFilter):
    """
    This filter removes structures that have certain species that are too close
    together.
    """

    def __init__(self, specie_and_min_dist_dict):
        """
        Args:
            specie_and_min_dist_dict:
                A species string to float mapping. For example, {"Na+": 1}
                means that all Na+ ions must be at least 1 Angstrom away from
                each other. Multiple species criteria can be applied. Note that
                the testing is done based on the actual object. If you have a
                structure with Element, you must use {"Na":1} instead to filter
                based on Element and not Specie.
        """
        self.specie_and_min_dist = {smart_element_or_specie(k): v
                                    for k, v
                                    in specie_and_min_dist_dict.items()}

    def test(self, structure):
        all_species = set(self.specie_and_min_dist.keys())
        for site in structure:
            species = site.species_and_occu.keys()
            sp_to_test = set(species).intersection(all_species)
            if sp_to_test:
                max_r = max([self.specie_and_min_dist[sp]
                             for sp in sp_to_test])
                nn = structure.get_neighbors(site, max_r)
                for sp in sp_to_test:
                    for (nnsite, dist) in nn:
                        if sp in nnsite.species_and_occu.keys():
                            if dist < self.specie_and_min_dist[sp]:
                                return False
        return True

    @property
    def to_dict(self):
        return {"version": __version__, "@module": self.__class__.__module__,
                "@class": self.__class__.__name__,
                "init_args": {"specie_and_min_dist_dict":
                              {str(sp): v
                              for sp, v in self.specie_and_min_dist.items()}}}

class RemoveDuplicatesFilter(AbstractStructureFilter):
    """
    This filter removes exact duplicate structures from the transmuter.
    """

    def __init__(self,
                 structure_matcher=StructureMatcher(
                     comparator=ElementComparator())):
        """
        Args:
            comparator:
                The comparator to be used in the matching
        """
        self._structure_list = []
        if isinstance(structure_matcher, dict):
            self._sm = StructureMatcher.from_dict(structure_matcher)
        else:
            self._sm = structure_matcher

    def test(self, structure):
        if not self._structure_list:
            self._structure_list.append(structure)
            return True

        for s in self._structure_list:
            if self._sm._comparator.get_structure_hash(structure) ==\
                    self._sm._comparator.get_structure_hash(s):
                if self._sm.fit(s, structure):
                    return False

        self._structure_list.append(structure)
        return True

    @property
    def to_dict(self):
        return {"version": __version__, "@module": self.__class__.__module__,
                "@class": self.__class__.__name__,
                "init_args": {"structure_matcher": self._sm.to_dict}}

class ChargeBalanceFilter(AbstractStructureFilter):
    """
     This filter removes structures that are not
     charge balanced from the transmuter.
     This only works if the structure is
     oxidation state decorated, as structures
     with only elemental sites are automatically
     assumed to have net charge of 0
    """
    def test(self, structure):
        if structure.charge == 0.0:
            return True
        else:
            return False

    @property
    def to_dict(self):
        return {"@module": self.__class__.__module__,
                "@class": self.__class__.__name__}
<|endoftext|>


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

import os
import rospy
import urllib2
import json
from lg_json_config.srv import JSONConfig
from lg_json_config.srv import JSONConfigMore

NODE_NAME = 'lg_json_config'
SRV_QUERY = '/'.join(('', NODE_NAME, 'query'))
SRV_MORE_QUERY = '/'.join(('', NODE_NAME, 'query_more'))
PARAM_URL = '~url'
DEFAULT_URL = 'http://lg-head/portal/config.json'

class ConfigRequestHandler():
    def get_url(self):
        return rospy.get_param(PARAM_URL, DEFAULT_URL)

    def get_config(self, url=None):
        if not url:
            url = self.get_url()
        try:
            response = urllib2.urlopen(url)
            r = response.read()
        except Exception, e:
            r = """{"result": "error": "reason": "%s"}""" % e
            return r

        try:
            json.loads(r)
        except ValueError:
            r = """{"result": "error": "reason": "Could not parse json from %s"}""" % url
            return r

        return r

    def handle_request(self, request):
        config = self.get_config()
        return JSONConfigResponse(config)

    def handle_request_more(self, request):
        # using the default url to avoid a security hole where a user could
        # input file:// as the ~url param, and get access to any file on the
        # local filesystem
        url = os.path.join(os.path.dirname(DEFAULT_URL), request.filename)
        config = self.get_config(url)
        return JSONConfigMoreResponse(config)

def main():
    rospy.init_node(NODE_NAME)

    handler = ConfigRequestHandler()

    rospy.Service(
        SRV_QUERY,
        JSONConfig,
        handler.handle_request
    )

    rospy.Service(
        SRV_MORE_QUERY,
        JSONConfigMore,
        handler.handle_request_more
    )

    rospy.spin()

if __name__ == '__main__':
    main()

# vim: tabstop=8 expandtab shiftwidth=4 softtabstop=4 smartindent
<commit_msg>Added proper imports<commit_after>#!/usr/bin/env python

import os
import rospy
import urllib2
import json
from lg_json_config.srv import JSONConfig
from lg_json_config.srv import JSONConfigMore
from lg_json_config.srv import JSONConfigResponse
from lg_json_config.srv import JSONConfigMoreResponse

NODE_NAME = 'lg_json_config'
SRV_QUERY = '/'.join(('', NODE_NAME, 'query'))
SRV_MORE_QUERY = '/'.join(('', NODE_NAME, 'query_more'))
PARAM_URL = '~url'
DEFAULT_URL = 'http://lg-head/portal/config.json'

class ConfigRequestHandler():
    def get_url(self):
        return rospy.get_param(PARAM_URL, DEFAULT_URL)

    def get_config(self, url=None):
        if not url:
            url = self.get_url()
        try:
            response = urllib2.urlopen(url)
            r = response.read()
        except Exception, e:
            r = """{"result": "error": "reason": "%s"}""" % e
            return r

        try:
            json.loads(r)
        except ValueError:
            r = """{"result": "error": "reason": "Could not parse json from %s"}""" % url
            return r

        return r

    def handle_request(self, request):
        config = self.get_config()
        return JSONConfigResponse(config)

    def handle_request_more(self, request):
        # using the default url to avoid a security hole where a user could
        # input file:// as the ~url param, and get access to any file on the
        # local filesystem
        url = os.path.join(os.path.dirname(DEFAULT_URL), request.filename)
        config = self.get_config(url)
        return JSONConfigMoreResponse(config)

def main():
    rospy.init_node(NODE_NAME)

    handler = ConfigRequestHandler()

    rospy.Service(
        SRV_QUERY,
        JSONConfig,
        handler.handle_request
    )

    rospy.Service(
        SRV_MORE_QUERY,
        JSONConfigMore,
        handler.handle_request_more
    )

    rospy.spin()

if __name__ == '__main__':
    main()

# vim: tabstop=8 expandtab shiftwidth=4 softtabstop=4 smartindent
<|endoftext|>


<commit_before># coding: utf-8

from __future__ import division, print_function, unicode_literals, absolute_import

"""
This module defines utility classes and functions.
"""

import six
from io import open
import os
import tempfile
from subprocess import Popen, PIPE

import numpy as np

try:
    import pybel as pb
except ImportError:
    pb = None
from pymatgen import Molecule
from pymatgen.core.operations import SymmOp
from pymatgen.util.coord_utils import get_angle
from pymatgen.io.babel import BabelMolAdaptor

from monty.os.path import which
from monty.dev import requires
from monty.tempfile import ScratchDir

__author__ = 'Kiran Mathew, Michael Humbert, Brandon Wood'
__email__ = 'kmathew@lbl.gov'

class Polymer(object):
    """
    Generate polymer molecule via Random walk. At each poisition there are
    a total of 5 possible moves(exclusing the previous directions).
    """
    def __init__(self, start_monomer, s_head, s_tail,
                 monomer, head, tail,
                 end_monomer, e_head, e_tail,
                 n_units, link_distance=1.0, linear_chain=False):
        """
        Args:
            start_monomer (Molecule): Starting molecule
            s_head (int): starting atom index of the start_monomer molecule
            s_tail (int): tail atom index of the start_monomer
            monomer (Molecule): The monomer
            head (int): index of the atom in the monomer that forms the head
            tail (int): tail atom index. monomers will be connected from
                tail to head
            end_monomer (Molecule): Terminal molecule
            e_head (int): starting atom index of the end_monomer molecule
            e_tail (int): tail atom index of the end_monomer
            n_units (int): number of monomer units excluding the start and
                terminal molecules
            link_distance (float): distance between consecutive monomers
            linear_chain (bool): linear or random walk polymer chain
        """
        self.start = s_head
        self.end = s_tail
        self.monomer = monomer
        self.n_units = n_units
        self.link_distance = link_distance
        self.linear_chain = linear_chain
        # translate monomers so that head atom is at the origin
        start_monomer.translate_sites(range(len(start_monomer)),
                                      - monomer.cart_coords[s_head])
        monomer.translate_sites(range(len(monomer)),
                                - monomer.cart_coords[head])
        end_monomer.translate_sites(range(len(end_monomer)),
                                    - monomer.cart_coords[e_head])
        self.mon_vector = monomer.cart_coords[tail] - monomer.cart_coords[head]
        self.moves = {1: [1, 0, 0],
                      2: [0, 1, 0],
                      3: [0, 0, 1],
                      4: [-1, 0, 0],
                      5: [0, -1, 0],
                      6: [0, 0, -1]}
        self.prev_move = 1
        # places the start monomer at the beginning of the chain
        self.molecule = start_monomer.copy()
        self.length = 1
        # create the chain
        self._create(self.monomer, self.mon_vector)
        # terminate the chain with the end_monomer
        self.n_units += 1
        end_mon_vector = end_monomer.cart_coords[e_tail] - \
                         end_monomer.cart_coords[e_head]
        self._create(end_monomer, end_mon_vector)

    def _create(self, monomer, mon_vector):
        """
        create the polymer from the monomer

        Args:
            monomer (Molecule)
            mon_vector (numpy.array): molecule vector that starts from the
                start atom index to the end atom index
        """
        while self.length != (self.n_units-1):
            if self.linear_chain:
                move_direction = np.array(mon_vector) / np.linalg.norm(mon_vector)
            else:
                move_direction = self._next_move_direction()
            self._add_monomer(monomer.copy(), mon_vector, move_direction)

    def _next_move_direction(self):
        """
        pick a move at random from the list of moves
        """
        nmoves = len(self.moves)
        move = np.random.randint(1, nmoves+1)
        while self.prev_move == (move + 3) % nmoves:
            move = np.random.randint(1, nmoves+1)
        self.prev_move = move
        return np.array(self.moves[move])

    def _align_monomer(self, monomer, mon_vector, move_direction):
        """
        rotate the monomer so that it is aligned along the move direction

        Args:
            monomer (Molecule)
            mon_vector (numpy.array): molecule vector that starts from the
                start atom index to the end atom index
            move_direction (numpy.array): the direction of the polymer chain
                extension
        """
        axis = np.cross(mon_vector, move_direction)
        origin = monomer[self.start].coords
        angle = get_angle(mon_vector, move_direction)
        op = SymmOp.from_origin_axis_angle(origin, axis, angle)
        monomer.apply_operation(op)

    def _add_monomer(self, monomer, mon_vector, move_direction):
        """
        extend the polymer molecule by adding a monomer along mon_vector direction

        Args:
            monomer (Molecule): monomer molecule
            mon_vector (numpy.array): monomer vector that points from head to tail.
            move_direction (numpy.array): direction along which the monomer
                will be positioned
        """
        translate_by = self.molecule.cart_coords[self.end] + \
                       self.link_distance * move_direction
        monomer.translate_sites(range(len(monomer)), translate_by)
        if not self.linear_chain:
            self._align_monomer(monomer, mon_vector, move_direction)
        # add monomer if there are no crossings
        does_cross = False
        for i, site in enumerate(monomer):
            try:
                self.molecule.append(site.specie, site.coords,
                                     properties=site.properties)
            except:
                does_cross = True
                polymer_length = len(self.molecule)
                self.molecule.remove_sites(
                    range(polymer_length - i, polymer_length))
                break
        if not does_cross:
            self.length += 1
            self.end += len(self.monomer)

class PackmolRunner(object):
    """
    Wrapper for the Packmol software that can be used to pack various types of
    molecules into a one single unit.
    """

    @requires(which(b'packmol'),
              "PackmolRunner requires the executable 'packmol' to be in "
              "the path. Please download packmol from "
              "https://github.com/leandromartinez98/packmol "
              "and follow the instructions in the README to compile. "
              "Don't forget to add the packmol binary to your path")
    def __init__(self, mols, param_list, input_file="pack.inp",
                 tolerance=2.0, filetype="xyz",
                 control_params={"maxit": 20, "nloop": 600},
                 auto_box=True, output_file="packed.xyz"):
        """
        Args:
              mols:
                   list of Molecules to pack
              input_file:
                        name of the packmol input file
              tolerance:
                        min distance between the atoms
              filetype:
                       input/output structure file type
              control_params:
                           packmol control parameters dictionary. Basically
                           all parameters other than structure/atoms
              param_list:
                    list of parameters containing dicts for each molecule
              auto_box:
                    put the molecule assembly in a box
              output_file:
                    output file name. The extension will be adjusted
                    according to the filetype
        """
        self.mols = mols
        self.param_list = param_list
        self.input_file = input_file
        self.boxit = auto_box
        self.control_params = control_params
        if not self.control_params.get("tolerance"):
            self.control_params["tolerance"] = tolerance
        if not self.control_params.get("filetype"):
            self.control_params["filetype"] = filetype
        if not self.control_params.get("output"):
            self.control_params["output"] = "{}.{}".format(
                output_file.split(".")[0], self.control_params["filetype"])
        if self.boxit:
            self._set_box()

    def _format_param_val(self, param_val):
        """
        Internal method to format values in the packmol parameter dictionaries

        Args:
              param_val:
                   Some object to turn into String

        Returns:
               string representation of the object
        """
        if isinstance(param_val, list):
            return ' '.join(str(x) for x in param_val)
        else:
            return str(param_val)

    def _set_box(self):
        """
        Set the box size for the molecular assembly
        """
        volume = 0.0
        for idx, mol in enumerate(self.mols):
            length = max([np.max(mol.cart_coords[:, i])-np.min(mol.cart_coords[:, i])
                           for i in range(3)]) + 2.0
            volume += (length**3.0) * float(self.param_list[idx]['number'])
        length = volume**(1.0/3.0)
        for idx, mol in enumerate(self.mols):
            self.param_list[idx]['inside box'] = '0.0 0.0 0.0 {} {} {}'.format(
                length, length, length)

    def run(self):
        """
        Runs packmol

        Returns:
                Molecule object
        """
        scratch = tempfile.gettempdir()
        with ScratchDir(scratch, copy_to_current_on_exit=True) as d:
            with open(os.path.join(d, self.input_file), 'wt',
                      encoding="utf-8") as inp:
                for k, v in six.iteritems(self.control_params):
                    inp.write('{} {}\n'.format(k, self._format_param_val(v)))
                for idx, mol in enumerate(self.mols):
                    a = BabelMolAdaptor(mol)
                    pm = pb.Molecule(a.openbabel_mol)
                    pm.write(self.control_params["filetype"],
                             filename=os.path.join(d, '{}.{}'.format(idx,
                                 self.control_params["filetype"])).encode("ascii"),
                             overwrite=True)
                    inp.write("\n")
                    inp.write(
                        "structure {}.{}\n".format(os.path.join(d, str(idx)),
                                                   self.control_params["filetype"]))
                    for k, v in six.iteritems(self.param_list[idx]):
                        inp.write(
                            '  {} {}\n'.format(k, self._format_param_val(v)))
                    inp.write('end structure\n')
            proc = Popen(['packmol'],
                         stdin=open(os.path.join(d, self.input_file), 'r'),
                         stdout=PIPE)
            (stdout, stderr) = proc.communicate()
            output_file = os.path.join(d, self.control_params["output"])
            if os.path.isfile(output_file):
                packed_mol = BabelMolAdaptor.from_file(output_file)
                print("packed molecule written to {}".format(
                    self.control_params["output"]))
                return packed_mol.pymatgen_mol
            else:
                print("Packmol execution failed")
                print(stdout, stderr)
                return None

if __name__ == '__main__':
    ethanol_coords = [[0.00720, -0.56870, 0.00000],
                      [-1.28540, 0.24990, 0.00000],
                      [1.13040, 0.31470, 0.00000],
                      [0.03920, -1.19720, 0.89000],
                      [0.03920, -1.19720, -0.89000],
                      [-1.31750, 0.87840, 0.89000],
                      [-1.31750, 0.87840, -0.89000],
                      [-2.14220, -0.42390, -0.00000],
                      [1.98570, -0.13650, -0.00000]]
    ethanol = Molecule(["C", "C", "O", "H", "H", "H", "H", "H", "H"],
                       ethanol_coords)
    water_coords = [[9.626, 6.787, 12.673],
                    [9.626, 8.420, 12.673],
                    [10.203, 7.604, 12.673]]
    water = Molecule(["H", "H", "O"], water_coords)
    pmr = PackmolRunner([ethanol, water],
                        [{"number": 1, "fixed": [0, 0, 0, 0, 0, 0],
                          "centerofmass": ""},
                         {"number": 15, "inside sphere": [0, 0, 0, 5]}],
                        input_file="packmol_input.inp", tolerance=2.0,
                        filetype="xyz",
                        control_params={"nloop": 1000},
                        auto_box=False, output_file="cocktail.xyz")
    s = pmr.run()
<commit_msg>refactor packmolrunner<commit_after># coding: utf-8

from __future__ import division, print_function, unicode_literals, absolute_import

"""
This module defines utility classes and functions.
"""

import six
from io import open
import os
import tempfile
from subprocess import Popen, PIPE

import numpy as np

try:
    import pybel as pb
except ImportError:
    pb = None
from pymatgen import Molecule
from pymatgen.core.operations import SymmOp
from pymatgen.util.coord_utils import get_angle
from pymatgen.io.babel import BabelMolAdaptor

from monty.os.path import which
from monty.dev import requires
from monty.tempfile import ScratchDir

__author__ = 'Kiran Mathew, Michael Humbert, Brandon Wood'
__email__ = 'kmathew@lbl.gov'

class Polymer(object):
    """
    Generate polymer molecule via Random walk. At each poisition there are
    a total of 5 possible moves(exclusing the previous directions).
    """
    def __init__(self, start_monomer, s_head, s_tail,
                 monomer, head, tail,
                 end_monomer, e_head, e_tail,
                 n_units, link_distance=1.0, linear_chain=False):
        """
        Args:
            start_monomer (Molecule): Starting molecule
            s_head (int): starting atom index of the start_monomer molecule
            s_tail (int): tail atom index of the start_monomer
            monomer (Molecule): The monomer
            head (int): index of the atom in the monomer that forms the head
            tail (int): tail atom index. monomers will be connected from
                tail to head
            end_monomer (Molecule): Terminal molecule
            e_head (int): starting atom index of the end_monomer molecule
            e_tail (int): tail atom index of the end_monomer
            n_units (int): number of monomer units excluding the start and
                terminal molecules
            link_distance (float): distance between consecutive monomers
            linear_chain (bool): linear or random walk polymer chain
        """
        self.start = s_head
        self.end = s_tail
        self.monomer = monomer
        self.n_units = n_units
        self.link_distance = link_distance
        self.linear_chain = linear_chain
        # translate monomers so that head atom is at the origin
        start_monomer.translate_sites(range(len(start_monomer)),
                                      - monomer.cart_coords[s_head])
        monomer.translate_sites(range(len(monomer)),
                                - monomer.cart_coords[head])
        end_monomer.translate_sites(range(len(end_monomer)),
                                    - monomer.cart_coords[e_head])
        self.mon_vector = monomer.cart_coords[tail] - monomer.cart_coords[head]
        self.moves = {1: [1, 0, 0],
                      2: [0, 1, 0],
                      3: [0, 0, 1],
                      4: [-1, 0, 0],
                      5: [0, -1, 0],
                      6: [0, 0, -1]}
        self.prev_move = 1
        # places the start monomer at the beginning of the chain
        self.molecule = start_monomer.copy()
        self.length = 1
        # create the chain
        self._create(self.monomer, self.mon_vector)
        # terminate the chain with the end_monomer
        self.n_units += 1
        end_mon_vector = end_monomer.cart_coords[e_tail] - \
                         end_monomer.cart_coords[e_head]
        self._create(end_monomer, end_mon_vector)

    def _create(self, monomer, mon_vector):
        """
        create the polymer from the monomer

        Args:
            monomer (Molecule)
            mon_vector (numpy.array): molecule vector that starts from the
                start atom index to the end atom index
        """
        while self.length != (self.n_units-1):
            if self.linear_chain:
                move_direction = np.array(mon_vector) / np.linalg.norm(mon_vector)
            else:
                move_direction = self._next_move_direction()
            self._add_monomer(monomer.copy(), mon_vector, move_direction)

    def _next_move_direction(self):
        """
        pick a move at random from the list of moves
        """
        nmoves = len(self.moves)
        move = np.random.randint(1, nmoves+1)
        while self.prev_move == (move + 3) % nmoves:
            move = np.random.randint(1, nmoves+1)
        self.prev_move = move
        return np.array(self.moves[move])

    def _align_monomer(self, monomer, mon_vector, move_direction):
        """
        rotate the monomer so that it is aligned along the move direction

        Args:
            monomer (Molecule)
            mon_vector (numpy.array): molecule vector that starts from the
                start atom index to the end atom index
            move_direction (numpy.array): the direction of the polymer chain
                extension
        """
        axis = np.cross(mon_vector, move_direction)
        origin = monomer[self.start].coords
        angle = get_angle(mon_vector, move_direction)
        op = SymmOp.from_origin_axis_angle(origin, axis, angle)
        monomer.apply_operation(op)

    def _add_monomer(self, monomer, mon_vector, move_direction):
        """
        extend the polymer molecule by adding a monomer along mon_vector direction

        Args:
            monomer (Molecule): monomer molecule
            mon_vector (numpy.array): monomer vector that points from head to tail.
            move_direction (numpy.array): direction along which the monomer
                will be positioned
        """
        translate_by = self.molecule.cart_coords[self.end] + \
                       self.link_distance * move_direction
        monomer.translate_sites(range(len(monomer)), translate_by)
        if not self.linear_chain:
            self._align_monomer(monomer, mon_vector, move_direction)
        # add monomer if there are no crossings
        does_cross = False
        for i, site in enumerate(monomer):
            try:
                self.molecule.append(site.specie, site.coords,
                                     properties=site.properties)
            except:
                does_cross = True
                polymer_length = len(self.molecule)
                self.molecule.remove_sites(
                    range(polymer_length - i, polymer_length))
                break
        if not does_cross:
            self.length += 1
            self.end += len(self.monomer)

class PackmolRunner(object):
    """
    Wrapper for the Packmol software that can be used to pack various types of
    molecules into a one single unit.
    """

    @requires(which(b'packmol'),
              "PackmolRunner requires the executable 'packmol' to be in "
              "the path. Please download packmol from "
              "https://github.com/leandromartinez98/packmol "
              "and follow the instructions in the README to compile. "
              "Don't forget to add the packmol binary to your path")
    def __init__(self, mols, param_list, input_file="pack.inp",
                 tolerance=2.0, filetype="xyz",
                 control_params={"maxit": 20, "nloop": 600},
                 auto_box=True, output_file="packed.xyz"):
        """
        Args:
              mols:
                   list of Molecules to pack
              input_file:
                        name of the packmol input file
              tolerance:
                        min distance between the atoms
              filetype:
                       input/output structure file type
              control_params:
                           packmol control parameters dictionary. Basically
                           all parameters other than structure/atoms
              param_list:
                    list of parameters containing dicts for each molecule
              auto_box:
                    put the molecule assembly in a box
              output_file:
                    output file name. The extension will be adjusted
                    according to the filetype
        """
        self.mols = mols
        self.param_list = param_list
        self.input_file = input_file
        self.boxit = auto_box
        self.control_params = control_params
        if not self.control_params.get("tolerance"):
            self.control_params["tolerance"] = tolerance
        if not self.control_params.get("filetype"):
            self.control_params["filetype"] = filetype
        if not self.control_params.get("output"):
            self.control_params["output"] = "{}.{}".format(
                output_file.split(".")[0], self.control_params["filetype"])
        if self.boxit:
            self._set_box()

    def _format_param_val(self, param_val):
        """
        Internal method to format values in the packmol parameter dictionaries

        Args:
              param_val:
                   Some object to turn into String

        Returns:
               string representation of the object
        """
        if isinstance(param_val, list):
            return ' '.join(str(x) for x in param_val)
        else:
            return str(param_val)

    def _set_box(self):
        """
        Set the box size for the molecular assembly
        """
        net_volume = 0.0
        for idx, mol in enumerate(self.mols):
            length = max([np.max(mol.cart_coords[:, i])-np.min(mol.cart_coords[:, i])
                           for i in range(3)]) + 2.0
            net_volume += (length**3.0) * float(self.param_list[idx]['number'])
        length = net_volume**(1.0/3.0)
        for idx, mol in enumerate(self.mols):
            self.param_list[idx]['inside box'] = '0.0 0.0 0.0 {} {} {}'.format(
                length, length, length)

    def _write_input(self, input_dir="."):
        """
        Write the packmol input file to the input directory.

        Args:
            input_dir (string): path to the input directory
        """
        with open(os.path.join(input_dir, self.input_file), 'wt', encoding="utf-8") as inp:
            for k, v in six.iteritems(self.control_params):
                inp.write('{} {}\n'.format(k, self._format_param_val(v)))
            # write the structures of the constituent molecules to file and set
            # the molecule id and the corresponding filename in the packmol
            # input file.
            for idx, mol in enumerate(self.mols):
                a = BabelMolAdaptor(mol)
                pm = pb.Molecule(a.openbabel_mol)
                filename = os.path.join(
                    input_dir, '{}.{}'.format(
                        idx, self.control_params["filetype"])).encode("ascii")
                pm.write(self.control_params["filetype"], filename=filename,
                         overwrite=True)
                inp.write("\n")
                inp.write(
                    "structure {}.{}\n".format(
                        os.path.join(input_dir, str(idx)),
                        self.control_params["filetype"]))
                for k, v in six.iteritems(self.param_list[idx]):
                    inp.write('  {} {}\n'.format(k, self._format_param_val(v)))
                inp.write('end structure\n')

    def run(self, copy_to_current_on_exit=False):
        """
        Write the input file to the scratch directory, run packmol and return
        the packed molecule.

        Args:
            copy_to_current_on_exit (bool): Whether or not to copy the packmol
                input/output files from the scratch directory to the current
                directory.

        Returns:
                Molecule object
        """
        scratch = tempfile.gettempdir()
        with ScratchDir(scratch, copy_to_current_on_exit=copy_to_current_on_exit) as scratch_dir:
            self._write_input(input_dir=scratch_dir)
            packmol_bin = ['packmol']
            packmol_input = open(os.path.join(scratch_dir, self.input_file), 'r')
            p = Popen(packmol_bin, stdin=packmol_input, stdout=PIPE, stderr=PIPE)
            p.wait()
            (stdout, stderr) = p.communicate()
            output_file = os.path.join(scratch_dir, self.control_params["output"])
            if os.path.isfile(output_file):
                packed_mol = BabelMolAdaptor.from_file(output_file)
                print("packed molecule written to {}".format(
                    self.control_params["output"]))
                return packed_mol.pymatgen_mol
            else:
                print("Packmol execution failed")
                print(stdout, stderr)
                return None

if __name__ == '__main__':
    ethanol_coords = [[0.00720, -0.56870, 0.00000],
                      [-1.28540, 0.24990, 0.00000],
                      [1.13040, 0.31470, 0.00000],
                      [0.03920, -1.19720, 0.89000],
                      [0.03920, -1.19720, -0.89000],
                      [-1.31750, 0.87840, 0.89000],
                      [-1.31750, 0.87840, -0.89000],
                      [-2.14220, -0.42390, -0.00000],
                      [1.98570, -0.13650, -0.00000]]
    ethanol = Molecule(["C", "C", "O", "H", "H", "H", "H", "H", "H"],
                       ethanol_coords)
    water_coords = [[9.626, 6.787, 12.673],
                    [9.626, 8.420, 12.673],
                    [10.203, 7.604, 12.673]]
    water = Molecule(["H", "H", "O"], water_coords)
    pmr = PackmolRunner([ethanol, water],
                        [{"number": 1, "fixed": [0, 0, 0, 0, 0, 0],
                          "centerofmass": ""},
                         {"number": 15, "inside sphere": [0, 0, 0, 5]}],
                        input_file="packmol_input.inp", tolerance=2.0,
                        filetype="xyz",
                        control_params={"nloop": 1000},
                        auto_box=False, output_file="cocktail.xyz")
    s = pmr.run()
<|endoftext|>


<commit_before>
import select, socket, struct, sys
import logging, time, threading
from Queue import Queue, Empty
from pymoku import Moku, FrameTimeout

import _instrument

from pymoku import NotDeployedException

log = logging.getLogger(__name__)

class FrameQueue(Queue):
	def put(self, item, block=True, timeout=None):
		""" Behaves the same way as default except that instead of raising Full, it
		    just pushes the item on to the deque anyway, throwing away old frames."""
		self.not_full.acquire()
		try:
			if self.maxsize > 0 and block:
				if timeout is None:
					while self._qsize() == self.maxsize:
						self.not_full.wait()
				elif timeout < 0:
					raise ValueError("'timeout' must be a non-negative number")
				else:
					endtime = _time() + timeout
					while self._qsize() == self.maxsize:
						remaining = endtime - _time()
						if remaining <= 0.0:
							break
						self.not_full.wait(remaining)
			self._put(item)
			self.unfinished_tasks += 1
			self.not_empty.notify()
		finally:
			self.not_full.release()

class DataFrame(object):
	"""
	.. mysterious undocumented magic.

	.. autoinstanceattribute:: pymoku._frame_instrument.DataFrame.ch1
		:annotation: = [CH1_DATA]

	.. autoinstanceattribute:: pymoku._frame_instrument.DataFrame.ch1
		:annotation: = [CH2_DATA]

	.. autoinstanceattribute:: pymoku._frame_instrument.DataFrame.frameid
		:annotation: = n

	.. autoinstanceattribute:: pymoku._frame_instrument.DataFrame.waveformid
		:annotation: = n
	"""
	def __init__(self):
		self.complete = False
		self.slices = [False] * 6

		#: Channel 1 data array. Present whether or not the channel is enabled, but the contents
		#: are undefined in the latter case.
		self.ch1 = []

		#: Channel 2 data array.
		self.ch2 = []

		self.stateid = None
		self.trigstate = None

		#: Frame number. Increments monotonically but wraps at 16-bits.
		self.frameid = 0

		#: Incremented once per trigger event. Wraps at 32-bits.
		self.waveformid = 0

		self.flags = None

	def add_packet(self, packet):
		hdr_len = 13
		smpls = 342
		d_len = smpls * 4
		if len(packet) != hdr_len + d_len:
			return

		data = struct.unpack('<BHBBBBBIB', packet[:hdr_len])
		frameid = data[1]
		instrid = data[2]
		chan = (data[3] >> 4) & 0x0F
		sliceid = data[3] & 0x0F

		self.stateid = data[4]
		self.trigstate = data[5]
		self.flags = data[6]
		self.waveformid = data[7]

		if self.frameid != frameid:
			self.frameid = frameid
			self.slices = [False] * 6

		dat = struct.unpack('<' + 'i' * smpls, packet[hdr_len:])
		dat = [ x if x != -0x80000000 else None for x in dat ]

		if chan == 0:
			self.slices[sliceid] = True
			self.ch1[sliceid * d_len : (sliceid + 1) * d_len] = dat
			self.ch1 = self.ch1[:1024]
		else:
			self.slices[sliceid + 3] = True
			self.ch2[sliceid * d_len : (sliceid + 1) * d_len] = dat
			self.ch2 = self.ch2[:1024]

		self.complete = (self.slices == [True, True, True, True, True, True])

# Revisit: Should this be a Mixin? Are there more instrument classifications of this type, recording ability, for example?
class FrameBasedInstrument(_instrument.MokuInstrument):
	""" FrameWaffle """
	def __init__(self):
		super(FrameBasedInstrument, self).__init__()
		self._buflen = 1
		self._queue = FrameQueue(maxsize=self._buflen)
		self._hb_forced = False

	def flush(self):
		""" Clear the Frame Buffer.
		This is normally not required as one can simply wait for the correctly-generated frames to propagate through
		using the appropriate arguments to :any:`get_frame`."""
		with self._queue.mutex:
			self._queue.queue.clear()

	def set_buffer_length(self, buflen):
		""" Set the internal frame buffer length."""
		self._buflen = buflen
		self._queue = FrameQueue(maxsize=buflen)

	def get_buffer_length(self):
		""" Return the current length of the internal frame buffer """
		return self._buflen

	def get_frame(self, timeout=None, wait=True):
		""" Get a :any:`DataFrame` from the internal frame buffer"""
		try:
			# Dodgy hack, infinite timeout gets translated in to just an exceedingly long one
			endtime = time.time() + (timeout or sys.maxint)
			while self._running:
				frame = self._queue.get(block=True, timeout=timeout)
				# Should really just wait for the new stateid to propagte through, but
				# at the moment we don't support stateid and stateid_alt being different;
				# i.e. we can't rerender already aquired data. Until we fix this, wait
				# for a trigger to propagate through so we don't at least render garbage
				if not wait or frame.trigstate == self._stateid:
					return frame
				elif time.time() > endtime:
					raise FrameTimeout()
		except Empty:
			raise FrameTimeout()

	def datalogger_start(self, duration=0):
		""" Start recording data with the current settings.
		It is up to the user to ensure that the current aquisition rate is sufficiently slow to not loose samples"""
		if self._moku is None: raise NotDeployedException()
		# TODO: rest of the options, handle errors
		self._moku._stream_start(end=duration)

	def datalogger_stop(self):
		""" Stop a recording session previously started with :py:func:`datalogger_start`"""
		if self._moku is None: raise NotDeployedException()
		# TODO: Handle errors
		self._moku._stream_stop()

	def datalogger_status(self):
		""" Return the status of the most recent recording session to be started.
		This is still valid after the stream has stopped, in which case the status will reflect that it's safe
		to start a new session"""
		if self._moku is None: raise NotDeployedException()
		return self._moku._stream_status()[0]

	def datalogger_transferred(self):
		""" Return the number of samples recorded in the most recent session.
		This is valid both for running sessions and one that has been completed."""
		if self._moku is None: raise NotDeployedException()
		return self._moku._stream_status()[1]

	def set_running(self, state):
		prev_state = self._running
		super(FrameBasedInstrument, self).set_running(state)
		if state and not prev_state:
			self._fr_worker = threading.Thread(target=self._frame_worker)
			self._hb_worker = threading.Thread(target=self._heartbeat_worker)
			self._fr_worker.start()
			self._hb_worker.start()
		elif not state and prev_state:
			self._fr_worker.join()
			self._hb_worker.join()

	def _send_heartbeat(self, hbs, port):
		try:
			d, a = hbs.recvfrom(1024)
			if len(d) >= 3 and d[0] == '@':
				self._hb_forced = True
		except socket.timeout:
			pass

		try:
			hdr = 0x41 if self._hb_forced else 0x40
			ts = int(time.time() % 2**16)
			hbs.sendto(struct.pack('<BH', hdr, ts), (self._moku._ip, port))
		except:
			# TODO: Catch only what we expect, which is moku disappeared and whatever
			# the network layer can throw at us
			log.exception("HB")

	def _frame_worker(self):
		import select

		fs = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
		fs.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
		fs.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1)
		fs.bind(('0.0.0.0', Moku.PORT))

		fr = DataFrame()

		try:
			while self._running:
				if fs in select.select([fs], [], [], 1.0)[0]:
					d, a = fs.recvfrom(4096)
					if a[0] != self._moku._ip:
						log.debug("Thowing away data from Moku %s", a[0])
						continue

					fr.add_packet(d)

					if fr.complete:
						self._queue.put_nowait(fr)
						fr = DataFrame()
		finally:
			fs.close()

	def _heartbeat_worker(self):
		hs = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
		hs.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
		hs.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1)
		hs.settimeout(0.1)
		hs.bind(('0.0.0.0', Moku.PORT + 1))

		try:
			while self._running:
				self._send_heartbeat(hs, Moku.PORT + 1)
				time.sleep(1.0)
		finally:
			hs.close()

<commit_msg>Frame: Move to new ZMQ frame socket. Requires Moku v1.1 or later<commit_after>
import select, socket, struct, sys
import logging, time, threading
from Queue import Queue, Empty
from pymoku import Moku, FrameTimeout

import _instrument

from pymoku import NotDeployedException

log = logging.getLogger(__name__)

class FrameQueue(Queue):
	def put(self, item, block=True, timeout=None):
		""" Behaves the same way as default except that instead of raising Full, it
		    just pushes the item on to the deque anyway, throwing away old frames."""
		self.not_full.acquire()
		try:
			if self.maxsize > 0 and block:
				if timeout is None:
					while self._qsize() == self.maxsize:
						self.not_full.wait()
				elif timeout < 0:
					raise ValueError("'timeout' must be a non-negative number")
				else:
					endtime = _time() + timeout
					while self._qsize() == self.maxsize:
						remaining = endtime - _time()
						if remaining <= 0.0:
							break
						self.not_full.wait(remaining)
			self._put(item)
			self.unfinished_tasks += 1
			self.not_empty.notify()
		finally:
			self.not_full.release()

class DataFrame(object):
	"""
	.. mysterious undocumented magic.

	.. autoinstanceattribute:: pymoku._frame_instrument.DataFrame.ch1
		:annotation: = [CH1_DATA]

	.. autoinstanceattribute:: pymoku._frame_instrument.DataFrame.ch1
		:annotation: = [CH2_DATA]

	.. autoinstanceattribute:: pymoku._frame_instrument.DataFrame.frameid
		:annotation: = n

	.. autoinstanceattribute:: pymoku._frame_instrument.DataFrame.waveformid
		:annotation: = n
	"""
	def __init__(self):
		self.complete = False
		self.chs_valid = [False, False]

		#: Channel 1 data array. Present whether or not the channel is enabled, but the contents
		#: are undefined in the latter case.
		self.ch1 = []

		#: Channel 2 data array.
		self.ch2 = []

		self.stateid = None
		self.trigstate = None

		#: Frame number. Increments monotonically but wraps at 16-bits.
		self.frameid = 0

		#: Incremented once per trigger event. Wraps at 32-bits.
		self.waveformid = 0

		self.flags = None

	def add_packet(self, packet):
		hdr_len = 13
		smpls = 1026
		d_len = smpls * 4
		if len(packet) != hdr_len + d_len:
			log.warning("Corrupt frame recevied")
			return

		data = struct.unpack('<BHBBBBBIB', packet[:hdr_len])
		frameid = data[1]
		instrid = data[2]
		chan = (data[3] >> 4) & 0x0F

		self.stateid = data[4]
		self.trigstate = data[5]
		self.flags = data[6]
		self.waveformid = data[7]

		if self.frameid != frameid:
			self.frameid = frameid
			self.chs_valid = [False, False]

		dat = struct.unpack('<' + 'i' * smpls, packet[hdr_len:])
		dat = [ x if x != -0x80000000 else None for x in dat ]

		# For historical reasons the data length is 1026 while there are only 1024
		# valid samples. Trim the fat.
		if chan == 0:
			self.chs_valid[0] = True
			self.ch1 = dat[:1024]
		else:
			self.chs_valid[1] = True
			self.ch2 = dat[:1024]

		self.complete = all(self.chs_valid)

# Revisit: Should this be a Mixin? Are there more instrument classifications of this type, recording ability, for example?
class FrameBasedInstrument(_instrument.MokuInstrument):
	""" FrameWaffle """
	def __init__(self):
		super(FrameBasedInstrument, self).__init__()
		self._buflen = 1
		self._queue = FrameQueue(maxsize=self._buflen)
		self._hb_forced = False

	def flush(self):
		""" Clear the Frame Buffer.
		This is normally not required as one can simply wait for the correctly-generated frames to propagate through
		using the appropriate arguments to :any:`get_frame`."""
		with self._queue.mutex:
			self._queue.queue.clear()

	def set_buffer_length(self, buflen):
		""" Set the internal frame buffer length."""
		self._buflen = buflen
		self._queue = FrameQueue(maxsize=buflen)

	def get_buffer_length(self):
		""" Return the current length of the internal frame buffer """
		return self._buflen

	def get_frame(self, timeout=None, wait=True):
		""" Get a :any:`DataFrame` from the internal frame buffer"""
		try:
			# Dodgy hack, infinite timeout gets translated in to just an exceedingly long one
			endtime = time.time() + (timeout or sys.maxint)
			while self._running:
				frame = self._queue.get(block=True, timeout=timeout)
				# Should really just wait for the new stateid to propagte through, but
				# at the moment we don't support stateid and stateid_alt being different;
				# i.e. we can't rerender already aquired data. Until we fix this, wait
				# for a trigger to propagate through so we don't at least render garbage
				if not wait or frame.trigstate == self._stateid:
					return frame
				elif time.time() > endtime:
					raise FrameTimeout()
		except Empty:
			raise FrameTimeout()

	def datalogger_start(self, duration=0):
		""" Start recording data with the current settings.
		It is up to the user to ensure that the current aquisition rate is sufficiently slow to not loose samples"""
		if self._moku is None: raise NotDeployedException()
		# TODO: rest of the options, handle errors
		self._moku._stream_start(end=duration)

	def datalogger_stop(self):
		""" Stop a recording session previously started with :py:func:`datalogger_start`"""
		if self._moku is None: raise NotDeployedException()
		# TODO: Handle errors
		self._moku._stream_stop()

	def datalogger_status(self):
		""" Return the status of the most recent recording session to be started.
		This is still valid after the stream has stopped, in which case the status will reflect that it's safe
		to start a new session"""
		if self._moku is None: raise NotDeployedException()
		return self._moku._stream_status()[0]

	def datalogger_transferred(self):
		""" Return the number of samples recorded in the most recent session.
		This is valid both for running sessions and one that has been completed."""
		if self._moku is None: raise NotDeployedException()
		return self._moku._stream_status()[1]

	def set_running(self, state):
		prev_state = self._running
		super(FrameBasedInstrument, self).set_running(state)
		if state and not prev_state:
			self._fr_worker = threading.Thread(target=self._frame_worker)
			self._hb_worker = threading.Thread(target=self._heartbeat_worker)
			self._fr_worker.start()
			self._hb_worker.start()
		elif not state and prev_state:
			self._fr_worker.join()
			self._hb_worker.join()

	def _send_heartbeat(self, hbs, port):
		try:
			d, a = hbs.recvfrom(1024)
			if len(d) >= 3 and d[0] == '@':
				self._hb_forced = True
		except socket.timeout:
			pass

		try:
			hdr = 0x41 if self._hb_forced else 0x40
			ts = int(time.time() % 2**16)
			hbs.sendto(struct.pack('<BH', hdr, ts), (self._moku._ip, port))
		except:
			# TODO: Catch only what we expect, which is moku disappeared and whatever
			# the network layer can throw at us
			log.exception("HB")

	def _frame_worker(self):
		import zmq

		ctx = zmq.Context()
		skt = ctx.socket(zmq.SUB)
		skt.connect("tcp://%s:27185" % self._moku._ip)
		skt.setsockopt_string(zmq.SUBSCRIBE, u'')

		fr = DataFrame()

		try:
			while self._running:
				if skt in zmq.select([skt], [], [], 1.0)[0]:
					d = skt.recv()
					fr.add_packet(d)

					if fr.complete:
						self._queue.put_nowait(fr)
						fr = DataFrame()
		finally:
			skt.close()

	def _heartbeat_worker(self):
		hs = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
		hs.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
		hs.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1)
		hs.settimeout(0.1)
		hs.bind(('0.0.0.0', Moku.PORT + 1))

		try:
			while self._running:
				self._send_heartbeat(hs, Moku.PORT + 1)
				time.sleep(1.0)
		finally:
			hs.close()

<|endoftext|>


<commit_before>import os
import argparse
import urllib2
import shutil
from distutils.file_util import copy_file, move_file
from distutils.dir_util import mkpath, remove_tree

from pyne.api import nuc_data
from pyne.dbgen.api import build_dir
from pyne.dbgen.decay import make_decay
from pyne.dbgen.atomic_weight import make_atomic_weight
from pyne.dbgen.scattering_lengths import make_scattering_lengths
from pyne.dbgen.simple_xs import make_simple_xs
from pyne.dbgen.cinder import make_cinder

# Thanks to http://patorjk.com/software/taag/
# and http://www.chris.com/ascii/index.php?art=creatures/dragons (Jeff Ferris)
# for ASCII art inspiriation

pyne_logo = """\

                                  /   \
 _                        )      ((   ))     (
(@)                      /|\      ))_((     /|\
|-|                     / | \    (/\|/\)   / | \                      (@)
| | -------------------/--|-voV---\`|'/--Vov-|--\---------------------|-|
|-|                         '^`   (o o)  '^`                          | |
| |                               `\Y/'                               |-|
|-|                                                                   | |
| |        /\             ___           __  __             /\         |-|
|-|       /^~\           / _ \_   _  /\ \ \/__\           /^~\        | |
| |       /^~\          / /_)/ | | |/  \/ /_\             /^~\        |-|
|-|       /^~\         / ___/| |_| / /\  //__             /^~\        | |
| |       ^||`         \/     \__, \_\ \/\__/             ^||`        |-|
|-|        ||                |____/                        ||         | |
| |       ====                                            ====        |-|
|-|                                                                   | |
| |                                                                   |-|
|-|___________________________________________________________________| |
(@)              l   /\ /         ( (       \ /\   l                `\|-|
                 l /   V           \ \       V   \ l                  (@)
                 l/                _) )_          \I
                                   `\ /'
                                     `
"""

def _fetch_prebuilt(args):
    nuc_data, build_dir = args.nuc_data, args.build_dir
    prebuilt_nuc_data = os.path.join(build_dir, 'prebuilt_nuc_data.h5')
    if os.path.exists(prebuilt_nuc_data):
        return

    pnd = urllib2.urlopen("<PUT ME IN>")
    with open(prebuilt_nuc_data, 'wb') as f:
        f.write(pnd.read())

    shutil.copyfile(prebuilt_nuc_data, nuc_data)

def main():
    """Entry point for nuc_data_make utility."""
    print pyne_logo

    make_funcs = [('atomic_weight', make_atomic_weight),
                  ('scattering_lengths', make_scattering_lengths),
                  ('decay', make_decay),
                  ('simple_xs', make_simple_xs),
                  ('cinder', make_cinder),
                  ]
    make_map = dict(make_funcs)

    # Parse the command line arguments
    parser = argparse.ArgumentParser(description='Make a nuclear data library.')
    parser.add_argument('-o', dest='nuc_data', action='store', default=nuc_data,
                        help='path to the output database file.')
    parser.add_argument('-b', dest='build_dir', action='store', default=build_dir,
                        help='path to the build directory.')
    parser.add_argument('--datapath', dest='datapath', action='store', default="",
                        help='MCNP DATAPATH.')
    parser.add_argument('--fetch-prebuilt', dest='fetch_prebuilt', action='store',
                        type=lambda s: 't' in s.lower() or 'y' in s.lower(),
                        default=True, help='grab partially assembled file [y/n].')
    parser.add_argument('-m', dest='make', action='store', default='all',
                        help='comma-separated parts of nuc_data to make: ' + \
                        ", ".join([mf[0] for mf in make_funcs]) + ', all, and none.')
    parser.add_argument('--clean', dest='clean', type=int, default=0,
                        help="""level to clean up existing files.
                                0: no cleaning (default).
                                1: clean nuc_data.
                                2: clean nuc_data and build_dir.""")
    args = parser.parse_args()

    # clean nuc data
    if args.clean in [1, 2]:
        print "removing nuc_data from {0}".format(args.nuc_data)
        try:
            os.remove(args.nuc_data)
        except OSError:
            pass

    # Make the build dir
    if args.clean == 2:
        print "removing build_dir from {0}".format(args.build_dir)
        remove_tree(args.build_dir)
    mkpath(args.build_dir)

    # Determine what to make
    if args.make == 'none':
        make_order = []
    elif args.make == 'all':
        make_order = [mf[0] for mf in make_funcs]
    else:
        make_order = args.make.replace(' ', "").split(',')

    # fetch prebuilt data library if possible
    if args.fetch_prebuilt:
        _fetch_prebuilt(args)

    # Make the various tables
    print "Making nuc_data at {0}".format(args.nuc_data)
    for mo in make_order:
        make_map[mo](args)

if __name__ == '__main__':
    main()
<commit_msg>Added --make-open-only to nuc_data_make.<commit_after>import os
import argparse
import urllib2
import shutil
from distutils.file_util import copy_file, move_file
from distutils.dir_util import mkpath, remove_tree

from pyne.api import nuc_data
from pyne.dbgen.api import build_dir
from pyne.dbgen.decay import make_decay
from pyne.dbgen.atomic_weight import make_atomic_weight
from pyne.dbgen.scattering_lengths import make_scattering_lengths
from pyne.dbgen.simple_xs import make_simple_xs
from pyne.dbgen.cinder import make_cinder

# Thanks to http://patorjk.com/software/taag/
# and http://www.chris.com/ascii/index.php?art=creatures/dragons (Jeff Ferris)
# for ASCII art inspiriation

pyne_logo = """\

                                  /   \
 _                        )      ((   ))     (
(@)                      /|\      ))_((     /|\
|-|                     / | \    (/\|/\)   / | \                      (@)
| | -------------------/--|-voV---\`|'/--Vov-|--\---------------------|-|
|-|                         '^`   (o o)  '^`                          | |
| |                               `\Y/'                               |-|
|-|                                                                   | |
| |        /\             ___           __  __             /\         |-|
|-|       /^~\           / _ \_   _  /\ \ \/__\           /^~\        | |
| |       /^~\          / /_)/ | | |/  \/ /_\             /^~\        |-|
|-|       /^~\         / ___/| |_| / /\  //__             /^~\        | |
| |       ^||`         \/     \__, \_\ \/\__/             ^||`        |-|
|-|        ||                |____/                        ||         | |
| |       ====                                            ====        |-|
|-|                                                                   | |
| |                                                                   |-|
|-|___________________________________________________________________| |
(@)              l   /\ /         ( (       \ /\   l                `\|-|
                 l /   V           \ \       V   \ l                  (@)
                 l/                _) )_          \I
                                   `\ /'
                                     `
"""

def _fetch_prebuilt(args):
    nuc_data, build_dir = args.nuc_data, args.build_dir
    prebuilt_nuc_data = os.path.join(build_dir, 'prebuilt_nuc_data.h5')
    prebuilt_nuc_data_url = "<PUT ME IN>"
    if os.path.exists(prebuilt_nuc_data):
        return

    print "Fetching pre-built nuc_data.h5 from " + prebuilt_nuc_data_url
    pnd = urllib2.urlopen(prebuilt_nuc_data_url)
    with open(prebuilt_nuc_data, 'wb') as f:
        f.write(pnd.read())

    shutil.copyfile(prebuilt_nuc_data, nuc_data)

def main():
    """Entry point for nuc_data_make utility."""
    print pyne_logo

    make_funcs = [('atomic_weight', make_atomic_weight),
                  ('scattering_lengths', make_scattering_lengths),
                  ('decay', make_decay),
                  ('simple_xs', make_simple_xs),
                  ('cinder', make_cinder),
                  ]
    make_map = dict(make_funcs)
    make_open = {'scattering_lengths', 'simple_xs'}

    # Parse the command line arguments
    parser = argparse.ArgumentParser(description='Make a nuclear data library.')
    parser.add_argument('-o', dest='nuc_data', action='store', default=nuc_data,
                        help='path to the output database file.')
    parser.add_argument('-b', dest='build_dir', action='store', default=build_dir,
                        help='path to the build directory.')
    parser.add_argument('--datapath', dest='datapath', action='store', default="",
                        help='MCNP DATAPATH.')
    parser.add_argument('--fetch-prebuilt', dest='fetch_prebuilt', action='store',
                        type=lambda s: 't' in s.lower() or 'y' in s.lower(),
                        default=True, help='grab partially assembled file [y/n].')
    parser.add_argument('--make-open-only', dest='make_open_only', action='store',
                        type=lambda s: 't' in s.lower() or 'y' in s.lower(),
                        default=False, help='only add open data to file [y/n].')
    parser.add_argument('-m', dest='make', action='store', default='all',
                        help='comma-separated parts of nuc_data to make: ' + \
                        ", ".join([mf[0] for mf in make_funcs]) + ', all, and none.')
    parser.add_argument('--clean', dest='clean', type=int, default=0,
                        help="""level to clean up existing files.
                                0: no cleaning (default).
                                1: clean nuc_data.
                                2: clean nuc_data and build_dir.""")
    args = parser.parse_args()

    # clean nuc data
    if args.clean in [1, 2]:
        print "removing nuc_data from {0}".format(args.nuc_data)
        try:
            os.remove(args.nuc_data)
        except OSError:
            pass

    # Make the build dir
    if args.clean == 2:
        print "removing build_dir from {0}".format(args.build_dir)
        remove_tree(args.build_dir)
    mkpath(args.build_dir)

    # Determine what to make
    if args.make == 'none':
        make_order = []
    elif args.make == 'all':
        make_order = [mf[0] for mf in make_funcs]
    else:
        make_order = args.make.replace(' ', "").split(',')

    if args.make_open_only:
        make_order = [mo for mo in make_order if mo in make_open]

    # fetch prebuilt data library if possible
    if args.fetch_prebuilt:
        _fetch_prebuilt(args)

    # Make the various tables
    print "Making nuc_data at {0}".format(args.nuc_data)
    for mo in make_order:
        make_map[mo](args)

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


<commit_before>
import os

from app_config import get_project_config_path, update_config_with_sections, get_config_section

class Status(object):
    class StatusType:
        UPLOAD_VIDEO = "upload_video"
        UPLOAD_HOMOGRAPHY = "upload_homography"
        CONFIGURATION_TEST = "configuration_test"
        FEATURE_TRACKING = "feature_tracking"
        OBJECT_TRACKING = "object_tracking"
        SAFETY_ANALYSIS = "safety_analysis"
        HIGHLIGHT_VIDEO = "highlight_video"

    class Status:
        INCOMPLETE = 0
        IN_PROGRESS = 1
        COMPLETE = 2

    @staticmethod
    def create_status_dict():
        return {
            self.Type.UPLOAD_VIDEO: 0,
            self.Type.UPLOAD_HOMOGRAPHY: 0,
            self.Type.CONFIGURATION_TEST: 0,
            self.Type.FEATURE_TRACKING: 0,
            self.Type.OBJECT_TRACKING: 0,
            self.Type.SAFETY_ANALYSIS: 0,
            self.Type.HIGHLIGHT_VIDEO: 0,
        }

class StatusHelper(object):
    @staticmethod
    def initalize_project(identifier):
        d = Status.create_status_dict()
        config_path = get_project_config_path(identifier)
        for (status_type, status) in d.iteritems():
            update_config_with_sections(config_path, "status", status_type, status)

    @staticmethod
    def set_status(identifier, status_type, status):
        config_path = get_project_config_path(identifier)
        update_config_with_sections(config_path, "status", status_type, status)

    @staticmethod
    def get_status(identifier):
        config_path = get_project_config_path(identifier)
        (success, value) = get_config_section(config_path, "status")
        if success:
            return value
        else:
            return None

<commit_msg>Optimizing statusHandler with failure code<commit_after>
import os

from app_config import get_project_config_path, update_config_with_sections, get_config_section

class Status(object):
    class StatusType:
        UPLOAD_HOMOGRAPHY = "upload_homography"
        CONFIGURATION_TEST = "configuration_test"
        FEATURE_TRACKING = "feature_tracking"
        OBJECT_TRACKING = "object_tracking"
        SAFETY_ANALYSIS = "safety_analysis"
        HIGHLIGHT_VIDEO = "highlight_video"

    class Status:
        FAILURE = -1
        INCOMPLETE = 0
        IN_PROGRESS = 1
        COMPLETE = 2

    @staticmethod
    def create_status_dict():
        return {
            self.Type.UPLOAD_HOMOGRAPHY: 0,
            self.Type.CONFIGURATION_TEST: 0,
            self.Type.FEATURE_TRACKING: 0,
            self.Type.OBJECT_TRACKING: 0,
            self.Type.SAFETY_ANALYSIS: 0,
            self.Type.HIGHLIGHT_VIDEO: 0,
        }

class StatusHelper(object):
    @staticmethod
    def initalize_project(identifier):
        d = Status.create_status_dict()
        config_path = get_project_config_path(identifier)
        for (status_type, status) in d.iteritems():
            update_config_with_sections(config_path, "status", status_type, status)

    @staticmethod
    def set_status(identifier, status_type, status):
        config_path = get_project_config_path(identifier)
        update_config_with_sections(config_path, "status", status_type, status)

    @staticmethod
    def get_status(identifier):
        config_path = get_project_config_path(identifier)
        (success, value) = get_config_section(config_path, "status")
        if success:
            return value
        else:
            return None

<|endoftext|>


<commit_before>#!/usr/bin/python2.7
# -*- encoding=utf-8 -*-

from argparse import ArgumentParser
import codecs
import gevent
from gevent import monkey
import json
from types import UnicodeType

from crawlers import Crawler
from crawlers.local.static import get_election_type_name
from utils import check_dir

def print_json(filename, data):
    with open(filename, 'w') as f:
        json.dump(data, f, encoding="UTF-8", indent=2)

def print_csv(filename, data):

    def transform(txt):
        if isinstance(txt, int):
            txt = str(txt)
        if isinstance(txt, list):
            txt = '||'.join(txt)
        txt = txt.replace(',', '|')
        if isinstance(txt, UnicodeType):
            txt = txt.encode('utf8')
        return txt

    attrs = ['assembly_no', 'district', 'cand_no', 'party', 'name_kr',
             'name_cn', 'sex', 'birthyear', 'birthmonth', 'birthday',
             'address', 'job', 'education', 'experience', 'recommend_priority',
             'votenum', 'voterate', 'elected']

    with open(filename, 'w') as f:
        f.write(codecs.BOM_UTF8)
        f.write(','.join(attrs))
        f.write('\n')
        for cand in data:
            values = (cand[attr] if attr in cand else '' for attr in attrs)
            values = (transform(value) for value in values)
            f.write(','.join(values))
            f.write('\n')

def crawl(target, _type, nth, printer, filename, level=None):
    crawler = Crawler(target, _type, nth, level)
    cand_list = crawler.crawl()
    printer(filename, cand_list)

def create_parser():
    parser = ArgumentParser()
    parser.add_argument('target', choices=['assembly', 'local', 'president'],\
            help="name of target election")
    parser.add_argument('type', choices=['candidates', 'elected', 'precandidates'],
            help="type of person")
    parser.add_argument('start', help="starting election id", type=float)
    parser.add_argument('end', help="ending election id", type=float,\
            nargs='?', default=None)
    parser.add_argument('-t', dest='test', action='store_true')
    parser.add_argument('-d', '--directory', help="specify data directory")

    # TODO: change to subparser
    parser.add_argument('-l', '--level', choices=[
        'pg', 'province_governor',
        'pm', 'province_member',
        'mg', 'municipality_governor',
        'mm', 'municipality_member',
        'eg', 'education_governor',
        'em', 'education_member'],
        help="Specify level for local elections")
    return  parser

def main(args):
    printer = print_csv if args.test else print_json
    filetype = 'csv' if args.test else 'json'
    datadir = args.directory if args.directory else '.'
    check_dir(datadir)

    if args.target=='local':
        if args.end:
            jobs = []
            args.level = get_election_type_name(args.level)
            for n in xrange(args.start, args.end+1):
                filename = '%s/%s-%s-%s-%d.%s'\
                    % (datadir, args.target, args.level, args.type, n, filetype)
                job = gevent.spawn(crawl, target=args.target, level=args.level,\
                    _type=args.type, nth=n, filename=filename, printer=printer)
                jobs.append(job)
            gevent.joinall(jobs)
        else:
            n = args.start
            args.level = get_election_type_name(args.level)
            filename = '%s/%s-%s-%s-%.01f.%s' %\
                    (datadir, args.target, args.level, args.type, n, filetype)
            crawl(target=args.target, level=args.level, _type=args.type, nth=n,\
                        filename=filename, printer=printer)
    else:
        if args.end:
            jobs = []
            for n in xrange(args.start, args.end+1):
                filename = '%s/%s-%s-%d.%s'\
                        % (datadir, args.target, args.type, n, filetype)
                job = gevent.spawn(crawl, target=args.target, _type=args.type, nth=n,\
                        filename=filename, printer=printer)
                jobs.append(job)
            gevent.joinall(jobs)
        else:
            n = args.start
            filename = '%s/%s-%s-%.01f.%s' %\
                    (datadir, args.target, args.type, n, filetype)
            crawl(target=args.target, _type=args.type, nth=n,\
                        filename=filename, printer=printer)
    print 'Data written to %s' % filename

if __name__ == '__main__':
    monkey.patch_all()
    parser = create_parser()
    args = parser.parse_args()
    main(args)
<commit_msg>Reformat election commission crawler help<commit_after>#!/usr/bin/python2.7
# -*- encoding=utf-8 -*-

from argparse import ArgumentParser, RawTextHelpFormatter
import codecs
import gevent
from gevent import monkey
import json
from types import UnicodeType

from crawlers import Crawler
from crawlers.local.static import get_election_type_name
from utils import check_dir

def print_json(filename, data):
    with open(filename, 'w') as f:
        json.dump(data, f, encoding="UTF-8", indent=2)

def print_csv(filename, data):

    def transform(txt):
        if isinstance(txt, int):
            txt = str(txt)
        if isinstance(txt, list):
            txt = '||'.join(txt)
        txt = txt.replace(',', '|')
        if isinstance(txt, UnicodeType):
            txt = txt.encode('utf8')
        return txt

    attrs = ['assembly_no', 'district', 'cand_no', 'party', 'name_kr',
             'name_cn', 'sex', 'birthyear', 'birthmonth', 'birthday',
             'address', 'job', 'education', 'experience', 'recommend_priority',
             'votenum', 'voterate', 'elected']

    with open(filename, 'w') as f:
        f.write(codecs.BOM_UTF8)
        f.write(','.join(attrs))
        f.write('\n')
        for cand in data:
            values = (cand[attr] if attr in cand else '' for attr in attrs)
            values = (transform(value) for value in values)
            f.write(','.join(values))
            f.write('\n')

def crawl(target, _type, nth, printer, filename, level=None):
    crawler = Crawler(target, _type, nth, level)
    cand_list = crawler.crawl()
    printer(filename, cand_list)

def create_parser():
    parser = ArgumentParser(formatter_class=RawTextHelpFormatter)
    parser.add_argument('target', choices=['assembly', 'local', 'president'],\
            help="name of target election")
    parser.add_argument('type', choices=['candidates', 'elected', 'precandidates'],
            help="type of person")
    parser.add_argument('start', help="starting election id", type=float)
    parser.add_argument('end', help="ending election id", type=float,\
            nargs='?', default=None)
    parser.add_argument('-t', '--test', dest='test', action='store_true',
            help="assign datatype to csv instead of json")
    parser.add_argument('-d', dest='directory', help="specify data directory")

    # TODO: change to subparser
    parser.add_argument('-l', choices=['pg', 'pm', 'pp', 'mg', 'mm', 'mp', 'eg', 'em'],
        dest="level",
        help="specify level for local elections.\n"
            "- 1st char: {p:province, m:municipality, e:education},\n"
            "- 2nd char: {g: governor, m: member}")
    return  parser

def main(args):
    printer = print_csv if args.test else print_json
    filetype = 'csv' if args.test else 'json'
    datadir = args.directory if args.directory else '.'
    check_dir(datadir)

    if args.target=='local':
        if args.end:
            jobs = []
            args.level = get_election_type_name(args.level)
            for n in xrange(args.start, args.end+1):
                filename = '%s/%s-%s-%s-%d.%s'\
                    % (datadir, args.target, args.level, args.type, n, filetype)
                job = gevent.spawn(crawl, target=args.target, level=args.level,\
                    _type=args.type, nth=n, filename=filename, printer=printer)
                jobs.append(job)
            gevent.joinall(jobs)
        else:
            n = args.start
            args.level = get_election_type_name(args.level)
            filename = '%s/%s-%s-%s-%.01f.%s' %\
                    (datadir, args.target, args.level, args.type, n, filetype)
            crawl(target=args.target, level=args.level, _type=args.type, nth=n,\
                        filename=filename, printer=printer)
    else:
        if args.end:
            jobs = []
            for n in xrange(args.start, args.end+1):
                filename = '%s/%s-%s-%d.%s'\
                        % (datadir, args.target, args.type, n, filetype)
                job = gevent.spawn(crawl, target=args.target, _type=args.type, nth=n,\
                        filename=filename, printer=printer)
                jobs.append(job)
            gevent.joinall(jobs)
        else:
            n = args.start
            filename = '%s/%s-%s-%.01f.%s' %\
                    (datadir, args.target, args.type, n, filetype)
            crawl(target=args.target, _type=args.type, nth=n,\
                        filename=filename, printer=printer)
    print 'Data written to %s' % filename

if __name__ == '__main__':
    monkey.patch_all()
    parser = create_parser()
    args = parser.parse_args()
    main(args)
<|endoftext|>