More regression with Bayes¶

Introduction¶

Still more regression examples from Bayesian Analysis with Python (Edition 3). This time, a dataset with a bad case of Heteroskedasticity. We have a dataset where we measure baby length: as babies get longer, so the variance grows (the gap between between the longest and shortest baby gets bigger as they get older)

We develop a model that tries to capture this behaviour.

I also learn a lot about PyMC extrapolation / prediction, and how to frame the model specification to accommodate this

Implementation¶

Imports¶

In [4]:

import pathlib

import arviz as az
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import preliz as pz
import pymc as pm
import xarray as xa

Load dataset¶

In [5]:

path = pathlib.Path("data/babies.csv")
df = pd.read_csv(path)
print(df.head(2))

   month     length
0      0  49.192857
1      0  51.221429

Visualize dataset¶

We can clearly see that the spread of baby lengths get bigger as they get older

In [6]:

plt.plot(df["month"], df["length"], "+")

Out[6]:

[<matplotlib.lines.Line2D at 0x22443cf1be0>]

No description has been provided for this image

Build and run linear model¶

In order to contrast with a more nuanced model, we first build the simplest linear model: constant variance, length linearly dependent on age.

Some points to note:

We use pm.Data() call to store our dataset. Later, we will change the month Data collection, so the model (PyMC) will recalculate predictions for the new month values.
We explicitly set the length of the vector of length predictions (l_pred) as being the length of the month dataset (shape=month.shape)
We explicitly set the dimension name for l_pred to be the same as the automatically generated name for month. This will make subsequent plotting and processing easier
l_err0 is our linear model, to which we add Normal noise to get the l_pred prediction
We have not normalized our data, and have chosen prior values for the linear model that make sense (by eye)

In [7]:

with pm.Model() as model_1:
    # store months and length in Data containers
    # later we will want to change months to get
    # out-of-sample predictions
    # the parameter shape=month.shape will allow us to change the months array,
    # to get predictions beyond the end of the data

    month = pm.Data("month", df.month)
    length = pm.Data("length", df.length)

    # set by eye
    intercept = pm.Normal("intercept", mu=55, sigma=100)

    # set by eye
    slope = pm.Normal("slope", mu=1.4, sigma=10)

    # set by eye
    error_sd = pm.HalfCauchy("error_sd", beta=10)

    # l_err0 is the line with no noise added
    l_err0 = pm.Deterministic("l_err0", intercept + slope * month)

    # l_pred is the linear model with added noise
    l_pred = pm.Normal(
        "l_pred",
        mu=l_err0,
        sigma=error_sd,
        observed=length,
        shape=month.shape,
        dims="month_dim_0",
    )

    idata_1 = pm.sample()
# end with

C:\Users\donrc\anaconda3\envs\bayes_env\Lib\site-packages\threadpoolctl.py:1226: RuntimeWarning: 
Found Intel OpenMP ('libiomp') and LLVM OpenMP ('libomp') loaded at
the same time. Both libraries are known to be incompatible and this
can cause random crashes or deadlocks on Linux when loaded in the
same Python program.
Using threadpoolctl may cause crashes or deadlocks. For more
information and possible workarounds, please see
    https://github.com/joblib/threadpoolctl/blob/master/multiple_openmp.md

  warnings.warn(msg, RuntimeWarning)
NUTS[nutpie]: [intercept, slope, error_sd]

Output()

Visualize the model (look OK)

In [8]:

pm.model_to_graphviz(model_1)

Out[8]:

Look for diagnostic problems (all OK)

In [9]:

problems = az.diagnose(idata_1)

Divergences
No divergent transitions found.

E-BFMI
E-BFMI satisfactory for all chains.

ESS
Effective sample size satisfactory for all parameters.

R-hat
R-hat values satisfactory for all parameters.

Processing complete, no problems detected.

Visualize the posterior distributions for the key parameters of the model

In [10]:

az.plot_dist(idata_1, var_names=["intercept", "slope", "error_sd"])

Out[10]:

<arviz_plots.plot_collection.PlotCollection at 0x224499ff4d0>

Examine the data structure resulting from our model run. This will be extended below to support more analysis. Note the absence of any Prior data

In [11]:

idata_1

Out[11]:

<xarray.DataTree>
Group: /
├── Group: /posterior
│       Dimensions:       (chain: 4, draw: 1000, l_err0_dim_0: 796)
│       Coordinates:
│         * chain         (chain) int64 32B 0 1 2 3
│         * draw          (draw) int64 8kB 0 1 2 3 4 5 6 ... 993 994 995 996 997 998 999
│         * l_err0_dim_0  (l_err0_dim_0) int64 6kB 0 1 2 3 4 5 ... 791 792 793 794 795
│       Data variables:
│           intercept     (chain, draw) float64 32kB 56.17 56.2 55.8 ... 55.51 55.93
│           slope         (chain, draw) float64 32kB 1.438 1.425 1.473 ... 1.476 1.445
│           error_sd      (chain, draw) float64 32kB 4.341 4.013 4.208 ... 4.288 4.187
│           l_err0        (chain, draw, l_err0_dim_0) float64 25MB 56.17 56.17 ... 90.61
│       Attributes:
│           created_at:                 2026-10-24T00:09:57.316014+00:00
│           creation_library:           ArviZ
│           creation_library_version:   1.3.0
│           creation_library_language:  Python
│           sample_dims:                ['chain', 'draw']
│           inference_library:          nutpie
│           inference_library_version:  0.16.11
│           sampling_time:              0.2003464698791504
│           tuning_steps:               400
├── Group: /sample_stats
│       Dimensions:                   (chain: 4, draw: 1000)
│       Coordinates:
│         * chain                     (chain) int64 32B 0 1 2 3
│         * draw                      (draw) int64 8kB 0 1 2 3 4 ... 995 996 997 998 999
│       Data variables: (12/20)
│           depth                     (chain, draw) uint64 32kB 2 2 2 2 2 ... 2 3 3 3 3
│           maxdepth_reached          (chain, draw) bool 4kB False False ... False False
│           step_size                 (chain, draw) float64 32kB 0.7544 ... 0.7452
│           transformation_update_id  (chain, draw) int64 32kB 0 0 0 0 0 0 ... 0 0 0 0 0
│           step_size_bar             (chain, draw) float64 32kB 0.7534 0.7534 ... 0.715
│           mean_tree_accept          (chain, draw) float64 32kB 0.9689 0.7949 ... 1.0
│           ...                        ...
│           fisher_distance           (chain, draw) float64 32kB 0.03575 2.186 ... 2.665
│           transformation_index      (chain, draw) int64 32kB 338 338 338 ... 338 338
│           diverging                 (chain, draw) bool 4kB False False ... False False
│           divergence_draw           (chain, draw) uint64 32kB 0 0 0 0 0 ... 0 0 0 0 0
│           divergence_message        (chain, draw) object 32kB None None ... None None
│           divergence_energy_error   (chain, draw) float64 32kB nan nan nan ... nan nan
│       Attributes:
│           created_at:                  2026-10-24T00:09:57.296718+00:00
│           creation_library:            ArviZ
│           creation_library_version:    1.3.0
│           creation_library_language:   Python
│           sample_dims:                 ['chain', 'draw']
│           inference_library:           nutpie
│           inference_library_version:   0.16.11
│           inference_library_settings:  {"sampler": "nuts", "adaptation": "diag", "s...
├── Group: /constant_data
│       Dimensions:      (month_dim_0: 796)
│       Coordinates:
│         * month_dim_0  (month_dim_0) int64 6kB 0 1 2 3 4 5 ... 790 791 792 793 794 795
│       Data variables:
│           month        (month_dim_0) int32 3kB 0 0 0 0 0 0 0 ... 24 24 24 24 24 24 24
│       Attributes:
│           created_at:                 2026-10-24T00:09:57.310432+00:00
│           creation_library:           ArviZ
│           creation_library_version:   1.3.0
│           creation_library_language:  Python
│           inference_library:          pymc
│           inference_library_version:  6.3.1
│           sample_dims:                []
└── Group: /observed_data
        Dimensions:      (month_dim_0: 796)
        Coordinates:
          * month_dim_0  (month_dim_0) int64 6kB 0 1 2 3 4 5 ... 790 791 792 793 794 795
        Data variables:
            l_pred       (month_dim_0) float64 6kB 49.19 51.22 51.22 ... 97.35 94.19
        Attributes:
            created_at:                 2026-10-24T00:09:57.314062+00:00
            creation_library:           ArviZ
            creation_library_version:   1.3.0
            creation_library_language:  Python
            inference_library:          pymc
            inference_library_version:  6.3.1
            sample_dims:                []

Get the key parameters in pandas DataFrame form, and then extract them from the pandas DataFrame

In [12]:

results = az.summary(idata_1, var_names=["intercept", "slope"])
results

Out[12]:

	mean	sd	eti89_lb	eti89_ub	ess_bulk	ess_tail	r_hat	mcse_mean	mcse_sd
intercept	56.15	0.272	56	57	1349	1681	1.00	0.0074	0.0049
slope	1.44	0.02	1.4	1.5	1332	1518	1.00	0.00055	0.00036

In [13]:

line_a = results.loc["intercept", "mean"]
line_b = results.loc["slope", "mean"]
line_a, line_b

Out[13]:

(np.float64(56.14674251118621), np.float64(1.440154543991303))

Prediction of in-sample values¶

Initial DataTree from model run¶

In [14]:

idata_1

Digression: Get predictions from Priors¶

This shows how to set up the data structures to support predictions using the Priors

Not used below, but shown here for completeness

In [15]:

with model_1:
    prior_samples = pm.sample_prior_predictive()
# end with

Sampling: [error_sd, intercept, l_pred, slope]

New data structures created by the sample_prior_predictive() call

In [16]:

prior_samples

Update our model DataTree with Prior sampling, and display DataTree structure

In [17]:

idata_1.update(prior_samples)
idata_1

End of digression¶

Get predictions from one chain of Posteriors¶

We want to see the spread of predicted length values for three positions on the x axis

First we have to create the needed data structures. Note that we are not considering new month values: we use all existing month values from the input dataset (a given month value: say, 10) will be repeated many times in the input dataset)

