Coverage for gpkit/interactive/plotting.py: 60%

1"""Plotting methods""" 

2from collections import Counter 

3import plotly.graph_objects as go 

4import matplotlib.pyplot as plt 

5import numpy as np 

6from .plot_sweep import assign_axes 

7from .. import GPCOLORS 

8 

9 

10def compare(models, sweeps, posys, tol=0.001): 

11 """Compares the values of posys over a sweep of several models. 

12 

13 If posys is of the same length as models, this will plot different 

14 variables from different models. 

15 

16 Currently only supports a single sweepvar. 

17 

18 Example Usage: 

19 compare([aec, fbc], {"R": (160, 300)}, 

20 ["cost", ("W_{\\rm batt}", "W_{\\rm fuel}")], tol=0.001) 

21 """ 

22 sols = [m.autosweep(sweeps, tol, verbosity=0) for m in models] 

23 posys, axes = assign_axes(sols[0].bst.sweptvar, posys, None) 

24 for posy, ax in zip(posys, axes): 

25 for i, sol in enumerate(sols): 

26 if hasattr(posy, "__len__") and len(posy) == len(sols): 

27 p = posy[i] 

28 else: 

29 p = posy 

30 color = GPCOLORS[i % len(GPCOLORS)] 

31 if sol._is_cost(p): # pylint: disable=protected-access 

32 ax.fill_between(sol.sampled_at, 

33 sol.cost_lb(), sol.cost_ub(), 

34 facecolor=color, edgecolor=color, 

35 linewidth=0.75) 

36 else: 

37 ax.plot(sol.sampled_at, sol(p), color=color) 

38 

39 

40def plot_convergence(model): 

41 """Plots the convergence of a signomial programming model 

42 

43 Arguments 

44 --------- 

45 model: Model 

46 Signomial programming model that has already been solved 

47 

48 Returns 

49 ------- 

50 matplotlib.pyplot Figure 

51 Plot of cost as functions of SP iteration # 

52 """ 

53 fig, ax = plt.subplots() 

54 

55 it = np.array([]) 

56 cost = np.array([]) 

57 for n, gp in enumerate(model.program.gps): 

58 try: 

59 cost = np.append(cost, gp.result['cost']) 

60 it = np.append(it, n+1) 

61 except TypeError: 

62 pass 

63 ax.plot(it, cost, '-o') 

64 ax.set_xlabel('Iteration') 

65 ax.set_ylabel('Cost') 

66 ax.set_xticks(range(1, len(model.program.gps)+1)) 

67 return fig, ax 

68 

69 

70def treemap(model, itemize="variables", sizebycount=False): 

71 """Plots model structure as Plotly TreeMap 

72 

73 Arguments 

74 --------- 

75 model: Model 

76 GPkit model object 

77 

78 itemize (optional): string, either "variables" or "constraints" 

79 Specify whether to iterate over the model varkeys or constraints 

80 

81 sizebycount (optional): bool 

82 Whether to size blocks by number of variables/constraints or use 

83 default sizing 

84 

85 Returns 

86 ------- 

87 plotly.graph_objects.Figure 

88 Plot of model hierarchy 

89 

90 """ 

91 modelnames = [] 

92 parents = [] 

93 sizes = [] 

94 

95 if itemize == "variables": 

96 lineagestrs = [l.lineagestr() or "Model" for l in model.varkeys] 

97 elif itemize == "constraints": 

98 lineagestrs = [l.lineagestr() or "Model" for l in model.flat()] 

99 

100 modelcount = Counter(lineagestrs) 

101 for modelname, count in modelcount.items(): 

102 modelnames.append(modelname) 

103 if "." in modelname: 

104 parent = modelname.rsplit(".", 1)[0] 

105 elif modelname != "Model": 

106 parent = "Model" 

107 else: 

108 parent = "" 

109 parents.append(parent) 

110 sizes.append(count) 

111 

112 for parent in parents: 

113 # pylint: disable=modified-iterating-list 

114 if parent not in modelnames: 

115 modelnames.append(parent) 

116 if "." in parent: 

117 grandparent = parent.rsplit(".", 1)[0] 

118 elif parent != "Model": 

119 grandparent = "Model" 

120 else: 

121 grandparent = "" 

122 parents.append(grandparent) # TODO: this seems dangeroud 

123 sizes.append(0) 

124 

125 fig = go.Figure(go.Treemap( 

126 ids=modelnames, 

127 labels=[modelname.split(".")[-1] for modelname in modelnames], 

128 parents=parents, 

129 values=sizes if sizebycount else None, 

130 )) 

131 return fig