Add paper example: Fully non parametric curve fit

[juanitorduz]

Hey 👋 ! It's me (again 🙈 ). I was reading the paper Bayesian additive regression trees for probabilistic programming and I was gladly surprised with the model presented in Section 4.4 Heteroscedasticity, in particular with the Code Block 6 where you model the mean and variance with a BART model of size 2. I think This example should be in the documentation.

The most interesting part would be to clarify the connection between the Y parameter (The response vector ) of the BART model with the likelihood. Note that in the code

with pm.Model() as model_marketing_full:
    w = pmb.BART("w", X, Y, m = 200, size =2 )
    y = pm.Normal("y", w[0], np.abs(w[1]), observed = Y)
    idata_marketing_full = pm.sample()

w[1] (i.e. the variance estimation) is estimated using Y. I am sometimes confused about the relationship between the Y parameter and the BART random variable (see for example #31). I think adding this example would benefit new users a lot.

I could draft a PR. Should it be added into the same notebook as in pymc-devs/pymc-examples#507 ?
Thanks 🙂

[aloctavodia]

Glad to see you here, again, haha.

I think it will be great to have this example in the docs (and many more! I will later open one or more issues about this). The example should be in a new notebook. Thanks, again for your help!

With the last release of BART (and pymc 5/pytensor) I think the model should be written like this.

with pm.Model() as model_marketing_full:
w = pmb.BART("w", X, Y, m = 200, shape= (200, 2))
y = pm.Normal("y", w[:,0], np.abs(w[:,1]), observed = Y)
idata_marketing_full = pm.sample()

[juanitorduz]

ok! I will give it a try and create a PR in the pymc-examples repo in the upcoming days :)

[juanitorduz]

I have started working on it and I just want to confirm I am not making a mistake. From what I see in the paper, the orinal matrix X has dimensions (n_obs, 2). Does it contain two copies of the youtube data cols:

df = pd.read_csv("../data/marketing.csv", sep=";", decimal=",")

n_obs = df.shape[0]

X = df["youtube"].to_numpy()
Y = df["sales"].to_numpy()

X_ext = np.stack([X, X], axis=1)

with pm.Model() as model_marketing_full:
    w = pmb.BART(name="w", X=X_ext, Y=Y, m = 200, shape=(n_obs, 2))
    y = pm.Normal(name="y", mu=w[:, 0], sigma = pm.math.abs(w[:, 1]), observed=Y)
    idata_marketing_full = pm.sample()

If so, the model takes like 15 min to fit. Is that expected?
Thanks :)

[aloctavodia]

X has to be 2D (as usual), but you should not repeat the "youtube" feature. With the syntax shape=(n_obs, 2) we are getting two independent BART variables, both variables will use the same info in X to define the tree partitions.

15 minutes seems to be a sign that something is wrong. I expect something around 1 or 2 minutes maybe

[juanitorduz]

Hey! Unfortunately, I have not been able to make it work.

The suggested solution:

X = df["youtube"].to_numpy().reshape(-1, 1)
Y = df["sales"].to_numpy()

with pm.Model() as model_marketing_full:
    w = pmb.BART(name="w", X=X, Y=Y, m=50, shape=(n_obs, 2))
    y = pm.Normal(name="y", mu=w[0, :], sigma=pm.math.abs(w[1, :]), observed=Y)


pm.model_to_graphviz(model=model_marketing_full)

outputs

This explains the long time sampling as it sees w as a 200 x 200 matrix, which is clearly wrong.

The closest attempt was:

with pm.Model() as model_marketing_full:
    w = pmb.BART(name="w", X=X, Y=Y, m=50, shape=(2, n_obs))
    y = pm.Normal(name="y", mu=w[0], sigma=pm.math.abs(w[1]), observed=Y)


pm.model_to_graphviz(model=model_marketing_full)

but I get an error at the end of the sampling

NotImplementedError                       Traceback (most recent call last)
File ~/opt/anaconda3/envs/pymc-examples-env/lib/python3.10/site-packages/pytensor/link/basic.py:111, in Container.__set__(self, value)
    109 try:
    110     # Use in-place filtering when/if possible
--> 111     self.storage[0] = self.type.filter_inplace(
    112         value, self.storage[0], **kwargs
    113     )
    114 except NotImplementedError:

File ~/opt/anaconda3/envs/pymc-examples-env/lib/python3.10/site-packages/pytensor/graph/type.py:130, in Type.filter_inplace(self, value, storage, strict, allow_downcast)
    111 """Return data or an appropriately wrapped/converted data by converting it in-place.
    112 
    113 This method allows one to reuse old allocated memory.  If this method
   (...)
    128 NotImplementedError
    129 """
--> 130 raise NotImplementedError()

NotImplementedError: 

During handling of the above exception, another exception occurred:

TypeError                                 Traceback (most recent call last)
Cell In[47], line 2
      1 with model_marketing_full:
----> 2     idata_marketing_full = pm.sample(draws=20)

File ~/opt/anaconda3/envs/pymc-examples-env/lib/python3.10/site-packages/pymc/sampling/mcmc.py:590, in sample(draws, step, init, n_init, initvals, trace, chains, cores, tune, progressbar, model, random_seed, discard_tuned_samples, compute_convergence_checks, callback, jitter_max_retries, return_inferencedata, keep_warning_stat, idata_kwargs, mp_ctx, **kwargs)
    587     n_draws = max(0, len(mtrace) - n_tune)
    589 if discard_tuned_samples:
--> 590     mtrace = mtrace[n_tune:]
    592 # save metadata in SamplerReport
    593 mtrace.report._n_tune = n_tune