In [18]:

# Sample posterior predictive and append it automatically
with model_1:
    pm.sample_posterior_predictive(idata_1, extend_inferencedata=True)
# end with

Sampling: [l_pred]

Output()

Display DataTree structure now: note the prior_predictive, and posterior_predictive data structures

In [19]:

idata_1

Extract data for plotting¶

We want the spread of sampled length values as the left end , middle, and right end of the input dataset

We start by looking at just one of our chains of sampling (chain=0)

Extract from DataTree by:

index into group "posterior_predicative"
the variable we want is predicted length ("l_pred"). PyMC would normally an index "l_pred_dim_0" for each prediction position, but we have overridden this to month_dim_0 , running from 0 to 795 (length on input data array = 796); so ask for:
- first (=0), middle (=200), and last (=795) position of x axis, using "month_dim_0" index
- first chain (=0)
- all draws (0 to 99)
get the "l_pred" data variable [ "l_pred" ]
get the values as numpy array (100 data values at each of three positions), using .values

In [20]:

pred_vals = (
    idata_1["posterior_predictive"]
    .sel({"month_dim_0": [0, 400, 795], "chain": 0, "draw": slice(0, 99)})["l_pred"]
    .values
)
pred_vals.shape

Out[20]:

(100, 3)

Plot results: we see there is a constant variance

In [21]:

labels = ["first", "middle", "last"]
for i in range(3):
    _ = plt.hist(pred_vals[:, i], histtype="step", label=labels[i])
# end for
plt.legend()
plt.title("Distribution of predicted length at position indicated")

Out[21]:

Text(0.5, 1.0, 'Distribution of predicted length at position indicated')

Get predictions from mean of all chains of Posteriors¶

In the same way, we now average across all chains

In [22]:

pred_vals = (
    idata_1["posterior_predictive"]
    .sel({"month_dim_0": [0, 400, 795], "draw": slice(0, 99)})
    .mean(dim="chain")["l_pred"]
    .values
)
type(pred_vals)

Out[22]:

numpy.ndarray

Plot the results. We expect with 4 chains, to have $1/2$ the standard deviation.

In [23]:

labels = ["first", "middle", "last"]
for i in range(3):
    _ = plt.hist(pred_vals[:, i], histtype="step", label=labels[i])
