<commit_before>from __future__ import unicode_literals
import frappe
import frappe.defaults
import unittest

from frappe.contacts.report.addresses_and_contacts.addresses_and_contacts import get_data

def get_custom_linked_doctype():
	if bool(frappe.get_all("DocType", filters={'name':'Test Custom Doctype'})):
		return

	doc = frappe.get_doc({
		"doctype": "DocType",
		"module": "Core",
		"custom": 1,
		"fields": [{
			"label": "Test Field",
			"fieldname": "test_field",
			"fieldtype": "Data"
		},
		{
			"label": "Contact HTML",
			"fieldname": "contact_html",
			"fieldtype": "HTML"
		},
		{
			"label": "Address HTML",
			"fieldname": "address_html",
			"fieldtype": "HTML"
		}],
		"permissions": [{
			"role": "System Manager",
			"read": 1
		}],
		"name": "Test Custom Doctype",
	})
	doc.insert()

def get_custom_doc_for_address_and_contacts():
	get_custom_linked_doctype()
	linked_doc = frappe.get_doc({
		"doctype": "Test Custom Doctype",
		"test_field": "Hello",
	}).insert()
	return linked_doc

def create_linked_address(link_list):
	if frappe.flags.test_address_created:
		return

	address = frappe.get_doc({
		"doctype": "Address",
		"address_title": "_Test Address",
		"address_type": "Billing",
		"address_line1": "test address line 1",
		"address_line2": "test address line 2",
		"city": "Milan",
		"country": "Italy"
	})

	for name in link_list:
		address.append("links",{
			'link_doctype': 'Test Custom Doctype',
			'link_name': name
		})

	address.insert()
	frappe.flags.test_address_created = True

def create_linked_contact(link_list):
	if frappe.flags.test_contact_created:
		return

	contact = frappe.get_doc({
		"doctype": "Contact",
		"salutation": "Mr",
		"email_id": "test_contact@example.com",
		"first_name": "_Test First Name",
		"last_name": "_Test Last Name",
		"is_primary_contact": 1,
		"phone": "+91 0000000000",
		"status": "Open"
	})

	for name in link_list:
		contact.append("links",{
			'link_doctype': 'Test Custom Doctype',
			'link_name': name
		})

	contact.insert()
	frappe.flags.test_contact_created = True

class TestAddressesAndContacts(unittest.TestCase):
	def test_get_data(self):
		linked_docs = [get_custom_doc_for_address_and_contacts()]
		links_list = [item.name for item in linked_docs]
		create_linked_contact(links_list)
		create_linked_address(links_list)
		report_data = get_data({"reference_doctype": "Test Custom Doctype"})
		for link in links_list:
			test_item = [link, 'test address line 1', 'test address line 2', 'Milan', None, None, 'Italy', 0, '_Test First Name', '_Test Last Name', '+91 0000000000', None, 'test_contact@example.com', 1]
			self.assertEquals(test_item, report_data)

	def tearDown(self):
		frappe.db.rollback()
<commit_msg>fix: test cases using assertListEqual<commit_after>from __future__ import unicode_literals
import frappe
import frappe.defaults
import unittest

from frappe.contacts.report.addresses_and_contacts.addresses_and_contacts import get_data

def get_custom_linked_doctype():
	if bool(frappe.get_all("DocType", filters={'name':'Test Custom Doctype'})):
		return

	doc = frappe.get_doc({
		"doctype": "DocType",
		"module": "Core",
		"custom": 1,
		"fields": [{
			"label": "Test Field",
			"fieldname": "test_field",
			"fieldtype": "Data"
		},
		{
			"label": "Contact HTML",
			"fieldname": "contact_html",
			"fieldtype": "HTML"
		},
		{
			"label": "Address HTML",
			"fieldname": "address_html",
			"fieldtype": "HTML"
		}],
		"permissions": [{
			"role": "System Manager",
			"read": 1
		}],
		"name": "Test Custom Doctype",
	})
	doc.insert()

def get_custom_doc_for_address_and_contacts():
	get_custom_linked_doctype()
	linked_doc = frappe.get_doc({
		"doctype": "Test Custom Doctype",
		"test_field": "Hello",
	}).insert()
	return linked_doc

def create_linked_address(link_list):
	if frappe.flags.test_address_created:
		return

	address = frappe.get_doc({
		"doctype": "Address",
		"address_title": "_Test Address",
		"address_type": "Billing",
		"address_line1": "test address line 1",
		"address_line2": "test address line 2",
		"city": "Milan",
		"country": "Italy"
	})

	for name in link_list:
		address.append("links",{
			'link_doctype': 'Test Custom Doctype',
			'link_name': name
		})

	address.insert()
	frappe.flags.test_address_created = True

def create_linked_contact(link_list):
	if frappe.flags.test_contact_created:
		return

	contact = frappe.get_doc({
		"doctype": "Contact",
		"salutation": "Mr",
		"email_id": "test_contact@example.com",
		"first_name": "_Test First Name",
		"last_name": "_Test Last Name",
		"is_primary_contact": 1,
		"phone": "+91 0000000000",
		"status": "Open"
	})

	for name in link_list:
		contact.append("links",{
			'link_doctype': 'Test Custom Doctype',
			'link_name': name
		})

	contact.insert()
	frappe.flags.test_contact_created = True

class TestAddressesAndContacts(unittest.TestCase):
	def test_get_data(self):
		linked_docs = [get_custom_doc_for_address_and_contacts()]
		links_list = [item.name for item in linked_docs]
		create_linked_contact(links_list)
		create_linked_address(links_list)
		report_data = get_data({"reference_doctype": "Test Custom Doctype"})
		for idx, link in enumerate(links_list):
			test_item = [link, 'test address line 1', 'test address line 2', 'Milan', None, None, 'Italy', 0, '_Test First Name', '_Test Last Name', '+91 0000000000', None, 'test_contact@example.com', 1]
			self.assertListEqual(test_item, report_data[idx])

	def tearDown(self):
		frappe.db.rollback()
<|endoftext|>


<commit_before>"""
Transform samples of heavy-tailed data so that it is approximately normally distributed

Based on the paper:
The Lambert Way to Gaussianize heavy tailed data with the inverse of Tukey's h transformation as a special case
Based on algorithm in Appendix C
Author generously provides code in R:
https://cran.r-project.org/web/packages/LambertW/

This code written by Greg Ver Steeg, 2015.
"""

import numpy as np
from scipy.special import lambertw
from scipy.stats import kurtosis
from scipy.optimize import fmin  # TODO: Explore efficacy of other opt. methods

np.seterr(all='warn')

class Lambert(object):
    """
    Gaussianize heavy-tailed data using Lambert's W.

    Conventions
    ----------
    This class is a wrapper that follows sklearn naming/style (e.g. fit(X) to train).
    In this code, x is the input, y is the output. But in the functions outside the class, I follow
    Georg's convention that Y is the input and X is the output (Gaussianized) data.

    Parameters
    ----------
    tol : float, default = 1e-4

    max_iter : int, default = 200
        Maximum number of iterations to search for correct parameters of Lambert transform.

    Attributes
    ----------
    taus : list of tuples
        For each variable, we have a tuple consisting of (mu, sigma, delta), corresponding to the parameters of the
        appropriate Lambert transform. Eq. 6 and 8 in the paper below.

    References
    ----------
    [1] Georg Goerg. The Lambert Way to Gaussianize heavy tailed data with
                        the inverse of Tukey's h transformation as a special case
    """

    def __init__(self, tol=1.22e-4, max_iter=100):
        self.tol = tol
        self.max_iter = max_iter
        self.taus = []  # Store tau for each transformed variable

    def fit(self, x):
        x = np.asarray(x)
        if len(x.shape) == 1:
            x = x[:, np.newaxis]
        elif len(x.shape) != 2:
            print "Data should be a 1-d list of samples to transform or a 2d array with samples as rows."
        for x_i in x.T:
            self.taus.append(igmm(x_i, tol=self.tol, max_iter=self.max_iter))

    def transform(self, x):
        x = np.asarray(x)
        if len(x.shape) == 1:
            x = x[:, np.newaxis]
        elif len(x.shape) != 2:
            print "Data should be a 1-d list of samples to transform or a 2d array with samples as rows."
        if x.shape[1] != len(self.taus):
            print "%d variables in test data, but %d variables were in training data." % (x.shape[1], len(self.taus))

        return np.array([w_t(x_i, tau_i) for x_i, tau_i in zip(x.T, self.taus)]).T

    def fit_transform(self, x):
        self.fit(x)
        return self.transform(x)

    def invert(self, y):
        return np.array([inverse(y_i, tau_i) for y_i, tau_i in zip(y.T, self.taus)]).T

    def qqplot(self, x, prefix='qq'):
        """Show qq plots compared to normal before and after the transform."""
        import pylab
        from scipy.stats import probplot
        y = self.transform(x)

        for i, (x_i, y_i) in enumerate(zip(x.T, y.T)):
            probplot(x_i, dist="norm", plot=pylab)
            pylab.savefig(prefix + '_%d_before.png' % i)
            pylab.clf()

            probplot(y_i, dist="norm", plot=pylab)
            pylab.savefig(prefix + '_%d_after.png' % i)
            pylab.clf()

