Fix pickle serialization of SOLPSSimulation and SOLPSMesh

Current implementation of __getstate__ and __setstate__ methods of SOLPSSimulation and SOLPSMesh breaks serialization of these objects with pickle. For example, SOLPSMesh.__getstate__() ignores the state of _cr, _cz, _nx, _ny and other SOLPSMesh attributes defined at initialization. Therefore, pickle.loads(pickle.dumps(solps_mesh)) returns a broken object. The same is true for SOLPSSimulation.

I suggest moving most of the code from the SOLPSMesh andSOLPSSimulation initializers to the _initial_setup() functions, which can be called in either __init __ () or __setstate__(). For example, the code for SOLPSMesh.__init__() and SOLPSMesh.__setstate__() would look like this:

    def __init__(self, r, z, vol, neighbix, neighbiy):

        if r.shape != z.shape:
            raise ValueError('Shape of r array: {0} mismatch the shape of z array: {1}.'.format(r.shape, z.shape))

        if vol.shape != r.shape[1:]:
            raise ValueError('Shape of vol array: {0} mismatch the grid dimentions: {1}.'.format(vol.shape, r.shape[1:]))

        if neighbix.shape != r.shape:
            raise ValueError('Shape of neighbix array must be {0}, but it is  {1}.'.format(r.shape, neighbix.shape))

        if neighbiy.shape != r.shape:
            raise ValueError('Shape of neighbix array must be {0}, but it is  {}.'.format(r.shape, neighbiy.shape))

        self._r = r
        self._z = z

        self._vol = vol

        self._neighbix = neighbix.astype(int)
        self._neighbiy = neighbiy.astype(int)

        self._initial_setup()

    def __setstate__(self, state):
        self._r = state['r']
        self._z = state['z']
        self._vol = state['vol']
        self._neighbix = state['neighbix']
        self._neighbiy = state['neighbiy']
        self._initial_setup()

Similar changes should be made to the SOLPSSimulation.__init__() and SOLPSSimulation.__setstate__().

Once this is fixed, saving and loading the SOLPSSimulation will become trivial:

    def save(self, filename):
        """
        Saves SOLPSSimulation to file.
        """
        with open(filename, 'wb') as file_handle:
            pickle.dump(self, file_handle)

    @classmethod
    def load(cls, filename):
        """
        Loads SOLPSSimulation from file.
        """
        with open(filename, 'rb') as file_handle:
            sim = pickle.load(file_handle)

        return sim

and the load_solps_from_pickle() function could be removed as redundant.

I'll make a draft PR to show what I'm proposing here.

cherab / solps

Fix pickle serialization of SOLPSSimulation and SOLPSMesh #53