Exposure of people to changing extreme precipitation

import pandas as pd
import numpy as np
import plotly.graph_objects as go
from plotly.subplots import make_subplots
from scipy.stats import genextreme
from scipy.spatial import cKDTree
import rioxarray as rxr
import geopandas as gpd
import regionmask
import os
import re
from pathlib import Path

# --- Hazard settings (identical to changing_hazard.ipynb) ---
admin_id       = 'ITF2'
baseline_start = 1976
baseline_end   = 2005
return_periods = np.array([2, 5, 10, 20, 50, 100])

future_periods = {
    'Near-term (2021-2050)': (2021, 2050),
    'Mid-term (2041-2070)':  (2041, 2070),
    'Long-term (2071-2100)': (2071, 2100),
}
future_period_names = list(future_periods.keys())

rcps        = ['rcp_2_6', 'rcp_4_5', 'rcp_8_5']
rcp_labels  = {'rcp_2_6': 'RCP 2.6', 'rcp_4_5': 'RCP 4.5', 'rcp_8_5': 'RCP 8.5'}
rcp_colors  = {'rcp_2_6': 'steelblue', 'rcp_4_5': 'darkorange', 'rcp_8_5': 'crimson'}

# --- Exposure-specific settings ---
rp_exposure      = 20    # return period to track [years]
factor_threshold = 1.5   # likelihood factor defining 'exposed'
window_len       = 30    # sliding window length [years]

# --- Paths ---
#workdir     = Path('/home/nejk/code/extreme_precip_exposure')
#os.chdir(workdir)
main_dir    = Path("../")
data_dir    = main_dir / 'data'
hist_file   = data_dir / admin_id / 'extreme_precip_exposure' / f'{admin_id}_precip_hist_amax_all_points.csv'
proj_file   = data_dir / admin_id / 'extreme_precip_exposure' / f'{admin_id}_precip_proj_amax_all_points.csv'
ghs_wup_dir = data_dir / 'population' / 'GHS_WUP_POP'
nuts_shp    = data_dir / 'regions' / 'NUTS_RG_20M_2024_4326' / 'NUTS_RG_20M_2024_4326.shp'
pop_csv     = data_dir / admin_id / 'ghs_population' / f'population_ghs_wup_{admin_id}.csv'

print(f'Region           : {admin_id} (Molise, Italy)')
print(f'Event analysed   : {rp_exposure}-year return period')
print(f'Exposure threshold: factor \u2265 {factor_threshold}\u00d7')

# Region boundary
nuts_gdf = gpd.read_file(nuts_shp)
sel_gdf  = nuts_gdf[nuts_gdf['NUTS_ID'] == admin_id]
lon_min, lat_min, lon_max, lat_max = sel_gdf.geometry.total_bounds
print(f'Region boundary loaded: {sel_gdf["NUTS_NAME"].values[0]} ({admin_id})')

# Load historical annual maxima
hist_df = pd.read_csv(hist_file)
hist_df = hist_df.dropna(subset=['rx1day'])
hist_df = hist_df.drop_duplicates(subset=['member_id', 'lat', 'lon', 'year'])
print(f'Historical data: {hist_df["member_id"].nunique()} models, '
      f'{hist_df[["lat","lon"]].drop_duplicates().shape[0]} grid points, '
      f'{hist_df["year"].min()}\u2013{hist_df["year"].max()}')

# Fit GEV to baseline period
baseline_df = hist_df[(hist_df['year'] >= baseline_start) & (hist_df['year'] <= baseline_end)]
gev_records = []
for (member_id, lat, lon), grp in baseline_df.groupby(['member_id', 'lat', 'lon']):
    amax = grp['rx1day'].dropna().values
    if len(amax) < 10:
        continue
    try:
        shape, loc, scale = genextreme.fit(amax)
        gev_records.append({'member_id': member_id, 'lat': lat, 'lon': lon,
                            'shape': shape, 'loc': loc, 'scale': scale})
    except Exception:
        continue
gev_df = pd.DataFrame(gev_records)

# Compute return levels for every (model, grid point, return period)
rl_records = []
for _, row in gev_df.iterrows():
    hist_sub = hist_df[
        (hist_df['member_id'] == row['member_id']) &
        (hist_df['lat']       == row['lat'])       &
        (hist_df['lon']       == row['lon'])
    ]
    for rp in return_periods:
        rl = float(genextreme.ppf(1 - 1/rp, row['shape'], row['loc'], row['scale']))
        rc = int((hist_sub['rx1day'] >= rl).sum())
        rl_records.append({
            'member_id': row['member_id'], 'lat': row['lat'], 'lon': row['lon'],
            'return_period': rp, 'return_level': rl, 'hist_count': rc,
        })
