Strategy API

flowdesign.strategy

Generative modeling utilities for FlowDesign.

EnergyBasedStrategy

Bases: IntermediateStrategy

Energy-based physics injection from Fig. 2b/2c.

Use t_cut=0.0 in run(...) for full-trajectory injection, or t_cut>0.0 for an unconditional prefix followed by physics-guided injection.

Source code in flowdesign/strategy/intermediate.py

class EnergyBasedStrategy(IntermediateStrategy):
    """
    Energy-based physics injection from Fig. 2b/2c.

    Use t_cut=0.0 in run(...) for full-trajectory injection, or t_cut>0.0
    for an unconditional prefix followed by physics-guided injection.
    """

    def run(self, *args, **kwargs) -> dict:
        result = super().run(*args, **kwargs)
        t_cut = float(kwargs.get("t_cut", 0.0))
        result["strategy"] = "energy_full" if t_cut == 0.0 else "energy_intermediate"
        return result

EnergyFullStrategy

Bases: EnergyBasedStrategy

Energy-based strategy with physics injection from t=0.

Source code in flowdesign/strategy/intermediate.py

class EnergyFullStrategy(EnergyBasedStrategy):
    """Energy-based strategy with physics injection from t=0."""

    def run(self, *args, **kwargs) -> dict:
        kwargs["t_cut"] = 0.0
        return super().run(*args, **kwargs)

IntermediateStrategy

Generic intermediate strategy for a pre-trained flow-matching model.

It is agnostic to the exact physical quantities: everything is delegated to physics_fn and loss_fn.

Pipeline: 1) (Optional) Unconditional rollout from t = 0 -> t_cut using flow matching only. 2) Physics-guided drift from t = t_cut -> 1.0 using gradients from loss_fn(physics_fn(x)).

Parameters

wrapped_vf : callable Velocity field wrapper, callable as wrapped_vf(x, t) -> [B, D]. physics_fn : callable Function mapping x -> AeroCoeffs, typically a WingPhysics instance. loss_fn : callable Function mapping AeroCoeffs -> (loss, metrics_dict). dim : int Dimensionality of the design variables (e.g., 16 for CST). device : torch.device or str, optional Device for computation. If None, automatically selects CUDA if available. lambda_guidance : float, optional Scaling factor for the physics gradient guidance term. base_output_dir : str or None, optional Directory for saving CSV outputs. If None, no files are written.

Source code in flowdesign/strategy/intermediate.py

class IntermediateStrategy:
    """
    Generic intermediate strategy for a pre-trained flow-matching model.

    It is agnostic to the exact physical quantities:
    everything is delegated to `physics_fn` and `loss_fn`.

    Pipeline:
    1) (Optional) Unconditional rollout from t = 0 -> t_cut using flow matching only.
    2) Physics-guided drift from t = t_cut -> 1.0 using gradients from loss_fn(physics_fn(x)).

    Parameters
    ----------
    wrapped_vf : callable
        Velocity field wrapper, callable as wrapped_vf(x, t) -> [B, D].
    physics_fn : callable
        Function mapping x -> AeroCoeffs, typically a WingPhysics instance.
    loss_fn : callable
        Function mapping AeroCoeffs -> (loss, metrics_dict).
    dim : int
        Dimensionality of the design variables (e.g., 16 for CST).
    device : torch.device or str, optional
        Device for computation. If None, automatically selects CUDA if available.
    lambda_guidance : float, optional
        Scaling factor for the physics gradient guidance term.
    base_output_dir : str or None, optional
        Directory for saving CSV outputs. If None, no files are written.
    """

    def __init__(
        self,
        wrapped_vf,
        physics_fn,
        loss_fn,
        dim: int,
        device=None,
        lambda_guidance: float = 10.0,
        base_output_dir: str | None = "results",
        max_grad_norm: float | None = None,
        clamp_abs: float | None = None,
    ):
        self.wrapped_vf = wrapped_vf
        self.physics_fn = physics_fn
        self.loss_fn = loss_fn
        self.dim = int(dim)

        self.device = (
            torch.device("cuda" if torch.cuda.is_available() else "cpu")
            if device is None
            else torch.device(device)
        )
        self.lambda_guidance = float(lambda_guidance)
        self.base_output_dir = base_output_dir
        self.max_grad_norm = None if max_grad_norm is None else float(max_grad_norm)
        self.clamp_abs = None if clamp_abs is None else float(clamp_abs)

        if self.base_output_dir is not None:
            os.makedirs(self.base_output_dir, exist_ok=True)

    @torch.no_grad()
    def _unconditional_rollout(self, batch_size: int, n_steps_uncond: int, t_cut: float):
        """
        Unconditional flow-matching rollout from t=0 -> t_cut.

        Returns
        -------
        x_now : torch.Tensor
            State at cutoff time, shape [B, D].
        x_traj : list[torch.Tensor]
            List of states along the trajectory (each on CPU).
        """
        z0 = torch.randn((batch_size, self.dim), device=self.device)
        x_now = z0
        x_traj = [x_now.detach().cpu()]

        if n_steps_uncond <= 0:
            return x_now, x_traj

        t_vals = torch.linspace(0.0, t_cut, steps=n_steps_uncond, device=self.device)
        dt = float(t_cut) / float(n_steps_uncond)

        for step, t in enumerate(t_vals):
            v = self.wrapped_vf(x_now, t.expand(batch_size))
            v = torch.nan_to_num(v, nan=0.0, posinf=0.0, neginf=0.0)
            x_now = x_now + dt * v
            if self.clamp_abs is not None:
                x_now = x_now.clamp(-self.clamp_abs, self.clamp_abs)
            x_traj.append(x_now.detach().cpu())

            if step % 10 == 0 or step == len(t_vals) - 1:
                print(f"[Uncond Step {step}/{n_steps_uncond}] t={t.item():.3f}")

        return x_now.detach(), x_traj

    def run(
        self,
        n_steps: int,
        batch_size: int,
        t_cut: float,
        save_tag: str | None = None,
        save_csv: bool = True,
    ) -> dict:
        """
        Run the intermediate strategy.

        Parameters
        ----------
        n_steps : int
            Total number of steps (unconditional + physics-guided).
        batch_size : int
            Number of samples to generate.
        t_cut : float
            Cutoff time for switching from unconditional rollout to physics-guided drift.
            0 <= t_cut <= 1.0.
        save_tag : str, optional
            Extra tag for filenames (e.g. "_tc0.60_run0").
        save_csv : bool, optional
            If True, save a summary CSV of metrics.

        Returns
        -------
        result_dict : dict
            Contains:
                - "x_out"           : [B, 1, D] final design variables
                - "loss"            : list of physics-guided losses
                - "alignment_history"
                - "total_time"
                - "avg_time_per_sample"
                - "metrics_last_step" : dict of metrics from the final step
                - "diversity_std"
                - "strategy"        : "intermediate"
                - "x_trajectory"    : [T, B, D] full trajectory
                - "trajectory_semantics": dict explaining the single inference
                  time axis used by intermediate/energy guidance.
        """
        assert 0.0 <= t_cut <= 1.0, "t_cut must be in [0, 1]."
        strategy = "intermediate"
        print(f"=== Running Intermediate Strategy (t_cut={t_cut:.2f}) ===")

        time_start = time.time()

        # -----------------------------
        # 1) Unconditional rollout
        # -----------------------------
        x_traj = []
        if t_cut > 0.0:
            n_steps_uncond = int(n_steps * t_cut)
            xt_uncond, x_traj_uncond = self._unconditional_rollout(
                batch_size=batch_size,
                n_steps_uncond=n_steps_uncond,
                t_cut=t_cut,
            )
            x_traj.extend(x_traj_uncond)
            xt = xt_uncond.clone().requires_grad_(True)
            print(f"✅ Unconditional inference to t={t_cut:.2f} done.")
        else:
            xt = torch.randn((batch_size, self.dim), device=self.device)
            xt = xt.detach().clone().requires_grad_(True)
            x_traj.append(xt.detach().cpu())
            print("⚡ t_cut=0: skipping unconditional rollout; start from noise.")

        # -----------------------------
        # 2) Physics-guided drift
        # -----------------------------
        lambda_guidance = self.lambda_guidance
        n_phys_steps = max(int(n_steps * (1.0 - t_cut)), 1)
        dt_drift = 1.0 / n_phys_steps
        t_vals_drift = torch.linspace(
            t_cut, 1.0, steps=n_phys_steps + 1, device=self.device
        )

        loss_history = []
        alignment_history = []
        last_metrics = {}

        for step in range(n_phys_steps):
            t_val = t_vals_drift[step].expand(batch_size)

            # Flow velocity (no grad)
            with torch.no_grad():
                v = self.wrapped_vf(xt, t_val)
                v = torch.nan_to_num(v, nan=0.0, posinf=0.0, neginf=0.0)

            # Physics + loss (with grad)
            xt.requires_grad_(True)
            coeffs = self.physics_fn(xt)                       # AeroCoeffs
            loss_phys, metrics = self.loss_fn(coeffs)          # scalar + dict

            grad_xt = torch.autograd.grad(loss_phys, xt)[0]
            grad_xt = torch.nan_to_num(grad_xt, nan=0.0, posinf=0.0, neginf=0.0)
            if self.max_grad_norm is not None:
                grad_flat = grad_xt.reshape(batch_size, -1)
                grad_norm = grad_flat.norm(dim=1).clamp_min(1e-12)
                scale = (self.max_grad_norm / grad_norm).clamp_max(1.0).reshape(batch_size, *([1] * (grad_xt.ndim - 1)))
                grad_xt = grad_xt * scale

            # Euler step with guidance
            xt = xt.detach() + dt_drift * (v - lambda_guidance * grad_xt)
            if self.clamp_abs is not None:
                xt = xt.clamp(-self.clamp_abs, self.clamp_abs)
            x_traj.append(xt.detach().cpu())

            # Alignment (cosine similarity)
            v_flat = v.view(batch_size, -1)
            g_flat = grad_xt.view(batch_size, -1)
            cos_sim = torch.nn.functional.cosine_similarity(
                v_flat, -lambda_guidance * g_flat, dim=1
            )
            alignment_score = cos_sim.mean().item()
            alignment_history.append(alignment_score)

            # Logging
            loss_val = loss_phys.item()
            loss_history.append(loss_val)
            last_metrics = {k: v.detach().cpu() for k, v in metrics.items()}

            if step % 10 == 0 or step == n_phys_steps - 1:
                cl_mean = metrics.get("cl_mean", None)
                cd_max = metrics.get("cd_max", None)
                cm_mean = metrics.get("cm_mean", None)
                print(
                    f"[Step {step}/{n_phys_steps}] Loss={loss_val:.3e} "
                    f"| CL={float(cl_mean):.3f} | CD_max={float(cd_max):.3f} | CM={float(cm_mean):.3f}"
                    if (cl_mean is not None and cd_max is not None and cm_mean is not None)
                    else f"[Step {step}/{n_phys_steps}] Loss={loss_val:.3e}"
                )

        # -----------------------------
        # 3) Wrap up results
        # -----------------------------
        total_time = time.time() - time_start
        x_out = xt.detach().view(batch_size, 1, self.dim)
        full_traj = torch.stack(x_traj, dim=0).to(self.device)  # [T, B, D]

        result_dict = {
            "x_out": x_out,
            "alignment_history": alignment_history,
            "loss": loss_history,
            "total_time": total_time,
            "avg_time_per_sample": total_time / batch_size,
            "metrics_last_step": last_metrics,
            "diversity_std": x_out.std().item(),
            "strategy": strategy,
            "x_trajectory": full_traj,
            "trajectory_semantics": {
                "x_trajectory": "[T+1, B, D], one inference-time trajectory; physics guidance starts at t_cut",
                "t_cut": float(t_cut),
                "n_steps": int(n_steps),
                "n_unconditional_steps": int(n_steps * t_cut) if t_cut > 0.0 else 0,
                "n_guided_steps": int(n_phys_steps),
            },
        }

        # -----------------------------
        # 4) Optional: save CSV summary
        # -----------------------------
        if self.base_output_dir is not None and save_csv:
            tag = save_tag if save_tag is not None else ""
            csv_path = os.path.join(
                self.base_output_dir, f"{strategy}{tag}_metrics.csv"
            )
            # Flatten metrics into a single row
            row = {
                "strategy": strategy,
                "t_cut": t_cut,
                "total_time": total_time,
                "avg_time_per_sample": total_time / batch_size,
                "diversity_std": result_dict["diversity_std"],
            }
            for k, v in last_metrics.items():
                # v is a tensor scalar
                row[k] = float(v)

            df = pd.DataFrame([row])
            df.to_csv(csv_path, index=False)
            print(f"📄 Metrics saved to {csv_path}")

        return result_dict

