<commit_before># Copyright 2016 by Raytheon BBN Technologies Corp.  All Rights Reserved.

# See get_sequence_function (the main entrypoint), and
# for a sample usage, see main.py.

import ast
import os
import sys

from pyqgl2.ast_util import ast2str, NodeError
from pyqgl2.quickcopy import quickcopy
from pyqgl2.qreg import is_qbit_create

class SequenceExtractor(object):
    """
    Create QGL1 code from a modified AST

    This class has the logic to take the input QGL2 AST as modified by the compiler,
    and produce a QGL1 function reference suitable for execution.

    Note: this assumes that the AST is for one function
    definition that has already been inlined, successfully
    flattened, grouped, and sequenced already.
    """

    def __init__(self, importer, allocated_qregs):

        self.importer = importer
        self.allocated_qregs = allocated_qregs

        self.qbits = set()
        self.qbit_creates = list() # expressions that create Qubits
        self.sequence = []

        # the imports we need to make in order to satisfy the stubs
        #
        # the key is the name of the module (i.e. something like
        # 'QGL.PulsePrimitives') and the values are sets of import
        # clauses (i.e. 'foo' or 'foo as bar')
        #
        self.stub_imports = dict()

    def find_imports(self, node):
        '''Fill in self.stub_imports with all the per module imports needed'''

        default_namespace = node.qgl_fname

        for subnode in ast.walk(node):
            if (isinstance(subnode, ast.Call) and
                    isinstance(subnode.func, ast.Name)):
                funcname = subnode.func.id

                # QRegister calls will be stripped by find_sequences, so skip
                if funcname == 'QRegister':
                    continue

                # If we created a node without an qgl_fname,
                # then use the default namespace instead.
                # FIXME: This is a hack, but it will work for now.
                #
                if not hasattr(subnode, 'qgl_fname'):
                    namespace = default_namespace
                else:
                    namespace = subnode.qgl_fname

                if hasattr(subnode, 'qgl_implicit_import'):
                    (sym_name, module_name, orig_name) = \
                            subnode.qgl_implicit_import
                else:
                    fdef = self.importer.resolve_sym(namespace, funcname)

                    if not fdef:
                        NodeError.error_msg(subnode,
                                'cannot find import info for [%s]' % funcname)
                        return False
                    elif not fdef.qgl_stub_import:
                        NodeError.error_msg(subnode,
                                'not a stub: [%s]' % funcname)
                        return False

                    (sym_name, module_name, orig_name) = fdef.qgl_stub_import

                if orig_name:
                    import_str = '%s as %s' % (orig_name, sym_name)
                else:
                    import_str = sym_name

                if module_name not in self.stub_imports:
                    self.stub_imports[module_name] = set()

                self.stub_imports[module_name].add(import_str)

        return True

    def create_imports_list(self):
        '''Using the stub_imports created by find_imports,
        create a list of import statement strings to include in the function.'''

        import_list = list()

        for module_name in sorted(self.stub_imports.keys()):
            for sym_name in sorted(self.stub_imports[module_name]):
                import_list.append(
                        'from %s import %s' % (module_name, sym_name))

        return import_list

    def qbits_from_qregs(self, allocated_qregs):
        qbits = set()
        for qreg in allocated_qregs.values():
            qbits.update(qreg.qubits)
        return qbits

    def expand_arg(self, arg):
        expanded_args = []
        if isinstance(arg, ast.Name) and arg.id in self.allocated_qregs:
            qreg = self.allocated_qregs[arg.id]
            # add an argument for each constituent qubit in the QRegister
            for n in range(len(qreg)):
                new_arg = ast.Name(id=qreg.use_name(n), ctx=ast.Load())
                expanded_args.append(new_arg)
        elif (isinstance(arg, ast.Subscript) and
                arg.value.id in self.allocated_qregs):
            # add an argument for the subset of qubits referred to
            # by the QReference
            # eval the subscript to extract the slice
            qreg = self.allocated_qregs[arg.value.id]
            try:
                qref = eval(ast2str(arg), None, self.allocated_qregs)
            except:
                NodeError(arg,
                    "Error evaluating QReference [%s]" % ast2str(arg))
            # convert the slice into a list of indices
            idx = range(len(qreg))[qref.idx]
            if not hasattr(idx, '__iter__'):
                idx = (idx,)
            for n in idx:
                new_arg = ast.Name(id=qreg.use_name(n), ctx=ast.Load())
                expanded_args.append(new_arg)
        elif isinstance(arg, (ast.Tuple, ast.List)):
            # expand tuples and lists to include all constituent qubits
            # FIXME would be safer to create a set() of qubits
            expanded_elts = []
            for elt in arg.elts:
                expanded_elts.extend(self.expand_arg(elt))
            new_arg = quickcopy(arg)
            new_arg.elts = expanded_elts
            expanded_args.append(new_arg)
        else:
            # don't expand it
            expanded_args.append(arg)
        return expanded_args

    def expand_qreg_call(self, node):
        new_stmnts = []
        expanded_args = [self.expand_arg(a) for a in node.value.args]
        # TODO verify that QRegister lengths match
        for args in zip(*expanded_args):
            stmnt = quickcopy(node)
            stmnt.value.args = args
            new_stmnts.append(stmnt)

        return new_stmnts

    def find_sequences(self, node):
        '''
        Input AST node is the main function definition.
        Strips out QRegister creation statements and builds
        a list of corresponding Qubit creation statements.
        Converts calls on QRegisters into calls on Qubits.'''

        if not isinstance(node, ast.FunctionDef):
            NodeError.fatal_msg(node, 'not a function definition')
            return False

        self.qbits = self.qbits_from_qregs(self.allocated_qregs)
        for q in self.qbits:
            stmnt = ast.parse("QBIT_{0} = QubitFactory('q{0}')".format(q))
            self.qbit_creates.append(stmnt)

        lineNo = -1
        while lineNo+1 < len(node.body):
            lineNo += 1
            # print("Line %d of %d" % (lineNo+1, len(node.body)))
            stmnt = node.body[lineNo]
            # print("Looking at stmnt %s" % stmnt)
            if is_qbit_create(stmnt):
                # drop it
                continue

            elif isinstance(stmnt, ast.Expr):
                # expand calls on QRegisters into calls on Qubits
                if (hasattr(stmnt, 'qgl2_type') and
                        (stmnt.qgl2_type == 'stub' or stmnt.qgl2_type == 'measurement')):
                    new_stmnts = self.expand_qreg_call(stmnt)
                    self.sequence.extend(new_stmnts)
                else:
                    self.sequence.append(stmnt)
            else:
                NodeError.error_msg(stmnt,
                                    'orphan statement %s' % ast.dump(stmnt))

        # print("Seqs: %s" % self.sequences)
        if not self.sequence:
            NodeError.warning_msg(node, "No qubit operations discovered")
            return False

        return True

    def emit_function(self, func_name='qgl1_main', setup=None):
        """
        Create a function that, when run, creates the context
        in which the sequence is evaluated, and evaluate it.

        Assumes find_imports and find_sequences have already
        been called.

        func_name is the name for the function, if provided.
        I'm not certain that the name has any significance
        or whether this function will be, for all practical
        purposes, a lambda.
        """

        # assumes a standard 4-space indent; if the surrounding
        # context uses some other indentation scheme, the interpreter
        # may gripe, and pep8 certainly will
        #
        indent = '    '

        # FIXME: Ugliness
        # In the proper namespace we need to import all the QGL1 functions
        # that this method is using / might use.
        # Here we include the imports matching stuff in qgl2.qgl1.py
        # as required by create_imports_list(), plus
        # extras as required.

        base_imports = indent + 'from QGL import QubitFactory\n'

        found_imports = ('\n' + indent).join(self.create_imports_list())

        # define the method preamble
        preamble = 'def %s():\n' % func_name
        preamble += base_imports
        preamble += indent + found_imports
        preamble += '\n\n'

        # FIXME: In below block, calls to ast2str are the slowest part
        # (78%) of calling get_sequence_function. Fixable?

        for node in self.qbit_creates:
            preamble += indent + ast2str(node).strip() + '\n'

        if setup:
            for setup_stmnt in setup:
                preamble += indent + ('%s\n' % ast2str(setup_stmnt).strip())

        sequence = [ast2str(item).strip() for item in self.sequence]

        # TODO there must be a more elegant way to indent this properly
        seq_str = indent + 'seq = [\n' + 2 * indent
        seq_str += (',\n' + 2 * indent).join(sequence)
        seq_str += '\n' + indent + ']\n'

        postamble = indent + 'return seq\n'

        res =  preamble + seq_str + postamble
        return res

def get_sequence_function(node, func_name, importer, allocated_qregs,
        intermediate_fout=None, saveOutput=False, filename=None,
        setup=None):
    """
    Create a function that encapsulates the QGL code
    from the given AST node, which is presumed to already
    be fully pre-processed.

    TODO: we don't test that the node is fully pre-processed.
    TODO: each step of the preprocessor should mark the nodes
    so that we know whether or not they've been processed.
    """

    builder = SequenceExtractor(importer, allocated_qregs)

    builder.find_sequences(node)
    builder.find_imports(node)
    code = builder.emit_function(func_name, setup)
    if intermediate_fout:
        print(('#start function\n%s\n#end function' % code),
              file=intermediate_fout, flush=True)
    if saveOutput and filename:
        newf = os.path.abspath(filename[:-3] + "qgl1.py")
        with open(newf, 'w') as compiledFile:
            compiledFile.write(code)
        print("Saved compiled code to %s" % newf)

    NodeError.diag_msg(
            node, 'generated code:\n#start\n%s\n#end code' % code)

    # TODO: we might want to pass in elements of the local scope
    scratch_scope = dict()
    eval(compile(code, '<none>', mode='exec'), globals(), scratch_scope)
    NodeError.halt_on_error()

    return scratch_scope[func_name]
<commit_msg>Add some comments.<commit_after># Copyright 2016 by Raytheon BBN Technologies Corp.  All Rights Reserved.

# See get_sequence_function (the main entrypoint), and
# for a sample usage, see main.py.

import ast
import os
import sys

from pyqgl2.ast_util import ast2str, NodeError
from pyqgl2.quickcopy import quickcopy
from pyqgl2.qreg import is_qbit_create

class SequenceExtractor(object):
    """
    Create QGL1 code from a modified AST

    This class has the logic to take the input QGL2 AST as modified by the compiler,
    and produce a QGL1 function reference suitable for execution.

    Note: this assumes that the AST is for one function
    definition that has already been inlined, successfully
    flattened, grouped, and sequenced already.
    """

    def __init__(self, importer, allocated_qregs):

        self.importer = importer
        self.allocated_qregs = allocated_qregs

        self.qbits = set()
        self.qbit_creates = list() # expressions that create Qubits
        self.sequence = []

        # the imports we need to make in order to satisfy the stubs
        #
        # the key is the name of the module (i.e. something like
        # 'QGL.PulsePrimitives') and the values are sets of import
        # clauses (i.e. 'foo' or 'foo as bar')
        #
        self.stub_imports = dict()

    def find_imports(self, node):
        '''Fill in self.stub_imports with all the per module imports needed'''

        default_namespace = node.qgl_fname

        for subnode in ast.walk(node):
            if (isinstance(subnode, ast.Call) and
                    isinstance(subnode.func, ast.Name)):
                funcname = subnode.func.id

                # QRegister calls will be stripped by find_sequences, so skip
                if funcname == 'QRegister':
                    continue

                # If we created a node without an qgl_fname,
                # then use the default namespace instead.
                # FIXME: This is a hack, but it will work for now.
                #
                if not hasattr(subnode, 'qgl_fname'):
                    namespace = default_namespace
                else:
                    namespace = subnode.qgl_fname

                if hasattr(subnode, 'qgl_implicit_import'):
                    (sym_name, module_name, orig_name) = \
                            subnode.qgl_implicit_import
                else:
                    fdef = self.importer.resolve_sym(namespace, funcname)

                    if not fdef:
                        NodeError.error_msg(subnode,
                                'cannot find import info for [%s]' % funcname)
                        return False
                    elif not fdef.qgl_stub_import:
                        NodeError.error_msg(subnode,
                                'not a stub: [%s]' % funcname)
                        return False

                    (sym_name, module_name, orig_name) = fdef.qgl_stub_import

                if orig_name:
                    import_str = '%s as %s' % (orig_name, sym_name)
                else:
                    import_str = sym_name

                if module_name not in self.stub_imports:
                    self.stub_imports[module_name] = set()

                self.stub_imports[module_name].add(import_str)

        return True

    def create_imports_list(self):
        '''Using the stub_imports created by find_imports,
        create a list of import statement strings to include in the function.'''

        import_list = list()

        for module_name in sorted(self.stub_imports.keys()):
            for sym_name in sorted(self.stub_imports[module_name]):
                import_list.append(
                        'from %s import %s' % (module_name, sym_name))

        return import_list

    def qbits_from_qregs(self, allocated_qregs):
        '''
        Returns the set of all qubits contained within allocated_qregs.
        '''
        qbits = set()
        for qreg in allocated_qregs.values():
            qbits.update(qreg.qubits)
        return qbits

    def expand_arg(self, arg):
        '''
        Expands a single argument to a stub call. QRegisters are expanded
        to a list of constituent Qubits. QRegister subscripts are similar,
        except that the slice selects which Qubits to return. Tuples and
        lists are different in that the exansion happens "within" the
        tuple. So, given
            a = QRegister(2)
            b = QRegister(1)
        we do (in AST shorthand):
            expand_arg("a") -> ["QBIT_1", "QBIT_2"]
            expand_arg("a[1]") -> ["QBIT_2"]
            expand_arg("(a,b)") -> [("QBIT_1", "QBIT_2", "QBIT_3")]
        '''
        expanded_args = []
        if isinstance(arg, ast.Name) and arg.id in self.allocated_qregs:
            qreg = self.allocated_qregs[arg.id]
            # add an argument for each constituent qubit in the QRegister
            for n in range(len(qreg)):
                new_arg = ast.Name(id=qreg.use_name(n), ctx=ast.Load())
                expanded_args.append(new_arg)
        elif (isinstance(arg, ast.Subscript) and
                arg.value.id in self.allocated_qregs):
            # add an argument for the subset of qubits referred to
            # by the QReference
            # eval the subscript to extract the slice
            qreg = self.allocated_qregs[arg.value.id]
            try:
                qref = eval(ast2str(arg), None, self.allocated_qregs)
            except:
                NodeError(arg,
                    "Error evaluating QReference [%s]" % ast2str(arg))
            # convert the slice into a list of indices
            idx = range(len(qreg))[qref.idx]
            if not hasattr(idx, '__iter__'):
                idx = (idx,)
            for n in idx:
                new_arg = ast.Name(id=qreg.use_name(n), ctx=ast.Load())
                expanded_args.append(new_arg)
        elif isinstance(arg, (ast.Tuple, ast.List)):
            # expand tuples and lists to include all constituent qubits
            # FIXME would be safer to create a set() of qubits
            expanded_elts = []
            for elt in arg.elts:
                expanded_elts.extend(self.expand_arg(elt))
            new_arg = quickcopy(arg)
            new_arg.elts = expanded_elts
            expanded_args.append(new_arg)
        else:
            # don't expand it
            expanded_args.append(arg)
        return expanded_args

    def expand_qreg_call(self, node):
        '''
        Expands calls on stubs to element-wise broadcast over QRegister
        arguments. So that given
            a = QRegister(2)
            b = QRegister(2)
        single-qubit calls expand like:
            X(a)
        becomes
            X(QBIT_1)
            X(QBIT_2)
        and two-qubit calls expand like:
            CNOT(a, b)
        becomes
            CNOT(QBIT_1, QBIT_3)
            CNOT(QBIT_2, QBIT_4)
        Calls that act on tuples or lists expand inside the tuple, i.e.
            Barrier("", (a, b))
        becomes
            Barrier("", (QBIT_1, QBIT_2, QBIT_3, QBIT_4))

        Returns a list of ast.Call nodes with appropriate argumnet
        substitutions.
        '''
        new_stmnts = []
        expanded_args = [self.expand_arg(a) for a in node.value.args]
        # TODO verify that QRegister lengths match
        for args in zip(*expanded_args):
            stmnt = quickcopy(node)
            stmnt.value.args = args
            new_stmnts.append(stmnt)

        return new_stmnts

    def find_sequences(self, node):
        '''
        Input AST node is the main function definition.
        Strips out QRegister creation statements and builds
        a list of corresponding Qubit creation statements.
        Converts calls on QRegisters into calls on Qubits.'''

        if not isinstance(node, ast.FunctionDef):
            NodeError.fatal_msg(node, 'not a function definition')
            return False

        self.qbits = self.qbits_from_qregs(self.allocated_qregs)
        for q in self.qbits:
            stmnt = ast.parse("QBIT_{0} = QubitFactory('q{0}')".format(q))
            self.qbit_creates.append(stmnt)

        lineNo = -1
        while lineNo+1 < len(node.body):
            lineNo += 1
            # print("Line %d of %d" % (lineNo+1, len(node.body)))
            stmnt = node.body[lineNo]
            # print("Looking at stmnt %s" % stmnt)
            if is_qbit_create(stmnt):
                # drop it
                continue

            elif isinstance(stmnt, ast.Expr):
                # expand calls on QRegisters into calls on Qubits
                if (hasattr(stmnt, 'qgl2_type') and
                        (stmnt.qgl2_type == 'stub' or stmnt.qgl2_type == 'measurement')):
                    new_stmnts = self.expand_qreg_call(stmnt)
                    self.sequence.extend(new_stmnts)
                else:
                    self.sequence.append(stmnt)
            else:
                NodeError.error_msg(stmnt,
                                    'orphan statement %s' % ast.dump(stmnt))

        # print("Seqs: %s" % self.sequences)
        if not self.sequence:
            NodeError.warning_msg(node, "No qubit operations discovered")
            return False

        return True

    def emit_function(self, func_name='qgl1_main', setup=None):
        """
        Create a function that, when run, creates the context
        in which the sequence is evaluated, and evaluate it.

        Assumes find_imports and find_sequences have already
        been called.

        func_name is the name for the function, if provided.
        I'm not certain that the name has any significance
        or whether this function will be, for all practical
        purposes, a lambda.
        """

        # assumes a standard 4-space indent; if the surrounding
        # context uses some other indentation scheme, the interpreter
        # may gripe, and pep8 certainly will
        #
        indent = '    '

        # FIXME: Ugliness
        # In the proper namespace we need to import all the QGL1 functions
        # that this method is using / might use.
        # Here we include the imports matching stuff in qgl2.qgl1.py
        # as required by create_imports_list(), plus
        # extras as required.

        base_imports = indent + 'from QGL import QubitFactory\n'

        found_imports = ('\n' + indent).join(self.create_imports_list())

        # define the method preamble
        preamble = 'def %s():\n' % func_name
        preamble += base_imports
        preamble += indent + found_imports
        preamble += '\n\n'

        # FIXME: In below block, calls to ast2str are the slowest part
        # (78%) of calling get_sequence_function. Fixable?

        for node in self.qbit_creates:
            preamble += indent + ast2str(node).strip() + '\n'

        if setup:
            for setup_stmnt in setup:
                preamble += indent + ('%s\n' % ast2str(setup_stmnt).strip())

        sequence = [ast2str(item).strip() for item in self.sequence]

        # TODO there must be a more elegant way to indent this properly
        seq_str = indent + 'seq = [\n' + 2 * indent
        seq_str += (',\n' + 2 * indent).join(sequence)
        seq_str += '\n' + indent + ']\n'

        postamble = indent + 'return seq\n'

        res =  preamble + seq_str + postamble
        return res

def get_sequence_function(node, func_name, importer, allocated_qregs,
        intermediate_fout=None, saveOutput=False, filename=None,
        setup=None):
    """
    Create a function that encapsulates the QGL code
    from the given AST node, which is presumed to already
    be fully pre-processed.

    TODO: we don't test that the node is fully pre-processed.
    TODO: each step of the preprocessor should mark the nodes
    so that we know whether or not they've been processed.
    """

    builder = SequenceExtractor(importer, allocated_qregs)

    builder.find_sequences(node)
    builder.find_imports(node)
    code = builder.emit_function(func_name, setup)
    if intermediate_fout:
        print(('#start function\n%s\n#end function' % code),
              file=intermediate_fout, flush=True)
    if saveOutput and filename:
        newf = os.path.abspath(filename[:-3] + "qgl1.py")
        with open(newf, 'w') as compiledFile:
            compiledFile.write(code)
        print("Saved compiled code to %s" % newf)

    NodeError.diag_msg(
            node, 'generated code:\n#start\n%s\n#end code' % code)

    # TODO: we might want to pass in elements of the local scope
    scratch_scope = dict()
    eval(compile(code, '<none>', mode='exec'), globals(), scratch_scope)
    NodeError.halt_on_error()

    return scratch_scope[func_name]