# end for
plt.legend()
plt.title("Distribution of predicted length at position indicated")

Out[23]:

Text(0.5, 1.0, 'Distribution of predicted length at position indicated')

Conclusions¶

We see that the variance is the same at all positions. What our model specified, but not what we want (see below)

Prediction of out-of-sample values¶

We change our months array, and ask for posterior predictions again

In [24]:

with model_1:
    # set out of sample months
    oos_months = [i for i in range(31)]
    pm.set_data(
        {"month": oos_months},
    )

    post_pred = pm.sample_posterior_predictive(
        idata_1,
    )

# end with

Sampling: [l_pred]

Output()

Get the predicted lengths, averaged across chains and draws.

In [25]:

l_preds = post_pred["posterior_predictive"]["l_pred"].mean(("draw", "chain"))

Plot the:

raw data (blue cross)
values from out-of-sample predictions (red line)
linear model ($y = a + b \times x$) (green points)

In [26]:

plt.plot(
    df["month"],
    df["length"],
    "+",
    label="raw data",
)
plt.plot(
    oos_months,
    l_preds,
    "-r",
    label="predicted",
)
plt.plot(
    df["month"],
    line_a + line_b * df["month"],
    "og",
    label="linear fit",
)
plt.legend()

Out[26]:

<matplotlib.legend.Legend at 0x2244a4dd6a0>

Digression: using plot_lm¶

Default plot

In [27]:

idata_lm = idata_1.copy()

pc = az.plot_lm(
    idata_lm,
)

Nicer looking plot

In [28]:

# do a linear model plot
# set Confidence Interval band  0.3 transparent
# set the scatter plot points to green, and make partially transparent
# set the figure size
# set the confidence interval type to HDI (High Density Interval)
# set the HDI band to contain 95% of estimated prob. densit
pc = az.plot_lm(
    idata_1,
    visuals={
        "ci_band": {"alpha": 0.3},
        "observed_scatter": {"color": "C2", "alpha": 0.6},
    },
    figure_kwargs={"figsize": (10, 6)},
    ci_kind="hdi",
    ci_prob=0.95,
)

# add a pale grid
pc.map(
    az.visuals.grid,
    axis="both",
    color="lightgray",
)

# label x axis
pc.map(
    az.visuals.labelled_x,
    text="Month",
)
# label y axis
pc.map(
    az.visuals.labelled_y,
    text="Length",
)


_ = pc.viz.figure.item().axes[0].set_title("Baby growth (linear model)")

End of digression¶

sqrt model, constant variance¶

This essentially the same model as before, except we use the sqrt of month (intercept + slope * month**0.5)

In [29]:

with pm.Model() as model_1:

    month = pm.Data("month", df.month)
    length = pm.Data("length", df.length)

    intercept = pm.Normal("intercept", mu=50, sigma=100)
    slope = pm.Normal("slope", mu=20, sigma=20)
    error_sd = pm.HalfCauchy("error_sd", beta=10)
    l_err0 = pm.Deterministic("l_err0", intercept + slope * month**0.5)
    l_pred = pm.Normal(
        "l_pred",
        mu=l_err0,
        sigma=error_sd,
        observed=length,
        shape=month.shape,
        dims="month_dim_0",
    )

    idata_1 = pm.sample()
# end with

C:\Users\donrc\anaconda3\envs\bayes_env\Lib\site-packages\threadpoolctl.py:1226: RuntimeWarning: 
Found Intel OpenMP ('libiomp') and LLVM OpenMP ('libomp') loaded at
the same time. Both libraries are known to be incompatible and this
can cause random crashes or deadlocks on Linux when loaded in the
same Python program.
Using threadpoolctl may cause crashes or deadlocks. For more
information and possible workarounds, please see
    https://github.com/joblib/threadpoolctl/blob/master/multiple_openmp.md

  warnings.warn(msg, RuntimeWarning)
NUTS[nutpie]: [intercept, slope, error_sd]

Output()

Visualize model

In [30]:

pm.model_to_graphviz(model_1)

Out[30]:

Check for model issues (None)

In [31]:

problems = az.diagnose(idata_1)

