<commit_before># Copyright (C) 2011 Linaro Limited
#
# Author: Zygmunt Krynicki <zygmunt.krynicki@linaro.org>
#
# This file is part of Launch Control.
#
# Launch Control is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License version 3
# as published by the Free Software Foundation
#
# Launch Control is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with Launch Control.  If not, see <http://www.gnu.org/licenses/>.

"""
DataViews: Encapsulated SQL query definitions.

Implementation of the following launchpad blueprint:
https://blueprints.launchpad.net/launch-control/+spec/other-linaro-n-data-views-for-launch-control
"""

from xml.sax.handler import ContentHandler
from xml.sax import parseString
import re

class DataView(object):
    """
    Data view, a container for SQL query and optional arguments
    """

    def __init__(self, name, sql, arguments, documentation, summary):
        self.name = name
        self.sql = sql
        self.arguments = arguments
        self.documentation = documentation
        self.summary = summary

    @classmethod
    def load_from_xml(self, xml_string):
        """
        Load a data view instance from XML description

        This raises ValueError in several error situations.
        TODO: check what kind of exceptions this can raise
        """
        handler = _DataViewHandler()
        parseString(xml_string, handler)
        return handler.dataview

class DataViewArgument(object):
    """
    Data view argument for SQL prepared statements
    """

    def __init__(self, name, type, default, help):
        self.name = name
        self.type = type
        self.default = default
        self.help = help

class _DataViewHandler(ContentHandler):
    """
    ContentHandler subclass for parsing DataView documents
    """

    class Object(object):
        """
        Simple container for named attributes
        """

    def startDocument(self):
        # Text can be None or a [] that accumulates all detected text
        self.text = None
        # Data view object
        self.obj = self.Object()
        self.obj.sql = {}
        self.obj.arguments = []

    def endDocument(self):
        self.dataview = DataView(self.obj.name,
            self.obj.sql, self.obj.arguments,
            self.obj.documentation, self.obj.summary)

    def startElement(self, name, attrs):
        if name == "data-view":
            self.obj.name = attrs["name"]
        elif name == "summary" or name == "documentation":
            self._start_text()
        elif name == "sql":
            self._start_text()
            self.obj.sql_backend = attrs.get("backend", "")
        elif name == "value":
            self.text.append("{" + attrs["name"] + "}")
        elif name == "argument":
            if "name" not in attrs:
                raise ValueError("Argument name missing")
            if "type" not in attrs:
                raise ValueError("Argument type missing")
            if attrs["type"] not in ("string", "number", "boolean", "timestamp"):
                raise ValueError("Invalid argument type")
            arg = DataViewArgument(name=attrs["name"], type=attrs["type"],
                                   default=attrs.get("default", None),
                                   help=attrs.get("help", None))
            self.obj.arguments.append(arg)

    def _end_text(self):
        """
        Stop collecting text and produce a stripped string with deduplicated whitespace
        """
        full_text = re.sub("\s+", " ", u''.join(self.text)).strip()
        self.text = None
        return full_text

    def _start_text(self):
        """
        Start collecting text
        """
        self.text = []

    def endElement(self, name):
        if name == "sql":
            self.obj.sql[self.obj.sql_backend] = self._end_text()
        elif name == "documentation":
            self.obj.documentation = self._end_text()
        elif name == "summary":
            self.obj.summary = self._end_text()

    def characters(self, content):
        if isinstance(self.text, list):
            self.text.append(content)

__all__ = ["DataView", "DataViewArgument"]
<commit_msg>Add DataViewRepository implementation<commit_after># Copyright (C) 2011 Linaro Limited
#
# Author: Zygmunt Krynicki <zygmunt.krynicki@linaro.org>
#
# This file is part of Launch Control.
#
# Launch Control is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License version 3
# as published by the Free Software Foundation
#
# Launch Control is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with Launch Control.  If not, see <http://www.gnu.org/licenses/>.

"""
DataViews: Encapsulated SQL query definitions.

Implementation of the following launchpad blueprint:
https://blueprints.launchpad.net/launch-control/+spec/other-linaro-n-data-views-for-launch-control
"""

from xml.sax import parseString
from xml.sax.handler import ContentHandler
import logging
import os
import re

class DataView(object):
    """
    Data view, a container for SQL query and optional arguments
    """

    def __init__(self, name, sql, arguments, documentation, summary):
        self.name = name
        self.sql = sql
        self.arguments = arguments
        self.documentation = documentation
        self.summary = summary

    @classmethod
    def load_from_xml(self, xml_string):
        """
        Load a data view instance from XML description

        This raises ValueError in several error situations.
        TODO: check what kind of exceptions this can raise
        """
        handler = _DataViewHandler()
        parseString(xml_string, handler)
        return handler.dataview

class DataViewArgument(object):
    """
    Data view argument for SQL prepared statements
    """

    def __init__(self, name, type, default, help):
        self.name = name
        self.type = type
        self.default = default
        self.help = help

class _DataViewHandler(ContentHandler):
    """
    ContentHandler subclass for parsing DataView documents
    """

    class Object(object):
        """
        Simple container for named attributes
        """

    def startDocument(self):
        # Text can be None or a [] that accumulates all detected text
        self.text = None
        # Data view object
        self.obj = self.Object()
        self.obj.sql = {}
        self.obj.arguments = []

    def endDocument(self):
        self.dataview = DataView(self.obj.name,
            self.obj.sql, self.obj.arguments,
            self.obj.documentation, self.obj.summary)

    def startElement(self, name, attrs):
        if name == "data-view":
            self.obj.name = attrs["name"]
        elif name == "summary" or name == "documentation":
            self._start_text()
        elif name == "sql":
            self._start_text()
            self.obj.sql_backend = attrs.get("backend", "")
        elif name == "value":
            self.text.append("{" + attrs["name"] + "}")
        elif name == "argument":
            if "name" not in attrs:
                raise ValueError("Argument name missing")
            if "type" not in attrs:
                raise ValueError("Argument type missing")
            if attrs["type"] not in ("string", "number", "boolean", "timestamp"):
                raise ValueError("Invalid argument type")
            arg = DataViewArgument(name=attrs["name"], type=attrs["type"],
                                   default=attrs.get("default", None),
                                   help=attrs.get("help", None))
            self.obj.arguments.append(arg)

    def _end_text(self):
        """
        Stop collecting text and produce a stripped string with deduplicated whitespace
        """
        full_text = re.sub("\s+", " ", u''.join(self.text)).strip()
        self.text = None
        return full_text

    def _start_text(self):
        """
        Start collecting text
        """
        self.text = []

    def endElement(self, name):
        if name == "sql":
            self.obj.sql[self.obj.sql_backend] = self._end_text()
        elif name == "documentation":
            self.obj.documentation = self._end_text()
        elif name == "summary":
            self.obj.summary = self._end_text()

    def characters(self, content):
        if isinstance(self.text, list):
            self.text.append(content)

class DataViewRepository(object):

    _instance = None

    def __init__(self):
        self.data_views = []

    def load_from_directory(self, directory):
        for name in os.listdir(directory):
            pathname = os.path.join(directory, name)
            if os.path.isfile(pathname) and pathname.endswith(".xml"):
                self.load_from_file(pathname)

    def load_from_file(self, pathname):
        try:
            with open(pathname, "rt") as stream:
                text = stream.read()
            data_view = DataView.load_from_xml(text)
            self.data_views.append(data_view)
        except Exception as exc:
            logging.error("Unable to load data view from %s: %s", pathname, exc)

    @classmethod
    def get_instance(cls):
        from django.conf import settings
        if cls._instance is None:
            cls._instance = cls()
            cls._instance.load_default()

        # I development mode always reload data views
        if getattr(settings, "DEBUG", False) is True:
            cls._instance.data_views = []
            cls._instance.load_default()
        return cls._instance

    def load_default(self):
        from django.conf import settings
        for dirname in getattr(settings, "DATAVIEW_DIRS", []):
            self.load_from_directory(dirname)

__all__ = ["DataView", "DataViewArgument", "DataViewRepository"]
<|endoftext|>


<commit_before>__all__ = ['shuffleMockCatalog', 'generate_upid']
import warnings
from itertools import izip
import numpy as np
from numpy.lib.recfunctions import rename_fields

def _iter_plateau_in_sorted_array(a):
    k = np.where(a[1:] != a[:-1])[0]
    k += 1
    i = 0
    for j in k:
        yield i, j
        i = j
    yield i, len(a)

def _iter_indices_in_bins(bins, a):
    s = a.argsort()
    k = np.searchsorted(a, bins, 'right', sorter=s)
    i = 0
    for j in k:
        yield s[i:j]
        i = j
    yield s[i:]

def _random_rotation_matrix(n=3):
    A = np.linalg.qr(np.random.randn(n,n))[0]
    if np.linalg.det(A) > 0:
        return A
    return _random_rotation_matrix(n)

_axes = list('xyz')

def _get_xyz(a, ax_type=float):
    return np.fromiter((a[ax] for ax in _axes), ax_type, 3)

