Coverage for gamdpy/configuration/colarray.py: 97%

1import numpy as np 

2 

3# Can be doctested: https://docs.python.org/3/library/doctest.html 

4# python3 -m doctest colarray.py # No output == No problem! 

5# python3 -m doctest -v colarray.py # The verbose version 

6 

7class colarray(): 

8 """ The Column array Class 

9 

10 A class storing several sets ('columns') of lengths with identical dimensions in a single numpy array. Strings are used as indices along the zeroth dimension corresponding to different columns of lengths. 

11  

12 Examples 

13 -------- 

14 

15 Storage for positions, velocities, and forces, for 1000 particles in 2 dimensions: 

16 

17 >>> ca = colarray(('r', 'v', 'f'), size=(1000,2)) 

18 >>> ca.shape 

19 (3, 1000, 2) 

20 >>> ca.column_names 

21 ('r', 'v', 'f') 

22 >>> # Data is accessed via string indices (similar to dataframes in pandas): 

23 >>> ca['r'] = np.ones((1000,2)) 

24 >>> ca['v'] = 2 # Broadcast by numpy to correct shape 

25 >>> print(ca['r'] + 0.01*ca['v']) 

26 [[1.02 1.02] 

27 [1.02 1.02] 

28 [1.02 1.02] 

29 ... 

30 [1.02 1.02] 

31 [1.02 1.02] 

32 [1.02 1.02]] 

33 """ 

34 

35 # Most error handling is left to be handled by numpy, as it gives usefull error messages  

36 # (illustrated in the documentation string above). 

37 

38 def __init__(self, column_names, size, dtype=np.float32, array=None): 

39 self.column_names = column_names 

40 self.dtype = dtype 

41 self.indices = {key:index for index,key in enumerate(column_names)} 

42 if type(array)==np.ndarray: # Used, e.g., when loading from file 

43 self.array = array 

44 else: 

45 self.array = np.zeros((len(column_names), *size), dtype=dtype) 

46 

47 self.shape = self.array.shape 

48 

49 def __setitem__(self, key, data): 

50 self.array[self.indices[key]] = data 

51 

52 def __getitem__(self, key): 

53 return self.array[self.indices[key]] 

54 

55 def __repr__(self): 

56 return 'colarray('+str(tuple(self.indices.keys()))+', '+self.array.shape[1:].__repr__()+')\n'+self.array.__repr__() 

57 

58 def copy(self): 

59 return colarray(self.column_names, self.shape, self.dtype, self.array.copy()) 

60 

61 

62 def save(self, filename): 

63 """ 

64 Save a colarray to disk. 

65 >>> ca = colarray(('r', 'v', 'f'), size=(1000,2)) 

66 >>> ca.save('my_colarray') 

67 >>> # Remove the files used for storage of the colarray: 

68 >>> colarray.remove_files('my_colarray') 

69 """ 

70 np.save(f'{filename}.npy', self.array) 

71 with open(f'{filename}.col', 'w') as f: # Use pickle / json 

72 f.write(str(len(self.column_names)) + '\n') 

73 for key in self.column_names: 

74 f.write(key + '\n') 

75 return 

76 

77 

78 def load(filename): 

79 """ 

80 Load a colarray from disk. 

81 >>> ca = colarray(('r', 'v', 'f'), size=(1000,2)) 

82 >>> ca['f'] = np.random.uniform(size=(1000,2))  

83 >>> ca.save('my_colarray') 

84 >>> ca2 = colarray.load('my_colarray') 

85 >>> for col in ca.column_names: 

86 ... print(np.all(ca2[col]==ca[col])) 

87 True 

88 True 

89 True 

90 

91 Remove the files used for storage of the colarray: 

92 >>> colarray.remove_files('my_colarray') 

93  

94 The file(s) needs to be present: 

95 >>> ca2 = colarray.load('my_colarray') 

96 Traceback (most recent call last): 

97 ... 

98 FileNotFoundError: [Errno 2] No such file or directory: 'my_colarray.col' 

99 """ 

100 

101 with open(f'{filename}.col', 'r') as f: 

102 num_columns = int(f.readline()) 

103 column_names = [] 

104 for i in range(num_columns): 

105 column_names.append(f.readline()[:-1]) # removing '\n' 

106 array = np.load(f'{filename}.npy') 

107 return colarray(column_names, array.shape[1:], array=array) 

108 

109 def remove_files(filename): 

110 """ 

111 Remove files storing a colarray 

112 >>> ca = colarray(('r', 'v', 'f'), size=(1000,2)) 

113 >>> ca.save('my_colarray') 

114 >>> colarray.remove_files('my_colarray') 

115 """ 

116 

117 import os 

118 os.remove(f'{filename}.col') 

119 os.remove(f'{filename}.npy') 

120