[FEATURE] - Add GWT LKT (Lake Transport) package to the API

[FTW12345]

Request
Feature request: Implement GWT LKT package (lake transport) in iMOD Python
Problem
With LAK in GWF, we need LKT in GWT to:
• set solute concentrations for rainfall and evaporation (and optionally runoff/ext-inflow),
• compute lake concentration,
• exchange mass with aquifer via the corresponding LAK flows.
Currently we hand-write lake.lkt and patch the GWT name file to add:
lkt6 C:/…/lake.lkt lake
(Works, but not ideal because of the extra steps).
Solution (sketch of desired API)
Import
from imod.mf6 import LakeTransport # new
Package
lkt = LakeTransport(
flow_package_name="lake", # or infer by package name match
starting_concentration=0.0, # per-lake 1D or scalar
status="ACTIVE", # per time/lake
rainfall_concentration=rain_ts, # (time,) or (time,lake)
evaporation_concentration=evap_ts, # (time,) or (time,lake)
runoff_concentration=None,
extinflow_concentration=None,
print_input=True,
print_concentration=True,
print_flows=True,
save_flows=True,
concentration_file=str(outdir/"lake_conc.bin"),
budget_file=str(outdir/"lake_lkt.bud"),
budgetcsv_file=str(outdir/"lake_lkt_budget.csv"),
)
gwt["lkt"] = lkt
Requirements

1. LKT dimensions derived from the linked LAK (NLAKES), so no explicit dimensions block in the API.
2. Support FLOW_PACKAGE_NAME and also implicit matching by package name.
3. Support timeseries semantics analogous to other packages (xr.DataArray with time coord; broadcasting over lake where needed).
4. Ensure writer uses forward slashes on Windows, consistent with other writers.
5. Validation should check presence of corresponding GWF LAK package and provide an error if missing.

Workaround / first version
• Write lake.lkt manually (see script below).
• Patch GWT .nam to insert lkt6 … lake inside the begin packages … end packages block, before oc6.
Script

#%% Create LKT Package with own python script
'# Why: Needed to complement LAK package and impod python API misses the LKT package.
'# =============================================================================

#%%% -- lake.lkt --------------------------------------------------------------
#%%%% --- CONFIGURE PATHS -----------------------------------------------------

gwt_dir   = RESULT_SEQ_DIR / f"{MODELNAME}_gwt"     # Folder for writing GWT nam(e) file
gwt_nam   = Path(gwt_dir) / f"{MODELNAME}_gwt.nam"    # File name for nam(e) file
lkt_fname = "lake.lkt"                                                         # Name for to be written file with specified concentrations for rainfall and evaporation

#%%%% --- SPECIFY INPUTS ------------------------------------------------------

flow_pkg_name = "lake"  # The GWF LAK package name

lake_strt_conc = 0.0    # Starting lake concentration: 0.0 gram TDS/L

'# Time series for lake concentrations rainfall and evaporation:   # 1D arrays with shape (NPER,)
'# - Length == NPER (number of stress periods (SP))
'# - rainfall_conc: lake rainfall concentration; Values: 0.016 gram TDS /l as scalar per SP  (or 1.049 gram TDS / l to align with recharge)
'# - evaporation_conc: lake evaporation concentration; Values: 0.0 gram TDS /l as scalar per SP (or 1.049 gram TDS / l to align with evapotranspiration)

rainfall_conc = np.full(NPER, 0.016, dtype=float)
evaporation_conc = np.full(NPER, 0.0, dtype=float)