Divergences
No divergent transitions found.

E-BFMI
E-BFMI satisfactory for all chains.

ESS
Effective sample size satisfactory for all parameters.

R-hat
R-hat values satisfactory for all parameters.

Processing complete, no problems detected.

In [32]:

az.plot_dist(idata_1, var_names=["intercept", "slope", "error_sd"])

Out[32]:

<arviz_plots.plot_collection.PlotCollection at 0x2244730d6e0>

Get model parameters for plotting

In [33]:

results = az.summary(idata_1, var_names=["intercept", "slope"])
results

Out[33]:

	mean	sd	eti89_lb	eti89_ub	ess_bulk	ess_tail	r_hat	mcse_mean	mcse_sd
intercept	47.46	0.33	47	48	941	1232	1.00	0.011	0.0075
slope	8.112	0.099	8	8.3	964	1238	1.00	0.0032	0.0024

In [34]:

line_a = results.loc["intercept", "mean"]
line_b = results.loc["slope", "mean"]
line_a, line_b

Out[34]:

(np.float64(47.45810407495688), np.float64(8.112174504009483))

Plot results - fit looks better

In [35]:

plt.plot(df["month"], df["length"], "+", label="data points")
plt.plot(df["month"], line_a + line_b * df["month"] ** 0.5, "-r", label="sqrt fit")
plt.legend()

Out[35]:

<matplotlib.legend.Legend at 0x2244c339be0>

Get predictions from Posteriors¶

As before, we want to see the spread of predicted length values for three positions on the x axis

First set up required data structures

In [36]:

# Sample posterior predictive and append it automatically
with model_1:
    pm.sample_posterior_predictive(idata_1, extend_inferencedata=True)
# end with

Sampling: [l_pred]

Output()

Get the data to plot

In [37]:

pred_vals = (
    idata_1["posterior_predictive"]
    .sel({"month_dim_0": [0, 400, 795], "draw": slice(0, 99)})
    .mean(dim="chain")["l_pred"]
    .values
)
type(pred_vals)

Out[37]:

numpy.ndarray

Do the plotting

In [38]:

labels = ["first", "middle", "last"]
for i in range(3):
    _ = plt.hist(pred_vals[:, i], histtype="step", label=labels[i])
# end for
plt.legend()
plt.title("Distribution of predicted length at position indicated")

Out[38]:

Text(0.5, 1.0, 'Distribution of predicted length at position indicated')

Still a constant variance

Final model¶

In our final model, we have length with a $sqrt (month)$ dependence, and noise standard deviation to grow linearly with month (error_0 + error_slope * month)

In [39]:

with pm.Model() as model_1:

    month = pm.Data("month", df.month)
    length = pm.Data("length", df.length)

    intercept = pm.Normal("intercept", mu=50, sigma=100)
    slope = pm.Normal("slope", mu=20, sigma=20)

    error_0 = pm.HalfNormal("error_0", sigma=10)
    error_slope = pm.HalfNormal("error_slope", sigma=10)

    l_err0 = pm.Deterministic("l_err0", intercept + slope * month**0.5)
    error_sd = pm.Deterministic("error_sd", error_0 + error_slope * month)

    l_pred = pm.Normal(
        "l_pred",
        mu=l_err0,
        sigma=error_sd,
        observed=length,
        shape=month.shape,
        dims="month_dim_0",
    )

    idata_1 = pm.sample()
# end with

C:\Users\donrc\anaconda3\envs\bayes_env\Lib\site-packages\threadpoolctl.py:1226: RuntimeWarning: 
Found Intel OpenMP ('libiomp') and LLVM OpenMP ('libomp') loaded at
the same time. Both libraries are known to be incompatible and this
can cause random crashes or deadlocks on Linux when loaded in the
same Python program.
Using threadpoolctl may cause crashes or deadlocks. For more
information and possible workarounds, please see
    https://github.com/joblib/threadpoolctl/blob/master/multiple_openmp.md

  warnings.warn(msg, RuntimeWarning)
NUTS[nutpie]: [intercept, slope, error_0, error_slope]

Output()

Visualize model

In [40]:

pm.model_to_graphviz(model_1)

Out[40]:

Check for issues (None)

In [41]:

problems = az.diagnose(idata_1)

Divergences
No divergent transitions found.

E-BFMI
E-BFMI satisfactory for all chains.

ESS
Effective sample size satisfactory for all parameters.

R-hat
R-hat values satisfactory for all parameters.

Processing complete, no problems detected.

Get model parameters

In [42]:

az.summary(idata_1, var_names=["intercept", "slope"])

Out[42]:

	mean	sd	eti89_lb	eti89_ub	ess_bulk	ess_tail	r_hat	mcse_mean	mcse_sd
intercept	48.05	0.234	48	48	1293	1580	1.00	0.0065	0.0044
slope	7.867	0.085	7.7	8	1265	1603	1.00	0.0024	0.0016

In [43]:

line_a = results.loc["intercept", "mean"]
line_b = results.loc["slope", "mean"]
line_a, line_b

Out[43]:

(np.float64(47.45810407495688), np.float64(8.112174504009483))

Visualize fit (looks OK)

In [44]:

plt.plot(df["month"], df["length"], "+", label="data points")
plt.plot(df["month"], line_a + line_b * df["month"] ** 0.5, "-r", label="sqrt fit")
plt.legend()

Out[44]:

<matplotlib.legend.Legend at 0x2244bb42900>

Check distributions of key parameters. We can see thet the length standard deviation grows by $0.11$ each month

In [45]:

az.plot_dist(
    idata_1,
    var_names=[
        "intercept",
        "slope",
    ],
)

Out[45]:

<arviz_plots.plot_collection.PlotCollection at 0x22446ddba80>

In [46]:

az.plot_dist(
    idata_1,
    var_names=[
        "error_0",
        "error_slope",
    ],
)

Out[46]:

<arviz_plots.plot_collection.PlotCollection at 0x2244744e7b0>

Check variance is actually increasing¶

Display initial model data structures

In [47]:

idata_1

Create the Posterior predictions (we have 100 samples at each of 3 x axis positions)

In [48]:

# Sample posterior predictive and append it automatically
with model_1:
    pm.sample_posterior_predictive(idata_1, extend_inferencedata=True)
# end with

pred_vals = (
    idata_1["posterior_predictive"]
    .sel({"month_dim_0": [0, 400, 795], "draw": slice(0, 99)})
    .mean(dim="chain")["l_pred"]
    .values
)

Sampling: [l_pred]

Output()

Plot the prediction distributions: we can see variance growing as we move left to right

In [49]:

labels = ["first", "middle", "last"]
for i in range(3):
    _ = plt.hist(pred_vals[:, i], histtype="step", label=labels[i])
# end for
plt.legend()
plt.title("Distribution of predicted length at position indicated")

Out[49]:

Text(0.5, 1.0, 'Distribution of predicted length at position indicated')

Show the data structures

In [50]:

idata_1

Visualize growth of standard deviation of length¶

Get data for plotting

In [51]:

post_vals = idata_1["posterior"].mean(dim=["draw", "chain"])["error_sd"].values

In [52]:

len(post_vals)

Out[52]:

In [53]:

plt.plot(df.month, post_vals, "+")

Out[53]:

[<matplotlib.lines.Line2D at 0x224516cf230>]

Indeed standard deviation grows linearly

Final model predictions¶

We now look at predictions for all months out to month 40

Change the month data structure to trigger calculations

In [54]:

with model_1:
    # set out of sample months
    oos_months = [i for i in range(41)]
    pm.set_data(
        {"month": oos_months},
    )

    post_pred = pm.sample_posterior_predictive(
        idata_1, var_names=["l_pred", "error_sd"]
    )

# end with

Sampling: [l_pred]

Output()

Show the new data structures

In [55]:

post_pred

Get the data for plotting

In [56]:

l_preds = post_pred["posterior_predictive"]["l_pred"].mean(("draw", "chain"))

error_sds = post_pred["posterior_predictive"]["error_sd"].mean(("draw", "chain"))

Get the model key parameters

In [57]:

results = az.summary(idata_1, var_names=["intercept", "slope"])
results

Out[57]:

	mean	sd	eti89_lb	eti89_ub	ess_bulk	ess_tail	r_hat	mcse_mean	mcse_sd
