<commit_before># Licensed under a 3-clause BSD style license - see LICENSE.rst
"""
Tools to build an ePSF.
"""

from __future__ import (absolute_import, division, print_function,
                        unicode_literals)

import numpy as np

from astropy.stats import SigmaClip

__all__ = ['EPSFBuilder']

class EPSFBuilder(object):
    def __init__(self, peak_fit_box=5, peak_search_box='fitbox',
                 recenter_accuracy=1.0e-4, recenter_max_iters=1000,
                 ignore_badfit_stars=True, stat='median',
                 sigma_clip=SigmaClip(sigma=3., iters=10),
                 smoothing_kernel='quar', fitter=EPSFFitter(residuals=True),
                 max_iters=50, accuracy=1e-4, epsf=None):

        self.peak_fit_box = peak_fit_box
        self.peak_search_box = peak_search_box
        self.recenter_accuracy = recenter_accuracy
        self.recenter_max_iters = recenter_max_iters
        self.ignore_badfit_stars = ignore_badfit_stars
        self.stat = stat
        self.sigma_clip = sigma_clip
        self.smoothing_kernel = smoothing_kernel
        self.fitter = fitter

        max_iters = int(max_iters)
        if max_iters < 0:
            raise ValueErro('max_iters must be non-negative.')
        self.max_iters = max_iters

        if accuracy <= 0.0:
            raise ValueError('accuracy must be a positive number.')
        self.accuracy = accuracy

        self.epsf = epsf

    def __call__(self, data, star_table):
        return self.build_psf(data, star_table)

    def _build_psf_step(self, psf=None):
        if len(stars) < 1:
            raise ValueError("'stars' must be a list containing at least one "
                            "'Star' object.")

        if psf is None:
            psf = init_psf(stars)
        elif isinstance(psf, type):
            psf = init_psf(stars, psf_cls=psf)
        else:
            psf = copy.deepcopy(psf)

        recenter_accuracy = float(recenter_accuracy)
        if recenter_accuracy <= 0.0:
            raise ValueError("Re-center accuracy must be a strictly positive "
                            "number.")

        recenter_nmax = int(recenter_nmax)
        if recenter_nmax < 0:
            raise ValueError("Maximum number of re-recntering iterations must be "
                            "a non-negative integer number.")

        cx, cy = psf.origin
        ny, nx = psf.shape
        pscale_x, pscale_y = psf.pixel_scale

        # get all stars including linked stars as a flat list:
        all_stars = []
        for s in stars:
            all_stars += s.get_linked_list()

        # allocate "accumulator" array (to store transformed PSFs):
        apsf = [[[] for k in range(nx)] for k in range(ny)]

        norm_psf_data = psf.normalized_data

        for s in all_stars:
            if s.ignore:
                continue

            if ignore_badfit_stars and s.fit_error_status is not None and \
            s.fit_error_status > 0:
                continue

            pixlist = s.centered_plist(normalized=True)

            # evaluate previous PSF model at star pixel location in the PSF grid:
            ovx = s.pixel_scale[0] / pscale_x
            ovy = s.pixel_scale[1] / pscale_y
            x = ovx * (pixlist[:, 0])
            y = ovy * (pixlist[:, 1])
            old_model_vals = psf.evaluate(x=x, y=y, flux=1.0, x_0=0.0, y_0=0.0)

            # find integer location of star pixels in the PSF grid
            # and compute residuals
            ix = py2round(x + cx).astype(np.int)
            iy = py2round(y + cy).astype(np.int)
            pv = pixlist[:, 2] / (ovx * ovy) - old_model_vals
            m = np.logical_and(
                np.logical_and(ix >= 0, ix < nx),
                np.logical_and(iy >= 0, iy < ny)
            )

            # add all pixel values to the corresponding accumulator:
            for i, j, v in zip(ix[m], iy[m], pv[m]):
                apsf[j][i].append(v)

        psfdata = np.empty((ny, nx), dtype=np.float)
        psfdata.fill(np.nan)

        for i in range(nx):
            for j in range(ny):
                psfdata[j, i] = _pixstat(
                    apsf[j][i], stat=stat, nclip=nclip, lsig=lsig, usig=usig,
                    default=np.nan
                )

        mask = np.isfinite(psfdata)
        if not np.all(mask):
            # fill in the "holes" (=np.nan) using interpolation
            # I. using cubic spline for inner holes
            psfdata = interpolate_missing_data(psfdata, method='cubic',
                                            mask=mask)

            # II. we fill outer points with zeros
            mask = np.isfinite(psfdata)
            psfdata[np.logical_not(mask)] = 0.0

        # add residuals to old PSF data:
        psfdata += norm_psf_data

        # apply a smoothing kernel to the PSF:
        psfdata = _smoothPSF(psfdata, ker)

        shift_x = 0
        shift_y = 0
        peak_eps2 = recenter_accuracy**2
        eps2_prev = None
        y, x = np.indices(psfdata.shape, dtype=np.float)
        ePSF = psf.make_similar_from_data(psfdata)
        ePSF.fill_value = 0.0

        for iteration in range(recenter_nmax):
            # find peak location:
            peak_x, peak_y = find_peak(
                psfdata, xmax=cx, ymax=cy, peak_fit_box=peak_fit_box,
                peak_search_box=peak_search_box, mask=None
            )

            dx = cx - peak_x
            dy = cy - peak_y

            eps2 = dx**2 + dy**2
            if (eps2_prev is not None and eps2 > eps2_prev) or eps2 < peak_eps2:
                break
            eps2_prev = eps2

            shift_x += dx
            shift_y += dy

            # Resample PSF data to a shifted grid such that the pick of the PSF is
            # at expected position.
            psfdata = ePSF.evaluate(x=x, y=y, flux=1.0,
                                    x_0=shift_x + cx, y_0=shift_y + cy)

        # apply final shifts and fill in any missing data:
        if shift_x != 0.0 or shift_y != 0.0:
            ePSF.fill_value = np.nan
            psfdata = ePSF.evaluate(x=x, y=y, flux=1.0,
                                    x_0=shift_x + cx, y_0=shift_y + cy)

            # fill in the "holes" (=np.nan) using 0 (no contribution to the flux):
            mask = np.isfinite(psfdata)
            psfdata[np.logical_not(mask)] = 0.0

        norm = np.abs(np.sum(psfdata, dtype=np.float64))
        psfdata /= norm

        # Create ePSF model and return:
        ePSF = psf.make_similar_from_data(psfdata)

        return ePSF

    def build_psf(self, data, star_table):
        """
        Iteratively build the psf.
        """

        self._extract_stars(data, star_table, common_catalog=None,
                            extract_size=11, recenter=False,
                            peak_fit_box=5, peak_search_box='fitbox',
                            catmap={'x': 'x', 'y': 'y', 'lon': 'lon',
                                    'lat': 'lat', 'weight': 'weight',
                                    'id': 'id'}, cat_name_kwd='name',
                            image_name_kwd='name')

        # get all stars including linked stars as a flat list:
        all_stars = []
        for s in stars:
            all_stars += s.get_linked_list()

        # create an array of star centers:
        prev_centers = np.asarray([s.center for s in stars], dtype=np.float)

        # initialize array for detection of scillatory behavior:
        oscillatory = np.zeros((len(all_stars), ), dtype=np.bool)
        prev_failed = np.zeros((len(all_stars), ), dtype=np.bool)

        niter = -1
        eps2 = 2.0 * acc2
        dxy = np.zeros((len(all_stars), 2), dtype=np.float)

        while niter < nmax and np.amax(eps2) >= acc2 and not np.all(oscillatory):
            niter += 1

            # improved PSF:
            psf = build_psf(
                stars=stars,
                psf=psf,
                peak_fit_box=peak_fit_box,
                peak_search_box=peak_search_box,
                recenter_accuracy=recenter_accuracy,
                recenter_nmax=recenter_nmax,
                ignore_badfit_stars=ignore_badfit_stars,
                stat=stat,
                nclip=nclip,
                lsig=lsig,
                usig=usig,
                ker=ker
            )

            # improved fit of the PSF to stars:
            stars = fit_stars(
                stars=stars,
                psf=psf,
                psf_fit_box=psf_fit_box,
                fitter=fitter,
                fitter_kwargs=fitter_kwargs,
                residuals=False
            )

            # get all stars including linked stars as a flat list:
            all_stars = []
            for s in stars:
                all_stars += s.get_linked_list()

            # create an array of star centers at this iteration:
            centers = np.asarray([s.center for s in stars], dtype=np.float)

            # detect oscillatory behavior
            failed = np.array([s.fit_error_status > 0 for s in stars],
                            dtype=np.bool)
            if niter > 2:  # allow a few iterations at the beginning
                oscillatory = np.logical_and(prev_failed, np.logical_not(failed))
                for s, osc in zip(stars, oscillatory):
                    s.ignore = bool(osc)
                prev_failed = failed

            # check termination criterion:
            good_mask = np.logical_not(np.logical_or(failed, oscillatory))
            dxy = centers - prev_centers
            mdxy = dxy[good_mask]
            eps2 = np.sum(mdxy * mdxy, axis=1, dtype=np.float64)
            prev_centers = centers

        # compute residuals:
        if residuals:
            res = compute_residuals(psf, all_stars)
        else:
            res = len(all_stars) * [None]

        for s, r in zip(all_stars, res):
            s.fit_residual = r

        # assign coordinate residuals of the iterative process:
        for s, (dx, dy) in zip(all_stars, dxy):
            s.iter_fit_eps = (float(dx), float(dy))

        return (psf, stars, niter)
<commit_msg>Add init_psf to epsf builder<commit_after># Licensed under a 3-clause BSD style license - see LICENSE.rst
"""
Tools to build an ePSF.
"""

from __future__ import (absolute_import, division, print_function,
                        unicode_literals)

import numpy as np
from astropy.stats import SigmaClip
from astropy.table import Table

from .epsf_fitter import EPSFFitter, compute_residuals

