<commit_before># Calculates OPR/DPR/CCWM
# For implementation details, see
# http://www.chiefdelphi.com/forums/showpost.php?p=484220&postcount=19

# M is n x n where n is # of teams
# s is n x 1 where n is # of teams
# solve [M][x]=[s] for [x]

# x is OPR and should be n x 1

from collections import defaultdict
import numpy as np

from google.appengine.api import memcache

from database import event_query
from models.event_team import EventTeam

class MatchstatsHelper(object):
    @classmethod
    def build_team_mapping(cls, matches):
        """
        Returns (team_list, team_id_map)
        team_list: A list of team_str such as '254' or '254B'
        team_id_map: A dict of key: team_str, value: row index in x_matrix that corresponds to the team
        """
        # Build team list
        team_list = set()
        for match in matches:
            if match.comp_level != 'qm':  # only consider quals matches
                continue
            for alliance_color in ['red', 'blue']:
                for team in match.alliances[alliance_color]['teams']:
                    team_list.add(team[3:])  # turns "frc254B" into "254B"

        team_list = list(team_list)
        team_id_map = {}
        for i, team in enumerate(team_list):
            team_id_map[team] = i

        return team_list, team_id_map

    @classmethod
    def build_Minv_matrix(cls, matches, team_id_map, played_only=False):
        n = len(team_id_map.keys())
        M = np.zeros([n, n])
        for match in matches:
            if match.comp_level != 'qm':  # only consider quals matches
                continue
            if played_only and not match.has_been_played:
                continue
            for alliance_color in ['red', 'blue']:
                alliance_teams = match.alliances[alliance_color]['teams']
                for team1 in alliance_teams:
                    team1_id = team_id_map[team1[3:]]
                    for team2 in alliance_teams:
                        M[team1_id, team_id_map[team2[3:]]] += 1
        return np.linalg.pinv(M)

    @classmethod
    def build_s_matrix(cls, matches, team_id_map, stat_type, init_stats=None, init_stats_default=0, limit_matches=None):
        n = len(team_id_map.keys())
        s = np.zeros([n, 1])
        for match in matches:
            if match.comp_level != 'qm':  # only consider quals matches
                continue

            treat_as_unplayed = limit_matches is not None and match.match_number > limit_matches

            for alliance_color in ['red', 'blue']:
                alliance_teams = [team[3:] for team in match.alliances[alliance_color]['teams']]
                stat = cls._get_stat(stat_type, match, alliance_color, alliance_teams, init_stats, init_stats_default, treat_as_unplayed)
                for team in alliance_teams:
                    s[team_id_map[team]] += stat
        return s

    @classmethod
    def calc_stat(cls, matches, team_list, team_id_map, Minv, stat_type, init_stats=None, init_stats_default=0, limit_matches=None):
        s = cls.build_s_matrix(matches, team_id_map, stat_type, init_stats=init_stats, init_stats_default=init_stats_default, limit_matches=limit_matches)
        x = np.dot(Minv, s)

        stat_dict = {}
        for team, stat in zip(team_list, x):
            stat_dict[team] = stat[0]
        return stat_dict

    @classmethod
    def _get_stat(cls, stat_type, match, alliance_color, alliance_teams, init_stats, init_stats_default, treat_as_unplayed):
        match_played = match.has_been_played and not treat_as_unplayed

        if match_played:
            if stat_type == 'oprs':
                return match.alliances[alliance_color]['score']
            elif stat_type == 'dprs':
                if alliance_color == 'red':
                    other_alliance_color = 'blue'
                else:
                    other_alliance_color = 'red'
                return match.alliances[other_alliance_color]['score']
            elif stat_type == 'ccwms':
                if alliance_color == 'red':
                    other_alliance_color = 'blue'
                else:
                    other_alliance_color = 'red'
                return match.alliances[alliance_color]['score'] - match.alliances[other_alliance_color]['score']

            # 2016 specific
            if stat_type == '2016autoPointsOPR':
                if match.score_breakdown and alliance_color in match.score_breakdown:
                    return match.score_breakdown[alliance_color]['autoPoints']
            elif stat_type == '2016bouldersOPR':
                if match.score_breakdown and alliance_color in match.score_breakdown:
                    return (match.score_breakdown[alliance_color]['autoBouldersLow'] +
                        match.score_breakdown[alliance_color]['autoBouldersHigh'] +
                        match.score_breakdown[alliance_color]['teleopBouldersLow'] +
                        match.score_breakdown[alliance_color]['teleopBouldersHigh'])
            elif stat_type == '2016crossingsOPR':
                if match.score_breakdown and alliance_color in match.score_breakdown:
                    return (match.score_breakdown[alliance_color]['position1crossings'] +
                        match.score_breakdown[alliance_color]['position2crossings'] +
                        match.score_breakdown[alliance_color]['position3crossings'] +
                        match.score_breakdown[alliance_color]['position4crossings'] +
                        match.score_breakdown[alliance_color]['position5crossings'])

        # None of the above cases were met. Return default.
        if init_stats and stat_type in init_stats:
            total = 0
            for team in alliance_teams:
                total += init_stats[stat_type].get(team, init_stats_default)
            return total
        else:
            total = 0
            for team in alliance_teams:
                total += init_stats_default
            return total

    @classmethod
    def calculate_matchstats(cls, matches, year):
        if not matches:
            return {}

        team_list, team_id_map = cls.build_team_mapping(matches)
        if not team_list:
            return {}
        last_event_stats = cls.get_last_event_stats(team_list, matches[0].event)

        Minv = cls.build_Minv_matrix(matches, team_id_map, played_only=True)

        oprs_dict = cls.calc_stat(matches, team_list, team_id_map, Minv, 'oprs')
        dprs_dict = cls.calc_stat(matches, team_list, team_id_map, Minv, 'dprs')
        ccwms_dict = cls.calc_stat(matches, team_list, team_id_map, Minv, 'ccwms')

        stats = {'oprs': oprs_dict, 'dprs': dprs_dict, 'ccwms': ccwms_dict}

        if year == 2016:
            # First ranking tiebreaker
            stats['2016autoPointsOPR'] = cls.calc_stat(matches, team_list, team_id_map, Minv, '2016autoPointsOPR')

            # For RP calculation
            stats['2016crossingsOPR'] = cls.calc_stat(matches, team_list, team_id_map, Minv, '2016crossingsOPR')
            stats['2016bouldersOPR'] = cls.calc_stat(matches, team_list, team_id_map, Minv, '2016bouldersOPR')

        return stats

    @classmethod
    def get_last_event_stats(cls, team_list, event_key):
        year = int(event_key.id()[:4])
        cur_event = event_key.get()

        # Check cache for stored OPRs
        cache_key = '{}:last_event_stats'.format(event_key.id())
        last_event_stats = memcache.get(cache_key)
        if last_event_stats is None:
            last_event_stats = defaultdict(dict)

        # Make necessary queries for missing stats
        futures = []
        for team in team_list:
            if team not in last_event_stats:
                events_future = event_query.TeamYearEventsQuery('frc{}'.format(team), year).fetch_async()
                futures.append((team, events_future))

        # Add missing stats to last_event_stats
        for team, events_future in futures:
            events = events_future.get_result()

            # Find last event before current event
            last_event = None
            last_event_start = None
            for event in events:
                if event.official and event.start_date < cur_event.start_date:
                    if last_event is None or event.start_date > last_event_start:
                        last_event = event
                        last_event_start = event.start_date

            if last_event is not None and last_event.matchstats:
                for stat, values in last_event.matchstats.items():
                    if values and team in values:
                        last_event_stats[stat][team] = values[team]

        memcache.set(cache_key, last_event_stats, 60*60*24)
        return last_event_stats
<commit_msg>Don't count Einstein for last event OPR. Fixes #1555<commit_after># Calculates OPR/DPR/CCWM
# For implementation details, see
# http://www.chiefdelphi.com/forums/showpost.php?p=484220&postcount=19

# M is n x n where n is # of teams
# s is n x 1 where n is # of teams
# solve [M][x]=[s] for [x]

# x is OPR and should be n x 1

from collections import defaultdict
import numpy as np

from google.appengine.api import memcache

from consts.event_type import EventType
from database import event_query
from models.event_team import EventTeam