def generate_upid(pid, id, recursive=True):
    """
    To generate (or to fix) the upid of a halo catalog.

    Parameters
    ----------
    pid : array_like
        An ndarray of integer that contains the parent IDs of each halo.
    id : array_like
        An ndarray of integer that contains the halo IDs.
    recursive : bool, optional
        Whether or not to run this function recursively. Default is True.

    Returns
    -------
    upid : array_like
        The ultimate parent IDs.

    Examples
    --------
    >>> halos['upid'] = generate_upid(halos['pid'], halos['id'])

    """
    pid = np.asanyarray(pid)
    id = np.asanyarray(id)
    s = pid.argsort()
    idx = np.fromiter(_iter_plateau_in_sorted_array(pid[s]), \
            np.dtype([('start', int), ('stop', int)]))
    unique_pid = pid[s[idx['start']]]
    if unique_pid[0] == -1:
        unique_pid = unique_pid[1:]
        idx = idx[1:]
    host_flag = (pid == -1)
    not_found = np.where(np.in1d(unique_pid, id[host_flag], True, True))[0]
    if not len(not_found):
        return pid
    sub_flag = np.where(~host_flag)[0]
    found = sub_flag[np.in1d(id[sub_flag], unique_pid[not_found], True)]
    found = found[id[found].argsort()]
    assert (id[found] == unique_pid[not_found]).all()
    del host_flag, sub_flag, unique_pid
    pid_old = pid.copy()
    for i, j in izip(found, not_found):
        pid[s[slice(*idx[j])]] = pid_old[i]
    del pid_old, idx, s, found, not_found
    return generate_upid(pid, id, True) if recursive else pid

def shuffleMockCatalog(mock_ids, halo_catalog, bin_width=None, bins=None,
        proxy='mvir', box_size=None, apply_rsd=False,
        shuffle_centrals=True, shuffle_satellites=True, rotate_satellites=False,
        return_structured_array=False):
    """
    Shuffle a mock catalog according to Zentner et al. (2014) [arXiv:1311.1818]

    Parameters
    ----------
    mock_ids : array_like
        Should be a 1-d array of int which contains the corresponding halo IDs
        for the galaxies in the mock catalog to be shuffled.
    halo_catalog : array_like
        Should be a 1-d structrued array which has the following fields:
        id, upid, x, y, z, vz (if `apply_rsd` it True), and the proxy.
    bin_width : float or None, optional
        The width of the bin, in dex.
    bins : int, array_like, or None, optional
        If an integer is provided, it is interpreted as the number of bins.
        If an array is provided, it is interpreted as the edges of the bins.
        The parameter _overwrites_ `bin_width`.
    proxy : string, optional
        The proxy to bin on. Must be present in the fields of `halo_catalog`.
    box_size : float or None, optional
        The side length of the box. Should be in the same unit as x, y, z.
    apply_rsd : bool, optional
        Whether or not to apply redshift space distortions on the z-axis.
        (Default is False)
    shuffle_centrals : bool, optional
        Whether or not to shuffle central galaxies (Default is True)
    shuffle_satellites : bool, optional
        Whether or not to shuffle satellite galaxies (Default is True)
    rotate_satellites : bool, optional
        Whether or not to apply a random rotation to satellite galaxies
        (Default is False)
    return_structured_array : bool, optional
        Whether to return a structured array that contains x, y, z
        or just a n-by-3 float array.

    Returns
    -------
    pos : array_like
        A ndarray that contains x, y, z of the shuffled positions.
    """

    # check necessary fields in halo_catalog
    fields = ['id', 'upid', proxy] + _axes
    if apply_rsd:
        fields.append('vz')
    if not all((f in halo_catalog.dtype.names for f in fields)):
        raise ValueError('`halo_catalog` should have the following fields: '+ \
                ', '.join(fields))

    # check dtype
    ax_type = halo_catalog['x'].dtype.type
    if any((halo_catalog[ax].dtype.type != ax_type for ax in 'yz')):
        raise ValueError('The types of fields x, y, z in `halo_catalog` ' \
                'must all be the same.')

    # check all mock_ids are in halo_catalog
    s = halo_catalog['id'].argsort()
    idx = np.searchsorted(halo_catalog['id'], mock_ids, sorter=s)
    try:
        idx = s[idx]
    except IndexError:
        raise ValueError('`mock_ids` must all present in `halo_catalog`')
    if not (halo_catalog['id'][idx] == mock_ids).all():
        raise ValueError('`mock_ids` must all present in `halo_catalog`')
    mock_idx = np.ones(len(halo_catalog), dtype=int)
    mock_idx *= -1
    mock_idx[idx] = np.arange(len(mock_ids))
    del s, idx

    # separate hosts and subs
    host_flag = (halo_catalog['upid'] == -1)
    hosts = rename_fields(halo_catalog[host_flag], {'upid':'mock_idx'})
    hosts['mock_idx'] = mock_idx[host_flag]
    subs = rename_fields(halo_catalog[~host_flag], {'id':'mock_idx'})
    subs['mock_idx'] = mock_idx[~host_flag]
    subs = subs[subs['mock_idx'] > -1]
    del host_flag, mock_idx

    # group subhalos
    subs.sort(order='upid')
    idx = np.fromiter(_iter_plateau_in_sorted_array(subs['upid']), \
            np.dtype([('start', int), ('stop', int)]))
    host_ids = subs['upid'][idx['start']]
    if not np.in1d(host_ids, hosts['id'], True).all():
        raise ValueError('Some subhalos associdated with the mock galaxies ' \
                'have no parent halos in `halo_catalog`. Consider using ' \
                '`generate_upid` to fix this.')
    subs_idx = np.zeros(len(hosts), dtype=idx.dtype)
    subs_idx[np.in1d(hosts['id'], host_ids, True)] = idx
    del idx, host_ids

    # check bins
    try:
        bin_width = float(bin_width)
    except (ValueError, TypeError):
        bin_width = None
    else:
        if bin_width <= 0:
            bin_width = None
    if bin_width is None:
        bin_width = 0.1

    mi = np.log10(hosts[proxy].min()*0.99999)
    ma = np.log10(hosts[proxy].max())

    if bins is None:
        bins = int(np.ceil((ma-mi)/bin_width))
        mi = ma - bin_width*bins
    try:
        bins = int(bins)
    except (ValueError, TypeError):
        bins = np.asarray(bins)
        if len(bins) < 2 or (bins[1:]<bins[:-1]).any():
            raise ValueError('Please specify a valid `bin` parameter.')
    else:
        bins = np.logspace(mi, ma, bins+1)

    # create the array for storing results
    pos = np.empty((len(mock_ids), 3), ax_type)
    pos.fill(np.nan)

    # loop of bins of proxy (e.g. mvir)
    for i, indices in enumerate(_iter_indices_in_bins(bins, hosts[proxy])):
        if not len(indices):
            continue

        if i==0 or i==len(bins):
            if (hosts['mock_idx'][indices] > -1).any() or \
                    any((subs_idx['start'][j] < subs_idx['stop'][j] \
                    for j in indices)):
                warnings.warn('Some halos associdated with the mock catalog ' \
                        'are outside the bin range.', RuntimeWarning)
            continue

        # shuffle satellites
        if shuffle_satellites:
            choices = indices.tolist()
        for j in indices:
            subs_this = subs[slice(*subs_idx[j])]
            if not len(subs_this):
                continue
            mock_idx_this = subs_this['mock_idx']
            pos[mock_idx_this] = subs_this[_axes].view((ax_type,3))
            if shuffle_satellites:
                k = choices.pop(np.random.randint(len(choices)))
                pos[mock_idx_this] -= _get_xyz(hosts[j], ax_type)
                if rotate_satellites:
                    pos[mock_idx_this] = np.dot(pos[mock_idx_this], \
                            _random_rotation_matrix())
                pos[mock_idx_this] += _get_xyz(hosts[k], ax_type)
                if apply_rsd:
                    pos[mock_idx_this,2] += (subs_this['vz'] \
                            + hosts['vz'][k] - hosts['vz'][j])/100.0
            else:
                if rotate_satellites:
                    host_pos = _get_xyz(hosts[j], ax_type)
                    pos[mock_idx_this] -= host_pos
                    pos[mock_idx_this] = np.dot(pos[mock_idx_this], \
                            _random_rotation_matrix())
                    pos[mock_idx_this] += host_pos
                if apply_rsd:
                    pos[mock_idx_this,2] += subs_this['vz']/100.0

        # shuffle hosts
        has_mock = indices[hosts['mock_idx'][indices] > -1]
        if not len(has_mock):
            continue
        mock_idx_this = hosts['mock_idx'][has_mock]
        if shuffle_centrals:
            has_mock = np.random.choice(indices, len(has_mock), False)
        pos[mock_idx_this] = hosts[_axes][has_mock].view((ax_type,3))
        if apply_rsd:
            pos[mock_idx_this,2] += hosts['vz'][has_mock]/100.0

    # sanity check
    if np.isnan(pos).any():
        warnings.warn('Some galaxies in the mock catalog have not been ' \
                'assigned a new position. Maybe the corresponding halo is ' \
                'outside the bin range.', RuntimeWarning)

    # wrap box
    if box_size is not None:
        pos = np.remainder(pos, box_size, pos)

    if return_structured_array:
        pos = pos.view(np.dtype(zip(_axes, [ax_type]*3)))

    return pos

<commit_msg>fix periodic b.c. when doing rotation<commit_after>__all__ = ['shuffleMockCatalog', 'generate_upid']
import warnings
from itertools import izip
import numpy as np
from numpy.lib.recfunctions import rename_fields

def _iter_plateau_in_sorted_array(a):
    k = np.where(a[1:] != a[:-1])[0]
    k += 1
    i = 0
    for j in k:
        yield i, j
        i = j
    yield i, len(a)

def _iter_indices_in_bins(bins, a):
    s = a.argsort()
    k = np.searchsorted(a, bins, 'right', sorter=s)
    i = 0
    for j in k:
        yield s[i:j]
        i = j
    yield s[i:]

def _random_rotation_matrix(n=3):
    A = np.linalg.qr(np.random.randn(n,n))[0]
    if np.linalg.det(A) > 0:
        return A
    return _random_rotation_matrix(n)

def _apply_rotation(pos, box_size):
    half_box_size = box_size * 0.5
    pos = np.copy(pos)
    pos[pos >  half_box_size] -= box_size
    pos[pos < -half_box_size] += box_size
    A = _random_rotation_matrix()
    return np.dot(pos, A, pos)

_axes = list('xyz')