return_levels_df = pd.DataFrame(rl_records)

n_models = gev_df['member_id'].nunique()
n_points = gev_df[['lat', 'lon']].drop_duplicates().shape[0]
print(f'GEV fits: {n_models} models \u00d7 {n_points} grid points = {len(gev_df)} combinations')

# Hazard grid and KD-tree
haz_grid   = return_levels_df[['lat', 'lon']].drop_duplicates().reset_index(drop=True)
haz_coords = haz_grid[['lat', 'lon']].values          # shape (N_haz, 2)
tree       = cKDTree(haz_coords)
print(f'Hazard grid: {len(haz_coords)} cells')

# Infer grid spacing from hazard points
lats_haz = sorted(haz_grid['lat'].unique())
lons_haz = sorted(haz_grid['lon'].unique())
dlat = float(np.median(np.diff(lats_haz))) if len(lats_haz) > 1 else 0.11
dlon = float(np.median(np.diff(lons_haz))) if len(lons_haz) > 1 else 0.11
print(f'Grid spacing  : \u0394lat={dlat:.4f}\u00b0, \u0394lon={dlon:.4f}\u00b0')

# Region mask for GHS rasters
admin_mask = regionmask.from_geopandas(sel_gdf, names='NUTS_ID')

# Load and regrid all GHS-WUP-POP rasters
print('\nLoading and regridding GHS-WUP-POP rasters ...')
tif_paths = sorted(ghs_wup_dir.glob('*/GHS_WUP_POP_E*_GLOBE_R2025A_4326_30ss_V1_0.tif'))
if not tif_paths:
    raise FileNotFoundError(f'No GHS-WUP-POP rasters found under {ghs_wup_dir}. \nPlease download data following the instructions in the How-To-Guide on Population data.')

pop_haz_records = []
for tif_path in tif_paths:
    m = re.search(r'E(\d{4})', tif_path.name)
    if not m:
        continue
    year = int(m.group(1))

    da = rxr.open_rasterio(tif_path, masked=True).squeeze('band', drop=True)
    da_clip = da.sel(
        x=slice(lon_min - 0.1, lon_max + 0.1),
        y=slice(lat_max + 0.1, lat_min - 0.1),   # y is descending in rasters
    )

    # Boolean region mask on the GHS grid
    raw_mask  = admin_mask.mask(da_clip.x, da_clip.y)
    in_region = ~np.isnan(raw_mask.values)         # True where inside the region

    lons_2d, lats_2d = np.meshgrid(da_clip.x.values, da_clip.y.values)
    pop_vals = np.nan_to_num(da_clip.values, nan=0.0)
    pop_vals = np.where(pop_vals < 0, 0.0, pop_vals)  # remove fill values

    pop_lat_in = lats_2d[in_region]
    pop_lon_in = lons_2d[in_region]
    pop_val_in = pop_vals[in_region]

    # Assign each GHS cell to nearest hazard cell and sum
    _, idx = tree.query(np.column_stack([pop_lat_in, pop_lon_in]))
    pop_per_haz = np.zeros(len(haz_coords))
    np.add.at(pop_per_haz, idx, pop_val_in)

    for i in range(len(haz_coords)):
        pop_haz_records.append({
            'year': year,
            'lat':  haz_coords[i, 0],
            'lon':  haz_coords[i, 1],
            'population': pop_per_haz[i],
        })
    print(f'  {year}: {pop_per_haz.sum():>12,.0f} people')

pop_haz_df = pd.DataFrame(pop_haz_records)
print(f'\nDone. {pop_haz_df["year"].nunique()} years \u00d7 {len(haz_coords)} hazard cells = {len(pop_haz_df):,} rows')

# Load projections
print('Loading projection data ...')
proj_df = pd.read_csv(proj_file)
proj_df = proj_df.dropna(subset=['rx1day'])
proj_df = proj_df.drop_duplicates(subset=['member_id', 'lat', 'lon', 'year', 'rcp'])
print(f'  {proj_df["member_id"].nunique()} models, '
      f'{proj_df[["lat","lon"]].drop_duplicates().shape[0]} grid points, '
      f'{proj_df["year"].min()}\u2013{proj_df["year"].max()}')

# Historical threshold for the chosen return period
threshold_rp = (
    return_levels_df[return_levels_df['return_period'] == rp_exposure]
    [['member_id', 'lat', 'lon', 'return_level']]
    .rename(columns={'return_level': 'hist_threshold'})
)

# Compute future exceedance counts and likelihood factors (same approach as changing_hazard)
future_records = []
for period_name, (pstart, pend) in future_periods.items():
    period_proj = proj_df[(proj_df['year'] >= pstart) & (proj_df['year'] <= pend)]
    wlen = pend - pstart + 1
    for rcp in rcps:
        rcp_proj = period_proj[period_proj['rcp'] == rcp]
        if rcp_proj.empty:
            continue
        merged = rcp_proj.merge(threshold_rp, on=['member_id', 'lat', 'lon'], how='inner')
        merged['exceeds'] = merged['rx1day'] >= merged['hist_threshold']
        counts = (
            merged.groupby(['member_id', 'lat', 'lon'])['exceeds']
            .sum()
            .reset_index(name='proj_count')
        )
        counts['period'] = period_name
        counts['rcp']    = rcp
        counts['factor'] = counts['proj_count'] * rp_exposure / wlen
        future_records.append(counts)

future_rp_df = pd.concat(future_records, ignore_index=True)
print(f'Future hazard table: {len(future_rp_df):,} rows')

# Pre-build interpolated population for arbitrary years (for reuse in timeseries section)
pop_haz_pivot = pop_haz_df.pivot_table(index=['lat', 'lon'], columns='year', values='population')
all_years_ghs = sorted(pop_haz_pivot.columns)
all_years_interp = list(range(all_years_ghs[0], all_years_ghs[-1] + 1))
pop_haz_pivot_full = (
    pop_haz_pivot
    .reindex(columns=sorted(set(all_years_ghs) | set(all_years_interp)))
    .interpolate(axis=1, method='linear', limit_direction='both')
    [all_years_interp]
)

def pop_at_year(year):
    """Return population Series indexed by (lat, lon) for the given year."""
    y = int(np.clip(year, all_years_interp[0], all_years_interp[-1]))
    return pop_haz_pivot_full[y]

