dataset

dataset

MarsRecon dataset package — HiRISE imagery and DTM datasets.

Public API re-exports the user-facing classes so callers can write:

from dataset import MarsHiRISE, MarsHiRISEDTM, HiRISEGeoSampler

Internal modules are organized into subpackages:

• core/ — dataset classes (base, rdr, dtm)
• sampling/ — geo-sampler and patch-packing geometry
• preprocessing/ — JP2/IMG → COG conversion
• stats/ — dataset statistics (per-channel, histograms)
• validation/ — diagnostic visualizations for sampler/coverage

MarsHiRISEBase

MarsHiRISEBase(root: Path, *, split: str = 'train', target: str | None = None, transforms: Callable[[Sample], Sample] | None = None, download: bool = False, bbox: tuple[float, float, float, float] | None = None, checksum: bool = False, reuse_cache: bool = True, res: float = 1.0 / 118502.26464032)

Bases: GeoDataset

Abstract base for Mars HiRISE-derived datasets.

Provides the common skeleton for index downloading/parsing, spatial filtering, parallel footprint extraction, async image downloading, coverage visualisation, and the verify/cache lifecycle.

Subclasses must implement the hooks marked @abstractmethod.

Source code in src/dataset/core/base.py

def __init__(
        self,
        root: Path,
        *,
        split: str = "train",
        target: str | None = None,
        transforms: Callable[[Sample], Sample] | None = None,
        download: bool = False,
        bbox: tuple[float, float, float, float] | None = None,
        checksum: bool = False,
        reuse_cache: bool = True,
        res: float = 1.0 / 118_502.26464032,
) -> None:
    super(GeoDataset, self).__init__()

    self.root = pathlib.Path(root)
    self.bbox = bbox
    self.split = split
    self.target = target
    self.transforms = transforms
    self.download = download
    self.checksum = checksum
    self.reuse_cache = reuse_cache

    self.res: float = res
    self._crs = MARS_GEOGRAPHIC_CRS

    self.index: gpd.GeoDataFrame | None = None
    self._raw_index: pd.DataFrame | None = None

    self._verify()

spatial_index_cache property

spatial_index_cache: Path

Path to the GeoPackage cache file.

Incorporates target, bbox, and subclass-specific suffix parts to avoid stale-cache collisions when configuration changes.

prefer_cog staticmethod

prefer_cog(path: Path | None) -> pathlib.Path | None

Return COG sidecar (.tif) for path if it exists, else path.

Source code in src/dataset/core/base.py

@staticmethod
def prefer_cog(path: pathlib.Path | None) -> pathlib.Path | None:
    """Return COG sidecar (.tif) for *path* if it exists, else *path*."""
    if path is None:
        return None
    cog = path.with_suffix(".tif")
    return cog if cog.exists() else path

merge_tiles staticmethod

merge_tiles(tiles: list[Tensor]) -> torch.Tensor

Mosaic co-registered tiles with a first-non-zero-wins strategy.

Source code in src/dataset/core/base.py

@staticmethod
def merge_tiles(tiles: list[torch.Tensor]) -> torch.Tensor:
    """Mosaic co-registered tiles with a first-non-zero-wins strategy."""
    if len(tiles) == 1:
        return tiles[0]

    max_h = max(t.shape[1] for t in tiles)
    max_w = max(t.shape[2] for t in tiles)
    n_ch = max(t.shape[0] for t in tiles)

    if not all(t.shape[0] == n_ch for t in tiles):
        logger.warning(
            "merge_tiles: channel mismatch (%d tiles, max %d ch) — padding.",
            len(tiles), n_ch,
        )
        padded = []
        for t in tiles:
            if t.shape[0] < n_ch:
                pad = torch.zeros(
                    n_ch - t.shape[0], t.shape[1], t.shape[2], dtype=t.dtype
                )
                t = torch.cat([t, pad])
            padded.append(t)
        tiles = padded

    merged = torch.zeros((n_ch, max_h, max_w), dtype=torch.float32)
    for tile in tiles:
        h, w = tile.shape[1], tile.shape[2]
        empty = merged[:, :h, :w] == 0.0
        merged[:, :h, :w][empty] = tile[empty]
    return merged

plot_coverage

plot_coverage(resolution: float | None = None, show_count: bool = True, suptitle: str | None = None) -> Figure

Visualise the spatial coverage of all entries in the index.

The plot is cropped to the actual extent of the entries, with a count heatmap plus per-entry bounding boxes.

Source code in src/dataset/core/base.py