<|endoftext|>


<commit_before>#!/usr/bin/python

from snappy import *

# Returns a VolumeData object for manifolds with the given names
# Volumes' precision is based on pari, so set it there
def get_volume_data(man_nms, engine = 'magma', retrieve = True):
    recs = dict()
    for nm in man_nms:
        try:
            var = Manifold(nm).ptolemy_variety(2,'all')
            try:
                if retrieve:
                    sols = var.retrieve_solutions()
                else:
                    raise Exception('Coding too lazy!')
            except:
                sols = var.compute_solutions(engine = engine)
            data = [(c.number_field(),c.volume_numerical()) for c in sols.flatten()]
            for d in data:
                for v in d[1]:
                    recs.setdefault(str(d[0]),list()).append((str(v),nm))
            for k in recs.keys():   # remove duplicates
                recs[k] = list(set(recs[k]))
        except Exception as e:
            print(str(e))+'; skipping '+nm
            continue
    return VolumeData(data = recs)

def get_potential_trace_fields(poly):
    pol = pari(poly)
    try:
        return [str(rec[0]) for rec in pol.nfsubfields()[1:] if _binmiss(rec[0].poldegree(),pol.poldegree())]   # poldegree returns int
    except: # we want cypari.gen.PariError, but no idea how to reference; fortunately, anything else will just raise again
        pol = pol.polredabs()
        return get_potential_trace_fields(str(pol))

def _binmiss(s,l):
    while s < l:
        s *= 2
    return s == l

# Wrapper for manipulating data on pseudo-volumes
class VolumeData:

    # structure: dict poly ---> list of (volume, manifold, ?fitted)
    def __init__(self, data = dict()):
        self.data = data

    def get_polys(self):
        return self.data.keys()

    def get_volumes(self,poly):
        return [rec[0] for rec in self.data[poly]]

    def get_manifolds(self,poly):
        return [rec[1] for rec in self.data[poly]]

    def get_volume_data(self,poly):
        return self.data[poly]

    # returns a VolumeData object containing the data from this and other; in case of a conflict, other's data takes precendence
    def combine_with(self,other):
        new_data = dict(self.data)
        for p in other.get_polys():
            new_data.setdefault(p,dict())
            new_data[p] = list(set(new_data[p].extend(other.get_volume_data(p))))
        return VolumeData(data = new_data)

    # given an (open, ready to write) file object or valid filename, writes the data
    def write_to_csv(self, output_file, seperator = ';', append = False):
        try:
            if type(output_file) == str:
                if append:
                    f = open(output_file,'a')
                else:
                    f = open(output_file,'w')
            else:
                f = output_file
            f.write('"TraceField"'+seperator+'"Volume"'+seperator+'"Manifold"'+'\n')
            for p in self.get_polys():
                for rec in self.get_volume_data(p):
                    f.write('"'+p+'"'+seperator)
                    f.write('"'+rec[0]+'"'+seperator)
                    f.write('"'+rec[1]+'"\n')
        finally:
            if type(output_file) == str and f:
                f.close()

    # given an (open, ready to read data) file object or valid filename, reads the file and returns a VolumeData that would write it
def read_volumedata_csv(infile, seperator = ';'):
    try:
        if type(infile) ==  str:
            f = open(infile,'r')
            f.readline()
        else:
            f = infile
        data = dict()
        for l in f.readlines():
            w = l.strip('\n').replace('"','').split(seperator)
            data.setdefault(w[0],list()).append(w[1:])
        return VolumeData(data = data)
    finally:
        if type(infile) == str and f:
            f.close()
<commit_msg>Fixed a related mistake.<commit_after>#!/usr/bin/python

from snappy import *

# Returns a VolumeData object for manifolds with the given names
# Volumes' precision is based on pari, so set it there
def get_volume_data(man_nms, engine = 'magma', retrieve = True):
    recs = dict()
    for nm in man_nms:
        try:
            var = Manifold(nm).ptolemy_variety(2,'all')
            try:
                if retrieve:
                    sols = var.retrieve_solutions()
                else:
                    raise Exception('Coding too lazy!')
            except:
                sols = var.compute_solutions(engine = engine)
            data = [(c.number_field(),c.volume_numerical()) for c in sols.flatten()]
            for d in data:
                for v in d[1]:
                    recs.setdefault(str(d[0]),list()).append((str(v),nm))
            for k in recs.keys():   # remove duplicates
                recs[k] = list(set(recs[k]))
        except Exception as e:
            print(str(e))+'; skipping '+nm
            continue
    return VolumeData(data = recs)

def get_potential_trace_fields(poly):
    pol = pari(poly)
    try:
        return [str(rec[0]) for rec in pol.nfsubfields()[1:] if _binmiss(rec[0].poldegree(),pol.poldegree())]   # poldegree returns int
    except: # we want cypari.gen.PariError, but no idea how to reference; fortunately, anything else will just raise again
        pol = pol.polredabs()
        return get_potential_trace_fields(str(pol))

def _binmiss(s,l):
    while s < l:
        s *= 2
    return s == l

# Wrapper for manipulating data on pseudo-volumes
class VolumeData:

    # structure: dict poly ---> list of (volume, manifold, ?fitted)
    def __init__(self, data = dict()):
        self.data = data

    def get_polys(self):
        return self.data.keys()

    def get_volumes(self,poly):
        return [rec[0] for rec in self.data[poly]]

    def get_manifolds(self,poly):
        return [rec[1] for rec in self.data[poly]]

    def get_volume_data(self,poly):
        return self.data[poly]

    # returns a VolumeData object containing the data from this and other; in case of a conflict, other's data takes precendence
    def combine_with(self,other):
        new_data = dict(self.data)
        for p in other.get_polys():
            new_data.setdefault(p,dict())
            new_data[p] = list(set(new_data[p].extend(other.get_volume_data(p))))
        return VolumeData(data = new_data)

    # given an (open, ready to write) file object or valid filename, writes the data
    def write_to_csv(self, output_file, seperator = ';', append = False):
        try:
            if type(output_file) == str:
                if append:
                    f = open(output_file,'a')
                else:
                    f = open(output_file,'w')
            else:
                f = output_file
            if not append:
                f.write('"TraceField"'+seperator+'"Volume"'+seperator+'"Manifold"'+'\n')
            for p in self.get_polys():
                for rec in self.get_volume_data(p):
                    f.write('"'+p+'"'+seperator)
                    f.write('"'+rec[0]+'"'+seperator)
                    f.write('"'+rec[1]+'"\n')
        finally:
            if type(output_file) == str and f:
                f.close()

    # given an (open, ready to read data) file object or valid filename, reads the file and returns a VolumeData that would write it
def read_volumedata_csv(infile, seperator = ';'):
    try:
        if type(infile) ==  str:
            f = open(infile,'r')
            f.readline()
        else:
            f = infile
        data = dict()
        for l in f.readlines():
            w = l.strip('\n').replace('"','').split(seperator)
            data.setdefault(w[0],list()).append(w[1:])
        return VolumeData(data = data)
    finally:
        if type(infile) == str and f:
            f.close()
<|endoftext|>


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

from openfisca_france.model.base import *  # noqa analysis:ignore

###################################################
## PLAN :
# 1) Gains taxés comme des salaires
# 2) Autres plus-values (exonérées, au barème, ...)
# 3) Abattements sur plus-values
# 4) Plus-values taxées forfaitairement
# 5) Autres variables
###################################################

# Gains taxés comme des salaires

class f1tv(Variable):
    cerfa_field = {
        0: u"1TV",
        1: u"1UV",
        }
    value_type = int
    unit = 'currency'
    entity = Individu
    label = u"Gains de levée d'options sur titres en cas de cession ou de conversion au porteur dans le délai d'indisponibilité: entre 1 et 2 ans"
    end = '2015-12-31'
    definition_period = YEAR

class f1tw(Variable):
    cerfa_field = {
        0: u"1TW",
        1: u"1UW",
        }
    value_type = int
    unit = 'currency'
    entity = Individu
    label = u"Gains de levée d'options sur titres en cas de cession ou de conversion au porteur dans le délai d'indisponibilité: entre 2 et 3 ans"
    end = '2015-12-31'
    definition_period = YEAR

class f1tx(Variable):
    cerfa_field = {
        0: u"1TX",
        1: u"1UX",
        }
    value_type = int
    unit = 'currency'
    entity = Individu
    label = u"Gains de levée d'options sur titres en cas de cession ou de conversion au porteur dans le délai d'indisponibilité: entre 3 et 4 ans"
    end = '2016-12-31'
    definition_period = YEAR

class f3vj(Variable):
    cerfa_field = {
        0: u"3VJ",
        1: u"3VK",
        }
    value_type = int
    unit = 'currency'
    entity = Individu
    label = u"Gains imposables sur option dans la catégorie des salaires"
    definition_period = YEAR

# Autres plus-values

class f3vg(Variable):
    cerfa_field = u"3VG"
    value_type = int
    unit = 'currency'
    entity = FoyerFiscal
    label = u"Plus-value imposable sur gains de cession de valeurs mobilières, de droits sociaux et gains assimilés"
    definition_period = YEAR

class f3vh(Variable):
    cerfa_field = u"3VH"
    value_type = int
    unit = 'currency'
    entity = FoyerFiscal
    label = u"Moins-value imposable su gains de cession de valeurs mobilières, de droits sociaux et gains assimilés"
    definition_period = YEAR

class f3vv(Variable):
    cerfa_field = u"3VV"
    value_type = int
    unit = 'currency'
    entity = FoyerFiscal
    label = u"Plus-values réalisées par les non-résidents: montant du prélèvement de 45 % déjà versé"
    # start_date = date(2013, 1, 1)
    end = '2013-12-31'
    definition_period = YEAR

