ins-pricing 0.1.6__py3-none-any.whl

ins_pricing v2/modelling/plotting/geo.py

@@ -0,0 +1,362 @@
+from __future__ import annotations
+
+from typing import Optional, Sequence, Tuple
+
+import numpy as np
+import pandas as pd
+import matplotlib.tri as mtri
+
+from .common import EPS, PlotStyle, finalize_figure, plt
+
+try:  # optional map basemap support
+    import contextily as cx
+except Exception:  # pragma: no cover - optional dependency
+    cx = None
+
+
+_MERCATOR_MAX_LAT = 85.05112878
+_MERCATOR_FACTOR = 20037508.34
+
+
+def _require_contextily(func_name: str) -> None:
+    if cx is None:
+        raise RuntimeError(
+            f"{func_name} requires contextily. Install it via 'pip install contextily'."
+        )
+
+
+def _lonlat_to_mercator(lon: np.ndarray, lat: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
+    lon = np.asarray(lon, dtype=float)
+    lat = np.asarray(lat, dtype=float)
+    lat = np.clip(lat, -_MERCATOR_MAX_LAT, _MERCATOR_MAX_LAT)
+    x = lon * _MERCATOR_FACTOR / 180.0
+    y = np.log(np.tan((90.0 + lat) * np.pi / 360.0)) * _MERCATOR_FACTOR / np.pi
+    return x, y
+
+
+def _apply_bounds(ax: plt.Axes, x: np.ndarray, y: np.ndarray, padding: float) -> None:
+    x_min, x_max = float(np.min(x)), float(np.max(x))
+    y_min, y_max = float(np.min(y)), float(np.max(y))
+    pad_x = (x_max - x_min) * padding
+    pad_y = (y_max - y_min) * padding
+    if pad_x == 0:
+        pad_x = 1.0
+    if pad_y == 0:
+        pad_y = 1.0
+    ax.set_xlim(x_min - pad_x, x_max + pad_x)
+    ax.set_ylim(y_min - pad_y, y_max + pad_y)
+
+
+def _resolve_basemap(source):
+    if cx is None or source is None:
+        return source
+    if isinstance(source, str):
+        provider = cx.providers
+        for part in source.split("."):
+            if isinstance(provider, dict):
+                provider = provider[part]
+            else:
+                provider = getattr(provider, part)
+        return provider
+    return source
+
+
+def _sanitize_geo(
+    df: pd.DataFrame,
+    x_col: str,
+    y_col: str,
+    value_col: str,
+    weight_col: Optional[str] = None,
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray]]:
+    x = pd.to_numeric(df[x_col], errors="coerce").to_numpy(dtype=float)
+    y = pd.to_numeric(df[y_col], errors="coerce").to_numpy(dtype=float)
+    z = pd.to_numeric(df[value_col], errors="coerce").to_numpy(dtype=float)
+    w = None
+    if weight_col:
+        w = pd.to_numeric(df[weight_col], errors="coerce").to_numpy(dtype=float)
+
+    if w is None:
+        mask = np.isfinite(x) & np.isfinite(y) & np.isfinite(z)
+    else:
+        mask = np.isfinite(x) & np.isfinite(y) & np.isfinite(z) & np.isfinite(w)
+        w = w[mask]
+    return x[mask], y[mask], z[mask], w
+
+
+def _downsample_points(
+    x: np.ndarray,
+    y: np.ndarray,
+    z: np.ndarray,
+    w: Optional[np.ndarray],
+    max_points: Optional[int],
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray]]:
+    if max_points is None:
+        return x, y, z, w
+    max_points = int(max_points)
+    if max_points <= 0 or len(x) <= max_points:
+        return x, y, z, w
+    rng = np.random.default_rng(13)
+    idx = rng.choice(len(x), size=max_points, replace=False)
+    if w is None:
+        return x[idx], y[idx], z[idx], None
+    return x[idx], y[idx], z[idx], w[idx]
+
+
+def plot_geo_heatmap(
+    df: pd.DataFrame,
+    *,
+    x_col: str,
+    y_col: str,
+    value_col: str,
+    weight_col: Optional[str] = None,
+    bins: int | Tuple[int, int] = 50,
+    agg: str = "mean",
+    cmap: str = "YlOrRd",
+    title: str = "Geo Heatmap",
+    ax: Optional[plt.Axes] = None,
+    show: bool = False,
+    save_path: Optional[str] = None,
+    style: Optional[PlotStyle] = None,
+) -> plt.Figure:
+    style = style or PlotStyle()
+    if agg not in {"mean", "sum"}:
+        raise ValueError("agg must be 'mean' or 'sum'.")
+    x, y, z, w = _sanitize_geo(df, x_col, y_col, value_col, weight_col)
+
+    if isinstance(bins, int):
+        bins = (bins, bins)
+
+    if w is None:
+        sum_z, x_edges, y_edges = np.histogram2d(x, y, bins=bins, weights=z)
+        if agg == "sum":
+            grid = sum_z
+        else:
+            count, _, _ = np.histogram2d(x, y, bins=bins)
+            grid = sum_z / np.maximum(count, 1.0)
+    else:
+        sum_w, x_edges, y_edges = np.histogram2d(x, y, bins=bins, weights=w)
+        sum_zw, _, _ = np.histogram2d(x, y, bins=bins, weights=z * w)
+        grid = sum_zw / np.maximum(sum_w, EPS)
+
+    created_fig = ax is None
+    if created_fig:
+        fig, ax = plt.subplots(figsize=style.figsize)
+    else:
+        fig = ax.figure
+
+    im = ax.imshow(
+        grid.T,
+        origin="lower",
+        extent=[x_edges[0], x_edges[-1], y_edges[0], y_edges[-1]],
+        aspect="auto",
+        cmap=cmap,
+    )
+    cbar = fig.colorbar(im, ax=ax)
+    cbar.set_label(value_col, fontsize=style.label_size)
+    cbar.ax.tick_params(labelsize=style.tick_size)
+
+    ax.set_xlabel(x_col, fontsize=style.label_size)
+    ax.set_ylabel(y_col, fontsize=style.label_size)
+    ax.set_title(title, fontsize=style.title_size)
+    ax.tick_params(axis="both", labelsize=style.tick_size)
+
+    if created_fig:
+        finalize_figure(fig, save_path=save_path, show=show, style=style)
+
+    return fig
+
+
+def plot_geo_contour(
+    df: pd.DataFrame,
+    *,
+    x_col: str,
+    y_col: str,
+    value_col: str,
+    weight_col: Optional[str] = None,
+    max_points: Optional[int] = None,
+    levels: int | Sequence[float] = 10,
+    cmap: str = "viridis",
+    title: str = "Geo Contour",
+    ax: Optional[plt.Axes] = None,
+    show_points: bool = False,
+    show: bool = False,
+    save_path: Optional[str] = None,
+    style: Optional[PlotStyle] = None,
+) -> plt.Figure:
+    style = style or PlotStyle()
+    x, y, z, w = _sanitize_geo(df, x_col, y_col, value_col, weight_col)
+    x, y, z, w = _downsample_points(x, y, z, w, max_points)
+
+    if w is not None:
+        z = z * w
+
+    triang = mtri.Triangulation(x, y)
+
+    created_fig = ax is None
+    if created_fig:
+        fig, ax = plt.subplots(figsize=style.figsize)
+    else:
+        fig = ax.figure
+
+    contour = ax.tricontourf(triang, z, levels=levels, cmap=cmap)
+    if show_points:
+        ax.scatter(x, y, s=6, c="k", alpha=0.2)
+    cbar = fig.colorbar(contour, ax=ax)
+    cbar.set_label(value_col, fontsize=style.label_size)
+    cbar.ax.tick_params(labelsize=style.tick_size)
+
+    ax.set_xlabel(x_col, fontsize=style.label_size)
+    ax.set_ylabel(y_col, fontsize=style.label_size)
+    ax.set_title(title, fontsize=style.title_size)
+    ax.tick_params(axis="both", labelsize=style.tick_size)
+
+    if created_fig:
+        finalize_figure(fig, save_path=save_path, show=show, style=style)
+
+    return fig
+
+
+def plot_geo_heatmap_on_map(
+    df: pd.DataFrame,
+    *,
+    lon_col: str,
+    lat_col: str,
+    value_col: str,
+    weight_col: Optional[str] = None,
+    bins: int | Tuple[int, int] = 100,
+    agg: str = "mean",
+    cmap: str = "YlOrRd",
+    alpha: float = 0.6,
+    basemap: Optional[object] = "CartoDB.Positron",
+    zoom: Optional[int] = None,
+    padding: float = 0.05,
+    title: str = "Geo Heatmap (Map)",
+    ax: Optional[plt.Axes] = None,
+    show_points: bool = False,
+    show: bool = False,
+    save_path: Optional[str] = None,
+    style: Optional[PlotStyle] = None,
+) -> plt.Figure:
+    _require_contextily("plot_geo_heatmap_on_map")
+    style = style or PlotStyle()
+    if agg not in {"mean", "sum"}:
+        raise ValueError("agg must be 'mean' or 'sum'.")
+    lon, lat, z, w = _sanitize_geo(df, lon_col, lat_col, value_col, weight_col)
+    x, y = _lonlat_to_mercator(lon, lat)
+
+    if isinstance(bins, int):
+        bins = (bins, bins)
+
+    if w is None:
+        sum_z, x_edges, y_edges = np.histogram2d(x, y, bins=bins, weights=z)
+        if agg == "sum":
+            grid = sum_z
+        else:
+            count, _, _ = np.histogram2d(x, y, bins=bins)
+            grid = sum_z / np.maximum(count, 1.0)
+    else:
+        sum_w, x_edges, y_edges = np.histogram2d(x, y, bins=bins, weights=w)
+        sum_zw, _, _ = np.histogram2d(x, y, bins=bins, weights=z * w)
+        grid = sum_zw / np.maximum(sum_w, EPS)
+
+    created_fig = ax is None
+    if created_fig:
+        fig, ax = plt.subplots(figsize=style.figsize)
+    else:
+        fig = ax.figure
+
+    _apply_bounds(ax, x, y, padding)
+    ax.set_aspect("equal", adjustable="box")
+
+    source = _resolve_basemap(basemap)
+    if source is not None:
+        if zoom is None:
+            cx.add_basemap(ax, source=source, crs="EPSG:3857")
+        else:
+            cx.add_basemap(ax, source=source, crs="EPSG:3857", zoom=zoom)
+
+    im = ax.imshow(
+        grid.T,
+        origin="lower",
+        extent=[x_edges[0], x_edges[-1], y_edges[0], y_edges[-1]],
+        aspect="auto",
+        cmap=cmap,
+        alpha=alpha,
+    )
+    if show_points:
+        ax.scatter(x, y, s=6, c="k", alpha=0.25)
+
+    cbar = fig.colorbar(im, ax=ax)
+    cbar.set_label(value_col, fontsize=style.label_size)
+    cbar.ax.tick_params(labelsize=style.tick_size)
+
+    ax.set_title(title, fontsize=style.title_size)
+    ax.tick_params(axis="both", labelsize=style.tick_size)
+
+    if created_fig:
+        finalize_figure(fig, save_path=save_path, show=show, style=style)
+
+    return fig
+
+
+def plot_geo_contour_on_map(
+    df: pd.DataFrame,
+    *,
+    lon_col: str,
+    lat_col: str,
+    value_col: str,
+    weight_col: Optional[str] = None,
+    max_points: Optional[int] = None,
+    levels: int | Sequence[float] = 10,
+    cmap: str = "viridis",
+    alpha: float = 0.6,
+    basemap: Optional[object] = "CartoDB.Positron",
+    zoom: Optional[int] = None,
+    padding: float = 0.05,
+    title: str = "Geo Contour (Map)",
+    ax: Optional[plt.Axes] = None,
+    show_points: bool = False,
+    show: bool = False,
+    save_path: Optional[str] = None,
+    style: Optional[PlotStyle] = None,
+) -> plt.Figure:
+    _require_contextily("plot_geo_contour_on_map")
+    style = style or PlotStyle()
+    lon, lat, z, w = _sanitize_geo(df, lon_col, lat_col, value_col, weight_col)
+    lon, lat, z, w = _downsample_points(lon, lat, z, w, max_points)
+    x, y = _lonlat_to_mercator(lon, lat)
+    if w is not None:
+        z = z * w
+
+    created_fig = ax is None
+    if created_fig:
+        fig, ax = plt.subplots(figsize=style.figsize)
+    else:
+        fig = ax.figure
+
+    _apply_bounds(ax, x, y, padding)
+    ax.set_aspect("equal", adjustable="box")
+
+    source = _resolve_basemap(basemap)
+    if source is not None:
+        if zoom is None:
+            cx.add_basemap(ax, source=source, crs="EPSG:3857")
+        else:
+            cx.add_basemap(ax, source=source, crs="EPSG:3857", zoom=zoom)
+
+    triang = mtri.Triangulation(x, y)
+    contour = ax.tricontourf(triang, z, levels=levels, cmap=cmap, alpha=alpha)
+    if show_points:
+        ax.scatter(x, y, s=6, c="k", alpha=0.25)
+
+    cbar = fig.colorbar(contour, ax=ax)
+    cbar.set_label(value_col, fontsize=style.label_size)
+    cbar.ax.tick_params(labelsize=style.tick_size)
+
+    ax.set_title(title, fontsize=style.title_size)
+    ax.tick_params(axis="both", labelsize=style.tick_size)
+
+    if created_fig:
+        finalize_figure(fig, save_path=save_path, show=show, style=style)
+
+    return fig

