<commit_before>from collections import defaultdict

from fabric.api import env, local, run, sudo, cd, hide, prefix
from fabric.context_managers import shell_env, prefix
from fabric.operations import put, get
from fabric.contrib.files import append, exists
env.use_ssh_config = True

def prepare():
    prepare_apt()
    checkout_cache()

def prepare_apt():
    sudo("apt-get -qq update")
    sudo("apt-get -y install git python3 make python-virtualenv zip python-dev")
    # Needed to build the docs
    sudo("apt-get -y install graphviz inkscape texlive texlive-xetex texlive-fonts-recommended texlive-latex-extra")

def remove_userspace():
    """
    Deletes (!) the SymPy changes. Use with great care.
    """
    run("rm -rf repos")

def checkout_cache():
    run("rm -rf sympy-cache.git")
    run("git clone --bare https://github.com/sympy/sympy.git sympy-cache.git")

def gitrepos(branch=None):
    if not branch:
        # Use the current branch (of this git repo, not the one in Vagrant)
        branch = local("git rev-parse --abbrev-ref HEAD", capture=True)
    run("mkdir -p repos")
    with cd("repos"):
        run("git clone --reference ../sympy-cache.git https://github.com/sympy/sympy.git")
        if branch != "master":
            with cd("sympy"):
                run("git checkout -t origin/%s" % branch)

def get_sympy_version():
    with cd("repos/sympy"):
        version = run('python -c "import sympy;print sympy.__version__"')
    assert '\n' not in version
    assert ' ' not in version
    assert '\t' not in version
    return version

def test_git():
    with cd("repos/sympy"):
        run("./setup.py test")

def test_tarball(release='2'):
    if release not in {'2', '3'}: # TODO: Add win32
        raise ValueError("release must be one of '2', '3', not %s" % release)
    run("virtualenv test-{release}-virtualenv".format(release=release))
    run("source test-{release}-virtualenv/bin/activate".format(release=release))
    if release == '2':
        run("cp /vagrant/release/{2} releasetar.tar".format(**tarball_formatter))
    if release == '3':
        run("cp /vagrant/release/{2} releasetar.tar".format(**tarball_formatter))
    run("tar xvf releasetar.tar")
    with cd("releastar"):
        run("python setup.py install")
        run('python -c "import sympy; print sympy.__version__"')

def release(branch=None):
    remove_userspace()
    gitrepos(branch)
    python2_tarball()
    python3_tarball()
    build_docs()
    copy_release_files()
    test_tarball('2')
    test_tarball('3')

def python2_tarball():
    with cd("repos/sympy"):
        run("git clean -dfx")
        run("./setup.py clean")
        run("./setup.py sdist")
        run("./setup.py bdist_wininst")
        run("mv dist/{2win32-orig} dist/{2win32}".format(**tarball_formatter))

def python3_tarball():
    with cd("repos/sympy"):
        run("bin/use2to3")
        with cd("py3k-sympy"):
            run("./setup.py clean")
            run("./setup.py sdist")
            # We have to have 3.2 and 3.3 tarballs to make things work in
            # pip. See https://groups.google.com/d/msg/sympy/JEwi4ohGB90/FfjVDxZIkSEJ.
            run("mv dist/{source-orig} dist/{3.2}".format(**tarball_formatter))
            run("cp dist/{3.2} dist/{3.3}".format(**tarball_formatter))
            # We didn't test this yet:
            #run("./setup.py bdist_wininst")

def build_docs():
    with cd("repos/sympy"):
        run("mkdir -p dist")
        run("virtualenv docs-virtualenv")
        run("source docs-virtualenv/bin/activate; pip install sphinx==1.1.3 numpy")
        with cd("doc"):
            run("make clean")
            run("source ../docs-virtualenv/bin/activate; make html-errors")
            with cd("_build"):
                run("mv html {html-nozip}".format(**tarball_formatter))
                run("zip -9lr {html} {html-nozip}".format(**tarball_formatter))
                run("cp {html} ../../dist/".format(**tarball_formatter))
            run("make clean")
            run("source ../docs-virtualenv/bin/activate; make latex")
            with cd("_build"):
                with cd("latex"):
                    run("make")
                    run("cp {pdf-orig} ../../../dist/{pdf}".format(**tarball_formatter))

def copy_release_files():
    with cd("repos/sympy"):
        run("mkdir -p /vagrant/release")
        run("cp dist/* /vagrant/release/")
        run("cp py3k-sympy/dist/* /vagrant/release/")

def show_files(file):
    """
    Show the contents of a tarball.

    The current options for file are

    2: The Python 2 tarball
    3: The Python 3 tarball
    2win: The Python 2 Windows installer (Not yet implemented!)
    3win: The Python 3 Windows installer (Not yet implemented!)
    html: The html docs zip
    """
    # TODO:
    # - Automatically check that Python 3 has the same files as Python 2
    # - List the files that are in git but not in the release
    # - List the files in the Windows installers
    if file == '2':
        local("tar tf release/{2}".format(**tarball_formatter))
    elif file == '3':
        py32 = "{3.2}".format(**tarball_formatter)
        py33 = "{3.3}".format(**tarball_formatter)
        assert md5(py32).split()[0] == md5(py33).split()[0]
        local("tar tf release/" + py32)
    elif file in {'2win', '3win'}:
        raise NotImplementedError("Windows installers")
    elif file == 'html':
        local("unzip -l release/{html}".format(**tarball_formatter))
    else:
        raise ValueError(file + " is not valid")

def md5(file='*'):
    out = local("md5sum release/" + file, capture=True)
    print out
    return out

def get_tarball_name(file):
    """
    Get the name of a tarball

    file should be one of

    source-orig: The original name of the source tarball
    2:           The Python 2 tarball (after renaming)
    3.2:         The Python 3.2 tarball (after renaming)
    3.3:         The Python 3.3 tarball (after renaming)
    2win32-orig: The original name of the Python 2 win32 installer
    2win32:      The name of the Python 2 win32 installer (after renaming)
    html:        The name of the html zip
    html-nozip:  The name of the html, without ".zip"
    pdf-orig:    The original name of the pdf file
    pdf:         The name of the pdf file (after renaming)
    """
    version = get_sympy_version()
    doctypename = defaultdict(str, {'html': 'zip', 'pdf': 'pdf'})
    winos = defaultdict(str, {'2win32': 'win32', '2win32-orig': 'linux-i686'})
    if file in {'source-orig', '2'}:
        name = 'sympy-{version}.tar.gz'
    elif file in {'3.2', '3.3'}:
        name = "sympy-{version}-py{pyversion}.tar.gz"
    elif file in {'2win32', '2win32-orig'}:
        name = "sympy-{version}.{wintype}.exe"
    elif file in {'html', 'pdf', 'html-nozip'}:
        name = "sympy-docs-{type}-{version}"
        if not file.endswith('nozip'):
            file += ".{extension}"
    elif file == 'pdf-orig':
        name = "sympy-{version}.pdf"
    else:
        raise ValueError(file + " is not a recognized argument")
    ret = name.format(version=version, pyversion=file, type=file,
        extension=doctypename[file], wintype=winos[file])
    print ret # REMOVE ME
    return ret

tarball_name_types = {
    'source-orig:',
    '2:',
    '3.2:',
    '3.3:',
    '2win32-orig:',
    '2win32:',
    'html:',
    'pdf-orig:',
    'pdf',
    }

tarball_formatter = {name: get_tarball_name(name) for name in tarball_name_types}

# ------------------------------------------------
# Vagrant related configuration

def vagrant():
    vc = _get_vagrant_config()
    # change from the default user to 'vagrant'
    env.user = vc['User']
    # connect to the port-forwarded ssh
    env.hosts = ['%s:%s' % (vc['HostName'], vc['Port'])]
    # use vagrant ssh key
    env.key_filename = vc['IdentityFile'].strip('"')
    # Forward the agent if specified:
    env.forward_agent = vc.get('ForwardAgent', 'no') == 'yes'

