<commit_before>"""
Utilities to extract features from images.
"""

# Authors: Emmanuelle Gouillart <emmanuelle.gouillart@normalesup.org>
#          Gael Varoquaux <gael.varoquaux@normalesup.org>
# License: BSD

import numpy as np
from scipy import sparse
from ..utils.fixes import in1d

################################################################################
# From an image to a graph

def _make_edges_3d(n_x, n_y, n_z=1):
    """Returns a list of edges for a 3D image.

    Parameters
    ===========
    n_x: integer
        The size of the grid in the x direction.
    n_y: integer
        The size of the grid in the y direction.
    n_z: integer, optional
        The size of the grid in the z direction, defaults to 1
    """
    vertices = np.arange(n_x*n_y*n_z).reshape((n_x, n_y, n_z))
    edges_deep = np.vstack((vertices[:, :, :-1].ravel(),
                            vertices[:, :, 1:].ravel()))
    edges_right = np.vstack((vertices[:, :-1].ravel(), vertices[:, 1:].ravel()))
    edges_down = np.vstack((vertices[:-1].ravel(), vertices[1:].ravel()))
    edges = np.hstack((edges_deep, edges_right, edges_down))
    return edges

def _compute_gradient_3d(edges, img):
    n_x, n_y, n_z = img.shape
    gradient = np.abs(img[edges[0]/(n_y*n_z), \
                                (edges[0] % (n_y*n_z))/n_z, \
                                (edges[0] % (n_y*n_z))%n_z] - \
                           img[edges[1]/(n_y*n_z), \
                                (edges[1] % (n_y*n_z))/n_z, \
                                (edges[1] % (n_y*n_z)) % n_z])
    return gradient

# XXX: Why mask the image after computing the weights?

def _mask_edges_weights(mask, edges, weights=None):
    """Apply a mask to edges (weighted or not)"""
    inds = np.arange(mask.size)
    inds = inds[mask.ravel()]
    ind_mask = np.logical_and(in1d(edges[0], inds),
                              in1d(edges[1], inds))
    edges = edges[:, ind_mask]
    if weights is not None:
        weights = weights[ind_mask]
    maxval = edges.max()
    order = np.searchsorted(np.unique(edges.ravel()), np.arange(maxval+1))
    edges = order[edges]
    if weights is None:
        return edges
    else:
        return edges, weights

def _to_graph(n_x, n_y, n_z, mask=None, img=None,
              return_as=sparse.coo_matrix, dtype=None):
    """Auxiliary function for img_to_graph and grid_to_graph
    """
    edges = _make_edges_3d(n_x, n_y, n_z)

    if dtype is None:
        if img is None:
            dtype = np.bool
        else:
            dtype = img.dtype

    if img is not None:
        img = np.atleast_3d(img)
        weights = _compute_gradient_3d(edges, img)
        if mask is not None:
            edges, weights = _mask_edges_weights(mask, edges, weights)
            diag = img.squeeze()[mask]
        else:
            diag = img.ravel()
        n_voxels = diag.size
    else:
        if mask is not None:
            edges = _mask_edges_weights(mask, edges)
        weights = np.ones(edges.shape[1], dtype=dtype)
        if mask is None:
            n_voxels = n_x * n_y * n_z
        else:
            n_voxels = np.sum(mask)
        diag = np.zeros(n_voxels)

    diag_idx = np.arange(n_voxels)
    i_idx = np.hstack((edges[0], edges[1]))
    j_idx = np.hstack((edges[1], edges[0]))
    graph = sparse.coo_matrix((np.hstack((weights, weights, diag)),
                              (np.hstack((i_idx, diag_idx)),
                               np.hstack((j_idx, diag_idx)))),
                              (n_voxels, n_voxels),
                              dtype=dtype)
    if return_as is np.ndarray:
        return graph.todense()
    return return_as(graph)

def img_to_graph(img, mask=None, return_as=sparse.coo_matrix, dtype=None):
    """Graph of the pixel-to-pixel gradient connections

    Edges are weighted with the gradient values.

    Parameters
    ===========
    img: ndarray, 2D or 3D
        2D or 3D image
    mask : ndarray of booleans, optional
        An optional mask of the image, to consider only part of the
        pixels.
    return_as: np.ndarray or a sparse matrix class, optional
        The class to use to build the returned adjacency matrix.
    dtype: None or dtype, optional
        The data of the returned sparse matrix. By default it is the
        dtype of img
    """
    img = np.atleast_3d(img)
    n_x, n_y, n_z = img.shape
    return _to_graph(n_x, n_y, n_z, mask, img, return_as, dtype)

def grid_to_graph(n_x, n_y, n_z=1, mask=None, return_as=sparse.coo_matrix,
                  dtype=np.bool):
    """Graph of the pixel-to-pixel connections

    Edges exist if 2 voxels are connected.

    Parameters
    ===========
    n_x: int
        Dimension in x axis
    n_y: int
        Dimension in y axis
    n_z: int, optional, default 1
        Dimension in z axis
    mask : ndarray of booleans, optional
        An optional mask of the image, to consider only part of the
        pixels.
    return_as: np.ndarray or a sparse matrix class, optional
        The class to use to build the returned adjacency matrix.
    dtype: dtype, optional, default bool
        The data of the returned sparse matrix. By default it is bool
    """
    return _to_graph(n_x, n_y, n_z, mask=mask, return_as=return_as, dtype=dtype)

<commit_msg>ENH : ones on diag in grid_to_graph + fix dtype<commit_after>"""
Utilities to extract features from images.
"""

# Authors: Emmanuelle Gouillart <emmanuelle.gouillart@normalesup.org>
#          Gael Varoquaux <gael.varoquaux@normalesup.org>
# License: BSD

import numpy as np
from scipy import sparse
from ..utils.fixes import in1d

################################################################################
# From an image to a graph

def _make_edges_3d(n_x, n_y, n_z=1):
    """Returns a list of edges for a 3D image.

    Parameters
    ===========
    n_x: integer
        The size of the grid in the x direction.
    n_y: integer
        The size of the grid in the y direction.
    n_z: integer, optional
        The size of the grid in the z direction, defaults to 1
    """
    vertices = np.arange(n_x*n_y*n_z).reshape((n_x, n_y, n_z))
    edges_deep = np.vstack((vertices[:, :, :-1].ravel(),
                            vertices[:, :, 1:].ravel()))
    edges_right = np.vstack((vertices[:, :-1].ravel(), vertices[:, 1:].ravel()))
    edges_down = np.vstack((vertices[:-1].ravel(), vertices[1:].ravel()))
    edges = np.hstack((edges_deep, edges_right, edges_down))
    return edges

def _compute_gradient_3d(edges, img):
    n_x, n_y, n_z = img.shape
    gradient = np.abs(img[edges[0]/(n_y*n_z), \
                                (edges[0] % (n_y*n_z))/n_z, \
                                (edges[0] % (n_y*n_z))%n_z] - \
                           img[edges[1]/(n_y*n_z), \
                                (edges[1] % (n_y*n_z))/n_z, \
                                (edges[1] % (n_y*n_z)) % n_z])
    return gradient

# XXX: Why mask the image after computing the weights?

def _mask_edges_weights(mask, edges, weights=None):
    """Apply a mask to edges (weighted or not)"""
    inds = np.arange(mask.size)
    inds = inds[mask.ravel()]
    ind_mask = np.logical_and(in1d(edges[0], inds),
                              in1d(edges[1], inds))
    edges = edges[:, ind_mask]
    if weights is not None:
        weights = weights[ind_mask]
    maxval = edges.max()
    order = np.searchsorted(np.unique(edges.ravel()), np.arange(maxval+1))
    edges = order[edges]
    if weights is None:
        return edges
    else:
        return edges, weights

def _to_graph(n_x, n_y, n_z, mask=None, img=None,
              return_as=sparse.coo_matrix, dtype=None):
    """Auxiliary function for img_to_graph and grid_to_graph
    """
    edges = _make_edges_3d(n_x, n_y, n_z)

    if dtype is None:
        if img is None:
            dtype = np.bool
        else:
            dtype = img.dtype

    if img is not None:
        img = np.atleast_3d(img)
        weights = _compute_gradient_3d(edges, img)
        if mask is not None:
            edges, weights = _mask_edges_weights(mask, edges, weights)
            diag = img.squeeze()[mask]
        else:
            diag = img.ravel()
        n_voxels = diag.size
    else:
        if mask is not None:
            edges = _mask_edges_weights(mask, edges)
        weights = np.ones(edges.shape[1], dtype=dtype)
        if mask is None:
            n_voxels = n_x * n_y * n_z
        else:
            n_voxels = np.sum(mask)
        diag = np.ones(n_voxels, dtype=dtype)

    diag_idx = np.arange(n_voxels)
    i_idx = np.hstack((edges[0], edges[1]))
    j_idx = np.hstack((edges[1], edges[0]))
    graph = sparse.coo_matrix((np.hstack((weights, weights, diag)),
                              (np.hstack((i_idx, diag_idx)),
                               np.hstack((j_idx, diag_idx)))),
                              (n_voxels, n_voxels),
                              dtype=dtype)
    if return_as is np.ndarray:
        return graph.todense()
    return return_as(graph)

def img_to_graph(img, mask=None, return_as=sparse.coo_matrix, dtype=None):
    """Graph of the pixel-to-pixel gradient connections

    Edges are weighted with the gradient values.

    Parameters
    ===========
    img: ndarray, 2D or 3D
        2D or 3D image
    mask : ndarray of booleans, optional
        An optional mask of the image, to consider only part of the
        pixels.
    return_as: np.ndarray or a sparse matrix class, optional
        The class to use to build the returned adjacency matrix.
    dtype: None or dtype, optional
        The data of the returned sparse matrix. By default it is the
        dtype of img
    """
    img = np.atleast_3d(img)
    n_x, n_y, n_z = img.shape
    return _to_graph(n_x, n_y, n_z, mask, img, return_as, dtype)

def grid_to_graph(n_x, n_y, n_z=1, mask=None, return_as=sparse.coo_matrix,
                  dtype=np.bool):
    """Graph of the pixel-to-pixel connections

    Edges exist if 2 voxels are connected.

    Parameters
    ===========
    n_x: int
        Dimension in x axis
    n_y: int
        Dimension in y axis
    n_z: int, optional, default 1
        Dimension in z axis
    mask : ndarray of booleans, optional
        An optional mask of the image, to consider only part of the
        pixels.
    return_as: np.ndarray or a sparse matrix class, optional
        The class to use to build the returned adjacency matrix.
    dtype: dtype, optional, default bool
        The data of the returned sparse matrix. By default it is bool
    """
    return _to_graph(n_x, n_y, n_z, mask=mask, return_as=return_as, dtype=dtype)

