| """Diagnostics for the CAMulator timescales tutorial: plain functions on xarray objects. |
| |
| Conventions: latitude/longitude coordinates named 'latitude' and 'longitude' (0-360 in the data). A region is a dict |
| {'lat': (lat_min, lat_max), 'lon': (lon_min, lon_max)} with longitudes in degrees east (negative = west, e.g. -20 for |
| 20°W); lon_min > lon_max means the region crosses the date line. The older tuple form (lat_min, lat_max, lon_min, |
| lon_max) in 0-360 is also accepted. |
| """ |
|
|
| import cftime |
| import numpy as np |
| import pandas as pd |
| import xarray as xr |
| from scipy.ndimage import minimum_filter |
|
|
| LAT, LON = "latitude", "longitude" |
| G = 9.80665 |
| MM_PER_DAY = 4000.0 |
| KM_PER_DEG = 111.2 |
|
|
|
|
| |
| def area_weights(lat): |
| """Relative cell area ∝ |Δ sin(lat)|, computed as CREDIT's conservation fixers do (cells are equal in longitude).""" |
| lat = np.asarray(lat) |
| return xr.DataArray(np.abs(np.gradient(np.sin(np.deg2rad(lat)), edge_order=2)), coords={LAT: lat}, dims=LAT) |
|
|
|
|
| def select(da, region=None): |
| """Subset to a region (see the module docstring); a region without 'lon' keeps all longitudes.""" |
| if region is None: |
| return da |
| if isinstance(region, dict): |
| la0, la1 = region["lat"] |
| lo0, lo1 = region.get("lon", (0, 360)) |
| if lo1 - lo0 >= 360: |
| return da.sel({LAT: slice(la0, la1)}) |
| lo0, lo1 = lo0 % 360, lo1 % 360 |
| else: |
| la0, la1, lo0, lo1 = region |
| da = da.sel({LAT: slice(la0, la1)}) |
| if lo0 <= lo1: |
| return da.sel({LON: slice(lo0, lo1)}) |
| return xr.concat([da.sel({LON: slice(lo0, 360)}), da.sel({LON: slice(0, lo1)})], dim=LON) |
|
|
|
|
| def area_mean(da, region=None, mask=None): |
| """Area-weighted mean over latitude and longitude (optionally in a region and/or where a boolean mask is True).""" |
| da = select(da, region) |
| w = area_weights(da[LAT].values) |
| if mask is not None: |
| w = w * select(mask, region).astype(float) |
| return da.weighted(w).mean([LAT, LON]) |
|
|
|
|
| def rms(da, region=None): |
| """Area-weighted root-mean-square over latitude and longitude.""" |
| return np.sqrt(area_mean(da**2, region)) |
|
|
|
|
| def zonal_mean(da): |
| return da.mean(LON) |
|
|
|
|
| def band_mean(zm, lat_min, lat_max): |
| """Area-weighted mean of a zonal-mean field over a latitude band.""" |
| sub = zm.sel({LAT: slice(lat_min, lat_max)}) |
| return sub.weighted(area_weights(sub[LAT].values)).mean(LAT) |
|
|
|
|
| def noleap(when): |
| """'1981-03-03T06' -> cftime no-leap datetime (the calendar of CAMulator's data and output).""" |
| t = pd.Timestamp(when) |
| return cftime.DatetimeNoLeap(t.year, t.month, t.day, t.hour) |
|
|
|
|
| def days_since(times, ref): |
| """Days from cftime `ref` to each of `times`.""" |
| return np.array([(t - ref).total_seconds() for t in np.asarray(times)]) / 86400.0 |
|
|
|
|
| |
| def layer_thickness(ds, hyai, hybi, like="Qtot", lev="level"): |
| """Pressure thickness Δp_k (Pa) of each model layer from the interfaces p = hyai + hybi * PS (hyai in Pa), float32, |
| with the dims ordered like ds[like].""" |
| ilev = hyai.dims[0] |
|
|
| def per_layer(c): |
| d = c.diff(ilev).rename({ilev: lev}).drop_vars(lev, errors="ignore") |
| return d.assign_coords({lev: ds[lev].values}).astype("float32") |
|
|
| return (per_layer(hyai) + per_layer(hybi) * ds["PS"].astype("float32")).transpose(*ds[like].dims) |
|
|
|
|
| def precipitable_water(ds, hyai, hybi): |
| """Column water PW = (1/g) Σ Qtot_k Δp_k (kg/m² = mm). Stays lazy if Qtot is lazy.""" |
| return (ds["Qtot"].astype("float32") * layer_thickness(ds, hyai, hybi)).sum("level") / G |
|
|
|
|
| |
| def mslp(ps, trefht, phis): |
| """Mean sea-level pressure (Pa) from PS, TREFHT and surface geopotential PHIS, using CREDIT's own Trenberth et al. |
| (1993) formula (credit.postblock.mslp, the function behind its mslp_diagnostic postblock).""" |
| import torch |
| from credit.postblock.mslp import mslp_from_surface_pressure |
|
|
| assert phis.shape == ps.shape[-2:], "PHIS must be on the same grid as PS" |
| out = mslp_from_surface_pressure( |
| torch.as_tensor(np.asarray(ps.values), dtype=torch.float64), |
| torch.as_tensor(np.asarray(trefht.values), dtype=torch.float64), |
| torch.as_tensor(np.asarray(phis.values), dtype=torch.float64), |
| ) |
| return xr.DataArray(out.numpy(), coords=ps.coords, dims=ps.dims, name="MSLP") |
|
|
|
|
| def great_circle_deg(lat1, lon1, lat2, lon2): |
| """Great-circle distance in degrees of arc.""" |
| r = np.deg2rad |
| c = np.sin(r(lat1)) * np.sin(r(lat2)) + np.cos(r(lat1)) * np.cos(r(lat2)) * np.cos(r(lon1) - r(lon2)) |
| return np.rad2deg(np.arccos(np.clip(c, -1, 1))) |
|
|
|
|
| def follow_low(slp, time, lat, lon, max_jump=6.0): |
| """Track one low forward and backward in time from (time, lat, lon): at each step, the nearest local minimum of slp |
| within max_jump degrees of the previous position. Returns DataFrame(time, lat, lon, p) sorted by time.""" |
| lats, lons = slp[LAT].values, slp[LON].values |
| times = list(slp["time"].values) |
|
|
| def minima(k): |
| f = slp.isel(time=k).values |
| i, j = np.nonzero(f == minimum_filter(f, size=(5, 5), mode=("nearest", "wrap"))) |
| return lats[i], lons[j], f[i, j] |
|
|
| def walk(k, a, o, step): |
| rows = [] |
| while 0 <= k < len(times): |
| la, lo, p = minima(k) |
| d = great_circle_deg(a, o, la, lo) |
| if not (d < max_jump).any(): |
| break |
| n = int(np.argmin(d)) |
| a, o = la[n], lo[n] |
| rows.append({"time": times[k], "lat": a, "lon": o, "p": p[n]}) |
| k += step |
| return rows |
|
|
| k0 = times.index(time) |
| rows = walk(k0, lat, lon, +1) |
| rows += walk(k0 - 1, rows[0]["lat"], rows[0]["lon"], -1) if rows else [] |
| return pd.DataFrame(rows).sort_values("time").reset_index(drop=True) |
|
|
|
|
| def track_near(slp, ref, radius=10.0): |
| """The storm in a forecast, located relative to a reference track: at each reference time the forecast covers, the |
| lowest slp within `radius` degrees of the reference position. If that lowest point is on the rim of the search disk |
| there is no closed low inside, and the row is NaN. Returns DataFrame(time, lat, lon, p).""" |
| lats, lons = slp[LAT].values, slp[LON].values |
| LA, LO = np.meshgrid(lats, lons, indexing="ij") |
| times = list(slp["time"].values) |
| rows = [] |
| for t, a, o in zip(ref["time"], ref["lat"], ref["lon"]): |
| if t not in times: |
| continue |
| f = slp.isel(time=times.index(t)).values |
| d = great_circle_deg(a, o, LA, LO) |
| i, j = np.unravel_index(np.argmin(np.where(d <= radius, f, np.inf)), f.shape) |
| if d[i, j] > radius - 2.0: |
| rows.append({"time": t, "lat": np.nan, "lon": np.nan, "p": np.nan}) |
| else: |
| rows.append({"time": t, "lat": LA[i, j], "lon": LO[i, j], "p": f[i, j]}) |
| return pd.DataFrame(rows, columns=["time", "lat", "lon", "p"]) |
|
|
|
|
| def storm_summary(track, ref): |
| """Deepest point of a forecast track, and its position error (km) at the time the reference track is deepest.""" |
| peak = ref.loc[ref["p"].idxmin()] |
| at_peak = track[track["time"] == peak["time"]] |
| if track["p"].notna().any(): |
| deepest = track.loc[track["p"].idxmin()] |
| min_p, t_min = deepest["p"], deepest["time"] |
| else: |
| min_p, t_min = np.nan, None |
| err = np.nan |
| if len(at_peak) and np.isfinite(at_peak["lat"].iloc[0]): |
| err = great_circle_deg(peak["lat"], peak["lon"], at_peak["lat"].iloc[0], at_peak["lon"].iloc[0]) * KM_PER_DEG |
| return {"min_p": min_p, "time_of_min": t_min, "peak_position_error_km": err} |
|
|
|
|
| |
| def doubling_time(lead, err, max_err, min_lead=1.0): |
| """Days per doubling from a straight-line fit of log2(err) against lead, using leads >= min_lead where |
| err <= max_err (the exponential phase, before saturation). NaN if fewer than 3 points qualify.""" |
| lead, err = np.asarray(lead, float), np.asarray(err, float) |
| ok = (lead >= min_lead) & (err <= max_err) & (err > 0) |
| if ok.sum() < 3: |
| return np.nan |
| slope = np.polyfit(lead[ok], np.log2(err[ok]), 1)[0] |
| return 1.0 / slope if slope > 0 else np.inf |
|
|
|
|
| def first_lead(lead, values, threshold): |
| """First lead at which values >= threshold (NaN if never).""" |
| hit = np.nonzero(np.asarray(values) >= threshold)[0] |
| return float(np.asarray(lead)[hit[0]]) if hit.size else np.nan |
|
|
|
|
| |
| def block_mean(da, spinup_days, block_days): |
| """Mean over consecutive block_days-long blocks after spin-up, using da's 'lead' coordinate (days since the start). |
| Dim 'block'; an incomplete last block is dropped.""" |
| lead = da["lead"].values |
| keep = np.nonzero(lead > spinup_days)[0] |
| idx = np.floor((lead[keep] - spinup_days - 1e-9) / block_days).astype(int) |
| counts = np.bincount(idx) |
| out = da.isel(time=keep).drop_vars("lead").assign_coords(time=idx).rename(time="block").groupby("block").mean() |
| return out.sel(block=np.nonzero(counts == counts.max())[0]) |
|
|
|
|
| def block_quantile(da, q, spinup_days, block_days, region=None): |
| """Quantile q of all values (time, lat, lon) in each block after spin-up, within a region.""" |
| da = select(da, region) |
| lead = da["lead"].values |
| idx = np.floor((lead - spinup_days - 1e-9) / block_days).astype(int) |
| blocks = [b for b in np.unique(idx[lead > spinup_days]) if (idx == b).sum() == np.bincount(idx[idx >= 0]).max()] |
| vals = [float(np.quantile(da.isel(time=np.nonzero(idx == b)[0]).values, q)) for b in blocks] |
| return xr.DataArray(vals, coords={"block": blocks}, dims="block") |
|
|
|
|
| def paired(exp, ctl, noise_a, noise_b): |
| """One-year Δ = block mean of (exp - ctl); the noise pair's own Δ (noise_b - noise_a); and the standard error of a |
| one-year Δ from the noise pair's block differences (blocks paired by date). Returns (delta, noise_delta, se).""" |
| noise = noise_b - noise_a |
| se = noise.std("block", ddof=1) / np.sqrt(noise.sizes["block"]) |
| return (exp - ctl).mean("block"), noise.mean("block"), se |
|
|
|
|
| def itcz(zm, band=20.0): |
| """ITCZ metrics from zonal-mean precipitation zm(..., latitude), on a 0.1° grid between band S and N: |
| centroid (latitude splitting the area-weighted precipitation in half; Frierson & Hwang 2012), spread |
| (precipitation-weighted standard deviation of latitude; smaller = narrower) and P_trop (band-mean precipitation).""" |
| fine = np.arange(-band, band + 1e-6, 0.1) |
| w = np.cos(np.deg2rad(fine)) |
| zm = zm.sortby(LAT) |
| flat = zm.transpose(..., LAT).values.reshape(-1, zm.sizes[LAT]) |
| rows = [] |
| for p in flat: |
| pw = np.interp(fine, zm[LAT].values, p) * w |
| cum = np.cumsum(pw) |
| mean_lat = (pw * fine).sum() / pw.sum() |
| rows.append((np.interp(0.5 * cum[-1], cum, fine), np.sqrt((pw * (fine - mean_lat) ** 2).sum() / pw.sum()), |
| pw.sum() / w.sum())) |
| rows = np.array(rows) |
| other = [d for d in zm.dims if d != LAT] |
| shape = [zm.sizes[d] for d in other] |
| coords = {d: zm[d] for d in other if d in zm.coords} |
| return xr.Dataset({k: (other, rows[:, i].reshape(shape)) for i, k in enumerate(["centroid", "spread", "P_trop"])}, |
| coords=coords) |
|
|
|
|
| |
| def zonal_spectrum(da, lat_min, lat_max): |
| """Power at each zonal wavenumber k (number of waves around a latitude circle), from an FFT along longitude, |
| area-weighted over a latitude band and averaged over every other dimension (e.g. time). DataArray over 'wavenumber'.""" |
| sub = da.sel({LAT: slice(lat_min, lat_max)}) |
| x = np.asarray(sub.transpose(..., LAT, LON).values, dtype=float) |
| p = np.abs(np.fft.rfft(x - x.mean(axis=-1, keepdims=True), axis=-1)) ** 2 |
| w = area_weights(sub[LAT].values).values |
| p = (p * w[:, None]).sum(axis=-2) / w.sum() |
| p = p.reshape(-1, p.shape[-1]).mean(axis=0) |
| return xr.DataArray(p, coords={"wavenumber": np.arange(p.size)}, dims="wavenumber") |
|
|
|
|
| def t95(n): |
| """Half-width of a 95% range in standard errors for a mean of n samples (Student's t, n - 1 degrees of freedom): |
| 4.30 for 3 samples, 2.23 for 11, approaching 1.96 for many.""" |
| from scipy.stats import t |
|
|
| return float(t.ppf(0.975, n - 1)) |
|
|
|
|
| def log_binned(spectrum, n_bins=18): |
| """Average a spectrum (over 'wavenumber') in log-spaced wavenumber bins from 1 to the largest wavenumber, which |
| smooths the noisy small scales. Returns (bin centres, binned power) as numpy arrays.""" |
| kmax = int(spectrum["wavenumber"].max()) |
| edges = np.unique(np.round(np.logspace(0, np.log10(kmax + 1), n_bins + 1)).astype(int)) |
| power = np.array([float(spectrum.sel(wavenumber=slice(lo, hi - 1)).mean()) for lo, hi in zip(edges[:-1], edges[1:])]) |
| return np.sqrt(edges[:-1] * np.maximum(edges[1:] - 1, edges[:-1])), power |
|
|
|
|
| |
| def east_west(da): |
| """Longitudes as -180..180 (negative = west), sorted west to east.""" |
| return da.assign_coords({LON: (da[LON] + 180) % 360 - 180}).sortby(LON) |
|
|
|
|
| def after_spinup(ds, days): |
| """Drop the first `days` of a run (uses the 'lead' coordinate).""" |
| return ds.isel(time=np.nonzero(ds["lead"].values > days)[0]) |
|
|
|
|
| def mean_and_range(values, dim): |
| """Mean over `dim` and the half-width of its 95 % range (Student's t, for few samples).""" |
| n = values.sizes[dim] |
| return values.mean(dim), t95(n) * values.std(dim, ddof=1) / np.sqrt(n) |
|
|
|
|
| |
| def plateau(drift, last_days=3): |
| """The level a drift curve settles at: its average over the last few days.""" |
| return float(drift.where(drift["lead"] > float(drift["lead"].max()) - last_days, drop=True).mean()) |
|
|
|
|
| def doubling_days(drift, unrelated): |
| """Days per doubling while the difference is still small (below a quarter of `unrelated`).""" |
| return doubling_time(drift["lead"], drift, 0.25 * unrelated) |
|
|
|
|
| def time_to_reach(drift, level): |
| """First lead (days) at which the drift reaches `level`.""" |
| return first_lead(drift["lead"].values, drift.values, level) |
|
|
|
|
| def equivalent_noise(model_error, copy_drift, copy_noise, day=2): |
| """Starting noise (K) that would give a copy as large a difference as the model error, at a given day. Early on, |
| differences scale with the starting noise, so this is the copy's noise times the ratio of the two curves.""" |
| k = int(np.argmin(abs(model_error["lead"].values - day))) |
| return copy_noise * float(model_error[k]) / float(copy_drift[k]) |
|
|
|
|
| def wave_diagram(v, region): |
| """Time-longitude diagram: wind averaged over the region's latitudes, longitudes -180..180.""" |
| return east_west(select(v, region).mean(LAT)).load() |
|
|
|
|
| def wave_slice(field, lon): |
| """One column of a wave diagram: the wind over time at the longitude nearest `lon`.""" |
| return field.sel({LON: lon}, method="nearest") |
|
|
|
|
| def wave_peak(series, when, window_days=1.25): |
| """Strongest northward wind in a wave_slice within +-window_days of a date (waves pass every ~3 days, so this |
| finds the one wave near `when`).""" |
| t = noleap(when) |
| series = series.sel(time=slice(t - pd.Timedelta(f"{window_days}D"), t + pd.Timedelta(f"{window_days}D"))) |
| k = int(series.argmax("time")) |
| return {"time": series["time"].values[k], "wind": float(series[k])} |
|
|
|
|
| def hours_between(later, earlier): |
| return (later - earlier).total_seconds() / 3600 |
|
|
|
|
| def weekly_correlation(a, b, start): |
| """Correlation between two wave diagrams, week by week from the start date.""" |
| days = days_since(a["time"].values, noleap(start)) |
| out = [] |
| for week in range(int(np.ceil(days.max() / 7))): |
| sel = np.nonzero((days > 7 * week) & (days <= 7 * (week + 1)))[0] |
| out.append(float(np.corrcoef(a.isel(time=sel).values.ravel(), b.isel(time=sel).values.ravel())[0, 1])) |
| return out |
|
|
|
|
| |
| def spectrum_ratio(forecast, truth, lead, lats, n_bins=18): |
| """Forecast power / truth power at each (binned) zonal wavenumber, over latitudes `lats` and the forecast times in a |
| window around `lead` (just the first step for lead < 1; +-1 day otherwise). Returns (wavenumbers, ratio) or None.""" |
| lo, hi = (0, 0.5) if lead < 1 else (lead - 1, lead + 1) |
| truth_times = set(truth["time"].values) |
| keep = [i for i, (ld, t) in enumerate(zip(forecast["lead"].values, forecast["time"].values)) |
| if lo < ld <= hi and t in truth_times] |
| if not keep: |
| return None |
| f = forecast.isel(time=keep) |
| k, pf = log_binned(zonal_spectrum(f, *lats), n_bins) |
| _, pt = log_binned(zonal_spectrum(truth.sel(time=f["time"]), *lats), n_bins) |
| return k, pf / pt |
|
|
|
|
| def effective_resolution(k, ratio, min_k=10): |
| """The wavenumber (>= min_k) beyond which the forecast has less than half of the truth's variation for good, |
| interpolated (in log wavenumber) between the last bin at or above one half and the next one. NaN if never.""" |
| k, ratio = np.asarray(k, dtype=float), np.asarray(ratio, dtype=float) |
| above = np.nonzero((k >= min_k) & (ratio >= 0.5))[0] |
| if len(above) == 0 or above[-1] == len(k) - 1: |
| return np.nan |
| i = above[-1] |
| fraction = (ratio[i] - 0.5) / (ratio[i] - ratio[i + 1]) |
| return float(np.exp(np.log(k[i]) + fraction * (np.log(k[i + 1]) - np.log(k[i])))) |
|
|
|
|
| |
| def weekly_difference(experiment, control, variable, region): |
| """Experiment minus control, member by member, as weekly means over a region: DataArray (member, week).""" |
| diffs = [] |
| for exp_run, ctl_run in zip(experiment, control): |
| e = block_mean(area_mean(exp_run[variable], region), 0, 7) |
| c = block_mean(area_mean(ctl_run[variable], region), 0, 7) |
| diffs.append((e - c).load()) |
| d = xr.concat(diffs, "member") |
| return d.assign_coords(block=d["block"] + 1).rename(block="week") |
|
|
|
|
| def noise_band(control, variable, region, from_week=3): |
| """How big a difference between two runs' weekly means gets from weather noise alone: the largest |control i - |
| control j| over all pairs of control members and all weeks from `from_week` on (by then the members' weather is |
| unrelated, section 1). Falls back to the last week if the run is shorter.""" |
| runs = [block_mean(area_mean(run[variable], region), 0, 7).load() for run in control] |
| pairs = [abs(runs[i] - runs[j]) for i in range(len(runs)) for j in range(i + 1, len(runs))] |
| weekly = xr.concat(pairs, "pair") |
| weeks = weekly["block"] + 1 |
| return float(weekly.where(weeks >= min(from_week, int(weeks.max()))).max()) |
|
|
|
|
| def outside_noise(weekly, band): |
| """Number of members whose difference is bigger than the noise band, each week. DataArray (week).""" |
| return (abs(weekly) > band).sum("member") |
|
|
|
|
| def map_difference(run, other, variable, days): |
| """Map of run minus other, averaged over the lead days (first, last].""" |
| def average(r): |
| lead = r["lead"].values |
| return r[variable].isel(time=np.nonzero((lead > days[0]) & (lead <= days[1]))[0]).mean("time") |
| return (average(run) - average(other)).load() |
|
|
|
|
| |
| SEASON_MONTHS = {"DJF": [12, 1, 2], "MAM": [3, 4, 5], "JJA": [6, 7, 8], "SON": [9, 10, 11]} |
|
|
|
|
| def clausius_clapeyron_rate(temperature): |
| """Clausius-Clapeyron: % increase in saturation vapour pressure per K, L / (R_v T^2), at temperature T (K).""" |
| return 100 * 2.501e6 / (461.5 * temperature**2) |
|
|
|
|
| def monthly(da): |
| """Calendar-month means (one year of output gives each month once).""" |
| return da.groupby("time.month").mean() |
|
|
|
|
| def percent_change_per_month(warm, control, region=None, mask=None): |
| """Monthly % change of a field's area mean, warm vs control: DataArray over month.""" |
| w = monthly(area_mean(warm, region, mask)) |
| c = monthly(area_mean(control, region, mask)) |
| return 100 * (w / c - 1), c |
|
|
|
|
| def seasonal_change(warm, control, region, season): |
| """% change of a region's rain in one season, its 95 % range (from the season's three months), and the control's |
| mean in that season.""" |
| pct, ctl = percent_change_per_month(warm, control, region) |
| months = SEASON_MONTHS[season] |
| if not set(months) <= set(pct["month"].values): |
| return np.nan, np.nan, np.nan |
| mean, half = mean_and_range(pct.sel(month=months), "month") |
| return float(mean), float(half), float(ctl.sel(month=months).mean()) |
|
|
|
|
| def days_per_tenfold(drift, useful_until_days): |
| """How many days of useful forecast each 10x smaller starting error buys, from the copies' 'useful until' days. |
| drift: {noise_K: drift curve}; useful_until_days: {noise_K: days}.""" |
| noises = sorted(useful_until_days) |
| gains = [] |
| for small, large in zip(noises[:-1], noises[1:]): |
| if np.isfinite(useful_until_days[small]) and np.isfinite(useful_until_days[large]): |
| gains.append((useful_until_days[small] - useful_until_days[large]) / np.log10(large / small)) |
| return float(np.mean(gains)) if gains else np.nan |
|
|
|
|
| def small_scale_detail(forecast, truth, lead, lats, min_k=20): |
| """% of the truth's small-scale variation (zonal wavenumbers >= min_k) that the forecast has, around `lead`.""" |
| result = spectrum_ratio(forecast, truth, lead, lats) |
| if result is None: |
| return np.nan |
| k, ratio = result |
| return 100 * float(ratio[k >= min_k].mean()) |
|
|