def _get_vagrant_config():
    """
    Parses vagrant configuration and returns it as dict of ssh parameters
    and their values
    """
    result = local('vagrant ssh-config', capture=True)
    conf = {}
    for line in iter(result.splitlines()):
        parts = line.split()
        conf[parts[0]] = ' '.join(parts[1:])
    return conf

def restart_network():
    """
    Do this if the VM won't connect to the internet.
    """
    run("sudo /etc/init.d/networking restart")

# ---------------------------------------
# Just a simple testing command:

def uname():
    run('uname -a')
<commit_msg>Fix the "No hosts found" issue<commit_after>from collections import defaultdict

from fabric.api import env, local, run, sudo, cd, hide, prefix
from fabric.context_managers import shell_env, prefix
from fabric.operations import put, get
from fabric.contrib.files import append, exists
env.use_ssh_config = True

def prepare():
    prepare_apt()
    checkout_cache()

def prepare_apt():
    sudo("apt-get -qq update")
    sudo("apt-get -y install git python3 make python-virtualenv zip python-dev")
    # Needed to build the docs
    sudo("apt-get -y install graphviz inkscape texlive texlive-xetex texlive-fonts-recommended texlive-latex-extra")

def remove_userspace():
    """
    Deletes (!) the SymPy changes. Use with great care.
    """
    run("rm -rf repos")

def checkout_cache():
    run("rm -rf sympy-cache.git")
    run("git clone --bare https://github.com/sympy/sympy.git sympy-cache.git")

def gitrepos(branch=None):
    if not branch:
        # Use the current branch (of this git repo, not the one in Vagrant)
        branch = local("git rev-parse --abbrev-ref HEAD", capture=True)
    run("mkdir -p repos")
    with cd("repos"):
        run("git clone --reference ../sympy-cache.git https://github.com/sympy/sympy.git")
        if branch != "master":
            with cd("sympy"):
                run("git checkout -t origin/%s" % branch)

def get_sympy_version():
    with cd("repos/sympy"):
        version = run('python -c "import sympy;print sympy.__version__"')
    assert '\n' not in version
    assert ' ' not in version
    assert '\t' not in version
    return version

def test_git():
    with cd("repos/sympy"):
        run("./setup.py test")

def test_tarball(release='2'):
    if release not in {'2', '3'}: # TODO: Add win32
        raise ValueError("release must be one of '2', '3', not %s" % release)
    run("virtualenv test-{release}-virtualenv".format(release=release))
    run("source test-{release}-virtualenv/bin/activate".format(release=release))
    if release == '2':
        run("cp /vagrant/release/{2} releasetar.tar".format(**tarball_formatter()))
    if release == '3':
        run("cp /vagrant/release/{2} releasetar.tar".format(**tarball_formatter()))
    run("tar xvf releasetar.tar")
    with cd("releastar"):
        run("python setup.py install")
        run('python -c "import sympy; print sympy.__version__"')

def release(branch=None):
    remove_userspace()
    gitrepos(branch)
    python2_tarball()
    python3_tarball()
    build_docs()
    copy_release_files()
    test_tarball('2')
    test_tarball('3')

def python2_tarball():
    with cd("repos/sympy"):
        run("git clean -dfx")
        run("./setup.py clean")
        run("./setup.py sdist")
        run("./setup.py bdist_wininst")
        run("mv dist/{2win32-orig} dist/{2win32}".format(**tarball_formatter()))

def python3_tarball():
    with cd("repos/sympy"):
        run("bin/use2to3")
        with cd("py3k-sympy"):
            run("./setup.py clean")
            run("./setup.py sdist")
            # We have to have 3.2 and 3.3 tarballs to make things work in
            # pip. See https://groups.google.com/d/msg/sympy/JEwi4ohGB90/FfjVDxZIkSEJ.
            run("mv dist/{source-orig} dist/{3.2}".format(**tarball_formatter()))
            run("cp dist/{3.2} dist/{3.3}".format(**tarball_formatter()))
            # We didn't test this yet:
            #run("./setup.py bdist_wininst")

def build_docs():
    with cd("repos/sympy"):
        run("mkdir -p dist")
        run("virtualenv docs-virtualenv")
        run("source docs-virtualenv/bin/activate; pip install sphinx==1.1.3 numpy")
        with cd("doc"):
            run("make clean")
            run("source ../docs-virtualenv/bin/activate; make html-errors")
            with cd("_build"):
                run("mv html {html-nozip}".format(**tarball_formatter()))
                run("zip -9lr {html} {html-nozip}".format(**tarball_formatter()))
                run("cp {html} ../../dist/".format(**tarball_formatter()))
            run("make clean")
            run("source ../docs-virtualenv/bin/activate; make latex")
            with cd("_build"):
                with cd("latex"):
                    run("make")
                    run("cp {pdf-orig} ../../../dist/{pdf}".format(**tarball_formatter()))

def copy_release_files():
    with cd("repos/sympy"):
        run("mkdir -p /vagrant/release")
        run("cp dist/* /vagrant/release/")
        run("cp py3k-sympy/dist/* /vagrant/release/")

def show_files(file):
    """
    Show the contents of a tarball.

    The current options for file are

    2: The Python 2 tarball
    3: The Python 3 tarball
    2win: The Python 2 Windows installer (Not yet implemented!)
    3win: The Python 3 Windows installer (Not yet implemented!)
    html: The html docs zip
    """
    # TODO:
    # - Automatically check that Python 3 has the same files as Python 2
    # - List the files that are in git but not in the release
    # - List the files in the Windows installers
    if file == '2':
        local("tar tf release/{2}".format(**tarball_formatter()))
    elif file == '3':
        py32 = "{3.2}".format(**tarball_formatter())
        py33 = "{3.3}".format(**tarball_formatter())
        assert md5(py32).split()[0] == md5(py33).split()[0]
        local("tar tf release/" + py32)
    elif file in {'2win', '3win'}:
        raise NotImplementedError("Windows installers")
    elif file == 'html':
        local("unzip -l release/{html}".format(**tarball_formatter()))
    else:
        raise ValueError(file + " is not valid")

def md5(file='*'):
    out = local("md5sum release/" + file, capture=True)
    print out
    return out

def get_tarball_name(file):
    """
    Get the name of a tarball

    file should be one of

    source-orig: The original name of the source tarball
    2:           The Python 2 tarball (after renaming)
    3.2:         The Python 3.2 tarball (after renaming)
    3.3:         The Python 3.3 tarball (after renaming)
    2win32-orig: The original name of the Python 2 win32 installer
    2win32:      The name of the Python 2 win32 installer (after renaming)
    html:        The name of the html zip
    html-nozip:  The name of the html, without ".zip"
    pdf-orig:    The original name of the pdf file
    pdf:         The name of the pdf file (after renaming)
    """
    version = get_sympy_version()
    doctypename = defaultdict(str, {'html': 'zip', 'pdf': 'pdf'})
    winos = defaultdict(str, {'2win32': 'win32', '2win32-orig': 'linux-i686'})
    if file in {'source-orig', '2'}:
        name = 'sympy-{version}.tar.gz'
    elif file in {'3.2', '3.3'}:
        name = "sympy-{version}-py{pyversion}.tar.gz"
    elif file in {'2win32', '2win32-orig'}:
        name = "sympy-{version}.{wintype}.exe"
    elif file in {'html', 'pdf', 'html-nozip'}:
        name = "sympy-docs-{type}-{version}"
        if not file.endswith('nozip'):
            file += ".{extension}"
    elif file == 'pdf-orig':
        name = "sympy-{version}.pdf"
    else:
        raise ValueError(file + " is not a recognized argument")
    ret = name.format(version=version, pyversion=file, type=file,
        extension=doctypename[file], wintype=winos[file])
    print ret # REMOVE ME
    return ret

tarball_name_types = {
    'source-orig:',
    '2:',
    '3.2:',
    '3.3:',
    '2win32-orig:',
    '2win32:',
    'html:',
    'pdf-orig:',
    'pdf',
    }

