Example #1
0
    def records(self,
                category=None,
                region=None,
                fields=None,
                cuts=None,
                include_weight=True,
                systematic='NOMINAL',
                return_idx=False,
                **kwargs):
        if include_weight and fields is not None:
            if 'weight' not in fields:
                fields = list(fields) + ['weight']
            if 'o1' not in fields:
                fields = list(fields) + ['o1']
            if 'dphi' not in fields:
                fields = list(fields) + ['dphi']

        selection = self.cuts(category, region) & cuts

        log.info("requesting table from Data %d" % self.year)
        log.debug("using selection: %s" % selection)

        # read the table with a selection
        if selection:
            rec = self.h5data.read_where(selection.where(), **kwargs)
        else:
            rec = self.h5data.read(**kwargs)

        # add weight field
        if include_weight:
            # data is not weighted
            weights = np.ones(rec.shape[0], dtype='f8')
            rec = recfunctions.rec_append_fields(rec,
                names='weight',
                data=weights,
                dtypes='f8')
        o1=self.optimate(rec)
        if o1 is not None:
            rec =recfunctions.rec_append_fields(
                rec, names = 'o1', data=o1, dtypes='f8')
        dphi=self.dphate(rec)
        if dphi is not None:
            rec =recfunctions.rec_append_fields(
                rec, names = 'dphi', data=dphi, dtypes='f8')

        if fields is not None:
            rec = rec[fields]

        if return_idx:
            # only valid if selection is non-empty
            idx = self.h5data.get_where_list(selection.where(), **kwargs)
            return [(rec, idx)]

        return [rec]
def retain_relevant_fields(data):
    #TODO read these from a conf file
    aggregate_fields = {}
    aggregate_fields['aggr0_6'] = data['ika0'] + data['ika1'] + data[
        'ika2'] + data['ika3'] + data['ika4'] + data['ika5'] + data['ika6']
    aggregate_fields['aggr7_12'] = data['ika7'] + data['ika8'] + data[
        'ika9'] + data['ika10'] + data['ika11'] + data['ika12']
    aggregate_fields['aggr13_17'] = data['ika13'] + data['ika14'] + data[
        'ika15'] + data['ika16'] + data['ika17']
    aggregate_fields['aggr18_29'] = data['ika18'] + data['ika19'] + data[
        'ika20'] + data['ika21'] + data['ika22'] + data['ika23'] + data[
            'ika24'] + data['ika25_29']
    aggregate_fields['aggr30_64'] = data['ika30_34'] + data['ika35_39'] + data[
        'ika40_44'] + data['ika45_49'] + data['ika50_54'] + data[
            'ika55_59'] + data['ika60_64']
    aggregate_fields['aggr64_'] = data['ika65_69'] + data['ika70_74'] + data[
        'ika75_79'] + data['ika80_84'] + data['ika85_89'] + data[
            'ika90_94'] + data['ika95_']

    #add the fields to data rec array
    augmented_data = rf.rec_append_fields(data, aggregate_fields.keys(),
                                          aggregate_fields.values())
    #...and add these fields later. they are here for their column names, but the line above would cause an exception if we added them before
    aggregate_fields['asyht'] = data['asyht']
    aggregate_fields['ruots'] = data['ruots']
    aggregate_fields['ekoord'] = data['ekoord']
    aggregate_fields['nkoord'] = data['nkoord']

    # drop all fields whose names are not in aggregate_fields
    fields2drop = [
        d for d in data.dtype.names if d not in aggregate_fields.keys()
    ]
    return rf.rec_drop_fields(augmented_data, fields2drop)
Example #3
0
def add_cols(
        base: NpRecArray, data: NpRecArray,
        cols: NpFields = None) -> NpRecArray:
    """Add columns from source table to target table.

    Wrapper function to numpy's `rec_append_fields`_.

    Args:
        base: Numpy record array with table like data
        data: Numpy record array storing the fields to add to the base.
        cols: String or sequence of strings corresponding to the names of the
            new columns. If cols is None, then all columns of the data table
            are appended. Default: None

    Returns:
        Numpy record array containing the base array, as well as the
        appended columns.

    """
    cols = cols or getattr(data, 'dtype').names
    check.has_type("'cols'", cols, (tuple, str))
    cols = list(cols) # make cols mutable

    # Append fields
    return nprec.rec_append_fields(base, cols, [data[c] for c in cols])
Example #4
0
def addcols(
        base: NpRecArray, data: NpRecArray,
        cols: NpFields = None) -> NpRecArray:
    """Add columns from source table to target table.

    Wrapper function to numpy's `rec_append_fields`_.

    Args:
        base: Numpy record array with table like data
        data: Numpy record array storing the fields to add to the base.
        cols: String or sequence of strings corresponding to the names of the
            new columns. If cols is None, then all columns of the data table
            are appended. Default: None

    Returns:
        Numpy record array containing the base array, as well as the
        appended columns.

    """
    cols = cols or getattr(data, 'dtype').names
    check.has_type("'cols'", cols, (tuple, str))
    cols = list(cols) # make cols mutable

    # Append fields
    return nprf.rec_append_fields(base, cols, [data[c] for c in cols])
Example #5
0
    async def _add_turnover(self, frame_type: FrameType):
        # 从bars对应的frame中提取天数
        if frame_type in TimeFrame.minute_level_frames:
            dates = sorted(set(map(lambda x: x.date(), self._bars["frame"])))
        else:
            dates = sorted(set(self._bars["frame"]))
        date = dates[-1]
        n = len(dates)

        cc_recs = await Valuation.get_circulating_cap(self.code, date, n)

        # align circulating_cap with self.bars
        if frame_type != FrameType.DAY:
            tmp_bars = accl.numpy_append_fields(
                self._bars,
                "join_key",
                [x.date() for x in self._bars["frame"]],
                dtypes=[("join_key", "O")],
            )
            cc_recs.dtype.names = ["join_key", "circulating_cap"]
            tmp_bars = accl.join_by_left("join_key", tmp_bars, cc_recs)

        else:
            # rename 'date' to frame, so we can align self._bars with circulating_cap
            cc_recs.dtype.names = "frame", "circulating_cap"
            tmp_bars = accl.join_by_left("frame", self._bars, cc_recs)
        # todo: 对非股票类证券,circulating_cap的单位不一定是手(即100股),此处需要调查
        self._bars = rfn.rec_append_fields(
            self._bars,
            "turnover",
            tmp_bars["volume"] / tmp_bars["circulating_cap"] / 100,
            [("<f4")],
        )
        return self._bars
Example #6
0
    def merged_records(self,
              category=None,
              region=None,
              fields=None,
              cuts=None,
              clf=None,
              clf_name='classifier',
              include_weight=True,
              systematic='NOMINAL'):

        recs = self.records(
            category=category,
            region=region,
            fields=fields,
            include_weight=include_weight,
            cuts=cuts,
            systematic=systematic)

        if include_weight and fields is not None:
            if 'weight' not in fields:
                fields = list(fields) + ['weight']
        rec = stack(recs, fields=fields)

        if clf is not None:
            scores, _ = clf.classify(
                self, category, region,
                cuts=cuts, systematic=systematic)
            rec = recfunctions.rec_append_fields(rec,
                names=clf_name,
                data=scores,
                dtypes='f4')

        return rec
Example #7
0
def add_quantities(d, coords=None,massE=0.511e6):
    '''
    Add physically interesting quantities to
    pext data.

    Parameters:
    -----------
      d      : pext array

    Keywords:
    ---------
      coords : sequence of positions for angle calculation. None
               or by default, calculate no angles.
               For 2D, takes the angle depending on the order passed,
               so this can be used for left-handed coordinate systems
               like LSP's xz.
      massE  : rest mass energy of particles.
    Returns a copy with the quantities appended.
    '''
    quants = {}
    quants['u_norm'] = np.sqrt(d['ux']**2+d['uy']**2+d['uz']**2)
    quants['KE']     =(np.sqrt(quants['u_norm']**2+1)-1)*massE;
    coords[:] = ['u'+coord for coord in coords];
    if not coords:
        pass;
    elif len(coords) == 3:
        quants['theta'] = np.arccos(d[coords[2]]/quants['u_norm']);
        quants['phi']   = np.arctan2(d[coords[1]],d[coords[0]]);
        quants['phi_n'] = np.arctan2(d[coords[2]],d[coords[0]]);
    elif len(coords) == 2:
        quants['phi']   = np.arctan2(d[coords[1]],d[coords[0]]);
    keys,items = zip(*quants.items());
    return rfn.rec_append_fields(
        d, keys, items);
