<commit_before>from rest_framework.authentication import SessionAuthentication, BasicAuthentication
from rest_framework_jwt.authentication import JSONWebTokenAuthentication
from rest_framework.permissions import IsAuthenticated
from django.contrib.auth.models import User
from rest_framework import viewsets
from .serializers import UserSerializer, TagSerializer, NotesSerializer, DiarySerializer, SecretSerializer
from .models import Tag, Diary, Notes, Secret
from rest_framework.decorators import list_route
from rest_framework.response import Response
from rest_framework import status

class UserViewSet(viewsets.ModelViewSet):
    """
    API endpoint that allows users to be viewed or edited.
    """
    queryset = User.objects.all().order_by('-date_joined')
    serializer_class = UserSerializer

class TagViewSet(viewsets.ModelViewSet):
    """
        API endpoint that allows tags to be created, viewed ,edited and deleted.
    """
    authentication_classes = (SessionAuthentication, BasicAuthentication, JSONWebTokenAuthentication)
    permission_classes = (IsAuthenticated,)
    queryset = Tag.objects.all()
    serializer_class = TagSerializer

    @list_route(methods=['get'])
    def cloud(self, request, *args, **kwargs):
        """
        Get tag cloud items
        """
        cloud = {}
        tags = Tag.objects.all()
        for i in tags:
            hiren = Diary.objects.filter(tag=i).count()
            cloud[i.name] = hiren
        return Response(cloud)

class NotesViewset(viewsets.ModelViewSet):
    """
        API endpoint that allows notes to be created, viewed ,edited and deleted.
    """
    authentication_classes = (SessionAuthentication, BasicAuthentication, JSONWebTokenAuthentication)
    permission_classes = (IsAuthenticated,)
    queryset = Notes.objects.all()
    serializer_class = NotesSerializer

class DiaryViewset(viewsets.ModelViewSet):
    """
        API endpoint that allows diary to be created, viewed ,edited and deleted.
    """
    authentication_classes = (SessionAuthentication, BasicAuthentication, JSONWebTokenAuthentication)
    permission_classes = (IsAuthenticated,)
    queryset = Diary.objects.all()
    serializer_class = DiarySerializer

class SecretViewset(viewsets.ModelViewSet):
    """
        API endpoint that allows secret key  to be created, viewed ,edited and deleted.
    """
    authentication_classes = (SessionAuthentication, BasicAuthentication, JSONWebTokenAuthentication)
    permission_classes = (IsAuthenticated,)
    queryset = Secret.objects.all()
    serializer_class = SecretSerializer

    def create(self, request, *args, **kwargs):
        """
        Check if secret key already created
        """
        count = Secret.objects.all().count()
        if count == 1:
            content = {'error': 'key already exits'}
            return Response(content, status.HTTP_403_FORBIDDEN)
        else:
            instance = self.get_object()
            Secret.objects.create(key=instance)
            serializer = self.get_serializer(instance)
            return Response(serializer.data, status.HTTP_201_CREATED)

    def destroy(self, request, pk=None):
        bunny = {'error': 'method not supported :/'}
        return Response(bunny, status.HTTP_403_FORBIDDEN)

<commit_msg>changed doc<commit_after>from rest_framework.authentication import SessionAuthentication, BasicAuthentication
from rest_framework_jwt.authentication import JSONWebTokenAuthentication
from rest_framework.permissions import IsAuthenticated
from django.contrib.auth.models import User
from rest_framework import viewsets
from .serializers import UserSerializer, TagSerializer, NotesSerializer, DiarySerializer, SecretSerializer
from .models import Tag, Diary, Notes, Secret
from rest_framework.decorators import list_route
from rest_framework.response import Response
from rest_framework import status

class UserViewSet(viewsets.ModelViewSet):
    """
    API endpoint that allows users to be viewed or edited.
    """
    queryset = User.objects.all().order_by('-date_joined')
    serializer_class = UserSerializer

class TagViewSet(viewsets.ModelViewSet):
    """
        API endpoint that allows tags to be created, viewed ,edited and deleted.
    """
    authentication_classes = (SessionAuthentication, BasicAuthentication, JSONWebTokenAuthentication)
    permission_classes = (IsAuthenticated,)
    queryset = Tag.objects.all()
    serializer_class = TagSerializer

    @list_route(methods=['get'])
    def cloud(self, request, *args, **kwargs):
        """
        Get tag cloud items
        """
        cloud = {}
        tags = Tag.objects.all()
        for i in tags:
            hiren = Diary.objects.filter(tag=i).count()
            cloud[i.name] = hiren
        return Response(cloud)

class NotesViewset(viewsets.ModelViewSet):
    """
        API endpoint that allows notes to be created, viewed ,edited and deleted.
    """
    authentication_classes = (SessionAuthentication, BasicAuthentication, JSONWebTokenAuthentication)
    permission_classes = (IsAuthenticated,)
    queryset = Notes.objects.all()
    serializer_class = NotesSerializer

class DiaryViewset(viewsets.ModelViewSet):
    """
        API endpoint that allows diary to be created, viewed ,edited and deleted.
    """
    authentication_classes = (SessionAuthentication, BasicAuthentication, JSONWebTokenAuthentication)
    permission_classes = (IsAuthenticated,)
    queryset = Diary.objects.all()
    serializer_class = DiarySerializer

class SecretViewset(viewsets.ModelViewSet):
    """
        API endpoint that allows secret key  to be created, viewed ,edited.
    """
    authentication_classes = (SessionAuthentication, BasicAuthentication, JSONWebTokenAuthentication)
    permission_classes = (IsAuthenticated,)
    queryset = Secret.objects.all()
    serializer_class = SecretSerializer

    def create(self, request, *args, **kwargs):
        """
        Check if secret key already exists
        """
        count = Secret.objects.all().count()
        if count == 1:
            content = {'error': 'key already exits'}
            return Response(content, status.HTTP_403_FORBIDDEN)
        else:
            instance = self.get_object()
            Secret.objects.create(key=instance)
            serializer = self.get_serializer(instance)
            return Response(serializer.data, status.HTTP_201_CREATED)

    def destroy(self, request, pk=None):
        bunny = {'error': 'method not supported :/'}
        return Response(bunny, status.HTTP_403_FORBIDDEN)

<|endoftext|>


<commit_before># Licensed under a 3-clause BSD style license - see LICENSE.rst
# -*- coding: utf-8 -*-
"""Test desidatamodel.stub functions
"""
import os
import unittest
from unittest.mock import patch
from pkg_resources import resource_filename
from astropy.io import fits
from collections import OrderedDict

from .datamodeltestcase import DataModelTestCase
from .. import DataModelError
from ..stub import (Stub, extrakey, file_size, fits_column_format,
                    extract_keywords, image_format, log)

class sim_comments(dict):
    """Simulate a dictionary.
    """
    def __getitem__(self, key):
        return "This is the comment on {0}.".format(key)

class sim_header(OrderedDict):
    """Simulate a FITS header object.
    """
    comments = sim_comments()

class sim_hdu(object):
    """Simulate a FITS HDU.
    """
    def __init__(self, header):
        self.header = header
        return

