<commit_before>from __future__ import (absolute_import, division, print_function)
__metaclass__ = type

DOCUMENTATION = """
"""

EXAMPLES = """
"""

RETURN = """
"""

from ansible.errors import AnsibleError
from ansible.plugins.lookup import LookupBase

try:
    from __main__ import display
except ImportError:
    from ansible.utils.display import Display
    display = Display()

class LookupModule(LookupBase):
    def run(self, terms, variables=None, **kwargs):
        try:
            assert 'group_names' in variables, "Missing 'group_names' in variables"
            assert len(terms) == 1, "Inheritance chain lookup plugin expects 1 term, got %s" % len(terms)
        except AssertionError as exc:
            raise AnsibleError(str(exc))
        _vars = variables or {}
        base = terms[0]
        r = []
        for prefix in ['all', 'group'] + [n + '_group' for n in _vars['group_names']] + ['host']:
            _var = prefix + '_' + base
            _val = _vars.get(_var, [])
            if _val:
                display.vvvv('%s = %s' % (_var, _val))
            try:
                t = self._templar.template(
                    _vars.get(_var, []), preserve_trailing_newlines=True,
                    convert_data=True, escape_backslashes=False)
            except Exception as exc:
                raise AnsibleError("Templating '%s' failed: %s" % (_var, str(exc)))
            if t != _val:
                display.vvvv('%s -> %s' % (_var, t))
            r.extend(t)
        if not r:
            raise AnsibleError("No inheritance chain variables exist for base '%s'" % base)
        return r
<commit_msg>Don't fail when no inheritance chain vars are found<commit_after>from __future__ import (absolute_import, division, print_function)
__metaclass__ = type

DOCUMENTATION = """
"""

EXAMPLES = """
"""

RETURN = """
"""

from ansible.errors import AnsibleError
from ansible.plugins.lookup import LookupBase

try:
    from __main__ import display
except ImportError:
    from ansible.utils.display import Display
    display = Display()

class LookupModule(LookupBase):
    def run(self, terms, variables=None, **kwargs):
        try:
            assert 'group_names' in variables, "Missing 'group_names' in variables"
            assert len(terms) == 1, "Inheritance chain lookup plugin expects 1 term, got %s" % len(terms)
        except AssertionError as exc:
            raise AnsibleError(str(exc))
        _vars = variables or {}
        base = terms[0]
        r = []
        for prefix in ['all', 'group'] + [n + '_group' for n in _vars['group_names']] + ['host']:
            _var = prefix + '_' + base
            _val = _vars.get(_var, [])
            if _val:
                display.vvvv('%s = %s' % (_var, _val))
            try:
                t = self._templar.template(
                    _vars.get(_var, []), preserve_trailing_newlines=True,
                    convert_data=True, escape_backslashes=False)
            except Exception as exc:
                raise AnsibleError("Templating '%s' failed: %s" % (_var, str(exc)))
            if t != _val:
                display.vvvv('%s -> %s' % (_var, t))
            r.extend(t)
        # this is actually not invalid if the name is valid, and there's no way
        # to default in the playbook, so we'll just return an empty list here
        #if not r:
        #    raise AnsibleError("No inheritance chain variables exist for base '%s'" % base)
        return r
<|endoftext|>


<commit_before>from collections import Counter

import unicodecsv

from django.core.management.base import BaseCommand
from django.utils import translation
from django.conf import settings

class Command(BaseCommand):
    help = "Calculate stats"

    def handle(self, filename, output, *args, **options):
        from froide.publicbody.models import PublicBody
        from froide.foirequest.models import FoiRequest
        translation.activate(settings.LANGUAGE_CODE)

        def get_status(status):
            KNOWN_STATUS = (
                'successful',
                'partially_successful',
                'refused',
            )
            MAP_STATUS = {}
            if status in KNOWN_STATUS:
                return status
            return MAP_STATUS.get(status, 'other')

        def convert_value(val):
            if not val:
                return 0
            else:
                return int(val)

        def stats_for_queryset(qs, key=None):
            status_counter = Counter()
            for r in qs:
                arg = key
                if arg is None:
                    arg = r.public_body.name
                status_counter[get_status(r.status)] += 1
            return status_counter

        reader = unicodecsv.DictReader(file(filename),
                encoding='utf-8')
        writer = unicodecsv.DictWriter(file(output, 'w'),
            ('id', 'name', 'type', 'official_count', 'fds_count'), encoding='utf-8')
        writer.writeheader()
        for row in reader:
            name = row['name']
            if 'Gesamt' in name:
                continue
            is_gb = False
            if ' GB' in name:
                is_gb = True
                name = name.replace(' GB', '')
            try:
                pb = PublicBody.objects.get(
                    other_names__startswith=name + ',')
            except PublicBody.DoesNotExist:
                continue
            if is_gb:
                pbs = PublicBody.objects.filter(root=pb)
                reqs = FoiRequest.objects.filter(
                    first_message__year=2012,
                    visibility__gt=0,
                    public_body__in=pbs, is_foi=True)
                stats = stats_for_queryset(reqs, u"GB: %s" % pb.name)
                name = '%s (GB)' % pb.name
                count = len(reqs)
            else:
                reqs = FoiRequest.objects.filter(first_message__year=2012,
                    visibility__gt=0,
                    public_body=pb, is_foi=True)
                stats = stats_for_queryset(reqs)
                name = pb.name
                count = len(reqs)
            writer.writerow({
                'id': row['name'],
                'name': name,
                'type': 'total',
                'official_count': convert_value(row['antraege']),
                'fds_count': count
            })
            writer.writerow({
                'id': row['name'],
                'name': name,
                'type': 'gewaehrt',
                'official_count': convert_value(row['gewaehrt']),
                'fds_count': stats['successful']
            })
            writer.writerow({
                'id': row['name'],
                'name': name,
                'type': 'teilweise_gewaehrt',
                'official_count': convert_value(row['teilweise_gewaehrt']),
                'fds_count': stats['partially_successful']
            })
            writer.writerow({
                'id': row['name'],
                'name': name,
                'type': 'abgelehnt',
                'official_count': convert_value(row['abgelehnt']),
                'fds_count': stats['refused']
            })
            writer.writerow({
                'id': row['name'],
                'name': name,
                'type': 'sonstige',
                'official_count': convert_value(row['sonstige']),
                'fds_count': stats['other']
            })
<commit_msg>Make calculate stats call less year dependent<commit_after>from collections import Counter

import unicodecsv

from django.core.management.base import BaseCommand
from django.utils import translation
from django.conf import settings

class Command(BaseCommand):
    help = "Calculate stats"

    def handle(self, filename, output=None, *args, **options):
        from froide.publicbody.models import PublicBody
        from froide.foirequest.models import FoiRequest
        translation.activate(settings.LANGUAGE_CODE)

        year = 2013

        def get_status(status):
            KNOWN_STATUS = (
                'successful',
                'partially_successful',
                'refused',
            )
            MAP_STATUS = {}
            if status in KNOWN_STATUS:
                return status
            return MAP_STATUS.get(status, 'other')

        def convert_value(val):
            if not val:
                return 0
            else:
                return int(val)

        def stats_for_queryset(qs, key=None):
            status_counter = Counter()
            for r in qs:
                arg = key
                if arg is None:
                    arg = r.public_body.name
                status_counter[get_status(r.status)] += 1
            return status_counter

        reader = unicodecsv.DictReader(file(filename), encoding='utf-8')
        if output == '-':
            output = self.stdout
        else:
            output = file(output, 'w')
        writer = unicodecsv.DictWriter(output,
            ('id', 'name', 'type', 'official_count', 'fds_count'), encoding='utf-8')
        writer.writeheader()
        for row in reader:
            name = row['name']
            if 'Gesamt' in name:
                continue
            is_gb = False
            if ' GB' in name:
                is_gb = True
                name = name.replace(' GB', '')
            try:
                pb = PublicBody.objects.get(
                    other_names__startswith=name + ',')
            except PublicBody.DoesNotExist:
                continue
            if is_gb:
                pbs = PublicBody.objects.filter(root=pb)
                reqs = FoiRequest.objects.filter(
                    first_message__year=year,
                    visibility__gt=0,
                    public_body__in=pbs, is_foi=True)
                stats = stats_for_queryset(reqs, u"GB: %s" % pb.name)
                name = '%s (GB)' % pb.name
                count = len(reqs)
            else:
                reqs = FoiRequest.objects.filter(first_message__year=year,
                    visibility__gt=0,
                    public_body=pb, is_foi=True)
                stats = stats_for_queryset(reqs)
                name = pb.name
                count = len(reqs)
            writer.writerow({
                'id': row['name'],
                'name': name,
                'type': 'total',
                'official_count': convert_value(row['antraege']),
                'fds_count': count
            })
            writer.writerow({
                'id': row['name'],
                'name': name,
                'type': 'gewaehrt',
                'official_count': convert_value(row['gewaehrt']),
                'fds_count': stats['successful']
            })
            writer.writerow({
                'id': row['name'],
                'name': name,
                'type': 'teilweise_gewaehrt',
                'official_count': convert_value(row['teilweise_gewaehrt']),
                'fds_count': stats['partially_successful']
            })
            writer.writerow({
                'id': row['name'],
                'name': name,
                'type': 'abgelehnt',
                'official_count': convert_value(row['abgelehnt']),
                'fds_count': stats['refused']
            })
            writer.writerow({
                'id': row['name'],
                'name': name,
                'type': 'sonstige',
                'official_count': convert_value(row['sonstige']),
                'fds_count': stats['other']
            })
<|endoftext|>


<commit_before>#!/usr/bin/python
"""Utilities to facilitate out-of-core learning in sklearn"""

import numpy as np
import sklearn
import random

def fit(estimator, data_sequence, batch_size=100, max_steps=None):
    '''Fit a model to a generator stream.

    :parameters:
      - estimator : sklearn.base.BaseEstimator
        The model object.  Must implement ``partial_fit()``

      - data_sequence : generator
        A generator that yields samples

      - batch_size : int
        Maximum number of samples to buffer before updating the model

      - max_steps : int or None
        If ``None``, run until the stream is exhausted.
        Otherwise, run until at most ``max_steps`` examples have been processed.
    '''

    # Is this a supervised or unsupervised learner?
    supervised = isinstance(estimator, sklearn.base.ClassifierMixin)

    # Does the learner support partial fit?
    assert(hasattr(estimator, 'partial_fit'))

    def _run(data, supervised):
        """Wrapper function to partial_fit()"""

        data_x = []
        data_y = []
        if supervised:
            for sample in data:
                data_x.append(sample[0])
                data_y.append(sample[-1])

            # Support sparsity here
            data_x = np.asarray(data_x)
            data_y = np.asarray(data_y)

            estimator.partial_fit(data_x, data_y)
        else:
            # Support sparsity here
            estimator.partial_fit(np.asarray(data))

    buf = []
    for i, x_new in enumerate(data_sequence):
        buf.append(x_new)

        # We've run too far, stop
        if max_steps is None and max_steps <= i:
            break

        # Buffer is full, do an update
        if len(buf) == batch_size:
            _run(buf, supervised)
            buf = []

    # Update on whatever's left over
    if len(buf) > 0:
        _run(buf, supervised)

