nilyin
/
wrf-python
forked from 3rdparty/wrf-python
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity
Browse Source

Added a new iterable class demo to reduce files

lon0
Bill Ladwig 7 years ago committed by Kevin Hallock
parent
ff2f561c51
commit
22c4efa8a9
1 changed files with 181 additions and 5 deletions
Whitespace
Show all changes Ignore whitespace when comparing lines Ignore changes in amount of whitespace Ignore changes in whitespace at EOL
Unified View
Diff Options
Show Stats Download Patch File Download Diff File
  1. 186
      test/ipynb/WRF_python_demo.ipynb

186
test/ipynb/WRF_python_demo.ipynb
Unescape View File

@ -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", "cell_type": "markdown",
"metadata": {}, "metadata": {},
@ -1227,21 +1403,21 @@
], ],
"metadata": { "metadata": {
"kernelspec": { "kernelspec": {
"display_name": "Python 3", "display_name": "Python 2",
"language": "python", "language": "python",
"name": "python3" "name": "python2"
}, },
"language_info": { "language_info": {
"codemirror_mode": { "codemirror_mode": {
"name": "ipython", "name": "ipython",
"version": 3 "version": 2
}, },
"file_extension": ".py", "file_extension": ".py",
"mimetype": "text/x-python", "mimetype": "text/x-python",
"name": "python", "name": "python",
"nbconvert_exporter": "python", "nbconvert_exporter": "python",
"pygments_lexer": "ipython3", "pygments_lexer": "ipython2",
"version": "3.6.4" "version": "2.7.15"
} }
}, },
"nbformat": 4, "nbformat": 4,

Write Preview
Loading…
Cancel
Save