def plot_coverage(
        self,
        resolution: float | None = None,
        show_count: bool = True,
        suptitle: str | None = None,
) -> Figure:
    """Visualise the spatial coverage of all entries in the index.

    The plot is cropped to the actual extent of the entries, with
    a count heatmap plus per-entry bounding boxes.
    """
    if self.index is None or len(self.index) == 0:
        raise RuntimeError("Spatial index is empty — nothing to plot.")

    all_bounds = self.index.geometry.bounds
    lon_min = float(all_bounds["minx"].min())
    lon_max = float(all_bounds["maxx"].max())
    lat_min = float(all_bounds["miny"].min())
    lat_max = float(all_bounds["maxy"].max())

    lon_span = lon_max - lon_min or 1.0
    lat_span = lat_max - lat_min or 1.0
    margin_lon = lon_span * 0.05
    margin_lat = lat_span * 0.05
    lon_min -= margin_lon
    lon_max += margin_lon
    lat_min -= margin_lat
    lat_max += margin_lat

    if resolution is None:
        resolution = min(lon_span, lat_span) / 20.0

    coverage, lon_edges, lat_edges = self._coverage_grid(
        resolution=resolution,
        lon_min=lon_min, lon_max=lon_max,
        lat_min=lat_min, lat_max=lat_max,
    )
    lon_centers = (lon_edges[:-1] + lon_edges[1:]) / 2
    lat_centers = (lat_edges[:-1] + lat_edges[1:]) / 2

    n_obs = len(self.index)
    covered_cells = int(np.count_nonzero(coverage))
    total_cells = coverage.size
    pct_covered = 100.0 * covered_cells / total_cells if total_cells else 0.0
    median_overlap = (
        float(np.median(coverage[coverage > 0])) if covered_cells else 0.0
    )
    max_overlap = int(coverage.max())

    fig = plt.figure(figsize=(16, 10))
    gs = fig.add_gridspec(
        2, 2,
        width_ratios=[4, 1], height_ratios=[1, 4],
        hspace=0.05, wspace=0.05,
    )
    ax_main = fig.add_subplot(gs[1, 0])
    ax_top = fig.add_subplot(gs[0, 0], sharex=ax_main)
    ax_right = fig.add_subplot(gs[1, 1], sharey=ax_main)

    display = np.ma.masked_equal(coverage, 0)
    norm = LogNorm(vmin=1, vmax=max(max_overlap, 1)) if show_count else None
    im = ax_main.pcolormesh(
        lon_edges, lat_edges, display,
        cmap="YlOrRd", norm=norm, shading="flat",
        rasterized=True, zorder=1,
    )
    ax_main.pcolormesh(
        lon_edges, lat_edges,
        np.ma.masked_not_equal(coverage, 0),
        cmap="Greys", vmin=0, vmax=1,
        shading="flat", rasterized=True, zorder=0,
    )

    cmap_boxes = matplotlib.colormaps.get_cmap("tab20")
    rects = []
    for geom in self.index.geometry:
        b = geom.bounds
        rects.append(Rectangle((b[0], b[1]), b[2] - b[0], b[3] - b[1]))

    box_colors = [cmap_boxes(i % 20) for i in range(n_obs)]
    pc = PatchCollection(
        rects,
        facecolors=[(r, g, b, 0.08) for r, g, b, _ in box_colors],
        edgecolors=[(r, g, b, 0.85) for r, g, b, _ in box_colors],
        linewidths=0.6, zorder=2,
    )
    ax_main.add_collection(pc)

    ax_main.set_xlim(lon_min, lon_max)
    ax_main.set_ylim(lat_min, lat_max)
    ax_main.set_xlabel("Longitude (°E, normalised to [−180, 180])")
    ax_main.set_ylabel("Latitude (°)")

    def _nice_ticks(lo: float, hi: float, n: int = 6) -> np.ndarray:
        span = hi - lo
        raw_step = span / n
        magnitude = 10 ** np.floor(np.log10(raw_step))
        for candidate in [1, 2, 2.5, 5, 10]:
            step = candidate * magnitude
            if span / step <= n + 1:
                break
        start = np.ceil(lo / step) * step
        return np.arange(start, hi + step * 0.5, step)

    ax_main.set_xticks(_nice_ticks(lon_min, lon_max))
    ax_main.set_yticks(_nice_ticks(lat_min, lat_max))
    ax_main.grid(color="white", linewidth=0.3, alpha=0.5, zorder=3)

    if show_count:
        cbar = fig.colorbar(im, ax=ax_main, pad=0.01, fraction=0.025)
        cbar.set_label("Entries per cell (log scale)", fontsize=8)

    legend_handles = [
        Patch(facecolor="#bbbbbb", label="No coverage"),
        Patch(facecolor=plt.cm.YlOrRd(0.15), label="Low overlap"),
        Patch(facecolor=plt.cm.YlOrRd(0.85), label="High overlap"),
    ]
    ax_main.legend(
        handles=legend_handles, loc="lower left",
        fontsize=7, framealpha=0.75,
    )

    ax_top.bar(
        lon_centers, coverage.sum(axis=0),
        width=resolution, color="steelblue", alpha=0.8, linewidth=0,
    )
    ax_top.set_ylabel("Sum", fontsize=8)
    ax_top.tick_params(labelbottom=False, labelsize=7)
    ax_top.grid(axis="y", linewidth=0.4, alpha=0.5)

    ax_right.barh(
        lat_centers, coverage.sum(axis=1),
        height=resolution, color="steelblue", alpha=0.8, linewidth=0,
    )
    ax_right.set_xlabel("Sum", fontsize=8)
    ax_right.tick_params(labelleft=False, labelsize=7)
    ax_right.grid(axis="x", linewidth=0.4, alpha=0.5)

    summary = (
        f"Entries : {n_obs:,}\n"
        f"Cells covered : {pct_covered:.1f}%  "
        f"({covered_cells:,} / {total_cells:,}  @  {resolution:.4f}°/cell)\n"
        f"Overlap — median : {median_overlap:.1f}×   max : {max_overlap}×"
    )
    ax_main.text(
        0.01, 0.99, summary,
        transform=ax_main.transAxes,
        va="top", ha="left", fontsize=8,
        bbox=dict(boxstyle="round,pad=0.3", facecolor="white", alpha=0.80),
        zorder=4,
    )

    title = suptitle or (
        f"{type(self).__name__} coverage — {n_obs:,} entries"
    )
    if self.target:
        title += f"  (filter: '{self.target}')"
    fig.suptitle(title, fontsize=12, y=1.005)

    return fig

plot_global_coverage

plot_global_coverage(output_pdf: str | Path, input_svg: str | Path = 'MarsTopography.svg', lat_bounds: tuple[float, float] = (-57.0, 57.0), target_id: str = 'topopgraphy') -> None

Injects dataset coordinates directly into an SVG and exports a vector PDF.

Handles SVGs where the map is embedded as a and instantiated via a tag.

Source code in src/dataset/core/base.py

