Skip to content

frontx.finite

frontx.finite

Finite-difference diffusion solver and data fitting.

This module provides a small public API to (a) integrate a radial 1D diffusion model with a state-dependent diffusivity D(theta), and (b) fit that model to spatio-temporal data. Internally, integration is performed with Diffrax and a centered finite-difference stencil.

Public API: - :class:Solution: wrapper for evaluating the simulated field and metadata. - :func:solve: integrate the PDE given initial/boundary conditions. - :func:fit: estimate model/parameters by minimizing data misfit.

Solution

Bases: Module

Simulation result wrapper for the finite-difference model.

A :class:Solution evaluates the predicted field :math:\theta(r, t) on-demand via dense output from the ODE integrator and linear interpolation across the spatial grid.

Attributes:

Name Type Description
D Callable[[float | Array | ndarray[Any, Any]], float | Array | ndarray[Any, Any]]

Diffusivity callable used in the simulation (accepts theta and returns D(theta) with broadcastable shape).
r1 float

Domain radius (the grid spans r in [0, r1]).
Source code in frontx/finite.py 
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
class Solution(eqx.Module):
    """Simulation result wrapper for the finite-difference model.

    A :class:`Solution` evaluates the predicted field :math:`\\theta(r, t)`
    on-demand via dense output from the ODE integrator and linear interpolation
    across the spatial grid.

    Attributes:
        D: Diffusivity callable used in the simulation (accepts ``theta`` and
            returns ``D(theta)`` with broadcastable shape).
        r1: Domain radius (the grid spans ``r in [0, r1]``).
    """

    D: Callable[
        [float | jax.Array | np.ndarray[Any, Any]],
        float | jax.Array | np.ndarray[Any, Any],
    ]
    r1: float
    _sol: diffrax.Solution

    def __call__(
        self,
        r: float | jax.Array | np.ndarray[Any, Any],
        t: float | jax.Array | np.ndarray[Any, Any],
    ) -> float | jax.Array | np.ndarray[Any, Any]:
        """Evaluate the simulated field at coordinates ``(r, t)``.

        Dense time output is queried from the solver and linearly interpolated
        over the spatial grid ``[0, r1]``.

        Args:
            r: Radial position(s) where the field is desired (scalar or array).
            t: Time(s) at which to evaluate (scalar or array). Supports
                broadcasting with ``r``.

        Returns:
            Values of ``theta(r, t)`` with a shape broadcastable from ``r`` and
            ``t``.
        """
        theta = self._sol.evaluate(t)
        return jnp.interp(r, jnp.linspace(0, self.r1, theta.size), theta)

    @property
    def t1(self) -> float | jax.Array | np.ndarray[Any, Any]:
        """Final integration time."""
        return self._sol.t1

    @property
    def result(self) -> RESULTS:
        """Diffrax integration status (see :data:`frontx.RESULTS`)."""
        return self._sol.result

D: Callable[[float | jax.Array | np.ndarray[Any, Any]], float | jax.Array | np.ndarray[Any, Any]] instance-attribute

r1: float instance-attribute

t1: float | jax.Array | np.ndarray[Any, Any] property

Final integration time.

result: RESULTS property

Diffrax integration status (see :data:frontx.RESULTS).

__call__(r: float | jax.Array | np.ndarray[Any, Any], t: float | jax.Array | np.ndarray[Any, Any]) -> float | jax.Array | np.ndarray[Any, Any]

Evaluate the simulated field at coordinates (r, t).

Dense time output is queried from the solver and linearly interpolated over the spatial grid [0, r1].

Parameters:

Name Type Description Default
r float | Array | ndarray[Any, Any]

Radial position(s) where the field is desired (scalar or array).

 required
t float | Array | ndarray[Any, Any]

Time(s) at which to evaluate (scalar or array). Supports broadcasting with r.

 required

Returns:

Type Description
float | Array | ndarray[Any, Any]

Values of theta(r, t) with a shape broadcastable from r and
float | Array | ndarray[Any, Any]

t.
Source code in frontx/finite.py 
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
def __call__(
    self,
    r: float | jax.Array | np.ndarray[Any, Any],
    t: float | jax.Array | np.ndarray[Any, Any],
) -> float | jax.Array | np.ndarray[Any, Any]:
    """Evaluate the simulated field at coordinates ``(r, t)``.

    Dense time output is queried from the solver and linearly interpolated
    over the spatial grid ``[0, r1]``.

    Args:
        r: Radial position(s) where the field is desired (scalar or array).
        t: Time(s) at which to evaluate (scalar or array). Supports
            broadcasting with ``r``.

    Returns:
        Values of ``theta(r, t)`` with a shape broadcastable from ``r`` and
        ``t``.
    """
    theta = self._sol.evaluate(t)
    return jnp.interp(r, jnp.linspace(0, self.r1, theta.size), theta)