<|endoftext|>


<commit_before>from django.contrib.auth.models import User
from django.db import models
from django.db.models.signals import post_save
from django.dispatch.dispatcher import receiver
from django.template.loader import render_to_string
from django.utils.translation import ugettext_lazy as _
from django_enumfield import enum

from osmaxx.conversion.converters.detail_levels import DETAIL_LEVEL_CHOICES, DETAIL_LEVEL_ALL
from osmaxx.conversion_api import coordinate_reference_systems as crs
from osmaxx.excerptexport.models.output_file import OutputFile
from .excerpt import Excerpt

# TODO: remove ExtractionOrderState,
# TODO:   since this is obsolete and replaced with osmaxx.conversion_api.statuses.STATUS_CHOICES
class ExtractionOrderState(enum.Enum):
    UNDEFINED = 0
    INITIALIZED = 1
    QUEUED = 2
    PROCESSING = 3
    FINISHED = 4
    FAILED = 6

FINAL_STATES = {ExtractionOrderState.FINISHED, ExtractionOrderState.FAILED}

CONVERSION_PROGRESS_TO_EXTRACTION_ORDER_STATE_MAPPING = {
    'new': ExtractionOrderState.INITIALIZED,
    'received': ExtractionOrderState.QUEUED,
    'started': ExtractionOrderState.PROCESSING,
    'successful': ExtractionOrderState.FINISHED,
    'error': ExtractionOrderState.FAILED,
}

def get_order_status_from_conversion_progress(progress):
    return CONVERSION_PROGRESS_TO_EXTRACTION_ORDER_STATE_MAPPING.get(progress, ExtractionOrderState.UNDEFINED)

class ExtractionOrder(models.Model):
    DOWNLOAD_STATUS_NOT_DOWNLOADED = 0
    DOWNLOAD_STATUS_DOWNLOADING = 1
    DOWNLOAD_STATUS_AVAILABLE = 2

    DOWNLOAD_STATUSES = (
        (DOWNLOAD_STATUS_NOT_DOWNLOADED, 'unknown'),
        (DOWNLOAD_STATUS_DOWNLOADING, 'downloading'),
        (DOWNLOAD_STATUS_AVAILABLE, 'received'),
    )

    state = enum.EnumField(ExtractionOrderState, default=ExtractionOrderState.INITIALIZED, verbose_name=_('state'))
    coordinate_reference_system = models.IntegerField(
        verbose_name=_('CRS'), choices=crs.CRS_CHOICES, default=crs.WGS_84
    )
    detail_level = models.IntegerField(
        verbose_name=_('detail level'), choices=DETAIL_LEVEL_CHOICES, default=DETAIL_LEVEL_ALL
    )
    process_id = models.TextField(blank=True, null=True, verbose_name=_('process link'))
    orderer = models.ForeignKey(User, related_name='extraction_orders', verbose_name=_('orderer'))
    excerpt = models.ForeignKey(Excerpt, related_name='extraction_orders', verbose_name=_('excerpt'), null=True)
    progress_url = models.URLField(verbose_name=_('progress URL'), null=True, blank=True)
    process_start_time = models.DateTimeField(verbose_name=_('process start time'), null=True, blank=True)
    download_status = models.IntegerField(
        _('file status'),
        choices=DOWNLOAD_STATUSES,
        default=DOWNLOAD_STATUS_NOT_DOWNLOADED
    )

    def forward_to_conversion_service(self, *, incoming_request):
        clipping_area_json = self.excerpt.send_to_conversion_service()
        jobs_json = [
            export.send_to_conversion_service(clipping_area_json, incoming_request)
            for export in self.exports.all()
        ]
        return jobs_json

    def __str__(self):
        return ', '.join(
            [
                '[{order_id}] orderer: {orderer_name}'.format(
                    order_id=self.id,
                    orderer_name=self.orderer.get_username(),
                ),
                'excerpt: {}'.format(str(self.excerpt_name)),
                'state: {}'.format(self.get_state_display()),
            ]
        )

    @property
    def excerpt_name(self):
        """
        Returns:
              user-given excerpt name for user-defined excerpts,
              country name for countries,
              None if order has no excerpt (neither country nor user-defined)
        """
        if self.excerpt:
            return self.excerpt.name

    @property
    def output_files(self):
        return OutputFile.objects.filter(export__extraction_order=self)

    @property
    def are_downloads_ready(self):
        return self.state == ExtractionOrderState.FINISHED

    @property
    def extraction_formats(self):
        return self.exports.values_list('file_format', flat=True)

    @extraction_formats.setter
    def extraction_formats(self, value):
        new_formats = frozenset(value)
        previous_formats = self.exports.values_list('file_format', flat=True)
        assert new_formats.issuperset(previous_formats)
        self._new_formats = new_formats  # Will be collected and cleaned up by attach_new_formats.
        if self.id is not None:
            attach_new_formats(self.__class__, instance=self)

    def set_status_from_conversion_progress(self, job_overall_progress):
        if self.state not in [ExtractionOrderState.FINISHED, ExtractionOrderState.FAILED]:
            self.state = get_order_status_from_conversion_progress(job_overall_progress)

    def get_absolute_url(self):
        from django.core.urlresolvers import reverse
        return reverse('excerptexport:export_list')

    def send_email_if_all_exports_done(self, incoming_request):
        if all(export.is_status_final for export in self.exports.all()):
            from osmaxx.utilities.shortcuts import Emissary
            emissary = Emissary(recipient=self.orderer)
            emissary.inform_mail(
                subject=self._get_all_exports_done_email_subject(),
                mail_body=self._get_all_exports_done_mail_body(incoming_request)
            )

    def _get_all_exports_done_email_subject(self):
        view_context = dict(
            extraction_order=self,
            successful_exports_count=self.exports.filter(output_file__isnull=False).count(),
            failed_exports_count=self.exports.filter(output_file__isnull=True).count(),
        )
        return render_to_string(
            'excerptexport/email/all_exports_of_extraction_order_done_subject.txt',
            context=view_context,
        ).strip()

    def _get_all_exports_done_mail_body(self, incoming_request):
        view_context = dict(
            extraction_order=self,
            successful_exports=self.exports.filter(output_file__isnull=False),
            failed_exports=self.exports.filter(output_file__isnull=True),
            request=incoming_request,
        )
        return render_to_string(
            'excerptexport/email/all_exports_of_extraction_order_done_body.txt',
            context=view_context,
        ).strip()

@receiver(post_save, sender=ExtractionOrder)
def attach_new_formats(sender, instance, **kwargs):
    if hasattr(instance, '_new_formats'):
        for format in instance._new_formats:
            instance.exports.get_or_create(file_format=format)
        del instance._new_formats
<commit_msg>remove old unused method<commit_after>from django.contrib.auth.models import User
from django.db import models
from django.db.models.signals import post_save
from django.dispatch.dispatcher import receiver
from django.template.loader import render_to_string
from django.utils.translation import ugettext_lazy as _
from django_enumfield import enum

from osmaxx.conversion.converters.detail_levels import DETAIL_LEVEL_CHOICES, DETAIL_LEVEL_ALL
from osmaxx.conversion_api import coordinate_reference_systems as crs
from osmaxx.excerptexport.models.output_file import OutputFile
from .excerpt import Excerpt

# TODO: remove ExtractionOrderState,
# TODO:   since this is obsolete and replaced with osmaxx.conversion_api.statuses.STATUS_CHOICES
class ExtractionOrderState(enum.Enum):
    UNDEFINED = 0
    INITIALIZED = 1
    QUEUED = 2
    PROCESSING = 3
    FINISHED = 4
    FAILED = 6

FINAL_STATES = {ExtractionOrderState.FINISHED, ExtractionOrderState.FAILED}

CONVERSION_PROGRESS_TO_EXTRACTION_ORDER_STATE_MAPPING = {
    'new': ExtractionOrderState.INITIALIZED,
    'received': ExtractionOrderState.QUEUED,
    'started': ExtractionOrderState.PROCESSING,
    'successful': ExtractionOrderState.FINISHED,
    'error': ExtractionOrderState.FAILED,
}

def get_order_status_from_conversion_progress(progress):
    return CONVERSION_PROGRESS_TO_EXTRACTION_ORDER_STATE_MAPPING.get(progress, ExtractionOrderState.UNDEFINED)

class ExtractionOrder(models.Model):
    DOWNLOAD_STATUS_NOT_DOWNLOADED = 0
    DOWNLOAD_STATUS_DOWNLOADING = 1
    DOWNLOAD_STATUS_AVAILABLE = 2

    DOWNLOAD_STATUSES = (
        (DOWNLOAD_STATUS_NOT_DOWNLOADED, 'unknown'),
        (DOWNLOAD_STATUS_DOWNLOADING, 'downloading'),
        (DOWNLOAD_STATUS_AVAILABLE, 'received'),
    )

    state = enum.EnumField(ExtractionOrderState, default=ExtractionOrderState.INITIALIZED, verbose_name=_('state'))
    coordinate_reference_system = models.IntegerField(
        verbose_name=_('CRS'), choices=crs.CRS_CHOICES, default=crs.WGS_84
    )
    detail_level = models.IntegerField(
        verbose_name=_('detail level'), choices=DETAIL_LEVEL_CHOICES, default=DETAIL_LEVEL_ALL
    )
    process_id = models.TextField(blank=True, null=True, verbose_name=_('process link'))
    orderer = models.ForeignKey(User, related_name='extraction_orders', verbose_name=_('orderer'))
    excerpt = models.ForeignKey(Excerpt, related_name='extraction_orders', verbose_name=_('excerpt'), null=True)
    progress_url = models.URLField(verbose_name=_('progress URL'), null=True, blank=True)
    process_start_time = models.DateTimeField(verbose_name=_('process start time'), null=True, blank=True)
    download_status = models.IntegerField(
        _('file status'),
        choices=DOWNLOAD_STATUSES,
        default=DOWNLOAD_STATUS_NOT_DOWNLOADED
    )

    def forward_to_conversion_service(self, *, incoming_request):
        clipping_area_json = self.excerpt.send_to_conversion_service()
        jobs_json = [
            export.send_to_conversion_service(clipping_area_json, incoming_request)
            for export in self.exports.all()
        ]
        return jobs_json

    def __str__(self):
        return ', '.join(
            [
                '[{order_id}] orderer: {orderer_name}'.format(
                    order_id=self.id,
                    orderer_name=self.orderer.get_username(),
                ),
                'excerpt: {}'.format(str(self.excerpt_name)),
                'state: {}'.format(self.get_state_display()),
            ]
        )

    @property
    def excerpt_name(self):
        """
        Returns:
              user-given excerpt name for user-defined excerpts,
              country name for countries,
              None if order has no excerpt (neither country nor user-defined)
        """
        if self.excerpt:
            return self.excerpt.name

    @property
    def output_files(self):
        return OutputFile.objects.filter(export__extraction_order=self)

    @property
    def are_downloads_ready(self):
        return self.state == ExtractionOrderState.FINISHED

    @property
    def extraction_formats(self):
        return self.exports.values_list('file_format', flat=True)

    @extraction_formats.setter
    def extraction_formats(self, value):
        new_formats = frozenset(value)
        previous_formats = self.exports.values_list('file_format', flat=True)
        assert new_formats.issuperset(previous_formats)
        self._new_formats = new_formats  # Will be collected and cleaned up by attach_new_formats.
        if self.id is not None:
            attach_new_formats(self.__class__, instance=self)

    def get_absolute_url(self):
        from django.core.urlresolvers import reverse
        return reverse('excerptexport:export_list')

    def send_email_if_all_exports_done(self, incoming_request):
        if all(export.is_status_final for export in self.exports.all()):
            from osmaxx.utilities.shortcuts import Emissary
            emissary = Emissary(recipient=self.orderer)
            emissary.inform_mail(
                subject=self._get_all_exports_done_email_subject(),
                mail_body=self._get_all_exports_done_mail_body(incoming_request)
            )

    def _get_all_exports_done_email_subject(self):
        view_context = dict(
            extraction_order=self,
            successful_exports_count=self.exports.filter(output_file__isnull=False).count(),
            failed_exports_count=self.exports.filter(output_file__isnull=True).count(),
        )
        return render_to_string(
            'excerptexport/email/all_exports_of_extraction_order_done_subject.txt',
            context=view_context,
        ).strip()

    def _get_all_exports_done_mail_body(self, incoming_request):
        view_context = dict(
            extraction_order=self,
            successful_exports=self.exports.filter(output_file__isnull=False),
            failed_exports=self.exports.filter(output_file__isnull=True),
            request=incoming_request,
        )
        return render_to_string(
            'excerptexport/email/all_exports_of_extraction_order_done_body.txt',
            context=view_context,
        ).strip()