ins_pricing v2/modelling/plotting/importance.py

@@ -0,0 +1,121 @@
+from __future__ import annotations
+
+from typing import Mapping, Optional, Sequence, Tuple
+
+import numpy as np
+import pandas as pd
+
+from .common import PlotStyle, finalize_figure, plt
+
+
+def _to_series(
+    importance: Mapping[str, float]
+    | Sequence[Tuple[str, float]]
+    | pd.Series
+    | np.ndarray,
+    feature_names: Optional[Sequence[str]] = None,
+) -> pd.Series:
+    if isinstance(importance, pd.Series):
+        return importance.copy()
+    if isinstance(importance, Mapping):
+        return pd.Series(dict(importance))
+    if isinstance(importance, np.ndarray):
+        if feature_names is None:
+            raise ValueError("feature_names is required when importance is an array.")
+        return pd.Series(importance, index=list(feature_names))
+    return pd.Series(dict(importance))
+
+
+def shap_importance(
+    shap_values: np.ndarray,
+    feature_names: Sequence[str],
+) -> pd.Series:
+    if shap_values.ndim == 3:
+        shap_values = shap_values[0]
+    if shap_values.ndim != 2:
+        raise ValueError("shap_values should be 2d (n_samples, n_features).")
+    scores = np.abs(shap_values).mean(axis=0)
+    return pd.Series(scores, index=list(feature_names))
+
+
+def plot_feature_importance(
+    importance: Mapping[str, float]
+    | Sequence[Tuple[str, float]]
+    | pd.Series
+    | np.ndarray,
+    *,
+    feature_names: Optional[Sequence[str]] = None,
+    top_n: int = 30,
+    title: str = "Feature Importance",
+    sort_by: str = "abs",
+    descending: bool = True,
+    show_values: bool = False,
+    ax: Optional[plt.Axes] = None,
+    show: bool = False,
+    save_path: Optional[str] = None,
+    style: Optional[PlotStyle] = None,
+) -> plt.Figure:
+    style = style or PlotStyle()
+    series = _to_series(importance, feature_names=feature_names)
+    series = series.replace([np.inf, -np.inf], np.nan).dropna()
+
+    if sort_by not in {"abs", "value"}:
+        raise ValueError("sort_by must be 'abs' or 'value'.")
+    sort_key = series.abs() if sort_by == "abs" else series
+    series = series.loc[sort_key.sort_values(ascending=not descending).index]
+
+    if top_n > 0:
+        series = series.head(int(top_n))
+
+    created_fig = ax is None
+    if created_fig:
+        height = max(3.0, 0.3 * len(series))
+        fig, ax = plt.subplots(figsize=(style.figsize[0], height))
+    else:
+        fig = ax.figure
+
+    y_pos = np.arange(len(series))
+    ax.barh(y_pos, series.values, color=style.palette[0])
+    ax.set_yticks(y_pos)
+    ax.set_yticklabels(series.index, fontsize=style.tick_size)
+    ax.invert_yaxis()
+    ax.set_title(title, fontsize=style.title_size)
+    ax.tick_params(axis="x", labelsize=style.tick_size)
+    if style.grid:
+        ax.grid(True, axis="x", linestyle=style.grid_style, alpha=style.grid_alpha)
+
+    if show_values:
+        for idx, val in enumerate(series.values):
+            ax.text(val, idx, f" {val:.3f}", va="center", fontsize=style.tick_size)
+
+    if created_fig:
+        finalize_figure(fig, save_path=save_path, show=show, style=style)
+
+    return fig
+
+
+def plot_shap_importance(
+    shap_values: np.ndarray,
+    feature_names: Sequence[str],
+    *,
+    top_n: int = 30,
+    title: str = "SHAP Importance",
+    show_values: bool = False,
+    ax: Optional[plt.Axes] = None,
+    show: bool = False,
+    save_path: Optional[str] = None,
+    style: Optional[PlotStyle] = None,
+) -> plt.Figure:
+    series = shap_importance(shap_values, feature_names)
+    return plot_feature_importance(
+        series,
+        top_n=top_n,
+        title=title,
+        sort_by="abs",
+        descending=True,
+        show_values=show_values,
+        ax=ax,
+        show=show,
+        save_path=save_path,
+        style=style,
+    )

