Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

QobjEvo tries to compress the Qobj elements #36

Open
BoxiLi opened this issue Mar 2, 2024 · 1 comment
Open

QobjEvo tries to compress the Qobj elements #36

BoxiLi opened this issue Mar 2, 2024 · 1 comment

Comments

@BoxiLi
Copy link
Member

BoxiLi commented Mar 2, 2024

jax.jit fails when I try to include more than one time-dependent elements in QobjEvo. This is because QobEvo is trying to compare the Qobjs in the elements and compress them if possible. For my directly calling of QobjEvo, I can set the argument compress=False, but QobjEvo is called also within sesolve:

H = QobjEvo(H, args=args, tlist=tlist)

Example

@jax.jit
def pulse_fun(t, T):
    return t/T

@jax.jit
def tmp(t, T):
    # Demonstrate a simple case with more than one time-dependent elements.
    H = qutip.destroy(3) * qutip.coefficient(pulse_fun, args={"T": T}) + qutip.create(3) * qutip.coefficient(pulse_fun, args={"T": T})
    result = qutip.sesolve(
    H,
    qutip.basis(3),
    [0, 2.],
    options={
        "method": "diffrax",
        "dt0":0.01,
        "progress_bar":""})
    return result.final_state.data_as("JaxArray")

tmp(1., 10.)
@Ericgig
Copy link
Member

Ericgig commented Mar 4, 2024

I think that all solver function calls sesolve, mesolve, etc. do not work.
The initialization step is not jax friendly.
Returning Qobj is also broken and returning Result is not supported.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

No branches or pull requests

2 participants