<commit_before># -*- coding: utf-8 -*-

"""
Project Index -- prepares an index of WikiProjects and their in-scope pages for use by scripts
Copyright (C) 2015 James Hare
Licensed under MIT License: http://mitlicense.org
"""

from collections import namedtuple
import re

from BTrees.IIBTree import IIBTree
from oursql import ProgrammingError

from reportsbot.task import Task
from reportsbot.util import to_wiki_format

__all__ = ["UpdateProjectIndex"]

Project = namedtuple("Project", ["id", "title", "categories"])
Page = namedtuple("Page", ["id", "talkid", "ns", "title", "isredir"])

class UpdateProjectIndex(Task):
    """Updates the index of articles associated with each WikiProject."""

    def __init__(self, bot):
        super().__init__(bot)
        self._page_table = bot.wikiid + "_page"
        self._project_table = bot.wikiid + "_project"
        self._index_table = bot.wikiid + "_index"

    def _create_tables(self, cursor):
        """Create this wiki's various tables, using base_* as references."""
        self._logger.info("Creating index tables for %s", self._bot.wikiid)

        query1 = "DROP TABLE IF EXISTS {}"
        query2 = "CREATE TABLE {} LIKE {}"

        cursor.execute(query1.format(self._index_table))
        cursor.execute(query1.format(self._project_table))
        cursor.execute(query1.format(self._page_table))

        cursor.execute(query2.format(self._page_table, "base_page"))
        cursor.execute(query2.format(self._project_table, "base_project"))
        cursor.execute(query2.format(self._index_table, "base_index"))

    def _ensure_tables(self):
        """Ensure that all necessary tables exist for this wiki."""
        query = "SELECT 1 FROM {} LIMIT 1"
        with self._bot.localdb as cursor:
            try:
                cursor.execute(query.format(self._page_table))
                cursor.execute(query.format(self._project_table))
                cursor.execute(query.format(self._index_table))
            except ProgrammingError:
                self._create_tables(cursor)

    def _get_project_categories(self):
        """Return a list of Wikiproject classification categories."""
        query = r"""SELECT page_title
            FROM page
            WHERE page_namespace = 14 AND (
                page_title LIKE "%-Class\_%\_articles" OR
                page_title LIKE "Unassessed\_%\_articles" OR
                page_title LIKE "WikiProject\_%\_articles"
            )"""

        with self._bot.wikidb as cursor:
            cursor.execute(query)
            return [cat.decode("utf8") for (cat,) in cursor.fetchall()]

    def _find_project(self, cursor, name):
        """Return the ID and title of the given WikiProject's root page."""
        query1 = """SELECT page_id, page_namespace, page_title,
                rd_namespace, rd_title
            FROM page LEFT JOIN redirect ON rd_from = page_id
            WHERE page_title IN (?, ?) AND page_namespace = 4
            ORDER BY CHAR_LENGTH(page_title) ASC LIMIT 1"""
        query2 = """SELECT page_id
            FROM page
            WHERE page_namespace = ? AND page_title = ?"""

        candidates = ("WikiProject_%s" % name, "WikiProject_%ss" % name)
        cursor.execute(query1, candidates)

        results = cursor.fetchall()
        if not results:
            raise ValueError(name)

        pid, ns, title, rns, rtitle = results[0]
        if rns and rtitle:
            cursor.execute(query2, (rns, rtitle))
            result = cursor.fetchall()
            if not result:  # Broken redirect?
                raise ValueError(name)
            pid = result[0][0]
            ns, title = rns, rtitle

        ns_name = self._bot.site.namespaces[ns].custom_name
        fulltitle = to_wiki_format(ns_name + ":" + title.decode("utf8"))
        return pid, fulltitle

    def _get_projects(self):
        """Return a list of valid WikiProjects.

        Each project is a 3-tuple: (project_id, project_title, categories).
        """
        self._logger.info("Building project list")
        categories = self._get_project_categories()
        catmap = {}

        # Extract and map project names to their categories:
        regex = r"""^
            (?:.*?-Class_|Unassessed_)?
            (?:.*?-importance_)?
            (?:WikiProject_)?
            (.*?)                               # Actual project name fragment
            (?:-related|_task_force)?
            _articles$"""
        for cat in categories:
            name = re.search(regex, cat, flags=re.VERBOSE).group(1)
            if not name:  # Shouldn't happen, but would break us if it did
                continue
            name = name[0].upper() + name[1:]  # Normalize
            if name in catmap:
                catmap[name].append(cat)
            else:
                catmap[name] = [cat]

        # Build the Project objects:
        projects = {}
        with self._bot.wikidb as cursor:
            for name, cats in catmap.items():
                try:
                    pid, title = self._find_project(cursor, name)
                except ValueError:
                    msg = "Rejecting project: %s (%s categories, first: %s)"
                    self._logger.debug(msg, name, len(cats), cats[0])
                    continue

                if pid in projects:
                    projects[pid].categories.extend(cats)
                else:
                    projects[pid] = Project(pid, title, cats)

        msg = "%s projects with %s total categories"
        total_cats = sum(len(proj.categories) for proj in projects.values())
        self._logger.info(msg, len(projects), total_cats)
        return list(projects.values())

    def _sync_projects(self, cursor, projects):
        """Synchronize the given projects with the database."""
        self._logger.info("Synchronizing projects")

        query1 = "SELECT project_id, project_title FROM {}"
        query2 = """DELETE {0}, {1}
            FROM {0} LEFT JOIN {1} ON project_id = index_project
            WHERE project_id = ?"""
        query3 = "INSERT INTO {} (project_id, project_title) VALUES (?, ?)"
        query4 = "UPDATE {} SET project_title = ? WHERE project_id = ?"

        query1 = query1.format(self._project_table)
        query2 = query2.format(self._project_table, self._index_table)
        query3 = query3.format(self._project_table)
        query4 = query4.format(self._project_table)

        cursor.execute(query1)
        old = dict(cursor.fetchall())
        new = {proj.id: proj.title for proj in projects}

        to_remove = old.keys() - new.keys()
        to_add = new.keys() - old.keys()
        to_update = [pid for pid in new.keys() & old.keys()
                     if new[pid] != old[pid]]

        msg = "Remove/add/update: %s/%s/%s"
        self._logger.info(msg, len(to_remove), len(to_add), len(to_update))

        cursor.executemany(query2, [(pid,) for pid in to_remove])
        cursor.executemany(query3, [(pid, new[pid]) for pid in to_add])
        cursor.executemany(query4, [(new[pid], pid) for pid in to_update])

    def _get_talkmap(self, cursor):
        """Return a dict mapping talk page IDs to subject page IDs.

        Actually, this is misleading. The map also stores whether the subject
        page is a redirect, and for maximum space efficiency, we do this by
        shifting the subject page ID right by one bit and storing the boolean
        in the least significant bit.
        """
        self._logger.info("Building talkmap")

        query = """SELECT talk.page_id, subj.page_id, subj.page_is_redirect
        FROM page AS talk
        INNER JOIN page AS subj ON talk.page_title = subj.page_title
            AND talk.page_namespace = subj.page_namespace + 1
        WHERE talk.page_namespace % 2 = 1"""

        talkmap = IIBTree()
        cursor.execute(query)
        self._logger.debug("Fetching result chunks")
        while True:
            resultset = cursor.fetchmany(100000)
            if not resultset:
                break

            self._logger.debug(
                "Fetched chunk (%s+%s)", len(talkmap), len(resultset))
            for talkid, subjectid, isredir in resultset:
                talkmap[talkid] = (subjectid << 1) | isredir

        return talkmap

    def _get_pages_in_project(self, cursor, talkmap, project):
        """Return a list of Page objects within the given project.

        Each page is a 5-tuple of (subject_page_id, talk_page_id,
        subject_page_ns, title, subject_is_redirect).
        """
        query = """SELECT DISTINCT page_id, page_namespace - 1, page_title
            FROM page
            JOIN categorylinks ON cl_from = page_id
            WHERE cl_type = "page" AND page_namespace % 2 = 1
            AND cl_to IN ({})"""

        query = query.format(", ".join("?" * len(project.categories)))
        cursor.execute(query, project.categories)
        pages = [Page(talkmap[talkid] >> 1, talkid, ns, title.decode("utf8"),
                      bool(talkmap[talkid] & 1))
                 for (talkid, ns, title) in cursor.fetchall()
                 if talkid in talkmap]

        self._logger.debug("    %s pages", len(pages))
        return pages

    def _save_modified_pages(self, cursor, to_remove, to_add, to_update):
        """Update modified pages in the database."""
        msg = "    remove/add/update: %s/%s/%s"
        self._logger.debug(msg, len(to_remove), len(to_add), len(to_update))

        query1 = """DELETE {0}, {1}
            FROM {0} LEFT JOIN {1} ON page_id = index_page
            WHERE page_talk_id = ?"""
        query2 = """INSERT INTO {}
            (page_id, page_talk_id, page_title, page_ns, page_is_redirect)
            VALUES (?, ?, ?, ?, ?)"""
        query3 = """UPDATE {}
            SET page_talk_id = ?, page_title = ?, page_ns = ?,
                page_is_redirect = ?
            WHERE page_id = ?"""

        query1 = query1.format(self._page_table, self._index_table)
        query2 = query2.format(self._page_table)
        query3 = query3.format(self._page_table)

        # TODO: optimization candidates:
        cursor.executemany(query1, to_remove)
        cursor.executemany(query2, to_add)
        cursor.executemany(query3, to_update)

    def _sync_pageset(self, cursor, pages):
        """Synchronize every page in the given list with the database."""
        query = """SELECT page_id, page_talk_id, page_title, page_ns,
                page_is_redirect
            FROM {}
            WHERE page_id IN (%s)""".format(self._page_table)

        pageids = [page.id for page in pages]
        cursor.execute(query % ", ".join("?" * len(pages)), pageids)
        current = {row[0]: row[1:] for row in cursor.fetchall()}

        to_remove = []
        to_add = []
        to_update = []

        for page in pages:
            new = (page.talkid, page.title, page.ns, page.isredir)
            if page.id in current:
                cur = current[page.id]
                if new != cur:
                    if page.talkid != cur[0]:
                        # The talk page of this subject page has changed IDs,
                        # indicating that stuff has been moved around; we need
                        # to ensure that the talk page's previous entry in the
                        # database is removed (if any) before we update it
                        # since talk IDs are uniquely indexed:
                        to_remove.append((page.talkid,))
                    to_update.append(new + (page.id,))
            else:
                to_add.append((page.id,) + new)

        self._save_modified_pages(cursor, to_remove, to_add, to_update)

    def _sync_pages(self, cursor, processed, pages):
        """Synchronize the given list of pages with the database.

        Pages whose IDs are in in the processed set are skipped. The other
        pages are added to the set after processing.
        """
        unprocessed = [page for page in pages if page.id not in processed]
        self._logger.debug("    %s unprocessed to check", len(unprocessed))
        chunksize = 10000

        for start in range(0, len(unprocessed), chunksize):
            chunk = unprocessed[start:start+chunksize]
            num = int(start / chunksize) + 1
            self._logger.debug("    sync chunk #%s: %s pages", num, len(chunk))
            self._sync_pageset(cursor, chunk)

        processed |= {page.id for page in unprocessed}

    def _get_current_index(self, cursor, project):
        """Return a list of page IDs currently indexed in the given project."""
        query = "SELECT index_page FROM {} WHERE index_project = ?"
        cursor.execute(query.format(self._index_table), (project.id,))
        return [pageid for (pageid,) in cursor.fetchall()]

    def _sync_index(self, cursor, project, newpages, oldids):
        """Synchronize the index table for the given project."""
        newids = [page.id for page in newpages]

        msg = "    sync index: %s -> %s (change: %s)"
        change = len(set(oldids) ^ set(newids))
        self._logger.debug(msg, len(oldids), len(newids), change)

        query1 = "DELETE FROM {} WHERE index_page = ? AND index_project = ?"
        query2 = "INSERT INTO {} (index_page, index_project) VALUES (?, ?)"

        query1 = query1.format(self._index_table)
        query2 = query2.format(self._index_table)

        to_remove = set(oldids) - set(newids)
        to_add = set(newids) - set(oldids)

        # TODO: optimization candidates:
        cursor.executemany(query1, [(pid, project.id) for pid in to_remove])
        cursor.executemany(query2, [(pid, project.id) for pid in to_add])

    def _clear_old_pages(self, cursor, valid):
        """Remove all pages from the database that aren't in the given set."""
        query1 = "SELECT page_id FROM {}"
        query2 = "DELETE FROM {} WHERE page_id = ?"

        query1 = query1.format(self._page_table)
        query2 = query2.format(self._page_table)

        cursor.execute(query1)
        current = {pageid for (pageid,) in cursor.fetchall()}
        to_remove = current - valid

        # TODO: optimization candidate:
        cursor.executemany(query2, [(pageid,) for pageid in to_remove])

    def _sync_pages_and_index(self, cursor, projects):
        """Synchronize the database's page and index tables."""
        with self._bot.wikidb as wcursor:
            talkmap = self._get_talkmap(wcursor)
            processed = set()

            self._logger.info("Synchronizing pages and index")
            for project in projects:
                self._logger.debug("Processing: %s", project.title)
                pages = self._get_pages_in_project(wcursor, talkmap, project)
                self._sync_pages(cursor, processed, pages)
                current = self._get_current_index(cursor, project)
                self._sync_index(cursor, project, pages, current)

            self._clear_old_pages(cursor, processed)

    def run(self):
        self._ensure_tables()
        projects = self._get_projects()

        with self._bot.localdb as cursor:
            self._sync_projects(cursor, projects)
            self._sync_pages_and_index(cursor, projects)
