From a13ac1e759813e7dde572bd31d41d583935c783a Mon Sep 17 00:00:00 2001
From: Aleksandra Apolinarska <aapolina@ethz.ch>
Date: Thu, 13 Feb 2025 14:45:09 +0100
Subject: [PATCH] #66 ghcomponent to merge datasets

---
 src/aixd_ara/api.py                           |  16 +++
 .../components/ara_DatasetsMerge/code.py      |  25 ++++
 .../components/ara_DatasetsMerge/icon.png     | Bin 0 -> 1688 bytes
 .../ara_DatasetsMerge/metadata.json           |  45 ++++++++
 src/aixd_ara/controller.py                    | 108 ++++++++++++++++++
 src/aixd_ara/gh_ui.py                         |  10 ++
 6 files changed, 204 insertions(+)
 create mode 100644 src/aixd_ara/components/ara_DatasetsMerge/code.py
 create mode 100644 src/aixd_ara/components/ara_DatasetsMerge/icon.png
 create mode 100644 src/aixd_ara/components/ara_DatasetsMerge/metadata.json

diff --git a/src/aixd_ara/api.py b/src/aixd_ara/api.py
index 8314243..4ca3ca6 100644
--- a/src/aixd_ara/api.py
+++ b/src/aixd_ara/api.py
@@ -376,6 +376,22 @@ def ara_welcome(*args):
     print()
 
 
+@app.route("/merge_datasets", methods=["POST"])
+def merge_datasets():
+    data = request.data
+    data = json.loads(data)
+    session_id = data["session_id"]
+    sc = SessionController.create(session_id)
+
+    result = sc.merge_datasets(
+        root_folder=data["root_folder"],
+        new_dataset_name=data["new_dataset_name"],
+        samples_per_file=data["samples_per_file"],
+    )
+    response = json.dumps(result, cls=DataEncoder)
+    return response
+
+
 if __name__ == "__main__":
     import sys
 
