Objectives

Objectives convert raw equation streams into stable named scalar losses. They do not combine weighted losses; that policy belongs to a separately selected balancer.

Contracts

The protocols are structural and runtime-checkable. Custom implementations do not need to inherit them.

ObjectiveTerm [source]

Bases: Protocol

Define one independently configurable group of named objective losses.

loss_names `property`

loss_names: Sequence[str]

Return names produced by this term.

Returns:

Type	Description
`Sequence[str]`	Ordered unique loss names.

batch_keys `property`

batch_keys: Sequence[str]

Return batch names required by this objective term.

Returns:

Type	Description
`Sequence[str]`	Ordered unique batch names.

losses [source]

losses(
    model_apply: ModelApply,
    model_state: Any,
    batches: NamedBatches,
) -> Mapping[str, jax.Array]

Evaluate this term's named scalar losses.

Parameters:

Name	Type	Description	Default
`model_apply`	`ModelApply`	Pure explicit-state model application callable.	required
`model_state`	`Any`	Differentiable model parameter PyTree.	required
`batches`	`NamedBatches`	Fixed-structure named input batches.	required

Returns:

Type	Description
`Mapping[str, Array]`	Scalar losses keyed by :attr:`loss_names`.

residual_stream [source]

residual_stream(
    name: str,
    model_apply: ModelApply,
    model_state: Any,
    batches: NamedBatches,
) -> jax.Array

Evaluate one raw stream for derivative-based balancing.

Parameters:

Name	Type	Description	Default
`name`	`str`	One entry from :attr:`loss_names`.	required
`model_apply`	`ModelApply`	Pure explicit-state model application callable.	required
`model_state`	`Any`	Differentiable model parameter PyTree.	required
`batches`	`NamedBatches`	Fixed-structure named residual batches.	required

Returns:

Type	Description
`Array`	Raw output or equation residual array.

Objective [source]

Bases: Protocol

Define the scalar-loss contract consumed by compiled training steps.

loss_names `property`

loss_names: Sequence[str]

Return the stable names used to align losses and balancer weights.

Returns:

Type	Description
`Sequence[str]`	Ordered unique loss names.

batch_keys `property`

batch_keys: Sequence[str]

Return batch names required by the complete objective.

Returns:

Type	Description
`Sequence[str]`	Ordered unique batch names.

losses [source]

losses(
    model_apply: ModelApply,
    model_state: Any,
    batches: NamedBatches,
) -> Mapping[str, jax.Array]

Evaluate named scalar losses.

Parameters:

Name	Type	Description	Default
`model_apply`	`ModelApply`	Pure explicit-state model application callable.	required
`model_state`	`Any`	Differentiable model parameter PyTree.	required
`batches`	`NamedBatches`	Fixed-structure named input batches.	required

Returns:

Type	Description
`Mapping[str, Array]`	Scalar loss arrays keyed by :attr:`loss_names`.

ResidualObjective [source]

Bases: Objective, Protocol

Extend an objective with one-at-a-time residual streams for derivative-based balancing.

residual_stream [source]

residual_stream(
    name: str,
    model_apply: ModelApply,
    model_state: Any,
    batches: NamedBatches,
) -> jax.Array

Evaluate one raw residual stream.

Parameters:

Name	Type	Description	Default
`name`	`str`	One entry from :attr:`loss_names`.	required
`model_apply`	`ModelApply`	Pure explicit-state model application callable.	required
`model_state`	`Any`	Differentiable model parameter PyTree.	required
`batches`	`NamedBatches`	Fixed-structure named residual batches.	required

Returns:

Type	Description
`Array`	Residual array whose parameter Jacobian defines the stream's empirical NTK diagnostics.

Residual terms

Each outer residual group becomes one scalar loss. ResidualTerm sums the mean square of every array in the group. For example, one phase-wrapped loss can combine cosine and sine residuals without a custom objective class.

If names is absent, metadata attached by residual_equation produces names such as pde/continuity.

from functools import partial

import math

from phijax.equations import burgers_1d
from phijax.objectives import ResidualTerm

pde_term = ResidualTerm(
    partial(burgers_1d, viscosity_coefficient=0.01 / math.pi),
    batch_key="pde",
)

ResidualTerm [source]

ResidualTerm(
    residual_fn: ResidualFunction,
    *,
    batch_key: str,
    names: Sequence[str] | None = None,
    ntk_stream: ResidualStream | None = None,
)

Reduce one configured grouped-residual function into named scalar losses.

One outer residual group corresponds to one entry in :attr:loss_names. When explicit names are omitted, the equation's static local residual names are prefixed with :attr:batch_key. Every array inside a group contributes its own mean-squared value to that scalar loss. This preserves composite residuals such as phase-wrapped cosine and sine errors without requiring an equation-specific objective class.

Attributes:

Name	Type	Description
`loss_names`		Stable names aligned with the residual groups returned by `residual_fn`.

Initialize a generic residual term.

Parameters:

Name	Type	Description	Default
`residual_fn`	`ResidualFunction`	Configured equation callable returning one residual group per name.	required
`batch_key`	`str`	Key selecting the equation batch from the named training batches.	required
`names`	`Sequence[str] \| None`	Optional non-empty unique scalar loss names ordered like the equation's residual groups. When omitted, names are formed as `batch_key/local_name` from metadata attached to `residual_fn` by :func:`phijax.equations.residual_equation`.	`None`
`ntk_stream`	`ResidualStream \| None`	Optional stream override. When omitted, the equation metadata selects `"residual"` or `"output"`.	`None`

Raises:

Type	Description
`TypeError`	If `residual_fn` is not callable.
`ValueError`	If names, inferred equation metadata, `batch_key`, or `ntk_stream` are invalid.

losses [source]

losses(
    model_apply: ModelApply,
    model_state: Any,
    batches: NamedBatches,
) -> dict[str, jax.Array]

Evaluate and reduce the configured equation residual groups.

Parameters:

Name	Type	Description	Default
`model_apply`	`ModelApply`	Pure explicit-state model application callable.	required
`model_state`	`Any`	Differentiable model parameter PyTree.	required
`batches`	`NamedBatches`	Mapping containing :attr:`batch_key`.	required

Returns:

Type	Description
`dict[str, Array]`	Scalar `float32` losses keyed by :attr:`loss_names`.

residual_stream [source]

residual_stream(
    name: str,
    model_apply: ModelApply,
    model_state: Any,
    batches: NamedBatches,
) -> jax.Array

Return one two-dimensional stream for derivative-based balancing.

Parameters:

Name	Type	Description	Default
`name`	`str`	One configured scalar loss name.	required
`model_apply`	`ModelApply`	Pure explicit-state model application callable.	required
`model_state`	`Any`	Differentiable model parameter PyTree.	required
`batches`	`NamedBatches`	Mapping containing :attr:`batch_key`.	required

Returns:

Type	Description
`Array`	Residual matrix with samples on axis zero and flattened residual features on axis one.

Raises:

Type	Description
`KeyError`	If `name` does not belong to this term.

Composite objectives

The term mapping establishes stable composition names and preserves declaration order. Loss names must be globally unique across terms.

Use from_equations() when each equation consumes the batch with the same mapping key. Equation metadata supplies local loss names and the default exact-NTK stream:

from phijax.equations import base_data_fidelity, burgers_1d
from phijax.objectives import CompositeObjective

objective = CompositeObjective.from_equations(
    {
        "initial": base_data_fidelity,
        "boundary": base_data_fidelity,
        "pde": burgers_1d,
    }
)

Use explicit terms when loss names, batch routing, or streams require overrides:

from phijax.objectives import CompositeObjective

objective = CompositeObjective(
    terms={
        "initial": initial_term,
        "boundary": boundary_term,
        "pde": pde_term,
    }
)

CompositeObjective [source]

CompositeObjective(terms: Mapping[str, ObjectiveTerm])

Compose an ordered mapping of independently configured objective terms.

The term mapping and loss names are resolved on the host before JAX traces a training step. Mapping keys provide stable Hydra override paths, while each value may produce one or more scalar losses and raw residual streams.

