{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":101849,"databundleVersionId":12846694,"isSourceIdPinned":false,"sourceType":"competition"}],"dockerImageVersionId":31089,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\nimport logging\nfrom typing import Tuple, List\nimport numpy as np\nimport pandas as pd\nimport tensorflow as tf\nfrom tensorflow import keras\nfrom tensorflow.keras import layers\nimport tensorflow_probability as tfp\nfrom tqdm import tqdm\nfrom multiprocessing import Pool, cpu_count\n\n# Force unbuffered output for real-time logging in Kaggle\nos.environ['PYTHONUNBUFFERED'] = '1'\n\n# Configure logging\nlogging.basicConfig(\n    level=logging.INFO,\n    format='%(asctime)s - %(levelname)s - %(message)s',\n    handlers=[logging.StreamHandler()]\n)\nlogger = logging.getLogger(__name__)\n\n# Check dependencies\ntry:\n    import tensorflow as tf\n    import tensorflow_probability as tfp\n    import pandas as pd\n    import numpy as np\n    import tqdm\n    if not tf.__version__.startswith('2.'):\n        logger.warning(f\"TensorFlow version {tf.__version__} detected. Version 2.17.0+ recommended.\")\n    if not tfp.__version__.startswith('0.'):\n        logger.warning(f\"TensorFlow Probability version {tfp.__version__} detected. Version 0.24.0+ recommended.\")\n    logger.info(f\"Using TensorFlow {tf.__version__} and TensorFlow Probability {tfp.__version__}\")\nexcept ImportError as e:\n    logger.error(f\"Missing dependency: {e}. Install: tensorflow==2.17.0, tensorflow-probability==0.24.0, pandas, numpy, tqdm\")\n    raise\n\n# Configuration\nclass Config:\n    \"\"\"Configuration for the Ariel Data Challenge 2025.\"\"\"\n    data_path: str = '/kaggle/input/ariel-data-challenge-2025'\n    train_spec_file: str = 'train.csv'\n    wavelengths_file: str = 'wavelengths.csv'\n    test_ids_file: str = 'sample_submission.csv'\n    train_dir: str = 'train'\n    test_dir: str = 'test'\n    batch_size: int = 32\n    epochs: int = 20\n    cache_dir: str = '/kaggle/working/cache'\n    checkpoint_dir: str = '/kaggle/working/checkpoints'\n    chunk_size: int = 100\n    mc_dropout_samples: int = 5\n    n_outputs: int = 0\n    default_n_outputs: int = 283\n\n# Custom Attention Layer\nclass AttentionLayer(layers.Layer):\n    \"\"\"Custom attention layer to focus on important spectral features.\"\"\"\n    def __init__(self):\n        super().__init__()\n        self.dense = layers.Dense(1, activation=None)\n\n    def call(self, inputs: tf.Tensor) -> tf.Tensor:\n        scores = self.dense(inputs)\n        attention_weights = tf.nn.softmax(scores, axis=1)\n        return inputs * attention_weights\n\n# Physics-Informed Loss Layer\nclass PhysicsLossLayer(layers.Layer):\n    \"\"\"Enforces smoothness and energy variance constraints for spectra.\"\"\"\n    def __init__(self, wavelengths: np.ndarray, phys_reg_weight: float, n_outputs: int):\n        super().__init__()\n        if n_outputs < 3:\n            raise ValueError(f\"n_outputs must be >= 3 for Conv1D with kernel_size=3, got {n_outputs}\")\n        self.wavelengths = tf.convert_to_tensor(wavelengths, dtype=tf.float32)\n        self.phys_reg_weight = phys_reg_weight\n        self.n_outputs = n_outputs\n        kernel = np.array([1., -2., 1.], dtype=np.float32).reshape(3, 1, 1)\n        self.conv = layers.Conv1D(1, 3, padding='valid', use_bias=False, trainable=False)\n        self.conv.build((None, self.n_outputs, 1))\n        self.conv.set_weights([kernel])\n\n    def call(self, inputs: tf.Tensor) -> tf.Tensor:\n        mu = inputs\n        mu_expanded = tf.expand_dims(mu, axis=-1)\n        d2 = self.conv(mu_expanded)\n        smooth_loss = tf.reduce_mean(tf.square(d2), axis=[1, 2])\n        E = tf.reduce_sum(mu / tf.math.pow(self.wavelengths, 4), axis=1)\n        energy_var = tfp.stats.variance(E)\n        phys_loss = self.phys_reg_weight * (smooth_loss + 0.1 * energy_var)\n        return phys_loss\n\n    def get_config(self) -> dict:\n        return {\n            'wavelengths': self.wavelengths.numpy(),\n            'phys_reg_weight': self.phys_reg_weight,\n            'n_outputs': self.n_outputs\n        }\n\n# Model Definition\nclass ArielModel(keras.Model):\n    \"\"\"CNN-based model with attention for spectral prediction and uncertainties.\"\"\"\n    def __init__(\n        self,\n        input_shape: tuple,\n        output_shape: int,\n        wavelengths: np.ndarray,\n        phys_reg_weight: float,\n        dropout_rate: float\n    ):\n        super().__init__()\n        self.conv1 = layers.Conv1D(64, 5, padding='same', activation='relu')\n        self.bn1 = layers.BatchNormalization()\n        self.conv2 = layers.Conv1D(32, 3, padding='same', activation='relu')\n        self.bn2 = layers.BatchNormalization()\n        self.attention = AttentionLayer()\n        self.flatten = layers.Flatten()\n        self.dense1 = layers.Dense(128, activation='relu')\n        self.dropout = layers.Dropout(dropout_rate)\n        self.mean_output = layers.Dense(output_shape, name='mean_output')\n        self.var_output = layers.Dense(output_shape, activation='softplus', name='var_output')\n        self.phys_loss = PhysicsLossLayer(wavelengths, phys_reg_weight, output_shape)\n\n    def call(self, inputs: tf.Tensor, training: bool = False) -> dict:\n        x = tf.expand_dims(inputs, axis=-1)\n        x = self.conv1(x)\n        x = self.bn1(x, training=training)\n        x = self.conv2(x)\n        x = self.bn2(x, training=training)\n        x = self.attention(x)\n        x = self.flatten(x)\n        x = self.dense1(x)\n        x = self.dropout(x, training=training)\n        mean = self.mean_output(x)\n        var = self.var_output(x)\n        phys = self.phys_loss(mean)\n        return {'mean_output': mean, 'var_output': var, 'phys_loss': phys}\n\n# Data Pipeline\nclass ArielDataPipeline:\n    \"\"\"Handles data loading and preprocessing for all planets.\"\"\"\n    def __init__(self, cfg: Config):\n        self.cfg = cfg\n        try:\n            os.makedirs(cfg.cache_dir, exist_ok=True)\n            os.makedirs(cfg.checkpoint_dir, exist_ok=True)\n        except Exception as e:\n            logger.error(f\"Failed to create directories {cfg.cache_dir} or {cfg.checkpoint_dir}: {e}\")\n            raise\n\n    def _load_observation(self, pid: str, is_train: bool = True) -> np.ndarray:\n        \"\"\"Load and average observations for a planet.\"\"\"\n        pattern = 'AIRS-CH0_signal_{}.parquet'\n        obs_dir = os.path.join(self.cfg.data_path, self.cfg.train_dir if is_train else self.cfg.test_dir, pid)\n        obs = []\n        n_outputs = self.cfg.n_outputs if self.cfg.n_outputs > 0 else self.cfg.default_n_outputs\n        for i in range(2):\n            obs_file = os.path.join(obs_dir, pattern.format(i))\n            if not os.path.exists(obs_file):\n                logger.warning(f\"File not found: {obs_file}, using zeros\")\n                obs.append(np.zeros(n_outputs, dtype=np.float32))\n                continue\n            try:\n                obs_df = pd.read_parquet(obs_file)\n                arr = obs_df.values.astype(np.float32).mean(axis=0)[:n_outputs]\n                if np.any(np.isnan(arr)) or np.any(arr < 0):\n                    logger.warning(f\"Invalid data in {obs_file}: NaNs or negative values\")\n                    mask = np.isnan(arr) | (arr < 0)\n                    if np.all(mask):\n                        logger.warning(f\"All values invalid in {obs_file}, using zeros\")\n                        arr = np.zeros_like(arr)\n                    else:\n                        arr[mask] = np.interp(np.where(mask)[0], np.where(~mask)[0], arr[~mask])\n                obs.append(arr)\n            except Exception as e:\n                logger.error(f\"Error loading {obs_file}: {e}\")\n                obs.append(np.zeros(n_outputs, dtype=np.float32))\n        return np.mean(obs, axis=0)\n\n    def _worker(self, args: Tuple[str, bool]) -> np.ndarray:\n        \"\"\"Worker function for parallel processing.\"\"\"\n        pid, is_train = args\n        return self._load_observation(pid, is_train)\n\n    def _prepare_cache(self, planet_ids: List[str], is_train: bool = True) -> np.ndarray:\n        \"\"\"Cache observations in parallel with chunked processing.\"\"\"\n        cache_file = os.path.join(self.cfg.cache_dir, f'cached_obs_{\"train\" if is_train else \"test\"}.npy')\n        if os.path.exists(cache_file):\n            logger.info(f\"Loading cached data from {cache_file}\")\n            try:\n                cached_obs = np.load(cache_file)\n                n_outputs = self.cfg.n_outputs if self.cfg.n_outputs > 0 else self.cfg.default_n_outputs\n                if cached_obs.shape[0] == len(planet_ids) and cached_obs.shape[1] == n_outputs:\n                    return cached_obs\n                logger.warning(f\"Cached data shape mismatch, regenerating cache: {cache_file}\")\n            except Exception as e:\n                logger.warning(f\"Failed to load cache {cache_file}: {e}, regenerating\")\n\n        logger.info(f\"Caching observations for {len(planet_ids)} planets...\")\n        cached_obs = []\n        for chunk_start in range(0, len(planet_ids), self.cfg.chunk_size):\n            chunk_end = min(chunk_start + self.cfg.chunk_size, len(planet_ids))\n            chunk_ids = planet_ids[chunk_start:chunk_end]\n            logger.info(f\"Processing chunk {chunk_start // self.cfg.chunk_size + 1}/{(len(planet_ids) - 1) // self.cfg.chunk_size + 1}\")\n            with Pool(processes=cpu_count()) as pool:\n                chunk_obs = np.array(list(tqdm.tqdm(\n                    pool.imap(self._worker, [(pid, is_train) for pid in chunk_ids]),\n                    total=len(chunk_ids),\n                    desc=f\"Caching chunk {chunk_start // self.cfg.chunk_size + 1}\"\n                )), dtype=np.float32)\n            cached_obs.append(chunk_obs)\n        cached_obs = np.concatenate(cached_obs, axis=0)\n        try:\n            np.save(cache_file, cached_obs)\n            logger.info(f\"Saved cached data to {cache_file}\")\n        except Exception as e:\n            logger.error(f\"Failed to save cache: {e}\")\n        return cached_obs\n\n    def prepare_datasets(self) -> Tuple[tf.data.Dataset, tf.data.Dataset, np.ndarray, List[str]]:\n        \"\"\"Prepare datasets for training and testing.\"\"\"\n        # Validate file paths\n        for file_path in [\n            os.path.join(self.cfg.data_path, self.cfg.train_spec_file),\n            os.path.join(self.cfg.data_path, self.cfg.wavelengths_file),\n            os.path.join(self.cfg.data_path, self.cfg.test_ids_file)\n        ]:\n            if not os.path.exists(file_path):\n                logger.error(f\"Required file not found: {file_path}\")\n                raise FileNotFoundError(f\"Required file not found: {file_path}\")\n\n        # Load wavelengths\n        logger.info(\"Loading wavelengths...\")\n        try:\n            wavelengths_df = pd.read_csv(os.path.join(self.cfg.data_path, self.cfg.wavelengths_file))\n            logger.info(f\"Loaded wavelengths_df with shape: {wavelengths_df.shape}\")\n            logger.info(f\"Wavelengths_df columns: {wavelengths_df.columns.tolist()}\")\n            logger.info(f\"First 5 rows of wavelengths_df:\\n{wavelengths_df.head().to_string()}\")\n            self.cfg.n_outputs = len(wavelengths_df)\n            if self.cfg.n_outputs < 3:\n                logger.warning(f\"Number of wavelengths ({self.cfg.n_outputs}) is too small. Using default n_outputs={self.cfg.default_n_outputs}.\")\n                self.cfg.n_outputs = self.cfg.default_n_outputs\n                wavelengths = np.linspace(0.5, 5.0, self.cfg.n_outputs).astype(np.float32)\n            else:\n                wavelengths = wavelengths_df.values.flatten().astype(np.float32)\n            logger.info(f\"Set n_outputs to {self.cfg.n_outputs}\")\n        except Exception as e:\n            logger.error(f\"Failed to load wavelengths: {e}. Using default n_outputs={self.cfg.default_n_outputs}.\")\n            self.cfg.n_outputs = self.cfg.default_n_outputs\n            wavelengths = np.linspace(0.5, 5.0, self.cfg.n_outputs).astype(np.float32)\n\n        # Load training data\n        logger.info(\"Loading training data...\")\n        try:\n            train_spec_df = pd.read_csv(os.path.join(self.cfg.data_path, self.cfg.train_spec_file))\n            train_ids = train_spec_df['planet_id'].astype(str).values\n            train_spec = train_spec_df.filter(regex='wl_').values[:, :self.cfg.n_outputs].astype(np.float32)\n            logger.info(f\"Loaded {len(train_ids)} training samples\")\n        except Exception as e:\n            logger.error(f\"Failed to load training data: {e}\")\n            raise\n\n        if len(train_ids) != len(train_spec):\n            raise ValueError(f\"Train IDs ({len(train_ids)}) and train spectra ({len(train_spec)}) mismatch\")\n        if train_spec.shape[1] != self.cfg.n_outputs:\n            raise ValueError(f\"Train spectra columns ({train_spec.shape[1]}) mismatch with n_outputs ({self.cfg.n_outputs})\")\n\n        # Load test data\n        logger.info(\"Loading test data...\")\n        try:\n            test_ids_df = pd.read_csv(os.path.join(self.cfg.data_path, self.cfg.test_ids_file))\n            test_ids = test_ids_df['planet_id'].astype(str).values\n            logger.info(f\"Loaded {len(test_ids)} test samples\")\n        except Exception as e:\n            logger.error(f\"Failed to load test IDs: {e}\")\n            raise\n\n        # Cache observations\n        logger.info(\"Caching training observations...\")\n        train_obs = self._prepare_cache(train_ids, is_train=True)\n        logger.info(\"Caching test observations...\")\n        test_obs = self._prepare_cache(test_ids, is_train=False)\n\n        # Validate shapes\n        if train_obs.shape[0] != len(train_ids) or train_obs.shape[1] != self.cfg.n_outputs:\n            raise ValueError(f\"Train observations shape {train_obs.shape} mismatch with expected ({len(train_ids)}, {self.cfg.n_outputs})\")\n        if test_obs.shape[0] != len(test_ids) or test_obs.shape[1] != self.cfg.n_outputs:\n            raise ValueError(f\"Test observations shape {test_obs.shape} mismatch with expected ({len(test_ids)}, {self.cfg.n_outputs})\")\n\n        # Create datasets with dictionary output for y_true\n        ds_train = tf.data.Dataset.from_tensor_slices((\n            train_obs,\n            {'mean_output': train_spec, 'var_output': tf.zeros_like(train_spec), 'phys_loss': tf.zeros((len(train_ids),))}\n        )).batch(self.cfg.batch_size).prefetch(tf.data.AUTOTUNE)\n        ds_test = tf.data.Dataset.from_tensor_slices(test_obs).batch(self.cfg.batch_size).prefetch(tf.data.AUTOTUNE)\n        logger.info(f\"Prepared datasets: {len(train_ids)} train samples, {len(test_ids)} test samples\")\n        return ds_train, ds_test, wavelengths, test_ids\n\n# Training and Prediction\ndef train_and_predict(\n    cfg: Config,\n    ds_train: tf.data.Dataset,\n    ds_test: tf.data.Dataset,\n    wavelengths: np.ndarray,\n    test_ids: List[str]\n) -> Tuple[keras.Model, pd.DataFrame]:\n    \"\"\"Train the model and generate predictions.\"\"\"\n    logger.info(\"Initializing model...\")\n    try:\n        model = ArielModel(\n            input_shape=(cfg.n_outputs,),\n            output_shape=cfg.n_outputs,\n            wavelengths=wavelengths,\n            phys_reg_weight=0.5,\n            dropout_rate=0.3\n        )\n        model.compile(\n            optimizer=keras.optimizers.AdamW(learning_rate=1e-3),\n            loss={\n                'mean_output': 'mse',\n                'var_output': 'mse',\n                'phys_loss': lambda y_true, y_pred: y_pred\n            },\n            loss_weights={'mean_output': 1.0, 'var_output': 0.1, 'phys_loss': 0.5},\n            metrics={'mean_output': 'mse', 'var_output': 'mse'}\n        )\n    except Exception as e:\n        logger.error(f\"Model initialization failed: {e}\")\n        raise\n\n    logger.info(\"Starting training...\")\n    try:\n        model.fit(\n            ds_train,\n            epochs=cfg.epochs,\n            verbose=1,\n            callbacks=[\n                keras.callbacks.ModelCheckpoint(\n                    f'{cfg.checkpoint_dir}/model.h5',\n                    save_best_only=True,\n                    monitor='loss',\n                    mode='min'\n                ),\n                keras.callbacks.EarlyStopping(\n                    monitor='loss',\n                    patience=5,\n                    restore_best_weights=True\n                )\n            ]\n        )\n    except Exception as e:\n        logger.error(f\"Training failed: {e}\")\n        raise\n\n    logger.info(\"Generating predictions with MC-Dropout...\")\n    mus, stds = [], []\n    try:\n        for x in tqdm.tqdm(ds_test, desc=\"Test batches\"):\n            preds = [model(x, training=True)['mean_output'] for _ in range(cfg.mc_dropout_samples)]\n            preds = np.array([p.numpy() for p in preds])\n            mus.append(np.mean(preds, axis=0))\n            stds.append(np.std(preds, axis=0))\n        mu_ens = np.concatenate(mus, axis=0)\n        sig_ens = np.concatenate(stds, axis=0)\n    except Exception as e:\n        logger.error(f\"Prediction failed: {e}\")\n        raise\n\n    # Validate predictions\n    if mu_ens.shape != (len(test_ids), cfg.n_outputs) or sig_ens.shape != (len(test_ids), cfg.n_outputs):\n        logger.error(f\"Prediction shape mismatch: expected ({len(test_ids)}, {cfg.n_outputs}), got mu_ens {mu_ens.shape}, sig_ens {sig_ens.shape}\")\n        raise ValueError(f\"Prediction shape mismatch\")\n    if np.all(mu_ens == 0) or np.all(sig_ens == 0):\n        logger.error(f\"All prediction values are zero: mu_ens mean={np.mean(mu_ens)}, sig_ens mean={np.mean(sig_ens)}\")\n        raise ValueError(\"Prediction values are all zero\")\n    if np.any(np.isnan(mu_ens)) or np.any(np.isnan(sig_ens)):\n        logger.warning(\"NaN values in predictions, replacing with zeros\")\n        mu_ens = np.nan_to_num(mu_ens, nan=0.0)\n        sig_ens = np.nan_to_num(sig_ens, nan=0.0)\n    logger.info(f\"Prediction stats - mu_ens mean: {np.mean(mu_ens):.4f}, sig_ens mean: {np.mean(sig_ens):.4f}\")\n\n    # Create submission DataFrame\n    logger.info(\"Creating submission DataFrame...\")\n    try:\n        sample_submission = pd.read_csv(os.path.join(cfg.data_path, cfg.test_ids_file))\n        expected_columns = sample_submission.columns.tolist()\n        logger.info(f\"Expected columns: {expected_columns}\")\n\n        if len(expected_columns) != 1 + 2 * cfg.n_outputs:\n            logger.error(f\"Expected {1 + 2 * cfg.n_outputs} columns (1 planet_id + {2 * cfg.n_outputs} predictions), got {len(expected_columns)}\")\n            raise ValueError(f\"Column count mismatch in sample_submission.csv\")\n\n        prediction_cols = expected_columns[1:]\n        logger.info(f\"Using prediction columns: {prediction_cols[:5]}... (total {len(prediction_cols)})\")\n\n        df_data = np.hstack((np.array(test_ids).reshape(-1, 1), mu_ens, sig_ens))\n        if df_data.shape[1] != len(expected_columns):\n            logger.error(f\"Data shape {df_data.shape} does not match expected columns {len(expected_columns)}\")\n            raise ValueError(f\"Data column count {df_data.shape[1]} does not match {len(expected_columns)}\")\n\n        df = pd.DataFrame(df_data, columns=expected_columns)\n        df['planet_id'] = df['planet_id'].astype(str)\n        for col in prediction_cols:\n            df[col] = df[col].astype(float)\n\n        if list(df.columns) != expected_columns:\n            logger.error(f\"Submission columns mismatch: expected {expected_columns[:5]}..., got {list(df.columns)[:5]}...\")\n            raise ValueError(f\"Submission columns mismatch\")\n        if df.iloc[0].isna().all():\n            logger.error(\"All values in DataFrame are NaN\")\n            raise ValueError(\"DataFrame contains only NaN values\")\n\n        logger.info(f\"Sample values - First row {prediction_cols[0]}: {df.iloc[0][prediction_cols[0]]:.4f}, {prediction_cols[cfg.n_outputs]}: {df.iloc[0][prediction_cols[cfg.n_outputs]]:.4f}\")\n\n    except Exception as e:\n        logger.error(f\"Failed to create submission DataFrame: {e}\")\n        raise\n\n    # Save submission\n    try:\n        df.to_csv('/kaggle/working/submission.csv', index=False)\n        logger.info(\"Submission saved to /kaggle/working/submission.csv\")\n        logger.info(f\"Submission preview (first 5 rows):\\n{df.head().to_string()}\")\n    except Exception as e:\n        logger.error(f\"Failed to save submission: {e}\")\n        raise\n\n    return model, df\n\n# Main Execution\ndef main():\n    \"\"\"Execute the pipeline for all planets.\"\"\"\n    logger.info(\"Starting execution for all planets...\")\n    tf.keras.utils.set_random_seed(42)\n    os.environ['TF_ENABLE_AUTO_MIXED_PRECISION'] = '1'\n    cfg = Config()\n    data_pipeline = ArielDataPipeline(cfg)\n\n    try:\n        ds_train, ds_test, wavelengths, test_ids = data_pipeline.prepare_datasets()\n        logger.info(f\"Model will be initialized with n_outputs={cfg.n_outputs}\")\n        model, submission_df = train_and_predict(cfg, ds_train, ds_test, wavelengths, test_ids)\n        logger.info(\"Execution completed successfully!\")\n    except Exception as e:\n        logger.error(f\"Execution failed: {e}\")\n        raise\n\nif __name__ == '__main__':\n    main()","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}