ins_pricing v2/modelling/run_logging.py

@@ -0,0 +1,133 @@
+from __future__ import annotations
+
+import logging
+import os
+import sys
+import tempfile
+from datetime import datetime
+from pathlib import Path
+from typing import Optional, TextIO
+
+_LOGGING_CONFIGURED = False
+_LOG_PATH: Optional[Path] = None
+_LOG_FILE: Optional[TextIO] = None
+
+_TRUTHY = {"1", "true", "yes", "y", "on"}
+
+
+class _TeeStream:
+    def __init__(self, primary: TextIO, secondary: TextIO) -> None:
+        self._primary = primary
+        self._secondary = secondary
+
+    def write(self, data: str) -> int:
+        if not data:
+            return 0
+        try:
+            self._primary.write(data)
+        except Exception:
+            pass
+        try:
+            self._secondary.write(data)
+        except Exception:
+            pass
+        return len(data)
+
+    def flush(self) -> None:
+        for stream in (self._primary, self._secondary):
+            try:
+                stream.flush()
+            except Exception:
+                pass
+
+    def isatty(self) -> bool:
+        return bool(getattr(self._primary, "isatty", lambda: False)())
+
+    def fileno(self) -> int:
+        return self._primary.fileno()
+
+    def __getattr__(self, name: str):
+        return getattr(self._primary, name)
+
+
+def _is_truthy(value: Optional[str]) -> bool:
+    return str(value).strip().lower() in _TRUTHY
+
+
+def _resolve_log_dir(log_dir: Optional[str | Path]) -> Optional[Path]:
+    candidates: list[Path] = []
+    if log_dir:
+        candidates.append(Path(log_dir).expanduser())
+    env_dir = os.environ.get("INS_PRICING_LOG_DIR")
+    if env_dir:
+        candidates.append(Path(env_dir).expanduser())
+    candidates.append(Path.cwd() / "logs")
+    candidates.append(Path.home() / ".ins_pricing" / "logs")
+    candidates.append(Path(tempfile.gettempdir()) / "ins_pricing_logs")
+
+    for cand in candidates:
+        try:
+            cand.mkdir(parents=True, exist_ok=True)
+            return cand
+        except Exception:
+            continue
+    return None
+
+
+def _build_log_filename(prefix: str) -> str:
+    ts = datetime.now().strftime("%Y%m%d_%H%M%S")
+    rank = (
+        os.environ.get("RANK")
+        or os.environ.get("LOCAL_RANK")
+        or os.environ.get("SLURM_PROCID")
+    )
+    suffix = f"r{rank}" if rank is not None else f"pid{os.getpid()}"
+    safe_prefix = "".join(
+        ch if ch.isalnum() or ch in "-_." else "_" for ch in prefix)
+    return f"{safe_prefix}_{ts}_{suffix}.log"
+
+
+def configure_run_logging(
+    *,
+    prefix: str = "ins_pricing",
+    log_dir: Optional[str | Path] = None,
+    level: int = logging.INFO,
+    announce: bool = True,
+) -> Optional[Path]:
+    global _LOGGING_CONFIGURED, _LOG_PATH, _LOG_FILE
+
+    if _LOGGING_CONFIGURED:
+        return _LOG_PATH
+    if _is_truthy(os.environ.get("INS_PRICING_LOG_DISABLE")):
+        return None
+
+    resolved_dir = _resolve_log_dir(log_dir)
+    if resolved_dir is None:
+        return None
+
+    log_path = resolved_dir / _build_log_filename(prefix)
+    try:
+        log_file = log_path.open("a", encoding="utf-8")
+    except Exception:
+        return None
+
+    sys.stdout = _TeeStream(sys.stdout, log_file)  # type: ignore[assignment]
+    sys.stderr = _TeeStream(sys.stderr, log_file)  # type: ignore[assignment]
+    _LOG_FILE = log_file
+    _LOG_PATH = log_path
+    _LOGGING_CONFIGURED = True
+
+    root = logging.getLogger()
+    if not root.handlers:
+        logging.basicConfig(
+            level=level,
+            format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
+            handlers=[logging.StreamHandler(sys.stdout)],
+        )
+    else:
+        root.setLevel(level)
+
+    if announce:
+        print(f"[ins_pricing] log saved to {log_path}", flush=True)
+
+    return log_path

