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 ..globals import NamedVariables, Vectorize 

8from ..varkey import VarKey 

9from ..small_classes import Strings, Numbers 

10from ..small_scripts import is_sweepvar 

11 

12 

13def addmodelstodescr(descr, addtonamedvars=None): 

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

15 if NamedVariables.lineage: 

16 descr["lineage"] = NamedVariables.lineage 

17 if addtonamedvars: 

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

19 

20 

21class Variable(Monomial): 

22 """A described singlet Monomial. 

23 

24 Arguments 

25 --------- 

26 *args : list 

27 may contain "name" (Strings) 

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

29 "units" (Strings) 

30 and/or "label" (Strings) 

31 **descr : dict 

32 VarKey description 

33 

34 Returns 

35 ------- 

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

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

38 """ 

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

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

41 self.key, = args 

42 else: 

43 for arg in args: 

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

45 descr["name"] = arg 

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

47 and "value" not in descr): 

48 descr["value"] = arg 

49 elif (isinstance(arg, Iterable) 

50 and not isinstance(arg, Strings)): 

51 if is_sweepvar(arg): 

52 descr["value"] = arg 

53 else: 

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

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

56 descr["units"] = arg 

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

58 descr["label"] = arg 

59 addmodelstodescr(descr, addtonamedvars=self) 

60 self.key = VarKey(**descr) 

61 Monomial.__init__(self, self.key.hmap) 

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

63 self.__class__ = Variable 

64 

65 __hash__ = NomialData.__hash__ 

66 

67 def to(self, units): 

68 "Create new Signomial converted to new units" 

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

70 

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

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

73 

74 Example 

75 ------- 

76 x = Variable('x') 

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

78 """ 

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

80 arg, = args 

81 if not isinstance(arg, dict): 

82 args = [{self: arg}] 

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

84 

85 

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

87 """A described vector of singlet Monomials. 

88 

89 Arguments 

90 --------- 

91 shape : int or tuple 

92 length or shape of resulting array 

93 *args : 

94 may contain "name" (Strings) 

95 "value" (Iterable) 

96 "units" (Strings) 

97 and/or "label" (Strings) 

98 **descr : 

99 VarKey description 

100 

101 Returns 

102 ------- 

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

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

105 """ 

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

107 if "idx" in descr: 

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

109 

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

111 if Vectorize.vectorization: 

112 shape += Vectorize.vectorization 

113 

114 descr["shape"] = shape 

115 

116 for arg in args: 

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

118 descr["name"] = arg 

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

120 and not isinstance(arg, Strings) 

121 and "value" not in descr): 

122 descr["value"] = arg 

123 elif hasattr(arg, "__call__"): 

124 descr["value"] = arg 

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

126 descr["units"] = arg 

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

128 descr["label"] = arg 

129 

130 if "name" not in descr: 

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

132 

133 value_option = None 

134 if "value" in descr: 

135 value_option = "value" 

136 if value_option: 

137 values = descr.pop(value_option) 

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

139 if Vectorize.vectorization: 

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

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

142 else: 

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

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

145 values = np.array(values) 

146 if values.shape != shape: 

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

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

149 

150 veckeydescr = descr.copy() 

151 addmodelstodescr(veckeydescr) 

152 if value_option: 

153 if hasattr(values, "__call__"): 

154 veckeydescr["original_fn"] = values 

155 veckeydescr[value_option] = values 

156 veckey = VarKey(**veckeydescr) 

157 

158 descr["veckey"] = veckey 

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

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

161 while not it.finished: 

162 i = it.multi_index 

163 it.iternext() 

164 descr["idx"] = i 

165 if value_option: 

166 if hasattr(values, "__call__"): 

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

168 else: 

169 descr[value_option] = values[i] 

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

171 

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

173 obj.key = veckey 

174 obj.units = obj.key.units 

175 return obj 

176 

177 

178def veclinkedfn(linkedfn, i): 

179 "Generate an indexed linking function." 

180 def newlinkedfn(c): 

181 "Linked function that pulls out a particular index" 

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

183 return newlinkedfn 

184 

185 

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

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

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

189 if Vectorize.vectorization: 

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

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

192 return Variable(*args, **descr)