Pazy wing static deflection sweep¶

Run a grid of Pazy cases with varying AoA and dynamic pressure in parallel using JAX's pmap. Using the provided discretisation, this can run 1024 cases in the order of a minute on an M2 MacBook Air. The results are plotted as a heatmap and line plots for each AoA.

Set XLA flags to set one host device per case. This must run before JAX is imported anywhere in the kernel, otherwise JAX will initialise the default device count and ignore these flags.

In [ ]:

import os

n_vel = 32  # number of velocities
n_alpha = 32  #number of angles of attack
n_case = n_vel * n_alpha
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 n_vel = 32 # number of velocities n_alpha = 32 #number of angles of attack n_case = n_vel * n_alpha 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

In [ ]:

import time
from typing import cast

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

from flapjax.aero.flowfields import ConstantFlowField
from flapjax.coupled import CoupledAeroelastic
from flapjax.models.pazy.straight.pazy_wing import generate_pazy_wing
from flapjax.utils.print_utils import set_verbosity

import time from typing import cast import jax from jax import numpy as jnp from matplotlib import pyplot as plt from matplotlib.colors import Normalize from flapjax.aero.flowfields import ConstantFlowField from flapjax.coupled import CoupledAeroelastic from flapjax.models.pazy.straight.pazy_wing import generate_pazy_wing from flapjax.utils.print_utils import set_verbosity

JAX configuration and verbosity

In [ ]:

set_verbosity("silent")  # suppress console prints as there would be a lot for 1k+ cases

# ensure that the XLA flags have worked
assert jax.device_count() == n_case, (
    f"Expected {n_case} devices, got {jax.device_count()}"
)

set_verbosity("silent") # suppress console prints as there would be a lot for 1k+ cases # ensure that the XLA flags have worked assert jax.device_count() == n_case, ( f"Expected {n_case} devices, got {jax.device_count()}" )

Build the AoA / velocity sweep and stack the individual CoupledAeroelastic cases into a single pytree with a leading device axis.

In [ ]:

rho = 1.225  # freestream density
u_inf_vec = jnp.linspace(1.0, 80.0, n_vel)  # freestream velocity sweep
alphas = jnp.deg2rad(jnp.linspace(-5.0, 15.0, n_alpha))  # angle of attack sweep

# create a list of all the cases. This means creating an instance of the Pazy wing for each condition.
cases = []
for u_inf in u_inf_vec:
    u_inf_mag = u_inf
    for alpha in alphas:
        cases.append(
            generate_pazy_wing(
                flowfield=ConstantFlowField(
                    u_inf=jnp.array(
                        (
                            u_inf_mag,
                            0.0,
                            0.0,
                        )
                    ),
                    rho=rho,
                    relative_motion=True,
                ),
                aoa=alpha,
                m=12,
                node_multiplier=2,
            )
        )

# stack the list of cases into a JAX pytree so that we can map across them
stacked_case = jax.tree_util.tree_map(lambda *xs: jnp.stack(xs), *cases)

rho = 1.225 # freestream density u_inf_vec = jnp.linspace(1.0, 80.0, n_vel) # freestream velocity sweep alphas = jnp.deg2rad(jnp.linspace(-5.0, 15.0, n_alpha)) # angle of attack sweep # create a list of all the cases. This means creating an instance of the Pazy wing for each condition. cases = [] for u_inf in u_inf_vec: u_inf_mag = u_inf for alpha in alphas: cases.append( generate_pazy_wing( flowfield=ConstantFlowField( u_inf=jnp.array( ( u_inf_mag, 0.0, 0.0, ) ), rho=rho, relative_motion=True, ), aoa=alpha, m=12, node_multiplier=2, ) ) # stack the list of cases into a JAX pytree so that we can map across them stacked_case = jax.tree_util.tree_map(lambda *xs: jnp.stack(xs), *cases)

Per-case solve function. Returns the mid-chord tip Z displacement normalised by the semi-span.

In [ ]:

def solve(case_: CoupledAeroelastic):
    static_sol = case_.static_solve(
        prescribed_dofs=tuple(range(6)),  # prescribe the 6 degress of freedom at the root (clamped)
        horseshoe=True,  # use a horseshoe wake as this means only one wake panel per strip
    )

    # compute the mid-chord deflection from the aerodynamic grid coordinate output zeta_b
    return (
            0.5
            * (
                    static_sol.aero.zeta_b[0][0, -1, 2]
                    + static_sol.aero.zeta_b[0][-1, -1, 2]
            )
            / 0.55
    )


