<commit_before>"""
doctr

A tool to automatically deploy docs to GitHub pages from Travis CI.

The doctr command is two commands in one. To use, first run

doctr configure

on your local machine. This will prompt for your GitHub credentials and the
name of the repo you want to deploy docs for. This will generate a secure key,
which you should insert into your .travis.yml.

Then, on Travis, for the build where you build your docs, add

    - doctr deploy

to the end of the build to deploy the docs to GitHub pages.  This will only
run on the master branch, and won't run on pull requests.

For more information, see https://drdoctr.github.io/doctr/docs/
"""

import sys
import os
import argparse
import shlex
import subprocess

from textwrap import dedent

from .local import (generate_GitHub_token, encrypt_variable, encrypt_file,
    upload_GitHub_deploy_key, generate_ssh_key, check_repo_exists, GitHub_login)
from .travis import (setup_GitHub_push, commit_docs, push_docs,
    get_current_repo, sync_from_log, find_sphinx_build_dir, run)
from . import __version__

def get_parser():
    # This uses RawTextHelpFormatter so that the description (the docstring of
    # this module) is formatted correctly. Unfortunately, that means that
    # parser help is not text wrapped (but all other help is).
    parser = argparse.ArgumentParser(description=__doc__,
        formatter_class=argparse.RawTextHelpFormatter, epilog="""
Run --help on the subcommands like 'doctr deploy --help' to see the
options available.
        """,
        )
    parser.add_argument('-V', '--version', action='version', version='doctr ' + __version__)

    subcommand = parser.add_subparsers(title='subcommand', dest='subcommand')
    deploy_parser = subcommand.add_parser('deploy', help="""Deploy the docs to GitHub from Travis.""")
    deploy_parser.set_defaults(func=deploy)
    deploy_parser.add_argument('--force', action='store_true', help="""Run the deploy command even
    if we do not appear to be on Travis.""")
    deploy_parser.add_argument('--token', action='store_true', default=False,
        help="""Push to GitHub using a personal access token. Use this if you
        used 'doctr configure --token'.""")
    deploy_parser.add_argument('--key-path', default='github_deploy_key.enc',
        help="""Path of the encrypted GitHub deploy key. The default is %(default)r.""")
    deploy_parser.add_argument('--built-docs', default=None,
        help="""Location of the built html documentation to be deployed to
        gh-pages. If not specified, Doctr will try to automatically detect build location""")
    deploy_parser.add_argument('--gh-pages-docs', default='docs',
        help="""Directory to deploy the html documentation to on gh-pages. The
        default is %(default)r.""")
    deploy_parser.add_argument('--tmp-dir', default=None,
        help=argparse.SUPPRESS)
    deploy_parser.add_argument('--deploy-repo', default=None, help="""Repo to
        deploy the docs to. By default, it deploys to the repo Doctr is run from.""")
    deploy_parser.add_argument('--no-require-master', dest='require_master', action='store_false',
        default=True, help="""Allow docs to be pushed from a branch other than master""")
    deploy_parser.add_argument('--command', default=None, help="""Command to
        be run before committing and pushing. If the command creates
        additional files that should be deployed, they should be added to the
        index.""")
    deploy_parser.add_argument('--no-sync', dest='sync', action='store_false',
        default=True, help="""Don't sync any files. This is generally used in
        conjunction with the --command flag, for instance, if the command syncs
        the files for you. Any files you wish to commit should be added to the
        index.""")
    deploy_parser.add_argument('--no-push', dest='push', action='store_false',
        default=True, help="Run all the steps except the last push step."
        "Useful for debugging")

    configure_parser = subcommand.add_parser('configure', help="Configure doctr. This command should be run locally (not on Travis).")
    configure_parser.set_defaults(func=configure)
    configure_parser.add_argument('--force', action='store_true', help="""Run the configure command even
    if we appear to be on Travis.""")
    configure_parser.add_argument('--token', action="store_true", default=False,
        help="""Generate a personal access token to push to GitHub. The default is to use a
        deploy key. WARNING: This will grant read/write access to all the
        public repositories for the user. This option is not recommended
        unless you are using a separate GitHub user for deploying.""")
    configure_parser.add_argument("--no-upload-key", action="store_false", default=True,
        dest="upload_key", help="""Don't automatically upload the deploy key to GitHub. If you select this
        option, you will not be prompted for your GitHub credentials, so this option is not compatible with
        private repositories.""")
    configure_parser.add_argument('--key-path', default='github_deploy_key',
        help="""Path to save the encrypted GitHub deploy key. The default is %(default)r.
    The .enc extension is added to the file automatically.""")

    return parser

def process_args(parser):
    args = parser.parse_args()

    if not args.subcommand:
        parser.print_usage()
        parser.exit(1)

    try:
        return args.func(args, parser)
    except RuntimeError as e:
        sys.exit("Error: " + e.args[0])

def on_travis():
    return os.environ.get("TRAVIS_JOB_NUMBER", '')

def deploy(args, parser):
    if not args.force and not on_travis():
        parser.error("doctr does not appear to be running on Travis. Use "
            "doctr deploy --force to run anyway.")

    if args.tmp_dir:
        parser.error("The --tmp-dir flag has been removed (doctr no longer uses a temporary directory when deploying).")

    build_repo = get_current_repo()
    deploy_repo = args.deploy_repo or build_repo

    current_commit = subprocess.check_output(['git', 'rev-parse', 'HEAD']).decode('utf-8').strip()
    try:
        can_push = setup_GitHub_push(deploy_repo, auth_type='token' if args.token else
                             'deploy_key', full_key_path=args.key_path,
                             require_master=args.require_master)

        if args.sync:
            built_docs = args.built_docs or find_sphinx_build_dir()

            log_file = os.path.join(args.gh_pages_docs, '.doctr-files')

            print("Moving built docs into place")
            added, removed = sync_from_log(src=built_docs,
                dst=args.gh_pages_docs, log_file=log_file)

        else:
            added, removed = [], []

        if args.command:
            run(shlex.split(args.command))

        changes = commit_docs(added=added, removed=removed)
        if changes:
            if can_push and args.push:
                push_docs()
            else:
                print("Don't have permission to push. Not trying.")
        else:
            print("The docs have not changed. Not updating")
    finally:
        subprocess.run(['git', 'checkout', current_commit])

class IncrementingInt:
    def __init__(self, i=0):
        self.i = i

    def __repr__(self):
        ret = repr(self.i)
        self.i += 1
        return ret

    __str__ = __repr__

def configure(args, parser):
    if not args.force and on_travis():
        parser.error("doctr appears to be running on Travis. Use "
            "doctr configure --force to run anyway.")

    if args.upload_key:
        login_kwargs = GitHub_login()
    else:
        login_kwargs = {'auth': None, 'headers': None}

    build_repo = input("What repo do you want to build the docs for (org/reponame, like 'drdoctr/doctr')? ")
    is_private = check_repo_exists(build_repo, service='github', **login_kwargs)
    check_repo_exists(build_repo, service='travis')

    deploy_repo = input("What repo do you want to deploy the docs to? [{build_repo}] ".format(build_repo=build_repo))
    if not deploy_repo:
        deploy_repo = build_repo

    if deploy_repo != build_repo:
        check_repo_exists(deploy_repo, service='github', **login_kwargs)

    N = IncrementingInt(1)

    header = "\n================== You should now do the following ==================\n"

    if args.token:
        token = generate_GitHub_token(**login_kwargs)['token']
        encrypted_variable = encrypt_variable("GH_TOKEN={token}".format(token=token).encode('utf-8'),
            build_repo=build_repo, is_private=is_private, **login_kwargs)
        print(dedent("""
        A personal access token for doctr has been created.

        You can go to https://github.com/settings/tokens to revoke it."""))

        print(header)
    else:
        ssh_key = generate_ssh_key("doctr deploy key for {deploy_repo}".format(deploy_repo=deploy_repo), keypath=args.key_path)
        key = encrypt_file(args.key_path, delete=True)
        encrypted_variable = encrypt_variable(b"DOCTR_DEPLOY_ENCRYPTION_KEY=" + key, build_repo=build_repo, is_private=is_private, **login_kwargs)

        deploy_keys_url = 'https://github.com/{deploy_repo}/settings/keys'.format(deploy_repo=deploy_repo)

        if args.upload_key:

            upload_GitHub_deploy_key(deploy_repo, ssh_key, **login_kwargs)

            print(dedent("""
            The deploy key has been added for {deploy_repo}.

            You can go to {deploy_keys_url} to revoke the deploy key.\
            """.format(deploy_repo=deploy_repo, deploy_keys_url=deploy_keys_url, keypath=args.key_path)))
            print(header)
        else:
            print(header)
            print(dedent("""\
            {N}. Go to {deploy_keys_url}
            and add the following as a new key:

            {ssh_key}
            Be sure to allow write access for the key.
            """.format(ssh_key=ssh_key, deploy_keys_url=deploy_keys_url, N=N)))

        # TODO: Should we just delete the public key?

        print(dedent("""\
        {N}. Commit the file {keypath}.enc.
        The file {keypath}.pub contains the public deploy key for GitHub. It
        does not need to be committed.
        """.format(keypath=args.key_path, N=N)))

    options = ''
    if args.key_path != 'github_deploy_key':
        options += ' --key-path {keypath}.enc'.format(keypath=args.key_path)
    if deploy_repo != build_repo:
        options += ' --deploy-repo {deploy_repo}'.format(deploy_repo=deploy_repo)

    print(dedent("""\
    {N}. Add

        - pip install doctr
        - doctr deploy{options}

    to the docs build of your .travis.yml.
    """.format(options=options, N=N)))

    print(dedent("""\
    {N}. Put

        env:
          global:
            - secure: "{encrypted_variable}"

    in your .travis.yml.
    """.format(encrypted_variable=encrypted_variable.decode('utf-8'), N=N)))

def main():
    return process_args(get_parser())

if __name__ == '__main__':
    sys.exit(main())
<commit_msg>Add the warnings from the README to the output of doctr configure<commit_after>"""
doctr

A tool to automatically deploy docs to GitHub pages from Travis CI.

The doctr command is two commands in one. To use, first run

doctr configure

on your local machine. This will prompt for your GitHub credentials and the
name of the repo you want to deploy docs for. This will generate a secure key,
which you should insert into your .travis.yml.

Then, on Travis, for the build where you build your docs, add

    - doctr deploy

to the end of the build to deploy the docs to GitHub pages.  This will only
run on the master branch, and won't run on pull requests.

For more information, see https://drdoctr.github.io/doctr/docs/
"""

import sys
import os
import argparse
import shlex
import subprocess

from textwrap import dedent

from .local import (generate_GitHub_token, encrypt_variable, encrypt_file,
    upload_GitHub_deploy_key, generate_ssh_key, check_repo_exists, GitHub_login)
from .travis import (setup_GitHub_push, commit_docs, push_docs,
    get_current_repo, sync_from_log, find_sphinx_build_dir, run)
from . import __version__