def _get_xyz(a, ax_type=float):
    return np.fromiter((a[ax] for ax in _axes), ax_type, 3)

def generate_upid(pid, id, recursive=True):
    """
    To generate (or to fix) the upid of a halo catalog.

    Parameters
    ----------
    pid : array_like
        An ndarray of integer that contains the parent IDs of each halo.
    id : array_like
        An ndarray of integer that contains the halo IDs.
    recursive : bool, optional
        Whether or not to run this function recursively. Default is True.

    Returns
    -------
    upid : array_like
        The ultimate parent IDs.

    Examples
    --------
    >>> halos['upid'] = generate_upid(halos['pid'], halos['id'])

    """
    pid = np.asanyarray(pid)
    id = np.asanyarray(id)
    s = pid.argsort()
    idx = np.fromiter(_iter_plateau_in_sorted_array(pid[s]), \
            np.dtype([('start', int), ('stop', int)]))
    unique_pid = pid[s[idx['start']]]
    if unique_pid[0] == -1:
        unique_pid = unique_pid[1:]
        idx = idx[1:]
    host_flag = (pid == -1)
    not_found = np.where(np.in1d(unique_pid, id[host_flag], True, True))[0]
    if not len(not_found):
        return pid
    sub_flag = np.where(~host_flag)[0]
    found = sub_flag[np.in1d(id[sub_flag], unique_pid[not_found], True)]
    found = found[id[found].argsort()]
    assert (id[found] == unique_pid[not_found]).all()
    del host_flag, sub_flag, unique_pid
    pid_old = pid.copy()
    for i, j in izip(found, not_found):
        pid[s[slice(*idx[j])]] = pid_old[i]
    del pid_old, idx, s, found, not_found
    return generate_upid(pid, id, True) if recursive else pid

def shuffleMockCatalog(mock_ids, halo_catalog, bin_width=None, bins=None,
        proxy='mvir', box_size=None, apply_rsd=False,
        shuffle_centrals=True, shuffle_satellites=True, rotate_satellites=False,
        return_structured_array=False):
    """
    Shuffle a mock catalog according to Zentner et al. (2014) [arXiv:1311.1818]

    Parameters
    ----------
    mock_ids : array_like
        Should be a 1-d array of int which contains the corresponding halo IDs
        for the galaxies in the mock catalog to be shuffled.
    halo_catalog : array_like
        Should be a 1-d structrued array which has the following fields:
        id, upid, x, y, z, vz (if `apply_rsd` it True), and the proxy.
    bin_width : float or None, optional
        The width of the bin, in dex.
    bins : int, array_like, or None, optional
        If an integer is provided, it is interpreted as the number of bins.
        If an array is provided, it is interpreted as the edges of the bins.
        The parameter _overwrites_ `bin_width`.
    proxy : string, optional
        The proxy to bin on. Must be present in the fields of `halo_catalog`.
    box_size : float or None, optional
        The side length of the box. Should be in the same unit as x, y, z.
    apply_rsd : bool, optional
        Whether or not to apply redshift space distortions on the z-axis.
        (Default is False)
    shuffle_centrals : bool, optional
        Whether or not to shuffle central galaxies (Default is True)
    shuffle_satellites : bool, optional
        Whether or not to shuffle satellite galaxies (Default is True)
    rotate_satellites : bool, optional
        Whether or not to apply a random rotation to satellite galaxies
        (Default is False)
    return_structured_array : bool, optional
        Whether to return a structured array that contains x, y, z
        or just a n-by-3 float array.

    Returns
    -------
    pos : array_like
        A ndarray that contains x, y, z of the shuffled positions.
    """

    # check necessary fields in halo_catalog
    fields = ['id', 'upid', proxy] + _axes
    if apply_rsd:
        fields.append('vz')
    if not all((f in halo_catalog.dtype.names for f in fields)):
        raise ValueError('`halo_catalog` should have the following fields: '+ \
                ', '.join(fields))

    # check dtype
    ax_type = halo_catalog['x'].dtype.type
    if any((halo_catalog[ax].dtype.type != ax_type for ax in 'yz')):
        raise ValueError('The types of fields x, y, z in `halo_catalog` ' \
                'must all be the same.')

    # check all mock_ids are in halo_catalog
    s = halo_catalog['id'].argsort()
    idx = np.searchsorted(halo_catalog['id'], mock_ids, sorter=s)
    try:
        idx = s[idx]
    except IndexError:
        raise ValueError('`mock_ids` must all present in `halo_catalog`')
    if not (halo_catalog['id'][idx] == mock_ids).all():
        raise ValueError('`mock_ids` must all present in `halo_catalog`')
    mock_idx = np.ones(len(halo_catalog), dtype=int)
    mock_idx *= -1
    mock_idx[idx] = np.arange(len(mock_ids))
    del s, idx

    # separate hosts and subs
    host_flag = (halo_catalog['upid'] == -1)
    hosts = rename_fields(halo_catalog[host_flag], {'upid':'mock_idx'})
    hosts['mock_idx'] = mock_idx[host_flag]
    subs = rename_fields(halo_catalog[~host_flag], {'id':'mock_idx'})
    subs['mock_idx'] = mock_idx[~host_flag]
    subs = subs[subs['mock_idx'] > -1]
    del host_flag, mock_idx

    # group subhalos
    subs.sort(order='upid')
    idx = np.fromiter(_iter_plateau_in_sorted_array(subs['upid']), \
            np.dtype([('start', int), ('stop', int)]))
    host_ids = subs['upid'][idx['start']]
    if not np.in1d(host_ids, hosts['id'], True).all():
        raise ValueError('Some subhalos associdated with the mock galaxies ' \
                'have no parent halos in `halo_catalog`. Consider using ' \
                '`generate_upid` to fix this.')
    subs_idx = np.zeros(len(hosts), dtype=idx.dtype)
    subs_idx[np.in1d(hosts['id'], host_ids, True)] = idx
    del idx, host_ids

    # check bins
    try:
        bin_width = float(bin_width)
    except (ValueError, TypeError):
        bin_width = None
    else:
        if bin_width <= 0:
            bin_width = None
    if bin_width is None:
        bin_width = 0.1

    mi = np.log10(hosts[proxy].min()*0.99999)
    ma = np.log10(hosts[proxy].max())

    if bins is None:
        bins = int(np.ceil((ma-mi)/bin_width))
        mi = ma - bin_width*bins
    try:
        bins = int(bins)
    except (ValueError, TypeError):
        bins = np.asarray(bins)
        if len(bins) < 2 or (bins[1:]<bins[:-1]).any():
            raise ValueError('Please specify a valid `bin` parameter.')
    else:
        bins = np.logspace(mi, ma, bins+1)

    # create the array for storing results
    pos = np.empty((len(mock_ids), 3), ax_type)
    pos.fill(np.nan)

    # loop of bins of proxy (e.g. mvir)
    for i, indices in enumerate(_iter_indices_in_bins(bins, hosts[proxy])):
        if not len(indices):
            continue

        if i==0 or i==len(bins):
            if (hosts['mock_idx'][indices] > -1).any() or \
                    any((subs_idx['start'][j] < subs_idx['stop'][j] \
                    for j in indices)):
                warnings.warn('Some halos associdated with the mock catalog ' \
                        'are outside the bin range.', RuntimeWarning)
            continue

        # shuffle satellites
        if shuffle_satellites:
            choices = indices.tolist()
        for j in indices:
            subs_this = subs[slice(*subs_idx[j])]
            if not len(subs_this):
                continue
            mock_idx_this = subs_this['mock_idx']
            pos[mock_idx_this] = subs_this[_axes].view((ax_type,3))
            if shuffle_satellites:
                k = choices.pop(np.random.randint(len(choices)))
                pos[mock_idx_this] -= _get_xyz(hosts[j], ax_type)
                if rotate_satellites:
                    pos[mock_idx_this] = \
                            _apply_rotation(pos[mock_idx_this], box_size)
                pos[mock_idx_this] += _get_xyz(hosts[k], ax_type)
                if apply_rsd:
                    pos[mock_idx_this,2] += (subs_this['vz'] \
                            + hosts['vz'][k] - hosts['vz'][j])/100.0
            else:
                if rotate_satellites:
                    host_pos = _get_xyz(hosts[j], ax_type)
                    pos[mock_idx_this] -= host_pos
                    pos[mock_idx_this] = \
                            _apply_rotation(pos[mock_idx_this], box_size)
                    pos[mock_idx_this] += host_pos
                if apply_rsd:
                    pos[mock_idx_this,2] += subs_this['vz']/100.0

        # shuffle hosts
        has_mock = indices[hosts['mock_idx'][indices] > -1]
        if not len(has_mock):
            continue
        mock_idx_this = hosts['mock_idx'][has_mock]
        if shuffle_centrals:
            has_mock = np.random.choice(indices, len(has_mock), False)
        pos[mock_idx_this] = hosts[_axes][has_mock].view((ax_type,3))
        if apply_rsd:
            pos[mock_idx_this,2] += hosts['vz'][has_mock]/100.0

    # sanity check
    if np.isnan(pos).any():
        warnings.warn('Some galaxies in the mock catalog have not been ' \
                'assigned a new position. Maybe the corresponding halo is ' \
                'outside the bin range.', RuntimeWarning)

    # wrap box
    if box_size is not None:
        pos = np.remainder(pos, box_size, pos)

    if return_structured_array:
        pos = pos.view(np.dtype(zip(_axes, [ax_type]*3)))

    return pos

<|endoftext|>


<commit_before># a class to manage SRTM surfaces

import json
import numpy as np
import os
from pylab import *
import random
import scipy.interpolate
import struct
import urllib
import zipfile

import navpy

