examples/torch_dataloader.py

# Example with all options:
#   python examples/torch_dataloader.py --root . --split all --batch-size 2 --num-workers 0 --tile-size 128 --patch-stride 128 --max-land-fraction 0.30 --date-start 20000101 --date-end 20000101 --max-dates 1 --include-salinity --metadata-cache-dir /tmp/depthdif_cache --require-argo --accepted-argo-qc-flags 1 2

from __future__ import annotations

import argparse
from pathlib import Path
import sys
from typing import Any

REPO_ROOT = Path(__file__).resolve().parents[1]
if str(REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(REPO_ROOT))

from depthdif_dataset import ArgoGeoTIFFGriddedPatchDataset, build_dataloader


def _shape_or_value(value: Any) -> Any:
    """Return tensor shapes for compact terminal output."""
    return tuple(value.shape) if hasattr(value, "shape") else value


def main() -> None:
    """Open the packaged dataset and print the first PyTorch batch."""
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--root", type=Path, default=REPO_ROOT)
    parser.add_argument("--split", choices=("all", "train", "val"), default="all")
    parser.add_argument("--batch-size", type=int, default=2)
    parser.add_argument("--num-workers", type=int, default=0)
    parser.add_argument("--tile-size", type=int, default=128)
    parser.add_argument("--patch-stride", type=int, default=128)
    parser.add_argument("--max-land-fraction", type=float, default=0.30)
    parser.add_argument("--date-start", type=int, default=None)
    parser.add_argument("--date-end", type=int, default=None)
    parser.add_argument("--max-dates", type=int, default=1)
    parser.add_argument("--include-salinity", action="store_true")
    parser.add_argument(
        "--metadata-cache-dir",
        type=Path,
        default=None,
        help="Optional cache directory for patch/date metadata CSVs.",
    )
    parser.add_argument(
        "--require-argo",
        action="store_true",
        help="Filter rows to patches with ARGO profiles; this may scan the compact ARGO store on first use.",
    )
    parser.add_argument(
        "--accepted-argo-qc-flags",
        type=int,
        nargs="+",
        default=(1, 2),
        help="Whole-profile ARGO QC codes accepted for compact temperature and salinity profiles.",
    )
    parser.add_argument(
        "--no-filter-bad-argo-quality",
        action="store_false",
        dest="filter_bad_argo_quality",
        help="Disable filtering by compact ARGO profile QC markers.",
    )
    args = parser.parse_args()

    dataset = ArgoGeoTIFFGriddedPatchDataset(
        geotiff_root_dir=args.root,
        split=args.split,
        tile_size=args.tile_size,
        patch_stride=args.patch_stride,
        max_land_fraction=args.max_land_fraction,
        date_start=args.date_start,
        date_end=args.date_end,
        max_dates=args.max_dates,
        include_salinity=args.include_salinity,
        require_argo_for_train=args.require_argo,
        require_argo_for_val=args.require_argo,
        require_argo_for_all=args.require_argo,
        count_argo_support=args.require_argo,
        filter_bad_argo_quality=args.filter_bad_argo_quality,
        accepted_argo_qc_flags=args.accepted_argo_qc_flags,
        metadata_cache_dir=args.metadata_cache_dir,
    )
    loader = build_dataloader(
        dataset,
        batch_size=args.batch_size,
        num_workers=args.num_workers,
        shuffle=True,
    )
    batch = next(iter(loader))

    print(f"dataset rows: {len(dataset)}")
    print(f"depth levels: {len(dataset.depth_axis_m)}")
    print(f"filter bad ARGO quality: {dataset.filter_bad_argo_quality}")
    print(f"accepted ARGO QC flags: {dataset.accepted_argo_qc_flags}")
    print(
        f"date coverage in this run: {dataset.available_dates[0]}..{dataset.available_dates[-1]}"
    )
    for key, value in batch.items():
        if key == "info":
            continue
        print(f"{key}: {_shape_or_value(value)}")


if __name__ == "__main__":
    main()