<commit_before>"""SSL tools."""

import os
import sys

from . import package
from . import utils

class CertificateBackend(object):
    """Base class."""

    def __init__(self, config):
        """Set path to certificates."""
        self.config = config

    def overwrite_existing_certificate(self):
        """Check if certificate already exists."""
        if os.path.exists(self.config.get("general", "tls_key_file")):
            if not self.config.getboolean("general", "force"):
                answer = utils.user_input(
                    "Overwrite the existing SSL certificate? (y/N) ")
                if not answer.lower().startswith("y"):
                    return False
        return True

class SelfSignedCertificate(CertificateBackend):
    """Create a self signed certificate."""

    def __init__(self, *args, **kwargs):
        """Sanity checks."""
        super(SelfSignedCertificate, self).__init__(*args, **kwargs)
        if self.config.has_option("general", "tls_key_file"):
            # Compatibility
            return
        for base_dir in ["/etc/pki/tls", "/etc/ssl"]:
            if os.path.exists(base_dir):
                self.config.set(
                    "general", "tls_key_file",
                    "{}/private/%(hostname)s.key".format(base_dir))
                self.config.set(
                    "general", "tls_cert_file",
                    "{}/certs/%(hostname)s.cert".format(base_dir))
                return
        raise RuntimeError("Cannot find a directory to store certificate")

    def generate_cert(self):
        """Create a certificate."""
        if not self.overwrite_existing_certificate():
            return
        utils.printcolor(
            "Generating new self-signed certificate", utils.YELLOW)
        utils.exec_cmd(
            "openssl req -new -newkey rsa:4096 -days 365 -nodes -x509 "
            "-subj '/CN={}' -keyout {} -out {}".format(
                self.config.get("general", "hostname"),
                self.config.get("general", "tls_key_file"),
                self.config.get("general", "tls_cert_file"))
        )

class LetsEncryptCertificate(CertificateBackend):
    """Create a certificate using letsencrypt."""

    def __init__(self, *args, **kwargs):
        """Update config."""
        super(LetsEncryptCertificate, self).__init__(*args, **kwargs)
        self.hostname = self.config.get("general", "hostname")
        self.config.set("general", "tls_cert_file", (
            "/etc/letsencrypt/live/{}/fullchain.pem".format(self.hostname)))
        self.config.set("general", "tls_key_file", (
            "/etc/letsencrypt/live/{}/privkey.pem".format(self.hostname)))

    def install_certbot(self):
        """Install certbot script to generate cert."""
        name, version, _id = utils.dist_info()
        name = name.lower()
        if name == "ubuntu":
            package.backend.update()
            package.backend.install("software-properties-common")
            utils.exec_cmd("add-apt-repository -y universe")
            if version == "18.04":
                utils.exec_cmd("add-apt-repository -y ppa:certbot/certbot")
            package.backend.update()
            package.backend.install("certbot")
        elif name == "debian":
            package.backend.update()
            package.backend.install("certbot")
        elif "centos" in name:
            package.backend.install("certbot")
        else:
            utils.printcolor("Failed to install certbot, aborting.", utils.RED)
            sys.exit(1)

    def generate_cert(self):
        """Create a certificate."""
        utils.printcolor(
            "Generating new certificate using letsencrypt", utils.YELLOW)
        self.install_certbot()
        utils.exec_cmd(
            "certbot certonly -n --standalone -d {} -m {} --agree-tos"
            .format(
                self.hostname, self.config.get("letsencrypt", "email")))
        with open("/etc/cron.d/letsencrypt", "w") as fp:
            fp.write("0 */12 * * * root certbot renew "
                     "--quiet --no-self-upgrade --force-renewal\n")
        cfg_file = "/etc/letsencrypt/renewal/{}.conf".format(self.hostname)
        pattern = "s/authenticator = standalone/authenticator = nginx/"
        utils.exec_cmd("perl -pi -e '{}' {}".format(pattern, cfg_file))

def get_backend(config):
    """Return the appropriate backend."""
    if not config.getboolean("certificate", "generate"):
        return None
    if config.get("certificate", "type") == "letsencrypt":
        return LetsEncryptCertificate(config)
    return SelfSignedCertificate(config)
<commit_msg>Fix #393<commit_after>"""SSL tools."""

import os
import sys

from . import package
from . import utils

class CertificateBackend(object):
    """Base class."""

    def __init__(self, config):
        """Set path to certificates."""
        self.config = config

    def overwrite_existing_certificate(self):
        """Check if certificate already exists."""
        if os.path.exists(self.config.get("general", "tls_key_file")):
            if not self.config.getboolean("general", "force"):
                answer = utils.user_input(
                    "Overwrite the existing SSL certificate? (y/N) ")
                if not answer.lower().startswith("y"):
                    return False
        return True

class SelfSignedCertificate(CertificateBackend):
    """Create a self signed certificate."""

    def __init__(self, *args, **kwargs):
        """Sanity checks."""
        super(SelfSignedCertificate, self).__init__(*args, **kwargs)
        if self.config.has_option("general", "tls_key_file"):
            # Compatibility
            return
        for base_dir in ["/etc/pki/tls", "/etc/ssl"]:
            if os.path.exists(base_dir):
                self.config.set(
                    "general", "tls_key_file",
                    "{}/private/%(hostname)s.key".format(base_dir))
                self.config.set(
                    "general", "tls_cert_file",
                    "{}/certs/%(hostname)s.cert".format(base_dir))
                return
        raise RuntimeError("Cannot find a directory to store certificate")

    def generate_cert(self):
        """Create a certificate."""
        if not self.overwrite_existing_certificate():
            return
        utils.printcolor(
            "Generating new self-signed certificate", utils.YELLOW)
        utils.exec_cmd(
            "openssl req -new -newkey rsa:4096 -days 365 -nodes -x509 "
            "-subj '/CN={}' -keyout {} -out {}".format(
                self.config.get("general", "hostname"),
                self.config.get("general", "tls_key_file"),
                self.config.get("general", "tls_cert_file"))
        )

class LetsEncryptCertificate(CertificateBackend):
    """Create a certificate using letsencrypt."""

    def __init__(self, *args, **kwargs):
        """Update config."""
        super(LetsEncryptCertificate, self).__init__(*args, **kwargs)
        self.hostname = self.config.get("general", "hostname")
        self.config.set("general", "tls_cert_file", (
            "/etc/letsencrypt/live/{}/fullchain.pem".format(self.hostname)))
        self.config.set("general", "tls_key_file", (
            "/etc/letsencrypt/live/{}/privkey.pem".format(self.hostname)))

    def install_certbot(self):
        """Install certbot script to generate cert."""
        name, version = utils.dist_info()
        name = name.lower()
        if name == "ubuntu":
            package.backend.update()
            package.backend.install("software-properties-common")
            utils.exec_cmd("add-apt-repository -y universe")
            if version == "18.04":
                utils.exec_cmd("add-apt-repository -y ppa:certbot/certbot")
            package.backend.update()
            package.backend.install("certbot")
        elif name.startswith("debian"):
            package.backend.update()
            package.backend.install("certbot")
        elif "centos" in name:
            package.backend.install("certbot")
        else:
            utils.printcolor("Failed to install certbot, aborting.", utils.RED)
            sys.exit(1)

    def generate_cert(self):
        """Create a certificate."""
        utils.printcolor(
            "Generating new certificate using letsencrypt", utils.YELLOW)
        self.install_certbot()
        utils.exec_cmd(
            "certbot certonly -n --standalone -d {} -m {} --agree-tos"
            .format(
                self.hostname, self.config.get("letsencrypt", "email")))
        with open("/etc/cron.d/letsencrypt", "w") as fp:
            fp.write("0 */12 * * * root certbot renew "
                     "--quiet --no-self-upgrade --force-renewal\n")
        cfg_file = "/etc/letsencrypt/renewal/{}.conf".format(self.hostname)
        pattern = "s/authenticator = standalone/authenticator = nginx/"
        utils.exec_cmd("perl -pi -e '{}' {}".format(pattern, cfg_file))

def get_backend(config):
    """Return the appropriate backend."""
    if not config.getboolean("certificate", "generate"):
        return None
    if config.get("certificate", "type") == "letsencrypt":
        return LetsEncryptCertificate(config)
    return SelfSignedCertificate(config)
<|endoftext|>


<commit_before>#!/usr/bin/env python3
"""Example for aiohttp.web basic server with cookies.
"""

from pprint import pformat

from aiohttp import web

tmpl = '''\

        Login

        Logout

{}

'''

async def root(request):
    resp = web.Response(content_type='text/html')
    resp.text = tmpl.format(pformat(request.cookies))
    return resp

async def login(request):
    resp = web.HTTPFound(location='/')
    resp.set_cookie('AUTH', 'secret')
    return resp

async def logout(request):
    resp = web.HTTPFound(location='/')
    resp.del_cookie('AUTH')
    return resp

def init():
    app = web.Application()
    app.router.add_get('/', root)
    app.router.add_get('/login', login)
    app.router.add_get('/logout', logout)
    return app

web.run_app(init())
<commit_msg>Fix ./examples/web_cookies.py to avoid a warning message<commit_after>#!/usr/bin/env python3
"""Example for aiohttp.web basic server with cookies.
"""

from pprint import pformat

from aiohttp import web

tmpl = '''\

        Login

        Logout

{}

'''

async def root(request):
    resp = web.Response(content_type='text/html')
    resp.text = tmpl.format(pformat(request.cookies))
    return resp

async def login(request):
    exc = web.HTTPFound(location='/')
    exc.set_cookie('AUTH', 'secret')
    raise exc

async def logout(request):
    exc = web.HTTPFound(location='/')
    exc.del_cookie('AUTH')
    raise exc

def init():
    app = web.Application()
    app.router.add_get('/', root)
    app.router.add_get('/login', login)
    app.router.add_get('/logout', logout)
    return app

web.run_app(init())
<|endoftext|>


<commit_before>import os
import sys

import arrow
from github import Github

from models import FileCommit, FileContent, Session, User

# Globals
#

GITHUB_API_TOKEN = os.environ.get('GITHUB_API_TOKEN', None)
if not GITHUB_API_TOKEN:
    print("warning: GITHUB_API_TOKEN is not defined. API calls are rate-limited.", file=sys.stderr)

