<commit_before>from __future__ import print_function, division
import numpy as np
import astropy.io.fits as pyfits
from astropy.time import Time
from astropy import units as u
from astropy.coordinates import SkyCoord
from pkg_resources import resource_filename
from operator import itemgetter
from desisurvey.utils import radec2altaz, mjd2lst, equ2gal_J2000, sort2arr, inLSTwindow
from desitarget.targetmask import obsconditions as obsbits
from desisurvey.exposurecalc import airMassCalculator
import copy
import os
import os.path

import warnings
warnings.simplefilter('error', RuntimeWarning)

class surveyPlan:
    """
    Main class for survey planning
    """

    def __init__(self, MJDstart, MJDend, surveycal, tilesubset=None):
        """Initialises survey by reading in the file desi_tiles.fits
        and populates the class members.

        Arguments:
            MJDstart: day of the (re-)start of the survey
            MJDend: day of the end of the survey
            surveycal: list of dictionnaries with times of sunset, twilight, etc

        Optional:
            tilesubset: array of integer tileids to use; ignore others
        """

        self.surveycal = surveycal
        # Read in DESI tile data
        # Columns are:
        #   'TILEID'; format = 'J'
        #   'RA'; format = 'D'
        #   'DEC'; format = 'D'
        #   'PASS'; format = 'I'
        #   'IN_DESI'; format = 'I'
        #   'EBV_MED'; format = 'E'
        #   'AIRMASS'; format = 'E'
        #   'STAR_DENSITY'; format = 'E'
        #   'EXPOSEFAC'; format = 'E'
        #   'PROGRAM'; format = '6A'
        #   'OBSCONDITIONS'; format = 'J'
        try:
            tiles_path = os.path.join(os.environ['DESIMODEL'], 'data', 'footprint', 'desi-tiles.fits')
        except KeyError:
            raise RuntimeError('DESIMODEL environment variable is not defined.')
        hdulist0 = pyfits.open(tiles_path)
        tiledata0 = hdulist0[1].data
        # This works because original table has index = tileID.
        if tilesubset is not None:
            tiledata1 = tiledata0[tilesubset]
        else:
            tiledata1 = tiledata0
        # Only keep DESI tiles
        in_desi = np.where(tiledata1.field('IN_DESI')==1)
        tb_temp = tiledata1[in_desi]
        # Ignore EXTRA tiles (passes 8 and 9)
        no_extra = np.where(tb_temp.field('PROGRAM') != 'EXTRA')
        tiledata = tb_temp[no_extra]
        hdulist0.close()

        # Dummy initilisations
        priority = 0
        status = 0
        ha = 0.0
        lstmin = 0.0
        lstmax = 0.0
        explen = 0.0

        # Loop over elements in table
        self.numtiles = len(tiledata)
        self.tiles = np.recarray((self.numtiles,),
                                 names = ('TILEID', 'RA', 'DEC', 'PASS', 'EBV_MED', 'PROGRAM', 'OBSCONDITIONS', 'GAL_CAP', 'SUBLIST', 'PRIORITY', 'STATUS', 'HA', 'LSTMIN', 'LSTMAX', 'EXPLEN'),
                                 formats = ['i4', 'f8', 'f8', 'i4', 'f4', 'a6', 'i2', 'i4', 'i4', 'i4', 'i4', 'f8', 'f8', 'f8', 'f8'])
        for i in range(len(tiledata)):
            #l, b = equ2gal_J2000(tiledata[i].field('RA'), tiledata[i].field('DEC'))
            #cap = b / np.abs(b)
            if tiledata[i].field('RA') > 75.0 and tiledata[i].field('RA') < 300.0:
                cap = 1.0
            else:
                cap = -1.0
            if tiledata[i].field('PASS') == 0:
                if self.inFirstYearFullDepthField(tiledata[i].field('DEC'), cap, True):
                    sublist = 0
                else:
                    sublist = 8
            elif tiledata[i].field('PASS') == 1:
                if self.inFirstYearFullDepthField(tiledata[i].field('DEC'), cap, False):
                    sublist = 1
                else:
                    sublist = 9
            elif tiledata[i].field('PASS') == 2:
                if self.inFirstYearFullDepthField(tiledata[i].field('DEC'), cap, False):
                    sublist = 2
                else:
                    sublist = 10
            elif tiledata[i].field('PASS') == 3:
                if self.inFirstYearFullDepthField(tiledata[i].field('DEC'), cap, False):
                    sublist = 3
                else:
                    sublist = 11
            elif tiledata[i].field('PASS') == 4:
                if self.inFirstYearFullDepthField(tiledata[i].field('DEC'), cap, True):
                    sublist = 4
                else:
                    sublist = 12
            elif tiledata[i].field('PASS') == 5:
                if self.inFirstYearFullDepthField(tiledata[i].field('DEC'), cap, True):
                    sublist = 5
                else:
                    sublist = 13
            elif tiledata[i].field('PASS') == 6:
                if self.inFirstYearFullDepthField(tiledata[i].field('DEC'), cap, False):
                    sublist = 6
                else:
                    sublist = 14
            elif tiledata[i].field('PASS') == 7:
                if self.inFirstYearFullDepthField(tiledata[i].field('DEC'), cap, False):
                    sublist = 7
                else:
                    sublist = 15
            self.tiles[i] = (tiledata[i].field('TILEID'),
                              tiledata[i].field('RA'),
                              tiledata[i].field('DEC'),
                              tiledata[i].field('PASS'),
                              tiledata[i].field('EBV_MED'),
                              tiledata[i].field('PROGRAM'),
                              tiledata[i].field('OBSCONDITIONS'),
                              cap,
                              sublist,
                              priority,
                              status,
                              ha,
                              lstmin,
                              lstmax,
                              explen)

        self.LSTres = 2.5 # bin width in degrees, starting with 10 minutes for now
        self.nLST = int(np.floor(360.0/self.LSTres))
        self.LSTbins = np.zeros(self.nLST)
        for i in range(self.nLST):
            self.LSTbins[i] = (float(i) + 0.5) * self.LSTres

        self.assignHA(MJDstart, MJDend)
        self.tiles.sort(axis=0, order=('SUBLIST', 'DEC'))

    def assignHA(self, MJDstart, MJDend, compute=False):
        """Assigns optimal hour angles for the DESI tiles;
        can be re-run at any point during the survey to
        reoptimise the schedule.
        """

        if compute:
            f = open("ha_check.dat", 'w')
            obs_dark = self.plan_ha(MJDstart, MJDend, False)
            obs_bright = self.plan_ha(MJDstart, MJDend, True)
            for tile in self.tiles:
                if (tile['OBSCONDITIONS'] & obsbits.mask('GRAY')) != 0 or (tile['OBSCONDITIONS'] & obsbits.mask('DARK')) != 0:
                    j=0
                    while j < len(obs_dark):
                        if tile['TILEID'] == obs_dark['tileid'][j]:
                            tile['HA'] = obs_dark['ha'][j]
                            tile['LSTMIN'] = obs_dark['beginobs'][j]
                            tile['LSTMAX'] = obs_dark['endobs'][j]
                            tile['EXPLEN'] = obs_dark['obstime'][j]
                            break
                        else:
                            j += 1
                else:
                    i=0
                    while i < len(obs_bright):
                        if tile['TILEID'] == obs_bright['tileid'][i]:
                            tile['HA'] = obs_bright['ha'][i]
                            tile['LSTMIN'] = obs_bright['beginobs'][i]<commit_msg>Simplify ctor loop and speed up from 30s to <0.1s<commit_after>from __future__ import print_function, division
