Coverage for gpkit/interactive/plot_sweep.py: 7%

69 statements  

« prev     ^ index     » next       coverage.py v7.4.0, created at 2024-01-07 22:56 -0500

1"Implements plot_sweep1d function" 

2import matplotlib.pyplot as plt 

3from ..exceptions import InvalidGPConstraint 

4 

5 

6def assign_axes(var, posys, axes): 

7 "Assigns axes to posys, creating and formatting if necessary" 

8 if not hasattr(posys, "__iter__"): 

9 posys = [posys] 

10 n = len(posys) 

11 if axes is None: 

12 _, axes = plt.subplots(n, 1, sharex="col", figsize=(4.5, 3+1.5*n)) 

13 if n == 1: 

14 axes = [axes] 

15 format_and_label_axes(var, posys, axes) 

16 elif n == 1 and not hasattr(axes, "__len__"): 

17 axes = [axes] 

18 return posys, axes 

19 

20 

21def format_and_label_axes(var, posys, axes, ylabel=True): 

22 "Formats and labels axes" 

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

24 if ylabel: 

25 if hasattr(posy, "key"): 

26 ustr = posy.key.unitstr(dimless="-") 

27 ylabel = (posy.key.descr.get("label", posy.key.name) 

28 + f" [{ustr}]") 

29 else: 

30 ylabel = str(posy) 

31 ax.set_ylabel(ylabel) 

32 ax.grid(color="0.6") 

33 # ax.set_frame_on(False) 

34 for item in [ax.xaxis.label, ax.yaxis.label]: 

35 item.set_fontsize(12) 

36 for item in ax.get_xticklabels() + ax.get_yticklabels(): 

37 item.set_fontsize(9) 

38 ax.tick_params(length=0) 

39 ax.spines['left'].set_visible(False) 

40 ax.spines['top'].set_visible(False) 

41 for i in ax.spines.values(): 

42 i.set_linewidth(0.6) 

43 i.set_color("0.6") 

44 i.set_linestyle("dotted") 

45 ustr = var.key.unitstr(dimless="-") 

46 xlabel = (var.key.descr.get("label", var.key.name) 

47 + f" [{ustr}]") 

48 ax.set_xlabel(xlabel) # pylint: disable=undefined-loop-variable 

49 plt.locator_params(nbins=4) 

50 plt.subplots_adjust(wspace=0.15) 

51 

52 

53# pylint: disable=too-many-locals,too-many-branches,too-many-statements 

54def plot_1dsweepgrid(model, sweeps, posys, origsol=None, tol=0.01, **solveargs): 

55 """Creates and plots a sweep from an existing model 

56 

57 Example usage: 

58 f, _ = plot_sweep_1d(m, {'x': np.linspace(1, 2, 5)}, 'y') 

59 f.savefig('mysweep.png') 

60 """ 

61 origsubs = {swept: model.substitutions[swept] for swept in sweeps 

62 if swept in model.substitutions} 

63 if origsubs and not origsol: 

64 try: 

65 origsol = model.solve(**solveargs) 

66 except InvalidGPConstraint: 

67 origsol = model.localsolve(**solveargs) 

68 if not hasattr(posys, "__iter__"): 

69 posys = [posys] 

70 

71 nposy, nsweep = len(posys), len(sweeps) 

72 f, axes = plt.subplots(nposy, nsweep, sharex='col', sharey='row', 

73 figsize=(4+2*nsweep, 4+2*nposy)) 

74 plt.subplots_adjust(hspace=0.15) 

75 

76 for i, (swept, swept_over) in enumerate(sweeps.items()): 

77 if isinstance(swept_over, tuple) and len(swept_over) == 2: 

78 sol = model.autosweep({swept: swept_over}, tol=tol, **solveargs) 

79 else: 

80 sol = model.sweep({swept: swept_over}, **solveargs) 

81 

82 if nsweep == 1: 

83 if nposy == 1: 

84 subaxes = [axes] 

85 else: 

86 subaxes = axes 

87 elif nposy == 1: 

88 subaxes = [axes[i]] 

89 else: 

90 subaxes = axes[:, i] 

91 

92 sol.plot(posys, subaxes) 

93 if origsubs: 

94 for posy, ax in zip(posys, subaxes): 

95 ax.plot(origsubs[swept], origsol(posy), "ko", markersize=4) 

96 format_and_label_axes(swept, posys, subaxes, ylabel=bool(i == 0)) 

97 model.substitutions.update(origsubs) 

98 

99 return f, axes