generate_samples.py

#!/usr/bin/env python3
"""
Generate sample VTI files for the examples directory.
"""

import numpy as np
from pathlib import Path

try:
    import vtk
    from vtk.util import numpy_support
    VTK_AVAILABLE = True
except ImportError:
    VTK_AVAILABLE = False


def create_gaussian_vti(output_path: Path, nx=30, ny=25, nz=20):
    """Create a simple Gaussian VTI file."""
    if not VTK_AVAILABLE:
        print("VTK not available")
        return False
    
    # Create image data
    image_data = vtk.vtkImageData()
    image_data.SetDimensions(nx, ny, nz)
    image_data.SetSpacing(0.1, 0.1, 0.1)
    image_data.SetOrigin(0.0, 0.0, 0.0)
    
    # Create coordinate arrays
    x = np.linspace(0, 3, nx)
    y = np.linspace(0, 2.5, ny)
    z = np.linspace(0, 2, nz)
    
    # Create meshgrid
    X, Y, Z = np.meshgrid(x, y, z, indexing='ij')
    
    # Create scalar data - 3D Gaussian
    center_x, center_y, center_z = 1.5, 1.25, 1.0
    sigma = 0.8
    
    gaussian = np.exp(-((X - center_x)**2 + (Y - center_y)**2 + (Z - center_z)**2) / (2 * sigma**2))
    
    # Create distance field
    distance = np.sqrt((X - center_x)**2 + (Y - center_y)**2 + (Z - center_z)**2)
    
    # Create simple vector field (gradient)
    vector_x = np.gradient(gaussian, axis=0)
    vector_y = np.gradient(gaussian, axis=1)
    vector_z = np.gradient(gaussian, axis=2)
    
    # Flatten arrays for VTK
    gaussian_flat = gaussian.flatten(order='F')
    distance_flat = distance.flatten(order='F')
    vector_flat = np.column_stack([
        vector_x.flatten(order='F'),
        vector_y.flatten(order='F'),
        vector_z.flatten(order='F')
    ])
    
    # Add scalar data
    gaussian_array = numpy_support.numpy_to_vtk(gaussian_flat)
    gaussian_array.SetName("temperature")
    image_data.GetPointData().AddArray(gaussian_array)
    image_data.GetPointData().SetActiveScalars("temperature")
    
    # Add distance field
    distance_array = numpy_support.numpy_to_vtk(distance_flat)
    distance_array.SetName("distance")
    image_data.GetPointData().AddArray(distance_array)
    
    # Add vector field
    vector_array = numpy_support.numpy_to_vtk(vector_flat)
    vector_array.SetName("velocity")
    image_data.GetPointData().AddArray(vector_array)
    
    # Write to file
    writer = vtk.vtkXMLImageDataWriter()
    writer.SetFileName(str(output_path))
    writer.SetInputData(image_data)
    writer.Write()
    
    return True


def create_wave_vti(output_path: Path, nx=40, ny=30, nz=25):
    """Create a wave pattern VTI file."""
    if not VTK_AVAILABLE:
        print("VTK not available")
        return False
    
    # Create image data
    image_data = vtk.vtkImageData()
    image_data.SetDimensions(nx, ny, nz)
    image_data.SetSpacing(0.05, 0.05, 0.05)
    image_data.SetOrigin(0.0, 0.0, 0.0)
    
    # Create coordinate arrays
    x = np.linspace(0, 2*np.pi, nx)
    y = np.linspace(0, 2*np.pi, ny)
    z = np.linspace(0, np.pi, nz)
    
    # Create meshgrid
    X, Y, Z = np.meshgrid(x, y, z, indexing='ij')
    
    # Create wave patterns
    wave1 = np.sin(X) * np.cos(Y) * np.sin(Z)
    wave2 = np.cos(X + Y) * np.sin(Z)
    pressure = wave1 + 0.5 * wave2
    
    # Create magnitude field
    magnitude = np.sqrt(X**2 + Y**2 + Z**2)
    
    # Flatten arrays for VTK
    pressure_flat = pressure.flatten(order='F')
    magnitude_flat = magnitude.flatten(order='F')
    
    # Add scalar data
    pressure_array = numpy_support.numpy_to_vtk(pressure_flat)
    pressure_array.SetName("pressure")
    image_data.GetPointData().AddArray(pressure_array)
    image_data.GetPointData().SetActiveScalars("pressure")
    
    # Add magnitude field
    magnitude_array = numpy_support.numpy_to_vtk(magnitude_flat)
    magnitude_array.SetName("magnitude")
    image_data.GetPointData().AddArray(magnitude_array)
    
    # Write to file
    writer = vtk.vtkXMLImageDataWriter()
    writer.SetFileName(str(output_path))
    writer.SetInputData(image_data)
    writer.Write()
    
    return True


def main():
    """Generate sample VTI files."""
    if not VTK_AVAILABLE:
        print("VTK not available. Cannot generate sample files.")
        return False
    
    # Create output directories
    examples_dir = Path("examples/sample_data")
    tests_dir = Path("tests/sample_data")
    
    examples_dir.mkdir(parents=True, exist_ok=True)
    tests_dir.mkdir(parents=True, exist_ok=True)
    
    print("Generating sample VTI files...")
    
    # Generate files
    files_created = []
    
    # Simple Gaussian (small file)
    gaussian_file = examples_dir / "gaussian_simple.vti"
    if create_gaussian_vti(gaussian_file, nx=20, ny=15, nz=12):
        files_created.append(gaussian_file)
        print(f"✅ Created: {gaussian_file}")
    
    # Larger Gaussian for tests
    gaussian_test = tests_dir / "sample.vti"
    if create_gaussian_vti(gaussian_test, nx=30, ny=25, nz=20):
        files_created.append(gaussian_test)
        print(f"✅ Created: {gaussian_test}")
    
    # Wave pattern
    wave_file = examples_dir / "wave_pattern.vti"
    if create_wave_vti(wave_file):
        files_created.append(wave_file)
        print(f"✅ Created: {wave_file}")
    
    print(f"\n🎉 Generated {len(files_created)} sample VTI files!")
    
    # Show file info
    for file_path in files_created:
        if file_path.exists():
            size_mb = file_path.stat().st_size / (1024 * 1024)
            print(f"   📁 {file_path.name}: {size_mb:.2f} MB")
    
    return True


if __name__ == "__main__":
    success = main()
    exit(0 if success else 1)