Merge pull request #246 from ArtesiaWater/label_layers

Add DatasetCrossSection.label_layers
gwmod · Aug 25, 2023 · 021fc6a · 021fc6a
2 parents 2800853 + 7f3d31c
commit 021fc6a
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 5 deletions.
diff --git a/nlmod/plot/dcs.py b/nlmod/plot/dcs.py
@@ -160,7 +160,14 @@ def line_intersect_grid(self, cs_line):
         xys = xys[xys[:, -1].argsort()]
         return xys
 
-    def plot_layers(self, colors=None, min_label_area=np.inf, fontsize=None, **kwargs):
+    def plot_layers(
+        self,
+        colors=None,
+        min_label_area=np.inf,
+        fontsize=None,
+        only_labels=False,
+        **kwargs,
+    ):
         if colors is None:
             cmap = plt.get_cmap("tab20")
             colors = [cmap(i) for i in range(len(self.layer))]
@@ -202,8 +209,9 @@ def plot_layers(self, colors=None, min_label_area=np.inf, fontsize=None, **kwarg
                 # xy = np.vstack((x, y)).T
                 color = colors[i]
                 pol = matplotlib.patches.Polygon(xy, facecolor=color, **kwargs)
-                self.ax.add_patch(pol)
-                polygons.append(pol)
+                if not only_labels:
+                    self.ax.add_patch(pol)
+                    polygons.append(pol)
 
                 if not np.isinf(min_label_area):
                     pols = Polygon(xy)
@@ -221,16 +229,25 @@ def plot_layers(self, colors=None, min_label_area=np.inf, fontsize=None, **kwarg
                             yp = list(reversed(y[int(len(x) / 2) :]))
                             yp2 = np.interp(xt, xp, yp)
                             yt = np.mean([yp1, yp2])
-                            self.ax.text(
+                            ht = self.ax.text(
                                 xt,
                                 yt,
                                 self.layer[i],
                                 ha="center",
                                 va="center",
                                 fontsize=fontsize,
                             )
+                            if only_labels:
+                                polygons.append(ht)
         return polygons
 
+    def label_layers(self, min_label_area=None):
+        if min_label_area is None:
+            # plot labels of layers with an average thickness of 1 meter
+            # in entire cross-section
+            min_label_area = self.line.length * 1
+        return self.plot_layers(min_label_area=min_label_area, only_labels=True)
+
     def plot_grid(
         self,
         edgecolor="k",

diff --git a/tests/test_011_dcs.py b/tests/test_011_dcs.py
@@ -8,6 +8,7 @@ def test_dcs_structured():
     line = [(0, 0), (1000, 1000)]
     dcs = nlmod.plot.DatasetCrossSection(ds, line)
     dcs.plot_layers()
+    dcs.label_layers()
     dcs.plot_array(ds["kh"], alpha=0.5)
     dcs.plot_grid()
 
@@ -17,5 +18,6 @@ def test_dcs_vertex():
     line = [(0, 0), (1000, 1000)]
     dcs = nlmod.plot.DatasetCrossSection(ds, line)
     dcs.plot_layers()
+    dcs.label_layers()
     dcs.plot_array(ds["kh"], alpha=0.5)
-    dcs.plot_grid()
+    dcs.plot_grid(vertical=False)