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.

Sign up for GitHub

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

Generate restart from nesting file #179

Open
wants to merge 3 commits into
base: develop
Choose a base branch
from
Open

Generate restart from nesting file #179

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
caf1ac8
scripts to calculate montg while running
AchrefAO Feb 7, 2023
e2b161e
script to generate a restart file from a nesting file
AchrefAO Feb 7, 2023
251e84d
Update topaz_archv2restart_TP5.csh
caglartac Aug 11, 2023
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
282 changes: 155 additions & 127 deletions bin/calc_montg1.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,17 +1,17 @@
#!/usr/bin/env python
import argparse
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot
import abfile.abfile as abf
import numpy
import modeltools.hycom
import logging
import re
import os
import subprocess
import argparse
import datetime
import os.path

import sys
import re
import logging
import abfile
import pickle
import numpy
# Set up logger
##_loglevel=logging.DEBUG
_loglevel=logging.INFO
logger = logging.getLogger(__name__)
logger.setLevel(_loglevel)
Expand All @@ -22,23 +22,21 @@
logger.addHandler(ch)
logger.propagate=False

# Note that this part can be easily implemented into the python code
# that aims to generate nesting [ab] files. I may not have time to manage it
# but evrything has been provided here or other earlier developed NERSC scripts.
# By Mostafa Bakhoday-Paskyabi
# July 4th 2019 (almost at last day of my work at NERSC)
#
# ../bin/calc_montg1.py ../nest/015/archv.2016_001_00.a ./
# ../bin/calc_montg1.py YOUR_INPUT_AFILE OUTPUT_PATH
# It saves with same name as a-file to outpput path

"""
usage:
python ../bin/topaz6_calc_montg1.py archv_file restart_file output_path
python ../bin/topaz6_calc_montg1.py ../nest/070/archv.2022_004_00.a ./data/TP6restart.2021_365_00_0000.a ./

cline1="HYCOM archive files from climatology generated from"
cline2="NEMO Relaxation fields on native grid, MBP2019"
cline3="Expt 01.5 nhybrd=50 nsigma= 0 ds00= 1.00 dp00= 1.00 dp00x= 450.0 dp00f=1.150"
python /cluster/home/achoth/pyScripts_AO_Betzy/ValidationFreq/Alf_calc_montg1.py /cluster/work/users/achoth/TP5a0.06/nest/080_NewMontg/archv.1997_001_00.b /cluster/work/users/achoth/TP5a0.06/expt_08.1/data/restart.1997_001_00_0000.b /cluster/work/users/achoth/TP5a0.06/nest/test_dir/


"""
### AO
# Get component of montgomery potential at surface that depends on pbavg
def montg1_pb(thstar,p) :
kdm=thstar.shape[0]
#AO thref=_constants.thref
thref=1e-3
# Part that depends on pbavg :
montgpb=-thstar[kdm-1,:,:]*thref**2
Expand All @@ -49,149 +47,179 @@ def montg1_pb(thstar,p) :




# Get component of montgomery potential at surface that does not depend on pbavg
def montg1_no_pb(psikk,thkk,thstar,p) :
kdm=thstar.shape[0]
##AO thref=_constants.thref
thref=1e-3
montgc=psikk+(p[kdm,:,:]*(thkk-thstar[kdm-1,:,:]))*thref**2
for k in reversed(range(kdm-1)) :
montgc [:,:]=montgc [:,:]+ p[k+1,:,:]*(thstar[k+1,:,:]-thstar[k,:])*thref**2
montgc [:,:]=montgc [:,:]+ p[k+1,:,:]*(thstar[k+1,:,:]-thstar[k,:])*thref**2
return montgc



def main(archv_files,opath,header_line1=cline1,header_line2=cline2,header_line3=cline3 ) :


### AO


