chore: import upstream snapshot with attribution

rdagent/scenarios/data_science/example/source_data/arf-12-hours-prediction-task/prepare.py

@@ -0,0 +1,69 @@
+import random
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+import sparse
+
+CURRENT_DIR = Path(__file__).resolve().parent
+ROOT_DIR = CURRENT_DIR.parent.parent
+
+raw_feature_path = CURRENT_DIR / "X.npz"
+raw_label_path = CURRENT_DIR / "ARF_12h.csv"
+
+public = ROOT_DIR / "arf-12-hours-prediction-task"
+private = ROOT_DIR / "eval" / "arf-12-hours-prediction-task"
+
+if not (public / "test").exists():
+    (public / "test").mkdir(parents=True, exist_ok=True)
+
+if not (public / "train").exists():
+    (public / "train").mkdir(parents=True, exist_ok=True)
+
+if not private.exists():
+    private.mkdir(parents=True, exist_ok=True)
+
+SEED = 42
+random.seed(SEED)
+np.random.seed(SEED)
+
+X_sparse = sparse.load_npz(raw_feature_path)  # COO matrix, shape: [N, D, T]
+df_label = pd.read_csv(raw_label_path)  # Contains column 'ARF_LABEL'
+N = X_sparse.shape[0]
+
+indices = np.arange(N)
+np.random.shuffle(indices)
+split = int(0.7 * N)
+train_idx, test_idx = indices[:split], indices[split:]
+
+X_train = X_sparse[train_idx]
+X_test = X_sparse[test_idx]
+
+df_train = df_label.iloc[train_idx].reset_index(drop=True)
+df_test = df_label.iloc[test_idx].reset_index(drop=True)
+
+submission_df = df_test.copy()
+submission_df["ARF_LABEL"] = 0
+submission_df.drop(submission_df.columns.difference(["ID", "ARF_LABEL"]), axis=1, inplace=True)
+submission_df.to_csv(public / "sample_submission.csv", index=False)
+
+df_test.to_csv(private / "submission_test.csv", index=False)
+
+df_test.drop(["ARF_LABEL"], axis=1, inplace=True)
+df_test.to_csv(public / "test" / "ARF_12h.csv", index=False)
+sparse.save_npz(public / "test" / "X.npz", X_test)
+
+sparse.save_npz(public / "train" / "X.npz", X_train)
+df_train.to_csv(public / "train" / "ARF_12h.csv", index=False)
+
+assert (
+    X_train.shape[0] == df_train.shape[0]
+), f"Mismatch: X_train rows ({X_train.shape[0]}) != df_train rows ({df_train.shape[0]})"
+assert (
+    X_test.shape[0] == df_test.shape[0]
+), f"Mismatch: X_test rows ({X_test.shape[0]}) != df_test rows ({df_test.shape[0]})"
+assert df_test.shape[1] == 2, "Public test set should have 2 columns"
+assert df_train.shape[1] == 3, "Public train set should have 3 columns"
+assert len(df_train) + len(df_test) == len(
+    df_label
+), "Length of new_train and new_test should equal length of old_train"

rdagent/scenarios/data_science/example/source_data/playground-series-s4e9/prepare.py

@@ -0,0 +1,42 @@
+from pathlib import Path
+
+import pandas as pd
+from sklearn.model_selection import train_test_split
+
+
+def prepare(raw: Path, public: Path, private: Path):
+
+    # Create train and test splits from train set
+    old_train = pd.read_csv(raw / "train.csv")
+    new_train, new_test = train_test_split(old_train, test_size=0.1, random_state=0)
+
+    # Create sample submission
+    sample_submission = new_test.copy()
+    sample_submission["price"] = 43878.016
+    sample_submission.drop(sample_submission.columns.difference(["id", "price"]), axis=1, inplace=True)
+    sample_submission.to_csv(public / "sample_submission.csv", index=False)
+
+    # Create private files
+    new_test.to_csv(private / "submission_test.csv", index=False)
+
+    # Create public files visible to agents
+    new_train.to_csv(public / "train.csv", index=False)
+    new_test.drop(["price"], axis=1, inplace=True)
+    new_test.to_csv(public / "test.csv", index=False)
+
+    # Checks
+    assert new_test.shape[1] == 12, "Public test set should have 12 columns"
+    assert new_train.shape[1] == 13, "Public train set should have 13 columns"
+    assert len(new_train) + len(new_test) == len(
+        old_train
+    ), "Length of new_train and new_test should equal length of old_train"
+
+
+if __name__ == "__main__":
+    competitions = "playground-series-s4e9"
+    raw = Path(__file__).resolve().parent
+    prepare(
+        raw=raw,
+        public=raw.parent.parent / competitions,
+        private=raw.parent.parent / "eval" / competitions,
+    )