Upload simple.py

0529de9 verified 20 days ago

15.7 kB

"""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 <work_dir>/rollout/climate/<name>. 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")

	"""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 <work_dir>/rollout/climate/<name>. 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")