gh = Github(GITHUB_API_TOKEN)

session = Session()

source_repo_name = 'sd17spring/ReadingJournal'
organization_name = source_repo_name.split('/')[0]

# get the repos
#

instructor_logins = {user.login for team in gh.get_organization(organization_name).get_teams() for user in team.get_members()}

source_repo = gh.get_repo(source_repo_name)
student_repos = sorted((repo for repo in source_repo.get_forks()
                        if repo.owner.login not in instructor_logins),
                       key=lambda repo: repo.owner.login.lower())
all_repos = [source_repo] + student_repos
[repo.owner.login for repo in all_repos]

# insert students
#

def upsert(instances, keyname):
    if not instances:
        return
    klass = instances[0].__class__
    instance_map = {getattr(instance, keyname): instance for instance in instances}
    for key, id in session.query(getattr(klass, keyname), klass.id).filter(getattr(klass, keyname).in_(instance_map.keys())):
        instance = instance_map.pop(key)
        instance.id = id
        session.merge(instance)
    session.add_all(instance_map.values())

students = [User(login=repo.owner.login, fullname=repo.owner.name, role='organization' if repo == source_repo else 'student')
            for repo in all_repos]
upsert(students, 'login')
session.commit()

# insert file hashes
#

file_hashes = {item.sha: (repo, item)
               for repo in all_repos
               for item in repo.get_git_tree(repo.get_commits()[0].sha, recursive=True).tree
               if item.type == 'blob'}

file_contents = [FileContent(sha=item.sha, content='') for repo, item in file_hashes.values()]

upsert(file_contents, 'sha')
session.commit()  # TODO remove this; commit with next transaction

# insert file commits
#

# TODO consolidate this with `upsert`, above
def upsert2(instances, key1, key2):
    MAX_ROWS = 200
    if not instances:
        return
    if len(instances) > MAX_ROWS:
        upsert2(instances[:MAX_ROWS], key1, key2)
        upsert2(instances[MAX_ROWS:], key1, key2)
        return
    klass = instances[0].__class__
    instance_map = {(getattr(instance, key1), getattr(instance, key2)): instance for instance in instances}
    rows = session.query(klass.id, getattr(klass, key1), getattr(klass, key2))
    rows = rows.filter(getattr(klass, key1).in_([k[0] for k in instance_map.keys()]))
    rows = rows.filter(getattr(klass, key2).in_([k[1] for k in instance_map.keys()]))
    for id, *keys in rows:
        if tuple(keys) not in instance_map:
            continue
        instance = instance_map.pop(tuple(keys))
        instance.id = id
        session.merge(instance)
    print('%s: updated %d records; added %d records' % (klass.__name__, len(instances) - len(instance_map), len(instance_map)))
    session.add_all(instance_map.values())

def parse_git_datetime(s):
    return arrow.get(s, 'ddd, DD MMM YYYY HH:mm:ss ZZZ').datetime

repo_commits = [(repo, commit)
                for repo in all_repos
                for commit in repo.get_commits()
                if repo == source_repo or (commit.author == repo.owner and len(commit.parents) == 1)]
len(repo_commits)

# Use a dict, to record only the latest commit for each file
file_commit_recs = {(repo.owner.login, item.filename): (item.sha, parse_git_datetime(commit.last_modified))
                    for repo, commit in reversed(repo_commits)
                    for item in commit.files}
len(file_commit_recs)

user_instances = list(session.query(User).filter(User.login.in_([repo.owner.login for repo in all_repos])))
user_instance_map = {instance.login: instance for instance in user_instances}  # FIXME there's surely some way to do this within the ORM

file_commits = [FileCommit(user_id=user_instance_map[login].id, path=path, mod_time=mod_time, sha=sha)
                for (login, path), (sha, mod_time) in file_commit_recs.items()]

upsert2(file_commits, 'user_id', 'path')
session.commit()
<commit_msg>consolidate upsert and upsert2<commit_after>import os
import sys

import arrow
from github import Github

from models import FileCommit, FileContent, Session, User

# Globals
#

GITHUB_API_TOKEN = os.environ.get('GITHUB_API_TOKEN', None)
if not GITHUB_API_TOKEN:
    print("warning: GITHUB_API_TOKEN is not defined. API calls are rate-limited.", file=sys.stderr)

gh = Github(GITHUB_API_TOKEN)

session = Session()

source_repo_name = 'sd17spring/ReadingJournal'
organization_name = source_repo_name.split('/')[0]

# get the repos
#

instructor_logins = {user.login for team in gh.get_organization(organization_name).get_teams() for user in team.get_members()}

source_repo = gh.get_repo(source_repo_name)
student_repos = sorted((repo for repo in source_repo.get_forks()
                        if repo.owner.login not in instructor_logins),
                       key=lambda repo: repo.owner.login.lower())
all_repos = [source_repo] + student_repos
[repo.owner.login for repo in all_repos]

# insert students
#

def upsert(session, instances, *keys):
    """Merge or add instances to the session.

    Instance are merged if the database contains a row with the same values for keys.
    """
    MAX_ROWS = 200
    if not instances:
        return
    if len(instances) > MAX_ROWS:
        upsert(session, instances[:MAX_ROWS], *keys)
        upsert(session, instances[MAX_ROWS:], *keys)
        return
    klass = instances[0].__class__
    instance_map = {tuple(getattr(instance, key) for key in keys): instance
                    for instance in instances}
    rows = session.query(klass.id, *(getattr(klass, key) for key in keys))
    for i, key in enumerate(keys):
        rows = rows.filter(getattr(klass, key).in_({k[i] for k in instance_map.keys()}))
    for id, *keys_values in rows:
        key = tuple(keys_values)
        if key not in instance_map:
            continue
        instance = instance_map.pop(key)
        instance.id = id
        session.merge(instance)
    print('%s: updated %d records; added %d records' % (klass.__name__, len(instances) - len(instance_map), len(instance_map)))
    session.add_all(instance_map.values())

students = [User(login=repo.owner.login, fullname=repo.owner.name, role='organization' if repo == source_repo else 'student')
            for repo in all_repos]
upsert(session, students, 'login')
session.commit()

# insert file hashes
#

file_hashes = {item.sha: (repo, item)
               for repo in all_repos
               for item in repo.get_git_tree(repo.get_commits()[0].sha, recursive=True).tree
               if item.type == 'blob'}

file_contents = [FileContent(sha=item.sha, content='') for repo, item in file_hashes.values()]

upsert(session, file_contents, 'sha')
session.commit()  # TODO remove this; commit with next transaction

# insert file commits
#

def parse_git_datetime(s):
    return arrow.get(s, 'ddd, DD MMM YYYY HH:mm:ss ZZZ').datetime

repo_commits = [(repo, commit)
                for repo in all_repos
                for commit in repo.get_commits()
                if repo == source_repo or (commit.author == repo.owner and len(commit.parents) == 1)]
len(repo_commits)

# Use a dict, to record only the latest commit for each file
file_commit_recs = {(repo.owner.login, item.filename): (item.sha, parse_git_datetime(commit.last_modified))
                    for repo, commit in reversed(repo_commits)
                    for item in commit.files}
len(file_commit_recs)

user_instances = list(session.query(User).filter(User.login.in_([repo.owner.login for repo in all_repos])))
user_instance_map = {instance.login: instance for instance in user_instances}  # FIXME there's surely some way to do this within the ORM

file_commits = [FileCommit(user_id=user_instance_map[login].id, path=path, mod_time=mod_time, sha=sha)
                for (login, path), (sha, mod_time) in file_commit_recs.items()]

upsert(session, file_commits, 'user_id', 'path')
session.commit()
<|endoftext|>


<commit_before># This file is part of Indico.
# Copyright (C) 2002 - 2018 European Organization for Nuclear Research (CERN).
#
# Indico is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License as
# published by the Free Software Foundation; either version 3 of the
# License, or (at your option) any later version.
#
# Indico 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 Indico; if not, see <http://www.gnu.org/licenses/>.

from __future__ import unicode_literals

from datetime import date, datetime, time

from dateutil.relativedelta import relativedelta
from flask import session
from sqlalchemy.orm import joinedload, load_only, raiseload

from indico.core.db import db
from indico.core.db.sqlalchemy.principals import PrincipalType
from indico.core.db.sqlalchemy.util.queries import db_dates_overlap, escape_like
from indico.modules.rb import rb_settings
from indico.modules.rb.models.equipment import EquipmentType, RoomEquipmentAssociation
from indico.modules.rb.models.favorites import favorite_room_table
from indico.modules.rb.models.principals import RoomPrincipal
from indico.modules.rb.models.reservation_occurrences import ReservationOccurrence
from indico.modules.rb.models.reservations import Reservation
from indico.modules.rb.models.room_features import RoomFeature
from indico.modules.rb.models.rooms import Room
from indico.modules.rb.statistics import calculate_rooms_occupancy
from indico.modules.rb.util import rb_is_admin
from indico.util.caching import memoize_redis

def _filter_coordinates(query, filters):
    try:
        sw_lat = filters['sw_lat']
        sw_lng = filters['sw_lng']
        ne_lat = filters['ne_lat']
        ne_lng = filters['ne_lng']
    except KeyError:
        return query

    return query.filter(Room.latitude >= sw_lat,
                        Room.latitude <= ne_lat,
                        Room.longitude >= sw_lng,
                        Room.longitude <= ne_lng)

def _make_room_text_filter(text):
    text = '%{}%'.format(escape_like(text))
    columns = ('site', 'division', 'building', 'floor', 'number', 'comments', 'full_name')
    return db.or_(getattr(Room, col).ilike(text) for col in columns)

def _query_managed_rooms(user):
    criteria = [db.and_(RoomPrincipal.type == PrincipalType.user,
                        RoomPrincipal.user_id == user.id,
                        RoomPrincipal.has_management_permission())]
    for group in user.local_groups:
        criteria.append(db.and_(RoomPrincipal.type == PrincipalType.local_group,
                                RoomPrincipal.local_group_id == group.id,
                                RoomPrincipal.has_management_permission()))
    for group in user.iter_all_multipass_groups():
        criteria.append(db.and_(RoomPrincipal.type == PrincipalType.multipass_group,
                                RoomPrincipal.multipass_group_provider == group.provider.name,
                                db.func.lower(RoomPrincipal.multipass_group_name) == group.name.lower(),
                                RoomPrincipal.has_management_permission()))
    return Room.query.filter(Room.is_active, Room.acl_entries.any(db.or_(*criteria)))

