Source code for imod.mf6.simulation

import collections
import pathlib
import os

import jinja2
import numpy as np
import xarray as xr

import imod


[docs]class Modflow6Simulation(collections.UserDict): def _initialize_template(self): loader = jinja2.PackageLoader("imod", "templates/mf6") env = jinja2.Environment(loader=loader) self._template = env.get_template("sim-nam.j2") def __init__(self, name): super(__class__, self).__init__() self.name = name self._initialize_template() def __setitem__(self, key, value): super(__class__, self).__setitem__(key, value)
[docs] def update(self, *args, **kwargs): for k, v in dict(*args, **kwargs).items(): self[k] = v
[docs] def time_discretization(self, times): """ Collect all unique times """ self.use_cftime = any( [model._use_cftime() for model in self.values() if model._pkg_id == "model"] ) times = [imod.wq.timeutil.to_datetime(time, self.use_cftime) for time in times] for model in self.values(): if model._pkg_id == "model": times.extend(model._yield_times()) # np.unique also sorts times = np.unique(np.hstack(times)) duration = imod.wq.timeutil.timestep_duration(times, self.use_cftime) # Generate time discretization, just rely on default arguments # Probably won't be used that much anyway? timestep_duration = xr.DataArray( duration, coords={"time": np.array(times)[:-1]}, dims=("time",) ) self["time_discretization"] = imod.mf6.TimeDiscretization( timestep_duration=timestep_duration )
[docs] def render(self): """Renders simulation namefile""" d = {} solvername = None models = [] modelnames = [] for key, value in self.items(): if value._pkg_id == "tdis": d["tdis6"] = f"{key}.tdis" elif value._pkg_id == "model": models.append(("gwf6", f"{key}/{key}.nam", key)) modelnames.append(key) elif value._pkg_id == "ims": solvername = key d["models"] = models if solvername is None: raise ValueError("No numerical solution") d["solutiongroups"] = [[("ims6", f"{solvername}.ims", modelnames)]] return self._template.render(d)
[docs] def write(self, directory="."): directory = pathlib.Path(directory) directory.mkdir(exist_ok=True, parents=True) # Write simulation namefile mfsim_content = self.render() mfsim_path = directory / "mfsim.nam" with open(mfsim_path, "w") as f: f.write(mfsim_content) # Write time discretization file self["time_discretization"].write(directory, "time_discretization") # Write individual models globaltimes = self["time_discretization"]["time"].values with imod.util.cd(directory): for key, value in self.items(): # skip timedis, exchanges if value._pkg_id == "model": value.write(key, globaltimes) elif value._pkg_id == "ims": value.write(".", key)