app.py

import gradio as gr
import numpy as np
import matplotlib.pyplot as plt

# Constants and helper functions (same as before)
G_EARTH = 9.81
R_EARTH = 6371000
MU_EARTH = 3.986e14
R_GAS = 287
ATM_PRESSURE_SEA = 101325
TEMP_SEA = 288.15
LAPSE_RATE = 0.0065

def atmospheric_conditions(altitude):
    temp = TEMP_SEA - LAPSE_RATE * altitude if altitude < 11000 else 216.65
    pressure = ATM_PRESSURE_SEA * (temp / TEMP_SEA) ** (-G_EARTH / (LAPSE_RATE * R_GAS))
    density = pressure / (R_GAS * temp)
    return temp, pressure, density

def orbital_velocity(altitude):
    return np.sqrt(MU_EARTH / (R_EARTH + altitude))

def drag_force(velocity, altitude, diameter, cd=0.3):
    _, _, rho = atmospheric_conditions(altitude)
    area = np.pi * (diameter / 2) ** 2
    return 0.5 * rho * velocity ** 2 * cd * area

material_density = {"Aluminum": 2700, "Titanium": 4500, "Composite": 1600}
material_strength = {"Aluminum": 310e6, "Titanium": 950e6, "Composite": 600e6}
tank_density = {"Aluminum": 2700, "Stainless Steel": 8000, "Composite": 1600}

cost_factors = {
    "engines": 200000,
    "structural_material": {"Aluminum": 50, "Titanium": 150, "Composite": 200},
    "tank_material": {"Aluminum": 40, "Stainless Steel": 60, "Composite": 180},
    "payload": 100
}

# Simulate rocket function (placeholder, same as before)
def simulate_rocket(
    engine_type, num_engines, isp_vac, thrust_vac, chamber_pressure, nozzle_diameter,
    structural_material, structural_thickness, diameter, height, staging_enabled,
    stage1_fuel_mass, stage1_oxidizer_mass, stage2_fuel_mass, stage2_oxidizer_mass,
    control_system, payload_mass, tank_material, insulation_thickness, target_altitude
):
    return {
        "Delta-v (km/s)": "10.0 km/s",
        "TWR (Sea Level)": "1.5",
        "Structural Integrity (%)": "85",
        "Cost ($)": "$1000000"
    }, plt.figure()

# Optimization function with multiple goals and constraints
def optimize_design(
    delta_v_goal, twr_goal, integrity_goal, cost_goal,
    payload_mass, structural_material, control_system, target_altitude,
    max_engines, max_thickness, staging_enabled
):
    # Define design space
    engine_types = ["Liquid", "Solid", "Hybrid"]
    num_engines_options = range(1, max_engines + 1)
    isp_options = [250, 300, 350, 400]
    thrust_options = [100000, 200000, 500000, 1000000]
    structural_thickness_options = [t for t in [2, 5, 10, 15] if t <= max_thickness]
    
    best_design = None
    best_score = -float('inf')
    best_metrics = {}
    
    for engine_type in engine_types:
        for num_engines in num_engines_options:
            for isp in isp_options:
                for thrust in thrust_options:
                    for thickness in structural_thickness_options:
                        # Calculate design metrics
                        chamber_pressure = 7e6
                        nozzle_diameter = 0.5
                        diameter = 3
                        height = 20
                        stage1_fuel_mass = 2000 if staging_enabled else 4000
                        stage1_oxidizer_mass = 2000 if staging_enabled else 4000
                        stage2_fuel_mass = 500 if staging_enabled else 0
                        stage2_oxidizer_mass = 500 if staging_enabled else 0
                        tank_material = "Aluminum"
                        insulation_thickness = 10
                        
                        results, _ = simulate_rocket(
                            engine_type, num_engines, isp, thrust, chamber_pressure, nozzle_diameter,
                            structural_material, thickness, diameter, height, staging_enabled,
                            stage1_fuel_mass, stage1_oxidizer_mass, stage2_fuel_mass, stage2_oxidizer_mass,
                            control_system, payload_mass, tank_material, insulation_thickness, target_altitude
                        )
                        
                        if "error" not in results:
                            delta_v = float(results["Delta-v (km/s)"].split()[0])
                            twr = float(results["TWR (Sea Level)"].split()[0])
                            integrity = float(results["Structural Integrity (%)"])
                            cost = float(results["Cost ($)"].replace("$", "").replace(",", ""))
                            
                            # Check constraints (example thresholds for now)
                            if twr > 1.1 and integrity >= 80:
                                # Calculate score based on goals
                                # Simple scoring: minimize deviation from goals
                                delta_v_diff = abs(delta_v - delta_v_goal) / 15  # Normalize
                                twr_diff = max(0, twr_goal - twr) / 2  # Normalize, penalize if below goal
                                integrity_diff = max(0, integrity_goal - integrity) / 100  # Normalize
                                cost_diff = max(0, cost - cost_goal) / 10000000  # Normalize
                                score = -(delta_v_diff + twr_diff + integrity_diff + cost_diff)
                                
                                if score > best_score:
                                    best_score = score
                                    best_design = {
                                        "Engine Type": engine_type,
                                        "Num Engines": num_engines,
                                        "ISP": isp,
                                        "Thrust": thrust,
                                        "Thickness": thickness,
                                        "Staging": staging_enabled
                                    }
                                    best_metrics = results
    
    if best_design:
        return best_design, best_metrics
    else:
        return "No feasible design found", {}

