Skip to content

Commit

Permalink
clean up
Browse files Browse the repository at this point in the history
Also: xarray 2024.5 has a bug in the interpolation function.
We need to use an older version for now.
  • Loading branch information
aaschwanden committed Jun 7, 2024
1 parent 90f97b1 commit e5b3b8a
Show file tree
Hide file tree
Showing 5 changed files with 269 additions and 160 deletions.
2 changes: 1 addition & 1 deletion .pylintrc
Original file line number Diff line number Diff line change
Expand Up @@ -3,4 +3,4 @@ jobs=4 #number of processes to use
[BASIC]
good-names=nameOfYourProject #names to be considered ok
[pre-commit-hook]
disable=bare-except, import-outside-toplevel, too-many-return-statements, line-too-long, invalid-name, too-many-arguments, dangerous-default-value, too-many-locals
disable=bare-except, import-outside-toplevel, too-many-return-statements, line-too-long, invalid-name, too-many-arguments, dangerous-default-value, too-many-locals, duplicate-code
98 changes: 89 additions & 9 deletions glacier_flow_tools/profiles.py
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,85 @@
matplotlib.use("agg")


def extract_profile(
profile,
obs_ds: xr.Dataset,
sim_ds: xr.Dataset,
result_dir: Union[str, Path, None] = None,
obs_normal_var: str = "obs_v_normal",
obs_normal_error_var: str = "obs_v_err_normal",
obs_normal_component_vars: dict = {"x": "vx", "y": "vy"},
obs_normal_component_error_vars: dict = {"x": "vx_err", "y": "vy_err"},
sim_normal_var: str = "sim_v_normal",
sim_normal_component_vars: dict = {"x": "uvelsurf", "y": "vvelsurf"},
compute_profile_normal: bool = True,
stats: List[str] = ["rmsd", "pearson_r"],
stats_kwargs: Dict = {},
) -> xr.Dataset:
"""
Extract and process profiles from observation and simulation datasets.
This function extracts profiles from the given observation and simulation datasets, processes them, and saves the processed profile as a netCDF file. It also merges the processed profile with a given GeoDataFrame on intersection.
Parameters
----------
profile : dict
The profile to be processed.
obs_ds : xr.Dataset
The observation dataset.
sim_ds : xr.Dataset
The simulation dataset.
result_dir : str or Path, optional
The directory where the result netCDF file will be saved, by default None, no file is saved.
obs_normal_var : str, optional
The variable name for the normal component of the observations, by default 'obs_v_normal'.
obs_normal_error_var : str, optional
The variable name for the error of the normal component of the observations, by default 'obs_v_err_normal'.
obs_normal_component_vars : dict, optional
The variable names for the components of the normal component of the observations, by default {"x": "vx", "y": "vy"}.
obs_normal_component_error_vars : dict, optional
The variable names for the errors of the components of the normal component of the observations, by default {"x": "vx_err", "y": "vy_err"}.
sim_normal_var : str, optional
The variable name for the normal component of the simulations, by default 'sim_v_normal'.
sim_normal_component_vars : dict, optional
The variable names for the components of the normal component of the simulations, by default {"x": "uvelsurf", "y": "vvelsurf"}.
compute_profile_normal : bool, optional
Whether to compute the normal component of the profile, by default True.
stats : list of str, optional
The statistics to be computed for the profile, by default ["rmsd", "pearson_r"].
stats_kwargs : dict, optional
Additional keyword arguments to pass to the function for computing the statistics, by default {}.
Returns
-------
tuple of xr.Dataset and gp.GeoDataFrame
The processed profile as an xr.Dataset and the merged GeoDataFrame.
Examples
--------
>>> extract_profiles(profile, profiles_df, obs_ds, sim_ds, stats=["rmsd", "pearson_r"], result_dir='.', obs_normal_var='obs_v_normal', obs_normal_error_var='obs_v_err_normal', obs_normal_component_vars={"x": "vx", "y": "vy"}, obs_normal_component_error_vars={"x": "vx_err", "y": "vy_err"}, sim_normal_var='sim_v_normal', sim_normal_component_vars={"x": "uvelsurf", "y": "vvelsurf"}, compute_profile_normal=True, stats_kwargs={})
"""
os_profile = process_profile(
profile,
obs_ds=obs_ds,
sim_ds=sim_ds,
compute_profile_normal=compute_profile_normal,
obs_normal_var=obs_normal_var,
obs_normal_error_var=obs_normal_error_var,
obs_normal_component_vars=obs_normal_component_vars,
obs_normal_component_error_vars=obs_normal_component_error_vars,
sim_normal_var=sim_normal_var,
sim_normal_component_vars=sim_normal_component_vars,
stats=stats,
stats_kwargs=stats_kwargs,
)

if result_dir is not None:
os_file = Path(result_dir) / f"""{profile["profile_name"]}_profile.nc"""
os_profile.to_netcdf(os_file, engine="h5netcdf")
return os_profile


