m-thesis-introduction/simulations/lib/location/location.py

import numpy as np
from functools import partial

class Location:
    """
    A location is a point designated by a spatial coordinate x.

    Locations are wrappers around a Numpy N-dimensional array.
    """

    def __init__(self, x):
        self.x = np.asarray(x)

    def __repr__(self):
        return "Location({})".format(repr(self.x))

    def __getitem__(self, key):
        return self.x[key]

    def __setitem__(self, key, val):
        self.x[key] = val

    # math
    def __add__(self, other):
        if isinstance(other, Location):
            other = other.x

        return self.__class__(self.x + other)

    def __sub__(self, other):
        if isinstance(other, Location):
            other = other.x

        return self.__class__(self.x - other)

    def __mul__(self, other):
        return self.__class__(self.x * other)

    def __eq__(self, other):
        if isinstance(other, Location):
            other = other.x

        return np.all(self.x == other)

    # math alias functions
    __radd__ = __add__
    __rsub__ = __sub__
    __rmul__ = __mul__