{"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":118765,"databundleVersionId":15231210,"sourceType":"competition"}],"dockerImageVersionId":31239,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nfrom scipy.ndimage import gaussian_filter1d\nfrom scipy.spatial.transform import Rotation\nfrom sklearn.decomposition import PCA\nimport warnings\nwarnings.filterwarnings('ignore')\n\n# ===============================\n# ENHANCED RNA 3D Model\n# ===============================\n\nclass EnhancedRNA3DModel:\n    def __init__(self):\n        # Standard A-form helix parameters\n        self.rise = 2.81  # Angstroms per base pair\n        self.twist = 32.7  # degrees per base pair\n        self.radius = 10.0  # Angstroms\n        self.phase_offset = np.pi / 2  # Start at 90 degrees\n        \n        # Base-specific parameters from experimental data\n        self.base_radii = {\n            'A': 9.8, 'U': 9.6, 'T': 9.6, 'G': 9.9, 'C': 9.7\n        }\n        self.base_rise_mod = {\n            'A': 0.98, 'U': 1.02, 'T': 1.02, 'G': 0.96, 'C': 1.04\n        }\n        self.base_twist_mod = {\n            'A': 1.02, 'U': 0.98, 'T': 0.98, 'G': 1.03, 'C': 0.97\n        }\n        \n        # Base-pairing effects\n        self.pairing_strength = {\n            'AU': 0.9, 'UA': 0.9, 'AT': 0.9, 'TA': 0.9,\n            'GC': 1.1, 'CG': 1.1, 'GU': 0.85, 'UG': 0.85\n        }\n        \n        # Stacking energies (relative)\n        self.stacking_energy = {\n            'AA': 1.0, 'AU': 1.2, 'AG': 1.1, 'AC': 0.9,\n            'UA': 1.2, 'UU': 1.0, 'UG': 1.3, 'UC': 1.1,\n            'GA': 1.1, 'GU': 1.3, 'GG': 1.4, 'GC': 1.2,\n            'CA': 0.9, 'CU': 1.1, 'CG': 1.2, 'CC': 1.0\n        }\n    \n    def calculate_base_params(self, sequence):\n        \"\"\"Calculate sequence-dependent parameters\"\"\"\n        n = len(sequence)\n        radii = np.zeros(n)\n        rises = np.zeros(n)\n        twists = np.zeros(n)\n        \n        for i, nt in enumerate(sequence):\n            radii[i] = self.base_radii.get(nt, 10.0)\n            rises[i] = self.rise * self.base_rise_mod.get(nt, 1.0)\n            twists[i] = np.deg2rad(self.twist * self.base_twist_mod.get(nt, 1.0))\n            \n            # Consider nearest neighbors\n            if i > 0:\n                pair = sequence[i-1] + nt\n                if pair in self.pairing_strength:\n                    radii[i] *= self.pairing_strength[pair]\n                    radii[i-1] *= self.pairing_strength[pair]\n                \n                if pair in self.stacking_energy:\n                    energy_factor = self.stacking_energy[pair] / 1.2  # normalize\n                    rises[i] *= (0.98 + 0.04 * energy_factor)\n        \n        return radii, rises, twists\n    \n    def build_helix_backbone(self, sequence):\n        \"\"\"Build initial helical backbone with sequence effects\"\"\"\n        n = len(sequence)\n        radii, rises, twists = self.calculate_base_params(sequence)\n        \n        coords = np.zeros((n, 3))\n        cumulative_angle = 0\n        \n        for i in range(n):\n            # Add sequence-dependent variations\n            radius_variation = 0.3 * np.sin(i * 0.5)  # minor periodic variation\n            angle_variation = 0.1 * np.cos(i * 0.3)\n            \n            angle = cumulative_angle + self.phase_offset + angle_variation\n            r = radii[i] * (1.0 + 0.05 * radius_variation)\n            \n            coords[i, 0] = r * np.cos(angle)\n            coords[i, 1] = r * np.sin(angle)\n            coords[i, 2] = i * np.mean(rises[:i+1]) if i > 0 else 0\n            \n            cumulative_angle += twists[i]\n        \n        return coords\n    \n    def add_thermal_fluctuations(self, coords, temperature=300):\n        \"\"\"Add realistic thermal fluctuations\"\"\"\n        n = coords.shape[0]\n        \n        # Low-frequency modes (global bending)\n        lf_noise = np.random.normal(0, 0.1, size=(n, 3))\n        for dim in range(3):\n            lf_noise[:, dim] = gaussian_filter1d(lf_noise[:, dim], sigma=3.0)\n        \n        # High-frequency modes (local dynamics)\n        hf_noise = np.random.normal(0, 0.05, size=(n, 3))\n        for dim in range(3):\n            hf_noise[:, dim] = gaussian_filter1d(hf_noise[:, dim], sigma=0.5)\n        \n        # Scale by temperature factor\n        temp_factor = np.sqrt(temperature / 300)\n        coords += (lf_noise + hf_noise) * temp_factor\n        \n        return coords\n    \n    def apply_global_bending(self, coords, bend_angle=0.0, bend_axis=None):\n        \"\"\"Apply global bending to mimic flexibility\"\"\"\n        if bend_axis is None:\n            bend_axis = np.random.randn(3)\n            bend_axis /= np.linalg.norm(bend_axis)\n        \n        n = coords.shape[0]\n        \n        # Create bending profile (more bend at ends for flexibility)\n        t = np.linspace(-1, 1, n)\n        bend_profile = bend_angle * (1 - t**2)  # parabolic profile\n        \n        # Apply progressive rotation\n        for i in range(n):\n            if i > 0:\n                rotation = Rotation.from_rotvec(bend_axis * bend_profile[i] / n)\n                coords[i:] = rotation.apply(coords[i:] - coords[i]) + coords[i]\n        \n        return coords\n    \n    def add_solvent_effects(self, coords, sequence):\n        \"\"\"Simulate solvent accessibility effects\"\"\"\n        n = len(sequence)\n        \n        # Hydrophobic bases (A, G) tend to be more buried\n        burial_factors = {'A': 0.9, 'G': 0.92, 'C': 1.0, 'U': 1.05, 'T': 1.05}\n        \n        for i, nt in enumerate(sequence):\n            factor = burial_factors.get(nt, 1.0)\n            # Apply radial compression for buried bases\n            r = np.linalg.norm(coords[i, :2])\n            if r > 0:\n                coords[i, :2] *= (0.95 + 0.05 * factor)\n        \n        return coords\n    \n    def optimize_geometry(self, coords):\n        \"\"\"Simple geometric optimization\"\"\"\n        # Center the structure\n        coords -= coords.mean(axis=0)\n        \n        # Align principal axes\n        pca = PCA(n_components=3)\n        coords_pca = pca.fit_transform(coords)\n        \n        # Reorient for consistency\n        # Make longest dimension along z-axis\n        variances = np.var(coords_pca, axis=0)\n        z_idx = np.argmax(variances)\n        \n        # Rearrange axes if needed\n        if z_idx != 2:\n            temp = coords_pca[:, z_idx].copy()\n            coords_pca[:, z_idx] = coords_pca[:, 2]\n            coords_pca[:, 2] = temp\n        \n        return coords_pca\n    \n    def generate_ensemble(self, sequence, n_models=5):\n        \"\"\"Generate ensemble of structures\"\"\"\n        ensemble = []\n        \n        for model_idx in range(n_models):\n            # Base helix\n            coords = self.build_helix_backbone(sequence)\n            \n            # Model-specific variations\n            if model_idx > 0:\n                # Vary parameters based on model index\n                bend_angle = np.random.uniform(0, np.deg2rad(10 * model_idx))\n                bend_axis = np.random.randn(3)\n                bend_axis /= np.linalg.norm(bend_axis)\n                \n                coords = self.apply_global_bending(coords, bend_angle, bend_axis)\n                \n                # Add thermal fluctuations (different temperature for each model)\n                temperature = 300 + model_idx * 50\n                coords = self.add_thermal_fluctuations(coords, temperature)\n                \n                # Solvent effects\n                coords = self.add_solvent_effects(coords, sequence)\n            \n            # Smoothing\n            for dim in range(3):\n                coords[:, dim] = gaussian_filter1d(coords[:, dim], sigma=1.0)\n            \n            # Geometry optimization\n            coords = self.optimize_geometry(coords)\n            \n            ensemble.append(coords)\n        \n        return ensemble\n    \n    def validate_geometry(self, coords):\n        \"\"\"Validate basic geometric constraints\"\"\"\n        # Check distances between consecutive bases\n        distances = np.linalg.norm(np.diff(coords, axis=0), axis=1)\n        avg_dist = np.mean(distances)\n        \n        # A-form RNA should have ~2.8-3.0 Å between consecutive bases\n        if avg_dist < 2.5 or avg_dist > 3.5:\n            # Rescale if needed\n            scale_factor = 2.8 / avg_dist\n            coords *= scale_factor\n        \n        return coords\n\n\n# ===============================\n# ENSEMBLE GENERATION WITH CONSISTENT SCALING\n# ===============================\n\nclass RNA3DEnsembleGenerator:\n    def __init__(self):\n        self.model = EnhancedRNA3DModel()\n        self.reference_scale = 1.0\n    \n    def generate_models(self, sequence, n_models=5):\n        \"\"\"Generate consistent ensemble of models\"\"\"\n        models = self.model.generate_ensemble(sequence, n_models)\n        \n        # Ensure consistent scaling across all models\n        all_coords = np.concatenate(models, axis=0)\n        global_mean = all_coords.mean(axis=0)\n        global_std = all_coords.std()\n        \n        if global_std > 0:\n            target_std = 12.0  # Reasonable scale for RNA structures\n            scale_factor = target_std / global_std\n            \n            for i in range(len(models)):\n                models[i] = (models[i] - global_mean) * scale_factor\n        \n        # Final validation\n        for i in range(len(models)):\n            models[i] = self.model.validate_geometry(models[i])\n        \n        return models\n\n\n# ===============================\n# KAGGLE SUBMISSION BUILDER\n# ===============================\n\ndef create_enhanced_submission(\n    test_csv_path=\"/kaggle/input/stanford-rna-3d-folding-2/test_sequences.csv\",\n    output_path=\"submission.csv\",\n    n_models=5\n):\n    \"\"\"\n    Create enhanced submission file for RNA 3D Folding challenge\n    \"\"\"\n    print(\"Loading test data...\")\n    test_df = pd.read_csv(test_csv_path)\n    \n    print(f\"Processing {len(test_df)} sequences...\")\n    generator = RNA3DEnsembleGenerator()\n    \n    all_rows = []\n    \n    for idx, row in test_df.iterrows():\n        target_id = row[\"target_id\"]\n        sequence = row[\"sequence\"]\n        \n        if idx % 100 == 0:\n            print(f\"Processing sequence {idx+1}/{len(test_df)}: {target_id}\")\n        \n        # Generate ensemble of models\n        models = generator.generate_models(sequence, n_models)\n        \n        # Create entries for each nucleotide\n        for nt_idx, nt in enumerate(sequence):\n            entry = {\n                \"ID\": f\"{target_id}_{nt_idx+1}\",\n                \"resname\": nt,\n                \"resid\": nt_idx + 1\n            }\n            \n            # Add coordinates from each model\n            for model_idx in range(n_models):\n                coords = models[model_idx][nt_idx]\n                entry[f\"x_{model_idx+1}\"] = coords[0]\n                entry[f\"y_{model_idx+1}\"] = coords[1]\n                entry[f\"z_{model_idx+1}\"] = coords[2]\n            \n            all_rows.append(entry)\n    \n    # Create DataFrame\n    print(\"Creating submission DataFrame...\")\n    submission_df = pd.DataFrame(all_rows)\n    \n    # Ensure correct column order\n    column_order = [\"ID\", \"resname\", \"resid\"]\n    for i in range(1, n_models + 1):\n        column_order.extend([f\"x_{i}\", f\"y_{i}\", f\"z_{i}\"])\n    \n    submission_df = submission_df[column_order]\n    \n    # Save with appropriate precision\n    print(f\"Saving to {output_path}...\")\n    submission_df.to_csv(output_path, index=False, float_format=\"%.3f\")\n    \n    print(f\"Submission created with {len(submission_df)} entries\")\n    print(f\"Shape: {submission_df.shape}\")\n    \n    # Print statistics\n    print(\"\\nSubmission Statistics:\")\n    for i in range(1, n_models + 1):\n        x_col = f\"x_{i}\"\n        mean_std = submission_df[[f\"x_{i}\", f\"y_{i}\", f\"z_{i}\"]].std().mean()\n        print(f\"Model {i}: Avg coordinate std = {mean_std:.3f}\")\n    \n    return submission_df\n\n\n# ===============================\n# QUALITY CONTROL FUNCTIONS\n# ===============================\n\ndef check_submission_quality(submission_df):\n    \"\"\"Perform quality checks on submission\"\"\"\n    print(\"\\n=== Quality Check ===\")\n    \n    # Check for NaN values\n    nan_count = submission_df.isna().sum().sum()\n    print(f\"NaN values: {nan_count}\")\n    \n    # Check coordinate ranges\n    coord_columns = [col for col in submission_df.columns if col.startswith(('x_', 'y_', 'z_'))]\n    coord_data = submission_df[coord_columns]\n    \n    print(f\"Coordinate min: {coord_data.min().min():.2f}\")\n    print(f\"Coordinate max: {coord_data.max().max():.2f}\")\n    print(f\"Coordinate mean: {coord_data.mean().mean():.2f}\")\n    print(f\"Coordinate std: {coord_data.std().mean():.2f}\")\n    \n    # Check sequential distances for first model\n    test_ids = submission_df['ID'].str.split('_').str[0].unique()[:3]\n    \n    for test_id in test_ids[:3]:  # Check first 3 sequences\n        seq_data = submission_df[submission_df['ID'].str.startswith(test_id)]\n        if len(seq_data) > 1:\n            coords = seq_data[['x_1', 'y_1', 'z_1']].values\n            distances = np.linalg.norm(np.diff(coords, axis=0), axis=1)\n            avg_dist = distances.mean()\n            print(f\"{test_id}: Avg consecutive distance = {avg_dist:.3f} Å\")\n    \n    return True\n\n\n# ===============================\n# MAIN EXECUTION\n# ===============================\n\nif __name__ == \"__main__\":\n    # Generate enhanced submission\n    submission = create_enhanced_submission(\n        test_csv_path=\"/kaggle/input/stanford-rna-3d-folding-2/test_sequences.csv\",\n        output_path=\"submission.csv\",\n        n_models=5\n    )\n    \n    # Perform quality check\n    check_submission_quality(submission)\n    \n    # Display first few rows\n    print(\"\\n=== First 10 rows of submission ===\")\n    print(submission.head(10))\n    \n    print(\"\\n✅ Enhanced submission created successfully!\")","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2026-01-09T07:02:37.595534Z","iopub.execute_input":"2026-01-09T07:02:37.595968Z","iopub.status.idle":"2026-01-09T07:02:47.065157Z","shell.execute_reply.started":"2026-01-09T07:02:37.595942Z","shell.execute_reply":"2026-01-09T07:02:47.063779Z"}},"outputs":[],"execution_count":null}]}