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

from os import environ

from requests import delete, get, post, put

def call(
        uri,
        request_data=None,
        method='GET',
        only_response_code=False,
        rest_url=environ.get('TFS_REST_URL', None),
        token=environ.get('TFS_TOKEN', None),
):
    """Makes a REST call to TFS rest api.
    May use three environment variables:
    * TFS_REST_URL, e.g. http://myorg.com/stash
    * TFS_TOKEN, the tfs token

    Args:
        uri (str): e.g. "{organisation}/{project}/_apis/build/definitions/{definition_id}?api-version=3.2"
        request_data (dict): the JSON request
        method: http method
        only_response_code: default False
        rest_url: default environ.get('TFS_REST_URL', None)
        token: default environ.get('TFS_TOKEN', None),

    Return:
          the JSON response
    """
    return call_url(f'{rest_url}{uri}', request_data, method, only_response_code, token)

def call_url(
        url,
        request_data=None,
        method='GET',
        only_response_code=False,
        token=environ.get('TFS_TOKEN', None),
):
    """Makes a REST call to TFS rest api.
    May use three environment variables:
    * TFS_TOKEN, the tfs token

    Args:
        url (str): e.g. "http://myorg:8080/tfs/{organisation}/{project}/_apis/build/definitions/{definition_id}?api-version=3.2"
        request_data (dict): the JSON request
        method: http method
        only_response_code: default False
        token: default environ.get('TFS_TOKEN', None),

    Return:
          the JSON response
    """
    m = {'DELETE': delete,
         'GET': get,
         'POST': post,
         'PUT': put,
         }[method]

    params = {'url': url, 'auth': ('', token)}
    if request_data:
        params['json'] = request_data

    response = m(**params)
    return response.status_code if only_response_code else response.json() if response.text else response.status_code()
<commit_msg>modified:   vang/tfs/api.py<commit_after>#!/usr/bin/env python3

from os import environ

from requests import delete, get, post, put

def call(
        uri,
        request_data=None,
        method='GET',
        only_response_code=False,
        rest_url=environ.get('TFS_REST_URL', None),
        token=environ.get('TFS_TOKEN', None),
):
    """Makes a REST call to TFS rest api.
    May use three environment variables:
    * TFS_REST_URL, e.g. http://myorg.com/stash
    * TFS_TOKEN, the tfs token

    Args:
        uri (str): e.g. "{organisation}/{project}/_apis/build/definitions/{definition_id}?api-version=3.2"
        request_data (dict): the JSON request
        method: http method
        only_response_code: default False
        rest_url: default environ.get('TFS_REST_URL', None)
        token: default environ.get('TFS_TOKEN', None),

    Return:
          the JSON response
    """
    return call_url(f'{rest_url}{uri}', request_data, method, only_response_code, token)

def call_url(
        url,
        request_data=None,
        method='GET',
        only_response_code=False,
        token=environ.get('TFS_TOKEN', None),
        verify_certificate=False,
):
    """Makes a REST call to TFS rest api.
    May use three environment variables:
    * TFS_TOKEN, the tfs token

    Args:
        url (str): e.g. "http://myorg:8080/tfs/{organisation}/{project}/_apis/build/definitions/{definition_id}?api-version=3.2"
        request_data (dict): the JSON request
        method: http method
        only_response_code: default False
        token: default environ.get('TFS_TOKEN', None),
        verify_certificate: True if https certificate should be verified,

    Return:
          the JSON response
    """
    m = {'DELETE': delete,
         'GET': get,
         'POST': post,
         'PUT': put,
         }[method]

    params = {'url': url, 'auth': ('', token), 'verify': verify_certificate}
    if request_data:
        params['json'] = request_data

    response = m(**params)
    return response.status_code if only_response_code else response.json() if response.text else response.status_code()
<|endoftext|>


<commit_before>import inspect

from coalib.misc.i18n import _
from coalib.misc.Enum import enum

class DocumentationComment:
    _ParseMode = enum("DESCRIPTION", "PARAM", "RETVAL")

    def __init__(self, desc, param_dict, retval_desc):
        """
        Represents a documentation comment of a python class or function.

        :param desc:        A description as string.
        :param param_dict:  A dictionary containing parameter names as key and
                            their description as value.
        :param retval_desc: A string describing the return value.
        """
        self.desc = desc
        self.param_dict = param_dict
        self.retval_desc = retval_desc

    @classmethod
    def from_docstring(cls, docstring):
        """
        Parses a python docstring. Usable attributes are:
        :param
        @param
        :return
        @return
        """
        if not isinstance(docstring, str):
            raise TypeError("Parameter docstring has to be a string.")

        lines = _(inspect.cleandoc(docstring)).split("\n")

        parse_mode = cls._ParseMode.DESCRIPTION
        cur_param = ""

        desc = ""
        param_dict = {}
        retval_desc = ""
        for line in lines:
            line = line.strip()

            if line == "":
                parse_mode = cls._ParseMode.DESCRIPTION

                continue

            if line.startswith(":param ") or line.startswith("@param "):
                parse_mode = cls._ParseMode.PARAM
                splitted = line[7:].split(":", 1)
                cur_param = splitted[0]
                param_dict[cur_param] = splitted[1].strip()

                continue

            if line.startswith(":return: ") or line.startswith("@return: "):
                parse_mode = cls._ParseMode.RETVAL
                retval_desc = line[9:].strip()

                continue

            if parse_mode == cls._ParseMode.RETVAL:
                retval_desc += " " + line
            elif parse_mode == cls._ParseMode.PARAM:
                param_dict[cur_param] += " " + line
            else:
                desc += " " + line

        return (cls(desc=desc.strip(),
                    param_dict=param_dict,
                    retval_desc=retval_desc.strip()))
<commit_msg>DocumentationComment: Remove param type check<commit_after>import inspect

from coalib.misc.i18n import _
from coalib.misc.Enum import enum

class DocumentationComment:
    _ParseMode = enum("DESCRIPTION", "PARAM", "RETVAL")

    def __init__(self, desc, param_dict, retval_desc):
        """
        Represents a documentation comment of a python class or function.

        :param desc:        A description as string.
        :param param_dict:  A dictionary containing parameter names as key and
                            their description as value.
        :param retval_desc: A string describing the return value.
        """
        self.desc = desc
        self.param_dict = param_dict
        self.retval_desc = retval_desc

    @classmethod
    def from_docstring(cls, docstring):
        """
        Parses a python docstring. Usable attributes are:
        :param
        @param
        :return
        @return
        """
        lines = _(inspect.cleandoc(docstring)).split("\n")

        parse_mode = cls._ParseMode.DESCRIPTION
        cur_param = ""

        desc = ""
        param_dict = {}
        retval_desc = ""
        for line in lines:
            line = line.strip()

            if line == "":
                parse_mode = cls._ParseMode.DESCRIPTION

                continue

            if line.startswith(":param ") or line.startswith("@param "):
                parse_mode = cls._ParseMode.PARAM
                splitted = line[7:].split(":", 1)
                cur_param = splitted[0]
                param_dict[cur_param] = splitted[1].strip()

                continue

            if line.startswith(":return: ") or line.startswith("@return: "):
                parse_mode = cls._ParseMode.RETVAL
                retval_desc = line[9:].strip()

                continue

            if parse_mode == cls._ParseMode.RETVAL:
                retval_desc += " " + line
            elif parse_mode == cls._ParseMode.PARAM:
                param_dict[cur_param] += " " + line
            else:
                desc += " " + line

        return (cls(desc=desc.strip(),
                    param_dict=param_dict,
                    retval_desc=retval_desc.strip()))
<|endoftext|>


<commit_before># Copyright 2016 - 2019  Ternaris.
# SPDX-License-Identifier: AGPL-3.0-only

import math
from itertools import count
from subprocess import PIPE

import numpy

import marv_api as marv
from marv_api.types import File
from marv_api.utils import popen
from marv_ros.img_tools import ImageConversionError, ImageFormatError, imgmsg_to_cv2

from .bag import get_message_type, messages

try:
    import cv2
except ImportError:
    cv2 = None

def ros2cv(msg, scale=1, offset=0):
    if 'FC' in msg.encoding:
        passimg = numpy.nan_to_num(imgmsg_to_cv2(msg))
        valscaled = cv2.convertScaleAbs(passimg, None, scale, offset)
        return valscaled

    mono = msg.encoding.startswith('mono') or msg.encoding[-1] in ['1', 'U', 'S', 'F']
    return imgmsg_to_cv2(msg, 'mono8' if mono else 'bgr8')

@marv.node(File)
@marv.input('stream', foreach=marv.select(messages, '*:sensor_msgs/Image'))
@marv.input('speed', default=4)
@marv.input('convert_32FC1_scale', default=1)
@marv.input('convert_32FC1_offset', default=0)
def ffmpeg(stream, speed, convert_32FC1_scale, convert_32FC1_offset):  # pylint: disable=invalid-name
    """Create video for each sensor_msgs/Image topic with ffmpeg."""
    # pylint: disable=too-many-locals

    yield marv.set_header(title=stream.topic)
    name = f"{stream.topic.replace('/', '_')[1:]}.webm"
    video = yield marv.make_file(name)
    duration = (stream.end_time - stream.start_time) * 1e-9
    framerate = (stream.msg_count / duration) if duration else 1

    pytype = get_message_type(stream)
    rosmsg = pytype()

    encoder = None
    while True:
        msg = yield marv.pull(stream)
        if msg is None:
            break
        rosmsg.deserialize(msg.data)
        try:
            img = ros2cv(rosmsg, convert_32FC1_scale, convert_32FC1_offset)
        except (ImageFormatError, ImageConversionError) as err:
            log = yield marv.get_logger()
            log.error('could not convert image from topic %s: %s ', stream.topic, err)
            return

        if not encoder:
            ffargs = [
                'ffmpeg',
                '-f', 'rawvideo',
                '-pixel_format', '%s' % 'gray' if len(img.shape) == 2 else 'bgr24',
                '-video_size', '%dx%d' % (rosmsg.width, rosmsg.height),
                '-framerate', '%s' % framerate,
                '-i', '-',
                '-c:v', 'libvpx-vp9',
                '-pix_fmt', 'yuv420p',
                '-loglevel', 'error',
                '-threads', '8',
                '-speed', str(speed),
                '-y',
                video.path,
            ]
            encoder = popen(ffargs, stdin=PIPE)

        encoder.stdin.write(img)

    encoder.stdin.close()
    encoder.wait()
    yield video