Example #8
0
def add_time_column(table, name='time', pop_start=True, pop_offset=True):
    """Append a column named 'time' by combining the gps_start and _offset

    Parameters
    ----------
    table : `EventTable`
        table of events to modify
    name : `str`, optional
        name of field to append, default: 'time'
    pop_start: `bool`, optional
        remove the 'gps_start' field when finished, default: `True`
    pop_offset: `bool`, optional
        remove the 'gps_offset' field when finished, default: `True`

    Returns
    -------
    mod : `recarray`, matches type of input
        a modified version of the input table with the new time field
    """
    type_ = type(table)
    t = table['gps_start'] + table['gps_offset']
    drop = []
    if pop_start:
        drop.append('gps_start')
    if pop_offset:
        drop.append('gps_offset')
    if drop:
        table = recfunctions.rec_drop_fields(table, drop)
    return recfunctions.rec_append_fields(table, [name], [t]).view(type_)
Example #9
0
def move_bad_fields_to_bottom(oldArray, orderedFieldList, orderedTypeList):
    """
    Move the given fields in a structured array to the bottom and change their type
    
    Input
    -----                                                           
    oldArray : numpy structured array
        previous array to modify               
    orderFieldList : list
        list of fields to move and change type 
    orderedTypeList : list
        list of new types for the fields
                           
    Returns an array with some fields moved to the bottom and with a different type
    """

    outArray = oldArray.copy()
    for name, typ in zip(orderedFieldList, orderedTypeList):
        #Remove field of interest from the array
        tmpArray = rec.rec_drop_fields(outArray, name)

        #Append the same field at the end of the array with the right data type
        outArray = rec.rec_append_fields(tmpArray,
                                         name,
                                         oldArray[name].copy(),
                                         dtypes=typ)
    return outArray
Example #10
0
 def parse_trkpts(self, keys=None, trkseg=-1):
     if (keys == None) or (len(keys) == 0):
         keys = self.get_trkpt_elements()
     row = self.get_trkpt_count(trkseg)
     self.d = np.ones(row, dtype={'names': ['ok'], 'formats': ['bool']})
     for key, typ in ([('lat', 'float'), ('lon', 'float')] + keys):
         self.d = npr.rec_append_fields(self.d, key, np.zeros(row), typ)
         self.scale[key] = 1.0
         self.unit[key] = "SI"
     idx = 0
     for trkpt in self.gpxdoc.iter('{*}trkpt'):
         self.d['lat'][idx] = float(trkpt.get(
             'lat'))  # lat and lon are the only mandatory elements
         self.d['lon'][idx] = float(trkpt.get(
             'lon'))  # lat and lon are the only mandatory elements
         for child in trkpt.findall('.//{*}*'):
             key = re.sub(r'\{.*?\}', '', child.tag)
             if key in self.d.dtype.names:
                 typ = dict(keys)[key]
                 if typ == 'float':
                     self.d[key][idx] = float(child.text)
                 elif typ == 'int':
                     self.d[key][idx] = int(child.text)
                 else:
                     self.d[key][idx] = child.text
         idx += 1
     self.append_column('idx', 'int')
     self['idx'] = np.arange(self.get_row_count())
Example #11
0
def filelist(flist, survey, stripndir=10, report=True):
    m, flag_a = qkeep_all(flist, problem=True)
    qkeepfun = {
        'mosaic': qkeep_mosaic,
        '90prime': qkeep_90prime,
        'decam': qkeep_decam
    }
    if survey in qkeepfun:
        ms, flag_s = qkeepfun[survey](flist, problem=True)
    else:
        raise ValueError('unknown survey!')
    m = m & ms
    flist['qkeep'] = m
    from numpy.lib.recfunctions import rec_append_fields
    flist = rec_append_fields(flist, ['flag1', 'flag2'], [flag_a, flag_s])
    if report:
        report_problems(flist)
    flist = flist[(flist['prodtype'] == 'image')
                  & (flist['proctype'] == 'InstCal')]
    m = flist['qkeep'] != 0
    m2 = most_recent_versions(flist[m])
    flist['qkeep'][m] = m2
    if survey == '90prime':
        # repair OBJECT keywords
        ind = numpy.flatnonzero(flist['object'] == '')
        for i, f in zip(ind, flist['filename'][ind]):
            newpath = f.replace('BOK_CP', 'BOK_Raw')
            newpath = newpath.replace('CP', '')
            newpath = newpath[:-17] + '_ori.fits.fz'
            flist['object'][i] = fits.getheader(newpath)['OBJECT']

    flist['filename'] = [
        '/'.join(f.split('/')[stripndir:]) for f in flist['filename']
    ]
    return flist
Example #12
0
def _add_weight_field(arr, tree, weight_name='weight', weight_dtype='f4'):
    weights = np.empty(arr.shape[0], dtype=weight_dtype)
    weights.fill(tree.GetWeight())
    return recfunctions.rec_append_fields(arr,
                                          names=weight_name,
                                          data=weights,
                                          dtypes=weight_dtype)
Example #13
0
def add_time_column(table, name='time', pop_start=True, pop_offset=True):
    """Append a column named 'time' by combining the gps_start and _offset

    Parameters
    ----------
    table : `EventTable`
        table of events to modify
    name : `str`, optional
        name of field to append, default: 'time'
    pop_start: `bool`, optional
        remove the 'gps_start' field when finished, default: `True`
    pop_offset: `bool`, optional
        remove the 'gps_offset' field when finished, default: `True`

    Returns
    -------
    mod : `recarray`, matches type of input
        a modified version of the input table with the new time field
    """
    type_ = type(table)
    t = table['gps_start'] + table['gps_offset']
    drop = []
    if pop_start:
        drop.append('gps_start')
    if pop_offset:
        drop.append('gps_offset')
    if drop:
        table = recfunctions.rec_drop_fields(table, drop)
    return recfunctions.rec_append_fields(table, [name], [t]).view(type_)
Example #14
0
def read_bliss_file(path):
    fits = fitsio.FITS(path)

    # The header for the image that this catalog was derived from is
    # stored in the first extension of the FITS file. However, it's format
    # is weird, so we parse it.

    # First we build a string
    hdrstr = '\n'.join(fits['LDAC_IMHEAD'].read()[0][0])
    # Then we use astropy to parse that string into a dict
    hdr = Header.fromstring(hdrstr, sep='\n')

    # Now we read the catalog
    observations = fits['LDAC_OBJECTS'].read()

    # The image header gives us access to image-level quantities, like
    # EXPNUM, CCDNUM, MJD-OBS, etc. Careful, these quantities may have a different byte order than the catalog data.
    EXPNUM = np.tile(hdr['EXPNUM'], len(observations))
    CCDNUM = np.tile(hdr['CCDNUM'], len(observations))

    # We can then append those quantities to the object array
    observations = rec_append_fields(observations,
                                     names=['EXPNUM', 'CCDNUM'],
                                     data=[EXPNUM, CCDNUM])
    return observations
Example #15
0
def extendDefaultTrainingDataset(names, arrays, base_default_featset='full'):
    """base : array
    Input array to extend.

    names : string, sequence
    String or sequence of strings corresponding to the names of the new fields.

    data : array or sequence of arrays
    Array or sequence of arrays storing the fields to add to the base.
    """

    training_dataset = getDefaultTrainingDataset()

    # select features from integrated dataset
    if base_default_featset is None:
        base_featset = []
    if base_default_featset in DEFAULT_FEATSETS:
        base_featset = DEFAULT_FEATSETS[base_default_featset]
    else:
        base_featset = list(base_default_featset)
    featset = ['SAV_coords', 'true_label'] + base_featset
    base_dataset = training_dataset[featset]

    # extend base training dataset
    fm = rfn.rec_append_fields(base_dataset, names, arrays, dtypes='float32')

    return fm
