sample_data.py

#!/usr/bin/env python3
"""
Sample data generation for computational chemistry template
Provides realistic molecular data for testing and demonstration
"""

import numpy as np
import matplotlib.pyplot as plt

def generate_molecular_coordinates(molecule_name="benzene"):
    """Generate sample molecular coordinates for common molecules"""

    coordinates = {
        "benzene": {
            "atoms": ["C", "C", "C", "C", "C", "C", "H", "H", "H", "H", "H", "H"],
            "coords": np.array([
                [0.000, 1.396, 0.000],    # C1
                [1.209, 0.698, 0.000],    # C2
                [1.209, -0.698, 0.000],   # C3
                [0.000, -1.396, 0.000],   # C4
                [-1.209, -0.698, 0.000],  # C5
                [-1.209, 0.698, 0.000],   # C6
                [0.000, 2.479, 0.000],    # H1
                [2.147, 1.240, 0.000],    # H2
                [2.147, -1.240, 0.000],   # H3
                [0.000, -2.479, 0.000],   # H4
                [-2.147, -1.240, 0.000],  # H5
                [-2.147, 1.240, 0.000]    # H6
            ])
        },
        "water": {
            "atoms": ["O", "H", "H"],
            "coords": np.array([
                [0.000, 0.000, 0.119],     # O
                [0.000, 0.757, -0.477],    # H1
                [0.000, -0.757, -0.477]    # H2
            ])
        },
        "methane": {
            "atoms": ["C", "H", "H", "H", "H"],
            "coords": np.array([
                [0.000, 0.000, 0.000],     # C
                [0.629, 0.629, 0.629],     # H1
                [-0.629, -0.629, 0.629],   # H2
                [-0.629, 0.629, -0.629],   # H3
                [0.629, -0.629, -0.629]    # H4
            ])
        }
    }

    return coordinates.get(molecule_name, coordinates["benzene"])

def generate_energy_profile(npoints=50, barrier_height=25.0, reaction_energy=-15.0):
    """Generate a realistic reaction energy profile"""

    x = np.linspace(0, 6, npoints)

    # Create a double-well potential with transition states
    energy = np.zeros_like(x)

    # Initial reactant state
    energy[x <= 1] = 0

    # First transition state
    mask1 = (x > 1) & (x <= 2)
    energy[mask1] = barrier_height * np.sin(np.pi * (x[mask1] - 1))**2

    # Intermediate state
    mask2 = (x > 2) & (x <= 3)
    energy[mask2] = reaction_energy * 0.6

    # Second transition state
    mask3 = (x > 3) & (x <= 4)
    energy[mask3] = reaction_energy * 0.6 + (barrier_height * 0.8) * np.sin(np.pi * (x[mask3] - 3))**2

    # Second intermediate
    mask4 = (x > 4) & (x <= 5)
    energy[mask4] = reaction_energy * 0.3

    # Final products
    energy[x > 5] = reaction_energy

    return x, energy

def generate_md_trajectory(nframes=1000, nresidues=50):
    """Generate sample molecular dynamics trajectory data"""

    time = np.linspace(0, 100, nframes)  # 100 ns simulation

    # RMSD with realistic behavior
    rmsd_backbone = 1.0 + 0.5 * (1 - np.exp(-time/10)) + 0.2 * np.random.normal(0, 1, nframes)
    rmsd_backbone = np.maximum(rmsd_backbone, 0.5)  # Ensure positive RMSD

    # RMSF per residue
    rmsf_values = np.random.gamma(1.5, 0.8, nresidues)

    # Temperature and pressure
    temperature = 300 + 5 * np.sin(time/5) + 2 * np.random.normal(0, 1, nframes)
    pressure = 1.0 + 0.1 * np.sin(time/3) + 0.05 * np.random.normal(0, 1, nframes)

    return {
        "time": time,
        "rmsd_backbone": rmsd_backbone,
        "rmsf_values": rmsf_values,
        "temperature": temperature,
        "pressure": pressure
    }

def generate_docking_results(ncompounds=1000):
    """Generate virtual screening/docking results"""

    np.random.seed(42)  # For reproducibility

    # Binding affinities with realistic distribution
    binding_scores = np.random.gamma(2, 2, ncompounds) * -1 - 4
    binding_scores = np.clip(binding_scores, -12, -4)

    # ADMET properties
    molecular_weight = np.random.normal(350, 80, ncompounds)
    lipophilicity = np.random.normal(3.5, 1.2, ncompounds)
    polar_surface_area = np.random.normal(70, 25, ncompounds)

    # Filter for hits
    hits_mask = binding_scores <= -8.0

    return {
        "binding_scores": binding_scores,
        "molecular_weight": molecular_weight,
        "lipophilicity": lipophilicity,
        "polar_surface_area": polar_surface_area,
        "hits_mask": hits_mask,
        "n_hits": np.sum(hits_mask)
    }

if __name__ == "__main__":
    # Generate and save sample datasets
    print("Generating sample computational chemistry datasets...")

    # Molecular coordinates
    benzene = generate_molecular_coordinates("benzene")
    print(f"Benzene molecule: {len(benzene['atoms'])} atoms")

    # Energy profile
    x, energy = generate_energy_profile()
    print(f"Energy profile: {len(x)} points, barrier = {max(energy):.1f} kcal/mol")

    # MD trajectory
    md_data = generate_md_trajectory()
    print(f"MD trajectory: {len(md_data['time'])} frames")

    # Docking results
    docking_data = generate_docking_results()
    print(f"Docking results: {len(docking_data['binding_scores'])} compounds, {docking_data['n_hits']} hits")

    print("Sample data generation complete!")