Split data

A split decides which rows belong to training, validation, and testing. Create it before configuring minibatches; shuffling batches never changes the held-out data. Selecting rows must preserve their likelihood contributions. The optimizer does not check this property.

This page uses a small Gaussian regression. Expand the setup to run the examples from top to bottom.

import logging

import jax.numpy as jnp
import numpy as np
import pandas as pd
import tensorflow_probability.substrates.jax.distributions as tfd

import liesel.model as lsl
import liesel.optim as opt

Show code cell source

Hide code cell source

logging.getLogger("liesel").setLevel(logging.WARNING)

rng = np.random.default_rng(42)
x = np.linspace(-1.0, 1.0, 128)
X = lsl.Var.new_obs(jnp.asarray(x), name="X")
beta = lsl.Var.new_param(0.0, lsl.Dist(tfd.Normal, 0.0, 5.0), name="beta")
log_sigma = lsl.Var.new_param(0.0, name="log_sigma")
sigma = lsl.Var.new_calc(jnp.exp, log_sigma, name="sigma")
mu = lsl.Var.new_calc(lambda X, beta: X * beta, X, beta, name="mu")
y = lsl.Var.new_obs(
    jnp.asarray(0.5 * x + rng.normal(scale=0.7, size=x.size)),
    lsl.Dist(tfd.Normal, mu, sigma),
    name="y",
)
model = lsl.Model(y)

Keep rows together

Split responses and their covariates together:

split = opt.PositionSplit.from_model(
    model,
    position_keys=["X", "y"],
    validate_axis_share=0.2,
    test_axis_share=0.1,
    seed=42,
)

pd.DataFrame(
    {"rows": [split.train_axis_size, split.validate_axis_size, split.test_axis_size]},
    index=["training", "validation", "test"],
)

	rows
training	91
validation	25
test	12

This puts 70% of rows in training, 20% in validation, and 10% in testing, subject to rounding. Splits with validation or test data shuffle by default; set shuffle=False for an ordered split. Without holdouts, splits preserve the original row order and ignore the seed, even with shuffle=True. This is also the behavior of LieselOptim’s automatic full-training split.

The split’s seed chooses which rows go into each part. The seed passed to LieselOptim controls batch sampling during fitting. Both default to 0. Starting parameter values come from the model; seed any random data or starting values separately.

Split several groups

By default, LieselOptim raises when observed arrays have different lengths. This gives you a chance to check which arrays share rows and which are shared data. For independent groups, opt in to grouping by length and inspect the result:

y_a = lsl.Var.new_obs(
    rng.normal(size=80),
    lsl.Dist(tfd.Normal, 0.0, 1.0),
    name="y_a",
)
y_b = lsl.Var.new_obs(
    rng.normal(size=40),
    lsl.Dist(tfd.Normal, 0.0, 1.0),
    name="y_b",
)
grouped_model = lsl.Model([y_a, y_b])
grouped_split = opt.PositionSplit.from_model(
    grouped_model, multi_size="manager", validate_axis_share=0.2, seed=42
)

pd.DataFrame(
    [
        {
            "keys": group.position_keys,
            "training": group.train_axis_size,
            "validation": group.validate_axis_size,
        }
        for group in grouped_split.splits
    ]
)

	keys	training	validation
0	[y_b]	32	8
1	[y_a]	64	16

Pass the checked split to LieselOptim as split=grouped_split. Matching lengths do not establish row alignment: flat or omitted position_keys group arrays by length only. To choose groups explicitly, use PositionSplitManager.from_model with nested keys, such as position_keys=[["X_a", "y_a"], ["X_b", "y_b"]]. The two-group tutorial shows a complete fit.

Arrays within a group share row indices. Different groups split independently, even if their lengths happen to match. Every group must have validation data if any group does. Keep shared arrays out of row groups with split_axes={key: None}, as described below.

Automatic grouping requires an observed likelihood in each group. For row data without a likelihood, supply an explicit nested group.

Keep shared data

Shared tables and scalar constants must stay unchanged in every split. For a lookup table z indexed by row-level group IDs, use PositionSplit.from_model(model, split_axes={"z": None}) and pass the result as LieselOptim(..., split=split). This also keeps z out of automatic batches, even if its length matches a response. Keep per-row covariates, weights, and offsets with their response.

For observations on an axis other than zero, set split_axes and the corresponding batch_axes when creating batches.

PositionSplit holds the split data. Split holds reusable row indices; call split_position() to apply them. Manager classes handle several groups. See from_model() for the factory options.

For reduced likelihoods or custom objectives, see Scale the loss. For computed observations such as copulas, see Batch weak responses.

Open the split API overview in a separate page.

Next, create minibatches or read how loss scaling handles held-out data.