Merge branch 'master' into pl_qualtran_prototype

Author: austingmhuang

doc/releases/changelog-dev.md

@@ -856,6 +856,10 @@ Here's a list of deprecations made this release. For a more detailed breakdown o
 
 <h3>Documentation 📝</h3>
 
+* The functions in `qml.qchem.vibrational` are updated to include additional information about the 
+  theory and input arguments.
+  [(#6918)](https://github.com/PennyLaneAI/pennylane/pull/6918)
+
 * The usage examples for `qml.decomposition.DecompositionGraph` have been updated.
   [(#7692)](https://github.com/PennyLaneAI/pennylane/pull/7692)
 

pennylane/compiler/python_compiler/jax_utils.py

@@ -18,6 +18,7 @@
 from typing import Callable, TypeAlias
 
 import jaxlib
+from catalyst import QJIT
 from jaxlib.mlir.dialects import stablehlo as jstablehlo  # pylint: disable=no-name-in-module
 from jaxlib.mlir.ir import Context as jContext  # pylint: disable=no-name-in-module
 from jaxlib.mlir.ir import Module as jModule  # pylint: disable=no-name-in-module
@@ -30,6 +31,7 @@
 from xdsl.dialects import tensor as xtensor
 from xdsl.dialects import transform as xtransform
 from xdsl.parser import Parser as xParser
+from xdsl.traits import SymbolTable as xSymbolTable
 
 from pennylane.compiler.python_compiler.quantum_dialect import QuantumDialect
 
@@ -42,9 +44,7 @@ def _module_inline(func: JaxJittedFunction, *args, **kwargs) -> jModule:
 
 
 def module(func: JaxJittedFunction) -> Callable[..., jModule]:
-    """
-    Decorator for _module_inline
-    """
+    """Decorator for _module_inline"""
 
     @wraps(func)
     def wrapper(*args, **kwargs) -> jModule:
@@ -54,18 +54,14 @@ def wrapper(*args, **kwargs) -> jModule:
 
 
 def _generic_inline(func: JaxJittedFunction, *args, **kwargs) -> str:  # pragma: no cover
-    """
-    Create the generic textual representation for the jax.jit'ed function
-    """
+    """Create the generic textual representation for the jax.jit'ed function"""
     lowered = func.lower(*args, **kwargs)
     mod = lowered.compiler_ir()
     return mod.operation.get_asm(binary=False, print_generic_op_form=True, assume_verified=True)
 
 
 def generic(func: JaxJittedFunction) -> Callable[..., str]:  # pragma: no cover
-    """
-    Decorator for _generic_inline.
-    """
+    """Decorator for _generic_inline."""
 
     @wraps(func)
     def wrapper(*args, **kwargs) -> str:
@@ -125,12 +121,47 @@ def wrapper(*_, **__):
 
 
 def xdsl_module(func: JaxJittedFunction) -> Callable[..., xbuiltin.ModuleOp]:  # pragma: no cover
-    """
-    Decorator for _xdsl_module_inline
-    """
+    """Decorator for _xdsl_module_inline"""
 
     @wraps(func)
     def wrapper(*args, **kwargs) -> xbuiltin.ModuleOp:
         return _xdsl_module_inline(func, *args, **kwargs)
 
     return wrapper
+
+
+def inline_module(
+    from_mod: xbuiltin.ModuleOp, to_mod: xbuiltin.ModuleOp, change_main_to: str = None
+) -> None:
+    """Inline the contents of one xDSL module into another xDSL module. The inlined body is appended
+    to the end of ``to_mod``."""
+    if change_main_to:
+        main = xSymbolTable.lookup_symbol(from_mod, "main")
+        if main is not None:
+            assert isinstance(main, xfunc.FuncOp)
+            main.properties["sym_name"] = xbuiltin.StringAttr(change_main_to)
+
+    for op in from_mod.body.ops:
+        xSymbolTable.insert_or_update(to_mod, op.clone())
+
+
+def inline_jit_to_module(func: JaxJittedFunction, mod: xbuiltin.ModuleOp, *args, **kwargs) -> None:
+    """Inline a ``jax.jit``-ed Python function to an xDSL module. The inlined body is appended
+    to the end of ``mod``."""
+    func_mod = _xdsl_module_inline(func, *args, **kwargs)
+    inline_module(func_mod, mod, change_main_to=func.__name__)
+
+
+def xdsl_from_qjit(func: QJIT) -> Callable[..., xbuiltin.ModuleOp]:
+    """Decorator to convert QJIT-ed functions into xDSL modules."""
+
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        func.jaxpr, *_ = func.capture(args, **kwargs)
+        mlir_module = func.generate_ir()
+        generic_str = mlir_module.operation.get_asm(
+            binary=False, print_generic_op_form=True, assume_verified=True
+        )
+        return parse_generic_to_xdsl_module(generic_str)
+
+    return wrapper

pennylane/labs/tests/vibrational/test_pes_generator.py

@@ -27,7 +27,7 @@
 
 h5py = pytest.importorskip("h5py")
 
-# pylint: disable=too-many-arguments, protected-access, too-many-positional-arguments
+# pylint: disable=too-many-arguments, protected-access, too-many-positional-arguments, unsubscriptable-object
 
 ref_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "test_ref_files")
 

pennylane/qchem/vibrational/localize_modes.py

@@ -174,26 +174,38 @@ def _localize_modes(freqs, vecs):
 
 
 def localize_normal_modes(freqs, vecs, bins=[2600]):
-    """
-    Localizes vibrational normal modes.
+    r"""Computes spatially localized vibrational normal modes.
 
-    The normal modes are localized by separating frequencies into specified ranges following the
-    procedure described in `J. Chem. Phys. 141, 104105 (2014)
+    The vibrational normal modes are localized using a localizing unitary following the procedure
+    described in `J. Chem. Phys. 141, 104105 (2014)
     <https://pubs.aip.org/aip/jcp/article-abstract/141/10/104105/74317/
-    Efficient-anharmonic-vibrational-spectroscopy-for?redirectedFrom=fulltext>`_.
+    Efficient-anharmonic-vibrational-spectroscopy-for?redirectedFrom=fulltext>`_. The localizing
+    unitary :math:`U` is defined in terms of the normal and local coordinates, :math:`q` and
+    :math:`\tilde{q}`, respectively as:
+
+    .. math::
+
+        \tilde{q} = \sum_{j=1}^M U_{ij} q_j,
+
+    where :math:`M` is the number of modes. The normal modes
+    can be separately localized, to prevent mixing between specific groups of normal modes, by
+    defining frequency ranges in ``bins``. For instance, ``bins = [2600]`` allows to separately
+    localize modes that have frequencies above and below :math:`2600` reciprocal centimetre (:math:`\text{cm}^{-1}`).
+    Similarly, ``bins = [1300, 2600]`` allows to separately localize modes in three groups that have
+    frequencies below :math:`1300`, between :math:`1300-2600` and above :math:`2600`.
 
     Args:
-        freqs (list[float]): normal mode frequencies in ``cm^-1``
-        vecs (TensorLike[float]): displacement vectors for normal modes
-        bins (list[float]): List of upper bound frequencies in ``cm^-1`` for creating separation bins .
-            Default is ``[2600]`` which means having one bin for all frequencies between ``0`` and  ``2600 cm^-1``.
+        freqs (TensorLike[float]): normal mode frequencies in reciprocal centimetre (:math:`\text{cm}^{-1}`).
+        vecs (TensorLike[float]): displacement vectors of the normal modes
+        bins (List[float]): grid of frequencies for grouping normal modes.
+            Default is ``[2600]``.
 
     Returns:
         tuple: A tuple containing the following:
-         - list[float] : localized frequencies
-         - TensorLike[float] : localized displacement vectors
-         - TensorLike[float] : localization matrix describing the relationship between
-           original and localized modes.
+         - TensorLike[float] : localized frequencies in reciprocal centimetre (:math:`\text{cm}^{-1}`).
+         - List[TensorLike[float]] : localized displacement vectors
+         - TensorLike[float] : localization matrix describing the relationship between the
+           original and the localized modes
 
     **Example**
 
@@ -209,7 +221,7 @@ def localize_normal_modes(freqs, vecs, bins=[2600]):
     ...                      [-5.49709883e-17,  7.49851221e-08, -2.77912798e-02]]])
     >>> freqs_loc, vecs_loc, uloc = qml.qchem.localize_normal_modes(freqs, vectors)
     >>> freqs_loc
-    array([1332.62008773, 2296.73455892, 2296.7346082 ])
+    array([1332.62013257, 2296.73453455, 2296.73460655])
 
     """
     if not bins: