<commit_before># coding: utf-8

from __future__ import division, unicode_literals

"""
Defines SymmetryGroup parent class and PointGroup and SpaceGroup classes.
Shyue Ping Ong thanks Marc De Graef for his generous sharing of his
SpaceGroup data as published in his textbook "Structure of Materials".
"""

__author__ = "Shyue Ping Ong"
__copyright__ = "Copyright 2013, The Materials Virtual Lab"
__version__ = "0.1"
__maintainer__ = "Shyue Ping Ong"
__email__ = "ongsp@ucsd.edu"
__date__ = "4/4/14"

import os
from itertools import product
from fractions import Fraction

import numpy as np

import json

with open(os.path.join(os.path.dirname(__file__), "symm_data.json")) as f:
    SYMM_DATA = json.load(f)

GENERATOR_MATRICES = SYMM_DATA["generator_matrices"]
POINT_GROUP_ENC = SYMM_DATA["point_group_encoding"]
SPACE_GROUP_ENC = SYMM_DATA["space_group_encoding"]
TRANSLATIONS = {k: Fraction(v) for k, v in SYMM_DATA["translations"].items()}

class SymmetryGroup(object):

    def get_orbit(self, p, tol=1e-5):
        """
        Returns the orbit for a point.

        Args:
            p: Point as a 3x1 array.
            tol: Tolerance for determining if sites are the same. 1e-5 should
                be sufficient for most purposes. Set to 0 for exact matching
                (and also needed for symbolic orbits).

        Returns:
            ([array]) Orbit for point.
        """
        orbit = []
        for o in self.symmetry_ops:
            pp = np.dot(o, p)
            pp = np.mod(pp, 1)
            if not in_array_list(orbit, pp, tol=tol):
                orbit.append(pp)
        return orbit

class PointGroup(SymmetryGroup):
    """
    Class representing a Point Group, with generators and symmetry operations.

    .. attribute:: symbol

        Full International or Hermann-Mauguin Symbol.

    .. attribute:: generators

        List of generator matrices. Note that 3x3 matrices are used for Point
        Groups.

    .. attribute:: symmetry_ops

        Full set of symmetry operations as matrices.
    """

    def __init__(self, int_symbol):
        """
        Initializes a Point Group from its international symbol.

        Args:
            int_symbol (str): International or Hermann-Mauguin Symbol.
        """
        self.symbol = int_symbol
        self.generators = [GENERATOR_MATRICES[c]
                           for c in POINT_GROUP_ENC[int_symbol]]
        self.symmetry_ops = self._generate_full_symmetry_ops()
        self.order = len(self.symmetry_ops)

    def _generate_full_symmetry_ops(self):
        symm_ops = list(self.generators)
        new_ops = self.generators
        while len(new_ops) > 0:
            gen_ops = []
            for g1, g2 in product(new_ops, symm_ops):
                op = np.dot(g1, g2)
                if not in_array_list(symm_ops, op):
                    gen_ops.append(op)
                    symm_ops.append(op)
            new_ops = gen_ops
        return symm_ops

class SpaceGroup(SymmetryGroup):
    """
    Class representing a SpaceGroup.

    .. attribute:: symbol

        Full International or Hermann-Mauguin Symbol.

    .. attribute:: int_number

        International number

    .. attribute:: generators

        List of generator matrices. Note that 4x4 matrices are used for Space
        Groups.

    .. attribute:: order

        Order of Space Group
    """

    #Contains the entire list of supported Space Group symbols.
    SG_SYMBOLS = tuple(SPACE_GROUP_ENC.keys())

    def __init__(self, int_symbol):
        """
        Initializes a Space Group from its *full* international symbol.

        Args:
            int_symbol (str): Full International or Hermann-Mauguin Symbol.
                The notation is a LaTeX-like string, with screw axes being
                represented by an underscore. For example, "P6_3/mmc". Note
                that for rhomohedral cells, the hexagonal setting can be
                accessed by adding a "H", e.g., "R-3mH".
        """
        self.symbol = int_symbol
        # TODO: Support different origin choices.
        enc = list(SPACE_GROUP_ENC[int_symbol]["enc"])
        inversion = int(enc.pop(0))
        ngen = int(enc.pop(0))
        symm_ops = [np.eye(4)]
        if inversion:
            symm_ops.append(np.array(
                [[-1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0],
                 [0, 0, 0, 1]]))
        for i in range(ngen):
            m = np.eye(4)
            m[:3, :3] = GENERATOR_MATRICES[enc.pop(0)]
            m[0, 3] = TRANSLATIONS[enc.pop(0)]
            m[1, 3] = TRANSLATIONS[enc.pop(0)]
            m[2, 3] = TRANSLATIONS[enc.pop(0)]
            symm_ops.append(m)
        self.generators = symm_ops
        self.int_number = SPACE_GROUP_ENC[int_symbol]["int_number"]
        self.order = SPACE_GROUP_ENC[int_symbol]["order"]
        self._symmetry_ops = None

    def _generate_full_symmetry_ops(self):
        symm_ops = np.array(self.generators)
        for op in symm_ops:
            op[0:3, 3] = np.mod(op[0:3, 3], 1)
        new_ops = symm_ops
        while len(new_ops) > 0 and len(symm_ops) < self.order:
            gen_ops = []
            for g in new_ops:
                temp_ops = np.einsum('ijk,kl', symm_ops, g)
                for op in temp_ops:
                    op[0:3, 3] = np.mod(op[0:3, 3], 1)
                    ind = np.where(np.abs(1 - op[0:3, 3]) < 1e-5)
                    op[ind, 3] = 0
                    if not in_array_list(symm_ops, op):
                        gen_ops.append(op)
                        symm_ops = np.append(symm_ops, [op], axis=0)
            new_ops = gen_ops
        assert len(symm_ops) == self.order
        return symm_ops

    @property
    def symmetry_ops(self):
        """
        Full set of symmetry operations as matrices. Lazily initialized as
        generation sometimes takes a bit of time.
        """
        if self._symmetry_ops is None:
            self._symmetry_ops = self._generate_full_symmetry_ops()
        return self._symmetry_ops

    @property
    def crystal_system(self):
        i = self.int_number
        if i <= 2:
            return "triclinic"
        elif i <= 15:
            return "monoclinic"
        elif i <= 74:
            return "orthorhombic"
        elif i <= 142:
            return "tetragonal"
        elif i <= 167:
            return "trigonal"
        elif i <= 194:
            return "hexagonal"
        else:
            return "cubic"

    def get_orbit(self, p, tol=1e-5):
        """
        Returns the orbit for a point.

        Args:
            p: Point as a 3x1 array.

        Returns:
            ([array]) Orbit for point.
        """
        p = np.append(p, [1])
        orbit = super(SpaceGroup, self).get_orbit(p, tol=tol)
        return np.delete(orbit, np.s_[-1:], 1)

    @classmethod
    def from_int_number(cls, int_number, hexagonal=True):
        """
        Obtains a SpaceGroup from its international number.

        Args:
            int_number (int): International number.
            hexagonal (bool): For rhombohedral groups, whether to return the
                hexagonal setting (default) or rhombohedral setting.

        Returns:
            (SpaceGroup)
        """
        return SpaceGroup(sg_symbol_from_int_number(int_number,
                                                    hexagonal=hexagonal))

    def __str__(self):
        return "Spacegroup %s with international number %d and order %d" % (
            self.symbol, self.int_number, len(self.symmetry_ops))

def sg_symbol_from_int_number(int_number, hexagonal=True):
    """
    Obtains a SpaceGroup name from its international number.

    Args:
        int_number (int): International number.
        hexagonal (bool): For rhombohedral groups, whether to return the
            hexagonal setting (default) or rhombohedral setting.

    Returns:
        (str) Spacegroup symbol
    """
    syms = []
    for n, v in SPACE_GROUP_ENC.items():
        if v["int_number"] == int_number:
            syms.append(n)
    if len(syms) == 0:
        raise ValueError("Invalid international number!")
    if len(syms) == 2:
        if hexagonal:
            syms = list(filter(lambda s: s.endswith("H"), syms))
        else:
            syms = list(filter(lambda s: not s.endswith("H"), syms))
    return syms.pop()

def in_array_list(array_list, a, tol=1e-5):
    """
    Extremely efficient nd-array comparison using numpy's broadcasting. This
    function checks if a particular array a, is present in a list of arrays.
    It works for arrays of any size, e.g., even matrix searches.

    Args:
        array_list ([array]): A list of arrays to compare to.
        a (array): The test array for comparison.
        tol (float): The tolerance. Defaults to 1e-5. If 0, an exact match is
            done.

    Returns:
        (bool)
    """
    if len(array_list) == 0:
        return False
    axes = tuple(range(1, a.ndim + 1))
    if not tol:
        return np.any(np.all(np.equal(array_list, a[None, :]), axes))
    else:
        return np.any(np.sum(np.abs(array_list - a[None, :]), axes) < tol)
<commit_msg>Use SymmOP instead of just simple matrices.<commit_after># coding: utf-8

from __future__ import division, unicode_literals

"""
Defines SymmetryGroup parent class and PointGroup and SpaceGroup classes.
Shyue Ping Ong thanks Marc De Graef for his generous sharing of his
SpaceGroup data as published in his textbook "Structure of Materials".
"""

__author__ = "Shyue Ping Ong"
__copyright__ = "Copyright 2013, The Materials Virtual Lab"
__version__ = "0.1"
__maintainer__ = "Shyue Ping Ong"
__email__ = "ongsp@ucsd.edu"
__date__ = "4/4/14"

import os
from itertools import product
from fractions import Fraction
import numpy as np
import json

from pymatgen.core.operations import SymmOp

with open(os.path.join(os.path.dirname(__file__), "symm_data.json")) as f:
    SYMM_DATA = json.load(f)

GENERATOR_MATRICES = SYMM_DATA["generator_matrices"]
POINT_GROUP_ENC = SYMM_DATA["point_group_encoding"]
SPACE_GROUP_ENC = SYMM_DATA["space_group_encoding"]
TRANSLATIONS = {k: Fraction(v) for k, v in SYMM_DATA["translations"].items()}

class SymmetryGroup(object):
    """
    This is simply a parent class for a symmetry group to make it easier to
    share code that is common between any group.
    """

    def get_orbit(self, p, tol=1e-5):
        """
        Returns the orbit for a point.

        Args:
            p: Point as a 3x1 array.
            tol: Tolerance for determining if sites are the same. 1e-5 should
                be sufficient for most purposes. Set to 0 for exact matching
                (and also needed for symbolic orbits).

        Returns:
            ([array]) Orbit for point.
        """
        orbit = []
        for o in self.symmetry_ops:
            pp = o.operate(p)
            pp = np.mod(pp, 1)
            if not in_array_list(orbit, pp, tol=tol):
                orbit.append(pp)
        return orbit

class PointGroup(SymmetryGroup):
    """
    Class representing a Point Group, with generators and symmetry operations.

    .. attribute:: symbol

        Full International or Hermann-Mauguin Symbol.

    .. attribute:: generators

        List of generator matrices. Note that 3x3 matrices are used for Point
        Groups.

    .. attribute:: symmetry_ops

        Full set of symmetry operations as matrices.
    """

    def __init__(self, int_symbol):
        """
        Initializes a Point Group from its international symbol.

        Args:
            int_symbol (str): International or Hermann-Mauguin Symbol.
        """
        self.symbol = int_symbol
        self.generators = [GENERATOR_MATRICES[c]
                           for c in POINT_GROUP_ENC[int_symbol]]
        self.symmetry_ops = [SymmOp.from_rotation_and_translation(m)
                             for m in self._generate_full_symmetry_ops()]
        self.order = len(self.symmetry_ops)

    def _generate_full_symmetry_ops(self):
        symm_ops = list(self.generators)
        new_ops = self.generators
        while len(new_ops) > 0:
            gen_ops = []
            for g1, g2 in product(new_ops, symm_ops):
                op = np.dot(g1, g2)
                if not in_array_list(symm_ops, op):
                    gen_ops.append(op)
                    symm_ops.append(op)
            new_ops = gen_ops
        return symm_ops