def mux(generators):
    '''Randomly multiplex over a list of generators until they are all exhausted.

    :parameters:
      - generators : list
        List of generators to randomly multiplex

    :yields:
      - iterates from each of the input generators
    '''

    # Loop until the generators list is empty
    while len(generators) > 0:
        # Pick a random generator
        i = random.randint(0, len(generators) - 1)

        try:
            # Pull from the i'th generator
            yield generators[i].next()

        except StopIteration:
            # This one's done.  Remove it from the list.
            generators.pop(i)

def mux_bounded(sample_generator, containers, working_size=10, max_iter=1, shuffle=True, **kwargs):
    '''Generate a sequence by multiplexing a generator function over subsets of data.

    1. Select a subset of ``working_size`` from containers
    2. Apply ``sample_generator`` to each item ``c`` in the subset: ``sample_generator(c, **kwargs)``
    3. Multiplex the results of each generator

    :parameters:
      - sample_generator : function
        A generator which takes as input a container (see below), and yields samples
      - containers: list-like
        A list of container objects
      - working_size : int > 0
        The maximum number of working containers to keep active at any time
      - max_iter : int > 0
        How many passes through the container set to take
      - shuffle : boolean
        Permute the containers at each iterations
      - kwargs : additional keyword arguments
        Parameters to pass throuh to sample_generator

    :yields:
      - sequence of items generated by multiplexing ``sample_generator()``
    '''

    idx = range(len(containers))

    for _ in range(max_iter):

        if shuffle:
            random.shuffle(idx)

        for i in range(0, len(idx), working_size):
            generators = [sample_generator(ctr, **kwargs) for ctr in containers[idx[i:i+working_size]]]

            for x in mux(generators):
                yield x

<commit_msg>fixed some slicing errors<commit_after>#!/usr/bin/python
"""Utilities to facilitate out-of-core learning in sklearn"""

import numpy as np
import sklearn
import random

def fit(estimator, data_sequence, batch_size=100, max_steps=None):
    '''Fit a model to a generator stream.

    :parameters:
      - estimator : sklearn.base.BaseEstimator
        The model object.  Must implement ``partial_fit()``

      - data_sequence : generator
        A generator that yields samples

      - batch_size : int
        Maximum number of samples to buffer before updating the model

      - max_steps : int or None
        If ``None``, run until the stream is exhausted.
        Otherwise, run until at most ``max_steps`` examples have been processed.
    '''

    # Is this a supervised or unsupervised learner?
    supervised = isinstance(estimator, sklearn.base.ClassifierMixin)

    # Does the learner support partial fit?
    assert(hasattr(estimator, 'partial_fit'))

    def _run(data, supervised):
        """Wrapper function to partial_fit()"""

        data_x = []
        data_y = []
        if supervised:
            for sample in data:
                data_x.append(sample[0])
                data_y.append(sample[-1])

            # Support sparsity here
            data_x = np.asarray(data_x)
            data_y = np.asarray(data_y)

            estimator.partial_fit(data_x, data_y)
        else:
            # Support sparsity here
            estimator.partial_fit(np.asarray(data))

    buf = []
    for i, x_new in enumerate(data_sequence):
        buf.append(x_new)

        # We've run too far, stop
        if max_steps is None and max_steps <= i:
            break

        # Buffer is full, do an update
        if len(buf) == batch_size:
            _run(buf, supervised)
            buf = []

    # Update on whatever's left over
    if len(buf) > 0:
        _run(buf, supervised)

def mux(generators):
    '''Randomly multiplex over a list of generators until they are all exhausted.

    :parameters:
      - generators : list
        List of generators to randomly multiplex

    :yields:
      - iterates from each of the input generators
    '''

    # Loop until the generators list is empty
    while len(generators) > 0:
        # Pick a random generator
        i = random.randint(0, len(generators) - 1)

        try:
            # Pull from the i'th generator
            yield generators[i].next()

        except StopIteration:
            # This one's done.  Remove it from the list.
            generators.pop(i)

def mux_bounded(sample_generator, containers, working_size=10, max_iter=1, shuffle=True, **kwargs):
    '''Generate a sequence by multiplexing a generator function over subsets of data.

    1. Select a subset of ``working_size`` from containers
    2. Apply ``sample_generator`` to each item ``c`` in the subset: ``sample_generator(c, **kwargs)``
    3. Multiplex the results of each generator

    :parameters:
      - sample_generator : function
        A generator which takes as input a container (see below), and yields samples
      - containers: list-like
        A list of container objects
      - working_size : int > 0
        The maximum number of working containers to keep active at any time
      - max_iter : int > 0
        How many passes through the container set to take
      - shuffle : boolean
        Permute the containers at each iterations
      - kwargs : additional keyword arguments
        Parameters to pass throuh to sample_generator

    :yields:
      - sequence of items generated by multiplexing ``sample_generator()``
    '''

    idx = range(len(containers))

    for _ in range(max_iter):

        if shuffle:
            random.shuffle(idx)

        for i in range(0, len(idx), working_size):
            generators = [sample_generator(containers[j], **kwargs) for j in idx[i:i+working_size]]

            for x in mux(generators):
                yield x

<|endoftext|>


<commit_before>import numpy as np
import sys
import astropy.table
from astropy import units as u
from astropy.io import fits
from astropy import wcs
import pdb
import pylab as plt

''' Uses astropy tables to translate raw DOLPHOT photometric catalogs into summary .fits files '''

def syntax():
	print("make_photfits.py <raw dolphot photometry file>, <output filename>,  'filter1 fitler2, ...', <reference .fits filename>")

def make_phot(datafile, refname, filters, outname, full=False, summary=True, gst=True):

	# no functionality for full photometry output yet
	# currently assumes only 2 filters present

	# relevant output labels
	global_labels = ['Number', 'RA', 'DEC', 'X', 'Y', 'OBJECT_TYPE']
	filter_labels = ['_VEGA', '_ERR', '_SNR', '_SHARP', '_ROUND', '_CROWD', '_FLAG']
	formats = ['f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8']

	filters = filters.split()
	nfilters = len(filters)

	# load raw DOLPHOT photometry
	print('Loading raw DOLPHOT file...')
	data = np.loadtxt(datafile)
	print('Loaded {0} objects'.format(len(data[:,0])))

	# number of objects
	num = np.arange(len(data[:,0])) + 1

	# Load the FITS hdulist using astropy.io.fits
	hdulist = fits.open(refname)

	# Parse the WCS keywords in the primary HDU
	w = wcs.WCS(hdulist[1].header)

	# convert x and y to RA and DEC using WCS from reference image
	world = w.wcs_pix2world(data[:,2], data[:,3], 1)

	if summary == True:
		t = astropy.table.Table()
		t.add_column(astropy.table.Column(name=global_labels[0], data=num)) # star number
		t.add_column(astropy.table.Column(name=global_labels[1], data=world[0])) # RA
		t.add_column(astropy.table.Column(name=global_labels[2], data=world[1])) # DEC
		t.add_column(astropy.table.Column(name=global_labels[3], data=data[:,2])) # X
		t.add_column(astropy.table.Column(name=global_labels[4], data=data[:,3])) # Y
		t.add_column(astropy.table.Column(name=global_labels[5], data=data[:,10])) # ObjType

		# other relevant columns: Instrumental Mag, Err, SNR, Sharp, Round, Crowd, Flag
		cols = np.int_(np.asarray((15, 17, 19, 20, 21, 22, 23)))

		# loops over number of filters, names of filters to generate output columns
		for i in range(nfilters):
			for j, k in enumerate(filter_labels):
				t.add_column(astropy.table.Column(name=filters[i]+k,
					data=data[:,cols[j] + (i*13) ]))

		t.write(sys.argv[1]+'.summary.fits', overwrite=True)

		if gst == True:

			# cull photometry on both bands
			# assumes only 2 bands for now
			snr = 5.
			sharp = 0.1
			crowd = 10.
			objtype = 1

			wgood = np.where( (t[filters[0]+'_SNR'] >= snr) & (t[filters[1]+'_SNR'] >= snr) &
						( t[filters[0]+'_SHARP']**2 < sharp ) +  ( t[filters[1]+'_SHARP']**2 < sharp ) &
						( (t[filters[0]+'_CROWD'] + t[filters[1]+'_CROWD']) < crowd ) &
						(t['OBJECT_TYPE'] <= objtype) )

			t1 = t[wgood]
			t1.write(sys.argv[1]+'.gst.fits', overwrite=True)

			'''
			plt.figure()
			plt.scatter(t[filters[0]+'_VEGA'] - t[filters[1]+'_VEGA'],t[filters[1]+'_VEGA'], color='k', edgecolor='None', s=0.5)
			plt.scatter(t1[filters[0]+'_VEGA'] - t1[filters[1]+'_VEGA'],t1[filters[1]+'_VEGA'], color='r', edgecolor='None', s=0.25)
			plt.xlim(-1,4)
			plt.ylim(30,16)
			'''

def main():
	datafile = sys.argv[1]
	outfile = sys.argv[2]
	filters = sys.argv[3]
	ref_image = sys.argv[4]

	make_phot(datafile, ref_image, filters, outfile)

main()

'''
t = astropy.table.Table()
t.add_column(astropy.table.Column(name=global_labels[0], data=data[:,0]))
t.add_column(astropy.table.Column(name=global_labels[1], data=data[:,1]))
t.add_column(astropy.table.Column(name=global_labels[2], data=data[:,2]))
k=3
for i in range(len(filters)):
    for j in range(len(filter_labels)):
        t.add_column(astropy.table.Column(name=filters[i]+filter_labels[j], data=data[:,k]))
        print data[:,k][-1]
        k+=1

t.write(sys.argv[1]+'.fits')
'''

<commit_msg>removed call to matplotlib<commit_after>import numpy as np
import sys
import astropy.table
from astropy import units as u
from astropy.io import fits
from astropy import wcs
import pdb
#import pylab as plt

''' Uses astropy tables to translate raw DOLPHOT photometric catalogs into summary .fits files '''

def syntax():
	print("make_photfits.py <raw dolphot photometry file>, <output filename>,  'filter1 fitler2, ...', <reference .fits filename>")

