sawnergy 1.0.5__py3-none-any.whl → 1.0.7__py3-none-any.whl

sawnergy/embedding/embedder.py

@@ -1,11 +1,11 @@
 from __future__ import annotations
 
 """
-Embedding orchestration for Skip-Gram with Negative Sampling (SGNS).
+Embedding orchestration for Skip-Gram (SG) and Skip-Gram with Negative Sampling (SGNS).
 
 This module consumes attractive/repulsive walk corpora produced by the walker
 pipeline and trains per-frame embeddings using either the PyTorch or PureML
-implementations of SGNS. The resulting embeddings can be persisted back into
+implementations of SG/SGNS. The resulting embeddings can be persisted back into
 an ``ArrayStorage`` archive along with rich metadata describing the training
 configuration.
 """
@@ -38,7 +38,8 @@ class Embedder:
                  WALKS_path: str | Path,
                  base: Literal["torch", "pureml"],
                  *,
-                 seed: int | None = None
+                 seed: int | None = None,
+                 objective: Literal["sgns", "sg"] = "sgns"
                 ) -> None:
         """Initialize the embedder and load walk tensors.
 
@@ -53,6 +54,8 @@ class Embedder:
             base: Which SGNS backend to use, either ``"torch"`` or ``"pureml"``.
             seed: Optional seed for the embedder's RNG. If ``None``, a random
                 32-bit seed is chosen.
+            objective: Training objective, either ``"sgns"`` (negative sampling)
+                or ``"sg"`` (plain full-softmax Skip-Gram).
 
         Raises:
             ValueError: If required metadata is missing or any loaded walk array
@@ -156,20 +159,25 @@ class Embedder:
 
         # MODEL HANDLE
         self.model_base: Literal["torch", "pureml"] = base
-        self.model_constructor = self._get_SGNS_constructor_from(base)
-        _logger.info("SGNS backend resolved: %s", getattr(self.model_constructor, "__name__", repr(self.model_constructor)))
+        self.objective: Literal["sgns", "sg"] = objective
+        self.model_constructor = self._get_SGNS_constructor_from(base, objective)
+        _logger.info(
+            "SG backend resolved: %s (objective=%s)",
+            getattr(self.model_constructor, "__name__", repr(self.model_constructor)), self.objective
+        )
 
     # -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=- PRIVATE -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=
 
     # HELPERS:
 
     @staticmethod
-    def _get_SGNS_constructor_from(base: Literal["torch", "pureml"]):
-        """Resolve the SGNS implementation class for the selected backend."""
+    def _get_SGNS_constructor_from(base: Literal["torch", "pureml"],
+                                   objective: Literal["sgns", "sg"]):
+        """Resolve the SG/SGNS implementation class for the selected backend."""
         if base == "torch":
             try:
-                from .SGNS_torch import SGNS_Torch
-                return SGNS_Torch
+                from .SGNS_torch import SGNS_Torch, SG_Torch
+                return SG_Torch if objective == "sg" else SGNS_Torch
             except Exception:
                 raise ImportError(
                     "PyTorch is not installed, but base='torch' was requested. "
@@ -178,8 +186,8 @@ class Embedder:
                 )
         elif base == "pureml":
             try:
-                from .SGNS_pml import SGNS_PureML
-                return SGNS_PureML
+                from .SGNS_pml import SGNS_PureML, SG_PureML
+                return SG_PureML if objective == "sg" else SGNS_PureML
             except Exception:
                 raise ImportError(
                     "PureML is not installed, but base='pureml' was requested. "
@@ -322,21 +330,24 @@ class Embedder:
     # -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-= PUBLIC -=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-=-= 
 
     def embed_frame(self,
-              frame_id: int,
-              RIN_type: Literal["attr", "repuls"],
-              using: Literal["RW", "SAW", "merged"],
-              window_size: int,
-              num_negative_samples: int,
-              num_epochs: int,
-              batch_size: int,
-              *,
-              shuffle_data: bool = True,
-              dimensionality: int = 128,
-              alpha: float = 0.75,
-              device: str | None = None,
-              sgns_kwargs: dict[str, object] | None = None,
-              _seed: int | None = None
-              ) -> np.ndarray:
+            frame_id: int,
+            RIN_type: Literal["attr", "repuls"],
+            using: Literal["RW", "SAW", "merged"],
+            window_size: int,
+            num_negative_samples: int,
+            num_epochs: int,
+            batch_size: int,
+            *,
+            lr_step_per_batch: bool = False,
+            shuffle_data: bool = True,
+            dimensionality: int = 128,
+            alpha: float = 0.75,
+            device: str | None = None,
+            sgns_kwargs: dict[str, object] | None = None,
+            kind: Literal["in", "out", "avg"] = "in",
+            _seed: int | None = None,
+            objective: Literal["sgns", "sg"] | None = None
+            ) -> np.ndarray:
         """Train embeddings for a single frame and return the input embedding matrix.
 
         Args:
