I'm having trouble getting a NumPy MaskedArray
subclass to round-trip through pickle and preserve the extra subclass attributes. Here is an example:
import numpy as np
import cPickle as pickle
from numpy import ma
class SubArray(np.ndarray):
"""Defines a generic np.ndarray subclass, that stores some metadata
in the dictionary `info`."""
def __new__(cls, arr, info={}):
x = np.asanyarray(arr).view(cls)
x.info = info
return x
def __array_finalize__(self, obj):
self.info = getattr(obj, 'info', {'ATTR': 'MISSING'})
return
class MSubArray(SubArray, ma.MaskedArray):
def __new__(cls, data, info={}, mask=ma.nomask, dtype=None):
subarr = SubArray(data, info)
_data = ma.MaskedArray.__new__(cls, data=subarr, mask=mask, dtype=dtype)
_data.info = subarr.info
return _data
def __array_finalize__(self, obj):
ma.MaskedArray.__array_finalize__(self, obj)
SubArray.__array_finalize__(self, obj)
return
ms = MSubArray([1, 2], info={'a': 1})
print('Pre-pickle:', ms.info, ms.data.info)
pkl = pickle.dumps(ms)
ms_from_pkl = pickle.loads(pkl)
print('Post-pickle:', ms_from_pkl.info, ms_from_pkl.data.info)
This produces:
Pre-pickle: {'a': 1} {'a': 1}
Post-pickle: {} {}
Any hints on what I'm doing wrong would be most appreciated!