def get_parser():
    # This uses RawTextHelpFormatter so that the description (the docstring of
    # this module) is formatted correctly. Unfortunately, that means that
    # parser help is not text wrapped (but all other help is).
    parser = argparse.ArgumentParser(description=__doc__,
        formatter_class=argparse.RawTextHelpFormatter, epilog="""
Run --help on the subcommands like 'doctr deploy --help' to see the
options available.
        """,
        )
    parser.add_argument('-V', '--version', action='version', version='doctr ' + __version__)

    subcommand = parser.add_subparsers(title='subcommand', dest='subcommand')
    deploy_parser = subcommand.add_parser('deploy', help="""Deploy the docs to GitHub from Travis.""")
    deploy_parser.set_defaults(func=deploy)
    deploy_parser.add_argument('--force', action='store_true', help="""Run the deploy command even
    if we do not appear to be on Travis.""")
    deploy_parser.add_argument('--token', action='store_true', default=False,
        help="""Push to GitHub using a personal access token. Use this if you
        used 'doctr configure --token'.""")
    deploy_parser.add_argument('--key-path', default='github_deploy_key.enc',
        help="""Path of the encrypted GitHub deploy key. The default is %(default)r.""")
    deploy_parser.add_argument('--built-docs', default=None,
        help="""Location of the built html documentation to be deployed to
        gh-pages. If not specified, Doctr will try to automatically detect build location""")
    deploy_parser.add_argument('--gh-pages-docs', default='docs',
        help="""Directory to deploy the html documentation to on gh-pages. The
        default is %(default)r.""")
    deploy_parser.add_argument('--tmp-dir', default=None,
        help=argparse.SUPPRESS)
    deploy_parser.add_argument('--deploy-repo', default=None, help="""Repo to
        deploy the docs to. By default, it deploys to the repo Doctr is run from.""")
    deploy_parser.add_argument('--no-require-master', dest='require_master', action='store_false',
        default=True, help="""Allow docs to be pushed from a branch other than master""")
    deploy_parser.add_argument('--command', default=None, help="""Command to
        be run before committing and pushing. If the command creates
        additional files that should be deployed, they should be added to the
        index.""")
    deploy_parser.add_argument('--no-sync', dest='sync', action='store_false',
        default=True, help="""Don't sync any files. This is generally used in
        conjunction with the --command flag, for instance, if the command syncs
        the files for you. Any files you wish to commit should be added to the
        index.""")
    deploy_parser.add_argument('--no-push', dest='push', action='store_false',
        default=True, help="Run all the steps except the last push step."
        "Useful for debugging")

    configure_parser = subcommand.add_parser('configure', help="Configure doctr. This command should be run locally (not on Travis).")
    configure_parser.set_defaults(func=configure)
    configure_parser.add_argument('--force', action='store_true', help="""Run the configure command even
    if we appear to be on Travis.""")
    configure_parser.add_argument('--token', action="store_true", default=False,
        help="""Generate a personal access token to push to GitHub. The default is to use a
        deploy key. WARNING: This will grant read/write access to all the
        public repositories for the user. This option is not recommended
        unless you are using a separate GitHub user for deploying.""")
    configure_parser.add_argument("--no-upload-key", action="store_false", default=True,
        dest="upload_key", help="""Don't automatically upload the deploy key to GitHub. If you select this
        option, you will not be prompted for your GitHub credentials, so this option is not compatible with
        private repositories.""")
    configure_parser.add_argument('--key-path', default='github_deploy_key',
        help="""Path to save the encrypted GitHub deploy key. The default is %(default)r.
    The .enc extension is added to the file automatically.""")

    return parser

def process_args(parser):
    args = parser.parse_args()

    if not args.subcommand:
        parser.print_usage()
        parser.exit(1)

    try:
        return args.func(args, parser)
    except RuntimeError as e:
        sys.exit("Error: " + e.args[0])

def on_travis():
    return os.environ.get("TRAVIS_JOB_NUMBER", '')

def deploy(args, parser):
    if not args.force and not on_travis():
        parser.error("doctr does not appear to be running on Travis. Use "
            "doctr deploy --force to run anyway.")

    if args.tmp_dir:
        parser.error("The --tmp-dir flag has been removed (doctr no longer uses a temporary directory when deploying).")

    build_repo = get_current_repo()
    deploy_repo = args.deploy_repo or build_repo

    current_commit = subprocess.check_output(['git', 'rev-parse', 'HEAD']).decode('utf-8').strip()
    try:
        can_push = setup_GitHub_push(deploy_repo, auth_type='token' if args.token else
                             'deploy_key', full_key_path=args.key_path,
                             require_master=args.require_master)

        if args.sync:
            built_docs = args.built_docs or find_sphinx_build_dir()

            log_file = os.path.join(args.gh_pages_docs, '.doctr-files')

            print("Moving built docs into place")
            added, removed = sync_from_log(src=built_docs,
                dst=args.gh_pages_docs, log_file=log_file)

        else:
            added, removed = [], []

        if args.command:
            run(shlex.split(args.command))

        changes = commit_docs(added=added, removed=removed)
        if changes:
            if can_push and args.push:
                push_docs()
            else:
                print("Don't have permission to push. Not trying.")
        else:
            print("The docs have not changed. Not updating")
    finally:
        subprocess.run(['git', 'checkout', current_commit])

class IncrementingInt:
    def __init__(self, i=0):
        self.i = i

    def __repr__(self):
        ret = repr(self.i)
        self.i += 1
        return ret

    __str__ = __repr__

def configure(args, parser):
    if not args.force and on_travis():
        parser.error("doctr appears to be running on Travis. Use "
            "doctr configure --force to run anyway.")

    if args.upload_key:
        login_kwargs = GitHub_login()
    else:
        login_kwargs = {'auth': None, 'headers': None}

    build_repo = input("What repo do you want to build the docs for (org/reponame, like 'drdoctr/doctr')? ")
    is_private = check_repo_exists(build_repo, service='github', **login_kwargs)
    check_repo_exists(build_repo, service='travis')

    deploy_repo = input("What repo do you want to deploy the docs to? [{build_repo}] ".format(build_repo=build_repo))
    if not deploy_repo:
        deploy_repo = build_repo

    if deploy_repo != build_repo:
        check_repo_exists(deploy_repo, service='github', **login_kwargs)

    N = IncrementingInt(1)

    header = "\n================== You should now do the following ==================\n"

    if args.token:
        token = generate_GitHub_token(**login_kwargs)['token']
        encrypted_variable = encrypt_variable("GH_TOKEN={token}".format(token=token).encode('utf-8'),
            build_repo=build_repo, is_private=is_private, **login_kwargs)
        print(dedent("""
        A personal access token for doctr has been created.

        You can go to https://github.com/settings/tokens to revoke it."""))

        print(header)
    else:
        ssh_key = generate_ssh_key("doctr deploy key for {deploy_repo}".format(deploy_repo=deploy_repo), keypath=args.key_path)
        key = encrypt_file(args.key_path, delete=True)
        encrypted_variable = encrypt_variable(b"DOCTR_DEPLOY_ENCRYPTION_KEY=" + key, build_repo=build_repo, is_private=is_private, **login_kwargs)

        deploy_keys_url = 'https://github.com/{deploy_repo}/settings/keys'.format(deploy_repo=deploy_repo)

        if args.upload_key:

            upload_GitHub_deploy_key(deploy_repo, ssh_key, **login_kwargs)

            print(dedent("""
            The deploy key has been added for {deploy_repo}.

            You can go to {deploy_keys_url} to revoke the deploy key.\
            """.format(deploy_repo=deploy_repo, deploy_keys_url=deploy_keys_url, keypath=args.key_path)))
            print(header)
        else:
            print(header)
            print(dedent("""\
            {N}. Go to {deploy_keys_url}
            and add the following as a new key:

            {ssh_key}
            Be sure to allow write access for the key.
            """.format(ssh_key=ssh_key, deploy_keys_url=deploy_keys_url, N=N)))

        # TODO: Should we just delete the public key?

        print(dedent("""\
        {N}. Commit the file {keypath}.enc.
        The file {keypath}.pub contains the public deploy key for GitHub. It
        does not need to be committed.
        """.format(keypath=args.key_path, N=N)))

    options = ''
    if args.key_path != 'github_deploy_key':
        options += ' --key-path {keypath}.enc'.format(keypath=args.key_path)
    if deploy_repo != build_repo:
        options += ' --deploy-repo {deploy_repo}'.format(deploy_repo=deploy_repo)

    print(dedent("""\
    {N}. Add

        script:
          - set -e
          - # Command to build your docs
          - pip install doctr
          - doctr deploy{options}

    to the docs build of your .travis.yml.  The 'set -e' prevents doctr from
    running when the docs build fails. Use the 'script' section so that if
    doctr fails it causes the build to fail.
    """.format(options=options, N=N)))

    print(dedent("""\
    {N}. Put

        env:
          global:
            - secure: "{encrypted_variable}"

    in your .travis.yml.
    """.format(encrypted_variable=encrypted_variable.decode('utf-8'), N=N)))

def main():
    return process_args(get_parser())

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


<commit_before># ----------------------------------------------------------------------
# Numenta Platform for Intelligent Computing (NuPIC)
# Copyright (C) 2013, Numenta, Inc.  Unless you have purchased from
# Numenta, Inc. a separate commercial license for this software code, the
# following terms and conditions apply:
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License version 3 as
# published by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
# See the GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see http://www.gnu.org/licenses.
#
# http://numenta.org/licenses/
# ----------------------------------------------------------------------

"""Implements the flat spatial pooler."""

import copy
import cPickle
import itertools
import numpy

from nupic.bindings.math import GetNTAReal
from nupic.research.spatial_pooler import SpatialPooler

realDType = GetNTAReal()

