{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "423cdddd",
   "metadata": {},
   "source": [
    "# Adding new data variables\n",
    "or more verbosely:\n",
    "# Scenario 1: Adding new data variables without altering the existing data structure\n",
    "\n",
    "If you want to write your own function or method for an `xsnowDataset`, simply code your function in a standard python module and decide whether you want to keep this module private to yourself or host it in a public repository.\n",
    "\n",
    "Here is a cheat sheet for the steps you have to take. You will find more detailed explanations and a demonstration further below:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6075c2c",
   "metadata": {},
   "source": [
    "```{admonition}  Recipe\n",
    ":class: note\n",
    "\n",
    " 1. Import: `from xsnow import xsnowDataset, register_xsnowDataset_method, registry`\n",
    " 2. Definition of function/method:\n",
    "            @register_xsnowDataset_method\n",
    "            def my_method(xs: xsnowDataset) -> Any:\n",
    " 3. (optional) If returning an xsnowDataset, `._rewrap()` the Dataset back into the correct class and make sure you created metadata for the newly added DataArray.\n",
    "\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74ea574c",
   "metadata": {},
   "source": [
    "Regarding 1.)\n",
    "\n",
    " * If you use or return an `xsnowDataset` you have to import the class; alternatively, you can `import xsnow` and then use the class as `xsnow.xsnowDataset`.\n",
    " * (optional) If you want to create a *method* (and not only a *function*), import `register_xsnowDataset_method`.\n",
    " * (optional) If you add a variable that may exist in the xsnow registry of known variables with their names, units, etc, then import it.\n",
    "\n",
    " Regarding 2.)\n",
    "\n",
    " * Define your function.\n",
    " * (optional) If you want to create a *method*, add `@register_xsnowDataset_method` in the line above your function signature and ensure that the first argument is an `xsnowDataset`.\n",
    "\n",
    " Regarding 3.)\n",
    "\n",
    " * (optional) If you want to return an `xsnowDataset`, ensure that the class of the returned object is the same class as the input object. To be compatible with *Scenario-2-extensions*, please do that with the pattern  \n",
    "        `xs_out = xs_in._rewrap(xs_modified)`\n",
    " * (optional) If you want to return an `xsnowDataset` and added a new variable, please ensure that the new variable contains metadata, such as `long_name`, `standard_name`, and `units`.\n",
    "\n",
    "The following code block shows an example from the [Critical crack length extension](../../api/_generated/xsnow.extensions.critical_crack_length), which applies all of these three steps. Feel free to use it as a template!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14b9decb",
   "metadata": {
    "tags": [
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "\n",
    "import numpy as np\n",
    "import xarray as xr\n",
    "\n",
    "from xsnow import xsnowDataset, register_xsnowDataset_method, registry\n",
    "from xsnow.extensions.critical_crack_length \\\n",
    "    import compute_rho_slab, compute_rc_flat_richter2019\n",
    "\n",
    "@register_xsnowDataset_method\n",
    "def compute_critical_crack_length(xs: xsnowDataset) -> xsnowDataset:\n",
    "    \"\"\"\n",
    "    Demo for customizing a new xsnowDataset method\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    xs : xsnowDataset\n",
    "        The first object always has to be an xsnowDataset if this function is\n",
    "        supposed to become a method!\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    new_xs : xsnowDataset\n",
    "        A new xsnowDataset with a new data variable `rc_flat` [m].\n",
    "\n",
    "    Examples\n",
    "    --------\n",
    "    import xsnow\n",
    "    import mymodule  # wherever your custom method lives\n",
    "    xs = xsnow.single_profile_timeseries()\n",
    "    xs_new = xs.compute_critical_crack_length()\n",
    "    \"\"\"\n",
    "\n",
    "    if xs.data is None:\n",
    "        raise ValueError(\"Cannot compute on an empty dataset.\")\n",
    "\n",
    "    required = ['density', 'grain_size', 'shear_strength', 'height']\n",
    "    for v in required:\n",
    "        if v not in xs.data:\n",
    "            raise ValueError(\n",
    "                f\"Missing required variable for rc_flat calculation: '{v}'\"\n",
    "            )\n",
    "\n",
    "    density_arr = xs.data['density'].to_numpy()\n",
    "    gsize_arr = xs.data['grain_size'].to_numpy()\n",
    "    shearstrength_arr = xs.data['shear_strength'].to_numpy()\n",
    "    height_arr = xs.data['height'].to_numpy()\n",
    "\n",
    "    # Derive per-layer thickness from cumulative height along 'layer'.\n",
    "    # layer=0 is bottom (height ~ 0), layer increases upward.\n",
    "    height_bottom = height_arr.shift(layer=1, fill_value=0)\n",
    "    thick_arr = height_arr - height_bottom  # [m]\n",
    "\n",
    "    # Compute density of slab above each layer\n",
    "    rho_slab_arr = compute_rho_slab(\n",
    "        density=density_arr,\n",
    "        thickness=thick_arr,\n",
    "        xp=np,\n",
    "    )\n",
    "\n",
    "    # Physics core: Richter et al. (2019) flat-terrain rc\n",
    "    rc_flat_arr = compute_rc_flat_richter2019(\n",
    "        rho_layer=density_arr,\n",
    "        grain_size_mm=gsize_arr,\n",
    "        shear_strength_kpa=shearstrength_arr,\n",
    "        rho_slab=rho_slab_arr,\n",
    "        xp=np,\n",
    "    )\n",
    "\n",
    "    # Wrap result back into xarray using same dims/coords as height\n",
    "    rc_flat_da = xr.DataArray(\n",
    "        rc_flat_arr,\n",
    "        dims=xs.data['height'].dims,\n",
    "        coords=xs.data['height'].coords,\n",
    "    )\n",
    "\n",
    "    # Attach registry metadata if available\n",
    "    if 'rc_flat' in registry:\n",
    "        reg_entry = registry['rc_flat']\n",
    "        rc_flat_da.attrs = {\n",
    "            'long_name': reg_entry.get('long_name', \n",
    "                                       'critical crack length (flat terrain)'),\n",
    "            'units': reg_entry.get('si_units', 'm'),\n",
    "            'standard_name': reg_entry.get('standard_name', 'rc_flat'),\n",
    "        }\n",
    "    else:\n",
    "        rc_flat_da.attrs = {\n",
    "            'long_name': 'critical crack length on flat terrain',\n",
    "            'units': 'm',\n",
    "            'standard_name': 'rc_flat',\n",
    "        }\n",
    "\n",
    "    ## create new xarray.Dataset and wrap back into correct class:\n",
    "    new_xs = xs.assign(rc_flat=rc_flat_da)\n",
    "    new_xs = xs._rewrap(new_xs)\n",
    "    return new_xs"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "xsnow-dev",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}