#%%%% --- WRITE ---------------------------------------------------------------
lkt_path = Path(gwt_dir) / lkt_fname
with open(lkt_path, "w", encoding="utf-8") as f:
    # OPTIONS Block
    f.write("BEGIN OPTIONS\n")
    f.write(f"  FLOW_PACKAGE_NAME {flow_pkg_name}\n")  # Link to GWF LAK
    f.write("  PRINT_INPUT\n")
    f.write("  PRINT_CONCENTRATION\n")
    f.write("  PRINT_FLOWS\n")
    f.write("  SAVE_FLOWS\n")
    f.write(f"  CONCENTRATION FILEOUT {Path(gwt_dir,'lake_conc.bin')}\n")
    f.write(f"  BUDGET FILEOUT {Path(gwt_dir,'lake_lkt.bud')}\n")
    f.write(f"  BUDGETCSV FILEOUT {Path(gwt_dir,'lake_lkt_budget.csv')}\n")
    f.write("END OPTIONS\n\n")

    # PACKAGEDATA Block: Assign <ifno> <strt> [<aux(naux)>] [<boundname>]
    f.write("BEGIN PACKAGEDATA\n")
    f.write(f"  1 {lake_strt_conc:.7g} LakeA\n") # ifno: 1 (NLAKES; equals the LAK package NLAKES)
    f.write("END PACKAGEDATA\n\n")

    # PERIOD blocks: Assign STATUS, RAINFALL and EVAPORATION concentrations each SP
    for kper in range(1, NPER + 1):
        f.write(f"BEGIN PERIOD {kper}\n")
        f.write("  1 STATUS ACTIVE\n")
        f.write(f"  1 RAINFALL {rainfall_conc[kper-1]:.9g}\n")
        f.write(f"  1 EVAPORATION {evaporation_conc[kper-1]:.9g}\n")
        f.write("END PERIOD\n\n")

print(f"Wrote LKT: {lkt_path}")

#%%% -- PATCH GWT NAME FILE so MODFLOW 6 loads LKT ----------------------------
# Iserts the line "lkt6 path/lake.lkt lake" (file type; filename; internal package name)


def PATCH_GWT_NAMEFILE(gwt_nam, lkt_path, flow_pkg_name="lake"):
    nam = Path(gwt_nam)
    if not nam.exists():
        raise FileNotFoundError(f"GWT name file not found: {gwt_nam}")

    # Normalize path to forward slashes (POSIX)
    lkt_posix = Path(lkt_path).as_posix()

    with nam.open("r", encoding="utf-8") as f:
        lines = f.readlines()

    # Find begin/end of packages block
    begin_idx = next((i for i, s in enumerate(lines)
                      if s.strip().lower().startswith("begin packages")), None)
    end_idx = next((i for i, s in enumerate(lines)
                    if i > (begin_idx or -1) and s.strip().lower().startswith("end packages")), None)

    if begin_idx is None or end_idx is None:
        raise RuntimeError("Could not locate 'begin packages' / 'end packages' block in GWT name file.")

    # Check if an lkt6 entry already exists in the block
    block = lines[begin_idx+1:end_idx]
    has_lkt = any(re.search(r'(^|\s)lkt6(\s|$)', s, re.IGNORECASE) for s in block)
    if has_lkt:
        print("lkt6 already present in GWT name file (skipped).")
        return

    # Build the line (match existing indentation style)
    new_line = f"  lkt6 {lkt_posix} {flow_pkg_name}\n"

    # Find oc6 line to insert before; else insert before end packages
    oc_idx_in_block = next((i for i, s in enumerate(block)
                            if re.match(r'\s*oc6\s', s, re.IGNORECASE)), None)

    insert_at_global = (begin_idx + 1 + (oc_idx_in_block if oc_idx_in_block is not None else len(block)))
    lines.insert(insert_at_global, new_line)

    with nam.open("w", encoding="utf-8") as f:
        f.writelines(lines)

    print(f"Inserted into {gwt_nam}:\n{new_line.strip()} (before oc6)" if oc_idx_in_block is not None
          else f"Inserted into {gwt_nam}:\n{new_line.strip()} (before end packages)")

PATCH_GWT_NAMEFILE(
    gwt_nam,
    lkt_path,
    flow_pkg_name="lake"
)