class FlatSpatialPooler(SpatialPooler):
	"""
	This class implements the flat spatial pooler. This version of the spatial
	pooler contains no toplogy information. It uses global coverage and global
	inhibition
	"""

	def __init__(self,
               inputShape=(32, 32),
               inputBorder=8,
               inputDensity=1.0,
               coincidencesShape=(48, 48),
               coincInputRadius=16,
               coincInputPoolPct=1.0,
               gaussianDist=False,
               commonDistributions=False,
               localAreaDensity=-1.0,
               numActivePerInhArea=10.0,
               stimulusThreshold=0,
               synPermInactiveDec=0.01,
               synPermActiveInc=0.1,
               synPermActiveSharedDec=0.0,
               synPermOrphanDec=0.0,
               synPermConnected=0.10,
               minPctDutyCycleBeforeInh=0.001,
               minPctDutyCycleAfterInh=0.001,
               dutyCyclePeriod=1000,
               maxFiringBoost=10.0,
               maxSSFiringBoost=2.0,
               maxSynPermBoost=10.0,
               minDistance=0.0,
               cloneMap=None,
               numCloneMasters=-1,
               seed=-1,
               spVerbosity=0,
               printPeriodicStats=0,
               testMode=False,
               globalInhibition=False,
               spReconstructionParam="unweighted_mean",
               useHighTier=True,
               randomSP=False,
              ):

		super(FlatSpatialPooler,self).__init__(
				inputDimensions=numpy.array(inputShape),
				columnDimensions=numpy.array(coincidencesShape),
				potentialRadius=coincInputRadius,
				potentialPct=coincInputPoolPct,
				globalInhibition=globalInhibition,
				localAreaDensity=localAreaDensity,
				numActiveColumnsPerInhArea=numActivePerInhArea,
				stimulusThreshold=stimulusThreshold,
        synPermInactiveDec=synPermInactiveDec,
        synPermActiveInc=synPermActiveInc,
        synPermConnected=synPermConnected,
        minPctOverlapDutyCycle=minPctDutyCycleBeforeInh,
        minPctActiveDutyCycle=minPctDutyCycleAfterInh,
        dutyCyclePeriod=dutyCyclePeriod,
        maxBoost=maxFiringBoost,
				seed=seed,
				spVerbosity=spVerbosity,
			)

		# save arguments
		self._numInputs = numpy.prod(numpy.array(inputShape))
		self._numColumns = numpy.prod(numpy.array(coincidencesShape))
		self._minDistance = minDistance
		self._randomSP = randomSP

		#set active duty cycles to ones, because they set anomaly scores to 0
		self._activeDutyCycles = numpy.ones(self._numColumns)

		# set of columns to be 'hungry' for learning
		self._boostFactors *= maxFiringBoost

	# This constructor is a minimal, stripped down version of the
	# constructure above. The constructor above is only used to
	# provid backwards compatibility to the old spatial pooler.
	# def __init__(self,
	# 						 numInputs,
	# 						 numColumns,
	# 						 localAreaDensity=0.1,
	# 						 numActiveColumnsPerInhArea=-1,
	# 						 stimulusThreshold=0,
	# 						 minDistance=0.0,
	# 						 maxBoost=10.0,
	# 						 seed=-1,
	# 						 spVerbosity=0,
	# 						 randomSP=False,
	# 						 ):

	# 	super(FlatSpatialPooler,self).__init__(
	# 			inputDimensions=numInputs,
	# 			columnDimensions=numColumns,
	# 			potentialRadius=numInputs,
	# 			potentialPct=0.5,
	# 			globalInhibition=True,
	# 			localAreaDensity=localAreaDensity,
	# 			numActiveColumnsPerInhArea=numActiveColumnsPerInhArea,
	# 			stimulusThreshold=stimulusThreshold,
	# 			seed=seed
	# 		)

	# 	#verify input is valid
	# 	assert(numColumns > 0)
	# 	assert(numInputs > 0)

	# 	# save arguments
	# 	self._numInputs = numInputs
	# 	self._numColumns = numColumns
	# 	self._minDistance = minDistance
	# 	self._randomSP = randomSP

	# 	#set active duty cycles to ones, because they set anomaly scores to 0
	# 	self._activeDutyCycles = numpy.ones(self._numColumns)

	# 	# set of columns to be 'hungry' for learning
	# 	self._boostFactors *= maxBoost

	def compute(self, inputVector, learn=True):
		if self._randomSP:
			learn=False

		assert (numpy.size(inputVector) == self._numInputs)
		self._updateBookeepingVars(learn)
		inputVector = numpy.array(inputVector, dtype=realDType)
		overlaps = self._calculateOverlap(inputVector)
		overlapsPct = self._calculateOverlapPct(overlaps)
		highTierColumns = self._selectHighTierColumns(overlapsPct)
		virginColumns = self._selectVirginColumns()

		if learn:
			vipOverlaps = self._boostFactors * overlaps
		else:
			vipOverlaps = overlaps.copy()

		vipBonus = max(vipOverlaps) + 1.0
		if learn:
			vipOverlaps[virginColumns] = vipBonus
		vipOverlaps[highTierColumns] += vipBonus
		activeColumns = self._inhibitColumns(vipOverlaps)

		if learn:
			self._adaptSynapses(inputVector, activeColumns)
			self._updateDutyCycles(overlaps, activeColumns)
			self._bumpUpWeakColumns()
			self._updateBoostFactors()

			if self._isUpdateRound():
				self._updateInhibitionRadius()
				self._updateMinDutyCycles()
		else:
      activeColumns = self._stripNeverLearned(activeColumns)

		activeArray = numpy.zeros(self._numColumns)
		activeArray[activeColumns] = 1
		return activeArray

	def _selectVirginColumns(self):
		return numpy.where(self._activeDutyCycles == 0)[0]

	def _selectHighTierColumns(self, overlapsPct):
		return numpy.where(overlapsPct >= (1.0 - self._minDistance))[0]

  <commit_msg>fixed indentation issue<commit_after># ----------------------------------------------------------------------
# Numenta Platform for Intelligent Computing (NuPIC)
# Copyright (C) 2013, Numenta, Inc.  Unless you have purchased from
# Numenta, Inc. a separate commercial license for this software code, the
# following terms and conditions apply:
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License version 3 as
# published by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
# See the GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see http://www.gnu.org/licenses.
#
# http://numenta.org/licenses/
# ----------------------------------------------------------------------

"""Implements the flat spatial pooler."""

import copy
import cPickle
import itertools
import numpy

from nupic.bindings.math import GetNTAReal
from nupic.research.spatial_pooler import SpatialPooler

realDType = GetNTAReal()

class FlatSpatialPooler(SpatialPooler):
	"""
	This class implements the flat spatial pooler. This version of the spatial
	pooler contains no toplogy information. It uses global coverage and global
	inhibition
	"""

	def __init__(self,
               inputShape=(32, 32),
               inputBorder=8,
               inputDensity=1.0,
               coincidencesShape=(48, 48),
               coincInputRadius=16,
               coincInputPoolPct=1.0,
               gaussianDist=False,
               commonDistributions=False,
               localAreaDensity=-1.0,
               numActivePerInhArea=10.0,
               stimulusThreshold=0,
               synPermInactiveDec=0.01,
               synPermActiveInc=0.1,
               synPermActiveSharedDec=0.0,
               synPermOrphanDec=0.0,
               synPermConnected=0.10,
               minPctDutyCycleBeforeInh=0.001,
               minPctDutyCycleAfterInh=0.001,
               dutyCyclePeriod=1000,
               maxFiringBoost=10.0,
               maxSSFiringBoost=2.0,
               maxSynPermBoost=10.0,
               minDistance=0.0,
               cloneMap=None,
               numCloneMasters=-1,
               seed=-1,
               spVerbosity=0,
               printPeriodicStats=0,
               testMode=False,
               globalInhibition=False,
               spReconstructionParam="unweighted_mean",
               useHighTier=True,
               randomSP=False,
              ):

		super(FlatSpatialPooler,self).__init__(
				inputDimensions=numpy.array(inputShape),
				columnDimensions=numpy.array(coincidencesShape),
				potentialRadius=coincInputRadius,
				potentialPct=coincInputPoolPct,
				globalInhibition=globalInhibition,
				localAreaDensity=localAreaDensity,
				numActiveColumnsPerInhArea=numActivePerInhArea,
				stimulusThreshold=stimulusThreshold,
        synPermInactiveDec=synPermInactiveDec,
        synPermActiveInc=synPermActiveInc,
        synPermConnected=synPermConnected,
        minPctOverlapDutyCycle=minPctDutyCycleBeforeInh,
        minPctActiveDutyCycle=minPctDutyCycleAfterInh,
        dutyCyclePeriod=dutyCyclePeriod,
        maxBoost=maxFiringBoost,
				seed=seed,
				spVerbosity=spVerbosity,
			)

		# save arguments
		self._numInputs = numpy.prod(numpy.array(inputShape))
		self._numColumns = numpy.prod(numpy.array(coincidencesShape))
		self._minDistance = minDistance
		self._randomSP = randomSP

		#set active duty cycles to ones, because they set anomaly scores to 0
		self._activeDutyCycles = numpy.ones(self._numColumns)

		# set of columns to be 'hungry' for learning
		self._boostFactors *= maxFiringBoost

	# This constructor is a minimal, stripped down version of the
	# constructure above. The constructor above is only used to
	# provid backwards compatibility to the old spatial pooler.
	# def __init__(self,
	# 						 numInputs,
	# 						 numColumns,
	# 						 localAreaDensity=0.1,
	# 						 numActiveColumnsPerInhArea=-1,
	# 						 stimulusThreshold=0,
	# 						 minDistance=0.0,
	# 						 maxBoost=10.0,
	# 						 seed=-1,
	# 						 spVerbosity=0,
	# 						 randomSP=False,
	# 						 ):

	# 	super(FlatSpatialPooler,self).__init__(
	# 			inputDimensions=numInputs,
	# 			columnDimensions=numColumns,
	# 			potentialRadius=numInputs,
	# 			potentialPct=0.5,
	# 			globalInhibition=True,
	# 			localAreaDensity=localAreaDensity,
	# 			numActiveColumnsPerInhArea=numActiveColumnsPerInhArea,
	# 			stimulusThreshold=stimulusThreshold,
	# 			seed=seed
	# 		)

	# 	#verify input is valid
	# 	assert(numColumns > 0)
	# 	assert(numInputs > 0)

	# 	# save arguments
	# 	self._numInputs = numInputs
	# 	self._numColumns = numColumns
	# 	self._minDistance = minDistance
	# 	self._randomSP = randomSP

	# 	#set active duty cycles to ones, because they set anomaly scores to 0
	# 	self._activeDutyCycles = numpy.ones(self._numColumns)

	# 	# set of columns to be 'hungry' for learning
	# 	self._boostFactors *= maxBoost

	def compute(self, inputVector, learn=True):
		if self._randomSP:
			learn=False

		assert (numpy.size(inputVector) == self._numInputs)
		self._updateBookeepingVars(learn)
		inputVector = numpy.array(inputVector, dtype=realDType)
		overlaps = self._calculateOverlap(inputVector)
		overlapsPct = self._calculateOverlapPct(overlaps)
		highTierColumns = self._selectHighTierColumns(overlapsPct)
		virginColumns = self._selectVirginColumns()

		if learn:
			vipOverlaps = self._boostFactors * overlaps
		else:
			vipOverlaps = overlaps.copy()

		vipBonus = max(vipOverlaps) + 1.0
		if learn:
			vipOverlaps[virginColumns] = vipBonus
		vipOverlaps[highTierColumns] += vipBonus
		activeColumns = self._inhibitColumns(vipOverlaps)

		if learn:
			self._adaptSynapses(inputVector, activeColumns)
			self._updateDutyCycles(overlaps, activeColumns)
			self._bumpUpWeakColumns()
			self._updateBoostFactors()

			if self._isUpdateRound():
				self._updateInhibitionRadius()
				self._updateMinDutyCycles()
		else:
			activeColumns = self._stripNeverLearned(activeColumns)

		activeArray = numpy.zeros(self._numColumns)
		activeArray[activeColumns] = 1
		return activeArray

	def _selectVirginColumns(self):
		return numpy.where(self._activeDutyCycles == 0)[0]

	def _selectHighTierColumns(self, overlapsPct):
		return numpy.where(overlapsPct >= (1.0 - self._minDistance))[0]

  <|endoftext|>


<commit_before>from __future__ import print_function
from pywinauto.application import Application
import time
import csv
import os
import sys