@receiver(post_save, sender=ExtractionOrder)
def attach_new_formats(sender, instance, **kwargs):
    if hasattr(instance, '_new_formats'):
        for format in instance._new_formats:
            instance.exports.get_or_create(file_format=format)
        del instance._new_formats
<|endoftext|>


<commit_before>from django.contrib import admin
from jawns.models import Content, They

class TheyAdmin(admin.ModelAdmin):
    model = They
    def get_context_data(admin.ModelAdmin):
        context = super(TheyAdmin, self).get_context_data(**kwargs)
        return context

class ContentAdmin(admin.ModelAdmin):
    model = Content
    def get_context_data(self, **kwargs):
        context = super(ContentAdmin, self).get_context_data(**kwargs)
        return context

admin.site.register(Content, ContentAdmin)
<commit_msg>Whoops<commit_after>from django.contrib import admin
from jawns.models import Content, They

class TheyAdmin(admin.ModelAdmin):
    model = They
    def get_context_data(self, **kwargs):
        context = super(TheyAdmin, self).get_context_data(**kwargs)
        return context

class ContentAdmin(admin.ModelAdmin):
    model = Content
    def get_context_data(self, **kwargs):
        context = super(ContentAdmin, self).get_context_data(**kwargs)
        return context

admin.site.register(Content, ContentAdmin)
<|endoftext|>


<commit_before># This code is original from jsmin by Douglas Crockford, it was translated to
# Python by Baruch Even. The original code had the following copyright and
# license.
#
# /* jsmin.c
#	2007-05-22
#
# Copyright (c) 2002 Douglas Crockford  (www.crockford.com)
#
# Permission is hereby granted, free of charge, to any person obtaining a copy of
# this software and associated documentation files (the "Software"), to deal in
# the Software without restriction, including without limitation the rights to
# use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies
# of the Software, and to permit persons to whom the Software is furnished to do
# so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# The Software shall be used for Good, not Evil.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
# */
from __future__ import unicode_literals

from StringIO import StringIO

def jsmin(js):
	ins = StringIO(js)
	outs = StringIO()
	JavascriptMinify().minify(ins, outs)
	str = outs.getvalue()
	if len(str) > 0 and str[0] == '\n':
		str = str[1:]
	return str

def isAlphanum(c):
	"""return true if the character is a letter, digit, underscore,
		   dollar sign, or non-ASCII character.
	"""
	return ((c >= 'a' and c <= 'z') or (c >= '0' and c <= '9') or
			(c >= 'A' and c <= 'Z') or c == '_' or c == '$' or c == '\\' or (c is not None and ord(c) > 126));

class UnterminatedComment(Exception):
	pass

class UnterminatedStringLiteral(Exception):
	pass

class UnterminatedRegularExpression(Exception):
	pass

class JavascriptMinify(object):

	def _outA(self):
		self.outstream.write(self.theA)
	def _outB(self):
		self.outstream.write(self.theB)

	def _get(self):
		"""return the next character from stdin. Watch out for lookahead. If
		   the character is a control character, translate it to a space or
		   linefeed.
		"""
		c = self.theLookahead
		self.theLookahead = None
		if c == None:
			c = self.instream.read(1)
		if c >= ' ' or c == '\n':
			return c
		if c == '': # EOF
			return '\000'
		if c == '\r':
			return '\n'
		return ' '

	def _peek(self):
		self.theLookahead = self._get()
		return self.theLookahead

	def _next(self):
		"""get the next character, excluding comments. peek() is used to see
		   if an unescaped '/' is followed by a '/' or '*'.
		"""
		c = self._get()
		if c == '/' and self.theA != '\\':
			p = self._peek()
			if p == '/':
				c = self._get()
				while c > '\n':
					c = self._get()
				return c
			if p == '*':
				c = self._get()
				while 1:
					c = self._get()
					if c == '*':
						if self._peek() == '/':
							self._get()
							return ' '
					if c == '\000':
						raise UnterminatedComment()

		return c

	def _action(self, action):
		"""do something! What you do is determined by the argument:
		   1   Output A. Copy B to A. Get the next B.
		   2   Copy B to A. Get the next B. (Delete A).
		   3   Get the next B. (Delete B).
		   action treats a string as a single character. Wow!
		   action recognizes a regular expression if it is preceded by ( or , or =.
		"""
		if action <= 1:
			self._outA()

		if action <= 2:
			self.theA = self.theB
			if self.theA == "'" or self.theA == '"':
				while 1:
					self._outA()
					self.theA = self._get()
					if self.theA == self.theB:
						break
					if self.theA <= '\n':
						raise UnterminatedStringLiteral()
					if self.theA == '\\':
						self._outA()
						self.theA = self._get()

		if action <= 3:
			self.theB = self._next()
			if self.theB == '/' and (self.theA == '(' or self.theA == ',' or
									 self.theA == '=' or self.theA == ':' or
									 self.theA == '[' or self.theA == '?' or
									 self.theA == '!' or self.theA == '&' or
									 self.theA == '|' or self.theA == ';' or
									 self.theA == '{' or self.theA == '}' or
									 self.theA == '\n'):
				self._outA()
				self._outB()
				while 1:
					self.theA = self._get()
					if self.theA == '/':
						break
					elif self.theA == '\\':
						self._outA()
						self.theA = self._get()
					elif self.theA <= '\n':
						raise UnterminatedRegularExpression()
					self._outA()
				self.theB = self._next()

	def _jsmin(self):
		"""Copy the input to the output, deleting the characters which are
		   insignificant to JavaScript. Comments will be removed. Tabs will be
		   replaced with spaces. Carriage returns will be replaced with linefeeds.
		   Most spaces and linefeeds will be removed.
		"""
		self.theA = '\n'
		self._action(3)

		while self.theA != '\000':
			if self.theA == ' ':
				if isAlphanum(self.theB):
					self._action(1)
				else:
					self._action(2)
			elif self.theA == '\n':
				if self.theB in ['{', '[', '(', '+', '-']:
					self._action(1)
				elif self.theB == ' ':
					self._action(3)
				else:
					if isAlphanum(self.theB):
						self._action(1)
					else:
						self._action(2)
			else:
				if self.theB == ' ':
					if isAlphanum(self.theA):
						self._action(1)
					else:
						self._action(3)
				elif self.theB == '\n':
					if self.theA in ['}', ']', ')', '+', '-', '"', '\'']:
						self._action(1)
					else:
						if isAlphanum(self.theA):
							self._action(1)
						else:
							self._action(3)
				else:
					self._action(1)

	def minify(self, instream, outstream):
		self.instream = instream
		self.outstream = outstream
		self.theA = '\n'
		self.theB = None
		self.theLookahead = None

		self._jsmin()
		self.instream.close()<commit_msg>[hotfix] removed unicode literals from minify.py<commit_after># This code is original from jsmin by Douglas Crockford, it was translated to
# Python by Baruch Even. The original code had the following copyright and
# license.
#
# /* jsmin.c
#	2007-05-22
#
# Copyright (c) 2002 Douglas Crockford  (www.crockford.com)
#
# Permission is hereby granted, free of charge, to any person obtaining a copy of
# this software and associated documentation files (the "Software"), to deal in
# the Software without restriction, including without limitation the rights to
# use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies
# of the Software, and to permit persons to whom the Software is furnished to do
# so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# The Software shall be used for Good, not Evil.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
# */

from StringIO import StringIO

def jsmin(js):
	ins = StringIO(js)
	outs = StringIO()
	JavascriptMinify().minify(ins, outs)
	str = outs.getvalue()
	if len(str) > 0 and str[0] == '\n':
		str = str[1:]
	return str

def isAlphanum(c):
	"""return true if the character is a letter, digit, underscore,
		   dollar sign, or non-ASCII character.
	"""
	return ((c >= 'a' and c <= 'z') or (c >= '0' and c <= '9') or
			(c >= 'A' and c <= 'Z') or c == '_' or c == '$' or c == '\\' or (c is not None and ord(c) > 126));

class UnterminatedComment(Exception):
	pass

class UnterminatedStringLiteral(Exception):
	pass

class UnterminatedRegularExpression(Exception):
	pass

