<commit_before>"""Provides class which implements a customizable pipeline."""

import os
import pickle
import random

import cv2
import numpy as np

class ImageDataPipeline:

    """Specifies a custom pipeline object which performs a sequence of image preprocessing ops."""

    def __init__(self,
                 rescale=[128, 128],
                 colorspace='CIELAB',
                 current_bounds=[[0, 255], [0, 255], [0, 255]],
                 desired_bounds=[[0, 1], [-1, 1], [-1, 1]],
                 mode='NCHW',
                 dtype='float32'):
        """
        Instance Attributes:
            - These attributes are used by the core pipeline function `preprocess_image` as
              arguments to the lower level methods.
                - This method is used as a building block for more complex pipelines.
                    - i.e. `preprocess_directory` and `preprocess_classes`.
            - For each attribute, see the methods listed for details.
                - rescale
                    - `resize_image`
                - colorspace
                    - `convert_colorspace`
                - mode
                    - `format_tensor`
                - current_bounds, desired_bounds, dtype
                    - `normalize_image`
        """
        self.rescale = rescale
        self.colorspace = colorspace
        self.current_bounds = current_bounds
        self.desired_bounds = desired_bounds
        self.mode = mode
        self.dtype = dtype

    def valid_file(self,
                   path):
        """
        Ensures path points to a file of a supported filetype by OpenCV.

        Parameters:
            - path (str)
                - Path to the file.
        Returns:
            - A boolean; true if valid, false if not.
        """
        supported_formats = [
            '.bmp',
            '.pbm',
            '.pgm',
            '.ppm',
            '.sr',
            '.ras',
            '.jpeg',
            '.jpg',
            '.jpe',
            '.jp2',
            '.tiff',
            '.tif',
            '.png'
        ]
        extension = os.path.splitext(path)[1].lower()
        if extension in supported_formats:
            return True
        return False

    def load_image(self,
                   path):
        """
        Loads an image's RGB values as a tensor.

        Parameters:
            - path (str)
                - Path to the image.
        Returns:
            - A tensor.
                - Encoded in RGB.
                - Formatted in HWC.
                - Datatype is uint8.
        """
        image = cv2.imread(path)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        return image

    def resize_image(self,
                     image,
                     rescale):
        """
        Resizes an image with bilinear interpolation.

        Parameters:
            - image (np.ndarray)
                - Encoded in RGB.
                - Formatted in HWC.
                - Datatype is uint8.
            - rescale (list of two ints)
                - Desired [height, width].
                - e.g. [1080, 1920]
        Returns:
            - An idential tensor, with a different height/width as per rescale arg.
        """
        return cv2.resize(image, tuple(reversed(rescale)), interpolation=cv2.INTER_LINEAR)

    def convert_colorspace(self,
                           image,
                           colorspace):
        """
        Converts an RGB image to a different color space.

        Parameters:
            - image (np.ndarray)
                - Encoded in RGB.
                - Formatted in HWC.
                - Datatype is uint8.
            - colorspace (str)
                - Options are: 'RGB' (unchanged), 'GRAYSCALE', 'RGB+GRAYSCALE', 'CIELAB', 'HSV'
                    - 'GRAYSCALE' is computed via OpenCV's implementation.
                        - https://bit.ly/2pUL2hR
                        - Output tensors will be HxWx1 in range [0, 1].
                    - 'RGB+GRAYSCALE' is simply RGB with a fourth channel—grayscale.
                        - Output tensors will be HxWx4 in range [0, 1].
                        - Grayscale is computed as per the same link above.
                    - 'CIELAB' is computed via OpenCV's implementation.
                        - https://bit.ly/2pUL2hR
                        - 'L' is bounded in [0, 1]; 'A' and 'B' are in [-1, 1].
                        - The white reference point is from the D65 illuminant; shape is HxWx3.
                    - 'HSV' is computed via OpenCV's implementation.
                        - https://bit.ly/2pUL2hR
                        - Output tensors will be HxWx3 in range [0, 1].
        Returns:
            - The converted tensor.
            - Still in uint8.
            - Still formatted in HWC, but may have different number of channels.
        """
        if colorspace == 'CIELAB':
            image = cv2.cvtColor(image, cv2.COLOR_RGB2LAB)
        elif colorspace == 'GRAYSCALE':
            image = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
        elif colorspace == 'RGB+GRAYSCALE':
            gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
            image = np.dstack((image, gray))
        elif colorspace == 'HSV':
            image = cv2.cvtColor(image, cv2.COLOR_RGB2HSV)
        return image

    def normalize_image(self,
                        image,
                        current_bounds,
                        desired_bounds,
                        dtype='float32'):
        """
        Normalizes an image CHANNELWISE.
        Changes the boundaries of the interval in which the image's numerical values lie.

        Parameters:
            - image (np.ndarray)
                - Formatted in HWC.
                - Datatype is uint8.
            - current_bounds (list of lists of two ints each)
                - e.g. For a uint8 image with 2 channels: [[0, 255], [0, 255]]
            - desired_bounds (list of lists of two ints each)
                - The desired boundaries for the new tensor.
            - dtype (str)
                - A numpy-compatible datatype.
        Returns:
            - The resulting tensor.
            - Still formatted in HWC.
            - Only difference is the datatype and range of allowed numbers.
        """
        image = image.astype(dtype)
        number_of_channels = image.shape[2]
        for channel in range(0, number_of_channels):
            image[:, :, channel] += -current_bounds[channel][0]
            image[:, :, channel] /= (current_bounds[channel][1] /
                                     (desired_bounds[channel][1] - desired_bounds[channel][0]))
            image[:, :, channel] += desired_bounds[channel][0]
        return image

    def format_tensor(self,
                      image,
                      mode):
        """
        Takes in a HWC tensor and converts to either NHWC or NCHW.

        Parameters:
            - image (np.ndarray)
                - The input tensor which represents the image.
                - Must be formatted in HWC.
            - mode (str)
                - Either 'NHWC' or 'NCHW'.
        """
        image = np.expand_dims(image, axis=0)
        if mode == 'NCHW':
            image = image.transpose([0, 3, 1, 2])
        return image

    def preprocess_image(self,
                         path):
        """
        Pipeline for complete preprocessing of an image.
        Loads image, rescales it, converts colorspace, normalizes, and converts datatype.
        Finalizes by converting HWC to NHWC.

        Parameters:
            - path (str)
                - Path to the image to preprocess.
        Returns:
            - The fully preprocessed image in NHWC format.
        """
        image = self.load_image(path)
        image = self.resize_image(image, self.rescale)
        image = self.convert_colorspace(image, self.colorspace)
        image = self.normalize_image(image, self.current_bounds, self.desired_bounds,
                                     dtype=self.dtype)
        image = self.format_tensor(image, self.mode)
        return image

    def preprocess_directory(self,
                             path):
        """
        An extension of the preprocess_image method for directories.
        Given a directory, preprocesses images in it with the preprocess_image method.
        Subdirectories and files of unsupported formats are ignored.

        Parameters:
            - path (str)
                - Path to the directory.
        Yields:
            - A tuple (filepath, preprocessed_image_array).
                - See the preprocess_image function in this module for details on the latter.
        """
        for filename in os.listdir(path):
            filepath = os.path.join(path, filename)
            if not self.valid_file(filepath):
                continue
            preprocessed_image = self.preprocess_image(filepath)
            yield filepath, preprocessed_image

    def preprocess_classes(self,
                           steps,
                           train_dir):
        """
        Given a directory of subdirectories of images, preprocesses an image from the 1st subdir,
        then the 2nd, then the Nth, and then loops back towards the 1st and gets another image,
        etc.

        After all images in a subdir have been preprocessed (given that `steps` is big enough),
        preprocessing will start over at the beginning of the subdir in question.

        The order of images within each subdir is randomized at the start, but not randomized again
        afterwards.

        Parameters:
            - steps (int)
                - Amount of step-input-label triplets to generate.
            - train_dir (str)
                - Path to the directory of classes.
                - e.g. 'data/train', where 'train' holds subdirs with images in them.
                - IMPORTANT: This directory should also hold a file called `label.pkl`.
                    - A (pickled) dictionary of numpy arrays.
                    - Maps the name of the subdirectory (class) to a label.
                        - e.g. {'cats': np.array([[1, 0]]), 'dogs': np.array([[0, 1]])}
                            - Each label must have the same shape!
                            - In this case, the two labels are of shape [1, 2].
        Yields:
            - A tuple (step, image_path, preprocessed_image_array, label_array) starting w/ step 1.
        """
        classes = [directory for directory in os.listdir(train_dir) if directory != 'label.pkl']
        cursors = {}
        images = {}
        for class_ in classes:
            cursors[class_] = 0
            images[class_] = os.listdir(os.path.join(train_dir, class_))
            random.shuffle(images[class_])
        label_path = os.path.join(train_dir, 'label.pkl')
        with open(label_path, 'rb') as f:
            label_dict = pickle.load(f)
        step = 0
        while True:
            for class_ in classes:
                if step < steps:
                    step += 1
                else:
                    return
                image_path = os.path.join(train_dir, class_, images[class_][cursors[class_]])
                if not self.valid_file(image_path):
                    continue
                preprocessed_image = self.preprocess_image(image_path)
                if cursors[class_] == (len(images[class_]) - 1):
                    cursors[class_] = 0
                else:
                    cursors[class_] += 1
                yield step, image_path, preprocessed_image, label_dict[class_]
