metadentify 0.1.0a0__py3-none-any.whl

metadentify/args.py

@@ -0,0 +1,59 @@
+import argparse
+
+
+def get_base_parser(description: str = "") -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description=description)
+
+    parser.add_argument("--experiment_setup_path", type=str, default="")
+    parser.add_argument("--config_path", type=str, default="config.py")
+    parser.add_argument("--wandb_project", type=str, default="")
+    parser.add_argument("--wandb_entity", type=str, default="")
+    parser.add_argument("--progress_bar", action="store_true")
+    parser.add_argument("--savelast", action="store_true")
+    parser.add_argument("--overwrite_data", default=True, action=argparse.BooleanOptionalAction)
+    parser.add_argument("--disable_wandb", action="store_true")
+
+    parser.add_argument("--experiment_name", type=str, required=True)
+    parser.add_argument("--query_name", type=str, default="PATE")
+    parser.add_argument("--baseline_setting", type=str, default=None)
+
+    parser.add_argument("--backbone_type", type=str, default="q-cnp")
+    parser.add_argument("--embed_dim", type=int, default=64)
+    parser.add_argument("--num_heads", type=int, default=4)
+    parser.add_argument("--num_layers", type=int, default=2)
+    parser.add_argument("--num_tau_samples", type=int, default=8)
+    parser.add_argument("--source_embed_dim", type=int, default=16)
+    parser.add_argument("--output_dim", type=int, default=1)
+    parser.add_argument("--num_inducing_points", type=int, default=32)
+    parser.add_argument("--dropout", type=float, default=0.0)
+
+    parser.add_argument("--standardize", action="store_true")
+    parser.add_argument("--num_pytorch_workers", type=int, default=0)
+    parser.add_argument("--num_datagen_workers", type=int, default=-2)
+    parser.add_argument("--dataset_size", type=int, default=1000)
+    parser.add_argument("--num_query_points", type=int, default=10)
+    parser.add_argument("--num_scms", type=int, default=None)
+    parser.add_argument("--num_train_tasks", type=int, default=10000)
+    parser.add_argument("--num_val_tasks", type=int, default=1000)
+    parser.add_argument("--num_test_tasks", type=int, default=1000)
+    parser.add_argument("--tasks_per_file", type=int, default=500)
+    parser.add_argument("--batch_size", type=int, default=64)
+    parser.add_argument("--prefetch_factor", type=int, default=2)
+    parser.add_argument("--online_data", action="store_true")
+    parser.add_argument("--total_steps", type=int, default=None)
+
+    parser.add_argument(
+        "--val_metrics_normalized", default=True, action=argparse.BooleanOptionalAction
+    )
+    parser.add_argument("--plot_results", action="store_true")
+    parser.add_argument("--lr", type=float, default=1e-3)
+    parser.add_argument("--lambda_crossing_penalty", type=float, default=1.0)
+    parser.add_argument("--weight_decay", type=float, default=0.0)
+    parser.add_argument("--patience", type=int, default=30)
+    parser.add_argument("--max_epochs", type=int, default=100)
+    parser.add_argument("--run_baselines", action="store_true")
+    parser.add_argument("--save_dir", type=str, default="data")
+    parser.add_argument("--checkpoint_dir", type=str, default="checkpoints")
+    parser.add_argument("--plot_diagnostics", action="store_true")
+
+    return parser

metadentify/baselines.py

