diff --git a/docs/steady/04benchmarks/test_polygon_areasink.ipynb b/docs/steady/04benchmarks/test_polygon_areasink.ipynb index dd49d6d2..8cff98cb 100644 --- a/docs/steady/04benchmarks/test_polygon_areasink.ipynb +++ b/docs/steady/04benchmarks/test_polygon_areasink.ipynb @@ -83,7 +83,7 @@ "d2hdx2 = (ml.head(x + d, y) - 2 * ml.head(x, y) + ml.head(x - d, y)) / (d**2)\n", "d2hdy2 = (ml.head(x, y + d) - 2 * ml.head(x, y) + ml.head(x, y - d)) / (d**2)\n", "d2hdx2 + d2hdy2\n", - "aqin = ml.aq.inhomlist[0]\n", + "aqin = ml.aq.inhomdict[p1.name]\n", "print(\"recharge from numerical derivative: \", np.sum(aqin.T * (d2hdx2 + d2hdy2)))\n", "h = ml.head(x, y)\n", "print(\"leakage from aq0 to aq1 from head difference: \", (h[1] - h[0]) / aqin.c[1])\n", @@ -92,6 +92,13 @@ " aqin.T[1] * (d2hdx2[1] + d2hdy2[1]),\n", ")" ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] } ], "metadata": { diff --git a/docs/transient/03xsections/1d_elements.ipynb b/docs/transient/03xsections/1d_elements.ipynb index ab501ccc..1e3ab637 100644 --- a/docs/transient/03xsections/1d_elements.ipynb +++ b/docs/transient/03xsections/1d_elements.ipynb @@ -152,7 +152,7 @@ "metadata": {}, "outputs": [], "source": [ - "mlconf = tft.ModelMaq(kaq=k, z=[0, -H], Saq=Ss, tmin=1e-3, tmax=1e3, topboundary=\"conf\")\n", + "mlconf = tft.ModelMaq(kaq=k, z=[0, -H], Saq=Ss, tmin=1e-1, tmax=10, topboundary=\"conf\")\n", "ls = tft.LineSink1D(mlconf, tsandq=[(0, Q)], layers=[0])\n", "mlconf.solve()" ] @@ -199,7 +199,7 @@ "metadata": {}, "outputs": [], "source": [ - "ml = tft.ModelMaq(kaq=k, z=[1, 0, -H], c=[10], tmin=1e-3, tmax=1e3, topboundary=\"semi\")\n", + "ml = tft.ModelMaq(kaq=k, z=[1, 0, -H], c=[10], tmin=1e-1, tmax=10, topboundary=\"semi\")\n", "ls = tft.LineSink1D(ml, tsandq=[(0, Q)], layers=[0])\n", "ml.solve()" ] diff --git a/docs/transient/03xsections/1d_inhom_elements.ipynb b/docs/transient/03xsections/1d_inhom_elements.ipynb index c68961e6..90eb4814 100644 --- a/docs/transient/03xsections/1d_inhom_elements.ipynb +++ b/docs/transient/03xsections/1d_inhom_elements.ipynb @@ -204,7 +204,7 @@ " ax1.plot(x, h[1].squeeze(), label=f\"t={t[i]:.0f} d\", ls=\"dashed\")\n", "\n", "hss = mlss.headalongline(x, y)\n", - "ax0.plot(x, hss[0].squeeze(), lw=1.5, color=\"k\", ls=\"dotted\", label=\"timflow\")\n", + "ax0.plot(x, hss[0].squeeze(), lw=1.5, color=\"k\", ls=\"dotted\", label=\"timflow.steady\")\n", "ax1.plot(x, hss[1].squeeze(), lw=1.5, color=\"k\", ls=\"dotted\")\n", "\n", "ax0.legend(loc=(0, 1), frameon=False, ncol=2)\n", @@ -509,7 +509,7 @@ " -np.inf,\n", " 0.0,\n", " kaq=k,\n", - " Saq=Saq,\n", + " Saq=S,\n", " z=z,\n", " c=c,\n", " topboundary=\"semi\",\n", @@ -520,7 +520,7 @@ " 0.0,\n", " np.inf,\n", " kaq=k,\n", - " Saq=Saq,\n", + " Saq=S,\n", " z=z,\n", " c=c,\n", " topboundary=\"semi\",\n", @@ -538,13 +538,7 @@ "metadata": {}, "outputs": [], "source": [ - "mlss = tfs.ModelMaq(\n", - " kaq=k,\n", - " z=z,\n", - " c=c,\n", - " topboundary=\"semi\",\n", - " hstar=0.0,\n", - ")\n", + "mlss = tfs.ModelXsection(naq=1)\n", "\n", "inhom_left_ss = tfs.XsectionMaq(\n", " mlss,\n", @@ -575,10 +569,20 @@ "metadata": {}, "outputs": [], "source": [ - "mlsemi.plots.xsection(\n", + "# plot transient xsec for parameters\n", + "ax = mlsemi.plots.xsection(\n", " xy=[(-100, 0), (100, 0)],\n", " params=True,\n", " labels=False,\n", + " hstar=np.nan,\n", + " boundaries=False,\n", + ")\n", + "# plot steady to get correct representation of hstar\n", + "mlss.plots.xsection(\n", + " xy=[(-100, 0), (100, 0)],\n", + " params=False,\n", + " labels=False,\n", + " ax=ax,\n", ");" ] }, diff --git a/docs/utils/calibrating_timflow_models.ipynb b/docs/utils/calibrating_timflow_models.ipynb new file mode 100644 index 00000000..4e1b08c4 --- /dev/null +++ b/docs/utils/calibrating_timflow_models.ipynb @@ -0,0 +1,969 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "9efa0ee9", + "metadata": {}, + "source": [ + "# Calibration of steady and transient models\n", + "\n", + "This notebook introduces the `timflow.Calibration` class, which can be used to\n", + "calibrate steady and/or transient models to head observations.\n", + "\n", + "## Contents\n", + "\n", + "- [Steady calibration](#steady-calibration)\n", + "- [Transient calibration](#transient-calibration)\n", + "- [Combined calibration](#combined-calibration)\n", + "- [Calibrating on time series with constant offsets](#calibrating-on-time-series-with-constant-offsets)\n", + "- [Calibrating on time series with time shifts](#calibrating-on-time-series-with-time-shifts)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e5d84887", + "metadata": {}, + "outputs": [], + "source": [ + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "\n", + "import timflow as tf\n", + "import timflow.steady as tfs\n", + "import timflow.transient as tft" + ] + }, + { + "cell_type": "markdown", + "id": "aaec704e", + "metadata": {}, + "source": [ + "## Steady calibration\n", + "\n", + "In this section we will test the calibration of a steady model. First, we build a model\n", + "to represent the \"truth\"." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c467f670", + "metadata": {}, + "outputs": [], + "source": [ + "ml_s = tfs.ModelXsection(naq=1)\n", + "river_s = tfs.XsectionMaq(\n", + " ml_s,\n", + " x1=-np.inf,\n", + " x2=0.0,\n", + " kaq=[10.0],\n", + " z=[0.1, 0, -10],\n", + " c=[1e-4],\n", + " npor=0.3,\n", + " topboundary=\"semi\",\n", + " hstar=2,\n", + " name=\"river\",\n", + ")\n", + "\n", + "polder_s = tfs.XsectionMaq(\n", + " ml_s,\n", + " x1=0.0,\n", + " x2=np.inf,\n", + " kaq=[10.0],\n", + " z=[1, 0, -10],\n", + " c=[500],\n", + " npor=0.3,\n", + " topboundary=\"semi\",\n", + " hstar=0,\n", + " name=\"polder\",\n", + ")\n", + "ml_s.solve()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a3b6c07c", + "metadata": {}, + "outputs": [], + "source": [ + "fig, ax = plt.subplots(figsize=(10, 3))\n", + "ax.set_ylim(-10.5, 3.5)\n", + "ml_s.plots.xsection(\n", + " xy=([-100, 0], [100, 0]), names=True, labels=True, params=True, ax=ax\n", + ");" + ] + }, + { + "cell_type": "markdown", + "id": "03ec6d3d", + "metadata": {}, + "source": [ + "Select 3 head observations that we will use in our calibration." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7b26cc78", + "metadata": {}, + "outputs": [], + "source": [ + "x = np.linspace(-50, 1000, 101)\n", + "h0 = ml_s.headalongline(x=x, y=np.zeros_like(x))\n", + "\n", + "x_obs = [10, 100, 500]\n", + "h_obs = ml_s.headalongline(x=x_obs, y=np.zeros_like(x_obs))[0]\n", + "\n", + "plt.figure(figsize=(10, 3))\n", + "plt.plot(x, h0[0], label='\"Truth\"')\n", + "plt.plot(x_obs, h_obs, \"kx\", label=\"observations\")\n", + "plt.legend(loc=(0, 1), frameon=False, ncol=2)\n", + "plt.ylabel(\"head [m]\")\n", + "plt.xlabel(\"distance from river [m]\")\n", + "plt.grid()" + ] + }, + { + "cell_type": "markdown", + "id": "ba384cd1", + "metadata": {}, + "source": [ + "Perform the calibration. We are using the same model we created before, but we are\n", + "purposely setting the initial parameter estimates to different values. " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "323a4f32", + "metadata": {}, + "outputs": [], + "source": [ + "cal = tf.Calibrate(steady_model=ml_s)\n", + "\n", + "for i in range(len(x_obs)):\n", + " cal.add_steady_head(name=f\"obs{i}\", x=x_obs[i], y=0.0, layer=[0], h=h_obs[i])\n", + "\n", + "cal.set_aquifer_parameter(\"kaq\", layers=[0], initial=2.0, inhoms=[\"river\", \"polder\"])\n", + "cal.set_aquifer_parameter(\"c\", layers=[0], initial=10.0, inhoms=[\"polder\"])\n", + "\n", + "cal.fit()" + ] + }, + { + "cell_type": "markdown", + "id": "d724ed73", + "metadata": {}, + "source": [ + "View the parameters dataframe:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d229242c", + "metadata": {}, + "outputs": [], + "source": [ + "cal.parameters.style.format(formatter=\"{:.2f}\", subset=[\"initial\", \"optimal\"])" + ] + }, + { + "cell_type": "markdown", + "id": "879e5d1e", + "metadata": {}, + "source": [ + "Plot the results and compare to the truth." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b9d37f8e", + "metadata": {}, + "outputs": [], + "source": [ + "hc = ml_s.headalongline(x=x, y=np.zeros_like(x))\n", + "\n", + "plt.figure(figsize=(10, 3))\n", + "plt.plot(x, h0[0], label='\"Truth\"')\n", + "plt.plot(x, hc[0], ls=\"dashed\", label=\"calibrated\")\n", + "plt.plot(x_obs, h_obs, \"kx\", label=\"observations\")\n", + "plt.legend(loc=(0, 1), frameon=False, ncol=3)\n", + "plt.ylabel(\"head [m]\")\n", + "plt.xlabel(\"distance from river [m]\")\n", + "plt.grid()" + ] + }, + { + "cell_type": "markdown", + "id": "3f5c3d6a", + "metadata": {}, + "source": [ + "## Transient calibration\n", + "\n", + "In this section, we will test the calibration of a transient model with a sudden rise\n", + "of 2 m in river water level after t=0.1 days." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d9d27971", + "metadata": {}, + "outputs": [], + "source": [ + "ml_t = tft.ModelXsection(naq=1, tmin=1e-3, tmax=100)\n", + "\n", + "river_t = tft.XsectionMaq(\n", + " ml_t,\n", + " x1=-np.inf,\n", + " x2=0.0,\n", + " kaq=[10.0],\n", + " z=[0.1, 0, -10],\n", + " c=[1e-4],\n", + " Saq=[1e-3],\n", + " topboundary=\"semi\",\n", + " tsandhstar=[(0.1, 2)],\n", + " name=\"river\",\n", + ")\n", + "\n", + "polder_t = tft.XsectionMaq(\n", + " ml_t,\n", + " x1=0.0,\n", + " x2=np.inf,\n", + " kaq=[10.0],\n", + " z=[1, 0, -10],\n", + " c=[500],\n", + " Saq=[1e-3],\n", + " topboundary=\"semi\",\n", + " name=\"polder\",\n", + ")\n", + "ml_t.solve()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "fd8c0c30", + "metadata": {}, + "outputs": [], + "source": [ + "fig, ax = plt.subplots(figsize=(10, 3))\n", + "ax.set_ylim(-10.5, 3.5)\n", + "ml_t.plots.xsection(\n", + " xy=([-100, 0], [100, 0]),\n", + " names=True,\n", + " labels=True,\n", + " params=True,\n", + " ax=ax,\n", + " sep=\"\\n\",\n", + " hstar=2.0,\n", + ");" + ] + }, + { + "cell_type": "markdown", + "id": "49195c34", + "metadata": {}, + "source": [ + "Plot the head over time along the cross-section. We will extract the heads at the three\n", + "observation points we used earlier as our observation time series." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e64558ee", + "metadata": {}, + "outputs": [], + "source": [ + "t = np.logspace(-1, 1, 11)\n", + "h0 = ml_t.headalongline(x=x, y=np.zeros_like(x), t=t)\n", + "\n", + "plt.figure(figsize=(10, 3))\n", + "for i in range(len(t)):\n", + " plt.plot(x, h0[0, i], label=f\"t={t[i]:.1f} d\")\n", + "plt.legend(loc=(0, 1), frameon=False, ncol=6, fontsize=\"small\")\n", + "plt.ylabel(\"head [m]\")\n", + "plt.xlabel(\"distance from river [m]\")\n", + "plt.grid()" + ] + }, + { + "cell_type": "markdown", + "id": "5c2e0f00", + "metadata": {}, + "source": [ + "Plot our observation data. Note the log-scale of the x-axis." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "bc1f0025", + "metadata": {}, + "outputs": [], + "source": [ + "h_obs_series = ml_t.headalongline(x=x_obs, y=np.zeros_like(x_obs), t=t)\n", + "\n", + "plt.figure(figsize=(10, 3))\n", + "for i in range(len(x_obs)):\n", + " plt.plot(t, h_obs_series[0, :, i], label=f\"obs{i} (x={x_obs[i]})\", marker=\"o\")\n", + "plt.legend(loc=(0, 1), frameon=False, ncol=3)\n", + "plt.ylabel(\"head [m]\")\n", + "plt.xlabel(\"time [d]\")\n", + "# plt.xscale(\"log\")\n", + "plt.grid()" + ] + }, + { + "cell_type": "markdown", + "id": "79efd10f", + "metadata": {}, + "source": [ + "Create the calibration class, add the observation time series, set the calibration\n", + "parameters, and calibrate the model." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "899a2c05", + "metadata": {}, + "outputs": [], + "source": [ + "cal = tf.Calibrate(transient_model=ml_t)\n", + "\n", + "for i in range(len(x_obs)):\n", + " cal.add_head_time_series(\n", + " name=f\"obs{i}\",\n", + " x=x_obs[i],\n", + " y=0.0,\n", + " layer=[0],\n", + " t=t,\n", + " h=h_obs_series[0, :, i],\n", + " normalized=True,\n", + " )\n", + "\n", + "cal.set_aquifer_parameter(\"kaq\", layers=[0], initial=2.0, inhoms=[\"river\", \"polder\"])\n", + "cal.set_aquifer_parameter(\"c\", layers=[0], initial=10.0, inhoms=[\"polder\"])\n", + "cal.set_aquifer_parameter(\"Saq\", layers=[0], initial=1e-4, inhoms=[\"polder\"])\n", + "\n", + "cal.fit()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e349649b", + "metadata": {}, + "outputs": [], + "source": [ + "cal.parameters.style.format(formatter=\"{:.2e}\", subset=[\"initial\", \"optimal\"])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "bb95ca5b", + "metadata": {}, + "outputs": [], + "source": [ + "plt.figure(figsize=(10, 3))\n", + "hm = ml_t.headalongline(x=x_obs, y=np.zeros_like(x_obs), t=t)\n", + "for i in range(len(x_obs)):\n", + " plt.plot(\n", + " t, h_obs_series[0, :, i], label=f\"obs{i} (x={x_obs[i]})\", marker=\"x\", ls=\"none\"\n", + " )\n", + " plt.plot(t, hm[0, :, i], label=f\"sim{i}\", color=f\"C{i}\")\n", + "\n", + "plt.xscale(\"log\")\n", + "plt.legend(loc=(0, 1), frameon=False, ncol=6, fontsize=\"small\")\n", + "plt.xlabel(\"time [days]\")\n", + "plt.ylabel(\"head [m]\")\n", + "plt.grid()" + ] + }, + { + "cell_type": "markdown", + "id": "3a05b9bf", + "metadata": {}, + "source": [ + "## Combined calibration\n", + "\n", + "Next, we want to attempt a combined calibration. This is done by superposition of a\n", + "steady model, representing the average situation (i.e. the average river stage), and a\n", + "transient model that captures the effect of a sudden change in the river water level.\n", + "The average river stage is 2 m+ref, the sudden rise is 2m at t=0.1 days.\n", + "\n", + "We build 2 models, a steady and transient model. Note that we give the `XSection`\n", + "elements the same names in both models. This allows the calibration class to share\n", + "parameters between zones (e.g. the polder) across both models.\n", + "\n", + "The steady model is added to the transient model. This means the transient model will\n", + "use the steady model to compute heads and flows: $h = h_\\text{steady} + h_\\text{transient}$." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "98217e4d", + "metadata": {}, + "outputs": [], + "source": [ + "ml_s = tfs.ModelXsection(naq=1)\n", + "river_s = tfs.XsectionMaq(\n", + " ml_s,\n", + " x1=-np.inf,\n", + " x2=0.0,\n", + " kaq=[10.0],\n", + " z=[1, 0, -10],\n", + " c=[1e-4],\n", + " npor=0.3,\n", + " topboundary=\"semi\",\n", + " hstar=2,\n", + " name=\"river\",\n", + ")\n", + "\n", + "polder_s = tfs.XsectionMaq(\n", + " ml_s,\n", + " x1=0.0,\n", + " x2=np.inf,\n", + " kaq=[10.0],\n", + " z=[1, 0, -10],\n", + " c=[500],\n", + " npor=0.3,\n", + " topboundary=\"semi\",\n", + " hstar=0,\n", + " name=\"polder\",\n", + ")\n", + "\n", + "# add steady model to transient model\n", + "ml_t = tft.ModelXsection(naq=1, tmin=1e-3, tmax=100, steady=ml_s)\n", + "\n", + "river_t = tft.XsectionMaq(\n", + " ml_t,\n", + " x1=-np.inf,\n", + " x2=0.0,\n", + " kaq=[10.0],\n", + " z=[1, 0, -10],\n", + " c=[1e-4],\n", + " Saq=[1e-3],\n", + " topboundary=\"semi\",\n", + " tsandhstar=[(0.1, 2)],\n", + " name=\"river\",\n", + ")\n", + "\n", + "polder_t = tft.XsectionMaq(\n", + " ml_t,\n", + " x1=0.0,\n", + " x2=np.inf,\n", + " kaq=[10.0],\n", + " z=[1, 0, -10],\n", + " c=[500],\n", + " Saq=[1e-3],\n", + " topboundary=\"semi\",\n", + " name=\"polder\",\n", + ")\n", + "ml_t.solve()" + ] + }, + { + "cell_type": "markdown", + "id": "9e3764d7", + "metadata": {}, + "source": [ + "Plot the change in head over time." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "403f0f6e", + "metadata": {}, + "outputs": [], + "source": [ + "t = np.logspace(-1, 1, 11)\n", + "h0 = ml_t.headalongline(x=x, y=np.zeros_like(x), t=t)\n", + "\n", + "plt.figure(figsize=(10, 3))\n", + "for i in range(len(t)):\n", + " plt.plot(x, h0[0, i], label=f\"t={t[i]:.1f} d\")\n", + "plt.legend(loc=(0, 1), frameon=False, ncol=6, fontsize=\"small\")\n", + "plt.ylabel(\"head [m]\")\n", + "plt.xlabel(\"distance from river [m]\")\n", + "plt.grid()" + ] + }, + { + "cell_type": "markdown", + "id": "b5c62ad3", + "metadata": {}, + "source": [ + "Now we get our head time series from the model to use as our observation time series in\n", + "the calibration." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0b4d5499", + "metadata": {}, + "outputs": [], + "source": [ + "h_obs_series = ml_t.headalongline(x=x_obs, y=np.zeros_like(x_obs), t=t)\n", + "\n", + "plt.figure(figsize=(10, 3))\n", + "for i in range(len(x_obs)):\n", + " plt.plot(t, h_obs_series[0, :, i], label=f\"obs{i} (x={x_obs[i]})\", marker=\"x\", ls=\"-\")\n", + "plt.legend(loc=(0, 1), frameon=False, ncol=3)\n", + "plt.ylabel(\"head [m]\")\n", + "plt.xlabel(\"time [d]\")\n", + "plt.grid()\n", + "plt.xscale(\"log\")" + ] + }, + { + "cell_type": "markdown", + "id": "4fe33bc6", + "metadata": {}, + "source": [ + "Now we will calibrate both models simultaneously by passing both models to the\n", + "calibrate class. We need to add the steady model separately so that the calibration\n", + "class can link the parameters between the models.\n", + "\n", + "We add the head observations as before. \n", + "\n", + "Note that when defining the calibration parameters, we can now set the names of the\n", + "`XSections` we want to target, as well as the `model`. The `model` can be `\"steady\"`,\n", + "`\"transient\"` or `\"both\"`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e5ba87ab", + "metadata": {}, + "outputs": [], + "source": [ + "cal = tf.Calibrate(transient_model=ml_t, steady_model=ml_s)\n", + "\n", + "for i in range(len(x_obs)):\n", + " cal.add_head_time_series(\n", + " name=f\"obs{i}\",\n", + " x=x_obs[i],\n", + " y=0.0,\n", + " layer=[0],\n", + " t=t,\n", + " h=h_obs_series[0, :, i],\n", + " normalized=True,\n", + " )\n", + "\n", + "cal.set_aquifer_parameter(\n", + " \"kaq\",\n", + " layers=[0],\n", + " initial=2.0,\n", + " pmin=2.0,\n", + " pmax=100.0,\n", + " inhoms=[\"polder\"],\n", + " model=\"both\",\n", + ")\n", + "cal.set_aquifer_parameter(\n", + " \"c\",\n", + " layers=[0],\n", + " initial=1000.0,\n", + " pmin=1.0,\n", + " pmax=10_000,\n", + " inhoms=[\"polder\"],\n", + " model=\"both\",\n", + ")\n", + "cal.set_aquifer_parameter(\n", + " \"Saq\",\n", + " layers=[0],\n", + " initial=1e-4,\n", + " inhoms=[\"polder\"],\n", + " pmin=1e-10,\n", + " pmax=1.0,\n", + " model=\"transient\",\n", + " log_scale=True,\n", + ")\n", + "\n", + "# cal.lmfit()\n", + "cal.fit()" + ] + }, + { + "cell_type": "markdown", + "id": "018aee81", + "metadata": {}, + "source": [ + "Let's view the parameters dataframe" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "629cb4f3", + "metadata": {}, + "outputs": [], + "source": [ + "cal.parameters.style.format(\n", + " formatter=\"{:.2e}\", subset=[\"initial\", \"optimal\", \"pmin\", \"pmax\"]\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "934f75a2", + "metadata": {}, + "outputs": [], + "source": [ + "plt.figure(figsize=(10, 3))\n", + "hm = ml_t.headalongline(x=x_obs, y=np.zeros_like(x_obs), t=t)\n", + "for i in range(len(x_obs)):\n", + " plt.plot(\n", + " t, h_obs_series[0, :, i], label=f\"obs{i} (x={x_obs[i]})\", marker=\"x\", ls=\"none\"\n", + " )\n", + " plt.plot(t, hm[0, :, i], label=f\"sim{i}\", color=f\"C{i}\")\n", + "\n", + "plt.xscale(\"log\")\n", + "plt.legend(loc=(0, 1), frameon=False, ncol=6)\n", + "plt.xlabel(\"time [days]\")\n", + "plt.ylabel(\"head [m]\")\n", + "plt.grid()" + ] + }, + { + "cell_type": "markdown", + "id": "1ddb5b4b", + "metadata": {}, + "source": [ + "## Calibrating on time series with constant offsets\n", + "\n", + "The Calibrate class supports adding unknown constants for each head time series. This\n", + "is useful when the data has been measured relative to some reference level, but we are\n", + "only interested in the variation in heads. This could be used for calibrating to head\n", + "changes, while ignoring the absolute levels. Or, it could be used to calibrate on head\n", + "observations that lie along a river, where the reference level might change slightly\n", + "as we move downstream along that river.\n", + "\n", + "In this example, we're simply adding an increasing constant value (1, 2 or 3 m) to\n", + "each head observation from the previous example.\n", + "\n", + "Note that when adding the head time series, we now specify the constant and provide it\n", + "with a tuple containing three values: `constant = (initial, pmin, pmax)`. This\n", + "represents the initial guess of the reference level, and the upper and lower bounds.\n", + "The upper and lower bounds are optional and may be ommitted by passing a single float." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "805185af", + "metadata": {}, + "outputs": [], + "source": [ + "cal = tf.Calibrate(transient_model=ml_t, steady_model=ml_s)\n", + "\n", + "for i in range(len(x_obs)):\n", + " cal.add_head_time_series(\n", + " name=f\"obs{i}\",\n", + " x=x_obs[i],\n", + " y=0.0,\n", + " layer=[0],\n", + " t=t,\n", + " h=h_obs_series[0, :, i] + (i + 1.0),\n", + " constant=(0.1, -10, 10),\n", + " )\n", + "\n", + "cal.set_aquifer_parameter(\n", + " \"kaq\",\n", + " layers=[0],\n", + " initial=2.0,\n", + " pmin=2.0,\n", + " pmax=100.0,\n", + " inhoms=[\"polder\"],\n", + " model=\"both\",\n", + ")\n", + "cal.set_aquifer_parameter(\n", + " \"c\",\n", + " layers=[0],\n", + " initial=1000.0,\n", + " pmin=1.0,\n", + " pmax=10_000,\n", + " inhoms=[\"polder\"],\n", + " model=\"both\",\n", + ")\n", + "cal.set_aquifer_parameter(\n", + " \"Saq\",\n", + " layers=[0],\n", + " initial=1e-4,\n", + " inhoms=[\"polder\"],\n", + " pmin=1e-10,\n", + " pmax=1.0,\n", + " model=\"transient\",\n", + " log_scale=True,\n", + ")\n", + "\n", + "# cal.lmfit()\n", + "cal.fit()" + ] + }, + { + "cell_type": "markdown", + "id": "91885d05", + "metadata": {}, + "source": [ + "The parameters dataframe shows that we were able to get the correct values of the\n", + "constants." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5852d1ed", + "metadata": {}, + "outputs": [], + "source": [ + "cal.parameters.style.format(\n", + " formatter=\"{:.2f}\", subset=[\"initial\", \"optimal\", \"pmin\", \"pmax\"]\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "9bcfb5e0", + "metadata": {}, + "source": [ + "Plot the calibration results. We now use the head time series objects contained within\n", + "the Calibration class to plot the observations. When we apply the constants, we see the\n", + "model fits perfectly with the data. If we set `apply_constant=False`, we would see that\n", + "the observations are shifted vertically relative to the model simulation." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "40ea3054", + "metadata": {}, + "outputs": [], + "source": [ + "plt.figure(figsize=(10, 3))\n", + "hm = ml_t.headalongline(x=x_obs, y=np.zeros_like(x_obs), t=t)\n", + "for i in range(len(x_obs)):\n", + " cal.observations_dict[f\"obs{i}\"].plot(\n", + " ax=plt.gca(), marker=\"x\", ls=\"none\", apply_constant=True\n", + " )\n", + " plt.plot(t, hm[0, :, i], label=f\"sim{i}\", color=f\"C{i}\")\n", + "\n", + "plt.xscale(\"log\")\n", + "plt.legend(loc=(0, 1), frameon=False, ncol=6)\n", + "plt.xlabel(\"time [days]\")\n", + "plt.ylabel(\"head [m]\")\n", + "plt.grid()" + ] + }, + { + "cell_type": "markdown", + "id": "c096114f", + "metadata": {}, + "source": [ + "## Calibrating on time series with time shifts\n", + "\n", + "Another option the new Calibration class provides is to set time shift parameters in the calibration. This allows the optimization to shift the head time series in time. This can be useful when there are phase shifts in the data that are not represented in our model, e.g. when the observations lie along a river, where a flood wave might arrive slightly earlier in upstream observation wells relative to the downstream observation wells. \n", + "\n", + "In this example, we're simply adding an increasing time shift (0.1, 0.2 and 0.3 days) to\n", + "each head observation series from the previous example.\n", + "\n", + "Note that when adding the head time series, we now specify the time shift and provide it\n", + "with a tuple containing three values: `time_shift = (initial, pmin, pmax)`. This\n", + "represents the initial guess of the time shift, and the upper and lower bounds.\n", + "The upper and lower bounds are optional and may be ommitted by passing a single float.\n", + "\n", + "We recreate the transient cross-section model here to reduce the order of the `tmin`\n", + "parameter, because shifting the observation time series can cause observations to lie\n", + "close to the changes in boundary conditions, which will introduce NaNs into the\n", + "simulation. These are filtered out, but the optimization might still be affected, so it\n", + "is recommended to also reduce the `tmin` to avoid this as much as possible." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1e6a8fbe", + "metadata": {}, + "outputs": [], + "source": [ + "# add steady model to transient model\n", + "ml_t = tft.ModelXsection(naq=1, tmin=1e-6, tmax=100, steady=ml_s)\n", + "\n", + "river_t = tft.XsectionMaq(\n", + " ml_t,\n", + " x1=-np.inf,\n", + " x2=0.0,\n", + " kaq=[10.0],\n", + " z=[1, 0, -10],\n", + " c=[1e-4],\n", + " Saq=[1e-3],\n", + " topboundary=\"semi\",\n", + " tsandhstar=[(0.1, 2)],\n", + " name=\"river\",\n", + ")\n", + "\n", + "polder_t = tft.XsectionMaq(\n", + " ml_t,\n", + " x1=0.0,\n", + " x2=np.inf,\n", + " kaq=[10.0],\n", + " z=[1, 0, -10],\n", + " c=[500],\n", + " Saq=[1e-3],\n", + " topboundary=\"semi\",\n", + " name=\"polder\",\n", + ")\n", + "ml_t.solve()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "bcdff213", + "metadata": {}, + "outputs": [], + "source": [ + "cal = tf.Calibrate(transient_model=ml_t, steady_model=ml_s)\n", + "\n", + "for i in range(len(x_obs)):\n", + " cal.add_head_time_series(\n", + " name=f\"obs{i}\",\n", + " x=x_obs[i],\n", + " y=0.0,\n", + " layer=[0],\n", + " t=t + 0.1 * (i + 1),\n", + " h=h_obs_series[0, :, i],\n", + " time_shift=(1 / 24.0, 0.0, 1.0), # initial guess is we're 1 hour off\n", + " )\n", + "\n", + "cal.set_aquifer_parameter(\n", + " \"kaq\",\n", + " layers=[0],\n", + " initial=2.0,\n", + " pmin=2.0,\n", + " pmax=100.0,\n", + " inhoms=[\"polder\"],\n", + " model=\"both\",\n", + ")\n", + "cal.set_aquifer_parameter(\n", + " \"c\",\n", + " layers=[0],\n", + " initial=1000.0,\n", + " pmin=1.0,\n", + " pmax=10_000,\n", + " inhoms=[\"polder\"],\n", + " model=\"both\",\n", + ")\n", + "cal.set_aquifer_parameter(\n", + " \"Saq\",\n", + " layers=[0],\n", + " initial=1e-4,\n", + " inhoms=[\"polder\"],\n", + " pmin=1e-10,\n", + " pmax=1.0,\n", + " model=\"transient\",\n", + " log_scale=True,\n", + ")\n", + "\n", + "# cal.lmfit()\n", + "cal.fit()" + ] + }, + { + "cell_type": "markdown", + "id": "e7bc239b", + "metadata": {}, + "source": [ + "We got a few warnings about the computation time lying to close to a change in boundary\n", + "condition, but the results looks good nonetheless. The parameters dataframe shows the\n", + "time shifts are estimated correctly." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3d2b392b", + "metadata": {}, + "outputs": [], + "source": [ + "cal.parameters.style.format(\n", + " formatter=\"{:.2f}\", subset=[\"initial\", \"optimal\", \"pmin\", \"pmax\"]\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "34a3b64c", + "metadata": {}, + "source": [ + "Plot the calibration results. We now use the head time series objects contained within\n", + "the Calibration class to plot the observations. When we apply the time shifts, we see the\n", + "model fits perfectly with the data. If we set `apply_time_shift=False`, we would see that\n", + "the observations are shifted in time relative to the model simulation." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6e237ba0", + "metadata": {}, + "outputs": [], + "source": [ + "plt.figure(figsize=(10, 3))\n", + "hm = ml_t.headalongline(x=x_obs, y=np.zeros_like(x_obs), t=t)\n", + "for i in range(len(x_obs)):\n", + " cal.observations_dict[f\"obs{i}\"].plot(\n", + " ax=plt.gca(), marker=\"x\", ls=\"none\", apply_time_shift=True\n", + " )\n", + " plt.plot(t, hm[0, :, i], label=f\"sim{i}\", color=f\"C{i}\")\n", + "\n", + "plt.xscale(\"log\")\n", + "plt.legend(loc=(0, 1), frameon=False, ncol=6)\n", + "plt.xlabel(\"time [days]\")\n", + "plt.ylabel(\"head [m]\")\n", + "plt.grid()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cda26f2f", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "timflow (3.13.11)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.13.11" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/timflow/__init__.py b/timflow/__init__.py index 2e684d84..b268c5be 100644 --- a/timflow/__init__.py +++ b/timflow/__init__.py @@ -9,4 +9,5 @@ # ruff : noqa: F401 from timflow import steady, transient +from timflow.calibrate import Calibrate from timflow.version import __version__, show_versions diff --git a/timflow/calibrate.py b/timflow/calibrate.py new file mode 100644 index 00000000..db4a42ec --- /dev/null +++ b/timflow/calibrate.py @@ -0,0 +1,1553 @@ +"""Unified calibration framework for transient and steady-state models. + +Supports independent and joint calibration where models share parameters. + +Example:: + + # setup joint calibration + cal = Calibrate(transient_model=ml_t, steady_model=ml_s) + + # add observations, either a steady head observation or a time series of heads + cal.add_head_time_series(name='obs1', x=0, y=0, layer=0, t=t, h=h) + cal.add_steady_head(name='obs2', x=0, y=0, layer=0, h=h_steady) + + # set calibration parameters + cal.set_aquifer_parameter('kaq', layers=0, initial=10, pmin=1, pmax=100, model="both") + + # calibrate model + cal.fit() +""" + +import warnings +from dataclasses import dataclass, field +from typing import Any, Iterable, Literal, Optional + +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +from scipy.linalg import LinAlgError, get_lapack_funcs, svd +from scipy.optimize import least_squares + +from timflow.steady.element import Element as SteadyElement +from timflow.steady.model import Model as SteadyModel +from timflow.transient.element import Element as TransientElement +from timflow.transient.model import Model as TransientModel + + +@dataclass +class ParameterTarget: + """A single array slice in a model that a Parameter controls.""" + + array: np.ndarray + key: slice # the slice into `array` + + def set(self, value: float) -> None: + self.array[self.key] = value + + def get(self) -> np.ndarray: + return self.array[self.key] + + +@dataclass +class Parameter: + """A single optimizable parameter that may affect multiple model arrays. + + Parameters can span both transient and steady-state models, enabling + joint calibration with shared parameters. + """ + + name: str + initial: float | None + pmin: float = -np.inf + pmax: float = np.inf + targets: list[ParameterTarget] = field(default_factory=list) + models: list[str] = field(default_factory=list) + inhoms: list[str] = field(default_factory=list) + log_scale: bool = False + + # filled after fitting + optimal: float = np.nan + std: Optional[float] = None + + @property + def effective_initial(self) -> float: + """Effective starting value for optimization. + + Returns ``initial`` if set explicitly, otherwise reads the current + value from the model parameter arrays. This allows calibration to + start from the current model state when ``initial=None``. + """ + if self.initial is not None: + return self.initial + return self.get() + + def set(self, value: float) -> None: + """Push value to all model arrays this parameter controls.""" + if self.initial is None: + initial = float(self.targets[0].get()[0]) # read current model value + else: + initial = self.initial + if value is None: + value = initial + if self.log_scale: + value = np.sign(initial) * 10**value + for target in self.targets: + target.set(value) + + def get(self) -> float: + """Read back current value from first target.""" + if self.targets or self.initial is None: + value = float(self.targets[0].get()[0]) + else: + value = self.initial + return value + + def add_target( + self, array: np.ndarray, key: slice, model: Any = None, inhom: Any = None + ) -> None: + self.targets.append(ParameterTarget(array, key)) + if model is not None: + self.models.append(model.model_type) + if inhom is not None: + self.inhoms.append(inhom.name) + + +@dataclass +class SteadyHead: + """A single steady-state head observation at a point in the aquifer. + + Attributes + ---------- + x, y : float + Observation coordinates. + layer : int + Aquifer layer index (0-based). + h : float + Observed head. + weight : float + Weight applied to the residual in the objective function. + """ + + x: float + y: float + layer: int + h: float + weight: float = 1.0 + model_key: str = field(default="steady", init=False) # 'steady' + + +@dataclass +class SteadyHeadInWell: + """A single steady-state head observation in a well. + + Attributes + ---------- + element : SteadyElement + Well element in which the head is observed. + x, y : float + Location of the element. + layer : int + Aquifer layer index (0-based). + h : float + Observed head. + weight : float + Weight applied to the residual in the objective function. + """ + + element: SteadyElement + x: float + y: float + layer: int + h: float + weight: float = 1.0 + model_key: str = field(default="steady", init=False) # 'steady' + + +@dataclass +class HeadSeries: + """A transient head observation time series at a point in the aquifer. + + Attributes + ---------- + x, y : float + Observation coordinates. + layer : int + Aquifer layer index (0-based). + t : np.ndarray + Observation times. + h : np.ndarray + Observed heads. + weights : float, np.ndarray, optional + Per time-series (float) or per-timestep weights (array). Defaults to uniform + weight of 1.0 if ``None``. + constant : float or (float, float, float), optional + If not ``None``, a constant offset is added as a calibration + parameter. Supply a float for the initial value (unbounded), or a + ``(initial, pmin, pmax)`` tuple to set bounds. + time_shift : float or (float, float, float), optional + If not ``None``, a time shift is added as a calibration parameter. + Supply a float for the initial value (unbounded), or a + ``(initial, pmin, pmax)`` tuple to set bounds. + """ + + x: float + y: float + layer: int + t: np.ndarray + h: np.ndarray + weights: Optional[np.ndarray] = None + normalized: bool = False + model_key: str = field(default="transient", init=False) # 'transient' + constant: float | tuple[float, float, float] | None = ( + None # constant parameter and optional bounds + ) + time_shift: float | tuple[float, float, float] | None = ( + None # time shift and optional bounds + ) + # placeholder for constant and time_shift parameters + _constant: np.ndarray = field(default_factory=lambda: np.zeros(1), init=False) + _time_shift: np.ndarray = field(default_factory=lambda: np.zeros(1), init=False) + + def plot( + self, + ax: plt.Axes | None = None, + apply_time_shift: bool = True, + apply_constant: bool = True, + **kwargs, + ) -> plt.Axes: + """Plot the observation time series. + + Parameters + ---------- + ax : matplotlib.axes.Axes, optional + Axes to plot on. A new figure is created if ``None``. + apply_time_shift : bool + Subtract the fitted time shift from ``t`` before plotting. + Has no effect if no time-shift parameter was added. + apply_constant : bool + Subtract the fitted constant from ``h`` before plotting. + Has no effect if no constant parameter was added. + **kwargs + Passed to :func:`matplotlib.axes.Axes.plot`. + + Returns + ------- + matplotlib.axes.Axes + """ + if ax is None: + _, ax = plt.subplots() + t = ( + self.t - self._time_shift + if (self.time_shift is not None) and apply_time_shift + else self.t + ) + h = ( + self.h - self._constant + if (self.constant is not None) and apply_constant + else self.h + ) + ax.plot(t, h, **kwargs) + return ax + + +@dataclass +class HeadSeriesInWell: + """A transient head observation time series in a well. + + Attributes + ---------- + element : TransientElement + Well element at which heads are observed. + t : np.ndarray + Observation times. + h : np.ndarray + Observed heads. + constant : float or (float, float, float), optional + If not ``None``, a constant offset is added as a calibration + parameter. Supply a float for the initial value (unbounded), or a + ``(initial, pmin, pmax)`` tuple to set bounds. + time_shift : float or (float, float, float), optional + If not ``None``, a time shift is added as a calibration parameter. + Supply a float for the initial value (unbounded), or a + ``(initial, pmin, pmax)`` tuple to set bounds. + """ + + element: TransientElement + t: np.ndarray + h: np.ndarray + model_key: str = field(default="transient", init=False) # 'transient' + constant: float | tuple[float, float, float] | None = ( + None # constant parameter and optional bounds + ) + time_shift: float | tuple[float, float, float] | None = ( + None # time shift and optional bounds + ) + # placeholder for constant and time_shift parameters + _constant: np.ndarray = field(default_factory=lambda: np.zeros(1), init=False) + _time_shift: np.ndarray = field(default_factory=lambda: np.zeros(1), init=False) + + +class Calibrate: + """Unified calibration class for transient and/or steady-state models. + + Supports independent and joint calibration where both model types share + parameters (e.g., hydraulic conductivity, aquitard resistance). + + + + Parameters + ---------- + transient_model : optional + Transient model instance. + steady_model : optional + Steady-state model instance. + reference_time : float, optional + Specify reference time to compute head changes relative to the head at that + time. The residuals are then computed with the following formula:: + + res = ((sim - sim(t_ref)) - (obs - obs(t_ref))) * w + + The default is None, which uses the following formula for the residuals:: + + res = (sim - obs) * w + + This can be useful for noisy aquifer test data, for example. + + Notes + ----- + When both models are provided, parameters registered via + :meth:`set_aquifer_parameter` default to ``model='both'``, linking them + across models so the optimizer updates both simultaneously. Pass + ``model='transient'`` or ``model='steady'`` to restrict a parameter to + one model. + + Examples + -------- + Transient-only calibration:: + + cal = tf.Calibrate(transient_model=ml_tr) + cal.add_head_time_series("obs1", x=100.0, y=0.0, layer=0, t=t, h=h) + cal.set_aquifer_parameter("kaq", layers=[0], initial=5.0, inhoms=["polder"]) + cal.fit() + + Joint steady/transient calibration with shared parameters:: + + cal = tf.Calibrate(transient_model=ml_tr, steady_model=ml_ss) + cal.add_steady_head("h_mean", x=100.0, y=0.0, layer=0, h=1.2) + cal.add_head_time_series("h_obs", x=100.0, y=0.0, layer=0, t=t, h=h) + cal.set_aquifer_parameter("kaq", layers=[0], initial=5.0, model="both") + cal.fit() + """ + + def __init__( + self, + transient_model: TransientModel | None = None, + steady_model: SteadyModel | None = None, + reference_time: float | None = None, + ): + if transient_model is None and steady_model is None: + raise ValueError("At least one model must be provided.") + self.transient_model = transient_model + self.reference_time = reference_time + self._parameters: dict[str, Parameter] = {} + self.observations_dict: dict[str, HeadSeries | SteadyHead] = {} + self.observations_in_well_dict: dict[ + str, (SteadyHeadInWell | HeadSeriesInWell) + ] = {} + + embedded_steady = getattr(transient_model, "steady", None) + if embedded_steady is not None: + if steady_model is not None and steady_model is not embedded_steady: + warnings.warn( + "The transient model already has a steady model embedded " + "(transient_model.steady). The steady_model argument will be " + "ignored and the embedded steady model will be used instead.", + stacklevel=2, + ) + elif steady_model is None: + warnings.warn( + "The transient model has an embedded steady model " + "(transient_model.steady). Using it as the steady model for " + "this calibration.", + stacklevel=2, + ) + self.steady_model = embedded_steady + else: + self.steady_model = steady_model + + def set_aquifer_parameter( + self, + name: str, + layers: int | Iterable[int], + initial: float | None = None, + pmin: float = -np.inf, + pmax: float = np.inf, + log_scale: bool = False, + inhoms: str | list[str] | None = None, + model: str = "both", + ) -> None: + """Register an aquifer parameter for optimization. + + Parameters + ---------- + name : str + Parameter type: 'kaq', 'Saq', 'c', 'Sll', or 'kzoverkh'. + layers : int or iterable of int + Layer(s) affected. Consecutive layers are grouped under one + scalar parameter. + initial : float + Starting value (in linear space). if None, uses current model parameter + value. In case of coupled parameters between inhoms or models, uses + first value it encounters. + pmin, pmax : float + Bounds for the optimizer. + log_scale : bool + Whether to optimize in log10 space (recommended for parameters like + hydraulic conductivity that can span orders of magnitude). + inhoms : str or list, optional + Inhomogeneity name(s) to target. ``None`` targets the background + aquifer. + model : {'both', 'transient', 'steady'} + Which model(s) this parameter applies to. + + Examples + -------- + Set calibration parameter for hydraulic conductivity:: + + cal.set_aquifer_parameter( + "kaq", layers=[0], initial=1.0, pmin=1.0, pmax=100.0 + ) + + See Also + -------- + set_parameter_by_reference : Register a parameter by supplying the array + reference directly, useful for non-aquifer parameters like well + resistance. + """ + assert isinstance(name, str), "name must be a string" + assert model in ("both", "transient", "steady"), ( + "model must be 'both', 'transient', or 'steady'" + ) + + from_lay, to_lay = self._parse_layers(name, layers) + models = self._resolve_models(model) + pname = self._make_pname(name, from_lay, to_lay, inhoms, models) + + # set up parameter and add target arrays + param = Parameter( + name=pname, initial=initial, pmin=pmin, pmax=pmax, log_scale=log_scale + ) + for ml in models: + for iaq in self._resolve_inhoms(ml, inhoms): + arr = self._get_aquifer_parameter_array(ml, iaq, name) + slc = slice(from_lay, to_lay + 1) + param.add_target(arr, slc, model=ml, inhom=iaq) + + if initial is not None: + if log_scale: + param.set(np.log10(np.abs(initial))) # initialise arrays immediately + else: + param.set(initial) # initialise arrays immediately + self._parameters[pname] = param + + def set_parameter_by_reference( + self, + name: str, + parameter: np.ndarray, + initial: float = 0.0, + pmin: float = -np.inf, + pmax: float = np.inf, + log_scale: bool = False, + ) -> None: + """Register a parameter by supplying the array reference directly. + + Useful for element-level parameters (e.g., well resistance) that are + not part of the standard aquifer parameter arrays. + + Parameters + ---------- + name : str + Unique name for the parameter. + parameter : np.ndarray + Array whose values will be updated during optimization. + initial : float + Initial parameter value. + pmin, pmax : float + Lower and upper bounds for the optimizer. + log_scale : bool + Optimize in log10 space. + + Examples + -------- + Set a calibration parameter by supplying the array reference directly:: + + cal.set_parameter_by_reference( + "well_resistance", well.res, initial=1.0, pmin=0.1, pmax=10.0 + ) + + See Also + -------- + set_aquifer_parameter : Register an aquifer parameter by specifying the type and + layer(s) it applies to. + """ + assert isinstance(parameter, np.ndarray), "parameter must be a numpy array" + if initial is None: + initial = parameter[0] + param = Parameter( + name=name, initial=initial, pmin=pmin, pmax=pmax, log_scale=log_scale + ) + param.add_target(parameter, slice(None)) + if log_scale: + param.set(np.log10(np.abs(initial))) + else: + param.set(initial) + self._parameters[name] = param + + def add_steady_head( + self, name: str, x: float, y: float, layer: int, h: float, weight: float = 1.0 + ) -> None: + """Add a steady-state head observation. + + Parameters + ---------- + name : str + Unique observation name. + x, y : float + Observation coordinates. + layer : int + Aquifer layer index (0-based). + h : float + Observed head. + weight : float, optional + Weight in the objective function. Default is 1. + + Examples + -------- + Add steady head observation at (x=100, y=0) in layer 0 with observed head of 1.5:: + + cal.add_steady_head("piezometer1", x=100.0, y=0.0, layer=0, h=1.5) + """ + if self.steady_model is None: + raise ValueError( + "Steady model must be provided to add steady head observations." + ) + self.observations_dict[name] = SteadyHead( + x=x, y=y, layer=layer, h=h, weight=weight + ) + + def add_steady_head_in_well( + self, name: str, well_element: SteadyElement, h: float, weight: float = 1.0 + ) -> None: + """Add a steady-state head observation inside a well. + + Parameters + ---------- + name : str + Unique observation name. + well_element : SteadyElement + Well element at which the head is observed. + h : float + Observed head. + weight : float, optional + Weight in the objective function. Default is 1. + + Examples + -------- + Add a steady head observation inside a well:: + + pump_well = timflow.steady.Well(...) + cal.add_steady_head_in_well("well1", well_element=pump_well, h=0.8) + """ + if self.steady_model is None: + raise ValueError( + "Steady model must be provided to add steady head observations." + ) + self.observations_in_well_dict[name] = SteadyHeadInWell( + element=well_element, + x=well_element.x, + y=well_element.y, + layer=well_element.layer, + h=h, + weight=weight, + ) + + @staticmethod + def _parse_optional_param( + arg: float | tuple[float, float, float] | None, + default_initial: float, + ) -> tuple[float, float, float] | None: + """Parse None | float | (initial, pmin, pmax) -> (initial, pmin, pmax) or None.""" + if arg is None: + return None + if isinstance(arg, bool): + return (default_initial, -np.inf, np.inf) + if isinstance(arg, (int, float)): + return (arg, -np.inf, np.inf) + return tuple(arg) # (initial, pmin, pmax) + + def _add_series_constant( + self, + name: str, + obs: Any, + initial: float, + pmin: float = -np.inf, + pmax: float = np.inf, + ) -> None: + constant = Parameter(name + "_constant", initial=initial, pmin=pmin, pmax=pmax) + constant.add_target(obs._constant, slice(None)) + self._parameters[name + "_constant"] = constant + + def _add_series_time_shift( + self, + name: str, + obs: Any, + initial: float, + pmin: float = -np.inf, + pmax: float = np.inf, + ) -> None: + time_shift = Parameter( + name + "_time_shift", initial=initial, pmin=pmin, pmax=pmax + ) + time_shift.add_target(obs._time_shift, slice(None)) + self._parameters[name + "_time_shift"] = time_shift + + def add_head_time_series( + self, + name: str, + x: float, + y: float, + layer: int, + t: np.ndarray, + h: np.ndarray, + weights: np.ndarray | None = None, + normalized: bool = False, + constant: float | tuple[float, float, float] | None = None, + time_shift: float | tuple[float, float, float] | None = None, + ) -> None: + """Add a transient head observation time series. + + Parameters + ---------- + name : str + Unique observation name. + x, y : float + Observation coordinates. + layer : int + Aquifer layer index (0-based). + t : array_like + Observation times. + h : array_like + Observed heads. + weights : float, np.ndarray, optional + Per time-series (float) or per-timestep weights (array). Defaults to + uniform weight of 1.0 if ``None``. + normalized : bool + Indicates whether head observations were normalized relative to some + reference level, e.g. heads fluctuate around 0, or whether heads were + provided in absolute values. + constant : float or (float, float, float), optional + Add a calibrated constant offset to this series. Supply a float + for the initial value (unbounded), or a ``(initial, pmin, pmax)`` + tuple to set bounds. ``None`` (default) disables this parameter. + time_shift : float or (float, float, float), optional + Add a calibrated time shift to this series. Supply a float for + the initial value (unbounded), or a ``(initial, pmin, pmax)`` + tuple to set bounds. ``None`` (default) disables this parameter. + + Examples + -------- + Basic usage:: + + cal.add_head_time_series("obs1", x=50.0, y=0.0, layer=0, t=t, h=h) + + With a bounded constant offset:: + + cal.add_head_time_series( + "obs1", x=50.0, y=0.0, layer=0, t=t, h=h, constant=(0.1, -1.0, 1.0) + ) + """ + if self.transient_model is None: + raise ValueError("Transient model must be provided to add head time series.") + obs = HeadSeries( + x=x, + y=y, + layer=layer, + t=t, + h=h, + weights=weights, + normalized=normalized, + constant=constant, + time_shift=time_shift, + ) + self.observations_dict[name] = obs + if constant is not None: + initial, pmin, pmax = self._parse_optional_param( + constant, default_initial=np.mean(h) + ) + self._add_series_constant(name, obs, initial=initial, pmin=pmin, pmax=pmax) + if time_shift is not None: + initial, pmin, pmax = self._parse_optional_param( + time_shift, default_initial=1 / 24.0 + ) + self._add_series_time_shift(name, obs, initial=initial, pmin=pmin, pmax=pmax) + + def add_head_time_series_in_well( + self, + name: str, + well_element: TransientElement, + t: np.ndarray, + h: np.ndarray, + constant: float | tuple[float, float, float] | None = None, + time_shift: float | tuple[float, float, float] | None = None, + ) -> None: + """Add a transient head observation time series inside a well. + + Parameters + ---------- + name : str + Unique observation name. + well_element : TransientElement + Well element in which heads are observed. + t : array_like + Observation times. + h : array_like + Observed heads. + constant : float or (float, float, float), optional + Add a calibrated constant offset to this series. Supply a float + for the initial value (unbounded), or a ``(initial, pmin, pmax)`` + tuple to set bounds. ``None`` (default) disables this parameter. + time_shift : float or (float, float, float), optional + Add a calibrated time shift to this series. Supply a float for + the initial value (unbounded), or a ``(initial, pmin, pmax)`` + tuple to set bounds. ``None`` (default) disables this parameter. + + Examples + -------- + Add observed head time series inside a well:: + + pump_well = timflow.transient.Well(...) + cal.add_head_time_series_in_well("well1", well_element=pump_well, t=t, h=h) + """ + if self.transient_model is None: + raise ValueError("Transient model must be provided to add head time series.") + obs = HeadSeriesInWell( + element=well_element, + t=t, + h=h, + constant=constant, + time_shift=time_shift, + ) + self.observations_in_well_dict[name] = obs + if constant is not None: + initial, pmin, pmax = self._parse_optional_param( + constant, default_initial=np.mean(h) + ) + self._add_series_constant(name, obs, initial=initial, pmin=pmin, pmax=pmax) + if time_shift is not None: + initial, pmin, pmax = self._parse_optional_param( + time_shift, default_initial=1 / 24.0 + ) + self._add_series_time_shift(name, obs, initial=initial, pmin=pmin, pmax=pmax) + + def residuals(self, p: np.ndarray, printdot: bool = False) -> np.ndarray: + """Compute residuals for parameter vector ``p``. + + Parameters + ---------- + p : np.ndarray + Current parameter values in the same order as + ``self.parameters``. + printdot : bool + Print a dot to stdout on each call, useful for tracking progress. + + Returns + ------- + np.ndarray + 1-D array of weighted residuals (observed minus simulated). + """ + if printdot: + print(".", end="", flush=True) + + # 1. Push parameter values into model arrays + for value, param in zip(p, self._parameters.values(), strict=True): + param.set(value) + + # 2. Solve whichever models are registered + needs_steady = any( + s.model_key == "steady" for s in self.observations_dict.values() + ) or any(s.model_key == "steady" for s in self.observations_in_well_dict.values()) + needs_transient = any( + s.model_key == "transient" for s in self.observations_dict.values() + ) or any( + s.model_key == "transient" for s in self.observations_in_well_dict.values() + ) + + if needs_steady and self.steady_model is not None: + self.steady_model.solve(silent=True) + if needs_transient and self.transient_model is not None: + # Solve the steady model first so transient_model.head() returns + # up-to-date h_transient + h_steady values. + if self.steady_model is not None: + self.steady_model.solve(silent=True) + self.transient_model.solve(silent=True) + + # 3. Accumulate residuals + rv = np.empty(0) + for obs in self.observations_dict.values(): + if obs.model_key == "transient": + dt = obs._time_shift if obs.time_shift is not None else 0.0 + _tr_has_steady = getattr(self.transient_model, "steady", None) is not None + if obs.normalized or _tr_has_steady: + # Normalized obs need no steady component; and when the + # transient model already embeds a steady model, its head() + # call already returns h_transient + h_steady, so there is + # nothing to add manually. + hsteady = 0.0 + elif self.steady_model is not None: + hsteady = self.steady_model.head(obs.x, obs.y, layers=obs.layer) + else: + warnings.warn( + f"Observation '{obs.x},{obs.y}' is not marked as normalized " + "but no steady model is provided and the transient model has " + "no embedded steady model. Residuals will be computed against " + "transient heads only.", + stacklevel=2, + ) + hsteady = 0.0 + h = ( + self.transient_model.head(obs.x, obs.y, obs.t - dt, layers=obs.layer) + + hsteady + ) + w = obs.weights if obs.weights is not None else np.ones_like(h) + c = obs._constant if obs.constant is not None else 0.0 + if self.reference_time is not None: + # get closest observation to reference time + tref_idx = np.abs(obs.t - self.reference_time).argmin() + closest_ref_time = obs.t[tref_idx] + htref = self.transient_model.head( + obs.x, obs.y, closest_ref_time, layers=obs.layer + ).squeeze() + res = ((obs.h - obs.h[tref_idx]) - (h - htref) - c) * w + else: + res = (obs.h - h - c) * w + rv = np.append(rv, res) + elif obs.model_key == "steady": + h = self.steady_model.head(obs.x, obs.y, layers=obs.layer) + w = obs.weight if obs.weight is not None else 1.0 + rv = np.append(rv, np.atleast_1d((obs.h - h) * w)) + + for obs in self.observations_in_well_dict.values(): + if obs.model_key == "transient": + dt = obs._time_shift if obs.time_shift is not None else 0.0 + t = obs.t - dt + h = obs.element.headinside(t)[0] + w = obs.weights if obs.weights is not None else np.ones_like(h) + c = obs._constant if obs.constant is not None else 0.0 + if self.reference_time is not None: + # get closest observation to reference time + tref_idx = np.abs(obs.t - self.reference_time).argmin() + closest_ref_time = obs.t[tref_idx] + htref = obs.element.headinside(closest_ref_time)[0] + res = ((obs.h - obs.h[tref_idx]) - (h - htref) - c) * w + else: + res = (obs.h - h - c) * w + rv = np.append(rv, res) + # fix for nans, when tmin is larger than timestep after change in bc + # not ideal but better than crashing the optimizer. Warnings are already + # printed by the model when this happens. + nan_mask = np.isnan(rv) + if nan_mask.any(): + rv[nan_mask] = np.interp(t[nan_mask], t[~nan_mask], rv[~nan_mask]) + elif obs.model_key == "steady": + h = obs.element.headinside() + w = obs.weight if obs.weight is not None else 1.0 + rv = np.append(rv, np.atleast_1d((obs.h - h) * w)) + return rv + + def residuals_lmfit(self, lmfitparams: Any, printdot: bool = False) -> np.ndarray: + """Compute residuals from an ``lmfit.Parameters`` object. + + Wrapper around :meth:`residuals` that extracts the parameter vector + from an ``lmfit.Parameters`` object, for use with + :func:`lmfit.minimize`. + + Parameters + ---------- + lmfitparams : lmfit.Parameters + Current parameter values. + printdot : bool + Print a dot on each call. + + Returns + ------- + np.ndarray + 1-D array of weighted residuals. + """ + p = np.array(list(lmfitparams.valuesdict().values())) + return self.residuals(p, printdot=printdot) + + def lmfit( + self, + printdot: bool = True, + report: bool = True, + initial: bool = True, + **kwargs, + ) -> None: + """Run the least-squares fit using ``lmfit``. + + Uses the Levenberg-Marquardt algorithm via :func:`lmfit.minimize`. + Results are stored in ``self.result`` and optimal values are written + back to each registered :class:`Parameter`. + + Parameters + ---------- + printdot : bool + Print a dot per function evaluation to indicate progress. + report : bool + Print a fit summary (message, parameter table, RMSE) on completion. + initial : bool + When ``True`` (default) the optimizer starts from the registered + ``initial`` values. Set to ``False`` to warm-start from the + previously fitted optimal values instead. If no optimal values + are available yet the method silently falls back to the initial + values. + **kwargs + Forwarded to :func:`lmfit.minimize`. + + Examples + -------- + Basic usage:: + + cal.lmfit() + + Warm-start from a previous fit:: + + cal.lmfit() # phase 1 + cal.lmfit(initial=False) # phase 2: start from phase-1 optimal + + See Also + -------- + :func:`Calibrate.fit` for fitting with ``scipy.optimize.least_squares``. + ``lmfit.minimize`` for more optimization options. + """ + import lmfit + + lmfitparams = lmfit.Parameters() + for name, p in self._parameters.items(): + sv = self._start_value(p, initial) + if p.log_scale: + lb, ub = self._log_scale_bounds(p.pmin, p.pmax, np.sign(sv)) + lmfitparams.add( + name, + value=np.log10(np.abs(sv)), + min=lb if np.isfinite(lb) else None, + max=ub if np.isfinite(ub) else None, + ) + else: + lmfitparams.add(name, value=sv, min=p.pmin, max=p.pmax) + fit_kws = {"epsfcn": 1e-4} # this is essential to specify step for the Jacobian + self.result = lmfit.minimize( + self.residuals_lmfit, + lmfitparams, + method="leastsq", + kws={"printdot": printdot}, + **fit_kws, + **kwargs, + ) + self.result.method = "lm" + print("", flush=True) + + if self.result.success: + for (_, popt), param in zip( + self.result.params.items(), self._parameters.values(), strict=True + ): + if param.log_scale: + sv = self._start_value(param, initial) + param.optimal = np.sign(sv) * 10**popt.value + else: + param.optimal = popt.value + + if report: + res = self.residuals_lmfit(self.result.params) + print(self.result.message) + print(self.parameters) + print(f"RMSE: {np.sqrt(np.mean(res**2)):.3e}") + + @staticmethod + def _start_value(p: "Parameter", use_initial: bool) -> float: + """Return the linear-space starting value for parameter ``p``. + + When *use_initial* is ``True`` (the default) the registered + ``initial`` value (or the current model-array value when + ``initial=None``) is used. When ``False`` the previously fitted + optimal is used instead; if no optimal is available yet the method + silently falls back to the initial value. + """ + if not use_initial and not np.isnan(p.optimal): + return p.optimal + return p.effective_initial + + def fit( + self, + method: str = "trf", + diff_step: float = 1e-4, + xtol: float = 1e-8, + report: bool = True, + printdot: bool = True, + initial: bool = True, + **kwargs, + ) -> None: + """Run the least-squares fit using :func:`scipy.optimize.least_squares`. + + Parameters + ---------- + method : {'lm', 'trf', 'dogbox'} + Optimization algorithm. ``'lm'`` (Levenberg-Marquardt) ignores + bounds; use ``'trf'`` or ``'dogbox'`` when ``pmin``/``pmax`` + constraints are set. Default is ``'trf'``. + diff_step : float + Relative step size used to compute the numerical Jacobian. + xtol : float + Convergence tolerance on the change in parameters. + report : bool + Print a fit summary (parameter table, RMSE) on completion. + printdot : bool + Print a dot per function evaluation to indicate progress. + initial : bool + When ``True`` (default) the optimizer starts from the registered + ``initial`` values (or the current model-array values when + ``initial=None`` was passed to :meth:`set_aquifer_parameter`). + Set to ``False`` to warm-start from the previously fitted optimal + values instead — useful for sequential calibration workflows (e.g. + calibrate on steady observations first, then continue with + transient observations). If no optimal values are available yet + the method silently falls back to the initial values. + **kwargs + Forwarded to :func:`scipy.optimize.least_squares`. + + Examples + -------- + Fit with method="trf" which supports bounded parameters:: + + cal.fit() + + Levenberg-Marquardt with looser tolerance (does not support bounds):: + + cal.fit(method="lm", xtol=1e-6) + + Warm-start from a previous fit (e.g. after calibrating steady-state + observations first):: + + cal.fit() # phase 1: steady calibration + cal.add_head_time_series(...) # add transient observations + cal.fit(initial=False) # phase 2: start from steady optimal + + See Also + -------- + :func:` Calibrate.lmfit` for fitting with ``lmfit``. + :func:`scipy.optimize.least_squares` for more optimization options. + """ + p0 = np.array( + [ + self._start_value(p, initial) + if not p.log_scale + else np.log10(np.abs(self._start_value(p, initial))) + for p in self._parameters.values() + ] + ) + lb = np.array( + [ + self._log_scale_bounds( + p.pmin, p.pmax, np.sign(self._start_value(p, initial)) + )[0] + if p.log_scale + else p.pmin + for p in self._parameters.values() + ] + ) + ub = np.array( + [ + self._log_scale_bounds( + p.pmin, p.pmax, np.sign(self._start_value(p, initial)) + )[1] + if p.log_scale + else p.pmax + for p in self._parameters.values() + ] + ) + bounds = (lb, ub) + + self.result = least_squares( + self.residuals, + p0, + args=(printdot,), + bounds=bounds, + method=method, + diff_step=diff_step, + xtol=xtol, + x_scale="jac", + **kwargs, + ) + self.result.method = method + print("", flush=True) + + # Push optimal values into models and record results + res = self.residuals(self.result.x) + for value, param in zip(self.result.x, self._parameters.values(), strict=True): + if param.log_scale: + param.optimal = np.sign(self._start_value(param, initial)) * 10**value + else: + param.optimal = value + + if report: + print(self.parameters) + print(f"RMSE: {np.sqrt(np.mean(res**2)):.3e}") + + def rmse(self) -> float: + """Return the root-mean-square error at the current optimal parameters. + + Returns + ------- + float + RMSE of the weighted residuals. + """ + result = getattr(self, "result", None) + if result is not None and getattr(result, "x", None) is not None: + params_vec = result.x + else: + # Fall back to reconstructing optimization-space parameters + values = [] + for p in self._parameters.values(): + if getattr(p, "log_scale", False): + values.append( + np.log10( + np.abs(p.effective_initial) + if np.isnan(p.optimal) + else np.abs(p.optimal) + ) + ) + else: + values.append( + p.effective_initial if np.isnan(p.optimal) else p.optimal + ) + params_vec = np.array(values, dtype=float) + + r = self.residuals(params_vec) + return float(np.sqrt(np.mean(r**2))) + + @staticmethod + def get_covariances( + jacobian: np.ndarray, + cost: float, + method: Literal["trf", "dogbox", "lm"] = "trf", + absolute_sigma: bool = False, + ) -> np.ndarray: + """ + Method to get the covariance matrix from the jacobian. + + Parameters + ---------- + jacobian : ArrayLike + The jacobian matrix with dimensions nobs, npar. + cost : float + The cost value of the scipy.optimize.OptimizeResult which is half + the sum of squares. That's why the cost is multiplied by a factor + of two internally to get the sum of squares. + method : Literal["trf", "dogbox", "lm"], optional + Algorithm with which the minimization is performed. Default is "trf". + absolute_sigma : bool, optional + If True, `sigma` is used in an absolute sense and the estimated + parameter covariance `pcov` reflects these absolute values. If + False (default), only the relative magnitudes of the `sigma` values + matter. The returned parameter covariance matrix `pcov` is based on + scaling `sigma` by a constant factor. This constant is set by + demanding that the reduced `chisq` for the optimal parameters + `popt` when using the *scaled* `sigma` equals unity. In other + words, `sigma` is scaled to match the sample variance of the + residuals after the fit. Default is False. + Mathematically, ``pcov(absolute_sigma=False) = + pcov(absolute_sigma=True) * chisq(popt)/(M-N)`` + + Returns + ------- + pcov: array_like + numpy array with the covariance matrix. + + Notes + ----- + This method is copied from Scipy: + https://github.com/scipy/scipy/blob/92d2a8592782ee19a1161d0bf3fc2241ba78bb63/scipy/optimize/_minpack_py.py + Please refer to the SciPy optimization module:: + https://docs.scipy.org/doc/scipy/reference/optimize.html + """ + nobs, npar = jacobian.shape + cost = 2 * cost # res.cost is half sum of squares! + s_sq = cost / (nobs - npar) # variance of the residuals + + if method == "lm": + # https://github.com/scipy/scipy/blob/92d2a8592782ee19a1161d0bf3fc2241ba78bb63/scipy/optimize/_minpack_py.py#L480C9-L499C38 + # fjac A permutation of the R matrix of a QR factorization of the + # final approximate Jacobian matrix. + _, fjac = np.linalg.qr(jacobian) + # leastsq expects the fjacobian to be in fortran order (npar, nobs) + # that why it is transposed in the original code + + ipvt = np.arange(1, npar + 1, dtype=int) + n = len(ipvt) + r = np.triu(fjac[:n, :]) + + # old method deprecated in scipy 1.10.0 since + # the explicit dot product was not necessary and sometimes + # the result was not symmetric positive definite. + # See https://github.com/scipy/scipy/issues/4555. + # old method + # perm = np.take(np.eye(n), ipvt - 1, 0) + # R = np.dot(r, perm) + # cov_x = np.linalg.inv(np.dot(np.transpose(R), R)) + + # new method: + perm = ipvt - 1 + inv_triu = get_lapack_funcs("trtri", (r,)) + try: + # inverse of permuted matrix is a permutation of matrix inverse + invR, trtri_info = inv_triu(r) # default: upper, non-unit diag + if trtri_info != 0: # explicit comparison for readability + raise LinAlgError + invR[perm] = invR.copy() + pcov = invR @ invR.T # cov_x in the original code + except (LinAlgError, ValueError): + pcov = None + else: + # https://github.com/scipy/scipy/blob/92d2a8592782ee19a1161d0bf3fc2241ba78bb63/scipy/optimize/_minpack_py.py#L1029-L1055 + # Do Moore-Penrose inverse discarding zero singular values. + _, s, VT = svd(jacobian, full_matrices=False) + threshold = np.finfo(float).eps * max(jacobian.shape) * s[0] + s = s[s > threshold] + VT = VT[: s.size] + pcov = np.dot(VT.T / s**2, VT) + + if pcov is None or np.isnan(pcov).any(): + # indeterminate covariance + pcov = np.full(shape=(npar, npar), fill_value=np.inf, dtype=float) + elif not absolute_sigma: + if nobs > npar: + pcov = pcov * s_sq + else: + pcov = np.full(shape=(npar, npar), fill_value=np.inf, dtype=float) + + return pcov + + def get_correlations(self) -> np.ndarray: + """Return the parameter correlation matrix. + + Uses the covariance matrix from the fit result. For lmfit results the + covariance is read directly from ``self.result.covar``; for scipy + results it is computed from the Jacobian via :meth:`get_covariances`. + + Returns + ------- + np.ndarray + Correlation matrix of shape ``(n_params, n_params)``. + """ + if hasattr(self.result, "covar") and self.result.covar is not None: + pcov = self.result.covar + else: + pcov = self.get_covariances( + jacobian=self.result.jac, + cost=self.result.cost, + method=self.result.method, + absolute_sigma=False, + ) + v = np.sqrt(np.diag(pcov)) + with np.errstate(divide="ignore", invalid="ignore"): + corr = pcov / np.outer(v, v) + corr[pcov == 0] = 0 + return corr + + @property + def parameters(self) -> pd.DataFrame: + """Summary of all registered calibration parameters. + + Returns + ------- + pd.DataFrame + DataFrame indexed by parameter name with columns: ``initial``, + ``optimal``, ``pmin``, ``pmax``, ``log_scaled``, ``n_targets``, + ``n_models``, ``n_inhoms``. + """ + rows = [] + for p in self._parameters.values(): + rows.append( + { + "name": p.name, + "initial": p.initial, + "optimal": p.optimal, + "pmin": p.pmin, + "pmax": p.pmax, + "log_scaled": p.log_scale, + "n_targets": len(p.targets), + "n_models": len(set(p.models)), + "n_inhoms": len(set(p.inhoms)), + } + ) + return pd.DataFrame(rows).set_index("name") + + def _parse_layers(self, name: str, layers: int | Iterable[int]) -> tuple[int, int]: + """Parse layer specification and return (from_layer, to_layer).""" + if isinstance(layers, Iterable) and not isinstance(layers, str): + layers = list(layers) + from_lay, to_lay = min(layers), max(layers) + if (np.diff(layers) > 1).any(): + warnings.warn( + f"Non-consecutive layers; setting '{name}' for " + f"layers {from_lay}–{to_lay}.", + stacklevel=3, + ) + elif isinstance(layers, int): + from_lay = to_lay = layers + else: + raise TypeError(f"layers must be int or iterable of int, got {type(layers)}") + return from_lay, to_lay + + def _resolve_models(self, model_key: str) -> list: + """Return list of model objects corresponding to model_key.""" + mapping = { + "transient": [self.transient_model], + "steady": [self.steady_model], + "both": [self.transient_model, self.steady_model], + } + return [ml for ml in mapping[model_key] if ml is not None] + + def _resolve_inhoms(self, model: Any, inhoms: str | list[str] | None) -> list: + """Return list of aquifer objects for given inhoms spec.""" + if inhoms is None: + return [model.aq] + if isinstance(inhoms, str): + inhoms = [inhoms] + elif not isinstance(inhoms, list): + inhoms = list(inhoms) + return [model.aq.inhomdict[i] if isinstance(i, str) else i for i in inhoms] + + @staticmethod + def _log_scale_bounds(pmin: float, pmax: float, sign: float) -> tuple[float, float]: + """Convert linear-space bounds to log10(abs) optimizer space. + + For a positive parameter (sign >= 0): + pmin > 0 → lb = log10(pmin), ub = log10(pmax) + For a negative parameter (sign < 0): + pmin ≤ pmax ≤ 0, so abs(pmax) ≤ abs(value) ≤ abs(pmin) + lb = log10(abs(pmax)), ub = log10(abs(pmin)) + Infinite or incompatible bounds are passed through as ±inf. + """ + if sign >= 0: + lb = np.log10(pmin) if (np.isfinite(pmin) and pmin > 0) else -np.inf + ub = np.log10(pmax) if (np.isfinite(pmax) and pmax > 0) else np.inf + else: + lb = np.log10(-pmax) if (np.isfinite(pmax) and pmax < 0) else -np.inf + ub = np.log10(-pmin) if (np.isfinite(pmin) and pmin < 0) else np.inf + return lb, ub + + @staticmethod + def _get_aquifer_parameter_array(model, aq, name: str) -> np.ndarray: + """Return reference to the appropriate parameter array in the aquifer object.""" + lookup = {"kaq": aq.kaq, "c": aq.c} + if hasattr(aq, "Saq"): + lookup["Saq"] = aq.Saq + if hasattr(aq, "Sll"): + lookup["Sll"] = aq.Sll + if hasattr(aq, "kzoverkh"): + lookup["kzoverkh"] = aq.kzoverkh + assert aq.model == model, "Aquifer does not belong to the given model" + for prefix, arr in lookup.items(): + if name.startswith(prefix): + return arr + raise ValueError( + f"Parameter '{name}' not recognised. Supported: {list(lookup.keys())}" + ) + + @staticmethod + def _make_pname( + name: str, + from_lay: int, + to_lay: int, + inhoms: str | list[str] | None, + models: list, + ) -> str: + """Construct a unique parameter name based on the specification.""" + base = f"{name}_{from_lay}_{to_lay}" + if inhoms is not None: + inhom_names = ( + [inhoms] + if isinstance(inhoms, str) + else [i if isinstance(i, str) else i.name for i in inhoms] + ) + base += "_" + "_".join(inhom_names) + return base + + @staticmethod + def _nse(h_obs: np.ndarray, h_mod: np.ndarray) -> float: + """Compute the Nash-Sutcliffe Efficiency (NSE). + + NSE = 1 - Σ(h_obs - h_mod)² / Σ(h_obs - mean(h_obs))² + + Parameters + ---------- + h_obs : np.ndarray + Observed heads. + h_mod : np.ndarray + Modeled heads. + + Returns + ------- + nse : float + NSE value, where 1 is a perfect fit, 0 means the model is as good as using + the mean of the observations, and negative values indicate a worse fit. + Returns NaN when all observations are identical (zero variance). + """ + denom = float(np.sum((h_obs - np.mean(h_obs)) ** 2)) + if denom == 0.0: + return np.nan + return float(1.0 - np.sum((h_obs - h_mod) ** 2) / denom) + + def plot_transient_results( + self, + tmin: float | None = None, + tmax: float | None = None, + figsize: tuple[float, float] = (10, 8), + obs_kwargs: dict | None = None, + model_kwargs: dict | None = None, + sharey: bool = False, + ) -> tuple[plt.Figure, np.ndarray]: + """Plot modeled vs observed transient head time series. + + Creates one subplot per transient observation with a shared x-axis, + comparing observed heads to the current model response. Call this + method before calibration to inspect the initial fit, or after + calibration to inspect the calibrated fit. + + Parameters + ---------- + tmin : float, optional + Start of the plotted time window. Defaults to the earliest + observation time across all series. + tmax : float, optional + End of the plotted time window. Defaults to the latest + observation time across all series. + figsize : tuple of float, optional + Figure size ``(width, height)`` in inches. Default is ``(10, 8)``. + obs_kwargs : dict, optional + Keyword arguments passed to :func:`matplotlib.axes.Axes.plot` for + the observed data. Default: black dots. + model_kwargs : dict, optional + Keyword arguments passed to :func:`matplotlib.axes.Axes.plot` for + the modeled response. Default: blue solid line. + sharey : bool, optional + If ``True``, all subplots share the same y-axis limits. + Default is ``False``. + + Returns + ------- + fig : matplotlib.figure.Figure + axes : np.ndarray of matplotlib.axes.Axes + One Axes per transient observation. + """ + # Collect transient observations in insertion order + transient_items = [ + (name, obs) + for name, obs in self.observations_dict.items() + if obs.model_key == "transient" + ] + transient_well_items = [ + (name, obs) + for name, obs in self.observations_in_well_dict.items() + if obs.model_key == "transient" + ] + all_items = transient_items + transient_well_items + n_obs = len(all_items) + + if n_obs == 0: + raise ValueError("No transient observations to plot.") + + # Default styling + obs_kw: dict = {"color": "k", "marker": ".", "linestyle": "none"} + obs_kw.update(obs_kwargs or {}) + + model_kw: dict = {"color": "tab:blue", "label": "model"} + model_kw.update(model_kwargs or {}) + + _tr_has_steady = getattr(self.transient_model, "steady", None) is not None + + # Solve the current model state, just in case, but should already be done usually. + if self.steady_model is not None: + self.steady_model.solve(silent=True) + if self.transient_model is not None: + self.transient_model.solve(silent=True) + + # Create subplots + fig, ax_array = plt.subplots( + n_obs, 1, sharex=True, sharey=sharey, figsize=figsize + ) + axes = np.atleast_1d(ax_array) + + for ax, (name, obs) in zip(axes, all_items, strict=True): + dt = obs._time_shift if obs.time_shift is not None else 0.0 + t_plot = obs.t - dt # shared time axis for observed and modeled + + # Build observed label, annotating any corrections + obs_label_parts = [] + if obs.time_shift is not None: + obs_label_parts.append(f"\u0394t={float(obs._time_shift[0]):.2f}") + if obs.constant is not None: + obs_label_parts.append(f"\u0394h={float(obs._constant[0]):.2f}") + obs_suffix = f" ({', '.join(obs_label_parts)})" if obs_label_parts else "" + obs_label = f"{name}{obs_suffix}" + + # Compute observed heads with corrections applied + if self.reference_time is not None: + tref_idx = np.abs(obs.t - self.reference_time).argmin() + h_obs_plot = obs.h - obs.h[tref_idx] + else: + c = float(obs._constant[0]) if obs.constant is not None else 0.0 + h_obs_plot = obs.h - c + + # Compute modeled heads + if name in dict(transient_items): + if obs.normalized or _tr_has_steady: + hsteady = 0.0 + elif self.steady_model is not None: + hsteady = float( + np.squeeze(self.steady_model.head(obs.x, obs.y, layers=obs.layer)) + ) + else: + hsteady = 0.0 + h_mod = ( + self.transient_model.head(obs.x, obs.y, obs.t - dt, layers=obs.layer) + + hsteady + ).squeeze() + else: + h_mod = obs.element.headinside(obs.t - dt)[0] + + if self.reference_time is not None: + h_mod = h_mod - h_mod[tref_idx] + + # Apply tmin/tmax window + mask = np.ones(len(t_plot), dtype=bool) + if tmin is not None: + mask &= t_plot >= tmin + if tmax is not None: + mask &= t_plot <= tmax + + nse = self._nse(h_obs_plot[mask], h_mod[mask]) + nse_str = f"NSE={nse:.2f}" if np.isfinite(nse) else "NSE=n/a" + model_label = f"{model_kw.get('label', 'model')} ({nse_str})" + + ax.plot(t_plot[mask], h_obs_plot[mask], label=obs_label, **obs_kw) + ax.plot(t_plot[mask], h_mod[mask], **{**model_kw, "label": model_label}) + ax.set_ylabel("head") + ax.grid(True) + ax.legend(loc=(0, 1), frameon=False, ncol=2) + ax.set_xlim(left=t_plot[mask][0], right=t_plot[mask][-1]) + + axes[-1].set_xlabel("time") + fig.tight_layout() + return fig, axes diff --git a/timflow/plots/plots.py b/timflow/plots/plots.py index 7aa30689..184d1c89 100644 --- a/timflow/plots/plots.py +++ b/timflow/plots/plots.py @@ -88,6 +88,7 @@ def xsection( fmt=None, units=None, hstar=None, + boundaries=True, horizontal_axis: Literal["x", "y", "s"] = "s", sep: Literal[", ", "\n"] = ", ", **kwargs, @@ -124,6 +125,9 @@ def xsection( override hstar value for plotting water level in transient 1D inhomogeneities that use hstar, useful for plotting pretty cross-sections when reference level is not equal to 0. + boundaries : bool, optional + whether to plot aquifer boundaries for cross-section models, + default is True sep : str Separator between parameters, either ", " or "\n" **kwargs @@ -151,6 +155,7 @@ def xsection( fmt=fmt, units=units, hstar=hstar, + boundaries=boundaries, sep=sep, ) @@ -270,6 +275,7 @@ def _xsection_simple_aquifer( fmt, units=None, hstar=None, + boundaries=True, sep: Literal[", ", "\n"] = ", ", ): """Handle cross-section plotting for SimpleAquifer models.""" @@ -324,7 +330,8 @@ def _xsection_simple_aquifer( if isinstance(e, HstarXsection): e.plot(ax=ax, hstar=hstar) else: - e.plot(ax=ax) + if not e.inhomelement or boundaries: + e.plot(ax=ax) return ax @@ -340,7 +347,7 @@ def _xsection_plot_inhoms( units, sep: Literal[", ", "\n"] = ", ", ): - """Plot inhomogeneities for SimpleAquifer models. + r"""Plot inhomogeneities for SimpleAquifer models. Parameters ---------- @@ -359,33 +366,21 @@ def _xsection_plot_inhoms( units : dict or None Dictionary of units for parameters (unused in transient, kept for compatibility) + sep : str, optional + Separator between parameters, either ", " or "\n" """ - if self._ml.model_type == "steady": - for inhom in self._ml.aq.inhomlist: - inhom.plot( - ax=ax, - labels=labels, - params=params, - names=names, - x1=x1, - x2=x2, - fmt=fmt, - units=units, - sep=sep, - ) - elif self._ml.model_type == "transient": - for inhom in self._ml.aq.inhomdict.values(): - inhom.plot( - ax=ax, - labels=labels, - params=params, - names=names, - x1=x1, - x2=x2, - fmt=fmt, - units=units, - sep=sep, - ) + for inhom in self._ml.aq.inhomdict.values(): + inhom.plot( + ax=ax, + labels=labels, + params=params, + names=names, + x1=x1, + x2=x2, + fmt=fmt, + units=units, + sep=sep, + ) def _get_xsection_line_params(self, xy, ax, horizontal_axis): """Get parameters for cross-section line. diff --git a/timflow/steady/aquifer.py b/timflow/steady/aquifer.py index 792b60eb..43207705 100644 --- a/timflow/steady/aquifer.py +++ b/timflow/steady/aquifer.py @@ -89,6 +89,9 @@ def __init__(self, model, kaq, c, z, npor, ltype, model3d=False): def initialize(self): self.elementlist = [] # Elementlist of aquifer + # Recompute T for when kaq is changed + self.T = self.kaq * self.Haq + self.Tcol = self.T.reshape(self.naq, 1) d0 = 1.0 / (self.c * self.T) d0[:-1] += 1.0 / (self.c[1:] * self.T[:-1]) dp1 = -1.0 / (self.c[1:] * self.T[1:]) @@ -156,7 +159,7 @@ def summary(self): add_cols = [] summary = pd.DataFrame( index=range(self.nlayers), - columns=["layer", "layer_type", "H", "k_h", "c"] + add_cols, + columns=["layer", "layer_type", "H", "k_h"] + add_cols + ["c"], ) summary.index.name = "#" layertype = {"a": "aquifer", "l": "leaky layer"} @@ -170,8 +173,8 @@ def summary(self): summary.loc[0, "c"] = np.nan # reset confined resistance to nan summary.loc[maskaq, "kzoverkh"] = self.kzoverkh else: - summary.iloc[~maskaq, 2] = self.Hll[1:] - summary.iloc[~maskaq, 4] = self.c[1:] + summary.loc[~maskaq, "H"] = self.Hll[1:] + summary.loc[~maskaq, "c"] = self.c[1:] else: summary.loc[maskaq, "H"] = self.Haq summary.loc[maskaq, "k_h"] = self.kaq @@ -180,8 +183,8 @@ def summary(self): summary.loc[maskaq, "c"] = self.c summary.loc[maskaq, "kzoverkh"] = self.kzoverkh else: - summary.iloc[~maskaq, 2] = self.Hll - summary.iloc[~maskaq, 4] = self.c + summary.loc[~maskaq, "H"] = self.Hll + summary.loc[~maskaq, "c"] = self.c summary.loc[:, "layer"] = self.layernumber return summary # .set_index("layer") @@ -193,25 +196,30 @@ class Aquifer(AquiferData): """ def __init__(self, model, kaq, c, z, npor, ltype, model3d=False): - super().__init__(model, kaq, c, z, npor, ltype, model3d) - self.inhomlist = [] + super().__init__(model, kaq, c, z, npor, ltype, model3d=model3d) + self.inhomdict = {} self.area = 1e300 # Needed to find smallest inhom def initialize(self): - # cause we are going to call initialize for inhoms - AquiferData.initialize(self) - for inhom in self.inhomlist: + super().initialize() + # 2 passes to ensure all data is present prior to creating elements + for inhom in self.inhomdict.values(): inhom.initialize() - for inhom in self.inhomlist: + for inhom in self.inhomdict.values(): inhom.create_elements() def add_inhom(self, inhom): - self.inhomlist.append(inhom) - return len(self.inhomlist) - 1 # returns number in the list + inhom_number = len(self.inhomdict) + if not hasattr(inhom, "name") or inhom.name is None: + inhom.name = f"inhom{inhom_number:02g}" + if inhom.name in self.inhomdict: + raise ValueError(f"Inhomogeneity name '{inhom.name}' already exists.") + self.inhomdict[inhom.name] = inhom + return inhom_number def find_aquifer_data(self, x, y): rv = self - for inhom in self.inhomlist: + for inhom in self.inhomdict.values(): if inhom.isinside(x, y): if inhom.area < rv.area: rv = inhom @@ -231,7 +239,7 @@ class SimpleAquifer(Aquifer): def __init__(self, naq): self.naq = naq - self.inhomlist = [] + self.inhomdict = {} self.area = 1e300 # Needed to find smallest inhomogeneity self.elementlist = [] @@ -239,7 +247,8 @@ def __repr__(self): return f"Simple Aquifer: {self.naq} aquifer(s)" def initialize(self): - for inhom in self.inhomlist: + # 2 passes to ensure all data is present prior to creating elements + for inhom in self.inhomdict.values(): inhom.initialize() - for inhom in self.inhomlist: + for inhom in self.inhomdict.values(): inhom.create_elements() diff --git a/timflow/steady/inhomogeneity.py b/timflow/steady/inhomogeneity.py index a9021cc3..b5caa554 100644 --- a/timflow/steady/inhomogeneity.py +++ b/timflow/steady/inhomogeneity.py @@ -42,7 +42,9 @@ class PolygonInhom(AquiferData): tiny = 1e-8 - def __init__(self, model, xy, kaq, c, z, npor, ltype, hstar, N, order, ndeg): + def __init__( + self, model, xy, kaq, c, z, npor, ltype, hstar, N, order, ndeg, name=None + ): # All input variables except model should be numpy arrays # That should be checked outside this function): AquiferData.__init__(self, model, kaq, c, z, npor, ltype) @@ -50,6 +52,7 @@ def __init__(self, model, xy, kaq, c, z, npor, ltype, hstar, N, order, ndeg): self.ndeg = ndeg self.hstar = hstar self.N = N + self.name = name self.inhom_number = self.model.aq.add_inhom(self) self.z1, self.z2 = compute_z1z2(xy) self.Nsides = len(self.z1) @@ -186,6 +189,8 @@ class PolygonInhomMaq(PolygonInhom): ndeg : int number of points used between two segments to numerically integrate normal discharge + name : string, optional + name of inhomogeneity """ tiny = 1e-8 @@ -203,6 +208,7 @@ def __init__( N=None, order=3, ndeg=3, + name=None, ): if c is None: c = [] @@ -220,7 +226,19 @@ def __init__( ltype, ) = param_maq(kaq, z, c, npor, topboundary) PolygonInhom.__init__( - self, model, xy, kaq, c, z, npor, ltype, hstar, N, order, ndeg + self, + model, + xy, + kaq, + c, + z, + npor, + ltype, + hstar, + N, + order, + ndeg, + name=name, ) diff --git a/timflow/steady/inhomogeneity1d.py b/timflow/steady/inhomogeneity1d.py index 71c3673d..68408c36 100644 --- a/timflow/steady/inhomogeneity1d.py +++ b/timflow/steady/inhomogeneity1d.py @@ -66,11 +66,8 @@ def __init__( self.x2 = x2 self.hstar = hstar self.N = N + self.name = name self.inhom_number = self.model.aq.add_inhom(self) - if name is None: - self.name = f"inhom{self.inhom_number:02g}" - else: - self.name = name self.addlinesinks = True # Set to False not to add line-sinks def __repr__(self): diff --git a/timflow/steady/linesink1d.py b/timflow/steady/linesink1d.py index 52be6d62..f5c5ec70 100644 --- a/timflow/steady/linesink1d.py +++ b/timflow/steady/linesink1d.py @@ -70,12 +70,12 @@ def initialize(self): self.cosnorm = np.cos(self.theta_norm_out) * np.ones(self.ncp) self.sinnorm = np.sin(self.theta_norm_out) * np.ones(self.ncp) if self.wh == "H": - self.wh = self.aq.Haq[self.layers] + self._wh = self.aq.Haq[self.layers] elif self.wh == "2H": - self.wh = 2.0 * self.aq.Haq[self.layers] + self._wh = 2.0 * self.aq.Haq[self.layers] elif np.isscalar(self.wh): - self.wh = self.wh * np.ones(self.nlayers) - self.resfac = self.aq.Haq[self.layers] * self.res / self.wh + self._wh = self.wh * np.ones(self.nlayers) + self.resfac = self.aq.Haq[self.layers] * self.res / self._wh def potinf(self, x, y, aq=None): if aq is None: diff --git a/timflow/steady/model.py b/timflow/steady/model.py index 8d78a0d9..ba18d686 100644 --- a/timflow/steady/model.py +++ b/timflow/steady/model.py @@ -59,7 +59,7 @@ def __init__(self, kaq, c, z, npor, ltype, model3d=False): # That should be checked outside this function self.elementlist = [] self.elementdict = {} # only elements that have a label - self.aq = Aquifer(self, kaq, c, z, npor, ltype, model3d) + self.aq = Aquifer(self, kaq, c, z, npor, ltype, model3d=model3d) self.modelname = "ml" # Used for writing out input self.name = "Model" self.model_type = "steady" # Model type for plotting and other purposes @@ -826,11 +826,16 @@ def aquifer_summary(self): pandas.DataFrame dataframe with summary of aquifer(s) parameters """ + if type(self) is Model: + raise NotImplementedError( + "aquifer_summary is not supported for the base Model class; " + "use ModelMaq, Model3D instead." + ) aqs = {} if not isinstance(self.aq, SimpleAquifer): aqs["background"] = self.aq.summary() - for i, iaq in enumerate(self.aq.inhomlist): - aqs[f"inhom{i}"] = iaq.summary() + for iaq in self.aq.inhomdict.values(): + aqs[iaq.name] = iaq.summary() if aqs: return pd.concat(aqs, axis=0) @@ -1050,14 +1055,14 @@ def check_inhoms(self): """ # check aquifers naqs = {} - for inhom in self.aq.inhomlist: + for inhom in self.aq.inhomdict.values(): naqs[inhom.name] = inhom.naq check = np.array(list(naqs.values())) == self.aq.naq if not check.all(): raise ValueError(f"Number of aquifers does not match {self.aq.naq}:\n{naqs}") # check -inf to inf - x1list = np.array([inhom.x1 for inhom in self.aq.inhomlist]) - x2list = np.array([inhom.x2 for inhom in self.aq.inhomlist]) + x1list = np.array([inhom.x1 for inhom in self.aq.inhomdict.values()]) + x2list = np.array([inhom.x2 for inhom in self.aq.inhomdict.values()]) xmin = x1list.min() xmax = x2list.max() if not (np.isinf(xmin) and np.sign(xmin) < 0): diff --git a/timflow/transient/__init__.py b/timflow/transient/__init__.py index 87c547cb..239a98d6 100644 --- a/timflow/transient/__init__.py +++ b/timflow/transient/__init__.py @@ -17,6 +17,7 @@ LeakyWallString, ) from timflow.transient.linedoublet1d import ( + ImpermeableWall1D, LeakyLineDoublet1D, # deprecated LeakyWall1D, ) diff --git a/timflow/transient/aquifer.py b/timflow/transient/aquifer.py index 6b036750..00db6abe 100644 --- a/timflow/transient/aquifer.py +++ b/timflow/transient/aquifer.py @@ -237,24 +237,44 @@ def findlayer(self, z): return layernumber, ltype, modellayer def summary(self): + if self.nlayers == self.naq: + model3d = True + add_cols = ["kzoverkh"] + else: + model3d = False + add_cols = [] summary = pd.DataFrame( index=range(self.nlayers), - columns=["layer", "layer_type", "k_h", "c", "S_s"], + columns=["layer", "layer_type", "H", "k_h"] + add_cols + ["c", "S_s"], ) summary.index.name = "#" layertype = {"a": "aquifer", "l": "leaky layer"} summary["layer_type"] = [layertype[lt] for lt in self.ltype] + maskaq = self.ltype == "a" if self.topboundary.startswith("con"): - summary.iloc[::2, 2] = self.kaq - summary.iloc[::2, 4] = self.Saq - summary.iloc[1::2, 3] = self.c - summary.iloc[1::2, 4] = self.Sll + summary.loc[maskaq, "H"] = self.Haq + summary.loc[maskaq, "k_h"] = self.kaq + summary.loc[maskaq, "S_s"] = self.Saq + if model3d: + summary.loc[maskaq, "c"] = self.c + summary.loc[0, "c"] = np.nan # reset confined resistance to nan + summary.loc[maskaq, "kzoverkh"] = self.kzoverkh + else: + summary.loc[~maskaq, "c"] = self.c + summary.loc[~maskaq, "S_s"] = self.Sll else: - summary.iloc[1::2, 2] = self.kaq - summary.iloc[1::2, 4] = self.Saq - summary.iloc[::2, 4] = self.Sll - summary.iloc[::2, 3] = self.c + summary.loc[maskaq, "H"] = self.Haq + summary.loc[maskaq, "k_h"] = self.kaq + summary.loc[maskaq, "S_s"] = self.Saq + if model3d: + summary.loc[~maskaq, "H"] = self.Hll[0] + summary.loc[maskaq, "c"] = self.c + summary.loc[maskaq, "kzoverkh"] = self.kzoverkh + else: + summary.loc[~maskaq, "H"] = self.Hll + summary.loc[~maskaq, "c"] = self.c + summary.loc[~maskaq, "S_s"] = self.Sll summary.loc[:, "layer"] = self.layernumber return summary # .set_index("layer") @@ -313,8 +333,11 @@ def __repr__(self): def initialize(self): super().initialize() + # 2 passes to ensure all data is present prior to creating elements for inhom in self.inhomdict.values(): inhom.initialize() + for inhom in self.inhomdict.values(): + inhom.create_elements() def is_inside(self, x, y): return True @@ -347,5 +370,8 @@ def __repr__(self): return f"Simple Aquifer: {self.naq} aquifer(s)" def initialize(self): + # 2 passes to ensure all data is present prior to creating elements for inhom in self.inhomdict.values(): inhom.initialize() + for inhom in self.inhomdict.values(): + inhom.create_elements() diff --git a/timflow/transient/inhom1d.py b/timflow/transient/inhom1d.py index 0c3cb4e7..57e91b45 100644 --- a/timflow/transient/inhom1d.py +++ b/timflow/transient/inhom1d.py @@ -156,10 +156,6 @@ def is_inside(self, x, _): """ return (x >= self.x1) and (x < self.x2) - def initialize(self): - super().initialize() - self.create_elements() - def create_elements(self): """Create linesinks to meet the continuity conditions the at the boundaries.""" if np.isfinite(self.x1) and np.isfinite(self.x2): @@ -200,8 +196,8 @@ def create_elements(self): self.model, self.x1, range(self.naq), aq=aqin, label=None ) if self.tsandN is not None: - assert self.topboundary[:3] == "con" or self.topboundary[:3] == "phr", ( - Exception("Infiltration can only be applied to a confined aquifer.") + assert self.topboundary == "con" or self.topboundary == "phr", Exception( + "Infiltration can only be applied to a confined aquifer." ) AreaSinkXsection(self.model, self.x1, self.x2, tsandN=self.tsandN) if self.tsandhstar is not None: @@ -269,7 +265,7 @@ def plot( r0 = x1 if labels or params: - lli = 1 if self.topboundary == "con" else 0 + lli = 1 if self.topboundary in ["con", "phr"] else 0 aqi = 0 else: lli = None diff --git a/timflow/transient/invlapnumba.py b/timflow/transient/invlapnumba.py index 2b8dd0cd..55e69700 100644 --- a/timflow/transient/invlapnumba.py +++ b/timflow/transient/invlapnumba.py @@ -214,8 +214,13 @@ def invlapcomp(time, pot, npint, M, tintervals, enumber, etstart, ebc, nlayers): it = np.argmax(t > 0) # find_first if t[it] < tintervals[0]: # there are times before first interval if print_tmin_warning: - print("Warning, some of the times are smaller than tmin after") - print("a change in boundary condition. nans are substituted") + print( + "Warning, some of the times (", + t[it], + ") are smaller than tmin =", + tintervals[0], + "after a change in boundary condition. Nans are substituted.", + ) print_tmin_warning = False if t[-1] < tintervals[0]: # all times before first interval itnew = len(t) diff --git a/timflow/transient/linedoublet1d.py b/timflow/transient/linedoublet1d.py index 14c81ae8..fded8af5 100644 --- a/timflow/transient/linedoublet1d.py +++ b/timflow/transient/linedoublet1d.py @@ -158,7 +158,7 @@ def plot(self, ax=None): class LeakyWall1D(LineDoublet1DBase, LeakyWallEquation): - r"""Leaky line doublet with specified resistance. + r"""Leaky wall with specified resistance. Parameters ---------- @@ -166,8 +166,9 @@ class LeakyWall1D(LineDoublet1DBase, LeakyWallEquation): model to which the element is added xld : float x-coordinate of the line doublet - res : float - resistance of the line doublet + res : float or str + resistance of the line doublet, default is "imp" for impermeable, + use float to specify a finite resistance. layers : int, array or list layer (int) or layers (list or array) in which line doublet is located label : string or None (default: None) @@ -194,6 +195,31 @@ def initialize(self): ) +class ImpermeableWall1D(LeakyWall1D): + """Impermeable wall. + + Parameters + ---------- + model : Model object + model to which the element is added + xld : float + x-coordinate of the line doublet + layers : int, array or list + layer (int) or layers (list or array) in which line doublet is located + label : string or None (default: None) + label of the element + """ + + def __init__(self, model, xld=0, layers=0, label=None): + super().__init__( + model, + xld, + res="imp", + layers=layers, + label=label, + ) + + class LeakyLineDoublet1D(LeakyWall1D): """Deprecated alias for :class:`.LeakyWall1D`. diff --git a/timflow/transient/linesink1d.py b/timflow/transient/linesink1d.py index 2b7abfa7..a112ee1c 100644 --- a/timflow/transient/linesink1d.py +++ b/timflow/transient/linesink1d.py @@ -317,6 +317,11 @@ def __init__( layers=0, label=None, ): + if isinstance(tsandh, str) and tsandh == "fixed": + tsandh = [(0, 0)] + etype = "z" + else: + etype = "v" super().__init__( model, xls, @@ -324,7 +329,7 @@ def __init__( res=res, wh=wh, layers=layers, - type="v", + type=etype, name="River1D", label=label, ) diff --git a/timflow/transient/model.py b/timflow/transient/model.py index 6b87bbdd..9829ce11 100644 --- a/timflow/transient/model.py +++ b/timflow/transient/model.py @@ -27,7 +27,7 @@ from timflow.version import check_tqdm_parallel -class TimModel: +class Model: def __init__( self, kaq=[1, 1], @@ -910,6 +910,11 @@ def aquifer_summary(self): pandas.DataFrame dataframe with summary of aquifer(s) parameters """ + if type(self) is Model: + raise NotImplementedError( + "aquifer_summary is not supported for the base Model class; " + "use ModelMaq, Model3D instead." + ) aqs = {} if not isinstance(self.aq, SimpleAquifer): aqs["background"] = self.aq.summary() @@ -918,7 +923,7 @@ def aquifer_summary(self): return pd.concat(aqs, axis=0) -class ModelMaq(TimModel): +class ModelMaq(Model): """Create model specifying a multi-aquifer sequence of aquifer-leakylayer-etc. Parameters @@ -1028,7 +1033,7 @@ def __init__( self.name = "ModelMaq" -class Model3D(TimModel): +class Model3D(Model): """Create a multi-layer model object consisting of many aquifer layers. The @@ -1157,7 +1162,7 @@ def __init__( self.name = "Model3D" -class ModelXsection(TimModel): +class ModelXsection(Model): r"""Model class for cross-section models. Parameters diff --git a/timflow/transient/plots.py b/timflow/transient/plots.py index 68dc91c9..0df4dc79 100644 --- a/timflow/transient/plots.py +++ b/timflow/transient/plots.py @@ -95,6 +95,86 @@ def head_along_line( ax.grid(True) return ax + def discharge_along_line( + self, + x1=0, + x2=1, + y1=0, + y2=0, + npoints=100, + t=1.0, + layers=0, + sstart=0, + color=None, + lw=1, + figsize=None, + ax=None, + legend=True, + grid=True, + ): + """Plot head along line. + + Parameters + ---------- + x1, x2, y1, y2 : float + start and end coordinates of line + npoints : int + number of points along line + t : scalar or array + times at which to plot discharge + layers : + layers for which to plot discharge + sstart : float + starting distance for cross-section + color : str + color of line + lw : float + line width + figsize : tuple of 2 values + size of figure + ax : matplotlib.Axes + axes to plot on, default is None which creates a new figure + legend : bool + add legend to plot + grid : bool + add grid to plot + + Returns + ------- + ax : matplotlib.Axes + axes with plot + """ + layers = np.atleast_1d(layers) + t = np.atleast_1d(t) + if ax is None: + _, ax = plt.subplots(figsize=figsize) + x = np.linspace(x1, x2, npoints) + y = np.linspace(y1, y2, npoints) + s = np.sqrt((x - x[0]) ** 2 + (y - y[0]) ** 2) + sstart + qx, qy = self._ml.disvecalongline(x, y, t, layers) + if x1 == x2: + direction = "y" + qs = qy + elif y1 == y2: + direction = "x" + qs = qx + else: + direction = "s" + qs = np.sqrt(qx**2 + qy**2) + nlayers, ntime, npoints = qx.shape + for i in range(nlayers): + for j in range(ntime): + lbl = f"q$_{direction}$ (t={t[j]}, layer={layers[i]})" + if color is None: + ax.plot(s, qs[i, j, :], lw=lw, label=lbl) + else: + ax.plot(s, qs[i, j, :], color, lw=lw, label=lbl) + if legend: + ax.legend(loc=(0, 1), ncol=3, frameon=False) + if grid: + ax.grid(True) + return ax + def contour( self, win, diff --git a/timflow/transient/stripareasink.py b/timflow/transient/stripareasink.py index 2a7ba20b..c8741b05 100644 --- a/timflow/transient/stripareasink.py +++ b/timflow/transient/stripareasink.py @@ -135,14 +135,14 @@ def __repr__(self): return f"{self.__class__.__name__}: " + str([self.x1, self.x2]) def initialize(self): - if np.isfinite(self.x1) and np.isfinite(self.x2): - self.xc = (self.x1 + self.x2) / 2.0 - if not np.isfinite(self.x1): + if np.isneginf(self.x1) and np.isposinf(self.x2): + self.xc = 0.0 + elif np.isneginf(self.x1): self.xc = self.x2 - 1e-5 - elif not np.isfinite(self.x2): + elif np.isposinf(self.x2): self.xc = self.x1 + 1e-5 else: - self.xc = 0.0 + self.xc = (self.x1 + self.x2) / 2.0 self.L = np.abs(self.x2 - self.x1) self.aq = self.model.aq.find_aquifer_data(self.xc, 0.0) self.setbc()