diff --git a/plotting.py b/plotting.py index e5c1863ee40081c2f13786abd6c6656666a00bf9..8e7850e7b0649b87cff9d4491797214f287f0ba3 100644 --- a/plotting.py +++ b/plotting.py @@ -296,7 +296,7 @@ class ErrorSleeve: def values(self, x): """Return the errors at the positions x.""" - return self.errorFunction(x) + return self.errorFunction(x) + self.func(x, *self.p)[:, np.newaxis] class BrokenAxis: @@ -334,6 +334,7 @@ class BrokenAxis: xaxisPos = None yaxisPos = None + lim = None def __init__(self, *args, **kwargs): self.init(*args, **kwargs) @@ -443,7 +444,7 @@ class BrokenAxis: return self - def setupAxes(self, dic): + def _get_grid_properties(self): if self.which == "x": nrows = 1 @@ -468,7 +469,17 @@ class BrokenAxis: sharey = [(0, 1), (2, 3)] else: raise Exception() + return nrows, ncols, height_ratios, width_ratios, sharex, sharey + def setupAxes(self, dic): + ( + nrows, + ncols, + height_ratios, + width_ratios, + sharex, + sharey, + ) = self._get_grid_properties() self.a = dic.pop("ax", []) self.f = dic.pop("fig", None) if self.f is None: @@ -852,6 +863,36 @@ class BrokenAxis: a.set_yscale(*args, **kwargs) return None + def axhline(self, *args, **kwargs): + if len(args) > 0: + y = args[0] + else: + y = kwargs.pop("y", None) + if y is None: + raise Exception("You need to give y in axhline.") + axes = self._on_which_axes_to_plot_this(y=y) + for a in axes: + a.axhline(*args, **kwargs) + return None + + def _on_which_axes_to_plot_this(self, x=None, y=None): + x_contained = [True for a in self.a] + y_contained = [True for a in self.a] + if x is not None: + if not iterTest(x): + x = [x] + for i, a in enumerate(self.a): + limits = a.get_xlim() + x_contained[i] = any([limits[0] < X < limits[1] for X in x]) + if y is not None: + if not iterTest(y): + y = [y] + for i, a in enumerate(self.a): + limits = a.get_ylim() + y_contained[i] = any([limits[0] < Y < limits[1] for Y in y]) + + return [a for i, a in enumerate(self.a) if x_contained[i] and y_contained[i]] + class brokenAxis(BrokenAxis):