sawnergy 1.0.6__tar.gz → 1.0.8__tar.gz

{sawnergy-1.0.6/sawnergy.egg-info → sawnergy-1.0.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sawnergy
-Version: 1.0.6
+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
@@ -39,19 +39,57 @@ Dynamic: summary
 ![Python](https://img.shields.io/badge/python-3.11%2B-blue)
 
 A toolkit for transforming molecular dynamics (MD) trajectories into rich graph representations, sampling
-random and self-avoiding walks, learning node embeddings, and visualising residue interaction networks (RINs). SAWNERGY
+random and self-avoiding walks, learning node embeddings, and visualizing residue interaction networks (RINs). SAWNERGY
 keeps the full workflow — from `cpptraj` output to skip-gram embeddings (node2vec approach) — inside Python, backed by efficient Zarr-based archives and optional GPU acceleration.
 
 ---
 
+## Installation
+
+   ```bash
+   pip install sawnergy
+   ```
+
+> **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 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 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 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.
+
+---
+
 ## Why SAWNERGY?
 
 - **Bridge simulations and graph ML**: Convert raw MD trajectories into residue interaction networks ready for graph
   algorithms and downstream machine learning tasks.
-- **Deterministic, shareable artefacts**: 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 serealization overhead; archives are written in chunked, compressed form for fast read/write.
-- **Flexible embedding backends**: Train skip-gram with negative sampling (SGNS) models using either PureML or PyTorch.
-- **Visualization out of the box**: Plot and animate residue networks without leaving Python, using the data produced by RINBuilder
+- **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.
 
 ---
 
@@ -91,9 +129,9 @@ 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, symmetrise, remove self-interactions, and L1-normalise the matrices,
-  - compute per-residue centres of mass (COM) over the same frames.
-* Outputs a compressed Zarr archive with transition matrices, optional prenormalised energies, COM snapshots, and rich
+  - 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.).
 * Supports parallel `cpptraj` execution, batch processing, and keeps temporary stores tidy via
   `ArrayStorage.compress_and_cleanup`.
@@ -103,7 +141,7 @@ node indexing, and RNG seeds stay consistent across the toolchain.
 * Opens RIN archives, resolves dataset names from attributes, and renders nodes plus attractive/repulsive edge bundles
   in 3D using Matplotlib.
 * Allows both static frame visualization and trajectory animation.
-* Handles backend selection (`Agg` fallback in headless environments) and offers convenient colour palettes via
+* Handles backend selection (`Agg` fallback in headless environments) and offers convenient color palettes via
   `visualizer_util`.
 
 ### `sawnergy.walks.Walker`
@@ -116,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
 
@@ -140,23 +175,13 @@ 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) |
+| **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.
-
----
-
-## Installation
-
-   ```bash
-   pip install sawnergy
-   ```
-
-> **Note:** RIN building requires `cpptraj` (AmberTools). Ensure it is discoverable via `$PATH` or the `CPPTRAJ`
-> environment variable.
+- In **Embeddings**, `alpha` and `num_negative_samples` apply to **SGNS** only and are ignored for `objective="sg"`.
 
 ---
 