<commit_msg>Better category fragment -> project name lookup.<commit_after># -*- coding: utf-8 -*-

"""
Project Index -- prepares an index of WikiProjects and their in-scope pages for use by scripts
Copyright (C) 2015 James Hare
Licensed under MIT License: http://mitlicense.org
"""

from collections import namedtuple
import re

from BTrees.IIBTree import IIBTree
from oursql import ProgrammingError

from reportsbot.task import Task
from reportsbot.util import to_wiki_format

__all__ = ["UpdateProjectIndex"]

Project = namedtuple("Project", ["id", "title", "categories"])
Page = namedtuple("Page", ["id", "talkid", "ns", "title", "isredir"])

class UpdateProjectIndex(Task):
    """Updates the index of articles associated with each WikiProject."""

    def __init__(self, bot):
        super().__init__(bot)
        self._page_table = bot.wikiid + "_page"
        self._project_table = bot.wikiid + "_project"
        self._index_table = bot.wikiid + "_index"

    def _create_tables(self, cursor):
        """Create this wiki's various tables, using base_* as references."""
        self._logger.info("Creating index tables for %s", self._bot.wikiid)

        query1 = "DROP TABLE IF EXISTS {}"
        query2 = "CREATE TABLE {} LIKE {}"

        cursor.execute(query1.format(self._index_table))
        cursor.execute(query1.format(self._project_table))
        cursor.execute(query1.format(self._page_table))

        cursor.execute(query2.format(self._page_table, "base_page"))
        cursor.execute(query2.format(self._project_table, "base_project"))
        cursor.execute(query2.format(self._index_table, "base_index"))

    def _ensure_tables(self):
        """Ensure that all necessary tables exist for this wiki."""
        query = "SELECT 1 FROM {} LIMIT 1"
        with self._bot.localdb as cursor:
            try:
                cursor.execute(query.format(self._page_table))
                cursor.execute(query.format(self._project_table))
                cursor.execute(query.format(self._index_table))
            except ProgrammingError:
                self._create_tables(cursor)

    def _get_project_categories(self):
        """Return a list of Wikiproject classification categories."""
        query = r"""SELECT page_title
            FROM page
            WHERE page_namespace = 14 AND (
                page_title LIKE "%-Class\_%\_articles" OR
                page_title LIKE "Unassessed\_%\_articles" OR
                page_title LIKE "WikiProject\_%\_articles"
            )"""

        with self._bot.wikidb as cursor:
            cursor.execute(query)
            return [cat.decode("utf8") for (cat,) in cursor.fetchall()]

    def _find_project(self, cursor, name):
        """Return the ID and title of the given WikiProject's root page."""
        query1 = """SELECT page_id, page_namespace, page_title,
                rd_namespace, rd_title
            FROM page LEFT JOIN redirect ON rd_from = page_id
            WHERE page_title = ? AND page_namespace = 4"""
        query2 = """SELECT page_id
            FROM page
            WHERE page_namespace = ? AND page_title = ?"""

        candidates = ("WikiProject_%s" % name, "WikiProject_%ss" % name, name)
        for candidate in candidates:
            cursor.execute(query1, candidate)
            results = cursor.fetchall()
            if results:
                break
        else:
            raise ValueError(name)

        pid, ns, title, rns, rtitle = results[0]
        if rns and rtitle:
            cursor.execute(query2, (rns, rtitle))
            result = cursor.fetchall()
            if not result:  # Broken redirect?
                raise ValueError(name)
            pid = result[0][0]
            ns, title = rns, rtitle

        ns_name = self._bot.site.namespaces[ns].custom_name
        fulltitle = to_wiki_format(ns_name + ":" + title.decode("utf8"))
        return pid, fulltitle

    def _get_projects(self):
        """Return a list of valid WikiProjects.

        Each project is a 3-tuple: (project_id, project_title, categories).
        """
        self._logger.info("Building project list")
        categories = self._get_project_categories()
        catmap = {}

        # Extract and map project names to their categories:
        regex = r"""^
            (?:.*?-Class_|Unassessed_)?
            (?:.*?-importance_)?
            (?:WikiProject_)?
            (.*?)                               # Actual project name fragment
            (?:-related|_task_force)?
            _articles$"""
        for cat in categories:
            name = re.search(regex, cat, flags=re.VERBOSE).group(1)
            if not name:  # Shouldn't happen, but would break us if it did
                continue
            name = name[0].upper() + name[1:]  # Normalize
            if name in catmap:
                catmap[name].append(cat)
            else:
                catmap[name] = [cat]

        # Build the Project objects:
        projects = {}
        with self._bot.wikidb as cursor:
            for name, cats in catmap.items():
                try:
                    pid, title = self._find_project(cursor, name)
                except ValueError:
                    msg = "Rejecting project: %s (%s categories, first: %s)"
                    self._logger.debug(msg, name, len(cats), cats[0])
                    continue

                if pid in projects:
                    projects[pid].categories.extend(cats)
                else:
                    projects[pid] = Project(pid, title, cats)

        msg = "%s projects with %s total categories"
        total_cats = sum(len(proj.categories) for proj in projects.values())
        self._logger.info(msg, len(projects), total_cats)
        return list(projects.values())

    def _sync_projects(self, cursor, projects):
        """Synchronize the given projects with the database."""
        self._logger.info("Synchronizing projects")

        query1 = "SELECT project_id, project_title FROM {}"
        query2 = """DELETE {0}, {1}
            FROM {0} LEFT JOIN {1} ON project_id = index_project
            WHERE project_id = ?"""
        query3 = "INSERT INTO {} (project_id, project_title) VALUES (?, ?)"
        query4 = "UPDATE {} SET project_title = ? WHERE project_id = ?"

        query1 = query1.format(self._project_table)
        query2 = query2.format(self._project_table, self._index_table)
        query3 = query3.format(self._project_table)
        query4 = query4.format(self._project_table)

        cursor.execute(query1)
        old = dict(cursor.fetchall())
        new = {proj.id: proj.title for proj in projects}

        to_remove = old.keys() - new.keys()
        to_add = new.keys() - old.keys()
        to_update = [pid for pid in new.keys() & old.keys()
                     if new[pid] != old[pid]]

        msg = "Remove/add/update: %s/%s/%s"
        self._logger.info(msg, len(to_remove), len(to_add), len(to_update))

        cursor.executemany(query2, [(pid,) for pid in to_remove])
        cursor.executemany(query3, [(pid, new[pid]) for pid in to_add])
        cursor.executemany(query4, [(new[pid], pid) for pid in to_update])

    def _get_talkmap(self, cursor):
        """Return a dict mapping talk page IDs to subject page IDs.

        Actually, this is misleading. The map also stores whether the subject
        page is a redirect, and for maximum space efficiency, we do this by
        shifting the subject page ID right by one bit and storing the boolean
        in the least significant bit.
        """
        self._logger.info("Building talkmap")

        query = """SELECT talk.page_id, subj.page_id, subj.page_is_redirect
        FROM page AS talk
        INNER JOIN page AS subj ON talk.page_title = subj.page_title
            AND talk.page_namespace = subj.page_namespace + 1
        WHERE talk.page_namespace % 2 = 1"""

        talkmap = IIBTree()
        cursor.execute(query)
        self._logger.debug("Fetching result chunks")
        while True:
            resultset = cursor.fetchmany(100000)
            if not resultset:
                break

            self._logger.debug(
                "Fetched chunk (%s+%s)", len(talkmap), len(resultset))
            for talkid, subjectid, isredir in resultset:
                talkmap[talkid] = (subjectid << 1) | isredir

        return talkmap

    def _get_pages_in_project(self, cursor, talkmap, project):
        """Return a list of Page objects within the given project.

        Each page is a 5-tuple of (subject_page_id, talk_page_id,
        subject_page_ns, title, subject_is_redirect).
        """
        query = """SELECT DISTINCT page_id, page_namespace - 1, page_title
            FROM page
            JOIN categorylinks ON cl_from = page_id
            WHERE cl_type = "page" AND page_namespace % 2 = 1
            AND cl_to IN ({})"""

        query = query.format(", ".join("?" * len(project.categories)))
        cursor.execute(query, project.categories)
        pages = [Page(talkmap[talkid] >> 1, talkid, ns, title.decode("utf8"),
                      bool(talkmap[talkid] & 1))
                 for (talkid, ns, title) in cursor.fetchall()
                 if talkid in talkmap]

        self._logger.debug("    %s pages", len(pages))
        return pages

    def _save_modified_pages(self, cursor, to_remove, to_add, to_update):
        """Update modified pages in the database."""
        msg = "    remove/add/update: %s/%s/%s"
        self._logger.debug(msg, len(to_remove), len(to_add), len(to_update))

        query1 = """DELETE {0}, {1}
            FROM {0} LEFT JOIN {1} ON page_id = index_page
            WHERE page_talk_id = ?"""
        query2 = """INSERT INTO {}
            (page_id, page_talk_id, page_title, page_ns, page_is_redirect)
            VALUES (?, ?, ?, ?, ?)"""
        query3 = """UPDATE {}
            SET page_talk_id = ?, page_title = ?, page_ns = ?,
                page_is_redirect = ?
            WHERE page_id = ?"""

        query1 = query1.format(self._page_table, self._index_table)
        query2 = query2.format(self._page_table)
        query3 = query3.format(self._page_table)

        # TODO: optimization candidates:
        cursor.executemany(query1, to_remove)
        cursor.executemany(query2, to_add)
        cursor.executemany(query3, to_update)

    def _sync_pageset(self, cursor, pages):
        """Synchronize every page in the given list with the database."""
        query = """SELECT page_id, page_talk_id, page_title, page_ns,
                page_is_redirect
            FROM {}
            WHERE page_id IN (%s)""".format(self._page_table)

        pageids = [page.id for page in pages]
        cursor.execute(query % ", ".join("?" * len(pages)), pageids)
        current = {row[0]: row[1:] for row in cursor.fetchall()}

        to_remove = []
        to_add = []
        to_update = []

        for page in pages:
            new = (page.talkid, page.title, page.ns, page.isredir)
            if page.id in current:
                cur = current[page.id]
                if new != cur:
                    if page.talkid != cur[0]:
                        # The talk page of this subject page has changed IDs,
                        # indicating that stuff has been moved around; we need
                        # to ensure that the talk page's previous entry in the
                        # database is removed (if any) before we update it
                        # since talk IDs are uniquely indexed:
                        to_remove.append((page.talkid,))
                    to_update.append(new + (page.id,))
            else:
                to_add.append((page.id,) + new)

        self._save_modified_pages(cursor, to_remove, to_add, to_update)

    def _sync_pages(self, cursor, processed, pages):
        """Synchronize the given list of pages with the database.

        Pages whose IDs are in in the processed set are skipped. The other
        pages are added to the set after processing.
        """
        unprocessed = [page for page in pages if page.id not in processed]
        self._logger.debug("    %s unprocessed to check", len(unprocessed))
        chunksize = 10000

        for start in range(0, len(unprocessed), chunksize):
            chunk = unprocessed[start:start+chunksize]
            num = int(start / chunksize) + 1
            self._logger.debug("    sync chunk #%s: %s pages", num, len(chunk))
            self._sync_pageset(cursor, chunk)

        processed |= {page.id for page in unprocessed}

    def _get_current_index(self, cursor, project):
        """Return a list of page IDs currently indexed in the given project."""
        query = "SELECT index_page FROM {} WHERE index_project = ?"
        cursor.execute(query.format(self._index_table), (project.id,))
        return [pageid for (pageid,) in cursor.fetchall()]

    def _sync_index(self, cursor, project, newpages, oldids):
        """Synchronize the index table for the given project."""
        newids = [page.id for page in newpages]

        msg = "    sync index: %s -> %s (change: %s)"
        change = len(set(oldids) ^ set(newids))
        self._logger.debug(msg, len(oldids), len(newids), change)

        query1 = "DELETE FROM {} WHERE index_page = ? AND index_project = ?"
        query2 = "INSERT INTO {} (index_page, index_project) VALUES (?, ?)"

        query1 = query1.format(self._index_table)
        query2 = query2.format(self._index_table)

        to_remove = set(oldids) - set(newids)
        to_add = set(newids) - set(oldids)

        # TODO: optimization candidates:
        cursor.executemany(query1, [(pid, project.id) for pid in to_remove])
        cursor.executemany(query2, [(pid, project.id) for pid in to_add])

    def _clear_old_pages(self, cursor, valid):
        """Remove all pages from the database that aren't in the given set."""
        query1 = "SELECT page_id FROM {}"
        query2 = "DELETE FROM {} WHERE page_id = ?"

        query1 = query1.format(self._page_table)
        query2 = query2.format(self._page_table)

        cursor.execute(query1)
        current = {pageid for (pageid,) in cursor.fetchall()}
        to_remove = current - valid

        # TODO: optimization candidate:
        cursor.executemany(query2, [(pageid,) for pageid in to_remove])

    def _sync_pages_and_index(self, cursor, projects):
        """Synchronize the database's page and index tables."""
        with self._bot.wikidb as wcursor:
            talkmap = self._get_talkmap(wcursor)
            processed = set()

            self._logger.info("Synchronizing pages and index")
            for project in projects:
                self._logger.debug("Processing: %s", project.title)
                pages = self._get_pages_in_project(wcursor, talkmap, project)
                self._sync_pages(cursor, processed, pages)
                current = self._get_current_index(cursor, project)
                self._sync_index(cursor, project, pages, current)

            self._clear_old_pages(cursor, processed)

    def run(self):
        self._ensure_tables()
        projects = self._get_projects()

        with self._bot.localdb as cursor:
            self._sync_projects(cursor, projects)
            self._sync_pages_and_index(cursor, projects)
