buffer.py 1.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354
  1. # -*- coding: utf-8 -*-
  2. from grass.pygrass.raster.raster_type import TYPE as RTYPE
  3. import ctypes
  4. import numpy as np
  5. _CELL = ('int', 'int0', 'int8', 'int16', 'int32', 'int64')
  6. CELL = tuple([getattr(np, attr) for attr in _CELL if hasattr(np, attr)])
  7. _FCELL = 'float', 'float16', 'float32'
  8. FCELL = tuple([getattr(np, attr) for attr in _FCELL if hasattr(np, attr)])
  9. _DCELL = 'float64', 'float128'
  10. DCELL = tuple([getattr(np, attr) for attr in _DCELL if hasattr(np, attr)])
  11. class Buffer(np.ndarray):
  12. """shape, mtype='FCELL', buffer=None, offset=0,
  13. strides=None, order=None
  14. """
  15. @property
  16. def mtype(self):
  17. if self.dtype in CELL:
  18. return 'CELL'
  19. elif self.dtype in FCELL:
  20. return 'FCELL'
  21. elif self.dtype in DCELL:
  22. return DCELL
  23. else:
  24. err = "Raster type: %r not supported by GRASS."
  25. raise TypeError(err % self.dtype)
  26. def __new__(cls, shape, mtype='FCELL', buffer=None, offset=0,
  27. strides=None, order=None):
  28. obj = np.ndarray.__new__(cls, shape, RTYPE[mtype]['numpy'],
  29. buffer, offset, strides, order)
  30. obj.pointer_type = ctypes.POINTER(RTYPE[mtype]['ctypes'])
  31. obj.p = obj.ctypes.data_as(obj.pointer_type)
  32. return obj
  33. def __array_finalize__(self, obj):
  34. if obj is None:
  35. return
  36. self.pointer_type = getattr(obj, 'pointer_type', None)
  37. self.p = getattr(obj, 'p', None)
  38. def __array_wrap__(self, out_arr, context=None):
  39. """See:
  40. http://docs.scipy.org/doc/numpy/user/
  41. basics.subclassing.html#array-wrap-for-ufuncs"""
  42. if out_arr.dtype == np.bool:
  43. # there is not support for boolean maps, so convert into integer
  44. out_arr = out_arr.astype(np.int32)
  45. out_arr.p = out_arr.ctypes.data_as(out_arr.pointer_type)
  46. return np.ndarray.__array_wrap__(self, out_arr, context)