Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Move _infer_meta_data and _parse_size to utils #6779

Merged
merged 4 commits into from
Jul 12, 2022
Merged

Move _infer_meta_data and _parse_size to utils #6779

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
944532f
Move dataset plot functions to utils
Illviljan Jul 12, 2022
55f10f8
move parse_size
Illviljan Jul 12, 2022
a83a2c3
move markersize
Illviljan Jul 12, 2022
4b2cb93
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Jul 12, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
131 changes: 2 additions & 129 deletions xarray/plot/dataset_plot.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,98 +10,12 @@
from .utils import (
_add_colorbar,
_get_nice_quiver_magnitude,
_is_numeric,
_infer_meta_data,
_parse_size,
_process_cmap_cbar_kwargs,
get_axis,
label_from_attrs,
)

# copied from seaborn
_MARKERSIZE_RANGE = np.array([18.0, 72.0])


def _infer_meta_data(ds, x, y, hue, hue_style, add_guide, funcname):
dvars = set(ds.variables.keys())
error_msg = " must be one of ({:s})".format(", ".join(dvars))

if x not in dvars:
raise ValueError("x" + error_msg)

if y not in dvars:
raise ValueError("y" + error_msg)

if hue is not None and hue not in dvars:
raise ValueError("hue" + error_msg)

if hue:
hue_is_numeric = _is_numeric(ds[hue].values)

if hue_style is None:
hue_style = "continuous" if hue_is_numeric else "discrete"

if not hue_is_numeric and (hue_style == "continuous"):
raise ValueError(
f"Cannot create a colorbar for a non numeric coordinate: {hue}"
)

if add_guide is None or add_guide is True:
add_colorbar = True if hue_style == "continuous" else False
add_legend = True if hue_style == "discrete" else False
else:
add_colorbar = False
add_legend = False
else:
if add_guide is True and funcname not in ("quiver", "streamplot"):
raise ValueError("Cannot set add_guide when hue is None.")
add_legend = False
add_colorbar = False

if (add_guide or add_guide is None) and funcname == "quiver":
add_quiverkey = True
if hue:
add_colorbar = True
if not hue_style:
hue_style = "continuous"
elif hue_style != "continuous":
raise ValueError(
"hue_style must be 'continuous' or None for .plot.quiver or "
".plot.streamplot"
)
else:
add_quiverkey = False

if (add_guide or add_guide is None) and funcname == "streamplot":
if hue:
add_colorbar = True
if not hue_style:
hue_style = "continuous"
elif hue_style != "continuous":
raise ValueError(
"hue_style must be 'continuous' or None for .plot.quiver or "
".plot.streamplot"
)

if hue_style is not None and hue_style not in ["discrete", "continuous"]:
raise ValueError("hue_style must be either None, 'discrete' or 'continuous'.")

if hue:
hue_label = label_from_attrs(ds[hue])
hue = ds[hue]
else:
hue_label = None
hue = None

return {
"add_colorbar": add_colorbar,
"add_legend": add_legend,
"add_quiverkey": add_quiverkey,
"hue_label": hue_label,
"hue_style": hue_style,
"xlabel": label_from_attrs(ds[x]),
"ylabel": label_from_attrs(ds[y]),
"hue": hue,
}


def _infer_scatter_data(ds, x, y, hue, markersize, size_norm, size_mapping=None):

Expand Down Expand Up @@ -134,47 +48,6 @@ def _infer_scatter_data(ds, x, y, hue, markersize, size_norm, size_mapping=None)
return data


# copied from seaborn
def _parse_size(data, norm):

import matplotlib as mpl

if data is None:
return None

data = data.values.flatten()

if not _is_numeric(data):
levels = np.unique(data)
numbers = np.arange(1, 1 + len(levels))[::-1]
else:
levels = numbers = np.sort(np.unique(data))

min_width, max_width = _MARKERSIZE_RANGE
# width_range = min_width, max_width

if norm is None:
norm = mpl.colors.Normalize()
elif isinstance(norm, tuple):
norm = mpl.colors.Normalize(*norm)
elif not isinstance(norm, mpl.colors.Normalize):
err = "``size_norm`` must be None, tuple, or Normalize object."
raise ValueError(err)

norm.clip = True
if not norm.scaled():
norm(np.asarray(numbers))
# limits = norm.vmin, norm.vmax

