<commit_before>import astropy
from astropy.io import fits
import numpy as np
from uvcal import UVCal
import datetime

class CALFITS(UVCal):
    """
    Defines a calfits-specific class for reading and writing uvfits files.
    """

    def _indexhdus(self, hdulist):
        # input a list of hdus
        # return a dictionary of table names
        tablenames = {}
        for i in range(len(hdulist)):
            try:
                tablenames[hdulist[i].header['EXTNAME']] = i
            except(KeyError):
                continue
        return tablenames

    def write_calfits(self, filename, spoof_nonessential=False,
                      run_check=True, run_check_acceptability=True, clobber=False):
        """
        Write the data to a uvfits file.

        Args:
            filename: The uvfits file to write to.
            spoof_nonessential: Option to spoof the values of optional
                UVParameters that are not set but are required for uvfits files.
                Default is False.
            run_check: Option to check for the existence and proper shapes of
                required parameters before writing the file. Default is True.
            run_check_acceptability: Option to check acceptability of the values of
                required parameters before writing the file. Default is True.

        """
        if run_check:
            self.check(run_check_acceptability=run_check_acceptability)

        if self.Njones > 1:
            jones_spacing = np.diff(self.jones_array)
            if np.min(jones_spacing) < np.max(jones_spacing):
                raise ValueError('The jones values are not evenly spaced.'
                                 'The calibration fits file format does not'
                                 ' support unevenly spaced polarizations.')

        today = datetime.date.today().strftime("Date: %d, %b %Y")
        prihdr = fits.Header()
        if self.cal_type != 'gain':
            sechdr = fits.Header()
            sechdr['EXTNAME'] = 'FLAGS'
        # Conforming to fits format
        prihdr['SIMPLE'] = True
        prihdr['BITPIX'] = 32
        prihdr['NAXIS'] = 5
        prihdr['TELESCOP'] = self.telescope_name
        prihdr['GNCONVEN'] = self.gain_convention
        prihdr['NTIMES'] = self.Ntimes
        prihdr['NFREQS'] = self.Nfreqs
        prihdr['NANTSDAT'] = self.Nants_data
        prihdr['NJONES'] = self.Njones
        prihdr['CALTYPE'] = self.cal_type
        prihdr['INTTIME'] = self.integration_time
        prihdr['CHWIDTH'] = self.channel_width
        prihdr['NANTSTEL'] = self.Nants_telescope
        prihdr['NSPWS'] = self.Nspws
        prihdr['XORIENT'] = self.x_orientation
        prihdr['FRQRANGE'] = ','.join(map(str, self.freq_range))
        prihdr['TMERANGE'] = ','.join(map(str, self.time_range))
        for line in self.history.splitlines():
            prihdr.add_history(line)

        for p in self.extra():
            ep = getattr(self, p)
            if ep.form is 'str':
                prihdr['{0}'.format(p.upper().replace('_', '')[:8])] = ep.value
            else:
                continue

        if self.observer:
            prihdr['OBSERVER'] = self.observer
        if self.git_origin_cal:
            prihdr['ORIGCAL'] = self.git_origin_cal
        if self.git_hash_cal:
            prihdr['HASHCAL'] = self.git_hash_cal

        if self.cal_type == 'unknown':
            raise ValueError("unknown calibration type. Do not know how to"
                             "store parameters")

        if self.cal_type == 'gain':
            # Set header variable for gain.
            prihdr['CTYPE4'] = ('FREQS', 'Frequency.')
            prihdr['CUNIT4'] = ('Hz', 'Units of frequecy.')
            prihdr['CRVAL4'] = self.freq_array[0][0]
            prihdr['CDELT4'] = self.channel_width

            prihdr['CTYPE5'] = ('ANTAXIS', 'See ANTARR in ANTENNA extension for values.')
            prihdr['CUNIT5'] = 'Integer'
            prihdr['CRVAL5'] = 0
            prihdr['CDELT5'] = -1

            # set the last axis for number of arrays.
            prihdr['CTYPE1'] = ('Narrays', 'Number of image arrays.')
            prihdr['CUNIT1'] = ('Integer', 'Number of image arrays. Increment.')
            prihdr['CDELT1'] = 1
            if self.input_flag_array is not None:
                prihdr['CRVAL1'] = (5, 'Number of image arrays.')
                pridata = np.concatenate([self.gain_array.real[:, :, :, :, np.newaxis],
                                          self.gain_array.imag[:, :, :, :, np.newaxis],
                                          self.flag_array[:, :, :, :, np.newaxis],
                                          self.input_flag_array[:, :, :, :, np.newaxis],
                                          self.quality_array[:, :, :, :, np.newaxis]],
                                         axis=-1)
            else:
                prihdr['CRVAL1'] = (4, 'Number of image arrays.')
                pridata = np.concatenate([self.gain_array.real[:, :, :, :, np.newaxis],
                                          self.gain_array.imag[:, :, :, :, np.newaxis],
                                          self.flag_array[:, :, :, :, np.newaxis],
                                          self.quality_array[:, :, :, :, np.newaxis]],
                                         axis=-1)

        if self.cal_type == 'delay':
            # Set header variable for delay.

            prihdr['CTYPE4'] = ('ANTAXIS', 'See ANTARR in ANTENNA extension for values.')
            prihdr['CUNIT4'] = 'Integer'
            prihdr['CRVAL4'] = 0
            prihdr['CDELT4'] = -1

            # set the last axis for number of arrays.
            prihdr['CTYPE1'] = ('Narrays', 'Number of image arrays.')
            prihdr['CUNIT1'] = ('Integer', 'Number of image arrays. Value.')
            prihdr['CRVAL1'] = (2, 'Number of image arrays.')
            prihdr['CDELT1'] = 1

            pridata = np.concatenate([self.delay_array[:, :, :, np.newaxis],
                                      self.quality_array[:, :, :, np.newaxis]],
                                     axis=-1)

            # Set headers for the second hdu containing the flags. Only in cal_type=delay.
            sechdr['CTYPE2'] = ('JONES', 'Jones matrix array')
            sechdr['CUNIT2'] = ('Integer', 'representative integer for polarization.')
            sechdr['CRVAL2'] = self.jones_array[0]  # always start with first jones.
            if self.Njones > 1:
                sechdr['CDELT2'] = jones_spacing[0]
            else:
                sechdr['CDELT2'] = -1

            sechdr['CTYPE3'] = ('TIME', 'Time axis.')
            sechdr['CUNIT3'] = ('JD', 'Time in julian date format')
            sechdr['CRVAL3'] = self.time_array[0]
            sechdr['CDELT3'] = self.integration_time

            sechdr['CTYPE4'] = ('FREQS', 'Valid frequencies to apply delay.')
            sechdr['CUNIT4'] = ('Hz', 'Units of frequecy.')
            sechdr['CRVAL4'] = self.freq_array[0][0]
            sechdr['CDELT4'] = self.channel_width

            sechdr['CTYPE5'] = ('ANTAXIS', 'See ANTARR in ANTENNA extension for values.')

            if self.input_flag_array is not None:
                secdata = np.concatenate([self.flag_array.astype(np.int64)[:, :, :, :, np.newaxis],
                                          self.input_flag_array.astype(np.int64)[:, :, :, :, np.newaxis]],
                                         axis=-1)
                sechdr['CTYPE1'] = ('Narrays', 'Number of image arrays.')
                sechdr['CUNIT1'] = ('Integer', 'Number of image arrays. Value.')
                sechdr['CRVAL1'] = (2, 'Number of image arrays.')
                sechdr['CDELT1'] = 1

            else:
                secdata = self.flag_array.astype(np.int64)[:, :, :, :, np.newaxis]  # Can't be bool
                sechdr['CTYPE1'] = ('Narrays', 'Number of image arrays.')
                sechdr['CUNIT1'] = ('Integer', 'Number of image arrays. Value.')
                sechdr['CRVAL1'] = (1, 'Number of image arrays.')
                sechdr['CDELT1'] = 1

        # primary header ctypes for NAXIS [ for both gain and delay cal_type.]
        # Check polarizations.
        prihdr['CTYPE2'] = ('JONES', 'Jones matrix array')
        prihdr['CUNIT2'] = ('Integer', 'representative integer for polarization.')
        prihdr['CRVAL2'] = self.jones_array[0]  # always start with first jones.
        if self.Njones > 1:
            prihdr['CDELT2'] = jones_spacing[0]
        else:
            prihdr['CDELT2'] = -1

        prihdr['CTYPE3'] = ('TIME', 'Time axis.')
        prihdr['CUNIT3'] = ('JD', 'Time in julian date format')
        prihdr['CRVAL3'] = self.time_array[0]
        prihdr['CDELT3'] = self.integration_time

        prihdu = fits.PrimaryHDU(data=pridata, header=prihdr)

        col1 = fits.Column(name='ANTNAME', format='8A',
                           array=self.antenna_names)
        col2 = fits.Column(name='ANTINDEX', format='D',
                           array=self.antenna_numbers)
        if self.Nants_data == self.Nants_telescope:
            col3 = fits.Column(name='ANTARR', format='D',
                               array=self.ant_array)
        else:
            ant_array_use = np.append(self.ant_array, np.zeros(self.Nants_telescope - self.Nants_data, dtype=np.int) - 1)
            col3 = fits.Column(name='ANTARR', format='D',
                               array=self.ant_array)
        cols = fits.ColDefs([col1, col2, col3])
        ant_hdu = fits.BinTableHDU.from_columns(cols)
        ant_hdu.header['EXTNAME'] = 'ANTENNAS'

        if self.cal_type != 'gain':
            prihdu = fits.PrimaryHDU(data=pridata, header=prihdr)
            sechdu = fits.ImageHDU(data=secdata, header=sechdr)
            hdulist = fits.HDUList([prihdu, ant_hdu, sechdu])

        else:
            prihdu = fits.PrimaryHDU(data=pridata, header=prihdr)
            hdulist = fits.HDUList([prihdu, ant_hdu])

        if float(astropy.__version__[0:3]) < 1.3:
            hdulist.writeto(filename, clobber=clobber)
        else:
            hdulist.writeto(filename, overwrite=clobber)

    def read_calfits(self, filename, run_check=True, run_check_acceptability=True):
        F = fits.open(filename)
        data = F[0].data
        hdr = F[0].header.copy()
        hdunames = self._indexhdus(F)

        anthdu = F[hdunames['ANTENNAS']]
        antdata = anthdu.data
        self.antenna_names = map(str, antdata['ANTNAME'])
        self.antenna_numbers = map(int, antdata['ANTINDEX'])
        self.ant_array = map(int, antdata['ANTARR'])
        if np.min(self.ant_array) < 0:
            self.ant_array = self.ant_array[np.where(self.ant_array >= 0)[0]]

        self.Nfreqs = hdr['NFREQS']
        self.Njones = hdr['NJONES']
        self.Ntimes = hdr['NTIMES']
        self.channel_width = hdr['CHWIDTH']
        self.integration_time = hdr['INTTIME']
        self.telescope_name = hdr['TELESCOP']
        self.history = str(hdr.get('HISTORY', ''))
        if self.pyuvdata_version_str not in self.history.replace('\n', ''):
            if self.history.endswith('\n'):
                self.history += self.pyuvdata_version_str
            else:
                self.history += '\n' + self.pyuvdata_version_str
        while 'HISTORY' in hdr.keys():
            hdr.remove('HISTORY')
        self.freq_range = map(float, hdr['FRQRANGE'].split(','))
        self.time_range = map(float, hdr['TMERANGE'].split(','))
        self.Nspws = hdr['NSPWS']
        self.Nants_data = hdr['NANTSDAT']
        self.Nants_telescope = hdr['NANTSTEL']
        self.gain_convention = hdr['GNCONVEN']
        self.x_orientation = hdr['XORIENT']
        self.cal_type = hdr['CALTYPE']
        try:
            self.observer = hdr['OBSERVER']
        except:
            pass
        try:
            self.git_origin_cal = hdr['ORIGCAL']
        except:
            pass
        try:
            self.git_hash_cal = hdr['HASHCAL']
        except:
            pass

        # get data.
        if self.cal_type == 'gain':
            self.set_gain()
            self.gain_array = data[:, :, :, :, 0] + 1j * data[:, :, :, :, 1]
            self.flag_array = data[:, :, :, :, 2].astype('bool')
            if hdr['CRVAL1'] == 5:
                self.input_flag_array = data[:, :, :, :, 3].astype('bool')
                self.quality_array = data[:, :, :, :, 4]
            else:
                self.quality_array = data[:, :, :, :, 3]

            # generate frequency array.
            self.freq_array = np.arange(self.Nfreqs).reshape(1, -1) * hdr['CDELT4'] + hdr['CRVAL4']

        if self.cal_type == 'delay':
            self.set_delay()
            self.delay_array = data[:, :, :, 0]
            self.quality_array = data[:, :, :, 1]
            sechdu = F[hdunames['FLAGS']]
            flag_data = sechdu.data
            flag_hdr = sechdu.header
            if sechdu.header['CRVAL1'] == 2:
                self.flag_array = flag_data[:, :, :, :, 0].astype('bool')
                self.input_flag_array = flag_data[:, :, :, :, 1]
            else:
                self.flag_array = flag_data[:, :, :, :, 0].astype('bool')

            # generate frequency array.
            self.freq_array = np.arange(self.Nfreqs).reshape(1, -1) * flag_hdr['CDELT4'] + flag_hdr['CRVAL4']

        # generate polarization, and time array.
        self.jones_array = np.arange(self.Njones) * hdr['CDELT2'] + hdr['CRVAL2']
        self.time_array = np.arange(self.Ntimes) * hdr['CDELT3'] + hdr['CRVAL3']

        if run_check:
            self.check(run_check_acceptability=run_check_acceptability)