def generateDataFile(t_interval, interface_name, file_name):
    # start Wireshark
    if (os.path.exists(r"C:\Program Files (x86)\Wireshark")):
        app = Application(backend='uia').start(r"C:\Program Files (x86)\Wireshark\Wireshark.exe")
    else:
        if (os.path.exists(r"C:\Program Files\Wireshark")):
            app = Application(backend='uia').start(r"C:\Program Files\Wireshark\Wireshark.exe")
        else:
            print("Can't find wireshark on your computer")

    win = app['The Wireshark Network Analyzer']

    if app.software_update.exists(timeout=10):
        app.software_update.skip_this_version.click()
        app.software_update.wait_not('visible')  # just to make sure it's closed
    win.wait('ready', timeout=15)

    # Try to find interface_name in TreeView interfaces list
    tree = win["Interface list Interface list"]
    interface_name = "\\" + interface_name
    try:
        tree.get_item(interface_name).double_click_input()
    except Exception:
        print('No such interface')
        exit()

    win = app['Dialog']

    # Wait while WireShark collect information
    win.wait('ready')
    time.sleep(t_interval)
    win.wait('ready')

    # Stop WireShark
    win = app['Dialog']
    win['Stop'].click()

    # open menu (File -> Export Packet Dissections -> As CSV)
    win['File Alt+F'].select()
    win = app.top_window()
    win['Export Packet Dissections'].select()
    win = app.top_window()
    win['As CSV...'].click_input()

    # Export FileDialog
    win = app['Dialog']
    child = win['Export File Dialog']

    check_file_exist = True
    try:
        open(file_name)
    except IOError:
        check_file_exist = False

    # input path to temporary file
    win.type_keys(file_name)
    child.window(best_match = "Save").click()
    # child['Save'].click()

    # if window "confirm Save As" pop up
    if (check_file_exist == True):
        child = win['Confirm Save As']
        child.window(best_match='yes').click()
        # child['yes'].click()

    # Quit
    win['File Alt+F'].select()
    win = app.top_window()
    win['Quit Ctrl+Q'].click_input()

    # if window "Quit without Saving" pop up
    try:
        win = app.top_window()
        win['Quit without Saving Alt+w'].click()
    except:
        pass

def parseFile(file_name):
    # parse csv file
    prot_dict = {}
    with open(file_name) as csvfile:
        text = csv.reader(csvfile)
        # skip first line with header
        next(text, None)
        # initialize dict with list
        for row in text:
            prot_dict[row[4]] = list()
        # return back to beginning
        csvfile.seek(0)
        next(text, None)
        # collect information about protocols
        for row in text:
            prot_dict[row[4]].append(int(row[5]))
    # calculate count packets, mean packer length and traffic size
    for key in prot_dict:
        temp_list = prot_dict[key]
        count_pack = len(temp_list)
        mean_pack_len = int(sum(temp_list)/len(temp_list))
        traffic_size = sum(temp_list)
        prot_dict[key] = {}
        prot_dict[key].append(count_pack)
        prot_dict[key].append(mean_pack_len)
        prot_dict[key].append(traffic_size)
    os.remove(file_name)
    return prot_dict

def printResult(result):
    print_order = list()
    keys = list(result.keys())

    for _ in range(0, len(keys)):
        max_ = 0
        temp_key = ''
        # find max key
        for key in keys:
            if result[key][2] > max_:
                max_ = result[key][2]
                temp_key = key
        print_order.append(temp_key)
        keys.remove(temp_key)

    print("Protocol  count protocols  mean length  traffic size")
    for key in print_order:
        string = key + "  " + str(result[key][0]) + "  "\
                 + str(result[key][1]) + "  " + str(result[key][2])
        print(string)

if (len(sys.argv) < 3):
    print("Expected time and interface_name")
    exit()
else:
    t_interval = int(sys.argv[1])
    interface_name = sys.argv[2]

file_name = os.path.dirname(os.path.abspath(__file__))
file_name = os.path.join(file_name, r'test.csv')

generateDataFile(t_interval,interface_name,file_name)
result = parseFile(file_name)
printResult(result)
<commit_msg>added doc string and fixed error<commit_after>"""
Example script captures network traffic using WireShark and prints a short summary for every protocol.
Requirements:
  - Wireshark 2.2.5
  - pywinauto 0.6.1+
This example opens "Wireshark", navigates to 'network_connection_name',
captures network traffis for 'capture_time', saves all data to temporary file,
parses it and shows short summer for every protocol.
"""

from __future__ import print_function
from pywinauto.application import Application
import time
import csv
import os
import sys

def generateDataFile(t_interval, interface_name, file_name):
    # start Wireshark
    if (os.path.exists(r"C:\Program Files (x86)\Wireshark")):
        app = Application(backend='uia').start(r"C:\Program Files (x86)\Wireshark\Wireshark.exe")
    else:
        if (os.path.exists(r"C:\Program Files\Wireshark")):
            app = Application(backend='uia').start(r"C:\Program Files\Wireshark\Wireshark.exe")
        else:
            print("Can't find wireshark on your computer")

    win = app['The Wireshark Network Analyzer']

    if app.software_update.exists(timeout=10):
        app.software_update.skip_this_version.click()
        app.software_update.wait_not('visible')  # just to make sure it's closed
    win.wait('ready', timeout=15)

    # Try to find interface_name in TreeView interfaces list
    tree = win["Interface list Interface list"]
    interface_name = "\\" + interface_name
    try:
        tree.get_item(interface_name).double_click_input()
    except Exception:
        print('No such interface')
        exit()

    win = app['Dialog']

    # Wait while WireShark collect information
    win.wait('ready')
    time.sleep(t_interval)
    win.wait('ready')

    # Stop WireShark
    win = app['Dialog']
    win['Stop'].click()

    # open menu (File -> Export Packet Dissections -> As CSV)
    win['File Alt+F'].select()
    win = app.top_window()
    win['Export Packet Dissections'].select()
    win = app.top_window()
    win['As CSV...'].click_input()

    # Export FileDialog
    win = app['Dialog']
    child = win['Export File Dialog']

    check_file_exist = True
    try:
        open(file_name)
    except IOError:
        check_file_exist = False

    # input path to temporary file
    win.type_keys(file_name)
    child.window(best_match="Save").click()
    # child['Save'].click()

    # if window "confirm Save As" pop up
    if (check_file_exist is True):
        child = win['Confirm Save As']
        child.window(best_match='yes').click()
        # child['yes'].click()

    # Quit
    win['File Alt+F'].select()
    win = app.top_window()
    win['Quit Ctrl+Q'].click_input()

    # if window "Quit without Saving" pop up
    try:
        win = app.top_window()
        win['Quit without Saving Alt+w'].click()
    except:
        pass

def parseFile(file_name):
    # parse csv file
    prot_dict = {}
    with open(file_name) as csvfile:
        text = csv.reader(csvfile)
        # skip first line with header
        next(text, None)
        # initialize dict with list
        for row in text:
            prot_dict[row[4]] = list()
        # return back to beginning
        csvfile.seek(0)
        next(text, None)
        # collect information about protocols
        for row in text:
            prot_dict[row[4]].append(int(row[5]))
    # calculate count packets, mean packer length and traffic size
    for key in prot_dict:
        temp_list = prot_dict[key]
        count_pack = len(temp_list)
        mean_pack_len = int(sum(temp_list) / len(temp_list))
        traffic_size = sum(temp_list)
        del prot_dict[key][:]
        prot_dict[key].append(count_pack)
        prot_dict[key].append(mean_pack_len)
        prot_dict[key].append(traffic_size)
    os.remove(file_name)
    return prot_dict

def printResult(result):
    print_order = list()
    keys = list(result.keys())

    for _ in range(0, len(keys)):
        max_ = 0
        temp_key = ''
        # find max key
        for key in keys:
            if result[key][2] > max_:
                max_ = result[key][2]
                temp_key = key
        print_order.append(temp_key)
        keys.remove(temp_key)

    print("Protocol  count protocols  mean length  traffic size")
    for key in print_order:
        string = key + "  " + str(result[key][0]) + "  "\
                 + str(result[key][1]) + "  " + str(result[key][2])
        print(string)

if (len(sys.argv) < 3):
    print('''This pywinauto example captures network traffic using WireShark and prints a short summary for every protocol.
Usage: python wireshark.py <capture_time (seconds)> <network_connection_name>
Example: python wireshark.py 5 Ethernet
It will capture all the packets from "Ethernet" interface during 5 seconds.''')
    exit()
else:
    t_interval = int(sys.argv[1])
    interface_name = sys.argv[2]

file_name = os.path.dirname(os.path.abspath(__file__))
file_name = os.path.join(file_name, r'test.csv')

generateDataFile(t_interval,interface_name,file_name)
result = parseFile(file_name)
printResult(result)
<|endoftext|>


<commit_before>""" Dotconf lexer and parser.
"""

import sys
from glob import glob

import ply.lex as lex
import ply.yacc as yacc

from dotconf.tree import ConfigSection, ConfigValue, Position

UNITS = {'k': 10 ** 3,
         'M': 10 ** 6,
         'G': 10 ** 9,
         'T': 10 ** 12,
         'P': 10 ** 15,
         'E': 10 ** 18,
         'Z': 10 ** 21,
         'Y': 10 ** 24,
         'Ki': 2 ** 10,
         'Mi': 2 ** 20,
         'Gi': 2 ** 30,
         'Ti': 2 ** 40,
         'Pi': 2 ** 50,
         'Ei': 2 ** 60,
         'Zi': 2 ** 70,
         'Yi': 2 ** 80}

class ParsingError(Exception):

    """ Error raised when a parsing error occurs.
    """

    def __init__(self, msg, position=None):
        super(ParsingError, self).__init__(msg)
        self.position = position

def default_external_opener(locator):
    """ The default locator used to open included external files.
    """
    parsed_externals = []
    for external in glob(locator):
        try:
            with open(external) as fexternal:
                external_data = fexternal.read()
        except IOError as err:
            raise ParsingError('Unable to open %s (%s)' % (external, err))
        parser = DotconfParser(external_data, debug=False, write_tables=False,
                               errorlog=yacc.NullLogger(), input_name=external,
                               external_opener=default_external_opener)
        parsed_externals.append(parser.parse())
    return parsed_externals

#
# Lexer
#

class DotconfLexer(object):

    """ Lexer for the DotConf format.

    :param \*\*kwargs: arguments to give to the ply lexer

    Usage example::

    >>> lexer = DotConfLexer()
    >>> lexer.input('test { key = yes }')
    >>> print lexer.next()
    """

    def __init__(self, encoding='utf-8', **kwargs):
        self._lexer = lex.lex(module=self, **kwargs)
        self._encoding = encoding

    #
    # Tokens definition
    #

    reserved = {'yes': 'YES', 'no': 'NO', 'include': 'INCLUDE'}
    reserved.update(dict(((kw, 'UNIT') for kw in UNITS)))
    tokens = ['LBRACE', 'RBRACE', 'NAME', 'TEXT', 'NUMBER',
              'ASSIGN', 'LIST_SEP'] + list(set(reserved.values()))

    t_LBRACE = '{'
    t_RBRACE = '}'
    t_ASSIGN = '='
    t_LIST_SEP = ','
    t_ignore = ' \t'

    def t_NAME(self, token):
        r'[a-zA-Z_][a-zA-Z0-9_-]*'
        token.type = self.reserved.get(token.value, 'NAME')
        # Handle the boolean case:
        if token.type == 'YES':
            token.value = True
        elif token.type == 'NO':
            token.value = False
        elif token.type == 'UNIT':
            token.value = UNITS[token.value]
        return token

    def t_TEXT(self, token):
        r'(["]([\\]["]|[^"]|)*["]|[\']([\\][\']|[^\'])*[\'])'
        value = token.value[1:-1].replace('\\' + token.value[0], token.value[0])
        token.value = value
        # Count the lines in the string:
        token.lexer.lineno += value.count('\n')
        return token

    def t_NUMBER(self, token):
        r'[-+]?[0-9]+(\.[0-9]+)?'
        if token.value.isdigit():
            token.value = int(token.value)
        else:
            token.value = float(token.value)
        return token

    def t_EOL(self, token):
        r'[\n]+'
        token.lexer.lineno += len(token.value)

    def t_COMMENT(self, token):
        r'[#].*'

    def t_error(self, token):
        position = Position(self._input_name,
                            self._lexer.lineno,
                            self._lexer.lexpos)
        raise ParsingError('Illegal character %r' % token.value[0], position)

    #
    # Public methods
    #

    def column(self, lexpos):
        """ Find the column according to the lexpos.
        """
        # This code is taken from the python-ply documentation
        # see: http://www.dabeaz.com/ply/ply.html section 4.6
        last_cr = self._current_input.rfind('\n', 0, lexpos)
        if last_cr < 0:
            last_cr = 0
        column = (lexpos - last_cr)
        return column

    #
    # Bindings to the internal _lexer object
    #

    def input(self, input):
        if sys.version_info[0] >= 3:
            if isinstance(input, bytes):
                input = input.decode(self._encoding)
        else:
            if isinstance(input, str):
                input = input.decode(self._encoding)
        self._current_input = input
        return self._lexer.input(input)

    def __getattr__(self, name):
        attr = getattr(self._lexer, name)
        if attr is None:
            raise AttributeError("'%s' object has no attribute '%s'" % (self, name))
        else:
            return attr

