import calendar
import datetime
import glob
import os
import shutil
import subprocess
import numpy as np
from ....utils import path
from logging import info, warning
from ....utils.netcdf import save_nc
from .io.outputs.fake_sim import write_sim, write_obs
from .io.outputs.read_sim import read_obs
from .io.outputs.fake_end import make_end
from .io.outputs.make_mod_out import make_mod_out
from .io.outputs.fetch_end import make_restart_adj
[docs]
def run(self, runsubdir, mode, workdir, ddi, do_simu=True,
approx_transf=False, ref_fwd_dir="", overlap=False,
**kwargs):
"""Run LMDZ model in forward or adjoint mode
Args:
runsubdir (str): working directory for the current run
mode (str): forward or backward
workdir (str): pycif working directory
do_simu (bool): if False, considers that the simulation was
already run
"""
# raise Exception
#raise Exception("stropped")
# Reset update obs to start next time with an empty obs.bin
self.iniobs[ddi] = False
self.reset_obs[ddi] = True
if not do_simu:
# if mode in ["fwd", "tl"]:
# # Keeps the running directory in memory for later adjoint
# # simulations
# self.adj_refdir = "{}/../".format(runsubdir)
return
# Cleaning the directory
path.remove("{}/all_good".format(runsubdir))
path.remove("{}/restart.nc".format(runsubdir))
path.remove("{}/obs_out.bin".format(runsubdir))
path.remove("{}/*_out.bin".format(runsubdir))
if mode != "adj":
path.remove("{}/traj*.bin".format(runsubdir))
if not approx_transf:
# In adjoint mode, setting sensitivities to zero if overlap period
if mode == "adj" and overlap:
obs_scheme = read_obs(runsubdir)
write_obs(obs_scheme, runsubdir)
# Running the model
info("Running sub-simulation in {}".format(runsubdir))
with open("{}/dispersion.log".format(runsubdir), "w") as log:
process = subprocess.Popen(
[self.mpirun] +
getattr(self, "run_options") +
([] if not hasattr(self, "nproc")
else ["-np", "{}".format(self.nproc)]) +
["{}/dispersion.e".format(runsubdir)],
cwd=runsubdir,
stdout=log,
stderr=subprocess.PIPE,
)
_, err = process.communicate()
if err != "" and not os.path.isfile(
"{}/all_good".format(runsubdir)
):
info("Exception in LMDZ")
info(str(err, "utf-8"))
raise Exception(
"LMDZ did not run properly in {}".format(runsubdir)
)
# Dumps to NetCDF if asked
if getattr(self, "dump", False):
info(
"Dumping LMDZ outputs to a NetCDF file in {}".format(runsubdir)
)
dump2nc(self, runsubdir)
# Putting all simulated concentrations to zero if overlap period
if overlap:
obs_scheme = read_obs(runsubdir)
write_sim(obs_scheme, runsubdir)
else:
info("NOT running sub-simulation in {} "
"due to approximating operator".format(runsubdir))
# Mimic all_good
os.system("touch {}/all_good".format(runsubdir))
# Make artificial outputs
if mode in ["fwd", "tl"]:
# Mimic obs_out.bin (with 0 for the simulation)
obs_scheme = read_obs(runsubdir)
write_sim(obs_scheme, runsubdir)
# Make empty restart.nc for safe linking
# Use end concentrations from reference forward directory
make_end(self, runsubdir, ddi, ref_fwd_dir)
else:
# Mimick mod_out files and restart
make_mod_out(self, runsubdir, ddi, ref_fwd_dir)
# Process here initial conditions for the next simulation
date = datetime.datetime.strptime(
os.path.basename(runsubdir), "%Y-%m-%d_%H-%M"
)
shutil.move(
"{}/restart.nc".format(runsubdir),
date.strftime(
"{}/../chain/restart_%Y%m%d%H%M.nc".format(runsubdir)
),
)
if mode == "tl":
restart_tl = "{}/restart_tl.bin".format(runsubdir)
if os.path.isfile(restart_tl):
shutil.move(
"{}/restart_tl.bin".format(runsubdir),
date.strftime(
"{}/../chain/restart_tl_%Y%m%d%H%M.bin".format(runsubdir)
),
)
else:
warning(
f"The file {restart_tl} was not generated. This happens when "
f"the tangent-linear is called with no increment in any of the inputs."
)
if mode != "adj":
# Saves trajectory files for later adjoint simulations
list_file = glob.glob("{}/traj*".format(runsubdir))
for traj_file in list_file:
rad, ext = os.path.splitext(os.path.basename(traj_file))
shutil.move(
"{}/{}{}".format(runsubdir, rad, ext),
date.strftime(
"{}/../chain/{}_%Y%m%d%H%M{}".format(
runsubdir, rad, ext
)
),
)
# Keeps the running directory in memory for later adjoint simulations
# self.adj_refdir = "{}/../".format(runsubdir)
[docs]
def dump2nc(self, runsubdir):
"""Dumps simulated concentration field to a netCDF file"""
# Defines reference dates for NetCDF dates
dref = datetime.datetime.strptime(
os.path.basename(os.path.normpath(runsubdir)), "%Y-%m-%d_%H-%M"
)
# Grid dimension
nlon = self.domain.nlon
nlat = self.domain.nlat
nlev = self.domain.nlev
ndates = calendar.monthrange(dref.year, dref.month)[1] * 8 + 1
# Dimensions
dimnames = ["time", "lev", "lat", "lon"]
dimlens = [None, nlev, nlat, nlon]
units = [dref.strftime("hours since %Y-%m-%d %H:%M:%S")]
vardims = [("time",)]
varnames = ["time"]
dtypes = ["d"]
variables = [3 * np.arange(ndates)]
# Loop over trajq files to extract concentrations
for spec in self.chemistry.emis_species.attributes:
traj_file = "{}/trajq_{}.bin".format(runsubdir, spec)
field = np.fromfile(traj_file, dtype="float")
field = field.reshape((-1, nlev, nlat, nlon), order="C")
vardims.append(("time", "lev", "lat", "lon"))
varnames.append(spec)
variables.append(field)
dtypes.append("d")
units.append("kg/kg")
# Add longitudes and latitudes
vardims.extend([("lat"), ("lon")])
varnames.extend(["lat", "lon"])
variables.append(self.domain.zlat[:, 0])
variables.append(self.domain.zlon[0, :])
dtypes.extend(2 * ["d"])
units.extend(["", ""])
# Saves to NetCDF
fileout = "{}/trajq.nc".format(runsubdir)
save_nc(
fileout, variables, varnames, vardims, dimnames, dimlens, units, dtypes
)
