<commit_before>import json
import requests
import tempfile
from django.core import files

from django.contrib.gis.geos import GEOSGeometry
from django.core.exceptions import PermissionDenied
from django.utils.html import strip_tags

from easy_thumbnails.files import get_thumbnailer

from rest_framework import serializers
from rest_framework_gis import serializers as geoserializers

from core.exceptions import MalformedRequestData
from categories.serializer import CategorySerializer
from categories.models import Category, Field
from users.serializers import UserSerializer

from .models import (
    Location, Observation, Comment, MediaFile, ImageFile, VideoFile
)

class LocationSerializer(geoserializers.GeoFeatureModelSerializer):
    class Meta:
        model = Location
        geo_field = 'geometry'
        fields = ('id', 'name', 'description', 'status', 'created_at')
        write_only_fields = ('status',)

class LocationContributionSerializer(serializers.ModelSerializer):
    class Meta:
        model = Location
        depth = 1
        fields = ('id', 'name', 'description', 'status', 'created_at')
        write_only_fields = ('status',)

class ObservationSerializer(serializers.ModelSerializer):
    creator = UserSerializer(fields=('id', 'display_name'))
    updator = UserSerializer(fields=('id', 'display_name'))
    category = serializers.SerializerMethodField('get_category')

    class Meta:
        model = Observation
        depth = 0
        fields = (
            'status', 'category', 'creator', 'updator', 'created_at', 'version'
        )

    def get_category(self, observation):
        return observation.category.id

class ContributionSerializer(object):
    """
    Serializes and deserializes contribution object from and to its GeoJSON
    conterparts.
    """
    def __init__(self, instance=None, data=None, many=False, context=None):
        self.many = many
        self.context = context

        if self.many:
            self.instance = instance
        else:
            self.instance = self.restore_object(instance, data)

    @property
    def data(self):
        if self.many:
            features = [self.to_native_min(obj) for obj in self.instance]
            return {
                "type": "FeatureCollection",
                "features": features
            }

        return self.to_native(self.instance)

    def restore_location(self, instance=None, data=None, geometry=None):
        if instance is not None:
            if data is not None:
                instance.name = data.get('name') or instance.name
                instance.description = (data.get('description') or
                                        instance.description)
                private = data.get('private') or instance.private
                private_for_project = (data.get('private_for_project') or
                                       instance.private_for_project)

            if geometry is not None:
                if type(geometry) is not unicode:
                    geometry = json.dumps(geometry)

                instance.geometry = GEOSGeometry(geometry)

            instance.save()
            return instance
        else:
            if (data is not None) and ('id' in data):
                try:
                    return Location.objects.get_single(
                        self.context.get('user'),
                        self.context.get('project').id,
                        data.get('id')
                    )
                except PermissionDenied, error:
                    raise MalformedRequestData(error)
            else:
                name = None
                description = None
                private_for_project = None
                private = False

                if data is not None:
                    name = strip_tags(data.get('name'))
                    description = strip_tags(data.get('description'))
                    private = data.get('private')
                    private_for_project = data.get('private_for_project')

                return Location(
                    name=name,
                    description=description,
                    geometry=GEOSGeometry(json.dumps(geometry)),
                    private=private,
                    private_for_project=private_for_project,
                    creator=self.context.get('user')
                )

    def restore_object(self, instance=None, data=None):
        if data is not None:
            properties = data.get('properties')
            location = properties.get('location')
            attributes = properties.get('attributes')
            user = self.context.get('user')

            status = properties.pop('status', None)

            if instance is not None:
                self.restore_location(
                    instance.location,
                    data=data.get('properties').pop('location', None),
                    geometry=data.pop('geometry', None)
                )

                return instance.update(
                    attributes=attributes,
                    updator=user,
                    status=status
                )
            else:
                project = self.context.get('project')

                try:
                    category = project.categories.get(
                        pk=properties.pop('category'))
                except Category.DoesNotExist:
                    raise MalformedRequestData('The category can not'
                                               'be used with the project or '
                                               'does not exist.')

                location = self.restore_location(
                    data=data.get('properties').pop('location', None),
                    geometry=data.get('geometry')
                )

                return Observation.create(
                    attributes=attributes,
                    creator=user,
                    location=location,
                    project=category.project,
                    category=category,
                    status=status
                )
        else:
            return instance

    def get_display_field(self, obj):
        if obj.display_field is not None:
            display_field = obj.display_field.split(':', 1)
            return {
                'key': display_field[0],
                'value': display_field[1]
            }
        else:
            return None

    def get_search_result(self, obj, q):
        search_matches = {}

        matcher = obj.search_matches.split('#####')

        for field in matcher:
            match = field.split(':', 1)
            if q.lower() in match[1].lower():
                search_matches[match[0]] = match[1]

        return search_matches

    def to_native_base(self, obj):
        location = obj.location

        isowner = False
        if not self.context.get('user').is_anonymous():
            isowner = obj.creator == self.context.get('user')

        updator = None
        if obj.updator is not None:
            updator = {
                'id': obj.updator.id,
                'display_name': obj.updator.display_name
            }

        json_object = {
            'id': obj.id,
            'type': 'Feature',
            'geometry': json.loads(obj.location.geometry.geojson),
            'properties': {
                'status': obj.status,
                'creator': {
                    'id': obj.creator.id,
                    'display_name': obj.creator.display_name
                },
                'updator': updator,
                'created_at': obj.created_at,
                'version': obj.version,
                'location': {
                    'id': location.id,
                    'name': location.name,
                    'description': location.description
                }
            },
            'isowner': isowner
        }

        return json_object

    def to_native_min(self, obj):
        json_object = self.to_native_base(obj)
        json_object['category'] = {
            'id': obj.category.id,
            'name': obj.category.name,
            'description': obj.category.description,
            'symbol': (obj.category.symbol.url
                       if obj.category.symbol else None),
            'colour': obj.category.colour
        }

        json_object['display_field'] = self.get_display_field(obj)

        q = self.context.get('search')
        if q is not None:
            json_object['search_matches'] = self.get_search_result(obj, q)

        return json_object

    def to_native(self, obj):
        json_object = self.to_native_base(obj)

        category_serializer = CategorySerializer(
            obj.category, context=self.context)
        json_object['category'] = category_serializer.data

        comment_serializer = CommentSerializer(
            obj.comments.filter(respondsto=None),
            many=True,
            context=self.context
        )
        json_object['comments'] = comment_serializer.data

        review_serializer = CommentSerializer(
            obj.comments.filter(review_status='open'),
            many=True,
            context=self.context
        )
        json_object['review_comments'] = review_serializer.data

        file_serializer = FileSerializer(
            obj.files_attached.all(),
            many=True,
            context=self.context
        )
        json_object['media'] = file_serializer.data

        attributes = {}
        for field in obj.category.fields.all():
            value = obj.attributes.get(field.key)
            if value is not None:
                attributes[field.key] = field.convert_from_string(value)

        json_object['properties']['attributes'] = attributes

        return json_object

class CommentSerializer(serializers.ModelSerializer):
    creator = UserSerializer(fields=('id', 'display_name'))
    isowner = serializers.SerializerMethodField('get_is_owner')

    class Meta:
        model = Comment
        fields = ('id', 'respondsto', 'created_at', 'text', 'isowner',
                  'creator', 'review_status')
        read_only = ('id', 'respondsto', 'created_at')

    def to_native(self, obj):
        native = super(CommentSerializer, self).to_native(obj)
        native['responses'] = CommentSerializer(
            obj.responses.all(),
            many=True,
            context=self.context
        ).data

        return native

    def get_is_owner(self, comment):
        if not self.context.get('user').is_anonymous():
            return comment.creator == self.context.get('user')
        else:
            return False

class FileSerializer(serializers.ModelSerializer):
    creator = UserSerializer(fields=('id', 'display_name'))
    isowner = serializers.SerializerMethodField('get_is_owner')
    url = serializers.SerializerMethodField('get_url')
    file_type = serializers.SerializerMethodField('get_type')
    thumbnail_url = serializers.SerializerMethodField('get_thumbnail_url')

    class Meta:
        model = MediaFile
        fields = (
            'id', 'name', 'description', 'created_at', 'creator', 'isowner',
            'url', 'thumbnail_url', 'file_type'
        )

    def get_type(self, obj):
        """
        Returns the type of the MediaFile
        """
        return obj.type_name

    def get_is_owner(self, obj):
        """
        Returns `True` if the user provided in the serializer context is the
        creator of this file
        """
        if not self.context.get('user').is_anonymous():
            return obj.creator == self.context.get('user')
        else:
            return False

    def get_url(self, obj):
        """
        Return the url to access this file based on its file type
        """
        if isinstance(obj, ImageFile):
            return obj.image.url
        elif isinstance(obj, VideoFile):
            return obj.youtube_link

    def _get_thumb(self, image, size=(300, 300)):
        thumbnailer = get_thumbnailer(image)
        thumb = thumbnailer.get_thumbnail({
            'crop': True,
            'size': size
        })
        return thumb

    def get_thumbnail_url(self, obj):
        """
        Creates and returns a thumbnail for the ImageFile object
        """
        if isinstance(obj, ImageFile):
            return self._get_thumb(obj.image).url

        elif isinstance(obj, VideoFile):
            if obj.thumbnail:
                # thumbnail has been downloaded, return the link
                return self._get_thumb(obj.thumbnail).url

            request = requests.get(
                'http://img.youtube.com/vi/%s/0.jpg' % obj.youtube_id,
                stream=True
            )

            if request.status_code != requests.codes.ok:
                # Image not found, return placeholder thumbnail
                return '/static/img/play.png'

            lf = tempfile.NamedTemporaryFile()
            # Read the streamed image in sections
            for block in request.iter_content(1024 * 8):

                # If no more file then stop
                if not block:
                    break

                # Write image block to temporary file
                lf.write(block)

            file_name = obj.youtube_id + '.jpg'
            obj.thumbnail.save(file_name, files.File(lf))

            from PIL import Image

            w, h = Image.open(obj.thumbnail).size

            thumb = self._get_thumb(obj.thumbnail, size=(h, h))
            obj.thumbnail.save(file_name, thumb)

            return self._get_thumb(obj.thumbnail).url
<commit_msg>Fixes #219<commit_after>import json
import requests
import tempfile
from django.core import files

from django.contrib.gis.geos import GEOSGeometry
from django.core.exceptions import PermissionDenied
from django.utils.html import strip_tags

from easy_thumbnails.files import get_thumbnailer

from rest_framework import serializers
from rest_framework_gis import serializers as geoserializers

from core.exceptions import MalformedRequestData
from categories.serializer import CategorySerializer
from categories.models import Category
from users.serializers import UserSerializer

from .models import (
    Location, Observation, Comment, MediaFile, ImageFile, VideoFile
)

class LocationSerializer(geoserializers.GeoFeatureModelSerializer):
    class Meta:
        model = Location
        geo_field = 'geometry'
        fields = ('id', 'name', 'description', 'status', 'created_at')
        write_only_fields = ('status',)

class LocationContributionSerializer(serializers.ModelSerializer):
    class Meta:
        model = Location
        depth = 1
        fields = ('id', 'name', 'description', 'status', 'created_at')
        write_only_fields = ('status',)