@marv.node(File)
@marv.input('stream', foreach=marv.select(messages, '*:sensor_msgs/Image'))
@marv.input('image_width', default=320)
@marv.input('max_frames', default=50)
@marv.input('convert_32FC1_scale', default=1)
@marv.input('convert_32FC1_offset', default=0)
def images(stream, image_width, max_frames, convert_32FC1_scale, convert_32FC1_offset):
    """Extract max_frames equidistantly spread images from each sensor_msgs/Image stream.

    Args:
        stream: sensor_msgs/Image stream
        image_width (int): Scale to image_width, keeping aspect ratio.
        max_frames (int): Maximum number of frames to extract.
        convert_32FC1_scale (float): Scale factor for FC image values.
        convert_32FC1_offset (float): Offset for FC image values.

    """
    # pylint: disable=invalid-name,too-many-locals

    yield marv.set_header(title=stream.topic)
    pytype = get_message_type(stream)
    rosmsg = pytype()
    interval = int(math.ceil(stream.msg_count / max_frames))
    digits = int(math.ceil(math.log(stream.msg_count) / math.log(10)))
    name_template = '%s-{:0%sd}.jpg' % (stream.topic.replace('/', ':')[1:], digits)
    counter = count()
    while True:
        msg = yield marv.pull(stream)
        if msg is None:
            break
        idx = next(counter)
        if idx % interval:
            continue

        rosmsg.deserialize(msg.data)
        try:
            img = ros2cv(rosmsg, convert_32FC1_scale, convert_32FC1_offset)
        except (ImageFormatError, ImageConversionError) as err:
            log = yield marv.get_logger()
            log.error('could not convert image from topic %s: %s ', stream.topic, err)
            return

        height = int(round(image_width * img.shape[0] / img.shape[1]))
        scaled_img = cv2.resize(img, (image_width, height),
                                interpolation=cv2.INTER_AREA)
        name = name_template.format(idx)
        imgfile = yield marv.make_file(name)
        cv2.imwrite(imgfile.path, scaled_img, (cv2.IMWRITE_JPEG_QUALITY, 60))
        yield imgfile
<commit_msg>Add support for CompressedImage in cam nodes<commit_after># Copyright 2016 - 2019  Ternaris.
# SPDX-License-Identifier: AGPL-3.0-only

import math
from itertools import count
from subprocess import PIPE

import numpy

import marv_api as marv
from marv_api.types import File
from marv_api.utils import popen
from marv_ros.img_tools import (ImageConversionError, ImageFormatError, compressed_imgmsg_to_cv2,
                                imgmsg_to_cv2)

from .bag import make_deserialize, messages

try:
    import cv2
except ImportError:
    cv2 = None

IMAGE_MSG_TYPES = ','.join(['*:sensor_msgs/Image', '*:sensor_msgs/CompressedImage'])

def ros2cv(msg, scale=1, offset=0):
    if hasattr(msg, 'format'):
        return compressed_imgmsg_to_cv2(msg)

    if 'FC' in msg.encoding:
        passimg = numpy.nan_to_num(imgmsg_to_cv2(msg))
        valscaled = cv2.convertScaleAbs(passimg, None, scale, offset)
        return valscaled

    mono = msg.encoding.startswith('mono') or msg.encoding[-1] in ['1', 'U', 'S', 'F']
    return imgmsg_to_cv2(msg, 'mono8' if mono else 'bgr8')

@marv.node(File)
@marv.input('stream', foreach=marv.select(messages, IMAGE_MSG_TYPES))
@marv.input('speed', default=4)
@marv.input('convert_32FC1_scale', default=1)
@marv.input('convert_32FC1_offset', default=0)
def ffmpeg(stream, speed, convert_32FC1_scale, convert_32FC1_offset):  # pylint: disable=invalid-name
    """Create video for each sensor_msgs/Image topic with ffmpeg."""
    # pylint: disable=too-many-locals

    yield marv.set_header(title=stream.topic)
    name = f"{stream.topic.replace('/', '_')[1:]}.webm"
    video = yield marv.make_file(name)
    duration = (stream.end_time - stream.start_time) * 1e-9
    framerate = (stream.msg_count / duration) if duration else 1

    deserialize = make_deserialize(stream)
    encoder = None
    while True:
        msg = yield marv.pull(stream)
        if msg is None:
            break
        rosmsg = deserialize(msg.data)
        try:
            img = ros2cv(rosmsg, convert_32FC1_scale, convert_32FC1_offset)
        except (ImageFormatError, ImageConversionError) as err:
            log = yield marv.get_logger()
            log.error('could not convert image from topic %s: %s ', stream.topic, err)
            return

        if not encoder:
            ffargs = [
                'ffmpeg',
                '-f', 'rawvideo',
                '-pixel_format', 'gray' if len(img.shape) == 2 else 'bgr24',
                '-video_size', f'{img.shape[1]}x{img.shape[0]}',
                '-framerate', str(framerate),
                '-i', '-',
                '-c:v', 'libvpx-vp9',
                '-pix_fmt', 'yuv420p',
                '-loglevel', 'error',
                '-threads', '8',
                '-speed', str(speed),
                '-y',
                video.path,
            ]
            encoder = popen(ffargs, stdin=PIPE)

        encoder.stdin.write(img)

    encoder.stdin.close()
    encoder.wait()
    yield video

@marv.node(File)
@marv.input('stream', foreach=marv.select(messages, IMAGE_MSG_TYPES))
@marv.input('image_width', default=320)
@marv.input('max_frames', default=50)
@marv.input('convert_32FC1_scale', default=1)
@marv.input('convert_32FC1_offset', default=0)
def images(stream, image_width, max_frames, convert_32FC1_scale, convert_32FC1_offset):
    """Extract max_frames equidistantly spread images from each sensor_msgs/Image stream.

    Args:
        stream: sensor_msgs/Image stream
        image_width (int): Scale to image_width, keeping aspect ratio.
        max_frames (int): Maximum number of frames to extract.
        convert_32FC1_scale (float): Scale factor for FC image values.
        convert_32FC1_offset (float): Offset for FC image values.

    """
    # pylint: disable=invalid-name,too-many-locals

    yield marv.set_header(title=stream.topic)
    deserialize = make_deserialize(stream)
    interval = int(math.ceil(stream.msg_count / max_frames))
    digits = int(math.ceil(math.log(stream.msg_count) / math.log(10)))
    name_template = '%s-{:0%sd}.jpg' % (stream.topic.replace('/', ':')[1:], digits)
    counter = count()
    while True:
        msg = yield marv.pull(stream)
        if msg is None:
            break
        idx = next(counter)
        if idx % interval:
            continue

        rosmsg = deserialize(msg.data)
        try:
            img = ros2cv(rosmsg, convert_32FC1_scale, convert_32FC1_offset)
        except (ImageFormatError, ImageConversionError) as err:
            log = yield marv.get_logger()
            log.error('could not convert image from topic %s: %s ', stream.topic, err)
            return

        height = int(round(image_width * img.shape[0] / img.shape[1]))
        scaled_img = cv2.resize(img, (image_width, height),
                                interpolation=cv2.INTER_AREA)
        name = name_template.format(idx)
        imgfile = yield marv.make_file(name)
        cv2.imwrite(imgfile.path, scaled_img, (cv2.IMWRITE_JPEG_QUALITY, 60))
        yield imgfile
<|endoftext|>


<commit_before># -*- coding:utf-8 -*-
"""
===================================================================
Provide coordinate file class which do operations on these files.
===================================================================
Written by PytLab <shaozhengjiang@gmail.com>, November 2014
Updated by PytLab <shaozhengjiang@gmail.com>, August 2015

==============================================================

"""
import numpy as np

from vaspy import VasPy, CarfileValueError
from functions import *

class AtomCo(VasPy):
    "Base class to be inherited by atomco classes."
    def __init__(self, filename):
        VasPy.__init__(self, filename)

    def __repr__(self):
        return self.get_content()

    def __str__(self):
        return self.__repr__()

    def __getattribute__(self, attr):
        '''
        确保atomco_dict能够及时根据data值的变化更新.
        '''
        if attr == 'atomco_dict':
            return self.get_atomco_dict(self.data)
        else:
            return object.__getattribute__(self, attr)

    def verify(self):
        if len(self.data) != self.ntot:
            raise CarfileValueError('Atom numbers mismatch!')

    def get_atomco_dict(self, data):
        "根据已获取的data和atoms, atoms_num, 获取atomco_dict"
        # [1, 1, 1, 16] -> [0, 1, 2, 3, 19]
        idx_list = [sum(self.atoms_num[:i]) for i in xrange(1, len(self.atoms)+1)]
        idx_list = [0] + idx_list
        data_list = data.tolist()
        atomco_dict = {}
        for atom, idx, next_idx in zip(self.atoms, idx_list[:-1], idx_list[1:]):
            atomco_dict.setdefault(atom, data_list[idx: next_idx])

        self.atomco_dict = atomco_dict

        return atomco_dict

class XyzFile(AtomCo):
    """
    Create a .xyz file class.

    Example:

    >>> a = XyzFile(filename='ts.xyz')

    Class attributes descriptions
    =======================================================================
      Attribute      Description
      ============  =======================================================
      filename       string, name of the file the direct coordiante data
                     stored in
      ntot           int, the number of total atom number
      step           int, STEP number in OUT.ANI file
      atoms          list of strings, atom types
      natoms         list of tuples, same shape with atoms.
                     (atom name, atom number)
                     atom number of atoms in atoms
      atomco_dict    dict, {atom name: coordinates}
      data           np.array, coordinates of atoms, dtype=float64
      ============  =======================================================
    """
    def __init__(self, filename):
        AtomCo.__init__(self, filename)
        self.load()
        self.verify()

    def load(self):
        "加载文件内容"
        with open(self.filename, 'r') as f:
            content_list = f.readlines()
        ntot = int(content_list[0].strip())  # total atom number
        step = int(str2list(content_list[1])[-1])  # iter step number

        #get atom coordinate and number info
        data_list = [str2list(line) for line in content_list[2:]]
        data_array = np.array(data_list)  # dtype=np.string
        atoms_list = list(data_array[:, 0])  # 1st column
        data = np.float64(data_array[:, 1:])  # rest columns

        #get atom number for each atom
        atoms = []
        for atom in atoms_list:
            if atom not in atoms:
                atoms.append(atom)
        atoms_num = [atoms_list.count(atom) for atom in atoms]
        natoms = zip(atoms, atoms_num)

        #set class attrs
        self.ntot = ntot
        self.step = step
        self.atoms = atoms
        self.atoms_num = atoms_num
        self.natoms = natoms
        self.data = data

        #get atomco_dict
        self.get_atomco_dict(data)

        return

    def coordinate_transfrom(self, axes=np.array([[1.0, 0.0, 0.0],
                                                  [0.0, 1.0, 0.0],
                                                  [0.0, 0.0, 1.0]])):
        "相对坐标和实坐标转换"
        "Use Ax=b to do coordinate transform. direct to cartesian"
        b = np.matrix(self.data.T)
        A = np.matrix(axes).T
        x = A.I*b

        return np.array(x.T)

    def get_content(self):
        "获取最新文件内容字符串"
        ntot = "%12d\n" % self.ntot
        step = "STEP =%9d\n" % self.step
        data = atomdict2str(self.atomco_dict, self.atoms)
        content = ntot + step + data

        return content

    def tofile(self, filename='atomco.xyz'):
        "XyzFile object to .xyz file."
        content = self.get_content()

        with open(filename, 'w') as f:
            f.write(content)

        return

