Metadata-Version: 2.5
Name: flapjax-model-gen
Version: 0.1.66
Summary: Match strip-wise aerodynamic corrections (cl0/cla/cm0/cma) to a NASTRAN-ready twist distribution using flapjax.
Author-email: Ben Preston <b.preston23@imperial.ac.uk>
License-Expression: MIT
License-File: LICENSE
Keywords: aeroelasticity,flapjax,nastran,vortex lattice method
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: Operating System :: OS Independent
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Topic :: Scientific/Engineering
Classifier: Topic :: Scientific/Engineering :: Physics
Requires-Python: >=3.12
Requires-Dist: flapjax>=1.3.12
Requires-Dist: jax>=0.8.1
Requires-Dist: matplotlib>=3.8.4
Requires-Dist: pynastran>=1.3.4
Description-Content-Type: text/markdown

# flapjax-model-gen

Generate a [flapjax](https://github.com/ben-l-p/flapjax) aeroelastic model — structure, mass, aerodynamic mesh,
control surfaces, half-models — directly from NASTRAN input files, using
[pyNastran](https://github.com/SteveDoyle2/pyNastran) to parse bulk data. Also ships a companion workflow for
matching strip-wise aerodynamic corrections to a NASTRAN-ready twist distribution.

This is a thin, fast-moving companion package to `flapjax`, kept separate so it can be iterated on and released
independently.

## Building a flapjax model from NASTRAN inputs

`flapjax_model_gen.nastran` converts a NASTRAN aircraft model into a fully populated flapjax
`CoupledAeroelastic`, ready for `.reference_configuration()`/`.static_solve()`. It expects three input files,
matching how these models are actually organised:

- **Stiffness model** (`.bdf`, `GRID` + `CBEAM` + `PBEAM`, referencing `MAT1`) — `parse_stiffness_bdf` builds
  node coordinates, connectivity, per-element orientation (`y_vector`) and per-element 6x6 stiffness.
- **Mass model** (one or more `.nsb` files, `CONM2` entries only — e.g. one structural, one fuel) —
  `parse_mass_files` sums them into one lumped 6x6 mass matrix per node.
- **Aerodynamic mesh** (`.bdf`, `CAERO1` panels + `SPLINE1`/`SET1`) — `read_aero_bdf` reads it, and
  `parse_spline_surfaces` groups CAERO1 panels into continuous surfaces using the model's own `SPLINE1`/`SET1`
  entries, instead of requiring you to hand-list which CAERO eids and structural GRID ids belong together.

Everything derivable from these files is handled automatically: node topology, stiffness, lumped mass, which
CAERO1 panels form one continuous surface (including auto-combining a chordwise split, e.g. a main surface +
control surface, and a spanwise split, e.g. inboard/outboard discretisation), each surface's aerodynamic grid
rediscretized to sit exactly at the structural node positions, the aero-to-structure `dof_mapping`, and (for a
named control surface) the `grid_func` that deflects it. Everything else — flight condition, wake truncation,
boundary constraints, gravity, mirror symmetry — has no NASTRAN-derivable value and is a plain keyword argument
to `build_coupled_aeroelastic`, mirroring the underlying flapjax constructor it's forwarded to; see that
function's own docstring for the full parameter list.

```python
import jax.numpy as jnp
from flapjax.aero import ConstantFlowField
from flapjax_model_gen.nastran import (
    parse_stiffness_bdf, parse_mass_files, read_aero_bdf, parse_spline_surfaces, reorder_by_surfaces,
    SurfaceGridSpec, build_coupled_aeroelastic,
)

beam_model = parse_stiffness_bdf("stiffness.bdf")
aero_model = read_aero_bdf("aero.bdf")
surfaces = parse_spline_surfaces(aero_model)  # one SplineSurface per SPLINE1/SET1 grouping

beam_model = reorder_by_surfaces(beam_model, [s.node_ids for s in surfaces], span_axis=1)
m_lumped, _ = parse_mass_files(
    ["structure.nsb", "fuel.nsb"], node_ids=beam_model.node_ids, node_coords=beam_model.coords
)

wing = SurfaceGridSpec(name="wing", surface=surfaces[0], span_axis=1, m=8, m_star=20)
case, hinge_offsets = build_coupled_aeroelastic(
    beam_model,
    [wing],
    aero_model,
    m_lumped=m_lumped,
    dt=0.01,
    flowfield=ConstantFlowField(u_inf=jnp.array([30.0, 0.0, 0.0]), rho=1.225, relative_motion=True),
)
```

`hinge_offsets` is `{surface_name: hinge_offset}` — zero for every surface here (no `extra_strips`/`SurfaceGridSpec.n_span` used), but always safe to pass straight through as `twist_from_polar.find_twist_distribution`'s own `hinge_offset` (see "Matching aerodynamic corrections to a twist distribution" below).

### Multiple aerodynamic surfaces

A model with more than one surface (wing, tail, fin, ...) just needs more `SurfaceGridSpec`s — no manual
`grid_shapes`/`x0_aero`/`dof_mapping` bookkeeping. Each spec names its surface (an arbitrary label, just used to
cross-reference control surfaces below) and its own `span_axis` — e.g. a wing/horizontal tail spans in y, a
vertical tail in z; a single shared axis can't sort both correctly:

```python
surfaces = parse_spline_surfaces(aero_model)
# parse_spline_surfaces's own order is ascending SET1-id, not semantically meaningful -- identify each
# surface by its caero_ids/node_ids (from your model), not by position.
wing_surface, htail_surface, vtail_surface = surfaces

beam_model = reorder_by_surfaces(
    beam_model,
    [wing_surface.node_ids, htail_surface.node_ids, vtail_surface.node_ids],
    span_axis=[1, 1, 2],
)

specs = [
    SurfaceGridSpec(name="wing", surface=wing_surface, span_axis=1, m=8, m_star=20),
    SurfaceGridSpec(name="htail", surface=htail_surface, span_axis=1, m=6, m_star=20),
    SurfaceGridSpec(name="vtail", surface=vtail_surface, span_axis=2, m=6, m_star=20),
]
case, _hinge_offsets = build_coupled_aeroelastic(beam_model, specs, aero_model, m_lumped=m_lumped, dt=..., flowfield=...)
```

`reorder_by_surfaces` only groups/orders the *structural* side; it's a plain `reorder_nodes(beam_model,
node_order)` under the hood, callable directly if you want to reorder by something other than a spline
grouping. Only `SPLINE1` is supported (raises `NotImplementedError` for `SPLINE2`-`5`, which pyNastran does
parse into an object but this package doesn't handle). `SPLINE6`/`SPLINE7` aren't implemented by pyNastran
itself at all -- they land in `model.reject_cards`, not `model.splines` -- and a `SPLINE7`'s surface can't be
recovered reliably from the raw card either, so any surface using one of these is skipped entirely, with a
`UserWarning`, rather than guessed at or left to raise.

### Control surfaces

Name a control surface with its spanwise extent and a chordwise hinge fraction at *each* end (root and tip --
e.g. a tapered aileron sitting at a different percentage of chord at either end); `build_coupled_aeroelastic`
patches the surface's own reference grid with the exact straight hinge line those two points imply (every
spanwise station in between, not just the two ends), derives the hinge rotation axis directly from them (no
longer assumed to run along the surface's own `span_axis`), and wires it all into a `grid_func` that exposes
the name as a deflection-angle keyword:

```python
from flapjax_model_gen.nastran import ControlSurfaceRequest

wing = SurfaceGridSpec(name="wing", surface=wing_surface, span_axis=1, m=20)

case, _hinge_offsets = build_coupled_aeroelastic(
    beam_model,
    [wing],
    aero_model,
    control_surfaces={
        # 73% chord at the root end of the aileron, 68% at the tip -- pass the same value for both for an
        # unswept, constant-percent-chord hinge.
        "aileron": ControlSurfaceRequest(
            surface_name="wing", span_lo=8.0, span_hi=10.0, hinge_frac_root=0.73, hinge_frac_tip=0.68
        ),
    },
    m_lumped=m_lumped,
    dt=...,
    flowfield=...,
)

deflected = case.aero.grid_func(case.aero.zeta_b0, aileron=jnp.deg2rad(5.0))
```

### Half (symmetric) models

For a half-model mirrored about a plane (e.g. the aircraft's x-z plane, normal to the global y axis), reduce
the full-model pieces first: `half_model` drops the negative-side structure (keeping the mirror-plane nodes)
and halves the stiffness of any element lying entirely in the plane; `half_model_lumped_mass` does the nodal
equivalent for lumped mass; `restrict_surface_nodes` pares each surface's node list to match. Each surface's
own CAERO1 panels are then filtered automatically (dropping the negative/mirror-plane side) via
`SurfaceGridSpec.mirror_axis`:

```python
from flapjax_model_gen.nastran import half_model, half_model_lumped_mass, reindex_lumped_mass, restrict_surface_nodes

full_node_ids = beam_model.node_ids
beam_model = half_model(beam_model, mirror_axis=1)
m_lumped, _ = reindex_lumped_mass(m_lumped, full_node_ids, beam_model.node_ids)
m_lumped = half_model_lumped_mass(m_lumped, beam_model.coords, mirror_axis=1)
surfaces = [restrict_surface_nodes(s, beam_model.node_ids) for s in surfaces]

wing = SurfaceGridSpec(name="wing", surface=surfaces[0], span_axis=1, m=8, mirror_axis=1)
case, _hinge_offsets = build_coupled_aeroelastic(
    beam_model,
    [wing],
    aero_model,
    m_lumped=m_lumped,
    dt=...,
    flowfield=...,
    mirror_point=jnp.zeros(3),
    mirror_normal=jnp.array([0.0, 1.0, 0.0]),  # tells flapjax's UVLM about the symmetry plane
)
```

### Auxiliary spline nodes

Some models don't splice a `SPLINE1` directly to the beam's own structural nodes -- instead the `SET1` lists a
separate set of auxiliary leading/trailing-edge `GRID` points (defined in their own file, e.g.
`spline_nodes.bdf`), each rigidly tied to one real structural node via an `RBE2` (any `PLOTEL` entries, used
only for plotting, are ignored). `read_spline_nodes_bdf`/`resolve_spline_surfaces` translate a
`parse_spline_surfaces` result's node ids through those `RBE2`s into real structural node ids before calling
`reorder_by_surfaces` -- an LE/TE pair rigidly tied to the same station collapses to that one node id, and an
id not covered by any `RBE2` (i.e. the `SET1` already lists a real structural node) passes through unchanged,
so it's always safe to call:

```python
from flapjax_model_gen.nastran import read_spline_nodes_bdf, resolve_spline_surfaces

surfaces = resolve_spline_surfaces(surfaces, read_spline_nodes_bdf("spline_nodes.bdf"))
```

### Stitching surfaces together

Two surfaces built independently -- e.g. a wingtip device bolted onto a wing, each splined to its own `SET1`
-- can leave a seam at their shared edge: a small coordinate mismatch (each one's own `SplineSurface`
geometry doesn't quite agree with the other's at the junction), and no shared `dof_mapping` entry tying them
to the same structural frame. A `StitchRequest` forces a portion of one named `SurfaceGridSpec`'s grid/
`dof_mapping` to exactly match another's -- both fields, not just coordinates, so the join stays coincident
even after the structure deforms (the two columns now share the same structural frame, not just the same
undeformed reference position):

```python
from flapjax_model_gen.nastran import StitchRequest

# the winglet's own first column (its root) is forced onto the wing's own last column (its tip) -- both the
# local grid coordinates and the structural dof_mapping entry, overriding the winglet's own independent
# geometry/attachment at that one column.
case, _hinge_offsets = build_coupled_aeroelastic(
    beam_model,
    [wing, winglet],
    aero_model,
    stitch=[
        StitchRequest(surface_name="winglet", n_slice=slice(0, 1), other_surface_name="wing", other_n_slice=slice(-1, None)),
    ],
    m_lumped=m_lumped,
    dt=...,
    flowfield=...,
)
```

Applied in list order right after each surface's own grid/`dof_mapping` are built, but before any
`control_surfaces`/`incidence_surfaces` are wired in -- avoid overlapping a control/incidence surface's span
range with a stitched column range, since the span lookup (`control_surface_indices`/`incidence_indices`)
still uses the stitched surface's own *original* node coordinates, not the donor's.

`n_slice`/`other_n_slice` (spanwise columns) is the common case above, for a wingtip/root join. `StitchRequest`
also takes `m_slice`/`other_m_slice` (chordwise rows), for the opposite kind of seam: two surfaces that are
meant to share one continuous *chord*, but were built as two entirely separate `SurfaceGridSpec`s -- e.g. a
wing's trailing edge and a separately-built control surface's leading edge, both already splined to the *same*
structural nodes (so `dof_mapping` already agrees, it's purely a geometry seam). Giving `m_slice` alone narrows
the patch to a chordwise band across every column; giving both `m_slice` and `n_slice` narrows it to a 2D
sub-block. Whether `dof_mapping` is reassigned then depends on whether `m_slice` covers the *whole* chord: a
whole-chord stitch (the default -- `m_slice` left `None`) transfers `dof_mapping` too, same as above; a partial
one only patches geometry, and instead checks that the targeted columns are already tied to the same
structural node as the donor's, raising rather than silently leaving the two surfaces' frames mismatched:

```python
case, _hinge_offsets = build_coupled_aeroelastic(
    beam_model,
    [wing, flap],  # both splined to the same structural nodes
    aero_model,
    stitch=[
        StitchRequest(
            surface_name="flap", other_surface_name="wing", m_slice=slice(0, 1), other_m_slice=slice(-1, None)
        ),
    ],
    m_lumped=m_lumped,
    dt=...,
    flowfield=...,
)
```

A third kind of seam comes up with a `transpose`d surface (`SurfaceGridSpec.transpose`/`aero_mesh.transpose_panel`
-- for a surface whose own span/chord axes run the opposite way from what `span_axis` expects, e.g. a fuselage
modelled as a flat, lengthwise strip): a fuselage beam section meets a wing-like fuselage section where the
shared edge is one whole *column* on one side
(`n_slice`, spanning the whole chord) but one whole *row* on the other (`other_m_slice`, spanning every column) --
the same physical line, described along swapped axes. `other_transpose=True` transposes the donor's selected
block before copying it in, so `m_slice` is matched against `other_n_slice` instead of `other_m_slice` (and
`n_slice` against `other_m_slice`). `dof_mapping` is *never* reassigned for this kind of stitch, regardless of
`m_slice` -- the donor's selected row generally spans several of its own structural nodes, so there's no single
one a whole target column could sensibly adopt. `build_coupled_aeroelastic` also resolves the donor's block
through global coordinates (via each surface's own node positions) before re-localizing it against the target's
*own* node, rather than copying local offsets straight across -- needed precisely because those local offsets
are only meaningful relative to whichever node they came from, and here that's no longer a single shared node:

```python
case, _hinge_offsets = build_coupled_aeroelastic(
    beam_model,
    [front_fuselage, mid_fuselage],
    aero_model,
    stitch=[
        StitchRequest(
            surface_name="front_fuselage", other_surface_name="mid_fuselage",
            n_slice=slice(-1, None), other_m_slice=slice(0, 1), other_transpose=True,
        ),
    ],
    m_lumped=m_lumped,
    dt=...,
    flowfield=...,
)
```

### Checking panel/surface groupings visually

`plot_caero1_panels` plots each `CAERO1`'s quadrilateral outline in 3D (matplotlib), labelled by its element
id, and `plot_beam_model` does the equivalent for the structural side -- both useful for sanity-checking
`span_axis`/`mirror_axis` choices and `parse_spline_surfaces` groupings before wiring up the full model:

```python
from flapjax_model_gen.nastran import parse_caero1_panels, plot_beam_model, plot_caero1_panels

ax = plot_caero1_panels(parse_caero1_panels(aero_model), surfaces=surfaces)  # colours each panel by its surface
ax.figure.show()  # or ax.figure.savefig("panels.png")

ax = plot_beam_model(beam_model)
ax.figure.show()
```

### Lower-level building blocks

`build_coupled_aeroelastic`/`SurfaceGridSpec` are a convenience layer over functions you can call directly for
finer control: `build_surface_grid` (one surface's aero grid, given its `SplineSurface` and node coordinates)
is itself `combine_surface_panels` (auto-detect/merge a chordwise split into one full-chord panel, order
root-to-tip) + `build_local_grid` (rediscretize). `control_surface_indices` computes the `m_slice`/`n_slice`/
`hinge_axis` flapjax's own `add_control_surface` needs, if you're assembling a `grid_func` by hand.
`dof_mapping_for_surfaces` resolves a surface's node ids to structural indices, if you're assembling a `UVLM`
by hand rather than through `build_coupled_aeroelastic`. `stitch_surface_grids` is the same patch
`StitchRequest` applies (see "Stitching surfaces together" above), for use directly on a grid/`dof_mapping`
pair if you're assembling `UVLM` by hand.

### What's deliberately out of scope / assumed

Each backed by a clear error or warning rather than a silent wrong answer if violated:

- Only `CAERO1`, `CBEAM`/`PBEAM` (not `PBEAML`/`PBCOMP`), `MAT1`, and `SPLINE1` are supported.
- `CBEAM`'s `OFFT` must be the default `'GGG'`, and end offsets (`WA`/`WB`) must be zero — no offset beam ends.
- `CONM2`'s `CID` must be the basic coordinate system (`0` or `-1`); mass is taken entirely from `CONM2` mass
  files, not `PBEAM`'s `NSM` or any distributed cross-section mass.
- A tapered `PBEAM` (multiple stations) is collapsed to one constant cross-section per element by averaging its
  station values — flapjax elements are constant-property, so this is an approximation for a genuinely tapered
  element, not a limitation you can configure around.
- `PBEAM`'s `I12` (bend-bend coupling) is not incorporated into `k_cs` — a warning fires if it's nonzero, since
  results will be approximate for that element.
- The aerodynamic grid's spanwise stations are placed exactly at the given structural node positions (flapjax's
  `dof_mapping` ties each aero column rigidly to one structural node, with no separate spline layer like
  NASTRAN's beam splines) — only the chordwise discretization (`m`/`x_over_c`) is a free choice.
- A control surface's spanwise extent must land exactly on existing structural nodes — there's no clipping at
  an arbitrary span location (its chordwise hinge fraction, at each end, has no such restriction — see
  "Control surfaces" above).
- There's no NASTRAN `SPC`/`SPC1` parsing, so boundary constraints (a clamped root, a multibody hinge) are
  built as flapjax constraint objects by hand and passed to `build_coupled_aeroelastic(constraints=...)`.
- The axis convention (`CBEAM`'s orientation vector = flapjax's `y_vector` directly; PBEAM `I1`→bending about
  local z, `I2`→about local y) was verified against the NASTRAN QRG and cross-checked against flapjax's own
  beam code, not assumed — see the docstring in `nastran/stiffness.py` for the derivation.

## Matching aerodynamic corrections to a twist distribution

A secondary workflow: given a strip-wise aerodynamic correction (`cl0`, `cla`, `cm0`, `cma`, sampled against a
*rigid* flight shape from CFD, with no knowledge of elastic deformation) already attached to a flapjax `UVLM`
surface — wired in via `polar_data`/`polar_function`, machinery that already lives in `flapjax` itself, see its
`models/cantilever_wing/polar_correction.ipynb` tutorial — find the per-station twist to bake into geometry so
that, once the jig-shape structure deforms to its self-consistent flight shape, that surface's own circulation-
derived local angle of attack reports back the correction's own `cl0`/`cm0` (and its affine extrapolation at a
second AoA) everywhere along the span, without needing the correction to run at solve time.

Everything this method needs lives in one file, `twist_from_polar.py` — see its module docstring (and
`find_twist_distribution`'s own) for the full outer-fixed-point-loop algorithm:

- `apply_twist_to_grid`/`make_twisted_grid` — rotate an already-built (or freshly-built) local aero grid by an
  independent twist angle per spanwise station. For a column built by `aero_mesh.insert_strips_at_stations`/
  `uniform_spanwise_grid` (a strip rigidly tied to a different column's structural node), pass that column's own
  `hinge_offset` (`station_coords - attach_coords`) so it twists about its own true span position rather than
  about its attach node — otherwise a large offset between the two introduces a non-smooth lever-arm error.
  `build_coupled_aeroelastic` already computes this for every surface and hands it back as its own
  `hinge_offset_by_name` return value, so it rarely needs recomputing by hand — see the example below.
- `TwistMatchResult`/`match_twist` — the one piece with no existing flapjax equivalent: a JAX-native
  Levenberg-Marquardt solver (autodiff Jacobian, `jax.lax.while_loop`) for the per-outer-iteration twist fit.
- `find_twist_distribution` — wires the above into a FIXED-length (`jax.lax.fori_loop`) outer loop against a
  real per-outer-iteration FSI re-solve, so the twist found accounts for the flight shape's own elastic
  response to it. Deliberately has no Python-level control flow (no convergence check, no early exit, no
  printing/plotting), so it's trace-safe end to end and can itself be `jax.jit`-ed — directly, or via
  `functools.partial` binding every argument except `coupled` first — to run against many aircraft mass/
  velocity cases without retracing per case, since `CoupledAeroelastic`/`UVLM`/`BeamStructure` are themselves
  `flapjax` pytrees with mass/flight-condition data as dynamic leaves.
- `plot_twist_match_result` — save `cl`/`cm`, twist and vertical-deflection diagnostic plots for a
  `find_twist_distribution` result, split out into its own, deliberately non-jit-able function (real
  matplotlib/file I/O) so `find_twist_distribution` doesn't have to be. Call it eagerly, after the fact, on
  whichever case's result you want to inspect.
- `find_twist_and_slope_multiplier_distribution` — as `find_twist_distribution`, but also fitting a per-strip
  multiplier onto `cla`/`cma` (shared between the two) JOINTLY with twist, every outer iteration, to close
  whatever `cl`/`cm` residual twist alone can't: twist can only shift each strip's local alpha, never rescale
  the slope the model actually responds with, so whenever that slope doesn't match the polar's own, twist
  alone settles for a compromise at both AoA samples instead of zeroing either. Fitting the multiplier alone,
  as a separate step AFTER twist has already converged, was tried first and doesn't work: that twist
  converged against the polar's ORIGINAL, unscaled slope, so once the multiplier rescales it, both the FSI
  equilibrium and the twist that's actually optimal for it have silently moved out from under the "converged"
  twist. Returns a `TwistAndSlopeMultiplier(twist, slope_multiplier)` — each outer iteration rebuilds the
  scaled polar from the current multiplier, re-solves the two AoA samples' anchor pass against THAT polar,
  jointly fits twist and multiplier, then re-solves the FSI equilibrium at the new twist AND new scaled polar
  together.

Both `find_twist_distribution` and `find_twist_and_slope_multiplier_distribution` take `aoa_sample_mode`,
controlling how the second AoA sample (`+delta_aoa`) is obtained. `"rotate"` (default) rigidly rotates
(`rotate_hg`) the already-converged zero-AoA flight shape — cheap, but can't capture how the structure itself
would respond differently if it were actually pitched up. `"resolve"` instead rebuilds the whole model at
`orientation_euler = (0, delta_aoa, 0)` (the same whole-airframe-AoA convention flapjax's own example models
use) and lets the FSI loop converge fresh at that pitched orientation — more physically complete, but roughly
doubles the FSI cost of every outer iteration. Assumes the model's own baseline `orientation_euler` is the
identity, true by construction for this whole workflow's jig-shape-at-0-degrees contract.

Two smaller, standalone pieces elsewhere in this package support it: `linear_polar.py` (`LinearPolar`/
`linear_polar_function`, a `PolarFunction` evaluating `cl = cl0 + cla * alpha`, `cm = cm0 + cma * alpha`,
`cd = 0`, so `cl0`/`cla`/`cm0`/`cma` data drops straight into `UVLM(polar_data=..., polar_function=...)`) and
`rigid.py` (`rotate_hg`, rigidly rotating a structure's reference SE(3) frames — used to represent the second
AoA sample as a rotation of the geometry rather than the freestream).

```python
import jax.numpy as jnp
from flapjax.aero.flowfields import ConstantFlowField
from flapjax_model_gen import LinearPolar, find_twist_distribution, linear_polar_function
from flapjax_model_gen.nastran import build_coupled_aeroelastic

# Your real NASTRAN parsing -- see "Building a flapjax model from NASTRAN inputs" above (parse_stiffness_bdf +
# reorder_by_surfaces, read_aero_bdf + parse_spline_surfaces, parse_mass_files). The structure must be in jig
# (unloaded, as-manufactured) shape.
beam_model = ...  # NastranBeamModel
surfaces = [...]  # Sequence[SurfaceGridSpec] -- must include the surface being twisted
aero_model = ...  # BDF, from read_aero_bdf
m_lumped = ...  # Array | None

surface_name = "wing"  # which SurfaceGridSpec.name to match twist on
delta_aoa = jnp.deg2rad(1.0)  # second sample's rigid rotation, and the polar's assumed affine extent

# Strip-wise correction on surface_name, from CFD on the rigid flight shape, with no knowledge of the structure.
polar = LinearPolar(cl0=jnp.array([...]), cla=jnp.array([...]), cm0=jnp.array([...]), cma=jnp.array([...]))

coupled, hinge_offsets = build_coupled_aeroelastic(
    beam_model,
    surfaces,
    aero_model,
    m_lumped=m_lumped,
    dt=0.01,
    flowfield=ConstantFlowField(u_inf=jnp.array([30.0, 0.0, 0.0]), rho=1.225, relative_motion=True),
    polar_data=[polar],
    polar_function=[linear_polar_function],
)
surf_idx = next(i for i, spec in enumerate(surfaces) if spec.name == surface_name)

prescribed_dofs = [...]  # e.g. clamped-root DOFs

# hinge_offsets[surface_name] is zero everywhere unless surface_name uses extra_strips/SurfaceGridSpec.n_span
# -- always safe to pass through regardless, so a strip twists about its own true span position rather than
# about whichever node it's rigidly tied to.
twist_distribution = find_twist_distribution(
    coupled, surf_idx, delta_aoa, prescribed_dofs, hinge_offset=hinge_offsets[surface_name]
)
```

`find_twist_distribution` also takes `fixed_zero_stations` — hold specific stations' twist at exactly zero
while still fitting the full-span `cl`/`cm` targets, e.g. a region you don't want washout applied to. To sanity-
check the fit (e.g. on one representative case out of a batch), save diagnostic plots separately:

```python
from flapjax_model_gen import plot_twist_match_result

plot_twist_match_result(
    coupled, surf_idx, delta_aoa, twist_distribution, prescribed_dofs,
    hinge_offset=hinge_offsets[surface_name],
)
```

To also close whatever `cl`/`cm` residual remains after twist alone (see above), use the combined routine
instead and build the corrected `LinearPolar` to actually use downstream from its `slope_multiplier`:

```python
from flapjax_model_gen import find_twist_and_slope_multiplier_distribution

result = find_twist_and_slope_multiplier_distribution(
    coupled, surf_idx, delta_aoa, prescribed_dofs, hinge_offset=hinge_offsets[surface_name]
)
corrected_polar = polar._replace(cla=polar.cla * result.slope_multiplier, cma=polar.cma * result.slope_multiplier)

# plot_twist_match_result's slope_multiplier argument makes the diagnostic plot reflect the corrected model
# (rescaled cla/cma) rather than the original polar's -- the closed-form targets themselves never change.
plot_twist_match_result(
    coupled, surf_idx, delta_aoa, result.twist, prescribed_dofs, slope_multiplier=result.slope_multiplier,
    hinge_offset=hinge_offsets[surface_name],
)
```

Pass `aoa_sample_mode="resolve"` to either routine for a genuine FSI re-solve at `+delta_aoa` instead of the
cheap rigid-rotation default (see above):

```python
twist_distribution = find_twist_distribution(
    coupled, surf_idx, delta_aoa, prescribed_dofs, aoa_sample_mode="resolve",
    hinge_offset=hinge_offsets[surface_name],
)
```

### Running against many mass/velocity cases

Since `find_twist_distribution` has no Python-level control flow, it can be `jax.jit`-ed wholesale — bind
everything except `coupled` with `functools.partial` first, then call the jitted function once per case:

```python
import functools
import jax

twist_fn = jax.jit(
    functools.partial(
        find_twist_distribution, surf_idx=surf_idx, delta_aoa=delta_aoa,
        prescribed_dofs=prescribed_dofs, n_outer=15,
    )
)
twist_distributions = [twist_fn(coupled) for coupled in many_coupled_cases]  # one `coupled` per mass/velocity
```

The first call traces and compiles; every later call with a different `coupled` (same surfaces/topology, just
different mass/flight-condition values) reuses that compiled executable instead of retracing.

## Development

```bash
uv sync --dev
uv run pytest
uv run ruff check src
```

Tests are deliberately lightweight (geometry/algebra invariants, synthetic NASTRAN fixtures, and a synthetic
least-squares problem for the optimizer) rather than full aeroelastic solves, so the suite stays fast to
iterate against.
