148 lines
3.9 KiB
Python
148 lines
3.9 KiB
Python
def dimensionality_decorator(func):
|
|
def decorated(self, *args):
|
|
if hasattr(self[0], "__len__"):
|
|
return Array([func(s_, *args) for s_ in self])
|
|
else:
|
|
return Array(func(self, *args))
|
|
return decorated
|
|
|
|
def indexing_decorator(func):
|
|
def decorated(self, v, *args):
|
|
nn = len(shape(v))
|
|
if (nn ==1):
|
|
if (type(v[0]) is bool):
|
|
return Array([s_ for s_, v_ in zip(self,v) if v_])
|
|
else:
|
|
return Array([self[v_] for v_ in v])
|
|
elif nn ==2:
|
|
if type(v[0]) is bool:
|
|
return Array([[s_ for s_, v_ in zip(self[n],v[n]) if v_ ] for n in range(nn)])
|
|
else:
|
|
ret = func(self, v, *args)
|
|
if hasattr(ret, "__len__"):
|
|
return Array(ret)
|
|
else:
|
|
return ret
|
|
return decorated
|
|
|
|
class Array(list):
|
|
def __init__(self, data):
|
|
super(Array, self).__init__(data)
|
|
self.getint = super(Array, self).__getitem__
|
|
@dimensionality_decorator
|
|
def __add__(self, v):
|
|
return add(self, v)
|
|
@dimensionality_decorator
|
|
def __sub__(self, v):
|
|
return sub(self, v)
|
|
@dimensionality_decorator
|
|
def __mul__(self, v):
|
|
return mul(self, v)
|
|
@dimensionality_decorator
|
|
def __div__(self, v):
|
|
return div(self, v)
|
|
@dimensionality_decorator
|
|
def interpolate(self, num):
|
|
return interpolate(self, num)
|
|
@dimensionality_decorator
|
|
def __lt__(self, v):
|
|
return lower_than(self, v)
|
|
@dimensionality_decorator
|
|
def __le__(self, v):
|
|
return lower_equal(self, v)
|
|
@dimensionality_decorator
|
|
def __gt__(self, v):
|
|
return greater_than(self, v)
|
|
@dimensionality_decorator
|
|
def __ge__(self, v):
|
|
return greater_equal(self, v)
|
|
@indexing_decorator
|
|
def __getitem__(self, v):
|
|
return super(Array, self).__getitem__(v)
|
|
@property
|
|
def shape(self):
|
|
return shape(self)
|
|
@property
|
|
def abs(self):
|
|
return arr_abs(self)
|
|
@property
|
|
def T(self):
|
|
return Array(transpose(self))
|
|
### math helper functions
|
|
def shape(self):
|
|
o = self
|
|
s = []
|
|
while hasattr(o, "__len__"):
|
|
s = s+[len(o)]
|
|
o = o[0]
|
|
return s
|
|
|
|
def transpose(a):
|
|
mm = len(a[0])
|
|
a_t = [[v[m] for v in a] for m in range(mm)]
|
|
return a_t
|
|
|
|
def arr_abs(a):
|
|
n = a.shape
|
|
if len(n)==1:
|
|
a_abs = [abs(e) for e in a]
|
|
elif len(n)==2:
|
|
a_abs = [[abs(e) for e in v] for v in a]
|
|
return Array(a_abs)
|
|
|
|
def interpolate(v,points_per_interval, round_to=2):
|
|
v_i=[[round(v[m]+1.*n/points_per_interval*(v[m+1]-v[m]), round_to) for n in range(points_per_interval)] for m in range(len(v)-1)]
|
|
ret = []
|
|
for i in v_i:
|
|
ret=ret+i
|
|
ret=ret+[v[-1]]
|
|
return ret
|
|
|
|
def add(list1, v):
|
|
if hasattr(v, "__len__"):
|
|
return [l+v_ for l, v_ in zip(list1, v)]
|
|
else:
|
|
return [l+v for l in list1]
|
|
|
|
def sub(list1, v):
|
|
if hasattr(v, "__len__"):
|
|
return [l-v_ for l, v_ in zip(list1, v)]
|
|
else:
|
|
return [l-v for l in list1]
|
|
|
|
def mul(list1, v):
|
|
if hasattr(v, "__len__"):
|
|
return [l*v_ for l, v_ in zip(list1, v)]
|
|
else:
|
|
return [l*v for l in list1]
|
|
|
|
def div(list1, v):
|
|
if hasattr(v, "__len__"):
|
|
return [1.*l/v_ for l, v_ in zip(list1, v)]
|
|
else:
|
|
return [1.*l/v for l in list1]
|
|
|
|
def lower_than(self, v):
|
|
if hasattr(v, "__len__"):
|
|
return [v_>s_ for v_, s_ in zip(v, self)]
|
|
else:
|
|
return [v>s_ for s_ in self]
|
|
|
|
def lower_equal(self, v):
|
|
if hasattr(v, "__len__"):
|
|
return [v_>=s_ for v_, s_ in zip(v, self)]
|
|
else:
|
|
return [v>=s_ for s_ in self]
|
|
|
|
def greater_than(self, v):
|
|
if hasattr(v, "__len__"):
|
|
return [s_>v_ for v_, s_ in zip(v, self)]
|
|
else:
|
|
return [s_>v for s_ in self]
|
|
|
|
def greater_equal(self, v):
|
|
if hasattr(v, "__len__"):
|
|
return [s_>=v_ for v_, s_ in zip(v, self)]
|
|
else:
|
|
return [s_>=v for s_ in self]
|