def approx_montg1(thstar,p,srfhgt,idm,jdm,rstr_file):
# estimate montg1---------------------------------------------
# Read input variables from lowest layer of a restart file
thref=1e-3
#psikk_file="./TP6restart.2021_041_00_0000.a"
psikk_file=rstr_file#[0]#[-30:]
logger.info("----->rstr_file :%s"%psikk_file)
m=re.match( "^(.*)(\.[ab])", psikk_file)
if m : psikk_file=m.group(1)
logger.info("----->rstr_file :%s"%psikk_file)
rfile=abfile.ABFileRestart(psikk_file,"r",idm=idm,jdm=jdm)
psikk=rfile.read_field("psikk",0)
thkk=rfile.read_field("thkk",0)
rfile.close()
#print(( "psikk.min(),psikk.max()=",psikk.min(),psikk.max())
#print( "thkk.min(),thkkk.max()=",thkk.min(),thkk.max()
#
#montg1pb = modeltools.hycom._montg_tools.montg1_pb(thstar,p)
#montg1c = modeltools.hycom._montg_tools.montg1_no_pb(psikk,thkk,thstar,p)
montg1pb = montg1_pb(thstar,p)
montg1c = montg1_no_pb(psikk,thkk,thstar,p)
print( "montg1c.min(),montg1c.max()=",montg1c.min(),montg1c.max())
print( "montg1pb.min(),montg1pb.max()=",montg1pb.min(),montg1pb.max() )
pbavg = (srfhgt - montg1c) / (montg1pb+thref)
print( "pbavg.min(),pbavg.max()=",pbavg.min(),pbavg.max() )
montg1 = montg1pb*pbavg + montg1c
logger.info("Estimated montg1 ")
print( " min max montg1=",montg1.min(),montg1.max() )
# end estimate montg1---------------------------------------------
return montg1



def main(archv_files,restart_file,out_path) :
#def main(dtlist,archv_files,restart_file,archv_basedir="") :
# main([],args.archive_file,archv_basedir=args.archv_basedir)

# Set paths using REGION.src environment (must be available at python invocation)

# Open blkdat files. Get some properties
header_line1="HYCOM nested archive files"
header_line2="Compute montg1 from Nemo ssh"
header_line3="Expt 01.7 nhybrd=50 nsigma= 0 ds00= 1.00 dp00= 1.00 dp00x= 450.0 dp00f=1.150"
#
bp=modeltools.hycom.BlkdatParser("blkdat.input")
idm = bp["idm"]
jdm = bp["jdm"]
kdm = bp["kdm"]
thflag = bp["thflag"]
thbase = numpy.float(bp["thbase"])
thbase = bp["thbase"]
kapref = bp["kapref"]
iversn = bp["iversn"]
iexpt = bp["iexpt"]
yrflag = bp["yrflag"]
#thref=modeltools.hycom.thref
iexpt = bp["iexpt"]
thref=1e-3
if kapref == -1 :
kapnum = 2
msg="Only kapref>=0 is implemented for now"
logger.error(msg)
raise ValueError(msg)
else :
kapnum = 1

if kapnum > 1 :
msg="Only kapnum=1 is implemented for now"
logger.error(msg)
raise ValueError(msg)


# hycom sigma and kappa, written in python. NB: sigver is not used here.
# Modify to use other equations of state. For now we assume sigver is:
# 1 (7-term eqs referenced to 0 bar)
# 2 (7-term eqs referenced to 2000 bar)
if thflag == 0 :
sigver=1
else :
sigver=2
sig = modeltools.hycom.Sigma(thflag)
if kapref > 0 : kappa = modeltools.hycom.Kappa(kapref,thflag*1000.0e4) #



# Loop through archive files
# First read nemo ssh from the meain files
#archvf_00=archv_files[0][-19:-4]+"00.a"

for archv_file in archv_files :
archv_file=archv_file[:-2]
logger.debug("Processing %s"%(archv_file))
arcfile=abf.ABFileArchv(archv_file,"r")