def make_phot(datafile, refname, filters, outname, full=False, summary=True, gst=True):

	# no functionality for full photometry output yet
	# currently assumes only 2 filters present

	# relevant output labels
	global_labels = ['Number', 'RA', 'DEC', 'X', 'Y', 'OBJECT_TYPE']
	filter_labels = ['_VEGA', '_ERR', '_SNR', '_SHARP', '_ROUND', '_CROWD', '_FLAG']
	formats = ['f8', 'f8', 'f8', 'f8', 'f8', 'f8', 'f8']

	filters = filters.split()
	nfilters = len(filters)

	# load raw DOLPHOT photometry
	print('Loading raw DOLPHOT file...')
	data = np.loadtxt(datafile)
	print('Loaded {0} objects'.format(len(data[:,0])))

	# number of objects
	num = np.arange(len(data[:,0])) + 1

	# Load the FITS hdulist using astropy.io.fits
	hdulist = fits.open(refname)

	# Parse the WCS keywords in the primary HDU
	w = wcs.WCS(hdulist[1].header)

	# convert x and y to RA and DEC using WCS from reference image
	world = w.wcs_pix2world(data[:,2], data[:,3], 1)

	if summary == True:
		t = astropy.table.Table()
		t.add_column(astropy.table.Column(name=global_labels[0], data=num)) # star number
		t.add_column(astropy.table.Column(name=global_labels[1], data=world[0])) # RA
		t.add_column(astropy.table.Column(name=global_labels[2], data=world[1])) # DEC
		t.add_column(astropy.table.Column(name=global_labels[3], data=data[:,2])) # X
		t.add_column(astropy.table.Column(name=global_labels[4], data=data[:,3])) # Y
		t.add_column(astropy.table.Column(name=global_labels[5], data=data[:,10])) # ObjType

		# other relevant columns: Instrumental Mag, Err, SNR, Sharp, Round, Crowd, Flag
		cols = np.int_(np.asarray((15, 17, 19, 20, 21, 22, 23)))

		# loops over number of filters, names of filters to generate output columns
		for i in range(nfilters):
			for j, k in enumerate(filter_labels):
				t.add_column(astropy.table.Column(name=filters[i]+k,
					data=data[:,cols[j] + (i*13) ]))

		t.write(sys.argv[1]+'.summary.fits', overwrite=True)

		if gst == True:

			# cull photometry on both bands
			# assumes only 2 bands for now
			snr = 5.
			sharp = 0.1
			crowd = 10.
			objtype = 1

			wgood = np.where( (t[filters[0]+'_SNR'] >= snr) & (t[filters[1]+'_SNR'] >= snr) &
						( t[filters[0]+'_SHARP']**2 < sharp ) +  ( t[filters[1]+'_SHARP']**2 < sharp ) &
						( (t[filters[0]+'_CROWD'] + t[filters[1]+'_CROWD']) < crowd ) &
						(t['OBJECT_TYPE'] <= objtype) )

			t1 = t[wgood]
			t1.write(sys.argv[1]+'.gst.fits', overwrite=True)

			'''
			plt.figure()
			plt.scatter(t[filters[0]+'_VEGA'] - t[filters[1]+'_VEGA'],t[filters[1]+'_VEGA'], color='k', edgecolor='None', s=0.5)
			plt.scatter(t1[filters[0]+'_VEGA'] - t1[filters[1]+'_VEGA'],t1[filters[1]+'_VEGA'], color='r', edgecolor='None', s=0.25)
			plt.xlim(-1,4)
			plt.ylim(30,16)
			'''

def main():
	datafile = sys.argv[1]
	outfile = sys.argv[2]
	filters = sys.argv[3]
	ref_image = sys.argv[4]

	make_phot(datafile, ref_image, filters, outfile)

main()

'''
t = astropy.table.Table()
t.add_column(astropy.table.Column(name=global_labels[0], data=data[:,0]))
t.add_column(astropy.table.Column(name=global_labels[1], data=data[:,1]))
t.add_column(astropy.table.Column(name=global_labels[2], data=data[:,2]))
k=3
for i in range(len(filters)):
    for j in range(len(filter_labels)):
        t.add_column(astropy.table.Column(name=filters[i]+filter_labels[j], data=data[:,k]))
        print data[:,k][-1]
        k+=1

t.write(sys.argv[1]+'.fits')
'''

<|endoftext|>


<commit_before># coding: utf-8

""" Classes for observing / reducing data. """

from __future__ import division, print_function

__author__ = "adrn <adrn@astro.columbia.edu>"

# Standard library
import os, sys
from collections import defaultdict
import glob
import logging
import json
import re
from datetime import datetime

# Third-party
from astropy.time import Time
from astropy.io import fits
from astropy.io.misc import fnpickle, fnunpickle
from astropy.modeling import models, fitting
import astropy.units as u
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.gridspec import GridSpec
from matplotlib import cm

# Project
from .util import *
from .arc import find_line_list, ArcSpectrum
from .ccd import *
from .pointing import TelescopePointing

# Create logger
logger = logging.getLogger(__name__)

_line_colors = ["red", "green", "blue", "magenta", "cyan", "yellow"]

class ObservingRun(object):

    def __init__(self, path, ccd, data_path=None, redux_path=None,
                 config_file=None):
        """ An object to store global properties of an observing run,
            such as the CCD configuration, a rough wavelength solution,
            various system paths.

            If only 'path' is specified, assumes that all data for that
            observing run live under path/data and reduction products
            will go under path/reduction.

            Parameters
            ----------
            path : str
                Path to the top level of observing run.
            ccd : CCD
                The CCD used to take data.
            data_path : str (optional)
                Path to data for this observing run. Defaults to path/data.
            redux_path : str (optional)
                Path to store reduction products for this observing run.
                Defaults to path/reduction.
            config_file : str (optional)
                Path to a queue config file. Default is in
                <data_path>/config.json

        """

        self.path = str(path)
        if not os.path.exists(self.path):
            raise IOError("Path {0} does not exist!".format(self.path))

        if not isinstance(ccd, CCD):
            raise TypeError("'ccd' must be a CCD instance.")
        self.ccd = ccd

        self.data_path = data_path
        if self.data_path is None:
            self.data_path = os.path.join(self.path, "data")

        if not os.path.exists(self.data_path):
            raise IOError("Data path {0} does not exist!"\
                          .format(self.data_path))

        self.redux_path = redux_path
        if self.redux_path is None:
            self.redux_path = os.path.join(self.path, "reduction")

        if not os.path.exists(self.redux_path):
            os.mkdir(self.redux_path)

        self.master_arc = None

        # find all nights from the run -- these are subdirectories in
        #   data_path that start with, e.g., m10**13
        self.nights = dict()
        pattr = re.compile("m([0-9]{6})")
        for m_path in glob.glob(os.path.join(self.data_path, "m*")):
            xx,m_date = os.path.split(m_path)
            d = pattr.search(m_date).groups()[0]
            month = int(float(d[:2]))
            day = int(float(d[2:4]))
            yr = int("20" + d[4:])

            utc = Time(datetime(yr,month,day+1), scale="utc")
            self.nights[m_date] = ObservingNight(utc=utc,
                                                 observing_run=self)

        # load the object/queue config
        if config_file is None:
            config_file = os.path.join(self.data_path, "config.json")

        if not os.path.exists(config_file):
            raise IOError("Queue spec/config file '{0}'' does not exist!"\
                          .format(config_file))

        with open(config_file) as f:
            config = json.loads(f.read())

        for night_str,night in self.nights.items():
            for obj in config[night_str]:
                # in json file, i only specify the exposure id. have
                #   to turn this into a filename
                _fn = os.path.join(night.data_path, "{0}.{1:04d}.fit")
                spec_files = [_fn.format(night_str, ii) for ii in obj["spec"]]
                arc_files = [_fn.format(night_str, ii) for ii in obj["arc"]]

                ptg = TelescopePointing(self, spec_files, arc_files)
                night.pointings.append(ptg)

    def make_master_arc(self, night, narcs=10, overwrite=False):
        """ Make a 'master' 1D arc for this observing run, cache it to
            a JSON file in night.data_path.

            Parameters
            ----------
            night : ObservingNight
            narcs : int (optional)
                Number of arc files to median when processing.
            overwrite : bool (optional)
                Overwrite the cache file or not.
        """
        plt.switch_backend("agg")

        narcs = int(narcs)
        cache_file = os.path.join(self.redux_path, "arc", \
                                  "master_arc.json")

        if os.path.exists(cache_file) and overwrite:
            os.remove(cache_file)

        if not os.path.exists(cache_file):
            # first find all COMP files in the data path
            comp_files = find_all_imagetyp(night.data_path, "COMP")
            if len(comp_files) < narcs:
                raise ValueError("Fewer than narcs={0} arcs found in {1}"\
                                 .format(narcs, self.data_path))

            comp_files = comp_files[:narcs]

            # make a 3D data structure to hold data from all comps
            all_comps = np.zeros(self.ccd.shape + (narcs, ))

            for ii,fn in enumerate(comp_files):
                all_comps[...,ii] = fits.getdata(fn, 0)

            # only select out the part of the read-out CCD to use for science
            all_comps = all_comps[self.ccd.regions["data"]]

            # take median over each individual exposure
            median_comp = np.median(all_comps, axis=-1)

            # now take median over the columns in the sub-image
            median_comp_1d = np.median(median_comp, axis=1)
            pix = np.arange(len(median_comp_1d))

            # TODO: pix isn't quite right if select a subsample in the disperson axis...

            with open(cache_file, "w") as f:
                s = dict()
                s["pixel"] = pix.tolist()
                s["counts"] = median_comp_1d.tolist()
                f.write(json.dumps(s))

        self.master_arc = ArcSpectrum.from_json(cache_file)

class ObservingNight(object):

    def __init__(self, utc, observing_run, data_path=None, redux_path=None):
        """ An object to store various properties of a night of an
            observing run, e.g, all bias frames, flats, objects, etc.
            Also, date of run, etc.

            Parameters
            ----------
            utc : astropy.time.Time
                The UT date of the night.
            observing_run : ObservingRun
                The parent ObservingRun object.
            data_path : str (optional)
                Path to data from this night. Defaults to
                <observing_run.data_path>/mMMDDYY where day is *not* the
                UT day of the run.
            redux_path : str (optional)

        """

        self.utc = utc
        self.observing_run = observing_run

        # convention is to use civil date at start of night for data,
        #   which is utc day - 1
        day = "{:02d}".format(utc.datetime.day-1)
        month = utc.datetime.strftime("%m")
        year = utc.datetime.strftime("%y")
        self._night_str = "m{0}{1}{2}".format(month, day, year)

        if data_path is None:
            data_path = os.path.join(self.observing_run.data_path,
                                     str(self))
        self.data_path = data_path

        if redux_path is None:
            redux_path = os.path.join(self.observing_run.redux_path,
                                      str(self))
        self.redux_path = redux_path

        if not os.path.exists(data_path):
            raise ValueError("Path to data ({0}) doesn't exist!"\
                             .format(self.data_path))

        if not os.path.exists(redux_path):
            os.mkdir(redux_path)

        # create a dict with all unique object names as keys, paths to
        #   files as values
        all_object_files = find_all_imagetyp(self.data_path, "OBJECT")
        if len(all_object_files) == 0:
            raise ValueError("No object files found in '{0}'"\
                             .format(self.data_path))

        object_dict = defaultdict(list)
        for filename in all_object_files:
            hdr = fits.getheader(filename)
            object_dict[hdr["OBJECT"]].append(filename)

        self.object_files = object_dict

        self.master_bias = self.make_master_bias()
        self.master_flat = self.make_master_flat()

        # all pointings of the telescope during this night
        self.pointings = []

    def __str__(self):
        return self._night_str

    def make_master_bias(self, overwrite=False):
        """ Make a master bias frame, store image.

            Parameters
            ----------
            overwrite : bool (optional)
        """

        cache_file = os.path.join(self.redux_path, "master_bias.fits")

        if os.path.exists(cache_file) and overwrite:
            os.remove(cache_file)

        if not os.path.exists(cache_file):
            all_bias_files = find_all_imagetyp(self.data_path, "ZERO")
            Nbias = len(all_bias_files)

            all_bias = np.zeros(self.observing_run.ccd.shape + (Nbias,))
            for ii,filename in enumerate(all_bias_files):
                hdu = fits.open(filename)[0]

                if hdu.header["EXPTIME"] > 0:
                    raise ValueError("Bias frame with EXPTIME > 0! ({0})"\
                                     .format(filename))

                all_bias[...,ii] = hdu.data

            master_bias = np.median(all_bias, axis=-1)
            hdu = fits.PrimaryHDU(master_bias)
            hdu.writeto(cache_file)
        else:
            master_bias = fits.getdata(cache_file)

        return master_bias

    def make_master_flat(self, overwrite=False):
        """ Make a master flat frame, store image.

            Parameters
            ----------
            overwrite : bool (optional)
        """

        bias = self.make_master_bias()
        cache_file = os.path.join(self.redux_path, "master_flat.fits")

        ccd = self.observing_run.ccd

        if os.path.exists(cache_file) and overwrite:
            os.remove(cache_file)

        all_flats = None
        if not os.path.exists(cache_file):
            all_flat_files = find_all_imagetyp(self.data_path, "FLAT")
            Nflat = len(all_flat_files)

            for ii,filename in enumerate(all_flat_files):
                data = fits.getdata(filename,0)
                corrected = ccd.bias_correct_frame(data, bias)

                if all_flats is None:
                    all_flats = np.zeros(corrected.shape + (Nflat,))

                all_flats[...,ii] = corrected

            master_flat = np.median(all_flats, axis=-1)
            hdu = fits.PrimaryHDU(master_flat)
            hdu.writeto(cache_file)
        else:
            master_flat = fits.getdata(cache_file)

        return master_flat

    def make_normalized_flat(self, overwrite=False):
        """ Make a normalized Flat frame by fitting and dividing out
            the flat lamp spectrum.

            TODO: right now, this just fits a 1D polynomial. could do
                  a 2D fit to the image instead...

            Parameters
            ----------
            overwrite : bool (optional)
        """

        # TODO: I can't get this 2D fit to work...
        # x, y = np.mgrid[:shp[0], :shp[1]]
        # p = models.Polynomial2DModel(degree=13)
        # fit = fitting.LinearLSQFitter(p)
        # fit(x, y, master_flat)
        # smooth_flat = p(x,y)

        cache_file = os.path.join(self.redux_path, "normalized_flat.fits")

        if os.path.exists(cache_file) and overwrite:
            os.remove(cache_file)

        if not os.path.exists(cache_file):
            # collapse the flat to 1D, fit a polynomial
            master_flat_1d = np.median(self.master_flat, axis=1)
            p = models.Polynomial1DModel(degree=15)
            fit = fitting.NonLinearLSQFitter(p)
            fit(pix, master_flat_1d)
            smooth_flat = p(pix)

            # normalize the flat
            n_flat = self.master_flat / smooth_flat[:,np.newaxis]
            hdu = fits.PrimaryHDU(n_flat)
            hdu.writeto(cache_file)
        else:
            n_flat = fits.getdata(cache_file)

        return n_flat