<commit_msg>improve comments<commit_after>import astropy
from astropy.io import fits
import numpy as np
from uvcal import UVCal
import datetime

class CALFITS(UVCal):
    """
    Defines a calfits-specific class for reading and writing uvfits files.
    """

    def _indexhdus(self, hdulist):
        # input a list of hdus
        # return a dictionary of table names
        tablenames = {}
        for i in range(len(hdulist)):
            try:
                tablenames[hdulist[i].header['EXTNAME']] = i
            except(KeyError):
                continue
        return tablenames

    def write_calfits(self, filename, spoof_nonessential=False,
                      run_check=True, run_check_acceptability=True, clobber=False):
        """
        Write the data to a uvfits file.

        Args:
            filename: The uvfits file to write to.
            spoof_nonessential: Option to spoof the values of optional
                UVParameters that are not set but are required for uvfits files.
                Default is False.
            run_check: Option to check for the existence and proper shapes of
                required parameters before writing the file. Default is True.
            run_check_acceptability: Option to check acceptability of the values of
                required parameters before writing the file. Default is True.

        """
        if run_check:
            self.check(run_check_acceptability=run_check_acceptability)

        if self.Njones > 1:
            jones_spacing = np.diff(self.jones_array)
            if np.min(jones_spacing) < np.max(jones_spacing):
                raise ValueError('The jones values are not evenly spaced.'
                                 'The calibration fits file format does not'
                                 ' support unevenly spaced polarizations.')

        today = datetime.date.today().strftime("Date: %d, %b %Y")
        prihdr = fits.Header()
        if self.cal_type != 'gain':
            sechdr = fits.Header()
            sechdr['EXTNAME'] = 'FLAGS'
        # Conforming to fits format
        prihdr['SIMPLE'] = True
        prihdr['BITPIX'] = 32
        prihdr['NAXIS'] = 5
        prihdr['TELESCOP'] = self.telescope_name
        prihdr['GNCONVEN'] = self.gain_convention
        prihdr['NTIMES'] = self.Ntimes
        prihdr['NFREQS'] = self.Nfreqs
        prihdr['NANTSDAT'] = self.Nants_data
        prihdr['NJONES'] = self.Njones
        prihdr['CALTYPE'] = self.cal_type
        prihdr['INTTIME'] = self.integration_time
        prihdr['CHWIDTH'] = self.channel_width
        prihdr['NANTSTEL'] = self.Nants_telescope
        prihdr['NSPWS'] = self.Nspws
        prihdr['XORIENT'] = self.x_orientation
        prihdr['FRQRANGE'] = ','.join(map(str, self.freq_range))
        prihdr['TMERANGE'] = ','.join(map(str, self.time_range))
        for line in self.history.splitlines():
            prihdr.add_history(line)

        for p in self.extra():
            ep = getattr(self, p)
            if ep.form is 'str':
                prihdr['{0}'.format(p.upper().replace('_', '')[:8])] = ep.value
            else:
                continue

        if self.observer:
            prihdr['OBSERVER'] = self.observer
        if self.git_origin_cal:
            prihdr['ORIGCAL'] = self.git_origin_cal
        if self.git_hash_cal:
            prihdr['HASHCAL'] = self.git_hash_cal

        if self.cal_type == 'unknown':
            raise ValueError("unknown calibration type. Do not know how to"
                             "store parameters")

        if self.cal_type == 'gain':
            # Set header variable for gain.
            # ANTAXIS axis number differs between delay and gain because there's no frequency axis in delay
            prihdr['CTYPE4'] = ('FREQS', 'Frequency.')
            prihdr['CUNIT4'] = ('Hz', 'Units of frequecy.')
            prihdr['CRVAL4'] = self.freq_array[0][0]
            prihdr['CDELT4'] = self.channel_width

            prihdr['CTYPE5'] = ('ANTAXIS', 'See ANTARR in ANTENNA extension for values.')
            prihdr['CUNIT5'] = 'Integer'
            prihdr['CRVAL5'] = 0
            prihdr['CDELT5'] = -1

            # set the last axis for number of arrays.
            prihdr['CTYPE1'] = ('Narrays', 'Number of image arrays.')
            prihdr['CUNIT1'] = ('Integer', 'Number of image arrays. Increment.')
            prihdr['CDELT1'] = 1
            if self.input_flag_array is not None:
                prihdr['CRVAL1'] = (5, 'Number of image arrays.')
                pridata = np.concatenate([self.gain_array.real[:, :, :, :, np.newaxis],
                                          self.gain_array.imag[:, :, :, :, np.newaxis],
                                          self.flag_array[:, :, :, :, np.newaxis],
                                          self.input_flag_array[:, :, :, :, np.newaxis],
                                          self.quality_array[:, :, :, :, np.newaxis]],
                                         axis=-1)
            else:
                prihdr['CRVAL1'] = (4, 'Number of image arrays.')
                pridata = np.concatenate([self.gain_array.real[:, :, :, :, np.newaxis],
                                          self.gain_array.imag[:, :, :, :, np.newaxis],
                                          self.flag_array[:, :, :, :, np.newaxis],
                                          self.quality_array[:, :, :, :, np.newaxis]],
                                         axis=-1)

        if self.cal_type == 'delay':
            # Set header variable for delay.

            # ANTAXIS axis number differs between delay and gain because there's no frequency axis in delay
            prihdr['CTYPE4'] = ('ANTAXIS', 'See ANTARR in ANTENNA extension for values.')
            prihdr['CUNIT4'] = 'Integer'
            prihdr['CRVAL4'] = 0
            prihdr['CDELT4'] = -1

            # set the last axis for number of arrays.
            prihdr['CTYPE1'] = ('Narrays', 'Number of image arrays.')
            prihdr['CUNIT1'] = ('Integer', 'Number of image arrays. Value.')
            prihdr['CRVAL1'] = (2, 'Number of image arrays.')
            prihdr['CDELT1'] = 1

            pridata = np.concatenate([self.delay_array[:, :, :, np.newaxis],
                                      self.quality_array[:, :, :, np.newaxis]],
                                     axis=-1)

            # Set headers for the second hdu containing the flags. Only in cal_type=delay.
            sechdr['CTYPE2'] = ('JONES', 'Jones matrix array')
            sechdr['CUNIT2'] = ('Integer', 'representative integer for polarization.')
            sechdr['CRVAL2'] = self.jones_array[0]  # always start with first jones.
            if self.Njones > 1:
                sechdr['CDELT2'] = jones_spacing[0]
            else:
                sechdr['CDELT2'] = -1

            sechdr['CTYPE3'] = ('TIME', 'Time axis.')
            sechdr['CUNIT3'] = ('JD', 'Time in julian date format')
            sechdr['CRVAL3'] = self.time_array[0]
            sechdr['CDELT3'] = self.integration_time

            sechdr['CTYPE4'] = ('FREQS', 'Valid frequencies to apply delay.')
            sechdr['CUNIT4'] = ('Hz', 'Units of frequecy.')
            sechdr['CRVAL4'] = self.freq_array[0][0]
            sechdr['CDELT4'] = self.channel_width

            sechdr['CTYPE5'] = ('ANTAXIS', 'See ANTARR in ANTENNA extension for values.')

            if self.input_flag_array is not None:
                secdata = np.concatenate([self.flag_array.astype(np.int64)[:, :, :, :, np.newaxis],
                                          self.input_flag_array.astype(np.int64)[:, :, :, :, np.newaxis]],
                                         axis=-1)
                sechdr['CTYPE1'] = ('Narrays', 'Number of image arrays.')
                sechdr['CUNIT1'] = ('Integer', 'Number of image arrays. Value.')
                sechdr['CRVAL1'] = (2, 'Number of image arrays.')
                sechdr['CDELT1'] = 1

            else:
                secdata = self.flag_array.astype(np.int64)[:, :, :, :, np.newaxis]  # Can't be bool
                sechdr['CTYPE1'] = ('Narrays', 'Number of image arrays.')
                sechdr['CUNIT1'] = ('Integer', 'Number of image arrays. Value.')
                sechdr['CRVAL1'] = (1, 'Number of image arrays.')
                sechdr['CDELT1'] = 1

        # primary header ctypes for NAXIS [ for both gain and delay cal_type.]
        # Check polarizations.
        prihdr['CTYPE2'] = ('JONES', 'Jones matrix array')
        prihdr['CUNIT2'] = ('Integer', 'representative integer for polarization.')
        prihdr['CRVAL2'] = self.jones_array[0]  # always start with first jones.
        if self.Njones > 1:
            prihdr['CDELT2'] = jones_spacing[0]
        else:
            prihdr['CDELT2'] = -1

        prihdr['CTYPE3'] = ('TIME', 'Time axis.')
        prihdr['CUNIT3'] = ('JD', 'Time in julian date format')
        prihdr['CRVAL3'] = self.time_array[0]
        prihdr['CDELT3'] = self.integration_time

        prihdu = fits.PrimaryHDU(data=pridata, header=prihdr)

        col1 = fits.Column(name='ANTNAME', format='8A',
                           array=self.antenna_names)
        col2 = fits.Column(name='ANTINDEX', format='D',
                           array=self.antenna_numbers)
        if self.Nants_data == self.Nants_telescope:
            col3 = fits.Column(name='ANTARR', format='D',
                               array=self.ant_array)
        else:
            # ant_array is shorter than the other columns.
            # Pad the extra rows with -1s. Need to undo on read.
            nants_add = self.Nants_telescope - self.Nants_data
            ant_array_use = np.append(self.ant_array,
                                      np.zeros(nants_add, dtype=np.int) - 1)
            col3 = fits.Column(name='ANTARR', format='D',
                               array=self.ant_array)
        cols = fits.ColDefs([col1, col2, col3])
        ant_hdu = fits.BinTableHDU.from_columns(cols)
        ant_hdu.header['EXTNAME'] = 'ANTENNAS'

        if self.cal_type != 'gain':
            prihdu = fits.PrimaryHDU(data=pridata, header=prihdr)
            sechdu = fits.ImageHDU(data=secdata, header=sechdr)
            hdulist = fits.HDUList([prihdu, ant_hdu, sechdu])

        else:
            prihdu = fits.PrimaryHDU(data=pridata, header=prihdr)
            hdulist = fits.HDUList([prihdu, ant_hdu])

        if float(astropy.__version__[0:3]) < 1.3:
            hdulist.writeto(filename, clobber=clobber)
        else:
            hdulist.writeto(filename, overwrite=clobber)

    def read_calfits(self, filename, run_check=True, run_check_acceptability=True):
        F = fits.open(filename)
        data = F[0].data
        hdr = F[0].header.copy()
        hdunames = self._indexhdus(F)

        anthdu = F[hdunames['ANTENNAS']]
        antdata = anthdu.data
        self.antenna_names = map(str, antdata['ANTNAME'])
        self.antenna_numbers = map(int, antdata['ANTINDEX'])
        self.ant_array = map(int, antdata['ANTARR'])
        if np.min(self.ant_array) < 0:
            # ant_array was shorter than the other columns, so it was padded with -1s.
            # Remove the padded entries.
            self.ant_array = self.ant_array[np.where(self.ant_array >= 0)[0]]

        self.Nfreqs = hdr['NFREQS']
        self.Njones = hdr['NJONES']
        self.Ntimes = hdr['NTIMES']
        self.channel_width = hdr['CHWIDTH']
        self.integration_time = hdr['INTTIME']
        self.telescope_name = hdr['TELESCOP']
        self.history = str(hdr.get('HISTORY', ''))
        if self.pyuvdata_version_str not in self.history.replace('\n', ''):
            if self.history.endswith('\n'):
                self.history += self.pyuvdata_version_str
            else:
                self.history += '\n' + self.pyuvdata_version_str
        while 'HISTORY' in hdr.keys():
            hdr.remove('HISTORY')
        self.freq_range = map(float, hdr['FRQRANGE'].split(','))
        self.time_range = map(float, hdr['TMERANGE'].split(','))
        self.Nspws = hdr['NSPWS']
        self.Nants_data = hdr['NANTSDAT']
        self.Nants_telescope = hdr['NANTSTEL']
        self.gain_convention = hdr['GNCONVEN']
        self.x_orientation = hdr['XORIENT']
        self.cal_type = hdr['CALTYPE']
        try:
            self.observer = hdr['OBSERVER']
        except:
            pass
        try:
            self.git_origin_cal = hdr['ORIGCAL']
        except:
            pass
        try:
            self.git_hash_cal = hdr['HASHCAL']
        except:
            pass

        # get data.
        if self.cal_type == 'gain':
            self.set_gain()
            self.gain_array = data[:, :, :, :, 0] + 1j * data[:, :, :, :, 1]
            self.flag_array = data[:, :, :, :, 2].astype('bool')
            if hdr['CRVAL1'] == 5:
                self.input_flag_array = data[:, :, :, :, 3].astype('bool')
                self.quality_array = data[:, :, :, :, 4]
            else:
                self.quality_array = data[:, :, :, :, 3]

            # generate frequency array from primary data unit.
            self.freq_array = np.arange(self.Nfreqs).reshape(1, -1) * hdr['CDELT4'] + hdr['CRVAL4']

        if self.cal_type == 'delay':
            self.set_delay()
            self.delay_array = data[:, :, :, 0]
            self.quality_array = data[:, :, :, 1]
            sechdu = F[hdunames['FLAGS']]
            flag_data = sechdu.data
            flag_hdr = sechdu.header
            if sechdu.header['CRVAL1'] == 2:
                self.flag_array = flag_data[:, :, :, :, 0].astype('bool')
                self.input_flag_array = flag_data[:, :, :, :, 1]
            else:
                self.flag_array = flag_data[:, :, :, :, 0].astype('bool')

            # generate frequency array from flag data unit (no freq axis in primary).
            self.freq_array = np.arange(self.Nfreqs).reshape(1, -1) * flag_hdr['CDELT4'] + flag_hdr['CRVAL4']

        # generate polarization and time array for either cal_type.
        self.jones_array = np.arange(self.Njones) * hdr['CDELT2'] + hdr['CRVAL2']
        self.time_array = np.arange(self.Ntimes) * hdr['CDELT3'] + hdr['CRVAL3']

        if run_check:
            self.check(run_check_acceptability=run_check_acceptability)