class PosCar(AtomCo):
    def __init__(self, filename='POSCAR'):
        """
        Class to generate POSCAR or CONTCAR-like objects.

        Example:

        >>> a = PosCar(filename='POSCAR')

        Class attributes descriptions
        =======================================================================
          Attribute      Description
          ============  =======================================================
          filename       string, name of the file the direct coordiante data
                         stored in
          axes_coeff     float, Scale Factor of axes
          axes           np.array, axes of POSCAR
          atoms          list of strings, atom types
          ntot           int, the number of total atom number
          natoms         list of int, same shape with atoms
                         atom number of atoms in atoms
          tf             list of list, T&F info of atoms
          data           np.array, coordinates of atoms, dtype=float64
          ============  =======================================================
        """
        AtomCo.__init__(self, filename)
        #load all data in file
        self.load()
        self.verify()

    def load(self):
        "获取文件数据信息"
        "Load all information in POSCAR."
        with open(self.filename, 'r') as f:
            content_list = f.readlines()
        #get scale factor
        axes_coeff = float(content_list[1])
        #axes
        axes = [str2list(axis) for axis in content_list[2:5]]
        #atom info
        atoms = str2list(content_list[5])
        atoms_num = str2list(content_list[6])  # atom number
        #data
        data, tf = [], []  # data and T or F info
        for line_str in content_list[9:]:
            line_list = str2list(line_str)
            data.append(line_list[:3])
            if len(line_list) > 3:
                tf.append(line_list[3:])
        #data type convertion
        axes = np.float64(np.array(axes))  # to float
        atoms_num = [int(i) for i in atoms_num]
        data = np.float64(np.array(data))
        tf = np.array(tf)

        #set class attrs
        self.axes_coeff = axes_coeff
        self.axes = axes
        self.atoms = atoms
        self.atoms_num = atoms_num
        self.ntot = sum(atoms_num)
        self.natoms = zip(atoms, atoms_num)
        self.data = data
        self.tf = tf

        #get atomco_dict
        self.get_atomco_dict(data)

        return

    def constrain_atom(self, atom, to='F', axis='all'):
        "修改某一类型原子的FT信息"
        # [1, 1, 1, 16] -> [0, 1, 2, 3, 19]
        idx_list = [sum(self.atoms_num[:i]) for i in xrange(1, len(self.atoms)+1)]
        idx_list = [0] + idx_list

        if to not in ['T', 'F']:
            raise CarfileValueError('Variable to must be T or F.')

        for atomtype, idx, next_idx in \
                zip(self.atoms, idx_list[:-1], idx_list[1:]):
            if atomtype == atom:
                if axis == 'x' or axis == 'X':
                    self.tf[idx:next_idx, 0] = to
                elif axis == 'y' or axis == 'Y':
                    self.tf[idx:next_idx, 1] = to
                elif axis == 'z' or axis == 'Z':
                    self.tf[idx:next_idx, 2] = to
                else:
                    self.tf[idx:next_idx, :] = to
                break

        return self.tf

    def get_content(self):
        "根据对象数据获取文件内容字符串"
        content = 'Created by VASPy\n'
        axe_coeff = " %.9f\n" % self.axes_coeff
        #axes
        axes_list = self.axes.tolist()
        axes = ''
        for axis in axes_list:
            axes += "%14.8f%14.8f%14.8f\n" % tuple(axis)
        #atom info
        atoms, atoms_num = zip(*self.natoms)
        atoms = ("%5s"*len(atoms)+"\n") % atoms
        atoms_num = ("%5d"*len(atoms_num)+"\n") % atoms_num
        #string
        info = "Selective Dynamics\nDirect\n"
        #data and tf
        data_tf = ''
        for data, tf in zip(self.data.tolist(), self.tf):
            data_tf += ("%18.12f"*3+"%5s"*3+"\n") % tuple(data+tf)
        #merge all strings
        content += axe_coeff+axes+atoms+atoms_num+info+data_tf

        return content

    def tofile(self, filename='POSCAR_c'):
        "生成文件"
        "PosCar object to POSCAR or CONTCAR."
        content = self.get_content()

        with open(filename, 'w') as f:
            f.write(content)

        return
<commit_msg>fix bug<commit_after># -*- coding:utf-8 -*-
"""
===================================================================
Provide coordinate file class which do operations on these files.
===================================================================
Written by PytLab <shaozhengjiang@gmail.com>, November 2014
Updated by PytLab <shaozhengjiang@gmail.com>, August 2015

==============================================================

"""
import numpy as np

from vaspy import VasPy, CarfileValueError
from functions import *

class AtomCo(VasPy):
    "Base class to be inherited by atomco classes."
    def __init__(self, filename):
        VasPy.__init__(self, filename)

    def __repr__(self):
        return self.get_content()

    def __str__(self):
        return self.__repr__()

    def __getattribute__(self, attr):
        '''
        确保atomco_dict能够及时根据data值的变化更新.
        '''
        if attr == 'atomco_dict':
            return self.get_atomco_dict(self.data)
        else:
            return object.__getattribute__(self, attr)

    def verify(self):
        if len(self.data) != self.ntot:
            raise CarfileValueError('Atom numbers mismatch!')

    def get_atomco_dict(self, data):
        "根据已获取的data和atoms, atoms_num, 获取atomco_dict"
        # [1, 1, 1, 16] -> [0, 1, 2, 3, 19]
        idx_list = [sum(self.atoms_num[:i]) for i in xrange(1, len(self.atoms)+1)]
        idx_list = [0] + idx_list
        data_list = data.tolist()
        atomco_dict = {}
        for atom, idx, next_idx in zip(self.atoms, idx_list[:-1], idx_list[1:]):
            atomco_dict.setdefault(atom, data_list[idx: next_idx])

        self.atomco_dict = atomco_dict

        return atomco_dict

class XyzFile(AtomCo):
    """
    Create a .xyz file class.

    Example:

    >>> a = XyzFile(filename='ts.xyz')

    Class attributes descriptions
    =======================================================================
      Attribute      Description
      ============  =======================================================
      filename       string, name of the file the direct coordiante data
                     stored in
      ntot           int, the number of total atom number
      step           int, STEP number in OUT.ANI file
      atoms          list of strings, atom types
      natoms         list of tuples, same shape with atoms.
                     (atom name, atom number)
                     atom number of atoms in atoms
      atomco_dict    dict, {atom name: coordinates}
      data           np.array, coordinates of atoms, dtype=float64
      ============  =======================================================
    """
    def __init__(self, filename):
        AtomCo.__init__(self, filename)
        self.load()
        self.verify()

    def load(self):
        "加载文件内容"
        with open(self.filename, 'r') as f:
            content_list = f.readlines()
        ntot = int(content_list[0].strip())  # total atom number
        step = int(str2list(content_list[1])[-1])  # iter step number

        #get atom coordinate and number info
        data_list = [str2list(line) for line in content_list[2:]]
        data_array = np.array(data_list)  # dtype=np.string
        atoms_list = list(data_array[:, 0])  # 1st column
        data = np.float64(data_array[:, 1:])  # rest columns

        #get atom number for each atom
        atoms = []
        for atom in atoms_list:
            if atom not in atoms:
                atoms.append(atom)
        atoms_num = [atoms_list.count(atom) for atom in atoms]
        natoms = zip(atoms, atoms_num)

        #set class attrs
        self.ntot = ntot
        self.step = step
        self.atoms = atoms
        self.atoms_num = atoms_num
        self.natoms = natoms
        self.data = data

        #get atomco_dict
        self.get_atomco_dict(data)

        return

    def coordinate_transfrom(self, axes=np.array([[1.0, 0.0, 0.0],
                                                  [0.0, 1.0, 0.0],
                                                  [0.0, 0.0, 1.0]])):
        "相对坐标和实坐标转换"
        "Use Ax=b to do coordinate transform. direct to cartesian"
        b = np.matrix(self.data.T)
        A = np.matrix(axes).T
        x = A.I*b

        return np.array(x.T)

    def get_content(self):
        "获取最新文件内容字符串"
        ntot = "%12d\n" % self.ntot
        step = "STEP =%9d\n" % self.step
        data = atomdict2str(self.atomco_dict, self.atoms)
        content = ntot + step + data

        return content

    def tofile(self, filename='atomco.xyz'):
        "XyzFile object to .xyz file."
        content = self.get_content()

        with open(filename, 'w') as f:
            f.write(content)

        return

class PosCar(AtomCo):
    def __init__(self, filename='POSCAR'):
        """
        Class to generate POSCAR or CONTCAR-like objects.

        Example:

        >>> a = PosCar(filename='POSCAR')

        Class attributes descriptions
        =======================================================================
          Attribute      Description
          ============  =======================================================
          filename       string, name of the file the direct coordiante data
                         stored in
          axes_coeff     float, Scale Factor of axes
          axes           np.array, axes of POSCAR
          atoms          list of strings, atom types
          ntot           int, the number of total atom number
          natoms         list of int, same shape with atoms
                         atom number of atoms in atoms
          tf             list of list, T&F info of atoms
          data           np.array, coordinates of atoms, dtype=float64
          ============  =======================================================
        """
        AtomCo.__init__(self, filename)
        #load all data in file
        self.load()
        self.verify()

    def load(self):
        "获取文件数据信息"
        "Load all information in POSCAR."
        with open(self.filename, 'r') as f:
            content_list = f.readlines()
        #get scale factor
        axes_coeff = float(content_list[1])
        #axes
        axes = [str2list(axis) for axis in content_list[2:5]]
        #atom info
        atoms = str2list(content_list[5])
        atoms_num = str2list(content_list[6])  # atom number
        #data
        data, tf = [], []  # data and T or F info
        for line_str in content_list[9:]:
            line_list = str2list(line_str)
            data.append(line_list[:3])
            if len(line_list) > 3:
                tf.append(line_list[3:])
        #data type convertion
        axes = np.float64(np.array(axes))  # to float
        atoms_num = [int(i) for i in atoms_num]
        data = np.float64(np.array(data))
        tf = np.array(tf)

        #set class attrs
        self.axes_coeff = axes_coeff
        self.axes = axes
        self.atoms = atoms
        self.atoms_num = atoms_num
        self.ntot = sum(atoms_num)
        self.natoms = zip(atoms, atoms_num)
        self.data = data
        self.tf = tf

        #get atomco_dict
        self.get_atomco_dict(data)

        return

    def constrain_atom(self, atom, to='F', axis='all'):
        "修改某一类型原子的FT信息"
        # [1, 1, 1, 16] -> [0, 1, 2, 3, 19]
        idx_list = [sum(self.atoms_num[:i]) for i in xrange(1, len(self.atoms)+1)]
        idx_list = [0] + idx_list

        if to not in ['T', 'F']:
            raise CarfileValueError('Variable to must be T or F.')

        for atomtype, idx, next_idx in \
                zip(self.atoms, idx_list[:-1], idx_list[1:]):
            if atomtype == atom:
                if axis == 'x' or axis == 'X':
                    self.tf[idx:next_idx, 0] = to
                elif axis == 'y' or axis == 'Y':
                    self.tf[idx:next_idx, 1] = to
                elif axis == 'z' or axis == 'Z':
                    self.tf[idx:next_idx, 2] = to
                else:
                    self.tf[idx:next_idx, :] = to
                break

        return self.tf

    def get_content(self):
        "根据对象数据获取文件内容字符串"
        content = 'Created by VASPy\n'
        axe_coeff = " %.9f\n" % self.axes_coeff
        #axes
        axes_list = self.axes.tolist()
        axes = ''
        for axis in axes_list:
            axes += "%14.8f%14.8f%14.8f\n" % tuple(axis)
        #atom info
        atoms, atoms_num = zip(*self.natoms)
        atoms = ("%5s"*len(atoms)+"\n") % atoms
        atoms_num = ("%5d"*len(atoms_num)+"\n") % atoms_num
        #string
        info = "Selective Dynamics\nDirect\n"
        #data and tf
        data_tf = ''
        for data, tf in zip(self.data.tolist(), self.tf.tolist()):
            data_tf += ("%18.12f"*3+"%5s"*3+"\n") % tuple(data+tf)
        #merge all strings
        content += axe_coeff+axes+atoms+atoms_num+info+data_tf

        return content

    def tofile(self, filename='POSCAR_c'):
        "生成文件"
        "PosCar object to POSCAR or CONTCAR."
        content = self.get_content()

        with open(filename, 'w') as f:
            f.write(content)

        return