class MatchstatsHelper(object):
    @classmethod
    def build_team_mapping(cls, matches):
        """
        Returns (team_list, team_id_map)
        team_list: A list of team_str such as '254' or '254B'
        team_id_map: A dict of key: team_str, value: row index in x_matrix that corresponds to the team
        """
        # Build team list
        team_list = set()
        for match in matches:
            if match.comp_level != 'qm':  # only consider quals matches
                continue
            for alliance_color in ['red', 'blue']:
                for team in match.alliances[alliance_color]['teams']:
                    team_list.add(team[3:])  # turns "frc254B" into "254B"

        team_list = list(team_list)
        team_id_map = {}
        for i, team in enumerate(team_list):
            team_id_map[team] = i

        return team_list, team_id_map

    @classmethod
    def build_Minv_matrix(cls, matches, team_id_map, played_only=False):
        n = len(team_id_map.keys())
        M = np.zeros([n, n])
        for match in matches:
            if match.comp_level != 'qm':  # only consider quals matches
                continue
            if played_only and not match.has_been_played:
                continue
            for alliance_color in ['red', 'blue']:
                alliance_teams = match.alliances[alliance_color]['teams']
                for team1 in alliance_teams:
                    team1_id = team_id_map[team1[3:]]
                    for team2 in alliance_teams:
                        M[team1_id, team_id_map[team2[3:]]] += 1
        return np.linalg.pinv(M)

    @classmethod
    def build_s_matrix(cls, matches, team_id_map, stat_type, init_stats=None, init_stats_default=0, limit_matches=None):
        n = len(team_id_map.keys())
        s = np.zeros([n, 1])
        for match in matches:
            if match.comp_level != 'qm':  # only consider quals matches
                continue

            treat_as_unplayed = limit_matches is not None and match.match_number > limit_matches

            for alliance_color in ['red', 'blue']:
                alliance_teams = [team[3:] for team in match.alliances[alliance_color]['teams']]
                stat = cls._get_stat(stat_type, match, alliance_color, alliance_teams, init_stats, init_stats_default, treat_as_unplayed)
                for team in alliance_teams:
                    s[team_id_map[team]] += stat
        return s

    @classmethod
    def calc_stat(cls, matches, team_list, team_id_map, Minv, stat_type, init_stats=None, init_stats_default=0, limit_matches=None):
        s = cls.build_s_matrix(matches, team_id_map, stat_type, init_stats=init_stats, init_stats_default=init_stats_default, limit_matches=limit_matches)
        x = np.dot(Minv, s)

        stat_dict = {}
        for team, stat in zip(team_list, x):
            stat_dict[team] = stat[0]
        return stat_dict

    @classmethod
    def _get_stat(cls, stat_type, match, alliance_color, alliance_teams, init_stats, init_stats_default, treat_as_unplayed):
        match_played = match.has_been_played and not treat_as_unplayed

        if match_played:
            if stat_type == 'oprs':
                return match.alliances[alliance_color]['score']
            elif stat_type == 'dprs':
                if alliance_color == 'red':
                    other_alliance_color = 'blue'
                else:
                    other_alliance_color = 'red'
                return match.alliances[other_alliance_color]['score']
            elif stat_type == 'ccwms':
                if alliance_color == 'red':
                    other_alliance_color = 'blue'
                else:
                    other_alliance_color = 'red'
                return match.alliances[alliance_color]['score'] - match.alliances[other_alliance_color]['score']

            # 2016 specific
            if stat_type == '2016autoPointsOPR':
                if match.score_breakdown and alliance_color in match.score_breakdown:
                    return match.score_breakdown[alliance_color]['autoPoints']
            elif stat_type == '2016bouldersOPR':
                if match.score_breakdown and alliance_color in match.score_breakdown:
                    return (match.score_breakdown[alliance_color]['autoBouldersLow'] +
                        match.score_breakdown[alliance_color]['autoBouldersHigh'] +
                        match.score_breakdown[alliance_color]['teleopBouldersLow'] +
                        match.score_breakdown[alliance_color]['teleopBouldersHigh'])
            elif stat_type == '2016crossingsOPR':
                if match.score_breakdown and alliance_color in match.score_breakdown:
                    return (match.score_breakdown[alliance_color]['position1crossings'] +
                        match.score_breakdown[alliance_color]['position2crossings'] +
                        match.score_breakdown[alliance_color]['position3crossings'] +
                        match.score_breakdown[alliance_color]['position4crossings'] +
                        match.score_breakdown[alliance_color]['position5crossings'])

        # None of the above cases were met. Return default.
        if init_stats and stat_type in init_stats:
            total = 0
            for team in alliance_teams:
                total += init_stats[stat_type].get(team, init_stats_default)
            return total
        else:
            total = 0
            for team in alliance_teams:
                total += init_stats_default
            return total

    @classmethod
    def calculate_matchstats(cls, matches, year):
        if not matches:
            return {}

        team_list, team_id_map = cls.build_team_mapping(matches)
        if not team_list:
            return {}
        last_event_stats = cls.get_last_event_stats(team_list, matches[0].event)

        Minv = cls.build_Minv_matrix(matches, team_id_map, played_only=True)

        oprs_dict = cls.calc_stat(matches, team_list, team_id_map, Minv, 'oprs')
        dprs_dict = cls.calc_stat(matches, team_list, team_id_map, Minv, 'dprs')
        ccwms_dict = cls.calc_stat(matches, team_list, team_id_map, Minv, 'ccwms')

        stats = {'oprs': oprs_dict, 'dprs': dprs_dict, 'ccwms': ccwms_dict}

        if year == 2016:
            # First ranking tiebreaker
            stats['2016autoPointsOPR'] = cls.calc_stat(matches, team_list, team_id_map, Minv, '2016autoPointsOPR')

            # For RP calculation
            stats['2016crossingsOPR'] = cls.calc_stat(matches, team_list, team_id_map, Minv, '2016crossingsOPR')
            stats['2016bouldersOPR'] = cls.calc_stat(matches, team_list, team_id_map, Minv, '2016bouldersOPR')

        return stats

    @classmethod
    def get_last_event_stats(cls, team_list, event_key):
        year = int(event_key.id()[:4])
        cur_event = event_key.get()

        # Check cache for stored OPRs
        cache_key = '{}:last_event_stats'.format(event_key.id())
        last_event_stats = memcache.get(cache_key)
        if last_event_stats is None:
            last_event_stats = defaultdict(dict)

        # Make necessary queries for missing stats
        futures = []
        for team in team_list:
            if team not in last_event_stats:
                events_future = event_query.TeamYearEventsQuery('frc{}'.format(team), year).fetch_async()
                futures.append((team, events_future))

        # Add missing stats to last_event_stats
        for team, events_future in futures:
            events = events_future.get_result()

            # Find last event before current event
            last_event = None
            last_event_start = None
            for event in events:
                if event.official and event.start_date < cur_event.start_date and event.event_type_enum != EventType.CMP_FINALS:
                    if last_event is None or event.start_date > last_event_start:
                        last_event = event
                        last_event_start = event.start_date

            if last_event is not None and last_event.matchstats:
                for stat, values in last_event.matchstats.items():
                    if values and team in values:
                        last_event_stats[stat][team] = values[team]

        memcache.set(cache_key, last_event_stats, 60*60*24)
        return last_event_stats
<|endoftext|>


<commit_before># -*- coding: utf8 -*-
import sys

from .base import *
from .holonet import *

TESTING = 'test' in sys.argv

if TESTING:
    from .test import *

try:
    from .local import *
except ImportError as e:
    raise ImportError("Couldn't load local settings holonet.settings.local")
<commit_msg>Not raise import error when local.py not exist. Required for docs.<commit_after># -*- coding: utf8 -*-
import sys

from .base import *
from .holonet import *

TESTING = 'test' in sys.argv

if TESTING:
    from .test import *

try:
    from .local import *
except ImportError as e:
    print('Unable to load local settings.')
<|endoftext|>


<commit_before># Copyright (c) 2009-2021 The Regents of the University of Michigan
# This file is part of the HOOMD-blue project, released under the BSD 3-Clause
# License.

"""Implement CustomWriter."""

from hoomd.custom import (CustomOperation, _InternalCustomOperation, Action)
from hoomd.operation import Writer

class _WriterProperty:

    @property
    def analyzer(self):
        return self._action

    @analyzer.setter
    def analyzer(self, analyzer):
        if isinstance(analyzer, Action):
            self._action = analyzer
        else:
            raise ValueError(
                "analyzer must be an instance of hoomd.custom.Action")

class CustomWriter(CustomOperation, _WriterProperty, Writer):
    """User-defined writer.

    Args:
        action (hoomd.custom.Action): The action to call.
        trigger (hoomd.trigger.Trigger): Select the timesteps to call the
          action.

    `CustomWriter` is a `hoomd.operation.Writer` that wraps a user-defined
    `hoomd.custom.Action` object so the action can be added to a
    `hoomd.Operations` instance for use with `hoomd.Simulation` objects.

    Writers may read the system state and generate output files or print to
    output streams. Writers should not modify the system state.

    See Also:
        The base class `hoomd.custom.CustomOperation`.

        `hoomd.update.CustomUpdater`

        `hoomd.tune.CustomTuner`
    """
    _cpp_list_name = 'analyzers'
    _cpp_class_name = 'PythonAnalyzer'

class _InternalCustomWriter(_InternalCustomOperation, Writer):
    _cpp_list_name = 'analyzers'
    _cpp_class_name = 'PythonAnalyzer'
<commit_msg>Rename CustomWriter action attribute to writer<commit_after># Copyright (c) 2009-2021 The Regents of the University of Michigan
# This file is part of the HOOMD-blue project, released under the BSD 3-Clause
# License.

"""Implement CustomWriter."""

from hoomd.custom import (CustomOperation, _InternalCustomOperation, Action)
from hoomd.operation import Writer

class _WriterProperty:

    @property
    def writer(self):
        return self._action

    @writer.setter
    def writer(self, analyzer):
        if isinstance(analyzer, Action):
            self._action = analyzer
        else:
            raise ValueError(
                "analyzer must be an instance of hoomd.custom.Action")

class CustomWriter(CustomOperation, _WriterProperty, Writer):
    """User-defined writer.

    Args:
        action (hoomd.custom.Action): The action to call.
        trigger (hoomd.trigger.Trigger): Select the timesteps to call the
          action.

    `CustomWriter` is a `hoomd.operation.Writer` that wraps a user-defined
    `hoomd.custom.Action` object so the action can be added to a
    `hoomd.Operations` instance for use with `hoomd.Simulation` objects.

    Writers may read the system state and generate output files or print to
    output streams. Writers should not modify the system state.

    See Also:
        The base class `hoomd.custom.CustomOperation`.

        `hoomd.update.CustomUpdater`

        `hoomd.tune.CustomTuner`
    """
    _cpp_list_name = 'analyzers'
    _cpp_class_name = 'PythonAnalyzer'

class _InternalCustomWriter(_InternalCustomOperation, Writer):
    _cpp_list_name = 'analyzers'
    _cpp_class_name = 'PythonAnalyzer'
<|endoftext|>


<commit_before>#!/usr/bin/env python
#
# Root command that delegates to all GreatFET subcommands.

# This file is part of GreatFET.

from __future__ import print_function

import difflib
import errno
import sys
import os

# The prefix which all greatfet-subcommands start with.
GREATFET_PREFIX = 'greatfet_'

def looks_like_valid_greatfet_subcommand(directory, executable):
    """
    Returns true iff the given directory/binary pair seem to represent
    a valid GreatFET subcommand.

    @param directory The directory in the PATH we're currently looking in.
    @param executable The binary name, which should be presnent in the given directory.
    @return True iff the given binary appares to be a valid GreatFET command.
    """

    full_path = os.path.join(directory, executable)

    # Valid GreatFET subcommands start with our prefix.
    if not executable.startswith(GREATFET_PREFIX):
        return False

    # Valid GreatFET subcommands are files.
    if not os.path.isfile(full_path):
        return False

    # Valid GreatFET subcommands are executable.
    if not os.access(full_path, os.X_OK):
        return False

    # If all of the commands above are met, we have what looks like a subcommand.
    return True

def find_all_subcommands():
    """
    Locates all GreatFET subcommands in the user's path -- that is, all
    binaries that start with GREATFET_PREFIX.
    """

    # If we don't have a PATH variable, we can't have subcommands.
    # Bail out.
    if 'PATH' not in os.environ:
        return []

    # Parse the system's PATH variable and get each of the relevant directories.
    sys_path = os.environ['PATH']
    path_entires = sys_path.split(os.pathsep)

    # Search each entry in the path for GreatFET subcommands.
    subcommands = {}
    for directory in path_entires:

        # Skip any path entries that don't seem to be represented on the real system...
        if not os.path.isdir(directory):
            continue

        # ... and search each entry that is a real directory.
        for executable in os.listdir(directory):

            # If the executable seems to be a GreatFET command, use it.
            if looks_like_valid_greatfet_subcommand(directory, executable):

                # Cache the relationships between subcommands and their executables.
                full_path = os.path.join(directory, executable)
                subcommand_name = executable[len(GREATFET_PREFIX):]
                subcommands[subcommand_name] = full_path

    return subcommands

def find_subcommand(name, allow_partial_matches=True, print_errors=True):
    """
    Returns the full path to the current subcommand, if one exists.

    @param name The name of the subcommand to look for.
    @param allow_partial_matches If set, this function will accept abbreviated subcommands.
    @param print_errors If set, errors will be printed.
    @return The full path to the subcommand, or None if none exists.
    """

    subcommands = find_all_subcommands()

    # If we have a direct match, return it.
    if name in subcommands:
        return subcommands[name]

    # If we're accepting partial matches, look for one.
    if allow_partial_matches:
        matches = [subcommand for subcommand in subcommands.keys() if subcommand.startswith(name)]

        # If we've found exactly one, use it.
        if len(matches) == 1:
            return subcommands[matches[0]]

        # Otherwise, print the error.
        elif matches and print_errors:
            matches = "\n\t".join(matches)
            print("Subcommand short-name '{}' is ambiguous; it could refer to:\n\t{}\n".format(name, matches))
            return False

    if print_errors:
        print("ERROR: Unsupported subcommand '{}'.\nCheck to ensure the package providing the " \
                "subcommand is installed.\n".format(name))

    return False

def find_corrections_message(name):
    """
    Generates a message that provides potential corrections to the user if
    their command doesn't match.
    """

    # Find a list of "did you mean" style corrections.
    corrections = difflib.get_close_matches(name, find_all_subcommands())

    # If we didn't find any, don't provide a message.
    if not corrections:
        return ''

    # Otherwise, generate a string that informs the user of their potential error.
    plural_suffix = "s are" if len(corrections) > 1 else " is"
    corrections.insert(0, "The most similar sub-command{}:".format(plural_suffix))
    return "\n\t".join(corrections)

def print_usage(argv):

    # If we don't have argument name information, assume this was called "greatfet"
    name = os.path.basename(argv[0]) if len(argv) else "greatfet"

    print("usage: {} <subcommand>\n".format(name))
    print("Top-level utility for working with GreatFET devices.\n")
    print("supported subcommands:")

    for subcommand in find_all_subcommands():
        print("\t{}".format(subcommand))

    print("\nTo get help for a subcommand, use '{} <subcommand> --help'.".format(name))
    print("For example, for help with the firmware subcommand, use '{} firmware --help'.\n\n".format(name))
    print("You can create and install your own subcommands. Simply create an executable with")
    print("the name greatfet_<subcommand>, and add it to your path. See the (forthcoming)")
    print("documentation for more information.\n")

def main():
    """
    Main routine that delegates all commands to GreatFET subcommands.
    """
    argv = sys.argv[:]

    # If we don't have a subcommand, abort.
    if len(argv) < 2:
        print_usage(argv)
        sys.exit(errno.EINVAL)

    # Get the subcommand name.
    subcommand_name = argv[1]

    # If help is passed as a special-case subcommand, and we have
    # more arguments, help the user out by generating a likely-valid help request.
    if subcommand_name == "help" and len(argv) > 2:
        subcommand_name = argv[2]
        argv[1], argv[2] = argv[2], "--help"

    # Find the binary to execute...
    binary = find_subcommand(subcommand_name)

    # If we couldn't find the relevant binary, print a message.
    if not binary:
        # If there are words that are similar to the relevant word, suggest them as corrections.
        corrections = find_corrections_message(subcommand_name)
        if corrections:
            print(corrections)
            print('')

        sys.exit(errno.EINVAL)

    # Compute the arguments that should be passed to the subcommand,
    # which include the subcommand name and all arguments beyond the subcommand.
    binary_name = os.path.basename(binary)
    arguments = [binary_name] + argv[2:]

    # Pass control entirely to the subcommand.
    sys.exit(os.execv(binary, arguments))

if __name__ == '__main__':
    main()
<commit_msg>Return once the child process *ends* (windows bug fix)<commit_after>#!/usr/bin/env python
#
# Root command that delegates to all GreatFET subcommands.

# This file is part of GreatFET.

from __future__ import print_function

import difflib
import errno
import sys
import os

# The prefix which all greatfet-subcommands start with.
GREATFET_PREFIX = 'greatfet_'

def looks_like_valid_greatfet_subcommand(directory, executable):
    """
    Returns true iff the given directory/binary pair seem to represent
    a valid GreatFET subcommand.

    @param directory The directory in the PATH we're currently looking in.
    @param executable The binary name, which should be presnent in the given directory.
    @return True iff the given binary appares to be a valid GreatFET command.
    """

    full_path = os.path.join(directory, executable)

    # Valid GreatFET subcommands start with our prefix.
    if not executable.startswith(GREATFET_PREFIX):
        return False

    # Valid GreatFET subcommands are files.
    if not os.path.isfile(full_path):
        return False

    # Valid GreatFET subcommands are executable.
    if not os.access(full_path, os.X_OK):
        return False

    # If all of the commands above are met, we have what looks like a subcommand.
    return True

def find_all_subcommands():
    """
    Locates all GreatFET subcommands in the user's path -- that is, all
    binaries that start with GREATFET_PREFIX.
    """

    # If we don't have a PATH variable, we can't have subcommands.
    # Bail out.
    if 'PATH' not in os.environ:
        return []

    # Parse the system's PATH variable and get each of the relevant directories.
    sys_path = os.environ['PATH']
    path_entires = sys_path.split(os.pathsep)

    # Search each entry in the path for GreatFET subcommands.
    subcommands = {}
    for directory in path_entires:

        # Skip any path entries that don't seem to be represented on the real system...
        if not os.path.isdir(directory):
            continue

        # ... and search each entry that is a real directory.
        for executable in os.listdir(directory):

            # If the executable seems to be a GreatFET command, use it.
            if looks_like_valid_greatfet_subcommand(directory, executable):

                # Cache the relationships between subcommands and their executables.
                full_path = os.path.join(directory, executable)
                subcommand_name = executable[len(GREATFET_PREFIX):]
                subcommands[subcommand_name] = full_path

    return subcommands

def find_subcommand(name, allow_partial_matches=True, print_errors=True):
    """
    Returns the full path to the current subcommand, if one exists.

    @param name The name of the subcommand to look for.
    @param allow_partial_matches If set, this function will accept abbreviated subcommands.
    @param print_errors If set, errors will be printed.
    @return The full path to the subcommand, or None if none exists.
    """

    subcommands = find_all_subcommands()

    # If we have a direct match, return it.
    if name in subcommands:
        return subcommands[name]

    # If we're accepting partial matches, look for one.
    if allow_partial_matches:
        matches = [subcommand for subcommand in subcommands.keys() if subcommand.startswith(name)]

        # If we've found exactly one, use it.
        if len(matches) == 1:
            return subcommands[matches[0]]

        # Otherwise, print the error.
        elif matches and print_errors:
            matches = "\n\t".join(matches)
            print("Subcommand short-name '{}' is ambiguous; it could refer to:\n\t{}\n".format(name, matches))
            return False

    if print_errors:
        print("ERROR: Unsupported subcommand '{}'.\nCheck to ensure the package providing the " \
                "subcommand is installed.\n".format(name))

    return False

def find_corrections_message(name):
    """
    Generates a message that provides potential corrections to the user if
    their command doesn't match.
    """

    # Find a list of "did you mean" style corrections.
    corrections = difflib.get_close_matches(name, find_all_subcommands())

    # If we didn't find any, don't provide a message.
    if not corrections:
        return ''

    # Otherwise, generate a string that informs the user of their potential error.
    plural_suffix = "s are" if len(corrections) > 1 else " is"
    corrections.insert(0, "The most similar sub-command{}:".format(plural_suffix))
    return "\n\t".join(corrections)

def print_usage(argv):

    # If we don't have argument name information, assume this was called "greatfet"
    name = os.path.basename(argv[0]) if len(argv) else "greatfet"

    print("usage: {} <subcommand>\n".format(name))
    print("Top-level utility for working with GreatFET devices.\n")
    print("supported subcommands:")

    for subcommand in find_all_subcommands():
        print("\t{}".format(subcommand))

    print("\nTo get help for a subcommand, use '{} <subcommand> --help'.".format(name))
    print("For example, for help with the firmware subcommand, use '{} firmware --help'.\n\n".format(name))
    print("You can create and install your own subcommands. Simply create an executable with")
    print("the name greatfet_<subcommand>, and add it to your path. See the (forthcoming)")
    print("documentation for more information.\n")

def main():
    """
    Main routine that delegates all commands to GreatFET subcommands.
    """
    argv = sys.argv[:]

    # If we don't have a subcommand, abort.
    if len(argv) < 2:
        print_usage(argv)
        sys.exit(errno.EINVAL)

    # Get the subcommand name.
    subcommand_name = argv[1]

    # If help is passed as a special-case subcommand, and we have
    # more arguments, help the user out by generating a likely-valid help request.
    if subcommand_name == "help" and len(argv) > 2:
        subcommand_name = argv[2]
        argv[1], argv[2] = argv[2], "--help"

    # Find the binary to execute...
    binary = find_subcommand(subcommand_name)

    # If we couldn't find the relevant binary, print a message.
    if not binary:
        # If there are words that are similar to the relevant word, suggest them as corrections.
        corrections = find_corrections_message(subcommand_name)
        if corrections:
            print(corrections)
            print('')

        sys.exit(errno.EINVAL)

    # Compute the arguments that should be passed to the subcommand,
    # which include the subcommand name and all arguments beyond the subcommand.
    binary_name = os.path.basename(binary)
    arguments = [binary_name] + argv[2:]

    # Pass control entirely to the subcommand.
    sys.exit(os.spawnv(os.P_WAIT, binary, arguments))

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


<commit_before>import time
import logging
import numpy as np
from os.path import splitext

from analysis.plotting.plotting import plot_occupancy, plot_fancy_occupancy, make_occupancy_hist
from daq.readout import get_col_row_array_from_data_record_array, convert_data_array, data_array_from_data_dict_iterable, is_data_record, is_data_from_channel

from scan.scan import ScanBase
from daq.readout import open_raw_data_file

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)-8s] (%(threadName)-10s) %(message)s")

local_configuration = {
    "cfg_name": 'noise_occupancy_tuning',  # the name of the new config with the tuning
    "occupancy_limit": 10 ** (-5),  # 0 will mask any pixel with occupancy greater than zero
    "triggers": 10000000,
    "trig_count": 1,
    "disable_for_mask": ['Enable'],
    "enable_for_mask": ['Imon'],
    "overwrite_mask": False,
    "col_span": [1, 80],
    "row_span": [1, 336],
    "timeout_no_data": 10
}

