qbhf2's picture
added NvidiaWarp and GarmentCode repos
66c9c8a
import math
from enum import Enum
import numpy as np
class Polynomial(Enum):
"""Polynomial family defining interpolation nodes over an interval"""
GAUSS_LEGENDRE = 0
"""Gauss--Legendre 1D polynomial family (does not include endpoints)"""
LOBATTO_GAUSS_LEGENDRE = 1
"""Lobatto--Gauss--Legendre 1D polynomial family (includes endpoints)"""
EQUISPACED_CLOSED = 2
"""Closed 1D polynomial family with uniformly distributed nodes (includes endpoints)"""
EQUISPACED_OPEN = 3
"""Open 1D polynomial family with uniformly distributed nodes (does not include endpoints)"""
def __str__(self):
return self.name
def is_closed(family: Polynomial):
"""Whether the polynomial roots include interval endpoints"""
return family == Polynomial.LOBATTO_GAUSS_LEGENDRE or family == Polynomial.EQUISPACED_CLOSED
def _gauss_legendre_quadrature_1d(n: int):
if n == 1:
coords = [0.0]
weights = [2.0]
elif n == 2:
coords = [-math.sqrt(1.0 / 3), math.sqrt(1.0 / 3)]
weights = [1.0, 1.0]
elif n == 3:
coords = [0.0, -math.sqrt(3.0 / 5.0), math.sqrt(3.0 / 5.0)]
weights = [8.0 / 9.0, 5.0 / 9.0, 5.0 / 9.0]
elif n == 4:
c_a = math.sqrt(3.0 / 7.0 - 2.0 / 7.0 * math.sqrt(6.0 / 5.0))
c_b = math.sqrt(3.0 / 7.0 + 2.0 / 7.0 * math.sqrt(6.0 / 5.0))
w_a = (18.0 + math.sqrt(30.0)) / 36.0
w_b = (18.0 - math.sqrt(30.0)) / 36.0
coords = [c_a, -c_a, c_b, -c_b]
weights = [w_a, w_a, w_b, w_b]
elif n == 5:
c_a = 1.0 / 3.0 * math.sqrt(5.0 - 2.0 * math.sqrt(10.0 / 7.0))
c_b = 1.0 / 3.0 * math.sqrt(5.0 + 2.0 * math.sqrt(10.0 / 7.0))
w_a = (322.0 + 13.0 * math.sqrt(70.0)) / 900.0
w_b = (322.0 - 13.0 * math.sqrt(70.0)) / 900.0
coords = [0.0, c_a, -c_a, c_b, -c_b]
weights = [128.0 / 225.0, w_a, w_a, w_b, w_b]
else:
raise NotImplementedError
# Shift from [-1, 1] to [0, 1]
weights = 0.5 * np.array(weights)
coords = 0.5 * np.array(coords) + 0.5
return coords, weights
def _lobatto_gauss_legendre_quadrature_1d(n: int):
if n == 2:
coords = [-1.0, 1.0]
weights = [1.0, 1.0]
elif n == 3:
coords = [-1.0, 0.0, 1.0]
weights = [1.0 / 3.0, 4.0 / 3.0, 1.0 / 3.0]
elif n == 4:
coords = [-1.0, -1.0 / math.sqrt(5.0), 1.0 / math.sqrt(5.0), 1.0]
weights = [1.0 / 6.0, 5.0 / 6.0, 5.0 / 6.0, 1.0 / 6.0]
elif n == 5:
coords = [-1.0, -math.sqrt(3.0 / 7.0), 0.0, math.sqrt(3.0 / 7.0), 1.0]
weights = [1.0 / 10.0, 49.0 / 90.0, 32.0 / 45.0, 49.0 / 90.0, 1.0 / 10.0]
else:
raise NotImplementedError
# Shift from [-1, 1] to [0, 1]
weights = 0.5 * np.array(weights)
coords = 0.5 * np.array(coords) + 0.5
return coords, weights
def _uniform_open_quadrature_1d(n: int):
step = 1.0 / (n + 1)
coords = np.linspace(step, 1.0 - step, n)
weights = np.full(n, 1.0 / (n + 1))
# Boundaries have 3/2 the weight
weights[0] = 1.5 / (n + 1)
weights[-1] = 1.5 / (n + 1)
return coords, weights
def _uniform_closed_quadrature_1d(n: int):
coords = np.linspace(0.0, 1.0, n)
weights = np.full(n, 1.0 / (n - 1))
# Boundaries have half the weight
weights[0] = 0.5 / (n - 1)
weights[-1] = 0.5 / (n - 1)
return coords, weights
def _open_newton_cotes_quadrature_1d(n: int):
step = 1.0 / (n + 1)
coords = np.linspace(step, 1.0 - step, n)
# Weisstein, Eric W. "Newton-Cotes Formulas." From MathWorld--A Wolfram Web Resource.
# https://mathworld.wolfram.com/Newton-CotesFormulas.html
if n == 1:
weights = np.array([1.0])
elif n == 2:
weights = np.array([0.5, 0.5])
elif n == 3:
weights = np.array([2.0, -1.0, 2.0]) / 3.0
elif n == 4:
weights = np.array([11.0, 1.0, 1.0, 11.0]) / 24.0
elif n == 5:
weights = np.array([11.0, -14.0, 26.0, -14.0, 11.0]) / 20.0
elif n == 6:
weights = np.array([611.0, -453.0, 562.0, 562.0, -453.0, 611.0]) / 1440.0
elif n == 7:
weights = np.array([460.0, -954.0, 2196.0, -2459.0, 2196.0, -954.0, 460.0]) / 945.0
else:
raise NotImplementedError
return coords, weights
def _closed_newton_cotes_quadrature_1d(n: int):
coords = np.linspace(0.0, 1.0, n)
# OEIS: A093735, A093736
if n == 2:
weights = np.array([1.0, 1.0]) / 2.0
elif n == 3:
weights = np.array([1.0, 4.0, 1.0]) / 3.0
elif n == 4:
weights = np.array([3.0, 9.0, 9.0, 3.0]) / 8.0
elif n == 5:
weights = np.array([14.0, 64.0, 24.0, 64.0, 14.0]) / 45.0
elif n == 6:
weights = np.array([95.0 / 288.0, 125.0 / 96.0, 125.0 / 144.0, 125.0 / 144.0, 125.0 / 96.0, 95.0 / 288.0])
elif n == 7:
weights = np.array([41, 54, 27, 68, 27, 54, 41], dtype=float) / np.array(
[140, 35, 140, 35, 140, 35, 140], dtype=float
)
elif n == 8:
weights = np.array(
[
5257,
25039,
343,
20923,
20923,
343,
25039,
5257,
]
) / np.array(
[
17280,
17280,
640,
17280,
17280,
640,
17280,
17280,
],
dtype=float,
)
else:
raise NotImplementedError
# Normalize with interval length
weights = weights / (n - 1)
return coords, weights
def quadrature_1d(point_count: int, family: Polynomial):
"""Return quadrature points and weights for the given family and point count"""
if family == Polynomial.GAUSS_LEGENDRE:
return _gauss_legendre_quadrature_1d(point_count)
if family == Polynomial.LOBATTO_GAUSS_LEGENDRE:
return _lobatto_gauss_legendre_quadrature_1d(point_count)
if family == Polynomial.EQUISPACED_CLOSED:
return _closed_newton_cotes_quadrature_1d(point_count)
if family == Polynomial.EQUISPACED_OPEN:
return _open_newton_cotes_quadrature_1d(point_count)
raise NotImplementedError
def lagrange_scales(coords: np.array):
"""Return the scaling factors for Lagrange polynomials with roots at coords"""
lagrange_scale = np.empty_like(coords)
for i in range(len(coords)):
deltas = coords[i] - coords
deltas[i] = 1.0
lagrange_scale[i] = 1.0 / np.prod(deltas)
return lagrange_scale