# This has to be a function, because you cannot call any function here at
# import time (before the vagrant() function is fun).
def tarball_formatter():
    return {name: get_tarball_name(name) for name in tarball_name_types}

# ------------------------------------------------
# Vagrant related configuration

def vagrant():
    vc = _get_vagrant_config()
    # change from the default user to 'vagrant'
    env.user = vc['User']
    # connect to the port-forwarded ssh
    env.hosts = ['%s:%s' % (vc['HostName'], vc['Port'])]
    # use vagrant ssh key
    env.key_filename = vc['IdentityFile'].strip('"')
    # Forward the agent if specified:
    env.forward_agent = vc.get('ForwardAgent', 'no') == 'yes'

def _get_vagrant_config():
    """
    Parses vagrant configuration and returns it as dict of ssh parameters
    and their values
    """
    result = local('vagrant ssh-config', capture=True)
    conf = {}
    for line in iter(result.splitlines()):
        parts = line.split()
        conf[parts[0]] = ' '.join(parts[1:])
    return conf

def restart_network():
    """
    Do this if the VM won't connect to the internet.
    """
    run("sudo /etc/init.d/networking restart")

# ---------------------------------------
# Just a simple testing command:

def uname():
    run('uname -a')
<|endoftext|>


<commit_before>from collections import defaultdict

from fabric.api import env, local, run, sudo, cd, hide, prefix
from fabric.context_managers import shell_env, prefix
from fabric.operations import put, get
from fabric.contrib.files import append, exists
env.use_ssh_config = True

def prepare():
    prepare_apt()
    checkout_cache()

def prepare_apt():
    sudo("apt-get -qq update")
    sudo("apt-get -y install git python3 make python-virtualenv zip python-dev")
    # Needed to build the docs
    sudo("apt-get -y install graphviz inkscape texlive texlive-xetex texlive-fonts-recommended texlive-latex-extra")

def remove_userspace():
    """
    Deletes (!) the SymPy changes. Use with great care.
    """
    run("rm -rf repos")

def checkout_cache():
    run("rm -rf sympy-cache.git")
    run("git clone --bare https://github.com/sympy/sympy.git sympy-cache.git")

def gitrepos(branch=None):
    if not branch:
        # Use the current branch (of this git repo, not the one in Vagrant)
        branch = local("git rev-parse --abbrev-ref HEAD", capture=True)
    run("mkdir -p repos")
    with cd("repos"):
        run("git clone --reference ../sympy-cache.git https://github.com/sympy/sympy.git")
        if branch != "master":
            with cd("sympy"):
                run("git checkout -t origin/%s" % branch)

def get_sympy_version():
    with cd("repos/sympy"):
        version = run('python -c "import sympy;print sympy.__version__"')
    assert '\n' not in version
    assert ' ' not in version
    assert '\t' not in version
    return version

def test():
    with cd("repos/sympy"):
        run("./setup.py test")

def release(branch=None):
    remove_userspace()
    gitrepos(branch)
    python2_tarball()
    python3_tarball()
    build_docs()
    copy_release_files()

def python2_tarball():
    with cd("repos/sympy"):
        run("git clean -dfx")
        run("./setup.py clean")
        run("./setup.py sdist")
        run("./setup.py bdist_wininst")
        run("mv dist/{2win32-orig} dist/{2win32}".format(**tarball_formatter))

def python3_tarball():
    with cd("repos/sympy"):
        run("bin/use2to3")
        with cd("py3k-sympy"):
            run("./setup.py clean")
            run("./setup.py sdist")
            # We have to have 3.2 and 3.3 tarballs to make things work in
            # pip. See https://groups.google.com/d/msg/sympy/JEwi4ohGB90/FfjVDxZIkSEJ.
            run("mv dist/{source-orig} dist/{3.2}".format(**tarball_formatter))
            run("cp dist/{3.2} dist/{3.3}".format(**tarball_formatter))
            # We didn't test this yet:
            #run("./setup.py bdist_wininst")

def build_docs():
    with cd("repos/sympy"):
        run("mkdir -p dist")
        run("virtualenv docs-virtualenv")
        run("source docs-virtualenv/bin/activate; pip install sphinx==1.1.3 numpy")
        with cd("doc"):
            run("make clean")
            run("source ../docs-virtualenv/bin/activate; make html-errors")
            with cd("_build"):
                run("mv html {html-nozip}".format(**tarball_formatter))
                run("zip -9lr {html} {html-nozip}".format(**tarball_formatter))
                run("cp {html} ../../dist/".format(**tarball_formatter))
            run("make clean")
            run("source ../docs-virtualenv/bin/activate; make latex")
            with cd("_build"):
                with cd("latex"):
                    run("make")
                    run("cp {pdf-orig} ../../../dist/{pdf}".format(**tarball_formatter))

def copy_release_files():
    with cd("repos/sympy"):
        run("mkdir -p /vagrant/release")
        run("cp dist/* /vagrant/release/")
        run("cp py3k-sympy/dist/* /vagrant/release/")

def show_files(file):
    """
    Show the contents of a tarball.

    The current options for file are

    2: The Python 2 tarball
    3: The Python 3 tarball
    2win: The Python 2 Windows installer (Not yet implemented!)
    3win: The Python 3 Windows installer (Not yet implemented!)
    html: The html docs zip
    """
    # TODO:
    # - Automatically check that Python 3 has the same files as Python 2
    # - List the files that are in git but not in the release
    # - List the files in the Windows installers
    if file == '2':
        local("tar tf release/{2}".format(**tarball_formatter))
    elif file == '3':
        py32 = "{3.2}".format(**tarball_formatter)
        py33 = "{3.3}".format(**tarball_formatter)
        assert md5(py32).split()[0] == md5(py33).split()[0]
        local("tar tf release/" + py32)
    elif file in {'2win', '3win'}:
        raise NotImplementedError("Windows installers")
    elif file == 'html':
        local("unzip -l release/{html}".format(**tarball_formatter))
    else:
        raise ValueError(file + " is not valid")

def md5(file='*'):
    out = local("md5sum release/" + file, capture=True)
    print out
    return out

def get_tarball_name(file):
    """
    Get the name of a tarball

    file should be one of

    source-orig: The original name of the source tarball
    2:           The Python 2 tarball (after renaming)
    3.2:         The Python 3.2 tarball (after renaming)
    3.3:         The Python 3.3 tarball (after renaming)
    2win32-orig: The original name of the Python 2 win32 installer
    2win32:      The name of the Python 2 win32 installer (after renaming)
    html:        The name of the html zip
    html-nozip:  The name of the html, without ".zip"
    pdf-orig:    The original name of the pdf file
    pdf:         The name of the pdf file (after renaming)
    """
    version = get_sympy_version()
    doctypename = defaultdict(str, {'html': 'zip', 'pdf': 'pdf'})
    winos = defaultdict(str, {'2win32': 'win32', '2win32-orig': 'linux-i686'})
    if file in {'source-orig', '2'}:
        name = 'sympy-{version}.tar.gz'
    elif file in {'3.2', '3.3'}:
        name = "sympy-{version}-py{pyversion}.tar.gz"
    elif file in {'2win32', '2win32-orig'}:
        name = "sympy-{version}.{wintype}.exe"
    elif file in {'html', 'pdf', 'html-nozip'}:
        name = "sympy-docs-{type}-{version}"
        if not file.endswith('nozip'):
            file += ".{extension}"
    elif file == 'pdf-orig':
        name = "sympy-{version}.pdf"
    else:
        raise ValueError(file + " is not a recognized argument")
    ret = name.format(version=version, pyversion=file, type=file,
        extension=doctypename[file], wintype=winos[file])
    print ret # REMOVE ME
    return ret

tarball_name_types = {
    'source-orig:',
    '2:',
    '3.2:',
    '3.3:',
    '2win32-orig:',
    '2win32:',
    'html:',
    'pdf-orig:',
    'pdf',
    }

tarball_formatter = {name: get_tarball_name(name) for name in tarball_name_types}

# ------------------------------------------------
# Vagrant related configuration