<|endoftext|>


<commit_before>#!/usr/bin/env python
#-*- coding: utf-8 -*-
"""extract_opensubs_dialogs.py:
   A script to extract text from OpenSubtitles.

   Copyright (c) 2017 Takaaki Hori  (thori@merl.com)

   This software is released under the MIT License.
   http://opensource.org/licenses/mit-license.php

"""

import xml.etree.ElementTree as ET
from gzip import GzipFile
import sys
import re
import glob
import random
from tqdm import tqdm
import six
import argparse

# regular expressions for cleaning up text
tag_re = re.compile(r'(<!--.*?-->|<[^>]*>|{[^}]*}|\([^\)]*\)|\[[^\]]*\])')
symbol_re = re.compile(u'([-–—]+$| [-–—]+|[-–—]+ |% %|#+|\'\'|``| \' |[\(\)\"])')
space_re = re.compile(r'\s+')
appos1_re = re.compile(r' \' (?=(em|im|s|t|bout|cause)\s+)')
appos2_re = re.compile(r'(?<=n) \' ')
title_re = re.compile(r'(?<=( mr| jr| ms| dr| st|mrs)) \.')
speaker_re = re.compile(r'^\s*[A-z]*\s*:')
ignore_re = re.compile(r'^(\s*|[\.\?$%!,:;])$')

def preprocess(sent):
    new_sent = tag_re.sub(' ', sent).replace('\\\'','\'').lower()
    new_sent = appos1_re.sub(' \'', new_sent)
    new_sent = appos2_re.sub('\' ', new_sent)
    new_sent = title_re.sub('. ', new_sent)
    new_sent = speaker_re.sub('', new_sent)
    new_sent = space_re.sub(' ', symbol_re.sub(' ', new_sent)).strip()
    if not ignore_re.match(new_sent):
        return new_sent.encode('utf-8')
    else:
        return ''

def extract(filePath, corpus):
    """extract text from an XML file and tokenize it
    """
    if filePath.endswith('.gz'):
        tree = ET.parse(GzipFile(filename=filePath))
    else:
        tree = ET.parse(filePath)

    root = tree.getroot()
    sent = ''
    for child in root:
        for elem in child:
            if elem.tag == 'w':
                sent += ' ' + elem.text

        if not sent.strip().endswith(':'):
            new_sent = preprocess(sent)
            #six.print_('%s -> %s' % (sent.encode('utf-8'),new_sent))
            if new_sent:
                corpus.append((new_sent, len(new_sent.split())))
            sent = ''

    corpus.append(('',0))

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--output', default=['train.txt'], nargs='+', type=str,
                        help='Filenames of data')
    parser.add_argument('--ratio', default=[0.01], nargs='+', type=float,
                        help='Extraction rate for each data set')
    parser.add_argument('--max-length', default=20, type=float,
                        help='Maximum length of sentences')
    parser.add_argument('--rootdir', default='.',
                        help='root directory of data source')

    args = parser.parse_args()

    if len(args.output) != len(args.ratio):
        raise Exception('The number of output files (%d) and the number extraction ratios (%d) should be the same.' % (len(args.output), len(args.partion)))

    rootdir = args.rootdir
    print('collecting files from ' + rootdir)
    xmlfiles = glob.glob(rootdir + '/*.xml.gz') + glob.glob(rootdir + '/*/*.xml.gz') + glob.glob(rootdir + '/*/*/*.xml.gz')

    random.shuffle(xmlfiles, random.random)
    total_ratio = sum(args.ratio)
    n_files = int( len(xmlfiles) * total_ratio )
    print('loading text from %d/%d files' % (n_files, len(xmlfiles)))
    corpus = []
    for n in tqdm(six.moves.range(n_files)):
        extract(xmlfiles[n], corpus)

    print('%d sentences loaded' % len(corpus))
    random.seed(1)
    indices = list(six.moves.range(len(corpus)-1))
    random.shuffle(indices, random.random)

    partition = [0]
    acc_ratio = 0.0
    for r in args.ratio:
        acc_ratio += r / total_ratio
        partition.append(int(acc_ratio * len(corpus)))

    for n in six.moves.range(len(partition)-1):
        print('writing %d sentence pairs to %s' % (partition[n+1]-partition[n], args.output[n]))
        with open(args.output[n],'w') as f:
            for idx in indices[partition[n]:partition[n+1]]:
                len1 = corpus[idx][1]
                len2 = corpus[idx+1][1]
                if 0 < len1 < args.max_length and 0 < len2 < args.max_length:
                    six.print_('U: %s\nS: %s\n' % (corpus[idx][0],corpus[idx+1][0]), file=f)

<commit_msg>Fixed a bug in python3 and modified random seed, regular expression, etc.<commit_after>#!/usr/bin/env python
#-*- coding: utf-8 -*-
"""extract_opensubs_dialogs.py:
   A script to extract text from OpenSubtitles.

   Copyright (c) 2017 Takaaki Hori  (thori@merl.com)

   This software is released under the MIT License.
   http://opensource.org/licenses/mit-license.php

"""

import xml.etree.ElementTree as ET
from gzip import GzipFile
import sys
import re
import glob
import random
from tqdm import tqdm
import six
import argparse

def preprocess(sent):
    """ text preprocessing using regular expressions
    """
    # remove tags
    new_sent = re.sub(r'(<!--.*?-->|<[^>]*>|{[^}]*}|\([^\)]*\)|\[[^\]]*\])',' ', sent)
    # replace apostrophe and convert letters to lower case
    new_sent = new_sent.replace('\\\'','\'').lower()
    # delete a space right after an isolated apostrophe
    new_sent = re.sub(r' \' (?=(em|im|s|t|bout|cause)\s+)', ' \'', new_sent)
    # delete a space right before an isolated apostrophe
    new_sent = re.sub(r'(?<=n) \' ', '\' ', new_sent)
    # delete a space right before a period for titles
    new_sent = re.sub(r'(?<=( mr| jr| ms| dr| st|mrs)) \.', '. ', new_sent)
    # remove speaker tag "xxx: "
    new_sent = re.sub(r'^\s*[A-z]*\s*:', '', new_sent)
    # remove unnecessary symbols
    new_sent = re.sub(u'([-–—]+$| [-–—]+|[-–—]+ |% %|#+|\'\'|``| \' |[\(\)\"])', ' ', new_sent)
    # convert i̇->i
    new_sent = re.sub(u'i̇','i', new_sent)
    # convert multiple spaces to a single space
    new_sent = re.sub(r'\s+', ' ', new_sent).strip()
    # ignore sentence with anly space or some symbols
    if not re.match(r'^(\s*|[\.\?$%!,:;])$', new_sent):
        return new_sent
    else:
        return ''

def extract(filePath, corpus):
    """extract text from an XML file and tokenize it
    """
    if filePath.endswith('.gz'):
        tree = ET.parse(GzipFile(filename=filePath))
    else:
        tree = ET.parse(filePath)

    root = tree.getroot()
    sent = ''
    for child in root:
        for elem in child:
            if elem.tag == 'w':
                sent += ' ' + elem.text

        if not sent.strip().endswith(':'):
            if six.PY2:
                new_sent = preprocess(sent).encode('utf-8')
            else:
                new_sent = preprocess(sent)

            if new_sent:
                corpus.append((new_sent, len(new_sent.split())))
            sent = ''

    corpus.append(('',0))

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--output', default=['train.txt'], nargs='+', type=str,
                        help='Filenames of data')
    parser.add_argument('--ratio', default=[0.01], nargs='+', type=float,
                        help='Extraction rate for each data set')
    parser.add_argument('--max-length', default=20, type=float,
                        help='Maximum length of sentences')
    parser.add_argument('--rootdir', default='.',
                        help='root directory of data source')

    args = parser.parse_args()

    if len(args.output) != len(args.ratio):
        raise Exception('The number of output files (%d) and the number extraction ratios (%d) should be the same.' % (len(args.output), len(args.partion)))
    random.seed(99)

    rootdir = args.rootdir
    print('collecting files from ' + rootdir)
    xmlfiles = glob.glob(rootdir + '/*.xml.gz') + glob.glob(rootdir + '/*/*.xml.gz') + glob.glob(rootdir + '/*/*/*.xml.gz')

    random.shuffle(xmlfiles, random.random)
    total_ratio = sum(args.ratio)
    n_files = int( len(xmlfiles) * total_ratio )
    print('loading text from %d/%d files' % (n_files, len(xmlfiles)))
    corpus = []
    for n in tqdm(six.moves.range(n_files)):
        extract(xmlfiles[n], corpus)

    print('%d sentences loaded' % len(corpus))
    indices = list(six.moves.range(len(corpus)-1))
    random.shuffle(indices, random.random)

    partition = [0]
    acc_ratio = 0.0
    for r in args.ratio:
        acc_ratio += r / total_ratio
        partition.append(int(acc_ratio * len(corpus)))

    for n in six.moves.range(len(partition)-1):
        print('writing %d sentence pairs to %s' % (partition[n+1]-partition[n], args.output[n]))
        with open(args.output[n],'w') as f:
            for idx in indices[partition[n]:partition[n+1]]:
                len1 = corpus[idx][1]
                len2 = corpus[idx+1][1]
                if 0 < len1 < args.max_length and 0 < len2 < args.max_length:
                    six.print_('U: %s\nS: %s\n' % (corpus[idx][0],corpus[idx+1][0]), file=f)

<|endoftext|>


<commit_before>import unittest
from random import Random

from fito import DictDataStore
from fito import SpecField
from fito.operations.decorate import as_operation
from fito.operations.operation import Operation
from fito.specs.base import InvalidSpecInstance, Spec, BaseSpecField, PrimitiveField

@as_operation()
def base_case(i=0):
    return i

@as_operation(op=SpecField)
def f(op, val=1): return val

class Test(object):
    @as_operation(method_type='instance', op=SpecField(base_type=base_case))
    def op1(self, op, val=1): return 1

    @as_operation(method_type='class', op=SpecField(base_type=base_case))
    def op2(cls, op): return 1

def get_test_operations():
    res = [
        base_case(),
        f(base_case()),
        Test().op1(base_case()),
        Test.op2(base_case()),
    ]

    input = DictDataStore()
    numbers = [Number(i) for i in xrange(10)]
    for number in numbers: input[number] = number.n * 2

    # rnd = Random(42)
    # for i in xrange(10):
    #     numbers.append(rnd.choice(numbers) + rnd.choice(numbers))

    return res + numbers

class Numeric(Operation):
    def __add__(self, other):
        return AddOperation(self, other)

class Number(Numeric):
    n = PrimitiveField(0)

    def apply(self, runner):
        return self.n

class AddOperation(Numeric):
    left = SpecField(0, base_type=Numeric)
    right = SpecField(1, base_type=Numeric)

    def apply(self, runner):
        return runner.execute(self.left) + runner.execute(self.right)

    def __repr__(self):
        return '{} + {}'.format(self.left, self.right)

class TestOperation(unittest.TestCase):
    def test_methods(self):
        for op in get_test_operations():
            assert isinstance(op, Operation)
            op.execute()
            repr(op)
<commit_msg>MemoryObject test<commit_after>import unittest
from random import Random