class ObservationSerializer(serializers.ModelSerializer):
    creator = UserSerializer(fields=('id', 'display_name'))
    updator = UserSerializer(fields=('id', 'display_name'))
    category = serializers.SerializerMethodField('get_category')

    class Meta:
        model = Observation
        depth = 0
        fields = (
            'status', 'category', 'creator', 'updator', 'created_at', 'version'
        )

    def get_category(self, observation):
        return observation.category.id

class ContributionSerializer(object):
    """
    Serializes and deserializes contribution object from and to its GeoJSON
    conterparts.
    """
    def __init__(self, instance=None, data=None, many=False, context=None):
        self.many = many
        self.context = context

        if self.many:
            self.instance = instance
        else:
            self.instance = self.restore_object(instance, data)

    @property
    def data(self):
        if self.many:
            features = [self.to_native_min(obj) for obj in self.instance]
            return {
                "type": "FeatureCollection",
                "features": features
            }

        return self.to_native(self.instance)

    def restore_location(self, instance=None, data=None, geometry=None):
        if instance is not None:
            if data is not None:
                instance.name = data.get('name') or instance.name
                instance.description = (data.get('description') or
                                        instance.description)
                private = data.get('private') or instance.private
                private_for_project = (data.get('private_for_project') or
                                       instance.private_for_project)

            if geometry is not None:
                if type(geometry) is not unicode:
                    geometry = json.dumps(geometry)

                instance.geometry = GEOSGeometry(geometry)

            instance.save()
            return instance
        else:
            if (data is not None) and ('id' in data):
                try:
                    return Location.objects.get_single(
                        self.context.get('user'),
                        self.context.get('project').id,
                        data.get('id')
                    )
                except PermissionDenied, error:
                    raise MalformedRequestData(error)
            else:
                name = None
                description = None
                private_for_project = None
                private = False

                if data is not None:
                    name = strip_tags(data.get('name'))
                    description = strip_tags(data.get('description'))
                    private = data.get('private')
                    private_for_project = data.get('private_for_project')

                return Location(
                    name=name,
                    description=description,
                    geometry=GEOSGeometry(json.dumps(geometry)),
                    private=private,
                    private_for_project=private_for_project,
                    creator=self.context.get('user')
                )

    def restore_object(self, instance=None, data=None):
        if data is not None:
            properties = data.get('properties')
            location = properties.get('location')
            attributes = properties.get('attributes')
            user = self.context.get('user')

            status = properties.pop('status', None)

            if instance is not None:
                self.restore_location(
                    instance.location,
                    data=data.get('properties').pop('location', None),
                    geometry=data.pop('geometry', None)
                )

                return instance.update(
                    attributes=attributes,
                    updator=user,
                    status=status
                )
            else:
                project = self.context.get('project')

                try:
                    category = project.categories.get(
                        pk=properties.pop('category'))
                except Category.DoesNotExist:
                    raise MalformedRequestData('The category can not'
                                               'be used with the project or '
                                               'does not exist.')

                location = self.restore_location(
                    data=data.get('properties').pop('location', None),
                    geometry=data.get('geometry')
                )

                return Observation.create(
                    attributes=attributes,
                    creator=user,
                    location=location,
                    project=category.project,
                    category=category,
                    status=status
                )
        else:
            return instance

    def get_display_field(self, obj):
        if obj.display_field is not None:
            display_field = obj.display_field.split(':', 1)
            value = display_field[1] if display_field[1] != 'None' else None
            return {
                'key': display_field[0],
                'value': value
            }
        else:
            return None

    def get_search_result(self, obj, q):
        search_matches = {}

        matcher = obj.search_matches.split('#####')

        for field in matcher:
            match = field.split(':', 1)
            if q.lower() in match[1].lower():
                search_matches[match[0]] = match[1]

        return search_matches

    def to_native_base(self, obj):
        location = obj.location

        isowner = False
        if not self.context.get('user').is_anonymous():
            isowner = obj.creator == self.context.get('user')

        updator = None
        if obj.updator is not None:
            updator = {
                'id': obj.updator.id,
                'display_name': obj.updator.display_name
            }

        json_object = {
            'id': obj.id,
            'type': 'Feature',
            'geometry': json.loads(obj.location.geometry.geojson),
            'properties': {
                'status': obj.status,
                'creator': {
                    'id': obj.creator.id,
                    'display_name': obj.creator.display_name
                },
                'updator': updator,
                'created_at': obj.created_at,
                'version': obj.version,
                'location': {
                    'id': location.id,
                    'name': location.name,
                    'description': location.description
                }
            },
            'isowner': isowner
        }

        return json_object

    def to_native_min(self, obj):
        json_object = self.to_native_base(obj)
        json_object['category'] = {
            'id': obj.category.id,
            'name': obj.category.name,
            'description': obj.category.description,
            'symbol': (obj.category.symbol.url
                       if obj.category.symbol else None),
            'colour': obj.category.colour
        }

        json_object['display_field'] = self.get_display_field(obj)

        q = self.context.get('search')
        if q is not None:
            json_object['search_matches'] = self.get_search_result(obj, q)

        return json_object

    def to_native(self, obj):
        json_object = self.to_native_base(obj)

        category_serializer = CategorySerializer(
            obj.category, context=self.context)
        json_object['category'] = category_serializer.data

        comment_serializer = CommentSerializer(
            obj.comments.filter(respondsto=None),
            many=True,
            context=self.context
        )
        json_object['comments'] = comment_serializer.data

        review_serializer = CommentSerializer(
            obj.comments.filter(review_status='open'),
            many=True,
            context=self.context
        )
        json_object['review_comments'] = review_serializer.data

        file_serializer = FileSerializer(
            obj.files_attached.all(),
            many=True,
            context=self.context
        )
        json_object['media'] = file_serializer.data

        attributes = {}
        for field in obj.category.fields.all():
            value = obj.attributes.get(field.key)
            if value is not None:
                attributes[field.key] = field.convert_from_string(value)

        json_object['properties']['attributes'] = attributes

        return json_object

class CommentSerializer(serializers.ModelSerializer):
    creator = UserSerializer(fields=('id', 'display_name'))
    isowner = serializers.SerializerMethodField('get_is_owner')

    class Meta:
        model = Comment
        fields = ('id', 'respondsto', 'created_at', 'text', 'isowner',
                  'creator', 'review_status')
        read_only = ('id', 'respondsto', 'created_at')

    def to_native(self, obj):
        native = super(CommentSerializer, self).to_native(obj)
        native['responses'] = CommentSerializer(
            obj.responses.all(),
            many=True,
            context=self.context
        ).data

        return native

    def get_is_owner(self, comment):
        if not self.context.get('user').is_anonymous():
            return comment.creator == self.context.get('user')
        else:
            return False

class FileSerializer(serializers.ModelSerializer):
    creator = UserSerializer(fields=('id', 'display_name'))
    isowner = serializers.SerializerMethodField('get_is_owner')
    url = serializers.SerializerMethodField('get_url')
    file_type = serializers.SerializerMethodField('get_type')
    thumbnail_url = serializers.SerializerMethodField('get_thumbnail_url')

    class Meta:
        model = MediaFile
        fields = (
            'id', 'name', 'description', 'created_at', 'creator', 'isowner',
            'url', 'thumbnail_url', 'file_type'
        )

    def get_type(self, obj):
        """
        Returns the type of the MediaFile
        """
        return obj.type_name

    def get_is_owner(self, obj):
        """
        Returns `True` if the user provided in the serializer context is the
        creator of this file
        """
        if not self.context.get('user').is_anonymous():
            return obj.creator == self.context.get('user')
        else:
            return False

    def get_url(self, obj):
        """
        Return the url to access this file based on its file type
        """
        if isinstance(obj, ImageFile):
            return obj.image.url
        elif isinstance(obj, VideoFile):
            return obj.youtube_link

    def _get_thumb(self, image, size=(300, 300)):
        thumbnailer = get_thumbnailer(image)
        thumb = thumbnailer.get_thumbnail({
            'crop': True,
            'size': size
        })
        return thumb

    def get_thumbnail_url(self, obj):
        """
        Creates and returns a thumbnail for the ImageFile object
        """
        if isinstance(obj, ImageFile):
            return self._get_thumb(obj.image).url

        elif isinstance(obj, VideoFile):
            if obj.thumbnail:
                # thumbnail has been downloaded, return the link
                return self._get_thumb(obj.thumbnail).url

            request = requests.get(
                'http://img.youtube.com/vi/%s/0.jpg' % obj.youtube_id,
                stream=True
            )

            if request.status_code != requests.codes.ok:
                # Image not found, return placeholder thumbnail
                return '/static/img/play.png'

            lf = tempfile.NamedTemporaryFile()
            # Read the streamed image in sections
            for block in request.iter_content(1024 * 8):

                # If no more file then stop
                if not block:
                    break

                # Write image block to temporary file
                lf.write(block)

            file_name = obj.youtube_id + '.jpg'
            obj.thumbnail.save(file_name, files.File(lf))

            from PIL import Image

            w, h = Image.open(obj.thumbnail).size

            thumb = self._get_thumb(obj.thumbnail, size=(h, h))
            obj.thumbnail.save(file_name, thumb)

            return self._get_thumb(obj.thumbnail).url
<|endoftext|>


<commit_before>from __future__ import annotations

from dataclasses import dataclass, replace
from math import pi
from typing import Callable, List, Optional, Sequence, Tuple, Union

from cairo import Context as CairoContext
from gaphas.geometry import Rectangle
from typing_extensions import TypedDict

from gaphor.diagram.style import Style, combined_style
from gaphor.diagram.text import (
    FontStyle,
    FontWeight,
    TextAlign,
    TextDecoration,
    VerticalAlign,
    focus_box_pos,
    text_draw,
    text_size,
)

class Padding:  # Enum
    TOP = 0
    RIGHT = 1
    BOTTOM = 2
    LEFT = 3

@dataclass(frozen=True)
class SizeContext:
    @classmethod
    def from_context(cls, context, style) -> SizeContext:
        return SizeContext(cairo=context.cairo, style=style)

    cairo: CairoContext
    style: Style

@dataclass(frozen=True)
class DrawContext:
    """
    Special context for draw()'ing the item. The draw-context contains
    stuff like the cairo context and flags like selected and
    focused.
    """

    @classmethod
    def from_context(cls, context, style) -> DrawContext:
        return DrawContext(
            cairo=context.cairo,
            style=style,
            selected=context.selected,
            focused=context.focused,
            hovered=context.hovered,
            dropzone=context.dropzone,
        )

    cairo: CairoContext
    style: Style
    selected: bool
    focused: bool
    hovered: bool
    dropzone: bool

class cairo_state:
    def __init__(self, cr):
        self._cr = cr

    def __enter__(self):
        self._cr.save()
        return self._cr

    def __exit__(self, _type, _value, _traceback):
        self._cr.restore()

def stroke(context: DrawContext, fill=True, highlight=False):
    style = context.style
    cr = context.cairo
    fill_color = style.get("background-color")
    if fill and fill_color:
        with cairo_state(cr):
            cr.set_source_rgba(*fill_color)
            cr.fill_preserve()

    if highlight:
        draw_highlight(context)

    with cairo_state(cr):
        stroke = style.get("color")
        if stroke:
            cr.set_source_rgba(*stroke)
        line_width = style.get("line-width")
        if line_width:
            cr.set_line_width(line_width)
        cr.stroke()