def _query_all_rooms_for_acl_check():
    return (Room.query
            .filter(Room.is_active)
            .options(load_only('id', 'protection_mode', 'reservations_need_confirmation'),
                     raiseload('owner'),
                     joinedload('acl_entries')))

@memoize_redis(900)
def has_managed_rooms(user):
    if user.can_get_all_multipass_groups:
        return _query_managed_rooms(user).has_rows()
    else:
        query = _query_all_rooms_for_acl_check()
        return any(r.can_manage(user, allow_admin=False) for r in query)

@memoize_redis(900)
def get_managed_room_ids(user):
    if user.can_get_all_multipass_groups:
        return {id_ for id_, in _query_managed_rooms(user).with_entities(Room.id)}
    else:
        query = _query_all_rooms_for_acl_check()
        return {r.id for r in query if r.can_manage(user, allow_admin=False)}

@memoize_redis(3600)
def get_room_statistics(room):
    data = {
        'count': {
            'id': 'times_booked',
            'values': [],
            'note': False
        },
        'percentage': {
            'id': 'occupancy',
            'values': [],
            'note': True
        }
    }
    ranges = [7, 30, 365]
    end_date = date.today()
    for days in ranges:
        start_date = date.today() - relativedelta(days=days)
        count = (ReservationOccurrence.query
                 .join(ReservationOccurrence.reservation)
                 .join(Reservation.room)
                 .filter(Room.id == room.id,
                         ReservationOccurrence.is_valid,
                         db_dates_overlap(ReservationOccurrence,
                                          'start_dt', datetime.combine(start_date, time()),
                                          'end_dt', datetime.combine(end_date, time.max)))
                 .count())
        percentage = calculate_rooms_occupancy([room], start_date, end_date) * 100
        if count > 0 or percentage > 0:
            data['count']['values'].append({'days': days, 'value': count})
            data['percentage']['values'].append({'days': days, 'value': percentage})
    return data

def search_for_rooms(filters, allow_admin=False, availability=None):
    """Search for a room, using the provided filters.

    :param filters: The filters, provided as a dictionary
    :param allow_admin: A boolean specifying whether admins have override privileges
    :param availability: A boolean specifying whether (un)available rooms should be provided,
                         or `None` in case all rooms should be returned.
    """
    query = (Room.query
             .outerjoin(favorite_room_table, db.and_(favorite_room_table.c.user_id == session.user.id,
                                                     favorite_room_table.c.room_id == Room.id))
             .reset_joinpoint()  # otherwise filter_by() would apply to the favorite table
             .options(raiseload('owner'))
             .filter(Room.is_active)
             .order_by(favorite_room_table.c.user_id.is_(None), db.func.indico.natsort(Room.full_name)))

    criteria = {}
    if 'capacity' in filters:
        query = query.filter(Room.capacity >= filters['capacity'])
    if 'building' in filters:
        criteria['building'] = filters['building']
    if 'division' in filters:
        criteria['division'] = filters['division']
    query = query.filter_by(**criteria)
    if 'text' in filters:
        query = query.filter(_make_room_text_filter(filters['text']))
    if filters.get('equipment'):
        subquery = (db.session.query(RoomEquipmentAssociation)
                    .with_entities(db.func.count(RoomEquipmentAssociation.c.room_id))
                    .filter(
                        RoomEquipmentAssociation.c.room_id == Room.id,
                        EquipmentType.name.in_(filters['equipment'])
                    )
                    .join(EquipmentType, RoomEquipmentAssociation.c.equipment_id == EquipmentType.id)
                    .correlate(Room)
                    .as_scalar())
        query = query.filter(subquery == len(filters['equipment']))
    if filters.get('features'):
        for feature in filters['features']:
            query = query.filter(Room.available_equipment.any(EquipmentType.features.any(RoomFeature.name == feature)))
    if filters.get('favorite'):
        query = query.filter(favorite_room_table.c.user_id.isnot(None))
    if filters.get('mine'):
        ids = get_managed_room_ids(session.user)
        query = query.filter(Room.id.in_(ids))
    query = _filter_coordinates(query, filters)

    if availability is None:
        return query

    start_dt, end_dt = filters['start_dt'], filters['end_dt']
    repeatability = (filters['repeat_frequency'], filters['repeat_interval'])
    include_blockings = not rb_is_admin(session.user)
    availability_filters = [Room.filter_available(start_dt, end_dt, repeatability, include_blockings=include_blockings,
                                                  include_pre_bookings=True)]
    if not (allow_admin and rb_is_admin(session.user)):
        selected_period_days = (filters['end_dt'] - filters['start_dt']).days
        booking_limit_days = db.func.coalesce(Room.booking_limit_days, rb_settings.get('booking_limit'))

        criterion = db.and_(Room.filter_bookable_hours(start_dt.time(), end_dt.time()),
                            Room.filter_nonbookable_periods(start_dt, end_dt),
                            db.or_(booking_limit_days.is_(None),
                            selected_period_days <= booking_limit_days))
        unbookable_ids = [room.id
                          for room in query.filter(db.and_(*availability_filters), ~criterion)
                          if not room.can_override(session.user, allow_admin=False)]
        availability_filters.append(~Room.id.in_(unbookable_ids))
    availability_criterion = db.and_(*availability_filters)
    if availability is False:
        availability_criterion = ~availability_criterion
    return query.filter(availability_criterion)
<commit_msg>Trim search text input<commit_after># This file is part of Indico.
# Copyright (C) 2002 - 2018 European Organization for Nuclear Research (CERN).
#
# Indico is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License as
# published by the Free Software Foundation; either version 3 of the
# License, or (at your option) any later version.
#
# Indico 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 Indico; if not, see <http://www.gnu.org/licenses/>.

from __future__ import unicode_literals

from datetime import date, datetime, time

from dateutil.relativedelta import relativedelta
from flask import session
from sqlalchemy.orm import joinedload, load_only, raiseload

from indico.core.db import db
from indico.core.db.sqlalchemy.principals import PrincipalType
from indico.core.db.sqlalchemy.util.queries import db_dates_overlap, escape_like
from indico.modules.rb import rb_settings
from indico.modules.rb.models.equipment import EquipmentType, RoomEquipmentAssociation
from indico.modules.rb.models.favorites import favorite_room_table
from indico.modules.rb.models.principals import RoomPrincipal
from indico.modules.rb.models.reservation_occurrences import ReservationOccurrence
from indico.modules.rb.models.reservations import Reservation
from indico.modules.rb.models.room_features import RoomFeature
from indico.modules.rb.models.rooms import Room
from indico.modules.rb.statistics import calculate_rooms_occupancy
from indico.modules.rb.util import rb_is_admin
from indico.util.caching import memoize_redis

def _filter_coordinates(query, filters):
    try:
        sw_lat = filters['sw_lat']
        sw_lng = filters['sw_lng']
        ne_lat = filters['ne_lat']
        ne_lng = filters['ne_lng']
    except KeyError:
        return query

    return query.filter(Room.latitude >= sw_lat,
                        Room.latitude <= ne_lat,
                        Room.longitude >= sw_lng,
                        Room.longitude <= ne_lng)

def _make_room_text_filter(text):
    text = '%{}%'.format(escape_like(text))
    columns = ('site', 'division', 'building', 'floor', 'number', 'comments', 'full_name')
    return db.or_(getattr(Room, col).ilike(text) for col in columns)

def _query_managed_rooms(user):
    criteria = [db.and_(RoomPrincipal.type == PrincipalType.user,
                        RoomPrincipal.user_id == user.id,
                        RoomPrincipal.has_management_permission())]
    for group in user.local_groups:
        criteria.append(db.and_(RoomPrincipal.type == PrincipalType.local_group,
                                RoomPrincipal.local_group_id == group.id,
                                RoomPrincipal.has_management_permission()))
    for group in user.iter_all_multipass_groups():
        criteria.append(db.and_(RoomPrincipal.type == PrincipalType.multipass_group,
                                RoomPrincipal.multipass_group_provider == group.provider.name,
                                db.func.lower(RoomPrincipal.multipass_group_name) == group.name.lower(),
                                RoomPrincipal.has_management_permission()))
    return Room.query.filter(Room.is_active, Room.acl_entries.any(db.or_(*criteria)))

def _query_all_rooms_for_acl_check():
    return (Room.query
            .filter(Room.is_active)
            .options(load_only('id', 'protection_mode', 'reservations_need_confirmation'),
                     raiseload('owner'),
                     joinedload('acl_entries')))

@memoize_redis(900)
def has_managed_rooms(user):
    if user.can_get_all_multipass_groups:
        return _query_managed_rooms(user).has_rows()
    else:
        query = _query_all_rooms_for_acl_check()
        return any(r.can_manage(user, allow_admin=False) for r in query)

@memoize_redis(900)
def get_managed_room_ids(user):
    if user.can_get_all_multipass_groups:
        return {id_ for id_, in _query_managed_rooms(user).with_entities(Room.id)}
    else:
        query = _query_all_rooms_for_acl_check()
        return {r.id for r in query if r.can_manage(user, allow_admin=False)}

@memoize_redis(3600)
def get_room_statistics(room):
    data = {
        'count': {
            'id': 'times_booked',
            'values': [],
            'note': False
        },
        'percentage': {
            'id': 'occupancy',
            'values': [],
            'note': True
        }
    }
    ranges = [7, 30, 365]
    end_date = date.today()
    for days in ranges:
        start_date = date.today() - relativedelta(days=days)
        count = (ReservationOccurrence.query
                 .join(ReservationOccurrence.reservation)
                 .join(Reservation.room)
                 .filter(Room.id == room.id,
                         ReservationOccurrence.is_valid,
                         db_dates_overlap(ReservationOccurrence,
                                          'start_dt', datetime.combine(start_date, time()),
                                          'end_dt', datetime.combine(end_date, time.max)))
                 .count())
        percentage = calculate_rooms_occupancy([room], start_date, end_date) * 100
        if count > 0 or percentage > 0:
            data['count']['values'].append({'days': days, 'value': count})
            data['percentage']['values'].append({'days': days, 'value': percentage})
    return data

