Merge pull request #66 from JakobBD/pivot_nda_to_df

Add option to to_df function to make one of the dims to columns

examples/example3.ipynb

@@ -3299,8 +3299,7 @@
     }
    ],
    "source": [
-    "stock_df = dynamic_stock.stock.to_df(index=False)\n",
-    "stock_df = stock_df.pivot(index=\"Time\", columns=\"Region\")[\"value\"]\n",
+    "stock_df = dynamic_stock.stock.to_df(dim_to_columns=\"Region\")\n",
     "\n",
     "fig = px.line(stock_df, title=\"In-use stocks of steel\")\n",
     "fig.show()"
@@ -6385,8 +6384,7 @@
     "inflow = dynamic_stock.inflow\n",
     "outflow = dynamic_stock.outflow\n",
     "with np.errstate(divide=\"ignore\"):\n",
-    "    ratio_df = (outflow / inflow).to_df(index=False)\n",
-    "ratio_df = ratio_df.pivot(index=\"Time\", columns=\"Region\")[\"value\"]\n",
+    "    ratio_df = (outflow / inflow).to_df(dim_to_columns=\"Region\")\n",
     "\n",
     "fig = px.line(ratio_df, title=\"Ratio outflow:inflow\")\n",
     "fig.show()"

examples/example3.py

@@ -163,8 +163,7 @@ def read_parameter_values(self, parameter: str, dims: DimensionSet) -> Parameter
 # First, we plot the steel stock in the different countries over time.
 
 # %%
-stock_df = dynamic_stock.stock.to_df(index=False)
-stock_df = stock_df.pivot(index="Time", columns="Region")["value"]
+stock_df = dynamic_stock.stock.to_df(dim_to_columns="Region")
 
 fig = px.line(stock_df, title="In-use stocks of steel")
 fig.show()
@@ -176,8 +175,7 @@ def read_parameter_values(self, parameter: str, dims: DimensionSet) -> Parameter
 inflow = dynamic_stock.inflow
 outflow = dynamic_stock.outflow
 with np.errstate(divide="ignore"):
-    ratio_df = (outflow / inflow).to_df(index=False)
-ratio_df = ratio_df.pivot(index="Time", columns="Region")["value"]
+    ratio_df = (outflow / inflow).to_df(dim_to_columns="Region")
 
 fig = px.line(ratio_df, title="Ratio outflow:inflow")
 fig.show()

examples/example5.ipynb

@@ -174,8 +174,8 @@
     "        self.stocks[\"in use\"].inflow[...] = self.flows[\"sysenv => market\"]\n",
     "        survival_model = NormalSurvival(\n",
     "            dims=self.stocks[\"in use\"].inflow.dims,\n",
-    "            lifetime_mean=self.parameters[\"vehicle lifetime\"].values,\n",
-    "            lifetime_std=self.parameters[\"vehicle lifetime\"].values * 0.3,\n",
+    "            lifetime_mean=self.parameters[\"vehicle lifetime\"],\n",
+    "            lifetime_std=self.parameters[\"vehicle lifetime\"] * 0.3,\n",
     "        )\n",
     "        if not isinstance(self.stocks[\"in use\"], InflowDrivenDSM):\n",
     "            self.stocks[\"in use\"] = self.stocks[\"in use\"].to_stock_type(\n",
@@ -4556,11 +4556,7 @@
    "source": [
     "np.nan_to_num(vehicle_mfa_2.flows[\"scrap => sysenv\"].values, copy=False)\n",
     "scrap_outflow = vehicle_mfa_2.flows[\"scrap => sysenv\"].sum_nda_over(sum_over_dims=(\"r\", \"m\"))\n",
-    "outflow_df = pd.DataFrame(\n",
-    "    data=scrap_outflow.values,\n",
-    "    columns=scrap_outflow.dims[\"w\"].items,\n",
-    "    index=scrap_outflow.dims[\"t\"].items,\n",
-    ")\n",
+    "outflow_df = scrap_outflow.to_df(dim_to_columns=\"waste\")\n",
     "outflow_df = outflow_df[outflow_df.index > 2017]\n",
     "fig = px.line(outflow_df, title=\"Scrap outflow\")\n",
     "fig.show()\n",

examples/example5.py

@@ -135,8 +135,8 @@ def compute_stock(self):
         self.stocks["in use"].inflow[...] = self.flows["sysenv => market"]
         survival_model = NormalSurvival(
             dims=self.stocks["in use"].inflow.dims,
-            lifetime_mean=self.parameters["vehicle lifetime"].values,
-            lifetime_std=self.parameters["vehicle lifetime"].values * 0.3,
+            lifetime_mean=self.parameters["vehicle lifetime"],
+            lifetime_std=self.parameters["vehicle lifetime"] * 0.3,
         )
         if not isinstance(self.stocks["in use"], InflowDrivenDSM):
             self.stocks["in use"] = self.stocks["in use"].to_stock_type(
@@ -349,11 +349,7 @@ def vehicle_new_registration(self, data, dims):
 # %%
 np.nan_to_num(vehicle_mfa_2.flows["scrap => sysenv"].values, copy=False)
 scrap_outflow = vehicle_mfa_2.flows["scrap => sysenv"].sum_nda_over(sum_over_dims=("r", "m"))
-outflow_df = pd.DataFrame(
-    data=scrap_outflow.values,
-    columns=scrap_outflow.dims["w"].items,
-    index=scrap_outflow.dims["t"].items,
-)
+outflow_df = scrap_outflow.to_df(dim_to_columns="waste")
 outflow_df = outflow_df[outflow_df.index > 2017]
 fig = px.line(outflow_df, title="Scrap outflow")
 fig.show()

sodym/named_dim_arrays.py

@@ -237,12 +237,18 @@ def __setitem__(self, keys, item):
         self.values[slice_obj.ids] = item.sum_values_to(slice_obj.dim_letters)
         return
 
-    def to_df(self, index: bool = True):
+    def to_df(self, index: bool = True, dim_to_columns: str = None) -> pd.DataFrame:
         multiindex = pd.MultiIndex.from_product([d.items for d in self.dims], names=self.dims.names)
         df = pd.DataFrame({"value": self.values.flatten()})
         df = df.set_index(multiindex)
+        if dim_to_columns is not None:
+            if dim_to_columns not in self.dims.names:
+                raise ValueError(f"Dimension name {dim_to_columns} not found in nda.dims.names")
+            df.reset_index(inplace=True)
+            index_names = [n for n in self.dims.names if n != dim_to_columns]
+            df = df.pivot(index=index_names, columns=dim_to_columns, values="value")
         if not index:
-            df = df.reset_index()
+            df.reset_index(inplace=True)
         return df
 
     def split(self, dim_letter: str) -> dict: