blnetwork 0.1.0__py3-none-any.whl

blnetwork/export.py

@@ -0,0 +1,269 @@
+from __future__ import annotations
+
+from typing import List, Optional, Tuple, Dict
+
+import numpy as np
+import torch
+import sys
+import io
+import torch.nn.functional as F
+
+def _fmt_num(x: float, ndigits: int = 4) -> str:
+    s = f"{x:.{ndigits}f}"
+    s = s.rstrip("0").rstrip(".")
+    return s if s != "" else "0"
+
+
+def _safe_numpy(t):
+    if t is None:
+        return None
+    if torch.is_tensor(t):
+        return t.detach().cpu().numpy()
+    return np.asarray(t)
+
+
+def _get_bl_unit(block):
+    unit = getattr(block, "unit", None)
+    if unit is None:
+        raise AttributeError("Block has no attribute 'unit' (expected BLBlock.unit).")
+    return unit
+
+
+def _get_lambdas(unit):
+    lam = unit.lam 
+    if getattr(unit, "constrain_lambda", False):
+        eps = float(getattr(unit, "eps", 1e-8))
+        lam = F.softplus(lam) + eps
+    
+    lam_u = _safe_numpy(lam[0])
+    lam_c = _safe_numpy(lam[1])
+    lam_t = _safe_numpy(lam[2])
+    
+    return lam_u, lam_c, lam_t
+
+
+def _get_backbone(model):
+    backbone = getattr(model, "backbone", None)
+    if backbone is not None:
+        return backbone
+    if hasattr(model, "blocks"):
+        return model
+    raise AttributeError("Cannot find backbone. Expected model.backbone or model.blocks.")
+
+
+def _get_blocks(model) -> list:
+    backbone = _get_backbone(model)
+    blocks = getattr(backbone, "blocks", None)
+    if blocks is None:
+        raise AttributeError("Backbone has no attribute 'blocks'.")
+    return list(blocks)
+
+
+def _get_hidden_dims(model) -> Optional[Tuple[int, ...]]:
+    backbone = _get_backbone(model)
+    v = getattr(backbone, "hidden_dims", None)
+    if v is None:
+        return None
+    return tuple(map(int, v))
+
+
+def _get_output_linears(model) -> Dict[str, torch.nn.Linear]:
+    outs: Dict[str, torch.nn.Linear] = {}
+
+    lin = getattr(model, "linear_out", None)
+    if lin is not None and hasattr(lin, "weight"):
+        outs["Output Layer (Discrete)"] = lin
+        return outs
+
+    head = getattr(model, "head", None)
+    if head is None:
+        return outs
+    
+    if isinstance(head, torch.nn.Linear):
+        outs["Output Layer"] = head
+        return outs
+
+    lin = getattr(head, "linear", None)
+    if lin is not None and hasattr(lin, "weight"):
+        outs["Output Layer"] = lin
+
+    return outs
+
+
+def _emit_part_lines(
+    part_name: str,
+    lam_val: float,
+    w_row: np.ndarray,
+    b_val: float,
+    feature_names: List[str],
+    ndigits: int = 4,
+    tol: float = 0.0,
+) -> List[str]:
+    lines = []
+    lines.append(f"{part_name}")
+    lines.append(f"lambda {_fmt_num(float(lam_val), ndigits)}")
+
+    for name, w in zip(feature_names, w_row):
+        w = float(w)
+        if abs(w) > tol:
+            lines.append(f"----{name}     {_fmt_num(w, ndigits)}")
+
+    if abs(float(b_val)) > tol:
+        lines.append(f"----C     {_fmt_num(float(b_val), ndigits)}")
+
+    return lines
+
+
+def _print_blocks(
+    block,
+    feature_names: List[str],
+    layer_idx: int,
+    ndigits: int = 4,
+    tol: float = 0.0,
+) -> int:
+
+    unit = _get_bl_unit(block)
+
+    num_basis = int(unit.lin_u.out_features)
+
+    lam_u, lam_c, lam_t = _get_lambdas(unit)
+
+    w_u = _safe_numpy(unit.lin_u.weight)
+    b_u = _safe_numpy(unit.lin_u.bias)
+    w_c = _safe_numpy(unit.lin_c.weight)
+    b_c = _safe_numpy(unit.lin_c.bias)
+    w_t = _safe_numpy(unit.lin_t.weight)
+    b_t = _safe_numpy(unit.lin_t.bias)
+
+    for j in range(num_basis):
+        block_id = j + 1
+        print(f"--B{layer_idx}{block_id}")
+
+        # U part
+        lines = _emit_part_lines(
+            "U",
+            lam_val=float(lam_u[j]),
+            w_row=w_u[j],
+            b_val=float(b_u[j]),
+            feature_names=feature_names,
+            ndigits=ndigits,
+            tol=tol,
+        )
+        for ln in lines:
+            print(ln)
+
+        # C part
+        lines = _emit_part_lines(
+            "C",
+            lam_val=float(lam_c[j]),
+            w_row=w_c[j],
+            b_val=float(b_c[j]),
+            feature_names=feature_names,
+            ndigits=ndigits,
+            tol=tol,
+        )
+        for ln in lines:
+            print(ln)
+
+        # T part
+        lines = _emit_part_lines(
+            "T",
+            lam_val=float(lam_t[j]),
+            w_row=w_t[j],
+            b_val=float(b_t[j]),
+            feature_names=feature_names,
+            ndigits=ndigits,
+            tol=tol,
+        )
+        for ln in lines:
+            print(ln)
+
+        print("")  
+
+    return num_basis
+
+
+def _print_core(model, blocks, hidden_dims, feat_names, ndigits, tol, title="BL Model Structure"):
+    print("=" * 72)
+    print(title)
+    print("=" * 72)
+
+    print(f"hidden_dims = {hidden_dims}")
+    print(f"feature_dim = {len(feat_names)}")
+    print("")
+
+    current_feat_names = feat_names
+    
+    for layer_idx, block in enumerate(blocks, start=1):
+        num_basis = _print_blocks(
+            block=block,
+            feature_names=current_feat_names,
+            layer_idx=layer_idx,
+            ndigits=ndigits,
+            tol=tol,
+        )
+
+        if layer_idx < len(blocks):
+            current_feat_names = [f"B{layer_idx}{i+1}" for i in range(num_basis)]
+
+    outs = _get_output_linears(model)
+    if len(outs) > 0:
+        print("=" * 72)
+        print("OUTPUT LINEAR(S)")
+        print("=" * 72)
+        for name, lin in outs.items():
+            w = _safe_numpy(lin.weight)
+            b = _safe_numpy(lin.bias) if lin.bias is not None else None
+            print(name)
+            if getattr(w, "size", 0) and w.size <= 16:
+                w_flat = [float(x) for x in w.reshape(-1)]
+                w_fmt = ", ".join(_fmt_num(x, ndigits) for x in w_flat)
+                print(f"weight values = [{w_fmt}]")
+
+            if b is None:
+                print("bias = None")
+            else:
+                if getattr(b, "size", 0) and b.size <= 16:
+                    b_flat = [float(x) for x in b.reshape(-1)]
+                    b_fmt = ", ".join(_fmt_num(x, ndigits) for x in b_flat)
+                    print(f"bias values = [{b_fmt}]")
+            print("")
+
+def export_structure(
+    model,
+    df=None,
+    feature_names=None,
+    txt_path: Optional[str] = None,
+    ndigits: int = 4,
+    tol: float = 0.0,
+    title: str = "BL Model Structure",
+):
+
+    blocks = _get_blocks(model)
+    hidden_dims = _get_hidden_dims(model)
+
+    if feature_names is not None:
+        feat_names = list(feature_names)
+    elif df is not None:
+        feat_names = list(df.columns)
+    else:
+        unit0 = _get_bl_unit(blocks[0])
+        if not hasattr(unit0, "lin_u"):
+            raise AttributeError("Cannot infer input dim from BLUnit. Expected 'lin_u'.")
+        in_dim = int(unit0.lin_u.in_features)
+        feat_names = [f"x{i+1}" for i in range(in_dim)]
+
+    if txt_path is None:
+        _print_core(model, blocks, hidden_dims, feat_names, ndigits, tol, title)
+        return
+
+    old = sys.stdout
+    buf = io.StringIO()
+    sys.stdout = buf
+    try:
+        _print_core(model, blocks, hidden_dims, feat_names, ndigits, tol, title)
+    finally:
+        sys.stdout = old
+
+    with open(txt_path, "w", encoding="utf-8") as f:
+        f.write(buf.getvalue())