def plot_profile(
ds: xr.Dataset,
result_dir: Path,
Expand Down Expand Up @@ -448,7 +527,8 @@ def process_profile(
This function extracts profiles from the observed and simulated datasets along the given profile, merges them, and calculates the specified statistics.
"""

x, y = map(np.asarray, profile["geometry"].xy)
coords = get_coordinates(profile["geometry"])
x, y = coords[:, 0], coords[:, 1]
profile_name = profile["profile_name"]
profile_id = profile["profile_id"]
stats_kwargs_copy = stats_kwargs.copy()
Expand Down Expand Up @@ -974,7 +1054,6 @@ def extract_profile(

ds = xr.merge(das)
ds["profile_id"] = [profile_id]

if compute_profile_normal:
a = [(v in ds.data_vars) for v in normal_component_vars.values()]
if np.all(np.array(a)):
Expand Down Expand Up @@ -1006,7 +1085,7 @@ def plot(
sigma: float = 1,
title: Union[str, None] = None,
obs_var: str = "v",
obs_error_var: str = "v_err",
obs_error_var: Union[str, None] = "v_err",
sim_var: str = "velsurf_mag",
palette: str = "Paired",
obs_kwargs: dict = {"color": "0", "lw": 0.75, "marker": "o", "ms": 1.5},
Expand Down Expand Up @@ -1069,12 +1148,13 @@ def plot(
n_exps = self._obj["exp_id"].size
fig = plt.figure(figsize=figsize)
ax = fig.add_subplot(111)
ax.fill_between(
self._obj["profile_axis"],
self._obj[obs_var] - sigma * self._obj[obs_error_var],
self._obj[obs_var] + sigma * self._obj[obs_error_var],
**obs_error_kwargs,
)
if obs_error_var is not None:
ax.fill_between(
self._obj["profile_axis"],
self._obj[obs_var] - sigma * self._obj[obs_error_var],
self._obj[obs_var] + sigma * self._obj[obs_error_var],
**obs_error_kwargs,
)
ax.plot(
self._obj["profile_axis"],
self._obj[obs_var],
Expand Down
109 changes: 109 additions & 0 deletions notebooks/Ice Thickness Datasets.ipynb
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,115 @@
"id": "968ac4c2-94b2-4f5e-9e48-6832f2fa3097",
"metadata": {},
"outputs": [],
"source": [
"import pylab as plt\n",
"import xarray as xr\n",
"import geopandas as gp\n",
"from functools import partial"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "c4fa6ce1-c349-41aa-ac60-1796578d7eba",
"metadata": {},
"outputs": [],
"source": [
"from glacier_flow_tools.profiles import process_profile\n",
"from glacier_flow_tools.utils import preprocess_nc"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "6d03506c-625c-47cd-9379-a12f6e3abcde",
"metadata": {},
"outputs": [],
"source": [
"profile_resolution = 500 # m"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "1b80a187-b497-42f1-8021-b50925855801",
"metadata": {},
"outputs": [],
"source": [
"# process_profile will use non-interactive \"Agg\", so we switch back to default\n",
"plt.switch_backend('module://matplotlib_inline.backend_inline')"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "7c6d75d7-4da1-4fc6-b037-c92e0b879218",
"metadata": {},
"outputs": [],
"source": [
"itslive_ds = xr.open_dataset(\"data/its_live_jak.nc\")\n",
"exp_ds = xr.open_mfdataset(\"../tests/data/velsurf_mag_gris_g9000m_id_*_0_50.nc\",\n",
" preprocess=partial(preprocess_nc,\n",
" drop_dims=[\"z\", \"zb\", \"nv4\"],\n",
" drop_vars=[\"timestamp\", \"shelfbtemp\", \"effective_ice_surface_temp\", \"ice_surface_temp\", \"hardav\", \"nv4\"]), \n",
" concat_dim=\"exp_id\",\n",
" combine=\"nested\",\n",
" chunks=\"auto\",\n",
" engine=\"h5netcdf\",\n",
" decode_times=False,\n",
" parallel=True\n",
" )\n",
"profiles_gp = gp.read_file(\"../glacier_flow_tools/data/greenland-flux-gates-5.gpkg\").rename(columns={\"id\": \"profile_id\", \"name\": \"profile_name\"})\n",
"geom = profiles_gp.segmentize(profile_resolution)\n",
"profiles_gp = gp.GeoDataFrame(profiles_gp, geometry=geom)\n",
"profiles_gp = profiles_gp[[\"profile_id\", \"profile_name\", \"geometry\"]]\n",
"\n",
"profile_gp = [p for _, p in profiles_gp.iterrows()][1]"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "5469888c-01c1-4a1c-94f3-8e74db65a6f4",
"metadata": {},
"outputs": [],
"source": [
"fig, ax = plt.subplots()\n",
"itslive_ds[\"v\"].plot(cmap=\"speed_colorblind\", vmin=0, vmax=1500, ax=ax, cbar_kwargs={\"shrink\": 0.75})\n",
"profiles_gp.plot(ax=ax, color=\"k\")\n",
"ax.set_title(\"Jakobshavn Flux Gate\\nObserved Speed\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "43254926-ef9f-4223-8da9-f5d066240d09",
"metadata": {},
"outputs": [],
"source": [
"jak_profile = process_profile(profile_gp,\n",
" itslive_ds,\n",
" exp_ds,\n",
" compute_profile_normal=True,\n",
" stats_kwargs={\"obs_var\": \"v\", \"sim_var\": \"velsurf_mag\"})"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "f6bec873-8faa-48d2-8293-f56ca303ff7e",
"metadata": {},
"outputs": [],
"source": [
"jak_profile.profiles.plot(obs_error_var=None)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "3c430400-68ac-4862-89af-d0985e52f3c1",
"metadata": {},
"outputs": [],
"source": []
}
],
Expand Down
Loading

0 comments on commit e5b3b8a

Please sign in to comment.