Source code for imod.mf6.simulation

import collections
import os
import pathlib

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, keep_trailing_newline=True) self._template = env.get_template("sim-nam.j2") def __init__(self, name): super(__class__, self).__init__() = 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 from boundary conditions and insert times from argument as well. Function creates TimeDiscretization object which is set to self["time_discretization"] Parameters ---------- times : xr.DataArray of datetime-likes times to be inserted into model time discretization. Note ---- To set the other parameters of the TimeDiscretization object, you have to set these to the object after calling this function. Example ------- >>> simulation = imod.mf6.Modflow6Simulation("example") >>> simulation.time_discretization(times=["2000-01-01", "2000-01-02"]) >>> # Set number of timesteps >>> simulation["time_discretization"]["n_timesteps"] = 5 """ 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 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)
[docs] def write_qgis_project(self, crs, directory=".", aggregate_layers=False): directory = pathlib.Path(directory) directory.mkdir(exist_ok=True, parents=True) with for key, value in self.items(): # skip timedis, exchanges if value._pkg_id == "model": value.write_qgis_project( key, crs, aggregate_layers=aggregate_layers )