#
# Parser
#

class DotconfParser(object):

    """ Parser for the Dotconf format.
    """

    tokens = DotconfLexer.tokens

    def __init__(self, input, **kwargs):
        self._input = input
        self._input_name = kwargs.pop('input_name', '<unknown>')
        self._external_opener = kwargs.pop('external_opener',
                                           default_external_opener)
        self._lexer = kwargs.pop('lexer', DotconfLexer())
        self._parser = yacc.yacc(module=self, **kwargs)
        self._old_line = 0

    def _check_line(self, current, lineno, pos, token):
        if self._old_line == current:
            pos = Position(self._input_name, lineno, pos)
            raise ParsingError('Syntax error near of "%s", '
                               'newline missing?' % token, pos)
        else:
            self._old_line = current

    #
    # Rules
    #

    start = 'top'

    def p_top(self, p):
        """top : section_content"""
        section = ConfigSection('__top__')
        for child in p[1]:
            if isinstance(child, ConfigSection):
                child.parent = section
            section.register(child)
        p[0] = section

    def p_assignation(self, p):
        """assignment : NAME ASSIGN value
                      | NAME ASSIGN list"""
        position = Position(self._input_name,
                            p.lineno(3),
                            p.lexer.column(p.lexpos(3)))
        value = ConfigValue(p[1], p[3], position=position)
        p[0] = value

    def p_value(self, p):
        """value : TEXT
                 | YES
                 | NO
                 | number"""
        p[0] = p[1]

    def p_number(self, p):
        """number : NUMBER
                  | NUMBER UNIT"""
        if len(p) == 2:
            p[0] = p[1]
        else:
            p[0] = p[1] * p[2]

    #
    # List management rules:
    #

    def p_list(self, p):
        """list : value LIST_SEP list_next
                | value LIST_SEP"""
        if len(p) == 3:
            p[0] = [p[1]]
        elif len(p) == 4:
            p[3].insert(0, p[1])
            p[0] = p[3]

    def p_list_next(self, p):
        """list_next : value
                     | value LIST_SEP
                     | value LIST_SEP list_next"""
        if len(p) in (2, 3):
            p[0] = [p[1]]
        elif len(p) == 4:
            p[3].insert(0, p[1])
            p[0] = p[3]

    #
    # Sections:
    #

    def p_section_content_empty(self, p):
        """section_content : empty"""
        p[0] = []

    def p_section_content_assignation(self, p):
        """section_content : section_content assignment
                           | section_content section"""
        self._check_line(p.lexer.lineno, p.lineno(2),
                         self._lexer.column(p.lexpos(2)), p[2].name)
        p[1].append(p[2])
        p[0] = p[1]

    def p_section_content_include(self, p):
        """section_content : section_content INCLUDE TEXT"""
        for external in self._external_opener(p[3]):
            p[1] += list(external.iterflatchildren())
        p[0] = p[1]

    def p_section(self, p):
        """section : NAME LBRACE section_content RBRACE
                   | NAME section_args LBRACE section_content RBRACE"""
        name = p[1]
        if len(p) == 5:
            args = None
            section_content = p[3]
        else:
            args = p[2]
            section_content = p[4]
        column = p.lexer.column(p.lexpos(1))
        position = Position(self._input_name, p.lineno(1), column)
        section = ConfigSection(name, args=args, position=position)
        for child in section_content:
            if isinstance(child, ConfigSection):
                child.parent = section
            section.register(child)
        p[0] = section

    def p_section_args(self, p):
        """section_args : value LIST_SEP section_args
                        | value"""
        if len(p) == 2:
            position = Position(self._input_name,
                                p.lineno(1),
                                p.lexer.column(p.lexpos(1)))
            p[0] = ConfigValue('<args>', [p[1]], position=position)
        else:
            p[3].value.insert(0, p[1])
            p[0] = p[3]

    def p_empty(self, p):
        """empty :"""
        pass

    def p_error(self, token):
        column = self._lexer.column(token.lexpos)
        pos = Position(self._input_name, token.lineno, column)
        raise ParsingError('Syntax error near of "%s"' % token.value, pos)

    #
    # Bindings to the internel _parser object
    #

    def parse(self):
        return self._parser.parse(self._input, self._lexer, tracking=True)

    def __getattr__(self, name):
        attr = getattr(self._parser, name)
        if attr is None:
            raise AttributeError("'%s' object has no attribute '%s'" % (self, name))
        else:
            return attr
<commit_msg>Added input_name for lexer<commit_after>""" Dotconf lexer and parser.
"""

import sys
from glob import glob

import ply.lex as lex
import ply.yacc as yacc

from dotconf.tree import ConfigSection, ConfigValue, Position

UNITS = {'k': 10 ** 3,
         'M': 10 ** 6,
         'G': 10 ** 9,
         'T': 10 ** 12,
         'P': 10 ** 15,
         'E': 10 ** 18,
         'Z': 10 ** 21,
         'Y': 10 ** 24,
         'Ki': 2 ** 10,
         'Mi': 2 ** 20,
         'Gi': 2 ** 30,
         'Ti': 2 ** 40,
         'Pi': 2 ** 50,
         'Ei': 2 ** 60,
         'Zi': 2 ** 70,
         'Yi': 2 ** 80}

class ParsingError(Exception):

    """ Error raised when a parsing error occurs.
    """

    def __init__(self, msg, position=None):
        super(ParsingError, self).__init__(msg)
        self.position = position

def default_external_opener(locator):
    """ The default locator used to open included external files.
    """
    parsed_externals = []
    for external in glob(locator):
        try:
            with open(external) as fexternal:
                external_data = fexternal.read()
        except IOError as err:
            raise ParsingError('Unable to open %s (%s)' % (external, err))
        parser = DotconfParser(external_data, debug=False, write_tables=False,
                               errorlog=yacc.NullLogger(), input_name=external,
                               external_opener=default_external_opener)
        parsed_externals.append(parser.parse())
    return parsed_externals

#
# Lexer
#

class DotconfLexer(object):

    """ Lexer for the DotConf format.

    :param \*\*kwargs: arguments to give to the ply lexer

    Usage example::

    >>> lexer = DotConfLexer()
    >>> lexer.input('test { key = yes }')
    >>> print lexer.next()
    """

    def __init__(self, encoding='utf-8', input_name='<unknown>', **kwargs):
        self._lexer = lex.lex(module=self, **kwargs)
        self._encoding = encoding
        self._input_name = input_name

    #
    # Tokens definition
    #

    reserved = {'yes': 'YES', 'no': 'NO', 'include': 'INCLUDE'}
    reserved.update(dict(((kw, 'UNIT') for kw in UNITS)))
    tokens = ['LBRACE', 'RBRACE', 'NAME', 'TEXT', 'NUMBER',
              'ASSIGN', 'LIST_SEP'] + list(set(reserved.values()))

    t_LBRACE = '{'
    t_RBRACE = '}'
    t_ASSIGN = '='
    t_LIST_SEP = ','
    t_ignore = ' \t'

    def t_NAME(self, token):
        r'[a-zA-Z_][a-zA-Z0-9_-]*'
        token.type = self.reserved.get(token.value, 'NAME')
        # Handle the boolean case:
        if token.type == 'YES':
            token.value = True
        elif token.type == 'NO':
            token.value = False
        elif token.type == 'UNIT':
            token.value = UNITS[token.value]
        return token

    def t_TEXT(self, token):
        r'(["]([\\]["]|[^"]|)*["]|[\']([\\][\']|[^\'])*[\'])'
        value = token.value[1:-1].replace('\\' + token.value[0], token.value[0])
        token.value = value
        # Count the lines in the string:
        token.lexer.lineno += value.count('\n')
        return token

    def t_NUMBER(self, token):
        r'[-+]?[0-9]+(\.[0-9]+)?'
        if token.value.isdigit():
            token.value = int(token.value)
        else:
            token.value = float(token.value)
        return token

    def t_EOL(self, token):
        r'[\n]+'
        token.lexer.lineno += len(token.value)

    def t_COMMENT(self, token):
        r'[#].*'

    def t_error(self, token):
        position = Position(self._input_name,
                            self._lexer.lineno,
                            self._lexer.lexpos)
        raise ParsingError('Illegal character %r' % token.value[0], position)

    #
    # Public methods
    #

    def column(self, lexpos):
        """ Find the column according to the lexpos.
        """
        # This code is taken from the python-ply documentation
        # see: http://www.dabeaz.com/ply/ply.html section 4.6
        last_cr = self._current_input.rfind('\n', 0, lexpos)
        if last_cr < 0:
            last_cr = 0
        column = (lexpos - last_cr)
        return column

    #
    # Bindings to the internal _lexer object
    #

    def input(self, input):
        if sys.version_info[0] >= 3:
            if isinstance(input, bytes):
                input = input.decode(self._encoding)
        else:
            if isinstance(input, str):
                input = input.decode(self._encoding)
        self._current_input = input
        return self._lexer.input(input)

    def __getattr__(self, name):
        attr = getattr(self._lexer, name)
        if attr is None:
            raise AttributeError("'%s' object has no attribute '%s'" % (self, name))
        else:
            return attr

#
# Parser
#