run(n_steps, batch_size, t_cut, save_tag=None, save_csv=True)

Run the intermediate strategy.

Parameters

n_steps : int Total number of steps (unconditional + physics-guided). batch_size : int Number of samples to generate. t_cut : float Cutoff time for switching from unconditional rollout to physics-guided drift. 0 <= t_cut <= 1.0. save_tag : str, optional Extra tag for filenames (e.g. "_tc0.60_run0"). save_csv : bool, optional If True, save a summary CSV of metrics.

Returns

result_dict : dict Contains: - "x_out" : [B, 1, D] final design variables - "loss" : list of physics-guided losses - "alignment_history" - "total_time" - "avg_time_per_sample" - "metrics_last_step" : dict of metrics from the final step - "diversity_std" - "strategy" : "intermediate" - "x_trajectory" : [T, B, D] full trajectory - "trajectory_semantics": dict explaining the single inference time axis used by intermediate/energy guidance.

Source code in flowdesign/strategy/intermediate.py

def run(
    self,
    n_steps: int,
    batch_size: int,
    t_cut: float,
    save_tag: str | None = None,
    save_csv: bool = True,
) -> dict:
    """
    Run the intermediate strategy.

    Parameters
    ----------
    n_steps : int
        Total number of steps (unconditional + physics-guided).
    batch_size : int
        Number of samples to generate.
    t_cut : float
        Cutoff time for switching from unconditional rollout to physics-guided drift.
        0 <= t_cut <= 1.0.
    save_tag : str, optional
        Extra tag for filenames (e.g. "_tc0.60_run0").
    save_csv : bool, optional
        If True, save a summary CSV of metrics.

    Returns
    -------
    result_dict : dict
        Contains:
            - "x_out"           : [B, 1, D] final design variables
            - "loss"            : list of physics-guided losses
            - "alignment_history"
            - "total_time"
            - "avg_time_per_sample"
            - "metrics_last_step" : dict of metrics from the final step
            - "diversity_std"
            - "strategy"        : "intermediate"
            - "x_trajectory"    : [T, B, D] full trajectory
            - "trajectory_semantics": dict explaining the single inference
              time axis used by intermediate/energy guidance.
    """
    assert 0.0 <= t_cut <= 1.0, "t_cut must be in [0, 1]."
    strategy = "intermediate"
    print(f"=== Running Intermediate Strategy (t_cut={t_cut:.2f}) ===")

    time_start = time.time()

    # -----------------------------
    # 1) Unconditional rollout
    # -----------------------------
    x_traj = []
    if t_cut > 0.0:
        n_steps_uncond = int(n_steps * t_cut)
        xt_uncond, x_traj_uncond = self._unconditional_rollout(
            batch_size=batch_size,
            n_steps_uncond=n_steps_uncond,
            t_cut=t_cut,
        )
        x_traj.extend(x_traj_uncond)
        xt = xt_uncond.clone().requires_grad_(True)
        print(f"✅ Unconditional inference to t={t_cut:.2f} done.")
    else:
        xt = torch.randn((batch_size, self.dim), device=self.device)
        xt = xt.detach().clone().requires_grad_(True)
        x_traj.append(xt.detach().cpu())
        print("⚡ t_cut=0: skipping unconditional rollout; start from noise.")

    # -----------------------------
    # 2) Physics-guided drift
    # -----------------------------
    lambda_guidance = self.lambda_guidance
    n_phys_steps = max(int(n_steps * (1.0 - t_cut)), 1)
    dt_drift = 1.0 / n_phys_steps
    t_vals_drift = torch.linspace(
        t_cut, 1.0, steps=n_phys_steps + 1, device=self.device
    )

    loss_history = []
    alignment_history = []
    last_metrics = {}

    for step in range(n_phys_steps):
        t_val = t_vals_drift[step].expand(batch_size)

        # Flow velocity (no grad)
        with torch.no_grad():
            v = self.wrapped_vf(xt, t_val)
            v = torch.nan_to_num(v, nan=0.0, posinf=0.0, neginf=0.0)

        # Physics + loss (with grad)
        xt.requires_grad_(True)
        coeffs = self.physics_fn(xt)                       # AeroCoeffs
        loss_phys, metrics = self.loss_fn(coeffs)          # scalar + dict

        grad_xt = torch.autograd.grad(loss_phys, xt)[0]
        grad_xt = torch.nan_to_num(grad_xt, nan=0.0, posinf=0.0, neginf=0.0)
        if self.max_grad_norm is not None:
            grad_flat = grad_xt.reshape(batch_size, -1)
            grad_norm = grad_flat.norm(dim=1).clamp_min(1e-12)
            scale = (self.max_grad_norm / grad_norm).clamp_max(1.0).reshape(batch_size, *([1] * (grad_xt.ndim - 1)))
            grad_xt = grad_xt * scale

        # Euler step with guidance
        xt = xt.detach() + dt_drift * (v - lambda_guidance * grad_xt)
        if self.clamp_abs is not None:
            xt = xt.clamp(-self.clamp_abs, self.clamp_abs)
        x_traj.append(xt.detach().cpu())

        # Alignment (cosine similarity)
        v_flat = v.view(batch_size, -1)
        g_flat = grad_xt.view(batch_size, -1)
        cos_sim = torch.nn.functional.cosine_similarity(
            v_flat, -lambda_guidance * g_flat, dim=1
        )
        alignment_score = cos_sim.mean().item()
        alignment_history.append(alignment_score)

        # Logging
        loss_val = loss_phys.item()
        loss_history.append(loss_val)
        last_metrics = {k: v.detach().cpu() for k, v in metrics.items()}

        if step % 10 == 0 or step == n_phys_steps - 1:
            cl_mean = metrics.get("cl_mean", None)
            cd_max = metrics.get("cd_max", None)
            cm_mean = metrics.get("cm_mean", None)
            print(
                f"[Step {step}/{n_phys_steps}] Loss={loss_val:.3e} "
                f"| CL={float(cl_mean):.3f} | CD_max={float(cd_max):.3f} | CM={float(cm_mean):.3f}"
                if (cl_mean is not None and cd_max is not None and cm_mean is not None)
                else f"[Step {step}/{n_phys_steps}] Loss={loss_val:.3e}"
            )

    # -----------------------------
    # 3) Wrap up results
    # -----------------------------
    total_time = time.time() - time_start
    x_out = xt.detach().view(batch_size, 1, self.dim)
    full_traj = torch.stack(x_traj, dim=0).to(self.device)  # [T, B, D]

    result_dict = {
        "x_out": x_out,
        "alignment_history": alignment_history,
        "loss": loss_history,
        "total_time": total_time,
        "avg_time_per_sample": total_time / batch_size,
        "metrics_last_step": last_metrics,
        "diversity_std": x_out.std().item(),
        "strategy": strategy,
        "x_trajectory": full_traj,
        "trajectory_semantics": {
            "x_trajectory": "[T+1, B, D], one inference-time trajectory; physics guidance starts at t_cut",
            "t_cut": float(t_cut),
            "n_steps": int(n_steps),
            "n_unconditional_steps": int(n_steps * t_cut) if t_cut > 0.0 else 0,
            "n_guided_steps": int(n_phys_steps),
        },
    }

    # -----------------------------
    # 4) Optional: save CSV summary
    # -----------------------------
    if self.base_output_dir is not None and save_csv:
        tag = save_tag if save_tag is not None else ""
        csv_path = os.path.join(
            self.base_output_dir, f"{strategy}{tag}_metrics.csv"
        )
        # Flatten metrics into a single row
        row = {
            "strategy": strategy,
            "t_cut": t_cut,
            "total_time": total_time,
            "avg_time_per_sample": total_time / batch_size,
            "diversity_std": result_dict["diversity_std"],
        }
        for k, v in last_metrics.items():
            # v is a tensor scalar
            row[k] = float(v)

        df = pd.DataFrame([row])
        df.to_csv(csv_path, index=False)
        print(f"📄 Metrics saved to {csv_path}")

    return result_dict

