pde_opt.numerics.functions.legendre

This module contains the Legendre polynomial expansion class for representing functions in PDEs.

Classes

ChemicalPotentialLegendrePolynomials(*args, ...)	Chemical potential Legendre polynomials.
DiffusionLegendrePolynomials(*args, **kwargs)	Diffusion Legendre polynomials.
FixedDegreeChemicalPotential(*args, **kwargs)	Chemical potential with fixed high degree (1000) but zeros for unused parameters.
LegendrePolynomialExpansion(*args, **kwargs)	Legendre polynomial expansion.
LegendrePolynomials(max_degree)
TestChemicalPotential(*args, **kwargs)	Test chemical potential.

class pde_opt.numerics.functions.legendre.LegendrePolynomialExpansion(*args: Any, **kwargs: Any)

Legendre polynomial expansion.

__init__(params: jax.Array)

params: jax.Array

max_degree: int

class pde_opt.numerics.functions.legendre.DiffusionLegendrePolynomials(*args: Any, **kwargs: Any)

Diffusion Legendre polynomials.

Uses exp to ensure positivity.

__init__(params: jax.Array)

expansion: LegendrePolynomialExpansion

class pde_opt.numerics.functions.legendre.ChemicalPotentialLegendrePolynomials(*args: Any, **kwargs: Any)

Chemical potential Legendre polynomials.

__init__(params: jax.Array, prior_fn: Callable | None = None)

expansion: LegendrePolynomialExpansion

prior_fn: Callable

class pde_opt.numerics.functions.legendre.FixedDegreeChemicalPotential(*args: Any, **kwargs: Any)

Chemical potential with fixed high degree (1000) but zeros for unused parameters.

This class always uses degree 1000 Legendre polynomials, but zeros out parameters beyond the specified number. This ensures constant forward pass time while allowing gradient computation scaling tests.

NOTE: THIS IS ONLY USED FOR TESTING PURPOSES.

__init__(params: jax.Array, prior_fn: Callable | None = None)

expansion: LegendrePolynomialExpansion

prior_fn: Callable

num_active_params: int

class pde_opt.numerics.functions.legendre.TestChemicalPotential(*args: Any, **kwargs: Any)

Test chemical potential.

__init__(params: jax.Array, prior_fn: Callable | None = None)

params: jax.Array

prior_fn: Callable

class pde_opt.numerics.functions.legendre.LegendrePolynomials(max_degree: int)

max_degree: int

T0(x)

T1(x)

__init__(max_degree: int) → None

T2(x)

T3(x)

T4(x)

T5(x)

T6(x)

T7(x)

T8(x)

T9(x)

T10(x)