Release: auto-parallel ensembles + remaining parallelization

climate_tookit/calculate_hazards/ensemble_hazards.py

@@ -361,13 +361,14 @@ def calculate_ensemble(crop: str, lat: float, lon: float,
                        fixed_season: Optional[str] = None,
                        soilcp: float = DEFAULT_SOILCP,
                        soilsat: float = DEFAULT_SOILSAT,
-                       max_workers: int = 8) -> Dict:
+                       max_workers: int = 0) -> Dict:
     mode = 'fixed_season' if fixed_season else 'auto_detect'
     fixed_w = (_expand_windows(start_year, end_year, _parse_fixed(fixed_season))
                if fixed_season else None)
 
     jobs = [(m, sc) for sc in scenarios for m in models]
-    workers = max(1, min(max_workers, len(jobs)))
+    # max_workers <= 0 -> auto: one worker per job, capped at 16.
+    workers = max(1, min(len(jobs), 16) if max_workers <= 0 else min(max_workers, len(jobs)))
 
     print(f"\nNEX-GDDP ensemble: {crop} at ({lat:.4f}, {lon:.4f})  "
           f"{start_year}-{end_year}")
@@ -842,9 +843,9 @@ def print_baseline_future_ltm(cmp: Dict, baseline_label: str,
     p.add_argument('--soilsat', type=float, default=DEFAULT_SOILSAT,
                    help=f'Extra soil water from field capacity to saturation, mm '
                         f'(water-balance NDWL0; default: {DEFAULT_SOILSAT})')
-    p.add_argument('--workers',      type=int, default=8,
+    p.add_argument('--workers',      type=int, default=0,
                    help='parallel GEE fetch workers across model x scenario '
-                        '(default: 8; use 1 to disable parallelism)')
+                        '(default: auto = one per job, capped at 16; use 1 to disable)')
     p.add_argument('--baseline-source', type=str, default=None,
                    choices=['era_5', 'agera_5', 'chirps+chirts'],
                    help='enable Baseline LTM vs Future LTM comparison using this '

climate_tookit/climate_statistics/ensemble_statistics.py

@@ -253,6 +253,7 @@ def analyze_ensemble_nex_gddp(
     exclude_models: Optional[List[str]] = None,
     extra_months:   int = 6,
     verbose:        bool = True,
+    max_workers:    int = 0,
 ) -> Dict[str, Any]:
     """
     Future LTM via NEX-GDDP CMIP6 ensemble.
@@ -284,44 +285,71 @@ def analyze_ensemble_nex_gddp(
         print(f"{'=' * 60}")
 
     # Per-model collection. The LTM list is keyed by model name so callers can verify the ensemble = mean of per-model values.
-    per_model_seasons: List[List[Dict[str, Any]]] = []
-    per_model_ltm_list: List[Dict[str, Any]]     = []  # ordered, used for step 2 below
-    per_model_ltm_by_name: Dict[str, Dict[str, Any]] = {}
-    per_model_annual:  List[Dict[str, Dict]]     = []
-    models_ok: List[str]            = []
-    failed:    List[Dict[str, str]] = []
-    for i, model in enumerate(active, 1):
-        if verbose:
-            print(f"\n  [{i:02d}/{len(active):02d}] {model}", flush=True)
+    # Each model's analyze_climate_statistics() is an independent, I/O-bound GEE
+    # run, so the models are fetched concurrently. Order is restored afterwards.
+    def _run_model(model: str):
+        """Run one model's statistics. Returns (model, result_dict, error)."""
         try:
-            # STEP 1 of the FUTURE LTM pipeline (per-model):
             r = analyze_climate_statistics(
                 location_coord=location_coord,
                 start_year=start_year, end_year=end_year,
                 source='nex_gddp', fixed_season=fixed_season,
                 model=model, scenario=scenario,
                 extra_months=extra_months,
             )
-            seasons    = r.get('season_statistics') or []
-            ltm_model  = r.get('ltm_season_summary') or {}  
-            annual     = r.get('annual_summary') or {}
+            seasons   = r.get('season_statistics') or []
+            ltm_model = r.get('ltm_season_summary') or {}
             if not seasons and not ltm_model.get('windows'):
-                failed.append({'model': model,
-                               'error': 'no seasons or LTM produced'})
-                if verbose:
-                    print("    x  no seasons or LTM produced")
-                continue
-            per_model_seasons.append(seasons)
-            per_model_ltm_list.append(ltm_model)
-            per_model_ltm_by_name[model] = ltm_model
-            per_model_annual.append(annual)
-            models_ok.append(model)
-            if verbose:
-                print("    ok")
+                return model, None, 'no seasons or LTM produced'
+            return model, r, None
         except Exception as exc:
-            failed.append({'model': model, 'error': str(exc)})
+            return model, None, str(exc)
+
+    raw_by_model: Dict[str, Dict[str, Any]] = {}
+    failed:    List[Dict[str, str]] = []
+    # max_workers <= 0 -> auto: one worker per model, capped at 16.
+    workers = max(1, min(len(active), 16) if max_workers <= 0 else min(max_workers, len(active)))
+
+    def _record(model, r, err, idx):
+        if err is not None or r is None:
+            failed.append({'model': model, 'error': err or 'unknown error'})
             if verbose:
-                print(f"    x  {exc}")
+                print(f"  [{idx:02d}/{len(active):02d}] {model}  x  {err}", flush=True)
+        else:
+            raw_by_model[model] = r
+            if verbose:
+                print(f"  [{idx:02d}/{len(active):02d}] {model}  ok", flush=True)
+
+    if workers == 1:
+        for i, model in enumerate(active, 1):
+            m, r, err = _run_model(model)
+            _record(m, r, err, i)
+    else:
+        from concurrent.futures import ThreadPoolExecutor, as_completed
+        order = {m: i for i, m in enumerate(active, 1)}
+        with ThreadPoolExecutor(max_workers=workers) as ex:
+            futs = {ex.submit(_run_model, m): m for m in active}
+            for fut in as_completed(futs):
+                m, r, err = fut.result()
+                _record(m, r, err, order[m])
+
+    # Restore deterministic (input) model order, then unpack into the
+    # downstream per-model lists exactly as the serial path produced them.
+    per_model_seasons: List[List[Dict[str, Any]]] = []
+    per_model_ltm_list: List[Dict[str, Any]]     = []
+    per_model_ltm_by_name: Dict[str, Dict[str, Any]] = {}
+    per_model_annual:  List[Dict[str, Dict]]     = []
+    models_ok: List[str]            = []
+    for model in active:
+        r = raw_by_model.get(model)
+        if r is None:
+            continue
+        per_model_seasons.append(r.get('season_statistics') or [])
+        ltm_model = r.get('ltm_season_summary') or {}
+        per_model_ltm_list.append(ltm_model)
+        per_model_ltm_by_name[model] = ltm_model
+        per_model_annual.append(r.get('annual_summary') or {})
+        models_ok.append(model)
 
     if not models_ok:
         return {'error': 'All models failed.', 'failed_models': failed}
@@ -537,6 +565,14 @@ def print_report(result: Dict[str, Any]) -> None:
 
 # CLI
 def main() -> None:
+    # Ensure Unicode output (arrows, box-drawing, degree sign) works on
+    # Windows consoles that default to cp1252.
+    try:
+        sys.stdout.reconfigure(encoding='utf-8')
+        sys.stderr.reconfigure(encoding='utf-8')
+    except (AttributeError, ValueError):
+        pass
+
     if "--list-models" in sys.argv:
         print("Available NEX-GDDP CMIP6 models:")
         for i, m in enumerate(NEX_GDDP_MODELS, 1):
@@ -592,6 +628,9 @@ def main() -> None:
                    help='Skip saving the JSON output')
     p.add_argument("--quiet",      action='store_true',
                    help='Suppress per-model progress prints')
+    p.add_argument("--workers",    type=int, default=0,
+                   help='Parallel GEE fetch workers across models '
+                        '(default: auto = one per model, capped at 16; use 1 to disable)')
     args = p.parse_args()
 
     try:
@@ -637,6 +676,7 @@ def main() -> None:
             exclude_models=excl,
             extra_months=args.extra_months,
             verbose=not args.quiet,
+            max_workers=args.workers,
         )
         all_results[scenario] = result
 