def w_d(z, delta):
    # Eq. 9
    if delta < 1e-6:
        return z
    return np.sign(z) * np.sqrt(np.real(lambertw(delta * z ** 2)) / delta)

def w_t(y, tau):
    # Eq. 8
    return tau[0] + tau[1] * w_d((y - tau[0]) / tau[1], tau[2])

def inverse(x, tau):
    # Eq. 6
    u = (x - tau[0]) / tau[1]
    return tau[0] + tau[1] * (u * np.exp(u * u * (tau[2] * 0.5)))

def igmm(y, tol=1.22e-4, max_iter=100):
    # Infer mu, sigma, delta using IGMM in Alg.2, Appendix C
    delta0 = delta_init(y)
    tau1 = (np.median(y), np.std(y) * (1. - 2. * delta0) ** 0.75, delta0)
    for k in range(max_iter):
        tau0 = tau1
        z = (y - tau1[0]) / tau1[1]
        delta1 = delta_gmm(z)
        x = tau0[0] + tau1[1] * w_d(z, delta1)
        mu1, sigma1 = np.mean(x), np.std(x)
        tau1 = (mu1, sigma1, delta1)

        if np.linalg.norm(np.array(tau1) - np.array(tau0)) < tol:
            break
        else:
            if k == max_iter - 1:
                print "Warning: No convergence after %d iterations. Increase max_iter." % max_iter
    return tau1

def delta_gmm(z):
    # Alg. 1, Appendix C
    delta0 = delta_init(z)

    def func(q):
        u = w_d(z, np.exp(q))
        if not np.all(np.isfinite(u)):
            return 0.
        else:
            k = kurtosis(u, fisher=True, bias=False)**2
            if not np.isfinite(k) or k > 1e10:
                return 1e10
            else:
                return k

    res = fmin(func, np.log(delta0), disp=0)
    return np.around(np.exp(res[-1]), 6)

def delta_init(z):
    gamma = kurtosis(z, fisher=False, bias=False)
    with np.errstate(all='ignore'):
        delta0 = np.clip(1. / 66 * (np.sqrt(66 * gamma - 162.) - 6.), 0.01, 0.25)
    if not np.isfinite(delta0):
        delta0 = 0.01
    return delta0

if __name__ == '__main__':
    # Command line interface
    # Sample commands:
    # python gaussianize.py test_data.csv
    import csv
    import sys, os
    import traceback
    from optparse import OptionParser, OptionGroup

    parser = OptionParser(usage="usage: %prog [options] data_file.csv \n"
                                "It is assumed that the first row and first column of the data CSV file are labels.\n"
                                "Use options to indicate otherwise.")
    group = OptionGroup(parser, "Input Data Format Options")
    group.add_option("-c", "--no_column_names",
                     action="store_true", dest="nc", default=False,
                     help="We assume the top row is variable names for each column. "
                          "This flag says that data starts on the first row and gives a "
                          "default numbering scheme to the variables (1,2,3...).")
    group.add_option("-r", "--no_row_names",
                     action="store_true", dest="nr", default=False,
                     help="We assume the first column is a label or index for each sample. "
                          "This flag says that data starts on the first column.")
    group.add_option("-d", "--delimiter",
                     action="store", dest="delimiter", type="string", default=",",
                     help="Separator between entries in the data, default is ','.")
    parser.add_option_group(group)

    group = OptionGroup(parser, "Output Options")
    group.add_option("-o", "--output",
                     action="store", dest="output", type="string", default="gaussian_output.csv",
                     help="Where to store gaussianized data.")
    group.add_option("-q", "--qqplots",
                     action="store_true", dest="q", default=False,
                     help="Produce qq plots for each variable before and after transform.")
    parser.add_option_group(group)

    (options, args) = parser.parse_args()
    if not len(args) == 1:
        print "Run with '-h' option for usage help."
        sys.exit()

    #Load data from csv file
    filename = args[0]
    with open(filename, 'rU') as csvfile:
        reader = csv.reader(csvfile, delimiter=options.delimiter)
        if options.nc:
            variable_names = None
        else:
            variable_names = reader.next()[(1 - options.nr):]
        sample_names = []
        data = []
        for row in reader:
            if options.nr:
                sample_names = None
            else:
                sample_names.append(row[0])
            data.append(row[(1 - options.nr):])

    try:
        X = -np.array(data, dtype=float)  # Data matrix in numpy format
    except:
        print "Incorrect data format.\nCheck that you've correctly specified options " \
              "such as continuous or not, \nand if there is a header row or column.\n" \
              "Run 'python gaussianize.py -h' option for help with options."
        traceback.print_exc(file=sys.stdout)
        sys.exit()

    out = Lambert()
    y = out.fit_transform(X)
    with open(options.output, 'w') as csvfile:
        writer = csv.writer(csvfile, delimiter=options.delimiter)
        if not options.nc:
            writer.writerow([""] * (1 - options.nr) + variable_names)
        for i, row in enumerate(y):
            if not options.nr:
                writer.writerow([sample_names[i]] + list(row))
            else:
                writer.writerow(row)

    if options.q:
        print 'Making qq plots'
        prefix = options.output.split('.')[0]
        if not os.path.exists(prefix+'_q'):
            os.makedirs(prefix+'_q')
        out.qqplot(X, prefix=prefix + '_q/q')<commit_msg>Fixed accidental sign error.<commit_after>"""
Transform samples of heavy-tailed data so that it is approximately normally distributed

Based on the paper:
The Lambert Way to Gaussianize heavy tailed data with the inverse of Tukey's h transformation as a special case
Based on algorithm in Appendix C
Author generously provides code in R:
https://cran.r-project.org/web/packages/LambertW/

This code written by Greg Ver Steeg, 2015.
"""

import numpy as np
from scipy.special import lambertw
from scipy.stats import kurtosis
from scipy.optimize import fmin  # TODO: Explore efficacy of other opt. methods

np.seterr(all='warn')

class Lambert(object):
    """
    Gaussianize heavy-tailed data using Lambert's W.

    Conventions
    ----------
    This class is a wrapper that follows sklearn naming/style (e.g. fit(X) to train).
    In this code, x is the input, y is the output. But in the functions outside the class, I follow
    Georg's convention that Y is the input and X is the output (Gaussianized) data.

    Parameters
    ----------
    tol : float, default = 1e-4

    max_iter : int, default = 200
        Maximum number of iterations to search for correct parameters of Lambert transform.

    Attributes
    ----------
    taus : list of tuples
        For each variable, we have a tuple consisting of (mu, sigma, delta), corresponding to the parameters of the
        appropriate Lambert transform. Eq. 6 and 8 in the paper below.

    References
    ----------
    [1] Georg Goerg. The Lambert Way to Gaussianize heavy tailed data with
                        the inverse of Tukey's h transformation as a special case
    """

    def __init__(self, tol=1.22e-4, max_iter=100):
        self.tol = tol
        self.max_iter = max_iter
        self.taus = []  # Store tau for each transformed variable

    def fit(self, x):
        x = np.asarray(x)
        if len(x.shape) == 1:
            x = x[:, np.newaxis]
        elif len(x.shape) != 2:
            print "Data should be a 1-d list of samples to transform or a 2d array with samples as rows."
        for x_i in x.T:
            self.taus.append(igmm(x_i, tol=self.tol, max_iter=self.max_iter))

    def transform(self, x):
        x = np.asarray(x)
        if len(x.shape) == 1:
            x = x[:, np.newaxis]
        elif len(x.shape) != 2:
            print "Data should be a 1-d list of samples to transform or a 2d array with samples as rows."
        if x.shape[1] != len(self.taus):
            print "%d variables in test data, but %d variables were in training data." % (x.shape[1], len(self.taus))

        return np.array([w_t(x_i, tau_i) for x_i, tau_i in zip(x.T, self.taus)]).T

    def fit_transform(self, x):
        self.fit(x)
        return self.transform(x)

    def invert(self, y):
        return np.array([inverse(y_i, tau_i) for y_i, tau_i in zip(y.T, self.taus)]).T

    def qqplot(self, x, prefix='qq'):
        """Show qq plots compared to normal before and after the transform."""
        import pylab
        from scipy.stats import probplot
        y = self.transform(x)

        for i, (x_i, y_i) in enumerate(zip(x.T, y.T)):
            probplot(x_i, dist="norm", plot=pylab)
            pylab.savefig(prefix + '_%d_before.png' % i)
            pylab.clf()

            probplot(y_i, dist="norm", plot=pylab)
            pylab.savefig(prefix + '_%d_after.png' % i)
            pylab.clf()

