1"Implement Variable and ArrayVariable classes" 

2from collections.abc import Iterable 

3import numpy as np 

4from .data import NomialData 

5from .array import NomialArray 

6from .math import Monomial 

7from .map import NomialMap 

8from ..globals import NamedVariables, Vectorize 

9from ..varkey import VarKey 

10from ..small_classes import Strings, Numbers, HashVector 

11from ..small_scripts import is_sweepvar 

12 

13 

14def addmodelstodescr(descr, addtonamedvars=None): 

15 "Add models to descr, optionally adding the second argument to NAMEDVARS" 

16 if NamedVariables.lineage: 

17 descr["lineage"] = NamedVariables.lineage 

18 if addtonamedvars: 

19 NamedVariables.namedvars[descr["lineage"]].append(addtonamedvars) 

20 

21 

22class Variable(Monomial): 

23 """A described singlet Monomial. 

24 

25 Arguments 

26 --------- 

27 *args : list 

28 may contain "name" (Strings) 

29 "value" (Numbers + Quantity) or (Iterable) for a sweep 

30 "units" (Strings) 

31 and/or "label" (Strings) 

32 **descr : dict 

33 VarKey description 

34 

35 Returns 

36 ------- 

37 Monomials containing a VarKey with the name '$name', 

38 where $name is the vector's name and i is the VarKey's index. 

39 """ 

40 def __init__(self, *args, **descr): 

41 if len(args) == 1 and isinstance(args[0], VarKey): 

42 self.key, = args 

43 else: 

44 for arg in args: 

45 if isinstance(arg, Strings) and "name" not in descr: 

46 descr["name"] = arg 

47 elif (isinstance(arg, Numbers) or hasattr(arg, "__call__") 

48 and "value" not in descr): 

49 descr["value"] = arg 

50 elif (isinstance(arg, Iterable) 

51 and not isinstance(arg, Strings)): 

52 if is_sweepvar(arg): 

53 descr["value"] = arg 

54 else: 

55 descr["value"] = ("sweep", arg) 

56 elif isinstance(arg, Strings) and "units" not in descr: 

57 descr["units"] = arg 

58 elif isinstance(arg, Strings) and "label" not in descr: 

59 descr["label"] = arg 

60 addmodelstodescr(descr, addtonamedvars=self) 

61 self.key = VarKey(**descr) 

62 hmap = NomialMap({HashVector({self.key: 1}): 1.0}) 

63 hmap.units = self.key.units 

64 Monomial.__init__(self, hmap) 

65 # NOTE: needed because Signomial.__init__ will change the class 

66 self.__class__ = Variable 

67 

68 __hash__ = NomialData.__hash__ 

69 

70 def to(self, units): 

71 "Create new Signomial converted to new units" 

72 return Monomial(self).to(units) # pylint: disable=no-member 

73 

74 def sub(self, *args, **kwargs): # pylint: disable=arguments-differ 

75 """Same as nomial substitution, but also allows single-argument calls 

76 

77 Example 

78 ------- 

79 x = Variable('x') 

80 assert x.sub(3) == Variable('x', value=3) 

81 """ 

82 if len(args) == 1 and "val" not in kwargs: 

83 arg, = args 

84 if not isinstance(arg, dict): 

85 args = [{self: arg}] 

86 return Monomial.sub(self, *args, **kwargs) 

87 

88 

89class ArrayVariable(NomialArray): # pylint: disable=too-many-locals 

90 """A described vector of singlet Monomials. 

91 

92 Arguments 

93 --------- 

94 shape : int or tuple 

95 length or shape of resulting array 

96 *args : 

97 may contain "name" (Strings) 

98 "value" (Iterable) 

99 "units" (Strings) 

100 and/or "label" (Strings) 

101 **descr : 

102 VarKey description 

103 

104 Returns 

105 ------- 

106 NomialArray of Monomials, each containing a VarKey with name '$name_{i}', 

107 where $name is the vector's name and i is the VarKey's index. 

108 """ 

109 def __new__(cls, shape, *args, **descr): # pylint: disable=too-many-branches, too-many-statements, arguments-differ 

110 if "idx" in descr: 

111 raise ValueError("the description field 'idx' is reserved") 

112 

113 shape = (shape,) if isinstance(shape, Numbers) else tuple(shape) 

114 if Vectorize.vectorization: 

115 shape += Vectorize.vectorization 

116 

117 descr["shape"] = shape 

118 

119 for arg in args: 

120 if isinstance(arg, Strings) and "name" not in descr: 

121 descr["name"] = arg 

122 elif (isinstance(arg, (Numbers, Iterable)) 

123 and not isinstance(arg, Strings) 

124 and "value" not in descr): 

125 descr["value"] = arg 

126 elif hasattr(arg, "__call__"): 

127 descr["value"] = arg 

128 elif isinstance(arg, Strings) and "units" not in descr: 

129 descr["units"] = arg 

130 elif isinstance(arg, Strings) and "label" not in descr: 

131 descr["label"] = arg 

132 

133 if "name" not in descr: 

134 descr["name"] = "\\fbox{%s}" % VarKey.unique_id() 

135 

136 value_option = None 

137 if "value" in descr: 

138 value_option = "value" 

139 if value_option: 

140 values = descr.pop(value_option) 

141 if value_option and not hasattr(values, "__call__"): 

142 if Vectorize.vectorization: 

143 if not hasattr(values, "shape"): 

144 values = np.full(shape, values, "f") 

145 else: 

146 values = np.broadcast_to(values, reversed(shape)).T 

147 elif not hasattr(values, "shape"): 

148 values = np.array(values) 

149 if values.shape != shape: 

150 raise ValueError("the value's shape %s is different than" 

151 " the vector's %s." % (values.shape, shape)) 

152 

153 veckeydescr = descr.copy() 

154 addmodelstodescr(veckeydescr) 

155 if value_option: 

156 if hasattr(values, "__call__"): 

157 veckeydescr["original_fn"] = values 

158 veckeydescr[value_option] = values 

159 veckey = VarKey(**veckeydescr) 

160 

161 descr["veckey"] = veckey 

162 vl = np.empty(shape, dtype="object") 

163 it = np.nditer(vl, flags=['multi_index', 'refs_ok']) 

164 while not it.finished: 

165 i = it.multi_index 

166 it.iternext() 

167 descr["idx"] = i 

168 if value_option: 

169 if hasattr(values, "__call__"): 

170 descr[value_option] = veclinkedfn(values, i) 

171 else: 

172 descr[value_option] = values[i] 

173 vl[i] = Variable(**descr) 

174 

175 obj = np.asarray(vl).view(NomialArray) 

176 obj.key = veckey 

177 obj.units = obj.key.units 

178 return obj 

179 

180 

181def veclinkedfn(linkedfn, i): 

182 "Generate an indexed linking function." 

183 def newlinkedfn(c): 

184 "Linked function that pulls out a particular index" 

185 return np.array(linkedfn(c))[i] 

186 return newlinkedfn 

187 

188 

189class VectorizableVariable(Variable, ArrayVariable): # pylint: disable=too-many-ancestors 

190 "A Variable outside a vectorized environment, an ArrayVariable within." 

191 def __new__(cls, *args, **descr): 

192 if Vectorize.vectorization: 

193 shape = descr.pop("shape", ()) 

194 return ArrayVariable.__new__(cls, shape, *args, **descr) 

195 return Variable(*args, **descr)