saymrwulf/NTT-learning
1
0
Fork 0
mirror of https://github.com/saymrwulf/NTT-learning.git synced 2026-07-30 20:07:54 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Make NTT lessons more visual and concrete

Browse source
This commit is contained in:
saymrwulf 2026-04-15 21:20:58 +02:00
parent 901d465e38
commit edf1b69cf5
11 changed files with 623 additions and 30 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

18
notebooks/butterfly_mechanics/03_fast_forward_ct/lecture.ipynb
View file

@ -36,7 +36,7 @@
}
},
"outputs": [],
"source": "# MANDATORY | difficulty 3 | Trace The Exact n=4 Paper Example\n\nfrom IPython.display import display\n\nfrom ntt_learning.toy_ntt import fast_ntt_psi_ct_trace, forward_ntt_psi\nfrom ntt_learning.visuals import interactive_trace, plot_trace_overview\n\nsignal = [1, 2, 3, 4]\nmodulus = 7681\npsi = 1925\ntrace = fast_ntt_psi_ct_trace(signal, modulus, psi)\n\nprint(\"raw CT output (BO):\", trace.raw_output)\nprint(\"bit-reversed back to NO:\", trace.normal_order_output)\nprint(\"direct NTT_psi:\", forward_ntt_psi(signal, modulus, psi))\ndisplay(plot_trace_overview(trace, title=\"CT overview for [1,2,3,4]\"))\ndisplay(interactive_trace(trace, title=\"CT forward trace\"))\n"
"source": "# MANDATORY | difficulty 3 | Trace The Exact n=4 Paper Example\n\nfrom IPython.display import display\n\nfrom ntt_learning.toy_ntt import fast_ntt_psi_ct_trace, forward_ntt_psi\nfrom ntt_learning.visuals import interactive_trace, plot_butterfly_network, plot_trace_overview\n\nsignal = [1, 2, 3, 4]\nmodulus = 7681\npsi = 1925\ntrace = fast_ntt_psi_ct_trace(signal, modulus, psi)\n\nprint(\"raw CT output (BO):\", trace.raw_output)\nprint(\"bit-reversed back to NO:\", trace.normal_order_output)\nprint(\"direct NTT_psi:\", forward_ntt_psi(signal, modulus, psi))\ndisplay(plot_trace_overview(trace, title=\"CT overview for [1,2,3,4]\"))\ndisplay(plot_butterfly_network(trace, title=\"Full CT network for [1,2,3,4]\"))\ndisplay(interactive_trace(trace, title=\"CT forward trace\"))\n"
},
{
"cell_type": "markdown",
@ -76,7 +76,21 @@
}
},
"outputs": [],
"source": "# MANDATORY | difficulty 3 | Go One Stage Deeper With n=8\n\nfrom IPython.display import display\n\nfrom ntt_learning.toy_ntt import fast_ntt_psi_ct_trace, find_psi\nfrom ntt_learning.visuals import interactive_trace, plot_trace_overview\n\nsignal = [0, 1, 2, 3, 4, 5, 6, 7]\nmodulus = 97\npsi = find_psi(8, modulus)\ntrace = fast_ntt_psi_ct_trace(signal, modulus, psi)\n\nprint(\"psi:\", psi)\nprint(\"BO output:\", trace.raw_output)\nprint(\"NO output:\", trace.normal_order_output)\ndisplay(plot_trace_overview(trace, title=\"Three CT stages for n=8\"))\ndisplay(interactive_trace(trace, title=\"n=8 CT stage explorer\"))\n"
"source": "# MANDATORY | difficulty 3 | Go One Stage Deeper With n=8\n\nfrom IPython.display import display\n\nfrom ntt_learning.toy_ntt import fast_ntt_psi_ct_trace, find_psi\nfrom ntt_learning.visuals import interactive_trace, plot_butterfly_network, plot_trace_overview\n\nsignal = [0, 1, 2, 3, 4, 5, 6, 7]\nmodulus = 97\npsi = find_psi(8, modulus)\ntrace = fast_ntt_psi_ct_trace(signal, modulus, psi)\n\nprint(\"psi:\", psi)\nprint(\"BO output:\", trace.raw_output)\nprint(\"NO output:\", trace.normal_order_output)\ndisplay(plot_trace_overview(trace, title=\"Three CT stages for n=8\"))\ndisplay(plot_butterfly_network(trace, title=\"Full CT network for n=8\"))\ndisplay(interactive_trace(trace, title=\"n=8 CT stage explorer\"))\n"
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"pedagogy": {
"role": "mandatory",
"difficulty": 3,
"kind": "demo",
"title": "See BO Output And NO Output Side By Side"
}
},
"outputs": [],
"source": "# MANDATORY | difficulty 3 | See BO Output And NO Output Side By Side\n\nfrom IPython.display import display\n\nfrom ntt_learning.toy_ntt import fast_ntt_psi_ct_trace\nfrom ntt_learning.visuals import plot_bit_reversal_mapping, plot_vector_comparison\n\ntrace = fast_ntt_psi_ct_trace([1, 2, 3, 4], 7681, 1925)\ndisplay(\n plot_vector_comparison(\n trace.raw_output,\n trace.normal_order_output,\n left_label=\"BO\",\n right_label=\"NO\",\n title=\"Same CT values, different ordering\",\n )\n)\ndisplay(plot_bit_reversal_mapping(4, title=\"Why BO becomes NO after bit-reversal\"))\n"
},
{
"cell_type": "markdown",

4
notebooks/butterfly_mechanics/03_fast_forward_ct/studio.ipynb
View file

@ -36,7 +36,7 @@
}
},
"outputs": [],
"source": "# MANDATORY | difficulty 3 | Compare Two CT Traces\n\nfrom IPython.display import display\n\nfrom ntt_learning.toy_ntt import fast_ntt_psi_ct_trace\nfrom ntt_learning.visuals import plot_trace_overview\n\ntrace_a = fast_ntt_psi_ct_trace([1, 2, 3, 4], 7681, 1925)\ntrace_b = fast_ntt_psi_ct_trace([5, 6, 7, 8], 7681, 1925)\n\ndisplay(plot_trace_overview(trace_a, title=\"CT trace A\"))\ndisplay(plot_trace_overview(trace_b, title=\"CT trace B\"))\n"
"source": "# MANDATORY | difficulty 3 | Compare Two CT Traces\n\nfrom IPython.display import display\n\nfrom ntt_learning.toy_ntt import fast_ntt_psi_ct_trace\nfrom ntt_learning.visuals import plot_butterfly_network, plot_trace_overview\n\ntrace_a = fast_ntt_psi_ct_trace([1, 2, 3, 4], 7681, 1925)\ntrace_b = fast_ntt_psi_ct_trace([5, 6, 7, 8], 7681, 1925)\n\ndisplay(plot_trace_overview(trace_a, title=\"CT trace A\"))\ndisplay(plot_trace_overview(trace_b, title=\"CT trace B\"))\ndisplay(plot_butterfly_network(trace_a, title=\"CT network A\"))\ndisplay(plot_butterfly_network(trace_b, title=\"CT network B\"))\n"
},
{
"cell_type": "markdown",
@ -62,7 +62,7 @@
}
},
"outputs": [],
"source": "# MANDATORY | difficulty 2 | See A Wrong-Order Comparison Failure\n\nfrom ntt_learning.toy_ntt import fast_ntt_psi_ct_trace, forward_ntt_psi\n\ntrace = fast_ntt_psi_ct_trace([1, 2, 3, 4], 7681, 1925)\ndirect = forward_ntt_psi([1, 2, 3, 4], 7681, 1925)\n\nprint(\"wrong comparison: CT BO output vs direct NO output\")\nprint(trace.raw_output, direct)\nprint(\"correct comparison: CT NO output vs direct NO output\")\nprint(trace.normal_order_output, direct)\n"
"source": "# MANDATORY | difficulty 2 | See A Wrong-Order Comparison Failure\n\nfrom IPython.display import display\n\nfrom ntt_learning.toy_ntt import fast_ntt_psi_ct_trace, forward_ntt_psi\nfrom ntt_learning.visuals import plot_vector_comparison\n\ntrace = fast_ntt_psi_ct_trace([1, 2, 3, 4], 7681, 1925)\ndirect = forward_ntt_psi([1, 2, 3, 4], 7681, 1925)\n\nprint(\"wrong comparison: CT BO output vs direct NO output\")\nprint(trace.raw_output, direct)\nprint(\"correct comparison: CT NO output vs direct NO output\")\nprint(trace.normal_order_output, direct)\ndisplay(\n plot_vector_comparison(\n trace.raw_output,\n direct,\n left_label=\"CT_BO\",\n right_label=\"direct_NO\",\n title=\"Wrong comparison: BO against NO\",\n )\n)\ndisplay(\n plot_vector_comparison(\n trace.normal_order_output,\n direct,\n left_label=\"CT_NO\",\n right_label=\"direct_NO\",\n title=\"Correct comparison after reordering\",\n )\n)\n"
},
{
"cell_type": "markdown",

4
notebooks/butterfly_mechanics/04_fast_inverse_gs/lecture.ipynb
View file

@ -36,7 +36,7 @@
}
},
"outputs": [],
"source": "# MANDATORY | difficulty 3 | Trace The Exact n=4 GS Paper Example\n\nfrom IPython.display import display\n\nfrom ntt_learning.toy_ntt import fast_intt_psi_gs_trace\nfrom ntt_learning.visuals import interactive_trace, plot_trace_overview\n\nbo_input = [1467, 3471, 2807, 7621]\ntrace = fast_intt_psi_gs_trace(bo_input, 7681, 1925)\n\nprint(\"unscaled NO output:\", trace.raw_output)\nprint(\"scaled NO output:\", trace.scaled_output)\ndisplay(plot_trace_overview(trace, title=\"GS overview for the n=4 paper example\"))\ndisplay(interactive_trace(trace, title=\"GS inverse trace\"))\n"
"source": "# MANDATORY | difficulty 3 | Trace The Exact n=4 GS Paper Example\n\nfrom IPython.display import display\n\nfrom ntt_learning.toy_ntt import fast_intt_psi_gs_trace\nfrom ntt_learning.visuals import interactive_trace, plot_butterfly_network, plot_trace_overview, plot_vector_comparison\n\nbo_input = [1467, 3471, 2807, 7621]\ntrace = fast_intt_psi_gs_trace(bo_input, 7681, 1925)\n\nprint(\"unscaled NO output:\", trace.raw_output)\nprint(\"scaled NO output:\", trace.scaled_output)\ndisplay(plot_trace_overview(trace, title=\"GS overview for the n=4 paper example\"))\ndisplay(plot_butterfly_network(trace, title=\"Full GS network for the n=4 paper example\"))\ndisplay(\n plot_vector_comparison(\n trace.raw_output,\n trace.scaled_output,\n left_label=\"unscaled\",\n right_label=\"scaled\",\n title=\"Why the final n^-1 scaling matters\",\n )\n)\ndisplay(interactive_trace(trace, title=\"GS inverse trace\"))\n"
},
{
"cell_type": "code",
@ -76,7 +76,7 @@
}
},
"outputs": [],
"source": "# MANDATORY | difficulty 3 | Full Forward And Inverse Round Trip\n\nfrom ntt_learning.toy_ntt import fast_intt_psi_gs_trace, fast_ntt_psi_ct_trace\n\nsignal = [1, 2, 3, 4]\nforward_trace = fast_ntt_psi_ct_trace(signal, 7681, 1925)\ninverse_trace = fast_intt_psi_gs_trace(forward_trace.raw_output, 7681, 1925)\n\nprint(\"forward BO output:\", forward_trace.raw_output)\nprint(\"inverse scaled output:\", inverse_trace.scaled_output)\n"
"source": "# MANDATORY | difficulty 3 | Full Forward And Inverse Round Trip\n\nfrom IPython.display import display\n\nfrom ntt_learning.toy_ntt import fast_intt_psi_gs_trace, fast_ntt_psi_ct_trace\nfrom ntt_learning.visuals import plot_vector_comparison\n\nsignal = [1, 2, 3, 4]\nforward_trace = fast_ntt_psi_ct_trace(signal, 7681, 1925)\ninverse_trace = fast_intt_psi_gs_trace(forward_trace.raw_output, 7681, 1925)\n\nprint(\"forward BO output:\", forward_trace.raw_output)\nprint(\"inverse scaled output:\", inverse_trace.scaled_output)\ndisplay(\n plot_vector_comparison(\n signal,\n inverse_trace.scaled_output,\n left_label=\"original\",\n right_label=\"recovered\",\n title=\"Forward CT followed by inverse GS\",\n )\n)\n"
},
{
"cell_type": "markdown",

4
notebooks/butterfly_mechanics/04_fast_inverse_gs/studio.ipynb
View file

@ -36,7 +36,7 @@
}
},
"outputs": [],
"source": "# MANDATORY | difficulty 3 | See CT Output Feed GS Input\n\nfrom ntt_learning.toy_ntt import fast_intt_psi_gs_trace, fast_ntt_psi_ct_trace\n\nsignal = [5, 6, 7, 8]\nforward_trace = fast_ntt_psi_ct_trace(signal, 7681, 1925)\ninverse_trace = fast_intt_psi_gs_trace(forward_trace.raw_output, 7681, 1925)\n\nprint(\"CT BO output:\", forward_trace.raw_output)\nprint(\"GS scaled output:\", inverse_trace.scaled_output)\n"
"source": "# MANDATORY | difficulty 3 | See CT Output Feed GS Input\n\nfrom IPython.display import display\n\nfrom ntt_learning.toy_ntt import fast_intt_psi_gs_trace, fast_ntt_psi_ct_trace\nfrom ntt_learning.visuals import plot_vector_comparison\n\nsignal = [5, 6, 7, 8]\nforward_trace = fast_ntt_psi_ct_trace(signal, 7681, 1925)\ninverse_trace = fast_intt_psi_gs_trace(forward_trace.raw_output, 7681, 1925)\n\nprint(\"CT BO output:\", forward_trace.raw_output)\nprint(\"GS scaled output:\", inverse_trace.scaled_output)\ndisplay(\n plot_vector_comparison(\n signal,\n inverse_trace.scaled_output,\n left_label=\"start\",\n right_label=\"after_CT_then_GS\",\n title=\"CT output cleanly feeds GS input\",\n )\n)\n"
},
{
"cell_type": "markdown",
@ -62,7 +62,7 @@
}
},
"outputs": [],
"source": "# MANDATORY | difficulty 2 | See A Missing-Scale Failure\n\nfrom ntt_learning.toy_ntt import fast_intt_psi_gs_trace\n\ntrace = fast_intt_psi_gs_trace([1467, 3471, 2807, 7621], 7681, 1925)\nprint(\"unscaled:\", trace.raw_output)\nprint(\"scaled:\", trace.scaled_output)\n"
"source": "# MANDATORY | difficulty 2 | See A Missing-Scale Failure\n\nfrom IPython.display import display\n\nfrom ntt_learning.toy_ntt import fast_intt_psi_gs_trace\nfrom ntt_learning.visuals import plot_vector_comparison\n\ntrace = fast_intt_psi_gs_trace([1467, 3471, 2807, 7621], 7681, 1925)\nprint(\"unscaled:\", trace.raw_output)\nprint(\"scaled:\", trace.scaled_output)\ndisplay(\n plot_vector_comparison(\n trace.raw_output,\n trace.scaled_output,\n left_label=\"missing_scale\",\n right_label=\"correct\",\n title=\"Missing n^-1 scale vs correct output\",\n )\n)\n"
},
{
"cell_type": "markdown",

4
notebooks/foundations/02_negative_wrapped_ntt/lecture.ipynb
View file

@ -36,7 +36,7 @@
}
},
"outputs": [],
"source": "# MANDATORY | difficulty 2 | Inspect \u03c9, \u03c8, And The Direct Transform Matrix\n\nfrom ntt_learning.toy_ntt import find_primitive_root, find_psi, ntt_psi_exponent_grid, ntt_psi_matrix\n\nmodulus = 17\nn = 4\nomega = find_primitive_root(n, modulus)\npsi = find_psi(n, modulus)\n\nprint(\"omega:\", omega)\nprint(\"psi:\", psi)\nprint(\"exponent grid:\")\nfor row in ntt_psi_exponent_grid(n):\n print(row)\nprint(\"NTT_psi matrix:\")\nfor row in ntt_psi_matrix(n, modulus, psi):\n print(row)\n"
"source": "# MANDATORY | difficulty 2 | Inspect \u03c9, \u03c8, And The Direct Transform Matrix\n\nfrom IPython.display import display\n\nfrom ntt_learning.toy_ntt import find_primitive_root, find_psi, ntt_psi_exponent_grid, ntt_psi_matrix\nfrom ntt_learning.visuals import plot_ntt_psi_exponent_heatmap, plot_ntt_psi_matrix_heatmap\n\nmodulus = 17\nn = 4\nomega = find_primitive_root(n, modulus)\npsi = find_psi(n, modulus)\n\nprint(\"omega:\", omega)\nprint(\"psi:\", psi)\nprint(\"exponent grid:\")\nfor row in ntt_psi_exponent_grid(n):\n print(row)\nprint(\"NTT_psi matrix:\")\nfor row in ntt_psi_matrix(n, modulus, psi):\n print(row)\n\ndisplay(plot_ntt_psi_exponent_heatmap(n, title=\"Exponents 2ij + i for n=4\"))\ndisplay(plot_ntt_psi_matrix_heatmap(n, modulus, psi, title=\"Concrete NTT_psi matrix in Z_17\"))\n"
},
{
"cell_type": "markdown",
@ -76,7 +76,7 @@
}
},
"outputs": [],
"source": "# MANDATORY | difficulty 3 | Use Direct NTT\u03c8 For Negacyclic Multiplication\n\nfrom ntt_learning.toy_ntt import find_psi, forward_ntt_psi, inverse_ntt_psi, negacyclic_multiply, pointwise_multiply\n\nleft = [1, 2, 3, 4]\nright = [5, 6, 7, 8]\nmodulus = 17\npsi = find_psi(4, modulus)\n\nleft_hat = forward_ntt_psi(left, modulus, psi)\nright_hat = forward_ntt_psi(right, modulus, psi)\nproduct_hat = pointwise_multiply(left_hat, right_hat, modulus)\n\nprint(\"NTT_psi(left):\", left_hat)\nprint(\"NTT_psi(right):\", right_hat)\nprint(\"pointwise product:\", product_hat)\nprint(\"inverse of pointwise product:\", inverse_ntt_psi(product_hat, modulus, psi))\nprint(\"schoolbook negacyclic:\", negacyclic_multiply(left, right, n=4, modulus=modulus))\n"
"source": "# MANDATORY | difficulty 3 | Use Direct NTT\u03c8 For Negacyclic Multiplication\n\nfrom IPython.display import display\n\nfrom ntt_learning.toy_ntt import find_psi, forward_ntt_psi, inverse_ntt_psi, negacyclic_multiply, pointwise_multiply\nfrom ntt_learning.visuals import plot_transform_pipeline\n\nleft = [1, 2, 3, 4]\nright = [5, 6, 7, 8]\nmodulus = 17\npsi = find_psi(4, modulus)\n\nleft_hat = forward_ntt_psi(left, modulus, psi)\nright_hat = forward_ntt_psi(right, modulus, psi)\nproduct_hat = pointwise_multiply(left_hat, right_hat, modulus)\n\nprint(\"NTT_psi(left):\", left_hat)\nprint(\"NTT_psi(right):\", right_hat)\nprint(\"pointwise product:\", product_hat)\nprint(\"inverse of pointwise product:\", inverse_ntt_psi(product_hat, modulus, psi))\nprint(\"schoolbook negacyclic:\", negacyclic_multiply(left, right, n=4, modulus=modulus))\ndisplay(plot_transform_pipeline(left, right, modulus=modulus, psi=psi, title=\"Direct negacyclic multiply pipeline\"))\n"
},
{
"cell_type": "markdown",

4
notebooks/foundations/02_negative_wrapped_ntt/studio.ipynb
View file

@ -36,7 +36,7 @@
}
},
"outputs": [],
"source": "# MANDATORY | difficulty 3 | Compare Positive-Wrapped And Negative-Wrapped Views\n\nfrom ntt_learning.toy_ntt import find_primitive_root, find_psi, forward_ntt, forward_ntt_psi\n\nsignal = [1, 2, 3, 4]\nmodulus = 17\nomega = find_primitive_root(4, modulus)\npsi = find_psi(4, modulus)\n\nprint(\"positive-wrapped NTT:\", forward_ntt(signal, modulus, omega))\nprint(\"negative-wrapped NTT_psi:\", forward_ntt_psi(signal, modulus, psi))\n"
"source": "# MANDATORY | difficulty 3 | Compare Positive-Wrapped And Negative-Wrapped Views\n\nfrom IPython.display import display\n\nfrom ntt_learning.toy_ntt import find_primitive_root, find_psi, forward_ntt, forward_ntt_psi\nfrom ntt_learning.visuals import plot_vector_comparison\n\nsignal = [1, 2, 3, 4]\nmodulus = 17\nomega = find_primitive_root(4, modulus)\npsi = find_psi(4, modulus)\n\npositive = forward_ntt(signal, modulus, omega)\nnegative = forward_ntt_psi(signal, modulus, psi)\n\nprint(\"positive-wrapped NTT:\", positive)\nprint(\"negative-wrapped NTT_psi:\", negative)\ndisplay(\n plot_vector_comparison(\n positive,\n negative,\n left_label=\"positive\",\n right_label=\"negative\",\n title=\"Same signal, different transform story\",\n )\n)\n"
},
{
"cell_type": "markdown",
@ -62,7 +62,7 @@
}
},
"outputs": [],
"source": "# MANDATORY | difficulty 2 | See A Wrong-Root Failure\n\nfrom ntt_learning.toy_ntt import find_primitive_root, find_psi, forward_ntt_psi\n\nsignal = [1, 2, 3, 4]\nmodulus = 17\nomega = find_primitive_root(4, modulus)\npsi = find_psi(4, modulus)\n\nprint(\"correct psi-based transform:\", forward_ntt_psi(signal, modulus, psi))\nprint(\"wrongly using omega as if it were psi:\", forward_ntt_psi(signal, modulus, omega))\n"
"source": "# MANDATORY | difficulty 2 | See A Wrong-Root Failure\n\nfrom IPython.display import display\n\nfrom ntt_learning.toy_ntt import find_primitive_root, find_psi, forward_ntt_psi\nfrom ntt_learning.visuals import plot_vector_comparison\n\nsignal = [1, 2, 3, 4]\nmodulus = 17\nomega = find_primitive_root(4, modulus)\npsi = find_psi(4, modulus)\n\ncorrect = forward_ntt_psi(signal, modulus, psi)\nwrong = forward_ntt_psi(signal, modulus, omega)\n\nprint(\"correct psi-based transform:\", correct)\nprint(\"wrongly using omega as if it were psi:\", wrong)\ndisplay(\n plot_vector_comparison(\n wrong,\n correct,\n left_label=\"wrong_root\",\n right_label=\"correct\",\n title=\"Wrong root vs correct root\",\n )\n)\n"
},
{
"cell_type": "markdown",

4
notebooks/kyber_mapping/05_kyber_ntt_and_base_multiplication/lecture.ipynb
View file

@ -36,7 +36,7 @@
}
},
"outputs": [],
"source": "# MANDATORY | difficulty 3 | Check The Kyber Root Reality Directly\n\nfrom ntt_learning.toy_ntt import find_primitive_root\n\nprint(\"3329 - 1 =\", 3329 - 1)\nprint(\"primitive 256-th root in Z_3329:\", find_primitive_root(256, 3329))\ntry:\n find_primitive_root(512, 3329)\nexcept Exception as exc:\n print(\"512-th root fails exactly because:\", exc)\n"
"source": "# MANDATORY | difficulty 3 | Check The Kyber Root Reality Directly\n\nfrom IPython.display import display\n\nfrom ntt_learning.toy_ntt import find_primitive_root\nfrom ntt_learning.visuals import plot_root_order_comparison\n\nprint(\"3329 - 1 =\", 3329 - 1)\nprint(\"primitive 256-th root in Z_3329:\", find_primitive_root(256, 3329))\ntry:\n find_primitive_root(512, 3329)\nexcept Exception as exc:\n print(\"512-th root fails exactly because:\", exc)\n\ndisplay(\n plot_root_order_comparison(\n [(4, 17), (4, 13), (8, 97), (256, 3329)],\n title=\"Which moduli allow n-th and 2n-th root stories?\",\n )\n)\n"
},
{
"cell_type": "markdown",
@ -88,7 +88,7 @@
}
},
"outputs": [],
"source": "# MANDATORY | difficulty 3 | See A Toy Base Multiplication\n\nfrom ntt_learning.toy_ntt import base_multiply_pair\n\nleft = [7, 11]\nright = [5, 13]\nzeta = 4\nmodulus = 17\n\nraw = [left[0] * right[0], left[0] * right[1] + left[1] * right[0], left[1] * right[1]]\nreduced = [(raw[0] + zeta * raw[2]) % modulus, raw[1] % modulus]\n\nprint(\"raw degree-2 product:\", raw)\nprint(\"reduce with x^2 = zeta:\", reduced)\nprint(\"base_multiply_pair:\", base_multiply_pair(left, right, zeta, modulus))\n"
"source": "# MANDATORY | difficulty 3 | See A Toy Base Multiplication\n\nfrom IPython.display import display\n\nfrom ntt_learning.toy_ntt import base_multiply_pair\nfrom ntt_learning.visuals import plot_base_multiply_pair_diagram\n\nleft = [7, 11]\nright = [5, 13]\nzeta = 4\nmodulus = 17\n\nraw = [left[0] * right[0], left[0] * right[1] + left[1] * right[0], left[1] * right[1]]\nreduced = [(raw[0] + zeta * raw[2]) % modulus, raw[1] % modulus]\n\nprint(\"raw degree-2 product:\", raw)\nprint(\"reduce with x^2 = zeta:\", reduced)\nprint(\"base_multiply_pair:\", base_multiply_pair(left, right, zeta, modulus))\ndisplay(\n plot_base_multiply_pair_diagram(\n left,\n right,\n zeta=zeta,\n modulus=modulus,\n title=\"Why Kyber-style multiplication stays in 2-slot blocks\",\n )\n)\n"
},
{
"cell_type": "markdown",

4
notebooks/kyber_mapping/05_kyber_ntt_and_base_multiplication/studio.ipynb
View file

@ -36,7 +36,7 @@
}
},
"outputs": [],
"source": "# MANDATORY | difficulty 3 | Toy Full \u03c8 Story vs Kyber Root Reality\n\nfrom ntt_learning.toy_ntt import find_psi\n\nprint(\"toy n=4, q=17 has psi:\", find_psi(4, 17))\ntry:\n find_psi(256, 3329)\nexcept Exception as exc:\n print(\"Kyber v3 does not have that full psi story:\", exc)\n"
"source": "# MANDATORY | difficulty 3 | Toy Full \u03c8 Story vs Kyber Root Reality\n\nfrom IPython.display import display\n\nfrom ntt_learning.toy_ntt import find_psi\nfrom ntt_learning.visuals import plot_root_order_comparison\n\nprint(\"toy n=4, q=17 has psi:\", find_psi(4, 17))\ntry:\n find_psi(256, 3329)\nexcept Exception as exc:\n print(\"Kyber v3 does not have that full psi story:\", exc)\n\ndisplay(\n plot_root_order_comparison(\n [(4, 17), (256, 7681), (256, 3329)],\n title=\"Toy full psi story vs Kyber modulus reality\",\n )\n)\n"
},
{
"cell_type": "markdown",
@ -62,7 +62,7 @@
}
},
"outputs": [],
"source": "# MANDATORY | difficulty 2 | See The Exact Obstruction Again\n\nq = 3329\nn = 256\nprint({\"q_minus_1\": q - 1, \"n\": n, \"2n\": 2 * n, \"q_minus_1_mod_n\": (q - 1) % n, \"q_minus_1_mod_2n\": (q - 1) % (2 * n)})\n"
"source": "# MANDATORY | difficulty 2 | See The Exact Obstruction Again\n\nfrom IPython.display import display\n\nfrom ntt_learning.visuals import plot_root_order_comparison\n\nq = 3329\nn = 256\nprint({\"q_minus_1\": q - 1, \"n\": n, \"2n\": 2 * n, \"q_minus_1_mod_n\": (q - 1) % n, \"q_minus_1_mod_2n\": (q - 1) % (2 * n)})\ndisplay(plot_root_order_comparison([(256, 3329)], title=\"Kyber v3 obstruction in one row\"))\n"
},
{
"cell_type": "markdown",

4
notebooks/professional/06_debugging_ntt_failures/lecture.ipynb
View file

@ -36,7 +36,7 @@
}
},
"outputs": [],
"source": "# MANDATORY | difficulty 3 | See Four Failure Modes Side By Side\n\nfrom ntt_learning.toy_ntt import (\n fast_intt_psi_gs_trace,\n fast_ntt_psi_ct_trace,\n forward_ntt_psi,\n negacyclic_reduce,\n schoolbook_convolution,\n)\n\nsignal = [1, 2, 3, 4]\nforward_trace = fast_ntt_psi_ct_trace(signal, 7681, 1925)\ninverse_trace = fast_intt_psi_gs_trace(forward_trace.raw_output, 7681, 1925)\n\nraw = schoolbook_convolution([1, 2, 3, 4], [5, 6, 7, 8])\nwrong_sign = [raw[0] + raw[4], raw[1] + raw[5], raw[2] + raw[6], raw[3]]\nwrong_order = list(forward_trace.raw_output)\nwrong_scale = list(inverse_trace.raw_output)\nwrong_root = forward_ntt_psi(signal, 7681, 3383)\n\nprint(\"wrong sign fold:\", wrong_sign)\nprint(\"correct sign fold:\", negacyclic_reduce(raw, n=4))\nprint(\"wrong BO-vs-NO comparison:\", wrong_order)\nprint(\"correct NO output:\", forward_trace.normal_order_output)\nprint(\"missing final scaling:\", wrong_scale)\nprint(\"wrong root in direct transform:\", wrong_root)\n"
"source": "# MANDATORY | difficulty 3 | See Four Failure Modes Side By Side\n\nfrom IPython.display import display\n\nfrom ntt_learning.toy_ntt import (\n fast_intt_psi_gs_trace,\n fast_ntt_psi_ct_trace,\n forward_ntt_psi,\n negacyclic_reduce,\n schoolbook_convolution,\n)\nfrom ntt_learning.visuals import plot_vector_comparison\n\nsignal = [1, 2, 3, 4]\nforward_trace = fast_ntt_psi_ct_trace(signal, 7681, 1925)\ninverse_trace = fast_intt_psi_gs_trace(forward_trace.raw_output, 7681, 1925)\n\nraw = schoolbook_convolution([1, 2, 3, 4], [5, 6, 7, 8])\nwrong_sign = [raw[0] + raw[4], raw[1] + raw[5], raw[2] + raw[6], raw[3]]\nwrong_order = list(forward_trace.raw_output)\nwrong_scale = list(inverse_trace.raw_output)\nwrong_root = forward_ntt_psi(signal, 7681, 3383)\n\nprint(\"wrong sign fold:\", wrong_sign)\nprint(\"correct sign fold:\", negacyclic_reduce(raw, n=4))\nprint(\"wrong BO-vs-NO comparison:\", wrong_order)\nprint(\"correct NO output:\", forward_trace.normal_order_output)\nprint(\"missing final scaling:\", wrong_scale)\nprint(\"wrong root in direct transform:\", wrong_root)\ndisplay(\n plot_vector_comparison(\n wrong_sign,\n negacyclic_reduce(raw, n=4),\n left_label=\"wrong_sign\",\n right_label=\"correct_sign\",\n title=\"Wrong sign vs correct negacyclic fold\",\n )\n)\ndisplay(\n plot_vector_comparison(\n wrong_order,\n forward_trace.normal_order_output,\n left_label=\"wrong_order\",\n right_label=\"correct_order\",\n title=\"Wrong BO/NO comparison\",\n )\n)\ndisplay(\n plot_vector_comparison(\n wrong_scale,\n inverse_trace.scaled_output,\n left_label=\"missing_scale\",\n right_label=\"correct_scale\",\n title=\"Missing scale vs corrected inverse\",\n )\n)\n"
},
{
"cell_type": "code",
@ -50,7 +50,7 @@
}
},
"outputs": [],
"source": "# MANDATORY | difficulty 3 | Interactive Failure Picker\n\nimport ipywidgets as widgets\nfrom IPython.display import display\n\nfrom ntt_learning.toy_ntt import fast_intt_psi_gs_trace, fast_ntt_psi_ct_trace, forward_ntt_psi\n\nforward_trace = fast_ntt_psi_ct_trace([1, 2, 3, 4], 7681, 1925)\ninverse_trace = fast_intt_psi_gs_trace(forward_trace.raw_output, 7681, 1925)\n\nfailures = {\n \"wrong_order\": list(forward_trace.raw_output),\n \"correct_order\": list(forward_trace.normal_order_output),\n \"missing_scale\": list(inverse_trace.raw_output),\n \"scaled\": list(inverse_trace.scaled_output),\n \"wrong_root\": forward_ntt_psi([1, 2, 3, 4], 7681, 3383),\n}\n\ndef preview(mode=\"wrong_order\"):\n print(mode, \"->\", failures[mode])\n\ndisplay(widgets.interact(preview, mode=sorted(failures)))\n"
"source": "# MANDATORY | difficulty 3 | Interactive Failure Picker\n\nimport ipywidgets as widgets\nfrom IPython.display import display\n\nfrom ntt_learning.toy_ntt import fast_intt_psi_gs_trace, fast_ntt_psi_ct_trace, forward_ntt_psi\nfrom ntt_learning.visuals import plot_vector_comparison\n\nforward_trace = fast_ntt_psi_ct_trace([1, 2, 3, 4], 7681, 1925)\ninverse_trace = fast_intt_psi_gs_trace(forward_trace.raw_output, 7681, 1925)\n\nfailures = {\n \"wrong_order\": list(forward_trace.raw_output),\n \"correct_order\": list(forward_trace.normal_order_output),\n \"missing_scale\": list(inverse_trace.raw_output),\n \"scaled\": list(inverse_trace.scaled_output),\n \"wrong_root\": forward_ntt_psi([1, 2, 3, 4], 7681, 3383),\n}\nreferences = {\n \"wrong_order\": list(forward_trace.normal_order_output),\n \"correct_order\": list(forward_trace.normal_order_output),\n \"missing_scale\": list(inverse_trace.scaled_output),\n \"scaled\": list(inverse_trace.scaled_output),\n \"wrong_root\": list(forward_ntt_psi([1, 2, 3, 4], 7681, 1925)),\n}\n\ndef preview(mode=\"wrong_order\"):\n print(mode, \"->\", failures[mode])\n display(\n plot_vector_comparison(\n failures[mode],\n references[mode],\n left_label=mode,\n right_label=\"reference\",\n title=f\"{mode} compared with the correct reference\",\n )\n )\n\ndisplay(widgets.interact(preview, mode=sorted(failures)))\n"
},
{
"cell_type": "markdown",

368
ntt_learning/visuals.py
View file

@ -11,7 +11,19 @@ matplotlib.use("Agg")
import matplotlib.pyplot as plt
from IPython.display import clear_output, display
from .toy_ntt import TransformStage, TransformTrace, pairwise_product_grid, wraparound_contributions
from .toy_ntt import (
TransformStage,
TransformTrace,
base_multiply_pair,
bit_reversed_order,
forward_ntt_psi,
inverse_ntt_psi,
ntt_psi_exponent_grid,
ntt_psi_matrix,
pairwise_product_grid,
pointwise_multiply,
wraparound_contributions,
)
def _value_colors(values: Sequence[int]) -> list[str]:
@ -46,6 +58,45 @@ def _draw_value_row(ax, values: Sequence[int], y: float, prefix: str) -> None:
)
def _annotate_grid(ax, grid: Sequence[Sequence[int]]) -> None:
for row, row_values in enumerate(grid):
for column, value in enumerate(row_values):
ax.text(
column,
row,
str(value),
ha="center",
va="center",
color="#101010",
fontsize=10,
family="monospace",
)
def plot_integer_grid(
grid: Sequence[Sequence[int]],
*,
title: str,
x_label: str,
y_label: str,
cmap: str = "YlGnBu",
):
"""Plot a heatmap with the exact integer values written in every cell."""
if not grid or not grid[0]:
raise ValueError("plot_integer_grid requires a non-empty rectangular grid")
fig, ax = plt.subplots(figsize=(max(6, len(grid[0]) * 1.2), max(4, len(grid) * 0.85)))
ax.imshow(grid, cmap=cmap, aspect="auto")
ax.set_title(title, fontsize=14, fontweight="bold")
ax.set_xlabel(x_label)
ax.set_ylabel(y_label)
ax.set_xticks(range(len(grid[0])))
ax.set_yticks(range(len(grid)))
_annotate_grid(ax, grid)
fig.tight_layout()
return fig
def plot_convolution_grid(
left: Sequence[int], right: Sequence[int], title: str = "Schoolbook Product Grid"
):
@ -70,9 +121,7 @@ def plot_convolution_grid(
heatmap_ax.set_xticks(range(len(right)))
heatmap_ax.set_yticks(range(len(left)))
for row, row_values in enumerate(grid):
for column, value in enumerate(row_values):
heatmap_ax.text(column, row, str(value), ha="center", va="center", color="#101010", fontsize=10)
_annotate_grid(heatmap_ax, grid)
sum_ax.axis("off")
sum_ax.set_title("Diagonal Sums = Convolution Coefficients", fontsize=12, fontweight="bold", pad=8)
@ -98,6 +147,28 @@ def plot_convolution_grid(
return fig
def plot_ntt_psi_exponent_heatmap(length: int, title: str = "NTT_psi Exponent Grid"):
"""Plot the exponent pattern 2ij + i used by the direct negative-wrapped NTT."""
return plot_integer_grid(
ntt_psi_exponent_grid(length),
title=title,
x_label="output index j",
y_label="input index i",
cmap="YlOrRd",
)
def plot_ntt_psi_matrix_heatmap(length: int, modulus: int, psi: int, title: str = "NTT_psi Matrix Values"):
"""Plot the concrete direct transform matrix over Z_q."""
return plot_integer_grid(
ntt_psi_matrix(length, modulus, psi),
title=title,
x_label="output index j",
y_label="input index i",
cmap="PuBuGn",
)
def plot_wraparound(
coefficients: Sequence[int],
n: int,
@ -150,6 +221,51 @@ def plot_wraparound(
return fig
def plot_vector_comparison(
left: Sequence[int],
right: Sequence[int],
*,
left_label: str = "left",
right_label: str = "right",
title: str = "Vector Comparison",
):
"""Plot two vectors slot-by-slot with explicit differences."""
if len(left) != len(right):
raise ValueError("plot_vector_comparison requires equal-length vectors")
differences = [int(right_value - left_value) for left_value, right_value in zip(left, right)]
fig, axes = plt.subplots(3, 1, figsize=(max(8, len(left) * 1.3), 7), height_ratios=[1, 1, 1])
labels = [left_label, right_label, "delta"]
rows = [left, right, differences]
row_colors = ["#edf6f9", "#fff3b0", "#f5cac3"]
for ax, label, values, row_color in zip(axes, labels, rows, row_colors):
ax.axis("off")
ax.set_title(label, fontsize=12, fontweight="bold", pad=6)
for index, value in enumerate(values):
ax.text(
index,
0,
f"{index}\n{value}",
ha="center",
va="center",
fontsize=10,
family="monospace",
bbox={
"boxstyle": "round,pad=0.32",
"facecolor": row_color if label != "delta" else _value_colors([value])[0],
"edgecolor": "#222222",
"linewidth": 1.1,
},
)
ax.set_xlim(-0.8, len(values) - 0.2)
ax.set_ylim(-0.8, 0.8)
fig.suptitle(title, fontsize=14, fontweight="bold")
fig.tight_layout()
return fig
def plot_bit_reversal_mapping(length: int, title: str = "Normal Order To Bit-Reversed Order"):
"""Plot the bit-reversal permutation as explicit wires."""
if length <= 0 or length & (length - 1):
@ -193,6 +309,77 @@ def plot_bit_reversal_mapping(length: int, title: str = "Normal Order To Bit-Rev
return fig
def plot_butterfly_network(trace: TransformTrace, title: str | None = None):
"""Plot the whole staged network with pair links visible at each stage."""
if title is None:
title = f"{trace.algorithm.upper()} Butterfly Network"
columns = [trace.input_values] + [stage.output_values for stage in trace.stages]
fig, ax = plt.subplots(figsize=(max(10, len(columns) * 2.4), max(5, len(trace.input_values) * 0.8)))
ax.set_title(title, fontsize=14, fontweight="bold")
ax.axis("off")
x_positions = [index * 2.4 for index in range(len(columns))]
for column_index, (x, values) in enumerate(zip(x_positions, columns)):
for row_index, value in enumerate(values):
ax.text(
x,
-row_index,
str(value),
ha="center",
va="center",
fontsize=10,
family="monospace",
bbox={
"boxstyle": "round,pad=0.26",
"facecolor": _value_colors([value])[0],
"edgecolor": "#222222",
"linewidth": 1.0,
},
)
label = "input" if column_index == 0 else f"s{column_index}"
ax.text(x, 1.1, label, ha="center", va="center", fontsize=11, fontweight="bold")
if column_index == 0:
continue
previous_x = x_positions[column_index - 1]
stage = trace.stages[column_index - 1]
colors = ["#264653", "#2a9d8f", "#e76f51", "#8d99ae", "#c1121f", "#3a86ff"]
for index in range(len(values)):
ax.plot([previous_x + 0.35, x - 0.35], [-index, -index], color="#b0b0b0", linewidth=0.9, alpha=0.65)
for pair_index, ((left, right), zeta) in enumerate(zip(stage.pairings, stage.zetas)):
color = colors[pair_index % len(colors)]
x_mid = (previous_x + x) / 2
ax.plot([previous_x + 0.35, x - 0.35], [-left, -left], color=color, linewidth=2.2)
ax.plot([previous_x + 0.35, x - 0.35], [-right, -right], color=color, linewidth=2.2)
ax.plot([x_mid, x_mid], [-left, -right], color=color, linewidth=2.6, alpha=0.95)
ax.text(
x_mid,
-((left + right) / 2),
f"zeta={zeta}",
ha="center",
va="center",
fontsize=8,
family="monospace",
bbox={
"boxstyle": "round,pad=0.18",
"facecolor": "#ffffff",
"edgecolor": color,
"linewidth": 1.0,
},
)
ax.set_xlim(-1.1, x_positions[-1] + 1.1)
ax.set_ylim(-len(trace.input_values) + 0.2, 1.8)
fig.tight_layout()
return fig
def plot_stage(stage: TransformStage, title: str | None = None):
"""Plot one explicit butterfly stage with input and output rows."""
if title is None:
@ -245,6 +432,179 @@ def plot_stage(stage: TransformStage, title: str | None = None):
return fig
def plot_transform_pipeline(
left: Sequence[int],
right: Sequence[int],
*,
modulus: int,
psi: int,
title: str = "Transform-Domain Multiply Pipeline",
):
"""Plot the end-to-end direct NTT_psi multiply pipeline."""
left_hat = forward_ntt_psi(left, modulus, psi)
right_hat = forward_ntt_psi(right, modulus, psi)
product_hat = pointwise_multiply(left_hat, right_hat, modulus)
recovered = inverse_ntt_psi(product_hat, modulus, psi)
lanes = [
("left", list(left)),
("right", list(right)),
("left_hat", left_hat),
("right_hat", right_hat),
("pointwise", product_hat),
("inverse", recovered),
]
fig, ax = plt.subplots(figsize=(max(10, len(lanes) * 2.15), max(4.6, len(left) * 0.85)))
ax.set_title(title, fontsize=14, fontweight="bold")
ax.axis("off")
for lane_index, (label, values) in enumerate(lanes):
x = lane_index * 2.2
for row_index, value in enumerate(values):
ax.text(
x,
-row_index,
str(value),
ha="center",
va="center",
fontsize=10,
family="monospace",
bbox={
"boxstyle": "round,pad=0.24",
"facecolor": _value_colors([value])[0],
"edgecolor": "#222222",
"linewidth": 1.0,
},
)
ax.text(x, 1.1, label, ha="center", va="center", fontsize=10, fontweight="bold")
if lane_index < len(lanes) - 1:
ax.annotate(
"",
xy=(x + 1.5, -len(values) / 2 + 0.4),
xytext=(x + 0.6, -len(values) / 2 + 0.4),
arrowprops={"arrowstyle": "->", "color": "#6c757d", "linewidth": 1.8},
)
ax.text(4.4, 1.6, "NTT_psi", ha="center", va="center", fontsize=10, family="monospace")
ax.text(8.8, 1.6, "slotwise *", ha="center", va="center", fontsize=10, family="monospace")
ax.text(11.0, 1.6, "INTT_psi", ha="center", va="center", fontsize=10, family="monospace")
ax.set_xlim(-1.0, (len(lanes) - 1) * 2.2 + 1.1)
ax.set_ylim(-len(left) + 0.2, 2.0)
fig.tight_layout()
return fig
def plot_base_multiply_pair_diagram(
left: Sequence[int],
right: Sequence[int],
*,
zeta: int,
modulus: int,
title: str = "Base Multiplication On A Degree-1 Pair",
):
"""Plot the two-term base multiplication block used in Kyber-style explanations."""
if len(left) != 2 or len(right) != 2:
raise ValueError("plot_base_multiply_pair_diagram expects two 2-entry vectors")
result = base_multiply_pair(left, right, zeta, modulus)
fig, ax = plt.subplots(figsize=(9, 4.6))
ax.set_title(title, fontsize=14, fontweight="bold")
ax.axis("off")
left_x = 0
right_x = 2.8
out_x = 6.8
ys = [0.9, -0.7]
for x, label, values, facecolor in [
(left_x, "left", left, "#edf6f9"),
(right_x, "right", right, "#fff3b0"),
(out_x, "out", result, "#d8f3dc"),
]:
ax.text(x, 1.8, label, ha="center", va="center", fontsize=12, fontweight="bold")
for index, (y, value) in enumerate(zip(ys, values)):
ax.text(
x,
y,
f"{label}[{index}] = {value}",
ha="center",
va="center",
fontsize=10,
family="monospace",
bbox={
"boxstyle": "round,pad=0.3",
"facecolor": facecolor,
"edgecolor": "#222222",
"linewidth": 1.0,
},
)
ax.annotate("", xy=(left_x + 0.9, 0.1), xytext=(out_x - 1.0, 0.9), arrowprops={"arrowstyle": "->", "linewidth": 2.0, "color": "#355070"})
ax.annotate("", xy=(left_x + 0.9, -0.7), xytext=(out_x - 1.0, -0.7), arrowprops={"arrowstyle": "->", "linewidth": 2.0, "color": "#355070"})
ax.annotate("", xy=(right_x + 0.9, 0.1), xytext=(out_x - 1.0, 0.9), arrowprops={"arrowstyle": "->", "linewidth": 2.0, "color": "#6d597a"})
ax.annotate("", xy=(right_x + 0.9, -0.7), xytext=(out_x - 1.0, -0.7), arrowprops={"arrowstyle": "->", "linewidth": 2.0, "color": "#6d597a"})
ax.text(
4.8,
1.05,
f"c0 = a0*b0 + zeta*a1*b1 mod {modulus}\n= {result[0]}",
ha="center",
va="center",
fontsize=10,
family="monospace",
bbox={"boxstyle": "round,pad=0.32", "facecolor": "#ffffff", "edgecolor": "#355070"},
)
ax.text(
4.8,
-1.0,
f"c1 = a0*b1 + a1*b0 mod {modulus}\n= {result[1]}",
ha="center",
va="center",
fontsize=10,
family="monospace",
bbox={"boxstyle": "round,pad=0.32", "facecolor": "#ffffff", "edgecolor": "#6d597a"},
)
ax.text(4.8, 0.0, f"zeta = {zeta}", ha="center", va="center", fontsize=10, family="monospace")
ax.set_xlim(-1.0, 8.2)
ax.set_ylim(-2.0, 2.2)
fig.tight_layout()
return fig
def plot_root_order_comparison(samples: Sequence[tuple[int, int]], title: str = "Root Existence Check"):
"""Plot which moduli allow n-th and 2n-th root stories."""
fig, ax = plt.subplots(figsize=(10, max(4.5, len(samples) * 0.75)))
ax.set_title(title, fontsize=14, fontweight="bold")
ax.axis("off")
headers = ["n", "q", "n | q-1", "2n | q-1"]
x_positions = [0, 2, 4.2, 6.8]
for x, header in zip(x_positions, headers):
ax.text(x, 1.2, header, ha="center", va="center", fontsize=11, fontweight="bold")
for row_index, (n, q) in enumerate(samples):
y = -row_index
statuses = [str(n), str(q), "yes" if (q - 1) % n == 0 else "no", "yes" if (q - 1) % (2 * n) == 0 else "no"]
for x, value in zip(x_positions, statuses):
facecolor = "#d8f3dc" if value == "yes" else "#f5cac3" if value == "no" else "#edf6f9"
ax.text(
x,
y,
value,
ha="center",
va="center",
fontsize=10,
family="monospace",
bbox={"boxstyle": "round,pad=0.25", "facecolor": facecolor, "edgecolor": "#222222"},
)
ax.set_xlim(-1.0, 8.0)
ax.set_ylim(-len(samples) + 0.2, 1.8)
fig.tight_layout()
return fig
def plot_trace_overview(trace: TransformTrace, title: str | None = None):
"""Plot every stage output as a column of values."""
if title is None:

235
tools/render_notebooks.py
View file

@ -805,7 +805,10 @@ def build_bundle_02() -> None:
"demo",
"Inspect ω, ψ, And The Direct Transform Matrix",
"""
from IPython.display import display
from ntt_learning.toy_ntt import find_primitive_root, find_psi, ntt_psi_exponent_grid, ntt_psi_matrix
from ntt_learning.visuals import plot_ntt_psi_exponent_heatmap, plot_ntt_psi_matrix_heatmap
modulus = 17
n = 4
@ -820,6 +823,9 @@ def build_bundle_02() -> None:
print("NTT_psi matrix:")
for row in ntt_psi_matrix(n, modulus, psi):
print(row)
display(plot_ntt_psi_exponent_heatmap(n, title="Exponents 2ij + i for n=4"))
display(plot_ntt_psi_matrix_heatmap(n, modulus, psi, title="Concrete NTT_psi matrix in Z_17"))
""",
),
markdown(
@ -856,7 +862,10 @@ def build_bundle_02() -> None:
"demo",
"Use Direct NTTψ For Negacyclic Multiplication",
"""
from IPython.display import display
from ntt_learning.toy_ntt import find_psi, forward_ntt_psi, inverse_ntt_psi, negacyclic_multiply, pointwise_multiply
from ntt_learning.visuals import plot_transform_pipeline
left = [1, 2, 3, 4]
right = [5, 6, 7, 8]
@ -872,6 +881,7 @@ def build_bundle_02() -> None:
print("pointwise product:", product_hat)
print("inverse of pointwise product:", inverse_ntt_psi(product_hat, modulus, psi))
print("schoolbook negacyclic:", negacyclic_multiply(left, right, n=4, modulus=modulus))
display(plot_transform_pipeline(left, right, modulus=modulus, psi=psi, title="Direct negacyclic multiply pipeline"))
""",
),
markdown(
@ -1160,15 +1170,30 @@ def build_bundle_02() -> None:
"demo",
"Compare Positive-Wrapped And Negative-Wrapped Views",
"""
from IPython.display import display
from ntt_learning.toy_ntt import find_primitive_root, find_psi, forward_ntt, forward_ntt_psi
from ntt_learning.visuals import plot_vector_comparison
signal = [1, 2, 3, 4]
modulus = 17
omega = find_primitive_root(4, modulus)
psi = find_psi(4, modulus)
print("positive-wrapped NTT:", forward_ntt(signal, modulus, omega))
print("negative-wrapped NTT_psi:", forward_ntt_psi(signal, modulus, psi))
positive = forward_ntt(signal, modulus, omega)
negative = forward_ntt_psi(signal, modulus, psi)
print("positive-wrapped NTT:", positive)
print("negative-wrapped NTT_psi:", negative)
display(
plot_vector_comparison(
positive,
negative,
left_label="positive",
right_label="negative",
title="Same signal, different transform story",
)
)
""",
),
markdown(
@ -1191,15 +1216,30 @@ def build_bundle_02() -> None:
"exercise",
"See A Wrong-Root Failure",
"""
from IPython.display import display
from ntt_learning.toy_ntt import find_primitive_root, find_psi, forward_ntt_psi
from ntt_learning.visuals import plot_vector_comparison
signal = [1, 2, 3, 4]
modulus = 17
omega = find_primitive_root(4, modulus)
psi = find_psi(4, modulus)
print("correct psi-based transform:", forward_ntt_psi(signal, modulus, psi))
print("wrongly using omega as if it were psi:", forward_ntt_psi(signal, modulus, omega))
correct = forward_ntt_psi(signal, modulus, psi)
wrong = forward_ntt_psi(signal, modulus, omega)
print("correct psi-based transform:", correct)
print("wrongly using omega as if it were psi:", wrong)
display(
plot_vector_comparison(
wrong,
correct,
left_label="wrong_root",
right_label="correct",
title="Wrong root vs correct root",
)
)
""",
),
markdown(
@ -1275,7 +1315,7 @@ def build_bundle_03() -> None:
from IPython.display import display
from ntt_learning.toy_ntt import fast_ntt_psi_ct_trace, forward_ntt_psi
from ntt_learning.visuals import interactive_trace, plot_trace_overview
from ntt_learning.visuals import interactive_trace, plot_butterfly_network, plot_trace_overview
signal = [1, 2, 3, 4]
modulus = 7681
@ -1286,6 +1326,7 @@ def build_bundle_03() -> None:
print("bit-reversed back to NO:", trace.normal_order_output)
print("direct NTT_psi:", forward_ntt_psi(signal, modulus, psi))
display(plot_trace_overview(trace, title="CT overview for [1,2,3,4]"))
display(plot_butterfly_network(trace, title="Full CT network for [1,2,3,4]"))
display(interactive_trace(trace, title="CT forward trace"))
""",
),
@ -1330,7 +1371,7 @@ def build_bundle_03() -> None:
from IPython.display import display
from ntt_learning.toy_ntt import fast_ntt_psi_ct_trace, find_psi
from ntt_learning.visuals import interactive_trace, plot_trace_overview
from ntt_learning.visuals import interactive_trace, plot_butterfly_network, plot_trace_overview
signal = [0, 1, 2, 3, 4, 5, 6, 7]
modulus = 97
@ -1341,9 +1382,34 @@ def build_bundle_03() -> None:
print("BO output:", trace.raw_output)
print("NO output:", trace.normal_order_output)
display(plot_trace_overview(trace, title="Three CT stages for n=8"))
display(plot_butterfly_network(trace, title="Full CT network for n=8"))
display(interactive_trace(trace, title="n=8 CT stage explorer"))
""",
),
code(
"mandatory",
3,
"demo",
"See BO Output And NO Output Side By Side",
"""
from IPython.display import display
from ntt_learning.toy_ntt import fast_ntt_psi_ct_trace
from ntt_learning.visuals import plot_bit_reversal_mapping, plot_vector_comparison
trace = fast_ntt_psi_ct_trace([1, 2, 3, 4], 7681, 1925)
display(
plot_vector_comparison(
trace.raw_output,
trace.normal_order_output,
left_label="BO",
right_label="NO",
title="Same CT values, different ordering",
)
)
display(plot_bit_reversal_mapping(4, title="Why BO becomes NO after bit-reversal"))
""",
),
markdown(
"mandatory",
2,
@ -1602,13 +1668,15 @@ def build_bundle_03() -> None:
from IPython.display import display
from ntt_learning.toy_ntt import fast_ntt_psi_ct_trace
from ntt_learning.visuals import plot_trace_overview
from ntt_learning.visuals import plot_butterfly_network, plot_trace_overview
trace_a = fast_ntt_psi_ct_trace([1, 2, 3, 4], 7681, 1925)
trace_b = fast_ntt_psi_ct_trace([5, 6, 7, 8], 7681, 1925)
display(plot_trace_overview(trace_a, title="CT trace A"))
display(plot_trace_overview(trace_b, title="CT trace B"))
display(plot_butterfly_network(trace_a, title="CT network A"))
display(plot_butterfly_network(trace_b, title="CT network B"))
""",
),
markdown(
@ -1631,7 +1699,10 @@ def build_bundle_03() -> None:
"exercise",
"See A Wrong-Order Comparison Failure",
"""
from IPython.display import display
from ntt_learning.toy_ntt import fast_ntt_psi_ct_trace, forward_ntt_psi
from ntt_learning.visuals import plot_vector_comparison
trace = fast_ntt_psi_ct_trace([1, 2, 3, 4], 7681, 1925)
direct = forward_ntt_psi([1, 2, 3, 4], 7681, 1925)
@ -1640,6 +1711,24 @@ def build_bundle_03() -> None:
print(trace.raw_output, direct)
print("correct comparison: CT NO output vs direct NO output")
print(trace.normal_order_output, direct)
display(
plot_vector_comparison(
trace.raw_output,
direct,
left_label="CT_BO",
right_label="direct_NO",
title="Wrong comparison: BO against NO",
)
)
display(
plot_vector_comparison(
trace.normal_order_output,
direct,
left_label="CT_NO",
right_label="direct_NO",
title="Correct comparison after reordering",
)
)
""",
),
markdown(
@ -1710,7 +1799,7 @@ def build_bundle_04() -> None:
from IPython.display import display
from ntt_learning.toy_ntt import fast_intt_psi_gs_trace
from ntt_learning.visuals import interactive_trace, plot_trace_overview
from ntt_learning.visuals import interactive_trace, plot_butterfly_network, plot_trace_overview, plot_vector_comparison
bo_input = [1467, 3471, 2807, 7621]
trace = fast_intt_psi_gs_trace(bo_input, 7681, 1925)
@ -1718,6 +1807,16 @@ def build_bundle_04() -> None:
print("unscaled NO output:", trace.raw_output)
print("scaled NO output:", trace.scaled_output)
display(plot_trace_overview(trace, title="GS overview for the n=4 paper example"))
display(plot_butterfly_network(trace, title="Full GS network for the n=4 paper example"))
display(
plot_vector_comparison(
trace.raw_output,
trace.scaled_output,
left_label="unscaled",
right_label="scaled",
title="Why the final n^-1 scaling matters",
)
)
display(interactive_trace(trace, title="GS inverse trace"))
""",
),
@ -1751,7 +1850,10 @@ def build_bundle_04() -> None:
"demo",
"Full Forward And Inverse Round Trip",
"""
from IPython.display import display
from ntt_learning.toy_ntt import fast_intt_psi_gs_trace, fast_ntt_psi_ct_trace
from ntt_learning.visuals import plot_vector_comparison
signal = [1, 2, 3, 4]
forward_trace = fast_ntt_psi_ct_trace(signal, 7681, 1925)
@ -1759,6 +1861,15 @@ def build_bundle_04() -> None:
print("forward BO output:", forward_trace.raw_output)
print("inverse scaled output:", inverse_trace.scaled_output)
display(
plot_vector_comparison(
signal,
inverse_trace.scaled_output,
left_label="original",
right_label="recovered",
title="Forward CT followed by inverse GS",
)
)
""",
),
markdown(
@ -2008,7 +2119,10 @@ def build_bundle_04() -> None:
"demo",
"See CT Output Feed GS Input",
"""
from IPython.display import display
from ntt_learning.toy_ntt import fast_intt_psi_gs_trace, fast_ntt_psi_ct_trace
from ntt_learning.visuals import plot_vector_comparison
signal = [5, 6, 7, 8]
forward_trace = fast_ntt_psi_ct_trace(signal, 7681, 1925)
@ -2016,6 +2130,15 @@ def build_bundle_04() -> None:
print("CT BO output:", forward_trace.raw_output)
print("GS scaled output:", inverse_trace.scaled_output)
display(
plot_vector_comparison(
signal,
inverse_trace.scaled_output,
left_label="start",
right_label="after_CT_then_GS",
title="CT output cleanly feeds GS input",
)
)
""",
),
markdown(
@ -2038,11 +2161,23 @@ def build_bundle_04() -> None:
"exercise",
"See A Missing-Scale Failure",
"""
from IPython.display import display
from ntt_learning.toy_ntt import fast_intt_psi_gs_trace
from ntt_learning.visuals import plot_vector_comparison
trace = fast_intt_psi_gs_trace([1467, 3471, 2807, 7621], 7681, 1925)
print("unscaled:", trace.raw_output)
print("scaled:", trace.scaled_output)
display(
plot_vector_comparison(
trace.raw_output,
trace.scaled_output,
left_label="missing_scale",
right_label="correct",
title="Missing n^-1 scale vs correct output",
)
)
""",
),
markdown(
@ -2117,7 +2252,10 @@ def build_bundle_05() -> None:
"demo",
"Check The Kyber Root Reality Directly",
"""
from IPython.display import display
from ntt_learning.toy_ntt import find_primitive_root
from ntt_learning.visuals import plot_root_order_comparison
print("3329 - 1 =", 3329 - 1)
print("primitive 256-th root in Z_3329:", find_primitive_root(256, 3329))
@ -2125,6 +2263,13 @@ def build_bundle_05() -> None:
find_primitive_root(512, 3329)
except Exception as exc:
print("512-th root fails exactly because:", exc)
display(
plot_root_order_comparison(
[(4, 17), (4, 13), (8, 97), (256, 3329)],
title="Which moduli allow n-th and 2n-th root stories?",
)
)
""",
),
markdown(
@ -2171,7 +2316,10 @@ def build_bundle_05() -> None:
"demo",
"See A Toy Base Multiplication",
"""
from IPython.display import display
from ntt_learning.toy_ntt import base_multiply_pair
from ntt_learning.visuals import plot_base_multiply_pair_diagram
left = [7, 11]
right = [5, 13]
@ -2184,6 +2332,15 @@ def build_bundle_05() -> None:
print("raw degree-2 product:", raw)
print("reduce with x^2 = zeta:", reduced)
print("base_multiply_pair:", base_multiply_pair(left, right, zeta, modulus))
display(
plot_base_multiply_pair_diagram(
left,
right,
zeta=zeta,
modulus=modulus,
title="Why Kyber-style multiplication stays in 2-slot blocks",
)
)
""",
),
markdown(
@ -2410,13 +2567,23 @@ def build_bundle_05() -> None:
"demo",
"Toy Full ψ Story vs Kyber Root Reality",
"""
from IPython.display import display
from ntt_learning.toy_ntt import find_psi
from ntt_learning.visuals import plot_root_order_comparison
print("toy n=4, q=17 has psi:", find_psi(4, 17))
try:
find_psi(256, 3329)
except Exception as exc:
print("Kyber v3 does not have that full psi story:", exc)
display(
plot_root_order_comparison(
[(4, 17), (256, 7681), (256, 3329)],
title="Toy full psi story vs Kyber modulus reality",
)
)
""",
),
markdown(
@ -2438,9 +2605,14 @@ def build_bundle_05() -> None:
"exercise",
"See The Exact Obstruction Again",
"""
from IPython.display import display
from ntt_learning.visuals import plot_root_order_comparison
q = 3329
n = 256
print({"q_minus_1": q - 1, "n": n, "2n": 2 * n, "q_minus_1_mod_n": (q - 1) % n, "q_minus_1_mod_2n": (q - 1) % (2 * n)})
display(plot_root_order_comparison([(256, 3329)], title="Kyber v3 obstruction in one row"))
""",
),
markdown(
@ -2512,6 +2684,8 @@ def build_bundle_06() -> None:
"demo",
"See Four Failure Modes Side By Side",
"""
from IPython.display import display
from ntt_learning.toy_ntt import (
fast_intt_psi_gs_trace,
fast_ntt_psi_ct_trace,
@ -2519,6 +2693,7 @@ def build_bundle_06() -> None:
negacyclic_reduce,
schoolbook_convolution,
)
from ntt_learning.visuals import plot_vector_comparison
signal = [1, 2, 3, 4]
forward_trace = fast_ntt_psi_ct_trace(signal, 7681, 1925)
@ -2536,6 +2711,33 @@ def build_bundle_06() -> None:
print("correct NO output:", forward_trace.normal_order_output)
print("missing final scaling:", wrong_scale)
print("wrong root in direct transform:", wrong_root)
display(
plot_vector_comparison(
wrong_sign,
negacyclic_reduce(raw, n=4),
left_label="wrong_sign",
right_label="correct_sign",
title="Wrong sign vs correct negacyclic fold",
)
)
display(
plot_vector_comparison(
wrong_order,
forward_trace.normal_order_output,
left_label="wrong_order",
right_label="correct_order",
title="Wrong BO/NO comparison",
)
)
display(
plot_vector_comparison(
wrong_scale,
inverse_trace.scaled_output,
left_label="missing_scale",
right_label="correct_scale",
title="Missing scale vs corrected inverse",
)
)
""",
),
code(
@ -2548,6 +2750,7 @@ def build_bundle_06() -> None:
from IPython.display import display
from ntt_learning.toy_ntt import fast_intt_psi_gs_trace, fast_ntt_psi_ct_trace, forward_ntt_psi
from ntt_learning.visuals import plot_vector_comparison
forward_trace = fast_ntt_psi_ct_trace([1, 2, 3, 4], 7681, 1925)
inverse_trace = fast_intt_psi_gs_trace(forward_trace.raw_output, 7681, 1925)
@ -2559,9 +2762,25 @@ def build_bundle_06() -> None:
"scaled": list(inverse_trace.scaled_output),
"wrong_root": forward_ntt_psi([1, 2, 3, 4], 7681, 3383),
}
references = {
"wrong_order": list(forward_trace.normal_order_output),
"correct_order": list(forward_trace.normal_order_output),
"missing_scale": list(inverse_trace.scaled_output),
"scaled": list(inverse_trace.scaled_output),
"wrong_root": list(forward_ntt_psi([1, 2, 3, 4], 7681, 1925)),
}
def preview(mode="wrong_order"):
print(mode, "->", failures[mode])
display(
plot_vector_comparison(
failures[mode],
references[mode],
left_label=mode,
right_label="reference",
title=f"{mode} compared with the correct reference",
)
)
display(widgets.interact(preview, mode=sorted(failures)))
""",