Merge pull request #1 from jr3cermak/esmf840

jsimkins2 · web-flow · commit b1bab6b884be · 2023-01-29T13:34:16.000-05:00
esmf840
diff --git a/setup/CONDA.md b/setup/CONDA.md
@@ -0,0 +1,89 @@
+# Conda Environment Instructions
+
+## Python 3.7.x
+
+1) `conda create -n xesmf_env python=3.7`
+2) `conda activate xesmf_env`
+3) `conda install -c conda-forge xesmf=0.3.0 esmpy=8.2.0 bottleneck=1.3.5`
+4) `conda install -c conda-forge dask=2021.10.0 netcdf4`
+
+## Python 3.9.x
+
+NOTE: This method takes a long time for the resolver to find all the
+right package combinations.
+
+`$ conda create -n xesmf_env_test -c conda-forge xesmf=0.3.0 esmpy=8.2.0 bottleneck=1.3.5 dask=2021.10.0 netcdf4`
+
+Last checked on Jan 26, 2023.
+
+## Python 3.10.x
+
+1) `conda create -n xesmf_env python=3.10.8`
+2) `conda activate xesmf_env`
+3) `conda install -c conda-forge xesmf=0.3.0 esmpy=8.2.0 bottleneck=1.3.5`
+4) `conda install -c conda-forge dask=2021.10.0 netcdf4`
+
+Last checked on Jan 26, 2023.
+
+## Python 3.11.x
+
+There are reported speed improvements to this version of python.
+
+## Jupyter
+
+5) `conda install -c conda-forge jupyter jupyterlab numba nodejs`
+
+## Flooding
+
+If flooding is required:
+
+6) Install HCtFlood
+
+```
+git clone https://github.com/raphaeldussin/HCtFlood
+cd src/HCtFlood
+python -m pip install -e .
+```
+
+If HCtFlood requires more iterations to converge on a solution
+here is the needed modification:
+```
+diff --git a/HCtFlood/kara.py b/HCtFlood/kara.py
+index 539050b..00201f0 100644
+--- a/HCtFlood/kara.py
++++ b/HCtFlood/kara.py
+@@ -106,7 +106,7 @@ def flood_kara_ma(masked_array, spval=1e+15):
+
+
+ @njit
+-def flood_kara_raw(field, mask, nmax=1000):
++def flood_kara_raw(field, mask, nmax=2000):
+     """Extrapolate land values onto land using the kara method
+     (https://doi.org/10.1175/JPO2984.1)
+```
+
+# Notes
+
+- ESMF packages must have mpi support, see below.
+- You must use `xesmf=0.3.0` to be able to create and reuse weight files.
+- There is also a problem with a later version of dask, recommend `dask=2021.10.0`
+
+After the conda packages are installed, run: `conda list | grep mpi`, the following
+packages should be similar in appearance:
+
+```
+esmf                      8.2.0           mpi_mpich_h4975321_100    conda-forge
+esmpy                     8.2.0           mpi_mpich_py37h7352969_101    conda-forge
+fftw                      3.3.10          nompi_h77c792f_102    conda-forge
+hdf5                      1.12.1          mpi_mpich_h9c45103_0    conda-forge
+libnetcdf                 4.8.1           mpi_mpich_h319fa22_1    conda-forge
+mpi                       1.0                       mpich    conda-forge
+mpi4py                    3.1.3            py37h52370cb_1    conda-forge
+mpich                     4.0.2              h846660c_100    conda-forge
+netcdf-fortran            4.5.4           mpi_mpich_h1364a43_0    conda-forge
+netcdf4                   1.5.8           nompi_py37hf784469_101    conda-forge
+```
+
+It is very important that esmf, esmpy, libnetcdf, hdf5 and netcdf-fortran have
+`mpi_mpich` within the build name (3rd column) of the package listing.  
+If they show up as `nompi` then ESMF will not work.
diff --git a/setup/README.md b/setup/README.md
@@ -1,15 +1,43 @@
 # Conda Environment Instructions
 
