From 5e4577902dacefc40327ab11da92c004b6f67898 Mon Sep 17 00:00:00 2001 From: PatReis Date: Tue, 12 Dec 2023 17:52:41 +0100 Subject: [PATCH] added layers for polynom.py --- kgcnn/layers/polynom.py | 280 +++++++++++++++++++++++++++++++++++-- kgcnn/layers/relational.py | 1 + 2 files changed, 266 insertions(+), 15 deletions(-) diff --git a/kgcnn/layers/polynom.py b/kgcnn/layers/polynom.py index 53aa18ac..dc6e4141 100644 --- a/kgcnn/layers/polynom.py +++ b/kgcnn/layers/polynom.py @@ -2,6 +2,7 @@ import scipy as sp import scipy.special from keras import ops +from keras.layers import Layer from scipy.optimize import brentq @@ -86,11 +87,11 @@ def tf_spherical_bessel_jn_explicit(x, n=0): \sum_{k=0}^{\left\lfloor(n-1)/2\right\rfloor}(-1)^{k}\frac{a_{2k+1}(n+\tfrac{1}{2})}{z^{2k+2}}.` Args: - x (tf.Tensor): Values to compute :math:`j_n(x)` for. + x (Tensor): Values to compute :math:`j_n(x)` for. n (int): Positive integer for the bessel order :math:`n`. Returns: - tf.Tensor: Spherical bessel function of order :math:`n` + Tensor: Spherical bessel function of order :math:`n` """ sin_x = ops.sin(x - n * np.pi / 2) cos_x = ops.cos(x - n * np.pi / 2) @@ -98,17 +99,92 @@ def tf_spherical_bessel_jn_explicit(x, n=0): sum_cos = ops.zeros_like(x) for k in range(int(np.floor(n / 2)) + 1): if 2 * k < n + 1: - prefactor = float(sp.special.factorial(n + 2 * k) / np.power(2, 2 * k) / sp.special.factorial( + prefactor_sin = float(sp.special.factorial(n + 2 * k) / np.power(2, 2 * k) / sp.special.factorial( 2 * k) / sp.special.factorial(n - 2 * k) * np.power(-1, k)) - sum_sin += prefactor * ops.power(x, - (2 * k + 1)) + sum_sin += prefactor_sin * ops.power(x, - (2 * k + 1)) for k in range(int(np.floor((n - 1) / 2)) + 1): if 2 * k + 1 < n + 1: - prefactor = float(sp.special.factorial(n + 2 * k + 1) / np.power(2, 2 * k + 1) / sp.special.factorial( + prefactor_cos = float(sp.special.factorial(n + 2 * k + 1) / np.power(2, 2 * k + 1) / sp.special.factorial( 2 * k + 1) / sp.special.factorial(n - 2 * k - 1) * np.power(-1, k)) - sum_cos += prefactor * ops.power(x, - (2 * k + 2)) + sum_cos += prefactor_cos * ops.power(x, - (2 * k + 2)) return sum_sin * sin_x + sum_cos * cos_x +class SphericalBesselJnExplicit(Layer): + r"""Compute spherical bessel functions :math:`j_n(x)` for constant positive integer :math:`n` explicitly. + TensorFlow has to cache the function for each :math:`n`. No gradient through :math:`n` or very large number + of :math:`n`'s is possible. + The spherical bessel functions and there properties can be looked up at + https://en.wikipedia.org/wiki/Bessel_function#Spherical_Bessel_functions. + For this implementation the explicit expression from https://dlmf.nist.gov/10.49 has been used. + The definition is: + + :math:`a_{k}(n+\tfrac{1}{2})=\begin{cases}\dfrac{(n+k)!}{2^{k}k!(n-k)!},&k=0,1,\dotsc,n\\ + 0,&k=n+1,n+2,\dotsc\end{cases}` + + :math:`\mathsf{j}_{n}\left(z\right)=\sin\left(z-\tfrac{1}{2}n\pi\right)\sum_{k=0}^{\left\lfloor n/2\right\rfloor} + (-1)^{k}\frac{a_{2k}(n+\tfrac{1}{2})}{z^{2k+1}}+\cos\left(z-\tfrac{1}{2}n\pi\right) + \sum_{k=0}^{\left\lfloor(n-1)/2\right\rfloor}(-1)^{k}\frac{a_{2k+1}(n+\tfrac{1}{2})}{z^{2k+2}}.` + """ + + def __init__(self, n=0, **kwargs): + r"""Initialize layer with constant n. + + Args: + n (int): Positive integer for the bessel order :math:`n`. + """ + super(SphericalBesselJnExplicit, self).__init__(**kwargs) + self.n = n + self._pre_factor_sin = [] + self._pre_factor_cos = [] + self._powers_sin = [] + self._powers_cos = [] + for k in range(int(np.floor(n / 2)) + 1): + if 2 * k < n + 1: + fac_sin = float(sp.special.factorial(n + 2 * k) / np.power(2, 2 * k) / sp.special.factorial( + 2 * k) / sp.special.factorial(n - 2 * k) * np.power(-1, k)) + pow_sin = - (2 * k + 1) + self._pre_factor_sin.append(fac_sin) + self._powers_sin.append(pow_sin) + for k in range(int(np.floor((n - 1) / 2)) + 1): + if 2 * k + 1 < n + 1: + fac_cos = float(sp.special.factorial(n + 2 * k + 1) / np.power(2, 2 * k + 1) / sp.special.factorial( + 2 * k + 1) / sp.special.factorial(n - 2 * k - 1) * np.power(-1, k)) + pow_cos = - (2 * k + 2) + self._pre_factor_cos.append(fac_cos) + self._powers_cos.append(pow_cos) + + def build(self, input_shape): + """Build layer.""" + super(SphericalBesselJnExplicit, self).build(input_shape) + + def call(self, x, **kwargs): + """Element-wise operation. + + Args: + x (Tensor): Values to compute :math:`j_n(x)` for. + + Returns: + Tensor: Spherical bessel function of order :math:`n` + """ + n = self.n + sin_x = ops.sin(x - n * np.pi / 2) + cos_x = ops.cos(x - n * np.pi / 2) + sum_sin = ops.zeros_like(x) + sum_cos = ops.zeros_like(x) + for a, r in zip(self._pre_factor_sin, self._powers_sin): + sum_sin += a * ops.power(x, r) + for b, s in zip(self._pre_factor_cos, self._powers_cos): + sum_cos += b * ops.power(x, s) + return sum_sin * sin_x + sum_cos * cos_x + + def get_config(self): + """Update layer config.""" + config = super(SphericalBesselJnExplicit, self).get_config() + config.update({"n": self.n}) + return config + + def tf_spherical_bessel_jn(x, n=0): r"""Compute spherical bessel functions :math:`j_n(x)` for constant positive integer :math:`n` via recursion. TensorFlow has to cache the function for each :math:`n`. No gradient through :math:`n` or very large number @@ -126,11 +202,11 @@ def tf_spherical_bessel_jn(x, n=0): :math:`j_{2}(x)=\left(\frac{3}{x^{2}} - 1\right)\frac{\sin x}{x} - \frac{3}{x}\frac{\cos x}{x}` Args: - x (tf.Tensor): Values to compute :math:`j_n(x)` for. + x (Tensor): Values to compute :math:`j_n(x)` for. n (int): Positive integer for the bessel order :math:`n`. Returns: - tf.tensor: Spherical bessel function of order :math:`n` + Tensor: Spherical bessel function of order :math:`n` """ if n < 0: raise ValueError("Order parameter must be >= 0 for this implementation of spherical bessel function.") @@ -158,11 +234,11 @@ def tf_legendre_polynomial_pn(x, n=0): :math:`P_n(x)=\sum_{k=0}^{\lfloor n/2\rfloor} (-1)^k \frac{(2n - 2k)! \, }{(n-k)! \, (n-2k)! \, k! \, 2^n} x^{n-2k}` Args: - x (tf.Tensor): Values to compute :math:`P_n(x)` for. + x (Tensor): Values to compute :math:`P_n(x)` for. n (int): Positive integer for :math:`n` in :math:`P_n(x)`. Returns: - tf.tensor: Legendre Polynomial of order :math:`n`. + Tensor: Legendre Polynomial of order :math:`n`. """ out_sum = ops.zeros_like(x) prefactors = [ @@ -174,6 +250,56 @@ def tf_legendre_polynomial_pn(x, n=0): return out_sum +class LegendrePolynomialPn(Layer): + r"""Compute the (non-associated) Legendre polynomial :math:`P_n(x)` for constant positive integer :math:`n` + via explicit formula. + TensorFlow has to cache the function for each :math:`n`. No gradient through :math:`n` or very large number + of :math:`n` is possible. + Closed form can be viewed at https://en.wikipedia.org/wiki/Legendre_polynomials. + + :math:`P_n(x)=\sum_{k=0}^{\lfloor n/2\rfloor} (-1)^k \frac{(2n - 2k)! \, }{(n-k)! \, (n-2k)! \, k! \, 2^n} x^{n-2k}` + + """ + + def __init__(self, n=0, **kwargs): + r"""Initialize layer with constant n. + + Args: + n (int): Positive integer for :math:`n` in :math:`P_n(x)`. + """ + super(LegendrePolynomialPn, self).__init__(**kwargs) + self.n = n + self._pre_factors = [ + float((-1) ** k * sp.special.factorial(2 * n - 2 * k) / sp.special.factorial(n - k) / sp.special.factorial( + n - 2 * k) / sp.special.factorial(k) / 2 ** n) for k in range(0, int(np.floor(n / 2)) + 1) + ] + self._powers = [float(n - 2 * k) for k in range(0, int(np.floor(n / 2)) + 1)] + + def build(self, input_shape): + """Build layer.""" + super(LegendrePolynomialPn, self).build(input_shape) + + def call(self, x, **kwargs): + """Element-wise operation. + + Args: + x (Tensor): Values to compute :math:`P_n(x)` for. + + Returns: + Tensor: Legendre Polynomial of order :math:`n`. + """ + out_sum = ops.zeros_like(x) + for a, r in zip(self._pre_factors, self._powers): + out_sum = out_sum + a * ops.power(x, r) + return out_sum + + def get_config(self): + """Update layer config.""" + config = super(LegendrePolynomialPn, self).get_config() + config.update({"n": self.n}) + return config + + def tf_spherical_harmonics_yl(theta, l=0): r"""Compute the spherical harmonics :math:`Y_{ml}(\cos\theta)` for :math:`m=0` and constant non-integer :math:`l`. TensorFlow has to cache the function for each :math:`l`. No gradient through :math:`l` or very large number @@ -188,11 +314,11 @@ def tf_spherical_harmonics_yl(theta, l=0): :math:`P_n(x)=\sum_{k=0}^{\lfloor n/2\rfloor} (-1)^k \frac{(2n - 2k)! \, }{(n-k)! \, (n-2k)! \, k! \, 2^n} x^{n-2k}` Args: - theta (tf.Tensor): Values to compute :math:`Y_l(\cos\theta)` for. + theta (Tensor): Values to compute :math:`Y_l(\cos\theta)` for. l (int): Positive integer for :math:`l` in :math:`Y_l(\cos\theta)`. Returns: - tf.tensor: Spherical harmonics for :math:`m=0` and constant non-integer :math:`l`. + Tensor: Spherical harmonics for :math:`m=0` and constant non-integer :math:`l`. """ x = ops.cos(theta) out_sum = ops.zeros_like(x) @@ -206,6 +332,62 @@ def tf_spherical_harmonics_yl(theta, l=0): return out_sum +class SphericalHarmonicsYl(Layer): + r"""Compute the spherical harmonics :math:`Y_{ml}(\cos\theta)` for :math:`m=0` and constant non-integer :math:`l`. + TensorFlow has to cache the function for each :math:`l`. No gradient through :math:`l` or very large number + of :math:`n` is possible. Uses a simplified formula with :math:`m=0` from + https://en.wikipedia.org/wiki/Spherical_harmonics: + + :math:`Y_{l}^{m}(\theta ,\phi)=\sqrt{\frac{(2l+1)}{4\pi} \frac{(l -m)!}{(l +m)!}} \, P_{l}^{m}(\cos{\theta }) \, + e^{i m \phi}` + + where the associated Legendre polynomial simplifies to :math:`P_l(x)` for :math:`m=0`: + + :math:`P_n(x)=\sum_{k=0}^{\lfloor n/2\rfloor} (-1)^k \frac{(2n - 2k)! \, }{(n-k)! \, (n-2k)! \, k! \, 2^n} x^{n-2k}` + + """ + + def __init__(self, l=0, **kwargs): + r"""Initialize layer with constant l. + + Args: + l (int): Positive integer for :math:`l` in :math:`Y_l(\cos\theta)`. + """ + super(SphericalHarmonicsYl, self).__init__(**kwargs) + self.l = l + self._pre_factors = [ + float((-1) ** k * sp.special.factorial(2 * l - 2 * k) / sp.special.factorial(l - k) / sp.special.factorial( + l - 2 * k) / sp.special.factorial(k) / 2 ** l) for k in range(0, int(np.floor(l / 2)) + 1)] + self._powers = [float(l - 2 * k) for k in range(0, int(np.floor(l / 2)) + 1)] + self._scale = float(np.sqrt((2 * l + 1) / 4 / np.pi)) + + def build(self, input_shape): + """Build layer.""" + super(SphericalHarmonicsYl, self).build(input_shape) + + def call(self, theta, **kwargs): + """Element-wise operation. + + Args: + theta (Tensor): Values to compute :math:`Y_l(\cos\theta)` for. + + Returns: + Tensor: Spherical harmonics for :math:`m=0` and constant non-integer :math:`l`. + """ + x = ops.cos(theta) + out_sum = ops.zeros_like(x) + for a, r in zip(self._pre_factors, self._powers): + out_sum = out_sum + a * ops.power(x, r) + out_sum = out_sum * self._scale + return out_sum + + def get_config(self): + """Update layer config.""" + config = super(SphericalHarmonicsYl, self).get_config() + config.update({"l": self.l}) + return config + + def tf_associated_legendre_polynomial(x, l=0, m=0): r"""Compute the associated Legendre polynomial :math:`P_{l}^{m}(x)` for :math:`m` and constant positive integer :math:`l` via explicit formula. @@ -215,12 +397,12 @@ def tf_associated_legendre_polynomial(x, l=0, m=0): \cdot \binom{l}{k}\binom{\frac{l+k-1}{2}}{l}`. Args: - x (tf.Tensor): Values to compute :math:`P_{l}^{m}(x)` for. + x (Tensor): Values to compute :math:`P_{l}^{m}(x)` for. l (int): Positive integer for :math:`l` in :math:`P_{l}^{m}(x)`. m (int): Positive/Negative integer for :math:`m` in :math:`P_{l}^{m}(x)`. Returns: - tf.tensor: Legendre Polynomial of order n. + Tensor: Legendre Polynomial of order n. """ if np.abs(m) > l: raise ValueError("Error: Legendre polynomial must have -l<= m <= l") @@ -239,4 +421,72 @@ def tf_associated_legendre_polynomial(x, l=0, m=0): sp.special.factorial(k) / sp.special.factorial(k - m) * sp.special.binom(l, k) * sp.special.binom((l + k - 1) / 2, l)) - return sum_out * x_prefactor * neg_m \ No newline at end of file + return sum_out * x_prefactor * neg_m + + +class AssociatedLegendrePolynomialPlm(Layer): + r"""Compute the associated Legendre polynomial :math:`P_{l}^{m}(x)` for :math:`m` and constant positive + integer :math:`l` via explicit formula. + Closed Form from taken from https://en.wikipedia.org/wiki/Associated_Legendre_polynomials. + + :math:`P_{l}^{m}(x)=(-1)^{m}\cdot 2^{l}\cdot (1-x^{2})^{m/2}\cdot \sum_{k=m}^{l}\frac{k!}{(k-m)!}\cdot x^{k-m} + \cdot \binom{l}{k}\binom{\frac{l+k-1}{2}}{l}`. + """ + + def __init__(self, l: int = 0, m: int = 0, **kwargs): + r"""Initialize layer with constant m, l. + + Args: + l (int): Positive integer for :math:`l` in :math:`P_{l}^{m}(x)`. + m (int): Positive/Negative integer for :math:`m` in :math:`P_{l}^{m}(x)`. + """ + super(AssociatedLegendrePolynomialPlm, self).__init__(**kwargs) + self.m = m + self.l = l + if np.abs(m) > l: + raise ValueError("Error: Legendre polynomial must have -l<= m <= l") + if l < 0: + raise ValueError("Error: Legendre polynomial must have l>=0") + if m < 0: + m = -m + neg_m = float(np.power(-1, m) * sp.special.factorial(l - m) / sp.special.factorial(l + m)) + else: + neg_m = 1 + self._m = m + self._neg_m = neg_m + self._x_pre_factor = float(np.power(-1, m) * np.power(2, l)) + self._powers = [] + self._pre_factors = [] + for k in range(m, l + 1): + self._powers.append(k - m) + fac = float( + sp.special.factorial(k) / sp.special.factorial(k - m) * sp.special.binom(l, k) * + sp.special.binom((l + k - 1) / 2, l)) + self._pre_factors.append(fac) + + def build(self, input_shape): + """Build layer.""" + super(AssociatedLegendrePolynomialPlm, self).build(input_shape) + + def call(self, x, **kwargs): + """Element-wise operation. + + Args: + x (Tensor): Values to compute :math:`P_{l}^{m}(x)` for. + + Returns: + Tensor: Legendre Polynomial of order n. + """ + neg_m = self._neg_m + m = self._m + x_pre_factor = ops.power(1 - ops.square(x), m / 2) * self._x_pre_factor + sum_out = ops.zeros_like(x) + for a, r in zip(self._pre_factors, self._powers): + sum_out += ops.power(x, r) * a + return sum_out * x_pre_factor * neg_m + + def get_config(self): + """Update layer config.""" + config = super(AssociatedLegendrePolynomialPlm, self).get_config() + config.update({"l": self.l, "m": self.m}) + return config diff --git a/kgcnn/layers/relational.py b/kgcnn/layers/relational.py index 32ea5a55..0fbc5aaa 100644 --- a/kgcnn/layers/relational.py +++ b/kgcnn/layers/relational.py @@ -242,6 +242,7 @@ def batch_dot(x, k): return ops.sum(ops.expand_dims(x, axis=-1) * k, axis=-2, keepdims=False) def get_config(self): + """Update layer config.""" config = super().get_config() config.update({ "units": self.units,