def plot_global_coverage(self,  # pragma: no cover
                         output_pdf: str | pathlib.Path,
                         input_svg: str | pathlib.Path = "MarsTopography.svg",
                         lat_bounds: tuple[float, float] = (-57.0, 57.0),
                         target_id: str = "topopgraphy"  # Note: matching your snippet's spelling
                         ) -> None:
    """Injects dataset coordinates directly into an SVG and exports a vector PDF.

    Handles SVGs where the map is embedded as a <def> and instantiated via a <use> tag.
    """
    import xml.etree.ElementTree as ET
    import math
    import cairosvg

    # 1. Parse the SVG and register namespaces
    # Keeping namespaces empty prevents ugly ns0: prefixes in the output
    ET.register_namespace('', "http://www.w3.org/2000/svg")
    ET.register_namespace('xlink', "http://www.w3.org/1999/xlink")
    tree = ET.parse(input_svg)
    root = tree.getroot()

    # Build a parent map so we can easily inject nodes side-by-side
    parent_map = {c: p for p in tree.iter() for c in p}

    # 2. Locate the <use> node
    use_node = tree.find(f".//*[@id='{target_id}']")
    if use_node is None:
        raise ValueError(f"Could not find an SVG node with id='{target_id}'")

    # 3. Find the referenced definition
    # It could be under 'href' (SVG 2) or 'xlink:href' (SVG 1.1)
    href = use_node.attrib.get("{http://www.w3.org/1999/xlink}href") or use_node.attrib.get("href")
    if not href or not href.startswith("#"):
        raise ValueError(f"The node '{target_id}' does not have a valid href pointing to a def.")

    source_id = href[1:]  # Strip the '#'
    source_node = tree.find(f".//*[@id='{source_id}']")
    if source_node is None:
        raise ValueError(f"Could not find the referenced <image> node with id='{source_id}'")

    # 4. Extract raw dimensions from the source definition
    img_w = float(source_node.attrib['width'])
    img_h = float(source_node.attrib['height'])

    # Sometimes <use> elements have their own x/y offsets in addition to the transform
    use_x = float(use_node.attrib.get('x', 0.0))
    use_y = float(use_node.attrib.get('y', 0.0))

    # 5. Setup Projection Math
    min_lat, max_lat = lat_bounds

    def lat_to_mercator(lat_deg: float) -> float:
        lat_deg = max(min(lat_deg, 89.9), -89.9)
        return math.log(math.tan(math.pi / 4.0 + math.radians(lat_deg) / 2.0))

    merc_top = lat_to_mercator(max_lat)
    merc_bottom = lat_to_mercator(min_lat)
    merc_range = merc_top - merc_bottom

    def get_svg_coords(lon: float, lat: float) -> tuple[float, float]:
        """Map lon/lat to the raw image pixel coordinates."""
        norm_x = (lon + 180.0) / 360.0
        cx = use_x + (norm_x * img_w)

        merc_y = lat_to_mercator(lat)
        norm_y = (merc_top - merc_y) / merc_range
        cy = use_y + (norm_y * img_h)

        return cx, cy

    # 6. Create SVG Groups
    data_group = ET.Element('g', id="hirise_data_layers")

    # CRITICAL STEP: Copy the transform matrix from the <use> node to our data group
    # This forces the SVG renderer to align our dots with the scaled/moved image
    if 'transform' in use_node.attrib:
        data_group.set('transform', use_node.attrib['transform'])

    # Note: Because the transform matrix scales everything down (e.g., ~0.47x),
    # we need to increase the circle radius so they don't become microscopic.

    # Use "fill-opacity" instead of "opacity" to prevent CairoSVG from rasterizing the vectors.
    raw_group = ET.SubElement(data_group, 'g', id="layer_raw_index", fill="blue", stroke="black",
                              **{"fill-opacity": "0.7"})
    filtered_group = ET.SubElement(data_group, 'g', id="layer_filtered_index", fill="red", stroke="black",
                                   **{"stroke-width": "1.0"})

    df = self._load_cum_index()

    # 7. Inject the raw cumulative index
    if df is not None and not df.empty:
        raw_lons = ((df["MINIMUM_LONGITUDE"].astype(float) + 180.0) % 360.0) - 180.0
        raw_lats = df["MINIMUM_LATITUDE"].astype(float)

        for lon, lat in zip(raw_lons, raw_lats):
            if not (min_lat <= lat <= max_lat):
                continue
            cx, cy = get_svg_coords(lon, lat)
            ET.SubElement(raw_group, 'circle', cx=f"{cx:.2f}", cy=f"{cy:.2f}", r="2.5")

    # 8. Inject the precise filtered index points
    if self.index is not None and not self.index.empty:
        centroids = self.index.to_crs(MARS_MERCATOR_CRS).geometry.centroid.to_crs(MARS_GEOGRAPHIC_CRS)
        for lon, lat in zip(centroids.x, centroids.y):
            if not (min_lat <= lat <= max_lat):
                continue
            cx, cy = get_svg_coords(lon, lat)
            ET.SubElement(filtered_group, 'circle', cx=f"{cx:.2f}", cy=f"{cy:.2f}", r="8")

    # 9. Append the new elements right after the <use> node
    parent = parent_map[use_node]
    use_index = list(parent).index(use_node)
    parent.insert(use_index + 1, data_group)

    # 10. Save out to PDF
    import logging
    logger = logging.getLogger(__name__)

    if isinstance(output_pdf, str):
        output_pdf = Path(output_pdf)

    logger.info(f"Writing vectorized PDF to {output_pdf}")
    svg_bytes = ET.tostring(root, encoding='utf-8', method='xml')

    with open(output_pdf.with_suffix(".svg"), "wb") as f:
        f.write(svg_bytes)

    cairosvg.svg2pdf(bytestring=svg_bytes, write_to=str(output_pdf), dpi=600)

ProductMeta dataclass

ProductMeta(scaling_factor: float = _DEFAULT_SCALING_FACTOR, offset: float = _DEFAULT_OFFSET, sample_bits: int = _DEFAULT_SAMPLE_BITS, effective_max_dn: int = (1 << _EFFECTIVE_BIT_DEPTH) - 1, filter_names: list[str] = (lambda: ['NEAR-INFRARED', 'RED', 'BLUE-GREEN'])(), bands: int = 3, incidence_angle: float = 45.0, solar_azimuth: float = 270.0, emission_angle: float | None = None, phase_angle: float | None = None, local_time: float | None = None, solar_longitude: float | None = None, sub_solar_azimuth: float | None = None, north_azimuth: float | None = None, observation_id: str | None = None, product_id: str | None = None, product_version_id: str | None = None, target_name: str | None = None, mission_phase_name: str | None = None, orbit_number: int | None = None, rationale_desc: str | None = None, start_time: str | None = None, stop_time: str | None = None, product_creation_time: str | None = None, map_scale: float | None = None, map_resolution: float | None = None, map_projection_type: str | None = None, center_latitude: float | None = None, center_longitude: float | None = None, center_filter_wavelength: float | None = None)

Radiometric metadata parsed from a per-product PDS3 LBL file.

Shared across RDR and DTM ortho products. The defaults correspond to representative COLOR RDR label values; DTM ortho images should parse their own LBL which typically has 8-bit samples and different offsets.

from_lbl classmethod

from_lbl(lbl_path: Path) -> ProductMeta

Parse a PDS3 LBL and return a populated instance.

Source code in src/dataset/core/base.py