class DotconfParser(object):

    """ Parser for the Dotconf format.
    """

    tokens = DotconfLexer.tokens

    def __init__(self, input, **kwargs):
        self._input = input
        self._input_name = kwargs.pop('input_name', '<unknown>')
        self._external_opener = kwargs.pop('external_opener',
                                           default_external_opener)
        self._lexer = kwargs.pop('lexer', DotconfLexer(input_name=self._input_name))
        self._parser = yacc.yacc(module=self, **kwargs)
        self._old_line = 0

    def _check_line(self, current, lineno, pos, token):
        if self._old_line == current:
            pos = Position(self._input_name, lineno, pos)
            raise ParsingError('Syntax error near of "%s", '
                               'newline missing?' % token, pos)
        else:
            self._old_line = current

    #
    # Rules
    #

    start = 'top'

    def p_top(self, p):
        """top : section_content"""
        section = ConfigSection('__top__')
        for child in p[1]:
            if isinstance(child, ConfigSection):
                child.parent = section
            section.register(child)
        p[0] = section

    def p_assignation(self, p):
        """assignment : NAME ASSIGN value
                      | NAME ASSIGN list"""
        position = Position(self._input_name,
                            p.lineno(3),
                            p.lexer.column(p.lexpos(3)))
        value = ConfigValue(p[1], p[3], position=position)
        p[0] = value

    def p_value(self, p):
        """value : TEXT
                 | YES
                 | NO
                 | number"""
        p[0] = p[1]

    def p_number(self, p):
        """number : NUMBER
                  | NUMBER UNIT"""
        if len(p) == 2:
            p[0] = p[1]
        else:
            p[0] = p[1] * p[2]

    #
    # List management rules:
    #

    def p_list(self, p):
        """list : value LIST_SEP list_next
                | value LIST_SEP"""
        if len(p) == 3:
            p[0] = [p[1]]
        elif len(p) == 4:
            p[3].insert(0, p[1])
            p[0] = p[3]

    def p_list_next(self, p):
        """list_next : value
                     | value LIST_SEP
                     | value LIST_SEP list_next"""
        if len(p) in (2, 3):
            p[0] = [p[1]]
        elif len(p) == 4:
            p[3].insert(0, p[1])
            p[0] = p[3]

    #
    # Sections:
    #

    def p_section_content_empty(self, p):
        """section_content : empty"""
        p[0] = []

    def p_section_content_assignation(self, p):
        """section_content : section_content assignment
                           | section_content section"""
        self._check_line(p.lexer.lineno, p.lineno(2),
                         self._lexer.column(p.lexpos(2)), p[2].name)
        p[1].append(p[2])
        p[0] = p[1]

    def p_section_content_include(self, p):
        """section_content : section_content INCLUDE TEXT"""
        for external in self._external_opener(p[3]):
            p[1] += list(external.iterflatchildren())
        p[0] = p[1]

    def p_section(self, p):
        """section : NAME LBRACE section_content RBRACE
                   | NAME section_args LBRACE section_content RBRACE"""
        name = p[1]
        if len(p) == 5:
            args = None
            section_content = p[3]
        else:
            args = p[2]
            section_content = p[4]
        column = p.lexer.column(p.lexpos(1))
        position = Position(self._input_name, p.lineno(1), column)
        section = ConfigSection(name, args=args, position=position)
        for child in section_content:
            if isinstance(child, ConfigSection):
                child.parent = section
            section.register(child)
        p[0] = section

    def p_section_args(self, p):
        """section_args : value LIST_SEP section_args
                        | value"""
        if len(p) == 2:
            position = Position(self._input_name,
                                p.lineno(1),
                                p.lexer.column(p.lexpos(1)))
            p[0] = ConfigValue('<args>', [p[1]], position=position)
        else:
            p[3].value.insert(0, p[1])
            p[0] = p[3]

    def p_empty(self, p):
        """empty :"""
        pass

    def p_error(self, token):
        column = self._lexer.column(token.lexpos)
        pos = Position(self._input_name, token.lineno, column)
        raise ParsingError('Syntax error near of "%s"' % token.value, pos)

    #
    # Bindings to the internel _parser object
    #

    def parse(self):
        return self._parser.parse(self._input, self._lexer, tracking=True)

    def __getattr__(self, name):
        attr = getattr(self._parser, name)
        if attr is None:
            raise AttributeError("'%s' object has no attribute '%s'" % (self, name))
        else:
            return attr
<|endoftext|>


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

"""
Implementation of a median subtraction algorithm for model PSF subtraction in
high-contrast imaging sequences. In the case of ADI, the algorithm is based on
[MAR06]_. The ADI+IFS method, is an extension of this basic idea to
multi-spectral cubes.

.. [MAR06]
   | Marois et al. 2006
   | **Angular Differential Imaging: A Powerful High-Contrast Imaging
     Technique**
   | *The Astrophysical Journal, Volume 641, Issue 1, pp. 556-564*
   | `https://arxiv.org/abs/astro-ph/0512335
     <https://arxiv.org/abs/astro-ph/0512335>`_

"""

from __future__ import division, print_function

__author__ = 'Carlos Alberto Gomez Gonzalez'
__all__ = ['median_sub']

import numpy as np
from multiprocessing import cpu_count
from ..conf import time_ini, timing
from ..var import get_annulus_segments, mask_circle
from ..preproc import (cube_derotate, cube_collapse, check_pa_vector,
                       check_scal_vector)
from ..preproc import cube_rescaling_wavelengths as scwave
from ..conf.utils_conf import pool_map, fixed, print_precision
from ..preproc.derotation import _find_indices_adi, _define_annuli
from ..preproc.rescaling import _find_indices_sdi

def median_sub(cube, angle_list, scale_list=None, fwhm=4, radius_int=0, asize=4,
               delta_rot=1, delta_sep=(0.1, 1), mode='fullfr', nframes=4,
               imlib='opencv', interpolation='lanczos4', collapse='median',
               nproc=1, full_output=False, verbose=True):
    """ Implementation of a median subtraction algorithm for model PSF
    subtraction in high-contrast imaging sequences. In the case of ADI, the
    algorithm is based on [MAR06]_. The ADI+IFS method is an extension of this
    basic idea to multi-spectral cubes.

    Parameters
    ----------
    cube : array_like, 3d
        Input cube.
    angle_list : array_like, 1d
        Corresponding parallactic angle for each frame.
    scale_list :
        Scaling factors in case of IFS data (ADI+mSDI cube). Usually, the
        scaling factors are the central channel wavelength divided by the
        shortest wavelength in the cube (more thorough approaches can be used
        to get the scaling factors). This scaling factors are used to re-scale
        the spectral channels and align the speckles.
    fwhm : float
        Known size of the FHWM in pixels to be used. Default is 4.
    radius_int : int, optional
        The radius of the innermost annulus. By default is 0, if >0 then the
        central circular area is discarded.
    asize : int, optional
        The size of the annuli, in pixels.
    delta_rot : int, optional
        Factor for increasing the parallactic angle threshold, expressed in
        FWHM. Default is 1 (excludes 1 FHWM on each side of the considered
        frame).
    delta_sep : float or tuple of floats, optional
        The threshold separation in terms of the mean FWHM (for ADI+mSDI data).
        If a tuple of two values is provided, they are used as the lower and
        upper intervals for the threshold (grows as a function of the
        separation).
    mode : {'fullfr', 'annular'}, str optional
        In ``fullfr`` mode only the median frame is subtracted, in ``annular``
        mode also the 4 closest frames given a PA threshold (annulus-wise) are
        subtracted.
    nframes : int or None, optional
        Number of frames (even value) to be used for building the optimized
        reference PSF when working in ``annular`` mode. None by default, which
        means that all frames, excluding the thresholded ones, are used.
    imlib : str, optional
        See the documentation of the ``vip_hci.preproc.frame_rotate`` function.
    interpolation : str, optional
        See the documentation of the ``vip_hci.preproc.frame_rotate`` function.
    collapse : {'median', 'mean', 'sum', 'trimmean'}, str optional
        Sets the way of collapsing the frames for producing a final image.
    nproc : None or int, optional
        Number of processes for parallel computing. If None the number of
        processes will be set to cpu_count()/2. By default the algorithm works
        in single-process mode.
    full_output: bool, optional
        Whether to return the final median combined image only or with other
        intermediate arrays.
    verbose : bool, optional
        If True prints to stdout intermediate info.

    Returns
    -------
    cube_out : array_like, 3d
        [full_output=True] The cube of residuals.
    cube_der : array_like, 3d
        [full_output=True] The derotated cube of residuals.
    frame : array_like, 2d
        Median combination of the de-rotated cube.

    """
    global ARRAY
    ARRAY = cube.copy()

    if not (ARRAY.ndim == 3 or ARRAY.ndim == 4):
        raise TypeError('Input array is not a 3d or 4d array')

    if verbose:
        start_time = time_ini()

    if nproc is None:
        nproc = cpu_count() // 2        # Hyper-threading doubles the # of cores

    angle_list = check_pa_vector(angle_list)

    if ARRAY.ndim == 3:
        n, y, _ = ARRAY.shape

        if ARRAY.shape[0] != angle_list.shape[0]:
            msg = 'Input vector or parallactic angles has wrong length'
            raise TypeError(msg)

        # The median frame is first subtracted from each frame
        model_psf = np.median(ARRAY, axis=0)
        ARRAY -= model_psf

        # Depending on the ``mode``
        if mode == 'fullfr':
            if radius_int > 0:
                cube_out = mask_circle(ARRAY, radius_int)
            else:
                cube_out = ARRAY
            if verbose:
                print('Median psf reference subtracted')

        elif mode == 'annular':
            if nframes is not None:
                if nframes % 2 != 0:
                    raise TypeError('`nframes` argument must be even value')

            n_annuli = int((y / 2 - radius_int) / asize)
            if verbose:
                print('N annuli = {}, FWHM = {}'.format(n_annuli, fwhm))

            res = pool_map(nproc, _median_subt_ann_adi, fixed(range(n_annuli)),
                           angle_list, n_annuli, fwhm, radius_int, asize,
                           delta_rot, nframes, verbose=verbose,
                           msg='Processing annuli:', progressbar_single=True)

            res = np.array(res)
            mres = res[:, 0]
            yy = res[:, 1]
            xx = res[:, 2]
            cube_out = np.zeros_like(ARRAY)
            for ann in range(n_annuli):
                cube_out[:, yy[ann], xx[ann]] = mres[ann]

            if verbose:
                print('Optimized median psf reference subtracted')

        else:
            raise RuntimeError('Mode not recognized')

        cube_der = cube_derotate(cube_out, angle_list, imlib=imlib,
                                 interpolation=interpolation)
        frame = cube_collapse(cube_der, mode=collapse)

    elif ARRAY.ndim == 4:
        z, n, y_in, x_in = ARRAY.shape

        if scale_list is None:
            raise ValueError('Scaling factors vector must be provided')
        else:
            if np.array(scale_list).ndim > 1:
                raise ValueError('Scaling factors vector is not 1d')
            if not scale_list.shape[0] == z:
                raise ValueError('Scaling factors vector has wrong length')

        # Exploiting spectral variability (radial movement)
        fwhm = int(np.round(np.mean(fwhm)))
        n_annuli = int((y_in / 2 - radius_int) / asize)

        if nframes is not None:
            if nframes % 2 != 0:
                raise TypeError('`nframes` argument must be even value')

        if verbose:
            print('{} spectral channels per IFS frame'.format(z))
            print('First median subtraction exploiting spectral variability')
            if mode == 'annular':
                print('N annuli = {}, mean FWHM = {:.3f}'.format(n_annuli,
                                                                 fwhm))

        res = pool_map(nproc, _median_subt_fr_sdi, fixed(range(n)),
                       scale_list, n_annuli, fwhm, radius_int, asize,
                       delta_sep, nframes, imlib, interpolation, collapse,
                       mode)
        residuals_cube_channels = np.array(res)

        if verbose:
            timing(start_time)
            print('{} ADI frames'.format(n))
            print('Median subtraction in the ADI fashion')

        if mode == 'fullfr':
            median_frame = np.median(residuals_cube_channels, axis=0)
            residuals_final = residuals_cube_channels - median_frame
            residuals_final_der = cube_derotate(residuals_final, angle_list,
                                                imlib=imlib,
                                                interpolation=interpolation)
            frame = cube_collapse(residuals_final_der, mode=collapse)
            if verbose:
                timing(start_time)

            cube_out = None
            cube_der = None

        elif mode == 'annular':
            if verbose:
                print('N annuli = {}, mean FWHM = {:.3f}'.format(n_annuli,
                                                                 fwhm))
            ARRAY = residuals_cube_channels

            res = pool_map(nproc, _median_subt_ann_adi, fixed(range(n_annuli)),
                           angle_list, n_annuli, fwhm, radius_int, asize,
                           delta_rot, nframes)

            res = np.array(res)
            mres = res[:, 0]
            yy = res[:, 1]
            xx = res[:, 2]
            pa_thrs = np.array(res[:, 3])
            if verbose:
                print('PA thresholds: ')
                print_precision(pa_thrs)

            cube_out = np.zeros_like(ARRAY)
            for ann in range(n_annuli):
                cube_out[:, yy[ann], xx[ann]] = mres[ann]

            cube_der = cube_derotate(cube_out, angle_list, imlib=imlib,
                                     interpolation=interpolation)
            frame = cube_collapse(cube_der, mode=collapse)

        else:
            raise RuntimeError('Mode not recognized')

    if verbose:
        print('Done derotating and combining')
        timing(start_time)
    if full_output:
        return cube_out, cube_der, frame
    else:
        return frame