# return the lower left corner of the 1x1 degree tile containing
# the specified lla coordinate
def lla_ll_corner(lat_deg, lon_deg):
    return int(floor(lat_deg)), int(floor(lon_deg))

# return the tile base name for the specified coordinate
def make_tile_name(lat, lon):
    ll_lat, ll_lon = lla_ll_corner(lat, lon)
    if ll_lat < 0:
        slat = "S%2d" % ll_lat
    else:
        slat = "N%2d" % ll_lat
    if ll_lon < 0:
        slon = "W%03d" % -ll_lon
    else:
        slon = "W%03d" % ll_lon
    return slat + slon

class SRTM():
    def __init__(self, lat, lon, dict_path):
        self.lat, self.lon = lla_ll_corner(lat, lon)
        self.srtm_dict = {}
        self.srtm_cache_dir = '/var/tmp' # unless set otherwise
        self.srtm_z = None
        self.i = None
        self.load_srtm_dict(dict_path)

    # load the directory download dictionary (mapping a desired file
    # to a download path.)
    def load_srtm_dict(self, dict_path):
        dict_file = dict_path + '/srtm.json'
        try:
            f = open(dict_file, 'r')
            self.srtm_dict = json.load(f)
            f.close()
        except:
            print "Notice: unable to read =", dict_file

    # if we'd like a persistant place to cache srtm files so they don't
    # need to be re-download every time we run anything
    def set_srtm_cache_dir(self, srtm_cache_dir):
        if not os.path.exists(srtm_cache_dir):
            print "Notice: srtm cache path doesn't exist =", srtm_cache_dir
        else:
            self.srtm_cache_dir = srtm_cache_dir

    # download and extract srtm file into cache directory
    def download_srtm(self, fileroot):
        if fileroot in self.srtm_dict:
            url = self.srtm_dict[fileroot]
            download_file = self.srtm_cache_dir + '/' + fileroot + '.hgt.zip'
            print "Notice: downloading:", url
            file = urllib.URLopener()
            print file.retrieve(url, download_file)
            return True
        else:
            print "Notice: requested srtm that is outside catalog"
            return False

    def parse(self):
        tilename = make_tile_name(self.lat, self.lon)
        cache_file = self.srtm_cache_dir + '/' + tilename + '.hgt.zip'
        if not os.path.exists(cache_file):
            if not self.download_srtm(tilename):
                return False
        print "Notice: parsing SRTM file:", cache_file
        #f = open(cache_file, "rb")
        zip = zipfile.ZipFile(cache_file)
        f = zip.open(tilename + '.hgt', 'r')
        contents = f.read()
        f.close()
        # read 1,442,401 (1201x1201) high-endian
        # signed 16-bit words into self.z
        self.srtm_z = struct.unpack(">1442401H", contents)
        return True

    def make_lla_interpolator(self):
        print "Notice: constructing LLA interpolator"

        srtm_pts = np.zeros((1201, 1201))
        for r in range(0,1201):
            for c in range(0,1201):
                idx = (1201*r)+c
                va = self.srtm_z[idx]
                if va == 65535 or va < 0 or va > 10000:
                    va = 0.0
                z = va
                srtm_pts[c,1200-r] = z
        x = np.linspace(self.lon, self.lon+1, 1201)
        y = np.linspace(self.lat, self.lat+1, 1201)
        #print x
        #print y
        self.lla_interp = scipy.interpolate.RegularGridInterpolator((x, y), srtm_pts, bounds_error=False, fill_value=-32768)
        #print self.lla_interp([-93.14530573529404, 45.220697421008396])
        #for i in range(20):
        #    x = -94.0 + random.random()
        #    y = 45 + random.random()
        #    z = self.lla_interp([x,y])
        #    print [x, y, z[0]]

    def lla_interpolate(self, point_list):
        return self.lla_interp(point_list)

    def ned_interpolate(self, point_list):
        return self.ned_interp(point_list)

    def plot_raw(self):
        zzz = np.zeros((1201,1201))
        for r in range(0,1201):
            for c in range(0,1201):
                va=self.srtm_z[(1201*r)+c]
                if (va==65535 or va<0 or va>10000):
                    va=0.0
                zzz[r][c]=float(va)

        #zz=np.log1p(zzz)
        imshow(zzz, interpolation='bilinear',cmap=cm.gray,alpha=1.0)
        grid(False)
        show()

# Build a gridded elevation interpolation table centered at lla_ref
# with width and height.  This is a little bit of quick feet dancing,
# but allows areas to span corners or edges of srtm tiles and attempts
# to stay on a fast path of regular grids, even though a regularly lla
# grid != a regular ned grid.
class NEDGround():
    def __init__(self, lla_ref, width_m, height_m, step_m):
        self.tile_dict = {}
        self.load_tiles(lla_ref, width_m, height_m)
        self.make_interpolator(lla_ref, width_m, height_m, step_m)

    def load_tiles(self, lla_ref, width_m, height_m):
        print "Notice: loading DEM tiles"
        ll_ned = np.array([[-height_m*0.5, -width_m*0.5, 0.0]])
        ur_ned = np.array([[height_m*0.5, width_m*0.5, 0.0]])
        ll_lla = navpy.ned2lla(ll_ned, lla_ref[0], lla_ref[1], lla_ref[2])
        ur_lla = navpy.ned2lla(ur_ned, lla_ref[0], lla_ref[1], lla_ref[2])
        #print ll_lla
        #print ur_lla
        lat1, lon1 = lla_ll_corner( ll_lla[0], ll_lla[1] )
        lat2, lon2 = lla_ll_corner( ur_lla[0], ur_lla[1] )
        for lat in range(lat1, lat2+1):
            for lon in range(lon1, lon2+1):
                srtm = SRTM(lat, lon, '../srtm')
                if srtm.parse():
                    srtm.make_lla_interpolator()
                    #srtm.plot_raw()
                    tile_name = make_tile_name(lat, lon)
                    self.tile_dict[tile_name] = srtm

    def make_interpolator(self, lla_ref, width_m, height_m, step_m):
        print "Notice: constructing NED area interpolator"
        rows = (height_m / step_m) + 1
        cols = (width_m / step_m) + 1
        #ned_pts = np.zeros((cols, rows))
        #for r in range(0,rows):
        #    for c in range(0,cols):
        #        idx = (cols*r)+c
        #        #va = self.srtm_z[idx]
        #        #if va == 65535 or va < 0 or va > 10000:
        #        #    va = 0.0
        #        #z = va
        #        ned_pts[c,r] = 0.0

        # build regularly gridded x,y coordinate list and ned_pts array
        n_list = np.linspace(-height_m*0.5, height_m*0.5, rows)
        e_list = np.linspace(-width_m*0.5, width_m*0.5, cols)
        #print "e's:", e_list
        #print "n's:", n_list
        ned_pts = []
        for e in e_list:
            for n in n_list:
                ned_pts.append( [n, e, 0] )

        # convert ned_pts list to lla coordinates (so it's not
        # necessarily an exact grid anymore, but we can now
        # interpolate elevations out of the lla interpolators for each
        # tile.
        navpy_pts = navpy.ned2lla(ned_pts, lla_ref[0], lla_ref[1], lla_ref[2])
        #print navpy_pts

        # build list of (lat, lon) points for doing actual lla
        # elevation lookup
        ll_pts = []
        for i in range( len(navpy_pts[0]) ):
            lat = navpy_pts[0][i]
            lon = navpy_pts[1][i]
            ll_pts.append( [ lon, lat ] )
        #print "ll_pts:", ll_pts

        # set all the elevations in the ned_ds list to the extreme
        # minimum value. (rows,cols) might seem funny, but (ne)d is
        # reversed from (xy)z ... don't think about it too much or
        # you'll get a headache. :-)
        ned_ds = np.zeros((rows,cols))
        ned_ds[:][:] = -32768
        #print "ned_ds:", ned_ds

        # for each tile loaded, interpolate as many elevation values
        # as we can, then copy the good values into ned_ds.  When we
        # finish all the loaded tiles, we should have elevations for
        # the entire range of points.
        for tile in self.tile_dict:
            zs = self.tile_dict[tile].lla_interpolate(np.array(ll_pts))
            #print zs
            # copy the good altitudes back to the corresponding ned points
            for r in range(0,rows):
                for c in range(0,cols):
                    idx = (rows*c)+r
                    if zs[idx] > -10000:
                        ned_ds[r,c] = zs[idx]

        # quick sanity check
        for r in range(0,rows):
            for c in range(0,cols):
                idx = (rows*c)+r
                if ned_ds[r,c] < -10000:
                    print "Problem interpolating elevation for:", ll_pts[idx]
                    ned_ds[r,c] = 0.0
        #print "ned_ds:", ned_ds

        # now finally build the actual grid interpolator with evenly
        # spaced ned n, e values and elevations interpolated out of
        # the srtm lla interpolator.
        self.interp = scipy.interpolate.RegularGridInterpolator((n_list, e_list), ned_ds, bounds_error=False, fill_value=-32768)

        do_plot = False
        if do_plot:
            imshow(ned_ds, interpolation='bilinear', origin='lower', cmap=cm.gray, alpha=1.0)
            grid(False)
            show()

        do_test = False
        if do_test:
            for i in range(40):
                ned = [(random.random()-0.5)*height_m,
                       (random.random()-0.5)*width_m,
                       0.0]
                lla = navpy.ned2lla(ned, lla_ref[0], lla_ref[1], lla_ref[2])
                #print "ned=%s, lla=%s" % (ned, lla)
                nedz = self.interp([ned[0], ned[1]])
                tile = make_tile_name(lla[0], lla[1])
                llaz = self.tile_dict[tile].lla_interpolate(np.array([lla[1], lla[0]]))
                print "nedz=%.2f llaz=%.2f" % (nedz, llaz)

    # while error > eps: find altitude at current point, new pt = proj
    # vector to current alt.
    def interpolate_vector(self, ned, v):
        #print ned
        p = ned[:] # copy hopefully

        # sanity check (always assume camera pose is above ground!)
        if v[2] <= 0.0:
            return p

        eps = 0.01
        count = 0
        #print "start:", p
        #print "vec:", v
        #print "ned:", ned
        ground = self.interp([p[0], p[1]])
        error = abs(p[2] + ground[0])
        #print "  p=%s ground=%s error=%s" % (p, ground, error)
        while error > eps and count < 25 and ground[0] > -32768:
            d_proj = -(ned[2] + ground[0])
            factor = d_proj / v[2]
            n_proj = v[0] * factor
            e_proj = v[1] * factor
            #print "proj = %s %s" % (n_proj, e_proj)
            p = [ ned[0] + n_proj, ned[1] + e_proj, ned[2] + d_proj ]
            #print "new p:", p
            ground = self.interp([p[0], p[1]])
            error = abs(p[2] + ground[0])
            #print "  p=%s ground=%.2f error = %.3f" % (p, ground, error)
            count += 1
        #print "ground:", ground[0]
        if ground[0] > -32768:
            return p
        else:
            #print " returning nans"
            return np.zeros(3)*np.nan

    # return a list of (3d) ground intersection points for the give
    # vector list and camera pose.  Vectors are already transformed
    # into ned orientation.
    def interpolate_vectors(self, ned, v_list):
        pt_list = []
        for v in v_list:
            p = self.interpolate_vector(ned, v.flatten())
            pt_list.append(p)
        return pt_list