@classmethod
def from_lbl(cls, lbl_path: pathlib.Path) -> ProductMeta:
    """Parse a PDS3 LBL and return a populated instance."""
    obj = cls()
    if not lbl_path.exists():
        logger.debug("LBL not found, using defaults: %s", lbl_path)
        return obj
    try:
        text = lbl_path.read_text(errors="replace")
    except OSError as exc:
        logger.warning("Could not read LBL %s: %s", lbl_path, exc)
        return obj

    def _float(pat: str) -> float | None:
        m = re.search(pat, text, re.MULTILINE)
        return float(m.group(1)) if m else None

    def _int(pat: str) -> int | None:
        m = re.search(pat, text, re.MULTILINE)
        return int(m.group(1)) if m else None

    def _str(pat: str) -> str | None:
        m = re.search(pat, text, re.MULTILINE)
        return m.group(1).strip().strip('"').strip("'") if m else None

    if (v := _float(r"^\s*INCIDENCE_ANGLE\s*=\s*([\d.eE+\-]+)")) is not None:
        obj.incidence_angle = v
    # Match SOLAR_AZIMUTH but not SUB_SOLAR_AZIMUTH
    if (v := _float(r"^\s*(?<!SUB_)SOLAR_AZIMUTH\s*=\s*([\d.eE+\-]+)")) is not None:
        obj.solar_azimuth = v
    if (v := _float(r"^\s*EMISSION_ANGLE\s*=\s*([\d.eE+\-]+)")) is not None:
        obj.emission_angle = v
    if (v := _float(r"^\s*PHASE_ANGLE\s*=\s*([\d.eE+\-]+)")) is not None:
        obj.phase_angle = v
    if (v := _float(r"^\s*LOCAL_TIME\s*=\s*([\d.eE+\-]+)")) is not None:
        obj.local_time = v
    if (v := _float(r"^\s*SOLAR_LONGITUDE\s*=\s*([\d.eE+\-]+)")) is not None:
        obj.solar_longitude = v
    if (v := _float(r"^\s*SUB_SOLAR_AZIMUTH\s*=\s*([\d.eE+\-]+)")) is not None:
        obj.sub_solar_azimuth = v
        # If solar_azimuth wasn't set from SOLAR_AZIMUTH, fall back here.
        if not re.search(r"^\s*(?<!SUB_)SOLAR_AZIMUTH\s*=", text, re.MULTILINE):
            obj.solar_azimuth = v
    if (v := _float(r"^\s*NORTH_AZIMUTH\s*=\s*([\d.eE+\-]+)")) is not None:
        obj.north_azimuth = v

    # Identification / timing
    obj.observation_id = _str(r"^\s*OBSERVATION_ID\s*=\s*\"?([^\"\n]+?)\"?\s*$")
    obj.product_id = _str(r"^\s*PRODUCT_ID\s*=\s*\"?([^\"\n]+?)\"?\s*$")
    obj.product_version_id = _str(r"^\s*PRODUCT_VERSION_ID\s*=\s*\"?([^\"\n]+?)\"?\s*$")
    obj.target_name = _str(r"^\s*TARGET_NAME\s*=\s*\"?([^\"\n]+?)\"?\s*$")
    obj.mission_phase_name = _str(r"^\s*MISSION_PHASE_NAME\s*=\s*\"?([^\"\n]+?)\"?\s*$")
    obj.rationale_desc = _str(r"^\s*RATIONALE_DESC\s*=\s*\"?([^\"\n]+?)\"?\s*$")
    obj.start_time = _str(r"^\s*START_TIME\s*=\s*([^\s\n]+)")
    obj.stop_time = _str(r"^\s*STOP_TIME\s*=\s*([^\s\n]+)")
    obj.product_creation_time = _str(
        r"^\s*PRODUCT_CREATION_TIME\s*=\s*([^\s\n]+)"
    )
    if (v := _int(r"^\s*ORBIT_NUMBER\s*=\s*(\d+)")) is not None:
        obj.orbit_number = v

    # Map projection
    if (v := _float(r"^\s*MAP_SCALE\s*=\s*([\d.eE+\-]+)")) is not None:
        obj.map_scale = v
    if (v := _float(r"^\s*MAP_RESOLUTION\s*=\s*([\d.eE+\-]+)")) is not None:
        obj.map_resolution = v
    obj.map_projection_type = _str(
        r"^\s*MAP_PROJECTION_TYPE\s*=\s*\"?([^\"\n]+?)\"?\s*$"
    )
    if (v := _float(r"^\s*CENTER_LATITUDE\s*=\s*([\d.eE+\-]+)")) is not None:
        obj.center_latitude = v
    if (v := _float(r"^\s*CENTER_LONGITUDE\s*=\s*([\d.eE+\-]+)")) is not None:
        obj.center_longitude = v
    if (v := _float(r"^\s*CENTER_FILTER_WAVELENGTH\s*=\s*([\d.eE+\-]+)")) is not None:
        obj.center_filter_wavelength = v

    if (v := _float(r"^\s*SCALING_FACTOR\s*=\s*([\d.eE+\-]+)")) is not None:
        obj.scaling_factor = v
    if (v := _float(r"^\s*OFFSET\s*=\s*([\d.eE+\-]+)")) is not None:
        obj.offset = v
    if (v := _int(r"^\s*SAMPLE_BITS\s*=\s*(\d+)")) is not None:
        obj.sample_bits = v
    if (v := _int(r"^\s*BANDS\s*=\s*(\d+)")) is not None:
        obj.bands = v
    if m := re.search(r"SAMPLE_BIT_MASK\s*=\s*2#([01]+)#", text):
        obj.effective_max_dn = (1 << m.group(1).count("1")) - 1
    if m := re.search(r"FILTER_NAME\s*=\s*\(([^)]+)\)", text, re.DOTALL):
        names = [s.strip().strip('"').strip("'") for s in m.group(1).split(",")]
        if names:
            obj.filter_names = names
    elif m := re.search(r'FILTER_NAME\s*=\s*"?(\w[\w-]*)"?', text):
        obj.filter_names = [m.group(1)]

    return obj

MarsHiRISEDTM

MarsHiRISEDTM(root: Path = '/scratch/mars_hirise_dtm', *, split: str = 'train', target: str | None = None, include_ortho: bool = True, ortho_type: OrthoType | list[OrthoType] = 'RED', ortho_scale: str | None = None, transforms: Callable[[Sample], Sample] | None = None, download: bool = False, bbox: tuple[float, float, float, float] | None = None, checksum: bool = False, reuse_cache: bool = True, normalize_elevation: bool = False, elevation_stats_path: str | None = None, return_meta: bool = False)

Bases: MarsHiRISEBase

Mars HiRISE Digital Terrain Model (DTM) dataset.

HiRISE DTMs <https://www.uahirise.org/dtm/about.php>__ are derived from stereo pairs of HiRISE observations. Each DTM set consists of:

• A DTM (.IMG) — 32-bit floating-point raster where each pixel value is an areoid elevation (metres) or planetary radius (metres). 1 DN = 1 m.
• Orthoimages (.JP2) — the left and right stereo observations orthorectified onto the DTM, in RED (1-band) and/or IRB (3-band: near-IR, RED, blue-green) colour content, at one or more resolutions (A = 0.25 m, B = 0.5 m, C = 1.0 m, D = 2.0 m).

This class indexes the PDS DTM cumulative index (DTMCUMINDEX.TAB) and groups all products by stereo pair.

Directory layout ~~~~~~~~~~~~~~~~ ::

<root>/
    DTMCUMINDEX.LBL
    DTMCUMINDEX.TAB
    images/
        DTEEC_011265_1560_011331_1560_U01.IMG
        ESP_011265_1560_RED_A_01_ORTHO.JP2
        ESP_011265_1560_RED_A_01_ORTHO.LBL
        ...

Sample dict ~~~~~~~~~~~ __getitem__ returns a dict containing:

• "elevation" — (1, H, W) float32 in metres. Nodata = NaN.
• "left_red" — (1, H, W) float32 I/F [0, 1] (if requested)
• "right_red" — same, for the right observation
• "left_irb" — (3, H, W) float32 I/F [0, 1] (if requested)
• "right_irb" — same, for the right observation
• "bounds" tensor, "crs" WKT string

Sampler units ~~~~~~~~~~~~~ self.crs is geographic; self.res is in degrees/pixel (~1.69e-5 deg/px ≈ 1 m at the equator by default).

Dataset homepage

https://www.uahirise.org/dtm/about.php

.. versionadded:: 0.8

Initialise the dataset.

Parameters:

Name	Type	Description	Default
root	Path	Root directory containing PDS index files and data.	'/scratch/mars_hirise_dtm'
split	str	Dataset split — informational.	'train'
target	str | None	Case-insensitive substring filter on character columns of the cumulative index (e.g. "Eberswalde").	None
include_ortho	bool	If True, load orthoimage patches alongside elevation.	True
ortho_type	OrthoType | list[OrthoType]	Which orthoimage colour content to load: "RED" (1-band) and/or "IRB" (3-band).	'RED'
ortho_scale	str | None	Preferred scale letter ("A"–"D"). If None, the finest resolution available is used.	None
transforms	Callable[[Sample], Sample] | None	Optional transform applied to each sample.	None
download	bool	Fetch from PDS if absent.	False
bbox	tuple[float, float, float, float] | None	(lon_min, lat_min, lon_max, lat_max) in degrees.	None
checksum	bool	Verify checksums (not yet implemented).	False
reuse_cache	bool	Reuse cached spatial index if present.	True
normalize_elevation	bool	Z-score normalise elevation using stats.	False
elevation_stats_path	str | None	JSON file with "mean" and "std".	None

Raises:

Type	Description
ValueError	If ortho_type contains an unrecognised value.
DatasetNotFoundError	If index is absent and download=False.

Source code in src/dataset/core/dtm.py

def __init__(
        self,
        root: Path = "/scratch/mars_hirise_dtm",
        *,
        split: str = "train",
        target: str | None = None,
        include_ortho: bool = True,
        ortho_type: OrthoType | list[OrthoType] = "RED",
        ortho_scale: str | None = None,
        transforms: Callable[[Sample], Sample] | None = None,
        download: bool = False,
        bbox: tuple[float, float, float, float] | None = None,
        checksum: bool = False,
        reuse_cache: bool = True,
        normalize_elevation: bool = False,
        elevation_stats_path: str | None = None,
        return_meta: bool = False,
) -> None:
    """Initialise the dataset.

    Args:
        root: Root directory containing PDS index files and data.
        split: Dataset split — informational.
        target: Case-insensitive substring filter on character columns
            of the cumulative index (e.g. ``"Eberswalde"``).
        include_ortho: If ``True``, load orthoimage patches alongside
            elevation.
        ortho_type: Which orthoimage colour content to load:
            ``"RED"`` (1-band) and/or ``"IRB"`` (3-band).
        ortho_scale: Preferred scale letter (``"A"``–``"D"``).  If
            ``None``, the finest resolution available is used.
        transforms: Optional transform applied to each sample.
        download: Fetch from PDS if absent.
        bbox: ``(lon_min, lat_min, lon_max, lat_max)`` in degrees.
        checksum: Verify checksums (not yet implemented).
        reuse_cache: Reuse cached spatial index if present.
        normalize_elevation: Z-score normalise elevation using stats.
        elevation_stats_path: JSON file with ``"mean"`` and ``"std"``.

    Raises:
        ValueError: If ortho_type contains an unrecognised value.
        DatasetNotFoundError: If index is absent and ``download=False``.
    """
    # ── Ortho configuration (before super().__init__ calls _verify) ──
    self.return_meta = return_meta
    self.include_ortho = include_ortho
    if isinstance(ortho_type, str):
        ortho_type = [ortho_type]
    invalid = set(ortho_type) - {"RED", "IRB"}
    if invalid:
        raise ValueError(
            f"Invalid ortho_type(s): {invalid}. Valid: 'RED', 'IRB'"
        )
    self.ortho_types: list[str] = list(ortho_type)
    self.ortho_scale = ortho_scale.upper() if ortho_scale else None

    # ── Elevation normalisation ──
    self.normalize_elevation = normalize_elevation
    self._elev_mean: float | None = None
    self._elev_std: float | None = None

    if self.normalize_elevation:
        path = pathlib.Path(elevation_stats_path or "")
        if not elevation_stats_path or not path.exists():
            raise ValueError(
                f"elevation_stats_path must point to a valid stats JSON "
                f"when normalize_elevation=True (got: {path.absolute()})"
            )
        with open(path) as f:
            stats = json.load(f)
        self._elev_mean = float(stats["mean"])
        self._elev_std = float(stats["std"])

    # DTM typical resolution ≈ 1 m/pix → ~1/(59 251) deg/pix at equator.
    super().__init__(
        root,
        split=split,
        target=target,
        transforms=transforms,
        download=download,
        bbox=bbox,
        checksum=checksum,
        reuse_cache=reuse_cache,
        res=1.0 / 59_251.13,
    )

plot

plot(sample: Sample, show_titles: bool = True, suptitle: str | None = None, **kwargs) -> Figure

Visualise a sample: elevation + ortho panels.

Source code in src/dataset/core/dtm.py

def plot(
        self,
        sample: Sample,
        show_titles: bool = True,
        suptitle: str | None = None,
        **kwargs,
) -> Figure:
    """Visualise a sample: elevation + ortho panels."""
    has_elev = "elevation" in sample
    ortho_keys = [
        k for k in ("left_red", "right_red", "left_irb", "right_irb")
        if k in sample
    ]
    n_panels = max(int(has_elev) + len(ortho_keys), 1)

    fig, axes = plt.subplots(1, n_panels, figsize=(6 * n_panels, 6))
    if n_panels == 1:
        axes = [axes]

    panel = 0

    # ── Elevation ─────────────────────────────────────────────────
    if has_elev:
        ax = axes[panel]
        elev = sample["elevation"]
        if elev.ndim == 4:
            elev = elev[0]
        elev_np = elev[0].numpy().copy()
        valid = np.isfinite(elev_np)
        if valid.any():
            vmin, vmax = np.nanpercentile(elev_np[valid], [2, 98])
            im = ax.imshow(
                np.where(valid, elev_np, np.nan),
                cmap="terrain", vmin=vmin, vmax=vmax,
                interpolation="nearest",
            )
            fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04,
                         label="Elevation (m)")
        else:
            ax.imshow(elev_np, cmap="terrain", interpolation="nearest")
        ax.axis("off")
        if show_titles:
            ax.set_title("Elevation (DTM)")
        panel += 1

    # ── Ortho panels ─────────────────────────────────────────────
    for key in ortho_keys:
        ax = axes[panel]
        img = sample[key]
        if img.ndim == 4:
            img = img[0]

        if img.shape[0] >= 3:
            img_np = img[:3].permute(1, 2, 0).numpy().copy()
            cmap = None
        else:
            img_np = img[0].numpy().copy()
            cmap = "grey"

        ax.imshow(
            self._percentile_stretch(img_np),
            cmap=cmap, interpolation="nearest",
        )
        ax.axis("off")
        if show_titles:
            ax.set_title(f"Ortho — {key.replace('_', ' ').title()}")
        panel += 1

    if suptitle is not None:
        fig.suptitle(suptitle)
    fig.tight_layout()
    return fig

plot3d

plot3d(sample: Sample, show_titles: bool = True, suptitle: str | None = None, **kwargs) -> Figure

Visualise a sample: 3D elevation panel.

Source code in src/dataset/core/dtm.py

def plot3d(
        self,
        sample: Sample,
        show_titles: bool = True,
        suptitle: str | None = None,
        **kwargs,
) -> Figure:
    """Visualise a sample: 3D elevation panel."""
    has_elev = "elevation" in sample

    fig, ax = plt.subplots(1, 1, figsize=(6, 6), subplot_kw={"projection": "3d"})
    ax: Axes3D

    if has_elev:
        elev = sample["elevation"]
        if elev.ndim == 4:
            elev = elev[0]
        elev_np = elev[0].numpy().copy()
        valid = np.isfinite(elev_np)

        x, y = np.meshgrid(range(elev_np.shape[1]), range(elev_np.shape[0]))

        if valid.any():
            stride = max(1, min(elev_np.shape[0], elev_np.shape[1]) // 100)

            im = ax.plot_surface(
                np.where(valid, x, np.nan),
                np.where(valid, y, np.nan),
                np.where(valid, elev_np, np.nan),
                cmap="terrain",
                rstride=stride,
                cstride=stride,
                linewidth=0,
                antialiased=False,
            )
            fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04,
                         label="Elevation (m)")

        ax.axis("off")
        if show_titles:
            ax.set_title("3D Elevation (DTM)")

    if suptitle is not None:
        fig.suptitle(suptitle)
    fig.tight_layout()
    return fig