class JavascriptMinify(object):

	def _outA(self):
		self.outstream.write(self.theA)
	def _outB(self):
		self.outstream.write(self.theB)

	def _get(self):
		"""return the next character from stdin. Watch out for lookahead. If
		   the character is a control character, translate it to a space or
		   linefeed.
		"""
		c = self.theLookahead
		self.theLookahead = None
		if c == None:
			c = self.instream.read(1)
		if c >= ' ' or c == '\n':
			return c
		if c == '': # EOF
			return '\000'
		if c == '\r':
			return '\n'
		return ' '

	def _peek(self):
		self.theLookahead = self._get()
		return self.theLookahead

	def _next(self):
		"""get the next character, excluding comments. peek() is used to see
		   if an unescaped '/' is followed by a '/' or '*'.
		"""
		c = self._get()
		if c == '/' and self.theA != '\\':
			p = self._peek()
			if p == '/':
				c = self._get()
				while c > '\n':
					c = self._get()
				return c
			if p == '*':
				c = self._get()
				while 1:
					c = self._get()
					if c == '*':
						if self._peek() == '/':
							self._get()
							return ' '
					if c == '\000':
						raise UnterminatedComment()

		return c

	def _action(self, action):
		"""do something! What you do is determined by the argument:
		   1   Output A. Copy B to A. Get the next B.
		   2   Copy B to A. Get the next B. (Delete A).
		   3   Get the next B. (Delete B).
		   action treats a string as a single character. Wow!
		   action recognizes a regular expression if it is preceded by ( or , or =.
		"""
		if action <= 1:
			self._outA()

		if action <= 2:
			self.theA = self.theB
			if self.theA == "'" or self.theA == '"':
				while 1:
					self._outA()
					self.theA = self._get()
					if self.theA == self.theB:
						break
					if self.theA <= '\n':
						raise UnterminatedStringLiteral()
					if self.theA == '\\':
						self._outA()
						self.theA = self._get()

		if action <= 3:
			self.theB = self._next()
			if self.theB == '/' and (self.theA == '(' or self.theA == ',' or
									 self.theA == '=' or self.theA == ':' or
									 self.theA == '[' or self.theA == '?' or
									 self.theA == '!' or self.theA == '&' or
									 self.theA == '|' or self.theA == ';' or
									 self.theA == '{' or self.theA == '}' or
									 self.theA == '\n'):
				self._outA()
				self._outB()
				while 1:
					self.theA = self._get()
					if self.theA == '/':
						break
					elif self.theA == '\\':
						self._outA()
						self.theA = self._get()
					elif self.theA <= '\n':
						raise UnterminatedRegularExpression()
					self._outA()
				self.theB = self._next()

	def _jsmin(self):
		"""Copy the input to the output, deleting the characters which are
		   insignificant to JavaScript. Comments will be removed. Tabs will be
		   replaced with spaces. Carriage returns will be replaced with linefeeds.
		   Most spaces and linefeeds will be removed.
		"""
		self.theA = '\n'
		self._action(3)

		while self.theA != '\000':
			if self.theA == ' ':
				if isAlphanum(self.theB):
					self._action(1)
				else:
					self._action(2)
			elif self.theA == '\n':
				if self.theB in ['{', '[', '(', '+', '-']:
					self._action(1)
				elif self.theB == ' ':
					self._action(3)
				else:
					if isAlphanum(self.theB):
						self._action(1)
					else:
						self._action(2)
			else:
				if self.theB == ' ':
					if isAlphanum(self.theA):
						self._action(1)
					else:
						self._action(3)
				elif self.theB == '\n':
					if self.theA in ['}', ']', ')', '+', '-', '"', '\'']:
						self._action(1)
					else:
						if isAlphanum(self.theA):
							self._action(1)
						else:
							self._action(3)
				else:
					self._action(1)

	def minify(self, instream, outstream):
		self.instream = instream
		self.outstream = outstream
		self.theA = '\n'
		self.theB = None
		self.theLookahead = None

		self._jsmin()
		self.instream.close()
<|endoftext|>