__all__ = ['EPSFBuilder']

class EPSFBuilder(object):
    def __init__(self, peak_fit_box=5, peak_search_box='fitbox',
                 recenter_accuracy=1.0e-4, recenter_max_iters=1000,
                 ignore_badfit_stars=True, stat='median',
                 sigma_clip=SigmaClip(sigma=3., iters=10),
                 smoothing_kernel='quar', fitter=EPSFFitter(residuals=True),
                 max_iters=50, accuracy=1e-4, epsf=None):

        self.peak_fit_box = peak_fit_box
        self.peak_search_box = peak_search_box
        self.recenter_accuracy = recenter_accuracy
        self.recenter_max_iters = recenter_max_iters
        self.ignore_badfit_stars = ignore_badfit_stars
        self.stat = stat
        self.sigma_clip = sigma_clip
        self.smoothing_kernel = smoothing_kernel
        self.fitter = fitter

        max_iters = int(max_iters)
        if max_iters <= 0:
            raise ValueError('max_iters must be a positive number.')
        self.max_iters = max_iters

        if accuracy <= 0.0:
            raise ValueError('accuracy must be a positive number.')
        self.accuracy2 = accuracy**2

        self.epsf = epsf

    def __call__(self, data, stars):
        return self.build_psf(data, stars)

    def _build_psf_step(self, psf=None):
        if len(stars) < 1:
            raise ValueError('stars must be a list containing at least '
                             'one Star object.')

        if psf is None:
            psf = init_psf(stars)
        elif isinstance(psf, type):
            psf = init_psf(stars, psf_cls=psf)
        else:
            psf = copy.deepcopy(psf)

        recenter_accuracy = float(recenter_accuracy)
        if recenter_accuracy <= 0.0:
            raise ValueError("Re-center accuracy must be a strictly positive "
                            "number.")

        recenter_nmax = int(recenter_nmax)
        if recenter_nmax < 0:
            raise ValueError("Maximum number of re-recntering iterations must be "
                            "a non-negative integer number.")

        cx, cy = psf.origin
        ny, nx = psf.shape
        pscale_x, pscale_y = psf.pixel_scale

        # get all stars including linked stars as a flat list:
        all_stars = []
        for s in stars:
            all_stars += s.get_linked_list()

        # allocate "accumulator" array (to store transformed PSFs):
        apsf = [[[] for k in range(nx)] for k in range(ny)]

        norm_psf_data = psf.normalized_data

        for s in all_stars:
            if s.ignore:
                continue

            if ignore_badfit_stars and s.fit_error_status is not None and \
            s.fit_error_status > 0:
                continue

            pixlist = s.centered_plist(normalized=True)

            # evaluate previous PSF model at star pixel location in the PSF grid:
            ovx = s.pixel_scale[0] / pscale_x
            ovy = s.pixel_scale[1] / pscale_y
            x = ovx * (pixlist[:, 0])
            y = ovy * (pixlist[:, 1])
            old_model_vals = psf.evaluate(x=x, y=y, flux=1.0, x_0=0.0, y_0=0.0)

            # find integer location of star pixels in the PSF grid
            # and compute residuals
            ix = py2round(x + cx).astype(np.int)
            iy = py2round(y + cy).astype(np.int)
            pv = pixlist[:, 2] / (ovx * ovy) - old_model_vals
            m = np.logical_and(
                np.logical_and(ix >= 0, ix < nx),
                np.logical_and(iy >= 0, iy < ny)
            )

            # add all pixel values to the corresponding accumulator:
            for i, j, v in zip(ix[m], iy[m], pv[m]):
                apsf[j][i].append(v)

        psfdata = np.empty((ny, nx), dtype=np.float)
        psfdata.fill(np.nan)

        for i in range(nx):
            for j in range(ny):
                psfdata[j, i] = _pixstat(
                    apsf[j][i], stat=stat, nclip=nclip, lsig=lsig, usig=usig,
                    default=np.nan
                )

        mask = np.isfinite(psfdata)
        if not np.all(mask):
            # fill in the "holes" (=np.nan) using interpolation
            # I. using cubic spline for inner holes
            psfdata = interpolate_missing_data(psfdata, method='cubic',
                                            mask=mask)

            # II. we fill outer points with zeros
            mask = np.isfinite(psfdata)
            psfdata[np.logical_not(mask)] = 0.0

        # add residuals to old PSF data:
        psfdata += norm_psf_data

        # apply a smoothing kernel to the PSF:
        psfdata = _smoothPSF(psfdata, ker)

        shift_x = 0
        shift_y = 0
        peak_eps2 = recenter_accuracy**2
        eps2_prev = None
        y, x = np.indices(psfdata.shape, dtype=np.float)
        ePSF = psf.make_similar_from_data(psfdata)
        ePSF.fill_value = 0.0

        for iteration in range(recenter_nmax):
            # find peak location:
            peak_x, peak_y = find_peak(
                psfdata, xmax=cx, ymax=cy, peak_fit_box=peak_fit_box,
                peak_search_box=peak_search_box, mask=None
            )

            dx = cx - peak_x
            dy = cy - peak_y

            eps2 = dx**2 + dy**2
            if (eps2_prev is not None and eps2 > eps2_prev) or eps2 < peak_eps2:
                break
            eps2_prev = eps2

            shift_x += dx
            shift_y += dy

            # Resample PSF data to a shifted grid such that the pick of the PSF is
            # at expected position.
            psfdata = ePSF.evaluate(x=x, y=y, flux=1.0,
                                    x_0=shift_x + cx, y_0=shift_y + cy)

        # apply final shifts and fill in any missing data:
        if shift_x != 0.0 or shift_y != 0.0:
            ePSF.fill_value = np.nan
            psfdata = ePSF.evaluate(x=x, y=y, flux=1.0,
                                    x_0=shift_x + cx, y_0=shift_y + cy)

            # fill in the "holes" (=np.nan) using 0 (no contribution to the flux):
            mask = np.isfinite(psfdata)
            psfdata[np.logical_not(mask)] = 0.0

        norm = np.abs(np.sum(psfdata, dtype=np.float64))
        psfdata /= norm

        # Create ePSF model and return:
        ePSF = psf.make_similar_from_data(psfdata)

        return ePSF

    def build_psf(self, data, stars, psf=None):
        """
        Iteratively build the psf.
        """

        if isinstance(stars, Table):
            self.stars = self._extract_stars(
                data, stars, common_catalog=None, extract_size=11,
                recenter=False, peak_fit_box=5, peak_search_box='fitbox',
                catmap={'x': 'x', 'y': 'y', 'lon': 'lon', 'lat': 'lat',
                        'weight': 'weight', 'id': 'id'}, cat_name_kwd='name',
                image_name_kwd='name')
        else:
            # TODO: check if stars is a list of Stars
            self.stars = stars

        # get all stars (including linked stars) as a flat list
        all_stars = []
        for s in self.stars:
            all_stars += s.get_linked_list()
        nstars = len(all_stars)

        # create an array of star centers
        prev_centers = np.array([s.center for s in stars], dtype=np.float)

        # initialize array for detection of oscillatory behavior
        oscillatory = np.zeros(nstars, dtype=bool)
        prev_failed = np.zeros(nstars, dtype=bool)
        dxy = np.zeros((nstars, 2), dtype=np.float)

        iter_num = 0
        eps2 = 2.0 * self.accuracy2
        while (iter_num < self.max_iters and np.amax(eps2) >= self.accuracy2
               and not np.all(oscillatory)):
            iter_num += 1

            # build/improve the PSF
            psf = self._build_psf_step(psf=psf)

            # fit the new PSF to the stars to find improved centers
            stars = self.fitter(stars, psf)

            # get all stars (including linked stars) as a flat list
            all_stars = []
            for s in stars:
                all_stars += s.get_linked_list()

            # create an array of star centers at this iteration
            centers = np.array([s.center for s in stars], dtype=np.float)

            # detect oscillatory behavior
            failed = np.array([s.fit_error_status > 0 for s in stars],
                              dtype=np.bool)

            # exclude oscillatory stars after 3 iterations
            if iter_num > 3:
                oscillatory = np.logical_and(prev_failed,
                                             np.logical_not(failed))
                for s, osc in zip(stars, oscillatory):
                    s.ignore = bool(osc)
                prev_failed = failed

            # check termination criterion
            good_mask = np.logical_not(np.logical_or(failed, oscillatory))
            dxy = centers - prev_centers
            mdxy = dxy[good_mask]
            eps2 = np.sum(mdxy * mdxy, axis=1, dtype=np.float64)
            prev_centers = centers

        # TODO: make compute_residuals a method of Stars
        # compute residuals
        #if residuals:
        #    res = compute_residuals(psf, all_stars)
        #else:
        #    res = len(all_stars) * [None]
        #
        #for s, r in zip(all_stars, res):
        #    s.fit_residual = r

        # assign coordinate residuals of the iterative process
        for s, (dx, dy) in zip(all_stars, dxy):
            s.iter_fit_eps = (float(dx), float(dy))

        # TODO: return iter_num as Stars attribute
        return psf, stars

