netemb 0.4.0__py3-none-any.whl

netemb/__init__.py

@@ -0,0 +1,15 @@
+from ._version import __version__
+from .estimator import NetworkEmbedder
+from .features import LaplacianSpectrum, FiedlerVector, SpectrumWithFiedler, FeatureExtractor
+from .generators import discover_generators, GeneratorProtocol
+
+__all__ = [
+    "__version__",
+    "NetworkEmbedder",
+    "LaplacianSpectrum",
+    "FiedlerVector",
+    "SpectrumWithFiedler",
+    "FeatureExtractor",
+    "discover_generators",
+    "GeneratorProtocol",
+]

netemb/_io.py

@@ -0,0 +1,68 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+import torch
+
+if TYPE_CHECKING:
+    from .estimator import NetworkEmbedder
+
+from platformdirs import user_data_dir
+
+from .model import EmbeddingNet
+from . import features as _features
+
+_DEFAULT_SAVE_PATH = Path(user_data_dir("netemb")) / "model.pt"
+
+
+def save(embedder: NetworkEmbedder, path: str | Path | None = None) -> Path:
+    """Persist model weights and hyperparameters to disk.
+
+    Parameters
+    ----------
+    embedder : NetworkEmbedder
+    path : str or Path, optional
+        Target file. Defaults to the platform user-data directory
+        (e.g. ~/Library/Application Support/netemb/model.pt on macOS).
+
+    Returns
+    -------
+    Path  — the path actually written to.
+    """
+    if embedder.model_ is None:
+        raise RuntimeError("Call fit() before save().")
+
+    target = Path(path) if path is not None else _DEFAULT_SAVE_PATH
+    target.parent.mkdir(parents=True, exist_ok=True)
+
+    torch.save({
+        'max_nodes':         embedder.max_nodes,
+        'emb_dim':           embedder.emb_dim,
+        'hidden_dims':       embedder.hidden_dims,
+        'activation':        embedder.activation,
+        'dropout':           embedder.dropout,
+        'feature_extractor': _features.serialize_feature_extractor(embedder.feature_extractor),
+        'input_dim':         embedder.model_.gates.numel(),
+        'state_dict':        embedder.model_.state_dict(),
+    }, target)
+    print(f"Model saved to {target}")
+    return target
+
+
+def load(path: str | Path | None = None) -> dict[str, Any]:
+    """Load a previously saved NetworkEmbedder from disk.
+
+    Parameters
+    ----------
+    path : str or Path, optional
+        File to read from. Defaults to the platform user-data directory.
+
+    Returns
+    -------
+    dict  — checkpoint dictionary (deserialized by the caller)
+    """
+    target = Path(path) if path is not None else _DEFAULT_SAVE_PATH
+    checkpoint = torch.load(target, weights_only=True)
+    print(f"Model loaded from {target}")
+    return checkpoint

netemb/_memcheck.py

@@ -0,0 +1,34 @@
+"""Lightweight RSS memory helpers for profiling netemb internals."""
+from __future__ import annotations
+import os
+import sys
+
+
+def mem_mb() -> float:
+    try:
+        import psutil
+    except ImportError:
+        raise ImportError(
+            "psutil is required for memory profiling. "
+            "Install with:  pip install psutil  or  pip install netemb[profile]"
+        ) from None
+    return psutil.Process(os.getpid()).memory_info().rss / 1024 / 1024
+
+
+def checkpoint(label: str) -> None:
+    print(f"[memcheck] {label}: {mem_mb():.1f} MB RSS", flush=True)
+
+
+def malloc_trim() -> None:
+    """Ask glibc to return freed heap pages to the OS (Linux only, no-op elsewhere).
+
+    On Linux, glibc's ptmalloc keeps freed memory in its internal pool instead
+    of returning it to the OS. Calling this after each generator run prevents
+    the RSS from growing monotonically across the dataset assembly loop.
+    """
+    if sys.platform.startswith("linux"):
+        try:
+            import ctypes
+            ctypes.CDLL("libc.so.6").malloc_trim(0)
+        except Exception:
+            pass

netemb/_train.py

