Userguide and tutorial JAX transformation of Pulse Scheduler #194

to24toro · 2023-02-27T09:09:56Z

Summary

This tutorial and user guide are related to #149.
In the user guide,

show some pulses using ScalableSymbolicPulse.
jax-jit these pulses.

In the tutorial,

show the optimization of a single X gate defined by pulse schedulers.
The tutorial is much overlapping with the configuration of https://qiskit.org/documentation/dynamics/tutorials/optimizing_pulse_sequence.html .
Whether or not merging this PR depends on the discussion.

Details and comments

Currently, a new class ScalableSymbolicPulse which is a subclass of SymbolicPuse and assumed to hava a scalable form was introduced. So I use this class in this PR. However, I use the SymbolicPulse in #149. I would like to discuss which we should use.

DanPuzzuoli

I've given direct comments on the user guide. I think it is good, but just needs to be a bit simplified.

For the tutorial, I think it can be much more simplified. I don't think it is necessary to have separate 2-A, ..., 6-A and 2-B, ..., 6-B sections. I was thinking the original 2-6 steps could be kept as they are, and a single additional section numbered 7 (with subsections) could be added for doing the ScalableSymbolicPulse optimization, using the same model and fidelity function as sections 2-6. I realize that you are wanting to use DRAG, the relevance of which requires using a transmon model, but for the simplicity of the presentation I think it makes sense to just use the qubit model of section 2.

DanPuzzuoli · 2023-03-01T20:43:51Z

docs/userguide/how_to_use_pulse_schedule_for_jax_jit.rst

+.. warning::
+    This is an advanced topic --- Some extend of knowledge about ``Qiskit-pulse`` is required.
+    See `Qiskit pulse API Reference <https://qiskit.org/documentation/apidoc/pulse.html>`__ for the detailed  API information.
+    Also, see `Qiskit tutorial <https://qiskit.org/documentation/tutorials/circuits_advanced/06_building_pulse_schedules.html>`__ 
+    for a tutorial about a pulse scheduler.


I don't think this warning is necessary. I think linking to documentation for using pulse is good, but I think after that we can assume that the user knows what pulse is.

DanPuzzuoli · 2023-03-01T20:45:58Z

docs/userguide/how_to_use_pulse_schedule_for_jax_jit.rst

+.. note::
+    At present, only the ``ScalableSymbolicPulse`` class is supported by JAX.
+    We cannot use any subclasses such as ``Gaussian`` and ``GaussianSquare``
+    for JAX.


Can you add a comment here about why you can only use ScalableSymbolicPulse? Something like:

Suggested change

.. note::

At present, only the ``ScalableSymbolicPulse`` class is supported by JAX.

We cannot use any subclasses such as ``Gaussian`` and ``GaussianSquare``

for JAX.

.. note::

At present, only the ``ScalableSymbolicPulse`` class is supported by JAX, as the validation

present in other pulse types, such as ``Gaussian``, is not JAX-compatible.

It makes sense. Modified at b7ef506.

DanPuzzuoli · 2023-03-01T20:47:44Z

docs/userguide/how_to_use_pulse_schedule_for_jax_jit.rst

+    Array.set_default_backend('jax')
+
+
+2. How to generate a variety of pulses by ``ScalableSymbolicPulse``


Do you think we need 3 examples here, or would one be sufficient? My initial thought is that only one is necessary, to show the user how to set it up.

Maybe if you want to use the Guassian pulse later in the next section, you can keep only that one?

Certainly, it may be too much to users to show 3 patterns specially. Modified at b7ef506 .

DanPuzzuoli · 2023-03-01T20:52:59Z

docs/userguide/how_to_use_pulse_schedule_for_jax_jit.rst

+4. How to execute jit compilation using ``ScalableSymbolicPulse``
+-----------------------------------------------------------------
+
+Using a Gaussian pulse as an example, jit compilation is executed as follows.


I think this section should be more specifically phrased to be about how the pulse -> signal converter is JAX-compatible using ScalableSymbolicPulses. E.g. something like

Suggested change

4. How to execute jit compilation using ``ScalableSymbolicPulse``

-----------------------------------------------------------------

Using a Gaussian pulse as an example, jit compilation is executed as follows.

4. JAX transforming Pulse to Signal conversion involving ``ScalableSymbolicPulse``

------------------------------------------------------------------------------------

Using a Gaussian pulse as an example, we show that code involving ``ScalableSymbolicPulse``

and the pulse to signal converter can be JAX-compiled (or more generally, JAX-transformed).

You could also add a reference to the optimization tutorial showing this process being performed with a function that involves a simulation.

Thanks for adding the sentence. Modified at b7ef506.

DanPuzzuoli

This current iteration is looking good. I've suggestions for edits to the text.

Aside from these suggestions, one thing you should do is add direct sphinx references to classes/functions in Terra. E.g. in the DynamicsBackend docs, there are references to objects in terra, such as:

:class:`~qiskit.providers.backend.BackendV2`

(As an aside, in case you haven't seen this before, this line in docs/conf.py sets up the intersphinx package for linking documentation to other packages. In this case it is only qiskit, but qiskit_experiments is also being added in #193 .)

So, e.g., all of the references to ScalableSymbolicPulse can be turned into links to the API documentation for the class. I think you should even be able to link to the qiskit pulse module using

:mod:`~qiskit.pulse`

(not sure if this is correct syntax, but something like this should work).

DanPuzzuoli · 2023-03-03T22:33:26Z

docs/tutorials/optimizing_pulse_sequence.rst

+on :ref:`How-to use JAX with qiskit-dynamics <how-to use jax>` to
 work with JAX.


Suggested change

on :ref:`How-to use JAX with qiskit-dynamics <how-to use jax>` to

work with JAX.

on :ref:`How-to use JAX with qiskit-dynamics <how-to use jax>`.

DanPuzzuoli · 2023-03-03T22:34:38Z

docs/tutorials/optimizing_pulse_sequence.rst

+7. Optimization of X-gate using pulse schedules.
+


Suggested change

7. Optimization of X-gate using pulse schedules.

7. Repeat the :math:`X`-gate optimization, alternatively using pulse schedules to specify the control sequence.

DanPuzzuoli · 2023-03-03T22:34:56Z

docs/tutorials/optimizing_pulse_sequence.rst

@@ -254,3 +257,129 @@ approximation analysis.
 .. jupyter-execute::

    opt_signal.samples.sum()
+
+
+7. Optimization of X-gate using pulse schedules


Suggested change

7. Optimization of X-gate using pulse schedules

7. Repeat the :math:`X`-gate optimization, alternatively using pulse schedules to specify the control sequence

DanPuzzuoli · 2023-03-03T22:36:17Z

docs/tutorials/optimizing_pulse_sequence.rst

+7. Optimization of X-gate using pulse schedules
+-----------------------------------------------
+
+Here, we introduce the example of the optimization of a gate generated by pulse schedules.


Suggested change

Here, we introduce the example of the optimization of a gate generated by pulse schedules.

Here, we perform the optimization again, however now we specify the parameterized control sequence to optimize as a pulse schedule.

DanPuzzuoli · 2023-03-03T22:39:59Z

docs/tutorials/optimizing_pulse_sequence.rst

+We construct the GaussianSquare pulse which sigma and width are parametrized directly from ``ScalableSymbolicPulse``.
+Although qiskit provides a ``GaussianSquare`` class that generates a GaussianSquare pulse, which is a subclass of ``ScalableSymbolicPulse``, 
+this class is currently not JAX-compatible.
+


Suggested change

We construct the GaussianSquare pulse which sigma and width are parametrized directly from ``ScalableSymbolicPulse``.

Although qiskit provides a ``GaussianSquare`` class that generates a GaussianSquare pulse, which is a subclass of ``ScalableSymbolicPulse``,

this class is currently not JAX-compatible.

We construct a Gaussian square pulse as a ``ScalableSymbolicPulse`` instance, parameterized by ``sigma`` and ``width``.

Although qiskit pulse provides a ``GaussianSquare``, this class is not JAX compatible.

Add a reference here to the user guide as a source of more information about how ScalableSymbolicPulse is JAX compatible but other pulses are not.

Also - I've changed your initial "GaussianSquare" pulse to just the words Guassian square, to try to differentiate it from the GaussianSquare class that exists in pulse.

DanPuzzuoli · 2023-03-03T22:52:26Z

docs/userguide/how_to_use_pulse_schedule_for_jax_jit.rst

+2. How to generate a gaussion pulse by ``ScalableSymbolicPulse``
+-------------------------------------------------------------------


Suggested change

2. How to generate a gaussion pulse by ``ScalableSymbolicPulse``

-------------------------------------------------------------------

2. How to define a Gaussian pulse using ``ScalableSymbolicPulse``

-----------------------------------------------------------------

As the standard ``Gaussian`` pulse is not JAX-compatible, to define a Gaussian pulse to use in optimization,

we need to instantiate a ``ScalableSymbolicPulse`` with a Gaussian parameterization.

First, define the symbolic representation in `sympy`.

DanPuzzuoli · 2023-03-03T22:55:53Z

docs/userguide/how_to_use_pulse_schedule_for_jax_jit.rst

+Here is an example of a gaussian pulse generated using ``ScalableSymbolicPulse``.
+
+
+A Gaussian pulse is generated by the following codes.
+It notes that values of amplifier, angle, sigma and duration are typical ones.


Suggested change

Here is an example of a gaussian pulse generated using ``ScalableSymbolicPulse``.

A Gaussian pulse is generated by the following codes.

It notes that values of amplifier, angle, sigma and duration are typical ones.

Next, define the ``ScalableSymbolicPulse`` using the above expression.

DanPuzzuoli · 2023-03-03T22:57:11Z

docs/userguide/how_to_use_pulse_schedule_for_jax_jit.rst

+    gaussian_pulse.draw()
+
+
+4. JAX transforming Pulse to Signal conversion involving ``ScalableSymbolicPulse``


Suggested change

4. JAX transforming Pulse to Signal conversion involving ``ScalableSymbolicPulse``

3. JAX transforming Pulse to Signal conversion involving ``ScalableSymbolicPulse``

DanPuzzuoli · 2023-03-03T22:57:29Z

docs/userguide/how_to_use_pulse_schedule_for_jax_jit.rst

+4. JAX transforming Pulse to Signal conversion involving ``ScalableSymbolicPulse``
+------------------------------------------------------------------------------------
+
+Using a Gaussian pulse as an example, we show that code involving ``ScalableSymbolicPulse``


Suggested change

Using a Gaussian pulse as an example, we show that code involving ``ScalableSymbolicPulse``

Using a Gaussian pulse as an example, we show that a function involving ``ScalableSymbolicPulse``

DanPuzzuoli · 2023-03-03T22:57:48Z

docs/userguide/how_to_use_pulse_schedule_for_jax_jit.rst

+
+.. jupyter-execute::
+
+    # use Amplifier as a variable


Suggested change

# use Amplifier as a variable

# use amplitude as the function argument

to24toro · 2023-03-06T14:18:04Z

I modified, improved English phrases and also add sphinx refs at 55ff5bc .

DanPuzzuoli

Some final small changes. After these I will approve.

DanPuzzuoli · 2023-03-07T02:48:05Z

docs/tutorials/optimizing_pulse_sequence.rst

+        return 1. - fidelity(Array(result[0].y[-1])).data
+
+
+We set the initial values of `sigma` and `width` for the optimization as :math:`initial_params = np.array([10, 10,])`.


Suggested change

We set the initial values of `sigma` and `width` for the optimization as :math:`initial_params = np.array([10, 10,])`.

We set the initial values of ``sigma`` and ``width`` for the optimization as ``initial_params = np.array([10, 10])``.

Sorry, I somehow missed the extra comma last time. Also - all mark up here should just be double ticks, as these are code references.

Fix at 83684a2.

DanPuzzuoli · 2023-03-07T02:48:35Z

docs/tutorials/optimizing_pulse_sequence.rst

+
+.. jupyter-execute::
+
+    initial_params = np.array([10, 10,])


Suggested change

initial_params = np.array([10, 10,])

initial_params = np.array([10, 10])

Fix at 83684a2.

DanPuzzuoli · 2023-03-07T02:54:17Z

docs/tutorials/optimizing_pulse_sequence.rst

+Here, we perform the optimization again, however now we specify the parameterized control sequence to optimize as a pulse schedule.
+
+We construct a Gaussian square pulse as a :class:`~qiskit.pulse.library.ScalableSymbolicPulse` instance, parameterized by ``sigma`` and ``width``.
+Although qiskit pulse provides a :class:`~qiskit.pulse.library.GaussianSquare`, this class is not JAX compatible.


Add a reference here to the how_to_use_pulse_schedule_for_jax_jit.rst user guide. E.g.

See the user guide entry on :ref:`JAX-compatible pulse schedules <pulse jax>`.

(where pulse jax is the reference name to add to the start of the user guide).

Fix at 83684a2.

DanPuzzuoli · 2023-03-07T03:01:25Z

docs/tutorials/optimizing_pulse_sequence.rst

+
+
+
+We can draw the optimized pulse, whose parameters are retrieved by :math:`opt_results.x`.


Suggested change

We can draw the optimized pulse, whose parameters are retrieved by :math:`opt_results.x`.

We can draw the optimized pulse, whose parameters are retrieved by ``opt_results.x``.

Fix at 83684a2.

to24toro added 4 commits February 22, 2023 12:18

add how_to_use_pulse_schedule_for_jax_jit

0ac4024

add optimizing_pulse_sequence_using_pulse_schedule

6d3e153

modify docs

cf3aa96

modify useguide how_to_use_pulse_schedule_for_jax_jit

bea7355

to24toro changed the title ~~Feature/add docs for jax jit~~ Userguide and tutorial JAX transformation of Pulse Scheduler Feb 27, 2023

to24toro added 4 commits February 28, 2023 16:35

delete last sentences in how_to_use_pulse_schedule_for_jax_jit

dcced58

modify tutorials and userguide

d79222c

delete optimizing_pulse_sequence_using_pulse_schedule from index

cd1861c

change a method for optimization

76109e1

DanPuzzuoli requested changes Mar 1, 2023

View reviewed changes

to24toro and others added 3 commits March 2, 2023 08:15

change config and comment of the userguide

b7ef506

improve config of optimizing_pulse_sequence

33881a8

Merge branch 'main' into feature/add_docs_for_jax_jit

c35ea49

to24toro marked this pull request as ready for review March 2, 2023 03:09

DanPuzzuoli requested changes Mar 3, 2023

View reviewed changes

to24toro added 3 commits March 6, 2023 22:27

add qiskit_experiments to intersphinx_mapping

5a9716a

add direct sphinx refs to classes and modify sentence expressions

55ff5bc

modify title underline too short

3b4a11a

DanPuzzuoli requested changes Mar 7, 2023

View reviewed changes

to24toro and others added 2 commits March 8, 2023 15:43

modify the representation of references

83684a2

Merge branch 'main' into feature/add_docs_for_jax_jit

3a4ed20

DanPuzzuoli approved these changes Mar 8, 2023

View reviewed changes

DanPuzzuoli merged commit 96c1a14 into qiskit-community:main Mar 8, 2023

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Userguide and tutorial JAX transformation of Pulse Scheduler #194

Userguide and tutorial JAX transformation of Pulse Scheduler #194

to24toro commented Feb 27, 2023 •

edited

Loading

DanPuzzuoli left a comment

DanPuzzuoli Mar 1, 2023

DanPuzzuoli Mar 1, 2023

to24toro Mar 1, 2023 •

edited

Loading

DanPuzzuoli Mar 1, 2023

DanPuzzuoli Mar 1, 2023

to24toro Mar 1, 2023 •

edited

Loading

DanPuzzuoli Mar 1, 2023

to24toro Mar 2, 2023

DanPuzzuoli left a comment

DanPuzzuoli Mar 3, 2023

DanPuzzuoli Mar 3, 2023

DanPuzzuoli Mar 3, 2023

DanPuzzuoli Mar 3, 2023

DanPuzzuoli Mar 3, 2023

DanPuzzuoli Mar 3, 2023

DanPuzzuoli Mar 3, 2023

DanPuzzuoli Mar 3, 2023

DanPuzzuoli Mar 3, 2023

DanPuzzuoli Mar 3, 2023

DanPuzzuoli Mar 3, 2023

to24toro commented Mar 6, 2023 •

edited

Loading

DanPuzzuoli left a comment

DanPuzzuoli Mar 7, 2023

to24toro Mar 8, 2023

DanPuzzuoli Mar 7, 2023

to24toro Mar 8, 2023

DanPuzzuoli Mar 7, 2023

to24toro Mar 8, 2023

DanPuzzuoli Mar 7, 2023

to24toro Mar 8, 2023

		Array.set_default_backend('jax')


		2. How to generate a variety of pulses by ``ScalableSymbolicPulse``

		on :ref:`How-to use JAX with qiskit-dynamics <how-to use jax>` to
		work with JAX.

	7. Optimization of X-gate using pulse schedules.
	7. Repeat the :math:`X`-gate optimization, alternatively using pulse schedules to specify the control sequence.

	Here, we introduce the example of the optimization of a gate generated by pulse schedules.
	Here, we perform the optimization again, however now we specify the parameterized control sequence to optimize as a pulse schedule.

		2. How to generate a gaussion pulse by ``ScalableSymbolicPulse``
		-------------------------------------------------------------------

-. How to generate a gaussion pulse by ``ScalableSymbolicPulse``
--------------------------------------------------------------------
+. How to define a Gaussian pulse using ``ScalableSymbolicPulse``
+-----------------------------------------------------------------
+As the standard ``Gaussian`` pulse is not JAX-compatible, to define a Gaussian pulse to use in optimization,
+we need to instantiate a ``ScalableSymbolicPulse`` with a Gaussian parameterization.
+First, define the symbolic representation in `sympy`.

		gaussian_pulse.draw()


		4. JAX transforming Pulse to Signal conversion involving ``ScalableSymbolicPulse``

	4. JAX transforming Pulse to Signal conversion involving ``ScalableSymbolicPulse``
	3. JAX transforming Pulse to Signal conversion involving ``ScalableSymbolicPulse``

	Using a Gaussian pulse as an example, we show that code involving ``ScalableSymbolicPulse``
	Using a Gaussian pulse as an example, we show that a function involving ``ScalableSymbolicPulse``

	# use Amplifier as a variable
	# use amplitude as the function argument

		return 1. - fidelity(Array(result[0].y[-1])).data


		We set the initial values of `sigma` and `width` for the optimization as :math:`initial_params = np.array([10, 10,])`.

	We set the initial values of `sigma` and `width` for the optimization as :math:`initial_params = np.array([10, 10,])`.
	We set the initial values of ``sigma`` and ``width`` for the optimization as ``initial_params = np.array([10, 10])``.

	initial_params = np.array([10, 10,])
	initial_params = np.array([10, 10])




		We can draw the optimized pulse, whose parameters are retrieved by :math:`opt_results.x`.

Userguide and tutorial JAX transformation of Pulse Scheduler #194

Userguide and tutorial JAX transformation of Pulse Scheduler #194

Conversation

to24toro commented Feb 27, 2023 • edited Loading

Summary

Details and comments

DanPuzzuoli left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

to24toro Mar 1, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

to24toro Mar 1, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

DanPuzzuoli left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

to24toro commented Mar 6, 2023 • edited Loading

DanPuzzuoli left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

to24toro commented Feb 27, 2023 •

edited

Loading

to24toro Mar 1, 2023 •

edited

Loading

to24toro Mar 1, 2023 •

edited

Loading

to24toro commented Mar 6, 2023 •

edited

Loading