Fitting with two planets using rebound throws a fatal error

[bailer-jones]

Summary
I think a 2D array needs to be recast as a 3D one when using rebound for the forward model.

The issue
I am using orbitize (version 2.2.2 in python 3.9.18 on MacOS 13.3.1) to fit a system of two planets with RA,Dec data using the rebound feature. When I do this the sampler (using MCMC) gives the following error:

 File "/opt/anaconda3/envs/rebound/lib/python3.9/site-packages/orbitize/system.py", line 540, in compute_model
 model[self.radec[body_num], 0] = raoff[self.radec[body_num], body_num, :]  # N_epochs x N_bodies x N_orbits
 IndexError: too many indices for array: array is 2-dimensional, but 3 were indexed

To Reproduce
sys = system.System(num_secondary_bodies=2,
use_rebound=True,
fit_secondary_mass=True,
data_table=dat,
tau_ref_epoch=timeref,
stellar_or_system_mass=mass_star/msun, mass_err=0,
plx=1.0, plx_err=0.0)
lab = sys.param_idx
mcmcpar = {
'num_temps': 1,
'num_walkers': 20,
'num_threads': 1,
'num_steps': 1000,
'num_burnin': 0,
'num_thin': 1
}
mySampler = sampler.MCMC(sys, num_temps=mcmcpar['num_temps'], num_walkers=mcmcpar['num_walkers'])
blah = mySampler.run_sampler(total_orbits=mcmcpar['num_steps']*mcmcpar['num_walkers'],
burn_steps=mcmcpar['num_burnin'], thin=mcmcpar['num_thin'])

In my case dat comprises 50 simulated astrometric measurements on just one planet, and various priors are set. but I don't think the exact details here are relevant.

Apparent cause
I think the issue is in the system.py file in the compute_all_orbits() function.
(Note that when this is called by the sampler it reports having n_orbits=1 inside.) Line 350, which is only called when we are opting to use rebound, is

raoff, deoff, vz = nbody.calc_orbit(epochs, sma, ecc, inc, argp, lan, tau, plx, mtot, tau_ref_epoch=self.tau_ref_epoch, m_pl=m_pl, output_star=True)

This returns raoff, deoff, (and vz) as 2D arrays of size N_epochs x N_bodies. These are passed back to line 505 of compute_model():

raoff, decoff, vz = self.compute_all_orbits(standard_params_arr, comp_rebound=True)

However, when this is used further down in that function on line 540:

model[self.radec[body_num], 0] = raoff[self.radec[body_num], body_num, :]  # N_epochs x N_bodies x N_orbits

this is expecting a 3D array, not a 2D one. This is where the error occurs.

Possible solution
I think it is correct that nbody.calc_orbit() in compute_all_orbits() is returning raoff and deoff for just a single orbit (n_orbit=1). But raoff (and deoff) then need to be cast as a 3D array to be compatible with its use on line 540. The following, inserted immediately after line 359, prevents the error and seems give the correct behaviour:

if(n_orbits==1):
     raoff = np.reshape(raoff, (raoff.shape[0],raoff.shape[1],1))
     deoff = np.reshape(deoff, (deoff.shape[0],deoff.shape[1],1))
     vz = np.reshape(vz, (vz.shape[0],vz.shape[1],1))

although I've not tested the code with vz data.

[sblunt]

Fixed in #360