def w_d(z, delta):
    # Eq. 9
    if delta < 1e-6:
        return z
    return np.sign(z) * np.sqrt(np.real(lambertw(delta * z ** 2)) / delta)

def w_t(y, tau):
    # Eq. 8
    return tau[0] + tau[1] * w_d((y - tau[0]) / tau[1], tau[2])

def inverse(x, tau):
    # Eq. 6
    u = (x - tau[0]) / tau[1]
    return tau[0] + tau[1] * (u * np.exp(u * u * (tau[2] * 0.5)))

def igmm(y, tol=1.22e-4, max_iter=100):
    # Infer mu, sigma, delta using IGMM in Alg.2, Appendix C
    delta0 = delta_init(y)
    tau1 = (np.median(y), np.std(y) * (1. - 2. * delta0) ** 0.75, delta0)
    for k in range(max_iter):
        tau0 = tau1
        z = (y - tau1[0]) / tau1[1]
        delta1 = delta_gmm(z)
        x = tau0[0] + tau1[1] * w_d(z, delta1)
        mu1, sigma1 = np.mean(x), np.std(x)
        tau1 = (mu1, sigma1, delta1)

        if np.linalg.norm(np.array(tau1) - np.array(tau0)) < tol:
            break
        else:
            if k == max_iter - 1:
                print "Warning: No convergence after %d iterations. Increase max_iter." % max_iter
    return tau1

def delta_gmm(z):
    # Alg. 1, Appendix C
    delta0 = delta_init(z)

    def func(q):
        u = w_d(z, np.exp(q))
        if not np.all(np.isfinite(u)):
            return 0.
        else:
            k = kurtosis(u, fisher=True, bias=False)**2
            if not np.isfinite(k) or k > 1e10:
                return 1e10
            else:
                return k

    res = fmin(func, np.log(delta0), disp=0)
    return np.around(np.exp(res[-1]), 6)

def delta_init(z):
    gamma = kurtosis(z, fisher=False, bias=False)
    with np.errstate(all='ignore'):
        delta0 = np.clip(1. / 66 * (np.sqrt(66 * gamma - 162.) - 6.), 0.01, 0.25)
    if not np.isfinite(delta0):
        delta0 = 0.01
    return delta0

if __name__ == '__main__':
    # Command line interface
    # Sample commands:
    # python gaussianize.py test_data.csv
    import csv
    import sys, os
    import traceback
    from optparse import OptionParser, OptionGroup

    parser = OptionParser(usage="usage: %prog [options] data_file.csv \n"
                                "It is assumed that the first row and first column of the data CSV file are labels.\n"
                                "Use options to indicate otherwise.")
    group = OptionGroup(parser, "Input Data Format Options")
    group.add_option("-c", "--no_column_names",
                     action="store_true", dest="nc", default=False,
                     help="We assume the top row is variable names for each column. "
                          "This flag says that data starts on the first row and gives a "
                          "default numbering scheme to the variables (1,2,3...).")
    group.add_option("-r", "--no_row_names",
                     action="store_true", dest="nr", default=False,
                     help="We assume the first column is a label or index for each sample. "
                          "This flag says that data starts on the first column.")
    group.add_option("-d", "--delimiter",
                     action="store", dest="delimiter", type="string", default=",",
                     help="Separator between entries in the data, default is ','.")
    parser.add_option_group(group)

    group = OptionGroup(parser, "Output Options")
    group.add_option("-o", "--output",
                     action="store", dest="output", type="string", default="gaussian_output.csv",
                     help="Where to store gaussianized data.")
    group.add_option("-q", "--qqplots",
                     action="store_true", dest="q", default=False,
                     help="Produce qq plots for each variable before and after transform.")
    parser.add_option_group(group)

    (options, args) = parser.parse_args()
    if not len(args) == 1:
        print "Run with '-h' option for usage help."
        sys.exit()

    #Load data from csv file
    filename = args[0]
    with open(filename, 'rU') as csvfile:
        reader = csv.reader(csvfile, delimiter=options.delimiter)
        if options.nc:
            variable_names = None
        else:
            variable_names = reader.next()[(1 - options.nr):]
        sample_names = []
        data = []
        for row in reader:
            if options.nr:
                sample_names = None
            else:
                sample_names.append(row[0])
            data.append(row[(1 - options.nr):])

    try:
        X = np.array(data, dtype=float)  # Data matrix in numpy format
    except:
        print "Incorrect data format.\nCheck that you've correctly specified options " \
              "such as continuous or not, \nand if there is a header row or column.\n" \
              "Run 'python gaussianize.py -h' option for help with options."
        traceback.print_exc(file=sys.stdout)
        sys.exit()

    out = Lambert()
    y = out.fit_transform(X)
    with open(options.output, 'w') as csvfile:
        writer = csv.writer(csvfile, delimiter=options.delimiter)
        if not options.nc:
            writer.writerow([""] * (1 - options.nr) + variable_names)
        for i, row in enumerate(y):
            if not options.nr:
                writer.writerow([sample_names[i]] + list(row))
            else:
                writer.writerow(row)

    if options.q:
        print 'Making qq plots'
        prefix = options.output.split('.')[0]
        if not os.path.exists(prefix+'_q'):
            os.makedirs(prefix+'_q')
        out.qqplot(X, prefix=prefix + '_q/q')<|endoftext|>


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

'''
This module performs list of transformations on list of entries, and provides a
consistent input vs output interface for transformations on db entries.
'''

from __future__ import division

__author__ = "Shyue Ping Ong"
__copyright__ = "Copyright 2012, The Materials Project"
__version__ = "1.0"
__maintainer__ = "Shyue Ping Ong"
__email__ = "shyue@mit.edu"
__date__ = "Mar 2, 2012"

import os
import re
import datetime
import json

from pymatgen.core.structure import Structure
from pymatgen.io.cifio import CifParser
from pymatgen.io.vaspio import Poscar
from pymatgen.transformations.standard_transformations import transformation_from_dict
from collections import Iterable