class f3vz(Variable):
    cerfa_field = u"3VZ"
    value_type = int
    unit = 'currency'
    entity = FoyerFiscal
    label = u"Plus-values imposables sur cessions d’immeubles ou de biens meubles"
    # start_date = date(2011, 1, 1)
    definition_period = YEAR

class f3vc(Variable):
    cerfa_field = u"3VC"
    value_type = int
    unit = 'currency'
    entity = FoyerFiscal
    label = u"Produits et plus-values exonérés provenant de structure de capital-risque"
    # start_date = date(2006, 1, 1)
    definition_period = YEAR

class f3vp(Variable):
    cerfa_field = u"3VP"
    value_type = int
    unit = 'currency'
    entity = FoyerFiscal
    label = u"Plus-values exonérées de cessions de titres de jeunes entreprises innovantes"
    # start_date = date(2007, 1, 1)
    end = '2013-12-31'
    definition_period = YEAR

class f3vy(Variable):
    cerfa_field = u"3VY"
    value_type = int
    unit = 'currency'
    entity = FoyerFiscal
    label = u"Plus-values exonérées de cessions de participations supérieures à 25% au sein du groupe familial"
    # start_date = date(2011, 1, 1)
    end = '2013-12-31'
    definition_period = YEAR

class f3ve(Variable):
    cerfa_field = u"3VE"
    value_type = int
    unit = 'currency'
    entity = FoyerFiscal
    label = u"Plus-values réalisées par les non-résidents pour lesquelles vous demandez le remboursement de l'excédent du prélèvement de 45 %"
    definition_period = YEAR

class f3sb(Variable):
    cerfa_field = u"3SB"
    value_type = int
    unit = 'currency'
    entity = FoyerFiscal
    label = u"Plus-values en report d'imposition, dont le report a expiré cette année"
    # start_date = date(2013, 1, 1)
    definition_period = YEAR

class f3wb(Variable):
    cerfa_field = u"3WB"
    value_type = int
    unit = 'currency'
    entity = FoyerFiscal
    label = u"Plus-values des individus transférant leur domicile fiscal hors de France; plus-values et créances ne bénéficiant pas du sursis de paiement; plus-values nettes imposables au barème"
    # start_date = date(2013, 1, 1)
    definition_period = YEAR

class f3wd(Variable):
    cerfa_field = u"3WD"
    value_type = int
    unit = 'currency'
    entity = FoyerFiscal
    label = u"Plus-values des individus transférant leur domicile fiscal hors de France; plus-values et créances ne bénéficiant pas du sursis de paiement; plus-values de base soumise aux prélèvements sociaux"
    # start_date = date(2013, 1, 1)
    definition_period = YEAR

# Abattements sur plus-values

class f3va(Variable):
    cerfa_field = u"3VA"
    value_type = int
    unit = 'currency'
    entity = FoyerFiscal
    label = u"Abattements nets (abattement pour durée de détention renforcé et abattement fixe spécial) appliqués sur des plus-values réalisées par les dirigeants de PME lors de leur départ à la retraite"
    # start_date = date(2006, 1, 1)
    definition_period = YEAR

class f3vb(Variable):
    cerfa_field = u"3VB"
    value_type = int
    unit = 'currency'
    entity = FoyerFiscal
    label = u"Abattements nets (abattement pour durée de détention renforcé et abattement fixe spécial) appliqués sur des moins-values réalisées par les dirigeants de PME lors de leur départ à la retraite"
    # start_date = date(2006, 1, 1)
    definition_period = YEAR

class f3sg(Variable):
    cerfa_field = u"3SG"
    value_type = int
    unit = 'currency'
    entity = FoyerFiscal
    label = u"Abattement net pour durée de détention : appliqué sur des plus-values"
    # start_date = date(2013, 1, 1)
    definition_period = YEAR

class f3sh(Variable):
    cerfa_field = u"3SH"
    value_type = int
    unit = 'currency'
    entity = FoyerFiscal
    label = u"Abattement net pour durée de détention : appliqué sur des moins-values"
    # start_date = date(2013, 1, 1)
    end = '2014-12-31'
    definition_period = YEAR

class f3sl(Variable):
    cerfa_field = u"3SL"
    value_type = int
    unit = 'currency'
    entity = FoyerFiscal
    label = u"Abattement net pour durée de détention renforcée : appliqué sur des plus-values"
    # start_date = date(2013, 1, 1)
    definition_period = YEAR

class f3sm(Variable):
    cerfa_field = u"3SM"
    value_type = int
    unit = 'currency'
    entity = FoyerFiscal
    label = u"Abattement net pour durée de détention renforcée : appliqué sur des moins-values"
    # start_date = date(2013, 1, 1)
    end = '2014-12-31'
    definition_period = YEAR

class abattement_net_retraite_dirigeant_pme(Variable):
    value_type = int
    unit = 'currency'
    entity = FoyerFiscal
    reference = u"http://bofip.impots.gouv.fr/bofip/2894-PGP"
    label = u"Abattement net pour durée de détention des titres en cas de départ à la retraite d'un dirigeant"
    definition_period = YEAR

    def formula_2006_01_01(foyer_fiscal, period):
        f3va = foyer_fiscal('f3va', period)
        f3vb = foyer_fiscal('f3va', period)

        return f3va - f3vb

class abattement_net_duree_detention(Variable):
    value_type = int
    unit = 'currency'
    entity = FoyerFiscal
    reference = u"http://bofip.impots.gouv.fr/bofip/9540-PGP"
    label = u"Abattement net pour durée de détention (de droit commun et renforcé)"
    definition_period = YEAR

    def formula_2013_01_01(foyer_fiscal, period):
        f3sg = foyer_fiscal('f3sg', period)
        f3sh = foyer_fiscal('f3sh', period)
        f3sl = foyer_fiscal('f3sl', period)
        f3sm = foyer_fiscal('f3sm', period)

        return max_(0, f3sg - f3sh) + max_(0, f3sl - f3sm)

    def formula_2015_01_01(foyer_fiscal, period):
        f3sg = foyer_fiscal('f3sg', period)
        f3sl = foyer_fiscal('f3sl', period)

        return f3sg + f3sl

# Plus values et gains taxables à des taux forfaitaires

class f3vd(Variable):
    cerfa_field = {
        0: u"3VD",
        1: u"3SD",
        }
    value_type = int
    unit = 'currency'
    entity = Individu
    label = u"Gains de levée d'options sur titres et gains d'acquisition d'actions taxables à 18 %"
    # start_date = date(2008, 1, 1)
    definition_period = YEAR

class f3vi(Variable):
    cerfa_field = {
        0: u"3VI",
        1: u"3SI",
        }
    value_type = int
    unit = 'currency'
    entity = Individu
    label = u"Gains de levée d'options sur titres et gains d'acquisition d'actions taxables à 30 %"
    definition_period = YEAR

class f3vf(Variable):
    cerfa_field = {
        0: u"3VF",
        1: u"3SF",
        }
    value_type = int
    unit = 'currency'
    entity = Individu
    label = u"Gains de levée d'options sur titres et gains d'acquisition d'actions taxables à 41 %"
    definition_period = YEAR

class f3vm(Variable):
    cerfa_field = u"3VM"
    value_type = int
    unit = 'currency'
    entity = FoyerFiscal
    label = u"Clôture du PEA avant l'expiration de la 2e année: gains taxables à 22.5 %"
    # start_date = date(2012, 1, 1)
    definition_period = YEAR

class f3vt(Variable):
    cerfa_field = u"3VT"
    value_type = int
    unit = 'currency'
    entity = FoyerFiscal
    label = u"Clôture du PEA  entre la 2e et la 5e année: gains taxables à 19 %"
    # start_date = date(2012, 1, 1)
    definition_period = YEAR

class f3vl(Variable):
    cerfa_field = u"3VL"
    value_type = int
    unit = 'currency'
    entity = FoyerFiscal
    label = u"Distributions par des sociétés de capital-risque taxables à 19 %"
    end = '2013-12-31'
    definition_period = YEAR

class f3sa(Variable):
    cerfa_field = u"3SA"
    value_type = int
    entity = FoyerFiscal
    label = u"Plus-values de cessions de titres réalisées par un entrepreneur, taxables à 19%"
    # start_date = date(2012, 1, 1)
    end = '2012-12-31'
    definition_period = YEAR

# Autres variables

class f3vv_end_2010(Variable):
    cerfa_field = u"3VV"
    value_type = int
    unit = 'currency'
    entity = FoyerFiscal
    label = u"Pertes ouvrant droit au crédit d’impôt de 19 % "
    # start_date = date(2010, 1, 1)
    end = '2010-12-31'
    definition_period = YEAR

class f3vu(Variable):
    value_type = int
    entity = FoyerFiscal
    end = '2009-12-31'
    definition_period = YEAR<commit_msg>Corrige une erreur dans la formule "abattement_net_retraite_dirigeant_pme"<commit_after># -*- coding: utf-8 -*-

from openfisca_france.model.base import *  # noqa analysis:ignore

###################################################
## PLAN :
# 1) Gains taxés comme des salaires
# 2) Autres plus-values (exonérées, au barème, ...)
# 3) Abattements sur plus-values
# 4) Plus-values taxées forfaitairement
# 5) Autres variables
###################################################

# Gains taxés comme des salaires

class f1tv(Variable):
    cerfa_field = {
        0: u"1TV",
        1: u"1UV",
        }
    value_type = int
    unit = 'currency'
    entity = Individu
    label = u"Gains de levée d'options sur titres en cas de cession ou de conversion au porteur dans le délai d'indisponibilité: entre 1 et 2 ans"
    end = '2015-12-31'
    definition_period = YEAR

class f1tw(Variable):
    cerfa_field = {
        0: u"1TW",
        1: u"1UW",
        }
    value_type = int
    unit = 'currency'
    entity = Individu
    label = u"Gains de levée d'options sur titres en cas de cession ou de conversion au porteur dans le délai d'indisponibilité: entre 2 et 3 ans"
    end = '2015-12-31'
    definition_period = YEAR

class f1tx(Variable):
    cerfa_field = {
        0: u"1TX",
        1: u"1UX",
        }
    value_type = int
    unit = 'currency'
    entity = Individu
    label = u"Gains de levée d'options sur titres en cas de cession ou de conversion au porteur dans le délai d'indisponibilité: entre 3 et 4 ans"
    end = '2016-12-31'
    definition_period = YEAR

