{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":87793,"databundleVersionId":11553390,"sourceType":"competition"},{"sourceId":11068753,"sourceType":"datasetVersion","datasetId":6896780}],"dockerImageVersionId":30918,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# RNA 3D Folding Data Processing and Modeling Pipeline\n\n\nThis codebase is designed for the processing, analysis, and modeling of RNA 3D folding data. It includes several distinct sections that handle everything from data optimization and loading to model evaluation and submission generation.\n\n\n## Overview\n\n\n- **Data Optimization and Loading:**  \n\n  Functions are provided to optimize DataFrame memory usage by downcasting numeric types and converting low-cardinality object columns into categorical types. CSV files containing RNA sequences and labels (3D coordinates) are loaded and concatenated as needed.\n\n\n- **Data Integrity and Duplicate Checks:**  \n\n  The code performs integrity checks to ensure that data transformations preserve the original information and identifies any duplicate entries.\n\n\n- **Data Analysis:**  \n\n  Separate routines analyze RNA sequences and 3D coordinates. This includes plotting sequence length distributions and nucleotide frequencies for sequences, as well as coordinate histograms for structural labels.\n\n- **Processed Data Creation:**  \n\n  The pipeline processes sequence and structure data into NumPy arrays (features and targets) suitable for model training and saves these arrays to disk.\n\n\n- **Structural Variation and TM-Score Calculations:**  \n\n  Functions are implemented to introduce structural variation (by adding controlled noise) and to calculate the TM-score, a metric used to evaluate the similarity between predicted and true RNA structures.\n\n\n- **Inference and Submission Generation:**  \n\n  A full inference pipeline is provided, including:\n\n  - Creating input JSON files for a given RNA sequence.\n\n  - Running inference using a pre-trained Protenix model.\n\n  - Extracting and processing predictions (including C1' atom coordinates).\n\n  - Generating multi-structure submissions by averaging predictions from several models.\n\n\n- **Model Search and Evaluation:**  \n\n  The code includes routines to iterate over multiple training configurations in search of the best performing model based on metrics such as MAE, MSE, and TM-score.","metadata":{}},{"cell_type":"markdown","source":"## Library Imports ","metadata":{}},{"cell_type":"code","source":"# Core Python Libraries\nimport os\nimport time\nimport gc\nimport glob\nimport json\nimport sys\nimport subprocess\nimport traceback\nfrom collections import Counter\nfrom typing import Optional, Any, Callable\nfrom tqdm import tqdm\n\n# Data Manipulation Libraries\nimport numpy as np\nimport pandas as pd\n\n# Machine Learning Libraries\nimport tensorflow as tf\nfrom sklearn.model_selection import train_test_split\n\n# Visualization Libraries\nimport matplotlib.pyplot as plt\nimport matplotlib.colors as mcolors\n\n# Configuration\nwarnings = __import__('warnings')  # Import warnings\nwarnings.filterwarnings('ignore')  # Suppress warnings\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-05T11:33:33.766011Z","iopub.status.idle":"2025-04-05T11:33:33.766272Z","shell.execute_reply":"2025-04-05T11:33:33.766164Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Data Exploration and Analysis","metadata":{}},{"cell_type":"code","source":"# Define constants and file paths\nDATA_DIR: str = os.getenv('DATA_DIR', '/kaggle/input/stanford-rna-3d-folding/')\nMAIN_FILES: list[str] = [\n    \"train_sequences.csv\", \n    \"train_labels.csv\", \n    \"validation_sequences.csv\", \n    \"validation_labels.csv\", \n    \"test_sequences.csv\",\n    \"sample_submission.csv\"\n]\nDEFAULT_THRESHOLD: float = 0.4  # Default threshold for category conversion\n\n# -------------------------\n# Data Optimization Section\n# -------------------------\ndef optimize_dataframe(\n    df: pd.DataFrame, \n    inplace: bool = False, \n    category_threshold: float = DEFAULT_THRESHOLD\n) -> pd.DataFrame:\n    \"\"\"\n    Optimizes the DataFrame to save memory by downcasting numeric types and\n    converting low-cardinality object columns to 'category'.\n    \"\"\"\n    if not 0 <= category_threshold <= 1:\n        raise ValueError(\"category_threshold must be between 0 and 1.\")\n    \n    df = df if inplace else df.copy()\n    \n    for col in df.columns:\n        col_type = df[col].dtype\n        \n        # Optimize integer columns\n        if np.issubdtype(col_type, np.integer):\n            c_min, c_max = df[col].min(), df[col].max()\n            if c_min > np.iinfo(np.int8).min and c_max < np.iinfo(np.int8).max:\n                df[col] = df[col].astype(np.int8)\n            elif c_min > np.iinfo(np.int16).min and c_max < np.iinfo(np.int16).max:\n                df[col] = df[col].astype(np.int16)\n            elif c_min > np.iinfo(np.int32).min and c_max < np.iinfo(np.int32).max:\n                df[col] = df[col].astype(np.int32)\n        \n        # Optimize floating point columns\n        elif np.issubdtype(col_type, np.floating):\n            if df[col].min() > np.finfo(np.float32).min and df[col].max() < np.finfo(np.float32).max:\n                df[col] = df[col].astype(np.float32)\n        \n        # Convert object columns to category if low cardinality\n        elif col_type == object:\n            unique_vals = len(df[col].unique())\n            if unique_vals / len(df) < category_threshold:\n                df[col] = df[col].astype('category')\n    \n    return df\n\n# -------------------------\n# Data Loading Section\n# -------------------------\ndef load_main_data(chunksize: int = 50000) -> dict[str, pd.DataFrame]:\n    \"\"\"\n    Loads CSV files from the main files list, optimizes them, and returns a dictionary\n    mapping file names to their respective DataFrames.\n    \"\"\"\n    data: dict[str, pd.DataFrame] = {}\n    for file_name in MAIN_FILES:\n        file_path: str = os.path.join(DATA_DIR, file_name)\n        if os.path.exists(file_path):\n            chunks = pd.read_csv(file_path, on_bad_lines='skip', low_memory=False, chunksize=chunksize)\n            dataframes = [optimize_dataframe(chunk, category_threshold=DEFAULT_THRESHOLD) for chunk in chunks]\n            data[file_name] = pd.concat(dataframes, ignore_index=True)\n        else:\n            print(f\"File {file_path} not found!\")\n    return data\n\n# -------------------------\n# Data Integrity and Duplicate Checks\n# -------------------------\ndef check_data_integrity(original_df: pd.DataFrame, optimized_df: pd.DataFrame) -> None:\n    \"\"\"\n    Checks whether the optimization has preserved the original data.\n    \"\"\"\n    try:\n        pd.testing.assert_frame_equal(original_df, optimized_df, check_like=True)\n        print(\"Integrity check passed: Data is unchanged after optimization.\")\n    except AssertionError as error:\n        print(f\"Data integrity check failed: {error}\")\n\ndef check_duplicates(df: pd.DataFrame) -> Optional[pd.DataFrame]:\n    \"\"\"\n    Checks and returns duplicate entries in the DataFrame.\n    \"\"\"\n    duplicates: pd.DataFrame = df[df.duplicated(keep=False)]\n    if not duplicates.empty:\n        print(f\"Warning: Found {duplicates.shape[0]} duplicate entries.\")\n        return duplicates\n    else:\n        print(\"No duplicates found.\")\n    return None\n\n# -------------------------\n# Threshold Testing and Plotting\n# -------------------------\ndef test_thresholds(df: pd.DataFrame) -> tuple[np.ndarray, list[float]]:\n    \"\"\"\n    Tests different thresholds for converting object columns to categories\n    and returns the thresholds and their corresponding memory usages.\n    \"\"\"\n    thresholds: np.ndarray = np.linspace(0.1, 0.9, 9)\n    memory_usages: list[float] = []\n    for threshold in thresholds:\n        optimized_df = optimize_dataframe(df.copy(), category_threshold=threshold)\n        mem_usage = optimized_df.memory_usage(deep=True).sum() / 1024**2\n        memory_usages.append(mem_usage)\n    return thresholds, memory_usages\n\ndef plot_memory_usage(thresholds: np.ndarray, memory_usages: list[float]) -> None:\n    \"\"\"\n    Plots memory usage against the threshold values.\n    \"\"\"\n    plt.figure(figsize=(10, 6))\n    plt.plot(thresholds, memory_usages, marker='o', linestyle='-')\n    plt.title(\"Memory Usage vs. Threshold\")\n    plt.xlabel(\"Threshold\")\n    plt.ylabel(\"Memory Usage (MB)\")\n    plt.grid(True)\n    plt.show()\n\n# -------------------------\n# Data Analysis Section\n# -------------------------\ndef analyze_sequence_data(df_sequences: pd.DataFrame) -> None:\n    \"\"\"\n    Analyzes RNA sequence data and plots a histogram of sequence lengths.\n    \"\"\"\n    print(f\"Total sequences: {len(df_sequences)}\")\n    print(f\"Columns available: {df_sequences.columns.tolist()}\")\n    \n    if 'sequence' in df_sequences.columns:\n        seq_lengths = df_sequences['sequence'].apply(len)\n        \n        # Plot histogram of sequence lengths\n        plt.figure(figsize=(10, 6))\n        plt.hist(seq_lengths, bins=30, edgecolor='black')\n        plt.title(\"Distribution of RNA Sequence Lengths\")\n        plt.xlabel(\"Sequence Length\")\n        plt.ylabel(\"Frequency\")\n        plt.show()\n        \n        print(f\"Sequence length statistics:\\nMin: {seq_lengths.min()}, Max: {seq_lengths.max()}, Average: {seq_lengths.mean():.2f}\")\n        \n        # Nucleotide count analysis\n        nucleotides = ['A', 'C', 'G', 'U']\n        nucleotide_counts = {n: df_sequences['sequence'].str.count(n).sum() for n in nucleotides}\n        total_nucleotides = sum(nucleotide_counts.values())\n        \n        print(\"\\nNucleotide distribution:\")\n        for nucleotide, count in nucleotide_counts.items():\n            print(f\"{nucleotide}: {count} ({(count / total_nucleotides * 100):.2f}%)\")\n        plot_nucleotide_distribution(df_sequences)\n\ndef plot_nucleotide_distribution(df_sequences: pd.DataFrame) -> None:\n    \"\"\"\n    Plots the nucleotide distribution as a bar chart showing counts and percentages\n    with different colors for each nucleotide.\n    \"\"\"\n    nucleotides = ['A', 'C', 'G', 'U']\n    nucleotide_counts = {n: df_sequences['sequence'].str.count(n).sum() for n in nucleotides}\n    total_nucleotides = sum(nucleotide_counts.values())\n    \n    # Prepare data for plotting\n    labels = list(nucleotide_counts.keys())\n    counts = list(nucleotide_counts.values())\n    percentages = [count / total_nucleotides * 100 for count in counts]\n    \n    # Define different colors for each bar\n    colors = ['red', 'green', 'blue', 'orange']\n    \n    # Create bar chart with different colors\n    plt.figure(figsize=(8, 6))\n    bars = plt.bar(labels, counts, color=colors, edgecolor='black')\n    plt.xlabel(\"Nucleotide\")\n    plt.ylabel(\"Count\")\n    plt.title(\"Nucleotide Distribution\")\n    \n    # Annotate each bar with its percentage\n    for bar, percentage in zip(bars, percentages):\n        height = bar.get_height()\n        plt.text(bar.get_x() + bar.get_width() / 2, height + 0.001 * total_nucleotides, f\"{percentage:.2f}%\", ha='center', va='bottom')\n\n\ndef analyze_label_data(df_labels: pd.DataFrame) -> None:\n    \"\"\"\n    Analyzes 3D coordinate data in the labels DataFrame and plots histograms for each coordinate.\n    \"\"\"\n    print(f\"Total label entries: {len(df_labels)}\")\n    print(f\"Columns available: {df_labels.columns.tolist()}\")\n    \n    coord_columns = [col for col in df_labels.columns if col.startswith(('x_', 'y_', 'z_'))]\n    if coord_columns:\n        print(f\"\\nFound {len(coord_columns)} coordinate columns.\")\n        # Plot histograms for each coordinate column\n        for col in coord_columns:\n            plt.figure(figsize=(10, 4))\n            plt.hist(df_labels[col].dropna(), bins=30, edgecolor='black')\n            plt.title(f\"Distribution of {col}\")\n            plt.xlabel(col)\n            plt.ylabel(\"Frequency\")\n            plt.show()\n            \n            print(f\"{col} - Mean: {df_labels[col].mean():.2f}, Std: {df_labels[col].std():.2f}\")\n\n# -------------------------\n# Submission Template Creation\n# -------------------------\ndef create_submission_template(\n    test_df: pd.DataFrame, \n    sample_submission_df: Optional[pd.DataFrame]\n) -> pd.DataFrame:\n    \"\"\"\n    Creates a submission template. If a sample submission is provided, it copies it;\n    otherwise, it generates a new template based on the test sequences.\n    \"\"\"\n    if sample_submission_df is None:\n        print(\"Sample submission file not found. Generating a new submission template.\")\n        ids, resnames, resids = [], [], []\n        \n        for _, row in test_df.iterrows():\n            sequence = row.get('sequence')\n            target_id = row.get('target_id')\n            if sequence and target_id:\n                for i, nucleotide in enumerate(sequence, start=1):\n                    ids.append(f\"{target_id}_{i}\")\n                    resnames.append(nucleotide)\n                    resids.append(i)\n        \n        submission_df = pd.DataFrame({\n            'ID': ids,\n            'resname': resnames,\n            'resid': resids\n        })\n        \n        # Initialize coordinate columns for five structures\n        for i in range(1, 6):\n            submission_df[f'x_{i}'] = 0.0\n            submission_df[f'y_{i}'] = 0.0\n            submission_df[f'z_{i}'] = 0.0\n    else:\n        submission_df = sample_submission_df.copy()\n        print(\"Submission template created based on the provided sample submission.\")\n    \n    return submission_df\n\n# -------------------------\n# Main Routine\n# -------------------------\ndef main() -> dict[str, pd.DataFrame]:\n    start_time = time.time()\n    \n    print(\"Loading main data files...\")\n    data: dict[str, pd.DataFrame] = load_main_data()\n    \n    # Report loaded file shapes\n    print(\"\\nLoaded files:\")\n    for file_name, df in data.items():\n        shape_info = df.shape if df is not None else \"Not found\"\n        print(f\"- {file_name}: {shape_info}\")\n    \n    # Analyze training sequences\n    if \"train_sequences.csv\" in data:\n        print(\"\\n===== Training Sequences Analysis =====\")\n        analyze_sequence_data(data[\"train_sequences.csv\"])\n    \n    # Analyze training labels\n    if \"train_labels.csv\" in data:\n        print(\"\\n===== Training Labels Analysis =====\")\n        analyze_label_data(data[\"train_labels.csv\"])\n    \n    # Check for duplicates in training sequences\n    if \"train_sequences.csv\" in data:\n        print(\"\\nChecking duplicates in training sequences...\")\n        check_duplicates(data[\"train_sequences.csv\"])\n    \n    \n    # Create submission template\n    if \"test_sequences.csv\" in data:\n        print(\"\\nCreating submission template...\")\n        submission_df = create_submission_template(\n            test_df=data[\"test_sequences.csv\"],\n            sample_submission_df=data.get(\"sample_submission.csv\")\n        )\n        print(f\"Submission template shape: {submission_df.shape}\")\n        print(submission_df.head())\n    \n    runtime = time.time() - start_time\n    print(f\"\\nRuntime: {runtime:.2f} seconds\")\n    \n    return data\n\nif __name__ == '__main__':\n    main_data = main()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-05T11:25:17.296902Z","iopub.execute_input":"2025-04-05T11:25:17.297201Z","iopub.status.idle":"2025-04-05T11:25:18.592746Z","shell.execute_reply.started":"2025-04-05T11:25:17.297178Z","shell.execute_reply":"2025-04-05T11:25:18.591939Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Data Preperation","metadata":{}},{"cell_type":"code","source":"# File paths\nDATA_DIR: str = \"/kaggle/input/stanford-rna-3d-folding/\"\nOUTPUT_DIR: str = \"/kaggle/working/\"\nos.makedirs(OUTPUT_DIR, exist_ok=True)\n\ndef load_data() -> dict[str, pd.DataFrame]:\n    \"\"\"\n    Loads the necessary CSV files for the competition.\n    \n    Returns:\n        A dictionary containing DataFrames for:\n          - train_seq: training sequences\n          - valid_seq: validation sequences\n          - test_seq: test sequences\n          - train_labels: training structures (3D coordinates)\n          - valid_labels: validation structures (3D coordinates)\n          - sample_submission: sample submission format\n    \"\"\"\n    data: dict[str, pd.DataFrame] = {}\n    data['train_seq'] = pd.read_csv(os.path.join(DATA_DIR, \"train_sequences.csv\"))\n    data['valid_seq'] = pd.read_csv(os.path.join(DATA_DIR, \"validation_sequences.csv\"))\n    data['test_seq'] = pd.read_csv(os.path.join(DATA_DIR, \"test_sequences.csv\"))\n    data['train_labels'] = pd.read_csv(os.path.join(DATA_DIR, \"train_labels.csv\"))\n    data['valid_labels'] = pd.read_csv(os.path.join(DATA_DIR, \"validation_labels.csv\"))\n    data['sample_submission'] = pd.read_csv(os.path.join(DATA_DIR, \"sample_submission.csv\"))\n    \n    return data\n\ndef create_mapping_valid(valid_seq_df: pd.DataFrame, valid_labels_df: pd.DataFrame) -> dict[str, dict[str, Any]]:\n    \"\"\"\n    Creates a mapping between validation sequences and their corresponding 3D structures.\n    \n    Args:\n        valid_seq_df: DataFrame with validation sequences.\n        valid_labels_df: DataFrame with validation labels (3D coordinates).\n        \n    Returns:\n        A dictionary where each key is a sequence ID and the value is a dictionary with:\n          - 'sequence': The RNA sequence.\n          - 'structures': A list of structures (each is a list of [x, y, z] coordinates).\n    \"\"\"\n    # Extract sequence ID from label ID (format: e.g., \"R1107_1\" -> \"R1107\")\n    valid_labels_df['seq_id'] = valid_labels_df['ID'].apply(lambda x: x.split('_')[0])\n    \n    seq_ids: set = set(valid_seq_df['target_id'])\n    label_seq_ids: set = set(valid_labels_df['seq_id'])\n    \n    overlap: set = seq_ids.intersection(label_seq_ids)\n    print(f\"Correspondence for validation: {len(overlap)} of {len(seq_ids)}\")\n    \n    mapping: dict[str, dict[str, Any]] = {}\n    for seq_id in overlap:\n        # Retrieves the sequence\n        seq: str = valid_seq_df[valid_seq_df['target_id'] == seq_id]['sequence'].iloc[0]\n    \n        # Get all residues for this sequence, sorted by residue number\n        residues: pd.DataFrame = valid_labels_df[valid_labels_df['seq_id'] == seq_id].sort_values('resid')\n        \n        # Determine the number of structures from coordinate columns (x_i)\n        num_structures: int = 1\n        for col in residues.columns:\n            if col.startswith('x_'):\n                struct_num = int(col.split('_')[1])\n                num_structures = max(num_structures, struct_num)\n        \n        structures: list[list[list[float]]] = []\n        \n        # Build a structure for each available prediction\n        for struct_idx in range(1, num_structures + 1):\n            coords: list[list[float]] = []\n            has_valid_coords: bool = False\n            if f'x_{struct_idx}' in residues.columns:\n                for _, row in residues.iterrows():\n                    x: float = row[f'x_{struct_idx}']\n                    y: float = row[f'y_{struct_idx}']\n                    z: float = row[f'z_{struct_idx}']\n                    # Validate coordinate values\n                    if abs(x) < 1.0e+17 and abs(y) < 1.0e+17 and abs(z) < 1.0e+17:\n                        coords.append([x, y, z])\n                        has_valid_coords = True\n                    else:\n                        coords.append([np.nan, np.nan, np.nan])\n            if has_valid_coords:\n                structures.append(coords)\n        \n        if structures:\n            mapping[seq_id] = {\n                'sequence': seq,\n                'structures': structures\n            }\n    \n    print(f\"Mapping created with {len(mapping)} valid sequences\")\n    return mapping\n\ndef create_processed_data(mapping: dict[str, dict[str, Any]], output_prefix: str) -> tuple[Optional[np.ndarray], Optional[np.ndarray], list[str]]:\n    \"\"\"\n    Processes the mapping to create feature (X) and target (y) arrays suitable for training.\n    It also saves these arrays to the output directory.\n    \n    Args:\n        mapping: Dictionary with mapping of sequences to structures.\n        output_prefix: A prefix for naming output files (e.g., 'train' or 'valid').\n        \n    Returns:\n        A tuple (X, y, ids) where:\n            X: A NumPy array of one-hot encoded sequences.\n            y: A NumPy array of coordinate targets.\n            ids: List of sequence IDs.\n    \"\"\"\n    if not mapping:\n        print(f\"WARNING: No valid mapping for {output_prefix}\")\n        return None, None, []\n    \n    X_data: list[np.ndarray] = []\n    y_data: list[np.ndarray] = []\n    ids: list[str] = []\n    \n    for seq_id, data in mapping.items():\n        seq: str = data['sequence']\n        structures: list[list[list[float]]] = data['structures']\n        if not structures:\n            continue\n        \n        # Use the first valid structure for training\n        structure: list[list[float]] = structures[0]\n        if len(structure) != len(seq):\n            print(f\"WARNING: Sequence length ({len(seq)}) and coordinate count ({len(structure)}) differ for {seq_id}\")\n            continue\n        \n        # One-hot encoding for the sequence (A, C, G, U, unknown)\n        features: list[list[int]] = []\n        for nucleotide in seq:\n            if nucleotide == 'A':\n                features.append([1, 0, 0, 0, 0])\n            elif nucleotide == 'C':\n                features.append([0, 1, 0, 0, 0])\n            elif nucleotide == 'G':\n                features.append([0, 0, 1, 0, 0])\n            elif nucleotide == 'U':\n                features.append([0, 0, 0, 1, 0])\n            else:\n                features.append([0, 0, 0, 0, 1])\n        \n        X_data.append(np.array(features))\n        y_data.append(np.array(structure))\n        ids.append(seq_id)\n    \n    if not X_data:\n        print(f\"WARNING: No valid processed data for {output_prefix}\")\n        return None, None, []\n    \n    # Pad sequences so all examples have the same length\n    max_length: int = max(x.shape[0] for x in X_data)\n    X_padded: list[np.ndarray] = []\n    y_padded: list[np.ndarray] = []\n    \n    for x, y in zip(X_data, y_data):\n        if x.shape[0] < max_length:\n            x_pad = np.zeros((max_length, 5))\n            x_pad[:x.shape[0], :] = x\n            y_pad = np.zeros((max_length, 3))\n            y_pad[:y.shape[0], :] = y\n            X_padded.append(x_pad)\n            y_padded.append(y_pad)\n        else:\n            X_padded.append(x)\n            y_padded.append(y)\n    \n    X: np.ndarray = np.array(X_padded)\n    y: np.ndarray = np.array(y_padded)\n    \n    # Save processed arrays\n    np.save(os.path.join(OUTPUT_DIR, f'X_{output_prefix}.npy'), X)\n    np.save(os.path.join(OUTPUT_DIR, f'y_{output_prefix}.npy'), y)\n    with open(os.path.join(OUTPUT_DIR, f'{output_prefix}_ids.txt'), 'w') as f:\n        for id_val in ids:\n            f.write(f\"{id_val}\\n\")\n    \n    print(f\"Processed data for {output_prefix}: X.shape = {X.shape}, y.shape = {y.shape}\")\n    return X, y, ids\n\ndef main() -> dict[str, Any]:\n    \"\"\"\n    Main routine for loading data, creating a mapping for validation sequences,\n    processing the data into features and targets, and saving the processed data.\n    \n    Returns:\n        A dictionary containing processed data arrays and mapping details.\n    \"\"\"\n    print(\"Loading data...\")\n    data_dict: dict[str, pd.DataFrame] = load_data()\n    \n    print(\"\\nCreating mapping for validation data...\")\n    valid_mapping: dict[str, dict[str, Any]] = create_mapping_valid(data_dict['valid_seq'], data_dict['valid_labels'])\n    \n    # Process and save validation data\n    X_valid, y_valid, valid_ids = create_processed_data(valid_mapping, 'valid')\n    \n    # In absence of a reliable training mapping, we use validation data for training\n    print(\"\\nUsing validation data as training (due to mapping constraints)...\")\n    X_train, y_train, train_ids = X_valid, y_valid, valid_ids\n    if X_train is not None:\n        np.save(os.path.join(OUTPUT_DIR, 'X_train.npy'), X_train)\n        np.save(os.path.join(OUTPUT_DIR, 'y_train.npy'), y_train)\n        with open(os.path.join(OUTPUT_DIR, 'train_ids.txt'), 'w') as f:\n            for id_val in train_ids:\n                f.write(f\"{id_val}\\n\")\n    \n    return {\n        'X_train': X_train,\n        'y_train': y_train,\n        'X_valid': X_valid,\n        'y_valid': y_valid,\n        'valid_mapping': valid_mapping,\n        'valid_ids': valid_ids\n    }\n\nif __name__ == \"__main__\":\n    processed_data = main()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-05T11:25:22.566458Z","iopub.execute_input":"2025-04-05T11:25:22.566827Z","iopub.status.idle":"2025-04-05T11:25:27.726458Z","shell.execute_reply.started":"2025-04-05T11:25:22.566799Z","shell.execute_reply":"2025-04-05T11:25:27.725605Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## RNA 3D Folding Pipeline","metadata":{}},{"cell_type":"markdown","source":"","metadata":{}},{"cell_type":"code","source":"# Define directories\nDATA_DIR = \"/kaggle/input/stanford-rna-3d-folding/\"\nOUTPUT_DIR = \"/kaggle/working/\"\nos.makedirs(OUTPUT_DIR, exist_ok=True)\n\n##############################################\n# 1. Function to generate structural variation\n##############################################\ndef sample_structural_variation(coords, noise_level=0.5, preserve_distance=True, \n                                use_global_movement=False, correlation=0.7):\n    \"\"\"\n    Enhanced version of structural variation sampling with better\n    handling of large RNAs and improved noise distribution.\n    \"\"\"\n    new_coords = coords.copy()\n    valid_mask = ~np.all(coords == 0, axis=1)\n    valid_indices = np.where(valid_mask)[0]\n    \n    if len(valid_indices) < 3:\n        return new_coords\n    \n    typical_bond_length = 3.8  # Typical RNA backbone distance\n    \n    if use_global_movement and len(valid_indices) > 20:\n        distances = []\n        for i in range(1, len(valid_indices)):\n            idx1 = valid_indices[i-1]\n            idx2 = valid_indices[i]\n            dist = np.linalg.norm(coords[idx1] - coords[idx2])\n            distances.append((i, dist))\n        distances.sort(key=lambda x: x[1], reverse=True)\n        num_hinges = min(2, len(distances)//3)\n        for h in range(num_hinges):\n            if h < len(distances):\n                hinge_point = distances[h][0]\n                if hinge_point < 5 or hinge_point > len(valid_indices) - 5:\n                    continue\n                hinge_idx = valid_indices[hinge_point]\n                angle = np.random.exponential(0.2)\n                if np.random.random() < 0.5:\n                    angle = -angle\n                sin_a, cos_a = np.sin(angle), np.cos(angle)\n                tilt = np.random.normal(0, 0.1)\n                rotation_matrix = np.array([\n                    [cos_a, -sin_a, 0],\n                    [sin_a, cos_a, tilt],\n                    [0, -tilt, 1]\n                ])\n                ref_point = new_coords[hinge_idx]\n                for i in valid_indices[hinge_point+1:]:\n                    vector = new_coords[i] - ref_point\n                    rotated = np.dot(vector, rotation_matrix)\n                    new_coords[i] = ref_point + rotated\n    \n    prev_noise = np.zeros(3)\n    correlation = 0.5  # High correlation for smoother variations\n    for i in range(1, len(coords)):\n        if not valid_mask[i] or not valid_mask[i-1]:\n            continue\n        vec = new_coords[i-1] - new_coords[i]\n        vec_length = np.linalg.norm(vec)\n        new_noise = np.random.normal(0, noise_level, size=3)\n        noise_vec = correlation * prev_noise + (1 - correlation) * new_noise\n        prev_noise = noise_vec.copy()\n        noise_norm = np.linalg.norm(noise_vec)\n        if noise_norm > 0:\n            noise_vec = noise_vec / noise_norm * (noise_level * vec_length)\n        new_vec = vec + noise_vec\n        if preserve_distance:\n            current_length = np.linalg.norm(new_vec)\n            if current_length > 0:\n                target_length = typical_bond_length * (1 + np.random.normal(0, 0.05))\n                new_vec = new_vec / current_length * target_length\n        new_coords[i] = new_coords[i-1] - new_vec\n    \n    return new_coords\n\ndef normalize_structure(coords):\n    \"\"\"\n    Centralizes and normalizes the structure.\n    \"\"\"\n    valid_mask = ~np.all(coords == 0, axis=1)\n    valid_coords = coords[valid_mask]\n    center = np.mean(valid_coords, axis=0)\n    centered_coords = coords.copy()\n    centered_coords[valid_mask] = valid_coords - center\n    return centered_coords\n\ndef check_structure_validity(coords, min_distance=0.8, max_distance=7.0, allow_clashes=0.05):\n    \"\"\"\n    Checks the structure for bond validity and residue clashes.\n    \"\"\"\n    valid = True\n    valid_mask = ~np.all(coords == 0, axis=1)\n    valid_coords = coords[valid_mask]\n    \n    if len(valid_coords) < 3:\n        return True\n    \n    invalid_bonds = 0\n    for i in range(1, len(valid_coords)):\n        dist = np.linalg.norm(valid_coords[i] - valid_coords[i-1])\n        if dist < min_distance or dist > max_distance:\n            invalid_bonds += 1\n    if invalid_bonds / len(valid_coords) > 0.1:\n        valid = False\n    \n    clashes = 0\n    total_pairs = 0\n    for i in range(len(valid_coords)):\n        for j in range(i+3, len(valid_coords)):\n            total_pairs += 1\n            dist = np.linalg.norm(valid_coords[i] - valid_coords[j])\n            if dist < min_distance:\n                clashes += 1\n    if total_pairs > 0 and clashes / total_pairs > allow_clashes:\n        valid = False\n    \n    return valid\n\n##############################################\n# 2. TM-score Calculation Functions\n##############################################\ndef calculate_tm_score(pred_coords, true_coords, d0_scale=1.24):\n    \"\"\"\n    Calculates an approximate TM-score between predicted and true coordinates.\n    \"\"\"\n    mask = ~np.all(true_coords == 0, axis=1)\n    pred = pred_coords[mask]\n    true = true_coords[mask]\n    L = len(true)\n    if L < 3:\n        return 0.0\n    if L >= 30:\n        d0 = 0.6 * np.sqrt(L - 0.5) - 2.5\n        d0 = max(0.1, d0)\n    elif L >= 24:\n        d0 = 0.7\n    elif L >= 20:\n        d0 = 0.6\n    elif L >= 16:\n        d0 = 0.5\n    elif L >= 12:\n        d0 = 0.4\n    else:\n        d0 = 0.3\n    distances = np.sqrt(np.sum((pred - true) ** 2, axis=1))\n    tm_terms = 1.0 / (1.0 + (distances / (d0 + 1e-8)) ** 2)\n    tm_score = np.sum(tm_terms) / L\n    return float(tm_score)\n\ndef calculate_tm_score_exact(pred_coords, true_coords):\n    \"\"\"\n    Calculates TM-score with multiple rotation schemes.\n    \"\"\"\n    mask = ~np.all(true_coords == 0, axis=1)\n    pred = pred_coords[mask]\n    true = true_coords[mask]\n    Lref = len(true)\n    if Lref < 3:\n        return 0.0\n    if Lref >= 30:\n        d0 = 0.6 * np.sqrt(Lref - 0.5) - 2.5\n    elif Lref >= 24:\n        d0 = 0.7\n    elif Lref >= 20:\n        d0 = 0.6\n    elif Lref >= 16:\n        d0 = 0.5\n    elif Lref >= 12:\n        d0 = 0.4\n    else:\n        d0 = 0.3\n    pred_centered = pred - np.mean(pred, axis=0)\n    true_centered = true - np.mean(true, axis=0)\n    best_tm_score = 0.0\n    fragment_lengths = [Lref, max(5, Lref//2), max(5, Lref//4)]\n    for frag_len in fragment_lengths:\n        for i in range(0, Lref - frag_len + 1, max(1, frag_len//2)):\n            pred_frag = pred_centered[i:i+frag_len]\n            for j in range(0, Lref - frag_len + 1, max(1, frag_len//2)):\n                true_frag = true_centered[j:j+frag_len]\n                covariance = np.dot(pred_frag.T, true_frag)\n                U, S, Vt = np.linalg.svd(covariance)\n                rotation = np.dot(U, Vt)\n                rotations_to_try = [\n                    rotation,\n                    np.dot(rotation, np.array([[0, 1, 0], [-1, 0, 0], [0, 0, 1]])),\n                    np.dot(rotation, np.array([[-1, 0, 0], [0, -1, 0], [0, 0, 1]]))\n                ]\n                for rot in rotations_to_try:\n                    pred_aligned = np.dot(pred_centered, rot)\n                    distances = np.sqrt(np.sum((pred_aligned - true_centered) ** 2, axis=1))\n                    tm_terms = 1.0 / (1.0 + (distances / d0) ** 2)\n                    tm_score = np.sum(tm_terms) / Lref\n                    best_tm_score = max(best_tm_score, tm_score)\n    return float(best_tm_score)\n\n##############################################\n# 3. Data Loading Function\n##############################################\ndef load_processed_data():\n    \"\"\"\n    Loads processed training and validation data.\n    \"\"\"\n    X_train = np.load(os.path.join(OUTPUT_DIR, 'X_train.npy'))\n    y_train = np.load(os.path.join(OUTPUT_DIR, 'y_train.npy'))\n    X_valid = np.load(os.path.join(OUTPUT_DIR, 'X_valid.npy'))\n    y_valid = np.load(os.path.join(OUTPUT_DIR, 'y_valid.npy'))\n    \n    print(f\"Data loaded - X_train: {X_train.shape}, y_train: {y_train.shape}\")\n    print(f\"Data loaded - X_valid: {X_valid.shape}, y_valid: {y_valid.shape}\")\n    \n    return X_train, y_train, X_valid, y_valid\n\n##############################################\n# 4. Reference Model (Baseline)\n##############################################\ndef reference_based_approach(X_ref, y_ref, geometric_sampling=False, noise_level=0.2, correlation=0.7):\n    try:\n        class ReferenceModel:\n            def __init__(self, geometric_sampling=False, base_noise_level=0.2, correlation=0.7):\n                self.geometric_sampling = geometric_sampling\n                self.base_noise_level = base_noise_level\n                self.correlation = correlation\n                \n            def fit(self, X, y):\n                self.reference_structures = np.nan_to_num(y, nan=0.0)\n                self.global_mean = np.nanmean(y, axis=(0, 1))\n                self.global_std = np.nanstd(y, axis=(0, 1))\n                self.global_mean = np.nan_to_num(self.global_mean, nan=0.0)\n                self.global_std = np.nan_to_num(self.global_std, nan=1.0)\n                self.size_groups = {}\n                for i in range(len(self.reference_structures)):\n                    valid_mask = ~np.all(self.reference_structures[i] == 0, axis=1)\n                    size = np.sum(valid_mask)\n                    if size < 120:\n                        group = \"small\"\n                    elif size < 200:\n                        group = \"medium\"\n                    else:\n                        group = \"large\"\n                    if group not in self.size_groups:\n                        self.size_groups[group] = []\n                    self.size_groups[group].append(i)\n                print(f\"Size distribution - Small: {len(self.size_groups.get('small', []))}, \"\n                      f\"Medium: {len(self.size_groups.get('medium', []))}, \"\n                      f\"Large: {len(self.size_groups.get('large', []))}\")\n                return self\n                \n            def predict(self, X):\n                batch_size = X.shape[0]\n                seq_length = X.shape[1]\n                predictions = np.zeros((batch_size, seq_length, 3))\n                for i in range(batch_size):\n                    valid_mask = ~np.all(X[i] == 0, axis=1)\n                    size = np.sum(valid_mask)\n                    if size < 120:\n                        noise_level = self.base_noise_level * 0.6\n                    elif size < 200:\n                        noise_level = self.base_noise_level * 1.0\n                    else:\n                        noise_level = self.base_noise_level * 0.4\n                    group = \"small\" if size < 120 else \"medium\" if size < 200 else \"large\"\n                    if group in self.size_groups:\n                        ref_idx = np.random.choice(self.size_groups[group])\n                        base_struct = self.reference_structures[ref_idx].copy()\n                        if self.geometric_sampling:\n                            predictions[i] = sample_structural_variation(\n                                base_struct, \n                                noise_level=noise_level,\n                                preserve_distance=True,\n                                use_global_movement=(size < 120),\n                                correlation=self.correlation\n                            )\n                        else:\n                            noise = np.random.normal(0, noise_level, base_struct.shape)\n                            predictions[i] = base_struct + noise\n                    else:\n                        sample = np.random.normal(self.global_mean, self.global_std, size=(seq_length, 3))\n                        if self.geometric_sampling:\n                            predictions[i] = sample_structural_variation(\n                                sample, \n                                noise_level=noise_level,\n                                preserve_distance=True,\n                                use_global_movement=(size < 120),\n                                correlation=self.correlation\n                            )\n                        else:\n                            predictions[i] = sample\n                return predictions\n        \n        model = ReferenceModel(geometric_sampling=geometric_sampling, \n                               base_noise_level=noise_level,\n                               correlation=correlation)\n        model.fit(X_ref, y_ref)\n        return model\n    \n    except Exception as e:\n        print(f\"Error in reference_based_approach: {str(e)}\")\n        return None\n\n##############################################\n# 5. Model Evaluation Function\n##############################################\ndef evaluate_model(model, X_valid, y_valid):\n    \"\"\"\n    Evaluates the model by calculating MAE, MSE, and TM-score for each structure.\n    \"\"\"\n    y_valid = np.nan_to_num(y_valid, nan=0.0)\n    y_pred = model.predict(X_valid)\n    y_pred = np.nan_to_num(y_pred, nan=0.0)\n    mae = np.mean(np.abs(y_pred - y_valid))\n    mse = np.mean((y_pred - y_valid)**2)\n    print(f\"Overall MAE: {mae:.4f}\")\n    print(f\"Overall MSE: {mse:.4f}\")\n    tm_scores = []\n    for i in range(len(X_valid)):\n        tm = calculate_tm_score(y_pred[i], y_valid[i])\n        tm_scores.append(tm)\n    avg_tm_score = np.mean(tm_scores)\n    print(f\"Average approximate TM-score: {avg_tm_score:.4f}\")\n    return {\n        'mae': mae,\n        'mse': mse,\n        'tm_scores': tm_scores,\n        'avg_tm_score': avg_tm_score\n    }\n\n##############################################\n# 6. Main Function to Search for the Best Model\n##############################################\ndef search_best_model(X_train, y_train, X_valid, y_valid, max_iterations=10, target_score=0.20):\n    \"\"\"\n    Searches for the best performing model over multiple iterations.\n    \"\"\"\n    best_model = None\n    best_metrics = None\n    best_score = 0.0\n    print(f\"Starting model search (max {max_iterations} iterations, target score: {target_score})\")\n    for iteration in range(max_iterations):\n        print(f\"\\n----- Iteration {iteration+1}/{max_iterations} -----\")\n        np.random.seed(iteration * 42)\n        model = reference_based_approach(X_valid, y_valid, geometric_sampling=True)\n        if model is None:\n            print(\"Model creation failed in this iteration, continuing...\")\n            continue\n        metrics = evaluate_model(model, X_valid, y_valid)\n        current_score = metrics['avg_tm_score']\n        print(f\"Iteration {iteration+1} TM-score: {current_score:.4f} (best so far: {best_score:.4f})\")\n        if current_score > best_score:\n            print(f\"New best model found! TM-score improved: {best_score:.4f} -> {current_score:.4f}\")\n            best_model = model\n            best_metrics = metrics\n            best_score = current_score\n            np.save(os.path.join(OUTPUT_DIR, 'best_predictions.npy'), model.predict(X_valid))\n        if current_score >= target_score:\n            print(f\"Target TM-score of {target_score} reached! Stopping search.\")\n            break\n    print(f\"\\nSearch completed. Best TM-score: {best_score:.4f}\")\n    return best_model, best_metrics\n\n##############################################\n# 7. Test Feature Preparation Function\n##############################################\ndef prepare_test_features(test_seq_df, max_length=720):\n    \"\"\"\n    Prepares test features by one-hot encoding the sequence.\n    \"\"\"\n    X_test = []\n    for _, row in test_seq_df.iterrows():\n        seq = row['sequence']\n        features = []\n        for nucleotide in seq:\n            if nucleotide == 'A':\n                features.append([1, 0, 0, 0, 0])\n            elif nucleotide == 'C':\n                features.append([0, 1, 0, 0, 0])\n            elif nucleotide == 'G':\n                features.append([0, 0, 1, 0, 0])\n            elif nucleotide == 'U':\n                features.append([0, 0, 0, 1, 0])\n            else:\n                features.append([0, 0, 0, 0, 1])\n        if len(features) < max_length:\n            features.extend([[0, 0, 0, 0, 0]] * (max_length - len(features)))\n        else:\n            features = features[:max_length]\n        X_test.append(features)\n    return np.array(X_test)\n\n##############################################\n# 8. Submission Generation \n##############################################\ndef generate_submission_multi_structures(model, test_seq_df, sample_submission_df, num_structs=5):\n    \"\"\"\n    Generates a submission file that includes multiple (num_structs) sets of (x,y,z)\n    coordinates for each residue: x_1, y_1, z_1, x_2, y_2, ..., x_num_structs, y_num_structs, z_num_structs.\n    \"\"\"\n    # 1) Convert the test sequences into model-friendly features\n    X_test = prepare_test_features(test_seq_df)\n    \n    # 2) We'll store multiple structures for each sequence in a dict\n    seq_to_coords = {}\n    \n    # 3) Predict coordinates. You have multiple options to get 5 different structures:\n    # Prepare a list to hold [predictions_run_1, predictions_run_2, ..., predictions_run_5],\n    # each run having shape (num_sequences, max_length, 3).\n    multi_predictions = []\n    for run_idx in range(num_structs):\n        # Change the seed so each run is different (if your model uses np.random)\n        np.random.seed(1000 + run_idx * 999)\n        preds = model.predict(X_test)\n        multi_predictions.append(preds)\n    \n    # 4) For each sequence in the test DataFrame, gather the 5 structures\n    for i, (_, row) in enumerate(test_seq_df.iterrows()):\n        target_id = row['target_id']\n        seq = row['sequence']\n        seq_length = len(seq)\n        \n        # Extract the i-th prediction from each run\n        structures_for_seq = []\n        for run_idx in range(num_structs):\n            # Take only the first seq_length residues from the run_idx-th predictions\n            coords = multi_predictions[run_idx][i][:seq_length].copy()\n            # Optional: normalize each structure\n            coords = normalize_structure(coords)\n            structures_for_seq.append(coords)\n        \n        seq_to_coords[target_id] = structures_for_seq\n    \n    # 5) Now fill in the submission DataFrame columns: x_1,y_1,z_1,...,x_5,y_5,z_5\n    submission_df = sample_submission_df.copy()\n    for i, row in submission_df.iterrows():\n        # Extract sequence ID and residue index from \"ID\" (e.g. \"R1107_3\" => seq_id=\"R1107\", residue_idx=2)\n        id_parts = row['ID'].split('_')\n        seq_id = id_parts[0]\n        residue_idx = int(id_parts[1]) - 1  # make 0-based index\n        if seq_id in seq_to_coords:\n            structures = seq_to_coords[seq_id]\n            # structures is a list of length num_structs, each shape (seq_length,3)\n            \n            # Check bounds\n            if residue_idx < len(structures[0]):\n                # Fill columns x_1,y_1,z_1,...,x_5,y_5,z_5\n                for struct_idx in range(num_structs):\n                    submission_df.at[i, f'x_{struct_idx+1}'] = structures[struct_idx][residue_idx, 0]\n                    submission_df.at[i, f'y_{struct_idx+1}'] = structures[struct_idx][residue_idx, 1]\n                    submission_df.at[i, f'z_{struct_idx+1}'] = structures[struct_idx][residue_idx, 2]\n    \n    # 6) Save the submission\n    submission_file = os.path.join(OUTPUT_DIR, 'submission1.csv')\n    submission_df.to_csv(submission_file, index=False)\n    print(f\"Multi-structure submission saved to {submission_file}\")\n    return submission_df\n\n##############################################\n# Main: Training and Submission Generation\n##############################################\nif __name__ == \"__main__\":\n    # Load training and validation data\n    X_train, y_train, X_valid, y_valid = load_processed_data()\n    \n    # Search for the best model (using multiple iterations)\n    best_model, best_metrics = search_best_model(X_train, y_train, X_valid, y_valid, max_iterations=10, target_score=0.20)\n    \n    # Load test data and sample submission template\n    test_seq_df = pd.read_csv(os.path.join(DATA_DIR, \"test_sequences.csv\"))\n    sample_submission_df = pd.read_csv(os.path.join(DATA_DIR, \"sample_submission.csv\"))\n    \n    # Generate submission using the best model's prediction (without ensemble averaging)\n    generate_submission_multi_structures(best_model, test_seq_df, sample_submission_df)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-05T11:38:14.411805Z","iopub.execute_input":"2025-04-05T11:38:14.412136Z","iopub.status.idle":"2025-04-05T11:38:16.265488Z","shell.execute_reply.started":"2025-04-05T11:38:14.412109Z","shell.execute_reply":"2025-04-05T11:38:16.264761Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Prepare Protenix Model","metadata":{}},{"cell_type":"code","source":"DATASET_PATH_PROTENIX = '/kaggle/input/required-presets'\n\n# Install dependencies from wheels\nwheel_path = os.path.join(DATASET_PATH_PROTENIX, 'wheels')\n!pip install -qqq --no-index --find-links {wheel_path} torch numpy pandas scipy rdkit protenix\n\n# Add Protenix to Python path\nprotenix_path = os.path.join(DATASET_PATH_PROTENIX, 'Protenix')\nsys.path.append(protenix_path)\nfrom biotite.structure.io import pdbx","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-05T11:25:32.205228Z","iopub.execute_input":"2025-04-05T11:25:32.205524Z","iopub.status.idle":"2025-04-05T11:28:01.402551Z","shell.execute_reply.started":"2025-04-05T11:25:32.205501Z","shell.execute_reply":"2025-04-05T11:28:01.401455Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Inference Protenix Model","metadata":{}},{"cell_type":"code","source":"def create_input_json(sequence: str, target_id: str) -> list[dict[str, Any]]:\n    \"\"\"\n    Create the input JSON for a single RNA sequence.\n    \"\"\"\n    input_json: list[dict[str, Any]] = [{\n        \"sequences\": [\n            {\n                \"rnaSequence\": {\n                    \"sequence\": sequence,\n                    \"count\": 1,\n                    \"modifications\": []\n                }\n            }\n        ],\n        \"name\": target_id,\n        \"covalent_bonds\": []\n    }]\n    return input_json\n\ndef run_inference(input_json_path: str, output_dir: str, target_id: str) -> bool:\n    \"\"\"\n    Run inference using the Protenix model with kaggle-specific paths.\n    \"\"\"\n    # Ensure output directory exists\n    os.makedirs(output_dir, exist_ok=True)\n    \n    # Define command with kaggle-specific paths\n    cmd: list[str] = [\n        \"python\", \"/kaggle/input/required-presets/Protenix/runner/inference.py\",\n        \"--seeds\", \"42\",\n        \"--dump_dir\", output_dir,\n        \"--input_json_path\", input_json_path,\n        \"--model.N_cycle\", \"10\",\n        \"--sample_diffusion.N_sample\", \"5\",\n        \"--sample_diffusion.N_step\", \"200\",\n        \"--load_checkpoint_path\", \"/kaggle/input/required-presets/Protenix/release_data/checkpoint/model_v0.2.0.pt\",\n        \"--use_deepspeed_evo_attention\", \"false\"\n    ]\n    \n    print(f\"Running inference for {target_id}\")\n    try:\n        result = subprocess.run(cmd, capture_output=True, text=True)\n        if result.returncode != 0:\n            print(f\"Error running inference for {target_id}: {result.stderr}\")\n            return False\n        return True\n    except Exception as e:\n        print(f\"Exception running inference for {target_id}: {e}\")\n        return False\n\ndef extract_c1_coordinates(cif_file_path: str) -> Optional[np.ndarray]:\n    \"\"\"\n    Extract C1' atom coordinates from a CIF file using biotite.\n    \"\"\"\n    try:\n        # Read the CIF file using the correct biotite method\n        with open(cif_file_path, 'r') as f:\n            cif_data = pdbx.CIFFile.read(f)\n        \n        # Get structure from CIF data\n        atom_array = pdbx.get_structure(cif_data, model=1)\n        \n        # Clean atom names and find C1' atoms\n        atom_names_clean = np.char.strip(atom_array.atom_name.astype(str))\n        mask_c1 = atom_names_clean == \"C1'\"\n        c1_atoms = atom_array[mask_c1]\n        \n        if len(c1_atoms) == 0:\n            print(f\"Warning: No C1' atoms found in {cif_file_path}\")\n            return None\n        \n        # Sort by residue ID and return coordinates\n        sort_indices = np.argsort(c1_atoms.res_id)\n        c1_atoms_sorted = c1_atoms[sort_indices]\n        c1_coords = c1_atoms_sorted.coord\n        \n        return c1_coords\n    except Exception as e:\n        print(f\"Error extracting C1' coordinates from {cif_file_path}: {e}\")\n        return None\n\ndef process_sequence(sequence: str, target_id: str, temp_dir: str, output_dir: str) -> Optional[list[np.ndarray]]:\n    \"\"\"\n    Process a single RNA sequence and return C1' coordinates.\n    \"\"\"\n    print(f\"Processing {target_id}: {sequence}\")\n    \n    # Create input JSON\n    input_json = create_input_json(sequence, target_id)\n    \n    # Save JSON to temporary file\n    os.makedirs(temp_dir, exist_ok=True)\n    input_json_path: str = os.path.join(temp_dir, f\"{target_id}_input.json\")\n    with open(input_json_path, \"w\") as f:\n        json.dump(input_json, f, indent=4)\n    \n    # Run inference\n    success: bool = run_inference(input_json_path, output_dir, target_id)\n    \n    if not success:\n        print(f\"Inference failed for {target_id}\")\n        return None\n    \n    # Find the CIF files for this target\n    target_prediction_dir: str = os.path.join(output_dir, target_id, \"seed_42\", \"predictions\")\n    if not os.path.exists(target_prediction_dir):\n        print(f\"Prediction directory not found for {target_id}\")\n        return None\n    \n    # Look for CIF files with the pattern {target_id}_seed_42_sample_*.cif\n    cif_files: list[str] = sorted(glob.glob(os.path.join(target_prediction_dir, f\"{target_id}_seed_42_sample_*.cif\")))\n    \n    # If no CIF files found, return None\n    if not cif_files:\n        print(f\"No CIF files found for {target_id}\")\n        return None\n    \n    print(f\"Found {len(cif_files)} CIF files for {target_id}\")\n    \n    # Extract C1' coordinates from each CIF file\n    all_coords: list[np.ndarray] = []\n    for cif_file in cif_files:\n        coords = extract_c1_coordinates(cif_file)\n        if coords is not None:\n            all_coords.append(coords)\n    \n    if not all_coords:\n        print(f\"No valid C1' coordinates found for {target_id}\")\n        return None\n    \n    # Ensure we have 5 models (if we have fewer, duplicate the last one)\n    while len(all_coords) < 5:\n        print(f\"Only {len(all_coords)} models found for {target_id}, duplicating last model\")\n        all_coords.append(all_coords[-1])\n    \n    return all_coords[:5]  # Ensure we only have 5 models\n\ndef create_submission(test_sequences_df: pd.DataFrame, c1_coords_dict: dict[str, Optional[list[np.ndarray]]], output_file: str) -> None:\n    \"\"\"\n    Create the submission CSV file with C1' coordinates.\n    \"\"\"\n    rows: list[dict[str, Any]] = []\n    \n    # Process each sequence\n    for _, row in test_sequences_df.iterrows():\n        target_id: str = row['target_id']\n        sequence: str = row['sequence']\n        \n        if target_id not in c1_coords_dict or c1_coords_dict[target_id] is None:\n            print(f\"No prediction found for {target_id}, using zeros\")\n            # Create empty predictions (all zeros)\n            for i, residue in enumerate(sequence):\n                row_data: dict[str, Any] = {\n                    'ID': f\"{target_id}_{i+1}\",\n                    'resname': residue,\n                    'resid': i+1\n                }\n                for model in range(1, 6):\n                    row_data[f'x_{model}'] = 0.0\n                    row_data[f'y_{model}'] = 0.0\n                    row_data[f'z_{model}'] = 0.0\n                rows.append(row_data)\n        else:\n            # Get the 5 models for this target\n            models: list[np.ndarray] = c1_coords_dict[target_id]  # type: ignore\n            # Create a row for each residue\n            for i, residue in enumerate(sequence):\n                row_data: dict[str, Any] = {\n                    'ID': f\"{target_id}_{i+1}\",\n                    'resname': residue,\n                    'resid': i+1\n                }\n                # Add coordinates for each model\n                for model_idx in range(5):\n                    if model_idx < len(models) and i < len(models[model_idx]):\n                        row_data[f'x_{model_idx+1}'] = models[model_idx][i][0]\n                        row_data[f'y_{model_idx+1}'] = models[model_idx][i][1]\n                        row_data[f'z_{model_idx+1}'] = models[model_idx][i][2]\n                    else:\n                        # If coordinates are not available, use zeros\n                        row_data[f'x_{model_idx+1}'] = 0.0\n                        row_data[f'y_{model_idx+1}'] = 0.0\n                        row_data[f'z_{model_idx+1}'] = 0.0\n                \n                rows.append(row_data)\n    \n    # Create DataFrame and save to CSV\n    df: pd.DataFrame = pd.DataFrame(rows)\n    df.to_csv(output_file, index=False)\n    print(f\"Created submission file: {output_file}\")\n\ndef main() -> None:\n    \"\"\"\n    Main function.\n    \"\"\"\n    # Set up required symlinks for CCD cache as in kaggle_inference.py\n    os.makedirs(\"/usr/local/lib/python3.10/dist-packages/release_data/ccd_cache\", exist_ok=True)\n    \n    source_ccd_file: str = \"/kaggle/input/required-presets/Protenix/release_data/ccd_cache/components.v20240608.cif\"\n    target_ccd_file: str = \"/usr/local/lib/python3.10/dist-packages/release_data/ccd_cache/components.v20240608.cif\"\n    \n    source_rdkit_file: str = \"/kaggle/input/required-presets/Protenix/release_data/ccd_cache/components.v20240608.cif.rdkit_mol.pkl\"\n    target_rdkit_file: str = \"/usr/local/lib/python3.10/dist-packages/release_data/ccd_cache/components.v20240608.cif.rdkit_mol.pkl\"\n    \n    # Create the symlinks if the source files exist\n    if os.path.exists(source_ccd_file) and not os.path.exists(target_ccd_file):\n        try:\n            os.symlink(source_ccd_file, target_ccd_file)\n            print(\"Created symlink for CCD file\")\n        except Exception as e:\n            print(f\"Error creating symlink for CCD file: {e}\")\n    \n    if os.path.exists(source_rdkit_file) and not os.path.exists(target_rdkit_file):\n        try:\n            os.symlink(source_rdkit_file, target_rdkit_file)\n            print(\"Created symlink for RDKIT file\")\n        except Exception as e:\n            print(f\"Error creating symlink for RDKIT file: {e}\")\n    \n    # Create directories\n    temp_dir: str = \"./input\"    # Same as in kaggle_inference.py\n    output_dir: str = \"./output\"  # Same as in kaggle_inference.py\n    os.makedirs(temp_dir, exist_ok=True)\n    os.makedirs(output_dir, exist_ok=True)\n    \n    # Load data using your preloaded function (this avoids re-reading the CSV)\n    data_dict: dict[str, Any] = load_data()  # Assumes load_data() is defined elsewhere\n    # Use the already loaded test sequences DataFrame:\n    test_sequences_df: pd.DataFrame = data_dict['test_seq']\n    print(f\"Loaded {len(test_sequences_df)} test sequences from preloaded data\")\n    \n    # Process each sequence\n    c1_coords_dict: dict[str, Optional[list[np.ndarray]]] = {}\n    for _, row in tqdm(test_sequences_df.iterrows(), total=len(test_sequences_df)):\n        target_id: str = row['target_id']\n        sequence: str = row['sequence']\n        \n        # Check if we already have predictions for this target\n        target_prediction_dir: str = os.path.join(output_dir, target_id, \"seed_42\", \"predictions\")\n        if os.path.exists(target_prediction_dir):\n            print(f\"Found existing prediction for {target_id}, loading coordinates\")\n            # Extract coordinates from existing predictions\n            cif_files: list[str] = sorted(glob.glob(os.path.join(target_prediction_dir, f\"{target_id}_seed_42_sample_*.cif\")))\n            \n            all_coords: list[np.ndarray] = []\n            for cif_file in cif_files:\n                coords = extract_c1_coordinates(cif_file)\n                if coords is not None:\n                    all_coords.append(coords)\n            \n            if all_coords:\n                # Ensure we have 5 models\n                while len(all_coords) < 5:\n                    all_coords.append(all_coords[-1])\n                c1_coords_dict[target_id] = all_coords[:5]\n                continue\n        \n        # Process the sequence if no existing prediction was found or was invalid\n        c1_coords = process_sequence(sequence, target_id, temp_dir, output_dir)\n        c1_coords_dict[target_id] = c1_coords\n    \n    # Create submission file\n    create_submission(test_sequences_df, c1_coords_dict, \"submission2.csv\")\n\nif __name__ == \"__main__\":\n    main()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-05T11:38:43.559594Z","iopub.execute_input":"2025-04-05T11:38:43.559901Z","iopub.status.idle":"2025-04-05T12:14:55.626880Z","shell.execute_reply.started":"2025-04-05T11:38:43.559880Z","shell.execute_reply":"2025-04-05T12:14:55.625926Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Ensemble Prediction","metadata":{}},{"cell_type":"code","source":"# List all submission file paths (adjust the pattern as needed)\nsubmission_files = glob.glob(\"submission*.csv\")\n\n# Load each submission file into a list of DataFrames\ndfs = [pd.read_csv(file) for file in submission_files]\n\n# Initialize a DataFrame to hold the combined predictions (using the first submission as a template)\ncombined = dfs[0].copy()\n\n# Identify coordinate columns (assuming they follow the pattern x_i, y_i, z_i for i in [1, 2, ..., num_structs])\ncoord_cols = [f'{axis}_{i}' for i in range(1, num_structs + 1) for axis in ['x', 'y', 'z']]\n\n# For each coordinate column, calculate the average across all submission DataFrames\nfor col in coord_cols:\n    # Create a list of the same column from each DataFrame\n    cols_to_average = [df[col] for df in dfs if col in df.columns]\n    # Compute the row-wise average and assign it back to the combined DataFrame\n    combined[col] = pd.concat(cols_to_average, axis=1).mean(axis=1)\n\n# Save the combined submission to a CSV file\ncombined.to_csv(\"submission.csv\", index=False)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-05T12:15:31.851827Z","iopub.execute_input":"2025-04-05T12:15:31.852126Z","iopub.status.idle":"2025-04-05T12:15:31.947543Z","shell.execute_reply.started":"2025-04-05T12:15:31.852104Z","shell.execute_reply":"2025-04-05T12:15:31.946647Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"pd.read_csv(\"submission.csv\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-05T12:15:36.663158Z","iopub.execute_input":"2025-04-05T12:15:36.663434Z","iopub.status.idle":"2025-04-05T12:15:36.700660Z","shell.execute_reply.started":"2025-04-05T12:15:36.663412Z","shell.execute_reply":"2025-04-05T12:15:36.699940Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}