class TestStub(DataModelTestCase):

    def test_Stub(self):
        """Test aspects of initialization of Stub objects.
        """
        with fits.open(resource_filename('desidatamodel.test',
                                         't/fits_file.fits')) as hdulist:
            stub = Stub(hdulist)
            self.assertEqual(stub.nhdr, 2)
        hdulist = list()
        hdr = sim_header()
        hdr['SIMPLE'] = True
        hdr['BITPIX'] = 8
        hdr['NAXIS'] = 0
        hdr['EXTEND'] = True
        hdulist.append(sim_hdu(hdr))
        for k in range(10):
            hdr = sim_header()
            hdr['XTENSION'] = 'IMAGE'
            hdr['BITPIX'] = -32
            hdr['NAXIS'] = 2
            hdr['NAXIS1'] = 10
            hdr['NAXIS2'] = 10
            hdr['EXTNAME'] = 'HDU{0:02d}'.format(k+1)
            hdulist.append(sim_hdu(hdr))
        stub = Stub(hdulist)
        self.assertEqual(stub.nhdr, 11)
        self.assertEqual(stub.hduname, 'HDU{0:02d}')
        hdulist = list()
        hdr = sim_header()
        hdr['SIMPLE'] = True
        hdr['BITPIX'] = 8
        hdr['NAXIS'] = 0
        hdr['EXTEND'] = True
        hdulist.append(sim_hdu(hdr))
        for k in range(100):
            hdr = sim_header()
            hdr['XTENSION'] = 'IMAGE'
            hdr['BITPIX'] = -32
            hdr['NAXIS'] = 2
            hdr['NAXIS1'] = 10
            hdr['NAXIS2'] = 10
            hdr['EXTNAME'] = 'HDU{0:02d}'.format(k+1)
            hdulist.append(sim_hdu(hdr))
        stub = Stub(hdulist)
        self.assertEqual(stub.nhdr, 101)
        self.assertEqual(stub.hduname, 'HDU{0:03d}')
        hdulist = list()
        hdr = sim_header()
        hdr['SIMPLE'] = True
        hdr['BITPIX'] = 8
        hdr['NAXIS'] = 0
        hdr['EXTEND'] = True
        hdulist.append(sim_hdu(hdr))
        hdr = sim_header()
        hdr['XTENSION'] = 'IMAGE'
        hdr['BITPIX'] = -32
        hdr['NAXIS'] = 2
        hdr['NAXIS1'] = 10
        hdr['NAXIS2'] = 10
        hdr['EXTNAME'] = 'HDU1'
        hdulist.append(sim_hdu(hdr))
        hdr = sim_header()
        hdr['XTENSION'] = 'TABLE   '
        hdr['BITPIX'] = 8
        hdr['NAXIS'] = 2
        hdr['NAXIS1'] = 10
        hdr['NAXIS2'] = 10
        hdr['EXTNAME'] = 'HDU2'
        hdulist.append(sim_hdu(hdr))
        stub = Stub(hdulist)
        self.assertEqual(stub.nhdr, 3)
        meta = stub.hdumeta
        self.assertLog(log, 0, 'HDU0 has no EXTNAME set!')
        self.assertLog(log, 1, 'Unknown extension type: TABLE.')
        self.assertEqual(meta[1]['format'],
                         'Data: FITS image [float32, 10x10]')
        self.assertEqual(meta[2]['format'],
                         'Unknown extension type: TABLE.')
        hdulist = list()
        hdr = sim_header()
        hdr['SIMPLE'] = True
        hdr['BITPIX'] = 8
        hdr['NAXIS'] = 0
        hdr['EXTEND'] = True
        hdr['EXTNAME'] = 'PRIMARY'
        hdulist.append(sim_hdu(hdr))
        stub = Stub(hdulist)
        sec = stub.section(0)
        expected_sec = ['HDU0',
                        '----',
                        '',
                        'EXTNAME = PRIMARY',
                        '',
                        '*Summarize the contents of this HDU.*',
                        '',
                        'This HDU has no non-standard required keywords.',
                        '',
                        'Empty HDU.',
                        '']
        self.assertEqual(len(sec), len(expected_sec))
        for k in range(len(sec)):
            self.assertEqual(sec[k], expected_sec[k])

    def test_image_format(self):
        """Test format string for image HDUs.
        """
        hdr = sim_header()
        hdr['SIMPLE'] = True
        hdr['BITPIX'] = 8
        hdr['NAXIS'] = 0
        hdr['EXTEND'] = True
        i = image_format(hdr)
        self.assertEqual(i, 'Empty HDU.')
        hdr = sim_header()
        hdr['SIMPLE'] = True
        hdr['BITPIX'] = 8
        hdr['NAXIS'] = 1
        hdr['NAXIS1'] = 1000
        i = image_format(hdr)
        self.assertEqual(i, 'Data: FITS image [char, 1000]')
        hdr = sim_header()
        hdr['SIMPLE'] = True
        hdr['BITPIX'] = 16
        hdr['NAXIS'] = 2
        hdr['NAXIS1'] = 1000
        hdr['NAXIS2'] = 1000
        i = image_format(hdr)
        self.assertEqual(i, 'Data: FITS image [int16, 1000x1000]')
        hdr = sim_header()
        hdr['SIMPLE'] = True
        hdr['BITPIX'] = 128
        hdr['NAXIS'] = 2
        hdr['NAXIS1'] = 1000
        hdr['NAXIS2'] = 1000
        i = image_format(hdr)
        self.assertEqual(i, 'Data: FITS image [BITPIX=128, 1000x1000]')
        hdr = sim_header()
        hdr['SIMPLE'] = True
        hdr['BITPIX'] = 128
        hdr['NAXIS'] = 2
        hdr['NAXIS1'] = 1000
        hdr['NAXIS2'] = 1000
        hdr['BUNIT'] = '10**-17 ergs/(cm**2*s*Angstrom)'
        with self.assertRaises(DataModelError) as e:
            i = image_format(hdr, True)
        self.assertEqual(str(e.exception), "'10**-17 ergs/(cm**2*s*Angstrom)' did not parse as fits unit: At col 8, Unit 'ergs' not supported by the FITS standard. Did you mean erg?")
        i = image_format(hdr, False)
        self.assertLog(log, 0, "'10**-17 ergs/(cm**2*s*Angstrom)' did not parse as fits unit: At col 8, Unit 'ergs' not supported by the FITS standard. Did you mean erg?")
        self.assertEqual(i, 'Data: FITS image [BITPIX=128, 1000x1000]')
        hdr = sim_header()
        hdr['SIMPLE'] = True
        hdr['BITPIX'] = 128
        hdr['NAXIS'] = 2
        hdr['NAXIS1'] = 1000
        hdr['NAXIS2'] = 1000
        hdr['BUNIT'] = '10**-17 erg/(cm**2*s*Angstrom)'
        i = image_format(hdr, False)
        self.assertLog(log, 1, "BUNIT   = '1e-17 erg / (Angstrom cm2 s)'")
        self.assertEqual(i, 'Data: FITS image [BITPIX=128, 1000x1000]')

    def test_extrakey(self):
        """Test the identification of non-boring keys.
        """
        extrakey_tests = {'TTYPE01': False,
                          'TFORM12': False,
                          'TUNIT00': False,
                          'TCOMM33': False,
                          'TDIM11': False,
                          'BITPIX': False,
                          'DEPNAM33': False,
                          'DEPVER33': False,
                          'ZTENSION': False,
                          'ZBITPIX': False,
                          'FOOBAR': True}
        for k in extrakey_tests:
            if extrakey_tests[k]:
                self.assertTrue(extrakey(k))
            else:
                self.assertFalse(extrakey(k))

    @patch('os.path.getsize')
    def test_file_size(self, mock_size):
        """Test the determination and formatting of file size.
        """
        sizes = {5: '5 bytes',
                 2048: '2 KB',
                 3145728: '3 MB',
                 6442450944: '6 GB',
                 7696581394432: '7.0 TB'}
        for s in sizes:
            mock_size.return_value = s
            self.assertEqual(file_size(s), sizes[s])

    def test_fits_column_format(self):
        """Test the translation of FITS column format strings.
        """
        formats = {
            '1PB': '8-bit stream',
            '1PI': '16-bit stream',
            '1PJ': '32-bit stream',
            'A': 'char[1]',
            'B': 'binary',
            'L': 'logical',
            'E': 'float32',
            'D': 'float64',
            'I': 'int16',
            'J': 'int32',
            'K': 'int64',
            '10D': 'float64[10]',
            '20J': 'int32[20]'}
        for f in formats:
            ff = fits_column_format(f)
            self.assertEqual(ff, formats[f])

    def test_extract_keywords(self):
        """Test the parsing of a full FITS HDU.
        """
        hdr = sim_header()
        hdr['SIMPLE'] = True
        hdr['BITPIX'] = 8
        hdr['NAXIS'] = 0
        hdr['EXTEND'] = True
        lines = extract_keywords(hdr)
        self.assertEqual(len(lines), 0)
        hdr['BOOLEAN'] = True
        hdr['VERSION'] = '0.1.2'
        hdr['INTEGER'] = 12345
        hdr['FLOAT'] = 3.14159
        hdr['UNDR_'] = 'underscore_'
        lines = extract_keywords(hdr)
        expected_lines = [('BOOLEAN', 'T', 'bool',
                           'This is the comment on BOOLEAN.'),
                          ('VERSION', '0.1.2', 'str',
                           'This is the comment on VERSION.'),
                          ('INTEGER', '12345', 'int',
                           'This is the comment on INTEGER.'),
                          ('FLOAT', '3.14159', 'float',
                           'This is the comment on FLOAT.'),
                          ('UNDR\\_', 'underscore\\_', 'str',
                           'This is the comment on UNDR\\_.')]
        for k in range(len(lines)):
            self.assertEqual(lines[k], expected_lines[k])
        hdr = sim_header()
        hdr['SIMPLE'] = True
        hdr['BITPIX'] = 8
        hdr['NAXIS'] = 0
        hdr['EXTEND'] = True
        # Artificial example, but needed to trigger an exception
        hdr['FOOBAR'] = (3.14159, 2.71828)
        lines = extract_keywords(hdr)
        expected_lines = [('FOOBAR', (3.14159, 2.71828),
                           'Unknown', 'This is the comment on FOOBAR.')]
        self.assertLog(log, 0, "Raised AttributeError on FOOBAR = (3.14159, 2.71828).")
        for k in range(len(lines)):
            self.assertEqual(lines[k], expected_lines[k])

    def test_process_file(self):
        """Full test of parsing a FITS file.
        """
        filename = resource_filename('desidatamodel.test', 't/fits_file.fits')
        modelfile = resource_filename('desidatamodel.test', 't/fits_file.rst')
        with open(modelfile) as m:
            modeldata = m.read()
        stub = Stub(filename)
        data = str(stub)
        self.assertEqual(stub.modelname, 'fits_file')
        modellines = modeldata.split('\n')
        for i, l in enumerate(data.split('\n')):
            self.assertEqual(l, modellines[i])

def test_suite():
    """Allows testing of only this module with the command::

        python setup.py test -m <modulename>
    """
    return unittest.defaultTestLoader.loadTestsFromName(__name__)
<commit_msg>test compressed HDUs<commit_after># Licensed under a 3-clause BSD style license - see LICENSE.rst
# -*- coding: utf-8 -*-
"""Test desidatamodel.stub functions
"""
import os
import unittest
from unittest.mock import patch
from pkg_resources import resource_filename
from astropy.io import fits
from collections import OrderedDict

from .datamodeltestcase import DataModelTestCase
from .. import DataModelError
from ..stub import (Stub, extrakey, file_size, fits_column_format,
                    extract_keywords, image_format, log)

class sim_comments(dict):
    """Simulate a dictionary.
    """
    def __getitem__(self, key):
        return "This is the comment on {0}.".format(key)

class sim_header(OrderedDict):
    """Simulate a FITS header object.
    """
    comments = sim_comments()

class sim_hdu(object):
    """Simulate a FITS HDU.
    """
    def __init__(self, header):
        self.header = header
        return