<commit_msg>fix linelist finder<commit_after># coding: utf-8

""" Classes for observing / reducing data. """

from __future__ import division, print_function

__author__ = "adrn <adrn@astro.columbia.edu>"

# Standard library
import os, sys
from collections import defaultdict
import glob
import logging
import json
import re
from datetime import datetime

# Third-party
from astropy.time import Time
from astropy.io import fits
from astropy.io.misc import fnpickle, fnunpickle
from astropy.modeling import models, fitting
import astropy.units as u
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.gridspec import GridSpec
from matplotlib import cm

# Project
from .util import *
from .arc import find_line_list, ArcSpectrum
from .ccd import *
from .pointing import TelescopePointing

# Create logger
logger = logging.getLogger(__name__)

_line_colors = ["red", "green", "blue", "magenta", "cyan", "yellow"]

class ObservingRun(object):

    def __init__(self, path, ccd, data_path=None, redux_path=None,
                 config_file=None):
        """ An object to store global properties of an observing run,
            such as the CCD configuration, a rough wavelength solution,
            various system paths.

            If only 'path' is specified, assumes that all data for that
            observing run live under path/data and reduction products
            will go under path/reduction.

            Parameters
            ----------
            path : str
                Path to the top level of observing run.
            ccd : CCD
                The CCD used to take data.
            data_path : str (optional)
                Path to data for this observing run. Defaults to path/data.
            redux_path : str (optional)
                Path to store reduction products for this observing run.
                Defaults to path/reduction.
            config_file : str (optional)
                Path to a queue config file. Default is in
                <data_path>/config.json

        """

        self.path = str(path)
        if not os.path.exists(self.path):
            raise IOError("Path {0} does not exist!".format(self.path))

        if not isinstance(ccd, CCD):
            raise TypeError("'ccd' must be a CCD instance.")
        self.ccd = ccd

        self.data_path = data_path
        if self.data_path is None:
            self.data_path = os.path.join(self.path, "data")

        if not os.path.exists(self.data_path):
            raise IOError("Data path {0} does not exist!"\
                          .format(self.data_path))

        self.redux_path = redux_path
        if self.redux_path is None:
            self.redux_path = os.path.join(self.path, "reduction")

        if not os.path.exists(self.redux_path):
            os.mkdir(self.redux_path)

        self.master_arc = None

        # find all nights from the run -- these are subdirectories in
        #   data_path that start with, e.g., m10**13
        self.nights = dict()
        pattr = re.compile("m([0-9]{6})")
        for m_path in glob.glob(os.path.join(self.data_path, "m*")):
            xx,m_date = os.path.split(m_path)
            d = pattr.search(m_date).groups()[0]
            month = int(float(d[:2]))
            day = int(float(d[2:4]))
            yr = int("20" + d[4:])

            utc = Time(datetime(yr,month,day+1), scale="utc")
            self.nights[m_date] = ObservingNight(utc=utc,
                                                 observing_run=self)

        # load the object/queue config
        if config_file is None:
            config_file = os.path.join(self.data_path, "config.json")

        if not os.path.exists(config_file):
            raise IOError("Queue spec/config file '{0}'' does not exist!"\
                          .format(config_file))

        with open(config_file) as f:
            config = json.loads(f.read())

        for night_str,night in self.nights.items():
            for obj in config[night_str]:
                # in json file, i only specify the exposure id. have
                #   to turn this into a filename
                _fn = os.path.join(night.data_path, "{0}.{1:04d}.fit")
                spec_files = [_fn.format(night_str, ii) for ii in obj["spec"]]
                arc_files = [_fn.format(night_str, ii) for ii in obj["arc"]]

                ptg = TelescopePointing(night, spec_files, arc_files)
                night.pointings.append(ptg)

    def make_master_arc(self, night, narcs=10, overwrite=False):
        """ Make a 'master' 1D arc for this observing run, cache it to
            a JSON file in night.data_path.

            Parameters
            ----------
            night : ObservingNight
            narcs : int (optional)
                Number of arc files to median when processing.
            overwrite : bool (optional)
                Overwrite the cache file or not.
        """
        plt.switch_backend("agg")

        narcs = int(narcs)
        cache_file = os.path.join(self.redux_path, "arc", \
                                  "master_arc.json")

        if os.path.exists(cache_file) and overwrite:
            os.remove(cache_file)

        if not os.path.exists(cache_file):
            # first find all COMP files in the data path
            comp_files = find_all_imagetyp(night.data_path, "COMP")
            if len(comp_files) < narcs:
                raise ValueError("Fewer than narcs={0} arcs found in {1}"\
                                 .format(narcs, self.data_path))

            comp_files = comp_files[:narcs]

            # make a 3D data structure to hold data from all comps
            all_comps = np.zeros(self.ccd.shape + (narcs, ))

            for ii,fn in enumerate(comp_files):
                all_comps[...,ii] = fits.getdata(fn, 0)

            # only select out the part of the read-out CCD to use for science
            all_comps = all_comps[self.ccd.regions["data"]]

            # take median over each individual exposure
            median_comp = np.median(all_comps, axis=-1)

            # now take median over the columns in the sub-image
            median_comp_1d = np.median(median_comp, axis=1)
            pix = np.arange(len(median_comp_1d))

            # TODO: pix isn't quite right if select a subsample in the disperson axis...

            with open(cache_file, "w") as f:
                s = dict()
                s["pixel"] = pix.tolist()
                s["counts"] = median_comp_1d.tolist()
                f.write(json.dumps(s))

        self.master_arc = ArcSpectrum.from_json(cache_file)

class ObservingNight(object):

    def __init__(self, utc, observing_run, data_path=None, redux_path=None):
        """ An object to store various properties of a night of an
            observing run, e.g, all bias frames, flats, objects, etc.
            Also, date of run, etc.

            Parameters
            ----------
            utc : astropy.time.Time
                The UT date of the night.
            observing_run : ObservingRun
                The parent ObservingRun object.
            data_path : str (optional)
                Path to data from this night. Defaults to
                <observing_run.data_path>/mMMDDYY where day is *not* the
                UT day of the run.
            redux_path : str (optional)

        """

        self.utc = utc
        self.observing_run = observing_run

        # convention is to use civil date at start of night for data,
        #   which is utc day - 1
        day = "{:02d}".format(utc.datetime.day-1)
        month = utc.datetime.strftime("%m")
        year = utc.datetime.strftime("%y")
        self._night_str = "m{0}{1}{2}".format(month, day, year)

        if data_path is None:
            data_path = os.path.join(self.observing_run.data_path,
                                     str(self))
        self.data_path = data_path

        if redux_path is None:
            redux_path = os.path.join(self.observing_run.redux_path,
                                      str(self))
        self.redux_path = redux_path

        if not os.path.exists(data_path):
            raise ValueError("Path to data ({0}) doesn't exist!"\
                             .format(self.data_path))

        if not os.path.exists(redux_path):
            os.mkdir(redux_path)

        # create a dict with all unique object names as keys, paths to
        #   files as values
        all_object_files = find_all_imagetyp(self.data_path, "OBJECT")
        if len(all_object_files) == 0:
            raise ValueError("No object files found in '{0}'"\
                             .format(self.data_path))

        object_dict = defaultdict(list)
        for filename in all_object_files:
            hdr = fits.getheader(filename)
            object_dict[hdr["OBJECT"]].append(filename)

        self.object_files = object_dict

        self.master_bias = self.make_master_bias()
        self.master_flat = self.make_master_flat()

        # all pointings of the telescope during this night
        self.pointings = []

    def __str__(self):
        return self._night_str

    def make_master_bias(self, overwrite=False):
        """ Make a master bias frame, store image.

            Parameters
            ----------
            overwrite : bool (optional)
        """

        cache_file = os.path.join(self.redux_path, "master_bias.fits")

        if os.path.exists(cache_file) and overwrite:
            os.remove(cache_file)

        if not os.path.exists(cache_file):
            all_bias_files = find_all_imagetyp(self.data_path, "ZERO")
            Nbias = len(all_bias_files)

            all_bias = np.zeros(self.observing_run.ccd.shape + (Nbias,))
            for ii,filename in enumerate(all_bias_files):
                hdu = fits.open(filename)[0]

                if hdu.header["EXPTIME"] > 0:
                    raise ValueError("Bias frame with EXPTIME > 0! ({0})"\
                                     .format(filename))

                all_bias[...,ii] = hdu.data

            master_bias = np.median(all_bias, axis=-1)
            hdu = fits.PrimaryHDU(master_bias)
            hdu.writeto(cache_file)
        else:
            master_bias = fits.getdata(cache_file)

        return master_bias

    def make_master_flat(self, overwrite=False):
        """ Make a master flat frame, store image.

            Parameters
            ----------
            overwrite : bool (optional)
        """

        bias = self.make_master_bias()
        cache_file = os.path.join(self.redux_path, "master_flat.fits")

        ccd = self.observing_run.ccd

        if os.path.exists(cache_file) and overwrite:
            os.remove(cache_file)

        all_flats = None
        if not os.path.exists(cache_file):
            all_flat_files = find_all_imagetyp(self.data_path, "FLAT")
            Nflat = len(all_flat_files)

            for ii,filename in enumerate(all_flat_files):
                data = fits.getdata(filename,0)
                corrected = ccd.bias_correct_frame(data, bias)

                if all_flats is None:
                    all_flats = np.zeros(corrected.shape + (Nflat,))

                all_flats[...,ii] = corrected

            master_flat = np.median(all_flats, axis=-1)
            hdu = fits.PrimaryHDU(master_flat)
            hdu.writeto(cache_file)
        else:
            master_flat = fits.getdata(cache_file)

        return master_flat

    def make_normalized_flat(self, overwrite=False):
        """ Make a normalized Flat frame by fitting and dividing out
            the flat lamp spectrum.

            TODO: right now, this just fits a 1D polynomial. could do
                  a 2D fit to the image instead...

            Parameters
            ----------
            overwrite : bool (optional)
        """

        # TODO: I can't get this 2D fit to work...
        # x, y = np.mgrid[:shp[0], :shp[1]]
        # p = models.Polynomial2DModel(degree=13)
        # fit = fitting.LinearLSQFitter(p)
        # fit(x, y, master_flat)
        # smooth_flat = p(x,y)

        cache_file = os.path.join(self.redux_path, "normalized_flat.fits")

        if os.path.exists(cache_file) and overwrite:
            os.remove(cache_file)

        if not os.path.exists(cache_file):
            # collapse the flat to 1D, fit a polynomial
            master_flat_1d = np.median(self.master_flat, axis=1)
            p = models.Polynomial1DModel(degree=15)
            fit = fitting.NonLinearLSQFitter(p)
            fit(pix, master_flat_1d)
            smooth_flat = p(pix)

            # normalize the flat
            n_flat = self.master_flat / smooth_flat[:,np.newaxis]
            hdu = fits.PrimaryHDU(n_flat)
            hdu.writeto(cache_file)
        else:
            n_flat = fits.getdata(cache_file)

        return n_flat
