<commit_before>import functools
import itertools
import operator

import xp

einsum_symbols = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz'
einsum_symbols_set = set(einsum_symbols)

options = {
    'has_numpy_issue_10926': True,
}

def _concat(lists):
    return sum(lists, [])

def _prod(xs):
    return functools.reduce(operator.mul, xs, 1)

def _parse_int_subscript(sub):
    subscripts = ""
    for s in sub:
        if s is Ellipsis:
            subscripts += "@"
        elif isinstance(s, int):
            subscripts += einsum_symbols[s]
        else:
            raise TypeError("For this input type lists must contain "
                            "either int or Ellipsis")
    return subscripts

def _parse_einsum_input(operands, parse_ellipsis=True):
    """Parse einsum operands.

    This function is based on `numpy.core.einsumfunc._parse_einsum_input`
    function in NumPy 1.14.

    Returns
    -------
    input_strings : str
        Parsed input strings
    output_string : str
        Parsed output string
    operands : list of array_like
        The operands to use in the numpy contraction

    Examples
    --------
    The operand list is simplified to reduce printing:

    >>> a = np.random.rand(4, 4)
    >>> b = np.random.rand(4, 4, 4)
    >>> __parse_einsum_input(('...a,...a->...', a, b))
    ('za,xza', 'xz', [a, b])

    >>> __parse_einsum_input((a, [Ellipsis, 0], b, [Ellipsis, 0]))
    ('za,xza', 'xz', [a, b])
    """

    if len(operands) == 0:
        raise ValueError("No input operands")

    if isinstance(operands[0], str):
        subscripts = operands[0].replace(" ", "")
        operands = list(operands[1:])

        # Ensure all characters are valid
        for s in subscripts:
            if s in '.,->':
                continue
            if s not in einsum_symbols:
                raise ValueError("Character %s is not a valid symbol." % s)

        # Parse "..."
        subscripts = subscripts.replace("...", "@")
        if "." in subscripts:
            raise ValueError("Invalid Ellipses.")

        # Parse "->"
        if ("-" in subscripts) or (">" in subscripts):
            # Check for proper "->"
            invalid = (subscripts.count("-") > 1) or (subscripts.count(">") > 1)
            subscripts = subscripts.split("->")
            if invalid or len(subscripts) != 2:
                raise ValueError("Subscripts can only contain one '->'.")
            input_subscripts, output_subscript = subscripts

        else:
            input_subscripts = subscripts
            output_subscript = None

        input_subscripts = input_subscripts.split(",")
        if len(input_subscripts) != len(operands):
            raise ValueError("Number of einsum subscripts must be equal to the "
                             "number of operands.")

    else:
        tmp_operands = list(operands)
        operands = []
        input_subscripts = []
        while len(tmp_operands) >= 2:
            operands.append(tmp_operands.pop(0))
            input_subscripts.append(_parse_int_subscript(
                tmp_operands.pop(0)))
        if tmp_operands:
            output_subscript = _parse_int_subscript(tmp_operands[0])
        else:
            output_subscript = None

    return input_subscripts, output_subscript, operands

def _chr(char):
    if char < 0:
        return "...[%d]" % char
    else:
        return chr(char)

def _parse_ellipsis_subscript(subscript, ndim=None, ellipsis_len=None):
    subs = subscript.split('@')
    if len(subs) == 1:
        sub, = subs
        if ndim is not None and len(sub) != ndim:
            # raise ValueError later
            return "Einstein sum subscript %s does not contain the correct number of indices " % subs
        return list(map(ord, sub))
    elif len(subs) == 2:
        left_sub, right_sub = subs
        if ndim is not None:
            ellipsis_len = ndim - (len(left_sub) + len(right_sub))
        if ellipsis_len < 0:
            # raise ValueError later
            return "Einstein sum subscript %s...%s does not contain the correct number of indices " % (left_sub, right_sub)
        return list(itertools.chain(
            map(ord, left_sub),
            range(-ellipsis_len, 0),
            map(ord, right_sub),
        ))
    else:
        # >= 2 ellipses for an operand
        raise ValueError("Invalid Ellipses.")

def einsum(*operands):
    input_subscripts, output_subscript, operands = _parse_einsum_input(operands)

    input_subscripts = [
        _parse_ellipsis_subscript(sub, ndim=arr.ndim)
        for sub, arr in zip(input_subscripts, operands)
    ]
    for i, sub_or_err in enumerate(input_subscripts):
        if isinstance(sub_or_err, str):
            raise ValueError(sub_or_err + "for operand %d." % i)

    # Get length of each unique dimension and ensure all dimensions are correct
    dimension_dict = {}
    for tnum, term in enumerate(input_subscripts):
        sh = operands[tnum].shape
        for cnum, char in enumerate(term):
            dim = sh[cnum]
            if char in dimension_dict.keys():
                # For broadcasting cases we always want the largest dim size
                if dimension_dict[char] == 1:
                    dimension_dict[char] = dim
                elif dim not in (1, dimension_dict[char]):
                    dim_old = dimension_dict[char]
                    raise ValueError("Size of label '%s' for operand %d (%d) "
                                     "does not match previous terms (%d)."
                                     % (_chr(char), tnum, dim, dim_old))
            else:
                dimension_dict[char] = dim

    if output_subscript is None:
        # Build output subscripts
        tmp_subscripts = _concat(input_subscripts)
        output_subscript = [
            s
            for s in sorted(set(tmp_subscripts))
            if s < 0 or tmp_subscripts.count(s) == 1
        ]
    else:
        if options['has_numpy_issue_10926']:
            if '@' not in output_subscript and -1 in dimension_dict:
                raise ValueError("output had too few broadcast dimensions")
        output_subscript = _parse_ellipsis_subscript(
            output_subscript,
            ellipsis_len=len(list(s for s in dimension_dict.keys() if s < 0))
        )

        # Make sure output subscripts are in the input
        tmp_subscripts = set(_concat(input_subscripts))
        for char in output_subscript:
            if char not in tmp_subscripts:
                raise ValueError(
                    "Output character %s did not appear in the input" % _chr(char))

    path = [(0, 1)] * (len(operands) - 1)  # TODO(kataoka): optimize

    # diagonal
    for num, sub in enumerate(input_subscripts):
        i = 0
        while i < len(sub):
            s = sub[i]
            if sub.count(s) > 1:
                indices = []
                sub = []
                for j, t in enumerate(input_subscripts[num]):
                    if j == i:
                        indices.append(j)
                        sub.append(t)
                    elif t == s:
                        indices.append(j)
                    else:
                        sub.append(t)
                input_subscripts[num] = sub

                diag_ndim = len(indices)
                op = operands[num]
                dim = op.shape[i]
                op = xp.moveaxis(
                    op,
                    tuple(indices), tuple(range(diag_ndim))
                )
                operands[num] = xp.moveaxis(
                    op[(xp.arange(dim),) * diag_ndim],
                    0, i
                )
            del s
            i += 1
        del i

    # unary sum
    for num, sub in enumerate(input_subscripts):
        other_subscripts = input_subscripts.copy()
        other_subscripts[num] = output_subscript
        other_subscripts = _concat(other_subscripts)
        sum_axes = tuple(
            i
            for i, s in enumerate(sub)
            if s not in other_subscripts
        )
        if sum_axes:
            input_subscripts[num] = [
                s
                for i, s in enumerate(sub)
                if i not in sum_axes
            ]
            op = operands[num]
            operands[num] = op.sum(axis=sum_axes).astype(op.dtype)

    for num in range(len(operands)):
        op = operands[num]
        if 1 in op.shape:
            squeeze_indices = []
            sub = []
            for i, s in enumerate(input_subscripts[num]):
                if op.shape[i] == 1:
                    squeeze_indices.append(i)
                else:
                    sub.append(s)
            input_subscripts[num] = sub
            operands[num] = xp.squeeze(op, axis=tuple(squeeze_indices))

    """
    count_dict = {k: 0 for k in dimension_dict}
    for sub in input_subscripts:
        for s in sub:
            count_dict[s] += 1
    """

    for idx0, idx1 in path:
        # repeat binary einsum
        assert idx0 < idx1
        sub1 = input_subscripts.pop(idx1)
        op1 = operands.pop(idx1)
        sub0 = input_subscripts.pop(idx0)
        op0 = operands.pop(idx0)

        set0 = set(sub0)
        set1 = set(sub1)
        assert len(set0) == len(sub0)
        assert len(set1) == len(sub1)

        set_out = set(_concat([output_subscript] + input_subscripts))
        shared = set0 & set1
        batch_dims = shared & set_out
        contract_dims = shared - batch_dims

        bs0, cs0, ts0 = _make_transpose_axes(sub0, batch_dims, contract_dims)
        bs1, cs1, ts1 = _make_transpose_axes(sub1, batch_dims, contract_dims)

        batch_size = _prod([dimension_dict[s] for s in batch_dims])
        contract_size = _prod([dimension_dict[s] for s in contract_dims])

        tmp0 = op0.transpose(bs0 + ts0 + cs0).reshape(batch_size, -1, contract_size)
        tmp1 = op1.transpose(bs1 + cs1 + ts1).reshape(batch_size, contract_size, -1)
        tmp_out = xp.matmul(tmp0, tmp1)

        sub_b = [sub0[i] for i in bs0]
        assert sub_b == [sub1[i] for i in bs1]
        sub_l = [sub0[i] for i in ts0]
        sub_r = [sub1[i] for i in ts1]

        sub_out = sub_b + sub_l + sub_r
        op_out = tmp_out.reshape([dimension_dict[s] for s in sub_out])

        input_subscripts.append(sub_out)
        operands.append(op_out)

    # unary einsum at last
    op0, = operands
    sub0, = input_subscripts

    transpose_axes = []
    for s in output_subscript:
        try:
            transpose_axes.append(sub0.index(s))
        except ValueError:
            pass

    return op0.transpose(transpose_axes).reshape([
        dimension_dict[s]
        for s in output_subscript
    ])

def _tuple_sorted_by_0(zs):
    return tuple(i for _, i in sorted(zs))

def _make_transpose_axes(sub, b_dims, c_dims):
    bs = []
    cs = []
    ts = []
    for i, s in enumerate(sub):
        if s in b_dims:
            bs.append((s, i))
        elif s in c_dims:
            cs.append((s, i))
        else:
            ts.append((s, i))
    return (
        _tuple_sorted_by_0(bs),
        _tuple_sorted_by_0(cs),
        _tuple_sorted_by_0(ts),
    )
"""
    if position == 0:
        it = itertools.chain(sorted(bs), sorted(ts), sorted(cs))
    else:
        it = itertools.chain(sorted(bs), sorted(cs), sorted(ts))
    return tuple(i for _, i in it)
"""
<commit_msg>Apply asanyarray to operands<commit_after>import functools
import itertools
import operator

import xp

einsum_symbols = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz'
einsum_symbols_set = set(einsum_symbols)

options = {
    'has_numpy_issue_10926': True,
}

def _concat(lists):
    return sum(lists, [])

def _prod(xs):
    return functools.reduce(operator.mul, xs, 1)

