work

Author: lockwo

.github/workflows/deploy-docs.yml

@@ -26,8 +26,7 @@ jobs:
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
-          python -m pip install .
-          python -m pip install mkdocs mkdocs-material mkdocstrings[python] pymdown-extensions
+          python -m pip install ".[docs]"
       # https://github.com/mhausenblas/mkdocs-deploy-gh-pages/blob/master/action.sh
       - name: Build docs
         run: |

examples/03_ising_thermodynamics.ipynb

@@ -344,7 +344,7 @@
     "#         jax.random.choice(k_init, jnp.array([1, -1]), shape=(num_odd,)),\n",
     "#     ]\n",
     "\n",
-    "#     model_with_beta = IsingModel(weights=beta * edge_weights, \n",
+    "#     model_with_beta = IsingModel(weights=beta * edge_weights,\n",
     "#                           biases=beta * node_biases)\n",
     "\n",
     "#     samples = sample_fn(\n",

pyproject.toml

@@ -41,6 +41,16 @@ testing = [
     "ruff==0.11.0",
     "pyright==1.1.399",
 ]
+docs = [
+    "hippogriffe==0.2.0",
+    "griffe==1.7.3",
+    "mkdocs==1.6.1",
+    "mkdocs-include-exclude-files==0.1.0",
+    "mkdocs-ipynb==0.1.0",
+    "mkdocs-material==9.6.7",
+    "mkdocstrings[python]==0.28.3",
+    "pymdown-extensions==10.14.3",
+]
 
 [tool.setuptools]
 packages = { find = { include = ["isax", "isax.*"] } }

tests/test_sample.py

@@ -1,7 +1,15 @@
 import unittest
 import jax
 import jax.numpy as jnp
-from isax import BlockGraph, Edge, Node, IsingModel, IsingSampler, SamplingArgs, sample_chain
+from isax import (
+    BlockGraph,
+    Edge,
+    Node,
+    IsingModel,
+    IsingSampler,
+    SamplingArgs,
+    sample_chain,
+)
 
 
 class TestSampleChain(unittest.TestCase):
@@ -17,23 +25,29 @@ def test_readme_example(self):
                 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]
+        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)
+        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)
+        samples = sample_chain(
+            init_state, [sampler, sampler], model, sampling_args, key
+        )
 
         self.assertEqual(len(samples), 2)
         self.assertEqual(samples[0].shape, (100, len(even)))