<|endoftext|>


<commit_before>import os
import random
import string
import pandas as pd

from interfaces.SentenceOperation import SentenceOperation
from tasks.TaskTypes import TaskType
from nltk import word_tokenize

"""
Multilingual Lexicon Perturbation

This perturbation translates words in the text from any supported languages (e.g., English) to other supported languages (e.g., German) by using a multilingual lexicon. It can be used to test the robustness of a model in a multilingual setting.

There are 100 langauges supported as listed below:
Afrikaans (af), Amharic (am), Arabic (ar), Asturian (ast), Azerbaijani (az), Bashkir (ba), Belarusian (be), Bulgarian (bg), Bengali (bn),
Breton (br), Bosnian (bs), Catalan; Valencian (ca), Cebuano (ceb), Czech (cs), Welsh (cy), Danish (da), German (de), Greeek (el), English (en),
Spanish (es), Estonian (et), Persian (fa), Fulah (ff), Finnish (fi), French (fr), Western Frisian (fy), Irish (ga), Gaelic; Scottish Gaelic (gd),
Galician (gl), Gujarati (gu), Hausa (ha), Hebrew (he), Hindi (hi), Croatian (hr), Haitian; Haitian Creole (ht), Hungarian (hu), Armenian (hy),
Indonesian (id), Igbo (ig), Iloko (ilo), Icelandic (is), Italian (it), Japanese (ja), Javanese (jv), Georgian (ka), Kazakh (kk), Central Khmer (km),
Kannada (kn), Korean (ko), Luxembourgish; Letzeburgesch (lb), Ganda (lg), Lingala (ln), Lao (lo), Lithuanian (lt), Latvian (lv), Malagasy (mg),
Macedonian (mk), Malayalam (ml), Mongolian (mn), Marathi (mr), Malay (ms), Burmese (my), Nepali (ne), Dutch; Flemish (nl), Norwegian (no),
Northern Sotho (ns), Occitan (post 1500) (oc), Oriya (or), Panjabi; Punjabi (pa), Polish (pl), Pushto; Pashto (ps), Portuguese (pt),
Romanian; Moldavian; Moldovan (ro), Russian (ru), Sindhi (sd), Sinhala; Sinhalese (si), Slovak (sk), Slovenian (sl), Somali (so), Albanian (sq),
Serbian (sr), Swati (ss), Sundanese (su), Swedish (sv), Swahili (sw), Tamil (ta), Thai (th), Tagalog (tl), Tswana (tn), Turkish (tr),
Ukrainian (uk), Urdu (ur), Uzbek (uz), Vietnamese (vi), Wolof (wo), Xhosa (xh), Yiddish (yi), Yoruba (yo), Chinese (zh), Zulu (zu)
"""

FOLDER_PATH = '/'.join(os.path.abspath(__file__).split('/')[:-1])

def translate(model, tokenizer, text, src_lang, target_lang):
    tokenizer.src_lang = src_lang
    encoded = tokenizer(text, return_tensors="pt")
    generated_tokens = model.generate(**encoded, forced_bos_token_id=tokenizer.get_lang_id(target_lang))
    res = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
    return res

def perturb_sentence(lexicon_df, text, prob_mix=0.5, src_lang="en", trg_lang="zh", seed=0):
    random.seed(seed)
    l_df = lexicon_df.set_index(src_lang)

    words = word_tokenize(text)
    mixed_text = ""
    for word in words:
        # Add space between word
        if mixed_text != "":
            mixed_text += " "

        if word.lower() not in l_df.index:
            mixed_text += word
        else:
            rand_prob = random.random()
            if rand_prob < prob_mix:
                plain_word = word.translate(str.maketrans('', '', string.punctuation)).strip().lower()

                if plain_word == "":
                    continue

                if not word[0].isupper(): # lower case
                    perturbed_word = l_df.loc[plain_word, trg_lang]
                else:
                    perturbed_word = l_df.loc[plain_word, trg_lang].capitalize()

                mixed_text += perturbed_word
            else:
                mixed_text += word

    # print(text, mixed_text)
    return mixed_text

class MultilingualLexiconPerturbation(SentenceOperation):
    tasks = [
        TaskType.TEXT_CLASSIFICATION,
        TaskType.TEXT_TO_TEXT_GENERATION
    ]
    supported_languages = [
        'af', 'am', 'ar', 'ast', 'az', 'ba', 'be', 'bg', 'bn', 'br', 'bs', 'ca', 'ceb', 'cs', 'cy',
        'da', 'de', 'el', 'en', 'es', 'et', 'fa', 'ff', 'fi', 'fr', 'fy', 'ga', 'gd', 'gl', 'gu',
        'ha', 'he', 'hi', 'hr', 'ht', 'hu', 'hy', 'id', 'ig', 'ilo', 'is', 'it', 'ja', 'jv', 'ka',
        'kk', 'km', 'kn', 'ko', 'lb', 'lg', 'ln', 'lo', 'lt', 'lv', 'mg', 'mk', 'ml', 'mn', 'mr',
        'ms', 'my', 'ne', 'nl', 'no', 'ns', 'oc', 'or', 'pa', 'pl', 'ps', 'pt', 'ro', 'ru', 'sd',
        'si', 'sk', 'sl', 'so', 'sq', 'sr', 'ss', 'su', 'sv', 'sw', 'ta', 'th', 'tl', 'tn', 'tr',
        'uk', 'ur', 'uz', 'vi', 'wo', 'xh', 'yi', 'yo', 'zh', 'zu'
    ]
    heavy = False

    def __init__(self, seed=0, prob_mix=0.5, src_lang="en", trg_lang="zh"):
        super().__init__(seed)
        if src_lang not in self.supported_languages:
            raise ValueError(f'Invalid `src_lang` value "{src_lang}". Supported languages: {supported_languages}')

        if trg_lang not in self.supported_languages:
            raise ValueError(f'Invalid `trg_lang` value "{trg_lang}". Supported languages: {supported_languages}')

        self.lexicon_df = pd.read_pickle(f'{FOLDER_PATH}/multilingual_lexicon_uncased.xz')

        self.prob_mix=prob_mix
        self.src_lang=src_lang
        self.trg_lang=trg_lang

    def generate(self, sentence: str):
        pertubed_sentence = perturb_sentence(lexicon_df=self.lexicon_df, text=sentence, prob_mix=self.prob_mix, src_lang=self.src_lang, trg_lang=self.trg_lang, seed=self.seed)
        return pertubed_sentence
<commit_msg>add download nltk punkt before importing word_tokenize<commit_after>import os
import random
import string
import pandas as pd

from interfaces.SentenceOperation import SentenceOperation
from tasks.TaskTypes import TaskType

import nltk
nltk.download('punkt')
from nltk import word_tokenize

"""
Multilingual Lexicon Perturbation

This perturbation translates words in the text from any supported languages (e.g., English) to other supported languages (e.g., German) by using a multilingual lexicon. It can be used to test the robustness of a model in a multilingual setting.

There are 100 langauges supported as listed below:
Afrikaans (af), Amharic (am), Arabic (ar), Asturian (ast), Azerbaijani (az), Bashkir (ba), Belarusian (be), Bulgarian (bg), Bengali (bn),
Breton (br), Bosnian (bs), Catalan; Valencian (ca), Cebuano (ceb), Czech (cs), Welsh (cy), Danish (da), German (de), Greeek (el), English (en),
Spanish (es), Estonian (et), Persian (fa), Fulah (ff), Finnish (fi), French (fr), Western Frisian (fy), Irish (ga), Gaelic; Scottish Gaelic (gd),
Galician (gl), Gujarati (gu), Hausa (ha), Hebrew (he), Hindi (hi), Croatian (hr), Haitian; Haitian Creole (ht), Hungarian (hu), Armenian (hy),
Indonesian (id), Igbo (ig), Iloko (ilo), Icelandic (is), Italian (it), Japanese (ja), Javanese (jv), Georgian (ka), Kazakh (kk), Central Khmer (km),
Kannada (kn), Korean (ko), Luxembourgish; Letzeburgesch (lb), Ganda (lg), Lingala (ln), Lao (lo), Lithuanian (lt), Latvian (lv), Malagasy (mg),
Macedonian (mk), Malayalam (ml), Mongolian (mn), Marathi (mr), Malay (ms), Burmese (my), Nepali (ne), Dutch; Flemish (nl), Norwegian (no),
Northern Sotho (ns), Occitan (post 1500) (oc), Oriya (or), Panjabi; Punjabi (pa), Polish (pl), Pushto; Pashto (ps), Portuguese (pt),
Romanian; Moldavian; Moldovan (ro), Russian (ru), Sindhi (sd), Sinhala; Sinhalese (si), Slovak (sk), Slovenian (sl), Somali (so), Albanian (sq),
Serbian (sr), Swati (ss), Sundanese (su), Swedish (sv), Swahili (sw), Tamil (ta), Thai (th), Tagalog (tl), Tswana (tn), Turkish (tr),
Ukrainian (uk), Urdu (ur), Uzbek (uz), Vietnamese (vi), Wolof (wo), Xhosa (xh), Yiddish (yi), Yoruba (yo), Chinese (zh), Zulu (zu)
"""

FOLDER_PATH = '/'.join(os.path.abspath(__file__).split('/')[:-1])

def translate(model, tokenizer, text, src_lang, target_lang):
    tokenizer.src_lang = src_lang
    encoded = tokenizer(text, return_tensors="pt")
    generated_tokens = model.generate(**encoded, forced_bos_token_id=tokenizer.get_lang_id(target_lang))
    res = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
    return res