class NoiseOccupancyScan(ScanBase):
    scan_id = "noise_occupancy_tuning"

    def scan(self, cfg_name='noise_occupancy_tuning', occupancy_limit=10 ** (-5), triggers=10000000, trig_count=1, disable_for_mask=['Enable'], enable_for_mask=['Imon'], overwrite_mask=False, col_span=[1, 80], row_span=[1, 336], timeout_no_data=10, **kwargs):
        '''Masking pixels with occupancy above certain limit.

        Parameters
        ----------
        cfg_name : string
            File name of the configuration file. If None or not given, use default file name.
        occupancy_limit : float
            Occupancy limit which is multiplied with measured number of hits for each pixel. Any pixel above 1 will be masked.
        triggers : int
            Total number of triggers sent to FE. From 1 to 4294967295 (32-bit unsigned int).
        trig_count : int
            FE global register Trig_Count.
        disable_for_mask : list, tuple
            List of masks for which noisy pixels will be disabled.
        enable_for_mask : list, tuple
            List of masks for which noisy pixels will be enabled.
        overwrite_mask : bool
            Overwrite masks (disable_for_mask, enable_for_mask) if set to true. If set to false, make a combination of existing mask (configuration file) and generated mask (selected columns and rows), otherwise only use generated mask.
        col_span : list, tuple
            Column range (from minimum to maximum value). From 1 to 80.
        row_span : list, tuple
            Row range (from minimum to maximum value). From 1 to 336.
        timeout_no_data : int
            In seconds; if no data, stop scan after given time.
        scan_timeout : int
            In seconds; stop scan after given time.

        Note
        ----
        The total number of trigger is triggers * consecutive_lvl1.
        Please note that a high trigger rate leads to an effective lower threshold.
        '''
        # create restore point
        self.register.create_restore_point()
        self.trig_count = trig_count
        if trig_count == 0:
            consecutive_lvl1 = (2 ** self.register.get_global_register_objects(name=['Trig_Count'])[0].bitlength)
        else:
            consecutive_lvl1 = trig_count
        if occupancy_limit * triggers * consecutive_lvl1 < 1.0:
            logging.warning('Number of triggers too low for given occupancy limit. Any noise hit will lead to a masked pixel.')

        commands = []
        commands.extend(self.register.get_commands("confmode"))
        mask = self.register_utils.make_box_pixel_mask_from_col_row(column=col_span, row=row_span)  # 1 for selected columns, else 0
        for pixel_reg in disable_for_mask:  # enabled pixels set to 1
            if overwrite_mask:
                pixel_mask = mask
            else:
                pixel_mask = np.logical_and(mask, self.register.get_pixel_register_value(pixel_reg))
            self.register.set_pixel_register_value(pixel_reg, pixel_mask)
            commands.extend(self.register.get_commands("wrfrontend", same_mask_for_all_dc=False, name=pixel_reg))
        mask = self.register_utils.make_box_pixel_mask_from_col_row(column=col_span, row=row_span, default=1, value=0)  # 0 for selected columns, else 1
        for pixel_reg in enable_for_mask:  # disabled pixels set to 1
            if overwrite_mask:
                pixel_mask = mask
            else:
                pixel_mask = np.logical_or(mask, self.register.get_pixel_register_value(pixel_reg))
            self.register.set_pixel_register_value(pixel_reg, pixel_mask)
            commands.extend(self.register.get_commands("wrfrontend", same_mask_for_all_dc=False, name=pixel_reg))
        # disable C_inj mask
        pixel_reg = "C_High"
        self.register.set_pixel_register_value(pixel_reg, 0)
        commands.extend(self.register.get_commands("wrfrontend", same_mask_for_all_dc=True, name=pixel_reg))
        pixel_reg = "C_Low"
        self.register.set_pixel_register_value(pixel_reg, 0)
        commands.extend(self.register.get_commands("wrfrontend", same_mask_for_all_dc=True, name=pixel_reg))
#         self.register.set_global_register_value("Trig_Lat", 232)  # set trigger latency
        self.register.set_global_register_value("Trig_Count", trig_count)  # set number of consecutive triggers
        commands.extend(self.register.get_commands("wrregister", name=["Trig_Count"]))
        # setting FE into runmode
        commands.extend(self.register.get_commands("runmode"))
        self.register_utils.send_commands(commands)

        self.col_arr = np.array([], dtype=np.dtype('>u1'))
        self.row_arr = np.array([], dtype=np.dtype('>u1'))

        with open_raw_data_file(filename=self.scan_data_filename, title=self.scan_id) as raw_data_file:
            self.readout.start()

            # preload command
            command_delay = 500  # <100kHz
            lvl1_command = self.register.get_commands("lv1")[0] + self.register.get_commands("zeros", length=command_delay)[0]
            commnd_lenght = lvl1_command.length()
            logging.info('Estimated scan time: %ds' % int(commnd_lenght * 25 * (10 ** -9) * triggers))
            logging.info('Please stand by...')
            self.register_utils.send_command(lvl1_command, repeat=triggers, wait_for_finish=False, set_length=True, clear_memory=False)

            wait_for_first_data = False
            last_iteration = time.time()
            saw_no_data_at_time = last_iteration
            saw_data_at_time = last_iteration
            no_data_at_time = last_iteration
            while not self.stop_thread_event.wait(self.readout.readout_interval):
                last_iteration = time.time()
                try:
                    data = (self.readout.data.popleft(), )
                    raw_data_file.append(data)
                    col_arr_tmp, row_arr_tmp = convert_data_array(data_array_from_data_dict_iterable(data), filter_func=is_data_record, converter_func=get_col_row_array_from_data_record_array)
                    self.col_arr = np.concatenate((self.col_arr, col_arr_tmp))
                    self.row_arr = np.concatenate((self.row_arr, row_arr_tmp))
                    #logging.info('data words')
                except IndexError:  # no data
                    #logging.info('no data words')
                    no_data_at_time = last_iteration
                    if self.register_utils.is_ready:
                        self.stop_thread_event.set()
                        logging.info('Finished sending %d triggers' % triggers)
                    elif wait_for_first_data is False and saw_no_data_at_time > (saw_data_at_time + timeout_no_data):
                        logging.info('Reached no data timeout. Stopping Scan...')
                        self.stop_thread_event.set()
                    elif wait_for_first_data == False:
                        saw_no_data_at_time = no_data_at_time
                    elif self.reaout_utils.is_ready:
                        self.stop_thread_event.set()

                    continue

                saw_data_at_time = last_iteration

                if wait_for_first_data is True:
                    logging.info('Taking data...')
                    wait_for_first_data = False

            self.readout.stop()

        self.register.restore()

        occ_hist, _, _ = np.histogram2d(self.col_arr, self.row_arr, bins=(80, 336), range=[[1, 80], [1, 336]])

        self.occ_mask = np.zeros(shape=occ_hist.shape, dtype=np.dtype('>u1'))
        # noisy pixels are set to 1
        self.occ_mask[occ_hist > occupancy_limit * triggers * consecutive_lvl1] = 1
        # make inverse
        self.inv_occ_mask = self.register_utils.invert_pixel_mask(self.occ_mask)
        self.disable_for_mask = disable_for_mask
        if overwrite_mask:
            for mask in disable_for_mask:
                self.register.set_pixel_register_value(mask, self.inv_occ_mask)
        else:
            for mask in disable_for_mask:
                enable_mask = np.logical_and(self.inv_occ_mask, self.register.get_pixel_register_value(mask))
                self.register.set_pixel_register_value(mask, enable_mask)

        self.enable_for_mask = enable_for_mask
        if overwrite_mask:
            for mask in enable_for_mask:
                self.register.set_pixel_register_value(mask, self.occ_mask)
        else:
            for mask in enable_for_mask:
                disable_mask = np.logical_or(self.occ_mask, self.register.get_pixel_register_value(mask))
                self.register.set_pixel_register_value(mask, disable_mask)

#         plot_occupancy(make_occupancy_hist(self.col_arr, self.row_arr), z_max=None, filename=self.scan_data_filename + "_occupancy.pdf")
        self.register.save_configuration(cfg_name if cfg_name else (splitext(self.device_configuration["configuration_file"])[0] + '_' + self.scan_id))

    def analyze(self):
        from analysis.analyze_raw_data import AnalyzeRawData
        output_file = self.scan_data_filename + "_interpreted.h5"
        with AnalyzeRawData(raw_data_file=scan.scan_data_filename, analyzed_data_file=output_file, create_pdf=True) as analyze_raw_data:
            analyze_raw_data.create_source_scan_hist = True
#             analyze_raw_data.create_hit_table = True
#             analyze_raw_data.interpreter.debug_events(0, 0, True)  # events to be printed onto the console for debugging, usually deactivated
            analyze_raw_data.interpreter.set_warning_output(False)
            analyze_raw_data.interpret_word_table()
            analyze_raw_data.interpreter.print_summary()
            analyze_raw_data.plot_histograms()
            plot_occupancy(self.occ_mask.T, title='Noisy Pixels', z_max=1, filename=analyze_raw_data.output_pdf)
            plot_fancy_occupancy(self.occ_mask.T, z_max=1, filename=analyze_raw_data.output_pdf)
            for mask in self.disable_for_mask:
                mask_name = self.register.get_pixel_register_attributes("full_name", do_sort=True, name=[mask])[0]
                plot_occupancy(self.register.get_pixel_register_value(mask).T, title='%s Mask' % mask_name, z_max=1, filename=analyze_raw_data.output_pdf)
            for mask in self.enable_for_mask:
                mask_name = self.register.get_pixel_register_attributes("full_name", do_sort=True, name=[mask])[0]
                plot_occupancy(self.register.get_pixel_register_value(mask).T, title='%s Mask' % mask_name, z_max=1, filename=analyze_raw_data.output_pdf)

if __name__ == "__main__":
    import configuration
    scan = NoiseOccupancyScan(**configuration.default_configuration)
    scan.start(use_thread=False, **local_configuration)
    scan.stop()
    scan.analyze()
<commit_msg>BUG: fix scan -> self ENH: use save_configuration<commit_after>import time
import logging
import numpy as np

from analysis.plotting.plotting import plot_occupancy, plot_fancy_occupancy, make_occupancy_hist
from daq.readout import get_col_row_array_from_data_record_array, convert_data_array, data_array_from_data_dict_iterable, is_data_record, is_data_from_channel

from scan.scan import ScanBase
from daq.readout import open_raw_data_file

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)-8s] (%(threadName)-10s) %(message)s")

local_configuration = {
    "cfg_name": 'noise_occupancy_tuning',  # the name of the new config with the tuning
    "occupancy_limit": 10 ** (-5),  # 0 will mask any pixel with occupancy greater than zero
    "triggers": 10000000,
    "trig_count": 1,
    "disable_for_mask": ['Enable'],
    "enable_for_mask": ['Imon'],
    "overwrite_mask": False,
    "col_span": [1, 80],
    "row_span": [1, 336],
    "timeout_no_data": 10
}