Example #16
0
    def estimate(self, nparrays: List[np.recarray]):
        """Compute the book leverage

        Parameters
        ----------
        nparrays : List[np.recarray]
            Input datasets, same order as specified in `DATASETS_REQUIRED`

        Returns
        -------
        Tuple[pd.DataFrame, Dict[str, str]]
            Output dataset and the `VARIABLE_LABELS`
        """
        nparray = filter_funda(nparrays[0])
        # debts = long-term debt + debt in current liabilities
        debts = np.nansum([nparray.dltt, nparray.dlc], axis=0)
        # assets = debts + common equity
        assets = np.nansum([debts, nparray.ceq], axis=0)
        # book leverage = debts / assets
        bleverage = np.true_divide(debts, assets, where=(assets != 0))
        # set book leverage to missing if common equity is somehow missing
        bleverage[np.isnan(nparray.ceq)] = np.nan
        # add book leverage to the result
        nparray = rfn.rec_append_fields(nparray, NAME, bleverage)
        # keep only useful columns
        cols = set(rfn.get_names_flat(nparray.dtype))
        nparray.sort(order=(keys := ["gvkey", "datadate"]))
Example #17
0
 def estimate(self, nparrays: List[np.recarray]):
     nparray = filter_funda(nparrays[0])
     size = np.log(nparray.at, where=(nparray.at > 0))
     size[np.isnan(nparray.at)] = np.nan
     nparray = rfn.rec_append_fields(nparray, NAME, size)
     # keep only useful columns
     cols = set(rfn.get_names_flat(nparray.dtype))
     nparray.sort(order=(keys := ["gvkey", "datadate"]))
Example #18
0
 def estimate(self, nparrays: List[np.recarray]):
     nparray = filter_funda(nparrays[0])
     roa = np.true_divide(nparray.ib, nparray.at, where=(nparray.at != 0))
     roa[np.isnan(nparray.at)] = np.nan
     nparray = rfn.rec_append_fields(nparray, NAME, roa)
     # keep only useful columns
     cols = set(rfn.get_names_flat(nparray.dtype))
     nparray.sort(order=(keys := ["gvkey", "datadate"]))
Example #19
0
def readHESE(fname):
    f = np.load(fname)
    mask = f['SignalTrackness'] > 0.1
    f = f[mask]
    f = rec_append_fields(
        f,
        'zenith',
        #utils.dec_to_zen(f[settings['dec_true']]),
        f[settings['dec_true']],
        dtypes=np.float64)
    f = rec_append_fields(f,
                          'azimuth',
                          f[settings['ra_true']],
                          dtypes=np.float64)
    f = rec_append_fields(
        f,
        'zen_reco',
        #utils.dec_to_zen(f['DirReco_splinempe_dec']),
        f['DirReco_splinempe_dec'],
        dtypes=np.float64)
    f = rec_append_fields(f,
                          'az_reco',
                          f['DirReco_splinempe_ra'],
                          dtypes=np.float64)
    f = rec_append_fields(f,
                          'astro',
                          f['EMinus2'] /
                          f['TrueEnergy']**(settings['gamma'] - 2.),
                          dtypes=np.float64)
    f = rec_append_fields(f,
                          'ow',
                          f['EMinus2'] * f['TrueEnergy']**2,
                          dtypes=np.float64)
    f = rec_append_fields(f, 'energy', f['TrueEnergy'], dtypes=np.float64)
    return f
Example #20
0
def append_feat(data, name, fieldarray):
    ''' Appends fieldarray to data with the name 'name'. This allows new
    features to be added easily. 
    Because all new features will be built differently, it is up to you to 
    construct the fieldarray properly. 
    This is basically just a recast of numpy.lib.recfunctions.rec_append_fields
    , so that I do not have to look up the function again.'''

    return rfunc.rec_append_fields(data, name, fieldarray)
Example #21
0
def append_feat(data, name, fieldarray):
    ''' Appends fieldarray to data with the name 'name'. This allows new
    features to be added easily. 
    Because all new features will be built differently, it is up to you to 
    construct the fieldarray properly. 
    This is basically just a recast of numpy.lib.recfunctions.rec_append_fields
    , so that I do not have to look up the function again.'''

    return rfunc.rec_append_fields(data, name, fieldarray)
Example #22
0
def _add_weight_field(arr, tree,
                      weight_name='weight',
                      weight_dtype='f4'):
    weights = np.empty(arr.shape[0], dtype=weight_dtype)
    weights.fill(tree.GetWeight())
    return recfunctions.rec_append_fields(
            arr, names=weight_name,
            data=weights,
            dtypes=weight_dtype)
Example #23
0
def add_vsig(dat):
    vsig = get_vsig(dat)

    if 'vsig' in dat.dtype.names:

        dat['vsig'] = vsig
        return dat

    return nprf.rec_append_fields(dat, 'vsig', vsig, dtypes='<f4')
Example #24
0
def _add_weight_field(arr, tree, include_weight=False,
                      weight_name='weight',
                      weight_dtype='f4'):
    if not include_weight:
        return arr
    weights = np.ones(arr.shape[0], dtype=weight_dtype)
    weights *= tree.GetWeight()
    return recfunctions.rec_append_fields(arr, names=weight_name,
                                          data=weights,
                                          dtypes=weight_dtype)
Example #25
0
def add_rsig(dat):

    rsig = get_rsig(dat)

    if 'rsig' in dat.dtype.names:

        dat['rsig'] = rsig
        return dat

    return nprf.rec_append_fields(dat, 'rsig', rsig, dtypes='<f4')
Example #26
0
def recarr_join(arr1, arr2):
    """
    Join to record arrays column wise.
    """
    arr1 = arr1.copy()
    for field in arr2.dtype.names:
        #original function, deprecated
        # arr1 = pl.mlab.rec_append_fields(arr1,field,arr2[field])
        #new function
        arr1 = rec_append_fields(arr1, field, arr2[field])
    return arr1
Example #27
0
def estimate(nparrays: List[np.recarray]):
    def filter_funda(x):
        return x[np.in1d(x.datafmt, ('STD')) & np.in1d(x.indfmt, ('INDL'))
                 & np.in1d(x.popsrc, ('D')) & np.in1d(x.consol, ('C'))]

    nparray = filter_funda(nparrays[0])
    capx = np.true_divide(nparray.capx, nparray.at, where=(nparray.at != 0))
    capx[np.isnan(nparray.at)] = np.nan
    nparray = rfn.rec_append_fields(nparray, name, capx)
    # keep only useful columns
    cols = set(rfn.get_names_flat(nparray.dtype))
    nparray.sort(order=(keys := ['gvkey', 'datadate']))
Example #28
0
def _add_weight_field(arr,
                      tree,
                      weight_name='weight',
                      weight_dtype='f4'):
    """Add a new column containing the tree weight.

    """
    weights = np.empty(arr.shape[0], dtype=weight_dtype)
    weights.fill(tree.GetWeight())
    return recfunctions.rec_append_fields(
        arr, names=weight_name,
        data=weights,
        dtypes=weight_dtype)
Example #29
0
 def estimate(self, nparrays: List[np.recarray]):
     nparray = filter_funda(nparrays[0])
     # market value at fiscal year
     mv = nparray.prcc_f * nparray.csho
     # market-to-book = market value of equity / common equity
     mtb = np.true_divide(mv, nparray.ceq, where=(nparray.ceq != 0))
     # set mtb to missing if common equity is somehow missing
     mtb[np.isnan(nparray.ceq)] = np.nan
     # add book leverage to the result
     nparray = rfn.rec_append_fields(nparray, NAME, mtb)
     # keep only useful columns
     cols = set(rfn.get_names_flat(nparray.dtype))
     nparray.sort(order=(keys := ["gvkey", "datadate"]))