def _median_subt_fr_sdi(fr, wl, n_annuli, fwhm, radius_int, annulus_width,
                        delta_sep, nframes, imlib, interpolation, collapse,
                        mode):
    """ Optimized median subtraction on a multi-spectral frame (IFS data).
    """
    z, n, y_in, x_in = ARRAY.shape
    scale_list = check_scal_vector(wl)
    multispec_fr = scwave(ARRAY[:, fr, :, :], scale_list, imlib=imlib,
                          interpolation=interpolation)[0]    # rescaled cube

    if mode == 'annular':
        cube_res = np.zeros_like(multispec_fr)  # shape (z, resc_y, resc_x)

        if isinstance(delta_sep, tuple):
            delta_sep_vec = np.linspace(delta_sep[0], delta_sep[1], n_annuli)
        else:
            delta_sep_vec = [delta_sep] * n_annuli

        for ann in range(n_annuli):
            if ann == n_annuli - 1:
                inner_radius = radius_int + (ann * annulus_width - 1)
            else:
                inner_radius = radius_int + ann * annulus_width
            ann_center = inner_radius + (annulus_width / 2)

            indices = get_annulus_segments(multispec_fr[0], inner_radius,
                                           annulus_width)[0]
            yy = indices[0]
            xx = indices[1]
            matrix = multispec_fr[:, yy, xx]  # shape (z, npx_annulus)

            for j in range(z):
                indices_left = _find_indices_sdi(wl, ann_center, j, fwhm,
                                                 delta_sep_vec[ann], nframes)
                matrix_masked = matrix[indices_left]
                ref_psf_opt = np.median(matrix_masked, axis=0)
                curr_wv = matrix[j]
                subtracted = curr_wv - ref_psf_opt
                cube_res[j, yy, xx] = subtracted

    elif mode == 'fullfr':
        median_frame = np.median(multispec_fr, axis=0)
        cube_res = multispec_fr - median_frame

    frame_desc = scwave(cube_res, scale_list, full_output=False, inverse=True,
                        y_in=y_in, x_in=x_in, imlib=imlib,
                        interpolation=interpolation, collapse=collapse)
    return frame_desc

def _median_subt_ann_adi(ann, angle_list, n_annuli, fwhm, radius_int,
                         annulus_width, delta_rot, nframes):
    """ Optimized median subtraction for a given annulus.
    """
    if ARRAY.ndim == 3:
        n = ARRAY.shape[0]
    elif ARRAY.ndim == 4:
        n = ARRAY.shape[1]

    # The annulus is built, and the corresponding PA thresholds for frame
    # rejection are calculated. The PA rejection is calculated at center of
    # the annulus
    pa_thr, inner_radius, _ = _define_annuli(angle_list, ann, n_annuli, fwhm,
                                             radius_int, annulus_width,
                                             delta_rot, 1, False)
    if ARRAY.ndim == 3:
        indices = get_annulus_segments(ARRAY[0], inner_radius, annulus_width)[0]
    elif ARRAY.ndim == 4:
        indices = get_annulus_segments(ARRAY[0, 0], inner_radius,
                                       annulus_width)[0]
    yy = indices[0]
    xx = indices[1]
    matrix = ARRAY[:, yy, xx]  # shape [n x npx_annulus]
    matrix_res = np.zeros_like(matrix)

    # A second optimized psf reference is subtracted from each frame.
    # For each frame we find ``nframes``, depending on the PA threshold,
    # to construct this optimized psf reference
    for frame in range(n):
        if pa_thr != 0:
            indices_left = _find_indices_adi(angle_list, frame, pa_thr, nframes)
            matrix_disc = matrix[indices_left]
        else:
            matrix_disc = matrix

        ref_psf_opt = np.median(matrix_disc, axis=0)
        curr_frame = matrix[frame]
        subtracted = curr_frame - ref_psf_opt
        matrix_res[frame] = subtracted

    return matrix_res, yy, xx, pa_thr

<commit_msg>FIX: rename missing variables<commit_after>#! /usr/bin/env python

"""
Implementation of a median subtraction algorithm for model PSF subtraction in
high-contrast imaging sequences. In the case of ADI, the algorithm is based on
[MAR06]_. The ADI+IFS method, is an extension of this basic idea to
multi-spectral cubes.

.. [MAR06]
   | Marois et al. 2006
   | **Angular Differential Imaging: A Powerful High-Contrast Imaging
     Technique**
   | *The Astrophysical Journal, Volume 641, Issue 1, pp. 556-564*
   | `https://arxiv.org/abs/astro-ph/0512335
     <https://arxiv.org/abs/astro-ph/0512335>`_

"""

from __future__ import division, print_function

__author__ = 'Carlos Alberto Gomez Gonzalez'
__all__ = ['median_sub']

import numpy as np
from multiprocessing import cpu_count
from ..conf import time_ini, timing
from ..var import get_annulus_segments, mask_circle
from ..preproc import (cube_derotate, cube_collapse, check_pa_vector,
                       check_scal_vector)
from ..preproc import cube_rescaling_wavelengths as scwave
from ..conf.utils_conf import pool_map, fixed, print_precision
from ..preproc.derotation import _find_indices_adi, _define_annuli
from ..preproc.rescaling import _find_indices_sdi

def median_sub(cube, angle_list, scale_list=None, fwhm=4, radius_int=0, asize=4,
               delta_rot=1, delta_sep=(0.1, 1), mode='fullfr', nframes=4,
               imlib='opencv', interpolation='lanczos4', collapse='median',
               nproc=1, full_output=False, verbose=True):
    """ Implementation of a median subtraction algorithm for model PSF
    subtraction in high-contrast imaging sequences. In the case of ADI, the
    algorithm is based on [MAR06]_. The ADI+IFS method is an extension of this
    basic idea to multi-spectral cubes.

    Parameters
    ----------
    cube : array_like, 3d
        Input cube.
    angle_list : array_like, 1d
        Corresponding parallactic angle for each frame.
    scale_list :
        Scaling factors in case of IFS data (ADI+mSDI cube). Usually, the
        scaling factors are the central channel wavelength divided by the
        shortest wavelength in the cube (more thorough approaches can be used
        to get the scaling factors). This scaling factors are used to re-scale
        the spectral channels and align the speckles.
    fwhm : float
        Known size of the FHWM in pixels to be used. Default is 4.
    radius_int : int, optional
        The radius of the innermost annulus. By default is 0, if >0 then the
        central circular area is discarded.
    asize : int, optional
        The size of the annuli, in pixels.
    delta_rot : int, optional
        Factor for increasing the parallactic angle threshold, expressed in
        FWHM. Default is 1 (excludes 1 FHWM on each side of the considered
        frame).
    delta_sep : float or tuple of floats, optional
        The threshold separation in terms of the mean FWHM (for ADI+mSDI data).
        If a tuple of two values is provided, they are used as the lower and
        upper intervals for the threshold (grows as a function of the
        separation).
    mode : {'fullfr', 'annular'}, str optional
        In ``fullfr`` mode only the median frame is subtracted, in ``annular``
        mode also the 4 closest frames given a PA threshold (annulus-wise) are
        subtracted.
    nframes : int or None, optional
        Number of frames (even value) to be used for building the optimized
        reference PSF when working in ``annular`` mode. None by default, which
        means that all frames, excluding the thresholded ones, are used.
    imlib : str, optional
        See the documentation of the ``vip_hci.preproc.frame_rotate`` function.
    interpolation : str, optional
        See the documentation of the ``vip_hci.preproc.frame_rotate`` function.
    collapse : {'median', 'mean', 'sum', 'trimmean'}, str optional
        Sets the way of collapsing the frames for producing a final image.
    nproc : None or int, optional
        Number of processes for parallel computing. If None the number of
        processes will be set to cpu_count()/2. By default the algorithm works
        in single-process mode.
    full_output: bool, optional
        Whether to return the final median combined image only or with other
        intermediate arrays.
    verbose : bool, optional
        If True prints to stdout intermediate info.

    Returns
    -------
    cube_out : array_like, 3d
        [full_output=True] The cube of residuals.
    cube_der : array_like, 3d
        [full_output=True] The derotated cube of residuals.
    frame : array_like, 2d
        Median combination of the de-rotated cube.

    """
    global ARRAY
    ARRAY = cube.copy()

    if not (ARRAY.ndim == 3 or ARRAY.ndim == 4):
        raise TypeError('Input array is not a 3d or 4d array')

    if verbose:
        start_time = time_ini()

    if nproc is None:
        nproc = cpu_count() // 2        # Hyper-threading doubles the # of cores

    angle_list = check_pa_vector(angle_list)

    if ARRAY.ndim == 3:
        n, y, _ = ARRAY.shape

        if ARRAY.shape[0] != angle_list.shape[0]:
            msg = 'Input vector or parallactic angles has wrong length'
            raise TypeError(msg)

        # The median frame is first subtracted from each frame
        model_psf = np.median(ARRAY, axis=0)
        ARRAY -= model_psf

        # Depending on the ``mode``
        if mode == 'fullfr':
            if radius_int > 0:
                cube_out = mask_circle(ARRAY, radius_int)
            else:
                cube_out = ARRAY
            if verbose:
                print('Median psf reference subtracted')

        elif mode == 'annular':
            if nframes is not None:
                if nframes % 2 != 0:
                    raise TypeError('`nframes` argument must be even value')

            n_annuli = int((y / 2 - radius_int) / asize)
            if verbose:
                print('N annuli = {}, FWHM = {}'.format(n_annuli, fwhm))

            res = pool_map(nproc, _median_subt_ann_adi, fixed(range(n_annuli)),
                           angle_list, n_annuli, fwhm, radius_int, asize,
                           delta_rot, nframes, verbose=verbose,
                           msg='Processing annuli:', progressbar_single=True)

            res = np.array(res)
            mres = res[:, 0]
            yy = res[:, 1]
            xx = res[:, 2]
            cube_out = np.zeros_like(ARRAY)
            for ann in range(n_annuli):
                cube_out[:, yy[ann], xx[ann]] = mres[ann]

            if verbose:
                print('Optimized median psf reference subtracted')

        else:
            raise RuntimeError('Mode not recognized')

        cube_der = cube_derotate(cube_out, angle_list, imlib=imlib,
                                 interpolation=interpolation)
        frame = cube_collapse(cube_der, mode=collapse)

    elif ARRAY.ndim == 4:
        z, n, y_in, x_in = ARRAY.shape

        if scale_list is None:
            raise ValueError('Scaling factors vector must be provided')
        else:
            if np.array(scale_list).ndim > 1:
                raise ValueError('Scaling factors vector is not 1d')
            if not scale_list.shape[0] == z:
                raise ValueError('Scaling factors vector has wrong length')

        # Exploiting spectral variability (radial movement)
        fwhm = int(np.round(np.mean(fwhm)))
        n_annuli = int((y_in / 2 - radius_int) / asize)

        if nframes is not None:
            if nframes % 2 != 0:
                raise TypeError('`nframes` argument must be even value')

        if verbose:
            print('{} spectral channels per IFS frame'.format(z))
            print('First median subtraction exploiting spectral variability')
            if mode == 'annular':
                print('N annuli = {}, mean FWHM = {:.3f}'.format(n_annuli,
                                                                 fwhm))

        res = pool_map(nproc, _median_subt_fr_sdi, fixed(range(n)),
                       scale_list, n_annuli, fwhm, radius_int, asize,
                       delta_sep, nframes, imlib, interpolation, collapse,
                       mode)
        residuals_cube_channels = np.array(res)

        if verbose:
            timing(start_time)
            print('{} ADI frames'.format(n))
            print('Median subtraction in the ADI fashion')

        if mode == 'fullfr':
            median_frame = np.median(residuals_cube_channels, axis=0)
            cube_out = residuals_cube_channels - median_frame
            cube_der = cube_derotate(cube_out, angle_list, imlib=imlib,
                                     interpolation=interpolation)
            frame = cube_collapse(cube_der, mode=collapse)
            if verbose:
                timing(start_time)

        elif mode == 'annular':
            if verbose:
                print('N annuli = {}, mean FWHM = {:.3f}'.format(n_annuli,
                                                                 fwhm))
            ARRAY = residuals_cube_channels

            res = pool_map(nproc, _median_subt_ann_adi, fixed(range(n_annuli)),
                           angle_list, n_annuli, fwhm, radius_int, asize,
                           delta_rot, nframes)

            res = np.array(res)
            mres = res[:, 0]
            yy = res[:, 1]
            xx = res[:, 2]
            pa_thrs = np.array(res[:, 3])
            if verbose:
                print('PA thresholds: ')
                print_precision(pa_thrs)

            cube_out = np.zeros_like(ARRAY)
            for ann in range(n_annuli):
                cube_out[:, yy[ann], xx[ann]] = mres[ann]

            cube_der = cube_derotate(cube_out, angle_list, imlib=imlib,
                                     interpolation=interpolation)
            frame = cube_collapse(cube_der, mode=collapse)

        else:
            raise RuntimeError('Mode not recognized')

    if verbose:
        print('Done derotating and combining')
        timing(start_time)
    if full_output:
        return cube_out, cube_der, frame
    else:
        return frame