<commit_msg>Removed whitespace.<commit_after>"""Provides class which implements a customizable pipeline."""

import os
import pickle
import random

import cv2
import numpy as np

class ImageDataPipeline:

    """Specifies a custom pipeline object which performs a sequence of image preprocessing ops."""

    def __init__(self,
                 rescale=[128, 128],
                 colorspace='CIELAB',
                 current_bounds=[[0, 255], [0, 255], [0, 255]],
                 desired_bounds=[[0, 1], [-1, 1], [-1, 1]],
                 mode='NCHW',
                 dtype='float32'):
        """
        Instance Attributes:
            - These attributes are used by the core pipeline function `preprocess_image` as
              arguments to the lower level methods.
                - This method is used as a building block for more complex pipelines.
                    - i.e. `preprocess_directory` and `preprocess_classes`.
            - For each attribute, see the methods listed for details.
                - rescale
                    - `resize_image`
                - colorspace
                    - `convert_colorspace`
                - mode
                    - `format_tensor`
                - current_bounds, desired_bounds, dtype
                    - `normalize_image`
        """
        self.rescale = rescale
        self.colorspace = colorspace
        self.current_bounds = current_bounds
        self.desired_bounds = desired_bounds
        self.mode = mode
        self.dtype = dtype

    def valid_file(self,
                   path):
        """
        Ensures path points to a file of a supported filetype by OpenCV.

        Parameters:
            - path (str)
                - Path to the file.
        Returns:
            - A boolean; true if valid, false if not.
        """
        supported_formats = [
            '.bmp',
            '.pbm',
            '.pgm',
            '.ppm',
            '.sr',
            '.ras',
            '.jpeg',
            '.jpg',
            '.jpe',
            '.jp2',
            '.tiff',
            '.tif',
            '.png'
        ]
        extension = os.path.splitext(path)[1].lower()
        if extension in supported_formats:
            return True
        return False

    def load_image(self,
                   path):
        """
        Loads an image's RGB values as a tensor.

        Parameters:
            - path (str)
                - Path to the image.
        Returns:
            - A tensor.
                - Encoded in RGB.
                - Formatted in HWC.
                - Datatype is uint8.
        """
        image = cv2.imread(path)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        return image

    def resize_image(self,
                     image,
                     rescale):
        """
        Resizes an image with bilinear interpolation.

        Parameters:
            - image (np.ndarray)
                - Encoded in RGB.
                - Formatted in HWC.
                - Datatype is uint8.
            - rescale (list of two ints)
                - Desired [height, width].
                - e.g. [1080, 1920]
        Returns:
            - An idential tensor, with a different height/width as per rescale arg.
        """
        return cv2.resize(image, tuple(reversed(rescale)), interpolation=cv2.INTER_LINEAR)

    def convert_colorspace(self,
                           image,
                           colorspace):
        """
        Converts an RGB image to a different color space.

        Parameters:
            - image (np.ndarray)
                - Encoded in RGB.
                - Formatted in HWC.
                - Datatype is uint8.
            - colorspace (str)
                - Options are: 'RGB' (unchanged), 'GRAYSCALE', 'RGB+GRAYSCALE', 'CIELAB', 'HSV'
                    - 'GRAYSCALE' is computed via OpenCV's implementation.
                        - https://bit.ly/2pUL2hR
                        - Output tensors will be HxWx1 in range [0, 1].
                    - 'RGB+GRAYSCALE' is simply RGB with a fourth channel—grayscale.
                        - Output tensors will be HxWx4 in range [0, 1].
                        - Grayscale is computed as per the same link above.
                    - 'CIELAB' is computed via OpenCV's implementation.
                        - https://bit.ly/2pUL2hR
                        - 'L' is bounded in [0, 1]; 'A' and 'B' are in [-1, 1].
                        - The white reference point is from the D65 illuminant; shape is HxWx3.
                    - 'HSV' is computed via OpenCV's implementation.
                        - https://bit.ly/2pUL2hR
                        - Output tensors will be HxWx3 in range [0, 1].
        Returns:
            - The converted tensor.
            - Still in uint8.
            - Still formatted in HWC, but may have different number of channels.
        """
        if colorspace == 'CIELAB':
            image = cv2.cvtColor(image, cv2.COLOR_RGB2LAB)
        elif colorspace == 'GRAYSCALE':
            image = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
        elif colorspace == 'RGB+GRAYSCALE':
            gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
            image = np.dstack((image, gray))
        elif colorspace == 'HSV':
            image = cv2.cvtColor(image, cv2.COLOR_RGB2HSV)
        return image

    def normalize_image(self,
                        image,
                        current_bounds,
                        desired_bounds,
                        dtype='float32'):
        """
        Normalizes an image CHANNELWISE.
        Changes the boundaries of the interval in which the image's numerical values lie.

        Parameters:
            - image (np.ndarray)
                - Formatted in HWC.
                - Datatype is uint8.
            - current_bounds (list of lists of two ints each)
                - e.g. For a uint8 image with 2 channels: [[0, 255], [0, 255]]
            - desired_bounds (list of lists of two ints each)
                - The desired boundaries for the new tensor.
            - dtype (str)
                - A numpy-compatible datatype.
        Returns:
            - The resulting tensor.
            - Still formatted in HWC.
            - Only difference is the datatype and range of allowed numbers.
        """
        image = image.astype(dtype)
        number_of_channels = image.shape[2]
        for channel in range(0, number_of_channels):
            image[:, :, channel] += -current_bounds[channel][0]
            image[:, :, channel] /= (current_bounds[channel][1] /
                                     (desired_bounds[channel][1] - desired_bounds[channel][0]))
            image[:, :, channel] += desired_bounds[channel][0]
        return image

    def format_tensor(self,
                      image,
                      mode):
        """
        Takes in a HWC tensor and converts to either NHWC or NCHW.

        Parameters:
            - image (np.ndarray)
                - The input tensor which represents the image.
                - Must be formatted in HWC.
            - mode (str)
                - Either 'NHWC' or 'NCHW'.
        """
        image = np.expand_dims(image, axis=0)
        if mode == 'NCHW':
            image = image.transpose([0, 3, 1, 2])
        return image

    def preprocess_image(self,
                         path):
        """
        Pipeline for complete preprocessing of an image.
        Loads image, rescales it, converts colorspace, normalizes, and converts datatype.
        Finalizes by converting HWC to NHWC.

        Parameters:
            - path (str)
                - Path to the image to preprocess.
        Returns:
            - The fully preprocessed image in NHWC format.
        """
        image = self.load_image(path)
        image = self.resize_image(image, self.rescale)
        image = self.convert_colorspace(image, self.colorspace)
        image = self.normalize_image(image, self.current_bounds, self.desired_bounds,
                                     dtype=self.dtype)
        image = self.format_tensor(image, self.mode)
        return image

    def preprocess_directory(self,
                             path):
        """
        An extension of the preprocess_image method for directories.
        Given a directory, preprocesses images in it with the preprocess_image method.
        Subdirectories and files of unsupported formats are ignored.

        Parameters:
            - path (str)
                - Path to the directory.
        Yields:
            - A tuple (filepath, preprocessed_image_array).
                - See the preprocess_image function in this module for details on the latter.
        """
        for filename in os.listdir(path):
            filepath = os.path.join(path, filename)
            if not self.valid_file(filepath):
                continue
            preprocessed_image = self.preprocess_image(filepath)
            yield filepath, preprocessed_image

    def preprocess_classes(self,
                           steps,
                           train_dir):
        """
        Given a directory of subdirectories of images, preprocesses an image from the 1st subdir,
        then the 2nd, then the Nth, and then loops back towards the 1st and gets another image,
        etc.

        After all images in a subdir have been preprocessed (given that `steps` is big enough),
        preprocessing will start over at the beginning of the subdir in question.

        The order of images within each subdir is randomized at the start, but not randomized again
        afterwards.

        Parameters:
            - steps (int)
                - Amount of step-input-label triplets to generate.
            - train_dir (str)
                - Path to the directory of classes.
                - e.g. 'data/train', where 'train' holds subdirs with images in them.
                - IMPORTANT: This directory should also hold a file called `label.pkl`.
                    - A (pickled) dictionary of numpy arrays.
                    - Maps the name of the subdirectory (class) to a label.
                        - e.g. {'cats': np.array([[1, 0]]), 'dogs': np.array([[0, 1]])}
                            - Each label must have the same shape!
                            - In this case, the two labels are of shape [1, 2].
        Yields:
            - A tuple (step, image_path, preprocessed_image_array, label_array) starting w/ step 1.
        """
        classes = [directory for directory in os.listdir(train_dir) if directory != 'label.pkl']
        cursors = {}
        images = {}
        for class_ in classes:
            cursors[class_] = 0
            images[class_] = os.listdir(os.path.join(train_dir, class_))
            random.shuffle(images[class_])
        label_path = os.path.join(train_dir, 'label.pkl')
        with open(label_path, 'rb') as f:
            label_dict = pickle.load(f)
        step = 0
        while True:
            for class_ in classes:
                if step < steps:
                    step += 1
                else:
                    return
                image_path = os.path.join(train_dir, class_, images[class_][cursors[class_]])
                if not self.valid_file(image_path):
                    continue
                preprocessed_image = self.preprocess_image(image_path)
                if cursors[class_] == (len(images[class_]) - 1):
                    cursors[class_] = 0
                else:
                    cursors[class_] += 1
                yield step, image_path, preprocessed_image, label_dict[class_]