def vagrant():
    vc = _get_vagrant_config()
    # change from the default user to 'vagrant'
    env.user = vc['User']
    # connect to the port-forwarded ssh
    env.hosts = ['%s:%s' % (vc['HostName'], vc['Port'])]
    # use vagrant ssh key
    env.key_filename = vc['IdentityFile'].strip('"')
    # Forward the agent if specified:
    env.forward_agent = vc.get('ForwardAgent', 'no') == 'yes'

def _get_vagrant_config():
    """
    Parses vagrant configuration and returns it as dict of ssh parameters
    and their values
    """
    result = local('vagrant ssh-config', capture=True)
    conf = {}
    for line in iter(result.splitlines()):
        parts = line.split()
        conf[parts[0]] = ' '.join(parts[1:])
    return conf

def restart_network():
    """
    Do this if the VM won't connect to the internet.
    """
    run("sudo /etc/init.d/networking restart")

# ---------------------------------------
# Just a simple testing command:

def uname():
    run('uname -a')
<commit_msg>Add a function to test if the tarball can be installed<commit_after>from collections import defaultdict

from fabric.api import env, local, run, sudo, cd, hide, prefix
from fabric.context_managers import shell_env, prefix
from fabric.operations import put, get
from fabric.contrib.files import append, exists
env.use_ssh_config = True

def prepare():
    prepare_apt()
    checkout_cache()

def prepare_apt():
    sudo("apt-get -qq update")
    sudo("apt-get -y install git python3 make python-virtualenv zip python-dev")
    # Needed to build the docs
    sudo("apt-get -y install graphviz inkscape texlive texlive-xetex texlive-fonts-recommended texlive-latex-extra")

def remove_userspace():
    """
    Deletes (!) the SymPy changes. Use with great care.
    """
    run("rm -rf repos")

def checkout_cache():
    run("rm -rf sympy-cache.git")
    run("git clone --bare https://github.com/sympy/sympy.git sympy-cache.git")

def gitrepos(branch=None):
    if not branch:
        # Use the current branch (of this git repo, not the one in Vagrant)
        branch = local("git rev-parse --abbrev-ref HEAD", capture=True)
    run("mkdir -p repos")
    with cd("repos"):
        run("git clone --reference ../sympy-cache.git https://github.com/sympy/sympy.git")
        if branch != "master":
            with cd("sympy"):
                run("git checkout -t origin/%s" % branch)

def get_sympy_version():
    with cd("repos/sympy"):
        version = run('python -c "import sympy;print sympy.__version__"')
    assert '\n' not in version
    assert ' ' not in version
    assert '\t' not in version
    return version

def test_git():
    with cd("repos/sympy"):
        run("./setup.py test")

def test_tarball(release='2'):
    if release not in {'2', '3'}: # TODO: Add win32
        raise ValueError("release must be one of '2', '3', not %s" % release)
    run("virtualenv test-{release}-virtualenv".format(release=release))
    run("source test-{release}-virtualenv/bin/activate".format(release=release))
    if release == '2':
        run("cp /vagrant/release/{2} releasetar.tar".format(**tarball_formatter))
    if release == '3':
        run("cp /vagrant/release/{2} releasetar.tar".format(**tarball_formatter))
    run("tar xvf releasetar.tar")
    with cd("releastar"):
        run("python setup.py install")
        run('python -c "import sympy; print sympy.__version__"')

def release(branch=None):
    remove_userspace()
    gitrepos(branch)
    python2_tarball()
    python3_tarball()
    build_docs()
    copy_release_files()
    test_tarball('2')
    test_tarball('3')

def python2_tarball():
    with cd("repos/sympy"):
        run("git clean -dfx")
        run("./setup.py clean")
        run("./setup.py sdist")
        run("./setup.py bdist_wininst")
        run("mv dist/{2win32-orig} dist/{2win32}".format(**tarball_formatter))

def python3_tarball():
    with cd("repos/sympy"):
        run("bin/use2to3")
        with cd("py3k-sympy"):
            run("./setup.py clean")
            run("./setup.py sdist")
            # We have to have 3.2 and 3.3 tarballs to make things work in
            # pip. See https://groups.google.com/d/msg/sympy/JEwi4ohGB90/FfjVDxZIkSEJ.
            run("mv dist/{source-orig} dist/{3.2}".format(**tarball_formatter))
            run("cp dist/{3.2} dist/{3.3}".format(**tarball_formatter))
            # We didn't test this yet:
            #run("./setup.py bdist_wininst")

def build_docs():
    with cd("repos/sympy"):
        run("mkdir -p dist")
        run("virtualenv docs-virtualenv")
        run("source docs-virtualenv/bin/activate; pip install sphinx==1.1.3 numpy")
        with cd("doc"):
            run("make clean")
            run("source ../docs-virtualenv/bin/activate; make html-errors")
            with cd("_build"):
                run("mv html {html-nozip}".format(**tarball_formatter))
                run("zip -9lr {html} {html-nozip}".format(**tarball_formatter))
                run("cp {html} ../../dist/".format(**tarball_formatter))
            run("make clean")
            run("source ../docs-virtualenv/bin/activate; make latex")
            with cd("_build"):
                with cd("latex"):
                    run("make")
                    run("cp {pdf-orig} ../../../dist/{pdf}".format(**tarball_formatter))

def copy_release_files():
    with cd("repos/sympy"):
        run("mkdir -p /vagrant/release")
        run("cp dist/* /vagrant/release/")
        run("cp py3k-sympy/dist/* /vagrant/release/")

def show_files(file):
    """
    Show the contents of a tarball.

    The current options for file are

    2: The Python 2 tarball
    3: The Python 3 tarball
    2win: The Python 2 Windows installer (Not yet implemented!)
    3win: The Python 3 Windows installer (Not yet implemented!)
    html: The html docs zip
    """
    # TODO:
    # - Automatically check that Python 3 has the same files as Python 2
    # - List the files that are in git but not in the release
    # - List the files in the Windows installers
    if file == '2':
        local("tar tf release/{2}".format(**tarball_formatter))
    elif file == '3':
        py32 = "{3.2}".format(**tarball_formatter)
        py33 = "{3.3}".format(**tarball_formatter)
        assert md5(py32).split()[0] == md5(py33).split()[0]
        local("tar tf release/" + py32)
    elif file in {'2win', '3win'}:
        raise NotImplementedError("Windows installers")
    elif file == 'html':
        local("unzip -l release/{html}".format(**tarball_formatter))
    else:
        raise ValueError(file + " is not valid")

def md5(file='*'):
    out = local("md5sum release/" + file, capture=True)
    print out
    return out

def get_tarball_name(file):
    """
    Get the name of a tarball

    file should be one of

    source-orig: The original name of the source tarball
    2:           The Python 2 tarball (after renaming)
    3.2:         The Python 3.2 tarball (after renaming)
    3.3:         The Python 3.3 tarball (after renaming)
    2win32-orig: The original name of the Python 2 win32 installer
    2win32:      The name of the Python 2 win32 installer (after renaming)
    html:        The name of the html zip
    html-nozip:  The name of the html, without ".zip"
    pdf-orig:    The original name of the pdf file
    pdf:         The name of the pdf file (after renaming)
    """
    version = get_sympy_version()
    doctypename = defaultdict(str, {'html': 'zip', 'pdf': 'pdf'})
    winos = defaultdict(str, {'2win32': 'win32', '2win32-orig': 'linux-i686'})
    if file in {'source-orig', '2'}:
        name = 'sympy-{version}.tar.gz'
    elif file in {'3.2', '3.3'}:
        name = "sympy-{version}-py{pyversion}.tar.gz"
    elif file in {'2win32', '2win32-orig'}:
        name = "sympy-{version}.{wintype}.exe"
    elif file in {'html', 'pdf', 'html-nozip'}:
        name = "sympy-docs-{type}-{version}"
        if not file.endswith('nozip'):
            file += ".{extension}"
    elif file == 'pdf-orig':
        name = "sympy-{version}.pdf"
    else:
        raise ValueError(file + " is not a recognized argument")
    ret = name.format(version=version, pyversion=file, type=file,
        extension=doctypename[file], wintype=winos[file])
    print ret # REMOVE ME
    return ret

