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[:2] == "b'": 

78 writ = writ[2:-1] 

79 if writ != "\n": 

80 self.append(writ.rstrip("\n")) 

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

82 self.output.write(writ) 

83 

84 

85class DictOfLists(dict): 

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

87 

88 def append(self, sol): 

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

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

91 _enlist_dict(sol, self) 

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

93 else: 

94 _append_dict(sol, self) 

95 

96 def atindex(self, i): 

97 "Indexes into each list independently." 

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

99 

100 def to_arrays(self): 

101 "Converts all lists into array." 

102 _enray(self, self) 

103 

104 

105def _enlist_dict(d_in, d_out): 

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

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

108 if isinstance(v, dict): 

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

110 else: 

111 d_out[k] = [v] 

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

113 return d_out 

114 

115 

116def _append_dict(d_in, d_out): 

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

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

119 if isinstance(v, dict): 

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

121 else: 

122 try: 

123 d_out[k].append(v) 

124 except KeyError as e: 

125 raise RuntimeWarning("Key `%s` was added after the first sweep." 

126 % k) from e 

127 return d_out 

128 

129 

130def _index_dict(idx, d_in, d_out): 

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

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

133 if isinstance(v, dict): 

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

135 else: 

136 try: 

137 d_out[k] = v[idx] 

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

139 d_out[k] = v 

140 return d_out 

141 

142 

143def _enray(d_in, d_out): 

144 "Recursively turns lists into numpy arrays." 

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

146 if isinstance(v, dict): 

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

148 else: 

149 if len(v) == 1: 

150 v, = v 

151 else: 

152 v = np.array(v) 

153 d_out[k] = v 

154 return d_out 

155 

156 

157class HashVector(dict): 

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

159 

160 The HashVector class supports element-wise arithmetic: 

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

162 

163 Arguments 

164 --------- 

165 arg : iterable 

166 

167 Example 

168 ------- 

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

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

171 """ 

172 hashvalue = None 

173 

174 def __hash__(self): 

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

176 if self.hashvalue is None: 

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

178 return self.hashvalue 

179 

180 def copy(self): 

181 "Return a copy of this" 

182 hv = self.__class__(self) 

183 hv.hashvalue = self.hashvalue 

184 return hv 

185 

186 def __pow__(self, other): 

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

188 if isinstance(other, Numbers): 

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

190 return NotImplemented 

191 

192 def __mul__(self, other): 

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

194 

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

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

197 try: 

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

199 except: # pylint: disable=bare-except 

200 return NotImplemented 

201 

202 def __add__(self, other): 

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

204 

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

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

207 if isinstance(other, Numbers): 

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

209 if isinstance(other, dict): 

210 sums = self.copy() 

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

212 if key in sums: 

213 svalue = sums[key] 

214 if value == -svalue: 

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

216 else: 

217 sums[key] = value + svalue 

218 else: 

219 sums[key] = value 

220 sums.hashvalue = None 

221 return sums 

222 return NotImplemented 

223 

224 # pylint: disable=multiple-statements 

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

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

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

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

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

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

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

232 

233 

234EMPTY_HV = HashVector()