from fito import DictDataStore
from fito import SpecField
from fito.operations.decorate import as_operation
from fito.operations.operation import Operation, MemoryObject
from fito.specs.base import InvalidSpecInstance, Spec, BaseSpecField, PrimitiveField

@as_operation()
def base_case(i=0):
    return i

@as_operation(op=SpecField)
def f(op, val=1): return val

class Test(object):
    @as_operation(method_type='instance', op=SpecField(base_type=base_case))
    def op1(self, op, val=1): return 1

    @as_operation(method_type='class', op=SpecField(base_type=base_case))
    def op2(cls, op): return 1

def get_test_operations():
    res = [
        base_case(),
        f(base_case()),
        Test().op1(base_case()),
        Test.op2(base_case()),
    ]

    input = DictDataStore()
    numbers = [Number(i) for i in xrange(10)]
    for number in numbers: input[number] = number.n * 2

    # rnd = Random(42)
    # for i in xrange(10):
    #     numbers.append(rnd.choice(numbers) + rnd.choice(numbers))

    return res + numbers

class Numeric(Operation):
    def __add__(self, other):
        return AddOperation(self, other)

class Number(Numeric):
    n = PrimitiveField(0)

    def apply(self, runner):
        return self.n

class AddOperation(Numeric):
    left = SpecField(0, base_type=Numeric)
    right = SpecField(1, base_type=Numeric)

    def apply(self, runner):
        return runner.execute(self.left) + runner.execute(self.right)

    def __repr__(self):
        return '{} + {}'.format(self.left, self.right)

class List(MemoryObject):
    l = PrimitiveField(0)

class TestOperation(unittest.TestCase):
    def test_methods(self):
        for op in get_test_operations():
            assert isinstance(op, Operation)
            op.execute()
            repr(op)

    def test_memory_object(self):
        l = List(range(10))
        d = l.to_dict()
        import ipdb;ipdb.set_trace()
        assert l.l == List.dict2spec(l.to_dict()).l
<|endoftext|>


<commit_before>from django.conf import settings
from django.core.urlresolvers import reverse
from django.contrib.sites.models import Site
from django.core.cache import cache

from spreedly.models import Plan, Subscription
from spreedly.pyspreedly.api import Client

def sync_plans():
    '''
    Sync subscription plans with Spreedly API
    '''
    client = Client(settings.SPREEDLY_AUTH_TOKEN_SECRET, settings.SPREEDLY_SITE_NAME)

    for plan in client.get_plans():
        p, created = Plan.objects.get_or_create(speedly_id=plan['speedly_id'])

        changed = False
        for k, v in plan.items():
            if hasattr(p, k) and not getattr(p, k) == v:
                setattr(p, k, v)
                changed = True
        if changed:
            p.save()

def get_subscription(user):
    cache_key = 'spreedly-subscription-%d' % user.id

    subscription = cache.get(cache_key)

    if not subscription:
        client = Client(settings.SPREEDLY_AUTH_TOKEN_SECRET, settings.SPREEDLY_SITE_NAME)
        data = client.get_info(user.id)

        subscription, created = Subscription.objects.get_or_create(
            user=user
        )
        for k, v in data.items():
            if hasattr(subscription, k):
                setattr(subscription, k, v)
        subscription.save()
        cache.set(cache_key, subscription, 60 * 60)

    return subscription

def create_subscription(user):
    client = Client(settings.SPREEDLY_AUTH_TOKEN_SECRET, settings.SPREEDLY_SITE_NAME)
    client.get_or_create_subscriber(user.id, user.username)
    return get_subscription(user)

def get_or_create_subscription(user):
    client = Client(settings.SPREEDLY_AUTH_TOKEN_SECRET, settings.SPREEDLY_SITE_NAME)
    data = client.get_or_create_subscriber(user.id, user.username)

    subscription, created = Subscription.objects.get_or_create(
        user=user
    )
    for k, v in data.items():
        if hasattr(subscription, k):
            setattr(subscription, k, v)
    subscription.save()
    return subscription

def check_trial_eligibility(plan, user):
    if plan.plan_type != 'free_trial':
        return False
    try:
        # make sure the user is trial eligable (they don't have a subscription yet, or they are trial_elegible)
        not_allowed = Subscription.objects.get(user=user, trial_elegible=False)
        return False
    except Subscription.DoesNotExist:
        return True

def start_free_trial(plan, user):
    if check_trial_eligibility(plan, user):
        client = Client(settings.SPREEDLY_AUTH_TOKEN_SECRET, settings.SPREEDLY_SITE_NAME)
        client.get_or_create_subscriber(user.id, user.username)
        client.subscribe(user.id, plan.pk, trial=True)
        get_subscription(user)
        return True
    else:
        return False

def complimentary_time_extension(user, duration_quantity, duration_units):
    client = Client(settings.SPREEDLY_AUTH_TOKEN_SECRET, settings.SPREEDLY_SITE_NAME)
    client.complimentary_time_extension(user.id, duration_quantity, duration_units)
    return get_subscription(user)

def complimentary_subscription(user, duration_quantity, duration_units, feature_level):
    client = Client(settings.SPREEDLY_AUTH_TOKEN_SECRET, settings.SPREEDLY_SITE_NAME)
    client.complimentary_subscription(user.id, duration_quantity, duration_units, feature_level)
    return get_subscription(user)

def lifetime_complimentary_subscription(user, feature_level):
    client = Client(settings.SPREEDLY_AUTH_TOKEN_SECRET, settings.SPREEDLY_SITE_NAME)
    client.lifetime_complimentary_subscription(user.id, feature_level)
    return get_subscription(user)

def add_store_credit(user, amount):
    client = Client(settings.SPREEDLY_AUTH_TOKEN_SECRET, settings.SPREEDLY_SITE_NAME)
    client.add_store_credit(user.id, amount)
    return get_subscription(user)

def stop_auto_renew(user):
    client = Client(settings.SPREEDLY_AUTH_TOKEN_SECRET, settings.SPREEDLY_SITE_NAME)
    client.stop_auto_renew(user.id)
    return get_subscription(user)

def return_url(user, plan=None, trial=False):
    args = [user.id]
    if plan:
        args.append(plan.pk)
    url = 'http://%s%s' % (Site.objects.get(id=settings.SITE_ID), reverse('spreedly_return', args=args))
    if trial:
        url = url + '?trial=true'
    return url

def allow_free_trial(user):
    client = Client(settings.SPREEDLY_AUTH_TOKEN_SECRET, settings.SPREEDLY_SITE_NAME)
    client.allow_free_trial(user.id)
    return get_subscription(user)

def subscription_url(plan, user, return_url):
    return 'https://spreedly.com/%(site_name)s/subscribers/%(user_id)s/subscribe/%(plan_id)s/%(user_username)s?email=%(user_email)s&return_url=%(return_url)s&first_name=%(first_name)s&last_name=%(last_name)s' % {
        'site_name': settings.SPREEDLY_SITE_NAME,
        'plan_id': plan.pk,
        'user_id': user.id,
        'user_username': user.username,
        'user_email': user.email,
        'first_name': user.first_name,
        'last_name': user.last_name,
        'return_url': return_url,
    }
<commit_msg>Subscribe user to plan<commit_after>from django.conf import settings
from django.core.urlresolvers import reverse
from django.contrib.sites.models import Site
from django.core.cache import cache

from spreedly.models import Plan, Subscription
from spreedly.pyspreedly.api import Client

def sync_plans():
    '''
    Sync subscription plans with Spreedly API
    '''
    client = Client(settings.SPREEDLY_AUTH_TOKEN_SECRET, settings.SPREEDLY_SITE_NAME)

    for plan in client.get_plans():
        p, created = Plan.objects.get_or_create(speedly_id=plan['speedly_id'])

        changed = False
        for k, v in plan.items():
            if hasattr(p, k) and not getattr(p, k) == v:
                setattr(p, k, v)
                changed = True
        if changed:
            p.save()

def get_subscription(user):
    cache_key = 'spreedly-subscription-%d' % user.id

    subscription = cache.get(cache_key)

    if not subscription:
        client = Client(settings.SPREEDLY_AUTH_TOKEN_SECRET, settings.SPREEDLY_SITE_NAME)
        data = client.get_info(user.id)

        subscription, created = Subscription.objects.get_or_create(
            user=user
        )
        for k, v in data.items():
            if hasattr(subscription, k):
                setattr(subscription, k, v)
        subscription.save()
        cache.set(cache_key, subscription, 60 * 60)

    return subscription

def create_subscription(user):
    client = Client(settings.SPREEDLY_AUTH_TOKEN_SECRET, settings.SPREEDLY_SITE_NAME)
    client.get_or_create_subscriber(user.id, user.username)
    return get_subscription(user)

def subscribe(plan, user):
    client = Client(settings.SPREEDLY_AUTH_TOKEN_SECRET, settings.SPREEDLY_SITE_NAME)
    client.subscribe(user.id, plan.id)
    return get_subscription(user)

def get_or_create_subscription(user):
    client = Client(settings.SPREEDLY_AUTH_TOKEN_SECRET, settings.SPREEDLY_SITE_NAME)
    data = client.get_or_create_subscriber(user.id, user.username)

    subscription, created = Subscription.objects.get_or_create(
        user=user
    )
    for k, v in data.items():
        if hasattr(subscription, k):
            setattr(subscription, k, v)
    subscription.save()
    return subscription

def check_trial_eligibility(plan, user):
    if plan.plan_type != 'free_trial':
        return False
    try:
        # make sure the user is trial eligable (they don't have a subscription yet, or they are trial_elegible)
        not_allowed = Subscription.objects.get(user=user, trial_elegible=False)
        return False
    except Subscription.DoesNotExist:
        return True

def start_free_trial(plan, user):
    if check_trial_eligibility(plan, user):
        client = Client(settings.SPREEDLY_AUTH_TOKEN_SECRET, settings.SPREEDLY_SITE_NAME)
        client.get_or_create_subscriber(user.id, user.username)
        client.subscribe(user.id, plan.pk, trial=True)
        get_subscription(user)
        return True
    else:
        return False

def complimentary_time_extension(user, duration_quantity, duration_units):
    client = Client(settings.SPREEDLY_AUTH_TOKEN_SECRET, settings.SPREEDLY_SITE_NAME)
    client.complimentary_time_extension(user.id, duration_quantity, duration_units)
    return get_subscription(user)

def complimentary_subscription(user, duration_quantity, duration_units, feature_level):
    client = Client(settings.SPREEDLY_AUTH_TOKEN_SECRET, settings.SPREEDLY_SITE_NAME)
    client.complimentary_subscription(user.id, duration_quantity, duration_units, feature_level)
    return get_subscription(user)

def lifetime_complimentary_subscription(user, feature_level):
    client = Client(settings.SPREEDLY_AUTH_TOKEN_SECRET, settings.SPREEDLY_SITE_NAME)
    client.lifetime_complimentary_subscription(user.id, feature_level)
    return get_subscription(user)

def add_store_credit(user, amount):
    client = Client(settings.SPREEDLY_AUTH_TOKEN_SECRET, settings.SPREEDLY_SITE_NAME)
    client.add_store_credit(user.id, amount)
    return get_subscription(user)

def stop_auto_renew(user):
    client = Client(settings.SPREEDLY_AUTH_TOKEN_SECRET, settings.SPREEDLY_SITE_NAME)
    client.stop_auto_renew(user.id)
    return get_subscription(user)

