fs.to_dataframe enhancement. estimates bug fix.

elphick · Nov 5, 2024 · a88979d · a88979d
1 parent 640e0fe
commit a88979d
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 5 deletions.
diff --git a/elphick/geomet/flowsheet/flowsheet.py b/elphick/geomet/flowsheet/flowsheet.py
@@ -231,7 +231,9 @@ def unhealthy_node_records(self) -> pd.DataFrame:
         unhealthy_nodes = [n for n in self.graph.nodes if
                            self.graph.nodes[n]['mc'].node_type == NodeType.BALANCE and not self.graph.nodes[n][
                                'mc'].is_balanced]
-        unhealthy_data: pd.DataFrame = pd.concat([self.graph.nodes[n]['mc'].unbalanced_records.assign(node=self.graph.nodes[n]['mc'].name) for n in unhealthy_nodes], axis=1)
+        unhealthy_data: pd.DataFrame = pd.concat(
+            [self.graph.nodes[n]['mc'].unbalanced_records.assign(node=self.graph.nodes[n]['mc'].name) for n in
+             unhealthy_nodes], axis=1)
         # move the last column to the front
         unhealthy_data = unhealthy_data[[unhealthy_data.columns[-1]] + list(unhealthy_data.columns[:-1])]
 
@@ -682,22 +684,34 @@ def table_plot(self,
 
         return fig
 
-    def to_dataframe(self, stream_names: Optional[list[str]] = None):
+    def to_dataframe(self, stream_names: Optional[list[str]] = None, tidy: bool = True,
+                     as_mass: bool = False) -> pd.DataFrame:
         """Return a tidy dataframe
 
         Adds the mc name to the index so indexes are unique.
 
         Args:
             stream_names: Optional List of names of Stream/MassComposition objects (network edges) for export
+            tidy: If True, the data will be returned in a tidy format, otherwise wide
+            as_mass: If True, the mass data will be returned instead of the mass-composition data
 
         Returns:
 
         """
         chunks: List[pd.DataFrame] = []
         for u, v, data in self.graph.edges(data=True):
             if (stream_names is None) or ((stream_names is not None) and (data['mc'].name in stream_names)):
-                chunks.append(data['mc'].data.assign(name=data['mc'].name))
-        return pd.concat(chunks, axis='index').set_index('name', append=True)
+                if as_mass:
+                    chunks.append(data['mc'].mass_data.assign(name=data['mc'].name))
+                else:
+                    chunks.append(data['mc'].data.assign(name=data['mc'].name))
+
+        results: pd.DataFrame = pd.concat(chunks, axis='index').set_index('name', append=True)
+        if not tidy:  # wide format
+            results = results.unstack(level='name')
+            results.columns = [f'{col[1]}_{col[0]}' for col in results.columns]
+
+        return results
 
     def plot_parallel(self,
                       names: Optional[str] = None,

diff --git a/elphick/geomet/utils/estimates.py b/elphick/geomet/utils/estimates.py
@@ -65,7 +65,7 @@ def coerce_estimates(estimate_stream: Stream, input_stream: Stream,
     estimate_stream = estimate_stream.balance_composition()
 
     # clip the recovery
-    estimate_stream = estimate_stream.clip_recovery(recovery_bounds=(0.01, 0.99))
+    estimate_stream = estimate_stream.clip_recovery(other=input_stream, recovery_bounds=(0.01, 0.99))
 
     if estimate_stream.status.ok is False:
         raise ValueError('Estimate stream is not OK - it should be after bounding recovery')