Coverage for /builds/ericyuan00000/ase/ase/utils/arraywrapper.py: 100.00%

1# fmt: off 

2 

3"""Module for wrapping an array without being an array. 

4 

5This can be desirable because we would like atoms.cell to be like an array, 

6but we don't like atoms.positions @ atoms.cell to also be a cell, and as far 

7as I/we know, it's not possible to subclass array without that kind of thing 

8happening. 

9 

10For most methods and attributes we can just bind the name as a property: 

11 

12 @property 

13 def reshape(self): 

14 return np.asarray(self).reshape 

15 

16For 'in-place' operations like += we need to make sure to return self 

17rather than the array though: 

18 

19 def __iadd__(self, other): 

20 np.asarray(self).__iadd__(other) 

21 return self 

22 

23This module provides the @arraylike decorator which does these things 

24for all the interesting ndarray methods. 

25""" 

26 

27 

28from functools import update_wrapper 

29 

30import numpy as np 

31 

32inplace_methods = ['__iadd__', '__imul__', '__ipow__', '__isub__', 

33 '__itruediv__', '__imatmul__'] 

34 

35forward_methods = ['__abs__', '__add__', '__contains__', '__eq__', 

36 '__ge__', '__getitem__', '__gt__', '__hash__', 

37 '__iter__', '__le__', '__len__', '__lt__', 

38 '__mul__', '__ne__', '__neg__', '__pos__', 

39 '__pow__', '__radd__', '__rmul__', '__rpow__', 

40 '__rsub__', '__rtruediv__', '__setitem__', 

41 '__sub__', '__truediv__'] 

42 

43default_methods = ['__eq__', '__le__', '__lt__', '__ge__', 

44 '__gt__', '__ne__', '__hash__'] 

45 

46if hasattr(np.ndarray, '__matmul__'): 

47 forward_methods += ['__matmul__', '__rmatmul__'] 

48 

49 

50def forward_inplace_call(name): 

51 arraymeth = getattr(np.ndarray, name) 

52 

53 def f(self, obj): 

54 a = self.__array__() 

55 arraymeth(a, obj) 

56 return self 

57 

58 update_wrapper(f, arraymeth) 

59 return f 

60 

61 

62def wrap_array_attribute(name): 

63 wrappee = getattr(np.ndarray, name) 

64 if wrappee is None: # For example, __hash__ is None 

65 assert name == '__hash__' 

66 return None 

67 

68 def attr(self): 

69 array = np.asarray(self) 

70 return getattr(array, name) 

71 

72 update_wrapper(attr, wrappee) 

73 

74 # We don't want to encourage too liberal use of the numpy methods, 

75 # nor do we want the web docs to explode with numpy docstrings or 

76 # break our own doctests. 

77 # 

78 # Therefore we cheat and remove the docstring: 

79 attr.__doc__ = None 

80 return property(attr) 

81 

82 

83def arraylike(cls): 

84 """Decorator for being like an array without being an array. 

85 

86 Poke decorators onto cls so that getting an attribute 

87 really gets that attribute from the wrapped ndarray. 

88 

89 Exceptions are made for in-place methods like +=, *=, etc. 

90 These must return self since otherwise myobj += 1 would 

91 magically turn into an ndarray.""" 

92 for name in inplace_methods: 

93 if hasattr(np.ndarray, name) and not hasattr(cls, name): 

94 meth = forward_inplace_call(name) 

95 setattr(cls, name, meth) 

96 

97 allnames = [name for name in dir(np.ndarray) if not name.startswith('_')] 

98 for name in forward_methods + allnames: 

99 if hasattr(cls, name) and name not in default_methods: 

100 continue # Was overridden -- or there's a conflict. 

101 

102 prop = wrap_array_attribute(name) 

103 setattr(cls, name, prop) 

104 return cls