Attributes:

Name	Type	Description
`loss_names`		Flattened stable loss-name ordering across every term.
`batch_keys`		Required DataModule batch names in first-declaration order.

Initialize a composite objective.

Parameters:

Name	Type	Description	Default
`terms`	`Mapping[str, ObjectiveTerm]`	Non-empty ordered objective terms keyed by stable configuration names.	required

Raises:

Type	Description
`TypeError`	If an entry does not implement :class:`ObjectiveTerm`.
`ValueError`	If no terms are supplied or loss names are empty or duplicated.

from_equations [source] `classmethod`

from_equations(
    equations: Mapping[str, ResidualFunction],
) -> CompositeObjective

Create one residual term per named equation.

Mapping keys route DataModule batches and prefix equation-local loss names. Equation metadata selects the default derivative-balancing stream.

Parameters:

Name	Type	Description	Default
`equations`	`Mapping[str, ResidualFunction]`	Non-empty ordered mapping from batch names to decorated residual equations.	required

Returns:

Type	Description
`CompositeObjective`	Composite objective with inferred names, batch routing, and balancing streams.

losses [source]

losses(
    model_apply: ModelApply,
    model_state: Any,
    batches: NamedBatches,
) -> dict[str, jax.Array]

Evaluate and merge scalar losses from every configured term.

Parameters:

Name	Type	Description	Default
`model_apply`	`ModelApply`	Pure explicit-state model application callable.	required
`model_state`	`Any`	Differentiable model parameter PyTree.	required
`batches`	`NamedBatches`	Fixed-structure named input batches.	required

Returns:

Type	Description
`dict[str, Array]`	Scalar losses in :attr:`loss_names` order.

residual_stream [source]

residual_stream(
    name: str,
    model_apply: ModelApply,
    model_state: Any,
    batches: NamedBatches,
) -> jax.Array

Dispatch one named residual stream to its owning term.

Parameters:

Name	Type	Description	Default
`name`	`str`	One entry from :attr:`loss_names`.	required
`model_apply`	`ModelApply`	Pure explicit-state model application callable.	required
`model_state`	`Any`	Differentiable model parameter PyTree.	required
`batches`	`NamedBatches`	Fixed-structure named residual batches.	required

Returns:

Type	Description
`Array`	Raw output or equation residual array for `name`.

Raises:

Type	Description
`KeyError`	If no configured term owns `name`.

residual_streams [source]

residual_streams(
    model_apply: ModelApply,
    model_state: Any,
    batches: NamedBatches,
) -> dict[str, jax.Array]

Evaluate every configured raw residual stream.

Parameters:

Name	Type	Description	Default
`model_apply`	`ModelApply`	Pure explicit-state model application callable.	required
`model_state`	`Any`	Differentiable model parameter PyTree.	required
`batches`	`NamedBatches`	Fixed-structure named residual batches.	required

Returns:

Type	Description
`dict[str, Array]`	Raw arrays in :attr:`loss_names` order.

components [source]

components(losses: Mapping[str, Array]) -> jax.Array

Stack scalar losses in stable configured order.

Parameters:

Name	Type	Description	Default
`losses`	`Mapping[str, Array]`	Scalar loss arrays keyed by every name in :attr:`loss_names`.	required

Returns:

Type	Description
`Array`	`float32` component vector aligned with :attr:`loss_names`.

Raises:

Type	Description
`KeyError`	If a required loss is absent.

See Building equations and objectives for a complete custom equation composition.

Objectives

Contracts

ObjectiveTerm [source]

loss_names property

batch_keys property

losses [source]

residual_stream [source]

Objective [source]

loss_names property

batch_keys property

losses [source]

ResidualObjective [source]

residual_stream [source]

Residual terms

ResidualTerm [source]

losses [source]

residual_stream [source]

Composite objectives

CompositeObjective [source]

from_equations [source] classmethod

losses [source]

residual_stream [source]

residual_streams [source]

components [source]

loss_names `property`

batch_keys `property`

loss_names `property`

batch_keys `property`

from_equations [source] `classmethod`