@@ -347,6 +358,7 @@ class Embedder:
                 ``"merged"`` (concatenates both if available).
             window_size: Symmetric skip-gram window size ``k``.
             num_negative_samples: Number of negative samples per positive pair.
+                Ignored when ``objective="sg"``.
             num_epochs: Number of passes over the pair dataset.
             batch_size: Mini-batch size for training.
             shuffle_data: Whether to shuffle pairs each epoch.
@@ -355,11 +367,15 @@ class Embedder:
             device: Optional device string for the Torch backend (e.g., ``"cuda"``).
             sgns_kwargs: Extra keyword arguments forwarded to the backend SGNS
                 constructor. For PureML, required keys are:
-                ``{"optim", "optim_kwargs", "lr_sched", "lr_sched_kwargs"}``.
+                ``{"optim", "optim_kwargs"}``; ``lr_sched`` is optional, but if
+                provided then ``lr_sched_kwargs`` must also be provided.
+            kind: Which embedding matrix to return: ``"in"``, ``"out"``, or ``"avg"``.
             _seed: Optional child seed for this frame's model initialization.
+            objective: Training objective override for this call (``"sgns"`` or
+                ``"sg"``). If ``None``, uses the value set at construction.
 
         Returns:
-            np.ndarray: Learned **input** embedding matrix of shape ``(V, D)``.
+            np.ndarray: Learned embedding matrix (selected by ``kind``) of shape ``(V, D)``.
 
         Raises:
             ValueError: If requested walks are missing, if no training pairs are
@@ -391,10 +407,14 @@ class Embedder:
 
         model_kwargs: dict[str, object] = dict(sgns_kwargs or {})
         if self.model_base == "pureml":
-            required = {"optim", "optim_kwargs", "lr_sched", "lr_sched_kwargs"}
+            required = {"optim", "optim_kwargs"}
             missing = required.difference(model_kwargs)
             if missing:
                 raise ValueError(f"PureML backend requires {sorted(missing)} in sgns_kwargs.")
