# from https://github.com/alan-turing-institute/mogp-emulator/blob/main/mogp_emulator/demos/projectile.py
import numpy as np
import torch
from scipy.integrate import solve_ivp
from autoemulate.core.types import NumpyLike, TensorLike
from autoemulate.simulations.base import Simulator
[docs]
class Projectile(Simulator):
"""
Simulator of projectile motion.
A projectile is launched from an initial height of 2 meters at an angle of 45
degrees and falls under the influence of gravity and air resistance. Drag is
proportional to the square of the velocity. We would like to determine the distance
travelled by the projectile as a function of the drag coefficient and the launch
velocity.
"""
def __init__(
self,
parameters_range: dict[str, tuple[float, float]] | None = None,
output_names: list[str] | None = None,
log_level: str = "progress_bar",
):
if parameters_range is None:
parameters_range = {"c": (-5.0, 1.0), "v0": (0.0, 1000)}
if output_names is None:
output_names = ["distance"]
super().__init__(parameters_range, output_names, log_level)
def _forward(self, x: TensorLike) -> TensorLike:
"""
Simulate the projectile motion and return the distance travelled.
Parameters
----------
x: TensorLike
Dictionary of input parameter values to simulate:
- `c`: the drag coefficient on a log scale
- `v0`: velocity
Returns
-------
TensorLike
Distance travelled by projectile.
"""
assert x.shape[0] == 1, (
f"Simulator._forward expects a single input, got {x.shape[0]}"
)
y = simulate_projectile(x.cpu().numpy()[0])
return torch.tensor([y]).view(-1, 1)
[docs]
class ProjectileMultioutput(Simulator):
"""
Multi-output simulator of projectile motion.
Simulator of projectile motion that outputs both the distance travelled by the
projectile and its velocity on impact.
"""
def __init__(
self,
parameters_range: dict[str, tuple[float, float]] | None = None,
output_names: list[str] | None = None,
log_level: str = "progress_bar",
):
if parameters_range is None:
parameters_range = {"c": (-5.0, 1.0), "v0": (0.0, 1000)}
if output_names is None:
output_names = ["distance", "impact_velocity"]
super().__init__(parameters_range, output_names, log_level)
def _forward(self, x: TensorLike) -> TensorLike:
"""
Simulate the projectile motion with multiple outputs.
Simulate projectile motion and return the distance travelled and impact
velocity.
Parameters
----------
x: TensorLike
Dictionary of input parameter values to simulate:
- `c`: the drag coefficient on a log scale
- `v0`: velocity
Returns
-------
TensorLike
Distance travelled by projectile and impact velocity.
"""
assert x.shape[0] == 1, (
f"Simulator._forward expects a single input, got {x.shape[0]}"
)
y = simulate_projectile_multioutput(x.cpu().numpy()[0])
return torch.tensor([y])
[docs]
def f(t: float, y: NumpyLike, c: float): # noqa: ARG001
"""
Compute RHS of system of differential equations, returning vector derivative.
Parameters
----------
t: float
Time variable (not used).
y: array
Array of dependent variables (vx, vy, x, y).
c: float
Drag coefficient (non-negative).
"""
# check inputs and extract
assert len(y) == 4
assert c >= 0.0
vx = y[0]
vy = y[1]
# calculate derivatives
dydt = np.zeros(4)
dydt[0] = -c * vx * np.sqrt(vx**2 + vy**2)
dydt[1] = -9.8 - c * vy * np.sqrt(vx**2 + vy**2)
dydt[2] = vx
dydt[3] = vy
return dydt
[docs]
def event(t: float, y: NumpyLike, c: float) -> float: # noqa: ARG001
"""
Event to trigger end of integration. Stops when projectile hits ground.
Parameters
----------
t: float
Time variable (not used).
y: array
Array of dependent variables (vx, vy, x, y).
c: float
Drag coefficient (non-negative).
Returns
-------
float
The height of the projectile.
"""
assert len(y) == 4
assert c >= 0.0
return y[3]
# indicate this event terminates the simulation
event.terminal = True # pyright: ignore[reportFunctionMemberAccess]
[docs]
def simulator_base(x: NumpyLike):
"""
Simulate ODE system for projectile motion with drag.
Returns distance projectile travels.
Parameters
----------
x: NumpyLike
Array of input parameters (c, v0).
Returns
-------
results: scipy.integrate.OdeResult
Results of ODE integration.
"""
# unpack values
assert len(x) == 2
assert x[1] > 0.0
c = 10.0 ** x[0]
v0 = x[1]
# set initial conditions
y0 = np.zeros(4)
y0[0] = v0 / np.sqrt(2.0)
y0[1] = v0 / np.sqrt(2.0)
y0[3] = 2.0
# run simulation
return solve_ivp(f, (0.0, 1.0e8), y0, events=event, args=(c,))
[docs]
def simulate_projectile(x: NumpyLike) -> float:
"""
Return the distance travelled by the projectile.
Distance is obtained by solving the ODE system for projectile motion with drag.
Parameters
----------
x: NumpyLike
Array of input parameters (c, v0).
Returns
-------
distance: float
Distance travelled by projectile.
"""
results = simulator_base(x)
return results.y_events[0][0][2]
[docs]
def simulate_projectile_multioutput(x: NumpyLike) -> tuple[float, float]:
"""
Return the distance travelled by the projectile and its impact velocity.
Simulator to solve ODE system with multiple outputs.
Parameters
----------
x: NumpyLike
Array of input parameters (c, v0).
Returns
-------
float, float
Distance travelled by projectile and its velocity on impact.
"""
results = simulator_base(x)
return (
results.y_events[0][0][2],
np.sqrt(results.y_events[0][0][0] ** 2 + results.y_events[0][0][1] ** 2),
)