class SpaceGroup(SymmetryGroup):
    """
    Class representing a SpaceGroup.

    .. attribute:: symbol

        Full International or Hermann-Mauguin Symbol.

    .. attribute:: int_number

        International number

    .. attribute:: generators

        List of generator matrices. Note that 4x4 matrices are used for Space
        Groups.

    .. attribute:: order

        Order of Space Group
    """

    #Contains the entire list of supported Space Group symbols.
    SG_SYMBOLS = tuple(SPACE_GROUP_ENC.keys())

    def __init__(self, int_symbol):
        """
        Initializes a Space Group from its *full* international symbol.

        Args:
            int_symbol (str): Full International or Hermann-Mauguin Symbol.
                The notation is a LaTeX-like string, with screw axes being
                represented by an underscore. For example, "P6_3/mmc". Note
                that for rhomohedral cells, the hexagonal setting can be
                accessed by adding a "H", e.g., "R-3mH".
        """
        self.symbol = int_symbol
        # TODO: Support different origin choices.
        enc = list(SPACE_GROUP_ENC[int_symbol]["enc"])
        inversion = int(enc.pop(0))
        ngen = int(enc.pop(0))
        symm_ops = [np.eye(4)]
        if inversion:
            symm_ops.append(np.array(
                [[-1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0],
                 [0, 0, 0, 1]]))
        for i in range(ngen):
            m = np.eye(4)
            m[:3, :3] = GENERATOR_MATRICES[enc.pop(0)]
            m[0, 3] = TRANSLATIONS[enc.pop(0)]
            m[1, 3] = TRANSLATIONS[enc.pop(0)]
            m[2, 3] = TRANSLATIONS[enc.pop(0)]
            symm_ops.append(m)
        self.generators = symm_ops
        self.int_number = SPACE_GROUP_ENC[int_symbol]["int_number"]
        self.order = SPACE_GROUP_ENC[int_symbol]["order"]
        self._symmetry_ops = None

    def _generate_full_symmetry_ops(self):
        symm_ops = np.array(self.generators)
        for op in symm_ops:
            op[0:3, 3] = np.mod(op[0:3, 3], 1)
        new_ops = symm_ops
        while len(new_ops) > 0 and len(symm_ops) < self.order:
            gen_ops = []
            for g in new_ops:
                temp_ops = np.einsum('ijk,kl', symm_ops, g)
                for op in temp_ops:
                    op[0:3, 3] = np.mod(op[0:3, 3], 1)
                    ind = np.where(np.abs(1 - op[0:3, 3]) < 1e-5)
                    op[ind, 3] = 0
                    if not in_array_list(symm_ops, op):
                        gen_ops.append(op)
                        symm_ops = np.append(symm_ops, [op], axis=0)
            new_ops = gen_ops
        assert len(symm_ops) == self.order
        return symm_ops

    @property
    def symmetry_ops(self):
        """
        Full set of symmetry operations as matrices. Lazily initialized as
        generation sometimes takes a bit of time.
        """
        if self._symmetry_ops is None:
            self._symmetry_ops = [
                SymmOp(m) for m in self._generate_full_symmetry_ops()]
        return self._symmetry_ops

    @property
    def crystal_system(self):
        i = self.int_number
        if i <= 2:
            return "triclinic"
        elif i <= 15:
            return "monoclinic"
        elif i <= 74:
            return "orthorhombic"
        elif i <= 142:
            return "tetragonal"
        elif i <= 167:
            return "trigonal"
        elif i <= 194:
            return "hexagonal"
        else:
            return "cubic"

    def get_orbit(self, p, tol=1e-5):
        """
        Returns the orbit for a point.

        Args:
            p: Point as a 3x1 array.

        Returns:
            ([array]) Orbit for point.
        """
        p = np.append(p, [1])
        orbit = super(SpaceGroup, self).get_orbit(p, tol=tol)
        return np.delete(orbit, np.s_[-1:], 1)

    @classmethod
    def from_int_number(cls, int_number, hexagonal=True):
        """
        Obtains a SpaceGroup from its international number.

        Args:
            int_number (int): International number.
            hexagonal (bool): For rhombohedral groups, whether to return the
                hexagonal setting (default) or rhombohedral setting.

        Returns:
            (SpaceGroup)
        """
        return SpaceGroup(sg_symbol_from_int_number(int_number,
                                                    hexagonal=hexagonal))

    def __str__(self):
        return "Spacegroup %s with international number %d and order %d" % (
            self.symbol, self.int_number, len(self.symmetry_ops))

def sg_symbol_from_int_number(int_number, hexagonal=True):
    """
    Obtains a SpaceGroup name from its international number.

    Args:
        int_number (int): International number.
        hexagonal (bool): For rhombohedral groups, whether to return the
            hexagonal setting (default) or rhombohedral setting.

    Returns:
        (str) Spacegroup symbol
    """
    syms = []
    for n, v in SPACE_GROUP_ENC.items():
        if v["int_number"] == int_number:
            syms.append(n)
    if len(syms) == 0:
        raise ValueError("Invalid international number!")
    if len(syms) == 2:
        if hexagonal:
            syms = list(filter(lambda s: s.endswith("H"), syms))
        else:
            syms = list(filter(lambda s: not s.endswith("H"), syms))
    return syms.pop()

def in_array_list(array_list, a, tol=1e-5):
    """
    Extremely efficient nd-array comparison using numpy's broadcasting. This
    function checks if a particular array a, is present in a list of arrays.
    It works for arrays of any size, e.g., even matrix searches.

    Args:
        array_list ([array]): A list of arrays to compare to.
        a (array): The test array for comparison.
        tol (float): The tolerance. Defaults to 1e-5. If 0, an exact match is
            done.

    Returns:
        (bool)
    """
    if len(array_list) == 0:
        return False
    axes = tuple(range(1, a.ndim + 1))
    if not tol:
        return np.any(np.all(np.equal(array_list, a[None, :]), axes))
    else:
        return np.any(np.sum(np.abs(array_list - a[None, :]), axes) < tol)
<|endoftext|>


<commit_before>#!/usr/bin/env python
#common
import time
import datetime as t
import sys
#rgb matrix modules
try:
    from samplebase import SampleBase
    from PIL import Image
    from PIL import ImageDraw
except:
    print("install libraries")

from wqTime2LayoutWrapper import WqTime2LayoutHandler

class wqTargetApp(SampleBase):

    def __init__(self,updatePeriod,*args, **kwargs):
        super(wqTargetApp, self).__init__(*args, **kwargs)
        self.timeHandler = WqTime2LayoutHandler()
        self.bitMapSReg = [[x for x in range(11)] for x in range(10)]
        self.debug = True
        self.period = updatePeriod
        self.init = False
        self.intensity = 50
        self.process()

    def setBitMapForShiftingReg(self,bitMapId, bit):
        self.bitMapSReg[bitMapId][bit] = 1

    def resetBitMapForShiftingReg(self,bitMapId,bit):
        self.bitMapSReg[bitMapId][bit] = 0

    def printBitMapForShiftingReg(self):
        if self.debug:
            image = Image.new("RGB", (32, 32))  # Can be larger than matrix if wanted!!
            draw = ImageDraw.Draw(image)  # Declare Draw instance before prims
            self.matrix.Clear()
            print(" ### Mask For Shifting Registers ###")
            rowIdx = 0
            for row in self.bitMapSReg:
                print(row)
                pixlIdx = 0
                for pixel in row:
                    if pixel != 0:
                        draw.rectangle((2*pixlIdx+pixlIdx, 2*pixlIdx+pixlIdx+1, *rowIdx+rowIdx, 2*rowIdx+rowIdx+1), fill=(self.intensity,self.intensity,self.intensity), outline=(self.intensity,self.intensity,self.intensity))
                    pixlIdx = pixlIdx + 1
                rowIdx = rowIdx + 1
            self.matrix.SetImage(image)
            print(" \n")

    def setLetterBitValues(self):
        self.timeHandler.updateTime()

        for i in range(self.timeHandler.getLetterFieldRowSize()):
            for j in range(self.timeHandler.getLetterFieldColSize(i)):
                #Here we would set the shift registers!!!!!
                self.resetBitMapForShiftingReg(i,j)
                if(self.timeHandler.setActiveByIndex(i,j)):
                    self.setBitMapForShiftingReg(i,j)
        self.printBitMapForShiftingReg()

    def run(self):
        if self.init and self.timeHandler.getTimeSinceLastUpdate() < t.timedelta(minutes=self.period):
            print(" ### Wait for "+str(self.period)+" minutes ###")
            time.sleep(self.period*100)
        else:
            self.init = True
            self.setLetterBitValues()

if __name__ == '__main__':

    app = wqTargetApp(1)

    while True:
        try:
            # Start loop
            app.run()
        except KeyboardInterrupt:
            print("Exiting\n")
            sys.exit(0)<commit_msg>resetting everything else<commit_after>#!/usr/bin/env python
#common
import time
import datetime as t
import sys
#rgb matrix modules
try:
    from samplebase import SampleBase
    from PIL import Image
    from PIL import ImageDraw
except:
    print("install libraries")

from wqTime2LayoutWrapper import WqTime2LayoutHandler

class wqTargetApp(SampleBase):

    def __init__(self,updatePeriod,*args, **kwargs):
        super(wqTargetApp, self).__init__(*args, **kwargs)
        self.timeHandler = WqTime2LayoutHandler()
        self.bitMapSReg = [[x for x in range(11)] for x in range(10)]
        self.debug = True
        self.period = updatePeriod
        self.init = False
        self.intensity = 50
        self.process()

    def setBitMapForShiftingReg(self,bitMapId, bit):
        self.bitMapSReg[bitMapId][bit] = 1

    def resetBitMapForShiftingReg(self,bitMapId,bit):
        self.bitMapSReg[bitMapId][bit] = 0

    def printBitMapForShiftingReg(self):
        if self.debug:
            image = Image.new("RGB", (32, 32))  # Can be larger than matrix if wanted!!
            draw = ImageDraw.Draw(image)  # Declare Draw instance before prims
            self.matrix.Clear()
            print(" ### Mask For Shifting Registers ###")
            rowIdx = 0
            for row in self.bitMapSReg:
                print(row)
                pixlIdx = 0
                for pixel in row:
                    if pixel != 0:
                        draw.rectangle((2*pixlIdx+pixlIdx, 2*pixlIdx+pixlIdx+1, 2*rowIdx+rowIdx, 2*rowIdx+rowIdx+1), fill=(self.intensity,self.intensity,self.intensity), outline=(self.intensity,self.intensity,self.intensity))
                    else:
                        draw.rectangle((pixlIdx, pixlIdx, rowIdx, rowIdx), fill=(0,0,0), outline=(0,0,0))
                    pixlIdx = pixlIdx + 1
                rowIdx = rowIdx + 1
            self.matrix.SetImage(image)
            print(" \n")

    def setLetterBitValues(self):
        self.timeHandler.updateTime()

        for i in range(self.timeHandler.getLetterFieldRowSize()):
            for j in range(self.timeHandler.getLetterFieldColSize(i)):
                #Here we would set the shift registers!!!!!
                self.resetBitMapForShiftingReg(i,j)
                if(self.timeHandler.setActiveByIndex(i,j)):
                    self.setBitMapForShiftingReg(i,j)
        self.printBitMapForShiftingReg()

    def run(self):
        if self.init and self.timeHandler.getTimeSinceLastUpdate() < t.timedelta(minutes=self.period):
            print(" ### Wait for "+str(self.period)+" minutes ###")
            time.sleep(self.period*100)
        else:
            self.init = True
            self.setLetterBitValues()