Example #30
0
    def lens (focalLength = [0], Z = [0], label = None):

        #   -- component.lens --
        # Create a lens component object.
        # Example:
        # component.lens(f,z,label);
        # This creates a lens component with focal length f at position
        # z. label is a string which is used to identify the component.

        numcomps = len(focalLength)

        if numcomps > 1:
            zlength = len(Z)

            if zlength != numcomps:
                if zlength != 1:
                    raise Exception ("List of focal lengths must be the\
                                    same length as list of z positions.")                
                Z = Z*numcomps

            if label is not None:
                lablength = len(label)

                if lablength != numcomps:
                    if lablength != 1:
                        raise Exception ("List of focal lengths must be the\
                                        same length as list of labels.") 
                    label = label*numcomps

            lenslist = [[0,0,0]] # make the count start from 1
            for n in range(numcomps):
                c = component()
                if label is not None:
                    cl = c.lens([focalLength[n]],[Z[n]],[label[n]])
                else:
                    cl = c.lens([focalLength[n]],[Z[n]])
                lenslist.append(cl)
            return lenslist 

        
        fl = -1./focalLength[0]
        M = np.matrix([[1,0],[fl,1]])        

        o = component(M, Z)
        o.type = 'lens'
        o.parameters = nprf.rec_append_fields(o.parameters,'focalLength',[fl],dtypes = [(float)])
        if label is not None:
            o.label = label[0]
        else:
            o.label = 'no label'
        return o
Example #31
0
    def curvedMirror (radiusOfCurvature = [0], Z = [0], label = None):

        # -- component.curvedMirror --
        # Create a curved mirror component object.
        # Example:
        # mylens = component.curvedMirror(ROC,z,label);
        # This creates a lens component with radius of curvature ROC at position
        # z. label is a string which is used to identify the component.

        numcomps = len(radiusOfCurvature)

        if numcomps > 1:
            zlength = len(Z)

            if zlength != numcomps:
                if zlength != 1:
                    raise Exception ("List of radii must be the\
                                    same length as list of z positions.")                
                Z = Z*numcomps

            if label is not None:
                lablength = len(label)

                if lablength != numcomps:
                    if lablength != 1:
                        raise Exception ("List of radii must be the\
                                        same length as list of labels.") 
                    label = label*numcomps

            curvedMirrorlist = [[0,0,0]] # make the count start from 1
            for n in range(numcomps):
                c = component()
                if label is not None:
                    ccm = c.curvedMirror([radiusOfCurvature[n]],[Z[n]],[label[n]])
                else:
                    ccm = c.curvedMirror([radiusOfCurvature[n]],[Z[n]])
                curvedMirrorlist.append(ccm)
            return curvedMirrorlist

        radii = radiusOfCurvature[0]
        M = np.matrix([[1,0],[-2./radii,1]])        

        o = component(M, Z)
        o.type = 'curved mirror'
        o.parameters = nprf.rec_append_fields(o.parameters,'ROC',[radii],dtypes = [(float)])
        if label is not None:
            o.label = label[0]
        else:
            o.label = 'no label'
        return o
Example #32
0
def filelist(flist,
             survey,
             report=True,
             rawdir='/global/cfs/cdirs/cosmo/staging'):
    m, flag_a = qkeep_all(flist, problem=True)
    qkeepfun = {
        'mosaic': qkeep_mosaic,
        '90prime': qkeep_90prime,
        'decam': qkeep_decam
    }
    if survey in qkeepfun:
        ms, flag_s = qkeepfun[survey](flist, problem=True)
    else:
        raise ValueError('unknown survey!')
    m = m & ms
    flist['qkeep'] = m
    from numpy.lib.recfunctions import rec_append_fields
    flist = rec_append_fields(flist, ['flag1', 'flag2'], [flag_a, flag_s])
    if report:
        report_problems(flist)
    flist = flist[(flist['prodtype'] == 'image')
                  & (flist['proctype'] == 'InstCal')]
    m = flist['qkeep'] != 0
    m2 = most_recent_versions(flist[m])
    flist['qkeep'][m] = m2
    if survey == '90prime':
        # repair OBJECT keywords
        ind = numpy.flatnonzero(flist['object'] == '')
        for i, f in zip(ind, flist['filename'][ind]):
            # newpath = f.replace('BOK_CP', 'BOK_Raw')
            # newpath = newpath.replace('CP', '')
            dirs = f.split('/')
            dirs = dirs[:-3] + dirs[-2:]  # get rid of V2.0
            dirs[-3] = 'BOK_Raw'
            dirs[-2] = dirs[-2][2:]
            dirs = [rawdir] + dirs[-4:]
            newpath = '/'.join(dirs)
            newpath = newpath.replace('work/legacysurvey/dr9/images',
                                      'staging')
            cut = newpath.find('_ooi_')
            newpath = newpath[:cut] + '_ori.fits.fz'
            flist['object'][i] = fits.getheader(newpath)['OBJECT']

    imagesind = numpy.flatnonzero(
        [f == 'images' for f in flist['filename'][0].split('/')])
    imagesind = imagesind[0]
    flist['filename'] = [
        '/'.join(f.split('/')[imagesind + 1:]) for f in flist['filename']
    ]
    return flist
Example #33
0
 def open_fit(self, filename):
     self.filename = filename
     a = FitFile(filename)
     a.parse()
     records = list(a.get_messages(name='record'))
     row = len(records)
     if row != 0:
         keys = []
         self.d = np.ones(row, dtype={'names': ['ok'], 'formats': ['bool']})
         #parse first record
         for field in records[0]:
             if field.type.name == 'date_time':
                 keys.append((field.name, 'a20'))
             else:
                 keys.append((field.name, 'float'))
         for key, typ in (keys):
             #some files edited by fitfiletools (https://www.fitfiletools.com/) (eg remover) may have some strange structure, with twice the same field!
             try:
                 self.d = npr.rec_append_fields(self.d, key, np.zeros(row),
                                                typ)
                 self.scale[key] = 1.0
                 self.unit[key] = "SI"
             except:
                 pass
     idx = 0
     for r in a.get_messages(name='record'):
         for f in r.fields:
             if f.type.name == 'date_time':
                 self.d[f.name][idx] = f.value.isoformat() + 'Z'
             else:
                 try:
                     self.d[f.name][idx] = float(f.value)
                 except:
                     #print("invalid data", f.name, idx,f.data)
                     self.d[f.name][idx] = -1.0
         idx += 1
     #we now have to convert lat and lon
     self.d['position_lat'] /= 11930464.71
     self.d['position_long'] /= 11930464.71
     #new version of fitparse gives speed in mm/s
     self.d['speed'] /= 1000
     self.move_column('position_lat', 'lat')
     self.move_column('position_long', 'lon')
     self.move_column('altitude', 'ele')
     self.move_column('timestamp', 'time')
     self.append_column('idx', 'int')
     self['idx'] = np.arange(self.get_row_count())
Example #34
0
    def propagator (DZ = 0, Z = 0, label = None):

        # -- component.propagator --
        # Create a free-space propagator component object.
        # Example:
        # mylens = component.propagator(dz,z,label);
        # This creates a propagator component with length dz at position
        # z. label is a string which is used to identify the component.

        M = np.matrix([[1, DZ], [0, 1]]) 

        o = component(M, Z)
        o.type = 'propagator'
        o.parameters = nprf.rec_append_fields(o.parameters,'length',[DZ],dtypes = [(float)])
        if label is not None:
            o.label = label
        else:
            o.label = 'no label'
        return o
Example #35
0
def estimate(nparrays: List[np.recarray]):
    def filter_funda(x):
        return x[np.in1d(x.datafmt, ('STD')) & np.in1d(x.indfmt, ('INDL'))
                 & np.in1d(x.popsrc, ('D')) & np.in1d(x.consol, ('C'))]

    nparray = filter_funda(nparrays[0])
    # debts = long-term debt + debt in current liabilities
    debts = np.nansum([nparray.dltt, nparray.dlc], axis=0)
    # assets = debts + common equity
    assets = np.nansum([debts, nparray.ceq], axis=0)
    # book leverage = debts / assets
    bleverage = np.true_divide(debts, assets, where=(assets != 0))
    # set book leverage to missing if common equity is somehow missing
    bleverage[np.isnan(nparray.ceq)] = np.nan
    # add book leverage to the result
    nparray = rfn.rec_append_fields(nparray, name, bleverage)
    # keep only useful columns
    cols = set(rfn.get_names_flat(nparray.dtype))
    nparray.sort(order=(keys := ['gvkey', 'datadate']))