+            has_sched = ("lr_sched" in model_kwargs and model_kwargs["lr_sched"] is not None)
+            has_sched_kwargs = ("lr_sched_kwargs" in model_kwargs and model_kwargs["lr_sched_kwargs"] is not None)
+            if has_sched and not has_sched_kwargs:
+                raise ValueError("When providing lr_sched for PureML, you must also provide lr_sched_kwargs.")
 
         child_seed = int(self._seed if _seed is None else _seed)
         model_kwargs.update({
@@ -406,12 +426,16 @@ class Embedder:
         if self.model_base == "torch" and device is not None:
             model_kwargs["device"] = device
 
+        # Resolve objective (call-level override beats constructor default)
+        obj = self.objective if objective is None else objective
+        self.model_constructor = self._get_SGNS_constructor_from(self.model_base, obj)
         self.model = self.model_constructor(**model_kwargs)
 
         _logger.info(
-            "Training SGNS base=%s constructor=%s frame=%d pairs=%d dim=%d epochs=%d batch=%d neg=%d shuffle=%s",
+            "Training SG base=%s constructor=%s objective=%s frame=%d pairs=%d dim=%d epochs=%d batch=%d neg=%d shuffle=%s",
             self.model_base,
             getattr(self.model_constructor, "__name__", repr(self.model_constructor)),
+            obj,
             frame_id,
             pairs.shape[0],
             dimensionality,
@@ -421,24 +445,46 @@ class Embedder:
             shuffle_data
         )
 
-        self.model.fit(
-            centers,
-            contexts,
-            num_epochs,
-            batch_size,
-            num_negative_samples,
-            noise_probs,
-            shuffle_data,
-            lr_step_per_batch=False
-        )
+        if obj == "sgns":
+            self.model.fit(
+                centers,
+                contexts,
+                num_epochs,
+                batch_size,
+                num_negative_samples,
+                noise_probs,
+                shuffle_data,
+                lr_step_per_batch
+            )
+        else:
+            self.model.fit(
+                centers,
+                contexts,
+                num_epochs,
+                batch_size,
+                shuffle_data,
+                lr_step_per_batch
+            )
+
+        # Select embedding matrix by kind
+        if kind == "in":
+            embeddings = getattr(self.model, "in_embeddings", None)
+            if embeddings is None:
+                embeddings = getattr(self.model, "embeddings", None)
+                if embeddings is None:
+                    params = getattr(self.model, "parameters", None)
+                    if isinstance(params, tuple) and params:
+                        embeddings = params[0]
+        elif kind == "out":
+            embeddings = getattr(self.model, "out_embeddings", None)
+        else:  # "avg"
+            embeddings = getattr(self.model, "avg_embeddings", None)
 
-        embeddings = getattr(self.model, "embeddings", None)
-        if embeddings is None:
-            params = getattr(self.model, "parameters", None)
-            if isinstance(params, tuple) and params:
-                embeddings = params[0]
         if embeddings is None:
-            raise AttributeError("SGNS model does not expose embeddings via '.embeddings' or '.parameters[0]'")
+            raise AttributeError(
+                "SG/SGNS model does not expose the requested embeddings: "
+                f"kind={kind!r}"
+            )
 
         embeddings = np.asarray(embeddings)
         _logger.info("Frame %d embeddings ready: shape=%s dtype=%s", frame_id, embeddings.shape, embeddings.dtype)
@@ -460,7 +506,9 @@ class Embedder:
         sgns_kwargs: dict[str, object] | None = None,
         output_path: str | Path | None = None,
         num_matrices_in_compressed_blocks: int = 20,
-        compression_level: int = 3):
+        compression_level: int = 3,
+        objective: Literal["sgns", "sg"] | None = None,
+        kind: Literal["in", "out", "avg"] = "in"):
         """Train embeddings for all frames and persist them to compressed storage.
 
         Iterates through all frames (``1..frame_count``), trains an SGNS model
@@ -472,6 +520,7 @@ class Embedder:
             using: Walk collections: ``"RW"``, ``"SAW"``, or ``"merged"``.
             window_size: Symmetric skip-gram window size ``k``.
             num_negative_samples: Number of negative samples per positive pair.
+                Ignored when ``objective="sg"``.
             num_epochs: Number of epochs for each frame.
             batch_size: Mini-batch size used during training.
             shuffle_data: Whether to shuffle pairs each epoch.
@@ -487,6 +536,10 @@ class Embedder:
             num_matrices_in_compressed_blocks: Number of per-frame matrices to
                 store per compressed chunk in the output archive.
             compression_level: Blosc Zstd compression level (0-9).
+            objective: Training objective for all frames (``"sgns"`` or ``"sg"``).
+                If ``None``, uses the value set at construction.
+            kind: Which embedding matrix to persist for each frame: ``"in"``,
+                ``"out"``, or ``"avg"``.
 
         Returns:
             str: Filesystem path to the written embeddings archive (``.zip``).
@@ -498,8 +551,8 @@ class Embedder:
 
         Notes:
             - A deterministic child seed is spawned per frame from the master
-              seed using ``np.random.SeedSequence`` to ensure reproducibility
-              across runs.
+            seed using ``np.random.SeedSequence`` to ensure reproducibility
+            across runs.
         """
         current_time = sawnergy_util.current_time()
         if output_path is None:
@@ -535,7 +588,9 @@ class Embedder:
                     alpha=alpha,
                     device=device,
                     sgns_kwargs=sgns_kwargs,
-                    _seed=child_seed
+                    _seed=child_seed,
+                    objective=objective,
+                    kind=kind
                 )
             )
 
@@ -568,6 +623,7 @@ class Embedder:
             storage.add_attr("frame_embeddings_name", block_name)
             storage.add_attr("arrays_per_chunk", int(num_matrices_in_compressed_blocks))
             storage.add_attr("compression_level", int(compression_level))
+            storage.add_attr("objective", self.objective if objective is None else objective)
 
         _logger.info("Embedding archive written to %s", output_path)
         return str(output_path)

sawnergy/embedding/visualizer.py

@@ -0,0 +1,247 @@
+from __future__ import annotations
+
+# third party
+import numpy as np
+import matplotlib as mpl
+
+# built-in
+from pathlib import Path
+from typing import Sequence
+import logging
+
+# local
+from ..visual import visualizer_util
+from .. import sawnergy_util
+
+# *----------------------------------------------------*
+#                        GLOBALS
+# *----------------------------------------------------*
+
+_logger = logging.getLogger(__name__)
+
+# *----------------------------------------------------*
+#                        HELPERS
+# *----------------------------------------------------*
+
+def _safe_svd_pca(X: np.ndarray, k: int) -> tuple[np.ndarray, np.ndarray]:
+    """Compute k principal directions via SVD and project onto them."""
+    if X.ndim != 2:
+        raise ValueError(f"PCA expects 2D array (N, D); got {X.shape}")
+    _, D = X.shape
+    if k not in (2, 3):
+        raise ValueError(f"PCA dimensionality must be 2 or 3; got {k}")
+    if D < k:
+        raise ValueError(f"Requested k={k} exceeds feature dim D={D}")
+    Xc = X - X.mean(axis=0, keepdims=True)
+    _, _, Vt = np.linalg.svd(Xc, full_matrices=False)
+    comps = Vt[:k].copy()
+    proj = Xc @ comps.T
+    return proj, comps
+
+def _set_equal_axes_3d(ax, xyz: np.ndarray, *, padding: float = 0.05) -> None:
+    if xyz.size == 0:
+        return
+    x, y, z = xyz[:, 0], xyz[:, 1], xyz[:, 2]
+    xmin, xmax = float(x.min()), float(x.max())
+    ymin, ymax = float(y.min()), float(y.max())
+    zmin, zmax = float(z.min()), float(z.max())
+    xr = xmax - xmin
+    yr = ymax - ymin
+    zr = zmax - zmin
+    r = max(xr, yr, zr)
+    pad = padding * (r if r > 0 else 1.0)
+    cx, cy, cz = (xmin + xmax) / 2.0, (ymin + ymax) / 2.0, (zmin + zmax) / 2.0
+    ax.set_xlim(cx - r / 2 - pad, cx + r / 2 + pad)
+    ax.set_ylim(cy - r / 2 - pad, cy + r / 2 + pad)
+    ax.set_zlim(cz - r / 2 - pad, cz + r / 2 + pad)
+    try:
+        ax.set_box_aspect([1, 1, 1])
+    except Exception:
+        pass
+
+# *----------------------------------------------------*
+#                        CLASS
+# *----------------------------------------------------*
+
+class Visualizer:
+    """3D PCA visualizer for per-frame embeddings"""
+
+    no_instances: bool = True
+
+    def __init__(
+        self,
+        EMB_path: str | Path,
+        figsize: tuple[int, int] = (9, 7),
+        default_node_color: str = visualizer_util.GRAY,
+        depthshade: bool = False,
+        antialiased: bool = False,
+        init_elev: float = 35,
+        init_azim: float = 45,
+        *,
+        show: bool = False
+    ) -> None:
+        # Backend & pyplot
+        visualizer_util.ensure_backend(show)
+        import matplotlib.pyplot as plt
+        self._plt = plt
+
+        if Visualizer.no_instances:
+            try:
+                visualizer_util.warm_start_matplotlib()
+            finally:
+                Visualizer.no_instances = False
+
+        # Load embeddings archive
+        EMB_path = Path(EMB_path)
+        with sawnergy_util.ArrayStorage(EMB_path, mode="r") as storage:
+            name = storage.get_attr("frame_embeddings_name")
+            E = storage.read(name, slice(None))
+        if E.ndim != 3:
+            raise ValueError(f"Expected embeddings of shape (T,N,D); got {E.shape}")
+        self.E = np.asarray(E)
+        self.T, self.N, self.D = map(int, self.E.shape)
+        _logger.info("Loaded embeddings: T=%d, N=%d, D=%d", self.T, self.N, self.D)
+
+        # Coloring normalizer (parity with RIN Visualizer)
+        self._residue_norm = mpl.colors.Normalize(0, max(1, self.N - 1))
+
+        # Figure / axes / artists
+        self._fig = self._plt.figure(figsize=figsize)
+        self._ax = None
+        self._scatter = None
+        self._marker_size = 30.0
+        self._init_elev = init_elev
+        self._init_azim = init_azim
+        self.default_node_color = default_node_color
+        self._antialiased = bool(antialiased)
+        self._depthshade = bool(depthshade)
+
+    # ------------------------------ PRIVATE ------------------------------ #
+
+    def _ensure_axes(self) -> None:
+        if self._ax is not None and self._scatter is not None:
+            return
+        self._fig.clf()
+        self._ax = self._fig.add_subplot(111, projection="3d")
+        self._ax.view_init(self._init_elev, self._init_azim)
+        self._scatter = self._ax.scatter(
+            [], [], [],
+            s=self._marker_size,
+            depthshade=self._depthshade,
+            edgecolors="none",
+            antialiased=self._antialiased,
+        )
+        try:
+            self._ax.set_axis_off()
+        except Exception:
+            pass
+
+    def _project3(self, X: np.ndarray) -> np.ndarray:
+        """Return a 3D PCA projection of embeddings (always 3 coordinates).
+
+        If the embedding dimensionality D < 3, the remaining coordinate(s) are set to 0
+        so that the returned array still has shape (N, 3).
+        """
+        k = 3 if X.shape[1] >= 3 else 2
+        P, _ = _safe_svd_pca(X, k)
+        if k == 2:
+            P = np.c_[P, np.zeros((P.shape[0], 1), dtype=P.dtype)]
+        return P
+
+    def _select_nodes(self, displayed_nodes: Sequence[int] | str | None) -> np.ndarray:
+        if displayed_nodes is None or displayed_nodes == "ALL":
+            return np.arange(self.N, dtype=np.int64)
+        idx = np.asarray(displayed_nodes)
+        if idx.dtype.kind not in "iu":
+            raise TypeError("displayed_nodes must be None, 'ALL', or an integer sequence.")
+        if idx.min() < 1 or idx.max() > self.N:
+            raise IndexError(f"displayed_nodes out of range [1,{self.N}]")
+        return idx.astype(np.int64) - 1
+
+    def _apply_colors(self, node_colors, idx: np.ndarray) -> np.ndarray:
+        # RIN Visualizer semantics:
+        if isinstance(node_colors, str):
+            node_cmap = self._plt.get_cmap(node_colors)
+            return node_cmap(self._residue_norm(idx))
+        if node_colors is None:
+            full = visualizer_util.map_groups_to_colors(
+                N=self.N, groups=None, default_color=self.default_node_color, one_based=True
+            )
+            return np.asarray(full)[idx]
+        arr = np.asarray(node_colors)
+        if arr.ndim == 2 and arr.shape[0] == self.N and arr.shape[1] in (3, 4):
+            return arr[idx]
+        full = visualizer_util.map_groups_to_colors(
+            N=self.N, groups=node_colors, default_color=self.default_node_color, one_based=True
+        )
+        return np.asarray(full)[idx]
+
+    # ------------------------------ PUBLIC ------------------------------- #
+
+    def build_frame(
+        self,
+        frame_id: int,
+        *,
+        node_colors: str | np.ndarray | None = "rainbow",
+        displayed_nodes: Sequence[int] | str | None = "ALL",
+        show_node_labels: bool = False,
+        show: bool = False
+    ) -> None:
+        """Render a single frame as a PCA **3D** scatter (matches RIN Visualizer API)."""
+        frame0 = int(frame_id) - 1
+        if not (0 <= frame0 < self.T):
+            raise IndexError(f"frame_id out of range [1,{self.T}]")
+        self._ensure_axes()
+
+        idx = self._select_nodes(displayed_nodes)
+        X = self.E[frame0, idx, :]   # (n, D)
+        P = self._project3(X)        # (n, 3)
+        colors = self._apply_colors(node_colors, idx)
+
+        x, y, z = P[:, 0], P[:, 1], P[:, 2]
+        self._scatter._offsets3d = (x, y, z)
+        self._scatter.set_facecolors(colors)
+        _set_equal_axes_3d(self._ax, P, padding=0.05)
+        self._ax.view_init(self._init_elev, self._init_azim)
+
+        if show_node_labels:
+            for txt in getattr(self, "_labels", []):
+                try:
+                    txt.remove()
+                except Exception:
+                    pass
+            self._labels = []
+            for p, nid in zip(P, idx + 1):
+                self._labels.append(self._ax.text(p[0], p[1], p[2], str(int(nid)), fontsize=8))
+
+        # Be friendly to test dummies (they may lack tight_layout/canvas)
+        try:
+            self._fig.tight_layout()
+        except Exception:
+            try:
+                self._fig.subplots_adjust()
+            except Exception:
+                pass
+        try:
+            self._fig.canvas.draw_idle()
+        except Exception:
+            pass
+
+        if show:
+            try:
+                self._plt.show(block=True)
+            except TypeError:
+                self._plt.show()
+
+    # convenience
+    def savefig(self, path: str | Path, *, dpi: int = 150) -> None:
+        self._fig.savefig(path, dpi=dpi)
+
+    def close(self) -> None:
+        try:
+            self._plt.close(self._fig)
+        except Exception:
+            pass
+
+
+__all__ = ["Visualizer"]