tarball_name_types = {
    'source-orig:',
    '2:',
    '3.2:',
    '3.3:',
    '2win32-orig:',
    '2win32:',
    'html:',
    'pdf-orig:',
    'pdf',
    }

tarball_formatter = {name: get_tarball_name(name) for name in tarball_name_types}

# ------------------------------------------------
# Vagrant related configuration

def vagrant():
    vc = _get_vagrant_config()
    # change from the default user to 'vagrant'
    env.user = vc['User']
    # connect to the port-forwarded ssh
    env.hosts = ['%s:%s' % (vc['HostName'], vc['Port'])]
    # use vagrant ssh key
    env.key_filename = vc['IdentityFile'].strip('"')
    # Forward the agent if specified:
    env.forward_agent = vc.get('ForwardAgent', 'no') == 'yes'

def _get_vagrant_config():
    """
    Parses vagrant configuration and returns it as dict of ssh parameters
    and their values
    """
    result = local('vagrant ssh-config', capture=True)
    conf = {}
    for line in iter(result.splitlines()):
        parts = line.split()
        conf[parts[0]] = ' '.join(parts[1:])
    return conf

def restart_network():
    """
    Do this if the VM won't connect to the internet.
    """
    run("sudo /etc/init.d/networking restart")

# ---------------------------------------
# Just a simple testing command:

def uname():
    run('uname -a')
<|endoftext|>


<commit_before>version = "0.1.0"
__version__ = version
<commit_msg>Set Repour (devel) version to be 0.11.0<commit_after>version = "0.11.0"
__version__ = version
<|endoftext|>


<commit_before>from model import Reservation
from mongoengine import *
from docopt import docopt
import pprint
import json
from cloudmesh_common.tables import array_dict_table_printer

def reservation_connect():
    """establishes connection to the reservation server"""
    try:
        db = connect('reservation', port=27777)
        return db
    except Exception, e:
        print "ERROR: could not establish a connection to mongo db"
        print e
        return None

def _print_reservations(reservations, format="table"):
    """prints a reservation"""
    reservation_list = []
    for reservation in reservations:
        reservation_list.append(reservation.to_json())
    if format == "json":
        print json.dumps(reservation_list, indent=4)
    if format == "table":
        print array_dict_table_printer(reservation_list,
                                       ["cm_id",
                                        "label",
                                        "project",
                                        "user",
                                        "host",
                                        "start_time",
                                        "end_time",
                                        "summary",
                                        ]
                                       )

def shell_command_reservation(arguments):
    """
      ::

        Usage:
            reservation --rst
            reservation --version
            reservation find [all]
                             [--user=USER_ID]
                             [--label=ID]
                             [--cm_id=ID]
            reservation list [--cm_id=CM_ID]
                             [--user=USER_ID]
                             [--project=PROJECT_ID]
                             [--label=STRING]
                             [--start=TIME_START]
                             [--end=TIME_END]
                             [--host=HOST]
                             [--summary=SUMMARY]
                             [--format=FORMAT]
            reservation duration [--cm_id=CM_ID]
            reservation delete [all]
                               [--cm_id=CM_ID]
                               [--user=USER_ID]
                               [--project=PROJECT_ID]
                               [--label=STRING]
                               [--start=TIME_START]
                               [--end=TIME_END]
                               [--host=HOST]
            reservation update [--cm_id=CM_ID]
                               [--user=USER_ID]
                               [--project=PROJECT_ID]
                               [--label=STRING]
                               [--start=TIME_START]
                               [--end=TIME_END]
                               [--host=HOST]
                               [--summary=SUMMARY]
                               [--cm_id=CM_ID]
            reservation add --cm_id=CM_ID --user=USER_ID --project=PROJECT_ID --label=STRING --start=TIME_START --end=TIME_END --host=HOST --summary=SUMMARY
            reservation add --file=FILE

        Options:
            --rst                 print an rst manul page
            --version             print the version
            --label=STRING        label id reservation
            --cm_id=CM_ID         reservation cloudmesh id
            --user=USER_ID        user id
            --project=PROJECT_ID  project id
            --start=TIME_START    Start time of the reservation, in
                                  YYYY/MM/DD HH:MM:SS format. [default: 1901-01-01]
            --end=TIME_END        End time of the reservation, in
                                  YYYY/MM/DD HH:MM:SS format. In addition a duration
                                  can be specified if the + sign is the first sign.
                                  The duration will than be added to
                                  the start time. [default: 2100-12-31]
            --host=HOST           host number
            --summary=SUMMARY     summary of the reservation
            --file=FILE           Adding multiple reservations from one file
            --format=FORMAT       Format is either table or jaon
                                  [default: table]
    """

    #print arguments["login"]
    if arguments["--rst"]:

        print 70 * "*"
        print "Manual Pages"
        print 70 * "*"
        print
        print "reservation"
        print 70 * "="
        print "\n::\n"
        lines = shell_command_reservation.__doc__.split("\n")
        for line in lines:
            print "  ", line

    elif arguments["--version"]:
        print reservation.__version__

    elif(arguments["list"]):
        reservations = Reservation()
        rsv= reservations.list(cm_id=arguments["--cm_id"],
                               user=arguments["--user"],
                               project=arguments["--project"],
                               label= arguments["--label"],
                               start_time= arguments["--start"],
                               end_time=arguments["--end"],
                               host=arguments["--host"],
                               summary=arguments["--summary"])
        _print_reservations(rsv, arguments["--format"])

    elif(arguments["find"]):
        reservations = Reservation()
        if(arguments["all"]):
            _print_reservations(reservations.find_all())
        elif(arguments["--user"]):
            _print_reservations(reservations.find_user(arguments["--user"]))
        elif(arguments["--label"]):
            _print_reservations(reservations.find_label(arguments["--label"]))
        elif(arguments["--cm_id"]):
            _print_reservations(reservations.find_id(arguments["--cm_id"]))
        else:
            _print_reservations(reservations.find_all())
    elif(arguments["duration"]):
        reservations = Reservation()
        print reservations.duration(arguments["--cm_id"])
    elif(arguments["delete"]):
        if(arguments["all"]):
            reservations = Reservation()
            reservations.delete_all()
        else:
            reservations = Reservation()
            reservations.delete_selection(cm_id=arguments["--cm_id"],
                                          user=arguments["--user"],
                                          project=arguments["--project"],
                                          label= arguments["--label"],
                                          start_time= arguments["--start"],
                                          end_time=arguments["--end"],
                                          host=arguments["--host"])
    elif(arguments["add"]):

        if arguments["--file"] is None:

            reservations = Reservation(label=arguments["--label"],
                                    user=arguments["--user"],
                                    project=arguments["--project"],
                                    start_time=arguments["--start"],
                                    end_time=arguments["--end"],
                                    cm_id=arguments["--cm_id"],
                                    host=arguments["--host"],
                                    summary=arguments["--summary"])
            reservations.add()
        else:
            try:
                with open(os.path.join(sys.path[0], arguments["--file"])) as file:
                    reader = csv.reader(file)
                    for row in reader:
                        reservations = Reservation(cm_id=row[0],
                                                label=row[1],
                                                user=row[2],
                                                project=row[3],
                                                start_time=row[4],
                                                end_time=row[5],
                                                host=row[6],
                                                summary=row[7])
                        reservations.add()
            except Exception as e:
                print "Error in adding from file. ", e
'''
    elif(arguments["update"]):
          reservations = Reservation()
          fromObj = [str(sys.argv[2]).split("=")[0].replace("--", ""),str(sys.argv[2]).split("=")[1]]
          toObj = [str(sys.argv[3]).split("=")[0].replace("--", ""),str(sys.argv[3]).split("=")[1]]
          fromBody = {cm_id=101, project = 20, user = "oliver"}
          toBody = {cm_id=101, project = 20, user = "oliver"}
          reservations.update_selection(cm_id=fromObj[1],project=toObj[1])
          print reservations.find_all()
'''

