Bug fix when nstates=1; when expects is None; small changes to plotting#11
Bug fix when nstates=1; when expects is None; small changes to plotting#11dkweiss31 merged 13 commits intodkweiss31:mainfrom
nstates=1; when expects is None; small changes to plotting#11Conversation
cyrus-baker
commented
May 18, 2025
- fix bug when plot_expects when expects is None;
- fix dim calculation in PropagatorInfidelity;
- fix bug when initial state is only one in CoherentInfidelity
…ug when plot expects in propagator where expects is None; 2. fix dim calculation in PropagatorInfidelity; 3. fix bug when initial state is only one in CoherentInfidelity
qontrol/cost.py
Outdated
There was a problem hiding this comment.
I'm not sure this is necessary, see here for the function this calls which already specifies the last two axes
There was a problem hiding this comment.
Oh yes, this bug is caused by squeeze, not by trace
dkweiss31
left a comment
There was a problem hiding this comment.
Thanks for this PR @littlebaker! I've left some comments. I'm wondering if you might be able to adjust your solution when expects is None to avoid just having an empty plot. Could we use different default plotters for the propogator functionality? Additionally, you've uncovered a few bugs, and I'm wondering if you can add test coverage for those use cases so that they shouldn't show up again?
|
Sure, I will remove the empty plot when |
|
Sure, I will add some test code |
| # check plotter | ||
| if plotter is None: | ||
| if ( | ||
| model.exp_ops is not None | ||
| and len(model.exp_ops) > 0 | ||
| and not isinstance(model, SEPropagatorModel) | ||
| and not isinstance(model, MEPropagatorModel) | ||
| ): | ||
| plotter = DefaultPlotter() | ||
| else: | ||
| plotter = Plotter([plot_fft, plot_controls]) | ||
|
|
There was a problem hiding this comment.
I like what you're going for here. I think what this is pointing to is that we need one more layer of abstraction. The base Model class shouldn't have exp_ops as an attribute, since the propagator models don't have it. I'm thinking of something like
class Model(eqx.Module):
H_function: callable
tsave_function: callable
def __call__(
self,
parameters: Array | dict,
method: Method = Tsit5(), # noqa B008
gradient: Gradient | None = None,
options: dq.Options = dq.Options(), # noqa B008
) -> tuple[Result, TimeQArray]:
raise NotImplementedError
class SolveModel(Model):
exp_ops: Array | None
class PropagatorModel(Model):
passWould you be able to make these changes to the model file? Then the check above could become much simpler.
There was a problem hiding this comment.
Sure, that will be a good idea.
|
I think we're getting close! One quick thing: there are a lot of .pyc files that are part of this PR: can you remove them? |
|
Ok, sure. I will remove them. |
|
Sorry, wrong operation, I did not finished. |
def plot_controls(
ax: Axes, _expects: Array | None, model: Model, parameters: Array | dict
) -> Axes:
"""Plot the Hamiltonian prefactors, usually corresponding to controls."""
ax.set_facecolor('none')
H = model.H_function(parameters)
tsave = model.tsave_function(parameters)
controls = get_controls(H, tsave)
H_labels = [f'$H_{idx}$' for idx in range(len(controls))]
for idx, control in enumerate(controls):
ax.plot(tsave, np.real(control) / 2 / np.pi, label=H_labels[idx])
ax.legend(loc='lower right', framealpha=0.0)
ax.set_ylabel('pulse amplitude [GHz]')
ax.set_xlabel('time [ns]')
return axwhy |
|
Right, this is a mostly arbitrary choice that the controls included the 2pis, so that I needed to divide by 2pi when plotting. You are right though that not everyone will do this and it can be unexpected for users, so I'm happy to change this default behavior (though one can of course use their own custom |
… plot.py and optimize.py
…l. 2. remove exp_ops in PropagatorModel. 3. fix plotter creation in `optimize.py` BREAKING CHANGE: Now `sepropagator_model` and `mepropagator_model` no longer accept `exp_ops`
|
@littlebaker is this PR ready for my review? (I do still see a .json and .pyc file that shouldn't be included) |
|
Oh, sorry for that. I will check it again. |
nstates=1; when expects is None; small changes to plotting
dkweiss31
left a comment
There was a problem hiding this comment.
Thanks @littlebaker ! Everything LGTM. I will let the ci run and then we can merge :). Congrats on your first contribution!!