Decouple DLC dataframe format from split_individuals in to_dlc_file (revives feature from #314, replaces stale PR #529)

docs/source/conf.py

@@ -6,6 +6,8 @@
 # -- Project information -----------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information
 
+"""Sphinx configuration for the movement documentation."""
+
 import os
 import sys
 
@@ -119,12 +121,14 @@
         "dependencies": ["../../.binder/requirements.txt"],
     },
     "reference_url": {"movement": None},
-    "default_thumb_file": "source/_static/data_icon.png",  # default thumbnail image
-    # Do not render config comments with the pattern # sphinx_gallery_config [= value]
+    "default_thumb_file": "source/_static/data_icon.png",
+    # default thumbnail image
+    # Do not render config comments with the pattern
+    # sphinx_gallery_config [= value]
     "remove_config_comments": True,
     # Mini-galleries config, see https://sphinx-gallery.github.io/stable/configuration.html#add-mini-galleries-for-api-documentation
-    "backreferences_dir": "api/backreferences",  # directory where function/class granular galleries are stored
-    "doc_module": ("movement",),  # module for which to generate mini-galleries
+    "backreferences_dir": "api/backreferences",
+    # directory where function/class granular galleries are stored
 }
 
 
@@ -274,7 +278,10 @@
 
 <p>Sorry, we couldn't find that page.</p>
 
-<p>We occasionally restructure the movement website, and some links may have broken.</p> 
+<p>
+We occasionally restructure the movement website,
+and some links may have broken.
+</p>
 
 <p>Try using the search box or go to the homepage.</p>
 """,

movement/io/save_poses.py

@@ -88,7 +88,9 @@
 
 
 def to_dlc_style_df(
-    ds: xr.Dataset, split_individuals: bool = False
+    ds: xr.Dataset,
+    split_individuals: bool = False,
+    dlc_df_format: Literal["single-animal", "multi-animal"] = "multi-animal",
 ) -> pd.DataFrame | dict[str, pd.DataFrame]:
     """Convert a ``movement`` dataset to DeepLabCut-style DataFrame(s).
 
@@ -99,73 +101,96 @@
         and associated metadata.
     split_individuals
         If True, return a dictionary of DataFrames per individual, with
-        individual names as keys. If False (default), return a single
-        DataFrame for all individuals (see Notes).
+        individual names as keys.
+    dlc_df_format
+        Specifies the DLC dataframe format. "single-animal" produces the
+        older DLC (<2.0) format without the "individuals" column level,
+        while "multi-animal" includes it (DLC >=2.0).
+        Default is "multi-animal".
 
     Returns
     -------
     pandas.DataFrame or dict
         DeepLabCut-style pandas DataFrame or dictionary of DataFrames.
 
-    Notes
-    -----
-    The DataFrame(s) will have a multi-index column with the following levels:
-    "scorer", "bodyparts", "coords" (if split_individuals is True),
-    or "scorer", "individuals", "bodyparts", "coords"
-    (if split_individuals is False).
-
-    For 2D data, regardless of the provenance of the points-wise confidence
-    scores, they will be referred to as "likelihood", and stored in
-    the "coords" level as DeepLabCut expects.
-
-    For 3D data, the "coords" level will only contain "x", "y", and "z",
-    as DeepLabCut does not currently provide 3D likelihoods.
-
-    See Also
-    --------
-    to_dlc_file : Save dataset directly to a DeepLabCut-style .h5 or .csv file.
-
     """
     ValidPosesInputs.validate(ds)
+
+    if dlc_df_format not in ["single-animal", "multi-animal"]:
+        raise logger.error(
+            ValueError(
+                f"Invalid value for 'dlc_df_format': {dlc_df_format}. "
+                "Expected 'single-animal' or 'multi-animal'."
+            )
+        )
+
     scorer = ["movement"]
     bodyparts = ds.coords["keypoints"].data.tolist()
     base_coords = ds.coords["space"].data.tolist()
+
     coords = (
         base_coords
         if "z" in ds.coords["space"]
         else base_coords + ["likelihood"]
     )
+
     individuals = ds.coords["individuals"].data.tolist()
+
     if split_individuals:
-        df_dict = {}
+        df_dict: dict[str, pd.DataFrame] = {}
+
         for individual in individuals:
             individual_data = ds.sel(individuals=individual)
+
             index_levels = ["scorer", "bodyparts", "coords"]
+
             columns = pd.MultiIndex.from_product(
                 [scorer, bodyparts, coords], names=index_levels
             )
+
             df = _ds_to_dlc_style_df(individual_data, columns)
+
             df_dict[individual] = df
-        logger.info(
-            "Converted poses dataset to DeepLabCut-style DataFrames "
-            "per individual."
-        )
+
+            logger.info(
+                f"Converted poses for individual {individual} to DataFrame."
+            )
+
         return df_dict
+
     else:
-        index_levels = ["scorer", "individuals", "bodyparts", "coords"]
-        columns = pd.MultiIndex.from_product(
-            [scorer, individuals, bodyparts, coords], names=index_levels
-        )
+        if dlc_df_format == "multi-animal":
+            index_levels = ["scorer", "individuals", "bodyparts", "coords"]
+
+            columns = pd.MultiIndex.from_product(
+                [scorer, individuals, bodyparts, coords],
+                names=index_levels,
+            )
+
+        else:  # single-animal format
+            index_levels = ["scorer", "bodyparts", "coords"]
+
+            columns = pd.MultiIndex.from_product(
+                [scorer, bodyparts, coords],
+                names=index_levels,
+            )
+
         df_all = _ds_to_dlc_style_df(ds, columns)
+
         logger.info("Converted poses dataset to DeepLabCut-style DataFrame.")
+
         return df_all
 
+        raise RuntimeError("Unexpected state in to_dlc_style_df.")
+
 
+# noqa: C901
 def to_dlc_file(
     ds: xr.Dataset,
     file_path: str | Path,
     split_individuals: bool | Literal["auto"] = "auto",
-) -> None:
+    dlc_df_format: Literal["single-animal", "multi-animal"] = "multi-animal",
+):
     """Save a ``movement`` dataset to DeepLabCut file(s).
 
     Parameters