class TestStub(DataModelTestCase):

    def test_Stub(self):
        """Test aspects of initialization of Stub objects.
        """
        with fits.open(resource_filename('desidatamodel.test',
                                         't/fits_file.fits')) as hdulist:
            stub = Stub(hdulist)
            self.assertEqual(stub.nhdr, 2)
        hdulist = list()
        hdr = sim_header()
        hdr['SIMPLE'] = True
        hdr['BITPIX'] = 8
        hdr['NAXIS'] = 0
        hdr['EXTEND'] = True
        hdulist.append(sim_hdu(hdr))
        for k in range(10):
            hdr = sim_header()
            hdr['XTENSION'] = 'IMAGE'
            hdr['BITPIX'] = -32
            hdr['NAXIS'] = 2
            hdr['NAXIS1'] = 10
            hdr['NAXIS2'] = 10
            hdr['EXTNAME'] = 'HDU{0:02d}'.format(k+1)
            hdulist.append(sim_hdu(hdr))
        stub = Stub(hdulist)
        self.assertEqual(stub.nhdr, 11)
        self.assertEqual(stub.hduname, 'HDU{0:02d}')
        hdulist = list()
        hdr = sim_header()
        hdr['SIMPLE'] = True
        hdr['BITPIX'] = 8
        hdr['NAXIS'] = 0
        hdr['EXTEND'] = True
        hdulist.append(sim_hdu(hdr))
        for k in range(100):
            hdr = sim_header()
            hdr['XTENSION'] = 'IMAGE'
            hdr['BITPIX'] = -32
            hdr['NAXIS'] = 2
            hdr['NAXIS1'] = 10
            hdr['NAXIS2'] = 10
            hdr['EXTNAME'] = 'HDU{0:02d}'.format(k+1)
            hdulist.append(sim_hdu(hdr))
        stub = Stub(hdulist)
        self.assertEqual(stub.nhdr, 101)
        self.assertEqual(stub.hduname, 'HDU{0:03d}')
        hdulist = list()
        hdr = sim_header()
        hdr['SIMPLE'] = True
        hdr['BITPIX'] = 8
        hdr['NAXIS'] = 0
        hdr['EXTEND'] = True
        hdulist.append(sim_hdu(hdr))
        hdr = sim_header()
        hdr['XTENSION'] = 'IMAGE'
        hdr['BITPIX'] = -32
        hdr['NAXIS'] = 2
        hdr['NAXIS1'] = 10
        hdr['NAXIS2'] = 10
        hdr['EXTNAME'] = 'HDU1'
        hdulist.append(sim_hdu(hdr))
        hdr = sim_header()
        hdr['XTENSION'] = 'TABLE   '
        hdr['BITPIX'] = 8
        hdr['NAXIS'] = 2
        hdr['NAXIS1'] = 10
        hdr['NAXIS2'] = 10
        hdr['EXTNAME'] = 'HDU2'
        hdulist.append(sim_hdu(hdr))
        stub = Stub(hdulist)
        self.assertEqual(stub.nhdr, 3)
        meta = stub.hdumeta
        self.assertLog(log, 0, 'HDU0 has no EXTNAME set!')
        self.assertLog(log, 1, 'Unknown extension type: TABLE.')
        self.assertEqual(meta[1]['format'],
                         'Data: FITS image [float32, 10x10]')
        self.assertEqual(meta[2]['format'],
                         'Unknown extension type: TABLE.')
        hdulist = list()
        hdr = sim_header()
        hdr['SIMPLE'] = True
        hdr['BITPIX'] = 8
        hdr['NAXIS'] = 0
        hdr['EXTEND'] = True
        hdr['EXTNAME'] = 'PRIMARY'
        hdulist.append(sim_hdu(hdr))
        stub = Stub(hdulist)
        sec = stub.section(0)
        expected_sec = ['HDU0',
                        '----',
                        '',
                        'EXTNAME = PRIMARY',
                        '',
                        '*Summarize the contents of this HDU.*',
                        '',
                        'This HDU has no non-standard required keywords.',
                        '',
                        'Empty HDU.',
                        '']
        self.assertEqual(len(sec), len(expected_sec))
        for k in range(len(sec)):
            self.assertEqual(sec[k], expected_sec[k])

    def test_image_format(self):
        """Test format string for image HDUs.
        """
        hdr = sim_header()
        hdr['SIMPLE'] = True
        hdr['BITPIX'] = 8
        hdr['NAXIS'] = 0
        hdr['EXTEND'] = True
        i = image_format(hdr)
        self.assertEqual(i, 'Empty HDU.')
        hdr = sim_header()
        hdr['SIMPLE'] = True
        hdr['BITPIX'] = 8
        hdr['NAXIS'] = 1
        hdr['NAXIS1'] = 1000
        i = image_format(hdr)
        self.assertEqual(i, 'Data: FITS image [char, 1000]')
        hdr = sim_header()
        hdr['SIMPLE'] = True
        hdr['BITPIX'] = 16
        hdr['NAXIS'] = 2
        hdr['NAXIS1'] = 1000
        hdr['NAXIS2'] = 1000
        i = image_format(hdr)
        self.assertEqual(i, 'Data: FITS image [int16, 1000x1000]')
        hdr = sim_header()
        hdr['SIMPLE'] = True
        hdr['BITPIX'] = 128
        hdr['NAXIS'] = 2
        hdr['NAXIS1'] = 1000
        hdr['NAXIS2'] = 1000
        i = image_format(hdr)
        self.assertEqual(i, 'Data: FITS image [BITPIX=128, 1000x1000]')
        hdr = sim_header()
        hdr['SIMPLE'] = True
        hdr['BITPIX'] = 128
        hdr['NAXIS'] = 2
        hdr['NAXIS1'] = 1000
        hdr['NAXIS2'] = 1000
        hdr['BUNIT'] = '10**-17 ergs/(cm**2*s*Angstrom)'
        with self.assertRaises(DataModelError) as e:
            i = image_format(hdr, True)
        self.assertEqual(str(e.exception), "'10**-17 ergs/(cm**2*s*Angstrom)' did not parse as fits unit: At col 8, Unit 'ergs' not supported by the FITS standard. Did you mean erg?")
        i = image_format(hdr, False)
        self.assertLog(log, 0, "'10**-17 ergs/(cm**2*s*Angstrom)' did not parse as fits unit: At col 8, Unit 'ergs' not supported by the FITS standard. Did you mean erg?")
        self.assertEqual(i, 'Data: FITS image [BITPIX=128, 1000x1000]')
        hdr = sim_header()
        hdr['SIMPLE'] = True
        hdr['BITPIX'] = 128
        hdr['NAXIS'] = 2
        hdr['NAXIS1'] = 1000
        hdr['NAXIS2'] = 1000
        hdr['BUNIT'] = '10**-17 erg/(cm**2*s*Angstrom)'
        i = image_format(hdr, False)
        self.assertLog(log, 1, "BUNIT   = '1e-17 erg / (Angstrom cm2 s)'")
        self.assertEqual(i, 'Data: FITS image [BITPIX=128, 1000x1000]')
        #
        # Check compressed HDU
        #
        hdr = sim_header()
        hdr['SIMPLE'] = True
        hdr['ZBITPIX'] = 16
        hdr['NAXIS'] = 2
        hdr['NAXIS1'] = 1000
        hdr['NAXIS2'] = 1000
        i = image_format(hdr)
        self.assertEqual(i, 'Data: FITS image [int16 (compressed), 1000x1000]')
        hdr = sim_header()
        hdr['SIMPLE'] = True
        hdr['ZBITPIX'] = 128
        hdr['NAXIS'] = 2
        hdr['NAXIS1'] = 1000
        hdr['NAXIS2'] = 1000
        i = image_format(hdr)
        self.assertEqual(i, 'Data: FITS image [BITPIX=128 (compressed), 1000x1000]')

    def test_extrakey(self):
        """Test the identification of non-boring keys.
        """
        extrakey_tests = {'TTYPE01': False,
                          'TFORM12': False,
                          'TUNIT00': False,
                          'TCOMM33': False,
                          'TDIM11': False,
                          'BITPIX': False,
                          'DEPNAM33': False,
                          'DEPVER33': False,
                          'ZTENSION': False,
                          'ZBITPIX': False,
                          'FOOBAR': True}
        for k in extrakey_tests:
            if extrakey_tests[k]:
                self.assertTrue(extrakey(k))
            else:
                self.assertFalse(extrakey(k))

    @patch('os.path.getsize')
    def test_file_size(self, mock_size):
        """Test the determination and formatting of file size.
        """
        sizes = {5: '5 bytes',
                 2048: '2 KB',
                 3145728: '3 MB',
                 6442450944: '6 GB',
                 7696581394432: '7.0 TB'}
        for s in sizes:
            mock_size.return_value = s
            self.assertEqual(file_size(s), sizes[s])

    def test_fits_column_format(self):
        """Test the translation of FITS column format strings.
        """
        formats = {
            '1PB': '8-bit stream',
            '1PI': '16-bit stream',
            '1PJ': '32-bit stream',
            'A': 'char[1]',
            'B': 'binary',
            'L': 'logical',
            'E': 'float32',
            'D': 'float64',
            'I': 'int16',
            'J': 'int32',
            'K': 'int64',
            '10D': 'float64[10]',
            '20J': 'int32[20]'}
        for f in formats:
            ff = fits_column_format(f)
            self.assertEqual(ff, formats[f])

    def test_extract_keywords(self):
        """Test the parsing of a full FITS HDU.
        """
        hdr = sim_header()
        hdr['SIMPLE'] = True
        hdr['BITPIX'] = 8
        hdr['NAXIS'] = 0
        hdr['EXTEND'] = True
        lines = extract_keywords(hdr)
        self.assertEqual(len(lines), 0)
        hdr['BOOLEAN'] = True
        hdr['VERSION'] = '0.1.2'
        hdr['INTEGER'] = 12345
        hdr['FLOAT'] = 3.14159
        hdr['UNDR_'] = 'underscore_'
        lines = extract_keywords(hdr)
        expected_lines = [('BOOLEAN', 'T', 'bool',
                           'This is the comment on BOOLEAN.'),
                          ('VERSION', '0.1.2', 'str',
                           'This is the comment on VERSION.'),
                          ('INTEGER', '12345', 'int',
                           'This is the comment on INTEGER.'),
                          ('FLOAT', '3.14159', 'float',
                           'This is the comment on FLOAT.'),
                          ('UNDR\\_', 'underscore\\_', 'str',
                           'This is the comment on UNDR\\_.')]
        for k in range(len(lines)):
            self.assertEqual(lines[k], expected_lines[k])
        hdr = sim_header()
        hdr['SIMPLE'] = True
        hdr['BITPIX'] = 8
        hdr['NAXIS'] = 0
        hdr['EXTEND'] = True
        # Artificial example, but needed to trigger an exception
        hdr['FOOBAR'] = (3.14159, 2.71828)
        lines = extract_keywords(hdr)
        expected_lines = [('FOOBAR', (3.14159, 2.71828),
                           'Unknown', 'This is the comment on FOOBAR.')]
        self.assertLog(log, 0, "Raised AttributeError on FOOBAR = (3.14159, 2.71828).")
        for k in range(len(lines)):
            self.assertEqual(lines[k], expected_lines[k])

    def test_process_file(self):
        """Full test of parsing a FITS file.
        """
        filename = resource_filename('desidatamodel.test', 't/fits_file.fits')
        modelfile = resource_filename('desidatamodel.test', 't/fits_file.rst')
        with open(modelfile) as m:
            modeldata = m.read()
        stub = Stub(filename)
        data = str(stub)
        self.assertEqual(stub.modelname, 'fits_file')
        modellines = modeldata.split('\n')
        for i, l in enumerate(data.split('\n')):
            self.assertEqual(l, modellines[i])

def test_suite():
    """Allows testing of only this module with the command::

        python setup.py test -m <modulename>
    """
    return unittest.defaultTestLoader.loadTestsFromName(__name__)
<|endoftext|>