MarsHiRISE

MarsHiRISE(root: Path = '/scratch/mars_hirise', split: str = 'train', target: str | None = None, channels: list[Channel] | None = None, transforms: Callable[[Sample], Sample] | None = None, download: bool = False, bbox: tuple[float, float, float, float] | None = None, checksum: bool = False, reuse_cache: bool = True, normalize: bool = False, normalization_path: str | None = None, return_meta: bool = False)

Bases: MarsHiRISEBase

Mars HiRISE Reduced Data Records (RDR) dataset.

HiRISE <https://pds-imaging.jpl.nasa.gov/volumes/mro.html>__ is the High Resolution Imaging Science Experiment aboard the Mars Reconnaissance Orbiter (MRO). This dataset wraps the RDR products — radiometrically- corrected images resampled to a standard map projection — hosted on the NASA PDS Imaging Node.

Directory layout ~~~~~~~~~~~~~~~~ Files are stored preserving the PDS directory hierarchy under root::

<root>/
    RDRCUMINDEX.LBL
    RDRCUMINDEX.TAB
    images/
        PSP_001430_1780_COLOR.JP2
        PSP_001430_1780_COLOR.LBL
        PSP_001430_1780_RED.JP2
        PSP_001430_1780_RED.LBL

Each HiRISE observation produces two JP2 product files:

_COLOR.JP2 Three-band mosaic — in-file band order: NEAR-INFRARED (band 1), RED (band 2), BLUE-GREEN (band 3).

_RED.JP2 Single-band full-strip RED image. Uses more TDI lines; highest- quality RED source when colour context is not needed.

Channel selection ~~~~~~~~~~~~~~~~~ * Only "RED" requested → _RED.JP2 (higher fidelity). * Any "NEAR-INFRARED" or "BLUE-GREEN" → _COLOR.JP2.

Radiometric calibration ~~~~~~~~~~~~~~~~~~~~~~~ I/F = DN * SCALING_FACTOR + OFFSET, clipped to [0, 1].

Sampler units ~~~~~~~~~~~~~ self.crs is a geographic CRS; self.res is in degrees/pixel. Pass size to :class:~torchgeo.samplers.RandomGeoSampler in degrees (e.g. size=0.01 ≈ 1 185 px ≈ 593 m at the equator).

Dataset homepage

https://hirise-pds.lpl.arizona.edu/PDS/AAREADME.TXT

.. versionadded:: 0.7

Initialise the dataset.

Parameters:

Name	Type	Description	Default
root	Path	Root directory. Must contain the PDS index files and the observation JP2/LBL files, or download=True must be set.	'/scratch/mars_hirise'
split	str	Dataset split — informational.	'train'
target	str | None	Optional case-insensitive substring filter on all character columns of the cumulative index (e.g. "Olympus").	None
channels	list[Channel] | None	Which channels to include. Valid values: "NEAR-INFRARED", "RED", "BLUE-GREEN". Output tensor band order always follows :attr:all_channels. Defaults to all three.	None
transforms	Callable[[Sample], Sample] | None	Optional callable applied to each :class:Sample.	None
download	bool	Fetch index and images from the PDS server if absent.	False
bbox	tuple[float, float, float, float] | None	(lon_min, lat_min, lon_max, lat_max) bounding box filter in degrees ([-180, 180] longitude convention).	None
checksum	bool	Verify checksums after download (not yet implemented).	False
reuse_cache	bool	Reuse cached spatial index if present.	True
normalize	bool	If True, apply per-channel z-score normalisation using statistics from normalization_path.	False
normalization_path	str | None	Path to a JSON file with "channels", "mean", and "std" keys.	None

Raises:

Type	Description
ValueError	If channels contains an unrecognised name.
DatasetNotFoundError	If index files are absent and download=False.

Source code in src/dataset/core/rdr.py

def __init__(
        self,
        root: Path = "/scratch/mars_hirise",
        split: str = "train",
        target: str | None = None,
        channels: list[Channel] | None = None,
        transforms: Callable[[Sample], Sample] | None = None,
        download: bool = False,
        bbox: tuple[float, float, float, float] | None = None,
        checksum: bool = False,
        reuse_cache: bool = True,
        normalize: bool = False,
        normalization_path: str | None = None,
        return_meta: bool = False,
) -> None:
    """Initialise the dataset.

    Args:
        root: Root directory.  Must contain the PDS index files and the
            observation JP2/LBL files, or ``download=True`` must be set.
        split: Dataset split — informational.
        target: Optional case-insensitive substring filter on all character
            columns of the cumulative index (e.g. ``"Olympus"``).
        channels: Which channels to include.  Valid values:
            ``"NEAR-INFRARED"``, ``"RED"``, ``"BLUE-GREEN"``.  Output
            tensor band order always follows :attr:`all_channels`.
            Defaults to all three.
        transforms: Optional callable applied to each :class:`Sample`.
        download: Fetch index and images from the PDS server if absent.
        bbox: ``(lon_min, lat_min, lon_max, lat_max)`` bounding box filter
            in degrees ([-180, 180] longitude convention).
        checksum: Verify checksums after download (not yet implemented).
        reuse_cache: Reuse cached spatial index if present.
        normalize: If ``True``, apply per-channel z-score normalisation
            using statistics from *normalization_path*.
        normalization_path: Path to a JSON file with ``"channels"``,
            ``"mean"``, and ``"std"`` keys.

    Raises:
        ValueError: If *channels* contains an unrecognised name.
        DatasetNotFoundError: If index files are absent and
            ``download=False``.
    """
    self.return_meta = return_meta

    # ── Channel validation (before super().__init__ triggers _verify) ──
    if channels is None:
        self.channels: list[str] = list(ALL_CHANNELS)
    else:
        invalid = set(channels) - set(ALL_CHANNELS)
        if invalid:
            raise ValueError(
                f"Invalid channel(s): {invalid}. "
                f"Valid choices are: {ALL_CHANNELS}"
            )
        self.channels = [c for c in ALL_CHANNELS if c in set(channels)]

    # ── Normalisation ──
    self.normalize = normalize
    self.normalization_path = normalization_path
    self._normalizer = None

    if self.normalize:
        path = pathlib.Path(self.normalization_path or "")
        if not self.normalization_path or not path.exists():
            raise ValueError(
                f"normalization_path must be a valid path to "
                f"dataset_stats.json when normalize=True, currently "
                f"pointing to {path.absolute()}"
            )
        with open(self.normalization_path) as f:
            stats = json.load(f)

        stat_channels = stats["channels"]
        mean_dict = dict(zip(stat_channels, stats["mean"]))
        std_dict = dict(zip(stat_channels, stats["std"]))
        try:
            mean = [mean_dict[ch] for ch in self.channels]
            std = [std_dict[ch] for ch in self.channels]
        except KeyError as e:
            raise ValueError(
                f"Channel {e} not found in normalization stats."
            )

        self._normalizer = Normalize(mean=mean, std=std)

    # Native HiRISE RDR resolution: 1 / 118 502.26 pix/deg ≈ 8.44e-6 deg/pix
    super().__init__(
        root,
        split=split,
        target=target,
        transforms=transforms,
        download=download,
        bbox=bbox,
        checksum=checksum,
        reuse_cache=reuse_cache,
        res=1.0 / 118_502.26464032,
    )

