add custom plot export to CSV

src/View/InitialConditions/Window.py

@@ -263,6 +263,7 @@ class InitialConditionsWindow(PamhyrWindow):
         workdir = os.path.dirname(self._study.filename)
 
         return self.file_dialog(
+            select_file=True,
             callback=lambda d: self._import_from_file(d[0]),
             directory=workdir,
             default_suffix=".BIN",

src/View/Results/CustomPlot/CustomPlotValuesSelectionDialog.py

@@ -45,6 +45,7 @@ class CustomPlotValuesSelectionDialog(PamhyrDialog):
         self.setup_check_boxs()
 
         self.value = None
+        self.export_to_csv = False
 
     def setup_radio_buttons(self):
         self._radio = []
@@ -95,4 +96,6 @@ class CustomPlotValuesSelectionDialog(PamhyrDialog):
         )
 
         self.value = x, y
+        self.export_to_csv = self.find(QCheckBox,
+                                       "checkBox_export").isChecked()
         super().accept()

src/View/Results/Window.py

@@ -20,6 +20,8 @@ import os
 import csv
 import logging
 
+from numpy import sqrt
+
 from datetime import datetime
 from tools import trace, timer, logger_exception
 
@@ -489,12 +491,17 @@ class ResultsWindow(PamhyrWindow):
         dlg = CustomPlotValuesSelectionDialog(parent=self)
         if dlg.exec():
             x, y = dlg.value
-            self.create_new_tab_custom_plot(x, y)
+            export = dlg.export_to_csv
+            self.create_new_tab_custom_plot(x, y, export)
 
-    def create_new_tab_custom_plot(self, x: str, y: list):
+    def create_new_tab_custom_plot(self, x: str, y: list, export: bool):
         name = f"{x}: {','.join(y)}"
         wname = f"tab_custom_{x}_{y}"
 
+        if export:
+            self.export(x, y)
+            return
+
         tab_widget = self.find(QTabWidget, f"tabWidget")
 
         # This plot already exists
@@ -534,6 +541,7 @@ class ResultsWindow(PamhyrWindow):
         grid.addWidget(canvas, 1, 0)
         widget.setLayout(grid)
         tab_widget.addTab(widget, name)
+        tab_widget.setCurrentWidget(widget)
 
     def _copy(self):
         logger.info("TODO: copy")
@@ -582,18 +590,28 @@ class ResultsWindow(PamhyrWindow):
             self._button_last.setEnabled(True)
             self._button_play.setIcon(self._icon_start)
 
-    def export(self):
+    def export(self, x: str, y: list):
+
+        logger.debug(
+            "Export custom plot for: " +
+            f"{x} -> {','.join(y)}"
+        )
         self.file_dialog(
-            select_file=False,
+            callback=lambda f: self.export_to(f[0], x, y),
-            callback=lambda d: self.export_to(d[0])
+            default_suffix=".csv",
+            file_filter=["CSV (*.csv)"],
         )
 
-    def export_to(self, directory):
+    def export_to(self, filename, x, y):
         timestamps = sorted(self._results.get("timestamps"))
-        for reach in self._results.river.reachs:
+        if x == "rk":
-            self.export_reach(reach, directory, timestamps)
+            timestamp = self._get_current_timestamp()
+            self._export_rk(timestamp, y, filename)
+        elif x == "time":
+            profile = self._get_current_profile()
+            self._export_time(profile, y, filename)
 
-    def export_reach(self, reach, directory, timestamps):
+    def export_all(self, reach, directory, timestamps):
         name = reach.name
         name = name.replace(" ", "-")
         if len(timestamps) == 1:
@@ -607,25 +625,6 @@ class ResultsWindow(PamhyrWindow):
         with open(file_name, 'w', newline='') as csvfile:
             writer = csv.writer(csvfile, delimiter=',',
                                 quotechar='|', quoting=csv.QUOTE_MINIMAL)
-            if len(timestamps) > 1:
-                writer.writerow(["name", "rk", "data-file"])
-                for profile in reach.profiles:
-                    p_file_name = os.path.join(
-                        directory,
-                        f"cs_{profile.geometry.id}.csv"
-                    )
-
-                    writer.writerow([
-                        profile.name,
-                        profile.rk,
-                        p_file_name
-                    ])
-
-                    self.export_profile(reach,
-                                        profile,
-                                        p_file_name,
-                                        timestamps)
-            else:
             ts = timestamps[0]
             writer.writerow(self._table["raw_data"]._headers)
             for row in range(self._table["raw_data"].rowCount()):