MultipleShootingFlowStrategy dataclass

Multiple-shooting guided generation in DFlow source space.

The interval [t0, t1] is split into shooting segments. The optimization variables are the node states z_0 ... z_M. Terminal design loss is applied at z_M, while segment defects penalize violations of z_{i+1} = Phi_i(z_i). The returned designs are generated by reintegrating from the optimized source node z_0.

Source code in flowdesign/strategy/multiple_shooting.py

@dataclass
class MultipleShootingFlowStrategy:
    """
    Multiple-shooting guided generation in DFlow source space.

    The interval [t0, t1] is split into shooting segments. The optimization
    variables are the node states z_0 ... z_M. Terminal design loss is applied
    at z_M, while segment defects penalize violations of
    z_{i+1} = Phi_i(z_i). The returned designs are generated by reintegrating
    from the optimized source node z_0.
    """

    num_samples: int = 256
    optimize_variable: str = "source"
    guidance_steps: int = 50
    guidance_lr: float = 1e-2
    segments: int = 4
    flow_steps_per_segment: int = 2
    t0: float = 0.0
    t1: float = 1.0
    integrator: str = "rk4"
    optimizer: str = "adam"
    defect_weight: float = 100.0
    diversity_weight: float = 0.0
    source_l2_weight: float = 0.0
    keep_best: int | None = None
    device: str | torch.device | None = None
    clamp_abs: float | None = None
    record_optimization_trajectory: bool = False
    record_inference_trajectory: bool = False
    flow_model_path: str | None = None
    flow_input_dim: int | None = None
    flow_hidden_dim: int = 512
    flow_device: str | torch.device | None = None

    def run(self, problem: DesignProblem, model: Any = None, sampler: Any = None, **kwargs: Any) -> DesignResult:
        params = _merged_dataclass_params(self, kwargs)
        device = torch.device(params["device"] or "cpu")
        segments = int(params["segments"])
        if segments < 1:
            raise ValueError("segments must be >= 1.")

        source = _initial_samples(problem, model, sampler, int(params["num_samples"]), device)
        vector_field = _resolve_vector_field(model)
        if vector_field is None and params.get("flow_model_path"):
            vector_field = load_flow_model(
                params["flow_model_path"],
                input_dim=int(params["flow_input_dim"] or source.reshape(source.shape[0], -1).shape[-1]),
                hidden_dim=int(params["flow_hidden_dim"]),
                device=torch.device(params["flow_device"] or device),
            )
        nodes = _initial_nodes(
            vector_field,
            source,
            segments=segments,
            steps_per_segment=int(params["flow_steps_per_segment"]),
            t0=float(params["t0"]),
            t1=float(params["t1"]),
            integrator=str(params["integrator"]),
        ).detach().clone().requires_grad_(True)
        optimizer_name = str(params["optimizer"])
        optimizer = _make_optimizer(optimizer_name, [nodes], float(params["guidance_lr"]))

        loss_history: list[float] = []
        design_loss_history: list[float] = []
        defect_history: list[float] = []
        optimization_terminals: list[torch.Tensor] | None = [] if params["record_optimization_trajectory"] else None
        optimization_inference_trajectories: list[torch.Tensor] | None = (
            [] if params["record_inference_trajectory"] else None
        )
        _record_multiple_shooting_flow_state(
            vector_field,
            nodes[0].detach(),
            params,
            segments,
            optimization_terminals,
            optimization_inference_trajectories,
        )
        times = torch.linspace(float(params["t0"]), float(params["t1"]), segments + 1, device=device)
        for step in range(int(params["guidance_steps"])):
            last: dict[str, torch.Tensor] = {}

            def closure() -> torch.Tensor:
                optimizer.zero_grad()
                defect = _defect_loss(
                    vector_field,
                    nodes,
                    times,
                    steps_per_segment=int(params["flow_steps_per_segment"]),
                    integrator=str(params["integrator"]),
                )
                terminal = nodes[-1]
                outputs = problem.evaluate(terminal)
                design_loss = problem.loss(terminal, outputs).mean()
                loss = design_loss + float(params["defect_weight"]) * defect
                if float(params["diversity_weight"]):
                    loss = loss - float(params["diversity_weight"]) * terminal.std(dim=0).mean()
                if float(params["source_l2_weight"]):
                    loss = loss + float(params["source_l2_weight"]) * nodes[0].square().mean()
                loss.backward()
                last["loss"] = loss.detach()
                last["design_loss"] = design_loss.detach()
                last["defect"] = defect.detach()
                return loss

            if optimizer_name.lower() == "lbfgs":
                optimizer.step(closure)
            else:
                closure()
                optimizer.step()
            if params["clamp_abs"] is not None:
                with torch.no_grad():
                    nodes.clamp_(-float(params["clamp_abs"]), float(params["clamp_abs"]))
            _record_multiple_shooting_flow_state(
                vector_field,
                nodes[0].detach(),
                params,
                segments,
                optimization_terminals,
                optimization_inference_trajectories,
            )
            loss_history.append(float(last["loss"].item()))
            design_loss_history.append(float(last["design_loss"].item()))
            defect_history.append(float(last["defect"].item()))
            if step % 10 == 0:
                print(
                    f"[MultipleShootingFlow step {step}] "
                    f"loss={loss_history[-1]:.6g} design_loss={design_loss_history[-1]:.6g} "
                    f"defect={defect_history[-1]:.6g}"
                )

        with torch.no_grad():
            samples = flow_map(
                vector_field,
                nodes[0].detach(),
                flow_steps=segments * int(params["flow_steps_per_segment"]),
                t0=float(params["t0"]),
                t1=float(params["t1"]),
                integrator=str(params["integrator"]),
            )
        return evaluate_result(
            problem,
            samples,
            strategy="multiple_shooting",
            keep_best=params["keep_best"],
            metrics={
                "loss_history": loss_history,
                "design_loss_history": design_loss_history,
                "defect_history": defect_history,
                "num_samples": int(samples.shape[0]),
                **_trajectory_metrics(optimization_terminals, optimization_inference_trajectories),
            },
            metadata={
                "flow": "identity" if vector_field is None else "model",
                "segments": segments,
                "flow_steps_per_segment": int(params["flow_steps_per_segment"]),
                "integrator": str(params["integrator"]),
            },
        )

