diff --git a/test/ipynb/WRF_python_demo.ipynb b/test/ipynb/WRF_python_demo.ipynb index 8f073d0..f18c733 100644 --- a/test/ipynb/WRF_python_demo.ipynb +++ b/test/ipynb/WRF_python_demo.ipynb @@ -271,6 +271,182 @@ " " ] }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import tempfile\n", + "import glob\n", + "import shutil\n", + "import os\n", + "from netCDF4 import Dataset\n", + "from wrf import getvar, ll_to_xy, CoordPair, GeoBounds\n", + "\n", + "\n", + "class FileReduce(object):\n", + " def __init__(self, filenames, geobounds, tempdir=None, delete=True, reuse=True):\n", + " \"\"\"An iterable object for cutting out geographic domains.\n", + " \n", + " Args:\n", + " \n", + " filenames (sequence): A sequence of full paths to the WRF files\n", + " \n", + " geobounds (GeoBounds): A GeoBounds object defining the region of interest\n", + " \n", + " tempdir (str): The location to store the temporary cropped data files. If None, tempfile.mkdtemp is used.\n", + " \n", + " delete (bool): Set to True to delete the cropped files when FileReduce is garbage collected.\n", + " \n", + " reuse (bool): Set to True when you want to resuse the files that were previously converted. *tempdir* \n", + " must be set to a specific directory that contains the converted files.\n", + " \n", + " \"\"\"\n", + " self._filenames = filenames\n", + " self._i = 0\n", + " self._geobounds = geobounds\n", + " self._delete = delete\n", + " self._cache = set()\n", + " self._own_data = True\n", + " self._reuse = reuse\n", + " \n", + " if tempdir is not None:\n", + " if not os.path.exists(tempdir):\n", + " os.makedirs(tempdir)\n", + " self._tempdir = tempdir\n", + " if self._reuse:\n", + " self._cache = set((os.path.join(self._tempdir, name) \n", + " for name in os.listdir(self._tempdir)))\n", + " else:\n", + " self._tempdir = tempfile.mkdtemp()\n", + "\n", + " print (\"temporary directory is: {}\".format(self._tempdir))\n", + " self._prev = None\n", + " self._set_extents()\n", + " \n", + " def _set_extents(self):\n", + " fname = list(self._filenames)[0]\n", + " with Dataset(fname) as ncfile:\n", + " lons = [self._geobounds.bottom_left.lon, self._geobounds.top_right.lon]\n", + " lats = [self._geobounds.bottom_left.lat, self._geobounds.top_right.lat]\n", + " orig_west_east = len(ncfile.dimensions[\"west_east\"])\n", + " orig_south_north = len(ncfile.dimensions[\"south_north\"])\n", + " \n", + " \n", + " # Note: Not handling the moving nest here\n", + " # Extra points included around the boundaries to ensure domain is fully included\n", + " x_y = ll_to_xy(ncfile, lats, lons, meta=False)\n", + " self._start_x = 0 if x_y[0,0] == 0 else x_y[0,0] - 1\n", + " self._end_x = orig_west_east - 1 if x_y[0,1] >= orig_west_east - 1 else x_y[0,1] + 1\n", + " self._start_y = 0 if x_y[1,0] == 0 else x_y[1,0] - 1\n", + " self._end_y = orig_south_north if x_y[1,1] >= orig_south_north - 1 else x_y[1,1] + 1\n", + " \n", + " self._west_east = self._end_x - self._start_x + 1\n", + " self._west_east_stag = self._west_east + 1\n", + " self._south_north = self._end_y - self._start_y + 1\n", + " self._south_north_stag = self._south_north + 1\n", + " \n", + " \n", + " def __iter__(self):\n", + " return self\n", + " \n", + " def __copy__(self):\n", + " cp = type(self).__new__(self.__class__)\n", + " cp.__dict__.update(self.__dict__)\n", + " cp._own_data = False\n", + " cp._delete = False\n", + " \n", + " return cp\n", + " \n", + " def __del__(self):\n", + " if self._delete:\n", + " shutil.rmtree(self._tempdir)\n", + " \n", + " def reduce(self, fname):\n", + " outfilename = os.path.join(self._tempdir, \"reduced_\" + os.path.basename(fname))\n", + " \n", + " # WRF-Python can iterate over sequences several times during a 'getvar', so a cache is used to \n", + " if outfilename in self._cache:\n", + " return Dataset(outfilename)\n", + " \n", + " # New dimension sizes\n", + " dim_d = {\"west_east\" : self._west_east,\n", + " \"west_east_stag\" : self._west_east_stag,\n", + " \"south_north\" : self._south_north,\n", + " \"south_north_stag\" : self._south_north_stag\n", + " }\n", + " \n", + " # Data slice sizes for the 2D dimensions\n", + " slice_d = {\"west_east\" : slice(self._start_x, self._end_x + 1),\n", + " \"west_east_stag\" : slice(self._start_x, self._end_x + 2),\n", + " \"south_north\" : slice(self._start_y, self._end_y + 1),\n", + " \"south_north_stag\" : slice(self._start_y, self._end_y + 2)\n", + " }\n", + " \n", + " with Dataset(fname) as infile, Dataset(outfilename, mode=\"w\") as outfile:\n", + " print (\"reduce getting called!\")\n", + " \n", + " # Copy the global attributes\n", + " outfile.setncatts(infile.__dict__)\n", + "\n", + " # Copy Dimensions, limiting south_north and west_east to desired domain\n", + " for name, dimension in infile.dimensions.items():\n", + " dimsize = dim_d.get(name, len(dimension))\n", + " outfile.createDimension(name, dimsize)\n", + "\n", + " # Copy Variables \n", + " for name, variable in infile.variables.iteritems():\n", + " \n", + " new_slices = tuple((slice_d.get(dimname, slice(None)) for dimname in variable.dimensions))\n", + "\n", + " outvar = outfile.createVariable(name, variable.datatype, variable.dimensions)\n", + "\n", + " outvar[:] = variable[new_slices]\n", + "\n", + " outvar.setncatts(variable.__dict__)\n", + " \n", + " \n", + " result = Dataset(outfilename)\n", + " \n", + " self._cache.add(outfilename)\n", + " \n", + " return result\n", + " \n", + " \n", + " def next(self):\n", + " if self._i >= len(self._filenames):\n", + " if self._prev is not None:\n", + " self._prev.close()\n", + " raise StopIteration\n", + " else:\n", + " fname = self._filenames[self._i]\n", + " reduced_file = self.reduce(fname)\n", + " if self._prev is not None:\n", + " self._prev.close()\n", + " self._prev = reduced_file\n", + " \n", + " self._i += 1\n", + " \n", + " return reduced_file\n", + " \n", + " # Python 3\n", + " def __next__(self):\n", + " return self.next()\n", + "\n", + "ll = CoordPair(lat=24.0, lon=-87.)\n", + "ur = CoordPair(lat=27.0, lon=-84)\n", + "bounds = GeoBounds(ll, ur)\n", + "reduced_files = FileReduce(glob.glob(\"/Users/ladwig/Documents/wrf_files/wrf_vortex_multi/wrfout_d02*\"),\n", + " bounds, tempdir=\"/Users/ladwig/mytemp\", delete=False, reuse=True)\n", + "\n", + "slp = getvar(reduced_files, \"slp\")\n", + "print(slp)\n", + "\n", + "\n", + "del (reduced_files)\n" + ] + }, { "cell_type": "markdown", "metadata": {}, @@ -1227,21 +1403,21 @@ ], "metadata": { "kernelspec": { - "display_name": "Python 3", + "display_name": "Python 2", "language": "python", - "name": "python3" + "name": "python2" }, "language_info": { "codemirror_mode": { "name": "ipython", - "version": 3 + "version": 2 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.6.4" + "pygments_lexer": "ipython2", + "version": "2.7.15" } }, "nbformat": 4,