def search_for_rooms(filters, allow_admin=False, availability=None):
    """Search for a room, using the provided filters.

    :param filters: The filters, provided as a dictionary
    :param allow_admin: A boolean specifying whether admins have override privileges
    :param availability: A boolean specifying whether (un)available rooms should be provided,
                         or `None` in case all rooms should be returned.
    """
    query = (Room.query
             .outerjoin(favorite_room_table, db.and_(favorite_room_table.c.user_id == session.user.id,
                                                     favorite_room_table.c.room_id == Room.id))
             .reset_joinpoint()  # otherwise filter_by() would apply to the favorite table
             .options(raiseload('owner'))
             .filter(Room.is_active)
             .order_by(favorite_room_table.c.user_id.is_(None), db.func.indico.natsort(Room.full_name)))

    criteria = {}
    if 'capacity' in filters:
        query = query.filter(Room.capacity >= filters['capacity'])
    if 'building' in filters:
        criteria['building'] = filters['building']
    if 'division' in filters:
        criteria['division'] = filters['division']
    query = query.filter_by(**criteria)
    if 'text' in filters:
        text = ' '.join(filters['text'].strip().split())
        query = query.filter(_make_room_text_filter(text))
    if filters.get('equipment'):
        subquery = (db.session.query(RoomEquipmentAssociation)
                    .with_entities(db.func.count(RoomEquipmentAssociation.c.room_id))
                    .filter(
                        RoomEquipmentAssociation.c.room_id == Room.id,
                        EquipmentType.name.in_(filters['equipment'])
                    )
                    .join(EquipmentType, RoomEquipmentAssociation.c.equipment_id == EquipmentType.id)
                    .correlate(Room)
                    .as_scalar())
        query = query.filter(subquery == len(filters['equipment']))
    if filters.get('features'):
        for feature in filters['features']:
            query = query.filter(Room.available_equipment.any(EquipmentType.features.any(RoomFeature.name == feature)))
    if filters.get('favorite'):
        query = query.filter(favorite_room_table.c.user_id.isnot(None))
    if filters.get('mine'):
        ids = get_managed_room_ids(session.user)
        query = query.filter(Room.id.in_(ids))
    query = _filter_coordinates(query, filters)

    if availability is None:
        return query

    start_dt, end_dt = filters['start_dt'], filters['end_dt']
    repeatability = (filters['repeat_frequency'], filters['repeat_interval'])
    include_blockings = not rb_is_admin(session.user)
    availability_filters = [Room.filter_available(start_dt, end_dt, repeatability, include_blockings=include_blockings,
                                                  include_pre_bookings=True)]
    if not (allow_admin and rb_is_admin(session.user)):
        selected_period_days = (filters['end_dt'] - filters['start_dt']).days
        booking_limit_days = db.func.coalesce(Room.booking_limit_days, rb_settings.get('booking_limit'))

        criterion = db.and_(Room.filter_bookable_hours(start_dt.time(), end_dt.time()),
                            Room.filter_nonbookable_periods(start_dt, end_dt),
                            db.or_(booking_limit_days.is_(None),
                            selected_period_days <= booking_limit_days))
        unbookable_ids = [room.id
                          for room in query.filter(db.and_(*availability_filters), ~criterion)
                          if not room.can_override(session.user, allow_admin=False)]
        availability_filters.append(~Room.id.in_(unbookable_ids))
    availability_criterion = db.and_(*availability_filters)
    if availability is False:
        availability_criterion = ~availability_criterion
    return query.filter(availability_criterion)
<|endoftext|>


<commit_before># Re-puts entities of a given type, to set newly added properties.
#
# To run this script:
#  - Make sure App Engine library (incl. yaml) is in PYTHONPATH.
#  - Make sure that the remote API is included in app.yaml.
#  - Run "tools/appengine_console.py APP_ID".
#  - Import this module.
#  - update_entities.run(models.Issue) updates issues.

import logging
from google.appengine.ext import db
from codereview import models

def run(model_class, batch_size=100):
    last_key = None
    while True:
      q = model_class.all()
      if last_key:
        q.filter('__key__ >', last_key)
      q.order('__key__')
      this_batch_size = batch_size

      while True:
        try:
          batch = q.fetch(this_batch_size)
          break
        except db.Timeout:
          logging.warn("Query timed out, retrying")
          if this_batch_size == 1:
            logging.critical("Unable to update entities, aborting")
            return
          this_batch_size //= 2

      if not batch:
        break

      keys = None
      while not keys:
        try:
          keys = db.put(batch)
        except db.Timeout:
          logging.warn("Put timed out, retrying")

      last_key = keys[-1]
      print "Updated %d records" % (len(keys),)
<commit_msg>Catch a different kind of timeouts too (urllib2.URLError).<commit_after># Re-puts entities of a given type, to set newly added properties.
#
# To run this script:
#  - Make sure App Engine library (incl. yaml) is in PYTHONPATH.
#  - Make sure that the remote API is included in app.yaml.
#  - Run "tools/appengine_console.py APP_ID".
#  - Import this module.
#  - Import models from codereview.
#  - update_entities.run(models.Issue) updates issues.

import logging
from google.appengine.ext import db
from codereview import models
import urllib2

def run(model_class, batch_size=100, last_key=None):
    while True:
      q = model_class.all()
      if last_key:
        q.filter('__key__ >', last_key)
      q.order('__key__')
      this_batch_size = batch_size

      while True:
        try:
          try:
            batch = q.fetch(this_batch_size)
          except urllib2.URLError, err:
            if 'timed out' in str(err):
              raise db.Timeout
            else:
              raise
          break
        except db.Timeout:
          logging.warn("Query timed out, retrying")
          if this_batch_size == 1:
            logging.critical("Unable to update entities, aborting")
            return
          this_batch_size //= 2

      if not batch:
        break

      keys = None
      while not keys:
        try:
          keys = db.put(batch)
        except db.Timeout:
          logging.warn("Put timed out, retrying")

      last_key = keys[-1]
      print "Updated %d records" % (len(keys),)
<|endoftext|>


<commit_before># Tim Cornwell <realtimcornwell@gmail.com>)

import numpy

from crocodile.synthesis import sortw, doimg, dopredict, wslicimg, wslicfwd

from crocodile.clean import overlapIndices, argmax

import matplotlib.pyplot as plt

def majorcycle(T2, L2,
               p, v,
               gain,
               nmajor,
               nminor,
               wstep,
               scales,
               thresh=0.0,
               fracthresh=0.1):
    """Major cycle for MultiScale Clean

    :param T2: Field of view in radians
    :param L2: Maximum uv (for setting image pixel sizes)
    :param p: UVWs of visibilities (m)
    :param v: Values of visibilities
    :param nmajor: Number of major cycles
    :param nminor: Number of minor cycles
    :param wstep: Step in w (pixels)
    :param wscales: Array of scales in pixels
    :print thresh: Stopping threshold (for scale=0)
    :print fracthresh: Minor Cycle stopping threshold (for scale=0) fraction of peak
    """

    # The model is added to each major cycle and then the visibilities are
    # calculated from the full model
    ps, vso = sortw(p, v)
    dirty,psf=doimg(T2, L2, ps, vso, lambda *x: wslicimg(*x, wstep=wstep))
    comps=0.0*dirty.copy()
    for i in range(nmajor):
        print("Start of major cycle %d" % (i))
        cc, res=msclean(dirty, psf, True, gain, thresh, nminor, scales, fracthresh)
        plt.clf()
        plt.imshow(res,cmap='rainbow', origin='lower')
        plt.colorbar()
        plt.show()
        comps+=cc
        # dopredict resorts the data
        pss, vsp=dopredict(T2, L2, ps, comps, lambda *x: wslicfwd(*x, wstep=wstep))
        vsr=vso-vsp
        dirty, psf=doimg(T2, L2, ps, vsr, lambda *x: wslicimg(*x, wstep=wstep))
        print("End of major cycle")
    print("End of major cycles")
    return ps, vsp, comps, dirty

def msclean(dirty,
           psf,
           window,
           gain,
           thresh,
           niter,
           scales,
           fracthresh):
    """
    Multiscale CLEAN minor cycle (IEEE Journal of Selected Topics in Sig Proc, 2008 vol. 2 pp. 793-801)

    :param dirty: The dirty image, i.e., the image to be deconvolved
    :param psf: The point spread-function
    :param window: Regions where clean components are allowed. If
    True, then all of the dirty image is assumed to be allowed for
    clean components
    :param gain: The "loop gain", i.e., the fraction of the brightest
    pixel that is removed in each iteration
    :param thresh: Cleaning stops when the maximum of the absolute
    deviation of the residual is less than this value
    :param niter: Maximum number of components to make if the
    threshold "thresh" is not hit
    :param scales: Scales (in pixels width) to be used
    :returns clean component image, residual image
    """
    assert gain > 0.0 and gain < 2.0
    assert niter > 0
    assert len(scales) > 0

    comps=numpy.zeros(dirty.shape)

    pmax=psf.max()
    assert pmax> 0.0

    psfpeak=argmax(numpy.fabs(psf))
    print ("Peak of PSF = %s at %s" % (pmax, psfpeak))
    dmax=dirty.max()
    dpeak=argmax(dirty)
    print ("Peak of Dirty = %s at %s" % (dmax, dpeak))
    lpsf=psf/pmax
    ldirty=dirty/pmax

    # Create the scale images and form all the various products we need. We
    # use a third dimension to hold the scale-related images. scalestack is a 3D
    # cube holding the different scale images. convolvestack will take a 2D image
    # and add a third dimension holding the scale-convolved versions.

    scaleshape=[ldirty.shape[0], ldirty.shape[1], len(scales)]
    scalescaleshape=[ldirty.shape[0], ldirty.shape[1], len(scales), len(scales)]
    scalestack=createscalestack(scaleshape,scales, norm=True)

    couplingMatrix=numpy.zeros([len(scales),len(scales)])
    psfscalestack=convolvescalestack(scalestack, numpy.array(lpsf))
    resscalestack=convolvescalestack(scalestack, numpy.array(ldirty))
    # Evaluate the coupling matrix between the various scale sizes.
    psfscalescalestack=numpy.zeros(scalescaleshape)
    for iscale in numpy.arange(len(scales)):
        psfscalescalestack[:,:,:,iscale]=convolvescalestack(scalestack,psfscalestack[:,:,iscale])
        psfscalescalestack[:,:,iscale,:]=psfscalescalestack[:,:,:,iscale]
    for iscale in numpy.arange(len(scales)):
        for iscale1 in numpy.arange(len(scales)):
            couplingMatrix[iscale,iscale1]=numpy.max(psfscalescalestack[:,:,iscale,iscale1])
    print ("Coupling matrix =\n %s" % couplingMatrix)

    # The window is scale dependent - we form it by smoothing and thresholding
    # the input window. This prevents components being placed too close to the
    # edge of the image.

    if window is True:
        windowstack=numpy.ones(scalestack.shape, numpy.bool)