def _parse_int_subscript(sub):
    subscripts = ""
    for s in sub:
        if s is Ellipsis:
            subscripts += "@"
        elif isinstance(s, int):
            subscripts += einsum_symbols[s]
        else:
            raise TypeError("For this input type lists must contain "
                            "either int or Ellipsis")
    return subscripts

def _parse_einsum_input(operands, parse_ellipsis=True):
    """Parse einsum operands.

    This function is based on `numpy.core.einsumfunc._parse_einsum_input`
    function in NumPy 1.14.

    Returns
    -------
    input_strings : str
        Parsed input strings
    output_string : str
        Parsed output string
    operands : list of array_like
        The operands to use in the numpy contraction

    Examples
    --------
    The operand list is simplified to reduce printing:

    >>> a = np.random.rand(4, 4)
    >>> b = np.random.rand(4, 4, 4)
    >>> __parse_einsum_input(('...a,...a->...', a, b))
    ('za,xza', 'xz', [a, b])

    >>> __parse_einsum_input((a, [Ellipsis, 0], b, [Ellipsis, 0]))
    ('za,xza', 'xz', [a, b])
    """

    if len(operands) == 0:
        raise ValueError("No input operands")

    if isinstance(operands[0], str):
        subscripts = operands[0].replace(" ", "")
        operands = list(map(xp.asanyarray, operands[1:]))

        # Ensure all characters are valid
        for s in subscripts:
            if s in '.,->':
                continue
            if s not in einsum_symbols:
                raise ValueError("Character %s is not a valid symbol." % s)

        # Parse "..."
        subscripts = subscripts.replace("...", "@")
        if "." in subscripts:
            raise ValueError("Invalid Ellipses.")

        # Parse "->"
        if ("-" in subscripts) or (">" in subscripts):
            # Check for proper "->"
            invalid = (subscripts.count("-") > 1) or (subscripts.count(">") > 1)
            subscripts = subscripts.split("->")
            if invalid or len(subscripts) != 2:
                raise ValueError("Subscripts can only contain one '->'.")
            input_subscripts, output_subscript = subscripts

        else:
            input_subscripts = subscripts
            output_subscript = None

        input_subscripts = input_subscripts.split(",")
        if len(input_subscripts) != len(operands):
            raise ValueError("Number of einsum subscripts must be equal to the "
                             "number of operands.")

    else:
        tmp_operands = list(operands)
        operands = []
        input_subscripts = []
        while len(tmp_operands) >= 2:
            operands.append(xp.asanyarray(tmp_operands.pop(0)))
            input_subscripts.append(_parse_int_subscript(
                tmp_operands.pop(0)))
        if tmp_operands:
            output_subscript = _parse_int_subscript(tmp_operands[0])
        else:
            output_subscript = None

    return input_subscripts, output_subscript, operands

def _chr(char):
    if char < 0:
        return "...[%d]" % char
    else:
        return chr(char)

def _parse_ellipsis_subscript(subscript, ndim=None, ellipsis_len=None):
    subs = subscript.split('@')
    if len(subs) == 1:
        sub, = subs
        if ndim is not None and len(sub) != ndim:
            # raise ValueError later
            return "Einstein sum subscript %s does not contain the correct number of indices " % subs
        return list(map(ord, sub))
    elif len(subs) == 2:
        left_sub, right_sub = subs
        if ndim is not None:
            ellipsis_len = ndim - (len(left_sub) + len(right_sub))
        if ellipsis_len < 0:
            # raise ValueError later
            return "Einstein sum subscript %s...%s does not contain the correct number of indices " % (left_sub, right_sub)
        return list(itertools.chain(
            map(ord, left_sub),
            range(-ellipsis_len, 0),
            map(ord, right_sub),
        ))
    else:
        # >= 2 ellipses for an operand
        raise ValueError("Invalid Ellipses.")

def einsum(*operands):
    input_subscripts, output_subscript, operands = _parse_einsum_input(operands)

    input_subscripts = [
        _parse_ellipsis_subscript(sub, ndim=arr.ndim)
        for sub, arr in zip(input_subscripts, operands)
    ]
    for i, sub_or_err in enumerate(input_subscripts):
        if isinstance(sub_or_err, str):
            raise ValueError(sub_or_err + "for operand %d." % i)

    # Get length of each unique dimension and ensure all dimensions are correct
    dimension_dict = {}
    for tnum, term in enumerate(input_subscripts):
        sh = operands[tnum].shape
        for cnum, char in enumerate(term):
            dim = sh[cnum]
            if char in dimension_dict.keys():
                # For broadcasting cases we always want the largest dim size
                if dimension_dict[char] == 1:
                    dimension_dict[char] = dim
                elif dim not in (1, dimension_dict[char]):
                    dim_old = dimension_dict[char]
                    raise ValueError("Size of label '%s' for operand %d (%d) "
                                     "does not match previous terms (%d)."
                                     % (_chr(char), tnum, dim, dim_old))
            else:
                dimension_dict[char] = dim

    if output_subscript is None:
        # Build output subscripts
        tmp_subscripts = _concat(input_subscripts)
        output_subscript = [
            s
            for s in sorted(set(tmp_subscripts))
            if s < 0 or tmp_subscripts.count(s) == 1
        ]
    else:
        if options['has_numpy_issue_10926']:
            if '@' not in output_subscript and -1 in dimension_dict:
                raise ValueError("output had too few broadcast dimensions")
        output_subscript = _parse_ellipsis_subscript(
            output_subscript,
            ellipsis_len=len(list(s for s in dimension_dict.keys() if s < 0))
        )

        # Make sure output subscripts are in the input
        tmp_subscripts = set(_concat(input_subscripts))
        for char in output_subscript:
            if char not in tmp_subscripts:
                raise ValueError(
                    "Output character %s did not appear in the input" % _chr(char))

    path = [(0, 1)] * (len(operands) - 1)  # TODO(kataoka): optimize

    # diagonal
    for num, sub in enumerate(input_subscripts):
        i = 0
        while i < len(sub):
            s = sub[i]
            if sub.count(s) > 1:
                indices = []
                sub = []
                for j, t in enumerate(input_subscripts[num]):
                    if j == i:
                        indices.append(j)
                        sub.append(t)
                    elif t == s:
                        indices.append(j)
                    else:
                        sub.append(t)
                input_subscripts[num] = sub

                diag_ndim = len(indices)
                op = operands[num]
                dim = op.shape[i]
                op = xp.moveaxis(
                    op,
                    tuple(indices), tuple(range(diag_ndim))
                )
                operands[num] = xp.moveaxis(
                    op[(xp.arange(dim),) * diag_ndim],
                    0, i
                )
            del s
            i += 1
        del i

    # unary sum
    for num, sub in enumerate(input_subscripts):
        other_subscripts = input_subscripts.copy()
        other_subscripts[num] = output_subscript
        other_subscripts = _concat(other_subscripts)
        sum_axes = tuple(
            i
            for i, s in enumerate(sub)
            if s not in other_subscripts
        )
        if sum_axes:
            input_subscripts[num] = [
                s
                for i, s in enumerate(sub)
                if i not in sum_axes
            ]
            op = operands[num]
            operands[num] = op.sum(axis=sum_axes).astype(op.dtype)

    for num in range(len(operands)):
        op = operands[num]
        if 1 in op.shape:
            squeeze_indices = []
            sub = []
            for i, s in enumerate(input_subscripts[num]):
                if op.shape[i] == 1:
                    squeeze_indices.append(i)
                else:
                    sub.append(s)
            input_subscripts[num] = sub
            operands[num] = xp.squeeze(op, axis=tuple(squeeze_indices))

    """
    count_dict = {k: 0 for k in dimension_dict}
    for sub in input_subscripts:
        for s in sub:
            count_dict[s] += 1
    """

    for idx0, idx1 in path:
        # repeat binary einsum
        assert idx0 < idx1
        sub1 = input_subscripts.pop(idx1)
        op1 = operands.pop(idx1)
        sub0 = input_subscripts.pop(idx0)
        op0 = operands.pop(idx0)

        set0 = set(sub0)
        set1 = set(sub1)
        assert len(set0) == len(sub0)
        assert len(set1) == len(sub1)

        set_out = set(_concat([output_subscript] + input_subscripts))
        shared = set0 & set1
        batch_dims = shared & set_out
        contract_dims = shared - batch_dims

        bs0, cs0, ts0 = _make_transpose_axes(sub0, batch_dims, contract_dims)
        bs1, cs1, ts1 = _make_transpose_axes(sub1, batch_dims, contract_dims)

        batch_size = _prod([dimension_dict[s] for s in batch_dims])
        contract_size = _prod([dimension_dict[s] for s in contract_dims])

        tmp0 = op0.transpose(bs0 + ts0 + cs0).reshape(batch_size, -1, contract_size)
        tmp1 = op1.transpose(bs1 + cs1 + ts1).reshape(batch_size, contract_size, -1)
        tmp_out = xp.matmul(tmp0, tmp1)

        sub_b = [sub0[i] for i in bs0]
        assert sub_b == [sub1[i] for i in bs1]
        sub_l = [sub0[i] for i in ts0]
        sub_r = [sub1[i] for i in ts1]

        sub_out = sub_b + sub_l + sub_r
        op_out = tmp_out.reshape([dimension_dict[s] for s in sub_out])

        input_subscripts.append(sub_out)
        operands.append(op_out)

    # unary einsum at last
    op0, = operands
    sub0, = input_subscripts

    transpose_axes = []
    for s in output_subscript:
        try:
            transpose_axes.append(sub0.index(s))
        except ValueError:
            pass

    return op0.transpose(transpose_axes).reshape([
        dimension_dict[s]
        for s in output_subscript
    ])

def _tuple_sorted_by_0(zs):
    return tuple(i for _, i in sorted(zs))

def _make_transpose_axes(sub, b_dims, c_dims):
    bs = []
    cs = []
    ts = []
    for i, s in enumerate(sub):
        if s in b_dims:
            bs.append((s, i))
        elif s in c_dims:
            cs.append((s, i))
        else:
            ts.append((s, i))
    return (
        _tuple_sorted_by_0(bs),
        _tuple_sorted_by_0(cs),
        _tuple_sorted_by_0(ts),
    )
"""
    if position == 0:
        it = itertools.chain(sorted(bs), sorted(ts), sorted(cs))
    else:
        it = itertools.chain(sorted(bs), sorted(cs), sorted(ts))
    return tuple(i for _, i in it)
"""
<|endoftext|>


<commit_before><commit_msg>Fix typo<commit_after><|endoftext|>


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

from plantmeter.isodates import isodate
from dateutil.relativedelta import relativedelta

def previousAmortizationDate(purchase_date, current_date):

    if isodate(current_date) >= isodate('2003-01-01'):
        return '2003-01-01'

    return None

def pendingAmortization(purchase_date, current_date, investment_amount, amortized_amount):

    years = relativedelta(
        isodate(current_date),
        isodate(purchase_date),
        ).years

    yearly_amortitzation = investment_amount / 25

    if years < 2:
        return 0

    if years >= 25:
        return investment_amount - amortized_amount

    toAmortize = (years-1)*yearly_amortitzation - amortized_amount
    return max(0, toAmortize)