class NoiseOccupancyScan(ScanBase):
    scan_id = "noise_occupancy_tuning"

    def scan(self, cfg_name='noise_occupancy_tuning', occupancy_limit=10 ** (-5), triggers=10000000, trig_count=1, disable_for_mask=['Enable'], enable_for_mask=['Imon'], overwrite_mask=False, col_span=[1, 80], row_span=[1, 336], timeout_no_data=10, **kwargs):
        '''Masking pixels with occupancy above certain limit.

        Parameters
        ----------
        cfg_name : string
            File name of the configuration file. If '', use a file name generated from scan ID and number. If None, overwrite configuration file.
        occupancy_limit : float
            Occupancy limit which is multiplied with measured number of hits for each pixel. Any pixel above 1 will be masked.
        triggers : int
            Total number of triggers sent to FE. From 1 to 4294967295 (32-bit unsigned int).
        trig_count : int
            FE global register Trig_Count.
        disable_for_mask : list, tuple
            List of masks for which noisy pixels will be disabled.
        enable_for_mask : list, tuple
            List of masks for which noisy pixels will be enabled.
        overwrite_mask : bool
            Overwrite masks (disable_for_mask, enable_for_mask) if set to true. If set to false, make a combination of existing mask (configuration file) and generated mask (selected columns and rows), otherwise only use generated mask.
        col_span : list, tuple
            Column range (from minimum to maximum value). From 1 to 80.
        row_span : list, tuple
            Row range (from minimum to maximum value). From 1 to 336.
        timeout_no_data : int
            In seconds; if no data, stop scan after given time.
        scan_timeout : int
            In seconds; stop scan after given time.

        Note
        ----
        The total number of trigger is triggers * consecutive_lvl1.
        Please note that a high trigger rate leads to an effective lower threshold.
        '''
        # create restore point
        self.register.create_restore_point()
        if trig_count == 0:
            consecutive_lvl1 = (2 ** self.register.get_global_register_objects(name=['Trig_Count'])[0].bitlength)
        else:
            consecutive_lvl1 = trig_count
        if occupancy_limit * triggers * consecutive_lvl1 < 1.0:
            logging.warning('Number of triggers too low for given occupancy limit. Any noise hit will lead to a masked pixel.')

        commands = []
        commands.extend(self.register.get_commands("confmode"))
        mask = self.register_utils.make_box_pixel_mask_from_col_row(column=col_span, row=row_span)  # 1 for selected columns, else 0
        for pixel_reg in disable_for_mask:  # enabled pixels set to 1
            if overwrite_mask:
                pixel_mask = mask
            else:
                pixel_mask = np.logical_and(mask, self.register.get_pixel_register_value(pixel_reg))
            self.register.set_pixel_register_value(pixel_reg, pixel_mask)
            commands.extend(self.register.get_commands("wrfrontend", same_mask_for_all_dc=False, name=pixel_reg))
        mask = self.register_utils.make_box_pixel_mask_from_col_row(column=col_span, row=row_span, default=1, value=0)  # 0 for selected columns, else 1
        for pixel_reg in enable_for_mask:  # disabled pixels set to 1
            if overwrite_mask:
                pixel_mask = mask
            else:
                pixel_mask = np.logical_or(mask, self.register.get_pixel_register_value(pixel_reg))
            self.register.set_pixel_register_value(pixel_reg, pixel_mask)
            commands.extend(self.register.get_commands("wrfrontend", same_mask_for_all_dc=False, name=pixel_reg))
        # disable C_inj mask
        pixel_reg = "C_High"
        self.register.set_pixel_register_value(pixel_reg, 0)
        commands.extend(self.register.get_commands("wrfrontend", same_mask_for_all_dc=True, name=pixel_reg))
        pixel_reg = "C_Low"
        self.register.set_pixel_register_value(pixel_reg, 0)
        commands.extend(self.register.get_commands("wrfrontend", same_mask_for_all_dc=True, name=pixel_reg))
#         self.register.set_global_register_value("Trig_Lat", 232)  # set trigger latency
        self.register.set_global_register_value("Trig_Count", trig_count)  # set number of consecutive triggers
        commands.extend(self.register.get_commands("wrregister", name=["Trig_Count"]))
        # setting FE into runmode
        commands.extend(self.register.get_commands("runmode"))
        self.register_utils.send_commands(commands)

        self.col_arr = np.array([], dtype=np.dtype('>u1'))
        self.row_arr = np.array([], dtype=np.dtype('>u1'))

        with open_raw_data_file(filename=self.scan_data_filename, title=self.scan_id) as raw_data_file:
            self.readout.start()

            # preload command
            command_delay = 500  # <100kHz
            lvl1_command = self.register.get_commands("lv1")[0] + self.register.get_commands("zeros", length=command_delay)[0]
            commnd_lenght = lvl1_command.length()
            logging.info('Estimated scan time: %ds' % int(commnd_lenght * 25 * (10 ** -9) * triggers))
            logging.info('Please stand by...')
            self.register_utils.send_command(lvl1_command, repeat=triggers, wait_for_finish=False, set_length=True, clear_memory=False)

            wait_for_first_data = False
            last_iteration = time.time()
            saw_no_data_at_time = last_iteration
            saw_data_at_time = last_iteration
            no_data_at_time = last_iteration
            while not self.stop_thread_event.wait(self.readout.readout_interval):
                last_iteration = time.time()
                try:
                    data = (self.readout.data.popleft(), )
                    raw_data_file.append(data)
                    col_arr_tmp, row_arr_tmp = convert_data_array(data_array_from_data_dict_iterable(data), filter_func=is_data_record, converter_func=get_col_row_array_from_data_record_array)
                    self.col_arr = np.concatenate((self.col_arr, col_arr_tmp))
                    self.row_arr = np.concatenate((self.row_arr, row_arr_tmp))
                    #logging.info('data words')
                except IndexError:  # no data
                    #logging.info('no data words')
                    no_data_at_time = last_iteration
                    if self.register_utils.is_ready:
                        self.stop_thread_event.set()
                        logging.info('Finished sending %d triggers' % triggers)
                    elif wait_for_first_data is False and saw_no_data_at_time > (saw_data_at_time + timeout_no_data):
                        logging.info('Reached no data timeout. Stopping Scan...')
                        self.stop_thread_event.set()
                    elif wait_for_first_data == False:
                        saw_no_data_at_time = no_data_at_time
                    elif self.reaout_utils.is_ready:
                        self.stop_thread_event.set()

                    continue

                saw_data_at_time = last_iteration

                if wait_for_first_data is True:
                    logging.info('Taking data...')
                    wait_for_first_data = False

            self.readout.stop()

        self.register.restore()

        occ_hist, _, _ = np.histogram2d(self.col_arr, self.row_arr, bins=(80, 336), range=[[1, 80], [1, 336]])

        self.occ_mask = np.zeros(shape=occ_hist.shape, dtype=np.dtype('>u1'))
        # noisy pixels are set to 1
        self.occ_mask[occ_hist > occupancy_limit * triggers * consecutive_lvl1] = 1
        # make inverse
        self.inv_occ_mask = self.register_utils.invert_pixel_mask(self.occ_mask)
        self.disable_for_mask = disable_for_mask
        if overwrite_mask:
            for mask in disable_for_mask:
                self.register.set_pixel_register_value(mask, self.inv_occ_mask)
        else:
            for mask in disable_for_mask:
                enable_mask = np.logical_and(self.inv_occ_mask, self.register.get_pixel_register_value(mask))
                self.register.set_pixel_register_value(mask, enable_mask)

        self.enable_for_mask = enable_for_mask
        if overwrite_mask:
            for mask in enable_for_mask:
                self.register.set_pixel_register_value(mask, self.occ_mask)
        else:
            for mask in enable_for_mask:
                disable_mask = np.logical_or(self.occ_mask, self.register.get_pixel_register_value(mask))
                self.register.set_pixel_register_value(mask, disable_mask)

#         plot_occupancy(make_occupancy_hist(self.col_arr, self.row_arr), z_max=None, filename=self.scan_data_filename + "_occupancy.pdf")
        self.save_configuration(scan.cfg_name)

    def analyze(self):
        from analysis.analyze_raw_data import AnalyzeRawData
        output_file = self.scan_data_filename + "_interpreted.h5"
        with AnalyzeRawData(raw_data_file=self.scan_data_filename, analyzed_data_file=output_file, create_pdf=True) as analyze_raw_data:
            analyze_raw_data.create_source_scan_hist = True
#             analyze_raw_data.create_hit_table = True
#             analyze_raw_data.interpreter.debug_events(0, 0, True)  # events to be printed onto the console for debugging, usually deactivated
            analyze_raw_data.interpreter.set_warning_output(False)
            analyze_raw_data.interpret_word_table()
            analyze_raw_data.interpreter.print_summary()
            analyze_raw_data.plot_histograms()
            plot_occupancy(self.occ_mask.T, title='Noisy Pixels', z_max=1, filename=analyze_raw_data.output_pdf)
            plot_fancy_occupancy(self.occ_mask.T, z_max=1, filename=analyze_raw_data.output_pdf)
            for mask in self.disable_for_mask:
                mask_name = self.register.get_pixel_register_attributes("full_name", do_sort=True, name=[mask])[0]
                plot_occupancy(self.register.get_pixel_register_value(mask).T, title='%s Mask' % mask_name, z_max=1, filename=analyze_raw_data.output_pdf)
            for mask in self.enable_for_mask:
                mask_name = self.register.get_pixel_register_attributes("full_name", do_sort=True, name=[mask])[0]
                plot_occupancy(self.register.get_pixel_register_value(mask).T, title='%s Mask' % mask_name, z_max=1, filename=analyze_raw_data.output_pdf)

if __name__ == "__main__":
    import configuration
    scan = NoiseOccupancyScan(**configuration.default_configuration)
    scan.start(use_thread=False, **local_configuration)
    scan.stop()
    scan.analyze()
<|endoftext|>


<commit_before>from __future__ import unicode_literals
from hubbot.moduleinterface import ModuleInterface
from hubbot.response import IRCResponse, ResponseType
from hubbot.Utils.RNN.model import Model
import cPickle
import os
import tensorflow as tf
import unicodedata

class RNNChatter(ModuleInterface):
    help = ""
    save_dir = os.path.join("hubbot", "data", "RNN")

    def __init__(self, bot):
        self.model = None
        self.chars = None
        self.vocab = None
        super(RNNChatter, self).__init__(bot)

    def on_load(self):
        try:
            saved_args = cPickle.load(os.path.join(self.save_dir, "config.pkl"))
            self.chars, self.vocab = cPickle.load(os.path.join(self.save_dir, "chars_vocab.pkl"))
            self.model = Model(saved_args, True)
        except:
            self.bot.logger.exception("Exception when loading module 'RNNChatter'")
            self.bot.modulehandler.unload_module("RNNChatter")

    def should_trigger(self, message):
        if message.type in self.accepted_types:
            return True

    def on_trigger(self, message):
        """
        @type message: hubbot.message.IRCMessage
        """
        if self.bot.nickname.lower() in message.message_string.lower() and len(message.message_list) > 1:
            reply = ""
            prime_string = message.message_string.replace(self.bot.nickname, "")
            with tf.Session() as sess:
                tf.global_variables_initializer().run()
                saver = tf.train.Saver(tf.global_variables())
                ckpt = tf.train.get_checkpoint_state(self.save_dir)
                if ckpt and ckpt.model_checkpoint_path:
                    saver.restore(sess, ckpt.model_checkpoint_path)
                    reply = self.model.sample(sess, self.chars, self.vocab, 100, prime_string, 2)
            nick_list = [nick.lower() for nick in self.bot.channels[message.reply_to].users.keys()]
            for nick in nick_list:
                if nick in reply.lower():
                    reply_list = reply.lower().split()
                    nick_index = self._index_containing_substring(reply_list, nick)
                    new_list = [item for item in reply_list if nick not in item]
                    new_list.insert(nick_index, message.user.name)
                    reply = " ".join(new_list)
            reply = self._clean_up_string(reply).replace("\n", " ")
            return IRCResponse(ResponseType.SAY, reply.capitalize(), message.reply_to)

    @staticmethod
    def _index_containing_substring(string_list, substring):
        """
        Given a list of strings, returns the index of the first element that contains a given substring
        If none exists, returns -1
        """
        for i, s in enumerate(string_list):
            if substring in s:
                return i
        return -1

    @staticmethod
    def _clean_up_string(string):
        new_string = "".join(c for c in string if ord(c) >= 0x20)
        while unicodedata.category(new_string[0]) not in ["Ll", "Lu"]:
            new_string = new_string[1:]
        return new_string.replace("(.+.+)", "")
<commit_msg>[RNNChatter] Use portable version of pickle if cPickle isn't available<commit_after>from __future__ import unicode_literals
from hubbot.moduleinterface import ModuleInterface
from hubbot.response import IRCResponse, ResponseType
from hubbot.Utils.RNN.model import Model
try:
    import cPickle as pickle
except ImportError:
    import pickle
import os
import tensorflow as tf
import unicodedata

class RNNChatter(ModuleInterface):
    help = ""
    save_dir = os.path.join("hubbot", "data", "RNN")

    def __init__(self, bot):
        self.model = None
        self.chars = None
        self.vocab = None
        super(RNNChatter, self).__init__(bot)

    def on_load(self):
        try:
            saved_args = pickle.load(os.path.join(self.save_dir, "config.pkl"))
            self.chars, self.vocab = pickle.load(os.path.join(self.save_dir, "chars_vocab.pkl"))
            self.model = Model(saved_args, True)
        except:
            self.bot.logger.exception("Exception when loading module 'RNNChatter'")
            self.bot.modulehandler.unload_module("RNNChatter")

    def should_trigger(self, message):
        if message.type in self.accepted_types:
            return True

    def on_trigger(self, message):
        """
        @type message: hubbot.message.IRCMessage
        """
        if self.bot.nickname.lower() in message.message_string.lower() and len(message.message_list) > 1:
            reply = ""
            prime_string = message.message_string.replace(self.bot.nickname, "")
            with tf.Session() as sess:
                tf.global_variables_initializer().run()
                saver = tf.train.Saver(tf.global_variables())
                ckpt = tf.train.get_checkpoint_state(self.save_dir)
                if ckpt and ckpt.model_checkpoint_path:
                    saver.restore(sess, ckpt.model_checkpoint_path)
                    reply = self.model.sample(sess, self.chars, self.vocab, 100, prime_string, 2)
            nick_list = [nick.lower() for nick in self.bot.channels[message.reply_to].users.keys()]
            for nick in nick_list:
                if nick in reply.lower():
                    reply_list = reply.lower().split()
                    nick_index = self._index_containing_substring(reply_list, nick)
                    new_list = [item for item in reply_list if nick not in item]
                    new_list.insert(nick_index, message.user.name)
                    reply = " ".join(new_list)
            reply = self._clean_up_string(reply).replace("\n", " ")
            return IRCResponse(ResponseType.SAY, reply.capitalize(), message.reply_to)

    @staticmethod
    def _index_containing_substring(string_list, substring):
        """
        Given a list of strings, returns the index of the first element that contains a given substring
        If none exists, returns -1
        """
        for i, s in enumerate(string_list):
            if substring in s:
                return i
        return -1

    @staticmethod
    def _clean_up_string(string):
        new_string = "".join(c for c in string if ord(c) >= 0x20)
        while unicodedata.category(new_string[0]) not in ["Ll", "Lu"]:
            new_string = new_string[1:]
        return new_string.replace("(.+.+)", "")
<|endoftext|>


<commit_before>from __future__ import absolute_import

import logging

from celery import shared_task, states
from celery.exceptions import Ignore

from django.conf import settings
from django.core.cache import cache

LOGGER = logging.getLogger(__name__)

REGISTRY_LIMIT_LAYERS = getattr(settings, 'REGISTRY_LIMIT_LAYERS', -1)
REGISTRY_SEARCH_URL = getattr(settings, 'REGISTRY_SEARCH_URL', None)

if REGISTRY_SEARCH_URL is None:
    SEARCH_ENABLED = False
    SEARCH_TYPE = None
    SEARCH_URL = None
else:
    SEARCH_ENABLED = True
    SEARCH_TYPE = REGISTRY_SEARCH_URL.split('+')[0]
    SEARCH_URL = REGISTRY_SEARCH_URL.split('+')[1]

if REGISTRY_LIMIT_LAYERS > 0:
    DEBUG_SERVICES = True
    DEBUG_LAYERS_NUMBER = REGISTRY_LIMIT_LAYERS
else:
    DEBUG_SERVICES = False
    DEBUG_LAYERS_NUMBER = -1

@shared_task(bind=True)
def check_all_services(self):
    from hypermap.aggregator.models import Service
    service_to_processes = Service.objects.filter(active=True)
    for service in service_to_processes:
        check_service.delay(service.id)

@shared_task(bind=True)
def check_service(self, service_id):
    from hypermap.aggregator.models import Service
    service = Service.objects.get(pk=service_id)

    # 1. update layers and check service
    if getattr(settings, 'REGISTRY_HARVEST_SERVICES', True):
        service.update_layers()

    layer_to_process = service.layer_set.all()
    if DEBUG_SERVICES:
        layer_to_process = layer_to_process[0:DEBUG_LAYERS_NUMBER]

    service.check_available()

    # 2. check layers if the service is monitored and the layer is monitored
    if service.is_monitored:
        for layer in layer_to_process:
            if layer.is_monitored:
                if not settings.REGISTRY_SKIP_CELERY:
                    check_layer.delay(layer.id)
                else:
                    check_layer(layer.id)

    # 3. index layers
    if getattr(settings, 'REGISTRY_HARVEST_SERVICES', True):
        if not settings.REGISTRY_SKIP_CELERY:
            index_service.delay(service.id)
        else:
            index_service(service.id)

@shared_task(bind=True, soft_time_limit=10)
def check_layer(self, layer_id):
    from hypermap.aggregator.models import Layer
    layer = Layer.objects.get(pk=layer_id)
    LOGGER.debug('Checking layer %s' % layer.name)
    success, message = layer.check_available()
    # every time a layer is checked it should be indexed
    # for now we remove indexing but we do it using a scheduled task unless SKIP_CELERY_TASK
    if success and SEARCH_ENABLED:
        if settings.REGISTRY_SKIP_CELERY:
            index_layer(layer.id)
        else:
            index_layer(layer.id, use_cache=True)

@shared_task(bind=True)
def index_cached_layers(self):
    """
    Index and unindex all layers in the Django cache (Index all layers who have been checked).
    """
    from hypermap.aggregator.models import Layer

    if SEARCH_TYPE == 'solr':
        from hypermap.aggregator.solr import SolrHypermap
        solrobject = SolrHypermap()
    else:
        from hypermap.aggregator.elasticsearch_client import ESHypermap
        from elasticsearch import helpers
        es_client = ESHypermap()

    layers_cache = cache.get('layers')
    deleted_layers_cache = cache.get('deleted_layers')

    # 1. added layers cache
    if layers_cache:
        layers_list = list(layers_cache)
        LOGGER.debug('There are %s layers in cache: %s' % (len(layers_list), layers_list))

        batch_size = settings.REGISTRY_SEARCH_BATCH_SIZE
        batch_lists = [layers_list[i:i+batch_size] for i in range(0, len(layers_list), batch_size)]

        for batch_list_ids in batch_lists:
            layers = Layer.objects.filter(id__in=batch_list_ids)

            if batch_size > len(layers):
                batch_size = len(layers)

            LOGGER.debug('Syncing %s/%s layers to %s: %s' % (batch_size, len(layers_cache), layers, SEARCH_TYPE))

            try:
                # SOLR
                if SEARCH_TYPE == 'solr':
                    success, layers_errors_ids = solrobject.layers_to_solr(layers)
                    if success:
                        # remove layers from cache here
                        layers_cache = layers_cache.difference(set(batch_list_ids))
                        LOGGER.debug('Removing layers with id %s from cache' % batch_list_ids)
                        cache.set('layers', layers_cache)
                # ES
                elif SEARCH_TYPE == 'elasticsearch':
                    with_bulk, success = True, False
                    layers_to_index = [es_client.layer_to_es(layer, with_bulk) for layer in layers]
                    message = helpers.bulk(es_client.es, layers_to_index)

                    # Check that all layers where indexed...if not, don't clear cache.
                    # TODO: Check why es does not index all layers at first.
                    len_indexed_layers = message[0]
                    if len_indexed_layers == len(layers):
                        LOGGER.debug('%d layers indexed successfully' % (len_indexed_layers))
                        success = True
                    if success:
                        # remove layers from cache here
                        layers_cache = layers_cache.difference(set(batch_list_ids))
                        cache.set('layers', layers_cache)
                else:
                    raise Exception("Incorrect SEARCH_TYPE=%s" % SEARCH_TYPE)
            except Exception as e:
                LOGGER.error('Layers were NOT indexed correctly')
                LOGGER.error(e, exc_info=True)
    else:
        LOGGER.debug('No cached layers to add in search engine.')

    # 2. deleted layers cache
    if deleted_layers_cache:
        layers_list = list(deleted_layers_cache)
        LOGGER.debug('There are %s layers in cache for deleting: %s' % (len(layers_list), layers_list))
        # TODO implement me: batch layer index deletion
        for layer_id in layers_list:
            # SOLR
            if SEARCH_TYPE == 'solr':
                if Layer.objects.filter(pk=layer_id).exists():
                    layer = Layer.objects.get(id=layer_id)
                    unindex_layer(layer.id, use_cache=False)
                    deleted_layers_cache = deleted_layers_cache.difference(set([layer_id]))
                    cache.set('deleted_layers', deleted_layers_cache)
            else:
                # TODO implement me
                raise NotImplementedError
    else:
        LOGGER.debug('No cached layers to remove in search engine.')

@shared_task(name="clear_index")
def clear_index():
    if SEARCH_TYPE == 'solr':
        LOGGER.debug('Clearing the solr indexes')
        from hypermap.aggregator.solr import SolrHypermap
        solrobject = SolrHypermap()
        solrobject.clear_solr()
    elif SEARCH_TYPE == 'elasticsearch':
        LOGGER.debug('Clearing the ES indexes')
        from hypermap.aggregator.elasticsearch_client import ESHypermap
        esobject = ESHypermap()
        esobject.clear_es()

@shared_task(bind=True)
def remove_service_checks(self, service_id):
    """
    Remove all checks from a service.
    """
    from hypermap.aggregator.models import Service
    service = Service.objects.get(id=service_id)

    service.check_set.all().delete()
    layer_to_process = service.layer_set.all()
    for layer in layer_to_process:
        layer.check_set.all().delete()

@shared_task(bind=True)
def index_service(self, service_id):
    """
    Index a service in search engine.
    """

    from hypermap.aggregator.models import Service
    service = Service.objects.get(id=service_id)

    if not service.is_valid:
        LOGGER.debug('Not indexing service with id %s in search engine as it is not valid' % service.id)
        return

    LOGGER.debug('Indexing service %s' % service.id)
    layer_to_process = service.layer_set.all()

    for layer in layer_to_process:
        if not settings.REGISTRY_SKIP_CELERY:
            index_layer(layer.id, use_cache=True)
        else:
            index_layer(layer.id)