PostOptStrategy

Post-Optimization strategy for a pre-trained flow-matching model.

Pipeline: 1) Sample initial designs x0 by rolling out the flow model from t = 1 -> 0. 2) For each sample, run local gradient-based optimization (Adam) on the physical loss starting from x0[i]. 3) Collect per-sample loss curves and final metrics.

This class is agnostic to the exact physical quantities: - physics_fn maps x -> AeroCoeffs (CL, CD, CM, ...) - loss_fn maps AeroCoeffs -> (loss, metrics_dict)

Parameters

wrapped_vf : callable Velocity field wrapper, callable as wrapped_vf(x, t) -> [B, D]. (e.g. flowdesign.utils.dflow_utils.WrappedModel) physics_fn : AeroPhysics (or compatible) Callable mapping x -> AeroCoeffs (CL, CD, CM). loss_fn : PhysicalLoss (or compatible) Callable mapping AeroCoeffs -> (loss, metrics_dict). dim : int Dimensionality of the design variables (e.g., 16 for CST). device : str or torch.device, optional Device for computation. If None, automatically selects CUDA if available. base_output_dir : str or None, optional Directory for saving optional outputs. If None, no files are written.

Source code in flowdesign/strategy/postopt.py

class PostOptStrategy:
    """
    Post-Optimization strategy for a pre-trained flow-matching model.

    Pipeline:
    1) Sample initial designs x0 by rolling out the flow model from t = 1 -> 0.
    2) For each sample, run local gradient-based optimization (Adam) on the
        physical loss starting from x0[i].
    3) Collect per-sample loss curves and final metrics.

    This class is agnostic to the exact physical quantities:
    - `physics_fn` maps x -> AeroCoeffs (CL, CD, CM, ...)
    - `loss_fn`   maps AeroCoeffs -> (loss, metrics_dict)

    Parameters
    ----------
    wrapped_vf : callable
        Velocity field wrapper, callable as wrapped_vf(x, t) -> [B, D].
        (e.g. flowdesign.utils.dflow_utils.WrappedModel)
    physics_fn : AeroPhysics (or compatible)
        Callable mapping x -> AeroCoeffs (CL, CD, CM).
    loss_fn : PhysicalLoss (or compatible)
        Callable mapping AeroCoeffs -> (loss, metrics_dict).
    dim : int
        Dimensionality of the design variables (e.g., 16 for CST).
    device : str or torch.device, optional
        Device for computation. If None, automatically selects CUDA if available.
    base_output_dir : str or None, optional
        Directory for saving optional outputs. If None, no files are written.
    """

    def __init__(
        self,
        wrapped_vf: Callable[[torch.Tensor, torch.Tensor], torch.Tensor],
        physics_fn: AeroPhysics,
        loss_fn: PhysicalLoss,
        dim: int,
        device: Optional[str | torch.device] = None,
        base_output_dir: Optional[str] = "results_postopt",
    ):
        self.wrapped_vf = wrapped_vf
        self.physics_fn = physics_fn
        self.loss_fn = loss_fn
        self.dim = int(dim)

        if device is None:
            self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        else:
            self.device = torch.device(device)

        self.base_output_dir = base_output_dir
        if self.base_output_dir is not None:
            os.makedirs(self.base_output_dir, exist_ok=True)

    # ------------------------------------------------------------------
    # 1) Rollout flow model t=0 -> 1 to obtain x0
    # ------------------------------------------------------------------
    @torch.no_grad()
    def _sample_from_flow(
        self,
        n_flow_steps: int,
        batch_size: int,
    ) -> Tuple[torch.Tensor, List[torch.Tensor]]:
        """
        Roll out the flow model from t = 0.0 -> 1.0 to obtain initial designs x0.

        This mirrors the original postopt branch:
            - Start from z0 ~ N(0, I)
            - Time grid: t in [0.0, 1.0]
            - x_{k+1} = x_k + dt * v(x_k, t_k)

        Returns
        -------
        x0 : torch.Tensor
            Final designs at t=0, shape [B, D].
        traj_list : list[torch.Tensor]
            List of intermediate states (each on CPU) including z0 and x0.
        """
        device = self.device
        B, D = batch_size, self.dim

        z0 = torch.randn((B, D), device=device)
        x_list: List[torch.Tensor] = [z0.detach().cpu()]

        # t from 0.0 -> 1.0
        N = n_flow_steps
        dt = 1.0 / N
        t_vals = torch.linspace(0.0, 1.0, steps=N + 1, device=device)

        x_now = z0
        for k in range(N):
            tk = t_vals[k].expand(B)
            vk = self.wrapped_vf(x_now, tk)  # [B, D]
            x_now = x_now + dt * vk
            x_list.append(x_now.detach().cpu())

        x0 = x_now.detach()
        return x0, x_list

    # ------------------------------------------------------------------
    # 2) Local optimization starting from x0
    # ------------------------------------------------------------------
    def run(
        self,
        n_flow_steps: int,
        batch_size: int,
        max_opt_steps: int = 200,
        lr: float = 1e-2,
        early_stop_patience: int = 100,
        early_stop_thresh: float = 1e-4,
        thresholds: Tuple[float, ...] = (1e-2, 1e-3, 1e-4),
        save_tag: Optional[str] = None,
        save_csv: bool = True,
    ) -> Dict[str, Any]:
        """
        Run the Post-Optimization strategy.

        Parameters
        ----------
        n_flow_steps : int
            Number of steps for the flow rollout (t=1->0) to get x0.
        batch_size : int
            Number of samples to generate and optimize.
        max_opt_steps : int, optional
            Maximum number of local optimization steps per sample.
        lr : float, optional
            Learning rate for Adam optimizer.
        early_stop_patience : int, optional
            Early stopping patience on loss (no significant improvement).
        early_stop_thresh : float, optional
            Early stopping threshold on best loss achieved.
        thresholds : tuple of float, optional
            Thresholds at which we record the step index (e.g., 1e-2, 1e-3, 1e-4).
        save_tag : str, optional
            Extra tag for filenames (e.g. "_flow1000").
        save_csv : bool, optional
            If True, save a summary CSV with aggregate metrics.

        Returns
        -------
        result_dict : dict
            Contains:
                - "x_out"              : [B, 1, D] final optimized designs
                - "loss"               : list of per-sample loss curves (length B)
                - "steps_to_thresholds": dict mapping threshold -> avg step index
                - "total_time"
                - "avg_time_per_sample"
                - "mean_cl_error", "std_cl_error"
                - "mean_cm_error", "std_cm_error"
                - "mean_cd_violation", "max_cd_violation"
                - "diversity_std"
                - "strategy"           : "postopt"
                - "x0_from_flow"       : initial designs x0 [B, D]
        """
        strategy = "postopt"
        device = self.device
        time_start = time.time()

        print(f"=== Running PostOptStrategy (batch_size={batch_size}, flow_steps={n_flow_steps}) ===")

        # --------------------------------------------------------------
        # Step 1: Rollout flow to get x0
        # --------------------------------------------------------------
        with torch.no_grad():
            x0, x_traj_flow = self._sample_from_flow(
                n_flow_steps=n_flow_steps,
                batch_size=batch_size,
            )

        # Optionally save x0
        if self.base_output_dir is not None:
            tag = save_tag or ""
            x0_path = os.path.join(self.base_output_dir, f"{strategy}{tag}_x0_from_flow.pt")
            torch.save(x0, x0_path)
            print(f"Saved initial x0 from flow rollout to: {x0_path}")

        # --------------------------------------------------------------
        # Step 2: Per-sample local optimization
        # --------------------------------------------------------------
        all_x_opt: List[torch.Tensor] = []
        all_loss_curves: List[List[float]] = []
        steps_to_thresholds: Dict[float, List[int]] = {thr: [] for thr in thresholds}

        for i in range(batch_size):
            xi = x0[i].clone().to(device).detach().requires_grad_(True)
            optimizer = torch.optim.Adam([xi], lr=lr)
            loss_curve: List[float] = []

            min_loss = float("inf")
            patience = 0

            for step in range(max_opt_steps):
                optimizer.zero_grad()

                # Physics + loss (batch size 1)
                coeffs: AeroCoeffs = self.physics_fn(xi.unsqueeze(0))   # [1, D] -> AeroCoeffs
                loss_phys, _ = self.loss_fn(coeffs)                     # scalar

                loss_phys.backward()
                optimizer.step()

                loss_val = float(loss_phys.item())
                loss_curve.append(loss_val)

                if step % 10 == 0:
                    print(f"[Sample {i}, Step {step}] Loss: {loss_val:.6e}")

                # Update best and patience
                if loss_val < min_loss - 1e-6:
                    min_loss = loss_val
                    patience = 0
                else:
                    patience += 1

                # Early stopping conditions
                if min_loss < early_stop_thresh or patience >= early_stop_patience:
                    print(
                        f"✅ Sample {i} early stopped at step {step} "
                        f"with min_loss = {min_loss:.6e}"
                    )
                    break

            all_x_opt.append(xi.detach().cpu())
            all_loss_curves.append(loss_curve)

            # Record step index when each threshold is first reached
            for thr in thresholds:
                reached_step = len(loss_curve)  # default: not reached within window
                for s_idx, val in enumerate(loss_curve):
                    if val <= thr:
                        reached_step = s_idx
                        break
                steps_to_thresholds[thr].append(reached_step)

        # --------------------------------------------------------------
        # Step 3: Final metrics and aggregation
        # --------------------------------------------------------------
        x_opt_out = torch.stack(all_x_opt).to(device)  # [B, D]
        x_opt_out_view = x_opt_out.view(batch_size, 1, self.dim)

        total_time = time.time() - time_start

        # Compute final aero coefficients at optimized designs
        with torch.no_grad():
            coeffs_final: AeroCoeffs = self.physics_fn(x_opt_out)  # [B, D] -> AeroCoeffs
            cl_final = coeffs_final.cl
            cd_final = coeffs_final.cd
            cm_final = coeffs_final.cm

        # Try to read targets from loss_fn (PhysicalLoss)
        target_cl = getattr(self.loss_fn, "target_cl", 0.0)
        target_cm = getattr(self.loss_fn, "target_cm", 0.0)
        cd_thresh = getattr(self.loss_fn, "cd_threshold", 0.0)

        target_cl_tensor = torch.full_like(cl_final, float(target_cl))
        target_cm_tensor = torch.full_like(cm_final, float(target_cm))

        cl_error = torch.abs(cl_final - target_cl_tensor)
        cm_error = torch.abs(cm_final - target_cm_tensor)
        cd_viol = torch.relu(cd_final - cd_thresh)

        # Aggregate threshold steps (average over samples)
        avg_steps_to_thresholds = {
            thr: int(sum(vals) // max(len(vals), 1)) for thr, vals in steps_to_thresholds.items()
        }

        result_dict: Dict[str, Any] = {
            "x_out": x_opt_out_view,           # [B, 1, D]
            "loss": all_loss_curves,          # list of length B
            "steps_to_thresholds": avg_steps_to_thresholds,
            "total_time": total_time,
            "avg_time_per_sample": total_time / float(batch_size),
            "mean_cl_error": float(cl_error.mean().item()),
            "std_cl_error": float(cl_error.std().item()),
            "mean_cm_error": float(cm_error.mean().item()),
            "std_cm_error": float(cm_error.std().item()),
            "mean_cd_violation": float(cd_viol.mean().item()),
            "max_cd_violation": float(cd_viol.max().item()),
            "diversity_std": float(x_opt_out.std().item()),
            "strategy": strategy,
            "x0_from_flow": x0.to(device),
        }

        # --------------------------------------------------------------
        # Step 4: Optional CSV summary
        # --------------------------------------------------------------
        if self.base_output_dir is not None and save_csv:
            tag = save_tag or ""
            csv_path = os.path.join(self.base_output_dir, f"{strategy}{tag}_metrics.csv")

            row: Dict[str, Any] = {
                "strategy": strategy,
                "total_time": result_dict["total_time"],
                "avg_time_per_sample": result_dict["avg_time_per_sample"],
                "diversity_std": result_dict["diversity_std"],
                "mean_cl_error": result_dict["mean_cl_error"],
                "std_cl_error": result_dict["std_cl_error"],
                "mean_cm_error": result_dict["mean_cm_error"],
                "std_cm_error": result_dict["std_cm_error"],
                "mean_cd_violation": result_dict["mean_cd_violation"],
                "max_cd_violation": result_dict["max_cd_violation"],
            }
            # include thresholds statistics
            for thr, avg_step in avg_steps_to_thresholds.items():
                key = f"steps_to_{thr:.0e}".replace("-", "m")  # e.g. 1e-4 -> "steps_to_1e-04"
                row[key] = avg_step

            df = pd.DataFrame([row])
            df.to_csv(csv_path, index=False)
            print(f"📄 PostOpt metrics saved to: {csv_path}")

        return result_dict

run(n_flow_steps, batch_size, max_opt_steps=200, lr=0.01, early_stop_patience=100, early_stop_thresh=0.0001, thresholds=(0.01, 0.001, 0.0001), save_tag=None, save_csv=True)

Run the Post-Optimization strategy.

Parameters

n_flow_steps : int Number of steps for the flow rollout (t=1->0) to get x0. batch_size : int Number of samples to generate and optimize. max_opt_steps : int, optional Maximum number of local optimization steps per sample. lr : float, optional Learning rate for Adam optimizer. early_stop_patience : int, optional Early stopping patience on loss (no significant improvement). early_stop_thresh : float, optional Early stopping threshold on best loss achieved. thresholds : tuple of float, optional Thresholds at which we record the step index (e.g., 1e-2, 1e-3, 1e-4). save_tag : str, optional Extra tag for filenames (e.g. "_flow1000"). save_csv : bool, optional If True, save a summary CSV with aggregate metrics.

Returns

result_dict : dict Contains: - "x_out" : [B, 1, D] final optimized designs - "loss" : list of per-sample loss curves (length B) - "steps_to_thresholds": dict mapping threshold -> avg step index - "total_time" - "avg_time_per_sample" - "mean_cl_error", "std_cl_error" - "mean_cm_error", "std_cm_error" - "mean_cd_violation", "max_cd_violation" - "diversity_std" - "strategy" : "postopt" - "x0_from_flow" : initial designs x0 [B, D]

Source code in flowdesign/strategy/postopt.py

def run(
    self,
    n_flow_steps: int,
    batch_size: int,
    max_opt_steps: int = 200,
    lr: float = 1e-2,
    early_stop_patience: int = 100,
    early_stop_thresh: float = 1e-4,
    thresholds: Tuple[float, ...] = (1e-2, 1e-3, 1e-4),
    save_tag: Optional[str] = None,
    save_csv: bool = True,
) -> Dict[str, Any]:
    """
    Run the Post-Optimization strategy.

    Parameters
    ----------
    n_flow_steps : int
        Number of steps for the flow rollout (t=1->0) to get x0.
    batch_size : int
        Number of samples to generate and optimize.
    max_opt_steps : int, optional
        Maximum number of local optimization steps per sample.
    lr : float, optional
        Learning rate for Adam optimizer.
    early_stop_patience : int, optional
        Early stopping patience on loss (no significant improvement).
    early_stop_thresh : float, optional
        Early stopping threshold on best loss achieved.
    thresholds : tuple of float, optional
        Thresholds at which we record the step index (e.g., 1e-2, 1e-3, 1e-4).
    save_tag : str, optional
        Extra tag for filenames (e.g. "_flow1000").
    save_csv : bool, optional
        If True, save a summary CSV with aggregate metrics.

    Returns
    -------
    result_dict : dict
        Contains:
            - "x_out"              : [B, 1, D] final optimized designs
            - "loss"               : list of per-sample loss curves (length B)
            - "steps_to_thresholds": dict mapping threshold -> avg step index
            - "total_time"
            - "avg_time_per_sample"
            - "mean_cl_error", "std_cl_error"
            - "mean_cm_error", "std_cm_error"
            - "mean_cd_violation", "max_cd_violation"
            - "diversity_std"
            - "strategy"           : "postopt"
            - "x0_from_flow"       : initial designs x0 [B, D]
    """
    strategy = "postopt"
    device = self.device
    time_start = time.time()

    print(f"=== Running PostOptStrategy (batch_size={batch_size}, flow_steps={n_flow_steps}) ===")

    # --------------------------------------------------------------
    # Step 1: Rollout flow to get x0
    # --------------------------------------------------------------
    with torch.no_grad():
        x0, x_traj_flow = self._sample_from_flow(
            n_flow_steps=n_flow_steps,
            batch_size=batch_size,
        )

    # Optionally save x0
    if self.base_output_dir is not None:
        tag = save_tag or ""
        x0_path = os.path.join(self.base_output_dir, f"{strategy}{tag}_x0_from_flow.pt")
        torch.save(x0, x0_path)
        print(f"Saved initial x0 from flow rollout to: {x0_path}")

    # --------------------------------------------------------------
    # Step 2: Per-sample local optimization
    # --------------------------------------------------------------
    all_x_opt: List[torch.Tensor] = []
    all_loss_curves: List[List[float]] = []
    steps_to_thresholds: Dict[float, List[int]] = {thr: [] for thr in thresholds}

    for i in range(batch_size):
        xi = x0[i].clone().to(device).detach().requires_grad_(True)
        optimizer = torch.optim.Adam([xi], lr=lr)
        loss_curve: List[float] = []

        min_loss = float("inf")
        patience = 0

        for step in range(max_opt_steps):
            optimizer.zero_grad()

            # Physics + loss (batch size 1)
            coeffs: AeroCoeffs = self.physics_fn(xi.unsqueeze(0))   # [1, D] -> AeroCoeffs
            loss_phys, _ = self.loss_fn(coeffs)                     # scalar

            loss_phys.backward()
            optimizer.step()

            loss_val = float(loss_phys.item())
            loss_curve.append(loss_val)

            if step % 10 == 0:
                print(f"[Sample {i}, Step {step}] Loss: {loss_val:.6e}")

            # Update best and patience
            if loss_val < min_loss - 1e-6:
                min_loss = loss_val
                patience = 0
            else:
                patience += 1

            # Early stopping conditions
            if min_loss < early_stop_thresh or patience >= early_stop_patience:
                print(
                    f"✅ Sample {i} early stopped at step {step} "
                    f"with min_loss = {min_loss:.6e}"
                )
                break

        all_x_opt.append(xi.detach().cpu())
        all_loss_curves.append(loss_curve)

        # Record step index when each threshold is first reached
        for thr in thresholds:
            reached_step = len(loss_curve)  # default: not reached within window
            for s_idx, val in enumerate(loss_curve):
                if val <= thr:
                    reached_step = s_idx
                    break
            steps_to_thresholds[thr].append(reached_step)

    # --------------------------------------------------------------
    # Step 3: Final metrics and aggregation
    # --------------------------------------------------------------
    x_opt_out = torch.stack(all_x_opt).to(device)  # [B, D]
    x_opt_out_view = x_opt_out.view(batch_size, 1, self.dim)

    total_time = time.time() - time_start

    # Compute final aero coefficients at optimized designs
    with torch.no_grad():
        coeffs_final: AeroCoeffs = self.physics_fn(x_opt_out)  # [B, D] -> AeroCoeffs
        cl_final = coeffs_final.cl
        cd_final = coeffs_final.cd
        cm_final = coeffs_final.cm

    # Try to read targets from loss_fn (PhysicalLoss)
    target_cl = getattr(self.loss_fn, "target_cl", 0.0)
    target_cm = getattr(self.loss_fn, "target_cm", 0.0)
    cd_thresh = getattr(self.loss_fn, "cd_threshold", 0.0)

    target_cl_tensor = torch.full_like(cl_final, float(target_cl))
    target_cm_tensor = torch.full_like(cm_final, float(target_cm))

    cl_error = torch.abs(cl_final - target_cl_tensor)
    cm_error = torch.abs(cm_final - target_cm_tensor)
    cd_viol = torch.relu(cd_final - cd_thresh)

    # Aggregate threshold steps (average over samples)
    avg_steps_to_thresholds = {
        thr: int(sum(vals) // max(len(vals), 1)) for thr, vals in steps_to_thresholds.items()
    }

    result_dict: Dict[str, Any] = {
        "x_out": x_opt_out_view,           # [B, 1, D]
        "loss": all_loss_curves,          # list of length B
        "steps_to_thresholds": avg_steps_to_thresholds,
        "total_time": total_time,
        "avg_time_per_sample": total_time / float(batch_size),
        "mean_cl_error": float(cl_error.mean().item()),
        "std_cl_error": float(cl_error.std().item()),
        "mean_cm_error": float(cm_error.mean().item()),
        "std_cm_error": float(cm_error.std().item()),
        "mean_cd_violation": float(cd_viol.mean().item()),
        "max_cd_violation": float(cd_viol.max().item()),
        "diversity_std": float(x_opt_out.std().item()),
        "strategy": strategy,
        "x0_from_flow": x0.to(device),
    }

    # --------------------------------------------------------------
    # Step 4: Optional CSV summary
    # --------------------------------------------------------------
    if self.base_output_dir is not None and save_csv:
        tag = save_tag or ""
        csv_path = os.path.join(self.base_output_dir, f"{strategy}{tag}_metrics.csv")

        row: Dict[str, Any] = {
            "strategy": strategy,
            "total_time": result_dict["total_time"],
            "avg_time_per_sample": result_dict["avg_time_per_sample"],
            "diversity_std": result_dict["diversity_std"],
            "mean_cl_error": result_dict["mean_cl_error"],
            "std_cl_error": result_dict["std_cl_error"],
            "mean_cm_error": result_dict["mean_cm_error"],
            "std_cm_error": result_dict["std_cm_error"],
            "mean_cd_violation": result_dict["mean_cd_violation"],
            "max_cd_violation": result_dict["max_cd_violation"],
        }
        # include thresholds statistics
        for thr, avg_step in avg_steps_to_thresholds.items():
            key = f"steps_to_{thr:.0e}".replace("-", "m")  # e.g. 1e-4 -> "steps_to_1e-04"
            row[key] = avg_step

        df = pd.DataFrame([row])
        df.to_csv(csv_path, index=False)
        print(f"📄 PostOpt metrics saved to: {csv_path}")

    return result_dict

ShootingFlowStrategy dataclass

Single-shooting guided generation in DFlow source space.

The strategy solves the inference-time boundary-value problem by optimizing a source state z0 and differentiating through the flow map Phi(z0). If no flow/vector-field model is supplied, Phi is the identity map; this preserves the current package examples where strategies optimize design variables directly.

Source code in flowdesign/strategy/shooting.py

@dataclass
class ShootingFlowStrategy:
    """
    Single-shooting guided generation in DFlow source space.

    The strategy solves the inference-time boundary-value problem by optimizing
    a source state z0 and differentiating through the flow map Phi(z0). If no
    flow/vector-field model is supplied, Phi is the identity map; this preserves
    the current package examples where strategies optimize design variables
    directly.
    """

    num_samples: int = 256
    optimize_variable: str = "source"
    guidance_steps: int = 50
    guidance_lr: float = 1e-2
    flow_steps: int = 8
    t0: float = 0.0
    t1: float = 1.0
    integrator: str = "rk4"
    optimizer: str = "adam"
    diversity_weight: float = 0.0
    source_l2_weight: float = 0.0
    keep_best: int | None = None
    device: str | torch.device | None = None
    clamp_abs: float | None = None
    record_optimization_trajectory: bool = False
    record_inference_trajectory: bool = False
    flow_model_path: str | None = None
    flow_input_dim: int | None = None
    flow_hidden_dim: int = 512
    flow_device: str | torch.device | None = None

    def run(self, problem: DesignProblem, model: Any = None, sampler: Any = None, **kwargs: Any) -> DesignResult:
        params = _merged_dataclass_params(self, kwargs)
        device = torch.device(params["device"] or "cpu")
        source0 = _initial_samples(problem, model, sampler, int(params["num_samples"]), device)
        source = source0.detach().clone().requires_grad_(True)
        vector_field = _resolve_vector_field(model)
        if vector_field is None and params.get("flow_model_path"):
            vector_field = load_flow_model(
                params["flow_model_path"],
                input_dim=int(params["flow_input_dim"] or source.reshape(source.shape[0], -1).shape[-1]),
                hidden_dim=int(params["flow_hidden_dim"]),
                device=torch.device(params["flow_device"] or device),
            )
        optimizer_name = str(params["optimizer"])
        optimizer = _make_optimizer(optimizer_name, [source], float(params["guidance_lr"]))

        loss_history: list[float] = []
        design_loss_history: list[float] = []
        optimization_terminals: list[torch.Tensor] | None = [] if params["record_optimization_trajectory"] else None
        optimization_inference_trajectories: list[torch.Tensor] | None = (
            [] if params["record_inference_trajectory"] else None
        )
        _record_flow_state(
            vector_field,
            source.detach(),
            params,
            optimization_terminals,
            optimization_inference_trajectories,
        )
        for step in range(int(params["guidance_steps"])):
            last: dict[str, torch.Tensor] = {}

            def closure() -> torch.Tensor:
                optimizer.zero_grad()
                terminal = flow_map(
                    vector_field,
                    source,
                    flow_steps=int(params["flow_steps"]),
                    t0=float(params["t0"]),
                    t1=float(params["t1"]),
                    integrator=str(params["integrator"]),
                )
                outputs = problem.evaluate(terminal)
                design_loss = problem.loss(terminal, outputs).mean()
                loss = design_loss
                if float(params["diversity_weight"]):
                    loss = loss - float(params["diversity_weight"]) * terminal.std(dim=0).mean()
                if float(params["source_l2_weight"]):
                    loss = loss + float(params["source_l2_weight"]) * source.square().mean()
                loss.backward()
                last["loss"] = loss.detach()
                last["design_loss"] = design_loss.detach()
                return loss

            if optimizer_name.lower() == "lbfgs":
                optimizer.step(closure)
            else:
                closure()
                optimizer.step()
            if params["clamp_abs"] is not None:
                with torch.no_grad():
                    source.clamp_(-float(params["clamp_abs"]), float(params["clamp_abs"]))
            _record_flow_state(
                vector_field,
                source.detach(),
                params,
                optimization_terminals,
                optimization_inference_trajectories,
            )
            loss_history.append(float(last["loss"].item()))
            design_loss_history.append(float(last["design_loss"].item()))
            if step % 10 == 0:
                print(f"[ShootingFlow step {step}] loss={loss_history[-1]:.6g} design_loss={design_loss_history[-1]:.6g}")

        with torch.no_grad():
            samples = flow_map(
                vector_field,
                source.detach(),
                flow_steps=int(params["flow_steps"]),
                t0=float(params["t0"]),
                t1=float(params["t1"]),
                integrator=str(params["integrator"]),
            )
        return evaluate_result(
            problem,
            samples,
            strategy="shooting",
            keep_best=params["keep_best"],
            metrics={
                "loss_history": loss_history,
                "design_loss_history": design_loss_history,
                "num_samples": int(samples.shape[0]),
                **_trajectory_metrics(optimization_terminals, optimization_inference_trajectories),
            },
            metadata={
                "flow": "identity" if vector_field is None else "model",
                "flow_steps": int(params["flow_steps"]),
                "integrator": str(params["integrator"]),
            },
        )

SolverGuidedDFlowStrategy dataclass

DFlow-style generation guided by solver-estimated losses.

This strategy is intended for non-autograd evaluators such as XFoil. It samples candidates from the geometry prior or sampler, assembles the standard DesignProblem loss, estimates gradients by central finite differences, and updates the design variables.

Source code in flowdesign/strategy/solver_guided.py

@dataclass
class SolverGuidedDFlowStrategy:
    """
    DFlow-style generation guided by solver-estimated losses.

    This strategy is intended for non-autograd evaluators such as XFoil.
    It samples candidates from the geometry prior or sampler, assembles the
    standard DesignProblem loss, estimates gradients by central finite
    differences, and updates the design variables.
    """

    num_samples: int = 32
    guidance_steps: int = 5
    guidance_lr: float = 2e-2
    fd_step: float = 1e-3
    keep_best: int | None = 8
    device: str | torch.device | None = None
    clamp_abs: float | None = 0.8
    failure_penalty: float = 1.0e3

    def run(self, problem: DesignProblem, model: Any = None, sampler: Any = None, **kwargs: Any) -> DesignResult:
        device = torch.device(self.device or kwargs.get("device") or "cpu")
        n = int(kwargs.get("num_samples", self.num_samples))
        if sampler is not None:
            x = sampler.sample(model=model, n=n, problem=problem, device=device)
        else:
            x = problem.geometry.sample_prior(n, device=device)
        x = x.detach().clone().to(device)

        loss_history: list[float] = []
        for step in range(int(kwargs.get("guidance_steps", self.guidance_steps))):
            losses = self._loss(problem, x)
            grad = self._finite_difference_grad(problem, x)
            x = x - float(kwargs.get("guidance_lr", self.guidance_lr)) * grad
            if self.clamp_abs is not None:
                x = x.clamp(-float(self.clamp_abs), float(self.clamp_abs))
            loss_history.append(float(losses.mean().item()))
            print(f"[SolverGuidedDFlow step {step}] mean_loss={loss_history[-1]:.6g}")

        outputs = problem.evaluate(x)
        losses = problem.loss(x, outputs)
        feasibility = problem.is_feasible(x, outputs)
        keep_best = self.keep_best
        if keep_best is not None and keep_best > 0 and losses.numel() > keep_best:
            indices = torch.argsort(losses)[: int(keep_best)]
            x = x[indices]
            outputs = {key: value[indices] for key, value in outputs.items()}
            losses = losses[indices]
            feasibility = feasibility[indices]
        return DesignResult(
            samples=x,
            outputs=outputs,
            losses=losses,
            feasibility=feasibility,
            metrics={
                "num_samples": int(x.shape[0]),
                "num_feasible": int(feasibility.sum().item()),
                "loss_history": loss_history,
            },
            metadata={
                "strategy": "solver_guided_dflow",
                "fd_step": float(kwargs.get("fd_step", self.fd_step)),
            },
        )

    def _loss(self, problem: DesignProblem, x: torch.Tensor) -> torch.Tensor:
        losses = []
        with torch.no_grad():
            for row in x:
                row_batch = row.reshape(1, -1)
                try:
                    outputs = problem.evaluate(row_batch)
                    loss = problem.loss(row_batch, outputs).reshape(-1)[0]
                except Exception:
                    loss = torch.tensor(self.failure_penalty, dtype=x.dtype, device=x.device)
                losses.append(loss)
        return torch.stack(losses)

    def _finite_difference_grad(self, problem: DesignProblem, x: torch.Tensor) -> torch.Tensor:
        eps = float(self.fd_step)
        grad = torch.zeros_like(x)
        for dim in range(x.shape[1]):
            perturb = torch.zeros_like(x)
            perturb[:, dim] = eps
            loss_plus = self._loss(problem, x + perturb)
            loss_minus = self._loss(problem, x - perturb)
            grad[:, dim] = (loss_plus - loss_minus) / (2.0 * eps)
        return grad

get_init_x(x_dims, ode_func, y, H, init_method='random', ode_opts={}, blend_coef=1.0, init_t=0.0, final_t=1.0)

Gets the initial state of the ODE.

Parameters:

Name	Type	Description	Default
x_dims	Size	The dimensions of the initial state.	required
ode_func	Callable	model wraper to use for the ODE. Expected API: ode_func(t, x) -> torch.Tensor	required
y	Tensor	observed signal / control value	required
H	Callable	forward operator	required
init_method	Any	The method to use to initialize the initial state. (default: 'random') 'random' - init_x ~ N(0,I) 'backward' - init_x = ODEsolve(y, [1,0]) 'backward_blend'- init_x = H(ODEsolve(y, [1,0]))sqrt(blend_coef) + zsqrt(1-blend_coef) , z~N(0,I) int - init_x = torch.randn(x_dims) with seed = init_method	'random'
ode_opts	Dict	A dictionary of options for the ODE solver.	{}
blend_coef	float	The blend coefficient for the initial state.	1.0
init_t	float	The initial time. (default: 0.)	0.0
final_t	float	The final time. (default: 1.)	1.0

Returns:

Type	Description
Tensor	The initial state of the ODE.

Source code in flowdesign/strategy/dflow.py

def get_init_x(x_dims: torch.Size,
            ode_func: Callable,
            y: torch.Tensor,
            H: Callable,
            init_method: Any = 'random',
            ode_opts: Dict = {},
            blend_coef: float = 1.,
            init_t: float = 0.,
            final_t: float = 1.
            ) -> torch.Tensor:
    """Gets the initial state of the ODE.

    Args:
        x_dims: The dimensions of the initial state.
        ode_func: model wraper to use for the ODE. Expected API: ode_func(t, x) -> torch.Tensor
        y: observed signal / control value
        H: forward operator
        init_method: The method to use to initialize the initial state. (default: 'random')
                    'random'        - init_x ~ N(0,I)
                    'backward'      - init_x = ODEsolve(y, [1,0])
                    'backward_blend'- init_x = H(ODEsolve(y, [1,0]))*sqrt(blend_coef) + z*sqrt(1-blend_coef) , z~N(0,I)
                    int             - init_x = torch.randn(x_dims) with seed = init_method
        ode_opts: A dictionary of options for the ODE solver.
        blend_coef: The blend coefficient for the initial state.
        init_t: The initial time.  (default: 0.)
        final_t: The final time.  (default: 1.)

    Returns:
        The initial state of the ODE.
    """

    if init_method == 'random':
        init_x = torch.randn(x_dims).to(y)

    elif init_method == 'backward' or init_method == 'backward_blend':
        print('solve back')
        with torch.no_grad():
            # when solving back, y needs to have same shape as x_dims.
            # For example, when H is a downsampling operator, one needs to pass an upsampled version of y:
            # y_upsampled = F.interpolate(y, scale_factor=scale_factor, mode='bicubic')
            init_x = ode_integrate(ode_func, y, ode_opts, init_t=final_t, final_t=init_t) # note we solve backward

        if init_method == 'backward_blend':
            noise = torch.randn(x_dims).to(y)
            init_x = init_x * np.sqrt(blend_coef) + noise * np.sqrt(1-blend_coef)
    elif type(init_method) is int:
        torch.manual_seed(init_method)
        init_x = torch.randn(x_dims).to(y)
    else:
        raise ValueError(f"Uknown init_method {init_method}")
    return init_x

ode_integrate(ode_func, init_x, ode_opts={}, t_eps=0, init_t=0.0, final_t=1.0, t_arr=None, intermediate_points=False)

Integrates an ODE using the given ode_func.

Parameters:

Name	Type	Description	Default
ode_func	Callable	model wraper to use for the ODE. Expected API: ode_func(t, x) -> torch.Tensor	required
init_x	Tensor	The initial state of the ODE.	required
ode_opts	Dict	A dictionary of options for the ODE solver.	{}
t_eps	float	A small value to add to the initial time.	0
init_t	float	The initial time.	0.0
final_t	float	The final time.	1.0
t_arr	Optional[List[float]]	An optional array of times to evaluate the ODE at.	None
intermediate_points	bool	Whether to return the intermediate points.	False

Returns:

Type	Description
Tensor	The solution to the ODE at the final time. (or at intermediate points if intermediate_points=True)

Source code in flowdesign/strategy/dflow.py

def ode_integrate(ode_func: Callable,
                init_x: torch.Tensor,
                ode_opts: Dict = {},
                t_eps: float = 0,
                init_t: float = 0.,
                final_t: float = 1.,
                t_arr: Optional[List[float]] = None,
                intermediate_points: bool = False
    ) -> torch.Tensor:

    """Integrates an ODE using the given ode_func.

    Args:
        ode_func: model wraper to use for the ODE. Expected API: ode_func(t, x) -> torch.Tensor
        init_x: The initial state of the ODE.
        ode_opts: A dictionary of options for the ODE solver.
        t_eps: A small value to add to the initial time.
        init_t: The initial time.
        final_t: The final time.
        t_arr: An optional array of times to evaluate the ODE at.
        intermediate_points: Whether to return the intermediate points.

    Returns:
        The solution to the ODE at the final time. (or at intermediate points if intermediate_points=True)
    """

    if t_arr is None:
        t = torch.FloatTensor([init_t-t_eps, final_t]).to(init_x.device)
    else:
        t = torch.FloatTensor(t_arr).to(init_x.device)

    z = odeint(ode_func, init_x, t, **{"atol": 1e-5, "rtol": 1e-5, **ode_opts})
    if not intermediate_points:
        z = z[-1]
    return z

flow_map(vector_field, x, flow_steps=8, t0=0.0, t1=1.0, integrator='rk4')

Differentiable fixed-step flow map used by shooting strategies.

Source code in flowdesign/strategy/shooting.py

def flow_map(
    vector_field: VectorField | None,
    x: torch.Tensor,
    flow_steps: int = 8,
    t0: float = 0.0,
    t1: float = 1.0,
    integrator: str = "rk4",
) -> torch.Tensor:
    """Differentiable fixed-step flow map used by shooting strategies."""
    if vector_field is None or int(flow_steps) <= 0 or float(t0) == float(t1):
        return x

    steps = int(flow_steps)
    dt = (float(t1) - float(t0)) / steps
    state = x
    time = float(t0)
    method = integrator.lower()
    for _ in range(steps):
        if method == "euler":
            state = state + dt * _call_vector_field(vector_field, time, state)
        elif method == "rk4":
            k1 = _call_vector_field(vector_field, time, state)
            k2 = _call_vector_field(vector_field, time + 0.5 * dt, state + 0.5 * dt * k1)
            k3 = _call_vector_field(vector_field, time + 0.5 * dt, state + 0.5 * dt * k2)
            k4 = _call_vector_field(vector_field, time + dt, state + dt * k3)
            state = state + (dt / 6.0) * (k1 + 2.0 * k2 + 2.0 * k3 + k4)
        else:
            raise ValueError(f"Unsupported integrator: {integrator}")
        time += dt
    return state

flow_trajectory(vector_field, x, flow_steps=8, t0=0.0, t1=1.0, integrator='rk4')

Return inference-time states [T+1, B, D] from source to terminal.

Source code in flowdesign/strategy/shooting.py

def flow_trajectory(
    vector_field: VectorField | None,
    x: torch.Tensor,
    flow_steps: int = 8,
    t0: float = 0.0,
    t1: float = 1.0,
    integrator: str = "rk4",
) -> torch.Tensor:
    """Return inference-time states [T+1, B, D] from source to terminal."""
    if vector_field is None or int(flow_steps) <= 0 or float(t0) == float(t1):
        return x.unsqueeze(0).expand(int(flow_steps) + 1, *x.shape).clone()

    states = [x]
    steps = int(flow_steps)
    dt = (float(t1) - float(t0)) / steps
    state = x
    time = float(t0)
    method = integrator.lower()
    for _ in range(steps):
        if method == "euler":
            state = state + dt * _call_vector_field(vector_field, time, state)
        elif method == "rk4":
            k1 = _call_vector_field(vector_field, time, state)
            k2 = _call_vector_field(vector_field, time + 0.5 * dt, state + 0.5 * dt * k1)
            k3 = _call_vector_field(vector_field, time + 0.5 * dt, state + 0.5 * dt * k2)
            k4 = _call_vector_field(vector_field, time + dt, state + dt * k3)
            state = state + (dt / 6.0) * (k1 + 2.0 * k2 + 2.0 * k3 + k4)
        else:
            raise ValueError(f"Unsupported integrator: {integrator}")
        states.append(state)
        time += dt
    return torch.stack(states, dim=0)