fix: add 2pi² to z-state diffusion computation

.github/workflows/simulation_test.yml

@@ -75,4 +75,101 @@ jobs:
 
       - name: Run Simulation Test
         run: |
-          python tests/simulation_test/test_simulation.py
+          python tests/simulation_test/test_simulation.py
+
+      - name: Plot comparison
+        if: always()
+        id: plot
+        env:
+          NOTEBOOK: ${{ matrix.notebook }}
+        run: |
+          python tests/simulation_test/plot_comparison.py "$NOTEBOOK" sim_report
+          echo "name=$(basename "$NOTEBOOK" .ipynb)" >> $GITHUB_OUTPUT
+
+      - name: Upload comparison plot
+        if: always() && steps.plot.outcome == 'success'
+        uses: actions/upload-artifact@v4
+        with:
+          name: sim-report-${{ steps.plot.outputs.name }}
+          path: sim_report/
+
+      - name: Upload raw simulation data
+        if: always() && steps.plot.outcome == 'success'
+        uses: actions/upload-artifact@v4
+        with:
+          name: sim-data-${{ steps.plot.outputs.name }}
+          path: |
+            tests/simulation_test/ref_files/
+            tests/simulation_test/actual_files/
+            tests/simulation_test/seq_files/
+          if-no-files-found: ignore
+          retention-days: 14
+
+  report:
+    needs: test-notebooks
+    # PRs from forks get a read-only token, they only get the artifacts
+    if: always() && github.event.pull_request.head.repo.full_name == github.repository
+    runs-on: ubuntu-latest
+    permissions:
+      contents: write        # push images to the sim-test-images branch
+      pull-requests: write   # post the comment
+    steps:
+      - uses: actions/checkout@v4
+
+      - uses: actions/download-artifact@v4
+        with:
+          pattern: sim-report-*
+          path: sim_report
+          merge-multiple: true
+
+      - name: Push images to sim-test-images branch
+        id: push
+        env:
+          GH_TOKEN: ${{ github.token }}
+          DEST: pr-${{ github.event.pull_request.number }}/run-${{ github.run_id }}-${{ github.run_attempt }}
+        run: |
+          REPORT="$PWD/sim_report"
+          cd "$(mktemp -d)"
+          git init -q
+          git config user.name "github-actions[bot]"
+          git config user.email "41898282+github-actions[bot]@users.noreply.github.com"
+          git remote add origin "https://x-access-token:${GH_TOKEN}@github.com/${{ github.repository }}.git"
+          if git fetch -q --depth=1 origin sim-test-images; then
+            git checkout -q -b sim-test-images FETCH_HEAD
+          else
+            git checkout -q --orphan sim-test-images
+          fi
+          mkdir -p "$DEST"
+          cp "$REPORT"/*.png "$DEST"/ 2>/dev/null || true
+          git add -A
+          git commit -q --allow-empty -m "Simulation test images for $DEST"
+          # Retry in case another PR pushed concurrently
+          for i in 1 2 3; do
+            git push -q origin sim-test-images && break
+            git pull -q --rebase origin sim-test-images
+          done
+          echo "sha=$(git rev-parse HEAD)" >> $GITHUB_OUTPUT
+          echo "dest=$DEST" >> $GITHUB_OUTPUT
+
+      - name: Build report
+        run: |
+          python tests/simulation_test/build_report.py sim_report \
+            "https://raw.githubusercontent.com/${{ github.repository }}/${{ steps.push.outputs.sha }}/${{ steps.push.outputs.dest }}" \
+            "${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}" \
+            > comment.md
+
+      - name: Post or update PR comment
+        uses: actions/github-script@v7
+        with:
+          script: |
+            const body = require('fs').readFileSync('comment.md', 'utf8');
+            const marker = '<!-- simulation-test-report -->';
+            const { owner, repo } = context.repo;
+            const issue_number = context.payload.pull_request.number;
+            const comments = await github.paginate(github.rest.issues.listComments, { owner, repo, issue_number });
+            const existing = comments.find(c => c.body.includes(marker));
+            if (existing) {
+              await github.rest.issues.updateComment({ owner, repo, comment_id: existing.id, body });
+            } else {
+              await github.rest.issues.createComment({ owner, repo, issue_number, body });
+            }

CHANGELOG.md

@@ -1,3 +1,5 @@
+- 1.1.2
+  - fix: diffusion for z states was missing a (2pi)² in the b-value computation
 - 1.1.1
   - `VoxelGridPhantom.build()`: voxel positions are always centered around the origin (index offset `- n // 2`); the affine translation is no longer applied, only rotation and scaling
 - 1.1.0

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "maturin"
 
 [project]
 name = "mrzerocore"
-version = "1.1.1"
+version = "1.1.2"
 description = "Core functionality of MRzero"
 authors = [
     {name = "Jonathan Endres", email = "jonathan.endres@uk-erlangen.de"},

python/MRzeroCore/simulation/main_pass.py

@@ -266,7 +266,7 @@ def calc_mag(ancestor: tuple) -> torch.Tensor:
                     dist.mag = dist.mag * torch.exp(2j * np.pi * motion_phase[-1, :])
                 dist.kt_vec = dist_traj[-1]
             else:  # z or z0
-                k = torch.linalg.vector_norm(dist.kt_vec[:3])
+                k = 2 * torch.pi * torch.linalg.vector_norm(dist.kt_vec[:3])
                 diffusion = torch.exp(-1e-9 * data.D * total_time * k**2)
                 dist.mag = dist.mag * r1 * diffusion
             if dist.dist_type == 'z0':

src/pre_pass.rs

@@ -93,10 +93,22 @@ pub fn comp_graph(
     // Make precision of state merging dependent on the units used in the seq.
     // For this, we use a fraction of the smallest used step.
     let min_kt_step = [
-        seq.iter().flat_map(|r| r.gradm_event.iter().map(|g| g[0].abs() as f64)).filter(|x| *x > 1e-3).min_by(f64::total_cmp),
-        seq.iter().flat_map(|r| r.gradm_event.iter().map(|g| g[1].abs() as f64)).filter(|x| *x > 1e-3).min_by(f64::total_cmp),
-        seq.iter().flat_map(|r| r.gradm_event.iter().map(|g| g[2].abs() as f64)).filter(|x| *x > 1e-3).min_by(f64::total_cmp),
-        seq.iter().flat_map(|r| r.event_time.iter().map(|t| *t as f64)).filter(|x| *x > 1e-6).min_by(f64::total_cmp)
+        seq.iter()
+            .flat_map(|r| r.gradm_event.iter().map(|g| g[0].abs() as f64))
+            .filter(|x| *x > 1e-3)
+            .min_by(f64::total_cmp),
+        seq.iter()
+            .flat_map(|r| r.gradm_event.iter().map(|g| g[1].abs() as f64))
+            .filter(|x| *x > 1e-3)
+            .min_by(f64::total_cmp),
+        seq.iter()
+            .flat_map(|r| r.gradm_event.iter().map(|g| g[2].abs() as f64))
+            .filter(|x| *x > 1e-3)
+            .min_by(f64::total_cmp),
+        seq.iter()
+            .flat_map(|r| r.event_time.iter().map(|t| *t as f64))
+            .filter(|x| *x > 1e-6)
+            .min_by(f64::total_cmp),
     ];
     let inv_kt_grid = [
         1.0 / (0.1 * min_kt_step[0].unwrap_or(1.0)).clamp(1e-6, 1.0),
@@ -126,7 +138,7 @@ pub fn comp_graph(
                 max_dist_count,
                 min_dist_mag,
                 avg_b1_trig,
-                inv_kt_grid
+                inv_kt_grid,
             );
             dists_p = _dists_p;
             dists_z = _dists_z;
@@ -161,7 +173,9 @@ pub fn comp_graph(
                 // gives 1/3 * (k1^2 + k1*k2 + k2^2)
                 // k's are in rotations / meter but we need rad / m -> * 2pi
                 use std::f32::consts::TAU;
-                let b = 1.0 / 3.0 * TAU * TAU
+                let b = 1.0 / 3.0
+                    * TAU
+                    * TAU
                     * dt
                     * ((k1[0] * k1[0] + k1[0] * k2[0] + k2[0] * k2[0])
                         + (k1[1] * k1[1] + k1[1] * k2[1] + k2[1] * k2[1])
@@ -180,8 +194,9 @@ pub fn comp_graph(
         let r1 = (-rep_time / t1).exp();
 
         for mut dist in dists_z.iter().map(|d| d.borrow_mut()) {
+            use std::f64::consts::TAU;
             let sqr = |x| x * x;
-            let k2 = sqr(dist.kt_vec[0]) + sqr(dist.kt_vec[1]) + sqr(dist.kt_vec[2]);
+            let k2 = TAU * TAU * (sqr(dist.kt_vec[0]) + sqr(dist.kt_vec[1]) + sqr(dist.kt_vec[2]));
 
             dist.mag *= r1 * (-d * rep_time * k2 as f32).exp();
         }
@@ -354,7 +369,8 @@ pub fn analyze_graph(graph: &mut Vec<Vec<RcDist>>) {
                 let tmp = anc.dist.borrow().latent_signal;
                 let tmp_unormalized = anc.dist.borrow().latent_signal_unormalized;
                 anc.dist.borrow_mut().latent_signal = f32::max(tmp, dist.latent_signal * contrib);
-                anc.dist.borrow_mut().latent_signal_unormalized = f32::max(tmp_unormalized, dist.latent_signal_unormalized * contrib);
+                anc.dist.borrow_mut().latent_signal_unormalized =
+                    f32::max(tmp_unormalized, dist.latent_signal_unormalized * contrib);
             }
         }
     }

tests/simulation_test/build_report.py

@@ -0,0 +1,56 @@
+"""
+Assemble the PR comment from the summaries written by plot_comparison.py.
+
+Usage: python build_report.py <report_folder> <image_base_url> <run_url>
+Prints the markdown comment to stdout.
+"""
+import glob
+import json
+import os
+import sys
+
+MARKER = "<!-- simulation-test-report -->"
+ICONS = {"passed": "✅", "failed": "❌", "error": "⚠️"}
+
+
+def main():
+    report_folder, image_base_url, run_url = sys.argv[1:4]
+    summaries = []
+    for path in sorted(glob.glob(os.path.join(report_folder, "*.json"))):
+        with open(path) as f:
+            summaries.append(json.load(f))
+
+    lines = [
+        MARKER,
+        "## Simulation test: main (before) vs. this PR (after)",
+        "",
+        "| Sequence | Status | max mag NRMSE | max phase RMSE | max runtime Δ |",
+        "|---|---|---|---|---|",
+    ]
+    for s in summaries:
+        if s["status"] == "error":
+            lines.append(f"| {s['name']} | {ICONS['error']} | – | – | – |")
+        else:
+            lines.append(
+                f"| {s['name']} | {ICONS[s['status']]} | {s['max_mag_nrmse']:.2e} "
+                f"| {s['max_phase_rmse']:.2e} | {s['max_dt_percent']:+.1f}% |"
+            )
+    lines.append("")
+
+    for s in summaries:
+        # Failing sequences are expanded, passing ones collapsed
+        open_attr = " open" if s["status"] != "passed" else ""
+        lines.append(f"<details{open_attr}><summary>{ICONS[s['status']]} {s['name']}: {s['message']}</summary>")
+        lines.append("")
+        if "image" in s:
+            lines.append(f"![{s['name']}]({image_base_url}/{s['image']})")
+        lines.append("")
+        lines.append("</details>")
+        lines.append("")
+
+    lines.append(f"[Workflow run]({run_url}) — raw `.npz` data is attached as artifacts.")
+    print("\n".join(lines))
+
+
+if __name__ == "__main__":
+    main()

tests/simulation_test/plot_comparison.py

@@ -0,0 +1,96 @@
+"""
+Plot reference (main branch) vs. actual (current branch) simulation results.
+
+Usage: python plot_comparison.py <notebook_path> <output_folder>
+
+Writes <name>.png and <name>.json (summary used for the PR comment) into the
+output folder. Tolerates missing data so it can run after failed steps.
+"""
+import json
+import os
+import sys
+import numpy as np
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+import config
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+from utils import load_reference
+from test_simulation import compare_seq_parameters, compute_metrics
+
+
+def plot_comparison(name, actual, reference, output_folder):
+    mag_rmse, phase_rmse, dt_diff = compute_metrics(actual, reference)
+    # Show the accuracy level where the signals differ the most
+    i = int(np.argmax(mag_rmse + phase_rmse))
+    acc = config.ACC_ARRAY[i]
+    ref = np.asarray(reference["signal"][i]).reshape(-1)
+    act = np.asarray(actual["signal"][i]).reshape(-1)
+
+    fig, axes = plt.subplots(2, 2, figsize=(14, 7), constrained_layout=True)
+    fig.suptitle(f"{name}  (acc={acc:g}, worst level)")
+
+    ax = axes[0, 0]
+    ax.plot(np.abs(ref), lw=1, label="main (before)")
+    ax.plot(np.abs(act), lw=1, ls="--", label="PR (after)")
+    ax.set_title("Magnitude")
+    ax.legend()
+
+    ax = axes[0, 1]
+    ax.plot(np.angle(ref), lw=1, label="main (before)")
+    ax.plot(np.angle(act), lw=1, ls="--", label="PR (after)")
+    ax.set_title("Phase")
+    ax.legend()
+
+    ax = axes[1, 0]
+    ax.plot(np.abs(act - ref), lw=1, color="C3", label="|after - before|")
+    ax.plot(np.abs(act) - np.abs(ref), lw=1, color="C2", alpha=0.7, label="|after| - |before|")
+    ax.set_title("Difference")
+    ax.set_xlabel("sample")
+    ax.legend()
+
+    ax = axes[1, 1]
+    ax.loglog(config.ACC_ARRAY, np.maximum(mag_rmse, 1e-12), "o-", label="mag NRMSE")
+    ax.loglog(config.ACC_ARRAY, np.maximum(phase_rmse, 1e-12), "s-", label="phase RMSE")
+    ax.axhline(config.MAG_NRMSE, color="gray", ls=":", label="threshold")
+    ax.invert_xaxis()
+    ax.set_title("Error vs. accuracy")
+    ax.set_xlabel("accuracy")
+    ax.legend()
+
+    fig.savefig(os.path.join(output_folder, f"{name}.png"), dpi=90)
+    plt.close(fig)
+
+    passed, message = compare_seq_parameters(actual, reference)
+    return {
+        "name": name,
+        "status": "passed" if passed else "failed",
+        "message": message,
+        "image": f"{name}.png",
+        "max_mag_nrmse": float(mag_rmse.max()),
+        "max_phase_rmse": float(phase_rmse.max()),
+        "max_dt_percent": float(dt_diff.max()),
+    }
+
+
+def main():
+    notebook_path, output_folder = sys.argv[1:3]
+    os.makedirs(output_folder, exist_ok=True)
+    name = os.path.splitext(os.path.basename(notebook_path))[0]
+    npz_name = f"{name}.seq.npz"
+
+    try:
+        reference = load_reference(os.path.join(config.REF_FOLDER, npz_name))
+        actual = load_reference(os.path.join(config.ACTUAL_FOLDER, npz_name))
+        summary = plot_comparison(name, actual, reference, output_folder)
+    except Exception as e:
+        summary = {"name": name, "status": "error", "message": f"No comparison possible: {e}"}
+
+    with open(os.path.join(output_folder, f"{name}.json"), "w") as f:
+        json.dump(summary, f)
+    print(f"{name}: {summary['status']} - {summary['message']}")
+
+
+if __name__ == "__main__":
+    main()