blnetwork/inference/__init__.py

@@ -0,0 +1,10 @@
+from __future__ import annotations
+
+from .continuous import predict_continuous
+from .discrete import predict_class_discrete, predict_proba_discrete
+
+__all__ = [
+	"predict_continuous",
+	"predict_class_discrete",
+	"predict_proba_discrete",
+]

blnetwork/inference/continuous.py

@@ -0,0 +1,21 @@
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+from ..training import utils as U  
+
+
+@torch.no_grad()
+def predict_continuous(
+    predictor: nn.Module,
+    x: torch.Tensor,
+    *,
+    device: str | torch.device | None = None,
+    return_cpu: bool = True,
+) -> torch.Tensor:
+
+    predictor.eval()
+    dev = U.resolve_device(model=predictor, tensor=x, device=device)
+    y_hat = predictor(x.to(dev))
+    return y_hat.detach().cpu() if return_cpu else y_hat.detach()

blnetwork/inference/discrete.py

@@ -0,0 +1,40 @@
+import torch
+import torch.nn.functional as F
+
+from ..training import utils as U  
+
+
+@torch.no_grad()
+def predict_proba_discrete(
+    model,
+    x: torch.Tensor,
+    *,
+    temperature: float = 1.0,
+    device: str | torch.device | None = None,
+    return_cpu: bool = False,
+) -> torch.Tensor:
+    
+    dev = U.resolve_device(model, x, device)
+    x = x.to(dev)
+
+    scores = model.logits(x) if hasattr(model, "logits") else model(x)
+    logits = scores / float(temperature)
+    probs = F.softmax(logits, dim=1)
+
+    return probs.cpu() if return_cpu else probs
+
+
+@torch.no_grad()
+def predict_class_discrete(
+    model,
+    x: torch.Tensor,
+    *,
+    temperature: float = 1.0,
+    device: str | torch.device | None = None,
+    return_cpu: bool = True,
+) -> torch.Tensor:
+    probs = predict_proba_discrete(
+        model, x, temperature=temperature, device=device, return_cpu=False
+    )
+    pred = probs.argmax(dim=1)
+    return pred.cpu() if return_cpu else pred

