diffract-core 0.2.1__py3-none-any.whl

diffract/viz/styling/theme/presets.py

@@ -0,0 +1,58 @@
+from __future__ import annotations
+
+from diffract.viz.styling.palettes import PaletteBundle
+from diffract.viz.styling.palettes.color import DarkColorPalette
+
+from .components import (
+    AxesStyle,
+    BackgroundStyle,
+    LayoutStyle,
+    LegendStyle,
+    TypographyStyle,
+)
+from .theme import Theme
+
+DEFAULT_THEME = Theme()
+
+
+DARK_THEME = Theme(
+    background=BackgroundStyle(
+        plot_bgcolor="#1e1e1e",
+        paper_bgcolor="#2d2d2d",
+    ),
+    axes=AxesStyle(
+        grid_color="#444444",
+        line_color="#666666",
+        mirror=False,
+    ),
+    legend=LegendStyle(
+        bgcolor="rgba(45,45,45,0.9)",
+        border_color="#666666",
+    ),
+    palettes=PaletteBundle(
+        color=DarkColorPalette(),
+    ),
+)
+
+
+MINIMAL_THEME = Theme(
+    layout=LayoutStyle(
+        margin={"l": 50, "r": 20, "t": 40, "b": 50},
+    ),
+    typography=TypographyStyle(
+        font_family="Arial",
+        title_font_size=14,
+        label_font_size=12,
+        tick_font_size=10,
+    ),
+    axes=AxesStyle(
+        show_grid=False,
+        show_line=False,
+        mirror=False,
+    ),
+    legend=LegendStyle(
+        bgcolor="rgba(255,255,255,0)",
+        border_color="rgba(0,0,0,0)",
+        border_width=0,
+    ),
+)

diffract/viz/styling/theme/theme.py

@@ -0,0 +1,36 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+from diffract.viz.styling.palettes import PaletteBundle
+
+from .components import (
+    AxesStyle,
+    BackgroundStyle,
+    ColorbarStyle,
+    LayoutStyle,
+    LegendStyle,
+    TypographyStyle,
+)
+
+
+@dataclass
+class Theme:
+    """Complete theme configuration with composition.
+
+    Theme provides figure-level styling and palette configuration.
+    Explicit user settings on Plot/Configurator fields take precedence.
+
+    Example:
+        >>> from diffract.viz.styling.theme import Theme, DARK_THEME
+        >>> plot = BoxPlot(x=FieldRef("model"), y=FieldRef("accuracy"))
+        >>> fig = plot.render(session, theme=DARK_THEME)
+    """
+
+    layout: LayoutStyle = field(default_factory=LayoutStyle)
+    typography: TypographyStyle = field(default_factory=TypographyStyle)
+    background: BackgroundStyle = field(default_factory=BackgroundStyle)
+    axes: AxesStyle = field(default_factory=AxesStyle)
+    legend: LegendStyle = field(default_factory=LegendStyle)
+    colorbar: ColorbarStyle = field(default_factory=ColorbarStyle)
+    palettes: PaletteBundle = field(default_factory=PaletteBundle)

diffract_core-0.2.1.dist-info/METADATA