import numpy as np
import astropy.io.fits as pyfits
import astropy.table
from astropy.time import Time
from astropy import units as u
from astropy.coordinates import SkyCoord
from pkg_resources import resource_filename
from operator import itemgetter
from desisurvey.utils import radec2altaz, mjd2lst, equ2gal_J2000, sort2arr, inLSTwindow
from desitarget.targetmask import obsconditions as obsbits
from desisurvey.exposurecalc import airMassCalculator
import copy
import os
import os.path

import warnings
#warnings.simplefilter('error', RuntimeWarning)

class surveyPlan:
    """
    Main class for survey planning
    """

    def __init__(self, MJDstart, MJDend, surveycal, tilesubset=None):
        """Initialises survey by reading in the file desi_tiles.fits
        and populates the class members.

        Arguments:
            MJDstart: day of the (re-)start of the survey
            MJDend: day of the end of the survey
            surveycal: list of dictionnaries with times of sunset, twilight, etc

        Optional:
            tilesubset: array of integer tileids to use; ignore others
        """

        self.surveycal = surveycal
        # Read in DESI tile data
        # Columns are:
        #   'TILEID'; format = 'J'
        #   'RA'; format = 'D'
        #   'DEC'; format = 'D'
        #   'PASS'; format = 'I'
        #   'IN_DESI'; format = 'I'
        #   'EBV_MED'; format = 'E'
        #   'AIRMASS'; format = 'E'
        #   'STAR_DENSITY'; format = 'E'
        #   'EXPOSEFAC'; format = 'E'
        #   'PROGRAM'; format = '6A'
        #   'OBSCONDITIONS'; format = 'J'
        try:
            tiles_path = os.path.join(
                os.environ['DESIMODEL'], 'data', 'footprint', 'desi-tiles.fits')
        except KeyError:
            raise RuntimeError('DESIMODEL environment variable is not defined.')
        tiles = astropy.table.Table.read(tiles_path, hdu=1)

        # Only use tables with IN_DESI set.
        keep = tiles['IN_DESI'] > 0

        # Ignore EXTRA tiles.
        keep[tiles['PROGRAM'] == 'EXTRA'] = False

        # Restrict to a subset of tiles if requested.
        if tilesubset is not None:
            keep[~tilesubset] = False

        # Trim un-used rows.
        tiles.remove_rows(np.where(~keep)[0])
        numtiles = len(tiles)

        # Drop un-needed columns.
        tiles.remove_columns(['IN_DESI', 'AIRMASS', 'STAR_DENSITY', 'EXPOSEFAC'])

        # Add new columns.
        for name, dtype in (('GAL_CAP', np.int8), ('SUBLIST', np.int8),
                            ('PRIORITY', np.int32), ('STATUS', np.int32),
                            ('HA', np.float64), ('LSTMIN', np.float64),
                            ('LSTMAX', np.float64), ('EXPLEN', np.float64)):
            tiles[name] = np.zeros(numtiles, dtype=dtype)

        # Determine which galactic cap each tile is in: -1=south, +1=north.
        tiles['GAL_CAP'] = -1
        tiles['GAL_CAP'][(tiles['RA'] > 75) & (tiles['RA'] < 300)] = +1

        # Assign a sublist to each tile equal to pass for tiles in the
        # first-year full depth field, or else equal to pass+8.  The full
        # depth field covers 15 deg <= dec <= 25 deg in the NGC,
        # padded by 3 deg for the first pass in each program.
        dec = tiles['DEC']
        passnum = tiles['PASS']
        first_pass = (passnum == 0) | (passnum == 4) | (passnum == 5)
        dec_min = np.full(numtiles, 15.)
        dec_max = np.full(numtiles, 25.)
        dec_min[first_pass] -= 3.
        dec_max[first_pass] += 3.
        tiles['SUBLIST'] = passnum
        tiles['SUBLIST'][
            (tiles['GAL_CAP'] < 0) | (dec < dec_min) | (dec > dec_max)] += 8

        self.LSTres = 2.5 # bin width in degrees, starting with 10 minutes for now
        self.nLST = int(np.floor(360.0/self.LSTres))
        self.LSTbins = np.zeros(self.nLST)
        for i in range(self.nLST):
            self.LSTbins[i] = (float(i) + 0.5) * self.LSTres

        self.tiles = tiles
        self.numtiles = numtiles
        self.assignHA(MJDstart, MJDend)

        self.tiles.sort(('SUBLIST', 'DEC'))


    def assignHA(self, MJDstart, MJDend, compute=False):
        """Assigns optimal hour angles for the DESI tiles;
        can be re-run at any point during the survey to
        reoptimise the schedule.
        """

        if compute:
            f = open("ha_check.dat", 'w')
            obs_dark = self.plan_ha(MJDstart, MJDend, False)
            obs_bright = self.plan_ha(MJDstart, MJDend, True)
            for tile in self.tiles:
                if (tile['OBSCONDITIONS'] & obsbits.mask('GRAY')) != 0 or (tile['OBSCONDITIONS'] & obsbits.mask('DARK')) != 0:
                    j=0
                    while j < len(obs_dark):
                        if tile['TILEID'] == obs_dark['tileid'][j]:
                            tile['HA'] = obs_dark['ha'][j]
                            tile['LSTMIN'] = obs_dark['beginobs'][j]
                            tile['LSTMAX'] = obs_dark['endobs'][j]
                            tile['EXPLEN'] = obs_dark['obstime'][j]
                            break
                        else:
                            j += 1
                else:
                    i=0
                    while i < len(obs_bright):
                        if tile['TILEID'] == obs_bright['tileid'][i]:
                            tile['HA'] = obs_bright['ha'][i]
                            tile['LSTMIN'] = obs_bright['beginobs'][i]
{'fraction_non_alphanumeric': 0.1193621867881549, 'fraction_numerical': 0.013743356112376613, 'mean_word_length': 1.6554435483870968, 'pattern_counts': {'":': 0, '<': 14, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4015507', 'n_tokens_mistral': 4216, 'n_tokens_neox': 3994, 'n_words': 1000}