#    windowstack=convolvescalestack(scalestack, window)>0.9
    window=numpy.ones(scalestack.shape, numpy.bool)

    """ The minor cycle
    """
    print("Max abs in dirty image = %.6f" % numpy.fabs(resscalestack[:,:,0]).max())
    absolutethresh=max(thresh,fracthresh*numpy.fabs(resscalestack[:,:,0]).max())
    print("Start of minor cycle")
    print("This minor cycle will stop at %d iterations or peak < %s" % (niter, absolutethresh))

    for i in range(niter):
        # Find peak over all smoothed images
        mx,my,mscale=findabsmaxstack(resscalestack, window, couplingMatrix)
        if mx == None or my == None or mscale == None:
            print("Error in finding peak")
            break

        # Find the values to subtract, accounting for the coupling matrix
        mval=numpy.zeros(len(scales))
        mval[mscale]=resscalestack[mx, my, mscale]/couplingMatrix[mscale,mscale]
        print ("Minor cycle %d, peak %s at [%d, %d, %d]" % \
            (i, resscalestack[mx, my, :], mx, my, mscale))
        if numpy.fabs(mval[mscale]) < absolutethresh:
            print("Absolute value of peak %.6f is below stopping threshold %.6f" \
                % (numpy.fabs(resscalestack[mx,my,mscale]), absolutethresh))
            break

        #  Update the cached residuals and add to the cached model.
        a1o, a2o=overlapIndices(dirty, psf, mx-psfpeak[0], my-psfpeak[1])
        if numpy.abs(mval[mscale])>0:
            # Cross subtract from other scales
            for iscale in range(len(scales)):
                resscalestack[a1o[0]:a1o[1],a1o[2]:a1o[3],iscale]-= \
                    psfscalescalestack[a2o[0]:a2o[1],a2o[2]:a2o[3],iscale,mscale] * \
                    gain * mval[mscale]
            comps[a1o[0]:a1o[1],a1o[2]:a1o[3]] += \
                scalestack[a2o[0]:a2o[1],a2o[2]:a2o[3],mscale] * \
                gain * mval[mscale]
        else:
            break
    print("End of minor cycles")
    return comps, pmax*resscalestack[:,:,0]

def createscalestack(scaleshape,scales, norm=False):
    """ Create a cube consisting of the scales

    :param scaleshape: desired shape of stack
    :param scales: scales (in pixels)
    """
    assert scaleshape[2] == len(scales)

    basis=numpy.zeros(scaleshape)
    nx=scaleshape[0]
    ny=scaleshape[1]
    xcen=int(numpy.ceil(float(nx)/2.0))
    ycen=int(numpy.ceil(float(ny)/2.0))
    for iscale in numpy.arange(0,len(scales)):
        halfscale=int(numpy.ceil(scales[iscale]/2.0))
        if scales[iscale]>0.0:
            rscale2=1.0/(float(scales[iscale])/2.0)**2
            x=range(xcen-halfscale-1,xcen+halfscale+1)
            fx=numpy.array(x, 'float')-float(xcen)
            for y in range(ycen-halfscale-1,ycen+halfscale+1):
                fy=float(y-ycen)
                r2=fx*fx+fy*fy
                basis[x,y,iscale]=(1.0-r2*rscale2)
            basis[basis<0.0]=0.0
            if norm:
                basis[:,:,iscale]/=numpy.sum(basis[:,:,iscale])
        else:
            basis[xcen,ycen,iscale]=1.0
    return basis

def convolvescalestack(scalestack, img):
    """Convolve img by the specified scalestack, returning the resulting stack

    :param scalestack: stack containing the scales
    :param img: image to be convolved
    """

    convolved=numpy.zeros(scalestack.shape)
    ximg=numpy.fft.fftshift(numpy.fft.fft2(numpy.fft.fftshift(img)))

    nscales=scalestack.shape[2]
    for iscale in range(nscales):
        xscale=numpy.fft.fftshift(numpy.fft.fft2(numpy.fft.fftshift(scalestack[:,:,iscale])))
        xmult=ximg*xscale
        convolved[:,:,iscale]=numpy.real(numpy.fft.ifftshift(numpy.fft.ifft2(numpy.fft.ifftshift(xmult))))
    return convolved

def findabsmaxstack(stack, window, couplingMatrix):
    """Find the location and value of the absolute maximum in this stack
    """
    pabsmax=0.0
    pscale=None
    px=None
    py=None
    pshape=[stack.shape[0],stack.shape[1]]
    for iscale in range(stack.shape[2]):
        mx, my=numpy.unravel_index(numpy.fabs((stack)[:,:,iscale]).argmax(), pshape)
        thisabsmax=stack[mx,my,iscale]/couplingMatrix[iscale,iscale]
        if abs(thisabsmax)>abs(pabsmax):
            px=mx
            py=my
            pscale=iscale
            pabsmax=stack[px,py,pscale]/couplingMatrix[iscale,iscale]
    return px, py, pscale
<commit_msg>Delete old msclean<commit_after><|endoftext|>


<commit_before>"""
Spglib interface for ASE
"""

import phonopy._spglib as spg
import numpy as np

def get_symmetry(bulk, symprec=1e-5, angle_tolerance=-1.0):
    """
    Return symmetry operations as hash.
    Hash key 'rotations' gives the numpy integer array
    of the rotation matrices for scaled positions
    Hash key 'translations' gives the numpy double array
    of the translation vectors in scaled positions
    """

    # Atomic positions have to be specified by scaled positions for spglib.
    positions = bulk.get_scaled_positions().copy()
    lattice = bulk.get_cell().T.copy()
    numbers = np.array(bulk.get_atomic_numbers(), dtype='intc').copy()

    # Get number of symmetry operations and allocate symmetry operations
    # multi = spg.multiplicity(cell, positions, numbers, symprec)
    multi = 48 * bulk.get_number_of_atoms()
    rotation = np.zeros((multi, 3, 3), dtype='intc')
    translation = np.zeros((multi, 3))

    # Get symmetry operations
    magmoms = bulk.get_magnetic_moments()
    if magmoms == None:
        num_sym = spg.symmetry(rotation,
                               translation,
                               lattice,
                               positions,
                               numbers,
                               symprec,
                               angle_tolerance)
    else:
        num_sym = spg.symmetry_with_collinear_spin(rotation,
                                                   translation,
                                                   lattice,
                                                   positions,
                                                   numbers,
                                                   magmoms,
                                                   symprec,
                                                   angle_tolerance)

    return {'rotations': rotation[:num_sym].copy(),
            'translations': translation[:num_sym].copy()}

def get_symmetry_dataset(bulk, symprec=1e-5, angle_tolerance=-1.0):
    """
    number: International space group number
    international: International symbol
    hall: Hall symbol
    transformation_matrix:
      Transformation matrix from lattice of input cell to Bravais lattice
      L^bravais = L^original * Tmat
    origin shift: Origin shift in the setting of 'Bravais lattice'
    rotations, translations:
      Rotation matrices and translation vectors
      Space group operations are obtained by
        [(r,t) for r, t in zip(rotations, translations)]
    wyckoffs:
      Wyckoff letters
    """
    positions = bulk.get_scaled_positions().copy()
    lattice = bulk.get_cell().T.copy()
    numbers = np.array(bulk.get_atomic_numbers(), dtype='intc').copy()
    keys = ('number',
            'international',
            'hall',
            'transformation_matrix',
            'origin_shift',
            'rotations',
            'translations',
            'wyckoffs',
            'equivalent_atoms')
    dataset = {}
    for key, data in zip(keys, spg.dataset(lattice,
                                           positions,
                                           numbers,
                                           symprec,
                                           angle_tolerance)):
        dataset[key] = data

    dataset['international'] = dataset['international'].strip()
    dataset['hall'] = dataset['hall'].strip()
    dataset['transformation_matrix'] = np.array(
        dataset['transformation_matrix'], dtype='double')
    dataset['origin_shift'] = np.array(dataset['origin_shift'],
                                       dtype='double')
    dataset['rotations'] = np.array(dataset['rotations'], dtype='intc')
    dataset['translations'] = np.array(dataset['translations'],
                                       dtype='double')
    letters = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
    dataset['wyckoffs'] = [letters[x] for x in dataset['wyckoffs']]
    dataset['equivalent_atoms'] = np.array(dataset['equivalent_atoms'],
                                           dtype='intc')

    return dataset

def get_spacegroup(bulk, symprec=1e-5, angle_tolerance=-1.0):
    """
    Return space group in international table symbol and number
    as a string.
    """
    # Atomic positions have to be specified by scaled positions for spglib.
    return spg.spacegroup(bulk.get_cell().T.copy(),
                          bulk.get_scaled_positions().copy(),
                          np.array(bulk.get_atomic_numbers(),
                                   dtype='intc').copy(),
                          symprec,
                          angle_tolerance)

def get_pointgroup(rotations):
    """
    Return point group in international table symbol and number.
    The symbols are mapped to the numbers as follows:
    1   "1    "
    2   "-1   "
    3   "2    "
    4   "m    "
    5   "2/m  "
    6   "222  "
    7   "mm2  "
    8   "mmm  "
    9   "4    "
    10  "-4   "
    11  "4/m  "
    12  "422  "
    13  "4mm  "
    14  "-42m "
    15  "4/mmm"
    16  "3    "
    17  "-3   "
    18  "32   "
    19  "3m   "
    20  "-3m  "
    21  "6    "
    22  "-6   "
    23  "6/m  "
    24  "622  "
    25  "6mm  "
    26  "-62m "
    27  "6/mmm"
    28  "23   "
    29  "m-3  "
    30  "432  "
    31  "-43m "
    32  "m-3m "
    """

    # (symbol, pointgroup_number, transformation_matrix)
    return spg.pointgroup(np.array(rotations, dtype='intc').copy())