class f3vj(Variable):
    cerfa_field = {
        0: u"3VJ",
        1: u"3VK",
        }
    value_type = int
    unit = 'currency'
    entity = Individu
    label = u"Gains imposables sur option dans la catégorie des salaires"
    definition_period = YEAR

# Autres plus-values

class f3vg(Variable):
    cerfa_field = u"3VG"
    value_type = int
    unit = 'currency'
    entity = FoyerFiscal
    label = u"Plus-value imposable sur gains de cession de valeurs mobilières, de droits sociaux et gains assimilés"
    definition_period = YEAR

class f3vh(Variable):
    cerfa_field = u"3VH"
    value_type = int
    unit = 'currency'
    entity = FoyerFiscal
    label = u"Moins-value imposable su gains de cession de valeurs mobilières, de droits sociaux et gains assimilés"
    definition_period = YEAR

class f3vv(Variable):
    cerfa_field = u"3VV"
    value_type = int
    unit = 'currency'
    entity = FoyerFiscal
    label = u"Plus-values réalisées par les non-résidents: montant du prélèvement de 45 % déjà versé"
    # start_date = date(2013, 1, 1)
    end = '2013-12-31'
    definition_period = YEAR

class f3vz(Variable):
    cerfa_field = u"3VZ"
    value_type = int
    unit = 'currency'
    entity = FoyerFiscal
    label = u"Plus-values imposables sur cessions d’immeubles ou de biens meubles"
    # start_date = date(2011, 1, 1)
    definition_period = YEAR

class f3vc(Variable):
    cerfa_field = u"3VC"
    value_type = int
    unit = 'currency'
    entity = FoyerFiscal
    label = u"Produits et plus-values exonérés provenant de structure de capital-risque"
    # start_date = date(2006, 1, 1)
    definition_period = YEAR

class f3vp(Variable):
    cerfa_field = u"3VP"
    value_type = int
    unit = 'currency'
    entity = FoyerFiscal
    label = u"Plus-values exonérées de cessions de titres de jeunes entreprises innovantes"
    # start_date = date(2007, 1, 1)
    end = '2013-12-31'
    definition_period = YEAR

class f3vy(Variable):
    cerfa_field = u"3VY"
    value_type = int
    unit = 'currency'
    entity = FoyerFiscal
    label = u"Plus-values exonérées de cessions de participations supérieures à 25% au sein du groupe familial"
    # start_date = date(2011, 1, 1)
    end = '2013-12-31'
    definition_period = YEAR

class f3ve(Variable):
    cerfa_field = u"3VE"
    value_type = int
    unit = 'currency'
    entity = FoyerFiscal
    label = u"Plus-values réalisées par les non-résidents pour lesquelles vous demandez le remboursement de l'excédent du prélèvement de 45 %"
    definition_period = YEAR

class f3sb(Variable):
    cerfa_field = u"3SB"
    value_type = int
    unit = 'currency'
    entity = FoyerFiscal
    label = u"Plus-values en report d'imposition, dont le report a expiré cette année"
    # start_date = date(2013, 1, 1)
    definition_period = YEAR

class f3wb(Variable):
    cerfa_field = u"3WB"
    value_type = int
    unit = 'currency'
    entity = FoyerFiscal
    label = u"Plus-values des individus transférant leur domicile fiscal hors de France; plus-values et créances ne bénéficiant pas du sursis de paiement; plus-values nettes imposables au barème"
    # start_date = date(2013, 1, 1)
    definition_period = YEAR

class f3wd(Variable):
    cerfa_field = u"3WD"
    value_type = int
    unit = 'currency'
    entity = FoyerFiscal
    label = u"Plus-values des individus transférant leur domicile fiscal hors de France; plus-values et créances ne bénéficiant pas du sursis de paiement; plus-values de base soumise aux prélèvements sociaux"
    # start_date = date(2013, 1, 1)
    definition_period = YEAR

# Abattements sur plus-values

class f3va(Variable):
    cerfa_field = u"3VA"
    value_type = int
    unit = 'currency'
    entity = FoyerFiscal
    label = u"Abattements nets (abattement pour durée de détention renforcé et abattement fixe spécial) appliqués sur des plus-values réalisées par les dirigeants de PME lors de leur départ à la retraite"
    # start_date = date(2006, 1, 1)
    definition_period = YEAR

class f3vb(Variable):
    cerfa_field = u"3VB"
    value_type = int
    unit = 'currency'
    entity = FoyerFiscal
    label = u"Abattements nets (abattement pour durée de détention renforcé et abattement fixe spécial) appliqués sur des moins-values réalisées par les dirigeants de PME lors de leur départ à la retraite"
    # start_date = date(2006, 1, 1)
    definition_period = YEAR

class f3sg(Variable):
    cerfa_field = u"3SG"
    value_type = int
    unit = 'currency'
    entity = FoyerFiscal
    label = u"Abattement net pour durée de détention : appliqué sur des plus-values"
    # start_date = date(2013, 1, 1)
    definition_period = YEAR

class f3sh(Variable):
    cerfa_field = u"3SH"
    value_type = int
    unit = 'currency'
    entity = FoyerFiscal
    label = u"Abattement net pour durée de détention : appliqué sur des moins-values"
    # start_date = date(2013, 1, 1)
    end = '2014-12-31'
    definition_period = YEAR

class f3sl(Variable):
    cerfa_field = u"3SL"
    value_type = int
    unit = 'currency'
    entity = FoyerFiscal
    label = u"Abattement net pour durée de détention renforcée : appliqué sur des plus-values"
    # start_date = date(2013, 1, 1)
    definition_period = YEAR

class f3sm(Variable):
    cerfa_field = u"3SM"
    value_type = int
    unit = 'currency'
    entity = FoyerFiscal
    label = u"Abattement net pour durée de détention renforcée : appliqué sur des moins-values"
    # start_date = date(2013, 1, 1)
    end = '2014-12-31'
    definition_period = YEAR

class abattement_net_retraite_dirigeant_pme(Variable):
    value_type = int
    unit = 'currency'
    entity = FoyerFiscal
    reference = u"http://bofip.impots.gouv.fr/bofip/2894-PGP"
    label = u"Abattement net pour durée de détention des titres en cas de départ à la retraite d'un dirigeant"
    definition_period = YEAR

    def formula_2006_01_01(foyer_fiscal, period):
        f3va = foyer_fiscal('f3va', period)
        f3vb = foyer_fiscal('f3vb', period)

        return f3va - f3vb

class abattement_net_duree_detention(Variable):
    value_type = int
    unit = 'currency'
    entity = FoyerFiscal
    reference = u"http://bofip.impots.gouv.fr/bofip/9540-PGP"
    label = u"Abattement net pour durée de détention (de droit commun et renforcé)"
    definition_period = YEAR

    def formula_2013_01_01(foyer_fiscal, period):
        f3sg = foyer_fiscal('f3sg', period)
        f3sh = foyer_fiscal('f3sh', period)
        f3sl = foyer_fiscal('f3sl', period)
        f3sm = foyer_fiscal('f3sm', period)

        return max_(0, f3sg - f3sh) + max_(0, f3sl - f3sm)

    def formula_2015_01_01(foyer_fiscal, period):
        f3sg = foyer_fiscal('f3sg', period)
        f3sl = foyer_fiscal('f3sl', period)

        return f3sg + f3sl

# Plus values et gains taxables à des taux forfaitaires

class f3vd(Variable):
    cerfa_field = {
        0: u"3VD",
        1: u"3SD",
        }
    value_type = int
    unit = 'currency'
    entity = Individu
    label = u"Gains de levée d'options sur titres et gains d'acquisition d'actions taxables à 18 %"
    # start_date = date(2008, 1, 1)
    definition_period = YEAR

class f3vi(Variable):
    cerfa_field = {
        0: u"3VI",
        1: u"3SI",
        }
    value_type = int
    unit = 'currency'
    entity = Individu
    label = u"Gains de levée d'options sur titres et gains d'acquisition d'actions taxables à 30 %"
    definition_period = YEAR

class f3vf(Variable):
    cerfa_field = {
        0: u"3VF",
        1: u"3SF",
        }
    value_type = int
    unit = 'currency'
    entity = Individu
    label = u"Gains de levée d'options sur titres et gains d'acquisition d'actions taxables à 41 %"
    definition_period = YEAR

class f3vm(Variable):
    cerfa_field = u"3VM"
    value_type = int
    unit = 'currency'
    entity = FoyerFiscal
    label = u"Clôture du PEA avant l'expiration de la 2e année: gains taxables à 22.5 %"
    # start_date = date(2012, 1, 1)
    definition_period = YEAR

class f3vt(Variable):
    cerfa_field = u"3VT"
    value_type = int
    unit = 'currency'
    entity = FoyerFiscal
    label = u"Clôture du PEA  entre la 2e et la 5e année: gains taxables à 19 %"
    # start_date = date(2012, 1, 1)
    definition_period = YEAR

class f3vl(Variable):
    cerfa_field = u"3VL"
    value_type = int
    unit = 'currency'
    entity = FoyerFiscal
    label = u"Distributions par des sociétés de capital-risque taxables à 19 %"
    end = '2013-12-31'
    definition_period = YEAR

class f3sa(Variable):
    cerfa_field = u"3SA"
    value_type = int
    entity = FoyerFiscal
    label = u"Plus-values de cessions de titres réalisées par un entrepreneur, taxables à 19%"
    # start_date = date(2012, 1, 1)
    end = '2012-12-31'
    definition_period = YEAR

# Autres variables

class f3vv_end_2010(Variable):
    cerfa_field = u"3VV"
    value_type = int
    unit = 'currency'
    entity = FoyerFiscal
    label = u"Pertes ouvrant droit au crédit d’impôt de 19 % "
    # start_date = date(2010, 1, 1)
    end = '2010-12-31'
    definition_period = YEAR

class f3vu(Variable):
    value_type = int
    entity = FoyerFiscal
    end = '2009-12-31'
    definition_period = YEAR<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Copyright 2010 Asidev s.r.l.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

import plac
from . interface import BaseInterface