@@ -635,20 +634,6 @@ class ResultsWindow(PamhyrWindow):
                     line.append(self._table["raw_data"].data(index))
                 writer.writerow(line)
 
-    def export_profile(self, reach, profile, file_name, timestamps):
-        with open(file_name, 'w', newline='') as csvfile:
-            writer = csv.writer(csvfile, delimiter=',',
-                                quotechar='|', quoting=csv.QUOTE_MINIMAL)
-
-            writer.writerow(["timestamp", "z", "q"])
-
-            for ts in timestamps:
-                writer.writerow([
-                    ts,
-                    profile.get_ts_key(ts, "Z"),
-                    profile.get_ts_key(ts, "Q"),
-                ])
-
     def export_current(self):
         self.file_dialog(
             select_file=False,
@@ -657,9 +642,144 @@ class ResultsWindow(PamhyrWindow):
 
     def export_current_to(self, directory):
         reach = self._results.river.reachs[self._get_current_reach()]
-        self.export_reach(reach, directory, [self._get_current_timestamp()])
+        self.export_all(reach, directory, [self._get_current_timestamp()])
 
     def delete_tab(self, index):
         tab_widget = self.find(QTabWidget, f"tabWidget")
         self._additional_plot.pop(tab_widget.tabText(index))
         tab_widget.removeTab(index)
+
+    def _export_rk(self, timestamp, y, filename):
+        reach = self._results.river.reachs[self._get_current_reach()]
+        rk = reach.geometry.get_rk()
+        my_dict = {}
+        if "elevation" in y:
+            my_dict["elevation"] = reach.geometry.get_z_min()
+        if "discharge" in y:
+            my_dict["discharge"] = list(
+                map(
+                    lambda p: p.get_ts_key(timestamp, "Q"),
+                    reach.profiles
+                )
+            )
+        if "water_elevation" in y:
+            my_dict["water_elevation"] = list(
+                map(
+                    lambda p: p.get_ts_key(timestamp, "Z"),
+                    reach.profiles
+                )
+            )
+        if "velocity" in y:
+            my_dict["velocity"] = list(
+                map(
+                    lambda p: p.geometry.speed(
+                        p.get_ts_key(timestamp, "Q"),
+                        p.get_ts_key(timestamp, "Z")),
+                    reach.profiles
+                )
+            )
+        if "depth" in y:
+            my_dict["depth"] = list(
+                map(
+                    lambda p: p.geometry.max_water_depth(
+                        p.get_ts_key(timestamp, "Z")),
+                    reach.profiles
+                )
+            )
+        if "mean_depth" in y:
+            my_dict["mean_depth"] = list(
+                map(
+                    lambda p: p.geometry.mean_water_depth(
+                        p.get_ts_key(timestamp, "Z")),
+                    reach.profiles
+                )
+            )
+        if "froude" in y:
+            my_dict["froude"] = list(
+                map(
+                    lambda p:
+                        p.geometry.speed(
+                            p.get_ts_key(timestamp, "Q"),
+                            p.get_ts_key(timestamp, "Z")) /
+                        sqrt(9.81 * (
+                            p.geometry.wet_area(
+                                p.get_ts_key(timestamp, "Z")) /
+                            p.geometry.wet_width(
+                                p.get_ts_key(timestamp, "Z"))
+                            )),
+                    reach.profiles
+                )
+            )
+        if "wet_area" in y:
+            my_dict["wet_area"] = list(
+                map(
+                    lambda p: p.geometry.wet_area(
+                        p.get_ts_key(timestamp, "Z")),
+                    reach.profiles
+                )
+            )
+
+        with open(filename, 'w', newline='') as csvfile:
+            writer = csv.writer(csvfile, delimiter=',',
+                                quotechar='|', quoting=csv.QUOTE_MINIMAL)
+            header = ["rk"] + y
+            writer.writerow(header)
+            for row in range(len(rk)):
+                line = [rk[row]]
+                for var in y:
+                    line.append(my_dict[var][row])
+                writer.writerow(line)
+
+    def _export_time(self, profile, y, filename):
+        reach = self._results.river.reachs[self._get_current_reach()]
+        profile = reach.profile(profile)
+        ts = list(self._results.get("timestamps"))
+        ts.sort()
+        my_dict = {}
+        z = profile.get_key("Z")
+        q = profile.get_key("Q")
+        if "elevation" in y:
+            my_dict["elevation"] = [profile.geometry.z_min()] * len(ts)
+        if "discharge" in y:
+            my_dict["discharge"] = q
+        if "water_elevation" in y:
+            my_dict["water_elevation"] = z
+        if "velocity" in y:
+            my_dict["velocity"] = list(
+                map(
+                    lambda q, z: profile.geometry.speed(q, z),
+                    q, z
+                )
+            )
+        if "depth" in y:
+            my_dict["depth"] = list(
+                map(lambda z: profile.geometry.max_water_depth(z), z)
+            )
+        if "mean_depth" in y:
+            my_dict["mean_depth"] = list(
+                map(lambda z: profile.geometry.mean_water_depth(z), z)
+            )
+        if "froude" in y:
+            my_dict["froude"] = list(
+                map(lambda z, q:
+                    profile.geometry.speed(q, z) /
+                    sqrt(9.81 * (
+                        profile.geometry.wet_area(z) /
+                        profile.geometry.wet_width(z))
+                    ), z, q)
+                )
+        if "wet_area" in y:
+            my_dict["wet_area"] = list(
+                map(lambda z: profile.geometry.wet_area(z), z)
+            )
+
+        with open(filename, 'w', newline='') as csvfile:
+            writer = csv.writer(csvfile, delimiter=',',
+                                quotechar='|', quoting=csv.QUOTE_MINIMAL)
+            header = ["time"] + y
+            writer.writerow(header)
+            for row in range(len(ts)):
+                line = [ts[row]]
+                for var in y:
+                    line.append(my_dict[var][row])
+                writer.writerow(line)

src/View/Tools/ASubWindow.py

@@ -85,7 +85,7 @@ class WindowToolKit(object):
 
         return header, values
 
-    def file_dialog(self, select_file=True,
+    def file_dialog(self, select_file=None,
                     callback=lambda x: None,
                     directory=None,
                     default_suffix=None,
@@ -107,16 +107,19 @@ class WindowToolKit(object):
 
         dialog = QFileDialog(self, options=options)
 
+        if select_file is not None:
             if select_file:
                 mode = QFileDialog.FileMode.ExistingFile
             else:
                 mode = QFileDialog.FileMode.Directory
+        else:
+            mode = QFileDialog.FileMode.AnyFile
 
         dialog.setFileMode(mode)
         if directory is not None:
             dialog.setDirectory(directory)
 
-        if select_file:
+        if select_file is not False:
             if default_suffix is not None:
                 dialog.setDefaultSuffix(default_suffix)
 

src/View/ui/CustomPlotValuesSelectionDialog.ui

@@ -7,7 +7,7 @@
     <x>0</x>
     <y>0</y>
     <width>414</width>
-    <height>482</height>
+    <height>81</height>
    </rect>
   </property>
   <property name="windowTitle">
@@ -44,6 +44,22 @@
     </widget>
    </item>
    <item row="1" column="0">
+    <widget class="Line" name="line">
+     <property name="orientation">
+      <enum>Qt::Horizontal</enum>
+     </property>
+    </widget>
+   </item>
+   <item row="4" column="0">
+    <layout class="QHBoxLayout" name="horizontalLayout">
+     <item>
+      <widget class="QCheckBox" name="checkBox_export">
+       <property name="text">
+        <string>Export to CSV</string>
+       </property>
+      </widget>
+     </item>
+     <item>
       <widget class="QDialogButtonBox" name="buttonBox">
        <property name="orientation">
         <enum>Qt::Horizontal</enum>
@@ -54,25 +70,11 @@
       </widget>
      </item>
     </layout>
+   </item>
+  </layout>
  </widget>
  <resources/>
  <connections>
-  <connection>
-   <sender>buttonBox</sender>
-   <signal>accepted()</signal>
-   <receiver>Dialog</receiver>
-   <slot>accept()</slot>
-   <hints>
-    <hint type="sourcelabel">
-     <x>248</x>
-     <y>254</y>
-    </hint>
-    <hint type="destinationlabel">
-     <x>157</x>
-     <y>274</y>
-    </hint>
-   </hints>
-  </connection>
   <connection>
    <sender>buttonBox</sender>
    <signal>rejected()</signal>
@@ -89,5 +91,21 @@
     </hint>
    </hints>
   </connection>
+  <connection>
+   <sender>buttonBox</sender>
+   <signal>accepted()</signal>
+   <receiver>Dialog</receiver>
+   <slot>accept()</slot>
+   <hints>
+    <hint type="sourcelabel">
+     <x>248</x>
+     <y>254</y>
+    </hint>
+    <hint type="destinationlabel">
+     <x>157</x>
+     <y>274</y>
+    </hint>
+   </hints>
+  </connection>
  </connections>
 </ui>