Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

calc/write lsc2 #127

Merged
merged 4 commits into from
May 29, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
155 changes: 152 additions & 3 deletions pyschism/mesh/vgrid.py
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
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