<commit_msg>SRTM.py modify make_tile_name to work in Australia<commit_after># a class to manage SRTM surfaces

import json
import numpy as np
import os
from pylab import *
import random
import scipy.interpolate
import struct
import urllib
import zipfile

import navpy

# return the lower left corner of the 1x1 degree tile containing
# the specified lla coordinate
def lla_ll_corner(lat_deg, lon_deg):
    return int(floor(lat_deg)), int(floor(lon_deg))

# return the tile base name for the specified coordinate
def make_tile_name(lat, lon):
    ll_lat, ll_lon = lla_ll_corner(lat, lon)
    if ll_lat < 0:
        slat = "S%2d" % -ll_lat
    else:
        slat = "N%2d" % ll_lat
    if ll_lon < 0:
        slon = "W%03d" % -ll_lon
    else:
        slon = "E%03d" % ll_lon
    return slat + slon

class SRTM():
    def __init__(self, lat, lon, dict_path):
        self.lat, self.lon = lla_ll_corner(lat, lon)
        self.srtm_dict = {}
        self.srtm_cache_dir = '/var/tmp' # unless set otherwise
        self.srtm_z = None
        self.i = None
        self.load_srtm_dict(dict_path)

    # load the directory download dictionary (mapping a desired file
    # to a download path.)
    def load_srtm_dict(self, dict_path):
        dict_file = dict_path + '/srtm.json'
        try:
            f = open(dict_file, 'r')
            self.srtm_dict = json.load(f)
            f.close()
        except:
            print "Notice: unable to read =", dict_file

    # if we'd like a persistant place to cache srtm files so they don't
    # need to be re-download every time we run anything
    def set_srtm_cache_dir(self, srtm_cache_dir):
        if not os.path.exists(srtm_cache_dir):
            print "Notice: srtm cache path doesn't exist =", srtm_cache_dir
        else:
            self.srtm_cache_dir = srtm_cache_dir

    # download and extract srtm file into cache directory
    def download_srtm(self, fileroot):
        if fileroot in self.srtm_dict:
            url = self.srtm_dict[fileroot]
            download_file = self.srtm_cache_dir + '/' + fileroot + '.hgt.zip'
            print "Notice: downloading:", url
            file = urllib.URLopener()
            print file.retrieve(url, download_file)
            return True
        else:
            print "Notice: requested srtm that is outside catalog"
            return False

    def parse(self):
        tilename = make_tile_name(self.lat, self.lon)
        cache_file = self.srtm_cache_dir + '/' + tilename + '.hgt.zip'
        if not os.path.exists(cache_file):
            if not self.download_srtm(tilename):
                return False
        print "Notice: parsing SRTM file:", cache_file
        #f = open(cache_file, "rb")
        zip = zipfile.ZipFile(cache_file)
        f = zip.open(tilename + '.hgt', 'r')
        contents = f.read()
        f.close()
        # read 1,442,401 (1201x1201) high-endian
        # signed 16-bit words into self.z
        self.srtm_z = struct.unpack(">1442401H", contents)
        return True

    def make_lla_interpolator(self):
        print "Notice: constructing LLA interpolator"

        srtm_pts = np.zeros((1201, 1201))
        for r in range(0,1201):
            for c in range(0,1201):
                idx = (1201*r)+c
                va = self.srtm_z[idx]
                if va == 65535 or va < 0 or va > 10000:
                    va = 0.0
                z = va
                srtm_pts[c,1200-r] = z
        x = np.linspace(self.lon, self.lon+1, 1201)
        y = np.linspace(self.lat, self.lat+1, 1201)
        #print x
        #print y
        self.lla_interp = scipy.interpolate.RegularGridInterpolator((x, y), srtm_pts, bounds_error=False, fill_value=-32768)
        #print self.lla_interp([-93.14530573529404, 45.220697421008396])
        #for i in range(20):
        #    x = -94.0 + random.random()
        #    y = 45 + random.random()
        #    z = self.lla_interp([x,y])
        #    print [x, y, z[0]]

    def lla_interpolate(self, point_list):
        return self.lla_interp(point_list)

    def ned_interpolate(self, point_list):
        return self.ned_interp(point_list)

    def plot_raw(self):
        zzz = np.zeros((1201,1201))
        for r in range(0,1201):
            for c in range(0,1201):
                va=self.srtm_z[(1201*r)+c]
                if (va==65535 or va<0 or va>10000):
                    va=0.0
                zzz[r][c]=float(va)

        #zz=np.log1p(zzz)
        imshow(zzz, interpolation='bilinear',cmap=cm.gray,alpha=1.0)
        grid(False)
        show()

# Build a gridded elevation interpolation table centered at lla_ref
# with width and height.  This is a little bit of quick feet dancing,
# but allows areas to span corners or edges of srtm tiles and attempts
# to stay on a fast path of regular grids, even though a regularly lla
# grid != a regular ned grid.
class NEDGround():
    def __init__(self, lla_ref, width_m, height_m, step_m):
        self.tile_dict = {}
        self.load_tiles(lla_ref, width_m, height_m)
        self.make_interpolator(lla_ref, width_m, height_m, step_m)

    def load_tiles(self, lla_ref, width_m, height_m):
        print "Notice: loading DEM tiles"
        ll_ned = np.array([[-height_m*0.5, -width_m*0.5, 0.0]])
        ur_ned = np.array([[height_m*0.5, width_m*0.5, 0.0]])
        ll_lla = navpy.ned2lla(ll_ned, lla_ref[0], lla_ref[1], lla_ref[2])
        ur_lla = navpy.ned2lla(ur_ned, lla_ref[0], lla_ref[1], lla_ref[2])
        #print ll_lla
        #print ur_lla
        lat1, lon1 = lla_ll_corner( ll_lla[0], ll_lla[1] )
        lat2, lon2 = lla_ll_corner( ur_lla[0], ur_lla[1] )
        for lat in range(lat1, lat2+1):
            for lon in range(lon1, lon2+1):
                srtm = SRTM(lat, lon, '../srtm')
                if srtm.parse():
                    srtm.make_lla_interpolator()
                    #srtm.plot_raw()
                    tile_name = make_tile_name(lat, lon)
                    self.tile_dict[tile_name] = srtm

    def make_interpolator(self, lla_ref, width_m, height_m, step_m):
        print "Notice: constructing NED area interpolator"
        rows = (height_m / step_m) + 1
        cols = (width_m / step_m) + 1
        #ned_pts = np.zeros((cols, rows))
        #for r in range(0,rows):
        #    for c in range(0,cols):
        #        idx = (cols*r)+c
        #        #va = self.srtm_z[idx]
        #        #if va == 65535 or va < 0 or va > 10000:
        #        #    va = 0.0
        #        #z = va
        #        ned_pts[c,r] = 0.0

        # build regularly gridded x,y coordinate list and ned_pts array
        n_list = np.linspace(-height_m*0.5, height_m*0.5, rows)
        e_list = np.linspace(-width_m*0.5, width_m*0.5, cols)
        #print "e's:", e_list
        #print "n's:", n_list
        ned_pts = []
        for e in e_list:
            for n in n_list:
                ned_pts.append( [n, e, 0] )

        # convert ned_pts list to lla coordinates (so it's not
        # necessarily an exact grid anymore, but we can now
        # interpolate elevations out of the lla interpolators for each
        # tile.
        navpy_pts = navpy.ned2lla(ned_pts, lla_ref[0], lla_ref[1], lla_ref[2])
        #print navpy_pts

        # build list of (lat, lon) points for doing actual lla
        # elevation lookup
        ll_pts = []
        for i in range( len(navpy_pts[0]) ):
            lat = navpy_pts[0][i]
            lon = navpy_pts[1][i]
            ll_pts.append( [ lon, lat ] )
        #print "ll_pts:", ll_pts

        # set all the elevations in the ned_ds list to the extreme
        # minimum value. (rows,cols) might seem funny, but (ne)d is
        # reversed from (xy)z ... don't think about it too much or
        # you'll get a headache. :-)
        ned_ds = np.zeros((rows,cols))
        ned_ds[:][:] = -32768
        #print "ned_ds:", ned_ds

        # for each tile loaded, interpolate as many elevation values
        # as we can, then copy the good values into ned_ds.  When we
        # finish all the loaded tiles, we should have elevations for
        # the entire range of points.
        for tile in self.tile_dict:
            zs = self.tile_dict[tile].lla_interpolate(np.array(ll_pts))
            #print zs
            # copy the good altitudes back to the corresponding ned points
            for r in range(0,rows):
                for c in range(0,cols):
                    idx = (rows*c)+r
                    if zs[idx] > -10000:
                        ned_ds[r,c] = zs[idx]

        # quick sanity check
        for r in range(0,rows):
            for c in range(0,cols):
                idx = (rows*c)+r
                if ned_ds[r,c] < -10000:
                    print "Problem interpolating elevation for:", ll_pts[idx]
                    ned_ds[r,c] = 0.0
        #print "ned_ds:", ned_ds

        # now finally build the actual grid interpolator with evenly
        # spaced ned n, e values and elevations interpolated out of
        # the srtm lla interpolator.
        self.interp = scipy.interpolate.RegularGridInterpolator((n_list, e_list), ned_ds, bounds_error=False, fill_value=-32768)

        do_plot = False
        if do_plot:
            imshow(ned_ds, interpolation='bilinear', origin='lower', cmap=cm.gray, alpha=1.0)
            grid(False)
            show()

        do_test = False
        if do_test:
            for i in range(40):
                ned = [(random.random()-0.5)*height_m,
                       (random.random()-0.5)*width_m,
                       0.0]
                lla = navpy.ned2lla(ned, lla_ref[0], lla_ref[1], lla_ref[2])
                #print "ned=%s, lla=%s" % (ned, lla)
                nedz = self.interp([ned[0], ned[1]])
                tile = make_tile_name(lla[0], lla[1])
                llaz = self.tile_dict[tile].lla_interpolate(np.array([lla[1], lla[0]]))
                print "nedz=%.2f llaz=%.2f" % (nedz, llaz)

    # while error > eps: find altitude at current point, new pt = proj
    # vector to current alt.
    def interpolate_vector(self, ned, v):
        #print ned
        p = ned[:] # copy hopefully

        # sanity check (always assume camera pose is above ground!)
        if v[2] <= 0.0:
            return p

        eps = 0.01
        count = 0
        #print "start:", p
        #print "vec:", v
        #print "ned:", ned
        ground = self.interp([p[0], p[1]])
        error = abs(p[2] + ground[0])
        #print "  p=%s ground=%s error=%s" % (p, ground, error)
        while error > eps and count < 25 and ground[0] > -32768:
            d_proj = -(ned[2] + ground[0])
            factor = d_proj / v[2]
            n_proj = v[0] * factor
            e_proj = v[1] * factor
            #print "proj = %s %s" % (n_proj, e_proj)
            p = [ ned[0] + n_proj, ned[1] + e_proj, ned[2] + d_proj ]
            #print "new p:", p
            ground = self.interp([p[0], p[1]])
            error = abs(p[2] + ground[0])
            #print "  p=%s ground=%.2f error = %.3f" % (p, ground, error)
            count += 1
        #print "ground:", ground[0]
        if ground[0] > -32768:
            return p
        else:
            #print " returning nans"
            return np.zeros(3)*np.nan

    # return a list of (3d) ground intersection points for the give
    # vector list and camera pose.  Vectors are already transformed
    # into ned orientation.
    def interpolate_vectors(self, ned, v_list):
        pt_list = []
        for v in v_list:
            p = self.interpolate_vector(ned, v.flatten())
            pt_list.append(p)
        return pt_list