if __name__ == '__main__':

    app = wqTargetApp(1)

    while True:
        try:
            # Start loop
            app.run()
        except KeyboardInterrupt:
            print("Exiting\n")
            sys.exit(0)<|endoftext|>


<commit_before>'''
Created on 14 Aug 2017

@author: Mathias Bucher
'''
import unittest
from mock import MagicMock
from ctr.Controller import Controller
from ctr.Log import Log
from model.Model import Model
from view.View import View
from model.ModelEntry import ModelEntry

class TestController(unittest.TestCase):

    def setUp(self):
        self.log = Log("testlog.txt")
        self.log.add = MagicMock()
        self.ctr = Controller( self.log )
        self.ctr.currentEntry = ModelEntry(self.log, "furniture")
        self.ctr.currentEntry.description = "This is the description of furniture"
        self.ctr.currentEntry.keywords.append("chair")
        self.ctr.currentEntry.keywords.append("table")
        self.ctr.view.drawEntry = MagicMock()
        self.ctr.view.drawSearch = MagicMock()
        self.ctr.view.removeEntry = MagicMock()
        self.ctr.model.updateNameOfEntry = MagicMock()
        self.ctr.model.updateContentOfEntry= MagicMock()
        self.ctr.model.addEntry = MagicMock()
        self.ctr.model.hasEntry= MagicMock()
        self.ctr.model.getEntries = MagicMock()

    def tearDown(self):
        pass

    def testEntryNameChangeAction(self):
        '''Checks if the right method of model is called'''
        self.ctr.entryNameChangeAction("roofs")
        self.ctr.model.updateNameOfEntry.assert_called_with(self.ctr.currentEntry, "roofs")
        self.assertEqual("roofs", self.ctr.currentEntry.name)
        self.ctr.view.removeEntry.assert_called_with(self.ctr.currentEntry)
        self.ctr.view.drawEntry.assert_called_with(self.ctr.currentEntry)

    def testEntryDescriptionChangeAction(self):
        '''Checks if the right method of model is called'''
        self.ctr.entryDescriptionChangeAction("new description")
        self.ctr.model.updateContentOfEntry.assert_called_with(self.ctr.currentEntry)
        self.assertEqual("new description", self.ctr.currentEntry.description)
        self.ctr.view.drawEntry.assert_called_with(self.ctr.currentEntry)

    def testEntryKeywordChangeAction(self):
        '''Checks if the right method of model is called'''
        self.ctr.entryKeywordChangeAction("window")
        self.ctr.model.updateContentOfEntry.assert_called_with(self.ctr.currentEntry)
        s = self.ctr.currentEntry.getStringFromKeywords(self.ctr.currentEntry.keywords)
        self.assertEqual("window", s)
        self.ctr.view.drawEntry.assert_called_with(self.ctr.currentEntry)

    def testNewEntryAction(self):
        '''Checks if an entry is added correctly'''
        self.ctr.model.hasEntry.return_value = False
        self.ctr.newEntryAction()
        self.assertEqual("enter name", self.ctr.currentEntry.name)
        self.ctr.model.addEntry.assert_called_with(self.ctr.currentEntry)
        self.ctr.view.drawEntry.assert_called_with(self.ctr.currentEntry)

    def testNewEntryActionIfEnterNameExists(self):
        '''In this test, the default enter name entry already existst'''
        self.ctr.model.hasEntry.side_effect = [True, True, False]
        self.ctr.newEntryAction()
        self.assertEqual("enter name2", self.ctr.currentEntry.name)
        self.ctr.model.addEntry.assert_called_with(self.ctr.currentEntry)
        self.ctr.view.drawEntry.assert_called_with(self.ctr.currentEntry)

    def testSearchActionMultipleMatches(self):
        '''Tests whether the search action calls the right methods'''
        e1 = ModelEntry(self.log, "hit1")
        e2 = ModelEntry(self.log, "hit2")
        e3 = ModelEntry(self.log, "hit3")
        e4 = ModelEntry(self.log, "hit4")

        found = {"name" : [e1, e2],
                 "keyword" : [e3],
                 "description" : [e4]}

        self.ctr.model.getEntries.return_value = found
        self.ctr.searchAction("search")
        self.ctr.model.getEntries.assert_called_with("search")
        self.ctr.view.drawSearch.assert_called_with(found)

    def testSearchActionSingleMatch(self):
        '''Tests whether the search action calls the right methods'''
        e1 = ModelEntry(self.log, "hit1")

        found = {"name" : [],
                 "keyword" : [e1],
                 "description" : []}

        self.ctr.model.getEntries.return_value = found
        self.ctr.searchAction("search")
        self.ctr.model.getEntries.assert_called_with("search")
        self.ctr.view.drawSearch.assert_called_with(found)

    def testEntryClickedInVSearch(self):
        '''Tests whether the search action calls the right methods'''
        e = ModelEntry(self.log, "entry")
        self.ctr.entryClickedInVSearch(e)
        self.ctr.view.drawEntry.assert_called_with(e)

if __name__ == "__main__":
    #import sys;sys.argv = ['', 'Test.testads']
    unittest.main()<commit_msg>Added tests to controller<commit_after>'''
Created on 14 Aug 2017

@author: Mathias Bucher
'''
import unittest
from mock import MagicMock
from ctr.Controller import Controller
from ctr.Log import Log
from model.Model import Model
from view.View import View
from model.ModelEntry import ModelEntry

class TestController(unittest.TestCase):

    def setUp(self):
        self.log = Log("testlog.txt")
        self.log.add = MagicMock()
        self.ctr = Controller( self.log )
        self.ctr.currentEntry = ModelEntry(self.log, "furniture")
        self.ctr.currentEntry.description = "This is the description of furniture"
        self.ctr.currentEntry.keywords.append("chair")
        self.ctr.currentEntry.keywords.append("table")
        self.ctr.view.drawEntry = MagicMock()
        self.ctr.view.drawSearch = MagicMock()
        self.ctr.view.removeEntry = MagicMock()
        self.ctr.model.updateNameOfEntry = MagicMock()
        self.ctr.model.updateContentOfEntry= MagicMock()
        self.ctr.model.addEntry = MagicMock()
        self.ctr.model.hasEntry= MagicMock()
        self.ctr.model.getEntries = MagicMock()

    def tearDown(self):
        pass

    def testEntryNameChangeAction(self):
        '''Checks if the right method of model is called'''
        self.ctr.entryNameChangeAction("roofs")
        self.ctr.model.updateNameOfEntry.assert_called_with(self.ctr.currentEntry, "roofs")
        self.assertEqual("roofs", self.ctr.currentEntry.name)
        self.ctr.view.removeEntry.assert_called_with(self.ctr.currentEntry)
        self.ctr.view.drawEntry.assert_called_with(self.ctr.currentEntry)

    def testEntryDescriptionChangeAction(self):
        '''Checks if the right method of model is called'''
        self.ctr.entryDescriptionChangeAction("new description")
        self.ctr.model.updateContentOfEntry.assert_called_with(self.ctr.currentEntry)
        self.assertEqual("new description", self.ctr.currentEntry.description)
        self.ctr.view.drawEntry.assert_called_with(self.ctr.currentEntry)

    def testEntryKeywordChangeAction(self):
        '''Checks if the right method of model is called'''
        self.ctr.entryKeywordChangeAction("window")
        self.ctr.model.updateContentOfEntry.assert_called_with(self.ctr.currentEntry)
        s = self.ctr.currentEntry.getStringFromKeywords(self.ctr.currentEntry.keywords)
        self.assertEqual("window", s)
        self.ctr.view.drawEntry.assert_called_with(self.ctr.currentEntry)

    def testNewEntryAction(self):
        '''Checks if an entry is added correctly'''
        self.ctr.model.hasEntry.return_value = False
        self.ctr.newEntryAction()
        self.assertEqual("enter name", self.ctr.currentEntry.name)
        self.ctr.model.addEntry.assert_called_with(self.ctr.currentEntry)
        self.ctr.view.drawEntry.assert_called_with(self.ctr.currentEntry)

    def testNewEntryActionIfEnterNameExists(self):
        '''In this test, the default enter name entry already existst'''
        self.ctr.model.hasEntry.side_effect = [True, True, False]
        self.ctr.newEntryAction()
        self.assertEqual("enter name2", self.ctr.currentEntry.name)
        self.ctr.model.addEntry.assert_called_with(self.ctr.currentEntry)
        self.ctr.view.drawEntry.assert_called_with(self.ctr.currentEntry)

    def testSearchActionMultipleMatches(self):
        '''Tests whether the search action calls the right methods'''
        e1 = ModelEntry(self.log, "hit1")
        e2 = ModelEntry(self.log, "hit2")
        e3 = ModelEntry(self.log, "hit3")
        e4 = ModelEntry(self.log, "hit4")

        found = {"name" : [e1, e2],
                 "keyword" : [e3],
                 "description" : [e4]}

        self.ctr.model.getEntries.return_value = found
        self.ctr.searchAction("search")
        self.ctr.model.getEntries.assert_called_with("search")
        self.ctr.view.drawSearch.assert_called_with(found)

    def testSearchActionSingleMatch(self):
        '''Tests whether the search action calls the right methods'''
        e1 = ModelEntry(self.log, "hit1")

        found = {"name" : [],
                 "keyword" : [e1],
                 "description" : []}

        self.ctr.model.getEntries.return_value = found
        self.ctr.searchAction("search")
        self.ctr.model.getEntries.assert_called_with("search")
        self.ctr.view.drawSearch.assert_called_with(found)

    def testEntryClickedInVSearch(self):
        '''Tests whether the search action calls the right methods'''
        e = ModelEntry(self.log, "entry")

        self.ctr.entryClickedInVSearch(e)
        self.assertEqual(e, self.ctr.model.activeEntries[0])
        self.ctr.view.drawEntry.assert_called_with(e)

    def testCloseTabAction(self):
        e = ModelEntry(self.log, "entry")
        self.ctr.currentEntry = e
        self.ctr.model.activeEntries.append(e)
        self.ctr.closeTabAction()
        self.assertEqual(0, self.ctr.model.activeEntries.__len__())
        self.ctr.view.removeEntry.assert_called_with(e)

    def testTabChangeAction(self):
        e1 = ModelEntry(self.log, "e1")
        e2 = ModelEntry(self.log, "e2")

        self.ctr.model.activeEntries.append(e1)
        self.ctr.model.activeEntries.append(e2)
        self.ctr.currentEntry = e1
        self.ctr.tabChangeAction(e2.name)
        self.assertEqual(e2, self.ctr.currentEntry)

        self.ctr.tabChangeAction(None)
        self.assertEqual(e2, self.ctr.currentEntry)

if __name__ == "__main__":
    #import sys;sys.argv = ['', 'Test.testads']
    unittest.main()<|endoftext|>


<commit_before>import os
import sys
import errno
import fcntl
import atexit
import signal
import logging
import tempfile

__version__ = "1.1.0"

DEFAULT_PID_DIR = "/var/run/"

class PidFileError(Exception):
    pass

class PidFileUnreadableError(PidFileError):
    pass

class PidFileAlreadyRunningError(PidFileError):
    pass

class PidFileAlreadyLockedError(PidFileError):
    pass

