diff --git a/climate_tookit/calculate_hazards/ensemble_hazards.py b/climate_tookit/calculate_hazards/ensemble_hazards.py index 1611a62..ef2bed2 100644 --- a/climate_tookit/calculate_hazards/ensemble_hazards.py +++ b/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 ' diff --git a/climate_tookit/climate_statistics/ensemble_statistics.py b/climate_tookit/climate_statistics/ensemble_statistics.py index 45c6202..d90637f 100644 --- a/climate_tookit/climate_statistics/ensemble_statistics.py +++ b/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,17 +285,11 @@ 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, @@ -302,26 +297,59 @@ def analyze_ensemble_nex_gddp( 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 diff --git a/climate_tookit/climatology/long_term_climatology.py b/climate_tookit/climatology/long_term_climatology.py index a502a0a..1e8893d 100644 --- a/climate_tookit/climatology/long_term_climatology.py +++ b/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, diff --git a/climate_tookit/compare_periods/ensemble_periods.py b/climate_tookit/compare_periods/ensemble_periods.py index 8a2f7b9..b3404a4 100644 --- a/climate_tookit/compare_periods/ensemble_periods.py +++ b/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) diff --git a/climate_tookit/season_analysis/ensemble.py b/climate_tookit/season_analysis/ensemble.py index 7d9cc23..b3be697 100644 --- a/climate_tookit/season_analysis/ensemble.py +++ b/climate_tookit/season_analysis/ensemble.py @@ -359,7 +359,7 @@ def aggregate_overall(model_results: List[Dict]): return ensemble, model_averages # Top-level orchestrator -def run_ensemble(lat, lon, start_year, end_year, scenarios, models, fixed_arg=None, verbose=True): +def run_ensemble(lat, lon, start_year, end_year, scenarios, models, fixed_arg=None, verbose=True, max_workers=0): results = {} mode = 'fixed' if fixed_arg else 'auto' @@ -370,34 +370,61 @@ def run_ensemble(lat, lon, start_year, end_year, scenarios, models, fixed_arg=No f"{start_year}–{end_year} | mode={mode}") print('=' * 70) - per_model = [] - for i, model in enumerate(models, 1): - if verbose: - print(f" [{i:02d}/{len(models):02d}] {model:<22}", end=' ', flush=True) + # Each model is an independent, I/O-bound GEE run; fetch them + # concurrently and restore the input order afterwards. + def _run_model(model): try: s_dict, a_dict, skip_info = analyze_one_model( lat, lon, start_year, end_year, model, scenario, fixed_arg) - per_model.append({ + return model, { 'model': model, 'seasons_dict': s_dict, 'annual_dict': a_dict, 'skip_info': skip_info, - }) - if verbose: - n_seasons = sum(len(v) for v in s_dict.values()) - n_years = sum(1 for v in s_dict.values() if v) - extra = '' - if mode == 'auto' and skip_info['perhumid_years']: - extra = f" [perhumid: {len(skip_info['perhumid_years'])}y]" - print(f"✓ {n_seasons} season(s) over {n_years} year(s){extra}") + }, None except Exception as exc: - per_model.append({ + return model, { 'model': model, 'seasons_dict': {}, 'annual_dict': {}, 'skip_info': {'perhumid_years': [], 'no_season_years': [], 'analyzed_years': []}, 'error': f"{type(exc).__name__}: {exc}", - }) - if verbose: - print(f"✗ {type(exc).__name__}: {exc}") + }, f"{type(exc).__name__}: {exc}" + + def _log(idx, entry, err): + if not verbose: + return + if err: + print(f" [{idx:02d}/{len(models):02d}] {model_name(entry):<22} ✗ {err}", flush=True) + else: + s_dict = entry['seasons_dict'] + n_seasons = sum(len(v) for v in s_dict.values()) + n_years = sum(1 for v in s_dict.values() if v) + extra = '' + if mode == 'auto' and entry['skip_info'].get('perhumid_years'): + extra = f" [perhumid: {len(entry['skip_info']['perhumid_years'])}y]" + print(f" [{idx:02d}/{len(models):02d}] {model_name(entry):<22} " + f"✓ {n_seasons} season(s) over {n_years} year(s){extra}", flush=True) + + def model_name(entry): + return entry.get('model', '?') + + by_model = {} + # max_workers <= 0 -> auto: one worker per model, capped at 16. + workers = max(1, min(len(models), 16) if max_workers <= 0 else min(max_workers, len(models))) + if workers == 1: + for i, model in enumerate(models, 1): + _, entry, err = _run_model(model) + by_model[model] = entry + _log(i, entry, err) + else: + from concurrent.futures import ThreadPoolExecutor, as_completed + order = {m: i for i, m in enumerate(models, 1)} + with ThreadPoolExecutor(max_workers=workers) as ex: + futs = {ex.submit(_run_model, m): m for m in models} + for fut in as_completed(futs): + m, entry, err = fut.result() + by_model[m] = entry + _log(order[m], entry, err) + per_model = [by_model[m] for m in models if m in by_model] ok = sum(1 for r in per_model if not r.get('error')) diagnostics = _aggregate_skip_info(per_model) @@ -610,6 +637,13 @@ def print_summary(results): def main(): global NEX_GDDP_SOURCE + # 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): @@ -633,6 +667,8 @@ def main(): p.add_argument('--list-models', action='store_true', help='Print models and exit') p.add_argument('--output', help='Save JSON result here') 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() if args.source_key: @@ -669,6 +705,7 @@ def main(): scenarios = scenarios, models = models, fixed_arg = args.fixed_season, verbose = not args.quiet, + max_workers = args.workers, ) if not args.quiet: