Basic converter from flopy classes for a DISV model

[Huite]

This isn't something we necessarily want to include in imod, given the complexities of MODFLOW 6 input (which are matched by Flopy) which don't fit the imod data models. But I'm currently sharing via e-mail with @jdelsman which isn't great either, and some other projects might benefit...

# %%
import abc
from collections import defaultdict

import xugrid as xu
import imod
import flopy
import numpy as np
import xarray as xr
import pandas as pd

# %%

class FlopyMapper(abc.ABC):
    """
    Map flopy vars to imod-python vars.
    """
    @classmethod
    def to_imod(cls, package, **kwargs):
        # Scalar data only
        variables = {}
        for imod_name, flopy_name in cls._VAR_MAPPING.items():
            data = getattr(package, flopy_name)
            variables[imod_name] = data
        return cls._IMOD_CLS(**variables)

    @classmethod
    def _create_dataarray(cls, data, grid, coords=None, dims=None):
        if np.isscalar(data):
            return data
        if dims is None:
            if data.ndim == 1:
                dims = (grid.face_dimension,)
            elif data.ndim == 2:
                dims = ("layer", grid.face_dimension,)
        if coords is None and data.ndim == 2:
            coords = {"layer": np.arange(1, len(data) + 1)}
        return xu.UgridDataArray(
            xr.DataArray(data=data, coords=coords, dims=dims),
            grid,
        )

    @classmethod
    def _stress_period_data(cls, package, flopy_name):
        # Make period_data consistent in form: always a dict.
        if not package.has_stress_period_data:
            data = np.atleast_2d(getattr(package, flopy_name).data)
            time_index = np.array([0])
        else:
            # TODO: not sure if this returns a dict of int => np.ndarray
            data = np.concatenate(list(getattr(package, flopy_name).stress_period_data.values()))
            time_index = np.array(list(package.stress_period_data.keys()))
        return time_index, data


class FlopyGridMapper(FlopyMapper, abc.ABC):
    @classmethod
    def to_imod(cls, package, grid, **kwargs):
        variables = {}
        for imod_name, flopy_name in cls._VAR_MAPPING.items():
            data = getattr(package, flopy_name).array
            variables[imod_name] = cls._create_dataarray(data, grid=grid)
        return cls._IMOD_CLS(**variables)


class FlopyListMapper(FlopyMapper, abc.ABC):
    @classmethod
    def to_imod(cls, package, grid, times):
        dataframes = package.stress_period_data.dataframe

        # Determine number of layers required
        layers = set()
        time_index = []
        systems = defaultdict(list)
        n_system = 1
        for time, dataframe in dataframes.items():
            time_index.append(time)
            layers |= set(dataframe["cellid_layer"].unique())
            # Are cell id's unique? Otherwise we have to create multiple "systems" per package.
            grouped = dataframe.groupby(["cellid_layer", "cellid_cell"])
            maxsize = grouped.size().max()
            n_system = max(n_system, maxsize)
            for i in range(maxsize):
                systems[i].append(grouped.nth(i))

        layers = np.array(list(layers))
        shape = (len(time_index), len(layers), grid.n_face)
        dims = ("time", "layer", grid.face_dimension)
        coords = {"time": times[np.array(time_index)], "layer": layers + 1}


        imod_packages = {}
        for system_number, system in systems.items():
            variables = {}
            for imod_name, flopy_name in cls._VAR_MAPPING.items():
                data = np.full(shape, np.nan)
                for index, dataframe in enumerate(system):
                    # Layers do not have to be consecutive, e.g. layer 1, 3, 4, 7 is allowed.
                    layer_index = np.searchsorted(layers, dataframe["cellid_layer"])
                    data[index, layer_index, dataframe["cellid_cell"]] = dataframe[flopy_name]
                    variables[imod_name] = cls._create_dataarray(data, grid=grid, coords=coords, dims=dims)

            # Only add a suffix if we actually have multiple packages as a result.
            if n_system == 1:
                suffix = ""
            else:
                suffix = f"_{system_number}"
            imod_packages[f"{package.name[0]}{suffix}"] = cls._IMOD_CLS(**variables) 

        return imod_packages


class RiverMapper(FlopyListMapper):
    _VAR_MAPPING = {
        "stage": "stage",
        "conductance": "cond",
        "bottom_elevation": "rbot"
    }
    _IMOD_CLS = imod.mf6.River


class DrainageMapper(FlopyListMapper):
    _VAR_MAPPING = {
        "elevation": "elev",
        "conductance": "cond",
    }
    _IMOD_CLS = imod.mf6.Drainage


class GeneralHeadBoundaryMapper(FlopyListMapper):
    _VAR_MAPPING = {
        "head": "bhead",
        "conductance": "cond",
    }
    _IMOD_CLS = imod.mf6.GeneralHeadBoundary


class ConstantHeadMapper(FlopyListMapper):
    _VAR_MAPPING = {
        "head": "head",
    }
    _IMOD_CLS = imod.mf6.ConstantHead


class RechargeArrayMapper(FlopyMapper):
    _VAR_MAPPING = {
        "rate": "recharge",
    }
    _IMOD_CLS = imod.mf6.Recharge

    @classmethod
    def to_imod(cls, package, grid, times):
        variables = {}
        for imod_name, flopy_name in cls._VAR_MAPPING.items():
            time_index, data = cls._stress_period_data(package, flopy_name)
            variables[imod_name] = cls._create_dataarray(
                data, grid=grid, coords={"time": times[time_index], "layer": 1}, dims=("time", grid.face_dimension),
            )
        return cls._IMOD_CLS(**variables)


class NodePropertyFlowMapper(FlopyGridMapper):
    _VAR_MAPPING = {
        "icelltype": "icelltype",
        "k": "k",
        "k33": "k33",
        "starting_head_as_confined_thickness": "thickstrt",
    }
    _IMOD_CLS = imod.mf6.NodePropertyFlow


class StorageMapper(FlopyGridMapper):
    @classmethod
    def to_imod(cls, package, grid, **kwargs):
        var_mapping = {
            "convertible": "iconvert",
            "specific_yield": "sy",
        }
        if package.storagecoefficient:
            imod_cls = imod.mf6.StorageCoefficient 
            var_mapping["storage_coefficient"] = "ss"
        else:
            imod_cls = imod.mf6.SpecificStorageCoefficient
            var_mapping["specific_storage"] = "ss"

        variables = {}
        for imod_name, flopy_name in var_mapping.items():
            data = getattr(package, flopy_name).array
            variables[imod_name] = cls._create_dataarray(data, grid=grid)

        variables["transient"] = package.transient is not None
        return imod_cls(**variables)


class VerticesDiscretizationMapper(FlopyGridMapper):
    _VAR_MAPPING = {
        "top": "top",
        "bottom": "botm",
        "idomain": "idomain",
    }
    _IMOD_CLS = imod.mf6.VerticesDiscretization

    @staticmethod
    def _process_faces(face_node_connectivity):
        # Flip around to counterclockwise orientation
        n, m = face_node_connectivity.shape
        isnode = face_node_connectivity != -1
        n_valid = isnode.sum(axis=1)
        j = np.arange(m)                              
        indices = np.where(
            j <  n_valid[:, None],
            n_valid[:, None] - 1 - j,
            j,
        )
        out = np.take_along_axis(face_node_connectivity, indices, axis=1)
        # Skip the last column, it's the same as the first vertex.
        # Set it to zero via the flat view.
        closing_column = n_valid - 1 + (np.arange(n) * m)
        out.ravel()[closing_column] = -1
        return out[:, :-1]

    @classmethod
    def to_imod(cls, package):
        xy_dataframe = pd.DataFrame.from_records(package.vertices.array)
        cell2d_dataframe = pd.DataFrame.from_records(package.cell2d.array).fillna(-1)
        # We want the face_node_connectivity.
        # Columns are: icell2d, xc, yc, ncvert; then the vertices.
        face_node_connectivity = cls._process_faces(
            cell2d_dataframe.iloc[:, 4:].to_numpy().astype(int)
        )
        grid = xu.Ugrid2d(
            node_x=xy_dataframe["xv"],
            node_y=xy_dataframe["yv"],
            fill_value=-1,
            face_node_connectivity=face_node_connectivity,
        )
        return super().to_imod(package=package, grid=grid)