temp = numpy.ma.zeros((jdm,idm)) # Only needed when calculating density
saln = numpy.ma.zeros((jdm,idm)) # Only needed when calculating density
archvf_00=archv_file
logger.info("----->Processing archv file :%s"%archvf_00)
logger.info("----->Restart file _pthkk :%s"%restart_file)
logger.info("Daily mean nemo file:%s"%archvf_00)
if os.path.isfile(archvf_00):
logger.info("file %s OK:"%archvf_00)
else:
logger.info("Ooobs file does not exist %s:"%archvf_00)
print( "" )
##logger.debug("Reading nemo_srfhgt from %s"%(archvf_00))
arcfile_in=abfile.ABFileArchv(archvf_00,"r")
nemo_srfhgt=arcfile_in.read_field("srfhgt",0)
#nemo_montg1=arcfile_in.read_field("montg1",thflag)
# ---------------------------------------------------
logger.info("compute thstar and p for montg1")
sig = modeltools.hycom.Sigma(thflag)
# Read all layers .. (TODO: If there are memory problems, read and estimate sequentially)
temp = numpy.ma.zeros((jdm,idm)) # Only needed when calculating density
saln = numpy.ma.zeros((jdm,idm)) # Only needed when calculating density
th3d =numpy.ma.zeros((kdm,jdm,idm))
thstar=numpy.ma.zeros((kdm,jdm,idm))
dp =numpy.ma.zeros((jdm,idm))
p =numpy.ma.zeros((kdm+1,jdm,idm))
pup =numpy.ma.zeros((jdm,idm))
montg=numpy.ma.zeros((jdm,idm))
logger.info("---------------")
logger.info("-----Reading layers to get thstar and p----------")
for k in range(kdm) :
logger.info("Reading layer %d from %s"%(k,archv_file))
temp =arcfile.read_field("temp",k+1)
saln =arcfile.read_field("salin",k+1)
dp [:,:]=arcfile.read_field("thknss",k+1)
th3d [k ,:,:]=sig.sig(temp,saln) - thbase
##logger.debug("Reading layer %d from %s"%(k,archvf_00))
temp =arcfile_in.read_field("temp",k+1)
saln =arcfile_in.read_field("salin",k+1)
#dp [k ,:,:]=arcfile.read_field("thknss",k+1)
dp [:,:]=arcfile_in.read_field("thknss",k+1)
th3d [k ,:,:]=sig.sig(temp,saln)-float(thbase)
#dens_sig=sig.sig(temp,saln)
p [k+1,:,:]= p[k,:,:] + dp[:,:]
if k>0 : thstarup =numpy.ma.copy(thstar[k ,:,:])
thstar[k ,:,:]=numpy.ma.copy(th3d [k ,:,:])
if kapref > 0 :
thstar[k ,:,:]=thstar [k ,:,:] + kappa.kappaf(temp[:,:],
saln[:,:],
th3d[k,:,:]+thbase,
p[k,:,:])
thstar[k ,:,:]=thstar [k ,:,:] + kappa.kappaf(
temp[:,:], saln[:,:], th3d[k,:,:]+float(thbase), p[k,:,:])
elif kapref < 0 :
msg="Only kapref>=0 is implemented for now"
logger.error(msg)
raise ValueError(msg)
if k > 0 :
montg = montg - p[k,:,:] * (thstar[k ,:,:] - thstarup) * thref**2

thkk = thstar[k ,:,:]
psikk = montg

# This part of montg1 does nto depend on pbavg
montg1c = montg1_pb(thstar,p)
# This part depends linearly on pbavg
montg1pb = montg1_no_pb(psikk,thkk,thstar,p)
# Barotropic pressure. NB: srfhgt is in geopotential height :
srfhgt=arcfile.read_field("srfhgt",0)
pbavg = (srfhgt - montg1c) / (montg1pb+thref)
montg1 = montg1pb*pbavg + montg1c
logger.info("Estimated montg1 ")

# Apply mask
#montg1=numpy.ma.masked_where(srfhgt.mask,montg1)
#montg1[~msk]=2**100
#montg1 = numpy.ma.masked_where(srfhgt<2.0**99,montg1)
msk=numpy.where(srfhgt>2.0**99,1,0)


# Open new archive file, and write montg1 to it.
fnameout = opath+str(archv_file[-17:])
arcfile_out=abf.ABFileArchv(fnameout,"w",iversn=iversn,yrflag=yrflag,iexpt=iexpt,mask=False,cline1=header_line1,cline2=header_line2,cline3=header_line3)

fields=arcfile.get_fields()
for key in sorted(fields.keys()) :
fieldname = fields[key]["field"]
time_step = fields[key]["step"]
model_day = fields[key]["day"]
k = fields[key]["k"]
dens = fields[key]["dens"]
fld =arcfile.read_field(fieldname,k)

