Read in ak, bk coefficients (#36)

* initial changes to read in ak bk * read ak/bk * add xfail * remove input dir * further changes to unit tests * finish up test * add history * commit uncommited files * fix test comment * add input to top * read in data * read in netcdf file in eta mod * remove txt file * test * modify test and fix generate.py * remove emacs backup file * driver tests pass * fix helper.py * fix fv3core tests * fix physics test * fix grid tests * nullcommconfig * cleanup input * remove driver input * remove top level input * fix circular import problems * modify eta_file readin for test_restart_serial * comment out 91 test * rm safety checks * revert diagnostics.py * restore driver.py * revert initialization.py * restore state.py * restore analytic_init.py * restore init_utils.py and analytic_init.py * restore c_sw.py * d2a2c_vect.py * restore fv3core/stensils * restore translate_fvdynamics * restore physics/stencils * restore stencils * remove circular dependency * use pytest parametrize * cleanup generation.py * fstrinngs * add eta_file to MetricTerm init * remove eta_file argument in new_from_metric_terms and geos_wrapper * use pytest parametrize for the xfail tests * use pytest parametrize for the xfail tests * fix geos_wrapper and grid * fix tests * fstring * add test comments * fix util/HISTORY.md * fix comments * remove __init__.py from tests/main/grid * add jupyter notebooks to generate eta files * generate ak,bk,ptop on metricterm init * fix tests * exploit np.all in eta mod * remove tests/main/grid/input * update ci * test * remove input * edit ci yaml * remove push --------- Co-authored-by: mlee03 <[email protected]>
NOAA-GFDL · Jan 30, 2024 · 34eeea4 · 34eeea4
1 parent 0c68c5f
commit 34eeea4
Show file tree

Hide file tree

Showing 19 changed files with 484 additions and 850 deletions.
diff --git a/.github/workflows/main_unit_tests.yml b/.github/workflows/main_unit_tests.yml
@@ -22,6 +22,10 @@ jobs:
           run: |
             python -m pip install --upgrade pip
             pip install -r requirements_dev.txt
+        - name: Clone datafiles
+          run: |
+            mkdir -p tests/main/input && cd tests/main/input
+            git clone -b store_files https://github.com/mlee03/pace.git tmp && mv tmp/*.nc . && rm -rf tmp
         - name: Run all main tests
           run: |
             pytest -x tests/main
diff --git a/CONTRIBUTORS.md b/CONTRIBUTORS.md
@@ -11,6 +11,7 @@ List format (alphabetical order):  Surname, Name. Employer/Affiliation
 * Fuhrer, Oliver. Allen Institute for AI.
 * George, Rhea. Allen Institute for AI.
 * Harris, Lucas. GFDL.
+* Lee, Mi Kyung. GFDL.
 * Kung, Chris. NASA.
 * McGibbon, Jeremy. Allen Institute for AI.
 * Niedermayr, Yannick. ETH.

diff --git a/driver/examples/configs/baroclinic_c12.yaml b/driver/examples/configs/baroclinic_c12.yaml
@@ -94,3 +94,8 @@ physics_config:
   hydrostatic: false
   nwat: 6
   do_qa: true
+
+grid_config:
+  type:  generated
+  config:
+    eta_file: 'tests/main/input/eta79.nc'
diff --git a/driver/examples/configs/baroclinic_c12_write_restart.yaml b/driver/examples/configs/baroclinic_c12_write_restart.yaml
@@ -92,3 +92,8 @@ physics_config:
   hydrostatic: false
   nwat: 6
   do_qa: true
+
+grid_config:
+  type:  generated
+  config:
+    eta_file: "tests/main/input/eta79.nc"
diff --git a/driver/pace/driver/grid.py b/driver/pace/driver/grid.py
@@ -99,6 +99,7 @@ class GeneratedGridConfig(GridInitializer):
     dx_const: Optional[float] = 1000.0
     dy_const: Optional[float] = 1000.0
     deglat: Optional[float] = 15.0
+    eta_file: str = "None"
 
     def get_grid(
         self,
@@ -112,6 +113,7 @@ def get_grid(
             dx_const=self.dx_const,
             dy_const=self.dy_const,
             deglat=self.deglat,
+            eta_file=self.eta_file,
         )
         if self.stretch_factor != 1:  # do horizontal grid transformation
             _transform_horizontal_grid(

diff --git a/examples/notebooks/generate_eta_file_netcdf.ipynb b/examples/notebooks/generate_eta_file_netcdf.ipynb
@@ -0,0 +1,148 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2c056479",
+   "metadata": {
+    "lines_to_next_cell": 0
+   },
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8c96fbff",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import netCDF4 as nc\n",
+    "import numpy as np"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6827b1b5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\"\"\"\n",
+    "This notebook uses the python netCDF4 module\n",
+    "to create an eta_file containg\n",
+    "ak and bk coefficients for km=79\n",
+    "\"\"\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "45d4a704",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#create a Dataset instance\n",
+    "coefficients = nc.Dataset(\"eta79.nc\", \"w\", format=\"NETCDF4\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b964a014",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#Set dimensionsion\n",
+    "km = coefficients.createDimension(\"km\", 80)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d51c395f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#Create ak and bk variables\n",
+    "ak = coefficients.createVariable(\"ak\", np.float64, (\"km\"))\n",
+    "bk = coefficients.createVariable(\"bk\", np.float64, (\"km\"))\n",
+    "ak.units=\"\"\n",
+    "bk.units=\"\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6723352e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#Assign and write out values of ak\n",
+    "ak[:] = np.array(\n",
+    "    [  3.000000e+02,  6.467159e+02,  1.045222e+03,  1.469188e+03,  1.897829e+03,\n",
+    "       2.325385e+03,  2.754396e+03,  3.191294e+03,  3.648332e+03,  4.135675e+03,\n",
+    "       4.668282e+03,  5.247940e+03,  5.876271e+03,  6.554716e+03,  7.284521e+03,\n",
+    "       8.066738e+03,  8.902188e+03,  9.791482e+03,  1.073499e+04,  1.162625e+04,\n",
+    "       1.237212e+04,  1.299041e+04,  1.349629e+04,  1.390277e+04,  1.422098e+04,\n",
+    "       1.446058e+04,  1.462993e+04,  1.473633e+04,  1.478617e+04,  1.478511e+04,\n",
+    "       1.473812e+04,  1.464966e+04,  1.452370e+04,  1.436382e+04,  1.417324e+04,\n",
+    "       1.395491e+04,  1.371148e+04,  1.344540e+04,  1.315890e+04,  1.285407e+04,\n",
+    "       1.253280e+04,  1.219685e+04,  1.184788e+04,  1.148739e+04,  1.111682e+04,\n",
+    "       1.073748e+04,  1.035062e+04,  9.957395e+03,  9.558875e+03,  9.156069e+03,\n",
+    "       8.749922e+03,  8.341315e+03,  7.931065e+03,  7.519942e+03,  7.108648e+03,\n",
+    "       6.698281e+03,  6.290007e+03,  5.884984e+03,  5.484372e+03,  5.089319e+03,\n",
+    "       4.700960e+03,  4.320421e+03,  3.948807e+03,  3.587201e+03,  3.236666e+03,\n",
+    "       2.898237e+03,  2.572912e+03,  2.261667e+03,  1.965424e+03,  1.685079e+03,\n",
+    "       1.421479e+03,  1.175419e+03,  9.476516e+02,  7.388688e+02,  5.497130e+02,\n",
+    "       3.807626e+02,  2.325417e+02,  1.054810e+02, -8.381903e-04,  0.000000e+00] )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "195c9ef5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#Assign and write out values of bk \n",
+    "bk[:] = np.array(\n",
+    "  [ 0.,         0.,         0.,         0.,         0.,         0.,\n",
+    "    0.,         0.,         0.,         0.,         0.,         0.,\n",
+    "    0.,         0.,         0.,         0.,         0.,         0.,\n",
+    "    0.,         0.00106595, 0.00412866, 0.00900663, 0.01554263, 0.02359921,\n",
+    "    0.03305481, 0.0438012,  0.05574095, 0.06878554, 0.08285347, 0.09786981,\n",
+    "    0.1137643,  0.130471,   0.1479275,  0.1660746,  0.1848558,  0.2042166,\n",
+    "    0.2241053,  0.2444716,  0.2652672,  0.286445,   0.3079604,  0.3297701,\n",
+    "    0.351832,   0.3741062,  0.3965532,  0.4191364,  0.4418194,  0.4645682,\n",
+    "    0.48735,    0.5101338,  0.5328897,  0.5555894,  0.5782067,  0.6007158,\n",
+    "    0.6230936,  0.6452944,  0.6672683,  0.6889648,  0.7103333,  0.7313231,\n",
+    "    0.7518838,  0.7719651,  0.7915173,  0.8104913,  0.828839,   0.846513,\n",
+    "    0.8634676,  0.8796583,  0.8950421,  0.9095779,  0.9232264,  0.9359506,\n",
+    "    0.9477157,  0.9584892,  0.9682413,  0.9769447,  0.9845753,  0.9911126,\n",
+    "    0.9965372,  1.        ] )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c0f3bd9d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#Close netcdf file\n",
+    "coefficients.close()"
+   ]
+  }
+ ],
+ "metadata": {
+  "jupytext": {
+   "cell_metadata_filter": "-all",
+   "executable": "/usr/bin/env python3",
+   "main_language": "python",
+   "notebook_metadata_filter": "-all"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/examples/notebooks/generate_eta_file_xarray.ipynb b/examples/notebooks/generate_eta_file_xarray.ipynb
@@ -0,0 +1,140 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6dc5fe4c",
+   "metadata": {
+    "lines_to_next_cell": 0
+   },
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "81be9a15",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import numpy as np\n",
+    "import xarray as xr"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c74c6c07",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\"\"\"\n",
+    "This notebook uses the python xarray module\n",
+    "to create an eta_file containg\n",
+    "ak and bk coefficients for km=79\n",
+    "\"\"\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f72c5d5b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#Assign ak data\n",
+    "ak=np.array(\n",
+    "    [ 3.000000e+02,  6.467159e+02,  1.045222e+03,  1.469188e+03,  1.897829e+03,\n",
+    "      2.325385e+03,  2.754396e+03,  3.191294e+03,  3.648332e+03,  4.135675e+03,\n",
+    "      4.668282e+03,  5.247940e+03,  5.876271e+03,  6.554716e+03,  7.284521e+03,\n",
+    "      8.066738e+03,  8.902188e+03,  9.791482e+03,  1.073499e+04,  1.162625e+04,\n",
+    "      1.237212e+04,  1.299041e+04,  1.349629e+04,  1.390277e+04,  1.422098e+04,\n",
+    "      1.446058e+04,  1.462993e+04,  1.473633e+04,  1.478617e+04,  1.478511e+04,\n",
+    "      1.473812e+04,  1.464966e+04,  1.452370e+04,  1.436382e+04,  1.417324e+04,\n",
+    "      1.395491e+04,  1.371148e+04,  1.344540e+04,  1.315890e+04,  1.285407e+04,\n",
+    "      1.253280e+04,  1.219685e+04,  1.184788e+04,  1.148739e+04,  1.111682e+04,\n",
+    "      1.073748e+04,  1.035062e+04,  9.957395e+03,  9.558875e+03,  9.156069e+03,\n",
+    "      8.749922e+03,  8.341315e+03,  7.931065e+03,  7.519942e+03,  7.108648e+03,\n",
+    "      6.698281e+03,  6.290007e+03,  5.884984e+03,  5.484372e+03,  5.089319e+03,\n",
+    "      4.700960e+03,  4.320421e+03,  3.948807e+03,  3.587201e+03,  3.236666e+03,\n",
+    "      2.898237e+03,  2.572912e+03,  2.261667e+03,  1.965424e+03,  1.685079e+03,\n",
+    "      1.421479e+03,  1.175419e+03,  9.476516e+02,  7.388688e+02,  5.497130e+02,\n",
+    "      3.807626e+02,  2.325417e+02,  1.054810e+02, -8.381903e-04,  0.000000e+00] )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "f5b85c7e",
+   "metadata": {
+    "lines_to_next_cell": 2
+   },
+   "outputs": [],
+   "source": [
+    "#Assign bk data\n",
+    "bk=np.array(\n",
+    "    [ 0.,         0.,         0.,         0.,         0.,         0.,\n",
+    "      0.,         0.,         0.,         0.,         0.,         0.,\n",
+    "      0.,         0.,         0.,         0.,         0.,         0.,\n",
+    "      0.,         0.00106595, 0.00412866, 0.00900663, 0.01554263, 0.02359921,\n",
+    "      0.03305481, 0.0438012,  0.05574095, 0.06878554, 0.08285347, 0.09786981,\n",
+    "      0.1137643,  0.130471,   0.1479275,  0.1660746,  0.1848558,  0.2042166,\n",
+    "      0.2241053,  0.2444716,  0.2652672,  0.286445,   0.3079604,  0.3297701,\n",
+    "      0.351832,   0.3741062,  0.3965532,  0.4191364,  0.4418194,  0.4645682,\n",
+    "      0.48735,    0.5101338,  0.5328897,  0.5555894,  0.5782067,  0.6007158,\n",
+    "      0.6230936,  0.6452944,  0.6672683,  0.6889648,  0.7103333,  0.7313231,\n",
+    "      0.7518838,  0.7719651,  0.7915173,  0.8104913,  0.828839,   0.846513,\n",
+    "      0.8634676,  0.8796583,  0.8950421,  0.9095779,  0.9232264,  0.9359506,\n",
+    "      0.9477157,  0.9584892,  0.9682413,  0.9769447,  0.9845753,  0.9911126,\n",
+    "      0.9965372,  1.        ] )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c5450f7f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#Create a Dataset instance\n",
+    "coefficients = xr.Dataset(\n",
+    "    { \"ak\": ([\"km1\"], ak),\n",
+    "      \"bk\": ([\"km1\"], bk) \n",
+    "    })"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5a0e5487",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#Set attributes for each variable\n",
+    "coefficients[\"ak\"].attrs[\"units\"]=\"\"\n",
+    "coefficients[\"bk\"].attrs[\"units\"]=\"\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "612b0134",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#Write netcdf file\n",
+    "coefficients.to_netcdf(\"eta79.nc\")"
+   ]
+  }
+ ],
+ "metadata": {
+  "jupytext": {
+   "cell_metadata_filter": "-all",
+   "executable": "/usr/bin/env python3",
+   "main_language": "python",
+   "notebook_metadata_filter": "-all"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/fv3core/pace/fv3core/wrappers/geos_wrapper.py b/fv3core/pace/fv3core/wrappers/geos_wrapper.py
@@ -124,7 +124,9 @@ def __init__(
 
         # set up the metric terms and grid data
         metric_terms = pace.util.grid.MetricTerms(
-            quantity_factory=quantity_factory, communicator=self.communicator
+            quantity_factory=quantity_factory,
+            communicator=self.communicator,
+            eta_file=namelist["grid_config"]["config"]["eta_file"],
         )
         grid_data = pace.util.grid.GridData.new_from_metric_terms(metric_terms)
 

diff --git a/tests/main/driver/test_restart_fortran.py b/tests/main/driver/test_restart_fortran.py
@@ -47,7 +47,9 @@ def test_state_from_fortran_restart():
         damping_coefficients,
         driver_grid_data,
         grid_data,
-    ) = pace.driver.GeneratedGridConfig(restart_path=restart_dir).get_grid(
+    ) = pace.driver.GeneratedGridConfig(
+        restart_path=restart_dir, eta_file=restart_dir + "/fv_core.res.nc"
+    ).get_grid(
         quantity_factory, null_communicator
     )
 

diff --git a/tests/main/driver/test_restart_serial.py b/tests/main/driver/test_restart_serial.py
@@ -66,11 +66,14 @@ def test_restart_save_to_disk():
             sizer=sizer, backend=backend
         )
 
+        eta_file = driver_config.grid_config.config.eta_file
         (
             damping_coefficients,
             driver_grid_data,
             grid_data,
-        ) = pace.driver.GeneratedGridConfig().get_grid(quantity_factory, communicator)
+        ) = pace.driver.GeneratedGridConfig(eta_file=eta_file).get_grid(
+            quantity_factory, communicator
+        )
         init = AnalyticInit()
         driver_state = init.get_driver_state(
             quantity_factory=quantity_factory,

diff --git a/tests/main/fv3core/test_cartesian_grid.py b/tests/main/fv3core/test_cartesian_grid.py
@@ -8,7 +8,7 @@
 
 @pytest.mark.parametrize("npx", [8])
 @pytest.mark.parametrize("npy", [8])
-@pytest.mark.parametrize("npz", [1])
+@pytest.mark.parametrize("npz", [79])
 @pytest.mark.parametrize("dx_const", [1e2, 1e3])
 @pytest.mark.parametrize("dy_const", [2e2, 3e3])
 @pytest.mark.parametrize("deglat", [0.0, 15.0])
@@ -35,6 +35,7 @@ def test_cartesian_grid_generation(
         dx_const=dx_const,
         dy_const=dy_const,
         deglat=deglat,
+        eta_file="tests/main/input/eta79.nc",
     )
     assert np.all(grid_generator.lat_agrid.data == deglat * PI / 180.0)
     assert np.all(grid_generator.lon_agrid.data == 0.0)