"""
Helper functions that are used by WallGo.
This includes the derivative functions, a function to compute the step size for solving
detonations and other small functions.
"""
from typing import Callable
import numpy as np
from scipy import integrate, special, optimize
FIRST_DERIV_COEFF = {
"2": np.array([[-0.5, 0.5], [-1, 1], [-1, 1]], dtype=float),
"4": np.array(
[
[1, -8, 8, -1],
[-4, -6, 12, -2],
[-22, 36, -18, 4],
[-4, 18, -36, 22],
[2, -12, 6, 4],
],
dtype=float,
)
/ 12,
}
SECOND_DERIV_COEFF = {
"2": np.array([[1, -2, 1], [1, -2, 1], [1, -2, 1]], dtype=float),
"4": np.array(
[
[-1, 16, -30, 16, -1],
[11, -20, 6, 4, -1],
[35, -104, 114, -56, 11],
[11, -56, 114, -104, 35],
[-1, 4, 6, -20, 11],
],
dtype=float,
)
/ 12,
}
FIRST_DERIV_POS = {
"2": np.array([[-1, 1], [0, 1], [-1, 0]], dtype=float),
"4": np.array(
[[-2, -1, 1, 2], [-1, 0, 1, 2], [0, 1, 2, 3], [-3, -2, -1, 0], [-2, -1, 0, 1]],
dtype=float,
),
}
SECOND_DERIV_POS = {
"2": np.array([[-1, 0, 1], [0, 1, 2], [-2, -1, 0]], dtype=float),
"4": np.array(
[
[-2, -1, 0, 1, 2],
[-1, 0, 1, 2, 3],
[0, 1, 2, 3, 4],
[-4, -3, -2, -1, 0],
[-3, -2, -1, 0, 1],
],
dtype=float,
),
}
HESSIAN_POS = {
"2": np.array([[1, 1, -1, -1], [1, -1, 1, -1]], dtype=float),
"4": np.array(
[[2, 2, 1, 1, -1, -1, -2, -2], [2, -2, 1, -1, 1, -1, 2, -2]], dtype=float
),
}
HESSIAN_COEFF = {
"2": np.array([1, -1, -1, 1], dtype=float) / 4,
"4": np.array([-1, 1, 16, -16, -16, 16, 1, -1], dtype=float) / 48,
}
[docs]
def derivative(
f: Callable[..., np.ndarray],
x: float | np.ndarray,
n: int=1,
order: int=4,
bounds: tuple[float,float] | None=None,
epsilon: float=1e-16,
scale: float=1.0,
dx: float | None=None,
args: list | None=None,
) -> np.ndarray:
r"""Computes numerical derivatives of a callable function. Uses the :py:data:`epsilon`
and :py:data:`scale` parameters to estimate the optimal value of :py:data:`dx`, if the latter is
not provided.
Parameters
----------
f : function
Function to differentiate. Should take a float or an array as argument
and return a float or array (the returned array can have a different
shape as the input, but the first axis must match).
x : float or array-like
The position at which to evaluate the derivative.
n : int, optional
The number of derivatives to take. Can be 0, 1, 2. The default is 1.
order : int, optional
The accuracy order of the scheme. Errors are of order
:math:`\mathcal{O}({\rm d}x^{\text{order}/(\text{order+n})})`. Can be 2 or 4.
Note that the order at the endpoints is reduced by 1 as it would require
more function evaluations to keep the same order. The default is 4.
bounds : tuple or None, optional
Interval in which f can be called. If None, can be evaluated anywhere.
The default is None.
epsilon : float, optional
Fractional accuracy at which f can be evaluated. If f is a simple
function, should be close to the machine precision. Default is 1e-16.
scale : float, optional
Typical scale at which f(x) change by order 1. Default is 1.
dx : float or None, optional
The magnitude of finite differences. If None, use epsilon and scale to
estimate the optimal dx. Default is None.
args: list, optional
List of other fixed arguments passed to the function :math:`f`.
Returns
-------
res : float
The value of the derivative of :py:data:`f` evaluated at :py:data:`x`.
"""
x = np.asarray(x)
boundsTuple: tuple
if bounds is None:
boundsTuple = (-np.inf, np.inf)
else:
boundsTuple = tuple(bounds)
if args is None:
args = []
assert (
isinstance(boundsTuple, tuple) and
len(boundsTuple) == 2 and
boundsTuple[1] > boundsTuple[0]
), "Derivative error: bounds must be a tuple of 2 elements or None."
assert n in (0, 1, 2), "Derivative error: n must be 0, 1 or 2."
assert order in (2, 4), "Derivative error: order must be 2 or 4."
assert np.all(x <= boundsTuple[1]) and np.all(
x >= boundsTuple[0]
), f"Derivative error: {x=} must be inside bounds."
if n == 0:
return f(x, *args)
# If dx is not provided, we estimate it from scale and epsilon by minimizing
# the total error ~ epsilon/dx**n + dx**order.
dxFloat: float
if dx is None:
assert isinstance(epsilon, float), "Derivative error: epsilon must be a float."
assert isinstance(scale, float), "Derivative error: scale must be a float."
dxFloat = scale * epsilon ** (1 / (n + order))
else:
dxFloat = float(dx)
# This step increases greatly the accuracy because it makes sure (x + dx) - x
# is exactly equal to dx (no precision error).
temp = x + dxFloat
dxFloat = temp - x
offset = np.zeros_like(x, dtype=int)
offset -= x + dxFloat > float(boundsTuple[1])
offset += x - dxFloat < float(boundsTuple[0])
if order == 4:
offset -= x + 2 * dxFloat > boundsTuple[1]
offset += x - 2 * dxFloat < boundsTuple[0]
if n == 1:
pos = x[None, ...] + FIRST_DERIV_POS[str(order)].T[:, offset.tolist()]*dxFloat
coeff = FIRST_DERIV_COEFF[str(order)].T[:, offset.tolist()] / dxFloat
elif n == 2:
pos = x[None, ...] + SECOND_DERIV_POS[str(order)].T[:, offset.tolist()]*dxFloat
coeff = SECOND_DERIV_COEFF[str(order)].T[:, offset.tolist()] / dxFloat**2
fx = f(pos, *args)
fxShapeLength = len(fx.shape)
coeffShapeLength = len(coeff.shape)
return np.asarray(np.sum(
coeff.reshape(coeff.shape + (fxShapeLength - coeffShapeLength) * (1,))
* f(pos, *args),
axis=0,
))
[docs]
def gradient(
f: Callable[..., np.ndarray],
x: float | np.ndarray,
order: int=4,
epsilon: float=1e-16,
scale: float | np.ndarray=1.0,
dx: float | np.ndarray | None=None,
axis: list | int | None=None,
args: list | None=None,
) -> np.ndarray:
r"""Computes the gradient of a callable function. Uses the :py:data:`epsilon`
and :py:data:`scale` parameters to estimate the optimal value of :py:data:`dx`, if the latter is
not provided.
Parameters
----------
f : function
Function to differentiate. Should take an array as argument
and return an array.
x : array-like
The position at which to evaluate the derivative. The size of the last
axis must correspond to the number of variables on which f depends.
order : int, optional
The accuracy order of the scheme. Errors are of order
:math:`\mathcal{O}({\rm d}x^{\text{order}/(\text{order}+1)})`.
Can be 2 or 4. The default is 4.
epsilon : float, optional
Fractional accuracy at which f can be evaluated. If f is a simple
function, should be close to the machine precision. Default is 1e-16.
scale : float or array-like, optional
Typical scale at which f(x) change by order 1. Can be an array, in which
case each element corresponds to the scale of a different variable.
Default is 1.
dx : float or np.ndarray or None, optional
The magnitude of finite differences. Can be an array, in which case
each element corresponds to the dx of a different variable.If None, use
epsilon and scale to estimate the optimal dx. Default is None.
axis : list, int or None, optional
Element of the gradient to return. If None, returns the whole gradient.
Default is None.
args: list, optional
List of other fixed arguments passed to the function :math:`f`.
Returns
-------
res : float
The value of the gradient of :py:data:`f` evaluated at :py:data:`x`.
"""
x = np.asarray(x)
nbrVariables = x.shape[-1]
if args is None:
args = []
axisList: list
if isinstance(axis, int):
axisList = [axis]
elif axis is None:
axisList = np.arange(nbrVariables).tolist()
else:
axisList = list(axis)
for i in axisList:
assert (
-nbrVariables <= i < nbrVariables
), "Gradient error: axis must be between -nbrVariables and "\
"nbrVariables-1 or None."
assert order in (2,4), "Gradient error: order must be 2 or 4."
# If dx is not provided, we estimate it from scale and epsilon by minimizing
# the total error ~ epsilon/dx**n + dx**order.
dxArray: np.ndarray
if dx is None:
assert isinstance(epsilon, float), "Gradient error: epsilon must be a float."
if isinstance(scale, float):
scale = scale * np.ones(nbrVariables)
else:
scale = np.asanyarray(scale)
assert (
scale.size == nbrVariables
), "Gradient error: scale must be a float or an array of size nbrVariables."
dxArray = scale * epsilon ** (1 / (1 + order))
elif isinstance(dx, float):
dxArray = dx * np.ones(nbrVariables)
else:
dxArray = np.asarray(dx)
assert (
dxArray.size == nbrVariables
), "Gradient error: dx must be None, a float or an array of size nbrVariables."
# This step increases greatly the accuracy because it makes sure (x + dx) - x
# is exactly equal to dx (no precision error).
temp = x + dxArray
dxArray = temp - x
pos = np.expand_dims(x, (-3, -2)) + FIRST_DERIV_POS[str(order)][
0, :, None, None
] * np.identity(nbrVariables)[axisList, :] * np.expand_dims(dxArray, (-3, -2))
shape = pos.shape[:-1]
pos = pos.reshape((int(pos.size / nbrVariables), nbrVariables))
coeff = FIRST_DERIV_COEFF[str(order)][0, :, None] / np.expand_dims(
dxArray[..., axisList], -2
)
fEvaluation = f(pos, *args).reshape(shape)
return np.asarray(np.sum(coeff * fEvaluation, axis=-2))
[docs]
def hessian(
f: Callable[..., np.ndarray],
x: float | np.ndarray,
order: int=4,
epsilon: float=1e-16,
scale: float | np.ndarray=1.0,
dx: float | np.ndarray | None=None,
xAxis: list | int | None=None,
yAxis: list | int | None=None,
args: list | None=None,
) -> np.ndarray:
r"""Computes the hessian of a callable function. Uses the :py:data:`epsilon`
and :py:data:`scale` parameters to estimate the optimal value of :py:data:`dx`, if the latter is
not provided.
Parameters
----------
f : function
Function to differentiate. Should take an array as argument
and return an array.
x : array-like
The position at which to evaluate the derivative. The size of the last
axis must correspond to the number of variables on which f depends.
order : int, optional
The accuracy order of the scheme. Errors are of order
:math:`\mathcal{O}({\rm d}x^{\text{order}/(\text{order}+2)})`.
Can be 2 or 4. The default is 4.
epsilon : float, optional
Fractional accuracy at which f can be evaluated. If f is a simple
function, should be close to the machine precision. Default is 1e-16.
scale : float, optional
Typical scale at which f(x) change by order 1. Default is 1.
dx : float or None, optional
The magnitude of finite differences. If None, use epsilon and scale to
estimate the optimal dx. Default is None.
xAxis : list, int or None, optional
Lines of the hessian matrix to return. If None, returns all the lines.
Default is None.
yAxis : list, int or None, optional
Columns of the hessian matrix to return. If None, returns all the columns.
Default is None.
args: list, optional
List of other fixed arguments passed to the function :math:`f`.
Returns
-------
res : float
The value of the hessian of :py:data:`f` evaluated at :py:data:`x`.
"""
x = np.asarray(x)
nbrVariables = x.shape[-1]
if args is None:
args = []
xAxisList: list
if isinstance(xAxis, int):
xAxisList = [xAxis]
elif xAxis is None:
xAxisList = np.arange(nbrVariables).tolist()
else:
xAxisList = list(xAxis)
for i in xAxisList:
assert (
-nbrVariables <= i < nbrVariables
), "Hessian error: axis must be between -nbrVariables and "\
"nbrVariables-1 or None."
yAxisList: list
if isinstance(yAxis, int):
yAxisList = [yAxis]
elif yAxis is None:
yAxisList = np.arange(nbrVariables).tolist()
else:
yAxisList = list(yAxis)
for i in yAxisList:
assert (
-nbrVariables <= i < nbrVariables
), "Hessian error: axis must be between -nbrVariables and "\
"nbrVariables-1 or None."
assert order in (2, 4), "Hessian error: order must be 2 or 4."
# If dx is not provided, we estimate it from scale and epsilon by minimizing
# the total error ~ epsilon/dx**n + dx**order.
dxArray: np.ndarray
if dx is None:
assert isinstance(epsilon, float), "Hessian error: epsilon must be a float."
if isinstance(scale, float):
scale = scale * np.ones(nbrVariables)
else:
scale = np.asanyarray(scale)
assert (
scale.size == nbrVariables
), "Hessian error: scale must be a float or an array of size nbrVariables."
dxArray = scale * epsilon ** (1 / (2 + order))
elif isinstance(dx, float):
dxArray = dx * np.ones(nbrVariables)
else:
dxArray = np.asarray(dx)
assert (
dxArray.size == nbrVariables
), "Hessian error: dx must be None, a float or an array of size nbrVariables."
# This step increases greatly the accuracy because it makes sure (x + dx) - x
# is exactly equal to dx (no precision error).
temp = x + dxArray
dxArray = temp - x
pos = (
np.expand_dims(x, (-4, -3, -2))
+ HESSIAN_POS[str(order)][0, :, None, None, None]
* np.identity(nbrVariables)[xAxisList, None, :]
* np.expand_dims(dxArray, (-4, -3, -2))
+ HESSIAN_POS[str(order)][1, :, None, None, None]
* np.identity(nbrVariables)[None, yAxisList, :]
* np.expand_dims(dxArray, (-4, -3, -2))
)
shape = pos.shape[:-1]
pos = pos.reshape((int(pos.size / nbrVariables), nbrVariables))
coeff = HESSIAN_COEFF[str(order)][:, None, None] / (
np.expand_dims(dxArray[..., yAxisList], (-3, -2))
* np.expand_dims(dxArray[..., xAxisList], (-3, -1))
)
fEvaluation = f(pos, *args).reshape(shape)
return np.asarray(np.sum(coeff * fEvaluation, axis=-3))
[docs]
def gammaSq(v: float) -> float:
r"""
Lorentz factor :math:`\gamma^2` corresponding to velocity :math:`v`
"""
return 1.0 / (1.0 - v * v)
[docs]
def boostVelocity(xi: float, v: float) -> float:
"""
Lorentz-transformed velocity
"""
return (xi - v) / (1.0 - xi * v)
[docs]
def nextStepDeton(
pos1: float,
pos2: float,
pressure1: float,
pressure2: float,
mean2ndDeriv: float,
std2ndDeriv: float,
pressureTol: float,
posMax: float,
overshootProb: float = 0.05,
) -> float:
"""
Function used in :py:class:`EOM` to find detonation solutions. It finds the next
point to be sampled to try to bracket a root in such a way that the probability of
overshooting a root is roughly equal to overshootProb.
To estimate the overshoot probability, it fits the pressure to a quadratic which is
equal to :py:data:`pressure2` at :py:data`pos2`, but with uncertain 1st and 2nd derivatives which are
assumed to be normally distributed. The mean of the 1st derivative is computed by
finite differences from the last 2 points.
Parameters
----------
pos1 : float
Position of the first sampled point.
pos2 : float
position of the second sampled point.
pressure1 : float
Pressure at pos1.
pressure2 : float
Pressure at pos2.
mean2ndDeriv : float
Estimate of the 2nd derivative at pos2.
std2ndDeriv : float
Uncertainty on the 2nd derivative at pos2.
pressureTol : float
Relative accuracy at which pressure1 and pressure2 are computed.
posMax : float
Maximal position that the next step can have.
overshootProb : float, optional
Desired overshoot probability. A smaller value will lead to smaller step sizes
which will take longer to evaluate, but with less chances of missing a root.
The default is 0.05.
Returns
-------
float
Position where the overshoot probability is overshootProb (or posMax if there is
no solution).
"""
assert pos2 > pos1, "Error: pos2 must be greater than pos1."
assert posMax > pos2, "Error: posMax must be greater than pos2."
assert std2ndDeriv >= 0, "Error: std2ndDeriv must be positive."
assert pressureTol >= 0, "Error: pressureTol must positive."
assert 0 < overshootProb < 1, "Error: overshootProb must be between 0 and 1."
# This function requires pressure2 to be negative. If that's not the case,
# we invert the y axis.
if pressure2 > 0:
pressure1 *= -1
pressure2 *= -1
mean2ndDeriv *= -1
# Use pressure units such that abs(pressure2)=1.
# Helps when integrate to get the prob
pressure1 /= abs(pressure2)
mean2ndDeriv /= abs(pressure2)
std2ndDeriv /= abs(pressure2)
pressure2 = -1
dx = pos2 - pos1
dp = pressure2 - pressure1
# Mean and variance of first derivative at pos2
# (second term due to finite difference error)
meanDeriv = dp / dx + dx * mean2ndDeriv / 2
# First term: variance due to pressure uncertainty
# Second term: variance due to finite difference error
varDeriv = (pressure1**2 + pressure2**2) * (pressureTol / dx) ** 2 + (
std2ndDeriv * dx / 2
) ** 2
stdDeriv = np.sqrt(varDeriv)
# Probability of overshooting a root when the 2nd derivative is positive:
def probPositive(pos: float) -> float:
if pos == pos2:
return 0
# Min derivative that can lead to overshooting
def derivMin(secondDeriv: float) -> float:
return -pressure2 / (pos - pos2) - secondDeriv * (pos - pos2) / 2
# Probability density of the 2nd derivative
def probDensity(secondDeriv: float) -> float:
return float(np.exp(-((secondDeriv - mean2ndDeriv) ** 2) / std2ndDeriv**2/2)
* special.erfc((derivMin(secondDeriv) - meanDeriv) / stdDeriv / np.sqrt(2))
/ (2 * std2ndDeriv * np.sqrt(2 * np.pi)))
# Integrate the probability density to get the probability
return float(integrate.quad(
probDensity, -2 * pressure2 / (pos - pos2) ** 2, np.inf, full_output=1
)[0])
# Probability of overshooting 2 roots when the 2nd derivative is negative:
def probNegative(pos: float) -> float:
if pos == pos2:
return 0
# Min and max derivatives that can lead to overshooting
def derivMin(secondDeriv: float) -> float:
return float(np.sqrt(2 * secondDeriv * pressure2))
def derivMax(secondDeriv: float) -> float:
return -pressure2 / (pos - pos2) - secondDeriv * (pos - pos2) / 2
# Probability density of the 2nd derivative
def probDensity(secondDeriv: float) -> float:
return float(np.exp(-((secondDeriv - mean2ndDeriv) ** 2) / std2ndDeriv**2/2)
* (
special.erf((derivMax(secondDeriv) - meanDeriv) / stdDeriv / np.sqrt(2))
- special.erf(
(derivMin(secondDeriv) - meanDeriv) / stdDeriv / np.sqrt(2)
)
)
/ (2 * std2ndDeriv * np.sqrt(2 * np.pi))
)
# Integrate the probability density to get the probability
return float(integrate.quad(
probDensity, -np.inf, 2 * pressure2 / (pos - pos2) ** 2, full_output=1
)[0])
try:
# Find a position for which the total probability is overshootProb
return float(optimize.root_scalar(
lambda pos: probPositive(pos) + probNegative(pos) - overshootProb,
bracket=[pos2, posMax],
).root)
except:
# If no solution is found, returns posMax
return posMax