class RedirectInterface(BaseInterface):

    commands = ['list', 'create', 'edit', 'info', 'delete']
    name = 'redirects'

    @plac.annotations(
        source=('Source domain to redirect', 'positional', None, str, None,
                'SOURCE_DOMAIN'),
        destination=('Destination instance', 'positional', None, str, None,
                     'DESTINATION_DOMAIN'),
        http_code=('HTTP code to issue', 'option', 'c'),
        target_path=('target path on the destination domain', 'option', 'p')
    )
    def create(self, source, destination, http_code=None, target_path=None):
        """ Create a new empty redirect """
        params = {'source': source,
                  'destination': destination}
        if http_code:
            params['http_code'] = http_code
        if target_path:
            params['target_path'] = target_path

        self.api.post(self.get_url(), params)

    @plac.annotations(
        source=('Source domain to redirect', 'positional', None, str, None,
                'SOURCE_DOMAIN'),
        destination=('Destination instance', 'option', 'd', str, None,
                     'DESTINATION_DOMAIN'),
        http_code=('HTTP code to issue', 'option', 'c'),
        target_path=('target path on the destination domain', 'option', 'p')
    )
    def edit(self, source, destination=None, http_code=None, target_path=None):
        params = {}
        if destination:
            params['destination'] = destination
        if http_code:
            params['http_code'] = http_code
        if not target_path is None:
            params['target_path'] = target_path

        self.api.put(self.get_url(source), params)

<commit_msg>Print additional info in redirect list<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Copyright 2010 Asidev s.r.l.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

import plac
from . interface import BaseInterface

class RedirectInterface(BaseInterface):

    commands = ['list', 'create', 'edit', 'info', 'delete']
    name = 'redirects'

    @plac.annotations(
        source=('Source domain to redirect', 'positional', None, str, None,
                'SOURCE_DOMAIN'),
        destination=('Destination instance', 'positional', None, str, None,
                     'DESTINATION_DOMAIN'),
        http_code=('HTTP code to issue', 'option', 'c'),
        target_path=('target path on the destination domain', 'option', 'p')
    )
    def create(self, source, destination, http_code=None, target_path=None):
        """ Create a new empty redirect """
        params = {'source': source,
                  'destination': destination}
        if http_code:
            params['http_code'] = http_code
        if target_path:
            params['target_path'] = target_path

        self.api.post(self.get_url(), params)

    @plac.annotations(
        source=('Source domain to redirect', 'positional', None, str, None,
                'SOURCE_DOMAIN'),
        destination=('Destination instance', 'option', 'd', str, None,
                     'DESTINATION_DOMAIN'),
        http_code=('HTTP code to issue', 'option', 'c'),
        target_path=('target path on the destination domain', 'option', 'p')
    )
    def edit(self, source, destination=None, http_code=None, target_path=None):
        params = {}
        if destination:
            params['destination'] = destination
        if http_code:
            params['http_code'] = http_code
        if not target_path is None:
            params['target_path'] = target_path

        self.api.put(self.get_url(source), params)

    @plac.annotations(
        full=('Get complete output', 'flag', 'f'),
        verbose=('Be verbose', 'flag', 'v')
    )
    def list(self, full=False, verbose=False):
        quiet = not verbose
        response, content = self.api.get(self.root_url, quiet=quiet)
        content = content or {}

        if not full:
            for source, data in content.iteritems():
                print " • {source} => {destination}{target_path} [{http_code}]"\
                        .format(**data)
        else:
            for source, data in content.iteritems():
                print " • {}:".format(source)
                self.print_info(data, prompt='   ° ')

<|endoftext|>


<commit_before>import os
import tempfile
import unittest

from py4j.protocol import Py4JJavaError
from pyspark.ml import Pipeline
from pyspark.sql import types as t

from mleap.pyspark.feature.string_map import StringMap
from mleap.pyspark.spark_support import SimpleSparkSerializer
from tests.pyspark.lib.assertions import assert_df
from tests.pyspark.lib.spark_session import spark_session

INPUT_SCHEMA = t.StructType([t.StructField('key_col', t.StringType(), False),
                             t.StructField('extra_col', t.StringType(), False)])

OUTPUT_SCHEMA = t.StructType([t.StructField('key_col', t.StringType(), False),
                              t.StructField('extra_col', t.StringType(), False),
                              t.StructField('value_col', t.DoubleType(), False)])