File ~/opt/anaconda3/envs/pymc-examples-env/lib/python3.10/site-packages/pymc/backends/base.py:313, in MultiTrace.__getitem__(self, idx)
    311 def __getitem__(self, idx):
    312     if isinstance(idx, slice):
--> 313         return self._slice(idx)
    315     try:
    316         return self.point(int(idx))

File ~/opt/anaconda3/envs/pymc-examples-env/lib/python3.10/site-packages/pymc/backends/base.py:523, in MultiTrace._slice(self, slice)
    521 def _slice(self, slice):
    522     """Return a new MultiTrace object sliced according to `slice`."""
...
    252     )
    254 if self.filter_checks_isfinite and not np.all(np.isfinite(data)):
    255     raise ValueError("Non-finite elements not allowed")

TypeError: ("The type's shape ((2, 200)) is not compatible with the data's ((200, 2))", 'Container name "w"')

Also, I look for a unit-test but I could not find this type of case 🙈 . Could you give me a hint, please? Thanks!

[aloctavodia]

hmmm, are you using the last release? Let me try it myself and I will come back to you

[juanitorduz]

mmm ...not the last one ..let me try updating to the latest versions.

Last updated: Thu Feb 02 2023

Python implementation: CPython
Python version       : 3.10.9
IPython version      : 8.8.0

pytensor: 2.8.11

matplotlib: 3.5.2
numpy     : 1.23.5
pandas    : 1.5.3
pymc      : 5.0.1
pymc_bart : 0.3.0
arviz     : 0.14.0
sys       : 3.10.9 (main, Jan 11 2023, 09:18:20) [Clang 14.0.6 ]

Watermark: 2.3.1

[juanitorduz]

Well, the basic example

with pm.Model() as model_marketing_full:
    w = pmb.BART(name="w", X=X, Y=Y, m=50, shape=(2, n_obs))
    y = pm.Normal(name="y", mu=w[0, :], sigma=pm.math.abs(w[1, :]), observed=Y)
    idata_marketing_full = pm.sample()

Fails with a different error

---------------------------------------------------------------------------
RemoteTraceback                           Traceback (most recent call last)
RemoteTraceback: 
"""
Traceback (most recent call last):
  File "/Users/juanitorduz/opt/anaconda3/envs/pymc-examples-env/lib/python3.10/site-packages/pymc/sampling/parallel.py", line 122, in run
    self._start_loop()
  File "/Users/juanitorduz/opt/anaconda3/envs/pymc-examples-env/lib/python3.10/site-packages/pymc/sampling/parallel.py", line 174, in _start_loop
    point, stats = self._step_method.step(self._point)
  File "/Users/juanitorduz/opt/anaconda3/envs/pymc-examples-env/lib/python3.10/site-packages/pymc/step_methods/arraystep.py", line 102, in step
    apoint, stats = self.astep(q)
  File "/Users/juanitorduz/opt/anaconda3/envs/pymc-examples-env/lib/python3.10/site-packages/pymc_bart/pgbart.py", line 148, in astep
    self.sum_trees_noi = self.sum_trees - self.all_particles[tree_id].tree._predict()
  File "/Users/juanitorduz/opt/anaconda3/envs/pymc-examples-env/lib/python3.10/site-packages/pymc_bart/tree.py", line 111, in _predict
    output[leaf_node.idx_data_points] = leaf_node.value
IndexError: index 2 is out of bounds for axis 0 with size 2
"""

The above exception was the direct cause of the following exception:

IndexError                                Traceback (most recent call last)
IndexError: index 2 is out of bounds for axis 0 with size 2

The above exception was the direct cause of the following exception:

ParallelSamplingError                     Traceback (most recent call last)
Cell In[36], line 2
      1 with model_marketing_full:
----> 2     idata_marketing_full = pm.sample(draws=20)

File ~/opt/anaconda3/envs/pymc-examples-env/lib/python3.10/site-packages/pymc/sampling/mcmc.py:546, in sample(draws, step, init, n_init, initvals, trace, chains, cores, tune, progressbar, model, random_seed, discard_tuned_samples, compute_convergence_checks, callback, jitter_max_retries, return_inferencedata, keep_warning_stat, idata_kwargs, mp_ctx, **kwargs)
    544 _print_step_hierarchy(step)
    545 try:
--> 546     _mp_sample(**sample_args, **parallel_args)
    547 except pickle.PickleError:
    548     _log.warning("Could not pickle model, sampling singlethreaded.")

File ~/opt/anaconda3/envs/pymc-examples-env/lib/python3.10/site-packages/pymc/sampling/mcmc.py:914, in _mp_sample(draws, tune, step, chains, cores, random_seed, start, progressbar, traces, model, callback, mp_ctx, **kwargs)
    912 try:
    913     with sampler:
--> 914         for draw in sampler:
    915             strace = traces[draw.chain]
    916             strace.record(draw.point, draw.stats)

File ~/opt/anaconda3/envs/pymc-examples-env/lib/python3.10/site-packages/pymc/sampling/parallel.py:449, in ParallelSampler.__iter__(self)
...
--> 330     raise error from old_error
    331 elif msg[0] == "writing_done":
    332     proc._readable = True

ParallelSamplingError: Chain 0 failed with: index 2 is out of bounds for axis 0 with size 2

with

Python implementation: CPython
Python version       : 3.10.9
IPython version      : 8.8.0

pytensor: 2.9.1

pandas    : 1.5.3
pymc      : 5.0.2
arviz     : 0.14.0
matplotlib: 3.5.2
sys       : 3.10.9 (main, Jan 11 2023, 09:18:20) [Clang 14.0.6 ]
pymc_bart : 0.3.1
numpy     : 1.23.5

Watermark: 2.3.1

I hope this helps.

[juanitorduz]

It is working as expected with the latest release (0.3.2), see pymc-devs/pymc-examples#519