parallel_func = jax.pmap(solve)  # create a function that maps the solve function across the stacked cases

def solve(case_: CoupledAeroelastic): static_sol = case_.static_solve( prescribed_dofs=tuple(range(6)), # prescribe the 6 degress of freedom at the root (clamped) horseshoe=True, # use a horseshoe wake as this means only one wake panel per strip ) # compute the mid-chord deflection from the aerodynamic grid coordinate output zeta_b return ( 0.5 * ( static_sol.aero.zeta_b[0][0, -1, 2] + static_sol.aero.zeta_b[0][-1, -1, 2] ) / 0.55 ) parallel_func = jax.pmap(solve) # create a function that maps the solve function across the stacked cases

Run the analysis for all cases in parallel and time the execution. The results are reshaped into a grid for plotting.

In [ ]:

t_start = time.time()
tip_z = parallel_func(stacked_case)  # evaluate the function
jax.block_until_ready(tip_z)
t_end = time.time()

print("Total time: ", t_end - t_start)
print("Time per case: ", (t_end - t_start) / n_case)

tip_z_grid = tip_z.reshape(n_vel, n_alpha).T  # [n_alpha, n_vel] for plotting

t_start = time.time() tip_z = parallel_func(stacked_case) # evaluate the function jax.block_until_ready(tip_z) t_end = time.time() print("Total time: ", t_end - t_start) print("Time per case: ", (t_end - t_start) / n_case) tip_z_grid = tip_z.reshape(n_vel, n_alpha).T # [n_alpha, n_vel] for plotting

Heatmap of tip Z displacement across AoA and freestream velocity.

In [ ]:

fig, ax = plt.subplots()
mesh = ax.pcolormesh(
    u_inf_vec, jnp.rad2deg(alphas), tip_z_grid, shading="auto", cmap="bwr"
)
fig.colorbar(mesh, ax=ax, label="Tip Z Displacement [z/b]")
ax.set_xlabel("Velocity [m/s]")
ax.set_ylabel("Angle of Attack [deg]")
ax.set_title("Pazy Tip Vertical Displacement")
plt.show()

fig, ax = plt.subplots() mesh = ax.pcolormesh( u_inf_vec, jnp.rad2deg(alphas), tip_z_grid, shading="auto", cmap="bwr" ) fig.colorbar(mesh, ax=ax, label="Tip Z Displacement [z/b]") ax.set_xlabel("Velocity [m/s]") ax.set_ylabel("Angle of Attack [deg]") ax.set_title("Pazy Tip Vertical Displacement") plt.show()

Plots of displacement versus AoA, colour-coded by angle of attack.

In [ ]:

alphas_deg = jnp.rad2deg(alphas)

# noinspection PyArgumentList
norm = cast(Normalize, Normalize(vmin=float(alphas_deg.min()), vmax=float(alphas_deg.max())))

cmap = plt.get_cmap("viridis")

fig, ax = plt.subplots()
for i_aoa, aoa in enumerate(alphas_deg):
    ax.plot(u_inf_vec, tip_z_grid[i_aoa], color=cmap(norm(float(aoa))))
sm = plt.cm.ScalarMappable(norm=norm, cmap=cmap)
fig.colorbar(sm, ax=ax, label="Angle of Attack [deg]")
ax.set_xlabel("Velocity [m/s]")
ax.set_ylabel("Tip Z Displacement [z/b]")
ax.set_title("Pazy Wing Tip Z Displacement")
plt.show()

alphas_deg = jnp.rad2deg(alphas) # noinspection PyArgumentList norm = cast(Normalize, Normalize(vmin=float(alphas_deg.min()), vmax=float(alphas_deg.max()))) cmap = plt.get_cmap("viridis") fig, ax = plt.subplots() for i_aoa, aoa in enumerate(alphas_deg): ax.plot(u_inf_vec, tip_z_grid[i_aoa], color=cmap(norm(float(aoa)))) sm = plt.cm.ScalarMappable(norm=norm, cmap=cmap) fig.colorbar(sm, ax=ax, label="Angle of Attack [deg]") ax.set_xlabel("Velocity [m/s]") ax.set_ylabel("Tip Z Displacement [z/b]") ax.set_title("Pazy Wing Tip Z Displacement") plt.show()