jaxprop.components.nozzle_model_solver module

class jaxprop.components.nozzle_model_solver.BVPSettings(num_points: 'int', rtol: 'jnp.ndarray', atol: 'jnp.ndarray', max_steps: 'int', jac_mode: 'str', verbose: 'bool', method: 'str' = 'GaussNewton', warmup_method: 'str' = 'LevenbergMarquardt', warmup_steps: 'int' = 0, solve_mode: 'str' = 'mach_crit')

Bases: Module

atol: Array

jac_mode: str

max_steps: int

method: str = 'GaussNewton'

num_points: int

rtol: Array

solve_mode: str = 'mach_crit'

verbose: bool

warmup_method: str = 'LevenbergMarquardt'

warmup_steps: int = 0

class jaxprop.components.nozzle_model_solver.IVPSettings(solver_name: str = 'Dopri5', adjoint_name: str = 'DirectAdjoint', number_of_points: int = 50, rtol: float = 1e-06, atol: float = 1e-06)

Bases: Module

Settings for marching initial value problem solvers.

adjoint_name: str = 'DirectAdjoint'

atol: float = 1e-06

number_of_points: int = 50

rtol: float = 1e-06

solver_name: str = 'Dopri5'

class jaxprop.components.nozzle_model_solver.NozzleParams(fluid: 'Any', geometry: 'Callable', p0_in: 'jnp.ndarray' = <factory>, d0_in: 'jnp.ndarray' = <factory>, D_in: 'jnp.ndarray' = <factory>, length: 'jnp.ndarray' = <factory>, roughness: 'jnp.ndarray' = <factory>, T_wall: 'jnp.ndarray' = <factory>, Ma_in: 'jnp.ndarray' = <factory>, Ma_low: 'jnp.ndarray' = <factory>, Ma_high: 'jnp.ndarray' = <factory>, Ma_target: 'jnp.ndarray' = <factory>, heat_transfer: 'jnp.ndarray' = <factory>, wall_friction: 'jnp.ndarray' = <factory>)

Bases: Module

D_in: Array

Ma_high: Array

Ma_in: Array

Ma_low: Array

Ma_target: Array

T_wall: Array

d0_in: Array

fluid: Any

geometry: Callable

heat_transfer: Array

length: Array

p0_in: Array

roughness: Array

wall_friction: Array

class jaxprop.components.nozzle_model_solver.ResidualParams(x: 'jnp.ndarray', Dx: 'jnp.ndarray', model: 'NozzleParams')

Bases: Module

Dx: Array

model: NozzleParams

x: Array

jaxprop.components.nozzle_model_solver.chebyshev_lobatto_basis(N: int, x1: float, x2: float)

Return:

x : (N+1,) physical nodes in [x1, x2] D : (N+1, N+1) differentiation s.t. (u_x)(x_i) ≈ sum_j D[i,j] * u(x_j)

Built from Trefethen’s formula. D acts on nodal values to produce nodal derivatives.

jaxprop.components.nozzle_model_solver.chebyshev_lobatto_interpolate(x_nodes, y_nodes, x_eval)

Evaluate the Chebyshev-Lobatto barycentric interpolant at one or more points.

This function is a thin wrapper around chebyshev_lobatto_interpolate_and_derivative that discards the derivative and returns only the interpolated value.

Parameters:

x_nodesarray_like, shape (N+1,)

The Chebyshev-Lobatto nodes in the physical domain [x_min, x_max].

y_nodesarray_like, shape (N+1,)

Function values at the Chebyshev-Lobatto nodes.

x_evalfloat or array_like

Point(s) in the domain where the interpolant should be evaluated.

Returns:

pfloat or ndarray

Interpolated value(s) p(x_eval).

jaxprop.components.nozzle_model_solver.chebyshev_lobatto_interpolate_and_derivative(x_nodes, y_nodes, x_eval)