# Compute exposure for the three future periods
exposure_records = []
for period_name, (pstart, pend) in future_periods.items():
    pop_mid = pop_at_year((pstart + pend) // 2)
    total_pop = pop_mid.sum()

    for rcp in rcps:
        cell_factor = (
            future_rp_df[
                (future_rp_df['period'] == period_name) &
                (future_rp_df['rcp']    == rcp)
            ]
            .groupby(['lat', 'lon'])['factor'].mean()
        )
        exposed_idx = cell_factor[cell_factor >= factor_threshold].index
        exposed_pop = pop_mid.reindex(exposed_idx).fillna(0).sum()

        exposure_records.append({
            'period':          period_name,
            'rcp':             rcp,
            'mid_year':        (pstart + pend) // 2,
            'total_pop':       total_pop,
            'exposed_pop':     exposed_pop,
            'pct_exposed':     100 * exposed_pop / total_pop if total_pop > 0 else 0,
            'n_exposed_cells': len(exposed_idx),
            'n_total_cells':   len(cell_factor),
        })

exposure_df = pd.DataFrame(exposure_records)
exposure_df
#print(f'Exposure results (factor \u2265 {factor_threshold}\u00d7 for {rp_exposure}-year event):')
#print(exposure_df[['period', 'rcp', 'mid_year', 'total_pop', 'exposed_pop', 'pct_exposed',
#                    'n_exposed_cells', 'n_total_cells']]
#      .to_string(index=False, float_format='{:,.0f}'.format))

# plot absolute exposed population and percentage, faceted by future period and coloured by RCP (overlaid on total population in grey)
fig = make_subplots(
    rows=1, cols=3,
    subplot_titles=future_period_names,
    shared_yaxes=True,
    horizontal_spacing=0.06,
)

for col_idx, period_name in enumerate(future_period_names, start=1):
    pdata = exposure_df[exposure_df['period'] == period_name].set_index('rcp')
    x_labels = [rcp_labels[r] for r in rcps]

    # Grey bar: total regional population
    fig.add_trace(go.Bar(
        x=x_labels,
        y=pdata.loc[rcps, 'total_pop'].values,
        name='Total regional population',
        marker_color='lightgrey',
        legendgroup='total',
        showlegend=(col_idx == 1),
        hovertemplate='Total pop: %{y:,.0f}<extra></extra>',
    ), row=1, col=col_idx)

    # Coloured bar: exposed population (overlaid)
    fig.add_trace(go.Bar(
        x=x_labels,
        y=pdata.loc[rcps, 'exposed_pop'].values,
        name='Exposed population',
        marker_color=[rcp_colors[r] for r in rcps],
        legendgroup='exposed',
        showlegend=(col_idx == 1),
        hovertemplate='Exposed: %{y:,.0f}<extra></extra>',
    ), row=1, col=col_idx)

    # Annotate % exposed
    for i, rcp in enumerate(rcps):
        pct = pdata.loc[rcp, 'pct_exposed']
        fig.add_annotation(
            x=x_labels[i],
            y=pdata.loc[rcp, 'exposed_pop'],
            text=f'{pct:.0f}%',
            showarrow=False,
            yshift=10,
            font=dict(size=10),
            row=1, col=col_idx,
        )

fig.update_yaxes(title_text='Population', row=1, col=1)
fig.update_layout(
    title=(
        f'Population exposed to a \u2265{factor_threshold}\u00d7 increase in the '
        f'{rp_exposure}-year event<br>'
        f'<sup>Coloured bars = exposed; grey bars = total regional population at period mid-year. '
        f'Percentages annotated above bars.</sup>'
    ),
    barmode='overlay',
    legend_title='',
    autosize=True, height=500,
)
fig.show(config={'displayModeBar': True, 'responsive': True})

# Pre-merge projection data with thresholds (avoids repeating inside the loop)
proj_merged = proj_df.merge(threshold_rp, on=['member_id', 'lat', 'lon'], how='inner')
proj_merged['exceeds'] = proj_merged['rx1day'] >= proj_merged['hist_threshold']

t_first = int(proj_df['year'].min()) + window_len - 1  # first window end year
t_last  = int(proj_df['year'].max())
print(f'Sliding window: {window_len}-year trailing window, {t_first}\u2013{t_last}')

ts_records = []
for t in range(t_first, t_last + 1):
    window_data = proj_merged[
        (proj_merged['year'] >= t - window_len + 1) & (proj_merged['year'] <= t)
    ]
    pop_t       = pop_at_year(t)
    total_pop_t = pop_t.sum()

    for rcp in rcps:
        rcp_data = window_data[window_data['rcp'] == rcp]
        if rcp_data.empty:
            continue

        counts = (
            rcp_data.groupby(['member_id', 'lat', 'lon'])['exceeds']
            .sum()
            .reset_index(name='count')
        )
        counts['factor'] = counts['count'] * rp_exposure / window_len

        # Per-model exposure — threshold applied per model, then aggregate
        model_exposed = []
        for _mid, grp in counts.groupby('member_id'):
            cell_f  = grp.set_index(['lat', 'lon'])['factor']
            exp_idx = cell_f[cell_f >= factor_threshold].index
            exp_pop = pop_t.reindex(exp_idx).fillna(0).sum()
            model_exposed.append(float(exp_pop))

        # All summary statistics derived from the same per-model values so
        # the mean always lies within the p10–p90 band
        exposed_pop = float(np.mean(model_exposed))
        p10 = float(np.percentile(model_exposed, 10))
        p90 = float(np.percentile(model_exposed, 90))

        # Number of exposed cells: use the multi-model mean factor for a
        # representative spatial summary (not used for the timeseries lines)
        cell_mean   = counts.groupby(['lat', 'lon'])['factor'].mean()
        exposed_idx = cell_mean[cell_mean >= factor_threshold].index

        ts_records.append({
            'year':            t,
            'rcp':             rcp,
            'exposed_pop':     exposed_pop,
            'p10_exposed_pop': p10,
            'p90_exposed_pop': p90,
            'total_pop':       float(total_pop_t),
            'pct_exposed':     100 * exposed_pop / float(total_pop_t) if total_pop_t > 0 else 0,
            'pct_p10':         100 * p10 / float(total_pop_t) if total_pop_t > 0 else 0,
            'pct_p90':         100 * p90 / float(total_pop_t) if total_pop_t > 0 else 0,
            'n_exposed_cells': len(exposed_idx),
            'n_total_cells':   len(cell_mean),
        })

ts_df = pd.DataFrame(ts_records)
print(f'Done. {len(ts_df):,} rows ({ts_df["year"].nunique()} years \u00d7 {len(rcps)} RCPs)')

# Semi-transparent band colours matching rcp_colors (steelblue / darkorange / crimson)
rcp_band_colors = {
    'rcp_2_6': 'rgba(70,130,180,0.15)',
    'rcp_4_5': 'rgba(255,140,0,0.15)',
    'rcp_8_5': 'rgba(220,20,60,0.15)',
}

fig = make_subplots(
    rows=2, cols=1,
    shared_xaxes=True,
    subplot_titles=[
        f'People exposed to \u2265{factor_threshold}\u00d7 increase in {rp_exposure}-year event',
        '% of regional population exposed',
    ],
    vertical_spacing=0.12,
)

for rcp in rcps:
    rcp_ts     = ts_df[ts_df['rcp'] == rcp].sort_values('year')
    color      = rcp_colors[rcp]
    band_color = rcp_band_colors[rcp]
    label      = rcp_labels[rcp]

    # ---- absolute exposure (row 1) ----
    # 10th–90th percentile uncertainty band
    fig.add_trace(go.Scatter(
        x=rcp_ts['year'], y=rcp_ts['p10_exposed_pop'],
        mode='lines', line=dict(width=0), showlegend=False,
        hoverinfo='skip', legendgroup=rcp,
    ), row=1, col=1)
    fig.add_trace(go.Scatter(
        x=rcp_ts['year'], y=rcp_ts['p90_exposed_pop'],
        mode='lines', line=dict(width=0),
        fill='tonexty', fillcolor=band_color,
        name=f'{label} (10th–90th pct)', legendgroup=rcp, showlegend=True,
        hovertemplate=f'{label} p90<br>Year: %{{x}}<br>%{{y:,.0f}}<extra></extra>',
    ), row=1, col=1)
    # Ensemble mean line
    fig.add_trace(go.Scatter(
        x=rcp_ts['year'], y=rcp_ts['exposed_pop'],
        mode='lines', name=label, legendgroup=rcp, showlegend=False,
        line=dict(color=color, width=2.5),
        hovertemplate=f'{label} mean<br>Year: %{{x}}<br>%{{y:,.0f}}<extra></extra>',
    ), row=1, col=1)

    # ---- % exposed (row 2) ----
    fig.add_trace(go.Scatter(
        x=rcp_ts['year'], y=rcp_ts['pct_p10'],
        mode='lines', line=dict(width=0), showlegend=False,
        hoverinfo='skip', legendgroup=rcp,
    ), row=2, col=1)
    fig.add_trace(go.Scatter(
        x=rcp_ts['year'], y=rcp_ts['pct_p90'],
        mode='lines', line=dict(width=0),
        fill='tonexty', fillcolor=band_color,
        name=f'{label} (10th–90th pct)', legendgroup=rcp, showlegend=False,
        hovertemplate=f'{label} p90<br>Year: %{{x}}<br>%{{y:.1f}}%<extra></extra>',
    ), row=2, col=1)
    fig.add_trace(go.Scatter(
        x=rcp_ts['year'], y=rcp_ts['pct_exposed'],
        mode='lines', name=label, legendgroup=rcp, showlegend=False,
        line=dict(color=color, width=2.5),
        hovertemplate=f'{label} mean<br>Year: %{{x}}<br>%{{y:.1f}}%<extra></extra>',
    ), row=2, col=1)

# Total regional population (dotted grey reference)
pop_total_ts = ts_df[ts_df['rcp'] == rcps[0]].sort_values('year')
fig.add_trace(go.Scatter(
    x=pop_total_ts['year'], y=pop_total_ts['total_pop'],
    mode='lines', name='Total regional population',
    line=dict(color='lightgrey', width=1.5, dash='dot'),
    hovertemplate='Total pop (year %{x}): %{y:,.0f}<extra></extra>',
), row=1, col=1)

# Mark future period mid-years
for period_name, (pstart, pend) in future_periods.items():
    mid         = (pstart + pend) // 2
    label_short = period_name.split(' ')[0]
    for row in [1, 2]:
        fig.add_vline(
            x=mid, line_dash='dot', line_color='darkgrey', line_width=1,
            row=row, col=1,
        )
    fig.add_annotation(
        x=mid, y=1.04, text=label_short,
        showarrow=False, font=dict(size=9, color='darkgrey'),
        xref='x', yref='paper',
    )

fig.update_yaxes(title_text='People', showgrid=True, gridcolor='lightgrey', row=1, col=1)
fig.update_yaxes(title_text='%',      showgrid=True, gridcolor='lightgrey', row=2, col=1)
fig.update_xaxes(
    title_text='Window end year (trailing 30-year window)',
    showgrid=True, gridcolor='lightgrey',
    row=2, col=1,
)
fig.update_layout(
    title=dict(
        text=(
            f'Moving-window timeseries: population exposed to extreme precipitation increase<br>'
            f'<sup>30-year trailing window | {rp_exposure}-year event | '
            f'factor \u2265 {factor_threshold}\u00d7 | solid line = multi-model mean | '
            f'shaded band = 10th\u201390th percentile across ensemble members</sup>'
        ),
        font=dict(size=13),
    ),
    legend=dict(title='RCP scenario', x=1.01, y=0.99, xanchor='left', yanchor='top'),
    autosize=True, height=650,
)
fig.show(config={'displayModeBar': True, 'responsive': True})