<commit_before>#!/usr/bin/env python
# Copyright 2015-2021 Yelp Inc.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Small utility to update the Prometheus adapter's config to match soaconfigs.
"""
import argparse
import logging
import re
import sys
from pathlib import Path
from typing import cast
from typing import Dict
from typing import List
from typing import Optional
from typing import Tuple

import ruamel.yaml as yaml
from kubernetes.client import V1ConfigMap
from kubernetes.client import V1DeleteOptions
from kubernetes.client import V1ObjectMeta
from kubernetes.client.rest import ApiException
from mypy_extensions import TypedDict

from paasta_tools.kubernetes_tools import ensure_namespace
from paasta_tools.kubernetes_tools import get_kubernetes_app_name
from paasta_tools.kubernetes_tools import KubeClient
from paasta_tools.kubernetes_tools import KubernetesDeploymentConfig
from paasta_tools.kubernetes_tools import V1Pod
from paasta_tools.long_running_service_tools import AutoscalingParamsDict
from paasta_tools.paasta_service_config_loader import PaastaServiceConfigLoader
from paasta_tools.utils import DEFAULT_SOA_DIR
from paasta_tools.utils import get_services_for_cluster

log = logging.getLogger(__name__)

PROMETHEUS_ADAPTER_CONFIGMAP_NAMESPACE = "custom-metrics"
PROMETHEUS_ADAPTER_POD_NAMESPACE = "custom-metrics"
PROMETHEUS_ADAPTER_CONFIGMAP_NAME = "adapter-config"
PROMETHEUS_ADAPTER_CONFIGMAP_FILENAME = "config.yaml"
PROMETHEUS_ADAPTER_POD_NAME_PREFIX = "custom-metrics-apiserver"
PROMETHEUS_ADAPTER_POD_PHASES_TO_REMOVE = (
    "Running",
    "Pending",
)

DEFAULT_SCRAPE_PERIOD_S = 10
DEFAULT_EXTRAPOLATION_PERIODS = 10
DEFAULT_EXTRAPOLATION_TIME = DEFAULT_SCRAPE_PERIOD_S * DEFAULT_EXTRAPOLATION_PERIODS

class PrometheusAdapterRule(TypedDict):
    """
    Typed version of the (minimal) set of Prometheus adapter rule configuration options that we use
    """

    # see https://github.com/DirectXMan12/k8s-prometheus-adapter/blob/master/docs/config.md
    # for more detailed information
    # used for discovering what resources should be scaled
    seriesQuery: str
    # configuration for how to expose this rule to the HPA
    name: Dict[str, str]
    # used to associate metrics with k8s resources
    resources: Dict[str, str]
    # the actual query we want to send to Prometheus to use for scaling
    metricsQuery: str

class PrometheusAdapterConfig(TypedDict):
    """
    Typed version of the Prometheus adapter configuration dictionary.
    """

    rules: List[PrometheusAdapterRule]

def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Syncs the Prometheus metric adapter config with soaconfigs.",
    )

    parser.add_argument(
        "-d",
        "--soa-dir",
        dest="soa_dir",
        metavar="SOA_DIR",
        default=Path(DEFAULT_SOA_DIR),
        help="Directory to read service configs from. Default is %(default)s.",
        type=Path,
    )
    # TODO: do we need to be able to pass multiple clusters in?
    parser.add_argument(
        "-c",
        "--cluster",
        dest="cluster",
        help="PaaSTA cluster to generate configs for.",
        required=True,
    )
    parser.add_argument(
        "-v",
        "--verbose",
        action="store_true",
        dest="verbose",
        default=False,
        help="Enable verbose logging.",
    )
    parser.add_argument(
        "--dry-run",
        dest="dry_run",
        action="store_true",
        default=False,
        help="Enable verbose logging.",
    )

    return parser.parse_args()

def should_create_uwsgi_scaling_rule(
    instance: str, autoscaling_config: AutoscalingParamsDict,
) -> Tuple[bool, Optional[str]]:
    """
    Determines whether we should configure the prometheus adapter for a given service.
    Returns a 2-tuple of (should_create, reason_to_skip)
    """
    if autoscaling_config["metrics_provider"] == "uwsgi":
        if not autoscaling_config.get("use_prometheus", False):
            return False, "requested uwsgi autoscaling, but not using Prometheus"

        return True, None

    return False, "did not request uwsgi autoscaling"

def create_instance_uwsgi_scaling_rule(
    service: str,
    instance: str,
    autoscaling_config: AutoscalingParamsDict,
    paasta_cluster: str,
) -> PrometheusAdapterRule:
    """
    Creates a Prometheus adapter rule config for a given service instance.
    """
    setpoint = autoscaling_config["setpoint"]
    moving_average_window = autoscaling_config["moving_average_window_seconds"]
    # this should always be set, but we default to 0 for safety as the worst thing that would happen
    # is that we take a couple more iterations than required to hit the desired setpoint
    offset = autoscaling_config.get("offset", 0)
    deployment_name = get_kubernetes_app_name(service=service, instance=instance)
    worker_filter_terms = f"paasta_cluster='{paasta_cluster}',paasta_service='{service}',paasta_instance='{instance}'"
    replica_filter_terms = (
        f"paasta_cluster='{paasta_cluster}',deployment='{deployment_name}'"
    )

    current_replicas = f"""
        (
            scalar(
                kube_deployment_spec_replicas{{{replica_filter_terms}}} >= 0
                or
                max_over_time(
                    kube_deployment_spec_replicas{{{replica_filter_terms}}}[{DEFAULT_EXTRAPOLATION_TIME}s]
                )
            )
        )
    """
    load_per_instance = f"""
        avg(
            uwsgi_worker_busy{{{worker_filter_terms}}}
        ) by (kube_pod, kube_deployment)
    """
    desired_instances_at_each_point_in_time = f"""
        sum(
            {load_per_instance}
        ) by (kube_deployment) / {setpoint - offset}
    """
    desired_instances = f"""
        avg_over_time(
            (
                {desired_instances_at_each_point_in_time}
            )[{moving_average_window}s:]
        )
    """
    metrics_query = f"""
        {desired_instances} / {current_replicas}
    """

    metric_name = f"{deployment_name}-uwsgi-prom"

    return {
        "name": {"as": metric_name},
        "seriesQuery": f"uwsgi_worker_busy{{{worker_filter_terms}}}",
        "resources": {"template": "kube_<<.Resource>>"},
        # our nicely formatted query has enough whitespace that collapsing said
        # whitespace cuts down the size of the string by a meaningful amount
        "metricsQuery": re.sub(pattern=r"\s+", repl=" ", string=metrics_query).strip(),
    }

def get_rules_for_service_instance(
    service_name: str,
    instance_name: str,
    autoscaling_config: AutoscalingParamsDict,
    paasta_cluster: str,
) -> List[PrometheusAdapterRule]:
    """
    Returns a list of Prometheus Adapter rules for a given service instance. For now, this
    will always be a 0 or 1-element list - but when we support scaling on multiple metrics
    we will return N rules for a given service instance.
    """
    rules: List[PrometheusAdapterRule] = []

    should_create_uwsgi, skip_uwsgi_reason = should_create_uwsgi_scaling_rule(
        instance=instance_name, autoscaling_config=autoscaling_config,
    )
    if should_create_uwsgi:
        rules.append(
            create_instance_uwsgi_scaling_rule(
                service=service_name,
                instance=instance_name,
                autoscaling_config=autoscaling_config,
                paasta_cluster=paasta_cluster,
            )
        )
    else:
        log.debug(
            "Skipping %s.%s - %s.", service_name, instance_name, skip_uwsgi_reason,
        )

    return rules

def create_prometheus_adapter_config(
    paasta_cluster: str, soa_dir: Path
) -> PrometheusAdapterConfig:
    """
    Given a paasta cluster and a soaconfigs directory, create the necessary Prometheus adapter
    config to autoscale services.
    Currently supports the following metrics providers:
        * uwsgi
    """
    rules: List[PrometheusAdapterRule] = []
    # get_services_for_cluster() returns a list of (service, instance) tuples, but this
    # is not great for us: if we were to iterate over that we'd end up getting duplicates
    # for every service as PaastaServiceConfigLoader does not expose a way to get configs
    # for a single instance by name. instead, we get the unique set of service names and then
    # let PaastaServiceConfigLoader iterate over instances for us later
    services = {
        service_name
        for service_name, _ in get_services_for_cluster(
            cluster=paasta_cluster, instance_type="kubernetes", soa_dir=str(soa_dir)
        )
    }
    for service_name in services:
        config_loader = PaastaServiceConfigLoader(
            service=service_name, soa_dir=str(soa_dir)
        )
        for instance_config in config_loader.instance_configs(
            cluster=paasta_cluster, instance_type_class=KubernetesDeploymentConfig,
        ):
            rules.extend(
                get_rules_for_service_instance(
                    service_name=service_name,
                    instance_name=instance_config.instance,
                    autoscaling_config=instance_config.get_autoscaling_params(),
                    paasta_cluster=paasta_cluster,
                )
            )

    return {
        # we sort our rules so that we can easily compare between two different configmaps
        # as otherwise we'd need to do fancy order-independent comparisons between the two
        # sets of rules later due to the fact that we're not iterating in a deterministic
        # way and can add rules in any arbitrary order
        "rules": sorted(rules, key=lambda rule: rule["name"]["as"]),
    }

def update_prometheus_adapter_configmap(
    kube_client: KubeClient, config: PrometheusAdapterConfig
) -> None:
    kube_client.core.replace_namespaced_config_map(
        name=PROMETHEUS_ADAPTER_CONFIGMAP_NAME,
        namespace=PROMETHEUS_ADAPTER_CONFIGMAP_NAMESPACE,
        body=V1ConfigMap(
            metadata=V1ObjectMeta(name=PROMETHEUS_ADAPTER_CONFIGMAP_NAME),
            data={
                PROMETHEUS_ADAPTER_CONFIGMAP_FILENAME: yaml.dump(
                    config,
                    default_flow_style=False,
                    explicit_start=True,
                    width=sys.maxsize,
                )
            },
        ),
    )

def create_prometheus_adapter_configmap(
    kube_client: KubeClient, config: PrometheusAdapterConfig
) -> None:
    kube_client.core.create_namespaced_config_map(
        namespace=PROMETHEUS_ADAPTER_CONFIGMAP_NAMESPACE,
        body=V1ConfigMap(
            metadata=V1ObjectMeta(name=PROMETHEUS_ADAPTER_CONFIGMAP_NAME),
            data={
                PROMETHEUS_ADAPTER_CONFIGMAP_FILENAME: yaml.dump(
                    config, default_flow_style=False, explicit_start=True
                )
            },
        ),
    )

def get_prometheus_adapter_configmap(
    kube_client: KubeClient,
) -> Optional[PrometheusAdapterConfig]:
    try:
        config = cast(
            # we cast since mypy infers the wrong type since the k8s clientlib is untyped
            V1ConfigMap,
            kube_client.core.read_namespaced_config_map(
                name=PROMETHEUS_ADAPTER_CONFIGMAP_NAME,
                namespace=PROMETHEUS_ADAPTER_CONFIGMAP_NAMESPACE,
            ),
        )
    except ApiException as e:
        if e.status == 404:
            return None
        else:
            raise

    if not config:
        return None

    return yaml.safe_load(config.data[PROMETHEUS_ADAPTER_CONFIGMAP_FILENAME])

def restart_prometheus_adapter(kube_client: KubeClient) -> None:
    log.info("Attempting to remove existing adapter pod(s).")
    all_pods = cast(
        # once again, we cast since the kubernetes python api isn't typed
        List[V1Pod],
        kube_client.core.list_namespaced_pod(
            namespace=PROMETHEUS_ADAPTER_POD_NAMESPACE
        ).items,
    )
    # there should only ever be one pod actually up, but we might as well enforce that here
    # just in case there are more
    pods_to_delete = [
        pod
        for pod in all_pods
        if pod.metadata.name.startswith(PROMETHEUS_ADAPTER_POD_NAME_PREFIX)
        and pod.status.phase in PROMETHEUS_ADAPTER_POD_PHASES_TO_REMOVE
    ]
    log.debug("Found the following pods to delete: %s", pods_to_delete)

    for pod in pods_to_delete:
        log.debug("Attempting to remove %s.", pod.metadata.name)
        kube_client.core.delete_namespaced_pod(
            name=pod.metadata.name,
            namespace=pod.metadata.namespace,
            body=V1DeleteOptions(),
            # background propagation with no grace period is equivalent to doing a force-delete from kubectl
            grace_period_seconds=0,
            propagation_policy="Background",
        )
        log.debug("Removed %s.", pod.metadata.name)

    log.info("Adapter restarted successfully")

def main() -> int:
    args = parse_args()
    if args.verbose:
        logging.basicConfig(level=logging.DEBUG)
    else:
        logging.basicConfig(level=logging.INFO)

    log.info("Generating adapter config from soaconfigs.")
    config = create_prometheus_adapter_config(
        paasta_cluster=args.cluster, soa_dir=args.soa_dir
    )
    log.info("Generated adapter config from soaconfigs.")
    if args.dry_run:
        log.info(
            "Generated the following config:\n%s",
            yaml.dump(
                config, default_flow_style=False, explicit_start=True, width=sys.maxsize
            ),
        )
        return 0  # everything after this point requires creds/updates state
    else:
        log.debug(
            "Generated the following config:\n%s",
            yaml.dump(
                config, default_flow_style=False, explicit_start=True, width=sys.maxsize
            ),
        )

    if not config["rules"]:
        log.error("Got empty rule configuration - refusing to continue.")
        return 0

    kube_client = KubeClient()
    if not args.dry_run:
        ensure_namespace(kube_client, namespace="paasta")
        ensure_namespace(kube_client, namespace="custom-metrics")

    existing_config = get_prometheus_adapter_configmap(kube_client=kube_client)
    if existing_config and existing_config != config:
        log.info("Existing config differs from soaconfigs - updating.")
        log.debug("Existing data: %s", existing_config)
        log.debug("Desired data: %s", config)
        update_prometheus_adapter_configmap(kube_client=kube_client, config=config)
        log.info("Updated adapter config.")
    elif existing_config:
        log.info("Existing config matches soaconfigs - exiting.")
        return 0
    else:
        log.info("No existing config - creating.")
        create_prometheus_adapter_configmap(kube_client=kube_client, config=config)
        log.info("Created adapter config.")

    # the prometheus adapter doesn't currently have a good way to reload on config changes
    # so we do the next best thing: restart the pod so that it picks up the new config.
    # see: https://github.com/DirectXMan12/k8s-prometheus-adapter/issues/104
    restart_prometheus_adapter(kube_client=kube_client)

    return 0

if __name__ == "__main__":
    sys.exit(main())
<commit_msg>Break desired_instances_at_each_point_in_time down further.<commit_after>#!/usr/bin/env python
# Copyright 2015-2021 Yelp Inc.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Small utility to update the Prometheus adapter's config to match soaconfigs.
"""
import argparse
import logging
import re
import sys
from pathlib import Path
from typing import cast
from typing import Dict
from typing import List
from typing import Optional
from typing import Tuple

import ruamel.yaml as yaml
from kubernetes.client import V1ConfigMap
from kubernetes.client import V1DeleteOptions
from kubernetes.client import V1ObjectMeta
from kubernetes.client.rest import ApiException
from mypy_extensions import TypedDict

from paasta_tools.kubernetes_tools import ensure_namespace
from paasta_tools.kubernetes_tools import get_kubernetes_app_name
from paasta_tools.kubernetes_tools import KubeClient
from paasta_tools.kubernetes_tools import KubernetesDeploymentConfig
from paasta_tools.kubernetes_tools import V1Pod
from paasta_tools.long_running_service_tools import AutoscalingParamsDict
from paasta_tools.paasta_service_config_loader import PaastaServiceConfigLoader
from paasta_tools.utils import DEFAULT_SOA_DIR
from paasta_tools.utils import get_services_for_cluster

log = logging.getLogger(__name__)

PROMETHEUS_ADAPTER_CONFIGMAP_NAMESPACE = "custom-metrics"
PROMETHEUS_ADAPTER_POD_NAMESPACE = "custom-metrics"
PROMETHEUS_ADAPTER_CONFIGMAP_NAME = "adapter-config"
PROMETHEUS_ADAPTER_CONFIGMAP_FILENAME = "config.yaml"
PROMETHEUS_ADAPTER_POD_NAME_PREFIX = "custom-metrics-apiserver"
PROMETHEUS_ADAPTER_POD_PHASES_TO_REMOVE = (
    "Running",
    "Pending",
)

DEFAULT_SCRAPE_PERIOD_S = 10
DEFAULT_EXTRAPOLATION_PERIODS = 10
DEFAULT_EXTRAPOLATION_TIME = DEFAULT_SCRAPE_PERIOD_S * DEFAULT_EXTRAPOLATION_PERIODS

class PrometheusAdapterRule(TypedDict):
    """
    Typed version of the (minimal) set of Prometheus adapter rule configuration options that we use
    """

    # see https://github.com/DirectXMan12/k8s-prometheus-adapter/blob/master/docs/config.md
    # for more detailed information
    # used for discovering what resources should be scaled
    seriesQuery: str
    # configuration for how to expose this rule to the HPA
    name: Dict[str, str]
    # used to associate metrics with k8s resources
    resources: Dict[str, str]
    # the actual query we want to send to Prometheus to use for scaling
    metricsQuery: str

class PrometheusAdapterConfig(TypedDict):
    """
    Typed version of the Prometheus adapter configuration dictionary.
    """

    rules: List[PrometheusAdapterRule]