plot

plot(sample: Sample, show_titles: bool = True, suptitle: str | None = None, eps: float = 1e-08, **kwargs) -> Figure

Visualise an RDR image sample with percentile stretch.

Source code in src/dataset/core/rdr.py

def plot(
        self,
        sample: Sample,
        show_titles: bool = True,
        suptitle: str | None = None,
        eps: float = 1e-8,
        **kwargs,
) -> Figure:
    """Visualise an RDR image sample with percentile stretch."""
    eps = abs(eps)

    image: torch.Tensor = sample["image"]
    if image.ndim == 4:
        image = image[0]

    ch = self.channels
    if set(ch) >= {"NEAR-INFRARED", "RED", "BLUE-GREEN"}:
        idx = [
            ch.index("NEAR-INFRARED"),
            ch.index("RED"),
            ch.index("BLUE-GREEN"),
        ]
        rgb = image[idx]
        # If only one channel has non-zero data (e.g. COLOR file
        # absent), fall back to grayscale.
        nonzero = [(rgb[i] > 0).any().item() for i in range(3)]
        if sum(nonzero) == 1:
            active_i = nonzero.index(True)
            active_name = ["NEAR-INFRARED", "RED", "BLUE-GREEN"][active_i]
            img_np = rgb[active_i].numpy()
            cmap = "grey"
            title = (
                f"MarsHiRISE — {active_name} (COLOR file unavailable)"
            )
        else:
            img_np = rgb.permute(1, 2, 0).numpy()
            cmap = None
            title = "MarsHiRISE — false colour (NIR→R, RED→G, BG→B)"
    else:
        img_np = image[0].numpy()
        cmap = "grey"
        title = f"MarsHiRISE — {ch[0]}"

    # Percentile stretch — only over non-zero (data) pixels
    img_out = img_np.copy()
    if img_out.ndim == 3:
        for c in range(img_out.shape[2]):
            band = img_out[..., c]
            mask = np.logical_or(band > eps, band < -eps)
            data_pixels = band[mask]
            if len(data_pixels) > 0:
                p2, p98 = np.percentile(data_pixels, [2, 98])
                if p98 > p2:
                    img_out[..., c] = np.clip(
                        (band - p2) / (p98 - p2), 0, 1
                    )
                    img_out[..., c][np.logical_not(mask)] = 0
    else:
        mask = np.logical_or(img_out > eps, img_out < -eps)
        data_pixels = img_out[mask]
        if len(data_pixels) > 0:
            p2, p98 = np.percentile(data_pixels, [2, 98])
            if p98 > p2:
                img_out = np.clip(
                    (img_out - p2) / (p98 - p2), 0, 1
                )
                img_out[np.logical_not(mask)] = 0

    fig, ax = plt.subplots(figsize=(8, 8))
    ax.imshow(img_out, cmap=cmap, interpolation="nearest")
    ax.axis("off")
    if show_titles:
        ax.set_title(title)
    if suptitle is not None:
        fig.suptitle(suptitle)
    fig.tight_layout()
    return fig

HiRISEGeoSampler

HiRISEGeoSampler(dataset: GeoDataset, size: float | tuple[float, float], *, split: Literal['train', 'test', 'val', 'all'] | None = 'train', split_fractions: tuple[float, float, float] = (0.8, 0.1, 0.1), split_method: str = 'geographic', split_axis: str = 'longitude', n_folds: int | None = None, fold_idx: int = 0, seed: int = 42, length: int | None = None, stride: float | tuple[float, float] | None = None, roi=None, toi=None, units: Units = Units.CRS, generator: Generator | None = None, min_overlap: float = 0.5, replacement: bool = False, reuse_cache: bool = True, center_mode: Literal['simple', 'optimal'] = 'optimal', patch_overlap: float = 0.0, packing_phase_steps: int = 20, valid_region_rays: int = 3)

Bases: GeoSampler

Sampler that restricts patches to within HiRISE strip polygon footprints, with built-in train/val/test splitting and K-fold cross-validation.

Construction pre-computes a set of candidate patch centres for every strip in dataset.index. Two strategies are available:

• center_mode="simple" — classical grid-over-bbox followed by an intersection-area filter. Controlled by size and stride.

• center_mode="optimal" — computes each strip's valid centre region analytically (the locus of centres where the patch is guaranteed to meet min_overlap) and packs centres inside it with a row/column sweep. Adjacent spacing is controlled by patch_overlap.

The split is performed at the stereo-pair level — entire strips are assigned to train, val, or test. This prevents spatial data leakage (nearby terrain never appears on both sides of the split).

Parameters:

Name	Type	Description	Default
dataset	GeoDataset	The :class:~MarsHiRISEDTM dataset to sample from.	required
size	float | tuple[float, float]	Patch height and width in CRS units (degrees when units=Units.CRS) or pixels (when units=Units.PIXELS). A single float sets both dimensions equal.	required
split	Literal['train', 'test', 'val', 'all'] | None	Which split to sample from: "train", "val", or "test".	'train'
split_fractions	tuple[float, float, float]	(train, val, test) fractions summing to 1.0.	(0.8, 0.1, 0.1)
split_method	str	"geographic" (spatially contiguous blocks) or "random" (uniformly at random).	'geographic'
split_axis	str	For geographic splits: "longitude" or "latitude".	'longitude'
n_folds	int | None	Number of folds for K-fold cross-validation. None disables K-fold and uses split_fractions directly.	None
fold_idx	int	Which fold to use as the test set (0 to n_folds - 1).	0
seed	int	Random seed for reproducible split assignment.	42
length	int | None	Number of patches to yield per epoch. Defaults to the total number of pre-computed valid centres for this split.	None
stride	float | tuple[float, float] | None	(simple mode only) Centre-to-centre grid spacing in the same units as size. Defaults to size (non-overlapping grid). Ignored when center_mode="optimal".	None
roi		Optional Shapely Polygon to further restrict the spatial domain.	None
toi		Optional :class:pandas.Interval to restrict the temporal domain.	None
units	Units	Whether size and stride are given in CRS units or pixels.	CRS
generator	Generator | None	Optional :class:torch.Generator for reproducible sampling.	None
min_overlap	float	Minimum fraction of patch area that must overlap the strip footprint to be considered valid (default 0.5).	0.5
replacement	bool	Sample with replacement if True.	False
reuse_cache	bool	Reuse cached split assignment if available.	True
center_mode	Literal['simple', 'optimal']	"simple" (default, backwards-compatible) or "optimal" (geometric packing).	'optimal'
patch_overlap	float	(optimal mode only) Fractional overlap between adjacent patches in [0, 1). 0.0 is edge-to-edge, 0.5 means 50% overlap in both axes. Ignored when center_mode="simple".	0.0
packing_phase_steps	int	(optimal mode only) Number of phase offsets to try per sweep direction when packing. Default 20.	20
valid_region_rays	int	(optimal mode only) Number of extra rays per polygon edge when approximating the valid centre region. Default 3.	3

