diff --git a/xbout/plotting/plotfuncs.py b/xbout/plotting/plotfuncs.py index 1e9a0dcd..32007c3b 100644 --- a/xbout/plotting/plotfuncs.py +++ b/xbout/plotting/plotfuncs.py @@ -864,6 +864,7 @@ def plot2d_polygon( add_colorbar = True, colorbar_label = None, separatrix = True, + separatrix_kwargs = {"color":"white", "linestyle":"-", "linewidth":1}, targets = False, add_limiter_hatching=True, grid_only = False, @@ -945,32 +946,7 @@ def plot2d_polygon( ax.set_title(da.name) if separatrix: - # plot_separatrices(da, ax, x = "R", y = "Z") + plot_separatrices(da, ax, x = "R", y = "Z", **separatrix_kwargs) - color = "white" - ls = "-" - lw = 2 - - m = da.attrs["metadata"] - - - - if m["topology"] == "connected-double-null": - - # TODO: test this for single null - if m["MXG"] > 0 or m["MYG"] > 0: - lhs = (m["ixseps1"] + 1 - m["MXG"], slice(0,m["ny_inner"] + m["MYG"] * int(4/2))) - rhs = (m["ixseps1"] + 1 - m["MXG"], slice(m["ny_inner"]+m["MYG"]*(4 - 1), None)) - else: - lhs = (m["ixseps1"]-m["MXG"]-1, slice(0,m["ny_inner"])) - rhs = (m["ixseps1"]-m["MXG"]-1, slice(m["ny_inner"], None)) - ax.plot(da["Rxy_lower_right_corners"].data[lhs], da["Zxy_lower_right_corners"].data[lhs], c = color, lw = lw, ls = ls) - ax.plot(da["Rxy_lower_right_corners"].data[rhs], da["Zxy_lower_right_corners"].data[rhs], c = color, lw = lw, ls = ls) - - if m["topology"] == "single-null": - points = (m["ixseps1"]-m["MXG"]-1, slice(0,None)) - ax.plot(da["Rxy_lower_right_corners"].data[points], da["Zxy_lower_right_corners"].data[points], c = color, lw = lw, ls = ls) - - if targets: plot_targets(da, ax, x = "R", y = "Z", hatching = add_limiter_hatching) \ No newline at end of file diff --git a/xbout/plotting/utils.py b/xbout/plotting/utils.py index db70b74a..0e19331e 100644 --- a/xbout/plotting/utils.py +++ b/xbout/plotting/utils.py @@ -78,8 +78,10 @@ def _is_core_only(da): return ix1 >= nx and ix2 >= nx -def plot_separatrices(da, ax, *, x="R", y="Z"): - """Plot separatrices""" +def plot_separatrices(da, ax, *, x="R", y="Z", **kwargs): + """ + Plot separatrices. Kwargs are passed to ax.plot(). + """ if not isinstance(da, dict): da_regions = _decompose_regions(da) @@ -116,7 +118,11 @@ def plot_separatrices(da, ax, *, x="R", y="Z"): y_sep = 0.5 * ( da_region[y].isel(**{xcoord: 0}) + da_inner[y].isel(**{xcoord: -1}) ) - ax.plot(x_sep, y_sep, "k--") + default_style = {"color": "black", "linestyle": "--"} + if any(x for x in kwargs if x in ["c", "ls"]): + raise ValueError("When passing separatrix plot style kwargs, use 'color' and 'linestyle' instead lf 'c' and 'ls'") + style = {**default_style, **kwargs} + ax.plot(x_sep, y_sep, **style) def plot_targets(da, ax, *, x="R", y="Z", hatching=True):