# vim: et ts=4 sw=4
<commit_msg>refactor use var instead lit<commit_after># -*- coding:utf8 -*-

from plantmeter.isodates import isodate
from dateutil.relativedelta import relativedelta

def previousAmortizationDate(purchase_date, current_date):

    firstAmortization = '2003-01-01'
    if isodate(current_date) >= isodate(firstAmortization):
        return firstAmortization

    return None

def pendingAmortization(purchase_date, current_date, investment_amount, amortized_amount):

    years = relativedelta(
        isodate(current_date),
        isodate(purchase_date),
        ).years

    yearly_amortitzation = investment_amount / 25

    if years < 2:
        return 0

    if years >= 25:
        return investment_amount - amortized_amount

    toAmortize = (years-1)*yearly_amortitzation - amortized_amount
    return max(0, toAmortize)

# vim: et ts=4 sw=4
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
'''
.. module:: genomics.mapreduce
   :synopsis: MapReduce framework
   :noindex:
   :copyright: Copyright 2014 by Tiago Antao
   :license: GNU Affero, see LICENSE for details

.. moduleauthor:: Tiago Antao <tra@popgen.net>

'''
import os
import pickle
import tempfile

import genomics

mr_dir = genomics.cfg.mr_dir

def pickle_in(datum):
    w = tempfile.NamedTemporaryFile(dir=mr_dir, delete=False)
    pickle.dump(datum, w)
    w.close()
    return w.name

def do_map(pname):
    print(pname)
<commit_msg>map - ongoing<commit_after># -*- coding: utf-8 -*-
'''
.. module:: genomics.mapreduce
   :synopsis: MapReduce framework
   :noindex:
   :copyright: Copyright 2014 by Tiago Antao
   :license: GNU Affero, see LICENSE for details

.. moduleauthor:: Tiago Antao <tra@popgen.net>

'''
import importlib
import os
import pickle
import tempfile

import genomics

mr_dir = genomics.cfg.mr_dir

def pickle_in(datum):
    w = tempfile.NamedTemporaryFile(dir=mr_dir, delete=False)
    pickle.dump(datum, w)
    w.close()
    return w.name

def do_map(pname):
    f = open(pname, 'rb')
    paras = pickle.load(f, encoding='bytes')
    f.close()
    fqdn_name = paras[0]
    fqdn_toks = fqdn_name.split('.')
    module = importlib.import_module('.'.join(fqdn_toks[:-1]))
    fun = getattr(module, fqdn_toks[-1])
    node = paras[1]
    args = paras[2:]
    fun(node, *args)
    #os.remove(pname)
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# Author: Jayant Jain <jayant@rare-technologies.com>
# Copyright (C) 2016 RaRe Technologies

"""This script using for extracting plain text out of a raw Wikipedia dump. Input is an xml.bz2 file provided
by MediaWiki that looks like <LANG>wiki-<YYYYMMDD>-pages-articles.xml.bz2 or <LANG>wiki-latest-pages-articles.xml.bz2
(e.g. 14 GB of https://dumps.wikimedia.org/enwiki/latest/enwiki-latest-pages-articles.xml.bz2).

It streams through all the XML articles using multiple cores (#cores - 1, by default),
decompressing on the fly and extracting plain text from the articles and their sections.

For each extracted article, it prints its title, section names and plain text section contents, in json-line format.

How to use
----------
#. Process Wikipedia dump with this script ::

    python -m gensim.scripts.segment_wiki -i -f enwiki-latest-pages-articles.xml.bz2 -o enwiki-latest.json.gz

#. Read output in simple way:

.. sourcecode:: pycon

    >>> from gensim import utils
    >>> import json
    >>>
    >>> # iterate over the plain text data we just created
    >>> with utils.open('enwiki-latest.json.gz', 'rb') as f:
    >>>     for line in f:
    >>>         # decode each JSON line into a Python dictionary object
    >>>         article = json.loads(line)
    >>>
    >>>         # each article has a "title", a mapping of interlinks and a list of "section_titles" and
    >>>         # "section_texts".
    >>>         print("Article title: %s" % article['title'])
    >>>         print("Interlinks: %s" + article['interlinks'])
    >>>         for section_title, section_text in zip(article['section_titles'], article['section_texts']):
    >>>             print("Section title: %s" % section_title)
    >>>             print("Section text: %s" % section_text)

Notes
-----
Processing the entire English Wikipedia dump takes 1.7 hours (about 3 million articles per hour,
or 10 MB of XML per second) on an 8 core Intel i7-7700 @3.60GHz.

Command line arguments
----------------------

.. program-output:: python -m gensim.scripts.segment_wiki --help
   :ellipsis: 0, -10

"""

import argparse
import json
import logging
import multiprocessing
import re
import sys
from xml.etree import cElementTree
from functools import partial

from gensim.corpora.wikicorpus import IGNORED_NAMESPACES, WikiCorpus, filter_wiki, find_interlinks, get_namespace, utils
import gensim.utils

logger = logging.getLogger(__name__)

def segment_all_articles(file_path, min_article_character=200, workers=None, include_interlinks=False):
    """Extract article titles and sections from a MediaWiki bz2 database dump.

    Parameters
    ----------
    file_path : str
        Path to MediaWiki dump, typical filename is <LANG>wiki-<YYYYMMDD>-pages-articles.xml.bz2
        or <LANG>wiki-latest-pages-articles.xml.bz2.

    min_article_character : int, optional
        Minimal number of character for article (except titles and leading gaps).

    workers: int or None
        Number of parallel workers, max(1, multiprocessing.cpu_count() - 1) if None.

    include_interlinks: bool
        Whether or not interlinks should be included in the output

    Yields
    ------
    (str, list of (str, str), (Optionally) list of (str, str))
        Structure contains (title, [(section_heading, section_content), ...],
        (Optionally) [(interlink_article, interlink_text), ...]).

    """
    with gensim.utils.open(file_path, 'rb') as xml_fileobj:
        wiki_sections_corpus = _WikiSectionsCorpus(
            xml_fileobj, min_article_character=min_article_character, processes=workers,
            include_interlinks=include_interlinks)
        wiki_sections_corpus.metadata = True
        wiki_sections_text = wiki_sections_corpus.get_texts_with_sections()

        for article in wiki_sections_text:
            yield article

def segment_and_write_all_articles(file_path, output_file, min_article_character=200, workers=None,
                                   include_interlinks=False):
    """Write article title and sections to `output_file` (or stdout, if output_file is None).

    The output format is one article per line, in json-line format with 4 fields::

        'title' - title of article,
        'section_titles' - list of titles of sections,
        'section_texts' - list of content from sections,
        (Optional) 'section_interlinks' - list of interlinks in the article.

    Parameters
    ----------
    file_path : str
        Path to MediaWiki dump, typical filename is <LANG>wiki-<YYYYMMDD>-pages-articles.xml.bz2
        or <LANG>wiki-latest-pages-articles.xml.bz2.

    output_file : str or None
        Path to output file in json-lines format, or None for printing to stdout.

    min_article_character : int, optional
        Minimal number of character for article (except titles and leading gaps).

    workers: int or None
        Number of parallel workers, max(1, multiprocessing.cpu_count() - 1) if None.

    include_interlinks: bool
        Whether or not interlinks should be included in the output
    """
    if output_file is None:
        outfile = getattr(sys.stdout, 'buffer', sys.stdout)  # we want write bytes, so for py3 we used 'buffer'
    else:
        outfile = gensim.utils.open(output_file, 'wb')

    try:
        article_stream = segment_all_articles(file_path, min_article_character, workers=workers,
                                              include_interlinks=include_interlinks)
        for idx, article in enumerate(article_stream):
            article_title, article_sections = article[0], article[1]
            if include_interlinks:
                interlinks = article[2]

            output_data = {
                "title": article_title,
                "section_titles": [],
                "section_texts": [],
            }
            if include_interlinks:
                output_data["interlinks"] = interlinks

            for section_heading, section_content in article_sections:
                output_data["section_titles"].append(section_heading)
                output_data["section_texts"].append(section_content)

            if (idx + 1) % 100000 == 0:
                logger.info("processed #%d articles (at %r now)", idx + 1, article_title)
            outfile.write((json.dumps(output_data) + "\n").encode('utf-8'))

    finally:
        if output_file is not None:
            outfile.close()

def extract_page_xmls(f):
    """Extract pages from a MediaWiki database dump.

    Parameters
    ----------
    f : file
        File descriptor of MediaWiki dump.

    Yields
    ------
    str
        XML strings for page tags.

    """
    elems = (elem for _, elem in cElementTree.iterparse(f, events=("end",)))

    elem = next(elems)
    namespace = get_namespace(elem.tag)
    ns_mapping = {"ns": namespace}
    page_tag = "{%(ns)s}page" % ns_mapping

    for elem in elems:
        if elem.tag == page_tag:
            yield cElementTree.tostring(elem)
            # Prune the element tree, as per
            # http://www.ibm.com/developerworks/xml/library/x-hiperfparse/
            # except that we don't need to prune backlinks from the parent
            # because we don't use LXML.
            # We do this only for <page>s, since we need to inspect the
            # ./revision/text element. The pages comprise the bulk of the
            # file, so in practice we prune away enough.
            elem.clear()

def segment(page_xml, include_interlinks=False):
    """Parse the content inside a page tag

    Parameters
    ----------
    page_xml : str
        Content from page tag.

    include_interlinks : bool
        Whether or not interlinks should be parsed.

    Returns
    -------
    (str, list of (str, str), (Optionally) list of (str, str))
        Structure contains (title, [(section_heading, section_content), ...],
        (Optionally) [(interlink_article, interlink_text), ...]).

    """
    elem = cElementTree.fromstring(page_xml)
    filter_namespaces = ('0',)
    namespace = get_namespace(elem.tag)
    ns_mapping = {"ns": namespace}
    text_path = "./{%(ns)s}revision/{%(ns)s}text" % ns_mapping
    title_path = "./{%(ns)s}title" % ns_mapping
    ns_path = "./{%(ns)s}ns" % ns_mapping
    lead_section_heading = "Introduction"
    top_level_heading_regex = r"\n==[^=].*[^=]==\n"
    top_level_heading_regex_capture = r"\n==([^=].*[^=])==\n"

    title = elem.find(title_path).text
    text = elem.find(text_path).text
    ns = elem.find(ns_path).text
    if ns not in filter_namespaces:
        text = None

    if text is not None:
        if include_interlinks:
            interlinks = find_interlinks(text)
        section_contents = re.split(top_level_heading_regex, text)
        section_headings = [lead_section_heading] + re.findall(top_level_heading_regex_capture, text)
        section_headings = [heading.strip() for heading in section_headings]
        assert len(section_contents) == len(section_headings)
    else:
        interlinks = []
        section_contents = []
        section_headings = []

    section_contents = [filter_wiki(section_content) for section_content in section_contents]
    sections = list(zip(section_headings, section_contents))

    if include_interlinks:
        return title, sections, interlinks
    else:
        return title, sections