def main():
    arguments = docopt(shell_command_reservation.__doc__)
    db = reservation_connect()
    shell_command_reservation(arguments)

if __name__ == "__main__":
    main()

<commit_msg>fixig the find all cli command with a proper print function<commit_after>from model import Reservation
from mongoengine import *
from docopt import docopt
import pprint
import json
from cloudmesh_common.tables import array_dict_table_printer

def reservation_connect():
    """establishes connection to the reservation server"""
    try:
        db = connect('reservation', port=27777)
        return db
    except Exception, e:
        print "ERROR: could not establish a connection to mongo db"
        print e
        return None

def _print_reservations(reservations, format="table"):
    """prints a reservation"""
    reservation_list = []
    for reservation in reservations:
        reservation_list.append(reservation.to_json())
    if format == "json":
        print json.dumps(reservation_list, indent=4)
    if format == "table":
        print array_dict_table_printer(reservation_list,
                                       ["cm_id",
                                        "label",
                                        "project",
                                        "user",
                                        "host",
                                        "start_time",
                                        "end_time",
                                        "summary",
                                        ]
                                       )

def shell_command_reservation(arguments):
    """
      ::

        Usage:
            reservation --rst
            reservation --version
            reservation find [all]
                             [--user=USER_ID]
                             [--label=ID]
                             [--cm_id=ID]
                             [--format=FORMAT]
            reservation list [--cm_id=CM_ID]
                             [--user=USER_ID]
                             [--project=PROJECT_ID]
                             [--label=STRING]
                             [--start=TIME_START]
                             [--end=TIME_END]
                             [--host=HOST]
                             [--summary=SUMMARY]
                             [--format=FORMAT]
            reservation duration [--cm_id=CM_ID]
            reservation delete [all]
                               [--cm_id=CM_ID]
                               [--user=USER_ID]
                               [--project=PROJECT_ID]
                               [--label=STRING]
                               [--start=TIME_START]
                               [--end=TIME_END]
                               [--host=HOST]
            reservation update [--cm_id=CM_ID]
                               [--user=USER_ID]
                               [--project=PROJECT_ID]
                               [--label=STRING]
                               [--start=TIME_START]
                               [--end=TIME_END]
                               [--host=HOST]
                               [--summary=SUMMARY]
                               [--cm_id=CM_ID]
            reservation add --cm_id=CM_ID --user=USER_ID --project=PROJECT_ID --label=STRING --start=TIME_START --end=TIME_END --host=HOST --summary=SUMMARY
            reservation add --file=FILE

        Options:
            --rst                 print an rst manul page
            --version             print the version
            --label=STRING        label id reservation
            --cm_id=CM_ID         reservation cloudmesh id
            --user=USER_ID        user id
            --project=PROJECT_ID  project id
            --start=TIME_START    Start time of the reservation, in
                                  YYYY/MM/DD HH:MM:SS format. [default: 1901-01-01]
            --end=TIME_END        End time of the reservation, in
                                  YYYY/MM/DD HH:MM:SS format. In addition a duration
                                  can be specified if the + sign is the first sign.
                                  The duration will than be added to
                                  the start time. [default: 2100-12-31]
            --host=HOST           host number
            --summary=SUMMARY     summary of the reservation
            --file=FILE           Adding multiple reservations from one file
            --format=FORMAT       Format is either table or jaon
                                  [default: table]
    """

    #print arguments["login"]
    if arguments["--rst"]:

        print 70 * "*"
        print "Manual Pages"
        print 70 * "*"
        print
        print "reservation"
        print 70 * "="
        print "\n::\n"
        lines = shell_command_reservation.__doc__.split("\n")
        for line in lines:
            print "  ", line

    elif arguments["--version"]:
        print reservation.__version__

    elif(arguments["list"]):
        db = Reservation()
        reservations = db.list(cm_id=arguments["--cm_id"],
                               user=arguments["--user"],
                               project=arguments["--project"],
                               label= arguments["--label"],
                               start_time= arguments["--start"],
                               end_time=arguments["--end"],
                               host=arguments["--host"],
                               summary=arguments["--summary"])
        _print_reservations(reservations, arguments["--format"])

    elif(arguments["find"]):
        db = Reservation()

        if(arguments["all"]):
            reservations = db.find_all()
        elif(arguments["--user"]):
            reservations = db.find_user(arguments["--user"])
        elif(arguments["--label"]):
            reservations = db.find_label(arguments["--label"])
        elif(arguments["--cm_id"]):
            reservations = db.find_id(arguments["--cm_id"])
        else:
            reservations = db.find_all()

        _print_reservations(reservations, arguments["--format"])

    elif(arguments["duration"]):
        reservations = Reservation()
        print reservations.duration(arguments["--cm_id"])
    elif(arguments["delete"]):
        if(arguments["all"]):
            reservations = Reservation()
            reservations.delete_all()
        else:
            reservations = Reservation()
            reservations.delete_selection(cm_id=arguments["--cm_id"],
                                          user=arguments["--user"],
                                          project=arguments["--project"],
                                          label= arguments["--label"],
                                          start_time= arguments["--start"],
                                          end_time=arguments["--end"],
                                          host=arguments["--host"])
    elif(arguments["add"]):

        if arguments["--file"] is None:

            reservations = Reservation(label=arguments["--label"],
                                    user=arguments["--user"],
                                    project=arguments["--project"],
                                    start_time=arguments["--start"],
                                    end_time=arguments["--end"],
                                    cm_id=arguments["--cm_id"],
                                    host=arguments["--host"],
                                    summary=arguments["--summary"])
            reservations.add()
        else:
            try:
                with open(os.path.join(sys.path[0], arguments["--file"])) as file:
                    reader = csv.reader(file)
                    for row in reader:
                        reservations = Reservation(cm_id=row[0],
                                                label=row[1],
                                                user=row[2],
                                                project=row[3],
                                                start_time=row[4],
                                                end_time=row[5],
                                                host=row[6],
                                                summary=row[7])
                        reservations.add()
            except Exception as e:
                print "Error in adding from file. ", e
'''
    elif(arguments["update"]):
          reservations = Reservation()
          fromObj = [str(sys.argv[2]).split("=")[0].replace("--", ""),str(sys.argv[2]).split("=")[1]]
          toObj = [str(sys.argv[3]).split("=")[0].replace("--", ""),str(sys.argv[3]).split("=")[1]]
          fromBody = {cm_id=101, project = 20, user = "oliver"}
          toBody = {cm_id=101, project = 20, user = "oliver"}
          reservations.update_selection(cm_id=fromObj[1],project=toObj[1])
          print reservations.find_all()
'''

def main():
    arguments = docopt(shell_command_reservation.__doc__)
    db = reservation_connect()
    shell_command_reservation(arguments)

if __name__ == "__main__":
    main()

<|endoftext|>


<commit_before>from __future__ import print_function, division, absolute_import, unicode_literals
from builtins import bytes, dict, object, range, map, input, str
from future.utils import itervalues, viewitems, iteritems, listvalues, listitems

import numpy as np
from numba import jit
from rfpipe import util

import logging
logger = logging.getLogger(__name__)

def flag_data(st, data):
    """ Identifies bad data and flags it to 0.
    Converts to masked array for flagging, but returns zeroed numpy array.
    """

    datam = np.ma.masked_values(data, 0j, copy=False, shrink=False)

    spwchans = st.spw_chan_select
    for flagparams in st.prefs.flaglist:
        if len(flagparams) == 3:
            mode, arg0, arg1 = flagparams
        else:
            mode, arg0 = flagparams

        if mode == 'blstd':
            flag_blstd(datam, arg0, arg1)
        elif mode == 'badchtslide':
            flag_badchtslide(datam, spwchans, arg0, arg1)
        elif mode == 'badspw':
            flag_badspw(datam, spwchans, arg0)
        else:
            logger.warning("Flaging mode {0} not available.".format(mode))

    return datam.filled(0)

