Repository: PennyLaneAI/pennylane
Pull request #3866 (2023-03-03 22:07:34)

Allow phase to be a callable on RydbergHamiltonian
**Context:**
Current implementation of the `RydbergHamiltonian` stores `phase` on the observable instead of as a parametrized coefficient, so it can't be a callable. However, on the QuEra device, phase is tuneable on the timescale of the pulse program, so we want to allow phase to be a callable rather than only a fixed value.

**Description of the Change:**

Primary change: A wrapper function is added that takes the amplitude and phase and combines them into a single function, stored on a callable class (so it can be accurately identified later). The exact form of the returned function depends on which, if any, of the amplitude and phase were callables. The `rydberg_drive` function is refactored to use this wrapper class. 

This introduces a new problem, however, which is that there is no longer a 1-to-1 correspondence between the parameters and the callables for this operator. The mismatch occurs in part because `phase_parameters` and `amplitude_parameters` now both go to a combined function, even though the user defined the functions and would expect to pass the parameters separately. It also occurs because now there are two separate functions that need both sets of parameters. The mapping is now 2-to-2 in the event that both amplitude and phase are callable, and 1-to-2 if only one of them is.

Therefore, a few additional changes are made:
1. A function is added called `_rydberg_reorder_parameters`, which loops through the coefficients and reorganizes the parameters input by the user. This allows the user to input parameters 1-to-1 with the callables they've defined, while ensuring that they are passed to the call function so they match with the wrapped and otherwise modified callables.
2. A `reorder_fn` attribute is added to the `ParametrizedHamiltonianPytree`, which in the case that H is a `RydbergHamiltonian` is set to `_rydberg_reorder_parameters`. If not `None`, it is called before passing the parameters to evaluate the Hamiltonian. Future objects like `RydbergHamiltonain` (for example an analogous Transmon system) should be able to use the same implementation with a different `reorder_fn`.
3. The `__add__` method on `RydbergHamiltonian` now returns a `RydbergHamiltonian` instead of `ParametrizedHamiltonian` when added to something other than a `RydbergHamiltonian`, to ensure the resulting object is still callable.

**Benefits:**
`phase` can be a callable.