@@ -174,40 +199,26 @@
         ``movement`` dataset containing pose tracks, confidence scores,
         and associated metadata.
     file_path
-        Path to the file to save the poses to. The file extension
-        must be either .h5 (recommended) or .csv.
+        Path to the file to save the poses to.
     split_individuals
-        Whether to save individuals to separate files or to the same file
-        (see Notes). Defaults to "auto".
-
-    Notes
-    -----
-    If ``split_individuals`` is True, each individual will be saved to a
-    separate file, formatted as in a single-animal DeepLabCut project
-    (without the "individuals" column level). The individual's name will be
-    appended to the file path, just before the file extension, e.g.
-    "/path/to/filename_individual1.h5". If False, all individuals will be
-    saved to the same file, formatted as in a multi-animal DeepLabCut project
-    (with the "individuals" column level). The file path will not be modified.
-    If "auto", the argument's value is determined based on the number of
-    individuals in the dataset: True if there is only one, False otherwise.
-
-    See Also
-    --------
-    to_dlc_style_df : Convert dataset to DeepLabCut-style DataFrame(s).
-
-    Examples
-    --------
-    >>> from movement.io import save_poses, load_poses
-    >>> ds = load_poses.from_sleap_file("/path/to/file_sleap.analysis.h5")
-    >>> save_poses.to_dlc_file(ds, "/path/to/file_dlc.h5")
+        Whether to save individuals to separate files or the same file.
+    dlc_df_format
+        Specifies the DLC dataframe format ("single-animal" or "multi-animal").
+        Default is "multi-animal".
 
     """
     valid_path = validate_file_path(
         file_path, permission="w", suffixes={".csv", ".h5"}
     )
 
-    # Sets default behaviour for the function
+    if dlc_df_format not in ["single-animal", "multi-animal"]:
+        raise logger.error(
+            ValueError(
+                f"Invalid value for 'dlc_df_format': {dlc_df_format}. "
+                "Expected 'single-animal' or 'multi-animal'."
+            )
+        )
+
     if split_individuals == "auto":
         split_individuals = _auto_split_individuals(ds)
 
@@ -219,21 +230,42 @@
             )
         )
 
+    # Handle single-individual edge case
+    if split_individuals and ds.sizes.get("individuals", 1) == 1:
+        import warnings
+
+        warnings.warn(
+            "split_individuals=True ignored because dataset contains "
+            "only one individual.",
+            stacklevel=2,
+        )
+        split_individuals = False
+
     if split_individuals:
-        # split the dataset into a dictionary of dataframes per individual
-        df_dict = to_dlc_style_df(ds, split_individuals=True)
+        df_dict = to_dlc_style_df(
+            ds,
+            split_individuals=True,
+            dlc_df_format="single-animal",
+        )
 
         for key, df in df_dict.items():
-            # the key is the individual's name
             filepath = f"{valid_path.with_suffix('')}_{key}{valid_path.suffix}"
+
             if isinstance(df, pd.DataFrame):
                 _save_dlc_df(Path(filepath), df)
-            logger.info(f"Saved poses for individual {key} to {valid_path}.")
+
+            logger.info(f"Saved poses for individual {key} to {filepath}.")
+
     else:
-        # convert the dataset to a single dataframe for all individuals
-        df_all = to_dlc_style_df(ds, split_individuals=False)
+        df_all = to_dlc_style_df(
+            ds,
+            split_individuals=False,
+            dlc_df_format=dlc_df_format,
+        )
+
         if isinstance(df_all, pd.DataFrame):
             _save_dlc_df(valid_path, df_all)
+
         logger.info(f"Saved poses dataset to {valid_path}.")
 
 

pyproject.toml

@@ -17,7 +17,7 @@ dependencies = [
   "numpy>=2.0.0",
   "pandas",
   "h5py",
-  "netCDF4<1.7.3",
+  "netCDF4>=1.7.2,<1.7.3",
   "tables>=3.10.1",
   "attrs",
   "pooch",