class TransformedStructure(object):
    """
    Container object for new structures that include history of transformations.

    Each transformed structure is made up of a sequence of structures with associated
    transformation history.
    """

    def __init__(self, structure, transformations, history = None):
        """
        Standard constructor for a TransformedStructure

        Args:
            structure:
                input structure
            transformations:
                sequence of transformations to be applied to the input structure.
            history:
                optional history for the input structure, which provides a way
                to track structures having undergone multiple series of
                transformations.
        """
        history = [] if history == None else history
        self._source = {}
        self._structures = []
        self._transformations = []
        self._redo_trans = []
        if len(history) > 0:
            self._source = history[0]
            for i in xrange(1, len(history)):
                self._structures.append(Structure.from_dict(history[i]['input_structure']))
                self._transformations.append(transformation_from_dict(history[i]))

        self._structures.append(structure)
        for t in transformations:
            self.append_transformation(t)

    def undo_last_transformation(self):
        """
        Undo the last transformation in the TransformedStructure.

        Raises:
            IndexError if already at the oldest change.
        """
        if len(self._transformations) == 0:
            raise IndexError("Can't undo. Already at oldest change.")
        self._structures.pop()
        self._redo_trans.append(self._transformations.pop())

    def redo_next_transformation(self):
        """
        Redo the last undone transformation in the TransformedStructure.

        Raises:
            IndexError if already at the latest change.
        """
        if len(self._redo_trans) == 0:
            raise IndexError("Can't undo. Already at latest change.")
        t = self._redo_trans.pop()
        self.append_transformation(t, False)

    def __getitem__(self, index):
        return (self._structures[index], self._transformations[0:index])

    def __len__(self):
        return len(self._structures)

    def append_transformation(self, transformation, clear_redo = True):
        """
        Appends a transformation to the TransformedStructure.

        Arguments:
            transformation:
                Transformation to append
            clear_redo:
                Boolean indicating whether to clear the redo list. By default,
                this is True, meaning any appends clears the history of undoing.
                However, when using append_transformation to do a redo, the redo
                list should not be cleared to allow multiple redos.
        """
        new_s = transformation.apply_transformation(self._structures[-1])
        if isinstance(new_s, Iterable):
            new_s = new_s[0]
        self._structures.append(new_s)
        self._transformations.append(transformation)
        if clear_redo:
            self._redo_trans = []

    def extend_transformations(self, transformations):
        """
        Extends a sequence of transformations to the TransformedStructure.

        Arguments:
            transformations:
                Sequence of Transformations
        """
        for t in transformations:
            self.append_transformation(t)

    def get_vasp_input(self, vasp_input_set, generate_potcar = True):
        """
        Returns VASP input as a dict of vaspio objects.

        Args:
            vasp_input_set:
                pymatgen.io.vaspio_set.VaspInputSet like object that creates
                vasp input files from structures
            generate_potcar:
                Set to False to generate a POTCAR.spec file instead of a POTCAR,
                which contains the POTCAR labels but not the actual POTCAR. Defaults
                to True.
        """
        d = vasp_input_set.get_all_vasp_input(self._structures[-1], generate_potcar)
        d['transformations.json'] = json.dumps(self.to_dict)
        return d

    def write_vasp_input(self, vasp_input_set, output_dir, create_directory = True):
        """
        Writes VASP input to an output_dir.

        Args:
            vasp_input_set:
                pymatgen.io.vaspio_set.VaspInputSet like object that creates vasp input files from structures
            output_dir:
                Directory to output files
            create_directory:
                Create the directory if not present. Defaults to True.
        """
        vasp_input_set.write_input(self._structures[-1], output_dir, make_dir_if_not_present = create_directory)
        with open(os.path.join(output_dir, 'transformations.json'), 'w') as fp:
            json.dump(self.to_dict, fp)

    def __str__(self):
        output = ["Current structure"]
        output.append("------------")
        output.append(str(self._structures[-1]))
        output.append("\nSource")
        output.append("------------")
        output.append(str(self._source))
        output.append("\nTransformation history")
        output.append("------------")
        for t in self._transformations:
            output.append(str(t.to_dict))
        return "\n".join(output)

    @property
    def structures(self):
        """
        Returns a copy of all structures in the TransformedStructure. A structure
        is stored after every single transformation.
        """
        return [s for s in self._structures]

    @property
    def transformations(self):
        """
        Returns a copy of all transformations in the TransformedStructure.
        """
        return [t for t in self._transformations]

    @property
    def final_structure(self):
        """
        Returns the final structure in the TransformedStructure.
        """
        return self._structures[-1]

    @staticmethod
    def from_dict(d):
        """
        Creates a TransformedStructure from a dict.
        """
        s = Structure.from_dict(d)
        return TransformedStructure(s, [], d['history'])

    @property
    def to_dict(self):
        """
        Returns a dict representation of the TransformedStructure.
        """
        d = self._structures[-1].to_dict
        history = [self._source]
        for i, t in enumerate(self._transformations):
            tdict = t.to_dict
            tdict['input_structure'] = self._structures[i].to_dict
            history.append(tdict)
        d['history'] = history
        d['version'] = __version__
        return d

class TransformedStructureTree(object):
    """
    Container object for groups of structures that include history of transformations. This is used for generating
    multiple output structures for a single or multiple input structures through transformations that have multiple
    outputs

    Each structure is stored as a transformed structure. Transformations that produce multiple output structures
    increase the number of structures in the tree. A pruning method removes duplicate structures from the tree.
    Care must be used to ensure that the tree doesn't get too big
    """

    def __init__(self, structures_list, transformations, histories = None):
        """
        Standard constructor for a TransformedStructureTree

        Args:
            structures_list:
                input structures (must be a list)
            transformations:
                sequence of transformations to be applied to the input structures.
            histories:
                optional history list for the input structures, which provides a way
                to track structures having undergone multiple series of
                transformations. Must be the same length as structures_list
        """
        self._transformed_structures = []

        if histories:
            if not len(structures_list) == len(histories):
                raise ValueError("histories must have same length as structures_list")
            for i in range(len(structures_list)):
                self._transformed_structures.append(TransformedStructure(structures_list[i], transformations, histories[i]))
        else:
            for i in range(len(structures_list)):
                self._transformed_structures.append(TransformedStructure(structures_list[i], transformations, None))

    def undo_last_transformation(self):
        """
        Undo the last transformation in the TransformedStructure.

        Raises:
            IndexError if already at the oldest change.
        """
        for x in self._transformed_structures:
            x.undo_last_transformation()

    def redo_next_transformation(self):
        """
        Redo the last undone transformation in the TransformedStructure.

        Raises:
            IndexError if already at the latest change.
        """
        for x in self._transformed_structures:
            x.redo_next_transformation()

    def __getitem__(self, index):
        return [x.__getitem__(self, index) for x in self._transformed_structures]

    def __len__(self):
        return len(self._transformed_structures)

    def append_transformation(self, transformation, clear_redo = True):
        """
        Appends a transformation to the TransformedStructure.

        Arguments:
            transformation:
                Transformation to append
            clear_redo:
                Boolean indicating whether to clear the redo list. By default,
                this is True, meaning any appends clears the history of undoing.
                However, when using append_transformation to do a redo, the redo
                list should not be cleared to allow multiple redos.
        """
        new_transformed_structures = []
        structures_to_delete = []
        for x in self._transformed_structures:
            new_structs = transformation.apply_transformation(x.final_structure)
            if isinstance(new_structs, Iterable):
                structures_to_delete.append(x)
                for new_struct in new_structs:
                    new_tstruct = x
                    new_tstruct._structures.append(new_struct)
                    new_tstruct._transformations.append(transformation)
                    new_transformed_structures.append(new_tstruct)
            else:
                x._structures.append(new_structs)
                x._transformations.append(transformation)
            if clear_redo:
                x._redo_trans = []
        for x in structures_to_delete:
            self._transformed_structures.remove(x)
        self._transformed_structures.extend(new_transformed_structures)

    def extend_transformations(self, transformations):
        """
        Extends a sequence of transformations to the TransformedStructure.

        Arguments:
            transformations:
                Sequence of Transformations
        """
        for t in transformations:
            self.append_transformation(t)

    def get_vasp_input(self, vasp_input_set, generate_potcar = True):
        """
        Returns VASP input as a dict of vaspio objects.

        Args:
            vasp_input_set:
                pymatgen.io.vaspio_set.VaspInputSet like object that creates
                vasp input files from structures
            generate_potcar:
                Set to False to generate a POTCAR.spec file instead of a POTCAR,
                which contains the POTCAR labels but not the actual POTCAR. Defaults
                to True.
        """
        return [x.__getitem__(self, vasp_input_set, generate_potcar) for x in self._transformed_structures]

    def write_vasp_input(self, vasp_input_set, output_dir, create_directory = True):
        """
        Writes VASP input to an output_dir.

        Args:
            vasp_input_set:
                pymatgen.io.vaspio_set.VaspInputSet like object that creates vasp input files from structures
            output_dir:
                Directory to output files
            create_directory:
                Create the directory if not present. Defaults to True.
        """
        vasp_input_set.write_input(self._structures[-1], output_dir, make_dir_if_not_present = create_directory)
        with open(os.path.join(output_dir, 'transformations.json'), 'w') as fp:
            json.dump(self.to_dict, fp)

    def __str__(self):
        output = ["Current structures"]
        output.append("------------")
        for x in self._transformed_structures:
            output.append(str(x._structures[-1]))
        return "\n".join(output)

#TODO: Code these following methods
'''
    @property
    def structures(self):
        """
        Returns a copy of all structures in the TransformedStructure. A structure
        is stored after every single transformation.
        """
        return [s for s in self._structures]

    @property
    def transformations(self):
        """
        Returns a copy of all transformations in the TransformedStructure.
        """
        return [t for t in self._transformations]

    @property
    def final_structure(self):
        """
        Returns the final structure in the TransformedStructure.
        """
        return self._structures[-1]

    @staticmethod
    def from_dict(d):
        """
        Creates a TransformedStructure from a dict.
        """
        s = Structure.from_dict(d)
        return TransformedStructure(s, [], d['history'])

    @property
    def to_dict(self):
        """
        Returns a dict representation of the TransformedStructure.
        """
        d = self._structures[-1].to_dict
        history = [self._source]
        for i, t in enumerate(self._transformations):
            tdict = t.to_dict
            tdict['input_structure'] = self._structures[i].to_dict
            history.append(tdict)
        d['history'] = history
        d['version'] = __version__
        return d'''

class CifTransformedStructure(TransformedStructure):
    """
    Extended version of TransformedStructure to generates new materials from cifs.
    """

    def __init__(self, cif_string, transformations, primitive = True):
        """
        Args:
            cif_string:
                Input cif string. Should contain only one structure. For cifs
                containing multiple structures, please use CifTransmuter.
            transformations:
                Sequence of transformations to be applied to the input structure.
            primitive:
                Option to set if the primitive cell should be extracted. Defaults
                to True. However, there are certain instances where you might want
                to use a non-primitive cell, e.g., if you are trying to generate
                all possible orderings of partial removals or order a disordered
                structure.
        """
        parser = CifParser.from_string(cif_string)
        raw_string = re.sub("'", "\"", cif_string)
        cif_dict = parser.to_dict
        cif_keys = cif_dict.keys()
        s = parser.get_structures(primitive)[0]
        partial_cif = cif_dict[cif_keys[0]]
        if '_database_code_ICSD' in partial_cif:
            source = partial_cif['_database_code_ICSD'] + "-ICSD"
        else:
            source = 'uploaded cif'
        source_info = {'source':source, 'datetime':str(datetime.datetime.utcnow()), 'original_file':raw_string, 'cif_data':cif_dict[cif_keys[0]]}
        super(CifTransformedStructure, self).__init__(s, transformations, [source_info])

