Skip to content

Core API

flowdesign.core

Constraint

Bases: Protocol

Protocol for hard or soft design constraints.

Source code in flowdesign/core/interfaces.py
class Constraint(Protocol):
    """Protocol for hard or soft design constraints."""

    mode: str

    def evaluate(self, x: torch.Tensor, outputs: OutputDict | None = None) -> torch.Tensor:
        ...

    def penalty(self, x: torch.Tensor, outputs: OutputDict | None = None) -> torch.Tensor:
        ...

    def is_satisfied(self, x: torch.Tensor, outputs: OutputDict | None = None) -> torch.Tensor:
        ...

Evaluator

Bases: Protocol

Protocol for surrogates and physics solvers.

Source code in flowdesign/core/interfaces.py
class Evaluator(Protocol):
    """Protocol for surrogates and physics solvers."""

    def __call__(self, x: torch.Tensor, conditions: dict[str, Any] | None = None) -> OutputDict:
        ...

Geometry

Bases: Protocol

Protocol for design parameterizations.

Source code in flowdesign/core/interfaces.py
class Geometry(Protocol):
    """Protocol for design parameterizations."""

    design_dim: int

    def sample_prior(self, n: int, device: str | torch.device | None = None) -> torch.Tensor:
        ...

    def decode(self, params: torch.Tensor) -> Any:
        ...

    def encode(self, raw_geometry: Any) -> torch.Tensor:
        ...

    def validate(self, params: torch.Tensor) -> torch.Tensor:
        ...

Loss

Bases: Protocol

Protocol for differentiable objective terms.

Source code in flowdesign/core/interfaces.py
class Loss(Protocol):
    """Protocol for differentiable objective terms."""

    def __call__(self, x: torch.Tensor, outputs: OutputDict | None = None) -> torch.Tensor:
        ...

DesignProblem dataclass

Central FlowDesign problem object: geometry + evaluator + objectives + constraints.

Source code in flowdesign/core/problem.py
@dataclass
class DesignProblem:
    """Central FlowDesign problem object: geometry + evaluator + objectives + constraints."""

    geometry: Geometry
    evaluator: Evaluator
    objectives: list[Loss] = field(default_factory=list)
    constraints: list[Constraint] = field(default_factory=list)
    conditions: dict[str, Any] | None = None

    def evaluate(self, x: torch.Tensor, conditions: dict[str, Any] | None = None) -> OutputDict:
        cond = self.conditions if conditions is None else conditions
        return ensure_output_dict(self.evaluator(x, conditions=cond))

    def loss(self, x: torch.Tensor, outputs: OutputDict | None = None) -> torch.Tensor:
        if outputs is None:
            outputs = self.evaluate(x)
        batch = x.shape[0] if x.ndim > 1 else 1
        device = x.device
        dtype = x.dtype if torch.is_floating_point(x) else torch.float32
        total = torch.zeros(batch, dtype=dtype, device=device)
        for objective in self.objectives:
            total = total + _as_batch(objective(x, outputs), batch, device, dtype)
        for constraint in self.constraints:
            if getattr(constraint, "mode", "soft") == "soft":
                total = total + _as_batch(constraint.penalty(x, outputs), batch, device, dtype)
        return total

    def check_constraints(self, x: torch.Tensor, outputs: OutputDict | None = None) -> dict[str, torch.Tensor]:
        if outputs is None:
            outputs = self.evaluate(x)
        return {constraint.__class__.__name__: constraint.is_satisfied(x, outputs) for constraint in self.constraints}

    def is_feasible(self, x: torch.Tensor, outputs: OutputDict | None = None) -> torch.Tensor:
        checks = self.check_constraints(x, outputs)
        if not checks:
            batch = x.shape[0] if x.ndim > 1 else 1
            return torch.ones(batch, dtype=torch.bool, device=x.device)
        feasible = None
        for value in checks.values():
            value = value.bool()
            feasible = value if feasible is None else feasible & value
        return feasible

DesignResult dataclass

Standard result object returned by FlowDesign strategies.

Source code in flowdesign/core/result.py
@dataclass
class DesignResult:
    """Standard result object returned by FlowDesign strategies."""

    samples: torch.Tensor
    outputs: dict[str, torch.Tensor] = field(default_factory=dict)
    losses: torch.Tensor | None = None
    feasibility: torch.Tensor | None = None
    metrics: dict[str, Any] = field(default_factory=dict)
    metadata: dict[str, Any] = field(default_factory=dict)

    def __post_init__(self) -> None:
        self.samples = _as_cpu(self.samples)
        self.outputs = _as_cpu(self.outputs)
        if self.losses is not None:
            self.losses = _as_cpu(self.losses).reshape(-1)
        if self.feasibility is not None:
            self.feasibility = _as_cpu(self.feasibility).reshape(-1).bool()

    def best(self, feasible_only: bool = False) -> dict[str, Any]:
        """Return the best sample row by minimum loss."""
        if self.samples.numel() == 0:
            raise ValueError("Cannot select best sample from an empty result.")
        indices = torch.arange(self.samples.shape[0])
        if feasible_only and self.feasibility is not None:
            mask = self.feasibility.bool()
            if mask.any():
                indices = indices[mask]
        if self.losses is None:
            best_index = int(indices[0].item())
        else:
            local = torch.argmin(self.losses[indices])
            best_index = int(indices[local].item())
        return self._row(best_index)

    def filter_feasible(self) -> "DesignResult":
        """Return a new result containing only feasible samples."""
        if self.feasibility is None:
            return self
        mask = self.feasibility.bool()
        outputs = {key: value[mask] for key, value in self.outputs.items()}
        losses = self.losses[mask] if self.losses is not None else None
        return DesignResult(
            samples=self.samples[mask],
            outputs=outputs,
            losses=losses,
            feasibility=self.feasibility[mask],
            metrics=dict(self.metrics),
            metadata=dict(self.metadata),
        )

    def to_dataframe(self):
        """Return a pandas DataFrame when pandas is available, otherwise rows."""
        rows = [self._row(idx) for idx in range(self.samples.shape[0])]
        try:
            import pandas as pd

            return pd.DataFrame(rows)
        except Exception:
            return rows

    def save(self, path: str) -> None:
        """Save tensors and tabular metadata to a directory."""
        os.makedirs(path, exist_ok=True)
        torch.save(
            {
                "samples": self.samples,
                "outputs": self.outputs,
                "losses": self.losses,
                "feasibility": self.feasibility,
                "metrics": self.metrics,
                "metadata": self.metadata,
            },
            os.path.join(path, "design_result.pt"),
        )
        with open(os.path.join(path, "metrics.json"), "w", encoding="utf-8") as handle:
            json.dump(_jsonable({"metrics": self.metrics, "metadata": self.metadata}), handle, indent=2)
        rows = self.to_dataframe()
        if hasattr(rows, "to_csv"):
            rows.to_csv(os.path.join(path, "design_result.csv"), index=False)
        else:
            import csv

            if rows:
                with open(os.path.join(path, "design_result.csv"), "w", newline="", encoding="utf-8") as handle:
                    writer = csv.DictWriter(handle, fieldnames=list(rows[0].keys()))
                    writer.writeheader()
                    writer.writerows(rows)

    def _row(self, idx: int) -> dict[str, Any]:
        row: dict[str, Any] = {"sample_id": idx}
        if self.losses is not None:
            row["loss"] = float(self.losses[idx].item())
        if self.feasibility is not None:
            row["feasible"] = bool(self.feasibility[idx].item())
        for key, value in self.outputs.items():
            if value.ndim == 0:
                row[key] = float(value.item())
            elif value.shape[0] == self.samples.shape[0]:
                item = value[idx]
                row[key] = float(item.item()) if item.numel() == 1 else item.detach().cpu().tolist()
        row["sample"] = self.samples[idx].detach().cpu().tolist()
        return row

best(feasible_only=False)

Return the best sample row by minimum loss.

Source code in flowdesign/core/result.py
def best(self, feasible_only: bool = False) -> dict[str, Any]:
    """Return the best sample row by minimum loss."""
    if self.samples.numel() == 0:
        raise ValueError("Cannot select best sample from an empty result.")
    indices = torch.arange(self.samples.shape[0])
    if feasible_only and self.feasibility is not None:
        mask = self.feasibility.bool()
        if mask.any():
            indices = indices[mask]
    if self.losses is None:
        best_index = int(indices[0].item())
    else:
        local = torch.argmin(self.losses[indices])
        best_index = int(indices[local].item())
    return self._row(best_index)

filter_feasible()

Return a new result containing only feasible samples.

Source code in flowdesign/core/result.py
def filter_feasible(self) -> "DesignResult":
    """Return a new result containing only feasible samples."""
    if self.feasibility is None:
        return self
    mask = self.feasibility.bool()
    outputs = {key: value[mask] for key, value in self.outputs.items()}
    losses = self.losses[mask] if self.losses is not None else None
    return DesignResult(
        samples=self.samples[mask],
        outputs=outputs,
        losses=losses,
        feasibility=self.feasibility[mask],
        metrics=dict(self.metrics),
        metadata=dict(self.metadata),
    )

to_dataframe()

Return a pandas DataFrame when pandas is available, otherwise rows.

Source code in flowdesign/core/result.py
def to_dataframe(self):
    """Return a pandas DataFrame when pandas is available, otherwise rows."""
    rows = [self._row(idx) for idx in range(self.samples.shape[0])]
    try:
        import pandas as pd

        return pd.DataFrame(rows)
    except Exception:
        return rows

save(path)

Save tensors and tabular metadata to a directory.

Source code in flowdesign/core/result.py
def save(self, path: str) -> None:
    """Save tensors and tabular metadata to a directory."""
    os.makedirs(path, exist_ok=True)
    torch.save(
        {
            "samples": self.samples,
            "outputs": self.outputs,
            "losses": self.losses,
            "feasibility": self.feasibility,
            "metrics": self.metrics,
            "metadata": self.metadata,
        },
        os.path.join(path, "design_result.pt"),
    )
    with open(os.path.join(path, "metrics.json"), "w", encoding="utf-8") as handle:
        json.dump(_jsonable({"metrics": self.metrics, "metadata": self.metadata}), handle, indent=2)
    rows = self.to_dataframe()
    if hasattr(rows, "to_csv"):
        rows.to_csv(os.path.join(path, "design_result.csv"), index=False)
    else:
        import csv

        if rows:
            with open(os.path.join(path, "design_result.csv"), "w", newline="", encoding="utf-8") as handle:
                writer = csv.DictWriter(handle, fieldnames=list(rows[0].keys()))
                writer.writeheader()
                writer.writerows(rows)