<s> import sys <EOL> import numpy as np <EOL> from . . column import Column , ColDefs , FITS2NUMPY <EOL> from . . fitsrec import FITS_rec , FITS_record <EOL> from . . util import lazyproperty , _is_int , _is_pseudo_unsigned , _unsigned_zero <EOL> from . base import DTYPE2BITPIX <EOL> from . image import PrimaryHDU <EOL> from . table import _TableLikeHDU <EOL> class Group ( FITS_record ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , input , row = <NUM_LIT:0> , start = None , end = None , step = None , <EOL> base = None ) : <EOL> super ( Group , self ) . __init__ ( input , row , start , end , step , base ) <EOL> @ property <EOL> def parnames ( self ) : <EOL> return self . array . parnames <EOL> @ property <EOL> def data ( self ) : <EOL> return self . field ( self . array . _coldefs . names [ - <NUM_LIT:1> ] ) <EOL> @ lazyproperty <EOL> def _unique ( self ) : <EOL> return _par_indices ( self . parnames ) <EOL> def par ( self , parname ) : <EOL> """<STR_LIT>""" <EOL> if _is_int ( parname ) : <EOL> result = self . array [ self . row ] [ parname ] <EOL> else : <EOL> indx = self . _unique [ parname . upper ( ) ] <EOL> if len ( indx ) == <NUM_LIT:1> : <EOL> result = self . array [ self . row ] [ indx [ <NUM_LIT:0> ] ] <EOL> else : <EOL> result = self . array [ self . row ] [ indx [ <NUM_LIT:0> ] ] . astype ( '<STR_LIT>' ) <EOL> for i in indx [ <NUM_LIT:1> : ] : <EOL> result += self . array [ self . row ] [ i ] <EOL> return result <EOL> def setpar ( self , parname , value ) : <EOL> """<STR_LIT>""" <EOL> if _is_int ( parname ) : <EOL> self . array [ self . row ] [ parname ] = value <EOL> else : <EOL> indx = self . _unique [ parname . upper ( ) ] <EOL> if len ( indx ) == <NUM_LIT:1> : <EOL> self . array [ self . row ] [ indx [ <NUM_LIT:0> ] ] = value <EOL> else : <EOL> if isinstance ( value , ( list , tuple ) ) and len ( indx ) == len ( value ) : <EOL> for i in range ( len ( indx ) ) : <EOL> self . array [ self . row ] [ indx [ i ] ] = value [ i ] <EOL> else : <EOL> raise ValueError ( '<STR_LIT>' <EOL> '<STR_LIT>' % len ( indx ) ) <EOL> class GroupData ( FITS_rec ) : <EOL> """<STR_LIT>""" <EOL> _record_type = Group <EOL> def __new__ ( cls , input = None , bitpix = None , pardata = None , parnames = [ ] , <EOL> bscale = None , bzero = None , parbscales = None , parbzeros = None ) : <EOL> """<STR_LIT>""" <EOL> if not isinstance ( input , FITS_rec ) : <EOL> if pardata is None : <EOL> npars = <NUM_LIT:0> <EOL> else : <EOL> npars = len ( pardata ) <EOL> if parbscales is None : <EOL> parbscales = [ None ] * npars <EOL> if parbzeros is None : <EOL> parbzeros = [ None ] * npars <EOL> if parnames is None : <EOL> parnames = [ '<STR_LIT>' % ( idx + <NUM_LIT:1> ) for idx in range ( npars ) ] <EOL> if len ( parnames ) != npars : <EOL> raise ValueError ( '<STR_LIT>' <EOL> '<STR_LIT>' ) <EOL> unique_parnames = _unique_parnames ( parnames + [ '<STR_LIT>' ] ) <EOL> if bitpix is None : <EOL> bitpix = DTYPE2BITPIX [ input . dtype . name ] <EOL> fits_fmt = GroupsHDU . _bitpix2tform [ bitpix ] <EOL> format = FITS2NUMPY [ fits_fmt ] <EOL> data_fmt = '<STR_LIT>' % ( str ( input . shape [ <NUM_LIT:1> : ] ) , format ) <EOL> formats = '<STR_LIT:U+002C>' . join ( ( [ format ] * npars ) + [ data_fmt ] ) <EOL> gcount = input . shape [ <NUM_LIT:0> ] <EOL> cols = [ Column ( name = unique_parnames [ idx ] , format = fits_fmt , <EOL> bscale = parbscales [ idx ] , bzero = parbzeros [ idx ] ) <EOL> for idx in range ( npars ) ] <EOL> cols . append ( Column ( name = unique_parnames [ - <NUM_LIT:1> ] , format = fits_fmt , <EOL> bscale = bscale , bzero = bzero ) ) <EOL> coldefs = ColDefs ( cols ) <EOL> self = FITS_rec . __new__ ( cls , <EOL> np . rec . array ( None , <EOL> formats = formats , <EOL> names = coldefs . names , <EOL> shape = gcount ) ) <EOL> self . _data_field = unique_parnames [ - <NUM_LIT:1> ] <EOL> self . _coldefs = coldefs <EOL> self . parnames = parnames <EOL> for idx , name in enumerate ( unique_parnames [ : - <NUM_LIT:1> ] ) : <EOL> column = coldefs [ idx ] <EOL> scale , zero = self . _get_scale_factors ( column ) [ <NUM_LIT:3> : <NUM_LIT:5> ] <EOL> if scale or zero : <EOL> self . _cache_field ( name , pardata [ idx ] ) <EOL> else : <EOL> np . rec . recarray . field ( self , idx ) [ : ] = pardata [ idx ] <EOL> column = coldefs [ self . _data_field ] <EOL> scale , zero = self . _get_scale_factors ( column ) [ <NUM_LIT:3> : <NUM_LIT:5> ] <EOL> if scale or zero : <EOL> self . _cache_field ( self . _data_field , input ) <EOL> else : <EOL> np . rec . recarray . field ( self , npars ) [ : ] = input <EOL> else : <EOL> self = FITS_rec . __new__ ( cls , input ) <EOL> self . parnames = None <EOL> return self <EOL> def __array_finalize__ ( self , obj ) : <EOL> super ( GroupData , self ) . __array_finalize__ ( obj ) <EOL> if isinstance ( obj , GroupData ) : <EOL> self . parnames = obj . parnames <EOL> elif isinstance ( obj , FITS_rec ) : <EOL> self . parnames = obj . _coldefs . names <EOL> def __getitem__ ( self , key ) : <EOL> out = super ( GroupData , self ) . __getitem__ ( key ) <EOL> if isinstance ( out , GroupData ) : <EOL> out . parnames = self . parnames <EOL> return out <EOL> @ property <EOL> def data ( self ) : <EOL> """<STR_LIT>""" <EOL> return self . field ( self . _coldefs . names [ - <NUM_LIT:1> ] ) <EOL> @ lazyproperty <EOL> def _unique ( self ) : <EOL> return _par_indices ( self . parnames ) <EOL> def par ( self , parname ) : <EOL> """<STR_LIT>""" <EOL> if _is_int ( parname ) : <EOL> result = self . field ( parname ) <EOL> else : <EOL> indx = self . _unique [ parname . upper ( ) ] <EOL> if len ( indx ) == <NUM_LIT:1> : <EOL> result = self . field ( indx [ <NUM_LIT:0> ] ) <EOL> else : <EOL> result = self . field ( indx [ <NUM_LIT:0> ] ) . astype ( '<STR_LIT>' ) <EOL> for i in indx [ <NUM_LIT:1> : ] : <EOL> result += self . field ( i ) <EOL> return result <EOL> class GroupsHDU ( PrimaryHDU , _TableLikeHDU ) : <EOL> """<STR_LIT>""" <EOL> _bitpix2tform = { <NUM_LIT:8> : '<STR_LIT:B>' , <NUM_LIT:16> : '<STR_LIT:I>' , <NUM_LIT:32> : '<STR_LIT>' , <NUM_LIT:64> : '<STR_LIT>' , - <NUM_LIT:32> : '<STR_LIT:E>' , - <NUM_LIT:64> : '<STR_LIT:D>' } <EOL> _data_type = GroupData <EOL> _data_field = '<STR_LIT>' <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , data = None , header = None ) : <EOL> super ( GroupsHDU , self ) . __init__ ( data = data , header = header ) <EOL> if len ( self . _axes ) <= <NUM_LIT:0> : <EOL> self . _axes = [ <NUM_LIT:0> ] <EOL> self . _header [ '<STR_LIT>' ] = <NUM_LIT:1> <EOL> self . _header . set ( '<STR_LIT>' , <NUM_LIT:0> , after = '<STR_LIT>' ) <EOL> @ classmethod <EOL> def match_header ( cls , header ) : <EOL> keyword = header . cards [ <NUM_LIT:0> ] . keyword <EOL> return ( keyword == '<STR_LIT>' and '<STR_LIT>' in header and <EOL> header [ '<STR_LIT>' ] == True ) <EOL> @ lazyproperty <EOL> def data ( self ) : <EOL> """<STR_LIT>""" <EOL> data = self . _get_tbdata ( ) <EOL> data . _coldefs = self . columns <EOL> data . parnames = self . parnames <EOL> del self . columns <EOL> return data <EOL> @ lazyproperty <EOL> def parnames ( self ) : <EOL> """<STR_LIT>""" <EOL> pcount = self . _header [ '<STR_LIT>' ] <EOL> return [ self . _header [ '<STR_LIT>' + str ( idx + <NUM_LIT:1> ) ] for idx in range ( pcount ) ] <EOL> @ lazyproperty <EOL> def columns ( self ) : <EOL> if self . _has_data and hasattr ( self . data , '<STR_LIT>' ) : <EOL> return self . data . _coldefs <EOL> format = self . _bitpix2tform [ self . _header [ '<STR_LIT>' ] ] <EOL> pcount = self . _header [ '<STR_LIT>' ] <EOL> parnames = [ ] <EOL> bscales = [ ] <EOL> bzeros = [ ] <EOL> for idx in range ( pcount ) : <EOL> bscales . append ( self . _header . get ( '<STR_LIT>' + str ( idx + <NUM_LIT:1> ) , None ) ) <EOL> bzeros . append ( self . _header . get ( '<STR_LIT>' + str ( idx + <NUM_LIT:1> ) , None ) ) <EOL> parnames . append ( self . _header [ '<STR_LIT>' + str ( idx + <NUM_LIT:1> ) ] ) <EOL> formats = [ format ] * len ( parnames ) <EOL> dim = [ None ] * len ( parnames ) <EOL> parnames . append ( '<STR_LIT>' ) <EOL> bscales . append ( self . _header . get ( '<STR_LIT>' ) ) <EOL> bzeros . append ( self . _header . get ( '<STR_LIT>' ) ) <EOL> data_shape = self . shape [ : - <NUM_LIT:1> ] <EOL> formats . append ( str ( int ( np . prod ( data_shape ) ) ) + format ) <EOL> dim . append ( data_shape ) <EOL> parnames = _unique_parnames ( parnames ) <EOL> self . _data_field = parnames [ - <NUM_LIT:1> ] <EOL> cols = [ Column ( name = name , format = fmt , bscale = bscale , bzero = bzero , <EOL> dim = dim ) <EOL> for name , fmt , bscale , bzero , dim in <EOL> zip ( parnames , formats , bscales , bzeros , dim ) ] <EOL> coldefs = ColDefs ( cols ) <EOL> return coldefs <EOL> @ property <EOL> def _nrows ( self ) : <EOL> if not self . _data_loaded : <EOL> return self . _header . get ( '<STR_LIT>' , <NUM_LIT:0> ) <EOL> else : <EOL> return len ( self . data ) <EOL> @ lazyproperty <EOL> def _theap ( self ) : <EOL> return <NUM_LIT:0> <EOL> @ property <EOL> def size ( self ) : <EOL> """<STR_LIT>""" <EOL> size = <NUM_LIT:0> <EOL> naxis = self . _header . get ( '<STR_LIT>' , <NUM_LIT:0> ) <EOL> if naxis > <NUM_LIT:1> : <EOL> size = <NUM_LIT:1> <EOL> for idx in range ( <NUM_LIT:1> , naxis ) : <EOL> size = size * self . _header [ '<STR_LIT>' + str ( idx + <NUM_LIT:1> ) ] <EOL> bitpix = self . _header [ '<STR_LIT>' ] <EOL> gcount = self . _header . get ( '<STR_LIT>' , <NUM_LIT:1> ) <EOL> pcount = self . _header . get ( '<STR_LIT>' , <NUM_LIT:0> ) <EOL> size = abs ( bitpix ) * gcount * ( pcount + size ) // <NUM_LIT:8> <EOL> return size <EOL> def update_header ( self ) : <EOL> old_naxis = self . _header . get ( '<STR_LIT>' , <NUM_LIT:0> ) <EOL> if self . _data_loaded : <EOL> if isinstance ( self . data , GroupData ) : <EOL> self . _axes = list ( self . data . data . shape ) [ <NUM_LIT:1> : ] <EOL> self . _axes . reverse ( ) <EOL> self . _axes = [ <NUM_LIT:0> ] + self . _axes <EOL> field0 = self . data . dtype . names [ <NUM_LIT:0> ] <EOL> field0_code = self . data . dtype . fields [ field0 ] [ <NUM_LIT:0> ] . name <EOL> elif self . data is None : <EOL> self . _axes = [ <NUM_LIT:0> ] <EOL> field0_code = '<STR_LIT>' <EOL> else : <EOL> raise ValueError ( '<STR_LIT>' ) <EOL> self . _header [ '<STR_LIT>' ] = DTYPE2BITPIX [ field0_code ] <EOL> self . _header [ '<STR_LIT>' ] = len ( self . _axes ) <EOL> for idx , axis in enumerate ( self . _axes ) : <EOL> if ( idx == <NUM_LIT:0> ) : <EOL> after = '<STR_LIT>' <EOL> else : <EOL> after = '<STR_LIT>' + str ( idx ) <EOL> self . _header . set ( '<STR_LIT>' + str ( idx + <NUM_LIT:1> ) , axis , after = after ) <EOL> for idx in range ( len ( self . _axes ) + <NUM_LIT:1> , old_naxis + <NUM_LIT:1> ) : <EOL> try : <EOL> del self . _header [ '<STR_LIT>' + str ( idx ) ] <EOL> except KeyError : <EOL> pass <EOL> if self . _has_data and isinstance ( self . data , GroupData ) : <EOL> self . _header . set ( '<STR_LIT>' , True , <EOL> after = '<STR_LIT>' + str ( len ( self . _axes ) ) ) <EOL> self . _header . set ( '<STR_LIT>' , len ( self . data . parnames ) , after = '<STR_LIT>' ) <EOL> self . _header . set ( '<STR_LIT>' , len ( self . data ) , after = '<STR_LIT>' ) <EOL> column = self . data . _coldefs [ self . _data_field ] <EOL> scale , zero = self . data . _get_scale_factors ( column ) [ <NUM_LIT:3> : <NUM_LIT:5> ] <EOL> if scale : <EOL> self . _header . set ( '<STR_LIT>' , column . bscale ) <EOL> if zero : <EOL> self . _header . set ( '<STR_LIT>' , column . bzero ) <EOL> for idx , name in enumerate ( self . data . parnames ) : <EOL> self . _header . set ( '<STR_LIT>' + str ( idx + <NUM_LIT:1> ) , name ) <EOL> column = self . data . _coldefs [ idx ] <EOL> scale , zero = self . data . _get_scale_factors ( column ) [ <NUM_LIT:3> : <NUM_LIT:5> ] <EOL> if scale : <EOL> self . _header . set ( '<STR_LIT>' + str ( idx + <NUM_LIT:1> ) , column . bscale ) <EOL> if zero : <EOL> self . _header . set ( '<STR_LIT>' + str ( idx + <NUM_LIT:1> ) , column . bzero ) <EOL> if '<STR_LIT>' in self . _header : <EOL> if len ( self . _axes ) : <EOL> after = '<STR_LIT>' + str ( len ( self . _axes ) ) <EOL> else : <EOL> after = '<STR_LIT>' <EOL> self . _header . set ( '<STR_LIT>' , after = after ) <EOL> def _writedata_internal ( self , fileobj ) : <EOL> """<STR_LIT>""" <EOL> size = <NUM_LIT:0> <EOL> if self . data is not None : <EOL> self . data . _scale_back ( ) <EOL> if sys . byteorder == '<STR_LIT>' : <EOL> swap_types = ( '<STR_LIT:<>' , '<STR_LIT:=>' ) <EOL> else : <EOL> swap_types = ( '<STR_LIT:<>' , ) <EOL> if _is_pseudo_unsigned ( self . data . dtype ) : <EOL> output = np . array ( <EOL> self . data - _unsigned_zero ( self . data . dtype ) , <EOL> dtype = '<STR_LIT>' % self . data . dtype . itemsize ) <EOL> should_swap = False <EOL> else : <EOL> output = self . data <EOL> fname = self . data . dtype . names [ <NUM_LIT:0> ] <EOL> byteorder = self . data . dtype . fields [ fname ] [ <NUM_LIT:0> ] . str [ <NUM_LIT:0> ] <EOL> should_swap = ( byteorder in swap_types ) <EOL> if not fileobj . simulateonly : <EOL> if should_swap : <EOL> output . byteswap ( True ) <EOL> try : <EOL> fileobj . writearray ( output ) <EOL> finally : <EOL> output . byteswap ( True ) <EOL> else : <EOL> fileobj . writearray ( output ) <EOL> size += output . size * output . itemsize <EOL> return size <EOL> def _verify ( self , option = '<STR_LIT>' ) : <EOL> errs = super ( GroupsHDU , self ) . _verify ( option = option ) <EOL> self . req_cards ( '<STR_LIT>' , <NUM_LIT:2> , <EOL> lambda v : ( _is_int ( v ) and v >= <NUM_LIT:1> and v <= <NUM_LIT> ) , <NUM_LIT:1> , <EOL> option , errs ) <EOL> self . req_cards ( '<STR_LIT>' , <NUM_LIT:3> , lambda v : ( _is_int ( v ) and v == <NUM_LIT:0> ) , <NUM_LIT:0> , <EOL> option , errs ) <EOL> after = self . _header [ '<STR_LIT>' ] + <NUM_LIT:3> <EOL> pos = lambda x : x >= after <EOL> self . req_cards ( '<STR_LIT>' , pos , _is_int , <NUM_LIT:1> , option , errs ) <EOL> self . req_cards ( '<STR_LIT>' , pos , _is_int , <NUM_LIT:0> , option , errs ) <EOL> self . req_cards ( '<STR_LIT>' , pos , lambda v : ( v == True ) , True , option , <EOL> errs ) <EOL> return errs <EOL> def _calculate_datasum ( self , blocking ) : <EOL> """<STR_LIT>""" <EOL> if self . _has_data : <EOL> byteorder = self . data . dtype . fields [ self . data . dtype . names [ <NUM_LIT:0> ] ] [ <NUM_LIT:0> ] . str [ <NUM_LIT:0> ] <EOL> if byteorder != '<STR_LIT:>>' : <EOL> byteswapped = True <EOL> d = self . data . byteswap ( True ) <EOL> d . dtype = d . dtype . newbyteorder ( '<STR_LIT:>>' ) <EOL> else : <EOL> byteswapped = False <EOL> d = self . data <EOL> byte_data = d . view ( type = np . ndarray , dtype = np . ubyte ) <EOL> cs = self . _compute_checksum ( byte_data , blocking = blocking ) <EOL> if byteswapped : <EOL> d . byteswap ( True ) <EOL> d . dtype = d . dtype . newbyteorder ( '<STR_LIT:<>' ) <EOL> return cs <EOL> else : <EOL> return super ( GroupsHDU , self ) . _calculate_datasum ( blocking = blocking ) <EOL> def _summary ( self ) : <EOL> summary = super ( GroupsHDU , self ) . _summary ( ) <EOL> name , classname , length , shape , format , gcount = summary <EOL> if shape : <EOL> shape = shape [ <NUM_LIT:1> : ] <EOL> if shape and all ( shape ) : <EOL> format = self . columns [ <NUM_LIT:0> ] . dtype . name <EOL> gcount = '<STR_LIT>' % ( self . _gcount , self . _pcount ) <EOL> return ( name , classname , length , shape , format , gcount ) <EOL> def _par_indices ( names ) : <EOL> """<STR_LIT>""" <EOL> unique = { } <EOL> for idx , name in enumerate ( names ) : <EOL> name = name . upper ( ) <EOL> if name in unique : <EOL> unique [ name ] . append ( idx ) <EOL> else : <EOL> unique [ name ] = [ idx ] <EOL> return unique <EOL> def _unique_parnames ( names ) : <EOL> """<STR_LIT>""" <EOL> upper_names = set ( ) <EOL> unique_names = [ ] <EOL> for name in names : <EOL> name_upper = name . upper ( ) <EOL> while name_upper in upper_names : <EOL> name = '<STR_LIT:_>' + name <EOL> name_upper = '<STR_LIT:_>' + name_upper <EOL> unique_names . append ( name ) <EOL> upper_names . add ( name_upper ) <EOL> return unique_names </s>