Users can do:
```
H_global = rydberg_drive(amp, phase, detuning, wires=[0])
```
where any or all of `amp`, `phase` and `detuning` are callalbes, and then pass the parameters as `[amplitude_parameters, phase_parameters, detuning_parameters]` (omitting any that aren't relevant if one of the variables in a float instead of a callable), and all of the mapping will take place in the call function, so we don't have to introduce any weird UI things.

github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [`doc/releases/changelog-dev.md`](/PennyLaneAI/pennylane/blob/master/doc/releases/changelog-dev.md) with:
* A one-to-two sentence description of the change. You may include a small working example for new features.
* A link back to this PR.
* Your name (or GitHub username) in the contributors section.
lillian542: Example code to use:

```
from jax import numpy as jnp
import jax

import pennylane as qml
from pennylane.pulse.rydberg_hamiltonian import rydberg_drive
```

```
def amp(p, t):
    return p[0] * jnp.exp(-(t-p[1])**2/(2*p[2]**2))

def phase(p, t):
    return p * t

def detuning(p, t):
    return jnp.cos(t) * p
```

```
H_global = rydberg_drive(amp, phase, detuning, wires=[0])

params = [jnp.array([1.2, 2.3, 3.4]), 4.5, 5.6]

H_global(params, 2)
```

```
dev = qml.device("default.qubit.jax", wires=range(15))

@jax.jit
@qml.qnode(dev, interface="jax")
def circuit(params):
    qml.evolve(H_global)(params, 2)
    return [qml.expval(qml.PauliZ(i)) for i in range(3)]

circuit(params)
```

You can make any of the 3 floats instead and adjust `params` accordingly.
lillian542: Ah, thanks! Yes, that will have to be included in the `__add__` method for `ParametrizedHamiltonian`.

Also I still need to add tests 😅 
lillian542: [sc-34635]
codecov[bot]: ## [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3866?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#3866](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3866?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (294815b) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/8cfbe6606881846d11be79620b4ec5ce0124d324?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (8cfbe66) will **decrease** coverage by `0.01%`.
> The diff coverage is `98.11%`.

```diff
@@            Coverage Diff             @@
##           master    #3866      +/-   ##
==========================================
- Coverage   99.70%   99.70%   -0.01%     
==========================================
  Files         342      342              
  Lines       29853    29935      +82     
==========================================
+ Hits        29766    29846      +80     
- Misses         87       89       +2     
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3866?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/pulse/parametrized\_hamiltonian.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3866?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3B1bHNlL3BhcmFtZXRyaXplZF9oYW1pbHRvbmlhbi5weQ==) | `98.80% <50.00%> (-1.20%)` | :arrow_down: |
| [pennylane/pulse/rydberg\_hamiltonian.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3866?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3B1bHNlL3J5ZGJlcmdfaGFtaWx0b25pYW4ucHk=) | `99.27% <98.92%> (-0.73%)` | :arrow_down: |
| [pennylane/pulse/parametrized\_hamiltonian\_pytree.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3866?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3B1bHNlL3BhcmFtZXRyaXplZF9oYW1pbHRvbmlhbl9weXRyZWUucHk=) | `100.00% <100.00%> (ø)` | |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)


Patch:
diff --git a/pennylane/pulse/parametrized_hamiltonian_pytree.py b/pennylane/pulse/parametrized_hamiltonian_pytree.py
index f8f162e3c..3882b6584 100644
--- a/pennylane/pulse/parametrized_hamiltonian_pytree.py
+++ b/pennylane/pulse/parametrized_hamiltonian_pytree.py
@@ -23,6 +23,7 @@ from jax.tree_util import register_pytree_node_class
 import pennylane as qml
 
 from .parametrized_hamiltonian import ParametrizedHamiltonian
+from .rydberg_hamiltonian import RydbergHamiltonian, _rydberg_reorder_parameters
 
 
 @register_pytree_node_class
@@ -33,6 +34,7 @@ class ParametrizedHamiltonianPytree:
     mat_fixed: Optional[Union[jnp.ndarray, sparse.CSR]]
     mats_parametrized: Tuple[Union[jnp.ndarray, sparse.CSR], ...]
     coeffs_parametrized: Tuple[Callable]
+    reorder_fn: Callable
 
     @staticmethod
     def from_hamiltonian(H: ParametrizedHamiltonian, *, dense: bool = False, wire_order=None):
@@ -57,7 +59,18 @@ class ParametrizedHamiltonianPytree:
         mats_parametrized = tuple(
             make_array(qml.matrix(op, wire_order=wire_order)) for op in H.ops_parametrized
         )
-        return ParametrizedHamiltonianPytree(mat_fixed, mats_parametrized, H.coeffs_parametrized)
+
+        if isinstance(H, RydbergHamiltonian):
+            return ParametrizedHamiltonianPytree(
+                mat_fixed,
+                mats_parametrized,
+                H.coeffs_parametrized,
+                reorder_fn=_rydberg_reorder_parameters,
+            )
+
+        return ParametrizedHamiltonianPytree(
+            mat_fixed, mats_parametrized, H.coeffs_parametrized, reorder_fn=None
+        )
 
     def __call__(self, pars, t):
         if self.mat_fixed is not None:
@@ -66,6 +79,9 @@ class ParametrizedHamiltonianPytree:
         else:
             ops = ()
             coeffs = ()
+
+        if self.reorder_fn:
+            pars = self.reorder_fn(pars, self.coeffs_parametrized)
         coeffs = coeffs + tuple(f(p, t) for f, p in zip(self.coeffs_parametrized, pars))
         return LazyDotPytree(coeffs, ops + self.mats_parametrized)
 
@@ -77,20 +93,23 @@ class ParametrizedHamiltonianPytree:
         """
         matrices = (self.mat_fixed, self.mats_parametrized)
         param_coeffs = self.coeffs_parametrized
-        return (matrices, param_coeffs)
+        reorder_fn = self.reorder_fn
+        return (matrices, param_coeffs, reorder_fn)
 
     @classmethod
-    def tree_unflatten(cls, param_coeffs: tuple, matrices: tuple):
+    def tree_unflatten(cls, param_coeffs: tuple, matrices: tuple, reorder_fn: callable):
         """Function used by ``jax`` to unflatten the JaxParametrizedOperator.
 
         Args:
             param_coeffs (tuple): tuple containing the parametrized coefficients of the class
             matrices (tuple): tuple containing the matrices of the class
+            reorder_fn(callable): callable or None indicating how parameters should be
+                re-orderd to pass to the __call__ method
 
         Returns:
             JaxParametrizedOperator: a JaxParametrizedOperator instance
         """
-        return cls(*matrices, param_coeffs)
+        return cls(*matrices, param_coeffs, reorder_fn)
 
 
 @register_pytree_node_class
diff --git a/pennylane/pulse/rydberg_hamiltonian.py b/pennylane/pulse/rydberg_hamiltonian.py
index b677bb7e1..b9e22226a 100644
--- a/pennylane/pulse/rydberg_hamiltonian.py
+++ b/pennylane/pulse/rydberg_hamiltonian.py
@@ -21,6 +21,9 @@ import numpy as np
 
 import pennylane as qml
 from pennylane.wires import Wires
+from pennylane.operation import Operator
+from pennylane.ops.qubit.hamiltonian import Hamiltonian
+
 
 from .parametrized_hamiltonian import ParametrizedHamiltonian
 
@@ -119,12 +122,11 @@ def rydberg_interaction(
     )
 
 
-def rydberg_drive(amplitude, detuning, phase, wires):
+def rydberg_drive(amplitude, phase, detuning, wires):
     r"""Returns a :class:`ParametrizedHamiltonian` representing the action of a driving laser
     field with the given rabi frequency, detuning and phase acting on the given wires
 
     .. math::
-
         \frac{1}{2} \Omega(t) \sum_{i \in \text{wires}} (\cos(\phi)\sigma_i^x - \sin(\phi)\sigma_i^y) -
         \frac{1}{2} \delta(t) \sum_{i \in \text{wires}} \sigma_i^z
 
@@ -138,9 +140,9 @@ def rydberg_drive(amplitude, detuning, phase, wires):
     Args:
         amplitude (Union[float, Callable]): float or callable returning the amplitude (in MHz) of a
             laser field
+        phase (Union[float, Callable]): float or callable returning the phase (in radians) of the laser field
         detuning (Union[float, Callable]): float or callable returning the detuning (in MHz) of a
             laser field
-        phase (float): float containing the phase (in radians) of the laser field
         wires (Union[int, List[int]]): integer or list containing wire values for the Rydberg atoms that
             the laser field acts on
 
@@ -163,21 +165,25 @@ def rydberg_drive(amplitude, detuning, phase, wires):
     .. code-block:: python
 
         atom_coordinates = [[0, 0], [0, 4], [4, 0], [4, 4]]
+        wires = [1, 2, 3, 4]
         H_i = qml.pulse.rydberg_interaction(atom_coordinates, wires)
 
         amplitude = lambda p, t: p * jnp.sin(jnp.pi * t)
-        detuning = 3 * jnp.pi / 4
         phase = jnp.pi / 2
-        wires = [0, 1, 2, 3]
-        H_d = qml.pulse.rydberg_drive(amplitude, detuning, phase, wires)
+        detuning = 3 * jnp.pi / 4
+        H_d = qml.pulse.rydberg_drive(amplitude, phase, detuning, wires)
 
     >>> H_i
     ParametrizedHamiltonian: terms=6
     >>> H_d
-    ParametrizedHamiltonian: terms=2
+    ParametrizedHamiltonian: terms=3
+
+    The first two terms of the drive Hamiltonian ``H_d`` correspond to the first sum (the sine and cosine terms),
+    describing drive between the ground and excited states. The third term corresponds to the shift term
+    due to detuning from resonance. This drive term corresponds to a global drive that acts on all 4 wires of
+    the device.
 
-    The two terms of the drive field correspond to the first and second sum, corresponding to the drive and the shift term.
-    This drive term corresponds to a global drive that acts on all wires of the device.
+    The full Hamiltonian can be evaluated:
 
     .. code-block:: python
 
@@ -200,9 +206,9 @@ def rydberg_drive(amplitude, detuning, phase, wires):
     .. code-block:: python
 
         amplitude_local = lambda p, t: p[0] * jnp.sin(2 * jnp.pi * t) + p[1]
-        detuning_local = lambda p, t: p * jnp.exp(-0.25 * t)
-        phase_local = jnp.pi / 4
-        H_local = qml.pulse.rydberg_drive(amplitude_local, detuning_local, phase_local, [0, 1])
+        phase_local = lambda p, t: p * jnp.exp(-0.25 * t)
+        detuning_local = jnp.pi / 4
+        H_local = qml.pulse.rydberg_drive(amplitude_local, phase_local, detuning_local, [0, 1])
 
         H = H_i + H_d + H_local
 
@@ -215,7 +221,7 @@ def rydberg_drive(amplitude, detuning, phase, wires):
     >>> params = [2.4, [1.3, -2.0]]
     >>> circuit_local(params)
     Array(0.45782223, dtype=float64)
-    >>> jax.grad(circuit_local_drives)(params)
+    >>> jax.grad(circuit_local)(params)
     [Array(-0.33522988, dtype=float64),
      [Array(0.40320718, dtype=float64, weak_type=True),
       Array(-0.12003976, dtype=float64, weak_type=True)]]
@@ -224,16 +230,21 @@ def rydberg_drive(amplitude, detuning, phase, wires):
         wires = [wires]
 
     # We compute the `coeffs` and `observables` of the laser field
-    coeffs = [amplitude, detuning]
-    amplitude_observable = sum(
-        qml.math.cos(phase) * qml.PauliX(wire) - qml.math.sin(phase) * qml.PauliY(wire)
-        for wire in wires
-    )
-    detuning_observable = sum(qml.PauliZ(wire) for wire in wires)
-    observables = [amplitude_observable, detuning_observable]
+    coeffs = [
+        amplitude_and_phase(qml.math.cos, amplitude, phase),
+        amplitude_and_phase(qml.math.sin, amplitude, phase),
+        detuning,
+    ]
+
+    drive_x_term = sum(qml.PauliX(wire) for wire in wires)
+    drive_y_term = sum(-qml.PauliY(wire) for wire in wires)
+    detuning_term = sum(qml.PauliZ(wire) for wire in wires)
+
+    observables = [drive_x_term, drive_y_term, detuning_term]
 
     # We convert the pulse data into a list of ``RydbergPulse`` objects
-    pulses = [RydbergPulse(amplitude, detuning, phase, wires)]
+    pulses = [RydbergPulse(amplitude, phase, detuning, wires)]
+
     return RydbergHamiltonian(coeffs, observables, pulses=pulses)
 
 
@@ -292,26 +303,75 @@ class RydbergHamiltonian(ParametrizedHamiltonian):
         self.interaction_coeff = interaction_coeff
         super().__init__(coeffs, observables)
 
+    def __call__(self, params, t):
+        params = _rydberg_reorder_parameters(params, self.coeffs_parametrized)
+        return super().__call__(params, t)
+
     def __add__(self, other):
-        if not isinstance(other, RydbergHamiltonian):
-            return super().__add__(other)
-
-        # Update coeffs, obs and hardware attributes
-        if self.register is not None:
-            if other.register is not None:
-                raise ValueError("We cannot add two Hamiltonians with an interaction term!")
-            if not self.wires.contains_wires(other.wires):
+        if isinstance(other, RydbergHamiltonian):
+            # Update coeffs, obs and hardware attributes
+            if self.register is not None:
+                if other.register is not None:
+                    raise ValueError("We cannot add two Hamiltonians with an interaction term!")
+                if not self.wires.contains_wires(other.wires):
+                    warnings.warn(
+                        "The wires of the laser fields are not present in the Rydberg ensemble."
+                    )
+            elif other.register is not None and not other.wires.contains_wires(self.wires):
                 warnings.warn(
                     "The wires of the laser fields are not present in the Rydberg ensemble."
                 )
-        elif other.register is not None and not other.wires.contains_wires(self.wires):
-            warnings.warn("The wires of the laser fields are not present in the Rydberg ensemble.")
 
-        new_register = self.register or other.register
-        new_pulses = self.pulses + other.pulses
-        new_ops = self.ops + other.ops
-        new_coeffs = self.coeffs + other.coeffs
-        return RydbergHamiltonian(new_coeffs, new_ops, register=new_register, pulses=new_pulses)
+            new_register = self.register or other.register
+            new_pulses = self.pulses + other.pulses
+            new_ops = self.ops + other.ops
+            new_coeffs = self.coeffs + other.coeffs
+            return RydbergHamiltonian(new_coeffs, new_ops, register=new_register, pulses=new_pulses)
+
+        ops = self.ops.copy()
+        coeffs = self.coeffs.copy()
+        register = self.register
+        pulses = self.pulses
+
+        if isinstance(other, (Hamiltonian, ParametrizedHamiltonian)):
+            new_coeffs = coeffs + list(other.coeffs.copy())
+            new_ops = ops + other.ops.copy()
+            return RydbergHamiltonian(new_coeffs, new_ops, register=register, pulses=pulses)
+
+        if isinstance(other, qml.ops.SProd):  # pylint: disable=no-member
+            new_coeffs = coeffs + [other.scalar]
+            new_ops = ops + [other.base]
+            return RydbergHamiltonian(new_coeffs, new_ops, register=register, pulses=pulses)
+
+        if isinstance(other, Operator):
+            new_coeffs = coeffs + [1]
+            new_ops = ops + [other]
+            return RydbergHamiltonian(new_coeffs, new_ops, register=register, pulses=pulses)
+
+        return NotImplemented
+
+    def __radd__(self, other):
+        """Deals with the special case where a RydbergHamiltonian is added to a
+        ParametrizedHamiltonian. Ensures that this returns a RydbergHamiltonian where
+        the order of the parametrized coefficients and operators matches the order of
+        the hamiltonians, i.e. that
+
+        ParametrizedHamiltonian + RydbergHamiltonian
+
+        returns a RydbergHamiltonian where the call expects params = [params_PH] + [params_RH]
+        """
+        if isinstance(other, ParametrizedHamiltonian):
+            ops = self.ops.copy()
+            coeffs = self.coeffs.copy()
+
+            new_coeffs = other.coeffs.copy() + coeffs
+            new_ops = other.ops.copy() + ops
+
+            return RydbergHamiltonian(
+                new_coeffs, new_ops, register=self.register, pulses=self.pulses
+            )
+
+        return self.__add__(other)
 
 
 @dataclass
@@ -322,16 +382,16 @@ class RydbergPulse:
     Args:
         amplitude (Union[float, Callable]): float or callable returning the amplitude (in MHz) of a laser
             field
+        phase (Union[float, Callable]): float containing the phase (in radians) of the laser field
         detuning (Union[float, Callable]): float or callable returning the detuning (in MHz) of a
             laser field
-        phase (float): float containing the phase (in radians) of the laser field
         wires (Union[int, List[int]]): integer or list containing wire values that the laser field
             acts on
     """
 
     amplitude: Union[float, Callable]
+    phase: Union[float, Callable]
     detuning: Union[float, Callable]
-    phase: float
     wires: List[Wires]
 
     def __post_init__(self):
@@ -344,3 +404,77 @@ class RydbergPulse:
             and self.detuning == other.detuning
             and self.wires == other.wires
         )
+
+
+def amplitude_and_phase(trig_fn, amp, phase):
+    """Wrapper function for combining amplitude and phase into a single callable
+    (or constant if neither amplitude nor phase are callable)."""
+    if not callable(amp) and not callable(phase):
+        return amp * trig_fn(phase)
+    return AmplitudeAndPhase(trig_fn, amp, phase)
+
+
+# pylint:disable = too-few-public-methods
+class AmplitudeAndPhase:
+    """Class storing combined amplitude and phase callable if either or both
+    of amplitude or phase are callable."""
+
+    def __init__(self, trig_fn, amp, phase):
+        self.amp_is_callable = callable(amp)
+        self.phase_is_callable = callable(phase)
+
+        def callable_amp_and_phase(params, t):
+            return amp(params[0], t) * trig_fn(phase(params[1], t))
+
+        def callable_amp(params, t):
+            return amp(params, t) * trig_fn(phase)
+
+        def callable_phase(params, t):
+            return amp * trig_fn(phase(params, t))
+
+        if self.amp_is_callable and self.phase_is_callable:
+            self.func = callable_amp_and_phase
+
+        elif self.amp_is_callable:
+            self.func = callable_amp
+
+        elif self.phase_is_callable:
+            self.func = callable_phase
+
+    def __call__(self, params, t):
+        return self.func(params, t)
+
+
+def _rydberg_reorder_parameters(params, coeffs_parametrized):
+    """Takes `params`, and reorganizes it based on whether the Hamiltonian has
+    callable phase and/or callable amplitude.
+
+    Consolidates phase and amplitude parameters in the case that both are callable,
+    and duplicates phase and/or amplitude parameters if either are callables, since
+    they will be passed to two operators in the Hamiltonian"""
+
+    reordered_params = []
+
+    coeff_idx = 0
+    params_idx = 0
+
+    for i, coeff in enumerate(coeffs_parametrized):
+        if i == coeff_idx:
+            if isinstance(coeff, AmplitudeAndPhase):
+                if coeff.phase_is_callable and coeff.amp_is_callable:
+                    # add the joined parameters twice, and skip an index
+                    reordered_params.append([params[params_idx], params[params_idx + 1]])
+                    reordered_params.append([params[params_idx], params[params_idx + 1]])
+                    coeff_idx += 2
+                    params_idx += 2
+                elif coeff.phase_is_callable or coeff.amp_is_callable:
+                    reordered_params.append(params[params_idx])
+                    reordered_params.append(params[params_idx])
+                    coeff_idx += 2
+                    params_idx += 1
+            else:
+                reordered_params.append(params[params_idx])
+                coeff_idx += 1
+                params_idx += 1
+
+    return reordered_params


Test patch:
diff --git a/tests/pulse/test_parametrized_hamiltonian_pytree.py b/tests/pulse/test_parametrized_hamiltonian_pytree.py
index 595594888..46328dcf6 100644
--- a/tests/pulse/test_parametrized_hamiltonian_pytree.py
+++ b/tests/pulse/test_parametrized_hamiltonian_pytree.py
@@ -24,8 +24,16 @@ try:
         LazyDotPytree,
         ParametrizedHamiltonianPytree,
     )
+    from pennylane.pulse.rydberg_hamiltonian import (
+        _rydberg_reorder_parameters,
+        amplitude_and_phase,
+        rydberg_drive,
+    )
 except ImportError:
-    pytestmark = pytest.mark.skip
+    pass
+
+# if this fails, test file will be skipped
+jnp = pytest.importorskip("jax.numpy")
 
 
 def f1(p, t):
@@ -33,19 +41,31 @@ def f1(p, t):
 
 
 def f2(p, t):
-    return p * np.cos(t**2)
+    return np.cos(p * t**2)
+
 
+PH = qml.dot([1, 2, f1, f2], [qml.PauliX(0), qml.PauliY(1), qml.PauliZ(2), qml.Hadamard(3)])
 
-params = [1.2, 2.3]
+RH = rydberg_drive(amplitude=f1, phase=f2, detuning=1, wires=[0, 1, 2])
 
-H = qml.dot([1, 2, f1, f2], [qml.PauliX(0), qml.PauliY(1), qml.PauliZ(2), qml.Hadamard(3)])
+# Hamiltonians and the parameters for the individual coefficients
+HAMILTONIANS_WITH_COEFF_PARAMETERS = [
+    (PH, None, [f1, f2], [1.2, 2.3]),
+    (
+        RH,
+        _rydberg_reorder_parameters,
+        [amplitude_and_phase(jnp.cos, f1, f2), amplitude_and_phase(jnp.sin, f1, f2)],
+        [[1.2, 2.3], [1.2, 2.3]],
+    ),
+]
 
 
 @pytest.mark.jax
 class TestParametrizedHamiltonianPytree:
     """Unit tests for the ParametrizedHamiltonianPytree class."""
 
-    def test_attributes(self):
+    @pytest.mark.parametrize("H, fn, coeffs, params", HAMILTONIANS_WITH_COEFF_PARAMETERS)
+    def test_attributes(self, H, fn, coeffs, params):
         """Test that the attributes of the ParametrizedHamiltonianPytree class are initialized
         correctly."""
         from jax.experimental import sparse
@@ -56,21 +76,46 @@ class TestParametrizedHamiltonianPytree:
 
         assert isinstance(H_pytree.mat_fixed, sparse.CSR)
         assert isinstance(H_pytree.mats_parametrized, tuple)
-        assert H_pytree.coeffs_parametrized == [f1, f2]
+        assert qml.math.allclose(
+            [c(p, 2) for c, p in zip(H_pytree.coeffs_parametrized, params)],
+            [c(p, 2) for c, p in zip(coeffs, params)],
+            atol=1e-7,
+        )
+        assert H_pytree.reorder_fn == fn
 
-    def test_call_method(self):
+    def test_call_method_parametrized_hamiltonian(self):
         """Test that the call method works correctly."""
         H_pytree = ParametrizedHamiltonianPytree.from_hamiltonian(
-            H, dense=False, wire_order=[2, 3, 1, 0]
+            PH, dense=False, wire_order=[2, 3, 1, 0]
         )
-
+        params = [1.2, 2.3]
         time = 4
         res = H_pytree(params, t=time)
 
         assert isinstance(res, LazyDotPytree)
         assert qml.math.allclose(res.coeffs, (1, f1(params[0], time), f2(params[1], time)))
 
-    def test_flatten_method(self):
+    def test_call_method_rydberg_hamiltonian(self):
+        """Test that the call method works correctly."""
+        H_pytree = ParametrizedHamiltonianPytree.from_hamiltonian(
+            RH, dense=False, wire_order=[2, 3, 1, 0]
+        )
+        params = [1.2, 2.3]
+        time = 4
+        res = H_pytree(params, t=time)
+
+        assert isinstance(res, LazyDotPytree)
+        assert qml.math.allclose(
+            res.coeffs,
+            (
+                1,
+                amplitude_and_phase(jnp.cos, f1, f2)(params, time),
+                amplitude_and_phase(jnp.sin, f1, f2)(params, time),
+            ),
+        )
+
+    @pytest.mark.parametrize("H, fn", [(PH, None), (RH, _rydberg_reorder_parameters)])
+    def test_flatten_method(self, H, fn):
         """Test the tree_flatten method."""
         H_pytree = ParametrizedHamiltonianPytree.from_hamiltonian(
             H, dense=False, wire_order=[2, 3, 1, 0]
@@ -82,9 +127,11 @@ class TestParametrizedHamiltonianPytree:
         assert flat_tree == (
             (H_pytree.mat_fixed, H_pytree.mats_parametrized),
             H_pytree.coeffs_parametrized,
+            fn,
         )
 
-    def test_unflatten_method(self):
+    @pytest.mark.parametrize("H, fn", [(PH, None), (RH, _rydberg_reorder_parameters)])
+    def test_unflatten_method(self, H, fn):
         """Test the tree_unflatten method."""
         H_pytree = ParametrizedHamiltonianPytree.from_hamiltonian(
             H, dense=False, wire_order=[2, 3, 1, 0]
@@ -92,11 +139,12 @@ class TestParametrizedHamiltonianPytree:
 
         flat_tree = H_pytree.tree_flatten()
 
-        new_H_pytree = H_pytree.tree_unflatten(flat_tree[1], flat_tree[0])
+        new_H_pytree = H_pytree.tree_unflatten(flat_tree[1], flat_tree[0], flat_tree[2])
 
         assert new_H_pytree.mat_fixed == H_pytree.mat_fixed
         assert new_H_pytree.mats_parametrized == H_pytree.mats_parametrized
         assert new_H_pytree.coeffs_parametrized == H_pytree.coeffs_parametrized
+        assert new_H_pytree.reorder_fn == fn
 
 
 @pytest.mark.jax
@@ -126,7 +174,6 @@ class TestLazyDotPytree:
 
     def test_rmul(self):
         """Test the __rmul__ method"""
-        import jax.numpy as jnp
 
         coeffs = [1, 2, 3]
         ops = [qml.PauliX(0), qml.PauliY(0), qml.PauliZ(0)]
diff --git a/tests/pulse/test_rydberg_hamiltonian.py b/tests/pulse/test_rydberg_hamiltonian.py
index 561c7de4b..f687c29aa 100644
--- a/tests/pulse/test_rydberg_hamiltonian.py
+++ b/tests/pulse/test_rydberg_hamiltonian.py
@@ -22,13 +22,29 @@ import pytest
 
 import pennylane as qml
 from pennylane.pulse import RydbergHamiltonian, rydberg_drive, rydberg_interaction
-from pennylane.pulse.rydberg_hamiltonian import RydbergPulse
+from pennylane.pulse.rydberg_hamiltonian import (
+    RydbergPulse,
+    AmplitudeAndPhase,
+    amplitude_and_phase,
+    _rydberg_reorder_parameters,
+)
 from pennylane.wires import Wires
 
 atom_coordinates = [[0, 0], [0, 5], [5, 0], [10, 5], [5, 10], [10, 10]]
 wires = [1, 6, 0, 2, 4, 3]
 
 
+def f1(p, t):
+    return p * np.sin(t) * (t - 1)
+
+
+def f2(p, t):
+    return p * np.cos(t**2)
+
+
+param = [1.2, 2.3]
+
+
 class TestRydbergHamiltonian:
     """Unit tests for the properties of the RydbergHamiltonian class."""
 
@@ -42,7 +58,7 @@ class TestRydbergHamiltonian:
         assert rm.wires == Wires([])
         assert rm.interaction_coeff == 862690
 
-    def test_add(self):
+    def test_add_rydberg_hamiltonian(self):
         """Test that the __add__ dunder method works correctly."""
         rm1 = RydbergHamiltonian(
             coeffs=[1, 2],
@@ -85,6 +101,7 @@ class TestRydbergHamiltonian:
         res1 = rh + ph
         res2 = ph + rh
 
+        assert isinstance(res1, RydbergHamiltonian)
         assert res1.coeffs_fixed == coeffs
         assert res1.coeffs_parametrized == []
         assert all(qml.equal(op1, op2) for op1, op2 in zip(res1.ops_fixed, ops))
@@ -95,6 +112,7 @@ class TestRydbergHamiltonian:
         ops.reverse()
         h_wires.reverse()
 
+        assert isinstance(res2, RydbergHamiltonian)
         assert res2.coeffs_fixed == coeffs
         assert res2.coeffs_parametrized == []
         assert all(qml.equal(op1, op2) for op1, op2 in zip(res2.ops_fixed, ops))
@@ -134,16 +152,59 @@ class TestRydbergHamiltonian:
         ):
             _ = Ht + Hd
 
+    def test_hamiltonian_callable_after_addition_right(self):
+        """Tests that if a ParametrizedHamiltonian is added onto a
+        RydbergHamiltonian with callable coefficients from the right, the
+        resulting object is a RydbergHamiltonian that can be called without
+        raising an error"""
 
-def f1(p, t):
-    return p * np.sin(t) * (t - 1)
+        def amp(p, t):
+            return p[0] * np.exp(-((t - p[1]) ** 2) / (2 * p[2] ** 2))
 
+        def phase(p, t):
+            return p * t
 
-def f2(p, t):
-    return p * np.cos(t**2)
+        def detuning(p, t):
+            return np.cos(t) * p
 
+        H_global = qml.pulse.rydberg_drive(amp, phase, detuning, wires=[0])
+        H_global += qml.pulse.rydberg_drive(amp, phase, detuning, wires=[0])
+        H_global += np.polyval * qml.PauliX(0)
 
-param = [1.2, 2.3]
+        # start with RydbergHamiltonians, then add ParametrizedHamiltonian
+        params = [np.array([1.2, 2.3, 3.4]), 4.5, 5.6]
+        params += [np.array([1.2, 2.3, 3.4]), 4.5, 5.6]
+        params += [np.ones(2)]
+
+        assert isinstance(H_global, RydbergHamiltonian)
+        H_global(params, 2)  # no error raised
+
+    def test_hamiltonian_callable_after_addition_left(self):
+        """Tests that if a ParametrizedHamiltonian is added onto a
+        RydbergHamiltonian with callable coefficients from the left, the
+        resulting object is a RydbergHamiltonian that can be called without
+        raising an error"""
+
+        def amp(p, t):
+            return p[0] * np.exp(-((t - p[1]) ** 2) / (2 * p[2] ** 2))
+
+        def phase(p, t):
+            return p * t
+
+        def detuning(p, t):
+            return np.cos(t) * p
+
+        # start with ParametrizedHamiltonian, add on RydbergHamiltonians
+        H_global = np.polyval * qml.PauliX(0)
+        H_global += qml.pulse.rydberg_drive(amp, phase, detuning, wires=[0])
+        H_global += qml.pulse.rydberg_drive(amp, phase, detuning, wires=[0])
+
+        params = [np.ones(2)]
+        params += [np.array([1.2, 2.3, 3.4]), 4.5, 5.6]
+        params += [np.array([1.2, 2.3, 3.4]), 4.5, 5.6]
+
+        assert isinstance(H_global, RydbergHamiltonian)
+        H_global(params, 2)  # no error raised
 
 
 class TestInteractionWithOperators:
@@ -153,29 +214,32 @@ class TestInteractionWithOperators:
     ops_with_coeffs = (
         (qml.Hamiltonian([2], [qml.PauliZ(0)]), 2),
         (qml.Hamiltonian([1.7], [qml.PauliZ(0)]), 1.7),
+        (3 * qml.PauliZ(0), 3),
         (qml.ops.SProd(3, qml.PauliZ(0)), 3),
     )
     ops = (
         qml.PauliX(2),
         qml.PauliX(2) @ qml.PauliX(3),
         qml.CNOT([0, 1]),
-    )  # ToDo: maybe add more operators to test here?
+    )
 
     @pytest.mark.parametrize("H, coeff", ops_with_coeffs)
     def test_add_special_operators(self, H, coeff):
         """Test that a Hamiltonian and SProd can be added to a RydbergHamiltonian, and
         will be incorporated in the H_fixed term, with their coefficients included in H_coeffs_fixed.
         """
-        R = rydberg_drive(amplitude=f1, detuning=f2, phase=0, wires=[0, 1])
+        R = rydberg_drive(amplitude=f1, phase=0, detuning=f2, wires=[0, 1])
         params = [1, 2]
         # Adding on the right
         new_pH = R + H
+        assert isinstance(new_pH, RydbergHamiltonian)
         assert R.H_fixed() == 0
         assert qml.equal(new_pH.H_fixed(), qml.s_prod(coeff, qml.PauliZ(0)))
         assert new_pH.coeffs_fixed[0] == coeff
         assert qml.math.allequal(new_pH(params, t=0.5).matrix(), qml.matrix(R(params, t=0.5) + H))
         # Adding on the left
         new_pH = H + R
+        assert isinstance(new_pH, RydbergHamiltonian)
         assert R.H_fixed() == 0
         assert qml.equal(new_pH.H_fixed(), qml.s_prod(coeff, qml.PauliZ(0)))
         assert new_pH.coeffs_fixed[0] == coeff
@@ -185,17 +249,27 @@ class TestInteractionWithOperators:
     def test_add_other_operators(self, op):
         """Test that a Hamiltonian, SProd, Tensor or Operator can be added to a
         ParametrizedHamiltonian, and will be incorporated in the H_fixed term"""
-        R = rydberg_drive(amplitude=f1, detuning=f2, phase=0, wires=[0, 1])
+        R = rydberg_drive(amplitude=f1, phase=0, detuning=f2, wires=[0, 1])
+        params = [1, 2]
 
         # Adding on the right
         new_pH = R + op
+        assert isinstance(new_pH, RydbergHamiltonian)
         assert R.H_fixed() == 0
         assert qml.equal(new_pH.H_fixed(), qml.s_prod(1, op))
+        new_pH(params, 2)  # confirm calling does not raise error
 
         # Adding on the left
         new_pH = op + R
+        assert isinstance(new_pH, RydbergHamiltonian)
         assert R.H_fixed() == 0
         assert qml.equal(new_pH.H_fixed(), qml.s_prod(1, op))
+        new_pH(params, 2)  # confirm calling does not raise error
+
+    def test_unknown_type_raises_error(self):
+        R = rydberg_drive(amplitude=f1, phase=0, detuning=f2, wires=[0, 1])
+        with pytest.raises(TypeError, match="unsupported operand type"):
+            R += 3
 
 
 class TestRydbergInteraction:
@@ -258,13 +332,14 @@ class TestRydbergDrive:
     def test_attributes_and_number_of_terms(self):
         """Test that the attributes and the number of terms of the ``ParametrizedHamiltonian`` returned by
         ``rydberg_drive`` are correct."""
-        Hd = rydberg_drive(amplitude=1, detuning=2, phase=3, wires=[1, 2])
+
+        Hd = rydberg_drive(amplitude=1, phase=2, detuning=3, wires=[1, 2])
 
         assert isinstance(Hd, RydbergHamiltonian)
         assert Hd.interaction_coeff == 862690
         assert Hd.wires == Wires([1, 2])
         assert Hd.register is None
-        assert len(Hd.ops) == 2  # amplitude and detuning terms of the Hamiltonian
+        assert len(Hd.ops) == 3  # 2 amplitude/phase terms and one detuning term of the Hamiltonian
         assert Hd.pulses == [RydbergPulse(1, 2, 3, [1, 2])]
 
     def test_multiple_local_drives(self):
@@ -276,45 +351,259 @@ class TestRydbergDrive:
         def fb(p, t):
             return np.cos(p * t)
 
-        H1 = rydberg_drive(amplitude=fa, detuning=3, phase=1, wires=[0, 3])
-        H2 = rydberg_drive(amplitude=1, detuning=fb, phase=3, wires=[1, 2])
+        H1 = rydberg_drive(amplitude=fa, phase=1, detuning=3, wires=[0, 3])
+        H2 = rydberg_drive(amplitude=1, phase=3, detuning=fb, wires=[1, 2])
         Hd = H1 + H2
 
         ops_expected = [
             qml.Hamiltonian([1, 1], [qml.PauliZ(0), qml.PauliZ(3)]),
-            qml.Hamiltonian(
-                [np.cos(3), -np.sin(3), np.cos(3), -np.sin(3)],
-                [qml.PauliX(1), qml.PauliY(1), qml.PauliX(2), qml.PauliY(2)],
-            ),
-            qml.Hamiltonian(
-                [np.cos(1), -np.sin(1), np.cos(1), -np.sin(1)],
-                [qml.PauliX(0), qml.PauliY(0), qml.PauliX(3), qml.PauliY(3)],
-            ),
+            qml.Hamiltonian([1, 1], [qml.PauliX(1), qml.PauliX(2)]),
+            qml.sum(qml.s_prod(-1, qml.PauliY(1)), qml.s_prod(-1, qml.PauliY(2))),
+            qml.Hamiltonian([1, 1], [qml.PauliX(0), qml.PauliX(3)]),
+            qml.sum(qml.s_prod(-1, qml.PauliY(0)), qml.s_prod(-1, qml.PauliY(3))),
             qml.Hamiltonian([1, 1], [qml.PauliZ(1), qml.PauliZ(2)]),
         ]
-        coeffs_expected = [3, 1, fa, fb]
-        H_expected = qml.dot(coeffs_expected, ops_expected)
+        coeffs_expected = [
+            3,
+            np.cos(3),
+            np.sin(3),
+            AmplitudeAndPhase(np.cos, fa, 1),
+            AmplitudeAndPhase(np.sin, fa, 1),
+            fb,
+        ]
+        H_expected = RydbergHamiltonian(coeffs_expected, ops_expected)
 
+        # structure of Hamiltonian is as expected
         assert isinstance(Hd, RydbergHamiltonian)
         assert Hd.interaction_coeff == 862690
         assert Hd.wires == Wires([0, 3, 1, 2])
         assert Hd.register is None
-        assert len(Hd.ops) == 4  # amplitude and detuning terms of the Hamiltonian
+        assert len(Hd.ops) == 6  # 2 terms for amplitude/phase and one detuning for each drive
+
+        # coefficients are correct
         # Callable coefficients are shifted to the end of the list.
-        assert Hd.coeffs[0:2] == [3, 1]
-        assert Hd.coeffs[2] is fa
-        assert Hd.coeffs[3] is fb
+        assert Hd.coeffs[0:3] == [3, np.cos(3), np.sin(3)]
+        assert isinstance(Hd.coeffs[3], AmplitudeAndPhase)
+        assert isinstance(Hd.coeffs[4], AmplitudeAndPhase)
+        assert Hd.coeffs[5] is fb
+
+        # pulses were added correctly
         assert len(Hd.pulses) == 2
         assert Hd.pulses == H1.pulses + H2.pulses
+
+        # Hamiltonian is as expected
         assert qml.equal(Hd([0.5, -0.5], t=5), H_expected([0.5, -0.5], t=5))
 
 
+def callable_amp(p, t):
+    return np.polyval(p, t)
+
+
+def callable_phase(p, t):
+    return p[0] * np.sin(p[1] * t)
+
+
+def sine_func(p, t):
+    return np.sin(p * t)
+
+
+def cosine_fun(p, t):
+    return np.cos(p * t)
+
+
+class TestAmplitudeAndPhase:
+    """Test the AmplitudeAndPhase class that provides callable
+    phase/amplitude combinations"""
+
+    def test_amplitude_and_phase_no_callables(self):
+        """Test that when calling amplitude_and_phase, if neither are callable,
+        a float is returned instead of an AmplitudeAndPhase object"""
+        f = amplitude_and_phase(np.sin, 3, 4)
+        expected_result = 3 * np.sin(4)
+
+        assert isinstance(f, float)
+        assert f == expected_result
+
+    def test_amplitude_and_phase_callable_phase(self):
+        """Test that when calling amplitude_and_phase, if only phase is callable,
+        an AmplitudeAndPhase object with callable phase and fixed amplitude is
+        correctly created"""
+        f = amplitude_and_phase(np.sin, 2.7, callable_phase)
+
+        # attributes are correct
+        assert isinstance(f, AmplitudeAndPhase)
+        assert f.amp_is_callable is False
+        assert f.phase_is_callable is True
+        assert f.func.__name__ == "callable_phase"
+
+        # calling yields expected result
+        expected_result = 2.7 * np.sin(callable_phase([1.3, 2.5], 2))
+        assert f([1.3, 2.5], 2) == expected_result
+
+    def test_amplitude_and_phase_callable_amplitude(self):
+        """Test that when calling amplitude_and_phase, if only amplitude is callable,
+        an AmplitudeAndPhase object with callable amplitude and fixed phase is
+        correctly created"""
+        f = amplitude_and_phase(np.sin, callable_amp, 0.7)
+
+        # attributes are correct
+        assert isinstance(f, AmplitudeAndPhase)
+        assert f.amp_is_callable is True
+        assert f.phase_is_callable is False
+        assert f.func.__name__ == "callable_amp"
+
+        # calling yields expected result
+        expected_result = callable_amp([1.7], 2) * np.sin(0.7)
+        assert f([1.7], 2) == expected_result
+
+    def test_amplitude_and_phase_both_callable(self):
+        """Test that when calling amplitude_and_phase, if both are callable,
+        an AmplitudeAndPhase object with callable amplitude and phase is
+        correctly created"""
+        f = amplitude_and_phase(np.sin, callable_amp, callable_phase)
+
+        # attributes are correct
+        assert isinstance(f, AmplitudeAndPhase)
+        assert f.amp_is_callable is True
+        assert f.phase_is_callable is True
+        assert f.func.__name__ == "callable_amp_and_phase"
+
+        # calling yields expected result
+        expected_result = callable_amp([1.7], 2) * np.sin(callable_phase([1.3, 2.5], 2))
+        assert f([[1.7], [1.3, 2.5]], 2) == expected_result
+
+    def test_callable_phase_and_amplitude_hamiltonian(self):
+        """Test that using callable amplitude and phase in rydberg_drive
+        creates AmplitudeAndPhase callables, and the resulting Hamiltonian
+        can be called successfully"""
+
+        detuning = 2
+
+        Hd = rydberg_drive(sine_func, cosine_fun, detuning, wires=[0, 1])
+
+        assert len(Hd.coeffs) == 3
+        assert isinstance(Hd.coeffs[1], AmplitudeAndPhase)
+        assert isinstance(Hd.coeffs[2], AmplitudeAndPhase)
+        t = 1.7
+
+        evaluated_H = Hd([3.4, 5.6], t)
+
+        expected_H_fixed = qml.s_prod(
+            detuning, qml.Hamiltonian([1, 1], [qml.PauliZ(0), qml.PauliZ(1)])
+        )
+
+        c1 = np.sin(3.4 * t) * np.cos(np.cos(5.6 * t))
+        c2 = np.sin(3.4 * t) * np.sin(np.cos(5.6 * t))
+        expected_H_parametrized = qml.sum(
+            qml.s_prod(c1, qml.Hamiltonian([1, 1], [qml.PauliX(0), qml.PauliX(1)])),
+            qml.s_prod(c2, qml.sum(qml.s_prod(-1, qml.PauliY(0)), qml.s_prod(-1, qml.PauliY(1)))),
+        )
+
+        assert qml.equal(evaluated_H[0], expected_H_fixed)
+        assert qml.equal(evaluated_H[1], expected_H_parametrized)
+
+    def test_callable_phase_hamiltonian(self):
+        """Test that using callable phase in rydberg_drive creates AmplitudeAndPhase
+        callables, and the resulting Hamiltonian can be called"""
+
+        detuning = 2
+
+        Hd = rydberg_drive(7.2, sine_func, detuning, wires=[0, 1])
+
+        assert len(Hd.coeffs) == 3
+        assert isinstance(Hd.coeffs[1], AmplitudeAndPhase)
+        assert isinstance(Hd.coeffs[2], AmplitudeAndPhase)
+        t = 1.7
+
+        evaluated_H = Hd([5.6], t)
+
+        expected_H_fixed = qml.s_prod(
+            detuning, qml.Hamiltonian([1, 1], [qml.PauliZ(0), qml.PauliZ(1)])
+        )
+
+        c1 = 7.2 * np.cos(np.sin(5.6 * t))
+        c2 = 7.2 * np.sin(np.sin(5.6 * t))
+        expected_H_parametrized = qml.sum(
+            qml.s_prod(c1, qml.Hamiltonian([1, 1], [qml.PauliX(0), qml.PauliX(1)])),
+            qml.s_prod(c2, qml.sum(qml.s_prod(-1, qml.PauliY(0)), qml.s_prod(-1, qml.PauliY(1)))),
+        )
+
+        assert qml.equal(evaluated_H[0], expected_H_fixed)
+        assert qml.equal(evaluated_H[1], expected_H_parametrized)
+
+    def test_callable_amplitude_hamiltonian(self):
+        """Test that using callable amplitude in rydberg_drive creates AmplitudeAndPhase
+        callables, and the resulting Hamiltonian can be called"""
+
+        detuning = 2
+
+        Hd = rydberg_drive(sine_func, 4.3, detuning, wires=[0, 1])
+
+        assert len(Hd.coeffs) == 3
+        assert isinstance(Hd.coeffs[1], AmplitudeAndPhase)
+        assert isinstance(Hd.coeffs[2], AmplitudeAndPhase)
+        t = 1.7
+
+        evaluated_H = Hd([3.4], t)
+
+        expected_H_fixed = qml.s_prod(
+            detuning, qml.Hamiltonian([1, 1], [qml.PauliZ(0), qml.PauliZ(1)])
+        )
+
+        c1 = np.sin(3.4 * t) * np.cos(4.3)
+        c2 = np.sin(3.4 * t) * np.sin(4.3)
+        expected_H_parametrized = qml.sum(
+            qml.s_prod(c1, qml.Hamiltonian([1, 1], [qml.PauliX(0), qml.PauliX(1)])),
+            qml.s_prod(c2, qml.sum(qml.s_prod(-1, qml.PauliY(0)), qml.s_prod(-1, qml.PauliY(1)))),
+        )
+
+        assert qml.equal(evaluated_H[0], expected_H_fixed)
+        assert qml.equal(evaluated_H[1], expected_H_parametrized)
+
+    COEFFS_AND_PARAMS = [
+        (
+            [AmplitudeAndPhase(np.cos, f1, f2), AmplitudeAndPhase(np.sin, f1, f2), f2],
+            [1, 2, 3],
+            [[1, 2], [1, 2], 3],
+        ),
+        (
+            [AmplitudeAndPhase(np.cos, f1, 6.3), AmplitudeAndPhase(np.sin, f1, 6.3), f2],
+            [1, 3],
+            [1, 1, 3],
+        ),
+        (
+            [AmplitudeAndPhase(np.cos, 6.3, f1), AmplitudeAndPhase(np.sin, 6.3, f1), f2],
+            [1, 3],
+            [1, 1, 3],
+        ),
+        (
+            [f1, AmplitudeAndPhase(np.cos, f1, 5.7), AmplitudeAndPhase(np.sin, f1, 5.7), f2],
+            [1, 2, 3],
+            [1, 2, 2, 3],
+        ),
+        (
+            [f1, AmplitudeAndPhase(np.cos, f1, f1), AmplitudeAndPhase(np.sin, f1, f1), f2],
+            [2, [3, 5], 8, [13, 21]],
+            [2, [[3, 5], 8], [[3, 5], 8], [13, 21]],
+        ),
+        ([f1, f2, f1, f2], [1, 2, 3, 4], [1, 2, 3, 4]),
+    ]
+
+    @pytest.mark.parametrize("coeffs, params, expected_output", COEFFS_AND_PARAMS)
+    def test_rydberg_reorder_parameters_all(self, coeffs, params, expected_output):
+        """Tests that the function organizing the parameters to pass to the
+        RydbergHamiltonian works as expected when AmplitudeAndPhase callables
+        are included"""
+
+        assert _rydberg_reorder_parameters(params, coeffs) == expected_output
+
+
 class TestRydbergPulse:
     """Unit tests for the ``RydbergPulse`` class."""
 
     def test_init(self):
         """Test the initialization of the ``RydbergPulse`` class."""
-        p = RydbergPulse(amplitude=4, phase=8, detuning=9, wires=[0, 4, 7])
+        p = RydbergPulse(amplitude=4, detuning=9, phase=8, wires=[0, 4, 7])
         assert p.amplitude == 4
         assert p.phase == 8
         assert p.detuning == 9
@@ -347,23 +636,30 @@ class TestIntegration:
         def fb(p, t):
             return p[0] * jnp.sin(p[1] * t)
 
-        Ht = rydberg_drive(amplitude=fa, detuning=fb, phase=0, wires=1)
+        Ht = rydberg_drive(amplitude=fa, phase=0, detuning=fb, wires=1)
 
         dev = qml.device("default.qubit", wires=wires)
 
         ts = jnp.array([0.0, 3.0])
         H_obj = sum(qml.PauliZ(i) for i in range(2))
 
-        @jax.jit
         @qml.qnode(dev, interface="jax")
         def qnode(params):
             qml.evolve(Hd + Ht)(params, ts)
             return qml.expval(H_obj)
 
+        @jax.jit
+        @qml.qnode(dev, interface="jax")
+        def qnode_jit(params):
+            qml.evolve(Hd + Ht)(params, ts)
+            return qml.expval(H_obj)
+
         params = (jnp.ones(5), jnp.array([1.0, jnp.pi]))
         res = qnode(params)
+        res_jit = qnode_jit(params)
 
         assert isinstance(res, jax.Array)
+        assert res == res_jit
 
     @pytest.mark.jax
     def test_jitted_qnode_multidrive(self):
@@ -383,21 +679,71 @@ class TestIntegration:
         def fc(p, t):
             return p[0] * jnp.sin(t) + jnp.cos(p[1] * t)
 
-        H1 = rydberg_drive(amplitude=fa, detuning=fb, phase=0, wires=1)
-        H2 = rydberg_drive(amplitude=fc, detuning=jnp.pi / 4, phase=3 * jnp.pi, wires=4)
+        H1 = rydberg_drive(amplitude=fa, phase=0, detuning=fb, wires=1)
+        H2 = rydberg_drive(amplitude=fc, phase=3 * jnp.pi, detuning=jnp.pi / 4, wires=4)
 
         dev = qml.device("default.qubit", wires=wires)
 
         ts = jnp.array([0.0, 3.0])
         H_obj = sum(qml.PauliZ(i) for i in range(2))
 
-        @jax.jit
         @qml.qnode(dev, interface="jax")
         def qnode(params):
             qml.evolve(Hd + H1 + H2)(params, ts)
             return qml.expval(H_obj)
 
+        @jax.jit
+        @qml.qnode(dev, interface="jax")
+        def qnode_jit(params):
+            qml.evolve(Hd + H1 + H2)(params, ts)
+            return qml.expval(H_obj)
+
+        params = (jnp.ones(5), jnp.array([1.0, jnp.pi]), jnp.array([jnp.pi / 2, 0.5]))
+        res = qnode(params)
+        res_jit = qnode_jit(params)
+
+        assert isinstance(res, jax.Array)
+        assert res == res_jit
+
+    @pytest.mark.jax
+    def test_jitted_qnode_all_coeffs_callable(self):
+        """Test that a ``RydbergHamiltonian`` class can be executed within a
+        jitted qnode when all coeffs are callable."""
+        import jax
+        import jax.numpy as jnp
+
+        H_drift = rydberg_interaction(register=atom_coordinates, wires=wires)
+
+        def fa(p, t):
+            return jnp.polyval(p, t)
+
+        def fb(p, t):
+            return p[0] * jnp.sin(p[1] * t)
+
+        def fc(p, t):
+            return p[0] * jnp.sin(t) + jnp.cos(p[1] * t)
+
+        H_drive = rydberg_drive(amplitude=fa, phase=fb, detuning=fc, wires=1)
+
+        dev = qml.device("default.qubit", wires=wires)
+
+        ts = jnp.array([0.0, 3.0])
+        H_obj = sum(qml.PauliZ(i) for i in range(2))
+
+        @qml.qnode(dev, interface="jax")
+        def qnode(params):
+            qml.evolve(H_drift + H_drive)(params, ts)
+            return qml.expval(H_obj)
+
+        @jax.jit
+        @qml.qnode(dev, interface="jax")
+        def qnode_jit(params):
+            qml.evolve(H_drift + H_drive)(params, ts)
+            return qml.expval(H_obj)
+
         params = (jnp.ones(5), jnp.array([1.0, jnp.pi]), jnp.array([jnp.pi / 2, 0.5]))
         res = qnode(params)
+        res_jit = qnode_jit(params)
 
         assert isinstance(res, jax.Array)
+        assert res == res_jit