Example #36
0
    def records(self,
                category=None,
                region=None,
                fields=None,
                cuts=None,
                include_weight=True,
                systematic='NOMINAL',
                return_idx=False,
                **kwargs):
        if include_weight and fields is not None:
            if 'weight' not in fields:
                fields = list(fields) + ['weight']

        selection = self.cuts(category, region) & cuts

        log.info("requesting table from Data %d" % self.year)
        log.debug("using selection: %s" % selection)

        # read the table with a selection
        if selection:
            rec = self.h5data.read_where(selection.where(), **kwargs)
        else:
            rec = self.h5data.read(**kwargs)

        # add weight field
        if include_weight:
            # data is not weighted
            weights = np.ones(rec.shape[0], dtype='f8')
            rec = recfunctions.rec_append_fields(rec,
                names='weight',
                data=weights,
                dtypes='f8')

        if fields is not None:
            rec = rec[fields]

        if return_idx:
            # only valid if selection is non-empty
            idx = self.h5data.get_where_list(selection.where(), **kwargs)
            return [(rec, idx)]

        return [rec]
Example #37
0
def addhash(frame,**kw):
    '''
    helper function to add hashes to the given frame
    given in the dictionary d returned from firsthash.

    Parameters:
    -----------
      frame :  frame to hash.

    Keywords:
    ---------
      same as genhash
    
    Returns frame with added hashes, although it will be added in
    place.
    '''
    hashes = genhash(frame,**kw);
    frame['data'] = rfn.rec_append_fields(
        frame['data'],'hash',hashes);
    return frame;
Example #38
0
def firstyeartilefile(tiles):
    tiles = numpy.array(tiles).copy()
    tiles.dtype.names = [n.lower() for n in tiles.dtype.names]
    newtiles = tiles.copy()
    permute = {0: 2, 1: 3, 2: 4, 3: 0, 4: 1}
    for npass, opass in permute.items():
        newtiles[tiles['pass'] == npass] = tiles[tiles['pass'] == opass]
    keepfields = ['tileid', 'pass']
    for field in keepfields:
        newtiles[field] = tiles[field]
    ntiles = numpy.sum(tiles['pass'] == 0)
    # ntiles*3 since permute[3] = 0, the original centers
    newtiles['centerid'] = (newtiles['tileid'] % ntiles) + ntiles * 3
    from numpy.lib.recfunctions import rec_append_fields
    myr1passes = firstyeartiles(newtiles, orig=False)
    myr1 = numpy.zeros(len(newtiles), dtype='bool')
    for m0 in myr1passes:
        myr1 = myr1 | m0
    newtiles = rec_append_fields(newtiles, 'year1', myr1)
    return newtiles
Example #39
0
    def dielectric (R1, R2, thickness = 0, n = 1, Z = 0, label = None):
        # -- component.dielectric --
        # Create a dielectric component object.
        # Example:
        # mylens = component.dielectric(R1, R2, thickness, n, Z, label);
        # This creates a dielectric (thick lens) component at position
        # z. label is a string which is used to identify the component.
        
        dist = np.matrix([[1, thickness],[0, 1]])
        refract1 = np.matrix([[1, 0], [(n-1)/R2, n]])
        refract2 = np.matrix([[1, 0], [(1-n)/(R1*n), 1./n]])
        M = refract1*dist*refract2 

        o = component(M, Z)
        o.type = 'dielectric'
        o.parameters = nprf.rec_append_fields(o.parameters,'length',[thickness],dtypes = [(float)])
        if label is not None:
            o.label = label
        else:
            o.label = 'no label'
        return o
Example #40
0
def addhash(frame, hashd, removedupes=False):
    '''
    [Copied from Gregory Ngirmang's lspreader.
    https://github.com/noobermin/lspreader]
    helper function to add hashes to the given frame
    given in the dictionary d returned from firsthash.
    Parameters:
    -----------
      frame :  frame to hash.
      d     :  hash specification generated from firsthash.
    Keywords:
    ---------
      removedupes: put -1 in duplicates
    
    Returns frame with added hashes, although it will be modified in
    place
    '''
    hashes = genhash(frame, hashd, removedupes)
    frame['data'] = rfn.rec_append_fields(
        frame['data'],'hash',hashes)
    return frame
Example #41
0
def addhash(frame,d,removedupes=False):
    '''
    helper function to add hashes to the given frame
    given in the dictionary d returned from firsthash.

    Parameters:
    -----------
      frame :  frame to hash.
      d     :  hash specification generated from firsthash.

    Keywords:
    ---------
      removedups: put -1 in duplicates
    
    Returns frame with added hashes, although it will be added in
    place.
    '''
    hashes = genhash(frame,d,removedupes);
    frame['data'] = rfn.rec_append_fields(
        frame['data'],'hash',hashes);
    return frame;
Example #42
0
def estimate(nparrays: List[np.recarray]):

    def filter_funda(x): return x[
        np.in1d(x.datafmt, ('STD')) &
        np.in1d(x.indfmt, ('INDL')) &
        np.in1d(x.popsrc, ('D')) &
        np.in1d(x.consol, ('C'))
    ]

    nparray = filter_funda(nparrays[0])
    # market value at fiscal year
    mv = nparray.prcc_f * nparray.csho
    # market-to-book = market value of equity / common equity
    mtb = np.true_divide(mv, nparray.ceq, where=(nparray.ceq != 0))
    # set mtb to missing if common equity is somehow missing
    mtb[np.isnan(nparray.ceq)] = np.nan
    # add book leverage to the result
    nparray = rfn.rec_append_fields(nparray, name, mtb)
    # keep only useful columns
    cols = set(rfn.get_names_flat(nparray.dtype))
    nparray.sort(order=(keys := ['gvkey', 'datadate']))
Example #43
0
    def estimate(self, nparrays: List[np.recarray]):
        """Compute the capital expenditures scaled by total assets

        Parameters
        ----------
        nparrays : List[np.recarray]
            Input datasets, same order as specified in `DATASETS_REQUIRED`

        Returns
        -------
        Tuple[pd.DataFrame, Dict[str, str]]
            Output dataset and the `VARIABLE_LABELS`
        """
        nparray = filter_funda(nparrays[0])
        capx = np.true_divide(nparray.capx,
                              nparray.at,
                              where=(nparray.at != 0))
        capx[np.isnan(nparray.at)] = np.nan
        nparray = rfn.rec_append_fields(nparray, NAME, capx)
        # keep only useful columns
        cols = set(rfn.get_names_flat(nparray.dtype))
        nparray.sort(order=(keys := ["gvkey", "datadate"]))
Example #44
0
def to_recarray(self, columns=None, on_attributeerror='raise',
                get_as_columns=False):
    """Convert this table to a structured `numpy.recarray`

    This returned `~numpy.recarray` is a blank data container, mapping
    columns in the original LIGO_LW table to fields in the output, but
    mapping none of the instance methods of the origin table.

    Parameters
    ----------
    columns : `list` of `str`, optional
        the columns to populate, if not given, all columns present in the
        table are mapped

    on_attributeerror : `str`, optional
        how to handle `AttributeError` when accessing rows, one of

        - 'raise' : raise normal exception
        - 'ignore' : skip over this column
        - 'warn' : print a warning instead of raising error

    get_as_columns : `bool`, optional
        convert all `get_xxx()` methods into fields in the
        `~numpy.recarray`; the default is to _not_ do this.
    """
    # get numpy-type columns
    if columns is None:
       columns = self.columnnames
    dtypes = [(str(c), NUMPY_TYPE[self.validcolumns[c]])
              for c in columns]
    # create array
    m = len(self)
    out = GWRecArray((len(self),), dtype=dtypes)
    # and fill it
    for column in columns:
        orig_type = self.validcolumns[column]
        try:
            if orig_type == 'ilwd:char':  # numpy tries long() which breaks
                out[column] = map(int, self.getColumnByName(column))
            else:
                out[column] = self.getColumnByName(column)
        except AttributeError as e:
            if on_attributeerror == 'ignore':
                pass
            elif on_attributeerror == 'warn':
                warnings.warn('Caught %s: %s' % (type(e).__name__, str(e)))
            else:
                raise
    # fill out get_xxx columns
    if get_as_columns:
        getters = filter(lambda x: x[0].startswith('get_'),
                         inspect.getmembers(self, predicate=inspect.ismethod))
        new = []
        for name, meth in getters:
            column = name.split('_', 1)[1]
            if column in columns:  # don't overwrite existing columns
                continue
            try:
                array = meth()
            except (AttributeError, ValueError, TypeError):
                continue
            else:
                try:
                    dtype = array.dtype
                except AttributeError:
                    try:
                        dtype = type(array[0])
                    except (TypeError, KeyError):
                        continue
                    except IndexError:
                        dtype = None
                if dtype == LIGOTimeGPS:
                    dtype = numpy.float64
                elif dtype in INVALID_REC_TYPES:
                    continue
                new.append((column, array, dtype))
        names, data, dtypes = zip(*new)
        if names:
            out = recfunctions.rec_append_fields(
                out, names, data, dtypes).view(type(out))

    return out