def init_psf(stars, shape=None, oversampling=4.0, pixel_scale=None,
             psf_cls=PSF2DModel, **kwargs):
    """
    Parameters
    ----------
    stars : Star, list of Star
        A list of :py:class:`~psfutils.catalogs.Star` objects
        containing image data of star cutouts that are used to "build" a PSF.

    psf : PSF2DModel, type, None, optional
        An existing approximation to the PSF which needs to be recomputed
        using new ``stars`` (with new parameters for center and flux)
        or a class indicating the type of the ``psf`` object to be created
        (preferebly a subclass of :py:class:`PSF2DModel`).
        If `psf` is `None`, a new ``psf`` will be computed using
        :py:class:`PSF2DModel`.

    shape : tuple, optional
        Numpy-style shape of the output PSF. If shape is not specified (i.e.,
        it is set to `None`), the shape will be derived from the sizes of the
        input star models. This is ignored if `psf` is not `None`.

    oversampling : float, tuple of float, list of (float, tuple of float), \
optional
        Oversampling factor of the PSF relative to star. It indicates how many
        times a pixel in the star's image should be sampled when creating PSF.
        If a single number is provided, that value will be used for both
        ``X`` and ``Y`` axes and for all stars. When a tuple is provided,
        first number will indicate oversampling along the ``X`` axis and the
        second number will indicate oversampling along the ``Y`` axis. It is
        also possible to have individualized oversampling factors for each star
        by providing a list of integers or tuples of integers.

    """
    # check parameters:
    if pixel_scale is None and oversampling is None:
        raise ValueError(
            "'oversampling' and 'pixel_scale' cannot be 'None' together. "
            "At least one of these two parameters must be provided."
        )

    # get all stars including linked stars as a flat list:
    all_stars = []
    for s in stars:
        all_stars += s.get_linked_list()

    # find pixel scale:
    if pixel_scale is None:
        ovx, ovy = _parse_tuple_pars(oversampling, name='oversampling',
                                     dtype=float)

        # compute PSF's pixel scale as the smallest scale of the stars
        # divided by the requested oversampling factor:
        pscale_x, pscale_y = all_stars[0].pixel_scale
        for s in all_stars[1:]:
            px, py = s.pixel_scale
            if px < pscale_x:
                pscale_x = px
            if py < pscale_y:
                pscale_y = py

        pscale_x /= ovx
        pscale_y /= ovy

    else:
        pscale_x, pscale_y = _parse_tuple_pars(pixel_scale, name='pixel_scale',
                                               dtype=float)

    # if shape is None, find the minimal shape that will include input star's
    # data:
    if shape is None:
        w = np.array([((s.x_center + 0.5) * s.pixel_scale[0] / pscale_x,
                       (s.nx - s.x_center - 0.5) * s.pixel_scale[0] / pscale_x)
                      for s in all_stars])
        h = np.array([((s.y_center + 0.5) * s.pixel_scale[1] / pscale_y,
                       (s.ny - s.y_center - 0.5) * s.pixel_scale[1] / pscale_y)
                      for s in all_stars])

        # size of the PSF in the input image pixels
        # (the image with min(pixel_scale)):
        nx = int(np.ceil(np.amax(w / ovx + 0.5)))
        ny = int(np.ceil(np.amax(h / ovy + 0.5)))

        # account for a maximum error of 1 pix in the initial star coordinates:
        nx += 2
        ny += 2

        # convert to oversampled pixels:
        nx = int(np.ceil(np.amax(nx * ovx + 0.5)))
        ny = int(np.ceil(np.amax(ny * ovy + 0.5)))

        # we prefer odd sized images
        nx += 1 - nx % 2
        ny += 1 - ny % 2
        shape = (ny, nx)

    else:
        (ny, nx) = _parse_tuple_pars(shape, name='shape', dtype=int)

    # center of the output grid:
    cx = (nx - 1) / 2.0
    cy = (ny - 1) / 2.0

    with warnings.catch_warnings(record=False):
        warnings.simplefilter("ignore", NonNormalizable)

        # filter out parameters set by us:
        kwargs = copy.deepcopy(kwargs)
        for kw in ['data', 'origin', 'normalize', 'pixel_scale']:
            if kw in kwargs:
                del kwargs[kw]

        data = np.zeros((ny, nx), dtype=np.float)
        psf = psf_cls(
            data=data, origin=(cx, cy), normalize=True,
            pixel_scale=(pscale_x, pscale_y), **kwargs
        )

    return psf
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.common.keys import Keys
from selenium.webdriver.support.ui import Select
from selenium.common.exceptions import NoSuchElementException
from selenium.common.exceptions import NoAlertPresentException
import unittest, time, re

class Search(unittest.TestCase):
    def setUp(self):
        self.driver = webdriver.Firefox()
        self.driver.implicitly_wait(30)
        self.base_url = "http://localhost"
        self.verificationErrors = []
        self.accept_next_alert = True

    def test_existed_film_local_searching(self):

        class NotExistedFilmException(Exception): pass

        search_target = u"Криминальное чтиво"

        driver = self.driver
        driver.get(self.base_url + "/php4dvd/")
        driver.implicitly_wait(3)

        #login
        driver.find_element_by_id("username").clear()
        driver.find_element_by_id("username").send_keys("admin")
        driver.find_element_by_name("password").clear()
        driver.find_element_by_name("password").send_keys("admin")
        driver.find_element_by_name("submit").click()

        searchbox = driver.find_element_by_class_name("inputDefault")
        searchbox.clear()
        searchbox.send_keys(search_target)
        searchbox.send_keys(Keys.ENTER)

        time.sleep(1)

        all_films = driver.find_elements_by_class_name("title")

        if len(all_films) == 0:
            raise NotExistedFilmException("Film was not found")

        for film in all_films:

            if film.get_attribute("textContent") == search_target:
                print "We have found film with title: ", film.get_attribute("textContent")
                break

        driver.find_element_by_link_text("Home").click()

    def test_not_existed_film_local_searching(self):

        class NotExistedFilmException(Exception): pass

        search_target = u"Красная шапочка"

        driver = self.driver
        driver.get(self.base_url + "/php4dvd/")
        driver.implicitly_wait(3)

        #login
        driver.find_element_by_id("username").clear()
        driver.find_element_by_id("username").send_keys("admin")
        driver.find_element_by_name("password").clear()
        driver.find_element_by_name("password").send_keys("admin")
        driver.find_element_by_name("submit").click()

        searchbox = driver.find_element_by_class_name("inputDefault")
        searchbox.clear()
        searchbox.send_keys(search_target)
        searchbox.send_keys(Keys.ENTER)

        time.sleep(1)

        all_films = driver.find_elements_by_class_name("title")

        if len(all_films) > 0:
            for film in all_films:

                if film.get_attribute("textContent") == search_target:
                    print "We have found film with title: ", film.get_attribute("textContent")
                    raise NotExistedFilmException("We have found not existed film")

        driver.find_element_by_link_text("Home").click()

    def is_element_present(self, how, what):
        try: self.driver.find_element(by=how, value=what)
        except NoSuchElementException as e: return False
        return True

    def is_alert_present(self):
        try: self.driver.switch_to_alert()
        except NoAlertPresentException as e: return False
        return True

    def close_alert_and_get_its_text(self):
        try:
            alert = self.driver.switch_to_alert()
            alert_text = alert.text
            if self.accept_next_alert:
                alert.accept()
            else:
                alert.dismiss()
            return alert_text
        finally: self.accept_next_alert = True

    def tearDown(self):
        self.driver.quit()
        self.assertEqual([], self.verificationErrors)

if __name__ == "__main__":
    unittest.main()
<commit_msg>Fixed film_searching test<commit_after># -*- coding: utf-8 -*-
from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.common.keys import Keys
from selenium.webdriver.support.ui import Select
from selenium.common.exceptions import NoSuchElementException
from selenium.common.exceptions import NoAlertPresentException
from selenium.webdriver.support.wait import WebDriverWait
from selenium.webdriver.support.expected_conditions import *
import unittest, time, re

class Search(unittest.TestCase):
    def setUp(self):
        self.driver = webdriver.Firefox()
        self.driver.implicitly_wait(30)
        self.base_url = "http://localhost"
        self.verificationErrors = []
        self.accept_next_alert = True

    def test_existed_film_local_searching(self):

        class NotExistedFilmException(Exception): pass

        search_target = u"Криминальное чтиво"
        wait = WebDriverWait(self.driver, 3)

        driver = self.driver
        driver.get(self.base_url + "/php4dvd/")
        driver.implicitly_wait(3)

        #login
        driver.find_element_by_id("username").clear()
        driver.find_element_by_id("username").send_keys("admin")
        driver.find_element_by_name("password").clear()
        driver.find_element_by_name("password").send_keys("admin")
        driver.find_element_by_name("submit").click()

        random_film = driver.find_element_by_class_name("title")

        searchbox = driver.find_element_by_class_name("inputDefault")
        searchbox.clear()
        searchbox.send_keys(search_target)
        searchbox.send_keys(Keys.ENTER)

        wait.until(invisibility_of_element_located((By.NAME, random_film)))

        time.sleep(1)

        films_after_search = driver.find_elements_by_class_name("title")

        if len(films_after_search) == 0:
            raise NotExistedFilmException("Film was not found")

        for film in films_after_search:

            if film.get_attribute("textContent") == search_target:
                print "We have found film with title: ", film.get_attribute("textContent")
                break

        driver.find_element_by_link_text("Home").click()

    def test_not_existed_film_local_searching(self):

        class NotExistedFilmException(Exception): pass

        search_target = u"Красная шапочка"
        wait = WebDriverWait(self.driver, 3)

        driver = self.driver
        driver.get(self.base_url + "/php4dvd/")
        driver.implicitly_wait(3)

        #login
        driver.find_element_by_id("username").clear()
        driver.find_element_by_id("username").send_keys("admin")
        driver.find_element_by_name("password").clear()
        driver.find_element_by_name("password").send_keys("admin")
        driver.find_element_by_name("submit").click()

        random_film = driver.find_element_by_class_name("title")

        searchbox = driver.find_element_by_class_name("inputDefault")
        searchbox.clear()
        searchbox.send_keys(search_target)
        searchbox.send_keys(Keys.ENTER)

        wait.until(invisibility_of_element_located((By.NAME, random_film)))

        time.sleep(1)

        films_after_search = driver.find_elements_by_class_name("title")

        if len(films_after_search) > 0:
            for film in films_after_search:

                if film.get_attribute("textContent") == search_target:
                    print "We have found film with title: ", film.get_attribute("textContent")
                    raise NotExistedFilmException("We have found not existed film")

        driver.find_element_by_link_text("Home").click()

    def is_element_present(self, how, what):
        try: self.driver.find_element(by=how, value=what)
        except NoSuchElementException as e: return False
        return True

    def is_alert_present(self):
        try: self.driver.switch_to_alert()
        except NoAlertPresentException as e: return False
        return True

    def close_alert_and_get_its_text(self):
        try:
            alert = self.driver.switch_to_alert()
            alert_text = alert.text
            if self.accept_next_alert:
                alert.accept()
            else:
                alert.dismiss()
            return alert_text
        finally: self.accept_next_alert = True

    def tearDown(self):
        self.driver.quit()
        self.assertEqual([], self.verificationErrors)

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


<commit_before># encoding: UTF-8

"""Library for IMPACT related lattice/model function."""

from __future__ import print_function

__copyright__ = "Copyright (c) 2015, Facility for Rare Isotope Beams"

__author__ = "Dylan Maxwell"

import os.path, shutil

import numpy as np

from ..common import DataError

def build_result(directory=None):
    """Convenience method to build IMPACT model result.
    """
    result_factory = ResultFactory(directory)

    return result_factory.build()

class ResultFactory(object):
    """A factory to get simulation results from IMPACT working directory.

    """

    def __init__(self, directory=None):
        """Initialize class, give IMPACT working directory.

        :param directory: IMPACT working directory (default is current directory)
        """
        self.directory = directory

    def build(self, runimpact=False, **kwargs):
        """ Build result from impact simulation.
        It needs IMPACT fort file, fort.18, fort.24, fort.25, and fort.26

        :param runimpact: flag to identify whether to execute a new run, False by default.

        :param IMPACT:  IMPACT-Z version, either "FRIB" or "LBL". "FRIB" by default.
        :param fort18:  impact energy file name
        :param fort24:  impact horizontal file name
        :param fort25:  impact vertical file name
        :param fort26:  impact longitudinal file name
        :param keep:    keep simulation results, True by default.
        :return:
        """

        if self.directory != None:
            wkdir = self.directory
        else:
            wkdir = os.getcwd()

        _IMPACT = kwargs.get("IMPACT", "FRIB")
        _fort18 = kwargs.get("fort18", "fort.18")
        _fort24 = kwargs.get("fort24", "fort.24")
        _fort25 = kwargs.get("fort25", "fort.25")
        _fort26 = kwargs.get("fort26", "fort.26")
        _keep = kwargs.get("keep", True)

        if runimpact:
            raise NotImplementedError("ResultFactory: Execution of IMPACT not implemented. See phylib.machine.frib.model.impact.ResultFactory.")

        fort18path = os.path.join(wkdir, _fort18)
        if not os.path.isfile(fort18path):
            raise RuntimeError("ResultFactory: IMPACT output not found: {}".format(fort18path))

        fort24path = os.path.join(wkdir, _fort24)
        if not os.path.isfile(fort24path):
            raise RuntimeError("ResultFactory: IMPACT output not found: {}".format(fort24path))

        fort25path = os.path.join(wkdir, _fort25)
        if not os.path.isfile(fort25path):
            raise RuntimeError("ResultFactory: IMPACT output not found: {}".format(fort25path))

        fort26path = os.path.join(wkdir, _fort26)
        if not os.path.isfile(fort25path):
            raise RuntimeError("ResultFactory: IMPACT output not found: {}".format(fort25path))

        # read data in if all data files are in place
        if _IMPACT == "FRIB":
            # z, phase (rad), energy (MeV), gamma, beta
            fort18 = np.loadtxt(fort18path, usecols=(0, 1, 3, 2, 4))
            # X0, X0', Xrms, X'rms, Ex, Alpha x, Beta x
            fort24 = np.loadtxt(fort24path, usecols=(1, 3, 2, 4, 7, 5, 6))
            # Y0, Y0', Yrms, Y'rms, Yp, Alpha y, Beta y
            fort25 = np.loadtxt(fort25path, usecols=(1, 3, 2, 4, 7, 5, 6))
            # Z0, Z0', Zrms, Z'rms, Zp, Alpha Z, Beta z
            fort26 = np.loadtxt(fort26path, usecols=(1, 3, 2, 4, 7, 5, 6))
        elif _IMPACT in ["LBL", "LBNL"]:
            # z, phase (rad), energy (MeV), gamma, beta
            fort18 = np.loadtxt(fort18path, usecols=(0, 1, 3, 2, 4))
            # X0, X0', Xrms, X'rms, Ex, Alpha x, Beta x
            fort24 = np.loadtxt(fort24path, usecols=(1, 3, 2, 4, 6, 5))
            # Y0, Y0', Yrms, Y'rms, Yp, Alpha y, Beta y
            fort25 = np.loadtxt(fort25path, usecols=(1, 3, 2, 4, 6, 5))
            # Z0, Z0', Zrms, Z'rms, Zp, Alpha Z, Beta z
            fort26 = np.loadtxt(fort26path, usecols=(1, 3, 2, 4, 6, 5))

        else:
            raise RuntimeError("Unknown IMPACT version. Cannot parse results.")

        if not _keep:
            shutil.rmtree(wkdir)

        return Result(fort18, fort24, fort25, fort26)

class Result(object):

    def __init__(self, fort18, fort24, fort25, fort26):
        """

        :param fort18: numpy N x 5 array
        :param fort24: numpy N x 7 array, or N x 6 for LBL version (no TWISS beta)
        :param fort25: numpy N x 7 array, or N x 6 for LBL version (no TWISS beta)
        :param fort26: numpy N x 7 array, or N x 6 for LBL version (no TWISS beta)
        :return:
        """
        # z, phase (rad), energy (MeV), gamma, beta
        self._fort18 = fort18
        # X0, X0', Xrms, X'rms, Ex, Alpha x, Beta x
        self._fort24 = fort24
        # Y0, Y0', Yrms, Y'rms, Yp, Alpha y, Beta y
        self._fort25 = fort25
        # Z0, Z0', Zrms, Z'rms, Zp, Alpha Z, Beta z
        self._fort26 = fort26

    def get_sposition(self):
        """

        :return: list of s position for each element
        """
        return self._fort18[:, 0]

    def get_orbit(self, plane="X"):
        """Get beam orbit

        :param plane:
        :return:
        """
        if plane.upper() in [ "X"]:
            data = np.empty((self._fort24.shape[0], 2))
            data[:, 0] = self._fort18[:, 0]
            data[:, 1] = self._fort24[:, 0]
        elif plane.upper() in [ "Y"]:
            data = np.empty((self._fort24.shape[0], 2))
            data[:, 0] = self._fort18[:, 0]
            data[:, 1] = self._fort25[:, 0]
        elif plane.upper() in [ "XY"]:
            data = np.empty((self._fort24.shape[0], 3))
            data[:, 0] = self._fort18[:, 0]
            data[:, 1] = self._fort24[:, 0]
            data[:, 2] = self._fort25[:, 0]
        else:
            raise RuntimeError("Result: Unknown plane to get beam orbit: {}".format(plane))

        return data

    def get_abs_phase(self):
        """Get accumulated beam phase in radian.

        :return:
        """
        return self._fort18[:, [0, 1]]

    def get_energy(self):
        """ Get beam energy in MeV.

        :return:
        """
        return self._fort18[:, [0, 2]]

    def get_twiss_alpha(self, plane="X"):
        """Get beam twiss alpha parameters.

        :return:
        """
        if plane.upper() in [ "X"]:
            data = np.empty((self._fort24.shape[0], 2))
            data[:, 0] = self._fort18[:, 0]
            data[:, 1] = self._fort24[:, 5]
        elif plane.upper() in [ "Y"]:
            data = np.empty((self._fort24.shape[0], 2))
            data[:, 0] = self._fort18[:, 0]
            data[:, 1] = self._fort25[:, 5]
        elif plane.upper() in [ "XY"]:
            data = np.empty((self._fort24.shape[0], 3))
            data[:, 0] = self._fort18[:, 0]
            data[:, 1] = self._fort24[:, 5]
            data[:, 2] = self._fort25[:, 5]
        elif plane.upper() in [ "Z"]:
            data = np.empty((self._fort24.shape[0], 2))
            data[:, 0] = self._fort18[:, 0]
            data[:, 1] = self._fort25[:, 5]
        else:
            raise RuntimeError("Result: Unknown plane to get Twiss Alpha: {}".format(plane))

        return data

    def get_twiss_beta(self, plane="X"):
        """Get beam twiss beta parameters.

        :return:
        """
        if plane.upper() in [ "X"]:
            if self._fort18.shape()[1] < 7:
                raise DataError("No Twiss Beta available")
            data = np.empty((self._fort24.shape[0], 2))
            data[:, 0] = self._fort18[:, 0]
            data[:, 1] = self._fort24[:, 6]
        elif plane.upper() in [ "Y"]:
            if self._fort18.shape()[1] < 7:
                raise DataError("No Twiss Beta available")
            data = np.empty((self._fort24.shape[0], 2))
            data[:, 0] = self._fort18[:, 0]
            data[:, 1] = self._fort25[:, 6]
        elif plane.upper() in [ "XY"]:
            if self._fort18.shape()[1] < 7:
                raise DataError("No Twiss Beta available")
            data = np.empty((self._fort24.shape[0], 3))
            data[:, 0] = self._fort18[:, 0]
            data[:, 1] = self._fort24[:, 6]
            data[:, 2] = self._fort25[:, 6]
        elif plane.upper() in [ "Z"]:
            if self._fort18.shape()[1] < 7:
                raise DataError("No Twiss Beta available")
            data = np.empty((self._fort24.shape[0], 2))
            data[:, 0] = self._fort18[:, 0]
            data[:, 1] = self._fort25[:, 6]
        else:
            raise RuntimeError("Result: Unknown plane to get Twiss beta: {}".format(plane))

        return data

    def get_beam_rms(self, plane="X"):
        """Get beam RMS size in X, Y, Z, or XY [meter].

        :return:
        """
        if plane.upper() in [ "X"]:
            data = np.empty((self._fort24.shape[0], 2))
            data[:, 0] = self._fort18[:, 0]
            data[:, 1] = self._fort24[:, 2]
        elif plane.upper() in [ "Y"]:
            data = np.empty((self._fort24.shape[0], 2))
            data[:, 0] = self._fort18[:, 0]
            data[:, 1] = self._fort25[:, 2]
        elif plane.upper() in [ "XY"]:
            data = np.empty((self._fort24.shape[0], 3))
            data[:, 0] = self._fort18[:, 0]
            data[:, 1] = self._fort24[:, 2]
            data[:, 2] = self._fort25[:, 2]
        elif plane.upper() in [ "Z"]:
            data = np.empty((self._fort24.shape[0], 2))
            data[:, 0] = self._fort18[:, 0]
            data[:, 1] = self._fort26[:, 2]
        else:
            raise RuntimeError("Result: Unknown plane to get beam RMS size: {}".format(plane))

        return data

    def get_beam_emittance(self, plane="X"):
        """Get beam normalized RMS emittance (m-rad for transverse and degree-MeV for longitudinal).

        :return:
        """
        if plane.upper() in [ "X"]:
            data = np.empty((self._fort24.shape[0], 2))
            data[:, 0] = self._fort18[:, 0]
            data[:, 1] = self._fort24[:, 4]
        elif plane.upper() in [ "Y"]:
            data = np.empty((self._fort24.shape[0], 2))
            data[:, 0] = self._fort18[:, 0]
            data[:, 1] = self._fort25[:, 4]
        elif plane.upper() in [ "XY"]:
            data = np.empty((self._fort24.shape[0], 3))
            data[:, 0] = self._fort18[:, 0]
            data[:, 1] = self._fort24[:, 4]
            data[:, 2] = self._fort25[:, 4]
        elif plane.upper() in [ "Z"]:
            data = np.empty((self._fort24.shape[0], 2))
            data[:, 0] = self._fort18[:, 0]
            data[:, 1] = self._fort26[:, 4]
        else:
            raise RuntimeError("Result: Unknown plane to get beam normalized RMS Emittance: {}".format(plane))

        return data
<commit_msg>API to access impact simulation result<commit_after># encoding: UTF-8

"""Library for IMPACT related lattice/model function."""

from __future__ import print_function

__copyright__ = "Copyright (c) 2015, Facility for Rare Isotope Beams"

__author__ = "Dylan Maxwell"

import os.path, shutil

import numpy as np

from ..common import DataError

def build_result(directory=None):
    """Convenience method to build IMPACT model result.
    """
    result_factory = ResultFactory(directory)

    return result_factory.build()

class ResultFactory(object):
    """A factory to get simulation results from IMPACT working directory.

    """

    def __init__(self, directory=None):
        """Initialize class, give IMPACT working directory.

        :param directory: IMPACT working directory (default is current directory)
        """
        self.directory = directory

    def build(self, runimpact=False, **kwargs):
        """ Build result from impact simulation.
        It needs IMPACT fort file, fort.18, fort.24, fort.25, and fort.26

        :param runimpact: flag to identify whether to execute a new run, False by default.

        :param IMPACT:  IMPACT-Z version, either "FRIB" or "LBL". "FRIB" by default.
        :param fort18:  impact energy file name
        :param fort24:  impact horizontal file name
        :param fort25:  impact vertical file name
        :param fort26:  impact longitudinal file name
        :param keep:    keep simulation results, True by default.
        :return:
        """

        if self.directory != None:
            wkdir = self.directory
        else:
            wkdir = os.getcwd()

        _IMPACT = kwargs.get("IMPACT", "FRIB")
        _fort18 = kwargs.get("fort18", "fort.18")
        _fort24 = kwargs.get("fort24", "fort.24")
        _fort25 = kwargs.get("fort25", "fort.25")
        _fort26 = kwargs.get("fort26", "fort.26")
        _keep = kwargs.get("keep", True)

        if runimpact:
            raise NotImplementedError("ResultFactory: Execution of IMPACT not implemented. See phylib.machine.frib.model.impact.ResultFactory.")

        fort18path = os.path.join(wkdir, _fort18)
        if not os.path.isfile(fort18path):
            raise RuntimeError("ResultFactory: IMPACT output not found: {}".format(fort18path))

        fort24path = os.path.join(wkdir, _fort24)
        if not os.path.isfile(fort24path):
            raise RuntimeError("ResultFactory: IMPACT output not found: {}".format(fort24path))

        fort25path = os.path.join(wkdir, _fort25)
        if not os.path.isfile(fort25path):
            raise RuntimeError("ResultFactory: IMPACT output not found: {}".format(fort25path))

        fort26path = os.path.join(wkdir, _fort26)
        if not os.path.isfile(fort25path):
            raise RuntimeError("ResultFactory: IMPACT output not found: {}".format(fort25path))

        # read data in if all data files are in place
        if _IMPACT == "FRIB":
            # z, phase (rad), energy (MeV), gamma, beta
            fort18 = np.loadtxt(fort18path, usecols=(0, 1, 3, 2, 4))
            # X0, X0', Xrms, X'rms, Ex, Alpha x, Beta x
            fort24 = np.loadtxt(fort24path, usecols=(1, 3, 2, 4, 7, 5, 6))
            # Y0, Y0', Yrms, Y'rms, Ey, Alpha y, Beta y
            fort25 = np.loadtxt(fort25path, usecols=(1, 3, 2, 4, 7, 5, 6))
            # Z0, Z0', Zrms, Z'rms, Ez, Alpha Z, Beta z
            fort26 = np.loadtxt(fort26path, usecols=(1, 3, 2, 4, 7, 5, 6))
        elif _IMPACT in ["LBL", "LBNL"]:
            # z, phase (rad), energy (MeV), gamma, beta
            fort18 = np.loadtxt(fort18path, usecols=(0, 1, 3, 2, 4))
            # X0, X0', Xrms, X'rms, Ex, Alpha x, Beta x
            fort24 = np.loadtxt(fort24path, usecols=(1, 3, 2, 4, 6, 5))
            # Y0, Y0', Yrms, Y'rms, Ey, Alpha y, Beta y
            fort25 = np.loadtxt(fort25path, usecols=(1, 3, 2, 4, 6, 5))
            # Z0, Z0', Zrms, Z'rms, Ez, Alpha Z, Beta z
            fort26 = np.loadtxt(fort26path, usecols=(1, 3, 2, 4, 6, 5))

        else:
            raise RuntimeError("Unknown IMPACT version. Cannot parse results.")

        if not _keep:
            shutil.rmtree(wkdir)

        return Result(fort18, fort24, fort25, fort26)