@@ -181,10 +206,10 @@ rin_builder.build_rin(
     molecule_of_interest=1,
     frame_range=(1, 100),
     frame_batch_size=10,
-    prune_low_energies_frac=0.3,
+    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
@@ -192,52 +217,43 @@ 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,
-    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`.
 
 ---
 
-## Visualisation
+## Visualization
 
 ```python
 from sawnergy.visual import Visualizer
 
-v = sawnergy.visual.Visualizer("./RIN_demo.zip")
+v = Visualizer("./RIN_demo.zip")
 v.build_frame(1,
     node_colors="rainbow",
     displayed_nodes="ALL",
@@ -250,14 +266,20 @@ v.build_frame(1,
 
 `Visualizer` lazily loads datasets and works even in headless environments (falls back to the `Agg` backend).
 
+```python
+from sawnergy.embedding import Visualizer
+
+viz = Visualizer("./EMBEDDINGS_demo.zip", normalize_rows=True)
+viz.build_frame(1, show=True)
+```
+
 ---
 
 ## Advanced Notes
 
 - **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.
 
 ---
@@ -268,8 +290,9 @@ v.build_frame(1,
 ├── 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
 │
@@ -278,7 +301,7 @@ v.build_frame(1,
 
 ---
 
-## 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.6 → sawnergy-1.0.8}/README.md

@@ -5,19 +5,57 @@
 ![Python](https://img.shields.io/badge/python-3.11%2B-blue)
 
 A toolkit for transforming molecular dynamics (MD) trajectories into rich graph representations, sampling
-random and self-avoiding walks, learning node embeddings, and visualising residue interaction networks (RINs). SAWNERGY
+random and self-avoiding walks, learning node embeddings, and visualizing residue interaction networks (RINs). SAWNERGY
 keeps the full workflow — from `cpptraj` output to skip-gram embeddings (node2vec approach) — inside Python, backed by efficient Zarr-based archives and optional GPU acceleration.
 
 ---
 
+## Installation
+
+   ```bash
+   pip install sawnergy
+   ```
+
+> **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 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 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 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.
+
+---
+
 ## Why SAWNERGY?
 
 - **Bridge simulations and graph ML**: Convert raw MD trajectories into residue interaction networks ready for graph
   algorithms and downstream machine learning tasks.
-- **Deterministic, shareable artefacts**: 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 serealization overhead; archives are written in chunked, compressed form for fast read/write.
-- **Flexible embedding backends**: Train skip-gram with negative sampling (SGNS) models using either PureML or PyTorch.
-- **Visualization out of the box**: Plot and animate residue networks without leaving Python, using the data produced by RINBuilder
+- **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.
 
 ---
 
@@ -57,9 +95,9 @@ 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, symmetrise, remove self-interactions, and L1-normalise the matrices,
-  - compute per-residue centres of mass (COM) over the same frames.
-* Outputs a compressed Zarr archive with transition matrices, optional prenormalised energies, COM snapshots, and rich
+  - 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.).
 * Supports parallel `cpptraj` execution, batch processing, and keeps temporary stores tidy via
   `ArrayStorage.compress_and_cleanup`.
@@ -69,7 +107,7 @@ node indexing, and RNG seeds stay consistent across the toolchain.
 * Opens RIN archives, resolves dataset names from attributes, and renders nodes plus attractive/repulsive edge bundles
   in 3D using Matplotlib.
 * Allows both static frame visualization and trajectory animation.
-* Handles backend selection (`Agg` fallback in headless environments) and offers convenient colour palettes via
+* Handles backend selection (`Agg` fallback in headless environments) and offers convenient color palettes via
   `visualizer_util`.
 
 ### `sawnergy.walks.Walker`
@@ -82,13 +120,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
 
@@ -106,23 +141,13 @@ 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) |
+| **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.
-
----
-
-## Installation
-
-   ```bash
-   pip install sawnergy
-   ```
-
-> **Note:** RIN building requires `cpptraj` (AmberTools). Ensure it is discoverable via `$PATH` or the `CPPTRAJ`
-> environment variable.
+- In **Embeddings**, `alpha` and `num_negative_samples` apply to **SGNS** only and are ignored for `objective="sg"`.
 
 ---
 
@@ -147,10 +172,10 @@ rin_builder.build_rin(
     molecule_of_interest=1,
     frame_range=(1, 100),
     frame_batch_size=10,
-    prune_low_energies_frac=0.3,
+    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
@@ -158,52 +183,43 @@ 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,
-    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`.
 
 ---
 
-## Visualisation
+## Visualization
 
 ```python
 from sawnergy.visual import Visualizer
 
-v = sawnergy.visual.Visualizer("./RIN_demo.zip")
+v = Visualizer("./RIN_demo.zip")
 v.build_frame(1,
     node_colors="rainbow",
     displayed_nodes="ALL",
@@ -216,14 +232,20 @@ v.build_frame(1,
 
 `Visualizer` lazily loads datasets and works even in headless environments (falls back to the `Agg` backend).
 
+```python
+from sawnergy.embedding import Visualizer
+
+viz = Visualizer("./EMBEDDINGS_demo.zip", normalize_rows=True)
+viz.build_frame(1, show=True)
+```
+
 ---
 
 ## Advanced Notes
 
 - **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.
 
 ---
@@ -234,8 +256,9 @@ v.build_frame(1,
 ├── 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
 │
@@ -244,7 +267,7 @@ v.build_frame(1,
 
 ---
 
-## 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.