Use lazy module imports

src/qldpc/_util.py

@@ -17,12 +17,40 @@
 
 from __future__ import annotations
 
+import importlib.util
+import sys
 from collections.abc import Callable
-from typing import TypeVar
+from types import ModuleType
+from typing import TYPE_CHECKING, TypeVar
 
 CallableType = TypeVar("CallableType", bound=Callable[..., object])
 
 
+def lazy_import(name: str) -> ModuleType:
+    """Import a module lazily, deferring its execution until its first attribute access.
+
+    Uses importlib.util.LazyLoader so a heavy but rarely-used dependency stays out of ``import
+    qldpc`` while call sites keep using ``module.attr`` exactly as if it had been imported eagerly.
+    """
+    if (module := sys.modules.get(name)) is not None:
+        return module
+    spec = importlib.util.find_spec(name)
+    assert spec is not None and spec.loader is not None
+    spec.loader = importlib.util.LazyLoader(spec.loader)
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+# networkx is loaded lazily to keep it and its ~110 ms import out of ``import qldpc``.
+# Call sites import it as ``from qldpc._util import networkx as nx``.
+if TYPE_CHECKING:
+    import networkx
+else:
+    networkx = lazy_import("networkx")
+
+
 def format_docstring(**substitutions: object) -> Callable[[CallableType], CallableType]:
     """Substitute named values into a function's docstring via str.format.
 

src/qldpc/_util_test.py

@@ -15,7 +15,22 @@
 limitations under the License.
 """
 
-from qldpc._util import format_docstring
+import sys
+from types import ModuleType
+
+from qldpc._util import format_docstring, lazy_import
+
+
+def test_lazy_import() -> None:
+    """A lazily imported module is only executed on first attribute access."""
+    name = "colorsys"  # a small stdlib module not otherwise imported by the test suite
+    sys.modules.pop(name, None)
+
+    module = lazy_import(name)  # uncached path
+    assert isinstance(module, ModuleType)
+    assert callable(module.rgb_to_hls)  # force the deferred execution
+
+    assert lazy_import(name) is module  # cached path returns the same module
 
 
 def test_format_docstring() -> None:

src/qldpc/circuits/common.py

@@ -19,6 +19,7 @@
 
 import functools
 from collections.abc import Callable, Mapping, Sequence
+from types import ModuleType
 from typing import TYPE_CHECKING, ParamSpec, TypeVar
 
 import galois
@@ -29,15 +30,26 @@
 from qldpc import codes, math
 
 ####################################################################################################
-# try to import tsim and define a circuit type that may be either of stim.Circuit or tsim.Circuit
-try:
+# define a circuit type that may be either a stim.Circuit or an (optional) tsim.Circuit
+
+if TYPE_CHECKING:
     import tsim
 
     stim_or_tsim_Circuit = TypeVar("stim_or_tsim_Circuit", stim.Circuit, tsim.Circuit)
-except ImportError:  # pragma: no cover
-    if not TYPE_CHECKING:
-        tsim = None
-        stim_or_tsim_Circuit = TypeVar("stim_or_tsim_Circuit", bound=stim.Circuit)
+else:
+    stim_or_tsim_Circuit = TypeVar("stim_or_tsim_Circuit", bound=stim.Circuit)
+
+
+def _load_tsim_if_installed() -> ModuleType | None:
+    """Lazily import the optional tsim package, returning None if it is not installed."""
+    try:
+        import tsim
+
+        return tsim
+    except ImportError:  # pragma: no cover
+        return None
+
+
 ####################################################################################################
 
 