solve(D: Callable[[float | jax.Array | np.ndarray[Any, Any]], float | jax.Array | np.ndarray[Any, Any]], r1: float, t1: float, *, i: jax.Array | np.ndarray[Any, Any], b: float | None = None, throw: bool = True) -> Solution

Integrate the finite-difference diffusion model.

Solves a radial 1D diffusion-like PDE on r ∈ [0, r1] up to time t1 with state-dependent diffusivity D(theta). The initial condition is theta(r, 0) = i[r] (grid-aligned). At the inner boundary, r = 0, a Neumann-like symmetry condition is applied by the stencil. At the outer boundary, r = r1, a Neumann condition is used; optionally if b is provided, the first cell is clamped to b in y0 to mimic a Dirichlet-type condition on the inner node.

Parameters:

Name Type Description Default
D Callable[[float | Array | ndarray[Any, Any]], float | Array | ndarray[Any, Any]]

Callable returning D(theta) (scalar or array) for the current state theta.

 required
r1 float

Domain radius (maximum r).

 required
t1 float

Final time to integrate to.

 required
i Array | ndarray[Any, Any]

Initial condition values on the spatial grid (1D array). Its length defines the number of spatial nodes.

 required
b float | None

Optional value to clamp the first grid node at initialization.

 None
throw bool

If True, Diffrax raises on integration failures.

 True

Returns:

Name Type Description
A Solution

class:Solution object that can be called as sol(r, t) and
Solution

exposes sol.t1 and sol.result.
Notes
  • Spatial discretization uses a centered three-point stencil with a harmonic-like averaging of D at faces.
  • Time integration uses Kvaerno5 with adaptive step control.
Source code in frontx/finite.py 
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
@eqx.filter_jit
def solve(  # noqa: PLR0913
    D: Callable[  # noqa: N803
        [float | jax.Array | np.ndarray[Any, Any]],
        float | jax.Array | np.ndarray[Any, Any],
    ],
    r1: float,
    t1: float,
    *,
    i: jax.Array | np.ndarray[Any, Any],
    b: float | None = None,
    throw: bool = True,
) -> Solution:
    """Integrate the finite-difference diffusion model.

    Solves a radial 1D diffusion-like PDE on ``r ∈ [0, r1]`` up to time ``t1``
    with state-dependent diffusivity ``D(theta)``. The initial condition is
    ``theta(r, 0) = i[r]`` (grid-aligned). At the inner boundary,
    ``r = 0``, a Neumann-like symmetry condition is applied by the stencil.
    At the outer boundary, ``r = r1``, a Neumann condition is used; optionally
    if ``b`` is provided, the first cell is clamped to ``b`` in ``y0`` to mimic
    a Dirichlet-type condition on the inner node.

    Args:
        D: Callable returning ``D(theta)`` (scalar or array) for the current
            state ``theta``.
        r1: Domain radius (maximum ``r``).
        t1: Final time to integrate to.
        i: Initial condition values on the spatial grid (1D array). Its length
            defines the number of spatial nodes.
        b: Optional value to clamp the first grid node at initialization.
        throw: If ``True``, Diffrax raises on integration failures.

    Returns:
        A :class:`Solution` object that can be called as ``sol(r, t)`` and
        exposes ``sol.t1`` and ``sol.result``.

    Notes:
        - Spatial discretization uses a centered three-point stencil with a
            harmonic-like averaging of ``D`` at faces.
        - Time integration uses ``Kvaerno5`` with adaptive step control.
    """
    i = jnp.asarray(i)
    dr = r1 / (i.size - 1)
    dr2 = dr**2

    @diffrax.ODETerm[jax.Array]
    def term(
        _: float | jax.Array | np.ndarray[Any, Any],
        theta: jax.Array,
        args: None,  # noqa: ARG001
    ) -> jax.Array:
        D_ = jnp.asarray(D(theta))  # noqa: N806
        if D_.ndim == 0:
            D_ = jnp.repeat(D_, theta.size)  # noqa: N806
        Df = (D_[1:] + D_[:-1]) / 2  # noqa: N806

        return jnp.concatenate(
            [
                jnp.array([Df[0] / dr2 * (theta[1] - theta[0]) if b is None else 0.0]),
                Df[:-1] / dr2 * theta[:-2]
                - (Df[1:] + Df[:-1]) / dr2 * theta[1:-1]
                + Df[1:] / dr2 * theta[2:],
                jnp.array([Df[-1] / dr2 * (theta[-2] - theta[-1])]),
            ]
        )

    sol = diffrax.diffeqsolve(
        term,
        solver=diffrax.Kvaerno5(),
        t0=0,
        t1=t1,
        dt0=None,
        y0=i if b is None else i.at[0].set(b),
        stepsize_controller=diffrax.PIDController(rtol=1e-3, atol=1e-6),
        saveat=diffrax.SaveAt(t0=True, t1=True, dense=True),
        throw=throw,
    )

    return Solution(
        D=D,
        r1=r1,
        _sol=sol,
    )