diff --git a/src/aixd_ara/components/ara_DatasetsMerge/code.py b/src/aixd_ara/components/ara_DatasetsMerge/code.py
new file mode 100644
index 0000000..f2ec77f
--- /dev/null
+++ b/src/aixd_ara/components/ara_DatasetsMerge/code.py
@@ -0,0 +1,25 @@
+# flake8: noqa
+from scriptcontext import sticky as st
+from Grasshopper.Kernel.GH_RuntimeMessageLevel import Error, Warning
+
+from aixd_ara.gh_ui import merge_datasets
+from aixd_ara.gh_ui_helper import clear_sticky
+from aixd_ara.gh_ui_helper import component_id
+from aixd_ara.gh_ui_helper import session_id
+
+if samples_per_file is None: samples_per_file = 1000
+assert samples_per_file>0, "samples_per_file must be a positive integer, got {}".format(samples_per_file)
+
+if root_folder:
+    cid = component_id(session_id(), ghenv.Component, "ProjectSetup")
+    
+    if merge:
+        st[cid] = merge_datasets(session_id(), root_folder, new_dataset_name,samples_per_file)
+    
+    if cid in st.keys():
+        status = st[cid]["status"]
+        msg = st[cid]["msg"]
+        # if status=="error":
+        #     ghenv.Component.AddRuntimeMessage(Error, msg)
+        # elif status=="warning":
+        #     ghenv.Component.AddRuntimeMessage(Warning, msg)
\ No newline at end of file
diff --git a/src/aixd_ara/components/ara_DatasetsMerge/icon.png b/src/aixd_ara/components/ara_DatasetsMerge/icon.png
new file mode 100644
index 0000000000000000000000000000000000000000..01f69de672948a9c8f1ed6e138bc87b34fb78d56
GIT binary patch
literal 1688
zcmV;J250$+P)<h;3K|Lk000e1NJLTq000;O000;W1^@s6;CDUv00009a7bBm000XT
z000XT0n*)m`~Uz3N=ZaPR7i<^Rc~w*#}$8Xc5Zj?&R3t$=R0GZJF7<6u1Rdg5~-0=
zI}vK7B#JPR5<94ZNR9eWY7Om|QmMc<QmHMiRMa+#3awN*{SZ+Jfm6|7M<J@(5@Tw>
zQ82zcY{J?2&iBW6ySF>@s#&n(kcyx5c4lYZzBli^-@HX4f`{>~@$)eve3lTmB0wDw
zAs~K1Z!<hAM1bpa^^^eoVa3<apBs{>bprrHfOV_tn@4)U&&<N+$%cdPzxz(Wmc};;
zSR;5*-471&f#C-)OE+$WtNFZX6;@5PY}SV@e?!=cb#|?OW}NQJx|PxKF@FVu`y^?N
z1jqo85r7(1Vz9itWX`?!kG9;cTU#C#MWns`Zu{Ufsk&&?q6dDw2n4VJNf`EWhyV+h
zFE^F4*-&L=<x6h0YC6SIoyHjJ>gd3Z&MqAI+S8CoB)l`G(rKEXpQo9b8JL`yY%XLn
zO*>wCX{sTSECWfvpPoM%$;y`35RhdqPsQfmyZD4#bGXJeY~9+bwIvd~E*R7t+t#dV
zm1&xWreWZ7&p%Jwwr`h0N+Be<u9RXvH7^ezK1_OjeR1b&2c}38a9jBLkpyNgUy03*
z|GhUBkKwB$uR{M<`*}DNGEGhM`!vn;G3Fy7c`l#BOeO=BVv$6#C>o;Cn(ex_<+|0D
zc--b(6IP)RVA3TO0k@^z3MARB$@e=0eiMH6i(f*0D8&6fUm&O(^?uGnmnJ9c7BU%M
zCY@#~rHN#cqR}Wd1OiUjFx;|~)!Vl-<v+jub}``dRWu*3ajF4I3B2kU4T`WNC9B#t
zrqe0b*VChCGq)qZK7PFAuVZ7zBkHQ}+qaKK4;?ZMhWyDL-BL>7+(_MY4XOeo7q}3_
z6lwK9z?P0<{QUTFeCmmw`oH|)&9-u>g!}jJhrz)?=<V(GMB&PnE1pbDOu*^Wr}6CB
zv&Omcarw&Ly_K13SGlZ|0(DR#hAwc&RTt_$$G5d_!zX)t*_&^?(I#9Mr>3ShjBdud
zvjYbX!06~G?Ao;pethVV@tfcMRwKtk9UQ2KRbt=(a~)!?3s^0bg9{6bI5IMfj$PB$
z{b*}zdlc<GC>^CcJ3EUj%ULje2=xK1M&O=|U`{f}Ayfj4%W`=Q-g@hg*cgpEmSuS#
zpEz*>QmNDz468M@r*xEVEEa?0V_Ts$-bevviAKCfQCKEkbG)X5+zePjT{kaX`VcMK
z65ZY1P%f9@;K75QC=3h?c;n>c<c6bUG6}7%tq_ScP*Y==^2>L~&uZ8R%ThxD+;O~q
zKtcj+Y1s<j`}XUw|Apst)hiYj7B-5xy1MF(3P+9{fp9pyQT*oSX2|7oczxz72zw2l
zeq2vSXoVO^<c=e~2sE?Knwy&kDFHObaB*=F&YU>|d-m-4zlC0hDs(oRg{Gz^0O`^<
z`kQWk*_>M@s{ja6ML@uiAkx*lxHx$}_Ol~LDNt8et5zy}XlMv}dU`x3>a?4oqEQE!
zo}LENG$D~_l{+7gXY7ogbMBN8sElO{mAR`TU<gKIOF8lNp?}wIUTwd7eOmVQ^}+7l
zyW!NSQ*i9qF;5<b%D+n2g$oyeL`|Suta3N!GPy=T7OX(*Q>IenxDtc_;K7I$@83O3
zTra5nD~>C+<nwvBc=6)rM6Ibk6`ks&j$_rZ^vON8Hec1;5^?FUH3+@0#rt0veDE*v
zAOXlAr}ptbKlsYn*qB$$&dyFx6eZV1QB&;?4-Z43P=LeVdG=p@TT&lmVW!ONoLqs}
z_3bZT`v`aZ==Vo8j=(tspO1mA0|a!$r4Pruubunbi?(I?*6Fx@@zlxF>9jWw22A(p
zcb|IqhcEi3q4bG`Ygu6nz8V}UomBhS@q?2`xlb9_Kvxb71VfV)Fa$<T7gH1CUB$(P
z=G@(E+_7v!5h7@6)%M3ipZ0F^-}%W%!!^yiBbl88bk>+7czy|6{*-*>W<@#R!j)cn
zRH3DYs)Yk+pt>trjxD>N*?Bc;bbJ)oi%zq#>O_4xw~?)idMrUe<`9avISAW_(&aYz
z{~H`Gol9;a0~eBjNdl?Fs{jL1&5sDw3`mgTw-lBXZdni%U19|RSOkQM2H@revZWde
i{wK7zkpCRf0N{U_gGW`#H)(SK0000<MNUMnLSTZiCN!Y{

literal 0
HcmV?d00001

diff --git a/src/aixd_ara/components/ara_DatasetsMerge/metadata.json b/src/aixd_ara/components/ara_DatasetsMerge/metadata.json
new file mode 100644
index 0000000..68b8333
--- /dev/null
+++ b/src/aixd_ara/components/ara_DatasetsMerge/metadata.json
@@ -0,0 +1,45 @@
+{
+    "name": "DatasetsMerge",
+    "nickname": "DatasetsMerge",
+    "category": "ARA",
+    "subcategory": "2 Dataset",
+    "description": "Merges multiple datasets into a single dataset. Requires that the datasets have the same schema (variable names, types, dimensions).",
+    "exposure": 2,
+    "ghpython": {
+        "isAdvancedMode": false,
+        "iconDisplay": 0,
+        "inputParameters": [
+            {
+                "name": "root_folder",
+                "description": "Path to the folder containing the datasets to merge.",
+                "typeHintID": "str",
+                "scriptParamAccess": 0
+            },
+            {
+                "name": "new_dataset_name",
+                "description": "Name of the merged dataset. (Optional, default: 'merged_dataset'.)",
+                "typeHintID": "str",
+                "scriptParamAccess": 0
+            },
+            {
+                "name": "samples_per_file",
+                "description": "Number of samples to be saved in each file of the new dataset. (Optional, default: 1000.)",
+                "typeHintID": "int",
+                "scriptParamAccess": 0
+            },
+            {
+                "name": "merge",
+                "description": "Triggers the merge process.",
+                "typeHintID": "bool",
+                "scriptParamAccess": 0
+            }
+        ],
+
+        "outputParameters": [
+            {
+                "name": "msg",
+                "description": "Message logs."
+            }
+        ]
+    }
+}
diff --git a/src/aixd_ara/controller.py b/src/aixd_ara/controller.py
index 8ca9b79..f0e6341 100644
--- a/src/aixd_ara/controller.py
+++ b/src/aixd_ara/controller.py
@@ -6,6 +6,8 @@
 import base64
 import os
 import random
+import shutil
+import pandas as pd
 
 import pytorch_lightning as pl
 import torch
@@ -608,6 +610,112 @@ def blocknames_from_dataobjects(self, dataobjects):
 
         return list(set(blocknames))
 
+    @staticmethod
+    def merge_datasets(root_folder: str, new_dataset_name: str, samples_per_file: int):
+        """
+        This function merges all datasets in the root_folder into a new dataset with the name new_dataset_name.
+        Prerequisites:
+        - All datasets are stored in subfolders of the root_folder
+        - Each dataset has the same design_parameters and performance_attributes (variable names, types, dimensions)
+
+        Parameters:
+        -----------
+        root_folder: str
+            Path to the root folder containing the datasets.
+        new_dataset_name: str
+            Name of the new dataset. The new dataset will be created in a subfolder with this name. Default is "merged_dataset".
+            If such a folder already exists, it will be overwritten.
+
+        Returns:
+        --------
+        status : str
+            Status of the process: "warning", "error" or None
+        msg : str
+            A message containing information about the process.
+        """
+        # TODO: check what happens with the domains, may need to get a union explicitly.
+
+        status = None
+        msg = ""
+
+        if not os.path.exists(root_folder):
+            txt = f"Folder {root_folder} does not exist.\n"
+            status = "error"
+            return status, txt
+
+        if not new_dataset_name:
+            new_dataset_name = "merged_dataset"
+
+        if os.path.exists(os.path.join(root_folder, new_dataset_name)):
+            txt = f"[WARNING] Folder {new_dataset_name} already exists. It will be overwritten.\n\n"
+            status = "warning"
+            print(txt)
+            msg += txt
+            shutil.rmtree(os.path.join(root_folder, new_dataset_name))
+
+        dataset_names = [
+            name
+            for name in os.listdir(root_folder)
+            if os.path.isdir(os.path.join(root_folder, name)) and name != new_dataset_name
+        ]
+        txt = f"Found following subfolders: {dataset_names}.\n\n"
+        print(txt)
+        msg += txt
+
+        def _load_df(root_folder, dataset_name):
+            # Load old sharded data from pickled dataframes
+
+            # DPs
+            directory = os.path.join(root_folder, dataset_name, "design_parameters")
+            df_dp_all = []
+
+            for filename in os.listdir(directory):
+                if filename.endswith(".pkl"):
+                    filepath = os.path.join(directory, filename)
+                    df = pd.read_pickle(filepath)
+                    df_dp_all.append(df)
+
+            df_dp_all = pd.concat(df_dp_all, axis=0)
+
+            # PAs
+            directory = os.path.join(root_folder, dataset_name, "performance_attributes")
+            df_pa_all = []
+
+            for filename in os.listdir(directory):
+                if filename.endswith(".pkl"):
+                    filepath = os.path.join(directory, filename)
+                    df = pd.read_pickle(filepath)
+                    df_pa_all.append(df)
+
+            df_pa_all = pd.concat(df_pa_all, axis=0)
+            df_all = pd.merge(df_dp_all, df_pa_all, how="inner", on=["uid"])
+            df_all = df_all.drop(columns=["uid"])
+            return df_all
+
+        # load all datasets into a single dataframe
+        dfs = []
+        for dataset_name in dataset_names:
+            df = _load_df(root_folder, dataset_name)
+            dfs.append(df)
+        df_all = pd.concat(dfs)
+
+        # create new Dataset, take one of the old datasets as a template
+        dataset_temp = Dataset.from_dataset_folder(os.path.join(root_folder, dataset_names[0]))
+
+        dataset_new = Dataset(
+            name=new_dataset_name,
+            root_path=root_folder,
+            file_format="json",
+            design_par=dataset_temp.design_par,
+            perf_attributes=dataset_temp.perf_attributes,
+            overwrite=True,
+        )
+
+        # load data
+        dataset_new.import_data_from_df(df_all, samples_perfile=samples_per_file, flag_fromscratch=True)
+        msg += f"New dataset with {len(df_all)} samples has been created in {os.path.join(root_folder, new_dataset_name)}.\n"
+        return {"status": status, "msg": msg}
+
 
 # --------------------------------------------------------------
 # helper methods
diff --git a/src/aixd_ara/gh_ui.py b/src/aixd_ara/gh_ui.py
index 22ee8d4..f028388 100644
--- a/src/aixd_ara/gh_ui.py
+++ b/src/aixd_ara/gh_ui.py
@@ -103,3 +103,13 @@ def model_input_output_dimensions(session_id):
 def request_designs(session_id, request, n_designs):
     data = {"session_id": session_id, "requested_values": request, "n_designs": n_designs}
     return http_post_request(action="request_designs", data=data)
+
+
+def merge_datasets(session_id, root_folder, new_dataset_name, samples_per_file):
+    data = {
+        "session_id": session_id,
+        "root_folder": root_folder,
+        "new_dataset_name": new_dataset_name,
+        "samples_per_file": samples_per_file,
+    }
+    return http_post_request(action="merge_datasets", data=data)