# Gradio interface
with gr.Blocks(title="Rocket Design Optimizer") as app:
    gr.Markdown(
        """
        # Rocket Design Optimizer
        Define your design goals (desired performance metrics) and constraints, and the app will find the optimal rocket design.
        """
    )
    
    with gr.Row():
        with gr.Column():
            gr.Markdown("### Design Goals (Performance Metrics)")
            delta_v_goal = gr.Slider(5, 15, value=10, label="Delta-v Goal (km/s)")
            twr_goal = gr.Slider(1.1, 2.0, value=1.5, label="TWR Goal (Sea Level)")
            integrity_goal = gr.Slider(80, 100, value=85, label="Structural Integrity Goal (%)")
            cost_goal = gr.Slider(100000, 10000000, value=5000000, label="Cost Goal ($)")
        
        with gr.Column():
            gr.Markdown("### Constraints")
            payload_mass = gr.Slider(10, 1000, value=200, label="Payload Mass (kg)")
            structural_material = gr.Dropdown(choices=["Aluminum", "Titanium", "Composite"], value="Aluminum", label="Structural Material")
            control_system = gr.Dropdown(choices=["Gimbaled", "TVS", "RCS"], value="Gimbaled", label="Control System")
            target_altitude = gr.Slider(100000, 2000000, value=400000, label="Target Altitude (m)")
            max_engines = gr.Slider(1, 4, value=4, label="Max Number of Engines")
            max_thickness = gr.Slider(2, 15, value=15, label="Max Structural Thickness (mm)")
            staging_enabled = gr.Checkbox(value=True, label="Enable Staging")
    
    with gr.Row():
        optimize_button = gr.Button("Optimize Design")
        design_output = gr.JSON(label="Optimized Design Parameters")
        performance_output = gr.JSON(label="Performance Metrics")
        radar_plot = gr.Plot(label="Performance Profile")
    
    def run_optimization(
        delta_v_goal, twr_goal, integrity_goal, cost_goal,
        payload_mass, structural_material, control_system, target_altitude,
        max_engines, max_thickness, staging_enabled
    ):
        design, metrics = optimize_design(
            delta_v_goal, twr_goal, integrity_goal, cost_goal,
            payload_mass, structural_material, control_system, target_altitude,
            max_engines, max_thickness, staging_enabled
        )
        if design == "No feasible design found":
            return design, {}, plt.figure()
        else:
            # Simulate the design to get full metrics
            results, fig = simulate_rocket(
                design["Engine Type"], design["Num Engines"], design["ISP"], design["Thrust"], 7e6, 0.5,
                structural_material, design["Thickness"], 3, 20, staging_enabled,
                2000 if staging_enabled else 4000, 2000 if staging_enabled else 4000,
                500 if staging_enabled else 0, 500 if staging_enabled else 0,
                control_system, payload_mass, "Aluminum", 10, target_altitude
            )
            return design, metrics, fig
    
    optimize_button.click(
        fn=run_optimization,
        inputs=[
            delta_v_goal, twr_goal, integrity_goal, cost_goal,
            payload_mass, structural_material, control_system, target_altitude,
            max_engines, max_thickness, staging_enabled
        ],
        outputs=[design_output, performance_output, radar_plot]
    )

app.launch()