def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Syncs the Prometheus metric adapter config with soaconfigs.",
    )

    parser.add_argument(
        "-d",
        "--soa-dir",
        dest="soa_dir",
        metavar="SOA_DIR",
        default=Path(DEFAULT_SOA_DIR),
        help="Directory to read service configs from. Default is %(default)s.",
        type=Path,
    )
    # TODO: do we need to be able to pass multiple clusters in?
    parser.add_argument(
        "-c",
        "--cluster",
        dest="cluster",
        help="PaaSTA cluster to generate configs for.",
        required=True,
    )
    parser.add_argument(
        "-v",
        "--verbose",
        action="store_true",
        dest="verbose",
        default=False,
        help="Enable verbose logging.",
    )
    parser.add_argument(
        "--dry-run",
        dest="dry_run",
        action="store_true",
        default=False,
        help="Enable verbose logging.",
    )

    return parser.parse_args()

def should_create_uwsgi_scaling_rule(
    instance: str, autoscaling_config: AutoscalingParamsDict,
) -> Tuple[bool, Optional[str]]:
    """
    Determines whether we should configure the prometheus adapter for a given service.
    Returns a 2-tuple of (should_create, reason_to_skip)
    """
    if autoscaling_config["metrics_provider"] == "uwsgi":
        if not autoscaling_config.get("use_prometheus", False):
            return False, "requested uwsgi autoscaling, but not using Prometheus"

        return True, None

    return False, "did not request uwsgi autoscaling"

def create_instance_uwsgi_scaling_rule(
    service: str,
    instance: str,
    autoscaling_config: AutoscalingParamsDict,
    paasta_cluster: str,
) -> PrometheusAdapterRule:
    """
    Creates a Prometheus adapter rule config for a given service instance.
    """
    setpoint = autoscaling_config["setpoint"]
    moving_average_window = autoscaling_config["moving_average_window_seconds"]
    # this should always be set, but we default to 0 for safety as the worst thing that would happen
    # is that we take a couple more iterations than required to hit the desired setpoint
    offset = autoscaling_config.get("offset", 0)
    deployment_name = get_kubernetes_app_name(service=service, instance=instance)
    worker_filter_terms = f"paasta_cluster='{paasta_cluster}',paasta_service='{service}',paasta_instance='{instance}'"
    replica_filter_terms = (
        f"paasta_cluster='{paasta_cluster}',deployment='{deployment_name}'"
    )

    current_replicas = f"""
        (
            scalar(
                kube_deployment_spec_replicas{{{replica_filter_terms}}} >= 0
                or
                max_over_time(
                    kube_deployment_spec_replicas{{{replica_filter_terms}}}[{DEFAULT_EXTRAPOLATION_TIME}s]
                )
            )
        )
    """
    load_per_instance = f"""
        avg(
            uwsgi_worker_busy{{{worker_filter_terms}}}
        ) by (kube_pod, kube_deployment)
    """
    total_load = f"""
        sum(
            {load_per_instance}
        ) by (kube_deployment)
    """
    desired_instances_at_each_point_in_time = f"""
        {total_load} / {setpoint - offset}
    """
    desired_instances = f"""
        avg_over_time(
            (
                {desired_instances_at_each_point_in_time}
            )[{moving_average_window}s:]
        )
    """
    metrics_query = f"""
        {desired_instances} / {current_replicas}
    """

    metric_name = f"{deployment_name}-uwsgi-prom"

    return {
        "name": {"as": metric_name},
        "seriesQuery": f"uwsgi_worker_busy{{{worker_filter_terms}}}",
        "resources": {"template": "kube_<<.Resource>>"},
        # our nicely formatted query has enough whitespace that collapsing said
        # whitespace cuts down the size of the string by a meaningful amount
        "metricsQuery": re.sub(pattern=r"\s+", repl=" ", string=metrics_query).strip(),
    }

def get_rules_for_service_instance(
    service_name: str,
    instance_name: str,
    autoscaling_config: AutoscalingParamsDict,
    paasta_cluster: str,
) -> List[PrometheusAdapterRule]:
    """
    Returns a list of Prometheus Adapter rules for a given service instance. For now, this
    will always be a 0 or 1-element list - but when we support scaling on multiple metrics
    we will return N rules for a given service instance.
    """
    rules: List[PrometheusAdapterRule] = []

    should_create_uwsgi, skip_uwsgi_reason = should_create_uwsgi_scaling_rule(
        instance=instance_name, autoscaling_config=autoscaling_config,
    )
    if should_create_uwsgi:
        rules.append(
            create_instance_uwsgi_scaling_rule(
                service=service_name,
                instance=instance_name,
                autoscaling_config=autoscaling_config,
                paasta_cluster=paasta_cluster,
            )
        )
    else:
        log.debug(
            "Skipping %s.%s - %s.", service_name, instance_name, skip_uwsgi_reason,
        )

    return rules

def create_prometheus_adapter_config(
    paasta_cluster: str, soa_dir: Path
) -> PrometheusAdapterConfig:
    """
    Given a paasta cluster and a soaconfigs directory, create the necessary Prometheus adapter
    config to autoscale services.
    Currently supports the following metrics providers:
        * uwsgi
    """
    rules: List[PrometheusAdapterRule] = []
    # get_services_for_cluster() returns a list of (service, instance) tuples, but this
    # is not great for us: if we were to iterate over that we'd end up getting duplicates
    # for every service as PaastaServiceConfigLoader does not expose a way to get configs
    # for a single instance by name. instead, we get the unique set of service names and then
    # let PaastaServiceConfigLoader iterate over instances for us later
    services = {
        service_name
        for service_name, _ in get_services_for_cluster(
            cluster=paasta_cluster, instance_type="kubernetes", soa_dir=str(soa_dir)
        )
    }
    for service_name in services:
        config_loader = PaastaServiceConfigLoader(
            service=service_name, soa_dir=str(soa_dir)
        )
        for instance_config in config_loader.instance_configs(
            cluster=paasta_cluster, instance_type_class=KubernetesDeploymentConfig,
        ):
            rules.extend(
                get_rules_for_service_instance(
                    service_name=service_name,
                    instance_name=instance_config.instance,
                    autoscaling_config=instance_config.get_autoscaling_params(),
                    paasta_cluster=paasta_cluster,
                )
            )

    return {
        # we sort our rules so that we can easily compare between two different configmaps
        # as otherwise we'd need to do fancy order-independent comparisons between the two
        # sets of rules later due to the fact that we're not iterating in a deterministic
        # way and can add rules in any arbitrary order
        "rules": sorted(rules, key=lambda rule: rule["name"]["as"]),
    }

def update_prometheus_adapter_configmap(
    kube_client: KubeClient, config: PrometheusAdapterConfig
) -> None:
    kube_client.core.replace_namespaced_config_map(
        name=PROMETHEUS_ADAPTER_CONFIGMAP_NAME,
        namespace=PROMETHEUS_ADAPTER_CONFIGMAP_NAMESPACE,
        body=V1ConfigMap(
            metadata=V1ObjectMeta(name=PROMETHEUS_ADAPTER_CONFIGMAP_NAME),
            data={
                PROMETHEUS_ADAPTER_CONFIGMAP_FILENAME: yaml.dump(
                    config,
                    default_flow_style=False,
                    explicit_start=True,
                    width=sys.maxsize,
                )
            },
        ),
    )

def create_prometheus_adapter_configmap(
    kube_client: KubeClient, config: PrometheusAdapterConfig
) -> None:
    kube_client.core.create_namespaced_config_map(
        namespace=PROMETHEUS_ADAPTER_CONFIGMAP_NAMESPACE,
        body=V1ConfigMap(
            metadata=V1ObjectMeta(name=PROMETHEUS_ADAPTER_CONFIGMAP_NAME),
            data={
                PROMETHEUS_ADAPTER_CONFIGMAP_FILENAME: yaml.dump(
                    config, default_flow_style=False, explicit_start=True
                )
            },
        ),
    )

def get_prometheus_adapter_configmap(
    kube_client: KubeClient,
) -> Optional[PrometheusAdapterConfig]:
    try:
        config = cast(
            # we cast since mypy infers the wrong type since the k8s clientlib is untyped
            V1ConfigMap,
            kube_client.core.read_namespaced_config_map(
                name=PROMETHEUS_ADAPTER_CONFIGMAP_NAME,
                namespace=PROMETHEUS_ADAPTER_CONFIGMAP_NAMESPACE,
            ),
        )
    except ApiException as e:
        if e.status == 404:
            return None
        else:
            raise

    if not config:
        return None

    return yaml.safe_load(config.data[PROMETHEUS_ADAPTER_CONFIGMAP_FILENAME])

def restart_prometheus_adapter(kube_client: KubeClient) -> None:
    log.info("Attempting to remove existing adapter pod(s).")
    all_pods = cast(
        # once again, we cast since the kubernetes python api isn't typed
        List[V1Pod],
        kube_client.core.list_namespaced_pod(
            namespace=PROMETHEUS_ADAPTER_POD_NAMESPACE
        ).items,
    )
    # there should only ever be one pod actually up, but we might as well enforce that here
    # just in case there are more
    pods_to_delete = [
        pod
        for pod in all_pods
        if pod.metadata.name.startswith(PROMETHEUS_ADAPTER_POD_NAME_PREFIX)
        and pod.status.phase in PROMETHEUS_ADAPTER_POD_PHASES_TO_REMOVE
    ]
    log.debug("Found the following pods to delete: %s", pods_to_delete)

    for pod in pods_to_delete:
        log.debug("Attempting to remove %s.", pod.metadata.name)
        kube_client.core.delete_namespaced_pod(
            name=pod.metadata.name,
            namespace=pod.metadata.namespace,
            body=V1DeleteOptions(),
            # background propagation with no grace period is equivalent to doing a force-delete from kubectl
            grace_period_seconds=0,
            propagation_policy="Background",
        )
        log.debug("Removed %s.", pod.metadata.name)

    log.info("Adapter restarted successfully")

def main() -> int:
    args = parse_args()
    if args.verbose:
        logging.basicConfig(level=logging.DEBUG)
    else:
        logging.basicConfig(level=logging.INFO)

    log.info("Generating adapter config from soaconfigs.")
    config = create_prometheus_adapter_config(
        paasta_cluster=args.cluster, soa_dir=args.soa_dir
    )
    log.info("Generated adapter config from soaconfigs.")
    if args.dry_run:
        log.info(
            "Generated the following config:\n%s",
            yaml.dump(
                config, default_flow_style=False, explicit_start=True, width=sys.maxsize
            ),
        )
        return 0  # everything after this point requires creds/updates state
    else:
        log.debug(
            "Generated the following config:\n%s",
            yaml.dump(
                config, default_flow_style=False, explicit_start=True, width=sys.maxsize
            ),
        )

    if not config["rules"]:
        log.error("Got empty rule configuration - refusing to continue.")
        return 0

    kube_client = KubeClient()
    if not args.dry_run:
        ensure_namespace(kube_client, namespace="paasta")
        ensure_namespace(kube_client, namespace="custom-metrics")

    existing_config = get_prometheus_adapter_configmap(kube_client=kube_client)
    if existing_config and existing_config != config:
        log.info("Existing config differs from soaconfigs - updating.")
        log.debug("Existing data: %s", existing_config)
        log.debug("Desired data: %s", config)
        update_prometheus_adapter_configmap(kube_client=kube_client, config=config)
        log.info("Updated adapter config.")
    elif existing_config:
        log.info("Existing config matches soaconfigs - exiting.")
        return 0
    else:
        log.info("No existing config - creating.")
        create_prometheus_adapter_configmap(kube_client=kube_client, config=config)
        log.info("Created adapter config.")

    # the prometheus adapter doesn't currently have a good way to reload on config changes
    # so we do the next best thing: restart the pod so that it picks up the new config.
    # see: https://github.com/DirectXMan12/k8s-prometheus-adapter/issues/104
    restart_prometheus_adapter(kube_client=kube_client)

    return 0

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


