sawnergy 1.0.6__tar.gz → 1.0.7__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sawnergy
-Version: 1.0.6
+Version: 1.0.7
 Summary: Toolkit for transforming molecular dynamics (MD) trajectories into rich graph representations
 Home-page: https://github.com/Yehor-Mishchyriak/SAWNERGY
 Author: Yehor Mishchyriak
@@ -39,18 +39,44 @@ 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 cpptraj executable on its own, so just run your code and don't worry about it.
+
+---
+
+# UPDATES:
+
+## v1.0.7 — What’s new:
+- **Added plain SkipGram 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.
+- **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.
+- **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 +117,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-normalise 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 +129,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`
@@ -140,23 +166,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` → **(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"` |
 
 **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.
 - 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,7 +197,7 @@ 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,
@@ -210,6 +226,7 @@ embeddings_path = embedder.embed_all(
     RIN_type="attr",
     using="merged",
     window_size=4,
+    objective="sgns",
     num_negative_samples=5,
     num_epochs=5,
     batch_size=1024,
@@ -232,12 +249,12 @@ print("Embeddings written to", embeddings_path)
 
 ---
 
-## 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,6 +267,13 @@ 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 = sawnergy.embedding.Visualizer("./EMBEDDINGS_demo.zip")
+viz.build_frame(1, show=True)
+```
+
 ---
 
 ## Advanced Notes

{sawnergy-1.0.6 → sawnergy-1.0.7}/README.md

@@ -5,18 +5,44 @@
 ![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 cpptraj executable on its own, so just run your code and don't worry about it.
+
+---
+
+# UPDATES:
+
+## v1.0.7 — What’s new:
+- **Added plain SkipGram 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.
+- **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.
+- **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 +83,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-normalise 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 +95,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`
@@ -106,23 +132,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` → **(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"` |
 
 **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.
 - 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,7 +163,7 @@ 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,
@@ -176,6 +192,7 @@ embeddings_path = embedder.embed_all(
     RIN_type="attr",
     using="merged",
     window_size=4,
+    objective="sgns",
     num_negative_samples=5,
     num_epochs=5,
     batch_size=1024,
@@ -198,12 +215,12 @@ print("Embeddings written to", embeddings_path)
 
 ---
 
-## 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,6 +233,13 @@ 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 = sawnergy.embedding.Visualizer("./EMBEDDINGS_demo.zip")
+viz.build_frame(1, show=True)
+```
+
 ---
 
 ## Advanced Notes

sawnergy-1.0.7/sawnergy/embedding/SGNS_pml.py

@@ -0,0 +1,370 @@
+from __future__ import annotations
+
+# third party
+import numpy as np
+from pureml.machinery import Tensor
+from pureml.layers import Embedding, Affine
+from pureml.losses import BCE, CCE
+from pureml.general_math import sum as t_sum
+from pureml.optimizers import Optim, LRScheduler
+from pureml.training_utils import TensorDataset, DataLoader, one_hot
+from pureml.base import NN
+
+# built-in
+import logging
+from typing import Type
+
+# *----------------------------------------------------*
+#                        GLOBALS
+# *----------------------------------------------------*
+
+_logger = logging.getLogger(__name__)
+
+# *----------------------------------------------------*
+#                        CLASSES
+# *----------------------------------------------------*
+
+class SGNS_PureML(NN):
+    """PureML implementation of Skip-Gram with Negative Sampling."""
+
+    def __init__(self,
+                V: int,
+                D: int,
+                *,
+                seed: int | None = None,
+                optim: Type[Optim],
+                optim_kwargs: dict,
+                lr_sched: Type[LRScheduler] | None = None,
+                lr_sched_kwargs: dict | None = None,
+                device: str | None = None):
+        """
+        Args:
+            V: Vocabulary size (number of nodes).
+            D: Embedding dimensionality.
+            seed: Optional RNG seed for negative sampling.
+            optim: Optimizer class to instantiate.
+            optim_kwargs: Keyword arguments for the optimizer (required).
+            lr_sched: Optional learning-rate scheduler class.
+            lr_sched_kwargs: Keyword arguments for the scheduler (required if lr_sched is provided).
+            device: Target device string (e.g. "cuda"); accepted for API parity, ignored by PureML.
+        """
+
+        if optim_kwargs is None:
+            raise ValueError("optim_kwargs must be provided")
+        if lr_sched is not None and lr_sched_kwargs is None:
+            raise ValueError("lr_sched_kwargs required when lr_sched is provided")
+
+        self.V, self.D = int(V), int(D)
+
+        # embeddings
+        self.in_emb  = Embedding(self.V, self.D)
+        self.out_emb = Embedding(self.V, self.D)
+
+        # seed + RNG for negative sampling
+        self.seed = None if seed is None else int(seed)
+        self._rng = np.random.default_rng(self.seed)
+        if self.seed is not None:
+            # optional: also set global NumPy seed for any non-RNG paths
+            np.random.seed(self.seed)
+
+        # API compatibility: PureML is CPU-only
+        self.device = "cpu"
+
+        # optimizer / scheduler
+        self.optim: Optim = optim(self.parameters, **optim_kwargs)
+        self.lr_sched: LRScheduler | None = (
+            lr_sched(optim=self.optim, **lr_sched_kwargs) if lr_sched is not None else None
+        )
+
+        _logger.info(
+            "SGNS_PureML init: V=%d D=%d device=%s seed=%s",
+            self.V, self.D, self.device, self.seed
+        )
+
+    def _sample_neg(self, B: int, K: int, dist: np.ndarray) -> np.ndarray:
+        return self._rng.choice(self.V, size=(B, K), replace=True, p=dist)
+
+    def predict(self, center: Tensor, pos: Tensor, neg: Tensor) -> tuple[Tensor, Tensor]:
+        """Compute positive/negative logits for SGNS.
+
+        Shapes:
+            center: (B,)
+            pos:    (B,)
+            neg:    (B, K)
+        Returns:
+            pos_logits: (B,)
+            neg_logits: (B, K)
+        """
+        c      = self.in_emb(center)      # (B, D)
+        pos_e  = self.out_emb(pos)        # (B, D)
+        neg_e  = self.out_emb(neg)        # (B, K, D)
+
+        pos_logits = t_sum(c * pos_e, axis=-1)                # (B,)
+        neg_logits = t_sum(c[:, None, :] * neg_e, axis=-1)    # (B, K)
+        return pos_logits, neg_logits
+
+    def fit(self,
+            centers: np.ndarray,
+            contexts: np.ndarray,
+            num_epochs: int,
+            batch_size: int,
+            num_negative_samples: int,
+            noise_dist: np.ndarray,
+            shuffle_data: bool,
+            lr_step_per_batch: bool):
+        """Train SGNS on the provided center/context pairs."""
+        _logger.info(
+            "SGNS_PureML fit: epochs=%d batch=%d negatives=%d shuffle=%s",
+            num_epochs, batch_size, num_negative_samples, shuffle_data
+        )
+
+        if noise_dist.ndim != 1 or noise_dist.size != self.V:
+            raise ValueError(f"noise_dist must be 1-D with length {self.V}; got {noise_dist.shape}")
+        dist = np.asarray(noise_dist, dtype=np.float64)
+        if np.any(dist < 0):
+            raise ValueError("noise_dist has negative entries")
+        s = dist.sum()
+        if not np.isfinite(s) or s <= 0:
+            raise ValueError("noise_dist must have positive finite sum")
+        if abs(s - 1.0) > 1e-6:
+            dist = dist / s
+
+        data = TensorDataset(centers, contexts)
+        for epoch in range(1, num_epochs + 1):
+            epoch_loss = 0.0
+            batches = 0
+
+            for cen, pos in DataLoader(data, batch_size=batch_size, shuffle=shuffle_data):
+                B = cen.data.shape[0] if isinstance(cen, Tensor) else len(cen)
+
+                neg_idx_np = self._sample_neg(B, num_negative_samples, dist)
+                neg = Tensor(neg_idx_np, requires_grad=False)
+                x_pos_logits, x_neg_logits = self(cen, pos, neg)
+
+                y_pos = Tensor(np.ones_like(x_pos_logits.numpy(copy=False)), requires_grad=False)
+                y_neg = Tensor(np.zeros_like(x_neg_logits.numpy(copy=False)), requires_grad=False)
+
+                K = int(neg.data.shape[1])
+                loss = (
+                    BCE(y_pos, x_pos_logits, from_logits=True)
+                    + K*BCE(y_neg, x_neg_logits, from_logits=True)
+                )
+
+                self.optim.zero_grad()
+                loss.backward()
+                self.optim.step()
+
+                if lr_step_per_batch and self.lr_sched is not None:
+                    self.lr_sched.step()
+
+                loss_value = float(np.asarray(loss.data))
+                epoch_loss += loss_value
+                batches += 1
+                _logger.debug("Epoch %d batch %d loss=%.6f", epoch, batches, loss_value)
+
+            if (not lr_step_per_batch) and (self.lr_sched is not None):
+                self.lr_sched.step()
+
+            mean_loss = epoch_loss / max(batches, 1)
+            _logger.info("Epoch %d/%d mean_loss=%.6f", epoch, num_epochs, mean_loss)
+
+    @property
+    def in_embeddings(self) -> np.ndarray:
+        W: Tensor = self.in_emb.parameters[0]   # (V, D)
+        if W.shape != (self.V, self.D):
+            raise RuntimeError(
+                "Wrong embedding matrix shape: "
+                "self.in_emb.parameters[0].shape != (V, D)"
+            )
+        return W.numpy(copy=True, readonly=True)
+
+    @property
+    def out_embeddings(self) -> np.ndarray:
+        W: Tensor = self.out_emb.parameters[0]  # (V, D)
+        if W.shape != (self.V, self.D):
+            raise RuntimeError(
+                "Wrong embedding matrix shape: "
+                "self.out_emb.parameters[0].shape != (V, D)"
+            )
+        return W.numpy(copy=True, readonly=True)
+
+    @property
+    def avg_embeddings(self) -> np.ndarray:
+        return 0.5 * (self.in_embeddings + self.out_embeddings)
+
+class SG_PureML(NN):
+    """Plain Skip-Gram (full softmax) in PureML.
+
+    Trains two affine layers to emulate the classic Skip-Gram objective with a
+    **full** softmax over the vocabulary (no negative sampling):
+
+        x = one_hot(center, V)                 # (B, V)
+        y = x @ W_in + b_in                    # (B, D)
+        logits = y @ W_out + b_out             # (B, V)
+        loss = CCE(one_hot(context, V), logits, from_logits=True)
+
+    The learnable “input” embeddings are the rows of `W_in` (shape `(V, D)`), and
+    the “output” embeddings are the rows of `W_outᵀ` (also `(V, D)`).
+    """
+
+    def __init__(self,
+                 V: int,
+                 D: int,
+                 *,
+                 seed: int | None = None,
+                 optim: Type[Optim],
+                 optim_kwargs: dict,
+                 lr_sched: Type[LRScheduler] | None = None,
+                 lr_sched_kwargs: dict | None = None,
+                 device: str | None = None):
+        """Initialize the plain Skip-Gram model (full softmax).
+
+        Args:
+            V: Vocabulary size (number of nodes/tokens).
+            D: Embedding dimensionality.
+            seed: Optional RNG seed (kept for API parity; not used in layer init).
+            optim: Optimizer class to instantiate (e.g., `Adam`, `SGD`).
+            optim_kwargs: Keyword arguments passed to the optimizer constructor.
+            lr_sched: Optional learning-rate scheduler class.
+            lr_sched_kwargs: Keyword arguments for the scheduler
+                (required if `lr_sched` is provided).
+            device: Device string (e.g., `"cuda"`). Accepted for parity, ignored
+                by PureML (CPU-only).
+
+        Notes:
+            The encoder/decoder are implemented as:
+              • `in_emb = Affine(V, D)` (acts on a one-hot center index)
+              • `out_emb = Affine(D, V)`
+            so forward pass produces vocabulary-sized logits.
+        """
+        if optim_kwargs is None:
+            raise ValueError("optim_kwargs must be provided")
+        if lr_sched is not None and lr_sched_kwargs is None:
+            raise ValueError("lr_sched_kwargs required when lr_sched is provided")
+
+        self.V, self.D = int(V), int(D)
+
+        # input/output “embedding” projections
+        self.in_emb  = Affine(self.V, self.D)
+        self.out_emb = Affine(self.D, self.V)
+
+        self.seed = None if seed is None else int(seed)
+
+        # API compatibility: PureML is CPU-only
+        self.device = "cpu"
+
+        # optimizer / scheduler
+        self.optim: Optim = optim(self.parameters, **optim_kwargs)
+        self.lr_sched: LRScheduler | None = (
+            lr_sched(optim=self.optim, **lr_sched_kwargs) if lr_sched is not None else None
+        )
+
+        _logger.info(
+            "SG_PureML init: V=%d D=%d device=%s seed=%s",
+            self.V, self.D, self.device, self.seed
+        )
+
+    def predict(self, center: Tensor) -> Tensor:
+        """Return vocabulary logits for each center index.
+
+        Args:
+            center: Tensor of center indices with shape `(B,)` and integer dtype.
+
+        Returns:
+            Tensor: Logits over the vocabulary with shape `(B, V)`.
+        """
+        c = one_hot(dims=self.V, label=center)  # (B, V)
+        y = self.in_emb(c)                      # (B, D)
+        z = self.out_emb(y)                     # (B, V)
+        return z
+
+    def fit(self,
+            centers: np.ndarray,
+            contexts: np.ndarray,
+            num_epochs: int,
+            batch_size: int,
+            shuffle_data: bool,
+            lr_step_per_batch: bool,
+            **_ignore):
+        """Train Skip-Gram with full softmax on center/context pairs.
+
+        Args:
+            centers: Array of center indices, shape `(N,)`, dtype integer in `[0, V)`.
+            contexts: Array of context (target) indices, shape `(N,)`, dtype integer.
+            num_epochs: Number of passes over the dataset.
+            batch_size: Mini-batch size.
+            shuffle_data: Whether to shuffle pairs each epoch.
+            lr_step_per_batch: If True, call `lr_sched.step()` after every batch
+                (when a scheduler is provided). If False, step once per epoch.
+            **_ignore: Ignored kwargs for API compatibility with SGNS.
+
+        Optimization:
+            Uses `CCE(one_hot(context), logits, from_logits=True)` where
+            `logits = predict(center)`. Scheduler stepping obeys `lr_step_per_batch`.
+        """
+        _logger.info(
+            "SG_PureML fit: epochs=%d batch=%d shuffle=%s",
+            num_epochs, batch_size, shuffle_data
+        )
+        data = TensorDataset(centers, contexts)
+
+        for epoch in range(1, num_epochs + 1):
+            epoch_loss = 0.0
+            batches = 0
+
+            for cen, ctx in DataLoader(data, batch_size=batch_size, shuffle=shuffle_data):
+                logits = self(cen)                          # (B, V)
+                y = one_hot(self.V, label=ctx)             # (B, V)
+                loss = CCE(y, logits, from_logits=True)    # scalar
+
+                self.optim.zero_grad()
+                loss.backward()
+                self.optim.step()
+
+                if lr_step_per_batch and self.lr_sched is not None:
+                    self.lr_sched.step()
+
+                loss_value = float(np.asarray(loss.data))
+                epoch_loss += loss_value
+                batches += 1
+                _logger.debug("Epoch %d batch %d loss=%.6f", epoch, batches, loss_value)
+
+            if (not lr_step_per_batch) and (self.lr_sched is not None):
+                self.lr_sched.step()
+
+            mean_loss = epoch_loss / max(batches, 1)
+            _logger.info("Epoch %d/%d mean_loss=%.6f", epoch, num_epochs, mean_loss)
+
+    @property
+    def in_embeddings(self) -> np.ndarray:
+        """Input embeddings matrix `W_in` as `(V, D)` (copy, read-only)."""
+        W = self.in_emb.parameters[0]              # (V, D)
+        if W.shape != (self.V, self.D):
+            raise RuntimeError(
+                "Wrong embedding matrix shape: "
+                "self.in_emb.parameters[0].shape != (V, D)"
+            )
+        return W.numpy(copy=True, readonly=True)   # (V, D)
+
+    @property
+    def out_embeddings(self) -> np.ndarray:
+        """Output embeddings matrix `W_outᵀ` as `(V, D)` (copy, read-only).
+        (`out_emb.parameters[0]` is `(D, V)`, so we transpose.)"""
+        W = self.out_emb.parameters[0]             # (D, V)
+        if W.shape != (self.D, self.V):
+            raise RuntimeError(
+                "Wrong embedding matrix shape: "
+                "self.out_emb.parameters[0].shape != (D, V)"
+            )
+        return W.numpy(copy=True, readonly=True).T # (V, D)
+    
+    @property
+    def avg_embeddings(self) -> np.ndarray:
+        """Elementwise average of input/output embeddings, shape `(V, D)`."""
+        return 0.5 * (self.in_embeddings + self.out_embeddings)  # (V, D)
+
+
+__all__ = ["SGNS_PureML", "SG_PureML"]
+
+if __name__ == "__main__":
+    pass