def perturb_sentence(lexicon_df, text, prob_mix=0.5, src_lang="en", trg_lang="zh", seed=0):
    random.seed(seed)
    l_df = lexicon_df.set_index(src_lang)

    words = word_tokenize(text)
    mixed_text = ""
    for word in words:
        # Add space between word
        if mixed_text != "":
            mixed_text += " "

        if word.lower() not in l_df.index:
            mixed_text += word
        else:
            rand_prob = random.random()
            if rand_prob < prob_mix:
                plain_word = word.translate(str.maketrans('', '', string.punctuation)).strip().lower()

                if plain_word == "":
                    continue

                if not word[0].isupper(): # lower case
                    perturbed_word = l_df.loc[plain_word, trg_lang]
                else:
                    perturbed_word = l_df.loc[plain_word, trg_lang].capitalize()

                mixed_text += perturbed_word
            else:
                mixed_text += word

    # print(text, mixed_text)
    return mixed_text

class MultilingualLexiconPerturbation(SentenceOperation):
    tasks = [
        TaskType.TEXT_CLASSIFICATION,
        TaskType.TEXT_TO_TEXT_GENERATION
    ]
    supported_languages = [
        'af', 'am', 'ar', 'ast', 'az', 'ba', 'be', 'bg', 'bn', 'br', 'bs', 'ca', 'ceb', 'cs', 'cy',
        'da', 'de', 'el', 'en', 'es', 'et', 'fa', 'ff', 'fi', 'fr', 'fy', 'ga', 'gd', 'gl', 'gu',
        'ha', 'he', 'hi', 'hr', 'ht', 'hu', 'hy', 'id', 'ig', 'ilo', 'is', 'it', 'ja', 'jv', 'ka',
        'kk', 'km', 'kn', 'ko', 'lb', 'lg', 'ln', 'lo', 'lt', 'lv', 'mg', 'mk', 'ml', 'mn', 'mr',
        'ms', 'my', 'ne', 'nl', 'no', 'ns', 'oc', 'or', 'pa', 'pl', 'ps', 'pt', 'ro', 'ru', 'sd',
        'si', 'sk', 'sl', 'so', 'sq', 'sr', 'ss', 'su', 'sv', 'sw', 'ta', 'th', 'tl', 'tn', 'tr',
        'uk', 'ur', 'uz', 'vi', 'wo', 'xh', 'yi', 'yo', 'zh', 'zu'
    ]
    heavy = False

    def __init__(self, seed=0, prob_mix=0.5, src_lang="en", trg_lang="zh"):
        super().__init__(seed)
        if src_lang not in self.supported_languages:
            raise ValueError(f'Invalid `src_lang` value "{src_lang}". Supported languages: {supported_languages}')

        if trg_lang not in self.supported_languages:
            raise ValueError(f'Invalid `trg_lang` value "{trg_lang}". Supported languages: {supported_languages}')

        self.lexicon_df = pd.read_pickle(f'{FOLDER_PATH}/multilingual_lexicon_uncased.xz')

        self.prob_mix=prob_mix
        self.src_lang=src_lang
        self.trg_lang=trg_lang

    def generate(self, sentence: str):
        pertubed_sentence = perturb_sentence(lexicon_df=self.lexicon_df, text=sentence, prob_mix=self.prob_mix, src_lang=self.src_lang, trg_lang=self.trg_lang, seed=self.seed)
        return pertubed_sentence
<|endoftext|>


<commit_before># Copyright 2018, OpenCensus Authors
#
# 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.

"""Compiler options for OpenCensus.

Flags specified here must not impact ABI. Code compiled with and without these
opts will be linked together, and in some cases headers compiled with and
without these options will be part of the same program.

We use the same flags as absl, plus turn some warnings into errors.
"""

load(
    "@com_google_absl//absl:copts/GENERATED_copts.bzl",
    "ABSL_GCC_FLAGS",
    "ABSL_GCC_TEST_FLAGS",
    "ABSL_LLVM_FLAGS",
    "ABSL_LLVM_TEST_FLAGS",
    "ABSL_MSVC_FLAGS",
    "ABSL_MSVC_TEST_FLAGS",
)

WERROR = ["-Werror=return-type", "-Werror=switch"]

DEFAULT_COPTS = select({
    "//opencensus:llvm_compiler": ABSL_LLVM_FLAGS + WERROR,
    "//opencensus:windows": ABSL_MSVC_FLAGS,
    "//conditions:default": ABSL_GCC_FLAGS + WERROR,
})