def refine_cell(bulk, symprec=1e-5, angle_tolerance=-1.0):
    """
    Return refined cell
    """
    # Atomic positions have to be specified by scaled positions for spglib.
    num_atom = bulk.get_number_of_atoms()
    lattice = bulk.get_cell().T.copy()
    pos = np.zeros((num_atom * 4, 3), dtype='double')
    pos[:num_atom] = bulk.get_scaled_positions()

    numbers = np.zeros(num_atom * 4, dtype='intc')
    numbers[:num_atom] = np.array(bulk.get_atomic_numbers(), dtype='intc')
    num_atom_bravais = spg.refine_cell(lattice,
                                       pos,
                                       numbers,
                                       num_atom,
                                       symprec,
                                       angle_tolerance)

    return (lattice.T.copy(),
            pos[:num_atom_bravais].copy(),
            numbers[:num_atom_bravais].copy())

def find_primitive(bulk, symprec=1e-5, angle_tolerance=-1.0):
    """
    A primitive cell in the input cell is searched and returned
    as an object of Atoms class.
    If no primitive cell is found, (None, None, None) is returned.
    """

    # Atomic positions have to be specified by scaled positions for spglib.
    positions = bulk.get_scaled_positions().copy()
    lattice = bulk.get_cell().T.copy()
    numbers = np.array(bulk.get_atomic_numbers(), dtype='intc').copy()

    # lattice is transposed with respect to the definition of Atoms class
    num_atom_prim = spg.primitive(lattice,
                                  positions,
                                  numbers,
                                  symprec,
                                  angle_tolerance)
    if num_atom_prim > 0:
        return (lattice.T.copy(),
                positions[:num_atom_prim].copy(),
                numbers[:num_atom_prim].copy())
    else:
        return None, None, None

def get_ir_kpoints(kpoint,
                   bulk,
                   is_time_reversal=True,
                   symprec=1e-5):
    """
    Retrun irreducible kpoints
    """
    mapping = np.zeros(kpoint.shape[0], dtype='intc')
    spg.ir_kpoints(mapping,
                   kpoint,
                   bulk.get_cell().T.copy(),
                   bulk.get_scaled_positions().copy(),
                   np.array(bulk.get_atomic_numbers(), dtype='intc').copy(),
                   is_time_reversal * 1,
                   symprec)
    return mapping

def get_ir_reciprocal_mesh(mesh,
                           bulk,
                           is_shift=np.zeros(3, dtype='intc'),
                           is_time_reversal=True,
                           symprec=1e-5):
    """
    Return k-points mesh and k-point map to the irreducible k-points
    The symmetry is serched from the input cell.
    is_shift=[0, 0, 0] gives Gamma center mesh.
    """

    mapping = np.zeros(np.prod(mesh), dtype='intc')
    mesh_points = np.zeros((np.prod(mesh), 3), dtype='intc')
    spg.ir_reciprocal_mesh(mesh_points,
                           mapping,
                           np.array(mesh, dtype='intc').copy(),
                           np.array(is_shift, dtype='intc').copy(),
                           is_time_reversal * 1,
                           bulk.get_cell().T.copy(),
                           bulk.get_scaled_positions().copy(),
                           np.array(bulk.get_atomic_numbers(),
                                    dtype='intc').copy(),
                           symprec)

    return mapping, mesh_points

def relocate_BZ_grid_address(grid_address,
                             mesh,
                             reciprocal_lattice, # column vectors
                             is_shift=np.zeros(3, dtype='intc')):
    bz_grid_address = np.zeros(
        ((mesh[0] + 1) * (mesh[1] + 1) * (mesh[2] + 1), 3), dtype='intc')
    bz_map = np.zeros(
        (2 * mesh[0] - 1) * (2 * mesh[1] - 1) * (2 * mesh[2] - 1), dtype='intc')
    num_bz_ir = spg.BZ_grid_address(
        bz_grid_address,
        bz_map,
        grid_address,
        np.array(mesh, dtype='intc').copy(),
        np.array(reciprocal_lattice, dtype='double').copy(),
        np.array(is_shift, dtype='intc').copy())

    return bz_grid_address[:num_bz_ir], bz_map

def get_stabilized_reciprocal_mesh(mesh,
                                   rotations,
                                   is_shift=np.zeros(3, dtype='intc'),
                                   is_time_reversal=True,
                                   qpoints=np.array([], dtype='double')):
    """
    Return k-point map to the irreducible k-points and k-point grid points .

    The symmetry is searched from the input rotation matrices in real space.

    is_shift=[0, 0, 0] gives Gamma center mesh and the values 1 give
    half mesh distance shifts.
    """

    mapping = np.zeros(np.prod(mesh), dtype='intc')
    mesh_points = np.zeros((np.prod(mesh), 3), dtype='intc')
    qpoints = np.array(qpoints, dtype='double').copy()
    if qpoints.shape == (3,):
        qpoints = np.array([qpoints], dtype='double')
    if qpoints.shape == (0,):
        qpoints = np.array([[0, 0, 0]], dtype='double')
    spg.stabilized_reciprocal_mesh(mesh_points,
                                   mapping,
                                   np.array(mesh, dtype='intc').copy(),
                                   np.array(is_shift, dtype='intc'),
                                   is_time_reversal * 1,
                                   np.array(rotations, dtype='intc').copy(),
                                   np.array(qpoints, dtype='double'))

    return mapping, mesh_points

def get_triplets_reciprocal_mesh_at_q(fixed_grid_number,
                                      mesh,
                                      rotations,
                                      is_time_reversal=True):

    weights = np.zeros(np.prod(mesh), dtype='intc')
    third_q = np.zeros(np.prod(mesh), dtype='intc')
    mesh_points = np.zeros((np.prod(mesh), 3), dtype='intc')

    spg.triplets_reciprocal_mesh_at_q(weights,
                                      mesh_points,
                                      third_q,
                                      fixed_grid_number,
                                      np.array(mesh, dtype='intc').copy(),
                                      is_time_reversal * 1,
                                      np.array(rotations, dtype='intc').copy())

    return weights, third_q, mesh_points

def get_BZ_triplets_at_q(grid_point,
                         bz_grid_address,
                         bz_map,
                         weights,
                         mesh):
    """grid_address is overwritten."""
    num_ir_tripltes = (weights > 0).sum()
    triplets = np.zeros((num_ir_tripltes, 3), dtype='intc')
    num_ir_ret = spg.BZ_triplets_at_q(triplets,
                                      grid_point,
                                      bz_grid_address,
                                      bz_map,
                                      weights,
                                      np.array(mesh, dtype='intc').copy())
    return triplets

<commit_msg>Spglib update<commit_after>"""
Spglib interface for ASE
"""

import phonopy._spglib as spg
import numpy as np

def get_symmetry(bulk, use_magmoms=False, symprec=1e-5, angle_tolerance=-1.0):
    """
    Return symmetry operations as hash.
    Hash key 'rotations' gives the numpy integer array
    of the rotation matrices for scaled positions
    Hash key 'translations' gives the numpy double array
    of the translation vectors in scaled positions
    """

    # Atomic positions have to be specified by scaled positions for spglib.
    positions = bulk.get_scaled_positions().copy()
    lattice = bulk.get_cell().T.copy()
    numbers = np.intc(bulk.get_atomic_numbers()).copy()

    # Get number of symmetry operations and allocate symmetry operations
    # multi = spg.multiplicity(cell, positions, numbers, symprec)
    multi = 48 * bulk.get_number_of_atoms()
    rotation = np.zeros((multi, 3, 3), dtype='intc')
    translation = np.zeros((multi, 3))

    # Get symmetry operations
    if use_magmoms:
        magmoms = bulk.get_magnetic_moments()
        num_sym = spg.symmetry_with_collinear_spin(rotation,
                                                   translation,
                                                   lattice,
                                                   positions,
                                                   numbers,
                                                   magmoms,
                                                   symprec,
                                                   angle_tolerance)
    else:
        num_sym = spg.symmetry(rotation,
                               translation,
                               lattice,
                               positions,
                               numbers,
                               symprec,
                               angle_tolerance)

    return {'rotations': rotation[:num_sym].copy(),
            'translations': translation[:num_sym].copy()}

def get_symmetry_dataset(bulk, symprec=1e-5, angle_tolerance=-1.0):
    """
    number: International space group number
    international: International symbol
    hall: Hall symbol
    transformation_matrix:
      Transformation matrix from lattice of input cell to Bravais lattice
      L^bravais = L^original * Tmat
    origin shift: Origin shift in the setting of 'Bravais lattice'
    rotations, translations:
      Rotation matrices and translation vectors
      Space group operations are obtained by
        [(r,t) for r, t in zip(rotations, translations)]
    wyckoffs:
      Wyckoff letters
    """
    positions = bulk.get_scaled_positions().copy()
    lattice = bulk.get_cell().T.copy()
    numbers = np.array(bulk.get_atomic_numbers(), dtype='intc').copy()
    keys = ('number',
            'international',
            'hall',
            'transformation_matrix',
            'origin_shift',
            'rotations',
            'translations',
            'wyckoffs',
            'equivalent_atoms')
    dataset = {}
    for key, data in zip(keys, spg.dataset(lattice,
                                           positions,
                                           numbers,
                                           symprec,
                                           angle_tolerance)):
        dataset[key] = data

    dataset['international'] = dataset['international'].strip()
    dataset['hall'] = dataset['hall'].strip()
    dataset['transformation_matrix'] = np.array(
        dataset['transformation_matrix'], dtype='double')
    dataset['origin_shift'] = np.array(dataset['origin_shift'],
                                       dtype='double')
    dataset['rotations'] = np.array(dataset['rotations'], dtype='intc')
    dataset['translations'] = np.array(dataset['translations'],
                                       dtype='double')
    letters = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
    dataset['wyckoffs'] = [letters[x] for x in dataset['wyckoffs']]
    dataset['equivalent_atoms'] = np.array(dataset['equivalent_atoms'],
                                           dtype='intc')

    return dataset

def get_spacegroup(bulk, symprec=1e-5, angle_tolerance=-1.0):
    """
    Return space group in international table symbol and number
    as a string.
    """
    # Atomic positions have to be specified by scaled positions for spglib.
    return spg.spacegroup(bulk.get_cell().T.copy(),
                          bulk.get_scaled_positions().copy(),
                          np.array(bulk.get_atomic_numbers(),
                                   dtype='intc').copy(),
                          symprec,
                          angle_tolerance)

def get_pointgroup(rotations):
    """
    Return point group in international table symbol and number.
    The symbols are mapped to the numbers as follows:
    1   "1    "
    2   "-1   "
    3   "2    "
    4   "m    "
    5   "2/m  "
    6   "222  "
    7   "mm2  "
    8   "mmm  "
    9   "4    "
    10  "-4   "
    11  "4/m  "
    12  "422  "
    13  "4mm  "
    14  "-42m "
    15  "4/mmm"
    16  "3    "
    17  "-3   "
    18  "32   "
    19  "3m   "
    20  "-3m  "
    21  "6    "
    22  "-6   "
    23  "6/m  "
    24  "622  "
    25  "6mm  "
    26  "-62m "
    27  "6/mmm"
    28  "23   "
    29  "m-3  "
    30  "432  "
    31  "-43m "
    32  "m-3m "
    """

    # (symbol, pointgroup_number, transformation_matrix)
    return spg.pointgroup(np.array(rotations, dtype='intc').copy())

def refine_cell(bulk, symprec=1e-5, angle_tolerance=-1.0):
    """
    Return refined cell
    """
    # Atomic positions have to be specified by scaled positions for spglib.
    num_atom = bulk.get_number_of_atoms()
    lattice = bulk.get_cell().T.copy()
    pos = np.zeros((num_atom * 4, 3), dtype='double')
    pos[:num_atom] = bulk.get_scaled_positions()

    numbers = np.zeros(num_atom * 4, dtype='intc')
    numbers[:num_atom] = np.array(bulk.get_atomic_numbers(), dtype='intc')
    num_atom_bravais = spg.refine_cell(lattice,
                                       pos,
                                       numbers,
                                       num_atom,
                                       symprec,
                                       angle_tolerance)

    return (lattice.T.copy(),
            pos[:num_atom_bravais].copy(),
            numbers[:num_atom_bravais].copy())

def find_primitive(bulk, symprec=1e-5, angle_tolerance=-1.0):
    """
    A primitive cell in the input cell is searched and returned
    as an object of Atoms class.
    If no primitive cell is found, (None, None, None) is returned.
    """

    # Atomic positions have to be specified by scaled positions for spglib.
    positions = bulk.get_scaled_positions().copy()
    lattice = bulk.get_cell().T.copy()
    numbers = np.array(bulk.get_atomic_numbers(), dtype='intc').copy()

    # lattice is transposed with respect to the definition of Atoms class
    num_atom_prim = spg.primitive(lattice,
                                  positions,
                                  numbers,
                                  symprec,
                                  angle_tolerance)
    if num_atom_prim > 0:
        return (lattice.T.copy(),
                positions[:num_atom_prim].copy(),
                numbers[:num_atom_prim].copy())
    else:
        return None, None, None

def get_ir_kpoints(kpoint,
                   bulk,
                   is_time_reversal=True,
                   symprec=1e-5):
    """
    Retrun irreducible kpoints
    """
    mapping = np.zeros(kpoint.shape[0], dtype='intc')
    spg.ir_kpoints(mapping,
                   kpoint,
                   bulk.get_cell().T.copy(),
                   bulk.get_scaled_positions().copy(),
                   np.array(bulk.get_atomic_numbers(), dtype='intc').copy(),
                   is_time_reversal * 1,
                   symprec)
    return mapping

def get_ir_reciprocal_mesh(mesh,
                           bulk,
                           is_shift=np.zeros(3, dtype='intc'),
                           is_time_reversal=True,
                           symprec=1e-5):
    """
    Return k-points mesh and k-point map to the irreducible k-points
    The symmetry is serched from the input cell.
    is_shift=[0, 0, 0] gives Gamma center mesh.
    """

    mapping = np.zeros(np.prod(mesh), dtype='intc')
    mesh_points = np.zeros((np.prod(mesh), 3), dtype='intc')
    spg.ir_reciprocal_mesh(mesh_points,
                           mapping,
                           np.array(mesh, dtype='intc').copy(),
                           np.array(is_shift, dtype='intc').copy(),
                           is_time_reversal * 1,
                           bulk.get_cell().T.copy(),
                           bulk.get_scaled_positions().copy(),
                           np.array(bulk.get_atomic_numbers(),
                                    dtype='intc').copy(),
                           symprec)

    return mapping, mesh_points

def relocate_BZ_grid_address(grid_address,
                             mesh,
                             reciprocal_lattice, # column vectors
                             is_shift=np.zeros(3, dtype='intc')):
    bz_grid_address = np.zeros(
        ((mesh[0] + 1) * (mesh[1] + 1) * (mesh[2] + 1), 3), dtype='intc')
    bz_map = np.zeros(
        (2 * mesh[0] - 1) * (2 * mesh[1] - 1) * (2 * mesh[2] - 1), dtype='intc')
    num_bz_ir = spg.BZ_grid_address(
        bz_grid_address,
        bz_map,
        grid_address,
        np.array(mesh, dtype='intc').copy(),
        np.array(reciprocal_lattice, dtype='double').copy(),
        np.array(is_shift, dtype='intc').copy())

    return bz_grid_address[:num_bz_ir], bz_map

def get_stabilized_reciprocal_mesh(mesh,
                                   rotations,
                                   is_shift=np.zeros(3, dtype='intc'),
                                   is_time_reversal=True,
                                   qpoints=np.array([], dtype='double')):
    """
    Return k-point map to the irreducible k-points and k-point grid points .

    The symmetry is searched from the input rotation matrices in real space.

    is_shift=[0, 0, 0] gives Gamma center mesh and the values 1 give
    half mesh distance shifts.
    """

    mapping = np.zeros(np.prod(mesh), dtype='intc')
    mesh_points = np.zeros((np.prod(mesh), 3), dtype='intc')
    qpoints = np.array(qpoints, dtype='double').copy()
    if qpoints.shape == (3,):
        qpoints = np.array([qpoints], dtype='double')
    if qpoints.shape == (0,):
        qpoints = np.array([[0, 0, 0]], dtype='double')
    spg.stabilized_reciprocal_mesh(mesh_points,
                                   mapping,
                                   np.array(mesh, dtype='intc').copy(),
                                   np.array(is_shift, dtype='intc'),
                                   is_time_reversal * 1,
                                   np.array(rotations, dtype='intc').copy(),
                                   np.array(qpoints, dtype='double'))

    return mapping, mesh_points

def get_triplets_reciprocal_mesh_at_q(fixed_grid_number,
                                      mesh,
                                      rotations,
                                      is_time_reversal=True):

    weights = np.zeros(np.prod(mesh), dtype='intc')
    third_q = np.zeros(np.prod(mesh), dtype='intc')
    mesh_points = np.zeros((np.prod(mesh), 3), dtype='intc')

    spg.triplets_reciprocal_mesh_at_q(weights,
                                      mesh_points,
                                      third_q,
                                      fixed_grid_number,
                                      np.array(mesh, dtype='intc').copy(),
                                      is_time_reversal * 1,
                                      np.array(rotations, dtype='intc').copy())

    return weights, third_q, mesh_points

def get_BZ_triplets_at_q(grid_point,
                         bz_grid_address,
                         bz_map,
                         weights,
                         mesh):
    """grid_address is overwritten."""
    num_ir_tripltes = (weights > 0).sum()
    triplets = np.zeros((num_ir_tripltes, 3), dtype='intc')
    num_ir_ret = spg.BZ_triplets_at_q(triplets,
                                      grid_point,
                                      bz_grid_address,
                                      bz_map,
                                      weights,
                                      np.array(mesh, dtype='intc').copy())
    return triplets

<|endoftext|>


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

from __future__ import division
import json
import sys
import argparse
import traceback
import collections
import numpy as np
import matplotlib.pyplot as plt

types = ('wall', 'sys', 'user', 'cpu', 'perc')
concurrency = 1
mpi_size = 1
min_time = 0.01

#=============================================================================#
def nested_dict():
    return collections.defaultdict(nested_dict)

#import matplotlib
#matplotlib.use("qt4agg")
#import matplotlib.pyplot as plt

class timing_data():

    def __init__(self):
        self.data = nested_dict()
        for key in types:
            self.data[key] = []

    def append(self, _data):
        n = 0
        for key in types:
            self.data[key].append(_data[n])
            n += 1

    def __add__(self, rhs):
        for key in types:
            self.data[key].extend(rhs.data[key])

    def reset(self):
        self.data = nested_dict()
        for key in types:
            self.data[key] = []

    def __getitem__(self, key):
        return self.data[key]

#=============================================================================#
class timing_function():

    def __init__(self):
        self.data = timing_data()

    def process(self, denom, start, stop):
        _wall = float(int(stop['wall']) - int(start['wall'])) / denom
        _user = float(int(stop['user']) - int(start['user'])) / denom
        _sys = float(int(stop['sys']) - int(start['sys'])) / denom
        _cpu = _user + _sys
        if _wall > 0.0:
            _perc = (_cpu / _wall) * 100.0
        else:
            _perc = 100.0
        if _wall > min_time:
            self.data.append([_wall, _sys, _user, _cpu, _perc])

    def __getitem__(self, key):
        return self.data[key]

    def length(self):
        return len(self.data['cpu'])

