Airfoil dataset generator

engibench/problems/airfoil/dataset_slurm_airfoil.py

@@ -4,7 +4,9 @@
 """
 
 from argparse import ArgumentParser
+import sys
 
+from datasets import concatenate_datasets
 from datasets import load_dataset
 import numpy as np
 from scipy.stats import qmc
@@ -42,6 +44,20 @@ def calculate_runtime(group_size, minutes_per_sim=6):
     """
     # Fetch command line arguments for render and simulate to know whether to run those functions
     parser = ArgumentParser()
+    parser.add_argument(
+        "-account",
+        "--hpc_account",
+        type=str,
+        required=True,
+        help="HPC account allocation to charge for job submission",
+    )
+    parser.add_argument(
+        "-type",
+        "--job_type",
+        type=str,
+        required=True,
+        help="Engibench job type (simulate or optimization) for submission",
+    )
     parser.add_argument(
         "-n_designs",
         "--num_designs",
@@ -63,6 +79,13 @@ def calculate_runtime(group_size, minutes_per_sim=6):
         default=2,
         help="How many simulations do you wish to batch within each individual slurm job?",
     )
+    parser.add_argument(
+        "-minutes_per_sim",
+        "--minutes_per_simulation",
+        type=int,
+        default=15,
+        help="How long will each individual SLURM job take (in minutes) for either simulate or optimize? This is used to calculate the runtime for each job submission.",
+    )
     parser.add_argument(
         "-n_slurm_array",
         "--num_slurm_array",
@@ -114,13 +137,22 @@ def calculate_runtime(group_size, minutes_per_sim=6):
     )
     args = parser.parse_args()
 
+    # HPC account for job submission
+    hpc_account = args.hpc_account
+
+    # Job type (simulate or optimize)
+    job_type = args.job_type
+    if job_type not in ["simulate", "optimize"]:
+        raise ValueError(f"Invalid job type: {job_type}. Must be 'simulate' or 'optimize'.")
+
     # Number of samples & flow conditions
     n_designs = args.num_designs
     n_conditions = args.num_flow_conditions
 
     # Slurm parameters
     group_size = args.group_size
     n_slurm_array = args.num_slurm_array
+    minutes_per_sim = args.minutes_per_simulation
 
     # Flow parameter and angle of attack ranges
     min_ma = args.min_mach_number
@@ -138,22 +170,14 @@ def calculate_runtime(group_size, minutes_per_sim=6):
     print(f"Reynolds number:  {min_re:.2e} to {max_re:.2e}")
     print(f"Angle of attack:  {min_aoa:.1f} to {max_aoa:.1f}")
 
-    # Load airfoil designs from HF Database
+    # --- Dataset Loading ---
     ds = load_dataset("IDEALLab/airfoil_v0")
-    designs = (
-        ds["train"]["initial_design"]
-        + ds["train"]["optimal_design"]
-        + ds["val"]["initial_design"]
-        + ds["val"]["optimal_design"]
-        + ds["test"]["initial_design"]
-        + ds["test"]["optimal_design"]
-    )
-
-    # Use specified number of designs
+    all_data = concatenate_datasets([ds[split] for split in ds])
+    designs = all_data["initial_design"] + all_data["optimal_design"]
     if n_designs < len(designs):
         designs = designs[:n_designs]
 
-    # Generate all simulation configurations
+    # --- Config Generation ---
     config_id = 0
     simulate_configs_designs = []
     for design in designs:
@@ -185,8 +209,8 @@ def calculate_runtime(group_size, minutes_per_sim=6):
 
     slurm_config = slurm.SlurmConfig(
         name="Airfoil_dataset_generation",
-        runtime=calculate_runtime(group_size, minutes_per_sim=15),
-        account="fuge-prj-jrl",
+        runtime=calculate_runtime(group_size, minutes_per_sim=minutes_per_sim),
+        account=hpc_account,
         ntasks=1,
         cpus_per_task=1,
         log_dir="./sim_logs/",
@@ -197,16 +221,20 @@ def calculate_runtime(group_size, minutes_per_sim=6):
         sim_batch_configs = simulate_configs_designs[
             ibatch * group_size * n_slurm_array : (ibatch + 1) * group_size * n_slurm_array
         ]
-        print(len(sim_batch_configs))
         print(f"Submitting batch {ibatch + 1}/{int(n_sbatch_maps)}")
 
-        job_array = slurm.sbatch_map(
-            f=simulate_slurm,
-            args=sim_batch_configs,
-            slurm_args=slurm_config,
-            group_size=group_size,  # Number of jobs to batch in sequence to reduce job array size
-            work_dir="scratch",
-        )
+        if job_type == "simulate":
+            job_array = slurm.sbatch_map(
+                f=simulate_slurm,
+                args=sim_batch_configs,
+                slurm_args=slurm_config,
+                group_size=group_size,  # Number of jobs to batch in sequence to reduce job array size
+                work_dir="scratch",
+            )
+        elif job_type == "optimize":
+            sys.exit(
+                'Optimize function not yet implemented for SLURM dataset generation script. Please use "simulate" for job_type argument.'
+            )
 
         # Save the job array reference
         submitted_jobs.append(job_array)

engibench/problems/airfoil/simulation_jobs.py

@@ -54,3 +54,24 @@ def simulate_slurm(problem_configuration: dict, configuration_id: int, design: l
         "problem_configuration": problem_configuration,
         "configuration_id": configuration_id,
     }
+
+
+def optimize_slurm(problem_configuration: dict, configuration_id: int, design: list):
+    """Takes starting point (design coordinate and angle of attack) and config (mach, reynolds, angle of attack), then runs the aerodynamic optimization.
+
+    Any arguments should be things that you want to change across the different jobs, and anything
+    that is the same/static across the runs should just be defined inside this function.
+
+    Args:
+        problem_configuration (dict): The specific configuration used to initialize the optimization being passed.
+            For the airfoil problem this includes Mach number, Reynolds number, and angle of attack.
+        configuration_id (int): A unique identifier for the job for later debugging or tracking.
+        design (list): list of lists defining x and y coordinates of airfoil geometry.
+
+    Returns:
+        "performance_dict": Dictionary of aerodynamic performance (lift & drag).
+        "optimization_time": The time taken to run this optimization job. Useful for aggregating
+            the time taken for dataset generation.
+        "optimized_configuration": Problem configuration parameters for optimized design (optimized coordinates and angle of attack)
+        "configuration_id": Identifier for specific simulation configurations
+    """