Plot Method Added to Western USA Live Fuel Moisture Dataset

tests/datasets/test_western_usa_live_fuel_moisture.py

@@ -4,9 +4,11 @@
 import os
 from pathlib import Path
 
+import matplotlib.pyplot as plt
 import pytest
 import torch
 import torch.nn as nn
+from matplotlib.figure import Figure
 from pytest import MonkeyPatch
 
 from torchgeo.datasets import DatasetNotFoundError, WesternUSALiveFuelMoisture
@@ -40,3 +42,16 @@ def test_already_downloaded(self, dataset: WesternUSALiveFuelMoisture) -> None:
     def test_not_downloaded(self, tmp_path: Path) -> None:
         with pytest.raises(DatasetNotFoundError, match='Dataset not found'):
             WesternUSALiveFuelMoisture(tmp_path)
+
+    def test_plot(self, dataset: WesternUSALiveFuelMoisture) -> None:
+        sample = dataset[0]
+
+        # Test with a single variable - likely one of the missing lines
+        fig = dataset.plot(sample, variables_to_plot=['vv'])
+        assert isinstance(fig, Figure)
+        plt.close()
+
+        # Test with both suptitle and show_titles=False
+        fig = dataset.plot(sample, show_titles=False, suptitle='Custom title')
+        assert isinstance(fig, Figure)
+        plt.close()

torchgeo/datasets/western_usa_live_fuel_moisture.py

@@ -9,8 +9,10 @@
 from collections.abc import Callable, Iterable
 from typing import Any
 
+import matplotlib.pyplot as plt
 import pandas as pd
 import torch
+from matplotlib.figure import Figure
 
 from .errors import DatasetNotFoundError
 from .geo import NonGeoDataset
@@ -297,3 +299,116 @@ def _download(self) -> None:
         os.makedirs(self.root, exist_ok=True)
         azcopy = which('azcopy')
         azcopy('sync', self.url, self.root, '--recursive=true')
+
+    def plot(
+        self,
+        sample: dict[str, Any],
+        variables_to_plot: list[str] | None = None,
+        show_titles: bool = True,
+        suptitle: str | None = None,
+    ) -> Figure:
+        """Plot a time series visualization of the LFMC sample.
+
+        Args:
+            sample: a sample returned by :meth:`__getitem__`
+            variables_to_plot: a list of valid variable to be drawn in the plot
+            show_titles: flag indicating whether to show titles above each panel
+            suptitle: optional suptitle to use for the Figure
+
+        Returns:
+            a matplotlib Figure with the rendered sample
+
+        .. versionadded:: 0.8
+        """
+        if not variables_to_plot:
+            variables_to_plot = [
+                'slope',
+                'elevation',
+                'canopy_height',
+                'forest_cover',
+                'silt',
+                'sand',
+                'clay',
+                'vv',
+                'vh',
+                'red',
+                'green',
+                'blue',
+                'swir',
+                'nir',
+                'ndvi',
+                'ndwi',
+                'nirv',
+                'vv_red',
+                'vv_green',
+                'vv_blue',
+                'vv_swir',
+                'vv_nir',
+                'vv_ndvi',
+                'vv_ndwi',
+                'vv_nirv',
+                'vh_red',
+                'vh_green',
+                'vh_blue',
+                'vh_swir',
+                'vh_nir',
+                'vh_ndvi',
+                'vh_ndwi',
+                'vh_nirv',
+                'vh_vv',
+            ]
+
+        input_data = sample['input'].numpy()
+
+        # Time points to display on x-axis
+        time_labels = ['t', 't-1', 't-2', 't-3']
+
+        fig, axs = plt.subplots(
+            len(variables_to_plot),
+            1,
+            figsize=(6, 1.5 * len(variables_to_plot)),
+            sharex=True,
+        )
+
+        # Handle single subplot case
+        if len(variables_to_plot) == 1:
+            axs = [axs]
+
+        for i, var_base_name in enumerate(variables_to_plot):
+            values = []
+
+            # Extract data for each time point (t, t-1, t-2, t-3)
+            for t_label in time_labels:
+                full_var_name = f'{var_base_name}({t_label})'
+                var_position = self.all_variable_names.index(full_var_name)
+                values.append(input_data[var_position])
+
+            axs[i].plot(range(len(time_labels)), values, 'o-')
+            axs[i].grid(True, alpha=0.3)
+
+            if show_titles:
+                axs[i].set_title(f'{var_base_name.upper()}')
+
+        axs[-1].set_xticks(range(len(time_labels)))
+        axs[-1].set_xticklabels(time_labels)
+
+        # add coordinate and label information below the plot
+        lon = input_data[-2]
+        lat = input_data[-1]
+        lfmc_value = sample['label'].item()
+
+        axs[-1].text(
+            x=0.5,
+            y=-0.7,
+            s=f'Live Fuel Moisture Content\nat {lon:.4f}, {lat:.4f}: {lfmc_value:.2f}%',
+            ha='center',
+            transform=axs[-1].transAxes,
+        )
+
+        if suptitle is not None:
+            fig.suptitle(t=suptitle, y=1.6, transform=axs[0].transAxes)
+
+        # fig.tight_layout()
+        plt.tight_layout()
+
+        return fig