moospread 0.1.0__py3-none-any.whl

moospread/utils/offline_utils/__init__.py

@@ -0,0 +1,10 @@
+from moospread.utils.offline_utils.proxies import (SingleModelBaseTrainer, 
+                                                MultipleModels,
+                                                SingleModel,
+                                                offdata_get_dataloader)
+from moospread.utils.offline_utils.handle_task import (offdata_min_max_normalize, 
+                                                    offdata_min_max_denormalize, 
+                                                    offdata_z_score_normalize,
+                                                    offdata_z_score_denormalize,
+                                                    offdata_to_integers, 
+                                                    offdata_to_logits)

moospread/utils/offline_utils/handle_task.py

@@ -0,0 +1,203 @@
+import torch
+import torch.nn.functional as F
+from typing import List, Tuple
+
+
+def one_hot(a: torch.Tensor, num_classes: int) -> torch.Tensor:
+    """Convert an integer tensor to a float32 one-hot tensor using PyTorch.
+
+    Args:
+        a: Long/Int tensor of arbitrary shape containing class indices.
+        num_classes: Total number of classes.
+    Returns:
+        Tensor with shape ``a.shape + (num_classes,)`` and dtype float32.
+    """
+    if torch.is_floating_point(a):
+        raise ValueError("cannot one-hot encode non-integers (got floating dtype)")
+    return F.one_hot(a.to(torch.long), num_classes=num_classes).to(torch.float32)
+
+
+def help_to_logits(
+    x: torch.Tensor, num_classes: int, soft_interpolation: float = 0.6
+) -> torch.Tensor:
+    """Convert integer labels to *logit* representation with a soft uniform prior.
+
+    Mirrors the original NumPy implementation but in PyTorch.
+
+    Args:
+        x: Int tensor of shape ``(n_samples, 1)``.
+        num_classes: Number of classes for this position.
+        soft_interpolation: Interpolate between one-hot (Dirac) and uniform prior.
+    Returns:
+        Float tensor of shape ``(n_samples, 1, num_classes-1)``.
+    """
+    if torch.is_floating_point(x):
+        raise ValueError("cannot convert non-integers to logits")
+
+    device = x.device
+
+    # one-hot (n, 1, C)
+    one_hot_x = one_hot(x, num_classes=num_classes)  # float32
+
+    # uniform prior to interpolate with
+    uniform_prior = torch.full_like(one_hot_x, 1.0 / float(num_classes), device=device)
+
+    # convex combination
+    soft_x = soft_interpolation * one_hot_x + (1.0 - soft_interpolation) * uniform_prior
+
+    # convert to log probabilities
+    log_p = torch.log(soft_x)
+
+    # remove one degree of freedom: subtract the first component and drop it
+    return (log_p[:, :, 1:] - log_p[:, :, :1]).to(torch.float32)
+
+
+def offdata_to_logits(
+    x: torch.Tensor,
+    num_classes_on_each_position: List[int],
+    soft_interpolation: float = 0.6,
+) -> torch.Tensor:
+    """Convert a sequence of categorical integers to concatenated logits.
+
+    For each sequence position ``i`` with ``k_i`` classes, we form a length ``k_i-1``
+    logit vector (after removing the redundant degree of freedom) and then concatenate
+    all positions along the last axis.
+
+    Args:
+        x: Int tensor of shape ``(n_samples, seq_len)``.
+        num_classes_on_each_position: list of class counts per position.
+        soft_interpolation: Interpolation factor in ``[0, 1]``.
+    Returns:
+        Float tensor of shape ``(n_samples, sum_i (k_i - 1))`` (matches the original).
+    """
+    # Original code adds +1, then for positions with exactly 1 class, adds another +1
+    # to introduce a dummy class. We reproduce that behavior.
+    num_classes = [c + 1 for c in num_classes_on_each_position]
+    num_classes = [c + 1 if c == 1 else c for c in num_classes]
+
+    logits = []
+    seq_len = len(num_classes)
+    for i in range(seq_len):
+        temp_x = x[:, i].reshape(-1, 1)
+        logits.append(help_to_logits(temp_x, num_classes[i], soft_interpolation))
+
+    # concatenate along the last dim, then squeeze (to mimic NumPy's .squeeze())
+    out = torch.cat(logits, dim=2).squeeze()
+    return out
+
+
+def help_to_integers(x: torch.Tensor, true_num_of_classes: int) -> torch.Tensor:
+    """Convert per-position logits back to class integers.
+
+    Args:
+        x: Float tensor of shape ``(n_samples, 1, k-1)`` for a position with ``k`` classes.
+        true_num_of_classes: The (possibly dummy-adjusted) number of classes ``k``.
+    Returns:
+        Int tensor of shape ``(n_samples, 1)`` with the selected classes.
+    """
+    if not torch.is_floating_point(x):
+        raise ValueError("cannot convert non-floats to integers")
+
+    # Special-case: if k == 1 (RFP-Exact-v0 path), always return zeros
+    if true_num_of_classes == 1:
+        return torch.zeros(x.shape[:-1], dtype=torch.int64, device=x.device)
+
+    # Pad a leading zero component and take argmax along class dim
+    # x shape: (n, 1, k-1) -> pad to (n, 1, k)
+    x_padded = F.pad(x, pad=(1, 0))  # pad last dim: (left=1, right=0)
+    return torch.argmax(x_padded, dim=-1).to(torch.int64)
+
+
+def offdata_to_integers(x: torch.Tensor, num_classes_on_each_position: List[int]) -> torch.Tensor:
+    """Invert ``to_logits``: recover integer classes for each sequence position.
+
+    Note: This follows the original slicing behavior where the concatenated
+    per-position logits are packed along the second dimension (after a squeeze).
+
+    Args:
+        x: Float tensor with shape ``(n_samples, total_concatenated)`` where
+           ``total_concatenated = sum_i (k_i - 1)`` after the same class-count
+           adjustments used in ``to_logits``.
+        num_classes_on_each_position: List of true class counts per position.
+    Returns:
+        Int tensor of shape ``(n_samples, seq_len)``.
+    """
+    # Reproduce the same class-count adjustment as in to_logits
+    true_num_classes = [c + 1 for c in num_classes_on_each_position]
+    num_classes = [c + 1 if c == 1 else c for c in true_num_classes]
+
+    integers = []
+    start = 0
+    for k in num_classes:
+        width = k - 1
+        # Slice along dim=1 to match the original implementation
+        temp_x = x[:, start : start + width].reshape(-1, 1, width)
+        integers.append(help_to_integers(temp_x, k))
+        start += width
+
+    # Concatenate along seq dimension (dim=1)
+    return torch.cat(integers, dim=1)
+
+
+def offdata_z_score_normalize(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Z-score normalize features columnwise (match NumPy semantics).
+
+    Args:
+        x: Float tensor of shape ``(n_samples, n_dim)``.
+    Returns:
+        (x_norm, mean, std) where each has shape compatible with broadcasting.
+    """
+    if not torch.is_floating_point(x):
+        raise ValueError("cannot normalize discrete design values")
+
+    mean = torch.mean(x, dim=0)
+    # NumPy's np.std uses population std (ddof=0) by default -> unbiased=False
+    std = torch.std(x, dim=0, unbiased=False)
+    x_norm = (x - mean) / std
+    return x_norm, mean, std
+
+
+def offdata_z_score_denormalize(x: torch.Tensor, 
+                                x_mean: torch.Tensor, 
+                                x_std: torch.Tensor) -> torch.Tensor:
+    """Invert z-score normalization.
+
+    Args:
+        x: Float tensor ``(n_samples, n_dim)``.
+        x_mean: Mean used during normalization.
+        x_std: Std used during normalization.
+    Returns:
+        Denormalized tensor.
+    """
+    if not torch.is_floating_point(x):
+        raise ValueError("cannot denormalize discrete design values")
+    return x * x_std.to(x.device) + x_mean.to(x.device)
+
+
+def offdata_min_max_normalize(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Min-max normalize features columnwise.
+
+    Args:
+        x: Tensor of shape ``(n_samples, n_dim)``.
+    Returns:
+        (x_norm, x_min, x_max)
+    """
+    x_min = torch.min(x, dim=0).values
+    x_max = torch.max(x, dim=0).values
+    x_norm = (x - x_min) / (x_max - x_min)
+    return x_norm, x_min, x_max
+
+
+def offdata_min_max_denormalize(x: torch.Tensor, 
+                                x_min: torch.Tensor, 
+                                x_max: torch.Tensor) -> torch.Tensor:
+    """Invert min-max normalization.
+
+    Args:
+        x: Tensor of shape ``(n_samples, n_dim)``.
+        x_min: Per-dimension min.
+        x_max: Per-dimension max.
+    Returns:
+        Denormalized tensor.
+    """
+    return x * (x_max.to(x.device) - x_min.to(x.device)) + x_min.to(x.device)

moospread/utils/offline_utils/proxies.py

@@ -0,0 +1,338 @@
+"""
+Adapted from: https://github.com/lamda-bbo/offline-moo/blob/main/off_moo_baselines/multiple_models/nets.py
+"""
+
+import os
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.optim import Adam
+from torch.utils.data import DataLoader, TensorDataset, random_split
+from tqdm import tqdm
+
+def offdata_get_dataloader(
+    X,
+    y,
+    train_ratio: float = 0.9,
+    batch_size: int = 32,
+):
+
+    tensor_dataset = TensorDataset(X, y)
+    lengths = [
+        int(train_ratio * len(tensor_dataset)),
+        len(tensor_dataset) - int(train_ratio * len(tensor_dataset)),
+    ]
+    train_dataset, val_dataset = random_split(tensor_dataset, lengths)
+
+    train_loader = DataLoader(
+        train_dataset, batch_size=batch_size, shuffle=True, drop_last=False
+    )
+
+    val_loader = DataLoader(
+        val_dataset, batch_size=batch_size * 4, shuffle=False, drop_last=False
+    )
+
+    return train_loader, val_loader
+
+def spearman_correlation(x, y):
+    n = x.size(0)
+    _, rank_x = x.sort(0)
+    _, rank_y = y.sort(0)
+
+    d = rank_x - rank_y
+    d_squared_sum = (d**2).sum(0).float()
+
+    rho = 1 - (6 * d_squared_sum) / (n * (n**2 - 1))
+    return rho
+
+class SingleModelBaseTrainer(nn.Module):
+
+    def __init__(self, model, which_obj, args):
+        super(SingleModelBaseTrainer, self).__init__()
+        self.args = args
+
+        self.forward_lr = args["proxies_lr"]
+        self.forward_lr_decay = args["proxies_lr_decay"]
+        self.n_epochs = args["proxies_epochs"]
+        self.device = args["device"]
+        self.verbose = args["verbose"]
+
+        self.model = model
+
+        self.which_obj = which_obj
+
+        self.forward_opt = Adam(model.parameters(), lr=args["proxies_lr"])
+        self.train_criterion = lambda yhat, y: (
+            torch.sum(torch.mean((yhat - y) ** 2, dim=1))
+        )
+        self.mse_criterion = nn.MSELoss()
+
+    def _evaluate_performance(self, statistics, epoch, train_loader, val_loader):
+        self.model.eval()
+        with torch.no_grad():
+            y_all = torch.zeros((0, self.n_obj)).to(self.device)
+            outputs_all = torch.zeros((0, self.n_obj)).to(self.device)
+            for (
+                batch_x,
+                batch_y,
+            ) in train_loader:
+                batch_x = batch_x.to(self.device)
+                batch_y = batch_y.to(self.device)
+
+                y_all = torch.cat((y_all, batch_y), dim=0)
+                outputs = self.model(batch_x)
+                outputs_all = torch.cat((outputs_all, outputs), dim=0)
+
+            train_mse = self.mse_criterion(outputs_all, y_all)
+            train_corr = spearman_correlation(outputs_all, y_all)
+            train_pcc = self.compute_pcc(outputs_all, y_all)
+
+            statistics[f"model_{self.which_obj}/train/mse"] = train_mse.item()
+            for i in range(self.n_obj):
+                statistics[f"model_{self.which_obj}/train/rank_corr_{i + 1}"] = (
+                    train_corr[i].item()
+                )
+            # if self.verbose:
+            #     print(
+            #         "Epoch [{}/{}], MSE: {:}, PCC: {:}".format(
+            #             epoch + 1, self.n_epochs, train_mse.item(), train_pcc.item()
+            #         )
+            #     )
+
+        with torch.no_grad():
+            y_all = torch.zeros((0, self.n_obj)).to(self.device)
+            outputs_all = torch.zeros((0, self.n_obj)).to(self.device)
+
+            for batch_x, batch_y in val_loader:
+                batch_x = batch_x.to(self.device)
+                batch_y = batch_y.to(self.device)
+
+                y_all = torch.cat((y_all, batch_y), dim=0)
+                outputs = self.model(batch_x)
+                outputs_all = torch.cat((outputs_all, outputs))
+
+            val_mse = self.mse_criterion(outputs_all, y_all)
+            val_corr = spearman_correlation(outputs_all, y_all)
+            val_pcc = self.compute_pcc(outputs_all, y_all)
+
+            statistics[f"model_{self.which_obj}/valid/mse"] = val_mse.item()
+            for i in range(self.n_obj):
+                statistics[f"model_{self.which_obj}/valid/rank_corr_{i + 1}"] = (
+                    val_corr[i].item()
+                )
+
+            # if self.verbose:
+            #     print(
+            #         "Valid MSE: {:}, Valid PCC: {:}".format(val_mse.item(), val_pcc.item())
+            #     )
+    
+            if val_pcc.item() > self.min_pcc:
+                # if self.verbose:
+                #     print("🌸 New best epoch! 🌸")
+                self.min_pcc = val_pcc.item()
+                self.model.save(val_pcc=self.min_pcc)
+        return statistics
+
+    def launch(
+        self,
+        train_loader=None,
+        val_loader=None,
+        retrain_model: bool = True,
+    ):
+
+        def update_lr(optimizer, lr):
+            for param_group in optimizer.param_groups:
+                param_group["lr"] = lr
+
+        if not retrain_model and os.path.exists(self.model.save_path):
+            self.model.load()
+            return
+
+        assert train_loader is not None
+        assert val_loader is not None
+
+        self.n_obj = None
+        self.min_pcc = -1.0
+        statistics = {}
+
+        with tqdm(
+            total=self.n_epochs,
+            desc=f"Proxy Training",
+            unit="epoch",
+        ) as pbar:
+            
+            for epoch in range(self.n_epochs):
+                self.model.train()
+
+                losses = []
+                epoch_loss = 0.0
+                for batch_x, batch_y in train_loader:
+                    batch_x = batch_x.to(self.device)
+                    batch_y = batch_y.to(self.device)
+                    if self.n_obj is None:
+                        self.n_obj = batch_y.shape[1]
+
+                    self.forward_opt.zero_grad()
+                    outputs = self.model(batch_x)
+                    loss = self.train_criterion(outputs, batch_y)
+                    losses.append(loss.item() / batch_x.size(0))
+                    loss.backward()
+                    self.forward_opt.step()
+                    epoch_loss += loss.item()
+
+                statistics[f"model_{self.which_obj}/train/loss/mean"] = np.array(
+                    losses
+                ).mean()
+                statistics[f"model_{self.which_obj}/train/loss/std"] = np.array(
+                    losses
+                ).std()
+                statistics[f"model_{self.which_obj}/train/loss/max"] = np.array(
+                    losses
+                ).max()
+
+                self._evaluate_performance(statistics, epoch, train_loader, val_loader)
+
+                statistics[f"model_{self.which_obj}/train/lr"] = self.forward_lr
+                self.forward_lr *= self.forward_lr_decay
+                update_lr(self.forward_opt, self.forward_lr)
+
+                epoch_loss /= len(train_loader)
+                pbar.set_postfix({
+                        "Loss": epoch_loss,
+                        })
+                pbar.update(1)
+
+    def compute_pcc(self, valid_preds, valid_labels):
+        vx = valid_preds - torch.mean(valid_preds)
+        vy = valid_labels - torch.mean(valid_labels)
+        pcc = torch.sum(vx * vy) / (
+            torch.sqrt(torch.sum(vx**2) + 1e-12) * torch.sqrt(torch.sum(vy**2) + 1e-12)
+        )
+        return pcc
+
+
+class MultipleModels(nn.Module):
+    def __init__(
+        self, n_dim, n_obj, hidden_size, train_mode, device,
+        save_dir=None, save_prefix=None
+    ):
+        super(MultipleModels, self).__init__()
+        self.n_dim = n_dim
+        self.n_obj = n_obj
+        self.device = device
+
+        self.obj2model = {}
+        self.hidden_size = hidden_size
+        self.train_mode = train_mode
+
+        self.save_dir = save_dir
+        self.save_prefix = save_prefix
+        if self.save_dir is not None:
+            os.makedirs(self.save_dir, exist_ok=True)
+
+        for obj in range(self.n_obj):
+            self.create_models(obj)
+
+    def create_models(self, learning_obj):
+        model = SingleModel
+        new_model = model(
+            self.n_dim,
+            self.hidden_size,
+            which_obj=learning_obj,
+            device=self.device,
+            save_dir=self.save_dir,
+            save_prefix=self.save_prefix,
+        )
+        self.obj2model[learning_obj] = new_model
+
+    def set_kwargs(self, device=None, dtype=torch.float32):
+        for model in self.obj2model.values():
+            model.set_kwargs(device=device, dtype=dtype)
+            model.to(device=device, dtype=dtype)
+
+    def forward(self, x, forward_objs=None):
+        if forward_objs is None:
+            forward_objs = list(self.obj2model.keys())
+        x = [self.obj2model[obj](x) for obj in forward_objs]
+        x = torch.cat(x, dim=1)
+        return x
+
+
+activate_functions = [nn.LeakyReLU(), nn.LeakyReLU()]
+
+
+class SingleModel(nn.Module):
+    def __init__(
+        self, input_size, hidden_size, which_obj, device,
+        save_dir=None, save_prefix=None
+    ):
+        super(SingleModel, self).__init__()
+        self.n_dim = input_size
+        self.n_obj = 1
+        self.which_obj = which_obj
+        self.activate_functions = activate_functions
+        self.device = device
+
+        layers = []
+        layers.append(nn.Linear(input_size, hidden_size[0]))
+        for i in range(len(hidden_size) - 1):
+            layers.append(nn.Linear(hidden_size[i], hidden_size[i + 1]))
+        layers.append(nn.Linear(hidden_size[len(hidden_size) - 1], 1))
+
+        self.layers = nn.Sequential(*layers)
+        self.hidden_size = hidden_size
+
+        self.save_path = os.path.join(save_dir, f"{save_prefix}-{which_obj}.pt")
+
+    def forward(self, x):
+        for i in range(len(self.hidden_size)):
+            x = self.layers[i](x)
+            x = self.activate_functions[i](x)
+
+        x = self.layers[len(self.hidden_size)](x)
+        out = x
+
+        return out
+
+    def set_kwargs(self, device=None, dtype=torch.float32):
+        self.to(device=device, dtype=dtype)
+
+    def check_model_path_exist(self, save_path=None):
+        assert (
+            self.save_path is not None or save_path is not None
+        ), "save path should be specified"
+        if save_path is None:
+            save_path = self.save_path
+        return os.path.exists(save_path)
+
+    def save(self, val_pcc=None, save_path=None):
+        assert (
+            self.save_path is not None or save_path is not None
+        ), "save path should be specified"
+        if save_path is None:
+            save_path = self.save_path
+
+        self = self.to("cpu")
+        checkpoint = {
+            "model_state_dict": self.state_dict(),
+        }
+        if val_pcc is not None:
+            checkpoint["valid_pcc"] = val_pcc
+
+        torch.save(checkpoint, save_path)
+        self = self.to(self.device)
+
+    def load(self, save_path=None):
+        assert (
+            self.save_path is not None or save_path is not None
+        ), "save path should be specified"
+        if save_path is None:
+            save_path = self.save_path
+
+        checkpoint = torch.load(save_path, weights_only=False)
+        self.load_state_dict(checkpoint["model_state_dict"])
+        valid_pcc = checkpoint["valid_pcc"]
+        print(
+            f"Successfully load trained proxy model from {save_path} "
+            f"with valid PCC = {valid_pcc}"
+        )

moospread/utils/spread_utils.py

@@ -0,0 +1,91 @@
+from torch.utils.data import DataLoader, TensorDataset
+import dis
+import inspect
+import ast
+import textwrap
+
+def get_ddpm_dataloader(X,
+                        y,
+                        validation_split=0.1, 
+                        batch_size=32):
+    val_dataloader = None
+    if validation_split > 0.0:
+        train_size = int(X.shape[0] - int(X.shape[0] * validation_split))
+        X_val = X[train_size:]
+        y_val = y[train_size:]
+        X = X[:train_size]
+        y = y[:train_size]
+
+        tensor_x_val = X_val.float() 
+        tensor_y_val = y_val.float() 
+        dataset_val = TensorDataset(tensor_x_val, tensor_y_val)
+        val_dataloader = DataLoader(
+                    dataset_val,
+                    batch_size=batch_size,
+                    shuffle=False,
+                    drop_last=True,
+                )
+    tensor_x = X.float() 
+    tensor_y = y.float() 
+    dataset_train = TensorDataset(tensor_x, tensor_y)
+    train_dataloader = DataLoader(
+                dataset_train,
+                batch_size=batch_size,
+                shuffle=True,
+                drop_last=True,
+            )
+    
+    return train_dataloader, val_dataloader
+
+def is_pass_function(func) -> bool:
+    """
+    Return True iff the function body is effectively just `pass`
+    (optionally with a docstring). Works on Python 3.8–3.12+.
+    """
+    # -------- Try AST first (most reliable) --------
+    try:
+        src = inspect.getsource(func)
+    except (OSError, TypeError):
+        src = None
+
+    if src:
+        mod = ast.parse(textwrap.dedent(src))
+        fn = next((n for n in ast.walk(mod)
+                   if isinstance(n, (ast.FunctionDef, ast.AsyncFunctionDef))), None)
+        if fn:
+            body = list(fn.body)
+            # drop docstring if present
+            if body and isinstance(body[0], ast.Expr) and \
+               isinstance(getattr(body[0], "value", None), ast.Constant) and \
+               isinstance(body[0].value.value, str):
+                body = body[1:]
+            # True if remaining stmts are all Pass (or empty)
+            return all(isinstance(n, ast.Pass) for n in body)
+
+    # -------- Fallback: bytecode pattern (version-tolerant) --------
+    instrs = list(dis.get_instructions(func))
+
+    # remove version-specific noise
+    noise = {"RESUME", "CACHE", "EXTENDED_ARG", "NOP"}
+    core = [i for i in instrs if i.opname not in noise]
+
+    # strip docstring store: LOAD_CONST <str>; STORE_* __doc__
+    i = 0
+    while i + 1 < len(core):
+        a, b = core[i], core[i + 1]
+        if (a.opname == "LOAD_CONST" and isinstance(a.argval, str)
+            and b.opname in {"STORE_NAME", "STORE_FAST"} and b.argval == "__doc__"):
+            del core[i:i+2]
+            continue
+        i += 1
+
+    # Accept either:
+    #  - LOAD_CONST None; RETURN_VALUE
+    #  - RETURN_CONST None   (3.12+)
+    if len(core) == 2 and core[0].opname == "LOAD_CONST" and core[0].argval is None \
+       and core[1].opname == "RETURN_VALUE":
+        return True
+    if len(core) == 1 and core[0].opname == "RETURN_CONST" and core[0].argval is None:
+        return True
+
+    return False