Add get samples function to InstructionToSignals for JAX-jit usage

qiskit_dynamics/pulse/__init__.py

@@ -81,5 +81,4 @@
 
    InstructionToSignals
 """
-
 from .pulse_to_signals import InstructionToSignals

qiskit_dynamics/pulse/pulse_to_signals.py

@@ -14,8 +14,11 @@
 Pulse schedule to Signals converter.
 """
 
-from typing import Dict, List, Optional
+from typing import Callable, Dict, List, Optional
+import functools
+
 import numpy as np
+import sympy as sym
 
 from qiskit.pulse import (
     Schedule,
@@ -30,8 +33,11 @@
     ControlChannel,
     AcquireChannel,
 )
+from qiskit.pulse.exceptions import PulseError
+from qiskit.pulse.library import SymbolicPulse
 from qiskit import QiskitError
 
+from qiskit_dynamics.array import Array
 from qiskit_dynamics.signals import DiscreteSignal
 
 
@@ -166,29 +172,33 @@ def get_signals(self, schedule: Schedule) -> List[DiscreteSignal]:
                 if isinstance(inst.pulse, Waveform):
                     inst_samples = inst.pulse.samples
                 else:
-                    inst_samples = inst.pulse.get_waveform().samples
+                    inst_samples = get_samples(inst.pulse)
 
                 # build sample array to append to signal
                 times = self._dt * (start_sample + np.arange(len(inst_samples)))
                 samples = inst_samples * np.exp(
-                    2.0j * np.pi * freq * times
-                    + 1.0j * phi
-                    + 2.0j * np.pi * phase_accumulations[chan]
+                    Array(
+                        2.0j * np.pi * freq * times
+                        + 1.0j * phi
+                        + 2.0j * np.pi * phase_accumulations[chan]
+                    )
                 )
                 signals[chan].add_samples(start_sample, samples)
 
             if isinstance(inst, ShiftPhase):
                 phases[chan] += inst.phase
 
             if isinstance(inst, ShiftFrequency):
-                frequency_shifts[chan] += inst.frequency
-                phase_accumulations[chan] -= inst.frequency * start_sample * self._dt
+                frequency_shifts[chan] = frequency_shifts[chan] + Array(inst.frequency)
+                phase_accumulations[chan] = (
+                    phase_accumulations[chan] - inst.frequency * start_sample * self._dt
+                )
 
             if isinstance(inst, SetPhase):
                 phases[chan] = inst.phase
 
             if isinstance(inst, SetFrequency):
-                phase_accumulations[chan] -= (
+                phase_accumulations[chan] = phase_accumulations[chan] - (
                     (inst.frequency - (frequency_shifts[chan] + signals[chan].carrier_freq))
                     * start_sample
                     * self._dt
@@ -301,3 +311,59 @@ def _get_channel(self, channel_name: str):
             raise QiskitError(
                 f"Invalid channel name {channel_name} given to {self.__class__.__name__}."
             ) from error
+
+
+def get_samples(pulse: SymbolicPulse) -> np.ndarray:
+    """Return samples filled according to the formula that the pulse
+    represents and the parameter values it contains.
+
+    Args:
+        pulse: SymbolicPulse class.
+    Returns:
+        Samples of the pulse.
+    Raises:
+        PulseError: When parameters are not assigned.
+        PulseError: When expression for pulse envelope is not assigned.
+        PulseError: When a free symbol value is not defined in the pulse instance parameters.
+    """
+    envelope = pulse.envelope
+    pulse_params = pulse.parameters
+    if pulse.is_parameterized():
+        raise PulseError("Unassigned parameter exists. All parameters must be assigned.")
+
+    if envelope is None:
+        raise PulseError("Pulse envelope expression is not assigned.")
+
+    args = []
+    for symbol in sorted(envelope.free_symbols, key=lambda s: s.name):
+        if symbol.name == "t":
+            times = Array(np.arange(0, pulse_params["duration"]) + 1 / 2)
+            args.insert(0, times.data)
+            continue
+        try:
+            args.append(pulse_params[symbol.name])
+        except KeyError as ex:
+            raise PulseError(
+                f"Pulse parameter '{symbol.name}' is not defined for this instance. "
+                "Please check your waveform expression is correct."
+            ) from ex
+    return _lru_cache_expr(envelope, Array.default_backend())(*args)
+
+
+@functools.lru_cache(maxsize=None)
+def _lru_cache_expr(expr: sym.Expr, backend) -> Callable:
+    """A helper function to get lambdified expression.
+
+    Args:
+        expr: Symbolic expression to evaluate.
+        backend: Array backend.
+    Returns:
+        lambdified expression.
+    """
+    params = []
+    for param in sorted(expr.free_symbols, key=lambda s: s.name):
+        if param.name == "t":
+            params.insert(0, param)
+            continue
+        params.append(param)
+    return sym.lambdify(params, expr, modules=backend)

qiskit_dynamics/solvers/solver_classes.py

@@ -54,7 +54,8 @@
 
 
 try:
-    from jax import jit
+    from jax import core, jit
+    import jax.numpy as jnp
 except ImportError:
     pass
 
@@ -522,6 +523,8 @@ def solve(
             Array.default_backend() == "jax"
             and (method == "jax_odeint" or _is_diffrax_method(method))
             and all(isinstance(x, Schedule) for x in signals_list)
+            # check if jit transformation is already performed.
+            and not (isinstance(jnp.array(0), core.Tracer))
         ):
             all_results = self._solve_schedule_list_jax(
                 t_span_list=t_span_list,

releasenotes/notes/get_samples_for_instruction_to_signals-e1de83e3e65a30de.yaml

@@ -0,0 +1,22 @@
+---
+features:
+  - |
+    The logic of :class:`.InstructionToSignals` has been updated. 
+    This change allows you to jit compile a function including 
+    the converter with input pulse schedule that contains :class:`.SymbolicPulse`.
+
+    .. code-block:: python
+
+      from functools import partial
+      import jax
+      from qiskit import pulse
+      from qiskit_dynamics.pulse import InstructionToSignals
+
+      @partial(jax.jit, static_argnums=(0, 1))
+      def run_simulation(amp, sigma):
+          converter = InstructionToSignals(dt=1, carriers=None)
+          with pulse.build() as schedule:
+              pulse.play(pulse.Gaussian(160, amp, sigma), pulse.DriveChannel(0))
+          signals = converter.get_signals(schedule)
+
+          # continue with simulations