#=============================================================================#
def plot(filename, title, timing_data_dict):

    ntics = len(timing_data_dict)
    ytics = []

    if ntics == 0:
        print ('{} had no timing data less than the minimum time ({} s)'.format(filename,
                                                                            min_time))
        return()

    avgs = nested_dict()
    stds = nested_dict()
    for key in types:
        avgs[key] = []
        stds[key] = []

    for func, obj in timing_data_dict.items():
        ytics.append(func)
        for key in types:
            data = obj[key]
            avgs[key].append(np.mean(data))
            if len(data) > 1:
                stds[key].append(np.std(data))
            else:
                stds[key].append(0.0)

    # the x locations for the groups
    ind = np.arange(ntics)
    # the thickness of the bars: can also be len(x) sequence
    thickness = 0.8

    f = plt.figure()
    ax = f.add_subplot(111)
    ax.yaxis.tick_right()
    f.subplots_adjust(left=0.05, right=0.75, bottom=0.05, top=0.90)

    ytics.reverse()
    for key in types:
        avgs[key].reverse()
        stds[key].reverse()

    iter_order = ('cpu', 'wall', 'sys')
    plots = []
    lk = None
    for key in iter_order:
        data = avgs[key]
        err = stds[key]
        p = None
        if lk is None:
            p = plt.barh(ind, data, thickness, xerr=err)
        else:
            p = plt.barh(ind, data, thickness, xerr=err, bottom=lk)
        #lk = avgs[key]
        plots.append(p)

    plt.grid()
    plt.xlabel('Time [seconds]')
    plt.title('Timing report for {}'.format(title))
    plt.yticks(ind, ytics, ha='left')
    plt.setp(ax.get_yticklabels(), fontsize='smaller')
    plt.legend(plots, iter_order)
    plt.show()

#=============================================================================#
def read(filename):

    print('Opening {}...'.format(filename))
    f = open(filename, "r")
    data_0 = json.load(f)
    global concurrency
    global mpi_size

    max_level = 0
    concurrency_sum = 0
    mpi_size = len(data_0['ranks'])
    for i in range(0, len(data_0['ranks'])):
        data_1 = data_0['ranks'][i]
        concurrency_sum += int(data_1['timing_manager']['omp_concurrency'])
        for j in range(0, len(data_1['timing_manager']['timers'])):
            data_2 = data_1['timing_manager']['timers'][j]
            nlaps = int(data_2['timer.ref']['laps'])
            indent = ""
            nlevel = int(data_2['timer.level'])
            max_level = max([max_level, nlevel])

    concurrency = concurrency_sum / mpi_size
    timing_functions = nested_dict()
    for i in range(0, len(data_0['ranks'])):
        data_1 = data_0['ranks'][i]
        for j in range(0, len(data_1['timing_manager']['timers'])):
            data_2 = data_1['timing_manager']['timers'][j]
            nlaps = int(data_2['timer.ref']['laps'])
            indent = ""
            nlevel = int(data_2['timer.level'])
            for n in range(0, nlevel):
                indent = ' {}'.format(indent)
            tag = '{} {} x {}'.format(indent, data_2['timer.tag'], nlaps)

            if not tag in timing_functions:
                timing_functions[tag] = timing_function()
            timing_func = timing_functions[tag]
            data_3 = data_2['timer.ref']

            for k in range(0, nlaps):
                data_4 = data_3['raw_history'][k]
                start = data_4['start']['time_since_epoch']['count']
                stop = data_4['stop']['time_since_epoch']['count']
                timing_func.process(data_3['to_seconds_ratio_den'], start, stop)

            if timing_func.length() == 0:
                del timing_functions[tag]

    return timing_functions

#=============================================================================#
def main(args):
    global concurrency
    global mpi_size

    parser = argparse.ArgumentParser()
    parser.add_argument("-f", "--files", nargs='*', help="File input")

    args = parser.parse_args()
    print('Files: {}'.format(args.files))

    file_data = dict()
    file_title = dict()
    for filename in args.files:
        print ('Reading {}...'.format(filename))
        file_data[filename] = read(filename)
        title = filename.replace('timing_report_', '')
        title = title.replace('json', 'py')
        title = '"{}"\n@ MPI procs = {}, Threads/proc = {}'.format(title, mpi_size,
                                                                int(concurrency))
        file_title[filename] = title

    for filename, data in file_data.items():
        print ('Plotting {}...'.format(filename))
        plot(filename, file_title[filename], data)

#=============================================================================#
if __name__ == "__main__":
    try:
        main(sys.argv[1:])
    except Exception as e:
        exc_type, exc_value, exc_traceback = sys.exc_info()
        traceback.print_exception(exc_type, exc_value, exc_traceback, limit=5)
        print ('Exception - {}'.format(e))

    print ('Done - {}'.format(sys.argv[0]))
    sys.exit(0)
<commit_msg>Updated timing_plot.py to new serialization format<commit_after>#!/usr/bin/env python3

from __future__ import division
import json
import sys
import argparse
import traceback
import collections
import numpy as np
import matplotlib.pyplot as plt

types = ('wall', 'sys', 'user', 'cpu', 'perc')
concurrency = 1
mpi_size = 1
min_time = 0.01

#=============================================================================#
def nested_dict():
    return collections.defaultdict(nested_dict)

#import matplotlib
#matplotlib.use("qt4agg")
#import matplotlib.pyplot as plt

class timing_data():

    def __init__(self):
        self.data = nested_dict()
        for key in types:
            self.data[key] = []

    def append(self, _data):
        n = 0
        for key in types:
            self.data[key].append(_data[n])
            n += 1

    def __add__(self, rhs):
        for key in types:
            self.data[key].extend(rhs.data[key])

    def reset(self):
        self.data = nested_dict()
        for key in types:
            self.data[key] = []

    def __getitem__(self, key):
        return self.data[key]

#=============================================================================#
class timing_function():

    def __init__(self):
        self.data = timing_data()
        self.laps = 0

    def process(self, denom, obj, nlap):
        _wall = obj['wall_elapsed'] / denom
        _user = obj['user_elapsed'] / denom
        _sys = obj['system_elapsed'] / denom
        _cpu = obj['cpu_elapsed'] / denom
        _perc = (_cpu / _wall) * 100.0 if _wall > 0.0 else 100.0
        if _wall > min_time:
            self.data.append([_wall, _sys, _user, _cpu, _perc])
        self.laps += nlap

    def __getitem__(self, key):
        return self.data[key]

    def length(self):
        return len(self.data['cpu'])

#=============================================================================#
def read(filename):

    print('Opening {}...'.format(filename))
    f = open(filename, "r")
    data_0 = json.load(f)
    global concurrency
    global mpi_size

    max_level = 0
    concurrency_sum = 0
    mpi_size = len(data_0['ranks'])
    for i in range(0, len(data_0['ranks'])):
        data_1 = data_0['ranks'][i]
        concurrency_sum += int(data_1['timing_manager']['omp_concurrency'])
        for j in range(0, len(data_1['timing_manager']['timers'])):
            data_2 = data_1['timing_manager']['timers'][j]
            nlaps = int(data_2['timer.ref']['laps'])
            indent = ""
            nlevel = int(data_2['timer.level'])
            max_level = max([max_level, nlevel])

    concurrency = concurrency_sum / mpi_size
    timing_functions = nested_dict()
    for i in range(0, len(data_0['ranks'])):
        data_1 = data_0['ranks'][i]
        for j in range(0, len(data_1['timing_manager']['timers'])):
            data_2 = data_1['timing_manager']['timers'][j]
            nlaps = int(data_2['timer.ref']['laps'])
            indent = ""
            nlevel = int(data_2['timer.level'])
            for n in range(0, nlevel):
                indent = ' {}'.format(indent)
            tag = '{} {}'.format(indent, data_2['timer.tag'])

            if not tag in timing_functions:
                timing_functions[tag] = timing_function()
            timing_func = timing_functions[tag]
            data_3 = data_2['timer.ref']
            timing_func.process(data_3['to_seconds_ratio_den'], data_3, nlaps)

            if timing_func.length() == 0:
                del timing_functions[tag]

    return timing_functions

#=============================================================================#
def plot(filename, title, timing_data_dict):

    ntics = len(timing_data_dict)
    ytics = []

    if ntics == 0:
        print ('{} had no timing data less than the minimum time ({} s)'.format(filename,
                                                                            min_time))
        return()

    avgs = nested_dict()
    stds = nested_dict()
    for key in types:
        avgs[key] = []
        stds[key] = []

    for func, obj in timing_data_dict.items():
        ytics.append('{} x [ {} counts ]'.format(func, obj.laps))
        for key in types:
            data = obj[key]
            avgs[key].append(np.mean(data))
            if len(data) > 1:
                stds[key].append(np.std(data))
            else:
                stds[key].append(0.0)

    # the x locations for the groups
    ind = np.arange(ntics)
    # the thickness of the bars: can also be len(x) sequence
    thickness = 0.8

    f = plt.figure()
    ax = f.add_subplot(111)
    ax.yaxis.tick_right()
    f.subplots_adjust(left=0.05, right=0.75, bottom=0.05, top=0.90)

    ytics.reverse()
    for key in types:
        avgs[key].reverse()
        stds[key].reverse()

    iter_order = ('cpu', 'wall', 'sys')
    plots = []
    lk = None
    for key in iter_order:
        data = avgs[key]
        err = stds[key]
        p = None
        if lk is None:
            p = plt.barh(ind, data, thickness, xerr=err)
        else:
            p = plt.barh(ind, data, thickness, xerr=err, bottom=lk)
        #lk = avgs[key]
        plots.append(p)

    plt.grid()
    plt.xlabel('Time [seconds]')
    plt.title('Timing report for {}'.format(title))
    plt.yticks(ind, ytics, ha='left')
    plt.setp(ax.get_yticklabels(), fontsize='smaller')
    plt.legend(plots, iter_order)
    plt.show()

#=============================================================================#
def main(args):
    global concurrency
    global mpi_size

    parser = argparse.ArgumentParser()
    parser.add_argument("-f", "--files", nargs='*', help="File input")

    args = parser.parse_args()
    print('Files: {}'.format(args.files))

    file_data = dict()
    file_title = dict()
    for filename in args.files:
        print ('Reading {}...'.format(filename))
        file_data[filename] = read(filename)
        title = filename.replace('timing_report_', '')
        title = title.replace('json', 'py')
        title = '"{}"\n@ MPI procs = {}, Threads/proc = {}'.format(title, mpi_size,
                                                                int(concurrency))
        file_title[filename] = title

    for filename, data in file_data.items():
        print ('Plotting {}...'.format(filename))
        plot(filename, file_title[filename], data)

#=============================================================================#
if __name__ == "__main__":
    try:
        main(sys.argv[1:])
    except Exception as e:
        exc_type, exc_value, exc_traceback = sys.exc_info()
        traceback.print_exception(exc_type, exc_value, exc_traceback, limit=5)
        print ('Exception - {}'.format(e))

    print ('Done - {}'.format(sys.argv[0]))
    sys.exit(0)
<|endoftext|>