fit(D: Callable[[float | jax.Array | np.ndarray[Any, Any]], float | jax.Array | np.ndarray[Any, Any]], r1: float, t1: float, r: jax.Array | np.ndarray[Any, Any], t: jax.Array | np.ndarray[Any, Any], theta: jax.Array | np.ndarray[Any, Any], /, sigma: float | jax.Array | np.ndarray[Any, Any] = 1, *, i: jax.Array | np.ndarray[Any, Any], b: float | None = None, max_steps: int = 15) -> Solution

Fit the finite-difference model to spatio-temporal observations.

This routine searches over candidate models/parameters (through :func:frontx._inverse.param.de_fit) to minimize the mean squared error between the simulated field and observations theta(r, t) with optional per-sample weighting sigma.

Parameters:

Name Type Description Default
D Callable[[float | Array | ndarray[Any, Any]], float | Array | ndarray[Any, Any]]

Initial diffusivity callable or parameterized model to start from.

 required
r1 float

Domain radius for the simulation.

 required
t1 float

Final time for the simulation.

 required
r Array | ndarray[Any, Any]

Radial coordinates of observations, shape (Nr,).

 required
t Array | ndarray[Any, Any]

Time coordinates of observations, shape (Nt,).

 required
theta Array | ndarray[Any, Any]

Observed values with shape broadcastable to (Nr, Nt) or exactly (Nr, Nt).

 required
sigma float | Array | ndarray[Any, Any]

Noise/weight (scalar or array broadcastable to theta), used in the weighted MSE.

 1
i Array | ndarray[Any, Any]

Initial condition on the grid (length defines spatial resolution).

 required
b float | None

Optional clamped value for the first grid node at initialization.

 None
max_steps int

Maximum number of differential-evolution iterations.

 15

Returns:

Name Type Description
A Solution

class:Solution corresponding to the best candidate found.
Notes
  • Candidates are generated/updated by de_fit; this function defines the simulation and cost evaluation only.
  • If an integration attempt is unsuccessful, its cost is +inf.
Source code in frontx/finite.py 
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
def fit(  # noqa: PLR0913
    D: Callable[  # noqa: N803
        [float | jax.Array | np.ndarray[Any, Any]],
        float | jax.Array | np.ndarray[Any, Any],
    ],
    r1: float,
    t1: float,
    r: jax.Array | np.ndarray[Any, Any],
    t: jax.Array | np.ndarray[Any, Any],
    theta: jax.Array | np.ndarray[Any, Any],
    /,
    sigma: float | jax.Array | np.ndarray[Any, Any] = 1,
    *,
    i: jax.Array | np.ndarray[Any, Any],
    b: float | None = None,
    max_steps: int = 15,
) -> Solution:
    """Fit the finite-difference model to spatio-temporal observations.

    This routine searches over candidate models/parameters (through
    :func:`frontx._inverse.param.de_fit`) to minimize the mean squared error
    between the simulated field and observations ``theta(r, t)`` with optional
    per-sample weighting ``sigma``.

    Args:
        D: Initial diffusivity callable or parameterized model to start from.
        r1: Domain radius for the simulation.
        t1: Final time for the simulation.
        r: Radial coordinates of observations, shape ``(Nr,)``.
        t: Time coordinates of observations, shape ``(Nt,)``.
        theta: Observed values with shape broadcastable to ``(Nr, Nt)`` or
            exactly ``(Nr, Nt)``.
        sigma: Noise/weight (scalar or array broadcastable to ``theta``),
            used in the weighted MSE.
        i: Initial condition on the grid (length defines spatial resolution).
        b: Optional clamped value for the first grid node at initialization.
        max_steps: Maximum number of differential-evolution iterations.

    Returns:
        A :class:`Solution` corresponding to the best candidate found.

    Notes:
        - Candidates are generated/updated by ``de_fit``; this function defines
            the simulation and cost evaluation only.
        - If an integration attempt is unsuccessful, its cost is ``+inf``.
    """

    def candidate(
        D: Callable[  # noqa: N803
            [float | jax.Array | np.ndarray[Any, Any]],
            float | jax.Array | np.ndarray[Any, Any],
        ],
    ) -> Solution:
        return solve(D, r1, t1, i=i, b=b)

    def cost(sol: Solution) -> float:
        return jax.lax.cond(
            sol.result == RESULTS.successful,
            lambda: jnp.mean(((sol(r, t[:, jnp.newaxis]) - theta) / sigma) ** 2),
            lambda: jnp.inf,
        )

    return de_fit(candidate, cost, initial=D, max_steps=max_steps)