Package gavo :: Package formats :: Module fitstable
[frames] | no frames]

Source Code for Module gavo.formats.fitstable

  1  """ 
  2  Writing data in FITS binary tables 
  3  """ 
  4   
  5  #c Copyright 2008-2019, the GAVO project 
  6  #c 
  7  #c This program is free software, covered by the GNU GPL.  See the 
  8  #c COPYING file in the source distribution. 
  9   
 10   
 11  import os 
 12  import tempfile 
 13  import time 
 14   
 15  import numpy 
 16   
 17  from gavo import base 
 18  from gavo import rsc 
 19  from gavo import utils 
 20  from gavo.formats import common 
 21  from gavo.utils import pyfits 
 22  from gavo.votable import common as votcommon 
 23   
 24   
 25  _fitsCodeMap = { 
 26          "short": "I", 
 27          "int": "J", 
 28          "long": "K", 
 29          "float": "E", 
 30          "double": "D", 
 31          "boolean": "L", 
 32          "char": "A", 
 33          "unicodeChar": "A", 
 34  } 
 35   
 36  _typeConstructors = { 
 37          "short": int, 
 38          "int": int, 
 39          "long": int, 
 40          "float": float, 
 41          "double": float, 
 42          "boolean": int, 
 43          "char": str, 
 44          "unicodeChar": str, 
 45  } 
 46   
 47   