@@ -0,0 +1,445 @@
+from __future__ import annotations
+
+import inspect
+import os
+import random
+import sys
+import time
+import warnings
+from collections.abc import Callable
+
+import numpy as np
+import torch
+
+from .features import FeatureExtractor
+from .loss import pairwise_distance_loss
+
+# Number of networks generated per parallel task. It is fixed (rather than
+# derived from n_jobs) so that, for a given random_state, the generated
+# data is identical whatever the number of parallel jobs.
+_CHUNK_SIZE = 10
+
+_THREAD_ENV_VARS = (
+    'OMP_NUM_THREADS', 'OPENBLAS_NUM_THREADS', 'MKL_NUM_THREADS',
+    'VECLIB_MAXIMUM_THREADS', 'NUMEXPR_NUM_THREADS',
+)
+
+
+def resolve_n_jobs(n_jobs: int, warn: bool = True) -> int:
+    """Validate n_jobs and turn it into an actual number of worker processes.
+
+    -1 means all CPUs this process may use; a positive int is used as is,
+    with a warning (if warn) when it exceeds the usable CPUs.
+    """
+    if isinstance(n_jobs, bool) or not isinstance(n_jobs, int) or (n_jobs < 1 and n_jobs != -1):
+        raise ValueError(f"n_jobs must be -1 or a positive int; got {n_jobs!r}")
+    if n_jobs == 1:
+        return 1
+    usable = usable_cpus()
+    if n_jobs == -1:
+        return usable
+    if warn and n_jobs > usable:
+        warnings.warn(
+            f"n_jobs={n_jobs}, but this process can only use {usable} CPU(s) "
+            f"(CPU affinity or cgroup/Slurm limits), so the {n_jobs} parallel jobs "
+            f"will share them. On Slurm, request the CPUs for a single task, "
+            f"e.g. --ntasks=1 --cpus-per-task={n_jobs}.",
+            RuntimeWarning, stacklevel=4,
+        )
+    return n_jobs
+
+
+def usable_cpus() -> int:
+    """Number of CPUs this process may actually use.
+
+    Accounts for CPU affinity and cgroup quotas (e.g. set by Slurm), which
+    can be far fewer than the cores of the machine.
+    """
+    from joblib import cpu_count
+    return cpu_count()
+
+
+def check_memory(feature_extractor: FeatureExtractor, max_nodes: int, n_jobs: int) -> None:
+    """Warn when n_jobs parallel eigendecompositions are unlikely to fit in RAM.
+
+    Only the dense diagonalisation of the built-in extractors is estimated
+    (Laplacian, eigenvectors and LAPACK workspace), sized for directed
+    graphs, whose magnetic Laplacian is complex (twice the memory of the
+    real Laplacian of undirected graphs); the memory of the NetworkX graph
+    itself comes on top of it.
+    """
+    from .features import LaplacianSpectrum, FiedlerVector, SpectrumWithFiedler
+    if n_jobs < 2 or not isinstance(feature_extractor, (LaplacianSpectrum, FiedlerVector, SpectrumWithFiedler)):
+        return
+    k = getattr(feature_extractor, 'k', None)
+    if k is not None and max_nodes > k + 1:
+        return  # sparse solver: memory does not scale with max_nodes**2
+    n_matrices = 2 if isinstance(feature_extractor, LaplacianSpectrum) else 4
+    per_job = n_matrices * 16 * max_nodes ** 2
+    available = _available_memory()
+    if available is None or n_jobs * per_job <= available:
+        return
+    suggested = max(1, int(available // per_job))
+    warnings.warn(
+        f"Creating networks with up to {max_nodes} nodes needs up to about "
+        f"{per_job / 1024**3:.1f} GB per parallel job (for directed networks), i.e. "
+        f"{n_jobs * per_job / 1024**3:.1f} GB for n_jobs={n_jobs}, but only "
+        f"{available / 1024**3:.1f} GB of memory is available. Consider "
+        f"n_jobs={suggested}.",
+        RuntimeWarning, stacklevel=3,
+    )
+
+
+def _available_memory() -> int | None:
+    try:
+        import psutil
+        return psutil.virtual_memory().available
+    except ImportError:
+        pass
+    try:
+        return os.sysconf('SC_PAGE_SIZE') * os.sysconf('SC_PHYS_PAGES')
+    except (ValueError, OSError, AttributeError):
+        return None
+
+
+def _threads_per_worker(n_workers: int) -> int:
+    """Maths-library threads each worker may use: its share of the usable CPUs.
+
+    Batch systems often set OMP_NUM_THREADS etc. to all the allocated CPUs;
+    inherited by every worker, that would run n_workers times too many
+    threads, which then mostly wait for each other. A lower value set by
+    the user is kept.
+    """
+    n_threads = max(1, usable_cpus() // n_workers)
+    for v in _THREAD_ENV_VARS:
+        try:
+            n_threads = min(n_threads, max(1, int(os.environ[v])))
+        except (KeyError, ValueError):
+            pass
+    return n_threads
+
+
+def _run_chunk(g: Callable, kwargs: dict, seed: int,
+               n_threads: int | None = None) -> tuple[np.ndarray, np.ndarray, float]:
+    """Generate one chunk of networks with a dedicated seed.
+
+    n_threads, when given (worker processes), caps the threads used by
+    torch and, if threadpoolctl is installed, by the BLAS/OpenMP libraries
+    already loaded in the worker, in addition to the environment variables
+    set when the worker starts.
+
+    Returns numpy arrays (cheaper and safer to send back from worker
+    processes than torch tensors) plus the compute time in seconds.
+    """
+    from contextlib import nullcontext
+    from ._memcheck import malloc_trim
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    limits = nullcontext()
+    if n_threads is not None:
+        torch.set_num_threads(n_threads)
+        try:
+            from threadpoolctl import threadpool_limits
+            limits = threadpool_limits(limits=n_threads)
+        except ImportError:
+            pass
+    t0 = time.perf_counter()
+    with limits:
+        X, y = g(**kwargs)
+    elapsed = time.perf_counter() - t0
+    malloc_trim()
+    return np.asarray(X), np.asarray(y), elapsed
+
+
+def _run_chunk_in_process(g: Callable, kwargs: dict, seed: int) -> tuple[np.ndarray, np.ndarray, float]:
+    """_run_chunk, restoring the caller's RNG states afterwards.
+
+    This keeps the main-process RNGs (used e.g. to initialise the network
+    weights) in the same state as when chunks run in worker processes.
+    """
+    states = random.getstate(), np.random.get_state(), torch.get_rng_state()
+    try:
+        return _run_chunk(g, kwargs, seed)
+    finally:
+        random.setstate(states[0])
+        np.random.set_state(states[1])
+        torch.set_rng_state(states[2])
+
+
+def _fmt_time(seconds: float) -> str:
+    seconds = int(round(seconds))
+    if seconds < 60:
+        return f"{seconds}s"
+    if seconds < 3600:
+        return f"{seconds // 60}m{seconds % 60:02d}s"
+    return f"{seconds // 3600}h{seconds % 3600 // 60:02d}m"
+
+
+class _Progress:
+    """Progress report for the creation of synthetic networks.
+
+    In a terminal a single line is updated in place; otherwise (log files,
+    some notebooks) a line is printed every 10% of the total.
+    """
+
+    def __init__(self, gen_totals: dict[tuple[str, str], int], enabled: bool):
+        self.enabled = enabled
+        self.gen_totals = gen_totals
+        self.gen_done = dict.fromkeys(gen_totals, 0)
+        self.gen_time = dict.fromkeys(gen_totals, 0.0)
+        self.total = sum(gen_totals.values())
+        self.done = 0
+        self.t0 = time.perf_counter()
+        self.tty = sys.stdout.isatty()
+        self.next_decile = 1
+
+    def update(self, key: tuple[str, str], n: int, compute_time: float) -> None:
+        self.done += n
+        self.gen_done[key] += n
+        self.gen_time[key] += compute_time
+        if not self.enabled:
+            return
+        if self.gen_done[key] == self.gen_totals[key]:
+            label, name = key
+            self._print(f"  [{label}] {name}: {self.gen_totals[key]} networks, "
+                        f"{_fmt_time(self.gen_time[key])} compute time", final=True)
+        if self.tty:
+            self._print(self._line(), final=False)
+        elif self.done * 10 >= self.next_decile * self.total:
+            while self.done * 10 >= self.next_decile * self.total:
+                self.next_decile += 1
+            self._print(self._line(), final=True)
+
+    def close(self) -> None:
+        if self.enabled and self.tty:
+            print(flush=True)
+
+    def _line(self) -> str:
+        elapsed = time.perf_counter() - self.t0
+        line = (f"Creating networks: {self.done}/{self.total} "
+                f"({100 * self.done / self.total:.0f}%) | elapsed {_fmt_time(elapsed)}")
+        if 0 < self.done < self.total:
+            line += f" | ~{_fmt_time(elapsed * (self.total - self.done) / self.done)} left"
+        return line
+
+    def _print(self, text: str, final: bool) -> None:
+        if self.tty:
+            print('\r\033[K' + text, end='\n' if final else '', flush=True)
+        else:
+            print(text, flush=True)
+
+
+def assemble_datasets(
+    generators: list[Callable],
+    sizes: list[int],
+    max_nodes: int,
+    feature_extractor: FeatureExtractor | None = None,
+    n_jobs: int = 1,
+    verbose: bool = False,
+    labels: list[str] | None = None,
+    profile_memory: bool = False,
+) -> list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]]:
+    """Generate one or more datasets (e.g. training and evaluation) from generators.
+
+    Each generator is called on chunks of _CHUNK_SIZE networks, each chunk
+    with its own seed drawn from numpy's global RNG, so that the result for
+    a given seed does not depend on n_jobs. All chunks of all datasets are
+    processed by the same pool of worker processes.
+
+    Parameters
+    ----------
+    generators : list of GeneratorProtocol
+    sizes : list of int
+        Number of graphs per generator, one entry per dataset.
+    max_nodes : int
+    feature_extractor : FeatureExtractor or None
+        Passed on to each generator that accepts a feature_extractor
+        keyword argument, so training data matches whatever extractor is
+        configured on the NetworkEmbedder. Generators that don't accept
+        it (e.g. custom extra_generators) fall back to their own default.
+    n_jobs : int, default 1
+        Number of worker processes; -1 uses all available cores and 1
+        generates networks one at a time in the current process.
+    verbose : bool, default False
+        Print the progress of the creation process.
+    labels : list of str or None
+        Names of the datasets used in progress messages.
+    profile_memory : bool
+
+    Returns
+    -------
+    list of (X, y, groups), one per dataset, with
+        X : torch.Tensor, shape (N, feature_len)
+        y : torch.Tensor, shape (N,)
+        groups : torch.Tensor, shape (N,), dtype long
+    """
+    n_jobs = resolve_n_jobs(n_jobs)
+    labels = labels if labels is not None else [f"set {i}" for i in range(len(sizes))]
+    if profile_memory:
+        from ._memcheck import checkpoint
+
+    # One task per (dataset, generator, chunk); seeds follow this canonical order.
+    tasks = []
+    for s, size in enumerate(sizes):
+        for gi, g in enumerate(generators):
+            kwargs = dict(max_nodes=max_nodes)
+            if feature_extractor is not None and 'feature_extractor' in inspect.signature(g).parameters:
+                kwargs['feature_extractor'] = feature_extractor
+            for c, start in enumerate(range(0, size, _CHUNK_SIZE)):
+                tasks.append((s, gi, c, dict(kwargs, n=min(_CHUNK_SIZE, size - start))))
+    seeds = np.random.SeedSequence(int(np.random.randint(0, 2**32))).spawn(len(tasks))
+    seeds = [int(ss.generate_state(1)[0]) for ss in seeds]
+    # Interleave generators, so that progress and time estimates are representative.
+    order = sorted(range(len(tasks)), key=lambda t: (tasks[t][0], tasks[t][2], tasks[t][1]))
+
+    n_workers = min(n_jobs, len(tasks))
+    progress = _Progress(
+        {(labels[s], g.__name__): size for s, size in enumerate(sizes) for g in generators},
+        enabled=verbose,
+    )
+    if verbose:
+        print(f"Creating {' + '.join(f'{len(generators) * n} {l}' for n, l in zip(sizes, labels))} "
+              f"networks with {n_workers} parallel job(s)"
+              + (f" on {usable_cpus()} usable CPU(s)" if n_workers > 1 else ""), flush=True)
+
+    results = {}
+
+    def collect(t, result):
+        s, gi, c, _ = tasks[t]
+        results[s, gi, c] = result[:2]
+        progress.update((labels[s], generators[gi].__name__), len(result[1]), result[2])
+        if profile_memory and progress.gen_done[labels[s], generators[gi].__name__] == sizes[s]:
+            checkpoint(f"  after generator {generators[gi].__name__} ({labels[s]})")
+
+    if n_workers <= 1:
+        for t in order:
+            s, gi, c, kwargs = tasks[t]
+            collect(t, _run_chunk_in_process(generators[gi], kwargs, seeds[t]))
+    else:
+        from concurrent.futures import as_completed
+        from joblib.externals.loky import get_reusable_executor
+        if profile_memory:
+            print("[memcheck] n_jobs > 1: eigendecomposition checkpoints run in worker "
+                  "processes and are not reported", flush=True)
+        n_threads = _threads_per_worker(n_workers)
+        env = {v: str(n_threads) for v in _THREAD_ENV_VARS}
+        executor = get_reusable_executor(max_workers=n_workers, env=env)
+        try:
+            futures = {executor.submit(_run_chunk, generators[tasks[t][1]], tasks[t][3], seeds[t], n_threads): t
+                       for t in order}
+            for future in as_completed(futures):
+                collect(futures[future], future.result())
+        finally:
+            # Release the workers' memory before training starts.
+            executor.shutdown(wait=True, kill_workers=True)
+    progress.close()
+
+    datasets = []
+    for s, size in enumerate(sizes):
+        n_chunks = len(range(0, size, _CHUNK_SIZE))
+        Xs, ys, gs = [], [], []
+        for gi in range(len(generators)):
+            for c in range(n_chunks):
+                X_c, y_c = results[s, gi, c]
+                Xs.append(X_c)
+                ys.append(y_c)
+                gs.append(np.full(len(y_c), gi, dtype=np.int64))
+        X = torch.from_numpy(np.concatenate(Xs, axis=0))
+        y = torch.from_numpy(np.concatenate(ys, axis=0))
+        groups = torch.from_numpy(np.concatenate(gs))
+        if profile_memory:
+            size_mb = X.element_size() * X.nelement() / 1024 / 1024
+            checkpoint(f"  assemble_dataset done ({labels[s]}) — X {tuple(X.shape)}, {size_mb:.1f} MB tensor")
+        datasets.append((X, y, groups))
+    return datasets
+
+
+def assemble_dataset(
+    generators: list[Callable],
+    n_per_group: int,
+    max_nodes: int,
+    feature_extractor: FeatureExtractor | None = None,
+    profile_memory: bool = False,
+    n_jobs: int = 1,
+    verbose: bool = False,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Generate a single training or evaluation dataset; see assemble_datasets()."""
+    return assemble_datasets(generators, [n_per_group], max_nodes,
+                             feature_extractor=feature_extractor, n_jobs=n_jobs,
+                             verbose=verbose, profile_memory=profile_memory)[0]
+
+
+def run_training_loop(
+    net: torch.nn.Module,
+    X: torch.Tensor,
+    y: torch.Tensor,
+    groups: torch.Tensor,
+    X_test: torch.Tensor,
+    y_test: torch.Tensor,
+    groups_test: torch.Tensor,
+    epochs: int,
+    lr: float,
+    verbose: bool,
+    profile_memory: bool = False,
+    cross_family_weight: float = 0.0,
+    cross_family_margin: float = 1.0,
+) -> tuple[list[float], list[float]]:
+    """Train EmbeddingNet in place and return (train_loss, eval_loss) lists.
+
+    Parameters
+    ----------
+    net : EmbeddingNet
+    X, y, groups : training tensors
+    X_test, y_test, groups_test : held-out evaluation tensors
+    epochs : int
+    lr : float
+    verbose : bool
+    profile_memory : bool
+    cross_family_weight : float, default 0.0
+        Strength of cross-family repulsion. 0.0 = original behaviour (no repulsion).
+    cross_family_margin : float, default 1.0
+        Minimum desired distance between different network families.
+
+    Returns
+    -------
+    train_loss : list of float
+    eval_loss : list of float
+    """
+    optimizer = torch.optim.Adam(net.parameters(), lr=lr)
+    train_loss = []
+    eval_loss  = []
+
+    if profile_memory:
+        from ._memcheck import checkpoint
+        checkpoint("training loop START")
+
+    for epoch in range(epochs):
+        net.train()
+        optimizer.zero_grad()
+        emb  = net(X)
+        loss = pairwise_distance_loss(
+            emb, y, groups,
+            cross_family_weight=cross_family_weight,
+            cross_family_margin=cross_family_margin,
+        )
+        loss.backward()
+        optimizer.step()
+
+        net.eval()
+        with torch.no_grad():
+            e_loss = pairwise_distance_loss(
+                net(X_test), y_test, groups_test,
+                cross_family_weight=cross_family_weight,
+                cross_family_margin=cross_family_margin,
+            )
+        train_loss.append(loss.item())
+        eval_loss.append(e_loss.item())
+
+        if epoch % 50 == 0 and verbose:
+            print(f"Epoch {epoch:4d} | Train loss: {loss.item():.4f} | Eval loss: {e_loss.item():.4f}")
+
+        if profile_memory and epoch % 100 == 0:
+            checkpoint(f"  epoch {epoch:4d}")
+
+    return train_loss, eval_loss

netemb/_version.py

@@ -0,0 +1 @@
+__version__ = "0.4.0"