[Feature Request] Is it possible to store/save a starry model?

[christinahedges]

I have generated a load of starry models by fitting to data, and I would like to save them so I can plot the model in future notebooks. I don't think there is a feature to do this, and it's not possible to use pickle on a starry.Map object.

I think the best way for me to do this currently is to save all the map amplitudes and dictionaries of the orbital parameters, and then re-initialize the model when I want to use it in a new notebook. I can totally do this for now, but if there is an easy way to store the models and then re-open them I'd love to use it!

Thanks C

[taylorbell57]

Upvoting this to allow for the combination of emcee and multiprocessing which requires a pickleable object or a new star and planet object to be created at each iteration which is a huge overhead time cost (2 ms for v0.3 or 8 ms for v1.0; that's 100-400% of my current likelihood computation time without using starry). On a related note, v1.0 takes an order of magnitude longer just to update the map (2 us for v0.3 versus 3 ms for v1.0 for a planet that has a ydeg=1). The "Pre-computing some matrices... Done." being repeated a million times when initialising a planet in each MCMC step is also pretty rough and is something I've commented out locally.

[dfm]

@taylorbell57: You shouldn't need to re-create the Star object in every call. Instead, you'll want to create the theano function that takes parameters as input and returns the starry model. @rodluger: can explain how to do that exactly, but I would say that this is actually a separate issue.

[dfm]

(also you should use PyMC3 instead of emcee... it'll be much faster!)

[rodluger]

A few comments:

• @dfm is right that you should never re-instantiate a starry map repeatedly, since there are a ton of things that are purposefully pre-computed at the beginning to save time in calls to flux.
• Can you show me a MWE of your multiprocessing setup? I don't understand why you need a new map object at each iteration. @dfm, perhaps there's a way to instantiate a single map object once for each thread? Or am I missing something?
• You can set starry.config.quiet = True to suppress that informational message.
• Agreed with @dfm that you should consider PyMC3. The new version of starry (1.0.0) is specifically designed to work with PyMC3 sampling. Because of all the theano bindings, it's a bit slower than the first version if you're just trying to evaluate the flux. But if you want to do gradient-based inference, it can be orders of magnitude faster.

All that said, @dfm's suggestion of a compiled theano function is a good one. Here's an example. First, let's time your current approach with a mock example:

import starry
import numpy as np
import time

# Greedy mode
starry.config.lazy = False
starry.config.quiet = True

# Instantiate the map, set the coeffs, compute the flux
def slow_flux(coeffs, theta):
    map = starry.Map(10)
    map[1:, :] = coeffs
    return map.flux(theta=theta)

# Function args
coeffs = np.random.randn(120)
theta = np.linspace(0, 360, 1000)

# Time the function
nruns = 25
tstart = time.time()
for k in range(nruns):
    slow_flux(coeffs, theta)
elapsed = (time.time() - tstart) / nruns
print("Call took {:.3f} seconds.".format(elapsed))

>>> Call took 0.128 seconds.

Now, instead of running starry in greedy mode, let's run it in lazy mode and compile our custom function ourselves. This gets rid of all the Python-side overhead, including what you mentioned about how much slower it is to set coefficients:

import starry
import numpy as np
import theano
import theano.tensor as tt
import time

# Lazy mode
starry.config.lazy = True
starry.config.quiet = True

# Instantiate the map, set the coeffs, compute the flux
def slow_flux(coeffs, theta):
    map = starry.Map(10)
    map[1:, :] = coeffs
    return map.flux(theta=theta)

# Compile `slow_flux` into a theano function
arg1 = tt.dvector()
arg2 = tt.dvector()
fast_flux = theano.function([arg1, arg2], slow_flux(arg1, arg2))

# Function args
coeffs = np.random.randn(120)
theta = np.linspace(0, 360, 1000)

# Time the function
nruns = 25
tstart = time.time()
for k in range(nruns):
    fast_flux(coeffs, theta)
elapsed = (time.time() - tstart) / nruns
print("Call took {:.3f} seconds.".format(elapsed))

Call took 0.002 seconds.

About 3 orders of magnitude faster, about the difference you were finding. Let me know if you have questions. I'll try to write up a tutorial on all this, since I think it will be useful for others.

[taylorbell57]

Thanks for the quick feedback all! I've heard a couple times that I should switch to PyMC3, so I gave it a few tries yesterday but had a lot of trouble fitting a full-orbit Spitzer phasecurve using a PLD decorrelation technique. It optimises incredibly fast to a good solution using exoplanet.optimize() (tiny fraction of the time it'd take without starry+PyMC3), but then I cannot get the NUTS sampler to converge. I'll come back around to PyMC3 another day soon, but I need to make some progress on the short term (want to update some results for a AAS iPoster).

As for my comment about re-making a star+planet+system for every MCMC step, my concern was that it could be possible for a race condition to occur if I were to use a global variable and access that within the log-likelihood function (without passing it into the function). The thought was there could be one walker that updates the system before another walker tries to compute the lightcurve and then chaos would ensue. I've tried pretty hard to get this to happen using starry v0.3 though (by adding an long pause between setting the planet's radius and reading its value), and I've not been able to get a race condition to occur (tested by making sure planet.r is equal to the value that I input). I've attached my minimal example (as a .txt file because I can't upload .py files) that I thought should cause a race condition below. Perhaps python/emcee/multiprocessing is smarter than I'd thought and one of them is actually locking global variables in some intelligent way? I'll go on using the global variable for now and will circle back around to getting PyMC3 to work in a week or two. Beyond accessing the global system variable, I can't think of a way of passing the system into the log-likelihood function since it is required that everything passed in through the emcee.EnsembleSampler args argument be pickleable (to avoid race conditions I assume).

RaceCondition.txt

[rodluger]

Hey @taylorbell57, theano should in principle support pickling:

http://deeplearning.net/software/theano/tutorial/loading_and_saving.html

Can you try compiling the fast_flux function as above and passing it as an argument to EnsembleSampler? If that works, I think that's your best bet: just compile a function that returns system.lightcurve and take that as an argument to your lnprob function.

[rodluger]

If that doesn't work, I would try using the compiled theano function as global variable.

[rodluger]

Closing this due to inactivity. Feel free to reopen.

[jvines]

Hi @rodluger , I don't mean to resurrect a dead issue, but I find myself in the same position, I'm trying to use emcee with multiple threads, cannot move over to PyMC3, and I am encountering the same Map object is not pickable issue. When I try to run the code you put for compiling a theano function it fails for me saying it's trying to set an array element with a sequence.

What would be the best way to get this done for, say, modeling an exoplanet transit?

As of now, my attempt was to create a system object and update the parameters in each MCMC step for the fit Also the link you provided is quite dead, or at least I cannot open it

Thanks for the help!

[jvines]

Well, after a clean install I got your snippet there working, and I coded up this:

def LC_model(theta, lds, time):
    """Evaluate light curve model.

    Parameters
    ----------
    theta: array_like
        Array with the orbital parameters.
    theta_ld: array_like
        Array with the kipping limb darkening coefficients.
    time: array_like
        The time array.
    model: starry.kepler.System
        The starry System object.
    law: str
        The limb darkening law.

    Returns
    -------
    flux: array
        The transit flux.
    """
    p = theta[0]
    e = theta[1]
    w = theta[2]
    tp = theta[3]
    rr = theta[4]
    b = theta[5]
    a = theta[6]

    inc = np.arccos(b / a) * 180 / np.pi
    star = starry.Primary(starry.Map(udeg=2))
    star.map[1] = lds[0]
    star.map[2] = lds[1]
    planet = starry.kepler.Secondary(starry.Map(amp=0),
                                     m=0, porb=p, inc=inc, ecc=e,
                                     r=rr, w=w, t0=tp)
    model = starry.System(star, planet)

    flux = model.flux(time)
    return flux

# Compile `slow_flux` into a theano function
starry.config.lazy = True
theta = tt.dvector()
lds = tt.dvector()
time = tt.dvector()
fast_flux = theano.function(
    [theta, lds, time],
    LC_model(theta, lds, time))

I'm unsure, however, if that's the correct approach here, so any guidance would be appreciated.