1"""Miscellaneous small classes""" 

2from operator import xor 

3from functools import reduce 

4import numpy as np 

5from .units import Quantity, qty # pylint: disable=unused-import 

6 

7Strings = (str,) 

8Numbers = (int, float, np.number, Quantity) 

9 

10 

11class FixedScalarMeta(type): 

12 "Metaclass to implement instance checking for fixed scalars" 

13 def __instancecheck__(cls, obj): 

14 return hasattr(obj, "hmap") and len(obj.hmap) == 1 and not obj.vks 

15 

16 

17class FixedScalar(metaclass=FixedScalarMeta): # pylint: disable=no-init 

18 "Instances of this class are scalar Nomials with no variables" 

19 

20 

21class Count: 

22 "Like python 2's itertools.count, for Python 3 compatibility." 

23 def __init__(self): 

24 self.count = -1 

25 

26 def next(self): 

27 "Increment self.count and return it" 

28 self.count += 1 

29 return self.count 

30 

31 

32def matrix_converter(name): 

33 "Generates conversion function." 

34 def to_(self): # used in tocoo, tocsc, etc below 

35 "Converts to another type of matrix." 

36 # pylint: disable=unused-variable 

37 return getattr(self.tocsr(), "to"+name)() 

38 return to_ 

39 

40 

41class CootMatrix: 

42 "A very simple sparse matrix representation." 

43 def __init__(self, row, col, data): 

44 self.row, self.col, self.data = row, col, data 

45 self.shape = [(max(self.row) + 1) if self.row else 0, 

46 (max(self.col) + 1) if self.col else 0] 

47 

48 def __eq__(self, other): 

49 return (self.row == other.row and self.col == other.col 

50 and self.data == other.data and self.shape == other.shape) 

51 

52 tocoo = matrix_converter("coo") 

53 tocsc = matrix_converter("csc") 

54 todia = matrix_converter("dia") 

55 todok = matrix_converter("dok") 

56 todense = matrix_converter("dense") 

57 

58 def tocsr(self): 

59 "Converts to a Scipy sparse csr_matrix" 

60 from scipy.sparse import csr_matrix 

61 return csr_matrix((self.data, (self.row, self.col))) 

62 

63 def dot(self, arg): 

64 "Returns dot product with arg." 

65 return self.tocsr().dot(arg) 

66 

67 

68class SolverLog(list): 

69 "Adds a `write` method to list so it's file-like and can replace stdout." 

70 def __init__(self, output=None, *, verbosity=0): 

71 list.__init__(self) 

72 self.verbosity = verbosity 

73 self.output = output 

74 

75 def write(self, writ): 

76 "Append and potentially write the new line." 

77 if writ != "\n": 

78 writ = writ.rstrip("\n") 

79 self.append(str(writ)) 

80 if self.verbosity > 0: # pragma: no cover 

81 self.output.write(writ) 

82 

83 

84class DictOfLists(dict): 

85 "A hierarchy of dicionaries, with lists at the bottom." 

86 

87 def append(self, sol): 

88 "Appends a dict (of dicts) of lists to all held lists." 

89 if not hasattr(self, "initialized"): 

90 _enlist_dict(sol, self) 

91 self.initialized = True # pylint: disable=attribute-defined-outside-init 

92 else: 

93 _append_dict(sol, self) 

94 

95 def atindex(self, i): 

96 "Indexes into each list independently." 

97 return self.__class__(_index_dict(i, self, self.__class__())) 

98 

99 def to_arrays(self): 

100 "Converts all lists into array." 

101 _enray(self, self) 

102 

103 

104def _enlist_dict(d_in, d_out): 

105 "Recursively copies d_in into d_out, placing non-dict items into lists." 

106 for k, v in d_in.items(): 

107 if isinstance(v, dict): 

108 d_out[k] = _enlist_dict(v, v.__class__()) 

109 else: 

110 d_out[k] = [v] 

111 assert set(d_in.keys()) == set(d_out.keys()) 

112 return d_out 

113 

114 

