diff --git a/pyschism/mesh/vgrid.py b/pyschism/mesh/vgrid.py index d7dd7dfa..8b5e6480 100644 --- a/pyschism/mesh/vgrid.py +++ b/pyschism/mesh/vgrid.py @@ -9,6 +9,8 @@ from pyschism.mesh.hgrid import Hgrid +from matplotlib.pyplot import * + def C_of_sigma(sigma, theta_b, theta_f): assert theta_b <= 0. and theta_b <= 1. @@ -106,8 +108,16 @@ def is3D(self): class LSC2(Vgrid): - def __init__(self, sigma): - self.sigma = sigma + def __init__(self, hsm, nv, h_c, theta_b, theta_f): + self.hsm = np.array(hsm) + self.nv = np.array(nv) + self.h_c = h_c + self.theta_b = theta_b + self.theta_f = theta_f + self.m_grid = None + self._znd = None + self._snd = None + self._nlayer = None def __str__(self): f = [ @@ -131,12 +141,147 @@ def __str__(self): return '\n'.join(f) def get_xyz(self, gr3, crs=None): + if type(gr3) == Hgrid: + gr3 = gr3 + else: + gr3=Hgrid.open(gr3) xy = gr3.get_xy(crs) z = gr3.values[:, None]*self.sigma x = np.tile(xy[:, 0], (z.shape[1],)) y = np.tile(xy[:, 0], (z.shape[1],)) return np.vstack([x, y, z.flatten()]).T + def calc_m_grid(self): + ''' + create master grid + Adapted from: + https://github.com/wzhengui/pylibs/blob/master/pyScripts/gen_vqs.py + ''' + if self.m_grid: + pass + else: + z_mas=np.ones([self.nhm,self.nvrt])*np.nan; eta=0.0 + for m, [hsmi,nvi] in enumerate(zip(self.hsm,self.nv)): + #strethcing funciton + hc=min(hsmi,self.h_c) + for k in np.arange(nvi): + sigma= k/(1-nvi) #zi=-sigma #original sigma coordiante + #compute zcoordinate + cs=(1-self.theta_b)*np.sinh(self.theta_f*sigma)/np.sinh(self.theta_f)+\ + self.theta_b*(np.tanh(self.theta_f*(sigma+0.5))-\ + np.tanh(self.theta_f*0.5))/2/np.tanh(self.theta_f*0.5) + z_mas[m,k]=eta*(1+sigma)+hc*sigma+(hsmi-hc)*cs + + #normalize z_mas + z_mas[m]=-(z_mas[m]-z_mas[m,0])*hsmi/(z_mas[m,nvi-1]-z_mas[m,0]) + s_mas=np.array([z_mas[i]/self.hsm[i] for i in np.arange(self.nhm)]) + + self.m_grid = z_mas + + def make_m_plot(self): + ''' + plot master grid + Adapted from: + https://github.com/wzhengui/pylibs/blob/master/pyScripts/gen_vqs.py + ''' + #check master grid + for i in np.arange(self.nhm-1): + if min(self.m_grid[i,:self.nv[i]]-self.m_grid[i+1,:self.nv[i]])<0: \ + print('check: master grid layer={}, hsm={}, nv={}'.\ + format(i+1,self.hsm[i+1],self.nv[i+1])) + + #plot master grid + figure(figsize=[10,5]) + for i in np.arange(self.nhm): plot(i*np.ones(self.nvrt),\ + self.m_grid[i],'k-',lw=0.3) + for k in np.arange(self.nvrt): plot(np.arange(self.nhm),\ + self.m_grid.T[k],'k-',lw=0.3) + setp(gca(),xlim=[-0.5,self.nhm-0.5],ylim=[-self.hsm[-1],0.5]) + gcf().tight_layout() + + def calc_lsc2_att(self, gr3, crs=None): + ''' + master grid to lsc2: + compute vertical layers at nodes + gr3 is either file or Hgrid Object + Adapted from: + https://github.com/wzhengui/pylibs/blob/master/pyScripts/gen_vqs.py + ''' + if type(gr3) == Hgrid: + gd = gr3 + else: + gd=Hgrid.open(gr3) + dp = gd.values*-1 + fpz=dpself.hsm[m-1])*(dp<=self.hsm[m]) + ind1[fp]=m-1 + ind2[fp]=m + rat[fp]=(dp[fp]-self.hsm[m-1])/(self.hsm[m]-self.hsm[m-1]) + nlayer[fp]=self.nv[m] + + #Find the last non NaN node and fills the NaN values with it + last_non_nan = (~np.isnan(self.m_grid)).cumsum(1).argmax(1) + z_mas=np.array([np.nan_to_num(z_mas_arr,nan=z_mas_arr[last_non_nan[i]])\ + for i, z_mas_arr in enumerate(self.m_grid)]) + znd=z_mas[ind1]*(1-rat[:,None])+z_mas[ind2]*rat[:,None]; #z coordinate + for i in np.arange(len(gd.nodes.id)): + znd[i,nlayer[i]-1]=-dp[i] + znd[i,nlayer[i]:]=np.nan + snd=znd/dp[:,None]; #sigma coordinate + + #check vgrid + for i in np.arange(len(gd.nodes.id)): + for k in np.arange(self.nvrt-1): + if znd[i,k]<=znd[i,k+1]: + raise TypeError(f'wrong vertical layers') + + self._znd = znd + self._snd = snd + self._nlayer = nlayer + + + def write(self, path, overwrite=False): + ''' + write mg2lsc2 into vgrid.in + ''' + path = pathlib.Path(path) + if path.is_file() and not overwrite: + raise Exception( + 'File exists, pass overwrite=True to allow overwrite.') + + with open(path, 'w') as fid: + fid.write(' 1 !average # of layers={:0.2f}\n {} !nvrt\n'.format\ + (np.mean(self._nlayer),self.nvrt)) + bli=[]#bottom level index + for i in np.arange(len(self._nlayer)): + nlayeri=self._nlayer[i]; si=np.flipud(self._snd[i,:nlayeri]) + bli.append(self.nvrt-nlayeri+1) + fstr=f" {self.nvrt-nlayeri+1:2}" + fid.write(fstr) + for i in range(self.nvrt): + fid.write(f'\n {i+1}') + for n,bl in enumerate(bli): + si=np.flipud(self._snd[n]) + if bl <= i+1: + fid.write(f" {si[i]:.6f}") + else: + fid.write(f" {-9.:.6f}") + fid.close() + + @classmethod def open(cls, path): @@ -174,7 +319,11 @@ def open(cls, path): @property def nvrt(self): - return self.sigma.shape[1] + return self.nv[-1] + + @property + def nhm(self): + return self.hsm.shape[0] class SZ(Vgrid): diff --git a/tests/_mesh/test_vgrid.py b/tests/_mesh/test_vgrid.py index 3c0948f3..772b7e3d 100644 --- a/tests/_mesh/test_vgrid.py +++ b/tests/_mesh/test_vgrid.py @@ -3,15 +3,105 @@ import pathlib import unittest from pyschism.mesh import Vgrid - +from pyschism.mesh.hgrid import Hgrid +from pyschism.mesh.vgrid import LSC2 +import numpy as np class VgridTestCase(unittest.TestCase): + def setUp(self): + nodes = { + '1': ((0., 0.), 1.5), + '2': ((.5, 0.), 2.5), + '3': ((1., 0.), 3.5), + '4': ((1., 1.), 0.), + '5': ((0., 1.), 1.), + '6': ((.5, 1.5), -9.), + '7': ((.33, .33), 1.), + '8': ((.66, .33), 2.), + '9': ((.5, .66), 3.), + '10': ((-1., 1.), 4.), + '11': ((-1., 0.), 5.), + } + elements = { + '1': ['5', '7', '9'], + '2': ['1', '2', '7'], + '3': ['2', '3', '8'], + '4': ['8', '7', '2'], + '5': ['3', '4', '8'], + '6': ['4', '9', '8'], + '7': ['4', '6', '5'], + '8': ['5', '10', '11', '1'], + '9': ['9', '4', '5'], + '10': ['5', '1', '7'] + } + # write hgrid + tmpdir = tempfile.TemporaryDirectory() + hgrid = pathlib.Path(tmpdir.name) / 'hgrid.gr3' + with open(hgrid, 'w') as f: + f.write('\n') + f.write(f'{len(elements):d} ') + f.write(f'{len(nodes):d}\n') + for id, ((x, y), z) in nodes.items(): + f.write(f"{id} ") + f.write(f"{x} ") + f.write(f"{y} ") + f.write(f"{z}\n") + for id, geom in elements.items(): + f.write(f"{id} ") + f.write(f"{len(geom)} ") + for idx in geom: + f.write(f"{idx} ") + f.write(f"\n") + + hsm=[1,2,3,6] + nv=[2,4,5,5] + theta_b=0 + theta_f=2.5 + h_c=2 + + self.hgrid=Hgrid.open(hgrid,crs=4326) + self.hsm=hsm + self.nv=nv + self.theta_b=theta_b + self.theta_f=theta_f + self.h_c=h_c def test_init(self): v = Vgrid() v.boilerplate_2D self.assertIsInstance(v, Vgrid) + def test_LSC2(self): + lsc2_obj = LSC2(self.hsm,self.nv,self.h_c,self.theta_b,self.theta_f) + self.assertIsInstance(lsc2_obj, LSC2) + + def test_calc_m_grid(self): + lsc2_obj = LSC2(self.hsm,self.nv,self.h_c,self.theta_b,self.theta_f) + lsc2_obj.calc_m_grid() + self.assertIsInstance(lsc2_obj.m_grid.shape, (4,5)) + + def test_make_m_plot(self): + lsc2_obj = LSC2(self.hsm,self.nv,self.h_c,self.theta_b,self.theta_f) + lsc2_obj.calc_m_grid() + lsc2_obj.make_m_plot() + + def test_calc_lsc2_att(self): + gr3=self.hgrid + lsc2_obj = LSC2(self.hsm,self.nv,self.h_c,self.theta_b,self.theta_f) + lsc2_obj.calc_m_grid() + lsc2_obj.calc_lsc2_att(gr3) + self.assertIsInstance(lsc2_obj._znd.shape,(11,5)) + self.assertIsInstance(lsc2_obj._snd.shape,(11,5)) + self.assertIsInstance(lsc2_obj._nlayer,np.array([4,5,5,2,2,2,2,4,5,5,5])) + + def test_lsc2_write(self): + tmpdir = tempfile.TemporaryDirectory() + gr3=self.hgrid + lsc2_obj = LSC2(self.hsm,self.nv,self.h_c,self.theta_b,self.theta_f) + lsc2_obj.calc_m_grid() + lsc2_obj.calc_lsc2_att(gr3) + lsc2_obj.write(pathlib.Path(tmpdir.name) / 'vgrid_lsc2.in') + def test_write(self): tmpdir = tempfile.TemporaryDirectory() v = Vgrid() diff --git a/tests/test_vgrid.py b/tests/test_vgrid.py new file mode 100644 index 00000000..0c09d557 --- /dev/null +++ b/tests/test_vgrid.py @@ -0,0 +1,267 @@ +#! /usr/bin/env python +import argparse +# from datetime import datetime +import logging +import os +import pathlib +# import shutil +import tarfile +import tempfile +import unittest +import urllib.request +from datetime import datetime, timedelta + +from pyschism.cmd.bctides import BctidesCli +from pyschism.driver import ModelConfig +from pyschism.mesh import Hgrid + + +logging.basicConfig(level=logging.INFO, force=True) + + +class ModelConfigurationTestCase(unittest.TestCase): + + def setUp(self): + hgrid2D = os.getenv('NWM_TEST_MESH') + if hgrid2D is None: + data_directory = pathlib.Path(__file__).parent.absolute() / 'data' + hgrid2D = data_directory / "GulfStreamDevel/hgrid.gr3" + url = "https://www.dropbox.com/s/mjaxaqeggy721um/" + url += "Gulf_Stream_develop.tar.gz?dl=1" + g = urllib.request.urlopen(url) + tmpfile = tempfile.NamedTemporaryFile() + with open(tmpfile.name, 'b+w') as f: + f.write(g.read()) + with tarfile.open(tmpfile.name, "r:gz") as tar: + tar.extractall(data_directory / "GulfStreamDevel") + self.hgrid2D = hgrid2D + + #def _test_bctides_cli_2d(self): + # import tempfile + # tmpdir = tempfile.TemporaryDirectory() + # parser = argparse.ArgumentParser() + # add_bctides(parser.add_subparsers(dest='mode')) + # args = parser.parse_args( + # [ + # 'bctides', + # f'{self.hgrid2D}', # hgrid path + # '--run-days=5', # rndays + # '--hgrid-crs=epsg:4326', + # f'-o={tmpdir.name}', + # '--overwrite' + # ]) + # BctidesCli(args) + # # with open(tmpdir.name + '/bctides.in') as f: + # # for line in f: + # # print(line, end='') + + #def _test_bctides_cli_2d_w_velo(self): + # import tempfile + # tmpdir = tempfile.TemporaryDirectory() + # parser = argparse.ArgumentParser() + # add_bctides(parser.add_subparsers(dest='mode')) + # args = parser.parse_args( + # [ + # 'bctides', + # f'{self.hgrid2D}', # hgrid path + # '--run-days=5', # rndays, + # '--include-velocity', + # '--hgrid-crs=epsg:4326', + # f'-o={tmpdir.name}', + # '--overwrite' + # ]) + # BctidesCli(args) + # # with open(tmpdir.name + '/bctides.in') as f: + # # for line in f: + # # print(line, end='') + + #def test_bctides_cli_2d_w_elev2d(self): + # import tempfile + # tmpdir = tempfile.TemporaryDirectory() + # parser = argparse.ArgumentParser() + # add_bctides(parser.add_subparsers(dest='mode')) + # args = parser.parse_args( + # [ + # 'bctides', + # f'{self.hgrid2D}', # hgrid path + # '--run-days=5', # rndays, + # '--iettype-5', + # # '--ifltype-5', + # '--hgrid-crs=epsg:4326', + # f'-o={tmpdir.name}', + # '--overwrite' + # ]) + # BctidesCli(args) + + #def _test_bctides_cli_3d(self): + # import tempfile + # tmpdir = tempfile.TemporaryDirectory() + # parser = argparse.ArgumentParser() + # add_bctides(parser.add_subparsers(dest='mode')) + # args = parser.parse_args( + # [ + # 'bctides', + # f'{self.hgrid}', # hgrid path + # '--run-days=5', # rndays + # '--hgrid-crs=epsg:4326', + # f'-o={tmpdir.name}', + # '--overwrite' + # ]) + # BctidesCli(args) + + def test_case1_setup(self): + # Inputs: Modify As Needed! + start_date = "2018091000" #in YYYYMMDDHH format + end_date = "2018091600" #in YYYYMMDDHH format + + dt = timedelta(seconds=150.) + nspool = timedelta(minutes=20.) + + start_date = datetime.strptime(start_date, "%Y%m%d%H") + end_date = datetime.strptime(end_date, "%Y%m%d%H") + rnday = end_date - start_date + + hgrid = Hgrid.open( + 'https://raw.githubusercontent.com/geomesh/test-data/main/NWM/hgrid.ll', + crs=4326 + ) + + config = ModelConfig( + hgrid=hgrid, + flags=[[3, 3, 0, 0]], + constituents = 'major', + database = 'tpxo', + ) + + coldstart = config.coldstart( + start_date=start_date, + end_date=start_date + rnday, + timestep=dt, + nspool=timedelta(hours=1), + elev=True, + dahv=True, + ) + + with tempfile.TemporaryDirectory() as dn: + tmpdir = pathlib.Path(dn) + coldstart.write(tmpdir, overwrite=True) + + def test_case2_setup(self): + # Inputs: Modify As Needed! + start_date = "2018091000" #in YYYYMMDDHH format + end_date = "2018091600" #in YYYYMMDDHH format + + dt = timedelta(seconds=150.) + nspool = timedelta(minutes=20.) + + start_date = datetime.strptime(start_date, "%Y%m%d%H") + end_date = datetime.strptime(end_date, "%Y%m%d%H") + rnday = end_date - start_date + + hgrid = Hgrid.open( + 'https://raw.githubusercontent.com/geomesh/test-data/main/NWM/hgrid.ll', + crs=4326 + ) + + config = ModelConfig( + hgrid=hgrid, + flags=[[5, 5, 0, 0]], + constituents = 'major', + database = 'tpxo', + ) + + coldstart = config.coldstart( + start_date=start_date, + end_date=start_date + rnday, + timestep=dt, + nspool=timedelta(hours=1), + elev=True, + dahv=True, + ) + + with tempfile.TemporaryDirectory() as dn: + tmpdir = pathlib.Path(dn) + coldstart.write(tmpdir, overwrite=True) + + def test_case3_setup(self): + # Inputs: Modify As Needed! + start_date = "2018091000" #in YYYYMMDDHH format + end_date = "2018091600" #in YYYYMMDDHH format + + dt = timedelta(seconds=150.) + nspool = timedelta(minutes=20.) + + start_date = datetime.strptime(start_date, "%Y%m%d%H") + end_date = datetime.strptime(end_date, "%Y%m%d%H") + rnday = end_date - start_date + + hgrid = Hgrid.open( + 'https://raw.githubusercontent.com/geomesh/test-data/main/NWM/hgrid.ll', + crs=4326 + ) + + config = ModelConfig( + hgrid=hgrid, + flags=[[5, 5, 2, 2]], #constant temperature and salinity + constituents = 'major', + database = 'tpxo', + tthconst = [10.0], #T value + sthconst = [0.0], #S value + tobc = [0.1], #nuding factor for T + sobc = [0.1], #nuding factor for S + ) + + coldstart = config.coldstart( + start_date=start_date, + end_date=start_date + rnday, + timestep=dt, + nspool=timedelta(hours=1), + elev=True, + dahv=True, + ) + + with tempfile.TemporaryDirectory() as dn: + tmpdir = pathlib.Path(dn) + coldstart.write(tmpdir, overwrite=True) + + def test_case4_setup(self): + # Inputs: Modify As Needed! + start_date = "2018091000" #in YYYYMMDDHH format + end_date = "2018091600" #in YYYYMMDDHH format + + dt = timedelta(seconds=150.) + nspool = timedelta(minutes=20.) + + start_date = datetime.strptime(start_date, "%Y%m%d%H") + end_date = datetime.strptime(end_date, "%Y%m%d%H") + rnday = end_date - start_date + + hgrid = Hgrid.open( + 'https://raw.githubusercontent.com/geomesh/test-data/main/NWM/hgrid.ll', + crs=4326 + ) + + config = ModelConfig( + hgrid=hgrid, + flags=[[5, 5, 4, 4]], #constant temperature and salinity + constituents = 'major', + database = 'tpxo', + tobc = [0.1], #nuding factor for T + sobc = [0.1], #nuding factor for S + ) + + coldstart = config.coldstart( + start_date=start_date, + end_date=start_date + rnday, + timestep=dt, + nspool=timedelta(hours=1), + elev=True, + dahv=True, + ) + + with tempfile.TemporaryDirectory() as dn: + tmpdir = pathlib.Path(dn) + coldstart.write(tmpdir, overwrite=True) + +if __name__ == '__main__': + unittest.main()