Repository: aazuspan/sankee (language: python)

Issue:
Interactive widgets
Instead of returning a plot object, returning a [FigureWidget](https://plotly.com/python/figurewidget/) would allow for interactive controls of the plot parameters, like buttons to turn each class on and off or a slider to set the max classes.


Discussion / hints:

Fix patch:
diff --git a/.gitignore b/.gitignore
index dcbef05..0856271 100644
--- a/.gitignore
+++ b/.gitignore
@@ -4,4 +4,5 @@
 *.egg-info/
 *build/
 *dist/
-Pipfile.lock
\ No newline at end of file
+Pipfile.lock
+.vscode/
\ No newline at end of file
diff --git a/sankee/__init__.py b/sankee/__init__.py
index cd9a64d..eb14702 100644
--- a/sankee/__init__.py
+++ b/sankee/__init__.py
@@ -1,5 +1,5 @@
 from sankee import datasets
-from sankee.core import sankify
+from sankee.plotting import sankify
 
 __version__ = "0.1.0"
 
diff --git a/sankee/core.py b/sankee/core.py
deleted file mode 100644
index a403602..0000000
--- a/sankee/core.py
+++ /dev/null
@@ -1,491 +0,0 @@
-import sys
-from typing import Any, Dict, List, Tuple, Union
-
-import ee
-import numpy as np
-import pandas as pd
-import plotly.graph_objects as go
-
-from sankee import utils
-
-# Temporary property name to store image labels in
-LABEL_PROPERTY = "sankee_label"
-
-
-def sankify(
-    image_list: List[ee.Image],
-    band: str,
-    labels: Dict[int, str],
-    palette: Dict[int, str],
-    region: Union[None, ee.Geometry] = None,
-    label_list: Union[None, List[str]] = None,
-    exclude: Union[None, List[int]] = None,
-    max_classes: Union[None, int] = None,
-    n: int = 500,
-    title: Union[None, str] = None,
-    scale: Union[None, int] = None,
-    seed: int = 0,
-    dataset: Any = None,
-) -> go.Figure:
-    """
-    Generate an interactive Sankey plot showing land cover change over time from a series of images.
-
-    Parameters
-    ----------
-    image_list : List[ee.Image]
-        An ordered list of images representing a time series of classified data. Each image will be
-        sampled to generate the Sankey plot. Any length of list is allowed, but lists with more than
-        3 or 4 images may produce unusable plots.
-    band : str
-        The name of the band in all images of image_list that contains classified data.
-    labels : dict
-        The labels associated with each value of all images in image_list. Every value in the images
-        must be included as a key in the labels dictionary.
-    palette : dict
-        The colors associated with each value of all images in image_list. Every value in the images
-        must be included as a key in the palette dictionary.
-    region : ee.Geometry, default None
-        A region to generate samples within. The region must overlap all images. If none is
-        provided, the geometry of the first image will be used. For this to work, images must be
-        bounded.
-    label_list : List[str], default None
-        An ordered list of labels corresponding to the images. The list must be the same length as
-        image_list. If none is provided, sequential numeric labels will be automatically assigned
-        starting at 0. Labels are displayed on-hover on the Sankey nodes.
-    exclude : list[int], default None
-        An optional list of pixel values to exclude from the plot. Excluded values must be raw pixel
-        values rather than class labels. This can be helpful if the region is dominated by one or
-        more unchanging classes and the goal is to visualize changes in smaller classes.
-    max_classes : int, default None
-        If a value is provided, small classes will be removed until max_classes remain. Class size
-        is calculated based on total times sampled in the time series.
-    n : int, default 500
-        The number of sample points to randomly generate for characterizing all images. More samples
-        will provide more representative data but will take longer to process.
-    title : str, default None
-        An optional title that will be displayed above the Sankey plot.
-    scale : int, default None
-        The scale in image units to perform sampling at. If none is provided, GEE will attempt to
-        use the image's nominal scale, which may cause errors depending on the image projection.
-    seed : int, default 0
-        The seed value used to generate repeatable results during random sampling.
-    dataset : None
-        Unused parameter that will be removed in a future release. If you have a Dataset object to
-        sankify, use `Dataset.sankify` instead.
-
-    Returns
-    -------
-    plotly.graph_objs._figure.Figure
-        An interactive Sankey plot.
-    """
-    # Older versions of geemap are incompatible with sankee v0.1.0 and pass the wrong positional
-    # arguments.
-    if isinstance(band, ee.Geometry) and "geemap" in sys.modules:
-        raise ValueError(
-            "Your versions of `geemap` and `sankee` are incompatible. Please update geemap to the "
-            "latest version (pip install -U geemap)."
-        ) from None
-
-    if dataset is not None:
-        raise ValueError(
-            "`sankee.sankify` no longer supports a `dataset` parameter. Use "
-            "`Dataset.sankify` instead."
-        )
-
-    if region is None:
-        region = image_list[0].geometry()
-
-    label_list = label_list if label_list is not None else list(range(len(image_list)))
-    label_list = [str(label) for label in label_list]
-    if len(label_list) != len(image_list):
-        raise ValueError("The number of labels must match the number of images.")
-    if len(set(label_list)) != len(label_list):
-        raise ValueError("All labels in the `label_list` must be unique.")
-
-    labeled_images = _label_images(image_list, label_list)
-    sample_data = _collect_sample_data(labeled_images, region, band, label_list, n, scale, seed)
-    cleaned_data = _clean_data(sample_data, exclude, max_classes)
-
-    check_data_is_compatible(labels, cleaned_data)
-
-    return _generate_sankey_plot(cleaned_data, labels, palette, title)
-
-
-def _label_images(image_list: List[ee.Image], label_list: List[str]) -> List[ee.Image]:
-    """
-    Set a label property for each image in a list using the corresponding label list.
-    """
-    labeled = []
-    for img, label in zip(image_list, label_list):
-        labeled.append(img.set(LABEL_PROPERTY, label))
-
-    return labeled
-
-
-def _collect_sample_data(
-    image_list: List[ee.Image],
-    region: ee.Geometry,
-    band: str,
-    label_list: List[str],
-    n: int = 500,
-    scale: Union[None, int] = None,
-    seed: int = 0,
-) -> pd.DataFrame:
-    """
-    Randomly sample values of a list of images to quantify change over time.
-
-    Parameters
-    ----------
-    image_list : List[ee.Image]
-        Labeled, classified images representing change over time.
-    region : ee.Geometry
-        A region to sample within.
-    band : str
-        The name of the band in all images of image_list that contains classified data.
-    label_list : List[str]
-        A list of labels associated with each image in the image_list. If none is provided,
-        sequential numeric labels will be automatically assigned starting at 0.
-    n : int, default 500
-        The number of sample points to randomly generate for characterizing all images. More samples
-        will provide more representative data but will take longer to process.
-    scale : int, default None
-        The scale in image units to perform sampling at. If none is provided, GEE will attempt to
-        use the image's nominal scale, which may cause errors depending on the image projection.
-    seed : int, default 0
-        The random seed used to generate sample points, for repeatability.
-
-    Returns
-    -------
-    pd.DataFrame
-        A dataframe in which each row represents a single sample point and columns represent the
-        class of that point in each image of the image list.
-    """
-    samples = ee.FeatureCollection.randomPoints(region, n, seed)
-
-    sample_data = _extract_values_from_images_at_points(image_list, samples, band, scale)
-
-    try:
-        data = utils.feature_collection_to_dataframe(sample_data)
-    except ee.EEException as e:
-        # ee may raise an error if the band name isn't valid. This is a bit of a hack to catch that
-        # since ee doesn't raise more specific errors.
-        if band in e.args[0]:
-            raise ValueError(
-                f'"{band}" is not a valid band name. Check that the dataset band name exists for '
-                "all images."
-            )
-        else:
-            raise e
-
-    _check_for_missing_samples(data, label_list)
-
-    return data[label_list]
-
-
-def _check_for_missing_samples(data: pd.DataFrame, label_list: List[str]) -> None:
-    """Check that the sampled data has a column for each label in the label list. If not, it may be
-    that sampling occurred outside of the image bounds.
-    """
-    missing_labels = [label for label in label_list if label not in data.columns]
-    if missing_labels:
-        raise Exception(
-            f"No valid samples were collected in these images: {missing_labels}. Check that the"
-            " sampling region overlaps the image bounds."
-        )
-
-
-def _extract_values_from_images_at_points(
-    image_list: List[ee.Image], sample_points: ee.FeatureCollection, band: str, scale: int
-) -> ee.FeatureCollection:
-    """Take a list of images and a collection of sample points and extract image values to each
-    sample point. The image values will be stored in a property based on the image label.
-    """
-
-    def extract_values_at_point(pt):
-        for img in image_list:
-            cover = img.reduceRegion(
-                reducer=ee.Reducer.first(), geometry=pt.geometry(), scale=scale
-            ).get(band)
-            label = img.get(LABEL_PROPERTY)
-            pt = ee.Feature(pt).set(label, cover)
-
-        return pt
-
-    return sample_points.map(extract_values_at_point)
-
-
-def _clean_data(
-    data: pd.DataFrame, exclude: List[int] = None, max_classes: Union[None, int] = None
-) -> pd.DataFrame:
-    """
-    Perform some cleaning on data before plotting by excluding unwanted classes and limiting the
-    number of classes.
-
-    Parameters
-    ----------
-    data : pd.DataFrame
-        A dataframe in which each row represents as single sample point and columns represent the
-        class of that point in each image of an image list.
-    exclude : List[int]
-        A list of class values to remove from the dataframe.
-    max_classes : int
-        The maximum number of unique classes to include in the dataframe. If more classes are
-        present, the smallest classes will be omitted from the plot. If max_classes is None, no
-        classes will be dropped.
-
-    Returns
-    -------
-    pd.DataFrame
-        A cleaned dataframe.
-    """
-    data = data.dropna()
-
-    if exclude:
-        data = data[~data.isin(exclude).any(axis=1)]
-    if max_classes:
-        data = utils.drop_small_classes(data, max_classes)
-
-    return data
-
-
-def check_data_is_compatible(labels: Dict[int, str], data: pd.DataFrame) -> None:
-    """Check for values that are present in data and are not present in labels and raise an error if
-    any are found.
-    """
-    missing_keys = []
-
-    for _, col in data.iteritems():
-        missing_keys += utils.get_missing_keys(col, labels)
-
-    if missing_keys:
-        raise Exception(
-            "The following values are present in the data and undefined in the labels and palette:"
-            f" {np.unique(missing_keys)}"
-        )
-
-
-def _generate_sankey_plot(
-    data: pd.DataFrame, labels: Dict[int, str], palette: Dict[int, str], title: str
-) -> go.Figure:
-    (
-        node_labels,
-        link_labels,
-        node_palette,
-        link_palette,
-        label,
-        source,
-        target,
-        value,
-    ) = _format_for_sankey(data, labels, palette)
-    return _plot(
-        node_labels,
-        link_labels,
-        node_palette,
-        link_palette,
-        label,
-        source,
-        target,
-        value,
-        title=title,
-    )
-
-
-def _format_for_sankey(
-    data: pd.DataFrame, labels: Dict[int, str], palette: Dict[int, str]
-) -> Tuple:
-    """
-    Take a dataframe of data representing classified sample points and return all parameters needed
-    to generate a Sankey plot. This is done by looping through columns in groups of two representing
-    start and end conditions and reformating data to match the Plotly Sankey input parameters.
-
-    Parameters
-    ----------
-    data : pd.DataFrame
-        A dataframe in which each row represents as single sample point and columns represent the
-        columns represent classes over an arbitrary number of time periods.
-    labels : Dict[int, str]
-        A dictionary mapping class values to their corresponding labels.
-    palette : Dict[int, str]
-        A dictionary mapping class values to their corresponding colors.
-
-    Returns
-    -------
-    Tuple
-        A tuple of the following parameters: node_labels, link_labels, node_palette, link_palette,
-        label, source, target, values.
-    """
-    formatted_data = _group_and_format_data(data, labels)
-
-    node_labels = _build_node_labels(formatted_data)
-    link_labels = _build_link_labels(formatted_data)
-    label = _build_labels(formatted_data)
-
-    source = formatted_data.source.tolist()
-    target = formatted_data.target.tolist()
-    value = formatted_data.value.tolist()
-
-    def get_color(label):
-        return palette[[k for k, v in labels.items() if v == label][0]]
-
-    node_palette = [get_color(lab) for lab in label]
-    link_palette = [get_color(i) for i in [label[j] for j in source]]
-
-    return (node_labels, link_labels, node_palette, link_palette, label, source, target, value)
-
-
-def _group_and_format_data(data: pd.DataFrame, labels: List[str]) -> pd.DataFrame:
-    """
-    Take raw data, group it into groups of two, and generate a formatted dataframe.
-    """
-    current_index = 0
-
-    dfs = []
-    for group in _group_columns(data):
-        sankified = _reformat_group(data, group, labels, start_index=current_index)
-        # The start index of the next column group will be the end index of this column group. This
-        # sets the index offset to achieve that.
-        current_index = sankified.target.min()
-        dfs.append(sankified)
-
-    return pd.concat(dfs)
-
-
-def _group_columns(data: pd.DataFrame) -> pd.DataFrame:
-    """
-    Yield all groups of two adjacent columns from a dataframe.
-    """
-    for i in range(data.columns.size - 1):
-        group_indexes = [i, i + 1]
-        yield data.iloc[:, group_indexes]
-
-
-def _reformat_group(
-    raw_data: pd.DataFrame, group_data: pd.DataFrame, labels: List[str], start_index: int = 0
-) -> pd.DataFrame:
-    column_list = group_data.columns.tolist()
-
-    # Transform the data to get counts of each combination of condition
-    sankey_data = group_data.groupby(group_data.columns.tolist()).size().reset_index(name="value")
-
-    # Calculate normalized change from start to end condition
-    sankey_data["change"] = utils.normalized_change(sankey_data, column_list[0], "value")
-
-    sankey_data = _assign_unique_indexes(raw_data, sankey_data, start_index)
-
-    # Assign labels to each source and target
-    sankey_data["source_label"] = sankey_data.iloc[:, 0].apply(lambda i: labels[i])
-    sankey_data["target_label"] = sankey_data.iloc[:, 1].apply(lambda i: labels[i])
-
-    # Store the labels of the time series periods
-    sankey_data["source_period"] = column_list[0]
-    sankey_data["target_period"] = column_list[1]
-
-    return sankey_data[
-        [
-            "source",
-            "target",
-            "value",
-            "change",
-            "source_label",
-            "target_label",
-            "source_period",
-            "target_period",
-        ]
-    ]
-
-
-def _assign_unique_indexes(
-    raw_data: pd.DataFrame, sankey_data: pd.DataFrame, start_index: int
-) -> pd.DataFrame:
-    column_list = sankey_data.columns.tolist()
-
-    # Get lists of unique source and target classes
-    unique_source = pd.unique(raw_data[column_list[0]].values.flatten()).tolist()
-    unique_target = pd.unique(raw_data[column_list[1]].values.flatten()).tolist()
-
-    # Generate a unique index for each source and target
-    sankey_data["source"] = sankey_data.iloc[:, 0].apply(
-        lambda x: unique_source.index(x) + start_index
-    )
-    # Offset the target IDs by the last source class to prevent overlap with source IDs
-    sankey_data["target"] = sankey_data.iloc[:, 1].apply(
-        lambda x: unique_target.index(x) + sankey_data.source.max() + 1
-    )
-
-    return sankey_data
-
-
-def _build_node_labels(data: pd.DataFrame) -> List[str]:
-    """
-    Build a list of period labels for nodes, corresponding to the link indexes.
-    """
-    labels = pd.concat([data.source, data.target])
-    periods = pd.concat([data.source_period, data.target_period])
-    combined = pd.DataFrame({"label": labels, "period": periods})
-    return combined.groupby("label").period.first().tolist()
-
-
-def _build_link_labels(data: pd.DataFrame) -> List[str]:
-    """
-    Build strings describing the change in states for each row of the dataframe.
-    """
-    link_labels = []
-    for _, row in data.iterrows():
-        verb = "remained" if row.source_label == row.target_label else "became"
-        link_label = f"{round(row.change * 100)}% of {row.source_label} {verb} {row.target_label}"
-        link_labels.append(link_label)
-    return link_labels
-
-
-def _build_labels(data: pd.DataFrame) -> List[str]:
-    max_id = data.target.max()
-    # Build empty placeholder labels for each class
-    label = [None for i in range(max_id + 1)]
-
-    for _, row in data.iterrows():
-        # Iteratively replace the placeholder labels with correct labels for each index
-        label[row.source] = row.source_label
-        label[row.target] = row.target_label
-
-    return label
-
-
-def _plot(
-    node_labels, link_labels, node_palette, link_palette, label, source, target, value, title=None
-) -> go.Figure:
-    """
-    Generate a Sankey plot of land cover change over an arbitrary number of time steps.
-    """
-    fig = go.Figure(
-        data=[
-            go.Sankey(
-                node=dict(
-                    pad=30,
-                    thickness=20,
-                    line=dict(color="#000000", width=1),
-                    customdata=node_labels,
-                    hovertemplate="%{customdata}<extra></extra>",
-                    label=label,
-                    color=node_palette,
-                ),
-                link=dict(
-                    source=source,
-                    target=target,
-                    line=dict(color="#909090", width=1),
-                    value=value,
-                    color=link_palette,
-                    customdata=link_labels,
-                    hovertemplate="%{customdata} <extra></extra>",
-                ),
-            )
-        ]
-    )
-
-    fig.update_layout(
-        title_text=f"<b>{title}</b>" if title else None,
-        font_size=16,
-        title_x=0.5,
-        paper_bgcolor="rgba(0, 0, 0, 0)",
-    )
-
-    return fig
diff --git a/sankee/datasets.py b/sankee/datasets.py
index 5206f54..00b1cce 100644
--- a/sankee/datasets.py
+++ b/sankee/datasets.py
@@ -4,7 +4,7 @@
 import pandas as pd
 import plotly.graph_objects as go
 
-from sankee.core import sankify
+from sankee.plotting import sankify
 
 
 class Dataset:
@@ -34,7 +34,7 @@ def __init__(
         years : List[int]
             All years available in this dataset.
         nodata : int
-            An optional no-data value to automatically exclude when sankifying.
+            An optional no-data value to automatically remove when plotting.
         """
         self.name = name
         self.id = id
@@ -89,12 +89,12 @@ def sankify(
         self,
         years: List[int],
         region: ee.Geometry,
-        exclude: Union[List[int], None] = None,
         max_classes: Union[None, int] = None,
         n: int = 500,
         title: Union[str, None] = None,
         scale: Union[int, None] = None,
         seed: int = 0,
+        exclude: None = None,
     ) -> go.Figure:
         """
         Generate an interactive Sankey plot showing land cover change over time from a series of
@@ -125,6 +125,8 @@ def sankify(
             projection.
         seed : int, default 0
             The seed value used to generate repeatable results during random sampling.
+        exclude : None
+            Unused parameter that will be removed in a future release.
 
         Returns
         -------
@@ -138,16 +140,20 @@ def sankify(
         years = sorted(years)
 
         imgs = [self.get_year(year) for year in years]
-        exclude = exclude if exclude is not None else []
-        exclude = exclude + [self.nodata] if self.nodata is not None else exclude
+
+        labels = self.labels.copy()
+        palette = self.palette.copy()
+        if self.nodata is not None:
+            labels.pop(self.nodata)
+            palette.pop(self.nodata)
+
         return sankify(
             image_list=imgs,
             label_list=years,
-            labels=self.labels,
+            labels=labels,
             band=self.band,
-            palette=self.palette,
+            palette=palette,
             region=region,
-            exclude=exclude,
             max_classes=max_classes,
             n=n,
             title=title,
diff --git a/sankee/example_data.py b/sankee/example_data.py
new file mode 100644
index 0000000..f96c4cd
--- /dev/null
+++ b/sankee/example_data.py
@@ -0,0 +1,58 @@
+import ee
+
+
+# These can only be loaded once ee is initialized, which is why we put these in a class
+class regions:
+    hobet_coal_mine = ee.Geometry.Polygon(
+        [
+            [
+                [-81.973301, 38.121953],
+                [-82.000086, 38.105746],
+                [-82.009701, 38.077104],
+                [-82.000429, 38.060617],
+                [-81.965403, 38.043044],
+                [-81.950637, 38.073321],
+                [-81.942052, 38.107907],
+                [-81.973301, 38.121953],
+            ]
+        ]
+    )
+
+    mt_st_helens = ee.Geometry.Polygon(
+        [
+            [
+                [-122.289907, 46.289229],
+                [-122.22328, 46.329545],
+                [-122.15528, 46.346138],
+                [-122.053623, 46.323855],
+                [-122.033704, 46.252682],
+                [-122.102391, 46.231786],
+                [-122.161462, 46.24176],
+                [-122.260371, 46.233686],
+                [-122.289907, 46.289229],
+            ]
+        ]
+    )
+
+    las_vegas = ee.Geometry.Polygon(
+        [
+            [
+                [-115.01184401606046, 36.24170785506492],
+                [-114.98849806879484, 36.29928186470082],
+                [-115.25628981684171, 36.35238941394592],
+                [-115.34692702387296, 36.310348922031565],
+                [-115.37988600824796, 36.160811202271944],
+                [-115.30298171137296, 36.03653336474891],
+                [-115.25628981684171, 36.05207884201088],
+                [-115.26590285395109, 36.226199908103695],
+                [-115.19174513910734, 36.25499793268206],
+            ]
+        ]
+    )
+
+    rim_fire = (
+        ee.FeatureCollection("projects/sat-io/open-datasets/MTBS/burned_area_boundaries")
+        .filter(ee.Filter.eq("Event_ID", "CA3785712008620130817"))
+        .first()
+        .geometry()
+    )
diff --git a/sankee/plotting.py b/sankee/plotting.py
new file mode 100644
index 0000000..f96bbe7
--- /dev/null
+++ b/sankee/plotting.py
@@ -0,0 +1,399 @@
+import sys
+from collections import namedtuple
+from typing import Dict, List, Union
+
+import ee
+import ipywidgets as widgets
+import pandas as pd
+import plotly.graph_objects as go
+
+from sankee import utils
+from sankee.sampling import collect_sankey_data
+
+SankeyParameters = namedtuple(
+    "SankeyParameters",
+    [
+        "node_labels",
+        "link_labels",
+        "node_palette",
+        "link_palette",
+        "label",
+        "source",
+        "target",
+        "value",
+    ],
+)
+
+
+def sankify(
+    image_list: List[ee.Image],
+    band: str,
+    labels: Dict[int, str],
+    palette: Dict[int, str],
+    region: Union[None, ee.Geometry] = None,
+    label_list: Union[None, List[str]] = None,
+    max_classes: Union[None, int] = None,
+    n: int = 500,
+    title: Union[None, str] = None,
+    scale: Union[None, int] = None,
+    seed: int = 0,
+) -> go.Figure:
+    """
+    Generate an interactive Sankey plot showing land cover change over time from a series of images.
+
+    Parameters
+    ----------
+    image_list : List[ee.Image]
+        An ordered list of images representing a time series of classified data. Each image will be
+        sampled to generate the Sankey plot. Any length of list is allowed, but lists with more than
+        3 or 4 images may produce unusable plots.
+    band : str
+        The name of the band in all images of image_list that contains classified data.
+    labels : dict
+        The labels associated with each value of all images in image_list. Any values not defined
+        in the labels will be dropped from the sampled data.
+    palette : dict
+        The colors associated with each value of all images in image_list.
+    region : ee.Geometry, default None
+        A region to generate samples within. The region must overlap all images. If none is
+        provided, the geometry of the first image will be used. For this to work, images must be
+        bounded.
+    label_list : List[str], default None
+        An ordered list of labels corresponding to the images. The list must be the same length as
+        image_list. If none is provided, sequential numeric labels will be automatically assigned
+        starting at 0. Labels are displayed on-hover on the Sankey nodes.
+    max_classes : int, default None
+        If a value is provided, small classes will be removed until max_classes remain. Class size
+        is calculated based on total times sampled in the time series.
+    n : int, default 500
+        The number of sample points to randomly generate for characterizing all images. More samples
+        will provide more representative data but will take longer to process.
+    title : str, default None
+        An optional title that will be displayed above the Sankey plot.
+    scale : int, default None
+        The scale in image units to perform sampling at. If none is provided, GEE will attempt to
+        use the image's nominal scale, which may cause errors depending on the image projection.
+    seed : int, default 0
+        The seed value used to generate repeatable results during random sampling.
+
+    Returns
+    -------
+    plotly.graph_objs._figure.Figure
+        An interactive Sankey plot.
+    """
+    # Older versions of geemap are incompatible with sankee v0.1.0 and pass the wrong positional
+    # arguments.
+    if isinstance(band, ee.Geometry) and "geemap" in sys.modules:
+        raise ValueError(
+            "Your versions of `geemap` and `sankee` are incompatible. Please update geemap to the "
+            "latest version (pip install -U geemap)."
+        ) from None
+
+    if region is None:
+        region = image_list[0].geometry()
+
+    label_list = label_list if label_list is not None else list(range(len(image_list)))
+    label_list = [str(label) for label in label_list]
+    if len(label_list) != len(image_list):
+        raise ValueError("The number of labels must match the number of images.")
+    if len(set(label_list)) != len(label_list):
+        raise ValueError("All labels in the `label_list` must be unique.")
+
+    sample_data = collect_sankey_data(
+        image_list=image_list,
+        image_labels=label_list,
+        region=region,
+        band=band,
+        n=n,
+        scale=scale,
+        seed=seed,
+        include=list(labels.keys()),
+        max_classes=max_classes,
+    )
+
+    return SankeyPlot(
+        data=sample_data,
+        labels=labels,
+        palette=palette,
+        title=title,
+    )
+
+
+class SankeyPlot(widgets.DOMWidget):
+    def __init__(
+        self,
+        data: pd.DataFrame,
+        labels: Dict[int, str],
+        palette: Dict[int, str],
+        title: str,
+    ):
+        self.data = data
+        self.labels = labels
+        self.palette = palette
+        self.title = title
+
+        self.hide = []
+        self.plot = self.generate_plot()
+        self.gui = self.generate_gui()
+
+    def get_sorted_classes(self) -> pd.Series:
+        """Return all unique class values, sorted by the total number of observations."""
+        df = self.generate_dataframe()
+        start_count = (
+            df.groupby("source")
+            .mean()
+            .reset_index()[["source", "total"]]
+            .rename(columns={"source": "class", "total": "count"})
+        )
+        end_count = (
+            df.groupby("target")
+            .sum()
+            .reset_index()[["target", "changed"]]
+            .rename(columns={"target": "class", "changed": "count"})
+        )
+        total_count = pd.concat([start_count, end_count]).groupby("class").sum().reset_index()
+
+        return total_count.sort_values(by="count", ascending=False)["class"].reset_index(drop=True)
+
+    def get_active_classes(self) -> pd.Series:
+        """Return all unique active, visibile class values after filtering."""
+        df = self.generate_dataframe()
+
+        return df[["source", "target"]].melt().value.unique()
+
+    def generate_plot_parameters(self) -> SankeyParameters:
+        """Generate Sankey plot parameters from a formatted, cleaned dataframe"""
+        df = self.generate_dataframe()
+
+        source_df = df[["source", "source_year"]].rename(
+            columns={"source": "class", "source_year": "year"}
+        )
+        target_df = df[["target", "target_year"]].rename(
+            columns={"target": "class", "target_year": "year"}
+        )
+        all_classes = pd.concat([source_df, target_df])
+
+        all_classes = all_classes.drop_duplicates().reset_index(drop=True)
+        all_classes["label"] = all_classes["class"].apply(lambda k: self.labels[k]).tolist()
+        all_classes["color"] = all_classes["class"].apply(lambda k: self.palette[k]).tolist()
+        all_classes["id"] = all_classes.groupby(["year", "class"], sort=False).ngroup()
+
+        # Join the sequential class-year IDs to the dataframe
+        df["source_id"] = pd.merge(
+            left=df,
+            right=all_classes,
+            how="left",
+            left_on=["source_year", "source"],
+            right_on=["year", "class"],
+        )["id"]
+        df["target_id"] = pd.merge(
+            left=df,
+            right=all_classes,
+            how="left",
+            left_on=["target_year", "target"],
+            right_on=["year", "class"],
+        )["id"]
+
+        return SankeyParameters(
+            node_labels=all_classes.year,
+            link_labels=df.link_label,
+            node_palette=all_classes.color,
+            link_palette=df.source_color,
+            label=all_classes.label,
+            source=df.source_id,
+            target=df.target_id,
+            value=df.changed,
+        )
+
+    def generate_dataframe(self) -> pd.DataFrame:
+        """Convert raw sampling data to a formatted dataframe"""
+        data = self.data.copy()
+
+        if self.hide:
+            hide_mask = pd.concat([(data == i).any(axis=1) for i in self.hide], axis=1).any(axis=1)
+            data = data[~hide_mask]
+
+        permutations = []
+        # Get all unique class-year combinations
+        for source, target in utils.pairwise(data.columns):
+            permutations += list(
+                zip([source] * len(data), [target] * len(data), data[source], data[target])
+            )
+        df = pd.DataFrame(permutations, columns=["source_year", "target_year", "source", "target"])
+
+        # Count the unique combinations of all four fields
+        df = (
+            df.groupby(["source_year", "target_year", "source", "target"])
+            .size()
+            .reset_index()
+            .rename(columns={0: "changed"})
+        )
+        # Count the total number of source samples in each year
+        df["total"] = df.groupby(["source_year", "source"]).changed.transform(sum)
+        # Calculate what percent of the source samples went into each target class
+        df["proportion"] = df["changed"] / df["total"]
+
+        # Join the class labels and colors to the class IDs
+        df["source_label"] = df.source.apply(lambda k: self.labels[k])
+        df["target_label"] = df.target.apply(lambda k: self.labels[k])
+        df["source_color"] = df.source.apply(lambda k: self.palette[k])
+        df["target_color"] = df.target.apply(lambda k: self.palette[k])
+
+        def build_link_label(row: pd.Series) -> str:
+            # Early exit in case all classes are excluded
+            if row.shape[0] == 0:
+                return ""
+
+            verb = "remained" if row.source == row.target else "became"
+            pct = f"{row.proportion:.0%}"
+            return f"<b>{pct}</b> of <b>{row.source_label}</b> {verb} <b>{row.target_label}</b>"
+
+        # Describe the class changes
+        df["link_label"] = df.apply(build_link_label, axis=1)
+
+        return df
+
+    @property
+    def _view_name(self):
+        """When the Sankey object is displayed by IPython, render the plot"""
+        return self.gui._view_name
+
+    @property
+    def _model_id(self):
+        """When the Sankey object is displayed by IPython, render the plot"""
+        return self.gui._model_id
+
+    def update_layout(self, *args, **kwargs):
+        """Pass layout changes to the plot. This is primarily kept for compatibility with geemap."""
+        self.plot.update_layout(*args, **kwargs)
+
+    def generate_gui(self):
+        BUTTON_HEIGHT = "24px"
+        BUTTON_WIDTH = "24px"
+
+        unique_classes = self.get_sorted_classes()
+
+        def toggle_button(button):
+            button.toggle()
+
+            class_name = button.tooltip
+            class_id = [key for key in self.labels.keys() if self.labels[key] == class_name][0]
+
+            if not button.state:
+                self.hide.append(class_id)
+            else:
+                self.hide.remove(class_id)
+
+            update_plot()
+
+        def update_plot():
+            """Swap new data into the plot"""
+            new_plot = self.generate_plot()
+            self.plot.data[0].link = new_plot.data[0].link
+            self.plot.data[0].node = new_plot.data[0].node
+
+        buttons = []
+        for i in unique_classes:
+            label = self.labels[i]
+            on_color = self.palette[i]
+            state = i in self.get_active_classes()
+
+            button = utils.ColorToggleButton(tooltip=label, on_color=on_color, state=state)
+            button.layout.width = BUTTON_WIDTH
+            button.layout.height = BUTTON_HEIGHT
+
+            button.on_click(toggle_button)
+            buttons.append(button)
+
+        def reset_plot(change):
+            for button in buttons:
+                if not button.state:
+                    button.click()
+
+        reset_button = widgets.Button(
+            icon="refresh",
+            tooltip="Reset plot",
+            layout=widgets.Layout(height=BUTTON_HEIGHT, width=BUTTON_WIDTH, padding="0 0 0 3px"),
+        )
+        reset_button.on_click(reset_plot)
+
+        open_button = widgets.Button(
+            icon="external-link",
+            tooltip="Open in new tab",
+            layout=widgets.Layout(height=BUTTON_HEIGHT, width=BUTTON_WIDTH, padding="0 0 0 3px"),
+        )
+        open_button.on_click(
+            lambda _: self.plot.update_layout(width=None, height=None).show(renderer="browser")
+        )
+
+        button_box = widgets.HBox([*buttons, widgets.Label("|"), reset_button, open_button])
+
+        gui = widgets.VBox(
+            [
+                self.plot,
+                widgets.VBox([button_box], layout=widgets.Layout(align_items="center")),
+            ]
+        )
+
+        return gui
+
+    def generate_plot(self) -> go.Figure:
+        params = self.generate_plot_parameters()
+
+        shadow_color = "#76777a"
+        label_style = f"""
+            color: #fff;
+            font-weight: 600;
+            letter-spacing: -1px;
+            text-shadow:
+                0 0 4px black,
+                -1px 1px 0 {shadow_color},
+                1px 1px 0 {shadow_color},
+                1px -1px 0 {shadow_color},
+                -1px -1px 0 {shadow_color};
+        """
+
+        title_style = """
+            color: #fff;
+            font-weight: 900;
+            word-spacing: 10px;
+            letter-spacing: 3px;
+            text-shadow:
+                0 0 1px black,
+                0 0 2px black,
+                0 0 4px black;
+        """
+
+        fig = go.Figure(
+            data=[
+                go.Sankey(
+                    node=dict(
+                        pad=30,
+                        thickness=10,
+                        line=dict(color="#505050", width=1.5),
+                        customdata=params.node_labels,
+                        hovertemplate="<b>%{customdata}</b><extra></extra>",
+                        label=[f"<span style='{label_style}'>{s}</span>" for s in params.label],
+                        color=params.node_palette,
+                    ),
+                    link=dict(
+                        source=params.source,
+                        target=params.target,
+                        line=dict(color="#909090", width=1),
+                        value=params.value,
+                        color=params.link_palette,
+                        customdata=params.link_labels,
+                        hovertemplate="%{customdata} <extra></extra>",
+                    ),
+                )
+            ]
+        )
+
+        fig.update_layout(
+            title_text=f"<span style='{title_style}'>{self.title}</span>" if self.title else None,
+            font_size=16,
+            title_x=0.5,
+            paper_bgcolor="rgba(0, 0, 0, 0)",
+        )
+
+        return go.FigureWidget(fig)
diff --git a/sankee/sampling.py b/sankee/sampling.py
new file mode 100644
index 0000000..9c6eb11
--- /dev/null
+++ b/sankee/sampling.py
@@ -0,0 +1,80 @@
+from typing import List, Union
+
+import ee
+import pandas as pd
+
+
+def collect_sankey_data(
+    *,
+    image_list: List[ee.Image],
+    image_labels: List[str],
+    region: ee.Geometry,
+    band: str,
+    n: int = 500,
+    scale: Union[None, int] = None,
+    seed: int = 0,
+    include: Union[None, List[int]] = None,
+    max_classes: Union[None, int] = None,
+) -> pd.DataFrame:
+    """Collect all the data needed to generate a Sankey diagram from a list of images.
+
+    The resulting dataframe has one column for each image and one row for each sample point.
+    """
+    points = ee.FeatureCollection.randomPoints(region=region, points=n, seed=seed)
+    samples = _extract_values(
+        image_list=image_list,
+        image_labels=image_labels,
+        sample_points=points,
+        band=band,
+        scale=scale,
+    )
+    try:
+        data = pd.DataFrame.from_dict(
+            [feat["properties"] for feat in samples.toList(samples.size()).getInfo()]
+        ).dropna()
+    except ee.EEException as e:
+        if band in str(e):
+            raise ValueError(f"The band `{band}` was not found in all images.") from None
+        else:
+            raise e
+
+    for image in image_labels:
+        if image not in data.columns:
+            raise ValueError(
+                f"Valid samples were not found for image `{image}`. Check that the"
+                " image overlaps the sampling region."
+            )
+
+    if include is not None:
+        data = data[data.isin(include).all(axis=1)]
+
+    if max_classes is not None:
+        class_counts = data.melt().value.value_counts()
+        keep_classes = class_counts[:max_classes].index.tolist()
+        data = data[data.isin(keep_classes).all(axis=1)]
+
+    return data
+
+
+def _extract_values(
+    *,
+    image_list: List[ee.Image],
+    image_labels: List[str],
+    sample_points: ee.FeatureCollection,
+    band: str,
+    scale: int,
+) -> ee.FeatureCollection:
+    """Take a list of images and a collection of sample points and extract image values to each
+    sample point. The image values will be stored in a property based on the image label.
+    """
+
+    def extract_values_at_point(pt):
+        for img, label in zip(image_list, image_labels):
+            cover = img.reduceRegion(
+                reducer=ee.Reducer.first(), geometry=pt.geometry(), scale=scale
+            ).get(band)
+            pt = ee.Feature(pt).set(label, cover)
+
+        return pt
+
+    return sample_points.map(extract_values_at_point)
diff --git a/sankee/utils.py b/sankee/utils.py
index 48b3b6d..b27bdeb 100644
--- a/sankee/utils.py
+++ b/sankee/utils.py
@@ -1,19 +1,6 @@
-import pandas as pd
+import itertools
 
-
-def feature_collection_to_dataframe(fc):
-    """
-    Extract and return FeatureCollection properties as a Pandas dataframe.
-    """
-    return pd.DataFrame.from_dict([feat["properties"] for feat in fc.toList(fc.size()).getInfo()])
-
-
-def normalized_change(data, group_col, count_col):
-    """
-    Perform group-wise data normalization to a column of values. Output values will be the
-    proportion of all values in the group.
-    """
-    return data.groupby(group_col)[count_col].apply(lambda x: x / x.sum())
+import ipywidgets as widgets
 
 
 def get_missing_keys(key_list, key_dict):
@@ -24,28 +11,36 @@ def get_missing_keys(key_list, key_dict):
     return [key for key in key_list if key not in key_dict.keys()]
 
 
-def drop_small_classes(data: pd.DataFrame, keep_classes: int) -> pd.DataFrame:
+def pairwise(iterable):
+    """Polyfill itertools.pairwise for pre 3.10.
+
+    https://docs.python.org/3/library/itertools.html#itertools.pairwise
+    """
+    a, b = itertools.tee(iterable)
+    next(b, None)
+    return zip(a, b)
+
+
+class ColorToggleButton(widgets.Button):
     """
-    Remove small classes until a maximum number of classes is reached.
-
-    Parameters
-    ----------
-    data : pd.DataFrame
-        A dataframe in which each row represents as single sample point and columns represent the
-        class at each point in time.
-    keep_classes : int
-        The maximum number of classes to keep. If more classes are present, the smallest will be
-        removed.
-
-    Returns
-    -------
-    pd.DataFrame
-        A dataframe with rows that belong to the largest classes.
+    The ipywidgets.ToggleButton doesn't support a `button_color` style, so this turns a standard
+    ipywidgets.Button into a toggle button.
     """
-    class_counts = data.melt().groupby("value").size().reset_index(name="n")
-    largest_classes = (
-        class_counts.sort_values(by="n", ascending=False).value[:keep_classes].tolist()
-    )
-    dropped_data = data[data.isin(largest_classes).all(axis=1)]
 
-    return dropped_data
+    def __init__(self, *args, on_color, off_color=None, state=True, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.state = state
+        self.on_color = on_color
+        # Set an alpha value for the default off color
+        self.off_color = off_color if off_color is not None else f"{on_color}20"
+        current_color = on_color if state else self.off_color
+        self.style.button_color = current_color
+        self.layout.border = f"1px dashed {self.on_color}"
+
+    def toggle(self):
+        self.state = not self.state
+        self.update()
+
+    def update(self):
+        color = self.on_color if self.state else self.off_color
+        self.style.button_color = color
diff --git a/setup.py b/setup.py
index f6a9768..e6754dc 100644
--- a/setup.py
+++ b/setup.py
@@ -35,6 +35,7 @@
         "numpy",
         "pandas",
         "plotly>=5.2.2",
+        "ipywidgets",
     ],
     project_urls={
         "Bug Reports": "https://github.com/aazuspan/sankee/issues",


Test patch:
diff --git a/tests/data.py b/tests/data.py
new file mode 100644
index 0000000..83c3e9d
--- /dev/null
+++ b/tests/data.py
@@ -0,0 +1,21 @@
+import ee
+import pandas as pd
+
+import sankee
+
+ee.Initialize()
+
+TEST_DATASET = sankee.datasets.LCMS_LU
+TEST_IMAGE_LIST = [
+    TEST_DATASET.get_year(1985),
+    TEST_DATASET.get_year(2010),
+]
+TEST_IMAGE_LABELS = ["1985", "2010"]
+TEST_REGION = ee.Geometry.Point([-122.192688, 46.25917]).buffer(1000)
+
+TEST_DATA = pd.DataFrame(
+    {
+        "start": [1, 1, 1, 2, 2, 4],
+        "end": [1, 1, 1, 2, 3, 4],
+    }
+)
diff --git a/tests/test.py b/tests/test.py
deleted file mode 100644
index 64a7598..0000000
--- a/tests/test.py
+++ /dev/null
@@ -1,382 +0,0 @@
-import unittest
-
-import ee
-import numpy as np
-import pandas as pd
-from pandas.testing import assert_frame_equal
-
-import sankee
-
-ee.Initialize()
-
-
-TEST_REGION = ee.Geometry.Polygon(
-    [
-        [
-            [-130.07633697303308, 55.018689630077546],
-            [-130.07633697303308, 54.340663930662444],
-            [-128.81565582068933, 54.340663930662444],
-            [-128.81565582068933, 55.018689630077546],
-        ]
-    ]
-)
-TEST_IMG_LIST = [ee.Image("MODIS/006/MCD12Q1/2001_01_01"), ee.Image("MODIS/006/MCD12Q1/2010_01_01")]
-TEST_LABEL_LIST = ["2001", "2010"]
-TEST_DATASET = sankee.datasets.MODIS_LC_TYPE1
-TEST_BAND = "LC_Type1"
-TEST_LABELS = {
-    1: "Evergreen conifer forest",
-    2: "Evergreen broadleaf forest",
-    3: "Deciduous conifer forest",
-    4: "Deciduous broadleaf forest",
-    5: "Mixed forest",
-    6: "Closed shrubland",
-    7: "Open shrubland",
-    8: "Woody savanna",
-    9: "Savanna",
-    10: "Grassland",
-    11: "Permanent wetland",
-    12: "Cropland",
-    13: "Urban",
-    14: "Cropland and natural vegetation",
-    15: "Permanent snow and ice",
-    16: "Barren",
-    17: "Water",
-}
-TEST_PALETTE = {
-    1: "#086a10",
-    2: "#dcd159",
-    3: "#54a708",
-    4: "#78d203",
-    5: "#009900",
-    6: "#c6b044",
-    7: "#dcd159",
-    8: "#dade48",
-    9: "#fbff13",
-    10: "#b6ff05",
-    11: "#27ff87",
-    12: "#c24f44",
-    13: "#a5a5a5",
-    14: "#ff6d4c",
-    15: "#69fff8",
-    16: "#f9ffa4",
-    17: "#1c0dff",
-}
-TEST_DATA = pd.DataFrame(
-    {
-        "start": [1, 1, 1, 2, 2, 4],
-        "end": [1, 1, 1, 2, 3, 4],
-    }
-)
-
-THREE_PERIOD_TEST_DATA = pd.DataFrame(
-    {
-        "start": [1, 1, 1, 2, 2, 4],
-        "mid": [3, 3, 1, 1, 2, 4],
-        "end": [1, 1, 1, 2, 3, 4],
-    }
-)
-
-
-class TestUtils(unittest.TestCase):
-    def test_drop_small_classes(self):
-        """
-        Test that the classes with the fewest observations are dropped by drop_small_classes
-        """
-        dropped = sankee.utils.drop_small_classes(TEST_DATA, keep_classes=1)
-        target = pd.DataFrame({"start": [1, 1, 1], "end": [1, 1, 1]})
-        assert_frame_equal(dropped, target)
-
-    def test_get_missing_keys(self):
-        """
-        Test that a missing key is found from a dictionary.
-        """
-        test_dict = {"a": 0, "b": 1, "c": 2}
-        test_keys = ["a", "b", "c", "d"]
-        missing_keys = sankee.utils.get_missing_keys(test_keys, test_dict)
-        self.assertListEqual(missing_keys, ["d"])
-
-    def test_normalize_change_sum_is_correct(self):
-        """
-        Test that the sum of normalized changes is equal to the number of input classes.
-        """
-        test_data = pd.DataFrame(
-            {
-                "source": [0, 1, 1, 2, 3, 4, 4],
-                "target": [5, 6, 7, 7, 8, 9, 10],
-                "value": [1, 2, 1, 3, 1, 5, 9],
-            }
-        )
-
-        num_classes = test_data.source.unique().size
-        normalized_values = sankee.utils.normalized_change(test_data, "source", "value").values
-        self.assertAlmostEqual(num_classes, normalized_values.sum())
-
-
-class TestSankee(unittest.TestCase):
-    def test_dropna(self):
-        """
-        Test that NA values will be dropped by _clean_data
-        """
-        test_input = pd.DataFrame({0: [0.0, np.nan, 2.0], 1: [1, 2, 3]})
-        cleaned = sankee.core._clean_data(test_input)
-        target = pd.DataFrame({0: [0.0, 2.0], 1: [1, 3]})
-
-        self.assertTrue(np.array_equal(cleaned, target))
-
-    def test_exclude(self):
-        """
-        Test that values are correctly excluded by _clean_data
-        """
-        cleaned = sankee.core._clean_data(TEST_DATA, exclude=[2])
-        target = pd.DataFrame(
-            {
-                "start": [1, 1, 1, 4],
-                "end": [1, 1, 1, 4],
-            }
-        )
-
-        self.assertTrue(np.array_equal(cleaned, target))
-
-    def test_max_classes(self):
-        """
-        Test that small classes are removed to match max_classes by _clean_data
-        """
-        cleaned = sankee.core._clean_data(TEST_DATA, max_classes=2)
-        target = pd.DataFrame(
-            {
-                "start": [1, 1, 1, 2],
-                "end": [1, 1, 1, 2],
-            }
-        )
-
-        self.assertTrue(np.array_equal(cleaned, target))
-
-    def test_mismatched_label_list(self):
-        """
-        If the label list is a different length than the image list, a ValueError should be raised.
-        """
-        with self.assertRaisesRegex(ValueError, "number of labels must match"):
-            sankee.sankify(
-                image_list=TEST_IMG_LIST,
-                region=TEST_REGION,
-                label_list=["2001", "2010", "2020"],
-                band=TEST_BAND,
-                labels=TEST_LABELS,
-                palette=TEST_PALETTE,
-            )
-        with self.assertRaisesRegex(ValueError, "number of labels must match"):
-            sankee.sankify(
-                image_list=TEST_IMG_LIST,
-                region=TEST_REGION,
-                label_list=["2001"],
-                band=TEST_BAND,
-                labels=TEST_LABELS,
-                palette=TEST_PALETTE,
-            )
-
-    def test_duplicate_label_list(self):
-        """
-        Duplicate values in the label list should raise an error
-        """
-        with self.assertRaisesRegex(ValueError, "must be unique"):
-            sankee.sankify(
-                image_list=TEST_IMG_LIST,
-                region=TEST_REGION,
-                label_list=["2001", "2001"],
-                band=TEST_BAND,
-                labels=TEST_LABELS,
-                palette=TEST_PALETTE,
-            )
-
-    def test_bad_band(self):
-        """
-        If a band name is passed to _collect_sample_data and it's not in the image, a ValueError should be raised.
-        """
-        labeled = sankee.core._label_images(TEST_IMG_LIST, TEST_LABEL_LIST)
-        with self.assertRaises(ValueError):
-            sankee.core._collect_sample_data(
-                labeled, region=TEST_REGION, band="bad_band", label_list=TEST_LABEL_LIST
-            )
-
-    def test_format_for_sankey_with_two_periods(self):
-        """
-        Test that mock table data for a two-period time series is correctly reformatted for Sankey plotting.
-        """
-        (
-            node_labels,
-            link_labels,
-            node_palette,
-            link_palette,
-            label,
-            source,
-            target,
-            value,
-        ) = sankee.core._format_for_sankey(TEST_DATA, TEST_LABELS, TEST_PALETTE)
-
-        self.assertEqual(node_labels, ["start", "start", "start", "end", "end", "end", "end"])
-        self.assertEqual(
-            link_labels,
-            [
-                r"100% of Evergreen conifer forest remained Evergreen conifer forest",
-                r"50% of Evergreen broadleaf forest remained Evergreen broadleaf forest",
-                r"50% of Evergreen broadleaf forest became Deciduous conifer forest",
-                r"100% of Deciduous broadleaf forest remained Deciduous broadleaf forest",
-            ],
-        )
-        self.assertEqual(
-            node_palette,
-            ["#086a10", "#dcd159", "#78d203", "#086a10", "#dcd159", "#54a708", "#78d203"],
-        )
-        self.assertEqual(link_palette, ["#086a10", "#dcd159", "#dcd159", "#78d203"])
-        self.assertEqual(
-            label,
-            [
-                "Evergreen conifer forest",
-                "Evergreen broadleaf forest",
-                "Deciduous broadleaf forest",
-                "Evergreen conifer forest",
-                "Evergreen broadleaf forest",
-                "Deciduous conifer forest",
-                "Deciduous broadleaf forest",
-            ],
-        )
-        self.assertEqual(source, [0, 1, 1, 2])
-        self.assertEqual(target, [3, 4, 5, 6])
-        self.assertEqual(value, [3, 1, 1, 1])
-
-    def test_format_for_sankey_with_three_periods(self):
-        """
-        Test that mock table data for a three-period time series is correctly reformatted for Sankey plotting.
-        """
-        (
-            node_labels,
-            link_labels,
-            node_palette,
-            link_palette,
-            label,
-            source,
-            target,
-            value,
-        ) = sankee.core._format_for_sankey(THREE_PERIOD_TEST_DATA, TEST_LABELS, TEST_PALETTE)
-
-        self.assertEqual(
-            node_labels,
-            ["start", "start", "start", "mid", "mid", "mid", "mid", "end", "end", "end", "end"],
-        )
-        self.assertEqual(
-            link_labels,
-            [
-                "33% of Evergreen conifer forest remained Evergreen conifer forest",
-                "67% of Evergreen conifer forest became Deciduous conifer forest",
-                "50% of Evergreen broadleaf forest became Evergreen conifer forest",
-                "50% of Evergreen broadleaf forest remained Evergreen broadleaf forest",
-                "100% of Deciduous broadleaf forest remained Deciduous broadleaf forest",
-                "50% of Evergreen conifer forest remained Evergreen conifer forest",
-                "50% of Evergreen conifer forest became Evergreen broadleaf forest",
-                "100% of Evergreen broadleaf forest became Deciduous conifer forest",
-                "100% of Deciduous conifer forest became Evergreen conifer forest",
-                "100% of Deciduous broadleaf forest remained Deciduous broadleaf forest",
-            ],
-        )
-        self.assertEqual(
-            node_palette,
-            [
-                "#086a10",
-                "#dcd159",
-                "#78d203",
-                "#54a708",
-                "#086a10",
-                "#dcd159",
-                "#78d203",
-                "#086a10",
-                "#dcd159",
-                "#54a708",
-                "#78d203",
-            ],
-        )
-        self.assertEqual(
-            link_palette,
-            [
-                "#086a10",
-                "#086a10",
-                "#dcd159",
-                "#dcd159",
-                "#78d203",
-                "#086a10",
-                "#086a10",
-                "#dcd159",
-                "#54a708",
-                "#78d203",
-            ],
-        )
-        self.assertEqual(
-            label,
-            [
-                "Evergreen conifer forest",
-                "Evergreen broadleaf forest",
-                "Deciduous broadleaf forest",
-                "Deciduous conifer forest",
-                "Evergreen conifer forest",
-                "Evergreen broadleaf forest",
-                "Deciduous broadleaf forest",
-                "Evergreen conifer forest",
-                "Evergreen broadleaf forest",
-                "Deciduous conifer forest",
-                "Deciduous broadleaf forest",
-            ],
-        )
-        self.assertEqual(source, [0, 0, 1, 1, 2, 4, 4, 5, 3, 6])
-        self.assertEqual(target, [4, 3, 4, 5, 6, 7, 8, 9, 7, 10])
-        self.assertEqual(value, [1, 2, 1, 1, 1, 1, 1, 1, 2, 1])
-
-    def test_build_labels(self):
-        """
-        The _build_labels function should not allow duplicate labels within periods
-        but should allow duplicate labels between periods.
-        """
-        data = pd.DataFrame(
-            {
-                "source": [0, 0, 1, 2],
-                "source_label": ["zero", "zero", "one", "two"],
-                "target": [3, 3, 4, 5],
-                "target_label": ["zero", "zero", "four", "five"],
-            }
-        )
-
-        labels = sankee.core._build_labels(data)
-
-        self.assertListEqual(labels, ["zero", "one", "two", "zero", "four", "five"])
-
-    def test_build_link_label(self):
-        """
-        Test that descriptive link labels are built correctly whether classes
-        changed or remained the same.
-        """
-        change_df = pd.DataFrame(
-            {"source_label": ["a", "a"], "target_label": ["b", "a"], "change": [0.8, 0.2]}
-        )
-        link_labels = sankee.core._build_link_labels(change_df)
-
-        self.assertEqual(link_labels[0], "80% of a became b")
-        self.assertEqual(link_labels[1], "20% of a remained a")
-
-    def test_build_node_labels(self):
-        """
-        The _build_node_labels function should return one period label for each
-        unique value in the source and target fields.
-        """
-        data = pd.DataFrame(
-            {
-                "source": [0, 0, 0, 0],
-                "source_period": ["start", "start", "start", "start"],
-                "target": [1, 2, 3, 4],
-                "target_period": ["end", "end", "end", "end"],
-            }
-        )
-        node_labels = sankee.core._build_node_labels(data)
-        self.assertListEqual(node_labels, ["start", "end", "end", "end", "end"])
-
-
-if __name__ == "__main__":
-    unittest.main()
diff --git a/tests/test_sampling.py b/tests/test_sampling.py
new file mode 100644
index 0000000..edc1fc8
--- /dev/null
+++ b/tests/test_sampling.py
@@ -0,0 +1,48 @@
+import ee
+import pytest
+
+import sankee.sampling
+
+from .data import TEST_DATASET, TEST_IMAGE_LABELS, TEST_IMAGE_LIST, TEST_REGION
+
+ee.Initialize()
+
+
+def test_collect_sankey_data():
+    """Test that the correct number of rows and columns are sampled."""
+    n = 10
+
+    data = sankee.sampling.collect_sankey_data(
+        image_list=TEST_IMAGE_LIST,
+        image_labels=TEST_IMAGE_LABELS,
+        region=TEST_REGION,
+        band=TEST_DATASET.band,
+        n=n,
+    )
+
+    assert len(data) == n
+    assert data.columns.tolist() == TEST_IMAGE_LABELS
+
+
+def test_collect_sankey_data_bad_band():
+    """Test that an error is thrown when sampling an invalid band."""
+    with pytest.raises(ValueError, match="band `foo` was not found"):
+        sankee.sampling.collect_sankey_data(
+            image_list=TEST_IMAGE_LIST,
+            image_labels=TEST_IMAGE_LABELS,
+            region=TEST_REGION,
+            band="foo",
+            n=10,
+        )
+
+
+def test_collect_sankey_data_bad_region():
+    """Test that an error is thrown when sampling occurs outside the dataset."""
+    with pytest.raises(ValueError, match="samples were not found"):
+        sankee.sampling.collect_sankey_data(
+            image_list=TEST_IMAGE_LIST,
+            image_labels=TEST_IMAGE_LABELS,
+            region=ee.Geometry.Point(0, 0).buffer(100),
+            band=TEST_DATASET.band,
+            n=10,
+        )