@shared_task(bind=True)
def index_layer(self, layer_id, use_cache=False):
    """
    Index a layer in the search backend.
    If cache is set, append it to the list, if it isn't send the transaction right away.
    cache needs memcached to be available.
    """

    from hypermap.aggregator.models import Layer
    layer = Layer.objects.get(id=layer_id)

    if not layer.is_valid:
        LOGGER.debug('Not indexing or removing layer with id %s in search engine as it is not valid' % layer.id)
        unindex_layer(layer.id, use_cache)
        return

    if layer.was_deleted:
        LOGGER.debug('Not indexing or removing layer with id %s in search engine as was_deleted is true' % layer.id)
        unindex_layer(layer.id, use_cache)
        return

    # 1. if we use cache
    if use_cache:
        LOGGER.debug('Caching layer with id %s for syncing with search engine' % layer.id)
        layers = cache.get('layers')
        if layers is None:
            layers = set([layer.id])
        else:
            layers.add(layer.id)
        cache.set('layers', layers)
        return

    # 2. if we don't use cache
    # TODO: Make this function more DRY
    # by abstracting the common bits.
    if SEARCH_TYPE == 'solr':
        from hypermap.aggregator.solr import SolrHypermap
        LOGGER.debug('Syncing layer %s to solr' % layer.name)
        solrobject = SolrHypermap()
        success, message = solrobject.layer_to_solr(layer)
        # update the error message if using celery
        if not settings.REGISTRY_SKIP_CELERY:
            if not success:
                self.update_state(
                    state=states.FAILURE,
                    meta=message
                    )
                raise Ignore()
    elif SEARCH_TYPE == 'elasticsearch':
        from hypermap.aggregator.elasticsearch_client import ESHypermap
        LOGGER.debug('Syncing layer %s to es' % layer.name)
        esobject = ESHypermap()
        success, message = esobject.layer_to_es(layer)
        # update the error message if using celery
        if not settings.REGISTRY_SKIP_CELERY:
            if not success:
                self.update_state(
                    state=states.FAILURE,
                    meta=message
                    )
                raise Ignore()

@shared_task(bind=True)
def unindex_layers_with_issues(self, use_cache=False):
    """
    Remove the index for layers in search backend, which are linked to an issue.
    """
    from hypermap.aggregator.models import Issue, Layer, Service
    from django.contrib.contenttypes.models import ContentType

    layer_type = ContentType.objects.get_for_model(Layer)
    service_type = ContentType.objects.get_for_model(Service)

    for issue in Issue.objects.filter(content_type__pk=layer_type.id):
        unindex_layer(issue.content_object.id, use_cache)

    for issue in Issue.objects.filter(content_type__pk=service_type.id):
        for layer in issue.content_object.layer_set.all():
            unindex_layer(layer.id, use_cache)

@shared_task(bind=True)
def unindex_layer(self, layer_id, use_cache=False):
    """
    Remove the index for a layer in the search backend.
    If cache is set, append it to the list of removed layers, if it isn't send the transaction right away.
    """

    from hypermap.aggregator.models import Layer
    layer = Layer.objects.get(id=layer_id)

    if use_cache:
        LOGGER.debug('Caching layer with id %s for being removed from search engine' % layer.id)
        deleted_layers = cache.get('deleted_layers')
        if deleted_layers is None:
            deleted_layers = set([layer.id])
        else:
            deleted_layers.add(layer.id)
        cache.set('deleted_layers', deleted_layers)
        return

    if SEARCH_TYPE == 'solr':
        from hypermap.aggregator.solr import SolrHypermap
        LOGGER.debug('Removing layer %s from solr' % layer.id)
        try:
            solrobject = SolrHypermap()
            solrobject.remove_layer(layer.uuid)
        except Exception:
            LOGGER.error('Layer NOT correctly removed from Solr')
    elif SEARCH_TYPE == 'elasticsearch':
        # TODO implement me
        pass

@shared_task(bind=True)
def index_all_layers(self):
    """
    Index all layers in search engine.
    """
    from hypermap.aggregator.models import Layer

    if not settings.REGISTRY_SKIP_CELERY:
        layers_cache = set(Layer.objects.filter(is_valid=True).values_list('id', flat=True))
        deleted_layers_cache = set(Layer.objects.filter(is_valid=False).values_list('id', flat=True))
        cache.set('layers', layers_cache)
        cache.set('deleted_layers', deleted_layers_cache)
    else:
        for layer in Layer.objects.all():
            index_layer(layer.id)

@shared_task(bind=True)
def update_last_wm_layers(self, service_id, num_layers=10):
    """
    Update and index the last added and deleted layers (num_layers) in WorldMap service.
    """
    from hypermap.aggregator.models import Service

    LOGGER.debug(
        'Updating the index the last %s added and %s deleted layers in WorldMap service'
        % (num_layers, num_layers)
                )
    service = Service.objects.get(id=service_id)
    # TODO raise error if service type is not WM type
    if service.type == 'Hypermap:WorldMapLegacy':
        from hypermap.aggregator.models import update_layers_wm_legacy as update_layers_wm
    if service.type == 'Hypermap:WorldMap':
        from hypermap.aggregator.models import update_layers_geonode_wm as update_layers_wm

    update_layers_wm(service, num_layers)

    # Remove in search engine last num_layers that were deleted
    LOGGER.debug('Removing the index for the last %s deleted layers' % num_layers)
    layer_to_unindex = service.layer_set.filter(was_deleted=True).order_by('-last_updated')[0:num_layers]
    for layer in layer_to_unindex:
        if not settings.REGISTRY_SKIP_CELERY:
            unindex_layer(layer.id, use_cache=True)
        else:
            unindex_layer(layer.id)

    # Add/Update in search engine last num_layers that were added
    LOGGER.debug('Adding/Updating the index for the last %s added layers' % num_layers)
    layer_to_index = service.layer_set.filter(was_deleted=False).order_by('-last_updated')[0:num_layers]
    for layer in layer_to_index:
        if not settings.REGISTRY_SKIP_CELERY:
            index_layer(layer.id, use_cache=True)
        else:
            index_layer(layer.id)

@shared_task(bind=True)
def update_endpoint(self, endpoint_id, greedy_opt=False):
    from hypermap.aggregator.utils import create_services_from_endpoint
    from hypermap.aggregator.models import Endpoint

    endpoint = Endpoint.objects.get(id=endpoint_id)

    LOGGER.debug('Processing endpoint with id %s: %s' % (endpoint.id, endpoint.url))

    # Override the greedy_opt var with the value from the endpoint list
    # if it's available.
    if endpoint.endpoint_list:
        greedy_opt = endpoint.endpoint_list.greedy

    imported, message = create_services_from_endpoint(endpoint.url, greedy_opt=greedy_opt, catalog=endpoint.catalog)

    # this update will not execute the endpoint_post_save signal.
    Endpoint.objects.filter(id=endpoint.id).update(
        imported=imported, message=message, processed=True
    )

@shared_task(bind=True)
def update_endpoints(self, endpoint_list_id):
    from hypermap.aggregator.models import EndpointList
    endpoint_list_to_process = EndpointList.objects.get(id=endpoint_list_id)
    # for now we process the enpoint even if they were already processed
    endpoint_to_process = endpoint_list_to_process.endpoint_set.filter(processed=False)
    if not settings.REGISTRY_SKIP_CELERY:
        for endpoint in endpoint_to_process:
            update_endpoint.delay(endpoint.id)
    else:
        for endpoint in endpoint_to_process:
            update_endpoint(endpoint.id)
    return True
<commit_msg>Prevent a layer to be checked if it is marded as was_deleted<commit_after>from __future__ import absolute_import

import logging

from celery import shared_task, states
from celery.exceptions import Ignore

from django.conf import settings
from django.core.cache import cache

LOGGER = logging.getLogger(__name__)

REGISTRY_LIMIT_LAYERS = getattr(settings, 'REGISTRY_LIMIT_LAYERS', -1)
REGISTRY_SEARCH_URL = getattr(settings, 'REGISTRY_SEARCH_URL', None)

if REGISTRY_SEARCH_URL is None:
    SEARCH_ENABLED = False
    SEARCH_TYPE = None
    SEARCH_URL = None
else:
    SEARCH_ENABLED = True
    SEARCH_TYPE = REGISTRY_SEARCH_URL.split('+')[0]
    SEARCH_URL = REGISTRY_SEARCH_URL.split('+')[1]

if REGISTRY_LIMIT_LAYERS > 0:
    DEBUG_SERVICES = True
    DEBUG_LAYERS_NUMBER = REGISTRY_LIMIT_LAYERS
else:
    DEBUG_SERVICES = False
    DEBUG_LAYERS_NUMBER = -1

@shared_task(bind=True)
def check_all_services(self):
    from hypermap.aggregator.models import Service
    service_to_processes = Service.objects.filter(active=True)
    for service in service_to_processes:
        check_service.delay(service.id)

@shared_task(bind=True)
def check_service(self, service_id):
    from hypermap.aggregator.models import Service
    service = Service.objects.get(pk=service_id)

    # 1. update layers and check service
    if getattr(settings, 'REGISTRY_HARVEST_SERVICES', True):
        service.update_layers()

    layer_to_process = service.layer_set.all()
    if DEBUG_SERVICES:
        layer_to_process = layer_to_process[0:DEBUG_LAYERS_NUMBER]

    service.check_available()

    # 2. check layers if the service is monitored and the layer is monitored
    if service.is_monitored:
        for layer in layer_to_process:
            if layer.is_monitored and not layer.was_deleted:
                if not settings.REGISTRY_SKIP_CELERY:
                    check_layer.delay(layer.id)
                else:
                    check_layer(layer.id)

    # 3. index layers
    if getattr(settings, 'REGISTRY_HARVEST_SERVICES', True):
        if not settings.REGISTRY_SKIP_CELERY:
            index_service.delay(service.id)
        else:
            index_service(service.id)

@shared_task(bind=True, soft_time_limit=10)
def check_layer(self, layer_id):
    from hypermap.aggregator.models import Layer
    layer = Layer.objects.get(pk=layer_id)
    if layer.was_deleted:
        LOGGER.debug('Cannot check layer %s as it is marked as deleted!' % layer.name)
    else:
        LOGGER.debug('Checking layer %s' % layer.name)
        success, message = layer.check_available()
        # every time a layer is checked it should be indexed
        # for now we remove indexing but we do it using a scheduled task unless SKIP_CELERY_TASK
        if success and SEARCH_ENABLED:
            if settings.REGISTRY_SKIP_CELERY:
                index_layer(layer.id)
            else:
                index_layer(layer.id, use_cache=True)

