1"""Defines the VarKey class""" 

2from .small_classes import HashVector, Count, qty 

3from .repr_conventions import ReprMixin 

4 

5 

6class VarKey(ReprMixin): # pylint:disable=too-many-instance-attributes 

7 """An object to correspond to each 'variable name'. 

8 

9 Arguments 

10 --------- 

11 name : str, VarKey, or Monomial 

12 Name of this Variable, or object to derive this Variable from. 

13 

14 **descr : 

15 Any additional attributes, which become the descr attribute (a dict). 

16 

17 Returns 

18 ------- 

19 VarKey with the given name and descr. 

20 """ 

21 unique_id = Count().next 

22 vars_of_a_name = {} 

23 subscripts = ("lineage", "idx") 

24 

25 def __init__(self, name=None, **descr): 

26 # NOTE: Python arg handling guarantees 'name' won't appear in descr 

27 self.descr = descr 

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

29 unitrepr = self.unitrepr or self.units 

30 if unitrepr in ["", "-", None]: # dimensionless 

31 self.descr["units"] = None 

32 self.descr["unitrepr"] = "-" 

33 else: 

34 self.descr["units"] = qty(unitrepr) 

35 self.descr["unitrepr"] = unitrepr 

36 

37 self.key = self 

38 fullstr = self.str_without(["modelnums", "vec"]) 

39 self.eqstr = fullstr + str(self.lineage) + self.unitrepr 

40 self.hashvalue = hash(self.eqstr) 

41 self.keys = set((self.name, fullstr)) 

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

43 self.hmap.units = self.units 

44 

45 if "idx" in self.descr: 

46 if "veckey" not in self.descr: 

47 vecdescr = self.descr.copy() 

48 del vecdescr["idx"] 

49 self.veckey = VarKey(**vecdescr) 

50 else: 

51 self.keys.add(self.veckey) 

52 self.keys.add(self.str_without(["idx", "modelnums"])) 

53 

54 def __getstate__(self): 

55 "Stores varkey as its metadata dictionary, removing functions" 

56 state = self.descr.copy() 

57 for key, value in state.items(): 

58 if getattr(value, "__call__", None): 

59 state[key] = "unpickleable function %s" % value 

60 return state 

61 

62 def __setstate__(self, state): 

63 "Restores varkey from its metadata dictionary" 

64 self.__init__(**state) 

65 

66 def str_without(self, excluded=()): 

67 "Returns string without certain fields (such as 'lineage')." 

68 name = self.name 

69 if ("lineage" not in excluded and self.lineage 

70 and ("unnecessary lineage" not in excluded 

71 or self.necessarylineage)): 

72 name = self.lineagestr("modelnums" not in excluded) + "." + name 

73 if "idx" not in excluded: 

74 if self.idx: 

75 name += "[%s]" % ",".join(map(str, self.idx)) 

76 elif "vec" not in excluded and self.shape: 

77 name += "[:]" 

78 return name 

79 

80 __repr__ = str_without 

81 

82 # pylint: disable=multiple-statements 

83 def __hash__(self): return self.hashvalue 

84 def __getattr__(self, attr): return self.descr.get(attr, None) 

85 

86 @property 

87 def models(self): 

88 "Returns a tuple of just the names of models in self.lineage" 

89 return list(zip(*self.lineage))[0] 

90 

91 def latex(self, excluded=()): 

92 "Returns latex representation." 

93 name = self.name 

94 if "vec" not in excluded and "idx" not in excluded and self.shape: 

95 name = "\\vec{%s}" % name 

96 if "idx" not in excluded and self.idx: 

97 name = "{%s}_{%s}" % (name, ",".join(map(str, self.idx))) 

98 if ("lineage" not in excluded and self.lineage 

99 and ("unnecessary lineage" not in excluded 

100 or self.necessarylineage)): 

101 name = "{%s}_{%s}" % (name, 

102 self.lineagestr("modelnums" not in excluded)) 

103 return name 

104 

105 def __eq__(self, other): 

106 if not hasattr(other, "descr"): 

107 return False 

108 return self.eqstr == other.eqstr 

109 

110from .nomials import NomialMap # pylint: disable=wrong-import-position