CustomGridGenerator.py

Go to the documentation of this file.

# This file is part of the ExaHyPE2 project. For conditions of distribution and
# use, please see the copyright notice at www.peano-framework.org
 
import os
import yaml
import numpy as np
import re
 
from ..Tree import Tree
 
"""
Custom Grid Generator for Peano Load Balancing API
 
This module provides functionality to create tree structures from YAML files generated by 
the `generate_tree_yaml()` function. It serves as an input generator that enables reading 
previously saved tree structures for analysis, modification, or further processing.
 
## Purpose and Context
 
The CustomGridGenerator is designed to complete the round-trip workflow of the offline 
load balancing API. Instead of starting from scratch or reading complex grid files, this 
module enables:
 
1. **Tree Reconstruction**: Read previously saved tree structures from YAML files
2. **Workflow Continuity**: Continue analysis from previous splitting results
3. **Configuration Management**: Load pre-defined tree configurations for testing
4. **Result Verification**: Re-import trees for validation and comparison studies
 
## Integration with Load Balancing API
 
This module integrates with other components of the load balancing API:
 
- **Tree Class**: Reconstructs Tree objects from YAML data (see Tree.py)
- **OfflineDecomposition**: Can be integrated into the main script for YAML-based input
- **SplitTrees**: Provides tree data for further splitting strategy application
- **TreeVisualizer**: Enables visualization of reconstructed tree structures
 
## YAML Format Support
 
The reader supports YAML files generated by `generate_tree_yaml()` with the following structure:
 
- **Tree Metadata**: Name, dimensions, domain boundaries, and leaf counts
- **Split Summary**: Distribution statistics across subtrees
- **Tree Structure**: Complete spatial arrays for each refinement level with (weight, subtree_id) tuples
 
## Workflow Integration
 
This module fits into extended offline decomposition workflows:
 
1. **Initial Processing**: Original mesh construction and optimization
2. **YAML Export**: Tree structure saved using `generate_tree_yaml()`
3. **Configuration Storage**: YAML files stored for later use
4. **Tree Reconstruction**: CustomGridGenerator recreates Tree objects
5. **Further Analysis**: Additional splitting, visualization, or modification
6. **Iterative Optimization**: Multiple rounds of analysis and refinement
 
## Usage Examples
 
```python
# Read from existing YAML file
tree = CustomGridGenerator("MyCustomTree", "path/to/tree_structure.yaml")
 
# Apply additional splitting
from strategies.SplitTrees import split_tree
split_tree(tree, number_of_subtrees=2)
 
# Compare with original structure
from visualize.TreeVisualizer import TreeVisualizer
TreeVisualizer(tree, dimensions=tree._dimension, filename="extended_splits")
```
 
## Comparison with Other Input Methods
 
| Aspect | CustomGridGenerator | RegularGridGenerator | GridPatchFileReader |
|--------|---------------------|----------------------|-------------------|
| **Data Source** | YAML configuration files | Artificially generated | Real simulation checkpoints |
| **Structure** | Previously processed trees | Uniform regular grids | Adaptive mesh refinement |
| **Flexibility** | User-defined configurations | Predictable patterns | Simulation-dependent |
| **Use Case** | Configuration management | Testing and benchmarking | Real-world optimization |
 
## Integration Potential
 
The function can be integrated into the main `OfflineDecomposition.py` script by adding:
 
```bash
# Process existing YAML configuration
python3 OfflineDecomposition.py -y /path/to/config.yaml -p
 
# Apply additional splits to loaded configuration
python3 OfflineDecomposition.py -y /path/to/config.yaml -s 2 -p
```
 
## File Format Validation
 
The function includes comprehensive validation for:
- YAML syntax and structure
- Required metadata fields
- Tuple format consistency
- Dimension compatibility
- Level structure integrity
 
See the load balancing API documentation for complete workflow integration details.
"""
 
