from collections.abc import Mapping
from typing import Any, Literal
import numpy as np
from array_api._2024_12 import Array, ArrayNamespaceFull
from array_api_compat import array_namespace
from gumerov_expansion_coefficients import translational_coefficients
from ultrasphere import SphericalCoordinates, get_child
from ultrasphere_harmonics._core._eigenfunction import Phase, minus_1_power
from ._core import concat_harmonics, expand_dims_harmonics
from ._core._flatten import (
flatten_harmonics,
index_array_harmonics,
)
from ._core._harmonics import harmonics
from ._expansion import (
expand,
)
from ._helmholtz import harmonics_regular_singular
def _harmonics_translation_coef_plane_wave[TCartesian, TSpherical](
c: SphericalCoordinates[TSpherical, TCartesian],
cartesian: Mapping[TCartesian, Array],
*,
n_end: int,
n_end_add: int,
phase: Phase,
k: Array,
) -> Array:
r"""
Translation coefficients between same type of elementary solutions.
Returns $(R|R) = (S|S)$, where
.. math::
R(x + t) = \sum_n (R|R)_n(t) R(x) \\
S(x + t) = \sum_n (S|S)_n(t) S(x)
Parameters
----------
c : SphericalCoordinates[TSpherical, TCartesian]
The spherical coordinates.
cartesian : Mapping[TCartesian, Array]
The translation vector in cartesian coordinates.
Each array must have the same shape (...,).
n_end : int
The maximum degree of the harmonic.
n_end_add : int
The maximum degree of the harmonic to be summed over.
phase : Phase
Adjust phase (±) of the spherical harmonics, mainly to match conventions.
See `Phase` for details.
k : Array
The wavenumber.
Returns
-------
Array
The translation coefficients of shape (..., N, N).
Axis -1 is to be summed over with the elementary solutions
to get translated elementary solution
which quantum number is axis -2 indices.
"""
xp = array_namespace(*[cartesian[k] for k in c.c_nodes])
dtype = cartesian[c.c_nodes[0]].dtype
device = cartesian[c.c_nodes[0]].device
_, k = xp.broadcast_arrays(cartesian[c.c_nodes[0]], k)
n = index_array_harmonics(
c, c.root, n_end=n_end, xp=xp, expand_dims=True, flatten=True, device=device
)[:, None]
ns = index_array_harmonics(
c,
c.root,
n_end=n_end_add,
xp=xp,
expand_dims=True,
flatten=True,
device=device,
)[None, :]
def to_expand(spherical: Mapping[TSpherical, Array]) -> Array:
# returns [spherical1,...,sphericalN,user1,...,userM,harmn]
# [spherical1,...,sphericalN,harmn]
Y = harmonics(
c,
spherical,
n_end=n_end,
phase=phase,
expand_dims=True,
concat=True,
flatten=True,
)
x = c.to_cartesian(spherical)
ndim_user = cartesian[c.c_nodes[0]].ndim
ndim_spherical = c.s_ndim
ip = xp.sum(
xp.stack(
xp.broadcast_arrays(
*[
cartesian[i][
(None,) * ndim_spherical + (slice(None),) * ndim_user
]
* x[i][(slice(None),) * ndim_spherical + (None,) * ndim_user]
for i in c.c_nodes
]
),
axis=0,
),
axis=0,
)
# [spherical1,...,sphericalN,user1,...,userM]
e = xp.exp(1j * k[(None,) * ndim_spherical + (slice(None),) * ndim_user] * ip)
result = (
Y[(slice(None),) * ndim_spherical + (None,) * ndim_user + (slice(None),)]
* e[(slice(None),) * (ndim_spherical + ndim_user) + (None,)]
)
return result
# returns [user1,...,userM,harmn,harmn']
return (-1j) ** (n - ns) * expand(
c,
to_expand,
does_f_support_separation_of_variables=False,
n=n_end + n_end_add - 1,
n_end=n_end_add,
phase=phase,
xp=xp,
dtype=dtype,
device=device,
)
[docs]
def harmonics_twins_expansion[TCartesian, TSpherical](
c: SphericalCoordinates[TSpherical, TCartesian],
*,
n_end_1: int,
n_end_2: int,
phase: Phase,
xp: ArrayNamespaceFull,
dtype: Any | None = None,
device: Any | None = None,
conj_1: bool = False,
conj_2: bool = False,
) -> Array:
"""
Expansion coefficients of the twins of the harmonics.
Parameters
----------
c : SphericalCoordinates[TSpherical, TCartesian]
The spherical coordinates.
n_end_1 : int
The maximum degree of the harmonic
for the first harmonics.
n_end_2 : int
The maximum degree of the harmonic
for the second harmonics.
phase : Phase
Adjust phase (±) of the spherical harmonics, mainly to match conventions.
See `Phase` for details.
xp : ArrayNamespaceFull
The array namespace.
dtype : Any | None
The dtype, by default None
device : Any | None
The device, by default None
conj_1 : bool
Whether to conjugate the first harmonics.
by default False
conj_2 : bool
Whether to conjugate the second harmonics.
by default False
Returns
-------
Array
The expansion coefficients of the twins of shape
[*1st quantum number, *2nd quantum number, *3rd quantum number]
and dim `3 * c.s_ndim` and of dtype float, not complex.
The n_end for 1st quantum number is `n_end_1`,
The n_end for 2nd quantum number is `n_end_2`,
The n_end for 3rd quantum number is `n_end_1 + n_end_2 - 1`.
(not `n_end_1` or `n_end_2`)
Notes
-----
To get ∫Y_{n1}(x)Y_{n2}(x)Y_{n3}(x)dx
(integral involving three harmonics),
one may use
`harmonics_twins_expansion(conj_1=True, conj_2=True)`
"""
def to_expand(spherical: Mapping[TSpherical, Array]) -> Array:
# returns [theta,n1,...,nN,nsummed1,...,nsummedN]
# Y(n)Y*(nsummed)
Y1 = harmonics(
c,
spherical,
n_end=n_end_1,
phase=phase,
expand_dims=True,
concat=False,
)
Y1 = {k: v[(...,) + (None,) * c.s_ndim] for k, v in Y1.items()}
if conj_1:
Y1 = {k: xp.conj(v) for k, v in Y1.items()}
Y2 = harmonics(
c,
spherical,
n_end=n_end_2,
phase=phase,
expand_dims=True,
concat=False,
)
Y2 = {
k: v[(...,) + (None,) * c.s_ndim + (slice(None),) * c.s_ndim]
for k, v in Y2.items()
}
if conj_2:
Y2 = {k: xp.conj(v) for k, v in Y2.items()}
return {k: Y1[k] * Y2[k] for k in c.s_nodes}
# returns [user1,...,userM,n1,...,nN,np1,...,npN]
result: Mapping[TSpherical, Array] = expand(
c,
to_expand,
does_f_support_separation_of_variables=True,
n=n_end_1 + n_end_2 - 1, # at least n_end + 2
n_end=n_end_1 + n_end_2 - 1,
phase=phase,
xp=xp,
dtype=dtype,
device=device,
)
result = expand_dims_harmonics(c, result)
result = concat_harmonics(c, result)
result = xp.real(result)
result = flatten_harmonics(c, result)
result = flatten_harmonics(
c,
result,
axis_end=-2, # , n_end=n_end_2, include_negative_m=True
)
result = flatten_harmonics(
c,
result,
axis_end=-3, # , n_end=n_end_1, include_negative_m=True
)
return result
def _harmonics_translation_coef_triplet[TCartesian, TSpherical](
c: SphericalCoordinates[TSpherical, TCartesian],
spherical: Mapping[TSpherical | Literal["r"], Array],
*,
n_end: int,
n_end_add: int,
phase: Phase,
k: Array,
is_type_same: bool,
) -> Array:
r"""
Translation coefficients between same or different type of elementary solutions.
If is_type_same is True, returns $(R|R) = (S|S)$.
If is_type_same is False, returns $(S|R)$.
.. math::
R(x + t) = \sum_n (R|R)_n(t) R(x) \\
S(x + t) = \sum_n (S|S)_n(t) S(x) \\
S(x + t) = \sum_n (S|R)_n(t) R(x)
Parameters
----------
c : SphericalCoordinates[TSpherical, TCartesian]
The spherical coordinates.
spherical : Mapping[TSpherical, Array]
The translation vector in spherical coordinates.
n_end : int
The maximum degree of the harmonic.
n_end_add : int
The maximum degree of the harmonic to be summed over.
phase : Phase
Adjust phase (±) of the spherical harmonics, mainly to match conventions.
See `Phase` for details.
k : Array
The wavenumber.
is_type_same : bool
Whether the type of the elementary solutions is same.
Returns
-------
Array
The translation coefficients of shape (..., ndim, ndim).
The last axis are to be summed over with the elementary solutions
to get translated elementary solution which quantum number corresponds to
the second last axis indices.
"""
xp = array_namespace(*[spherical[k] for k in c.s_nodes])
dtype = spherical["r"].dtype
device = spherical["r"].device
# [user1,...,userM,n1,...,nN,nsummed1,...,nsummedN,ntemp1,...,ntempN]
n = index_array_harmonics(
c, c.root, n_end=n_end, expand_dims=True, xp=xp, flatten=True, device=device
)[:, None, None]
ns = index_array_harmonics(
c,
c.root,
n_end=n_end_add,
expand_dims=True,
xp=xp,
flatten=True,
device=device,
)[None, :, None]
ntemp = index_array_harmonics(
c,
c.root,
n_end=n_end + n_end_add - 1,
expand_dims=True,
xp=xp,
flatten=True,
device=device,
)[None, None, :]
# returns [user1,...,userM,n1,...,nN,np1,...,npN]
coef = (2 * xp.pi) ** (c.c_ndim / 2) * np.sqrt(2 / np.pi)
t_RS = harmonics_regular_singular(
c,
spherical,
n_end=n_end + n_end_add - 1,
phase=phase,
expand_dims=True,
concat=True,
k=k,
type="regular" if is_type_same else "singular",
flatten=True,
)
expansion = harmonics_twins_expansion(
c,
n_end_1=n_end,
n_end_2=n_end_add,
phase=phase,
conj_1=False,
conj_2=True,
xp=xp,
dtype=dtype,
device=device,
)
return coef * xp.sum(
(-1j) ** (n - ns - ntemp) * t_RS[..., None, None, :] * expansion,
axis=-1,
)
[docs]
def harmonics_translation_coef[TCartesian, TSpherical](
c: SphericalCoordinates[TSpherical, TCartesian],
spherical: Mapping[TSpherical | Literal["r"], Array],
*,
n_end: int,
n_end_add: int,
phase: Phase,
k: Array,
is_type_same: bool,
method: Literal["gumerov", "plane_wave", "triplet"] | None = None,
) -> Array:
r"""
Translation coefficients between same or different type of elementary solutions.
If is_type_same is True, returns $(R|R) = (S|S)$.
If is_type_same is False, returns $(S|R)$.
.. math::
R(x + t) = \sum_n (R|R)_n(t) R(x) \\
S(x + t) = \sum_n (S|S)_n(t) S(x) \\
S(x + t) = \sum_n (S|R)_n(t) R(x)
Parameters
----------
c : SphericalCoordinates[TSpherical, TCartesian]
The spherical coordinates.
spherical : Mapping[TSpherical, Array]
The translation vector in spherical coordinates.
n_end : int
The maximum degree of the harmonic.
n_end_add : int
The maximum degree of the harmonic to be summed over.
phase : Phase
Adjust phase (±) of the spherical harmonics, mainly to match conventions.
See `Phase` for details.
k : Array
The wavenumber.
is_type_same : bool
Whether the type of the elementary solutions is same.
method : Literal["gumerov", "plane_wave", "triplet"] | None
The method to compute the translation coefficients.
If None, the fastest method is chosen automatically.
"gumerov" is only available for branching type "ba".
"plane_wave" is only available when `is_type_same` is True.
Returns
-------
Array
The translation coefficients of shape (..., ndim, ndim).
The last axis are to be summed over with the elementary solutions
to get translated elementary solution which quantum number corresponds to
the second last axis indices.
Example
-------
>>> from array_api_compat import numpy as np
>>> from ultrasphere import create_spherical
>>> c = create_spherical()
>>> t = np.asarray([2, -7, 1])
>>> coef = harmonics_translation_coef(
... c,
... c.from_cartesian(t),
... n_end=2,
... n_end_add=2,
... phase=0,
... k=np.asarray(1.0),
... is_type_same=True,
... )
>>> np.round(coef, 2)
array([[ 0.12+0.j , 0.01+0.j , 0.02+0.06j, 0.02-0.06j],
[-0.01+0.j , -0.01+0.j , 0.01+0.04j, 0.01-0.04j],
[-0.02+0.06j, 0.01-0.04j, 0.18+0.j , -0.17-0.11j],
[-0.02-0.06j, 0.01+0.04j, -0.17+0.11j, 0.18+0.j ]])
>>> x = np.asarray([-1.0, 1.0, 0.0])
>>> y = x + t
>>> coef = harmonics_translation_coef(
... c,
... c.from_cartesian(t),
... n_end=2,
... n_end_add=6,
... phase=0,
... k=np.asarray(1.0),
... is_type_same=True,
... )
>>> R_x = harmonics_regular_singular(
... c,
... c.from_cartesian(x),
... n_end=6,
... phase=0,
... k=np.asarray(1.0),
... type="regular",
... concat=True,
... )
>>> R_y_approx = np.sum(coef * R_x[..., None, :], axis=-1)
>>> R_y = harmonics_regular_singular(
... c,
... c.from_cartesian(y),
... n_end=2,
... phase=0,
... k=np.asarray(1.0),
... type="regular",
... concat=True,
... )
>>> np.round(R_y_approx, 8)
array([-0.00541026-0.j , -0.01301699-0.j ,
-0.00919779+0.05521981j, -0.00919779-0.05521981j])
>>> np.round(R_y, 8)
array([-0.00542242+0.j , -0.01301453+0.j ,
-0.00920266+0.05521598j, -0.00920266-0.05521598j])
"""
xp = array_namespace(*[spherical[k] for k in c.s_nodes], k)
phase = Phase(phase)
device = spherical["r"].device
if method is None:
if c.branching_types_expression_str in ["a", "ba"]:
method = "gumerov"
elif is_type_same:
method = "plane_wave"
else:
method = "triplet"
if method == "gumerov":
if c.branching_types_expression_str == "a":
SR = harmonics_regular_singular(
c,
spherical,
n_end=n_end + n_end_add - 1,
phase=phase,
expand_dims=True,
concat=True,
k=k,
type="regular" if is_type_same else "singular",
flatten=True,
)
n = index_array_harmonics(
c, c.root, n_end=n_end, xp=xp, flatten=True, device=device
)[:, None]
n_add = index_array_harmonics(
c, c.root, n_end=n_end_add, xp=xp, flatten=True, device=device
)[None, :]
result = 2 * SR[..., n - n_add]
if Phase.NEGATIVE_LEGENDRE in phase:
# (- n) at the last means (- n - n_add - (n - n_add)) // 2
result *= minus_1_power(
(xp.abs(n) + xp.abs(n_add) + xp.abs(n - n_add)) // 2 - n
)
return result
elif c.branching_types_expression_str == "ba":
result = translational_coefficients(
k * spherical["r"],
spherical[c.root],
spherical[get_child(c.G, c.root, "sin")],
n_end=max(n_end, n_end_add),
same=is_type_same,
)
result = xp.moveaxis(result, -1, -2)[..., : n_end**2, : n_end_add**2]
if phase == Phase.CONDON_SHORTLEY:
return result
m = index_array_harmonics(
c,
get_child(c.G, c.root, "sin"),
n_end=n_end,
xp=xp,
flatten=True,
device=device,
)[:, None]
m_add = index_array_harmonics(
c,
get_child(c.G, c.root, "sin"),
n_end=n_end_add,
xp=xp,
flatten=True,
device=device,
)[None, :]
if phase == Phase(0):
result *= minus_1_power(m + m_add)
elif phase == Phase.NEGATIVE_LEGENDRE:
result *= minus_1_power(
(xp.abs(m) + m) // 2 + (xp.abs(m_add) + m_add) // 2
)
elif phase == (Phase.NEGATIVE_LEGENDRE | Phase.CONDON_SHORTLEY): # type: ignore[unreachable]
result *= minus_1_power(
(xp.abs(m) - m) // 2 + (xp.abs(m_add) - m_add) // 2
)
return result
else:
raise NotImplementedError()
elif method == "plane_wave":
if not is_type_same:
raise NotImplementedError(
"plane_wave method only supports is_type_same=True"
)
return _harmonics_translation_coef_plane_wave(
c,
cartesian=c.to_cartesian(spherical, as_array=True),
n_end=n_end,
n_end_add=n_end_add,
phase=phase,
k=k,
)
elif method == "triplet":
return _harmonics_translation_coef_triplet(
c,
spherical,
n_end=n_end,
n_end_add=n_end_add,
phase=phase,
k=k,
is_type_same=is_type_same,
)
else:
raise ValueError(f"Invalid method {method}.")