@@ -0,0 +1,460 @@
+from typing import Any
+
+import numpy as np
+from econml.dml import CausalForestDML, LinearDML
+from econml.iv.dml import DMLIV
+from sklearn.linear_model import LinearRegression, Ridge
+from sklearn.neural_network import MLPRegressor
+
+
+def confounder_linear_baseline(
+    x: np.ndarray,
+    t: np.ndarray,
+    y: np.ndarray,
+    query_x: np.ndarray | None = None,
+    x_sources: np.ndarray | None = None,
+) -> float | np.ndarray:
+    dataset_x = np.c_[t, x]
+    linmod = LinearRegression().fit(X=dataset_x, y=y)
+    estimate = linmod.coef_[0]
+
+    if query_x is not None:
+        estimate = np.ones((query_x.shape[0],)) * estimate
+
+    return estimate
+
+
+def confounder_ridge_baseline(
+    x: np.ndarray,
+    t: np.ndarray,
+    y: np.ndarray,
+    query_x: np.ndarray | None = None,
+    x_sources: np.ndarray | None = None,
+) -> float | np.ndarray:
+    dataset_x = np.c_[t, x]
+    linmod = Ridge().fit(X=dataset_x, y=y)
+    estimate = linmod.coef_[0]
+
+    if query_x is not None:
+        estimate = np.ones((query_x.shape[0],)) * estimate
+
+    return estimate
+
+
+def confounder_mlp_baseline(
+    x: np.ndarray,
+    t: np.ndarray,
+    y: np.ndarray,
+    query_x: np.ndarray | None = None,
+    x_sources: np.ndarray | None = None,
+) -> float | np.ndarray:
+    dataset_x = np.c_[t, x]
+    mlp = MLPRegressor(max_iter=3000).fit(X=dataset_x, y=y)
+    if np.array_equal(np.unique(t), [0, 1]):
+        t_sample = t * 0
+    else:
+        t_sample = np.random.uniform(low=t.min(), high=t.max(), size=x.shape[0])
+    control_x = np.c_[t_sample, x]
+    treatment_x = np.c_[t_sample + 1, x]
+    estimate = np.mean(mlp.predict(X=treatment_x) - mlp.predict(X=control_x))
+
+    if query_x is not None:
+        estimate = np.ones((query_x.shape[0],)) * estimate
+
+    return estimate
+
+
+def treatment_only_linear_baseline(
+    x: np.ndarray,
+    t: np.ndarray,
+    y: np.ndarray,
+    query_x: np.ndarray | None = None,
+    x_sources: np.ndarray | None = None,
+) -> float | np.ndarray:
+    linmod = LinearRegression().fit(X=t.reshape(-1, 1), y=y)
+    estimate = linmod.coef_[0]
+
+    if query_x is not None:
+        estimate = np.ones((query_x.shape[0],)) * estimate
+
+    return estimate
+
+
+def treatment_only_ridge_baseline(
+    x: np.ndarray,
+    t: np.ndarray,
+    y: np.ndarray,
+    query_x: np.ndarray | None = None,
+    x_sources: np.ndarray | None = None,
+) -> float | np.ndarray:
+    linmod = Ridge().fit(X=t.reshape(-1, 1), y=y)
+    estimate = linmod.coef_[0]
+
+    if query_x is not None:
+        estimate = np.ones((query_x.shape[0],)) * estimate
+
+    return estimate
+
+
+def tsls_linear_baseline(
+    x: np.ndarray,
+    t: np.ndarray,
+    y: np.ndarray,
+    query_x: np.ndarray | None = None,
+    x_sources: np.ndarray | None = None,
+) -> float | np.ndarray:
+    stage1linmod = LinearRegression().fit(X=x, y=t)
+    predicted_t = stage1linmod.predict(X=x)
+    stagt2linmod = LinearRegression().fit(X=predicted_t.reshape(-1, 1), y=y)
+    estimate = stagt2linmod.coef_[0]
+    return estimate
+
+
+def tsls_ridge_baseline(
+    x: np.ndarray,
+    t: np.ndarray,
+    y: np.ndarray,
+    query_x: np.ndarray | None = None,
+    x_sources: np.ndarray | None = None,
+) -> float | np.ndarray:
+    stage1linmod = Ridge().fit(X=x, y=t)
+    predicted_t = stage1linmod.predict(X=x)
+    stagt2linmod = Ridge().fit(X=predicted_t.reshape(-1, 1), y=y)
+    estimate = stagt2linmod.coef_[0]
+    return estimate
+
+
+def tsls_mlp_baseline(
+    x: np.ndarray,
+    t: np.ndarray,
+    y: np.ndarray,
+    query_x: np.ndarray | None = None,
+    x_sources: np.ndarray | None = None,
+) -> float | np.ndarray:
+    stage1model = MLPRegressor().fit(X=x, y=t.ravel())
+    predicted_t = stage1model.predict(X=x)
+    predicted_t = np.expand_dims(predicted_t, -1)
+    stage2model = Ridge().fit(X=predicted_t, y=y)
+    estimate = stage2model.coef_[0]
+    return estimate
+
+
+def proxy_linear_baseline(
+    x: np.ndarray,
+    t: np.ndarray,
+    y: np.ndarray,
+    query_x: np.ndarray | None = None,
+    x_sources: np.ndarray | None = None,
+) -> float | np.ndarray:
+    proxy_1 = x[:, [0]]
+    proxy_2 = x[:, [1]]
+    treatment_and_proxy_2 = np.c_[proxy_2, t]
+    stage1model = LinearRegression().fit(X=treatment_and_proxy_2, y=proxy_1)
+    predicted_proxy_1 = stage1model.predict(X=treatment_and_proxy_2)
+    treatment_and_predicted_proxy = np.c_[t, predicted_proxy_1]
+    stage2model = LinearRegression().fit(X=treatment_and_predicted_proxy, y=y)
+    estimate = stage2model.coef_[0]
+    return estimate
+
+
+def proxy_ridge_baseline(
+    x: np.ndarray,
+    t: np.ndarray,
+    y: np.ndarray,
+    query_x: np.ndarray | None = None,
+    x_sources: np.ndarray | None = None,
+) -> float | np.ndarray:
+    proxy_1 = x[:, [0]]
+    proxy_2 = x[:, [1]]
+    treatment_and_proxy_2 = np.c_[proxy_2, t]
+    stage1model = Ridge().fit(X=treatment_and_proxy_2, y=proxy_1)
+    predicted_proxy_1 = stage1model.predict(X=treatment_and_proxy_2)
+    treatment_and_predicted_proxy = np.c_[t, predicted_proxy_1]
+    stage2model = Ridge().fit(X=treatment_and_predicted_proxy, y=y)
+    estimate = stage2model.coef_[0]
+    return estimate
+
+
+def proxy_mlp_baseline(
+    x: np.ndarray,
+    t: np.ndarray,
+    y: np.ndarray,
+    query_x: np.ndarray | None = None,
+    x_sources: np.ndarray | None = None,
+) -> float | np.ndarray:
+    proxy_1 = x[:, [0]]
+    proxy_2 = x[:, [1]]
+    treatment_and_proxy_2 = np.c_[proxy_2, t]
+    stage1model = MLPRegressor().fit(X=treatment_and_proxy_2, y=proxy_1)
+    predicted_proxy_1 = stage1model.predict(X=treatment_and_proxy_2)
+    treatment_and_predicted_proxy = np.c_[t, predicted_proxy_1]
+    stage2model = Ridge().fit(X=treatment_and_predicted_proxy, y=y)
+    estimate = stage2model.coef_[0]
+    return estimate
+
+
+def confounder_dml_baseline(
+    x: np.ndarray,
+    t: np.ndarray,
+    y: np.ndarray,
+    query_x: np.ndarray | None = None,
+    x_sources: np.ndarray | None = None,
+    type: str = "forest",
+    binary_t: bool = False,
+) -> float | np.ndarray:
+    if type == "forest":
+        model = CausalForestDML(
+            discrete_treatment=binary_t,
+            discrete_outcome=False,
+            n_jobs=None,
+        )
+    elif type == "linear":
+        model = LinearDML(
+            discrete_treatment=binary_t,
+            discrete_outcome=False,
+        )
+    model.fit(Y=y, T=t, X=x, W=None)
+    estimate = model.ate(X=x)
+    return estimate
+
+
+def iv_dml_baseline(
+    x: np.ndarray,
+    t: np.ndarray,
+    y: np.ndarray,
+    query_x: np.ndarray | None = None,
+    x_sources: np.ndarray | None = None,
+    type: str = "forest",
+    binary_t: bool = False,
+) -> float | np.ndarray:
+    if type == "forest":
+        model = DMLIV(
+            discrete_treatment=binary_t,
+            discrete_instrument=False,
+            model_y_xw="forest",
+            model_t_xw="forest",
+            model_t_xwz="forest",
+        )
+    elif type == "linear":
+        model = DMLIV(
+            discrete_treatment=binary_t,
+            discrete_instrument=False,
+            model_y_xw="linear",
+            model_t_xw="linear",
+            model_t_xwz="linear",
+        )
+    model.fit(Y=y, T=t, Z=x, X=None, W=None)
+    estimate = model.ate()
+    return estimate
+
+
+def cate_confounder_linear_baseline(
+    x: np.ndarray,
+    t: np.ndarray,
+    y: np.ndarray,
+    query_x: np.ndarray,
+    x_sources: np.ndarray | None = None,
+) -> float | np.ndarray:
+    dataset_x = np.c_[x * t, t]
+    linmod = LinearRegression().fit(X=dataset_x, y=y)
+    estimate = linmod.coef_[0] * query_x + linmod.coef_[1]
+    return estimate
+
+
+def cate_confounder_ridge_baseline(
+    x: np.ndarray,
+    t: np.ndarray,
+    y: np.ndarray,
+    query_x: np.ndarray,
+    x_sources: np.ndarray | None = None,
+) -> float | np.ndarray:
+    dataset_x = np.c_[x * t, t]
+    linmod = Ridge().fit(X=dataset_x, y=y)
+    estimate = query_x * linmod.coef_[0] + linmod.coef_[1]
+    return estimate
+
+
+def cate_confounder_mlp_baseline(
+    x: np.ndarray,
+    t: np.ndarray,
+    y: np.ndarray,
+    query_x: np.ndarray,
+    num_t_samples: int = 1000,
+    x_sources: np.ndarray | None = None,
+) -> list[float]:
+    dataset_x = np.c_[t, x]
+    mlp = MLPRegressor(max_iter=3000).fit(X=dataset_x, y=y)
+    if np.array_equal(np.unique(t), [0, 1]):
+        t_sample = np.zeros((num_t_samples, 1))
+    else:
+        t_sample = np.random.uniform(low=t.min(), high=t.max(), size=(num_t_samples, 1))
+
+    estimates = []
+    for i in range(query_x.shape[0]):
+        query_row = np.array(query_x[i]).reshape(1, -1)
+        query_point = np.repeat(query_row, repeats=num_t_samples, axis=0)
+        control_x = np.c_[t_sample, query_point]
+        treatment_x = np.c_[t_sample + 1, query_point]
+        estimate = np.mean(mlp.predict(X=treatment_x) - mlp.predict(X=control_x))
+        estimates.append(estimate)
+
+    return estimates
+
+
+def get_baselines(baseline_setting: str) -> list[dict[str, Any]]:
+    t_only_dict = {
+        "model": treatment_only_ridge_baseline,
+        "name": "t_only_ridge",
+        "alias": "T-Only-Ridge",
+    }
+
+    baselines = [t_only_dict]
+
+    if baseline_setting in [
+        "t-only",
+        "t-only-double-invalid-ocp",
+        "t-only-iv-invalid-tcp",
+    ]:
+        return baselines
+
+    elif baseline_setting == "confounder":
+        baseline_1_model = confounder_ridge_baseline
+        baseline_1_name = "ridge_regression"
+        baseline_1_alias = "Reg-Ridge"
+
+        baseline_2_model = confounder_mlp_baseline
+        baseline_2_name = "mlp_regression"
+        baseline_2_alias = "Reg-MLP"
+
+    elif baseline_setting == "iv":
+        baseline_1_model = tsls_ridge_baseline
+        baseline_1_name = "tsls_ridge"
+        baseline_1_alias = "TSLS-Ridge"
+
+        baseline_2_model = tsls_mlp_baseline
+        baseline_2_name = "mlp_tsls"
+        baseline_2_alias = "TSLS-MLP"
+
+    elif baseline_setting in ["proxy", "double-invalid-ocp"]:
+        baseline_1_model = proxy_ridge_baseline
+        baseline_1_name = "proxy_tsls_ridge"
+        baseline_1_alias = "PrTSLS-Ridge"
+
+        baseline_2_model = proxy_mlp_baseline
+        baseline_2_name = "mlp_proxy_tsls"
+        baseline_2_alias = "PrTSLS-MLP"
+
+    elif baseline_setting == "confounder-linreg":
+        baseline_1_model = confounder_linear_baseline
+        baseline_1_name = "linear_regression"
+        baseline_1_alias = "Reg-Lin"
+
+        baseline_2_model = confounder_mlp_baseline
+        baseline_2_name = "mlp_regression"
+        baseline_2_alias = "Reg-MLP"
+
+    elif baseline_setting == "iv-linreg":
+        baseline_1_model = tsls_linear_baseline
+        baseline_1_name = "tsls_linear"
+        baseline_1_alias = "TSLS-Lin"
+
+        baseline_2_model = tsls_mlp_baseline
+        baseline_2_name = "mlp_tsls"
+        baseline_2_alias = "TSLS-MLP"
+
+    elif baseline_setting in ["proxy-linreg", "double-invalid-ocp-linreg"]:
+        baseline_1_model = proxy_linear_baseline
+        baseline_1_name = "proxy_tsls_linear"
+        baseline_1_alias = "PrTSLS-Lin"
+
+        baseline_2_model = proxy_mlp_baseline
+        baseline_2_name = "mlp_proxy_tsls"
+        baseline_2_alias = "PrTSLS-MLP"
+
+    elif baseline_setting == "confounder-and-iv":
+        pass
+
+    elif baseline_setting == "iv-invalid-tcp":
+        pass
+
+    elif baseline_setting == "confounder-binary-t-dml":
+
+        def baseline_1_model(x, t, y, query_x):
+            return confounder_dml_baseline(x, t, y, query_x, type="linear", binary_t=True)
+
+        baseline_1_name = "linear_dml"
+        baseline_1_alias = "DML-Lin"
+
+        def baseline_2_model(x, t, y, query_x):
+            return confounder_dml_baseline(x, t, y, query_x, type="forest", binary_t=True)
+
+        baseline_2_name = "forest_dml"
+        baseline_2_alias = "DML-RF"
+
+    elif baseline_setting == "confounder-continuous-t-dml":
+
+        def baseline_1_model(x, t, y, query_x):
+            return confounder_dml_baseline(x, t, y, query_x, type="linear", binary_t=False)
+
+        baseline_1_name = "linear_dml"
+        baseline_1_alias = "DML-Lin"
+
+        def baseline_2_model(x, t, y, query_x):
+            return confounder_dml_baseline(x, t, y, query_x, type="forest", binary_t=False)
+
+        baseline_2_name = "forest_dml"
+        baseline_2_alias = "DML-RF"
+
+    elif baseline_setting == "iv-binary-t-dml":
+
+        def baseline_1_model(x, t, y, query_x):
+            return iv_dml_baseline(x, t, y, query_x, type="linear", binary_t=True)
+
+        baseline_1_name = "linear_dmliv"
+        baseline_1_alias = "DMLIV-Lin"
+
+        def baseline_2_model(x, t, y, query_x):
+            return iv_dml_baseline(x, t, y, query_x, type="forest", binary_t=True)
+
+        baseline_2_name = "forest_dmliv"
+        baseline_2_alias = "DMLIV-RF"
+
+    elif baseline_setting == "iv-continuous-t-dml":
+
+        def baseline_1_model(x, t, y, query_x):
+            return iv_dml_baseline(x, t, y, query_x, type="linear", binary_t=False)
+
+        baseline_1_name = "linear_dmliv"
+        baseline_1_alias = "DMLIV-Lin"
+
+        def baseline_2_model(x, t, y, query_x):
+            return iv_dml_baseline(x, t, y, query_x, type="forest", binary_t=False)
+
+        baseline_2_name = "forest_dmliv"
+        baseline_2_alias = "DMLIV-RF"
+
+    elif baseline_setting == "cate-confounder-linear":
+        baseline_1_model = cate_confounder_ridge_baseline
+        baseline_1_name = "cate_ridge"
+        baseline_1_alias = "CATE-Ridge"
+
+        baseline_2_model = cate_confounder_mlp_baseline
+        baseline_2_name = "cate_mlp"
+        baseline_2_alias = "CATE-MLP"
+
+    elif baseline_setting == "cate-confounder-nonlinear":
+        baseline_1_model = cate_confounder_ridge_baseline
+        baseline_1_name = "cate_ridge"
+        baseline_1_alias = "CATE-Ridge"
+
+        baseline_2_model = cate_confounder_mlp_baseline
+        baseline_2_name = "cate_mlp"
+        baseline_2_alias = "CATE-MLP"
+
+    else:
+        raise ValueError(f"Unsupported baseline_setting: {baseline_setting}")
+
+    baselines.append(dict(model=baseline_1_model, name=baseline_1_name, alias=baseline_1_alias))
+    baselines.append(dict(model=baseline_2_model, name=baseline_2_name, alias=baseline_2_alias))
+
+    return baselines