class PidFile(object):
    __slots__ = ("pid", "pidname", "piddir", "enforce_dotpid_postfix", "register_term_signal_handler",
                 "term_signal_handler", "filename", "fh", "logger", "lock_pidfile", "chmod", "uid",
                 "gid", "force_tmpdir", "lazy")

    def __init__(self, pidname=None, piddir=None, enforce_dotpid_postfix=True,
                 register_term_signal_handler=True, term_signal_handler=None,
                 lock_pidfile=True, chmod=0o644, uid=-1, gid=-1, force_tmpdir=False,
                 lazy=False):
        self.pidname = pidname
        self.piddir = piddir
        self.enforce_dotpid_postfix = enforce_dotpid_postfix
        self.register_term_signal_handler = register_term_signal_handler
        self.term_signal_handler = term_signal_handler
        self.lock_pidfile = lock_pidfile
        self.chmod = chmod
        self.uid = uid
        self.gid = gid
        self.force_tmpdir = force_tmpdir
        self.lazy = lazy

        self.fh = None
        self.filename = None
        self.pid = None

        if not self.lazy:
            self._setup()

    def _setup(self):
        self.logger = logging.getLogger("PidFile")
        self.logger.debug("%r entering setup", self)
        if self.filename is None:
            self.pid = os.getpid()
            self.filename = self._make_filename()
            self._register_term_signal()

        # setup should only be performed once
        self.lazy = False

    def _make_filename(self):
        pidname = self.pidname
        piddir = self.piddir
        if pidname is None:
            pidname = "%s.pid" % os.path.basename(sys.argv[0])
        if self.enforce_dotpid_postfix and not pidname.endswith(".pid"):
            pidname = "%s.pid" % pidname
        if piddir is None:
            if os.path.isdir(DEFAULT_PID_DIR) and self.force_tmpdir is False:
                piddir = DEFAULT_PID_DIR
            else:
                piddir = tempfile.gettempdir()

        if not os.path.isdir(piddir):
            os.makedirs(piddir)
        if not os.access(piddir, os.R_OK):
            raise IOError("Pid file directory '%s' cannot be read" % piddir)
        if not os.access(piddir, os.W_OK):
            raise IOError("Pid file directory '%s' cannot be written to" % piddir)

        filename = os.path.abspath(os.path.join(piddir, pidname))
        return filename

    def _register_term_signal(self):
        if self.register_term_signal_handler:
            # Register TERM signal handler to make sure atexit runs on TERM signal

            term_signal_handler = self.term_signal_handler
            if term_signal_handler is None:
                def sigterm_noop_handler(*args, **kwargs):
                    raise SystemExit(1)
                term_signal_handler = sigterm_noop_handler

            if signal.getsignal(signal.SIGTERM) is not term_signal_handler:
                signal.signal(signal.SIGTERM, term_signal_handler)

    def check(self):

        def inner_check(fh):
            try:
                fh.seek(0)
                pid_str = fh.read(16).split("\n", 1)[0].strip()
                if not pid_str:
                    return None
                pid = int(pid_str)
            except (IOError, ValueError) as exc:
                self.close(fh=fh)
                raise PidFileUnreadableError(exc)
            try:
                os.kill(pid, 0)
            except OSError as exc:
                if exc.errno == errno.ESRCH:
                    # this pid is not running
                    return None
                self.close(fh=fh, cleanup=False)
                raise PidFileAlreadyRunningError(exc)
            self.close(fh=fh, cleanup=False)
            raise PidFileAlreadyRunningError("Program already running with pid: %d" % pid)

        if self.lazy:
            self._setup()

        self.logger.debug("%r check pidfile: %s", self, self.filename)

        if self.fh is None:
            if self.filename and os.path.isfile(self.filename):
                with open(self.filename, "r") as fh:
                    inner_check(fh)
        else:
            inner_check(self.fh)

    def create(self):
        if self.lazy:
            self._setup()

        self.logger.debug("%r create pidfile: %s", self, self.filename)
        self.fh = open(self.filename, 'a+')
        if self.lock_pidfile:
            try:
                fcntl.flock(self.fh.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
            except IOError as exc:
                self.close(cleanup=False)
                raise PidFileAlreadyLockedError(exc)
        self.check()
        if self.chmod:
            os.fchmod(self.fh.fileno(), self.chmod)
        if self.uid >= 0 or self.gid >= 0:
            os.fchown(self.fh.fileno(), self.uid, self.gid)
        self.fh.seek(0)
        self.fh.truncate()
        # pidfile must consist of the pid and a newline character
        self.fh.write("%d\n" % self.pid)
        self.fh.flush()
        self.fh.seek(0)
        atexit.register(self.close)

    def close(self, fh=None, cleanup=True):
        if not fh:
            fh = self.fh
        try:
            if fh is None or self.lazy is True:
                return
            self.logger.debug("%r closing pidfile: %s", self, self.filename)
            fh.close()
        except IOError as exc:
            # ignore error when file was already closed
            if exc.errno != errno.EBADF:
                raise
        finally:
            if self.filename and os.path.isfile(self.filename) and cleanup:
                os.remove(self.filename)

    def __enter__(self):
        self.create()
        return self

    def __exit__(self, exc_type=None, exc_value=None, exc_tb=None):
        self.close()
<commit_msg>default to lazy setup for all instantiations<commit_after>import os
import sys
import errno
import fcntl
import atexit
import signal
import logging
import tempfile

__version__ = "1.1.0"

DEFAULT_PID_DIR = "/var/run/"

class PidFileError(Exception):
    pass

class PidFileUnreadableError(PidFileError):
    pass

class PidFileAlreadyRunningError(PidFileError):
    pass

class PidFileAlreadyLockedError(PidFileError):
    pass

class PidFile(object):
    __slots__ = ("pid", "pidname", "piddir", "enforce_dotpid_postfix", "register_term_signal_handler",
                 "term_signal_handler", "filename", "fh", "logger", "lock_pidfile", "chmod", "uid",
                 "gid", "force_tmpdir", "lazy")

    def __init__(self, pidname=None, piddir=None, enforce_dotpid_postfix=True,
                 register_term_signal_handler=True, term_signal_handler=None,
                 lock_pidfile=True, chmod=0o644, uid=-1, gid=-1, force_tmpdir=False,
                 lazy=True):
        self.pidname = pidname
        self.piddir = piddir
        self.enforce_dotpid_postfix = enforce_dotpid_postfix
        self.register_term_signal_handler = register_term_signal_handler
        self.term_signal_handler = term_signal_handler
        self.lock_pidfile = lock_pidfile
        self.chmod = chmod
        self.uid = uid
        self.gid = gid
        self.force_tmpdir = force_tmpdir
        self.lazy = lazy

        self.fh = None
        self.filename = None
        self.pid = None

        if not self.lazy:
            self._setup()

    def _setup(self):
        self.logger = logging.getLogger("PidFile")
        self.logger.debug("%r entering setup", self)
        if self.filename is None:
            self.pid = os.getpid()
            self.filename = self._make_filename()
            self._register_term_signal()

        # setup should only be performed once
        self.lazy = False

    def _make_filename(self):
        pidname = self.pidname
        piddir = self.piddir
        if pidname is None:
            pidname = "%s.pid" % os.path.basename(sys.argv[0])
        if self.enforce_dotpid_postfix and not pidname.endswith(".pid"):
            pidname = "%s.pid" % pidname
        if piddir is None:
            if os.path.isdir(DEFAULT_PID_DIR) and self.force_tmpdir is False:
                piddir = DEFAULT_PID_DIR
            else:
                piddir = tempfile.gettempdir()

        if not os.path.isdir(piddir):
            os.makedirs(piddir)
        if not os.access(piddir, os.R_OK):
            raise IOError("Pid file directory '%s' cannot be read" % piddir)
        if not os.access(piddir, os.W_OK):
            raise IOError("Pid file directory '%s' cannot be written to" % piddir)

        filename = os.path.abspath(os.path.join(piddir, pidname))
        return filename

    def _register_term_signal(self):
        if self.register_term_signal_handler:
            # Register TERM signal handler to make sure atexit runs on TERM signal

            term_signal_handler = self.term_signal_handler
            if term_signal_handler is None:
                def sigterm_noop_handler(*args, **kwargs):
                    raise SystemExit(1)
                term_signal_handler = sigterm_noop_handler

            if signal.getsignal(signal.SIGTERM) is not term_signal_handler:
                signal.signal(signal.SIGTERM, term_signal_handler)

    def check(self):

        def inner_check(fh):
            try:
                fh.seek(0)
                pid_str = fh.read(16).split("\n", 1)[0].strip()
                if not pid_str:
                    return None
                pid = int(pid_str)
            except (IOError, ValueError) as exc:
                self.close(fh=fh)
                raise PidFileUnreadableError(exc)
            try:
                os.kill(pid, 0)
            except OSError as exc:
                if exc.errno == errno.ESRCH:
                    # this pid is not running
                    return None
                self.close(fh=fh, cleanup=False)
                raise PidFileAlreadyRunningError(exc)
            self.close(fh=fh, cleanup=False)
            raise PidFileAlreadyRunningError("Program already running with pid: %d" % pid)

        if self.lazy:
            self._setup()

        self.logger.debug("%r check pidfile: %s", self, self.filename)

        if self.fh is None:
            if self.filename and os.path.isfile(self.filename):
                with open(self.filename, "r") as fh:
                    inner_check(fh)
        else:
            inner_check(self.fh)

    def create(self):
        if self.lazy:
            self._setup()

        self.logger.debug("%r create pidfile: %s", self, self.filename)
        self.fh = open(self.filename, 'a+')
        if self.lock_pidfile:
            try:
                fcntl.flock(self.fh.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
            except IOError as exc:
                self.close(cleanup=False)
                raise PidFileAlreadyLockedError(exc)
        self.check()
        if self.chmod:
            os.fchmod(self.fh.fileno(), self.chmod)
        if self.uid >= 0 or self.gid >= 0:
            os.fchown(self.fh.fileno(), self.uid, self.gid)
        self.fh.seek(0)
        self.fh.truncate()
        # pidfile must consist of the pid and a newline character
        self.fh.write("%d\n" % self.pid)
        self.fh.flush()
        self.fh.seek(0)
        atexit.register(self.close)

    def close(self, fh=None, cleanup=True):
        if not fh:
            fh = self.fh
        try:
            if fh is None or self.lazy is True:
                return
            self.logger.debug("%r closing pidfile: %s", self, self.filename)
            fh.close()
        except IOError as exc:
            # ignore error when file was already closed
            if exc.errno != errno.EBADF:
                raise
        finally:
            if self.filename and os.path.isfile(self.filename) and cleanup:
                os.remove(self.filename)

    def __enter__(self):
        self.create()
        return self

    def __exit__(self, exc_type=None, exc_value=None, exc_tb=None):
        self.close()
<|endoftext|>


<commit_before><commit_msg>Fixed realtime loading (tested with demo)<commit_after><|endoftext|>


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

#Copyright (C) 2013 by Glenn Hickey
#
# Contains some fragments of code from sklearn/tests/test_hmm.py
# (2010 - 2013, scikit-learn developers (BSD License))
#
#Released under the MIT license, see LICENSE.txt
import unittest
import sys
import os
import math
from numpy.testing import assert_array_equal, assert_array_almost_equal

from teHmm.tracksInfo import TracksInfo
from teHmm.track import *
from teHmm.trackIO import readBedIntervals
from teHmm.hmm import MultitrackHmm
from teHmm.cfg import MultitrackCfg
from teHmm.emission import IndependentMultinomialEmissionModel, PairEmissionModel

from teHmm.tests.common import getTestDirPath
from teHmm.tests.common import TestBase
from teHmm.tests.bedTrackTest import getTracksInfoPath
from teHmm.tests.emissionTest import getBedStates

class TestCase(TestBase):

    def setUp(self):
        super(TestCase, self).setUp()

    def tearDown(self):
        super(TestCase, self).tearDown()

    def testInit(self):
        emissionModel = IndependentMultinomialEmissionModel(
            10, [3], zeroAsMissingData=False)
        pairModel = PairEmissionModel(emissionModel, [1.0] *
                                      emissionModel.getNumStates())
        cfg = MultitrackCfg(emissionModel, pairModel)
        cfg = MultitrackCfg(emissionModel, [3,8])
        cfg.validate()

    def testDefaultVsHmm(self):
        emissionModel = IndependentMultinomialEmissionModel(
            10, [3], zeroAsMissingData=False)
        hmm = MultitrackHmm(emissionModel)
        pairModel = PairEmissionModel(emissionModel, [1.0] *
                                      emissionModel.getNumStates())
        cfg = MultitrackCfg(emissionModel, pairModel)

    def testDefaultHmmViterbi(self):
        emissionModel = IndependentMultinomialEmissionModel(
            5, [3], zeroAsMissingData=False)
        hmm = MultitrackHmm(emissionModel)
        pairModel = PairEmissionModel(emissionModel, [1.0] *
                                      emissionModel.getNumStates())
        cfg = MultitrackCfg(emissionModel, pairModel)

        obs = np.array([[0],[0],[1],[2]], dtype=np.int16)
        hmmProb, hmmStates = hmm.decode(obs)
        cfgProb, cfgStates = cfg.decode(obs)
        assert_array_almost_equal(hmmProb, cfgProb)

    def testTraceBack(self):
        # a model with 2 states.  state 0 has a .75 chance of emitting 0
        # state 1 has a 0.95 chance of emitting 1
        emissionModel = IndependentMultinomialEmissionModel(
            2, [2], zeroAsMissingData=False)
        emProbs = np.zeros((1, 2, 2), dtype=np.float)
        emProbs[0,0] = [0.75, 0.25]
        emProbs[0,1] = [0.05, 0.95]
        emissionModel.logProbs = np.log(emProbs)

        hmm = MultitrackHmm(emissionModel)
        pairModel = PairEmissionModel(emissionModel, [1.0] *
                                      emissionModel.getNumStates())
        cfg = MultitrackCfg(emissionModel, pairModel)

        obs = np.array([[0],[0],[1],[0]], dtype=np.int16)
        hmmProb, hmmStates = hmm.decode(obs)
        cfgProb, cfgStates = cfg.decode(obs)
        assert_array_almost_equal(hmmProb, cfgProb)
        assert_array_almost_equal(hmmStates, [0, 0, 1, 0])
        assert_array_almost_equal(hmmStates, cfgStates)

    def testBasicNesting(self):
        # a model with 3 states.  state 0 has a .75 chance of emitting 0
        # state 1 has a 0.95 chance of emitting 1
        # state 2 has a 0.90 chance of emitting 1
        emissionModel = IndependentMultinomialEmissionModel(
            3, [2], zeroAsMissingData=False)
        emProbs = np.zeros((1, 3, 2), dtype=np.float)
        emProbs[0,0] = [0.75, 0.25]
        emProbs[0,1] = [0.05, 0.95]
        emProbs[0,2] = [0.01, 0.90]
        emissionModel.logProbs = np.log(emProbs)

        # state 1 is a nested pair state!
        pairModel = PairEmissionModel(emissionModel, [1.0] *
                                      emissionModel.getNumStates())
        cfg = MultitrackCfg(emissionModel, pairModel, nestStates = [1])

        obs = np.array([[0],[0],[1],[0]], dtype=np.int16)
        cfgProb, cfgStates = cfg.decode(obs)
        # 1 is a pair only state.  no way it should be here
        assert 1 not in cfgStates
        assert_array_equal(cfgStates, [0,0,2,0])

        obs = np.array([[1],[0],[0],[1]], dtype=np.int16)
        cfgProb, cfgStates = cfg.decode(obs)
        assert_array_equal(cfgStates, [2,0,0,2])

        alignment = np.array([[1],[0],[0],[1]], dtype=np.int16)
        cfgProb, cfgStates = cfg.decode(obs, alignmentTrack = alignment,
                                        defAlignmentSymbol=0)
        assert_array_equal(cfgStates, [1,0,0,1])

        alignment = np.array([[1],[0],[0],[2]], dtype=np.int16)
        cfgProb, cfgStates = cfg.decode(obs, alignmentTrack = alignment,
                                        defAlignmentSymbol=0)
        assert_array_equal(cfgStates, [2,0,0,2])

    def testSupervisedLearn(self):
        intervals = readBedIntervals(getTestDirPath("truth.bed"), ncol=4)
        truthIntervals = []
        for i in intervals:
            truthIntervals.append((i[0], i[1], i[2], int(i[3])))

        allIntervals = [(truthIntervals[0][0],
                        truthIntervals[0][1],
                        truthIntervals[-1][2])]
        trackData = TrackData()
        trackData.loadTrackData(getTracksInfoPath(3), allIntervals)
        assert len(trackData.getTrackTableList()) == 1

        em = IndependentMultinomialEmissionModel(
            4, trackData.getNumSymbolsPerTrack())
        hmm = MultitrackHmm(em)
        hmm.supervisedTrain(trackData, truthIntervals)
        hmm.validate()
        pairModel = PairEmissionModel(em, [1.0] *
                                      em.getNumStates())
        cfg = MultitrackCfg(em, pairModel, nestStates = [1])

        cfg.supervisedTrain(trackData, truthIntervals)
        cfg.validate()

def main():
    sys.argv = sys.argv[:1]
    unittest.main()

if __name__ == '__main__':
    main()

<commit_msg>touch up supervised test<commit_after>#!/usr/bin/env python

#Copyright (C) 2013 by Glenn Hickey
#
# Contains some fragments of code from sklearn/tests/test_hmm.py
# (2010 - 2013, scikit-learn developers (BSD License))
#
#Released under the MIT license, see LICENSE.txt
import unittest
import sys
import os
import math
from numpy.testing import assert_array_equal, assert_array_almost_equal

from teHmm.tracksInfo import TracksInfo
from teHmm.track import *
from teHmm.trackIO import readBedIntervals
from teHmm.hmm import MultitrackHmm
from teHmm.cfg import MultitrackCfg
from teHmm.emission import IndependentMultinomialEmissionModel, PairEmissionModel

from teHmm.tests.common import getTestDirPath
from teHmm.tests.common import TestBase
from teHmm.tests.bedTrackTest import getTracksInfoPath
from teHmm.tests.emissionTest import getBedStates

class TestCase(TestBase):

    def setUp(self):
        super(TestCase, self).setUp()

    def tearDown(self):
        super(TestCase, self).tearDown()

    def testInit(self):
        emissionModel = IndependentMultinomialEmissionModel(
            10, [3], zeroAsMissingData=False)
        pairModel = PairEmissionModel(emissionModel, [1.0] *
                                      emissionModel.getNumStates())
        cfg = MultitrackCfg(emissionModel, pairModel)
        cfg = MultitrackCfg(emissionModel, [3,8])
        cfg.validate()

    def testDefaultVsHmm(self):
        emissionModel = IndependentMultinomialEmissionModel(
            10, [3], zeroAsMissingData=False)
        hmm = MultitrackHmm(emissionModel)
        pairModel = PairEmissionModel(emissionModel, [1.0] *
                                      emissionModel.getNumStates())
        cfg = MultitrackCfg(emissionModel, pairModel)

    def testDefaultHmmViterbi(self):
        emissionModel = IndependentMultinomialEmissionModel(
            5, [3], zeroAsMissingData=False)
        hmm = MultitrackHmm(emissionModel)
        pairModel = PairEmissionModel(emissionModel, [1.0] *
                                      emissionModel.getNumStates())
        cfg = MultitrackCfg(emissionModel, pairModel)

        obs = np.array([[0],[0],[1],[2]], dtype=np.int16)
        hmmProb, hmmStates = hmm.decode(obs)
        cfgProb, cfgStates = cfg.decode(obs)
        assert_array_almost_equal(hmmProb, cfgProb)

    def testTraceBack(self):
        # a model with 2 states.  state 0 has a .75 chance of emitting 0
        # state 1 has a 0.95 chance of emitting 1
        emissionModel = IndependentMultinomialEmissionModel(
            2, [2], zeroAsMissingData=False)
        emProbs = np.zeros((1, 2, 2), dtype=np.float)
        emProbs[0,0] = [0.75, 0.25]
        emProbs[0,1] = [0.05, 0.95]
        emissionModel.logProbs = np.log(emProbs)

        hmm = MultitrackHmm(emissionModel)
        pairModel = PairEmissionModel(emissionModel, [1.0] *
                                      emissionModel.getNumStates())
        cfg = MultitrackCfg(emissionModel, pairModel)

        obs = np.array([[0],[0],[1],[0]], dtype=np.int16)
        hmmProb, hmmStates = hmm.decode(obs)
        cfgProb, cfgStates = cfg.decode(obs)
        assert_array_almost_equal(hmmProb, cfgProb)
        assert_array_almost_equal(hmmStates, [0, 0, 1, 0])
        assert_array_almost_equal(hmmStates, cfgStates)

    def testBasicNesting(self):
        # a model with 3 states.  state 0 has a .75 chance of emitting 0
        # state 1 has a 0.95 chance of emitting 1
        # state 2 has a 0.90 chance of emitting 1
        emissionModel = IndependentMultinomialEmissionModel(
            3, [2], zeroAsMissingData=False)
        emProbs = np.zeros((1, 3, 2), dtype=np.float)
        emProbs[0,0] = [0.75, 0.25]
        emProbs[0,1] = [0.05, 0.95]
        emProbs[0,2] = [0.01, 0.90]
        emissionModel.logProbs = np.log(emProbs)

        # state 1 is a nested pair state!
        pairModel = PairEmissionModel(emissionModel, [1.0] *
                                      emissionModel.getNumStates())
        cfg = MultitrackCfg(emissionModel, pairModel, nestStates = [1])

        obs = np.array([[0],[0],[1],[0]], dtype=np.int16)
        cfgProb, cfgStates = cfg.decode(obs)
        # 1 is a pair only state.  no way it should be here
        assert 1 not in cfgStates
        assert_array_equal(cfgStates, [0,0,2,0])

        obs = np.array([[1],[0],[0],[1]], dtype=np.int16)
        cfgProb, cfgStates = cfg.decode(obs)
        assert_array_equal(cfgStates, [2,0,0,2])

        alignment = np.array([[1],[0],[0],[1]], dtype=np.int16)
        cfgProb, cfgStates = cfg.decode(obs, alignmentTrack = alignment,
                                        defAlignmentSymbol=0)
        assert_array_equal(cfgStates, [1,0,0,1])

        alignment = np.array([[1],[0],[0],[2]], dtype=np.int16)
        cfgProb, cfgStates = cfg.decode(obs, alignmentTrack = alignment,
                                        defAlignmentSymbol=0)
        assert_array_equal(cfgStates, [2,0,0,2])

    def testHmmSupervisedLearn(self):
        """ Pretty much copied from the HMM unit test.  We try to recapitualte
        all results with a CFG with no nest states, which should be same as
        HMM"""
        intervals = readBedIntervals(getTestDirPath("truth.bed"), ncol=4)
        truthIntervals = []
        for i in intervals:
            truthIntervals.append((i[0], i[1], i[2], int(i[3])))

        allIntervals = [(truthIntervals[0][0],
                        truthIntervals[0][1],
                        truthIntervals[-1][2])]
        trackData = TrackData()
        trackData.loadTrackData(getTracksInfoPath(3), allIntervals)
        assert len(trackData.getTrackTableList()) == 1
        # set the fudge to 1 since when the test was written this was
        # hardcoded default
        em = IndependentMultinomialEmissionModel(
            4, trackData.getNumSymbolsPerTrack(),
			  fudge = 1.0)
        hmm = MultitrackHmm(em)
        hmm.supervisedTrain(trackData, truthIntervals)
        hmm.validate()
        pairModel = PairEmissionModel(em, [1.0] *
                                      em.getNumStates())
        # Test validates with neststate just for fun
        cfg = MultitrackCfg(em, pairModel, nestStates = [1])

        cfg.supervisedTrain(trackData, truthIntervals)
        cfg.validate()

        # Then reload as an hmm-equivalent
        cfg = MultitrackCfg(em, pairModel, nestStates = [])

        cfg.supervisedTrain(trackData, truthIntervals)
        cfg.validate()

        # check emissions, they should basically be binary.
        trackList = cfg.getTrackList()
        emp = np.exp(em.getLogProbs())
        ltrTrack = trackList.getTrackByName("ltr")
        track = ltrTrack.getNumber()
        cmap = ltrTrack.getValueMap()
        s0 = cmap.getMap(None)
        s1 = cmap.getMap(0)
        # we add 1 to all frequencies like emission trainer
        assert_array_almost_equal(emp[track][0][s0], 36. / 37.)
        assert_array_almost_equal(emp[track][0][s1], 1 - 36. / 37.)
        assert_array_almost_equal(emp[track][1][s0], 1 - 6. / 7.)
        assert_array_almost_equal(emp[track][1][s1], 6. / 7.)
        assert_array_almost_equal(emp[track][2][s0], 26. / 27.)
        assert_array_almost_equal(emp[track][2][s1], 1. - 26. / 27.)
        assert_array_almost_equal(emp[track][3][s0], 1. - 6. / 7.)
        assert_array_almost_equal(emp[track][3][s1], 6. / 7.)

        insideTrack = trackList.getTrackByName("inside")
        track = insideTrack.getNumber()
        cmap = insideTrack.getValueMap()
        s0 = cmap.getMap(None)
        s1 = cmap.getMap("Inside")
        assert_array_almost_equal(emp[track][0][s0], 36. / 37.)
        assert_array_almost_equal(emp[track][0][s1], 1 - 36. / 37.)
        assert_array_almost_equal(emp[track][1][s0], 6. / 7.)
        assert_array_almost_equal(emp[track][1][s1], 1 - 6. / 7.)
        assert_array_almost_equal(emp[track][2][s0], 1. - 26. / 27.)
        assert_array_almost_equal(emp[track][2][s1], 26. / 27.)
        assert_array_almost_equal(emp[track][3][s0], 6. / 7.)
        assert_array_almost_equal(emp[track][3][s1], 1. - 6. / 7.)

        # crappy check for start probs.  need to test transition too!
        freq = [0.0] * em.getNumStates()
        total = 0.0
        for interval in truthIntervals:
           state = interval[3]
           freq[state] += float(interval[2]) - float(interval[1])
           total += float(interval[2]) - float(interval[1])

        sprobs = cfg.getStartProbs()
        assert len(sprobs) == em.getNumStates()
        for state in xrange(em.getNumStates()):
            assert_array_almost_equal(freq[state] / total, sprobs[state])

        # transition probabilites
        # from eyeball:
        #c	0	5	0   0->0 +4   0->1 +1    0-> +5
        #c	5	10	1   1->1 +4   1->2 +1    1-> +5
        #c	10	35	2   2->2 +24  2->3 +1    2-> +25
        #c	35	40	3   3->3 +4   3->0 +1    3-> +5
        #c	40	70	0   0->0 +29             0-> +19
        realTransProbs = np.array([
            [33. / 34., 1. / 34., 0., 0.],
            [0., 4. / 5., 1. / 5., 0.],
            [0., 0., 24. / 25., 1. / 25.],
            [1. / 5., 0., 0., 4. / 5.]
            ])

        tprobs = np.exp(cfg.getLogProbTables()[0])
        assert tprobs.shape == (em.getNumStates(), em.getNumStates(),
                                em.getNumStates())
        for i in xrange(em.getNumStates()):
            for j in xrange(em.getNumStates()):
                fbTot = tprobs[i, i, j]
                if i != j:
                    fbTot += tprobs[i, j, i]
                assert_array_almost_equal(fbTot, realTransProbs[i,j])
        prob, states = cfg.decode(trackData.getTrackTableList()[0])
        for truthInt in truthIntervals:
            for i in xrange(truthInt[1], truthInt[2]):
                # gah, just realized that ltr track is binary, which means
                # ltr states can be either 1 or 3.  need to fix test properly
                # but just relax comparison for now.
                if truthInt[3] == 1 or truthInt[3] == 3:
                    assert states[i] == 1 or states[i] == 3
                else:
                    assert states[i] == truthInt[3]

def main():
    sys.argv = sys.argv[:1]
    unittest.main()

if __name__ == '__main__':
    main()

<|endoftext|>


<commit_before>#! /bin/env python
import sys
import xml.dom.minidom

import numpy as np
from six.moves import xrange

from .encoders import encode
from .vtktypes import np_to_vtk_type, sys_to_vtk_endian

class VtkExtent(object):
    def __init__(self, shape):
        if len(shape) > 3:
            raise ValueError("number of dimension must be <= 3")

        self._shape = shape

        self._extent = []
        for n in shape:
            self._extent.append((0, n - 1))
        for n in xrange(3 - len(shape)):
            self._extent.append((0, 0))

        self._extent_str = " ".join(["%d %d" % x for x in self._extent])

    def __str__(self):
        return self._extent_str

class VtkOrigin(object):
    def __init__(self, origin, spacing):
        if len(shape) > 3:
            raise ValueError("number of dimension must be <= 3")
        if len(spacing) != len(origin):
            raise ValueError("dimension mismatch for shape and origin")

        self._spacing = spacing
        self._cell_origin = origin

        self._cell_origin = []
        for (dx, x0) in zip(spacing, origin):
            self._cell_origin.append(x0 - dx * 0.5)

        for _ in xrange(3 - len(origin)):
            self._cell_origin.append(0.0)

        self._origin_str = " ".join(["%f" % x for x in self._cell_origin])

    def __str__(self):
        return self._origin_str

class VtkSpacing(object):
    def __init__(self, spacing):
        if len(spacing) > 3:
            raise ValueError("number of dimension must be <= 3")

        self._spacing = spacing

        self._padded_spacing = []
        for dx in spacing:
            self._padded_spacing.append(dx)
        for _ in xrange(3 - len(spacing)):
            self._padded_spacing.append(0.0)

        self._spacing_str = " ".join(["%f" % x for x in self._padded_spacing])

    def __str__(self):
        return self._spacing_str

# class VtkElement (object, xml.dom.minidom.Element):
class VtkElement(xml.dom.minidom.Element):
    def __init__(self, name, **kwargs):
        self.ownerDocument = None
        self.parentNode = None
        self.nextSibling = self.previousSibling = None
        xml.dom.minidom.Element.__init__(self, str(name), namespaceURI="VTK")
        self.setAttributes(**kwargs)

    def setAttributes(self, **kwargs):
        for (attr, value) in kwargs.items():
            self.setAttribute(attr, str(value))

    def getAttributes(self, names=None):
        attrs = {}

        if names is None:
            for i in range(self.attributes.length):
                attr = self.attributes.item(i)
                attrs[attr.localName] = attr.value
        else:
            for i in range(self.attributes.length):
                attr = self.attributes.item(i)
                if attr.localName in names:
                    attrs[attr.localName] = attr.value

        return attrs

class VtkTextElement(xml.dom.minidom.Text):
    def __init__(self, text):
        self.parentNode = None
        self.nextSibling = self.previousSibling = None
        self.replaceWholeText(text)

class VtkDataArrayElement(VtkElement):
    def __init__(self, **kwargs):
        VtkElement.__init__(self, "DataArray", **kwargs)

    def addData(self, data_string):
        self.appendChild(VtkTextElement(data_string))

class VtkDataElement(VtkElement):
    format = "ascii"

    def __init__(self, name, **kwargs):
        VtkElement.__init__(self, name)

    def addData(self, data, name, append=None, encoding="ascii", **kwargs):
        data_string = encode(data, encoding=encoding)
        data_array = VtkDataArrayElement(
            Name=name, type=np_to_vtk_type[str(data.dtype)], **kwargs
        )
        self.appendChild(data_array)

        if append is not None:
            data_array.setAttributes(offset=append.offset(), format="appended")
            append.addData(data_string)
        else:
            data_array.setAttributes(format=self.format)
            data_array.addData(data_string)

class VtkRootElement(VtkElement):
    def __init__(self, root_type):
        VtkElement.__init__(
            self,
            "VTKFile",
            type=root_type,
            version="0.1",
            byte_order=str(sys_to_vtk_endian[sys.byteorder]),
        )

class VtkGridElement(VtkElement):
    def __init__(self, name, **kwargs):
        VtkElement.__init__(self, name, **kwargs)

class VtkPieceElement(VtkElement):
    def __init__(self, **kwargs):
        VtkElement.__init__(self, "Piece", **kwargs)

class VtkAppendedDataElement(VtkElement):
    def __init__(self, data, **kwargs):
        VtkElement.__init__(self, "AppendedData", **kwargs)
        self.appendChild(VtkTextElement("_" + data))

    def addData(self, data):
        self.firstChild.appendData(str(data))

    def offset(self):
        return self.firstChild.length - 1

class VtkPointsElement(VtkDataElement):
    def __init__(self, coords, **kwargs):
        n_components = 3
        xyz = []
        for i in xrange(n_components):
            try:
                xyz.append(coords[i])
            except IndexError:
                xyz.append(np.array(coords[0]) * 0)
        xyz = np.vstack(xyz).transpose().flatten()
        VtkDataElement.__init__(self, "Points", **kwargs)
        self.addData(xyz, "Coordinates", NumberOfComponents=n_components, **kwargs)

class VtkCoordinatesElement(VtkDataElement):
    def __init__(self, xyz, **kwargs):
        VtkDataElement.__init__(self, "Coordinates", **kwargs)
        for (i, label) in enumerate(["x", "y", "z"]):
            try:
                self.addData(
                    xyz[i], label + "_coordinates", NumberOfComponents=1, **kwargs
                )
            except IndexError:
                self.addData(
                    np.array(xyz[0]) * 0.0,
                    label + "_coordinates",
                    NumberOfComponents=1,
                    **kwargs
                )

        # self.addData (xyz[0], 'x_coordinates', NumberOfComponents=1, **kwargs)
        # self.addData (xyz[1], 'y_coordinates', NumberOfComponents=1, **kwargs)
        # self.addData (xyz[2], 'z_coordinates', NumberOfComponents=1, **kwargs)

class VtkCellsElement(VtkDataElement):
    def __init__(self, connectivity, offset, types, **kwargs):
        VtkDataElement.__init__(self, "Cells", **kwargs)
        self.addData(connectivity, "connectivity", **kwargs)
        self.addData(offset, "offsets", **kwargs)
        self.addData(types, "types", **kwargs)

class VtkPointDataElement(VtkDataElement):
    def __init__(self, values, **kwargs):
        VtkDataElement.__init__(self, "PointData", **kwargs)
        for (name, value) in values.items():
            self.addData(value, name, NumberOfComponents=1, **kwargs)

class VtkCellDataElement(VtkDataElement):
    def __init__(self, values, **kwargs):
        VtkDataElement.__init__(self, "CellData", **kwargs)
        for (name, value) in values.items():
            self.addData(value, name, NumberOfComponents=1, **kwargs)
<commit_msg>Check both len of spacing and origin are <= 3.<commit_after>#! /bin/env python
import sys
import xml.dom.minidom

import numpy as np
from six.moves import xrange

from .encoders import encode
from .vtktypes import np_to_vtk_type, sys_to_vtk_endian

class VtkExtent(object):
    def __init__(self, shape):
        if len(shape) > 3:
            raise ValueError("number of dimension must be <= 3")

        self._shape = shape

        self._extent = []
        for n in shape:
            self._extent.append((0, n - 1))
        for n in xrange(3 - len(shape)):
            self._extent.append((0, 0))

        self._extent_str = " ".join(["%d %d" % x for x in self._extent])

    def __str__(self):
        return self._extent_str

class VtkOrigin(object):
    def __init__(self, origin, spacing):
        if len(origin) > 3 or len(spacing) > 3:
            raise ValueError("number of dimension must be <= 3")
        if len(spacing) != len(origin):
            raise ValueError("dimension mismatch for shape and origin")

        self._spacing = spacing
        self._cell_origin = origin

        self._cell_origin = []
        for (dx, x0) in zip(spacing, origin):
            self._cell_origin.append(x0 - dx * 0.5)

        for _ in xrange(3 - len(origin)):
            self._cell_origin.append(0.0)

        self._origin_str = " ".join(["%f" % x for x in self._cell_origin])

    def __str__(self):
        return self._origin_str

class VtkSpacing(object):
    def __init__(self, spacing):
        if len(spacing) > 3:
            raise ValueError("number of dimension must be <= 3")

        self._spacing = spacing

        self._padded_spacing = []
        for dx in spacing:
            self._padded_spacing.append(dx)
        for _ in xrange(3 - len(spacing)):
            self._padded_spacing.append(0.0)

        self._spacing_str = " ".join(["%f" % x for x in self._padded_spacing])

    def __str__(self):
        return self._spacing_str

# class VtkElement (object, xml.dom.minidom.Element):
class VtkElement(xml.dom.minidom.Element):
    def __init__(self, name, **kwargs):
        self.ownerDocument = None
        self.parentNode = None
        self.nextSibling = self.previousSibling = None
        xml.dom.minidom.Element.__init__(self, str(name), namespaceURI="VTK")
        self.setAttributes(**kwargs)

    def setAttributes(self, **kwargs):
        for (attr, value) in kwargs.items():
            self.setAttribute(attr, str(value))

    def getAttributes(self, names=None):
        attrs = {}

        if names is None:
            for i in range(self.attributes.length):
                attr = self.attributes.item(i)
                attrs[attr.localName] = attr.value
        else:
            for i in range(self.attributes.length):
                attr = self.attributes.item(i)
                if attr.localName in names:
                    attrs[attr.localName] = attr.value

        return attrs

class VtkTextElement(xml.dom.minidom.Text):
    def __init__(self, text):
        self.parentNode = None
        self.nextSibling = self.previousSibling = None
        self.replaceWholeText(text)

class VtkDataArrayElement(VtkElement):
    def __init__(self, **kwargs):
        VtkElement.__init__(self, "DataArray", **kwargs)

    def addData(self, data_string):
        self.appendChild(VtkTextElement(data_string))

class VtkDataElement(VtkElement):
    format = "ascii"

    def __init__(self, name, **kwargs):
        VtkElement.__init__(self, name)

    def addData(self, data, name, append=None, encoding="ascii", **kwargs):
        data_string = encode(data, encoding=encoding)
        data_array = VtkDataArrayElement(
            Name=name, type=np_to_vtk_type[str(data.dtype)], **kwargs
        )
        self.appendChild(data_array)

        if append is not None:
            data_array.setAttributes(offset=append.offset(), format="appended")
            append.addData(data_string)
        else:
            data_array.setAttributes(format=self.format)
            data_array.addData(data_string)

class VtkRootElement(VtkElement):
    def __init__(self, root_type):
        VtkElement.__init__(
            self,
            "VTKFile",
            type=root_type,
            version="0.1",
            byte_order=str(sys_to_vtk_endian[sys.byteorder]),
        )

class VtkGridElement(VtkElement):
    def __init__(self, name, **kwargs):
        VtkElement.__init__(self, name, **kwargs)

class VtkPieceElement(VtkElement):
    def __init__(self, **kwargs):
        VtkElement.__init__(self, "Piece", **kwargs)

class VtkAppendedDataElement(VtkElement):
    def __init__(self, data, **kwargs):
        VtkElement.__init__(self, "AppendedData", **kwargs)
        self.appendChild(VtkTextElement("_" + data))

    def addData(self, data):
        self.firstChild.appendData(str(data))

    def offset(self):
        return self.firstChild.length - 1

class VtkPointsElement(VtkDataElement):
    def __init__(self, coords, **kwargs):
        n_components = 3
        xyz = []
        for i in xrange(n_components):
            try:
                xyz.append(coords[i])
            except IndexError:
                xyz.append(np.array(coords[0]) * 0)
        xyz = np.vstack(xyz).transpose().flatten()
        VtkDataElement.__init__(self, "Points", **kwargs)
        self.addData(xyz, "Coordinates", NumberOfComponents=n_components, **kwargs)

class VtkCoordinatesElement(VtkDataElement):
    def __init__(self, xyz, **kwargs):
        VtkDataElement.__init__(self, "Coordinates", **kwargs)
        for (i, label) in enumerate(["x", "y", "z"]):
            try:
                self.addData(
                    xyz[i], label + "_coordinates", NumberOfComponents=1, **kwargs
                )
            except IndexError:
                self.addData(
                    np.array(xyz[0]) * 0.0,
                    label + "_coordinates",
                    NumberOfComponents=1,
                    **kwargs
                )

        # self.addData (xyz[0], 'x_coordinates', NumberOfComponents=1, **kwargs)
        # self.addData (xyz[1], 'y_coordinates', NumberOfComponents=1, **kwargs)
        # self.addData (xyz[2], 'z_coordinates', NumberOfComponents=1, **kwargs)

class VtkCellsElement(VtkDataElement):
    def __init__(self, connectivity, offset, types, **kwargs):
        VtkDataElement.__init__(self, "Cells", **kwargs)
        self.addData(connectivity, "connectivity", **kwargs)
        self.addData(offset, "offsets", **kwargs)
        self.addData(types, "types", **kwargs)

class VtkPointDataElement(VtkDataElement):
    def __init__(self, values, **kwargs):
        VtkDataElement.__init__(self, "PointData", **kwargs)
        for (name, value) in values.items():
            self.addData(value, name, NumberOfComponents=1, **kwargs)

class VtkCellDataElement(VtkDataElement):
    def __init__(self, values, **kwargs):
        VtkDataElement.__init__(self, "CellData", **kwargs)
        for (name, value) in values.items():
            self.addData(value, name, NumberOfComponents=1, **kwargs)
<|endoftext|>


<commit_before>from collections import defaultdict

from django.core.management.base import CommandError

from corehq.apps.hqcase.utils import bulk_update_cases
from corehq.form_processor.interfaces.dbaccessors import CaseAccessors
from corehq.form_processor.models import CommCareCaseSQL
from custom.enikshay.case_utils import (
    get_person_case_from_occurrence,
    CASE_TYPE_DRUG_RESISTANCE,
)
from custom.enikshay.management.commands.base_model_reconciliation import (
    BaseModelReconciliationCommand,
    DOMAIN,
)
from custom.enikshay.exceptions import ENikshayCaseNotFound
from custom.enikshay.const import (
    DRUG_RESISTANCE_SENSITIVITY,
    DRUG_RESISTANCE_SENSITIVITY_RESISTANT,
    DRUG_RESISTANCE_SENSITIVITY_SENSITIVE,
)

class Command(BaseModelReconciliationCommand):
    email_subject = "Drug Resistance Reconciliation Report"
    result_file_name_prefix = "drug_resistance_reconciliation_report"
    result_file_headers = [
        "occurrence_case_id",
        "person_case_id",
        "drug_id",
        "retain_case_id",
        "retain_reason",
        "closed_case_ids",
        "closed_extension_case_ids"
    ]

    def handle(self, *args, **options):
        self.commit = options.get('commit')
        self.log_progress = options.get('log_progress')
        self.recipient = (options.get('recipient') or 'mkangia@dimagi.com')
        self.recipient = list(self.recipient) if not isinstance(self.recipient, basestring) else [self.recipient]
        self.result_file_name = self.setup_result_file()
        self.drug_id_values = []
        self.sensitivity_values = []
        self.case_accessor = CaseAccessors(DOMAIN)
        # iterate all occurrence cases
        for occurrence_case_id in self._get_open_occurrence_case_ids_to_process():
            # Need to consider only public app cases so skip updates if enrolled in private
            if self.public_app_case(occurrence_case_id):
                self.reconcile_case(occurrence_case_id)

        if self.log_progress:
            print("All drug id values found %s" % self.drug_id_values)
            print("All sensitivity values found %s" % self.sensitivity_values)

        self.email_report()

    def public_app_case(self, occurrence_case_id):
        try:
            person_case = get_person_case_from_occurrence(DOMAIN, occurrence_case_id)
        except ENikshayCaseNotFound:
            return False
        self.person_case_id = person_case.case_id
        return super(Command, self).public_app_case(person_case)

    def reconcile_case(self, occurrence_case_id):
        # get all open drug resistance cases
        # group by drug_id
        # if any drug_id has more than one corresponding drug resistance case, fix them
        drug_resistance_cases = get_open_drug_resistance_cases_from_occurrence(occurrence_case_id)
        drug_resistance_cases_by_drug_id = defaultdict(list)
        for drug_resistance_case in drug_resistance_cases:
            drug_id = drug_resistance_case.get_case_property('drug_id')
            if drug_id not in self.drug_id_values:
                if self.log_progress:
                    print("New drug_id value found, %s" % drug_id)
                self.drug_id_values.append(drug_id)
            drug_resistance_cases_by_drug_id[drug_id].append(drug_resistance_case)
        for drug_id, drug_resistance_cases_for_drug in drug_resistance_cases_by_drug_id.items():
            if len(drug_resistance_cases_for_drug) > 1:
                self.reconcile_drug_resistance_cases(drug_resistance_cases_for_drug, drug_id, occurrence_case_id)

    def reconcile_drug_resistance_cases(self, drug_resistance_cases, drug_id, occurrence_case_id):
        """
        Take cases with the same drug_id and then keep only one and close others.
        Find the one in this order
        - that has sensitivity resistant
        - that has sensitivity sensitive
        - the first ever opened case
        :param drug_resistance_cases: under an occurrence with the same drug_id
        """
        # sanity check that we got more than one case so we should be considering closing cases
        # and confirm that all have the same drug id
        if (len(drug_resistance_cases) < 2 or
            any(drug_resistance_case.get_case_property('drug_id') != drug_id
                for drug_resistance_case in drug_resistance_cases)):
            raise CommandError("Asked to reconcile cases when not needed for occurrence case %s"
                               % occurrence_case_id)

        resistant_drug_resistance_cases = []
        sensitive_drug_resistance_cases = []

        # group cases by their sensitivity
        # possible values are resistant/ sensitive/ unknown
        for drug_resistance_case in drug_resistance_cases:
            sensitivity = drug_resistance_case.get_case_property(DRUG_RESISTANCE_SENSITIVITY)
            if sensitivity == DRUG_RESISTANCE_SENSITIVITY_RESISTANT:
                resistant_drug_resistance_cases.append(drug_resistance_case)
            elif sensitivity == DRUG_RESISTANCE_SENSITIVITY_SENSITIVE:
                sensitive_drug_resistance_cases.append(drug_resistance_case)
            else:
                if sensitivity not in self.sensitivity_values:
                    self.sensitivity_values.append(sensitivity)
                    print("New sensitivity value found, %s" % sensitivity)
        resistant_drug_resistance_cases_count = len(resistant_drug_resistance_cases)
        sensitive_drug_resistance_cases_count = len(resistant_drug_resistance_cases)

        # if there is more than one case with sensitivity resistant, keep the one
        # opened first.
        if resistant_drug_resistance_cases_count > 1:
            retain_case = self.get_first_opened_case(resistant_drug_resistance_cases)
            self.close_cases(drug_resistance_cases, occurrence_case_id, drug_id, retain_case,
                             "More than one resistance drug cases, picked first opened.")
        # if there is only one case with sensitivity resistant, keep that, close rest.
        elif resistant_drug_resistance_cases_count == 1:
            retain_case = resistant_drug_resistance_cases[0]
            self.close_cases(drug_resistance_cases, occurrence_case_id, drug_id, retain_case,
                             "Only one resistance drug case found.")
        # if there is more than one case with sensitivity sensitive, keep the one
        # opened first.
        elif sensitive_drug_resistance_cases_count > 1:
            retain_case = self.get_first_opened_case(sensitive_drug_resistance_cases)
            self.close_cases(drug_resistance_cases, occurrence_case_id, drug_id, retain_case,
                             "More than one sensitive drug cases, picked first opened.")
        # if there is only one case with sensitivity sensitive, keep that, close rest.
        elif sensitive_drug_resistance_cases_count == 1:
            retain_case = sensitive_drug_resistance_cases[0]
            self.close_cases(drug_resistance_cases, occurrence_case_id, drug_id, retain_case,
                             "Only one sensitive drug case found.")
        # No case has sensitivity resistant and sensitive. Probably multiple cases with unknown status
        # keep the one opened first, close rest.
        else:
            retain_case = self.get_first_opened_case(drug_resistance_cases)
            self.close_cases(drug_resistance_cases, occurrence_case_id, drug_id, retain_case,
                             "Picked first opened.")

    def _get_open_occurrence_case_ids_to_process(self):
        from corehq.sql_db.util import get_db_aliases_for_partitioned_query
        dbs = get_db_aliases_for_partitioned_query()
        for db in dbs:
            case_ids = (
                CommCareCaseSQL.objects
                .using(db)
                .filter(domain=DOMAIN, type="occurrence", closed=False)
                .values_list('case_id', flat=True)
            )
            num_case_ids = len(case_ids)
            if self.log_progress:
                print("processing %d docs from db %s" % (num_case_ids, db))
            for i, case_id in enumerate(case_ids):
                yield case_id
                if i % 1000 == 0 and self.log_progress:
                    print("processed %d / %d docs from db %s" % (i, num_case_ids, db))

    def close_cases(self, all_cases, occurrence_case_id, drug_id, retain_case, retain_reason):
        # remove duplicates in case ids to remove so that we dont retain and close
        # the same case by mistake
        all_case_ids = {case.case_id for case in all_cases}
        retain_case_id = retain_case.case_id
        case_ids_to_close = all_case_ids.copy()
        case_ids_to_close.remove(retain_case_id)

        case_accessor = CaseAccessors(DOMAIN)
        closing_extension_case_ids = case_accessor.get_extension_case_ids(case_ids_to_close)

        self.writerow({
            "occurrence_case_id": occurrence_case_id,
            "person_case_id": self.person_case_id,
            "drug_id": drug_id,
            "retain_case_id": retain_case_id,
            "retain_reason": retain_reason,
            "closed_case_ids": ','.join(map(str, case_ids_to_close)),
            "closed_extension_case_ids": ','.join(map(str, closing_extension_case_ids))
        })
        if self.commit:
            updates = [(case_id, {'close_reason': "duplicate_reconciliation"}, True)
                       for case_id in case_ids_to_close]
            bulk_update_cases(DOMAIN, updates, self.__module__)

def get_open_drug_resistance_cases_from_occurrence(occurrence_case_id):
    case_accessor = CaseAccessors(DOMAIN)
    all_cases = case_accessor.get_reverse_indexed_cases([occurrence_case_id])
    return [case for case in all_cases
            if not case.closed and case.type == CASE_TYPE_DRUG_RESISTANCE]
<commit_msg>check the correct set of cases for count [skip ci]<commit_after>from collections import defaultdict

from django.core.management.base import CommandError

from corehq.apps.hqcase.utils import bulk_update_cases
from corehq.form_processor.interfaces.dbaccessors import CaseAccessors
from corehq.form_processor.models import CommCareCaseSQL
from custom.enikshay.case_utils import (
    get_person_case_from_occurrence,
    CASE_TYPE_DRUG_RESISTANCE,
)
from custom.enikshay.management.commands.base_model_reconciliation import (
    BaseModelReconciliationCommand,
    DOMAIN,
)
from custom.enikshay.exceptions import ENikshayCaseNotFound
from custom.enikshay.const import (
    DRUG_RESISTANCE_SENSITIVITY,
    DRUG_RESISTANCE_SENSITIVITY_RESISTANT,
    DRUG_RESISTANCE_SENSITIVITY_SENSITIVE,
)

class Command(BaseModelReconciliationCommand):
    email_subject = "Drug Resistance Reconciliation Report"
    result_file_name_prefix = "drug_resistance_reconciliation_report"
    result_file_headers = [
        "occurrence_case_id",
        "person_case_id",
        "drug_id",
        "retain_case_id",
        "retain_reason",
        "closed_case_ids",
        "closed_extension_case_ids"
    ]

    def handle(self, *args, **options):
        self.commit = options.get('commit')
        self.log_progress = options.get('log_progress')
        self.recipient = (options.get('recipient') or 'mkangia@dimagi.com')
        self.recipient = list(self.recipient) if not isinstance(self.recipient, basestring) else [self.recipient]
        self.result_file_name = self.setup_result_file()
        self.drug_id_values = []
        self.sensitivity_values = []
        self.case_accessor = CaseAccessors(DOMAIN)
        # iterate all occurrence cases
        for occurrence_case_id in self._get_open_occurrence_case_ids_to_process():
            # Need to consider only public app cases so skip updates if enrolled in private
            if self.public_app_case(occurrence_case_id):
                self.reconcile_case(occurrence_case_id)

        if self.log_progress:
            print("All drug id values found %s" % self.drug_id_values)
            print("All sensitivity values found %s" % self.sensitivity_values)

        self.email_report()

    def public_app_case(self, occurrence_case_id):
        try:
            person_case = get_person_case_from_occurrence(DOMAIN, occurrence_case_id)
        except ENikshayCaseNotFound:
            return False
        self.person_case_id = person_case.case_id
        return super(Command, self).public_app_case(person_case)

    def reconcile_case(self, occurrence_case_id):
        # get all open drug resistance cases
        # group by drug_id
        # if any drug_id has more than one corresponding drug resistance case, fix them
        drug_resistance_cases = get_open_drug_resistance_cases_from_occurrence(occurrence_case_id)
        drug_resistance_cases_by_drug_id = defaultdict(list)
        for drug_resistance_case in drug_resistance_cases:
            drug_id = drug_resistance_case.get_case_property('drug_id')
            if drug_id not in self.drug_id_values:
                if self.log_progress:
                    print("New drug_id value found, %s" % drug_id)
                self.drug_id_values.append(drug_id)
            drug_resistance_cases_by_drug_id[drug_id].append(drug_resistance_case)
        for drug_id, drug_resistance_cases_for_drug in drug_resistance_cases_by_drug_id.items():
            if len(drug_resistance_cases_for_drug) > 1:
                self.reconcile_drug_resistance_cases(drug_resistance_cases_for_drug, drug_id, occurrence_case_id)

    def reconcile_drug_resistance_cases(self, drug_resistance_cases, drug_id, occurrence_case_id):
        """
        Take cases with the same drug_id and then keep only one and close others.
        Find the one in this order
        - that has sensitivity resistant
        - that has sensitivity sensitive
        - the first ever opened case
        :param drug_resistance_cases: under an occurrence with the same drug_id
        """
        # sanity check that we got more than one case so we should be considering closing cases
        # and confirm that all have the same drug id
        if (len(drug_resistance_cases) < 2 or
            any(drug_resistance_case.get_case_property('drug_id') != drug_id
                for drug_resistance_case in drug_resistance_cases)):
            raise CommandError("Asked to reconcile cases when not needed for occurrence case %s"
                               % occurrence_case_id)

        resistant_drug_resistance_cases = []
        sensitive_drug_resistance_cases = []

        # group cases by their sensitivity
        # possible values are resistant/ sensitive/ unknown
        for drug_resistance_case in drug_resistance_cases:
            sensitivity = drug_resistance_case.get_case_property(DRUG_RESISTANCE_SENSITIVITY)
            if sensitivity == DRUG_RESISTANCE_SENSITIVITY_RESISTANT:
                resistant_drug_resistance_cases.append(drug_resistance_case)
            elif sensitivity == DRUG_RESISTANCE_SENSITIVITY_SENSITIVE:
                sensitive_drug_resistance_cases.append(drug_resistance_case)
            else:
                if sensitivity not in self.sensitivity_values:
                    self.sensitivity_values.append(sensitivity)
                    print("New sensitivity value found, %s" % sensitivity)
        resistant_drug_resistance_cases_count = len(resistant_drug_resistance_cases)
        sensitive_drug_resistance_cases_count = len(sensitive_drug_resistance_cases)

        # if there is more than one case with sensitivity resistant, keep the one
        # opened first.
        if resistant_drug_resistance_cases_count > 1:
            retain_case = self.get_first_opened_case(resistant_drug_resistance_cases)
            self.close_cases(drug_resistance_cases, occurrence_case_id, drug_id, retain_case,
                             "More than one resistance drug cases, picked first opened.")
        # if there is only one case with sensitivity resistant, keep that, close rest.
        elif resistant_drug_resistance_cases_count == 1:
            retain_case = resistant_drug_resistance_cases[0]
            self.close_cases(drug_resistance_cases, occurrence_case_id, drug_id, retain_case,
                             "Only one resistance drug case found.")
        # if there is more than one case with sensitivity sensitive, keep the one
        # opened first.
        elif sensitive_drug_resistance_cases_count > 1:
            retain_case = self.get_first_opened_case(sensitive_drug_resistance_cases)
            self.close_cases(drug_resistance_cases, occurrence_case_id, drug_id, retain_case,
                             "More than one sensitive drug cases, picked first opened.")
        # if there is only one case with sensitivity sensitive, keep that, close rest.
        elif sensitive_drug_resistance_cases_count == 1:
            retain_case = sensitive_drug_resistance_cases[0]
            self.close_cases(drug_resistance_cases, occurrence_case_id, drug_id, retain_case,
                             "Only one sensitive drug case found.")
        # No case has sensitivity resistant and sensitive. Probably multiple cases with unknown status
        # keep the one opened first, close rest.
        else:
            retain_case = self.get_first_opened_case(drug_resistance_cases)
            self.close_cases(drug_resistance_cases, occurrence_case_id, drug_id, retain_case,
                             "Picked first opened.")

    def _get_open_occurrence_case_ids_to_process(self):
        from corehq.sql_db.util import get_db_aliases_for_partitioned_query
        dbs = get_db_aliases_for_partitioned_query()
        for db in dbs:
            case_ids = (
                CommCareCaseSQL.objects
                .using(db)
                .filter(domain=DOMAIN, type="occurrence", closed=False)
                .values_list('case_id', flat=True)
            )
            num_case_ids = len(case_ids)
            if self.log_progress:
                print("processing %d docs from db %s" % (num_case_ids, db))
            for i, case_id in enumerate(case_ids):
                yield case_id
                if i % 1000 == 0 and self.log_progress:
                    print("processed %d / %d docs from db %s" % (i, num_case_ids, db))

    def close_cases(self, all_cases, occurrence_case_id, drug_id, retain_case, retain_reason):
        # remove duplicates in case ids to remove so that we dont retain and close
        # the same case by mistake
        all_case_ids = {case.case_id for case in all_cases}
        retain_case_id = retain_case.case_id
        case_ids_to_close = all_case_ids.copy()
        case_ids_to_close.remove(retain_case_id)

        case_accessor = CaseAccessors(DOMAIN)
        closing_extension_case_ids = case_accessor.get_extension_case_ids(case_ids_to_close)

        self.writerow({
            "occurrence_case_id": occurrence_case_id,
            "person_case_id": self.person_case_id,
            "drug_id": drug_id,
            "retain_case_id": retain_case_id,
            "retain_reason": retain_reason,
            "closed_case_ids": ','.join(map(str, case_ids_to_close)),
            "closed_extension_case_ids": ','.join(map(str, closing_extension_case_ids))
        })
        if self.commit:
            updates = [(case_id, {'close_reason': "duplicate_reconciliation"}, True)
                       for case_id in case_ids_to_close]
            bulk_update_cases(DOMAIN, updates, self.__module__)

def get_open_drug_resistance_cases_from_occurrence(occurrence_case_id):
    case_accessor = CaseAccessors(DOMAIN)
    all_cases = case_accessor.get_reverse_indexed_cases([occurrence_case_id])
    return [case for case in all_cases
            if not case.closed and case.type == CASE_TYPE_DRUG_RESISTANCE]
<|endoftext|>