<commit_before>import logging

from django.contrib.contenttypes.models import ContentType
from django.core.exceptions import ObjectDoesNotExist
from django.db import transaction

from waldur_core.structure import models as structure_models
from waldur_mastermind.marketplace import callbacks
from waldur_mastermind.marketplace import models as marketplace_models
from waldur_mastermind.packages import models as package_models
from waldur_openstack.openstack import models as openstack_models
from waldur_openstack.openstack.apps import OpenStackConfig
from waldur_openstack.openstack_tenant import apps as openstack_tenant_apps
from waldur_openstack.openstack_tenant import models as openstack_tenant_models

from . import INSTANCE_TYPE, PACKAGE_TYPE, VOLUME_TYPE, utils

logger = logging.getLogger(__name__)

def create_template_for_plan(sender, instance, created=False, **kwargs):
    plan = instance

    if plan.scope:
        return

    if not created:
        return

    if plan.offering.type != PACKAGE_TYPE:
        return

    if not isinstance(plan.offering.scope, structure_models.ServiceSettings):
        logger.warning('Skipping plan synchronization because offering scope is not service settings. '
                       'Plan ID: %s', plan.id)
        return

    if plan.offering.scope.type != OpenStackConfig.service_name:
        logger.warning('Skipping plan synchronization because service settings type is not OpenStack. '
                       'Plan ID: %s', plan.id)
        return

    with transaction.atomic():
        template = package_models.PackageTemplate.objects.create(
            service_settings=plan.offering.scope,
            name=plan.name,
            description=plan.description,
            product_code=plan.product_code,
            article_code=plan.article_code,
        )
        plan.scope = template
        plan.save()

def synchronize_plan_component(sender, instance, created=False, **kwargs):
    component = instance

    if not created:
        return

    if component.plan.offering.type != PACKAGE_TYPE:
        return

    template = component.plan.scope
    if not template:
        logger.warning('Skipping plan component synchronization because offering does not have scope. '
                       'Offering ID: %s', component.plan.offering.id)
        return

    if not package_models.PackageComponent.objects.filter(
            template=template, type=component.component.type).exists():
        package_models.PackageComponent.objects.create(
            template=template,
            type=component.component.type,
            amount=component.amount,
            price=component.price,
        )

def create_offering_from_tenant(sender, instance, created=False, **kwargs):
    if created:
        return

    if not instance.tracker.has_changed('state'):
        return

    if instance.tracker.previous('state') != instance.States.CREATING:
        return

    if instance.state != instance.States.OK:
        return

    tenant = instance

    try:
        resource = marketplace_models.Resource.objects.get(scope=tenant)
    except ObjectDoesNotExist:
        logger.debug('Skipping offering creation for tenant because order '
                     'item does not exist. OpenStack tenant ID: %s', tenant.id)
        return

    try:
        service_settings = structure_models.ServiceSettings.objects.get(
            scope=tenant,
            type=openstack_tenant_apps.OpenStackTenantConfig.service_name,
        )
    except ObjectDoesNotExist:
        logger.debug('Skipping offering creation for tenant because service settings '
                     'object does not exist. OpenStack tenant ID: %s', tenant.id)
        return

    parent_offering = resource.offering
    for offering_type in (INSTANCE_TYPE, VOLUME_TYPE):
        try:
            category, offering_name = utils.get_category_and_name_for_offering_type(
                offering_type, service_settings)
        except ObjectDoesNotExist:
            logger.warning('Skipping offering creation for tenant because category '
                           'for instances and volumes is not yet defined.')
            continue
        payload = dict(
            type=offering_type,
            name=offering_name,
            scope=service_settings,
            shared=False,
            category=category,
        )

        fields = (
            'state',
            'customer',
            'attributes',
            'thumbnail',
            'vendor_details',
            'geolocations',
        )
        for field in fields:
            payload[field] = getattr(parent_offering, field)

        with transaction.atomic():
            offering = marketplace_models.Offering.objects.create(**payload)
            offering.allowed_customers.add(tenant.service_project_link.project.customer)

def archive_offering(sender, instance, **kwargs):
    service_settings = instance

    if service_settings.type == openstack_tenant_apps.OpenStackTenantConfig.service_name and \
            service_settings.content_type == ContentType.objects.get_for_model(openstack_models.Tenant):
        marketplace_models.Offering.objects.filter(scope=service_settings).update(
            state=marketplace_models.Offering.States.ARCHIVED
        )

def change_order_item_state(sender, instance, created=False, **kwargs):
    if created or not instance.tracker.has_changed('state'):
        return

    try:
        resource = marketplace_models.Resource.objects.get(scope=instance)
    except ObjectDoesNotExist:
        logger.warning('Skipping OpenStack resource state synchronization '
                       'because marketplace resource is not found. '
                       'Resource ID: %s', instance.id)
    else:
        callbacks.sync_resource_state(instance, resource)

def terminate_resource(sender, instance, **kwargs):
    try:
        resource = marketplace_models.Resource.objects.get(scope=instance)
    except ObjectDoesNotExist:
        logger.debug('Skipping resource terminate for OpenStack resource '
                     'because marketplace resource does not exist. '
                     'Resource ID: %s', instance.id)
    else:
        callbacks.resource_deletion_succeeded(resource)

def synchronize_volume_metadata(sender, instance, created=False, **kwargs):
    volume = instance
    if not created and not set(volume.tracker.changed()) & {'size', 'instance_id'}:
        return

    try:
        resource = marketplace_models.Resource.objects.get(scope=volume)
    except ObjectDoesNotExist:
        logger.debug('Skipping resource synchronization for OpenStack volume '
                     'because marketplace resource does not exist. '
                     'Resource ID: %s', instance.id)
        return

    utils.import_volume_metadata(resource)

def synchronize_instance_metadata(sender, instance, created=False, **kwargs):
    if not created and not set(instance.tracker.changed()) & {'name'}:
        return

    for volume in instance.volumes.all():
        try:
            resource = marketplace_models.Resource.objects.get(scope=volume)
        except ObjectDoesNotExist:
            logger.debug('Skipping resource synchronization for OpenStack volume '
                         'because marketplace resource does not exist. '
                         'Resource ID: %s', instance.id)
            continue

        resource.backend_metadata['instance_name'] = volume.instance.name
        resource.save(update_fields=['backend_metadata'])

def synchronize_internal_ips(sender, instance, created=False, **kwargs):
    internal_ip = instance
    if not created and not set(internal_ip.tracker.changed()) & {'ip4_address', 'instance_id'}:
        return

    vms = {vm for vm in (internal_ip.instance_id, internal_ip.tracker.previous('instance_id')) if vm}

    for vm in vms:
        try:
            scope = openstack_tenant_models.Instance.objects.get(id=vm)
            resource = marketplace_models.Resource.objects.get(scope=scope)
        except ObjectDoesNotExist:
            logger.debug('Skipping resource synchronization for OpenStack instance '
                         'because marketplace resource does not exist. '
                         'Resource ID: %s', vm)
            return

        utils.import_instance_metadata(resource)

def synchronize_floating_ips(sender, instance, created=False, **kwargs):
    floating_ip = instance
    if not created and not set(instance.tracker.changed()) & {'address', 'internal_ip_id'}:
        return

    internal_ips = {ip for ip in (floating_ip.internal_ip_id, floating_ip.tracker.previous('internal_ip_id')) if ip}
    for ip_id in internal_ips:
        try:
            scope = openstack_tenant_models.Instance.objects.get(internal_ips_set__id=ip_id)
            resource = marketplace_models.Resource.objects.get(scope=scope)
        except ObjectDoesNotExist:
            logger.debug('Skipping resource synchronization for OpenStack instance '
                         'because marketplace resource does not exist. '
                         'Resource ID: %s', ip_id)
            return

        utils.import_instance_metadata(resource)

def synchronize_resource_metadata(sender, instance, created=False, **kwargs):
    fields = {'action', 'action_details', 'state', 'runtime_state'}
    if not created and not set(instance.tracker.changed()) & fields:
        return

    try:
        resource = marketplace_models.Resource.objects.get(scope=instance)
    except ObjectDoesNotExist:
        logger.debug('Skipping resource synchronization for OpenStack resource '
                     'because marketplace resource does not exist. '
                     'Resource ID: %s', instance.id)
        return

    utils.import_resource_metadata(resource)
<commit_msg>Allow to continue resource processing even if previous item has been skipped.<commit_after>import logging

from django.contrib.contenttypes.models import ContentType
from django.core.exceptions import ObjectDoesNotExist
from django.db import transaction

from waldur_core.structure import models as structure_models
from waldur_mastermind.marketplace import callbacks
from waldur_mastermind.marketplace import models as marketplace_models
from waldur_mastermind.packages import models as package_models
from waldur_openstack.openstack import models as openstack_models
from waldur_openstack.openstack.apps import OpenStackConfig
from waldur_openstack.openstack_tenant import apps as openstack_tenant_apps
from waldur_openstack.openstack_tenant import models as openstack_tenant_models

from . import INSTANCE_TYPE, PACKAGE_TYPE, VOLUME_TYPE, utils

logger = logging.getLogger(__name__)

def create_template_for_plan(sender, instance, created=False, **kwargs):
    plan = instance

    if plan.scope:
        return

    if not created:
        return

    if plan.offering.type != PACKAGE_TYPE:
        return

    if not isinstance(plan.offering.scope, structure_models.ServiceSettings):
        logger.warning('Skipping plan synchronization because offering scope is not service settings. '
                       'Plan ID: %s', plan.id)
        return

    if plan.offering.scope.type != OpenStackConfig.service_name:
        logger.warning('Skipping plan synchronization because service settings type is not OpenStack. '
                       'Plan ID: %s', plan.id)
        return

    with transaction.atomic():
        template = package_models.PackageTemplate.objects.create(
            service_settings=plan.offering.scope,
            name=plan.name,
            description=plan.description,
            product_code=plan.product_code,
            article_code=plan.article_code,
        )
        plan.scope = template
        plan.save()

def synchronize_plan_component(sender, instance, created=False, **kwargs):
    component = instance

    if not created:
        return

    if component.plan.offering.type != PACKAGE_TYPE:
        return

    template = component.plan.scope
    if not template:
        logger.warning('Skipping plan component synchronization because offering does not have scope. '
                       'Offering ID: %s', component.plan.offering.id)
        return

    if not package_models.PackageComponent.objects.filter(
            template=template, type=component.component.type).exists():
        package_models.PackageComponent.objects.create(
            template=template,
            type=component.component.type,
            amount=component.amount,
            price=component.price,
        )

def create_offering_from_tenant(sender, instance, created=False, **kwargs):
    if created:
        return

    if not instance.tracker.has_changed('state'):
        return

    if instance.tracker.previous('state') != instance.States.CREATING:
        return

    if instance.state != instance.States.OK:
        return

    tenant = instance

    try:
        resource = marketplace_models.Resource.objects.get(scope=tenant)
    except ObjectDoesNotExist:
        logger.debug('Skipping offering creation for tenant because order '
                     'item does not exist. OpenStack tenant ID: %s', tenant.id)
        return

    try:
        service_settings = structure_models.ServiceSettings.objects.get(
            scope=tenant,
            type=openstack_tenant_apps.OpenStackTenantConfig.service_name,
        )
    except ObjectDoesNotExist:
        logger.debug('Skipping offering creation for tenant because service settings '
                     'object does not exist. OpenStack tenant ID: %s', tenant.id)
        return

    parent_offering = resource.offering
    for offering_type in (INSTANCE_TYPE, VOLUME_TYPE):
        try:
            category, offering_name = utils.get_category_and_name_for_offering_type(
                offering_type, service_settings)
        except ObjectDoesNotExist:
            logger.warning('Skipping offering creation for tenant because category '
                           'for instances and volumes is not yet defined.')
            continue
        payload = dict(
            type=offering_type,
            name=offering_name,
            scope=service_settings,
            shared=False,
            category=category,
        )

        fields = (
            'state',
            'customer',
            'attributes',
            'thumbnail',
            'vendor_details',
            'geolocations',
        )
        for field in fields:
            payload[field] = getattr(parent_offering, field)

        with transaction.atomic():
            offering = marketplace_models.Offering.objects.create(**payload)
            offering.allowed_customers.add(tenant.service_project_link.project.customer)

