Composition

Domains can be combined into product domains (e.g. space-time) and then used to define domain-aware functions and constraints.

Dataset factors (\(\Omega_{\text{data}}\))

Many operator-learning workflows decompose the domain as a product

\[ \Omega = \Omega_{\text{data}} \times \Omega_x \times \Omega_t \times \cdots, \]

where \(\Omega_{\text{data}}\) indexes a finite dataset of input functions/fields (e.g. forcing terms, initial conditions, material parameters) and the remaining factors are geometric/scalar coordinates.

DatasetDomain is a unary domain that stores an in-memory dataset (a PyTree of arrays with a shared leading dataset axis) and samples by random indexing. This integrates cleanly with:

• ProductDomain composition via @,
• structured sampling (dense_structure=ProductStructure((("data",),))),
• Domain.Model(...) for building operator-learning models that take (data, coords...).

Measure semantics

DatasetDomain(..., measure=...) controls the measure used by integral/mean reductions:

• measure="probability": \(\mu(\Omega_{\text{data}})=1\) (treat as an expectation),
• measure="count": \(\mu(\Omega_{\text{data}})=N\) where \(N\) is dataset size (treat as a finite-sum domain).

Example

import jax.numpy as jnp
import phydrax as phx

data = jnp.ones((128, 64))     # N=128 samples, each with 64 features
Omega = phx.domain.DatasetDomain(data, label="data") @ phx.domain.Interval1d(0.0, 1.0)

phydrax.domain.ProductDomain

A labeled Cartesian product of unary domains.

Given unary domains \(\Omega_1,\dots,\Omega_k\) with labels \(\ell_1,\dots,\ell_k\), a ProductDomain represents the product

\[ \Omega = \Omega_{\ell_1}\times\cdots\times\Omega_{\ell_k}. \]

Labels must be unique unless colliding factors are structurally equivalent, in which case they are de-duplicated. This allows joining domains that share the same label (e.g. two identical TimeIntervals) without ambiguity.

factors property

Return the unary factors of the product domain.

labels property

Return the label tuple \((\ell_1,\ldots,\ell_k)\) for this product domain.

__init__(*domains: _AbstractDomain)

Create a product domain from unary domains (or nested product domains).

factor(label: str) -> _AbstractUnaryDomain

Return the unary factor corresponding to label.

equivalent(other: object) -> bool

Return whether other is structurally equivalent to this domain.

boundary()

Return the boundary as a DomainComponentUnion.

This constructs a union over boundary terms for each factor:

• For geometry factors, one term with that label set to Boundary().
• For scalar factors (e.g. time), two terms corresponding to FixedStart() and FixedEnd().

This mirrors the common decomposition of \(\partial(\Omega_x\times[t_0,t_1])\) into spatial boundary and initial/final time slices.

---

phydrax.domain.DatasetDomain

A unary domain over a finite in-memory dataset.

A DatasetDomain stores a PyTree of arrays where every leaf has a leading dataset axis of the same length N. Sampling draws a batch of indices uniformly and returns the corresponding slice from each leaf.

This is intended for product domains like Omega_data @ Omega_x, where data samples are paired/broadcast with spatial points.

size property

measure property

__init__(data: PyTree[ArrayLike], /, *, label: str = 'data', measure: Literal[probability, count] = 'probability')

sample(num_points: int, *, sampler: str = 'uniform', key: Key[Array, ''] = jr.key(0)) -> PyTree[Array]