class Result(object):

    def __init__(self, fort18, fort24, fort25, fort26, impact="LBL"):
        """

        :param fort18: numpy N x 5 array
        :param fort24: numpy N x 7 array, or N x 6 for LBL version (no TWISS beta)
        :param fort25: numpy N x 7 array, or N x 6 for LBL version (no TWISS beta)
        :param fort26: numpy N x 7 array, or N x 6 for LBL version (no TWISS beta)
        :return:
        """
        # z, phase (rad), energy (MeV), gamma, beta
        self._fort18 = fort18
        # X0, X0', Xrms, X'rms, Ex, Alpha x, Beta x
        self._fort24 = fort24
        # Y0, Y0', Yrms, Y'rms, Yp, Alpha y, Beta y
        self._fort25 = fort25
        # Z0, Z0', Zrms, Z'rms, Zp, Alpha Z, Beta z
        self._fort26 = fort26
        self.elements = len(self._fort18[:, 0])

        self.impact = impact

    def __checkElements(self, elemIdx):
        """Check whether required element index is over the range.

        :param elemIdx: element index

        :raise: RuntimeError
        """
        if elemIdx > self.elements:
            raise RuntimeError("Required element index {0} over maximum {1}.".format(elemIdx, self.elements))

    def getSPosition(self, elemIdx=None):
        """Get element s position at the end if elemIdx is given, or list of s position for all elements

        :param elemIdx: index number of given element

        :return: s position or list
        :raise: RuntimeError
        """

        if elemIdx is None:
            return self._fort18[:, 0]
        else:
            self.__checkElements(elemIdx)
        return self._fort18[:, 0][elemIdx]

    def getAbsPhase(self, elemIdx=None):
        """Get accumulated beam phase in radian at the end if elemIdx is given,
        or a list for all elements

        :param elemIdx: index number of given element

        :return: accumulated beam phase or list
        :raise: RuntimeError
        """

        if elemIdx is None:
            return self._fort18[:, 1]
        else:
            self.__checkElements(elemIdx)
        return self._fort18[:, 1][elemIdx]

    def getEnergy(self, elemIdx=None):
        """Get beam energy in MeV/u at the end if elemIdx is given,
        or a list for all elements

        :param elemIdx: index number of given element

        :return: beam energy or list
        :raise: RuntimeError
        """

        if elemIdx is None:
            return self._fort18[:, 2]
        else:
            self.__checkElements(elemIdx)
        return self._fort18[:, 2][elemIdx]

    def getBetaGamma(self, elemIdx=None):
        """Get beam beta*gamma at the end if elemIdx is given, or a list for all elements

        :param elemIdx: index number of given element

        :return: beta*gamma or list
        :raise: RuntimeError
        """

        if elemIdx is None:
            return self._fort18[:, 3]*self._fort18[:, 4]
        else:
            self.__checkElements(elemIdx)
        return self._fort18[:, 3][elemIdx]*self._fort18[:, 4][elemIdx]

    def __getData(self, data, data2=None, elemIdx=None, col=0):
        """Common interface to get simulation data.

        :param elemIdx: element index
        :param data:    first data set
        :param data2:   2nd data set
        :param col:     column in data/data2

        :return: value at given location if elemIdx is `None`, or 1D array if data2 is `None`, otherwise 2D array

        :raise: RuntimeError
        """
        result = None
        if elemIdx is None:
            if data2 is None:
                result = data[:, col]
            else:
                result = np.empty((data.shape[0], 2))
                result[:, 0] = data[:, col]
                result[:, 1] = data2[:, col]
        else:
            self.__checkElements(elemIdx)
            if data2 is None:
                result = data[:, col][elemIdx]
            else:
                result = [data[:, col][elemIdx], data2[:, col][elemIdx]]
        return result

    def getOrbit(self, plane="X", elemIdx=None):
        """Get beam position at the end of an element if elemIdx is given, or beam orbit at all elements.
        Current implementation returns all position information from simulation, and does not separate BPM
        from other devices like magnet and other diagnostic devices like profile monitor.

        :param plane:    beam plane, either "X", "Y", or "XY"
        :param elemIdx:  element index, `None` by default

        :return: beam position at given location, or at all elements
        """
        if plane.upper() == "X":
            return self.__getData(self._fort24, elemIdx=elemIdx, col=0)
        elif plane.upper() == "Y":
            return self.__getData(self._fort25, elemIdx=elemIdx, col=0)
        elif plane.upper() == "XY":
            return self.__getData(self._fort24, data2=self._fort25, elemIdx=elemIdx, col=0)
        else:
            raise RuntimeError("Result: Unknown plane for beam orbit on: {}".format(plane))

    def getTwissAlpha(self, plane="X", elemIdx=None):
        """Get beam twiss alpha parameters at the end of an element if elemIdx is given, or at all elements.
        Current implementation returns all position information from simulation, and does not separate BPM
        from other devices like magnet and other diagnostic devices like profile monitor.

        :param plane:    beam plane, either "X", "Y", "Z", or "XY"
        :param elemIdx:  element index, `None` by default

        :return: beam twiss alpha at given location, or at all elements
        """
        if plane.upper() == "X":
            return self.__getData(self._fort24, elemIdx=elemIdx, col=5)
        elif plane.upper() == "Y":
            return self.__getData(self._fort25, elemIdx=elemIdx, col=5)
        elif plane.upper() == "Z":
            return self.__getData(self._fort26, elemIdx=elemIdx, col=5)
        elif plane.upper() == "XY":
            return self.__getData(self._fort24, data2=self._fort25, elemIdx=elemIdx, col=5)
        else:
            raise RuntimeError("Result: Unknown plane for beam twiss alpha on: {}".format(plane))

    def getTwissBeta(self, plane="X", elemIdx=None):
        """Get beam twiss beta parameters at the end of an element if elemIdx is given, or at all elements.
        Current implementation returns all position information from simulation, and does not separate BPM
        from other devices like magnet and other diagnostic devices like profile monitor.

        :param plane:    beam plane, either "X", "Y", "Z", or "XY"
        :param elemIdx:  element index, `None` by default

        :return: beam twiss beta at given location, or at all elements
        """
        if plane.upper() == "X":
            if self._fort24.shape()[1] < 7:
                raise DataError("No Twiss Beta available")
            return self.__getData(self._fort24, elemIdx=elemIdx, col=6)
        elif plane.upper() == "Y":
            if self._fort25.shape()[1] < 7:
                raise DataError("No Twiss Beta available")
            return self.__getData(self._fort25, elemIdx=elemIdx, col=6)
        elif plane.upper() == "Z":
            if self._fort26.shape()[1] < 7:
                raise DataError("No Twiss Beta available")
            return self.__getData(self._fort26, elemIdx=elemIdx, col=6)
        elif plane.upper() == "XY":
            if self._fort24.shape()[1] < 7 or self._fort25.shape()[1] < 7:
                raise DataError("No Twiss Beta available")
            return self.__getData(self._fort24, data2=self._fort25, elemIdx=elemIdx, col=6)
        else:
            raise RuntimeError("Result: Unknown plane for beam twiss beta on: {}".format(plane))

    def geBeamRms(self, plane="X", elemIdx=None):
        """Get beam RMS parameters at the end of an element if elemIdx is given, or at all elements.
        Current implementation returns all position information from simulation, and does not separate BPM
        from other devices like magnet and other diagnostic devices like profile monitor.

        :param plane:    beam plane, either "X", "Y", "Z", or "XY"
        :param elemIdx:  element index, `None` by default

        :return: beam twiss beta at given location, or at all elements
        """
        if plane.upper() == "X":
            return self.__getData(self._fort24, elemIdx=elemIdx, col=2)
        elif plane.upper() == "Y":
            return self.__getData(self._fort25, elemIdx=elemIdx, col=2)
        elif plane.upper() == "Z":
            return self.__getData(self._fort26, elemIdx=elemIdx, col=2)
        elif plane.upper() == "XY":
            return self.__getData(self._fort24, data2=self._fort25, elemIdx=elemIdx, col=2)
        else:
            raise RuntimeError("Result: Unknown plane for beam RMS on: {}".format(plane))

    def getEmittance(self, plane="X", elemIdx=None):
        """Get beam normalized emittance (m-rad for transverse and degree-MeV for longitudinal)
        at the end of an element if elemIdx is given, or at all elements.
        Current implementation returns all position information from simulation, and does not separate BPM
        from other devices like magnet and other diagnostic devices like profile monitor.

        :param plane:    beam plane, either "X", "Y", "Z", or "XY"
        :param elemIdx:  element index, `None` by default

        :return: beam emittance at given location, or at all elements
        """
        if plane.upper() == "X":
            return self.__getData(self._fort24, elemIdx=elemIdx, col=4)
        elif plane.upper() == "Y":
            return self.__getData(self._fort25, elemIdx=elemIdx, col=4)
        elif plane.upper() == "Z":
            return self.__getData(self._fort26, elemIdx=elemIdx, col=4)
        elif plane.upper() == "XY":
            return self.__getData(self._fort24, data2=self._fort25, elemIdx=elemIdx, col=4)
        else:
            raise RuntimeError("Result: Unknown plane for beam RMS on: {}".format(plane))

    def getBeamMomentumCentroid(self, plane="X", elemIdx=None):
        """Get beam centroid momentum (radian for transverse and MeV for longitudinal)
        at the end of an element if elemIdx is given, or at all elements.
        Current implementation returns all position information from simulation, and does not separate BPM
        from other devices like magnet and other diagnostic devices like profile monitor.

        :param plane:    beam plane, either "X", "Y", "Z", or "XY"
        :param elemIdx:  element index, `None` by default

        :return: beam emittance at given location, or at all elements
        """
        if plane.upper() == "X":
            return self.__getData(self._fort24, elemIdx=elemIdx, col=1)
        elif plane.upper() == "Y":
            return self.__getData(self._fort25, elemIdx=elemIdx, col=1)
        elif plane.upper() == "Z":
            return self.__getData(self._fort26, elemIdx=elemIdx, col=1)
        elif plane.upper() == "XY":
            return self.__getData(self._fort24, data2=self._fort25, elemIdx=elemIdx, col=1)
        else:
            raise RuntimeError("Result: Unknown plane for beam RMS on: {}".format(plane))

    def getMomentumRms(self, plane="X", elemIdx=None):
        """Get beam RMS momentum (radian for transverse and MeV for longitudinal)
        at the end of an element if elemIdx is given, or at all elements.
        Current implementation returns all position information from simulation, and does not separate BPM
        from other devices like magnet and other diagnostic devices like profile monitor.

        :param plane:    beam plane, either "X", "Y", "Z", or "XY"
        :param elemIdx:  element index, `None` by default

        :return: beam emittance at given location, or at all elements
        """
        if plane.upper() == "X":
            return self.__getData(self._fort24, elemIdx=elemIdx, col=3)
        elif plane.upper() == "Y":
            return self.__getData(self._fort25, elemIdx=elemIdx, col=3)
        elif plane.upper() == "Z":
            return self.__getData(self._fort26, elemIdx=elemIdx, col=3)
        elif plane.upper() == "XY":
            return self.__getData(self._fort24, data2=self._fort25, elemIdx=elemIdx, col=3)
        else:
            raise RuntimeError("Result: Unknown plane for beam RMS on: {}".format(plane))
<|endoftext|>


<commit_before>import base64
import hadoopy_hbase
import msgpack
import numpy as np
import json
import scipy.cluster.vq
import scipy as sp
import random
import picarus_takeout
import tables
import os
from driver import PicarusManager

def _setup(start_stop_rows, inputs):
    thrift = hadoopy_hbase.connect()  # TODO: Need to pass in thrift server/port
    manager = PicarusManager(thrift=thrift)
    slices = [base64.b64encode(start_row) + ',' + base64.b64encode(stop_row)
              for start_row, stop_row in start_stop_rows]
    os.nice(5)  # These are background tasks, don't let the CPU get too crazy
    return thrift, manager, slices, map(base64.b64encode, inputs)

def classifier_sklearn(queue, params, inputs, schema, start_stop_rows, table, owner):
    thrift, manager, slices, inputsb64 = _setup(start_stop_rows, inputs)
    label_features = {0: [], 1: []}
    for start_row, stop_row in start_stop_rows:
        row_cols = hadoopy_hbase.scanner(thrift, table,
                                         columns=[inputs['feature'], inputs['meta']],
                                         start_row=start_row, stop_row=stop_row)
        for row, cols in row_cols:
            try:
                label = int(cols[inputs['meta']] == params['class_positive'])
                label_features[label].append(cols[inputs['feature']])
            except KeyError:
                continue
    labels = [0] * len(label_features[0]) + [1] * len(label_features[1])
    features = label_features[0] + label_features[1]
    features = np.asfarray([msgpack.loads(x)[0] for x in features])
    import sklearn.svm
    classifier = sklearn.svm.LinearSVC()
    classifier.fit(features, np.asarray(labels))
    factory_info = {'slices': slices, 'num_rows': len(features), 'data': 'slices', 'params': params, 'inputs': inputsb64}
    model_link = {'name': 'picarus.LinearClassifier', 'kw': {'coefficients': classifier.coef_.tolist()[0],
                                                             'intercept': classifier.intercept_[0]}}
    model_chain = tables._takeout_model_chain_from_key(manager, inputs['feature']) + [model_link]
    queue.put(manager.input_model_param_to_key(**{'input': inputs['feature'], 'model_link': model_link, 'model_chain': model_chain, 'input_type': 'feature',
                                                  'output_type': 'binary_class_confidence', 'email': owner, 'name': manager.model_to_name(model_link),
                                                  'factory_info': json.dumps(factory_info)}))

def classifier_localnbnn(queue, params, inputs, schema, start_stop_rows, table, owner):
    thrift, manager, slices, inputsb64 = _setup(start_stop_rows, inputs)
    features = []
    indeces = []
    num_features = 0
    feature_size = 0
    labels_dict = {}
    labels = []
    for start_row, stop_row in start_stop_rows:
        row_cols = hadoopy_hbase.scanner(thrift, table,
                                         columns=[inputs['multi_feature'], inputs['meta']], start_row=start_row, stop_row=stop_row)
        for _, cols in row_cols:
            try:
                label = cols[inputs['meta']]
                f, s = msgpack.loads(cols[inputs['multi_feature']])
                if label not in labels_dict:
                    labels_dict[label] = len(labels_dict)
                    labels.append(label)
                feature_size = s[1]
                num_features += s[0]
                features += f
                indeces += [labels_dict[label]] * s[0]
            except KeyError:
                pass
    factory_info = {'slices': slices, 'data': 'slices', 'params': params, 'inputs': inputsb64}
    model_link = {'name': 'picarus.LocalNBNNClassifier', 'kw': {'features': features, 'indeces': indeces, 'labels': labels,
                                                                'feature_size': feature_size, 'max_results': params['max_results']}}
    model_chain = tables._takeout_model_chain_from_key(manager, inputs['multi_feature']) + [model_link]
    queue.put(manager.input_model_param_to_key(**{'input': inputs['multi_feature'], 'model_link': model_link, 'model_chain': model_chain,
                                                  'input_type': 'multi_feature', 'output_type': 'multi_class_distance',
                                                  'email': owner, 'name': manager.model_to_name(model_link), 'factory_info': json.dumps(factory_info)}))

def feature_bovw_mask(queue, params, inputs, schema, start_stop_rows, table, owner):
    thrift, manager, slices, inputsb64 = _setup(start_stop_rows, inputs)
    features = []
    for start_row, stop_row in start_stop_rows:
        row_cols = hadoopy_hbase.scanner(thrift, table,
                                         columns=[inputs['mask_feature']],
                                         start_row=start_row, stop_row=stop_row)
        for row, cols in row_cols:
            cur_feature = msgpack.loads(cols[inputs['mask_feature']])
            cur_feature = np.array(cur_feature[0]).reshape((-1, cur_feature[1][2]))
            features += random.sample(cur_feature, min(len(cur_feature), params['max_per_row']))
            print(len(features))
    features = np.asfarray(features)
    clusters = sp.cluster.vq.kmeans(features, params['num_clusters'])[0]
    num_clusters = clusters.shape[0]
    factory_info = {'slices': slices, 'num_features': len(features), 'data': 'slices', 'params': params, 'inputs': inputsb64}
    model_link = {'name': 'picarus.BOVWImageFeature', 'kw': {'clusters': clusters.ravel().tolist(), 'num_clusters': num_clusters,
                                                             'levels': params['levels']}}
    model_chain = tables._takeout_model_chain_from_key(manager, inputs['mask_feature']) + [model_link]
    queue.put(manager.input_model_param_to_key(**{'input': inputs['mask_feature'], 'model_link': model_link, 'model_chain': model_chain, 'input_type': 'feature',
                                                  'output_type': 'feature', 'email': owner, 'name': manager.model_to_name(model_link),
                                                  'factory_info': json.dumps(factory_info)}))

def hasher_spherical(queue, params, inputs, schema, start_stop_rows, table, owner):
    thrift, manager, slices, inputsb64 = _setup(start_stop_rows, inputs)
    features = []
    for start_row, stop_row in start_stop_rows:
        row_cols = hadoopy_hbase.scanner(thrift, table,
                                         columns=[inputs['feature']],
                                         start_row=start_row, stop_row=stop_row)
        for row, cols in row_cols:
            cur_feature = msgpack.loads(cols[inputs['feature']])
            features.append(np.array(cur_feature[0]))
    print('num_features[%d]' % len(features))
    features = np.asfarray(features)
    out = picarus_takeout.spherical_hasher_train(features, params['num_pivots'], params['eps_m'], params['eps_s'], params['max_iters'])
    out = {'pivots': out['pivots'].ravel().tolist(),
           'threshs': out['threshs'].tolist()}
    #out = picarus.modules.spherical_hash.train_takeout(features, params['num_pivots'], params['eps_m'], params['eps_s'], params['max_iters'])
    factory_info = {'slices': slices, 'num_features': len(features), 'data': 'slices', 'params': params, 'inputs': inputsb64}
    model_link = {'name': 'picarus.SphericalHasher', 'kw': out}
    model_chain = tables._takeout_model_chain_from_key(manager, inputs['feature']) + [model_link]
    queue.put(manager.input_model_param_to_key(**{'input': inputs['feature'], 'model_link': model_link, 'model_chain': model_chain, 'input_type': 'feature',
                                                  'output_type': 'hash', 'email': owner, 'name': manager.model_to_name(model_link),
                                                  'factory_info': json.dumps(factory_info)}))

def index_spherical(queue, params, inputs, schema, start_stop_rows, table, owner):
    thrift, manager, slices, inputsb64 = _setup(start_stop_rows, inputs)
    hashes = []
    labels = []
    for start_row, stop_row in start_stop_rows:
        row_cols = hadoopy_hbase.scanner(thrift, table,
                                         columns=[inputs['hash']],
                                         start_row=start_row, stop_row=stop_row)
        for row, cols in row_cols:
            hashes.append(cols[inputs['hash']])
            labels.append(row)
    hashes = ''.join(hashes)
    factory_info = {'slices': slices, 'num_hashes': len(labels), 'data': 'slices', 'params': params, 'inputs': inputsb64}
    model_link = {'name': 'picarus.SphericalHashIndex', 'kw': {'hashes': hashes,
                                                               'indeces': range(len(labels)), 'labels': labels,
                                                               'max_results': params['max_results']}}
    model_chain = tables._takeout_model_chain_from_key(manager, inputs['hash']) + [model_link]
    queue.put(manager.input_model_param_to_key(**{'input': inputs['hash'], 'model_link': model_link, 'model_chain': model_chain, 'input_type': 'hash',
                                                  'output_type': 'distance_image_rows', 'email': owner, 'name': manager.model_to_name(model_link),
                                                  'factory_info': json.dumps(factory_info)}))

def index_hamming_feature2d(queue, params, inputs, schema, start_stop_rows, table, owner):
    thrift, manager, slices, inputsb64 = _setup(start_stop_rows, inputs)
    hashes = []
    #keypoints = []
    labels = []
    indeces = []
    for start_row, stop_row in start_stop_rows:
        row_cols = hadoopy_hbase.scanner(thrift, table,
                                         columns=[inputs['feature2d_binary']],
                                         start_row=start_row, stop_row=stop_row)
        for row, cols in row_cols:
            f = msgpack.loads(cols[inputs['feature2d_binary']])
            print(f[2][0])
            hashes.append(f[0])
            #keypoints += f[1]
            indeces += [len(labels)] * f[2][0]
            labels.append(row)
            print(len(labels))
    hashes = ''.join(hashes)
    factory_info = {'slices': slices, 'num_hashes': len(indeces), 'num_images': len(labels), 'data': 'slices', 'params': params, 'inputs': inputsb64}
    #'keypoints': keypoints,
    model_link = {'name': 'picarus.HammingFeature2dHashIndex', 'kw': {'hashes': hashes,
                                                                      'indeces': indeces, 'labels': labels,
                                                                      'max_results': params['max_results'],
                                                                      'max_keypoint_results': params['max_keypoint_results'],
                                                                      'hamming_thresh': params['hamming_thresh']}}
    model_chain = tables._takeout_model_chain_from_key(manager, inputs['feature2d_binary']) + [model_link]
    queue.put(manager.input_model_param_to_key(**{'input': inputs['feature2d_binary'], 'model_link': model_link, 'model_chain': model_chain, 'input_type': 'feature2d_binary',
                                                  'output_type': 'distance_image_rows', 'email': owner, 'name': manager.model_to_name(model_link),
                                                  'factory_info': json.dumps(factory_info)}))

FACTORIES = {'factory/classifier/svmlinear': classifier_sklearn,
             'factory/classifier/localnbnn': classifier_localnbnn,
             'factory/feature/bovw': feature_bovw_mask,
             'factory/hasher/spherical': hasher_spherical,
             'factory/index/spherical': index_spherical,
             'factory/index/hamming_feature_2d': index_hamming_feature2d}
<commit_msg>Updated<commit_after>import base64
import hadoopy_hbase
import msgpack
import numpy as np
import json
import scipy.cluster.vq
import scipy as sp
import random
import picarus_takeout
import tables
import os
from driver import PicarusManager

def _setup(start_stop_rows, inputs):
    thrift = hadoopy_hbase.connect()  # TODO: Need to pass in thrift server/port
    manager = PicarusManager(thrift=thrift)
    slices = [base64.b64encode(start_row) + ',' + base64.b64encode(stop_row)
              for start_row, stop_row in start_stop_rows]
    os.nice(5)  # These are background tasks, don't let the CPU get too crazy
    return thrift, manager, slices, {k: base64.b64encode(v) for k, v in inputs.items()}

