diff --git a/demo.ipynb b/demo.ipynb index 7fae5a3c..55f0e9ad 100644 --- a/demo.ipynb +++ b/demo.ipynb @@ -18,6 +18,7 @@ "import subprocess\n", "from scipy.ndimage import binary_fill_holes\n", "from importlib import reload\n", + "# os.chdir(\"PATH TO WHERE YOU GIT CLONED\") ## Make sure your kernel is in the same directory as the regional_library.py file\n", "import regional_library as ml\n", "from dask.distributed import Client\n", "client = Client()\n", @@ -29,22 +30,31 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "# Demonstration of Regional Pipeline\n", + "# What does this package do?\n", + "\n", + "Setting up a regional model in MOM6 is a pain. The goal of this package is that users should spend their debugging time fixing a model that's running and doing weird things, rather than puzzling over a model that won't even start.\n", + "\n", + "In running this notebook, you'll hopefully have a running mom6 regional model. There will still be a lot of fiddling to do with the MOM_input file to make sure that the parameters are set up right for your domain, and you might want to manually edit some of the input files. BUT, this package should help you bypass most of the woes of regridding, encoding and understanding the arcane arts of the mom6 boundary segment files. \n" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ "\n", "This notebook is designed to showcase where we're up to so far. By the end you should have a running mom6 experiment on the domain of your choice. To make a stable test case:\n", "\n", - "* Keep your domain fairly small (my test case is a rectangle around Tasmania). If you go bigger you'll need to do some FRE tool shenannigens as explained in step 5 to get it working. \n", "* Avoid any regions with ice\n", "* Avoid regions near the north pole\n", "* Although the default configuration is meant to be RYF, I've not fixed up the calendar and encoding to run longer than a year just yet\n", "* If you choose to do OM2-01 forcing, set your start date to 1990-01-01 which is what I've got it hardcoded to in step 2 option 2. \n", "\n", - "Also hgrid is currently **not** mercator. It's equally spaced lat/long. To be updated very soon.\n", - "\n", + "Also hgrid is currently **not** mercator. It's equally spaced lat/long, with square cells at the centre of your domain and decrease in cell area away from equator. The pipeline is modular however, so if another hgrid generation function is written, it will be easy to pass say a gridtype=\"mercator\" to the experiment class.\n", "\n", "Input Type | Source\n", "---|---\n", - "Surface | JRA\n", + "Surface | JRA (or ERA5 - see end of notebook)\n", "Ocean | GLORYS reanalysis product OR ACCESS OM2-01\n", "Bathymetry | Gebco" ] @@ -54,7 +64,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "# Step 1: Choose our domain, define workspace paths" + "# Step 0: Your personal environment variables" ] }, { @@ -63,23 +73,57 @@ "metadata": {}, "outputs": [], "source": [ - "reload(ml)\n", - "expt_name = \"democosima\"\n", + "scratch = \"/scratch/v45/ab8992\"\n", + "home = \"/home/149/ab8992\"\n", + "## If using GLORYs, you'll need an email and password to access their database. make an account here: https://www.copernicus.eu/en/user/login?\n", + "pwd = \"YOUR COPERNICUS PASSWORD\" \n", + "usr = \"YOUR COPERNICUS USERNAME\" " + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Step 1: Choose our domain, define workspace paths\n", + "\n", + "To make sure that things are working I'd recommend starting with the default example defined below. If this runs ok, then change to a domain of your choice and hopefully it runs ok too! There's some troubleshooting you can do if not (check readme / readthedocs)\n", + "\n", + "To find the lat/lon of the domain you want to test you can use this GUI and copy paste below" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "expt_name = \"tasmania-20thdeg\"\n", "\n", "## Choose your coordinates and the name of your experiment\n", - "yextent = [-48,-38.95] ## latitude\n", - "xextent = [143, 150] ## longitude\n", + "yextent = [-48,-38.95] ## latitude\n", + "xextent = [143,150] ## longitude\n", "\n", - "daterange = [\"2003-01-01 00:00:00\", \"2003-01-05 00:00:00\"]\n", + "daterange = [\"2003-01-01 00:00:00\", \"2003-01-05 00:00:00\"] ## 2003 is a good compimise for GLORYs and JRA forcing as they overlap. JRA ends in 2012, GLORYS starts in 1993\n", "\n", "## Place where all your input files go\n", - "inputdir = f\"/scratch/v45/ab8992/mom6/regional_configs/{expt_name}/\"\n", + "inputdir = f\"{scratch}/mom6_regional_configs/{expt_name}/\"\n", "\n", "## Directory where you'll run the experiment from\n", - "rundir = f\"/home/149/ab8992/mom6_rundirs/{expt_name}/\"\n", + "rundir = f\"{home}/mom6_rundirs/{expt_name}/\"\n", "\n", "## Directory where fre tools are stored\n", - "toolpath = \"/home/157/ahg157/repos/mom5/src/tools/\" ## Compiled tools needed for construction of mask tables\n" + "toolpath = \"/home/157/ahg157/repos/mom5/src/tools/\" ## Compiled tools needed for construction of mask tables\n", + "\n", + "## Directory where raw downloads go before processing\n", + "tmpdir = f\"{scratch}/regional_tmp/{expt_name}\"\n", + "\n", + "for i in [rundir,tmpdir,inputdir]:\n", + " if not os.path.exists(i):\n", + " subprocess.run(f\"mkdir {i} -p\",shell=True)\n", + "\n", + "tmpdir = f\"{scratch}/regional_tmp/reanalysis-small\"\n", + "\n" ] }, { @@ -89,7 +133,25 @@ "source": [ "# Step 2: Prepare ocean forcing data\n", "\n", - "We need to cut out our ocean forcing. The pipeline expects an initial condition and one time-dependent segment per non-land boundary. Naming convention is \"east_unprocessed\" and \"ic_inprocessed\" for initial condition. Execute either of the following cells to pick GLORYs reanalysis or ACCESS OM2-01" + "We need to cut out our ocean forcing. The pipeline expects an initial condition and one time-dependent segment per non-land boundary. Naming convention is \"east_unprocessed\" and \"ic_unprocessed\" for initial condition. Execute either of the following cells to pick GLORYs reanalysis or ACCESS OM2-01" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## For default 'Tassie' domain:\n", + "You can just read in the boundaries I've already downloaded. Overwrite your tmpdir and continue with the notebook without generating ocean forcing files" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "tmpdir = \"/g/data/v45/ab8992/tassie-glorys\"" ] }, { @@ -102,7 +164,7 @@ "\n", "To do this you'll need to register with the Copernicus data centre to get a username and password. Fill these in below.\n", "\n", - "After executing, navigate to this directory in your terminal and run 'bash get_oceanfiles.sh'. Wait for all of your forcing segments to appear before continuing with the 'ocean forcing' step" + "After executing, navigate to this directory in your terminal and double check that all the files are there! Sometimes the data centre hangs and only retrieves a couple of files. In thise case, comment out the completed segments in `get_oceanfiles.sh` and run it again from terminal." ] }, { @@ -111,17 +173,16 @@ "metadata": {}, "outputs": [], "source": [ - "## Directory where raw downloads go before processing\n", - "tmpdir = f\"/scratch/v45/ab8992/regional_tmp/{expt_name}/\"\n", - "\n", - "pwd = \"YOUR COPERNICUS PASSWORD\" \n", - "usr = \"YOUR COPERNICUS USERNAME\" \n", "\n", "file = open(f\"{tmpdir}/get_oceanfiles.sh\",\"w\")\n", "file.write(\n", " ml.motu_requests(xextent, yextent, daterange, tmpdir, usr, pwd,[\"north\",\"south\",\"east\",\"west\"])\n", ")\n", - "file.close()\n" + "file.close()\n", + "\n", + "subprocess.run(\n", + " f\"bash {tmpdir}/get_oceanfiles.sh\",shell=True\n", + ")\n" ] }, { @@ -131,9 +192,11 @@ "source": [ "## Option 2: ACCESS OM2-01\n", "\n", - "If you have access to where it's located on Gadi, you can execute the following cell to cut out and save your segments and use these instead.\n", + "If you have access to where it's located on Gadi, you can execute the following cell to cut out and save your segments and use these instead. The default I've set it at below is to cut out 3 months. To cut out a year, uncomment the code above which concatenates several input files together. Keep in mind that these input files are HUGE and they'll take a while to open and processes. To do a whole year, you'll want to run with a whole node and go make yourself a cup of coffee (and maybe read the paper for a bit). \n", "\n", - "**NOTE: I haven't automated this properly. You'll need to fiddle around with the 'for i in range(1077,1082)' line to choose the right year. Could maybe use COSIMA cookbook for this step instead?**" + "The advantage of doing this though is that the input files that the pipeline has to deal with are a lot smaller, making subsequent computation a lot quicker. An older iteration of the boundary brushcutter was to read data directly from the huge datasets, but this required some very careful chunking to not break your kernel. \n", + "\n", + "**NOTE: I haven't automated this properly and it's hardcoded for the year of 1990, which corresponds to files 1077 - 1082. Could maybe use COSIMA cookbook for this step instead?**" ] }, { @@ -142,14 +205,20 @@ "metadata": {}, "outputs": [], "source": [ - "tmpdir = f\"/scratch/v45/ab8992/regional_tmp/{expt_name}/\"\n", "\n", - "om2_input = xr.concat(\n", - " [xr.open_mfdataset(f\"/g/data/ik11/outputs/access-om2-01/01deg_jra55v13_ryf9091/output{i}/ocean/ocean_daily*\",decode_times = False,parallel=True,chunks='auto') for i in range(1077,1082)]\n", - ")\n", + "########## TWO OPTIONS: ############################\n", + "\n", + "## Use this if you want to do a quick test for up to 3 months\n", + "om2_input = xr.open_mfdataset(f\"/g/data/ik11/outputs/access-om2-01/01deg_jra55v13_ryf9091/output1077/ocean/ocean_daily*\",decode_times = False,parallel=True,chunks='auto')\n", + "\n", + "## Use this to cut out entire year \n", + "# om2_input = xr.concat(\n", + "# [xr.open_mfdataset(f\"/g/data/ik11/outputs/access-om2-01/01deg_jra55v13_ryf9091/output{i}/ocean/ocean_daily*\",decode_times = False,parallel=True,chunks='auto') for i in range(1077,1082)],\n", + "# \"time\"\n", + "# )\n", "#! for i in range(1077,1082) is hardcoded to choose the year of 1990 Jan -> Dec 31. \n", + "#######################################################\n", "\n", - "reload(ml)\n", "## Cut out initial condition and save\n", "ic = om2_input[[\"u\",\"v\",\"salt\",\"temp\",\"eta_t\"]].sel( \n", " yu_ocean = slice(yextent[0] - 0.2,yextent[1] + 0.2),\n", @@ -160,6 +229,7 @@ "ic = ml.nicer_slicer(ic,[xextent[0],xextent[1]],[\"xu_ocean\",\"xt_ocean\"])\n", "ic.to_netcdf(tmpdir + \"/ic_unprocessed\")\n", "\n", + "## Cut out East and West segments. Does lat slice first then uses nicer slicer for lon slice\n", "eastwest = om2_input[[\"u\",\"v\",\"salt\",\"temp\",\"eta_t\"]].sel( \n", " yu_ocean = slice(yextent[0] - 0.2,yextent[1] + 0.2),\n", " yt_ocean = slice(yextent[0] - 0.2,yextent[1] + 0.2)\n", @@ -167,6 +237,7 @@ "ml.nicer_slicer(eastwest,[xextent[1],xextent[1]],[\"xu_ocean\",\"xt_ocean\"]).to_netcdf(tmpdir + \"/east_unprocessed\")\n", "ml.nicer_slicer(eastwest,[xextent[0],xextent[0]],[\"xu_ocean\",\"xt_ocean\"]).to_netcdf(tmpdir + \"/west_unprocessed\")\n", "\n", + "## Cut out North and South segments\n", "northsouth = ml.nicer_slicer(om2_input[[\"u\",\"v\",\"salt\",\"temp\",\"eta_t\"]],[xextent[0],xextent[1]],[\"xu_ocean\",\"xt_ocean\"])\n", "northsouth.sel(\n", " yu_ocean = slice(yextent[1] - 0.2,yextent[1] + 0.2),\n", @@ -186,16 +257,22 @@ "# Step 3: Make experiment object\n", "This object keeps track of your domain basics, as well as generating the hgrid, vgrid and setting up the folder structures. \n", "\n", - "After running you can have a look at your grids by calling expt.hgrid and expt.vgrid\n", + "After running you can have a look at your grids by calling `expt.hgrid` and `expt.vgrid`\n", "\n", - "Plotting vgrid with marker = '.' option lets you see the spacing, or plotting np.diff(expt.hgrid.zl).plot(marker = '.') shows you the vertical spacing profile.\n", + "Plotting vgrid with marker = '.' option lets you see the spacing, or plotting \n", + "```python\n", + "np.diff(expt.hgrid.zl).plot(marker = '.')\n", + "```\n", + " shows you the vertical spacing profile.\n", "\n", "## Modular workflow!\n", "\n", "After constructing your expt object, if you don't like my lazy default hgrid and vgrid you can simply modify and overwrite them. However, you'll also need to save them to disk again as I've not automated this just yet. For example:\n", "\n", + "```python\n", "expt.hgrid = custom_hgrid\n", - "expt.hgrid.to_netcdf(f\"{inputdir}/hgrid.nc\")" + "expt.hgrid.to_netcdf(f\"{inputdir}/hgrid.nc\")\n", + "```" ] }, { @@ -204,6 +281,7 @@ "metadata": {}, "outputs": [], "source": [ + "reload(ml)\n", "expt = ml.experiment(\n", " xextent,\n", " yextent,\n", @@ -223,7 +301,49 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "# Step 4: Handle the ocean forcing.\n", + "# Step 4: Set up bathymetry\n", + "\n", + "Similarly to ocean forcing, we point our 'bathymetry' method at the location of the file of choice, and pass it a dictionary mapping variable names. This time we don't need to preprocess the topography since it's just a 2D field and easier to deal with. Afterwards you can run `expt.topog` and have a look at your domain. After running this cell, your input directory will contain other topography - adjacent things like the ocean mosaic and mask table too. This defaults to a 10x10 layout which can be updated later." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "expt.bathymetry(\n", + " '/g/data/ik11/inputs/GEBCO_2022/GEBCO_2022.nc',\n", + " {\"xh\":\"lon\",\n", + " \"yh\":\"lat\",\n", + " \"elevation\":\"elevation\"}, ## Again this dictionary just maps mom6 variable names to what they are in your topog.\n", + " minimum_layers = 1\n", + " )" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Check out your domain:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "expt.topog.depth.plot()" + ] + }, + { + "attachments": {}, + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Step 5: Handle the ocean forcing - where the magic happens\n", "\n", "This cuts out and interpolates the initial condition as well as all boundaries (unless you don't pass it boundaries).\n", "\n", @@ -231,7 +351,7 @@ "\n", "If one of your segments is land, you can delete its string from the 'boundaries' list. You'll need to update MOM_input to reflect this though so it knows how many segments to look for, and their orientations. \n", "\n", - "### **Note: Only run one of the two cells below according to what forcing you chose!" + "### **Note: Only run one of the two cells below according to what forcing you chose!**" ] }, { @@ -283,37 +403,6 @@ " )" ] }, - { - "attachments": {}, - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Step 5: Set up bathymetry\n", - "\n", - "### **This is the most annoying step!** Will fix in future update\n", - "Currently my library calls 'make_topog_parallel' as a subprocess. For large domains, this will simply hang as it won't have enough memory. If it doesn't print after 3 mins just cancel it. In this case you need to go to your input directory, start an interactive job and run:\n", - "\n", - "PATH_TO_EXECUTABLE/make_topog_parallel --mosaic ocean_mosaic.nc --topog_type realistic --topog_file bathy_original.nc --topog_field 'elevation' --scale_factor -1 --output topog_raw.nc\n", - "\n", - "A path to the executable is /g/data/v45/jr5971/FRE-NCtools/build3_up_MAXXGRID/tools/make_topog/ if you have access.\n", - "\n", - "After this, you need to run expt.bathymetry again, this time passing the flag 'maketopog = False'. It will continue from after the make_topog step and finish the job. " - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "expt.bathymetry(\n", - " '/g/data/ik11/inputs/GEBCO_2022/GEBCO_2022.nc',\n", - " {\"xh\":\"lon\",\n", - " \"yh\":\"lat\",\n", - " \"elevation\":\"elevation\"}, ## Again this dictionary just maps mom6 variable names to what they are in your topog.\n", - " )" - ] - }, { "attachments": {}, "cell_type": "markdown", @@ -321,7 +410,7 @@ "source": [ "# Step 6 (optional) Select number of processors \n", "\n", - "This is just a wrapper for check_mask FRE tool. Choose the number of processors in the X and Y directions respectively" + "This is just a wrapper for check_mask FRE tool. Choose the number of processors in the X and Y directions respectively. Having run `.bathymetry()`, you already have one set up for 10x10 by default." ] }, { @@ -340,7 +429,7 @@ "source": [ "## Step 7 (optional) Regrid the runoff \n", "\n", - "### This step will be removed in a future update when this functionality is added to rest of pipeline. Currently it calls a function from the legacy regional_model_scripts file. Just execute cell to give your domain runoff from JRA\n" + " This step will be removed in a future update when this functionality is added to rest of pipeline. Currently it calls a function from the legacy regional_model_scripts file. Just execute cell to give your domain runoff from JRA\n" ] }, { @@ -367,7 +456,8 @@ "source": [ "# Step 8: Modify the default input directory to make a (hopefully) runnable configuration out of the box\n", "\n", - "This cell just copies a default run directory and modifies it to match your configuration." + "This cell just copies a default run directory and modifies it to match your configuration.\n", + "\n" ] }, { @@ -376,7 +466,8 @@ "metadata": {}, "outputs": [], "source": [ - "subprocess.run(f\"cp default_rundir/jra_surface* {rundir}\",shell = True)\n", + "subprocess.run(f\"cp default_rundir/jra_surface/* {rundir} -r\",shell = True)\n", + "# subprocess.run(f\"cp default_rundir/era5_surface/* {rundir} -r\",shell = True)\n", "subprocess.run(f\"ln -s {inputdir} {rundir}/inputdir\",shell=True)\n", "\n", "hgrid = xr.open_dataset(f\"{inputdir}/hgrid.nc\")\n", @@ -402,6 +493,9 @@ " lines[i] = f'MASKTABLE = \"{mask_table}\"\\n'\n", " else:\n", " lines[i] = \"# MASKTABLE = no mask table\"\n", + " # if \"LAYOUT\" in lines[i]:\n", + " # lines[i] = f'LAYOUT = {expt.layout[0]},{expt.layout[1]}\\n'\n", + "\n", " if \"NIGLOBAL\" in lines[i]: \n", " # lines[i] = f\"NIGLOBAL = {str(x_indices_centre[1] - x_indices_centre[0])}\\n\"\n", " lines[i] = f\"NIGLOBAL = {hgrid.nx.shape[0]//2}\\n\"\n", @@ -427,7 +521,8 @@ " if \"NIGLOBAL\" in lines[i]:\n", " # lines[i] = f\"NIGLOBAL = {str(x_indices_centre[1] - x_indices_centre[0])}\\n\"\n", " lines[i] = f\"NIGLOBAL = {hgrid.nx.shape[0]//2}\\n\"\n", - " \n", + " # if \"LAYOUT\" in lines[i]:\n", + " # lines[i] = f'LAYOUT = {expt.layout[0]},{expt.layout[1]}\\n'\n", " if \"NJGLOBAL\" in lines[i]:\n", " # lines[i] = f\"NJGLOBAL = {str(y_indices_centre[1] - y_indices_centre[0])}\\n\"\n", " lines[i] = f\"NJGLOBAL = {hgrid.ny.shape[0]//2}\\n\"\n", @@ -448,9 +543,24 @@ " \n", " if \"input:\" in lines[i]:\n", " lines[i + 1] = f\" - {inputdir}\\n\"\n", - " \n", + "\n", "inputfile = open(f\"{rundir}/config.yaml\",'w')\n", "inputfile.writelines(lines)\n", + "inputfile.close()\n", + "\n", + "\n", + "# Modify input.nml \n", + "inputfile = open(f\"{rundir}/input.nml\",'r')\n", + "lines = inputfile.readlines()\n", + "inputfile.close()\n", + "for i in range(len(lines)):\n", + " if \"current_date\" in lines[i]:\n", + " tmp = daterange[0].split(\" \")[0].split(\"-\")\n", + " lines[i] = f\"{lines[i].split(' = ')[0]} = {int(tmp[0])},{int(tmp[1])},{int(tmp[2])},0,0,0,\\n\"\n", + "\n", + " \n", + "inputfile = open(f\"{rundir}/input.nml\",'w')\n", + "inputfile.writelines(lines)\n", "inputfile.close()\n" ] }, @@ -461,7 +571,7 @@ "source": [ "# BONUS! Want to use ERA5 surface forcing instead?\n", "\n", - "This is WIP and not tested but thought I'd include it" + "This is WIP and not well tested but thought I'd include it" ] }, { @@ -499,14 +609,13 @@ "metadata": {}, "outputs": [], "source": [ - "#! Messy Work in progress for now but works. I think there's an issue with specific humidity units?\n", - "\n", "erapath = \"/g/data/rt52/era5/single-levels/reanalysis\"\n", "\n", "## Firstly just open all raw data\n", "rawdata = {}\n", "for fname , vname in zip([\"2t\",\"10u\",\"10v\",\"sp\",\"2d\"] , [\"t2m\",\"u10\",\"v10\",\"sp\",\"d2m\"]):\n", "\n", + " ## Cut out this variable to our domain size\n", " rawdata[fname] = ml.nicer_slicer(\n", " xr.open_mfdataset(f\"{erapath}/{fname}/{daterange[0].split('-')[0]}/{fname}*\",decode_times = False,chunks = {\"longitude\":100,\"latitude\":100}),\n", " xextent,\n", @@ -532,7 +641,7 @@ " ## Calculate specific humidity from dewpoint temperature \n", " q = xr.Dataset(\n", " data_vars= {\n", - " \"q\": 0.001 * (0.622 / rawdata[\"sp\"][\"sp\"]) * (10**(8.07131 - 1730.63 / (233.426 + rawdata[\"2d\"][\"d2m\"]))) * 101325 / 760\n", + " \"q\": (0.622 / rawdata[\"sp\"][\"sp\"]) * (10**(8.07131 - 1730.63 / (233.426 + rawdata[\"2d\"][\"d2m\"] - 273.15) )) * 101325 / 760\n", " }\n", "\n", " )\n", @@ -544,11 +653,7 @@ "\n", "## Update the data table to match:\n", "\n", - "subprocess.run(f\"cp default_rundir/era5_surface/data_table {rundir}/data_table\",shell = True)\n", - "\n", - "\n", - "\n", - "\n" + "subprocess.run(f\"cp default_rundir/era5_surface/data_table {rundir}/data_table\",shell = True)" ] } ], @@ -558,6 +663,18 @@ "language": "python", "name": "conda-env-analysis3-py" }, + "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.9.15" + }, "orig_nbformat": 4 }, "nbformat": 4, diff --git a/regional_library.py b/regional_library.py index 6163c033..26b9a7d7 100644 --- a/regional_library.py +++ b/regional_library.py @@ -299,7 +299,7 @@ class experiment: Knows everything about your regional experiment! Methods in this class will generate the various input files you need to generate a MOM6 experiment forced with Open Boundary Conditions. It's written agnostic to your choice of boundary forcing,topography and surface forcing - you need to tell it what your variables are all called via mapping dictionaries where keys are mom6 variable / coordinate names, and entries are what they're called in your dataset. """ - def __init__(self, xextent,yextent,daterange,resolution,vlayers,dz_ratio,depth,mom_run_dir,mom_input_dir,toolpath): + def __init__(self, xextent,yextent,daterange,resolution,vlayers,dz_ratio,depth,mom_run_dir,mom_input_dir,toolpath,gridtype = "even_spacing"): try: os.mkdir(mom_run_dir) except: @@ -320,8 +320,9 @@ def __init__(self, xextent,yextent,daterange,resolution,vlayers,dz_ratio,depth,m self.mom_run_dir = mom_run_dir self.mom_input_dir = mom_input_dir self.toolpath = toolpath - self.hgrid = self._make_hgrid() + self.hgrid = self._make_hgrid(gridtype) self.vgrid = self._make_vgrid() + self.gridtype = gridtype # if "temp" not in os.listdir(inputdir): # os.mkdir(inputdir + "temp") @@ -337,19 +338,28 @@ def __init__(self, xextent,yextent,daterange,resolution,vlayers,dz_ratio,depth,m return - def _make_hgrid(self,method = "even_spacing"): + def _make_hgrid(self,gridtype): """ Sets up hgrid based on users specification of domain. Default behaviour leaves latitude and longitude evenly spaced. If user specifies a resolution of 0.1 degrees, longitude is spaced this way and latitude spaced with 0.1 cos(mean_latitude). This way, grids in the centre of the domain are perfectly square, but decrease monatonically in area away from the equator """ - if method == "even_spacing": + if gridtype == "even_spacing": # longitudes will just be evenly spaced, based only on resolution and bounds - x = np.linspace(self.xextent[0],self.xextent[1],int((self.xextent[1] - self.xextent[0])/(self.res / 2)) + 1) + nx = int((self.xextent[1] - self.xextent[0])/(self.res / 2)) + if nx %2 != 1: + nx += 1 + + x = np.linspace(self.xextent[0],self.xextent[1],nx) # Latitudes evenly spaced by dx * cos(mean_lat) res_y = self.res * np.cos(np.mean(self.yextent) * np.pi / 180) - y = np.linspace(self.yextent[0],self.yextent[1],int((self.yextent[1] - self.yextent[0])/(res_y / 2)) + 1) + + ny = int((self.yextent[1] - self.yextent[0])/(res_y / 2)) + 1 + if ny %2 != 1: + ny += 1 + + y = np.linspace(self.yextent[0],self.yextent[1],ny) hgrid = rectangular_hgrid(x,y) hgrid.to_netcdf(self.mom_input_dir + "/hgrid.nc") @@ -641,61 +651,73 @@ def bathymetry(self,bathy_path,varnames,fill_channels = False,minimum_layers = 3 """ - - - ## Determine whether we need to adjust bathymetry longitude to match model grid. - # - # eg if bathy is 0->360 but self.hgrid is -180->180, longitude slice becomes - ## - - # if bathy[varnames["xh"]].values[0] < 0: - - - if maketopog == True: - bathy = xr.open_dataset(bathy_path,chunks="auto")[varnames["elevation"]] + bathy = xr.open_dataset(bathy_path,chunks={varnames["xh"]:500,varnames["yh"]:500})[varnames["elevation"]] bathy = bathy.sel({ - varnames["yh"]:slice(self.yextent[0],self.yextent[1]) + varnames["yh"]:slice(self.yextent[0] - 0.1,self.yextent[1] + 0.1) } ).astype("float") - bathy = nicer_slicer(bathy,self.xextent,varnames["xh"]) - + bathy = nicer_slicer(bathy,np.array(self.xextent) + np.array([-0.1,0.1]),varnames["xh"]) bathy.attrs['missing_value'] = -1e20 # This is what FRE tools expects I guess? - bathy.to_netcdf(f"{self.mom_input_dir}bathy_original.nc", engine='netcdf4') + bathy = xr.Dataset({"elevation":bathy}) - # #! New code to test: Can we regrid first to pass make_topog a smaller dataset to handle? - # tgrid = xr.Dataset( - # {"lon":(["lon"],self.hgrid.x.isel(nxp=slice(1, None, 2), nyp=1).values), - # "lat":(["lat"],self.hgrid.y.isel(nxp=1, nyp=slice(1, None, 2)).values) - # } - # ) - # regridder_t = xe.Regridder( - # bathy, tgrid, "bilinear", - # ) + bathy.lon.attrs["units"] = "degrees_east" + bathy.lat.attrs["units"] = "degrees_north" + bathy.lon.attrs["_FillValue"] = 1e20 + bathy.elevation.attrs["_FillValue"] = 1e20 + bathy.elevation.attrs["units"] = "m" + bathy.elevation.attrs["standard_name"] = "height_above_reference_ellipsoid" + bathy.elevation.attrs["long_name"] = "Elevation relative to sea level" + bathy.elevation.attrs["coordinates"] = "lon lat" - # bathy_regrid.to_netcdf(f"{self.mom_input_dir}bathy_regrid.nc", engine='netcdf4') - # #! End new test code - ## Now pass bathymetry through the FRE tools + bathy.to_netcdf(f"{self.mom_input_dir}bathy_original.nc", mode = "w",engine='netcdf4') - ## Make Topog - args = f"--mosaic ocean_mosaic.nc --topog_type realistic --topog_file bathy_original.nc --topog_field {varnames['elevation']} --scale_factor -1 --output topog_raw.nc".split(" ") - print( - "FRE TOOLS: make topog parallel\n\n", - subprocess.run(["/g/data/v45/jr5971/FRE-NCtools/build3_up_MAXXGRID/tools/make_topog/make_topog_parallel"] + args,cwd = self.mom_input_dir) + tgrid = xr.Dataset( + {"lon":(["lon"],self.hgrid.x.isel(nxp=slice(1, None, 2), nyp=1).values), + "lat":(["lat"],self.hgrid.y.isel(nxp=1, nyp=slice(1, None, 2)).values) + } ) + tgrid.lon.attrs["units"] = "degrees_east" + tgrid.lon.attrs["_FillValue"] = 1e20 + tgrid.lat.attrs["units"] = "degrees_north" + # tgrid.to_netcdf(f"{self.mom_input_dir}tgrid.nc", mode = "w",engine='netcdf4') + tgrid.to_netcdf(f"{self.mom_input_dir}topog_raw.nc", mode = "w",engine='netcdf4') + + #! Hardcoded for whole node notebook. + # topog_raw file is the 'target' grid used for gridgen. This is then overweitten by the second ESMF function (needs a blank netcdf to overwrite as the output) + if subprocess.run( + "mpirun ESMF_Regrid -s bathy_original.nc -d topog_raw.nc -m bilinear --src_var elevation --dst_var elevation --netcdf4 --src_regional --dst_regional", + shell = True,cwd = self.mom_input_dir).returncode != 0: + raise RuntimeError("Regridding of bathymetry failed! This is probably because mpirun was initialised earlier by xesmf doing some other regridding. Try restarting the kernel, then calling .bathymetry() before any other methods.") + ## reopen topography to modify topog = xr.open_dataset(self.mom_input_dir + "topog_raw.nc", engine='netcdf4') - ## Remove inland lakes + ## Ensure correct encoding + topog = xr.Dataset( + {"depth":(["ny","nx"],topog[varnames["elevation"]].values)} + ) + topog.attrs["depth"] = "meters" + topog.attrs["standard_name"] = "topographic depth at T-cell centers" + topog.attrs["coordinates"] = "zi" + + topog.expand_dims("tiles",0) + + if topog.depth.mean() < 0: + ## Ensure that coordinate is positive down! + topog["depth"] *= -1 + + + ## REMOVE INLAND LAKES min_depth = self.vgrid.zi[minimum_layers] @@ -743,16 +765,12 @@ def bathymetry(self,bathy_path,varnames,fill_channels = False,minimum_layers = 3 forward = True - + newmask = xr.where(newmask > 0 , 1,0) changed = np.max(newmask) == 1 land_mask += newmask - # land_mask.to_netcdf(self.mom_input_dir + "land_mask.nc") ocean_mask = np.abs(land_mask - 1) - # ocean_mask.to_netcdf(self.mom_input_dir + "ocean_mask.nc") - - # ocean_mask = ocean_mask.where(ocean_mask == 1,-1e30) topog["depth"] *= ocean_mask @@ -764,13 +782,13 @@ def bathymetry(self,bathy_path,varnames,fill_channels = False,minimum_layers = 3 subprocess.run("mv topog_deseas.nc topog.nc",shell=True,cwd=self.mom_input_dir) - ## Now run the remaining FRE tools to construct masks based on our topography + ## FINAL STEP: run the remaining FRE tools to construct masks based on our topography args = "--num_tiles 1 --dir . --mosaic_name ocean_mosaic --tile_file hgrid.nc".split(" ") print("MAKE SOLO MOSAIC",subprocess.run( self.toolpath + "make_solo_mosaic/make_solo_mosaic --num_tiles 1 --dir . --mosaic_name ocean_mosaic --tile_file hgrid.nc", - shell=True, - cwd = self.mom_input_dir),sep = "\n\n") + shell=True, + cwd = self.mom_input_dir),sep = "\n\n") @@ -780,25 +798,27 @@ def bathymetry(self,bathy_path,varnames,fill_channels = False,minimum_layers = 3 ,cwd = self.mom_input_dir),sep = "\n\n") self.processor_mask((10,10)) - return + self.topog = topog + return def processor_mask(self,layout): - """ - Just a wrapper for FRE Tools check_mask. User provides processor layout tuple of processing units. - """ - - if "topog.nc" not in os.listdir(self.mom_input_dir): - print("No topography file! Need to run make_bathymetry first") - return - try: - os.remove("mask_table*") ## Removes old mask table so as not to clog up inputdir - except: - pass - print("CHECK MASK" , subprocess.run( - self.toolpath + f"check_mask/check_mask --grid_file ocean_mosaic.nc --ocean_topog topog.nc --layout {layout[0]},{layout[1]} --halo 4", - shell=True, - cwd = self.mom_input_dir)) + """ + Just a wrapper for FRE Tools check_mask. User provides processor layout tuple of processing units. + """ + + if "topog.nc" not in os.listdir(self.mom_input_dir): + print("No topography file! Need to run make_bathymetry first") return + try: + os.remove("mask_table*") ## Removes old mask table so as not to clog up inputdir + except: + pass + print("CHECK MASK" , subprocess.run( + self.toolpath + f"check_mask/check_mask --grid_file ocean_mosaic.nc --ocean_topog topog.nc --layout {layout[0]},{layout[1]} --halo 4", + shell=True, + cwd = self.mom_input_dir)) + self.layout = layout + return