data_generation_example.py

#!/usr/bin/env python3
"""
Example script demonstrating the data generation module.

This script shows how to use the DataGenerator class to:
1. Generate synthetic datasets
2. Download benchmark datasets
3. Preprocess data
4. Convert to PyTorch Geometric format
"""

import sys
from pathlib import Path

# Add src to path for imports
sys.path.insert(0, str(Path(__file__).parent.parent / "src"))

from topogeonet.data import DataGenerator, generate_synthetic_dataset


def main():
    """Main example function."""
    print("🚀 TopoGeoNet Data Generation Example")
    print("=" * 50)
    
    # Initialize data generator
    generator = DataGenerator(
        data_dir="./example_data",
        cache_dir="./example_cache"
    )
    
    # Example 1: Generate synthetic graph dataset
    print("\n📊 Example 1: Generating Synthetic Graph Dataset")
    print("-" * 40)
    
    graphs, labels = generator.generate_synthetic_graphs(
        num_graphs=100,
        nodes_per_graph=(15, 30),
        feature_dim=16,
        edge_prob=0.2,
        graph_types=['random', 'community', 'scale_free'],
        save_path="./example_data/synthetic_graphs.pkl"
    )
    
    print(f"✅ Generated {len(graphs)} graphs with {len(set(labels))} classes")
    print(f"   Graph sizes: {[g['num_nodes'] for g in graphs[:5]]}...")
    print(f"   Labels distribution: {dict(zip(*np.unique(labels, return_counts=True)))}")
    
    # Example 2: Generate synthetic point cloud dataset
    print("\n🎯 Example 2: Generating Synthetic Point Cloud Dataset")
    print("-" * 40)
    
    point_clouds, pc_labels = generator.generate_geometric_point_clouds(
        num_samples=200,
        num_points=(100, 300),
        dimensions=3,
        shapes=['sphere', 'cube', 'cylinder', 'torus'],
        noise_level=0.02,
        save_path="./example_data/synthetic_point_clouds.pkl"
    )
    
    print(f"✅ Generated {len(point_clouds)} point clouds with {len(set(pc_labels))} shapes")
    print(f"   Point cloud sizes: {[pc.shape[0] for pc in point_clouds[:5]]}...")
    print(f"   Shape distribution: {dict(zip(*np.unique(pc_labels, return_counts=True)))}")
    
    # Example 3: Preprocess and convert to PyG format
    print("\n🔄 Example 3: Preprocessing and Converting to PyG Format")
    print("-" * 40)
    
    # Preprocess graph dataset
    preprocessing_config = {
        'normalize_features': True,
        'normalize_graph': True,
        'add_self_loops': True
    }
    
    processed_data = generator.preprocess_dataset(
        raw_data_path="./example_data/synthetic_graphs.pkl",
        dataset_type='graph',
        preprocessing_config=preprocessing_config,
        output_path="./example_data/processed_graphs.pkl"
    )
    
    print("✅ Graph dataset preprocessed")
    
    # Convert to PyTorch Geometric format
    try:
        pyg_path = generator.to_pytorch_geometric(
            data=processed_data,
            dataset_type='graph',
            target_dir="./example_data/graphs_pyg"
        )
        print(f"✅ Converted to PyG format: {pyg_path}")
    except ImportError as e:
        print(f"⚠️  PyG conversion skipped: {e}")
    
    # Example 4: Download benchmark datasets
    print("\n📥 Example 4: Downloading Benchmark Datasets")
    print("-" * 40)
    
    try:
        benchmark_paths = generator.download_benchmark_datasets()
        print(f"✅ Downloaded {len(benchmark_paths)} benchmark datasets:")
        for name, path in benchmark_paths.items():
            print(f"   - {name}: {path}")
    except Exception as e:
        print(f"⚠️  Benchmark download failed: {e}")
    
    # Example 5: Using convenience functions
    print("\n🎉 Example 5: Using Convenience Functions")
    print("-" * 40)
    
    # Generate dataset using convenience function
    try:
        simple_graphs, simple_labels = generate_synthetic_dataset(
            'graph',
            num_graphs=50,
            nodes_per_graph=20,
            feature_dim=8
        )
        print(f"✅ Generated {len(simple_graphs)} graphs using convenience function")
    except Exception as e:
        print(f"⚠️  Convenience function failed: {e}")
    
    # Example 6: Dataset information and management
    print("\n📋 Example 6: Dataset Information and Management")
    print("-" * 40)
    
    # List available datasets
    available_datasets = generator.list_available_datasets()
    print(f"📁 Available datasets: {available_datasets}")
    
    # Get dataset info
    if available_datasets:
        dataset_name = available_datasets[0]
        info = generator.get_dataset_info(generator.get_dataset_path(dataset_name))
        print(f"📊 Dataset '{dataset_name}' info:")
        print(f"   - Path: {info['path']}")
        print(f"   - Size: {info['size']} bytes")
        print(f"   - Type: {info['type']}")
    
    # Cleanup cache
    removed_files = generator.cleanup_cache()
    print(f"🧹 Cleaned up {removed_files} cache files")
    
    print("\n🎯 Data Generation Examples Completed!")
    print("=" * 50)
    print("Generated data is saved in './example_data/' directory")
    print("Check the files to see the results!")


if __name__ == "__main__":
    main()