def draw_border(box, context: DrawContext, bounding_box: Rectangle):
    cr = context.cairo
    d = context.style["border-radius"]
    x, y, width, height = bounding_box

    cr.move_to(x, d)
    cr.set_dash(context.style["dash-style"], 0)
    if d:
        x1 = width + x
        y1 = height + y
        cr.arc(d, d, d, pi, 1.5 * pi)
        cr.line_to(x1 - d, y)
        cr.arc(x1 - d, d, d, 1.5 * pi, y)
        cr.line_to(x1, y1 - d)
        cr.arc(x1 - d, y1 - d, d, 0, 0.5 * pi)
        cr.line_to(d, y1)
        cr.arc(d, y1 - d, d, 0.5 * pi, pi)
    else:
        cr.rectangle(x, y, width, height)

    cr.close_path()

    stroke(context)

def draw_top_separator(box: Box, context: DrawContext, bounding_box: Rectangle):
    x, y, w, h = bounding_box
    cr = context.cairo
    cr.move_to(x, y)
    cr.line_to(x + w, y)

    stroke(context, fill=False)

def draw_highlight(context: DrawContext):
    if not context.dropzone:
        return
    with cairo_state(context.cairo) as cr:
        highlight_color = context.style["highlight-color"]
        cr.set_source_rgba(*highlight_color)
        cr.set_line_width(cr.get_line_width() * 3.141)
        cr.stroke_preserve()

class Box:
    """
    A box like shape.

    Style properties:
    - min-height
    - min-width
    - padding: a tuple (top, right, bottom, left)
    - vertical-align: alignment of child shapes
    - border-radius
    """

    def __init__(
        self,
        *children,
        style: Style = {},
        draw: Optional[Callable[[Box, DrawContext, Rectangle], None]] = None
    ):
        self.children = children
        self.sizes: List[Tuple[int, int]] = []
        self._inline_style = style
        self._draw_border = draw

    def __len__(self):
        return len(self.children)

    def __iter__(self):
        return iter(self.children)

    def __getitem__(self, index):
        return self.children[index]

    def size(self, context: SizeContext):
        style: Style = combined_style(context.style, self._inline_style)
        min_width = style["min-width"]
        min_height = style["min-height"]
        padding = style["padding"]
        self.sizes = sizes = [c.size(context) for c in self.children]
        if sizes:
            widths, heights = list(zip(*sizes))
            return (
                max(
                    min_width,
                    max(widths) + padding[Padding.RIGHT] + padding[Padding.LEFT],
                ),
                max(
                    min_height,
                    sum(heights) + padding[Padding.TOP] + padding[Padding.BOTTOM],
                ),
            )
        else:
            return min_width, min_height

    def draw(self, context: DrawContext, bounding_box: Rectangle):
        style: Style = combined_style(context.style, self._inline_style)
        new_context = replace(context, style=style)
        padding = style["padding"]
        valign = style["vertical-align"]
        height = sum(h for _w, h in self.sizes)

        if self._draw_border:
            self._draw_border(self, new_context, bounding_box)
        x = bounding_box.x + padding[Padding.LEFT]
        if valign is VerticalAlign.MIDDLE:
            y = (
                bounding_box.y
                + padding[Padding.TOP]
                + (max(height, bounding_box.height - padding[Padding.TOP]) - height) / 2
            )
        elif valign is VerticalAlign.BOTTOM:
            y = bounding_box.y + bounding_box.height - height - padding[Padding.BOTTOM]
        else:
            y = bounding_box.y + padding[Padding.TOP]
        w = bounding_box.width - padding[Padding.RIGHT] - padding[Padding.LEFT]
        for c, (_w, h) in zip(self.children, self.sizes):
            c.draw(context, Rectangle(x, y, w, h))
            y += h

class IconBox:
    """
    A special type of box: the icon element is given the full width/height and
    all other shapes are drawn below the main icon shape.

    Style properties:
    - min-height
    - min-width
    - vertical-spacing: spacing between icon and children
    - padding: a tuple (top, right, bottom, left)
    """

    def __init__(self, icon, *children, style: Style = {}):
        self.icon = icon
        self.children = children
        self.sizes: List[Tuple[int, int]] = []
        self._inline_style = style

    def size(self, context: SizeContext):
        style = combined_style(context.style, self._inline_style)
        min_width = style["min-width"]
        min_height = style["min-height"]
        padding = style["padding"]
        self.sizes = [c.size(context) for c in self.children]
        width, height = self.icon.size(context)
        return (
            max(min_width, width + padding[Padding.RIGHT] + padding[Padding.LEFT]),
            max(min_height, height + padding[Padding.TOP] + padding[Padding.BOTTOM]),
        )

    def draw(self, context: DrawContext, bounding_box: Rectangle):
        style = combined_style(context.style, self._inline_style)
        new_context = replace(context, style=style)
        padding = style["padding"]
        vertical_spacing = style["vertical-spacing"]
        x = bounding_box.x + padding[Padding.LEFT]
        y = bounding_box.y + padding[Padding.TOP]
        w = bounding_box.width - padding[Padding.RIGHT] - padding[Padding.LEFT]
        h = bounding_box.height - padding[Padding.TOP] - padding[Padding.BOTTOM]
        self.icon.draw(new_context, Rectangle(x, y, w, h))
        y = y + bounding_box.height + vertical_spacing
        for c, (cw, ch) in zip(self.children, self.sizes):
            mw = max(w, cw)
            c.draw(context, Rectangle(x - (mw - w) / 2, y, mw, ch))
            y += ch

class Text:
    def __init__(self, text=lambda: "", width=lambda: -1, style: Style = {}):
        self._text = text if callable(text) else lambda: text
        self.width = width if callable(width) else lambda: width
        self._inline_style = style

    def text(self):
        try:
            return self._text()
        except AttributeError:
            return ""

    def size(self, context: SizeContext):
        style = combined_style(context.style, self._inline_style)
        min_w = style["min-width"]
        min_h = style["min-height"]
        padding = style["padding"]

        width, height = text_size(context.cairo, self.text(), style, self.width())  # type: ignore[type-var]
        return (
            max(min_w, width + padding[Padding.RIGHT] + padding[Padding.LEFT]),
            max(min_h, height + padding[Padding.TOP] + padding[Padding.BOTTOM]),
        )

    def text_box(self, style: Style, bounding_box: Rectangle) -> Rectangle:
        """Add padding to a bounding box."""
        padding = style["padding"]
        return Rectangle(
            bounding_box.x + padding[Padding.LEFT],
            bounding_box.y + padding[Padding.TOP],
            bounding_box.width - padding[Padding.RIGHT] - padding[Padding.LEFT],
            bounding_box.height - padding[Padding.TOP] - padding[Padding.BOTTOM],
        )

    def draw(
        self, context: DrawContext, bounding_box: Rectangle
    ) -> Tuple[int, int, int, int]:
        """Draw the text, return the location and size."""
        style = combined_style(context.style, self._inline_style)
        min_w = max(style["min-width"], bounding_box.width)
        min_h = max(style["min-height"], bounding_box.height)
        text_align = style["text-align"]
        text_box = self.text_box(style, bounding_box)

        with cairo_state(context.cairo) as cr:
            text_color = style.get("text-color")
            if text_color:
                cr.set_source_rgba(*text_color)

            x, y, w, h = text_draw(
                cr,
                self.text(),
                style,
                lambda w, h: (bounding_box.x, bounding_box.y),
                width=text_box.width,
                default_size=(min_w, min_h),
                text_align=text_align,
            )
        return x, y, w, h

class EditableText(Text):
    def __init__(self, text=lambda: "", width=lambda: -1, style: Style = {}):
        super().__init__(text, width, {"min-width": 30, "min-height": 14, **style})  # type: ignore[misc]
        self.bounding_box = Rectangle()
        self.text_size = (0, 0)

    def size(self, cr):
        s = super().size(cr)
        self.text_size = s
        return s

    def draw(
        self, context: DrawContext, bounding_box: Rectangle
    ) -> Tuple[int, int, int, int]:
        """Draw the editable text."""
        style = combined_style(context.style, self._inline_style)
        x, y, w, h = super().draw(context, bounding_box)
        text_box = self.text_box(style, bounding_box)
        text_align = style["text-align"]
        vertical_align = style["vertical-align"]
        x, y = focus_box_pos(text_box, self.text_size, text_align, vertical_align)
        text_draw_focus_box(context, x, y, *self.text_size)
        self.bounding_box = Rectangle(x, y, width=w, height=h)
        return x, y, w, h

def draw_default_head(context: DrawContext):
    """
    Default head drawer: move cursor to the first handle.
    """
    context.cairo.move_to(0, 0)

def draw_default_tail(context: DrawContext):
    """
    Default tail drawer: draw line to the last handle.
    """
    context.cairo.line_to(0, 0)

def draw_arrow_head(context: DrawContext):
    cr = context.cairo
    cr.set_dash((), 0)
    cr.move_to(15, -6)
    cr.line_to(0, 0)
    cr.line_to(15, 6)
    cr.stroke()
    cr.move_to(0, 0)

def draw_arrow_tail(context: DrawContext):
    cr = context.cairo
    cr.line_to(0, 0)
    cr.stroke()
    cr.move_to(15, -6)
    cr.line_to(0, 0)
    cr.line_to(15, 6)

def text_draw_focus_box(context, x, y, w, h):
    if context.hovered or context.focused:
        with cairo_state(context.cairo) as cr:
            # cr.set_dash(() if context.focused else (2.0, 2.0), 0)
            cr.set_dash((), 0)
            if context.focused:
                cr.set_source_rgb(0.6, 0.6, 0.6)
            else:
                cr.set_source_rgb(0.8, 0.8, 0.8)
            cr.set_line_width(0.5)
            cr.rectangle(x, y, w, h)
            cr.stroke()
<commit_msg>Make use of padding simpler in shapes<commit_after>from __future__ import annotations

from dataclasses import dataclass, replace
from math import pi
from typing import Callable, List, Optional, Sequence, Tuple, Union

from cairo import Context as CairoContext
from gaphas.geometry import Rectangle
from typing_extensions import TypedDict

from gaphor.diagram.style import Style, combined_style
from gaphor.diagram.text import (
    FontStyle,
    FontWeight,
    TextAlign,
    TextDecoration,
    VerticalAlign,
    focus_box_pos,
    text_draw,
    text_size,
)

@dataclass(frozen=True)
class SizeContext:
    @classmethod
    def from_context(cls, context, style) -> SizeContext:
        return SizeContext(cairo=context.cairo, style=style)

    cairo: CairoContext
    style: Style

@dataclass(frozen=True)
class DrawContext:
    """
    Special context for draw()'ing the item. The draw-context contains
    stuff like the cairo context and flags like selected and
    focused.
    """

    @classmethod
    def from_context(cls, context, style) -> DrawContext:
        return DrawContext(
            cairo=context.cairo,
            style=style,
            selected=context.selected,
            focused=context.focused,
            hovered=context.hovered,
            dropzone=context.dropzone,
        )

    cairo: CairoContext
    style: Style
    selected: bool
    focused: bool
    hovered: bool
    dropzone: bool