intercept	48.05	0.234	48	48	1293	1580	1.00	0.0065	0.0044
slope	7.867	0.085	7.7	8	1265	1603	1.00	0.0024	0.0016

In [58]:

line_a = results.loc["intercept", "mean"]
line_b = results.loc["slope", "mean"]
line_a, line_b

Out[58]:

(np.float64(48.045530742250314), np.float64(7.866914155716518))

Plot the result

In [59]:

plt.plot(
    df["month"],
    df["length"],
    "+",
    label="raw data",
)
plt.plot(
    oos_months,
    l_preds,
    "-r",
    label="predicted",
)
plt.plot(
    df["month"],
    line_a + line_b * df["month"] ** 0.5,
    "og",
    label="sqrt fit",
)
plt.legend()

Out[59]:

<matplotlib.legend.Legend at 0x22451a67e00>

Enhanced plot¶

We draw lines at 1.96 standard deviations above and below our fitted growth curve (95% lines)

In [60]:

plt.plot(
    df["month"],
    df["length"],
    "+",
    label="raw data",
)
plt.plot(
    oos_months,
    l_preds,
    "-r",
    label="predicted",
)

upper_line = [x + 1.96 * y for x, y in zip(l_preds.values, error_sds.values)]
lower_line = [x - 1.96 * y for x, y in zip(l_preds.values, error_sds.values)]
plt.plot(
    oos_months,
    upper_line,
    "-g",
    label="upper_band (95%)",
)
plt.plot(
    oos_months,
    lower_line,
    "-y",
    label="lower_band (95%)",
)
plt.legend()
plt.title("sqrt length, linear std ")

Out[60]:

Text(0.5, 1.0, 'sqrt length, linear std ')

Check model execution (did all of possible parameter space get explored?)¶

In a perfect world, the two curves should overlap. However, I have used the default sampling parameters (idata_1 = pm.sample()), so I am not too worried

In [61]:

pc = az.plot_energy(
    idata_1,
)

Conclusions¶

I spent a lot more time fighting PyMC and ArViz documentation (especially the behavior of dimension names), more so than actually understanding the model. What didn't help was that Google AI and search are totally poisoned by the significant changes in function names, etc in the history of the packages. Be aware and make sure you are reading the latest version of the manuals.

Also don't be lazy (like me), and do read the PyMC document on data structures, to make sure you understand how to extract data, and how to create ancillary data structures

Reproducability¶

In [62]:

%load_ext watermark

The watermark extension is already loaded. To reload it, use:
  %reload_ext watermark

In [63]:

%watermark

Last updated: 2026-10-24T10:10:44.734542+10:00

Python implementation: CPython
Python version       : 3.14.7
IPython version      : 9.17.1

Compiler    : MSC v.1944 64 bit (AMD64)
OS          : Windows
Release     : 11
Machine     : AMD64
Processor   : Intel64 Family 6 Model 170 Stepping 4, GenuineIntel
CPU cores   : 22
Architecture: 64bit

In [64]:

%watermark -h -iv -co

conda environment: bayes_env

Hostname: INSPIRON16

arviz     : 1.3.0
matplotlib: 3.11.1
numpy     : 2.4.6
pandas    : 3.0.5
preliz    : 0.28.0
pymc      : 6.3.1
xarray    : 2026.7.0

In [ ]:

Bayesian analysis of Heteroskedasticity

More regression with Bayes¶

Introduction¶

Implementation¶

Imports¶

Load dataset¶

Visualize dataset¶

Build and run linear model¶

Prediction of in-sample values¶

Initial DataTree from model run¶

Digression: Get predictions from Priors¶

End of digression¶

Get predictions from one chain of Posteriors¶

Extract data for plotting¶

Get predictions from mean of all chains of Posteriors¶

Conclusions¶

Prediction of out-of-sample values¶

Digression: using plot_lm¶

End of digression¶

sqrt model, constant variance¶

Get predictions from Posteriors¶

Final model¶

Check variance is actually increasing¶

Visualize growth of standard deviation of length¶

Final model predictions¶

Enhanced plot¶

Check model execution (did all of possible parameter space get explored?)¶

Conclusions¶

Reproducability¶

Comments

About