<commit_before>import unittest
from sqlalchemy_utils import database_exists, drop_database
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from sqlalchemy.dialects.postgresql import JSONB
from dateutil.parser import parse
from datetime import datetime
import importlib.util

from meerkat_abacus import data_management as manage
from meerkat_abacus import model, util, task_queue
from meerkat_abacus import config
spec = importlib.util.spec_from_file_location(
    "country_test",
    config.config_directory + config.country_config["country_tests"])
country_test = importlib.util.module_from_spec(spec)
spec.loader.exec_module(country_test)

class DbTest(unittest.TestCase):
    """
    Test setting up database functionality
    """
    def setUp(self):
        manage.create_db(config.DATABASE_URL,
                         model.Base,
                         drop=True)
        self.engine = create_engine(config.DATABASE_URL)
        self.Session = sessionmaker(bind=self.engine)
        self.session = self.Session()

    def tearDown(self):
        self.session.commit()
        self.session.close()
        if database_exists(config.DATABASE_URL):
            drop_database(config.DATABASE_URL)

    def test_create_db(self):
        if database_exists(config.DATABASE_URL):
            drop_database(config.DATABASE_URL)
        self.assertFalse(database_exists(config.DATABASE_URL))
        manage.create_db(config.DATABASE_URL,
                         model.Base,
                         drop=True)
        self.assertTrue(database_exists(config.DATABASE_URL))

    def test_locations(self):
        old_dir = manage.config.config_directory
        manage.config.config_directory = "meerkat_abacus/test/test_data/"
        manage.config.data_directory = "meerkat_abacus/test/test_data/"
        old_locs = manage.country_config["locations"]
        manage.country_config["locations"] = {
            "clinics": "demo_clinics.csv",
            "districts": "demo_districts.csv",
            "regions": "demo_regions.csv"
            }
        manage.import_locations(self.engine, self.session)
        results = self.session.query(model.Locations).all()
        self.assertEqual(len(results), 12)  # Important as we should not merge the final Clinic 1
        for r in results:
            #Checking that devideids are handled properly
            if r.name == "Clinic 1":
                if r.parent_location == 4:
                    self.assertEqual(r.deviceid, "1,6")
                else:
                    self.assertEqual(r.deviceid, "7")
        # This file has a clinic with a non existent district
        old_clinic_file = manage.country_config["locations"]["clinics"]
        manage.country_config["locations"]["clinics"] = "demo_clinics_error.csv"
        with self.assertRaises(KeyError):
            manage.import_locations(self.engine, self.session)

        #Clean Up
        manage.country_config["locations"]["clinics"] = old_clinic_file
        manage.country_config["locations"] = old_locs
        manage.config.config_directory = old_dir

    def test_table_data_from_csv(self):
        """Test table_data_from_csv"""

        manage.table_data_from_csv("demo_case", model.form_tables["case"],
                                   "meerkat_abacus/test/test_data/",
                                   self.session, self.engine,
                                   deviceids=["1", "2", "3", "4", "5", "6"],
                                   table_name=config.country_config["tables"]["case"])
        results = self.session.query(model.form_tables["case"]).all()
        self.assertEqual(len(results), 6)  # Only 6 of the cases have deviceids in 1-6
        for r in results:
            self.assertIn(r.uuid, ["1", "2", "3", "4", "5", "6"])

    def test_links(self):
        deviceids = ["1", "2", "3", "4", "5", "6", "7", "8"]

        manage.create_db(config.DATABASE_URL,
                         model.Base,
                         drop=True)

        manage.table_data_from_csv("demo_case", model.form_tables["case"],
                                   "meerkat_abacus/test/test_data/",
                                   self.session, self.engine,
                                   table_name=config.country_config["tables"]["case"],
                                   deviceids=deviceids)
        manage.table_data_from_csv("demo_alert", model.form_tables["alert"],
                                   "meerkat_abacus/test/test_data/",
                                   self.session, self.engine,
                                   deviceids=deviceids,
                                   table_name=config.country_config["tables"]["alert"])
        link_def = {
            "id": "test",
            "name": "Test",
            "from_table": "form_tables.case",
            "from_column": "meta/instanceID",
            "from_condition": "intro./visit:new",
            "from_date": "start",
            "to_table": "alert",
            "to_column": "link_to",
            "to_date": "end",
            "to_condition": "condition:yes",
            "which": "last",
            "data": {
                "status": {
                    "A": {"column": "letter",
                          "condition": "A"},
                    "B": {"column": "letter",
                          "condition": "B"},
                    "C": {"column": ["index", "letter"],
                          "condition": "7"}
                }
            }
        }
        old_links = manage.config.links.links
        manage.config.links.links = [link_def]
        manage.import_links(self.session)
        self.session.query(model.Links).delete()
        self.session.commit()

        manage.add_new_links()

        links = self.session.query(model.Links).all()
        # Should be 5 links
        self.assertEqual(5, len(links))
        for link in links:
            if link.link_value == "1":
                # Check that it got the latest
                self.assertEqual(link.data, {"status": "B"})
            if link.link_value == "8":
                self.assertEqual(link.data, {"status": "C"})
        #Clean up
        manage.config.links.links = old_links

    def test_db_setup(self):

        task_queue.set_up_db.apply().get()
        self.assertTrue(database_exists(config.DATABASE_URL))
        engine = self.engine
        session = self.session
        #Locations
        results = session.query(model.Locations)
        country_test.test_locations(results)

        if config.fake_data:
            for table in model.form_tables:
                results = session.query(model.form_tables[table])
                self.assertEqual(len(results.all()), )
        #Import variables
        agg_var = session.query(model.AggregationVariables).filter(
            model.AggregationVariables.id == "tot_1").first()
        self.assertEqual(agg_var.name, "Total")

        link_defs = session.query(model.LinkDefinitions)
        self.assertEqual(link_defs.first().name, "Alert Investigation")
        #To codes
        agg_var_female = session.query(model.AggregationVariables).filter(
            model.AggregationVariables.name == "Female").first()
        results = session.query(model.Data)
        sec_condition = agg_var.secondary_condition.split(":")
        number_of_totals = 0
        number_of_female = 0
        for row in results:
            if "tot_1" in row.variables.keys():
                number_of_totals += 1
            if str(agg_var_female.id) in row.variables.keys():
                number_of_female += 1
        total = session.query(model.form_tables["case"]).filter(
            model.form_tables["case"].data.contains(
                {sec_condition[0]: sec_condition[1]}))
        female = session.query(model.form_tables["case"]).filter(
            model.form_tables["case"].data.contains(
                {sec_condition[0]: sec_condition[1],
                 agg_var_female.db_column: agg_var_female.condition}))
        self.assertEqual(number_of_totals, len(total.all()))
        self.assertEqual(number_of_female, len(female.all()))
        #Add links
        manage.add_new_links()
        link_query = session.query(model.Links).filter(
            model.Links.link_def == "alert_investigation")
        links = {}
        for link in link_query:
            links[link.link_value] = link

        alert_query = session.query(model.Alerts)
        alerts = {}
        for a in alert_query:
            alerts[a.id] = a
        alert_inv_query = session.query(model.form_tables["alert"])
        alert_invs = {}
        for a in alert_inv_query:
            alert_invs.setdefault(a.data["pt./alert_id"], [])
            alert_invs[a.data["pt./alert_id"]].append(a)

        for alert_id in alerts.keys():
            if alert_id in alert_invs.keys():
                self.assertIn(alert_id, links.keys())
                if len(alert_invs[alert_id]) == 1:
                    self.assertEqual(links[alert_id].to_date,
                                     parse(alert_invs[alert_id][0].data["end"]))
                    labs = (alert_invs[alert_id][0]
                            .data["alert_labs./return_lab"])
                    country_test.test_alert_status(labs, links[alert_id])

                else:
                    investigations = alert_invs[alert_id]
                    largest_date = datetime(2015, 1, 1)
                    largest_inv = None
                    for inv in investigations:
                        if parse(inv.data["end"]) > largest_date:
                            largest_date = parse(inv.data["end"])
                            largest_inv = inv
                    self.assertEqual(links[alert_id].to_date, largest_date)
                    labs = (largest_inv
                            .data["alert_labs./return_lab"])
                    country_test.test_alert_status(labs, links[alert_id])
            else:
                self.assertNotIn(alert_id, links.keys())
        session.close()

        self.assertFalse(
            manage.set_up_everything(True, False, 100)
            )

    def test_get_proccess_data(self):
        old_fake = task_queue.config.fake_data
        old_s3 = task_queue.config.get_data_from_s3
        task_queue.config.fake_data = True
        task_queue.config.get_data_from_s3 = False
        manage.create_db(config.DATABASE_URL,
                         model.Base,
                         drop=True)

        numbers = {}
        manage.import_locations(self.engine, self.session)
        manage.add_fake_data(self.session, N=500, append=False)
        task_queue.get_proccess_data.apply().get()
        for table in model.form_tables:
            res = self.session.query(model.form_tables[table])
            numbers[table] = len(res.all())
        task_queue.get_proccess_data.apply().get()
        for table in model.form_tables:
            res = self.session.query(model.form_tables[table])
            self.assertEqual(numbers[table] + 5, len(res.all()))

        #Clean up
        task_queue.config.fake_data = old_fake
        task_queue.config.get_data_from_s3 = old_s3

if __name__ == "__main__":
    unittest.main()
<commit_msg>db_test bug<commit_after>import unittest
from sqlalchemy_utils import database_exists, drop_database
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from sqlalchemy.dialects.postgresql import JSONB
from dateutil.parser import parse
from datetime import datetime
import importlib.util

from meerkat_abacus import data_management as manage
from meerkat_abacus import model, util, task_queue
from meerkat_abacus import config
spec = importlib.util.spec_from_file_location(
    "country_test",
    config.config_directory + config.country_config["country_tests"])
country_test = importlib.util.module_from_spec(spec)
spec.loader.exec_module(country_test)

