plaknit 0.3.2__tar.gz → 0.3.3__tar.gz

{plaknit-0.3.2/plaknit.egg-info → plaknit-0.3.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: plaknit
-Version: 0.3.2
+Version: 0.3.3
 Summary: Processing Large-Scale PlanetScope Data
 Author-email: Dryver Finch <dryver2206@gmail.com>
 License: MIT License

{plaknit-0.3.2 → plaknit-0.3.3}/plaknit/__init__.py

@@ -10,7 +10,7 @@ from .models.rf import predict_rf, smooth_probs, train_rf
 
 __author__ = """Dryver Finch"""
 __email__ = "dryver2206@gmail.com"
-__version__ = "0.3.2"
+__version__ = "0.3.3"
 
 __all__ = [
     "BRTEnsemble",

{plaknit-0.3.2 → plaknit-0.3.3}/plaknit/cli/ensemble_cli.py

@@ -263,7 +263,18 @@ def main(argv: Optional[Sequence[str]] = None) -> int:
         "--jobs",
         type=int,
         default=1,
-        help="Parallel workers for block prediction (default: 1). Use -1 for all cores.",
+        help=(
+            "Parallel workers for block prediction (default: 1). Use -1 for all "
+            "cores on CPU, or up to the number of visible GPUs with --gpu."
+        ),
+    )
+    predict_parser.add_argument(
+        "--gpu",
+        action="store_true",
+        help=(
+            "Use CUDA for XGBoost prediction when available. With --jobs > 1, "
+            "models are distributed across persistent GPU workers."
+        ),
     )
     predict_parser.add_argument(
         "--block-overlap",
@@ -338,6 +349,7 @@ def main(argv: Optional[Sequence[str]] = None) -> int:
             feature_importance_out=args.feature_importance_out,
             block_overlap=args.block_overlap,
             jobs=args.jobs,
+            gpu=args.gpu,
         )
         return 0
 

{plaknit-0.3.2 → plaknit-0.3.3}/plaknit/data/sampling.py

@@ -189,33 +189,63 @@ def _collect_pseudo_absence_candidate_pool(
     *,
     buffer_meters: float,
 ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
-    """Collect valid background features and coordinates eligible for sampling."""
+    """Collect valid background features and coordinates eligible for sampling.
+
+    Raster data is read in bounded windows so large VRTs do not require a full
+    raster-sized float32 allocation during candidate collection.
+    """
 
     if buffer_meters < 0:
         raise ValueError("training buffer must be non-negative.")
 
-    full_window = windows.Window(0, 0, stack.width, stack.height)
-    data = stack.read(window=full_window, out_dtype="float32")
-    samples = data.reshape(stack.count, -1).T
-    valid = ~_nodata_pixel_mask(samples, stack.nodata_values)
-
     geometries = _buffer_training_geometries(stack, gdf, buffer_meters)
-    if geometries:
-        excluded = features.rasterize(
-            [(geom, 1) for geom in geometries],
-            out_shape=(stack.height, stack.width),
-            transform=stack.transform,
-            fill=0,
-            dtype="uint8",
-        ).reshape(-1)
-        valid &= excluded == 0
-
-    candidate_indices = np.flatnonzero(valid)
-    rows, cols = np.unravel_index(candidate_indices, (stack.height, stack.width))
+    feature_chunks: List[np.ndarray] = []
+    row_chunks: List[np.ndarray] = []
+    col_chunks: List[np.ndarray] = []
+    block_size = 512
+    for row_off in range(0, stack.height, block_size):
+        for col_off in range(0, stack.width, block_size):
+            win = windows.Window(
+                col_off=col_off,
+                row_off=row_off,
+                width=min(block_size, stack.width - col_off),
+                height=min(block_size, stack.height - row_off),
+            )
+            data = stack.read(window=win, out_dtype="float32")
+            samples = data.reshape(stack.count, -1).T
+            valid = ~_nodata_pixel_mask(samples, stack.nodata_values)
+
+            if geometries:
+                block_transform = windows.transform(win, stack.transform)
+                excluded = features.rasterize(
+                    [(geom, 1) for geom in geometries],
+                    out_shape=(int(win.height), int(win.width)),
+                    transform=block_transform,
+                    fill=0,
+                    dtype="uint8",
+                ).reshape(-1)
+                valid &= excluded == 0
+
+            candidate_indices = np.flatnonzero(valid)
+            if candidate_indices.size == 0:
+                continue
+            rows, cols = np.unravel_index(
+                candidate_indices, (int(win.height), int(win.width))
+            )
+            feature_chunks.append(samples[candidate_indices])
+            row_chunks.append((rows + int(win.row_off)).astype("int32", copy=False))
+            col_chunks.append((cols + int(win.col_off)).astype("int32", copy=False))
+
+    if not feature_chunks:
+        return (
+            np.empty((0, stack.count), dtype="float32"),
+            np.empty(0, dtype="int32"),
+            np.empty(0, dtype="int32"),
+        )
     return (
-        samples[candidate_indices],
-        rows.astype("int32", copy=False),
-        cols.astype("int32", copy=False),
+        np.vstack(feature_chunks),
+        np.concatenate(row_chunks),
+        np.concatenate(col_chunks),
     )
 
 

{plaknit-0.3.2 → plaknit-0.3.3}/plaknit/models/brt.py

@@ -151,12 +151,38 @@ def _prepare_brt_training_data(
 
 def _detect_gpu_available() -> bool:
     """Detect if GPU is available for XGBoost."""
+    return bool(_available_cuda_devices())
+
+
+def _available_cuda_devices() -> List[int]:
+    """Return logical CUDA device IDs visible to the current process."""
     try:
-        import cupy as cp  # noqa: F401
+        import cupy as cp
+
+        count = int(cp.cuda.runtime.getDeviceCount())
+    except (ImportError, RuntimeError, AttributeError):
+        return []
+    return list(range(max(0, count)))
+
 
-        return True
-    except ImportError:
-        return False
+def _configure_prediction_device(
+    model: xgb.XGBClassifier, gpu: bool, device_id: Optional[int] = None
+) -> None:
+    """Configure the loaded XGBoost model for CPU or CUDA inference."""
+    if not gpu:
+        return
+    if not _detect_gpu_available():
+        raise RuntimeError(
+            "GPU prediction requested, but CuPy/CUDA is unavailable. "
+            "Install the GPU dependencies or omit --gpu."
+        )
+    device = "cuda" if device_id is None else f"cuda:{device_id}"
+    try:
+        model.set_params(device=device, tree_method="hist", n_jobs=1)
+    except (AttributeError, TypeError, ValueError):
+        model.set_params(
+            predictor="gpu_predictor", tree_method="gpu_hist", n_jobs=1
+        )
 
 
 def train_brt(
@@ -658,6 +684,7 @@ def _init_predict_worker(
     out_dtype: str,
     nodata_value: Union[int, float],
     return_probs: bool,
+    gpu: bool,
 ) -> None:
     global _PREDICT_STACK
     global _PREDICT_MODEL
@@ -669,6 +696,7 @@ def _init_predict_worker(
     stack.__enter__()
     _PREDICT_STACK = stack
     _PREDICT_MODEL = joblib.load(model_path)
+    _configure_prediction_device(_PREDICT_MODEL, gpu)
     _PREDICT_OUT_DTYPE = out_dtype
     _PREDICT_NODATA = nodata_value
     _PREDICT_RETURN_PROBS = return_probs
@@ -706,6 +734,7 @@ def predict_brt(
     binary_out: Optional[PathLike] = None,
     block_overlap: int = 0,
     jobs: int = 1,
+    gpu: bool = False,
 ) -> Path:
     """Apply a trained BRT (XGBoost) model to a raster stack and write a classified GeoTIFF.
 
@@ -724,6 +753,7 @@ def predict_brt(
 
     _log("[bold cyan]Loading model...")
     model: xgb.XGBClassifier = joblib.load(model_path)
+    _configure_prediction_device(model, gpu)
     classes = getattr(model, "classes_", None)
     classes_dtype = getattr(classes, "dtype", np.int32)
     if np.issubdtype(classes_dtype, np.integer):
@@ -828,6 +858,7 @@ def predict_brt(
                         out_dtype,
                         nodata_value,
                         return_probs,
+                        gpu,
                     ),
                 ) as executor:
                     futures: Dict[concurrent.futures.Future, windows.Window] = {}

{plaknit-0.3.2 → plaknit-0.3.3}/plaknit/models/ensemble.py

@@ -7,7 +7,9 @@ import concurrent.futures
 import contextlib
 import csv
 import json
+import multiprocessing
 import os
+import queue
 import tempfile
 import warnings
 from pathlib import Path
@@ -35,6 +37,20 @@ from ..processing.evaluation import (
 
 PathLike = Union[str, Path]
 WindowTuple = Tuple[int, int, int, int]
+_WorkerResult = Tuple[int, List[int], np.ndarray]
+
+
+def _assign_models_to_workers(n_models: int, n_workers: int) -> List[List[int]]:
+    """Assign every model exactly once across workers in round-robin order."""
+    if n_models < 1:
+        raise ValueError("n_models must be at least 1.")
+    if n_workers < 1:
+        raise ValueError("n_workers must be at least 1.")
+    worker_count = min(n_models, n_workers)
+    return [
+        [model_idx for model_idx in range(n_models) if model_idx % worker_count == worker_id]
+        for worker_id in range(worker_count)
+    ]
 
 
 def _write_model_summary_csv(
@@ -131,6 +147,23 @@ def _write_feature_importance_csv(path: Path, models: Sequence[Any]) -> None:
             writer.writerow(row)
 
 
+def _aggregate_ensemble_probabilities(
+    member_probs: np.ndarray,
+    ci_t_crit: Optional[float],
+) -> Tuple[np.ndarray, Optional[np.ndarray], Optional[np.ndarray]]:
+    """Aggregate complete ensemble probabilities for one output window."""
+    mean = member_probs.mean(axis=0)
+    if ci_t_crit is None or member_probs.shape[0] < 2:
+        return mean, None, None
+    std = member_probs.std(axis=0, ddof=1)
+    half_width = ci_t_crit * std / np.sqrt(member_probs.shape[0])
+    return (
+        mean,
+        np.clip(mean - half_width, 0.0, 1.0),
+        np.clip(mean + half_width, 0.0, 1.0),
+    )
+
+
 def _predict_ensemble_block(
     stack: _RasterStack,
     models: List[Any],
@@ -192,7 +225,9 @@ def _predict_ensemble_block(
         ],
         axis=0,
     )
-    mean = probs_stack.mean(axis=0)
+    mean, lower, upper = _aggregate_ensemble_probabilities(
+        probs_stack, ci_t_crit
+    )
 
     def _scatter(values: np.ndarray) -> np.ndarray:
         full = np.full((samples.shape[0], num_classes), np.nan, dtype="float32")
@@ -203,12 +238,8 @@ def _predict_ensemble_block(
 
     mean_cube = _scatter(mean)
 
-    lower_cube = upper_cube = None
-    if n_models > 1 and ci_t_crit is not None:
-        std = probs_stack.std(axis=0, ddof=1)
-        half_width = ci_t_crit * std / np.sqrt(n_models)
-        lower_cube = _scatter(np.clip(mean - half_width, 0.0, 1.0))
-        upper_cube = _scatter(np.clip(mean + half_width, 0.0, 1.0))
+    lower_cube = _scatter(lower) if lower is not None else None
+    upper_cube = _scatter(upper) if upper is not None else None
 
     member_cube = np.full(
         (n_models, samples.shape[0], num_classes), np.nan, dtype="float32"
@@ -251,6 +282,7 @@ def _init_ensemble_predict_worker(
     model_paths: List[str],
     classes: List[Any],
     ci_t_crit: Optional[float],
+    gpu: bool,
 ) -> None:
     global _ENSEMBLE_PREDICT_STACK
     global _ENSEMBLE_PREDICT_MODELS
@@ -261,6 +293,8 @@ def _init_ensemble_predict_worker(
     stack.__enter__()
     _ENSEMBLE_PREDICT_STACK = stack
     _ENSEMBLE_PREDICT_MODELS = [joblib.load(path) for path in model_paths]
+    for model in _ENSEMBLE_PREDICT_MODELS:
+        brt._configure_prediction_device(model, gpu)
     _ENSEMBLE_PREDICT_CLASSES = np.asarray(classes)
     _ENSEMBLE_PREDICT_CI_T_CRIT = ci_t_crit
     atexit.register(_close_ensemble_predict_worker)
@@ -285,6 +319,43 @@ def _predict_ensemble_block_worker(
     )
 
 
+def _ensemble_gpu_worker(
+    image_paths: List[str],
+    band_indices: Optional[Sequence[int]],
+    model_paths: List[str],
+    model_indices: List[int],
+    classes: List[Any],
+    device_id: int,
+    task_queue: Any,
+    result_queue: Any,
+) -> None:
+    """Run assigned ensemble models on one persistent CUDA worker."""
+    try:
+        with _open_raster_stack(image_paths, band_indices=band_indices) as stack:
+            models = [joblib.load(model_paths[idx]) for idx in model_indices]
+            for model in models:
+                brt._configure_prediction_device(model, True, device_id=device_id)
+
+            while True:
+                task = task_queue.get()
+                if task is None:
+                    return
+                window_id, win_tuple, block_overlap = task
+                member_probs, _, _, _ = _predict_ensemble_block(
+                    stack,
+                    models,
+                    brt._tuple_to_window(win_tuple),
+                    classes=np.asarray(classes),
+                    block_overlap=block_overlap,
+                    ci_t_crit=None,
+                )
+                result_queue.put(
+                    ("result", window_id, model_indices, member_probs)
+                )
+    except BaseException as exc:
+        result_queue.put(("error", repr(exc)))
+
+
 class BRTEnsemble:
     """Ensemble of independent BRT models with probability summaries.
 
@@ -508,6 +579,7 @@ class BRTEnsemble:
         feature_importance_out: Optional[PathLike] = None,
         block_overlap: int = 0,
         jobs: int = 1,
+        gpu: bool = False,
     ) -> Path:
         """Write unweighted ensemble probability summaries for a raster stack.
 
@@ -534,7 +606,11 @@ class BRTEnsemble:
         block_overlap
             Number of pixels to overlap blocks for smoothing edge artifacts.
         jobs
-            Number of parallel worker processes.
+            Number of parallel worker processes. In GPU mode, this is capped by
+            the number of visible CUDA devices and ensemble models.
+        gpu
+            If True, configure each XGBoost model for CUDA prediction. Worker
+            processes run concurrently when ``jobs`` is greater than one.
 
         Returns
         -------
@@ -566,6 +642,14 @@ class BRTEnsemble:
             )
 
         models = [joblib.load(path) for path in model_paths]
+        cuda_devices = brt._available_cuda_devices() if gpu else []
+        if gpu and not cuda_devices:
+            raise RuntimeError(
+                "GPU prediction requested, but no visible CUDA devices were found."
+            )
+        if gpu and jobs <= 1:
+            for model in models:
+                brt._configure_prediction_device(model, True, device_id=0)
         self.models_ = models
 
         # Verify class consistency
@@ -676,7 +760,104 @@ class BRTEnsemble:
                         if upper_dst is not None and upper_probs is not None:
                             upper_dst.write(upper_probs, window=win)
 
-                    if jobs > 1:
+                    if gpu and jobs > 1:
+                        worker_count = min(jobs, len(models), len(cuda_devices))
+                        assignments = _assign_models_to_workers(
+                            len(models), worker_count
+                        )
+                        context = multiprocessing.get_context("spawn")
+                        task_queues = [context.Queue(maxsize=2) for _ in assignments]
+                        result_queue = context.Queue()
+                        workers = [
+                            context.Process(
+                                target=_ensemble_gpu_worker,
+                                args=(
+                                    [str(path) for path in stack.paths],
+                                    band_indices_to_use,
+                                    [str(path) for path in model_paths],
+                                    assignment,
+                                    classes_arr.tolist(),
+                                    cuda_devices[worker_id],
+                                    task_queue,
+                                    result_queue,
+                                ),
+                            )
+                            for worker_id, (assignment, task_queue) in enumerate(
+                                zip(assignments, task_queues)
+                            )
+                        ]
+                        for worker in workers:
+                            worker.start()
+                        try:
+                            for window_id, win in enumerate(window_iter):
+                                task = (
+                                    window_id,
+                                    brt._window_to_tuple(win),
+                                    block_overlap,
+                                )
+                                for task_queue in task_queues:
+                                    task_queue.put(task)
+
+                                partial_results: Dict[int, np.ndarray] = {}
+                                for _ in workers:
+                                    try:
+                                        message = result_queue.get(timeout=300)
+                                    except queue.Empty as exc:
+                                        dead = [
+                                            worker.pid
+                                            for worker in workers
+                                            if not worker.is_alive()
+                                        ]
+                                        raise RuntimeError(
+                                            "GPU prediction worker stopped before "
+                                            f"completing window {window_id}; pids={dead}."
+                                        ) from exc
+                                    if message[0] == "error":
+                                        raise RuntimeError(
+                                            f"GPU prediction worker failed: {message[1]}"
+                                        )
+                                    _, returned_window_id, indices, member_probs = message
+                                    if returned_window_id != window_id:
+                                        raise RuntimeError(
+                                            "GPU prediction returned an unexpected window."
+                                        )
+                                    for local_idx, model_idx in enumerate(indices):
+                                        if model_idx in partial_results:
+                                            raise RuntimeError(
+                                                f"Duplicate GPU prediction for model {model_idx}."
+                                            )
+                                        partial_results[model_idx] = member_probs[local_idx]
+
+                                if len(partial_results) != len(models):
+                                    raise RuntimeError(
+                                        "GPU prediction returned an incomplete ensemble "
+                                        f"for window {window_id}."
+                                    )
+                                full_member_probs = np.stack(
+                                    [partial_results[idx] for idx in range(len(models))],
+                                    axis=0,
+                                )
+                                mean_probs, lower_probs, upper_probs = (
+                                    _aggregate_ensemble_probabilities(
+                                        full_member_probs, ci_t_crit
+                                    )
+                                )
+                                _write_block_result(
+                                    win,
+                                    full_member_probs,
+                                    mean_probs,
+                                    lower_probs,
+                                    upper_probs,
+                                )
+                        finally:
+                            for task_queue in task_queues:
+                                task_queue.put(None)
+                            for worker in workers:
+                                worker.join(timeout=10)
+                                if worker.is_alive():
+                                    worker.terminate()
+                                    worker.join()
+                    elif jobs > 1:
                         max_workers = jobs
                         max_pending = max_workers * 2
                         image_paths = [str(path) for path in stack.paths]
@@ -690,6 +871,7 @@ class BRTEnsemble:
                                 model_path_strs,
                                 classes_arr.tolist(),
                                 ci_t_crit,
+                                gpu,
                             ),
                         ) as executor:
                             futures: Dict[concurrent.futures.Future, windows.Window] = (

{plaknit-0.3.2 → plaknit-0.3.3/plaknit.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: plaknit
-Version: 0.3.2
+Version: 0.3.3
 Summary: Processing Large-Scale PlanetScope Data
 Author-email: Dryver Finch <dryver2206@gmail.com>
 License: MIT License

{plaknit-0.3.2 → plaknit-0.3.3}/plaknit.egg-info/SOURCES.txt

@@ -31,6 +31,7 @@ plaknit/processing/__init__.py
 plaknit/processing/evaluation.py
 tests/test_classify.py
 tests/test_ensemble_cli.py
+tests/test_ensemble_prediction.py
 tests/test_geometry.py
 tests/test_mosaic.py
 tests/test_orders.py

{plaknit-0.3.2 → plaknit-0.3.3}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "plaknit"
-version = "0.3.2"
+version = "0.3.3"
 description = "Processing Large-Scale PlanetScope Data"
 readme = "README.md"
 requires-python = ">=3.8"
@@ -43,7 +43,7 @@ exclude = ["docs*"]
 dependencies = { file = ["requirements.txt"] }
 
 [tool.bumpversion]
-current_version = "0.3.2"
+current_version = "0.3.3"
 commit = true
 tag = true
 tag_name = "v{new_version}"

plaknit-0.3.3/tests/test_ensemble_prediction.py

@@ -0,0 +1,47 @@
+import numpy as np
+import pytest
+
+from plaknit.models.ensemble import (
+    _aggregate_ensemble_probabilities,
+    _assign_models_to_workers,
+)
+
+
+def test_assign_models_to_workers_round_robin() -> None:
+    assert _assign_models_to_workers(5, 3) == [[0, 3], [1, 4], [2]]
+
+
+def test_assign_models_to_workers_caps_workers_at_model_count() -> None:
+    assert _assign_models_to_workers(2, 8) == [[0], [1]]
+
+
+def test_assign_models_to_workers_rejects_empty_inputs() -> None:
+    with pytest.raises(ValueError):
+        _assign_models_to_workers(0, 1)
+    with pytest.raises(ValueError):
+        _assign_models_to_workers(1, 0)
+
+
+def test_aggregate_ensemble_probabilities_matches_sample_ci() -> None:
+    member_probs = np.asarray(
+        [
+            [[[0.1, 0.9]]],
+            [[[0.2, 0.8]]],
+            [[[0.3, 0.7]]],
+        ],
+        dtype="float32",
+    )
+    mean, lower, upper = _aggregate_ensemble_probabilities(member_probs, 2.0)
+    expected_mean = member_probs.mean(axis=0)
+    expected_half_width = 2.0 * member_probs.std(axis=0, ddof=1) / np.sqrt(3)
+    np.testing.assert_allclose(mean, expected_mean)
+    np.testing.assert_allclose(lower, np.clip(expected_mean - expected_half_width, 0, 1))
+    np.testing.assert_allclose(upper, np.clip(expected_mean + expected_half_width, 0, 1))
+
+
+def test_aggregate_single_model_has_no_interval() -> None:
+    member_probs = np.asarray([[[[0.4]]]], dtype="float32")
+    mean, lower, upper = _aggregate_ensemble_probabilities(member_probs, None)
+    np.testing.assert_array_equal(mean, member_probs[0])
+    assert lower is None
+    assert upper is None