sawnergy 1.0.0__py3-none-any.whl

sawnergy/walks/walker_util.py

@@ -0,0 +1,384 @@
+# third-party
+import numpy as np
+# built-in
+from multiprocessing import shared_memory
+import logging
+
+# *----------------------------------------------------*
+#                        GLOBALS
+# *----------------------------------------------------*
+
+_logger = logging.getLogger(__name__)
+
+# *----------------------------------------------------*
+#                        CLASSES
+# *----------------------------------------------------*
+
+class SharedNDArray:
+    """NumPy-facing wrapper over a raw :class:`multiprocessing.shared_memory.SharedMemory`.
+
+    This class does **not** own any data itself; it wraps an OS-level shared
+    memory segment and exposes it as a NumPy array via zero-copy views
+    (shape/dtype provided by the caller). The underlying buffer is just a
+    flat byte block; dimensionality and strides come from the views you
+    construct.
+
+    Usage model:
+      - Create a segment in the parent with :meth:`create`, optionally seeding
+        from an existing array (copied once, C-contiguous).
+      - Pass ``(name, shape, dtype)`` to workers and attach with :meth:`attach`.
+      - Obtain a view with :py:meth:`view` or the :py:attr:`array` property.
+        Views are read-only by default unless ``default_readonly=False`` or
+        ``view(readonly=False)`` is requested.
+      - Every process that opened the segment must call :meth:`close`.
+        Exactly one process should call :meth:`unlink` after all others have
+        closed to destroy the OS resource.
+
+    Indexing:
+      - ``__getitem__`` strictly supports **axis-0** basic indexing
+        (``None``, ``slice``, or ``int``). This guarantees **no-copy** views.
+        Fancy indexing (index arrays/boolean masks) is intentionally disallowed.
+      - For 1D arrays, ``int`` indexing would yield a NumPy scalar (not a view),
+        so it is rejected; use ``slice(i, i+1)`` for a one-row view instead.
+
+    Concurrency:
+      - Multiple readers are safe by design.
+      - If multiple writers may overlap, synchronize externally (e.g., a
+        :class:`multiprocessing.Lock`). The class does not implement locking.
+
+    Notes:
+      - The writeability flag is **per-view**. Marking one view read-only does
+        not prevent other processes (or other views) from writing.
+      - Shape/dtype are trusted by :meth:`attach`—they must match what was used
+        at creation time; no runtime validation is performed here.
+    """
+
+    def __init__(self,
+                shm: shared_memory.SharedMemory,
+                shape: tuple[int, ...],
+                dtype: np.dtype,
+                *,
+                default_readonly: bool = True):
+        """Construct a wrapper over an existing shared memory handle.
+
+        Args:
+            shm: An open :class:`SharedMemory` handle (already created/attached).
+            shape: Target array shape used for all views into this buffer.
+            dtype: Target NumPy dtype used for all views into this buffer.
+            default_readonly: If ``True``, views returned by :py:attr:`array`
+                are marked read-only; override per-call via :py:meth:`view`.
+
+        Remarks:
+            This constructor does not allocate memory; it only stores metadata.
+            Use :meth:`create` to allocate a new segment, or :meth:`attach`
+            to connect to an existing one by name.
+        """
+        self.shm   = shm
+        self.shape = tuple(shape)
+        self.dtype = np.dtype(dtype)
+        self._default_readonly = default_readonly
+        _logger.debug(
+            "SharedNDArray.__init__(name=%r, shape=%s, dtype=%s, default_readonly=%s)",
+            getattr(self.shm, "name", None), self.shape, self.dtype, self._default_readonly
+        )
+
+    def __len__(self) -> int:
+        """Return the size of axis 0 (NumPy semantics).
+
+        Returns:
+            The number of elements along the first dimension.
+
+        Raises:
+            TypeError: If the wrapped array is 0-D (unsized).
+        """
+        if len(self.shape) == 0:
+            _logger.error("len() called on 0-D array (shape=%s)", self.shape)
+            raise TypeError("len() of unsized object")
+        length = self.shape[0]
+        _logger.debug("__len__ -> %d", length)
+        return length
+
+    def __repr__(self):
+        """Debug-friendly representation showing name/shape/dtype."""
+        return f"SharedNDArray(name={self.name!r}, shape={self.shape}, dtype={self.dtype})"
+
+    def __getitem__(self, ids: int | slice | None = None):
+        """Axis-0 only, no-copy guaranteed.
+        - None      -> full view
+        - slice     -> view
+        - int       -> view (requires ndim >= 2); for 1D, use slice(i, i+1)
+        """
+        _logger.debug("__getitem__(ids=%r)", ids)
+        arr = self.array 
+        if ids is None:
+            _logger.debug("__getitem__: returning full view")
+            return arr
+        if isinstance(ids, slice):
+            _logger.debug("__getitem__: slice=%s", ids)
+            return arr[ids, ...]
+        if isinstance(ids, int):
+            if arr.ndim == 1:
+                _logger.error(
+                    "__getitem__: 1D int indexing requested (idx=%r) -> would copy; raising",
+                    ids
+                )
+                raise TypeError(
+                    "No-copy view for 1D int indexing is impossible. "
+                    "Use slice(i, i+1) to get a 1-row view."
+                )
+            _logger.debug("__getitem__: int=%d", ids)
+            return arr[ids, ...]
+        _logger.error("__getitem__: unsupported key type %s", type(ids).__name__)
+        raise TypeError("Only axis-0 int/slice/None are allowed for no-copy access.")
+
+    @classmethod
+    def attach(cls, name: str, shape, dtype):
+        """Attach to an existing shared memory segment by name.
+
+        Args:
+            name: System-wide shared memory name (as returned by :py:attr:`name`).
+            shape: Shape to interpret the buffer with (must match creator).
+            dtype: Dtype to interpret the buffer with (must match creator).
+
+        Returns:
+            A :class:`SharedNDArray` bound to the named segment.
+
+        Raises:
+            FileNotFoundError: If no segment with ``name`` exists.
+            PermissionError: If the segment exists but cannot be opened.
+
+        Notes:
+            This method trusts ``shape`` and ``dtype``; it does not verify that
+            they match the original settings. Passing inconsistent metadata
+            results in undefined views.
+        """
+        _logger.debug("SharedNDArray.attach(name=%r, shape=%s, dtype=%s)", name, shape, np.dtype(dtype))
+        shm = shared_memory.SharedMemory(name=name, create=False)
+        obj = cls(shm, shape, dtype)
+        _logger.debug("Attached to shared memory: name=%r", obj.name)
+        return obj
+
+    @classmethod
+    def create(cls, shape, dtype, *, from_array=None, name: str | None = None):
+        """Create a new shared memory segment and wrap it.
+
+        The allocated buffer is sized exactly as ``prod(shape) * dtype.itemsize``.
+        If ``from_array`` is provided, its contents are copied into the buffer
+        after being coerced to a C-contiguous array of ``dtype``. Otherwise the
+        buffer is zero-initialized.
+
+        Args:
+            shape: Desired array shape.
+            dtype: Desired NumPy dtype.
+            from_array: Optional source array to seed the buffer. Must match
+                ``shape`` after coercion to ``dtype``; copied as C-contiguous.
+            name: Optional OS-visible name for the segment. If omitted, a unique
+                name is generated.
+
+        Returns:
+            A :class:`SharedNDArray` bound to the newly created segment.
+
+        Raises:
+            ValueError: If ``from_array`` shape does not match ``shape`` after
+                dtype coercion.
+        """
+        dtype = np.dtype(dtype)
+        nbytes = int(np.prod(shape)) * dtype.itemsize
+        _logger.debug("SharedNDArray.create(shape=%s, dtype=%s, name=%r, nbytes=%d)", shape, dtype, name, nbytes)
+        shm = shared_memory.SharedMemory(create=True, size=nbytes, name=name)
+        obj = cls(shm, shape, dtype)
+
+        view = np.ndarray(shape, dtype=dtype, buffer=shm.buf)
+        if from_array is not None:
+            src = np.ascontiguousarray(from_array, dtype=dtype)
+            if src.shape != tuple(shape):
+                _logger.error("create: source shape %s does not match %s", src.shape, shape)
+                raise ValueError(f"shape mismatch: {src.shape} vs {shape}")
+            view[...] = src
+            _logger.debug("create: seeded from array (shape=%s, dtype=%s)", src.shape, src.dtype)
+        else:
+            view.fill(0)
+            _logger.debug("create: zero-initialized buffer")
+        _logger.debug("create: created shared segment name=%r", obj.name)
+        return obj
+
+    def close(self) -> None:
+        """Detach this process from the shared memory segment.
+
+        Call this in **every** process that opened/attached the segment.
+        After closing, any existing views into the buffer must **not** be used
+        unless you first copy them (e.g., ``np.array(view, copy=True)``).
+        """
+        _logger.debug("close(): name=%r", self.name)
+        self.shm.close()
+
+    def unlink(self) -> None:
+        """Destroy the shared memory segment (OS resource).
+
+        Call exactly **once** globally after all participating processes have
+        called :meth:`close`. After unlinking, the ``name`` may be reused by
+        the OS for new segments.
+        """
+        _logger.debug("unlink(): name=%r", self.name)
+        self.shm.unlink()
+
+    def view(self, *, readonly: bool | None = None) -> np.ndarray: # if readonly is False, arr is mutable
+        """Return a zero-copy NumPy view over the shared buffer.
+
+        Args:
+            readonly: If ``True``, the returned view is marked read-only.
+                If ``False``, the view is writable. If ``None`` (default),
+                the behavior follows ``self._default_readonly``.
+
+        Returns:
+            A NumPy ndarray that directly references the shared bytes using
+            the stored ``shape`` and ``dtype``.
+
+        Notes:
+            - The writeability flag is **per-view**; it does not affect other
+              views or other processes.
+            - Basic slicing of the returned array yields further views that
+              inherit the writeability flag; fancy indexing creates copies.
+        """
+        arr = np.ndarray(self.shape, dtype=self.dtype, buffer=self.shm.buf)
+        ro = self._default_readonly if readonly is None else readonly
+        _logger.debug("view(readonly=%r) -> resolved_readonly=%r", readonly, ro)
+        if ro:
+            arr.flags.writeable = False
+        return arr
+
+    @property
+    def name(self) -> str:
+        """System-wide name of the underlying shared memory segment."""
+        return self.shm.name
+
+    @property
+    def array(self) -> np.ndarray:
+        """Default zero-copy view honoring ``default_readonly``."""
+        _logger.debug("array property accessed (default_readonly=%r)", self._default_readonly)
+        return self.view(readonly=self._default_readonly)
+
+# *----------------------------------------------------*
+#                       FUNCTIONS
+# *----------------------------------------------------*
+
+def l1_norm(X: np.ndarray) -> np.ndarray:
+    """Return an L1-normalized copy of ``X`` (sum to 1), or zeros if invalid.
+
+    Args:
+        X (np.ndarray): Array of nonnegative weights/probabilities (any shape).
+            It is coerced with ``np.asarray(X, dtype=float)``.
+
+    Returns:
+        np.ndarray: Array with the same shape as ``X`` whose entries sum to 1
+        (within FP error). If the total mass is non-finite or <= 0, returns
+        an array of zeros with the same shape/dtype.
+
+    Notes:
+        - If ``X`` contains NaNs or Infs, the sum becomes non-finite and a
+          zeros array is returned.
+        - Works for any shape; normalization is over all elements.
+    """
+    X = np.asarray(X, dtype=float)
+    s = float(np.sum(X))
+    if not np.isfinite(s) or s <= 0.0:
+        return np.zeros_like(X)
+    return X / s
+
+
+def apply_on_axis0(X: np.ndarray, func):
+    """Apply a function independently to each slice ``X[i]`` along axis 0.
+
+    ``func`` is called once per ``i`` with a view/copy of ``X[i]`` and its
+    first result is used to allocate the output array.
+
+    Args:
+        X (np.ndarray): Input array of shape ``(N, ...)`` where ``N >= 1``.
+        func (Callable): Function taking ``X[i]`` (shape ``X.shape[1:]``) and
+            returning an array-like object. All returns must be broadcast-
+            compatible and have identical shape.
+
+    Returns:
+        np.ndarray: Stacked results with shape ``(N,) + out0.shape``, where
+        ``out0`` is ``func(X[0])``. The dtype matches ``np.asarray(out0).dtype``.
+
+    Raises:
+        IndexError: If ``X`` is empty along axis 0 (i.e., ``X.shape[0] == 0``).
+
+    Notes:
+        The first call to ``func`` determines the output dtype and shape.
+    """
+    X = np.asarray(X)
+    out0 = func(X[0])
+    # the 0th axis has to have as many dims as the X array has along the 0th axis;
+    # as for the other axes, they coincide in dimensionality with the output of func
+    out = np.empty((X.shape[0],) + np.shape(out0), dtype=np.asarray(out0).dtype)
+    out[0] = out0
+    for i in range(1, X.shape[0]):
+        out[i] = func(X[i])
+    return out
+
+
+def cosine_similarity(A: np.ndarray, eps: float = 1e-12):
+    """Create a callable that computes cosine similarity to a fixed array ``A``.
+
+    The returned function takes an array ``B`` (same shape as ``A``), computes
+    the cosine similarity between ``A`` and ``B`` (using flattened views),
+    and maps it from ``[-1, 1]`` to ``[0, 1]`` via ``(cos + 1) / 2``.
+    If either vector has norm below ``eps``, it returns ``0.0``.
+
+    Args:
+        A (np.ndarray): Reference array. Coerced with ``np.asarray``.
+        eps (float, optional): Small threshold to guard against division by
+            near-zero norms. Defaults to ``1e-12``.
+
+    Returns:
+        Callable[[np.ndarray], float]: Function ``inner(B)`` that returns a
+        similarity score in ``[0, 1]``.
+
+    Raises:
+        ValueError: If the input ``B`` provided to the returned function does
+            not match the shape of ``A``.
+    """
+
+    def inner(B: np.ndarray):
+        """Compute cosine similarity between the captured ``A`` and input ``B``.
+
+        Args:
+            B (np.ndarray): Array with the same shape as ``A``.
+
+        Returns:
+            float: Cosine similarity mapped to ``[0, 1]``. Returns ``0.0`` if
+            the product of the norms is below ``eps``.
+
+        Raises:
+            ValueError: If ``A.shape != B.shape``.
+        """
+        nonlocal A
+        nonlocal eps
+
+        A = np.asarray(A)
+        B = np.asarray(B)
+        if A.shape != B.shape:
+            raise ValueError(f"shapes must match, got {A.shape} vs {B.shape}")
+
+        a = A.ravel()
+        b = B.ravel()
+
+        denom = np.linalg.norm(a) * np.linalg.norm(b)
+        if denom < eps:
+            return 0.0
+        return (float(a @ b / denom) + 1) / 2 # translate from [-1, 1] to [0, 2] to [0, 1]
+
+    return inner
+
+
+__all__ = [
+    "SharedNDArray",
+    "l1_norm",
+    "apply_on_axis0",
+    "cosine_similarity"
+]
+
+if __name__ == "__main__":
+    pass

