work

lockwo · lockwo · commit 55888e6e9d8f · 2025-11-25T16:37:52.000-07:00
diff --git a/.github/FUNDING.yml b/.github/FUNDING.yml
@@ -0,0 +1 @@
+github: [lockwo]
diff --git a/.github/workflows/style-check.yml b/.github/workflows/style-check.yml
@@ -32,8 +32,7 @@ jobs:
           python-version: ${{ matrix.python-version }}
       - name: Run styles & lint check
         run: |
-          pip install -e .
-          pip install ruff pyright
+          pip install -e .[testing]
           ruff format --check isax/ tests/ examples/
           ruff check isax/ tests/ examples/
           pyright isax/
diff --git a/README.md b/README.md
@@ -12,6 +12,40 @@ pip install -e .
 
 Requires Python 3.10+.
 
+## Quick Example
+
+```python
+import jax
+import jax.numpy as jnp
+from isax import BlockGraph, Edge, Node, IsingModel, IsingSampler, SamplingArgs, sample_chain
+
+L = 4
+nodes = [Node() for _ in range(L * L)]
+
+edges = []
+for x in range(L):
+    for y in range(L):
+        i = x * L + y
+        edges.append(Edge(nodes[i], nodes[(x * L + (y + 1) % L)]))
+        edges.append(Edge(nodes[i], nodes[((x + 1) % L) * L + y]))
+
+even = [nodes[x * L + y] for x in range(L) for y in range(L) if (x + y) % 2 == 0]
+odd = [nodes[x * L + y] for x in range(L) for y in range(L) if (x + y) % 2 == 1]
+
+graph = BlockGraph([even, odd], edges)
+params = graph.get_sampling_params()
+
+model = IsingModel(weights=jnp.ones(len(edges)), biases=jnp.zeros(L * L))
+sampler = IsingSampler()
+sampling_args = SamplingArgs(gibbs_steps=100, blocks_to_sample=[0, 1], data=params)
+
+key = jax.random.key(0)
+init_state = [jax.random.choice(key, jnp.array([-1, 1]), (len(even),)),
+              jax.random.choice(key, jnp.array([-1, 1]), (len(odd),))]
+
+samples = sample_chain(init_state, [sampler, sampler], model, sampling_args, key)
+```
+
 ## Documentation
 
 Available at https://lockwo.github.io/isax
@@ -22,6 +56,7 @@ Available at https://lockwo.github.io/isax
 - [x] cleaner interface
 - [ ] improve example documentation/math background
 - [ ] add tests
+- [ ] runtime sampling params
 - [ ] generic block typing
 - [ ] generalize pytree typing for states
 - [ ] support non-gibbs samplers (wolff, mh, etc.)
diff --git a/docs/api/block.md b/docs/api/block.md
@@ -1,36 +1,22 @@
-# Block Graph Module
+# Block Graph Tools
 
 ::: isax.block.Node
     options:
-        show_source: false
-        show_root_heading: true
-        members:
-            - __init__
-            - __lt__
+        members: false
 
 ::: isax.block.Edge
     options:
-        show_source: false
-        show_root_heading: true
         members:
             - __init__
 
 ::: isax.block.EqxGraph
     options:
-        show_source: false
-        show_root_heading: true
-        members:
-            - node_to_global
-            - node_to_local
-            - block_to_global
+        members: false
+
 
 ::: isax.block.BlockGraph
     options:
-        show_source: false
-        show_root_heading: true
         members:
             - __init__
-            - get_edge_info
+            - get_edge_structure
             - get_sampling_params
-            - node_states_to_block_states
-            - block_states_to_node_states
diff --git a/docs/api/metrics.md b/docs/api/metrics.md
@@ -1,7 +1,5 @@
-# Metrics Module
+# Metrics
 
 Functions for computing physical observables and metrics from sampled states.
 
-::: isax.metrics.magnetization_per_site
-    options:
-        show_source: false
+::: isax.metrics.magnetization_per_site
diff --git a/docs/api/sample.md b/docs/api/sample.md
@@ -1,40 +1,27 @@
-# Sampling Module
-
-This module contains the core sampling functionality including models, samplers, and utilities.
+# Sampling
 
 ## Models
 
 ::: isax.sample.IsingModel
     options:
-        show_source: false
-        show_root_heading: true
         members:
-            - __init__
             - energy
             - to_sample_params
 
 ## Samplers
 
 ::: isax.sample.AbstractSampler
     options:
-        show_source: false
-        show_root_heading: true
-        members:
-            - sample
-            - initialize_state
+        members: false
 
 ::: isax.sample.IsingSampler
     options:
-        show_source: false
-        show_root_heading: true
         members:
             - sample
             - initialize_state
 
 ::: isax.sample.AnnealedIsingSampler
     options:
-        show_source: false
-        show_root_heading: true
         members:
             - __init__
             - sample
@@ -44,19 +31,8 @@ This module contains the core sampling functionality including models, samplers,
 
 ::: isax.sample.SamplingArgs
     options:
-        show_source: false
-        show_root_heading: true
         members:
             - __init__
 
 ::: isax.sample.sample_chain
-    options:
-        show_source: false
-
-::: isax.sample.sample_blocks
-    options:
-        show_source: false
 
-::: isax.sample.concat_state
-    options:
-        show_source: false
diff --git a/docs/index.md b/docs/index.md
@@ -1,82 +1,16 @@
-# ISAX
+# Getting Started
 
-Block-parallel sampling for Ising models on arbitrary graphs using JAX.
-
-ISAX is a [JAX](https://github.com/google/jax)-based library for sampling from Ising models using block-parallel Gibbs sampling. It supports hypergraphs, custom sampling schedules, and JAX transformations.
-
-## Quick Example
-
-```python
-import jax
-import jax.numpy as jnp
-from isax import IsingModel, IsingSampler, BlockGraph, sample_chain, SamplingArgs
-
-# Create a 2D lattice Ising model
-n_nodes = 100
-edges = [(i, j) for i in range(n_nodes) for j in range(i+1, n_nodes)
-         if abs(i-j) == 1 or abs(i-j) == 10]  # 2D grid connectivity
-
-# Initialize model with random weights
-key = jax.random.PRNGKey(0)
-weight_key, sample_key = jax.random.split(key)
-weights = jax.random.normal(weight_key, (len(edges),))
-biases = jnp.zeros(n_nodes)
-
-model = IsingModel(weights=weights, biases=biases)
-
-# Create block structure for parallel sampling
-graph = BlockGraph(n_nodes, edges, n_blocks=4)
-
-# Sample using Gibbs sampler
-sampler = IsingSampler()
-initial_state = jax.random.choice(sample_key, jnp.array([-1, 1]), (n_nodes,))
-
-# Run sampling
-samples = sample_chain(
-    initial_state,
-    [sampler] * graph.n_blocks,
-    model,
-    SamplingArgs(gibbs_steps=1000, blocks_to_sample=list(range(4)), data=graph.get_sampling_params()),
-    sample_key
-)
-```
+isax is a [JAX](https://github.com/google/jax)-based library for sampling from Ising models using blocked Gibbs sampling. It supports hypergraphs, flexible sampling/modeling, and all the usual JAX transformations. isax is heavily inspired by [thrml](https://docs.thrml.ai/en/latest/).
 
 ## Installation
 
-```bash
-pip install isax
-```
-
-Or install from source:
-
 ```bash
 git clone https://github.com/lockwo/isax
 cd isax
 pip install -e .
 ```
 
-## Features
-
-Traditional MCMC sampling for Ising models can be slow for large graphs. ISAX accelerates this through:
-
-1. **Block decomposition**: Divide the graph into blocks that can be updated in parallel
-2. **JAX compilation**: JIT-compile the sampling loops for maximum performance
-3. **Vectorization**: Process multiple chains simultaneously with `vmap`
-4. **Hardware acceleration**: Automatic GPU/TPU support through JAX
-
-## Navigation
-
-- [API Reference](api/block.md) - Complete API documentation
-- [Examples](examples/01_ising_model_physics.ipynb) - Jupyter notebooks with practical examples
 
 ## Citation
 
 If you use ISAX in your research, please cite:
-
-```bibtex
-@software{isax2025,
-  title = {ISAX: Block-parallel sampling for Ising models},
-  year = {2025},
-  url = {https://github.com/lockwo/isax}
-}
-```
diff --git a/examples/03_ising_thermodynamics.ipynb b/examples/03_ising_thermodynamics.ipynb
diff --git a/examples/04_paoa_optimization.ipynb b/examples/04_paoa_optimization.ipynb
@@ -6,7 +6,7 @@
    "source": [
     "# Probabilistic Approximate Optimization Algorithm\n",
     "\n",
-    "In this example, we recreate Figure 3(c) of https://arxiv.org/abs/2507.07420."
+    "In this example, we recreate Figure 3(c) of [https://arxiv.org/abs/2507.07420](https://arxiv.org/abs/2507.07420)."
    ]
   },
   {
diff --git a/mkdocs.yml b/mkdocs.yml
@@ -8,8 +8,6 @@ theme:
         - content.code.copy
     palette:
         # Light mode / dark mode
-        # We deliberately don't automatically use `media` to check a user's preferences. We default to light mode as
-        # (a) it looks more professional, and (b) is more obvious about the fact that it offers a (dark mode) toggle.
         - scheme: default
           primary: white
           accent: amber
@@ -26,9 +24,9 @@ theme:
         repo: fontawesome/brands/github  # GitHub logo in top right
         logo: "material/graph-outline"
 
-site_name: ISAX
+site_name: isax
 site_description: Block-parallel sampling for Ising models on arbitrary graphs using JAX.
-site_author: ISAX Contributors
+site_author: lockwo
 site_url: https://github.com/lockwo/isax
 
 repo_url: https://github.com/lockwo/isax 
diff --git a/tests/__init__.py b/tests/__init__.py
diff --git a/tests/test_metrics.py b/tests/test_metrics.py
@@ -0,0 +1,30 @@
+import unittest
+import jax.numpy as jnp
+from isax.metrics import magnetization_per_site
+
+
+class TestMagnetizationPerSite(unittest.TestCase):
+    def test_all_up_spins(self):
+        """All +1 spins should give magnetization of 1."""
+        state = [jnp.ones((10, 5), dtype=jnp.int32)]
+        result = magnetization_per_site(state)
+        self.assertEqual(len(result), 1)
+        self.assertTrue(jnp.allclose(result[0], jnp.ones(5)))
+
+    def test_all_down_spins(self):
+        """All -1 spins should give magnetization of -1."""
+        state = [jnp.full((10, 5), -1, dtype=jnp.int32)]
+        result = magnetization_per_site(state)
+        self.assertTrue(jnp.allclose(result[0], jnp.full(5, -1.0)))
+
+    def test_balanced_spins(self):
+        """Half +1 and half -1 should give magnetization of 0."""
+        up = jnp.ones((5, 4), dtype=jnp.int32)
+        down = jnp.full((5, 4), -1, dtype=jnp.int32)
+        state = [jnp.concatenate([up, down], axis=0)]
+        result = magnetization_per_site(state)
+        self.assertTrue(jnp.allclose(result[0], jnp.zeros(4)))
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/test_sample.py b/tests/test_sample.py
@@ -0,0 +1,46 @@
+import unittest
+import jax
+import jax.numpy as jnp
+from isax import BlockGraph, Edge, Node, IsingModel, IsingSampler, SamplingArgs, sample_chain
+
+
+class TestSampleChain(unittest.TestCase):
+    def test_readme_example(self):
+        """Run the README example to ensure it works."""
+        L = 4
+        nodes = [Node() for _ in range(L * L)]
+
+        edges = []
+        for x in range(L):
+            for y in range(L):
+                i = x * L + y
+                edges.append(Edge(nodes[i], nodes[(x * L + (y + 1) % L)]))
+                edges.append(Edge(nodes[i], nodes[((x + 1) % L) * L + y]))
+
+        even = [nodes[x * L + y] for x in range(L) for y in range(L) if (x + y) % 2 == 0]
+        odd = [nodes[x * L + y] for x in range(L) for y in range(L) if (x + y) % 2 == 1]
+
+        graph = BlockGraph([even, odd], edges)
+        params = graph.get_sampling_params()
+
+        model = IsingModel(weights=jnp.ones(len(edges)), biases=jnp.zeros(L * L))
+        sampler = IsingSampler()
+        sampling_args = SamplingArgs(gibbs_steps=100, blocks_to_sample=[0, 1], data=params)
+
+        key = jax.random.key(0)
+        init_state = [
+            jax.random.choice(key, jnp.array([-1, 1]), (len(even),)),
+            jax.random.choice(key, jnp.array([-1, 1]), (len(odd),)),
+        ]
+
+        samples = sample_chain(init_state, [sampler, sampler], model, sampling_args, key)
+
+        self.assertEqual(len(samples), 2)
+        self.assertEqual(samples[0].shape, (100, len(even)))
+        self.assertEqual(samples[1].shape, (100, len(odd)))
+        self.assertTrue(jnp.all((samples[0] == 1) | (samples[0] == -1)))
+        self.assertTrue(jnp.all((samples[1] == 1) | (samples[1] == -1)))
+
+
+if __name__ == "__main__":
+    unittest.main()

Original file line number	Diff line number	Diff line change
`@@ -6,7 +6,7 @@`
`6`	`6`	`"source": [`
`7`	`7`	`"# Probabilistic Approximate Optimization Algorithm\n",`
`8`	`8`	`"\n",`
`9`		`- "In this example, we recreate Figure 3(c) of https://arxiv.org/abs/2507.07420."`
	`9`	`+ "In this example, we recreate Figure 3(c) of [https://arxiv.org/abs/2507.07420](https://arxiv.org/abs/2507.07420)."`
`10`	`10`	`]`
`11`	`11`	`},`
`12`	`12`	`{`