<commit_before>"""
Automate the GALFIT analysis procedure. GALFIT
is a free, closed-source software that can be found
here:
https://users.obs.carnegiescience.edu/peng/work/galfit/galfit.html

Here is the user manual:
https://users.obs.carnegiescience.edu/peng/work/galfit/README.pdf

This module is mainly for a python wrapper to automate
using GALFIT. It also contains additional functions
to parse the output automatically.
"""

import numpy as np
import pandas as pd
import sys, os, subprocess
import warnings
from astropy.io import fits
from astropy.stats import sigma_clipped_stats
from astropy.table import Table
from astropy.wcs import WCS

def genconf(imgfile:str, psffile:str,
            configfile:str=None, outfile:str=None,
            finesample:int = 1, badpix:str = "none",
            constraints:str = "none",
            region:tuple = None, convobox:tuple = (100,100),
            zeropoint:float = 25.0, platescale:float = 0.125,
            position:tuple = None, int_mag:float = None,
            r_e:float = None, n:float = 1.0, axis_ratio:float = 0.5,
            pa:float = 0, skip_sky:bool = False):
    """
    Creates a configuration file for GALFIT. Conventionally,
    GALFIT is run using the command: `galfit <config-file>`.
    Args
    ----
    imgfile (str): path to the image fits file.
    psffile (str): path to the PSF model fits file.
    configfile (str, optional): path to configuration file to
        be created via this function.
    outfile (str, optional): path to GALFIT's output fits file.
        Defaults to `out.fits`
    finesample (int, optional): The PSF fine-sampling factor.
        Assumes no fine-sampling (i.e. a factor of 1) by default.
    badpix (str, optional): File containing a list of bad pixels.
        Assumes all pixels a re fine by default.
    constraints (str, optional): File containing fit parameter
        bounds. Check out this example constraints file
        https://users.obs.carnegiescience.edu/peng/work/galfit/EXAMPLE.CONSTRAINTS
        to learn how to use one.
    region (tuple, optional): Pixel coordinate bounds for the fitting.
        Required format: (xmin, xmax, ymin, ymax). All of the
        elements must be ints!
    convobox (tuple, optional): Size of the convolution box
        used to assess the model fit chi-squared value.
        Required format: (x_box:int , y_box:int)
    zeropoint (float, optional): Zeropoint of the image
        to compute the model magnitude.
    platescale (float, optional): Size of pixel in arcsec.
        Assumes 0.125''/pixel by default.
    position (tuple, optional): Guess for the centroid
        of the model fit. Required format (x_cen:float, y_cen:float).
        Assumes the center of the image is the initial guess by
        default.
    int_mag (float, optional): Initial guess for the
        integrated magnitude. Taken as the magnitude
        corresponding to the sum of all counts in the
        region to be fit.
    r_e (float, optional): Initial guess for
        the half light radius in pixels. Assumes
        half the size of the fitting region by default.
    n (float, optional): initial guess for the Sersic
        index.
    axis_ratio (float, optional): Initial guess for
        the ratio of minor to major axis of the fit model.
    PA (float, optional): Initial guess for the
        angle of the major axis counter clockwise
        relative to the vertical.
    skip_sky (bool, optional): Do you also want
        to fit a constant sky background? Set to
        false if your sky background is 0.
    Returns
    -------
    configfile (str): Path to the configuration
        file.
    """
    if configfile is None:
        warnings.warn("Creating a configuration file here")
        configfile = "galfit.feedme"
    with open(configfile,"w+") as fstream:
        #Image parameters.
        fstream.write("""===============================================================================\n
        # IMAGE and GALFIT CONTROL PARAMETERS\n
        """)
        assert os.path.isfile(imgfile), "Invalid image file path {:s}".format(imgfile)
        fstream.write("A) {:s}  # Input data image (FITS file)\n".format(imgfile))
        if outfile is None:
            warnings.warn("Creating output file here")
            outfile = "out.fits"
        fstream.write("B) {:s}  # Output data image block\n".format(outfile))
        fstream.write("C) none  # Sigma image name (made from data if blank or 'none')\n")
        assert os.path.isfile(psffile), "Invalid psf file path {:s}".format(psffile)
        fstream.write("D) {:s}  # Input PSF file\n".format(psffile))
        fstream.write("E) {:d}  # PSF fine sampling factor\n".format(finesample))
        fstream.write("F) {:s}  #Bad pixel mask\n".format(badpix))
        fstream.write("G) {:s}  # File with parameter constraints (ASCII file)\n".format(constraints))
        if region is None:
            input_str = input("Input image region to fit (xmin xmax ymin ymax): ").split()
            while len(input_str)!=4:
                input_str = input("Invalid region input. Insert again (xmin xmax ymin ymax): ").split()
            region = [int(bound) for bound in input_str]
        fstream.write("H) {:d} {:d} {:d} {:d}   # Image region to fit (xmin xmax ymin ymax)\n".format(region[0],region[1],region[2], region[3]))
        assert len(convobox)==2, "Only two integers specify convolution box dimensions"
        fstream.write("I) {:d} {:d} # Size of convolution box (x y)\n".format(convobox[0],convobox[1]))
        fstream.write("J) {:f}  # Photometric zeropoint (mag)\n".format(zeropoint))
        fstream.write("K) {:f} {:f} # Plate scale (dx dy) [arcsec/pixel]\n".format(platescale,platescale))
        fstream.write("O) regular   # Display type (regular, curses, both)\n")
        fstream.write("P) 0 # Choose: 0=optimize, 1=model, 2=imgblock, 3=subcomps\n")
        fstream.write("\n# INITIAL FITTING PARAMETERS\n")
        fstream.write("# Component number: 1\n")
        fstream.write("0) sersic # Component type\n")
        img = fits.getdata(imgfile)
        if position is None:
            pos_x, pos_y = img.shape/2
        else:
            pos_x, pos_y = position
        fstream.write("1) {:d} {:d} 1 1 # position x y\n".format(pos_x,pos_y))
        if int_mag is None:
            int_mag = zeropoint-2.5*np.log10(np.sum(img[region[2]:region[3],region[0]:region[1]]))
        fstream.write("3) {:f} 1 # Integrated magnitude\n".format(int_mag))
        if r_e is None:
            r_e = (region[3]-region[2]+region[1]-region[0])/2/3 #A third of the average region dimension
        fstream.write("4) {:f} 1 # effective radius (pix)\n".format(r_e))
        fstream.write("5) {:f} 1 # sersic index\n".format(n))
        for num in [6,7,8]:
            fstream.write("{:d}) 0.0000 0 # ----\n".format(num))
        fstream.write("9) {:f} 1 # Axis ratio (b/a)\n".format(axis_ratio))
        fstream.write("10) {:f} 1 # Position angle (PA) [deg: Up=0, left =90]\n".format(pa))
        fstream.write("Z) 0 # Skip this model? (yes=1,no=0)\n\n")
        if not skip_sky:
            fstream.write("# Component number: 2\n")
            fstream.write("0) sky # component type\n")
            _, median, _ = sigma_clipped_stats(img)
            fstream.write("1) {:f} 1 # Sky background\n".format(median))
            fstream.write("2) 0 0 # dsky/dx\n")
            fstream.write("3) 0 0 # dsky/dy\n")
            fstream.write("Z) 0 # Skip this model\n")
        fstream.write("================================================================================\n")
    return configfile

def run(imgfile, psffile,**kwargs):
    """
    Run galfit. The arguments are
    the same as genconf.
    """
    configfile = genconf(imgfile,psffile,**kwargs)
    os.system("galfit {:s}".format(configfile))
    return 0

def pix2coord(sersic_tab, wcs, platescale):
    """
    Takes the output table from galfit's
    fit.log file and converts all pixel
    measurements to physical measurements.
    Args
    ----
    sersic_tab (Table): Raw sersic
        fit parameter table
    wcs (WCS): transformation to
        go from pixels to angular units
    Returns
    -------
    sky_tab (Table): Sersic table
        translated to angular units
        on the sky.
    """

    sky_tab = Table()

    # First read the invariants
    sky_tab['axis_ratio'] = sersic_tab['axis_ratio']
    sky_tab['axis_ratio_err'] = sersic_tab['axis_ratio_err']
    sky_tab['n'] = sersic_tab['n']
    sky_tab['n_err'] = sersic_tab['n_err']
    sky_tab['PA'] = sersic_tab['PA']
    sky_tab['PA_err'] = sersic_tab['PA_err']

    # Read in centroids next
    xpix, ypix = sersic_tab['x']-1, sersic_tab['y']-1
    centr_coords = wcs.pixel_to_world(xpix, ypix)
    sky_tab['ra'] = centr_coords.ra.value
    sky_tab['dec'] = centr_coords.dec.value

    sky_tab['ra_err'] = sersic_tab['x_err']*platescale
    sky_tab['dec_err'] = sersic_tab['y_err']*platescale

    # Read in r_e next.
    sky_tab['r_e'] = sersic_tab['r_e']*platescale
    sky_tab['r_e_err'] = sersic_tab['r_e_err']*platescale

    # Reorder columns
    sky_tab = sky_tab['ra','ra_err', 'dec','dec_err', 'r_e', 'r_e_err', 'n', 'n_err','axis_ratio','axis_ratio_err','PA','PA_err']

    return sky_tab

def surf_brightness(coord, sersic_tab):
    """
    Estimates the surface brightness from
    the sersic model fit at the input
    coordinate location.
    Args
    ----
    coord (SkyCoord): Target coordinate
    sersic_tab (Table): Table of best fit
        sersic model parameters.
    wcs ()
    """<commit_msg>function to read fit.log<commit_after>"""
Automate the GALFIT analysis procedure. GALFIT
is a free, closed-source software that can be found
here:
https://users.obs.carnegiescience.edu/peng/work/galfit/galfit.html

Here is the user manual:
https://users.obs.carnegiescience.edu/peng/work/galfit/README.pdf

This module is mainly for a python wrapper to automate
using GALFIT. It also contains additional functions
to parse the output automatically.
"""

