feat(basis): add_block for structured Θ(a)·column blocks

[jkitchin]

Closes #18.

Adds BasisLibrary.add_block(), which copies every function of another library onto this library's feature space, optionally multiplied by a named data column. Design matrices of the form Φ = [Θ(c) ⊙ y_x | Θ(c)] become a supported idiom instead of a pile of hand-written add_custom calls, and a coefficient selected in the first block is literally a term of the unknown coefficient function.

theta = (BasisLibrary(n_features=1, feature_names=["c"])
         .add_constant().add_linear().add_polynomials(max_degree=2))

library = (BasisLibrary(n_features=2, feature_names=["c", "y_x"])
           .add_block(theta, multiply_by="y_x", block_name="horizontal")
           .add_block(theta, block_name="vertical"))

library.names    # ['y_x', 'c*y_x', 'c^2*y_x', '1', 'c', 'c^2']
library.blocks   # {'horizontal': [0, 1, 2], 'vertical': [3, 4, 5]}

What it does

• Names are generated as <basis>*<column>; the constant collapses (1*y_x → y_x) and a bare-sum source name is parenthesized (1+q → (1+q)*y_x), so the *-splitting parsers in regressor.to_callable() / simplify stay correct.
• Complexity is inherited from the source, plus 1 for the multiplication, plus an optional complexity_offset. This reproduces the hand-written scores in the issue exactly (y_x→1, q*y_x→2, q^2*y_x→3, the parametric term→4).
• Feature remapping: the source is written against its own columns and is re-expressed here, matching features by name (feature_map={"src": "target"} when names differ). Mapped feature_indices plus the multiplier are recorded. The gather is skipped when the two feature spaces coincide, so copied functions stay as cheap and as jit-friendly as the originals.
• Parametric terms pass through unchanged — registered via add_parametric, preserving bounds, log_scale and the name template, so profile-likelihood optimisation still applies inside the block.
• Source functions are copied, not shared, so one theta can seed several blocks.

Block bookkeeping

• BasisFunction.block field (serialized only when set), new func_type="block" with a from_dict branch that points at add_block() — same posture as custom.
• blocks property, filter_by_block(include=, exclude=) (an unknown block name raises rather than silently returning nothing), and without_blocks(*names) returning a new library with parametric bookkeeping re-indexed — the "did the vertical block earn its place at all?" diagnostic.

Tests

25 new tests in tests/test_basis.py covering name generation, evaluation, complexity, feature remapping, feature_map, parametric pass-through, block queries, without_blocks re-indexing, serialization and the error paths — plus an end-to-end case recovering s'(c) = 1 + 2c from y_c = s'(c)·y_x as {y_x, c*y_x} with the right coefficients.

• Full suite: 648 passed, 10 skipped on JAX
• NumPy backend (scripts/test_under_numpy.py): 648 passed, 10 skipped — no new JAX API surface beyond jnp.asarray and advanced indexing
• black --check and ruff check clean

Docs

New add_block section in the skill's guides/basis-library.md (verified against the actual behaviour, not written from memory) plus a routing row and guidance bullet in SKILL.md; src/jaxsr/skill/ re-synced from .claude/skills/jaxsr/.

🤖 Generated with Claude Code

https://claude.ai/code/session_01StMr2PeXc8wP67YvLCDqdQ

---

Generated by Claude Code