<|endoftext|>


<commit_before>"""
   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
"""

#!/usr/bin/python
import sys
import urllib2
from configuration import *
from BaseHTTPServer import BaseHTTPRequestHandler,HTTPServer
import webbrowser

clients = {}

#tarocco to speedup server reply
def _bare_address_string(self):
    host, port = self.client_address[:2]
    return str(host)
BaseHTTPRequestHandler.address_string = \
        _bare_address_string

def getMethodBy(rootNode,idname):
	if idname.isdigit():
	  return getMethodById(rootNode,idname)
	return getMethodByName(rootNode,idname)

def getMethodByName(rootNode,name,maxIter=5):
	val = None
	if hasattr(rootNode,name):
	 val = getattr(rootNode,name)
	else:
		pass
	 	#print("the method '" + name + "' is not part of this instance type:" + str(rootNode))
	return val

def getMethodById(rootNode,_id,maxIter=5):
	for i in runtimeInstances:
		if id(i)==_id:
			return i

	maxIter = maxIter - 1
	if maxIter<0:
	 	return None
	_id = int(_id)

	if id(rootNode)==_id:
		return rootNode

	try:
	 	if hasattr(rootNode,'children'):
	 	 	for i in rootNode.children.values():
	 	 	 	#print(str(i))
		 	 	val = getMethodById(i,_id, maxIter )
		 	 	if val!=None:
			 	 	return val
	except:
	 	pass
	return None

#This class will handles any incoming request from the browser
#The main application class can subclass this
#	In the do_POST and do_GET methods it is expected to receive requests such as:
#		- function calls with parameters
#		- file requests
class BaseApp(BaseHTTPRequestHandler,object):

	#this method is used to get the Application instance previously created
	#	managing on this, it is possible to switch to "single instance for multiple clients" or "multiple instance for multiple clients" execution way
	def instance(self):
		k = self.client_address[0]
		if not MULTIPLE_INSTANCE:
			k = 0 #overwrite the key value, so all clients will point the same instance
		if not(k in clients.keys()):
			runtimeInstances.append(self)

			clients[k]=self
			clients[k].attachments = "<script>var sendCommand = function (url,params) { var request = new XMLHttpRequest(); request.open('POST', url, false); request.send(params); };</script>"
		self.client = clients[k]

	def do_POST(self):
		self.instance()
		varLen = int(self.headers['Content-Length'])
		postVars = self.rfile.read(varLen)
		postVars = str(urllib2.unquote(postVars).decode('utf8'))
		postVars = postVars.split("&")
		paramDict = {}
		for s in postVars:
			paramDict[s.split("=")[0]] = s.split("=")[1]

		function = str(urllib2.unquote(self.path).decode('utf8'))
		self.processAll(function, paramDict)

	#Handler for the GET requests
	def do_GET(self):
		self.instance()
		params = str(urllib2.unquote(self.path).decode('utf8'))

		params = params.split("?")
		paramDict = dict()
		function = ""

		for p in params:
			if p.count("=")==0:
				function = p
				continue
			paramDict[p.split("=")[0]]=p.split("=")[1]

			if paramDict[p.split("=")[0]].count("'")==0 and paramDict[p.split("=")[0]].count('"')==0 and paramDict[p.split("=")[0]].count('.')==1:
				paramDict[p.split("=")[0]] = float(paramDict[p.split("=")[0]])
			else:
				if paramDict[p.split("=")[0]].isdigit() and paramDict[p.split("=")[0]].count("'")==0 and paramDict[p.split("=")[0]].count('"')==0 and paramDict[p.split("=")[0]].count('.')==0:
					paramDict[p.split("=")[0]] = int(paramDict[p.split("=")[0]])

		function=function[1:]
		self.processAll(function, paramDict)

		return

	def processAll(self, function, paramDict):
		ispath = True
		snake = None
		doNotCallMain = False
		if len(function.split("/"))>1:
			for attr in function.split("/"):
				if len(attr)==0:
					continue
				if ispath==False:
					break
				if snake==None:
					snake = getMethodBy(self.client.root,attr)
					ispath = ispath and (None!=snake)
					continue
				snake = getMethodBy(snake,attr)
				ispath = ispath and (None!=snake)
		else:
			function=function.replace('/','')
			if len(function)>0:
				snake = getMethodBy(self.client,function)
				ispath = ispath and (None!=snake)
			else:
				doNotCallMain = hasattr( self.client, "root" )
				ispath = True
				snake = self.main

		if ispath:
			ret = None
			if not doNotCallMain:
				ret = snake(**paramDict)

			if ret==None:
				self.send_response(200)
				self.send_header('Content-type','text/html')
				self.end_headers()

				self.wfile.write(self.client.attachments);
				self.wfile.write("<link href='" + BASE_ADDRESS + "style.css' rel='stylesheet' />");

				self.wfile.write( repr(self.client.root) )
			else:
				#here is the function that should return the content type
				self.send_response(200)
				self.send_header('Content-type',ret[1])
				self.end_headers()
				self.wfile.write( ret[0] )

		else:
			self.send_response(200)
			self.send_header('Content-type','text/html')
			self.end_headers()

			f = open("./"+function,"r+b")
			content = "".join(f.readlines())
			f.close()
			self.wfile.write(content)

#this method starts the webserver with a specific BaseApp subclass
def start(mainGuiClass):
	try:
		#Create a web server and define the handler to manage the incoming request
		server = HTTPServer(('', PORT_NUMBER), mainGuiClass)
		print( 'Started httpserver on port ' , PORT_NUMBER )

		webbrowser.open(BASE_ADDRESS)
		server.serve_forever()

	except KeyboardInterrupt:
		print( '^C received, shutting down the web server' )
		server.socket.close()

<commit_msg>Update server.py<commit_after>"""
   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
"""

#!/usr/bin/python
import sys
import urllib2
from configuration import *
from BaseHTTPServer import BaseHTTPRequestHandler,HTTPServer
import webbrowser

clients = {}

#tarocco to speedup server reply
def _bare_address_string(self):
    host, port = self.client_address[:2]
    return str(host)
BaseHTTPRequestHandler.address_string = \
        _bare_address_string

def getMethodBy(rootNode,idname):
	if idname.isdigit():
	  return getMethodById(rootNode,idname)
	return getMethodByName(rootNode,idname)

def getMethodByName(rootNode,name,maxIter=5):
	val = None
	if hasattr(rootNode,name):
	 val = getattr(rootNode,name)
	else:
		pass
	 	#print("the method '" + name + "' is not part of this instance type:" + str(rootNode))
	return val

def getMethodById(rootNode,_id,maxIter=5):
	for i in runtimeInstances:
		if id(i)==_id:
			return i

	maxIter = maxIter - 1
	if maxIter<0:
	 	return None
	_id = int(_id)

	if id(rootNode)==_id:
		return rootNode

	try:
	 	if hasattr(rootNode,'children'):
	 	 	for i in rootNode.children.values():
	 	 	 	#print(str(i))
		 	 	val = getMethodById(i,_id, maxIter )
		 	 	if val!=None:
			 	 	return val
	except:
	 	pass
	return None

