work

Author: lockwo

.gitignore

@@ -205,3 +205,5 @@ cython_debug/
 marimo/_static/
 marimo/_lsp/
 __marimo__/
+
+.DS_Store

README.md

@@ -1,2 +1,10 @@
 # isax
+
 Ising (like) models in JAX
+
+todo:
+- annealing
+- cleaner interface
+- generic block typing
+- generalize pytree typing for states
+- support non-gibbs samplers (wolff, mh, etc.)

examples/01_ising_model_physics.ipynb

@@ -7,7 +7,7 @@
    "source": [
     "# Flips Per Second\n",
     "\n",
-    "In this example we will be looking as Flips Per Second (FPS) as a metric for evaluating the speed of a Gibbs sampling program. This highlights the power of GPUs for workloads which have many parallel Gibbs chains to be sampled."
+    "In this example we will be looking as Flips Per Second (FPS) as a metric for evaluating the speed of a Gibbs sampling program."
    ]
   },
   {
@@ -19,12 +19,20 @@
    "source": [
     "import time\n",
     "\n",
+    "import dwave_networkx\n",
     "import jax\n",
     "import jax.numpy as jnp\n",
-    "import dwave_networkx\n",
+    "import matplotlib.pyplot as plt\n",
     "import networkx as nx\n",
-    "from isax.block import BlockGraph, Edge, AbstractNode\n",
-    "from isax.sample import IsingModel, IsingSampler, SamplingArgs, sample_chain"
+    "from isax import (\n",
+    "    BlockGraph,\n",
+    "    Edge,\n",
+    "    IsingModel,\n",
+    "    IsingSampler,\n",
+    "    Node,\n",
+    "    sample_chain,\n",
+    "    SamplingArgs,\n",
+    ")"
    ]
   },
   {
@@ -35,39 +43,28 @@
    "outputs": [],
    "source": [
     "def create_dwave_pegasus_graph(pegasus_size, key):\n",
-    "    \"\"\"Create DWave Pegasus graph like thrml does\"\"\"\n",
-    "    # Create the DWave Pegasus graph\n",
     "    graph = dwave_networkx.pegasus_graph(pegasus_size)\n",
-    "    \n",
-    "    # Convert coordinates to AbstractNode objects (isax uses AbstractNode, not SpinNode)\n",
-    "    coord_to_node = {coord: AbstractNode() for coord in graph.nodes}\n",
+    "    coord_to_node = {coord: Node() for coord in graph.nodes}\n",
     "    nx.relabel_nodes(graph, coord_to_node, copy=False)\n",
-    "    \n",
-    "    # Convert to our format\n",
     "    nodes = list(graph.nodes)\n",
     "    edges = [Edge(u, v) for u, v in graph.edges()]\n",
-    "    \n",
-    "    # Use graph coloring to create blocks (similar to thrml approach)\n",
+    "    # todo: dwave ships colorings already?\n",
     "    coloring = nx.coloring.greedy_color(graph, strategy=\"DSATUR\")\n",
     "    n_colors = max(coloring.values()) + 1\n",
     "    blocks = [[] for _ in range(n_colors)]\n",
-    "    \n",
-    "    # Form color groups\n",
     "    for node in graph.nodes:\n",
     "        blocks[coloring[node]].append(node)\n",
-    "    \n",
-    "    # Create BlockGraph with corrected argument order\n",
     "    block_graph = BlockGraph(blocks, edges)\n",
-    "    \n",
-    "    # Generate random parameters\n",
     "    key1, key2 = jax.random.split(key)\n",
     "    biases = jax.random.uniform(key1, (len(nodes),), minval=-0.1, maxval=0.1)\n",
     "    weights = jax.random.uniform(key2, (len(edges),), minval=-0.1, maxval=0.1)\n",
-    "    \n",
+    "\n",
     "    model = IsingModel(weights=weights, biases=biases)\n",
-    "    \n",
-    "    print(f\"Created Pegasus graph with {len(nodes)} nodes, {len(edges)} edges, {n_colors} color blocks\")\n",
-    "    \n",
+    "\n",
+    "    print(\n",
+    "        f\"Created Pegasus graph with {len(nodes)} nodes, {len(edges)} edges, {n_colors} color blocks\"\n",
+    "    )\n",
+    "\n",
     "    return model, block_graph, nodes, blocks"
    ]
   },
