"""Helpers for the simple notebook (CAMulator_climate_simple.ipynb): one-year runs, monthly means, and its figures. Every year run starts on Dec 1 1981, runs 396 days and is analysed for the 12 months Jan-Dec 1982 (December 1981 is spin-up, so the run has forgotten its starting weather). Runs live in /rollout/climate/. Each run is saved every 6 hours and then reduced to 12 monthly means of rain, 2-m temperature, OLR and column water vapour (monthly.nc in its run directory); the raw output is deleted to save disk space (~8 GB per run before, ~10 MB after). """ import shutil import matplotlib.pyplot as plt import numpy as np import xarray as xr from . import diagnostics as dg from . import edits from . import plotting as pl # --- places -------------------------------------------------------------------------------------------------------- AFRICA = {"lat": (-38, 38), "lon": (-20, 55)} REGIONS = { "West Africa (Sahel)": {"lat": (10, 18), "lon": (-15, 20)}, "East Africa": {"lat": (-5, 5), "lon": (33, 42)}, "Southern Africa": {"lat": (-30, -15), "lon": (15, 35)}, } # --- variables: display name, units, colour maps for values and for differences ------------------------------------ VARIABLES = { "rain": ("Rain", "mm/day", "Blues", "BrBG"), "temperature": ("2-m temperature", "K", "RdYlBu_r", "RdBu_r"), "olr": ("OLR", "W/m²", "Greys", "BrBG_r"), # outgoing longwave radiation: lower = more deep, cold cloud "vapour": ("Column water vapour", "mm", "YlGnBu", "BrBG"), } MONTHS = ["Jan", "Feb", "Mar", "Apr", "May", "Jun", "Jul", "Aug", "Sep", "Oct", "Nov", "Dec"] SEASONS = {"DJF": [12, 1, 2], "MAM": [3, 4, 5], "JJA": [6, 7, 8], "SON": [9, 10, 11], "year": list(range(1, 13))} SEASON_NAMES = {"DJF": "December-February", "MAM": "March-May", "JJA": "June-August", "SON": "September-November", "year": "whole year"} def label(variable): name, units, _, _ = VARIABLES[variable] return f"{name} ({units})" # --- the experiments ----------------------------------------------------------------------------------------------- EXPERIMENTS = { "warm_ocean_2K": ("The whole ocean 2 K warmer", edits.add_sst(2.0)), "warm_ocean_4K": ("The whole ocean 4 K warmer", edits.add_sst(4.0)), "el_nino": ("El Niño all year: the eastern equatorial Pacific up to 3 K warmer", edits.add_sst_patches([{"peak": 3.0, "lat": 0, "lon": -110, "lat_width": 7.5, "lon_width": 30}])), "indian_ocean_dipole": ("Positive Indian Ocean Dipole all year: west 1.5 K warmer, east 1.5 K cooler", edits.add_sst_patches([{"peak": 1.5, "lat": 0, "lon": 60, "lat_width": 10, "lon_width": 15}, {"peak": -1.5, "lat": -5, "lon": 100, "lat_width": 7, "lon_width": 12}])), "warm_gulf_of_guinea": ("Atlantic Niño all year: the Gulf of Guinea up to 2 K warmer", edits.add_sst_patches([{"peak": 2.0, "lat": 0, "lon": -8, "lat_width": 4, "lon_width": 10}])), "no_arctic_sea_ice": ("No Arctic sea ice, all year", edits.remove_sea_ice(lat_min=0.0)), } YEAR_START, YEAR_LENGTH, YEAR = "1981-12-01", "396d", 1982 YEAR_OUTPUTS = ["PRECT", "TREFHT", "FLUT", "PS", "Qtot"] # --- running and loading ------------------------------------------------------------------------------------------- def run_year(tut, name, experiment=None, starting_errors=0.0, fast_math=True): """Run one year (see the module docstring) and keep its monthly means. experiment: a key of EXPERIMENTS (None: the average forcing); starting_errors: K of random temperature errors added to the starting state. A run that already exists is not repeated.""" if (tut.run_dir(f"climate/{name}") / "monthly.nc").exists(): print(f"{name}: already run, using the saved result") return forcing = tut.forcing("climo", YEAR_START, YEAR_LENGTH) if experiment is not None: forcing = tut.perturb_forcing(forcing, f"year_{experiment}", YEAR_START, YEAR_LENGTH, EXPERIMENTS[experiment][1]) config = tut.member(f"climate/{name}", start=YEAR_START, length=YEAR_LENGTH, forcing=forcing, outputs=YEAR_OUTPUTS, amplitude=starting_errors, seed=1) tut.run([config], fast_math=fast_math) _keep_monthly_means(tut, name) def _keep_monthly_means(tut, name): run = tut.load(f"climate/{name}") run = run.isel(time=np.nonzero(run["time.year"].values == YEAR)[0]) out = xr.Dataset({"rain": run["PRECT"], "temperature": run["TREFHT"], "olr": run["FLUT"], "vapour": tut.precipitable_water(run)}) out = out.groupby("time.month").mean().astype("float32").load() out = out.drop_vars([c for c in out.coords if c not in ("month", "latitude", "longitude")]) out.to_netcdf(tut.run_dir(f"climate/{name}") / "monthly.nc") for raw in tut.run_dir(f"climate/{name}").glob("*_*Z"): shutil.rmtree(raw) def load_year(tut, name): """The 12 monthly means of a year run: rain (mm/day), temperature (K), olr (W/m²), vapour (mm).""" return xr.open_dataset(tut.run_dir(f"climate/{name}") / "monthly.nc").load() def load_truth(tut, filename): """CESM's monthly means, in the same form as load_year (see maintainer/make_truth_monthly.py).""" return xr.open_dataset(tut.data_dir / filename).load() def season(ds, which): """Average over the months of a season ('DJF', 'MAM', 'JJA', 'SON') or the whole 'year'.""" return ds.sel(month=SEASONS[which]).mean("month") def global_mean(da): """Area-weighted global mean (one value per month if da has months).""" return dg.area_mean(da) def noise_level(noise, variable, region=None): """The typical size of the noise in a region's average: the root-mean-square, over the 12 months, of the region-average difference between the two control runs (noise = control_repeat - control). Using all 12 months makes it steadier than one season's difference, and a season's average is less noisy than one month, so it errs on the side of caution.""" return float(np.sqrt((dg.area_mean(noise[variable], region) ** 2).mean())) # --- figures ------------------------------------------------------------------------------------------------------- def show_maps(panels, variable, region=None, difference=False, title=None, columns=2, boxes=None): """One map per entry of panels {subtitle: field}, all on one colour scale. region: None for the globe (centred on Africa) or e.g. AFRICA. difference=True: a colour scale centred on zero. boxes: {name: region} outlined on every map (e.g. REGIONS).""" name, units, cmap, diff_cmap = VARIABLES[variable] values = [np.asarray(dg.select(field, region) if region else field) for field in panels.values()] if difference: high = max(float(np.nanquantile(np.abs(v), 0.98)) for v in values) or 1.0 low, cmap = -high, diff_cmap else: low = min(float(np.nanquantile(v, 0.02)) for v in values) high = max(float(np.nanquantile(v, 0.98)) for v in values) rows = int(np.ceil(len(panels) / columns)) if region is None: width, height = 9, 4.6 else: # follow the region's shape, so the maps fill their space (la0, la1), (lo0, lo1) = region["lat"], region["lon"] width = 6 height = width * (la1 - la0) / ((lo1 - lo0) % 360 or 360) + 0.8 fig = plt.figure(figsize=(width * columns + 1.5, height * rows), layout="constrained") axes = [] for i, (subtitle, field) in enumerate(panels.items()): if region is None: ax = pl.map_axis(fig, rows, columns, i + 1, central_longitude=0) ax.set_global() else: ax = pl.map_axis(fig, rows, columns, i + 1, region=region) mesh = ax.pcolormesh(field["longitude"], field["latitude"], field, cmap=cmap, vmin=low, vmax=high, transform=pl.DATA_CRS) _draw_boxes(ax, boxes) ax.set_title(subtitle) axes.append(ax) fig.colorbar(mesh, ax=axes, shrink=0.8, label=("Change in " if difference else "") + f"{name} ({units})") if title: fig.suptitle(title, fontsize=14, fontweight="bold") plt.show() def show_comparison(model, truth, variable, seasons, names=("CAMulator", "CESM"), region=None, boxes=None, title=None): """One row per season ('year', 'JJA', ...): the model, the truth (one colour scale) and model minus truth (its own colour scale centred on zero; for rain, green = the model is wetter). model, truth: monthly datasets.""" name, units, cmap, diff_cmap = VARIABLES[variable] rows = [(season(model, s)[variable], season(truth, s)[variable]) for s in seasons] cut = (lambda f: dg.select(f, region)) if region else (lambda f: f) values = [np.asarray(cut(f)) for pair in rows for f in pair] low = min(float(np.nanquantile(v, 0.02)) for v in values) high = max(float(np.nanquantile(v, 0.98)) for v in values) diff_high = max(float(np.nanquantile(np.abs(np.asarray(cut(m - t))), 0.98)) for m, t in rows) or 1.0 if region is None: width, height = 6, 3.4 else: # follow the region's shape, so the maps fill their space (la0, la1), (lo0, lo1) = region["lat"], region["lon"] width = 4.5 height = width * (la1 - la0) / ((lo1 - lo0) % 360 or 360) + 0.6 fig = plt.figure(figsize=(width * 3 + 1, height * len(seasons) + 1.5), layout="constrained") value_axes, diff_axes = [], [] for row, (when, (m, t)) in enumerate(zip(seasons, rows)): when_text = SEASON_NAMES[when] panels = [(f"{names[0]}, {when_text}", m, cmap, low, high, value_axes), (f"{names[1]}, {when_text}", t, cmap, low, high, value_axes), (f"{names[0]} minus {names[1]}", m - t, diff_cmap, -diff_high, diff_high, diff_axes)] for col, (subtitle, field, colours, vmin, vmax, group) in enumerate(panels): index = row * 3 + col + 1 ax = pl.map_axis(fig, len(seasons), 3, index, region=region) if region else \ pl.map_axis(fig, len(seasons), 3, index, central_longitude=0) if region is None: ax.set_global() mesh = ax.pcolormesh(field["longitude"], field["latitude"], field, cmap=colours, vmin=vmin, vmax=vmax, transform=pl.DATA_CRS) _draw_boxes(ax, boxes) ax.set_title(subtitle) group.append((ax, mesh)) fig.colorbar(value_axes[0][1], ax=[a for a, _ in value_axes], orientation="horizontal", shrink=0.5, label=f"{name} ({units})") fig.colorbar(diff_axes[0][1], ax=[a for a, _ in diff_axes], orientation="horizontal", shrink=0.9, label=f"Difference in {name.lower() if name != 'OLR' else name} ({units})") if title: fig.suptitle(title, fontsize=14, fontweight="bold") plt.show() def _draw_boxes(ax, boxes): """Outline each region {name: region} in red, with its name.""" for box_name, box in (boxes or {}).items(): (la0, la1), (lo0, lo1) = box["lat"], box["lon"] ax.plot([lo0, lo1, lo1, lo0, lo0], [la0, la0, la1, la1, la0], color="C3", lw=1.5, transform=pl.DATA_CRS) ax.text(lo0 + 0.5, la1 + 0.8, box_name, color="C3", fontsize=8, transform=pl.DATA_CRS) def show_months(lines, variable, regions=None, title=None, grey=()): """Month by month, one panel per region: one line per entry of lines {legend label: monthly dataset}. regions: {name: region} (default REGIONS); {'Whole globe': None} for the global mean. Labels in `grey` are drawn as thin grey dashed lines (e.g. a repeat run, to show the noise).""" regions = REGIONS if regions is None else regions fig, axes = plt.subplots(1, len(regions), figsize=(max(6.5 * len(regions), 10), 4.2), layout="constrained", squeeze=False) for ax, (region_name, region) in zip(axes[0], regions.items()): for i, (line_label, ds) in enumerate(lines.items()): values = dg.area_mean(ds[variable], region) if line_label in grey: ax.plot(range(1, 13), values, color="0.55", lw=1.5, ls="--", label=line_label) else: style = dict(color="k", lw=2.5) if line_label.startswith("CESM") else dict(color=f"C{i}", lw=2) ax.plot(range(1, 13), values, marker="o", ms=4, label=line_label, **style) ax.set_xticks(range(1, 13), MONTHS) ax.set_title(region_name if region is None else f"{region_name} ({pl.region_text(region)})") ax.set_ylabel(label(variable)) ax.grid(alpha=0.3) axes[0][0].legend() if title: fig.suptitle(title, fontsize=14, fontweight="bold") plt.show() def show_sst_change(tut, experiment): """Map of what an experiment changes (on Jan 1): the SST, or for the sea-ice experiment the ice that is removed.""" first = tut.first_forcing_step(tut.forcing("climo", YEAR_START, YEAR_LENGTH), f"{YEAR}-01-01") edited = EXPERIMENTS[experiment][1](first.copy(deep=True)) fig = plt.figure(figsize=(12, 5.5), layout="constrained") if experiment == "no_arctic_sea_ice": removed = (first["ICEFRAC"] - edited["ICEFRAC"]).isel(time=0) ax = pl.polar_axis(fig, 1, 1, 1, lat_min=45) mesh = ax.pcolormesh(removed["longitude"], removed["latitude"], removed, cmap="Blues", vmin=0, vmax=1, transform=pl.DATA_CRS) fig.colorbar(mesh, ax=ax, shrink=0.8, label="Sea-ice cover removed (fraction of each grid box)") else: change = (edited["SST"] - first["SST"]).isel(time=0) high = max(float(np.abs(change).max()), 0.5) ax = pl.map_axis(fig, 1, 1, 1, central_longitude=0) ax.set_global() mesh = ax.pcolormesh(change["longitude"], change["latitude"], change, cmap="RdBu_r", vmin=-high, vmax=high, transform=pl.DATA_CRS) fig.colorbar(mesh, ax=ax, shrink=0.8, label="SST change (K)") ax.set_title(EXPERIMENTS[experiment][0], fontsize=14, fontweight="bold") plt.show() # --- part 1: how far two runs drift apart ------------------------------------------------------------------------- def drift(run, other): """Typical (global root-mean-square) difference in surface pressure between two runs (hPa), every 6 hours.""" difference = run["PS"].load() - other["PS"].sel(time=run["time"]).load() return dg.rms(difference) # --- part 4: water vapour and rain by latitude --------------------------------------------------------------------- def by_latitude(tut, control, warm): """Ocean-only (60°S-60°N) water vapour change in % per K of local warming, with the Clausius-Clapeyron rate at each latitude's temperature; and the change in rain (mm/day) by latitude, all longitudes.""" ocean = ~tut.land_mask(control) zonal = {} for key, ds in (("control", control), ("warm", warm)): zonal[key] = annual_zonal(ds.where(ocean)).sel(latitude=slice(-60, 60)) local_warming = zonal["warm"]["temperature"] - zonal["control"]["temperature"] vapour_rate = 100 * (zonal["warm"]["vapour"] / zonal["control"]["vapour"] - 1) / local_warming return {"vapour_rate": vapour_rate, "clausius_clapeyron": dg.clausius_clapeyron_rate(zonal["control"]["temperature"]), "rain_change": annual_zonal(warm)["rain"] - annual_zonal(control)["rain"], "control_rain": annual_zonal(control)["rain"]} def annual_zonal(ds): return ds.mean("month").mean("longitude")