def _median_subt_fr_sdi(fr, wl, n_annuli, fwhm, radius_int, annulus_width,
                        delta_sep, nframes, imlib, interpolation, collapse,
                        mode):
    """ Optimized median subtraction on a multi-spectral frame (IFS data).
    """
    z, n, y_in, x_in = ARRAY.shape
    scale_list = check_scal_vector(wl)
    multispec_fr = scwave(ARRAY[:, fr, :, :], scale_list, imlib=imlib,
                          interpolation=interpolation)[0]    # rescaled cube

    if mode == 'annular':
        cube_res = np.zeros_like(multispec_fr)  # shape (z, resc_y, resc_x)

        if isinstance(delta_sep, tuple):
            delta_sep_vec = np.linspace(delta_sep[0], delta_sep[1], n_annuli)
        else:
            delta_sep_vec = [delta_sep] * n_annuli

        for ann in range(n_annuli):
            if ann == n_annuli - 1:
                inner_radius = radius_int + (ann * annulus_width - 1)
            else:
                inner_radius = radius_int + ann * annulus_width
            ann_center = inner_radius + (annulus_width / 2)

            indices = get_annulus_segments(multispec_fr[0], inner_radius,
                                           annulus_width)[0]
            yy = indices[0]
            xx = indices[1]
            matrix = multispec_fr[:, yy, xx]  # shape (z, npx_annulus)

            for j in range(z):
                indices_left = _find_indices_sdi(wl, ann_center, j, fwhm,
                                                 delta_sep_vec[ann], nframes)
                matrix_masked = matrix[indices_left]
                ref_psf_opt = np.median(matrix_masked, axis=0)
                curr_wv = matrix[j]
                subtracted = curr_wv - ref_psf_opt
                cube_res[j, yy, xx] = subtracted

    elif mode == 'fullfr':
        median_frame = np.median(multispec_fr, axis=0)
        cube_res = multispec_fr - median_frame

    frame_desc = scwave(cube_res, scale_list, full_output=False, inverse=True,
                        y_in=y_in, x_in=x_in, imlib=imlib,
                        interpolation=interpolation, collapse=collapse)
    return frame_desc

def _median_subt_ann_adi(ann, angle_list, n_annuli, fwhm, radius_int,
                         annulus_width, delta_rot, nframes):
    """ Optimized median subtraction for a given annulus.
    """
    if ARRAY.ndim == 3:
        n = ARRAY.shape[0]
    elif ARRAY.ndim == 4:
        n = ARRAY.shape[1]

    # The annulus is built, and the corresponding PA thresholds for frame
    # rejection are calculated. The PA rejection is calculated at center of
    # the annulus
    pa_thr, inner_radius, _ = _define_annuli(angle_list, ann, n_annuli, fwhm,
                                             radius_int, annulus_width,
                                             delta_rot, 1, False)
    if ARRAY.ndim == 3:
        indices = get_annulus_segments(ARRAY[0], inner_radius, annulus_width)[0]
    elif ARRAY.ndim == 4:
        indices = get_annulus_segments(ARRAY[0, 0], inner_radius,
                                       annulus_width)[0]
    yy = indices[0]
    xx = indices[1]
    matrix = ARRAY[:, yy, xx]  # shape [n x npx_annulus]
    matrix_res = np.zeros_like(matrix)

    # A second optimized psf reference is subtracted from each frame.
    # For each frame we find ``nframes``, depending on the PA threshold,
    # to construct this optimized psf reference
    for frame in range(n):
        if pa_thr != 0:
            indices_left = _find_indices_adi(angle_list, frame, pa_thr, nframes)
            matrix_disc = matrix[indices_left]
        else:
            matrix_disc = matrix

        ref_psf_opt = np.median(matrix_disc, axis=0)
        curr_frame = matrix[frame]
        subtracted = curr_frame - ref_psf_opt
        matrix_res[frame] = subtracted

    return matrix_res, yy, xx, pa_thr

<|endoftext|>


<commit_before>"""Tests for module gramcore.filters.edges"""
import os
import numpy

from nose.tools import assert_equal

from skimage import io

from gramcore.filters import edges

def setup():
    """Create image fixture

    It creates a 20x40 image with black background and paints an 11x11 white
    rectangle inside. Choose odd dimensions for the filters to work more
    predictively.

    """
    img = numpy.zeros((20, 40))
    img[5:16, 15:26] = 255
    io.imsave('white-square.tif', img.astype('uint8'))

def teardown():
    """Delete fixture"""
    #os.remove('white-square.tif')

def test_canny():
    """Apply canny to grey image and check return values

    .. warning::

        This seems to produce some artifacts. You expect you get 44 (4*11)
        pixels of edges. Instead it gets 46, when sigma is 1 and 40 when sigma
        is 2. In both cases the shape is correct, but a few edge pixels are
        missing or instroduced at a wrong positions.

    """
    img = io.imread('white-square.tif')

    parameters = {'data': [img], 'sigma': 1.0}

    result = edges.canny(parameters)

    # this should be 44 check the resulting image with
    #result *= 255
    #io.imsave('result.tif', result)
    assert_equal(result.sum(), 46)

def test_prewitt():
    """Apply prewitt to grey image and check return values

    .. note::

        This produces correct shape though it shrinks it by 2 pixels, there
        are no edge pixels on the corners and each edge has a width of 2
        pixels. Based on the original rectangle size, which is 11x11, and the
        above issues it returns 4*9*2 = 72 edge pixels.

    """
    img = io.imread('white-square.tif')

    parameters = {'data': [img]}

    result = edges.prewitt(parameters)
    result = result.astype('uint8')

    assert_equal(result.sum(), 72)

def test_sobel():
    """Apply sobel to grey image and check return values

    .. note::

        This produces correct shape though it shrinks it by 2 pixels and each
        edge has a width of 2 pixels. Based on the original rectangle size,
        which is 11x11, and the above issues it returns 4*9*2 + 4 = 76 edge
        pixels.

    """
    img = io.imread('white-square.tif')

    parameters = {'data': [img]}

    result = edges.sobel(parameters)
    result = result.astype('uint8')

    assert_equal(result.sum(), 76)
<commit_msg>Delete fixture image on tearDown<commit_after>"""Tests for module gramcore.filters.edges"""
import os
import numpy

from nose.tools import assert_equal

from skimage import io

from gramcore.filters import edges

def setup():
    """Create image fixture

    It creates a 20x40 image with black background and paints an 11x11 white
    rectangle inside. Choose odd dimensions for the filters to work more
    predictively.

    """
    img = numpy.zeros((20, 40))
    img[5:16, 15:26] = 255
    io.imsave('white-square.tif', img.astype('uint8'))

def teardown():
    """Delete fixture"""
    os.remove('white-square.tif')

def test_canny():
    """Apply canny to grey image and check return values

    .. warning::

        This seems to produce some artifacts. You expect you get 44 (4*11)
        pixels of edges. Instead it gets 46, when sigma is 1 and 40 when sigma
        is 2. In both cases the shape is correct, but a few edge pixels are
        missing or instroduced at a wrong positions.

    """
    img = io.imread('white-square.tif')

    parameters = {'data': [img], 'sigma': 1.0}

    result = edges.canny(parameters)

    # this should be 44 check the resulting image with
    #result *= 255
    #io.imsave('result.tif', result)
    assert_equal(result.sum(), 46)

def test_prewitt():
    """Apply prewitt to grey image and check return values

    .. note::

        This produces correct shape though it shrinks it by 2 pixels, there
        are no edge pixels on the corners and each edge has a width of 2
        pixels. Based on the original rectangle size, which is 11x11, and the
        above issues it returns 4*9*2 = 72 edge pixels.

    """
    img = io.imread('white-square.tif')

    parameters = {'data': [img]}

    result = edges.prewitt(parameters)
    result = result.astype('uint8')

    assert_equal(result.sum(), 72)

def test_sobel():
    """Apply sobel to grey image and check return values

    .. note::

        This produces correct shape though it shrinks it by 2 pixels and each
        edge has a width of 2 pixels. Based on the original rectangle size,
        which is 11x11, and the above issues it returns 4*9*2 + 4 = 76 edge
        pixels.

    """
    img = io.imread('white-square.tif')

    parameters = {'data': [img]}

    result = edges.sobel(parameters)
    result = result.astype('uint8')

    assert_equal(result.sum(), 76)
<|endoftext|>