<jupyter_start><jupyter_text>Mapping aerobic habitat in trait space<jupyter_code>%load_ext autoreload
%autoreload 2
import cartopy
import cartopy.crs as ccrs
import cmocean
import data_collections as dc
import funnel
import intake
import matplotlib.pyplot as plt
import numpy as np
import util
import xarray as xr
from matplotlib import colors
sub_spec = dict(
    name='drift-corrected',
    experiment=['20C', 'RCP85'],
    member_id=dc.ocean_bgc_member_ids,
)
catalog = funnel.to_intake_esm(agg_member_id=True).search(**sub_spec)
catalog
cat = catalog.search(variable='vol_habitat_trait_spc')
cat.df.head()
try:
    cluster
    client
except:
    cluster, client = util.get_ClusterClient(memory='8GB')
    cluster.scale(6)
client
cat = catalog.search(variable='vol_habitat_trait_spc')
dsets = cat.to_dataset_dict(zarr_kwargs={'use_cftime': True})
keys = list(dsets.keys())

exp_keys = ['20C.ocn.pop.h.drift-corrected', 'RCP85.ocn.pop.h.drift-corrected']
ds = xr.concat([dsets[k] for k in exp_keys], dim='time', coords='minimal').compute()

ds.vol_habitat_trait_spc.data *= 1e-6 * 1e-6  # convert from cm^3 to 10^6 m^6
ds.vol_habitat_trait_spc.attrs['units'] = '10$^6$ m$^3$'
ds.vol_habitat_trait_spc.attrs['long_name'] = 'Aerobic habitat volume'

ds
yrfrac = util.year_frac(ds.time)
tndx_ref = np.where(yrfrac < 1966)[0]
tndx_2100 = np.where(yrfrac > 2080)[0]
with xr.set_options(keep_attrs=True):
    vol_hab_ref = ds.vol_habitat_trait_spc.isel(time=tndx_ref).mean(['member_id', 'time'])
    vol_hab_ref_std = ds.vol_habitat_trait_spc.isel(time=tndx_ref).std(['member_id', 'time'])

    vol_hab_2100 = ds.vol_habitat_trait_spc.isel(time=tndx_2100).mean(['member_id', 'time'])
    vol_hab_2100_std = ds.vol_habitat_trait_spc.isel(time=tndx_2100).std(['member_id', 'time'])

    # vol_hab_ref_percent = 100.0 * vol_hab_ref / total_volume
    # vol_hab_ref_percent.attrs['long_name'] = 'Ocean volume'
    # vol_hab_ref_percent.attrs['units'] = '%'

vol_hab_2100
vol_hab_ref.plot.contourf(levels=30)
vol_hab_2100.plot()
(vol_hab_2100 - vol_hab_ref).plot.contourf(levels=30)
change_percent = 100.0 * (vol_hab_2100 - vol_hab_ref) / vol_hab_ref.where(vol_hab_ref > 1e5)

mx = np.floor(change_percent.max())
mn = -24.0
divnorm = colors.TwoSlopeNorm(vmin=mn, vcenter=0, vmax=mx)
dx = 1.0

change_percent.plot.contourf(norm=divnorm, levels=np.arange(mn, mx + dx, dx), extend='both');
change_percent = 100.0 * (vol_hab_2100 - vol_hab_ref) / vol_hab_ref.where(vol_hab_ref > 1e5)
mx = np.floor(change_percent.max())
mn = -24.0
dx = 1.0

cf = plt.contourf(
    change_percent.Eo,
    1.0 / change_percent.Ac,
    change_percent,
    levels=np.arange(mn, mx + dx, dx),
    norm=colors.TwoSlopeNorm(vmin=mn, vcenter=0, vmax=mx),
    extend='both',
    cmap=cmocean.cm.curl_r,
)
plt.ylim(0, 25)
plt.colorbar(cf)<jupyter_output><empty_output>
StarcoderdataJupyterStructuredCleanDedup