<|endoftext|>


<commit_before>import django_filters

from nodeconductor.core import filters as core_filters

from . import models

class ExpertProviderFilter(django_filters.FilterSet):
    customer = core_filters.URLFilter(view_name='customer-detail', name='customer__uuid')
    customer_uuid = django_filters.UUIDFilter(name='customer__uuid')

    class Meta(object):
        model = models.ExpertProvider
        fields = []

class ExpertRequestFilter(django_filters.FilterSet):
    name = django_filters.CharFilter(lookup_expr='icontains')
    project = core_filters.URLFilter(view_name='project-detail', name='project__uuid')
    project_uuid = django_filters.UUIDFilter(name='project__uuid')

    class Meta(object):
        model = models.ExpertRequest
        fields = []

class ExpertBidFilter(django_filters.FilterSet):
    request = core_filters.URLFilter(view_name='expert-request-detail', name='request__uuid')
    request_uuid = django_filters.UUIDFilter(name='request__uuid')

    class Meta(object):
        model = models.ExpertBid
        fields = []
<commit_msg>Allow to filter expert requests by state [WAL-1041]<commit_after>import django_filters

from nodeconductor.core import filters as core_filters

from . import models

class ExpertProviderFilter(django_filters.FilterSet):
    customer = core_filters.URLFilter(view_name='customer-detail', name='customer__uuid')
    customer_uuid = django_filters.UUIDFilter(name='customer__uuid')

    class Meta(object):
        model = models.ExpertProvider
        fields = []

class ExpertRequestFilter(django_filters.FilterSet):
    name = django_filters.CharFilter(lookup_expr='icontains')
    project = core_filters.URLFilter(view_name='project-detail', name='project__uuid')
    project_uuid = django_filters.UUIDFilter(name='project__uuid')

    class Meta(object):
        model = models.ExpertRequest
        fields = ['state']

class ExpertBidFilter(django_filters.FilterSet):
    request = core_filters.URLFilter(view_name='expert-request-detail', name='request__uuid')
    request_uuid = django_filters.UUIDFilter(name='request__uuid')

    class Meta(object):
        model = models.ExpertBid
        fields = []
<|endoftext|>


<commit_before>import sys, re, os.path, numpy as np
from IPython import embed
from teafacto.util import argprun, tokenize, ticktock
from teafacto.blocks.match import MatchScore, CosineDistance, DotDistance, EuclideanDistance
from teafacto.blocks.lang.wordvec import Glove
from teafacto.blocks.seqproc import SeqEncoder, SimpleSeq2Vec
from teafacto.core.base import Block
import pickle

def run(
        epochs=10,
        numbats=100,
        negrate=1,
        lr=0.1,
        embdim=50,
        encdim=50,
        wreg=0.00005,
        marginloss=False,
        margin=1.,
        cosine=False,
        bidir=False,
    ):
    tt = ticktock("script")
    # get glove words
    g = Glove(encdim)
    words = g.D.keys()
    maxwordlen = 0
    for word in words:
        maxwordlen = max(maxwordlen, len(word))
    chars = set("".join(words))
    chars.add(" ")
    print "{} words, maxlen {}, {} characters in words".format(len(words), maxwordlen, len(chars))
    # get char word matrix
    chardic = dict(zip(chars, range(len(chars))))
    pickle.dump(chardic, open("glove2c2w.chardic.pkl", "w"))
    charwordmat = -np.ones((len(words)+1, maxwordlen), dtype="int32")
    charwordmat[0, 0] = chardic[" "]
    for i in range(0, len(words)):
        word = words[i]
        charwordmat[i+1, :len(word)] = [chardic[x] for x in word]
    print charwordmat[0]
    # encode characters
    cwenc = SimpleSeq2Vec(indim=len(chars),
                          inpembdim=embdim,
                          innerdim=encdim/2 if bidir else encdim,
                          maskid=-1,
                          bidir=bidir)
    dist = CosineDistance() if cosine else EuclideanDistance() #DotDistance()
    print "using " + str(dist)
    scorer = MatchScore(cwenc, g.block, scorer=dist)

    '''
    scorer.train([charwordmat, np.arange(len(words)+1)], np.ones((charwordmat.shape[0],), dtype="int32") * (-1 if cosine else 1))\
        .linear_objective().adagrad(lr=lr).l2(wreg)\
        .train(numbats=numbats, epochs=epochs)

    # NEGATIVE SAMPLING ::::::::::::::::::::::::::::::::::
    sys.exit()  # don't

    #embed()
    '''
    class NegIdxGen(object):
        def __init__(self, rng):
            self.min = 0
            self.max = rng

        def __call__(self, datas, gold):
            return datas, np.random.randint(self.min, self.max, gold.shape).astype("int32")

    if marginloss:
        obj = lambda p, n: (n-p+margin).clip(0, np.infty)
    else:
        obj = lambda p, n: n - p

    nscorer = scorer.nstrain([charwordmat, np.arange(len(words)+1)])\
        .negsamplegen(NegIdxGen(len(words))).negrate(negrate)\
        .objective(obj).adagrad(lr=lr).l2(wreg)\
        .train(numbats=numbats, epochs=epochs)

    cwenc.save("glove2c2w.block")

if __name__ == "__main__":
    argprun(run)<commit_msg>char2word ns<commit_after>import sys, re, os.path, numpy as np
from IPython import embed
from teafacto.util import argprun, tokenize, ticktock
from teafacto.blocks.match import MatchScore, CosineDistance, DotDistance, EuclideanDistance
from teafacto.blocks.lang.wordvec import Glove
from teafacto.blocks.seqproc import SeqEncoder, SimpleSeq2Vec
from teafacto.core.base import Block
import pickle

def run(
        epochs=10,
        numbats=100,
        negrate=1,
        lr=0.1,
        embdim=50,
        encdim=50,
        wreg=0.00005,
        marginloss=False,
        margin=1.,
        cosine=False,
        bidir=False,
    ):
    tt = ticktock("script")
    # get glove words
    g = Glove(encdim)
    words = g.D.keys()
    maxwordlen = 0
    for word in words:
        maxwordlen = max(maxwordlen, len(word))
    chars = set("".join(words))
    chars.add(" ")
    print "{} words, maxlen {}, {} characters in words".format(len(words), maxwordlen, len(chars))
    # get char word matrix
    chardic = dict(zip(chars, range(len(chars))))
    pickle.dump(chardic, open("glove2c2w.chardic.pkl", "w"))
    charwordmat = -np.ones((len(words)+1, maxwordlen), dtype="int32")
    charwordmat[0, 0] = chardic[" "]
    for i in range(0, len(words)):
        word = words[i]
        charwordmat[i+1, :len(word)] = [chardic[x] for x in word]
    print charwordmat[0]
    # encode characters
    cwenc = SimpleSeq2Vec(indim=len(chars),
                          inpembdim=embdim,
                          innerdim=encdim/2 if bidir else encdim,
                          maskid=-1,
                          bidir=bidir)
    dist = CosineDistance() if cosine else EuclideanDistance() #DotDistance()
    print "using " + str(dist)
    scorer = MatchScore(cwenc, g.block, scorer=dist)

    '''
    scorer.train([charwordmat, np.arange(len(words)+1)], np.ones((charwordmat.shape[0],), dtype="int32") * (-1 if cosine else 1))\
        .linear_objective().adagrad(lr=lr).l2(wreg)\
        .train(numbats=numbats, epochs=epochs)

    #embed()
    '''
    class NegIdxGen(object):
        def __init__(self, rng):
            self.min = 0
            self.max = rng

        def __call__(self, datas, gold):
            return datas, np.random.randint(self.min, self.max, gold.shape).astype("int32")

    if marginloss:
        obj = lambda p, n: (n - p + margin).clip(0, np.infty)
    else:
        obj = lambda p, n: n - p

    nscorer = scorer.nstrain([charwordmat, np.arange(len(words)+1)])\
        .negsamplegen(NegIdxGen(len(words))).negrate(negrate)\
        .objective(obj).adagrad(lr=lr).l2(wreg)\
        .train(numbats=numbats, epochs=epochs)

    cwenc.save("glove2c2w.block")

