Repository: PennyLaneAI/pennylane
Pull request #2081 (2022-01-10 11:31:17)

Generalize Rotosolve to arbitray-frequency QNodes
This implements the generalization of Rotosolve to arbitrary frequencies, further extending its use cases.
The last generalization #1489  allowed for arbitrary _integer-valued_ frequencies, now any frequency spectra can be tackled.

`RotosolveOptimizer` makes use of `qml.fourier.reconstruct` to obtain analytical reconstructions of univariate restrictions of the QNode of interest, and then minimizes this analytical reconstruction, using numerical methods in general and an analytic expression for the minimum for single-frequency functions.

This implements #2058.


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.
dwierichs: [sc-9758]
codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2081?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#2081](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2081?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (43a83d9) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/c9211f8a7d0bfd6865b760eeda8c07e66d4a78c3?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (c9211f8) will **decrease** coverage by `2.13%`.
> The diff coverage is `98.03%`.

[![Impacted file tree graph](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2081/graphs/tree.svg?width=650&height=150&src=pr&token=9S3YzlUgGk&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2081?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)

```diff
@@            Coverage Diff             @@
##           master    #2081      +/-   ##
==========================================
- Coverage   99.18%   97.04%   -2.14%     
==========================================
  Files         226      226              
  Lines       17380    17383       +3     
==========================================
- Hits        17239    16870     -369     
- Misses        141      513     +372     
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2081?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/optimize/rotosolve.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2081/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wdGltaXplL3JvdG9zb2x2ZS5weQ==) | `98.38% <97.91%> (-1.62%)` | :arrow_down: |
| [pennylane/fourier/reconstruct.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2081/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ZvdXJpZXIvcmVjb25zdHJ1Y3QucHk=) | `100.00% <100.00%> (ø)` | |
| [pennylane/math/multi\_dispatch.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2081/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21hdGgvbXVsdGlfZGlzcGF0Y2gucHk=) | `87.06% <100.00%> (-12.94%)` | :arrow_down: |
| [pennylane/interfaces/batch/jax.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2081/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvYmF0Y2gvamF4LnB5) | `0.00% <0.00%> (-100.00%)` | :arrow_down: |
| [pennylane/interfaces/jax.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2081/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvamF4LnB5) | `0.00% <0.00%> (-97.44%)` | :arrow_down: |
| [pennylane/devices/default\_qubit\_jax.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2081/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RldmljZXMvZGVmYXVsdF9xdWJpdF9qYXgucHk=) | `0.00% <0.00%> (-96.63%)` | :arrow_down: |
| [pennylane/qnn/cost.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2081/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3Fubi9jb3N0LnB5) | `42.85% <0.00%> (-57.15%)` | :arrow_down: |
| [pennylane/transforms/classical\_jacobian.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2081/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvY2xhc3NpY2FsX2phY29iaWFuLnB5) | `80.00% <0.00%> (-20.00%)` | :arrow_down: |
| [pennylane/math/utils.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2081/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21hdGgvdXRpbHMucHk=) | `81.81% <0.00%> (-18.19%)` | :arrow_down: |
| ... and [16 more](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2081/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2081?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2081?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Last update [c9211f8...43a83d9](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2081?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).


Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 3c49a70f5..4a6d1393b 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -4,7 +4,74 @@
 
 <h3>New features since last release</h3>
 
-* A tensor network templates module has been added. Quantum circuits with the shape of a matrix product state tensor network can now be easily implemented. Motivation and theory can be found in [arXiv:1803.11537](https://arxiv.org/abs/1803.11537). [(#1871)](https://github.com/PennyLaneAI/pennylane/pull/1871)
+* The `RotosolveOptimizer` has been generalized to arbitrary frequency spectra
+  in the cost function. Also note the changes in behaviour listed under *Breaking
+  changes*.
+  [(#2081)](https://github.com/PennyLaneAI/pennylane/pull/2081)
+
+  Previously, the `RotsolveOptimizer` was available for cost functions with
+  frequency spectra that only contained integers, and the maximal frequency
+  (instead of the number of frequencies) determined the cost of the optimization.
+  Now arbitrary frequencies are supported.
+
+  Consider the QNode
+  ```python
+  dev = qml.device("default.qubit", wires=2)
+
+  @qml.qnode(dev)
+  def qnode(x, Y):
+      qml.RX(2.5 * x, wires=0)
+      qml.CNOT(wires=[0, 1])
+      qml.RZ(0.3 * Y[0], wires=0)
+      qml.CRY(1.1 * Y[1], wires=[1, 0])
+      return qml.expval(qml.PauliX(0) @ qml.PauliZ(1))
+
+  x = np.array(0.8, requires_grad=True)
+  Y = np.array([-0.2, 1.5], requires_grad=True)
+  ```
+
+  Its frequency spectra can be easily obtained via `qml.fourier.qnode_spectrum`:
+  ```pycon
+  >>> spectra = qml.fourier.qnode_spectrum(qnode)(x, Y)
+  >>> spectra
+  {'x': {(): [-2.5, 0.0, 2.5]},
+   'Y': {(0,): [-0.3, 0.0, 0.3], (1,): [-1.1, -0.55, 0.0, 0.55, 1.1]}}
+  ```
+
+  We may then initialize the `RotosolveOptimizer` and minimize the QNode cost function
+  by providing this information about the frequency spectra. We also compare the cost at
+  each step to the initial cost.
+  ```pycon
+  >>> cost_init = qnode(x, Y)
+  >>> opt = qml.RotosolveOptimizer()
+  >>> for _ in range(2):
+  ...     x, Y = opt.step(qnode, x, Y, spectra=spectra)
+  ...     print(f"New cost: {np.round(qnode(x, Y), 3)}; Initial cost: {np.round(cost_init, 3)}")
+  New cost: 0.0; Initial cost: 0.706
+  New cost: -1.0; Initial cost: 0.706
+  ```
+
+  The optimization with `RotosolveOptimizer` is performed in substeps. The minimal cost
+  of these substeps can be retrieved by setting `full_output=True`.
+  ```pycon
+  >>> x = np.array(0.8, requires_grad=True)
+  >>> Y = np.array([-0.2, 1.5], requires_grad=True)
+  >>> opt = qml.RotosolveOptimizer()
+  >>> for _ in range(2):
+  ...     (x, Y), history = opt.step(qnode, x, Y, spectra=spectra, full_output=True)
+  ...     print(f"New cost: {np.round(qnode(x, Y), 3)} reached via substeps {np.round(history, 3)}")
+  New cost: 0.0 reached via substeps [-0.  0.  0.]
+  New cost: -1.0 reached via substeps [-0.276 -0.276 -1.   ]
+  ```
+  However, note that these intermediate minimal values are evaluations of the
+  *reconstructions* that Rotosolve creates and uses internally for the optimization,
+  and not of the original objective function. For noisy cost functions, these intermediate
+  evaluations may differ significantly from evaluations of the original cost function.
+
+* A tensor network templates module has been added. Quantum circuits with the shape
+  of a matrix product state tensor network can now be easily implemented.
+  Motivation and theory can be found in [arXiv:1803.11537](https://arxiv.org/abs/1803.11537).
+  [(#1871)](https://github.com/PennyLaneAI/pennylane/pull/1871)
 
   An example circuit that uses the `MPS` template is:
   ```python
@@ -36,7 +103,8 @@
   2: ────────────────╰X──RY(0.2)──╭C──RY(-0.15)──┤
   3: ─────────────────────────────╰X──RY(0.5)────┤ ⟨Z⟩
   ```
-* Added a template for tree tensor networks (TTN). [(#2043)](https://github.com/PennyLaneAI/pennylane/pull/2043)
+* Added a template for tree tensor networks (TTN).
+  [(#2043)](https://github.com/PennyLaneAI/pennylane/pull/2043)
   An example circuit that uses the `TTN` template is:
   ```python
   import pennylane as qml
@@ -100,7 +168,7 @@
   computes the metric tensor using four copies of the state vector and
   a number of operations that scales quadratically in the number of trainable
   parameters (see below for details).
-  
+
   Note that as it makes use of state cloning, it is inherently classical
   and can only be used with statevector simulators and `shots=None`.
 
@@ -111,7 +179,7 @@
 
   ```python
   dev = qml.device("default.qubit", wires=3)
-  
+
   @qml.qnode(dev)
   def circuit(x, y):
       qml.Rot(*x[0], wires=0)
@@ -139,8 +207,8 @@
   The adjoint method uses :math:`2P^2+4P+1` gates and state cloning operations if the circuit
   is composed only of trainable gates, where :math:`P` is the number of trainable operations.
   If non-trainable gates are included, each of them is applied about :math:`n^2-n` times, where
-  :math:`n` is the number of trainable operations that follow after the respective 
-  non-trainable operation in the circuit. This means that non-trainable gates later in the 
+  :math:`n` is the number of trainable operations that follow after the respective
+  non-trainable operation in the circuit. This means that non-trainable gates later in the
   circuit are executed less often, making the adjoint method a bit cheaper if such gates
   appear later.
   The adjoint method requires memory for 4 independent state vectors, which corresponds roughly
@@ -170,7 +238,7 @@
 
   params = list(range(1, 3))
   ```
-  
+
   The produced state is
 
   ```pycon
@@ -186,8 +254,8 @@
   [ 0.47415988+0.j          0.         0.73846026j  0.         0.25903472j
    -0.40342268+0.j        ]
   ```
-  
-* A precision argument has been added to the tape's ``to_openqasm`` function 
+
+* A precision argument has been added to the tape's ``to_openqasm`` function
   to control the precision of parameters.
   [(#2071)](https://github.com/PennyLaneAI/pennylane/pull/2071)
 
@@ -201,19 +269,42 @@
   uses RetworkX for its internal representation. This results in significant speedup
   for algorithms that rely on a directed acyclic graph representation.
   [(#1791)](https://github.com/PennyLaneAI/pennylane/pull/1791)
-  
+
 * The QAOA module now accepts both NetworkX and RetworkX graphs as function inputs.
   [(#1791)](https://github.com/PennyLaneAI/pennylane/pull/1791)
 
 <h3>Breaking changes</h3>
 
+* The behaviour of `RotosolveOptimizer` has been changed regarding
+  its keyword arguments.
+  [(#2081)](https://github.com/PennyLaneAI/pennylane/pull/2081)
+
+  The keyword arguments `optimizer` and `optimizer_kwargs` for the
+  `RotosolveOptimizer` have been renamed to `substep_optimizer`
+  and `substep_kwargs`, respectively. Furthermore they have been
+  moved from `step` and `step_and_cost` to the initialization `__init__`.
+
+  The keyword argument `num_freqs` has been renamed to `nums_frequency`
+  and is expected to take a different shape now:
+  Previously, it was expected to be an `int` or a list of entries, with
+  each entry in turn being either an `int` or a `list` of `int` entries.
+  Now the expected structure is a nested dictionary, matching the
+  formatting expected by
+  [qml.fourier.reconstruct](https://pennylane.readthedocs.io/en/stable/code/api/pennylane.fourier.reconstruct.html)
+  This also matches the expected formatting of the new keyword arguments
+  `spectra` and `shifts`.
+
+  For more details, see the
+  [RotosolveOptimizer documentation](https://pennylane.readthedocs.io/en/stable/code/api/pennylane.RotosolveOptimizer.html).
+
+
 <h3>Bug fixes</h3>
 
 * The available `diff_method` options for QNodes has been corrected in both the
   error messages and the documentation.
   [(#2078)](https://github.com/PennyLaneAI/pennylane/pull/2078)
 
-* Fixes a bug in `DefaultQubit` where the second derivative of QNodes at 
+* Fixes a bug in `DefaultQubit` where the second derivative of QNodes at
   positions corresponding to vanishing state vector amplitudes is wrong.
   [(#2057)](https://github.com/PennyLaneAI/pennylane/pull/2057)
 
@@ -232,8 +323,8 @@
   through the `unitary_to_rot` optimization transform.
   [(#2015)](https://github.com/PennyLaneAI/pennylane/pull/2015)
 
-* Fixes a bug which allows using `jax.jit` to be compatible with circuits 
-  which return `qml.probs` when the `default.qubit.jax` is provided with a custom shot 
+* Fixes a bug which allows using `jax.jit` to be compatible with circuits
+  which return `qml.probs` when the `default.qubit.jax` is provided with a custom shot
   vector.
   [(#2028)](https://github.com/PennyLaneAI/pennylane/pull/2028)
 
diff --git a/pennylane/fourier/reconstruct.py b/pennylane/fourier/reconstruct.py
index 3e10f2daa..b5c4773c4 100644
--- a/pennylane/fourier/reconstruct.py
+++ b/pennylane/fourier/reconstruct.py
@@ -371,7 +371,8 @@ def reconstruct(qnode, ids=None, nums_frequency=None, spectra=None, shifts=None)
             given. Ignored if ``nums_frequency!=None``.
 
     Returns:
-        function: Function which accepts the same arguments as the QNode.
+        function: Function which accepts the same arguments as the QNode and one additional
+        keyword argument ``f0`` to provide the QNode value at the given arguments.
         When called, this function will return a dictionary of dictionaries,
         formatted like ``nums_frequency`` or ``spectra`` ,
         that contains the univariate reconstructions per QNode parameter.
@@ -616,15 +617,14 @@ def reconstruct(qnode, ids=None, nums_frequency=None, spectra=None, shifts=None)
 
     atol = 1e-8
     ids, recon_fn, jobs, need_f0 = _prepare_jobs(ids, nums_frequency, spectra, shifts, atol)
-    arg_names = list(signature(qnode.func).parameters.keys())
+    sign_fn = qnode.func if isinstance(qnode, qml.QNode) else qnode
+    arg_names = list(signature(sign_fn).parameters.keys())
     arg_idx_from_names = {arg_name: i for i, arg_name in enumerate(arg_names)}
 
     @wraps(qnode)
-    def wrapper(*args, **kwargs):
-        if need_f0:
+    def wrapper(*args, f0=None, **kwargs):
+        if f0 is None and need_f0:
             f0 = qnode(*args, **kwargs)
-        else:
-            f0 = None
 
         interface = qml.math.get_interface(args[0])
 
diff --git a/pennylane/math/multi_dispatch.py b/pennylane/math/multi_dispatch.py
index dcce7574e..e75a419ec 100644
--- a/pennylane/math/multi_dispatch.py
+++ b/pennylane/math/multi_dispatch.py
@@ -532,6 +532,8 @@ def scatter_element_add(tensor, index, value, like=None):
     tensor([[ 0.1000, 20.2000,  0.3000],
             [ 0.4000,  0.5000, 10.6000]])
     """
+    if len(np.shape(tensor)) == 0 and index == ():
+        return tensor + value
 
     interface = like or _multi_dispatch([tensor, value])
     return np.scatter_element_add(tensor, index, value, like=interface)
diff --git a/pennylane/optimize/rotosolve.py b/pennylane/optimize/rotosolve.py
index f4f4a59f1..513a487ab 100644
--- a/pennylane/optimize/rotosolve.py
+++ b/pennylane/optimize/rotosolve.py
@@ -14,78 +14,66 @@
 """Rotosolve gradient free optimizer"""
 # pylint: disable=too-many-branches,cell-var-from-loop
 
+from inspect import signature
 import numpy as np
 from scipy.optimize import brute, shgo
-from pennylane.utils import _flatten, unflatten
 
-
-def _assert_integer(x):
-    """Raises a ValueError if x is not of a sub-datatype of np.integer."""
-    x_type = type(x)
-    if not np.issubdtype(x_type, np.integer):
-        raise ValueError(
-            f"The numbers of frequencies are expected to be integers. Received {x_type}."
-        )
+import pennylane as qml
 
 
-def _validate_num_freqs(num_freqs, requires_grad):
-    """Checks whether a number of frequencies input is valid
-    Args:
-        num_freqs (int or array[int]): Number of frequencies data
-        requires_grad (list[bool]): Information whether an argument is to be optimized, per arg
-    Returns:
-        array[int]: Parsed input value to correct shape or simply return input if array
-        array[int]: List of flattened inputs, with each input being unwrapped if it has length 1
-    """
-    if num_freqs is None:
-        num_freqs = [1] * sum(requires_grad)
-        num_freqs_flat = num_freqs
-    elif np.isscalar(num_freqs):
-        _assert_integer(num_freqs)
-        num_freqs = [num_freqs] * sum(requires_grad)
-        num_freqs_flat = num_freqs
-    else:
-        num_freqs_flat = []
-        for num_frequency in num_freqs:
-            if np.isscalar(num_frequency):
-                _assert_integer(num_frequency)
-                num_frequency_flat = num_frequency
-            else:
-                num_frequency_flat = list(_flatten(num_frequency))
-                for _num_freq in num_frequency_flat:
-                    _assert_integer(_num_freq)
-                num_frequency_flat = np.array(num_frequency_flat, dtype=int)
-            num_freqs_flat.append(num_frequency_flat)
-        if len(num_freqs) != sum(requires_grad):
-            raise ValueError(
-                "The length of the provided numbers of frequencies "
-                f"({len(num_freqs)}) does not match the number of function arguments "
-                f"({sum(requires_grad)})."
-            )
-    return num_freqs, num_freqs_flat
-
-
-def _brute_optimizer(fun, num_steps, **kwargs):
-    r"""Brute force optimizer, wrapper of scipy.optimizer.brute that repeats it
-    ``num_steps`` times. Signature is as expected by ``RotosolveOptimizer._rotosolve``
-    below, providing a scalar minimal position and the function value at that position."""
-    width = 2 * np.pi
-    x_min = 0.0
+def _brute_optimizer(fun, num_steps, bounds=None, **kwargs):
+    r"""Brute force optimizer, wrapper of scipy.optimize.brute that repeats it
+    ``num_steps`` times. Signature is as expected by ``RotosolveOptimizer._min_numeric``
+    below, returning a scalar minimal position and the function value at that position."""
     Ns = kwargs.pop("Ns")
+    width = bounds[0][1] - bounds[0][0]
+    center = (bounds[0][1] + bounds[0][0]) / 2
     for _ in range(num_steps):
-        _range = (x_min - width / 2, x_min + width / 2)
-        x_min, y_min, *_ = brute(fun, ranges=(_range,), full_output=True, Ns=Ns, **kwargs)
+        range_ = (center - width / 2, center + width / 2)
+        center, y_min, *_ = brute(fun, ranges=(range_,), full_output=True, Ns=Ns, **kwargs)
+        # We only ever use this function for 1D optimization
+        center = center[0]
         width /= Ns
 
-    return x_min, y_min
+    return center, y_min
 
 
 def _shgo_optimizer(fun, **kwargs):
     r"""Wrapper for ``scipy.optimize.shgo`` (Simplicial Homology global optimizer).
-    Signature is as expected by ``RotosolveOptimizer._rotosolve`` below, providing
+    Signature is as expected by ``RotosolveOptimizer._min_numeric`` below, providing
     a scalar minimal position and the function value at that position."""
     opt_res = shgo(fun, **kwargs)
-    return opt_res.x, opt_res.fun
+    return opt_res.x[0], opt_res.fun
+
+
+def _restrict_to_univariate(fn, arg_idx, par_idx, args, kwargs):
+    r"""Restrict a function to a univariate function for given argument
+    and parameter indices.
+
+    Args:
+        fn (callable): Multivariate function
+        arg_idx (int): Index of the argument that contains the parameter to restrict
+        par_idx (tuple[int]): Index of the parameter to restrict to within the argument
+        args (tuple): Arguments at which to restrict the function.
+        kwargs (dict): Keyword arguments at which to restrict the function.
+
+    Returns:
+        callable: Univariate restriction of ``fn``. That is, this callable takes
+        a single float value as input and has the same return type as ``fn``.
+        All arguments are set to the given ``args`` and the input value to this
+        function is added to the marked parameter.
+    """
+    the_arg = args[arg_idx]
+    if len(qml.math.shape(the_arg)) == 0:
+        shift_vec = qml.math.ones_like(the_arg)
+    else:
+        shift_vec = qml.math.zeros_like(the_arg)
+        shift_vec = qml.math.scatter_element_add(shift_vec, par_idx, 1.0)
+
+    def _univariate_fn(x):
+        return fn(*args[:arg_idx], the_arg + shift_vec * x, *args[arg_idx + 1 :], **kwargs)
+
+    return _univariate_fn
 
 
 class RotosolveOptimizer:
@@ -97,9 +85,42 @@ class RotosolveOptimizer:
     separately reconstructing the cost function with respect to each circuit parameter,
     while keeping all other parameters fixed.
 
+    Args:
+        substep_optimizer (str or callable): Optimizer to use for the substeps of Rotosolve
+            that carries out a univariate (i.e., single-parameter) global optimization.
+            It must take as inputs:
+
+            - A function ``fn`` that maps scalars to scalars,
+
+            - the (keyword) argument ``bounds``, and
+
+            - optional keyword arguments.
+
+            It must return two scalars:
+
+            - The input value ``x_min`` for which ``fn`` is minimal, and
+
+            - the minimal value ``y_min=fn(x_min)`` or ``None``.
+
+            Alternatively, the following optimizers are built-in and can be chosen by
+            passing their name:
+
+            - ``"brute"``: An iterative version of
+              `SciPy's brute force optimizer <https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.brute.html>`_.
+              It evaluates the function at ``Ns`` equidistant points across the range
+              :math:`[-\pi, \pi]` and iteratively refines the range around the point
+              with the smallest cost value for ``num_steps`` times.
+
+            - ``"shgo"``: `SciPy's SHGO optimizer <https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.shgo.html>`_.
+
+        substep_kwargs (dict): Keyword arguments to be passed to the ``substep_optimizer``
+            callable. For ``substep_optimizer="shgo"``, the original keyword arguments of
+            the SciPy implementation are available, for ``substep_optimizer="brute"`` the
+            keyword arguments ``ranges``, ``Ns`` and ``num_steps`` are useful.
+
     For each parameter, a purely classical one-dimensional global optimization over the
-    interval :math:`(-\pi,\pi]` is performed, which can be replaced by a closed-form expression for
-    the optimal value if the :math:`d^{th}` parametrized gate has only two eigenvalues. In this
+    interval :math:`(-\pi,\pi]` is performed, which is replaced by a closed-form expression for
+    the optimal value if the :math:`d\text{th}` parametrized gate has only two eigenvalues. In this
     case, the optimal value :math:`\theta^*_d` is given by
 
     .. math::
@@ -112,18 +133,36 @@ class RotosolveOptimizer:
     where :math:`\left<H\right>_{\theta_d}` is the expectation value of the objective function
     restricted to only depend on the parameter :math:`\theta_d`.
 
+    .. warning ::
+
+        The built-in one-dimensional optimizers ``"brute"`` and ``"shgo"`` for the substeps
+        of a Rotosolve optimization step use the interval :math:`(-\pi,\pi]`, rescaled with
+        the inverse smallest frequency as default domain to optimize over. For complicated
+        cost functions, this domain might not be suitable for the substep optimization.
+
     The algorithm is described in further detail in
-    `Vidal and Theis (2018) <https://arxiv.org/abs/1812.06323>`_ and
-    `Ostaszewski et al. (2019) <https://arxiv.org/abs/1905.09692>`_, and the reconstruction
-    method used for more general operations is described in
+    `Vidal and Theis (2018) <https://arxiv.org/abs/1812.06323>`_,
+    `Nakanishi, Fujii and Todo (2019) <https://journals.aps.org/prresearch/abstract/10.1103/PhysRevResearch.2.043158>`_,
+    `Parrish et al. (2019) <https://arxiv.org/abs/1904.03206>`_,
+    and
+    `Ostaszewski et al. (2019) <https://quantum-journal.org/papers/q-2021-01-28-391/>`_,
+    and the reconstruction method used for more general operations is described in
     `Wierichs et al. (2021) <https://arxiv.org/abs/2107.12390>`_.
 
     **Example:**
 
     Initialize the optimizer and set the number of steps to optimize over.
-
-    >>> opt = qml.optimize.RotosolveOptimizer()
-    >>> num_steps = 10
+    Recall that the optimization with ``RotosolveOptimizer`` uses global optimization substeps
+    of univariate functions. The optimization technique for these substeps can be chosen via the
+    ``substep_optimizer`` and ``substep_kwargs`` keyword arguments.
+    Here we use the built-in iterative version of
+    `SciPy's brute force optimizer <https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.brute.html>`_
+    with four iterations.
+    We will run Rotosolve itself for three iterations.
+
+    >>> opt_kwargs = {"num_steps": 4}
+    >>> opt = qml.optimize.RotosolveOptimizer(substep_optimizer="brute", substep_kwargs=opt_kwargs)
+    >>> num_steps = 3
 
     Next, we create a QNode we wish to optimize:
 
@@ -132,340 +171,490 @@ class RotosolveOptimizer:
         dev = qml.device('default.qubit', wires=3, shots=None)
 
         @qml.qnode(dev)
-        def cost_function(rot_param, layer_par, crot_param):
-            for i, par in enumerate(rot_param):
+        def cost_function(rot_param, layer_par, crot_param, rot_weights=None, crot_weights=None):
+            for i, par in enumerate(rot_param * rot_weights):
                 qml.RX(par, wires=i)
-
             for w in dev.wires:
                 qml.RX(layer_par, wires=w)
-
-            for i, par in enumerate(crot_param):
+            for i, par in enumerate(crot_param*crot_weights):
                 qml.CRY(par, wires=[i, (i+1)%3])
-
             return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1) @ qml.PauliZ(2))
 
     This QNode is defined simply by measuring the expectation value of the tensor
     product of ``PauliZ`` operators on all qubits.
     It takes three parameters:
 
-    - ``rot_param`` controls three Pauli rotations with three parameters (one frequency each),
-    - ``layer_par`` feeds into a layer of rotations with a single parameter (three frequencies), and
-    - ``crot_param`` feeds three parameters into three controlled Pauli rotations (two frequencies
-      each).
+    - ``rot_param`` controls three Pauli rotations with three parameters, multiplied with ``rot_weights``,
+    - ``layer_par`` feeds into a layer of rotations with a single parameter, and
+    - ``crot_param`` feeds three parameters, multiplied with ``crot_weights``, into
+      three controlled Pauli rotations.
 
-    We also initialize a set of parameters for all these operations, and summarize
-    the numbers of frequencies in ``num_freqs``.
+    We also initialize a set of parameters for all these operations, and start with
+    uniform weights, i.e., all ``rot_weights`` and ``crot_weights`` are set to one.
+    This means that all frequencies with which the parameters in ``rot_param`` and
+    ``crot_param`` enter the QNode are integer-valued.
+    The number of frequencies per parameter are summarized in ``nums_frequency``.
 
     .. code-block :: python
 
-        init_param = [
+        init_param = (
             np.array([0.3, 0.2, 0.67], requires_grad=True),
             np.array(1.1, requires_grad=True),
             np.array([-0.2, 0.1, -2.5], requires_grad=True),
-        ]
+        )
+        rot_weights = np.ones(3)
+        crot_weights = np.ones(3)
 
-        num_freqs = [[1, 1, 1], 3, [2, 2, 2]]
+        nums_frequency = {
+            "rot_param": {(0,): 1, (1,): 1, (2,): 1},
+            "layer_par": {(): 3},
+            "crot_param": {(0,): 2, (1,): 2, (2,): 2},
+        }
 
     The keyword argument ``requires_grad`` can be used to determine whether the respective
     parameter should be optimized or not, following the behaviour of gradient computations and
-    gradient-based optimizers when using Autograd.
+    gradient-based optimizers when using Autograd or Torch.
+    With TensorFlow, a ``tf.Variable`` inside a ``tf.GradientTape`` may be used to
+    mark variables as trainable.
 
-    In addition, the optimization technique for the Rotosolve substeps can be chosen via the
-    ``optimizer`` and ``optimizer_kwargs`` keyword arguments.
-    As an extra feature, the minimized cost of the intermediate univariate reconstructions can
+    Now we carry out the optimization.
+    The minimized cost of the intermediate univariate reconstructions can
     be read out via ``full_output``, including the cost *after* the full Rotosolve step:
 
-    .. code-block :: python
-
-        param = init_param.copy()
-        cost_rotosolve = []
-        for step in range(num_steps):
-            param, cost, sub_cost = opt.step_and_cost(
-                cost_function,
-                *param,
-                num_freqs=num_freqs,
-                full_output=True,
-            )
-            print(f"Cost before step: {cost}")
-            print(f"Minimization substeps: {np.round(sub_cost, 6)}")
-            cost_rotosolve.extend(sub_cost)
-
-    The optimized values for ``x`` are now stored in ``param`` and (sub)steps-vs-cost can be
-    assessed by plotting ``cost_rotosolve``.
+    >>> param = init_param
+    >>> cost_rotosolve = []
+    >>> for step in range(num_steps):
+    ...     param, cost, sub_cost = opt.step_and_cost(
+    ...         cost_function,
+    ...         *param,
+    ...         nums_frequency=nums_frequency,
+    ...         full_output=True,
+    ...         rot_weights=rot_weights,
+    ...         crot_weights=crot_weights,
+    ...     )
+    ...     print(f"Cost before step: {cost}")
+    ...     print(f"Minimization substeps: {np.round(sub_cost, 6)}")
+    ...     cost_rotosolve.extend(sub_cost)
+    Cost before step: 0.04200821039253547
+    Minimization substeps: [-0.230905 -0.863336 -0.980072 -0.980072 -1.       -1.       -1.      ]
+    Cost before step: -0.9999999990681161
+    Minimization substeps: [-1. -1. -1. -1. -1. -1. -1.]
+    Cost before step: -0.9999999999999996
+    Minimization substeps: [-1. -1. -1. -1. -1. -1. -1.]
+
+    The optimized values for the parameters are now stored in ``param``
+    and the optimization behaviour can be assessed by plotting ``cost_rotosolve``,
+    which include the substeps of the Rotosolve optimization.
     The ``full_output`` feature is available for both, ``step`` and ``step_and_cost``.
 
-    The most general form ``RotosolveOptimizer`` is designed to tackle currently is any
-    trigonometric cost function with integer frequencies up to the given value
-    of ``num_freqs`` per parameter. Not all of the integers up to ``num_freqs`` have to
-    be present in the frequency spectrum. In order to tackle equidistant but non-integer
-    frequencies, we recommend rescaling the argument of the function of interest.
+    In general, the frequencies in a QNode will not be integer-valued, requiring us
+    to provide the ``RotosolveOptimizer`` not only with the number of frequencies
+    but their concrete values. For the example QNode above, this happens if the
+    weights are no longer one:
+
+    >>> rot_weights = np.array([0.4, 0.8, 1.2], requires_grad=False)
+    >>> crot_weights = np.array([0.5, 1.0, 1.5], requires_grad=False)
+    >>> spectrum_fn = qml.fourier.qnode_spectrum(cost_function)
+    >>> spectra = spectrum_fn(*param, rot_weights=rot_weights, crot_weights=crot_weights)
+    >>> spectra["rot_param"]
+    {(0,): [-0.4, 0.0, 0.4], (1,): [-0.8, 0.0, 0.8], (2,): [-1.2, 0.0, 1.2]}
+    >>> spectra["crot_param"]
+    {(0,): [-0.5, -0.25, 0.0, 0.25, 0.5], (1,): [-1.0, -0.5, 0.0, 0.5, 1.0], (2,): [-1.5, -0.75, 0.0, 0.75, 1.5]}
+
+    We may provide these spectra instead of ``nums_frequency`` to Rotosolve to
+    enable the optimization of the QNode at these weights:
+
+    >>> param = init_param
+    >>> for step in range(num_steps):
+    ...     param, cost, sub_cost = opt.step_and_cost(
+    ...         cost_function,
+    ...         *param,
+    ...         spectra=spectra,
+    ...         full_output=True,
+    ...         rot_weights = rot_weights,
+    ...         crot_weights = crot_weights,
+    ...     )
+    ...     print(f"Cost before step: {cost}")
+    ...     print(f"Minimization substeps: {np.round(sub_cost, 6)}")
+    Cost before step: 0.09299359486191039
+    Minimization substeps: [-0.268008 -0.713209 -0.24993  -0.871989 -0.907672 -0.907892 -0.940474]
+    Cost before step: -0.9404742138557066
+    Minimization substeps: [-0.940474 -1.       -1.       -1.       -1.       -1.       -1.      ]
+    Cost before step: -1.0
+    Minimization substeps: [-1. -1. -1. -1. -1. -1. -1.]
+
+    As we can see, while the optimization got a bit harder and the optimizer takes a bit longer
+    to converge than previously, Rotosolve was able to adapt to the more complicated
+    dependence on the input arguments and still found the global minimum successfully.
     """
     # pylint: disable=too-few-public-methods
 
+    def __init__(self, substep_optimizer="brute", substep_kwargs=None):
+        self.substep_kwargs = {} if substep_kwargs is None else substep_kwargs
+        if substep_optimizer == "brute":
+            self.substep_optimizer = _brute_optimizer
+            self.substep_kwargs.setdefault("num_steps", 4)
+            self.substep_kwargs.setdefault("Ns", 100)
+        elif substep_optimizer == "shgo":
+            self.substep_optimizer = _shgo_optimizer
+        else:
+            self.substep_optimizer = substep_optimizer
+
+    def _validate_inputs(self, requires_grad, args, nums_frequency, spectra):
+        """Checks that for each trainable argument either the number of
+        frequencies or the frequency spectrum is given."""
+
+        for arg, (arg_name, _requires_grad) in zip(args, requires_grad.items()):
+            if _requires_grad:
+                _nums_frequency = nums_frequency.get(arg_name, {})
+                _spectra = spectra.get(arg_name, {})
+                all_keys = set(_nums_frequency) | set(_spectra)
+
+                shape = qml.math.shape(arg)
+                indices = np.ndindex(shape) if len(shape) > 0 else [()]
+                for par_idx in indices:
+                    if par_idx not in all_keys:
+                        raise ValueError(
+                            "Neither the number of frequencies nor the frequency spectrum "
+                            f"was provided for the entry {par_idx} of argument {arg_name}."
+                        )
+
     def step_and_cost(
         self,
         objective_fn,
         *args,
-        num_freqs=None,
-        optimizer=None,
-        optimizer_kwargs=None,
+        nums_frequency=None,
+        spectra=None,
+        shifts=None,
         full_output=False,
         **kwargs,
     ):
         r"""Update args with one step of the optimizer and return the corresponding objective
-        function value prior to the step.
+        function value prior to the step. Each step includes multiple substeps, one per
+        parameter.
 
         Args:
             objective_fn (function): the objective function for optimization. It should take a
                 sequence of the values ``*args`` and a list of the gates ``generators`` as inputs,
                 and return a single value.
-            *args : variable length sequence containing the initial
-                values of the variables to be optimized over or a single float with the initial
-                value.
-            num_freqs (int or array[int]): The number of frequencies in the ``objective_fn`` per
-                parameter. If an ``int``, the same number is used for all parameters; if
-                ``array[int]``, the shape of ``args`` and ``num_freqs`` has to coincide.
-                Defaults to ``num_freqs=1``, corresponding to Pauli rotation gates.
-            optimizer (callable or str): the optimization method used for the univariate
-                minimization if there is more than one frequency with respect to the respective
-                parameter. If a callable, should have the signature
-                ``(fun, **kwargs) -> x_min, y_min``, where ``y_min`` is tracked and returned
-                if ``full_output==True`` but is not relevant to the optimization.
-                If ``"brute"`` or ``"shgo"``, the corresponding global optimizer of SciPy is used.
-                Defaults to ``"brute"``.
-            optimizer_kwargs : keyword arguments for the ``optimizer``. For ``"brute"`` and
-                ``"shgo"``, these kwargs are passed to the respective SciPy implementation.
-                Has to be given as one dictionary, *not variable length*.
+            *args (Sequence): variable length sequence containing the initial values of the
+                variables to be optimized over or a single float with the initial value.
+            nums_frequency (dict[dict]): The number of frequencies in the ``objective_fn`` per
+                parameter. The keys must correspond to argument names of the objective
+                function, the values must be dictionaries that map parameter indices (``tuple``)
+                in the argument to the number of frequencies with which it enters the objective
+                function (``int``).
+                The parameter index for a scalar QNode argument is ``()``, for
+                one-dimensional array QNode arguments, it takes the form ``(i,)`` for the
+                i-th parameter in the argument.
+            spectra (dict[dict]): Frequency spectra in the ``objective_fn`` per parameter.
+                The formatting is the same as for ``nums_frequency``, but the values
+                of the inner dictionaries must be sequences of frequencies
+                (``Sequence[float]``).
+                For each parameter, ``num_frequency`` take precedence over ``spectra``.
+            shifts (dict[dict]): Shift angles for the reconstruction per QNode parameter.
+                The keys have to be argument names of ``qnode`` and the inner dictionaries have to
+                be mappings from parameter indices to the respective shift angles to be used for
+                that parameter. For :math:`R` non-zero frequencies, there must be :math:`2R+1`
+                shifts given. Ignored if ``nums_frequency`` gives a number of frequencies
+                for the respective parameter in the QNode argument.
             full_output (bool): whether to return the intermediate minimized energy values from
-                the univariate optimization steps.
+                the univariate optimization substeps.
             **kwargs : variable length keyword arguments for the objective function.
 
         Returns:
-            tuple(list [array] or array, float): the new variable values :math:`x^{(t+1)}` and
-            the objective function output prior to the step.
+            list [array] or array: the new variable values :math:`x^{(t+1)}`.
             If a single arg is provided, list [array] is replaced by array.
-            list [float]: the intermediate energy values, only returned if ``full_output=True``.
+            float: the objective function output prior to the step.
+            list [float]: the intermediate objective values, only returned if
+            ``full_output=True``.
+
+        The optimization step consists of multiple substeps.
+
+        For each substep,
+        one of the parameters in one of the QNode arguments is singled out, and the
+        objective function is considered as univariate function (i.e., function that
+        depends on a single scalar) of that parameter.
+
+        If ``nums_frequency`` states that there is only a single frequency, or ``spectra``
+        only contains one positive frequency, for a parameter, an analytic formula is
+        used to return the minimum of the univariate restriction.
+
+        For multiple frequencies, :func:`.fourier.reconstruct` is used to reconstruct
+        the univariate restriction and a numeric minimization is performed instead.
+        The latter minimization is performed using the ``substep_optimizer`` passed to
+        ``RotosolveOptimizer`` at initialization.
+
+        .. note::
+
+            One of ``nums_frequency`` and ``spectra`` must contain information
+            about each parameter that is to be trained with ``RotosolveOptimizer``.
+            For each univariate reconstruction, the data in ``nums_frequency`` takes
+            precedence over the information in ``spectra``.
+
+        .. warning::
+
+            ``RotosolveOptimizer`` will only update parameters that are *explicitly*
+            markes as trainable. Either via ``requires_grad`` if using Autograd or PyTorch,
+            or by using `tf.Variable` tensors inside a `GradientTape` if using TensorFlow.
+
         """
-        _requires_grad = [getattr(arg, "requires_grad", True) for arg in args]
-        num_freqs, num_freqs_flat = _validate_num_freqs(num_freqs, _requires_grad)
-
-        optimizer_kwargs = optimizer_kwargs or {}
-        optimizer = optimizer or "brute"
-
-        if optimizer == "brute":
-            optimizer = _brute_optimizer
-            optimizer_kwargs.setdefault("num_steps", 4)
-            optimizer_kwargs.setdefault("Ns", 100)
-        elif optimizer == "shgo":
-            optimizer = _shgo_optimizer
-            optimizer_kwargs.setdefault("bounds", ((-np.pi, np.pi),))
-
-        # will single out one arg to change at a time
-        # these hold the arguments not getting updated
+        # todo: does this signature call cover all cases?
+        sign_fn = objective_fn.func if isinstance(objective_fn, qml.QNode) else objective_fn
+        arg_names = list(signature(sign_fn).parameters.keys())
+        requires_grad = {
+            arg_name: qml.math.requires_grad(arg) for arg_name, arg in zip(arg_names, args)
+        }
+        nums_frequency = nums_frequency or {}
+        spectra = spectra or {}
+        self._validate_inputs(requires_grad, args, nums_frequency, spectra)
+
+        # we will single out one arg to change at a time
+        # the following hold the arguments not getting updated
         before_args = []
         after_args = list(args)
-        # mutable version of args to get updated
-        args_new = list(args)
 
         # Prepare intermediate minimization results cache
         if full_output:
             y_output = []
         # Compute the very first evaluation in order to be able to cache it
         fun_at_zero = objective_fn(*args, **kwargs)
+        first_substep_in_step = True
 
-        train_arg_index = 0
-        for arg_index, arg in enumerate(args):
+        for arg_idx, (arg, arg_name) in enumerate(zip(args, arg_names)):
             del after_args[0]
 
-            if _requires_grad[arg_index]:
-                num_frequency_flat = num_freqs_flat[train_arg_index]
-                x_flat = np.fromiter(_flatten(arg), dtype=float)
-                num_params = len(x_flat)
-                if np.isscalar(num_frequency_flat):
-                    num_frequency_flat = [num_frequency_flat] * num_params
-                else:
-                    if len(num_frequency_flat) != num_params:
-                        raise ValueError(
-                            "The number of the frequency counts "
-                            f"({len(num_frequency_flat)}) for the {arg_index}th argument does "
-                            f"not match the number of parameters in that argument ({num_params})."
-                        )
-                shift_vecs = np.eye(num_params)
-
-                # Iterate over current arg:
-                for par_index, _num_frequency in enumerate(num_frequency_flat):
-                    shift_vec = shift_vecs[par_index]
-                    # univariate objective function
-                    univariate_obj_fn = lambda x: objective_fn(
-                        *before_args,
-                        unflatten(x_flat + shift_vec * x, arg),
-                        *after_args,
-                        **kwargs,
+            if not requires_grad[arg_name]:
+                before_args.append(arg)
+                continue
+            shape = qml.math.shape(arg)
+            indices = np.ndindex(shape) if len(shape) > 0 else [()]
+            for par_idx in indices:
+                _fun_at_zero = fun_at_zero if first_substep_in_step else None
+                # Set a single parameter in a single argument to be reconstructed
+                num_freq = nums_frequency.get(arg_name, {}).get(par_idx, None)
+                spectrum = spectra.get(arg_name, {}).get(par_idx, None)
+                if spectrum is not None:
+                    spectrum = np.array(spectrum)
+
+                if num_freq == 1 or (spectrum is not None and len(spectrum[spectrum > 0])) == 1:
+                    _args = before_args + [arg] + after_args
+                    univariate = _restrict_to_univariate(
+                        objective_fn, arg_idx, par_idx, _args, kwargs
                     )
-                    x_min, y_min = self._rotosolve(
-                        univariate_obj_fn,
-                        _num_frequency,
-                        optimizer,
-                        optimizer_kwargs,
-                        full_output,
-                        fun_at_zero=(fun_at_zero if train_arg_index + par_index == 0 else None),
-                    )
-                    x_flat += shift_vec * x_min
-                    if full_output:
-                        y_output.append(y_min)
-
-                args_new[arg_index] = unflatten(x_flat, arg)
-                train_arg_index += 1
+                    freq = 1.0 if num_freq is not None else spectrum[spectrum > 0][0]
+                    x_min, y_min = self.min_analytic(univariate, freq, _fun_at_zero)
+                    arg = qml.math.scatter_element_add(arg, par_idx, x_min)
 
-            # updating before_args for next loop
-            before_args.append(args_new[arg_index])
+                else:
+                    ids = {arg_name: (par_idx,)}
+                    _nums_frequency = (
+                        {arg_name: {par_idx: num_freq}} if num_freq is not None else None
+                    )
+                    _spectra = {arg_name: {par_idx: spectrum}} if spectrum is not None else None
 
+                    # Set up the reconstruction function
+                    recon_fn = qml.fourier.reconstruct(
+                        objective_fn, ids, _nums_frequency, _spectra, shifts
+                    )
+                    # Perform the reconstruction
+                    recon = recon_fn(*before_args, arg, *after_args, f0=_fun_at_zero, **kwargs)[
+                        arg_name
+                    ][par_idx]
+                    if spectrum is None:
+                        spectrum = list(range(num_freq + 1))
+                    x_min, y_min = self._min_numeric(recon, spectrum)
+
+                    # Update the currently treated argument
+                    arg = qml.math.scatter_element_add(arg, par_idx, x_min - arg[par_idx])
+                first_substep_in_step = False
+
+                if full_output:
+                    y_output.append(y_min)
+
+            # updating before_args for next argument
+            before_args.append(arg)
+
+        # All arguments have been updated and/or passed to before_args
+        args = before_args
         # unwrap arguments if only one, backward compatible and cleaner
-        if len(args_new) == 1:
-            args_new = args_new[0]
+        if len(args) == 1:
+            args = args[0]
 
         if full_output:
-            return args_new, fun_at_zero, y_output
-        return args_new, fun_at_zero
+            return args, fun_at_zero, y_output
+
+        return args, fun_at_zero
 
     def step(
         self,
         objective_fn,
         *args,
-        num_freqs=None,
-        optimizer=None,
-        optimizer_kwargs=None,
+        nums_frequency=None,
+        spectra=None,
+        shifts=None,
         full_output=False,
         **kwargs,
     ):
-        r"""Update args with one step of the optimizer.
+        r"""Update args with one step of the optimizer. Each step includes
+        multiple substeps, one per parameter.
 
         Args:
             objective_fn (function): the objective function for optimization. It should take a
                 sequence of the values ``*args`` and a list of the gates ``generators`` as inputs,
                 and return a single value.
-            *args : variable length sequence containing the initial
-                values of the variables to be optimized over or a single float with the initial
-                value.
-            num_freqs (int or array[int]): The number of frequencies in the ``objective_fn`` per
-                parameter. If an ``int``, the same number is used for all parameters; if
-                ``array[int]``, the shape of ``args`` and ``num_freqs`` has to coincide.
-                Defaults to ``num_freqs=1``, corresponding to Pauli rotation gates.
-            optimizer (callable or str): the optimization method used for the univariate
-                minimization if there is more than one frequency with respect to the respective
-                parameter. If a callable, should have the signature
-                ``(fun, **kwargs) -> x_min, y_min``, where ``y_min`` is tracked and returned
-                if ``full_output==True`` but is not relevant to the optimization.
-                If ``"brute"`` or ``"shgo"``, the corresponding global optimizer of SciPy is used.
-                Defaults to ``"brute"``.
-            optimizer_kwargs : keyword arguments for the ``optimizer``. For ``"brute"`` and
-                ``"shgo"``, these kwargs are passed to the respective SciPy implementation.
-                Has to be given as one dictionary, *not variable length*.
+            *args (Sequence): variable length sequence containing the initial values of the
+                variables to be optimized over or a single float with the initial value.
+            nums_frequency (dict[dict]): The number of frequencies in the ``objective_fn`` per
+                parameter. The keys must correspond to argument names of the objective
+                function, the values must be dictionaries that map parameter indices (``tuple``)
+                in the argument to the number of frequencies with which it enters the objective
+                function (``int``).
+                The parameter index for a scalar QNode argument is ``()``, for
+                one-dimensional array QNode arguments, it takes the form ``(i,)`` for the
+                i-th parameter in the argument.
+            spectra (dict[dict]): Frequency spectra in the ``objective_fn`` per parameter.
+                The formatting is the same as for ``nums_frequency``, but the values
+                of the inner dictionaries must be sequences of frequencies
+                (``Sequence[float]``).
+                For each parameter, ``num_frequency`` take precedence over ``spectra``.
+            shifts (dict[dict]): Shift angles for the reconstruction per QNode parameter.
+                The keys have to be argument names of ``qnode`` and the inner dictionaries have to
+                be mappings from parameter indices to the respective shift angles to be used for
+                that parameter. For :math:`R` non-zero frequencies, there must be :math:`2R+1`
+                shifts given. Ignored if ``nums_frequency`` gives a number of frequencies
+                for the respective parameter in the QNode argument.
             full_output (bool): whether to return the intermediate minimized energy values from
-                the univariate optimization steps.
+                the univariate optimization substeps.
             **kwargs : variable length keyword arguments for the objective function.
 
         Returns:
-            list [array]: the new variable values :math:`x^{(t+1)}`.
-            If a single arg is provided, list [array] is replaced by array.
-            list [float]: the intermediate energy values, only returned if ``full_output=True``.
+            list [array] or array: the new variable values :math:`x^{(t+1)}`.
+                If a single arg is provided, list [array] is replaced by array.
+            list [float]: the intermediate objective values, only returned if
+                ``full_output=True``.
+
+        The optimization step consists of multiple substeps.
+
+        For each substep,
+        one of the parameters in one of the QNode arguments is singled out, and the
+        objective function is considered as univariate function (i.e., function that
+        depends on a single scalar) of that parameter.
+
+        If ``nums_frequency`` states that there is only a single frequency, or ``spectra``
+        only contains one positive frequency, for a parameter, an analytic formula is
+        used to return the minimum of the univariate restriction.
+
+        For multiple frequencies, :func:`.fourier.reconstruct` is used to reconstruct
+        the univariate restriction and a numeric minimization is performed instead.
+        The latter minimization is performed using the ``substep_optimizer`` passed to
+        ``RotosolveOptimizer`` at initialization.
+
+        .. note::
+
+            One of ``nums_frequency`` and ``spectra`` must contain information
+            about each parameter that is to be trained with ``RotosolveOptimizer``.
+            For each univariate reconstruction, the data in ``nums_frequency`` takes
+            precedence over the information in ``spectra``.
+
+        .. warning::
+
+            ``RotosolveOptimizer`` will only update parameters that are *explicitly*
+            markes as trainable. Either via ``requires_grad`` if using Autograd or PyTorch,
+            or by using `tf.Variable` tensors inside a `GradientTape` if using TensorFlow.
         """
         x_new, _, *y_output = self.step_and_cost(
             objective_fn,
             *args,
-            num_freqs=num_freqs,
-            optimizer=optimizer,
-            optimizer_kwargs=optimizer_kwargs,
+            nums_frequency=nums_frequency,
+            spectra=spectra,
+            shifts=shifts,
             full_output=full_output,
             **kwargs,
         )
         if full_output:
-            return x_new, y_output
+            # For full_output=True, y_output was wrapped in an outer list due
+            # to the dynamic unpacking
+            return x_new, y_output[0]
+
         return x_new
 
-    @staticmethod
-    def full_reconstruction_equ(fun, num_frequency, fun_at_zero=None):
-        r"""Reconstruct a univariate trigonometric function using trigonometric interpolation.
-        See `Vidal and Theis (2018) <https://arxiv.org/abs/1812.06323>`_ or
-        `Wierichs et al. (2021) <https://arxiv.org/abs/2107.12390>`_.
+    def _min_numeric(self, objective_fn, spectrum):
+        r"""Numerically minimize a trigonometric function that depends on a
+        single parameter. Uses potentially large numbers of function evaluations,
+        depending on the used substep_optimizer. The optimization method and
+        options are stored in ``RotosolveOptimizer.substep_optimizer``
+        and ``RotosolveOptimizer.substep_kwargs``.
 
         Args:
-            fun (callable): the function to reconstruct
-            num_frequency (int): the number of (integer) frequencies present in ``fun``.
-            fun_at_zero (float): The value of ``fun`` at 0. Computed if not provided.
+            objective_fn (callable): Trigonometric function to minimize
 
         Returns:
-            callable: The reconstruction function with ``num_frequency`` frequencies,
-            coinciding with ``fun`` on the same number of points.
+            float: Position of the minimum of ``objective_fn``
+            float: Value of the minimum of ``objective_fn``
+
+        The returned position is guaranteed to lie within :math:`(-\pi, \pi]`.
         """
-        fun_at_zero = float(fun(0.0)) if fun_at_zero is None else fun_at_zero
-        mus = np.arange(1, num_frequency + 1)
-        shifts_pos = 2 * mus * np.pi / (2 * num_frequency + 1)
-        shifts_neg = -shifts_pos[::-1]
-        evals = list(map(fun, shifts_neg)) + [fun_at_zero] + list(map(fun, shifts_pos))
-        shifts = np.concatenate([shifts_neg, [0.0], shifts_pos])
-        a, b = (num_frequency + 0.5) / np.pi, 0.5 / np.pi
-        reconstruction = lambda x: np.sum(
-            np.array(
-                [
-                    _eval * np.sinc(a * (x - shift)) / np.sinc(b * (x - shift))
-                    for _eval, shift in zip(evals, shifts)
-                ]
-            )
-        )
-        return reconstruction
+        opt_kwargs = self.substep_kwargs.copy()
+        if "bounds" not in self.substep_kwargs:
+            spectrum = qml.math.array(spectrum)
+            half_width = np.pi / qml.math.min(spectrum[spectrum > 0])
+            opt_kwargs["bounds"] = ((-half_width, half_width),)
 
-    @staticmethod
-    def _rotosolve(
-        objective_fn, num_frequency, optimizer, optimizer_kwargs, full_output, fun_at_zero=None
-    ):
-        r"""The rotosolve step for a univariate (restriction of a) cost function.
+        x_min, y_min = self.substep_optimizer(objective_fn, **opt_kwargs)
+        if y_min is None:
+            y_min = objective_fn(x_min)
 
-        Updates the parameter of the ``objective_fn`` based on Equation 1 in
-        `Ostaszewski et al. (2019) <https://arxiv.org/abs/1905.09692>`_ if ``num_frequency==1``,
-        or based on a reconstruction and global minimization subroutine if ``num_frequency>1``.
+        return x_min, y_min
+
+    @staticmethod
+    def min_analytic(objective_fn, freq, f0):
+        r"""Analytically minimize a trigonometric function that depends on a
+        single parameter and has a single frequency. Uses two or
+        three function evaluations.
 
         Args:
-            objective_fn (function): the objective function for optimization. It should take a
-            num_frequency (int): the number of frequencies in the ``objective_fn``.
-            optimizer (callable or str): the optimization method used if ``num_frequency>1``.
-                If a callable, should have the signature
-                ``(fun, **kwargs) -> x_min, y_min``, where ``y_min`` is tracked and returned
-                if ``full_output==True`` but is not relevant to the optimization.
-                If ``"brute"`` or ``"shgo"``, the corresponding global optimizer of SciPy is used.
-            optimizer_kwargs : keyword arguments for the ``optimizer``. For ``"brute"`` and
-                ``"shgo"``, these kwargs are passed to the respective SciPy implementation.
-                Has to be given as one dictionary, *not variable length*.
-            full_output (bool): Whether to track the intermediate minimized energy values.
-            fun_at_zero (float): The value of ``fun`` at 0. Computed if not provided.
+            objective_fn (callable): Trigonometric function to minimize
+            freq (float): Frequency :math:`f` in the ``objective_fn``
+            f0 (float): Value of the ``objective_fn`` at zero. Reduces the
+                number of calls to the function from three to two if given.
 
         Returns:
-            x_min (float): the minimizing input of ``objective_fn`` within :math:`(-\pi, \pi]`.
-            y_min (float): the minimal value of ``objective_fn``.
+            float: Position of the minimum of ``objective_fn``
+            float: Value of the minimum of ``objective_fn``
+
+        The closed form expression used here was derived in
+        `Vidal & Theis (2018) <https://arxiv.org/abs/1812.06323>`__ ,
+        `Parrish et al (2019) <https://arxiv.org/abs/1904.03206>`__ and
+        `Ostaszewski et al (2019) <https://arxiv.org/abs/1905.09692>`__.
+        We use the notation of Appendix A of the last of these references,
+        although we allow for an arbitrary frequency instead of restricting
+        to :math:`f=1`.
+        The returned position is guaranteed to lie within :math:`(-\pi/f, \pi/f]`.
+
+        The used formula for the minimization of the :math:`d-\text{th}`
+        parameter then reads
+
+        .. math::
+
+            \theta^*_d &= \underset{\theta_d}{\text{argmin}}\left<H\right>_{\theta_d}\\
+                  &= -\frac{\pi}{2f} - \frac{1}{f}\text{arctan2}\left(2\left<H\right>_{\theta_d=0}
+                  - \left<H\right>_{\theta_d=\pi/(2f)} - \left<H\right>_{\theta_d=-\pi/(2f)},
+                  \left<H\right>_{\theta_d=\pi/(2f)} - \left<H\right>_{\theta_d=-\pi/(2f)}\right),
+
         """
-        # pylint: disable=too-many-arguments
-        fun_at_zero = float(objective_fn(0.0)) if fun_at_zero is None else fun_at_zero
-        # Use closed form expression from Ostaszewski et al., using notation of App. A
-        if num_frequency == 1:
-            H_p = float(objective_fn(0.5 * np.pi))
-            H_m = float(objective_fn(-0.5 * np.pi))
-            C = 0.5 * (H_p + H_m)
-            B = np.arctan2(2 * (fun_at_zero - C), H_p - H_m)
-            x_min = -np.pi / 2 - B
-            if full_output:
-                A = np.sqrt((fun_at_zero - C) ** 2 + 0.25 * (H_p - H_m) ** 2)
-                y_min = -A + C
-            else:
-                y_min = None
-        else:
-            reconstruction = RotosolveOptimizer.full_reconstruction_equ(
-                objective_fn, num_frequency, fun_at_zero
-            )
-            x_min, y_min = optimizer(reconstruction, **optimizer_kwargs)
-            if y_min is None and full_output:
-                y_min = reconstruction(x_min)
-
-        if x_min <= -np.pi:
-            x_min += 2 * np.pi
+        if f0 is None:
+            f0 = objective_fn(0.0)
+        shift = 0.5 * np.pi / freq
+        fp = objective_fn(shift)
+        fm = objective_fn(-shift)
+        C = 0.5 * (fp + fm)
+        B = np.arctan2(2 * f0 - fp - fm, fp - fm)
+        x_min = -shift - B / freq
+        A = np.sqrt((f0 - C) ** 2 + 0.25 * (fp - fm) ** 2)
+        y_min = -A + C
+
+        if x_min <= -2 * shift:
+            x_min = x_min + 4 * shift
 
         return x_min, y_min


Test patch:
diff --git a/tests/optimize/test_optimize.py b/tests/optimize/test_optimize.py
index 057c47cfa..6605dc0bd 100644
--- a/tests/optimize/test_optimize.py
+++ b/tests/optimize/test_optimize.py
@@ -26,7 +26,6 @@ from pennylane.optimize import (
     AdagradOptimizer,
     RMSPropOptimizer,
     AdamOptimizer,
-    RotosolveOptimizer,
 )
 
 
@@ -52,11 +51,8 @@ def opt(opt_name):
     if opt_name == "adam":
         return AdamOptimizer(stepsize, beta1=gamma, beta2=delta)
 
-    if opt_name == "roto":
-        return RotosolveOptimizer()
 
-
-@pytest.mark.parametrize("opt_name", ["gd", "moment", "nest", "ada", "rms", "adam", "roto"])
+@pytest.mark.parametrize("opt_name", ["gd", "moment", "nest", "ada", "rms", "adam"])
 class TestOverOpts:
     """Tests keywords, multiple arguements, and non-training arguments in relevant optimizers"""
 
@@ -83,12 +79,8 @@ class TestOverOpts:
 
         args3, kwargs3 = spy.call_args_list[-1]
 
-        if opt_name != "roto":
-            assert args2 == (x,)
-            assert args3 == (x,)
-        else:
-            assert not np.allclose(x_new_two, x, atol=tol)
-            assert not np.allclose(x_new_three_wc, x, atol=tol)
+        assert args2 == (x,)
+        assert args3 == (x,)
 
         assert kwargs2 == {"c": 2.0}
         assert kwargs3 == {"c": 3.0}
@@ -118,13 +110,8 @@ class TestOverOpts:
         self.reset(opt)
         args_called2, kwargs2 = spy.call_args_list[-1]  # just take last call
 
-        if opt_name != "roto":
-            assert args_called1 == (x, y, z)
-            assert args_called2 == (x_new, y_new, z_new)
-        else:
-            assert not np.allclose(x_new, x, atol=tol)
-            assert not np.allclose(y_new, y, atol=tol)
-            assert not np.allclose(z_new, z, atol=tol)
+        assert args_called1 == (x, y, z)
+        assert args_called2 == (x_new, y_new, z_new)
 
         assert kwargs1 == {}
         assert kwargs2 == {}
diff --git a/tests/optimize/test_rotosolve.py b/tests/optimize/test_rotosolve.py
index 2a7148e0c..1f10e5764 100644
--- a/tests/optimize/test_rotosolve.py
+++ b/tests/optimize/test_rotosolve.py
@@ -14,6 +14,7 @@
 """
 Unit tests for the ``RotosolveOptimizer``.
 """
+import functools
 import pytest
 from scipy.optimize import shgo
 
@@ -53,72 +54,43 @@ def successive_params(par1, par2):
     return walking_param
 
 
-@pytest.mark.parametrize(
-    "fun, param, num_freq",
-    zip(
-        [
-            lambda x: np.sin(x),
-            lambda x: np.sin(x),
-            lambda x, y: np.sin(x) * np.sin(y),
-            lambda x, y: np.sin(x) * np.sin(y[0]) * np.sin(y[1]),
-        ],
-        [[0.5], [0.5], [0.5, 0.2], [0.5, [0.2, 0.4]]],
-        [[], [1, 1], [1], [1, 1, [1, 2]]],
-    ),
-)
-def test_wrong_len_num_freqs(fun, param, num_freq):
-    """Test that an error is raised for a different number of
-    numbers of frequencies than number of function arguments."""
+def test_error_missing_frequency_info():
+    """Test that an error is raised if neither nums_frequency nor spectra is given."""
 
     opt = RotosolveOptimizer()
+    fun = lambda x: x
+    x = np.array(0.5, requires_grad=True)
 
-    with pytest.raises(ValueError, match="The length of the provided numbers of frequencies"):
-        opt.step(fun, *param, num_freqs=num_freq)
+    with pytest.raises(ValueError, match="Neither the number of frequencies nor the"):
+        opt.step(fun, x)
 
 
-@pytest.mark.parametrize(
-    "fun, param, num_freq",
-    zip(
-        [
-            lambda x: np.sin(x),
-            lambda x: np.sin(x),
-            lambda x, y: np.sin(x) * np.sin(y),
-            lambda x, y: np.sin(x) * np.sin(y[0]) * np.sin(y[1]),
-        ],
-        [[0.5], [0.5], [0.5, 0.2], [0.5, [0.2, 0.4]]],
-        [[[1, 1]], [[]], [[1], [1, 1]], [[1], [1]]],
-    ),
-)
-def test_wrong_num_of_num_freqs_per_parameter(fun, param, num_freq):
-    """Test that an error is raised for a different number of
-    numbers of frequencies than number of function arguments."""
+def test_no_error_missing_frequency_info_untrainable():
+    """Test that no error is raised if neither nums_frequency nor spectra
+    is given for a parameter not marked as trainable."""
 
     opt = RotosolveOptimizer()
+    fun = lambda x, y: x
+    x = np.array(0.5, requires_grad=True)
+    y = np.array(0.1, requires_grad=False)
+    nums_frequency = {"x": {(): 1}}
 
-    with pytest.raises(ValueError, match="The number of the frequency counts"):
-        opt.step(fun, *param, num_freqs=num_freq)
+    opt.step(fun, x, y, nums_frequency=nums_frequency)
 
 
-@pytest.mark.parametrize(
-    "fun, param, num_freq",
-    zip(
-        [
-            lambda x: np.sin(x),
-            lambda x: np.sin(x),
-            lambda x, y: np.sin(x) * np.sin(y),
-            lambda x, y: np.sin(x) * np.sin(y[0]) * np.sin(y[1]),
-        ],
-        [[0.5], [0.5], [0.5, 0.2], [0.5, [0.2, 0.4]]],
-        [[0.1], [1.0], [1, 1.0], [1, [1, 2 + 1j]]],
-    ),
-)
-def test_wrong_typed_num_freqs(fun, param, num_freq):
-    """Test that an error is raised for a non-integer entry in the numbers of frequencies."""
+def test_error_missing_frequency_info_single_par():
+    """Test that an error is raised if neither nums_frequency nor spectra is given
+    for one of the function arguments."""
 
     opt = RotosolveOptimizer()
+    fun = lambda x: qml.math.sum(x)
+    x = np.arange(4, requires_grad=True)
+    nums_frequency = {"x": {(0,): 1, (1,): 1}}
+    spectra = {"x": {(0,): [0.0, 1.0], (2,): [0.0, 1.0]}}
 
-    with pytest.raises(ValueError, match="The numbers of frequencies are expected to be integers."):
-        opt.step(fun, *param, num_freqs=num_freq)
+    # For the first three entries either nums_frequency or spectra is provided
+    with pytest.raises(ValueError, match=r"was provided for the entry \(3,\)"):
+        opt.step(fun, x, nums_frequency=nums_frequency, spectra=spectra)
 
 
 classical_functions = [
@@ -133,12 +105,15 @@ classical_functions = [
         - np.cos(2 * y[1] - 1.35) * 0.111
     ),
     lambda x, y, z: -np.cos(x + 0.12) * 0.872 + np.sin(y - 2.01) - np.cos(z - 1.35) * 0.111,
-    lambda x, y, z: -np.cos(x + 0.06)
-    - np.cos(2 * x + 0.12) * 0.872
-    + np.sin(y - 2.01 / 3 - np.pi / 3)
-    + np.sin(3 * y - 2.01)
-    - np.cos(z - 1.35) * 0.111,
+    lambda x, y, z: (
+        -np.cos(x + 0.06)
+        - np.cos(2 * x + 0.12) * 0.872
+        + np.sin(y - 2.01 / 3 - np.pi / 3)
+        + np.sin(3 * y - 2.01)
+        - np.cos(z - 1.35) * 0.111
+    ),
 ]
+
 classical_minima = [
     (-np.pi / 2 - 0.124,),
     ([-0.12, -np.pi / 2 + 2.01, 1.35],),
@@ -147,6 +122,7 @@ classical_minima = [
     (-0.12, -np.pi / 2 + 2.01, 1.35),
     (-0.12 / 2, (-np.pi / 2 + 2.01) / 3, 1.35),
 ]
+
 classical_params = [
     (0.24,),
     ([0.2, -0.3, 0.1],),
@@ -155,72 +131,104 @@ classical_params = [
     (0.1, 0.2, 0.5),
     (0.9, 0.7, 0.2),
 ]
-classical_num_freqs = [[1], [[1, 1, 1]], [1, [1, 1]], [1, 2], 1, [2, 3, 1]]
+
+classical_nums_frequency = [
+    {"x": {(): 1}},
+    {"x": {(0,): 1, (1,): 1, (2,): 1}},
+    {"x": {(): 1}, "y": {(0,): 1, (1,): 1}},
+    {"x": {(): 1}, "y": {(0,): 2, (1,): 2}},
+    {"x": {(): 1}, "y": {(): 1}, "z": {(): 1}},
+    {"x": {(): 2}, "y": {(): 3}, "z": {(): 1}},
+]
+
+classical_expected_num_calls = [3, 9, 9, 13, 9, 15]
 
 
 def custom_optimizer(fun, **kwargs):
     r"""Wrapper for ``scipy.optimize.shgo`` that does not return y_min."""
     opt_res = shgo(fun, **kwargs)
-    return opt_res.x, None
+    return opt_res.x[0], None
 
 
-optimizers = [None, "brute", "shgo", custom_optimizer]
-optimizer_kwargs = [
+substep_optimizers = ["brute", "shgo", custom_optimizer]
+all_substep_kwargs = [
     {"Ns": 93, "num_steps": 3},
-    None,
     {"bounds": ((-1.0, 1.0),), "n": 512},
     {"bounds": ((-1.1, 1.4),)},
 ]
 
 
 @pytest.mark.parametrize(
-    "fun, x_min, param, num_freq",
-    list(zip(classical_functions, classical_minima, classical_params, classical_num_freqs)),
+    "fun, x_min, param, nums_freq, exp_num_calls",
+    list(
+        zip(
+            classical_functions,
+            classical_minima,
+            classical_params,
+            classical_nums_frequency,
+            classical_expected_num_calls,
+        )
+    ),
 )
 @pytest.mark.parametrize(
-    "optimizer, optimizer_kwargs",
-    list(zip(optimizers, optimizer_kwargs)),
+    "substep_optimizer, substep_kwargs",
+    list(zip(substep_optimizers, all_substep_kwargs)),
 )
-class TestWithClassicalFunctions:
+class TestWithClassicalFunction:
     def test_number_of_function_calls(
-        self, fun, x_min, param, num_freq, optimizer, optimizer_kwargs
+        self, fun, x_min, param, nums_freq, exp_num_calls, substep_optimizer, substep_kwargs
     ):
         """Tests that per parameter 2R+1 function calls are used for an update step."""
         global num_calls
         num_calls = 0
 
+        @functools.wraps(fun)
         def _fun(*args, **kwargs):
             global num_calls
             num_calls += 1
             return fun(*args, **kwargs)
 
-        opt = RotosolveOptimizer()
-        new_param = opt.step(
-            _fun,
-            *param,
-            num_freqs=num_freq,
-            optimizer=optimizer,
-            optimizer_kwargs=optimizer_kwargs,
-        )
+        opt = RotosolveOptimizer(substep_optimizer, substep_kwargs)
 
-        expected_num_calls = np.sum(
-            np.fromiter(_flatten(expand_num_freq(num_freq, param)), dtype=int) * 2 + 1
-        )
-        assert num_calls == expected_num_calls
+        # Parameters are not marked as trainable -> Expect only one execution for f0
+        new_param = opt.step(_fun, *param, nums_frequency=nums_freq)
+        assert num_calls == 1
+        num_calls = 0
+
+        # Parameters are now marked as trainable -> Expect full number of executions
+        param = tuple(np.array(p, requires_grad=True) for p in param)
+        new_param = opt.step(_fun, *param, nums_frequency=nums_freq)
+        assert num_calls == exp_num_calls
 
     def test_single_step_convergence(
-        self, fun, x_min, param, num_freq, optimizer, optimizer_kwargs
+        self, fun, x_min, param, nums_freq, exp_num_calls, substep_optimizer, substep_kwargs
     ):
         """Tests convergence for easy classical functions in a single Rotosolve step.
         Includes testing of the parameter output shape and the old cost when using step_and_cost."""
-        opt = RotosolveOptimizer()
+        opt = RotosolveOptimizer(substep_optimizer, substep_kwargs)
 
+        # Without trainable parameters, param should not be changed
         new_param_step = opt.step(
             fun,
             *param,
-            num_freqs=num_freq,
-            optimizer=optimizer,
-            optimizer_kwargs=optimizer_kwargs,
+            nums_frequency=nums_freq,
+        )
+        # The following accounts for the unpacking functionality for length-1 param
+        if len(param) == 1:
+            new_param_step = (new_param_step,)
+
+        assert np.allclose(
+            np.fromiter(_flatten(param), dtype=float),
+            np.fromiter(_flatten(new_param_step), dtype=float),
+            atol=1e-5,
+        )
+
+        # With trainable parameters, training should happen
+        param = tuple(np.array(p, requires_grad=True) for p in param)
+        new_param_step = opt.step(
+            fun,
+            *param,
+            nums_frequency=nums_freq,
         )
         # The following accounts for the unpacking functionality for length-1 param
         if len(param) == 1:
@@ -233,12 +241,11 @@ class TestWithClassicalFunctions:
             atol=1e-5,
         )
 
+        # Now with step_and_cost and trainable params
         new_param_step_and_cost, old_cost = opt.step_and_cost(
             fun,
             *param,
-            num_freqs=num_freq,
-            optimizer=optimizer,
-            optimizer_kwargs=optimizer_kwargs,
+            nums_frequency=nums_freq,
         )
         # The following accounts for the unpacking functionality for length-1 param
         if len(param) == 1:
@@ -252,30 +259,29 @@ class TestWithClassicalFunctions:
         )
         assert np.isclose(old_cost, fun(*param))
 
-    def test_full_output(self, fun, x_min, param, num_freq, optimizer, optimizer_kwargs):
+    def test_full_output(
+        self, fun, x_min, param, nums_freq, exp_num_calls, substep_optimizer, substep_kwargs
+    ):
         """Tests the ``full_output`` feature of Rotosolve, delivering intermediate cost
         function values at the univariate optimization substeps."""
-        opt = RotosolveOptimizer()
+        param = tuple(np.array(p, requires_grad=True) for p in param)
+        opt = RotosolveOptimizer(substep_optimizer, substep_kwargs)
 
         _, y_output_step = opt.step(
             fun,
             *param,
-            num_freqs=num_freq,
-            optimizer=optimizer,
-            optimizer_kwargs=optimizer_kwargs,
+            nums_frequency=nums_freq,
             full_output=True,
         )
         new_param, old_cost, y_output_step_and_cost = opt.step_and_cost(
             fun,
             *param,
-            num_freqs=num_freq,
-            optimizer=optimizer,
-            optimizer_kwargs=optimizer_kwargs,
+            nums_frequency=nums_freq,
             full_output=True,
         )
         # The following accounts for the unpacking functionality for length-1 param
         if len(param) == 1:
-            new_param_step = (new_param,)
+            new_param = (new_param,)
         expected_intermediate_x = successive_params(param, new_param)
         expected_y_output = [fun(*par) for par in expected_intermediate_x[1:]]
 
@@ -286,23 +292,22 @@ class TestWithClassicalFunctions:
 
 @pytest.mark.parametrize(
     "fun, x_min, param, num_freq",
-    list(zip(classical_functions, classical_minima, classical_params, classical_num_freqs)),
+    list(zip(classical_functions, classical_minima, classical_params, classical_nums_frequency)),
 )
 def test_multiple_steps(fun, x_min, param, num_freq):
     """Tests that repeated steps execute as expected."""
-    opt = RotosolveOptimizer()
+    param = tuple(np.array(p, requires_grad=True) for p in param)
+    substep_optimizer = "brute"
+    substep_kwargs = None
+    opt = RotosolveOptimizer(substep_optimizer, substep_kwargs)
 
-    optimizer = "brute"
-    optimizer_kwargs = None
     for _ in range(3):
         param = opt.step(
             fun,
             *param,
-            num_freqs=num_freq,
-            optimizer=optimizer,
-            optimizer_kwargs=optimizer_kwargs,
+            nums_frequency=num_freq,
         )
-        # The following accounts for the unpacking functionality for length-1 param
+        # The following accounts for the unpacking functionality for length-one param
         if len(x_min) == 1:
             param = (param,)
 
@@ -326,10 +331,18 @@ classical_minima_deact = [
 ]
 classical_params_deact = [
     (np.array(0.24, requires_grad=False),),
-    (0.3, np.array([0.8, 0.1], requires_grad=False)),
-    (0.1, np.array(0.2, requires_grad=False), 0.5),
+    (np.array(0.3, requires_grad=True), np.array([0.8, 0.1], requires_grad=False)),
+    (
+        np.array(0.1, requires_grad=True),
+        np.array(0.2, requires_grad=False),
+        np.array(0.5, requires_grad=True),
+    ),
+]
+classical_nums_frequency_deact = [
+    {"x": {(): 1}},
+    {"x": {(): 1}, "y": {(0,): 1, (1,): 1}},
+    {"x": {(): 1}, "y": {(): 1}, "z": {(): 1}},
 ]
-classical_num_freqs_deact = [[], [1], 1]
 
 
 @pytest.mark.parametrize(
@@ -339,7 +352,7 @@ classical_num_freqs_deact = [[], [1], 1]
             classical_functions_deact,
             classical_minima_deact,
             classical_params_deact,
-            classical_num_freqs_deact,
+            classical_nums_frequency_deact,
         )
     ),
 )
@@ -348,13 +361,14 @@ class TestDeactivatedTrainingWithClassicalFunctions:
         """Tests convergence for easy classical functions in a single Rotosolve step
         with some arguments deactivated for training.
         Includes testing of the parameter output shape and the old cost when using step_and_cost."""
-        opt = RotosolveOptimizer()
+        substep_optimizer = "brute"
+        substep_kwargs = None
+        opt = RotosolveOptimizer(substep_optimizer, substep_kwargs)
 
         new_param_step = opt.step(
             fun,
             *param,
-            num_freqs=num_freq,
-            optimizer="brute",
+            nums_frequency=num_freq,
         )
         # The following accounts for the unpacking functionality for length-1 param
         if len(param) == 1:
@@ -370,8 +384,7 @@ class TestDeactivatedTrainingWithClassicalFunctions:
         new_param_step_and_cost, old_cost = opt.step_and_cost(
             fun,
             *param,
-            num_freqs=num_freq,
-            optimizer="brute",
+            nums_frequency=num_freq,
         )
         # The following accounts for the unpacking functionality for length-1 param
         if len(param) == 1:
@@ -433,32 +446,44 @@ qnode_params = [
     (np.array([0.1, -0.3, 2.9]), 1.3, [0.2, 0.1]),
     (1.2, -2.3, -0.2),
 ]
-qnode_num_freqs = [
-    [num_wires],
-    [1, 2 * num_wires, [1, 2]],
-    num_wires,
+qnode_nums_frequency = [
+    {"x": {(): num_wires}},
+    {"x": {(0,): 1, (1,): 1, (2,): 1}, "y": {(): 2 * num_wires}, "z": {(0,): 1, (1,): 2}},
+    None,
+]
+qnode_spectra = [
+    None,
+    None,
+    {
+        "x": {(): list(range(num_wires + 1))},
+        "y": {(): list(range(num_wires + 1))},
+        "z": {(): list(range(num_wires + 1))},
+    },
 ]
 
 
 @pytest.mark.parametrize(
-    "qnode, param, num_freq",
-    list(zip(qnodes, qnode_params, qnode_num_freqs)),
+    "qnode, param, nums_frequency, spectra",
+    list(zip(qnodes, qnode_params, qnode_nums_frequency, qnode_spectra)),
 )
 @pytest.mark.parametrize(
-    "optimizer, optimizer_kwargs",
-    list(zip(optimizers, optimizer_kwargs)),
+    "substep_optimizer, substep_kwargs",
+    list(zip(substep_optimizers, all_substep_kwargs)),
 )
 class TestWithQNodes:
-    def test_single_step(self, qnode, param, num_freq, optimizer, optimizer_kwargs):
-        opt = RotosolveOptimizer()
+    def test_single_step(
+        self, qnode, param, nums_frequency, spectra, substep_optimizer, substep_kwargs
+    ):
+        """Test executing a single step of the RotosolveOptimizer on a QNode."""
+        param = tuple(np.array(p, requires_grad=True) for p in param)
+        opt = RotosolveOptimizer(substep_optimizer, substep_kwargs)
 
         repack_param = len(param) == 1
         new_param_step = opt.step(
             qnode,
             *param,
-            num_freqs=num_freq,
-            optimizer=optimizer,
-            optimizer_kwargs=optimizer_kwargs,
+            nums_frequency=nums_frequency,
+            spectra=spectra,
         )
         if repack_param:
             new_param_step = (new_param_step,)
@@ -469,9 +494,8 @@ class TestWithQNodes:
         new_param_step_and_cost, old_cost = opt.step_and_cost(
             qnode,
             *param,
-            num_freqs=num_freq,
-            optimizer=optimizer,
-            optimizer_kwargs=optimizer_kwargs,
+            nums_frequency=nums_frequency,
+            spectra=spectra,
         )
         if repack_param:
             new_param_step_and_cost = (new_param_step_and_cost,)
@@ -482,8 +506,15 @@ class TestWithQNodes:
         )
         assert np.isclose(qnode(*param), old_cost)
 
-    def test_multiple_steps(self, qnode, param, num_freq, optimizer, optimizer_kwargs):
-        opt = RotosolveOptimizer()
+    def test_multiple_steps(
+        self, qnode, param, nums_frequency, spectra, substep_optimizer, substep_kwargs
+    ):
+        """Test executing multiple steps of the RotosolveOptimizer on a QNode."""
+        param = tuple(np.array(p, requires_grad=True) for p in param)
+        # For the following 1D substep_optimizer, the bounds need to be expanded for these QNodes
+        if substep_optimizer in ["shgo", custom_optimizer]:
+            substep_kwargs["bounds"] = ((-2.0, 2.0),)
+        opt = RotosolveOptimizer(substep_optimizer, substep_kwargs)
 
         repack_param = len(param) == 1
         initial_cost = qnode(*param)
@@ -492,9 +523,8 @@ class TestWithQNodes:
             param = opt.step(
                 qnode,
                 *param,
-                num_freqs=num_freq,
-                optimizer=optimizer,
-                optimizer_kwargs=optimizer_kwargs,
+                nums_frequency=nums_frequency,
+                spectra=spectra,
             )
             # The following accounts for the unpacking functionality for length-1 param
             if repack_param: