import warnings
from jax import numpy as jnp
import numpy as np
from typing import Union, Tuple
from einops import repeat
from gensbi.utils.math import _expand_dims
from gensbi.diffusion.path import EDMPath
from gensbi.diffusion.path.scheduler import (
EDMScheduler,
VEEdmScheduler,
VPEdmScheduler,
)
from gensbi.diffusion.path.sm_path import SMPath
from gensbi.diffusion.path.scheduler import VPSmScheduler, VESmScheduler
[docs]
def init_ids_joint(dim_obs: int, dim_cond: int):
dim_joint = dim_obs + dim_cond
node_ids = jnp.arange(dim_joint).reshape((1, -1, 1))
obs_ids = jnp.arange(dim_obs).reshape((1, -1, 1)) # observation ids
cond_ids = jnp.arange(dim_obs, dim_joint).reshape((1, -1, 1)) # conditional ids
return node_ids, obs_ids, cond_ids
[docs]
def init_ids_1d(dim: int, semantic_id: Union[int, None] = None):
"""Build 1D positional IDs, returning ``(ids, dim)``.
``ids`` is ``(1, dim, 1)`` when ``semantic_id is None`` (position only), or
``(1, dim, 2)`` otherwise, with the position at axis 0 and the semantic id
at axis 1.
FIXME (axis-order footgun): this places the semantic id on the LAST axis,
which is the *reverse* of :func:`init_ids_2d` (semantic at axis 0, then h, w
-- the established convention). The two are safe in isolation, but they do
NOT line up if 1D and 2D ids are ever fed into the same RoPE grid (e.g.
FieldDiT Phase-2 obs+cond co-tokenization with a shared ``EmbedND``, where
``axes_dim[i]`` is matched to ids axis ``i``). This should be unified to the
2D convention (semantic at axis 0). It is not changed here because callers
that pass ``semantic_id`` -- notably the Flux1 ``rope1d`` path via
:func:`init_ids` -- and any code indexing ``ids[..., k]`` must be updated in
lockstep.
"""
if semantic_id is None:
ids = np.zeros((1, dim, 1), dtype=np.int32)
else:
ids = np.zeros((1, dim, 2), dtype=np.int32)
ids[..., 1] = semantic_id
ids[0, :, 0] = np.arange(dim)
return jnp.array(ids, dtype=jnp.int32), dim
[docs]
def healpix_rope_theta(nside: int) -> int:
"""Suggested RoPE ``theta`` for a full-sky HEALPix token grid at ``nside``.
Follows the project convention ``theta = 10 * token count`` (the same rule
:class:`~gensbi.models.flux1.model.Flux1Params` applies by default via
``10 * (dim_obs + dim_cond)``): a full-sky grid has ``12 * nside**2``
tokens. Exposed so spherical models can derive theta from the intuitive
knob (``nside``, always known after the encoder) instead of setting it by
hand.
"""
return 10 * 12 * nside**2
[docs]
def _validate_base_pixels(base_pixels):
"""Validate a base-pixel subset spec; return it as a list of ints."""
base_pixels = list(base_pixels)
if not base_pixels:
raise ValueError(
"base_pixels must be non-empty; omit it (or pass None) for full sky"
)
if any(not isinstance(b, (int, np.integer)) for b in base_pixels):
raise ValueError(f"base_pixels entries must be integers, got {base_pixels}")
if any(b < 0 or b > 11 for b in base_pixels) or len(set(base_pixels)) != len(
base_pixels
):
raise ValueError(
f"base_pixels must be unique integers in [0, 11], got {base_pixels}"
)
return base_pixels
[docs]
def init_ids_healpix(nside: int, base_pixels=None):
"""Build spherical RoPE ids for tokens on a HEALPix grid, returning
``(ids, num_tokens)``.
Method: standard N-dimensional RoPE (RoFormer, arXiv:2104.09864 — the
mechanism implemented by Flux1's ``EmbedND``) applied uniformly, on all
three axes and all frequency bands, to the 3D Cartesian coordinates of
HEALPix pixel centers on the unit sphere. Each token maps to its
pixel-center unit vector (``healpy.pix2vec``, NEST ordering), scaled to
pixel units so adjacent tokens differ by ~1 in coordinate (radius
``nside * sqrt(3/pi)`` = 1/pixel angular size), which keeps ``theta``'s
semantics identical to 2D-image usage (see :func:`healpix_rope_theta`).
Attention scores then depend on positions only through the chord vector
``n_q - n_k``, whose norm ``2 sin(gamma/2)`` is strictly monotone in
great-circle distance ``gamma`` — geodesic geometry with no projection
step, hence no face-seam or polar artifacts, and any ``base_pixels``
subset works by construction. Caveat: ``d(chord)/d(gamma) -> 0`` at
antipodes, so resolution among near-antipodal separations is mildly
compressed (benign for near/far attention). This is NOT an adaptation of
SpheRoPE (arXiv:2606.32033 — closest prior work; ERP grid, pretrained
constraints); see also StereoRoPE (arXiv:2606.31248, documents the
failure of index-based RoPE on HEALPix) and Unlu (arXiv:2310.04454, an
SO(3) feature-rotation alternative not adopted). Full rationale:
``docs/superpowers/specs/2026-07-19-healpix-rope-design.md``.
Use with Flux1 via ``id_embedding_strategy=("absolute", "rope")`` and a
3-entry ``axes_dim`` (each even, summing to the per-head dim, e.g.
``(22, 22, 20)`` for 64). Obs/theta-stream tokens automatically get
origin (0, 0, 0) rope ids — the identity rotation, i.e. an exactly
isotropic positional readout of the conditioning tokens.
Parameters
----------
nside : int
HEALPix resolution of the *token* grid (power of 2). With HEAL-SWIN
style encoders this is the bottleneck nside; tokens must correspond
to single HEALPix pixels (power-of-4 pixels-per-token upstream).
base_pixels : sequence of int, optional
Base pixels (0..11) covered by the token grid, for partial-sky
models. ``None`` (default) means full sky. Tokens are ordered by
base pixel as given, NEST within each.
Returns
-------
ids : jax.Array
``(1, num_tokens, 3)`` float32 scaled pixel-center coordinates.
num_tokens : int
``len(base_pixels) * nside**2``.
"""
if nside < 1 or (nside & (nside - 1)) != 0:
raise ValueError(f"nside must be a power of 2, got {nside}")
if base_pixels is None:
base_pixels = range(12)
base_pixels = _validate_base_pixels(base_pixels)
import healpy as hp # lazy: healpy pulls matplotlib, keep import light
face_len = nside**2
pix = np.concatenate(
[b * face_len + np.arange(face_len) for b in base_pixels]
)
x, y, z = hp.pix2vec(nside, pix, nest=True) # float64 host-side
radius = nside * np.sqrt(3.0 / np.pi) # 1 / pixel angular size
ids = radius * np.stack([x, y, z], axis=-1)[None, ...]
return jnp.asarray(ids, dtype=jnp.float32), ids.shape[1]
[docs]
def _normalize_patch_size(size):
"""Normalize a patch-size spec into an ``(obs, cond)`` tuple.
Parameters
----------
size : int or tuple of int
A single int is broadcast to both inputs (``8 -> (8, 8)``). A
length-2 tuple is taken as ``(obs_size, cond_size)`` so the two
inputs can use different patch sizes. Use ``1`` for an input that
is not patchified.
Returns
-------
tuple of int
``(obs_size, cond_size)``.
"""
if isinstance(size, int):
return (size, size)
size = tuple(size)
if len(size) != 2:
raise ValueError(
f"size must be an int or a length-2 (obs, cond) tuple, got {size!r}"
)
return size
[docs]
def init_ids_2d(dim: Tuple[int, int], semantic_id: int = 0, size: int = 2):
"""Build 2D positional IDs for a patchified image grid.
The grid has one entry per patch, i.e. ``(dim[0] // size, dim[1] // size)``,
matching ``patchify_2d(x, size=size)``. ``size`` is the patch edge length;
use ``size=1`` for no patchification (one token per pixel).
"""
img_ids = np.zeros((dim[0] // size, dim[1] // size, 3), dtype=np.int32)
img_ids[..., 0] = semantic_id
img_ids[..., 1] = img_ids[..., 1] + np.arange(dim[0] // size)[:, None]
img_ids[..., 2] = img_ids[..., 2] + np.arange(dim[1] // size)[None, :]
img_ids = repeat(img_ids, "h w c -> b (h w) c", b=1)
dim = (dim[0] // size) * (dim[1] // size)
return jnp.array(img_ids, dtype=jnp.int32), dim
[docs]
def _require_channel(x, name="input"):
"""Enforce a tabular channel axis (B, dim, C); reject a bare (B, dim)."""
x = jnp.asarray(x)
if x.ndim < 3:
raise ValueError(
f"{name} must carry a channel axis (B, dim, C); got shape "
f"{tuple(x.shape)}. A bare (B, dim) is not accepted — add a trailing "
f"channel axis (e.g. x[..., None] for C=1).")
return x
[docs]
def _single_obs(x_o, *, channel, name="x_o"):
"""Canonicalize a single conditioning observation, then enforce batch == 1.
Shape handling comes FIRST so a misshaped input can never be misread as a
batch (e.g. ``(dim, C)`` read as ``dim`` observations):
- ``channel="require"``: tabular flow-pipeline contract — input must
already carry batch and channel axes ``(1, dim, C)``; channel-less input
raises ``ValueError`` (same :func:`_require_channel` as training).
- ``channel="promote"``: FM-pipeline contract — lenient promotion:
``(dim,) -> (1, dim, 1)`` and ``(B, dim) -> (B, dim, 1)``.
- ``channel="none"``: structured inputs — the model owns the trailing
shape; only a leading batch axis (``ndim >= 2``) is required.
A leading batch axis > 1 then raises ``ValueError``: single-observation
methods never silently discard observations — use ``sample_batched``.
Returns the canonicalized array with its size-1 batch axis kept.
"""
x_o = jnp.asarray(x_o)
if channel == "require":
x_o = _require_channel(x_o, name)
elif channel == "promote":
orig_shape = tuple(x_o.shape)
x_o = _expand_dims(x_o)
if x_o.ndim < 3:
raise ValueError(
f"{name} must be at least 1-D (dim,); got shape {orig_shape}.")
elif channel == "none":
if x_o.ndim < 2:
raise ValueError(
f"{name} must carry a leading batch axis (e.g. (1,) + "
f"per_observation_shape); got shape {tuple(x_o.shape)}.")
else:
raise ValueError(f"unknown channel mode {channel!r}")
if x_o.shape[0] > 1:
raise ValueError(
f"{name} has a leading batch axis of size {x_o.shape[0]} > 1, but "
"this method conditions on a single observation and will not "
"silently discard the rest. Use sample_batched() for a batch of "
"conditions.")
return x_o
[docs]
def scale_lr(batch_size, base_lr=1e-4, reference_batch_size=256):
"""Scale learning rate based on batch size using square root scaling.
Parameters
----------
batch_size : int
The current batch size.
base_lr : float
The base learning rate for the reference batch size.
reference_batch_size : int, optional
The reference batch size. Defaults to 256.
Returns
-------
float
The adjusted learning rate.
"""
import math
return base_lr * math.sqrt(batch_size / reference_batch_size)
[docs]
_EMBEDDINGS_1D = {"absolute", "pos1d", "rope1d"}
[docs]
_EMBEDDINGS_2D = {"pos2d", "rope2d"}
[docs]
def _resolve_embedding_ids(dim, strategy: str, semantic_id: int, size: int = 2):
"""Resolve ID embeddings by strategy name.
Parameters
----------
dim : int or tuple of int
Dimension specification (number of tokens, or (H, W) for 2D images).
strategy : str or IdStrategy
Embedding strategy name (e.g., "absolute", "pos1d", "rope1d",
"pos2d", "rope2d") or a strategy object with ``build(dim)`` (e.g.
:class:`gensbi.recipes.HealpixRope`). NOTE these pipeline-side
builder names are a different vocabulary from the model-side
``id_embedding_strategy`` strings (where "rope" means "apply RoPE
to the ids the pipeline provides") — see
:mod:`gensbi.recipes.id_strategies`.
semantic_id : int
Semantic identifier for the token group (0=obs, 1=cond).
size : int, optional
Patch edge length for 2D strategies (default 2). Ignored for 1D
strategies. Use 1 for no patchification.
Returns
-------
ids : Array
Token ID array.
resolved_dim : int
Resolved flat dimension.
Raises
------
ValueError
If ``strategy`` is not recognized.
"""
if hasattr(strategy, "build"):
return strategy.build(dim)
if strategy in _EMBEDDINGS_1D:
return init_ids_1d(dim, semantic_id=semantic_id)
elif strategy in _EMBEDDINGS_2D:
return init_ids_2d(dim, semantic_id=semantic_id, size=size)
else:
raise ValueError(
f"Unknown id embedding strategy: {strategy!r}. Expected one of "
f"{sorted(_EMBEDDINGS_1D | _EMBEDDINGS_2D)} or an IdStrategy object "
"with a build(dim) method (e.g. gensbi.recipes.HealpixRope)."
)
[docs]
def build_edm_path(sde: str, config: dict) -> EDMPath:
"""Build an EDM-family diffusion path from an SDE type string and config.
Parameters
----------
sde : str
SDE type: ``"EDM"``, ``"VE"``, or ``"VP"``.
config : dict
Training configuration dict; scheduler hyperparameters are read from
here with sensible defaults.
Returns
-------
EDMPath
Configured diffusion path.
Raises
------
ValueError
If ``sde`` is not one of ``{"EDM", "VE", "VP"}``.
"""
if sde == "EDM":
return EDMPath(
scheduler=EDMScheduler(
sigma_min=config.get("sigma_min", 0.002),
sigma_max=config.get("sigma_max", 80.0),
)
)
elif sde == "VE":
return EDMPath(
scheduler=VEEdmScheduler(
sigma_min=config.get("sigma_min", 0.02),
sigma_max=config.get("sigma_max", 100.0),
)
)
elif sde == "VP":
return EDMPath(
scheduler=VPEdmScheduler(
beta_min=config.get("beta_min", 0.1),
beta_max=config.get("beta_max", 19.9),
)
)
else:
raise ValueError(f"Unknown sde type: {sde}")
[docs]
def build_sm_path(sde_type: str, config: dict) -> SMPath:
"""Build a score-matching path from an SDE type string and config.
Parameters
----------
sde_type : str
SDE type: ``"VP"`` or ``"VE"``.
config : dict
Training configuration dict; scheduler hyperparameters are read from
here with sensible defaults.
Returns
-------
SMPath
Configured score-matching path.
Raises
------
ValueError
If ``sde_type`` is not one of ``{"VP", "VE"}``.
"""
if sde_type == "VP":
return SMPath(
VPSmScheduler(
beta_min=config.get("beta_min", 0.001),
beta_max=config.get("beta_max", 3.0),
)
)
elif sde_type == "VE":
return SMPath(
VESmScheduler(
sigma_min=config.get("sigma_min", 0.001),
sigma_max=config.get("sigma_max", 15.0),
)
)
else:
raise ValueError(f"sde_type must be one of ['VP', 'VE'], got {sde_type}.")
[docs]
def parse_training_config(config_path: str):
"""Parse training and optimizer configuration from a YAML config file.
Reads the ``training`` and ``optimizer`` sections of the config and
returns a flat dictionary consumed by :class:`AbstractPipeline`.
Parameters
----------
config_path : str
Path to the YAML configuration file.
Returns
-------
training_config : dict
Parsed training configuration dictionary.
"""
import yaml
with open(config_path, "r") as f:
config = yaml.safe_load(f)
# Training parameters
train_params = config.get("training", {})
multistep = train_params.get("multistep", 1)
training_config = {
"nsteps": train_params.get("nsteps", 30000) * multistep,
"ema_decay": train_params.get("ema_decay", 0.999),
"multistep": multistep,
"experiment_id": train_params.get("experiment_id", 1),
"early_stopping": train_params.get("early_stopping", True),
"val_every": train_params.get("val_every", 100) * multistep,
"val_error_ratio": train_params.get("val_error_ratio", 1.3),
# Optional method-specific parameters (override strategy defaults)
"sigma_min": train_params.get("sigma_min", 0.002),
"sigma_max": train_params.get("sigma_max", 80.0),
}
# Optimizer parameters
opt_params = config.get("optimizer", {})
MAX_LR = opt_params.get("max_lr", 1e-3)
MIN_LR = opt_params.get("min_lr", 0.0)
training_config["max_lr"] = MAX_LR
training_config["min_lr"] = MIN_LR
training_config["min_scale"] = MIN_LR / MAX_LR if MAX_LR > 0 else 0.0
training_config["warmup_steps"] = opt_params.get("warmup_steps", 500)
training_config["decay_transition"] = opt_params.get("decay_transition", 0.85)
# ema_decay can also be specified in optimizer section (backward compat)
if "ema_decay" in opt_params:
training_config["ema_decay"] = opt_params["ema_decay"]
return training_config
[docs]
_MOVED_TO_PATCHING = ("patchify_2d", "depatchify_2d")
[docs]
def __getattr__(name):
# Deprecated aliases: the functions moved to gensbi.models.core.patching,
# but main's published docs teach this import path. Keep one release cycle.
if name in _MOVED_TO_PATCHING:
warnings.warn(
f"gensbi.recipes.utils.{name} has moved to "
"gensbi.models.core.patching; this alias will be removed in a "
"future release.",
DeprecationWarning, stacklevel=2)
from gensbi.models.core import patching
return getattr(patching, name)
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