#This class will handles any incoming request from the browser
#The main application class can subclass this
#	In the do_POST and do_GET methods it is expected to receive requests such as:
#		- function calls with parameters
#		- file requests
class BaseApp(BaseHTTPRequestHandler,object):

	#this method is used to get the Application instance previously created
	#	managing on this, it is possible to switch to "single instance for multiple clients" or "multiple instance for multiple clients" execution way
	def instance(self):
		k = self.client_address[0]
		if not MULTIPLE_INSTANCE:
			k = 0 #overwrite the key value, so all clients will point the same instance
		if not(k in clients.keys()):
			runtimeInstances.append(self)

			clients[k]=self
			clients[k].attachments = "<script>var sendCommand = function (url,params) { var request = new XMLHttpRequest(); request.open('POST', url, false); request.send(params); };</script>"
		self.client = clients[k]

	def do_POST(self):
		self.instance()
		varLen = int(self.headers['Content-Length'])
		postVars = self.rfile.read(varLen)
		postVars = str(urllib2.unquote(postVars).decode('utf8'))
		postVars = postVars.split("&")
		paramDict = {}
		for s in postVars:
			paramDict[s.split("=")[0]] = s.split("=")[1]

		function = str(urllib2.unquote(self.path).decode('utf8'))
		self.processAll(function, paramDict)

	#Handler for the GET requests
	def do_GET(self):
		self.instance()
		params = str(urllib2.unquote(self.path).decode('utf8'))

		params = params.split("?")
		paramDict = dict()
		function = ""

		for p in params:
			if p.count("=")==0:
				function = p
				continue
			paramDict[p.split("=")[0]]=p.split("=")[1]

			if paramDict[p.split("=")[0]].count("'")==0 and paramDict[p.split("=")[0]].count('"')==0 and paramDict[p.split("=")[0]].count('.')==1:
				paramDict[p.split("=")[0]] = float(paramDict[p.split("=")[0]])
			else:
				if paramDict[p.split("=")[0]].isdigit() and paramDict[p.split("=")[0]].count("'")==0 and paramDict[p.split("=")[0]].count('"')==0 and paramDict[p.split("=")[0]].count('.')==0:
					paramDict[p.split("=")[0]] = int(paramDict[p.split("=")[0]])

		function=function[1:]
		self.processAll(function, paramDict)

		return

	def processAll(self, function, paramDict):
		ispath = True
		snake = None
		doNotCallMain = False
		if len(function.split("/"))>1:
			for attr in function.split("/"):
				if len(attr)==0:
					continue
				if ispath==False:
					break
				if snake==None:
					snake = getMethodBy(self.client.root,attr)
					ispath = ispath and (None!=snake)
					continue
				snake = getMethodBy(snake,attr)
				ispath = ispath and (None!=snake)
		else:
			function=function.replace('/','')
			if len(function)>0:
				snake = getMethodBy(self.client,function)
				ispath = ispath and (None!=snake)
			else:
				doNotCallMain = hasattr( self.client, "root" )
				ispath = True
				snake = self.main

		if ispath:
			ret = None
			if not doNotCallMain:
				ret = snake(**paramDict)

			if ret==None:
				self.send_response(200)
				self.send_header('Content-type','text/html')
				self.end_headers()

				self.wfile.write(self.client.attachments);
				self.wfile.write("<link href='" + BASE_ADDRESS + "style.css' rel='stylesheet' />");

				self.wfile.write( repr(self.client.root) )
			else:
				#here is the function that should return the content type
				self.send_response(200)
				self.send_header('Content-type',ret[1])
				self.end_headers()
				self.wfile.write( ret[0] )

		else:
			self.send_response(200)
			self.send_header('Content-type','text/html')
			self.end_headers()

			f = open("./"+function,"r+b")
			content = "".join(f.readlines())
			f.close()
			self.wfile.write(content)

#this method starts the webserver with a specific BaseApp subclass
def start(mainGuiClass):
	try:
		#Create a web server and define the handler to manage the incoming request
		server = HTTPServer(('', PORT_NUMBER), mainGuiClass)
		print( 'Started httpserver on port ' , PORT_NUMBER )
		try:
			import android
			android.webbrowser.open( BASE_ADDRESS )
		except:
			webbrowser.open( BASE_ADDRESS )
		server.serve_forever()

	except KeyboardInterrupt:
		print( '^C received, shutting down the web server' )
		server.socket.close()

<|endoftext|>


<commit_before>from custom.ilsgateway.tanzania.handlers.keyword import KeywordHandler
from custom.ilsgateway.tanzania.reminders import HELP_REGISTERED, HELP_UNREGISTERED

class HelpHandler(KeywordHandler):

    def help(self):
        if self.user:
            self.respond(HELP_REGISTERED)
        else:
            self.respond(HELP_UNREGISTERED)

        return True
<commit_msg>Fixed problem with help handler<commit_after>from custom.ilsgateway.tanzania.handlers.keyword import KeywordHandler
from custom.ilsgateway.tanzania.reminders import HELP_REGISTERED, HELP_UNREGISTERED

class HelpHandler(KeywordHandler):

    def help(self):
        if self.user:
            self.respond(HELP_REGISTERED)
        else:
            self.respond(HELP_UNREGISTERED)

        return True

    def handle(self):
        return self.help()
<|endoftext|>


<commit_before>from __future__ import division
import unittest
import pymatgen
import sys
from pymatgen.io.vaspio import Poscar
from pymatgen.io.vaspio import Vasprun
from pymatgen.io.cifio import CifWriter
from pymatgen.io.cifio import CifParser
from pymatgen.core.lattice import Lattice
from pymatgen.core.structure import Structure
from pymatgen.transformations.standard_transformations import *
import numpy as np
import os

__status__ = "Test"

class DeformGeometry(object):
	"""
	A class which takes a CIF-file, generates a pymatgen core structure
	and applies a range of lattice deformations. The deformed lattices
	are to be used as inputs to an ab initio code such as VASP to compute
	the stress tensor and from this, the full elastic tensor of rank 4
	"""

	def __init__(self, path_to_cif):
		"""
		Constructor to initialize path to CIF-file

		Args:
			path_to_cif - string with full path to CIF-file, e.g.
			/home/MDEJONG1/pythonplayground/pymatgen/classes/7048.cif
		"""

		self._path = path_to_cif

	def CIF2struct(self):
		"""
		Parse CIF to pymatgen core structure for easy manipulation

		Args:
			none
		"""

		coords = list()
		species = []
		myCIF = CifParser(self._path).get_structures()[0]

		for k in range(len(myCIF._sites)):
			coords.append(Lattice(np.linalg.inv(myCIF._lattice._mc2d)).get_fractional_coords(myCIF._sites[k]._coords))
			species.append(myCIF._sites[k]._species.keys()[0])

		lattice = Lattice(myCIF._lattice._matrix)
		struct = Structure(lattice,species,coords)

		self.base_struct = struct
		self._coords = coords
		self._lattice = myCIF._lattice._matrix
		self._species = species

		return self.base_struct

	def get_residual_stress(self):
		"""
		Given a crystal structure, this method starts up a static calculation of the stress tensor. These
		are residual stresses which need to be subtracted later on.

		Args:
			none
		"""

		self._residual_stress = np.zeros((3,3))

		return self._residual_stress

	def deform(self, nd=0.02, ns=0.02, m=6, n=6): # this class still requires some user input
		"""
		Take original geometry (from CIF-file) and apply a range of deformations.
		Default values generally work well for metals. However, one might need to
		make changes for material such as oxides.

		Args:
			nd - maximum amount of normal strain
			ns - maximum amount of shear strain
			m - number of normal deformations used for fitting elastic constants per deformation mode, even integer required
			n - number of shear deformations used for fitting elastic constants per deformation mode, even integer required
		"""

		self.msteps = m
		self.nsteps = n

		if m%2!=0:
			raise ValueError("m has to be even.")

		if n%2!=0:
			raise ValueError("n has to be even.")

		mystrains = np.zeros((3, 3, np.int(m*3) + np.int(n*3)))

		defs = np.linspace(-nd, nd, num=m+1)
		defs = np.delete(defs, np.int(m/2), 0)
		defstructures = []

		counter = 0
		# First apply non-shear deformations
		# Depending on crystal symmetry, these may yield shear stresses, however
		for i1 in range(0, 3):

			for i2 in range(0, len(defs)):

				F = np.eye(3)
				F[i1, i1] = F[i1, i1] + defs[i2]
				C = np.transpose(F)*F
				E = 0.5*(C-np.eye(3)) # Green-Lagrange strain tensor
				Lp = self._lattice*F  # Deformation mapping of lattice vectors
				mystrains[:, :, counter] = E

				lat = Lattice(Lp)
				struc = Structure(lat,self._species,self._coords)
				defstructures.append(struc)
				counter += 1

		# Now apply shear deformations #
		sheardef = np.linspace(-ns, ns, num=n+1)
		sheardef = np.delete(sheardef, np.int(n/2), 0)

		F_index = [[0, 1], [0, 2], [1, 2]]
		for j1 in range(0, 3):

			for j2 in range(0, len(sheardef)):

				F = np.zeros((3,3))
				F = np.eye(3)
				F[F_index[j1][0], F_index[j1][1]] = F[F_index[j1][0], F_index[j1][1]] + sheardef[j2]
				F = np.matrix(F)

				C = np.transpose(F)*F
				E = 0.5*(C-np.eye(3)) # Green-Lagrange strain tensor
				Lp = self._lattice*F  # Deformation mapping of lattice vectors
				mystrains[:, :, counter] = E

				lat = Lattice(Lp)
				struc = Structure(lat,self._species,self._coords)
				defstructures.append(struc)
				counter += 1

		self.defstructures = defstructures
		self.mystrains = mystrains

		return self.defstructures, self.mystrains

	def get_stress_tensors(self, path_list):
		"""
		This method takes a list of paths, which contain output files (currently only vasprun.xml's are supported)
		containing the stress tensor, and stores these in a 3D-matrix.

		Args:
			path_list - list of paths that contain vasprun.xml's which contain calculated stress tensor for every
			deformed geometry
		"""

		self.path_list = path_list

		if np.int((self.msteps + self.nsteps)*3) != len(path_list):
			raise ValueError("Number of stress tensors does not match number of strain tensors.")

		stress_tensors = np.zeros((3, 3, np.int(self.msteps*3) + np.int(self.nsteps*3)))

		counter = 0
		for stress_path in path_list:

			A = Vasprun(stress_path)
			stressmatrix = A.ionic_steps[-1]['stress']
			stress_tensors[:, :, counter] = stressmatrix - self._residual_stress
			counter += 1

		self.stress_tensors = stress_tensors

	def fit_cij(self, tol1=0.02, sym=False):
		"""
		Take stress and strain tensors and fit elastic constants. By default, the elastic tensor is not symmetrized.

		Args:
			tol1 - tolerance used for comparing linear and quadratic fits of elastic constants
			sym - specifies whether or not the elastic tensor is symmetrized
		"""

Q = DeformGeometry('/home/MDEJONG1/pythonplayground/pymatgen/classes/7048.cif')

## Calling sequence ##
Q.CIF2struct()
Q.deform(0.02, 0.02, 4, 4)
Q.get_residual_stress()

# Perform VASP-calculations in specified paths #
Q.get_stress_tensors(['/home/MDEJONG1/pythonplayground/pymatgen/classes/vasprun.xml']*24)

#print Q.__dict__.keys()
#print Q.defstructures
#print Q.mystrains

#A = Vasprun('vasprun.xml')
#print A.ionic_steps[-1]['stress']

<commit_msg>updated genstrain.py<commit_after>from __future__ import division
import unittest
import pymatgen
import sys
from pymatgen.io.vaspio import Poscar
from pymatgen.io.vaspio import Vasprun
from pymatgen.io.cifio import CifWriter
from pymatgen.io.cifio import CifParser
from pymatgen.core.lattice import Lattice
from pymatgen.core.structure import Structure
from pymatgen.transformations.standard_transformations import *
import numpy as np
import os