class cairo_state:
    def __init__(self, cr):
        self._cr = cr

    def __enter__(self):
        self._cr.save()
        return self._cr

    def __exit__(self, _type, _value, _traceback):
        self._cr.restore()

def stroke(context: DrawContext, fill=True, highlight=False):
    style = context.style
    cr = context.cairo
    fill_color = style.get("background-color")
    if fill and fill_color:
        with cairo_state(cr):
            cr.set_source_rgba(*fill_color)
            cr.fill_preserve()

    if highlight:
        draw_highlight(context)

    with cairo_state(cr):
        stroke = style.get("color")
        if stroke:
            cr.set_source_rgba(*stroke)
        line_width = style.get("line-width")
        if line_width:
            cr.set_line_width(line_width)
        cr.stroke()

def draw_border(box, context: DrawContext, bounding_box: Rectangle):
    cr = context.cairo
    d = context.style["border-radius"]
    x, y, width, height = bounding_box

    cr.move_to(x, d)
    cr.set_dash(context.style["dash-style"], 0)
    if d:
        x1 = width + x
        y1 = height + y
        cr.arc(d, d, d, pi, 1.5 * pi)
        cr.line_to(x1 - d, y)
        cr.arc(x1 - d, d, d, 1.5 * pi, y)
        cr.line_to(x1, y1 - d)
        cr.arc(x1 - d, y1 - d, d, 0, 0.5 * pi)
        cr.line_to(d, y1)
        cr.arc(d, y1 - d, d, 0.5 * pi, pi)
    else:
        cr.rectangle(x, y, width, height)

    cr.close_path()

    stroke(context)

def draw_top_separator(box: Box, context: DrawContext, bounding_box: Rectangle):
    x, y, w, h = bounding_box
    cr = context.cairo
    cr.move_to(x, y)
    cr.line_to(x + w, y)

    stroke(context, fill=False)

def draw_highlight(context: DrawContext):
    if not context.dropzone:
        return
    with cairo_state(context.cairo) as cr:
        highlight_color = context.style["highlight-color"]
        cr.set_source_rgba(*highlight_color)
        cr.set_line_width(cr.get_line_width() * 3.141)
        cr.stroke_preserve()

class Box:
    """
    A box like shape.

    Style properties:
    - min-height
    - min-width
    - padding: a tuple (top, right, bottom, left)
    - vertical-align: alignment of child shapes
    - border-radius
    """

    def __init__(
        self,
        *children,
        style: Style = {},
        draw: Optional[Callable[[Box, DrawContext, Rectangle], None]] = None
    ):
        self.children = children
        self.sizes: List[Tuple[int, int]] = []
        self._inline_style = style
        self._draw_border = draw

    def __len__(self):
        return len(self.children)

    def __iter__(self):
        return iter(self.children)

    def __getitem__(self, index):
        return self.children[index]

    def size(self, context: SizeContext):
        style: Style = combined_style(context.style, self._inline_style)
        min_width = style["min-width"]
        min_height = style["min-height"]
        padding_top, padding_right, padding_bottom, padding_left = style["padding"]
        self.sizes = sizes = [c.size(context) for c in self.children]
        if sizes:
            widths, heights = list(zip(*sizes))
            return (
                max(min_width, max(widths) + padding_right + padding_left,),
                max(min_height, sum(heights) + padding_top + padding_bottom,),
            )
        else:
            return min_width, min_height

    def draw(self, context: DrawContext, bounding_box: Rectangle):
        style: Style = combined_style(context.style, self._inline_style)
        new_context = replace(context, style=style)
        padding_top, padding_right, padding_bottom, padding_left = style["padding"]
        valign = style["vertical-align"]
        height = sum(h for _w, h in self.sizes)

        if self._draw_border:
            self._draw_border(self, new_context, bounding_box)
        x = bounding_box.x + padding_left
        if valign is VerticalAlign.MIDDLE:
            y = (
                bounding_box.y
                + padding_top
                + (max(height, bounding_box.height - padding_top) - height) / 2
            )
        elif valign is VerticalAlign.BOTTOM:
            y = bounding_box.y + bounding_box.height - height - padding_bottom
        else:
            y = bounding_box.y + padding_top
        w = bounding_box.width - padding_right - padding_left
        for c, (_w, h) in zip(self.children, self.sizes):
            c.draw(context, Rectangle(x, y, w, h))
            y += h

class IconBox:
    """
    A special type of box: the icon element is given the full width/height and
    all other shapes are drawn below the main icon shape.

    Style properties:
    - min-height
    - min-width
    - vertical-spacing: spacing between icon and children
    - padding: a tuple (top, right, bottom, left)
    """

    def __init__(self, icon, *children, style: Style = {}):
        self.icon = icon
        self.children = children
        self.sizes: List[Tuple[int, int]] = []
        self._inline_style = style

    def size(self, context: SizeContext):
        style = combined_style(context.style, self._inline_style)
        min_width = style["min-width"]
        min_height = style["min-height"]
        padding_top, padding_right, padding_bottom, padding_left = style["padding"]
        self.sizes = [c.size(context) for c in self.children]
        width, height = self.icon.size(context)
        return (
            max(min_width, width + padding_right + padding_left),
            max(min_height, height + padding_top + padding_bottom),
        )

    def draw(self, context: DrawContext, bounding_box: Rectangle):
        style = combined_style(context.style, self._inline_style)
        new_context = replace(context, style=style)
        padding_top, padding_right, padding_bottom, padding_left = style["padding"]
        vertical_spacing = style["vertical-spacing"]
        x = bounding_box.x + padding_left
        y = bounding_box.y + padding_top
        w = bounding_box.width - padding_right - padding_left
        h = bounding_box.height - padding_top - padding_bottom
        self.icon.draw(new_context, Rectangle(x, y, w, h))
        y = y + bounding_box.height + vertical_spacing
        for c, (cw, ch) in zip(self.children, self.sizes):
            mw = max(w, cw)
            c.draw(context, Rectangle(x - (mw - w) / 2, y, mw, ch))
            y += ch

class Text:
    def __init__(self, text=lambda: "", width=lambda: -1, style: Style = {}):
        self._text = text if callable(text) else lambda: text
        self.width = width if callable(width) else lambda: width
        self._inline_style = style

    def text(self):
        try:
            return self._text()
        except AttributeError:
            return ""

    def size(self, context: SizeContext):
        style = combined_style(context.style, self._inline_style)
        min_w = style["min-width"]
        min_h = style["min-height"]
        padding_top, padding_right, padding_bottom, padding_left = style["padding"]

        width, height = text_size(context.cairo, self.text(), style, self.width())  # type: ignore[type-var]
        return (
            max(min_w, width + padding_right + padding_left),
            max(min_h, height + padding_top + padding_bottom),
        )

    def text_box(self, style: Style, bounding_box: Rectangle) -> Rectangle:
        """Add padding to a bounding box."""
        padding_top, padding_right, padding_bottom, padding_left = style["padding"]
        return Rectangle(
            bounding_box.x + padding_left,
            bounding_box.y + padding_top,
            bounding_box.width - padding_right - padding_left,
            bounding_box.height - padding_top - padding_bottom,
        )

    def draw(
        self, context: DrawContext, bounding_box: Rectangle
    ) -> Tuple[int, int, int, int]:
        """Draw the text, return the location and size."""
        style = combined_style(context.style, self._inline_style)
        min_w = max(style["min-width"], bounding_box.width)
        min_h = max(style["min-height"], bounding_box.height)
        text_align = style["text-align"]
        text_box = self.text_box(style, bounding_box)

        with cairo_state(context.cairo) as cr:
            text_color = style.get("text-color")
            if text_color:
                cr.set_source_rgba(*text_color)

            x, y, w, h = text_draw(
                cr,
                self.text(),
                style,
                lambda w, h: (bounding_box.x, bounding_box.y),
                width=text_box.width,
                default_size=(min_w, min_h),
                text_align=text_align,
            )
        return x, y, w, h

class EditableText(Text):
    def __init__(self, text=lambda: "", width=lambda: -1, style: Style = {}):
        super().__init__(text, width, {"min-width": 30, "min-height": 14, **style})  # type: ignore[misc]
        self.bounding_box = Rectangle()
        self.text_size = (0, 0)

    def size(self, cr):
        s = super().size(cr)
        self.text_size = s
        return s

    def draw(
        self, context: DrawContext, bounding_box: Rectangle
    ) -> Tuple[int, int, int, int]:
        """Draw the editable text."""
        style = combined_style(context.style, self._inline_style)
        x, y, w, h = super().draw(context, bounding_box)
        text_box = self.text_box(style, bounding_box)
        text_align = style["text-align"]
        vertical_align = style["vertical-align"]
        x, y = focus_box_pos(text_box, self.text_size, text_align, vertical_align)
        text_draw_focus_box(context, x, y, *self.text_size)
        self.bounding_box = Rectangle(x, y, width=w, height=h)
        return x, y, w, h

def draw_default_head(context: DrawContext):
    """
    Default head drawer: move cursor to the first handle.
    """
    context.cairo.move_to(0, 0)

def draw_default_tail(context: DrawContext):
    """
    Default tail drawer: draw line to the last handle.
    """
    context.cairo.line_to(0, 0)

def draw_arrow_head(context: DrawContext):
    cr = context.cairo
    cr.set_dash((), 0)
    cr.move_to(15, -6)
    cr.line_to(0, 0)
    cr.line_to(15, 6)
    cr.stroke()
    cr.move_to(0, 0)

def draw_arrow_tail(context: DrawContext):
    cr = context.cairo
    cr.line_to(0, 0)
    cr.stroke()
    cr.move_to(15, -6)
    cr.line_to(0, 0)
    cr.line_to(15, 6)

def text_draw_focus_box(context, x, y, w, h):
    if context.hovered or context.focused:
        with cairo_state(context.cairo) as cr:
            # cr.set_dash(() if context.focused else (2.0, 2.0), 0)
            cr.set_dash((), 0)
            if context.focused:
                cr.set_source_rgb(0.6, 0.6, 0.6)
            else:
                cr.set_source_rgb(0.8, 0.8, 0.8)
            cr.set_line_width(0.5)
            cr.rectangle(x, y, w, h)
            cr.stroke()
<|endoftext|>


<commit_before>import errno
import json
import os
import socket
import struct
from subprocess import call

from flask import Flask, request

ALLOWED_OWNERS = set(['mikeboers', 'FluentImage'])

def ip_to_int(ip):
    "convert decimal dotted quad string to long integer"
    if isinstance(ip, int):
        return ip
    return struct.unpack('!I', socket.inet_aton(ip))[0]

def network_mask(ip, bits):
    "Convert a network address to a long integer"
    return ip_to_int(ip) & ((2 << bits - 1) - 1)

def addr_in_network(ip, net):
    "Is an address in a network"
    return ip_to_int(ip) & net == net

GITHUB_NETMASK = network_mask('192.30.252.0', 22)