class _WikiSectionsCorpus(WikiCorpus):
    """Treat a wikipedia articles dump (<LANG>wiki-<YYYYMMDD>-pages-articles.xml.bz2
    or <LANG>wiki-latest-pages-articles.xml.bz2) as a (read-only) corpus.

    The documents are extracted on-the-fly, so that the whole (massive) dump can stay compressed on disk.

    """

    def __init__(self, fileobj, min_article_character=200, processes=None,
                 lemmatize=utils.has_pattern(), filter_namespaces=('0',), include_interlinks=False):
        """
        Parameters
        ----------
        fileobj : file
            File descriptor of MediaWiki dump.
        min_article_character : int, optional
            Minimal number of character for article (except titles and leading gaps).
        processes : int, optional
            Number of processes, max(1, multiprocessing.cpu_count() - 1) if None.
        lemmatize : bool, optional
            If `pattern` package is installed, use fancier shallow parsing to get token lemmas.
            Otherwise, use simple regexp tokenization.
        filter_namespaces : tuple of int, optional
            Enumeration of namespaces that will be ignored.
        include_interlinks: bool
            Whether or not interlinks should be included in the output

        """
        self.fileobj = fileobj
        self.filter_namespaces = filter_namespaces
        self.metadata = False
        if processes is None:
            processes = max(1, multiprocessing.cpu_count() - 1)
        self.processes = processes
        self.lemmatize = lemmatize
        self.min_article_character = min_article_character
        self.include_interlinks = include_interlinks

    def get_texts_with_sections(self):
        """Iterate over the dump, returning titles and text versions of all sections of articles.

        Notes
        -----
        Only articles of sufficient length are returned (short articles & redirects
        etc are ignored).

        Note that this iterates over the **texts**; if you want vectors, just use
        the standard corpus interface instead of this function:

        .. sourcecode:: pycon

            >>> for vec in wiki_corpus:
            >>>     print(vec)

        Yields
        ------
        (str, list of (str, str), list of (str, str))
            Structure contains (title, [(section_heading, section_content), ...],
            (Optionally)[(interlink_article, interlink_text), ...]).

        """
        skipped_namespace, skipped_length, skipped_redirect = 0, 0, 0
        total_articles, total_sections = 0, 0
        page_xmls = extract_page_xmls(self.fileobj)
        pool = multiprocessing.Pool(self.processes)
        # process the corpus in smaller chunks of docs, because multiprocessing.Pool
        # is dumb and would load the entire input into RAM at once...
        for group in utils.chunkize(page_xmls, chunksize=10 * self.processes, maxsize=1):
            for article in pool.imap(partial(segment, include_interlinks=self.include_interlinks),
                                     group):  # chunksize=10): partial(merge_names, b='Sons')
                article_title, sections = article[0], article[1]

                # article redirects are pruned here
                if any(article_title.startswith(ignore + ':') for ignore in IGNORED_NAMESPACES):  # filter non-articles
                    skipped_namespace += 1
                    continue
                if not sections or sections[0][1].lstrip().lower().startswith("#redirect"):  # filter redirect
                    skipped_redirect += 1
                    continue
                if sum(len(body.strip()) for (_, body) in sections) < self.min_article_character:
                    # filter stubs (incomplete, very short articles)
                    skipped_length += 1
                    continue
                total_articles += 1
                total_sections += len(sections)

                if self.include_interlinks:
                    interlinks = article[2]
                    yield (article_title, sections, interlinks)
                else:
                    yield (article_title, sections)

        logger.info(
            "finished processing %i articles with %i sections (skipped %i redirects, %i stubs, %i ignored namespaces)",
            total_articles, total_sections, skipped_redirect, skipped_length, skipped_namespace)
        pool.terminate()
        self.length = total_articles  # cache corpus length

if __name__ == "__main__":
    logging.basicConfig(format='%(asctime)s - %(module)s - %(levelname)s - %(message)s', level=logging.INFO)
    parser = argparse.ArgumentParser(formatter_class=argparse.RawTextHelpFormatter, description=__doc__[:-136])
    default_workers = max(1, multiprocessing.cpu_count() - 1)
    parser.add_argument('-f', '--file', help='Path to MediaWiki database dump (read-only).', required=True)
    parser.add_argument(
        '-o', '--output',
        help='Path to output file (stdout if not specified). If ends in .gz or .bz2, '
             'the output file will be automatically compressed (recommended!).')
    parser.add_argument(
        '-w', '--workers',
        help='Number of parallel workers for multi-core systems. Default: %(default)s.',
        type=int,
        default=default_workers
    )
    parser.add_argument(
        '-m', '--min-article-character',
        help="Ignore articles with fewer characters than this (article stubs). Default: %(default)s.",
        default=200
    )
    parser.add_argument(
        '-i', '--include-interlinks',
        help='Include a mapping for interlinks to other articles in the dump. The mappings format is: '
             '"interlinks": [("article_title_1", "interlink_text_1"), ("article_title_2", "interlink_text_2"), ...]',
        action='store_true'
    )
    args = parser.parse_args()

    logger.info("running %s", " ".join(sys.argv))
    segment_and_write_all_articles(
        args.file, args.output,
        min_article_character=args.min_article_character,
        workers=args.workers,
        include_interlinks=args.include_interlinks
    )

    logger.info("finished running %s", sys.argv[0])
<commit_msg>Fix TypeError when using the -m flag (#2734)<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# Author: Jayant Jain <jayant@rare-technologies.com>
# Copyright (C) 2016 RaRe Technologies

"""This script using for extracting plain text out of a raw Wikipedia dump. Input is an xml.bz2 file provided
by MediaWiki that looks like <LANG>wiki-<YYYYMMDD>-pages-articles.xml.bz2 or <LANG>wiki-latest-pages-articles.xml.bz2
(e.g. 14 GB of https://dumps.wikimedia.org/enwiki/latest/enwiki-latest-pages-articles.xml.bz2).

It streams through all the XML articles using multiple cores (#cores - 1, by default),
decompressing on the fly and extracting plain text from the articles and their sections.

For each extracted article, it prints its title, section names and plain text section contents, in json-line format.

How to use
----------
#. Process Wikipedia dump with this script ::

    python -m gensim.scripts.segment_wiki -i -f enwiki-latest-pages-articles.xml.bz2 -o enwiki-latest.json.gz

#. Read output in simple way:

.. sourcecode:: pycon

    >>> from gensim import utils
    >>> import json
    >>>
    >>> # iterate over the plain text data we just created
    >>> with utils.open('enwiki-latest.json.gz', 'rb') as f:
    >>>     for line in f:
    >>>         # decode each JSON line into a Python dictionary object
    >>>         article = json.loads(line)
    >>>
    >>>         # each article has a "title", a mapping of interlinks and a list of "section_titles" and
    >>>         # "section_texts".
    >>>         print("Article title: %s" % article['title'])
    >>>         print("Interlinks: %s" + article['interlinks'])
    >>>         for section_title, section_text in zip(article['section_titles'], article['section_texts']):
    >>>             print("Section title: %s" % section_title)
    >>>             print("Section text: %s" % section_text)

Notes
-----
Processing the entire English Wikipedia dump takes 1.7 hours (about 3 million articles per hour,
or 10 MB of XML per second) on an 8 core Intel i7-7700 @3.60GHz.

Command line arguments
----------------------

.. program-output:: python -m gensim.scripts.segment_wiki --help
   :ellipsis: 0, -10

"""

import argparse
import json
import logging
import multiprocessing
import re
import sys
from xml.etree import cElementTree
from functools import partial

from gensim.corpora.wikicorpus import IGNORED_NAMESPACES, WikiCorpus, filter_wiki, find_interlinks, get_namespace, utils
import gensim.utils

logger = logging.getLogger(__name__)

def segment_all_articles(file_path, min_article_character=200, workers=None, include_interlinks=False):
    """Extract article titles and sections from a MediaWiki bz2 database dump.

    Parameters
    ----------
    file_path : str
        Path to MediaWiki dump, typical filename is <LANG>wiki-<YYYYMMDD>-pages-articles.xml.bz2
        or <LANG>wiki-latest-pages-articles.xml.bz2.

    min_article_character : int, optional
        Minimal number of character for article (except titles and leading gaps).

    workers: int or None
        Number of parallel workers, max(1, multiprocessing.cpu_count() - 1) if None.

    include_interlinks: bool
        Whether or not interlinks should be included in the output

    Yields
    ------
    (str, list of (str, str), (Optionally) list of (str, str))
        Structure contains (title, [(section_heading, section_content), ...],
        (Optionally) [(interlink_article, interlink_text), ...]).

    """
    with gensim.utils.open(file_path, 'rb') as xml_fileobj:
        wiki_sections_corpus = _WikiSectionsCorpus(
            xml_fileobj, min_article_character=min_article_character, processes=workers,
            include_interlinks=include_interlinks)
        wiki_sections_corpus.metadata = True
        wiki_sections_text = wiki_sections_corpus.get_texts_with_sections()

        for article in wiki_sections_text:
            yield article

def segment_and_write_all_articles(file_path, output_file, min_article_character=200, workers=None,
                                   include_interlinks=False):
    """Write article title and sections to `output_file` (or stdout, if output_file is None).

    The output format is one article per line, in json-line format with 4 fields::

        'title' - title of article,
        'section_titles' - list of titles of sections,
        'section_texts' - list of content from sections,
        (Optional) 'section_interlinks' - list of interlinks in the article.

    Parameters
    ----------
    file_path : str
        Path to MediaWiki dump, typical filename is <LANG>wiki-<YYYYMMDD>-pages-articles.xml.bz2
        or <LANG>wiki-latest-pages-articles.xml.bz2.

    output_file : str or None
        Path to output file in json-lines format, or None for printing to stdout.

    min_article_character : int, optional
        Minimal number of character for article (except titles and leading gaps).

    workers: int or None
        Number of parallel workers, max(1, multiprocessing.cpu_count() - 1) if None.

    include_interlinks: bool
        Whether or not interlinks should be included in the output
    """
    if output_file is None:
        outfile = getattr(sys.stdout, 'buffer', sys.stdout)  # we want write bytes, so for py3 we used 'buffer'
    else:
        outfile = gensim.utils.open(output_file, 'wb')

    try:
        article_stream = segment_all_articles(file_path, min_article_character, workers=workers,
                                              include_interlinks=include_interlinks)
        for idx, article in enumerate(article_stream):
            article_title, article_sections = article[0], article[1]
            if include_interlinks:
                interlinks = article[2]

            output_data = {
                "title": article_title,
                "section_titles": [],
                "section_texts": [],
            }
            if include_interlinks:
                output_data["interlinks"] = interlinks

            for section_heading, section_content in article_sections:
                output_data["section_titles"].append(section_heading)
                output_data["section_texts"].append(section_content)

            if (idx + 1) % 100000 == 0:
                logger.info("processed #%d articles (at %r now)", idx + 1, article_title)
            outfile.write((json.dumps(output_data) + "\n").encode('utf-8'))

    finally:
        if output_file is not None:
            outfile.close()

def extract_page_xmls(f):
    """Extract pages from a MediaWiki database dump.

    Parameters
    ----------
    f : file
        File descriptor of MediaWiki dump.

    Yields
    ------
    str
        XML strings for page tags.

    """
    elems = (elem for _, elem in cElementTree.iterparse(f, events=("end",)))

    elem = next(elems)
    namespace = get_namespace(elem.tag)
    ns_mapping = {"ns": namespace}
    page_tag = "{%(ns)s}page" % ns_mapping

    for elem in elems:
        if elem.tag == page_tag:
            yield cElementTree.tostring(elem)
            # Prune the element tree, as per
            # http://www.ibm.com/developerworks/xml/library/x-hiperfparse/
            # except that we don't need to prune backlinks from the parent
            # because we don't use LXML.
            # We do this only for <page>s, since we need to inspect the
            # ./revision/text element. The pages comprise the bulk of the
            # file, so in practice we prune away enough.
            elem.clear()

def segment(page_xml, include_interlinks=False):
    """Parse the content inside a page tag

    Parameters
    ----------
    page_xml : str
        Content from page tag.

    include_interlinks : bool
        Whether or not interlinks should be parsed.

    Returns
    -------
    (str, list of (str, str), (Optionally) list of (str, str))
        Structure contains (title, [(section_heading, section_content), ...],
        (Optionally) [(interlink_article, interlink_text), ...]).

    """
    elem = cElementTree.fromstring(page_xml)
    filter_namespaces = ('0',)
    namespace = get_namespace(elem.tag)
    ns_mapping = {"ns": namespace}
    text_path = "./{%(ns)s}revision/{%(ns)s}text" % ns_mapping
    title_path = "./{%(ns)s}title" % ns_mapping
    ns_path = "./{%(ns)s}ns" % ns_mapping
    lead_section_heading = "Introduction"
    top_level_heading_regex = r"\n==[^=].*[^=]==\n"
    top_level_heading_regex_capture = r"\n==([^=].*[^=])==\n"

    title = elem.find(title_path).text
    text = elem.find(text_path).text
    ns = elem.find(ns_path).text
    if ns not in filter_namespaces:
        text = None

    if text is not None:
        if include_interlinks:
            interlinks = find_interlinks(text)
        section_contents = re.split(top_level_heading_regex, text)
        section_headings = [lead_section_heading] + re.findall(top_level_heading_regex_capture, text)
        section_headings = [heading.strip() for heading in section_headings]
        assert len(section_contents) == len(section_headings)
    else:
        interlinks = []
        section_contents = []
        section_headings = []

    section_contents = [filter_wiki(section_content) for section_content in section_contents]
    sections = list(zip(section_headings, section_contents))

    if include_interlinks:
        return title, sections, interlinks
    else:
        return title, sections

class _WikiSectionsCorpus(WikiCorpus):
    """Treat a wikipedia articles dump (<LANG>wiki-<YYYYMMDD>-pages-articles.xml.bz2
    or <LANG>wiki-latest-pages-articles.xml.bz2) as a (read-only) corpus.

    The documents are extracted on-the-fly, so that the whole (massive) dump can stay compressed on disk.

    """

    def __init__(self, fileobj, min_article_character=200, processes=None,
                 lemmatize=utils.has_pattern(), filter_namespaces=('0',), include_interlinks=False):
        """
        Parameters
        ----------
        fileobj : file
            File descriptor of MediaWiki dump.
        min_article_character : int, optional
            Minimal number of character for article (except titles and leading gaps).
        processes : int, optional
            Number of processes, max(1, multiprocessing.cpu_count() - 1) if None.
        lemmatize : bool, optional
            If `pattern` package is installed, use fancier shallow parsing to get token lemmas.
            Otherwise, use simple regexp tokenization.
        filter_namespaces : tuple of int, optional
            Enumeration of namespaces that will be ignored.
        include_interlinks: bool
            Whether or not interlinks should be included in the output

        """
        self.fileobj = fileobj
        self.filter_namespaces = filter_namespaces
        self.metadata = False
        if processes is None:
            processes = max(1, multiprocessing.cpu_count() - 1)
        self.processes = processes
        self.lemmatize = lemmatize
        self.min_article_character = min_article_character
        self.include_interlinks = include_interlinks

    def get_texts_with_sections(self):
        """Iterate over the dump, returning titles and text versions of all sections of articles.

        Notes
        -----
        Only articles of sufficient length are returned (short articles & redirects
        etc are ignored).

        Note that this iterates over the **texts**; if you want vectors, just use
        the standard corpus interface instead of this function:

        .. sourcecode:: pycon

            >>> for vec in wiki_corpus:
            >>>     print(vec)

        Yields
        ------
        (str, list of (str, str), list of (str, str))
            Structure contains (title, [(section_heading, section_content), ...],
            (Optionally)[(interlink_article, interlink_text), ...]).

        """
        skipped_namespace, skipped_length, skipped_redirect = 0, 0, 0
        total_articles, total_sections = 0, 0
        page_xmls = extract_page_xmls(self.fileobj)
        pool = multiprocessing.Pool(self.processes)
        # process the corpus in smaller chunks of docs, because multiprocessing.Pool
        # is dumb and would load the entire input into RAM at once...
        for group in utils.chunkize(page_xmls, chunksize=10 * self.processes, maxsize=1):
            for article in pool.imap(partial(segment, include_interlinks=self.include_interlinks),
                                     group):  # chunksize=10): partial(merge_names, b='Sons')
                article_title, sections = article[0], article[1]

                # article redirects are pruned here
                if any(article_title.startswith(ignore + ':') for ignore in IGNORED_NAMESPACES):  # filter non-articles
                    skipped_namespace += 1
                    continue
                if not sections or sections[0][1].lstrip().lower().startswith("#redirect"):  # filter redirect
                    skipped_redirect += 1
                    continue
                if sum(len(body.strip()) for (_, body) in sections) < self.min_article_character:
                    # filter stubs (incomplete, very short articles)
                    skipped_length += 1
                    continue
                total_articles += 1
                total_sections += len(sections)

                if self.include_interlinks:
                    interlinks = article[2]
                    yield (article_title, sections, interlinks)
                else:
                    yield (article_title, sections)

        logger.info(
            "finished processing %i articles with %i sections (skipped %i redirects, %i stubs, %i ignored namespaces)",
            total_articles, total_sections, skipped_redirect, skipped_length, skipped_namespace)
        pool.terminate()
        self.length = total_articles  # cache corpus length

if __name__ == "__main__":
    logging.basicConfig(format='%(asctime)s - %(module)s - %(levelname)s - %(message)s', level=logging.INFO)
    parser = argparse.ArgumentParser(formatter_class=argparse.RawTextHelpFormatter, description=__doc__[:-136])
    default_workers = max(1, multiprocessing.cpu_count() - 1)
    parser.add_argument('-f', '--file', help='Path to MediaWiki database dump (read-only).', required=True)
    parser.add_argument(
        '-o', '--output',
        help='Path to output file (stdout if not specified). If ends in .gz or .bz2, '
             'the output file will be automatically compressed (recommended!).')
    parser.add_argument(
        '-w', '--workers',
        help='Number of parallel workers for multi-core systems. Default: %(default)s.',
        type=int,
        default=default_workers
    )
    parser.add_argument(
        '-m', '--min-article-character',
        help="Ignore articles with fewer characters than this (article stubs). Default: %(default)s.",
        type=int,
        default=200
    )
    parser.add_argument(
        '-i', '--include-interlinks',
        help='Include a mapping for interlinks to other articles in the dump. The mappings format is: '
             '"interlinks": [("article_title_1", "interlink_text_1"), ("article_title_2", "interlink_text_2"), ...]',
        action='store_true'
    )
    args = parser.parse_args()

    logger.info("running %s", " ".join(sys.argv))
    segment_and_write_all_articles(
        args.file, args.output,
        min_article_character=args.min_article_character,
        workers=args.workers,
        include_interlinks=args.include_interlinks
    )

    logger.info("finished running %s", sys.argv[0])
<|endoftext|>


<commit_before>"""
.. module:: VisVTK
    :platform: Unix, Windows
    :synopsis: VTK visualization component for NURBS-Python

.. moduleauthor:: Onur Rauf Bingol <orbingol@gmail.com>

"""

from random import random
from . import vis
from . import vtk_helpers as vtkh
import numpy as np
from vtk.util.numpy_support import numpy_to_vtk
from vtk import VTK_FLOAT

class VisConfig(vis.VisConfigAbstract):
    """ Configuration class for VTK visualization module.

    This class is only required when you would like to change the visual defaults of the plots and the figure.

    The ``VisVTK`` module has the following configuration variables:

    * ``ctrlpts`` (bool): Control points polygon/grid visibility. *Default: True*
    * ``evalpts`` (bool): Curve/surface points visibility. *Default: True*
    * ``trims`` (bool): Trim curve visibility. *Default: True*
    * ``figure_size`` (list): Size of the figure in (x, y). *Default: (800, 600)*
    * ``line_width`` (int): Thickness of the lines on the figure. *Default: 1.0*
    """
    def __init__(self, **kwargs):
        super(VisConfig, self).__init__(**kwargs)
        self._bg = (  # background colors
            (0.5, 0.5, 0.5), (0.2, 0.2, 0.2), (0.25, 0.5, 0.75), (1.0, 1.0, 0.0),
            (1.0, 0.5, 0.0), (0.5, 0.0, 1.0), (0.0, 0.0, 0.0), (1.0, 1.0, 1.0)
        )
        self._bg_id = 0  # used for keeping track of the background numbering
        self.display_ctrlpts = kwargs.get('ctrlpts', True)
        self.display_evalpts = kwargs.get('evalpts', True)
        self.display_trims = kwargs.get('trims', True)
        self.figure_size = kwargs.get('figure_size', (800, 600))  # size of the render window
        self.line_width = kwargs.get('line_width', 1.0)

    def keypress_callback(self, obj, ev):
        """ VTK callback for keypress events.

        Keypress events:
            * ``e``: exit the application
            * ``p``: pick object (hover the mouse and then press to pick)
            * ``l``: change color of the picked object
            * ``f``: fly to point (click somewhere in the window and press to fly)
            * ``r``: reset the camera
            * ``s`` and ``w``: switch between solid and wireframe modes
            * ``b``: change background color
            * ``arrow keys``: pan the model
            * ``d``: print debug information (of picked object, point, etc.)

        Please refer to `vtkInteractorStyle <https://vtk.org/doc/nightly/html/classvtkInteractorStyle.html>`_ class
        reference for more details.

        :param obj: render window interactor
        :type obj: vtkRenderWindowInteractor
        :param ev: event name
        :type ev: str
        """
        key = obj.GetKeySym()  # pressed key (as str)
        render_window = obj.GetRenderWindow()  # vtkRenderWindow
        renderer = render_window.GetRenderers().GetFirstRenderer()  # vtkRenderer

        # Custom keypress events
        if key == 'Up':
            camera = renderer.GetActiveCamera()  # vtkCamera
            camera.Pitch(-2.5)
        if key == 'Down':
            camera = renderer.GetActiveCamera()  # vtkCamera
            camera.Pitch(2.5)
        if key == 'Left':
            camera = renderer.GetActiveCamera()  # vtkCamera
            camera.Yaw(-2.5)
        if key == 'Right':
            camera = renderer.GetActiveCamera()  # vtkCamera
            camera.Yaw(2.5)
        if key == 'b':
            if self._bg_id >= len(self._bg):
                self._bg_id = 0
            renderer.SetBackground(*self._bg[self._bg_id])
            self._bg_id += 1
        if key == 'l':
            picker = obj.GetPicker()  # vtkPropPicker
            actor = picker.GetActor()  # vtkActor
            if actor is not None:
                actor.GetProperty().SetColor(random(), random(), random())
        if key == 'd':
            picker = obj.GetPicker()  # vtkPropPicker
            actor = picker.GetActor()  # vtkActor
            if actor is not None:
                print("Name:", actor.GetMapper().GetArrayName())
                print("Index:", actor.GetMapper().GetArrayId())
            print("Selected point:", picker.GetSelectionPoint()[0:2])

        # Update render window
        render_window.Render()

class VisCurve3D(vis.VisAbstract):
    """ VTK visualization module for curves. """
    def __init__(self, config=VisConfig()):
        super(VisCurve3D, self).__init__(config=config)

    def render(self, **kwargs):
        """ Plots the curve and the control points polygon. """
        # Calling parent function
        super(VisCurve3D, self).render(**kwargs)

        # Initialize a list to store VTK actors
        vtk_actors = []

        # Start plotting
        for plot in self._plots:
            # Plot control points
            if plot['type'] == 'ctrlpts' and self.vconf.display_ctrlpts:
                # Points as spheres
                pts = np.array(plot['ptsarr'], dtype=np.float)
                # Handle 2-dimensional data
                if pts.shape[1] == 2:
                    pts = np.c_[pts, np.zeros(pts.shape[0], dtype=np.float)]
                vtkpts = numpy_to_vtk(pts, deep=False, array_type=VTK_FLOAT)
                vtkpts.SetName(plot['name'])
                actor1 = vtkh.create_actor_pts(pts=vtkpts, color=vtkh.create_color(plot['color']),
                                               name=plot['name'], idx=plot['idx'])
                vtk_actors.append(actor1)
                # Lines
                actor2 = vtkh.create_actor_polygon(pts=vtkpts, color=vtkh.create_color(plot['color']),
                                                   name=plot['name'], index=plot['idx'], size=self.vconf.line_width)
                vtk_actors.append(actor2)

            # Plot evaluated points
            if plot['type'] == 'evalpts' and self.vconf.display_evalpts:
                pts = np.array(plot['ptsarr'], dtype=np.float)
                # Handle 2-dimensional data
                if pts.shape[1] == 2:
                    pts = np.c_[pts, np.zeros(pts.shape[0], dtype=np.float)]
                vtkpts = numpy_to_vtk(pts, deep=False, array_type=VTK_FLOAT)
                vtkpts.SetName(plot['name'])
                actor1 = vtkh.create_actor_polygon(pts=vtkpts, color=vtkh.create_color(plot['color']),
                                                   name=plot['name'], index=plot['idx'], size=self.vconf.line_width * 2)
                vtk_actors.append(actor1)

        # Render actors
        vtkh.create_render_window(vtk_actors, dict(KeyPressEvent=(self.vconf.keypress_callback, 1.0)),
                                  figure_size=self.vconf.figure_size)

# It is easier to plot 2-dimensional curves with VisCurve3D
VisCurve2D = VisCurve3D

class VisSurface(vis.VisAbstract):
    """ VTK visualization module for surfaces. """
    def __init__(self, config=VisConfig()):
        super(VisSurface, self).__init__(config=config)
        self._module_config['ctrlpts'] = "quads"
        self._module_config['evalpts'] = "triangles"

    def render(self, **kwargs):
        """ Plots the surface and the control points grid. """
        # Calling parent function
        super(VisSurface, self).render(**kwargs)

        # Initialize a list to store VTK actors
        vtk_actors = []

        # Start plotting
        for plot in self._plots:
            # Plot control points
            if plot['type'] == 'ctrlpts' and self.vconf.display_ctrlpts:
                vertices = [v.data for v in plot['ptsarr'][0]]
                faces = [q.data for q in plot['ptsarr'][1]]
                # Points as spheres
                pts = np.array(vertices, dtype=np.float)
                vtkpts = numpy_to_vtk(pts, deep=False, array_type=VTK_FLOAT)
                vtkpts.SetName(plot['name'])
                actor1 = vtkh.create_actor_pts(pts=vtkpts, color=vtkh.create_color(plot['color']),
                                               name=plot['name'], index=plot['idx'])
                vtk_actors.append(actor1)
                # Quad mesh
                lines = np.array(faces, dtype=np.int)
                actor2 = vtkh.create_actor_mesh(pts=vtkpts, lines=lines, color=vtkh.create_color(plot['color']),
                                                name=plot['name'], index=plot['idx'], size=self.vconf.line_width)
                vtk_actors.append(actor2)

            # Plot evaluated points
            if plot['type'] == 'evalpts' and self.vconf.display_evalpts:
                vertices = [v.data for v in plot['ptsarr'][0]]
                vtkpts = numpy_to_vtk(vertices, deep=False, array_type=VTK_FLOAT)
                vtkpts.SetName(plot['name'])
                faces = [t.vertex_ids_zero for t in plot['ptsarr'][1]]
                tris = np.array(faces, dtype=np.int)
                actor1 = vtkh.create_actor_tri(pts=vtkpts, tris=tris, color=vtkh.create_color(plot['color']),
                                               name=plot['name'], index=plot['idx'])
                vtk_actors.append(actor1)

            # Plot trim curves
            if self.vconf.display_trims:
                if plot['type'] == 'trimcurve':
                    pts = np.array(plot['ptsarr'], dtype=np.float)
                    vtkpts = numpy_to_vtk(pts, deep=False, array_type=VTK_FLOAT)
                    vtkpts.SetName(plot['name'])
                    actor1 = vtkh.create_actor_polygon(pts=vtkpts, color=vtkh.create_color(plot['color']),
                                                       name=plot['name'], index=plot['idx'])
                    vtk_actors.append(actor1)

        # Render actors
        vtkh.create_render_window(vtk_actors, dict(KeyPressEvent=(self.vconf.keypress_callback, 1.0)),
                                  figure_size=self.vconf.figure_size)

class VisVolume(vis.VisAbstract):
    """ VTK visualization module for volumes. """
    def __init__(self, config=VisConfig()):
        super(VisVolume, self).__init__(config=config)
        self._module_config['ctrlpts'] = "points"
        self._module_config['evalpts'] = "points"

    def render(self, **kwargs):
        """ Plots the volume and the control points. """
        # Calling parent function
        super(VisVolume, self).render(**kwargs)

        # Initialize a list to store VTK actors
        vtk_actors = []

        # Start plotting
        for plot in self._plots:
            # Plot control points
            if plot['type'] == 'ctrlpts' and self.vconf.display_ctrlpts:
                # Points as spheres
                pts = np.array(plot['ptsarr'], dtype=np.float)
                vtkpts = numpy_to_vtk(pts, deep=False, array_type=VTK_FLOAT)
                vtkpts.SetName(plot['name'])
                temp_actor = vtkh.create_actor_pts(pts=vtkpts, color=vtkh.create_color(plot['color']),
                                                   name=plot['name'], index=plot['idx'])
                vtk_actors.append(temp_actor)

            # Plot evaluated points
            if plot['type'] == 'evalpts' and self.vconf.display_evalpts:
                pts = np.array(plot['ptsarr'], dtype=np.float)
                vtkpts = numpy_to_vtk(pts, deep=False, array_type=VTK_FLOAT)
                vtkpts.SetName(plot['name'])
                temp_actor = vtkh.create_actor_pts(pts=vtkpts, color=vtkh.create_color(plot['color']),
                                                   name=plot['name'], index=plot['idx'])
                vtk_actors.append(temp_actor)

        # Render actors
        vtkh.create_render_window(vtk_actors, dict(KeyPressEvent=(self.vconf.keypress_callback, 1.0)),
                                  figure_size=self.vconf.figure_size)

class VisVoxel(vis.VisAbstract):
    """ VTK visualization module for voxel representation of the volumes. """
    def __init__(self, config=VisConfig()):
        super(VisVoxel, self).__init__(config=config)
        self._module_config['ctrlpts'] = "points"
        self._module_config['evalpts'] = "voxels"

    def render(self, **kwargs):
        """ Plots the volume and the control points. """
        # Calling parent function
        super(VisVoxel, self).render(**kwargs)

        # Initialize a list to store VTK actors
        vtk_actors = []

        # Start plotting
        for plot in self._plots:
            # Plot control points
            if plot['type'] == 'ctrlpts' and self.vconf.display_ctrlpts:
                # Points as spheres
                pts = np.array(plot['ptsarr'], dtype=np.float)
                vtkpts = numpy_to_vtk(pts, deep=False, array_type=VTK_FLOAT)
                vtkpts.SetName(plot['name'])
                temp_actor = vtkh.create_actor_pts(pts=vtkpts, color=vtkh.create_color(plot['color']),
                                                   name=plot['name'], index=plot['idx'])
                vtk_actors.append(temp_actor)

            # Plot evaluated points
            if plot['type'] == 'evalpts' and self.vconf.display_evalpts:
                faces = np.array(plot['ptsarr'][1], dtype=np.float)
                filled = np.array(plot['ptsarr'][2], dtype=np.int)
                grid_filled = faces[filled == 1]
                temp_actor = vtkh.create_actor_hexahedron(grid=grid_filled, color=vtkh.create_color(plot['color']),
                                                          name=plot['name'], index=plot['idx'])
                vtk_actors.append(temp_actor)

        # Render actors
        vtkh.create_render_window(vtk_actors, dict(KeyPressEvent=(self.vconf.keypress_callback, 1.0)),
                                  figure_size=self.vconf.figure_size)
<commit_msg>Issue #46 assign key "h" for changing object visibility<commit_after>"""
.. module:: VisVTK
    :platform: Unix, Windows
    :synopsis: VTK visualization component for NURBS-Python

.. moduleauthor:: Onur Rauf Bingol <orbingol@gmail.com>

"""

from random import random
from . import vis
from . import vtk_helpers as vtkh
import numpy as np
from vtk.util.numpy_support import numpy_to_vtk
from vtk import VTK_FLOAT

class VisConfig(vis.VisConfigAbstract):
    """ Configuration class for VTK visualization module.

    This class is only required when you would like to change the visual defaults of the plots and the figure.

    The ``VisVTK`` module has the following configuration variables:

    * ``ctrlpts`` (bool): Control points polygon/grid visibility. *Default: True*
    * ``evalpts`` (bool): Curve/surface points visibility. *Default: True*
    * ``trims`` (bool): Trim curve visibility. *Default: True*
    * ``figure_size`` (list): Size of the figure in (x, y). *Default: (800, 600)*
    * ``line_width`` (int): Thickness of the lines on the figure. *Default: 1.0*
    """
    def __init__(self, **kwargs):
        super(VisConfig, self).__init__(**kwargs)
        self._bg = (  # background colors
            (0.5, 0.5, 0.5), (0.2, 0.2, 0.2), (0.25, 0.5, 0.75), (1.0, 1.0, 0.0),
            (1.0, 0.5, 0.0), (0.5, 0.0, 1.0), (0.0, 0.0, 0.0), (1.0, 1.0, 1.0)
        )
        self._bg_id = 0  # used for keeping track of the background numbering
        self.display_ctrlpts = kwargs.get('ctrlpts', True)
        self.display_evalpts = kwargs.get('evalpts', True)
        self.display_trims = kwargs.get('trims', True)
        self.figure_size = kwargs.get('figure_size', (800, 600))  # size of the render window
        self.line_width = kwargs.get('line_width', 1.0)

    def keypress_callback(self, obj, ev):
        """ VTK callback for keypress events.

        Keypress events:
            * ``e``: exit the application
            * ``p``: pick object (hover the mouse and then press to pick)
            * ``l``: change color of the picked object
            * ``f``: fly to point (click somewhere in the window and press to fly)
            * ``r``: reset the camera
            * ``s`` and ``w``: switch between solid and wireframe modes
            * ``b``: change background color
            * ``arrow keys``: pan the model
            * ``h``: change object visibility
            * ``d``: print debug information (of picked object, point, etc.)

        Please refer to `vtkInteractorStyle <https://vtk.org/doc/nightly/html/classvtkInteractorStyle.html>`_ class
        reference for more details.

        :param obj: render window interactor
        :type obj: vtkRenderWindowInteractor
        :param ev: event name
        :type ev: str
        """
        key = obj.GetKeySym()  # pressed key (as str)
        render_window = obj.GetRenderWindow()  # vtkRenderWindow
        renderer = render_window.GetRenderers().GetFirstRenderer()  # vtkRenderer
        picker = obj.GetPicker()  # vtkPropPicker
        actor = picker.GetActor()  # vtkActor

        # Custom keypress events
        if key == 'Up':
            camera = renderer.GetActiveCamera()  # vtkCamera
            camera.Pitch(-2.5)
        if key == 'Down':
            camera = renderer.GetActiveCamera()  # vtkCamera
            camera.Pitch(2.5)
        if key == 'Left':
            camera = renderer.GetActiveCamera()  # vtkCamera
            camera.Yaw(-2.5)
        if key == 'Right':
            camera = renderer.GetActiveCamera()  # vtkCamera
            camera.Yaw(2.5)
        if key == 'b':
            if self._bg_id >= len(self._bg):
                self._bg_id = 0
            renderer.SetBackground(*self._bg[self._bg_id])
            self._bg_id += 1
        if key == 'l':
            if actor is not None:
                actor.GetProperty().SetColor(random(), random(), random())
        if key == 'd':
            if actor is not None:
                print("Name:", actor.GetMapper().GetArrayName())
                print("Index:", actor.GetMapper().GetArrayId())
            print("Selected point:", picker.GetSelectionPoint()[0:2])
        if key == 'h':
            if actor is not None:
                actor.SetVisibility(not actor.GetVisibility())

        # Update render window
        render_window.Render()

class VisCurve3D(vis.VisAbstract):
    """ VTK visualization module for curves. """
    def __init__(self, config=VisConfig()):
        super(VisCurve3D, self).__init__(config=config)

    def render(self, **kwargs):
        """ Plots the curve and the control points polygon. """
        # Calling parent function
        super(VisCurve3D, self).render(**kwargs)

        # Initialize a list to store VTK actors
        vtk_actors = []

        # Start plotting
        for plot in self._plots:
            # Plot control points
            if plot['type'] == 'ctrlpts' and self.vconf.display_ctrlpts:
                # Points as spheres
                pts = np.array(plot['ptsarr'], dtype=np.float)
                # Handle 2-dimensional data
                if pts.shape[1] == 2:
                    pts = np.c_[pts, np.zeros(pts.shape[0], dtype=np.float)]
                vtkpts = numpy_to_vtk(pts, deep=False, array_type=VTK_FLOAT)
                vtkpts.SetName(plot['name'])
                actor1 = vtkh.create_actor_pts(pts=vtkpts, color=vtkh.create_color(plot['color']),
                                               name=plot['name'], idx=plot['idx'])
                vtk_actors.append(actor1)
                # Lines
                actor2 = vtkh.create_actor_polygon(pts=vtkpts, color=vtkh.create_color(plot['color']),
                                                   name=plot['name'], index=plot['idx'], size=self.vconf.line_width)
                vtk_actors.append(actor2)

            # Plot evaluated points
            if plot['type'] == 'evalpts' and self.vconf.display_evalpts:
                pts = np.array(plot['ptsarr'], dtype=np.float)
                # Handle 2-dimensional data
                if pts.shape[1] == 2:
                    pts = np.c_[pts, np.zeros(pts.shape[0], dtype=np.float)]
                vtkpts = numpy_to_vtk(pts, deep=False, array_type=VTK_FLOAT)
                vtkpts.SetName(plot['name'])
                actor1 = vtkh.create_actor_polygon(pts=vtkpts, color=vtkh.create_color(plot['color']),
                                                   name=plot['name'], index=plot['idx'], size=self.vconf.line_width * 2)
                vtk_actors.append(actor1)

        # Render actors
        vtkh.create_render_window(vtk_actors, dict(KeyPressEvent=(self.vconf.keypress_callback, 1.0)),
                                  figure_size=self.vconf.figure_size)

# It is easier to plot 2-dimensional curves with VisCurve3D
VisCurve2D = VisCurve3D

class VisSurface(vis.VisAbstract):
    """ VTK visualization module for surfaces. """
    def __init__(self, config=VisConfig()):
        super(VisSurface, self).__init__(config=config)
        self._module_config['ctrlpts'] = "quads"
        self._module_config['evalpts'] = "triangles"

    def render(self, **kwargs):
        """ Plots the surface and the control points grid. """
        # Calling parent function
        super(VisSurface, self).render(**kwargs)

        # Initialize a list to store VTK actors
        vtk_actors = []

        # Start plotting
        for plot in self._plots:
            # Plot control points
            if plot['type'] == 'ctrlpts' and self.vconf.display_ctrlpts:
                vertices = [v.data for v in plot['ptsarr'][0]]
                faces = [q.data for q in plot['ptsarr'][1]]
                # Points as spheres
                pts = np.array(vertices, dtype=np.float)
                vtkpts = numpy_to_vtk(pts, deep=False, array_type=VTK_FLOAT)
                vtkpts.SetName(plot['name'])
                actor1 = vtkh.create_actor_pts(pts=vtkpts, color=vtkh.create_color(plot['color']),
                                               name=plot['name'], index=plot['idx'])
                vtk_actors.append(actor1)
                # Quad mesh
                lines = np.array(faces, dtype=np.int)
                actor2 = vtkh.create_actor_mesh(pts=vtkpts, lines=lines, color=vtkh.create_color(plot['color']),
                                                name=plot['name'], index=plot['idx'], size=self.vconf.line_width)
                vtk_actors.append(actor2)

            # Plot evaluated points
            if plot['type'] == 'evalpts' and self.vconf.display_evalpts:
                vertices = [v.data for v in plot['ptsarr'][0]]
                vtkpts = numpy_to_vtk(vertices, deep=False, array_type=VTK_FLOAT)
                vtkpts.SetName(plot['name'])
                faces = [t.vertex_ids_zero for t in plot['ptsarr'][1]]
                tris = np.array(faces, dtype=np.int)
                actor1 = vtkh.create_actor_tri(pts=vtkpts, tris=tris, color=vtkh.create_color(plot['color']),
                                               name=plot['name'], index=plot['idx'])
                vtk_actors.append(actor1)

            # Plot trim curves
            if self.vconf.display_trims:
                if plot['type'] == 'trimcurve':
                    pts = np.array(plot['ptsarr'], dtype=np.float)
                    vtkpts = numpy_to_vtk(pts, deep=False, array_type=VTK_FLOAT)
                    vtkpts.SetName(plot['name'])
                    actor1 = vtkh.create_actor_polygon(pts=vtkpts, color=vtkh.create_color(plot['color']),
                                                       name=plot['name'], index=plot['idx'])
                    vtk_actors.append(actor1)

        # Render actors
        vtkh.create_render_window(vtk_actors, dict(KeyPressEvent=(self.vconf.keypress_callback, 1.0)),
                                  figure_size=self.vconf.figure_size)

class VisVolume(vis.VisAbstract):
    """ VTK visualization module for volumes. """
    def __init__(self, config=VisConfig()):
        super(VisVolume, self).__init__(config=config)
        self._module_config['ctrlpts'] = "points"
        self._module_config['evalpts'] = "points"

    def render(self, **kwargs):
        """ Plots the volume and the control points. """
        # Calling parent function
        super(VisVolume, self).render(**kwargs)

        # Initialize a list to store VTK actors
        vtk_actors = []

        # Start plotting
        for plot in self._plots:
            # Plot control points
            if plot['type'] == 'ctrlpts' and self.vconf.display_ctrlpts:
                # Points as spheres
                pts = np.array(plot['ptsarr'], dtype=np.float)
                vtkpts = numpy_to_vtk(pts, deep=False, array_type=VTK_FLOAT)
                vtkpts.SetName(plot['name'])
                temp_actor = vtkh.create_actor_pts(pts=vtkpts, color=vtkh.create_color(plot['color']),
                                                   name=plot['name'], index=plot['idx'])
                vtk_actors.append(temp_actor)

            # Plot evaluated points
            if plot['type'] == 'evalpts' and self.vconf.display_evalpts:
                pts = np.array(plot['ptsarr'], dtype=np.float)
                vtkpts = numpy_to_vtk(pts, deep=False, array_type=VTK_FLOAT)
                vtkpts.SetName(plot['name'])
                temp_actor = vtkh.create_actor_pts(pts=vtkpts, color=vtkh.create_color(plot['color']),
                                                   name=plot['name'], index=plot['idx'])
                vtk_actors.append(temp_actor)

        # Render actors
        vtkh.create_render_window(vtk_actors, dict(KeyPressEvent=(self.vconf.keypress_callback, 1.0)),
                                  figure_size=self.vconf.figure_size)

class VisVoxel(vis.VisAbstract):
    """ VTK visualization module for voxel representation of the volumes. """
    def __init__(self, config=VisConfig()):
        super(VisVoxel, self).__init__(config=config)
        self._module_config['ctrlpts'] = "points"
        self._module_config['evalpts'] = "voxels"

    def render(self, **kwargs):
        """ Plots the volume and the control points. """
        # Calling parent function
        super(VisVoxel, self).render(**kwargs)

        # Initialize a list to store VTK actors
        vtk_actors = []

        # Start plotting
        for plot in self._plots:
            # Plot control points
            if plot['type'] == 'ctrlpts' and self.vconf.display_ctrlpts:
                # Points as spheres
                pts = np.array(plot['ptsarr'], dtype=np.float)
                vtkpts = numpy_to_vtk(pts, deep=False, array_type=VTK_FLOAT)
                vtkpts.SetName(plot['name'])
                temp_actor = vtkh.create_actor_pts(pts=vtkpts, color=vtkh.create_color(plot['color']),
                                                   name=plot['name'], index=plot['idx'])
                vtk_actors.append(temp_actor)

            # Plot evaluated points
            if plot['type'] == 'evalpts' and self.vconf.display_evalpts:
                faces = np.array(plot['ptsarr'][1], dtype=np.float)
                filled = np.array(plot['ptsarr'][2], dtype=np.int)
                grid_filled = faces[filled == 1]
                temp_actor = vtkh.create_actor_hexahedron(grid=grid_filled, color=vtkh.create_color(plot['color']),
                                                          name=plot['name'], index=plot['idx'])
                vtk_actors.append(temp_actor)

        # Render actors
        vtkh.create_render_window(vtk_actors, dict(KeyPressEvent=(self.vconf.keypress_callback, 1.0)),
                                  figure_size=self.vconf.figure_size)
<|endoftext|>


<commit_before>import os

HERE = os.path.abspath(os.path.dirname(__file__))

DEBUG = False
TEMPLATE_DEBUG = True

TIME_ZONE = 'Europe/Paris'

LANGUAGE_CODE = 'en-us'

SITE_ID = 1

USE_I18N = True
USE_L10N = True

MEDIA_ROOT = os.path.join(HERE, 'media')
MEDIA_URL = '/media/'

STATIC_ROOT = os.path.join(HERE, 'static')
STATIC_URL = '/static/'

STATICFILES_FINDERS = (
    'django.contrib.staticfiles.finders.FileSystemFinder',
    'django.contrib.staticfiles.finders.AppDirectoriesFinder',
)

TEMPLATE_LOADERS = (
    ('django.template.loaders.cached.Loader', (
        'django.template.loaders.filesystem.Loader',
        'django.template.loaders.app_directories.Loader',
    )),
)

MIDDLEWARE_CLASSES = (
    'django.middleware.common.CommonMiddleware',
    'django.contrib.sessions.middleware.SessionMiddleware',
    'django.middleware.csrf.CsrfViewMiddleware',
    'django.contrib.auth.middleware.AuthenticationMiddleware',
    'django.contrib.messages.middleware.MessageMiddleware',
)

ROOT_URLCONF = 'geoportail.urls'

TEMPLATE_DIRS = (
    os.path.join(HERE, 'templates'),
)

INSTALLED_APPS = (
    'django.contrib.auth',
    'django.contrib.contenttypes',
    'django.contrib.sessions',
    'django.contrib.messages',
    'django.contrib.staticfiles',
    'django.contrib.admin',
    'django.contrib.gis',
    'floppyforms',
    'raven',

    'geoportail.map',
    'geoportail.geonames',
)

LOGGING = {
    'version': 1,
    'disable_existing_loggers': False,
    'formatters': {
        'simple': {
            'format': '%(asctime)s %(levelname)s: %(message)s'
        },
    },
    'handlers': {
        'console': {
            'level': 'DEBUG',
            'class': 'logging.StreamHandler',
            'formatter': 'simple',
        },
        'sentry': {
            'level': 'INFO',
            'class': 'raven.contrib.django.handlers.SentryHandler',
        },
    },
    'loggers': {
        'django.request': {
            'handlers': ['console'],
            'level': 'ERROR',
            'propagate': True,
        },
        'geoportail': {
            'handlers': ['console', 'sentry'],
            'level': 'DEBUG',
        },
        'raven': {
            'handlers': ['console'],
            'level': 'DEBUG',
            'propagate': False,
        },
        'sentry.errors': {
            'handlers': ['console'],
            'level': 'DEBUG',
            'propagate': False,
        },
    }
}

# This is a public key from http://api.ign.fr/tech-docs-js/examples/
# Works for development but need a real key for deployment
GEOPORTAL_API_KEY = '1711091050407331029'
<commit_msg>Opt-in raven if SENTRY_DSN in env variables<commit_after>import os

HERE = os.path.abspath(os.path.dirname(__file__))

DEBUG = False
TEMPLATE_DEBUG = True

TIME_ZONE = 'Europe/Paris'

LANGUAGE_CODE = 'en-us'

SITE_ID = 1

USE_I18N = True
USE_L10N = True

MEDIA_ROOT = os.path.join(HERE, 'media')
MEDIA_URL = '/media/'

STATIC_ROOT = os.path.join(HERE, 'static')
STATIC_URL = '/static/'

STATICFILES_FINDERS = (
    'django.contrib.staticfiles.finders.FileSystemFinder',
    'django.contrib.staticfiles.finders.AppDirectoriesFinder',
)

TEMPLATE_LOADERS = (
    ('django.template.loaders.cached.Loader', (
        'django.template.loaders.filesystem.Loader',
        'django.template.loaders.app_directories.Loader',
    )),
)

MIDDLEWARE_CLASSES = (
    'django.middleware.common.CommonMiddleware',
    'django.contrib.sessions.middleware.SessionMiddleware',
    'django.middleware.csrf.CsrfViewMiddleware',
    'django.contrib.auth.middleware.AuthenticationMiddleware',
    'django.contrib.messages.middleware.MessageMiddleware',
)

ROOT_URLCONF = 'geoportail.urls'

TEMPLATE_DIRS = (
    os.path.join(HERE, 'templates'),
)

INSTALLED_APPS = (
    'django.contrib.auth',
    'django.contrib.contenttypes',
    'django.contrib.sessions',
    'django.contrib.messages',
    'django.contrib.staticfiles',
    'django.contrib.admin',
    'django.contrib.gis',
    'floppyforms',

    'geoportail.map',
    'geoportail.geonames',
)

if 'SENTRY_DSN' in os.environ:
    INSTALLED_APPS += (
        'raven',
    )

LOGGING = {
    'version': 1,
    'disable_existing_loggers': False,
    'formatters': {
        'simple': {
            'format': '%(asctime)s %(levelname)s: %(message)s'
        },
    },
    'handlers': {
        'console': {
            'level': 'DEBUG',
            'class': 'logging.StreamHandler',
            'formatter': 'simple',
        },
        'sentry': {
            'level': 'INFO',
            'class': 'raven.contrib.django.handlers.SentryHandler',
        },
    },
    'loggers': {
        'django.request': {
            'handlers': ['console'],
            'level': 'ERROR',
            'propagate': True,
        },
        'geoportail': {
            'handlers': ['console', 'sentry'],
            'level': 'DEBUG',
        },
        'raven': {
            'handlers': ['console'],
            'level': 'DEBUG',
            'propagate': False,
        },
        'sentry.errors': {
            'handlers': ['console'],
            'level': 'DEBUG',
            'propagate': False,
        },
    }
}

# This is a public key from http://api.ign.fr/tech-docs-js/examples/
# Works for development but need a real key for deployment
GEOPORTAL_API_KEY = '1711091050407331029'
<|endoftext|>


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

"""
State codes
"""

STATE_CODE_MAP = {
    'BW': 'Baden-Württemberg',
    'BY': 'Bayern',
    'BE': 'Berlin',
    'BB': 'Brandenburg',
    'HB': 'Bremen',
    'HH': 'Hamburg',
    'HE': 'Hessen',
    'MV': 'Mecklenburg-Vorpommern',
    'NI': 'Niedersachsen',
    'NW': 'Nordrhein-Westfalen',
    'RP': 'Rheinland-Pfalz',
    'SL': 'Saarland',
    'SN': 'Sachsen',
    'ST': 'Sachsen-Anhalt',
    'SH': 'Schleswig-Holstein',
    'TH': 'Thüringen',
}

INVERSE_STATE_CODE_MAP = {
    v: k for k, v in STATE_CODE_MAP.items()
}

INVERSE_STATE_CODE_MAP.update(
    {
        'Baden-Württemberg': 'BW',
        'Bavaria': 'BY',
        'Berlin': 'BE',
        'Brandenburg': 'BB',
        'Bremen': 'HB',
        'Hamburg': 'HH',
        'Hesse': 'HE',
        'Lower Saxony': 'NI',
        'Mecklenburg-Western Pomerania': 'MV',
        'North Rhine-Westphalia': 'NRW',
        'Northrhine-Westphalia': 'NRW',
        'Rhineland-Palatinate': 'RP',
        'Saarland': 'SL',
        'Saxony': 'SN',
        'Saxony-Anhalt': 'ST',
        'Schleswig-Holstein': 'SH',
        'Thuringia': 'TH',
    }
)

for k, v in tuple(INVERSE_STATE_CODE_MAP.items()):
    keys = [
        k.lower(),
        k.upper(),
        k.replace('-', ' '),
        k.lower().replace('-', ' '),
        k.upper().replace('-', ' '),
    ]
    for k in keys:
        INVERSE_STATE_CODE_MAP[k] = v

class StateCodeError(KeyError):
    def __init__(self, message='', state_code=None):
        if state_code:
            code_str = '\n'.join('{} => {}'.format(k, v)
                                 for k, v in STATE_CODE_MAP.items())
            state_mode_msg = (
                'State code {} is not valid. Must be one of:\n{}'
            ).format(state_code, code_str)
            message += state_mode_msg
        super(StateCodeError, self).__init__(message)
<commit_msg>BUGFIX: wrong state code NRW<commit_after>#!/usr/bin/env python

"""
State codes
"""

STATE_CODE_MAP = {
    'BW': 'Baden-Württemberg',
    'BY': 'Bayern',
    'BE': 'Berlin',
    'BB': 'Brandenburg',
    'HB': 'Bremen',
    'HH': 'Hamburg',
    'HE': 'Hessen',
    'MV': 'Mecklenburg-Vorpommern',
    'NI': 'Niedersachsen',
    'NW': 'Nordrhein-Westfalen',
    'RP': 'Rheinland-Pfalz',
    'SL': 'Saarland',
    'SN': 'Sachsen',
    'ST': 'Sachsen-Anhalt',
    'SH': 'Schleswig-Holstein',
    'TH': 'Thüringen',
}

INVERSE_STATE_CODE_MAP = {
    v: k for k, v in STATE_CODE_MAP.items()
}

INVERSE_STATE_CODE_MAP.update(
    {
        'Baden-Württemberg': 'BW',
        'Bavaria': 'BY',
        'Berlin': 'BE',
        'Brandenburg': 'BB',
        'Bremen': 'HB',
        'Hamburg': 'HH',
        'Hesse': 'HE',
        'Lower Saxony': 'NI',
        'Mecklenburg-Western Pomerania': 'MV',
        'North Rhine-Westphalia': 'NW',
        'Northrhine-Westphalia': 'NW',
        'Rhineland-Palatinate': 'RP',
        'Saarland': 'SL',
        'Saxony': 'SN',
        'Saxony-Anhalt': 'ST',
        'Schleswig-Holstein': 'SH',
        'Thuringia': 'TH',
    }
)

for k, v in tuple(INVERSE_STATE_CODE_MAP.items()):
    keys = [
        k.lower(),
        k.upper(),
        k.replace('-', ' '),
        k.lower().replace('-', ' '),
        k.upper().replace('-', ' '),
    ]
    for k in keys:
        INVERSE_STATE_CODE_MAP[k] = v

class StateCodeError(KeyError):
    def __init__(self, message='', state_code=None):
        if state_code:
            code_str = '\n'.join('{} => {}'.format(k, v)
                                 for k, v in STATE_CODE_MAP.items())
            state_mode_msg = (
                'State code {} is not valid. Must be one of:\n{}'
            ).format(state_code, code_str)
            message += state_mode_msg
        super(StateCodeError, self).__init__(message)
<|endoftext|>