class PoscarTransformedStructure(TransformedStructure):
    """
    Generates a transformed structure from Poscar.
    """

    def __init__(self, poscar_string, transformations):
        """
        Args:
            poscar_string:
                Input POSCAR string.
            transformations:
                Sequence of transformations to be applied to the input structure.
        """
        p = Poscar.from_string(poscar_string)
        if not p.true_names:
            raise ValueError("Transformation can be craeted only from POSCAR strings with proper VASP5 element symbols.")
        raw_string = re.sub("'", "\"", poscar_string)
        s = p.struct
        source_info = {'source': "uploaded POSCAR", 'datetime':str(datetime.datetime.utcnow()), 'original_file':raw_string}
        super(PoscarTransformedStructure, self).__init__(s, transformations, [source_info])

<commit_msg>Implemented basic functionality for TransformedStructureCollection<commit_after>#!/usr/bin/env python

'''
This module performs list of transformations on list of entries, and provides a
consistent input vs output interface for transformations on db entries.
'''

from __future__ import division

__author__ = "Shyue Ping Ong"
__copyright__ = "Copyright 2012, The Materials Project"
__version__ = "1.0"
__maintainer__ = "Shyue Ping Ong"
__email__ = "shyue@mit.edu"
__date__ = "Mar 2, 2012"

import os
import re
import datetime
import json
import collections

from pymatgen.core.structure import Structure
from pymatgen.io.cifio import CifParser
from pymatgen.io.vaspio import Poscar
from pymatgen.transformations.standard_transformations import transformation_from_dict

class TransformedStructure(object):
    """
    Container object for new structures that include history of transformations.

    Each transformed structure is made up of a sequence of structures with associated
    transformation history.
    """

    def __init__(self, structure, transformations, history = None):
        """
        Standard constructor for a TransformedStructure

        Args:
            structure:
                input structure
            transformations:
                sequence of transformations to be applied to the input structure.
            history:
                optional history for the input structure, which provides a way
                to track structures having undergone multiple series of
                transformations.
        """
        history = [] if history == None else history
        self._source = {}
        self._structures = []
        self._transformations = []
        self._redo_trans = []
        if len(history) > 0:
            self._source = history[0]
            for i in xrange(1, len(history)):
                self._structures.append(Structure.from_dict(history[i]['input_structure']))
                self._transformations.append(transformation_from_dict(history[i]))

        self._structures.append(structure)
        for t in transformations:
            self.append_transformation(t)

    def undo_last_transformation(self):
        """
        Undo the last transformation in the TransformedStructure.

        Raises:
            IndexError if already at the oldest change.
        """
        if len(self._transformations) == 0:
            raise IndexError("Can't undo. Already at oldest change.")
        self._structures.pop()
        self._redo_trans.append(self._transformations.pop())

    def redo_next_transformation(self):
        """
        Redo the last undone transformation in the TransformedStructure.

        Raises:
            IndexError if already at the latest change.
        """
        if len(self._redo_trans) == 0:
            raise IndexError("Can't undo. Already at latest change.")
        t = self._redo_trans.pop()
        self.append_transformation(t, False)

    def __getitem__(self, index):
        return (self._structures[index], self._transformations[0:index])

    def __len__(self):
        return len(self._structures)

    def _alternative_transformed_structures(self, transformation, structures):
        if len(structures)>0:
            self._transformations.append(transformation)
            history = self.history
            self._transformations.pop()
        for x in structures:
            yield TransformedStructure(x, [], history)

    def append_transformation(self, transformation, clear_redo = True, return_alternatives = False):
        """
        Appends a transformation to the TransformedStructure.

        Arguments:
            transformation:
                Transformation to append
            clear_redo:
                Boolean indicating whether to clear the redo list. By default,
                this is True, meaning any appends clears the history of undoing.
                However, when using append_transformation to do a redo, the redo
                list should not be cleared to allow multiple redos.
        """

        new_s = transformation.apply_transformation(self._structures[-1])
        alternative_s = None

        if not isinstance(new_s, Structure):
            alternative_s = self._alternative_transformed_structures(transformation, new_s[1:])
            new_s = new_s[0]

        self._transformations.append(transformation)
        self._structures.append(new_s)

        if clear_redo:
            self._redo_trans = []

        if return_alternatives:
            return alternative_s

    def extend_transformations(self, transformations):
        """
        Extends a sequence of transformations to the TransformedStructure.

        Arguments:
            transformations:
                Sequence of Transformations
        """
        for t in transformations:
            self.append_transformation(t)

    def get_vasp_input(self, vasp_input_set, generate_potcar = True):
        """
        Returns VASP input as a dict of vaspio objects.

        Args:
            vasp_input_set:
                pymatgen.io.vaspio_set.VaspInputSet like object that creates
                vasp input files from structures
            generate_potcar:
                Set to False to generate a POTCAR.spec file instead of a POTCAR,
                which contains the POTCAR labels but not the actual POTCAR. Defaults
                to True.
        """
        d = vasp_input_set.get_all_vasp_input(self._structures[-1], generate_potcar)
        d['transformations.json'] = json.dumps(self.to_dict)
        return d

    def write_vasp_input(self, vasp_input_set, output_dir, create_directory = True):
        """
        Writes VASP input to an output_dir.

        Args:
            vasp_input_set:
                pymatgen.io.vaspio_set.VaspInputSet like object that creates vasp input files from structures
            output_dir:
                Directory to output files
            create_directory:
                Create the directory if not present. Defaults to True.
        """
        vasp_input_set.write_input(self._structures[-1], output_dir, make_dir_if_not_present = create_directory)
        with open(os.path.join(output_dir, 'transformations.json'), 'w') as fp:
            json.dump(self.to_dict, fp)

    def __str__(self):
        output = ["Current structure"]
        output.append("------------")
        output.append(str(self._structures[-1]))
        output.append("\nSource")
        output.append("------------")
        output.append(str(self._source))
        output.append("\nTransformation history")
        output.append("------------")
        for t in self._transformations:
            output.append(str(t.to_dict))
        return "\n".join(output)

    @property
    def structures(self):
        """
        Returns a copy of all structures in the TransformedStructure. A structure
        is stored after every single transformation.
        """
        return [s for s in self._structures]

    @property
    def transformations(self):
        """
        Returns a copy of all transformations in the TransformedStructure.
        """
        return [t for t in self._transformations]

    @property
    def final_structure(self):
        """
        Returns the final structure in the TransformedStructure.
        """
        return self._structures[-1]

    @staticmethod
    def from_dict(d):
        """
        Creates a TransformedStructure from a dict.
        """
        s = Structure.from_dict(d)
        return TransformedStructure(s, [], d['history'])

    @property
    def history(self):
        history = [self._source]
        for i, t in enumerate(self._transformations):
            tdict = t.to_dict
            tdict['input_structure'] = self._structures[i].to_dict
            history.append(tdict)
        return history

    @property
    def to_dict(self):
        """
        Returns a dict representation of the TransformedStructure.
        """
        d = self._structures[-1].to_dict
        d['history'] = self.history
        d['version'] = __version__
        return d