app = Flask(__name__)
@app.route('/', methods=['GET', 'POST'])
def main():

    if not app.debug and not addr_in_network(request.remote_addr, GITHUB_NETMASK):
        return 'You are not GitHub.'

    if request.method != 'POST':
        return 'must POST'

    payload = json.loads(request.form['payload'])
    owner_name = payload['repository']['owner']['name']
    repo_name = payload['repository']['name']

    if owner_name not in ALLOWED_OWNERS:
        return 'owner not permitted'

    local_path = os.path.abspath(os.path.join(
        __file__, '..', 'var', 'repositories', owner_name, repo_name + '.git'
    ))
    remote_url = 'git@github.com:%s/%s.git' % (owner_name, repo_name)

    if not os.path.exists(local_path):
        os.makedirs(local_path)
        call(['git', 'clone', '--mirror', remote_url, local_path])
    else:
        call(['git', '--git-dir', local_path, 'remote', 'set-url', remote_url])
        call(['git', '--git-dir', local_path, 'fetch', 'origin'])

    return 'ok'

app.run(host='', debug=bool(os.environ.get('DEBUG')), port=int(os.environ.get('PORT', 5000)))

<commit_msg>Assume that origin won't change<commit_after>import errno
import json
import os
import socket
import struct
from subprocess import call

from flask import Flask, request

ALLOWED_OWNERS = set(['mikeboers', 'FluentImage'])

def ip_to_int(ip):
    "convert decimal dotted quad string to long integer"
    if isinstance(ip, int):
        return ip
    return struct.unpack('!I', socket.inet_aton(ip))[0]

def network_mask(ip, bits):
    "Convert a network address to a long integer"
    return ip_to_int(ip) & ((2 << bits - 1) - 1)

def addr_in_network(ip, net):
    "Is an address in a network"
    return ip_to_int(ip) & net == net

GITHUB_NETMASK = network_mask('192.30.252.0', 22)

app = Flask(__name__)
@app.route('/', methods=['GET', 'POST'])
def main():

    if not app.debug and not addr_in_network(request.remote_addr, GITHUB_NETMASK):
        return 'You are not GitHub.'

    if request.method != 'POST':
        return 'must POST'

    payload = json.loads(request.form['payload'])
    owner_name = payload['repository']['owner']['name']
    repo_name = payload['repository']['name']

    if owner_name not in ALLOWED_OWNERS:
        return 'owner not permitted'

    local_path = os.path.abspath(os.path.join(
        __file__, '..', 'var', 'repositories', owner_name, repo_name + '.git'
    ))
    remote_url = 'git@github.com:%s/%s.git' % (owner_name, repo_name)

    if not os.path.exists(local_path):
        os.makedirs(local_path)
        call(['git', 'clone', '--mirror', remote_url, local_path])
    else:
        call(['git', '--git-dir', local_path, 'fetch', 'origin'])

    return 'ok'

app.run(host='', debug=bool(os.environ.get('DEBUG')), port=int(os.environ.get('PORT', 5000)))

<|endoftext|>


<commit_before># coding: utf-8

""" Potential used in Price-Whelan et al. (in prep.) TODO """

from __future__ import division, print_function

__author__ = "adrn <adrn@astro.columbia.edu>"

# Third-party
from astropy.constants import G
import astropy.units as u
import numpy as np

# Project
from .cpotential import CCompositePotential
from .cbuiltin import HernquistPotential, MiyamotoNagaiPotential, \
    LeeSutoTriaxialNFWPotential, SphericalNFWPotential, LogarithmicPotential
from ..units import galactic

__all__ = ['PW14Potential', 'LM10Potential', 'OblateMWPotential']

class PW14Potential(CCompositePotential):

    def __init__(self, m_disk=6.5E10, a=6.5, b=0.26,
                 m_spher=2E10, c=0.3,
                 q1=1.4, q2=1., q3=0.6,
                 v_c=0.247, r_s=30.,
                 phi=np.pi/2., theta=np.pi/2., psi=np.pi/2.,
                 units=galactic):

        # Choice of v_h sets circular velocity at Sun to 220 km/s
        self.units = units

        kwargs = dict()
        kwargs["disk"] = MiyamotoNagaiPotential(units=units,
                                                m=m_disk, a=a, b=b)

        kwargs["bulge"] = HernquistPotential(units=units,
                                             m=m_spher, c=c)

        if q1 == 1 and q2 == 1 and q3 == 1:
            kwargs["halo"] = SphericalNFWPotential(units=units,
                                                   v_c=v_c, r_s=r_s)
        else:
            kwargs["halo"] = LeeSutoTriaxialNFWPotential(units=units,
                                                         a=q1, b=q2, c=q3,
                                                         v_c=v_c, r_s=r_s,
                                                         phi=phi, theta=theta, psi=psi)
        super(PW14Potential,self).__init__(**kwargs)
        self.c_instance.G = G.decompose(units).value

class LM10Potential(CCompositePotential):

    def __init__(self, m_disk=1E11, a=6.5, b=0.26,
                 m_spher=3.4E10, c=0.7,
                 q1=1.38, q2=1., q3=1.36, phi=(97*u.degree).to(u.radian).value,
                 v_c=np.sqrt(2)*(121.858*u.km/u.s).to(u.kpc/u.Myr).value, r_h=12.,
                 units=galactic):

        kwargs = dict()
        kwargs["disk"] = MiyamotoNagaiPotential(units=units,
                                                m=m_disk, a=a, b=b)

        kwargs["bulge"] = HernquistPotential(units=units,
                                             m=m_spher, c=c)

        kwargs["halo"] = LogarithmicPotential(units=units,
                                              q1=q1, q2=q2, q3=q3,
                                              phi=phi, v_c=v_c, r_h=r_h)
        super(LM10Potential,self).__init__(**kwargs)
        self.c_instance.G = G.decompose(units).value

class OblateMWPotential(CCompositePotential):

    def __init__(self, m_disk=7E10, a=3.5, b=0.14,
                 m_spher=1E10, c=1.1,
                 q1=1., q2=0.75, q3=0.55,
                 v_c=0.239225, r_s=30.,
                 phi=0., theta=0., psi=0.,
                 units=galactic):
        """ Axis ratio values taken from Jing & Suto (2002). Other
            parameters come from a by-eye fit to Bovy's MW2014Potential.
        """

        # Choice of v_h sets circular velocity at Sun to 220 km/s
        self.units = units

        kwargs = dict()
        kwargs["disk"] = MiyamotoNagaiPotential(units=units,
                                                m=m_disk, a=a, b=b)

        kwargs["bulge"] = HernquistPotential(units=units,
                                             m=m_spher, c=c)

        kwargs["halo"] = LeeSutoTriaxialNFWPotential(units=units,
                                                     a=q1, b=q2, c=q3,
                                                     v_c=v_c, r_s=r_s,
                                                     phi=phi, theta=theta, psi=psi)
        super(OblateMWPotential,self).__init__(**kwargs)
        self.c_instance.G = G.decompose(units).value
<commit_msg>change name of oblate potential<commit_after># coding: utf-8

""" Potential used in Price-Whelan et al. (in prep.) TODO """

from __future__ import division, print_function

__author__ = "adrn <adrn@astro.columbia.edu>"

# Third-party
from astropy.constants import G
import astropy.units as u
import numpy as np

# Project
from .cpotential import CCompositePotential
from .cbuiltin import HernquistPotential, MiyamotoNagaiPotential, \
    LeeSutoTriaxialNFWPotential, SphericalNFWPotential, LogarithmicPotential
from ..units import galactic

__all__ = ['PW14Potential', 'LM10Potential', 'TriaxialMWPotential']

class PW14Potential(CCompositePotential):

    def __init__(self, m_disk=6.5E10, a=6.5, b=0.26,
                 m_spher=2E10, c=0.3,
                 q1=1.4, q2=1., q3=0.6,
                 v_c=0.247, r_s=30.,
                 phi=np.pi/2., theta=np.pi/2., psi=np.pi/2.,
                 units=galactic):

        # Choice of v_h sets circular velocity at Sun to 220 km/s
        self.units = units

        kwargs = dict()
        kwargs["disk"] = MiyamotoNagaiPotential(units=units,
                                                m=m_disk, a=a, b=b)

        kwargs["bulge"] = HernquistPotential(units=units,
                                             m=m_spher, c=c)

        if q1 == 1 and q2 == 1 and q3 == 1:
            kwargs["halo"] = SphericalNFWPotential(units=units,
                                                   v_c=v_c, r_s=r_s)
        else:
            kwargs["halo"] = LeeSutoTriaxialNFWPotential(units=units,
                                                         a=q1, b=q2, c=q3,
                                                         v_c=v_c, r_s=r_s,
                                                         phi=phi, theta=theta, psi=psi)
        super(PW14Potential,self).__init__(**kwargs)
        self.c_instance.G = G.decompose(units).value

class LM10Potential(CCompositePotential):

    def __init__(self, m_disk=1E11, a=6.5, b=0.26,
                 m_spher=3.4E10, c=0.7,
                 q1=1.38, q2=1., q3=1.36, phi=(97*u.degree).to(u.radian).value,
                 v_c=np.sqrt(2)*(121.858*u.km/u.s).to(u.kpc/u.Myr).value, r_h=12.,
                 units=galactic):

        kwargs = dict()
        kwargs["disk"] = MiyamotoNagaiPotential(units=units,
                                                m=m_disk, a=a, b=b)

        kwargs["bulge"] = HernquistPotential(units=units,
                                             m=m_spher, c=c)

        kwargs["halo"] = LogarithmicPotential(units=units,
                                              q1=q1, q2=q2, q3=q3,
                                              phi=phi, v_c=v_c, r_h=r_h)
        super(LM10Potential,self).__init__(**kwargs)
        self.c_instance.G = G.decompose(units).value

class TriaxialMWPotential(CCompositePotential):

    def __init__(self, m_disk=7E10, a=3.5, b=0.14,
                 m_spher=1E10, c=1.1,
                 q1=1., q2=0.75, q3=0.55,
                 v_c=0.239225, r_s=30.,
                 phi=0., theta=0., psi=0.,
                 units=galactic):
        """ Axis ratio values taken from Jing & Suto (2002). Other
            parameters come from a by-eye fit to Bovy's MW2014Potential.
        """

        # Choice of v_h sets circular velocity at Sun to 220 km/s
        self.units = units

        kwargs = dict()
        kwargs["disk"] = MiyamotoNagaiPotential(units=units,
                                                m=m_disk, a=a, b=b)

        kwargs["bulge"] = HernquistPotential(units=units,
                                             m=m_spher, c=c)

        kwargs["halo"] = LeeSutoTriaxialNFWPotential(units=units,
                                                     a=q1, b=q2, c=q3,
                                                     v_c=v_c, r_s=r_s,
                                                     phi=phi, theta=theta, psi=psi)
        super(TriaxialMWPotential,self).__init__(**kwargs)
        self.c_instance.G = G.decompose(units).value
<|endoftext|>


<commit_before>from collections import Counter
import pickle
import re
import numpy as np
import pandas as pd
from tqdm.auto import tqdm
from rdkit import Chem
from rdkit.Chem.BRICS import reactionDefs

import moses
from moses.metrics.utils import fragmenter
from moses.utils import mapper

isotope_re = re.compile(r'\[([0-9]+)[ab]?\*\]')

