generateRandomdata.py

import os
import numpy as np
import torch
import hydra
import time
import GPUtil
import random
from omegaconf import DictConfig

# =========================================================================
# 1. Reproducibility & Imports
# =========================================================================
def set_all_seeds(seed=42):
    """
    Set all possible random seeds for reproducibility.
    """
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

# --- Import Custom Models and Utilities ---
from model.carreau import carreau_yasuda_viscosity
from utils.tensors import generate_base_L_tensor, flatten_symmetric_tensors
from utils.write_sampledata1 import write_sampledata_file 
from utils.invariants import filter_admissible_region
from utils.plot_stage_progression1 import plot_stage_progression
from utils.stabletensorrandom import (
    generate_stable_sample_maxwell,
    generate_stable_sample_oldroyd,
    generate_stable_sample_ptt_exponential
)


# =========================================================================
# 2. Helper: Binning Logic (Condition Number -> Stage Tag)
# =========================================================================
def get_stage_tag(cond_val, mode="single_stage"):
    """
    Classifies a condition number into a specific folder/stage tag.
    """
    # 1. Single Stage: Broad range
    if mode == "single_stage":
        if 1.0 <= cond_val <= 2.4:
            return "1.0_2.4"
    
    # 2. Multi Stage: Granular bins
    elif mode == "multi_stage":
        # Stage 1: Foundation (Linear & Weakly Non-Linear)
        # Catches strict 1.0 and ranges up to 1.2
        if 1.0 <= cond_val < 1.2: return "1.0_1.2"
        
        # Stage 2: Transition
        if 1.2 <= cond_val < 1.4: return "1.2_1.4"

        # Stage 3: Transition
        if 1.4 <= cond_val < 1.6: return "1.4_1.6"

        # Stage 4: Transition
        if 1.6 <= cond_val < 1.8: return "1.6_1.8"
        
        # Stage 5: Challenge
        if 1.8 <= cond_val < 2.0: return "1.8_2.0"

        # Stage 6: Challenge
        if 2.0 <= cond_val < 2.2: return "2.0_2.2"
        
        # Stage 7: Edge Case
        if 2.2 <= cond_val <= 2.4: return "2.2_2.4"
        
        # Safety catch for float precision issues near 1.0
        if 0.99 <= cond_val < 1.0: return "1.0_1.2"

    return None

# =========================================================================
# 3. Main Hydra Entry Point
# =========================================================================
@hydra.main(config_path="config/data", config_name="dataConfig1", version_base=None)
def main(cfg: DictConfig) -> None:
    """
    Main RANDOM data generation routine.
    
    Logic:
    1. Define Bins.
    2. Generate GLOBAL pool of samples (to ensure natural distribution).
    3. Sort samples into Bins.
    4. Save and Analyze each Bin.
    """
    
    # ---------------------------------------------------------------------
    # A. Configure Modes and Bins
    # ---------------------------------------------------------------------
    if cfg.mode == "single_stage":
        bin_keys = ["1.0_2.4"]
        folder_mode = "single_stage"
        gen_range = (1.0, 2.4)
    elif cfg.mode == "multi_stage":
        # We use the detailed 8-bin structure for high resolution
        bin_keys = ["1.0_1.2", "1.2_1.4", "1.4_1.6", "1.6_1.8", "1.8_2.0", "2.0_2.2", "2.2_2.4"]
        folder_mode = "multi_stage"
        # --- NEW: Allow Custom Targeted Generation ---
        # If the user passes +custom_min and +custom_max, we focus only on that range.
        # Otherwise, we generate strictly 1.0 to 2.4 as usual.
        if "custom_min" in cfg and "custom_max" in cfg:
            c_min = float(cfg.custom_min)
            c_max = float(cfg.custom_max)
            gen_range = (c_min, c_max)
            print(f"🎯 TARGETED GENERATION: Focusing exclusively on Condition Numbers {gen_range}")
        else:
            gen_range = (1.0, 2.4)
    else:
        raise ValueError(f"Unknown cfg.mode: {cfg.mode}")

    print(f"=== Generating Random {cfg.constitutive_eq.replace('_', ' ').title()} | Mode: {cfg.mode} ===")
    print(f"   -> Goal: {cfg.n_samples} total samples distributed across bins.")

    set_all_seeds(cfg.seed)

    # Initialize Storage for Bins
    # Format: { tag: [ (L, Y, cond, W, resid), ... ] }
    data_bins = {k: [] for k in bin_keys}
    
    # Special lists for Carreau (which doesn't bin by stability)
    carreau_L, carreau_Y = [], []

    total_collected = 0
    attempts = 0
    max_attempts = cfg.n_samples * 500 # Safety limit

    # ---------------------------------------------------------------------
    # B. Global Generation Loop
    # ---------------------------------------------------------------------
    gen_start = time.time()
    
    while total_collected < cfg.n_samples and attempts < max_attempts:
        attempts += 1
        
        # --- CASE 1: Carreau-Yasuda (No Stability Check) ---
        if cfg.constitutive_eq == "carreau_yasuda":
            v_ratio = np.random.uniform(0, cfg.max_vorticity_ratio)
            L0 = generate_base_L_tensor(dim=cfg.dim, vorticity_ratio=v_ratio)
            nu = carreau_yasuda_viscosity(
                L0, nu_0=5.28e-5, nu_inf=3.30e-6,
                lambda_val=1.902, n=0.22, a=1.25
            )
            carreau_L.append(L0)
            carreau_Y.append(nu)
            total_collected += 1

        # --- CASE 2: Stability Models (Maxwell, Oldroyd, PTT) ---
        else:
            try:
                # 1. Generate Sample in Global Range
                if cfg.constitutive_eq == "maxwell_B":
                    L0, D, W, T, condA, resid = generate_stable_sample_maxwell(
                        cfg.dim, cfg.eta0, cfg.lam, target_cond=None, stage_range=gen_range
                    )
                elif cfg.constitutive_eq == "oldroyd_B":
                    L0, D, W, T, condA = generate_stable_sample_oldroyd(
                        cfg.dim, cfg.eta0, cfg.lam, cfg.lam_r, target_cond=None, stage_range=gen_range
                    )
                    resid = 0.0
                elif cfg.constitutive_eq == "ptt_exponential":
                    L0, D, W, T, condA = generate_stable_sample_ptt_exponential(
                        cfg.dim, cfg.eta0, cfg.lam, alpha=1.0, target_cond=None, stage_range=gen_range
                    )
                    resid = 0.0
                
                # 2. Determine Bin
                tag = get_stage_tag(condA, cfg.mode)
                
                # 3. Store
                if tag in data_bins:
                    data_bins[tag].append((L0, T, condA, W, resid))
                    total_collected += 1
            
            except Exception:
                pass # Skip failed generations

    gen_time = time.time() - gen_start
    print(f"   ✅ Generation finished in {gen_time:.2f}s. Collected: {total_collected}")

    # ---------------------------------------------------------------------
    # C. Processing & Saving Loop (Per Bin)
    # ---------------------------------------------------------------------
    
    # Handle Carreau separately (Single Bin)
    if cfg.constitutive_eq == "carreau_yasuda":
        # Treat as one big stage for saving
        # (This adapts the logic to fit the loop structure below)
        data_bins = {bin_keys[0]: [(l, y, 0.0, np.zeros((3,3)), 0.0) for l, y in zip(carreau_L, carreau_Y)]}

    # Iterate through bins
    for stage_tag, items in data_bins.items():
        if len(items) == 0:
            continue
            
        stage_start = time.time() # <--- Start Timer for this Stage
        print(f"   -> Processing Bin {stage_tag}: {len(items)} candidates")

        # 1. Unzip Data
        L0_list = [x[0] for x in items]
        Y_list  = [x[1] for x in items]
        condA_list = [x[2] for x in items]
        W_list  = [x[3] for x in items]
        residual_list = [x[4] for x in items]

        # 2. Filter Admissible Region (Skip for Carreau)
        if cfg.constitutive_eq != "carreau_yasuda":
            try:
                filtered_L0, kept_mask = filter_admissible_region(L0_list)
                Y_list = [x for x, k in zip(Y_list, kept_mask) if k]
                condA_list = [x for x, k in zip(condA_list, kept_mask) if k]
                W_list = [x for x, k in zip(W_list, kept_mask) if k]
                residual_list = [x for x, k in zip(residual_list, kept_mask) if k]
                L0_list = filtered_L0
            except Exception as e:
                print(f"      ⚠️ Filter error in {stage_tag}: {e}")
                continue

        # 3. Save Data if samples exist
        if len(L0_list) > 0:

            # --- NEW: Define the size folder name (e.g., "10ksamples" or "200ksamples") ---
            if cfg.n_samples >= 1000:
                size_folder = f"{int(cfg.n_samples/1000)}ksamples"
            else:
                size_folder = f"{cfg.n_samples}samples"
            # Paths
            stage_data_path = os.path.join(
                cfg.paths.data, "random", folder_mode, f"seed_{cfg.seed}", stage_tag, size_folder
            )
            stage_images_path = os.path.join(
                cfg.paths.images, "random", folder_mode, f"seed_{cfg.seed}", stage_tag, size_folder, cfg.constitutive_eq
            )
            os.makedirs(stage_data_path, exist_ok=True)
            os.makedirs(stage_images_path, exist_ok=True)

            # Convert & Save Tensors
            X_np = np.array(L0_list)
            Y_np = np.array(Y_list)
            
            # Flatten logic depends on model output shape
            if cfg.constitutive_eq == "carreau_yasuda":
                X_flat = X_np.reshape(X_np.shape[0], -1)
                Y_flat = Y_np.reshape(-1, 1) # Scalar viscosity
            else:
                X_flat = X_np.reshape(X_np.shape[0], -1)
                Y_flat = flatten_symmetric_tensors(Y_np) # Symmetric tensor 6-comp

            suffix = "_stage"
            torch.save(torch.tensor(X_flat, dtype=torch.float32), 
                       os.path.join(stage_data_path, f"X_{cfg.dim}D_{cfg.constitutive_eq}{suffix}.pt"))
            torch.save(torch.tensor(Y_flat, dtype=torch.float32), 
                       os.path.join(stage_data_path, f"Y_{cfg.dim}D_{cfg.constitutive_eq}{suffix}.pt"))
            
            print(f"      ✅ Saved {len(L0_list)} samples to {stage_tag}")

            # 4. Analysis Plots
            try:
                write_sampledata_file(
                    model_name=cfg.constitutive_eq,
                    stability_status=cfg.mode,
                    L_list=L0_list, T_list=Y_list,
                    condA_list=condA_list, lam=cfg.lam,
                    save_root=stage_images_path,
                    cfg=cfg, X_flat=X_flat, Y_flat=Y_flat,
                    stage_tag=stage_tag, 
                    residual_list=residual_list
                )
            except Exception as e:
                print(f"      ⚠️ Plotting error: {e}")
            
            # Carreau Parameter Study Plot
            if cfg.constitutive_eq == "carreau_yasuda":
                 # Logic for extra plot... (simplified for brevity, logic remains same as upload)
                 pass

        else:
            print(f"      ⚠️ Bin {stage_tag} empty after filtering.")

        # ---------------------------------------------------------------------
        # Timer & Device Info (Per Stage)
        # ---------------------------------------------------------------------
        elapsed_time = time.time() - stage_start
        gpus = GPUtil.getGPUs()
        if gpus:
            print(f"      ⏱ Stage {stage_tag} time: {elapsed_time:.2f}s on {gpus[0].name}")
        else:
            print(f"      ⏱ Stage {stage_tag} time: {elapsed_time:.2f}s on CPU")

    # ---------------------------------------------------------------------
    # D. Final Summary Plots
    # ---------------------------------------------------------------------
    if cfg.mode == "multi_stage":
        plot_stage_progression(
            data_root=os.path.join(cfg.paths.data, "random", folder_mode, f"seed_{cfg.seed}"),
            images_root=os.path.join(cfg.paths.images, "random", folder_mode, f"seed_{cfg.seed}"),
            model_name=cfg.constitutive_eq,
            suffix="_stage",
            stages=bin_keys, n_samples=cfg.n_samples
        )

if __name__ == "__main__":
    main()