115def _append_dict(d_in, d_out): 

116 "Recursively travels dict d_out and appends items found in d_in." 

117 for k, v in d_in.items(): 

118 if isinstance(v, dict): 

119 d_out[k] = _append_dict(v, d_out[k]) 

120 else: 

121 d_out[k].append(v) 

122 return d_out 

123 

124 

125def _index_dict(idx, d_in, d_out): 

126 "Recursively travels dict d_in, placing items at idx into dict d_out." 

127 for k, v in d_in.items(): 

128 if isinstance(v, dict): 

129 d_out[k] = _index_dict(idx, v, v.__class__()) 

130 else: 

131 try: 

132 d_out[k] = v[idx] 

133 except (IndexError, TypeError): # if not an array, return as is 

134 d_out[k] = v 

135 return d_out 

136 

137 

138def _enray(d_in, d_out): 

139 "Recursively turns lists into numpy arrays." 

140 for k, v in d_in.items(): 

141 if isinstance(v, dict): 

142 d_out[k] = _enray(v, v.__class__()) 

143 else: 

144 if len(v) == 1: 

145 v, = v 

146 else: 

147 v = np.array(v) 

148 d_out[k] = v 

149 return d_out 

150 

151 

152class HashVector(dict): 

153 """A simple, sparse, string-indexed vector. Inherits from dict. 

154 

155 The HashVector class supports element-wise arithmetic: 

156 any undeclared variables are assumed to have a value of zero. 

157 

158 Arguments 

159 --------- 

160 arg : iterable 

161 

162 Example 

163 ------- 

164 >>> x = gpkit.nomials.Monomial("x") 

165 >>> exp = gpkit.small_classes.HashVector({x: 2}) 

166 """ 

167 hashvalue = None 

168 

169 def __hash__(self): 

170 "Allows HashVectors to be used as dictionary keys." 

171 if self.hashvalue is None: 

172 self.hashvalue = reduce(xor, map(hash, self.items()), 0) 

173 return self.hashvalue 

174 

175 def copy(self): 

176 "Return a copy of this" 

177 hv = self.__class__(self) 

178 hv.hashvalue = self.hashvalue 

179 return hv 

180 

181 def __pow__(self, other): 

182 "Accepts scalars. Return Hashvector with each value put to a power." 

183 if isinstance(other, Numbers): 

184 return self.__class__({k: v**other for (k, v) in self.items()}) 

185 return NotImplemented 

186 

187 def __mul__(self, other): 

188 """Accepts scalars and dicts. Returns with each value multiplied. 

189 

190 If the other object inherits from dict, multiplication is element-wise 

191 and their key's intersection will form the new keys.""" 

192 try: 

193 return self.__class__({k: v*other for (k, v) in self.items()}) 

194 except: # pylint: disable=bare-except 

195 return NotImplemented 

196 

197 def __add__(self, other): 

198 """Accepts scalars and dicts. Returns with each value added. 

199 

200 If the other object inherits from dict, addition is element-wise 

201 and their key's union will form the new keys.""" 

202 if isinstance(other, Numbers): 

203 return self.__class__({k: v + other for (k, v) in self.items()}) 

204 if isinstance(other, dict): 

205 sums = self.copy() 

206 for key, value in other.items(): 

207 if key in sums: 

208 svalue = sums[key] 

209 if value == -svalue: 

210 del sums[key] # remove zeros created by addition 

211 else: 

212 sums[key] = value + svalue 

213 else: 

214 sums[key] = value 

215 sums.hashvalue = None 

216 return sums 

217 return NotImplemented 

218 

219 # pylint: disable=multiple-statements 

220 def __neg__(self): return -1*self 

221 def __sub__(self, other): return self + -other 

222 def __rsub__(self, other): return other + -self 

223 def __radd__(self, other): return self + other 

224 def __truediv__(self, other): return self * other**-1 

225 def __rtruediv__(self, other): return other * self**-1 

226 def __rmul__(self, other): return self * other 

227 

228 

229EMPTY_HV = HashVector()