"""
Prepare binary connection rules
"""
binary_defs = {}
connection_rules = set()
for row in reactionDefs:
    for a, b, t in row:
        if a in ['7a', '7b']:
            binary_defs[7] = 2**7
            connection_rules.add((7, 7))
        else:
            a = int(a)
            b = int(b)
            binary_defs[a] = binary_defs.get(a, 0) | 2**b
            binary_defs[b] = binary_defs.get(b, 0) | 2**a
            connection_rules.add((a, b))
            connection_rules.add((b, a))

class CombinatorialGenerator:
    def __init__(self, n_jobs=1, mode=0):
        """
        Combinatorial Generator randomly connects BRICS fragments

        Arguments:
            n_jobs: number of processes for training
            mode: sampling mode
                0: Sample fragment then connection
                1: Sample connection point then fragments
        """
        self.n_jobs = n_jobs
        self.set_mode(mode)
        self.fitted = False

    def fit(self, data):
        """
        Collects fragment frequencies in a training set

        Arguments:
            data: list of SMILES, training dataset

        """
        # Split molecules from dataset into BRICS fragments
        fragments = mapper(self.n_jobs)(fragmenter, data)

        # Compute fragment frequencies
        counts = Counter()
        for mol_frag in fragments:
            counts.update(mol_frag)
        counts = pd.DataFrame(
            pd.Series(counts).items(),
            columns=['fragment', 'count']
        )
        counts['attachment_points'] = [
            fragment.count('*')
            for fragment in counts['fragment'].values
        ]
        counts['connection_rules'] = [
            self.get_connection_rule(fragment)
            for fragment in counts['fragment'].values
        ]
        counts['frequency'] = counts['count'] / counts['count'].sum()
        self.fragment_counts = counts

        # Compute number of fragments distribution
        fragments_count_distribution = Counter([len(f) for f in fragments])
        total = sum(fragments_count_distribution.values())
        for k in fragments_count_distribution:
            fragments_count_distribution[k] /= total
        self.fragments_count_distribution = fragments_count_distribution
        self.fitted = True
        return self

    def save(self, path):
        """
        Saves a model using pickle

        Arguments:
            path: path to .pkl file for saving
        """
        if not self.fitted:
            raise RuntimeError("Can't save empty model."
                               " Fit the model first")
        data = {
            'fragment_counts': self.fragment_counts,
            'fragments_count_distribution': self.fragments_count_distribution,
            'n_jobs': self.n_jobs,
            'mode': self.mode
        }
        with open(path, 'wb') as f:
            pickle.dump(data, f)

    @classmethod
    def load(cls, path):
        """
        Loads saved model

        Arguments:
            path: path to saved .pkl file

        Returns:
            Loaded CombinatorialGenerator
        """
        with open(path, "rb") as f:
            data = pickle.load(f)
        model = cls()
        model.fragment_counts = data['fragment_counts']
        model.fragments_count_distribution = \
            data['fragments_count_distribution']
        model.n_jobs = data['n_jobs']
        model.mode = data['mode']
        model.fitted = True
        return model

    def set_mode(self, mode):
        if mode not in [0, 1]:
            raise ValueError('Incorrect mode value: %s' % mode)
        self.mode = mode

    def generate_one(self, seed=None):
        """
        Generates a SMILES string using fragment frequencies

        Arguments:
            seed: if specified, will set numpy seed before sampling

        Returns:
            SMILES string
        """
        if not self.fitted:
            raise RuntimeError("Fit the model before generating")
        if seed is not None:
            np.random.seed(seed)
        mol = None

        # Sample the number of fragments
        count_values, count_probs = zip(
            *self.fragments_count_distribution.items())
        total_fragments = np.random.choice(count_values, p=count_probs)

        counts = self.fragment_counts
        current_attachments = 0
        max_attachments = total_fragments - 1
        connections_mol = None
        for i in range(total_fragments):
            # Enforce lower and upper limit on the number of connection points
            counts_masked = counts[
                counts['attachment_points'] <= max_attachments
            ]
            if total_fragments == 1:
                min_attachments = 0
            elif i != 0 and current_attachments == 1 and \
                    total_fragments > i + 1:
                min_attachments = 2
            else:
                min_attachments = 1
            counts_masked = counts_masked[
                counts_masked['attachment_points'] >= min_attachments
            ]

            if mol is None:
                mol = self.sample_fragment(counts_masked)
            else:
                if self.mode == 1:  # Choose connection atom first
                    atom_mol = np.random.choice(connections_mol)
                    connections_mol = [atom_mol]

                con_filter = self.get_connection_filter(connections_mol)
                # Mask fragments with possible reactions
                counts_masked = counts_masked[
                    counts_masked['connection_rules'] & con_filter > 0
                ]
                fragment = self.sample_fragment(counts_masked)
                connections_fragment = self.get_connection_points(fragment)
                possible_connections = self.filter_connections(
                    connections_mol,
                    connections_fragment
                )

                c_i = np.random.choice(len(possible_connections))
                a1, a2 = possible_connections[c_i]

                # Connect a new fragment to the molecule
                mol = self.connect_mols(mol, fragment, a1, a2)

            connections_mol = self.get_connection_points(mol)
            current_attachments = len(connections_mol)
            max_attachments = (
                total_fragments - i - current_attachments
            )
        smiles = Chem.MolToSmiles(mol)
        return smiles

    def generate(self, n, seed=1, mode=0, verbose=False):
        self.set_mode(mode)
        seeds = range((seed - 1) * n, seed * n)
        if verbose:
            print('generating...')
            seeds = tqdm(seeds, total=n)
        samples = mapper(self.n_jobs)(self.generate_one, seeds)
        return samples

    def get_connection_rule(self, fragment):
        """
        return OR combination for possible incoming reactions

        Arguments:
            fragment: fragment smiles

        Returns:
            int
        """
        rule = 0
        for i in map(int, set(isotope_re.findall(fragment))):
            rule |= binary_defs[i]
        return rule

    @staticmethod
    def sample_fragment(counts):
        new_fragment = counts.sample(
            weights=counts['frequency']
        )
        new_fragment = dict(new_fragment.iloc[0])
        fragment = Chem.MolFromSmiles(new_fragment['fragment'])
        return fragment

    @staticmethod
    def get_connection_filter(atoms):
        """
        Return OR(2**isotopes)
        """
        connection_rule = 0
        for atom in atoms:
            connection_rule |= 2**atom.GetIsotope()
        return connection_rule

    @staticmethod
    def get_connection_points(mol):
        """
        Return connection points

        Arguments:
            mol: ROMol

        Returns:
            atom list
        """
        atoms = []
        for atom in mol.GetAtoms():
            if atom.GetSymbol() == '*':
                atoms.append(atom)
        return atoms

    @staticmethod
    def filter_connections(atoms1, atoms2, unique=True):
        possible_connections = []
        for a1 in atoms1:
            i1 = a1.GetIsotope()
            for a2 in atoms2:
                i2 = a2.GetIsotope()
                if (i1, i2) in connection_rules:
                    possible_connections.append((a1, a2))
        return possible_connections

    @staticmethod
    def connect_mols(mol1, mol2, atom1, atom2):
        combined = Chem.CombineMols(mol1, mol2)
        emol = Chem.EditableMol(combined)
        neighbor1_idx = atom1.GetNeighbors()[0].GetIdx()
        neighbor2_idx = atom2.GetNeighbors()[0].GetIdx()
        atom1_idx = atom1.GetIdx()
        atom2_idx = atom2.GetIdx()
        bond_order = atom2.GetBonds()[0].GetBondType()
        emol.AddBond(neighbor1_idx,
                     neighbor2_idx + mol1.GetNumAtoms(),
                     order=bond_order)
        emol.RemoveAtom(atom2_idx + mol1.GetNumAtoms())
        emol.RemoveAtom(atom1_idx)
        mol = emol.GetMol()
        return mol

def reproduce(seed, samples_path=None, metrics_path=None,
              n_jobs=1, device='cpu', verbose=False,
              samples=30000):
    train = moses.get_dataset('train')
    model = CombinatorialGenerator(n_jobs=n_jobs)

    if verbose:
        print("Training...")
    model.fit(train)

    if verbose:
        print(f"Sampling for seed {seed}")
    seeds = list(range(
        (seed - 1) * samples, seed * samples
    ))
    samples = mapper(n_jobs)(model.generate_one, seeds)
    if samples_path is not None:
        with open(samples_path, 'w') as f:
            f.write('SMILES\n')
            for sample in samples:
                f.write(sample+'\n')
    if verbose:
        print(f"Computing metrics for seed {seed}")
    metrics = moses.get_all_metrics(
        samples, n_jobs=n_jobs, device=device)
    if metrics_path is not None:
        with open(metrics_path, 'w') as f:
            for key, value in metrics.items():
                f.write("%s,%f\n" % (key, value))
    return samples, metrics

if __name__ == "__main__":
    import argparse
    import os

    parser = argparse.ArgumentParser(
        "Reproduce combinatorial experiment")
    parser.add_argument(
        '--n_jobs', type=int, required=False,
        default=1, help='Number of threads for computing metrics')
    parser.add_argument(
        '--device', type=str, required=False,
        default='cpu', help='Device for computing metrics')
    parser.add_argument(
        '--metrics_path', type=str, required=False,
        default='.', help='Path to save metrics')
    args = parser.parse_known_args()[0]

    for seed in [1, 2, 3]:
        filename = f'combinatorial_metrics_{seed}.csv'
        metrics_path = os.path.join(args.metrics_path, filename)

        reproduce(seed=seed, n_jobs=args.n_jobs,
                  device=args.device, verbose=True,
                  metrics_path=metrics_path)
<commit_msg>print test error<commit_after>from collections import Counter
import pickle
import re
import numpy as np
import pandas as pd
from tqdm.auto import tqdm
from rdkit import Chem
from rdkit.Chem.BRICS import reactionDefs

import moses
from moses.metrics.utils import fragmenter
from moses.utils import mapper

isotope_re = re.compile(r'\[([0-9]+)[ab]?\*\]')

"""
Prepare binary connection rules
"""
binary_defs = {}
connection_rules = set()
for row in reactionDefs:
    for a, b, t in row:
        if a in ['7a', '7b']:
            binary_defs[7] = 2**7
            connection_rules.add((7, 7))
        else:
            a = int(a)
            b = int(b)
            binary_defs[a] = binary_defs.get(a, 0) | 2**b
            binary_defs[b] = binary_defs.get(b, 0) | 2**a
            connection_rules.add((a, b))
            connection_rules.add((b, a))