@shared_task(bind=True)
def index_cached_layers(self):
    """
    Index and unindex all layers in the Django cache (Index all layers who have been checked).
    """
    from hypermap.aggregator.models import Layer

    if SEARCH_TYPE == 'solr':
        from hypermap.aggregator.solr import SolrHypermap
        solrobject = SolrHypermap()
    else:
        from hypermap.aggregator.elasticsearch_client import ESHypermap
        from elasticsearch import helpers
        es_client = ESHypermap()

    layers_cache = cache.get('layers')
    deleted_layers_cache = cache.get('deleted_layers')

    # 1. added layers cache
    if layers_cache:
        layers_list = list(layers_cache)
        LOGGER.debug('There are %s layers in cache: %s' % (len(layers_list), layers_list))

        batch_size = settings.REGISTRY_SEARCH_BATCH_SIZE
        batch_lists = [layers_list[i:i+batch_size] for i in range(0, len(layers_list), batch_size)]

        for batch_list_ids in batch_lists:
            layers = Layer.objects.filter(id__in=batch_list_ids)

            if batch_size > len(layers):
                batch_size = len(layers)

            LOGGER.debug('Syncing %s/%s layers to %s: %s' % (batch_size, len(layers_cache), layers, SEARCH_TYPE))

            try:
                # SOLR
                if SEARCH_TYPE == 'solr':
                    success, layers_errors_ids = solrobject.layers_to_solr(layers)
                    if success:
                        # remove layers from cache here
                        layers_cache = layers_cache.difference(set(batch_list_ids))
                        LOGGER.debug('Removing layers with id %s from cache' % batch_list_ids)
                        cache.set('layers', layers_cache)
                # ES
                elif SEARCH_TYPE == 'elasticsearch':
                    with_bulk, success = True, False
                    layers_to_index = [es_client.layer_to_es(layer, with_bulk) for layer in layers]
                    message = helpers.bulk(es_client.es, layers_to_index)

                    # Check that all layers where indexed...if not, don't clear cache.
                    # TODO: Check why es does not index all layers at first.
                    len_indexed_layers = message[0]
                    if len_indexed_layers == len(layers):
                        LOGGER.debug('%d layers indexed successfully' % (len_indexed_layers))
                        success = True
                    if success:
                        # remove layers from cache here
                        layers_cache = layers_cache.difference(set(batch_list_ids))
                        cache.set('layers', layers_cache)
                else:
                    raise Exception("Incorrect SEARCH_TYPE=%s" % SEARCH_TYPE)
            except Exception as e:
                LOGGER.error('Layers were NOT indexed correctly')
                LOGGER.error(e, exc_info=True)
    else:
        LOGGER.debug('No cached layers to add in search engine.')

    # 2. deleted layers cache
    if deleted_layers_cache:
        layers_list = list(deleted_layers_cache)
        LOGGER.debug('There are %s layers in cache for deleting: %s' % (len(layers_list), layers_list))
        # TODO implement me: batch layer index deletion
        for layer_id in layers_list:
            # SOLR
            if SEARCH_TYPE == 'solr':
                if Layer.objects.filter(pk=layer_id).exists():
                    layer = Layer.objects.get(id=layer_id)
                    unindex_layer(layer.id, use_cache=False)
                    deleted_layers_cache = deleted_layers_cache.difference(set([layer_id]))
                    cache.set('deleted_layers', deleted_layers_cache)
            else:
                # TODO implement me
                raise NotImplementedError
    else:
        LOGGER.debug('No cached layers to remove in search engine.')

@shared_task(name="clear_index")
def clear_index():
    if SEARCH_TYPE == 'solr':
        LOGGER.debug('Clearing the solr indexes')
        from hypermap.aggregator.solr import SolrHypermap
        solrobject = SolrHypermap()
        solrobject.clear_solr()
    elif SEARCH_TYPE == 'elasticsearch':
        LOGGER.debug('Clearing the ES indexes')
        from hypermap.aggregator.elasticsearch_client import ESHypermap
        esobject = ESHypermap()
        esobject.clear_es()

@shared_task(bind=True)
def remove_service_checks(self, service_id):
    """
    Remove all checks from a service.
    """
    from hypermap.aggregator.models import Service
    service = Service.objects.get(id=service_id)

    service.check_set.all().delete()
    layer_to_process = service.layer_set.all()
    for layer in layer_to_process:
        layer.check_set.all().delete()

@shared_task(bind=True)
def index_service(self, service_id):
    """
    Index a service in search engine.
    """

    from hypermap.aggregator.models import Service
    service = Service.objects.get(id=service_id)

    if not service.is_valid:
        LOGGER.debug('Not indexing service with id %s in search engine as it is not valid' % service.id)
        return

    LOGGER.debug('Indexing service %s' % service.id)
    layer_to_process = service.layer_set.all()

    for layer in layer_to_process:
        if not settings.REGISTRY_SKIP_CELERY:
            index_layer(layer.id, use_cache=True)
        else:
            index_layer(layer.id)

@shared_task(bind=True)
def index_layer(self, layer_id, use_cache=False):
    """
    Index a layer in the search backend.
    If cache is set, append it to the list, if it isn't send the transaction right away.
    cache needs memcached to be available.
    """

    from hypermap.aggregator.models import Layer
    layer = Layer.objects.get(id=layer_id)

    if not layer.is_valid:
        LOGGER.debug('Not indexing or removing layer with id %s in search engine as it is not valid' % layer.id)
        unindex_layer(layer.id, use_cache)
        return

    if layer.was_deleted:
        LOGGER.debug('Not indexing or removing layer with id %s in search engine as was_deleted is true' % layer.id)
        unindex_layer(layer.id, use_cache)
        return

    # 1. if we use cache
    if use_cache:
        LOGGER.debug('Caching layer with id %s for syncing with search engine' % layer.id)
        layers = cache.get('layers')
        if layers is None:
            layers = set([layer.id])
        else:
            layers.add(layer.id)
        cache.set('layers', layers)
        return

    # 2. if we don't use cache
    # TODO: Make this function more DRY
    # by abstracting the common bits.
    if SEARCH_TYPE == 'solr':
        from hypermap.aggregator.solr import SolrHypermap
        LOGGER.debug('Syncing layer %s to solr' % layer.name)
        solrobject = SolrHypermap()
        success, message = solrobject.layer_to_solr(layer)
        # update the error message if using celery
        if not settings.REGISTRY_SKIP_CELERY:
            if not success:
                self.update_state(
                    state=states.FAILURE,
                    meta=message
                    )
                raise Ignore()
    elif SEARCH_TYPE == 'elasticsearch':
        from hypermap.aggregator.elasticsearch_client import ESHypermap
        LOGGER.debug('Syncing layer %s to es' % layer.name)
        esobject = ESHypermap()
        success, message = esobject.layer_to_es(layer)
        # update the error message if using celery
        if not settings.REGISTRY_SKIP_CELERY:
            if not success:
                self.update_state(
                    state=states.FAILURE,
                    meta=message
                    )
                raise Ignore()

@shared_task(bind=True)
def unindex_layers_with_issues(self, use_cache=False):
    """
    Remove the index for layers in search backend, which are linked to an issue.
    """
    from hypermap.aggregator.models import Issue, Layer, Service
    from django.contrib.contenttypes.models import ContentType

    layer_type = ContentType.objects.get_for_model(Layer)
    service_type = ContentType.objects.get_for_model(Service)

    for issue in Issue.objects.filter(content_type__pk=layer_type.id):
        unindex_layer(issue.content_object.id, use_cache)

    for issue in Issue.objects.filter(content_type__pk=service_type.id):
        for layer in issue.content_object.layer_set.all():
            unindex_layer(layer.id, use_cache)

@shared_task(bind=True)
def unindex_layer(self, layer_id, use_cache=False):
    """
    Remove the index for a layer in the search backend.
    If cache is set, append it to the list of removed layers, if it isn't send the transaction right away.
    """

    from hypermap.aggregator.models import Layer
    layer = Layer.objects.get(id=layer_id)

    if use_cache:
        LOGGER.debug('Caching layer with id %s for being removed from search engine' % layer.id)
        deleted_layers = cache.get('deleted_layers')
        if deleted_layers is None:
            deleted_layers = set([layer.id])
        else:
            deleted_layers.add(layer.id)
        cache.set('deleted_layers', deleted_layers)
        return

    if SEARCH_TYPE == 'solr':
        from hypermap.aggregator.solr import SolrHypermap
        LOGGER.debug('Removing layer %s from solr' % layer.id)
        try:
            solrobject = SolrHypermap()
            solrobject.remove_layer(layer.uuid)
        except Exception:
            LOGGER.error('Layer NOT correctly removed from Solr')
    elif SEARCH_TYPE == 'elasticsearch':
        # TODO implement me
        pass

@shared_task(bind=True)
def index_all_layers(self):
    """
    Index all layers in search engine.
    """
    from hypermap.aggregator.models import Layer

    if not settings.REGISTRY_SKIP_CELERY:
        layers_cache = set(Layer.objects.filter(is_valid=True).values_list('id', flat=True))
        deleted_layers_cache = set(Layer.objects.filter(is_valid=False).values_list('id', flat=True))
        cache.set('layers', layers_cache)
        cache.set('deleted_layers', deleted_layers_cache)
    else:
        for layer in Layer.objects.all():
            index_layer(layer.id)

@shared_task(bind=True)
def update_last_wm_layers(self, service_id, num_layers=10):
    """
    Update and index the last added and deleted layers (num_layers) in WorldMap service.
    """
    from hypermap.aggregator.models import Service

    LOGGER.debug(
        'Updating the index the last %s added and %s deleted layers in WorldMap service'
        % (num_layers, num_layers)
                )
    service = Service.objects.get(id=service_id)
    # TODO raise error if service type is not WM type
    if service.type == 'Hypermap:WorldMapLegacy':
        from hypermap.aggregator.models import update_layers_wm_legacy as update_layers_wm
    if service.type == 'Hypermap:WorldMap':
        from hypermap.aggregator.models import update_layers_geonode_wm as update_layers_wm

    update_layers_wm(service, num_layers)

    # Remove in search engine last num_layers that were deleted
    LOGGER.debug('Removing the index for the last %s deleted layers' % num_layers)
    layer_to_unindex = service.layer_set.filter(was_deleted=True).order_by('-last_updated')[0:num_layers]
    for layer in layer_to_unindex:
        if not settings.REGISTRY_SKIP_CELERY:
            unindex_layer(layer.id, use_cache=True)
        else:
            unindex_layer(layer.id)

    # Add/Update in search engine last num_layers that were added
    LOGGER.debug('Adding/Updating the index for the last %s added layers' % num_layers)
    layer_to_index = service.layer_set.filter(was_deleted=False).order_by('-last_updated')[0:num_layers]
    for layer in layer_to_index:
        if not settings.REGISTRY_SKIP_CELERY:
            index_layer(layer.id, use_cache=True)
        else:
            index_layer(layer.id)

@shared_task(bind=True)
def update_endpoint(self, endpoint_id, greedy_opt=False):
    from hypermap.aggregator.utils import create_services_from_endpoint
    from hypermap.aggregator.models import Endpoint

    endpoint = Endpoint.objects.get(id=endpoint_id)

    LOGGER.debug('Processing endpoint with id %s: %s' % (endpoint.id, endpoint.url))

    # Override the greedy_opt var with the value from the endpoint list
    # if it's available.
    if endpoint.endpoint_list:
        greedy_opt = endpoint.endpoint_list.greedy

    imported, message = create_services_from_endpoint(endpoint.url, greedy_opt=greedy_opt, catalog=endpoint.catalog)

    # this update will not execute the endpoint_post_save signal.
    Endpoint.objects.filter(id=endpoint.id).update(
        imported=imported, message=message, processed=True
    )

@shared_task(bind=True)
def update_endpoints(self, endpoint_list_id):
    from hypermap.aggregator.models import EndpointList
    endpoint_list_to_process = EndpointList.objects.get(id=endpoint_list_id)
    # for now we process the enpoint even if they were already processed
    endpoint_to_process = endpoint_list_to_process.endpoint_set.filter(processed=False)
    if not settings.REGISTRY_SKIP_CELERY:
        for endpoint in endpoint_to_process:
            update_endpoint.delay(endpoint.id)
    else:
        for endpoint in endpoint_to_process:
            update_endpoint(endpoint.id)
    return True
<|endoftext|>