mpl-interactions suggestion

[ianhi]

Hi there!

Developer of mpl-interactions here - happy to see that you're using it :)

Just wanted to pop by and say that I added some features that may be useful to you. Specifically the ability to use controls as contextmanager: https://mpl-interactions.readthedocs.io/en/stable/examples/context.html

which would allow you simplify this: https://github.com/tomhepz/diatom-notebooks/blob/8b8786ceac4f602b9074106591b140a0d7f6884d/scripts/diatom-zeeman.py#L442-L454

to

# set current matplotlib axes
plt.sca(ax_rabi_osc)
with controls["bi"]:
    iplt.plot(times * 1e6, lambda t, bi: np.abs(state_vectors[bi,:,0])**2,c=STATE_CMAP[0])
    iplt.plot(times * 1e6, lambda t, bi: np.abs(state_vectors[bi,:,1])**2,c=STATE_CMAP[1])
    iplt.plot(times * 1e6, lambda t, bi: np.abs(state_vectors[bi,:,2])**2,c=STATE_CMAP[2])
    iplt.plot(times * 1e6, lambda t, bi: np.abs(state_vectors[bi,:,3])**2,c=STATE_CMAP[3])
    iplt.plot(times * 1e6, lambda t, bi: np.abs(state_vectors[bi,:,4])**2,c=STATE_CMAP[4])
    iplt.plot(times * 1e6, lambda t, bi: np.abs(state_vectors[bi,:,5])**2,c=STATE_CMAP[5])
    iplt.plot(times * 1e6, lambda t, bi: np.abs(state_vectors[bi,:,6])**2,c=STATE_CMAP[6])
    iplt.plot(times * 1e6, lambda t, bi: np.abs(state_vectors[bi,:,7])**2,c=STATE_CMAP[7])
    iplt.plot(times * 1e6, lambda t, bi: np.abs(state_vectors[bi,:,8])**2,c=STATE_CMAP[8])
    iplt.plot(times * 1e6, lambda t, bi: np.abs(state_vectors[bi,:,9])**2,c=STATE_CMAP[9])
    iplt.plot(times * 1e6, lambda t, bi: np.abs(state_vectors[bi,:,10])**2, c=STATE_CMAP[10])

    iplt.title(lambda bi: f"$E_0$={E_0[bi]:.2f} V/m")

although you can acutally go even further:

# set current matplotlib axes
plt.sca(ax_rabi_osc)
with controls["bi"]:
    for i in range(10):
        iplt.plot(times * 1e6, lambda t, bi: np.abs(state_vectors[bi,:,i])**2,c=STATE_CMAP[i])

    iplt.title(lambda bi: f"$E_0$={E_0[bi]:.2f} V/m")

also if you have ever have any features requests for mpl-interactions feel to open an issue here :) https://github.com/ianhi/mpl-interactions/issues/new/choose

[tomhepz]

Hi Ian, thanks for the mention, I hadn't dug into that bit of the documentation! Unfortunately, the for loop never seems to work, hence why I've done them manually. With the for loop, they display fine until the slider updates and then they disappear! I thought this was a garbage collection issue as described somewhere in the documentation, however I can keep a reference to them alive and it still doesn't work.

Without the for loop (as in your second displayed code segment) works exactly as expected. With the for loop, you can see below in the bottom left plot after moving the slider all but one of the plots dissapears! Any ideas on this one?

And thank you for the wondeful package!

Before:

After:

[ianhi]

That for loop behavior seems to be due to an intricacy of python scoping. See this SO answer: https://stackoverflow.com/a/2295368/835607

A smaller reproducible example that has the same collapsing behavior is

%matplotlib ipympl
import matplotlib.pyplot as plt
from mpl_interactions import ipyplot as iplt
from mpl_interactions.controller import Controls

ctrls = Controls(tau=(.5,1,100))
display(ctrls)
x = np.linspace(-5,5)
fig, ax = plt.subplots()
with ctrls:
    for i in range(10):
        iplt.plot(x, lambda t, tau: np.sin(tau*x)**i, lw=10-i)

because of how i set the linewidths you can see how the functions are all still there, but they are all plotting the same thing because i now equals 10 for all the lamdas:

from SO it seems that an easy fix is to just add i=i to the creation of the lambda: iplt.plot(x, lambda t, tau: np.sin(tau*x)**i, lw=10-i) which works for me.

that was very interesting! I learned something new there :)

[ianhi]

Another canonical way to do this sort of thing is to use functools.partial or make a function that returns functions like so:

%matplotlib ipympl
import matplotlib.pyplot as plt
from mpl_interactions import ipyplot as iplt
from mpl_interactions.controller import Controls

ctrls = Controls(tau=(.5,1,100))
display(ctrls)
x = np.linspace(-5,5)
fig, ax = plt.subplots()

def make_func(i):
    def f(t, tau):
        return np.sin(tau*x)**i
    return f

with ctrls:
    for i in range(10):
        iplt.plot(x, make_func(i), lw=10-i)

[tomhepz]

Thanks for the investigation, I never would've considered them all being there, only ontop of each other! I should've known rather than thinking the first was just special.

The simplest way I found to get this working with functools was:

%matplotlib ipympl
import matplotlib.pyplot as plt
from mpl_interactions import ipyplot as iplt
from mpl_interactions.controller import Controls
from functools import partial

ctrls = Controls(tau=(.5,1,100))
display(ctrls)
x = np.linspace(-5,5)
fig, ax = plt.subplots()

def to_plot(t,tau,i):
    return np.sin(tau*x)**i

with ctrls:
    for i in range(10):
        partial_fixed_parameter = partial(to_plot, i=i)
        iplt.plot(x, partial_fixed_parameter, lw=10-i)