<|endoftext|>


<commit_before>import sieve.grammar
import sieve.handler
import sieve.utils

__all__ = ('Rule', 'RuleSyntaxError',)

class RuleSyntaxError(Exception):
    pass

class Rule(object):

    @classmethod
    def register(cls):
        try:
            sieve.handler.register(cls.RULE_TYPE, cls.RULE_IDENTIFIER, cls)
        except AttributeError:
            # this method shouldn't be called on the SieveRule class directly,
            # only on subclasses that implement specific rules
            raise NotImplementedError

    def __init__(self, arguments=None, tests=None):
        if arguments is None:
            self.arguments = []
        else:
            self.arguments = arguments
        if tests is None:
            self.tests = []
        else:
            self.tests = tests

    def __str__(self):
        s = [ "%s" % self.RULE_IDENTIFIER, ]
        for arg in self.arguments:
            s.append(" %s" % arg)
        s.append('\n')
        for test in self.tests:
            s.append("(\n%s)\n" % sieve.utils.indent_string(str(test), 2))
        return ''.join(s)

    def validate_arguments(self, tagged_args=None, positional_args=None):
        if tagged_args is None:
            tagged_args = {}
        if positional_args is None:
            positional_args = []

        seen_args = {}
        i, n = 0, len(self.arguments)
        while i < n:
            if not isinstance(self.arguments[i], SieveTag):
                break
            num_valid_args = 0
            for arg_name, arg_validator in tagged_args.iteritems():
                num_valid_args = arg_validator.validate(self.arguments, i)
                if num_valid_args > 0:
                    if arg_name in seen_args:
                        raise SieveRuleSyntaxError(
                                "%s argument to %s was already seen earlier: %s"
                                % (arg_name, self.RULE_IDENTIFIER,
                                   self.arguments[i])
                                )
                    seen_args[arg_name] = self.arguments[i:i+num_valid_args]
                    i += num_valid_args
                    break
        # TODO: make sure all non-optional tagged arguments were seen

        if len(positional_args) != (n - i):
            raise SieveRuleSyntaxError(
                    "%s requires %d positional arguments but %d were "
                    "supplied"
                    % (self.RULE_IDENTIFIER, len(positional_args), n - i))

        for arg_position, arg_validator in enumerate(positional_args):
            if arg_validator.validate(self.arguments, i + arg_position) == 0:
                raise SieveRuleSyntaxError(
                        "positional argument #%d to %s was not in the "
                        "expected format"
                        % (arg_position + 1, self.RULE_IDENTIFIER))

        return (seen_args, self.arguments[i:])

    def validate_arguments_size(self, min_args, max_args=None):
        if max_args is None:
            max_args = min_args
        if len(self.arguments) < min_args or len(self.arguments) > max_args:
            if max_args == min_args:
                msg = "%d" % min_args
            else:
                msg = "between %d and %d" % (min_args, max_args)
            raise RuleSyntaxError("%s takes %s arguments" % (
                self.RULE_IDENTIFIER, msg))

    def validate_tests_size(self, min_tests, max_tests=None):
        if max_tests is None:
            max_tests = min_tests
        if len(self.tests) < min_tests or len(self.tests) > max_tests:
            if max_tests == min_tests:
                msg = "%d" % min_tests
            else:
                msg = "between %d and %d" % (min_tests, max_tests)
            raise RuleSyntaxError("%s takes %s tests" % (
                self.RULE_IDENTIFIER, msg))

    def validate_arg_is_stringlist(self, index, length=None):
        if not (isinstance(self.arguments[index], list)
                and all(isinstance(arg, basestring)
                        for arg in self.arguments[index])):
            raise RuleSyntaxError(
                    "%s requires argument %d to be a string or list of strings"
                    % (self.RULE_IDENTIFIER, index)
                    )
        if length is not None and len(self.arguments[index]) != length:
            if length == 1:
                msg = "a single string or list of one string"
            else:
                msg = "a list of %d strings" % length
            raise RuleSyntaxError("%s requires argument %d to be %s" % (
                self.RULE_IDENTIFIER, index, msg))

    def validate_arg_is_tag(self, index, allowed_tags=None):
        if not isinstance(self.arguments[index], sieve.grammar.Tag):
            raise RuleSyntaxError("%s requires argument %d to be a tag" %
                    (self.RULE_IDENTIFIER, index))
        if (allowed_tags is not None
                and self.arguments[index] not in allowed_tags):
            raise SieveRuleSyntaxError(
                    "%s requires argument %d to be one of these tags: %s"
                    % (self.RULE_IDENTIFIER, index,
                        ', '.join([ ':'+tag for tag in allowed_tags]))
                    )

    def validate_arg_is_number(self, index):
        try:
            long(self.arguments[index])
        except TypeError:
            raise RuleSyntaxError("%s requires argument %d to be a number"
                    % (self.RULE_IDENTIFIER, index))

    def validate_arg_is_comparator(self, index):
        self.validate_arg_is_stringlist(index, 1)
        if not sieve.handler.get('comparator', self.arguments[index][0]):
            raise RuleSyntaxError(
                    "'%s' comparator is unknown/unsupported"
                    % self.arguments[index][0])

    def validate_arg_is_match_type(self, index):
        self.validate_arg_is_tag(index, ('IS', 'CONTAINS', 'MATCHES'))

    def validate_arg_is_address_part(self, index):
        self.validate_arg_is_tag(index, ('LOCALPART', 'DOMAIN', 'ALL'))

    def evaluate(self, message, state):
        raise NotImplementedError

<commit_msg>update validation fn to use new sieve.grammar structure<commit_after>import sieve.grammar
import sieve.handler
import sieve.utils

__all__ = ('Rule', 'RuleSyntaxError',)

class RuleSyntaxError(Exception):
    pass