class InitialConditionsMapper(FlopyGridMapper):
    _VAR_MAPPING = {
        "start": "strt",
    }
    _IMOD_CLS = imod.mf6.InitialConditions


class OutputControlMapper(FlopyMapper):
    @classmethod
    def to_imod(cls, package, grid, times):
        time_indices = []
        save_head = []
        for time_index, recarray in package.saverecord.data.items():
            rows = [dict(zip(recarray.dtype.names,x)) for x in recarray]
            for row in rows:
                if row["rtype"] == "head":
                    time_indices.append(time_index)
                    setting = row["ocsetting"]
                    if setting == "frequency":
                        save_head.append(row["ocsetting_data"])
                    else:
                        save_head.append(setting)
        return imod.mf6.OutputControl(save_head=xr.DataArray(data=save_head, coords={"time": times[time_indices]}))


class WellMapper(FlopyMapper):
    @classmethod
    def to_imod(cls, package, grid, times):
        dataframes = package.stress_period_data.dataframe
        nperiod = len(dataframes)
        dfs = []
        df_lengths = []
        for time_index, dataframe in dataframes.items():
            df = dataframe.copy()
            nwell = len(df)
            df["time"] = times[time_index]
            df["id"] = np.arange(nwell) 
            dfs.append(df)
        total_df = pd.concat(dfs)

        # Single period, no need to identify what the same well is over time.
        if nperiod == 1:
            x, y = grid.face_coordinates[total_df["cellid_cell"]].T
            layer = total_df["cellid_layer"].to_numpy() + 1
            rate = (
                total_df
                .reset_index()
                .set_index(["time", "index"])["q"]
                .to_xarray()
            )
        else:
            try:
            # Maybe we can identify wells by cellid...
            # This only works with a single well per cell;
            # But results in a nice structured result
                ds = total_df.set_index(["time", "cellid_layer", "cellid_cell"])[["q"]].to_xarray()
                x, y = grid.face_coordinates[ds["cellid_cell"]].T
                layer = ds["cellid_layer"].to_numpy() + 1
                rate = ds.stack(index=("cellid_layer", "cellid_cell")).reset_index("index")["q"]
            except ValueError:
                # ValueError("cannot convert a DataFrame with a non-unique MultiIndex into xarray")
                # Treat every well as independent?
                ds = total_df.set_index(["time", "id"]).to_xarray().rename({"id": "index"})
                # In case of irregular number of wells, the conversion to xarray will have turned the
                # layer variable into float, with NaNs; imod-python demands an integer array.
                layer = ds["cellid_layer"].fillna(1).astype(int) + 1
                rate = ds["q"]
                # Get a time varying x, y
                # Fill so we don't index with bad values
                # I think imod-python will filter these away based on the NaNs in the rate variable...
                xy = grid.face_coordinates[ds["cellid_cell"].fillna(-1).astype(int)]
                x = rate.copy(data=xy[..., 0])
                y = rate.copy(data=xy[..., 1])

        return imod.mf6.LayeredWel(
            x=x,
            y=y,
            layer=layer,
            rate=rate,
        )

MAPPERS = {
    flopy.mf6.ModflowGwfnpf: NodePropertyFlowMapper,
    flopy.mf6.ModflowGwfsto: StorageMapper,
    flopy.mf6.ModflowGwfic: InitialConditionsMapper,
    flopy.mf6.ModflowGwfchd: ConstantHeadMapper,
    flopy.mf6.ModflowGwfrcha: RechargeArrayMapper,
    flopy.mf6.ModflowGwfriv: RiverMapper,
    flopy.mf6.ModflowGwfdrn: DrainageMapper,
    flopy.mf6.ModflowGwfghb: GeneralHeadBoundaryMapper,
    flopy.mf6.ModflowGwfoc: OutputControlMapper,
    flopy.mf6.ModflowGwfwel: WellMapper,
}