sawnergy/logging_util.py

@@ -6,7 +6,7 @@ from datetime import datetime
 
 def configure_logging(
     logs_dir: Path | str,
-    file_level: int = logging.DEBUG,
+    file_level: int = logging.WARNING,
     console_level: int = logging.WARNING
 ) -> None:
     """

sawnergy/rin/rin_builder.py

@@ -669,7 +669,7 @@ class RINBuilder:
         molecule_of_interest: int,
         frame_range: tuple[int, int] | None = None,
         frame_batch_size: int = -1,
-        prune_low_energies_frac: float = 0.3,
+        prune_low_energies_frac: float = 0.85,
         output_path: str | Path | None = None,
         keep_prenormalized_energies: bool = True,
         *,

sawnergy/visual/visualizer.py

@@ -107,7 +107,7 @@ class Visualizer:
         visualizer_util.ensure_backend(show)
         import matplotlib.pyplot as plt
         self._plt = plt
-    # ---------- WARM UP MPL ------------ #
+        # ---------- WARM UP MPL ------------ #
         _logger.debug("Visualizer.__init__ start | RIN_path=%s, figsize=%s, node_size=%s, edge_width=%s, depthshade=%s, antialiased=%s, init_view=(%s,%s)",
                       RIN_path, figsize, node_size, edge_width, depthshade, antialiased, init_elev, init_azim)
         if Visualizer.no_instances:
@@ -116,7 +116,7 @@ class Visualizer:
         else:
             _logger.debug("Skipping warm-start (no_instances=False).")
     
-    # ---------- LOAD THE DATA ---------- #
+        # ---------- LOAD THE DATA ---------- #
         with sawnergy_util.ArrayStorage(RIN_path, mode="r") as storage:
             com_name = storage.get_attr("com_name")
             attr_energies_name = storage.get_attr("attractive_energies_name")
@@ -135,7 +135,7 @@ class Visualizer:
         self.N = np.size(self.COM_coords[0], axis=0)
         _logger.debug("Computed N=%d", self.N)
 
-    # - SET UP THE CANVAS AND THE AXES - #
+        # - SET UP THE CANVAS AND THE AXES - #
         self._fig = plt.figure(figsize=figsize, num="SAWNERGY")
         self._ax  = self._fig.add_subplot(111, projection="3d")
         self._fig.subplots_adjust(left=0, right=1, bottom=0, top=1)
@@ -145,14 +145,14 @@ class Visualizer:
         self._ax.set_axis_off()
         _logger.debug("Figure and 3D axes initialized.")
 
-    # ------ SET UP PLOT ELEMENTS ------ #
+        # ------ SET UP PLOT ELEMENTS ------ #
         self._scatter: PathCollection  = self._ax.scatter([], [], [], s=node_size, depthshade=depthshade, edgecolors="none")
         self._attr: Line3DCollection   = Line3DCollection(np.empty((0,2,3)), linewidths=edge_width, antialiased=antialiased)
         self._repuls: Line3DCollection = Line3DCollection(np.empty((0,2,3)), linewidths=edge_width, antialiased=antialiased)
         self._ax.add_collection3d(self._attr); self._ax.add_collection3d(self._repuls) # set pointers to the attractive and repulsive collections
         _logger.debug("Artists created | scatter(empty), attr_lines(empty), repuls_lines(empty).")
 
-    # ---------- HELPER FIELDS --------- #
+        # ---------- HELPER FIELDS --------- #
         # NOTE: 'under the hood' everything is 0-base indexed,
         # BUT, from the API point of view, the indexing is 1-base,
         # because amino acid residues are 1-base indexed.
@@ -160,7 +160,7 @@ class Visualizer:
         self.default_node_color = default_node_color
         _logger.debug("Helper fields set | residue_norm=[0,%d], default_node_color=%s", self.N-1, self.default_node_color)
 
-    # DISALLOW MPL WARM-UP IN THE FUTURE
+        # DISALLOW MPL WARM-UP IN THE FUTURE
         Visualizer.no_instances = False
         _logger.debug("Visualizer.no_instances set to False.")
 

sawnergy/visual/visualizer_util.py

@@ -319,6 +319,9 @@ def build_line_segments(
     kept = edge_weights >= thresh
     rows, cols = rows[kept], cols[kept]
 
+    nz = weights[rows, cols] > 0.0
+    rows, cols = rows[nz], cols[nz]
+
     if rows.size == 0:
         _logger.debug("build_line_segments: no edges kept after threshold; returning empties.")
         return (np.empty((0, 2, 3), dtype=float),