if __name__ == "__main__":
    argprun(run)<|endoftext|>


<commit_before>from __future__ import unicode_literals

import datetime
from decimal import Decimal

from django.conf import settings
from django.core.validators import MinValueValidator, MaxValueValidator
from django.db import models
from django.utils import timezone
from django.utils.encoding import python_2_unicode_compatible
from django.utils.translation import ugettext_lazy as _
from jsonfield import JSONField
from model_utils import FieldTracker

from nodeconductor.core import models as core_models, utils as core_utils
from nodeconductor.core.exceptions import IncorrectStateException
from nodeconductor.structure import models as structure_models

from nodeconductor_assembly_waldur.packages import models as package_models
from . import utils

@python_2_unicode_compatible
class Invoice(core_models.UuidMixin, models.Model):
    """ Invoice describes billing information about purchased packages for customers on a monthly basis """

    class Permissions(object):
        customer_path = 'customer'

    class Meta(object):
        unique_together = ('customer', 'month', 'year')

    class States(object):
        PENDING = 'pending'
        CREATED = 'created'
        PAID = 'paid'
        CANCELED = 'canceled'

        CHOICES = ((PENDING, 'Pending'), (CREATED, 'Created'), (PAID, 'Paid'), (CANCELED, 'Canceled'))

    month = models.PositiveSmallIntegerField(default=utils.get_current_month,
                                             validators=[MinValueValidator(1), MaxValueValidator(12)])
    year = models.PositiveSmallIntegerField(default=utils.get_current_year)
    state = models.CharField(max_length=30, choices=States.CHOICES, default=States.PENDING)
    customer = models.ForeignKey(structure_models.Customer, related_name='+')
    tax_percent = models.DecimalField(default=0, max_digits=4, decimal_places=2,
                                      validators=[MinValueValidator(0), MaxValueValidator(100)])
    invoice_date = models.DateField(null=True, blank=True,
                                    help_text='Date then invoice moved from state pending to created.')

    tracker = FieldTracker()

    @property
    def tax(self):
        return self.price * self.tax_percent / 100

    @property
    def total(self):
        return self.price + self.tax

    @property
    def price(self):
        return sum((item.price for item in self.openstack_items.iterator()))

    @property
    def due_date(self):
        if self.invoice_date:
            return self.invoice_date + datetime.timedelta(days=settings.INVOICES['PAYMENT_INTERVAL'])

    @property
    def number(self):
        return 100000 + self.id

    def set_created(self):
        """
        Performs following actions:
            - Freeze all invoice items
            - Change state from pending to billed
        """
        if self.state != self.States.PENDING:
            raise IncorrectStateException('Invoice must be in pending state.')

        # XXX: Consider refactoring when different types of packages will be exposed.
        items = self.openstack_items.select_related('package').all()
        for item in items:
            if item.package:
                item.freeze()

        self.state = self.States.CREATED
        self.invoice_date = timezone.now().date()
        self.save(update_fields=['state', 'invoice_date'])

    def register_package(self, package, start=None):
        if start is None:
            start = timezone.now()

        end = core_utils.month_end(start)
        overlapping_item = OpenStackItem.objects.filter(
            invoice=self,
            end__day=start.day,
        ).order_by('-daily_price').first()

        daily_price = package.template.price
        if overlapping_item:
            if overlapping_item.daily_price > daily_price:
                if overlapping_item.end.day == utils.get_current_month_end().day:
                    overlapping_item.extend_to_the_end_of_the_day()
                    return

                start = start + timezone.timedelta(days=1)
            else:
                overlapping_item.shift_backward()

        OpenStackItem.objects.create(
            package=package,
            daily_price=daily_price,
            invoice=self,
            start=start,
            end=end)

    def __str__(self):
        return '%s | %s-%s' % (self.customer, self.year, self.month)

@python_2_unicode_compatible
class OpenStackItem(models.Model):
    """ OpenStackItem stores details for invoices about purchased OpenStack packages """

    invoice = models.ForeignKey(Invoice, related_name='openstack_items')

    package = models.ForeignKey(package_models.OpenStackPackage, on_delete=models.SET_NULL, null=True, related_name='+')
    package_details = JSONField(default={}, blank=True, help_text='Stores data about package')
    daily_price = models.DecimalField(max_digits=22, decimal_places=7,
                                      validators=[MinValueValidator(Decimal('0'))],
                                      default=0,
                                      help_text='Price per day.')
    start = models.DateTimeField(default=utils.get_current_month_start,
                                 help_text='Date and time when package usage has started.')
    end = models.DateTimeField(default=utils.get_current_month_end,
                               help_text='Date and time when package usage has ended.')

    @property
    def name(self):
        if self.package:
            return '%s (%s)' % (self.package.tenant.name, self.package.template.name)

        return '%s (%s)' % (self.package_details.get('tenant_name'), self.package_details.get('template_name'))

    @property
    def tax(self):
        return self.price * self.invoice.tax_percent / 100

    @property
    def total(self):
        return self.price + self.tax

    @property
    def price(self):
        return self.daily_price * self.usage_days

    @property
    def usage_days(self):
        """
        Returns the number of days package was used from the time
        it was purchased or from the start of current month
        """
        full_days = utils.get_full_days(self.start, self.end)
        return full_days

    def freeze(self, end=None, package_deletion=False):
        """
        Performs following actions:
            - Save tenant and package template names in "package_details"
            - On package deletion set "end" field as "end" and
              recalculate price based on the new "end" field.
        """
        self.package_details['tenant_name'] = self.package.tenant.name
        self.package_details['template_name'] = self.package.template.name
        update_fields = ['package_details']

        if package_deletion:
            self.end = end or timezone.now()
            update_fields.extend(['end'])

        self.save(update_fields=update_fields)

    def shift_backward(self, days=1):
        """
        Shifts end date to N 'days' ago.
        If N is larger than it lasts - zero length will be set.
        :param days: number of days to shift end date
        """
        if (self.end - self.start).days > days:
            end = self.end - timezone.timedelta(days=1)
        else:
            end = self.start

        self.end = end
        self.save()

    def extend_to_the_end_of_the_day(self):
        self.end = self.end.replace(hour=23, minute=59, second=59)
        self.save()

    def __str__(self):
        return self.name

@python_2_unicode_compatible
class PaymentDetails(core_models.UuidMixin, models.Model):
    """ Customer payment details """

    class Permissions(object):
        customer_path = 'customer'

    class Meta(object):
        verbose_name = _('Payment details')
        verbose_name_plural = _('Payment details')

    customer = models.OneToOneField(structure_models.Customer, related_name='payment_details')
    company = models.CharField(blank=True, max_length=150)
    type = models.CharField(blank=True, max_length=150)
    address = models.CharField(blank=True, max_length=300)
    country = models.CharField(blank=True, max_length=50)
    email = models.EmailField(blank=True, max_length=75)
    postal = models.CharField(blank=True, max_length=20)
    phone = models.CharField(blank=True, max_length=20)
    bank = models.CharField(blank=True, max_length=150)
    account = models.CharField(blank=True, max_length=50)
    default_tax_percent = models.DecimalField(default=0, max_digits=4, decimal_places=2,
                                              validators=[MinValueValidator(0), MaxValueValidator(100)])

    @classmethod
    def get_url_name(cls):
        return 'payment-details'

    def __str__(self):
        return 'PaymentDetails for %s' % self.customer
<commit_msg>Document packages registration conflicts [WAL-390]<commit_after>from __future__ import unicode_literals

import datetime
from decimal import Decimal

from django.conf import settings
from django.core.validators import MinValueValidator, MaxValueValidator
from django.db import models
from django.utils import timezone
from django.utils.encoding import python_2_unicode_compatible
from django.utils.translation import ugettext_lazy as _
from jsonfield import JSONField
from model_utils import FieldTracker

from nodeconductor.core import models as core_models, utils as core_utils
from nodeconductor.core.exceptions import IncorrectStateException
from nodeconductor.structure import models as structure_models

from nodeconductor_assembly_waldur.packages import models as package_models
from . import utils

@python_2_unicode_compatible
class Invoice(core_models.UuidMixin, models.Model):
    """ Invoice describes billing information about purchased packages for customers on a monthly basis """

    class Permissions(object):
        customer_path = 'customer'

    class Meta(object):
        unique_together = ('customer', 'month', 'year')

    class States(object):
        PENDING = 'pending'
        CREATED = 'created'
        PAID = 'paid'
        CANCELED = 'canceled'

        CHOICES = ((PENDING, 'Pending'), (CREATED, 'Created'), (PAID, 'Paid'), (CANCELED, 'Canceled'))

    month = models.PositiveSmallIntegerField(default=utils.get_current_month,
                                             validators=[MinValueValidator(1), MaxValueValidator(12)])
    year = models.PositiveSmallIntegerField(default=utils.get_current_year)
    state = models.CharField(max_length=30, choices=States.CHOICES, default=States.PENDING)
    customer = models.ForeignKey(structure_models.Customer, related_name='+')
    tax_percent = models.DecimalField(default=0, max_digits=4, decimal_places=2,
                                      validators=[MinValueValidator(0), MaxValueValidator(100)])
    invoice_date = models.DateField(null=True, blank=True,
                                    help_text='Date then invoice moved from state pending to created.')

    tracker = FieldTracker()

    @property
    def tax(self):
        return self.price * self.tax_percent / 100

    @property
    def total(self):
        return self.price + self.tax

    @property
    def price(self):
        return sum((item.price for item in self.openstack_items.iterator()))

    @property
    def due_date(self):
        if self.invoice_date:
            return self.invoice_date + datetime.timedelta(days=settings.INVOICES['PAYMENT_INTERVAL'])

    @property
    def number(self):
        return 100000 + self.id

    def set_created(self):
        """
        Performs following actions:
            - Freeze all invoice items
            - Change state from pending to billed
        """
        if self.state != self.States.PENDING:
            raise IncorrectStateException('Invoice must be in pending state.')

        # XXX: Consider refactoring when different types of packages will be exposed.
        items = self.openstack_items.select_related('package').all()
        for item in items:
            if item.package:
                item.freeze()

        self.state = self.States.CREATED
        self.invoice_date = timezone.now().date()
        self.save(update_fields=['state', 'invoice_date'])

    def register_package(self, package, start=None):
        if start is None:
            start = timezone.now()

        end = core_utils.month_end(start)
        overlapping_item = OpenStackItem.objects.filter(
            invoice=self,
            end__day=start.day,
        ).order_by('-daily_price').first()

        daily_price = package.template.price
        if overlapping_item:
            """
            If there is an item that overlaps with current one as shown below:
            |--01.03.2017-|-********-|----?--|
                                     |----?--|-01.06.2017-|-******-|
            we have to make next steps:
            1) If item is more expensive -> use it for price calculation
                and register new package starting from next day [-01.06.2017-]
            |--01.03.2017-|-********-|-*****-|
                                     |-------|-01.06.2017-|-******-|

            2) If item is more expensive and it is the end of the month
            repeat step 1 but do not register new package. It will be registered from new month.
            3) If item is cheaper do exactly the opposite and shift its end date to yesterday,
            so new package will be registered today
            |--01.03.2017-|-********-|-------|
                                     |-*****-|-01.06.2017-|-******-|
            """
            if overlapping_item.daily_price > daily_price:
                if overlapping_item.end.day == utils.get_current_month_end().day:
                    overlapping_item.extend_to_the_end_of_the_day()
                    return

                start = start + timezone.timedelta(days=1)
            else:
                overlapping_item.shift_backward()

        OpenStackItem.objects.create(
            package=package,
            daily_price=daily_price,
            invoice=self,
            start=start,
            end=end)

    def __str__(self):
        return '%s | %s-%s' % (self.customer, self.year, self.month)