TEST_COPTS = DEFAULT_COPTS + select({
    "//opencensus:llvm_compiler": ABSL_LLVM_TEST_FLAGS + WERROR,
    "//opencensus:windows": ABSL_MSVC_TEST_FLAGS,
    "//conditions:default": ABSL_GCC_TEST_FLAGS + WERROR,
})
<commit_msg>Disable signed/unsigned comparison warnings. (#454)<commit_after># Copyright 2018, OpenCensus Authors
#
# 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.

"""Compiler options for OpenCensus.

Flags specified here must not impact ABI. Code compiled with and without these
opts will be linked together, and in some cases headers compiled with and
without these options will be part of the same program.

We use the same flags as absl, plus turn some warnings into errors.
"""

load(
    "@com_google_absl//absl:copts/GENERATED_copts.bzl",
    "ABSL_GCC_FLAGS",
    "ABSL_GCC_TEST_FLAGS",
    "ABSL_LLVM_FLAGS",
    "ABSL_LLVM_TEST_FLAGS",
    "ABSL_MSVC_FLAGS",
    "ABSL_MSVC_TEST_FLAGS",
)

WERROR = ["-Werror=return-type", "-Werror=switch"]
WARN_FLAGS = ["-Wno-sign-compare"]

DEFAULT_COPTS = select({
    "//opencensus:llvm_compiler": ABSL_LLVM_FLAGS + WERROR + WARN_FLAGS,
    "//opencensus:windows": ABSL_MSVC_FLAGS,
    "//conditions:default": ABSL_GCC_FLAGS + WERROR + WARN_FLAGS,
})

TEST_COPTS = DEFAULT_COPTS + select({
    "//opencensus:llvm_compiler": ABSL_LLVM_TEST_FLAGS + WERROR + WARN_FLAGS,
    "//opencensus:windows": ABSL_MSVC_TEST_FLAGS,
    "//conditions:default": ABSL_GCC_TEST_FLAGS + WERROR + WARN_FLAGS,
})
<|endoftext|>


<commit_before>import unittest, os
from test import test_support

import warnings
warnings.filterwarnings(
    "ignore",
    category=DeprecationWarning,
    message=".*complex divmod.*are deprecated"
)

from random import random

# These tests ensure that complex math does the right thing

class ComplexTest(unittest.TestCase):

    def assertAlmostEqual(self, a, b):
        if isinstance(a, complex):
            if isinstance(b, complex):
                unittest.TestCase.assertAlmostEqual(self, a.real, b.real)
                unittest.TestCase.assertAlmostEqual(self, a.imag, b.imag)
            else:
                unittest.TestCase.assertAlmostEqual(self, a.real, b)
                unittest.TestCase.assertAlmostEqual(self, a.imag, 0.)
        else:
            if isinstance(b, complex):
                unittest.TestCase.assertAlmostEqual(self, a, b.real)
                unittest.TestCase.assertAlmostEqual(self, 0., b.imag)
            else:
                unittest.TestCase.assertAlmostEqual(self, a, b)

    def assertCloseAbs(self, x, y, eps=1e-9):
        """Return true iff floats x and y "are close\""""
        # put the one with larger magnitude second
        if abs(x) > abs(y):
            x, y = y, x
        if y == 0:
            return abs(x) < eps
        if x == 0:
            return abs(y) < eps
        # check that relative difference < eps
        self.assert_(abs((x-y)/y) < eps)

    def assertClose(self, x, y, eps=1e-9):
        """Return true iff complexes x and y "are close\""""
        self.assertCloseAbs(x.real, y.real, eps)
        self.assertCloseAbs(x.imag, y.imag, eps)

    def assertIs(self, a, b):
        self.assert_(a is b)

    def check_div(self, x, y):
        """Compute complex z=x*y, and check that z/x==y and z/y==x."""
        z = x * y
        if x != 0:
            q = z / x
            self.assertClose(q, y)
            q = z.__truediv__(x)
            self.assertClose(q, y)
        if y != 0:
            q = z / y
            self.assertClose(q, x)
            q = z.__truediv__(y)
            self.assertClose(q, x)

    def test_truediv(self):
        simple_real = [float(i) for i in range(-5, 6)]
        simple_complex = [complex(x, y) for x in simple_real for y in simple_real]
        for x in simple_complex:
            for y in simple_complex:
                self.check_div(x, y)

        # A naive complex division algorithm (such as in 2.0) is very prone to
        # nonsense errors for these (overflows and underflows).
        self.check_div(complex(1e200, 1e200), 1+0j)
        self.check_div(complex(1e-200, 1e-200), 1+0j)

        # Just for fun.
        for i in range(100):
            self.check_div(complex(random(), random()),
                           complex(random(), random()))

        self.assertRaises(ZeroDivisionError, complex.__truediv__, 1+1j, 0+0j)
        # FIXME: The following currently crashes on Alpha
        # self.assertRaises(OverflowError, pow, 1e200+1j, 1e200+1j)

    def test_truediv(self):
        self.assertAlmostEqual(complex.__truediv__(2+0j, 1+1j), 1-1j)
        self.assertRaises(ZeroDivisionError, complex.__truediv__, 1+1j, 0+0j)

    def test_floordiv(self):
        self.assertAlmostEqual(complex.__floordiv__(3+0j, 1.5+0j), 2)
        self.assertRaises(ZeroDivisionError, complex.__floordiv__, 3+0j, 0+0j)

    def test_richcompare(self):
        self.assertRaises(OverflowError, complex.__eq__, 1+1j, 1<<10000)
        self.assertEqual(complex.__lt__(1+1j, None), NotImplemented)
        self.assertIs(complex.__eq__(1+1j, 1+1j), True)
        self.assertIs(complex.__eq__(1+1j, 2+2j), False)
        self.assertIs(complex.__ne__(1+1j, 1+1j), False)
        self.assertIs(complex.__ne__(1+1j, 2+2j), True)
        self.assertRaises(TypeError, complex.__lt__, 1+1j, 2+2j)
        self.assertRaises(TypeError, complex.__le__, 1+1j, 2+2j)
        self.assertRaises(TypeError, complex.__gt__, 1+1j, 2+2j)
        self.assertRaises(TypeError, complex.__ge__, 1+1j, 2+2j)

    def test_mod(self):
        self.assertRaises(ZeroDivisionError, (1+1j).__mod__, 0+0j)

        a = 3.33+4.43j
        try:
            a % 0
        except ZeroDivisionError:
            pass
        else:
            self.fail("modulo parama can't be 0")

    def test_divmod(self):
        self.assertRaises(ZeroDivisionError, divmod, 1+1j, 0+0j)

    def test_pow(self):
        self.assertAlmostEqual(pow(1+1j, 0+0j), 1.0)
        self.assertAlmostEqual(pow(0+0j, 2+0j), 0.0)
        self.assertRaises(ZeroDivisionError, pow, 0+0j, 1j)
        self.assertAlmostEqual(pow(1j, -1), 1/1j)
        self.assertAlmostEqual(pow(1j, 200), 1)
        self.assertRaises(ValueError, pow, 1+1j, 1+1j, 1+1j)

        a = 3.33+4.43j
        self.assertEqual(a ** 0j, 1)
        self.assertEqual(a ** 0.+0.j, 1)

        self.assertEqual(3j ** 0j, 1)
        self.assertEqual(3j ** 0, 1)

        try:
            0j ** a
        except ZeroDivisionError:
            pass
        else:
            self.fail("should fail 0.0 to negative or complex power")

        try:
            0j ** (3-2j)
        except ZeroDivisionError:
            pass
        else:
            self.fail("should fail 0.0 to negative or complex power")

        # The following is used to exercise certain code paths
        self.assertEqual(a ** 105, a ** 105)
        self.assertEqual(a ** -105, a ** -105)
        self.assertEqual(a ** -30, a ** -30)

        self.assertEqual(0.0j ** 0, 1)

        b = 5.1+2.3j
        self.assertRaises(ValueError, pow, a, b, 0)

    def test_boolcontext(self):
        for i in range(100):
            self.assert_(complex(random() + 1e-6, random() + 1e-6))
        self.assert_(not complex(0.0, 0.0))

    def test_conjugate(self):
        self.assertClose(complex(5.3, 9.8).conjugate(), 5.3-9.8j)

    def test_constructor(self):
        class OS:
            def __init__(self, value): self.value = value
            def __complex__(self): return self.value
        class NS(object):
            def __init__(self, value): self.value = value
            def __complex__(self): return self.value
        self.assertEqual(complex(OS(1+10j)), 1+10j)
        self.assertEqual(complex(NS(1+10j)), 1+10j)
        self.assertRaises(TypeError, complex, OS(None))
        self.assertRaises(TypeError, complex, NS(None))

        self.assertAlmostEqual(complex("1+10j"), 1+10j)
        self.assertAlmostEqual(complex(10), 10+0j)
        self.assertAlmostEqual(complex(10.0), 10+0j)
        self.assertAlmostEqual(complex(10), 10+0j)
        self.assertAlmostEqual(complex(10+0j), 10+0j)
        self.assertAlmostEqual(complex(1,10), 1+10j)
        self.assertAlmostEqual(complex(1,10), 1+10j)
        self.assertAlmostEqual(complex(1,10.0), 1+10j)
        self.assertAlmostEqual(complex(1,10), 1+10j)
        self.assertAlmostEqual(complex(1,10), 1+10j)
        self.assertAlmostEqual(complex(1,10.0), 1+10j)
        self.assertAlmostEqual(complex(1.0,10), 1+10j)
        self.assertAlmostEqual(complex(1.0,10), 1+10j)
        self.assertAlmostEqual(complex(1.0,10.0), 1+10j)
        self.assertAlmostEqual(complex(3.14+0j), 3.14+0j)
        self.assertAlmostEqual(complex(3.14), 3.14+0j)
        self.assertAlmostEqual(complex(314), 314.0+0j)
        self.assertAlmostEqual(complex(314), 314.0+0j)
        self.assertAlmostEqual(complex(3.14+0j, 0j), 3.14+0j)
        self.assertAlmostEqual(complex(3.14, 0.0), 3.14+0j)
        self.assertAlmostEqual(complex(314, 0), 314.0+0j)
        self.assertAlmostEqual(complex(314, 0), 314.0+0j)
        self.assertAlmostEqual(complex(0j, 3.14j), -3.14+0j)
        self.assertAlmostEqual(complex(0.0, 3.14j), -3.14+0j)
        self.assertAlmostEqual(complex(0j, 3.14), 3.14j)
        self.assertAlmostEqual(complex(0.0, 3.14), 3.14j)
        self.assertAlmostEqual(complex("1"), 1+0j)
        self.assertAlmostEqual(complex("1j"), 1j)
        self.assertAlmostEqual(complex(),  0)
        self.assertAlmostEqual(complex("-1"), -1)
        self.assertAlmostEqual(complex("+1"), +1)
        self.assertAlmostEqual(complex("(1+2j)"), 1+2j)
        self.assertAlmostEqual(complex("(1.3+2.2j)"), 1.3+2.2j)

        class complex2(complex): pass
        self.assertAlmostEqual(complex(complex2(1+1j)), 1+1j)
        self.assertAlmostEqual(complex(real=17, imag=23), 17+23j)
        self.assertAlmostEqual(complex(real=17+23j), 17+23j)
        self.assertAlmostEqual(complex(real=17+23j, imag=23), 17+46j)
        self.assertAlmostEqual(complex(real=1+2j, imag=3+4j), -3+5j)

        c = 3.14 + 1j
        self.assert_(complex(c) is c)
        del c

        self.assertRaises(TypeError, complex, "1", "1")
        self.assertRaises(TypeError, complex, 1, "1")

        self.assertEqual(complex("  3.14+J  "), 3.14+1j)

        # SF bug 543840:  complex(string) accepts strings with \0
        # Fixed in 2.3.
        self.assertRaises(ValueError, complex, '1+1j\0j')

        self.assertRaises(TypeError, int, 5+3j)
        self.assertRaises(TypeError, int, 5+3j)
        self.assertRaises(TypeError, float, 5+3j)
        self.assertRaises(ValueError, complex, "")
        self.assertRaises(TypeError, complex, None)
        self.assertRaises(ValueError, complex, "\0")
        self.assertRaises(ValueError, complex, "3\09")
        self.assertRaises(TypeError, complex, "1", "2")
        self.assertRaises(TypeError, complex, "1", 42)
        self.assertRaises(TypeError, complex, 1, "2")
        self.assertRaises(ValueError, complex, "1+")
        self.assertRaises(ValueError, complex, "1+1j+1j")
        self.assertRaises(ValueError, complex, "--")
        self.assertRaises(ValueError, complex, "(1+2j")
        self.assertRaises(ValueError, complex, "1+2j)")
        self.assertRaises(ValueError, complex, "1+(2j)")
        self.assertRaises(ValueError, complex, "(1+2j)123")
        self.assertRaises(ValueError, complex, "1"*500)
        self.assertRaises(ValueError, complex, "x")

        class EvilExc(Exception):
            pass

        class evilcomplex:
            def __complex__(self):
                raise EvilExc

        self.assertRaises(EvilExc, complex, evilcomplex())

        class float2:
            def __init__(self, value):
                self.value = value
            def __float__(self):
                return self.value

        self.assertAlmostEqual(complex(float2(42.)), 42)
        self.assertAlmostEqual(complex(real=float2(17.), imag=float2(23.)), 17+23j)
        self.assertRaises(TypeError, complex, float2(None))

        class complex0(complex):
            """Test usage of __complex__() when inheriting from 'complex'"""
            def __complex__(self):
                return 42j

        class complex1(complex):
            """Test usage of __complex__() with a __new__() method"""
            def __new__(self, value=0j):
                return complex.__new__(self, 2*value)
            def __complex__(self):
                return self

        class complex2(complex):
            """Make sure that __complex__() calls fail if anything other than a
            complex is returned"""
            def __complex__(self):
                return None

        self.assertAlmostEqual(complex(complex0(1j)), 42j)
        self.assertAlmostEqual(complex(complex1(1j)), 2j)
        self.assertRaises(TypeError, complex, complex2(1j))

    def test_hash(self):
        for x in range(-30, 30):
            self.assertEqual(hash(x), hash(complex(x, 0)))
            x /= 3.0    # now check against floating point
            self.assertEqual(hash(x), hash(complex(x, 0.)))

    def test_abs(self):
        nums = [complex(x/3., y/7.) for x in range(-9,9) for y in range(-9,9)]
        for num in nums:
            self.assertAlmostEqual((num.real**2 + num.imag**2)  ** 0.5, abs(num))

    def test_repr(self):
        self.assertEqual(repr(1+6j), '(1+6j)')
        self.assertEqual(repr(1-6j), '(1-6j)')

        self.assertNotEqual(repr(-(1+0j)), '(-1+-0j)')

        self.assertEqual(1-6j,complex(repr(1-6j)))
        self.assertEqual(1+6j,complex(repr(1+6j)))
        self.assertEqual(-6j,complex(repr(-6j)))
        self.assertEqual(6j,complex(repr(6j)))

    def test_neg(self):
        self.assertEqual(-(1+6j), -1-6j)

    def test_file(self):
        a = 3.33+4.43j
        b = 5.1+2.3j

        fo = None
        try:
            fo = open(test_support.TESTFN, "wb")
            print(a, b, file=fo)
            fo.close()
            fo = open(test_support.TESTFN, "rb")
            self.assertEqual(fo.read(), bytes("%s %s\n" % (a, b)))
        finally:
            if (fo is not None) and (not fo.closed):
                fo.close()
            try:
                os.remove(test_support.TESTFN)
            except (OSError, IOError):
                pass

def test_main():
    test_support.run_unittest(ComplexTest)

if __name__ == "__main__":
    test_main()
<commit_msg>Make test_complex pass again now that floordiv and mod are illegal.<commit_after>import unittest, os
from test import test_support

import warnings
warnings.filterwarnings(
    "ignore",
    category=DeprecationWarning,
    message=".*complex divmod.*are deprecated"
)

from random import random

# These tests ensure that complex math does the right thing

class ComplexTest(unittest.TestCase):

    def assertAlmostEqual(self, a, b):
        if isinstance(a, complex):
            if isinstance(b, complex):
                unittest.TestCase.assertAlmostEqual(self, a.real, b.real)
                unittest.TestCase.assertAlmostEqual(self, a.imag, b.imag)
            else:
                unittest.TestCase.assertAlmostEqual(self, a.real, b)
                unittest.TestCase.assertAlmostEqual(self, a.imag, 0.)
        else:
            if isinstance(b, complex):
                unittest.TestCase.assertAlmostEqual(self, a, b.real)
                unittest.TestCase.assertAlmostEqual(self, 0., b.imag)
            else:
                unittest.TestCase.assertAlmostEqual(self, a, b)

    def assertCloseAbs(self, x, y, eps=1e-9):
        """Return true iff floats x and y "are close\""""
        # put the one with larger magnitude second
        if abs(x) > abs(y):
            x, y = y, x
        if y == 0:
            return abs(x) < eps
        if x == 0:
            return abs(y) < eps
        # check that relative difference < eps
        self.assert_(abs((x-y)/y) < eps)

    def assertClose(self, x, y, eps=1e-9):
        """Return true iff complexes x and y "are close\""""
        self.assertCloseAbs(x.real, y.real, eps)
        self.assertCloseAbs(x.imag, y.imag, eps)

    def assertIs(self, a, b):
        self.assert_(a is b)

    def check_div(self, x, y):
        """Compute complex z=x*y, and check that z/x==y and z/y==x."""
        z = x * y
        if x != 0:
            q = z / x
            self.assertClose(q, y)
            q = z.__truediv__(x)
            self.assertClose(q, y)
        if y != 0:
            q = z / y
            self.assertClose(q, x)
            q = z.__truediv__(y)
            self.assertClose(q, x)

    def test_truediv(self):
        simple_real = [float(i) for i in range(-5, 6)]
        simple_complex = [complex(x, y) for x in simple_real for y in simple_real]
        for x in simple_complex:
            for y in simple_complex:
                self.check_div(x, y)

        # A naive complex division algorithm (such as in 2.0) is very prone to
        # nonsense errors for these (overflows and underflows).
        self.check_div(complex(1e200, 1e200), 1+0j)
        self.check_div(complex(1e-200, 1e-200), 1+0j)

        # Just for fun.
        for i in range(100):
            self.check_div(complex(random(), random()),
                           complex(random(), random()))

        self.assertRaises(ZeroDivisionError, complex.__truediv__, 1+1j, 0+0j)
        # FIXME: The following currently crashes on Alpha
        # self.assertRaises(OverflowError, pow, 1e200+1j, 1e200+1j)

    def test_truediv(self):
        self.assertAlmostEqual(complex.__truediv__(2+0j, 1+1j), 1-1j)
        self.assertRaises(ZeroDivisionError, complex.__truediv__, 1+1j, 0+0j)

    def test_floordiv(self):
        self.assertRaises(TypeError, complex.__floordiv__, 3+0j, 1.5+0j)
        self.assertRaises(TypeError, complex.__floordiv__, 3+0j, 0+0j)

    def test_richcompare(self):
        self.assertRaises(OverflowError, complex.__eq__, 1+1j, 1<<10000)
        self.assertEqual(complex.__lt__(1+1j, None), NotImplemented)
        self.assertIs(complex.__eq__(1+1j, 1+1j), True)
        self.assertIs(complex.__eq__(1+1j, 2+2j), False)
        self.assertIs(complex.__ne__(1+1j, 1+1j), False)
        self.assertIs(complex.__ne__(1+1j, 2+2j), True)
        self.assertRaises(TypeError, complex.__lt__, 1+1j, 2+2j)
        self.assertRaises(TypeError, complex.__le__, 1+1j, 2+2j)
        self.assertRaises(TypeError, complex.__gt__, 1+1j, 2+2j)
        self.assertRaises(TypeError, complex.__ge__, 1+1j, 2+2j)

    def test_mod(self):
        # % is no longer supported on complex numbers
        self.assertRaises(TypeError, (1+1j).__mod__, 0+0j)
        self.assertRaises(TypeError, lambda: (3.33+4.43j) % 0)

    def test_divmod(self):
        self.assertRaises(TypeError, divmod, 1+1j, 1+0j)
        self.assertRaises(TypeError, divmod, 1+1j, 0+0j)

    def test_pow(self):
        self.assertAlmostEqual(pow(1+1j, 0+0j), 1.0)
        self.assertAlmostEqual(pow(0+0j, 2+0j), 0.0)
        self.assertRaises(ZeroDivisionError, pow, 0+0j, 1j)
        self.assertAlmostEqual(pow(1j, -1), 1/1j)
        self.assertAlmostEqual(pow(1j, 200), 1)
        self.assertRaises(ValueError, pow, 1+1j, 1+1j, 1+1j)

        a = 3.33+4.43j
        self.assertEqual(a ** 0j, 1)
        self.assertEqual(a ** 0.+0.j, 1)

        self.assertEqual(3j ** 0j, 1)
        self.assertEqual(3j ** 0, 1)

        try:
            0j ** a
        except ZeroDivisionError:
            pass
        else:
            self.fail("should fail 0.0 to negative or complex power")

        try:
            0j ** (3-2j)
        except ZeroDivisionError:
            pass
        else:
            self.fail("should fail 0.0 to negative or complex power")

        # The following is used to exercise certain code paths
        self.assertEqual(a ** 105, a ** 105)
        self.assertEqual(a ** -105, a ** -105)
        self.assertEqual(a ** -30, a ** -30)

        self.assertEqual(0.0j ** 0, 1)

        b = 5.1+2.3j
        self.assertRaises(ValueError, pow, a, b, 0)

    def test_boolcontext(self):
        for i in range(100):
            self.assert_(complex(random() + 1e-6, random() + 1e-6))
        self.assert_(not complex(0.0, 0.0))

    def test_conjugate(self):
        self.assertClose(complex(5.3, 9.8).conjugate(), 5.3-9.8j)

    def test_constructor(self):
        class OS:
            def __init__(self, value): self.value = value
            def __complex__(self): return self.value
        class NS(object):
            def __init__(self, value): self.value = value
            def __complex__(self): return self.value
        self.assertEqual(complex(OS(1+10j)), 1+10j)
        self.assertEqual(complex(NS(1+10j)), 1+10j)
        self.assertRaises(TypeError, complex, OS(None))
        self.assertRaises(TypeError, complex, NS(None))

        self.assertAlmostEqual(complex("1+10j"), 1+10j)
        self.assertAlmostEqual(complex(10), 10+0j)
        self.assertAlmostEqual(complex(10.0), 10+0j)
        self.assertAlmostEqual(complex(10), 10+0j)
        self.assertAlmostEqual(complex(10+0j), 10+0j)
        self.assertAlmostEqual(complex(1,10), 1+10j)
        self.assertAlmostEqual(complex(1,10), 1+10j)
        self.assertAlmostEqual(complex(1,10.0), 1+10j)
        self.assertAlmostEqual(complex(1,10), 1+10j)
        self.assertAlmostEqual(complex(1,10), 1+10j)
        self.assertAlmostEqual(complex(1,10.0), 1+10j)
        self.assertAlmostEqual(complex(1.0,10), 1+10j)
        self.assertAlmostEqual(complex(1.0,10), 1+10j)
        self.assertAlmostEqual(complex(1.0,10.0), 1+10j)
        self.assertAlmostEqual(complex(3.14+0j), 3.14+0j)
        self.assertAlmostEqual(complex(3.14), 3.14+0j)
        self.assertAlmostEqual(complex(314), 314.0+0j)
        self.assertAlmostEqual(complex(314), 314.0+0j)
        self.assertAlmostEqual(complex(3.14+0j, 0j), 3.14+0j)
        self.assertAlmostEqual(complex(3.14, 0.0), 3.14+0j)
        self.assertAlmostEqual(complex(314, 0), 314.0+0j)
        self.assertAlmostEqual(complex(314, 0), 314.0+0j)
        self.assertAlmostEqual(complex(0j, 3.14j), -3.14+0j)
        self.assertAlmostEqual(complex(0.0, 3.14j), -3.14+0j)
        self.assertAlmostEqual(complex(0j, 3.14), 3.14j)
        self.assertAlmostEqual(complex(0.0, 3.14), 3.14j)
        self.assertAlmostEqual(complex("1"), 1+0j)
        self.assertAlmostEqual(complex("1j"), 1j)
        self.assertAlmostEqual(complex(),  0)
        self.assertAlmostEqual(complex("-1"), -1)
        self.assertAlmostEqual(complex("+1"), +1)
        self.assertAlmostEqual(complex("(1+2j)"), 1+2j)
        self.assertAlmostEqual(complex("(1.3+2.2j)"), 1.3+2.2j)

        class complex2(complex): pass
        self.assertAlmostEqual(complex(complex2(1+1j)), 1+1j)
        self.assertAlmostEqual(complex(real=17, imag=23), 17+23j)
        self.assertAlmostEqual(complex(real=17+23j), 17+23j)
        self.assertAlmostEqual(complex(real=17+23j, imag=23), 17+46j)
        self.assertAlmostEqual(complex(real=1+2j, imag=3+4j), -3+5j)

        c = 3.14 + 1j
        self.assert_(complex(c) is c)
        del c

        self.assertRaises(TypeError, complex, "1", "1")
        self.assertRaises(TypeError, complex, 1, "1")

        self.assertEqual(complex("  3.14+J  "), 3.14+1j)

        # SF bug 543840:  complex(string) accepts strings with \0
        # Fixed in 2.3.
        self.assertRaises(ValueError, complex, '1+1j\0j')

        self.assertRaises(TypeError, int, 5+3j)
        self.assertRaises(TypeError, int, 5+3j)
        self.assertRaises(TypeError, float, 5+3j)
        self.assertRaises(ValueError, complex, "")
        self.assertRaises(TypeError, complex, None)
        self.assertRaises(ValueError, complex, "\0")
        self.assertRaises(ValueError, complex, "3\09")
        self.assertRaises(TypeError, complex, "1", "2")
        self.assertRaises(TypeError, complex, "1", 42)
        self.assertRaises(TypeError, complex, 1, "2")
        self.assertRaises(ValueError, complex, "1+")
        self.assertRaises(ValueError, complex, "1+1j+1j")
        self.assertRaises(ValueError, complex, "--")
        self.assertRaises(ValueError, complex, "(1+2j")
        self.assertRaises(ValueError, complex, "1+2j)")
        self.assertRaises(ValueError, complex, "1+(2j)")
        self.assertRaises(ValueError, complex, "(1+2j)123")
        self.assertRaises(ValueError, complex, "1"*500)
        self.assertRaises(ValueError, complex, "x")

        class EvilExc(Exception):
            pass

        class evilcomplex:
            def __complex__(self):
                raise EvilExc

        self.assertRaises(EvilExc, complex, evilcomplex())

        class float2:
            def __init__(self, value):
                self.value = value
            def __float__(self):
                return self.value

        self.assertAlmostEqual(complex(float2(42.)), 42)
        self.assertAlmostEqual(complex(real=float2(17.), imag=float2(23.)), 17+23j)
        self.assertRaises(TypeError, complex, float2(None))

        class complex0(complex):
            """Test usage of __complex__() when inheriting from 'complex'"""
            def __complex__(self):
                return 42j

        class complex1(complex):
            """Test usage of __complex__() with a __new__() method"""
            def __new__(self, value=0j):
                return complex.__new__(self, 2*value)
            def __complex__(self):
                return self

        class complex2(complex):
            """Make sure that __complex__() calls fail if anything other than a
            complex is returned"""
            def __complex__(self):
                return None

        self.assertAlmostEqual(complex(complex0(1j)), 42j)
        self.assertAlmostEqual(complex(complex1(1j)), 2j)
        self.assertRaises(TypeError, complex, complex2(1j))

    def test_hash(self):
        for x in range(-30, 30):
            self.assertEqual(hash(x), hash(complex(x, 0)))
            x /= 3.0    # now check against floating point
            self.assertEqual(hash(x), hash(complex(x, 0.)))

    def test_abs(self):
        nums = [complex(x/3., y/7.) for x in range(-9,9) for y in range(-9,9)]
        for num in nums:
            self.assertAlmostEqual((num.real**2 + num.imag**2)  ** 0.5, abs(num))

    def test_repr(self):
        self.assertEqual(repr(1+6j), '(1+6j)')
        self.assertEqual(repr(1-6j), '(1-6j)')

        self.assertNotEqual(repr(-(1+0j)), '(-1+-0j)')

        self.assertEqual(1-6j,complex(repr(1-6j)))
        self.assertEqual(1+6j,complex(repr(1+6j)))
        self.assertEqual(-6j,complex(repr(-6j)))
        self.assertEqual(6j,complex(repr(6j)))

    def test_neg(self):
        self.assertEqual(-(1+6j), -1-6j)

    def test_file(self):
        a = 3.33+4.43j
        b = 5.1+2.3j

        fo = None
        try:
            fo = open(test_support.TESTFN, "wb")
            print(a, b, file=fo)
            fo.close()
            fo = open(test_support.TESTFN, "rb")
            self.assertEqual(fo.read(), bytes("%s %s\n" % (a, b)))
        finally:
            if (fo is not None) and (not fo.closed):
                fo.close()
            try:
                os.remove(test_support.TESTFN)
            except (OSError, IOError):
                pass

def test_main():
    test_support.run_unittest(ComplexTest)

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