{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Subset operator with kerchunk\n",
    "\n",
    "* https://projectpythia.org/kerchunk-cookbook/notebooks/using_references/Datatree.html\n",
    "* https://guide.cloudnativegeo.org/kerchunk/kerchunk-in-practice.html"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1",
   "metadata": {},
   "source": [
    "## Import clisops tools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import clisops.utils.dataset_utils as clidu\n",
    "from clisops.ops.subset import subset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3",
   "metadata": {},
   "source": [
    "## Example with testdata from CEDA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import clisops.utils.testing as clite\n",
    "\n",
    "mini_esgf_data = clite.get_kerchunk_datasets()\n",
    "mini_esgf_data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5",
   "metadata": {},
   "source": [
    "### Open remote dataset with clisops"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# using .json\n",
    "\n",
    "ds = clidu.open_xr_dataset(mini_esgf_data[\"CMIP6_KERCHUNK_HTTPS_OPEN_JSON\"])\n",
    "ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# using .zst with compression\n",
    "\n",
    "ds = clidu.open_xr_dataset(mini_esgf_data[\"CMIP6_KERCHUNK_HTTPS_OPEN_ZST\"])\n",
    "ds"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8",
   "metadata": {},
   "source": [
    "### Subset remote dataset with clisops"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "outputs = subset(\n",
    "    ds=ds,\n",
    "    time=\"1900-01-01/1900-12-31\",\n",
    "    area=(0.0, 10.0, 175.0, 90.0),\n",
    "    output_type=\"xarray\",\n",
    ")\n",
    "\n",
    "print(f\"There are {len(outputs)} outputs.\")\n",
    "outputs[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10",
   "metadata": {},
   "source": [
    "## Example from project Pythia\n",
    "\n",
    "Open pythia dataset with xarray.\n",
    "\n",
    "* https://projectpythia.org/kerchunk-cookbook/notebooks/using_references/Datatree.html"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11",
   "metadata": {},
   "source": [
    "### Test dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "url = \"s3://carbonplan-share/nasa-nex-reference/references_prod/ACCESS-CM2_historical/reference.parquet\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13",
   "metadata": {},
   "source": [
    "### Open remote s3 dataset with xarray open_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14",
   "metadata": {},
   "outputs": [],
   "source": [
    "import xarray as xr\n",
    "from fsspec.implementations.reference import ReferenceFileSystem\n",
    "\n",
    "fs = ReferenceFileSystem(\n",
    "    url,\n",
    "    remote_protocol=\"s3\",\n",
    "    target_protocol=\"s3\",\n",
    "    remote_options={\"anon\": True},\n",
    "    target_options={\"anon\": True},\n",
    "    lazy=True,\n",
    ")\n",
    "\n",
    "ds = xr.open_dataset(\n",
    "    fs.get_mapper(),\n",
    "    engine=\"zarr\",\n",
    "    backend_kwargs={\n",
    "        \"consolidated\": False,\n",
    "        \"zarr_format\": 2,\n",
    "    },\n",
    "    chunks={\"time\": 3},\n",
    ")\n",
    "\n",
    "ds"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15",
   "metadata": {},
   "source": [
    "### Open dataset with xarray open_zarr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16",
   "metadata": {},
   "outputs": [],
   "source": [
    "import fsspec\n",
    "\n",
    "mapper = fsspec.get_mapper(\n",
    "    \"reference://\",\n",
    "    fo=url,\n",
    "    target_options={\"anon\": True},\n",
    "    remote_options={\"anon\": True},\n",
    "    remote_protocol=\"s3\",\n",
    "    target_protocol=\"s3\",\n",
    ")\n",
    "\n",
    "ds = xr.open_zarr(mapper, consolidated=False, zarr_format=2)\n",
    "ds"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17",
   "metadata": {},
   "source": [
    "### Open Pythia s3 dataset with clisops"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18",
   "metadata": {},
   "outputs": [],
   "source": [
    "options = {\n",
    "    \"remote_protocol\": \"s3\",\n",
    "    \"target_protocol\": \"s3\",\n",
    "    \"remote_options\": {\"anon\": True},\n",
    "    \"target_options\": {\"anon\": True},\n",
    "}\n",
    "\n",
    "ds = clidu.open_xr_dataset(url, **options)\n",
    "ds"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19",
   "metadata": {},
   "source": [
    "### Subset Python S3 dataset with clisops"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20",
   "metadata": {},
   "outputs": [],
   "source": [
    "outputs = subset(\n",
    "    ds=ds,\n",
    "    time=\"2000-01-01/2000-01-31\",\n",
    "    area=(0.0, 10.0, 175.0, 90.0),\n",
    "    output_type=\"xarray\",\n",
    ")\n",
    "\n",
    "print(f\"There are {len(outputs)} outputs.\")\n",
    "outputs[0]"
   ]
  }
 ],
 "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.13.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}