6 files changed, 73 insertions, 55 deletions

diff --git a/analysis/scalarmults/epare/error_model.py b/analysis/scalarmults/epare/error_model.py
index 2445dda..f05e415 100644
--- a/analysis/scalarmults/epare/error_model.py
+++ b/analysis/scalarmults/epare/error_model.py
@@ -22,18 +22,6 @@ def check_affine(k, q):
     return k % q == 0
 
 
-def check_any(*checks, q=None):
-    """Merge multiple checks together. The returned check function no longer takes the `q` parameter."""
-
-    def check_func(k, l):
-        for check in checks:
-            if check(k, l, q):
-                return True
-        return False
-
-    return check_func
-
-
 # These checks can be applied to add formulas. See the formulas notebook for background on them.
 checks_add = {
     "equal_multiples": check_equal_multiples,
@@ -84,21 +72,22 @@ class ErrorModel:
         object.__setattr__(self, "check_condition", check_condition)
         object.__setattr__(self, "precomp_to_affine", precomp_to_affine)
 
-    def check_add(self, q):
-        """Get the add formula check function for the given q."""
-        if self.checks == {"affine"}:
-            return lambda k, l: False
-        return check_any(
-            *map(
-                lambda name: checks_add[name],
-                filter(lambda check: check in checks_add, self.checks),
-            ),
-            q=q,
-        )
+    def make_check_funcs(self, q: int):
+        """Make the check_funcs for a given q."""
+        def affine(k):
+            return check_affine(k, q)
+
+        add_checks = [checks_add[check] for check in self.checks if check in checks_add]
+        def add(k, l):
+            for check in add_checks:
+                if check(k, l, q):
+                    return True
+            return False
 
-    def check_affine(self, q):
-        """Get the to-affine check function for the given q."""
-        return partial(check_affine, q=q)
+        return {
+            "affine": affine,
+            "add": add
+        }
 
     def __lt__(self, other):
         if not isinstance(other, ErrorModel):
diff --git a/analysis/scalarmults/epare/simulate.py b/analysis/scalarmults/epare/simulate.py
index 24e41b6..82c0331 100644
--- a/analysis/scalarmults/epare/simulate.py
+++ b/analysis/scalarmults/epare/simulate.py
@@ -7,7 +7,7 @@ from functools import partial
 
 from .config import Config
 from .mult_results import MultResults
-from .prob_map import ProbMap
+from .prob_map import ProbMap, hash_divisors
 
 from pyecsca.ec.params import DomainParameters
 from pyecsca.ec.mod import mod
@@ -81,15 +81,14 @@ def evaluate_multiples(
     use_multiply: bool = True,
 ):
     """
-    Takes MultIdent and MultResults and a set of divisors (base point orders `q`) and
-    evaluates them using the error model from the MultIdent. Note that the MultIdent
+    Takes a Config and MultResults and a set of divisors (base point orders `q`) and
+    evaluates them using the error model from the Config. Note that the Config
     must have an error model in this case. Returns the ProbMap.
     """
+    if not mult.has_error_model:
+        raise ValueError("Invalid config")
     errors = {divisor: 0 for divisor in divisors}
-    samples = len(res)
-    divisors_hash = hashlib.blake2b(
-        str(sorted(divisors)).encode(), digest_size=8
-    ).digest()
+    check_funcs = {q: mult.error_model.make_check_funcs(q) for q in divisors}
     for precomp_ctx, full_ctx, out in res:
         check_inputs = graph_to_check_inputs(
             precomp_ctx,
@@ -102,10 +101,7 @@ def evaluate_multiples(
         )
         for q in divisors:
             error = evaluate_checks(
-                check_funcs={
-                    "add": mult.error_model.check_add(q),
-                    "affine": mult.error_model.check_affine(q),
-                },
+                check_funcs=check_funcs[q],
                 check_inputs=check_inputs,
             )
             errors[q] += error
@@ -114,7 +110,9 @@ def evaluate_multiples(
     for q, error in errors.items():
         if error != 0:
             probs[q] = error / samples
-    return ProbMap(probs, divisors_hash, samples)
+    samples = len(res)
+    dhash = hash_divisors(divisors)
+    return ProbMap(probs, dhash, samples)
 
 
 def evaluate_multiples_direct(
diff --git a/analysis/scalarmults/epare/standalone/make_multiples.py b/analysis/scalarmults/epare/standalone/make_multiples.py
index dbc9041..5f4a2b4 100644
--- a/analysis/scalarmults/epare/standalone/make_multiples.py
+++ b/analysis/scalarmults/epare/standalone/make_multiples.py
@@ -57,7 +57,7 @@ def main(temp, workers, seed, samples):
             for mult, future in tqdm(pool.as_completed(), desc="Computing multiple graphs.", total=len(pool.tasks)):
                 i += 1
                 if error := future.exception():
-                    print("Error!", mult, error)
+                    click.echo("Error!", mult, error)
                     continue
                 fpath = future.result()
                 with fpath.open("rb") as f:
diff --git a/analysis/scalarmults/epare/standalone/make_probs.py b/analysis/scalarmults/epare/standalone/make_probs.py
index b92a215..f836853 100644
--- a/analysis/scalarmults/epare/standalone/make_probs.py
+++ b/analysis/scalarmults/epare/standalone/make_probs.py
@@ -45,11 +45,10 @@ def main(temp, workers, seed):
     in_path = Path(f"multiples_{seed}.zpickle")
     out_path = Path(f"probs_{seed}.zpickle")
 
-    with zstd.open(in_path, "rb") as f, zstd.open(out_path, "wb") as h, TaskExecutor(max_workers=workers) as pool, tqdm(total=len(all_configs), desc=f"Generating probability maps.", smoothing=0) as bar:
+    with zstd.open(in_path, "rb") as f, zstd.open(out_path, "wb") as h, TaskExecutor(max_workers=workers) as pool, tqdm(total=len(all_configs) * len(all_error_models), desc=f"Generating probability maps.", smoothing=0) as bar:
         while True:
             try:
                 start = f.tell()
-                bar.update(1)
                 mult, _ = pickle.load(f)
                 for error_model in all_error_models:
                     full = mult.with_error_model(error_model)
@@ -59,19 +58,21 @@ def main(temp, workers, seed):
                                      full, in_path, start, divisor_map["all"], use_init, use_multiply)
                 gc.collect()
                 for full, future in pool.as_completed(wait=False):
+                    bar.update(1)
                     if error := future.exception():
-                        print("Error!", full, error)
+                        click.echo("Error!", full, error)
                         continue
                     res = future.result()
                     pickle.dump((full, res), h)
             except EOFError:
                 break
             except pickle.UnpicklingError:
-                print("Bad unpickling, the multiples file is likely truncated.")
+                click.echo("Bad unpickling, the multiples file is likely truncated.")
                 break
         for full, future in pool.as_completed():
+            bar.update(1)
             if error := future.exception():
-                print("Error!", full, error)
+                click.echo("Error!", full, error)
                 continue
             res = future.result()
             pickle.dump((full, res), h)
diff --git a/analysis/scalarmults/epare/standalone/merge_probs.py b/analysis/scalarmults/epare/standalone/merge_probs.py
index 6748a3d..d40945b 100644
--- a/analysis/scalarmults/epare/standalone/merge_probs.py
+++ b/analysis/scalarmults/epare/standalone/merge_probs.py
@@ -11,7 +11,7 @@ from pathlib import Path
 
 from tqdm import tqdm
 
-from .. import ProbMap
+from .. import all_configs, all_error_models, ProbMap
 
 
 if sys.version_info >= (3, 14):
@@ -21,13 +21,31 @@ else:
 
 
 @click.command()
-def main():
+@click.option("-z", "--compressed", is_flag=True, help="Whether to load the probmaps compressed.")
+def main(compressed):
     maps = {}
-    for file in tqdm(Path().glob("probs_*.zpickle"), desc="Merging probmaps.", smoothing=0):
-        with zstd.open(file, "rb") as h:
+    if compressed:
+        glob = "probs_*.zpickle"
+        out = "merged.zpickle"
+        opener = zstd.open
+    else:
+        glob = "probs_*.pickle"
+        out = "merged.pickle"
+        opener = open
+
+    found = list(Path().glob(glob))
+    click.echo(f"Found {len(found)} probmap files.")
+    if not found:
+        return
+    
+    for file in tqdm(found, desc="Merging probmaps.", smoothing=0):
+        with opener(file, "rb") as h, tqdm(total=len(all_configs) * len(all_error_models), desc=f"Loading probmap {file}.", smoothing=0, leave=False) as bar:
+            i = 0
             while True:
                 try:
                     full, prob_map = pickle.load(h)
+                    bar.update(1)
+                    i += 1
                     if full not in maps:
                         maps[full] = prob_map
                     else:
@@ -35,9 +53,10 @@ def main():
                 except EOFError:
                     break
                 except pickle.UnpicklingError:
-                    print(f"Bad unpickling, the probs file {file} is likely truncated.")
+                    click.echo(f"Bad unpickling, the probs file {file} is likely truncated.")
                     break
-    with open("merged.pickle", "wb") as f:
+            click.echo(f"Loaded {i} probmaps from {file}.")
+    with opener(out, "wb") as f:
         pickle.dump(maps, f)
 
 
diff --git a/analysis/scalarmults/simulate.ipynb b/analysis/scalarmults/simulate.ipynb
index 2d46555..1039f0c 100644
--- a/analysis/scalarmults/simulate.ipynb
+++ b/analysis/scalarmults/simulate.ipynb
@@ -25,7 +25,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 29,
    "id": "b4386513-cc14-434b-a748-2863f8657452",
    "metadata": {},
    "outputs": [],
@@ -80,7 +80,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 30,
    "id": "3463a7bd-34d8-458b-8ceb-dddf99de21dc",
    "metadata": {},
    "outputs": [],
@@ -97,10 +97,21 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 31,
    "id": "170c11fc-86cf-4eb1-bf4e-b2e44b2d7ac5",
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Scalar multipliers considered:  65\n",
+      "Scalar multipliers (with a combination of up-to two countermeasures) considered:  5265\n",
+      "Error models considered:  32\n",
+      "Total configurations considered:  168480\n"
+     ]
+    }
+   ],
    "source": [
     "nmults = len(all_mults)\n",
     "nmults_ctr = len(all_mults_with_ctr)\n",
@@ -369,7 +380,7 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "a07cf1d0-14d8-4c9b-a86c-6a7437d4e9dc",
+   "id": "40d129c0-a80b-441e-a5d4-5e338ff3568c",
    "metadata": {},
    "outputs": [],
    "source": []