def flag_blstd(data, sigma, convergence):
    """ Use data (4d) to calculate (int, chan, pol) to be flagged.
    Masked arrays assumed as input.
    """

    sh = data.shape

    blstd = util.blstd(data.data, data.mask)  # uses mask
    blstd = np.ma.masked_equal(blstd, 0)
#    blstd = np.ma.std(data, axis=1)

    # iterate to good median and std values
    blstdmednew = np.ma.median(blstd)
    blstdstdnew = np.ma.std(blstd)
    blstdstd = blstdstdnew*2  # TODO: is this initialization used?
    while (blstdstd-blstdstdnew)/blstdstd > convergence:
        blstdstd = blstdstdnew
        blstdmed = blstdmednew
        blstd = np.ma.masked_where(blstd > blstdmed + sigma*blstdstd, blstd, copy=False)
        blstdmednew = np.ma.median(blstd)
        blstdstdnew = np.ma.std(blstd)

    # flag blstd too high
    badt, badch, badpol = np.where(blstd > blstdmednew + sigma*blstdstdnew)
    logger.info("flagged by blstd: {0} of {1} total channel/time/pol cells."
                .format(len(badt), sh[0]*sh[2]*sh[3]))

    for i in range(len(badt)):
        data.mask[badt[i], :, badch[i], badpol[i]] = True

def flag_badchtslide(data, spwchans, sigma, win):
    """ Use data (4d) to calculate (int, chan, pol) to be flagged
    """

    sh = data.shape

    meanamp = np.abs(data).mean(axis=1)

    # calc badch as deviation from median of window
    spec = meanamp.mean(axis=0)
#    specmed = slidedev(spec, win)
    specmed = np.concatenate([spec[chans] - np.ma.median(spec[chans]) for chans in spwchans])
    badch = np.where(specmed > sigma*np.ma.std(specmed, axis=0))

    # calc badt as deviation from median of window
    lc = meanamp.mean(axis=1)
    lcmed = slidedev(lc, win)
    badt = np.where(lcmed > sigma*np.ma.std(lcmed, axis=0))

    badtcnt = len(np.ma.unique(badt))
    badchcnt = len(np.ma.unique(badch))
    logger.info("flagged by badchtslide: {0}/{1} pol-times and {2}/{3} pol-chans."
                .format(badtcnt, sh[0]*sh[3], badchcnt, sh[2]*sh[3]))

    for i in range(len(badch[0])):
        data.mask[:, :, badch[0][i], badch[1][i]] = True

    for i in range(len(badt[0])):
        data.mask[badt[0][i], :, :, badt[1][i]] = True

def flag_badspw(data, spwchans, sigma):
    """ Use data median variance between spw to flag spw
    Best to use this after flagging bad channels.
    """

    nspw = len(spwchans)

    if nspw >= 4:
        # calc badspw
        spec = np.abs(data).mean(axis=3).mean(axis=1).mean(axis=0)
        deviations = []
        for chans in spwchans:
            if spec[chans].count() > 3:
                deviations.append(np.ma.std(spec[chans]))
            else:
                deviations.append(0)
        deviations = np.ma.masked_equal(np.nan_to_num(deviations), 0)
        logger.info("badspw flagging finds deviations per spw: {0}"
                    .format(deviations))

        badspw = []
        badspwnew = np.where(deviations > sigma*np.ma.median(deviations))[0]
        while len(badspwnew) > len(badspw):
            badspw = badspwnew
            goodspw = [spw for spw in range(nspw) if spw not in badspw]
            badspwnew = np.where(deviations > sigma*np.ma.median(deviations.take(goodspw)))[0]

        badspw = np.concatenate((badspw, np.where(deviations.mask)[0])).astype(int)

        logger.info("flagged {0}/{1} spw ({2})"
                    .format(len(badspw), nspw, badspw))

        for i in badspw:
            data.mask[:, :, spwchans[i], :] = True

    else:
        logger.warning("Fewer than 4 spw. Not performing badspw detetion.")

