Source code for pennylane.labs.trotter_error.realspace.ho_state
# Copyright 2025 Xanadu Quantum Technologies Inc.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# http://www.apache.org/licenses/LICENSE-2.0
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Contains the HOState class which represents a wavefunction in the harmonic oscillator basis"""
from __future__ import annotations
from typing import Dict, Sequence, Tuple
import numpy as np
from scipy.sparse import csr_array, identity, kron
[docs]
class HOState:
"""Represent a wavefunction in the harmonic oscillator basis.
Args:
modes (int): the number of vibrational modes
gridpoints (int): the number of gridpoints used to discretize the state
state: (Union[scipy.sparse.csr_array, Dict[Tuple[int], float]]): a sparse state vector for the full wavefunction or a dictionary containing the interacting modes and their non-zero coefficients
**Examples**
Building an :class:`~.pennylane.labs.trotter_error.HOState` from a dictionary
>>> from pennylane.labs.trotter_error import HOState
>>> n_modes = 3
>>> gridpoints = 5
>>> state_dict = {(1, 2, 3): 1, (0, 3, 2): 1}
>>> HOState(n_modes, gridpoints, state_dict)
HOState(modes=3, gridpoints=5, <Compressed Sparse Row sparse array of dtype 'int64'
with 2 stored elements and shape (125, 1)>
Coords Values
(17, 0) 1
(38, 0) 1)
Building an :class:`~.pennylane.labs.trotter_error.HOState` from a ``scipy.sparse.csr_array``
>>> from scipy.sparse import csr_array
>>> import numpy as np
>>> gridpoints = 2
>>> n_modes = 2
>>> state_vector = csr_array(np.array([0, 1, 0, 0]))
>>> HOState(n_modes, gridpoints, state_vector)
HOState(modes=2, gridpoints=2, <COOrdinate sparse array of dtype 'int64'
with 1 stored elements and shape (4, 1)>
Coords Values
(1, 0) 1)
"""
def __init__(self, modes: int, gridpoints: int, state: csr_array):
if isinstance(state, csr_array):
if state.shape == (gridpoints**modes,):
state = state.reshape((gridpoints**modes, 1))
if state.shape != (gridpoints**modes, 1):
raise ValueError(
f"Dimension mismatch. Expected vector of shape {(gridpoints ** modes, 1)} but got shape {state.shape}."
)
self.vector = state
elif isinstance(state, dict):
self.vector = self._vector_from_dict(modes, gridpoints, state)
else:
raise TypeError(f"state must be of type csr_array or dict, got {type(state)}.")
self.gridpoints = gridpoints
self.modes = modes
self.dim = gridpoints**modes
@classmethod
def _vector_from_dict(
cls, modes: int, gridpoints: int, coeffs: Dict[Tuple[int], float]
) -> csr_array:
"""Construct an :class:`~.pennylane.labs.trotter_error.HOState` from a dictionary.
Args:
modes (int): the number of vibrational modes
gridpoints (int): the number of gridpoints used to discretize the state
coeffs (Dict[Tuple[int]]): a dictionary representation of the state
Returns:
HOState: an :class:`~.pennylane.labs.trotter_error.HOState` representation of the state vector
**Example**
>>> from pennylane.labs.trotter_error import HOState
>>> n_modes = 3
>>> gridpoints = 5
>>> state_dict = {(1, 2, 3): 1, (0, 3, 2): 1}
>>> HOState.from_dict(n_modes, gridpoints, state_dict)
HOState(modes=3, gridpoints=5, <Compressed Sparse Row sparse array of dtype 'int64'
with 2 stored elements and shape (125, 1)>
Coords Values
(17, 0) 1
(38, 0) 1)
"""
rows, cols, vals = [], [], []
for index, val in coeffs.items():
if len(index) != modes:
raise ValueError(
f"Number of modes given was {modes}, but {index} contains {len(index)} modes."
)
row = sum(i * (gridpoints**exp) for exp, i in enumerate(reversed(index)))
rows.append(row)
vals.append(val)
rows = np.array(rows)
cols = np.zeros_like(rows)
vals = np.array(vals)
return csr_array((vals, (rows, cols)), shape=(gridpoints**modes, 1))
[docs]
@classmethod
def zero_state(cls, modes: int, gridpoints: int) -> HOState:
"""Construct an :class:`~.pennylane.labs.trotter_error.HOState` whose vector is zero.
Args:
modes (int): the number of vibrational modes
gridpoints(int): the number of gridpoints used to discretize the state
Returns:
HOState: an :class:`~.pennylane.labs.trotter_error.HOState` representing the zero state
**Example**
>>> from pennylane.labs.trotter_error import HOState
>>> HOState.zero_state(5, 10)
HOState(modes=5, gridpoints=10, <Compressed Sparse Row sparse array of dtype 'float64'
with 0 stored elements and shape (100000, 1)>)
"""
return cls(modes, gridpoints, csr_array((gridpoints**modes, 1)))
def __add__(self, other: HOState) -> HOState:
if not isinstance(other, HOState):
raise TypeError(f"Can only add HOState with another HOState, got {type(other)}.")
return HOState(self.modes, self.gridpoints, self.vector + other.vector)
def __mul__(self, scalar: float) -> HOState:
return HOState(self.modes, self.gridpoints, scalar * self.vector)
__rmul__ = __mul__
[docs]
def dot(self, other: HOState) -> float:
"""Return the dot product of two :class:`~.pennylane.labs.trotter_error.HOState` objects.
Args:
other (HOState): the state to take the dot product with
Returns:
float: the dot product of the two states
**Example**
>>> from pennylane.labs.trotter_error import HOState
>>> n_modes = 3
>>> gridpoints = 5
>>> state_dict = {(1, 2, 3): 1, (0, 3, 2): 1}
>>> state1 = HOState(n_modes, gridpoints, state_dict)
>>> state1.dot(state1)
2
"""
if self.dim != other.dim:
raise ValueError(
f"Dimension mismatch. Attempting to dot product vectors of dimension {self.dim} and {other.dim}."
)
return ((self.vector.T).dot(other.vector))[0, 0]
def __repr__(self):
return f"HOState(modes={self.modes}, gridpoints={self.gridpoints}, {self.vector})"
[docs]
class VibronicHO:
"""Represent the tensor product of harmonic oscillator states.
Args:
states (int): the number of electronic states
modes (int): the number of vibrational modes
gridpoints (int): the number of gridpoints used to discretize the state
ho_states (Sequence[HOState]): a sequence of :class:`~.pennylane.labs.trotter_error.HOState` objects representing the harmonic oscillator states
**Example**
>>> from pennylane.labs.trotter_error import HOState, VibronicHO
>>> n_modes = 3
>>> n_states = 2
>>> gridpoints = 5
>>> state_dict = {(1, 2, 3): 1, (0, 3, 2): 1}
>>> state = HOState(n_modes, gridpoints, state_dict)
>>> VibronicHO(n_states, n_modes, gridpoints, [state, state])
VibronicHO([HOState(modes=3, gridpoints=5, <Compressed Sparse Row sparse array of dtype 'int64'
with 2 stored elements and shape (125, 1)>
Coords Values
(17, 0) 1
(38, 0) 1), HOState(modes=3, gridpoints=5, <Compressed Sparse Row sparse array of dtype 'int64'
with 2 stored elements and shape (125, 1)>
Coords Values
(17, 0) 1
(38, 0) 1)])
"""
def __init__(self, states: int, modes: int, gridpoints: int, ho_states: Sequence[HOState]):
if len(ho_states) != states:
raise ValueError(
f"Got {len(ho_states)} harmonic oscillator states, but expected {states}."
)
for ho_state in ho_states:
if ho_state.modes != modes:
raise ValueError(
f"Mode mismatch: given {modes} modes, but found an HOState on {ho_state.modes} modes."
)
if ho_state.gridpoints != gridpoints:
raise ValueError(
f"Gridpoint mismatch: given {gridpoints} gridpoints, but found an HOState on {ho_state.gridpoints} gridpoints."
)
self.states = states
self.modes = modes
self.gridpoints = gridpoints
self.ho_states = ho_states
def __repr__(self):
return f"VibronicHO({self.ho_states})"
def __add__(self, other: VibronicHO) -> VibronicHO:
if self.states != other.states:
raise ValueError(
f"Cannot add VibronicHO on {self.states} states with VibronicHO on {other.states} states."
)
if self.modes != other.modes:
raise ValueError(
f"Cannot add VibronicHO on {self.modes} modes with VibronicHO on {other.modes} modes."
)
if self.gridpoints != other.gridpoints:
raise ValueError(
f"Cannot add VibronicHO on {self.gridpoints} gridpoints with VibronicHO on {other.gridpoints} gridpoints."
)
return VibronicHO(
states=self.states,
modes=self.modes,
gridpoints=self.gridpoints,
ho_states=[x + y for x, y in zip(self.ho_states, other.ho_states)],
)
def __mul__(self, scalar: float) -> VibronicHO:
return VibronicHO(
states=self.states,
modes=self.modes,
gridpoints=self.gridpoints,
ho_states=[scalar * ho for ho in self.ho_states],
)
__rmul__ = __mul__
[docs]
@classmethod
def zero_state(cls, states: int, modes: int, gridpoints: int) -> VibronicHO:
"""Construct a :class:`~.pennylane.labs.trotter_error.VibronicHO` representing the zero state.
Args:
states (int): the number of electronic states
modes (int): the number of vibrational modes
gridpoints(int): the number of gridpoints used to discretize the state
Returns:
VibronicHO: a :class:`~.pennylane.labs.trotter_error.VibronicHO` representing the zero state
**Example**
>>> from pennylane.labs.trotter_error import VibronicHO
>>> VibronicHO.zero_state(2, 3, 5)
VibronicHO([HOState(modes=3, gridpoints=5, <Compressed Sparse Row sparse array of dtype 'float64'
with 0 stored elements and shape (125, 1)>), HOState(modes=3, gridpoints=5, <Compressed Sparse Row sparse array of dtype 'float64'
with 0 stored elements and shape (125, 1)>)])
"""
return cls(
states=states,
modes=modes,
gridpoints=gridpoints,
ho_states=[HOState.zero_state(modes, gridpoints)] * states,
)
[docs]
def dot(self, other: VibronicHO):
"""Return the dot product of two :class:`~.pennylane.labs.trotter_error.VibronicHO` objects.
Args:
other (VibronicHO): the state to take the dot product with
Returns:
float: the dot product of the two states
**Example**
>>> from pennylane.labs.trotter_error import HOState, VibronicHO
>>> n_modes = 3
>>> n_states = 2
>>> gridpoints = 5
>>> state_dict = {(1, 2, 3): 1, (0, 3, 2): 1}
>>> state = HOState(n_modes, gridpoints, state_dict)
>>> vo_state = VibronicHO(n_states, n_modes, gridpoints, [state, state])
>>> vo_state.dot(vo_state)
4
"""
return np.real(sum(x.dot(y) for x, y in zip(self.ho_states, other.ho_states)))
def _tensor_with_identity(op: csr_array, gridpoints: int, n_modes: int, mode: int) -> csr_array:
"""Return the tensor product of ``op`` with the ``gridpoints**mode`` dimensional identity matrix on the left, and the ``gridpoints ** (n_modes - mode - 1)`` dimensional identity matrix on the right."""
if mode == 0:
return kron(op, identity(gridpoints ** (n_modes - 1)))
if mode == n_modes - 1:
return kron(identity(gridpoints ** (n_modes - 1)), op)
id_left = identity(gridpoints**mode)
id_right = identity(gridpoints ** (n_modes - mode - 1))
return kron(id_left, kron(op, id_right))
_modules/pennylane/labs/trotter_error/realspace/ho_state
Download Python script
Download Notebook
View on GitHub