class CombinatorialGenerator:
    def __init__(self, n_jobs=1, mode=0):
        """
        Combinatorial Generator randomly connects BRICS fragments

        Arguments:
            n_jobs: number of processes for training
            mode: sampling mode
                0: Sample fragment then connection
                1: Sample connection point then fragments
        """
        self.n_jobs = n_jobs
        self.set_mode(mode)
        self.fitted = False

    def fit(self, data):
        """
        Collects fragment frequencies in a training set

        Arguments:
            data: list of SMILES, training dataset

        """
        # Split molecules from dataset into BRICS fragments
        fragments = mapper(self.n_jobs)(fragmenter, data)

        # Compute fragment frequencies
        counts = Counter()
        for mol_frag in fragments:
            counts.update(mol_frag)
        counts = pd.DataFrame(
            pd.Series(counts).items(),
            columns=['fragment', 'count']
        )
        counts['attachment_points'] = [
            fragment.count('*')
            for fragment in counts['fragment'].values
        ]
        counts['connection_rules'] = [
            self.get_connection_rule(fragment)
            for fragment in counts['fragment'].values
        ]
        counts['frequency'] = counts['count'] / counts['count'].sum()
        self.fragment_counts = counts

        # Compute number of fragments distribution
        fragments_count_distribution = Counter([len(f) for f in fragments])
        total = sum(fragments_count_distribution.values())
        for k in fragments_count_distribution:
            fragments_count_distribution[k] /= total
        self.fragments_count_distribution = fragments_count_distribution
        self.fitted = True
        return self

    def save(self, path):
        """
        Saves a model using pickle

        Arguments:
            path: path to .pkl file for saving
        """
        if not self.fitted:
            raise RuntimeError("Can't save empty model."
                               " Fit the model first")
        data = {
            'fragment_counts': self.fragment_counts,
            'fragments_count_distribution': self.fragments_count_distribution,
            'n_jobs': self.n_jobs,
            'mode': self.mode
        }
        with open(path, 'wb') as f:
            pickle.dump(data, f)

    @classmethod
    def load(cls, path):
        """
        Loads saved model

        Arguments:
            path: path to saved .pkl file

        Returns:
            Loaded CombinatorialGenerator
        """
        with open(path, "rb") as f:
            data = pickle.load(f)
        model = cls()
        model.fragment_counts = data['fragment_counts']
        model.fragments_count_distribution = \
            data['fragments_count_distribution']
        model.n_jobs = data['n_jobs']
        model.mode = data['mode']
        model.fitted = True
        return model

    def set_mode(self, mode):
        if mode not in [0, 1]:
            raise ValueError('Incorrect mode value: %s' % mode)
        self.mode = mode

    def generate_one(self, seed=None):
        """
        Generates a SMILES string using fragment frequencies

        Arguments:
            seed: if specified, will set numpy seed before sampling

        Returns:
            SMILES string
        """
        if not self.fitted:
            raise RuntimeError("Fit the model before generating")
        if seed is not None:
            np.random.seed(seed)
        mol = None

        # Sample the number of fragments
        count_values, count_probs = zip(
            *self.fragments_count_distribution.items())
        total_fragments = np.random.choice(count_values, p=count_probs)

        counts = self.fragment_counts
        current_attachments = 0
        max_attachments = total_fragments - 1
        connections_mol = None
        for i in range(total_fragments):
            # Enforce lower and upper limit on the number of connection points
            counts_masked = counts[
                counts['attachment_points'] <= max_attachments
            ]
            if total_fragments == 1:
                min_attachments = 0
            elif i != 0 and current_attachments == 1 and \
                    total_fragments > i + 1:
                min_attachments = 2
            else:
                min_attachments = 1
            counts_masked = counts_masked[
                counts_masked['attachment_points'] >= min_attachments
            ]

            if mol is None:
                mol = self.sample_fragment(counts_masked)
            else:
                if self.mode == 1:  # Choose connection atom first
                    connections_mol = [np.random.choice(connections_mol)]

                con_filter = self.get_connection_filter(connections_mol)
                # Mask fragments with possible reactions
                counts_masked = counts_masked[
                    counts_masked['connection_rules'] & con_filter > 0
                ]
                fragment = self.sample_fragment(counts_masked)
                connections_fragment = self.get_connection_points(fragment)
                possible_connections = self.filter_connections(
                    connections_mol,
                    connections_fragment
                )

                c_i = np.random.choice(len(possible_connections))
                a1, a2 = possible_connections[c_i]

                # Connect a new fragment to the molecule
                mol = self.connect_mols(mol, fragment, a1, a2)

            connections_mol = self.get_connection_points(mol)
            current_attachments = len(connections_mol)
            max_attachments = (
                total_fragments - i - current_attachments
            )
        smiles = Chem.MolToSmiles(mol)
        return smiles

    def generate(self, n, seed=1, mode=0, verbose=False):
        self.set_mode(mode)
        seeds = range((seed - 1) * n, seed * n)
        if verbose:
            print('generating...')
            seeds = tqdm(seeds, total=n)
        samples = mapper(self.n_jobs)(self.generate_one, seeds)
        return samples

    def get_connection_rule(self, fragment):
        """
        return OR combination for possible incoming reactions

        Arguments:
            fragment: fragment smiles

        Returns:
            int
        """
        rule = 0
        for i in map(int, set(isotope_re.findall(fragment))):
            rule |= binary_defs[i]
        return rule

    @staticmethod
    def sample_fragment(counts):
        try:
            new_fragment = counts.sample(
                weights=counts['frequency']
            )
        except:
            print(counts)
            raise
        new_fragment = dict(new_fragment.iloc[0])
        fragment = Chem.MolFromSmiles(new_fragment['fragment'])
        return fragment

    @staticmethod
    def get_connection_filter(atoms):
        """
        Return OR(2**isotopes)
        """
        connection_rule = 0
        for atom in atoms:
            connection_rule |= 2**atom.GetIsotope()
        return connection_rule

    @staticmethod
    def get_connection_points(mol):
        """
        Return connection points

        Arguments:
            mol: ROMol

        Returns:
            atom list
        """
        atoms = []
        for atom in mol.GetAtoms():
            if atom.GetSymbol() == '*':
                atoms.append(atom)
        return atoms

    @staticmethod
    def filter_connections(atoms1, atoms2, unique=True):
        possible_connections = []
        for a1 in atoms1:
            i1 = a1.GetIsotope()
            for a2 in atoms2:
                i2 = a2.GetIsotope()
                if (i1, i2) in connection_rules:
                    possible_connections.append((a1, a2))
        return possible_connections

    @staticmethod
    def connect_mols(mol1, mol2, atom1, atom2):
        combined = Chem.CombineMols(mol1, mol2)
        emol = Chem.EditableMol(combined)
        neighbor1_idx = atom1.GetNeighbors()[0].GetIdx()
        neighbor2_idx = atom2.GetNeighbors()[0].GetIdx()
        atom1_idx = atom1.GetIdx()
        atom2_idx = atom2.GetIdx()
        bond_order = atom2.GetBonds()[0].GetBondType()
        emol.AddBond(neighbor1_idx,
                     neighbor2_idx + mol1.GetNumAtoms(),
                     order=bond_order)
        emol.RemoveAtom(atom2_idx + mol1.GetNumAtoms())
        emol.RemoveAtom(atom1_idx)
        mol = emol.GetMol()
        return mol

def reproduce(seed, samples_path=None, metrics_path=None,
              n_jobs=1, device='cpu', verbose=False,
              samples=30000):
    train = moses.get_dataset('train')
    model = CombinatorialGenerator(n_jobs=n_jobs)

    if verbose:
        print("Training...")
    model.fit(train)

    if verbose:
        print(f"Sampling for seed {seed}")
    seeds = list(range(
        (seed - 1) * samples, seed * samples
    ))
    samples = mapper(n_jobs)(model.generate_one, seeds)
    if samples_path is not None:
        with open(samples_path, 'w') as f:
            f.write('SMILES\n')
            for sample in samples:
                f.write(sample+'\n')
    if verbose:
        print(f"Computing metrics for seed {seed}")
    metrics = moses.get_all_metrics(
        samples, n_jobs=n_jobs, device=device)
    if metrics_path is not None:
        with open(metrics_path, 'w') as f:
            for key, value in metrics.items():
                f.write("%s,%f\n" % (key, value))
    return samples, metrics

if __name__ == "__main__":
    import argparse
    import os

    parser = argparse.ArgumentParser(
        "Reproduce combinatorial experiment")
    parser.add_argument(
        '--n_jobs', type=int, required=False,
        default=1, help='Number of threads for computing metrics')
    parser.add_argument(
        '--device', type=str, required=False,
        default='cpu', help='Device for computing metrics')
    parser.add_argument(
        '--metrics_path', type=str, required=False,
        default='.', help='Path to save metrics')
    args = parser.parse_known_args()[0]

    for seed in [1, 2, 3]:
        filename = f'combinatorial_metrics_{seed}.csv'
        metrics_path = os.path.join(args.metrics_path, filename)

        reproduce(seed=seed, n_jobs=args.n_jobs,
                  device=args.device, verbose=True,
                  metrics_path=metrics_path)
<|endoftext|>


<commit_before>from copy import deepcopy
from math import fsum

from numpy import dot, empty, inf, isfinite, log, maximum, sqrt, zeros
from numpy.linalg import norm

from .posterior import Posterior
from .site import Site

MAX_ITERS = 100

