{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "932b31db-40a6-40dc-8cd9-e1d64b67fe84",
   "metadata": {},
   "source": [
    "# Example with ARGO data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5af732dd-5924-4d73-b28b-a127e5f62a37",
   "metadata": {},
   "source": [
    "We first need to download some ARGO data for example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6adf1cc-b88d-488f-80c2-e55750983d82",
   "metadata": {},
   "outputs": [],
   "source": [
    "import requests\n",
    "\n",
    "url = \"https://www.ncei.noaa.gov/data/oceans/argo/gadr/data/coriolis/69022/nodc_69022_prof.nc\"\n",
    "r = requests.get(url, allow_redirects=True)\n",
    "with open(\"ARGO_example.nc\", \"wb\") as f:\n",
    "    f.write(r.content)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32e5cf88-fbed-4c1c-a688-bed8d67322d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ruff: noqa: F401\n",
    "# ruff: noqa: I001\n",
    "import xarray as xr\n",
    "import cf_xarray\n",
    "import gsw_xarray as gsw"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56008322-1dcc-499a-8300-2f6e958f6318",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = xr.open_dataset(\"ARGO_example.nc\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "884b1e28-4fa5-4125-92bd-3b5ab2909935",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5d11d4ed-9788-4968-b100-5b3ff20dbcf2",
   "metadata": {},
   "source": [
    "We can rely on cf-xarray to see what variables have standard names in our dataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3f4cc86-bebd-43fb-88cd-0f15d199f2d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds.cf"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2e90a57-b352-42f4-8abe-477c4ae30bbd",
   "metadata": {},
   "source": [
    "The dataset contains multiple time the same variable (e.g. 'pres_adjusted' and 'pres' both have the standard name 'sea_water_pressure'). For the accessor to work, only 1 variable or each standard name must be present, explicitely stated when calling the function, or the gsw option `set_cf_name_preference` must be set.\n",
    "For this example we will only retain the adjusted variables. We set the global option, but we could also use context, i.e."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78d61954-33d8-43fb-ab50-e1ed4e9f9f4c",
   "metadata": {},
   "outputs": [],
   "source": [
    "with gsw.set_cf_name_preference(\n",
    "    sea_water_pressure=\"pres_adjusted\",\n",
    "    sea_water_practical_salinity=\"psal_adjusted\",\n",
    "    sea_water_temperature=\"temp_adjusted\",\n",
    "):\n",
    "    # do the computation\n",
    "    pass\n",
    "\n",
    "gsw.set_cf_name_preference(\n",
    "    sea_water_pressure=\"pres_adjusted\",\n",
    "    sea_water_practical_salinity=\"psal_adjusted\",\n",
    "    sea_water_temperature=\"temp_adjusted\",\n",
    ")\n",
    "\n",
    "# We can check the options we have set:\n",
    "gsw.get_options()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54017ff8-aefe-4215-99f6-bd1b1209e02c",
   "metadata": {},
   "source": [
    "In the following sections we will demonstrate each features of gsw-xarray. We will focus on computing the potential density anomaly."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "58eec720-ff4a-4030-a9c7-682c7cb239e6",
   "metadata": {},
   "source": [
    "## Basic usage as drop in replacement of gsw"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "200b24bc-b5cf-4ff6-9888-e68f7ecc3507",
   "metadata": {},
   "outputs": [],
   "source": [
    "help(gsw.sigma0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45a8b454-3f5a-418d-9517-1ac1b77690f2",
   "metadata": {},
   "source": [
    "We need Absolute Salinity and Conservative Temperature, so 1st we need to do some conversions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bfed7078-6e0c-41ec-acc1-a91cab0f295b",
   "metadata": {},
   "outputs": [],
   "source": [
    "SA = gsw.SA_from_SP(\n",
    "    SP=ds.psal_adjusted, p=ds.pres_adjusted, lon=ds.longitude, lat=ds.latitude\n",
    ")\n",
    "CT = gsw.CT_from_t(SA=SA, t=ds.temp_adjusted, p=ds.pres_adjusted)\n",
    "sigma0 = gsw.sigma0(SA=SA, CT=CT)\n",
    "sigma0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c91e3276-8fe1-420d-a973-65f646045047",
   "metadata": {},
   "source": [
    "## Using Pint and pint-xarray to handle units"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "376ba05c-4099-45ca-a9da-a9730aafa17f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pint_xarray\n",
    "import cf_xarray.units"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b1e28ce-43b3-4cba-bb95-25bd1f29f829",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_pint = ds.pint.quantify()\n",
    "ds_pint"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "69633a73-78e1-45cd-a7cf-2fffd61f797d",
   "metadata": {},
   "source": [
    "We compute again sigma0, using the `ds_pint` dataset, i.e. variables have a physical dimension"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a3c9493-477b-40a5-b78c-d7e44cd01e34",
   "metadata": {},
   "outputs": [],
   "source": [
    "SA = gsw.SA_from_SP(\n",
    "    SP=ds_pint.psal_adjusted,\n",
    "    p=ds_pint.pres_adjusted,\n",
    "    lon=ds_pint.longitude,\n",
    "    lat=ds_pint.latitude,\n",
    ")\n",
    "CT = gsw.CT_from_t(SA=SA, t=ds_pint.temp_adjusted, p=ds_pint.pres_adjusted)\n",
    "sigma0 = gsw.sigma0(SA=SA, CT=CT)\n",
    "sigma0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "987dac5f-bfc2-4ec9-8c92-9d0230f7975e",
   "metadata": {},
   "source": [
    "gsw-xarray converts the units (if necessary) when using pint quantities:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91a5a7d7-01bc-47de-b14f-b191ae58a3c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# start to convert the pressure into Pascal\n",
    "pressure_in_pascal = ds_pint.pres_adjusted.pint.to(\"Pa\")\n",
    "pressure_in_pascal"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20cc79d6-a870-49a3-ab20-6b2fb1af0d19",
   "metadata": {},
   "source": [
    "Compute again density, using the pressure in Pascal. No worries as the conversion to dbar is automatic!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1830ff3e-6387-4dde-ac8a-e96d0d924ddf",
   "metadata": {},
   "outputs": [],
   "source": [
    "SA = gsw.SA_from_SP(\n",
    "    SP=ds_pint.psal_adjusted,\n",
    "    p=pressure_in_pascal,\n",
    "    lon=ds_pint.longitude,\n",
    "    lat=ds_pint.latitude,\n",
    ")\n",
    "CT = gsw.CT_from_t(SA=SA, t=ds_pint.temp_adjusted, p=ds_pint.pres_adjusted)\n",
    "sigma0 = gsw.sigma0(SA=SA, CT=CT)\n",
    "sigma0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ded85225-5d7c-4c7c-812e-666a6cb722a4",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "7cca13a3-113e-4c11-9a23-5daacf24053b",
   "metadata": {},
   "source": [
    "## Using the accessor to simplify the workflow\n",
    "### Common case\n",
    "\n",
    "gsw-xarray adds ths `gsw` accessor to datasets. This accessor makes it easy to run the gsw functions on variables from a dataset."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c55e1386-d32a-4400-8e5a-e1718762990c",
   "metadata": {},
   "source": [
    "A first solution is to use the accessor exactly as when using gsw:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6226e913-4068-460a-ac97-8437b50f0ce8",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds.gsw.SA_from_SP(\n",
    "    SP=ds.psal_adjusted, p=ds.pres_adjusted, lon=ds.longitude, lat=ds.latitude\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "29fe070c-f87d-4db1-8da0-31a014913e2e",
   "metadata": {},
   "source": [
    "This is however not very useful... A better option is to simply give the name of the variables from the dataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b3b9855-7707-45ff-8b7d-0febbf5b11bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds.gsw.SA_from_SP(\n",
    "    SP=\"psal_adjusted\", p=\"pres_adjusted\", lon=\"longitude\", lat=\"latitude\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ecd1ac5c-531d-47e3-ab67-0584975f689c",
   "metadata": {},
   "source": [
    "It is even possible to go 1 step further and rely on the usage of standard name! In this case, you don't need to provide any argument for the variables with the proper standard name."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e897e151-99af-440c-86c6-47ae21eb2738",
   "metadata": {},
   "source": [
    "With this method, it is way faster to compute the density:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd7ccc59-f029-49bb-9d87-10fc4ca70c6a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# WITHOUT any detection\n",
    "SA = gsw.SA_from_SP(\n",
    "    SP=ds.psal_adjusted, p=ds.pres_adjusted, lon=ds.longitude, lat=ds.latitude\n",
    ")\n",
    "CT = gsw.CT_from_t(SA=SA, t=ds.temp_adjusted, p=ds.pres_adjusted)\n",
    "sigma0 = gsw.sigma0(SA=SA, CT=CT)\n",
    "\n",
    "# WITH autodetection\n",
    "ds = ds.merge(ds.gsw.SA_from_SP())\n",
    "ds = ds.merge(ds.gsw.CT_from_t())\n",
    "ds = ds.merge(ds.gsw.sigma0())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64bd1555-0be0-4c93-8704-ef117600ed3e",
   "metadata": {},
   "source": [
    "You can also use brackets to get either 1 or multiple variables computed:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "362fa683-b83a-4a00-a4c3-4d7d37c7e0bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds.gsw[\"sigma0\"]  # Returns a DataArray\n",
    "ds.gsw[[\"sigma0\"]]  # Returns a Dataset\n",
    "ds.gsw[[\"sigma0\", \"alpha\", \"beta\", \"sigma1\", \"rho\"]]  # With multiple outputs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4eeac5d0-ef80-4ef8-a1e7-6ff20168287b",
   "metadata": {},
   "source": [
    "Of course any kind of mixture between all the solutions is possible:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6af51df8-dc5f-4b0f-95da-69d2c87cb44e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Give a value for SP\n",
    "# Take p from dataset\n",
    "# Automatically get lon and lat based on standard names\n",
    "ds.gsw.SA_from_SP(SP=35, p=\"pres_adjusted\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25b37979-ee4c-4cd3-bbc6-ead741324e2f",
   "metadata": {},
   "source": [
    "And it is also possible to use automatic discovery of argument with pint datasets:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae3c41f1-5f12-4904-823b-85ddca6bc915",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds_pint.gsw.SA_from_SP()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9406a41-eb4e-4a96-8b7a-02f747b3d587",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "18d4a562-9007-44b1-a396-7840f2346c8c",
   "metadata": {},
   "source": [
    "### Case with some argument without standard name\n",
    "\n",
    "Some functions have argument without any standard name. In this case, it is possible to refer to these arguments using gsw-xarray options.\n",
    "\n",
    "Let's take the function `gsw.SP_salinometer` that has 2 arguments: `Rt` without standard name, and `t` the in situ temperature.\n",
    "\n",
    "A 1st option is to explicitely provide a value or the name from the dataset (we will create some fake data for the purpose of this example):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1fb9dfb7-8da7-4600-b17c-6714150b0875",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds[\"salinometer_Rt\"] = 35"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "479d1acc-b9e3-42e6-b975-233f42b4a582",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds.gsw.SP_salinometer(Rt=\"salinometer_Rt\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5807dd7c-d4de-467c-9b1a-2088fa6c8b84",
   "metadata": {},
   "source": [
    "A 2nd solution is to use gsw-xarray option with function `set_non_cf_name`. This way if you need to compute multiple times functions that use these arguments without a standard name, you only need to provide once the mapping."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7f80c8b-5fe5-4043-a4cd-696594e876a7",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "help(gsw.set_non_cf_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5d2d4f47-a2b9-4f4a-8a20-ef2cfa522c27",
   "metadata": {},
   "source": [
    "Using a context manager:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ef5fb23-6699-4544-9f09-d6d6380a26fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "with gsw.set_non_cf_name(Rt=\"salinometer_Rt\"):\n",
    "    ds.gsw.SP_salinometer()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0b97c10-fea3-44c6-b8d3-51d5e0884202",
   "metadata": {},
   "source": [
    "Or globally:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1b3f9fd-383f-45bf-a14f-6e1923ee194a",
   "metadata": {},
   "outputs": [],
   "source": [
    "gsw.set_non_cf_name(Rt=\"salinometer_Rt\")\n",
    "ds.gsw.SP_salinometer()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62e92b68-863a-4e80-b764-0cce276efe3e",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "ce93863a-36f4-4b18-af9e-eff87d484141",
   "metadata": {},
   "source": [
    "## Note on performance\n",
    "\n",
    "It is still very slightly faster to provide the arguments than rely on autodetect, however for large data sets this difference should be negligible compared to the internal computation time of the gsw functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aed2eb9c-cc29-44aa-a9ed-233e88a4bc5a",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "gsw.SA_from_SP(\n",
    "    SP=ds.psal_adjusted, p=ds.pres_adjusted, lon=ds.longitude, lat=ds.latitude\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d288bf2-d5f2-4e86-a93f-36ac1f62a969",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "ds.gsw.SA_from_SP(\n",
    "    SP=\"psal_adjusted\", p=\"pres_adjusted\", lon=\"longitude\", lat=\"latitude\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "064490ba-a5bd-42c7-8217-f2f9b132e935",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "ds.gsw.SA_from_SP()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3de7f43-bdc5-4093-ac86-1885df3b040a",
   "metadata": {},
   "source": [
    "Compare with upstream gsw"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d774bd8-1138-4d4a-b5ea-89fb2f31857e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import gsw as gsw_upstream"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39d99ceb-b488-4277-8713-2795290bd3f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit\n",
    "gsw_upstream.SA_from_SP(\n",
    "    SP=ds.psal_adjusted, p=ds.pres_adjusted, lon=ds.longitude, lat=ds.latitude\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}