class TransformedStructureCollection(object):

    def __init__(self, transformed_structures, transformations, extend_collection = False):
        self._extend_collection = extend_collection
        self._transformed_structures = transformed_structures
        for trans in transformations:
            self.append_transformation(trans)

    def __getitem__(self, index):
        return self._transformed_structures[index]

    def undo_last_transformation(self):
        """
        Undo the last transformation in the TransformedStructure.

        Raises:
            IndexError if already at the oldest change.
        """
        for x in self._transformed_structures:
            x.undo_last_transformation()

    def redo_next_transformation(self):
        """
        Redo the last undone transformation in the TransformedStructure.

        Raises:
            IndexError if already at the latest change.
        """
        for x in self._transformed_structures:
            x.redo_next_transformation()

    def __len__(self):
        return len(self._transformed_structures)

    def append_transformation(self, transformation, clear_redo = True):
        """
        TODO: clean this up a lot

        Appends a transformation to the TransformedStructure.

        Arguments:
            transformation:
                Transformation to append
            clear_redo:
                Boolean indicating whether to clear the redo list. By default,
                this is True, meaning any appends clears the history of undoing.
                However, when using append_transformation to do a redo, the redo
                list should not be cleared to allow multiple redos.
        """
        new_structures = []

        for x in self._transformed_structures:
            new = x.append_transformation(transformation, clear_redo, return_alternatives = self._extend_collection)
            if new:
                new_structures.extend(new)
        self._transformed_structures.extend(new_structures)

    def extend_transformations(self, transformations):
        """
        Extends a sequence of transformations to the TransformedStructure.

        Arguments:
            transformations:
                Sequence of Transformations
        """
        for t in transformations:
            self.append_transformation(t)

    def batch_write_vasp_input(self, vasp_input_set, output_dir, create_directory = True):
        """
        Batch write vasp input for a sequence of transformed structures to output_dir,
        following the format output_dir/{formula}_{number}.

        Args:
            transformed_structures:
                Sequence of TransformedStructures.
            vasp_input_set:
                pymatgen.io.vaspio_set.VaspInputSet like object that creates
                vasp input files from structures
            output_dir:
                Directory to output files
            create_directory:
                Create the directory if not present. Defaults to True.
        """
        dnames_count = collections.defaultdict(int)
        for s in self._transformed_structures:
            formula = re.sub("\s+", "", s.final_structure.formula)
            dirname = os.path.join(output_dir, '{}_{}'.format(formula, dnames_count[formula] + 1))
            s.write_vasp_input(vasp_input_set, dirname, create_directory = True)
            dnames_count[formula] += 1

    def __str__(self):
        output = ["Current structures"]
        output.append("------------")
        for x in self._transformed_structures:
            output.append(str(x._structures[-1]))
        return "\n".join(output)

    def remove_duplicates(self):
        '''
        TODO: write this method
        '''
        pass

    @staticmethod
    def from_structures(structures_list, transformations = [], histories = None, extend_collection = False):
        """
        Constructor for a TransformedStructureCollection from list of structures

        Args:
            structures_list:
                input structures (must be a list)
            transformations:
                sequence of transformations to be applied to the input structures.
            histories:
                optional history list for the input structures, which provides a way
                to track structures having undergone multiple series of
                transformations. Must be the same length as structures_list
        """
        transformed_structures = []

        if histories:
            if not len(structures_list) == len(histories):
                raise ValueError("histories must have same length as structures_list")
            for i in range(len(structures_list)):
                transformed_structures.append(TransformedStructure(structures_list[i], None, histories[i]))
        else:
            for i in range(len(structures_list)):
                transformed_structures.append(TransformedStructure(structures_list[i], None, None))
        return TransformedStructureCollection(transformed_structures, transformations, extend_collection)

    @staticmethod
    def from_cifs(cif_filenames, transformations = [], primitive = True, extend_collection = False):
        '''
        Args:
            cif_filenames:
                List of strings of the cif files
        '''

        def cif_transformed_parser(cif_string, primitive = True):
            """
            Args:
                cif_string:
                    Input cif string. Should contain only one structure. For cifs
                    containing multiple structures, please use CifTransmuter.
                transformations:
                    Sequence of transformations to be applied to the input structure.
                primitive:
                    Option to set if the primitive cell should be extracted. Defaults
                    to True. However, there are certain instances where you might want
                    to use a non-primitive cell, e.g., if you are trying to generate
                    all possible orderings of partial removals or order a disordered
                    structure.
            """
            parser = CifParser.from_string(cif_string)
            raw_string = re.sub("'", "\"", cif_string)
            cif_dict = parser.to_dict
            cif_keys = cif_dict.keys()
            s = parser.get_structures(primitive)[0]
            partial_cif = cif_dict[cif_keys[0]]
            if '_database_code_ICSD' in partial_cif:
                source = partial_cif['_database_code_ICSD'] + "-ICSD"
            else:
                source = 'uploaded cif'
            source_info = {'source':source, 'datetime':str(datetime.datetime.utcnow()), 'original_file':raw_string, 'cif_data':cif_dict[cif_keys[0]]}
            return TransformedStructure(s, [], [source_info])

        transformed_structures = []
        for filename in cif_filenames:
            with open(filename, "r") as f:
                structure_data = []
                read_data = False
                for line in f:
                    if re.match("^\s*data", line):
                        structure_data.append([])
                        read_data = True
                    if read_data:
                        structure_data[-1].append(line)
                transformed_structures.extend([cif_transformed_parser("".join(data), primitive) for data in structure_data])
        return TransformedStructureCollection(transformed_structures, transformations, extend_collection)

    @staticmethod
    def from_poscars(poscar_filenames, transformations = [], extend_collection = False):
        def poscar_transformed_parser(poscar_string):
            """
            Args:
                poscar_string:
                    Input POSCAR string.
            """
            p = Poscar.from_string(poscar_string)
            if not p.true_names:
                raise ValueError("Transformation can be craeted only from POSCAR strings with proper VASP5 element symbols.")
            raw_string = re.sub("'", "\"", poscar_string)
            s = p.struct
            source_info = {'source': "uploaded POSCAR", 'datetime':str(datetime.datetime.utcnow()), 'original_file':raw_string}
            return TransformedStructure(s, [], [source_info])

        transformed_structures = []
        for filename in poscar_filenames:
            with open(filename, "r") as f:
                transformed_structures.append(poscar_transformed_parser(f.read()))
        return TransformedStructureCollection(transformed_structures, transformations, extend_collection)

<|endoftext|>


<commit_before>import unittest
import datetime as dt
import sqlite3
import sjcx_payments.payments as payments
import os

class Payments(unittest.TestCase):

    def test_add_balances(self):
        conn = sqlite3.connect('data/test_rewards.db')
        cursor = conn.cursor()
        test_date = dt.datetime(2015, 11, 1, 0, 0, 0)
        payments.add_balances(conn, cursor, test_date)
        cursor.execute('''SELECT MAX(balance) FROM rewards WHERE payment_date=?''',
                       (str(test_date),))
        max_balance =cursor.fetchone()[0]
        self.assertTrue(max_balance > 0)
        conn.close()

    def test_assign_points(self):
        conn = sqlite3.connect('data/test_rewards.db')
        cursor = conn.cursor()
        test_date = dt.datetime(2015, 11, 1, 0, 0, 0)
        payments.assign_points(conn, cursor, test_date)
        cursor.execute('''SELECT MAX(points) FROM rewards WHERE payment_date=?''',
                       (str(test_date),))
        points = cursor.fetchone()[0]
        self.assertTrue(points > 0)
        conn.close()

    def test_distribute_sjcx(self):
        conn = sqlite3.connect('data/test_rewards.db')
        cursor = conn.cursor()
        test_date = dt.datetime(2015, 11, 1, 0, 0, 0)
        payments.distribute_sjcx(conn, cursor, test_date)
        cursor.execute('''SELECT SUM(sjcx_reward) FROM rewards WHERE
                          payment_date = ?''', (str(test_date),))
        total = cursor.fetchone()[0]
        self.assertAlmostEqual(total, 100000)
        conn.close()

    def test_update_total_rewards(self):
        conn = sqlite3.connect('data/test_rewards.db')
        cursor = conn.cursor()
        test_date = dt.datetime(2015, 11, 1, 0, 0, 0)
        payments.update_total_rewards(conn, cursor, test_date)
        cursor.execute('''SELECT MAX(total_usd_reward) FROM rewards WHERE
                          payment_date = ?''', (str(test_date),))
        usd_reward = cursor.fetchone()[0]
        self.assertTrue(usd_reward > 0)
        cursor.execute('''UPDATE rewards SET balance = 0, sjcx_reward = 0,
                          points = 0, usd_reward = 0, total_usd_reward = 0''')
        conn.commit()

    def test_make_csv(self):
        conn = sqlite3.connect('data/test_rewards.db')
        cursor = conn.cursor()
        test_date = dt.datetime(2015, 11, 1, 0, 0, 0)
        payments.make_points_csv(cursor, test_date)
        dir = os.path.abspath('sjcx_rewards_2015_11_1.csv')
        boolean = os.path.exists(dir)
        self.assertTrue(boolean)

    def test_create_whitelist(self):
        whitelist = payments.create_whitelist()
        self.assertTrue(len(whitelist) > 0)

    def test_uptime_logistic_function(self):
        uptime = 0.5
        value = payments.uptime_logistic_function(uptime)
        self.assertTrue(isinstance(value, float))

    def test_height_function(self):
        height = 0.5
        value = payments.height_function(height)
        self.assertTrue(isinstance(value, float))

if __name__ == '__main__':
    unittest.main()
<commit_msg>create test<commit_after>import unittest
import datetime as dt
import sqlite3
import sjcx_payments.payments as payments
import os

