Skip to content

Commit

Permalink
Merge pull request #127 from felicio93/calc_write_lsc2
Browse files Browse the repository at this point in the history 
calc/write lsc2
  • Loading branch information
@cuill
cuill authored May 29, 2024
2 parents 6006c2c + bc1ba69 commit d07e726
Show file tree
Hide file tree
Showing 3 changed files with 510 additions and 4 deletions.
155 changes: 152 additions & 3 deletions pyschism/mesh/vgrid.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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 = [
Expand All @@ -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):

Expand Down Expand Up @@ -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):
Expand Down
92 changes: 91 additions & 1 deletion tests/_mesh/test_vgrid.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
Loading

0 comments on commit d07e726

Please sign in to comment.