Parallelised Gust Response of a Simple HALE Aircraft¶

This case demonstrates parallelisation of the simple_hale model using jax.pmap. It performs a parameter sweep over gust length (and optionally gust amplitude), solving one dynamic aeroelastic response per case in parallel and plotting the root strains. On a MacBook Air M2, this runs at roughly 7 seconds per dynamic case when parallelised over 32 cases.

Configure JAX host device count¶

jax.pmap requires one device per case. To parallelise the sweep on CPU, we force XLA to expose n_case host devices via the XLA_FLAGS environment variable. This must be set before importing jax.

In [ ]:

import os

# number of cases for the sweep
n_gust_length: int = 32
n_gust_amplitude: int = 1
n_case: int = n_gust_length * n_gust_amplitude

# force XLA to give one host device per case
# it isn't essential to have one core/thread per case, but it will likely be faster than having multiple cases share a single core/thread
os.environ["XLA_FLAGS"] = (
        os.environ.get("XLA_FLAGS", "")
        + f" --xla_force_host_platform_device_count={n_case} --xla_cpu_multi_thread_eigen=true"
).strip()

import os # number of cases for the sweep n_gust_length: int = 32 n_gust_amplitude: int = 1 n_case: int = n_gust_length * n_gust_amplitude # force XLA to give one host device per case # it isn't essential to have one core/thread per case, but it will likely be faster than having multiple cases share a single core/thread os.environ["XLA_FLAGS"] = ( os.environ.get("XLA_FLAGS", "") + f" --xla_force_host_platform_device_count={n_case} --xla_cpu_multi_thread_eigen=true" ).strip()

Imports¶

JAX imports must occur after XLA_FLAGS is set.

In [ ]:

import time
from copy import deepcopy
from typing import cast

import jax
import matplotlib.pyplot as plt
from jax import numpy as jnp
from matplotlib.colors import Normalize

from flapjax.aero.flowfields import ConstantFlowField, OneMinusCosineFlowField
from flapjax.coupled import CoupledAeroelastic
from flapjax.models.simple_hale.simple_hale import generate_simple_hale

assert jax.device_count() == n_case, (
    f"Expected {n_case} devices, got {jax.device_count()}"
)

import time from copy import deepcopy from typing import cast import jax import matplotlib.pyplot as plt from jax import numpy as jnp from matplotlib.colors import Normalize from flapjax.aero.flowfields import ConstantFlowField, OneMinusCosineFlowField from flapjax.coupled import CoupledAeroelastic from flapjax.models.simple_hale.simple_hale import generate_simple_hale assert jax.device_count() == n_case, ( f"Expected {n_case} devices, got {jax.device_count()}" )

Case parametrisation¶

Define the free-stream velocity, the physical simulation time, and the sweep ranges. If n_gust_amplitude > 1, amplitudes will also be swept.

In [ ]:

u_inf_mag: float = 10.0  # free-stream velocity magnitude
physical_time: float = 10.0

gust_lengths = jnp.linspace(0.1 * u_inf_mag, 5.0 * u_inf_mag, n_gust_length)
gust_amplitudes = (
    jnp.linspace(0.0, 0.5, n_gust_amplitude) if n_gust_amplitude > 1 else [0.3]
)

u_inf_mag: float = 10.0 # free-stream velocity magnitude physical_time: float = 10.0 gust_lengths = jnp.linspace(0.1 * u_inf_mag, 5.0 * u_inf_mag, n_gust_length) gust_amplitudes = ( jnp.linspace(0.0, 0.5, n_gust_amplitude) if n_gust_amplitude > 1 else [0.3] )

Build the trimmed base aircraft¶

Construct a HALE aircraft in a constant free-stream and trim it. This is used as the initial condition for all gust cases. Note that the compilation can take ~1 minute.

In [ ]:

const_flowfield = ConstantFlowField(
    u_inf=jnp.array((u_inf_mag, 0.0, 0.0)),
    rho=1.225,
    relative_motion=True,
)

base_hale = generate_simple_hale(flowfield=const_flowfield, sigma_wing=1.5)
n_tstep = int(physical_time / float(base_hale.aero.dt)) + 1

static_sol, trim_vars = base_hale.trim(
    prescribed_dofs=jnp.arange(6),
    zero_force_dofs=(0, 2, 4),  # balance drag, lift, and pitching moment
    trim_cs="elevator",
    thrust_nodes="thrust",
    trim_orientation="y",
    method="finite_difference",
    horseshoe=False,
)