48 -def _makeStringArray(values, colInd, colDesc):
49 """returns a pyfits-capable column array for strings stored in the colInd-th 50 column of values. 51 """ 52 try: 53 arr = numpy.array([str(v[colInd]) for v in values], dtype=numpy.str) 54 except UnicodeEncodeError: 55 56 def _(s): 57 if s is None: 58 return "" 59 else: 60 return s.encode("utf-8")
61 62 arr = numpy.array([_(v[colInd]) for v in values], dtype=numpy.str) 63 return "%dA"%arr.itemsize, arr 64 65
66 -def _getNullValue(colDesc):
67 """returns a null value we consider ok for a column described by colDesc. 68 69 This is supposed to be in the column data type. 70 """ 71 if votcommon.getLength(colDesc["arraysize"])>1: 72 # TODO: figure out nullvalue interpretation in FITS arrays. 73 return None 74 75 nullValue = colDesc["nullvalue"] 76 if nullValue is None: 77 # enter some reasonable defaults 78 if (colDesc["datatype"]=="float" 79 or colDesc["datatype"]=="double"): 80 nullValue = float("NaN") 81 elif colDesc["datatype"]=="text": 82 nullValue = "" 83 else: 84 nullValue = _typeConstructors[colDesc["datatype"]](nullValue) 85 86 return nullValue
87 88
89 -def _computeTypeCode(colDesc):
90 """returns a FITS type code suitable for colDesc. 91 92 This is deficient in that we don't do variable length arrays 93 and the like for now. This may need extension when we want 94 to get fancy with FITS tables. 95 """ 96 baseCode = _fitsCodeMap[colDesc["datatype"]] 97 length = votcommon.getLength(colDesc["arraysize"]) 98 99 if length==1: 100 return length, baseCode 101 102 elif length is None: 103 # NOTE: Strings have special handling in _makeExtension. 104 # Here, we just bail out for now. 105 # There actually is support for variable-length arrays 106 # in FITS. Let's see if we need it at some point 107 raise common.CannotSerializeIn("fits") 108 109 else: 110 return length, "%d%s"%(length, baseCode)
111 112
113 -def _makeValueArray(values, colInd, colDesc):
114 """returns a pyfits-capable column array for non-string values 115 stored in the colInd-th column of values. 116 """ 117 length, typecode = _computeTypeCode(colDesc) 118 nullValue = _getNullValue(colDesc) 119 nan = float("NaN") 120 121 if length>1: 122 # numpy 1:1.12.1-3 and pyfits from astropy 1.3-8 apparently 123 # have trouble with masked arrays here, so we exploit 124 # that in the VO and for floats NaN and NULL is treated equivalently. 125 # That doesn't help for ints. Let's hope we can go back to masked 126 # arrays here until then. 127 nullArray = numpy.array([nan]*length) 128 def mkval(v): 129 if v is None: 130 return nullArray 131 else: 132 return numpy.array(v)
133 134 else: 135 def mkval(v): 136 if v is None: 137 if nullValue is None: 138 raise ValueError("While serializing a FITS table: NULL" 139 " detected in column '%s' but no null value declared"% 140 colDesc["name"]) 141 return nullValue 142 else: 143 return v 144 145 arr = numpy.array([mkval(v[colInd]) for v in values]) 146 return typecode, arr 147 148
149 -def _makeExtension(serMan):
150 """returns a pyfits hdu for the valuemappers.SerManager instance table. 151 """ 152 values = list(serMan.getMappedTuples()) 153 columns = [] 154 utypes = [] 155 descriptions = [] 156 157 for colInd, colDesc in enumerate(serMan): 158 descriptions.append(colDesc["description"]) 159 if colDesc["datatype"]=="char" or colDesc["datatype"]=="unicodeChar": 160 makeArray = _makeStringArray 161 else: 162 makeArray = _makeValueArray 163 164 typecode, arr = makeArray(values, colInd, colDesc) 165 if typecode[-1] in 'ED': 166 nullValue = None # (NaN implied) 167 else: 168 nullValue = _getNullValue(colDesc) 169 170 columns.append(pyfits.Column(name=str(colDesc["name"]), 171 unit=str(colDesc["unit"]), format=typecode, 172 null=nullValue, array=arr)) 173 if colDesc["utype"]: 174 utypes.append((colInd, str(colDesc["utype"].lower()))) 175 176 hdu = pyfits.BinTableHDU.from_columns(pyfits.ColDefs(columns)) 177 for colInd, utype in utypes: 178 hdu.header.set("TUTYP%d"%(colInd+1), utype) 179 180 cards = hdu.header.cards 181 for colInd, desc in enumerate(descriptions): 182 cards["TTYPE%d"%(colInd+1)].comment = desc.encode("ascii", "ignore") 183 hdu.header.set("TCOMM%d"%(colInd+1), desc.encode("ascii", "ignore")) 184 185 if not hasattr(serMan.table, "IgnoreTableParams"): 186 for param in serMan.table.iterParams(): 187 if param.value is None: 188 continue 189 190 key, value, comment = str(param.name), param.value, param.description 191 if isinstance(value, unicode): 192 value = value.encode('ascii', "xmlcharrefreplace") 193 if isinstance(comment, unicode): 194 comment = comment.encode('ascii', "xmlcharrefreplace") 195 if len(key)>8: 196 key = "hierarch "+key 197 198 try: 199 hdu.header.set(key, value=value, comment=comment) 200 except ValueError as ex: 201 # do not fail just because some header couldn't be serialised 202 base.ui.notifyWarning( 203 "Failed to serialise param %s to a FITS header (%s)"%( 204 param.name, 205 utils.safe_str(ex))) 206 207 return hdu
208 209
210 -def _makeFITSTableNOLOCK(dataSet, acquireSamples=True):
211 """returns a hdulist containing extensions for the tables in dataSet. 212 213 You must make sure that this function is only executed once 214 since pyfits is not thread-safe. 215 """ 216 tables = [base.SerManager(table, acquireSamples=acquireSamples) 217 for table in dataSet.tables.values()] 218 extensions = [_makeExtension(table) for table in tables] 219 primary = pyfits.PrimaryHDU() 220 primary.header.set("DATE", time.strftime("%Y-%m-%d"), 221 "Date file was written") 222 return pyfits.HDUList([primary]+extensions)
223 224
225 -def makeFITSTable(dataSet, acquireSamples=False):
226 """returns a hdulist containing extensions for the tables in dataSet. 227 228 This function may block basically forever. Never call this from 229 the main server, always use threads or separate processes (until 230 pyfits is fixed to be thread-safe). 231 232 This will add table parameters as header cards on the resulting FITS 233 header. 234 """ 235 with utils.fitsLock(): 236 return _makeFITSTableNOLOCK(dataSet, acquireSamples)
237 238
239 -def writeFITSTableFile(hdulist):
240 """returns the name of a temporary file containing the FITS data for 241 hdulist. 242 """ 243 handle, pathname = tempfile.mkstemp(".fits", prefix="fitstable", 244 dir=base.getConfig("tempDir")) 245 246 # if there's more than the primary HDU, EXTEND=True is mandatory; let's 247 # be defensive here 248 if len(hdulist)>1: 249 hdulist[0].header.set("EXTEND", True, "More exts following") 250 251 with utils.silence(): 252 hdulist.writeto(pathname, clobber=1) 253 os.close(handle) 254 return pathname
255 256
257 -def makeFITSTableFile(dataSet, acquireSamples=True):
258 """returns the name of a temporary file containing a fits file 259 representing dataSet. 260 261 The caller is responsible to remove the file. 262 """ 263 hdulist = makeFITSTable(dataSet, acquireSamples) 264 return writeFITSTableFile(hdulist)
265 266
267 -def writeDataAsFITS(data, outputFile, acquireSamples=False):
268 """a formats.common compliant data writer. 269 270 This will write out table params as header cards. To serialise 271 those yourself (as is required for spectral data model compliant 272 tables), set an attribute IgnoreTableParams (with an arbitrary 273 value) on the table. 274 """ 275 data = rsc.wrapTable(data) 276 fitsName = makeFITSTableFile(data, acquireSamples) 277 try: 278 src = open(fitsName) 279 utils.cat(src, outputFile) 280 src.close() 281 finally: 282 os.unlink(fitsName)
283 284 common.registerDataWriter("fits", writeDataAsFITS, "application/fits", 285 "FITS Binary Table", ".fits") 286