-These scripts have been run successfully on Antares using a Python 3.7 conda environment located here - `/home/james/anaconda3/envs/xesmf/bin`. Here is a recipe for creating this. 
+These scripts have been run successfully using a Python 3.7 conda environment.
+Here is a recipe for creating this. 
 
 1) `conda create -n xesmf_env python=3.7`
 2) `conda activate xesmf_env`
 3) `conda install -c conda-forge xesmf esmpy=8.2.0`
-4) `conda install -c conda-forge dask netCDF4`
+4) `conda install -c conda-forge dask netcdf4`
 
+Please refer to [CONDA](CONDA.md) link for expanded information for installation of
+esmf, esmpy and xesmf for python.
 
+## Notes
 
-Notes:
+- Note that `esmpy=8.2.0` must be [installed in the same instance](https://github.com/JiaweiZhuang/xESMF/issues/47#issuecomment-665516640) of `xesmf` installation.  The packages must be built with MPI (either `mpich` or `openmpi`)
 
-- Note that `esmpy=8.2.0` must be [installed in the same instance](https://github.com/JiaweiZhuang/xESMF/issues/47#issuecomment-665516640) of `xesmf` installation.
-- If you're running on antares, my environment for this can be found at `/home/james/anaconda3/envs/xesmf`
+## MPI
+
+To ensure the installed conda packages are installed with MPI, look for `mpich` or `openmpi` in the `Build`
+column.  Here is some examples using `conda search esmpy`:
+
+```
+# Name                       Version           Build              Channel
+esmpy                          8.4.0 mpi_mpich_py311hf216de5_101  conda-forge
+esmpy                          8.4.0 mpi_openmpi_py311hcdf3ef6_1  conda-forge
+esmpy                          8.4.0 nompi_py311h8e2db7d_1  conda-forge
+```
+
+In this example, the first two packages are built with MPI since ('mpich` and `openmpi`) appear in 
+the build title.  The last package does not have MPI support and will not work if you attempt to
+utilize operations that require MPI to work.
+
+# Flooding
+
+If you require flooding of OBCs, then you also need to install:
+ * https://github.com/raphaeldussin/HCtFlood
+
+# Site specific notes
+
+## antares
+
+- An example installed environment for this can be found at `/home/james/anaconda3/envs/xesmf`
diff --git a/setup/boundary/boundary.py b/setup/boundary/boundary.py
@@ -299,30 +299,43 @@ def __init__(self, num, border, hgrid, in_degrees=True, output_dir='.', regrid_d
 
     @property
     def coords(self):
+        # Rename nxp and nyp to locations
         if self.border == 'south':
-            return xarray.Dataset({
+            rcoord = xarray.Dataset({
                 'lon': self.hgrid['x'].isel(nyp=0),
                 'lat': self.hgrid['y'].isel(nyp=0),
                 'angle': self.hgrid['angle_dx'].isel(nyp=0)
             })
+            rcoord = rcoord.rename_dims({'nxp': 'locations'})
         elif self.border == 'north':
-            return xarray.Dataset({
+            rcoord = xarray.Dataset({
                 'lon': self.hgrid['x'].isel(nyp=-1),
                 'lat': self.hgrid['y'].isel(nyp=-1),
                 'angle': self.hgrid['angle_dx'].isel(nyp=-1)
             })
+            rcoord = rcoord.rename_dims({'nxp': 'locations'})
         elif self.border == 'west':
-            return xarray.Dataset({
+            rcoord = xarray.Dataset({
                 'lon': self.hgrid['x'].isel(nxp=0),
                 'lat': self.hgrid['y'].isel(nxp=0),
                 'angle': self.hgrid['angle_dx'].isel(nxp=0)
             })
+            rcoord = rcoord.rename_dims({'nyp': 'locations'})
         elif self.border == 'east':
-            return xarray.Dataset({
+            rcoord = xarray.Dataset({
                 'lon': self.hgrid['x'].isel(nxp=-1),
                 'lat': self.hgrid['y'].isel(nxp=-1),
                 'angle': self.hgrid['angle_dx'].isel(nxp=-1)
             })
+            rcoord = rcoord.rename_dims({'nyp': 'locations'})
+
+        # Make lat and lon coordinates
+        rcoord = rcoord.assign_coords(
+            lat=rcoord['lat'],
+            lon=rcoord['lon']
+        )
+
+        return rcoord
 
     @property
     def nx(self):
@@ -496,18 +509,14 @@ def regrid_velocity(
         udest = uregrid(usource)
         vdest = vregrid(vsource)
 
-        # if lat and lon are variables in u/vsource, u/vdest will be dataset
         if isinstance(udest, xarray.Dataset):
             udest = udest.to_array().squeeze()
         if isinstance(vdest, xarray.Dataset):
             vdest = vdest.to_array().squeeze()
 
         # Rotate velocities to be model-relative.
         if rotate:
-            if self.border in ['south', 'north']:
-                angle = self.coords['angle'].rename({'nxp': 'locations'})
-            elif self.border in ['west', 'east']:
-                angle = self.coords['angle'].rename({'nyp': 'locations'})
+            angle = self.coords['angle']
             udest, vdest = rotate_uv(udest, vdest, angle)
 
         ds_uv = xarray.Dataset({
@@ -582,7 +591,12 @@ def regrid_tracer(
         tdest = regrid(tsource)
 
         if not isinstance(tdest, xarray.Dataset):
-            tdest = tdest.to_dataset()
+            # If tdest has a name use it,
+            # otherwise use the tsource.name
+            if tdest.name:
+                tdest = tdest.to_dataset()
+            else:
+                tdest = tdest.to_dataset(name=tsource.name)
             
         if 'z' in tsource.coords:
             tdest = fill_missing(tdest, fill=fill)
diff --git a/setup/boundary/validate_OBC.py b/setup/boundary/validate_OBC.py
@@ -0,0 +1,41 @@
+#!/usr/bin/env python
+
+import os
+import numpy as np
+import xarray as xr
+
+obc_dir = '/home/cermak/workdir/configs/nwa25/OBC'
+f1 = os.path.join(obc_dir, 'indiv1/uv_001_1996.nc')
+f2 = os.path.join(obc_dir, 'indiv.ans/uv_001_1996.nc')
+
+ds1 = xr.open_dataset(f1)
+ds2 = xr.open_dataset(f2)
+
+allvars = list(ds1.variables)
+coords = list(ds1.coords)
+onlyvars = set(allvars) - set(coords)
+
+notEqual = {
+        'count': 0,
+        'coord': [],
+        'variable': [],
+}
+
+for i in coords:
+    if not(np.all(ds1[i] == ds2[i])):
+        notEqual['coord'].append(i)
+        notEqual['count'] = notEqual['count'] + 1
+
+for i in onlyvars:
+    if not(np.all(ds1[i] == ds2[i])):
+        notEqual['variable'].append(i)
+        notEqual['count'] = notEqual['count'] + 1
+
+if notEqual['count'] > 0:
+    print("Some items do not match:")
+    if len(notEqual['coord']) > 0:
+        print("  Coordinates:\n    %s" % (", ".join(notEqual['coord'])))
+    if len(notEqual['variable']) > 0:
+        print("  Variables:\n    %s" % (", ".join(notEqual['variable'])))
+else:
+    print("NetCDF files are equivalent.")
diff --git a/setup/boundary/write_glorys_boundary.py b/setup/boundary/write_glorys_boundary.py
@@ -0,0 +1,96 @@
+from glob import glob
+from subprocess import run
+from os import path
+import xarray
+from boundary import Segment
+import os
+import numpy
+
+# xarray gives a lot of unnecessary warnings
+import warnings
+warnings.filterwarnings('ignore')
+
+
+def write_year(year, glorys_dir, segments, is_first_year=False):
+    glorys = (
+        xarray.open_mfdataset(sorted(glob(os.path.join(glorys_dir, f'{year}/GLORYS_REANALYSIS_{year}-*.nc'))))
+        .rename({'latitude': 'lat', 'longitude': 'lon', 'depth': 'z'})
+    )
+
+    # Floor first time down to midnight so that it matches initial conditions
+    if is_first_year:
+       glorys.time.data[0] = numpy.datetime64(f'{year}-01-01T00:00:00.000000000')
+    for seg in segments:
+        seg.regrid_velocity(glorys['uo'], glorys['vo'], suffix=year, flood=False)
+        for tr in ['thetao', 'so']:
+            seg.regrid_tracer(glorys[tr], suffix=year, flood=False)
+        seg.regrid_tracer(glorys['zos'], suffix=year, flood=False)
+
+
+# this is an xarray based way to concatenate the obc yearly files into one file (per variable of output)
+# the former way to do this was based on ncrcat from NCO tools
+def ncrcat_rename(nsegments, ncrcat_outdir,output_dir, delete_old_files=False):
+    rename_dict={'thetao':'temp', 'so':'salt', 'zos':'zeta', 'uv':'uv'}
+    for var in ['thetao', 'so', 'zos', 'uv']:
+        for seg in range(1, nsegments+1):
+            comb = xarray.open_mfdataset(f'{output_dir}{var}_00{seg}_*')
+            if var!='uv':
+                comb=comb.rename({f'{var}_segment_00{seg}':f'{rename_dict[var]}_segment_00{seg}'})
+                if var!='zos':
+                    comb=comb.rename({f'dz_{var}_segment_00{seg}':f'dz_{rename_dict[var]}_segment_00{seg}'})
+            # Fix output metadata, including removing all _FillValues.
+            all_vars = list(comb.data_vars.keys()) + list(comb.coords.keys())
+            encodings = {v: {'_FillValue': None} for v in all_vars}
+            encodings['time'].update({'dtype':'float64', 'calendar': 'gregorian'})
+            year1 = str(comb.time.values[2])[0:4]
+            year_end = str(comb.time.values[-2])[0:4]
+            print(year1,year_end)
+            comb.to_netcdf(f'{ncrcat_outdir}{rename_dict[var]}_00{seg}_{year1}_{year_end}.nc',
+                           encoding=encodings,
+                           unlimited_dims='time',
+                           format='NETCDF3_64BIT')
+            print(f'concatenated and renamed {rename_dict[var]}_00{seg}.nc')
+            del(comb)
+            if delete_old_files==True:
+                os.remove(f'{output_dir}{var}_00{seg}_*')
+
+
+def main():
+    first_year = 1996
+
+    # Original
+    #glorys_dir = '/Volumes/A1/workdir/james/glorys/'
+    #output_dir = '/Volumes/A1/workdir/james/nwa25_input/boundary/indiv_years/'
+    #ncrcat_outdir = '/Volumes/A1/workdir/james/nwa25_input/boundary/boundary_final/'
+    #hgrid = xarray.open_dataset('/home/james/gridInfo/nwa25/ocean_hgrid.nc')    
+    
+    # Rob
+    glorys_dir = '/Volumes/A1/workdir/james/glorys/'
+    output_dir = '/home/cermak/workdir/configs/nwa25/OBC/indiv/'
+    ncrcat_outdir = '/home/cermak/workdir/configs/nwa25/OBC/final/'
+    hgrid = xarray.open_dataset('/home/cermak/workdir/configs/nwa25/INPUT/ocean_hgrid.nc')
+    
+    segments = [
+        Segment(1, 'south', hgrid, output_dir=output_dir),
+        Segment(2, 'east', hgrid, output_dir=output_dir),
+        Segment(3, 'north', hgrid, output_dir=output_dir)
+    ]
+
+    for y in range(1996, 2018):
+        print(y)
+        write_year(y, glorys_dir, segments, is_first_year=y==first_year)
+    
+    # Original
+    #output_dir = '/Volumes/A1/workdir/james/nwa25_input/boundary/indiv_years/'
+    #ncrcat_outdir = '/Volumes/A1/workdir/james/nwa25_input/boundary/boundary_final/'
+    # Rob
+    output_dir = '/home/cermak/workdir/configs/nwa25/OBC/indiv/'
+    ncrcat_outdir = '/home/cermak/workdir/configs/nwa25/OBC/final/'
+
+    ncrcat_rename(3, ncrcat_outdir, output_dir)
+
+
+
+if __name__ == '__main__':
+    main()
+
diff --git a/setup/boundary/write_glorys_boundary_verify.py b/setup/boundary/write_glorys_boundary_verify.py