class Payments(unittest.TestCase):

    def test_add_balances(self):
        conn = sqlite3.connect('data/test_rewards.db')
        cursor = conn.cursor()
        test_date = dt.datetime(2015, 11, 1, 0, 0, 0)
        payments.add_balances(conn, cursor, test_date)
        cursor.execute('''SELECT MAX(balance) FROM rewards WHERE payment_date=?''',
                       (str(test_date),))
        max_balance =cursor.fetchone()[0]
        self.assertTrue(max_balance > 0)
        conn.close()

    def test_assign_points(self):
        conn = sqlite3.connect('data/test_rewards.db')
        cursor = conn.cursor()
        test_date = dt.datetime(2015, 11, 1, 0, 0, 0)
        payments.assign_points(conn, cursor, test_date)
        cursor.execute('''SELECT MAX(points) FROM rewards WHERE payment_date=?''',
                       (str(test_date),))
        points = cursor.fetchone()[0]
        self.assertTrue(points > 0)
        conn.close()

    def test_create_table(self):
        conn = sqlite3.connect('unittest.db')
        cursor = conn.cursor()
        payments.create_table(conn, cursor)
        dir = os.path.abspath('unittest.db')
        boolean = os.path.exists(dir)
        self.assertTrue(boolean)

    def test_distribute_sjcx(self):
        conn = sqlite3.connect('data/test_rewards.db')
        cursor = conn.cursor()
        test_date = dt.datetime(2015, 11, 1, 0, 0, 0)
        payments.distribute_sjcx(conn, cursor, test_date)
        cursor.execute('''SELECT SUM(sjcx_reward) FROM rewards WHERE
                          payment_date = ?''', (str(test_date),))
        total = cursor.fetchone()[0]
        self.assertAlmostEqual(total, 100000)
        conn.close()

    def test_update_total_rewards(self):
        conn = sqlite3.connect('data/test_rewards.db')
        cursor = conn.cursor()
        test_date = dt.datetime(2015, 11, 1, 0, 0, 0)
        payments.update_total_rewards(conn, cursor, test_date)
        cursor.execute('''SELECT MAX(total_usd_reward) FROM rewards WHERE
                          payment_date = ?''', (str(test_date),))
        usd_reward = cursor.fetchone()[0]
        self.assertTrue(usd_reward > 0)
        cursor.execute('''UPDATE rewards SET balance = 0, sjcx_reward = 0,
                          points = 0, usd_reward = 0, total_usd_reward = 0''')
        conn.commit()

    def test_make_csv(self):
        conn = sqlite3.connect('data/test_rewards.db')
        cursor = conn.cursor()
        test_date = dt.datetime(2015, 11, 1, 0, 0, 0)
        payments.make_points_csv(cursor, test_date)
        dir = os.path.abspath('sjcx_rewards_2015_11_1.csv')
        boolean = os.path.exists(dir)
        self.assertTrue(boolean)

    def test_create_whitelist(self):
        whitelist = payments.create_whitelist()
        self.assertTrue(len(whitelist) > 0)

    def test_uptime_logistic_function(self):
        uptime = 0.5
        value = payments.uptime_logistic_function(uptime)
        self.assertTrue(isinstance(value, float))

    def test_height_function(self):
        height = 0.5
        value = payments.height_function(height)
        self.assertTrue(isinstance(value, float))

if __name__ == '__main__':
    unittest.main()
<|endoftext|>


<commit_before># This file is part of QuTiP: Quantum Toolbox in Python.
#
#    Copyright (c) 2011 and later, Paul D. Nation and Robert J. Johansson.
#    All rights reserved.
#
#    Redistribution and use in source and binary forms, with or without
#    modification, are permitted provided that the following conditions are
#    met:
#
#    1. Redistributions of source code must retain the above copyright notice,
#       this list of conditions and the following disclaimer.
#
#    2. Redistributions in binary form must reproduce the above copyright
#       notice, this list of conditions and the following disclaimer in the
#       documentation and/or other materials provided with the distribution.
#
#    3. Neither the name of the QuTiP: Quantum Toolbox in Python nor the names
#       of its contributors may be used to endorse or promote products derived
#       from this software without specific prior written permission.
#
#    THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
#    "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
#    LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
#    PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
#    HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
#    SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
#    LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
#    DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
#    THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
#    (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
#    OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
###############################################################################

import numpy as np
from numpy.testing import assert_equal, run_module_suite
from scipy.sparse import isspmatrix_csr
from qutip import (commutator, position, momentum, create, destroy, displace,
                   jmat, num, phase, qdiags, qeye, identity, qutrit_ops,
                   squeeze, squeezing, zero_oper)

def test_commutator_type():
    "Operator CSR Type: commutator"
    op = commutator(position(5), momentum(5))
    assert_equal(isspmatrix_csr(op.data), True)

def test_create_type():
    "Operator CSR Type: create"
    op = create(5)
    assert_equal(isspmatrix_csr(op.data), True)

def test_destroy_type():
    "Operator CSR Type: destroy"
    op = destroy(5)
    assert_equal(isspmatrix_csr(op.data), True)

def test_displace_type():
    "Operator CSR Type: displace"
    op = displace(5, 0.1)
    assert_equal(isspmatrix_csr(op.data), True)

def test_jmat_type():
    "Operator CSR Type: spin ops"
    for k in ['x', 'y', 'z', '+', '-']:
        op = jmat(1/2, k)
        assert_equal(isspmatrix_csr(op.data), True)

def test_momentum_type():
    "Operator CSR Type: momentum"
    op = momentum(5)
    assert_equal(isspmatrix_csr(op.data), True)

def test_num_type():
    "Operator CSR Type: num"
    op = num(5)
    assert_equal(isspmatrix_csr(op.data), True)

def test_phase_type():
    "Operator CSR Type: phase"
    op = phase(5, np.pi)
    assert_equal(isspmatrix_csr(op.data), True)

def test_position_type():
    "Operator CSR Type: position"
    op = position(5)
    assert_equal(isspmatrix_csr(op.data), True)

def test_qdiags_type():
    "Operator CSR Type: qdiags"
    op = qdiags(np.sqrt(range(1, 4)), 1)
    assert_equal(isspmatrix_csr(op.data), True)

def test_qeye_type():
    "Operator CSR Type: qeye/identity"
    op = qeye(5)
    assert_equal(isspmatrix_csr(op.data), True)
    op = identity(5)
    assert_equal(isspmatrix_csr(op.data), True)

def test_qtrit_type():
    "Operator CSR Type: qutrit ops"
    ops = qutrit_ops()
    for k in ops:
        assert_equal(isspmatrix_csr(k.data), True)

def test_squeeze_type():
    "Operator CSR Type: squeeze"
    op = squeeze(5, 0.1j)
    assert_equal(isspmatrix_csr(op.data), True)

def test_squeezing_type():
    "Operator CSR Type: squeezing"
    op = squeezing(destroy(5), qeye(5), -0.1j)
    assert_equal(isspmatrix_csr(op.data), True)

def test_zero_type():
    "Operator CSR Type: zero_oper"
    op = zero_oper(5, [[5], [5]])
    assert_equal(isspmatrix_csr(op.data), True)

if __name__ == "__main__":
    run_module_suite()
<commit_msg>updated for change to qzero<commit_after># This file is part of QuTiP: Quantum Toolbox in Python.
#
#    Copyright (c) 2011 and later, Paul D. Nation and Robert J. Johansson.
#    All rights reserved.
#
#    Redistribution and use in source and binary forms, with or without
#    modification, are permitted provided that the following conditions are
#    met:
#
#    1. Redistributions of source code must retain the above copyright notice,
#       this list of conditions and the following disclaimer.
#
#    2. Redistributions in binary form must reproduce the above copyright
#       notice, this list of conditions and the following disclaimer in the
#       documentation and/or other materials provided with the distribution.
#
#    3. Neither the name of the QuTiP: Quantum Toolbox in Python nor the names
#       of its contributors may be used to endorse or promote products derived
#       from this software without specific prior written permission.
#
#    THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
#    "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
#    LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
#    PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
#    HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
#    SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
#    LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
#    DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
#    THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
#    (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
#    OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
###############################################################################

import numpy as np
from numpy.testing import assert_equal, run_module_suite
from scipy.sparse import isspmatrix_csr
from qutip import (commutator, position, momentum, create, destroy, displace,
                   jmat, num, phase, qdiags, qeye, identity, qutrit_ops,
                   squeeze, squeezing, qzero)

def test_commutator_type():
    "Operator CSR Type: commutator"
    op = commutator(position(5), momentum(5))
    assert_equal(isspmatrix_csr(op.data), True)

def test_create_type():
    "Operator CSR Type: create"
    op = create(5)
    assert_equal(isspmatrix_csr(op.data), True)

