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

sawnergy/embedding/visualizer.py

@@ -23,7 +23,7 @@ _logger = logging.getLogger(__name__)
 #                        HELPERS
 # *----------------------------------------------------*
 
-def _safe_svd_pca(X: np.ndarray, k: int) -> tuple[np.ndarray, np.ndarray]:
+def _safe_svd_pca(X: np.ndarray, k: int, *, row_l2: bool = False) -> 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}")
@@ -33,6 +33,9 @@ def _safe_svd_pca(X: np.ndarray, k: int) -> tuple[np.ndarray, np.ndarray]:
     if D < k:
         raise ValueError(f"Requested k={k} exceeds feature dim D={D}")
     Xc = X - X.mean(axis=0, keepdims=True)
+    if row_l2:
+        norms = np.linalg.norm(Xc, axis=1, keepdims=True)
+        Xc = Xc / np.clip(norms, 1e-9, None)
     _, _, Vt = np.linalg.svd(Xc, full_matrices=False)
     comps = Vt[:k].copy()
     proj = Xc @ comps.T
@@ -78,7 +81,8 @@ class Visualizer:
         init_elev: float = 35,
         init_azim: float = 45,
         *,
-        show: bool = False
+        show: bool = False,
+        normalize_rows: bool = False,
     ) -> None:
         # Backend & pyplot
         visualizer_util.ensure_backend(show)
@@ -106,7 +110,7 @@ class Visualizer:
         self._residue_norm = mpl.colors.Normalize(0, max(1, self.N - 1))
 
         # Figure / axes / artists
-        self._fig = self._plt.figure(figsize=figsize)
+        self._fig = self._plt.figure(figsize=figsize, num="SAWNERGY")
         self._ax = None
         self._scatter = None
         self._marker_size = 30.0
@@ -115,6 +119,7 @@ class Visualizer:
         self.default_node_color = default_node_color
         self._antialiased = bool(antialiased)
         self._depthshade = bool(depthshade)
+        self._normalize_rows = bool(normalize_rows)
 
     # ------------------------------ PRIVATE ------------------------------ #
 
@@ -143,7 +148,7 @@ class Visualizer:
         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)
+        P, _ = _safe_svd_pca(X, k, row_l2=self._normalize_rows)
         if k == 2:
             P = np.c_[P, np.zeros((P.shape[0], 1), dtype=P.dtype)]
         return P
@@ -214,7 +219,6 @@ class Visualizer:
             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:

{sawnergy-1.0.7.dist-info → sawnergy-1.0.8.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sawnergy
-Version: 1.0.7
+Version: 1.0.8
 Summary: Toolkit for transforming molecular dynamics (MD) trajectories into rich graph representations
 Home-page: https://github.com/Yehor-Mishchyriak/SAWNERGY
 Author: Yehor Mishchyriak
@@ -52,19 +52,31 @@ keeps the full workflow — from `cpptraj` output to skip-gram embeddings (node2
 
 > **Optional:** For GPU training, install PyTorch separately (e.g., `pip install torch`).
 > **Note:** RIN building requires `cpptraj` (AmberTools). Ensure it is discoverable via `$PATH` or the `CPPTRAJ`
-> environment variable. Probably the easiest solution: install AmberTools via conda, activate the environment, and SAWNERGY will find cpptraj executable on its own, so just run your code and don't worry about it.
+> environment variable. Probably the easiest solution: install AmberTools via Conda, activate the environment, and SAWNERGY will find the cpptraj executable on its own, so just run your code and don't worry about it.
 
 ---
 
 # UPDATES:
 
+## v1.0.8 — What’s new:
+- **Temporary deprecation of `SGNS_Torch`**
+  - `sawnergy.embedding.SGNS_Torch` currently produces noisy embeddings in practice. The issue likely stems from **weight initialization**, although the root cause has not yet been conclusively determined.
+  - **Action:** The class and its `__init__` docstring now carry a deprecation notice. Constructing the class emits a **`DeprecationWarning`** and logs a **warning**.
+  - **Use instead:** Prefer **`SG_Torch`** (plain Skip-Gram with full softmax) or the PureML backends **`SGNS_PureML`** / **`SG_PureML`**.
+  - **Compatibility:** No breaking API changes; imports remain stable. PureML backends are unaffected.
+- **Embedding visualizer update**
+  - Now you can L2 normalize your embeddings before display.
+- **Small improvements in the embedding module**
+  - Improved API with a lot of good defaults in place to ease usage out of the box.
+  - Small internal model tweaks.
+
 ## v1.0.7 — What’s new:
-- **Added plain SkipGram model**
+- **Added plain Skip-Gram model**
   - Now, the user can choose if they want to apply the negative sampling technique (two binary classifiers) or train a single classifier over the vocabulary (full softmax). For more detail, see: [node2vec](https://arxiv.org/pdf/1607.00653), [word2vec](https://arxiv.org/pdf/1301.3781), and [negative_sampling](https://arxiv.org/pdf/1402.3722).
 - **Set a harsher default for low interaction energies pruning during RIN construction**
   - Now we zero out 85% of the lowest interaction energies as opposed to the past 30% default, leading to more meaningful embeddings.
 - **BUG FIX: Visualizer**
-  - Previously, the visualizer would silently draw edges of 0 magnitude, meaning they were actually being drawn but were invisible due to full transparency and 0 width. As a result, the displayed image / animation would be very laggy. Now, this was fixed, and given high pruning default, the displayed interaction networks are clean and smooth under rotations, dragging, etc.
+  - Previously, the visualizer would silently draw edges of 0 magnitude, meaning they were actually being drawn but were invisible due to full transparency and 0 width. As a result, the displayed image/animation would be very laggy. Now, this was fixed, and given the higher pruning default, the displayed interaction networks are clean and smooth under rotations, dragging, etc.
 - **New Embedding Visualizer (3D)**
   - New lightweight viewer for per-frame embeddings that projects embeddings with PCA to a **3D** scatter. Supports the same node coloring semantics, optional node labels, and the same antialiasing/depthshade controls. Works in headless setups using the same backend guard and uses a blocking `show=True` for scripts.
 
@@ -77,7 +89,7 @@ keeps the full workflow — from `cpptraj` output to skip-gram embeddings (node2
 - **Deterministic, shareable artifacts**: Every stage produces compressed Zarr archives that contain both data and metadata so runs can be reproduced, shared, or inspected later.
 - **High-performance data handling**: Heavy arrays live in shared memory during walk sampling to allow parallel processing without serialization overhead; archives are written in chunked, compressed form for fast read/write.
 - **Flexible objectives & backends**: Train Skip-Gram with **negative sampling** (`objective="sgns"`) or **plain Skip-Gram** (`objective="sg"`), using either **PureML** (default) or **PyTorch**.
-- **Visualization out of the box**: Plot and animate residue networks without leaving Python, using the data produced by RINBuilder
+- **Visualization out of the box**: Plot and animate residue networks without leaving Python, using the data produced by RINBuilder.
 
 ---
 
@@ -117,7 +129,7 @@ node indexing, and RNG seeds stay consistent across the toolchain.
 * Wraps the AmberTools `cpptraj` executable to:
   - compute per-frame electrostatic (EMAP) and van der Waals (VMAP) energy matrices at the atomic level,
   - project atom–atom interactions to residue–residue interactions using compositional masks,
-  - prune, symmetrize, remove self-interactions, and L1-normalise the matrices,
+  - prune, symmetrize, remove self-interactions, and L1-normalize the matrices,
   - compute per-residue centers of mass (COM) over the same frames.
 * Outputs a compressed Zarr archive with transition matrices, optional pre-normalized energies, COM snapshots, and rich
   metadata (frame range, pruning quantile, molecule ID, etc.).
@@ -142,13 +154,10 @@ node indexing, and RNG seeds stay consistent across the toolchain.
 
 ### `sawnergy.embedding.Embedder`
 
-* Consumes walk archives, generates skip-gram pairs, and normalises them to 0-based indices.
-* Provides a unified interface to SGNS implementations:
-  - **PureML backend** (`SGNS_PureML`): works with the `pureml` ecosystem, optimistic for CPU training.
-  - **PyTorch backend** (`SGNS_Torch`): uses `torch.nn.Embedding` plays nicely with GPUs.
-* Both `SGNS_PureML` and `SGNS_Torch` accept training hyperparameters such as batch_size, LR, optimizer and LR_scheduler, etc.
-* Exposes `embed_frame` (single frame) and `embed_all` (all frames, deterministic seeding per frame) which return the
-  learned input embedding matrices and write them to disk when requested.
+* Consumes walk archives, generates skip-gram pairs, and normalizes them to 0-based indices.
+* Selects skip-gram (SG / SGNS) backends dynamically via `model_base="pureml"|"torch"` with per-backend overrides supplied through `model_kwargs`.
+* Handles deterministic per-frame seeding and returns the requested embedding `kind` (`"in"`, `"out"`, or `"avg"`) from `embed_frame` and `embed_all`.
+* Persists per-frame matrices with rich provenance (walk metadata, objective, hyperparameters, RNG seeds) when `embed_all` targets an output archive.
 
 ### Supporting Utilities
 
@@ -166,11 +175,11 @@ node indexing, and RNG seeds stay consistent across the toolchain.
 |---|---|---|
 | **RIN** | `ATTRACTIVE_transitions` → **(T, N, N)**, float32  •  `REPULSIVE_transitions` → **(T, N, N)**, float32 (optional)  •  `ATTRACTIVE_energies` → **(T, N, N)**, float32 (optional)  •  `REPULSIVE_energies` → **(T, N, N)**, float32 (optional)  •  `COM` → **(T, N, 3)**, float32 | `time_created` (ISO) • `com_name` = `"COM"` • `molecule_of_interest` (int) • `frame_range` = `(start, end)` inclusive • `frame_batch_size` (int) • `prune_low_energies_frac` (float in [0,1]) • `attractive_transitions_name` / `repulsive_transitions_name` (dataset names or `None`) • `attractive_energies_name` / `repulsive_energies_name` (dataset names or `None`) |
 | **Walks** | `ATTRACTIVE_RWs` → **(T, N·num_RWs, L+1)**, int32 (optional)  •  `REPULSIVE_RWs` → **(T, N·num_RWs, L+1)**, int32 (optional)  •  `ATTRACTIVE_SAWs` → **(T, N·num_SAWs, L+1)**, int32 (optional)  •  `REPULSIVE_SAWs` → **(T, N·num_SAWs, L+1)**, int32 (optional)  <br/>_Note:_ node IDs are **1-based**.| `time_created` (ISO) • `seed` (int) • `rng_scheme` = `"SeedSequence.spawn_per_batch_v1"` • `num_workers` (int) • `in_parallel` (bool) • `batch_size_nodes` (int) • `num_RWs` / `num_SAWs` (ints) • `node_count` (N) • `time_stamp_count` (T) • `walk_length` (L) • `walks_per_node` (int) • `attractive_RWs_name` / `repulsive_RWs_name` / `attractive_SAWs_name` / `repulsive_SAWs_name` (dataset names or `None`) • `walks_layout` = `"time_leading_3d"` |
-| **Embeddings** | `FRAME_EMBEDDINGS` → **(frames_written, vocab_size, D)**, typically float32 | `time_created` (ISO) • `seed` (int) • `rng_scheme` = `"SeedSequence.spawn_per_frame_v1"` • `source_walks_path` (str) • `model_base` = `"torch"` or `"pureml"` • `rin_type` = `"attr"` or `"repuls"` • `using_mode` = `"RW"|"SAW"|"merged"` • `window_size` (int) • `alpha` (float; noise exponent) • `dimensionality` = D • `num_negative_samples` (int) • `num_epochs` (int) • `batch_size` (int) • `shuffle_data` (bool) • `frames_written` (int) • `vocab_size` (int) • `frame_count` (int) • `embedding_dtype` (str) • `frame_embeddings_name` = `"FRAME_EMBEDDINGS"` • `arrays_per_chunk` (int) • `compression_level` (int) • `objective` = `"sgns"` or `"sg"` |
+| **Embeddings** | `FRAME_EMBEDDINGS` → **(T, N, D)**, float32 | `created_at` (ISO) • `frame_embeddings_name` = `"FRAME_EMBEDDINGS"` • `time_stamp_count` = T • `node_count` = N • `embedding_dim` = D • `model_base` = `"torch"` or `"pureml"` • `embedding_kind` = `"in"|"out"|"avg"` • `objective` = `"sgns"` or `"sg"` • `negative_sampling` (bool) • `num_negative_samples` (int) • `num_epochs` (int) • `batch_size` (int) • `window_size` (int) • `alpha` (float) • `lr_step_per_batch` (bool) • `shuffle_data` (bool) • `device_hint` (str) • `model_kwargs_repr` (repr string) • `RIN_type` = `"attr"` or `"repuls"` • `using` = `"RW"|"SAW"|"merged"` • `source_WALKS_path` (str) • `walk_length` (int) • `num_RWs` / `num_SAWs` (ints) • `attractive_*_name` / `repulsive_*_name` (dataset names or `None`) • `master_seed` (int) • `per_frame_seeds` (list[int]) • `arrays_per_chunk` (int) • `compression_level` (int) |
 
 **Notes**
 
-- In **RIN**, `T` equals the number of frame **batches** written (i.e., `frame_range` swept in steps of `frame_batch_size`). `ATTRACTIVE/REPULSIVE_energies` are **pre-normalised** absolute energies (written only when `keep_prenormalized_energies=True`), whereas `ATTRACTIVE/REPULSIVE_transitions` are the **row-wise L1-normalised** versions used for sampling.
+- In **RIN**, `T` equals the number of frame **batches** written (i.e., `frame_range` swept in steps of `frame_batch_size`). `ATTRACTIVE/REPULSIVE_energies` are **pre-normalized** absolute energies (written only when `keep_prenormalized_energies=True`), whereas `ATTRACTIVE/REPULSIVE_transitions` are the **row-wise L1-normalized** versions used for sampling.
 - All archives are Zarr v3 groups. ArrayStorage also maintains per-block metadata in root attrs: `array_chunk_size_in_block`, `array_shape_in_block`, and `array_dtype_in_block` (dicts keyed by dataset name). You’ll see these in every archive.
 - In **Embeddings**, `alpha` and `num_negative_samples` apply to **SGNS** only and are ignored for `objective="sg"`.
 
@@ -200,7 +209,7 @@ rin_builder.build_rin(
     prune_low_energies_frac=0.85,
     output_path=rin_path,
     include_attractive=True,
-    include_repulsive=False,
+    include_repulsive=False
 )
 
 # 2. Sample walks from the RIN
@@ -208,44 +217,34 @@ walker = Walker(rin_path, seed=123)
 walks_path = Path("./WALKS_demo.zip")
 walker.sample_walks(
     walk_length=16,
-    walks_per_node=32,
+    walks_per_node=100,
     saw_frac=0.25,
     include_attractive=True,
     include_repulsive=False,
     time_aware=False,
     output_path=walks_path,
-    in_parallel=False,
+    in_parallel=False
 )
 walker.close()
 
 # 3. Train embeddings per frame (PyTorch backend)
 import torch
 
-embedder = Embedder(walks_path, base="torch", seed=999)
+embedder = Embedder(walks_path, seed=999)
 embeddings_path = embedder.embed_all(
     RIN_type="attr",
     using="merged",
+    num_epochs=10,
+    negative_sampling=False,
     window_size=4,
-    objective="sgns",
-    num_negative_samples=5,
-    num_epochs=5,
-    batch_size=1024,
-    dimensionality=128,
-    shuffle_data=True,
-    output_path="./EMBEDDINGS_demo.zip",
-    sgns_kwargs={
-        "optim": torch.optim.Adam,
-        "optim_kwargs": {"lr": 1e-3},
-        "lr_sched": torch.optim.lr_scheduler.LambdaLR,
-        "lr_sched_kwargs": {"lr_lambda": lambda _: 1.0},
-        "device": "cuda" if torch.cuda.is_available() else "cpu",
-    },
+    device="cuda" if torch.cuda.is_available() else "cpu",
+    model_base="torch",
+    output_path="./EMBEDDINGS_demo.zip"
 )
 print("Embeddings written to", embeddings_path)
 ```
 
-> For the PureML backend, supply the relevant optimiser and scheduler via `sgns_kwargs`
-> (for example `optim=pureml.optimizers.Adam`, `lr_sched=pureml.optimizers.CosineAnnealingLR`).
+> For the PureML backend, set `model_base="pureml"` and pass the optimizer / scheduler classes inside `model_kwargs`.
 
 ---
 
@@ -270,7 +269,7 @@ v.build_frame(1,
 ```python
 from sawnergy.embedding import Visualizer
 
-viz = sawnergy.embedding.Visualizer("./EMBEDDINGS_demo.zip")
+viz = Visualizer("./EMBEDDINGS_demo.zip", normalize_rows=True)
 viz.build_frame(1, show=True)
 ```
 
@@ -280,8 +279,7 @@ viz.build_frame(1, show=True)
 
 - **Time-aware walks**: Set `time_aware=True`, provide `stickiness` and `on_no_options` when calling `Walker.sample_walks`.
 - **Shared memory lifecycle**: Call `Walker.close()` (or use a context manager) to release shared-memory segments.
-- **PureML vs PyTorch**: Choose the backend via `Embedder(..., base="pureml"|"torch")` and provide backend-specific
-  constructor kwargs through `sgns_kwargs` (optimizer, scheduler, device).
+- **PureML vs PyTorch**: Select the backend at call time with `model_base="pureml"|"torch"` (defaults to `"pureml"`) and pass optimizer / scheduler overrides through `model_kwargs`.
 - **ArrayStorage utilities**: Use `ArrayStorage` directly to peek into archives, append arrays, or manage metadata.
 
 ---
@@ -292,8 +290,9 @@ viz.build_frame(1, show=True)
 ├── sawnergy/
 │   ├── rin/           # RINBuilder and cpptraj integration helpers
 │   ├── walks/         # Walker class and shared-memory utilities
-│   ├── embedding/     # Embedder + SGNS backends (PureML / PyTorch)
+│   ├── embedding/     # Embedder + SG/SGNS backends (PureML / PyTorch)
 │   ├── visual/        # Visualizer and palette utilities
+│   │
 │   ├── logging_util.py
 │   └── sawnergy_util.py
 │
@@ -302,7 +301,7 @@ viz.build_frame(1, show=True)
 
 ---
 
-## Acknowledgements
+## Acknowledgments
 
 SAWNERGY builds on the AmberTools `cpptraj` ecosystem, NumPy, Matplotlib, Zarr, and PyTorch (for GPU acceleration if necessary; PureML is available by default).
 Big thanks to the upstream communities whose work makes this toolkit possible.

{sawnergy-1.0.7.dist-info → sawnergy-1.0.8.dist-info}/RECORD

@@ -1,11 +1,11 @@
 sawnergy/__init__.py,sha256=Dq1U38ah6nPRFEDKN41mYphcTynKfnItca6QkYkpSbs,248
 sawnergy/logging_util.py,sha256=mfYw8IsYtOfCXayjkd4g9jHuupluxRNbqyFegRkiAhQ,1476
 sawnergy/sawnergy_util.py,sha256=Htx9wr0S8TXt5aHT2mtEdYf1TCo_BC1IUwNNuZdIR-4,49432
-sawnergy/embedding/SGNS_pml.py,sha256=LfZDlIF3-KnWUAjhwOT5ggGl2OoReM8_L0TCVYs6GJ0,14299
-sawnergy/embedding/SGNS_torch.py,sha256=NIr-RlOmXlEPe3m8Z6XvuG0b8MGTidETfugigcQTwFs,11232
+sawnergy/embedding/SGNS_pml.py,sha256=-S7K7qwbDGUO_KW4gnA3dGyxuezN1ZK-WikPm7krEvs,14291
+sawnergy/embedding/SGNS_torch.py,sha256=NgVQnMtRSYY0IsPhB3XV7K1-uVSah0P77a8ID8zZ7Qw,13940
 sawnergy/embedding/__init__.py,sha256=T1YXb7S5Zyy_kIqlarDSX3imd_FGFH6nDuvLQ3hMKsE,1764
-sawnergy/embedding/embedder.py,sha256=K9I6HYYQFH7SHpgxeTCf8_MMvyLxVaAaltoMwJbgyqo,28749
-sawnergy/embedding/visualizer.py,sha256=bweituYNj5dOzFhvU4n_E-RbzZiUKw6bfJchFLfjFD4,8625
+sawnergy/embedding/embedder.py,sha256=02pcf3ies3Nuo19sCoJdMAYg7BFUHj4-wf4AZ5R6PAE,32492
+sawnergy/embedding/visualizer.py,sha256=x0BiSG9_nk9AUQm9RsZ2syKeCiaxX1gTlC85aYycMXY,8830
 sawnergy/rin/__init__.py,sha256=z19hLfEIp3bwzY-eCHQBQf0NRTCJzVz_FLIpVV5q0W4,162
 sawnergy/rin/rin_builder.py,sha256=d1cC4KKY9zzNlqhxHWTFM-QyXRXubd2zlCrSM-dV5pc,44624
 sawnergy/rin/rin_util.py,sha256=5TKywA5qfm76Gl4Cyz7oBPasmE5chclR7UM4hawwQOg,14939
@@ -15,9 +15,9 @@ sawnergy/visual/visualizer_util.py,sha256=7y3kWjHxDQMoG0dmimceHKTC5veVChoyvW7d0q
 sawnergy/walks/__init__.py,sha256=Z_Kaffhn3oUX13z9jbY0V5Ncdwj9Cnr--n9D-s7gh5k,250
 sawnergy/walks/walker.py,sha256=scvfZFrSL4AwpmspD0Jb0uhnrVIRRwE_hPCE3bG6zpg,37729
 sawnergy/walks/walker_util.py,sha256=ETdyPNIDwDQCA8Z5t38keBhYBJ56_ksT_0NhOCY-tHE,15361
-sawnergy-1.0.7.dist-info/licenses/LICENSE,sha256=cElK4bCsDhyAEON3H05s35bQZvxBcXBiCOrOdiUhDCY,11346
-sawnergy-1.0.7.dist-info/licenses/NOTICE,sha256=eVTbuSasZrmMJVtKoWOzsKyu4ZNm7Ks7dzI3Tx5tEHc,109
-sawnergy-1.0.7.dist-info/METADATA,sha256=x0PQa0JilbayBcgywmnCL8IZZwTylzz8gOGnvwJHeDc,15433
-sawnergy-1.0.7.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-sawnergy-1.0.7.dist-info/top_level.txt,sha256=-67FQD6FD9Gjt74WTmO9hNYA3MLB4HaSxci0sEKC5Lo,9
-sawnergy-1.0.7.dist-info/RECORD,,
+sawnergy-1.0.8.dist-info/licenses/LICENSE,sha256=cElK4bCsDhyAEON3H05s35bQZvxBcXBiCOrOdiUhDCY,11346
+sawnergy-1.0.8.dist-info/licenses/NOTICE,sha256=eVTbuSasZrmMJVtKoWOzsKyu4ZNm7Ks7dzI3Tx5tEHc,109
+sawnergy-1.0.8.dist-info/METADATA,sha256=_0u1smFM5oMqaO0xuc4ZX094B6F2swQqUrolOkpikVM,16084
+sawnergy-1.0.8.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+sawnergy-1.0.8.dist-info/top_level.txt,sha256=-67FQD6FD9Gjt74WTmO9hNYA3MLB4HaSxci0sEKC5Lo,9
+sawnergy-1.0.8.dist-info/RECORD,,