ins_pricing v2/modelling/tests/conftest.py

@@ -0,0 +1,8 @@
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parents[1]
+if str(ROOT) not in sys.path:
+    sys.path.insert(0, str(ROOT))

ins_pricing v2/modelling/tests/test_cross_val_generic.py

@@ -0,0 +1,66 @@
+import types
+
+import numpy as np
+import pandas as pd
+import pytest
+
+pytest.importorskip("torch")
+pytest.importorskip("optuna")
+pytest.importorskip("statsmodels")
+pytest.importorskip("xgboost")
+
+from ins_pricing.bayesopt.trainers import TrainerBase
+
+
+class DummyTrainer(TrainerBase):
+    def __init__(self):
+        ctx = types.SimpleNamespace(prop_test=0.2, rand_seed=123)
+        super().__init__(ctx, "Dummy", "Dummy")
+
+    def train(self) -> None:  # pragma: no cover - not used
+        raise NotImplementedError
+
+
+def test_cross_val_generic_iterates_all_splits():
+    trainer = DummyTrainer()
+
+    X = pd.DataFrame({"x": np.arange(12, dtype=float)})
+    y = pd.Series(np.arange(12, dtype=float))
+    w = pd.Series(np.ones(12, dtype=float))
+
+    def data_provider():
+        return X, y, w
+
+    class DummyModel:
+        def fit(self, X_train, y_train, sample_weight=None):
+            return self
+
+        def predict(self, X_val):
+            return np.zeros(len(X_val))
+
+    def model_builder(_params):
+        return DummyModel()
+
+    calls = []
+
+    def metric_fn(y_true, y_pred, weight):
+        calls.append(len(y_true))
+        return float(np.mean(y_pred))
+
+    splits = [
+        (np.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9]), np.array([10, 11])),
+        (np.array([2, 3, 4, 5, 6, 7, 8, 9, 10, 11]), np.array([0, 1])),
+        (np.array([0, 2, 4, 6, 8, 10]), np.array([1, 3, 5, 7, 9, 11])),
+    ]
+
+    result = trainer.cross_val_generic(
+        trial=object(),
+        hyperparameter_space={"p": lambda _t: 1.0},
+        data_provider=data_provider,
+        model_builder=model_builder,
+        metric_fn=metric_fn,
+        splitter=splits,
+    )
+
+    assert result == 0.0
+    assert len(calls) == len(splits)

ins_pricing v2/modelling/tests/test_distributed_utils.py

@@ -0,0 +1,18 @@
+import pytest
+
+pytest.importorskip("torch")
+
+from ins_pricing.bayesopt.utils import DistributedUtils
+
+
+def test_setup_ddp_without_env(monkeypatch):
+    monkeypatch.delenv("RANK", raising=False)
+    monkeypatch.delenv("WORLD_SIZE", raising=False)
+    monkeypatch.delenv("LOCAL_RANK", raising=False)
+
+    ok, local_rank, rank, world_size = DistributedUtils.setup_ddp()
+
+    assert ok is False
+    assert local_rank == 0
+    assert rank == 0
+    assert world_size == 1