<|endoftext|>


<commit_before>import math
from numbskull.inference import FACTORS
from scipy import sparse
from snorkel.learning.gen_learning import GenerativeModel, DEP_EXCLUSIVE, DEP_REINFORCING, DEP_FIXING, DEP_SIMILAR
import unittest
import random
import numpy as np

class TestCategorical(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        pass

    @classmethod
    def tearDownClass(cls):
        pass

    def _test_categorical(self, scoped_categorical=False, cardinality=4):
        # A set of true priors
        tol = 0.1
        cardinality = 4
        LF_acc_priors = [0.75, 0.75, 0.75, 0.75, 0.9]
        LF_acc_prior_weights = map(lambda x: 0.5 * np.log((cardinality - 1.0) * x / (1 - x)), LF_acc_priors)
        label_prior = 1

        def get_lf(label, cardinality, acc):
            if random.random() < acc:
                return label + 1
            lf = random.randint(0, cardinality - 2)
            if (lf >= label):
                lf += 1
            return lf + 1

        # Defines a label matrix
        n = 10000
        L = sparse.lil_matrix((n, 5), dtype=np.int64)

        # Store the supervised gold labels separately
        labels = np.zeros(n, np.int64)

        for i in range(n):
            y = random.randint(0, cardinality - 1)
            # First four LFs always vote, and have decent acc
            L[i, 0] = get_lf(y, cardinality, LF_acc_priors[0])
            L[i, 1] = get_lf(y, cardinality, LF_acc_priors[1])
            L[i, 2] = get_lf(y, cardinality, LF_acc_priors[2])
            L[i, 3] = get_lf(y, cardinality, LF_acc_priors[3])

            # The fifth LF is very accurate but has a much smaller coverage
            if random.random() < 0.2:
                L[i, 4] = get_lf(y, cardinality, LF_acc_priors[4])

            # The sixth LF is a small supervised set
            if random.random() < 0.1:
                labels[i] = y + 1

        # Test with priors -- first check init vals are correct
        print("Testing init:")
        gen_model = GenerativeModel(lf_propensity=True)
        gen_model.train(
            L,
            LF_acc_prior_weights=LF_acc_prior_weights,
            labels=labels,
            reg_type=2,
            reg_param=1,
            epochs=0,
            scoped_categorical=scoped_categorical
        )
        stats = gen_model.learned_lf_stats()
        accs = stats["Accuracy"]
        print(accs)
        print(gen_model.weights.lf_propensity)
        priors = np.array(LF_acc_priors + [label_prior])
        self.assertTrue(np.all(np.abs(accs - priors) < tol))

        # Now test that estimated LF accs are not too far off
        print("\nTesting estimated LF accs (TOL=%s)" % tol)
        gen_model.train(
            L,
            LF_acc_prior_weights=LF_acc_prior_weights,
            labels=labels,
            reg_type=0,
            reg_param=0.0,
            scoped_categorical=scoped_categorical
        )
        stats = gen_model.learned_lf_stats()
        accs = stats["Accuracy"]
        coverage = stats["Coverage"]
        print(accs)
        print(coverage)
        priors = np.array(LF_acc_priors + [label_prior])
        self.assertTrue(np.all(np.abs(accs - priors) < tol))
        self.assertTrue(np.all(np.abs(coverage - np.array([1, 1, 1, 1, 0.2, 0.1]) < tol)))

        # Test without supervised
        print("\nTesting without supervised")
        gen_model = GenerativeModel(lf_propensity=True)
        gen_model.train(L, reg_type=0, scoped_categorical=scoped_categorical)
        stats = gen_model.learned_lf_stats()
        accs = stats["Accuracy"]
        coverage = stats["Coverage"]
        print(accs)
        print(coverage)
        priors = np.array(LF_acc_priors)
        self.assertTrue(np.all(np.abs(accs - priors) < tol))
        self.assertTrue(np.all(np.abs(coverage - np.array([1, 1, 1, 1, 0.2]) < tol)))

        # Test with supervised
        print("\nTesting with supervised, without priors")
        gen_model = GenerativeModel(lf_propensity=True)
        gen_model.train(
            L,
            labels=labels,
            reg_type=0,
            scoped_categorical=scoped_categorical
        )
        stats = gen_model.learned_lf_stats()
        accs = stats["Accuracy"]
        coverage = stats["Coverage"]
        print(accs)
        print(coverage)
        priors = np.array(LF_acc_priors + [label_prior])
        self.assertTrue(np.all(np.abs(accs - priors) < tol))
        self.assertTrue(np.all(np.abs(coverage - np.array([1, 1, 1, 1, 0.2, 0.1]) < tol)))

        # Test without supervised, and (intentionally) bad priors, but weak strength
        print("\nTesting without supervised, with bad priors (weak)")
        gen_model = GenerativeModel(lf_propensity=True)
        bad_prior = [0.9, 0.8, 0.7, 0.6, 0.5]
        bad_prior_weights = map(lambda x: 0.5 * np.log((cardinality - 1.0) * x / (1 - x)), bad_prior)
        gen_model.train(
            L,
            LF_acc_prior_weights=bad_prior_weights,
            reg_type=0,
            scoped_categorical=scoped_categorical
        )
        stats = gen_model.learned_lf_stats()
        accs = stats["Accuracy"]
        coverage = stats["Coverage"]
        print(accs)
        print(coverage)
        priors = np.array(LF_acc_priors)
        self.assertTrue(np.all(np.abs(accs - priors) < tol))

        # Test without supervised, and (intentionally) bad priors
        print("\nTesting without supervised, with bad priors (strong)")
        gen_model = GenerativeModel(lf_propensity=True)
        gen_model.train(
            L,
            LF_acc_prior_weights=bad_prior_weights,
            reg_type=2,
            reg_param=100 * n,
            scoped_categorical=scoped_categorical
        )
        stats = gen_model.learned_lf_stats()
        accs = stats["Accuracy"]
        coverage = stats["Coverage"]
        print(accs)
        self.assertTrue(np.all(np.abs(accs - np.array(bad_prior)) < tol))

    def test_categorical(self):
        self._test_categorical()

    def test_scoped_categorical(self):
        # Repeat previous tests with scoped_categorical=True
        print("\n\nTesting scoped categorical with cardinality=4")
        self._test_categorical(scoped_categorical=True)
        print("\n\nTesting scoped categorical with cardinality=100")
        self._test_categorical(scoped_categorical=True, cardinality=100)

if __name__ == '__main__':
    unittest.main()
<commit_msg>slightly higher tolerance for scoped categorical<commit_after>import math
from numbskull.inference import FACTORS
from scipy import sparse
from snorkel.learning.gen_learning import GenerativeModel, DEP_EXCLUSIVE, DEP_REINFORCING, DEP_FIXING, DEP_SIMILAR
import unittest
import random
import numpy as np

class TestCategorical(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        pass

    @classmethod
    def tearDownClass(cls):
        pass

    def _test_categorical(self, scoped_categorical=False, cardinality=4):
        # A set of true priors
        tol = 0.1
        if scoped_categorical:
            tol = 0.15
        cardinality = 4
        LF_acc_priors = [0.75, 0.75, 0.75, 0.75, 0.9]
        LF_acc_prior_weights = map(lambda x: 0.5 * np.log((cardinality - 1.0) * x / (1 - x)), LF_acc_priors)
        label_prior = 1

        def get_lf(label, cardinality, acc):
            if random.random() < acc:
                return label + 1
            lf = random.randint(0, cardinality - 2)
            if (lf >= label):
                lf += 1
            return lf + 1

        # Defines a label matrix
        n = 10000
        L = sparse.lil_matrix((n, 5), dtype=np.int64)

        # Store the supervised gold labels separately
        labels = np.zeros(n, np.int64)

        for i in range(n):
            y = random.randint(0, cardinality - 1)
            # First four LFs always vote, and have decent acc
            L[i, 0] = get_lf(y, cardinality, LF_acc_priors[0])
            L[i, 1] = get_lf(y, cardinality, LF_acc_priors[1])
            L[i, 2] = get_lf(y, cardinality, LF_acc_priors[2])
            L[i, 3] = get_lf(y, cardinality, LF_acc_priors[3])

            # The fifth LF is very accurate but has a much smaller coverage
            if random.random() < 0.2:
                L[i, 4] = get_lf(y, cardinality, LF_acc_priors[4])

            # The sixth LF is a small supervised set
            if random.random() < 0.1:
                labels[i] = y + 1

        # Test with priors -- first check init vals are correct
        print("Testing init:")
        gen_model = GenerativeModel(lf_propensity=True)
        gen_model.train(
            L,
            LF_acc_prior_weights=LF_acc_prior_weights,
            labels=labels,
            reg_type=2,
            reg_param=1,
            epochs=0,
            scoped_categorical=scoped_categorical
        )
        stats = gen_model.learned_lf_stats()
        accs = stats["Accuracy"]
        print(accs)
        print(gen_model.weights.lf_propensity)
        priors = np.array(LF_acc_priors + [label_prior])
        self.assertTrue(np.all(np.abs(accs - priors) < tol))

        # Now test that estimated LF accs are not too far off
        print("\nTesting estimated LF accs (TOL=%s)" % tol)
        gen_model.train(
            L,
            LF_acc_prior_weights=LF_acc_prior_weights,
            labels=labels,
            reg_type=0,
            reg_param=0.0,
            scoped_categorical=scoped_categorical
        )
        stats = gen_model.learned_lf_stats()
        accs = stats["Accuracy"]
        coverage = stats["Coverage"]
        print(accs)
        print(coverage)
        priors = np.array(LF_acc_priors + [label_prior])
        self.assertTrue(np.all(np.abs(accs - priors) < tol))
        self.assertTrue(np.all(np.abs(coverage - np.array([1, 1, 1, 1, 0.2, 0.1]) < tol)))

        # Test without supervised
        print("\nTesting without supervised")
        gen_model = GenerativeModel(lf_propensity=True)
        gen_model.train(L, reg_type=0, scoped_categorical=scoped_categorical)
        stats = gen_model.learned_lf_stats()
        accs = stats["Accuracy"]
        coverage = stats["Coverage"]
        print(accs)
        print(coverage)
        priors = np.array(LF_acc_priors)
        self.assertTrue(np.all(np.abs(accs - priors) < tol))
        self.assertTrue(np.all(np.abs(coverage - np.array([1, 1, 1, 1, 0.2]) < tol)))

        # Test with supervised
        print("\nTesting with supervised, without priors")
        gen_model = GenerativeModel(lf_propensity=True)
        gen_model.train(
            L,
            labels=labels,
            reg_type=0,
            scoped_categorical=scoped_categorical
        )
        stats = gen_model.learned_lf_stats()
        accs = stats["Accuracy"]
        coverage = stats["Coverage"]
        print(accs)
        print(coverage)
        priors = np.array(LF_acc_priors + [label_prior])
        self.assertTrue(np.all(np.abs(accs - priors) < tol))
        self.assertTrue(np.all(np.abs(coverage - np.array([1, 1, 1, 1, 0.2, 0.1]) < tol)))

        # Test without supervised, and (intentionally) bad priors, but weak strength
        print("\nTesting without supervised, with bad priors (weak)")
        gen_model = GenerativeModel(lf_propensity=True)
        bad_prior = [0.9, 0.8, 0.7, 0.6, 0.5]
        bad_prior_weights = map(lambda x: 0.5 * np.log((cardinality - 1.0) * x / (1 - x)), bad_prior)
        gen_model.train(
            L,
            LF_acc_prior_weights=bad_prior_weights,
            reg_type=0,
            scoped_categorical=scoped_categorical
        )
        stats = gen_model.learned_lf_stats()
        accs = stats["Accuracy"]
        coverage = stats["Coverage"]
        print(accs)
        print(coverage)
        priors = np.array(LF_acc_priors)
        self.assertTrue(np.all(np.abs(accs - priors) < tol))

        # Test without supervised, and (intentionally) bad priors
        print("\nTesting without supervised, with bad priors (strong)")
        gen_model = GenerativeModel(lf_propensity=True)
        gen_model.train(
            L,
            LF_acc_prior_weights=bad_prior_weights,
            reg_type=2,
            reg_param=100 * n,
            scoped_categorical=scoped_categorical
        )
        stats = gen_model.learned_lf_stats()
        accs = stats["Accuracy"]
        coverage = stats["Coverage"]
        print(accs)
        self.assertTrue(np.all(np.abs(accs - np.array(bad_prior)) < tol))

    def test_categorical(self):
        self._test_categorical()

    def test_scoped_categorical(self):
        # Repeat previous tests with scoped_categorical=True
        print("\n\nTesting scoped categorical with cardinality=4")
        self._test_categorical(scoped_categorical=True)
        print("\n\nTesting scoped categorical with cardinality=100")
        self._test_categorical(scoped_categorical=True, cardinality=100)

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


<commit_before>from itertools import chain
from xml.etree import ElementTree as ET
import requests
from nala.structures.data import Document, Part
from nala.utils.cache import Cacheable

__author__ = 'Aleksandar'

class UniprotDocumentSelector(Cacheable):
    """
    Selects a list of pubmed IDs (articles) that are likely to have mutation mentions.

    Outline of the selection procedure:
    1. Select proteins given a uniprot query (by default Swiss-Prot human proteins)
    2. For each protein search for articles showing evidence of sequence variant of mutagenesis
    3. Return pubmed IDs (articles) associated with the evidence
    """

    def __init__(self):
        super().__init__()
        self.processed = set()
        self.uniprot_url = 'http://www.uniprot.org/uniprot/'

    def _get_uniprot_ids(self, query=None):
        if not query:
            query = '(annotation:(type:natural_variations) OR annotation:(type:mutagen))' \
                    ' AND reviewed:yes AND organism:"Homo sapiens (Human) [9606]"'
        params = {'query': query,
                  'columns': 'id',
                  'format': 'tab'}

        if query in self.cache:
            lines = self.cache[query]
        else:
            req = requests.get(self.uniprot_url, params)
            lines = req.text.splitlines()
            self.cache[query] = lines

        for uniprot_id in lines[1:]:  # skip first line
            yield uniprot_id

    def _get_pubmed_ids_for_protein(self, uniprot_id):
        if uniprot_id in self.cache:
            text = self.cache[uniprot_id]
        else:
            req = requests.get(self.uniprot_url + '{}.xml'.format(uniprot_id))
            text = req.text
            self.cache[uniprot_id] = text

        xml = ET.fromstring(text)
        ns = {'u': 'http://uniprot.org/uniprot'}  # namespace

        evidence_ids = []
        for elem in xml.findall('.//u:feature[@evidence]', ns):
            if elem.attrib['type'] in ('sequence variant', 'mutagenesis site'):
                evidence_ids.append(elem.attrib['evidence'])

        for elem in xml.findall('.//u:evidence[@key="{}"]/u:source/u:dbReference[@type="PubMed"]'.format(18), ns):
            pubmed_id = elem.attrib['id']
            # check to see if we have seen this id before
            if pubmed_id not in self.processed:
                self.processed.add(pubmed_id)
                yield pubmed_id

    def get_pubmed_ids(self):
        for uniprot_id in self._get_uniprot_ids():
            for pubmed_id in self._get_pubmed_ids_for_protein(uniprot_id):
                yield pubmed_id

class DownloadArticle(Cacheable):
    """
    A utility generator that for a given iterable of PMIDs generates Document objects
    created by downloading the articles associated with the pmid.
    """

    def __init__(self, one_part=False):
        super().__init__()
        self.one_part = one_part
        """whether to put everything (title, abstract, etc.) under the same part joined with new line"""
        self.pubmed_url = 'http://eutils.ncbi.nlm.nih.gov/entrez/eutils/efetch.fcgi'

    def download(self, pmids):
        for pmid in pmids:
            if pmid in self.cache:
                xml = ET.fromstring(self.cache[pmid])
            else:
                req = requests.get(self.pubmed_url, {'db': 'pubmed', 'retmode': 'xml', 'id': pmid})
                text = req.text
                xml = ET.fromstring(text)
                self.cache[pmid] = text

            doc = Document()

            if self.one_part:
                joined_text = '\n'.join(element.text for element in
                                        chain(xml.findall('.//ArticleTitle'), xml.findall('.//AbstractText')))
                doc.parts['title_and_abstract'] = Part(joined_text)
            else:
                counter = 0
                # for now only include title and abstract
                for elem in xml.findall('.//ArticleTitle'):
                    doc.parts['part_{}'.format(counter)] = Part(elem.text)
                    counter += 1
                for elem in xml.findall('.//AbstractText'):
                    doc.parts['part_{}'.format(counter)] = Part(elem.text)
                    counter += 1

            # yield the document but only if you found anything
            if len(doc.parts) > 0:
                yield pmid, doc

def generate_documents(n):
    """
    Generate a given number of documents for bootstrapping applying the default filters.

    :param n: how many documents do you want
    :type n: int
    :returns: nala.structures.data.Dataset
    """
    from nala.structures.data import Dataset
    from nala.structures.selection_pipelines import DocumentSelectorPipeline
    from itertools import count
    c = count(1)

    dataset = Dataset()
    with DocumentSelectorPipeline() as dsp:
        for pmid, document in dsp.execute():
            dataset.documents[pmid] = document
            # if we have generated enough documents stop
            if next(c) == n:
                break

    return dataset<commit_msg> [#141] Change article download<commit_after>from itertools import chain
from xml.etree import ElementTree as ET
import requests
from nala.structures.data import Document, Part
from nala.utils.cache import Cacheable

__author__ = 'Aleksandar'

class UniprotDocumentSelector(Cacheable):
    """
    Selects a list of pubmed IDs (articles) that are likely to have mutation mentions.

    Outline of the selection procedure:
    1. Select proteins given a uniprot query (by default Swiss-Prot human proteins)
    2. For each protein search for articles showing evidence of sequence variant of mutagenesis
    3. Return pubmed IDs (articles) associated with the evidence
    """

    def __init__(self):
        super().__init__()
        self.processed = set()
        self.uniprot_url = 'http://www.uniprot.org/uniprot/'

    def _get_uniprot_ids(self, query=None):
        if not query:
            query = '(annotation:(type:natural_variations) OR annotation:(type:mutagen))' \
                    ' AND reviewed:yes AND organism:"Homo sapiens (Human) [9606]"'
        params = {'query': query,
                  'columns': 'id',
                  'format': 'tab'}

        if query in self.cache:
            lines = self.cache[query]
        else:
            req = requests.get(self.uniprot_url, params)
            lines = req.text.splitlines()
            self.cache[query] = lines

        for uniprot_id in lines[1:]:  # skip first line
            yield uniprot_id

    def _get_pubmed_ids_for_protein(self, uniprot_id):
        if uniprot_id in self.cache:
            text = self.cache[uniprot_id]
        else:
            req = requests.get(self.uniprot_url + '{}.xml'.format(uniprot_id))
            text = req.text
            self.cache[uniprot_id] = text

        xml = ET.fromstring(text)
        ns = {'u': 'http://uniprot.org/uniprot'}  # namespace

        evidence_ids = []
        for elem in xml.findall('.//u:feature[@evidence]', ns):
            if elem.attrib['type'] in ('sequence variant', 'mutagenesis site'):
                evidence_ids.append(elem.attrib['evidence'])

        for elem in xml.findall('.//u:evidence[@key="{}"]/u:source/u:dbReference[@type="PubMed"]'.format(18), ns):
            pubmed_id = elem.attrib['id']
            # check to see if we have seen this id before
            if pubmed_id not in self.processed:
                self.processed.add(pubmed_id)
                yield pubmed_id

    def get_pubmed_ids(self):
        for uniprot_id in self._get_uniprot_ids():
            for pubmed_id in self._get_pubmed_ids_for_protein(uniprot_id):
                yield pubmed_id

class DownloadArticle(Cacheable):
    """
    A utility generator that for a given iterable of PMIDs generates Document objects
    created by downloading the articles associated with the pmid.
    """

    def __init__(self, one_part=False):
        super().__init__()
        self.one_part = one_part
        """whether to put everything (title, abstract, etc.) under the same part joined with new line"""
        self.pubmed_url = 'http://eutils.ncbi.nlm.nih.gov/entrez/eutils/efetch.fcgi'

    def download(self, pmids):
        for pmid in pmids:
            if pmid in self.cache:
                xml = ET.fromstring(self.cache[pmid])
            else:
                req = requests.get(self.pubmed_url, {'db': 'pubmed', 'retmode': 'xml', 'id': pmid})
                text = req.text
                xml = ET.fromstring(text)
                self.cache[pmid] = text

            doc = Document()

            if self.one_part:
                joined_text = '\n'.join(element.text for element in
                                        chain(xml.findall('.//ArticleTitle'), xml.findall('.//AbstractText')))
                doc.parts['title_and_abstract'] = Part(joined_text)
            else:
                # for now only include title and abstract
                doc.parts['title'] = Part(xml.find('.//ArticleTitle').text)
                abstract = []
                for elem in xml.findall('.//AbstractText'):
                    if 'Label' in elem.attrib:
                        abstract.append('{}: {}'.format(elem.attrib['Label'], elem.text))
                    else:
                        abstract.append(elem.text)

                doc.parts['abstract'] = Part(' '.join(abstract))

            # yield the document but only if you found anything
            if len(doc.parts) > 0:
                yield pmid, doc

def generate_documents(n):
    """
    Generate a given number of documents for bootstrapping applying the default filters.

    :param n: how many documents do you want
    :type n: int
    :returns: nala.structures.data.Dataset
    """
    from nala.structures.data import Dataset
    from nala.structures.selection_pipelines import DocumentSelectorPipeline
    from itertools import count
    c = count(1)

    dataset = Dataset()
    with DocumentSelectorPipeline() as dsp:
        for pmid, document in dsp.execute():
            dataset.documents[pmid] = document
            # if we have generated enough documents stop
            if next(c) == n:
                break

    return dataset<|endoftext|>


<commit_before>''' SmartAPI Configuration '''

from biothings.web.settings.default import ES_RESULT_TRANSFORM, QUERY_KWARGS

from config_key import *

# *****************************************************************************
# Credentials
# *****************************************************************************
# Define in <project_folder>/config_key.py:
#   COOKIE_SECRET = '<Any Random String>'
#   GITHUB_CLIENT_ID = '<your Github application Client ID>'
#   GITHUB_CLIENT_SECRET = '<your Github application Client Secret>'

# *****************************************************************************
# User Input Control
# *****************************************************************************
QUERY_KWARGS['*']['authors'] = {'type': list, 'group': 'esqb'}
QUERY_KWARGS['*']['tags'] = {'type': list, 'group': 'esqb'}

# *****************************************************************************
# Elasticsearch
# *****************************************************************************
ES_INDICES = {
    'metadata': 'smartapi_docs'
}

# *****************************************************************************
# Tornado URL Patterns
# *****************************************************************************
APP_LIST = [
    (r'/api/?', 'handlers.api.APIHandler'),
    (r'/api/query/?', 'biothings.web.handlers.QueryHandler', {"biothing_type": "metadata"}),
    (r'/api/validate/?', 'handlers.api.ValidateHandler'),
    (r'/api/metadata/?', 'handlers.api.APIHandler'),
    (r'/api/metadata/(.+)/?', 'handlers.api.APIHandler'),
    (r'/api/suggestion/?', 'handlers.api.ValueSuggestionHandler'),
]

# biothings web tester will read this
API_VERSION = ''
API_PREFIX = 'api'

# *****************************************************************************
# Biothings SDK Settings
# *****************************************************************************
ACCESS_CONTROL_ALLOW_METHODS = 'HEAD,GET,POST,DELETE,PUT,OPTIONS'
ES_QUERY_BUILDER = "pipeline.SmartAPIQueryBuilder"
ES_RESULT_TRANSFORM = "pipeline.SmartAPIResultTransform"
DISABLE_CACHING = True
<commit_msg>Remove unused import<commit_after>''' SmartAPI Configuration '''

from biothings.web.settings.default import QUERY_KWARGS

from config_key import *

# *****************************************************************************
# Credentials
# *****************************************************************************
# Define in <project_folder>/config_key.py:
#   COOKIE_SECRET = '<Any Random String>'
#   GITHUB_CLIENT_ID = '<your Github application Client ID>'
#   GITHUB_CLIENT_SECRET = '<your Github application Client Secret>'

# *****************************************************************************
# User Input Control
# *****************************************************************************
QUERY_KWARGS['*']['authors'] = {'type': list, 'group': 'esqb'}
QUERY_KWARGS['*']['tags'] = {'type': list, 'group': 'esqb'}

# *****************************************************************************
# Elasticsearch
# *****************************************************************************
ES_INDICES = {
    'metadata': 'smartapi_docs'
}

# *****************************************************************************
# Tornado URL Patterns
# *****************************************************************************
APP_LIST = [
    (r'/api/?', 'handlers.api.APIHandler'),
    (r'/api/query/?', 'biothings.web.handlers.QueryHandler', {"biothing_type": "metadata"}),
    (r'/api/validate/?', 'handlers.api.ValidateHandler'),
    (r'/api/metadata/?', 'handlers.api.APIHandler'),
    (r'/api/metadata/(.+)/?', 'handlers.api.APIHandler'),
    (r'/api/suggestion/?', 'handlers.api.ValueSuggestionHandler'),
]

# biothings web tester will read this
API_VERSION = ''
API_PREFIX = 'api'

# *****************************************************************************
# Biothings SDK Settings
# *****************************************************************************
ACCESS_CONTROL_ALLOW_METHODS = 'HEAD,GET,POST,DELETE,PUT,OPTIONS'
ES_QUERY_BUILDER = "pipeline.SmartAPIQueryBuilder"
ES_RESULT_TRANSFORM = "pipeline.SmartAPIResultTransform"
DISABLE_CACHING = True
<|endoftext|>


<commit_before># ===============================================================================
# Copyright 2011 Jake Ross
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ===============================================================================

# ============= standard library imports ========================
from __future__ import absolute_import
from __future__ import print_function

import socket
import time

from six.moves import range
# ============= enthought library imports =======================
from traits.api import Float

# ============= local library imports  ==========================
from pychron.globals import globalv
from pychron.hardware.core.checksum_helper import computeCRC
from pychron.hardware.core.communicators.communicator import Communicator, process_response
from pychron.regex import IPREGEX

class MessageFrame(object):
    def __init__(self, message_len=False, nmessage_len=4, checksum=False, nchecksum=4):
        self.nchecksum = nchecksum
        self.checksum = checksum
        self.nmessage_len = nmessage_len
        self.message_len = message_len

    def set_str(self, s):
        """
        L4,-,C4
        """
        if s:
            args = s.split(',')
            if len(args) == 3:
                ml = args[0]
                cs = args[2]
                self.nmessage_len = int(ml[1:])
                self.nchecksum = int(cs[1:])
                self.checksum = True
                self.message_len = True

class Handler(object):
    sock = None
    datasize = 2 ** 12
    address = None
    message_frame = None

    def set_frame(self, f):
        self.message_frame = MessageFrame()
        if f:
            self.message_frame.set_str(f)

    def get_packet(self):
        raise NotImplementedError

    def send_packet(self, p):
        raise NotImplementedError

    def end(self):
        pass

    # private
    def _recvall(self, recv, datasize=None, frame=None):
        """
        recv: callable that accepts 1 argument (datasize). should return a str
        """
        # ss = []
        sum = 0

        # disable message len checking
        # msg_len = 1
        # if self.use_message_len_checking:
        # msg_len = 0

        msg_len = 1
        nm = -1

        if frame is None:
            frame = self.message_frame

        if frame.message_len:
            msg_len = 0
            nm = frame.nmessage_len

        if datasize is None:
            datasize = self.datasize

        data = b''
        while 1:
            s = recv(datasize)
            if not s:
                break

            if not msg_len:
                msg_len = int(s[:nm], 16)

            sum += len(s)
            data += s
            if sum >= msg_len:
                break

        if frame.message_len:
            # trim off header
            data = data[nm:]

        if frame.checksum:
            nc = frame.nchecksum
            checksum = data[-nc:]
            data = data[:-nc]
            comp = computeCRC(data)
            if comp != checksum:
                print('checksum fail computed={}, expected={}'.format(comp, checksum))
                return

        return data.decode('utf-8')

class TCPHandler(Handler):
    def open_socket(self, addr, timeout=1.0):
        self.address = addr
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        if globalv.communication_simulation:
            timeout = 0.01

        self.sock.settimeout(timeout)
        self.sock.connect(addr)

    def get_packet(self, datasize=None, message_frame=None):
        return self._recvall(self.sock.recv, datasize=datasize, frame=message_frame)

    def send_packet(self, p):
        self.sock.send(p.encode('utf-8'))

    def end(self):
        self.sock.close()

class UDPHandler(Handler):
    def open_socket(self, addr, timeout=1.0):
        self.address = addr
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        if globalv.communication_simulation:
            timeout = 0.01
        self.sock.settimeout(timeout)

    def get_packet(self, **kw):
        def recv(ds):
            rx, _ = self.sock.recvfrom(ds)
            return rx

        return self._recvall(recv)

    def send_packet(self, p):
        self.sock.sendto(p.encode('utf-8'), self.address)

class EthernetCommunicator(Communicator):
    """
    Communicator of UDP or TCP.
    """
    host = None
    port = None
    read_port = None
    handler = None
    kind = 'UDP'
    test_cmd = None
    use_end = False
    verbose = False
    error_mode = False
    message_frame = ''
    timeout = Float(1.0)

    default_timeout = 3

    _comms_report_attrs = ('host', 'port', 'kind', 'timeout')

    @property
    def address(self):
        return '{}://{}:{}'.format(self.kind, self.host, self.port)

    def load(self, config, path):
        """
        """
        super(EthernetCommunicator, self).load(config, path)

        self.host = self.config_get(config, 'Communications', 'host')
        if self.host != 'localhost' and not IPREGEX.match(self.host):
            result = socket.getaddrinfo(self.host, 0, 0, 0, 0)
            if result:
                for family, kind, a, b, host in result:
                    if family == socket.AF_INET and kind == socket.SOCK_STREAM:
                        self.host = host[0]

        # self.host = 'localhost'
        self.port = self.config_get(config, 'Communications', 'port', cast='int')
        self.read_port = self.config_get(config, 'Communications', 'read_port', cast='int', optional=True)
        self.timeout = self.config_get(config, 'Communications', 'timeout', cast='float', optional=True, default=1.0)
        self.kind = self.config_get(config, 'Communications', 'kind', optional=True)
        self.test_cmd = self.config_get(config, 'Communications', 'test_cmd', optional=True, default='')
        self.use_end = self.config_get(config, 'Communications', 'use_end', cast='boolean', optional=True,
                                       default=False)
        self.message_frame = self.config_get(config, 'Communications', 'message_frame', optional=True, default='')
        self.default_timeout = self.config_get(config, 'Communications', 'default_timeout', cast='int',
                                               optional=True, default=3)

        if self.kind is None:
            self.kind = 'UDP'

        return True

    def open(self, *args, **kw):

        for k in ('host', 'port', 'message_frame', 'kind'):
            if k in kw:
                setattr(self, k, kw[k])

        return self.test_connection()

    def test_connection(self):
        self.simulation = False

        with self._lock:
            handler = self.get_handler()

        # send a test command so see if wer have connection
        cmd = self.test_cmd

        if cmd:
            self.debug('sending test command {}'.format(cmd))
            r = self.ask(cmd)
            if r is None:
                self.simulation = True

                # if handler:
                #     if handler.send_packet(cmd):
                #         r = handler.get_packet(cmd)
                #         if r is None:
                #             self.simulation = True
                #     else:
                #         self.simulation = True
                # else:
                #     self.simulation = True
        ret = not self.simulation and handler is not None
        return ret

    def get_read_handler(self, handler):
        if self.read_port:
            handler = self.get_handler(addrs=(self.host, self.read_port))

        return handler

    def get_handler(self, addrs=None, timeout=None):
        if timeout is None:
            timeout = self.timeout
        if addrs is None:
            addrs = (self.host, self.port)

        try:
            h = self.handler
            if h is None:
                if self.kind.lower() == 'udp':
                    h = UDPHandler()
                else:
                    h = TCPHandler()

                # self.debug('get handler cmd={}, {},{} {}'.format(cmd.strip() if cmd is not None else '---', self.host,
                #                                                  self.port, timeout))
                h.open_socket(addrs, timeout=timeout)
                h.set_frame(self.message_frame)
                self.handler = h
            return h
        except socket.error as e:
            print('ewafs', e, self.host, self.port)
            self.debug('Get Handler {}. timeout={}. comms simulation={}'.format(str(e),
                                                                                timeout,
                                                                                globalv.communication_simulation))
            self.error_mode = True
            self.handler = None

    def ask(self, cmd, retries=3, verbose=True, quiet=False, info=None, timeout=None,
            message_frame=None, delay=None, use_error_mode=True, *args, **kw):
        """
        @param cmd: ASCII text to send
        @param retries: number of retries if command fails
        @param verbose: add to log
        @param quiet: if true do not log the response
        @param info: str to add to response
        @param timeout: timeout in seconds
        @param message_frame: MessageFrame object
        @param delay: delay in seconds to wait before a `cmd` is sent

        """

        if self.simulation:
            if verbose:
                self.info('no handle    {}'.format(cmd.strip()))
            return

        # print self.write_terminator
        cmd = '{}{}'.format(cmd, self.write_terminator)
        # print cmd
        # cmd = '{}\n'.format(cmd)
        r = None
        with self._lock:
            if use_error_mode and self.error_mode:
                retries = 2

            re = 'ERROR: Connection refused: {}, timeout={}'.format(self.address, timeout)
            for i in range(retries):
                r = self._ask(cmd, timeout=timeout, message_frame=message_frame, delay=delay,
                              use_error_mode=use_error_mode)
                if r is not None:
                    break
                else:
                    time.sleep(0.025)
                    self.debug('doing retry {}'.format(i))
                    # else:
                    #     self._reset_connection()

            if r is not None:
                re = process_response(r)
            # else:
            #     self.error_mode = True

            if self.use_end:
                # self.debug('ending connection. Handler: {}, cmd={}'.format(self.handler, cmd))
                if self.handler:
                    self.handler.end()
                self._reset_connection()

            if verbose or (self.verbose and not quiet):
                self.log_response(cmd, re, info)

        return r

    def reset(self):
        if self.handler:
            self.handler.end()
        self._reset_connection()

    def read(self, datasize=None, *args, **kw):
        with self._lock:
            handler = self.get_handler()
            if handler:
                return handler.get_packet(datasize=datasize)

    def tell(self, cmd, verbose=True, quiet=False, info=None):
        with self._lock:
            handler = self.get_handler()
            try:
                cmd = '{}{}'.format(cmd, self.write_terminator)
                handler.send_packet(cmd)
                if verbose or self.verbose and not quiet:
                    self.log_tell(cmd, info)
            except socket.error as e:
                self.warning('tell. send packet. error: {}'.format(e))
                self.error_mode = True

    # private
    def _reset_connection(self):
        self.handler = None
        self.error_mode = False

    def _ask(self, cmd, timeout=None, message_frame=None, delay=None, use_error_mode=True):
        if self.error_mode:
            self.handler = None
            if use_error_mode:
                timeout = 0.25

        if timeout is None:
            timeout = self.default_timeout

        self.error_mode = False

        handler = self.get_handler(timeout=timeout)
        if not handler:
            return

        try:
            handler.send_packet(cmd)

            if delay:
                time.sleep(delay)

            handler = self.get_read_handler(handler, timeout=timeout)
            try:
                return handler.get_packet(message_frame=message_frame)
            except socket.error as e:
                self.warning('ask. get packet. error: {} address: {}'.format(e, self.address))
                self.error_mode = True
        except socket.error as e:
            self.warning('ask. send packet. error: {} address: {}'.format(e, self.address))
            self.error_mode = True

# ============= EOF ====================================
<commit_msg>fixed get_read_handler<commit_after># ===============================================================================
# Copyright 2011 Jake Ross
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ===============================================================================

# ============= standard library imports ========================
from __future__ import absolute_import
from __future__ import print_function

import socket
import time

from six.moves import range
# ============= enthought library imports =======================
from traits.api import Float

# ============= local library imports  ==========================
from pychron.globals import globalv
from pychron.hardware.core.checksum_helper import computeCRC
from pychron.hardware.core.communicators.communicator import Communicator, process_response
from pychron.regex import IPREGEX

class MessageFrame(object):
    def __init__(self, message_len=False, nmessage_len=4, checksum=False, nchecksum=4):
        self.nchecksum = nchecksum
        self.checksum = checksum
        self.nmessage_len = nmessage_len
        self.message_len = message_len

    def set_str(self, s):
        """
        L4,-,C4
        """
        if s:
            args = s.split(',')
            if len(args) == 3:
                ml = args[0]
                cs = args[2]
                self.nmessage_len = int(ml[1:])
                self.nchecksum = int(cs[1:])
                self.checksum = True
                self.message_len = True

class Handler(object):
    sock = None
    datasize = 2 ** 12
    address = None
    message_frame = None

    def set_frame(self, f):
        self.message_frame = MessageFrame()
        if f:
            self.message_frame.set_str(f)

    def get_packet(self):
        raise NotImplementedError

    def send_packet(self, p):
        raise NotImplementedError

    def end(self):
        pass

    # private
    def _recvall(self, recv, datasize=None, frame=None):
        """
        recv: callable that accepts 1 argument (datasize). should return a str
        """
        # ss = []
        sum = 0

        # disable message len checking
        # msg_len = 1
        # if self.use_message_len_checking:
        # msg_len = 0

        msg_len = 1
        nm = -1

        if frame is None:
            frame = self.message_frame

        if frame.message_len:
            msg_len = 0
            nm = frame.nmessage_len

        if datasize is None:
            datasize = self.datasize

        data = b''
        while 1:
            s = recv(datasize)
            if not s:
                break

            if not msg_len:
                msg_len = int(s[:nm], 16)

            sum += len(s)
            data += s
            if sum >= msg_len:
                break

        if frame.message_len:
            # trim off header
            data = data[nm:]

        if frame.checksum:
            nc = frame.nchecksum
            checksum = data[-nc:]
            data = data[:-nc]
            comp = computeCRC(data)
            if comp != checksum:
                print('checksum fail computed={}, expected={}'.format(comp, checksum))
                return

        return data.decode('utf-8')

class TCPHandler(Handler):
    def open_socket(self, addr, timeout=1.0):
        self.address = addr
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        if globalv.communication_simulation:
            timeout = 0.01

        self.sock.settimeout(timeout)
        self.sock.connect(addr)

    def get_packet(self, datasize=None, message_frame=None):
        return self._recvall(self.sock.recv, datasize=datasize, frame=message_frame)

    def send_packet(self, p):
        self.sock.send(p.encode('utf-8'))

    def end(self):
        self.sock.close()

class UDPHandler(Handler):
    def open_socket(self, addr, timeout=1.0):
        self.address = addr
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        if globalv.communication_simulation:
            timeout = 0.01
        self.sock.settimeout(timeout)

    def get_packet(self, **kw):
        def recv(ds):
            rx, _ = self.sock.recvfrom(ds)
            return rx

        return self._recvall(recv)

    def send_packet(self, p):
        self.sock.sendto(p.encode('utf-8'), self.address)

class EthernetCommunicator(Communicator):
    """
    Communicator of UDP or TCP.
    """
    host = None
    port = None
    read_port = None
    handler = None
    kind = 'UDP'
    test_cmd = None
    use_end = False
    verbose = False
    error_mode = False
    message_frame = ''
    timeout = Float(1.0)

    default_timeout = 3

    _comms_report_attrs = ('host', 'port', 'kind', 'timeout')

    @property
    def address(self):
        return '{}://{}:{}'.format(self.kind, self.host, self.port)

    def load(self, config, path):
        """
        """
        super(EthernetCommunicator, self).load(config, path)

        self.host = self.config_get(config, 'Communications', 'host')
        if self.host != 'localhost' and not IPREGEX.match(self.host):
            result = socket.getaddrinfo(self.host, 0, 0, 0, 0)
            if result:
                for family, kind, a, b, host in result:
                    if family == socket.AF_INET and kind == socket.SOCK_STREAM:
                        self.host = host[0]

        # self.host = 'localhost'
        self.port = self.config_get(config, 'Communications', 'port', cast='int')
        self.read_port = self.config_get(config, 'Communications', 'read_port', cast='int', optional=True)
        self.timeout = self.config_get(config, 'Communications', 'timeout', cast='float', optional=True, default=1.0)
        self.kind = self.config_get(config, 'Communications', 'kind', optional=True)
        self.test_cmd = self.config_get(config, 'Communications', 'test_cmd', optional=True, default='')
        self.use_end = self.config_get(config, 'Communications', 'use_end', cast='boolean', optional=True,
                                       default=False)
        self.message_frame = self.config_get(config, 'Communications', 'message_frame', optional=True, default='')
        self.default_timeout = self.config_get(config, 'Communications', 'default_timeout', cast='int',
                                               optional=True, default=3)

        if self.kind is None:
            self.kind = 'UDP'

        return True

    def open(self, *args, **kw):

        for k in ('host', 'port', 'message_frame', 'kind'):
            if k in kw:
                setattr(self, k, kw[k])

        return self.test_connection()

    def test_connection(self):
        self.simulation = False

        with self._lock:
            handler = self.get_handler()

        # send a test command so see if wer have connection
        cmd = self.test_cmd

        if cmd:
            self.debug('sending test command {}'.format(cmd))
            r = self.ask(cmd)
            if r is None:
                self.simulation = True

                # if handler:
                #     if handler.send_packet(cmd):
                #         r = handler.get_packet(cmd)
                #         if r is None:
                #             self.simulation = True
                #     else:
                #         self.simulation = True
                # else:
                #     self.simulation = True
        ret = not self.simulation and handler is not None
        return ret

    def get_read_handler(self, handler, **kw):
        if self.read_port:
            handler = self.get_handler(addrs=(self.host, self.read_port), **kw)

        return handler

    def get_handler(self, addrs=None, timeout=None):
        if timeout is None:
            timeout = self.timeout
        if addrs is None:
            addrs = (self.host, self.port)

        try:
            h = self.handler
            if h is None:
                if self.kind.lower() == 'udp':
                    h = UDPHandler()
                else:
                    h = TCPHandler()

                # self.debug('get handler cmd={}, {},{} {}'.format(cmd.strip() if cmd is not None else '---', self.host,
                #                                                  self.port, timeout))
                h.open_socket(addrs, timeout=timeout)
                h.set_frame(self.message_frame)
                self.handler = h
            return h
        except socket.error as e:
            print('ewafs', e, self.host, self.port)
            self.debug('Get Handler {}. timeout={}. comms simulation={}'.format(str(e),
                                                                                timeout,
                                                                                globalv.communication_simulation))
            self.error_mode = True
            self.handler = None

    def ask(self, cmd, retries=3, verbose=True, quiet=False, info=None, timeout=None,
            message_frame=None, delay=None, use_error_mode=True, *args, **kw):
        """
        @param cmd: ASCII text to send
        @param retries: number of retries if command fails
        @param verbose: add to log
        @param quiet: if true do not log the response
        @param info: str to add to response
        @param timeout: timeout in seconds
        @param message_frame: MessageFrame object
        @param delay: delay in seconds to wait before a `cmd` is sent

        """

        if self.simulation:
            if verbose:
                self.info('no handle    {}'.format(cmd.strip()))
            return

        # print self.write_terminator
        cmd = '{}{}'.format(cmd, self.write_terminator)
        # print cmd
        # cmd = '{}\n'.format(cmd)
        r = None
        with self._lock:
            if use_error_mode and self.error_mode:
                retries = 2

            re = 'ERROR: Connection refused: {}, timeout={}'.format(self.address, timeout)
            for i in range(retries):
                r = self._ask(cmd, timeout=timeout, message_frame=message_frame, delay=delay,
                              use_error_mode=use_error_mode)
                if r is not None:
                    break
                else:
                    time.sleep(0.025)
                    self.debug('doing retry {}'.format(i))
                    # else:
                    #     self._reset_connection()

            if r is not None:
                re = process_response(r)
            # else:
            #     self.error_mode = True

            if self.use_end:
                # self.debug('ending connection. Handler: {}, cmd={}'.format(self.handler, cmd))
                if self.handler:
                    self.handler.end()
                self._reset_connection()

            if verbose or (self.verbose and not quiet):
                self.log_response(cmd, re, info)

        return r

    def reset(self):
        if self.handler:
            self.handler.end()
        self._reset_connection()

    def read(self, datasize=None, *args, **kw):
        with self._lock:
            handler = self.get_handler()
            if handler:
                return handler.get_packet(datasize=datasize)

    def tell(self, cmd, verbose=True, quiet=False, info=None):
        with self._lock:
            handler = self.get_handler()
            try:
                cmd = '{}{}'.format(cmd, self.write_terminator)
                handler.send_packet(cmd)
                if verbose or self.verbose and not quiet:
                    self.log_tell(cmd, info)
            except socket.error as e:
                self.warning('tell. send packet. error: {}'.format(e))
                self.error_mode = True

    # private
    def _reset_connection(self):
        self.handler = None
        self.error_mode = False

    def _ask(self, cmd, timeout=None, message_frame=None, delay=None, use_error_mode=True):
        if self.error_mode:
            self.handler = None
            if use_error_mode:
                timeout = 0.25

        if timeout is None:
            timeout = self.default_timeout

        self.error_mode = False

        handler = self.get_handler(timeout=timeout)
        if not handler:
            return

        try:
            handler.send_packet(cmd)

            if delay:
                time.sleep(delay)

            handler = self.get_read_handler(handler, timeout=timeout)
            try:
                return handler.get_packet(message_frame=message_frame)
            except socket.error as e:
                self.warning('ask. get packet. error: {} address: {}'.format(e, self.address))
                self.error_mode = True
        except socket.error as e:
            self.warning('ask. send packet. error: {} address: {}'.format(e, self.address))
            self.error_mode = True

# ============= EOF ====================================
<|endoftext|>