def test_destroy_type():
    "Operator CSR Type: destroy"
    op = destroy(5)
    assert_equal(isspmatrix_csr(op.data), True)

def test_displace_type():
    "Operator CSR Type: displace"
    op = displace(5, 0.1)
    assert_equal(isspmatrix_csr(op.data), True)

def test_jmat_type():
    "Operator CSR Type: spin ops"
    for k in ['x', 'y', 'z', '+', '-']:
        op = jmat(1/2, k)
        assert_equal(isspmatrix_csr(op.data), True)

def test_momentum_type():
    "Operator CSR Type: momentum"
    op = momentum(5)
    assert_equal(isspmatrix_csr(op.data), True)

def test_num_type():
    "Operator CSR Type: num"
    op = num(5)
    assert_equal(isspmatrix_csr(op.data), True)

def test_phase_type():
    "Operator CSR Type: phase"
    op = phase(5, np.pi)
    assert_equal(isspmatrix_csr(op.data), True)

def test_position_type():
    "Operator CSR Type: position"
    op = position(5)
    assert_equal(isspmatrix_csr(op.data), True)

def test_qdiags_type():
    "Operator CSR Type: qdiags"
    op = qdiags(np.sqrt(range(1, 4)), 1)
    assert_equal(isspmatrix_csr(op.data), True)

def test_qeye_type():
    "Operator CSR Type: qeye/identity"
    op = qeye(5)
    assert_equal(isspmatrix_csr(op.data), True)
    op = identity(5)
    assert_equal(isspmatrix_csr(op.data), True)

def test_qtrit_type():
    "Operator CSR Type: qutrit ops"
    ops = qutrit_ops()
    for k in ops:
        assert_equal(isspmatrix_csr(k.data), True)

def test_squeeze_type():
    "Operator CSR Type: squeeze"
    op = squeeze(5, 0.1j)
    assert_equal(isspmatrix_csr(op.data), True)

def test_squeezing_type():
    "Operator CSR Type: squeezing"
    op = squeezing(destroy(5), qeye(5), -0.1j)
    assert_equal(isspmatrix_csr(op.data), True)

def test_zero_type():
    "Operator CSR Type: zero_oper"
    op = qzero(5, [[5], [5]])
    assert_equal(isspmatrix_csr(op.data), True)

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


<commit_before><commit_msg>change time<commit_after><|endoftext|>


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

import RPi.GPIO as GPIO
import time
import random
import signal
import sys

RED_LIGHT_TIME = 15
YELLO_LIGHT_TIME = 1.8
GREEN_LIGTH_TIME = 15

RANDOM_TIME = 5

RED_PIN = 17
YELLOW_PIN = 27
GREEN_PIN = 22

def setup():
  GPIO.setmode(GPIO.BCM)
  GPIO.setwarnings(False)
  GPIO.setup([RED_PIN, YELLOW_PIN, GREEN_PIN], GPIO.OUT)
  GPIO.output([RED_PIN, YELLOW_PIN, GREEN_PIN], GPIO.LOW)

def loop():
  turn_on(GREEN_PIN, GREEN_LIGTH_TIME + random.randint(0, RANDOM_TIME))
  turn_on(YELLOW_PIN, YELLO_LIGHT_TIME)
  turn_on(RED_PIN, RED_LIGHT_TIME + random.randint(0, RANDOM_TIME))

def finished():
  GPIO.cleanup()

def turn_on(pin, duration):
  print 'turn on %d (%f)' % (pin, duration)
  GPIO.output(pin, GPIO.HIGH)
  time.sleep(duration)
  print 'turn off'
  GPIO.output(pin, GPIO.LOW)

def signal_handler(signal, frame):
  finished()
  sys.exit(0)

if __name__ == '__main__':
  signal.signal(signal.SIGTERM, signal_handler)
  signal.signal(signal.SIGINT, signal_handler)
  setup()
  while (True):
    loop()

<commit_msg>drafted full traffic light program (test pending)<commit_after>#!/usr/bin/env python

import RPi.GPIO as GPIO
import time
import random
import signal
import sys
import threading

RED_PIN = 17
YELLOW_PIN = 27
GREEN_PIN = 22

HUMAN_RED_PIN = 18
HUMAN_GREEN_PIN = 23

PRESS_LIGHT_PIN = 24
WAIT_LIGHT_PIN = 25

BUTTON_PIN = 4

BUTTON_DELAY_TIME = 5

pin_map = {
  'r': RED_PIN,
  'y': YELLOW_PIN,
  'g': GREEN_PIN,
  'R': HUMAN_RED_PIN,
  'G': HUMAN_GREEN_PIN,
}

states = [
  { 'name': 'all-green', 'on_time': 15, 'off_time': 0, 'repeat': 0, 'on_lights': 'gG', 'interruptable': False, 'rand_time': 3, 'enable_button': False },
  { 'name': 'blink-green', 'on_time': 0.5, 'off_time': 0.5, 'repeat': 5, 'on_lights': 'gG', 'off_lights': 'G', 'interruptable': True, 'rand_time': 0, 'enable_button': True },
  { 'name': 'human-red', 'on_time': 2, 'off_time': 0, 'repeat': 0, 'lights': 'gR', 'interruptable': False, 'rand_time': 0, 'enable_button': True },
  { 'name': 'yellow', 'on_time': 1.7, 'off_time': 0, 'repeat': 0, 'lights': 'yR', 'interruptable': False, 'rand_time': 0, 'enable_button': True },
  { 'name': 'all-red-no-interrupt', 'on_time': 3, 'off_time': 0, 'repeat': 0, 'lights': 'rR', 'interruptable': False, 'rand_time': 3, 'enable_button': True },
  { 'name': 'all-red', 'on_time': 12, 'off_time': 0, 'repeat': 0, 'lights': 'rR', 'interruptable': True, 'rand_time': 3, 'enable_button': True },
]

curr_state = None
button_event = threading.Event()
button_state = 0  # x0b - diabled, x1b - enabled, 1xb - pressed, 0xb - not pressed

def setup():
  GPIO.setmode(GPIO.BCM)
  GPIO.setup([
      RED_PIN,
      YELLOW_PIN,
      GREEN_PIN,
      HUMAN_GREEN_PIN,
      HUMAN_RED_PIN,
      PRESS_LIGHT_PIN,
      WAIT_LIGHT_PIN
    ],
    GPIO.OUT)
  GPIO.setup(BUTTON_PIN, GPIO.IN, pull_up_down=GPIO.PUD_UP)
  GPIO.output([
      RED_PIN,
      YELLOW_PIN,
      GREEN_PIN,
      HUMAN_GREEN_PIN,
      HUMAN_RED_PIN,
      PRESS_LIGHT_PIN,
      WAIT_LIGHT_PIN
    ],
    GPIO.LOW)
  GPIO.add_interrupt_callback(
      BUTTON_PIN,
      button_pressed,
      edge='falling',
      pull_up_down=GPIO.PUD_UP,
      debounce_timeout_ms=20)

def loop():
  for state in states:
    process_state(state)

def finished():
  GPIO.cleanup()

def enable_button(enable):
  if enable:
    if button_state & 1 == 0:
      GPIO.add_event_detect(BUTTON_PIN, GPIO.FALLING, callback=button_pressed, bouncetime=200)
    button_state |= 1
  else:
    if button_state & 1:
      GPIO.remove_event_detect(BUTTON_PIN)
    button_state = 0
    button_event.clear()
  update_button_light()

def update_button_light():
  GPIO.output(
      (PRESS_LIGHT_PIN, WAIT_LIGHT_PIN),
      (button_state == 1, button_state == 3))

def button_pressed():
  time.sleep(BUTTON_DELAY_TIME)
  button_event.set()

def process_state(state):
  curr_state = state
  enable_button(state.get('enable_button'))
  all_lights, pins = zip(*pin_map.iteritems())
  for _ in xrange(max(1, state['repeat'])):
    # turn on lights
    on_lights = state['on_lights']
    GPIO.output(pins, map(lambda l: l in on_lights, all_lights))
    on_time = state['on_time'] + random.randint(0, state.get('rand_time') or 0)
    wait(state, on_time)

    # turn off lights
    off_lights = state.get('off_lights')
    if state.get('off_time'):
      GPIO.output(map(lambda l: pin_map[l], off_lights), False)
      wait(state, on_time)

def wait(state, duration):
  if state.get('interruptable'):
    button_event.wait(duration)
  else:
    time.sleep(duration)

def signal_handler(signal, frame):
  finished()
  sys.exit(0)

if __name__ == '__main__':
  signal.signal(signal.SIGTERM, signal_handler)
  signal.signal(signal.SIGINT, signal_handler)
  setup()
  while (True):
    loop()

<|endoftext|>