sawnergy-1.0.0.dist-info/METADATA

@@ -0,0 +1,290 @@
+Metadata-Version: 2.4
+Name: sawnergy
+Version: 1.0.0
+Summary: Toolkit for transforming molecular dynamics (MD) trajectories into rich graph representations
+Home-page: https://github.com/Yehor-Mishchyriak/SAWNERGY
+Author: Yehor Mishchyriak
+License: Apache-2.0
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE
+License-File: NOTICE
+Requires-Dist: numpy>=2.0
+Requires-Dist: zarr>=3.0
+Requires-Dist: threadpoolctl>=3.0
+Requires-Dist: matplotlib>=3.7
+Requires-Dist: psutil>=5.9
+Requires-Dist: ym-pure-ml>=1.2.0
+Dynamic: author
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: license
+Dynamic: license-file
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
+
+# SAWNERGY
+
+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
+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.
+
+---
+
+## 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
+
+---
+
+## Pipeline at a Glance
+
+```
+MD Trajectory + Topology
+          │
+          ▼
+      RINBuilder 
+          │   →  RIN archive (.zip/.zarr) → Visualizer (display/animate RINs)
+          ▼
+        Walker
+          │   →  Walks archive (RW/SAW per frame)
+          ▼
+       Embedder
+          │   →  Embedding archive (frame × vocab × dim)
+          ▼
+     Downstream ML
+```
+
+Each stage consumes the archive produced by the previous one. Metadata embedded in the archives ensures frame order,
+node indexing, and RNG seeds stay consistent across the toolchain.
+
+---
+
+## Core Components
+
+### `sawnergy.rin.RINBuilder`
+
+* 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
+  metadata (frame range, pruning quantile, molecule ID, etc.).
+* Supports parallel `cpptraj` execution, batch processing, and keeps temporary stores tidy via
+  `ArrayStorage.compress_and_cleanup`.
+
+### `sawnergy.visual.Visualizer`
+
+* 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
+  `visualizer_util`.
+
+### `sawnergy.walks.Walker`
+
+* Attaches to the RIN archive and loads attractive/repulsive transition matrices into shared memory using
+  `walker_util.SharedNDArray` so multiple processes can sample without copying.
+* Samples random walks (RW) and self-avoiding walks (SAW), optionally time-aware, that is, walks move through transition matrices with transition probabilities proportional to cosine similarity between the current and next frame. Randomness is controlled by the seed passed to the class constructor.
+* Persists walks as `(time, walk_id, length+1)` tensors (1-based node indices) alongside metadata such as
+  `walk_length`, `walks_per_node`, and RNG scheme.
+
+### `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.
+
+### Supporting Utilities
+
+* `sawnergy.sawnergy_util`
+  - `ArrayStorage`: thin wrapper over Zarr v3 with helpers for chunk management, attribute coercion to JSON, and transparent compression to `.zip` archives.
+  - Parallel helpers (`elementwise_processor`, `compose_steps`, etc.), temporary file management, logging, and runtime
+    inspection utilities.
+* `sawnergy.logging_util.configure_logging`: configure rotating file/console logging consistently across scripts.
+
+---
+
+## Archive Layouts
+
+| Archive | Key datasets (name → shape, dtype) | Important attributes (root `attrs`) |
+|---|---|---|
+| **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) |
+
+**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.
+
+---
+
+## Quick Start
+
+```python
+from pathlib import Path
+from sawnergy.logging_util import configure_logging
+from sawnergy.rin import RINBuilder
+from sawnergy.walks import Walker
+from sawnergy.embedding import Embedder
+
+import logging
+configure_logging("./logs", file_level=logging.WARNING, console_level=logging.INFO)
+
+# 1. Build a Residue Interaction Network archive
+rin_path = Path("./RIN_demo.zip")
+rin_builder = RINBuilder()
+rin_builder.build_rin(
+    topology_file="system.prmtop",
+    trajectory_file="trajectory.nc",
+    molecule_of_interest=1,
+    frame_range=(1, 100),
+    frame_batch_size=10,
+    prune_low_energies_frac=0.3,
+    output_path=rin_path,
+    include_attractive=True,
+    include_repulsive=False,
+)
+
+# 2. Sample walks from the RIN
+walker = Walker(rin_path, seed=123)
+walks_path = Path("./WALKS_demo.zip")
+walker.sample_walks(
+    walk_length=16,
+    walks_per_node=32,
+    saw_frac=0.25,
+    include_attractive=True,
+    include_repulsive=False,
+    time_aware=False,
+    output_path=walks_path,
+    in_parallel=False,
+)
+walker.close()
+
+# 3. Train embeddings per frame (PyTorch backend)
+import torch
+
+embedder = Embedder(walks_path, base="torch", seed=999)
+embeddings_path = embedder.embed_all(
+    RIN_type="attr",
+    using="merged",
+    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",
+    },
+)
+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`).
+
+---
+
+## Visualisation
+
+```python
+from sawnergy.visual import Visualizer
+
+v = sawnergy.visual.Visualizer("./RIN_demo.zip")
+v.build_frame(1,
+    node_colors="rainbow",
+    displayed_nodes="ALL",
+    displayed_pairwise_attraction_for_nodes="DISPLAYED_NODES",
+    displayed_pairwise_repulsion_for_nodes="DISPLAYED_NODES",
+    show_node_labels=True,
+    show=True
+)
+```
+
+`Visualizer` lazily loads datasets and works even in headless environments (falls back to the `Agg` backend).
+
+---
+
+## 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).
+- **ArrayStorage utilities**: Use `ArrayStorage` directly to peek into archives, append arrays, or manage metadata.
+
+---
+
+## Testing & Quality Assurance
+
+The automated test suite (`pytest`) synthesises deterministic cpptraj outputs and exercises the entire workflow:
+
+- RIN parsing, residue aggregation, and metadata verification.
+- Random/self-avoiding walk sampling and probability consistency with the RIN.
+- Embedding orchestration, frame ordering, SGNS pair generation property tests.
+- PureML and PyTorch SGNS smoke tests verifying finite weights and decreasing loss.
+- Visualiser smoke tests that cover data loading and artist updates.
+
+Run the suite (inside the project virtualenv):
+
+```bash
+python -m pytest
+```
+
+---
+
+## Project Structure
+
+```
+├── sawnergy/
+│   ├── rin/           # RINBuilder and cpptraj integration helpers
+│   ├── walks/         # Walker class and shared-memory utilities
+│   ├── embedding/     # Embedder + SGNS backends (PureML / PyTorch)
+│   ├── visual/        # Visualizer and palette utilities
+│   ├── logging_util.py
+│   └── sawnergy_util.py
+├── tests/             # Synthetic end-to-end tests (pytest)
+│
+└── README.md
+```
+
+---
+
+## Acknowledgements
+
+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.