@@ -78,10 +90,15 @@ def with_remapped_qubits(
     Returns:
         The input circuit with remapped qubits.
     """
+    tsim = _load_tsim_if_installed()
     if tsim is not None and isinstance(circuit, tsim.Circuit):
         output = with_remapped_qubits(circuit.stim_circuit, qubit_map, inverse=inverse)
         return tsim.Circuit.from_stim_program(output)
 
+    # tsim loads lazily, so `tsim.Circuit` above is untyped and does not narrow `circuit`; having
+    # ruled out a tsim circuit, the remaining possibility for the type variable is a stim.Circuit
+    assert isinstance(circuit, stim.Circuit)
+
     qubit_map = (
         qubit_map
         if isinstance(qubit_map, Mapping)

src/qldpc/circuits/memory/syndrome_measurement.py

@@ -20,10 +20,10 @@
 import abc
 import collections
 
-import networkx as nx
 import stim
 
 from qldpc import codes
+from qldpc._util import networkx as nx
 from qldpc.objects import Pauli
 
 from ..bookkeeping import MeasurementRecord, QubitIDs

src/qldpc/circuits/noise_model.py

@@ -118,7 +118,7 @@ def bad_qubit_noise(op: stim.CircuitInstruction) -> NoiseRule | None:
 
 from qldpc._util import format_docstring
 
-from .common import stim_or_tsim_Circuit, tsim, with_remapped_qubits
+from .common import _load_tsim_if_installed, stim_or_tsim_Circuit, with_remapped_qubits
 
 ####################################################################################################
 # global constants
@@ -232,6 +232,7 @@ def op_type(op_name: str) -> str:
 
 def as_noiseless_circuit(circuit: stim_or_tsim_Circuit) -> stim_or_tsim_Circuit:
     """Wrap a circuit in a noiseless, one-repitition stim.CircuitRepeatBlock."""
+    tsim = _load_tsim_if_installed()
     if tsim is not None and isinstance(circuit, tsim.Circuit):
         output = as_noiseless_circuit(circuit.stim_circuit)
         return tsim.Circuit.from_stim_program(output)
@@ -1031,6 +1032,7 @@ def noisy_circuit(
         Returns:
             The input circuit with added noise.
         """
+        tsim = _load_tsim_if_installed()
         if tsim is not None and isinstance(circuit, tsim.Circuit):
             output = self.noisy_circuit(
                 circuit.stim_circuit,

src/qldpc/codes/classical.py

@@ -21,12 +21,12 @@
 from collections.abc import Sequence
 
 import galois
-import networkx as nx
 import numpy as np
 import numpy.typing as npt
 import sympy
 
 from qldpc import abstract
+from qldpc._util import networkx as nx
 
 from .common import ClassicalCode
 

src/qldpc/codes/common.py

@@ -28,7 +28,6 @@
 from typing import Any, TypeVar, cast
 
 import galois
-import networkx as nx
 import numpy as np
 import numpy.typing as npt
 import scipy.linalg
@@ -37,6 +36,7 @@
 from typing_extensions import Self
 
 from qldpc import abstract, decoders, external, math
+from qldpc._util import networkx as nx
 from qldpc.math import IntegerArray
 from qldpc.objects import PAULIS_XZ, Node, Pauli, PauliXZ, QuditPauli
 

src/qldpc/codes/quantum.py

@@ -27,7 +27,6 @@
 from typing import TypeVar
 
 import galois
-import networkx as nx
 import numpy as np
 import numpy.typing as npt
 import scipy
@@ -36,6 +35,7 @@
 
 import qldpc
 from qldpc import abstract
+from qldpc._util import networkx as nx
 from qldpc.objects import CayleyComplex, ChainComplex, Node, Pauli, PauliXZ, QuditPauli
 
 from .classical import (

src/qldpc/decoders/custom.py

@@ -22,9 +22,8 @@
 import itertools
 import warnings
 from collections.abc import Callable, Collection, Iterator, Sequence
-from typing import Any, Protocol
+from typing import TYPE_CHECKING, Any, Protocol
 
-import cvxpy
 import galois
 import numpy as np
 import numpy.typing as npt
@@ -37,6 +36,9 @@
 
 from .dems import DetectorErrorModelArrays
 
+if TYPE_CHECKING:
+    import cvxpy
+
 PLACEHOLDER_ERROR_RATE = 1e-3  # required for some decoding methods
 
 
@@ -593,6 +595,8 @@ class ILPDecoder:
     """
 
     def __init__(self, matrix: IntegerArray, **decoder_args: object) -> None:
+        import cvxpy
+
         self.modulus = type(matrix).order if isinstance(matrix, galois.FieldArray) else 2
         if not galois.is_prime(self.modulus):
             raise ValueError("ILP decoding only supports prime number fields")
@@ -623,6 +627,7 @@ def __init__(self, matrix: IntegerArray, **decoder_args: object) -> None:
 
     def decode(self, syndrome: npt.NDArray[np.int_]) -> npt.NDArray[np.int_]:
         """Decode an error syndrome and return an inferred error."""
+        import cvxpy
 
         # identify all constraints
         constraints = self.variable_constraints + self.cvxpy_constraints_for_syndrome(syndrome)
@@ -649,6 +654,8 @@ def cvxpy_constraints_for_syndrome(
         to
         `expression = val + sum_j q^j s_j`.
         """
+        import cvxpy
+
         syndrome = np.asarray(syndrome, dtype=int) % self.modulus
 
         constraints = []

src/qldpc/decoders/retrieval.py

@@ -22,10 +22,8 @@
 from collections.abc import Sequence
 
 import galois
-import ldpc
 import numpy as np
 import numpy.typing as npt
-import pymatching
 import scipy.sparse
 import stim
 
@@ -123,6 +121,8 @@ def get_decoder_BP_OSD(
     - Documentation: https://software.roffe.eu/ldpc/quantum_decoder.html
     - Reference: https://arxiv.org/abs/2005.07016
     """
+    import ldpc
+
     pcm, error_channel = _to_ldpc_inputs(pcm_or_dem, error_rate, error_channel)
     return ldpc.BpOsdDecoder(pcm, error_channel=error_channel, **decoder_args)
 
@@ -155,6 +155,8 @@ def get_decoder_BP_LSD(
     - Documentation: https://software.roffe.eu/ldpc/quantum_decoder.html
     - Reference: https://arxiv.org/abs/2406.18655
     """
+    import ldpc
+
     pcm, error_channel = _to_ldpc_inputs(pcm_or_dem, error_rate, error_channel)
     return ldpc.bplsd_decoder.BpLsdDecoder(pcm, error_channel=error_channel, **decoder_args)
 
@@ -190,6 +192,8 @@ def get_decoder_BF(
       - https://arxiv.org/abs/2103.08049
       - https://arxiv.org/abs/2209.01180
     """
+    import ldpc
+
     pcm, error_channel = _to_ldpc_inputs(pcm_or_dem, error_rate, error_channel)
     return ldpc.BeliefFindDecoder(pcm, error_channel=error_channel, **decoder_args)
 
@@ -265,6 +269,8 @@ def get_decoder_MWPM(
         )
 
     # retrieve a matching decoder from pymatching
+    import pymatching
+
     return pymatching.Matching.from_check_matrix(pcm, **decoder_args)
 
 

src/qldpc/objects.py

@@ -25,11 +25,11 @@
 from typing import Literal
 
 import galois
-import networkx as nx
 import numpy as np
 import numpy.typing as npt
 
 from qldpc import abstract
+from qldpc._util import networkx as nx
 
 
 class Pauli(enum.Enum):