re/zvp.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3ff024c8-7e90-4094-ab07-a051a601d9cd",
   "metadata": {},
   "source": [
    "# ZVP-based reverse-engineering\n",
    "This notebook showcases the ZVP-based reverse-engineering technique for addition formulas.\n",
    "\n",
    " - [Exploration](#Exploration)\n",
    " - [Reverse-engineering](#Reverse-engineering)\n",
    "   - [Computing addition chains](#Computing-addition-chains)\n",
    "   - [Loading the formulas](#Loading-the-formulas)\n",
    "   - [Computing the factor sets](#Computing-the-factor-sets)\n",
    "   - [Accumulating polynomials](#Accumulating-polynomials)\n",
    "   - [Computing ZVP points](#Computing-ZVP-points)\n",
    "   - [Remapping](#Remapping)\n",
    "   - [Distinguishing map and distinguishing tree building](#Distinguishing-map-and-distinguishing-tree-building)\n",
    "   - [Evaluation](#Evaluation)\n"
   ]
  },
  {
   "cell_type": "code",
   "id": "17a9e580-5f1e-45c9-8afd-fc35e833b0b0",
   "metadata": {},
   "source": [
    "import io\n",
    "import holoviews as hv\n",
    "import pickle\n",
    "import multiprocessing\n",
    "import inspect\n",
    "import tempfile\n",
    "import sys\n",
    "import re\n",
    "from collections import Counter\n",
    "from matplotlib import pyplot as plt\n",
    "from sympy import FF, ZZ, symbols, Poly\n",
    "from contextlib import contextmanager\n",
    "from importlib import import_module, invalidate_caches\n",
    "from pathlib import Path\n",
    "from itertools import product\n",
    "from IPython.display import display\n",
    "from tqdm.notebook import tqdm, trange\n",
    "from anytree import LevelOrderIter, PreOrderIter\n",
    "\n",
    "\n",
    "from pyecsca.ec.model import ShortWeierstrassModel\n",
    "from pyecsca.ec.coordinates import AffineCoordinateModel\n",
    "from pyecsca.ec.curve import EllipticCurve\n",
    "from pyecsca.ec.params import DomainParameters, load_params_ecgen\n",
    "from pyecsca.ec.formula import AdditionFormula, DoublingFormula\n",
    "from pyecsca.ec.point import Point\n",
    "from pyecsca.ec.mod import mod, SymbolicMod\n",
    "from pyecsca.ec.mult import LTRMultiplier, AccumulationOrder\n",
    "from pyecsca.ec.error import UnsatisfiedAssumptionError\n",
    "from pyecsca.sca.re.tree import Map, Tree\n",
    "from pyecsca.sca.re.zvp import zvp_points, compute_factor_set, addition_chain, eliminate_y\n",
    "from pyecsca.misc.cfg import getconfig\n",
    "from pyecsca.misc.utils import TaskExecutor\n",
    "\n",
    "from eval import (eval_tree_symmetric1, eval_tree_asymmetric1, eval_tree_binomial1,\n",
    "                    success_rate_symmetric, success_rate_asymmetric, success_rate_binomial,\n",
    "                    query_rate_symmetric, query_rate_asymmetric, query_rate_binomial,\n",
    "                    amount_rate_symmetric, amount_rate_asymmetric, amount_rate_binomial,\n",
    "                    precise_rate_symmetric, precise_rate_asymmetric, precise_rate_binomial,\n",
    "                    success_rate_vs_majority_symmetric, success_rate_vs_majority_asymmetric,\n",
    "                    success_rate_vs_query_rate_symmetric, load, store)\n",
    "\n",
    "\n",
    "# Allow to use \"spawn\" multiprocessing method for function defined in a Jupyter notebook.\n",
    "# https://neuromancer.sk/article/35\n",
    "@contextmanager\n",
    "def enable_spawn(func):\n",
    "    invalidate_caches()\n",
    "    source = inspect.getsource(func)\n",
    "    with tempfile.NamedTemporaryFile(suffix=\".py\", mode=\"w\") as f:\n",
    "        f.write(source)\n",
    "        f.flush()\n",
    "        path = Path(f.name)\n",
    "        directory = str(path.parent)\n",
    "        sys.path.append(directory)\n",
    "        module = import_module(str(path.stem))\n",
    "        yield getattr(module, func.__name__)\n",
    "        sys.path.remove(directory)\n",
    "\n",
    "spawn_context = multiprocessing.get_context(\"spawn\")\n",
    "\n",
    "%matplotlib ipympl\n",
    "hv.extension(\"bokeh\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "faca31e8-fe51-4017-a41d-ba31b15c548d",
   "metadata": {},
   "source": [
    "model = ShortWeierstrassModel()\n",
    "coordsaff = AffineCoordinateModel(model)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "60547b40-9b8e-409b-8b7e-b3dc266c01d4",
   "metadata": {},
   "source": [
    "## Exploration\n",
    "First lets explore the behavior of addition formulas. The following two cells pick a coordinate model along with some formulas and symbolically unroll a scalar multiplication (assuming a simple LTR multiplier)."
   ]
  },
  {
   "cell_type": "code",
   "id": "8de124f1-9498-4e55-8b7f-1bd291ccf3fe",
   "metadata": {},
   "source": [
    "which = \"jacobian\"\n",
    "coords = model.coordinates[which]"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "2193ab53-9dc9-4ba4-89d3-e331e5cc20de",
   "metadata": {},
   "source": [
    "getconfig().ec.mod_implementation = \"symbolic\"\n",
    "x, y, z = symbols(\"x y z\")\n",
    "\n",
    "# A 64-bit prime order curve for testing things out\n",
    "p = 0xc50de883f0e7b167\n",
    "field = FF(p)\n",
    "a = SymbolicMod(Poly(0x4833d7aa73fa6694, x, y, z, domain=field), p)\n",
    "b = SymbolicMod(Poly(0xa6c44a61c5323f6a, x, y, z, domain=field), p)\n",
    "gx = SymbolicMod(Poly(0x5fd1f7d38d4f2333, x, y, z, domain=field), p)\n",
    "gy = SymbolicMod(Poly(0x21f43957d7e20ceb, x, y, z, domain=field), p)\n",
    "n = 0xc50de885003b80eb\n",
    "h = 1\n",
    "\n",
    "infty = Point(coords, X=mod(0, p), Y=mod(1, p), Z=mod(0, p))\n",
    "g = Point(coords, X=gx, Y=gy, Z=mod(1, p))\n",
    "\n",
    "curve = EllipticCurve(model, coords, p, infty, dict(a=a,b=b))\n",
    "params = DomainParameters(curve, g, n, h)\n",
    "\n",
    "\n",
    "add = coords.formulas[\"add-2007-bl\"]\n",
    "dbl = coords.formulas[\"dbl-2007-bl\"]\n",
    "mult = LTRMultiplier(add, dbl, None, False, AccumulationOrder.PeqRP, True, True)\n",
    "\n",
    "\n",
    "point = Point(coords,\n",
    "              X=SymbolicMod(Poly(x, x, y, z, domain=field), params.curve.prime),\n",
    "              Y=SymbolicMod(Poly(y, x, y, z, domain=field), params.curve.prime),\n",
    "              Z=SymbolicMod(Poly(z, x, y, z, domain=field), params.curve.prime))\n",
    "mult.init(params, point)\n",
    "res = mult.multiply(5)\n",
    "\n",
    "x_poly = Poly(res.X.x, domain=field)\n",
    "y_poly = Poly(res.Y.x, domain=field)\n",
    "z_poly = Poly(res.Z.x, domain=field)\n",
    "display(x_poly, y_poly, z_poly)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "0bcc8b9e-39ad-4b53-8f45-a7552b05baa2",
   "metadata": {},
   "source": [
    "The result is a Point with coordinates that are polynomials in the input coordinates and curve parameters. We now switch back to concrete representation."
   ]
  },
  {
   "cell_type": "code",
   "id": "d8ea0c4d-86e1-46af-ac20-569e6ef5439d",
   "metadata": {},
   "source": [
    "cfg = getconfig()\n",
    "cfg.ec.mod_implementation = \"gmp\""
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "c3b3ae94-242b-4f36-81e3-e8a64b5af33c",
   "metadata": {},
   "source": [
    "## Reverse-engineering\n",
    "Now, lets look at using the ZVP attack for reverse-engineering. First pick 10 curves per group, some random some with $a \\in \\{0, -1, -3 \\}$. The curves are otherwise not special in any way and just serve to randomize the process, as the existence of ZVP points for a given intermediate value polynomial depends on the curve."
   ]
  },
  {
   "cell_type": "code",
   "id": "0252f8fa-012c-453c-ba23-b40f9686fcce",
   "metadata": {},
   "source": [
    "curves = list(map(lambda spec: load_params_ecgen(io.BytesIO(spec.encode()), \"affine\"), [\n",
    "    # Random\n",
    "    \"\"\"[{\"field\":{\"p\":\"0xddf438409fc35161\"},\"a\":\"0x94d919b72f7dc6d8\",\"b\":\"0x9f39032abb23f62a\",\"order\":\"0xddf4383ffa8e6de7\",\"subgroups\":[{\"x\":\"0xd5673b3fe176fc6b\",\"y\":\"0x2d5b0a5bb2141317\",\"order\":\"0xddf4383ffa8e6de7\",\"cofactor\":\"0x1\",\"points\":[{\"x\":\"0xd5673b3fe176fc6b\",\"y\":\"0x2d5b0a5bb2141317\",\"order\":\"0xddf4383ffa8e6de7\"}]}]}]\"\"\",\n",
    "    \"\"\"[{\"field\":{\"p\":\"0xa42c1467a1ed04f3\"},\"a\":\"0x55d07340a4572f2d\",\"b\":\"0x0a938c37dfb0b6d5\",\"order\":\"0xa42c14689284d3a7\",\"subgroups\":[{\"x\":\"0x8633981c83ed43a2\",\"y\":\"0x7b5374e9d7997199\",\"order\":\"0xa42c14689284d3a7\",\"cofactor\":\"0x1\",\"points\":[{\"x\":\"0x8633981c83ed43a2\",\"y\":\"0x7b5374e9d7997199\",\"order\":\"0xa42c14689284d3a7\"}]}]}]\"\"\",\n",
    "    \"\"\"[{\"field\":{\"p\":\"0xea0d9cead19016ab\"},\"a\":\"0xcbbfe501c4ef6d92\",\"b\":\"0x5762de777a6d9178\",\"order\":\"0xea0d9cea8cd2c857\",\"subgroups\":[{\"x\":\"0xe7daa3e061c3111b\",\"y\":\"0x56ee59a6845c5e93\",\"order\":\"0xea0d9cea8cd2c857\",\"cofactor\":\"0x1\",\"points\":[{\"x\":\"0xe7daa3e061c3111b\",\"y\":\"0x56ee59a6845c5e93\",\"order\":\"0xea0d9cea8cd2c857\"}]}]}]\"\"\",\n",
    "    \"\"\"[{\"field\":{\"p\":\"0x9c7e7216decb71a7\"},\"a\":\"0x324ef48887401a87\",\"b\":\"0x3ce6f35a00280102\",\"order\":\"0x9c7e72175ebfe709\",\"subgroups\":[{\"x\":\"0x34683229b405418d\",\"y\":\"0x308c923cae004514\",\"order\":\"0x9c7e72175ebfe709\",\"cofactor\":\"0x1\",\"points\":[{\"x\":\"0x34683229b405418d\",\"y\":\"0x308c923cae004514\",\"order\":\"0x9c7e72175ebfe709\"}]}]}]\"\"\",\n",
    "    \"\"\"[{\"field\":{\"p\":\"0xeb5779f0bbf1ef5b\"},\"a\":\"0x2419e8bbc7b5f8f2\",\"b\":\"0xe74e5d3064a4f2e3\",\"order\":\"0xeb5779f21320c2e9\",\"subgroups\":[{\"x\":\"0x3b6c269560abeb00\",\"y\":\"0x29d157628e75e1c0\",\"order\":\"0xeb5779f21320c2e9\",\"cofactor\":\"0x1\",\"points\":[{\"x\":\"0x3b6c269560abeb00\",\"y\":\"0x29d157628e75e1c0\",\"order\":\"0xeb5779f21320c2e9\"}]}]}]\"\"\",\n",
    "    \"\"\"[{\"field\":{\"p\":\"0x97b6ea097868b95d\"},\"a\":\"0x550a41d65e4bcd13\",\"b\":\"0x47c5e527113b261c\",\"order\":\"0x97b6ea094947a76b\",\"subgroups\":[{\"x\":\"0x1e669fe19c865bd9\",\"y\":\"0x05a6bb891920440f\",\"order\":\"0x97b6ea094947a76b\",\"cofactor\":\"0x1\",\"points\":[{\"x\":\"0x1e669fe19c865bd9\",\"y\":\"0x05a6bb891920440f\",\"order\":\"0x97b6ea094947a76b\"}]}]}]\"\"\",\n",
    "    \"\"\"[{\"field\":{\"p\":\"0xa00629e6522032f7\"},\"a\":\"0x896f04a7ae302922\",\"b\":\"0x6bc03365b1f1cb50\",\"order\":\"0xa00629e5c03cf913\",\"subgroups\":[{\"x\":\"0x14b7b48954936d4e\",\"y\":\"0x670dc776273bf899\",\"order\":\"0xa00629e5c03cf913\",\"cofactor\":\"0x1\",\"points\":[{\"x\":\"0x14b7b48954936d4e\",\"y\":\"0x670dc776273bf899\",\"order\":\"0xa00629e5c03cf913\"}]}]}]\"\"\",\n",
    "    \"\"\"[{\"field\":{\"p\":\"0xd47ec1d03a62686d\"},\"a\":\"0xd00a3ee0f5c86b02\",\"b\":\"0x457a5b6c47db38d8\",\"order\":\"0xd47ec1d107db7d6f\",\"subgroups\":[{\"x\":\"0x41ebc3b763f3cd1b\",\"y\":\"0x3d6925f214620e0c\",\"order\":\"0xd47ec1d107db7d6f\",\"cofactor\":\"0x1\",\"points\":[{\"x\":\"0x41ebc3b763f3cd1b\",\"y\":\"0x3d6925f214620e0c\",\"order\":\"0xd47ec1d107db7d6f\"}]}]}]\"\"\",\n",
    "    \"\"\"[{\"field\":{\"p\":\"0xb1c9115c6f40d755\"},\"a\":\"0x79d3ceefafc44ce9\",\"b\":\"0x8316af84264df42b\",\"order\":\"0xb1c9115d17f84a45\",\"subgroups\":[{\"x\":\"0x8b0a274089b53fe5\",\"y\":\"0x3508d33c4beba5ad\",\"order\":\"0xb1c9115d17f84a45\",\"cofactor\":\"0x1\",\"points\":[{\"x\":\"0x8b0a274089b53fe5\",\"y\":\"0x3508d33c4beba5ad\",\"order\":\"0xb1c9115d17f84a45\"}]}]}]\"\"\",\n",
    "    \"\"\"[{\"field\":{\"p\":\"0x8f738fda18cd5dff\"},\"a\":\"0x4747f2f9b8628cbf\",\"b\":\"0x586cdb9378a1389f\",\"order\":\"0x8f738fd8fc7ebed3\",\"subgroups\":[{\"x\":\"0x7ad306c73b64c1b5\",\"y\":\"0x69e3ca555190da4b\",\"order\":\"0x8f738fd8fc7ebed3\",\"cofactor\":\"0x1\",\"points\":[{\"x\":\"0x7ad306c73b64c1b5\",\"y\":\"0x69e3ca555190da4b\",\"order\":\"0x8f738fd8fc7ebed3\"}]}]}]\"\"\",\n",
    "    # a = -1\n",
    "    \"\"\"[{\"field\":{\"p\":\"0xcfef393139c3007f\"},\"a\":\"0xcfef393139c3007e\",\"b\":\"0x950312812acb155f\",\"order\":\"0xcfef39320179387b\",\"subgroups\":[{\"x\":\"0xae2d2f58ca5b5cf7\",\"y\":\"0xc3a4bf3a1dc10005\",\"order\":\"0xcfef39320179387b\",\"cofactor\":\"0x1\",\"points\":[{\"x\":\"0xae2d2f58ca5b5cf7\",\"y\":\"0xc3a4bf3a1dc10005\",\"order\":\"0xcfef39320179387b\"}]}]}]\"\"\",\n",
    "    \"\"\"[{\"field\":{\"p\":\"0xb0461c0e4946cbd5\"},\"a\":\"0xb0461c0e4946cbd4\",\"b\":\"0x082c3722016def51\",\"order\":\"0xb0461c0e07e3e1bf\",\"subgroups\":[{\"x\":\"0x5142200263be1fe3\",\"y\":\"0x14984b7551ed21a9\",\"order\":\"0xb0461c0e07e3e1bf\",\"cofactor\":\"0x1\",\"points\":[{\"x\":\"0x5142200263be1fe3\",\"y\":\"0x14984b7551ed21a9\",\"order\":\"0xb0461c0e07e3e1bf\"}]}]}]\"\"\",\n",
    "    \"\"\"[{\"field\":{\"p\":\"0xeff607c2dc4f278b\"},\"a\":\"0xeff607c2dc4f278a\",\"b\":\"0x26fd03674f5092d2\",\"order\":\"0xeff607c30ab8c50d\",\"subgroups\":[{\"x\":\"0x004d4a5a9bb849fe\",\"y\":\"0x80eb7ef89110c149\",\"order\":\"0xeff607c30ab8c50d\",\"cofactor\":\"0x1\",\"points\":[{\"x\":\"0x004d4a5a9bb849fe\",\"y\":\"0x80eb7ef89110c149\",\"order\":\"0xeff607c30ab8c50d\"}]}]}]\"\"\",\n",
    "    \"\"\"[{\"field\":{\"p\":\"0xedc14fda51686379\"},\"a\":\"0xedc14fda51686378\",\"b\":\"0xb3973a86901e3364\",\"order\":\"0xedc14fda0cdbc199\",\"subgroups\":[{\"x\":\"0xc76f0776feb59336\",\"y\":\"0x625adaf0fb44ab9f\",\"order\":\"0xedc14fda0cdbc199\",\"cofactor\":\"0x1\",\"points\":[{\"x\":\"0xc76f0776feb59336\",\"y\":\"0x625adaf0fb44ab9f\",\"order\":\"0xedc14fda0cdbc199\"}]}]}]\"\"\",\n",
    "    \"\"\"[{\"field\":{\"p\":\"0xfc6ee07288f1b78f\"},\"a\":\"0xfc6ee07288f1b78e\",\"b\":\"0xe18821a83ca2ca30\",\"order\":\"0xfc6ee0713e07f37f\",\"subgroups\":[{\"x\":\"0x339d01a4b0db428e\",\"y\":\"0x68100d42e5ffd979\",\"order\":\"0xfc6ee0713e07f37f\",\"cofactor\":\"0x1\",\"points\":[{\"x\":\"0x339d01a4b0db428e\",\"y\":\"0x68100d42e5ffd979\",\"order\":\"0xfc6ee0713e07f37f\"}]}]}]\"\"\",\n",
    "    \"\"\"[{\"field\":{\"p\":\"0xa03c03a0072f69b3\"},\"a\":\"0xa03c03a0072f69b2\",\"b\":\"0x3208ecebb633b82c\",\"order\":\"0xa03c039ff31e37a7\",\"subgroups\":[{\"x\":\"0x8134208d53e6f6c0\",\"y\":\"0x6245db54032630a6\",\"order\":\"0xa03c039ff31e37a7\",\"cofactor\":\"0x1\",\"points\":[{\"x\":\"0x8134208d53e6f6c0\",\"y\":\"0x6245db54032630a6\",\"order\":\"0xa03c039ff31e37a7\"}]}]}]\"\"\",\n",
    "    \"\"\"[{\"field\":{\"p\":\"0xbc8c6e7ce26746d9\"},\"a\":\"0xbc8c6e7ce26746d8\",\"b\":\"0xb7e2b4fb2d769c4e\",\"order\":\"0xbc8c6e7ba032dda7\",\"subgroups\":[{\"x\":\"0x8e3c9cd771e7ffd8\",\"y\":\"0x4dd02403ca890c5a\",\"order\":\"0xbc8c6e7ba032dda7\",\"cofactor\":\"0x1\",\"points\":[{\"x\":\"0x8e3c9cd771e7ffd8\",\"y\":\"0x4dd02403ca890c5a\",\"order\":\"0xbc8c6e7ba032dda7\"}]}]}]\"\"\",\n",
    "    \"\"\"[{\"field\":{\"p\":\"0x9ccda4c062b95787\"},\"a\":\"0x9ccda4c062b95786\",\"b\":\"0x31fcbb278951e3b9\",\"order\":\"0x9ccda4bfae73e4f5\",\"subgroups\":[{\"x\":\"0x303ac583c81644e3\",\"y\":\"0x76713f6f470e94a0\",\"order\":\"0x9ccda4bfae73e4f5\",\"cofactor\":\"0x1\",\"points\":[{\"x\":\"0x303ac583c81644e3\",\"y\":\"0x76713f6f470e94a0\",\"order\":\"0x9ccda4bfae73e4f5\"}]}]}]\"\"\",\n",
    "    \"\"\"[{\"field\":{\"p\":\"0xa339e3745518416f\"},\"a\":\"0xa339e3745518416e\",\"b\":\"0x52d39a67f2401673\",\"order\":\"0xa339e3743950389b\",\"subgroups\":[{\"x\":\"0x6b8986f706afac58\",\"y\":\"0x5c901b1afa0b64da\",\"order\":\"0xa339e3743950389b\",\"cofactor\":\"0x1\",\"points\":[{\"x\":\"0x6b8986f706afac58\",\"y\":\"0x5c901b1afa0b64da\",\"order\":\"0xa339e3743950389b\"}]}]}]\"\"\",\n",
    "    \"\"\"[{\"field\":{\"p\":\"0x8b7d2baee4e47311\"},\"a\":\"0x8b7d2baee4e47310\",\"b\":\"0x21ab23afb5a9e2ca\",\"order\":\"0x8b7d2baf201f2bdd\",\"subgroups\":[{\"x\":\"0x797c1dec0d73ec64\",\"y\":\"0x28f90926ea9c6b33\",\"order\":\"0x8b7d2baf201f2bdd\",\"cofactor\":\"0x1\",\"points\":[{\"x\":\"0x797c1dec0d73ec64\",\"y\":\"0x28f90926ea9c6b33\",\"order\":\"0x8b7d2baf201f2bdd\"}]}]}]\"\"\",\n",
    "    # a = -3\n",
    "    \"\"\"[{\"field\":{\"p\":\"0x8d79ca36cee026a7\"},\"a\":\"0x8d79ca36cee026a4\",\"b\":\"0x0478c1f80ce2c9c6\",\"order\":\"0x8d79ca35a428c76f\",\"subgroups\":[{\"x\":\"0x2e94a3e38f8b345e\",\"y\":\"0x83e6c6f0cb8f69c4\",\"order\":\"0x8d79ca35a428c76f\",\"cofactor\":\"0x1\",\"points\":[{\"x\":\"0x2e94a3e38f8b345e\",\"y\":\"0x83e6c6f0cb8f69c4\",\"order\":\"0x8d79ca35a428c76f\"}]}]}]\"\"\",\n",
    "    \"\"\"[{\"field\":{\"p\":\"0x48e1a125250323a7\"},\"a\":\"0x48e1a125250323a4\",\"b\":\"0x02a4d99f41d23210\",\"order\":\"0x48e1a124f895db6d\",\"subgroups\":[{\"x\":\"0x409e15d65fcae55a\",\"y\":\"0x207e142056d62d07\",\"order\":\"0x48e1a124f895db6d\",\"cofactor\":\"0x1\",\"points\":[{\"x\":\"0x409e15d65fcae55a\",\"y\":\"0x207e142056d62d07\",\"order\":\"0x48e1a124f895db6d\"}]}]}]\"\"\",\n",
    "    \"\"\"[{\"field\":{\"p\":\"0xcb5aa8a7a10aa06b\"},\"a\":\"0xcb5aa8a7a10aa068\",\"b\":\"0x31fe9c57c570174f\",\"order\":\"0xcb5aa8a6cf812191\",\"subgroups\":[{\"x\":\"0x84c75d46fc687ff1\",\"y\":\"0x7424362ac73df187\",\"order\":\"0xcb5aa8a6cf812191\",\"cofactor\":\"0x1\",\"points\":[{\"x\":\"0x84c75d46fc687ff1\",\"y\":\"0x7424362ac73df187\",\"order\":\"0xcb5aa8a6cf812191\"}]}]}]\"\"\",\n",
    "    \"\"\"[{\"field\":{\"p\":\"0xba965ca9c8aa0a1b\"},\"a\":\"0xba965ca9c8aa0a18\",\"b\":\"0x676535a1eaf5c605\",\"order\":\"0xba965caae5741b6f\",\"subgroups\":[{\"x\":\"0x313d58c47b8ed95f\",\"y\":\"0x991ba98cbbb0fe9f\",\"order\":\"0xba965caae5741b6f\",\"cofactor\":\"0x1\",\"points\":[{\"x\":\"0x313d58c47b8ed95f\",\"y\":\"0x991ba98cbbb0fe9f\",\"order\":\"0xba965caae5741b6f\"}]}]}]\"\"\",\n",
    "    \"\"\"[{\"field\":{\"p\":\"0xbfb7747454a17d15\"},\"a\":\"0xbfb7747454a17d12\",\"b\":\"0x611b69881db4ce69\",\"order\":\"0xbfb7747547fd57d3\",\"subgroups\":[{\"x\":\"0x3385044d698640fc\",\"y\":\"0x50cee623251b559e\",\"order\":\"0xbfb7747547fd57d3\",\"cofactor\":\"0x1\",\"points\":[{\"x\":\"0x3385044d698640fc\",\"y\":\"0x50cee623251b559e\",\"order\":\"0xbfb7747547fd57d3\"}]}]}]\"\"\",\n",
    "    \"\"\"[{\"field\":{\"p\":\"0x99235f1ed44b3959\"},\"a\":\"0x99235f1ed44b3956\",\"b\":\"0x5d8dda19dbe804d4\",\"order\":\"0x99235f1d975f376d\",\"subgroups\":[{\"x\":\"0x4fed262974c1d800\",\"y\":\"0x27590c454edd59ca\",\"order\":\"0x99235f1d975f376d\",\"cofactor\":\"0x1\",\"points\":[{\"x\":\"0x4fed262974c1d800\",\"y\":\"0x27590c454edd59ca\",\"order\":\"0x99235f1d975f376d\"}]}]}]\"\"\",\n",
    "    \"\"\"[{\"field\":{\"p\":\"0xa7ff74a0dc8c161d\"},\"a\":\"0xa7ff74a0dc8c161a\",\"b\":\"0x583b968bb611b284\",\"order\":\"0xa7ff74a06811ee75\",\"subgroups\":[{\"x\":\"0x5f5c76454edf12e7\",\"y\":\"0x4c73cbfc44f41508\",\"order\":\"0xa7ff74a06811ee75\",\"cofactor\":\"0x1\",\"points\":[{\"x\":\"0x5f5c76454edf12e7\",\"y\":\"0x4c73cbfc44f41508\",\"order\":\"0xa7ff74a06811ee75\"}]}]}]\"\"\",\n",
    "    \"\"\"[{\"field\":{\"p\":\"0xb52c62ca8703a063\"},\"a\":\"0xb52c62ca8703a060\",\"b\":\"0x0baec43a07b54c21\",\"order\":\"0xb52c62c963037121\",\"subgroups\":[{\"x\":\"0x6fe4a521a29bc1ab\",\"y\":\"0x3fca7180021f8f0f\",\"order\":\"0xb52c62c963037121\",\"cofactor\":\"0x1\",\"points\":[{\"x\":\"0x6fe4a521a29bc1ab\",\"y\":\"0x3fca7180021f8f0f\",\"order\":\"0xb52c62c963037121\"}]}]}]\"\"\",\n",
    "    \"\"\"[{\"field\":{\"p\":\"0xb8921f25b6ce5267\"},\"a\":\"0xb8921f25b6ce5264\",\"b\":\"0xa575c9f97563265d\",\"order\":\"0xb8921f2592b6b39f\",\"subgroups\":[{\"x\":\"0x7eb120fada47765c\",\"y\":\"0x64ef4e51d4159304\",\"order\":\"0xb8921f2592b6b39f\",\"cofactor\":\"0x1\",\"points\":[{\"x\":\"0x7eb120fada47765c\",\"y\":\"0x64ef4e51d4159304\",\"order\":\"0xb8921f2592b6b39f\"}]}]}]\"\"\",\n",
    "    \"\"\"[{\"field\":{\"p\":\"0xc591b8c4df0afc19\"},\"a\":\"0xc591b8c4df0afc16\",\"b\":\"0x0a1eb46a6e647f0a\",\"order\":\"0xc591b8c3eb07239f\",\"subgroups\":[{\"x\":\"0x1963bfb862cb0bf3\",\"y\":\"0x30da8bb7fa77277d\",\"order\":\"0xc591b8c3eb07239f\",\"cofactor\":\"0x1\",\"points\":[{\"x\":\"0x1963bfb862cb0bf3\",\"y\":\"0x30da8bb7fa77277d\",\"order\":\"0xc591b8c3eb07239f\"}]}]}]\"\"\",\n",
    "    # a = 0\n",
    "    \"\"\"[{\"field\":{\"p\":\"0xceaf446a53f14bc1\"},\"a\":\"0x0000000000000000\",\"b\":\"0x326539376260f173\",\"order\":\"0xceaf446aae275419\",\"subgroups\":[{\"x\":\"0x98fe44948c3f8678\",\"y\":\"0x3d440ee959a912d7\",\"order\":\"0xceaf446aae275419\",\"cofactor\":\"0x1\",\"points\":[{\"x\":\"0x98fe44948c3f8678\",\"y\":\"0x3d440ee959a912d7\",\"order\":\"0xceaf446aae275419\"}]}]}]\"\"\",\n",
    "    \"\"\"[{\"field\":{\"p\":\"0xb3c2beca75d66de3\"},\"a\":\"0x0000000000000000\",\"b\":\"0x46069225826b51aa\",\"order\":\"0xb3c2bec95881b695\",\"subgroups\":[{\"x\":\"0x81500c226efa0d5a\",\"y\":\"0x674e09d296452eee\",\"order\":\"0xb3c2bec95881b695\",\"cofactor\":\"0x1\",\"points\":[{\"x\":\"0x81500c226efa0d5a\",\"y\":\"0x674e09d296452eee\",\"order\":\"0xb3c2bec95881b695\"}]}]}]\"\"\",\n",
    "    \"\"\"[{\"field\":{\"p\":\"0xd6097c1ce207aae7\"},\"a\":\"0x0000000000000000\",\"b\":\"0x7adaab54e7dfd564\",\"order\":\"0xd6097c1b407eb413\",\"subgroups\":[{\"x\":\"0x151da8fb1f83201e\",\"y\":\"0x8bfeb90ec1177a91\",\"order\":\"0xd6097c1b407eb413\",\"cofactor\":\"0x1\",\"points\":[{\"x\":\"0x151da8fb1f83201e\",\"y\":\"0x8bfeb90ec1177a91\",\"order\":\"0xd6097c1b407eb413\"}]}]}]\"\"\",\n",
    "    \"\"\"[{\"field\":{\"p\":\"0x97a3e2d617a2309d\"},\"a\":\"0x0000000000000000\",\"b\":\"0x7f311cba46652247\",\"order\":\"0x97a3e2d712ffd715\",\"subgroups\":[{\"x\":\"0x46d725812af15870\",\"y\":\"0x727f88365dbd0e80\",\"order\":\"0x97a3e2d712ffd715\",\"cofactor\":\"0x1\",\"points\":[{\"x\":\"0x46d725812af15870\",\"y\":\"0x727f88365dbd0e80\",\"order\":\"0x97a3e2d712ffd715\"}]}]}]\"\"\",\n",
    "    \"\"\"[{\"field\":{\"p\":\"0xd8d7f39f545a5da7\"},\"a\":\"0x0000000000000000\",\"b\":\"0x09097ddcd4be8d55\",\"order\":\"0xd8d7f3a0c4115b4b\",\"subgroups\":[{\"x\":\"0xc5c771a9827a3251\",\"y\":\"0x64bf52041ac05b23\",\"order\":\"0xd8d7f3a0c4115b4b\",\"cofactor\":\"0x1\",\"points\":[{\"x\":\"0xc5c771a9827a3251\",\"y\":\"0x64bf52041ac05b23\",\"order\":\"0xd8d7f3a0c4115b4b\"}]}]}]\"\"\",\n",
    "    \"\"\"[{\"field\":{\"p\":\"0x847de883ab4fbf4d\"},\"a\":\"0x0000000000000000\",\"b\":\"0x09866366b3b45c2d\",\"order\":\"0x847de8837e6d4477\",\"subgroups\":[{\"x\":\"0x62fd7b4bc7c9acb4\",\"y\":\"0x2d0942774607106b\",\"order\":\"0x847de8837e6d4477\",\"cofactor\":\"0x1\",\"points\":[{\"x\":\"0x62fd7b4bc7c9acb4\",\"y\":\"0x2d0942774607106b\",\"order\":\"0x847de8837e6d4477\"}]}]}]\"\"\",\n",
    "    \"\"\"[{\"field\":{\"p\":\"0xf0d617c3c47b7c77\"},\"a\":\"0x0000000000000000\",\"b\":\"0xd856b3dcb95764a2\",\"order\":\"0xf0d617c5512cec85\",\"subgroups\":[{\"x\":\"0xeaf9b352a3daac45\",\"y\":\"0x4e4e557f9fc3febc\",\"order\":\"0xf0d617c5512cec85\",\"cofactor\":\"0x1\",\"points\":[{\"x\":\"0xeaf9b352a3daac45\",\"y\":\"0x4e4e557f9fc3febc\",\"order\":\"0xf0d617c5512cec85\"}]}]}]\"\"\",\n",
    "    \"\"\"[{\"field\":{\"p\":\"0xce920b656c80b373\"},\"a\":\"0x0000000000000000\",\"b\":\"0xb4a07dfae71ddc62\",\"order\":\"0xce920b65eee38015\",\"subgroups\":[{\"x\":\"0x7895c02b3c5205b5\",\"y\":\"0x2926be6446b98d62\",\"order\":\"0xce920b65eee38015\",\"cofactor\":\"0x1\",\"points\":[{\"x\":\"0x7895c02b3c5205b5\",\"y\":\"0x2926be6446b98d62\",\"order\":\"0xce920b65eee38015\"}]}]}]\"\"\",\n",
    "    \"\"\"[{\"field\":{\"p\":\"0xb6d4e25f76cc9df7\"},\"a\":\"0x0000000000000000\",\"b\":\"0x18e95a2283692623\",\"order\":\"0xb6d4e25ea270ed03\",\"subgroups\":[{\"x\":\"0x2da7a97d5d899bc5\",\"y\":\"0x17d27fd34562e3d9\",\"order\":\"0xb6d4e25ea270ed03\",\"cofactor\":\"0x1\",\"points\":[{\"x\":\"0x2da7a97d5d899bc5\",\"y\":\"0x17d27fd34562e3d9\",\"order\":\"0xb6d4e25ea270ed03\"}]}]}]\"\"\",\n",
    "    \"\"\"[{\"field\":{\"p\":\"0xe7cd1979ebed69ed\"},\"a\":\"0x0000000000000000\",\"b\":\"0x278e92b83191a7da\",\"order\":\"0xe7cd197966893365\",\"subgroups\":[{\"x\":\"0xc4de44402da5b9a6\",\"y\":\"0x2b45e7f32e3701ba\",\"order\":\"0xe7cd197966893365\",\"cofactor\":\"0x1\",\"points\":[{\"x\":\"0xc4de44402da5b9a6\",\"y\":\"0x2b45e7f32e3701ba\",\"order\":\"0xe7cd197966893365\"}]}]}]\"\"\"\n",
    "    # b = 0 (causes more issues than gain) but the (w12-0) coords actually require it...\n",
    "    #\"\"\"[{\"field\":{\"p\":\"0x9d9119957f02fe3f\"},\"a\":\"0x0106903196d88df9\",\"b\":\"0x0000000000000000\",\"order\":\"0x9d9119957f02fe40\",\"subgroups\":[{\"x\":\"0x191a36b9cd81de96\",\"y\":\"0x10f2c6bded391aa9\",\"order\":\"0x9d9119957f02fe40\",\"cofactor\":\"0x1\",\"points\":[{\"x\":\"0x0000000000000000\",\"y\":\"0x0000000000000000\",\"order\":\"0x2\"},{\"x\":\"0x95913fae9065da0f\",\"y\":\"0x5eeddeee7152d6fb\",\"order\":\"0x276446655fc0bf9\"}]}]}]\"\"\"\n",
    "]))\n",
    "\n",
    "for i, params in enumerate(curves):\n",
    "    curve = params.curve\n",
    "    if curve.parameters[\"a\"] == -3:\n",
    "        params.name = \"a=-3\"\n",
    "    elif curve.parameters[\"a\"] == -1:\n",
    "        params.name = \"a=-1\"\n",
    "    elif curve.parameters[\"a\"] == 0:\n",
    "        params.name = \"a=0\"\n",
    "    else:\n",
    "        params.name = \"random\"\n",
    "    params.name += f\"[{i}]\""
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "4276de4c-78f4-4cdb-b60e-8c24eabfa00d",
   "metadata": {},
   "source": [
    "### Computing addition chains\n",
    "\n",
    "First lets fix some scalars, go over the curves and compute the addition chain to obtain information about which multiples of the input point will go into the formulas. The i-th scalar will be used with the i-th curve as defined above. There are 10 unique scalars, so each curve group will share those."
   ]
  },
  {
   "cell_type": "code",
   "id": "7ab00a5c-2873-40c0-a81c-873350897e46",
   "metadata": {},
   "source": [
    "scalars = [0b1111110, 0b1011101, 0b110110, 0b100100, 0b1000110, 0b1001101, 0b101001, 0b1100100, 0b1010110, 0b101010] * 4\n",
    "\n",
    "chains = []\n",
    "scalar_map = {}\n",
    "chain_map = {}\n",
    "ops = set()\n",
    "for scalar, params in zip(scalars, curves):\n",
    "    chain = addition_chain(scalar, params, LTRMultiplier, lambda add,dbl: LTRMultiplier(add, dbl, None, False, AccumulationOrder.PeqRP, True, True))\n",
    "    chains.append(chain)\n",
    "    scalar_map[params] = scalar\n",
    "    chain_map[params] = chain\n",
    "    ops.update(chain)\n",
    "print(sorted(list(ops)))"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "8d1df0a5-ac90-45ee-b54a-1fa25dd86c74",
   "metadata": {},
   "source": [
    "### Loading the formulas\n",
    "\n",
    "Now lets load the formulas, either just those from the EFD or also load the expanded library formulas if they are available. See the [formulas](formulas.ipynb) notebook."
   ]
  },
  {
   "cell_type": "code",
   "id": "da4e2981-cf3d-40ea-9d1e-9b6c57a6e01d",
   "metadata": {},
   "source": [
    "load_expanded = False\n",
    "\n",
    "formula_classes = [AdditionFormula, DoublingFormula]\n",
    "formula_groups = {}\n",
    "for coord_name, coords in tqdm(model.coordinates.items(), desc=f\"Loading {'expanded' if load_expanded else 'EFD'} formulas\"):\n",
    "    groups = []\n",
    "    for formula_class in formula_classes:\n",
    "        expanded_path = Path(f\"sw_{coord_name}_{formula_class.shortname}s.pickle\")\n",
    "        if load_expanded:\n",
    "            if not expanded_path.exists():\n",
    "                raise ValueError(f\"Expanded formulas do not exist {expanded_path}.\")\n",
    "            with expanded_path.open(\"rb\") as f:\n",
    "                expanded = pickle.load(f)\n",
    "            formula_group = list(expanded)\n",
    "        else:\n",
    "            formula_group = list(filter(lambda formula: isinstance(formula, formula_class) and (formula.name.startswith(\"add\") or formula.name.startswith(\"dbl\")), coords.formulas.values()))\n",
    "        groups.append(formula_group)\n",
    "    formula_groups[coords] = groups\n",
    "\n",
    "print(f\"Loaded {sum(sum(len(group) for group in pair) for pair in formula_groups.values())} formulas.\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "12646534-8ca5-48c1-a4ae-1ad62575821f",
   "metadata": {},
   "source": [
    "### Computing the factor sets\n",
    "\n",
    "Now, compute the factor sets of the formulas **in parallel**. There are two options available. `xonly_fsets` builds the factor sets \"x\"-only by eliminating y-coords using the curve equation. `filter_nonhomo` specifies whether to filter out non-homogenous polynomials.\n",
    "\n",
    "<div class=\"alert alert-warning\">\n",
    "\n",
    "This is a resource intensive cell that uses parallelism. Set the `num_workers` variable to something reasonable, like the number of cores of your machine minus two.\n",
    "\n",
    "</div>"
   ]
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": "num_workers = 30",
   "id": "70a18871d5b46df4",
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "f71a9b68-6873-4bee-b1c2-3166ef74897b",
   "metadata": {},
   "source": [
    "xonly_fsets = False\n",
    "filter_nonhomo = False\n",
    "\n",
    "def compute_fsets(formula_group, formula_class, fh, xo):\n",
    "    from pyecsca.sca.re.zvp import compute_factor_set\n",
    "    from pyecsca.ec.formula import DoublingFormula\n",
    "    fsets = []\n",
    "    for formula in formula_group:\n",
    "        fset = compute_factor_set(formula, filter_nonhomo=fh, xonly=xo)\n",
    "\n",
    "        # Fix the factor set polynomials for the case of doubling.\n",
    "        # TODO: Investigate how this plays with xonly an filter_nonhomo arguments.\n",
    "        if formula_class == DoublingFormula:\n",
    "            new_fset = set()\n",
    "            for poly in fset:\n",
    "                pl = poly.copy()\n",
    "                for symbol in poly.free_symbols:\n",
    "                    original = str(symbol)\n",
    "                    if original.endswith(\"1\"):\n",
    "                        new = original.replace(\"1\", \"2\")\n",
    "                        pl = pl.subs(original, new)\n",
    "                new_fset.add(pl)\n",
    "            fset = new_fset\n",
    "        fsets.append(fset)\n",
    "    return fsets\n",
    "\n",
    "factor_sets = {}\n",
    "with TaskExecutor(max_workers=num_workers, mp_context=spawn_context) as pool, enable_spawn(compute_fsets) as target:\n",
    "    for coord_name, coords in model.coordinates.items():\n",
    "        for formula_group, formula_class in zip(formula_groups[coords], formula_classes):\n",
    "            pool.submit_task((coords, formula_group, formula_class),\n",
    "                             target, formula_group, formula_class, filter_nonhomo, xonly_fsets)\n",
    "\n",
    "    for (coords, formula_group, formula_class), future in tqdm(pool.as_completed(), desc=\"Computing factor sets\", total=len(pool.tasks)):\n",
    "        if error := future.exception():\n",
    "            print(coords, formula_class.shortname, error)\n",
    "            raise error\n",
    "        else:\n",
    "            fsets = future.result()\n",
    "            print(f\"Got {coords.name} {formula_class.shortname}s {len(fsets)}.\")\n",
    "            for formula, fset in zip(formula_group, fsets):\n",
    "                factor_sets[formula] = fset"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "c7a753be-741e-4687-b1a9-4688281d2c71",
   "metadata": {},
   "source": [
    "You can now store the (or load the previously computed) factor sets."
   ]
  },
  {
   "cell_type": "code",
   "id": "aef30807-e9eb-44f2-8837-f027c54eb724",
   "metadata": {},
   "source": [
    "with open(\"factor_sets.pickle\", \"wb\") as f:\n",
    "    pickle.dump(factor_sets, f)\n",
    "    print(f\"Stored {len(factor_sets)} factor sets.\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "86df44e6-aa3c-48eb-8561-36b08eaa781c",
   "metadata": {},
   "source": [
    "with open(\"factor_sets.pickle\", \"rb\") as f:\n",
    "    factor_sets = pickle.load(f)\n",
    "    print(f\"Loaded {len(factor_sets)} factor sets.\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "5781ae2b-28ce-4926-a2ab-b3d4689df29c",
   "metadata": {},
   "source": [
    "### Accumulating polynomials\n",
    "\n",
    "We now accumulate all of the polynomials for the adds and the dbls. We do so for all the compatible curves from our curves list. We will be looking for ZVP points on all of these curves, to make sure at least some lead to solutions."
   ]
  },
  {
   "cell_type": "code",
   "id": "564b2c16-6655-4ddb-8657-611d1d489e6a",
   "metadata": {},
   "source": [
    "polynomials = {}\n",
    "for coord_name, coords in tqdm(model.coordinates.items(), desc=\"Accumulating\"):\n",
    "    coord_adds = 0\n",
    "    coord_dbls = 0\n",
    "    for chain, affine_params in zip(chains, curves):\n",
    "        try:\n",
    "            params = affine_params.to_coords(coords)\n",
    "        except UnsatisfiedAssumptionError:\n",
    "            continue\n",
    "        add_polynomials = set()\n",
    "        dbl_polynomials = set()\n",
    "        for formula_group, formula_class in zip(formula_groups[coords], formula_classes):\n",
    "            for formula in formula_group:\n",
    "                if formula_class == AdditionFormula:\n",
    "                    add_polynomials.update(factor_sets.get(formula, []))\n",
    "                else:\n",
    "                    dbl_polynomials.update(factor_sets.get(formula, []))\n",
    "        polynomials[params] = {\n",
    "            \"add\": add_polynomials,\n",
    "            \"dbl\": dbl_polynomials\n",
    "        }\n",
    "        coord_adds += len(add_polynomials)\n",
    "        coord_dbls += len(dbl_polynomials)\n",
    "    print(f\"Got {coord_adds} add polys and {coord_dbls} dbl polys for {coord_name}.\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "04e7245b-a36a-43dc-9e73-626241112923",
   "metadata": {},
   "source": [
    "### Computing ZVP points\n",
    "Now, lets compute the sets of ZVP points, going over all the polynomials. Also filter the points such that for each \"polynomial, curve category, k\" we have only one point, as more are unnecessary.\n",
    "\n",
    "<div class=\"alert alert-warning\">\n",
    "\n",
    "This is a resource intensive cell that uses parallelism. Recall the `num_workers` variable.\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "id": "8a637899-8c3c-43a1-86fa-f0322d41ce8c",
   "metadata": {},
   "source": [
    "# bound is the maximal dlog in the hard case of the DCP to be solved\n",
    "bound = 100\n",
    "# Note that if you do not have the \"pari\" extra dependency installed (\"cysignals\", \"cypari2\") this bound\n",
    "# will have to be limited very low and the memory usage will be significant.\n",
    "\n",
    "all_points = set()\n",
    "all_points_filtered = {}\n",
    "dk = set()\n",
    "with TaskExecutor(max_workers=num_workers, mp_context=spawn_context) as pool:\n",
    "    for coord_name, coords in tqdm(model.coordinates.items(), desc=\"Submitting\"):\n",
    "        for chain, affine_params in zip(chains, curves):\n",
    "            try:\n",
    "                params = affine_params.to_coords(coords)\n",
    "            except UnsatisfiedAssumptionError:\n",
    "                continue\n",
    "            unique = set()\n",
    "            for op, ks in chain:\n",
    "                if len(ks) == 1:\n",
    "                    k = ks[0]\n",
    "                else:\n",
    "                    # zvp_points assumes (P, [k]P)\n",
    "                    ks_mod = list(map(lambda v: mod(v, params.order), ks))\n",
    "                    k = int(ks_mod[1] / ks_mod[0])\n",
    "                polys = polynomials[params][op]\n",
    "                for poly in polys:\n",
    "                    only_1 = all((not str(gen).endswith(\"2\")) for gen in poly.gens)\n",
    "                    only_2 = all((not str(gen).endswith(\"1\")) for gen in poly.gens)\n",
    "                    # This is the hard case where a dlog needs to be substituted, so bound it.\n",
    "                    if not (only_1 or only_2) and k > bound:\n",
    "                        continue\n",
    "                    unique.add((poly, k))\n",
    "            for poly, k in unique:\n",
    "                pool.submit_task((poly, affine_params, k),\n",
    "                                 zvp_points, poly, params.curve, k, params.order)\n",
    "    for (poly, affine_params, k), future in tqdm(pool.as_completed(), desc=\"Computing\", total=len(pool.tasks), smoothing=0):\n",
    "        params_name_match = re.match(r\"(.+)\\[([0-9]+)\\]\", affine_params.name)\n",
    "        params_category = params_name_match.group(1)\n",
    "        if error := future.exception():\n",
    "            print(error)\n",
    "        elif result := future.result():\n",
    "            for point in result:\n",
    "                all_points.add((point, affine_params))\n",
    "            all_points_filtered[(poly, params_category, k)] = (result, affine_params)\n",
    "\n",
    "print(f\"Got {len(all_points)} points.\")\n",
    "print(f\"Got {len(all_points_filtered)} filtered points\")#, but just {len(set(all_points_filtered.values()))} unique.\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "8d571e64-5261-4cda-b490-5b85dd6663d2",
   "metadata": {},
   "source": [
    "You can now store the (or load the previously computed) point sets."
   ]
  },
  {
   "cell_type": "code",
   "id": "82bc0305-b302-4240-a34c-81e13fe646b5",
   "metadata": {},
   "source": [
    "with open(\"all_points.pickle\", \"wb\") as f:\n",
    "    pickle.dump((all_points, all_points_filtered), f)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "411832e9-26a8-46a4-82e6-5f6afe879188",
   "metadata": {},
   "source": [
    "with open(\"all_points.pickle\", \"rb\") as f:\n",
    "    all_points, all_points_filtered = pickle.load(f)\n",
    "    \n",
    "print(f\"Got {len(all_points)} points.\")\n",
    "print(f\"Got {len(all_points_filtered)} filtered points\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "ed2d4b51-573f-4fef-b1a4-d30b3be07423",
   "metadata": {},
   "source": [
    "### Remapping\n",
    "\n",
    "Our ZVP points might (due to the bounds above thus the incompleteness of our analysis) lead to more zeros than we attribute to them (in more configurations), i.e. \"false negatives\". They might also be erroneous and not lead to zeros if the argument `filter_nonhomo` is False, as non-homogenous intermediate polynomials are not filtered out of the analysis. They introduce \"false positives\" but also some true positives.\n",
    "\n",
    "Thus, we perform a remapping step where we execute the scalar multiplication with given points and trace whether they introduce the zeros. This gives us a new distinguishing map `remapped_hit_point_map`, now without \"false negatives\" or \"false positives\".\n",
    "\n",
    "Note that we also construct two additional maps `remapped_count_point_map` and `remapped_position_point_map` which represent the results of remapping using an oracle counting the zeros and an oracle giving the positions of the zeros in the computation, respectively.\n",
    "\n",
    "<div class=\"alert alert-warning\">\n",
    "\n",
    "This is a resource intensive cell that uses parallelism. Recall the `num_workers` variable.\n",
    "\n",
    "</div>"
   ]
  },
  {
   "cell_type": "code",
   "id": "ca932717-2b0a-4dea-ae84-988b7d9af233",
   "metadata": {},
   "source": [
    "def remap(coords, chunk, points, scalar_map):\n",
    "    import numpy as np\n",
    "    from pyecsca.ec.mult import LTRMultiplier, AccumulationOrder\n",
    "    from pyecsca.ec.context import local, DefaultContext\n",
    "    from pyecsca.ec.error import UnsatisfiedAssumptionError\n",
    "    from pyecsca.ec.formula import FormulaAction\n",
    "    lp = len(points)\n",
    "    lc = len(chunk)\n",
    "    counts = np.full((lc, lp), -1, dtype=np.int16)\n",
    "    positions = np.full((lc, lp), None, dtype=object)\n",
    "    for i, formulas in enumerate(chunk):\n",
    "        mult = LTRMultiplier(*formulas, None, False, AccumulationOrder.PeqRP, True, True)\n",
    "        \n",
    "        for j, entry in enumerate(points):\n",
    "            if entry is None:\n",
    "                continue\n",
    "            point, params = entry\n",
    "            mult.init(params, point)\n",
    "            scalar = scalar_map[params]\n",
    "            with local(DefaultContext()) as ctx:\n",
    "                try:\n",
    "                    mult.multiply(scalar)\n",
    "                except UnsatisfiedAssumptionError:\n",
    "                    continue\n",
    "\n",
    "            zeros = []\n",
    "            \n",
    "            def callback(action):\n",
    "                if isinstance(action, FormulaAction):\n",
    "                    for intermediate in action.op_results:\n",
    "                        zeros.append(int(intermediate.value) == 0)\n",
    "            ctx.actions[0].walk(callback)\n",
    "            count = sum(zeros)\n",
    "            counts[i, j] = count\n",
    "            positions[i, j] = tuple(zeros)        \n",
    "    return counts, positions\n",
    "\n",
    "remapped_hit_point_map = {}\n",
    "remapped_count_point_map = {}\n",
    "remapped_position_point_map = {}\n",
    "#all_points_list = list(all_points)\n",
    "all_points_list = list(set((list(res[0])[0], res[1]) for res in all_points_filtered.values()))\n",
    "\n",
    "with TaskExecutor(max_workers=num_workers, mp_context=spawn_context) as pool, enable_spawn(remap) as remap_spawn:\n",
    "    for coord_name, coords in tqdm(model.coordinates.items()):\n",
    "        param_map = {}\n",
    "        points_mapped = []\n",
    "        scalars_mapped = {}\n",
    "        mapped = 0\n",
    "        for point, params in tqdm(all_points_list, desc=f\"Map points to {coord_name}\", leave=False):\n",
    "            if params not in param_map:\n",
    "                try:\n",
    "                    param_map[params] = params.to_coords(coords)\n",
    "                except UnsatisfiedAssumptionError:\n",
    "                    param_map[params] = None\n",
    "            if param_map[params] is None:\n",
    "                points_mapped.append(None)\n",
    "            else:\n",
    "                mapped += 1\n",
    "                points_mapped.append((point.to_model(param_map[params].curve.coordinate_model, param_map[params].curve), param_map[params]))\n",
    "        print(f\"{coord_name}: {mapped} are compatible. Remapping...\")\n",
    "        for params, scalar in scalar_map.items():\n",
    "            if params not in param_map or param_map[params] is None:\n",
    "                continue\n",
    "            scalars_mapped[param_map[params]] = scalar\n",
    "        \n",
    "        pairs = list(product(*formula_groups[coords]))\n",
    "        chunk_size = 10\n",
    "        chunks = 0\n",
    "        for i in trange(0, len(pairs), chunk_size, desc=f\"Chunking {coord_name}\", smoothing=0, leave=False):\n",
    "            chunk = pairs[i:i+chunk_size]\n",
    "            chunks += 1\n",
    "            pool.submit_task(chunk,\n",
    "                             remap_spawn, coords, chunk, points_mapped, scalars_mapped)\n",
    "\n",
    "        for chunk, future in tqdm(pool.as_completed(), total=chunks, unit_scale=chunk_size, desc=f\"Remapping {coord_name} ({len(pairs)} formula pairs)\", leave=False, smoothing=0):\n",
    "            error = future.exception()\n",
    "            if error:\n",
    "                print(error)\n",
    "            else:\n",
    "                counts, positions = future.result()\n",
    "                for cfg, counts_row, positions_row in zip(chunk, counts, positions):\n",
    "                    hit_set = set()\n",
    "                    hit_map = {}\n",
    "                    count_map = {}\n",
    "                    position_map = {}\n",
    "                    for pp_tuple, count, position in zip(all_points_list, counts_row, positions_row):\n",
    "                        hit_map[pp_tuple] = None if count == -1 else (count > 0)\n",
    "                        count_map[pp_tuple] = count\n",
    "                        position_map[pp_tuple] = position\n",
    "                    remapped_hit_point_map[cfg] = hit_map\n",
    "                    remapped_count_point_map[cfg] = count_map\n",
    "                    remapped_position_point_map[cfg] = position_map\n",
    "        print(f\"{coord_name}: Remapped.\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "07755fca-07cd-4384-8581-64e701659361",
   "metadata": {},
   "source": [
    "You can now store the (or load the previously computed) remmapped maps."
   ]
  },
  {
   "cell_type": "code",
   "id": "95c3693b-166c-4515-9153-e3e00d222835",
   "metadata": {},
   "source": [
    "with open(\"remapped.pickle\", \"wb\") as f:\n",
    "    pickle.dump((remapped_hit_point_map, remapped_count_point_map, remapped_position_point_map), f)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "3914501e-269c-4a88-ba5b-ad3c6dd28cfb",
   "metadata": {},
   "source": [
    "with open(\"remapped.pickle\", \"rb\") as f:\n",
    "    remapped_hit_point_map, remapped_count_point_map, remapped_position_point_map = pickle.load(f)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "3c0d710b-19fd-4ff7-a39f-f38b1b8856f8",
   "metadata": {},
   "source": [
    "### Distinguishing map and distinguishing tree building\n",
    "Finally, we can build a tree using the remapped distinguishing map. Let's first wrap the raw mapping with a distinguishing Map object."
   ]
  },
  {
   "cell_type": "code",
   "id": "236d40a0-7a9c-48a0-b932-a0dfe1110a8c",
   "metadata": {},
   "source": [
    "cfgs = set()\n",
    "for coord_name, coords in model.coordinates.items():\n",
    "    cfgs.update(product(*formula_groups[coords]))\n",
    "\n",
    "param_categories = {\n",
    "    \"a=-1\": [\"projective-1\"],\n",
    "    \"a=-3\": [\"projective-3\", \"jacobian-3\", \"xyzz-3\"],\n",
    "    \"a=0\": [\"jacobian-0\"],\n",
    "    \"generic\": [\"jacobian\", \"projective\", \"modified\", \"xyzz\", \"xz\"],\n",
    "    \"b=0\": [\"w12-0\"]\n",
    "}\n",
    "cfg_categories = {}\n",
    "for name, coord_names in param_categories.items():\n",
    "    category_cfgs = set()\n",
    "    for coord_name in coord_names:\n",
    "        coords = model.coordinates[coord_name]\n",
    "        category_cfgs.update(filter(lambda cfg: cfg[0].coordinate_model == coords and cfg[1].coordinate_model == coords, cfgs))\n",
    "    cfg_categories[name] = category_cfgs\n",
    "category_map = {cfg: {\"category\": name} for name, category_cfgs in cfg_categories.items() for cfg in category_cfgs}\n",
    "dmap_categories = Map.from_io_maps(cfgs, category_map)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "db1ffe84-f58e-465c-b417-ad1d563c8377",
   "metadata": {},
   "source": [
    "dmap_remapped = Map.from_io_maps(cfgs, remapped_hit_point_map)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "48baa442-549f-4ffa-b296-4f964c8efc77",
   "metadata": {},
   "source": [
    "from copy import deepcopy\n",
    "dmap_dedup = deepcopy(dmap_remapped)\n",
    "dmap_dedup.deduplicate()\n",
    "print(f\"Rows before: {len(dmap_remapped.mapping)}, rows after deduplication: {len(dmap_dedup.mapping)}.\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "ce87cea8-ea88-49ee-a7dd-ead62dff5255",
   "metadata": {},
   "source": [
    "Let's now watch the tree getting built."
   ]
  },
  {
   "cell_type": "code",
   "id": "5ffc5564-a42b-4837-94bb-4ecb890b38b2",
   "metadata": {
    "scrolled": true
   },
   "source": [
    "tree_categories = Tree.build(cfgs, dmap_categories)\n",
    "tree_remapped = tree_categories.expand(dmap_dedup)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "bdb02213-4743-488c-b6ab-a934a409df3b",
   "metadata": {},
   "source": [
    "The tree is built, we can examine its properties, such as:\n",
    " - the total number of configurations (formula pairs),\n",
    " - its depth,\n",
    " - number of leaves (configurations that cannot be further distinguished),\n",
    " - average leaf size\n",
    " - mean result size (if this tree were to be used for reverse-enginnering)."
   ]
  },
  {
   "cell_type": "code",
   "id": "b921aadd-ccc6-4bc2-9cb0-e3905ff4f6ab",
   "metadata": {},
   "source": [
    "print(tree_remapped.describe())"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "c3571008-1898-41e0-87d5-848707345d62",
   "metadata": {},
   "source": [
    "print(tree_remapped.render_basic())"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "27f4be08-76fd-437a-bde7-8579b38fc686",
   "metadata": {},
   "source": [
    "We can also investigate the other oracles and the distinguishing trees they can build:"
   ]
  },
  {
   "cell_type": "code",
   "id": "5a680f71-f2de-4358-8352-67ec99d63704",
   "metadata": {},
   "source": [
    "dmap_count = Map.from_io_maps(cfgs, remapped_count_point_map)\n",
    "dmap_count.deduplicate()"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "764b6dec-f1d5-4e8f-8585-b35d9261b1ca",
   "metadata": {},
   "source": [
    "print(dmap_count.describe())"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "cee72bcb-2486-4285-8dc7-80c101ca1760",
   "metadata": {
    "scrolled": true
   },
   "source": [
    "tree_count = tree_categories.expand(dmap_count)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "17de9a81-102e-4804-906a-02d2baef42d5",
   "metadata": {},
   "source": [
    "dmap_position = Map.from_io_maps(cfgs, remapped_position_point_map)\n",
    "dmap_position.deduplicate()"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "bc3db76f-192f-4e70-b2c4-8aa5a5f4e078",
   "metadata": {},
   "source": [
    "print(dmap_position.describe())"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "db92360c-cb96-4382-8255-d7773772b07d",
   "metadata": {
    "scrolled": true
   },
   "source": [
    "tree_position = tree_categories.expand(dmap_position)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "b81d013a-7ff0-4b02-8c89-d6748d42874b",
   "metadata": {},
   "source": [
    "print(\"Zero hit\")\n",
    "print(tree_remapped.describe())\n",
    "print(\"\\nZero count\")\n",
    "print(tree_count.describe())\n",
    "print(\"\\nZero position\")\n",
    "print(tree_position.describe())"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "eff08e04-20bc-4fc7-978d-9cbfc06179b6",
   "metadata": {},
   "source": [
    "### Evaluation\n",
    "\n",
    "The following cells evaluate the ZVP-RE method and its behavior under various levels of noise. The cells with the `store` and `load` calls can be used to store the results so that different plots can be printed without re-running the evaluation.\n",
    "\n",
    "<div class=\"alert alert-warning\">\n",
    "\n",
    "These are resource intensive cells that use parallelism. Recall the `num_workers` variable.\n",
    "\n",
    "</div>"
   ]
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "First, lets look at the binary oracle in presence of symmetric noise.",
   "id": "7d61aa75df8e76b9"
  },
  {
   "cell_type": "code",
   "id": "3ef78642-fe65-46c7-8233-11dc43525991",
   "metadata": {},
   "source": "correct_rate, precise_rate, amount_rate, query_rate = eval_tree_symmetric1(cfgs, [tree_remapped], num_tries=100, num_cores=num_workers)",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "Let's save the results for later and then plot them.",
   "id": "3e9d4f11cc322958"
  },
  {
   "cell_type": "code",
   "id": "34b5adf6-c4d1-4534-b95c-7e85e780e66f",
   "metadata": {},
   "source": [
    "store(\"zvp_re_symmetric.nc\", correct_rate, precise_rate, amount_rate, query_rate)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "0a443d93-cf71-4c6e-bc11-998a31a47656",
   "metadata": {},
   "source": [
    "correct_rate, precise_rate, amount_rate, query_rate = load(\"zvp_re_symmetric.nc\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "3bef547e-7c8e-4ec4-a842-75889cf6acd9",
   "metadata": {},
   "source": [
    "success_rate_symmetric(correct_rate, None).savefig(\"zvp_re_success_rate_symmetric.pdf\", bbox_inches=\"tight\")\n",
    "precise_rate_symmetric(precise_rate).savefig(\"zvp_re_precise_rate_symmetric.pdf\", bbox_inches=\"tight\")\n",
    "query_rate_symmetric(query_rate).savefig(\"zvp_re_query_rate_symmetric.pdf\", bbox_inches=\"tight\")\n",
    "amount_rate_symmetric(amount_rate).savefig(\"zvp_re_amount_rate_symmetric.pdf\", bbox_inches=\"tight\")\n",
    "success_rate_vs_query_rate_symmetric(query_rate, correct_rate).savefig(\"zvp_re_scatter_symmetric.pdf\", bbox_inches=\"tight\")\n",
    "success_rate_vs_majority_symmetric(correct_rate).savefig(\"zvp_re_plot_symmetric.pdf\", bbox_inches=\"tight\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "Now, lets look at the binary oracle in presence of asymmetric noise.",
   "id": "4ef7396c82d760df"
  },
  {
   "cell_type": "code",
   "id": "064e76e9-a299-4a5c-ae44-adfe07d86901",
   "metadata": {},
   "source": "correct_rate_b, precise_rate_b, amount_rate_b, query_rate_b = eval_tree_asymmetric1(cfgs, [tree_remapped], num_tries=100, num_cores=num_workers)",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "Let's save the results for later and then plot them.",
   "id": "335e6e29d6e92030"
  },
  {
   "cell_type": "code",
   "id": "ccdebf35-6401-497f-8731-b205eab672f8",
   "metadata": {},
   "source": [
    "store(\"zvp_re_asymmetric.nc\", correct_rate_b, precise_rate_b, amount_rate_b, query_rate_b)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "8924c80d-2d65-40a6-b589-85bbe93d1094",
   "metadata": {},
   "source": [
    "correct_rate_b, precise_rate_b, amount_rate_b, query_rate_b = load(\"zvp_re_asymmetric.nc\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "fa23ea1b-2b95-4f22-b237-bd4d75980681",
   "metadata": {},
   "source": [
    "success_rate_asymmetric(correct_rate_b, None).savefig(\"zvp_re_success_rate_asymmetric.pdf\", bbox_inches=\"tight\")\n",
    "precise_rate_asymmetric(precise_rate_b).savefig(\"zvp_re_precise_rate_asymmetric.pdf\", bbox_inches=\"tight\")\n",
    "query_rate_asymmetric(query_rate_b).savefig(\"zvp_re_query_rate_asymmetric.pdf\", bbox_inches=\"tight\")\n",
    "amount_rate_asymmetric(amount_rate_b).savefig(\"zvp_re_amount_rate_asymmetric.pdf\", bbox_inches=\"tight\")\n",
    "success_rate_vs_majority_asymmetric(correct_rate_b).savefig(\"zvp_re_plot_asymmetric.pdf\", bbox_inches=\"tight\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "Finally, lets analyze the count oracle in the presence of noise.",
   "id": "2ba367834f6a62fb"
  },
  {
   "cell_type": "code",
   "id": "d0566c3d-fb55-496b-aeb3-17b674d32ca2",
   "metadata": {},
   "source": "correct_rate_c, precise_rate_c, amount_rate_c, query_rate_c = eval_tree_binomial1(cfgs, [tree_count], num_tries=100, num_cores=num_workers)",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "Let's save the results for later and then plot them.",
   "id": "580d9adeb81b6b4b"
  },
  {
   "cell_type": "code",
   "id": "ea781b4d-0073-4c83-9854-0f7b104f7b14",
   "metadata": {},
   "source": [
    "store(\"zvp_re_binomial.nc\", correct_rate_c, precise_rate_c, amount_rate_c, query_rate_c)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "088a9776-0713-402d-b351-994ba2e24d4d",
   "metadata": {},
   "source": [
    "correct_rate_c, precise_rate_c, amount_rate_c, query_rate_c = load(\"zvp_re_binomial.nc\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "1eedce32-a74a-4baf-a5ab-a08cd8b96a6e",
   "metadata": {},
   "source": [
    "success_rate_binomial(correct_rate_c, None).savefig(\"zvp_re_success_rate_binomial.pdf\", bbox_inches=\"tight\")\n",
    "precise_rate_binomial(precise_rate_c).savefig(\"zvp_re_precise_rate_binomial.pdf\", bbox_inches=\"tight\")\n",
    "query_rate_binomial(query_rate_c).savefig(\"zvp_re_query_rate_binomial.pdf\", bbox_inches=\"tight\")\n",
    "amount_rate_binomial(amount_rate_c).savefig(\"zvp_re_amount_rate_binomial.pdf\", bbox_inches=\"tight\")"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "c2c3dbc1-738e-4b71-b48d-d82488995e5a",
   "metadata": {},
   "source": [
    "### Factor sets"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e50daa23-25b8-4970-a69b-e31d8b000204",
   "metadata": {},
   "source": [
    "Now, lets examine the representation of factor sets and the distinguishing trees they can build."
   ]
  },
  {
   "cell_type": "code",
   "id": "d258bc29-b24f-4d57-a183-cb5d58e3f20f",
   "metadata": {},
   "source": [
    "fset_map = {}\n",
    "fset_nonhomo_map = {}\n",
    "factor_sets = {}\n",
    "factor_sets_nonhomo = {}\n",
    "for coord_name, coords in tqdm(model.coordinates.items()):\n",
    "    for formula_group in formula_groups[coords]:\n",
    "        for formula in tqdm(formula_group, leave=False):\n",
    "            factor_sets[formula] = compute_factor_set(formula)\n",
    "            factor_sets_nonhomo[formula] = compute_factor_set(formula, filter_nonhomo=False)\n",
    "    formula_combinations = list(product(*formula_groups[coords]))\n",
    "    for formulas in tqdm(formula_combinations, leave=False):\n",
    "        fset = set()\n",
    "        fset_nonhomo = set()\n",
    "        for formula in formulas:\n",
    "            fset.update(factor_sets[formula])\n",
    "            fset_nonhomo.update(factor_sets_nonhomo[formula])\n",
    "        fset_map[formulas] = fset\n",
    "        fset_nonhomo_map[formulas] = fset_nonhomo"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "fde74d6a-9895-4fed-83a1-29a5c06ad5a6",
   "metadata": {},
   "source": [
    "with open(\"factor_sets_extended.pickle\", \"wb\") as f:\n",
    "    pickle.dump((factor_sets, fset_map), f)\n",
    "with open(\"factor_sets_nonhomo_extended.pickle\", \"wb\") as f:\n",
    "    pickle.dump((factor_sets_nonhomo, fset_nonhomo_map), f)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "e513a996-3f61-4ae7-9089-e5e46de6b863",
   "metadata": {},
   "source": [
    "with open(\"factor_sets_extended.pickle\", \"rb\") as f:\n",
    "    factor_sets, fset_map = pickle.load(f)\n",
    "with open(\"factor_sets_nonhomo_extended.pickle\", \"rb\") as f:\n",
    "    factor_sets_nonhomo, fset_nonhomo_map = pickle.load(f)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "e1d50275-2917-4882-be29-28c67e58f23a",
   "metadata": {},
   "source": [
    "Build two trees, one with the filtered factor sets and one with unfilitered factor sets (with non-homogenous polynomials present)."
   ]
  },
  {
   "cell_type": "code",
   "id": "8076b999-0c93-4748-aa3d-3986f05c62e3",
   "metadata": {},
   "source": [
    "dmap_fset = Map.from_sets(set(fset_map.keys()), fset_map, deduplicate=True)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "47bc3406-6f60-4706-a59d-19feb8f50dea",
   "metadata": {},
   "source": [
    "print(dmap_fset.describe())"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "6d4f33f3-557a-4e79-aadb-10fd73aba1d0",
   "metadata": {},
   "source": [
    "tree_fset = Tree.build(dmap_fset.cfgs, dmap_fset)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "7e539028-c7af-446e-b523-97f08fa3e39d",
   "metadata": {},
   "source": [
    "dmap_fset_nonhomo = Map.from_sets(set(fset_nonhomo_map.keys()), fset_nonhomo_map, deduplicate=True)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "afdc58ef-28c2-469a-9a2d-6db44824915e",
   "metadata": {},
   "source": [
    "print(dmap_fset_nonhomo.describe())"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "660bf7d1-e25d-40c4-bcd7-d520c566c432",
   "metadata": {},
   "source": [
    "tree_fset_nonhomo = Tree.build(dmap_fset_nonhomo.cfgs, dmap_fset_nonhomo)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "93a67306-e507-422c-92e2-f1439437060f",
   "metadata": {},
   "source": [
    "print(\"Factor sets\")\n",
    "print(tree_fset.describe())\n",
    "print(\"\\nFactor sets (unfiltered)\")\n",
    "print(tree_fset_nonhomo.describe())"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "a60a099e-cb8f-4e9f-9b94-591957908d55",
   "metadata": {},
   "source": [
    "As we can see, our technique is able to distinguish formulas better than the filtered factor sets, but not better than unfiltered factor sets.\n",
    "\n",
    "We can further analyze where unused possibilities of distinguishing remain but expanding the tree using the distinguishing map built out of factor sets."
   ]
  },
  {
   "cell_type": "code",
   "id": "366ba814-c8ae-4567-8e26-27130f9dbfc9",
   "metadata": {},
   "source": [
    "expanded = tree_remapped.expand(dmap_fset)"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "code",
   "id": "bf8ad92e-a3e9-43c6-8c8a-c5e6d42eee6b",
   "metadata": {},
   "source": [
    "print(\"Zero hit + factor sets\")\n",
    "print(expanded.describe())"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "cell_type": "markdown",
   "id": "dc60901f-694b-47f3-b607-ef04646f41b9",
   "metadata": {},
   "source": [
    "The tree improves its distinguishing abilities, however, upon closer analysis, the polynomials that it uses to further split the configurations never actually have solutions on the curves with the assumed properties (e.g. `a = 0`)."
   ]
  },
  {
   "cell_type": "code",
   "id": "892e02c7-8c39-4467-91ff-b6ae265684aa",
   "metadata": {},
   "source": [],
   "outputs": [],
   "execution_count": null
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}