class DbTest(unittest.TestCase):
    """
    Test setting up database functionality
    """
    def setUp(self):
        manage.create_db(config.DATABASE_URL,
                         model.Base,
                         drop=True)
        self.engine = create_engine(config.DATABASE_URL)
        self.Session = sessionmaker(bind=self.engine)
        self.session = self.Session()

    def tearDown(self):
        self.session.commit()
        self.session.close()
        if database_exists(config.DATABASE_URL):
            drop_database(config.DATABASE_URL)

    def test_create_db(self):
        if database_exists(config.DATABASE_URL):
            drop_database(config.DATABASE_URL)
        self.assertFalse(database_exists(config.DATABASE_URL))
        manage.create_db(config.DATABASE_URL,
                         model.Base,
                         drop=True)
        self.assertTrue(database_exists(config.DATABASE_URL))

    def test_locations(self):
        old_dir = manage.config.config_directory
        manage.config.config_directory = "meerkat_abacus/test/test_data/"
        manage.config.data_directory = "meerkat_abacus/test/test_data/"
        old_locs = manage.country_config["locations"]
        manage.country_config["locations"] = {
            "clinics": "demo_clinics.csv",
            "districts": "demo_districts.csv",
            "regions": "demo_regions.csv"
            }
        manage.import_locations(self.engine, self.session)
        results = self.session.query(model.Locations).all()
        self.assertEqual(len(results), 12)  # Important as we should not merge the final Clinic 1
        for r in results:
            #Checking that devideids are handled properly
            if r.name == "Clinic 1":
                if r.parent_location == 4:
                    self.assertEqual(r.deviceid, "1,6")
                else:
                    self.assertEqual(r.deviceid, "7")
        # This file has a clinic with a non existent district
        old_clinic_file = manage.country_config["locations"]["clinics"]
        manage.country_config["locations"]["clinics"] = "demo_clinics_error.csv"
        with self.assertRaises(KeyError):
            manage.import_locations(self.engine, self.session)

        #Clean Up
        manage.country_config["locations"]["clinics"] = old_clinic_file
        manage.country_config["locations"] = old_locs
        manage.config.config_directory = old_dir

    def test_table_data_from_csv(self):
        """Test table_data_from_csv"""

        manage.table_data_from_csv("demo_case", model.form_tables["case"],
                                   "meerkat_abacus/test/test_data/",
                                   self.session, self.engine,
                                   deviceids=["1", "2", "3", "4", "5", "6"],
                                   table_name=config.country_config["tables"]["case"])
        results = self.session.query(model.form_tables["case"]).all()
        self.assertEqual(len(results), 6)  # Only 6 of the cases have deviceids in 1-6
        for r in results:
            self.assertIn(r.uuid, ["1", "2", "3", "4", "5", "6"])

    def test_links(self):
        deviceids = ["1", "2", "3", "4", "5", "6", "7", "8"]

        manage.create_db(config.DATABASE_URL,
                         model.Base,
                         drop=True)

        manage.table_data_from_csv("demo_case", model.form_tables["case"],
                                   "meerkat_abacus/test/test_data/",
                                   self.session, self.engine,
                                   table_name=config.country_config["tables"]["case"],
                                   deviceids=deviceids)
        manage.table_data_from_csv("demo_alert", model.form_tables["alert"],
                                   "meerkat_abacus/test/test_data/",
                                   self.session, self.engine,
                                   deviceids=deviceids,
                                   table_name=config.country_config["tables"]["alert"])
        link_def = {
            "id": "test",
            "name": "Test",
            "from_table": "form_tables.case",
            "from_column": "meta/instanceID",
            "from_condition": "intro./visit:new",
            "from_date": "start",
            "to_table": "alert",
            "to_column": "link_to",
            "to_date": "end",
            "to_condition": "condition:yes",
            "which": "last",
            "data": {
                "status": {
                    "A": {"column": "letter",
                          "condition": "A"},
                    "B": {"column": "letter",
                          "condition": "B"},
                    "C": {"column": ["index", "letter"],
                          "condition": "7"}
                }
            }
        }
        old_links = manage.config.links.links
        manage.config.links.links = [link_def]
        manage.import_links(self.session)
        self.session.query(model.Links).delete()
        self.session.commit()

        manage.add_new_links()

        links = self.session.query(model.Links).all()
        # Should be 5 links
        self.assertEqual(5, len(links))
        for link in links:
            if link.link_value == "1":
                # Check that it got the latest
                self.assertEqual(link.data, {"status": "B"})
            if link.link_value == "8":
                self.assertEqual(link.data, {"status": "C"})
        #Clean up
        manage.config.links.links = old_links

    def test_db_setup(self):

        task_queue.set_up_db.apply().get()
        self.assertTrue(database_exists(config.DATABASE_URL))
        engine = self.engine
        session = self.session
        #Locations
        results = session.query(model.Locations)
        country_test.test_locations(results)

        if config.fake_data:
            for table in model.form_tables:
                results = session.query(model.form_tables[table])
                self.assertEqual(len(results.all()), 500)
        #Import variables
        agg_var = session.query(model.AggregationVariables).filter(
            model.AggregationVariables.id == "tot_1").first()
        self.assertEqual(agg_var.name, "Total")

        link_defs = session.query(model.LinkDefinitions)
        self.assertEqual(link_defs.first().name, "Alert Investigation")
        #To codes
        agg_var_female = session.query(model.AggregationVariables).filter(
            model.AggregationVariables.name == "Female").first()
        results = session.query(model.Data)
        sec_condition = agg_var.secondary_condition.split(":")
        number_of_totals = 0
        number_of_female = 0
        for row in results:
            if "tot_1" in row.variables.keys():
                number_of_totals += 1
            if str(agg_var_female.id) in row.variables.keys():
                number_of_female += 1
        total = session.query(model.form_tables["case"]).filter(
            model.form_tables["case"].data.contains(
                {sec_condition[0]: sec_condition[1]}))
        female = session.query(model.form_tables["case"]).filter(
            model.form_tables["case"].data.contains(
                {sec_condition[0]: sec_condition[1],
                 agg_var_female.db_column: agg_var_female.condition}))
        self.assertEqual(number_of_totals, len(total.all()))
        self.assertEqual(number_of_female, len(female.all()))
        #Add links
        manage.add_new_links()
        link_query = session.query(model.Links).filter(
            model.Links.link_def == "alert_investigation")
        links = {}
        for link in link_query:
            links[link.link_value] = link

        alert_query = session.query(model.Alerts)
        alerts = {}
        for a in alert_query:
            alerts[a.id] = a
        alert_inv_query = session.query(model.form_tables["alert"])
        alert_invs = {}
        for a in alert_inv_query:
            alert_invs.setdefault(a.data["pt./alert_id"], [])
            alert_invs[a.data["pt./alert_id"]].append(a)

        for alert_id in alerts.keys():
            if alert_id in alert_invs.keys():
                self.assertIn(alert_id, links.keys())
                if len(alert_invs[alert_id]) == 1:
                    self.assertEqual(links[alert_id].to_date,
                                     parse(alert_invs[alert_id][0].data["end"]))
                    labs = (alert_invs[alert_id][0]
                            .data["alert_labs./return_lab"])
                    country_test.test_alert_status(labs, links[alert_id])

                else:
                    investigations = alert_invs[alert_id]
                    largest_date = datetime(2015, 1, 1)
                    largest_inv = None
                    for inv in investigations:
                        if parse(inv.data["end"]) > largest_date:
                            largest_date = parse(inv.data["end"])
                            largest_inv = inv
                    self.assertEqual(links[alert_id].to_date, largest_date)
                    labs = (largest_inv
                            .data["alert_labs./return_lab"])
                    country_test.test_alert_status(labs, links[alert_id])
            else:
                self.assertNotIn(alert_id, links.keys())
        session.close()

        self.assertFalse(
            manage.set_up_everything(True, False, 100)
            )

    def test_get_proccess_data(self):
        old_fake = task_queue.config.fake_data
        old_s3 = task_queue.config.get_data_from_s3
        task_queue.config.fake_data = True
        task_queue.config.get_data_from_s3 = False
        manage.create_db(config.DATABASE_URL,
                         model.Base,
                         drop=True)

        numbers = {}
        manage.import_locations(self.engine, self.session)
        manage.add_fake_data(self.session, N=500, append=False)
        task_queue.get_proccess_data.apply().get()
        for table in model.form_tables:
            res = self.session.query(model.form_tables[table])
            numbers[table] = len(res.all())
        task_queue.get_proccess_data.apply().get()
        for table in model.form_tables:
            res = self.session.query(model.form_tables[table])
            self.assertEqual(numbers[table] + 5, len(res.all()))

        #Clean up
        task_queue.config.fake_data = old_fake
        task_queue.config.get_data_from_s3 = old_s3

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


<commit_before>from bullet_train import BulletTrain
from django.contrib.contenttypes.models import ContentType
from django.contrib.auth.models import Permission
from django.conf import settings

def create_mc_permission(permission):
    ct, _ = ContentType.objects.get_or_create(
        app_label="mc",
        model=permission.content_type.model)
    new_perm, _ = Permission.objects.get_or_create(
        name=permission.name,
        content_type=ct,
        codename=permission.codename)
    for group in permission.group_set.all():
        group.permissions.add(new_perm)
    for user in permission.user_set.all():
        user.user_permissions.add(new_perm)

def flag_smith_has_feature(feature_name):
    if settings.FLAG_SMITH_API_KEY:
        flag_smith_key = settings.FLAG_SMITH_API_KEY.strip('"')
    else:
        flag_smith_key = ''
    bt = BulletTrain(environment_id=flag_smith_key)
    if bt:
        if bt.has_feature(feature_name):
            return bt.feature_enabled(feature_name)
    return False

class UserAlert(object):
    alert_type = ''  # identifies the type/class of an alert
    alert_style = 'info'  # info, error, success, warning
    message = ''  # message to be displayed to the user (can be HTML)
<commit_msg>[AC-8759-ian] add copy_m2m_fields util function<commit_after>from bullet_train import BulletTrain
from django.contrib.contenttypes.models import ContentType
from django.contrib.auth.models import Permission
from django.conf import settings

def create_mc_permission(permission):
    ct, _ = ContentType.objects.get_or_create(
        app_label="mc",
        model=permission.content_type.model)
    new_perm, _ = Permission.objects.get_or_create(
        name=permission.name,
        content_type=ct,
        codename=permission.codename)
    for group in permission.group_set.all():
        group.permissions.add(new_perm)
    for user in permission.user_set.all():
        user.user_permissions.add(new_perm)

def flag_smith_has_feature(feature_name):
    if settings.FLAG_SMITH_API_KEY:
        flag_smith_key = settings.FLAG_SMITH_API_KEY.strip('"')
    else:
        flag_smith_key = ''
    bt = BulletTrain(environment_id=flag_smith_key)
    if bt:
        if bt.has_feature(feature_name):
            return bt.feature_enabled(feature_name)
    return False

