{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":70367,"databundleVersionId":9188054,"sourceType":"competition"}],"dockerImageVersionId":30746,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Ariel Data Challenge 2024 - Subset Creation\n\n## Purpose\n\nThis notebook aims to create a smaller subset of the Ariel Data Challenge 2024 dataset while preserving the original folder structure. This subset will be used for initial data exploration and analysis. For people like me that likes to work locally without having 174 gb data installed.","metadata":{}},{"cell_type":"code","source":"import os\nimport pandas as pd\nimport pyarrow.parquet as pq\nimport zipfile\n\n# Define data directories and paths\ndata_dir = '/kaggle/input/ariel-data-challenge-2024'\nsample_dir = '/kaggle/working/sample_data'\nzip_file_path = '/kaggle/working/sample_data.zip'\n\nadc_info_path = os.path.join(data_dir, 'train_adc_info.csv')\nlabels_path = os.path.join(data_dir, 'train_labels.csv')\naxis_info_path = os.path.join(data_dir, 'axis_info.parquet')\nwavelength_path = os.path.join(data_dir, 'wavelengths.csv')\n\n# Read input data\nadc_info = pd.read_csv(adc_info_path)\nlabels = pd.read_csv(labels_path)\naxis_info = pq.read_table(axis_info_path).to_pandas()\nwavelengths = pd.read_csv(wavelength_path)\n\n# Ensure sample directory exists\nos.makedirs(sample_dir, exist_ok=True)\n\n# Number of planets to sample\nnum_of_planets = 10\n\n# Sample planet IDs\nsample_planet_ids = labels['planet_id'].sample(n=num_of_planets).tolist()\nprint(\"Sampled Planet IDs:\", sample_planet_ids)\n\n# Save sampled metadata\nadc_info_sample = adc_info[adc_info['planet_id'].isin(sample_planet_ids)]\nadc_info_sample.to_csv(os.path.join(sample_dir, 'train_adc_info.csv'), index=False)\n\nlabels_sample = labels[labels['planet_id'].isin(sample_planet_ids)]\nlabels_sample.to_csv(os.path.join(sample_dir, 'train_labels.csv'), index=False)\nprint(\"Sampled Metadata Saved\")\n\n# Function to save sample data for a planet and instrument\ndef save_sample_data(planet_id, instrument):\n    signal_path = os.path.join(data_dir, f'train/{planet_id}/{instrument}_signal.parquet')\n    signal_sample = pq.read_table(signal_path).to_pandas()\n    signal_sample_dir = os.path.join(sample_dir, f'train/{planet_id}')\n    os.makedirs(signal_sample_dir, exist_ok=True)\n    signal_sample.to_parquet(os.path.join(signal_sample_dir, f'{instrument}_signal.parquet'))\n\n    calibration_files = ['dead', 'linear_corr', 'read', 'flat', 'dark']\n    for calib in calibration_files:\n        calib_path = os.path.join(data_dir, f'train/{planet_id}/{instrument}_calibration/{calib}.parquet')\n        calib_sample = pq.read_table(calib_path).to_pandas()\n        calib_sample_dir = os.path.join(sample_dir, f'train/{planet_id}/{instrument}_calibration')\n        os.makedirs(calib_sample_dir, exist_ok=True)\n        calib_sample.to_parquet(os.path.join(calib_sample_dir, f'{calib}.parquet'))\n\n# Save sample data for each selected planet and instrument\ninstruments = ['AIRS-CH0', 'FGS1']\nfor planet_id in sample_planet_ids:\n    for instrument in instruments:\n        save_sample_data(planet_id, instrument)\nprint(\"Sampled Data Saved for Each Planet and Instrument\")\n\n# Copy the full small files\nos.system(f'cp {axis_info_path} {sample_dir}/axis_info.parquet')\nos.system(f'cp {wavelength_path} {sample_dir}/wavelengths.csv')\nprint(\"Additional Files Copied\")\n\n# Zip the sample data directory\nwith zipfile.ZipFile(zip_file_path, 'w') as zipf:\n    for root, dirs, files in os.walk(sample_dir):\n        for file in files:\n            file_path = os.path.join(root, file)\n            arcname = os.path.relpath(file_path, start=sample_dir)\n            zipf.write(file_path, arcname)\n\nprint(f\"Sample data saved and zipped to {zip_file_path}\")","metadata":{"execution":{"iopub.status.busy":"2024-08-07T00:13:49.772574Z","iopub.execute_input":"2024-08-07T00:13:49.773327Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}