def parse_tuple_string(tuple_str):
    """
    Parse a string representation of a (weight, subtree_id) tuple.
    
    Supports various formats:
    - "(1,0)" - standard format
    - "(1, 0)" - with spaces
    - "1,0" - without parentheses
    
    Args:
        tuple_str (str): String representation of the tuple
        
    Returns:
        tuple: (weight, subtree_id) as integers
        
    Raises:
        ValueError: If the string cannot be parsed as a valid tuple
    """
    # Remove parentheses and whitespace, then split by comma
    clean_str = tuple_str.strip().strip('()').replace(' ', '')
    
    try:
        parts = clean_str.split(',')
        if len(parts) != 2:
            raise ValueError(f"Expected 2 values, got {len(parts)}")
        
        weight = int(parts[0])
        subtree_id = int(parts[1])
        
        return (weight, subtree_id)
        
    except (ValueError, IndexError) as e:
        raise ValueError(f"Cannot parse tuple from '{tuple_str}': {e}")
 
def CustomGridGenerator(name, yaml_file_path):
    """
    Generate a Tree object by reading from a YAML file created by generate_tree_yaml().
    
    This function reconstructs the complete tree structure from a YAML file, including
    all spatial data, metadata, and subtree assignments. It provides a way to reload
    previously processed trees for further analysis or modification.
    
    ## YAML Structure Processing
    
    The function expects YAML files with the structure created by `generate_tree_yaml()`:
    
    ```yaml
    title: "Complete Tree Structure from [source]"
    description: "Spatial representation of tree with (weight, subtree_id) tuples"
    version: "2.0"
    tree_info:
      name: [tree_name]
      dimension: [2 or 3]
      max_level: [number_of_levels]
      domain_offset: [x, y] or [x, y, z]
      domain_size: [width, height] or [width, height, depth]
      total_leaf_count: [total_leaves]
      subtree_count: [number_of_subtrees]
    split_summary:
      subtree_distribution: [count_per_subtree]
    tree_structure:
      - level: 1
        size: 3
        data:
          0: "(1,0), (1,0), (1,1)"
          ...
    ```
    
    ## Data Reconstruction
    
    The function rebuilds the tree arrays by:
    1. Parsing metadata to determine tree structure
    2. Creating numpy arrays for each refinement level
    3. Parsing tuple strings to extract (weight, subtree_id) values
    4. Populating arrays according to spatial indexing
    5. Updating tree statistics and metadata
    
    ## Error Handling
    
    Comprehensive validation includes:
    - File existence and readability
    - YAML syntax validation
    - Required field presence
    - Data type consistency
    - Tuple format validation
    - Dimensional compatibility
    
    Args:
        name (str): Descriptive name for the loaded tree (overrides YAML name if desired)
        yaml_file_path (str): Path to the YAML file containing tree structure
        
    Returns:
        Tree: Reconstructed Tree object with complete structure
        
    Raises:
        FileNotFoundError: If YAML file doesn't exist or isn't readable
        yaml.YAMLError: If YAML file has syntax errors
        KeyError: If required fields are missing from YAML structure
        ValueError: If data formats are inconsistent or invalid
        
    Example:
        ```python
        # Load a previously saved tree structure
        tree = CustomGridGenerator(
            name="ReloadedTree",
            yaml_file_path="TreeStructure_l4_s3_split.yaml"
        )
        
        print(f"Loaded tree: {tree._name}")
        print(f"Dimensions: {tree._dimension}D")
        print(f"Levels: {tree._max_level}")
        print(f"Leaf count: {tree._leaf_count}")
        print(f"Subtrees: {tree._subtree_count}")
        
        # Apply additional processing
        from strategies.SplitTrees import split_tree
        split_tree(tree, number_of_subtrees=2)
        ```
        
    Note:
        - The function preserves all original spatial relationships and metadata
        - Tree name can be overridden with the name parameter
        - Large YAML files may require significant memory during reconstruction
        - The reconstructed tree is functionally identical to the original
    """
    
    print(f"Loading custom tree '{name}' from YAML file: {yaml_file_path}")
    
    # Validate file existence
    if not os.path.exists(yaml_file_path):
        raise FileNotFoundError(f"YAML file not found: {yaml_file_path}")
    
    if not os.path.isfile(yaml_file_path):
        raise ValueError(f"Path is not a file: {yaml_file_path}")
    
    # Load and parse YAML file
    try:
        with open(yaml_file_path, 'r') as file:
            yaml_data = yaml.safe_load(file)
            print(f"Successfully loaded YAML file")
    except yaml.YAMLError as e:
        raise yaml.YAMLError(f"Error parsing YAML file: {e}")
    except IOError as e:
        raise IOError(f"Error reading YAML file: {e}")
    
    # Validate required YAML structure
    required_fields = ['tree_info', 'tree_structure']
    for field in required_fields:
        if field not in yaml_data:
            raise KeyError(f"Required field '{field}' missing from YAML file")
    
    # Extract tree metadata
    tree_info = yaml_data['tree_info']
    required_tree_fields = ['dimension', 'max_level', 'domain_offset', 'domain_size']
    for field in required_tree_fields:
        if field not in tree_info:
            raise KeyError(f"Required tree_info field '{field}' missing from YAML file")
    
    dimension = tree_info['dimension']
    max_level = tree_info['max_level']
    domain_offset = tree_info['domain_offset']
    domain_size = tree_info['domain_size']
    
    # Validate dimension
    if dimension not in [2, 3]:
        raise ValueError(f"Invalid dimension: {dimension}. Must be 2 or 3")
    
    # Validate domain information
    expected_coords = dimension
    if len(domain_offset) != expected_coords or len(domain_size) != expected_coords:
        raise ValueError(f"Domain offset and size must have {expected_coords} coordinates for {dimension}D")
    
    print(f"  Dimension: {dimension}D")
    print(f"  Max level: {max_level}")
    print(f"  Domain offset: {domain_offset}")
    print(f"  Domain size: {domain_size}")
    
    # Create Tree object
    peano_tree = Tree(name, dimension, domain_offset, domain_size)
    
    # Process tree structure levels
    tree_structure = yaml_data['tree_structure']
    
    if len(tree_structure) != max_level:
        raise ValueError(f"Expected {max_level} levels, got {len(tree_structure)} in tree_structure")
    
    # Initialize arrays for all levels
    artificial_arrays = []
    
    for level_data in tree_structure:
        # Validate level structure
        if 'level' not in level_data or 'size' not in level_data or 'data' not in level_data:
            raise KeyError("Each level must have 'level', 'size', and 'data' fields")
        
        level_num = level_data['level']
        size = level_data['size']
        data = level_data['data']
        
        # Validate level number and size
        expected_level = len(artificial_arrays) + 1
        expected_size = 3 ** expected_level
        
        if level_num != expected_level:
            print(f"Warning: Level number {level_num} doesn't match expected {expected_level}")
        
        if size != expected_size:
            raise ValueError(f"Level {level_num}: expected size {expected_size}, got {size}")
        
        print(f"  Processing level {level_num}: {size}^{dimension} cells")
        
        # Create array for this level
        if dimension == 2:
            array = np.empty((size, size), dtype=object)
            
            # Parse 2D data
            for row_idx in range(size):
                if row_idx not in data:
                    raise KeyError(f"Level {level_num}: missing row {row_idx}")
                
                row_str = data[row_idx]
                # Split by ", " to separate complete tuples
                tuple_strings = [s.strip() for s in row_str.split(', ')]
                
                if len(tuple_strings) != size:
                    raise ValueError(f"Level {level_num}, row {row_idx}: expected {size} tuples, got {len(tuple_strings)}")
                
                for col_idx, tuple_str in enumerate(tuple_strings):
                    try:
                        weight, subtree_id = parse_tuple_string(tuple_str)
                        array[row_idx][col_idx] = (weight, subtree_id)
                    except ValueError as e:
                        raise ValueError(f"Level {level_num}, row {row_idx}, col {col_idx}: {e}")
        
        elif dimension == 3:
            array = np.empty((size, size, size), dtype=object)
            
            # Parse 3D data (planes and rows)
            # YAML structure: data[plane_a][row_b] = "tuple_c0, tuple_c1, tuple_c2"
            # Array structure: array[a][b][c] = (weight, subtree_id)
            for plane_idx in range(size):
                if plane_idx not in data:
                    raise KeyError(f"Level {level_num}: missing plane {plane_idx}")
                
                plane_data = data[plane_idx]
                if not isinstance(plane_data, dict):
                    raise ValueError(f"Level {level_num}, plane {plane_idx}: expected dictionary, got {type(plane_data)}")
                
                for row_idx in range(size):
                    if row_idx not in plane_data:
                        raise KeyError(f"Level {level_num}, plane {plane_idx}: missing row {row_idx}")
                    
                    row_str = plane_data[row_idx]
                    # Split by ", " to separate complete tuples
                    tuple_strings = [s.strip() for s in row_str.split(', ')]
                    
                    if len(tuple_strings) != size:
                        raise ValueError(f"Level {level_num}, plane {plane_idx}, row {row_idx}: expected {size} tuples, got {len(tuple_strings)}")
                    
                    for col_idx, tuple_str in enumerate(tuple_strings):
                        try:
                            weight, subtree_id = parse_tuple_string(tuple_str)
                            # Map YAML indices to array indices: data[a][b] -> array[a][b][c]
                            array[plane_idx][row_idx][col_idx] = (weight, subtree_id)
                        except ValueError as e:
                            raise ValueError(f"Level {level_num}, plane {plane_idx}, row {row_idx}, col {col_idx}: {e}")
        
        artificial_arrays.append(array)
    
    # Determine the actual subtree count from the data before populating tree
    max_subtree_id = 0
    for array in artificial_arrays:
        if dimension == 2:
            for i in range(array.shape[0]):
                for j in range(array.shape[1]):
                    weight, subtree_id = array[i][j]
                    if weight > 0:  # Only count existing cells
                        max_subtree_id = max(max_subtree_id, subtree_id)
        elif dimension == 3:
            for i in range(array.shape[0]):
                for j in range(array.shape[1]):
                    for k in range(array.shape[2]):
                        weight, subtree_id = array[i][j][k]
                        if weight > 0:  # Only count existing cells
                            max_subtree_id = max(max_subtree_id, subtree_id)
    
    # Set the correct subtree count and max level first
    peano_tree._subtree_count = max_subtree_id + 1
    peano_tree._max_level = max_level
    
    # Directly set the tree arrays and update metadata manually
    # (avoiding read_in_tree_from_artificial which hardcodes subtree_count = 1)
    peano_tree._tree_arrays = artificial_arrays
    peano_tree.update_split_pattern()
    
    if False:  # weighted is False for our case
        peano_tree.update_weight()
    
    print(f"Custom tree generation completed successfully")
    print(f"  Tree: {peano_tree._name}")
    print(f"  Leaf count: {peano_tree._leaf_count}")
    print(f"  Subtree count: {peano_tree._subtree_count}")
    
    # Validate against YAML metadata if available
    if 'total_leaf_count' in tree_info:
        expected_leaf_count = tree_info['total_leaf_count']
        if peano_tree._leaf_count != expected_leaf_count:
            print(f"Warning: Reconstructed leaf count ({peano_tree._leaf_count}) doesn't match YAML metadata ({expected_leaf_count})")
    
    if 'subtree_count' in tree_info:
        expected_subtree_count = tree_info['subtree_count']
        if peano_tree._subtree_count != expected_subtree_count:
            print(f"Warning: Reconstructed subtree count ({peano_tree._subtree_count}) doesn't match YAML metadata ({expected_subtree_count})")
    
    return peano_tree