{"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":"none","dataSources":[{"sourceId":101849,"databundleVersionId":13093295,"isSourceIdPinned":false,"sourceType":"competition"},{"sourceId":9629432,"sourceType":"datasetVersion","datasetId":5846888}],"dockerImageVersionId":31192,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-11-10T10:10:35.412746Z","iopub.execute_input":"2025-11-10T10:10:35.413062Z","iopub.status.idle":"2025-11-10T10:10:59.316213Z","shell.execute_reply.started":"2025-11-10T10:10:35.413015Z","shell.execute_reply":"2025-11-10T10:10:59.315299Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# install pqdm for parallel processing\n!pip install --no-index --find-links=/kaggle/input/ariel-2024-pqdm pqdm\n\nimport pandas as pd\nimport numpy as np\nfrom tqdm import tqdm\nfrom pqdm.threads import pqdm\nfrom scipy.optimize import minimize\nfrom scipy.signal import savgol_filter\nfrom astropy.stats import sigma_clip\nimport time\nfrom dataclasses import dataclass, field\nfrom typing import Dict, Tuple, List, Optional\nimport logging\n\n# Set up logging\nlogging.basicConfig(level=logging.INFO)\nlogger = logging.getLogger(__name__)\n\n@dataclass\nclass Config:\n    \"\"\"Configuration class for the Ariel data processing pipeline.\"\"\"\n    DATA_PATH: str = '/kaggle/input/ariel-data-challenge-2025'\n    DATASET: str = \"test\"\n    SCALE: float = 0.95\n    SIGMA: float = 0.0009\n    CUT_INF: int = 39\n    CUT_SUP: int = 321\n    MODEL_PHASE_DETECTION_SLICE: slice = field(default_factory=lambda: slice(30, 140))\n    MODEL_OPTIMIZATION_DELTA: int = 11\n    MODEL_POLYNOMIAL_DEGREE: int = 3\n    N_JOBS: int = 3\n    LOG_HOT_STATS: bool = False\n    \n    # Sensor configurations\n    SENSOR_CONFIG: Dict = field(default_factory=lambda: {\n        \"AIRS-CH0\": {\n            \"raw_shape\": [11250, 32, 356],\n            \"calibrated_shape\": [1, 32, 321 - 39],  # CUT_SUP - CUT_INF\n            \"linear_corr_shape\": (6, 32, 356),\n            \"dt_pattern\": (0.1, 4.5), \n            \"binning\": 30,\n            \"roi_y\": (10, 22),\n            \"roi_x\": None\n        },\n        \"FGS1\": {\n            \"raw_shape\": [135000, 32, 32],\n            \"calibrated_shape\": [1, 32, 32],\n            \"linear_corr_shape\": (6, 32, 32),\n            \"dt_pattern\": (0.1, 0.1),\n            \"binning\": 30 * 12,\n            \"roi_y\": (10, 22),\n            \"roi_x\": (10, 22)\n        }\n    })\n\n\nclass SignalProcessor:\n    \"\"\"Processes sensor signals for multiple planets.\"\"\"\n    \n    def __init__(self, config: Config):\n        self.cfg = config\n        self.adc_info = pd.read_csv(f\"{self.cfg.DATA_PATH}/adc_info.csv\")\n        self.planet_ids = pd.read_csv(\n            f'{self.cfg.DATA_PATH}/{self.cfg.DATASET}_star_info.csv', \n            index_col='planet_id'\n        ).index.astype(int)\n        self.stats = []  # For logging statistics\n\n    def _apply_linear_corr(self, linear_corr: np.ndarray, signal: np.ndarray) -> np.ndarray:\n        \"\"\"\n        Apply linearity correction to the signal.\n        \n        Args:\n            linear_corr: Coefficients for linearity correction\n            signal: Input signal to correct\n            \n        Returns:\n            Corrected signal\n        \"\"\"\n        coeffs = np.flip(linear_corr, axis=0)  # shape: (D, X, Y), D - highest degree first\n        x = signal.astype(np.float64, copy=False)\n        out = np.empty_like(x, dtype=np.float64)\n        out[...] = coeffs[0]  # broadcast (X,Y) -> (T,X,Y)\n        \n        for k in range(1, coeffs.shape[0]):\n            np.multiply(out, x, out=out)  # in-place multiplication\n            out += coeffs[k]  # broadcast (X,Y)\n\n        return out.astype(signal.dtype, copy=False)\n\n    def _calibrate_single_signal(self, planet_id: int, sensor: str) -> np.ndarray:\n        \"\"\"\n        Calibrate a single sensor signal for a planet.\n        \n        Args:\n            planet_id: ID of the planet to process\n            sensor: Sensor name ('AIRS-CH0' or 'FGS1')\n            \n        Returns:\n            Calibrated signal\n        \"\"\"\n        sensor_cfg = self.cfg.SENSOR_CONFIG[sensor]\n        \n        # Load data\n        signal = pd.read_parquet(\n            f\"{self.cfg.DATA_PATH}/{self.cfg.DATASET}/{planet_id}/{sensor}_signal_0.parquet\"\n        ).to_numpy()\n        \n        dark = pd.read_parquet(\n            f\"{self.cfg.DATA_PATH}/{self.cfg.DATASET}/{planet_id}/{sensor}_calibration_0/dark.parquet\"\n        ).to_numpy()\n        \n        dead = pd.read_parquet(\n            f\"{self.cfg.DATA_PATH}/{self.cfg.DATASET}/{planet_id}/{sensor}_calibration_0/dead.parquet\"\n        ).to_numpy()\n        \n        flat = pd.read_parquet(\n            f\"{self.cfg.DATA_PATH}/{self.cfg.DATASET}/{planet_id}/{sensor}_calibration_0/flat.parquet\"\n        ).to_numpy()\n        \n        linear_corr = pd.read_parquet(\n            f\"{self.cfg.DATA_PATH}/{self.cfg.DATASET}/{planet_id}/{sensor}_calibration_0/linear_corr.parquet\"\n        ).values.astype(np.float64).reshape(sensor_cfg[\"linear_corr_shape\"])\n\n        # Reshape & apply ADC correction\n        signal = signal.reshape(sensor_cfg[\"raw_shape\"])\n        gain = self.adc_info[f\"{sensor}_adc_gain\"].iloc[0]\n        offset = self.adc_info[f\"{sensor}_adc_offset\"].iloc[0]\n        signal = signal / gain + offset\n\n        # Identify hot pixels (for monitoring only, not masking)\n        hot = sigma_clip(dark, sigma=5, maxiters=5).mask\n\n        # Apply sensor-specific cropping\n        if sensor == \"AIRS-CH0\":\n            signal = signal[:, :, self.cfg.CUT_INF : self.cfg.CUT_SUP]\n            linear_corr = linear_corr[:, :, self.cfg.CUT_INF : self.cfg.CUT_SUP]\n            dark = dark[:, self.cfg.CUT_INF : self.cfg.CUT_SUP]\n            dead = dead[:, self.cfg.CUT_INF : self.cfg.CUT_SUP]\n            flat = flat[:, self.cfg.CUT_INF : self.cfg.CUT_SUP]\n            hot = hot[:, self.cfg.CUT_INF : self.cfg.CUT_SUP]\n\n        elif sensor == \"FGS1\":\n            y0, y1 = sensor_cfg[\"roi_y\"]\n            x0, x1 = sensor_cfg[\"roi_x\"]\n            signal = signal[:, y0:y1, x0:x1]\n            dark = dark[y0:y1, x0:x1]\n            dead = dead[y0:y1, x0:x1]\n            flat = flat[y0:y1, x0:x1]\n            linear_corr = linear_corr[:, y0:y1, x0:x1]\n            hot = hot[y0:y1, x0:x1]\n\n        # Non-negative clamp before linearity correction\n        np.maximum(signal, 0, out=signal)\n\n        # Apply linearity correction\n        if sensor == \"FGS1\":\n            signal = self._apply_linear_corr(linear_corr, signal)\n        elif sensor == \"AIRS-CH0\":\n            y0, y1 = sensor_cfg[\"roi_y\"]\n            sl = (slice(None), slice(y0, y1), slice(None))  # T, Y, λ\n            signal[sl] = self._apply_linear_corr(linear_corr[:, y0:y1, :], signal[sl])\n\n        # Dark subtraction with integration pattern consideration\n        base_dt, increment = sensor_cfg[\"dt_pattern\"]\n        even_scale = base_dt\n        odd_scale = base_dt + increment\n        signal[::2] -= dark * even_scale\n        signal[1::2] -= dark * odd_scale\n\n        # Apply flat field correction (excluding hot pixels from mask)\n        if sensor == \"FGS1\":\n            flat_roi = flat.astype(signal.dtype, copy=False).copy()\n            bad = dead | ~np.isfinite(flat_roi) | (flat_roi == 0)\n            flat_roi[bad] = np.nan\n            signal /= flat_roi\n\n        elif sensor == \"AIRS-CH0\":\n            y0, y1 = sensor_cfg[\"roi_y\"]\n            flat_roi = flat[y0:y1, :].astype(signal.dtype, copy=False).copy()\n            bad = dead[y0:y1, :] | ~np.isfinite(flat_roi) | (flat_roi == 0)\n            flat_roi[bad] = np.nan\n            signal[:, y0:y1, :] /= flat_roi\n\n        # Log statistics if enabled\n        if self.cfg.LOG_HOT_STATS:\n            self.stats.append({\n                \"planet_id\": int(planet_id),\n                \"sensor\": sensor,\n                \"hot_frac\": float(np.mean(hot)),\n                \"dead_frac\": float(np.mean(dead)),\n            })\n\n        return signal\n\n    def _preprocess_calibrated_signal(self, calibrated_signal: np.ndarray, sensor: str) -> np.ndarray:\n        \"\"\"\n        Preprocess calibrated signal by binning and applying weights.\n        \n        Args:\n            calibrated_signal: Calibrated signal from sensor\n            sensor: Sensor name\n            \n        Returns:\n            Preprocessed signal\n        \"\"\"\n        sensor_cfg = self.cfg.SENSOR_CONFIG[sensor]\n        binning = sensor_cfg[\"binning\"]\n\n        # Extract region of interest\n        if sensor == \"AIRS-CH0\":\n            y0, y1 = sensor_cfg[\"roi_y\"]\n            signal_roi = calibrated_signal[:, y0:y1, :]\n        elif sensor == \"FGS1\":\n            y0, y1 = sensor_cfg[\"roi_y\"]\n            x0, x1 = sensor_cfg[\"roi_x\"]\n            signal_roi = calibrated_signal[:, y0:y1, x0:x1]\n            signal_roi = signal_roi.reshape(signal_roi.shape[0], -1)\n\n        # Calculate mean signal and CDS\n        mean_signal = np.nanmean(signal_roi, axis=1)\n        cds_signal = mean_signal[1::2] + mean_signal[0::2]\n\n        # Bin the signal\n        n_bins = cds_signal.shape[0] // binning\n        binned = np.array([\n            cds_signal[j*binning : (j+1)*binning].mean(axis=0) \n            for j in range(n_bins)\n        ])\n\n        # Apply winsorization for AIRS\n        if sensor == \"AIRS-CH0\":\n            q_lo = np.nanpercentile(binned, 5.0, axis=1, keepdims=True)\n            q_hi = np.nanpercentile(binned, 95.0, axis=1, keepdims=True)\n            np.clip(binned, q_lo, q_hi, out=binned)\n\n        # Reshape FGS1 data\n        if sensor == \"FGS1\":\n            binned = binned.reshape((binned.shape[0], 1))\n\n        # Apply inverse variance weighting for AIRS\n        if sensor == \"AIRS-CH0\":\n            var = np.nanvar(binned, axis=0, ddof=1)\n            med = np.nanmedian(var)\n            safe_var = np.where(~np.isfinite(var) | (var <= 0), med if (np.isfinite(med) and med > 0) else 1.0, var)\n            w = 1.0 / safe_var\n\n            # Clip weights to prevent dominance by single channel\n            lo, hi = np.nanpercentile(w, 5.0), np.nanpercentile(w, 95.0)\n            if np.isfinite(lo) and np.isfinite(hi) and lo < hi:\n                w = np.clip(w, lo, hi)\n\n            # Normalize weights\n            M = binned.shape[1]\n            s = np.nansum(w)\n            if np.isfinite(s) and s > 0:\n                w = w * (M / s)\n            else:\n                w = np.ones_like(w)\n\n            # Apply weights\n            binned *= w[None, :]\n\n        return binned\n\n    def _process_planet_sensor(self, args: dict) -> np.ndarray:\n        \"\"\"Wrapper function for parallel processing.\"\"\"\n        planet_id, sensor = args['planet_id'], args['sensor']\n        calibrated = self._calibrate_single_signal(planet_id, sensor)\n        preprocessed = self._preprocess_calibrated_signal(calibrated, sensor)\n        return preprocessed\n\n    def process_all_data(self) -> np.ndarray:\n        \"\"\"Process all data for all planets and sensors.\"\"\"\n        logger.info(\"Processing FGS1 data...\")\n        args_fgs1 = [{\"planet_id\": planet_id, \"sensor\": \"FGS1\"} for planet_id in self.planet_ids]\n        preprocessed_fgs1 = pqdm(args_fgs1, self._process_planet_sensor, n_jobs=self.cfg.N_JOBS, desc=\"FGS1 Processing\")\n\n        logger.info(\"Processing AIRS-CH0 data...\")\n        args_airs_ch0 = [{\"planet_id\": planet_id, \"sensor\": \"AIRS-CH0\"} for planet_id in self.planet_ids]\n        preprocessed_airs_ch0 = pqdm(args_airs_ch0, self._process_planet_sensor, n_jobs=self.cfg.N_JOBS, desc=\"AIRS Processing\")\n\n        # Combine processed data\n        preprocessed_signal = np.concatenate(\n            [np.stack(preprocessed_fgs1), np.stack(preprocessed_airs_ch0)], axis=2\n        )\n        \n        return preprocessed_signal\n\n\nclass TransitModel:\n    \"\"\"Models transit signals to estimate depth.\"\"\"\n    \n    def __init__(self, config: Config):\n        self.cfg = config\n\n    def _phase_detector(self, signal: np.ndarray) -> Tuple[int, int]:\n        \"\"\"\n        Detect transit phases in the signal.\n        \n        Args:\n            signal: Input signal\n            \n        Returns:\n            Start and end indices of the transit\n        \"\"\"\n        search_slice = self.cfg.MODEL_PHASE_DETECTION_SLICE\n        min_index = np.argmin(signal[search_slice]) + search_slice.start\n        \n        signal1 = signal[:min_index]\n        signal2 = signal[min_index:]\n\n        grad1 = np.gradient(signal1)\n        if grad1.max() != 0:\n            grad1 /= grad1.max()\n        \n        grad2 = np.gradient(signal2)\n        if grad2.max() != 0:\n            grad2 /= grad2.max()\n\n        phase1 = np.argmin(grad1)\n        phase2 = np.argmax(grad2) + min_index\n\n        return phase1, phase2\n    \n    def _objective_function(self, s: float, signal: np.ndarray, phase1: int, phase2: int) -> float:\n        \"\"\"\n        Objective function for transit depth optimization.\n        \n        Args:\n            s: Transit depth parameter\n            signal: Input signal\n            phase1: Start of transit\n            phase2: End of transit\n            \n        Returns:\n            Error metric\n        \"\"\"\n        delta = self.cfg.MODEL_OPTIMIZATION_DELTA\n        power = self.cfg.MODEL_POLYNOMIAL_DEGREE\n\n        # Adjust delta if phases are too close to boundaries\n        if phase1 - delta <= 0 or phase2 + delta >= len(signal) or phase2 - delta - (phase1 + delta) < 5:\n            delta = min(2, phase1, len(signal) - phase2 - 1)\n\n        # Create modified signal with transit applied\n        y = np.concatenate([\n            signal[: phase1 - delta],\n            signal[phase1 + delta : phase2 - delta] * (1 + s),\n            signal[phase2 + delta :]\n        ])\n        x = np.arange(len(y))\n\n        # Fit polynomial and calculate error\n        coeffs = np.polyfit(x, y, deg=power)\n        poly = np.poly1d(coeffs)\n        error = np.abs(poly(x) - y).mean()\n        \n        return error\n\n    def predict(self, single_preprocessed_signal: np.ndarray) -> float:\n        \"\"\"\n        Predict transit depth for a single signal.\n        \n        Args:\n            single_preprocessed_signal: Preprocessed signal\n            \n        Returns:\n            Estimated transit depth\n        \"\"\"\n        signal_1d = single_preprocessed_signal[:, 1:].mean(axis=1)\n        signal_1d = savgol_filter(signal_1d, 23, 2)\n        \n        phase1, phase2 = self._phase_detector(signal_1d)\n\n        # Ensure phases are within valid bounds\n        phase1 = max(self.cfg.MODEL_OPTIMIZATION_DELTA, phase1)\n        phase2 = min(len(signal_1d) - self.cfg.MODEL_OPTIMIZATION_DELTA - 1, phase2)\n\n        # Optimize transit depth\n        result = minimize(\n            fun=self._objective_function,\n            x0=[0.0001],\n            args=(signal_1d, phase1, phase2),\n            method=\"Nelder-Mead\",\n            options={'xatol': 1e-8, 'fatol': 1e-8}\n        )\n        \n        return result.x[0]\n\n    def predict_all(self, preprocessed_signals: List[np.ndarray]) -> np.ndarray:\n        \"\"\"\n        Predict transit depths for all signals.\n        \n        Args:\n            preprocessed_signals: List of preprocessed signals\n            \n        Returns:\n            Array of predicted transit depths\n        \"\"\"\n        predictions = [\n            self.predict(preprocessed_signal)\n            for preprocessed_signal in tqdm(preprocessed_signals, desc=\"Transit Depth Prediction\")\n        ]\n        return np.array(predictions) * self.cfg.SCALE\n\n\ndef estimate_sigma_fgs(preprocessed_data: List[np.ndarray], cfg: Config) -> np.ndarray:\n    \"\"\"\n    Estimate sigma values for FGS1 sensor.\n    \n    Args:\n        preprocessed_data: List of preprocessed signals\n        cfg: Configuration object\n        \n    Returns:\n        Array of sigma values for FGS1\n    \"\"\"\n    sig_rel = []\n    delta = cfg.MODEL_OPTIMIZATION_DELTA\n    eps = 1e-12\n    \n    for single in preprocessed_data:\n        # Use AIRS white curve for phase detection\n        air_white = savgol_filter(single[:, 1:].mean(axis=1), 20, 2)\n        p1, p2 = _phase_detector_signal(air_white, cfg)\n        p1 = max(delta, p1)\n        p2 = min(len(air_white) - delta - 1, p2)\n\n        fgs = single[:, 0]\n        oot = (fgs[: p1 - delta] if p1 - delta > 0 else np.empty(0, fgs.dtype))\n        if p2 + delta < fgs.size:\n            oot = np.concatenate([oot, fgs[p2 + delta :]])\n        inn = fgs[p1 + delta : max(p1 + delta, p2 - delta)]\n\n        if oot.size == 0 or inn.size == 0:\n            sig_rel.append(np.nan)\n            continue\n\n        n_oot, n_in = len(oot), len(inn)\n        var_oot = np.nanvar(oot, ddof=1)\n        var_in = np.nanvar(inn, ddof=1)\n        oot_mean = float(np.nanmean(oot)) if np.isfinite(np.nanmean(oot)) else float(np.nanmean(fgs))\n        \n        # Relative uncertainty of depth\n        sigma_rel = np.sqrt(var_oot / max(n_oot, 1) + var_in / max(n_in, 1)) / max(oot_mean, eps)\n        sig_rel.append(sigma_rel)\n\n    s = np.asarray(sig_rel, dtype=float)\n    mask = np.isfinite(s) & (s > 0)\n    med = float(np.nanmedian(s[mask])) if mask.any() else 1.0\n\n    # Soft multiplier with clipping\n    k = np.ones_like(s)\n    if med > 0 and np.isfinite(med):\n        k[mask] = np.sqrt(s[mask] / med)\n    k = np.clip(k, 0.8, 1.25)\n\n    return k * cfg.SIGMA\n\n\ndef estimate_sigma_air(preprocessed_data: List[np.ndarray], cfg: Config) -> np.ndarray:\n    \"\"\"\n    Estimate sigma values for AIRS sensor.\n    \n    Args:\n        preprocessed_data: List of preprocessed signals\n        cfg: Configuration object\n        \n    Returns:\n        Array of sigma values for AIRS\n    \"\"\"\n    sig_rel = []\n    delta = cfg.MODEL_OPTIMIZATION_DELTA\n    eps = 1e-12\n\n    for single in preprocessed_data:\n        # White curve from AIRS binned data\n        white = np.nanmean(single[:, 1:], axis=1)\n        white_s = savgol_filter(white, 20, 2)\n\n        p1, p2 = _phase_detector_signal(white_s, cfg)\n        p1 = max(delta, p1)\n        p2 = min(len(white) - delta - 1, p2)\n\n        oot_left = white[: p1 - delta] if p1 - delta > 0 else np.empty(0, white.dtype)\n        oot_right = white[p2 + delta :] if (p2 + delta) < white.size else np.empty(0, white.dtype)\n        oot = np.concatenate([oot_left, oot_right]) if (oot_left.size + oot_right.size) else oot_left\n        inn = white[p1 + delta : max(p1 + delta, p2 - delta)]\n\n        if oot.size == 0 or inn.size == 0:\n            sig_rel.append(np.nan)\n            continue\n\n        n_oot, n_in = len(oot), len(inn)\n        var_oot = np.nanvar(oot, ddof=1)\n        var_in = np.nanvar(inn, ddof=1)\n        oot_mean = float(np.nanmean(oot)) if np.isfinite(np.nanmean(oot)) else float(np.nanmean(white))\n\n        sigma_rel = np.sqrt(var_oot / max(n_oot, 1) + var_in / max(n_in, 1)) / max(oot_mean, eps)\n        sig_rel.append(sigma_rel)\n\n    s = np.asarray(sig_rel, dtype=float)\n    mask = np.isfinite(s) & (s > 0)\n    med = float(np.nanmedian(s[mask])) if mask.any() else 1.0\n\n    # Soft multiplier with clipping\n    k = np.ones_like(s)\n    if med > 0 and np.isfinite(med):\n        k[mask] = np.sqrt(s[mask] / med)\n    k = np.clip(k, 0.90, 1.20)\n\n    return k * cfg.SIGMA\n\n\ndef _phase_detector_signal(signal: np.ndarray, cfg: Config) -> Tuple[int, int]:\n    \"\"\"\n    Detect phases in a signal.\n    \n    Args:\n        signal: Input signal\n        cfg: Configuration object\n        \n    Returns:\n        Start and end indices of detected feature\n    \"\"\"\n    sl = cfg.MODEL_PHASE_DETECTION_SLICE\n    min_idx = int(np.argmin(signal[sl])) + sl.start\n    s1 = signal[:min_idx]\n    s2 = signal[min_idx:]\n    \n    if s1.size < 3 or s2.size < 3:\n        return 0, len(signal) - 1\n        \n    g1 = np.gradient(s1)\n    g2 = np.gradient(s2)\n    \n    g1_max = np.max(g1) if np.size(g1) else 0.0\n    g2_max = np.max(g2) if np.size(g2) else 0.0\n    \n    if g1_max != 0: \n        g1 /= g1_max\n    if g2_max != 0: \n        g2 /= g2_max\n        \n    phase1 = int(np.argmin(g1))\n    phase2 = int(np.argmax(g2)) + min_idx\n    \n    return phase1, phase2\n\n\nclass SubmissionGenerator:\n    \"\"\"Generates submission file from predictions.\"\"\"\n    \n    def __init__(self, config: Config):\n        self.cfg = config\n        self.sample_submission = pd.read_csv(\n            \"/kaggle/input/ariel-data-challenge-2025/sample_submission.csv\", \n            index_col=\"planet_id\"\n        )\n\n    def create(self, predictions: np.ndarray, sigma_fgs: Optional[np.ndarray] = None, \n               sigma_air: Optional[np.ndarray] = None) -> pd.DataFrame:\n        \"\"\"\n        Create submission DataFrame.\n        \n        Args:\n            predictions: Array of predicted transit depths\n            sigma_fgs: Optional array of sigma values for FGS1\n            sigma_air: Optional array of sigma values for AIRS\n            \n        Returns:\n            Submission DataFrame\n        \"\"\"\n        planet_ids = self.sample_submission.index\n        n_mu = self.sample_submission.shape[1] // 2  # 283\n\n        preds = np.asarray(predictions, dtype=float).reshape(-1)\n        mu = np.tile(preds.reshape(-1, 1), (1, n_mu))\n        mu = np.clip(mu, 0, None)\n\n        sigmas = np.full_like(mu, self.cfg.SIGMA, dtype=float)\n        \n        if sigma_fgs is not None:\n            sigma_fgs = np.asarray(sigma_fgs, dtype=float).reshape(-1)\n            sigmas[:, 0] = np.clip(sigma_fgs, 1e-6, 0.1)\n            \n        if sigma_air is not None:\n            sigma_air = np.asarray(sigma_air, dtype=float).reshape(-1, 1)\n            sigmas[:, 1:] = np.clip(sigma_air, 1e-6, 0.1)\n\n        submission_df = pd.DataFrame(\n            np.concatenate([mu, sigmas], axis=1),\n            columns=self.sample_submission.columns,\n            index=planet_ids\n        )\n        submission_df.to_csv(\"submission.csv\")\n        return submission_df\n\n\ndef main():\n    \"\"\"Main execution function.\"\"\"\n    __t0 = time.perf_counter()\n    \n    # Initialize configuration and components\n    config = Config()\n    signal_processor = SignalProcessor(config)\n    \n    # Process data\n    logger.info(\"Starting data processing...\")\n    preprocessed_data = signal_processor.process_all_data()\n    \n    # Model predictions\n    logger.info(\"Starting transit modeling...\")\n    model = TransitModel(config)\n    predictions = model.predict_all(preprocessed_data)\n    \n    # Estimate sigma values\n    logger.info(\"Estimating sigma values...\")\n    sigma_fgs_vec = estimate_sigma_fgs(preprocessed_data, config)\n    sigma_air_vec = estimate_sigma_air(preprocessed_data, config)\n    \n    # Generate submission\n    logger.info(\"Generating submission...\")\n    submission_generator = SubmissionGenerator(config)\n    submission = submission_generator.create(predictions, sigma_fgs=sigma_fgs_vec, sigma_air=sigma_air_vec)\n    \n    # Print timing information\n    __t1 = time.perf_counter()\n    elapsed = __t1 - __t0\n    logger.info(f\"Total runtime: {elapsed:.2f} s ({elapsed/60:.2f} min)\")\n    \n    return submission\n\n\nif __name__ == \"__main__\":\n    submission = main()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-10T10:11:11.040935Z","iopub.execute_input":"2025-11-10T10:11:11.041269Z","iopub.status.idle":"2025-11-10T10:11:21.396981Z","shell.execute_reply.started":"2025-11-10T10:11:11.041245Z","shell.execute_reply":"2025-11-10T10:11:21.396213Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import matplotlib.pyplot as plt\nimport pandas as pd\nimport numpy as np\nfrom scipy.signal import savgol_filter\nfrom scipy.optimize import minimize\nimport logging\nimport warnings\n\n# --- New Helper Functions for Plotting ---\n\ndef plot_heatmap(data: np.ndarray, title: str, xlabel: str, ylabel: str, aspect: str = 'auto', cmap: str = 'viridis'):\n    \"\"\"Plots a 2D heatmap of sensor data.\"\"\"\n    with warnings.catch_warnings():\n        warnings.simplefilter(\"ignore\", UserWarning) # Ignore \"matplotlib.pyplot.figure\" warning\n        plt.figure(figsize=(12, 6))\n        \n        # Calculate robust color limits using 1st and 99th percentiles\n        v_min, v_max = np.nanpercentile(data, [1, 99])\n        \n        plt.imshow(data, aspect=aspect, cmap=cmap, interpolation='nearest', vmin=v_min, vmax=v_max)\n        plt.colorbar(label='Signal / Weight')\n        plt.title(title, fontsize=16)\n        plt.xlabel(xlabel)\n        plt.ylabel(ylabel)\n        plt.tight_layout()\n        plt.show()\n\ndef plot_light_curve(time: np.ndarray, flux: np.ndarray, title: str, ylabel: str, color: str = 'blue'):\n    \"\"\"Plots a 1D light curve.\"\"\"\n    with warnings.catch_warnings():\n        warnings.simplefilter(\"ignore\", UserWarning)\n        plt.figure(figsize=(12, 5))\n        plt.plot(time, flux, 'o', markersize=3, alpha=0.6, color=color)\n        plt.title(title, fontsize=16)\n        plt.xlabel('Binned Time Step')\n        plt.ylabel(ylabel)\n        plt.grid(True, alpha=0.2)\n        plt.tight_layout()\n        plt.show()\n\ndef plot_weights(weights: np.ndarray, title: str):\n    \"\"\"Plots a bar chart of the AIRS channel weights.\"\"\"\n    with warnings.catch_warnings():\n        warnings.simplefilter(\"ignore\", UserWarning)\n        plt.figure(figsize=(12, 5))\n        plt.bar(np.arange(len(weights)), weights, color='cyan')\n        plt.title(title, fontsize=16)\n        plt.xlabel('AIRS Wavelength Channel Index')\n        plt.ylabel('Calculated Weight (Trust)')\n        plt.grid(True, alpha=0.2, axis='y')\n        plt.tight_layout()\n        plt.show()\n\ndef plot_final_transit_fit(model: TransitModel, signal_1d: np.ndarray, title: str):\n    \"\"\"Re-creates and plots the final transit model fit.\"\"\"\n    cfg = model.cfg\n    \n    # 1. Get the smoothed signal\n    signal_savgol = savgol_filter(signal_1d, 23, 2)\n    \n    # 2. Get the phases\n    phase1, phase2 = model._phase_detector(signal_savgol)\n    phase1 = max(cfg.MODEL_OPTIMIZATION_DELTA, phase1)\n    phase2 = min(len(signal_savgol) - cfg.MODEL_OPTIMIZATION_DELTA - 1, phase2)\n\n    # 3. Get the final optimized depth 's'\n    result = minimize(\n        fun=model._objective_function,\n        x0=[0.0001],\n        args=(signal_savgol, phase1, phase2),\n        method=\"Nelder-Mead\",\n        options={'xatol': 1e-8, 'fatol': 1e-8}\n    )\n    final_s = result.x[0]\n\n    # 4. Get the polynomial drift component\n    delta = cfg.MODEL_OPTIMIZATION_DELTA\n    if phase1 - delta <= 0 or phase2 + delta >= len(signal_savgol) or phase2 - delta - (phase1 + delta) < 5:\n        delta = min(2, phase1, len(signal_savgol) - phase2 - 1)\n        \n    y_for_poly = np.concatenate([\n        signal_savgol[: phase1 - delta],\n        signal_savgol[phase1 + delta : phase2 - delta] * (1 + final_s),\n        signal_savgol[phase2 + delta :]\n    ])\n    x_for_poly = np.arange(len(y_for_poly))\n    \n    coeffs = np.polyfit(x_for_poly, y_for_poly, deg=cfg.MODEL_POLYNOMIAL_DEGREE)\n    poly_drift = np.poly1d(coeffs)\n    \n    # 5. Create the full transit model\n    time_steps = np.arange(len(signal_savgol))\n    full_model_fit = poly_drift(time_steps)\n    full_model_fit[phase1 + delta : phase2 - delta] /= (1 + final_s) # Create the \"dip\" in the poly\n\n    # 6. Plot everything\n    with warnings.catch_warnings():\n        warnings.simplefilter(\"ignore\", UserWarning)\n        plt.figure(figsize=(14, 7))\n    \n    # Plot 1: The full, complex fit\n    plt.subplot(2, 1, 1)\n    plt.plot(time_steps, signal_savgol, 'o', color='gray', markersize=3, alpha=0.5, label='Smoothed Data')\n    plt.plot(time_steps, poly_drift(time_steps), 'r-', linewidth=2, label=f'Polynomial Drift (Deg {cfg.MODEL_POLYNOMIAL_DEGREE})')\n    plt.plot(time_steps, full_model_fit, 'c-', linewidth=3, label=f'Full Transit Model (s={final_s:.6f})')\n    plt.axvspan(phase1, phase2, color='blue', alpha=0.1, label='Detected Transit Phase')\n    plt.title(title, fontsize=16)\n    plt.ylabel('Flux')\n    plt.legend()\n    plt.grid(True, alpha=0.2)\n    \n    # Plot 2: The de-trended light curve\n    plt.subplot(2, 1, 2)\n    detrended_flux = signal_savgol / poly_drift(time_steps)\n    detrended_model = full_model_fit / poly_drift(time_steps)\n    \n    plt.plot(time_steps, detrended_flux, 'o', color='gray', markersize=3, alpha=0.7, label='De-trended Data')\n    plt.plot(time_steps, detrended_model, 'c-', linewidth=3, label='Transit Model Shape')\n    plt.axvspan(phase1, phase2, color='blue', alpha=0.1)\n    plt.title(\"De-trended Light Curve\", fontsize=16)\n    plt.xlabel('Binned Time Step')\n    plt.ylabel('Normalized Flux')\n    plt.legend()\n    plt.grid(True, alpha=0.2)\n    \n    plt.tight_layout()\n    plt.show()\n\ndef get_airs_weights(binned_data: np.ndarray) -> np.ndarray:\n    \"\"\"Helper to recalculate and return the AIRS weights.\"\"\"\n    var = np.nanvar(binned_data, axis=0, ddof=1)\n    med = np.nanmedian(var)\n    safe_var = np.where(~np.isfinite(var) | (var <= 0), med if (np.isfinite(med) and med > 0) else 1.0, var)\n    w = 1.0 / safe_var\n\n    lo, hi = np.nanpercentile(w, 5.0), np.nanpercentile(w, 95.0)\n    if np.isfinite(lo) and np.isfinite(hi) and lo < hi:\n        w = np.clip(w, lo, hi)\n\n    M = binned_data.shape[1]\n    s = np.nansum(w)\n    if np.isfinite(s) and s > 0:\n        w = w * (M / s)\n    else:\n        w = np.ones_like(w)\n    return w\n\n# --- Main Visualization Function ---\n\ndef visualize_full_pipeline(planet_id: int, config: Config):\n    \"\"\"\n    Runs the full processing pipeline for a SINGLE planet and plots\n    the results at each major step.\n    \"\"\"\n    try:\n        print(f\"\\n--- [VISUALIZER] Starting for Planet {planet_id} ---\")\n        logger.info(f\"--- Visualizing Pipeline for Planet {planet_id} ---\")\n        processor = SignalProcessor(config)\n        model = TransitModel(config)\n        \n        # --- (NEW) Step 0: Plot Calibration Files ---\n        print(\"[VISUALIZER] Step 0: Loading and plotting calibration files...\")\n        dark_airs = pd.read_parquet(f\"{config.DATA_PATH}/{config.DATASET}/{planet_id}/AIRS-CH0_calibration_0/dark.parquet\").to_numpy()\n        flat_airs = pd.read_parquet(f\"{config.DATA_PATH}/{config.DATASET}/{planet_id}/AIRS-CH0_calibration_0/flat.parquet\").to_numpy()\n        \n        # Crop them just as the pipeline does\n        dark_airs = dark_airs[:, config.CUT_INF:config.CUT_SUP]\n        flat_airs = flat_airs[:, config.CUT_INF:config.CUT_SUP]\n        \n        plot_heatmap(dark_airs, f\"Planet {planet_id} - AIRS Dark Frame (What We Subtract)\", \"Wavelength\", \"Y-Pixel\", aspect='auto', cmap='inferno')\n        plot_heatmap(flat_airs, f\"Planet {planet_id} - AIRS Flat Frame (What We Divide By)\", \"Wavelength\", \"Y-Pixel\", aspect='auto', cmap='viridis')\n        print(\"[VISUALIZER] Step 0 COMPLETE.\")\n\n\n        # --- Step 1: Calibrated Frame ---\n        print(\"\\n[VISUALIZER] Step 1: Calibrating AIRS signal...\")\n        logger.info(\"Visualizing Step 1: Calibration (AIRS-CH0)\")\n        calibrated_airs = processor._calibrate_single_signal(planet_id, \"AIRS-CH0\")\n        plot_heatmap(calibrated_airs[0], f\"Planet {planet_id} - Calibrated AIRS Frame (Time 0)\", \"Wavelength\", \"Y-Pixel\", aspect='auto')\n        print(\"[VISUALIZER] Step 1 COMPLETE.\")\n\n        # --- Step 2: Preprocessing & Light Curves ---\n        print(\"\\n[VISUALIZER] Step 2: Preprocessing and creating light curves...\")\n        logger.info(\"Visualizing Step 2: Preprocessing (FGS1 & AIRS-CH0)\")\n        \n        # Process both sensors\n        p_fgs1 = processor._process_planet_sensor({\"planet_id\": planet_id, \"sensor\": \"FGS1\"})\n        p_airs_ch0 = processor._process_planet_sensor({\"planet_id\": planet_id, \"sensor\": \"AIRS-CH0\"})\n        \n        # (NEW) Plot the FGS1 light curve\n        fgs_light_curve = p_fgs1.flatten() # It's just (N, 1)\n        plot_light_curve(\n            np.arange(len(fgs_light_curve)),\n            fgs_light_curve,\n            f\"Planet {planet_id} - Binned FGS1 Light Curve\",\n            \"Binned Flux\",\n            color='green'\n        )\n        \n        # (NEW) Plot the AIRS 2D Spectroscopic Light Curve\n        plot_heatmap(\n            p_airs_ch0,\n            f\"Planet {planet_id} - Binned AIRS 2D Light Curve (Spectroscopic)\",\n            \"Wavelength Channel\",\n            \"Binned Time Step\",\n            aspect='auto',\n            cmap='viridis'\n        )\n        \n        # (NEW) Plot the AIRS weights\n        # We need the *un-weighted* binned data to calculate the weights\n        unweighted_binned_airs = processor._preprocess_calibrated_signal(calibrated_airs, \"AIRS-CH0\")\n        airs_weights = get_airs_weights(unweighted_binned_airs)\n        plot_weights(airs_weights, f\"Planet {planet_id} - AIRS Channel Weights (Trust)\")\n        \n        print(\"[VISUALIZER] Step 2 COMPLETE.\")\n\n        # --- Step 3: Modeling ---\n        print(\"\\n[VISUALIZER] Step 3: Fitting the final transit model...\")\n        logger.info(\"Visualizing Step 3: Transit Model Fit\")\n        \n        # Create the combined signal as used in the main script\n        combined_signal = np.concatenate([p_fgs1, p_airs_ch0], axis=1)\n        signal_1d_for_model = combined_signal[:, 1:].mean(axis=1)\n        \n        plot_final_transit_fit(\n            model,\n            signal_1d_for_model,\n            f\"Planet {planet_id} - Final Transit Model Fit (on AIRS White Light)\"\n        )\n        print(\"[VISUALIZER] Step 3 COMPLETE.\")\n        \n        logger.info(f\"--- Visualization Complete for Planet {planet_id} ---\")\n        print(f\"\\n--- [VISUALIZER] Finished for Planet {planet_id} ---\")\n\n    except Exception as e:\n        print(f\"\\n--- [VISUALIZER] ERROR! ---\")\n        print(f\"An error occurred: {e}\")\n        logger.error(f\"Error during visualization: {e}\")\n        import traceback\n        traceback.print_exc()\n\n# --- Run the Visualization ---\n# (This assumes you have already run Block 1 so all classes are in memory)\ntry:\n    config = Config()\n    PLANET_TO_VISUALIZE = pd.read_csv(\n        f'{config.DATA_PATH}/{config.DATASET}_star_info.csv'\n    ).planet_id.iloc[0]\n    \n    visualize_full_pipeline(int(PLANET_TO_VISUALIZE), config)\nexcept Exception as e:\n    logger.error(f\"Failed to start visualization: {e}\")\n    print(f\"Failed to start visualization. Make sure Block 1 has been run. Error: {e}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-10T10:23:04.705556Z","iopub.execute_input":"2025-11-10T10:23:04.706262Z","iopub.status.idle":"2025-11-10T10:23:15.098555Z","shell.execute_reply.started":"2025-11-10T10:23:04.706233Z","shell.execute_reply":"2025-11-10T10:23:15.09791Z"}},"outputs":[],"execution_count":null}]}