class StringMapTest(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        cls.spark = spark_session()

    @classmethod
    def tearDownClass(cls):
        cls.spark.stop()

    def setUp(self):
        self.input = StringMapTest.spark.createDataFrame([['a', 'b']], INPUT_SCHEMA)

    def test_map(self):
        result = StringMap(
            labels={'a': 1.0},
            inputCol='key_col',
            outputCol='value_col',
        ).transform(self.input)
        expected = StringMapTest.spark.createDataFrame([['a', 'b', 1.0]], OUTPUT_SCHEMA)
        assert_df(expected, result)

    def test_map_default_value(self):
        result = StringMap(
            labels={'z': 1.0},
            inputCol='key_col',
            outputCol='value_col',
            handleInvalid='keep',
        ).transform(self.input)
        expected = StringMapTest.spark.createDataFrame([['a', 'b', 0.0]], OUTPUT_SCHEMA)
        assert_df(expected, result)

    def test_map_custom_default_value(self):
        result = StringMap(
            labels={'z': 1.0},
            inputCol='key_col',
            outputCol='value_col',
            handleInvalid='keep',
            defaultValue=-1.0
        ).transform(self.input)
        expected = StringMapTest.spark.createDataFrame([['a', 'b', -1.0]], OUTPUT_SCHEMA)
        assert_df(expected, result)

    def test_map_missing_value_error(self):
        with self.assertRaises(Py4JJavaError) as error:
            StringMap(
                labels={'z': 1.0},
                inputCol='key_col',
                outputCol='value_col'
            ).transform(self.input).collect()
        self.assertIn('java.util.NoSuchElementException: Missing label: a', str(error.exception))

    def test_map_from_dataframe(self):
        labels_df = StringMapTest.spark.createDataFrame([['a', 1.0]], 'key_col: string, value_col: double')
        result = StringMap.from_dataframe(
            labels_df=labels_df,
            inputCol='key_col',
            outputCol='value_col'
        ).transform(self.input)
        expected = StringMapTest.spark.createDataFrame([['a', 'b', 1.0]], OUTPUT_SCHEMA)
        assert_df(expected, result)

    def test_serialize_to_bundle(self):
        print('anca')
        print(mleap.__version__)
        string_map = StringMap(
            labels={'a': 1.0},
            inputCol='key_col',
            outputCol='value_col',
        )
        pipeline = Pipeline(stages=[string_map]).fit(self.input)
        serialization_dataset = pipeline.transform(self.input)

        jar_file_path = _serialize_to_file(pipeline, serialization_dataset)
        deserialized_pipeline = _deserialize_from_file(jar_file_path)

        result = deserialized_pipeline.transform(self.input)
        expected = StringMapTest.spark.createDataFrame([['a', 'b', 1.0]], OUTPUT_SCHEMA)
        assert_df(expected, result)

    @staticmethod
    def test_validate_handleInvalid_ok():
        StringMap(labels={}, handleInvalid='error')

    def test_validate_handleInvalid_bad(self):
        with self.assertRaises(AssertionError):
            StringMap(labels=None, inputCol=dict(), outputCol=None, handleInvalid='invalid')

    def test_validate_labels_type_fails(self):
        with self.assertRaises(AssertionError):
            StringMap(labels=None, inputCol=set(), outputCol=None)

    def test_validate_labels_key_fails(self):
        with self.assertRaises(AssertionError):
            StringMap(labels=None, inputCol={False: 0.0}, outputCol=None)

    def test_validate_labels_value_fails(self):
        with self.assertRaises(AssertionError):
            StringMap(labels=None, inputCol={'valid_key_type': 'invalid_value_type'}, outputCol=None)

def _serialize_to_file(model, df_for_serializing):
    jar_file_path = _to_jar_file_path(
        os.path.join(tempfile.mkdtemp(), 'test_serialize_to_bundle-pipeline.zip'))
    SimpleSparkSerializer().serializeToBundle(model, jar_file_path, df_for_serializing)
    return jar_file_path

def _to_jar_file_path(path):
    return "jar:file:" + path

def _deserialize_from_file(path):
    return SimpleSparkSerializer().deserializeFromBundle(path)
<commit_msg>test pt2<commit_after>import os
import tempfile
import unittest
import mleap

from py4j.protocol import Py4JJavaError
from pyspark.ml import Pipeline
from pyspark.sql import types as t

from mleap.pyspark.feature.string_map import StringMap
from mleap.pyspark.spark_support import SimpleSparkSerializer
from tests.pyspark.lib.assertions import assert_df
from tests.pyspark.lib.spark_session import spark_session

INPUT_SCHEMA = t.StructType([t.StructField('key_col', t.StringType(), False),
                             t.StructField('extra_col', t.StringType(), False)])

OUTPUT_SCHEMA = t.StructType([t.StructField('key_col', t.StringType(), False),
                              t.StructField('extra_col', t.StringType(), False),
                              t.StructField('value_col', t.DoubleType(), False)])

class StringMapTest(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        cls.spark = spark_session()

    @classmethod
    def tearDownClass(cls):
        cls.spark.stop()

    def setUp(self):
        self.input = StringMapTest.spark.createDataFrame([['a', 'b']], INPUT_SCHEMA)

    def test_map(self):
        result = StringMap(
            labels={'a': 1.0},
            inputCol='key_col',
            outputCol='value_col',
        ).transform(self.input)
        expected = StringMapTest.spark.createDataFrame([['a', 'b', 1.0]], OUTPUT_SCHEMA)
        assert_df(expected, result)

    def test_map_default_value(self):
        result = StringMap(
            labels={'z': 1.0},
            inputCol='key_col',
            outputCol='value_col',
            handleInvalid='keep',
        ).transform(self.input)
        expected = StringMapTest.spark.createDataFrame([['a', 'b', 0.0]], OUTPUT_SCHEMA)
        assert_df(expected, result)

    def test_map_custom_default_value(self):
        result = StringMap(
            labels={'z': 1.0},
            inputCol='key_col',
            outputCol='value_col',
            handleInvalid='keep',
            defaultValue=-1.0
        ).transform(self.input)
        expected = StringMapTest.spark.createDataFrame([['a', 'b', -1.0]], OUTPUT_SCHEMA)
        assert_df(expected, result)

    def test_map_missing_value_error(self):
        with self.assertRaises(Py4JJavaError) as error:
            StringMap(
                labels={'z': 1.0},
                inputCol='key_col',
                outputCol='value_col'
            ).transform(self.input).collect()
        self.assertIn('java.util.NoSuchElementException: Missing label: a', str(error.exception))

    def test_map_from_dataframe(self):
        labels_df = StringMapTest.spark.createDataFrame([['a', 1.0]], 'key_col: string, value_col: double')
        result = StringMap.from_dataframe(
            labels_df=labels_df,
            inputCol='key_col',
            outputCol='value_col'
        ).transform(self.input)
        expected = StringMapTest.spark.createDataFrame([['a', 'b', 1.0]], OUTPUT_SCHEMA)
        assert_df(expected, result)

    def test_serialize_to_bundle(self):
        print('anca')
        print(mleap.__version__)
        string_map = StringMap(
            labels={'a': 1.0},
            inputCol='key_col',
            outputCol='value_col',
        )
        pipeline = Pipeline(stages=[string_map]).fit(self.input)
        serialization_dataset = pipeline.transform(self.input)

        jar_file_path = _serialize_to_file(pipeline, serialization_dataset)
        deserialized_pipeline = _deserialize_from_file(jar_file_path)

        result = deserialized_pipeline.transform(self.input)
        expected = StringMapTest.spark.createDataFrame([['a', 'b', 1.0]], OUTPUT_SCHEMA)
        assert_df(expected, result)

    @staticmethod
    def test_validate_handleInvalid_ok():
        StringMap(labels={}, handleInvalid='error')

    def test_validate_handleInvalid_bad(self):
        with self.assertRaises(AssertionError):
            StringMap(labels=None, inputCol=dict(), outputCol=None, handleInvalid='invalid')

    def test_validate_labels_type_fails(self):
        with self.assertRaises(AssertionError):
            StringMap(labels=None, inputCol=set(), outputCol=None)

    def test_validate_labels_key_fails(self):
        with self.assertRaises(AssertionError):
            StringMap(labels=None, inputCol={False: 0.0}, outputCol=None)

    def test_validate_labels_value_fails(self):
        with self.assertRaises(AssertionError):
            StringMap(labels=None, inputCol={'valid_key_type': 'invalid_value_type'}, outputCol=None)

def _serialize_to_file(model, df_for_serializing):
    jar_file_path = _to_jar_file_path(
        os.path.join(tempfile.mkdtemp(), 'test_serialize_to_bundle-pipeline.zip'))
    SimpleSparkSerializer().serializeToBundle(model, jar_file_path, df_for_serializing)
    return jar_file_path

def _to_jar_file_path(path):
    return "jar:file:" + path

def _deserialize_from_file(path):
    return SimpleSparkSerializer().deserializeFromBundle(path)
<|endoftext|>


<commit_before>import pandas as pd
from pandas.util.testing import assert_frame_equal
from nose.tools import (assert_equal, assert_not_equal, assert_items_equal,
                        assert_almost_equal, raises)
from nose.plugins.skip import Skip, SkipTest
from test_helpers import (
    true_func, assert_equal_array_array, make_1d_traj, data_filename
)

import openpathsampling as paths
from openpathsampling.high_level.interface_set import GenericVolumeInterfaceSet

import logging
logging.getLogger('openpathsampling.initialization').setLevel(logging.CRITICAL)
logging.getLogger('openpathsampling.ensemble').setLevel(logging.CRITICAL)
logging.getLogger('openpathsampling.storage').setLevel(logging.CRITICAL)
logging.getLogger('openpathsampling.netcdfplus').setLevel(logging.CRITICAL)

class TestResamplingStatistics(object):
    # NOTE: we test the mean_df and std_df functions within this
    def setup(self):
        # order of the column/index labels should not matter
        self.list_AB = ['A', 'B']
        self.list_BA = ['B', 'A']
        self.df1 = pd.DataFrame([[1.5, 2.0], [2.1, 3.0]],
                                columns=['A', 'B'], index=['A', 'B'])
        self.df2 = pd.DataFrame([[2.5, 1.0], [3.5, 2.4]],
                                columns=['B', 'A'], index=['A', 'B'])
        self.df3 = pd.DataFrame([[2.25, 3.25], [1.25, 2.25]],
                                columns=['A', 'B'], index=['B', 'A'])
        self.df4 = pd.DataFrame([[3.25, 2.25], [2.25, 1.25]],
                                columns=['B', 'A'], index=['B', 'A'])
        self.inputs = [self.df1, self.df2, self.df3, self.df4]

    def test_std(self):
        # extra test to get the one line that isn't otherwise used
        from openpathsampling.numerics.resampling_statistics import std_df
        expected_std = pd.DataFrame(
            [[0.17677669529663689, 0.17677669529663689],
             [0.10606601717798207, 0.17677669529663689]],
            columns=['A', 'B'], index=['A', 'B']
        )
        assert_frame_equal(std_df(self.inputs), expected_std)

    def test_initialization(self):
        stats = paths.numerics.ResamplingStatistics(
            function=lambda x: x,
            inputs=self.inputs
        )

        for (truth, beauty) in zip(self.inputs, stats.results):
            assert_frame_equal(beauty, truth)

        expected_mean = pd.DataFrame([[1.25, 2.25], [2.25, 3.25]],
                                     columns=['A', 'B'], index=['A', 'B'])
        assert_frame_equal(stats.mean, expected_mean)
        expected_std = pd.DataFrame(
            [[0.17677669529663689, 0.17677669529663689],
             [0.10606601717798207, 0.17677669529663689]],
            columns=['A', 'B'], index=['A', 'B']
        )
        assert_frame_equal(stats.std, expected_std)

    def test_percentile(self):
        # TODO: this would benefit from more tests (with more input frames)
        stats = paths.numerics.ResamplingStatistics(
            function=lambda x: x,
            inputs=self.inputs
        )

        assert_frame_equal(
            stats.percentile(0),
            pd.DataFrame([[1.0, 2.0], [2.1, 3.0]],
                         index=self.list_AB, columns=self.list_AB),
            check_dtype=False  # better if this wasn't needed
        )

        assert_frame_equal(
            stats.percentile(100),
            pd.DataFrame([[1.5, 2.5], [2.4, 3.5]],
                         index=self.list_AB, columns=self.list_AB),
            check_dtype=False
        )

        assert_frame_equal(
            stats.percentile(50),
            pd.DataFrame([[1.25, 2.25], [2.25, 3.25]],
                         index=self.list_AB, columns=self.list_AB),
            check_dtype=False
        )

class TestBlockResampling(object):
    def setup(self):
        self.samples = list(range(100))

    def test_default_initialization(self):
        resampler = paths.numerics.BlockResampling(self.samples)
        assert_equal(resampler.n_total_samples, 100)
        assert_equal(resampler.n_blocks, 20)
        assert_equal(resampler.n_per_block, 5)
        assert_equal(len(resampler.blocks), 20)
        assert_equal(resampler.n_resampled, 100)
        assert_equal(resampler.unassigned, [])
        assert_items_equal(resampler.blocks[0], list(range(0, 5)))
        assert_items_equal(resampler.blocks[1], list(range(5, 10)))
        assert_items_equal(resampler.blocks[10], list(range(50, 55)))
        assert_items_equal(resampler.blocks[-1], list(range(95, 100)))

    def test_n_blocks_initialization(self):
        resampler = paths.numerics.BlockResampling(self.samples, n_blocks=10)
        assert_equal(resampler.n_total_samples, 100)
        assert_equal(resampler.n_blocks, 10)
        assert_equal(resampler.n_per_block, 10)
        assert_equal(len(resampler.blocks), 10)
        assert_equal(resampler.n_resampled, 100)
        assert_equal(resampler.unassigned, [])
        assert_items_equal(resampler.blocks[0], list(range(0, 10)))
        assert_items_equal(resampler.blocks[1], list(range(10, 20)))
        assert_items_equal(resampler.blocks[5], list(range(50, 60)))
        assert_items_equal(resampler.blocks[-1], list(range(90, 100)))

        resampler_2 = paths.numerics.BlockResampling(self.samples, n_blocks=8)
        assert_equal(resampler_2.n_total_samples, 100)
        assert_equal(resampler_2.n_blocks, 8)
        assert_equal(resampler_2.n_per_block, 12)
        assert_equal(resampler_2.n_resampled, 96)
        assert_items_equal(resampler_2.blocks[-1], list(range(84, 96)))
        assert_items_equal(resampler_2.unassigned, list(range(96, 100)))

    def test_n_per_block_initialization(self):
        resampler = paths.numerics.BlockResampling(self.samples, n_per_block=10)
        assert_equal(resampler.n_total_samples, 100)
        assert_equal(resampler.n_blocks, 10)
        assert_equal(resampler.n_per_block, 10)
        assert_equal(len(resampler.blocks), 10)
        assert_equal(resampler.n_resampled, 100)
        assert_equal(resampler.unassigned, [])
        assert_items_equal(resampler.blocks[0], list(range(0, 10)))
        assert_items_equal(resampler.blocks[1], list(range(10, 20)))
        assert_items_equal(resampler.blocks[5], list(range(50, 60)))
        assert_items_equal(resampler.blocks[-1], list(range(90, 100)))

        resampler_2 = paths.numerics.BlockResampling(self.samples,
                                                     n_per_block=8)
        assert_equal(resampler_2.n_total_samples, 100)
        assert_equal(resampler_2.n_blocks, 12)
        assert_equal(resampler_2.n_resampled, 96)
        assert_items_equal(resampler_2.blocks[-1], list(range(88, 96)))
        assert_items_equal(resampler_2.unassigned, list(range(96, 100)))

    def test_n_blocks_and_n_per_block_initialization(self):
        resampler = paths.numerics.BlockResampling(self.samples, n_blocks=9,
                                                   n_per_block=10)
        assert_equal(resampler.n_total_samples, 100)
        assert_equal(resampler.n_blocks, 9)
        assert_equal(resampler.n_per_block, 10)
        assert_equal(len(resampler.blocks), 9)
        assert_equal(resampler.n_resampled, 90)
        assert_items_equal(resampler.unassigned, list(range(90, 100)))
        assert_items_equal(resampler.blocks[0], list(range(0, 10)))
        assert_items_equal(resampler.blocks[1], list(range(10, 20)))
        assert_items_equal(resampler.blocks[5], list(range(50, 60)))
        assert_items_equal(resampler.blocks[-1], list(range(80, 90)))
<commit_msg>Move to test_helpers.assert_items_equal<commit_after>import pandas as pd
from pandas.util.testing import assert_frame_equal
from nose.tools import (assert_equal, assert_not_equal,
                        assert_almost_equal, raises)
from nose.plugins.skip import Skip, SkipTest
from test_helpers import (
    true_func, assert_equal_array_array, make_1d_traj, data_filename,
    assert_items_equal
)

import openpathsampling as paths
from openpathsampling.high_level.interface_set import GenericVolumeInterfaceSet

import logging
logging.getLogger('openpathsampling.initialization').setLevel(logging.CRITICAL)
logging.getLogger('openpathsampling.ensemble').setLevel(logging.CRITICAL)
logging.getLogger('openpathsampling.storage').setLevel(logging.CRITICAL)
logging.getLogger('openpathsampling.netcdfplus').setLevel(logging.CRITICAL)

class TestResamplingStatistics(object):
    # NOTE: we test the mean_df and std_df functions within this
    def setup(self):
        # order of the column/index labels should not matter
        self.list_AB = ['A', 'B']
        self.list_BA = ['B', 'A']
        self.df1 = pd.DataFrame([[1.5, 2.0], [2.1, 3.0]],
                                columns=['A', 'B'], index=['A', 'B'])
        self.df2 = pd.DataFrame([[2.5, 1.0], [3.5, 2.4]],
                                columns=['B', 'A'], index=['A', 'B'])
        self.df3 = pd.DataFrame([[2.25, 3.25], [1.25, 2.25]],
                                columns=['A', 'B'], index=['B', 'A'])
        self.df4 = pd.DataFrame([[3.25, 2.25], [2.25, 1.25]],
                                columns=['B', 'A'], index=['B', 'A'])
        self.inputs = [self.df1, self.df2, self.df3, self.df4]

    def test_std(self):
        # extra test to get the one line that isn't otherwise used
        from openpathsampling.numerics.resampling_statistics import std_df
        expected_std = pd.DataFrame(
            [[0.17677669529663689, 0.17677669529663689],
             [0.10606601717798207, 0.17677669529663689]],
            columns=['A', 'B'], index=['A', 'B']
        )
        assert_frame_equal(std_df(self.inputs), expected_std)

    def test_initialization(self):
        stats = paths.numerics.ResamplingStatistics(
            function=lambda x: x,
            inputs=self.inputs
        )

        for (truth, beauty) in zip(self.inputs, stats.results):
            assert_frame_equal(beauty, truth)

        expected_mean = pd.DataFrame([[1.25, 2.25], [2.25, 3.25]],
                                     columns=['A', 'B'], index=['A', 'B'])
        assert_frame_equal(stats.mean, expected_mean)
        expected_std = pd.DataFrame(
            [[0.17677669529663689, 0.17677669529663689],
             [0.10606601717798207, 0.17677669529663689]],
            columns=['A', 'B'], index=['A', 'B']
        )
        assert_frame_equal(stats.std, expected_std)

    def test_percentile(self):
        # TODO: this would benefit from more tests (with more input frames)
        stats = paths.numerics.ResamplingStatistics(
            function=lambda x: x,
            inputs=self.inputs
        )

        assert_frame_equal(
            stats.percentile(0),
            pd.DataFrame([[1.0, 2.0], [2.1, 3.0]],
                         index=self.list_AB, columns=self.list_AB),
            check_dtype=False  # better if this wasn't needed
        )

        assert_frame_equal(
            stats.percentile(100),
            pd.DataFrame([[1.5, 2.5], [2.4, 3.5]],
                         index=self.list_AB, columns=self.list_AB),
            check_dtype=False
        )

        assert_frame_equal(
            stats.percentile(50),
            pd.DataFrame([[1.25, 2.25], [2.25, 3.25]],
                         index=self.list_AB, columns=self.list_AB),
            check_dtype=False
        )

class TestBlockResampling(object):
    def setup(self):
        self.samples = list(range(100))

    def test_default_initialization(self):
        resampler = paths.numerics.BlockResampling(self.samples)
        assert_equal(resampler.n_total_samples, 100)
        assert_equal(resampler.n_blocks, 20)
        assert_equal(resampler.n_per_block, 5)
        assert_equal(len(resampler.blocks), 20)
        assert_equal(resampler.n_resampled, 100)
        assert_equal(resampler.unassigned, [])
        assert_items_equal(resampler.blocks[0], list(range(0, 5)))
        assert_items_equal(resampler.blocks[1], list(range(5, 10)))
        assert_items_equal(resampler.blocks[10], list(range(50, 55)))
        assert_items_equal(resampler.blocks[-1], list(range(95, 100)))

    def test_n_blocks_initialization(self):
        resampler = paths.numerics.BlockResampling(self.samples, n_blocks=10)
        assert_equal(resampler.n_total_samples, 100)
        assert_equal(resampler.n_blocks, 10)
        assert_equal(resampler.n_per_block, 10)
        assert_equal(len(resampler.blocks), 10)
        assert_equal(resampler.n_resampled, 100)
        assert_equal(resampler.unassigned, [])
        assert_items_equal(resampler.blocks[0], list(range(0, 10)))
        assert_items_equal(resampler.blocks[1], list(range(10, 20)))
        assert_items_equal(resampler.blocks[5], list(range(50, 60)))
        assert_items_equal(resampler.blocks[-1], list(range(90, 100)))

        resampler_2 = paths.numerics.BlockResampling(self.samples, n_blocks=8)
        assert_equal(resampler_2.n_total_samples, 100)
        assert_equal(resampler_2.n_blocks, 8)
        assert_equal(resampler_2.n_per_block, 12)
        assert_equal(resampler_2.n_resampled, 96)
        assert_items_equal(resampler_2.blocks[-1], list(range(84, 96)))
        assert_items_equal(resampler_2.unassigned, list(range(96, 100)))

    def test_n_per_block_initialization(self):
        resampler = paths.numerics.BlockResampling(self.samples, n_per_block=10)
        assert_equal(resampler.n_total_samples, 100)
        assert_equal(resampler.n_blocks, 10)
        assert_equal(resampler.n_per_block, 10)
        assert_equal(len(resampler.blocks), 10)
        assert_equal(resampler.n_resampled, 100)
        assert_equal(resampler.unassigned, [])
        assert_items_equal(resampler.blocks[0], list(range(0, 10)))
        assert_items_equal(resampler.blocks[1], list(range(10, 20)))
        assert_items_equal(resampler.blocks[5], list(range(50, 60)))
        assert_items_equal(resampler.blocks[-1], list(range(90, 100)))

        resampler_2 = paths.numerics.BlockResampling(self.samples,
                                                     n_per_block=8)
        assert_equal(resampler_2.n_total_samples, 100)
        assert_equal(resampler_2.n_blocks, 12)
        assert_equal(resampler_2.n_resampled, 96)
        assert_items_equal(resampler_2.blocks[-1], list(range(88, 96)))
        assert_items_equal(resampler_2.unassigned, list(range(96, 100)))

    def test_n_blocks_and_n_per_block_initialization(self):
        resampler = paths.numerics.BlockResampling(self.samples, n_blocks=9,
                                                   n_per_block=10)
        assert_equal(resampler.n_total_samples, 100)
        assert_equal(resampler.n_blocks, 9)
        assert_equal(resampler.n_per_block, 10)
        assert_equal(len(resampler.blocks), 9)
        assert_equal(resampler.n_resampled, 90)
        assert_items_equal(resampler.unassigned, list(range(90, 100)))
        assert_items_equal(resampler.blocks[0], list(range(0, 10)))
        assert_items_equal(resampler.blocks[1], list(range(10, 20)))
        assert_items_equal(resampler.blocks[5], list(range(50, 60)))
        assert_items_equal(resampler.blocks[-1], list(range(80, 90)))
<|endoftext|>


<commit_before>"""
Microbenchmark on the "copy builtin to constant".

The benchmark doesn't use fatoptimize, but specialize explicitly the function.
"""
import fat
import sys
from fatoptimizer.benchmark import bench, format_dt, compared_dt

def func(obj):
    return len(obj)

if fat.get_specialized(func):
    print("ERROR: func() was already specialized")
    sys.exit(1)

def func_cst(obj):
    return 'LEN'(obj)
func_cst.__code__ = fat.replace_consts(func_cst.__code__, {'LEN': len})

dt = bench('func("abc")')
print("original bytecode (LOAD_GLOBAL dict): %s" % format_dt(dt))

dt2 = bench('func_cst("abc")')
print("LOAD_CONST: %s" % compared_dt(dt2, dt))

fat.specialize(func, func_cst, [fat.GuardBuiltins(('dict',))])
assert fat.get_specialized(func)
dt3 = bench('func("abc")')
print("specialized (LOAD_CONST): %s" % compared_dt(dt3, dt))

print("cost of GuardBuiltins: %s" % format_dt(dt3 - dt2, sign=True))
<commit_msg>rewrite bench_copy_cst with my benchmark library<commit_after>"""
Microbenchmark on the "copy builtin to constant".

The benchmark doesn't use fatoptimize, but specialize explicitly the function.
"""
import fat
import sys

mylen = len

def func(obj):
    return mylen(obj)

if fat.get_specialized(func):
    print("ERROR: func() was already specialized")
    sys.exit(1)

def func_cst(obj):
    return 'MYLEN'(obj)
func_cst.__code__ = fat.replace_consts(func_cst.__code__, {'MYLEN': mylen})

def run_benchmark(bench):
    bench.timeit(stmt='func("abc")',
                 globals=globals(),
                 name='original bytecode (LOAD_GLOBAL)')

    bench.timeit(stmt='func_cst("abc")',
                 globals=globals(),
                 name='LOAD_CONST')

    fat.specialize(func, func_cst, [fat.GuardGlobals(('mylen',))])
    assert fat.get_specialized(func)

    bench.timeit(stmt='func("abc")',
                 globals=globals(),
                 name='LOAD_CONST with guard')

<|endoftext|>