@@ -0,0 +1,530 @@
+Metadata-Version: 2.4
+Name: diffract-core
+Version: 0.2.1
+Summary: Spectral analysis of neural network weights and their evolution over training
+Project-URL: Homepage, https://github.com/Risk-AI-Research/diffract
+Project-URL: Documentation, https://risk-ai-research.github.io/diffract/
+Project-URL: Repository, https://github.com/Risk-AI-Research/diffract
+Project-URL: Issues, https://github.com/Risk-AI-Research/diffract/issues
+Project-URL: Changelog, https://github.com/Risk-AI-Research/diffract/blob/main/CHANGELOG.md
+Project-URL: Paper, https://openreview.net/forum?id=XBUHoiAGDE
+Author: Risk AI Research
+Maintainer-email: Nikita Borodin <borodinik.s@gmail.com>
+License-Expression: Apache-2.0
+License-File: LICENSE
+License-File: NOTICE
+Keywords: deep-learning,llm,neural-networks,random-matrix-theory,spectral-analysis
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Classifier: Typing :: Typed
+Requires-Python: <3.13,>=3.12
+Requires-Dist: dependency-injector>=4.44.0
+Requires-Dist: h5py>=3.9.0
+Requires-Dist: numpy<3,>=1.26
+Requires-Dist: ordered-set>=4.1.0
+Requires-Dist: powerlaw<2,>=1.5.0
+Requires-Dist: rootutils<2,>=1.0.7
+Requires-Dist: scipy>=1.12.0
+Requires-Dist: tqdm>=4.65.0
+Requires-Dist: typing-extensions>=4.8.0
+Provides-Extra: all
+Requires-Dist: hydra-core>=1.3.0; extra == 'all'
+Requires-Dist: kaleido>=0.2.1; extra == 'all'
+Requires-Dist: omegaconf>=2.3.0; extra == 'all'
+Requires-Dist: pandas>=2.1.0; extra == 'all'
+Requires-Dist: plotly>=6.0.0; extra == 'all'
+Requires-Dist: polars>=0.20.0; extra == 'all'
+Requires-Dist: pyyaml>=6.0.1; extra == 'all'
+Requires-Dist: taichi>=1.7.3; extra == 'all'
+Requires-Dist: torch>=2.7.0; extra == 'all'
+Provides-Extra: common
+Requires-Dist: hydra-core>=1.3.0; extra == 'common'
+Requires-Dist: kaleido>=0.2.1; extra == 'common'
+Requires-Dist: omegaconf>=2.3.0; extra == 'common'
+Requires-Dist: pandas>=2.1.0; extra == 'common'
+Requires-Dist: plotly>=6.0.0; extra == 'common'
+Requires-Dist: polars>=0.20.0; extra == 'common'
+Requires-Dist: pyyaml>=6.0.1; extra == 'common'
+Provides-Extra: dev
+Requires-Dist: bandit[toml]>=1.7.5; extra == 'dev'
+Requires-Dist: fsspec>=2024.2.0; extra == 'dev'
+Requires-Dist: hydra-core>=1.3.0; extra == 'dev'
+Requires-Dist: interrogate>=1.7.0; extra == 'dev'
+Requires-Dist: kaleido>=0.2.1; extra == 'dev'
+Requires-Dist: mypy>=1.10.1; extra == 'dev'
+Requires-Dist: omegaconf>=2.3.0; extra == 'dev'
+Requires-Dist: pandas>=2.1.0; extra == 'dev'
+Requires-Dist: plotly>=6.0.0; extra == 'dev'
+Requires-Dist: polars>=0.20.0; extra == 'dev'
+Requires-Dist: pre-commit>=3.6.0; extra == 'dev'
+Requires-Dist: pytest-cov>=4.0.0; extra == 'dev'
+Requires-Dist: pytest-mock>=3.12.0; extra == 'dev'
+Requires-Dist: pytest-timeout>=2.4.0; extra == 'dev'
+Requires-Dist: pytest-watch>=4.2.0; extra == 'dev'
+Requires-Dist: pytest-xdist>=3.3.0; extra == 'dev'
+Requires-Dist: pytest>=7.4.0; extra == 'dev'
+Requires-Dist: pyyaml>=6.0.1; extra == 'dev'
+Requires-Dist: ruff>=0.9.0; extra == 'dev'
+Requires-Dist: sybil>=6.0.0; extra == 'dev'
+Requires-Dist: types-pyyaml>=6.0.0; extra == 'dev'
+Requires-Dist: types-requests>=2.31.0; extra == 'dev'
+Requires-Dist: zarr>=2.18.0; extra == 'dev'
+Provides-Extra: docs
+Requires-Dist: furo>=2024.8.0; extra == 'docs'
+Requires-Dist: myst-parser>=2.0.0; extra == 'docs'
+Requires-Dist: sphinx-autobuild>=2024.10.0; extra == 'docs'
+Requires-Dist: sphinx-autodoc-typehints>=1.25.0; extra == 'docs'
+Requires-Dist: sphinx-copybutton>=0.5.0; extra == 'docs'
+Requires-Dist: sphinx-design>=0.6.0; extra == 'docs'
+Requires-Dist: sphinx>=7.0.0; extra == 'docs'
+Requires-Dist: sphinxext-opengraph>=0.9.0; extra == 'docs'
+Provides-Extra: flax
+Requires-Dist: flax>=0.10.0; extra == 'flax'
+Provides-Extra: frameworks
+Requires-Dist: flax>=0.10.0; extra == 'frameworks'
+Requires-Dist: onnx>=1.16.0; extra == 'frameworks'
+Requires-Dist: tensorflow>=2.16.0; extra == 'frameworks'
+Provides-Extra: notebooks
+Requires-Dist: einops>=0.7.0; extra == 'notebooks'
+Requires-Dist: ipywidgets>=8.1.0; extra == 'notebooks'
+Requires-Dist: jupyter>=1.0.0; extra == 'notebooks'
+Requires-Dist: matplotlib>=3.8.0; extra == 'notebooks'
+Requires-Dist: safetensors>=0.4.0; extra == 'notebooks'
+Provides-Extra: onnx
+Requires-Dist: onnx>=1.16.0; extra == 'onnx'
+Provides-Extra: pandas
+Requires-Dist: pandas>=2.1.0; extra == 'pandas'
+Provides-Extra: polars
+Requires-Dist: polars>=0.20.0; extra == 'polars'
+Provides-Extra: redis
+Requires-Dist: redis>=5.1.0; extra == 'redis'
+Provides-Extra: taichi
+Requires-Dist: taichi>=1.7.3; extra == 'taichi'
+Provides-Extra: tensorflow
+Requires-Dist: tensorflow>=2.16.0; extra == 'tensorflow'
+Provides-Extra: torch
+Requires-Dist: torch>=2.7.0; extra == 'torch'
+Provides-Extra: viz
+Requires-Dist: hydra-core>=1.3.0; extra == 'viz'
+Requires-Dist: kaleido>=0.2.1; extra == 'viz'
+Requires-Dist: omegaconf>=2.3.0; extra == 'viz'
+Requires-Dist: plotly>=6.0.0; extra == 'viz'
+Requires-Dist: pyyaml>=6.0.1; extra == 'viz'
+Provides-Extra: zarr
+Requires-Dist: fsspec>=2024.2.0; extra == 'zarr'
+Requires-Dist: s3fs>=2024.2.0; extra == 'zarr'
+Requires-Dist: zarr>=3.0.0; extra == 'zarr'
+Description-Content-Type: text/markdown
+
+# Diffract: Deep Neural Network Weight Analysis Library
+
+[![CI](https://github.com/Risk-AI-Research/diffract/actions/workflows/ci.yml/badge.svg)](https://github.com/Risk-AI-Research/diffract/actions/workflows/ci.yml)
+[![PyPI](https://img.shields.io/pypi/v/diffract-core)](https://pypi.org/project/diffract-core/)
+[![Python](https://img.shields.io/pypi/pyversions/diffract-core)](https://pypi.org/project/diffract-core/)
+[![Docs](https://img.shields.io/badge/docs-github.io-blue)](https://risk-ai-research.github.io/diffract/)
+[![License: Apache 2.0](https://img.shields.io/badge/License-Apache_2.0-blue.svg)](https://opensource.org/licenses/Apache-2.0)
+
+Official library of ["Diffract: Spectral View of LLM Domain Adaptation"](https://openreview.net/forum?id=XBUHoiAGDE)
+(**ICML 2026 Oral**).
+
+Diffract is a Python package for analyzing deep neural network weights and tracking
+their evolution over the course of training.
+
+With a straightforward API and a functional design centered on reusable _kernels_,
+Diffract automatically resolves dependencies, builds computation graphs, and schedules
+calculations. Parameters and results are persisted across sessions.
+
+It accepts models from PyTorch, TensorFlow, Flax, and ONNX, as well as plain
+dictionaries of NumPy weight matrices.
+
+<br>
+
+## 🚀 Quick Start
+
+Diffract requires Python 3.12. The core package installs without any deep
+learning framework; heavy dependencies are opt-in extras:
+
+```bash
+pip install diffract-core                # core: extraction, spectral metrics, storage
+pip install "diffract-core[torch]"       # + PyTorch model loading (CUDA wheels on Linux, ~2-3 GB)
+pip install "diffract-core[viz]"         # + Plotly visualization and YAML plot configs
+pip install "diffract-core[taichi]"      # + accelerated heavy-tailed fits and p-value kernels
+pip install "diffract-core[all]"         # torch + viz + taichi + pandas/polars exports
+```
+
+Further extras: `frameworks` (TensorFlow, Flax, ONNX), `pandas` / `polars`
+(DataFrame exports), `zarr` (cloud storage), `redis` (shared cache),
+`notebooks` (tooling for the example notebooks). The quotes around
+`"diffract-core[...]"` matter in zsh.
+
+### Development Install
+
+```bash
+# Install uv if you haven't already
+curl -LsSf https://astral.sh/uv/install.sh | sh
+
+# Clone and install (uv provisions Python 3.12 automatically)
+git clone https://github.com/Risk-AI-Research/diffract.git
+cd diffract
+uv sync --extra dev --extra torch
+```
+
+Then use it in your code:
+
+```python
+from diffract import Session
+
+# Quick experiments (in-memory, no persistence)
+session = Session(profile="ram")
+
+# Persistent local storage (SQLite in .diffract/)
+session = Session(profile="local")
+
+with session:
+    session.models.add(torch_model, model_id="bert-base")
+    session.compute.apply("frob_norm", "stable_rank", "log_norm")
+    # per-parameter metrics
+    metrics_df = session.results.export_metrics(
+        "frob_norm", "stable_rank", export_format="pandas"
+    )
+    # model-level aggregates
+    aggregates_df = session.results.export_aggregates(
+        "log_norm", export_format="pandas"
+    )
+```
+
+Check out
+[example notebooks](https://github.com/Risk-AI-Research/diffract/tree/main/examples) and
+[plot configurations](https://github.com/Risk-AI-Research/diffract/tree/main/examples/configs)
+for more examples. The
+[notebooks](https://github.com/Risk-AI-Research/diffract/tree/main/notebooks) directory
+contains `compare_two_checkpoints.ipynb`, an end-to-end walkthrough.
+
+<br>
+
+## 🤔 Why Diffract?
+
+Neural networks often feel like black boxes. Diffract provides tools to analyze their
+internal structure:
+
+- **Training Insights**: Track how weights evolve across training epochs.
+- **Architecture Analysis**: Compare different model architectures objectively.
+- **Initialization Studies**: Evaluate the impact of initialization methods.
+- **Spectral Analysis**: Compute empirical spectral distributions, ranks, and norms.
+- **Heavy-Tailed Distributions**: Detect power-law and exponential tails in weight
+  spectra.
+
+<br>
+
+## 🔑 Key Features
+
+- **Session-based API**: Simple `models.add`, `compute.apply`, and
+  `results.export_metrics` workflow.
+
+- **Kernels**: Reusable functions that compute model characteristics—ranks, norms,
+  spectral properties—stored as named _fields_ on each parameter. Dependencies are
+  resolved automatically.
+
+- **Persistent Storage**: Parameters and results survive between sessions. Supports
+  HDF5, SQLite, Zarr, and hybrid backends.
+
+- **Kernel Apply Levels**: Kernels can work at multiple levels:
+  - **PARAMETER** - Operate on individual weight matrices.
+  - **IN_MODEL** - Aggregate within a single model.
+  - **CROSS_MODEL** - Compare or aggregate across models.
+
+- **Built-in Visualization**: Publication-ready Plotly plots with theming support.
+
+- **Export Formats**: Get results as `pandas`, `polars`, `dict`, `json`, or `list`.
+
+<br>
+
+## ✨ Core Functionality
+
+### Adding Models
+
+Add models from various frameworks to a session:
+
+<!-- skip: next -->
+
+```python
+from diffract import Session
+
+session = Session()
+
+with session:
+    session.models.add(torch_model)  # torch.nn.Module
+    session.models.add(torch_state_dict, model_id="checkpoint")  # Dict[str, torch.Tensor]
+    session.models.add(numpy_weights, model_id="raw-weights")  # Dict[str, np.ndarray]
+    session.models.add(onnx_model, model_id="onnx-model")  # onnx.ModelProto
+    session.models.add(flax_model, model_id="flax-model")  # flax.linen.Module
+    session.models.add(tf_model, model_id="tf-model")  # TensorFlow model
+```
+
+### Computing Metrics
+
+Dependencies are resolved automatically:
+
+```python
+session.compute.apply("frob_norm", "stable_rank")
+session.compute.apply("pl_ks")  # has many dependencies—all resolved automatically
+```
+
+### Filtering Parameters
+
+Filter computations by model, parameter type, or name:
+
+```python
+from diffract import ParameterOverrides, ParameterType, Session
+
+session = Session()
+
+# Assign custom types and names during extraction
+overrides = {
+    "model.layers.0.attn.q_proj.weight": ParameterOverrides(name="q", ptype="attn"),
+    "model.layers.0.attn.k_proj.weight": ParameterOverrides(name="k", ptype="attn"),
+    "model.layers.0.mlp.fc1.weight": ParameterOverrides(ptype="mlp"),
+}
+
+with session:
+    session.models.add(model, model_id="gpt", parameter_overrides=overrides)
+    session.compute.apply("frob_norm")
+
+# Scope work to a subset with session.filter(...)
+gpt = session.filter(model_ids=["gpt"])
+with gpt:
+    gpt.compute.apply("frob_norm")
+
+attn = session.filter(param_types=[ParameterType.from_string("attn")])
+with attn:
+    attn.results.export_metrics("frob_norm")
+```
+
+### Retrieving Results
+
+Export results in various formats (`pandas`, `polars`, `dict`, `json`, or `list`):
+
+```python
+scalars_df = session.results.export_metrics("stable_rank", export_format="pandas")
+aggregates_df = session.results.export_aggregates(
+    "stable_rank", export_format="pandas"
+)
+
+# Other formats work the same way
+results = session.results.export_metrics("stable_rank", export_format="polars")
+results = session.results.export_metrics("stable_rank", export_format="dict")
+results = session.results.export_metrics("stable_rank", export_format="json")
+results = session.results.export_metrics("stable_rank", export_format="list")
+```
+
+### Visualization
+
+Create publication-ready Plotly plots:
+
+<!-- skip: next -->
+
+```python
+from diffract.viz import DEFAULT_THEME
+
+# Ergonomic helpers on session.viz accept field names directly
+session.viz.box(y="stable_rank", x="model_id", theme=DEFAULT_THEME).show()
+session.viz.scatter(x="frob_norm", y="stable_rank").show()
+```
+
+### YAML-Driven Plotting
+
+Define complex visualizations via Hydra configs:
+
+<!-- skip: next -->
+
+```python
+session.viz.draw(config_path="examples/configs/boxplot_stable_rank.yaml").show()
+```
+
+### Kernel Configuration
+
+List and configure kernels at runtime:
+
+```python
+session.compute.list_available_kernels(verbose=True)
+session.compute.list_available_metrics(verbose=True)
+session.compute.configure_kernel("hard_rank", threshold=1e-6)
+```
+
+### Session Management
+
+Parameters and results stay available whenever the session is reopened. `Session()`
+defaults to the in-memory `ram` profile; pick a persistent profile such as `local` or
+`hybrid` to keep data across separate runs:
+
+```python
+from diffract import Session
+
+session = Session()  # in-memory; use profile="local" to persist across runs
+
+# Add models and compute
+with session:
+    session.models.add(model, model_id="my-model")
+    session.compute.apply("frob_norm")
+
+# Reopen the session: earlier results are still there
+with session:
+    results = session.results.export_metrics("frob_norm", export_format="pandas")
+    session.models.list()
+    session.models.erase("my-model")
+```
+
+### Custom Kernels
+
+Implement your own research metrics using the session kernel decorator:
+
+```python
+from diffract import Session
+
+session = Session()
+
+with session:
+    # Define and register a custom kernel
+    @session.compute.kernel()
+    def my_custom_metric(frob_norm: float, *, scaling_factor: float = 1.0) -> float:
+        """Custom metric that scales the Frobenius norm."""
+        return frob_norm * scaling_factor
+
+    session.models.add(my_model)
+    session.compute.configure_kernel("my_custom_metric", scaling_factor=2.0)
+    session.compute.apply("my_custom_metric")
+```
+
+You can also override the registered name and output fields:
+
+```python
+with session:
+    @session.compute.kernel(name="scaled_metric", produce_fields=["scaled_result"])
+    def custom_analysis(frob_norm: float, stable_rank: float, *, weight: float = 0.5) -> float:
+        """Custom analysis combining multiple metrics."""
+        return weight * frob_norm + (1 - weight) * stable_rank
+```
+
+### Available Kernels
+
+Diffract includes kernels for norms, ranks, spectral analysis, heavy-tailed fits, and
+more. Run `session.compute.list_available_kernels(verbose=True)` to list them all.
+
+### Merging Sessions
+
+Merge parameters and results from another session:
+
+<!-- skip: next -->
+
+```python
+from diffract import Session
+
+session1 = Session(config_path="config1.ini")
+session2 = Session(config_path="config2.ini")
+
+with session1:
+    session1.models.add(model1, model_id="model-a")
+    session1.compute.apply("frob_norm")
+
+with session2:
+    session2.models.add(model2, model_id="model-b")
+    session2.utils.merge_other_session(session1, fields=["frob_norm"])
+```
+
+### Configuration
+
+Diffract offers built-in **profiles** for common setups:
+
+| Profile  | Storage       | Cache      | Use case                          |
+| -------- | ------------- | ---------- | --------------------------------- |
+| `ram`    | RAM           | None       | Quick experiments, no persistence |
+| `local`  | SQLite        | Simple LRU | Local development, persistent     |
+| `hybrid` | SQLite + HDF5 | Simple LRU | Large models, optimized arrays    |
+
+```python
+from diffract import Session
+
+# Use a profile (recommended for most users)
+session = Session(profile="ram")      # fast, temporary
+session = Session(profile="local")    # persistent, simple
+session = Session(profile="hybrid")   # persistent, optimized for large arrays
+
+# Or use a custom config file for full control
+session = Session(config_path="my_config.ini")
+```
+
+**Tip**: Start with a profile, then switch to a config file when you need
+reproducibility or custom settings.
+
+#### Advanced Configuration
+
+For production or reproducible experiments, use INI config files. See
+`src/diffract/configs/` for examples:
+
+```ini
+[storage]
+backend = "sqlite"
+
+[storage.sqlite]
+path = "data/diffract.db"
+
+[cache]
+backend = "simple"
+
+[parallel.thread_pool]
+max_workers = 4
+```
+
+#### Storage Backends
+
+- **RAM**: In-memory (no persistence)
+- **SQLite**: Lightweight database for metadata and arrays
+- **HDF5**: Optimized for large numerical arrays with compression
+- **Zarr**: Cloud-optimized array storage for large-scale data
+- **Hybrid**: SQLite (metadata) + HDF5/Zarr (arrays)
+
+#### Cache Backends
+
+- **Simple**: In-memory LRU cache
+- **Redis**: Distributed caching (requires `redis` extra)
+- **None**: Disable caching
+
+<br>
+
+## 📚 Documentation
+
+The documentation site is sourced from `docs/` and built with Sphinx + MyST. Install the
+tooling via `uv sync --extra docs` and run `make docs` to render the HTML locally.
+
+<br>
+
+## 📝 Citation
+
+If you use Diffract or build on the paper, please cite:
+
+```bibtex
+@inproceedings{borodin2026diffract,
+  title     = {Diffract: Spectral View of {LLM} Domain Adaptation},
+  author    = {Nikita Borodin and Maria Krylova and Artem Zabolotnyi and Dmitry Aspisov and Egor Shikov and Nikita Tyuplyaev and Oleg Travkin and Roman Alferov and Dmitry Vinichenko},
+  booktitle = {Forty-third International Conference on Machine Learning},
+  year      = {2026},
+  url       = {https://openreview.net/forum?id=XBUHoiAGDE}
+}
+```
+
+<br>
+
+## ❤️ Contributions
+
+Contributions are welcome! Fork the repo, create a feature branch, and submit a PR. Use
+`make lint` and `make test` to validate your changes.
+
+<br>
+
+## 📄 License
+
+Licensed under the Apache License 2.0 — see [LICENSE](LICENSE). Copyright 2026 Risk AI
+Research.