climate_tookit/climatology/long_term_climatology.py

@@ -901,7 +901,7 @@ def calculate_climatology_ensemble(
     exclude_models: Optional[List[str]] = None,
     output_dir: Optional[str] = None,
     verbose: bool = True,
-    max_workers: int = 8,
+    max_workers: int = 0,
 ) -> Dict[str, Any]:
     """
     NEX-GDDP-CMIP6 ensemble climatology.
@@ -945,7 +945,8 @@ def _run_model(model: str) -> Tuple[str, Optional[Dict[str, Any]], Optional[str]
 
     per_model_results: Dict[str, Dict[str, Any]] = {}
     failed: List[Dict[str, str]] = []
-    workers = max(1, min(max_workers, len(active)))
+    # max_workers <= 0 -> auto: one worker per model, capped at 16.
+    workers = max(1, min(len(active), 16) if max_workers <= 0 else min(max_workers, len(active)))
 
     def _record(model: str, r: Optional[Dict[str, Any]], err: Optional[str], idx: int) -> None:
         if err is not None or r is None:
@@ -1266,9 +1267,9 @@ def main():
                             '(default: all 16).')
     parser.add_argument('--exclude-models', type=str, default=None,
                        help='NEX-GDDP only. Comma-separated CMIP6 models to drop.')
-    parser.add_argument('--workers', type=int, default=8,
+    parser.add_argument('--workers', type=int, default=0,
                        help='NEX-GDDP only. Parallel fetch workers across models '
-                            '(default: 8; use 1 to disable parallelism).')
+                            '(default: auto = one per model, capped at 16; use 1 to disable).')
     parser.add_argument('--format', choices=['text', 'json'], default='text',
                        help='Output format (default: text)')
     parser.add_argument('--output', type=str,

climate_tookit/compare_periods/ensemble_periods.py

@@ -605,6 +605,7 @@ def ensemble_compare(
     exclude_models: Optional[List[str]] = None,
     focal_summary: Optional[Dict[str, Any]] = None,
     verbose:        bool = True,
+    max_workers:    int = 0,
 ) -> Dict[str, Any]:
     """
     Run the future-period-vs-baseline-period comparison once per NEX-GDDP model, then average across models.
