Sample packing fails

[stephtdouglas]

When trying to pack samples, I am running into a bug within samples.py. Here's a MWE:

import astropy.units as u
import thejoker as tj

t = [0., 49.452, 95.393, 127.587, 190.408]
rv = [38.77, 39.70, 37.45, 38.31, 38.31] * u.km/u.s
err = [0.184, 0.261, 0.112, 0.155, 0.223] * u.km/u.s
data = tj.RVData(t=t, rv=rv, rv_err=err)

prior = tj.JokerPrior.default(
    P_min=2*u.day, P_max=1e3*u.day,
    sigma_K0=3*u.km/u.s,
    sigma_v=38*u.km/u.s
    )

prior_samples = prior.sample(size=250_000,
                             random_state=rnd)

packed_samples, units = prior_samples.pack(u.km/u.s)

And I get the error

AttributeError                            Traceback (most recent call last)
Cell In [13], line 1
----> 1 packed_samples, units = prior_samples.pack(u.km/u.s)

File ~/anaconda3/envs/thejoker/lib/python3.8/site-packages/thejoker/samples.py:368, in JokerSamples.pack(self, units, names, nonlinear_only)
    366     units = dict()
    367 out_units = OrderedDict()
--> 368 [units.setdefault(k, v) for k, v in _nonlinear_internal_units.items()]
    370 if names is None:
    371     if nonlinear_only:

File ~/anaconda3/envs/thejoker/lib/python3.8/site-packages/thejoker/samples.py:368, in <listcomp>(.0)
    366     units = dict()
    367 out_units = OrderedDict()
--> 368 [units.setdefault(k, v) for k, v in _nonlinear_internal_units.items()]
    370 if names is None:
    371     if nonlinear_only:

AttributeError: 'CompositeUnit' object has no attribute 'setdefault'

For context, I'm doing this because I'm still trying to do some runs with TwoFace for a proposal, and that wants to do the prior sampling in batches rather than all at once. Hence I'm doing this and not using the built-in writer.

[adrn]

Hm! Try doing the .pack() without a unit passed in?

packed_samples, units = prior_samples.pack()

[stephtdouglas]

That worked! The old version must have just only wanted units for velocity (or jitter?), instead of the option to submit multiple units via the ordered dict. Thanks!