class UserAlert(object):
    alert_type = ''  # identifies the type/class of an alert
    alert_style = 'info'  # info, error, success, warning
    message = ''  # message to be displayed to the user (can be HTML)

def copy_m2m_fields(old_obj, new_obj, m2m_fields):
    for field in m2m_fields:
        field_pk = set(getattr(old_obj, field).all().values_list('pk', flat=True))
        getattr(new_obj, field).add(*field_pk)
<|endoftext|>


<commit_before>from __future__ import print_function

import time

import numpy as np
from keras.callbacks import Callback

class BatchTiming(Callback):
    """
    It measure robust stats for timing of batches and epochs.
    Useful for measuring the training process.

    For each epoch it prints median batch time and total epoch time.
    After training it prints overall median batch time and median epoch time.

    Usage: model.fit(X_train, Y_train, callbacks=[BatchTiming()])

    All times are in seconds.

    More info: https://keras.io/callbacks/
    """
    def on_train_begin(self, logs={}):
        self.all_batch_times = []
        self.all_epoch_times = []

    def on_epoch_begin(self, epoch, logs={}):
        self.epoch_batch_times = []

    def on_batch_begin(self, batch, logs={}):
        self.start_time = time.time()

    def on_batch_end(self, batch, logs={}):
        end_time = time.time()
        elapsed_time = end_time - self.start_time
        self.epoch_batch_times.append(elapsed_time)
        self.all_batch_times.append(elapsed_time)

    def on_epoch_end(self, epoch, logs={}):
        epoch_time = np.sum(self.epoch_batch_times)
        self.all_epoch_times.append(epoch_time)
        median_batch_time = np.median(self.epoch_batch_times)
        print('Epoch timing - batch (median): %0.5f, epoch: %0.5f (sec)' % \
            (median_batch_time, epoch_time))

    def on_train_end(self, logs={}):
        median_batch_time = np.median(self.all_batch_times)
        median_epoch_time = np.median(self.all_epoch_times)
        print('Overall - batch (median): %0.5f, epoch (median): %0.5f (sec)' % \
            (median_batch_time, median_epoch_time))

class SamplesPerSec(Callback):
    def __init__(self, batch_size):
        self.batch_size = batch_size

    def on_train_begin(self, logs={}):
        self.all_samples_per_sec = []

    def on_batch_begin(self, batch, logs={}):
        self.start_time = time.time()
        # self.batch_size = logs['size']

    def on_batch_end(self, batch, logs={}):
        end_time = time.time()
        elapsed_time = end_time - self.start_time
        samples_per_sec = self.batch_size / elapsed_time
        self.all_samples_per_sec.append(samples_per_sec)

    def on_epoch_end(self, epoch, logs={}):
        self.print_results()

    def print_results(self):
        print('Samples/sec: %0.2f' % np.median(self.all_samples_per_sec))

"""
Enables CUDA profiling (for usage in nvprof) just for a few batches.

The reasons are:

- profiling outputs are big (easily 100s MB - GBs) and repeating
- without a proper stop the outputs sometimes fail to save

Since initially the TensorFlow runtime may take time to optimize the graph we
skip a few epochs and then enable profiling for a few batches within the next
epoch.

It requires the `cudaprofile` package.
"""
class CudaProfile(Callback):
    import cudaprofile

    def __init__(self, warmup_epochs=0, batches_to_profile=None):
        self.warmup_epochs = warmup_epochs
        self.batches_to_profile = batches_to_profile
        self.enabled = False

    def set_params(self, params):
        self.params = params

    def on_epoch_begin(self, epoch, logs={}):
        if epoch == self.warmup_epochs:
            cudaprofile.start()
            self.enabled = True

    def on_batch_end(self, batch, logs={}):
        if self.enabled && batch >= batches_to_profile:
            cudaprofile.stop()
<commit_msg>Implement StagingAreaCallback for GPU prefetching.<commit_after>from __future__ import print_function

import time

import numpy as np
from keras.callbacks import Callback

class StagingAreaCallback(Callback):
    """
    It allows to prefetch input batches to GPU using TensorFlow StagingArea,
    making a simple asynchronous pipeline.

    The classic mechanism of copying input data to GPU in Keras with TensorFlow
    is `feed_dict`: a numpy array is synchronously copied from Python to TF memory
    and then using a host-to-device memcpy to GPU memory. The computation,
    however has to wait, which is wasteful.

    This class makes the HtoD memcpy asynchronous using a GPU-resident queue
    of size two (implemented by StaginArea). The mechanism is as follows:

    - at the beginning of an epoch one batch is `put()` into the queue
    - during each training step another is is `put()` into the queue and in
      parallel the batch already present at the GPU is `get()` from the queue
      at provide as tesnor input to the Keras model (this runs within a single
      `tf.Session.run()`)

    The input numpy arrays (features and targets) are provided via this
    callback and sliced into batches inside it. The last batch might be of
    smaller size without any problem (the StagingArea supports variable-sized
    batches and allows to enforce constant data sample shape). In the last
    batch zero-length slice is still put into the queue to keep the get+put
    operation uniform across all batches.

    Since it's hard to modify Keras to add more data to `feed_dict`, the data
    from numpy is fed into StagingArea in another `tf.Session.run()` before each
    training step via an intermediate `tf.Variable` and `feed_dict`. It is still
    synchronous. A better, though more complicated way would be to use TF queues
    (depracated) or Dataset API.

    In order to provide extra put() operation to `fetches`, we depend on a fork
    of Keras (https://github.com/bzamecnik/keras/tree/tf-function-session-run-args).
    A pull request to upstream will be made soon.

    Example usage:

    ```
    staging_area_callback = StagingAreaCallback(x_train, y_train, batch_size)

    image = Input(tensor=staging_area_callback.input_tensor)
    x = Dense(512, activation='relu')(image)
    digit = Dense(num_classes, activation='softmax')(x)
    model = Model(inputs=image, outputs=digit)

    model.compile(optimizer='sgd', loss='categorical_crossentropy',
        target_tensors=[staging_area_callback.target_tensor],
        fetches=staging_area_callback.extra_ops)

    model.fit(steps_per_epoch=steps_per_epoch, epochs=2,
        callbacks=[staging_area_callback])
    ```

    Full example: https://gist.github.com/bzamecnik/b520e2b1e199b193b715477929e39b22
    """
    def __init__(self, x, y, batch_size, prefetch_count=1):
        self.x = x
        self.y = y
        self.batch_size = batch_size
        self.prefetch_count = prefetch_count

        features_shape = (None,) + x.shape[1:]
        labels_shape = (None,) + y.shape[1:]

        with tf.device('/cpu:0'):
            # for feeding inputs to the the StagingArea
            # Let's try to decouple feeding data to StagingArea.put()
            # from the training batch session.run()
            # https://www.tensorflow.org/api_guides/python/reading_data#Preloaded_data
            self.features_batch_next_value = tf.placeholder(dtype=x.dtype, shape=features_shape)
            # - prevent the variable to be used as a model parameter: trainable=False, collections=[]
            # - allow dynamic variable shape (for the last batch): validate_shape=False
            features_batch_next = tf.Variable(self.features_batch_next_value, trainable=False, collections=[], validate_shape=False)
            self.labels_batch_next_value = tf.placeholder(dtype=y.dtype, shape=labels_shape)
            labels_batch_next = tf.Variable(self.labels_batch_next_value, trainable=False, collections=[], validate_shape=False)
        self.assign_next_batch = tf.group(features_batch_next.initializer, labels_batch_next.initializer)

        # will be used for prefetching to GPU
        area = tf.contrib.staging.StagingArea(
            dtypes=[x.dtype, y.dtype],
            shapes=[features_shape, labels_shape])

        self.area_put = area.put([features_batch_next.value(), labels_batch_next.value()])
        area_get_features, area_get_labels = area.get()
        self.area_size = area.size()
        self.area_clear = area.clear()

        self.input_tensor = area_get_features
        self.target_tensor = area_get_labels
        self.extra_ops = [self.area_put]

    def set_params(self, params):
        super().set_params(params)
        self.steps_per_epoch = self.params['steps']

    def _slice_batch(self, i):
        start = i * self.batch_size
        end = start + self.batch_size
        return (self.x[start:end], self.y[start:end])

    def _assign_batch(self, session, data):
        x_batch, y_batch = data
        session.run(self.assign_next_batch, feed_dict={
            self.features_batch_next_value: x_batch,
            self.labels_batch_next_value: y_batch})

    def on_epoch_begin(self, epoch, logs=None):
        sess = K.get_session()
        for i in range(self.prefetch_count):
            self._assign_batch(sess, self._slice_batch(i))
            sess.run(self.area_put)

    def on_batch_begin(self, batch, logs=None):
        sess = K.get_session()
        # Slice for `prefetch_count` last batches is empty.
        # It serves as a dummy value which is put into StagingArea
        # but never read.
        data = self._slice_batch(batch + self.prefetch_count)
        self._assign_batch(sess, data)

    def on_epoch_end(self, epoch, logs=None):
        sess = K.get_session()
        sess.run(self.area_clear)

class BatchTiming(Callback):
    """
    It measure robust stats for timing of batches and epochs.
    Useful for measuring the training process.

    For each epoch it prints median batch time and total epoch time.
    After training it prints overall median batch time and median epoch time.

    Usage: model.fit(X_train, Y_train, callbacks=[BatchTiming()])

    All times are in seconds.

    More info: https://keras.io/callbacks/
    """
    def on_train_begin(self, logs={}):
        self.all_batch_times = []
        self.all_epoch_times = []

    def on_epoch_begin(self, epoch, logs={}):
        self.epoch_batch_times = []

    def on_batch_begin(self, batch, logs={}):
        self.start_time = time.time()

    def on_batch_end(self, batch, logs={}):
        end_time = time.time()
        elapsed_time = end_time - self.start_time
        self.epoch_batch_times.append(elapsed_time)
        self.all_batch_times.append(elapsed_time)

    def on_epoch_end(self, epoch, logs={}):
        epoch_time = np.sum(self.epoch_batch_times)
        self.all_epoch_times.append(epoch_time)
        median_batch_time = np.median(self.epoch_batch_times)
        print('Epoch timing - batch (median): %0.5f, epoch: %0.5f (sec)' % \
            (median_batch_time, epoch_time))

    def on_train_end(self, logs={}):
        median_batch_time = np.median(self.all_batch_times)
        median_epoch_time = np.median(self.all_epoch_times)
        print('Overall - batch (median): %0.5f, epoch (median): %0.5f (sec)' % \
            (median_batch_time, median_epoch_time))