def archive_offering(sender, instance, **kwargs):
    service_settings = instance

    if service_settings.type == openstack_tenant_apps.OpenStackTenantConfig.service_name and \
            service_settings.content_type == ContentType.objects.get_for_model(openstack_models.Tenant):
        marketplace_models.Offering.objects.filter(scope=service_settings).update(
            state=marketplace_models.Offering.States.ARCHIVED
        )

def change_order_item_state(sender, instance, created=False, **kwargs):
    if created or not instance.tracker.has_changed('state'):
        return

    try:
        resource = marketplace_models.Resource.objects.get(scope=instance)
    except ObjectDoesNotExist:
        logger.warning('Skipping OpenStack resource state synchronization '
                       'because marketplace resource is not found. '
                       'Resource ID: %s', instance.id)
    else:
        callbacks.sync_resource_state(instance, resource)

def terminate_resource(sender, instance, **kwargs):
    try:
        resource = marketplace_models.Resource.objects.get(scope=instance)
    except ObjectDoesNotExist:
        logger.debug('Skipping resource terminate for OpenStack resource '
                     'because marketplace resource does not exist. '
                     'Resource ID: %s', instance.id)
    else:
        callbacks.resource_deletion_succeeded(resource)

def synchronize_volume_metadata(sender, instance, created=False, **kwargs):
    volume = instance
    if not created and not set(volume.tracker.changed()) & {'size', 'instance_id'}:
        return

    try:
        resource = marketplace_models.Resource.objects.get(scope=volume)
    except ObjectDoesNotExist:
        logger.debug('Skipping resource synchronization for OpenStack volume '
                     'because marketplace resource does not exist. '
                     'Resource ID: %s', instance.id)
        return

    utils.import_volume_metadata(resource)

def synchronize_instance_metadata(sender, instance, created=False, **kwargs):
    if not created and not set(instance.tracker.changed()) & {'name'}:
        return

    for volume in instance.volumes.all():
        try:
            resource = marketplace_models.Resource.objects.get(scope=volume)
        except ObjectDoesNotExist:
            logger.debug('Skipping resource synchronization for OpenStack volume '
                         'because marketplace resource does not exist. '
                         'Resource ID: %s', instance.id)
            continue

        resource.backend_metadata['instance_name'] = volume.instance.name
        resource.save(update_fields=['backend_metadata'])

def synchronize_internal_ips(sender, instance, created=False, **kwargs):
    internal_ip = instance
    if not created and not set(internal_ip.tracker.changed()) & {'ip4_address', 'instance_id'}:
        return

    vms = {vm for vm in (internal_ip.instance_id, internal_ip.tracker.previous('instance_id')) if vm}

    for vm in vms:
        try:
            scope = openstack_tenant_models.Instance.objects.get(id=vm)
            resource = marketplace_models.Resource.objects.get(scope=scope)
        except ObjectDoesNotExist:
            logger.debug('Skipping resource synchronization for OpenStack instance '
                         'because marketplace resource does not exist. '
                         'Resource ID: %s', vm)
            continue

        utils.import_instance_metadata(resource)

def synchronize_floating_ips(sender, instance, created=False, **kwargs):
    floating_ip = instance
    if not created and not set(instance.tracker.changed()) & {'address', 'internal_ip_id'}:
        return

    internal_ips = {ip for ip in (floating_ip.internal_ip_id, floating_ip.tracker.previous('internal_ip_id')) if ip}
    for ip_id in internal_ips:
        try:
            scope = openstack_tenant_models.Instance.objects.get(internal_ips_set__id=ip_id)
            resource = marketplace_models.Resource.objects.get(scope=scope)
        except ObjectDoesNotExist:
            logger.debug('Skipping resource synchronization for OpenStack instance '
                         'because marketplace resource does not exist. '
                         'Resource ID: %s', ip_id)
            continue

        utils.import_instance_metadata(resource)

def synchronize_resource_metadata(sender, instance, created=False, **kwargs):
    fields = {'action', 'action_details', 'state', 'runtime_state'}
    if not created and not set(instance.tracker.changed()) & fields:
        return

    try:
        resource = marketplace_models.Resource.objects.get(scope=instance)
    except ObjectDoesNotExist:
        logger.debug('Skipping resource synchronization for OpenStack resource '
                     'because marketplace resource does not exist. '
                     'Resource ID: %s', instance.id)
        return

    utils.import_resource_metadata(resource)
<|endoftext|>


<commit_before>"""Implicit plotting module for Diofant

The module implements a data series called ImplicitSeries which is used by
``Plot`` class to plot implicit plots for different backends.

See Also
========

diofant.plotting.plot
"""

from ..core import Dummy, Eq, Symbol, Tuple, sympify
from ..core.relational import (Equality, GreaterThan, LessThan, Relational,
                               StrictGreaterThan, StrictLessThan)
from ..external import import_module
from ..logic.boolalg import BooleanFunction
from ..polys.polyutils import _sort_gens
from ..utilities import flatten, lambdify
from ..utilities.decorator import doctest_depends_on
from .plot import BaseSeries, Plot

class ImplicitSeries(BaseSeries):
    """ Representation for Implicit plot """

    is_implicit = True

    def __init__(self, expr, var_start_end_x, var_start_end_y,
                 has_equality, use_interval_math, depth, nb_of_points,
                 line_color):
        super(ImplicitSeries, self).__init__()
        self.expr = sympify(expr)
        self.var_x = sympify(var_start_end_x[0])
        self.start_x = float(var_start_end_x[1])
        self.end_x = float(var_start_end_x[2])
        self.var_y = sympify(var_start_end_y[0])
        self.start_y = float(var_start_end_y[1])
        self.end_y = float(var_start_end_y[2])
        self.get_points = self.get_raster

        # If the expression has equality, i.e. Eq, Greaterthan, LessThan.
        self.has_equality = has_equality

        self.nb_of_points = nb_of_points
        self.use_interval_math = use_interval_math
        self.depth = 4 + depth
        self.line_color = line_color

    def get_raster(self):
        return self._get_meshes_grid()

    def _get_meshes_grid(self):
        """Generates the mesh for generating a contour.

        In the case of equality, ``contour`` function of matplotlib can
        be used. In other cases, matplotlib's ``contourf`` is used.
        """
        equal = False
        if isinstance(self.expr, Equality):
            expr = self.expr.lhs - self.expr.rhs
            equal = True

        elif isinstance(self.expr, (GreaterThan, StrictGreaterThan)):
            expr = self.expr.lhs - self.expr.rhs

        elif isinstance(self.expr, (LessThan, StrictLessThan)):
            expr = self.expr.rhs - self.expr.lhs
        else:
            raise NotImplementedError("The expression is not supported for "
                                      "plotting in uniform meshed plot.")
        np = import_module('numpy')
        xarray = np.linspace(self.start_x, self.end_x, self.nb_of_points)
        yarray = np.linspace(self.start_y, self.end_y, self.nb_of_points)
        x_grid, y_grid = np.meshgrid(xarray, yarray)

        func = lambdify((self.var_x, self.var_y), expr, "numpy")
        z_grid = func(x_grid, y_grid)
        z_grid[np.ma.where(z_grid < 0)] = -1
        z_grid[np.ma.where(z_grid > 0)] = 1
        if equal:
            return xarray, yarray, z_grid, 'contour'
        else:
            return xarray, yarray, z_grid, 'contourf'

@doctest_depends_on(modules=('matplotlib',))
def plot_implicit(expr, x_var=None, y_var=None, **kwargs):
    """A plot function to plot implicit equations / inequalities.

    Parameters
    ==========

    ``expr`` : Expr
        The equation / inequality that is to be plotted.
    ``x_var`` : symbol or tuple, optional
        symbol to plot on x-axis or tuple giving symbol and range
        as ``(symbol, xmin, xmax)``
    ``y_var`` : symbol or tuple, optional
        symbol to plot on y-axis or tuple giving symbol and range
        as ``(symbol, ymin, ymax)``

    If neither ``x_var`` nor ``y_var`` are given then the free symbols in the
    expression will be assigned in the order they are sorted.

    The following keyword arguments can also be used:

    ``adaptive`` : Boolean, optional
        The default value is set to True. It has to be
        set to False if you want to use a mesh grid.

    ``depth`` : integer
        The depth of recursion for adaptive mesh grid.
        Default value is 0. Takes value in the range (0, 4).

    ``points`` : integer
        The number of points if adaptive mesh grid is not
        used. Default value is 200.

    ``title`` : str
        The title for the plot.

    ``xlabel`` : str
        The label for the x-axis

    ``ylabel`` : string
        The label for the y-axis

    Aesthetics options:

    ``line_color`` : float or str
        Specifies the color for the plot.  See ``Plot`` to see how to
        set color for the plots.

    plot_implicit, by default, uses interval arithmetic to plot functions. If
    the expression cannot be plotted using interval arithmetic, it defaults to
    a generating a contour using a mesh grid of fixed number of points. By
    setting adaptive to False, you can force plot_implicit to use the mesh
    grid. The mesh grid method can be effective when adaptive plotting using
    interval arithmetic, fails to plot with small line width.

    Examples
    ========

    Plot expressions:

    >>> from diofant import plot_implicit, cos, sin, symbols, Eq, And
    >>> x, y = symbols('x y')

    Without any ranges for the symbols in the expression

    >>> p1 = plot_implicit(Eq(x**2 + y**2, 5))

    With the range for the symbols

    >>> p2 = plot_implicit(Eq(x**2 + y**2, 3),
    ...         (x, -3, 3), (y, -3, 3))

    With depth of recursion as argument.

    >>> p3 = plot_implicit(Eq(x**2 + y**2, 5),
    ...         (x, -4, 4), (y, -4, 4), depth = 2)

    Using mesh grid and not using adaptive meshing.

    >>> p4 = plot_implicit(Eq(x**2 + y**2, 5),
    ...         (x, -5, 5), (y, -2, 2), adaptive=False)

    Using mesh grid with number of points as input.

    >>> p5 = plot_implicit(Eq(x**2 + y**2, 5),
    ...         (x, -5, 5), (y, -2, 2),
    ...         adaptive=False, points=400)

    Plotting regions.

    >>> p6 = plot_implicit(y > x**2)

    When plotting an expression with a single variable (y - 1, for example),
    specify the x or the y variable explicitly:

    >>> p8 = plot_implicit(y - 1, y_var=y)
    >>> p9 = plot_implicit(x - 1, x_var=x)
    """

    # Represents whether the expression contains an Equality,
    # GreaterThan or LessThan
    has_equality = False

    def arg_expand(bool_expr):
        """
        Recursively expands the arguments of an Boolean Function
        """
        for arg in bool_expr.args:
            if isinstance(arg, BooleanFunction):
                arg_expand(arg)
            elif isinstance(arg, Relational):
                arg_list.append(arg)

    arg_list = []
    if isinstance(expr, BooleanFunction):
        arg_expand(expr)

    # Check whether there is an equality in the expression provided.
        if any(isinstance(e, (Equality, GreaterThan, LessThan))
               for e in arg_list):
            has_equality = True

    elif not isinstance(expr, Relational):
        expr = Eq(expr, 0)
        has_equality = True
    elif isinstance(expr, (Equality, GreaterThan, LessThan)):
        has_equality = True

    xyvar = [i for i in (x_var, y_var) if i is not None]
    free_symbols = expr.free_symbols
    range_symbols = Tuple(*flatten(xyvar)).free_symbols
    undeclared = free_symbols - range_symbols
    if len(free_symbols & range_symbols) > 2:
        raise NotImplementedError("Implicit plotting is not implemented for "
                                  "more than 2 variables")

    # Create default ranges if the range is not provided.
    default_range = Tuple(-5, 5)

    def _range_tuple(s):
        if isinstance(s, (Dummy, Symbol)):
            return Tuple(s) + default_range
        if len(s) == 3:
            return Tuple(*s)
        raise ValueError('symbol or `(symbol, min, max)` expected but got %s' % s)

    if len(xyvar) == 0:
        xyvar = list(_sort_gens(free_symbols))
    var_start_end_x = _range_tuple(xyvar[0])
    x = var_start_end_x[0]
    if len(xyvar) != 2:
        if x in undeclared or not undeclared:
            xyvar.append(Dummy('f(%s)' % x.name))
        else:
            xyvar.append(undeclared.pop())
    var_start_end_y = _range_tuple(xyvar[1])

    use_interval = kwargs.pop('adaptive', False)
    nb_of_points = kwargs.pop('points', 300)
    depth = kwargs.pop('depth', 0)
    line_color = kwargs.pop('line_color', "blue")
    # Check whether the depth is greater than 4 or less than 0.
    if depth > 4:
        depth = 4
    elif depth < 0:
        depth = 0

    series_argument = ImplicitSeries(expr, var_start_end_x, var_start_end_y,
                                     has_equality, use_interval, depth,
                                     nb_of_points, line_color)
    show = kwargs.pop('show', True)

    # set the x and y limits
    kwargs['xlim'] = tuple(float(x) for x in var_start_end_x[1:])
    kwargs['ylim'] = tuple(float(y) for y in var_start_end_y[1:])
    # set the x and y labels
    kwargs.setdefault('xlabel', var_start_end_x[0].name)
    kwargs.setdefault('ylabel', var_start_end_y[0].name)
    p = Plot(series_argument, **kwargs)
    if show:
        p.show()
    return p