def slidedev(arr, win):
    """ Given a (len x 2) array, calculate the deviation from the median per pol.
    Calculates median over a window, win.
    """

    med = np.zeros_like(arr)
    for i in range(len(arr)):
        inds = list(range(max(0, i-win//2), i)) + list(range(i+1, min(i+win//2, len(arr))))
        med[i] = np.ma.median(arr.take(inds, axis=0), axis=0)

    return arr-med

def getonlineflags(st, segment):
    """ Gets antenna flags for a given segment from either sdm or mcaf server.
    Returns an array of flags (1: good, 0: bad) for each baseline.
    """

    t0, t1 = st.segmenttimes[segment]
    if st.metadata.datasource == 'sdm':
        sdm = util.getsdm(st.metadata.filename, bdfdir=st.metadata.bdfdir)
        scan = sdm.scan(st.metadata.scan)
        takebls = [st.metadata.blarr_orig.tolist().index(list(bl)) for bl in st.blarr]
        flags = scan.flags([t0, t1])[:, takebls].all(axis=0)
    elif st.metadata.datasource == 'vys':
        try:
            from realfast.mcaf_servers import getblflags
            flags = getblflags(st.metadata.datasetId, st.blarr,
                               startTime=t0, endTime=t1)
        except (ImportError, Exception):
            logger.warning("No mcaf antenna flag server flags available")
            flags = np.ones(st.nbl)

    if not flags.all():
        logger.info('Found antennas to flag in time range {0}-{1} '
                    .format(t0, t1))
    else:
        logger.info('No flagged antennas in time range {0}-{1} '
                    .format(t0, t1))

    return flags

def flag_data_rtpipe(st, data):
    """ Flagging data in single process
    Deprecated.
    """
    try:
        import rtlib_cython as rtlib
    except ImportError:
        logger.error("rtpipe not installed. Cannot import rtlib for flagging.")

    # **hack!**
    d = {'dataformat': 'sdm', 'ants': [int(ant.lstrip('ea')) for ant in st.ants], 'excludeants': st.prefs.excludeants, 'nants': len(st.ants)}

    for flag in st.prefs.flaglist:
        mode, sig, conv = flag
        for spw in st.spw:
            chans = np.arange(st.metadata.spw_nchan[spw]*spw, st.metadata.spw_nchan[spw]*(1+spw))
            for pol in range(st.npol):
                status = rtlib.dataflag(data, chans, pol, d, sig, mode, conv)
                logger.info(status)

    # hack to get rid of bad spw/pol combos whacked by rfi
    if st.prefs.badspwpol:
        logger.info('Comparing overall power between spw/pol. Removing those with {0} times typical value'.format(st.prefs.badspwpol))
        spwpol = {}
        for spw in st.spw:
            chans = np.arange(st.metadata.spw_nchan[spw]*spw, st.metadata.spw_nchan[spw]*(1+spw))
            for pol in range(st.npol):
                spwpol[(spw, pol)] = np.abs(data[:, :, chans, pol]).std()

        meanstd = np.mean(list(spwpol.values()))
        for (spw, pol) in spwpol:
            if spwpol[(spw, pol)] > st.prefs.badspwpol*meanstd:
                logger.info('Flagging all of (spw %d, pol %d) for excess noise.' % (spw, pol))
                chans = np.arange(st.metadata.spw_nchan[spw]*spw, st.metadata.spw_nchan[spw]*(1+spw))
                data[:, :, chans, pol] = 0j

    return data
<commit_msg>changing calc for bad spw spectrum input for mem/compute efficiency<commit_after>from __future__ import print_function, division, absolute_import, unicode_literals
from builtins import bytes, dict, object, range, map, input, str
from future.utils import itervalues, viewitems, iteritems, listvalues, listitems

import numpy as np
from numba import jit
from rfpipe import util

import logging
logger = logging.getLogger(__name__)

def flag_data(st, data):
    """ Identifies bad data and flags it to 0.
    Converts to masked array for flagging, but returns zeroed numpy array.
    """

    datam = np.ma.masked_values(data, 0j, copy=False, shrink=False)

    spwchans = st.spw_chan_select
    for flagparams in st.prefs.flaglist:
        if len(flagparams) == 3:
            mode, arg0, arg1 = flagparams
        else:
            mode, arg0 = flagparams

        if mode == 'blstd':
            flag_blstd(datam, arg0, arg1)
        elif mode == 'badchtslide':
            flag_badchtslide(datam, spwchans, arg0, arg1)
        elif mode == 'badspw':
            flag_badspw(datam, spwchans, arg0)
        else:
            logger.warning("Flaging mode {0} not available.".format(mode))

    return datam.filled(0)

def flag_blstd(data, sigma, convergence):
    """ Use data (4d) to calculate (int, chan, pol) to be flagged.
    Masked arrays assumed as input.
    """

    sh = data.shape

    blstd = util.blstd(data.data, data.mask)  # uses mask
    blstd = np.ma.masked_equal(blstd, 0)
#    blstd = np.ma.std(data, axis=1)

    # iterate to good median and std values
    blstdmednew = np.ma.median(blstd)
    blstdstdnew = np.ma.std(blstd)
    blstdstd = blstdstdnew*2  # TODO: is this initialization used?
    while (blstdstd-blstdstdnew)/blstdstd > convergence:
        blstdstd = blstdstdnew
        blstdmed = blstdmednew
        blstd = np.ma.masked_where(blstd > blstdmed + sigma*blstdstd, blstd, copy=False)
        blstdmednew = np.ma.median(blstd)
        blstdstdnew = np.ma.std(blstd)

    # flag blstd too high
    badt, badch, badpol = np.where(blstd > blstdmednew + sigma*blstdstdnew)
    logger.info("flagged by blstd: {0} of {1} total channel/time/pol cells."
                .format(len(badt), sh[0]*sh[2]*sh[3]))

    for i in range(len(badt)):
        data.mask[badt[i], :, badch[i], badpol[i]] = True

def flag_badchtslide(data, spwchans, sigma, win):
    """ Use data (4d) to calculate (int, chan, pol) to be flagged
    """

    sh = data.shape

    meanamp = np.abs(data).mean(axis=1)

    # calc badch as deviation from median of window
    spec = meanamp.mean(axis=0)
#    specmed = slidedev(spec, win)
    specmed = np.concatenate([spec[chans] - np.ma.median(spec[chans]) for chans in spwchans])
    badch = np.where(specmed > sigma*np.ma.std(specmed, axis=0))

    # calc badt as deviation from median of window
    lc = meanamp.mean(axis=1)
    lcmed = slidedev(lc, win)
    badt = np.where(lcmed > sigma*np.ma.std(lcmed, axis=0))

    badtcnt = len(np.ma.unique(badt))
    badchcnt = len(np.ma.unique(badch))
    logger.info("flagged by badchtslide: {0}/{1} pol-times and {2}/{3} pol-chans."
                .format(badtcnt, sh[0]*sh[3], badchcnt, sh[2]*sh[3]))

    for i in range(len(badch[0])):
        data.mask[:, :, badch[0][i], badch[1][i]] = True

    for i in range(len(badt[0])):
        data.mask[badt[0][i], :, :, badt[1][i]] = True

def flag_badspw(data, spwchans, sigma):
    """ Use data median variance between spw to flag spw
    Best to use this after flagging bad channels.
    """

    nspw = len(spwchans)

    if nspw >= 4:
        # calc badspw
#        spec = np.abs(data).mean(axis=3).mean(axis=1).mean(axis=0)  # memory intensive and slow
        spec = np.abs(datam.mean(axis=0).mean(axis=2)).mean(axis=0)  # probably better
        deviations = []
        for chans in spwchans:
            if spec[chans].count() > 3:
                deviations.append(np.ma.std(spec[chans]))
            else:
                deviations.append(0)
        deviations = np.ma.masked_equal(np.nan_to_num(deviations), 0)
        logger.info("badspw flagging finds deviations per spw: {0}"
                    .format(deviations))

        badspw = []
        badspwnew = np.where(deviations > sigma*np.ma.median(deviations))[0]
        while len(badspwnew) > len(badspw):
            badspw = badspwnew
            goodspw = [spw for spw in range(nspw) if spw not in badspw]
            badspwnew = np.where(deviations > sigma*np.ma.median(deviations.take(goodspw)))[0]

        badspw = np.concatenate((badspw, np.where(deviations.mask)[0])).astype(int)

        logger.info("flagged {0}/{1} spw ({2})"
                    .format(len(badspw), nspw, badspw))

        for i in badspw:
            data.mask[:, :, spwchans[i], :] = True

    else:
        logger.warning("Fewer than 4 spw. Not performing badspw detetion.")

def slidedev(arr, win):
    """ Given a (len x 2) array, calculate the deviation from the median per pol.
    Calculates median over a window, win.
    """

    med = np.zeros_like(arr)
    for i in range(len(arr)):
        inds = list(range(max(0, i-win//2), i)) + list(range(i+1, min(i+win//2, len(arr))))
        med[i] = np.ma.median(arr.take(inds, axis=0), axis=0)

    return arr-med

def getonlineflags(st, segment):
    """ Gets antenna flags for a given segment from either sdm or mcaf server.
    Returns an array of flags (1: good, 0: bad) for each baseline.
    """

    t0, t1 = st.segmenttimes[segment]
    if st.metadata.datasource == 'sdm':
        sdm = util.getsdm(st.metadata.filename, bdfdir=st.metadata.bdfdir)
        scan = sdm.scan(st.metadata.scan)
        takebls = [st.metadata.blarr_orig.tolist().index(list(bl)) for bl in st.blarr]
        flags = scan.flags([t0, t1])[:, takebls].all(axis=0)
    elif st.metadata.datasource == 'vys':
        try:
            from realfast.mcaf_servers import getblflags
            flags = getblflags(st.metadata.datasetId, st.blarr,
                               startTime=t0, endTime=t1)
        except (ImportError, Exception):
            logger.warning("No mcaf antenna flag server flags available")
            flags = np.ones(st.nbl)

    if not flags.all():
        logger.info('Found antennas to flag in time range {0}-{1} '
                    .format(t0, t1))
    else:
        logger.info('No flagged antennas in time range {0}-{1} '
                    .format(t0, t1))

    return flags

def flag_data_rtpipe(st, data):
    """ Flagging data in single process
    Deprecated.
    """
    try:
        import rtlib_cython as rtlib
    except ImportError:
        logger.error("rtpipe not installed. Cannot import rtlib for flagging.")

    # **hack!**
    d = {'dataformat': 'sdm', 'ants': [int(ant.lstrip('ea')) for ant in st.ants], 'excludeants': st.prefs.excludeants, 'nants': len(st.ants)}

    for flag in st.prefs.flaglist:
        mode, sig, conv = flag
        for spw in st.spw:
            chans = np.arange(st.metadata.spw_nchan[spw]*spw, st.metadata.spw_nchan[spw]*(1+spw))
            for pol in range(st.npol):
                status = rtlib.dataflag(data, chans, pol, d, sig, mode, conv)
                logger.info(status)

    # hack to get rid of bad spw/pol combos whacked by rfi
    if st.prefs.badspwpol:
        logger.info('Comparing overall power between spw/pol. Removing those with {0} times typical value'.format(st.prefs.badspwpol))
        spwpol = {}
        for spw in st.spw:
            chans = np.arange(st.metadata.spw_nchan[spw]*spw, st.metadata.spw_nchan[spw]*(1+spw))
            for pol in range(st.npol):
                spwpol[(spw, pol)] = np.abs(data[:, :, chans, pol]).std()

        meanstd = np.mean(list(spwpol.values()))
        for (spw, pol) in spwpol:
            if spwpol[(spw, pol)] > st.prefs.badspwpol*meanstd:
                logger.info('Flagging all of (spw %d, pol %d) for excess noise.' % (spw, pol))
                chans = np.arange(st.metadata.spw_nchan[spw]*spw, st.metadata.spw_nchan[spw]*(1+spw))
                data[:, :, chans, pol] = 0j

    return data
<|endoftext|>