buffer.py 1.8 KB

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