Evaluate the Chebyshev-Lobatto barycentric interpolant and its derivative.

Parameters:

x_nodesarray_like, shape (N+1,)

The Chebyshev-Lobatto nodes in the physical domain [x_min, x_max].

y_nodesarray_like, shape (N+1,)

Function values at the Chebyshev-Lobatto nodes.

x_evalfloat or array_like

Point(s) in the domain where the interpolant and its derivative should be evaluated.

Returns:

pfloat or ndarray

Interpolated value(s) p(x_eval).

dpfloat or ndarray

First derivative p’(x_eval) with respect to x.

Notes

• Uses the barycentric interpolation formula, which is numerically stable even for high-degree polynomials.

• Correctly handles the case where x_eval coincides with one of the nodes, returning the exact nodal value and the exact derivative at that node.

• Works for scalar or vector x_eval.

jaxprop.components.nozzle_model_solver.compute_static_state(p0, d0, Ma, fluid)

solve h0 - h(p,s0) - 0.5 a(p,s0)^2 Ma^2 = 0 for p

jaxprop.components.nozzle_model_solver.evaluate_ode_rhs(x, z, args)

Vectorized full-model eval at all nodes from z=[u, ln(rho), ln(p)].

jaxprop.components.nozzle_model_solver.f64(value)

jaxprop.components.nozzle_model_solver.find_maximum_mach(x_nodes, y_nodes, newton_steps=50, rtol=1e-10, atol=1e-10)

Locate the single interior maximum of the Chebyshev-Lobatto interpolant.

Uses a smooth soft-argmax to pick an initial guess (avoids non-diff argmax), then runs a fixed number of Newton iterations on p’(x) = 0.

Parameters:

x_nodes(N+1,) array

Chebyshev-Lobatto nodes in the domain.

y_nodes(N+1,) array

Function values at the nodes (e.g., Mach number).

newton_stepsint, optional

Maximum Newton iterations. Default 50.

rtol, atolfloat, optional

Relative and absolute tolerances for Newton.

Returns:

x_starfloat

Location of the maximum in [x1, x2].

p_starfloat

Value of the interpolant at x_star.

jaxprop.components.nozzle_model_solver.get_residual_mach_critical(z, params: ResidualParams)

jaxprop.components.nozzle_model_solver.get_residual_mach_inlet(z, params: ResidualParams)

jaxprop.components.nozzle_model_solver.initialize_flowfield(num_points, params, Ma_min=0.1, Ma_max=0.2)

Generate an initial guess for the flowfield using a concave Mach number profile.

The Mach profile is defined as a symmetric parabola with its maximum (Ma_max) at the domain midpoint and its minimum (Ma_min) at both inlet and outlet. The corresponding velocity, density, and pressure fields are computed from the specified inlet stagnation state.

Parameters:

num_pointsint

Number of interior collocation points. The Chebyshev-Lobatto grid will contain num_points + 1 points in total.

paramsdict

Must contain:

p0_infloat

Inlet stagnation pressure.

d0_infloat

Inlet stagnation density.

fluidobject

Fluid property object

Ma_minfloat, optional

Minimum Mach number at the inlet and outlet. Default is 0.1.

Ma_maxfloat, optional

Maximum Mach number at the domain midpoint. Default is 0.5.

Returns:

z0ndarray, shape (3*(num_points+1),)

Initial guess vector at collocation points, concatenated as: [velocity, ln_density, ln_pressure].

jaxprop.components.nozzle_model_solver.replace_param(obj, field, value)

Return a copy of obj with a single field replaced.

jaxprop.components.nozzle_model_solver.solve_nozzle_model_collocation(initial_guess, params_model, params_solver)

Solves the collocation system using a warmup solver (e.g., Levenberg-Marquardt) followed by a main solver (e.g., Gauss-Newton). Returns the evaluated flowfield and the final solver result.

jaxprop.components.nozzle_model_solver.split_z(z, num_points)