initialize_rain_zarr.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b9850e43-c73f-4a4f-8fe9-0417e4c22f17",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import xarray as xr\n",
    "import pandas as pd\n",
    "\n",
    "import dask\n",
    "import dask.array as da\n",
    "from dask import delayed\n",
    "\n",
    "from glob import glob\n",
    "\n",
    "from datetime import datetime, timedelta\n",
    "import os\n",
    "import wrf\n",
    "import time\n",
    "\n",
    "import zarr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "e271de42-462e-4790-a064-c5b55f6ca042",
   "metadata": {},
   "outputs": [],
   "source": [
    "geog=xr.open_mfdataset('/eagle/climate_severe/geog/geo_em.d01.nc')\n",
    "\n",
    "dummy_shape=geog.HGT_M[0].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "24eb04a8-26bf-495c-bdd1-e70647fc881f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "AFWA_TOTPRECIP\n",
      "AFWA_TOTPRECIP\n",
      "Done\n"
     ]
    }
   ],
   "source": [
    "#%%time\n",
    "\n",
    "for sim in ['historical','end_of_century_8p5']:\n",
    "    for v in ['AFWA_TOTPRECIP']:\n",
    "        print(v)\n",
    "        #set zarr storage location\n",
    "        zarr_store='/eagle/climate_severe/bwallace_scratch/COARSENED/'+sim+'/'+v\n",
    "        #set variables to write to file\n",
    "        store_vars=[v,'Time']\n",
    "    \n",
    "        #establish path to pull data from\n",
    "        base_path='/eagle/climate_severe/Derived_Subsets/AFWA_TOTPRECIP/'+sim+'/'\n",
    "        #gather up the water year folders\n",
    "        basePathFiles=glob(base_path+'*')\n",
    "        basePathFiles.sort()\n",
    "\n",
    "        npList=np.array(basePathFiles)\n",
    "        \n",
    "        #listSize=npList.shape[0]\n",
    "        #listChunks=np.array_split(npList,15)\n",
    "\n",
    "        time_dims=np.array([1]*(npList.shape[0]*96))\n",
    "        full_shape=(time_dims.sum(),)+(dummy_shape)\n",
    "        offsets=np.cumsum(time_dims)\n",
    "        offsets-=offsets[0]\n",
    "\n",
    "        #define dictionary of attrs\n",
    "        ATTRS={\n",
    "            'Time':[\"Time\"],\n",
    "            'south_north':[\"south_north\"], #ydim\n",
    "            'west_east':[\"west_east\"], #xdim\n",
    "            v:[\"Time\",\"south_north\",\"west_east\"], #var is time, ydim, xdim\n",
    "    \n",
    "        }\n",
    "    \n",
    "        #store expected shapes of each var in dictionary\n",
    "        shp={\n",
    "            'Time':(time_dims.sum(),),\n",
    "            v:full_shape,\n",
    "        }\n",
    "    \n",
    "        #set chunksize. don't think Time matters here but i set it to match the variable\n",
    "        #below chunk configuration is chunked every 24 timesteps and continuous in space\n",
    "        chnks={\n",
    "            'Time':(24,),\n",
    "            v:(24,)+full_shape[1:],\n",
    "        }   \n",
    "        \n",
    "        #assign attrs to individual variables\n",
    "        attrs={}\n",
    "        attrs[v]={}\n",
    "        attrs[v][\"_ARRAY_DIMENSIONS\"]=ATTRS[v]\n",
    "    \n",
    "        #time is weird, need to do this or xarray screws up reading in the zarr\n",
    "        attrs['Time']={}\n",
    "        attrs['Time']['_ARRAY_DIMENSIONS']=ATTRS['Time']\n",
    "        attrs[\"Time\"].update(\n",
    "            calendar=\"proleptic_gregorian\",units='minutes since 1980-01-01', #make sure this is consistent with function\n",
    "        )\n",
    "\n",
    "        #create an empty zarr group \n",
    "        group=zarr.group(zarr_store)\n",
    "    \n",
    "        #allocate each variable in the zarr group\n",
    "        for varname in store_vars:\n",
    "            #can play around with the compression stuff, this had a nice tradeoff between compressibility and reaad times\n",
    "            v=group.empty(varname,shape=shp[varname],chunks=chnks[varname],dtype='float32',\n",
    "                          compressor=(zarr.Blosc(cname='zstd', clevel=3)),overwrite=True)\n",
    "            v.attrs.update(attrs[varname])\n",
    "        \n",
    "print('Done')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "92046ff4-2e97-42fa-950b-898a99dcc139",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<zarr.core.Array '/AFWA_TOTPRECIP' (525600, 899, 1399) float32>"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "group['AFWA_TOTPRECIP']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "56bd4a75-a971-4658-b7fb-f1d980a77d3e",
   "metadata": {},
   "outputs": [],
   "source": [
    "#zarr_store='/eagle/climate_severe/bwallace_scratch/COARSENED/end_of_century_8p5/AFWA_TOTPRECIP'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "ae72ffca-4ec7-46b6-a8a2-61ebfee5c79e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<zarr.hierarchy.Group '/'>"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#zarr.consolidate_metadata(zarr_store)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "81b68597-4bab-4e73-b946-4d8ac57b430a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "jhub_analysis",
   "language": "python",
   "name": "jhub_analysis"
  },
  "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.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}