Skip to content

Commit 744c3cb

Browse files
authored
ENH: Function vectorized speed-up and refactor (#1049)
* MNT: deduplicate function arithmetic logic. ENH: modularize function arithmetic and source dispatch. MNT: refactor function math modules. MNT: optimize function math speed. * MNT: enhance parachute trigger evaluation for speed. * MNT: improve typing stack on Function evaluation. * MNT: architecture corrections and optimizations. * TST: attempt at tests re-run. * MNT: minor fixes and docstring update.
1 parent ba9d130 commit 744c3cb

28 files changed

Lines changed: 4839 additions & 1730 deletions

CHANGELOG.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@ Attention: The newest changes should be on top -->
3232

3333
### Added
3434

35+
- ENH: Function vectorized speed-up and refactor [#1049](https://github.com/RocketPy-Team/RocketPy/pull/1049)
3536
- ENH: REV: Revert PR #958 [#1063](https://github.com/RocketPy-Team/RocketPy/pull/1063)
3637
- ENH: MNT: remove duplicate and merge-sync entries from Unreleased changelog [#1062](https://github.com/RocketPy-Team/RocketPy/pull/1062)
3738
### Changed

rocketpy/__init__.py

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
from . import utilities
12
from .control import _Controller
23
from .environment import Environment, EnvironmentAnalysis
34
from .exceptions import (
@@ -68,8 +69,3 @@
6869
StochasticTail,
6970
StochasticTrapezoidalFins,
7071
)
71-
72-
# Imported last: utilities pulls in Environment/Rocket/encoders, which are only
73-
# fully available once the imports above have run. Exposes
74-
# ``rocketpy.utilities`` (including ``enable_logging``) on ``import rocketpy``.
75-
from . import utilities

rocketpy/environment/environment.py

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -498,7 +498,7 @@ def __set_barometric_height_function(self, source):
498498
interpolation="linear",
499499
extrapolation="natural",
500500
)
501-
if callable(self.barometric_height.source):
501+
if not self.barometric_height.is_array_source():
502502
# discretize to speed up flight simulation
503503
self.barometric_height.set_discrete(
504504
0,
@@ -581,7 +581,7 @@ def __set_wind_heading_function(self, source):
581581
def __reset_barometric_height_function(self):
582582
# NOTE: this assumes self.pressure and max_expected_height are already set.
583583
self.barometric_height = self.pressure.inverse_function()
584-
if callable(self.barometric_height.source):
584+
if not self.barometric_height.is_array_source():
585585
# discretize to speed up flight simulation
586586
self.barometric_height.set_discrete(
587587
0,
@@ -900,7 +900,7 @@ def set_gravity_model(self, gravity=None):
900900
>>> g_0 = 9.80665
901901
>>> env_cte_g = Environment(gravity=g_0)
902902
>>> env_cte_g.gravity([0, 100, 1000])
903-
[np.float64(9.80665), np.float64(9.80665), np.float64(9.80665)]
903+
array([9.80665, 9.80665, 9.80665])
904904
905905
It's also possible to variate the gravity acceleration by defining
906906
its function of height:
@@ -1572,7 +1572,7 @@ def process_custom_atmosphere(
15721572
self.__reset_barometric_height_function()
15731573

15741574
# Check maximum height of custom pressure input
1575-
if not callable(self.pressure.source):
1575+
if self.pressure.is_array_source():
15761576
max_expected_height = max(self.pressure[-1, 0], max_expected_height)
15771577

15781578
# Save temperature profile
@@ -1582,14 +1582,14 @@ def process_custom_atmosphere(
15821582
else:
15831583
self.__set_temperature_function(temperature)
15841584
# Check maximum height of custom temperature input
1585-
if not callable(self.temperature.source):
1585+
if self.temperature.is_array_source():
15861586
max_expected_height = max(self.temperature[-1, 0], max_expected_height)
15871587

15881588
# Save wind profile
15891589
self.__set_wind_velocity_x_function(wind_u)
15901590
self.__set_wind_velocity_y_function(wind_v)
15911591
# Check maximum height of custom wind input
1592-
if not callable(self.wind_velocity_x.source):
1592+
if self.wind_velocity_x.is_array_source():
15931593
max_expected_height = max(self.wind_velocity_x[-1, 0], max_expected_height)
15941594

15951595
def wind_heading_func(h): # TODO: create another custom reset for heading
Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
"""Interpolation and extrapolation submodule for the Function class.
2+
3+
This package provides factory functions that build strategy objects for
4+
interpolation and extrapolation.
5+
6+
The registry helper :func:`build_interpolation_evaluator` is the main
7+
entry point consumed by :class:`rocketpy.mathutils.function.Function`.
8+
"""
9+
10+
from .registry import build_interpolation_evaluator
11+
12+
__all__ = ["build_interpolation_evaluator"]
Lines changed: 265 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,265 @@
1+
"""Coefficient fitting routines for 1-D interpolation methods.
2+
3+
Every public function in this module takes sorted ``x`` and ``y`` arrays and
4+
returns a data structure (NumPy array / tuple) that the corresponding
5+
evaluation factory can close over.
6+
"""
7+
8+
import warnings
9+
10+
import numpy as np
11+
from scipy import linalg
12+
13+
14+
def fit_polynomial(x, y):
15+
"""Fit a single polynomial of degree ``len(x) - 1`` through all points.
16+
17+
Parameters
18+
----------
19+
x : ndarray, shape (n,)
20+
y : ndarray, shape (n,)
21+
22+
Returns
23+
-------
24+
coeffs : ndarray, shape (n,)
25+
Coefficients in *ascending* power order: ``c[0] + c[1]*x + …``.
26+
"""
27+
degree = len(x) - 1
28+
if np.amax(np.abs(x)) ** degree > 1e308:
29+
warnings.warn(
30+
"Polynomial interpolation of too many points can't be done. "
31+
"Once the degree is too high, numbers get too large. "
32+
"Falling back to spline coefficients.",
33+
stacklevel=3,
34+
)
35+
return None # caller should fall back to spline
36+
V = np.vander(x, increasing=True)
37+
return np.linalg.solve(V, y)
38+
39+
40+
def fit_spline(x, y):
41+
r"""Fit a natural cubic spline in local coordinates.
42+
43+
Returns coefficients ``[a, b, c, d]`` for each interval stored as a
44+
``(4, n-1)`` array, where the polynomial for segment *i* is
45+
46+
.. math::
47+
p_i(t) = a_i + b_i t + c_i t^2 + d_i t^3, \quad t = x - x_i
48+
49+
Parameters
50+
----------
51+
x : ndarray, shape (n,)
52+
y : ndarray, shape (n,)
53+
54+
Returns
55+
-------
56+
coeffs : ndarray, shape (4, n-1)
57+
Rows are ``[a, b, c, d]``.
58+
"""
59+
n = len(x)
60+
h = np.diff(x)
61+
62+
# Build tridiagonal system for the c coefficients (natural BC: c[0]=c[-1]=0)
63+
banded = np.zeros((3, n))
64+
banded[1, 0] = banded[1, -1] = 1.0
65+
banded[2, :-2] = h[:-1]
66+
banded[1, 1:-1] = 2.0 * (h[:-1] + h[1:])
67+
banded[0, 2:] = h[1:]
68+
69+
rhs = np.zeros(n)
70+
rhs[1:-1] = 3.0 * ((y[2:] - y[1:-1]) / h[1:] - (y[1:-1] - y[:-2]) / h[:-1])
71+
72+
c = linalg.solve_banded((1, 1), banded, rhs, overwrite_ab=True, overwrite_b=True)
73+
74+
b = (y[1:] - y[:-1]) / h - h * (2.0 * c[:-1] + c[1:]) / 3.0
75+
d = (c[1:] - c[:-1]) / (3.0 * h)
76+
77+
return np.vstack([y[:-1], b, c[:-1], d])
78+
79+
80+
def fit_akima(x, y):
81+
r"""Fit an Akima spline in local coordinates.
82+
83+
The slopes at each knot are computed using the original Akima weighting
84+
of adjacent finite differences, extended at the boundaries by reflected
85+
values. This is fully vectorised — no Python loop.
86+
87+
Returns coefficients ``[a, b, c, d]`` for each interval in a ``(4, n-1)``
88+
array identical in layout to :func:`fit_spline`.
89+
90+
Parameters
91+
----------
92+
x : ndarray, shape (n,)
93+
y : ndarray, shape (n,)
94+
95+
Returns
96+
-------
97+
coeffs : ndarray, shape (4, n-1)
98+
"""
99+
n = len(x)
100+
h = np.diff(x)
101+
m = np.diff(y) / h # slopes of segments, shape (n-1,)
102+
103+
# Extend m with 2 ghost values on each side (Akima boundary reflection)
104+
# m_ext indices: 0..n+2 (n-1 interior + 2 left + 2 right)
105+
m_ext = np.empty(n + 3)
106+
m_ext[2:-2] = m
107+
m_ext[1] = 2.0 * m[0] - m[1]
108+
m_ext[0] = 2.0 * m_ext[1] - m[0]
109+
m_ext[-2] = 2.0 * m[-1] - m[-2]
110+
m_ext[-1] = 2.0 * m_ext[-2] - m[-1]
111+
112+
# Weights: |m_{i+1} - m_i|
113+
dm = np.abs(np.diff(m_ext)) # shape (n+2,)
114+
115+
# Akima slope at each knot: weighted average of adjacent segment slopes.
116+
# w1 = |m_{i+1} - m_i|, w2 = |m_{i-1} - m_{i-2}|
117+
# t_i = (w1 * m_{i-1} + w2 * m_i) / (w1 + w2)
118+
# When w1 + w2 == 0, fall back to simple average.
119+
w1 = dm[2:] # |m_{i+1} - m_i| for i in 0..n-1
120+
w2 = dm[:-2] # |m_{i-1} - m_{i-2}| for i in 0..n-1
121+
m_left = m_ext[1:-2] # m_{i-1}
122+
m_right = m_ext[2:-1] # m_i
123+
124+
wsum = w1 + w2
125+
t = 0.5 * (m_left + m_right)
126+
np.divide(
127+
w1 * m_left + w2 * m_right,
128+
wsum,
129+
out=t,
130+
where=wsum > 0,
131+
)
132+
133+
# Hermite-to-local coefficients:
134+
# a = y_i
135+
# b = t_i
136+
# c = (3*m_i - 2*t_i - t_{i+1}) / h_i
137+
# d = (t_i + t_{i+1} - 2*m_i) / h_i^2
138+
a = y[:-1]
139+
b = t[:-1]
140+
c_coeff = (3.0 * m - 2.0 * t[:-1] - t[1:]) / h
141+
d_coeff = (t[:-1] + t[1:] - 2.0 * m) / (h * h)
142+
143+
return np.vstack([a, b, c_coeff, d_coeff])
144+
145+
146+
def fit_pchip(x, y):
147+
r"""Fit a PCHIP (Piecewise Cubic Hermite Interpolating Polynomial).
148+
149+
Uses the Fritsch–Carlson method to compute monotone-preserving slopes
150+
and converts them to local-coordinate cubic coefficients identical in
151+
layout to :func:`fit_spline`.
152+
153+
Parameters
154+
----------
155+
x : ndarray, shape (n,)
156+
y : ndarray, shape (n,)
157+
158+
Returns
159+
-------
160+
coeffs : ndarray, shape (4, n-1)
161+
"""
162+
n = len(x)
163+
h = np.diff(x)
164+
m = np.diff(y) / h
165+
166+
# Compute slopes at each knot
167+
t = np.zeros(n)
168+
169+
# Interior knots
170+
if n > 2:
171+
# Harmonic mean weighted by segment lengths
172+
w1 = 2.0 * h[1:] + h[:-1]
173+
w2 = h[1:] + 2.0 * h[:-1]
174+
175+
# Where signs differ or either is zero, slope is zero (monotonicity)
176+
sign_change = (m[:-1] * m[1:]) <= 0
177+
pos_mask = ~sign_change
178+
179+
# Safe harmonic mean
180+
t[1:-1] = np.where(
181+
pos_mask,
182+
(w1 + w2) / (w1 / m[:-1] + w2 / m[1:]),
183+
0.0,
184+
)
185+
186+
# Boundary slopes: one-sided, clamped for monotonicity
187+
t[0] = _pchip_edge_slope(
188+
h[0], h[1] if n > 2 else h[0], m[0], m[1] if n > 2 else m[0]
189+
)
190+
t[-1] = _pchip_edge_slope(
191+
h[-1], h[-2] if n > 2 else h[-1], m[-1], m[-2] if n > 2 else m[-1]
192+
)
193+
194+
# Convert to local Hermite coefficients (same as akima)
195+
a = y[:-1]
196+
b = t[:-1]
197+
c_coeff = (3.0 * m - 2.0 * t[:-1] - t[1:]) / h
198+
d_coeff = (t[:-1] + t[1:] - 2.0 * m) / (h * h)
199+
200+
return np.vstack([a, b, c_coeff, d_coeff])
201+
202+
203+
def _pchip_edge_slope(h0, h1, m0, m1):
204+
"""Compute boundary slope for PCHIP (non-centred 3-point formula,
205+
clamped to preserve monotonicity)."""
206+
t = ((2.0 * h0 + h1) * m0 - h0 * m1) / (h0 + h1)
207+
if np.sign(t) != np.sign(m0):
208+
t = 0.0
209+
elif np.sign(m0) != np.sign(m1) and np.abs(t) > 3.0 * np.abs(m0):
210+
t = 3.0 * m0
211+
return t
212+
213+
214+
def precompute_linear_deriv_integral(x, y):
215+
"""Pre-compute data for analytical linear derivative and integral.
216+
217+
Parameters
218+
----------
219+
x : ndarray, shape (n,)
220+
y : ndarray, shape (n,)
221+
222+
Returns
223+
-------
224+
slopes : ndarray, shape (n-1,)
225+
Piecewise-constant derivative values.
226+
cum_integrals : ndarray, shape (n,)
227+
Cumulative trapezoid integral from x[0] to x[i].
228+
"""
229+
h = np.diff(x)
230+
slopes = np.diff(y) / h
231+
# Cumulative trapezoid: area of each trapezoid is 0.5*(y[i]+y[i+1])*h[i]
232+
trap_areas = 0.5 * (y[:-1] + y[1:]) * h
233+
cum_integrals = np.empty(len(x))
234+
cum_integrals[0] = 0.0
235+
np.cumsum(trap_areas, out=cum_integrals[1:])
236+
return slopes, cum_integrals
237+
238+
239+
def precompute_cubic_cumulative_integrals(x, coeffs):
240+
"""Pre-compute cumulative integrals over full intervals for cubic methods.
241+
242+
For piece *i* with local polynomial ``a + b*t + c*t² + d*t³``,
243+
the integral over the full interval ``[0, h_i]`` is
244+
``a*h + b*h²/2 + c*h³/3 + d*h⁴/4``.
245+
246+
Parameters
247+
----------
248+
x : ndarray, shape (n,)
249+
coeffs : ndarray, shape (4, n-1)
250+
251+
Returns
252+
-------
253+
cum : ndarray, shape (n,)
254+
``cum[0] = 0``; ``cum[i] = ∫_{x[0]}^{x[i]} p(t) dt``.
255+
"""
256+
h = np.diff(x)
257+
a, b, c, d = coeffs
258+
h2 = h * h
259+
h3 = h2 * h
260+
h4 = h3 * h
261+
piece_integrals = a * h + b * h2 / 2.0 + c * h3 / 3.0 + d * h4 / 4.0
262+
cum = np.empty(len(x))
263+
cum[0] = 0.0
264+
np.cumsum(piece_integrals, out=cum[1:])
265+
return cum

0 commit comments

Comments
 (0)