__status__ = "Test"

class DeformGeometry(object):
	"""
	A class which takes a CIF-file, generates a pymatgen core structure
	and applies a range of lattice deformations. The deformed lattices
	are to be used as inputs to an ab initio code such as VASP to compute
	the stress tensor and from this, the full elastic tensor of rank 4
	"""

	def __init__(self, path_to_cif):
		"""
		Constructor to initialize path to CIF-file

		Args:
			path_to_cif - string with full path to CIF-file, e.g.
			/home/MDEJONG1/pythonplayground/pymatgen/classes/7048.cif
		"""

		self._path = path_to_cif

	def CIF2struct(self):
		"""
		Parse CIF to pymatgen core structure for easy manipulation

		Args:
			none
		"""

		coords = list()
		species = []
		myCIF = CifParser(self._path).get_structures()[0]

#		check this part, cifio-code does this already!

		print myCIF
		for k in range(len(myCIF._sites)):
			coords.append(Lattice(np.linalg.inv(myCIF._lattice._mc2d)).get_fractional_coords(myCIF._sites[k]._coords))
			species.append(myCIF._sites[k]._species.keys()[0])

		lattice = Lattice(myCIF._lattice._matrix)
		struct = Structure(lattice,species,coords)
		print struct

		self.base_struct = struct
#		self._coords = coords
#		self._lattice = myCIF._lattice._matrix
#		self._species = species

		return self.base_struct

	def get_residual_stress(self):
		"""
		Given a crystal structure, this method starts up a static calculation of the stress tensor. These
		are residual stresses which need to be subtracted later on.

		Args:
			none
		"""

		self._residual_stress = np.zeros((3,3))

		return self._residual_stress

	def deform(self, nd=0.02, ns=0.02, m=6, n=6): # this class still requires some user input
		"""
		Take original geometry (from CIF-file) and apply a range of deformations.
		Default values generally work well for metals. However, one might need to
		make changes for material such as oxides.

		Args:
			nd - maximum amount of normal strain
			ns - maximum amount of shear strain
			m - number of normal deformations used for fitting elastic constants per deformation mode, even integer required
			n - number of shear deformations used for fitting elastic constants per deformation mode, even integer required
		"""

		self.msteps = m
		self.nsteps = n

		if m%2!=0:
			raise ValueError("m has to be even.")

		if n%2!=0:
			raise ValueError("n has to be even.")

		mystrains = np.zeros((3, 3, np.int(m*3) + np.int(n*3)))

		defs = np.linspace(-nd, nd, num=m+1)
		defs = np.delete(defs, np.int(m/2), 0)
		defstructures = []

		counter = 0
		# First apply non-shear deformations
		# Depending on crystal symmetry, these may yield shear stresses, however
		for i1 in range(0, 3):

			for i2 in range(0, len(defs)):

				F = np.eye(3)
				F[i1, i1] = F[i1, i1] + defs[i2]
				C = np.transpose(F)*F
				E = 0.5*(C-np.eye(3)) # Green-Lagrange strain tensor
				Lp = self._lattice*F  # Deformation mapping of lattice vectors
				mystrains[:, :, counter] = E

				lat = Lattice(Lp)
				struc = Structure(lat,self._species,self._coords)
				defstructures.append(struc)
				counter += 1

		# Now apply shear deformations #
		sheardef = np.linspace(-ns, ns, num=n+1)
		sheardef = np.delete(sheardef, np.int(n/2), 0)

		F_index = [[0, 1], [0, 2], [1, 2]]
		for j1 in range(0, 3):

			for j2 in range(0, len(sheardef)):

				F = np.zeros((3,3))
				F = np.eye(3)
				F[F_index[j1][0], F_index[j1][1]] = F[F_index[j1][0], F_index[j1][1]] + sheardef[j2]
				F = np.matrix(F)

				C = np.transpose(F)*F
				E = 0.5*(C-np.eye(3)) # Green-Lagrange strain tensor
				Lp = self._lattice*F  # Deformation mapping of lattice vectors
				mystrains[:, :, counter] = E

				lat = Lattice(Lp)
				struc = Structure(lat,self._species,self._coords)
				defstructures.append(struc)
				counter += 1

		self.defstructures = defstructures
		self.mystrains = mystrains

		return self.defstructures, self.mystrains

	def get_stress_tensors(self, path_list):
		"""
		This method takes a list of paths, which contain output files (currently only vasprun.xml's are supported)
		containing the stress tensor, and stores these in a 3D-matrix.

		Args:
			path_list - list of paths that contain vasprun.xml's which contain calculated stress tensor for every
			deformed geometry
		"""

		self.path_list = path_list

		if np.int((self.msteps + self.nsteps)*3) != len(path_list):
			raise ValueError("Number of stress tensors does not match number of strain tensors.")

		stress_tensors = np.zeros((3, 3, np.int(self.msteps*3) + np.int(self.nsteps*3)))

		counter = 0
		for stress_path in path_list:

			A = Vasprun(stress_path)
			stressmatrix = A.ionic_steps[-1]['stress']
			stress_tensors[:, :, counter] = stressmatrix - self._residual_stress
			counter += 1

		self.stress_tensors = stress_tensors

	def fit_cij(self, tol1=0.02, sym=False):
		"""
		Take stress and strain tensors and fit elastic constants. By default, the elastic tensor is not symmetrized.

		Args:
			tol1 - tolerance used for comparing linear and quadratic fits of elastic constants
			sym - specifies whether or not the elastic tensor is symmetrized
		"""

Q = DeformGeometry('/home/MDEJONG1/pythonplayground/pymatgen/classes/7048.cif')

## Calling sequence ##
Q.CIF2struct()
#Q.deform(0.02, 0.02, 4, 4)
#Q.get_residual_stress()

# Perform VASP-calculations in specified paths #
#Q.get_stress_tensors(['/home/MDEJONG1/pythonplayground/pymatgen/classes/vasprun.xml']*24)

#print Q.__dict__.keys()
#print Q.defstructures
#print Q.mystrains

#A = Vasprun('vasprun.xml')
#print A.ionic_steps[-1]['stress']

<|endoftext|>


<commit_before>from unittest import TestCase
from unittest.mock import patch

from email_parser import renderer, errors, cmd
from email_parser.reader import Template
from email_parser.fs import Email

class TestTextRenderer(TestCase):

    def setUp(self):
        self.email = Email('name', 'locale', 'path', 'full_path')
        self.renderer = renderer.TextRenderer([], {}, self.email)

    def test_happy_path(self):
        placeholders = {'content': 'dummy content'}

        actual = self.renderer.render(placeholders)

        self.assertEqual('dummy content', actual)

    def test_concat_multiple_placeholders(self):
        placeholders = {'content1': 'dummy content', 'content2': 'dummy content'}
        expected = renderer.TEXT_EMAIL_PLACEHOLDER_SEPARATOR.join(placeholders.values())

        actual = self.renderer.render(placeholders)

        self.assertEqual(expected, actual)

    def test_ignore_subject(self):
        placeholders = {'content': 'dummy content', 'subject': 'dummy subject'}

        actual = self.renderer.render(placeholders)

        self.assertEqual('dummy content', actual)

    def test_ignore_empty_placeholders(self):
        placeholders = {'content': 'dummy content', 'empty': ''}

        actual = self.renderer.render(placeholders)

        self.assertEqual('dummy content', actual)

    def test_ignored_placeholders(self):
        placeholders = {'content': 'dummy content', 'ignore': 'test'}

        r = renderer.TextRenderer(['ignore'], {}, self.email)

        actual = r.render(placeholders)

        self.assertEqual('dummy content', actual)

    def test_use_text_and_url_for_links(self):
        placeholders = {'content': 'dummy [link_text](http://link_url) content'}

        actual = self.renderer.render(placeholders)

        self.assertEqual('dummy link_text (http://link_url) content', actual)

    def test_default_link_locale_for_links(self):
        placeholders = {'content': 'dummy [link_text](http://link_url?locale={link_locale}) content'}

        actual = self.renderer.render(placeholders)

        self.assertEqual('dummy link_text (http://link_url?locale=en) content', actual)

    def test_link_locale_for_links(self):
        placeholders = {'content': 'dummy [link_text](http://link_url?locale={link_locale}) content'}

        r = renderer.TextRenderer(['ignore'], {'locale': 'zh'}, self.email)
        actual = r.render(placeholders)

        self.assertEqual('dummy link_text (http://link_url?locale=zh) content', actual)

    def test_use_text_if_href_is_empty(self):
        placeholders = {'content': 'dummy [http://link_url]() content'}

        actual = self.renderer.render(placeholders)

        self.assertEqual('dummy http://link_url content', actual)

    def test_use_href_if_text_is_same(self):
        placeholders = {'content': 'dummy [http://link_url](http://link_url) content'}

        actual = self.renderer.render(placeholders)

        self.assertEqual('dummy http://link_url content', actual)

    def test_url_with_params(self):
        placeholders = {'content': 'dummy [param_link](https://something.com/thing?id=mooo) content'}

        actual = self.renderer.render(placeholders)

        self.assertEqual('dummy param_link (https://something.com/thing?id=mooo) content', actual)

    def test_unordered_list(self):
        placeholders = {'content': '- one\n- two\n- three'}

        actual = self.renderer.render(placeholders)

        self.assertEqual('- one\n- two\n- three', actual.strip())

    def test_ordered_list(self):
        placeholders = {'content': '1. one\n2. two\n3. three'}

        actual = self.renderer.render(placeholders)

        self.assertEqual('1. one\n2. two\n3. three', actual.strip())

class TestSubjectRenderer(TestCase):

    def setUp(self):
        self.renderer = renderer.SubjectRenderer()

    def test_happy_path(self):
        placeholders = {'content': 'dummy content', 'subject': 'dummy subject'}

        actual = self.renderer.render(placeholders)

        self.assertEqual('dummy subject', actual[0])

    def test_ab(self):
        placeholders = {'content': 'dummy content', 'subject': 'dummy subject',
                        'subject_b': 'bbb', 'subject_a': 'aaa'}

        actual = self.renderer.render(placeholders)

        self.assertEqual(['dummy subject', 'aaa', 'bbb', None], actual)

    def test_raise_error_for_missing_subject(self):
        placeholders = {'content': 'dummy content'}

        with self.assertRaises(errors.MissingSubjectError):
            self.renderer.render(placeholders)