<commit_msg>Uniform meshed plots should allow everything, except Eq's<commit_after>"""Implicit plotting module for Diofant

The module implements a data series called ImplicitSeries which is used by
``Plot`` class to plot implicit plots for different backends.

See Also
========

diofant.plotting.plot
"""

from ..core import Dummy, Eq, Symbol, Tuple, sympify
from ..core.relational import Equality, GreaterThan, LessThan, Relational
from ..external import import_module
from ..logic.boolalg import BooleanFunction
from ..polys.polyutils import _sort_gens
from ..utilities import flatten, lambdify
from ..utilities.decorator import doctest_depends_on
from .plot import BaseSeries, Plot

class ImplicitSeries(BaseSeries):
    """ Representation for Implicit plot """

    is_implicit = True

    def __init__(self, expr, var_start_end_x, var_start_end_y,
                 has_equality, use_interval_math, depth, nb_of_points,
                 line_color):
        super(ImplicitSeries, self).__init__()
        self.expr = sympify(expr)
        self.var_x = sympify(var_start_end_x[0])
        self.start_x = float(var_start_end_x[1])
        self.end_x = float(var_start_end_x[2])
        self.var_y = sympify(var_start_end_y[0])
        self.start_y = float(var_start_end_y[1])
        self.end_y = float(var_start_end_y[2])
        self.get_points = self.get_raster

        # If the expression has equality, i.e. Eq, Greaterthan, LessThan.
        self.has_equality = has_equality

        self.nb_of_points = nb_of_points
        self.use_interval_math = use_interval_math
        self.depth = 4 + depth
        self.line_color = line_color

    def get_raster(self):
        return self._get_meshes_grid()

    def _get_meshes_grid(self):
        """Generates the mesh for generating a contour.

        In the case of equality, ``contour`` function of matplotlib can
        be used. In other cases, matplotlib's ``contourf`` is used.
        """
        equal = False
        if isinstance(self.expr, Equality):
            expr = self.expr.lhs - self.expr.rhs
            equal = True
        elif self.expr.has(Equality):  # pragma: no cover
            raise NotImplementedError("The expression is not supported for "
                                      "plotting in uniform meshed plot.")
        else:
            expr = self.expr
        np = import_module('numpy')
        xarray = np.linspace(self.start_x, self.end_x, self.nb_of_points)
        yarray = np.linspace(self.start_y, self.end_y, self.nb_of_points)
        x_grid, y_grid = np.meshgrid(xarray, yarray)

        func = lambdify((self.var_x, self.var_y), expr, "numpy")
        z_grid = func(x_grid, y_grid)
        z_grid[np.ma.where(z_grid < 0)] = -1
        z_grid[np.ma.where(z_grid > 0)] = 1
        if equal:
            return xarray, yarray, z_grid, 'contour'
        else:
            return xarray, yarray, z_grid, 'contourf'

@doctest_depends_on(modules=('matplotlib',))
def plot_implicit(expr, x_var=None, y_var=None, **kwargs):
    """A plot function to plot implicit equations / inequalities.

    Parameters
    ==========

    ``expr`` : Expr
        The equation / inequality that is to be plotted.
    ``x_var`` : symbol or tuple, optional
        symbol to plot on x-axis or tuple giving symbol and range
        as ``(symbol, xmin, xmax)``
    ``y_var`` : symbol or tuple, optional
        symbol to plot on y-axis or tuple giving symbol and range
        as ``(symbol, ymin, ymax)``

    If neither ``x_var`` nor ``y_var`` are given then the free symbols in the
    expression will be assigned in the order they are sorted.

    The following keyword arguments can also be used:

    ``adaptive`` : Boolean, optional
        The default value is set to True. It has to be
        set to False if you want to use a mesh grid.

    ``depth`` : integer
        The depth of recursion for adaptive mesh grid.
        Default value is 0. Takes value in the range (0, 4).

    ``points`` : integer
        The number of points if adaptive mesh grid is not
        used. Default value is 200.

    ``title`` : str
        The title for the plot.

    ``xlabel`` : str
        The label for the x-axis

    ``ylabel`` : string
        The label for the y-axis

    Aesthetics options:

    ``line_color`` : float or str
        Specifies the color for the plot.  See ``Plot`` to see how to
        set color for the plots.

    plot_implicit, by default, uses interval arithmetic to plot functions. If
    the expression cannot be plotted using interval arithmetic, it defaults to
    a generating a contour using a mesh grid of fixed number of points. By
    setting adaptive to False, you can force plot_implicit to use the mesh
    grid. The mesh grid method can be effective when adaptive plotting using
    interval arithmetic, fails to plot with small line width.

    Examples
    ========

    Plot expressions:

    >>> from diofant import plot_implicit, cos, sin, symbols, Eq, And
    >>> x, y = symbols('x y')

    Without any ranges for the symbols in the expression

    >>> p1 = plot_implicit(Eq(x**2 + y**2, 5))

    With the range for the symbols

    >>> p2 = plot_implicit(Eq(x**2 + y**2, 3),
    ...         (x, -3, 3), (y, -3, 3))

    With depth of recursion as argument.

    >>> p3 = plot_implicit(Eq(x**2 + y**2, 5),
    ...         (x, -4, 4), (y, -4, 4), depth = 2)

    Using mesh grid and not using adaptive meshing.

    >>> p4 = plot_implicit(Eq(x**2 + y**2, 5),
    ...         (x, -5, 5), (y, -2, 2), adaptive=False)

    Using mesh grid with number of points as input.

    >>> p5 = plot_implicit(Eq(x**2 + y**2, 5),
    ...         (x, -5, 5), (y, -2, 2),
    ...         adaptive=False, points=400)

    Plotting regions.

    >>> p6 = plot_implicit(y > x**2)

    Plotting Using boolean conjunctions.

    >>> p7 = plot_implicit(And(y > x, y > -x))

    When plotting an expression with a single variable (y - 1, for example),
    specify the x or the y variable explicitly:

    >>> p8 = plot_implicit(y - 1, y_var=y)
    >>> p9 = plot_implicit(x - 1, x_var=x)
    """

    # Represents whether the expression contains an Equality,
    # GreaterThan or LessThan
    has_equality = False

    def arg_expand(bool_expr):
        """
        Recursively expands the arguments of an Boolean Function
        """
        for arg in bool_expr.args:
            if isinstance(arg, BooleanFunction):
                arg_expand(arg)
            elif isinstance(arg, Relational):
                arg_list.append(arg)

    arg_list = []
    if isinstance(expr, BooleanFunction):
        arg_expand(expr)

    # Check whether there is an equality in the expression provided.
        if any(isinstance(e, (Equality, GreaterThan, LessThan))
               for e in arg_list):
            has_equality = True

    elif not isinstance(expr, Relational):
        expr = Eq(expr, 0)
        has_equality = True
    elif isinstance(expr, (Equality, GreaterThan, LessThan)):
        has_equality = True

    xyvar = [i for i in (x_var, y_var) if i is not None]
    free_symbols = expr.free_symbols
    range_symbols = Tuple(*flatten(xyvar)).free_symbols
    undeclared = free_symbols - range_symbols
    if len(free_symbols & range_symbols) > 2:
        raise NotImplementedError("Implicit plotting is not implemented for "
                                  "more than 2 variables")

    # Create default ranges if the range is not provided.
    default_range = Tuple(-5, 5)

    def _range_tuple(s):
        if isinstance(s, (Dummy, Symbol)):
            return Tuple(s) + default_range
        if len(s) == 3:
            return Tuple(*s)
        raise ValueError('symbol or `(symbol, min, max)` expected but got %s' % s)

    if len(xyvar) == 0:
        xyvar = list(_sort_gens(free_symbols))
    var_start_end_x = _range_tuple(xyvar[0])
    x = var_start_end_x[0]
    if len(xyvar) != 2:
        if x in undeclared or not undeclared:
            xyvar.append(Dummy('f(%s)' % x.name))
        else:
            xyvar.append(undeclared.pop())
    var_start_end_y = _range_tuple(xyvar[1])

    use_interval = kwargs.pop('adaptive', False)
    nb_of_points = kwargs.pop('points', 300)
    depth = kwargs.pop('depth', 0)
    line_color = kwargs.pop('line_color', "blue")
    # Check whether the depth is greater than 4 or less than 0.
    if depth > 4:
        depth = 4
    elif depth < 0:
        depth = 0

    series_argument = ImplicitSeries(expr, var_start_end_x, var_start_end_y,
                                     has_equality, use_interval, depth,
                                     nb_of_points, line_color)
    show = kwargs.pop('show', True)

    # set the x and y limits
    kwargs['xlim'] = tuple(float(x) for x in var_start_end_x[1:])
    kwargs['ylim'] = tuple(float(y) for y in var_start_end_y[1:])
    # set the x and y labels
    kwargs.setdefault('xlabel', var_start_end_x[0].name)
    kwargs.setdefault('ylabel', var_start_end_y[0].name)
    p = Plot(series_argument, **kwargs)
    if show:
        p.show()
    return p
<|endoftext|>


<commit_before>from __future__ import division, print_function, absolute_import

from ...utils.six.moves import xrange

import numpy as np
import nose
from dipy.io.bvectxt import orientation_from_string
from dipy.tracking.utils import (_rmi, connectivity_matrix, density_map,
                                 move_streamlines, ndbincount, reduce_labels,
                                 reorder_voxels_affine, streamline_mapping)
from numpy.testing import assert_array_almost_equal, assert_array_equal
from nose.tools import assert_equal, assert_raises, assert_true

def make_streamlines():
    streamlines = [ np.array([[0, 0, 0],
                              [1, 1, 1],
                              [2, 2, 2],
                              [5, 10, 12]], 'float'),
                    np.array([[1, 2, 3],
                              [3, 2, 0],
                              [5, 20, 33],
                              [40, 80, 120]], 'float') ]
    return streamlines

def test_density_map():
    #One streamline diagonal in volume
    streamlines = [np.array([np.arange(10)]*3).T]
    shape = (10, 10, 10)
    x = np.arange(10)
    expected = np.zeros(shape)
    expected[x, x, x] = 1.
    dm = density_map(streamlines, vol_dims=shape, voxel_size=(1, 1, 1))
    assert_array_equal(dm, expected)

    #add streamline, make voxel_size smaller. Each streamline should only be
    #counted once, even if multiple points lie in a voxel
    streamlines.append(np.ones((5, 3)))
    shape = (5, 5, 5)
    x = np.arange(5)
    expected = np.zeros(shape)
    expected[x, x, x] = 1.
    expected[0, 0, 0] += 1
    dm = density_map(streamlines, vol_dims=shape, voxel_size=(2, 2, 2))
    assert_array_equal(dm, expected)
    #should work with a generator
    dm = density_map(iter(streamlines), vol_dims=shape, voxel_size=(2, 2, 2))
    assert_array_equal(dm, expected)

    # Test passing affine
    affine = np.diag([2, 2, 2, 1.])
    affine[:3, 3] = 1.
    dm = density_map(streamlines, shape, affine=affine)
    assert_array_equal(dm, expected)

    # Shift the image by 2 voxels, ie 4mm
    affine[:3, 3] -= 4.
    expected_old = expected
    new_shape = [i + 2 for i in shape]
    expected = np.zeros(new_shape)
    expected[2:, 2:, 2:] = expected_old
    dm = density_map(streamlines, new_shape, affine=affine)
    assert_array_equal(dm, expected)

def test_connectivity_matrix():
    label_volume = np.array([[[3, 0, 0],
                              [0, 0, 0],
                              [0, 0, 4]]])
    streamlines = [np.array([[0,0,0],[0,0,0],[0,2,2]], 'float'),
                   np.array([[0,0,0],[0,1,1],[0,2,2]], 'float'),
                   np.array([[0,2,2],[0,1,1],[0,0,0]], 'float')]
    expected = np.zeros((5, 5), 'int')
    expected[3, 4] = 2
    expected[4, 3] = 1
    # Check basic Case
    matrix = connectivity_matrix(streamlines, label_volume, (1, 1, 1))
    assert_array_equal(matrix, expected)
    # Test mapping
    matrix, mapping = connectivity_matrix(streamlines, label_volume, (1, 1, 1),
                                          return_mapping=True)
    assert_array_equal(matrix, expected)
    assert_equal(mapping[3, 4], [0, 1])
    assert_equal(mapping[4, 3], [2])
    assert_equal(mapping.get((0, 0)), None)
    # Test mapping and symmetric
    matrix, mapping = connectivity_matrix(streamlines, label_volume, (1, 1, 1),
                                          symmetric=True, return_mapping=True)
    assert_equal(mapping[3, 4], [0, 1, 2])
    # When symmetric only (3,4) is a key, not (4, 3)
    assert_equal(mapping.get((4, 3)), None)
    # expected output matrix is symmetric version of expected
    expected = expected + expected.T
    assert_array_equal(matrix, expected)
    # Test mapping_as_streamlines, mapping dict has lists of streamlines
    matrix, mapping = connectivity_matrix(streamlines, label_volume, (1, 1, 1),
                                          return_mapping=True,
                                          mapping_as_streamlines=True)
    assert_true(mapping[3, 4][0] is streamlines[0])
    assert_true(mapping[3, 4][1] is streamlines[1])
    assert_true(mapping[4, 3][0] is streamlines[2])

    # Test passing affine to connectivity_matrix
    expected = matrix
    affine = np.diag([-1, -1, -1, 1.])
    streamlines = [-i for i in streamlines]
    matrix = connectivity_matrix(streamlines, label_volume, affine=affine)

def test_ndbincount():
    def check(expected):
        assert_equal(bc[0, 0], expected[0])
        assert_equal(bc[0, 1], expected[1])
        assert_equal(bc[1, 0], expected[2])
        assert_equal(bc[2, 2], expected[3])
    x = np.array([[0, 0], [0, 0], [0, 1], [0, 1], [1, 0], [2, 2]]).T
    expected = [2, 2, 1, 1]
    #count occurrences in x
    bc = ndbincount(x)
    assert_equal(bc.shape, (3, 3))
    check(expected)
    #pass in shape
    bc = ndbincount(x, shape=(4, 5))
    assert_equal(bc.shape, (4, 5))
    check(expected)
    #pass in weights
    weights = np.arange(6.)
    weights[-1] = 1.23
    expeceted = [1., 5., 4., 1.23]
    bc = ndbincount(x, weights=weights)
    check(expeceted)
    #raises an error if shape is too small
    assert_raises(ValueError, ndbincount, x, None, (2, 2))

def test_reduce_labels():
    shape = (4, 5, 6)
    #labels from 100 to 220
    labels = np.arange(100, np.prod(shape)+100).reshape(shape)
    #new labels form 0 to 120, and lookup maps range(0,120) to range(100, 220)
    new_labels, lookup = reduce_labels(labels)
    assert_array_equal(new_labels, labels-100)
    assert_array_equal(lookup, labels.ravel())

def test_move_streamlines():
    streamlines = make_streamlines()
    affine = np.eye(4)
    new_streamlines = move_streamlines(streamlines, affine)
    for i, test_sl in enumerate(new_streamlines):
        assert_array_equal(test_sl, streamlines[i])

    affine[:3,3] += (4,5,6)
    new_streamlines = move_streamlines(streamlines, affine)
    for i, test_sl in enumerate(new_streamlines):
        assert_array_equal(test_sl, streamlines[i]+(4, 5, 6))

    affine = np.eye(4)
    affine = affine[[2,1,0,3]]
    new_streamlines = move_streamlines(streamlines, affine)
    for i, test_sl in enumerate(new_streamlines):
        assert_array_equal(test_sl, streamlines[i][:, [2, 1, 0]])

def test_voxel_ornt():
    sh = (40,40,40)
    sz = (1, 2, 3)
    I4 = np.eye(4)

    ras = orientation_from_string('ras')
    sra = orientation_from_string('sra')
    lpi = orientation_from_string('lpi')
    srp = orientation_from_string('srp')

    affine = reorder_voxels_affine(ras, ras, sh, sz)
    assert_array_equal(affine, I4)
    affine = reorder_voxels_affine(sra, sra, sh, sz)
    assert_array_equal(affine, I4)
    affine = reorder_voxels_affine(lpi, lpi, sh, sz)
    assert_array_equal(affine, I4)
    affine = reorder_voxels_affine(srp, srp, sh, sz)
    assert_array_equal(affine, I4)

    streamlines = make_streamlines()
    box = np.array(sh)*sz

    sra_affine = reorder_voxels_affine(ras, sra, sh, sz)
    toras_affine = reorder_voxels_affine(sra, ras, sh, sz)
    assert_array_equal(np.dot(toras_affine, sra_affine), I4)
    expected_sl = (sl[:, [2, 0, 1]] for sl in streamlines)
    test_sl = move_streamlines(streamlines, sra_affine)
    for ii in xrange(len(streamlines)):
        assert_array_equal(next(test_sl), next(expected_sl))

    lpi_affine = reorder_voxels_affine(ras, lpi, sh, sz)
    toras_affine = reorder_voxels_affine(lpi, ras, sh, sz)
    assert_array_equal(np.dot(toras_affine, lpi_affine), I4)
    expected_sl = (box - sl for sl in streamlines)
    test_sl = move_streamlines(streamlines, lpi_affine)
    for ii in xrange(len(streamlines)):
        assert_array_equal(next(test_sl), next(expected_sl))

    srp_affine = reorder_voxels_affine(ras, srp, sh, sz)
    toras_affine = reorder_voxels_affine(srp, ras, (40,40,40), (3,1,2))
    assert_array_equal(np.dot(toras_affine, srp_affine), I4)
    expected_sl = [sl.copy() for sl in streamlines]
    for sl in expected_sl:
        sl[:, 1] = box[1] - sl[:, 1]
    expected_sl = (sl[:, [2, 0, 1]] for sl in expected_sl)
    test_sl = move_streamlines(streamlines, srp_affine)
    for ii in xrange(len(streamlines)):
        assert_array_equal(next(test_sl), next(expected_sl))

def test_streamline_mapping():
    streamlines = [np.array([[0,0,0],[0,0,0],[0,2,2]], 'float'),
                   np.array([[0,0,0],[0,1,1],[0,2,2]], 'float'),
                   np.array([[0,2,2],[0,1,1],[0,0,0]], 'float')]
    mapping = streamline_mapping(streamlines, (1,1,1))
    expected = {(0,0,0):[0,1,2], (0,2,2):[0,1,2], (0,1,1):[1,2]}
    assert_equal(mapping, expected)

    mapping = streamline_mapping(streamlines, (1,1,1),
                                 mapping_as_streamlines=True)
    expected = dict((k, [streamlines[i] for i in indices])
                    for k, indices in expected.items())
    assert_equal(mapping, expected)

def test_rmi():

    I1 = _rmi([3, 4], [10, 10])
    assert_equal(I1, 34)
    I1 = _rmi([0, 0], [10, 10])
    assert_equal(I1, 0)
    assert_raises(ValueError, _rmi, [10, 0], [10, 10])

    try:
        from numpy import ravel_multi_index
    except ImportError:
        raise nose.SkipTest()

    A, B, C, D = np.random.randint(0, 1000, size=[4, 100])

    I1 = _rmi([A, B], dims=[1000, 1000])
    I2 = ravel_multi_index([A, B], dims=[1000, 1000])
    assert_array_equal(I1, I2)

    I1 = _rmi([A, B, C, D], dims=[1000]*4)
    I2 = ravel_multi_index([A, B, C, D], dims=[1000]*4)
    assert_array_equal(I1, I2)

<commit_msg>TST - added test for affine with streamline_mapping<commit_after>from __future__ import division, print_function, absolute_import

from ...utils.six.moves import xrange

import numpy as np
import nose
from dipy.io.bvectxt import orientation_from_string
from dipy.tracking.utils import (_rmi, connectivity_matrix, density_map,
                                 move_streamlines, ndbincount, reduce_labels,
                                 reorder_voxels_affine, streamline_mapping)
from numpy.testing import assert_array_almost_equal, assert_array_equal
from nose.tools import assert_equal, assert_raises, assert_true

def make_streamlines():
    streamlines = [ np.array([[0, 0, 0],
                              [1, 1, 1],
                              [2, 2, 2],
                              [5, 10, 12]], 'float'),
                    np.array([[1, 2, 3],
                              [3, 2, 0],
                              [5, 20, 33],
                              [40, 80, 120]], 'float') ]
    return streamlines

def test_density_map():
    #One streamline diagonal in volume
    streamlines = [np.array([np.arange(10)]*3).T]
    shape = (10, 10, 10)
    x = np.arange(10)
    expected = np.zeros(shape)
    expected[x, x, x] = 1.
    dm = density_map(streamlines, vol_dims=shape, voxel_size=(1, 1, 1))
    assert_array_equal(dm, expected)

    #add streamline, make voxel_size smaller. Each streamline should only be
    #counted once, even if multiple points lie in a voxel
    streamlines.append(np.ones((5, 3)))
    shape = (5, 5, 5)
    x = np.arange(5)
    expected = np.zeros(shape)
    expected[x, x, x] = 1.
    expected[0, 0, 0] += 1
    dm = density_map(streamlines, vol_dims=shape, voxel_size=(2, 2, 2))
    assert_array_equal(dm, expected)
    #should work with a generator
    dm = density_map(iter(streamlines), vol_dims=shape, voxel_size=(2, 2, 2))
    assert_array_equal(dm, expected)

    # Test passing affine
    affine = np.diag([2, 2, 2, 1.])
    affine[:3, 3] = 1.
    dm = density_map(streamlines, shape, affine=affine)
    assert_array_equal(dm, expected)

    # Shift the image by 2 voxels, ie 4mm
    affine[:3, 3] -= 4.
    expected_old = expected
    new_shape = [i + 2 for i in shape]
    expected = np.zeros(new_shape)
    expected[2:, 2:, 2:] = expected_old
    dm = density_map(streamlines, new_shape, affine=affine)
    assert_array_equal(dm, expected)

def test_connectivity_matrix():
    label_volume = np.array([[[3, 0, 0],
                              [0, 0, 0],
                              [0, 0, 4]]])
    streamlines = [np.array([[0,0,0],[0,0,0],[0,2,2]], 'float'),
                   np.array([[0,0,0],[0,1,1],[0,2,2]], 'float'),
                   np.array([[0,2,2],[0,1,1],[0,0,0]], 'float')]
    expected = np.zeros((5, 5), 'int')
    expected[3, 4] = 2
    expected[4, 3] = 1
    # Check basic Case
    matrix = connectivity_matrix(streamlines, label_volume, (1, 1, 1))
    assert_array_equal(matrix, expected)
    # Test mapping
    matrix, mapping = connectivity_matrix(streamlines, label_volume, (1, 1, 1),
                                          return_mapping=True)
    assert_array_equal(matrix, expected)
    assert_equal(mapping[3, 4], [0, 1])
    assert_equal(mapping[4, 3], [2])
    assert_equal(mapping.get((0, 0)), None)
    # Test mapping and symmetric
    matrix, mapping = connectivity_matrix(streamlines, label_volume, (1, 1, 1),
                                          symmetric=True, return_mapping=True)
    assert_equal(mapping[3, 4], [0, 1, 2])
    # When symmetric only (3,4) is a key, not (4, 3)
    assert_equal(mapping.get((4, 3)), None)
    # expected output matrix is symmetric version of expected
    expected = expected + expected.T
    assert_array_equal(matrix, expected)
    # Test mapping_as_streamlines, mapping dict has lists of streamlines
    matrix, mapping = connectivity_matrix(streamlines, label_volume, (1, 1, 1),
                                          return_mapping=True,
                                          mapping_as_streamlines=True)
    assert_true(mapping[3, 4][0] is streamlines[0])
    assert_true(mapping[3, 4][1] is streamlines[1])
    assert_true(mapping[4, 3][0] is streamlines[2])

    # Test passing affine to connectivity_matrix
    expected = matrix
    affine = np.diag([-1, -1, -1, 1.])
    streamlines = [-i for i in streamlines]
    matrix = connectivity_matrix(streamlines, label_volume, affine=affine)

def test_ndbincount():
    def check(expected):
        assert_equal(bc[0, 0], expected[0])
        assert_equal(bc[0, 1], expected[1])
        assert_equal(bc[1, 0], expected[2])
        assert_equal(bc[2, 2], expected[3])
    x = np.array([[0, 0], [0, 0], [0, 1], [0, 1], [1, 0], [2, 2]]).T
    expected = [2, 2, 1, 1]
    #count occurrences in x
    bc = ndbincount(x)
    assert_equal(bc.shape, (3, 3))
    check(expected)
    #pass in shape
    bc = ndbincount(x, shape=(4, 5))
    assert_equal(bc.shape, (4, 5))
    check(expected)
    #pass in weights
    weights = np.arange(6.)
    weights[-1] = 1.23
    expeceted = [1., 5., 4., 1.23]
    bc = ndbincount(x, weights=weights)
    check(expeceted)
    #raises an error if shape is too small
    assert_raises(ValueError, ndbincount, x, None, (2, 2))

def test_reduce_labels():
    shape = (4, 5, 6)
    #labels from 100 to 220
    labels = np.arange(100, np.prod(shape)+100).reshape(shape)
    #new labels form 0 to 120, and lookup maps range(0,120) to range(100, 220)
    new_labels, lookup = reduce_labels(labels)
    assert_array_equal(new_labels, labels-100)
    assert_array_equal(lookup, labels.ravel())

def test_move_streamlines():
    streamlines = make_streamlines()
    affine = np.eye(4)
    new_streamlines = move_streamlines(streamlines, affine)
    for i, test_sl in enumerate(new_streamlines):
        assert_array_equal(test_sl, streamlines[i])

    affine[:3,3] += (4,5,6)
    new_streamlines = move_streamlines(streamlines, affine)
    for i, test_sl in enumerate(new_streamlines):
        assert_array_equal(test_sl, streamlines[i]+(4, 5, 6))

    affine = np.eye(4)
    affine = affine[[2,1,0,3]]
    new_streamlines = move_streamlines(streamlines, affine)
    for i, test_sl in enumerate(new_streamlines):
        assert_array_equal(test_sl, streamlines[i][:, [2, 1, 0]])

def test_voxel_ornt():
    sh = (40,40,40)
    sz = (1, 2, 3)
    I4 = np.eye(4)

    ras = orientation_from_string('ras')
    sra = orientation_from_string('sra')
    lpi = orientation_from_string('lpi')
    srp = orientation_from_string('srp')

    affine = reorder_voxels_affine(ras, ras, sh, sz)
    assert_array_equal(affine, I4)
    affine = reorder_voxels_affine(sra, sra, sh, sz)
    assert_array_equal(affine, I4)
    affine = reorder_voxels_affine(lpi, lpi, sh, sz)
    assert_array_equal(affine, I4)
    affine = reorder_voxels_affine(srp, srp, sh, sz)
    assert_array_equal(affine, I4)

    streamlines = make_streamlines()
    box = np.array(sh)*sz

    sra_affine = reorder_voxels_affine(ras, sra, sh, sz)
    toras_affine = reorder_voxels_affine(sra, ras, sh, sz)
    assert_array_equal(np.dot(toras_affine, sra_affine), I4)
    expected_sl = (sl[:, [2, 0, 1]] for sl in streamlines)
    test_sl = move_streamlines(streamlines, sra_affine)
    for ii in xrange(len(streamlines)):
        assert_array_equal(next(test_sl), next(expected_sl))

    lpi_affine = reorder_voxels_affine(ras, lpi, sh, sz)
    toras_affine = reorder_voxels_affine(lpi, ras, sh, sz)
    assert_array_equal(np.dot(toras_affine, lpi_affine), I4)
    expected_sl = (box - sl for sl in streamlines)
    test_sl = move_streamlines(streamlines, lpi_affine)
    for ii in xrange(len(streamlines)):
        assert_array_equal(next(test_sl), next(expected_sl))

    srp_affine = reorder_voxels_affine(ras, srp, sh, sz)
    toras_affine = reorder_voxels_affine(srp, ras, (40,40,40), (3,1,2))
    assert_array_equal(np.dot(toras_affine, srp_affine), I4)
    expected_sl = [sl.copy() for sl in streamlines]
    for sl in expected_sl:
        sl[:, 1] = box[1] - sl[:, 1]
    expected_sl = (sl[:, [2, 0, 1]] for sl in expected_sl)
    test_sl = move_streamlines(streamlines, srp_affine)
    for ii in xrange(len(streamlines)):
        assert_array_equal(next(test_sl), next(expected_sl))

def test_streamline_mapping():
    streamlines = [np.array([[0,0,0],[0,0,0],[0,2,2]], 'float'),
                   np.array([[0,0,0],[0,1,1],[0,2,2]], 'float'),
                   np.array([[0,2,2],[0,1,1],[0,0,0]], 'float')]
    mapping = streamline_mapping(streamlines, (1,1,1))
    expected = {(0,0,0):[0,1,2], (0,2,2):[0,1,2], (0,1,1):[1,2]}
    assert_equal(mapping, expected)

    mapping = streamline_mapping(streamlines, (1,1,1),
                                 mapping_as_streamlines=True)
    expected = dict((k, [streamlines[i] for i in indices])
                    for k, indices in expected.items())
    assert_equal(mapping, expected)

    # Test passing affine
    affine = np.eye(4)
    affine[:3, 3] = .5
    mapping = streamline_mapping(streamlines, affine=affine,
                                 mapping_as_streamlines=True)
    assert_equal(mapping, expected)

    # Make the voxel size smaller
    affine = np.diag([.5, .5, .5, 1.])
    affine[:3, 3] = .25
    expected = dict((tuple(i*2 for i in key), value)
                    for key, value in expected.items())
    mapping = streamline_mapping(streamlines, affine=affine,
                                 mapping_as_streamlines=True)
    assert_equal(mapping, expected)

def test_rmi():

    I1 = _rmi([3, 4], [10, 10])
    assert_equal(I1, 34)
    I1 = _rmi([0, 0], [10, 10])
    assert_equal(I1, 0)
    assert_raises(ValueError, _rmi, [10, 0], [10, 10])

    try:
        from numpy import ravel_multi_index
    except ImportError:
        raise nose.SkipTest()

    A, B, C, D = np.random.randint(0, 1000, size=[4, 100])

    I1 = _rmi([A, B], dims=[1000, 1000])
    I2 = ravel_multi_index([A, B], dims=[1000, 1000])
    assert_array_equal(I1, I2)

    I1 = _rmi([A, B, C, D], dims=[1000]*4)
    I2 = ravel_multi_index([A, B, C, D], dims=[1000]*4)
    assert_array_equal(I1, I2)

<|endoftext|>


<commit_before>import os
import sys
import unittest

import intelmq.lib.test as test
if os.environ.get('INTELMQ_TEST_EXOTIC'):
    from intelmq.bots.outputs.misp.output_api import MISPAPIOutputBot

# This file is a stub
# We cannot do much more as we are missing a mock MISP instance to use
# to initialise pymisp
<commit_msg>STY: fix codestyle of test_output_api<commit_after>import os
#  import unittest

#  import intelmq.lib.test as test
if os.environ.get('INTELMQ_TEST_EXOTIC'):
    from intelmq.bots.outputs.misp.output_api import MISPAPIOutputBot  # noqa

# This file is a stub
# We cannot do much more as we are missing a mock MISP instance to use
# to initialise pymisp
<|endoftext|>


<commit_before># Copyright (c) 2005-2006 Ralph Meijer
# See LICENSE for details

import re, time

import feedparser, simplejson

from zope.interface import Interface, implements

from twisted.application import service
from twisted.internet import reactor
from twisted.python import components, failure
from twisted.words.xish import domish

from mimir.common import extension, pubsub

SGMLTAG = re.compile('<.+?>', re.DOTALL)
NS_ATOM = 'http://www.w3.org/2005/Atom'

class FeedParserEncoder(simplejson.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, time.struct_time):
            return dict([(attr, getattr(obj, attr)) for attr in dir(obj)
                                                  if attr.startswith('tm_')])
        else:
            return simplejson.JSONEncoder.default(self, obj)

class NoNotify(Exception):
    pass

class INewsService(Interface):

    def process(channel, items):
        """
        Process new items for channel

        @param items: atom entries
        @type items: L{domish.Element}
        """

class NewsService(service.Service):

    implements(INewsService)

    def __init__(self, presenceMonitor, dbpool):
        self._dbpool = dbpool
        presenceMonitor.register_callback(self.onPresenceChange)

    def error(self, failure):
        print failure

    def onPresenceChange(self, entity, available, show):
        if available:
            if show not in ['away', 'xa', 'dnd', 'chat']:
                show = 'online'
        else:
            show = 'offline'

        print "  presence change to %s for %s", (show, repr(entity.full()))
        reactor.callLater(5, self.pageNotify, entity, show)

    def pageNotify(self, entity, show):
        d = self._dbpool.runInteraction(self._checkNotify, entity, show)
        d.addCallback(self._doNotify, entity)
        d.addCallback(self._setNotified)
        d.addErrback(lambda failure: failure.trap(NoNotify))
        d.addErrback(self.error)

    def _checkNotify(self, cursor, entity, show):
        cursor.execute("""SELECT user_id, message_type, ssl,
                                 count(news_id) as count
                          FROM auth_user
                          NATURAL JOIN news_prefs
                          NATURAL JOIN news_notify_presences
                          NATURAL JOIN news_page
                          NATURAL JOIN news_flags
                          NATURAL JOIN news
                          WHERE jid=%s AND
                                NOT suspended AND
                                presence=%s AND
                                NOT notified AND
                                date>last_visit
                          GROUP BY user_id, message_type, ssl""",
                       (entity.userhost(),
                        show))
        result = cursor.fetchone()

        if not result:
            return failure.Failure(NoNotify())

        return result

    def _doNotify(self, result, entity):
        userId, messageType, ssl, count = result

        title = u'New news on Mim\xedr!'
        link = "%s://mimir.ik.nu/news" % (ssl and 'https' or 'http')
        description = 'There '
        if count == 1:
            description += 'is 1 new item'
        else:
            description += 'are %s new items' % count
        description += ' on your news page'

        self.notifier.sendNotification(entity.full(), True, messageType,
                                       title, link, description)
        return userId

    def _setNotified(self, userId):
        return self._dbpool.runOperation("""UPDATE news_page SET notified=true
                                            WHERE user_id=%s""", userId)

    def process(self, channel, items):
        print "Got entries: %r" % items
        d = self._dbpool.runInteraction(self._processItems, channel, items)
        d.addCallback(self.notify)
        d.addErrback(self.error)

    def _processItems(self, cursor, channel, items):

        # Get channel title
        cursor.execute("SELECT title from channels WHERE channel=%s",
                       channel)
        title = cursor.fetchone()[0]

        print "Channel title: %s" % repr(title)

        feedDocument = domish.Element((NS_ATOM, 'feed'))
        for item in items:
            feedDocument.addChild(item)

        feed = feedparser.parse(feedDocument.toXml().encode('utf-8'))
        entries = feed.entries

        # Get notify list, including preferences
        cursor.execute("""SELECT user_id, jid, notify,
                                 description_in_notify, message_type,
                                 store_offline, notify_items
                          FROM news_prefs
                            NATURAL JOIN news_subscriptions
                            NATURAL JOIN news_notify
                          WHERE NOT suspended AND channel=%s""",
                       channel)

        notifyList = cursor.fetchall()

        # split notify list into a list of entities that are to be notified
        # and a list of entities that will get items marked as unread
        notifications = []
        markUnread = []
        for userId, jid, notify, descriptionInNotify, messageType, \
            storeOffline, notifyItems in \
                notifyList:
            if notify and notifyItems:
                notifications.append((jid,
                                      descriptionInNotify,
                                      messageType))
            elif storeOffline:
                markUnread.append(userId)

        # store items and mark unread
        notifyItems = []
        for entry in entries:
            newsId = self._storeItem(cursor, channel, entry)
            if newsId:
                for userId in markUnread:
                    cursor.execute("""INSERT INTO news_flags
                                      (user_id, news_id, unread) VALUES
                                      (%s, %s, true)""",
                                   (userId, newsId))
                    cursor.execute("""UPDATE news_page SET notified=false
                                      WHERE user_id=%s""",
                                   userId)

                notifyItems.append(entry)
        return (title, notifications, notifyItems)

    def _extractBasics(self, entry):
        if 'title' in entry:
            title = entry.title

            if 'source' in entry and 'title' in entry.source:
                content = "%s: %s" % (entry.source.title, title)

            if entry.title_detail == 'text/plain':
                title = feedparser._xmlescape(title)
        else:
            title = u''

        if 'link' in entry:
            link = entry.get('feedburner_origlink', entry.link)
        else:
            link = u''

        # Find a description. First try full text, then summary.
        content = None
        if 'content' in entry:
            content = entry.content[0]
        elif 'summary' in entry:
            content = entry.summary_detail

        if content:
            value = content.value
            if content.type == 'text/plain':
                value = feedparser._xmlescape(value)
            description = value
        else:
            description = u''

        date = None
        for attribute in 'updated', 'published', 'created':
            if attribute in entry:
                date = getattr(entry, '%s_parsed' % attribute)

        date = time.strftime("%Y-%m-%d %H:%M:%Sz", date or time.gmtime())

        return title, link, description, date

    def _storeItem(self, cursor, channel, entry):
        title, link, description, date = self._extractBasics(entry)
        json = simplejson.dumps(entry, cls=FeedParserEncoder)

        print "Storing item: %s" % repr(entry)

        cursor.execute("""UPDATE news
                          SET title=%s, description=%s, date=%s, parsed=%s
                          WHERE channel=%s AND link=%s""",
                       (title, description, date, json, channel, link))

        if cursor.rowcount == 1:
            print "UPDATE"
            return None

        cursor.execute("""INSERT INTO news
                          (channel, title, link, description, date, parsed)
                          VALUES
                          (%s, %s, %s, %s, %s, %s)""",
                       (channel, title, link, description, date, json))

        print "INSERT",

        cursor.execute("""SELECT news_id FROM news
                          WHERE channel=%s and link=%s""",
                          (channel, link))

        newsId = cursor.fetchone()[0]

        print newsId
        return newsId

    def notify(self, result):
        channelTitle, notifications, entries = result

        for entry in entries:
            title = "%s: %s" % (channelTitle,
                                entry.get('title', u'-- no title --'))
            link = entry.get('link', '')
            description = entry.get('summary')

            if description:
                description = SGMLTAG.sub('', description)
                description = domish.unescapeFromXml(description)
                description = description.rstrip() or None

            for jid, descriptionInNotify, messageType in notifications:
                self.notifier.sendNotification(jid, descriptionInNotify,
                                               messageType, title, link,
                                               description)

class XMPPHandlerFromService(pubsub.PubSubClient):

    def sendNotification(self, jid, descriptionInNotify, messageType,
                         title, link, description):

        message = domish.Element((None, 'message'))
        message['to'] = jid
        message['type'] = messageType

        if messageType == 'chat':
            body = "%s\n%s" % (title, link)
            if description and descriptionInNotify:
                body += "\n\n%s\n\n" % description

            message.addElement('body', None, body)
        elif messageType == 'headline':
            message.addElement('subject', None, title)
            if description:
                message.addElement('body', None, description)
            oob = message.addElement(('jabber:x:oob', 'x'))
            oob.addElement('url', None, link)
            oob.addElement('desc', None, title)

        print "Sending: %s" % repr(message.toXml())
        self.send(message)

    def itemsReceived(self, node, items):
        m = re.match(r"^mimir/news/(.+)$", node)

        if not m:
            return

        channel = m.group(1)

        entries = (item.entry for item in items
                              if item.entry and item.entry.uri == NS_ATOM)

        self.service.process(channel, entries)

components.registerAdapter(XMPPHandlerFromService,
                           INewsService,
                           extension.IExtensionProtocol)
<commit_msg>Check if content field actually has content.<commit_after># Copyright (c) 2005-2006 Ralph Meijer
# See LICENSE for details

import re, time

import feedparser, simplejson

from zope.interface import Interface, implements

from twisted.application import service
from twisted.internet import reactor
from twisted.python import components, failure
from twisted.words.xish import domish

from mimir.common import extension, pubsub

SGMLTAG = re.compile('<.+?>', re.DOTALL)
NS_ATOM = 'http://www.w3.org/2005/Atom'

class FeedParserEncoder(simplejson.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, time.struct_time):
            return dict([(attr, getattr(obj, attr)) for attr in dir(obj)
                                                  if attr.startswith('tm_')])
        else:
            return simplejson.JSONEncoder.default(self, obj)

class NoNotify(Exception):
    pass

class INewsService(Interface):

    def process(channel, items):
        """
        Process new items for channel

        @param items: atom entries
        @type items: L{domish.Element}
        """

class NewsService(service.Service):

    implements(INewsService)

    def __init__(self, presenceMonitor, dbpool):
        self._dbpool = dbpool
        presenceMonitor.register_callback(self.onPresenceChange)

    def error(self, failure):
        print failure

    def onPresenceChange(self, entity, available, show):
        if available:
            if show not in ['away', 'xa', 'dnd', 'chat']:
                show = 'online'
        else:
            show = 'offline'

        print "  presence change to %s for %s", (show, repr(entity.full()))
        reactor.callLater(5, self.pageNotify, entity, show)

    def pageNotify(self, entity, show):
        d = self._dbpool.runInteraction(self._checkNotify, entity, show)
        d.addCallback(self._doNotify, entity)
        d.addCallback(self._setNotified)
        d.addErrback(lambda failure: failure.trap(NoNotify))
        d.addErrback(self.error)

    def _checkNotify(self, cursor, entity, show):
        cursor.execute("""SELECT user_id, message_type, ssl,
                                 count(news_id) as count
                          FROM auth_user
                          NATURAL JOIN news_prefs
                          NATURAL JOIN news_notify_presences
                          NATURAL JOIN news_page
                          NATURAL JOIN news_flags
                          NATURAL JOIN news
                          WHERE jid=%s AND
                                NOT suspended AND
                                presence=%s AND
                                NOT notified AND
                                date>last_visit
                          GROUP BY user_id, message_type, ssl""",
                       (entity.userhost(),
                        show))
        result = cursor.fetchone()

        if not result:
            return failure.Failure(NoNotify())

        return result

    def _doNotify(self, result, entity):
        userId, messageType, ssl, count = result

        title = u'New news on Mim\xedr!'
        link = "%s://mimir.ik.nu/news" % (ssl and 'https' or 'http')
        description = 'There '
        if count == 1:
            description += 'is 1 new item'
        else:
            description += 'are %s new items' % count
        description += ' on your news page'

        self.notifier.sendNotification(entity.full(), True, messageType,
                                       title, link, description)
        return userId

    def _setNotified(self, userId):
        return self._dbpool.runOperation("""UPDATE news_page SET notified=true
                                            WHERE user_id=%s""", userId)

    def process(self, channel, items):
        print "Got entries: %r" % items
        d = self._dbpool.runInteraction(self._processItems, channel, items)
        d.addCallback(self.notify)
        d.addErrback(self.error)

    def _processItems(self, cursor, channel, items):

        # Get channel title
        cursor.execute("SELECT title from channels WHERE channel=%s",
                       channel)
        title = cursor.fetchone()[0]

        print "Channel title: %s" % repr(title)

        feedDocument = domish.Element((NS_ATOM, 'feed'))
        for item in items:
            feedDocument.addChild(item)

        feed = feedparser.parse(feedDocument.toXml().encode('utf-8'))
        entries = feed.entries

        # Get notify list, including preferences
        cursor.execute("""SELECT user_id, jid, notify,
                                 description_in_notify, message_type,
                                 store_offline, notify_items
                          FROM news_prefs
                            NATURAL JOIN news_subscriptions
                            NATURAL JOIN news_notify
                          WHERE NOT suspended AND channel=%s""",
                       channel)

        notifyList = cursor.fetchall()

        # split notify list into a list of entities that are to be notified
        # and a list of entities that will get items marked as unread
        notifications = []
        markUnread = []
        for userId, jid, notify, descriptionInNotify, messageType, \
            storeOffline, notifyItems in \
                notifyList:
            if notify and notifyItems:
                notifications.append((jid,
                                      descriptionInNotify,
                                      messageType))
            elif storeOffline:
                markUnread.append(userId)

        # store items and mark unread
        notifyItems = []
        for entry in entries:
            newsId = self._storeItem(cursor, channel, entry)
            if newsId:
                for userId in markUnread:
                    cursor.execute("""INSERT INTO news_flags
                                      (user_id, news_id, unread) VALUES
                                      (%s, %s, true)""",
                                   (userId, newsId))
                    cursor.execute("""UPDATE news_page SET notified=false
                                      WHERE user_id=%s""",
                                   userId)

                notifyItems.append(entry)
        return (title, notifications, notifyItems)

    def _extractBasics(self, entry):
        if 'title' in entry:
            title = entry.title

            if 'source' in entry and 'title' in entry.source:
                content = "%s: %s" % (entry.source.title, title)

            if entry.title_detail == 'text/plain':
                title = feedparser._xmlescape(title)
        else:
            title = u''

        if 'link' in entry:
            link = entry.get('feedburner_origlink', entry.link)
        else:
            link = u''

        # Find a description. First try full text, then summary.
        content = None
        if 'content' in entry and entry.content[0].value:
            content = entry.content[0]
        elif 'summary' in entry:
            content = entry.summary_detail

        if content:
            value = content.value
            if content.type == 'text/plain':
                value = feedparser._xmlescape(value)
            description = value
        else:
            description = u''

        date = None
        for attribute in 'updated', 'published', 'created':
            if attribute in entry:
                date = getattr(entry, '%s_parsed' % attribute)

        date = time.strftime("%Y-%m-%d %H:%M:%Sz", date or time.gmtime())

        return title, link, description, date

    def _storeItem(self, cursor, channel, entry):
        title, link, description, date = self._extractBasics(entry)
        json = simplejson.dumps(entry, cls=FeedParserEncoder)

        print "Storing item: %s" % repr(entry)

        cursor.execute("""UPDATE news
                          SET title=%s, description=%s, date=%s, parsed=%s
                          WHERE channel=%s AND link=%s""",
                       (title, description, date, json, channel, link))

        if cursor.rowcount == 1:
            print "UPDATE"
            return None

        cursor.execute("""INSERT INTO news
                          (channel, title, link, description, date, parsed)
                          VALUES
                          (%s, %s, %s, %s, %s, %s)""",
                       (channel, title, link, description, date, json))

        print "INSERT",

        cursor.execute("""SELECT news_id FROM news
                          WHERE channel=%s and link=%s""",
                          (channel, link))

        newsId = cursor.fetchone()[0]

        print newsId
        return newsId

    def notify(self, result):
        channelTitle, notifications, entries = result

        for entry in entries:
            title = "%s: %s" % (channelTitle,
                                entry.get('title', u'-- no title --'))
            link = entry.get('link', '')
            description = entry.get('summary')

            if description:
                description = SGMLTAG.sub('', description)
                description = domish.unescapeFromXml(description)
                description = description.rstrip() or None

            for jid, descriptionInNotify, messageType in notifications:
                self.notifier.sendNotification(jid, descriptionInNotify,
                                               messageType, title, link,
                                               description)

class XMPPHandlerFromService(pubsub.PubSubClient):

    def sendNotification(self, jid, descriptionInNotify, messageType,
                         title, link, description):

        message = domish.Element((None, 'message'))
        message['to'] = jid
        message['type'] = messageType

        if messageType == 'chat':
            body = "%s\n%s" % (title, link)
            if description and descriptionInNotify:
                body += "\n\n%s\n\n" % description

            message.addElement('body', None, body)
        elif messageType == 'headline':
            message.addElement('subject', None, title)
            if description:
                message.addElement('body', None, description)
            oob = message.addElement(('jabber:x:oob', 'x'))
            oob.addElement('url', None, link)
            oob.addElement('desc', None, title)

        print "Sending: %s" % repr(message.toXml())
        self.send(message)

    def itemsReceived(self, node, items):
        m = re.match(r"^mimir/news/(.+)$", node)

        if not m:
            return

        channel = m.group(1)

        entries = (item.entry for item in items
                              if item.entry and item.entry.uri == NS_ATOM)

        self.service.process(channel, entries)

components.registerAdapter(XMPPHandlerFromService,
                           INewsService,
                           extension.IExtensionProtocol)
<|endoftext|>


<commit_before>'''
Created on Aug 14, 2015

@author: Connor
'''

class Module(object):
    def __init__(self, mod_params, mod_tasks=[]):
        """ Make a unique mod name """
        self.name = mod_params['name']

        """ Modules match and execute in prioritized order """
        self.priority = mod_params['priority']

        """ Greedy mods stop future mods from being matched """
        self.greedy = True
        if 'greedy' in mod_params:
            self.greedy = mod_params['greedy']

        """ True if the mod is enabled """
        self.enabled = True
        if 'enabled' in mod_params:
            self.enabled = mod_params['enabled']

        """ Tasks find input text patterns and perform an action """
        self.tasks = mod_tasks
        self.tasks.sort(key=lambda task: task.priority, reverse=True)<commit_msg>Quick module fix<commit_after>'''
Created on Aug 14, 2015

@author: Connor
'''

class Module(object):
    def __init__(self, mod_params, mod_tasks=[]):
        """ Make a unique mod name """
        self.name = mod_params['name']

        """ Modules match and execute in prioritized order """
        self.priority = 0
        if 'priority' in mod_params:
            self.priority = mod_params['priority']

        """ Greedy mods stop future mods from being matched """
        self.greedy = True
        if 'greedy' in mod_params:
            self.greedy = mod_params['greedy']

        """ True if the mod is enabled """
        self.enabled = True
        if 'enabled' in mod_params:
            self.enabled = mod_params['enabled']

        """ Tasks find input text patterns and perform an action """
        self.tasks = mod_tasks
        self.tasks.sort(key=lambda task: task.priority, reverse=True)<|endoftext|>


<commit_before>import GPflow
import tensorflow as tf
import numpy as np
import unittest
from .reference import referenceRbfKernel

def referenceUnivariateLogMarginalLikelihood(y, K, noiseVariance):
    return -0.5 * y * y / (K + noiseVariance) - 0.5 * np.log(K + noiseVariance) - 0.5 * np.log(np.pi * 2.)

def referenceUnivariatePosterior(y, K, noiseVariance):
    mean = K * y / (K + noiseVariance)
    variance = K - K / (K + noiseVariance)
    return mean, variance

def referenceUnivariatePriorKL(meanA, meanB, varA, varB):
    # KL[ qA | qB ] = E_{qA} \log [qA / qB] where qA and qB are univariate normal distributions.
    return 0.5 * (np.log(varB) - np.log(varA) - 1. + varA/varB + (meanB-meanA) * (meanB - meanA) / varB)

def referenceMultivariatePriorKL(meanA, covA, meanB, covB):
    # KL[ qA | qB ] = E_{qA} \log [qA / qB] where qA and aB are K dimensional multivariate normal distributions.
    # Analytically tractable and equal to...
    # 0.5 * (Tr(covB^{-1} covA) + (meanB - meanA)^T covB^{-1} (meanB - meanA) - K + log(det(covB)) - log (det(covA)))
    K = covA.shape[0]
    traceTerm = 0.5 * np.trace(np.linalg.solve(covB, covA))
    delta = meanB - meanA
    mahalanobisTerm = 0.5 * np.dot(delta.T, np.linalg.solve(covB, delta))
    constantTerm = -0.5 * K
    priorLogDeterminantTerm = 0.5*np.linalg.slogdet(covB)[1]
    variationalLogDeterminantTerm = -0.5 * np.linalg.slogdet(covA)[1]
    return traceTerm + mahalanobisTerm + constantTerm + priorLogDeterminantTerm + variationalLogDeterminantTerm

def kernel(kernelVariance=1, lengthScale=1.):
    kern = GPflow.kernels.RBF(1)
    kern.variance = kernelVariance
    kern.lengthscales = lengthScale
    return kern

class VariationalUnivariateTest(unittest.TestCase):
    def setUp(self):
        tf.reset_default_graph()
        self.y_real = 2.
        self.K = 1.
        self.noiseVariance = 0.5
        self.univariate = 1
        self.oneLatentFunction = 1
        self.meanZero = 0.
        self.X = np.atleast_2d(np.array([0.]))
        self.Y = np.atleast_2d(np.array([self.y_real]))
        self.Z = self.X.copy()
        self.lik = GPflow.likelihoods.Gaussian()
        self.lik.variance = self.noiseVariance
        self.posteriorMean, self.posteriorVariance = referenceUnivariatePosterior(y=self.y_real, K=self.K,
                                                                                noiseVariance=self.noiseVariance)
        self.posteriorStd = np.sqrt(self.posteriorVariance)

    def get_model(self, is_diagonal, is_whitened):
        m = GPflow.svgp.SVGP(X=self.X, Y=self.Y,
                             kern=kernel(kernelVariance=self.K),
                             likelihood=self.lik, Z=self.Z, q_diag=is_diagonal, whiten=is_whitened)
        if is_diagonal:
            m.q_sqrt = np.ones((self.univariate, self.oneLatentFunction))*self.posteriorStd
        else:
            m.q_sqrt = np.ones((self.univariate, self.univariate, self.oneLatentFunction)) * self.posteriorStd
        m.q_mu = np.ones((self.univariate, self.oneLatentFunction)) * self.posteriorMean
        m.make_tf_array(m._free_vars)
        return m

    def test_prior_KL(self):
        meanA = self.posteriorMean
        varA = self.posteriorVariance
        meanB = self.meanZero  # Assumes a zero
        varB = self.K

        referenceKL = referenceUnivariatePriorKL(meanA, meanB, varA, varB)

        for is_diagonal in [True, False]:
            for is_whitened in [True, False]:
                model = self.get_model(is_diagonal, is_whitened)
                with model.tf_mode():
                    prior_KL_function = model.build_prior_KL()
                test_prior_KL = prior_KL_function.eval(session=model._session,
                                                       feed_dict={model._free_vars: model.get_free_state()})
                self.assertTrue(np.abs(referenceKL - test_prior_KL) < 1e-4)

    def test_build_likelihood(self):
        # reference marginal likelihood
        log_marginal_likelihood = referenceUnivariateLogMarginalLikelihood(y=self.y_real,
                                                                           K=self.K, noiseVariance=self.noiseVariance)

        for is_diagonal in [True, False]:
            for is_whitened in [True, False]:
                model = self.get_model(is_diagonal, is_whitened)
                model_likelihood = model.compute_log_likelihood()
                self.assertTrue(np.abs(model_likelihood - log_marginal_likelihood) < 1e-4)

    def testUnivariateConditionals(self):
        for is_diagonal in [True, False]:
            for is_whitened in [True, False]:
                m = self.get_model(is_diagonal, is_whitened)
                with m.tf_mode():
                    if is_whitened:
                        args = (self.X, self.Z, m.kern, m.q_mu, m.q_sqrt, self.oneLatentFunction)
                        fmean_func, fvar_func = GPflow.conditionals.gaussian_gp_predict_whitened(*args)
                    else:
                        args = (self.X, self.Z, m.kern, m.q_mu, m.q_sqrt, self.oneLatentFunction)
                        fmean_func, fvar_func = GPflow.conditionals.gaussian_gp_predict(*args)
                mean_value = fmean_func.eval(session=m._session, feed_dict={m._free_vars: m.get_free_state()})[0, 0]
                var_value = fvar_func.eval(session=m._session, feed_dict={m._free_vars: m.get_free_state()})[0, 0]
                self.assertTrue(np.abs(mean_value - self.posteriorMean) < 1e-4)
                self.assertTrue(np.abs(var_value - self.posteriorVariance) < 1e-4)

class VariationalMultivariateTest(unittest.TestCase):
    def setUp(self):
        tf.reset_default_graph()
        self.nDimensions = 3
        self.rng = np.random.RandomState(1)
        self.Y = self.rng.randn(self.nDimensions, 1)
        self.X = self.rng.randn(self.nDimensions, 1)
        self.Z = self.X.copy()
        self.noiseVariance = 0.5
        self.signalVariance = 1.5
        self.lengthScale = 1.7
        self.oneLatentFunction = 1
        self.lik = GPflow.likelihoods.Gaussian()
        self.lik.variance = self.noiseVariance
        self.q_mean = self.rng.randn(self.nDimensions, self.oneLatentFunction)
        self.q_sqrt_diag = self.rng.rand(self.nDimensions, self.oneLatentFunction)
        self.q_sqrt_full = np.tril(self.rng.rand(self.nDimensions, self.nDimensions))

    def getModel(self, is_diagonal, is_whitened):
        model = GPflow.svgp.SVGP(X=self.X, Y=self.Y,
                                 kern=kernel(kernelVariance=self.signalVariance,
                                             lengthScale=self.lengthScale),
                                 likelihood=self.lik, Z=self.Z, q_diag=is_diagonal, whiten=is_whitened)
        if is_diagonal:
            model.q_sqrt = self.q_sqrt_diag
        else:
            model.q_sqrt = self.q_sqrt_full[:, :, None]
        model.q_mu = self.q_mean
        return model

    def test_refrence_implementation_consistency(self):
        rng = np.random.RandomState(10)
        qMean = rng.randn()
        qCov = rng.rand()
        pMean = rng.rand()
        pCov = rng.rand()
        univariate_KL = referenceUnivariatePriorKL(qMean, pMean, qCov, pCov)
        multivariate_KL = referenceMultivariatePriorKL(np.array([[qMean]]), np.array([[qCov]]),
                                                       np.array([[pMean]]), np.array([[pCov]]))
        self.assertTrue(np.abs(univariate_KL - multivariate_KL) < 1e-4)

    def test_prior_KL_fullQ(self):
        covQ = np.dot(self.q_sqrt_full, self.q_sqrt_full.T)
        mean_prior = np.zeros((self.nDimensions, 1))

        for is_whitened in [True, False]:
            m = self.getModel(False, is_whitened)

            if is_whitened:
                cov_prior = np.eye(self.nDimensions)
            else:
                cov_prior = referenceRbfKernel(self.X, self.lengthScale, self.signalVariance)

            referenceKL = referenceMultivariatePriorKL(self.q_mean, covQ, mean_prior, cov_prior)
            # now get test KL.
            test_prior_KL = GPflow.param.AutoFlow()(m.build_prior_KL.__func__)(m)
            self.assertTrue(np.abs(referenceKL - test_prior_KL) < 1e-4)

if __name__ == "__main__":
    unittest.main()
<commit_msg>fixes in test_variational<commit_after>import GPflow
import tensorflow as tf
import numpy as np
import unittest
from .reference import referenceRbfKernel

def referenceUnivariateLogMarginalLikelihood(y, K, noiseVariance):
    return -0.5 * y * y / (K + noiseVariance) - 0.5 * np.log(K + noiseVariance) - 0.5 * np.log(np.pi * 2.)

def referenceUnivariatePosterior(y, K, noiseVariance):
    mean = K * y / (K + noiseVariance)
    variance = K - K / (K + noiseVariance)
    return mean, variance

def referenceUnivariatePriorKL(meanA, meanB, varA, varB):
    # KL[ qA | qB ] = E_{qA} \log [qA / qB] where qA and qB are univariate normal distributions.
    return 0.5 * (np.log(varB) - np.log(varA) - 1. + varA/varB + (meanB-meanA) * (meanB - meanA) / varB)

def referenceMultivariatePriorKL(meanA, covA, meanB, covB):
    # KL[ qA | qB ] = E_{qA} \log [qA / qB] where qA and aB are K dimensional multivariate normal distributions.
    # Analytically tractable and equal to...
    # 0.5 * (Tr(covB^{-1} covA) + (meanB - meanA)^T covB^{-1} (meanB - meanA) - K + log(det(covB)) - log (det(covA)))
    K = covA.shape[0]
    traceTerm = 0.5 * np.trace(np.linalg.solve(covB, covA))
    delta = meanB - meanA
    mahalanobisTerm = 0.5 * np.dot(delta.T, np.linalg.solve(covB, delta))
    constantTerm = -0.5 * K
    priorLogDeterminantTerm = 0.5*np.linalg.slogdet(covB)[1]
    variationalLogDeterminantTerm = -0.5 * np.linalg.slogdet(covA)[1]
    return traceTerm + mahalanobisTerm + constantTerm + priorLogDeterminantTerm + variationalLogDeterminantTerm

def kernel(kernelVariance=1, lengthScale=1.):
    kern = GPflow.kernels.RBF(1)
    kern.variance = kernelVariance
    kern.lengthscales = lengthScale
    return kern

class VariationalUnivariateTest(unittest.TestCase):
    def setUp(self):
        tf.reset_default_graph()
        self.y_real = 2.
        self.K = 1.
        self.noiseVariance = 0.5
        self.univariate = 1
        self.oneLatentFunction = 1
        self.meanZero = 0.
        self.X = np.atleast_2d(np.array([0.]))
        self.Y = np.atleast_2d(np.array([self.y_real]))
        self.Z = self.X.copy()
        self.lik = GPflow.likelihoods.Gaussian()
        self.lik.variance = self.noiseVariance
        self.posteriorMean, self.posteriorVariance = referenceUnivariatePosterior(y=self.y_real, K=self.K,
                                                                                noiseVariance=self.noiseVariance)
        self.posteriorStd = np.sqrt(self.posteriorVariance)

    def get_model(self, is_diagonal, is_whitened):
        m = GPflow.svgp.SVGP(X=self.X, Y=self.Y,
                             kern=kernel(kernelVariance=self.K),
                             likelihood=self.lik, Z=self.Z, q_diag=is_diagonal, whiten=is_whitened)
        if is_diagonal:
            m.q_sqrt = np.ones((self.univariate, self.oneLatentFunction))*self.posteriorStd
        else:
            m.q_sqrt = np.ones((self.univariate, self.univariate, self.oneLatentFunction)) * self.posteriorStd
        m.q_mu = np.ones((self.univariate, self.oneLatentFunction)) * self.posteriorMean
        return m

    def test_prior_KL(self):
        meanA = self.posteriorMean
        varA = self.posteriorVariance
        meanB = self.meanZero  # Assumes a zero
        varB = self.K

        referenceKL = referenceUnivariatePriorKL(meanA, meanB, varA, varB)

        for is_diagonal in [True, False]:
            for is_whitened in [True, False]:
                m = self.get_model(is_diagonal, is_whitened)

                test_prior_KL = GPflow.param.AutoFlow()(m.build_prior_KL.__func__)(m)
                self.assertTrue(np.abs(referenceKL - test_prior_KL) < 1e-4)

    def test_build_likelihood(self):
        # reference marginal likelihood
        log_marginal_likelihood = referenceUnivariateLogMarginalLikelihood(y=self.y_real,
                                                                           K=self.K, noiseVariance=self.noiseVariance)

        for is_diagonal in [True, False]:
            for is_whitened in [True, False]:
                model = self.get_model(is_diagonal, is_whitened)
                model_likelihood = model.compute_log_likelihood()
                self.assertTrue(np.abs(model_likelihood - log_marginal_likelihood) < 1e-4)

    def testUnivariateConditionals(self):
        for is_diagonal in [True, False]:
            for is_whitened in [True, False]:
                m = self.get_model(is_diagonal, is_whitened)
                free_vars = tf.placeholder(tf.float64)
                session = tf.Session()
                m.make_tf_array(free_vars)
                with m.tf_mode():
                    if is_whitened:
                        args = (self.X, self.Z, m.kern, m.q_mu, m.q_sqrt, self.oneLatentFunction)
                        fmean_func, fvar_func = GPflow.conditionals.gaussian_gp_predict_whitened(*args)
                    else:
                        args = (self.X, self.Z, m.kern, m.q_mu, m.q_sqrt, self.oneLatentFunction)
                        fmean_func, fvar_func = GPflow.conditionals.gaussian_gp_predict(*args)
                mean_value = fmean_func.eval(session=session, feed_dict={free_vars: m.get_free_state()})[0, 0]
                var_value = fvar_func.eval(session=session, feed_dict={free_vars: m.get_free_state()})[0, 0]
                self.assertTrue(np.abs(mean_value - self.posteriorMean) < 1e-4)
                self.assertTrue(np.abs(var_value - self.posteriorVariance) < 1e-4)

class VariationalMultivariateTest(unittest.TestCase):
    def setUp(self):
        tf.reset_default_graph()
        self.nDimensions = 3
        self.rng = np.random.RandomState(1)
        self.Y = self.rng.randn(self.nDimensions, 1)
        self.X = self.rng.randn(self.nDimensions, 1)
        self.Z = self.X.copy()
        self.noiseVariance = 0.5
        self.signalVariance = 1.5
        self.lengthScale = 1.7
        self.oneLatentFunction = 1
        self.lik = GPflow.likelihoods.Gaussian()
        self.lik.variance = self.noiseVariance
        self.q_mean = self.rng.randn(self.nDimensions, self.oneLatentFunction)
        self.q_sqrt_diag = self.rng.rand(self.nDimensions, self.oneLatentFunction)
        self.q_sqrt_full = np.tril(self.rng.rand(self.nDimensions, self.nDimensions))

    def getModel(self, is_diagonal, is_whitened):
        model = GPflow.svgp.SVGP(X=self.X, Y=self.Y,
                                 kern=kernel(kernelVariance=self.signalVariance,
                                             lengthScale=self.lengthScale),
                                 likelihood=self.lik, Z=self.Z, q_diag=is_diagonal, whiten=is_whitened)
        if is_diagonal:
            model.q_sqrt = self.q_sqrt_diag
        else:
            model.q_sqrt = self.q_sqrt_full[:, :, None]
        model.q_mu = self.q_mean
        return model

    def test_refrence_implementation_consistency(self):
        rng = np.random.RandomState(10)
        qMean = rng.randn()
        qCov = rng.rand()
        pMean = rng.rand()
        pCov = rng.rand()
        univariate_KL = referenceUnivariatePriorKL(qMean, pMean, qCov, pCov)
        multivariate_KL = referenceMultivariatePriorKL(np.array([[qMean]]), np.array([[qCov]]),
                                                       np.array([[pMean]]), np.array([[pCov]]))
        self.assertTrue(np.abs(univariate_KL - multivariate_KL) < 1e-4)

    def test_prior_KL_fullQ(self):
        covQ = np.dot(self.q_sqrt_full, self.q_sqrt_full.T)
        mean_prior = np.zeros((self.nDimensions, 1))

        for is_whitened in [True, False]:
            m = self.getModel(False, is_whitened)

            if is_whitened:
                cov_prior = np.eye(self.nDimensions)
            else:
                cov_prior = referenceRbfKernel(self.X, self.lengthScale, self.signalVariance)

            referenceKL = referenceMultivariatePriorKL(self.q_mean, covQ, mean_prior, cov_prior)
            # now get test KL.
            test_prior_KL = GPflow.param.AutoFlow()(m.build_prior_KL.__func__)(m)
            self.assertTrue(np.abs(referenceKL - test_prior_KL) < 1e-4)

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