def return_url(user, plan=None, trial=False):
    args = [user.id]
    if plan:
        args.append(plan.pk)
    url = 'http://%s%s' % (Site.objects.get(id=settings.SITE_ID), reverse('spreedly_return', args=args))
    if trial:
        url = url + '?trial=true'
    return url

def allow_free_trial(user):
    client = Client(settings.SPREEDLY_AUTH_TOKEN_SECRET, settings.SPREEDLY_SITE_NAME)
    client.allow_free_trial(user.id)
    return get_subscription(user)

def subscription_url(plan, user, return_url):
    return 'https://spreedly.com/%(site_name)s/subscribers/%(user_id)s/subscribe/%(plan_id)s/%(user_username)s?email=%(user_email)s&return_url=%(return_url)s&first_name=%(first_name)s&last_name=%(last_name)s' % {
        'site_name': settings.SPREEDLY_SITE_NAME,
        'plan_id': plan.pk,
        'user_id': user.id,
        'user_username': user.username,
        'user_email': user.email,
        'first_name': user.first_name,
        'last_name': user.last_name,
        'return_url': return_url,
    }
<|endoftext|>


<commit_before>import os
from pathlib import Path
import sys

if 'VIRTUAL_ENV' in os.environ:
    major, minor = sys.version_info.major, sys.version_info.minor
    python = 'python{major}.{minor}'.format(major=major, minor=minor)
    envroot = Path(os.environ['VIRTUAL_ENV'])
    env_site_packages = envroot / 'lib' / python / 'site-packages'
    sys.path.append(str(env_site_packages))

import coverage
import nose

def main():
    python_path = str(Path(__file__).parent.parent)
    sys.path.append(python_path)
    nose.run(argv=[__file__])

if __name__ == '__main__':
    cov = coverage.coverage()
    cov.start()
    main()
    cov.stop()
    cov.save()
<commit_msg>Moved test environment checks in run_tests.py<commit_after>import os
from pathlib import Path
import sys

# Add path for modules installed by pip

if 'VIRTUAL_ENV' in os.environ:
    major, minor = sys.version_info.major, sys.version_info.minor
    python = 'python{major}.{minor}'.format(major=major, minor=minor)
    envroot = Path(os.environ['VIRTUAL_ENV'])
    env_site_packages = envroot / 'lib' / python / 'site-packages'
    sys.path.append(str(env_site_packages))

import coverage
import nose

# Add blender repo path for tests module

blender_addon_repo_root_path = str(Path(__file__).parent.parent)
sys.path.append(blender_addon_repo_root_path)

from tests.utils import PREVIZ_API_ROOT_ENVVAR, PREVIZ_API_TOKEN_ENVVAR

def error_message():
    return '\n'.join(
        [
            '\n',
            78*'*',
            'Testing environment not set. Set the variables :',
            '\t-{}'.format(PREVIZ_API_ROOT_ENVVAR),
            '\t-{}'.format(PREVIZ_API_TOKEN_ENVVAR),
            78*'*',
            '\n'
        ]
    )

def is_environment_valid():
    def envvar_is_valid(envvar):
        return len(os.environ.get(envvar, '')) > 0

    return envvar_is_valid(PREVIZ_API_ROOT_ENVVAR) \
       and envvar_is_valid(PREVIZ_API_TOKEN_ENVVAR)

def main():
    cov = coverage.coverage()
    cov.start()
    nose.run(argv=[__file__])
    cov.stop()
    cov.save()

if __name__ == '__main__':
    if not is_environment_valid():
        print(error_message())
        raise EnvironmentError
    main()
<|endoftext|>


<commit_before>from __future__ import division

from cStringIO import StringIO
from libtbx.phil import command_line
from libtbx.utils import Sorry
import iotbx.phil
# from dials.util.command_line import Importer
from dials.util.options import OptionParser
from dials.util.options import flatten_reflections
from dials.util.options import flatten_experiments
from dials.array_family import flex

help_message = '''

This program takes as input the output of dials.index, i.e. experiments.json
and indexed.pickle files. Full refinement of the crystal and experimental
geometry parameters will be performed (by default) in all Bravais settings
that are consistent with the input primitive unit cell. A table is printed
containing various information for each potential Bravais setting, including
the metric fit (a measure of the deviation from the triclinic cell),
the root-mean-square-deviations (rmsd), in mm, between the observed and
predicted spot centroids, the refined unit cell parameters in each Bravais
setting, and the change of basis operator to transform from the triclinic cell
to each Bravais setting.

The program also generates a .json file for each Bravais setting, e.g.
bravais_setting_1.json, which is equivalent to the input experiments.json, but
with the crystal model refined in the chosen Bravais setting. These
bravais_setting_*.json files are suitable as input to dials.refine or
dials.integrate, although the indexed.pickle file will need to be re-indexed
using dials.reindex if the change of basis operator (cb_op) for the chosen
Bravais setting is not the identity operator (a,b,c).

Examples::

  dials.refine_bravais_settings experiments.json indexed.pickle

  dials.refine_bravais_settings experiments.json indexed.pickle nproc=4

'''

phil_scope = iotbx.phil.parse("""
lepage_max_delta = 5
  .type = float
verbosity = 0
  .type = int(value_min=0)
nproc = Auto
  .type = int(value_min=1)
crystal_id = None
  .type = int(value_min=0)
normalise = False
  .type = bool
  .help = "Normalise intensities before calculating correlation coefficients."
normalise_bins = 0
  .type = int
  .help = "Number of resolution bins for normalisation"
cc_n_bins = None
  .type = int(value_min=1)
  .help = "Number of resolution bins to use for calculation of correlation coefficients"
output {
  directory = "."
    .type = str
  log = dials.refine_bravais_settings.log
    .type = str
  debug_log = dials.refine_bravais_settings.debug.log
    .type = str
}

include scope dials.algorithms.refinement.refiner.phil_scope
""", process_includes=True)

# override default refinement parameters
phil_scope = phil_scope.fetch(source=iotbx.phil.parse(
  """\
refinement {
  reflections {
    reflections_per_degree=100
  }
}
"""))

def bravais_lattice_to_space_groups(chiral_only=True):
  from cctbx import sgtbx
  from cctbx.sgtbx import bravais_types
  from libtbx.containers import OrderedDict
  bravais_lattice_to_sg = OrderedDict()
  for sgn in range(230):
    sg = sgtbx.space_group_info(number=sgn+1).group()
    if (not chiral_only) or (sg.is_chiral()):
      bravais_lattice = bravais_types.bravais_lattice(group=sg)
      bravais_lattice_to_sg.setdefault(str(bravais_lattice), [])
      bravais_lattice_to_sg[str(bravais_lattice)].append(sg)
  return bravais_lattice_to_sg

def bravais_lattice_to_space_group_table(chiral_only=True):
  from logging import info
  bravais_lattice_to_sg = bravais_lattice_to_space_groups(
    chiral_only=chiral_only)
  info('Chiral space groups corresponding to each Bravais lattice:')
  for bravais_lattice, space_groups in bravais_lattice_to_sg.iteritems():
    info(': '.join(
      [bravais_lattice,
       ' '.join(
         [sg.type().lookup_symbol().replace(' 1', '').replace(' ', '')
          for sg in space_groups])]))

def run(args):
  from dials.util import log
  from logging import info
  import libtbx.load_env
  usage = "%s experiments.json indexed.pickle [options]" %libtbx.env.dispatcher_name

  parser = OptionParser(
    usage=usage,
    phil=phil_scope,
    read_experiments=True,
    read_reflections=True,
    check_format=False,
    epilog=help_message)

  params, options = parser.parse_args(show_diff_phil=False)

  # Configure the logging
  log.config(info=params.output.log, debug=params.output.debug_log)

  from dials.util.version import dials_version
  info(dials_version())

  # Log the diff phil
  diff_phil = parser.diff_phil.as_str()
  if diff_phil is not '':
    info('The following parameters have been modified:\n')
    info(diff_phil)

  experiments = flatten_experiments(params.input.experiments)
  reflections = flatten_reflections(params.input.reflections)
  assert(len(reflections) == 1)
  reflections = reflections[0]

  if len(experiments) == 0:
    parser.print_help()
    return
  elif len(experiments.crystals()) > 1:
    if params.crystal_id is not None:
      assert params.crystal_id < len(experiments.crystals())
      experiment_ids = experiments.where(crystal=experiments.crystals()[params.crystal_id])
      from dxtbx.model.experiment.experiment_list import ExperimentList
      experiments = ExperimentList([experiments[i] for i in experiment_ids])
      refl_selections = [reflections['id'] == i for i in experiment_ids]
      reflections['id'] = flex.int(len(reflections), -1)
      for i, sel in enumerate(refl_selections):
        reflections['id'].set_selected(sel, i)
      reflections = reflections.select(reflections['id'] > -1)
    else:
      raise Sorry("Only one crystal can be processed at a time: set crystal_id to choose experiment.")

  if params.refinement.reflections.outlier.algorithm in ('auto', libtbx.Auto):
    if experiments[0].goniometer is None:
      params.refinement.reflections.outlier.algorithm = 'sauter_poon'
    else:
      # different default to dials.refine
      # tukey is faster and more appropriate at the indexing step
      params.refinement.reflections.outlier.algorithm = 'tukey'

  from dials.algorithms.indexing.symmetry \
       import refined_settings_factory_from_refined_triclinic

  bravais_lattice_to_space_group_table()

  cb_op_to_primitive = experiments[0].crystal.get_space_group().info()\
    .change_of_basis_op_to_primitive_setting()
  if experiments[0].crystal.get_space_group().n_ltr() > 1:
    effective_group = experiments[0].crystal.get_space_group()\
      .build_derived_reflection_intensity_group(anomalous_flag=True)
    sys_absent_flags = effective_group.is_sys_absent(
      reflections['miller_index'])
    reflections = reflections.select(~sys_absent_flags)
  experiments[0].crystal.update(experiments[0].crystal.change_basis(cb_op_to_primitive))
  miller_indices = reflections['miller_index']
  miller_indices = cb_op_to_primitive.apply(miller_indices)
  reflections['miller_index'] = miller_indices

  Lfat = refined_settings_factory_from_refined_triclinic(
    params, experiments, reflections, lepage_max_delta=params.lepage_max_delta,
    nproc=params.nproc, refiner_verbosity=params.verbosity)
  s = StringIO()
  Lfat.labelit_printout(out=s)
  info(s.getvalue())
  from json import dumps
  from os.path import join
  open(join(params.output.directory, 'bravais_summary.json'), 'wb').write(dumps(Lfat.as_dict()))
  from dxtbx.serialize import dump
  import copy
  for subgroup in Lfat:
    expts = copy.deepcopy(experiments)
    for expt in expts:
      expt.crystal.update(subgroup.refined_crystal)
      expt.detector = subgroup.detector
      expt.beam = subgroup.beam
    dump.experiment_list(
      expts, join(params.output.directory, 'bravais_setting_%i.json' % (int(subgroup.setting_number))))
  return

if __name__ == '__main__':
  import sys
  from libtbx.utils import show_times_at_exit
  show_times_at_exit()
  run(sys.argv[1:])
<commit_msg>P1 is a special case<commit_after>from __future__ import division

from cStringIO import StringIO
from libtbx.phil import command_line
from libtbx.utils import Sorry
import iotbx.phil
# from dials.util.command_line import Importer
from dials.util.options import OptionParser
from dials.util.options import flatten_reflections
from dials.util.options import flatten_experiments
from dials.array_family import flex