Example #45
0
with open(lspf,"r") as f:
    lsp=f.read();
if not opts['<output>']:
    outname = re.search("(.*)\.lsp$",lspf).group(1)+"-pext";
else:
    outname = opts['<output>'];
dim=getdim(lsp);
pexts = getpexts(lsp);
latetime = float(opts['--late-time']) if opts['--late-time'] else None;
vprint('reading in files');
d = [ rd.read(name)
      for name in pext ];

d[:] = [
    rfn.rec_append_fields(
        id, 'species',
        np.ones(len(id)).astype(int)*pexts[i]['species'])
    for id,i in zip(d,key) ];
vprint('length of d={}'.format(len(d)));

d = [ i for i in d if i['t'].shape[0] > 0];
vprint('length of d={} after remove empties'.format(len(d)));
vprint('cutting out duplicate times');
if len(d) > 1:
    d = np.concatenate(d);
elif d == []:
    print("empty pext");
    quit();
else:
    d = d[0];
vprint('sorting by times')
Example #46
0
q1005  = readtxt('/home/ntejos/catalogs/Finn_2012/q1005_fuv_HI.sort',names=names)
q0209  = readtxt('/home/ntejos/catalogs/Finn_2012/q0209_fuv_HI.sort',names=names)
q1357  = readtxt('/home/ntejos/catalogs/Finn_2012/q1357_fuv_HI.sort',names=names)
q0107a = readtxt('/home/ntejos/catalogs/Q0107/A_HI.txt',names=names)
q0107b = readtxt('/home/ntejos/catalogs/Q0107/B_HI.txt',names=names)
q0107c = readtxt('/home/ntejos/catalogs/Q0107/C_HI.txt',names=names)
q1022  = readtxt('/home/ntejos/catalogs/J1022/HI.txt',names=names)

#read spectrum
names = 'wa,fl,er,norm'
q1005_fuv = readtxt('/home/ntejos/COS/q1005/FUV/q1005_fuv_norm.txt',names=names)


#append RA and DEC fields to QSO sightlines
radec = s2dec('10 05 35.24','01 34 45.7')
q1005 = rf.rec_append_fields(q1005,'RA' ,radec[0]*np.ones(len(q1005)))
q1005 = rf.rec_append_fields(q1005,'DEC',radec[1]*np.ones(len(q1005)))
radec = s2dec('02 09 30.7','-04 38 26')
q0209 = rf.rec_append_fields(q0209,'RA' ,radec[0]*np.ones(len(q0209)))
q0209 = rf.rec_append_fields(q0209,'DEC',radec[1]*np.ones(len(q0209)))
radec = s2dec('13 57 26.27','04 35 41.4')
q1357 = rf.rec_append_fields(q1357,'RA' ,radec[0]*np.ones(len(q1357)))
q1357 = rf.rec_append_fields(q1357,'DEC',radec[1]*np.ones(len(q1357)))
radec = s2dec('01 10 13.1','-02 19 52')
q0107a = rf.rec_append_fields(q0107a,'RA' ,radec[0]*np.ones(len(q0107a)))
q0107a = rf.rec_append_fields(q0107a,'DEC',radec[1]*np.ones(len(q0107a)))
radec  = s2dec('01 10 16.2','-02 18 50')
q0107b = rf.rec_append_fields(q0107b,'RA' ,radec[0]*np.ones(len(q0107b)))
q0107b = rf.rec_append_fields(q0107b,'DEC',radec[1]*np.ones(len(q0107b)))
radec  = s2dec('01 10 14.52','-02 16 57.5')
q0107c = rf.rec_append_fields(q0107c,'RA' ,radec[0]*np.ones(len(q0107c)))
Example #47
0
    def read_gals(
        self,
        model: Model,
        file_num: int,
        snapshot: int,
        pbar: Optional[tqdm] = None,
        plot_galaxies: bool = False,
        debug: bool = False
    ):
        """
        Reads the galaxies of a model file at snapshot specified by
        :py:attr:`~sage_analysis.model.Model.snapshot`.

        Parameters
        ----------

        model: :py:class:`~sage_analysis.model.Model` class
            The :py:class:`~sage_analysis.model.Model` we're reading data for.

        file_num: int
            Suffix number of the file we're reading.

        pbar: ``tqdm`` class instance, optional
            Bar showing the progress of galaxy reading.  If ``None``, progress bar will
            not show.

        plot_galaxies: bool, optional
            If set, plots and saves the 3D distribution of galaxies for this file.

        debug: bool, optional
            If set, prints out extra useful debug information.

        Returns
        -------
        gals : ``numpy`` structured array with format given by :py:method:`~_get_galaxy_struct`
            The galaxies for this file.

        Notes
        -----
        ``tqdm`` does not play nicely with printing to stdout. Hence we disable
        the ``tqdm`` progress bar if ``debug=True``.
        """

        fname = self._check_for_file(model, file_num)
        if fname is None:
            logger.debug(f"File\t{fname} \tdoes not exist!")
            return None

        with open(fname, "rb") as f:
            # First read the header information.
            Ntrees = np.fromfile(f, np.dtype(np.int32), 1)[0]
            num_gals = np.fromfile(f, np.dtype(np.int32), 1)[0]
            _ = np.fromfile(f, np.dtype((np.int32, Ntrees)), 1)

            # If there aren't any galaxies, exit here.
            if num_gals == 0:
                return None

            # Then the actual galaxies.
            gals = np.fromfile(f, self.galaxy_struct, num_gals)

            # If we're using the `tqdm` package, update the progress bar.
            if pbar is not None:
                pbar.set_postfix(file=fname, refresh=False)
                pbar.update(num_gals)

        if debug:
            print("")
            print(f"File {fname} contained {Ntrees} trees with {num_gals} galaxies")

            w = np.where(gals["StellarMass"] > 1.0)[0]
            print(f"{len(w)} of these galaxies have mass greater than 10^10Msun/h")

        if plot_galaxies:

            from sage_analysis.plots import plot_spatial_3d

            # Show the distribution of galaxies in 3D.
            pos = gals["Pos"][:]
            output_file = f"./galaxies_{file_num}.{model.plot_output_format}"
            plot_spatial_3d(pos, output_file, self.box_size)

        # For the HDF5 file, some data sets have dimensions Nx1 rather than Nx3
        # (e.g., Position). To ensure the galaxy data format is identical to the binary
        # output, we will split the binary fields into Nx1. This is simpler than creating
        # a new dataset within the HDF5 regime.
        from numpy.lib import recfunctions as rfn
        multidim_fields = ["Pos", "Vel", "Spin"]
        dim_names = ["x", "y", "z"]

        for field in multidim_fields:
            for dim_num, dim_name in enumerate(dim_names):
                dim_field = f"{field}{dim_name}"
                gals = rfn.rec_append_fields(gals, dim_field,
                                             gals[field][:, dim_num])

        return gals
