{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":87793,"databundleVersionId":11228175,"sourceType":"competition"}],"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# 1. Imports, Mixed Precision, and XLA\nimport os\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nimport tensorflow as tf\nfrom tensorflow.keras.models import Model\nfrom tensorflow.keras.layers import (Input, Embedding, Conv1D, Dropout, Add,\n                                     Activation, LayerNormalization)\nfrom tensorflow.keras.preprocessing.sequence import pad_sequences\nfrom tensorflow.keras.callbacks import EarlyStopping, ModelCheckpoint\n\n# Enable XLA for additional performance (if supported)\ntf.config.optimizer.set_jit(True)\n# Enable mixed precision for faster training\ntf.keras.mixed_precision.set_global_policy('mixed_float16')\n\n# Set seeds for reproducibility\nnp.random.seed(42)\ntf.random.set_seed(42)\n\n# 2. Data Loading and Exploration\nTRAIN_SEQ_PATH = '/kaggle/input/stanford-rna-3d-folding/train_sequences.csv'\nTRAIN_LABELS_PATH = '/kaggle/input/stanford-rna-3d-folding/train_labels.csv'\nVALID_SEQ_PATH = '/kaggle/input/stanford-rna-3d-folding/validation_sequences.csv'\nVALID_LABELS_PATH = '/kaggle/input/stanford-rna-3d-folding/validation_labels.csv'\nTEST_SEQ_PATH  = '/kaggle/input/stanford-rna-3d-folding/test_sequences.csv'\nSAMPLE_SUB_PATH = '/kaggle/input/stanford-rna-3d-folding/sample_submission.csv'\n\ntrain_sequences = pd.read_csv(TRAIN_SEQ_PATH)\ntrain_labels = pd.read_csv(TRAIN_LABELS_PATH)\nvalid_sequences = pd.read_csv(VALID_SEQ_PATH)\nvalid_labels = pd.read_csv(VALID_LABELS_PATH)\ntest_sequences = pd.read_csv(TEST_SEQ_PATH)\nsample_submission = pd.read_csv(SAMPLE_SUB_PATH)\n\n# Fill missing label values with 0\ntrain_labels.fillna(0, inplace=True)\nvalid_labels.fillna(0, inplace=True)\n\nprint(\"Train Sequences Shape:\", train_sequences.shape)\nprint(\"Train Labels Shape:\", train_labels.shape)\nprint(\"Validation Sequences Shape:\", valid_sequences.shape)\nprint(\"Validation Labels Shape:\", valid_labels.shape)\nprint(\"Test Sequences Shape:\", test_sequences.shape)\n\n# 3. Data Preprocessing\n\n## 3.1 Sequence Encoding\n# Map nucleotides: A:1, C:2, G:3, U:4; unknown -> 0\nnucleotide_map = {'A': 1, 'C': 2, 'G': 3, 'U': 4}\ndef encode_sequence(seq):\n    return [nucleotide_map.get(ch, 0) for ch in seq]\n\ntrain_sequences['encoded'] = train_sequences['sequence'].apply(encode_sequence)\nvalid_sequences['encoded'] = valid_sequences['sequence'].apply(encode_sequence)\ntest_sequences['encoded'] = test_sequences['sequence'].apply(encode_sequence)\n\n## 3.2 Processing Label Data\ndef process_labels(labels_df):\n    label_dict = {}\n    for idx, row in labels_df.iterrows():\n        parts = row['ID'].split('_')\n        target_id = \"_\".join(parts[:-1])\n        resid = int(parts[-1])\n        coord = np.array([row['x_1'], row['y_1'], row['z_1']], dtype=np.float32)\n        if target_id not in label_dict:\n            label_dict[target_id] = []\n        label_dict[target_id].append((resid, coord))\n    for key in label_dict:\n        sorted_coords = sorted(label_dict[key], key=lambda x: x[0])\n        coords = np.stack([c for r, c in sorted_coords])\n        label_dict[key] = coords\n    return label_dict\n\ntrain_labels_dict = process_labels(train_labels)\nvalid_labels_dict = process_labels(valid_labels)\n\n## 3.3 Creating Datasets and Padding\ndef create_dataset(sequences_df, labels_dict):\n    X, y, target_ids = [], [], []\n    for idx, row in sequences_df.iterrows():\n        tid = row['target_id']\n        if tid in labels_dict:\n            X.append(row['encoded'])\n            y.append(labels_dict[tid])\n            target_ids.append(tid)\n    return X, y, target_ids\n\nX_train, y_train, train_ids = create_dataset(train_sequences, train_labels_dict)\nX_valid, y_valid, valid_ids = create_dataset(valid_sequences, valid_labels_dict)\n\n# Determine maximum sequence length from training set (this is fixed)\nmax_len = max(len(seq) for seq in X_train)\nprint(\"Maximum sequence length (train):\", max_len)\n\nX_train_pad = pad_sequences(X_train, maxlen=max_len, padding='post', value=0)\nX_valid_pad = pad_sequences(X_valid, maxlen=max_len, padding='post', value=0)\n\ndef pad_coordinates(coord_array, max_len):\n    L = coord_array.shape[0]\n    if L < max_len:\n        pad_width = ((0, max_len - L), (0, 0))\n        return np.pad(coord_array, pad_width, mode='constant', constant_values=0)\n    else:\n        return coord_array\n\ny_train_pad = np.array([pad_coordinates(arr, max_len) for arr in y_train])\ny_valid_pad = np.array([pad_coordinates(arr, max_len) for arr in y_valid])\n\nprint(\"Any NaN in y_train_pad?\", np.isnan(y_train_pad).any())\nprint(\"X_train_pad shape:\", X_train_pad.shape)\nprint(\"y_train_pad shape:\", y_train_pad.shape)\n\n# Create and cache tf.data datasets\nbatch_size = 16\ntrain_dataset = tf.data.Dataset.from_tensor_slices((X_train_pad, y_train_pad)) \\\n                              .cache().shuffle(1000).batch(batch_size).prefetch(tf.data.AUTOTUNE)\nval_dataset = tf.data.Dataset.from_tensor_slices((X_valid_pad, y_valid_pad)) \\\n                            .cache().batch(batch_size).prefetch(tf.data.AUTOTUNE)\n\n# 4. Positional Encoding Layer\nclass PositionalEncoding(tf.keras.layers.Layer):\n    def __init__(self, max_len, d_model, **kwargs):\n        super(PositionalEncoding, self).__init__(**kwargs)\n        self.max_len = max_len\n        self.d_model = d_model\n        self.pos_encoding = self.positional_encoding(max_len, d_model)\n    \n    def get_config(self):\n        config = super(PositionalEncoding, self).get_config()\n        config.update({\"max_len\": self.max_len, \"d_model\": self.d_model})\n        return config\n    \n    def positional_encoding(self, max_len, d_model):\n        angle_rads = self.get_angles(np.arange(max_len)[:, np.newaxis],\n                                     np.arange(d_model)[np.newaxis, :],\n                                     d_model)\n        # apply sin to even indices in the array; 2i\n        angle_rads[:, 0::2] = np.sin(angle_rads[:, 0::2])\n        # apply cos to odd indices in the array; 2i+1\n        angle_rads[:, 1::2] = np.cos(angle_rads[:, 1::2])\n        pos_encoding = angle_rads[np.newaxis, ...]\n        return tf.cast(pos_encoding, dtype=tf.float32)\n    \n    def get_angles(self, pos, i, d_model):\n        angle_rates = 1 / np.power(10000, (2 * (i//2)) / np.float32(d_model))\n        return pos * angle_rates\n    \n    def call(self, inputs):\n        seq_len = tf.shape(inputs)[1]\n        # Cast the positional encoding to the same dtype as the inputs\n        pos_encoding = tf.cast(self.pos_encoding, inputs.dtype)\n        return inputs + pos_encoding[:, :seq_len, :]\n    \n    def compute_output_shape(self, input_shape):\n        return input_shape\n\n# 5. Custom Windowed Self-Attention Layer (Optimized)\nclass WindowedSelfAttention(tf.keras.layers.Layer):\n    def __init__(self, window_size, num_heads, key_dim, **kwargs):\n        super(WindowedSelfAttention, self).__init__(**kwargs)\n        self.window_size = window_size\n        self.num_heads = num_heads\n        self.key_dim = key_dim\n        self.mha = tf.keras.layers.MultiHeadAttention(num_heads=num_heads, key_dim=key_dim)\n    \n    def call(self, inputs):\n        # Use static shape information (max_len is fixed)\n        input_shape = inputs.shape  # (batch, max_len, channels)\n        seq_len_static = input_shape[1]\n        channels_static = input_shape[-1]\n        seq_len = seq_len_static  # fixed = max_len\n        remainder = seq_len % self.window_size\n        if remainder != 0:\n            pad_len = self.window_size - remainder\n            padded_inputs = tf.pad(inputs, [[0,0], [0, pad_len], [0,0]])\n            new_seq_len = seq_len + pad_len\n        else:\n            padded_inputs = inputs\n            new_seq_len = seq_len\n        num_windows = new_seq_len // self.window_size\n        channels = channels_static if channels_static is not None else tf.shape(inputs)[-1]\n        windows = tf.reshape(padded_inputs, (-1, num_windows, self.window_size, channels))\n        batch_size = tf.shape(windows)[0]\n        windows_reshaped = tf.reshape(windows, (batch_size * num_windows, self.window_size, channels))\n        attn_output = self.mha(windows_reshaped, windows_reshaped)\n        attn_windows = tf.reshape(attn_output, (batch_size, num_windows, self.window_size, channels))\n        output = tf.reshape(attn_windows, (batch_size, new_seq_len, channels))\n        output = output[:, :seq_len, :]\n        return output\n\n    def compute_output_shape(self, input_shape):\n        return input_shape\n\n# 6. Advanced Model Building with Positional Encoding & Layer Normalization\nvocab_size = max(nucleotide_map.values()) + 1  # +1 for padding token\nembedding_dim = 32  # d_model for positional encoding\nnum_filters = 128\nkernel_size = 3\ndrop_rate = 0.3\n\ndef residual_block(x, filters, kernel_size, dropout_rate, block_name, dilation_rate=1):\n    shortcut = x\n    x = Conv1D(filters, kernel_size, padding='same', activation='relu',\n               dilation_rate=dilation_rate, name=f'{block_name}_conv1')(x)\n    x = LayerNormalization(name=f'{block_name}_ln1')(x)\n    x = Dropout(dropout_rate, name=f'{block_name}_drop1')(x)\n    x = Conv1D(filters, kernel_size, padding='same', activation='linear',\n               dilation_rate=dilation_rate, name=f'{block_name}_conv2')(x)\n    x = LayerNormalization(name=f'{block_name}_ln2')(x)\n    x = Add(name=f'{block_name}_add')([shortcut, x])\n    x = Activation('relu', name=f'{block_name}_out')(x)\n    return x\n\ninput_seq = Input(shape=(max_len,), name='input_seq')\n# Embedding + Positional Encoding\nx = Embedding(input_dim=vocab_size, output_dim=embedding_dim, mask_zero=True, name='embedding')(input_seq)\nx = PositionalEncoding(max_len=max_len, d_model=embedding_dim, name='pos_encoding')(x)\n# Projection to higher dimension\nx = Conv1D(num_filters, 1, padding='same', activation='relu', name='proj_conv')(x)\n# Residual blocks\nx = residual_block(x, num_filters, kernel_size, drop_rate, block_name='resblock1')\nx = residual_block(x, num_filters, kernel_size, drop_rate, block_name='resblock2', dilation_rate=2)\n# Windowed Self-Attention\nwindow_size = 256\nx = WindowedSelfAttention(window_size=window_size, num_heads=4, key_dim=32, name='window_attn')(x)\nx = residual_block(x, num_filters, kernel_size, drop_rate, block_name='resblock3')\n# Additional convolution for further feature extraction\nx = Conv1D(num_filters, kernel_size, padding='same', activation='relu', name='conv_final')(x)\nx = LayerNormalization(name='ln_final')(x)\nx = Dropout(drop_rate, name='drop_final')(x)\n# Output: Predict 3 coordinates per residue (x, y, z)\noutput_coords = Conv1D(3, 1, padding='same', activation='linear', name='predicted_coords')(x)\n\n# Cosine decay learning rate schedule\nlr_schedule = tf.keras.optimizers.schedules.CosineDecay(initial_learning_rate=1e-4, decay_steps=1000)\noptimizer = tf.keras.optimizers.Adam(learning_rate=lr_schedule)\n\nmodel = Model(inputs=input_seq, outputs=output_coords)\nmodel.compile(optimizer=optimizer, loss='mse')\nmodel.summary()\n\n# 7. Model Training with Optimized Callbacks\nearly_stop = EarlyStopping(monitor='val_loss', patience=7, restore_best_weights=True, verbose=1)\ncheckpoint = ModelCheckpoint(\"best_model.keras\", monitor='val_loss', save_best_only=True, verbose=1)\n\nhistory = model.fit(train_dataset, validation_data=val_dataset,\n                    epochs=100,\n                    callbacks=[early_stop, checkpoint],\n                    verbose=1)\n\nplt.figure(figsize=(8, 5))\nplt.plot(history.history['loss'], label='Train Loss')\nplt.plot(history.history['val_loss'], label='Validation Loss')\nplt.xlabel(\"Epoch\")\nplt.ylabel(\"MSE Loss\")\nplt.title(\"Optimized Model Training vs. Validation Loss\")\nplt.legend()\nplt.show()\n\n# 8. Generating Predictions with Monte Carlo Dropout\nX_test = test_sequences['encoded'].tolist()\nX_test_pad = pad_sequences(X_test, maxlen=max_len, padding='post', value=0)\n\nnum_MC = 5\nmc_predictions = []\nfor i in range(num_MC):\n    preds = model(X_test_pad, training=True).numpy()\n    mc_predictions.append(preds)\n\n# 9. Building the Submission File\nsubmission_rows = []\nfor idx, row in test_sequences.iterrows():\n    target_id = row['target_id']\n    seq_encoded = row['encoded']\n    seq_length = len(seq_encoded)\n    drop_preds = [mc_predictions[m][idx, :seq_length, :] for m in range(num_MC)]\n    for i in range(seq_length):\n        coords = [drop_preds[m][i] for m in range(num_MC)]\n        row_dict = {\n            'ID': f\"{target_id}_{i+1}\",\n            'resname': row['sequence'][i],\n            'resid': i+1\n        }\n        for j in range(num_MC):\n            row_dict[f\"x_{j+1}\"] = coords[j][0]\n            row_dict[f\"y_{j+1}\"] = coords[j][1]\n            row_dict[f\"z_{j+1}\"] = coords[j][2]\n        submission_rows.append(row_dict)\n\nsubmission_df = pd.DataFrame(submission_rows)\nprint(\"Submission DataFrame shape:\", submission_df.shape)\nprint(submission_df.head(10))\n\nsubmission_df.to_csv(\"submission.csv\", index=False)\nprint(\"Submission file saved as submission.csv\")\n","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null}]}