class Rule(object):

    @classmethod
    def register(cls):
        try:
            sieve.handler.register(cls.RULE_TYPE, cls.RULE_IDENTIFIER, cls)
        except AttributeError:
            # this method shouldn't be called on the SieveRule class directly,
            # only on subclasses that implement specific rules
            raise NotImplementedError

    def __init__(self, arguments=None, tests=None):
        if arguments is None:
            self.arguments = []
        else:
            self.arguments = arguments
        if tests is None:
            self.tests = []
        else:
            self.tests = tests

    def __str__(self):
        s = [ "%s" % self.RULE_IDENTIFIER, ]
        for arg in self.arguments:
            s.append(" %s" % arg)
        s.append('\n')
        for test in self.tests:
            s.append("(\n%s)\n" % sieve.utils.indent_string(str(test), 2))
        return ''.join(s)

    def validate_arguments(self, tagged_args=None, positional_args=None):
        if tagged_args is None:
            tagged_args = {}
        if positional_args is None:
            positional_args = []

        seen_args = {}
        i, n = 0, len(self.arguments)
        while i < n:
            if not isinstance(self.arguments[i], sieve.grammar.Tag):
                break
            num_valid_args = 0
            for arg_name, arg_validator in tagged_args.iteritems():
                num_valid_args = arg_validator.validate(self.arguments, i)
                if num_valid_args > 0:
                    if arg_name in seen_args:
                        raise RuleSyntaxError(
                                "%s argument to %s was already seen earlier: %s"
                                % (arg_name, self.RULE_IDENTIFIER,
                                   self.arguments[i])
                                )
                    seen_args[arg_name] = self.arguments[i:i+num_valid_args]
                    i += num_valid_args
                    break
        # TODO: make sure all non-optional tagged arguments were seen

        if len(positional_args) != (n - i):
            raise RuleSyntaxError(
                    "%s requires %d positional arguments but %d were "
                    "supplied"
                    % (self.RULE_IDENTIFIER, len(positional_args), n - i))

        for arg_position, arg_validator in enumerate(positional_args):
            if arg_validator.validate(self.arguments, i + arg_position) == 0:
                raise RuleSyntaxError(
                        "positional argument #%d to %s was not in the "
                        "expected format"
                        % (arg_position + 1, self.RULE_IDENTIFIER))

        return (seen_args, self.arguments[i:])

    def validate_arguments_size(self, min_args, max_args=None):
        if max_args is None:
            max_args = min_args
        if len(self.arguments) < min_args or len(self.arguments) > max_args:
            if max_args == min_args:
                msg = "%d" % min_args
            else:
                msg = "between %d and %d" % (min_args, max_args)
            raise RuleSyntaxError("%s takes %s arguments" % (
                self.RULE_IDENTIFIER, msg))

    def validate_tests_size(self, min_tests, max_tests=None):
        if max_tests is None:
            max_tests = min_tests
        if len(self.tests) < min_tests or len(self.tests) > max_tests:
            if max_tests == min_tests:
                msg = "%d" % min_tests
            else:
                msg = "between %d and %d" % (min_tests, max_tests)
            raise RuleSyntaxError("%s takes %s tests" % (
                self.RULE_IDENTIFIER, msg))

    def validate_arg_is_stringlist(self, index, length=None):
        if not (isinstance(self.arguments[index], list)
                and all(isinstance(arg, basestring)
                        for arg in self.arguments[index])):
            raise RuleSyntaxError(
                    "%s requires argument %d to be a string or list of strings"
                    % (self.RULE_IDENTIFIER, index)
                    )
        if length is not None and len(self.arguments[index]) != length:
            if length == 1:
                msg = "a single string or list of one string"
            else:
                msg = "a list of %d strings" % length
            raise RuleSyntaxError("%s requires argument %d to be %s" % (
                self.RULE_IDENTIFIER, index, msg))

    def validate_arg_is_tag(self, index, allowed_tags=None):
        if not isinstance(self.arguments[index], sieve.grammar.Tag):
            raise RuleSyntaxError("%s requires argument %d to be a tag" %
                    (self.RULE_IDENTIFIER, index))
        if (allowed_tags is not None
                and self.arguments[index] not in allowed_tags):
            raise SieveRuleSyntaxError(
                    "%s requires argument %d to be one of these tags: %s"
                    % (self.RULE_IDENTIFIER, index,
                        ', '.join([ ':'+tag for tag in allowed_tags]))
                    )

    def validate_arg_is_number(self, index):
        try:
            long(self.arguments[index])
        except TypeError:
            raise RuleSyntaxError("%s requires argument %d to be a number"
                    % (self.RULE_IDENTIFIER, index))

    def validate_arg_is_comparator(self, index):
        self.validate_arg_is_stringlist(index, 1)
        if not sieve.handler.get('comparator', self.arguments[index][0]):
            raise RuleSyntaxError(
                    "'%s' comparator is unknown/unsupported"
                    % self.arguments[index][0])

    def validate_arg_is_match_type(self, index):
        self.validate_arg_is_tag(index, ('IS', 'CONTAINS', 'MATCHES'))

    def validate_arg_is_address_part(self, index):
        self.validate_arg_is_tag(index, ('LOCALPART', 'DOMAIN', 'ALL'))

    def evaluate(self, message, state):
        raise NotImplementedError

<|endoftext|>


<commit_before>__version__ = '1.10.17'

class ValidationTimeout(Exception):
    """Validation has timed out.

    May be replaced by the exception type raised by an external timeout
    handler when run in a server environment."""

    def __init__(self, timeout):
        self.timeout = timeout

    def __str__(self):
        return 'Validation timeout after %d seconds' % self.timeout
<commit_msg>Tag 1.10.18<commit_after>__version__ = '1.10.18'

class ValidationTimeout(Exception):
    """Validation has timed out.

    May be replaced by the exception type raised by an external timeout
    handler when run in a server environment."""

    def __init__(self, timeout):
        self.timeout = timeout

    def __str__(self):
        return 'Validation timeout after %d seconds' % self.timeout
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""
RDFa extractor

Based on pyrdfa3 and rdflib
"""
import json
import logging
from xml.dom import Node
from xml.dom.minidom import Attr, NamedNodeMap

from lxml.etree import ElementBase, _ElementStringResult, _ElementUnicodeResult, XPath, tostring
from lxml.html import fromstring, HTMLParser, HtmlElementClassLookup
from rdflib import Graph, logger as rdflib_logger
from rdflib.plugins.parsers.pyRdfa import pyRdfa as PyRdfa, Options, logger as pyrdfa_logger
from rdflib.plugins.parsers.pyRdfa.initialcontext import initial_context

from copy import deepcopy, copy

# silence rdflib/PyRdfa INFO logs
rdflib_logger.setLevel(logging.ERROR)
pyrdfa_logger.setLevel(logging.ERROR)

initial_context["http://www.w3.org/2011/rdfa-context/rdfa-1.1"].ns.update({
    "twitter": "https://dev.twitter.com/cards#",
    "fb": "http://ogp.me/ns/fb#"
})

class DomElementUnicodeResult(object):
    CDATA_SECTION_NODE = Node.CDATA_SECTION_NODE
    ELEMENT_NODE = Node.ELEMENT_NODE
    TEXT_NODE = Node.TEXT_NODE

    def __init__(self, text):
        self.text = text
        self.nodeType = Node.TEXT_NODE

    @property
    def data(self):
        if isinstance(self.text, _ElementUnicodeResult):
            return self.text
        else:
            raise RuntimeError

class DomTextNode(object):
    CDATA_SECTION_NODE = Node.CDATA_SECTION_NODE
    ELEMENT_NODE = Node.ELEMENT_NODE
    TEXT_NODE = Node.TEXT_NODE

    def __init__(self, text):
        self.data = text
        self.nodeType = Node.TEXT_NODE

def lxmlDomNodeType(node):
    if isinstance(node, ElementBase):
        return Node.ELEMENT_NODE

    elif isinstance(node, (_ElementStringResult, _ElementUnicodeResult)):
        if node.is_attribute:
            return Node.ATTRIBUTE_NODE
        else:
            return Node.TEXT_NODE
    else:
        return Node.NOTATION_NODE

class DomHtmlMixin(object):
    CDATA_SECTION_NODE = Node.CDATA_SECTION_NODE
    ELEMENT_NODE = Node.ELEMENT_NODE
    TEXT_NODE = Node.TEXT_NODE

    _xp_childrennodes = XPath('child::node()')

    @property
    def documentElement(self):
        return self.getroottree().getroot()

    @property
    def nodeType(self):
        return Node.ELEMENT_NODE

    @property
    def nodeName(self):
        # FIXME: this is a simpification
        return self.tag

    @property
    def tagName(self):
        return self.tag

    @property
    def localName(self):
        return self.xpath('local-name(.)')

    def hasAttribute(self, name):
        return name in self.attrib

    def getAttribute(self, name):
        return self.get(name)

    def setAttribute(self, name, value):
        self.set(name, value)

    def cloneNode(self, deep):
        return deepcopy(self) if deep else copy(self)

    @property
    def attributes(self):
        attrs = {}
        for name, value in self.attrib.items():
            a = Attr(name)
            a.value = value
            attrs[name] = a
        return NamedNodeMap(attrs, {}, self)

    @property
    def parentNode(self):
        return self.getparent()

    @property
    def childNodes_xpath(self):
        for n in self._xp_childrennodes(self):

            if isinstance(n, ElementBase):
                yield n

            elif isinstance(n, (_ElementStringResult, _ElementUnicodeResult)):

                if isinstance(n, _ElementUnicodeResult):
                    n = DomElementUnicodeResult(n)
                else:
                    n.nodeType = Node.TEXT_NODE
                    n.data = n
                yield n

    @property
    def childNodes(self):
        if self.text:
            yield DomTextNode(self.text)
        for n in self.iterchildren():
            yield n
            if n.tail:
                yield DomTextNode(n.tail)

    def getElementsByTagName(self, name):
        return self.iterdescendants(name)

    def getElementById(self, i):
        return self.get_element_by_id(i)

    @property
    def data(self):
        if isinstance(self, (_ElementStringResult, _ElementUnicodeResult)):
            return self
        else:
            raise RuntimeError

    def toxml(self, encoding=None):
        return tostring(self, encoding=encoding)

class DomHtmlElementClassLookup(HtmlElementClassLookup):
    _lookups = {}
    def lookup(self, node_type, document, namespace, name):
        k = (node_type, document, namespace, name)
        t = self._lookups.get(k)
        if t is None:
            cur = super(DomHtmlElementClassLookup, self).lookup(node_type, document, namespace, name)
            newtype = type('Dom'+cur.__name__, (cur, DomHtmlMixin), {})
            self._lookups[k] = newtype
            return newtype
        else:
            return t

class XmlDomHTMLParser(HTMLParser):
    """An HTML parser that is configured to return XmlDomHtmlElement
    objects, compatible with xml.dom API
    """
    def __init__(self, **kwargs):
        super(HTMLParser, self).__init__(**kwargs)
        parser_lookup = DomHtmlElementClassLookup()
        self.set_element_class_lookup(parser_lookup)

class RDFaExtractor(object):

    def extract(self, htmlstring, url='http://www.example.com/', encoding="UTF-8",
            expanded=True):

        domparser = XmlDomHTMLParser()
        tree = fromstring(htmlstring.encode('utf-8'), parser=domparser)

        options = Options(output_processor_graph=True,
                          embedded_rdf=False,
                          space_preserve=True,
                          vocab_expansion=False,
                          vocab_cache=False,
                          vocab_cache_report=False,
                          refresh_vocab_cache=False,
                          check_lite=False)

        g = PyRdfa(options, base=url).graph_from_DOM(tree, graph=Graph(), pgraph=Graph())
        jsonld_string = g.serialize(format='json-ld', auto_compact=not expanded).decode('utf-8')
        return json.loads(jsonld_string)
<commit_msg>Add extract_items() to RDFa extractor<commit_after># -*- coding: utf-8 -*-
"""
RDFa extractor

