diff --git a/analysis/config_types.py b/analysis/config_types.py index 66efcdc7..b4b10b56 100644 --- a/analysis/config_types.py +++ b/analysis/config_types.py @@ -133,6 +133,7 @@ def encode_ramped_trotter(self, **kwargs): self.trotter_order = kwargs.get("trotter_order", None) self.trotter_steps = kwargs.get("trotter_steps", None) self.phase_scale_factor = kwargs.get("phase_scale_factor", 1.01) + self.tensor_contraction_method = kwargs.get("tensor_contraction_method", None) # Validate trotter_steps if self.trotter_steps is not None: if not isinstance(self.trotter_steps, int) or self.trotter_steps < 1: @@ -147,6 +148,13 @@ def encode_ramped_trotter(self, **kwargs): f"phase_scale_factor = {self.phase_scale_factor} < 1.0 will map eigenvalue phases " f"outside the range [-π, π], which may cause phase wrapping issues. " f"Values >= 1.0 are recommended (default: 1.01).") + # Validate tensor_contraction_method + if self.tensor_contraction_method is not None: + valid_methods = ["auto", "incremental", "structured", "qualtran"] + if self.tensor_contraction_method not in valid_methods: + raise ValueError( + f"tensor_contraction_method must be one of {valid_methods}, " + f"got '{self.tensor_contraction_method}'") def encode_pauli_lcu(self, **kwargs): self._only_once() self.method = "pauli lcu" @@ -167,6 +175,7 @@ def _generate_TOML_table(self): self.save_if_present(table, "trotter_order") self.save_if_present(table, "trotter_steps") self.save_if_present(table, "phase_scale_factor") + self.save_if_present(table, "tensor_contraction_method") return table # ------------------------------------------------------------------------------------------------- diff --git a/analysis/unitary.py b/analysis/unitary.py index c957b8eb..9b132862 100644 --- a/analysis/unitary.py +++ b/analysis/unitary.py @@ -325,16 +325,18 @@ def encode_ramped_trotter( # Import the appropriate implementation if trotter_impl == 'flattened': from qhat.common.trotter_flattened import build_ramped_trotterized_unitary - # Pass combine_terms parameter only for flattened implementation + # Pass combine_terms and tensor_contraction_method parameters only for flattened implementation + tensor_contraction_method = getattr(config_unitary, 'tensor_contraction_method', None) return build_ramped_trotterized_unitary( pauli_strings.items(), method, tevol_hbar, Nsteps, - combine_terms=config_unitary.trotter_combine_terms) + combine_terms=config_unitary.trotter_combine_terms, + tensor_contraction_method=tensor_contraction_method) else: # 'original' from qhat.common.trotter_original import build_ramped_trotterized_unitary - # Original implementation doesn't support combine_terms parameter + # Original implementation doesn't support combine_terms or tensor_contraction_method parameters return build_ramped_trotterized_unitary( pauli_strings.items(), method, diff --git a/common/tests/test_trotter_flattened.py b/common/tests/test_trotter_flattened.py index ad816d77..0de969d7 100644 --- a/common/tests/test_trotter_flattened.py +++ b/common/tests/test_trotter_flattened.py @@ -1118,3 +1118,171 @@ def test_coefficient_symmetry_properties(self): # For these methods, verify basic properties # All coefficients should be real assert all(isinstance(c, (int, float)) for c in coeffs) + + +# ================================================================================== +# Test: Tensor contraction method selection +# ================================================================================== + +class TestTensorContractionMethods: + """Test that different tensor contraction methods work correctly and produce identical results.""" + + @pytest.mark.parametrize("num_steps", [1, 2, 3, 10, 50]) + @pytest.mark.parametrize("method", ["qualtran", "incremental", "structured", None]) + def test_all_methods_produce_unitary(self, num_steps, method): + """Test that all tensor contraction methods produce valid unitary matrices.""" + pauli_terms = [('X', 1.0), ('Z', 0.5), ('Y', 0.3)] + + trotter = Trotterization.from_method( + pauli_terms=pauli_terms, + method='second order', + time=0.1, + num_steps=num_steps, + tensor_contraction_method=method + ) + + U = trotter.tensor_contract() + + # Verify unitarity: U† U = I + U_dag_U = U.conj().T @ U + identity = np.eye(U.shape[0]) + assert np.allclose(U_dag_U, identity, atol=1e-10), \ + f"Method {method} with num_steps={num_steps} does not produce unitary matrix" + + @pytest.mark.parametrize("num_steps", [1, 2, 3, 10, 50]) + def test_all_methods_agree(self, num_steps): + """Test that all tensor contraction methods produce identical results.""" + pauli_terms = [('X', 1.0), ('Z', 0.5), ('Y', 0.3)] + methods = ['qualtran', 'incremental', 'structured', None] + + # Compute matrices for all methods + results = {} + for method in methods: + trotter = Trotterization.from_method( + pauli_terms=pauli_terms, + method='second order', + time=0.1, + num_steps=num_steps, + tensor_contraction_method=method + ) + results[str(method)] = trotter.tensor_contract() + + # Use incremental as reference + U_ref = results['incremental'] + + # Check that all methods produce the same matrix (within numerical precision) + for method_name, U in results.items(): + if method_name == 'incremental': + continue + max_diff = np.max(np.abs(U - U_ref)) + assert max_diff < 1e-14, \ + f"Method {method_name} differs from incremental by {max_diff:.2e} at num_steps={num_steps}" + + @pytest.mark.parametrize("num_steps", [1, 2, 3]) + def test_structured_works_for_small_num_steps(self, num_steps): + """Test that structured method now works for num_steps = 1, 2, 3 (after Option 1+2).""" + pauli_terms = [('X', 1.0), ('Z', 0.5)] + + # This should not raise an error anymore + trotter = Trotterization.from_method( + pauli_terms=pauli_terms, + method='second order', + time=0.1, + num_steps=num_steps, + tensor_contraction_method='structured' + ) + + U = trotter.tensor_contract() + + # Verify it's unitary + U_dag_U = U.conj().T @ U + identity = np.eye(U.shape[0]) + assert np.allclose(U_dag_U, identity, atol=1e-10) + + def test_pattern_detection_degenerate_case(self): + """Test that pattern detection handles num_steps=1 correctly.""" + pauli_terms = [('X', 1.0), ('Z', 0.5)] + + trotter = Trotterization.from_method( + pauli_terms=pauli_terms, + method='second order', + time=0.1, + num_steps=1 + ) + + pattern_info = trotter._detect_repeating_pattern() + + # Should return a pattern with the entire sequence + assert 'num_repeats' in pattern_info + assert pattern_info['num_repeats'] == 1 + assert 'repeating_terms' in pattern_info + assert len(pattern_info['repeating_terms']) == len(trotter.expanded_sequence) + assert len(pattern_info['start_terms']) == 0 + assert len(pattern_info['end_terms']) == 0 + + def test_auto_selection_behavior(self): + """Test that auto-selection chooses the right method based on num_steps.""" + pauli_terms = [('X', 1.0), ('Z', 0.5)] + + # For num_steps < 10, should use incremental + # For num_steps >= 10, should use structured (if pattern detected) + + # Small num_steps: verify it doesn't crash (auto selects incremental) + for num_steps in [1, 2, 5, 9]: + trotter = Trotterization.from_method( + pauli_terms=pauli_terms, + method='second order', + time=0.1, + num_steps=num_steps, + tensor_contraction_method=None # Auto + ) + U = trotter.tensor_contract() + assert U.shape == (2, 2) + + # Large num_steps: verify it doesn't crash (auto may select structured) + for num_steps in [10, 50]: + trotter = Trotterization.from_method( + pauli_terms=pauli_terms, + method='second order', + time=0.1, + num_steps=num_steps, + tensor_contraction_method=None # Auto + ) + U = trotter.tensor_contract() + assert U.shape == (2, 2) + + def test_qualtran_method_accessible(self): + """Test that the 'qualtran' method successfully calls Bloq.tensor_contract().""" + pauli_terms = [('X', 1.0), ('Z', 0.5)] + + trotter = Trotterization.from_method( + pauli_terms=pauli_terms, + method='second order', + time=0.1, + num_steps=5, + tensor_contraction_method='qualtran' + ) + + U = trotter.tensor_contract() + + # Should produce a valid unitary + U_dag_U = U.conj().T @ U + identity = np.eye(U.shape[0]) + assert np.allclose(U_dag_U, identity, atol=1e-10) + + def test_invalid_method_raises_error(self): + """Test that an invalid tensor_contraction_method is handled correctly.""" + # Note: Validation happens in analysis/config_types.py when used through the config system + # Here we test that an invalid method in tensor_contract() raises an error + pauli_terms = [('X', 1.0)] + + trotter = Trotterization( + pauli_terms=tuple(pauli_terms), + coefficients=(0.5, 0.5), + time=0.1, + num_steps=5, + tensor_contraction_method='invalid_method' # Invalid method + ) + + with pytest.raises(ValueError, match="Invalid tensor_contraction_method"): + trotter.tensor_contract() diff --git a/common/trotter_flattened.py b/common/trotter_flattened.py index 4d002825..82647bb7 100644 --- a/common/trotter_flattened.py +++ b/common/trotter_flattened.py @@ -9,7 +9,7 @@ import logging from functools import cached_property from math import cbrt -from typing import List, Sequence, Tuple, Union +from typing import List, Optional, Sequence, Tuple, Union import attrs import numpy as np @@ -324,6 +324,11 @@ class Trotterization(Bloq): num_steps: Number of Trotterization steps (time is divided by this). hbar: Reduced Planck constant (default 1.0). combine_terms: If True (default), combine adjacent identical terms. If False, keep all terms separate. + tensor_contraction_method: Optional method to force for tensor contraction. + None or "auto" = auto-select based on step count + "incremental" = force O(n) incremental contraction + "structured" = force O(log n) structured contraction (raises error if pattern not detected) + "qualtran" = use Qualtran's inherited Bloq.tensor_contract() method Example: >>> # Second-order Trotterization using named method @@ -353,6 +358,7 @@ class Trotterization(Bloq): num_steps: int hbar: float = 1.0 combine_terms: bool = True + tensor_contraction_method: Optional[str] = None def __attrs_post_init__(self): """Validate inputs.""" @@ -383,7 +389,8 @@ def from_method( time: float, num_steps: int, hbar: float = 1.0, - combine_terms: bool = True + combine_terms: bool = True, + tensor_contraction_method: Optional[str] = None ): """Create Trotterization using a named method. @@ -394,6 +401,10 @@ def from_method( num_steps: Number of Trotterization steps hbar: Reduced Planck constant (default 1.0) combine_terms: If True (default), combine adjacent identical terms. If False, keep all terms separate. + tensor_contraction_method: Optional method to force for tensor contraction. + None or "auto" = auto-select based on step count + "incremental" = force O(n) incremental contraction + "structured" = force O(log n) structured contraction (raises error if pattern not detected) Returns: Trotterization instance @@ -413,7 +424,8 @@ def from_method( time=time, num_steps=num_steps, hbar=hbar, - combine_terms=combine_terms + combine_terms=combine_terms, + tensor_contraction_method=tensor_contraction_method ) @cached_property @@ -507,21 +519,13 @@ def _t_complexity_(self) -> TComplexity: def tensor_contract(self) -> np.ndarray: """ - Optimized tensor contraction exploiting repeating structure. + Optimized tensor contraction - Phase 2 optimization: For multi-step Trotter, the expanded_sequence - has a repeating pattern due to term combining across step boundaries: - [start_block] + [repeating_block]^(n_repeats) + [end_block] - - We exploit this by: - 1. Detecting the repeating pattern - 2. Building matrices for start, repeating, and end blocks - 3. Using matrix exponentiation for the repeating part: O(log n) - 4. Combining: U_total = U_end @ U_repeat^n @ U_start - - Falls back to O(n) incremental approach if: - - num_steps is small (< 10) - overhead not worth it - - No clear repeating pattern detected + The tensor_contraction_method attribute can force a specific method: + - None or "auto": use auto-selection logic (default behavior) + - "incremental": force O(n) incremental contraction + - "structured": force O(log n) structured contraction (raises error if pattern not detected) + - "qualtran": use Qualtran's inherited Bloq.tensor_contract() method Returns: The full unitary matrix as a numpy array with shape (2^n_qubits, 2^n_qubits) @@ -543,31 +547,56 @@ def tensor_contract(self) -> np.ndarray: logger.debug(f"-- num_terms = {len(self.pauli_terms)}") logger.debug(f"-- expanded_sequence length = {len(self.expanded_sequence)}") + # Check if a specific method is forced via configuration + forced_method = self.tensor_contraction_method + if forced_method is not None and forced_method != "auto": + if forced_method == "qualtran": + logger.verbose(f"Using Qualtran Bloq.tensor_contract() method") + return super().tensor_contract() + elif forced_method == "incremental": + logger.verbose(f"Using FORCED O(n) incremental contraction") + return self._incremental_contraction() + elif forced_method == "structured": + logger.verbose(f"Using FORCED O(log n) structured contraction") + # Must detect pattern - do not fall back + logger.debug(f"Attempting to detect repeating pattern...") + pattern_info = self._detect_repeating_pattern() + if pattern_info['num_repeats'] < 1: + raise RuntimeError( + f"Forced structured tensor contraction but num_repeats={pattern_info['num_repeats']} < 1. " + f"Invalid pattern structure. " + f"Either use auto-selection or force incremental method." + ) + logger.debug(f"-- pattern_length = {pattern_info['pattern_length']}") + logger.debug(f"-- num_repeats = {pattern_info['num_repeats']}") + logger.debug(f"-- start_block length = {len(pattern_info['start_terms'])}") + logger.debug(f"-- repeating_block length = {len(pattern_info['repeating_terms'])}") + logger.debug(f"-- end_block length = {len(pattern_info['end_terms'])}") + return self._structured_contraction(pattern_info) + else: + raise ValueError( + f"Invalid tensor_contraction_method '{forced_method}'. " + f"Must be None, 'auto', 'incremental', 'structured', or 'qualtran'." + ) + + # Auto-selection logic (original behavior) # For small num_steps, incremental is faster (less overhead) if self.num_steps < 10: logger.verbose(f"Using O(n) incremental contraction (num_steps < 10)") return self._incremental_contraction() # Try to detect repeating pattern - logger.debug(f"Attempting to detect repeating pattern...") + logger.debug(f"Searching for repeating pattern...") pattern_info = self._detect_repeating_pattern() # If pattern detected, use O(log n) structured approach - if pattern_info['has_pattern'] and pattern_info['num_repeats'] >= 2: - logger.verbose(f"Pattern detected! Using O(log n) structured contraction") - logger.debug(f"-- pattern_length = {pattern_info['pattern_length']}") - logger.debug(f"-- num_repeats = {pattern_info['num_repeats']}") - logger.debug(f"-- start_block length = {len(pattern_info['start_terms'])}") - logger.debug(f"-- repeating_block length = {len(pattern_info['repeating_terms'])}") - logger.debug(f"-- end_block length = {len(pattern_info['end_terms'])}") - return self._structured_contraction(pattern_info) - else: - # Fall back to O(n) incremental approach - if pattern_info.get('has_pattern', False): - logger.verbose(f"Pattern detected but num_repeats < 2, using O(n) incremental contraction") - else: - logger.verbose(f"No repeating pattern detected, using O(n) incremental contraction") - return self._incremental_contraction() + logger.debug(f"Pattern for O(log n) structured contraction:") + logger.debug(f"-- pattern_length = {pattern_info['pattern_length']}") + logger.debug(f"-- num_repeats = {pattern_info['num_repeats']}") + logger.debug(f"-- start_block length = {len(pattern_info['start_terms'])}") + logger.debug(f"-- repeating_block length = {len(pattern_info['repeating_terms'])}") + logger.debug(f"-- end_block length = {len(pattern_info['end_terms'])}") + return self._structured_contraction(pattern_info) def _detect_repeating_pattern(self) -> dict: """ @@ -581,7 +610,6 @@ def _detect_repeating_pattern(self) -> dict: Returns: Dictionary with: - - has_pattern: bool - whether a pattern was detected - start_terms: list - terms before the repeating section - repeating_terms: list - the repeating unit - end_terms: list - terms after the repeating section @@ -590,10 +618,6 @@ def _detect_repeating_pattern(self) -> dict: """ seq = self.expanded_sequence - # Need at least 2 steps for a repeating pattern - if self.num_steps < 2: - return {'has_pattern': False} - # For P Pauli strings and C coefficients per step: # With term combining at boundaries, the repeating pattern length is: # P × C - C - 1 = C(P - 1) - 1 @@ -631,13 +655,14 @@ def _detect_repeating_pattern(self) -> dict: else: break - # Need at least 2 repetitions to be worth using matrix power - if num_repeats >= 2: + # Accept patterns with at least 1 repetition + # Note: num_repeats=1 provides no performance benefit over incremental, + # but allows structured method to work for edge cases (e.g., num_steps=2) + if num_repeats >= 1: # Found a valid pattern! end_block = seq[idx:] if idx < len(seq) else [] return { - 'has_pattern': True, 'start_terms': start_block, 'repeating_terms': candidate_pattern, 'end_terms': end_block, @@ -646,7 +671,15 @@ def _detect_repeating_pattern(self) -> dict: } # No repeating pattern found - return {'has_pattern': False} + # - repeating block is entire sequence with single repetition + # - start and end terms are empty + return { + 'start_terms': [], + 'repeating_terms': seq, + 'end_terms': [], + 'num_repeats': 1, + 'pattern_length': len(seq) + } def _structured_contraction(self, pattern_info: dict) -> np.ndarray: """ @@ -724,10 +757,10 @@ def _build_matrix_from_terms(self, terms: List[Tuple[int, float]]) -> np.ndarray def _incremental_contraction(self) -> np.ndarray: """ - Fallback O(n) incremental contraction (Phase 1 implementation). + O(n) incremental contraction Iterates through expanded_sequence and multiplies matrices one at a time. - Used when pattern detection fails or for small num_steps. + Used for small num_steps. Returns: Full unitary matrix @@ -811,7 +844,8 @@ def build_ramped_trotterized_unitary( method, timestep: float, numsteps: int, - combine_terms: bool = True + combine_terms: bool = True, + tensor_contraction_method: Optional[str] = None ): """Build a Trotterization using the old interface for backward compatibility. @@ -825,6 +859,11 @@ def build_ramped_trotterized_unitary( timestep: Time step for evolution numsteps: Number of Trotterization steps combine_terms: If True (default), combine adjacent identical terms. If False, keep all terms separate. + tensor_contraction_method: Optional method to force for tensor contraction. + None or "auto" = auto-select based on step count + "incremental" = force O(n) incremental contraction + "structured" = force O(log n) structured contraction (raises error if pattern not detected) + "qualtran" = use Qualtran's inherited Bloq.tensor_contract() method Returns: Trotterization instance @@ -845,5 +884,6 @@ def build_ramped_trotterized_unitary( time=timestep, num_steps=numsteps, hbar=1.0, - combine_terms=combine_terms + combine_terms=combine_terms, + tensor_contraction_method=tensor_contraction_method )