class EP(object):
    """
    Expectation Propagation algorithm.

    Let ::

        𝒩(𝐳｜𝐦, K)

    be the prior distribution.
    This class estimates the log of the marginal likelihood ::

        p(𝐲) = ∫∏ᵢp(yᵢ｜μᵢ = g(zᵢ))𝒩(𝐳｜𝐦, K) d𝐳

    via Expectation Propagation and provides its gradient.

    Attributes
    ----------
        _site: site-likelihood.
        _psite: previous site-likelihood.
        _cav: cavity distribution.
        _posterior: posterior distribution.
        _moments: moments for KL moment matching.
    """

    def __init__(self, nsites, posterior_type=Posterior, rtol=1.49e-05, atol=1.49e-08):
        self._site = Site(nsites)
        self._psite = Site(nsites)
        self._rtol = rtol
        self._atol = atol

        self._cav = dict(tau=zeros(nsites), eta=zeros(nsites))
        self._posterior = posterior_type(self._site)

        self._moments = {
            "log_zeroth": empty(nsites),
            "mean": empty(nsites),
            "variance": empty(nsites),
        }

        self._need_update = True
        self._compute_moments = None
        self._cache = dict(lml=None, grad=None)
        self.verbose = False

    def __copy__(self):
        cls = self.__class__
        ep = cls.__new__(cls)

        ep.__dict__["_site"] = deepcopy(self.site)
        ep.__dict__["_psite"] = deepcopy(self.__dict__["_psite"])

        ep.__dict__["_rtol"] = self._rtol
        ep.__dict__["_atol"] = self._atol

        ep.__dict__["_cav"] = deepcopy(self.cav)

        posterior_type = type(self.posterior)
        ep.__dict__["_posterior"] = posterior_type(ep.site)

        ep.__dict__["_moments"] = deepcopy(self.moments)

        ep.__dict__["_need_update"] = self._need_update
        ep.__dict__["_compute_moments"] = None
        ep.__dict__["_cache"] = deepcopy(self._cache)
        ep.__dict__["verbose"] = deepcopy(self.verbose)

        return ep

    def _update(self):
        if not self._need_update:
            return

        self._posterior.update()

        i = 0
        tol = -inf
        n1 = norm(self._site.tau)
        while i < MAX_ITERS and norm(self._site.tau - self._psite.tau) > tol:
            self._psite.tau[:] = self._site.tau
            self._psite.eta[:] = self._site.eta

            self._cav["tau"][:] = maximum(self._posterior.tau - self._site.tau, 0)
            self._cav["eta"][:] = self._posterior.eta - self._site.eta
            self._compute_moments(self._cav["eta"], self._cav["tau"], self._moments)

            self._site.update(
                self._moments["mean"], self._moments["variance"], self._cav
            )

            self._posterior.update()

            n0 = n1
            n1 = norm(self._site.tau)
            tol = self._rtol * min(n0, n1) + self._atol
            i += 1
            if i % 10 == 9:
                self._rtol *= 10

        if i == MAX_ITERS:
            msg = "Maximum number of EP iterations has" + " been attained."
            msg += " Last EP step was: %.10f." % norm(self._site.tau - self._psite.tau)
            raise ValueError(msg)

        if self.verbose:
            print("{} EP iterations".format(i))
        self._need_update = False

    @property
    def cav(self):
        return self._cav

    def lml(self):
        from numpy_sugar.linalg import cho_solve

        if self._cache["lml"] is not None:
            return self._cache["lml"]

        self._update()

        L = self._posterior.L()
        Q, S = self._posterior.cov["QS"]
        Q = Q[0]
        ttau = self._site.tau
        teta = self._site.eta
        ctau = self._cav["tau"]
        ceta = self._cav["eta"]
        m = self._posterior.mean

        TS = ttau + ctau

        lml = [
            -log(L.diagonal()).sum(),
            -0.5 * sum(log(S)),
            # lml += 0.5 * sum(log(ttau)),
            +0.5 * dot(teta, dot(Q, cho_solve(L, dot(Q.T, teta)))),
            -0.5 * dot(teta, teta / TS),
            +dot(m, teta) - 0.5 * dot(m, ttau * m),
            -0.5 * dot(m * ttau, dot(Q, cho_solve(L, dot(Q.T, 2 * teta - ttau * m)))),
            +sum(self._moments["log_zeroth"]),
            +0.5 * sum(log(TS)),
            # lml -= 0.5 * sum(log(ttau)),
            -0.5 * sum(log(ctau)),
            +0.5 * dot(ceta / TS, ttau * ceta / ctau - 2 * teta),
        ]
        lml = fsum(lml)

        if not isfinite(lml):
            raise ValueError("LML should not be %f." % lml)

        self._cache["lml"] = lml
        return lml

    def lml_derivative_over_cov(self, dQS):
        from numpy_sugar.linalg import cho_solve, ddot, dotd

        self._update()

        L = self._posterior.L()
        Q = self._posterior.cov["QS"][0][0]
        ttau = self._site.tau
        teta = self._site.eta

        diff = teta - ttau * self._posterior.mean

        v0 = dot(dQS[0][0], dQS[1] * dot(dQS[0][0].T, diff))
        v1 = ttau * dot(Q, cho_solve(L, dot(Q.T, diff)))
        dlml = 0.5 * dot(diff, v0)
        dlml -= dot(v0, v1)
        dlml += 0.5 * dot(v1, dot(dQS[0][0], dQS[1] * dot(dQS[0][0].T, v1)))
        dqs = ddot(dQS[1], dQS[0][0].T, left=True)
        diag = dotd(dQS[0][0], dqs)
        dlml -= 0.5 * sum(ttau * diag)

        tmp = cho_solve(L, dot(ddot(Q.T, ttau, left=False), dQS[0][0]))
        dlml += 0.5 * sum(ttau * dotd(Q, dot(tmp, dqs)))

        return dlml

    def lml_derivative_over_mean(self, dm):
        from numpy_sugar.linalg import cho_solve

        self._update()

        L = self._posterior.L()
        Q = self._posterior.cov["QS"][0][0]
        ttau = self._site.tau
        teta = self._site.eta

        diff = teta - ttau * self._posterior.mean

        dlml = dot(diff, dm)
        dlml -= dot(diff, dot(Q, cho_solve(L, dot(Q.T, (ttau * dm.T).T))))

        return dlml

    @property
    def moments(self):
        return self._moments

    @property
    def posterior(self):
        return self._posterior

    def set_compute_moments(self, cm):
        self._compute_moments = cm

    @property
    def site(self):
        return self._site

    def set_prior(self, mean, covariance):
        self._posterior.mean = mean
        self._posterior.cov = covariance
        self._need_update = True
        for k in self._cache:
            self._cache[k] = None
<commit_msg>flake 8<commit_after>from copy import deepcopy
from math import fsum

from numpy import dot, empty, inf, isfinite, log, maximum, zeros
from numpy.linalg import norm

from .posterior import Posterior
from .site import Site

MAX_ITERS = 100

class EP(object):
    """
    Expectation Propagation algorithm.

    Let ::

        𝒩(𝐳｜𝐦, K)

    be the prior distribution.
    This class estimates the log of the marginal likelihood ::

        p(𝐲) = ∫∏ᵢp(yᵢ｜μᵢ = g(zᵢ))𝒩(𝐳｜𝐦, K) d𝐳

    via Expectation Propagation and provides its gradient.

    Attributes
    ----------
        _site: site-likelihood.
        _psite: previous site-likelihood.
        _cav: cavity distribution.
        _posterior: posterior distribution.
        _moments: moments for KL moment matching.
    """

    def __init__(self, nsites, posterior_type=Posterior, rtol=1.49e-05, atol=1.49e-08):
        self._site = Site(nsites)
        self._psite = Site(nsites)
        self._rtol = rtol
        self._atol = atol

        self._cav = dict(tau=zeros(nsites), eta=zeros(nsites))
        self._posterior = posterior_type(self._site)

        self._moments = {
            "log_zeroth": empty(nsites),
            "mean": empty(nsites),
            "variance": empty(nsites),
        }

        self._need_update = True
        self._compute_moments = None
        self._cache = dict(lml=None, grad=None)
        self.verbose = False

    def __copy__(self):
        cls = self.__class__
        ep = cls.__new__(cls)

        ep.__dict__["_site"] = deepcopy(self.site)
        ep.__dict__["_psite"] = deepcopy(self.__dict__["_psite"])

        ep.__dict__["_rtol"] = self._rtol
        ep.__dict__["_atol"] = self._atol

        ep.__dict__["_cav"] = deepcopy(self.cav)

        posterior_type = type(self.posterior)
        ep.__dict__["_posterior"] = posterior_type(ep.site)

        ep.__dict__["_moments"] = deepcopy(self.moments)

        ep.__dict__["_need_update"] = self._need_update
        ep.__dict__["_compute_moments"] = None
        ep.__dict__["_cache"] = deepcopy(self._cache)
        ep.__dict__["verbose"] = deepcopy(self.verbose)

        return ep

    def _update(self):
        if not self._need_update:
            return

        self._posterior.update()

        i = 0
        tol = -inf
        n1 = norm(self._site.tau)
        while i < MAX_ITERS and norm(self._site.tau - self._psite.tau) > tol:
            self._psite.tau[:] = self._site.tau
            self._psite.eta[:] = self._site.eta

            self._cav["tau"][:] = maximum(self._posterior.tau - self._site.tau, 0)
            self._cav["eta"][:] = self._posterior.eta - self._site.eta
            self._compute_moments(self._cav["eta"], self._cav["tau"], self._moments)

            self._site.update(
                self._moments["mean"], self._moments["variance"], self._cav
            )

            self._posterior.update()

            n0 = n1
            n1 = norm(self._site.tau)
            tol = self._rtol * min(n0, n1) + self._atol
            i += 1
            if i % 10 == 9:
                self._rtol *= 10

        if i == MAX_ITERS:
            msg = "Maximum number of EP iterations has" + " been attained."
            msg += " Last EP step was: %.10f." % norm(self._site.tau - self._psite.tau)
            raise ValueError(msg)

        if self.verbose:
            print("{} EP iterations".format(i))
        self._need_update = False

    @property
    def cav(self):
        return self._cav

    def lml(self):
        from numpy_sugar.linalg import cho_solve

        if self._cache["lml"] is not None:
            return self._cache["lml"]

        self._update()

        L = self._posterior.L()
        Q, S = self._posterior.cov["QS"]
        Q = Q[0]
        ttau = self._site.tau
        teta = self._site.eta
        ctau = self._cav["tau"]
        ceta = self._cav["eta"]
        m = self._posterior.mean

        TS = ttau + ctau

        lml = [
            -log(L.diagonal()).sum(),
            -0.5 * sum(log(S)),
            # lml += 0.5 * sum(log(ttau)),
            +0.5 * dot(teta, dot(Q, cho_solve(L, dot(Q.T, teta)))),
            -0.5 * dot(teta, teta / TS),
            +dot(m, teta) - 0.5 * dot(m, ttau * m),
            -0.5 * dot(m * ttau, dot(Q, cho_solve(L, dot(Q.T, 2 * teta - ttau * m)))),
            +sum(self._moments["log_zeroth"]),
            +0.5 * sum(log(TS)),
            # lml -= 0.5 * sum(log(ttau)),
            -0.5 * sum(log(ctau)),
            +0.5 * dot(ceta / TS, ttau * ceta / ctau - 2 * teta),
        ]
        lml = fsum(lml)

        if not isfinite(lml):
            raise ValueError("LML should not be %f." % lml)

        self._cache["lml"] = lml
        return lml

    def lml_derivative_over_cov(self, dQS):
        from numpy_sugar.linalg import cho_solve, ddot, dotd

        self._update()

        L = self._posterior.L()
        Q = self._posterior.cov["QS"][0][0]
        ttau = self._site.tau
        teta = self._site.eta

        diff = teta - ttau * self._posterior.mean

        v0 = dot(dQS[0][0], dQS[1] * dot(dQS[0][0].T, diff))
        v1 = ttau * dot(Q, cho_solve(L, dot(Q.T, diff)))
        dlml = 0.5 * dot(diff, v0)
        dlml -= dot(v0, v1)
        dlml += 0.5 * dot(v1, dot(dQS[0][0], dQS[1] * dot(dQS[0][0].T, v1)))
        dqs = ddot(dQS[1], dQS[0][0].T, left=True)
        diag = dotd(dQS[0][0], dqs)
        dlml -= 0.5 * sum(ttau * diag)

        tmp = cho_solve(L, dot(ddot(Q.T, ttau, left=False), dQS[0][0]))
        dlml += 0.5 * sum(ttau * dotd(Q, dot(tmp, dqs)))

        return dlml

    def lml_derivative_over_mean(self, dm):
        from numpy_sugar.linalg import cho_solve

        self._update()

        L = self._posterior.L()
        Q = self._posterior.cov["QS"][0][0]
        ttau = self._site.tau
        teta = self._site.eta

        diff = teta - ttau * self._posterior.mean

        dlml = dot(diff, dm)
        dlml -= dot(diff, dot(Q, cho_solve(L, dot(Q.T, (ttau * dm.T).T))))

        return dlml

    @property
    def moments(self):
        return self._moments

    @property
    def posterior(self):
        return self._posterior

    def set_compute_moments(self, cm):
        self._compute_moments = cm

    @property
    def site(self):
        return self._site

    def set_prior(self, mean, covariance):
        self._posterior.mean = mean
        self._posterior.cov = covariance
        self._need_update = True
        for k in self._cache:
            self._cache[k] = None
<|endoftext|>