Package csb :: Package bio :: Package io :: Module mrc |
[frames] | no frames]

# Source Code for Module csb.bio.io.mrc

      1  """
      2  Cryo-EM density map I/O
      3
      4  @warning: dragons ahead, this module is experimental
      5  """
      6
      7  import numpy
      8  import struct
      9
     10
     11 -class DensityMapFormatError(ValueError):
     12      pass
     13
     14
     15 -class ByteOrder(object):
     16
     17      NATIVE = '='
     18      LITTLE = '<'
     19      BIG = '>'
     20
     21
     22 -class DensityInfo(object):
     23
     24 -    def __init__(self, data, spacing, origin, shape=None, header=None, axes=None):
     25
     26          self.data = data
     27          self.spacing = spacing
     28          self.origin = origin
     29          self.header = header
     30          self.shape = shape
     31          self.axes = axes
     32
     33          if shape is None and data is not None:
     34              self.shape = self.data.shape
     35
     36 -class HeaderInfo(object):
     37
     38 -    def __init__(self, fields):
     39
     40          fields = tuple(fields)
     41          if not len(fields) == 25:
     42              raise ValueError(fields)
     43
     44          self._fields = fields
     45
     46 -    def __getitem__(self, i):
     47          return self._fields[i]
     48
     49 -    def __iter__(self):
     50          return iter(self._fields)
     51
     52      @property
     53 -    def nc(self):
     54          return self._fields[0]
     55      @property
     56 -    def nr(self):
     57          return self._fields[1]
     58      @property
     59 -    def ns(self):
     60          return self._fields[2]
     61      @property
     62 -    def mode(self):
     63          return self._fields[3]
     64      @property
     65 -    def ncstart(self):
     66          return self._fields[4]
     67      @property
     68 -    def nrstart(self):
     69          return self._fields[5]
     70      @property
     71 -    def nsstart(self):
     72          return self._fields[6]
     73      @property
     74 -    def nx(self):
     75          return self._fields[7]
     76      @property
     77 -    def ny(self):
     78          return self._fields[8]
     79      @property
     80 -    def nz(self):
     81          return self._fields[9]
     82      @property
     83 -    def x(self):
     84          return self._fields[10]
     85      @property
     86 -    def y(self):
     87          return self._fields[11]
     88      @property
     89 -    def z(self):
     90          return self._fields[12]
     91      @property
     92 -    def alpha(self):
     93          return self._fields[13]
     94      @property
     95 -    def beta(self):
     96          return self._fields[14]
     97      @property
     98 -    def gamma(self):
     99          return self._fields[15]
    100      @property
    101 -    def mapc(self):
    102          return self._fields[16]
    103      @property
    104 -    def mapr(self):
    105          return self._fields[17]
    106      @property
    107 -    def maps(self):
    108          return self._fields[18]
    109      @property
    110 -    def amin(self):
    111          return self._fields[19]
    112      @property
    113 -    def amax(self):
    114          return self._fields[20]
    115      @property
    116 -    def amean(self):
    117          return self._fields[21]
    118      @property
    119 -    def ispg(self):
    120          return self._fields[22]
    121      @property
    122 -    def nsymbt(self):
    123          return self._fields[23]
    124      @property
    125 -    def lskflg(self):
    126          return self._fields[24]
    127
    128
    129 -class DensityMapReader(object):
    130      """
    131      Binary MRC density map reader.
    132
    133      @param filename: input MRC file name
    134      @type filename: str
    135      """
    136
    137      HEADER_SIZE = 1024
    138
    139 -    def __init__(self, filename):
    140          self._filename = filename
    141
    142      @property
    143 -    def filename(self):
    144          """
    145          Input MRC file name
    146          @rtype: str
    147          """
    148          return self._filename
    149
    150 -    def _rawheader(self, stream):
    151          """
    152          Read and return the raw binary header.
    153          """
    154          raw = stream.read(DensityMapReader.HEADER_SIZE)
    155          return bytes(raw)
    156
    157 -    def _inspect(self, rawheader, order):
    158          """
    159          Parse a raw binary header.
    160          """
    161          format = '{0}10l6f3l3f3l'.format(order)
    162          fields = struct.unpack(format, rawheader[:4 * 25])
    163
    164          return HeaderInfo(fields)
    165
    166 -    def _spacing(self, header):
    167
    168          if header.nx != 0 and header.ny != 0 and header.nz != 0:
    169              return (header.x / header.nx, header.y / header.ny, header.z / header.nz)
    170          else:
    171              return (0, 0, 0)
    172
    173 -    def _origin(self, header, spacing=None):
    174
    175          if spacing is None:
    176              spacing = self._spacing(header)
    177
    178          origin = header.ncstart, header.nrstart, header.nsstart
    179          return [origin[i] * spacing[i] for i in range(3)]
    180
    181 -    def _shape(self, header):
    182          return (header.ns, header.nr, header.nc)
    183
    184 -    def inspect_header(self, order=ByteOrder.NATIVE):
    185          """
    186          Parse the raw binary header of the density map.
    187
    188          @param order: byte order (defaults to L{ByteOrder.NATIVE})
    189          @type order: str
    190
    191          @return: header information
    192          @rtype: L{HeaderInfo}
    193          """
    194          with open(self.filename, 'rb') as stream:
    195              raw = self._rawheader(stream)
    196              return self._inspect(raw, order)
    197
    198 -    def read_header(self):
    199          """
    200          Read the header of the density map only.
    201
    202          @return: density info without any actual data (density.data is None)
    203          @rtype: L{DensityInfo}
    204          """
    205          with open(self.filename, 'rb') as stream:
    206
    207              raw = self._rawheader(stream)
    208
    209              header = self._inspect(raw, ByteOrder.NATIVE)
    210              spacing = self._spacing(header)
    211              origin = self._origin(header, spacing)
    212              shape = self._shape(header)
    213
    214              return DensityInfo(None, spacing, origin, shape=shape, header=raw)
    215
    216 -    def read(self):
    217          """
    218          Read the entire density map.
    219
    220          @return: complete density info
    221          @rtype: L{DensityInfo}
    222          """
    223          with open(self.filename, 'rb') as stream:
    224
    225              raw = self._rawheader(stream)
    226              header = self._inspect(raw, ByteOrder.NATIVE)
    227
    228              if header.mode == 2 or header.mode == 1:
    229                  byte_order = ByteOrder.NATIVE
    230
    231              elif header.mode == 33554432:
    232                  header = self._inspect(raw, ByteOrder.BIG)
    233                  byte_order = ByteOrder.BIG
    234
    235                  if header.mode == 33554432:
    236                      header = self._inspect(raw, ByteOrder.LITTLE)
    237                      byte_order = ByteOrder.LITTLE
    238
    239              else:
    240                  raise DensityMapFormatError("Not a mode 2 CCP4 map file")
    241
    242              stream.read(header.nsymbt)  # symmetry_data
    243
    244              count = header.ns * header.nr * header.nc
    245              map_data = stream.read(4 * count)
    246
    247              if byte_order == ByteOrder.NATIVE:
    248                  array = numpy.fromstring(map_data, numpy.float32, count)
    249              else:
    250                  array = numpy.zeros((count,), numpy.float32)
    251                  index = 0
    252                  while len(map_data) >= 4 * 10000:
    253                      values = struct.unpack(byte_order + '10000f', map_data[:4 * 10000])
    254                      array[index:index + 10000] = numpy.array(values, numpy.float32)
    255                      index += 10000
    256                      map_data = map_data[4 * 10000:]
    257                  values = struct.unpack(byte_order + '%df' % (len(map_data) / 4), map_data)
    258                  array[index:] = numpy.array(values, numpy.float32)
    259
    260              del map_data
    261
    262              array.shape = self._shape(header)
    263              data = array.T
    264
    265              spacing = self._spacing(header)
    266              origin = self._origin(header, spacing)
    267
    268              return DensityInfo(data, spacing, origin, header=raw)
    269
    270
    271 -class DensityMapWriter(object):
    272      """
    273      Binary MRC density map writer.
    274      """
    275
    276 -    def reconstruct_header(self, density):
    277          """
    278          Attempt to reconstruct the header, given L{DensityInfo}'s
    279          data shape, spacing and origin.
    280
    281          @param density: density info
    282          @type density: L{DensityInfo}
    283
    284          @return: reconstructed binary header
    285          @rtype: bytes
    286          """
    287          N = list(density.data.shape)
    288          MODE = 2
    289
    290          if isinstance(density.spacing, float):
    291              spacing = 3 * [density.spacing]
    292          else:
    293              spacing = density.spacing
    294
    295          if density.origin is None:
    296              origin = 3 * [0.]
    297          else:
    298              origin = density.origin
    299
    300          if density.axes is None:
    301              MAP = list(range(1, 4))
    302          else:
    303              MAP = list(density.axes)
    304
    305          start = [int(round(origin[i] / spacing[i], 0)) for i in range(3)]
    306          M = [density.data.shape[i] for i in range(3)]
    307          cella = [density.data.shape[i] * spacing[i] for i in range(3)]
    308          cellb = 3 * [90.]
    309
    310          stats = [density.data.min(), density.data.max(), density.data.mean()]
    311
    312          ISPG = 0
    313          NSYMBT = 0
    314          LSKFLG = 0
    315
    316          JUNK = [0] * 25
    317          ORIGIN = [0., 0., 0.]
    318          MACHST = 0
    319
    320          args = N + [MODE] + start + M + cella + cellb + \
    321                 MAP + stats + [ISPG, NSYMBT, LSKFLG] + JUNK + \
    322                 ORIGIN + [0, MACHST, 0., 0] + [b' ' * 796]
    323
    324          return struct.pack('=10l6f3l3f3l25l3f2l1f1l796s', *args)
    325
    326 -    def write(self, stream, density):
    327          """
    328          Write C{density} to a binary C{stream}.
    329
    330          @param stream: destination binary stream
    331          @type stream: stream
    332          @param density: input density info
    333          @type density: L{DensityInfo}
    334          """
    335          if density.header is not None:
    336              stream.write(density.header)
    337          else:
    338              stream.write(self.reconstruct_header(density))
    339
    340          data = density.data.T.flatten()
    341          x = struct.pack('=%df' % len(data), *data.tolist())
    342          stream.write(x)
    343
    344 -    def write_file(self, filename, density):
    345          """
    346          Write C{density} to a binary file.
    347
    348          @param filename: destination file name
    349          @type filename: str
    350          @param density: input density info
    351          @type density: L{DensityInfo}
    352          """
    353          with open(filename, 'wb') as stream:
    354              self.write(stream, density)
    355