@@ -82,42 +79,44 @@
     "    model, block_graph, nodes, blocks, chain_len, batch_size, n_reps, device\n",
     "):\n",
     "    key = jax.random.key(42)\n",
-    "    \n",
+    "\n",
     "    (adjs, masks, edge_infos), eqx_graph = block_graph.get_sampling_params()\n",
-    "    \n",
-    "    # Create samplers for each block\n",
+    "\n",
     "    samplers = [IsingSampler() for _ in range(len(blocks))]\n",
-    "    \n",
+    "\n",
     "    sampling_args = SamplingArgs(\n",
     "        gibbs_steps=chain_len,\n",
     "        blocks_to_sample=list(range(len(blocks))),\n",
     "        adjs=adjs,\n",
     "        masks=masks,\n",
     "        edge_info=edge_infos,\n",
-    "        eqx_graph=eqx_graph\n",
+    "        eqx_graph=eqx_graph,\n",
     "    )\n",
-    "    \n",
+    "\n",
     "    def sample_batch(key):\n",
     "        keys = jax.random.split(key, batch_size)\n",
-    "        \n",
+    "\n",
     "        def sample_single(single_key):\n",
     "            k_init, k_run = jax.random.split(single_key)\n",
-    "            \n",
-    "            # Initialize each block with random spins\n",
+    "\n",
     "            init_state = []\n",
     "            for block in blocks:\n",
-    "                block_state = jax.random.bernoulli(k_init, 0.5, (len(block),)).astype(jnp.int32) * 2 - 1\n",
+    "                block_state = (\n",
+    "                    jax.random.bernoulli(k_init, 0.5, (len(block),)).astype(jnp.int32)\n",
+    "                    * 2\n",
+    "                    - 1\n",
+    "                )\n",
     "                init_state.append(block_state)\n",
-    "            \n",
+    "\n",
     "            samples = sample_chain(init_state, samplers, model, sampling_args, k_run)\n",
     "            return samples\n",
-    "        \n",
+    "\n",
     "        return jax.vmap(sample_single)(keys)\n",
-    "    \n",
+    "\n",
     "    jit_sample_batch = jax.jit(sample_batch, device=device)\n",
-    "    \n",
+    "\n",
     "    keys = jax.random.split(key, n_reps)\n",
-    "    \n",
+    "\n",
     "    start_time = time.time()\n",
     "    _ = jax.block_until_ready(jit_sample_batch(keys[0]))\n",
     "    time_with_compile = time.time() - start_time\n",
@@ -132,9 +131,7 @@
     "    thruput = chain_len * batch_size * len(nodes)\n",
     "    flips_per_ns = thruput / (time_without_compile * 1e9)\n",
     "\n",
-    "    print(\n",
-    "        f\"chain_len: {chain_len}, batch_size: {batch_size}\"\n",
-    "    )\n",
+    "    print(f\"chain_len: {chain_len}, batch_size: {batch_size}\")\n",
     "    print(\n",
     "        f\"Time with compile: {time_with_compile:.4f}s, \"\n",
     "        f\"time without compile: {time_without_compile:.4f}s, \"\n",
@@ -241,21 +238,10 @@
     }
    ],
    "source": [
-    "import matplotlib.pyplot as plt\n",
-    "\n",
     "plt.figure(figsize=(6, 6))\n",
-    "plt.plot(batch_sizes, flips_per_ns_cpu, label=\"isax CPU\", marker=\"s\")\n",
     "if has_gpu:\n",
     "    plt.plot(batch_sizes, flips_per_ns_gpu, label=\"isax GPU\", marker=\"o\")\n",
-    "# plt.axhline(\n",
-    "#     143.80, linestyle=\"--\", color=\"red\", label=\"FPGA\"\n",
-    "# )\n",
-    "# plt.axhline(\n",
-    "#     12.88, linestyle=\"--\", color=\"green\", label=\"Single TPU\"\n",
-    "# )\n",
-    "# plt.axhline(\n",
-    "#     60, linestyle=\"--\", color=\"orange\", label=\"FPGA Hardware (thrml reference)\"\n",
-    "# )\n",
+    "plt.plot(batch_sizes, flips_per_ns_cpu, label=\"isax CPU\", marker=\"s\")\n",
     "plt.legend()\n",
     "plt.xlabel(\"Batch size\")\n",
     "plt.xscale(\"log\")\n",