metadentify/launch_sweep.py

@@ -0,0 +1,153 @@
+import argparse
+import os
+import re
+import subprocess
+
+from dotenv import load_dotenv
+
+from metadentify.sbatch import format_sbatch_script
+
+
+def initialize_wandb_sweep(yaml_config_path):
+    print(f"Initializing W&B sweep from {yaml_config_path} inside Singularity...")
+
+    singularity_env_path = os.environ.get("SINGULARITY_ENV_PATH")
+    singularity_container_path = os.environ.get("SINGULARITY_CONTAINER_PATH")
+
+    inner_command = f"source /ext3/env.sh; wandb sweep {yaml_config_path}"
+
+    singularity_cmd = [
+        "singularity",
+        "exec",
+        "--nv",
+        "--overlay",
+        singularity_env_path,
+        singularity_container_path,
+        "/bin/bash",
+        "-c",
+        inner_command,
+    ]
+
+    result = subprocess.run(singularity_cmd, capture_output=True, text=True)
+    output = result.stdout + result.stderr
+
+    if result.returncode != 0:
+        print(f"Error initializing sweep:\n{output}")
+        raise RuntimeError("Failed to create W&B sweep.")
+
+    match = re.search(r"wandb agent\s+([^\s]+)", output)
+    if not match:
+        print(f"Could not parse sweep ID from W&B output:\n{output}")
+        raise ValueError("Sweep ID regex match failed.")
+
+    sweep_id = match.group(1)
+    print(f"Successfully created sweep: {sweep_id}")
+    return sweep_id
+
+
+def launch_sweep_agents(
+    sweep_id,
+    num_agents,
+    num_nodes,
+    tasks_per_node,
+    cpus_per_task,
+    num_hours,
+    gigs_memory,
+    num_gpus,
+    slurm_directory="slurm_scripts",
+):
+    os.makedirs(slurm_directory, exist_ok=True)
+    base_sweep_id = sweep_id.split("/")[-1]
+
+    sweep_out_dir = os.path.join("out", base_sweep_id)
+    os.makedirs(sweep_out_dir, exist_ok=True)
+
+    wandb_entity = os.environ.get("WANDB_ENTITY")
+    wandb_project = os.environ.get("WANDB_PROJECT")
+
+    local_run_command = f"wandb agent {sweep_id} --entity {wandb_entity} --project {wandb_project}"
+
+    for i in range(1, num_agents + 1):
+        task_name = f"swp_{base_sweep_id}_{i}"
+        job_name = f"swp_{base_sweep_id[:6]}_{i}"
+
+        sbatch_script = format_sbatch_script(
+            local_run_command=local_run_command,
+            out_dir=sweep_out_dir,
+            task_name=task_name,
+            job_name=job_name,
+            num_nodes=num_nodes,
+            tasks_per_node=tasks_per_node,
+            cpus_per_task=cpus_per_task,
+            num_hours=num_hours,
+            gigs_memory=gigs_memory,
+            num_gpus=num_gpus,
+        )
+
+        task_path = os.path.join(slurm_directory, f"{task_name}.sbatch")
+        with open(task_path, "w") as out_file:
+            out_file.write(sbatch_script)
+
+        print(f"[{i}/{num_agents}] Submitting SLURM task: {task_name}")
+        subprocess.run(f"sbatch {task_path}", shell=True)
+
+    print(f"All {num_agents} agents submitted to queue")
+    print(f"SLURM logs will be saved to: {sweep_out_dir}/")
+
+
+def main():
+    load_dotenv()
+    required_keys = [
+        "SINGULARITY_ENV_PATH",
+        "SINGULARITY_CONTAINER_PATH",
+        "WANDB_ENTITY",
+        "WANDB_PROJECT",
+    ]
+    for key in required_keys:
+        if not os.environ.get(key):
+            raise ValueError(f"Missing {key} in .env file.")
+
+    parser = argparse.ArgumentParser(description="W&B Sweep Launcher for SLURM")
+
+    group = parser.add_mutually_exclusive_group(required=True)
+    group.add_argument(
+        "--config", type=str, help="Path to your sweep_config.yaml to start a NEW sweep"
+    )
+    group.add_argument(
+        "--sweep_id",
+        type=str,
+        help="Existing W&B sweep ID (e.g., entity/project/abcd1234) to RESUME an old sweep",
+    )
+    parser.add_argument(
+        "--num_jobs", type=int, default=4, help="Number of parallel SLURM jobs/agents to launch"
+    )
+
+    parser.add_argument("--nodes", type=int, default=1, help="Number of nodes per agent")
+    parser.add_argument("--tasks_per_node", type=int, default=1, help="Tasks per node")
+    parser.add_argument("--cpus", type=int, default=14, help="CPUs per task")
+    parser.add_argument("--hours", type=int, default=5, help="Wall time in hours")
+    parser.add_argument("--memory", type=int, default=32, help="Memory in GB")
+    parser.add_argument("--gpus", type=int, default=1, help="GPUs per agent")
+
+    args = parser.parse_args()
+
+    if args.config:
+        target_sweep_id = initialize_wandb_sweep(args.config)
+    else:
+        target_sweep_id = args.sweep_id
+        print(f"Attaching {args.num_jobs} new agents to existing sweep: {target_sweep_id}")
+
+    launch_sweep_agents(
+        sweep_id=target_sweep_id,
+        num_agents=args.num_jobs,
+        num_nodes=args.nodes,
+        tasks_per_node=args.tasks_per_node,
+        cpus_per_task=args.cpus,
+        num_hours=args.hours,
+        gigs_memory=args.memory,
+        num_gpus=args.gpus,
+    )
+
+
+if __name__ == "__main__":
+    main()