class SamplesPerSec(Callback):
    def __init__(self, batch_size):
        self.batch_size = batch_size

    def on_train_begin(self, logs={}):
        self.all_samples_per_sec = []

    def on_batch_begin(self, batch, logs={}):
        self.start_time = time.time()
        # self.batch_size = logs['size']

    def on_batch_end(self, batch, logs={}):
        end_time = time.time()
        elapsed_time = end_time - self.start_time
        samples_per_sec = self.batch_size / elapsed_time
        self.all_samples_per_sec.append(samples_per_sec)

    def on_epoch_end(self, epoch, logs={}):
        self.print_results()

    def print_results(self):
        print('Samples/sec: %0.2f' % np.median(self.all_samples_per_sec))

"""
Enables CUDA profiling (for usage in nvprof) just for a few batches.

The reasons are:

- profiling outputs are big (easily 100s MB - GBs) and repeating
- without a proper stop the outputs sometimes fail to save

Since initially the TensorFlow runtime may take time to optimize the graph we
skip a few epochs and then enable profiling for a few batches within the next
epoch.

It requires the `cudaprofile` package.
"""
class CudaProfile(Callback):

    def __init__(self, warmup_epochs=0, batches_to_profile=None):
        self.warmup_epochs = warmup_epochs
        self.batches_to_profile = batches_to_profile
        self.enabled = False

    def set_params(self, params):
        self.params = params

    def on_epoch_begin(self, epoch, logs={}):
        import cudaprofile
        if epoch == self.warmup_epochs:
            cudaprofile.start()
            self.enabled = True

    def on_batch_end(self, batch, logs={}):
        import cudaprofile
        if self.enabled and batch >= batches_to_profile:
            cudaprofile.stop()
<|endoftext|>


<commit_before>"""
Create data products from loop simulations
"""
import os
import warnings
import logging
import itertools
import toolz

import numpy as np
from scipy.interpolate import splev, splprep, interp1d
import scipy.ndimage
import astropy.units as u
from sunpy.sun import constants
from sunpy.coordinates.frames import HeliographicStonyhurst
import h5py
try:
    import distributed
except ImportError:
    warnings.warn('Dask distributed scheduler required for parallel execution')

class Observer(object):
    """
    Class for assembling AR from loops and creating data products from 2D projections.

    Parameters
    ----------
    field : `~synthesizAR.Field`
    instruments : `list`
    parallel : `bool`

    Examples
    --------
    Notes
    -----
    """

    def __init__(self, field, instruments, parallel=False):
        self.parallel = parallel
        self.field = field
        self.instruments = instruments
        self._channels_setup()

    def _channels_setup(self):
        """
        Tell each channel of each detector which wavelengths fall in it.
        """
        for instr in self.instruments:
            for channel in instr.channels:
                if channel['wavelength_range'] is not None:
                    channel['model_wavelengths'] = []
                    for wvl in self.field.loops[0].resolved_wavelengths:
                        if channel['wavelength_range'][0] <= wvl <= channel['wavelength_range'][-1]:
                            channel['model_wavelengths'].append(wvl)
                    if channel['model_wavelengths']:
                        channel['model_wavelengths'] = u.Quantity(channel['model_wavelengths'])

    @u.quantity_input
    def _interpolate_loops(self, ds: u.cm):
        """
        Interpolate loops to common resolution

        Interpolate all loops to a resolution (`ds`) below the minimum bin width
        of all of the instruments. This ensures that the image isn't 'patchy'
        when it is binned.
        """
        # Interpolate all loops in HEEQ coordinates
        total_coordinates = []
        interpolated_loop_coordinates = []
        for loop in self.field.loops:
            n_interp = int(np.ceil((loop.full_length/ds).decompose()))
            interpolated_s = np.linspace(loop.field_aligned_coordinate.value[0],
                                         loop.field_aligned_coordinate.value[-1], n_interp)
            interpolated_loop_coordinates.append(interpolated_s)
            nots, _ = splprep(loop.coordinates.value.T)
            total_coordinates.append(np.array(splev(np.linspace(0, 1, n_interp), nots)).T)

        total_coordinates = np.vstack(total_coordinates) * loop.coordinates.unit

        return total_coordinates, interpolated_loop_coordinates

    def build_detector_files(self, savedir, ds, **kwargs):
        """
        Create files to store interpolated counts before binning.

        Note
        ----
        After creating the instrument objects and passing them to the observer,
        it is always necessary to call this method.
        """
        file_template = os.path.join(savedir, '{}_counts.h5')
        total_coordinates, self._interpolated_loop_coordinates = self._interpolate_loops(ds)
        interp_s_shape = (int(np.median([s.shape for s in self._interpolated_loop_coordinates])),)
        for instr in self.instruments:
            # If no observer coordinate specified, use that of the magnetogram
            if instr.observer_coordinate is None:
                # FIXME: Setting attributes of other classes like this is bad!
                instr.observer_coordinate = (self.field.magnetogram.observer_coordinate
                                             .transform_to(HeliographicStonyhurst))
            chunks = kwargs.get('chunks', instr.observing_time.shape + interp_s_shape)
            dset_shape = instr.observing_time.shape + (len(total_coordinates),)
            instr.build_detector_file(file_template, dset_shape, chunks, self.field,
                                      parallel=self.parallel, **kwargs)
            with h5py.File(instr.counts_file, 'a') as hf:
                if 'coordinates' not in hf:
                    dset = hf.create_dataset('coordinates', data=total_coordinates.value)
                    dset.attrs['units'] = total_coordinates.unit.to_string()

    def flatten_detector_counts(self, **kwargs):
        """
        Calculate intensity for each loop, interpolate it to the appropriate spatial and temporal
        resolution, and store it. This is done either in serial or parallel.
        """
        if self.parallel:
            return self._flatten_detector_counts_parallel(**kwargs)
        else:
            self._flatten_detector_counts_serial(**kwargs)

    def _flatten_detector_counts_serial(self, **kwargs):
        emission_model = kwargs.get('emission_model', None)
        for instr in self.instruments:
            with h5py.File(instr.counts_file, 'a', driver=kwargs.get('hdf5_driver', None)) as hf:
                start_index = 0
                for interp_s, loop in zip(self._interpolated_loop_coordinates, self.field.loops):
                    instr.commit(instr.interpolate_and_store(
                        instr.observing_time.value, loop.velocity_x, loop, interp_s),
                        hf['velocity_x'], start_index)
                    instr.commit(instr.interpolate_and_store(
                        instr.observing_time.value, loop.velocity_y, loop, interp_s),
                        hf['velocity_y'], start_index)
                    instr.commit(instr.interpolate_and_store(
                        instr.observing_time.value, loop.velocity_z, loop, interp_s),
                        hf['velocity_z'], start_index)
                    instr.commit(instr.interpolate_and_store(
                        instr.observing_time.value, loop.electron_temperature, loop, interp_s),
                        hf['electron_temperature'], start_index)
                    instr.commit(instr.interpolate_and_store(
                        instr.observing_time.value, loop.ion_temperature, loop, interp_s),
                        hf['ion_temperature'], start_index)
                    instr.commit(instr.interpolate_and_store(
                        instr.observing_time.value, loop.density, loop, interp_s),
                        hf['density'], start_index)
                    start_index += interp_s.shape[0]
                instr.flatten_serial(self.field.loops, self._interpolated_loop_coordinates, hf,
                                     emission_model=emission_model)

    def _flatten_detector_counts_parallel(self, **kwargs):
        """
        Build custom Dask graph interpolating quantities for each in loop in time and space.
        """
        client = distributed.get_client()
        emission_model = kwargs.get('emission_model', None)
        futures = {}
        start_indices = np.insert(np.array(
            [s.shape[0] for s in self._interpolated_loop_coordinates]).cumsum()[:-1], 0, 0)
        for instr in self.instruments:
            # Create temporary files where interpolated results will be written
            tmp_dir = os.path.join(os.path.dirname(instr.counts_file), 'tmp_parallel_files')
            if not os.path.exists(tmp_dir):
                os.makedirs(tmp_dir)
            interp_futures = []
            for q in ['velocity_x', 'velocity_y', 'velocity_z', 'electron_temperature',
                      'ion_temperature', 'density']:
                partial_interp = toolz.curry(instr.interpolate_and_store)(
                    q, save_dir=tmp_dir, dset_name=q)
                loop_futures = client.map(partial_interp, self.field.loops,
                                          self._interpolated_loop_coordinates, start_indices)
                # Block until complete
                distributed.client.wait([loop_futures])
                interp_futures += loop_futures

            # Calculate and interpolate channel counts for instrument
            counts_futures = instr.flatten_parallel(self.field.loops,
                                                    self._interpolated_loop_coordinates,
                                                    tmp_dir, emission_model=emission_model)
            # Assemble into file and clean up
            assemble_future = client.submit(instr.assemble_arrays, interp_futures+counts_futures,
                                            instr.counts_file)
            futures[f'{instr.name}'] = client.submit(self._cleanup, assemble_future)

        return futures

    @staticmethod
    def _cleanup(filenames):
        for f in filenames:
            os.remove(f)
        os.rmdir(os.path.dirname(f))

    @staticmethod
    def assemble_map(observed_map, filename, time):
        observed_map.meta['tunit'] = time.unit.to_string()
        observed_map.meta['t_obs'] = time.value
        observed_map.save(filename)

    def bin_detector_counts(self, savedir, **kwargs):
        """
        Assemble pipelines for building maps at each timestep.

        Build pipeline for computing final synthesized data products. This can be done
        either in serial or parallel.

        Parameters
        ----------
        savedir : `str`
            Top level directory to save data products in
        """
        if self.parallel:
            futures = {instr.name: {} for instr in self.instruments}
            client = distributed.get_client()
        else:
            futures = None
        file_path_template = os.path.join(savedir, '{}', '{}', 'map_t{:06d}.fits')
        for instr in self.instruments:
            bins, bin_range = instr.make_detector_array(self.field)
            with h5py.File(instr.counts_file, 'r') as hf:
                reference_time = u.Quantity(hf['time'], hf['time'].attrs['units'])
            indices_time = [np.where(reference_time == time)[0][0] for time in instr.observing_time]
            for channel in instr.channels:
                header = instr.make_fits_header(self.field, channel)
                file_paths = [file_path_template.format(instr.name, channel['name'], i_time)
                              for i_time in indices_time]
                if not os.path.exists(os.path.dirname(file_paths[0])):
                    os.makedirs(os.path.dirname(file_paths[0]))
                # Parallel
                if self.parallel:
                    partial_detect = toolz.curry(instr.detect)(
                        channel, header=header, bins=bins, bin_range=bin_range)
                    map_futures = client.map(partial_detect, indices_time)
                    futures[instr.name][channel['name']] = client.map(
                        self.assemble_map, map_futures, file_paths, instr.observing_time)
                # Serial
                else:
                    for i, i_time in indices_time:
                        raw_map = instr.detect(channel, i_time, header, bins, bin_range)
                        self.assemble_map(raw_map, file_paths[i], instr.observing_time[i])

        return futures