@@ -637,12 +638,10 @@ def ensemble_compare(
         print(f"  Models    : {len(active)}")
         print(f"{'=' * 60}")
 
-    per_model: List[Dict[str, Any]] = []
-    failed:    List[Dict[str, str]] = []
-
-    for i, model in enumerate(active, 1):
-        if verbose:
-            print(f"\n  [{i:02d}/{len(active):02d}] {model}", flush=True)
+    # Each model runs an independent baseline + future comparison (two
+    # I/O-bound GEE span-fetches). Run models concurrently; restore order after.
+    def _run_model(model: str):
+        """Compare one model. Returns (model, result_dict, error)."""
         try:
             r = _compare_one_model(
                 location=location,
@@ -655,15 +654,42 @@ def ensemble_compare(
                 scenario=scenario,
             )
             if "error" in r:
-                if verbose: print(f"    x  {r['error']}")
-                failed.append({"model": model, "error": r["error"]})
-                continue
-            r["_model"] = model
-            per_model.append(r)
-            if verbose: print("    ok")
+                return model, None, r["error"]
+            return model, r, None
         except Exception as exc:
-            if verbose: print(f"    x  {exc}")
-            failed.append({"model": model, "error": str(exc)})
+            return model, None, str(exc)
+
+    by_model: Dict[str, Dict[str, Any]] = {}
+    failed:   List[Dict[str, str]] = []
+    # max_workers <= 0 -> auto: one worker per model, capped at 16.
+    workers = max(1, min(len(active), 16) if max_workers <= 0 else min(max_workers, len(active)))
+
+    def _record(model, r, err, idx):
+        if err is not None or r is None:
+            failed.append({"model": model, "error": err or "unknown error"})
+            if verbose:
+                print(f"  [{idx:02d}/{len(active):02d}] {model}  x  {err}", flush=True)
+        else:
+            r["_model"] = model
+            by_model[model] = r
+            if verbose:
+                print(f"  [{idx:02d}/{len(active):02d}] {model}  ok", flush=True)
+
+    if workers == 1:
+        for i, model in enumerate(active, 1):
+            m, r, err = _run_model(model)
+            _record(m, r, err, i)
+    else:
+        from concurrent.futures import ThreadPoolExecutor, as_completed
+        order = {m: i for i, m in enumerate(active, 1)}
+        with ThreadPoolExecutor(max_workers=workers) as ex:
+            futs = {ex.submit(_run_model, m): m for m in active}
+            for fut in as_completed(futs):
+                m, r, err = fut.result()
+                _record(m, r, err, order[m])
+
+    # Deterministic input-model order regardless of completion order.
+    per_model: List[Dict[str, Any]] = [by_model[m] for m in active if m in by_model]
 
     if not per_model:
         return {"error": "All models failed.", "failed_models": failed}
@@ -897,6 +923,13 @@ def print_report(r: Dict[str, Any]) -> None:
 
 # CLI
 def main() -> None:
+    # Ensure Unicode output works on Windows consoles that default to cp1252.
+    try:
+        sys.stdout.reconfigure(encoding='utf-8')
+        sys.stderr.reconfigure(encoding='utf-8')
+    except (AttributeError, ValueError):
+        pass
+
     if "--list-models" in sys.argv:
         print("Available NEX-GDDP-CMIP6 models:")
         for i, m in enumerate(NEX_GDDP_MODELS, 1):
@@ -945,6 +978,9 @@ def main() -> None:
                         f"{', '.join(sorted(SUPPORTED))}")
     p.add_argument("--output",         default=None, help="Save JSON result to this path")
     p.add_argument("--quiet",          action="store_true")
+    p.add_argument("--workers",        type=int, default=0,
+                   help="Parallel GEE fetch workers across models "
+                        "(default: auto = one per model, capped at 16; use 1 to disable)")
     args = p.parse_args()
 
     try:
@@ -1022,6 +1058,7 @@ def main() -> None:
             exclude_models=exclude,
             focal_summary=scenario_focal,
             verbose=not args.quiet,
+            max_workers=args.workers,
         )
         all_results[scenario] = result
         print_report(result)