class TestHtmlRenderer(TestCase):

    def _get_renderer(self, template_html, template_placeholders, **kwargs):
        template = Template(name='template_name',
                            styles=['template_style'],
                            content=template_html,
                            placeholders_order=template_placeholders)
        return renderer.HtmlRenderer(template,
                                     kwargs.get('link_locale_mappings', {}),
                                     kwargs.get('email', self.email),
                                     kwargs.get('settings', self.settings))

    def setUp(self):
        settings = cmd.default_settings()._asdict()
        settings['templates'] = 'dummy_templates'
        settings['images'] = 'dummy_images'
        self.settings = cmd.Settings(**settings)

        self.email = Email('name', 'locale', 'path', 'full_path')
        self.global_email = '<?xml version="1.0" encoding="UTF-8" ?><resources></resources>'

    @patch('email_parser.fs.read_file')
    def test_happy_path(self, mock_read):
        html = '{{content1}}'
        html_placeholders = ['content1']
        placeholders = {'content1': 'text1'}
        mock_read.side_effect = iter(['body {}'])

        renderer = self._get_renderer(html, html_placeholders)
        actual = renderer.render(placeholders)

        self.assertEqual('
text1
', actual)

    @patch('email_parser.fs.read_file')
    def test_empty_style(self, mock_read):
        html = '{{content}}'
        html_placeholders = ['content']
        placeholders = {'content': 'dummy_content'}
        mock_read.side_effect = iter([''])

        renderer = self._get_renderer(html, html_placeholders)
        actual = renderer.render(placeholders)

        self.assertEqual('
dummy_content
', actual)

    @patch('email_parser.fs.read_file')
    def test_include_raw_subject(self, mock_read):
        html = '{{subject}}'
        html_placeholders = ['subject']
        placeholders = {'subject': 'dummy_subject'}
        mock_read.side_effect = iter([''])

        renderer = self._get_renderer(html, html_placeholders)
        actual = renderer.render(placeholders)

        self.assertEqual('dummy_subject', actual)

    @patch('email_parser.fs.read_file')
    def test_include_base_url(self, mock_read):
        html = '{{base_url}}'
        html_placeholders = ['base_url']
        placeholders = {}
        mock_read.side_effect = iter([''])

        renderer = self._get_renderer(html, html_placeholders)
        actual = renderer.render(placeholders)

        self.assertEqual('dummy_images', actual)

    @patch('email_parser.fs.read_file')
    def test_ignore_missing_placeholders(self, mock_read):
        html = '{{content}}{{missing}}'
        html_placeholders = ['content', 'missing']
        placeholders = {'content': 'dummy_content'}
        mock_read.side_effect = iter([''])
        settings = self.settings._asdict()
        settings['strict'] = False
        settings = cmd.Settings(**settings)

        renderer = self._get_renderer(html, html_placeholders, settings=settings)
        actual = renderer.render(placeholders)

        self.assertEqual('
dummy_content
', actual)

    @patch('email_parser.fs.read_file')
    def test_fail_on_missing_placeholders(self, mock_read):
        html = '{{content}}{{missing}}'
        html_placeholders = ['content', 'missing']
        placeholders = {'content': 'dummy_content'}
        mock_read.side_effect = iter([''])

        renderer = self._get_renderer(html, html_placeholders)
        with self.assertRaises(errors.MissingTemplatePlaceholderError):
            renderer.render(placeholders)

    @patch('email_parser.fs.read_file')
    def test_rtl_locale(self, mock_read):
        html = '{{content}}'
        html_placeholders = ['content']
        placeholders = {'content': 'dummy_content'}
        mock_read.side_effect = iter([''])
        email_dict = self.email._asdict()
        email_dict['locale'] = 'ar'

        renderer = self._get_renderer(html, html_placeholders, email=Email(**email_dict))
        actual = renderer.render(placeholders)

        self.assertEqual('\n
\n  dummy_content\n
\n', actual)

    @patch('email_parser.fs.read_file')
    def test_rtl_two_placeholders(self, mock_read):
        html = '
{{content1}}

{{content2}}
'
        html_placeholders = ['content1', 'content2']
        placeholders = {'content1': 'dummy_content1', 'content2': 'dummy_content2'}
        mock_read.side_effect = iter([''])
        email_dict = self.email._asdict()
        email_dict['locale'] = 'ar'

        renderer = self._get_renderer(html, html_placeholders, email=Email(**email_dict))
        actual = renderer.render(placeholders)
        expected = '\n
\n
\n   dummy_content1\n
\n
\n
\n
\n   dummy_content2\n
\n
\
\n'
        self.assertEqual(expected, actual)

    @patch('email_parser.fs.read_file')
    def test_inline_styles(self, mock_read):
        html = '{{content}}'
        html_placeholders = ['content']
        placeholders = {'content': 'dummy_content'}
        mock_read.side_effect = iter(['p {color:red;}'])

        renderer = self._get_renderer(html, html_placeholders)
        actual = renderer.render(placeholders)

        self.assertEqual('
dummy_content
', actual)

    def test_no_tracking(self):
        placeholders = {'content': Placeholder('content', '[link_title](!http://link.com)')}
        template = Template('dummy', '', '{{content}}', ['content'])
        r = renderer.HtmlRenderer(template, self.email)
        expected = """

      link_title

"""

        actual = r.render(placeholders)
        self.assertEqual(expected, actual)
<commit_msg>fix tests<commit_after>from unittest import TestCase
from unittest.mock import patch

from email_parser import renderer, errors, cmd
from email_parser.reader import Template
from email_parser.fs import Email

class TestTextRenderer(TestCase):
    def setUp(self):
        self.email = Email('name', 'locale', 'path', 'full_path')
        self.renderer = renderer.TextRenderer([], {}, self.email)

    def test_happy_path(self):
        placeholders = {'content': 'dummy content'}

        actual = self.renderer.render(placeholders)

        self.assertEqual('dummy content', actual)

    def test_concat_multiple_placeholders(self):
        placeholders = {'content1': 'dummy content', 'content2': 'dummy content'}
        expected = renderer.TEXT_EMAIL_PLACEHOLDER_SEPARATOR.join(placeholders.values())

        actual = self.renderer.render(placeholders)

        self.assertEqual(expected, actual)

    def test_ignore_subject(self):
        placeholders = {'content': 'dummy content', 'subject': 'dummy subject'}

        actual = self.renderer.render(placeholders)

        self.assertEqual('dummy content', actual)

    def test_ignore_empty_placeholders(self):
        placeholders = {'content': 'dummy content', 'empty': ''}

        actual = self.renderer.render(placeholders)

        self.assertEqual('dummy content', actual)

    def test_ignored_placeholders(self):
        placeholders = {'content': 'dummy content', 'ignore': 'test'}

        r = renderer.TextRenderer(['ignore'], {}, self.email)

        actual = r.render(placeholders)

        self.assertEqual('dummy content', actual)

    def test_use_text_and_url_for_links(self):
        placeholders = {'content': 'dummy [link_text](http://link_url) content'}

        actual = self.renderer.render(placeholders)

        self.assertEqual('dummy link_text (http://link_url) content', actual)

    def test_default_link_locale_for_links(self):
        placeholders = {'content': 'dummy [link_text](http://link_url?locale={link_locale}) content'}

        actual = self.renderer.render(placeholders)

        self.assertEqual('dummy link_text (http://link_url?locale=en) content', actual)

    def test_link_locale_for_links(self):
        placeholders = {'content': 'dummy [link_text](http://link_url?locale={link_locale}) content'}

        r = renderer.TextRenderer(['ignore'], {'locale': 'zh'}, self.email)
        actual = r.render(placeholders)

        self.assertEqual('dummy link_text (http://link_url?locale=zh) content', actual)

    def test_use_text_if_href_is_empty(self):
        placeholders = {'content': 'dummy [http://link_url]() content'}

        actual = self.renderer.render(placeholders)

        self.assertEqual('dummy http://link_url content', actual)

    def test_use_href_if_text_is_same(self):
        placeholders = {'content': 'dummy [http://link_url](http://link_url) content'}

        actual = self.renderer.render(placeholders)

        self.assertEqual('dummy http://link_url content', actual)

    def test_url_with_params(self):
        placeholders = {'content': 'dummy [param_link](https://something.com/thing?id=mooo) content'}

        actual = self.renderer.render(placeholders)

        self.assertEqual('dummy param_link (https://something.com/thing?id=mooo) content', actual)

    def test_unordered_list(self):
        placeholders = {'content': '- one\n- two\n- three'}

        actual = self.renderer.render(placeholders)

        self.assertEqual('- one\n- two\n- three', actual.strip())

    def test_ordered_list(self):
        placeholders = {'content': '1. one\n2. two\n3. three'}

        actual = self.renderer.render(placeholders)

        self.assertEqual('1. one\n2. two\n3. three', actual.strip())

class TestSubjectRenderer(TestCase):
    def setUp(self):
        self.renderer = renderer.SubjectRenderer()

    def test_happy_path(self):
        placeholders = {'content': 'dummy content', 'subject': 'dummy subject'}

        actual = self.renderer.render(placeholders)

        self.assertEqual('dummy subject', actual[0])

    def test_ab(self):
        placeholders = {'content': 'dummy content', 'subject': 'dummy subject', 'subject_b': 'bbb', 'subject_a': 'aaa'}

        actual = self.renderer.render(placeholders)

        self.assertEqual(['dummy subject', 'aaa', 'bbb', None], actual)

    def test_raise_error_for_missing_subject(self):
        placeholders = {'content': 'dummy content'}

        with self.assertRaises(errors.MissingSubjectError):
            self.renderer.render(placeholders)

class TestHtmlRenderer(TestCase):
    def _get_renderer(self, template_html, template_placeholders, **kwargs):
        template = Template(
            name='template_name',
            styles=['template_style'],
            content=template_html,
            placeholders_order=template_placeholders)
        return renderer.HtmlRenderer(template,
                                     kwargs.get('link_locale_mappings', {}),
                                     kwargs.get('email', self.email), kwargs.get('settings', self.settings))

    def setUp(self):
        settings = cmd.default_settings()._asdict()
        settings['templates'] = 'dummy_templates'
        settings['images'] = 'dummy_images'
        self.settings = cmd.Settings(**settings)

        self.email = Email('name', 'locale', 'path', 'full_path')
        self.global_email = '<?xml version="1.0" encoding="UTF-8" ?><resources></resources>'

    @patch('email_parser.fs.read_file')
    def test_happy_path(self, mock_read):
        html = '{{content1}}'
        html_placeholders = ['content1']
        placeholders = {'content1': 'text1'}
        mock_read.side_effect = iter(['body {}'])

        renderer = self._get_renderer(html, html_placeholders)
        actual = renderer.render(placeholders)

        self.assertEqual('
text1
', actual)

    @patch('email_parser.fs.read_file')
    def test_empty_style(self, mock_read):
        html = '{{content}}'
        html_placeholders = ['content']
        placeholders = {'content': 'dummy_content'}
        mock_read.side_effect = iter([''])

        renderer = self._get_renderer(html, html_placeholders)
        actual = renderer.render(placeholders)

        self.assertEqual('
dummy_content
', actual)

    @patch('email_parser.fs.read_file')
    def test_include_raw_subject(self, mock_read):
        html = '{{subject}}'
        html_placeholders = ['subject']
        placeholders = {'subject': 'dummy_subject'}
        mock_read.side_effect = iter([''])

        renderer = self._get_renderer(html, html_placeholders)
        actual = renderer.render(placeholders)

        self.assertEqual('dummy_subject', actual)

    @patch('email_parser.fs.read_file')
    def test_include_base_url(self, mock_read):
        html = '{{base_url}}'
        html_placeholders = ['base_url']
        placeholders = {}
        mock_read.side_effect = iter([''])

        renderer = self._get_renderer(html, html_placeholders)
        actual = renderer.render(placeholders)

        self.assertEqual('dummy_images', actual)

    @patch('email_parser.fs.read_file')
    def test_ignore_missing_placeholders(self, mock_read):
        html = '{{content}}{{missing}}'
        html_placeholders = ['content', 'missing']
        placeholders = {'content': 'dummy_content'}
        mock_read.side_effect = iter([''])
        settings = self.settings._asdict()
        settings['strict'] = False
        settings = cmd.Settings(**settings)

        renderer = self._get_renderer(html, html_placeholders, settings=settings)
        actual = renderer.render(placeholders)

        self.assertEqual('
dummy_content
', actual)

    @patch('email_parser.fs.read_file')
    def test_fail_on_missing_placeholders(self, mock_read):
        html = '{{content}}{{missing}}'
        html_placeholders = ['content', 'missing']
        placeholders = {'content': 'dummy_content'}
        mock_read.side_effect = iter([''])

        renderer = self._get_renderer(html, html_placeholders)
        with self.assertRaises(errors.MissingTemplatePlaceholderError):
            renderer.render(placeholders)

    @patch('email_parser.fs.read_file')
    def test_rtl_locale(self, mock_read):
        html = '{{content}}'
        html_placeholders = ['content']
        placeholders = {'content': 'dummy_content'}
        mock_read.side_effect = iter([''])
        email_dict = self.email._asdict()
        email_dict['locale'] = 'ar'

        renderer = self._get_renderer(html, html_placeholders, email=Email(**email_dict))
        actual = renderer.render(placeholders)

        self.assertEqual('\n
\n  dummy_content\n
\n', actual)

    @patch('email_parser.fs.read_file')
    def test_rtl_two_placeholders(self, mock_read):
        html = '
{{content1}}

{{content2}}
'
        html_placeholders = ['content1', 'content2']
        placeholders = {'content1': 'dummy_content1', 'content2': 'dummy_content2'}
        mock_read.side_effect = iter([''])
        email_dict = self.email._asdict()
        email_dict['locale'] = 'ar'

        renderer = self._get_renderer(html, html_placeholders, email=Email(**email_dict))
        actual = renderer.render(placeholders)
        expected = '\n
\n
\n   dummy_content1\n
\n
\n
\n
\n   dummy_content2\n
\n
\
\n'

        self.assertEqual(expected, actual)

    @patch('email_parser.fs.read_file')
    def test_inline_styles(self, mock_read):
        html = '{{content}}'
        html_placeholders = ['content']
        placeholders = {'content': 'dummy_content'}
        mock_read.side_effect = iter(['p {color:red;}'])

        renderer = self._get_renderer(html, html_placeholders)
        actual = renderer.render(placeholders)

        self.assertEqual('
dummy_content
', actual)

    @patch('email_parser.fs.read_file')
    def test_no_tracking(self, mock_read):
        html = '{{content}}'
        html_placeholders = ['content']
        placeholders = {'content': '[link_title](!http://link.com)'}
        mock_read.side_effect = iter([''])
        expected = """

      link_title

"""

        r = self._get_renderer(html, html_placeholders)
        actual = r.render(placeholders)

        self.assertEqual(expected, actual)
<|endoftext|>


<commit_before>import struct
from datetime import datetime, time, timedelta
from twisted.python import log
from twisted.internet.protocol import DatagramProtocol
from sqlalchemy import distinct
from sqlalchemy.exc import SQLAlchemyError
from sqlalchemy.sql.expression import and_, or_, desc
import transaction
from skylines.model import DBSession, User, TrackingFix, Location, Follower
from skylines.tracking.crc import check_crc, set_crc

# More information about this protocol can be found in the XCSoar
# source code, source file src/Tracking/SkyLines/Protocol.hpp

MAGIC = 0x5df4b67b
TYPE_PING = 1
TYPE_ACK = 2
TYPE_FIX = 3
TYPE_TRAFFIC_REQUEST = 4
TYPE_TRAFFIC_RESPONSE = 5
TYPE_USER_NAME_REQUEST = 6
TYPE_USER_NAME_RESPONSE = 7

FLAG_ACK_BAD_KEY = 0x1

FLAG_LOCATION = 0x1
FLAG_TRACK = 0x2
FLAG_GROUND_SPEED = 0x4
FLAG_AIRSPEED = 0x8
FLAG_ALTITUDE = 0x10
FLAG_VARIO = 0x20
FLAG_ENL = 0x40

# for TYPE_TRAFFIC_REQUEST
TRAFFIC_FLAG_FOLLOWEES = 0x1
TRAFFIC_FLAG_CLUB = 0x2

USER_FLAG_NOT_FOUND = 0x1

class TrackingServer(DatagramProtocol):
    def pingReceived(self, host, port, key, payload):
        if len(payload) != 8: return
        id, reserved, reserved2 = struct.unpack('!HHI', payload)

        flags = 0

        pilot = User.by_tracking_key(key)
        if not pilot:
            flags |= FLAG_ACK_BAD_KEY

        data = struct.pack('!IHHQHHI', MAGIC, 0, TYPE_ACK, 0,
                           id, 0, flags)
        data = set_crc(data)
        self.transport.write(data, (host, port))

    def fixReceived(self, host, key, payload):
        if len(payload) != 32: return

        pilot = User.by_tracking_key(key)
        if not pilot:
            log.err("No such pilot: %x" % key)
            return

        data = struct.unpack('!IIiiIHHHhhH', payload)

        fix = TrackingFix()
        fix.ip = host
        fix.pilot = pilot

        # import the time stamp from the packet if it's within a
        # certain range
        time_of_day_ms = data[1] % (24 * 3600 * 1000)
        time_of_day_s = time_of_day_ms / 1000
        time_of_day = time(time_of_day_s / 3600,
                           (time_of_day_s / 60) % 60,
                           time_of_day_s % 60,
                           (time_of_day_ms % 1000) * 1000)
        now = datetime.utcnow()
        now_s = ((now.hour * 60) + now.minute) * 60 + now.second
        if now_s - 1800 < time_of_day_s < now_s + 180:
            fix.time = datetime.combine(now.date(), time_of_day)
        elif now_s < 1800 and time_of_day_s > 23 * 3600:
            # midnight rollover occurred
            fix.time = (datetime.combine(now.date(), time_of_day) -
                        timedelta(days=1))
        else:
            log.msg("ignoring time stamp from FIX packet: " + str(time_of_day))

        flags = data[0]
        if flags & FLAG_LOCATION:
            fix.location = Location(latitude=data[2] / 1000000.,
                                    longitude=data[3] / 1000000.)

        if flags & FLAG_TRACK:
            fix.track = data[5]

        if flags & FLAG_GROUND_SPEED:
            fix.ground_speed = data[6] / 16.

        if flags & FLAG_AIRSPEED:
            fix.airspeed = data[7] / 16.

        if flags & FLAG_ALTITUDE:
            fix.altitude = data[8]

        if flags & FLAG_VARIO:
            fix.vario = data[9] / 256.

        if flags & FLAG_ENL:
            fix.engine_noise_level = data[10]

        log.msg("{} {} {} {}".format(
            fix.time and fix.time.time(), host,
            unicode(pilot).encode('utf8', 'ignore'), fix.location))

        DBSession.add(fix)
        try:
            transaction.commit()
        except SQLAlchemyError, e:
            log.err(e, 'database error')
            transaction.abort()

    def trafficRequestReceived(self, host, port, key, payload):
        if len(payload) != 8: return

        pilot = User.by_tracking_key(key)
        if pilot is None:
            log.err("No such pilot: %d" % key)
            return

        data = struct.unpack('!II', payload)
        or_filters = []

        flags = data[0]
        if flags & TRAFFIC_FLAG_FOLLOWEES:
            subq = DBSession \
                .query(Follower.destination_id) \
                .filter(Follower.source_id == pilot.id) \
                .subquery()

            or_filters.append(TrackingFix.pilot_id.in_(subq))

        if flags & TRAFFIC_FLAG_CLUB:
            subq = DBSession \
                .query(User.id) \
                .filter(User.club_id == pilot.club_id) \
                .subquery()

            or_filters.append(TrackingFix.pilot_id.in_(subq))

        if len(or_filters) == 0:
            return

        query = DBSession.query(TrackingFix) \
            .distinct(TrackingFix.pilot_id) \
            .filter(and_(TrackingFix.time >= datetime.utcnow() - timedelta(hours=2),
                         TrackingFix.pilot_id != pilot.id,
                         TrackingFix.location_wkt != None,
                         TrackingFix.altitude != None,
                         or_(*or_filters))) \
            .order_by(TrackingFix.pilot_id, desc(TrackingFix.time)) \
            .limit(32)

        response = ''
        count = 0
        for fix in query:
            location = fix.location
            if location is None: continue

            t = fix.time
            t = t.hour * 3600000 + t.minute * 60000 + t.second * 1000 + t.microsecond / 1000
            response += struct.pack('!IIiihHI', fix.pilot_id, t,
                                    int(location.latitude * 1000000),
                                    int(location.longitude * 1000000),
                                    int(fix.altitude), 0, 0)
            count += 1

        response = struct.pack('!HBBI', 0, 0, count, 0) + response
        response = struct.pack('!IHHQ', MAGIC, 0, TYPE_TRAFFIC_RESPONSE, 0) + response
        response = set_crc(response)
        self.transport.write(response, (host, port))

    def userNameRequestReceived(self, host, port, key, payload):
        """The client asks for the display name of a user account."""

        if len(payload) != 8: return

        pilot = User.by_tracking_key(key)
        if pilot is None:
            log.err("No such pilot: %d" % key)
            return

        data = struct.unpack('!II', payload)
        user_id = data[0]

        user = DBSession.query(User).get(user_id)
        if user is None:
            response = struct.pack('!IHHQIIIBBBBII', MAGIC, 0, TYPE_USER_NAME_RESPONSE, 0,
                                   user_id, USER_FLAG_NOT_FOUND, 0,
                                   0, 0, 0, 0, 0, 0)
            response = set_crc(response)
            self.transport.write(response, (host, port))
            return

        name = user.display_name[:64].encode('utf8', 'ignore')
        club_id = user.club_id or 0

        response = struct.pack('!IHHQIIIBBBBII', MAGIC, 0, TYPE_USER_NAME_RESPONSE, 0,
                               user_id, 0, club_id,
                               len(name), 0, 0, 0, 0, 0)
        response += name
        response = set_crc(response)
        self.transport.write(response, (host, port))

    def datagramReceived(self, data, (host, port)):
        if len(data) < 16: return

        header = struct.unpack('!IHHQ', data[:16])
        if header[0] != MAGIC: return
        if not check_crc(data): return

        if header[2] == TYPE_FIX:
            self.fixReceived(host, header[3], data[16:])
        elif header[2] == TYPE_PING:
            self.pingReceived(host, port, header[3], data[16:])
        elif header[2] == TYPE_TRAFFIC_REQUEST:
            self.trafficRequestReceived(host, port, header[3], data[16:])
        elif header[2] == TYPE_TRAFFIC_REQUEST:
            self.trafficRequestReceived(host, port, header[3], data[16:])
        elif header[2] == TYPE_USER_NAME_REQUEST:
            self.userNameRequestReceived(host, port, header[3], data[16:])
<commit_msg>tracking/server: Import cleanup<commit_after>import struct
from datetime import datetime, time, timedelta
from twisted.python import log
from twisted.internet.protocol import DatagramProtocol
from sqlalchemy.exc import SQLAlchemyError
from sqlalchemy.sql.expression import and_, or_, desc
import transaction
from skylines.model import DBSession, User, TrackingFix, Location, Follower
from skylines.tracking.crc import check_crc, set_crc

# More information about this protocol can be found in the XCSoar
# source code, source file src/Tracking/SkyLines/Protocol.hpp

MAGIC = 0x5df4b67b
TYPE_PING = 1
TYPE_ACK = 2
TYPE_FIX = 3
TYPE_TRAFFIC_REQUEST = 4
TYPE_TRAFFIC_RESPONSE = 5
TYPE_USER_NAME_REQUEST = 6
TYPE_USER_NAME_RESPONSE = 7

FLAG_ACK_BAD_KEY = 0x1

FLAG_LOCATION = 0x1
FLAG_TRACK = 0x2
FLAG_GROUND_SPEED = 0x4
FLAG_AIRSPEED = 0x8
FLAG_ALTITUDE = 0x10
FLAG_VARIO = 0x20
FLAG_ENL = 0x40

# for TYPE_TRAFFIC_REQUEST
TRAFFIC_FLAG_FOLLOWEES = 0x1
TRAFFIC_FLAG_CLUB = 0x2

USER_FLAG_NOT_FOUND = 0x1

class TrackingServer(DatagramProtocol):
    def pingReceived(self, host, port, key, payload):
        if len(payload) != 8: return
        id, reserved, reserved2 = struct.unpack('!HHI', payload)

        flags = 0

        pilot = User.by_tracking_key(key)
        if not pilot:
            flags |= FLAG_ACK_BAD_KEY

        data = struct.pack('!IHHQHHI', MAGIC, 0, TYPE_ACK, 0,
                           id, 0, flags)
        data = set_crc(data)
        self.transport.write(data, (host, port))

    def fixReceived(self, host, key, payload):
        if len(payload) != 32: return

        pilot = User.by_tracking_key(key)
        if not pilot:
            log.err("No such pilot: %x" % key)
            return

        data = struct.unpack('!IIiiIHHHhhH', payload)

        fix = TrackingFix()
        fix.ip = host
        fix.pilot = pilot

        # import the time stamp from the packet if it's within a
        # certain range
        time_of_day_ms = data[1] % (24 * 3600 * 1000)
        time_of_day_s = time_of_day_ms / 1000
        time_of_day = time(time_of_day_s / 3600,
                           (time_of_day_s / 60) % 60,
                           time_of_day_s % 60,
                           (time_of_day_ms % 1000) * 1000)
        now = datetime.utcnow()
        now_s = ((now.hour * 60) + now.minute) * 60 + now.second
        if now_s - 1800 < time_of_day_s < now_s + 180:
            fix.time = datetime.combine(now.date(), time_of_day)
        elif now_s < 1800 and time_of_day_s > 23 * 3600:
            # midnight rollover occurred
            fix.time = (datetime.combine(now.date(), time_of_day) -
                        timedelta(days=1))
        else:
            log.msg("ignoring time stamp from FIX packet: " + str(time_of_day))

        flags = data[0]
        if flags & FLAG_LOCATION:
            fix.location = Location(latitude=data[2] / 1000000.,
                                    longitude=data[3] / 1000000.)

        if flags & FLAG_TRACK:
            fix.track = data[5]

        if flags & FLAG_GROUND_SPEED:
            fix.ground_speed = data[6] / 16.

        if flags & FLAG_AIRSPEED:
            fix.airspeed = data[7] / 16.

        if flags & FLAG_ALTITUDE:
            fix.altitude = data[8]

        if flags & FLAG_VARIO:
            fix.vario = data[9] / 256.

        if flags & FLAG_ENL:
            fix.engine_noise_level = data[10]

        log.msg("{} {} {} {}".format(
            fix.time and fix.time.time(), host,
            unicode(pilot).encode('utf8', 'ignore'), fix.location))

        DBSession.add(fix)
        try:
            transaction.commit()
        except SQLAlchemyError, e:
            log.err(e, 'database error')
            transaction.abort()

    def trafficRequestReceived(self, host, port, key, payload):
        if len(payload) != 8: return

        pilot = User.by_tracking_key(key)
        if pilot is None:
            log.err("No such pilot: %d" % key)
            return

        data = struct.unpack('!II', payload)
        or_filters = []

        flags = data[0]
        if flags & TRAFFIC_FLAG_FOLLOWEES:
            subq = DBSession \
                .query(Follower.destination_id) \
                .filter(Follower.source_id == pilot.id) \
                .subquery()

            or_filters.append(TrackingFix.pilot_id.in_(subq))

        if flags & TRAFFIC_FLAG_CLUB:
            subq = DBSession \
                .query(User.id) \
                .filter(User.club_id == pilot.club_id) \
                .subquery()

            or_filters.append(TrackingFix.pilot_id.in_(subq))

        if len(or_filters) == 0:
            return

        query = DBSession.query(TrackingFix) \
            .distinct(TrackingFix.pilot_id) \
            .filter(and_(TrackingFix.time >= datetime.utcnow() - timedelta(hours=2),
                         TrackingFix.pilot_id != pilot.id,
                         TrackingFix.location_wkt != None,
                         TrackingFix.altitude != None,
                         or_(*or_filters))) \
            .order_by(TrackingFix.pilot_id, desc(TrackingFix.time)) \
            .limit(32)

        response = ''
        count = 0
        for fix in query:
            location = fix.location
            if location is None: continue

            t = fix.time
            t = t.hour * 3600000 + t.minute * 60000 + t.second * 1000 + t.microsecond / 1000
            response += struct.pack('!IIiihHI', fix.pilot_id, t,
                                    int(location.latitude * 1000000),
                                    int(location.longitude * 1000000),
                                    int(fix.altitude), 0, 0)
            count += 1

        response = struct.pack('!HBBI', 0, 0, count, 0) + response
        response = struct.pack('!IHHQ', MAGIC, 0, TYPE_TRAFFIC_RESPONSE, 0) + response
        response = set_crc(response)
        self.transport.write(response, (host, port))

    def userNameRequestReceived(self, host, port, key, payload):
        """The client asks for the display name of a user account."""

        if len(payload) != 8: return

        pilot = User.by_tracking_key(key)
        if pilot is None:
            log.err("No such pilot: %d" % key)
            return

        data = struct.unpack('!II', payload)
        user_id = data[0]

        user = DBSession.query(User).get(user_id)
        if user is None:
            response = struct.pack('!IHHQIIIBBBBII', MAGIC, 0, TYPE_USER_NAME_RESPONSE, 0,
                                   user_id, USER_FLAG_NOT_FOUND, 0,
                                   0, 0, 0, 0, 0, 0)
            response = set_crc(response)
            self.transport.write(response, (host, port))
            return

        name = user.display_name[:64].encode('utf8', 'ignore')
        club_id = user.club_id or 0

        response = struct.pack('!IHHQIIIBBBBII', MAGIC, 0, TYPE_USER_NAME_RESPONSE, 0,
                               user_id, 0, club_id,
                               len(name), 0, 0, 0, 0, 0)
        response += name
        response = set_crc(response)
        self.transport.write(response, (host, port))

    def datagramReceived(self, data, (host, port)):
        if len(data) < 16: return

        header = struct.unpack('!IHHQ', data[:16])
        if header[0] != MAGIC: return
        if not check_crc(data): return

        if header[2] == TYPE_FIX:
            self.fixReceived(host, header[3], data[16:])
        elif header[2] == TYPE_PING:
            self.pingReceived(host, port, header[3], data[16:])
        elif header[2] == TYPE_TRAFFIC_REQUEST:
            self.trafficRequestReceived(host, port, header[3], data[16:])
        elif header[2] == TYPE_TRAFFIC_REQUEST:
            self.trafficRequestReceived(host, port, header[3], data[16:])
        elif header[2] == TYPE_USER_NAME_REQUEST:
            self.userNameRequestReceived(host, port, header[3], data[16:])
<|endoftext|>


<commit_before>#!/usr/bin/env python3

import pygame as pg
import sys
from game.Gameengine import Gameengine
from gui.Sidemenu import Sidemenu
# import level?
# import wave?
# non-steering bullets ?

framerate = 60
path_data   = "../res/path_points.txt" # dann in map oder level je nachdem ...
# wave_data   = "../res/.txt" # Lars' Wavemanager!

if __name__ == "__main__":
    pg.init()
    pg.display.set_caption("TowerDefense_KIT_PY'17")
    window_size = (1280,720)
    #main_surface = pg.display.set_mode(window_size, pg.FULLSCREEN)
    main_surface = pg.display.set_mode(window_size)
    clock = pg.time.Clock()

    gameengine = Gameengine( path_data,(window_size[0]-100,window_size[1]))

    sidemenu = Sidemenu(window_size[0]-100,0,100,window_size[1], pg.Color(128,128,128), gameengine)

    running = True
    doublespeed = False

    while True:
        deltatime = clock.tick(framerate)
        if gameengine.health <= 0 or not running:
            deltatime = 0
        if doublespeed:
            deltatime *=2

        for e in pg.event.get():
            if e.type == pg.QUIT:
                sys.exit()
            elif e.type == pg.KEYDOWN:
                if e.key == pg.K_q:
                    sys.exit()
                elif e.key == pg.K_ESCAPE:
                    gameengine.idle_tower=None
                elif e.key == pg.K_SPACE:
                    running = not running
                elif e.key == pg.K_PLUS or e.unicode == '+' or e.key == pg.K_h:
                    doublespeed = not doublespeed

            elif e.type == pg.MOUSEBUTTONDOWN:
                #button.check_press(pg.mouse.get_pos())
                #gameengine.towers[-1].place_down()
                #gameengine.add_tower("turret")
                #gameengine.add_enemy("normal")
                sidemenu.check_press(e.pos)
                gameengine.placeIdleTower()
            elif e.type == pg.MOUSEBUTTONUP:
                sidemenu.release()

        # all the updating
        gameengine.update(deltatime)

        # all the drawing
        gameengine.draw(main_surface)
        sidemenu.draw(main_surface)

#        # TEST
#        test_img = pg.image.load(enemy_path)
#        test_img = pg.transform.scale(test_img,(100,100))
#        render_img = pg.transform.rotate(test_img,pg.mouse.get_pos()[0])
#        main_surface.blit(render_img,(150,150))
#        # TEST END

        pg.display.flip() # Update the full display Surface to the screen \\ Update portions of the screen for software displays
<commit_msg>mousecontrol added<commit_after>#!/usr/bin/env python3

import pygame as pg
import sys
from game.Gameengine import Gameengine
from gui.Sidemenu import Sidemenu
# import level?
# import wave?
# non-steering bullets ?

framerate = 60
path_data   = "../res/path_points.txt" # dann in map oder level je nachdem ...
# wave_data   = "../res/.txt" # Lars' Wavemanager!

if __name__ == "__main__":
    pg.init()
    pg.display.set_caption("TowerDefense_KIT_PY'17")
    window_size = (1280,720)
    #main_surface = pg.display.set_mode(window_size, pg.FULLSCREEN)
    main_surface = pg.display.set_mode(window_size)
    clock = pg.time.Clock()

    gameengine = Gameengine( path_data,(window_size[0]-100,window_size[1]))

    sidemenu = Sidemenu(window_size[0]-100,0,100,window_size[1], pg.Color(128,128,128), gameengine)

    running = True
    doublespeed = False

    while True:
        deltatime = clock.tick(framerate)
        if gameengine.health <= 0 or not running:
            deltatime = 0
        if doublespeed:
            deltatime *=2

        for e in pg.event.get():
            if e.type == pg.QUIT:
                sys.exit()
            elif e.type == pg.KEYDOWN:
                if e.key == pg.K_q:
                    sys.exit()
                elif e.key == pg.K_ESCAPE:
                    gameengine.idle_tower=None
                elif e.key == pg.K_SPACE:
                    running = not running
                elif e.key == pg.K_PLUS or e.unicode == '+' or e.key == pg.K_h:
                    doublespeed = not doublespeed

            elif e.type == pg.MOUSEBUTTONDOWN:
                if e.button == 1:
                #button.check_press(pg.mouse.get_pos())
                #gameengine.towers[-1].place_down()
                #gameengine.add_tower("turret")
                #gameengine.add_enemy("normal")
                    sidemenu.check_press(e.pos)
                    gameengine.placeIdleTower()
                if e.button == 3:
                    gameengine.idle_tower=None
            elif e.type == pg.MOUSEBUTTONUP:
                sidemenu.release()

        # all the updating
        gameengine.update(deltatime)

        # all the drawing
        gameengine.draw(main_surface)
        sidemenu.draw(main_surface)

#        # TEST
#        test_img = pg.image.load(enemy_path)
#        test_img = pg.transform.scale(test_img,(100,100))
#        render_img = pg.transform.rotate(test_img,pg.mouse.get_pos()[0])
#        main_surface.blit(render_img,(150,150))
#        # TEST END

        pg.display.flip() # Update the full display Surface to the screen \\ Update portions of the screen for software displays
<|endoftext|>