def run_gcm(self, optim_iter):
"""run the MITgcm at the given optimization iteration"""
# iter0
# 0. create sim
run_dir = self.get_run_dir(optim_iter)
m1qn3_dir = self.get_m1qn3_dir()
sim = rp.Simulation(run_dir=run_dir, **self.dsim)
# 1. write data.ctrl
self.write_ctrl_namelist(optim_iter)
# 2. write data.optim
self.write_optim_namelist(optim_iter)
# 3. link and launch
sim.link_to_run_dir()
if optim_iter > 0:
fname = self.ctrl_packname + f'_MIT_CE_000.opt{optim_iter:04d}'
_symlink(join(m1qn3_dir, fname), join(run_dir, fname))
sim.write_slurm_script()
jid = sim.submit_slurm(**self.slurm)
self.submit_next_stage(next_stage='m1qn3',
optim_iter=optim_iter,
jid_depends=jid,
mysim=sim)
def run_gcm(self, optim_iter, gcm_type):
"""run the MITgcm at the given optimization iteration"""
# iter0
# 0. create sim
run_dir = self.get_run_dir(optim_iter, gcm_type)
m1qn3_dir = self.get_m1qn3_dir()
dsim = self.adm_sim if gcm_type =='adm' else self.fwd_sim
# --- Pickups 1: first kill the pickup_dir if needed
# then let the sim create the run directory, etc
# fwd: always uses same pickups to do original spinup+this iter's ctrl
# adm: use what fwd just gave you
if (optim_iter > 0 and gcm_type == 'adm'):
dsim['pickup_dir'] = None
sim = rp.Simulation(run_dir=run_dir,**dsim)
# --- Pickups 2: now copy over
if (optim_iter > 0 and gcm_type == 'adm'):
previous_fwd_dir = self.get_run_dir(optim_iter, 'fwd')
for suffix in ['meta','data']:
old_pickup = os.path.join(previous_fwd_dir,
f'pickup.{self.nIterSpinFinal:010d}.{suffix}')
new_pickup = os.path.join(run_dir,
f'pickup.{self.nIter0:010d}.{suffix}')
copyfile(old_pickup, new_pickup)
# 1. write data.ctrl
self.write_ctrl_namelist(optim_iter, gcm_type)
# 2. write data.optim
self.write_optim_namelist(optim_iter, gcm_type)
# 3. link and launch
sim.link_to_run_dir()
if optim_iter>0:
fname=self.ctrl_packname+f'_MIT_CE_000.opt{optim_iter:04d}'
_symlink(join(m1qn3_dir,fname), join(run_dir,fname))
sim.write_slurm_script()
jid = sim.submit_slurm(**self.slurm)
next_stage = 'm1qn3' if gcm_type == 'adm' else 'adm'
self.submit_next_stage(next_stage=next_stage,optim_iter=optim_iter,
jid_depends=jid,mysim=sim)
def start(self,main_dir,
adm_sim,fwd_sim,
stage='adm',optim_iter=0,
**kwargs):
"""Start the experiment by writing everything we need to file,
to be read in later by "pickup"
Parameters
----------
main_dir : str
main directory where all the mitgcm run directories will launch from
adm_sim, fwd_sim : dict
containing the base parameters for a rosypig.Simulation
with everything but run_dir
see rosypig.Simulation for details
kwargs
Are passed to override the default class attributes
Additionally:
optimx : str, optional
path to m1qn3 executable, default is main_dir/build.m1qn3/optim.x
"""
# --- make some dirs
slurm_dir = os.path.join(main_dir,'slurm')
for mydir in ['json',main_dir,slurm_dir]:
_dir(mydir);
dirs={'main_dir':main_dir,'slurm_dir':slurm_dir}
self.write_json(dirs,'_dirs.json')
self.write_json(adm_sim,'_adm_sim.json')
self.write_json(fwd_sim,'_fwd_sim.json')
# Requiring these arguments, but throwing them into kwargs b/c it's easy
if kwargs !={}:
self.write_json(kwargs,'_kwargs.json')
# --- "pickup" experiment at stage
self.pickup()
adm_sim.pop('name')
startsim = rp.Simulation('startmeup',**adm_sim)
self.submit_next_stage(next_stage=stage,optim_iter=optim_iter,
mysim=startsim)
def range_approx_one(self):
"""Define the Laplacian-like operator A = delta - div(kappa grad( ))
and use the MITgcm to:
1. generate n_samples gaussian random samples, g_i
2. solve the linear system for z_i: Az_i = g_i
to get the range of A^{-1}
3. use to compute the sample variance of A^{-1}
... pass to the next step
"""
jid_list = []
for Nx in self.NxList:
for xi in self.xiList:
_, write_dir, run_dir = self._get_dirs('range_approx_one', Nx,
xi)
self.smooth_writer(write_dir,
xi=xi,
smooth_apply=False,
num_inputs=self.n_samples)
matern = MaternField(xdalike=self.mymodel,
n_range=Nx,
horizontal_factor=xi,
isotropic=self.isotropic)
matern.write_binaries(write_dir=write_dir,
smoothOpNb=self.smoothOpNb)
sim = rp.Simulation(name=f'{Nx:02}dx_{xi:02}xi_oi_ra1',
run_dir=run_dir,
obs_dir=write_dir,
**self.dsim)
# launch job
sim.link_to_run_dir()
sim.write_slurm_script()
jid = sim.submit_slurm(**self.slurm)
jid_list.append(jid)
def start(self,
dirs,
dsim,
mymodel,
ctrl_ds,
startat='range_approx_one',
**kwargs):
"""Start the experiment by writing everything we need to file,
to be read in later by "pickup"
Parameters
----------
dirs : dict
containing the key directories for the experiment
'main_run' : the directory where MITgcm files are written
and simulations are started
'netcdf' : where to save netcdf files when finished
Note: a separate tmp/ directory is created inside for
intermittent saves
dsim : dict
containing the base parameters for a rosypig.Simulation
'machine', 'n_procs', 'exe_file', 'binary_dir', 'pickup_dir', 'time'
see rosypig.Simulation for details
mymodel : xarray DataArray
a mask field with True denoting domain wet points, e.g. maskC
Note: the dimension order should be appropriate for MITgcm read/write,
with vertical going from top to bottom
Note: dimensions are assumed to be ordered like
(vertical, meridional, zonal)
m0 : xarray DataArray
with the initial guess
Note: must have attribute "standard_name" with identifying name, e.g.
m0.attrs['standard_name'] = 'zeros'
ctrl_ds : xarray DataArray
dataset with the FULL coordinate dataset, for evaluating the laplacian
obs_mask : xarray DataArray
mask field denoting where observations are taken, on an "observation grid"
Note: dimensions are assumed to be ordered like
(vertical, meridional, zonal)
obs_mean, obs_std : xarray DataArray
containing observation mean and uncertainties
(of course, at the mask points!)
startat : str, optional
stage to start at
kwargs
Are passed to override the default class attributes
"""
# --- Add tmp netcdf directory
dirs['nctmp'] = dirs['netcdf'] + '/tmp'
for mydir in ['json', dirs['netcdf'], dirs['nctmp'], dirs['main_run']]:
_dir(mydir)
# --- Write the directories
def write_json(mydict, mysuff):
json_file = f'json/{self.experiment}' + mysuff
with open(json_file, 'w') as f:
json.dump(mydict, f)
write_json(dirs, '_dirs.json')
write_json(dsim, '_sim.json')
if kwargs != {}:
write_json(kwargs, '_kwargs.json')
# --- Write ctrl and obs datasets to netcdf
myctrl = xr.Dataset({'mymodel': mymodel})
myctrl.to_netcdf(dirs['nctmp'] + f'/{self.experiment}_ctrl.nc')
ctrl_ds.to_netcdf(dirs['nctmp'] + f'/{self.experiment}_cds.nc')
# --- "pickup" experiment at startat
self.pickup()
startsim = rp.Simulation('startmeup', **dsim)
self.submit_next_stage(next_stage=startat, mysim=startsim)
def run_m1qn3(self,optim_iter):
run_dir = self.get_run_dir(optim_iter, gcm_type='m1qn3')
m1qn3_dir = self.get_m1qn3_dir()
# link optim.x to run directory
exe = os.path.basename(self.optimx)
_symlink(self.optimx, join(m1qn3_dir,exe))
# link over cost vector
fname = self.cost_packname+f'_MIT_CE_000.opt{optim_iter:04d}'
_symlink(join(run_dir,fname), join(m1qn3_dir,fname))
# iter 0: get ctrl vector too
if optim_iter == 0:
fname = self.ctrl_packname+f'_MIT_CE_000.opt{optim_iter:04d}'
_symlink(join(run_dir,fname), join(m1qn3_dir,fname))
# copy data.ctrl, data.optim
for fname in ['data.ctrl','data.optim']:
copyfile(join(run_dir,fname), join(m1qn3_dir,fname))
# move to run dir and run
pwd = os.getenv('PWD')
os.chdir(m1qn3_dir)
run_cmd = f'./{exe} > stdout.{optim_iter:04d} 2> stderr.{optim_iter:04d}'
subprocess.run(run_cmd,shell=True)
# check stderr
fname=f'stderr.{optim_iter:04d}'
with open(fname,'r') as f:
stderr = f.read()
if len(stderr)>0:
raise RuntimeError(f'm1qn3 failed, see: {os.path.realpath(fname)}')
os.chdir(pwd)
# read gradient norm
with open(f'{m1qn3_dir}/m1qn3_output.txt','r') as f:
for line in f.readlines():
if 'two-norm of g' in line:
g_norm = float(line.split('=')[1][:-1].replace('D','E'))
# set ||g0||
if optim_iter==0:
self.g0 = g_norm
self.write_json
kwargs = self.read_json('_kwargs.json')
kwargs = kwargs if kwargs is not None else {}
kwargs['g0'] = g_norm
self.write_json(kwargs,'_kwargs.json')
# submit next optim iter
if optim_iter+1 < self.optim_iter_max and g_norm > self.epsg*self.g0:
sim=rp.Simulation(**self.fwd_sim)
self.submit_next_stage(next_stage='fwd',optim_iter=optim_iter+1, mysim=sim)
else:
if g_norm < self.epsg * self.g0:
print('\t\t---\t\tCongratulations!!\t\t---')
print('')
print(f' ||g_{optim_iter}||_2 / ||g_0||_2 < eps_g = {self.epsg}')
print('')
print(f' g_norm = {g_norm:1.6e}')
print(f' optim_iter = {optim_iter}')
else:
print(f' --- Reached Maximum Optimization Iterations: {self.optim_iter_max} ---')
print(' exiting ...')