if fieldname == "montg1" :
logger.info("Writing field %10s at level %3d to %s (modified)"%(fieldname,k,fnameout))
arcfile_out.write_field(montg1,None,fieldname,time_step,model_day,k,dens)
else :
arcfile_out.write_field(fld ,None,fieldname,time_step,model_day,k,dens)

logger.info("Finished writing to %s"%fnameout)
# ---------------------------------------------------
#logger.info("Adding nemo ssh to the tide elevation:%s"%archv_file_in_string)
logger.info("---------------------------------")
archv_file_out = out_path+str(archvf_00[-19:])
logger.info("----->Adding montg1 archv_file_out:%s"%archv_file_out)
logger.info("Output file %s"%(archv_file_out))
arcfile_out=abfile.ABFileArchv(archv_file_out,"w",iversn=iversn,yrflag=yrflag,iexpt=iexpt,mask=False,cline1=header_line1,cline2=header_line2,cline3=header_line3)
#
###AA regr = open("montg_regress_tide.pckl",'rb')
###AA a,b,nemomeandt = pickle.load(regr)
###AA tide_montg1 = (nemo_srfhgt - nemomeandt) * a + b
estimated_montg1=approx_montg1(thstar,p,nemo_srfhgt,idm,jdm,restart_file)
#---#
#read all keys and write updated version
for keys in sorted(arcfile_in.fields.keys()) :
fieldname = arcfile_in.fields[keys]["field"]
time_step = arcfile_in.fields[keys]["step"]
model_day = arcfile_in.fields[keys]["day"]
k = arcfile_in.fields[keys]["k"]
dens = arcfile_in.fields[keys]["dens"]
fld = arcfile_in.read_field(fieldname,k)
#logger.info("Current field %10s at level %3d, and dens %s"%(fieldname,k,dens))
if fieldname == "montg1" :
logger.info("Writing field %10s at level %3d to %s (modified)"%(fieldname,k,archv_file_out))
arcfile_out.write_field(estimated_montg1,None,fieldname,time_step,model_day,k,dens)
#elif fieldname == "srfhgt" :
# logger.info("Writing field %10s at level %3d to %s (modified)"%(fieldname,k,archv_file_out))
# #arcfile_out.write_field(tide_srfhgt,None,fieldname,time_step,model_day,k,dens)
else :
arcfile_out.write_field(fld ,None,fieldname,time_step,model_day,k,dens)

arcfile_in.close()
arcfile_out.close()
arcfile.close()

# Remove archive files lying around
#logger.info("Finished - files placed in location %s"%os.path.abspath("./"))





if __name__ == "__main__" :
basecmd=os.path.basename(sys.argv[0])

parser = argparse.ArgumentParser(description='Calculate montgomery potential in 1st layer (used in archv files)')
parser.add_argument('--header_line1', type=str, help="First header line, applicable only for archm & archv",default=cline1)
parser.add_argument('--header_line2', type=str, help="Montgomery with or without barotropic pressure",default=cline2)
parser.add_argument('--header_line3', type=str, help="Montgomery with or without barotropic pressure",default=cline3)
class DateTimeParseAction(argparse.Action) :
def __call__(self, parser, args, values, option_string=None):
tmp = datetime.datetime.strptime( values, "%Y-%m-%dT%H:%M:%S")
setattr(args, self.dest, tmp)


parser = argparse.ArgumentParser(description='')
#parser.add_argument("--archv-basedir",type = str, help = "",default="")
parser.add_argument('archive_file', type=str, nargs="+")
parser.add_argument('restart_file', type=str, help="Restart file produced by model run")
parser.add_argument('out_path', type=str, help='Output files path')
#parser_opt2.set_defaults(subparser_name="2")

parser.add_argument('archive_file', type=str, help='Files to add montg1 to',nargs="+")
parser.add_argument('opath', type=str, help='Output files path')
args = parser.parse_args()
main(args.archive_file,args.opath,header_line1=args.header_line1,header_line2=args.header_line2,header_line3=args.header_line3)

main(args.archive_file,args.restart_file,args.out_path)
Loading