More caching in automorphism group calculation (#210)

* cache automorphism group data

* even more caching

* fix cache coverage

* move group generator caching to abstract.py

* linting fix

qldpc/abstract.py

@@ -412,6 +412,10 @@ def from_name(name: str, field: int | None = None) -> Group:
         generators = [GroupMember(gen) for gen in external.groups.get_generators(standardized_name)]
         return Group(*generators, name=standardized_name, field=field)
 
+    def hashable_generators(self) -> tuple[tuple[int, ...], ...]:
+        """Generators of this group in a hashable form."""
+        return tuple(tuple(generator) for generator in self.generators)
+
 
 ################################################################################
 # elements of a group algebra

qldpc/abstract_test.py

@@ -46,6 +46,8 @@ def test_permutation_group() -> None:
         gen = galois.GF(2)([[1]])
         abstract.Group.from_generating_mats(gen, field=3)
 
+    assert isinstance(hash(group.hashable_generators()), int)
+
 
 def test_trivial_group() -> None:
     """Trivial group tests."""

qldpc/cache.py

@@ -37,7 +37,10 @@ def get_disk_cache(cache_name: str, *, cache_dir: str | None = None) -> diskcach
 
 
 def use_disk_cache(
-    cache_name: str, *, cache_dir: str | None = None
+    cache_name: str,
+    *,
+    cache_dir: str | None = None,
+    key_func: Callable[..., Hashable] | None = None,
 ) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
     """Decorator to cache results to disk."""
 
@@ -49,7 +52,10 @@ def decorator(function: Callable[..., Any]) -> Callable[..., Any]:
         def function_with_cache(*args: Hashable, **kwargs: Hashable) -> Any:
             # retrieve results from cache, if available
             cache = get_disk_cache(cache_name, cache_dir=cache_dir)
-            key = args + tuple(kwargs.items())
+            if key_func is not None:
+                key = key_func(*args, **kwargs)
+            else:
+                key = args + tuple(kwargs.items())
             if key in cache:
                 return cache[key]
 

qldpc/cache_test.py

@@ -49,3 +49,10 @@ def get_five(_: str) -> int:
         get_five("test_arg")  # save results to cache
         assert cache == {("test_arg",): 5}  # check cache
         assert cache[("test_arg",)] == get_five("test_arg")  # retrieve results
+
+        @qldpc.cache.use_disk_cache("test_name", key_func=lambda _: None)
+        def get_six(_: str) -> int:
+            return 6
+
+        assert get_six("test_arg") == 6
+        assert cache == {("test_arg",): 5, None: 6}

qldpc/circuits.py

@@ -24,9 +24,11 @@
 import numpy as np
 import stim
 
-from qldpc import abstract, codes
+from qldpc import abstract, cache, codes
 from qldpc.objects import Pauli, conjugate_xz, op_to_string
 
+CACHE_NAME = "qldpc_automorphisms"
+
 
 def restrict_to_qubits(func: Callable[..., stim.Circuit]) -> Callable[..., stim.Circuit]:
     """Restrict a circuit constructor to qubit-based codes."""
@@ -242,8 +244,20 @@ def get_transversal_automorphism_group(
     group_gates = abstract.Group(*map(abstract.GroupMember, column_perms))
 
     # intersect the groups above to find the group generated by a transversal gate set
-    group_aut_sympy = group_code.to_sympy().subgroup_search(group_gates.to_sympy().contains)
-    return abstract.Group.from_sympy(group_aut_sympy, field=code.field.order)
+    group_aut_gens = _sympy_group_intersection_generators(group_code, group_gates)
+    return abstract.Group(*map(abstract.GroupMember, group_aut_gens))
+
+
+@cache.use_disk_cache(
+    CACHE_NAME,
+    key_func=lambda xx, yy: (xx.hashable_generators(), yy.hashable_generators()),
+)
+def _sympy_group_intersection_generators(
+    group_a: abstract.Group, group_b: abstract.Group
+) -> tuple[tuple[int, ...], ...]:
+    """Get the generators of the intersection of two Sympy permutation groups."""
+    group_sympy = group_a.to_sympy().subgroup_search(group_b.to_sympy().contains)
+    return abstract.Group(group_sympy).hashable_generators()
 
 
 @restrict_to_qubits