import numpy as np
import pandas as pd
import sys, os, subprocess
import warnings
from astropy.io import fits
from astropy.stats import sigma_clipped_stats
from astropy.table import Table
from astropy.wcs import WCS

def genconf(imgfile:str, psffile:str,
            configfile:str=None, outfile:str=None,
            finesample:int = 1, badpix:str = "none",
            constraints:str = "none",
            region:tuple = None, convobox:tuple = (100,100),
            zeropoint:float = 25.0, platescale:float = 0.125,
            position:tuple = None, int_mag:float = None,
            r_e:float = None, n:float = 1.0, axis_ratio:float = 0.5,
            pa:float = 0, skip_sky:bool = False):
    """
    Creates a configuration file for GALFIT. Conventionally,
    GALFIT is run using the command: `galfit <config-file>`.
    Args
    ----
    imgfile (str): path to the image fits file.
    psffile (str): path to the PSF model fits file.
    configfile (str, optional): path to configuration file to
        be created via this function.
    outfile (str, optional): path to GALFIT's output fits file.
        Defaults to `out.fits`
    finesample (int, optional): The PSF fine-sampling factor.
        Assumes no fine-sampling (i.e. a factor of 1) by default.
    badpix (str, optional): File containing a list of bad pixels.
        Assumes all pixels a re fine by default.
    constraints (str, optional): File containing fit parameter
        bounds. Check out this example constraints file
        https://users.obs.carnegiescience.edu/peng/work/galfit/EXAMPLE.CONSTRAINTS
        to learn how to use one.
    region (tuple, optional): Pixel coordinate bounds for the fitting.
        Required format: (xmin, xmax, ymin, ymax). All of the
        elements must be ints!
    convobox (tuple, optional): Size of the convolution box
        used to assess the model fit chi-squared value.
        Required format: (x_box:int , y_box:int)
    zeropoint (float, optional): Zeropoint of the image
        to compute the model magnitude.
    platescale (float, optional): Size of pixel in arcsec.
        Assumes 0.125''/pixel by default.
    position (tuple, optional): Guess for the centroid
        of the model fit. Required format (x_cen:float, y_cen:float).
        Assumes the center of the image is the initial guess by
        default.
    int_mag (float, optional): Initial guess for the
        integrated magnitude. Taken as the magnitude
        corresponding to the sum of all counts in the
        region to be fit.
    r_e (float, optional): Initial guess for
        the half light radius in pixels. Assumes
        half the size of the fitting region by default.
    n (float, optional): initial guess for the Sersic
        index.
    axis_ratio (float, optional): Initial guess for
        the ratio of minor to major axis of the fit model.
    PA (float, optional): Initial guess for the
        angle of the major axis counter clockwise
        relative to the vertical.
    skip_sky (bool, optional): Do you also want
        to fit a constant sky background? Set to
        false if your sky background is 0.
    Returns
    -------
    configfile (str): Path to the configuration
        file.
    """
    if configfile is None:
        warnings.warn("Creating a configuration file here")
        configfile = "galfit.feedme"
    with open(configfile,"w+") as fstream:
        #Image parameters.
        fstream.write("""===============================================================================\n
        # IMAGE and GALFIT CONTROL PARAMETERS\n
        """)
        assert os.path.isfile(imgfile), "Invalid image file path {:s}".format(imgfile)
        fstream.write("A) {:s}  # Input data image (FITS file)\n".format(imgfile))
        if outfile is None:
            warnings.warn("Creating output file here")
            outfile = "out.fits"
        fstream.write("B) {:s}  # Output data image block\n".format(outfile))
        fstream.write("C) none  # Sigma image name (made from data if blank or 'none')\n")
        assert os.path.isfile(psffile), "Invalid psf file path {:s}".format(psffile)
        fstream.write("D) {:s}  # Input PSF file\n".format(psffile))
        fstream.write("E) {:d}  # PSF fine sampling factor\n".format(finesample))
        fstream.write("F) {:s}  #Bad pixel mask\n".format(badpix))
        fstream.write("G) {:s}  # File with parameter constraints (ASCII file)\n".format(constraints))
        if region is None:
            input_str = input("Input image region to fit (xmin xmax ymin ymax): ").split()
            while len(input_str)!=4:
                input_str = input("Invalid region input. Insert again (xmin xmax ymin ymax): ").split()
            region = [int(bound) for bound in input_str]
        fstream.write("H) {:d} {:d} {:d} {:d}   # Image region to fit (xmin xmax ymin ymax)\n".format(region[0],region[1],region[2], region[3]))
        assert len(convobox)==2, "Only two integers specify convolution box dimensions"
        fstream.write("I) {:d} {:d} # Size of convolution box (x y)\n".format(convobox[0],convobox[1]))
        fstream.write("J) {:f}  # Photometric zeropoint (mag)\n".format(zeropoint))
        fstream.write("K) {:f} {:f} # Plate scale (dx dy) [arcsec/pixel]\n".format(platescale,platescale))
        fstream.write("O) regular   # Display type (regular, curses, both)\n")
        fstream.write("P) 0 # Choose: 0=optimize, 1=model, 2=imgblock, 3=subcomps\n")
        fstream.write("\n# INITIAL FITTING PARAMETERS\n")
        fstream.write("# Component number: 1\n")
        fstream.write("0) sersic # Component type\n")
        img = fits.getdata(imgfile)
        if position is None:
            pos_x, pos_y = img.shape/2
        else:
            pos_x, pos_y = position
        fstream.write("1) {:d} {:d} 1 1 # position x y\n".format(pos_x,pos_y))
        if int_mag is None:
            int_mag = zeropoint-2.5*np.log10(np.sum(img[region[2]:region[3],region[0]:region[1]]))
        fstream.write("3) {:f} 1 # Integrated magnitude\n".format(int_mag))
        if r_e is None:
            r_e = (region[3]-region[2]+region[1]-region[0])/2/3 #A third of the average region dimension
        fstream.write("4) {:f} 1 # effective radius (pix)\n".format(r_e))
        fstream.write("5) {:f} 1 # sersic index\n".format(n))
        for num in [6,7,8]:
            fstream.write("{:d}) 0.0000 0 # ----\n".format(num))
        fstream.write("9) {:f} 1 # Axis ratio (b/a)\n".format(axis_ratio))
        fstream.write("10) {:f} 1 # Position angle (PA) [deg: Up=0, left =90]\n".format(pa))
        fstream.write("Z) 0 # Skip this model? (yes=1,no=0)\n\n")
        if not skip_sky:
            fstream.write("# Component number: 2\n")
            fstream.write("0) sky # component type\n")
            _, median, _ = sigma_clipped_stats(img)
            fstream.write("1) {:f} 1 # Sky background\n".format(median))
            fstream.write("2) 0 0 # dsky/dx\n")
            fstream.write("3) 0 0 # dsky/dy\n")
            fstream.write("Z) 0 # Skip this model\n")
        fstream.write("================================================================================\n")
    return configfile

def run(imgfile, psffile,**kwargs):
    """
    Run galfit. The arguments are
    the same as genconf.
    """
    configfile = genconf(imgfile,psffile,**kwargs)
    os.system("galfit {:s}".format(configfile))
    return 0

def pix2coord(sersic_tab, wcs, platescale):
    """
    Takes the output table from galfit's
    fit.log file and converts all pixel
    measurements to physical measurements.
    Args
    ----
    sersic_tab (Table): Raw sersic
        fit parameter table
    wcs (WCS): transformation to
        go from pixels to angular units
    Returns
    -------
    sky_tab (Table): Sersic table
        translated to angular units
        on the sky.
    """

    sky_tab = Table()

    # First read the invariants
    sky_tab['axis_ratio'] = sersic_tab['axis_ratio']
    sky_tab['axis_ratio_err'] = sersic_tab['axis_ratio_err']
    sky_tab['n'] = sersic_tab['n']
    sky_tab['n_err'] = sersic_tab['n_err']
    sky_tab['PA'] = sersic_tab['PA']
    sky_tab['PA_err'] = sersic_tab['PA_err']

    # Read in centroids next
    xpix, ypix = sersic_tab['x']-1, sersic_tab['y']-1
    centr_coords = wcs.pixel_to_world(xpix, ypix)
    sky_tab['ra'] = centr_coords.ra.value
    sky_tab['dec'] = centr_coords.dec.value

    sky_tab['ra_err'] = sersic_tab['x_err']*platescale
    sky_tab['dec_err'] = sersic_tab['y_err']*platescale

    # Read in r_e next.
    sky_tab['r_e'] = sersic_tab['r_e']*platescale
    sky_tab['r_e_err'] = sersic_tab['r_e_err']*platescale

    # Reorder columns
    sky_tab = sky_tab['ra','ra_err', 'dec','dec_err', 'r_e', 'r_e_err', 'n', 'n_err','axis_ratio','axis_ratio_err','PA','PA_err']

    return sky_tab

def read_fitlog(outfile, plate_scale, initfile):
    lines = [line.rstrip('\n') for line in open(outfile)]

    instance = []  # going to put all instances of use of input file here
    for kk, line in enumerate(lines):  # for all lines in fit.log
        if line == 'Init. par. file : ' + str(initfile) + ' ':  # if the it is from the input file used
            #print('yes/')
            for pp, item in enumerate(lines[kk:kk + 10]):
                if item[1:7] == 'sersic':  # look for the sersic fit model
                    #print(item, pp)
                    instance.append((item, pp, lines[kk:kk + 10][pp + 1]))  # and keep the model and error

    fit_out, where, err_out = instance[-1]  # only keep the latest one

    galfit = {} #store values in dict
    #Store values
    prss = fit_out.split(' ')
    keepp = [obj for obj in prss if obj != '']  # Remove white spaces
    galfit['PA'] = float(keepp[-1])
    galfit['b/a'] = float(keepp[-2])
    galfit['n'] = float(keepp[-3])
    galfit['reff_ang'] = float(keepp[-4]) * plate_scale

    #Store errors
    prss = err_out.split(' ')
    keepp = [obj for obj in prss if obj != '']  # Remove white spaces
    galfit['PA_err'] = float(keepp[-1])
    galfit['b/a_err'] = float(keepp[-2])
    galfit['n_err'] = float(keepp[-3])
    galfit['reff_ang_err'] = float(keepp[-4]) * plate_scale

    return galfit

def surf_brightness(coord, sersic_tab):
    """
    Estimates the surface brightness from
    the sersic model fit at the input
    coordinate location.
    Args
    ----
    coord (SkyCoord): Target coordinate
    sersic_tab (Table): Table of best fit
        sersic model parameters.
    wcs ()
    """<|endoftext|>