blnetwork/model/__init__.py

@@ -0,0 +1,7 @@
+from __future__ import annotations
+
+from .bldeep import BLDeep
+
+__all__ = [
+    "BLDeep"
+]

blnetwork/model/bldeep.py

@@ -0,0 +1,229 @@
+from typing import List, Sequence, Optional
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from . import utils as U
+
+
+class BLUnit(nn.Module):
+    def __init__(
+        self,
+        in_dim: int,
+        num_basis: int,
+        second_act_func: str = "relu",
+        third_act_func: str = "abs",
+        eps: float = 1e-8,
+        constrain_lambda: bool = True,
+        init_lambda: float = 1.0,
+        beta: float = 1.0,
+    ) -> None:
+        super().__init__()
+        self.in_dim = int(in_dim)
+        self.num_basis = int(num_basis)
+        self.second_act_func = str(second_act_func)
+        self.third_act_func = str(third_act_func)
+        self.eps = float(eps)
+        self.constrain_lambda = bool(constrain_lambda)
+        self.beta = float(beta)
+
+        init_lambda = float(init_lambda)
+
+        self.lin_u = nn.Linear(self.in_dim, self.num_basis, bias=True)
+        self.lin_c = nn.Linear(self.in_dim, self.num_basis, bias=True)
+        self.lin_t = nn.Linear(self.in_dim, self.num_basis, bias=True)
+
+        self.lam = nn.Parameter(torch.full((3, self.num_basis), init_lambda))
+
+    def forward(self, z: torch.Tensor) -> torch.Tensor:
+        u = torch.tanh(self.lin_u(z))
+        c = U.second_activation(self.lin_c(z), self.second_act_func, beta=self.beta)
+        t = U.third_activation(self.lin_t(z), self.third_act_func)
+
+        if self.constrain_lambda:
+            lam = F.softplus(self.lam) + self.eps
+        else:
+            lam = self.lam
+
+        lam_u = lam[0] 
+        lam_c = lam[1]  
+        lam_t = lam[2]  
+
+        return lam_u * u - lam_c * c - lam_t * t
+
+
+class BLBlock(nn.Module):
+    def __init__(
+        self,
+        in_dim: int,
+        num_basis: int,
+        second_act_func: str = "relu",
+        third_act_func: str = "abs",
+        constrain_lambda: bool = True,  
+        init_lambda: float = 1.0,
+        beta: float = 1.0,
+    ) -> None:
+        super().__init__()
+        self.unit = BLUnit(
+            in_dim=in_dim,
+            num_basis=num_basis,
+            second_act_func=second_act_func,
+            third_act_func=third_act_func,
+            constrain_lambda=constrain_lambda,  
+            init_lambda=init_lambda,
+            beta=beta,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.unit(x)
+
+
+class BLDeepBackbone(nn.Module):
+    def __init__(
+        self,
+        in_dim: int,
+        hidden_dims: Sequence[int],
+        second_act_func: str = "relu",
+        third_act_func: str = "abs",
+        constrain_lambda: bool = True,  
+        init_lambda: float = 1.0,
+        beta: float = 1.0,
+    ) -> None:
+        super().__init__()
+
+        self.in_dim = int(in_dim)
+        self.hidden_dims = list(hidden_dims)
+        dims: List[int] = [int(in_dim)] + list(hidden_dims[:-1])
+
+        self.blocks = nn.ModuleList(
+            BLBlock(
+                dims[i],
+                num_basis=hidden_dims[i],
+                second_act_func=second_act_func,
+                third_act_func=third_act_func,
+                constrain_lambda=constrain_lambda,  
+                init_lambda=init_lambda,
+                beta=beta,
+            )
+            for i in range(len(hidden_dims))
+        )
+        self.out_dim = int(hidden_dims[-1])
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        for blk in self.blocks:
+            x = blk(x)
+        return x
+
+
+class BLDeep(nn.Module):
+    def __init__(
+        self,
+        hidden_dims: Sequence[int],
+        second_act_func: str = "relu",
+        third_act_func: str = "abs",
+        head_bias: bool = True,
+        num_classes: Optional[int] = None,
+        task: str = "continuous",
+        constrain_lambda: bool = True, 
+        init_lambda: float = 1.0,
+        beta: float = 1.0,
+    ) -> None:
+        super().__init__()
+        self.hidden_dims = list(hidden_dims)
+        self.second_act_func = str(second_act_func)
+        self.third_act_func = third_act_func
+        self.head_bias = head_bias
+        self.num_classes = num_classes
+
+        if task not in {"continuous", "discrete"}:
+            raise ValueError(f"task must be 'continuous' or 'discrete', got '{task}'")
+        self.task = task
+        self.constrain_lambda = bool(constrain_lambda)  
+        self.init_lambda = float(init_lambda)
+        self.beta = float(beta)
+
+        self.x_dim: Optional[int] = None
+        self.y_dim: Optional[int] = None
+        self.backbone: Optional[nn.Module] = None
+        self.head: Optional[nn.Module] = None
+
+    def _build_architecture(self, x: torch.Tensor) -> None:
+        self.backbone = BLDeepBackbone(
+            in_dim=self.x_dim + self.y_dim,
+            hidden_dims=self.hidden_dims,
+            second_act_func=self.second_act_func,
+            third_act_func=self.third_act_func,
+            constrain_lambda=self.constrain_lambda,  
+            init_lambda=self.init_lambda,
+            beta=self.beta,
+        )
+        self.head = nn.Linear(self.backbone.out_dim, 1, bias=self.head_bias)
+
+        device, dtype = x.device, x.dtype
+        self.to(device=device, dtype=dtype)
+
+    def build(self, X: torch.Tensor, y: torch.Tensor) -> None:
+        self.x_dim = X.shape[1]
+
+        if self.task == "discrete":
+            m, y_idx = U.infer_num_classes(y)
+            self.num_classes = int(m)
+            self.y_dim = int(m)
+            if not torch.equal(y_idx, y.long()):
+                raise ValueError(
+                    f"Discrete labels must be in range [0..K-1]. "
+                    f"Got non-continuous labels. Please remap your labels to [0, 1, 2, ...] first."
+                )
+        else:
+            self.y_dim = 1 if y.ndim == 1 else y.shape[1]
+
+        self._build_architecture(X)
+    
+    def build_for_discrete_inference(self, x: torch.Tensor, num_classes: int) -> None:
+        if x.ndim != 2:
+            raise ValueError(f"x must be 2D (B, x_dim), got shape {tuple(x.shape)}")
+        if self.task != "discrete":
+            raise ValueError(f"This function only works with task='discrete', got task='{self.task}'")
+
+        self.x_dim = int(x.shape[1])
+        self.num_classes = int(num_classes)
+        self.y_dim = int(self.num_classes)
+
+        self._build_architecture(x)
+        
+    def score(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+        if self.backbone is None:
+            self.build(x, y)
+
+        if y.ndim == 1:
+            if self.task == "discrete":
+                y = F.one_hot(y.long(), num_classes=self.num_classes).to(device=x.device, dtype=x.dtype)
+            else:
+                y = y.unsqueeze(1)
+        elif y.ndim == 2:
+            if self.task == "discrete" and y.shape[1] == 1:
+                raise ValueError(
+                    f"For discrete task, y should be 1D class indices (shape (B,)), "
+                    f"but got shape {tuple(y.shape)}. Use y.squeeze(1) to convert (B,1) -> (B)."
+                )
+
+        z = torch.cat([x, y.to(device=x.device, dtype=x.dtype)], dim=1)
+        feats = self.backbone(z)
+        return self.head(feats)
+
+    def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+        return self.score(x, y)
+
+    def logits(self, x: torch.Tensor, num_classes: int | None = None) -> torch.Tensor:
+        if self.task != "discrete":
+            raise RuntimeError("logits() is only available when task='discrete'.")
+
+        m = int(num_classes or (self.num_classes or 0))
+        if m <= 0:
+            raise RuntimeError(
+                "Unknown num_classes for discrete inference. "
+                "Pass num_classes=K, or initialize BLDeep(..., num_classes=K), "
+            )
+        if self.backbone is None:
+            self.build_for_discrete_inference(x, num_classes=m)
+
+        return U.enumerate_onehot_logits(self.score, x, m=m)

blnetwork/model/utils.py

@@ -0,0 +1,68 @@
+import torch
+from typing import Callable, Tuple
+
+__all__ = [
+    "second_activation",
+    "third_activation",
+    "infer_num_classes",
+    "onehot_candidates",
+    "enumerate_onehot_logits",
+]
+
+def second_activation(z: torch.Tensor, second_act_func: str, beta: float = 1.0) -> torch.Tensor:
+    if second_act_func == "relu":
+        return torch.relu(z)
+    if second_act_func == "softplus":
+        return torch.nn.functional.softplus(z, beta=beta)
+    raise ValueError(f"Unknown second_act_func='{second_act_func}'. Use 'relu' or 'softplus'.")
+
+def third_activation(z: torch.Tensor, third_act_func: str) -> torch.Tensor:
+    if third_act_func == "abs":
+        return torch.abs(z)
+    if third_act_func == "square":
+        return z ** 2
+    raise ValueError(f"Unknown third_act_func='{third_act_func}'. Use 'abs' or 'square'.")
+
+def infer_num_classes(y: torch.Tensor) -> Tuple[int, torch.Tensor]:
+
+    if y.ndim == 2:
+        K = int(y.shape[1])
+        y_idx = torch.argmax(y, dim=1).long()
+        return K, y_idx
+
+    if y.ndim != 1:
+        raise ValueError(f"y must be 1D class-index or 2D one-hot, got shape={tuple(y.shape)}")
+
+    if y.dtype.is_floating_point:
+        if not torch.allclose(y, y.round()):
+            raise ValueError("Discrete y looks continuous (non-integer floats).")
+        y = y.round()
+
+    y = y.long()
+    classes, y_idx = torch.unique(y, sorted=True, return_inverse=True)
+    K = int(classes.numel())
+    return K, y_idx
+
+
+@torch.no_grad()
+def onehot_candidates(m: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
+    return torch.eye(int(m), device=device, dtype=dtype)
+
+
+def enumerate_onehot_logits(
+    score_fn: Callable[[torch.Tensor, torch.Tensor], torch.Tensor],
+    x: torch.Tensor,
+    m: int,
+) -> torch.Tensor:
+
+    B = x.shape[0]
+    m = int(m)
+    device, dtype = x.device, x.dtype
+
+    X_rep = x.repeat_interleave(m, dim=0)      
+    Y = onehot_candidates(m, device, dtype)    
+    Y_rep = Y.repeat(B, 1)                     
+
+    e = score_fn(X_rep, Y_rep)      
+    e = e.view(B, m)
+    return e

blnetwork/training/__init__.py

@@ -0,0 +1,17 @@
+from __future__ import annotations
+
+from .base import TrainConfig
+from .utils import OptimConfig, ExportConfig
+from .continuous import ContinuousTrainer
+from .discrete import DiscreteTrainer
+from .amortized import AmortizedConfig, fit_amortized_predictor
+
+__all__ = [
+    "OptimConfig",
+    "TrainConfig",
+    "ExportConfig",
+    "ContinuousTrainer",
+    "DiscreteTrainer",
+    "AmortizedConfig",
+    "fit_amortized_predictor"
+]