Example::

from hirise_sampler import HiRISEGeoSampler
from torchgeo.samplers import Units

# Legacy behaviour (unchanged; reuses existing caches)
train_sampler = HiRISEGeoSampler(
    dataset, size=0.005, split="train",
    split_fractions=(0.8, 0.1, 0.1),
    seed=42,
)

# Dense optimal packing with 50% patch overlap
train_sampler = HiRISEGeoSampler(
    dataset, size=0.005, split="train",
    split_fractions=(0.8, 0.1, 0.1),
    seed=42,
    center_mode="optimal",
    patch_overlap=0.5,
)

# 5-fold cross-validation, fold 0 as test
train_sampler = HiRISEGeoSampler(
    dataset, size=0.005, split="train",
    n_folds=5, fold_idx=0, seed=42,
)

Source code in src/dataset/sampling/sampler.py

def __init__(
        self,
        dataset: GeoDataset,
        size: float | tuple[float, float],
        *,
        split: Literal["train", "test", "val", "all"] | None = "train",
        split_fractions: tuple[float, float, float] = (0.8, 0.1, 0.1),
        split_method: str = "geographic",
        split_axis: str = "longitude",
        n_folds: int | None = None,
        fold_idx: int = 0,
        seed: int = 42,
        length: int | None = None,
        stride: float | tuple[float, float] | None = None,
        roi=None,
        toi=None,
        units: Units = Units.CRS,
        generator: torch.Generator | None = None,
        min_overlap: float = 0.5,
        replacement: bool = False,
        reuse_cache: bool = True,
        # ── New optimal-mode parameters ──
        center_mode: Literal["simple", "optimal"] = "optimal",
        patch_overlap: float = 0.0,
        packing_phase_steps: int = 20,
        valid_region_rays: int = 3,
) -> None:
    super().__init__(dataset, roi, toi)

    # Use the whole dataset
    if split is None:
        split = "all"

    if split == "all":
        split = "train"
        split_fractions = (1.0, 0.0, 0.0)

    # ── Validate split parameters ──
    if split not in VALID_SPLITS:
        raise ValueError(
            f"Invalid split '{split}'. Must be one of {sorted(VALID_SPLITS)}."
        )
    if split_method not in VALID_SPLIT_METHODS:
        raise ValueError(
            f"Invalid split_method '{split_method}'. "
            f"Must be one of {sorted(VALID_SPLIT_METHODS)}."
        )
    if center_mode not in VALID_CENTER_MODES:
        raise ValueError(
            f"Invalid center_mode '{center_mode}'. "
            f"Must be one of {sorted(VALID_CENTER_MODES)}."
        )
    if not (0.0 <= patch_overlap < 1.0):
        raise ValueError(
            f"patch_overlap must be in [0, 1); got {patch_overlap}"
        )

    train_f, val_f, test_f = split_fractions
    if n_folds is None and abs(train_f + val_f + test_f - 1.0) > 1e-6:
        raise ValueError(
            f"split_fractions must sum to 1.0, got "
            f"{train_f} + {val_f} + {test_f} = {train_f + val_f + test_f}"
        )

    self.split = split
    self.split_fractions = split_fractions
    self.split_method = split_method
    self.split_axis = split_axis
    self.n_folds = n_folds
    self.fold_idx = fold_idx
    self.seed = seed
    self.replacement = replacement
    self.min_overlap = min_overlap

    # Optimal-mode state
    self.center_mode = center_mode
    self.patch_overlap = float(patch_overlap)
    self.packing_phase_steps = int(packing_phase_steps)
    self.valid_region_rays = int(valid_region_rays)

    # ── Resolve size / stride to (height_deg, width_deg) ──
    size_h, size_w = _to_tuple(size)
    if units == Units.PIXELS:
        xres, yres = dataset.res
        size_h *= yres
        size_w *= xres

    if stride is None:
        stride_h = size_h * (1.0 - self.patch_overlap)
        stride_w = size_w * (1.0 - self.patch_overlap)
    else:
        stride_h, stride_w = _to_tuple(stride)
        if units == Units.PIXELS:
            stride_h *= yres
            stride_w *= xres

    self.size = (size_h, size_w)
    self.stride = (stride_h, stride_w)
    self.generator = generator

    # ── Compute or load split assignments ──
    self.cache_hash = ''
    self._assignments = self._get_split_assignments(dataset, reuse_cache)

    # Log split distribution
    split_counts = {}
    for s in VALID_SPLITS:
        split_counts[s] = sum(1 for v in self._assignments.values() if v == s)
    n_total = len(self._assignments)
    logger.info(
        "Split assignment: %d stereo pairs → train=%d, val=%d, test=%d "
        "(method=%s, seed=%d%s)",
        n_total,
        split_counts.get("train", 0),
        split_counts.get("val", 0),
        split_counts.get("test", 0),
        split_method,
        seed,
        f", fold={fold_idx}/{n_folds}" if n_folds else "",
    )

    # ── Pre-compute valid patch centres for this split only ──
    self._centers: list[tuple[float, float, pd.Interval]] = []
    # Per-strip diagnostics, populated during centre building.  Lets
    # downstream visualisation / evaluation code compare the two modes.
    self._per_strip_stats: list[dict[str, Any]] = []

    if self.center_mode == "simple":
        self._build_valid_centers_simple()
    else:
        self._build_valid_centers_optimal()

    n = len(self._centers)
    self.length = length if length is not None else max(1, n)

    if not self.replacement and self.length > n:
        logger.warning(
            "length (%d) > available centres (%d) with replacement=False; "
            "capping to %d. Use replacement=True for oversampling.",
            self.length, n, n,
        )
        self.length = n

    if not self._centers:
        logger.warning(
            "HiRISEGeoSampler: No valid patch centres found for split='%s'. "
            "Check that strip polygons are larger than the requested patch "
            "size (%.6f° × %.6f°) and that the split has data assigned.",
            self.split,
            size_h,
            size_w,
        )

split_summary property

split_summary: dict[str, str | int | dict[str, int] | Any]

Return the number of stereo pairs and patch centres per split.

per_strip_stats property

per_strip_stats: list[dict[str, Any]]

Per-strip diagnostics collected during centre building.

Useful for comparing "simple" vs "optimal" coverage on a per-strip basis. Each entry is a dict containing at minimum pair_idx and n_centers.