<commit_msg>block between sets of channel maps [skip ci]<commit_after>"""
Create data products from loop simulations
"""
import os
import warnings
import logging
import itertools
import toolz

import numpy as np
from scipy.interpolate import splev, splprep, interp1d
import scipy.ndimage
import astropy.units as u
from sunpy.sun import constants
from sunpy.coordinates.frames import HeliographicStonyhurst
import h5py
try:
    import distributed
except ImportError:
    warnings.warn('Dask distributed scheduler required for parallel execution')

class Observer(object):
    """
    Class for assembling AR from loops and creating data products from 2D projections.

    Parameters
    ----------
    field : `~synthesizAR.Field`
    instruments : `list`
    parallel : `bool`

    Examples
    --------
    Notes
    -----
    """

    def __init__(self, field, instruments, parallel=False):
        self.parallel = parallel
        self.field = field
        self.instruments = instruments
        self._channels_setup()

    def _channels_setup(self):
        """
        Tell each channel of each detector which wavelengths fall in it.
        """
        for instr in self.instruments:
            for channel in instr.channels:
                if channel['wavelength_range'] is not None:
                    channel['model_wavelengths'] = []
                    for wvl in self.field.loops[0].resolved_wavelengths:
                        if channel['wavelength_range'][0] <= wvl <= channel['wavelength_range'][-1]:
                            channel['model_wavelengths'].append(wvl)
                    if channel['model_wavelengths']:
                        channel['model_wavelengths'] = u.Quantity(channel['model_wavelengths'])

    @u.quantity_input
    def _interpolate_loops(self, ds: u.cm):
        """
        Interpolate loops to common resolution

        Interpolate all loops to a resolution (`ds`) below the minimum bin width
        of all of the instruments. This ensures that the image isn't 'patchy'
        when it is binned.
        """
        # Interpolate all loops in HEEQ coordinates
        total_coordinates = []
        interpolated_loop_coordinates = []
        for loop in self.field.loops:
            n_interp = int(np.ceil((loop.full_length/ds).decompose()))
            interpolated_s = np.linspace(loop.field_aligned_coordinate.value[0],
                                         loop.field_aligned_coordinate.value[-1], n_interp)
            interpolated_loop_coordinates.append(interpolated_s)
            nots, _ = splprep(loop.coordinates.value.T)
            total_coordinates.append(np.array(splev(np.linspace(0, 1, n_interp), nots)).T)

        total_coordinates = np.vstack(total_coordinates) * loop.coordinates.unit

        return total_coordinates, interpolated_loop_coordinates

    def build_detector_files(self, savedir, ds, **kwargs):
        """
        Create files to store interpolated counts before binning.

        Note
        ----
        After creating the instrument objects and passing them to the observer,
        it is always necessary to call this method.
        """
        file_template = os.path.join(savedir, '{}_counts.h5')
        total_coordinates, self._interpolated_loop_coordinates = self._interpolate_loops(ds)
        interp_s_shape = (int(np.median([s.shape for s in self._interpolated_loop_coordinates])),)
        for instr in self.instruments:
            # If no observer coordinate specified, use that of the magnetogram
            if instr.observer_coordinate is None:
                # FIXME: Setting attributes of other classes like this is bad!
                instr.observer_coordinate = (self.field.magnetogram.observer_coordinate
                                             .transform_to(HeliographicStonyhurst))
            chunks = kwargs.get('chunks', instr.observing_time.shape + interp_s_shape)
            dset_shape = instr.observing_time.shape + (len(total_coordinates),)
            instr.build_detector_file(file_template, dset_shape, chunks, self.field,
                                      parallel=self.parallel, **kwargs)
            with h5py.File(instr.counts_file, 'a') as hf:
                if 'coordinates' not in hf:
                    dset = hf.create_dataset('coordinates', data=total_coordinates.value)
                    dset.attrs['units'] = total_coordinates.unit.to_string()

    def flatten_detector_counts(self, **kwargs):
        """
        Calculate intensity for each loop, interpolate it to the appropriate spatial and temporal
        resolution, and store it. This is done either in serial or parallel.
        """
        if self.parallel:
            return self._flatten_detector_counts_parallel(**kwargs)
        else:
            self._flatten_detector_counts_serial(**kwargs)

    def _flatten_detector_counts_serial(self, **kwargs):
        emission_model = kwargs.get('emission_model', None)
        for instr in self.instruments:
            with h5py.File(instr.counts_file, 'a', driver=kwargs.get('hdf5_driver', None)) as hf:
                start_index = 0
                for interp_s, loop in zip(self._interpolated_loop_coordinates, self.field.loops):
                    instr.commit(instr.interpolate_and_store(
                        instr.observing_time.value, loop.velocity_x, loop, interp_s),
                        hf['velocity_x'], start_index)
                    instr.commit(instr.interpolate_and_store(
                        instr.observing_time.value, loop.velocity_y, loop, interp_s),
                        hf['velocity_y'], start_index)
                    instr.commit(instr.interpolate_and_store(
                        instr.observing_time.value, loop.velocity_z, loop, interp_s),
                        hf['velocity_z'], start_index)
                    instr.commit(instr.interpolate_and_store(
                        instr.observing_time.value, loop.electron_temperature, loop, interp_s),
                        hf['electron_temperature'], start_index)
                    instr.commit(instr.interpolate_and_store(
                        instr.observing_time.value, loop.ion_temperature, loop, interp_s),
                        hf['ion_temperature'], start_index)
                    instr.commit(instr.interpolate_and_store(
                        instr.observing_time.value, loop.density, loop, interp_s),
                        hf['density'], start_index)
                    start_index += interp_s.shape[0]
                instr.flatten_serial(self.field.loops, self._interpolated_loop_coordinates, hf,
                                     emission_model=emission_model)

    def _flatten_detector_counts_parallel(self, **kwargs):
        """
        Build custom Dask graph interpolating quantities for each in loop in time and space.
        """
        client = distributed.get_client()
        emission_model = kwargs.get('emission_model', None)
        futures = {}
        start_indices = np.insert(np.array(
            [s.shape[0] for s in self._interpolated_loop_coordinates]).cumsum()[:-1], 0, 0)
        for instr in self.instruments:
            # Create temporary files where interpolated results will be written
            tmp_dir = os.path.join(os.path.dirname(instr.counts_file), 'tmp_parallel_files')
            if not os.path.exists(tmp_dir):
                os.makedirs(tmp_dir)
            interp_futures = []
            for q in ['velocity_x', 'velocity_y', 'velocity_z', 'electron_temperature',
                      'ion_temperature', 'density']:
                partial_interp = toolz.curry(instr.interpolate_and_store)(
                    q, save_dir=tmp_dir, dset_name=q)
                loop_futures = client.map(partial_interp, self.field.loops,
                                          self._interpolated_loop_coordinates, start_indices)
                # Block until complete
                distributed.client.wait([loop_futures])
                interp_futures += loop_futures

            # Calculate and interpolate channel counts for instrument
            counts_futures = instr.flatten_parallel(self.field.loops,
                                                    self._interpolated_loop_coordinates,
                                                    tmp_dir, emission_model=emission_model)
            # Assemble into file and clean up
            assemble_future = client.submit(instr.assemble_arrays, interp_futures+counts_futures,
                                            instr.counts_file)
            futures[f'{instr.name}'] = client.submit(self._cleanup, assemble_future)

        return futures

    @staticmethod
    def _cleanup(filenames):
        for f in filenames:
            os.remove(f)
        os.rmdir(os.path.dirname(f))

    @staticmethod
    def assemble_map(observed_map, filename, time):
        observed_map.meta['tunit'] = time.unit.to_string()
        observed_map.meta['t_obs'] = time.value
        observed_map.save(filename)

    def bin_detector_counts(self, savedir, **kwargs):
        """
        Assemble pipelines for building maps at each timestep.

        Build pipeline for computing final synthesized data products. This can be done
        either in serial or parallel.

        Parameters
        ----------
        savedir : `str`
            Top level directory to save data products in
        """
        if self.parallel:
            futures = {instr.name: {} for instr in self.instruments}
            client = distributed.get_client()
        else:
            futures = None
        file_path_template = os.path.join(savedir, '{}', '{}', 'map_t{:06d}.fits')
        for instr in self.instruments:
            bins, bin_range = instr.make_detector_array(self.field)
            with h5py.File(instr.counts_file, 'r') as hf:
                reference_time = u.Quantity(hf['time'], hf['time'].attrs['units'])
            indices_time = [np.where(reference_time == time)[0][0] for time in instr.observing_time]
            for channel in instr.channels:
                header = instr.make_fits_header(self.field, channel)
                file_paths = [file_path_template.format(instr.name, channel['name'], i_time)
                              for i_time in indices_time]
                if not os.path.exists(os.path.dirname(file_paths[0])):
                    os.makedirs(os.path.dirname(file_paths[0]))
                # Parallel
                if self.parallel:
                    partial_detect = toolz.curry(instr.detect)(
                        channel, header=header, bins=bins, bin_range=bin_range)
                    map_futures = client.map(partial_detect, indices_time)
                    futures[instr.name][channel['name']] = client.map(
                        self.assemble_map, map_futures, file_paths, instr.observing_time)
                    distributed.client.wait(futures[instr.name][channel['name']])
                # Serial
                else:
                    for i, i_time in indices_time:
                        raw_map = instr.detect(channel, i_time, header, bins, bin_range)
                        self.assemble_map(raw_map, file_paths[i], instr.observing_time[i])

        return futures
<|endoftext|>