Merge pull request #127 from felicio93/calc_write_lsc2

calc/write lsc2

Author: cuill

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=dp<self.hsm[0]
+        dp[fpz]=self.hsm[0]
+        
+        #find hsm index for all points
+        rat=np.ones(len(gd.nodes.id))*np.nan
+        nlayer=np.zeros(len(gd.nodes.id)).astype('int')
+        ind1=np.zeros(len(gd.nodes.id)).astype('int')
+        ind2=np.zeros(len(gd.nodes.id)).astype('int')
+        for m, hsmi in enumerate(self.hsm):
+            if m==0:
+                fp=dp<=self.hsm[m]
+                ind1[fp]=0; ind2[fp]=0
+                rat[fp]=0; nlayer[fp]=self.nv[0]
+            else:
+                fp=(dp>self.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):

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()