scl = norm(numbers)
widths = np.asarray(min_width + scl * (max_width - min_width))
if scl.mask.any():
widths[scl.mask] = 0
sizes = dict(zip(levels, widths))

return pd.Series(sizes)


class _Dataset_PlotMethods:
"""
Enables use of xarray.plot functions as attributes on a Dataset.
Expand Down
127 changes: 127 additions & 0 deletions xarray/plot/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,9 @@

ROBUST_PERCENTILE = 2.0

# copied from seaborn
_MARKERSIZE_RANGE = np.array([18.0, 72.0])


def import_matplotlib_pyplot():
"""import pyplot"""
Expand Down Expand Up @@ -1141,3 +1144,127 @@ def _adjust_legend_subtitles(legend):
# The sutbtitles should have the same font size
# as normal legend titles:
text.set_size(font_size)


def _infer_meta_data(ds, x, y, hue, hue_style, add_guide, funcname):
dvars = set(ds.variables.keys())
error_msg = " must be one of ({:s})".format(", ".join(dvars))

if x not in dvars:
raise ValueError("x" + error_msg)

if y not in dvars:
raise ValueError("y" + error_msg)

if hue is not None and hue not in dvars:
raise ValueError("hue" + error_msg)

if hue:
hue_is_numeric = _is_numeric(ds[hue].values)

if hue_style is None:
hue_style = "continuous" if hue_is_numeric else "discrete"

if not hue_is_numeric and (hue_style == "continuous"):
raise ValueError(
f"Cannot create a colorbar for a non numeric coordinate: {hue}"
)

if add_guide is None or add_guide is True:
add_colorbar = True if hue_style == "continuous" else False
add_legend = True if hue_style == "discrete" else False
else:
add_colorbar = False
add_legend = False
else:
if add_guide is True and funcname not in ("quiver", "streamplot"):
raise ValueError("Cannot set add_guide when hue is None.")
add_legend = False
add_colorbar = False

if (add_guide or add_guide is None) and funcname == "quiver":
add_quiverkey = True
if hue:
add_colorbar = True
if not hue_style:
hue_style = "continuous"
elif hue_style != "continuous":
raise ValueError(
"hue_style must be 'continuous' or None for .plot.quiver or "
".plot.streamplot"
)
else:
add_quiverkey = False

if (add_guide or add_guide is None) and funcname == "streamplot":
if hue:
add_colorbar = True
if not hue_style:
hue_style = "continuous"
elif hue_style != "continuous":
raise ValueError(
"hue_style must be 'continuous' or None for .plot.quiver or "
".plot.streamplot"
)

if hue_style is not None and hue_style not in ["discrete", "continuous"]:
raise ValueError("hue_style must be either None, 'discrete' or 'continuous'.")

if hue:
hue_label = label_from_attrs(ds[hue])
hue = ds[hue]
else:
hue_label = None
hue = None

return {
"add_colorbar": add_colorbar,
"add_legend": add_legend,
"add_quiverkey": add_quiverkey,
"hue_label": hue_label,
"hue_style": hue_style,
"xlabel": label_from_attrs(ds[x]),
"ylabel": label_from_attrs(ds[y]),
"hue": hue,
}


# copied from seaborn
def _parse_size(data, norm):

import matplotlib as mpl

if data is None:
return None

data = data.values.flatten()

if not _is_numeric(data):
levels = np.unique(data)
numbers = np.arange(1, 1 + len(levels))[::-1]
else:
levels = numbers = np.sort(np.unique(data))

min_width, max_width = _MARKERSIZE_RANGE
# width_range = min_width, max_width

if norm is None:
norm = mpl.colors.Normalize()
elif isinstance(norm, tuple):
norm = mpl.colors.Normalize(*norm)
elif not isinstance(norm, mpl.colors.Normalize):
err = "``size_norm`` must be None, tuple, or Normalize object."
raise ValueError(err)

norm.clip = True
if not norm.scaled():
norm(np.asarray(numbers))
# limits = norm.vmin, norm.vmax

scl = norm(numbers)
widths = np.asarray(min_width + scl * (max_width - min_width))
if scl.mask.any():
widths[scl.mask] = 0
sizes = dict(zip(levels, widths))

return pd.Series(sizes)