const_flowfield = ConstantFlowField( u_inf=jnp.array((u_inf_mag, 0.0, 0.0)), rho=1.225, relative_motion=True, ) base_hale = generate_simple_hale(flowfield=const_flowfield, sigma_wing=1.5) n_tstep = int(physical_time / float(base_hale.aero.dt)) + 1 static_sol, trim_vars = base_hale.trim( prescribed_dofs=jnp.arange(6), zero_force_dofs=(0, 2, 4), # balance drag, lift, and pitching moment trim_cs="elevator", thrust_nodes="thrust", trim_orientation="y", method="finite_difference", horseshoe=False, )

Convert to a dynamic free-flying initial condition¶

For the dynamic gust runs we switch from a static-aircraft/dynamic-freestream formulation to a dynamic-aircraft/static-freestream one, so that the aircraft is free to translate and rotate through the (stationary) gust.

In [ ]:

dynamic_init = base_hale.initialise_dynamic(static_case=static_sol, prescribed_dofs=())

dynamic_init = base_hale.initialise_dynamic(static_case=static_sol, prescribed_dofs=())

Build the stacked parallel case¶

For each gust length/amplitude combination, we build a copy of the aircraft with a OneMinusCosineFlowField gust flowfield. The gust is positioned upstream of the aircraft by setting the gust offset with gust_x0, such so that all cases begin with the aircraft outside the gust. The individual case pytrees are then stacked along a leading axis, producing a single pytree that jax.pmap will split across devices.

In [ ]:

cases = []
for gust_length in gust_lengths:
    for gust_amplitude in gust_amplitudes:
        gust_flowfield = OneMinusCosineFlowField(
            u_inf=jnp.array((u_inf_mag, 0.0, 0.0)),
            rho=1.225,
            relative_motion=False,
            gust_length=gust_length,
            gust_amplitude=gust_amplitude * u_inf_mag,
            gust_x0=jnp.array((-2.0 * gust_length - 10.0, 0.0, 0.0)),
        )
        new_case = deepcopy(base_hale)
        new_case.aero.flowfield = gust_flowfield
        cases.append(new_case)

stacked_case = jax.tree_util.tree_map(lambda *xs: jnp.stack(xs), *cases)

cases = [] for gust_length in gust_lengths: for gust_amplitude in gust_amplitudes: gust_flowfield = OneMinusCosineFlowField( u_inf=jnp.array((u_inf_mag, 0.0, 0.0)), rho=1.225, relative_motion=False, gust_length=gust_length, gust_amplitude=gust_amplitude * u_inf_mag, gust_x0=jnp.array((-2.0 * gust_length - 10.0, 0.0, 0.0)), ) new_case = deepcopy(base_hale) new_case.aero.flowfield = gust_flowfield cases.append(new_case) stacked_case = jax.tree_util.tree_map(lambda *xs: jnp.stack(xs), *cases)

Parallel dynamic solve¶

solve runs a single dynamic case and returns the transient rotational (torsion + in/out-of-plane bending) strains at the root element. The use of jax.pmap maps the function over the leading axis of stacked_case, so all cases execute concurrently. The reported time includes JIT compilation on the first call, which in some cases can be significant. Subsequent calls will be faster.

The choice to only return the rotational strains is arbitrary. However, returning just a small subset of the full solution is often a good idea to reduce memory, as for instance this solution does not need to store the history of the aerodynamic states.

Note that we here leave console prints on (which can be disabled using set_verbosity()), which shows the progress of all cases. These prints are often of limited use in parallel, as they can be interleaved and hard to read. As we here use jax.pmap instead of jax.vmap, some cases may finish considerably before others.

In [ ]:

def solve(hale_: CoupledAeroelastic):
    # need to avoid type inspection as there is a bug in Paraview
    # noinspection PyTypeChecker
    dynamic_sol = hale_.dynamic_solve(
        init_case=dynamic_init, prescribed_dofs=(), n_tstep=n_tstep
    )
    # root strains: element 0, rotational strains only
    return dynamic_sol.structure.eps[:, 0, 3:]


t_start = time.time()
root_strains = jax.pmap(solve)(stacked_case)
jax.block_until_ready(root_strains)
t_end = time.time()

print(f"Time taken for {n_case} cases: {t_end - t_start:.2f} seconds")
print(f"Time per case: {(t_end - t_start) / n_case:.2f} seconds")

def solve(hale_: CoupledAeroelastic): # need to avoid type inspection as there is a bug in Paraview # noinspection PyTypeChecker dynamic_sol = hale_.dynamic_solve( init_case=dynamic_init, prescribed_dofs=(), n_tstep=n_tstep ) # root strains: element 0, rotational strains only return dynamic_sol.structure.eps[:, 0, 3:] t_start = time.time() root_strains = jax.pmap(solve)(stacked_case) jax.block_until_ready(root_strains) t_end = time.time() print(f"Time taken for {n_case} cases: {t_end - t_start:.2f} seconds") print(f"Time per case: {(t_end - t_start) / n_case:.2f} seconds")

Plot the root strains¶

For each strain component (torsion, in-plane bending, out-of-plane bending), plot the time history over all cases, coloured by gust length.

In [ ]:

t = jnp.arange(n_tstep) * base_hale.aero.dt
cmap = plt.get_cmap("viridis")

# noinspection PyArgumentList
norm = cast(Normalize, Normalize(vmin=float(gust_lengths[0]), vmax=float(gust_lengths[-1])))

titles = ["Torsional", "In-plane bending", "Out-of-plane bending"]

for i_dir in range(3):
    fig, ax = plt.subplots()
    for i in range(n_case):
        ax.plot(
            t,
            root_strains[i, :, i_dir],
            color=cmap(i / max(n_case - 1, 1)),
        )
    ax.set_xlabel("Time [s]")
    ax.set_ylabel("Root Strain [m/m]")
    ax.set_title(f"{titles[i_dir]} strain at root")

    sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
    fig.colorbar(sm, ax=ax, label="Gust length [m]")
    plt.show()

t = jnp.arange(n_tstep) * base_hale.aero.dt cmap = plt.get_cmap("viridis") # noinspection PyArgumentList norm = cast(Normalize, Normalize(vmin=float(gust_lengths[0]), vmax=float(gust_lengths[-1]))) titles = ["Torsional", "In-plane bending", "Out-of-plane bending"] for i_dir in range(3): fig, ax = plt.subplots() for i in range(n_case): ax.plot( t, root_strains[i, :, i_dir], color=cmap(i / max(n_case - 1, 1)), ) ax.set_xlabel("Time [s]") ax.set_ylabel("Root Strain [m/m]") ax.set_title(f"{titles[i_dir]} strain at root") sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm) fig.colorbar(sm, ax=ax, label="Gust length [m]") plt.show()