def get_times(tdis, default_start_time = "2020-01-01", default_time_unit="days"):
    start_time = tdis.start_date_time.data
    if start_time is None:
        start_time = pd.to_datetime(default_start_time)
    else:
        start_time = pd.to_datetime(start_time)

    time_unit = tdis.time_units.data
    if time_unit is None:
        time_unit = default_time_unit

    tdis_dataframe = pd.DataFrame.from_records(tdis.perioddata.get_data())
    duration = np.insert(tdis_dataframe["perlen"], 0, 0)[:-1]
    return start_time + pd.to_timedelta(duration, unit=time_unit).to_numpy()


def convert_groundwatermodel(flopy_model, times):
    newton = flopy_model.newtonoptions.data[0][0]
    try:
        under_relaxation = flopy_model.newtonoptions.data[0][1].lower() == "under_relaxation"
    except:
        pass

    model = imod.mf6.GroundwaterFlowModel(
        newton=newton,
        under_relaxation=under_relaxation,
        print_input=flopy_model.print_input is True,
        print_flows=flopy_model.print_flows is True, 
        save_flows=flopy_model.save_flows is True,
    )
    # Process DIS first for topology.
    disv = VerticesDiscretizationMapper.to_imod(package=flopy_model.disv)
    disv_name = flopy_model.disv.name[0]
    grid = disv.dataset.ugrid.grid
    model[disv_name] = disv

    packages = flopy_model.packagelist
    for package in packages:
        if isinstance(package, flopy.mf6.ModflowGwfdisv):
            continue

        converted = MAPPERS[type(package)].to_imod(package=package, grid=grid, times=times)
        # A river, gbh, etc. might require multiple systems.
        if isinstance(converted, dict):
            for key, value in converted.items():
                model[key] = value
        else:
            model[package.name[0]] = converted

    return model


def flopy_simulation_to_imod(sim):
    simulation = imod.mf6.Modflow6Simulation(name=sim.name)
    times = get_times(tdis=sim.tdis)
    modelnames = []
    for modelname, flopy_model in sim.model_dict.items():
        simulation[modelname] = convert_groundwatermodel(flopy_model=flopy_model, times=times)
        modelnames.append(modelname)

    simulation["ims"] = imod.mf6.Solution(
        modelnames=modelnames,
        outer_dvclose=1.0e-4,
        outer_maximum=50,
        inner_maximum=200,
        inner_dvclose=1.0e-5,
        inner_rclose=0.5,
        linear_acceleration="bicgstab",
        relaxation_factor=0.97,
    )
    return simulation

# %% 

sim = flopy.mf6.MFSimulation.load(sim_ws=".")

# %%

simulation = flopy_simulation_to_imod(sim)

# %%
# Manual cleanup

gwf = simulation["no_permits"]
idomain = gwf["disv"]["idomain"]
gwf["npf"]["k"] = gwf["npf"]["k"].where(idomain > 0)
gwf["npf"]["k33"] = gwf["npf"]["k33"].where(idomain > 0)
gwf["sto"]["storage_coefficient"] = gwf["sto"]["storage_coefficient"].where(idomain > 0)
gwf["sto"]["specific_yield"] = gwf["sto"]["specific_yield"].where(idomain > 0)
gwf["ic"]["start"] = gwf["ic"]["start"].where(idomain > 0)

# MF6 fails on GHB and RIV below bottom when "drying" is allowed:
gwf["sto"]["convertible"] *= 0 
gwf["npf"]["icelltype"] *= 0

simulation.create_time_discretization(["2020-01-02 00:00:00"])  # end time
# %%
simulation.write("test-huite", validate=False)
# %%
# Read output and visualize heads

head = imod.mf6.open_hds(r"z:\tmp\test-huite\no_permits\no_permits.hds", r"z:\tmp\test-huite\no_permits\disv.disv.grb")
head.isel(time=0, layer=0).ugrid.plot(levels=np.arange(-5, 25))