devlib overhaul

[jacione]

It just occurred to me that a much cleaner and more readable way to implement the various backend libraries might be to have operations as methods of a custom array class. For example:

from abc import ABC
import operator

class DevLibArray(ABC):
    def __init__(self, arr):
        self._arr = self._asarray(arr)

    @staticmethod
    @abstractmethod
    def _asarray(arr):
        """Backend-specific coercion."""

    @property
    def arr(self):
        return self._arr

    @classmethod
    def _wrap(cls, raw):
        """Wrap an already-native array without re-coercing."""
        obj = object.__new__(cls)
        obj._arr = raw
        return obj

    def _unwrap(self, other):
        if isinstance(other, DevLibArray):
            # don't silently mix backends
            if type(other) is not type(self):
                return TypeError
            return other._arr
        return other  # skip this for scalars, native arrays

    # Implement computational functions as methods.
    @abstractmethod
    def fft(self):
        pass


def _make_binop(op, reflected=False):
    def method(self, other):
        rhs = self._unwrap(other)
        if rhs is NotImplemented:
            return NotImplemented
        raw = op(rhs, self._arr) if reflected else op(self._arr, rhs)
        return type(self)._wrap(raw)
    return method

# Loop over all useful operators
for name, op in [('add', operator.add), ('sub', operator.sub), ]:
    setattr(DevLibArray, f'__{name}__', _make_binop(op))
    setattr(DevLibArray, f'__r{name}__', _make_binop(op, reflected=True))


class NumPyArray(DevLibArray):
    _asarray = staticmethod(np.asarray)

    def fft(self):
        return self._wrap(np.fft.fftn(self._arr))


class CuPyArray(DevLibArray):
    @staticmethod
    def _asarray(arr):
        return cp.asarray(arr)

    def fft(self):
        return self._wrap(cp.fft.fftn(self._arr))

The _wrap() and _unwrap() methods should keep it pretty efficient, and the operator methods would allow it to work like an array for things like simple arithmetic. There are probably other edge cases I haven't considered, but I think overall it could be a lot easier to read the actual controller code, etc. In implementation, it would look like this:

# old version
self.ds_image = devlib.fft(self.rs_image)

# new version
self.ds_image = self.rs_image.fft()

This would be a massive overhaul. Given my graduation timeline, I certainly won't have time to implement it myself. But I thought I'd share the idea for whoever comes (and/or stays) after me. Do with it what you will :)