help_message = '''

This program takes as input the output of dials.index, i.e. experiments.json
and indexed.pickle files. Full refinement of the crystal and experimental
geometry parameters will be performed (by default) in all Bravais settings
that are consistent with the input primitive unit cell. A table is printed
containing various information for each potential Bravais setting, including
the metric fit (a measure of the deviation from the triclinic cell),
the root-mean-square-deviations (rmsd), in mm, between the observed and
predicted spot centroids, the refined unit cell parameters in each Bravais
setting, and the change of basis operator to transform from the triclinic cell
to each Bravais setting.

The program also generates a .json file for each Bravais setting, e.g.
bravais_setting_1.json, which is equivalent to the input experiments.json, but
with the crystal model refined in the chosen Bravais setting. These
bravais_setting_*.json files are suitable as input to dials.refine or
dials.integrate, although the indexed.pickle file will need to be re-indexed
using dials.reindex if the change of basis operator (cb_op) for the chosen
Bravais setting is not the identity operator (a,b,c).

Examples::

  dials.refine_bravais_settings experiments.json indexed.pickle

  dials.refine_bravais_settings experiments.json indexed.pickle nproc=4

'''

phil_scope = iotbx.phil.parse("""
lepage_max_delta = 5
  .type = float
verbosity = 0
  .type = int(value_min=0)
nproc = Auto
  .type = int(value_min=1)
crystal_id = None
  .type = int(value_min=0)
normalise = False
  .type = bool
  .help = "Normalise intensities before calculating correlation coefficients."
normalise_bins = 0
  .type = int
  .help = "Number of resolution bins for normalisation"
cc_n_bins = None
  .type = int(value_min=1)
  .help = "Number of resolution bins to use for calculation of correlation coefficients"
output {
  directory = "."
    .type = str
  log = dials.refine_bravais_settings.log
    .type = str
  debug_log = dials.refine_bravais_settings.debug.log
    .type = str
}

include scope dials.algorithms.refinement.refiner.phil_scope
""", process_includes=True)

# override default refinement parameters
phil_scope = phil_scope.fetch(source=iotbx.phil.parse(
  """\
refinement {
  reflections {
    reflections_per_degree=100
  }
}
"""))

def bravais_lattice_to_space_groups(chiral_only=True):
  from cctbx import sgtbx
  from cctbx.sgtbx import bravais_types
  from libtbx.containers import OrderedDict
  bravais_lattice_to_sg = OrderedDict()
  for sgn in range(230):
    sg = sgtbx.space_group_info(number=sgn+1).group()
    if (not chiral_only) or (sg.is_chiral()):
      bravais_lattice = bravais_types.bravais_lattice(group=sg)
      bravais_lattice_to_sg.setdefault(str(bravais_lattice), [])
      bravais_lattice_to_sg[str(bravais_lattice)].append(sg)
  return bravais_lattice_to_sg

def bravais_lattice_to_space_group_table(chiral_only=True):
  from logging import info
  bravais_lattice_to_sg = bravais_lattice_to_space_groups(
    chiral_only=chiral_only)
  info('Chiral space groups corresponding to each Bravais lattice:')
  for bravais_lattice, space_groups in bravais_lattice_to_sg.iteritems():
    info(': '.join(
      [bravais_lattice,
       ' '.join(
         [short_space_group_name(sg) for sg in space_groups])]))

def short_space_group_name(space_group):
  sgt = space_group.type()
  symbol = sgt.lookup_symbol()
  if sgt.number() > 1:
    symbol = symbol.replace(' 1', '')
  return symbol.replace(' ', '')

def run(args):
  from dials.util import log
  from logging import info
  import libtbx.load_env
  usage = "%s experiments.json indexed.pickle [options]" %libtbx.env.dispatcher_name

  parser = OptionParser(
    usage=usage,
    phil=phil_scope,
    read_experiments=True,
    read_reflections=True,
    check_format=False,
    epilog=help_message)

  params, options = parser.parse_args(show_diff_phil=False)

  # Configure the logging
  log.config(info=params.output.log, debug=params.output.debug_log)

  from dials.util.version import dials_version
  info(dials_version())

  # Log the diff phil
  diff_phil = parser.diff_phil.as_str()
  if diff_phil is not '':
    info('The following parameters have been modified:\n')
    info(diff_phil)

  experiments = flatten_experiments(params.input.experiments)
  reflections = flatten_reflections(params.input.reflections)
  assert(len(reflections) == 1)
  reflections = reflections[0]

  if len(experiments) == 0:
    parser.print_help()
    return
  elif len(experiments.crystals()) > 1:
    if params.crystal_id is not None:
      assert params.crystal_id < len(experiments.crystals())
      experiment_ids = experiments.where(crystal=experiments.crystals()[params.crystal_id])
      from dxtbx.model.experiment.experiment_list import ExperimentList
      experiments = ExperimentList([experiments[i] for i in experiment_ids])
      refl_selections = [reflections['id'] == i for i in experiment_ids]
      reflections['id'] = flex.int(len(reflections), -1)
      for i, sel in enumerate(refl_selections):
        reflections['id'].set_selected(sel, i)
      reflections = reflections.select(reflections['id'] > -1)
    else:
      raise Sorry("Only one crystal can be processed at a time: set crystal_id to choose experiment.")

  if params.refinement.reflections.outlier.algorithm in ('auto', libtbx.Auto):
    if experiments[0].goniometer is None:
      params.refinement.reflections.outlier.algorithm = 'sauter_poon'
    else:
      # different default to dials.refine
      # tukey is faster and more appropriate at the indexing step
      params.refinement.reflections.outlier.algorithm = 'tukey'

  from dials.algorithms.indexing.symmetry \
       import refined_settings_factory_from_refined_triclinic

  bravais_lattice_to_space_group_table()

  cb_op_to_primitive = experiments[0].crystal.get_space_group().info()\
    .change_of_basis_op_to_primitive_setting()
  if experiments[0].crystal.get_space_group().n_ltr() > 1:
    effective_group = experiments[0].crystal.get_space_group()\
      .build_derived_reflection_intensity_group(anomalous_flag=True)
    sys_absent_flags = effective_group.is_sys_absent(
      reflections['miller_index'])
    reflections = reflections.select(~sys_absent_flags)
  experiments[0].crystal.update(experiments[0].crystal.change_basis(cb_op_to_primitive))
  miller_indices = reflections['miller_index']
  miller_indices = cb_op_to_primitive.apply(miller_indices)
  reflections['miller_index'] = miller_indices

  Lfat = refined_settings_factory_from_refined_triclinic(
    params, experiments, reflections, lepage_max_delta=params.lepage_max_delta,
    nproc=params.nproc, refiner_verbosity=params.verbosity)
  s = StringIO()
  Lfat.labelit_printout(out=s)
  info(s.getvalue())
  from json import dumps
  from os.path import join
  open(join(params.output.directory, 'bravais_summary.json'), 'wb').write(dumps(Lfat.as_dict()))
  from dxtbx.serialize import dump
  import copy
  for subgroup in Lfat:
    expts = copy.deepcopy(experiments)
    for expt in expts:
      expt.crystal.update(subgroup.refined_crystal)
      expt.detector = subgroup.detector
      expt.beam = subgroup.beam
    dump.experiment_list(
      expts, join(params.output.directory, 'bravais_setting_%i.json' % (int(subgroup.setting_number))))
  return

if __name__ == '__main__':
  import sys
  from libtbx.utils import show_times_at_exit
  show_times_at_exit()
  run(sys.argv[1:])
<|endoftext|>


<commit_before># nhlib: A New Hazard Library
# Copyright (C) 2012 GEM Foundation
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
import unittest

import numpy

from nhlib.site import Site, SiteCollection
from nhlib.geo.point import Point

class SiteTestCase(unittest.TestCase):
    def _assert_creation(self, error=None, **kwargs):
        default_kwargs = {
            'location': Point(10, 20),
            'vs30': 10,
            'vs30measured': False,
            'z1pt0': 20,
            'z2pt5': 30
        }
        default_kwargs.update(kwargs)
        kwargs = default_kwargs
        if error is not None:
            with self.assertRaises(ValueError) as ar:
                Site(**kwargs)
            self.assertEqual(ar.exception.message, error)
        else:
            site = Site(**kwargs)
            for attr in kwargs:
                self.assertEqual(getattr(site, attr), kwargs[attr])

    def test_wrong_vs30(self):
        error = 'vs30 must be positive'
        self._assert_creation(error=error, vs30=0)
        self._assert_creation(error=error, vs30=-1)

    def test_wrong_z1pt0(self):
        error = 'z1pt0 must be positive'
        self._assert_creation(error=error, z1pt0=0)
        self._assert_creation(error=error, z1pt0=-1)

    def test_wrong_z2pt5(self):
        error = 'z2pt5 must be positive'
        self._assert_creation(error=error, z2pt5=0)
        self._assert_creation(error=error, z2pt5=-1)

    def test_successful_creation(self):
        self._assert_creation()

class SiteCollectionCreationTestCase(unittest.TestCase):
    def test(self):
        s1 = Site(location=Point(10, 20, 30),
                  vs30=1.2, vs30measured=True,
                  z1pt0=3.4, z2pt5=5.6)
        s2 = Site(location=Point(-1.2, -3.4, -5.6),
                  vs30=55.4, vs30measured=False,
                  z1pt0=66.7, z2pt5=88.9)
        cll = SiteCollection([s1, s2])
        self.assertTrue((cll.vs30 == [1.2, 55.4]).all())
        self.assertTrue((cll.vs30measured == [True, False]).all())
        self.assertTrue((cll.z1pt0 == [3.4, 66.7]).all())
        self.assertTrue((cll.z2pt5 == [5.6, 88.9]).all())
        self.assertTrue((cll.mesh.lons == [10, -1.2]).all())
        self.assertTrue((cll.mesh.lats == [20, -3.4]).all())
        self.assertIs(cll.mesh.depths, None)
        for arr in (cll.vs30, cll.z1pt0, cll.z2pt5):
            self.assertIsInstance(arr, numpy.ndarray)
            self.assertEqual(arr.flags.writeable, False)
            self.assertEqual(arr.dtype, float)
        self.assertIsInstance(cll.vs30measured, numpy.ndarray)
        self.assertEqual(cll.vs30measured.flags.writeable, False)
        self.assertEqual(cll.vs30measured.dtype, bool)

        self.assertEqual(len(cll), 2)