Example #48
0
def get_triggers(channel, etg, segments, cache=None, snr=None, frange=None,
                 columns=None, raw=False, **kwargs):
    """Get triggers for the given channel
    """
    # get table from etg
    try:
        Table = TABLE[etg.lower()]
    except KeyError as e:
        e.args = ('Unknown ETG %r, cannot map to LIGO_LW Table class' % etg,)
        raise
    tablename = strip_table_name(Table.tableName)
    # get default columns for this table
    if columns is None:
        for key in COLUMNS:
            if issubclass(Table, key):
                columns = COLUMNS[key][:]
                break
    if 'channel' in columns:
        columns.pop('channel')

    # find triggers
    if cache is None:
        cache = find_trigger_files(channel, etg, segments, **kwargs)

    # read cache
    trigs = lsctables.New(Table, columns=columns)
    cache = cache.unique()
    cache.sort(key=lambda x: x.segment[0])
    for segment in segments:
        if len(cache.sieve(segment=segment)):
            if tablename.endswith('_inspiral'):
                filt = lambda t: float(t.get_end()) in segment
            else:
                filt = lambda t: float(t.get_peak()) in segment
            trigs.extend(Table.read(cache.sieve(segment=segment), filt=filt))

    # format table as numpy.recarray
    recarray = trigs.to_recarray(columns=columns)

    # filter
    if snr is not None:
        recarray = recarray[recarray['snr'] >= snr]
    if tablename.endswith('_burst') and frange is not None:
        recarray = recarray[
            (recarray['peak_frequency'] >= frange[0]) &
            (recarray['peak_frequency'] < frange[1])]

    # return basic table if 'raw'
    if raw:
        return recarray

    # otherwise spend the rest of this function converting functions to
    # something useful for the hveto core analysis
    addfields = {}
    dropfields = []

    # append channel to all events
    columns.append('channel')
    addfields['channel'] = numpy.repeat(channel, recarray.shape[0])

    # rename frequency column
    if tablename.endswith('_burst'):
        recarray = recfunctions.rename_fields(
            recarray, {'peak_frequency': 'frequency'})
        idx = columns.index('peak_frequency')
        columns.pop(idx)
        columns.insert(idx, 'frequency')

    # map time to its own column
    if tablename.endswith('_inspiral'):
        tcols = ['end_time', 'end_time_ns']
    elif tablename.endswith('_burst'):
        tcols = ['peak_time', 'peak_time_ns']
    else:
        tcols = None
    if tcols:
        times = recarray[tcols[0]] + recarray[tcols[1]] * 1e-9
        addfields['time'] = times
        dropfields.extend(tcols)
        columns = ['time'] + columns[2:]

    # add and remove fields as required
    if addfields:
        names, data = zip(*addfields.items())
        recarray = recfunctions.rec_append_fields(recarray, names, data)
        recarray = recfunctions.rec_drop_fields(recarray, dropfields)

    return recarray[columns]
Example #49
0
File: mc.py Project: sagittaeri/htt
    def records(self,
                category=None,
                region=None,
                fields=None,
                cuts=None,
                include_weight=True,
                systematic='NOMINAL',
                scale=1.,
                return_idx=False,
                **kwargs):

        from .ztautau import Ztautau

        if include_weight and fields is not None:
            if 'weight' not in fields:
                fields = list(fields) + ['weight']

        selection = self.cuts(category, region, systematic) & cuts
        table_selection = selection.where()

        if systematic == 'NOMINAL':
            log.info("requesting table from %s" %
                     (self.__class__.__name__))
        else:
            log.info("requesting table from %s for systematic %s " %
                     (self.__class__.__name__, systematic_name(systematic)))
        log.debug("using selection: %s" % selection)

        # TODO: handle cuts in weight expressions
        weight_branches = self.get_weight_branches(systematic, no_cuts=True)
        if systematic in SYSTEMATICS_BY_WEIGHT:
            systematic = 'NOMINAL'

        recs = []
        if return_idx:
            idxs = []
        for ds, _, sys_tables, sys_events, xs, kfact, effic in self.datasets:

            try:
                table = sys_tables[systematic]
                events = sys_events[systematic]
            except KeyError:
                log.debug(
                    "table for %s not present for %s "
                    "using NOMINAL" % (systematic, ds.name))
                table = sys_tables['NOMINAL']
                events = sys_events['NOMINAL']

            actual_scale = self.scale
            if isinstance(self, Ztautau):
                if systematic == ('ZFIT_UP',):
                    log.debug("scaling up for ZFIT_UP")
                    actual_scale += self.scale_error
                elif systematic == ('ZFIT_DOWN',):
                    log.debug("scaling down for ZFIT_DOWN")
                    actual_scale -= self.scale_error

            weight = (
                scale * actual_scale *
                LUMI[self.year] *
                xs * kfact * effic / events)

            # read the table with a selection
            try:
                if table_selection:
                    rec = table.read_where(table_selection, **kwargs)
                else:
                    rec = table.read(**kwargs)
            except Exception as e:
                print table
                print e
                continue
                #raise

            if return_idx:
                # only valid if table_selection is non-empty
                idx = table.get_where_list(table_selection, **kwargs)
                idxs.append(idx)

            # add weight field
            if include_weight:
                weights = np.empty(rec.shape[0], dtype='f8')
                weights.fill(weight)
                # merge the weight fields
                weights *= reduce(np.multiply,
                    [rec[br] for br in weight_branches])
                # drop other weight fields
                rec = recfunctions.rec_drop_fields(rec, weight_branches)
                # add the combined weight
                rec = recfunctions.rec_append_fields(rec,
                    names='weight',
                    data=weights,
                    dtypes='f8')
                if rec['weight'].shape[0] > 1 and rec['weight'].sum() == 0:
                    log.warning("{0}: weights sum to zero!".format(table.name))

            if fields is not None:
                try:
                    rec = rec[fields]
                except Exception as e:
                    print table
                    print rec.shape
                    print rec.dtype
                    print e
                    raise
            recs.append(rec)

        if return_idx:
            return zip(recs, idxs)
        return recs
Example #50
0
    def calculate(self, infile, field=1, simple=False):
        logger.info("Calculating magnitude limit from %s"%infile)

        #manglefile = self.config['mangle']['infile_%i'%field]
        #footfile = self.config['data']['footprint']
        #try: 
        #    footprint = fitsio.read(footfile)['I'].ravel()
        #except:
        #    logger.warn("Couldn't open %s; will try again."%footfile)
        #    footprint = footfile

        mag_column = self.config['catalog']['mag_%i_field'%field]
        magerr_column = self.config['catalog']['mag_err_%i_field'%field]

        # For simple maglims
        release = self.config['data']['release'].lower()
        band    = self.config['catalog']['mag_%i_band'%field]
        pixel_pix_name = 'PIX%i'%self.nside_pixel         

        # If the data already has a healpix pixel assignment then use it
        # Otherwise recalculate...
        try:
            data = fitsio.read(infile,columns=[pixel_pix_name])
        except ValueError as e:
            logger.info(str(e))
            columns=[self.config['catalog']['lon_field'],
                     self.config['catalog']['lat_field']]
            data = fitsio.read(infile,columns=columns)[columns]
            pix = ang2pix(self.nside_pixel,data[columns[0]],data[columns[1]])
            data = recfuncs.rec_append_fields(data,pixel_pix_name,pix)
            
        #mask_pixels = np.arange( hp.nside2npix(self.nside_mask), dtype='int')
        mask_maglims = np.zeros(hp.nside2npix(self.nside_mask))
         
        out_pixels = np.zeros(0,dtype='int')
        out_maglims = np.zeros(0)
         
        # Find the objects in each pixel
        pixel_pix = data[pixel_pix_name]
        mask_pix = ugali.utils.skymap.superpixel(pixel_pix,self.nside_pixel,self.nside_mask)
        count = Counter(mask_pix)
        pixels = sorted(count.keys())
        pix_digi = np.digitize(mask_pix,pixels).argsort()
        idx = 0
        min_num = 500
        signal_to_noise = 10.
        magerr_lim = 1/signal_to_noise
        for pix in pixels:
            # Calculate the magnitude limit in each pixel
            num = count[pix]
            objs = data[pix_digi[idx:idx+num]]
            idx += num
            if simple:
                # Set constant magnitude limits
                logger.debug("Simple magnitude limit for %s"%infile)
                mask_maglims[pix] = MAGLIMS[release][band]
            elif num < min_num:
                logger.info('Found <%i objects in pixel %i'%(min_num,pix))
                mask_maglims[pix] = 0
            else:
                mag = objs[mag_column]
                magerr = objs[magerr_column]
                # Estimate the magnitude limit as suggested by:
                # https://deswiki.cosmology.illinois.edu/confluence/display/DO/SVA1+Release+Document
                # (https://desweb.cosmology.illinois.edu/confluence/display/Operations/SVA1+Doc)
                maglim = np.median(mag[(magerr>0.9*magerr_lim)&(magerr<1.1*magerr_lim)])
         
                # Alternative method to estimate the magnitude limit by fitting median
                #mag_min, mag_max = mag.min(),mag.max()
                #mag_bins = np.arange(mag_min,mag_max,0.1) #0.1086?
                #x,y = ugali.utils.binning.binnedMedian(mag,magerr,mag_bins)
                #x,y = x[~np.isnan(y)],y[~np.isnan(y)]
                #magerr_med = interp1d(x,y)
                #mag0 = np.median(x) 
                #maglim = brentq(lambda a: magerr_med(a)-magerr_lim,x.min(),x.max(),disp=False)
                # Median from just objects near magerr cut
         
                mask_maglims[pix] = maglim

            logger.debug("%i (n=%i): maglim=%g"%(pix,num,mask_maglims[pix]))
            subpix = ugali.utils.skymap.subpixel(pix, self.nside_mask, self.nside_pixel)
            maglims = np.zeros(len(subpix)) + mask_maglims[pix] 
            out_pixels = np.append(out_pixels,subpix)
            out_maglims = np.append(out_maglims,maglims)
         
        # Remove empty pixels
        logger.info("Removing empty pixels")
        idx = np.nonzero(out_maglims > 0)[0]
        out_pixels  = out_pixels[idx]
        out_maglims = out_maglims[idx]
         
        # Remove pixels outside the footprint
        if self.footfile:
            logger.info("Checking footprint against %s"%self.footfile)
            lon,lat = pix2ang(self.nside_pixel,out_pixels)
            if self.config['coords']['coordsys'] == 'gal':
                ra,dec = gal2cel(lon,lat)
            else:    
                ra,dec = lon,lat
            footprint = inFootprint(self.footprint,ra,dec)
            idx = np.nonzero(footprint)[0]
            out_pixels = out_pixels[idx]
            out_maglims = out_maglims[idx]
         
        logger.info("MAGLIM = %.3f +/- %.3f"%(np.mean(out_maglims),np.std(out_maglims)))         
        return out_pixels,out_maglims
