sawnergy 1.0.4__tar.gz → 1.0.6__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sawnergy
-Version: 1.0.4
+Version: 1.0.6
 Summary: Toolkit for transforming molecular dynamics (MD) trajectories into rich graph representations
 Home-page: https://github.com/Yehor-Mishchyriak/SAWNERGY
 Author: Yehor Mishchyriak

{sawnergy-1.0.4 → sawnergy-1.0.6}/sawnergy/__init__.py

@@ -3,11 +3,13 @@ from . import logging_util
 from . import rin
 from . import visual
 from . import walks
+from . import embedding
 
 __all__ = [
     "sawnergy_util",
     "logging_util",
     "rin",
     "visual",
-    "walks"
+    "walks",
+    "embedding"
 ]

{sawnergy-1.0.4 → sawnergy-1.0.6}/sawnergy/embedding/SGNS_pml.py

@@ -34,45 +34,76 @@ class SGNS_PureML(NN):
                 seed: int | None = None,
                 optim: Type[Optim],
                 optim_kwargs: dict,
-                lr_sched: Type[LRScheduler],
-                lr_sched_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: PureML optimizer class.
-            optim_kwargs: Keyword arguments forwarded to the optimizer.
-            lr_sched: PureML learning-rate scheduler class.
-            lr_sched_kwargs: Keyword arguments forwarded to the scheduler.
+            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)
-        self.in_emb  = Embedding(V, D)
-        self.out_emb = Embedding(V, 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
+        )
 
-        self.optim: Optim          = optim(self.parameters, **optim_kwargs)
-        self.lr_sched: LRScheduler = lr_sched(**lr_sched_kwargs)
-        _logger.info("SGNS_PureML init: V=%d D=%d seed=%s", self.V, self.D, self.seed)
+        _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):
+    def _sample_neg(self, B: int, K: int, dist: np.ndarray) -> np.ndarray:
         """Draw negative samples according to the provided unigram distribution."""
         if dist.ndim != 1 or dist.size != self.V:
             raise ValueError(f"noise_dist must be 1-D with length {self.V}; got {dist.shape}")
         return self._rng.choice(self.V, size=(B, K), replace=True, p=dist)
 
-    def predict(self, center: Tensor, pos: Tensor, neg: Tensor) -> Tensor:
-        """Compute positive/negative logits for SGNS."""
-        c      = self.in_emb(center)
-        pos_e  = self.out_emb(pos)
-        neg_e  = self.out_emb(neg)
-        pos_logits = t_sum(c * pos_e, axis=-1)
-        neg_logits = t_sum(c[:, None, :] * neg_e, axis=-1)
-        #                       ^^^
-        # (B,1,D) * (B,K,D) → (B,K,D) → sum D → (B,K)
+    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,
@@ -94,29 +125,35 @@ class SGNS_PureML(NN):
         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):
-                neg = self._sample_neg(batch_size, num_negative_samples, noise_dist)
+                B = cen.data.shape[0] if isinstance(cen, Tensor) else len(cen)
 
+                neg_idx_np = self._sample_neg(B, num_negative_samples, noise_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.data))
                 y_neg = Tensor(np.zeros_like(x_neg_logits.data))
 
-                loss = BCE(y_pos, x_pos_logits, from_logits=True) + BCE(y_neg, x_neg_logits, from_logits=True)
+                loss = (
+                    BCE(y_pos, x_pos_logits, from_logits=True)
+                    + BCE(y_neg, x_neg_logits, from_logits=True)
+                )
 
                 self.optim.zero_grad()
                 loss.backward()
                 self.optim.step()
-                
-                if lr_step_per_batch:
+
+                if lr_step_per_batch and self.lr_sched is not None:
                     self.lr_sched.step()
 
-                loss_value = float(np.asarray(loss.data).mean())
+                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:
+            if (not lr_step_per_batch) and (self.lr_sched is not None):
                 self.lr_sched.step()
 
             mean_loss = epoch_loss / max(batches, 1)
@@ -124,7 +161,7 @@ class SGNS_PureML(NN):
 
     @property
     def embeddings(self) -> np.ndarray:
-        """Return the input embedding matrix as a NumPy array."""
+        """Return the input embedding matrix as a NumPy array (V, D)."""
         W: Tensor = self.in_emb.parameters[0]
         return np.asarray(W.data)
 

{sawnergy-1.0.4 → sawnergy-1.0.6}/sawnergy/embedding/embedder.py

@@ -330,6 +330,7 @@ class Embedder:
               num_epochs: int,
               batch_size: int,
               *,
+              lr_step_per_batch: bool = False,
               shuffle_data: bool = True,
               dimensionality: int = 128,
               alpha: float = 0.75,
@@ -355,7 +356,8 @@ 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.
             _seed: Optional child seed for this frame's model initialization.
 
         Returns:
@@ -391,10 +393,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({
@@ -429,7 +435,7 @@ class Embedder:
             num_negative_samples,
             noise_probs,
             shuffle_data,
-            lr_step_per_batch=False
+            lr_step_per_batch
         )
 
         embeddings = getattr(self.model, "embeddings", None)

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sawnergy
-Version: 1.0.4
+Version: 1.0.6
 Summary: Toolkit for transforming molecular dynamics (MD) trajectories into rich graph representations
 Home-page: https://github.com/Yehor-Mishchyriak/SAWNERGY
 Author: Yehor Mishchyriak