Based on pyrdfa3 and rdflib
"""
import json
import logging
from xml.dom import Node
from xml.dom.minidom import Attr, NamedNodeMap

from lxml.etree import ElementBase, _ElementStringResult, _ElementUnicodeResult, XPath, tostring
from lxml.html import fromstring, HTMLParser, HtmlElementClassLookup
from rdflib import Graph, logger as rdflib_logger
from rdflib.plugins.parsers.pyRdfa import pyRdfa as PyRdfa, Options, logger as pyrdfa_logger
from rdflib.plugins.parsers.pyRdfa.initialcontext import initial_context

from copy import deepcopy, copy

# silence rdflib/PyRdfa INFO logs
rdflib_logger.setLevel(logging.ERROR)
pyrdfa_logger.setLevel(logging.ERROR)

initial_context["http://www.w3.org/2011/rdfa-context/rdfa-1.1"].ns.update({
    "twitter": "https://dev.twitter.com/cards#",
    "fb": "http://ogp.me/ns/fb#"
})

class DomElementUnicodeResult(object):
    CDATA_SECTION_NODE = Node.CDATA_SECTION_NODE
    ELEMENT_NODE = Node.ELEMENT_NODE
    TEXT_NODE = Node.TEXT_NODE

    def __init__(self, text):
        self.text = text
        self.nodeType = Node.TEXT_NODE

    @property
    def data(self):
        if isinstance(self.text, _ElementUnicodeResult):
            return self.text
        else:
            raise RuntimeError

class DomTextNode(object):
    CDATA_SECTION_NODE = Node.CDATA_SECTION_NODE
    ELEMENT_NODE = Node.ELEMENT_NODE
    TEXT_NODE = Node.TEXT_NODE

    def __init__(self, text):
        self.data = text
        self.nodeType = Node.TEXT_NODE

def lxmlDomNodeType(node):
    if isinstance(node, ElementBase):
        return Node.ELEMENT_NODE

    elif isinstance(node, (_ElementStringResult, _ElementUnicodeResult)):
        if node.is_attribute:
            return Node.ATTRIBUTE_NODE
        else:
            return Node.TEXT_NODE
    else:
        return Node.NOTATION_NODE

class DomHtmlMixin(object):
    CDATA_SECTION_NODE = Node.CDATA_SECTION_NODE
    ELEMENT_NODE = Node.ELEMENT_NODE
    TEXT_NODE = Node.TEXT_NODE

    _xp_childrennodes = XPath('child::node()')

    @property
    def documentElement(self):
        return self.getroottree().getroot()

    @property
    def nodeType(self):
        return Node.ELEMENT_NODE

    @property
    def nodeName(self):
        # FIXME: this is a simpification
        return self.tag

    @property
    def tagName(self):
        return self.tag

    @property
    def localName(self):
        return self.xpath('local-name(.)')

    def hasAttribute(self, name):
        return name in self.attrib

    def getAttribute(self, name):
        return self.get(name)

    def setAttribute(self, name, value):
        self.set(name, value)

    def cloneNode(self, deep):
        return deepcopy(self) if deep else copy(self)

    @property
    def attributes(self):
        attrs = {}
        for name, value in self.attrib.items():
            a = Attr(name)
            a.value = value
            attrs[name] = a
        return NamedNodeMap(attrs, {}, self)

    @property
    def parentNode(self):
        return self.getparent()

    @property
    def childNodes_xpath(self):
        for n in self._xp_childrennodes(self):

            if isinstance(n, ElementBase):
                yield n

            elif isinstance(n, (_ElementStringResult, _ElementUnicodeResult)):

                if isinstance(n, _ElementUnicodeResult):
                    n = DomElementUnicodeResult(n)
                else:
                    n.nodeType = Node.TEXT_NODE
                    n.data = n
                yield n

    @property
    def childNodes(self):
        if self.text:
            yield DomTextNode(self.text)
        for n in self.iterchildren():
            yield n
            if n.tail:
                yield DomTextNode(n.tail)

    def getElementsByTagName(self, name):
        return self.iterdescendants(name)

    def getElementById(self, i):
        return self.get_element_by_id(i)

    @property
    def data(self):
        if isinstance(self, (_ElementStringResult, _ElementUnicodeResult)):
            return self
        else:
            raise RuntimeError

    def toxml(self, encoding=None):
        return tostring(self, encoding=encoding)

class DomHtmlElementClassLookup(HtmlElementClassLookup):
    _lookups = {}
    def lookup(self, node_type, document, namespace, name):
        k = (node_type, document, namespace, name)
        t = self._lookups.get(k)
        if t is None:
            cur = super(DomHtmlElementClassLookup, self).lookup(node_type, document, namespace, name)
            newtype = type('Dom'+cur.__name__, (cur, DomHtmlMixin), {})
            self._lookups[k] = newtype
            return newtype
        else:
            return t

class XmlDomHTMLParser(HTMLParser):
    """An HTML parser that is configured to return XmlDomHtmlElement
    objects, compatible with xml.dom API
    """
    def __init__(self, **kwargs):
        super(HTMLParser, self).__init__(**kwargs)
        parser_lookup = DomHtmlElementClassLookup()
        self.set_element_class_lookup(parser_lookup)

class RDFaExtractor(object):

    def extract(self, htmlstring, url='http://www.example.com/', encoding="UTF-8",
            expanded=True):

        domparser = XmlDomHTMLParser()
        tree = fromstring(htmlstring.encode('utf-8'), parser=domparser)
        return self.extract_items(tree, url, expanded=expanded)

    def extract_items(self, document, url, expanded=True):
        options = Options(output_processor_graph=True,
                          embedded_rdf=False,
                          space_preserve=True,
                          vocab_expansion=False,
                          vocab_cache=False,
                          vocab_cache_report=False,
                          refresh_vocab_cache=False,
                          check_lite=False)

        g = PyRdfa(options, base=url).graph_from_DOM(document, graph=Graph(), pgraph=Graph())
        jsonld_string = g.serialize(format='json-ld', auto_compact=not expanded).decode('utf-8')
        return json.loads(jsonld_string)
<|endoftext|>


<commit_before>from parserator.manual_labeling import consoleLabel, naiveConsoleLabel
import name_parser
from lxml import etree
import sys
import os.path
import data_prep_utils
import re

if __name__ == '__main__' :

    import csv
    from argparse import ArgumentParser
    import unidecode

    labels = name_parser.config.LABELS

    parser = ArgumentParser(description="Label some strings")
    parser.add_argument(dest="infile",
                        help="input csv", metavar="FILE")
    parser.add_argument(dest="outfile",
                        help="output csv", metavar="FILE")
    parser.add_argument("-n",
                        help="-n for naive labeling (if there isn't an existing .crfsuite settings file)", action="store_true")
    args = parser.parse_args()

    file_slug = re.sub('(.*/)|(.csv)|(unlabeled_)', '', args.infile)

    # Check to make sure we can write to outfile
    if os.path.isfile(args.outfile):
        with open(args.outfile, 'r+' ) as f:
            try :
                tree = etree.parse(f)
            except :
                raise ValueError("%s does not seem to be a valid xml file"
                                 % args.outfile)

    with open(args.infile, 'rU') as infile :
        reader = csv.reader(infile)

        strings = set([unidecode.unidecode(row[0]) for row in reader])

    p = name_parser.Parser()

    if args.n :
        labeled_list, raw_strings_left = naiveConsoleLabel(strings, labels, p)
    else:
        labeled_list, raw_strings_left = consoleLabel(strings, labels, p)

    data_prep_utils.appendListToXMLfile(labeled_list, args.outfile)
    data_prep_utils.list2file(raw_strings_left, 'training/data_prep/unlabeled_data/unlabeled_'+file_slug+'.csv')

<commit_msg>remove more stuff from manual labeling<commit_after>from parserator.manual_labeling import getArgumentParser, label
import name_parser
from lxml import etree
import sys
import os.path
import data_prep_utils
import re

if __name__ == '__main__' :

    arg_parser = getArgumentParser()
    args = arg_parser.parse_args()
    p = name_parser.Parser()

    label(p, args, 'training/data_prep/unlabeled_data/')<|endoftext|>