Example #51
0
def write_combined_file(
    target_file,
    output_dir='/Users/ptracy/Box Sync/00_postdoc_projects/Wind-STICS/zephyrus_transfer/output_dir/'
):

    STICS_data_product = 'DF'

    #Retrieve STICS data
    STICS_3D_data_format = np.dtype([('year', float), ('doy', float),
                                     ('sector', int), ('telescope', int),
                                     ('eoq', float), ('ion', 'S5'),
                                     (STICS_data_product, float),
                                     (STICS_data_product + '_error', float),
                                     ('delT', float)])
    STICS_data = np.loadtxt(target_file,
                            skiprows=4,
                            dtype=STICS_3D_data_format)
    #gather ion info
    m, q = ion_mq_stats(STICS_data['ion'][0])

    #Calculate the counts (from the stated error)
    counts_value = np.floor(
        np.round((STICS_data['DF'] / STICS_data['DF_error'])**2))
    counts_error_value = counts_value * (1.0 / np.sqrt(counts_value))

    #fix divide by zero errors
    small_num = 1e-10
    zero_ind = np.where((STICS_data['DF'] < small_num)
                        & (STICS_data['DF'] > -1.0 * small_num))
    counts_value[zero_ind] = 0.0
    counts_error_value[zero_ind] = 0.0

    #Calculate the differential flux (PSD[s^3/km^6]=1.076*m[amu]^2/(2*E[keV] ) * dJ[1/(cm^2*sr*s*keV)]
    dJ_value = STICS_data['DF'] / (
        1.076 * (m**2) / (2.0 * STICS_data['eoq'] * q))  #1/(cm^2*sr*s*keV)
    dJ_error_value = dJ_value * (1.0 / np.sqrt(counts_value))

    #fix divide by zero errors
    dJ_value[zero_ind] = 0.0
    dJ_error_value[zero_ind] = 0.0

    #append counts data to overall STICS_data
    #STICS_data=matplotlib.mlab.rec_append_fields(STICS_data, ('Counts', 'Counts_error','dJ', 'dJ_error'),
    #    (counts_value, counts_error_value, dJ_value, dJ_error_value))
    STICS_data = recfunctions.rec_append_fields(
        STICS_data, ('Counts', 'Counts_error', 'dJ', 'dJ_error'),
        (counts_value, counts_error_value, dJ_value,
         dJ_error_value))  #don't have matplotlib version on zephyrus (yet)

    temp = target_file.split('/')
    just_filename = temp[-1]  #remove directory info
    just_filename = just_filename.replace('wtdc', '')
    just_filename = just_filename.replace('DF', 'VDF')
    just_filename = just_filename.replace('_Lite', '')
    just_filename = just_filename.replace('_M-MOQ', '')
    just_filename = 'wstics_' + just_filename
    outfile1 = output_dir + just_filename

    header_string = ('year, doy, sector, telescope, eoq, ion, DF, DF_error,'
                     ' counts, counts_error, dJ, dJ_error, delT')
    #format_string='%7.2f %8.4f %2d %1d %8.4f %4s %8.2e %8.2e %8.2e %8.2e %8.2e %8.2e %8.2e'
    format_string = [
        '%7.2f', '%8.4f', '%2d', '%1d', '%8.4f', '%4s', '%8.2e', '%8.2e',
        '%8.2e', '%8.2e', '%8.2e', '%8.2e', '%8.2e'
    ]  #need to specify format in this way to use delimiter option in np.savetxt
    with open(outfile1, 'w') as file_handle:
        #np.savetxt(file_handle, STICS_data, header='Year, doyfrac, Telescope, Sec 0, Sec 1, Sec 2, Sec 3, '
        #    'Sec 4, Sec 5, Sec 6, Sec 7, Sec 8, Sec 9, Sec 10, Sec 11, Sec 12, Sec 13, Sec 14, Sec 15',
        #    fmt=('%d %f %d %e %e %e %e %e %e %e %e %e %e %e %e %e %e %e %e'))
        file_handle.write(
            '#  Wind / STICS 3D velocity distributions for ion ' +
            STICS_data['ion'][0] + '.\n')
        file_handle.write('# Created: ' + str(datetime.datetime.now()) + '\n')
        file_handle.write('# Modified to have C, dJ, and DF\n')
        np.savetxt(
            file_handle,
            STICS_data[[
                'year', 'doy', 'sector', 'telescope', 'eoq', 'ion', 'DF',
                'DF_error', 'Counts', 'Counts_error', 'dJ', 'dJ_error', 'delT'
            ]],  #fancy indexing of structured array
            header=header_string,
            fmt=format_string,
            delimiter='\t')

    #Remove old DF file after creation of the new one
    os.remove(target_file)  #remove file from directory after run

    print "Finished reprocessing wtdcLV2_Lite distribution file normally"  #Add this so I can search for it in the output files.
Example #52
0
modelnames = [name   for name,data,masses in modeldata]
datalist   = [data   for name,data,masses in modeldata]
masslists  = [masses for name,data,masses in modeldata]

# Also need the production cross sections
dataXS = np.loadtxt("data/TNGridxsSum.dat")
# The cross sections only vary with stop mass; make an interpolating function out of it.
uniqM,sorti = np.unique(dataXS[:,0],return_index=True) #get indices of elements with unique stop masses
XSifunc = interp.interp1d(uniqM,dataXS[sorti,2])

# Append to data arrays
newdatalist = []
for data in datalist:
   
   data = recfunctions.rec_append_fields(data,"XS",XSifunc(data['Mstop'])*1000) #convert units to fb

   # Append entries for signal yield predictions in each signal region
   Nprod_0lep = L_0lep*data["XS"] # Number of stop pairs produced (integrated luminosity * production cross section)
   Nprod_1lep = L_1lep*data["XS"]
   Nprod_2lep = L_2lep*data["XS"]
   Nprod_2b   = L_2b*data["XS"]  
   
   # Fraction of stop pairs that decayed into events ID'd as top + neutralino
   #  =(e.g.) numSR1/N_0lep (number of events counted / total number of stop pairs simulated)
   
   # Mean numbers of signal events predicted in the various signal regions:
   # 0 lep:
   SR_sdata_0lep = ["SR1_s","SR2_s","SR3_s"]
   data = recfunctions.rec_append_fields(data,\
            SR_sdata_0lep,\