@python_2_unicode_compatible
class OpenStackItem(models.Model):
    """ OpenStackItem stores details for invoices about purchased OpenStack packages """

    invoice = models.ForeignKey(Invoice, related_name='openstack_items')

    package = models.ForeignKey(package_models.OpenStackPackage, on_delete=models.SET_NULL, null=True, related_name='+')
    package_details = JSONField(default={}, blank=True, help_text='Stores data about package')
    daily_price = models.DecimalField(max_digits=22, decimal_places=7,
                                      validators=[MinValueValidator(Decimal('0'))],
                                      default=0,
                                      help_text='Price per day.')
    start = models.DateTimeField(default=utils.get_current_month_start,
                                 help_text='Date and time when package usage has started.')
    end = models.DateTimeField(default=utils.get_current_month_end,
                               help_text='Date and time when package usage has ended.')

    @property
    def name(self):
        if self.package:
            return '%s (%s)' % (self.package.tenant.name, self.package.template.name)

        return '%s (%s)' % (self.package_details.get('tenant_name'), self.package_details.get('template_name'))

    @property
    def tax(self):
        return self.price * self.invoice.tax_percent / 100

    @property
    def total(self):
        return self.price + self.tax

    @property
    def price(self):
        return self.daily_price * self.usage_days

    @property
    def usage_days(self):
        """
        Returns the number of days package was used from the time
        it was purchased or from the start of current month
        """
        full_days = utils.get_full_days(self.start, self.end)
        return full_days

    def freeze(self, end=None, package_deletion=False):
        """
        Performs following actions:
            - Save tenant and package template names in "package_details"
            - On package deletion set "end" field as "end" and
              recalculate price based on the new "end" field.
        """
        self.package_details['tenant_name'] = self.package.tenant.name
        self.package_details['template_name'] = self.package.template.name
        update_fields = ['package_details']

        if package_deletion:
            self.end = end or timezone.now()
            update_fields.extend(['end'])

        self.save(update_fields=update_fields)

    def shift_backward(self, days=1):
        """
        Shifts end date to N 'days' ago.
        If N is larger than it lasts - zero length will be set.
        :param days: number of days to shift end date
        """
        if (self.end - self.start).days > days:
            end = self.end - timezone.timedelta(days=1)
        else:
            end = self.start

        self.end = end
        self.save()

    def extend_to_the_end_of_the_day(self):
        self.end = self.end.replace(hour=23, minute=59, second=59)
        self.save()

    def __str__(self):
        return self.name

@python_2_unicode_compatible
class PaymentDetails(core_models.UuidMixin, models.Model):
    """ Customer payment details """

    class Permissions(object):
        customer_path = 'customer'

    class Meta(object):
        verbose_name = _('Payment details')
        verbose_name_plural = _('Payment details')

    customer = models.OneToOneField(structure_models.Customer, related_name='payment_details')
    company = models.CharField(blank=True, max_length=150)
    type = models.CharField(blank=True, max_length=150)
    address = models.CharField(blank=True, max_length=300)
    country = models.CharField(blank=True, max_length=50)
    email = models.EmailField(blank=True, max_length=75)
    postal = models.CharField(blank=True, max_length=20)
    phone = models.CharField(blank=True, max_length=20)
    bank = models.CharField(blank=True, max_length=150)
    account = models.CharField(blank=True, max_length=50)
    default_tax_percent = models.DecimalField(default=0, max_digits=4, decimal_places=2,
                                              validators=[MinValueValidator(0), MaxValueValidator(100)])

    @classmethod
    def get_url_name(cls):
        return 'payment-details'

    def __str__(self):
        return 'PaymentDetails for %s' % self.customer
<|endoftext|>


<commit_before>from __future__ import unicode_literals

from decimal import Decimal

from django.core.validators import MinValueValidator
from django.db import models
from django.utils.encoding import python_2_unicode_compatible

from nodeconductor.core import models as core_models
from nodeconductor.structure import models as structure_models
from nodeconductor_openstack import models as openstack_models

@python_2_unicode_compatible
class PackageTemplate(core_models.UuidMixin,
                      core_models.NameMixin,
                      core_models.UiDescribableMixin):
    service_settings = models.ForeignKey(structure_models.ServiceSettings)

    @property
    def price(self):
        return self.components.aggregate(total=models.Sum(
            models.F('price') * models.F('amount'),
            output_field=models.DecimalField(max_digits=13, decimal_places=7)))['total'] or Decimal('0')

    @staticmethod
    def get_required_component_types():
        return (PackageComponent.Types.RAM,
                PackageComponent.Types.CORES,
                PackageComponent.Types.STORAGE)

    def __str__(self):
        return '%s | %s' % (self.name, self.service_settings.type)

@python_2_unicode_compatible
class PackageComponent(models.Model):
    class Meta(object):
        unique_together = ('type', 'template')

    class Types(object):
        RAM = 'ram'
        CORES = 'cores'
        STORAGE = 'storage'

        CHOICES = ((RAM, 'RAM'), (CORES, 'Cores'), (STORAGE, 'Storage'))

    type = models.CharField(max_length=50, choices=Types.CHOICES)
    amount = models.PositiveIntegerField(default=0)
    price = models.DecimalField(default=0, max_digits=13, decimal_places=7,
                                validators=[MinValueValidator(Decimal('0'))], help_text='The price per unit of amount')
    template = models.ForeignKey(PackageTemplate, related_name='components')

    def __str__(self):
        return '%s | %s' % (self.type, self.template.name)

@python_2_unicode_compatible
class OpenStackPackage(core_models.UuidMixin, models.Model):
    """ OpenStackPackage allows to create OpenStack tenant based on PackageTemplate """
    template = models.ForeignKey(PackageTemplate, help_text='Tenant will be created based on this template.')
    tenant = models.ForeignKey(openstack_models.Tenant)
    service = models.ForeignKey(openstack_models.OpenStackService, null=True)

    def __str__(self):
        return 'Package "%s" for customer %s' % (self.template, self.customer)
<commit_msg>Connect OpenStackPackage with service settings<commit_after>from __future__ import unicode_literals

from decimal import Decimal

from django.core.validators import MinValueValidator
from django.db import models
from django.utils.encoding import python_2_unicode_compatible

from nodeconductor.core import models as core_models
from nodeconductor.structure import models as structure_models
from nodeconductor_openstack import models as openstack_models

@python_2_unicode_compatible
class PackageTemplate(core_models.UuidMixin,
                      core_models.NameMixin,
                      core_models.UiDescribableMixin):
    service_settings = models.ForeignKey(structure_models.ServiceSettings, related_name='+')

    @property
    def price(self):
        return self.components.aggregate(total=models.Sum(
            models.F('price') * models.F('amount'),
            output_field=models.DecimalField(max_digits=13, decimal_places=7)))['total'] or Decimal('0')

    @staticmethod
    def get_required_component_types():
        return (PackageComponent.Types.RAM,
                PackageComponent.Types.CORES,
                PackageComponent.Types.STORAGE)

    def __str__(self):
        return '%s | %s' % (self.name, self.service_settings.type)

@python_2_unicode_compatible
class PackageComponent(models.Model):
    class Meta(object):
        unique_together = ('type', 'template')

    class Types(object):
        RAM = 'ram'
        CORES = 'cores'
        STORAGE = 'storage'

        CHOICES = ((RAM, 'RAM'), (CORES, 'Cores'), (STORAGE, 'Storage'))

    type = models.CharField(max_length=50, choices=Types.CHOICES)
    amount = models.PositiveIntegerField(default=0)
    price = models.DecimalField(default=0, max_digits=13, decimal_places=7,
                                validators=[MinValueValidator(Decimal('0'))], help_text='The price per unit of amount')
    template = models.ForeignKey(PackageTemplate, related_name='components')

    def __str__(self):
        return '%s | %s' % (self.type, self.template.name)

@python_2_unicode_compatible
class OpenStackPackage(core_models.UuidMixin, models.Model):
    """ OpenStackPackage allows to create OpenStack tenant based on PackageTemplate """
    template = models.ForeignKey(PackageTemplate, help_text='Tenant will be created based on this template.')
    tenant = models.ForeignKey(openstack_models.Tenant)
    service_settings = models.ForeignKey(structure_models.ServiceSettings, related_name='+')

    def __str__(self):
        return 'Package "%s" for customer %s' % (self.template, self.customer)
<|endoftext|>


<commit_before>import asyncio
import base64
import datetime

from .. import types, functions
from ... import events

class QRLogin:
    """
    QR login information.

    Most of the time, you will present the `url` as a QR code to the user,
    and while it's being shown, call `wait`.
    """
    def __init__(self, client, ignored_ids):
        self._client = client
        self._request = functions.auth.ExportLoginTokenRequest(
            self._client.api_id, self._client.api_hash, ignored_ids)
        self._resp = None

    async def recreate(self):
        """
        Generates a new token and URL for a new QR code, useful if the code
        has expired before it was imported.
        """
        self._resp = await self._client(self._request)

    @property
    def token(self) -> bytes:
        """
        The binary data representing the token.

        It can be used by a previously-authorized client in a call to
        :tl:`auth.importLoginToken` to log the client that originally
        requested the QR login.
        """
        return self._resp.token

    @property
    def url(self) -> str:
        """
        The ``tg://login`` URI with the token. When opened by a Telegram
        application where the user is logged in, it will import the login
        token.

        If you want to display a QR code to the user, this is the URL that
        should be launched when the QR code is scanned (the URL that should
        be contained in the QR code image you generate).

        Whether you generate the QR code image or not is up to you, and the
        library can't do this for you due to the vast ways of generating and
        displaying the QR code that exist.

        The URL simply consists of `token` base64-encoded.
        """
        return 'tg://login?token={}'.format(base64.b64encode(self._resp.token))

    @property
    def expires(self) -> datetime.datetime:
        """
        The `datetime` at which the QR code will expire.

        If you want to try again, you will need to call `recreate`.
        """
        return self._resp.expires

    async def wait(self, timeout: float = None):
        """
        Waits for the token to be imported by a previously-authorized client,
        either by scanning the QR, launching the URL directly, or calling the
        import method.

        This method **must** be called before the QR code is scanned, and
        must be executing while the QR code is being scanned. Otherwise, the
        login will not complete.

        Will raise `asyncio.TimeoutError` if the login doesn't complete on
        time.

        Arguments
            timeout (float):
                The timeout, in seconds, to wait before giving up. By default
                the library will wait until the token expires, which is often
                what you want.

        Returns
            On success, an instance of :tl:`User`. On failure it will raise.
        """
        if timeout is None:
            timeout = (self._resp.expires - datetime.datetime.now(tz=datetime.timezone.utc)).total_seconds()

        event = asyncio.Event()

        async def handler(_update):
            event.set()

        self._client.add_event_handler(handler, events.Raw(types.UpdateLoginToken))

        try:
            # Will raise timeout error if it doesn't complete quick enough,
            # which we want to let propagate
            await asyncio.wait_for(event.wait(), timeout=timeout)
        finally:
            self._client.remove_event_handler(handler)

        # We got here without it raising timeout error, so we can proceed
        resp = await self._client(self._request)
        if isinstance(resp, types.auth.LoginTokenMigrateTo):
            await self._client._switch_dc(resp.dc_id)
            resp = await self._client(functions.auth.ImportLoginTokenRequest(resp.token))
            # resp should now be auth.loginTokenSuccess

        if isinstance(resp, types.auth.LoginTokenSuccess):
            user = resp.authorization.user
            self._client._on_login(user)
            return user

        raise TypeError('Login token response was unexpected: {}'.format(resp))