def classifier_sklearn(queue, params, inputs, schema, start_stop_rows, table, owner):
    thrift, manager, slices, inputsb64 = _setup(start_stop_rows, inputs)
    label_features = {0: [], 1: []}
    for start_row, stop_row in start_stop_rows:
        row_cols = hadoopy_hbase.scanner(thrift, table,
                                         columns=[inputs['feature'], inputs['meta']],
                                         start_row=start_row, stop_row=stop_row)
        for row, cols in row_cols:
            try:
                label = int(cols[inputs['meta']] == params['class_positive'])
                label_features[label].append(cols[inputs['feature']])
            except KeyError:
                continue
    labels = [0] * len(label_features[0]) + [1] * len(label_features[1])
    features = label_features[0] + label_features[1]
    features = np.asfarray([msgpack.loads(x)[0] for x in features])
    import sklearn.svm
    classifier = sklearn.svm.LinearSVC()
    classifier.fit(features, np.asarray(labels))
    factory_info = {'slices': slices, 'num_rows': len(features), 'data': 'slices', 'params': params, 'inputs': inputsb64}
    model_link = {'name': 'picarus.LinearClassifier', 'kw': {'coefficients': classifier.coef_.tolist()[0],
                                                             'intercept': classifier.intercept_[0]}}
    model_chain = tables._takeout_model_chain_from_key(manager, inputs['feature']) + [model_link]
    queue.put(manager.input_model_param_to_key(**{'input': inputs['feature'], 'model_link': model_link, 'model_chain': model_chain, 'input_type': 'feature',
                                                  'output_type': 'binary_class_confidence', 'email': owner, 'name': manager.model_to_name(model_link),
                                                  'factory_info': json.dumps(factory_info)}))

def classifier_localnbnn(queue, params, inputs, schema, start_stop_rows, table, owner):
    thrift, manager, slices, inputsb64 = _setup(start_stop_rows, inputs)
    features = []
    indeces = []
    num_features = 0
    feature_size = 0
    labels_dict = {}
    labels = []
    for start_row, stop_row in start_stop_rows:
        row_cols = hadoopy_hbase.scanner(thrift, table,
                                         columns=[inputs['multi_feature'], inputs['meta']], start_row=start_row, stop_row=stop_row)
        for _, cols in row_cols:
            try:
                label = cols[inputs['meta']]
                f, s = msgpack.loads(cols[inputs['multi_feature']])
                if label not in labels_dict:
                    labels_dict[label] = len(labels_dict)
                    labels.append(label)
                feature_size = s[1]
                num_features += s[0]
                features += f
                indeces += [labels_dict[label]] * s[0]
            except KeyError:
                pass
    factory_info = {'slices': slices, 'data': 'slices', 'params': params, 'inputs': inputsb64}
    model_link = {'name': 'picarus.LocalNBNNClassifier', 'kw': {'features': features, 'indeces': indeces, 'labels': labels,
                                                                'feature_size': feature_size, 'max_results': params['max_results']}}
    model_chain = tables._takeout_model_chain_from_key(manager, inputs['multi_feature']) + [model_link]
    queue.put(manager.input_model_param_to_key(**{'input': inputs['multi_feature'], 'model_link': model_link, 'model_chain': model_chain,
                                                  'input_type': 'multi_feature', 'output_type': 'multi_class_distance',
                                                  'email': owner, 'name': manager.model_to_name(model_link), 'factory_info': json.dumps(factory_info)}))

def feature_bovw_mask(queue, params, inputs, schema, start_stop_rows, table, owner):
    thrift, manager, slices, inputsb64 = _setup(start_stop_rows, inputs)
    features = []
    for start_row, stop_row in start_stop_rows:
        row_cols = hadoopy_hbase.scanner(thrift, table,
                                         columns=[inputs['mask_feature']],
                                         start_row=start_row, stop_row=stop_row)
        for row, cols in row_cols:
            cur_feature = msgpack.loads(cols[inputs['mask_feature']])
            cur_feature = np.array(cur_feature[0]).reshape((-1, cur_feature[1][2]))
            features += random.sample(cur_feature, min(len(cur_feature), params['max_per_row']))
            print(len(features))
    features = np.asfarray(features)
    clusters = sp.cluster.vq.kmeans(features, params['num_clusters'])[0]
    num_clusters = clusters.shape[0]
    factory_info = {'slices': slices, 'num_features': len(features), 'data': 'slices', 'params': params, 'inputs': inputsb64}
    model_link = {'name': 'picarus.BOVWImageFeature', 'kw': {'clusters': clusters.ravel().tolist(), 'num_clusters': num_clusters,
                                                             'levels': params['levels']}}
    model_chain = tables._takeout_model_chain_from_key(manager, inputs['mask_feature']) + [model_link]
    queue.put(manager.input_model_param_to_key(**{'input': inputs['mask_feature'], 'model_link': model_link, 'model_chain': model_chain, 'input_type': 'feature',
                                                  'output_type': 'feature', 'email': owner, 'name': manager.model_to_name(model_link),
                                                  'factory_info': json.dumps(factory_info)}))

def hasher_spherical(queue, params, inputs, schema, start_stop_rows, table, owner):
    thrift, manager, slices, inputsb64 = _setup(start_stop_rows, inputs)
    features = []
    for start_row, stop_row in start_stop_rows:
        row_cols = hadoopy_hbase.scanner(thrift, table,
                                         columns=[inputs['feature']],
                                         start_row=start_row, stop_row=stop_row)
        for row, cols in row_cols:
            cur_feature = msgpack.loads(cols[inputs['feature']])
            features.append(np.array(cur_feature[0]))
    print('num_features[%d]' % len(features))
    features = np.asfarray(features)
    out = picarus_takeout.spherical_hasher_train(features, params['num_pivots'], params['eps_m'], params['eps_s'], params['max_iters'])
    out = {'pivots': out['pivots'].ravel().tolist(),
           'threshs': out['threshs'].tolist()}
    #out = picarus.modules.spherical_hash.train_takeout(features, params['num_pivots'], params['eps_m'], params['eps_s'], params['max_iters'])
    factory_info = {'slices': slices, 'num_features': len(features), 'data': 'slices', 'params': params, 'inputs': inputsb64}
    model_link = {'name': 'picarus.SphericalHasher', 'kw': out}
    model_chain = tables._takeout_model_chain_from_key(manager, inputs['feature']) + [model_link]
    queue.put(manager.input_model_param_to_key(**{'input': inputs['feature'], 'model_link': model_link, 'model_chain': model_chain, 'input_type': 'feature',
                                                  'output_type': 'hash', 'email': owner, 'name': manager.model_to_name(model_link),
                                                  'factory_info': json.dumps(factory_info)}))

def index_spherical(queue, params, inputs, schema, start_stop_rows, table, owner):
    thrift, manager, slices, inputsb64 = _setup(start_stop_rows, inputs)
    hashes = []
    labels = []
    for start_row, stop_row in start_stop_rows:
        row_cols = hadoopy_hbase.scanner(thrift, table,
                                         columns=[inputs['hash']],
                                         start_row=start_row, stop_row=stop_row)
        for row, cols in row_cols:
            hashes.append(cols[inputs['hash']])
            labels.append(row)
    hashes = ''.join(hashes)
    factory_info = {'slices': slices, 'num_hashes': len(labels), 'data': 'slices', 'params': params, 'inputs': inputsb64}
    model_link = {'name': 'picarus.SphericalHashIndex', 'kw': {'hashes': hashes,
                                                               'indeces': range(len(labels)), 'labels': labels,
                                                               'max_results': params['max_results']}}
    model_chain = tables._takeout_model_chain_from_key(manager, inputs['hash']) + [model_link]
    queue.put(manager.input_model_param_to_key(**{'input': inputs['hash'], 'model_link': model_link, 'model_chain': model_chain, 'input_type': 'hash',
                                                  'output_type': 'distance_image_rows', 'email': owner, 'name': manager.model_to_name(model_link),
                                                  'factory_info': json.dumps(factory_info)}))

def index_hamming_feature2d(queue, params, inputs, schema, start_stop_rows, table, owner):
    thrift, manager, slices, inputsb64 = _setup(start_stop_rows, inputs)
    hashes = []
    #keypoints = []
    labels = []
    indeces = []
    for start_row, stop_row in start_stop_rows:
        row_cols = hadoopy_hbase.scanner(thrift, table,
                                         columns=[inputs['feature2d_binary']],
                                         start_row=start_row, stop_row=stop_row)
        for row, cols in row_cols:
            f = msgpack.loads(cols[inputs['feature2d_binary']])
            print(f[2][0])
            hashes.append(f[0])
            #keypoints += f[1]
            indeces += [len(labels)] * f[2][0]
            labels.append(row)
            print(len(labels))
    hashes = ''.join(hashes)
    factory_info = {'slices': slices, 'num_hashes': len(indeces), 'num_images': len(labels), 'data': 'slices', 'params': params, 'inputs': inputsb64}
    #'keypoints': keypoints,
    model_link = {'name': 'picarus.HammingFeature2dHashIndex', 'kw': {'hashes': hashes,
                                                                      'indeces': indeces, 'labels': labels,
                                                                      'max_results': params['max_results'],
                                                                      'max_keypoint_results': params['max_keypoint_results'],
                                                                      'hamming_thresh': params['hamming_thresh']}}
    model_chain = tables._takeout_model_chain_from_key(manager, inputs['feature2d_binary']) + [model_link]
    queue.put(manager.input_model_param_to_key(**{'input': inputs['feature2d_binary'], 'model_link': model_link, 'model_chain': model_chain, 'input_type': 'feature2d_binary',
                                                  'output_type': 'distance_image_rows', 'email': owner, 'name': manager.model_to_name(model_link),
                                                  'factory_info': json.dumps(factory_info)}))

FACTORIES = {'factory/classifier/svmlinear': classifier_sklearn,
             'factory/classifier/localnbnn': classifier_localnbnn,
             'factory/feature/bovw': feature_bovw_mask,
             'factory/hasher/spherical': hasher_spherical,
             'factory/index/spherical': index_spherical,
             'factory/index/hamming_feature_2d': index_hamming_feature2d}
<|endoftext|>