class SiteCollectionFilterTestCase(unittest.TestCase):
    SITES = [
        Site(location=Point(10, 20, 30), vs30=1.2, vs30measured=True,
             z1pt0=3, z2pt5=5),
        Site(location=Point(11, 12, 13), vs30=55.4, vs30measured=False,
             z1pt0=6, z2pt5=8),
        Site(location=Point(0, 2, 0), vs30=2, vs30measured=True,
             z1pt0=9, z2pt5=17),
        Site(location=Point(1, 1, 3), vs30=4, vs30measured=False,
             z1pt0=22, z2pt5=11)
    ]

    def test_filter(self):
        col = SiteCollection(self.SITES)
        filtered = col.filter(numpy.array([True, False, True, False]))
        self.assertIsInstance(filtered, SiteCollection)
        arreq = numpy.testing.assert_array_equal
        arreq(filtered.vs30, [1.2, 2])
        arreq(filtered.vs30measured, [True, True])
        arreq(filtered.z1pt0, [3, 9])
        arreq(filtered.z2pt5, [5, 17])
        arreq(filtered.mesh.lons, [10, 0])
        arreq(filtered.mesh.lats, [20, 2])
        self.assertIs(filtered.mesh.depths, None)

        filtered = col.filter(numpy.array([False, True, True, True]))
        self.assertIsInstance(filtered, SiteCollection)
        arreq(filtered.vs30, [55.4, 2, 4])
        arreq(filtered.vs30measured, [False, True, False])
        arreq(filtered.z1pt0, [6, 9, 22])
        arreq(filtered.z2pt5, [8, 17, 11])
        arreq(filtered.mesh.lons, [11, 0, 1])
        arreq(filtered.mesh.lats, [12, 2, 1])
        self.assertIs(filtered.mesh.depths, None)

    def test_filter_all_out(self):
        col = SiteCollection(self.SITES)
        filtered = col.filter(numpy.zeros(len(self.SITES), bool))
        self.assertIs(filtered, None)

    def test_filter_all_in(self):
        col = SiteCollection(self.SITES)
        filtered = col.filter(numpy.ones(len(self.SITES), bool))
        self.assertIs(filtered, col)

    def test_double_filter(self):
        col = SiteCollection(self.SITES)
        filtered = col.filter(numpy.array([True, False, True, True]))
        filtered2 = filtered.filter(numpy.array([False, True, False]))
        arreq = numpy.testing.assert_array_equal
        arreq(filtered2.vs30, [2])
        arreq(filtered2.vs30measured, [True])
        arreq(filtered2.z1pt0, [9])
        arreq(filtered2.z2pt5, [17])
        arreq(filtered2.mesh.lons, [0])
        arreq(filtered2.mesh.lats, [2])
        self.assertIs(filtered2.mesh.depths, None)
        arreq(filtered.indices, [0, 2, 3])
        arreq(filtered2.indices, [2])
        filtered2 = filtered.filter(numpy.array([True, False, True]))
        arreq(filtered2.indices, [0, 3])

    def test_expand_2d(self):
        col = SiteCollection(self.SITES)
        col.indices = numpy.array([1, 3, 5, 6])
        data_condensed = numpy.array([
            [1, 2, 3],
            [5, 6, 7],
            [10, 11, 12],
            [15, 16, 17]
        ])
        data_expanded = col.expand(data_condensed, total_sites=8,
                                   placeholder=-1)
        data_expanded_expected = numpy.array([
            [-1, -1, -1],
            [1, 2, 3],
            [-1, -1, -1],
            [5, 6, 7],
            [-1, -1, -1],
            [10, 11, 12],
            [15, 16, 17],
            [-1, -1, -1]
        ])
        numpy.testing.assert_array_equal(data_expanded, data_expanded_expected)

    def test_expand_1d(self):
        col = SiteCollection(self.SITES)
        col = col.filter(numpy.array([1, 0, 1, 1]))
        data_condensed = numpy.array([5, 6, 7])
        data_expanded = col.expand(data_condensed, total_sites=4,
                                   placeholder=100)
        data_expanded_expected = numpy.array([5, 100, 6, 7])
        numpy.testing.assert_array_equal(data_expanded, data_expanded_expected)

    def test_expand_no_filtering(self):
        col = SiteCollection(self.SITES)
        data_condensed = numpy.array([3, 2, 1, 0])
        data_expanded = col.expand(data_condensed, total_sites=4,
                                   placeholder=100)
        data_expanded_expected = data_condensed
        numpy.testing.assert_array_equal(data_expanded, data_expanded_expected)
<commit_msg>tests/site: added a test for pickling Sites<commit_after># nhlib: A New Hazard Library
# Copyright (C) 2012 GEM Foundation
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
import pickle
import unittest

import numpy

from nhlib.site import Site, SiteCollection
from nhlib.geo.point import Point

class SiteTestCase(unittest.TestCase):
    def _assert_creation(self, error=None, **kwargs):
        default_kwargs = {
            'location': Point(10, 20),
            'vs30': 10,
            'vs30measured': False,
            'z1pt0': 20,
            'z2pt5': 30
        }
        default_kwargs.update(kwargs)
        kwargs = default_kwargs
        if error is not None:
            with self.assertRaises(ValueError) as ar:
                Site(**kwargs)
            self.assertEqual(ar.exception.message, error)
        else:
            site = Site(**kwargs)
            for attr in kwargs:
                self.assertEqual(getattr(site, attr), kwargs[attr])

    def test_wrong_vs30(self):
        error = 'vs30 must be positive'
        self._assert_creation(error=error, vs30=0)
        self._assert_creation(error=error, vs30=-1)

    def test_wrong_z1pt0(self):
        error = 'z1pt0 must be positive'
        self._assert_creation(error=error, z1pt0=0)
        self._assert_creation(error=error, z1pt0=-1)

    def test_wrong_z2pt5(self):
        error = 'z2pt5 must be positive'
        self._assert_creation(error=error, z2pt5=0)
        self._assert_creation(error=error, z2pt5=-1)

    def test_successful_creation(self):
        self._assert_creation()

class SiteCollectionCreationTestCase(unittest.TestCase):
    def test(self):
        s1 = Site(location=Point(10, 20, 30),
                  vs30=1.2, vs30measured=True,
                  z1pt0=3.4, z2pt5=5.6)
        s2 = Site(location=Point(-1.2, -3.4, -5.6),
                  vs30=55.4, vs30measured=False,
                  z1pt0=66.7, z2pt5=88.9)
        cll = SiteCollection([s1, s2])
        self.assertTrue((cll.vs30 == [1.2, 55.4]).all())
        self.assertTrue((cll.vs30measured == [True, False]).all())
        self.assertTrue((cll.z1pt0 == [3.4, 66.7]).all())
        self.assertTrue((cll.z2pt5 == [5.6, 88.9]).all())
        self.assertTrue((cll.mesh.lons == [10, -1.2]).all())
        self.assertTrue((cll.mesh.lats == [20, -3.4]).all())
        self.assertIs(cll.mesh.depths, None)
        for arr in (cll.vs30, cll.z1pt0, cll.z2pt5):
            self.assertIsInstance(arr, numpy.ndarray)
            self.assertEqual(arr.flags.writeable, False)
            self.assertEqual(arr.dtype, float)
        self.assertIsInstance(cll.vs30measured, numpy.ndarray)
        self.assertEqual(cll.vs30measured.flags.writeable, False)
        self.assertEqual(cll.vs30measured.dtype, bool)

        self.assertEqual(len(cll), 2)

class SiteCollectionFilterTestCase(unittest.TestCase):
    SITES = [
        Site(location=Point(10, 20, 30), vs30=1.2, vs30measured=True,
             z1pt0=3, z2pt5=5),
        Site(location=Point(11, 12, 13), vs30=55.4, vs30measured=False,
             z1pt0=6, z2pt5=8),
        Site(location=Point(0, 2, 0), vs30=2, vs30measured=True,
             z1pt0=9, z2pt5=17),
        Site(location=Point(1, 1, 3), vs30=4, vs30measured=False,
             z1pt0=22, z2pt5=11)
    ]

    def test_filter(self):
        col = SiteCollection(self.SITES)
        filtered = col.filter(numpy.array([True, False, True, False]))
        self.assertIsInstance(filtered, SiteCollection)
        arreq = numpy.testing.assert_array_equal
        arreq(filtered.vs30, [1.2, 2])
        arreq(filtered.vs30measured, [True, True])
        arreq(filtered.z1pt0, [3, 9])
        arreq(filtered.z2pt5, [5, 17])
        arreq(filtered.mesh.lons, [10, 0])
        arreq(filtered.mesh.lats, [20, 2])
        self.assertIs(filtered.mesh.depths, None)

        filtered = col.filter(numpy.array([False, True, True, True]))
        self.assertIsInstance(filtered, SiteCollection)
        arreq(filtered.vs30, [55.4, 2, 4])
        arreq(filtered.vs30measured, [False, True, False])
        arreq(filtered.z1pt0, [6, 9, 22])
        arreq(filtered.z2pt5, [8, 17, 11])
        arreq(filtered.mesh.lons, [11, 0, 1])
        arreq(filtered.mesh.lats, [12, 2, 1])
        self.assertIs(filtered.mesh.depths, None)

    def test_filter_all_out(self):
        col = SiteCollection(self.SITES)
        filtered = col.filter(numpy.zeros(len(self.SITES), bool))
        self.assertIs(filtered, None)

    def test_filter_all_in(self):
        col = SiteCollection(self.SITES)
        filtered = col.filter(numpy.ones(len(self.SITES), bool))
        self.assertIs(filtered, col)

    def test_double_filter(self):
        col = SiteCollection(self.SITES)
        filtered = col.filter(numpy.array([True, False, True, True]))
        filtered2 = filtered.filter(numpy.array([False, True, False]))
        arreq = numpy.testing.assert_array_equal
        arreq(filtered2.vs30, [2])
        arreq(filtered2.vs30measured, [True])
        arreq(filtered2.z1pt0, [9])
        arreq(filtered2.z2pt5, [17])
        arreq(filtered2.mesh.lons, [0])
        arreq(filtered2.mesh.lats, [2])
        self.assertIs(filtered2.mesh.depths, None)
        arreq(filtered.indices, [0, 2, 3])
        arreq(filtered2.indices, [2])
        filtered2 = filtered.filter(numpy.array([True, False, True]))
        arreq(filtered2.indices, [0, 3])

    def test_expand_2d(self):
        col = SiteCollection(self.SITES)
        col.indices = numpy.array([1, 3, 5, 6])
        data_condensed = numpy.array([
            [1, 2, 3],
            [5, 6, 7],
            [10, 11, 12],
            [15, 16, 17]
        ])
        data_expanded = col.expand(data_condensed, total_sites=8,
                                   placeholder=-1)
        data_expanded_expected = numpy.array([
            [-1, -1, -1],
            [1, 2, 3],
            [-1, -1, -1],
            [5, 6, 7],
            [-1, -1, -1],
            [10, 11, 12],
            [15, 16, 17],
            [-1, -1, -1]
        ])
        numpy.testing.assert_array_equal(data_expanded, data_expanded_expected)

    def test_expand_1d(self):
        col = SiteCollection(self.SITES)
        col = col.filter(numpy.array([1, 0, 1, 1]))
        data_condensed = numpy.array([5, 6, 7])
        data_expanded = col.expand(data_condensed, total_sites=4,
                                   placeholder=100)
        data_expanded_expected = numpy.array([5, 100, 6, 7])
        numpy.testing.assert_array_equal(data_expanded, data_expanded_expected)

    def test_expand_no_filtering(self):
        col = SiteCollection(self.SITES)
        data_condensed = numpy.array([3, 2, 1, 0])
        data_expanded = col.expand(data_condensed, total_sites=4,
                                   placeholder=100)
        data_expanded_expected = data_condensed
        numpy.testing.assert_array_equal(data_expanded, data_expanded_expected)

class SitePickleTest(unittest.TestCase):
    # Tests for pickling Sites.

    def test_dumps_and_loads(self):
        point = Point(1, 2, 3)
        site1 = Site(point, 760.0, True, 100.0, 5.0)
        site2 = pickle.loads(pickle.dumps(site1))

        self.assertEqual(site1, site2)
<|endoftext|>