<commit_msg>Fix url property in QRLogin (#1494)<commit_after>import asyncio
import base64
import datetime

from .. import types, functions
from ... import events

class QRLogin:
    """
    QR login information.

    Most of the time, you will present the `url` as a QR code to the user,
    and while it's being shown, call `wait`.
    """
    def __init__(self, client, ignored_ids):
        self._client = client
        self._request = functions.auth.ExportLoginTokenRequest(
            self._client.api_id, self._client.api_hash, ignored_ids)
        self._resp = None

    async def recreate(self):
        """
        Generates a new token and URL for a new QR code, useful if the code
        has expired before it was imported.
        """
        self._resp = await self._client(self._request)

    @property
    def token(self) -> bytes:
        """
        The binary data representing the token.

        It can be used by a previously-authorized client in a call to
        :tl:`auth.importLoginToken` to log the client that originally
        requested the QR login.
        """
        return self._resp.token

    @property
    def url(self) -> str:
        """
        The ``tg://login`` URI with the token. When opened by a Telegram
        application where the user is logged in, it will import the login
        token.

        If you want to display a QR code to the user, this is the URL that
        should be launched when the QR code is scanned (the URL that should
        be contained in the QR code image you generate).

        Whether you generate the QR code image or not is up to you, and the
        library can't do this for you due to the vast ways of generating and
        displaying the QR code that exist.

        The URL simply consists of `token` base64-encoded.
        """
        return 'tg://login?token={}'.format(base64.b64encode(self._resp.token).decode('utf-8'))

    @property
    def expires(self) -> datetime.datetime:
        """
        The `datetime` at which the QR code will expire.

        If you want to try again, you will need to call `recreate`.
        """
        return self._resp.expires

    async def wait(self, timeout: float = None):
        """
        Waits for the token to be imported by a previously-authorized client,
        either by scanning the QR, launching the URL directly, or calling the
        import method.

        This method **must** be called before the QR code is scanned, and
        must be executing while the QR code is being scanned. Otherwise, the
        login will not complete.

        Will raise `asyncio.TimeoutError` if the login doesn't complete on
        time.

        Arguments
            timeout (float):
                The timeout, in seconds, to wait before giving up. By default
                the library will wait until the token expires, which is often
                what you want.

        Returns
            On success, an instance of :tl:`User`. On failure it will raise.
        """
        if timeout is None:
            timeout = (self._resp.expires - datetime.datetime.now(tz=datetime.timezone.utc)).total_seconds()

        event = asyncio.Event()

        async def handler(_update):
            event.set()

        self._client.add_event_handler(handler, events.Raw(types.UpdateLoginToken))

        try:
            # Will raise timeout error if it doesn't complete quick enough,
            # which we want to let propagate
            await asyncio.wait_for(event.wait(), timeout=timeout)
        finally:
            self._client.remove_event_handler(handler)

        # We got here without it raising timeout error, so we can proceed
        resp = await self._client(self._request)
        if isinstance(resp, types.auth.LoginTokenMigrateTo):
            await self._client._switch_dc(resp.dc_id)
            resp = await self._client(functions.auth.ImportLoginTokenRequest(resp.token))
            # resp should now be auth.loginTokenSuccess

        if isinstance(resp, types.auth.LoginTokenSuccess):
            user = resp.authorization.user
            self._client._on_login(user)
            return user

        raise TypeError('Login token response was unexpected: {}'.format(resp))
<|endoftext|>


<commit_before># coding=utf-8
__author__ = "Daniel Arroyo. 3DaGogo, Inc <daniel@3dagogo.com>"
__license__ = 'GNU Affero General Public License http://www.gnu.org/licenses/agpl.html'

import json
import socket
import asynchat
import asyncore
import threading
import logging

from octoprint.events import eventManager, Events

from octoprint.settings import settings

class AstroprintBoxRouter(asynchat.async_chat):
	def __init__(self):
		asynchat.async_chat.__init__(self)
		self._ibuffer = []
		self.set_terminator("\n")
		self._settings = settings()
		self._logger = logging.getLogger(__name__)
		self._eventManager = eventManager()
		self.connected = False

		addr = self._settings .get(['cloudSlicer','boxrouter'])

		if addr:
			if ":" in addr:
				addr = addr.split(':')
				self._address = addr[0]
				self._port = int(addr[1])
			else:
				self._address = addr
				self._port = 80

			self.boxrouter_connect()

			self._listener = threading.Thread(target=asyncore.loop)
			self._listener.daemon = True
			self._listener.start()

		else:
			self._logger.error('boxrouter address not specified in config')

	def boxrouter_connect(self):
		try:
			self.create_socket(socket.AF_INET, socket.SOCK_STREAM)
			self.connect( (self._address, self._port) )
			self._eventManager.fire(Events.ASTROPRINT_STATUS,'connecting');
		except Exception as e:
			self._logger.error(e)

	def handle_error(self):
		self._eventManager.fire(Events.ASTROPRINT_STATUS,'error');
		self._logger.error('Unkonwn error connecting to the AstroPrint service')

	#def handle_connect(self):
	#	pass

	def handle_close(self):
		self.close()
		self.connected = False
		self._eventManager.fire(Events.ASTROPRINT_STATUS,'disconnected');

	def collect_incoming_data(self, data):
		self._ibuffer.append(data)

	def found_terminator(self):
		self.onMessage()
		self._ibuffer = []

	def onMessage(self):
		for msg in self._ibuffer:
			msg = json.loads(msg)

			self._logger.info("Received message: %s" % msg)

			if msg['type'] == 'auth':
				self.processAuthenticate(msg['data'] if 'data' in msg else None)

	def processAuthenticate(self, data):
		if data:
			if 'error' in data:
				self._logger.warn(data['message'] if 'message' in data else 'Unkonwn authentication error')
			elif 'success' in data:
				self._logger.info("Connected to astroprint service")
				self.connected = True;
				self._eventManager.fire(Events.ASTROPRINT_STATUS,'connected');

		else:
			from octoprint.server import VERSION, networkManager

			self.push(json.dumps({
				'type': 'auth',
				'data': {
					'boxId': networkManager.getMacAddress(),
					'boxName': networkManager.getHostname(),
					'swVersion': VERSION,
					'publicKey': self._settings.get(['cloudSlicer', 'publicKey']),
					'privateKey': self._settings .get(['cloudSlicer', 'privateKey'])
				}
			}))<commit_msg>retry sequence for boxrouter connection<commit_after># coding=utf-8
__author__ = "Daniel Arroyo. 3DaGogo, Inc <daniel@3dagogo.com>"
__license__ = 'GNU Affero General Public License http://www.gnu.org/licenses/agpl.html'

import json
import socket
import asynchat
import asyncore
import threading
import logging

from time import sleep
from octoprint.events import eventManager, Events

from octoprint.settings import settings

class AstroprintBoxRouter(asynchat.async_chat):
	def __init__(self):
		asynchat.async_chat.__init__(self)
		self.MAX_RETRIES = 5
		self.WAIT_BETWEEN_RETRIES = 3 #seconds
		self._ibuffer = []
		self.set_terminator("\n")
		self._settings = settings()
		self._logger = logging.getLogger(__name__)
		self._eventManager = eventManager()
		self._retries = 0
		self.connected = False

		addr = self._settings .get(['cloudSlicer','boxrouter'])

		if addr:
			if ":" in addr:
				addr = addr.split(':')
				self._address = addr[0]
				self._port = int(addr[1])
			else:
				self._address = addr
				self._port = 80

			self.boxrouter_connect()

			self._listener = threading.Thread(target=asyncore.loop)
			self._listener.daemon = True
			self._listener.start()

		else:
			self._logger.error('boxrouter address not specified in config')

	def boxrouter_connect(self):
		try:
			self.create_socket(socket.AF_INET, socket.SOCK_STREAM)
			self.connect( (self._address, self._port) )
			self._eventManager.fire(Events.ASTROPRINT_STATUS,'connecting');
		except Exception as e:
			self._logger.error(e)

	def handle_error(self):
		self._eventManager.fire(Events.ASTROPRINT_STATUS,'error');
		self._logger.error('Unkonwn error connecting to the AstroPrint service')
		self.close()
		self._doRetry()

	#def handle_connect(self):
	#	pass

	def handle_close(self):
		self.close()
		self.connected = False
		self._eventManager.fire(Events.ASTROPRINT_STATUS,'disconnected');
		self._doRetry()

	def collect_incoming_data(self, data):
		self._ibuffer.append(data)

	def found_terminator(self):
		self.onMessage()
		self._ibuffer = []

	def _doRetry(self):
		if self._retries < self.MAX_RETRIES:
			self._retries += 1
			self._logger.info('Retrying boxrouter connection. Retry #%d' % self._retries)
			sleep(self.WAIT_BETWEEN_RETRIES)
			self.boxrouter_connect()
		else:
			self._logger.info('No more retries. Giving up...')

	def onMessage(self):
		for msg in self._ibuffer:
			msg = json.loads(msg)

			self._logger.info("Received message: %s" % msg)

			if msg['type'] == 'auth':
				self.processAuthenticate(msg['data'] if 'data' in msg else None)

	def processAuthenticate(self, data):
		if data:
			if 'error' in data:
				self._logger.warn(data['message'] if 'message' in data else 'Unkonwn authentication error')
			elif 'success' in data:
				self._logger.info("Connected to astroprint service")
				self.connected = True;
				self._retries = 0;
				self._eventManager.fire(Events.ASTROPRINT_STATUS,'connected');

		else:
			from octoprint.server import VERSION, networkManager

			self.push(json.dumps({
				'type': 'auth',
				'data': {
					'boxId': networkManager.getMacAddress(),
					'boxName': networkManager.getHostname(),
					'swVersion': VERSION,
					'publicKey': self._settings.get(['cloudSlicer', 'publicKey']),
					'privateKey': self._settings .get(['cloudSlicer', 'privateKey'])
				}
			}))<|endoftext|>


<commit_before>from scipy.optimize import curve_fit
import numpy as np
from numpy.fft import fft
from scipy.linalg import svd, eig, inv, pinv
from enum import Enum

def hilbert(signal):
    # construct the Hilbert transform of the signal via the FFT
    # in essense, we just want to set negative frequency components to zero
    spectrum = np.fft.fft(signal)
    n = len(signal)
    midpoint = int(np.ceil(n/2))

    kernel = np.zeros(n)
    kernel[0] = 1
    if n%2 == 0:
        kernel[midpoint] = 1
    kernel[1:midpoint] = 2
    return np.fft.ifft(kernel * spectrum)

def KT_estimation(data, times, order):
    # Find the hilbert transform
    analytic_signal = hilbert(data)
    time_step = times[1]-times[0]

    # Create the Hankel matrix
    N = len(analytic_signal)
    K = order
    M = (N//2)-1
    L = N-M+1
    H = np.zeros((L, M), dtype=np.complex128)
    for ct in range(M):
        H[:,ct] = analytic_signal[ct:ct+L]

    #Try and seperate the signal and noise subspace via the svd
    U,S,V = svd(H, False) # V is not transposed/conjugated in numpy svd

    #Reconstruct the approximate Hankel matrix with the first K singular values
    #Here we can iterate and modify the singular values
    S_k = np.diag(S[:K])

    #Estimate the variance from the rest of the singular values
    varEst = (1/((M-K)*L)) * np.sum(S[K:]**2)
    Sfilt = np.matmul(S_k**2 - L*varEst*np.eye(K), inv(S_k))
    Hbar = np.matmul(np.matmul(U[:,:K], Sfilt), V[:K,:])

    #Reconstruct the data from the averaged anti-diagonals
    cleanedData = np.zeros(N, dtype=np.complex128)
    tmpMat = np.flip(Hbar,1)
    idx = -L+1
    for ct in range(N-1,-1,-1):
        cleanedData[ct] = np.mean(np.diag(tmpMat,idx))
        idx += 1

    #Create a cleaned Hankel matrix
    cleanedH = np.empty_like(H)
    cleanedAnalyticSig = hilbert(cleanedData)
    for ct in range(M):
        cleanedH[:,ct] = cleanedAnalyticSig[ct:ct+L]

    #Compute Q with total least squares
    #U_K1*Q = U_K2
    U = svd(cleanedH, False)[0]
    U_K = U[:,0:K]
    tmpMat = np.hstack((U_K[:-1,:],U_K[1:,:]))
    V = svd(tmpMat, False)[2].T.conj()
    n = np.size(U_K,1)
    V_AB = V[:n,n:]
    V_BB = V[n:,n:]
    Q = -np.matmul(V_AB,inv(V_BB)) # Julia conversion gotcha...

    #Now poles are eigenvalues of Q
    poles, _ = eig(Q)

    #Take the log and return the decay constant and frequency
    freqs = np.zeros(K)
    Tcs   = np.zeros(K)
    for ct in range(K):
        sk = np.log(poles[ct])
        freqs[ct] = np.imag(sk)/(2*np.pi*time_step)
        Tcs[ct] = -1.0/np.real(sk)*time_step

    #Refit the data to get the amplitude
    A = np.zeros((N, K), dtype=np.complex128)
    for ct in range(K):
        A[:,ct] = np.power(poles[ct], range(0,N))

    amps = np.matmul(pinv(A),cleanedData)

    return freqs, Tcs, amps

def fit_rabi(xdata, ydata):
    """Analyze Rabi amplitude data to find pi-pulse amplitude and phase offset.
        Arguments:
            xdata: ndarray of calibration amplitudes. length should be even.
            ydata: measurement amplitudes
        Returns:
            pi_amp: Fitted amplitude of pi pulsed
            offset: Fitted mixer offset
            fit_pts: Fitted points."""

    def rabi_model(x, *p):
        return p[0] - p[1]*np.cos(2*np.pi*p[2]*(x - p[3]))

    #seed Rabi frequency from largest FFT component
    N = len(ydata)
    yfft = fft(ydata)
    f_max_ind = np.argmax(np.abs(yfft[1:N//2]))
    f_0 = 0.5 * max([1, f_max_ind]) / xdata[-1]
    amp_0 = 0.5*(ydata.max() - ydata.min())
    offset_0 = np.mean(ydata)
    phase_0 = 0
    if ydata[N//2 - 1] > offset_0:
        amp_0 = -amp_0
    popt, _ = curve_fit(rabi_model, xdata, ydata, [offset_0, amp_0, f_0, phase_0])
    f_rabi = np.abs(popt[2])
    pi_amp = 0.5/f_rabi
    offset = popt[3]
    return pi_amp, offset, rabi_model(xdata, *popt)

def fit_ramsey(xdata, ydata, two_freqs = False):
    if two_freqs:
        # Initial KT estimation
        freqs, Tcs, amps = KT_estimation(ydata, xdata, 2)
        p0 = [*freqs, *amps.real, *Tcs, 0, 0, 0]
        popt, pcov = curve_fit(ramsey_2f, xdata, ydata, p0 = p0)
        fopt = [popt[0], popt[1]]
    else:
        # Initial KT estimation
        freqs, Tcs, amps = KT_estimation(ydata, xdata, 1)
        p0 = [freqs[0], amps.real[0], Tcs[0], 0, 0]
        popt, pcov = curve_fit(ramsey_1f, xdata, ydata, p0 = p0)
        fopt = [popt[0]]
    perr = np.sqrt(np.diag(pcov))
    fopt = popt[:two_freqs+1]
    ferr = perr[:two_freqs+1]
    return fopt, ferr, popt

def ramsey_1f(x, f, A, tau, phi, y0):
    return A*np.exp(-x/tau)*np.cos(2*np.pi*f*x + phi) + y0

def ramsey_2f(x, f1, f2, A1, A2, tau1, tau2, phi1, phi2, y0):
    return ramsey_1f(x, f1, A1, tau1, phi1, y0/2) + ramsey_1f(x, f2, A2, tau2, phi2, y0/2)

def fit_drag(data, DRAG_vec, pulse_vec):
    """Fit calibration curves vs DRAG parameter, for variable number of pulses"""
    num_DRAG = len(DRAG_vec)
    num_seqs = len(pulse_vec)
    xopt_vec = np.zeros(num_seqs)
    data = norm_data(data).reshape(num_DRAG, len(data)/num_DRAG)
    #first fit sine to lowest n, for the full range
    data_n = data[:, 1]
    T0 = 2*(DRAG_vec[np.argmax(data_n)] - DRAG_vec[np.argmin(data_n)]) #rough estimate of period

    p0 = [0, 1, T0, 0]
    popt, pcov = curve_fit(sinf, DRAG_vec, data_n, p0 = p0)
    x_fine = np.linspace(min(DRAG_vec), max(DRAG_vec), 1001)
    xopt_vec[0] = x_fine[np.argmin(sinf(x_fine, *popt))]

    for ct in range(1, len(pulse_vec)):
        #quadratic fit for subsequent steps, narrower range
        data_n = data[:, ct]
        p0 = [1, xopt_vec[ct-1], 0]
        #recenter for next fit
        closest_ind =np.argmin(abs(DRAG_vec - x0))
        fit_range = np.round(0.5*num_DRAG*pulse_vec[0]/pulse_vec[ct])
        curr_DRAG_vec = DRAG_vec[max(0, closest_ind - fit_range) : min(num_DRAG-1, closest_ind + fit_range)]
        reduced_data_n = data_n[max(0, closest_ind - fit_range) : min(num_DRAG-1, closest_ind + fit_range)]
        #quadratic fit
        popt, pcov = curve_fit(quadf, curr_DRAG_vec, reduced_data_n, p0 = p0)
        perr = np.sqrt(np.diag(pcov))
        x_fine = np.linspace(min(curr_DRAG_vec), max(curr_DRAG_vec), 1001)
        x0 = x_fine[np.argmin(quadf(x_fine, *popt))]
    return xopt_vec[-1], perr[1]

def sinf(x, f, A, phi, y0):
    return A*np.sin(2*np.pi*f*x + phi) + y0

def quadf(x, A, x0, b):
    return A*(x-x0)**2+b

class CR_cal_type(Enum):
    LENGTH = 1
    PHASE = 2
    AMPLITUDE = 3

def fit_CR(xpoints, data, cal_type):
    """Fit CR calibration curves for variable pulse length, phase, or amplitude"""
    data0 = data[:len(data)/2]
    data1 = data[len(data)/2:]
    x_fine = np.linspace(min(xpoints), max(xpoints), 1001)
    if cal_type == CR_cal_type.LENGTH:
        p0 = [1, 2*xpoints[-1], -np.pi/2, 0]
        popt0, _ = curve_fit(sinf, xpoints, data0, p0 = p0)
        popt1, _ = curve_fit(sinf, xpoints, data1, p0 = p0)
        #find the first zero crossing
        yfit0 = sinf(x_fine[:round(abs(popt0[1])/2/(xpoints[1]-xpoints[0]))], *popt0)
        yfit1 = sinf(x_fine[:round(abs(popt1[1])/2/(xpoints[1]-xpoints[0]))], *popt1)
        #average between the two qc states, rounded to 10 ns
        xopt = round((x_fine[np.argmin(abs(yfit0))] + x_fine[np.argmin(abs(yfit1))])/2/10e-9)*10e-9
        print('CR length = %f ns'%xopt*1e9)
    elif cal_type == CR_cal_type.PHASE:
        p0 = [1, xpoints[-1], np.pi, 0]
        popt0, _ = curve_fit(sinf, x_fine, data0, p0 = p0)
        popt1, _ = curve_fit(sinf, x_fine, data1, p0 = p0)
        #find the phase for maximum contrast
        contrast = (sinf(x_fine, *popt0) - sinf(x_fine, *popt1))/2
        print('CR contrast = %f'%max(contrast))
        xopt = x_fine[np.argmax(contrast)] - np.pi
    elif cal_type == CR_cal_type.AMPLITUDE:
        popt0 = np.polyfit(xpoints, data0, 1)
        popt1 = np.polyfit(xpoints, data1, 1)
        yfit0 = popt0[0]*x_fine+popt0[1]
        yfit1 = popt1[0]*x_fine+popt1[1]
        #average between optimum amplitudes
        xopt = -(popt0[1]/popt0[0] + popt1[1]/popt1[0])/2
        print('CR amplitude = %f'%xopt)
    return xopt, popt0, popt1
<commit_msg>Use np.linalg.leastsq instead of explicit inversion/pseudoinversion<commit_after>from scipy.optimize import curve_fit
import numpy as np
from numpy.fft import fft
from scipy.linalg import svd, eig, inv, pinv
from enum import Enum

def hilbert(signal):
    # construct the Hilbert transform of the signal via the FFT
    # in essense, we just want to set negative frequency components to zero
    spectrum = np.fft.fft(signal)
    n = len(signal)
    midpoint = int(np.ceil(n/2))

    kernel = np.zeros(n)
    kernel[0] = 1
    if n%2 == 0:
        kernel[midpoint] = 1
    kernel[1:midpoint] = 2
    return np.fft.ifft(kernel * spectrum)

def KT_estimation(data, times, order):
    # Find the hilbert transform
    analytic_signal = hilbert(data)
    time_step = times[1]-times[0]

    # Create the Hankel matrix
    N = len(analytic_signal)
    K = order
    M = (N//2)-1
    L = N-M+1
    H = np.zeros((L, M), dtype=np.complex128)
    for ct in range(M):
        H[:,ct] = analytic_signal[ct:ct+L]

    #Try and seperate the signal and noise subspace via the svd
    U,S,V = svd(H, False) # V is not transposed/conjugated in numpy svd

    #Reconstruct the approximate Hankel matrix with the first K singular values
    #Here we can iterate and modify the singular values
    S_k = np.diag(S[:K])

    #Estimate the variance from the rest of the singular values
    varEst = (1/((M-K)*L)) * np.sum(S[K:]**2)
    Sfilt = np.matmul(S_k**2 - L*varEst*np.eye(K), inv(S_k))
    Hbar = np.matmul(np.matmul(U[:,:K], Sfilt), V[:K,:])

    #Reconstruct the data from the averaged anti-diagonals
    cleanedData = np.zeros(N, dtype=np.complex128)
    tmpMat = np.flip(Hbar,1)
    idx = -L+1
    for ct in range(N-1,-1,-1):
        cleanedData[ct] = np.mean(np.diag(tmpMat,idx))
        idx += 1

    #Create a cleaned Hankel matrix
    cleanedH = np.empty_like(H)
    cleanedAnalyticSig = hilbert(cleanedData)
    for ct in range(M):
        cleanedH[:,ct] = cleanedAnalyticSig[ct:ct+L]

    #Compute Q with total least squares
    #U_K1*Q = U_K2
    U = svd(cleanedH, False)[0]
    U_K = U[:,0:K]
    tmpMat = np.hstack((U_K[:-1,:],U_K[1:,:]))
    V = svd(tmpMat, False)[2].T.conj()
    n = np.size(U_K,1)
    V_AB = V[:n,n:]
    V_BB = V[n:,n:]
    Q = np.linalg.lstsq(V_BB.conj().T, -V_AB.conj().T)[0].conj().T

    #Now poles are eigenvalues of Q
    poles, _ = eig(Q)

    #Take the log and return the decay constant and frequency
    freqs = np.zeros(K)
    Tcs   = np.zeros(K)
    for ct in range(K):
        sk = np.log(poles[ct])
        freqs[ct] = np.imag(sk)/(2*np.pi*time_step)
        Tcs[ct] = -1.0/np.real(sk)*time_step

    #Refit the data to get the amplitude
    A = np.zeros((N, K), dtype=np.complex128)
    for ct in range(K):
        A[:,ct] = np.power(poles[ct], range(0,N))

    amps = np.linalg.lstsq(A, cleanedData)[0]

    return freqs, Tcs, amps

def fit_rabi(xdata, ydata):
    """Analyze Rabi amplitude data to find pi-pulse amplitude and phase offset.
        Arguments:
            xdata: ndarray of calibration amplitudes. length should be even.
            ydata: measurement amplitudes
        Returns:
            pi_amp: Fitted amplitude of pi pulsed
            offset: Fitted mixer offset
            fit_pts: Fitted points."""

    def rabi_model(x, *p):
        return p[0] - p[1]*np.cos(2*np.pi*p[2]*(x - p[3]))

    #seed Rabi frequency from largest FFT component
    N = len(ydata)
    yfft = fft(ydata)
    f_max_ind = np.argmax(np.abs(yfft[1:N//2]))
    f_0 = 0.5 * max([1, f_max_ind]) / xdata[-1]
    amp_0 = 0.5*(ydata.max() - ydata.min())
    offset_0 = np.mean(ydata)
    phase_0 = 0
    if ydata[N//2 - 1] > offset_0:
        amp_0 = -amp_0
    popt, _ = curve_fit(rabi_model, xdata, ydata, [offset_0, amp_0, f_0, phase_0])
    f_rabi = np.abs(popt[2])
    pi_amp = 0.5/f_rabi
    offset = popt[3]
    return pi_amp, offset, rabi_model(xdata, *popt)

def fit_ramsey(xdata, ydata, two_freqs = False):
    if two_freqs:
        # Initial KT estimation
        freqs, Tcs, amps = KT_estimation(ydata, xdata, 2)
        p0 = [*freqs, *amps.real, *Tcs, 0, 0, 0]
        popt, pcov = curve_fit(ramsey_2f, xdata, ydata, p0 = p0)
        fopt = [popt[0], popt[1]]
    else:
        # Initial KT estimation
        freqs, Tcs, amps = KT_estimation(ydata, xdata, 1)
        p0 = [freqs[0], amps.real[0], Tcs[0], 0, 0]
        popt, pcov = curve_fit(ramsey_1f, xdata, ydata, p0 = p0)
        fopt = [popt[0]]
    perr = np.sqrt(np.diag(pcov))
    fopt = popt[:two_freqs+1]
    ferr = perr[:two_freqs+1]
    return fopt, ferr, popt

def ramsey_1f(x, f, A, tau, phi, y0):
    return A*np.exp(-x/tau)*np.cos(2*np.pi*f*x + phi) + y0

def ramsey_2f(x, f1, f2, A1, A2, tau1, tau2, phi1, phi2, y0):
    return ramsey_1f(x, f1, A1, tau1, phi1, y0/2) + ramsey_1f(x, f2, A2, tau2, phi2, y0/2)

def fit_drag(data, DRAG_vec, pulse_vec):
    """Fit calibration curves vs DRAG parameter, for variable number of pulses"""
    num_DRAG = len(DRAG_vec)
    num_seqs = len(pulse_vec)
    xopt_vec = np.zeros(num_seqs)
    data = norm_data(data).reshape(num_DRAG, len(data)/num_DRAG)
    #first fit sine to lowest n, for the full range
    data_n = data[:, 1]
    T0 = 2*(DRAG_vec[np.argmax(data_n)] - DRAG_vec[np.argmin(data_n)]) #rough estimate of period

    p0 = [0, 1, T0, 0]
    popt, pcov = curve_fit(sinf, DRAG_vec, data_n, p0 = p0)
    x_fine = np.linspace(min(DRAG_vec), max(DRAG_vec), 1001)
    xopt_vec[0] = x_fine[np.argmin(sinf(x_fine, *popt))]

    for ct in range(1, len(pulse_vec)):
        #quadratic fit for subsequent steps, narrower range
        data_n = data[:, ct]
        p0 = [1, xopt_vec[ct-1], 0]
        #recenter for next fit
        closest_ind =np.argmin(abs(DRAG_vec - x0))
        fit_range = np.round(0.5*num_DRAG*pulse_vec[0]/pulse_vec[ct])
        curr_DRAG_vec = DRAG_vec[max(0, closest_ind - fit_range) : min(num_DRAG-1, closest_ind + fit_range)]
        reduced_data_n = data_n[max(0, closest_ind - fit_range) : min(num_DRAG-1, closest_ind + fit_range)]
        #quadratic fit
        popt, pcov = curve_fit(quadf, curr_DRAG_vec, reduced_data_n, p0 = p0)
        perr = np.sqrt(np.diag(pcov))
        x_fine = np.linspace(min(curr_DRAG_vec), max(curr_DRAG_vec), 1001)
        x0 = x_fine[np.argmin(quadf(x_fine, *popt))]
    return xopt_vec[-1], perr[1]

def sinf(x, f, A, phi, y0):
    return A*np.sin(2*np.pi*f*x + phi) + y0

def quadf(x, A, x0, b):
    return A*(x-x0)**2+b

class CR_cal_type(Enum):
    LENGTH = 1
    PHASE = 2
    AMPLITUDE = 3

def fit_CR(xpoints, data, cal_type):
    """Fit CR calibration curves for variable pulse length, phase, or amplitude"""
    data0 = data[:len(data)/2]
    data1 = data[len(data)/2:]
    x_fine = np.linspace(min(xpoints), max(xpoints), 1001)
    if cal_type == CR_cal_type.LENGTH:
        p0 = [1, 2*xpoints[-1], -np.pi/2, 0]
        popt0, _ = curve_fit(sinf, xpoints, data0, p0 = p0)
        popt1, _ = curve_fit(sinf, xpoints, data1, p0 = p0)
        #find the first zero crossing
        yfit0 = sinf(x_fine[:round(abs(popt0[1])/2/(xpoints[1]-xpoints[0]))], *popt0)
        yfit1 = sinf(x_fine[:round(abs(popt1[1])/2/(xpoints[1]-xpoints[0]))], *popt1)
        #average between the two qc states, rounded to 10 ns
        xopt = round((x_fine[np.argmin(abs(yfit0))] + x_fine[np.argmin(abs(yfit1))])/2/10e-9)*10e-9
        print('CR length = %f ns'%xopt*1e9)
    elif cal_type == CR_cal_type.PHASE:
        p0 = [1, xpoints[-1], np.pi, 0]
        popt0, _ = curve_fit(sinf, x_fine, data0, p0 = p0)
        popt1, _ = curve_fit(sinf, x_fine, data1, p0 = p0)
        #find the phase for maximum contrast
        contrast = (sinf(x_fine, *popt0) - sinf(x_fine, *popt1))/2
        print('CR contrast = %f'%max(contrast))
        xopt = x_fine[np.argmax(contrast)] - np.pi
    elif cal_type == CR_cal_type.AMPLITUDE:
        popt0 = np.polyfit(xpoints, data0, 1)
        popt1 = np.polyfit(xpoints, data1, 1)
        yfit0 = popt0[0]*x_fine+popt0[1]
        yfit1 = popt1[0]*x_fine+popt1[1]
        #average between optimum amplitudes
        xopt = -(popt0[1]/popt0[0] + popt1[1]/popt1[0])/2
        print('CR amplitude = %f'%xopt)
    return xopt, popt0, popt1
<|endoftext|>


<commit_before>import os
from tempfile import TemporaryFile
from copy import copy
from urlparse import urlparse

from boto.s3.key import Key
from boto.s3.connection import S3Connection

from . import StorageBackend

class S3StorageBackend(StorageBackend):

    def __init__(self, url=None, aws_key=None, aws_secret=None,
                 s3connection=None, bucket=None, prefix=None, **kwargs):

        super(S3StorageBackend, self).__init__(**kwargs)

        assert s3connection or (aws_key and aws_secret)
        self._conn = s3connection or S3Connection(aws_key, aws_secret)

        assert bucket or url
        bucket_name = bucket or urlparse(url).netloc

        self._bucket = self._conn.get_bucket(bucket_name)
        self._prefix = prefix

    @classmethod
    def valid_url(cls, url):
        return urlparse(url).scheme in ('s3')

    def bind_to_catalog(self, id=None):
        assert id
        cpy = copy(self)
        cpy._prefix = id
        return cpy

    def _get_keyname(self, subpath):
        if self._prefix:
            return os.path.join(self._prefix, subpath)
        else:
            return subpath

    def get(self, subpath):
        key = self._bucket.get_key(self._get_keyname(subpath))
        if key is not None:
            t = TemporaryFile()
            t.write(key.read())
            t.seek(0)
            return t
        else:
            return None

    def get_meta(self, subpath):
        key = self._bucket.lookup(self._get_keyname(subpath))
        if key is None:
            return None
        meta = dict()
        meta['size'] = key.size
        return meta

    def put(self, subpath, fileobj, max_age=None, **kwargs):
        k = Key(self._bucket)
        k.key = self._get_keyname(subpath)
        if max_age is not None:
            k.set_metadata('Cache-Control', 'max-age=%d' % (max_age))
        k.set_contents_from_file(fileobj)

    def list(self, prefix=None):
        contents = []
        subpath = self._get_keyname(prefix)
        for k in self._bucket.list(prefix=subpath):
            if k.name.endswith("/"):
                # skip "directory" keys
                continue
            rel_path = k.name[len(subpath) + 1:]
            contents.append(rel_path)
        return contents

    def delete(self, subpath):
        self._bucket.delete_key(self._get_keyname(subpath))

<commit_msg>Removing blank linke<commit_after>import os
from tempfile import TemporaryFile
from copy import copy
from urlparse import urlparse

from boto.s3.key import Key
from boto.s3.connection import S3Connection

from . import StorageBackend

class S3StorageBackend(StorageBackend):

    def __init__(self, url=None, aws_key=None, aws_secret=None,
                 s3connection=None, bucket=None, prefix=None, **kwargs):

        super(S3StorageBackend, self).__init__(**kwargs)

        assert s3connection or (aws_key and aws_secret)
        self._conn = s3connection or S3Connection(aws_key, aws_secret)

        assert bucket or url
        bucket_name = bucket or urlparse(url).netloc

        self._bucket = self._conn.get_bucket(bucket_name)
        self._prefix = prefix

    @classmethod
    def valid_url(cls, url):
        return urlparse(url).scheme in ('s3')

    def bind_to_catalog(self, id=None):
        assert id
        cpy = copy(self)
        cpy._prefix = id
        return cpy

    def _get_keyname(self, subpath):
        if self._prefix:
            return os.path.join(self._prefix, subpath)
        else:
            return subpath

    def get(self, subpath):
        key = self._bucket.get_key(self._get_keyname(subpath))
        if key is not None:
            t = TemporaryFile()
            t.write(key.read())
            t.seek(0)
            return t
        else:
            return None

    def get_meta(self, subpath):
        key = self._bucket.lookup(self._get_keyname(subpath))
        if key is None:
            return None
        meta = dict()
        meta['size'] = key.size
        return meta

    def put(self, subpath, fileobj, max_age=None, **kwargs):
        k = Key(self._bucket)
        k.key = self._get_keyname(subpath)
        if max_age is not None:
            k.set_metadata('Cache-Control', 'max-age=%d' % (max_age))
        k.set_contents_from_file(fileobj)

    def list(self, prefix=None):
        contents = []
        subpath = self._get_keyname(prefix)
        for k in self._bucket.list(prefix=subpath):
            if k.name.endswith("/"):
                # skip "directory" keys
                continue
            rel_path = k.name[len(subpath) + 1:]
            contents.append(rel_path)
        return contents

    def delete(self, subpath):
        self._bucket.delete_key(self._get_keyname(subpath))
<|endoftext|>