gene.py

from django_pandas.io import read_frame
from drf_yasg import openapi
from drf_yasg.utils import swagger_auto_schema
from marshmallow.exceptions import ValidationError
from rest_framework import filters
from rest_framework.decorators import action
from rest_framework.response import Response
from rest_framework.status import HTTP_204_NO_CONTENT, HTTP_422_UNPROCESSABLE_ENTITY

from metagenedb.apps.catalog.models import Gene
from metagenedb.api.catalog.qparams_validators.gene import GeneLengthQueryParams
from metagenedb.apps.catalog.serializers import GeneSerializer
from metagenedb.common.utils.df_operations import get_mask

from .bulk_viewset import BulkViewSet


class DocGeneLength(object):
    """
    Define response for API documentation of gene length distribution method
    {
    "results": {
        "counts": [
            0,
            887,
        ],
        "labels": [
            "0-9999",
            "10000-19999",
            "20000-29999",
        ]
    }
}
    """
    window_size_param = openapi.Parameter('window_size', in_=openapi.IN_QUERY, description='Size of the window.',
                                          type=openapi.TYPE_INTEGER, default=10000)
    counts = openapi.Schema(type="array", items=openapi.Schema(type="int"),
                            description="Counts for every window_size")
    labels = openapi.Schema(type="array", items=openapi.Schema(type="char"),
                            description="Corresponding windows")
    results = openapi.Schema(type="object", properties={'counts': counts, 'labels': labels},
                             description="results of your request")
    gene_length_schema = openapi.Schema(type="object", properties={'results': results})
    gene_length_response = openapi.Response('Get the distribution of gene length for a given window size',
                                            schema=gene_length_schema)


class GeneViewSet(BulkViewSet):
    search_fields = ['gene_name']
    filter_backends = (filters.SearchFilter,)
    queryset = Gene.objects.select_related('taxonomy').prefetch_related('functions').all()
    serializer_class = GeneSerializer
    lookup_field = 'gene_id'
    GENE_LENGTH_COL = 'length'
    DEFAULT_WINDOW_SIZE = 1000
    DEFAULT_STOP_AT = 10000

    def get_permissions(self):
        return super(self.__class__, self).get_permissions()

    def _count_windows(self, df, window_size=DEFAULT_WINDOW_SIZE, window_col=GENE_LENGTH_COL, stop_at=DEFAULT_STOP_AT):
        """
        Count how many line of the df belong to each windows defined by the window_size for the window_col
        :param df:
        :param window_col: column concerned by the window
        :param window_size: size of the window
        :return: {'data': COUNTS_BY_WINDOW, 'labels': START-END}
        """
        length_max = df[window_col].max()
        stop_at = length_max if length_max < stop_at else stop_at
        all_ranges = [[i, i + window_size] for i in range(0, stop_at + 1, window_size)]
        all_ranges[-1][1] = length_max + 1  # last should contain all above the stop_at
        data = []
        labels = []
        for rg in all_ranges:
            labels.append(f"{rg[0]/1000}k-{rg[1]/1000}k")
            data.append(df[get_mask(df, rg, window_col)].count()[window_col])
        # Change labels
        labels[0] = f"<{labels[0].split('-')[1]}"
        labels[-1] = f">{labels[-1].split('-')[0]}"
        return {
            'counts': data,
            'labels': labels
        }

    @swagger_auto_schema(
        manual_parameters=[DocGeneLength.window_size_param],
        responses={
            '200': DocGeneLength.gene_length_response,
            '204': 'No genes on the catalog to build the distribution'
        },
        operation_id='Gene length distribution',
    )
    @action(methods=['get'], detail=False)
    def gene_length(self, request):
        try:
            query_params = GeneLengthQueryParams().load(request.query_params)
        except ValidationError as validation_error:
            error_message = validation_error.normalized_messages()
            error_message.update({
                'allowed_query_params': ', '.join(GeneLengthQueryParams().declared_fields.keys())
            })
            return Response(error_message, status=HTTP_422_UNPROCESSABLE_ENTITY)

        window_size = query_params.get('window_size', self.DEFAULT_WINDOW_SIZE)
        stop_at = query_params.get('stop_at', self.DEFAULT_STOP_AT)
        df = read_frame(Gene.objects.all(), fieldnames=[self.GENE_LENGTH_COL])
        if df.empty:
            return Response(
                {'results': {}},
                status=HTTP_204_NO_CONTENT
            )
        return Response(
            {'results': self._count_windows(df, window_size=window_size, stop_at=stop_at)}
        )