This tutorial provide a stand-alone implementation of spin-2 SH with order $l=2$ and relavent S2L2 representations of Stokes vectors. It will be helpful for users about to reproduce our S2L2 representation method.
from typing import Literal
import math
import numpy as np
from scipy.spatial.transform import Rotation
where ${}_2Y_{lm}\coloneqq \left[\overset\leftrightarrow Y_{lm}\right]^{{\vec{\mathbf F}}_{\theta\phi}}$.
def S2L2_basis_wrtTP(theta, phi):
"""
Parameters
----------
theta: ArrayLike[*]
In the radian units. having a broadcastable shape with `phi`.
phi: ArrayLike[*]
In the radian units. having a broadcastable shape with `theta`.
Return
------
s2Y_2m: np.ndarray[*, 5], complex
Complex values for spin-2 SH, with respect to the $\theta\phi$-frame field, with $l=2$ at given values of the spherical coordinates.
`s2Y_2m[..., 0]` to `s2Y_2m[..., 4]` indicate $m=-2,\cdots,2$, respectively.
"""
theta, phi = np.broadcast_arrays(theta, phi)
cos = np.cos(theta)
sin = np.sin(theta)
cos_p1 = cos + 1
cos_m1 = cos - 1
s2Y_2m = np.full((*theta.shape, 5), 1/8*np.sqrt(5/np.pi), dtype=complex)
s2Y_2m[..., 0] *= (cos_p1)**2
s2Y_2m[..., 1] *= -2*sin*cos_p1
s2Y_2m[..., 2] *= np.sqrt(6) * (sin**2)
s2Y_2m[..., 3] *= 2*sin*cos_m1
s2Y_2m[..., 4] *= (cos_m1)**2
del cos_p1, cos_m1, sin
m = np.arange(-2, 3)
s2Y_2m[...] *= np.exp(m * 1j * phi[..., None])
return s2Y_2m
In many cases, we already has our measurement frame (local frames) $\vec{\mathbf F}\in\vec{\mathbb F}^3$ and then want Stokes parameters with respect to such frame $\vec{\mathbf F}$. Here we introduce two ways.
The first way is utilizing ${}_2Y_{lm}\left(\theta,\phi\right)$, which indicates complex-values Stokes parameters $s_1+is_2$ with respect to the $\theta\phi$-frame field $\vec{\mathbf F}_{\theta\phi}$. Then we can convert it with respect to the frame $\vec{\mathbf F}$ using the angle between these two frames, i.e.,
$$ \begin{align} {}_2Y_{lm}\left(\vec {\mathbf F}\right) &\coloneqq \left[\overset\leftrightarrow Y_{lm}\left(\vec{\mathbf F}\hat z_g\right)\right]^{{\vec{\mathbf F}}} = {}_2Y_{lm}\left(\theta,\phi\right) e^{-2i\psi}, \\ \text{where } \vec{\mathbf F} &= \vec {\mathbf F}_{\theta\phi} \begin{bmatrix} \cos\psi & -\sin\psi & 0 \\ \sin\psi & \cos\psi & 0 \\ 0 & 0 & 1 \end{bmatrix}. \end{align} $$Note that $\vec{\mathbf F}_g=\begin{bmatrix}\hat x_g & \hat y_g & \hat z_g\end{bmatrix}$ denotes the global (world) frame, so that $\vec{\mathbf F}\hat z_g$ indicates the ray direction (local z-axis) of the measurement frame $\vec {\mathbf F}$ and $\theta$, $\phi$, and $\psi$ can be obtained from ZYZ Euler angles of the rotation $\vec{\mathbf F}$ with respect to $\vec{\mathbf F}_g$, i.e.,
$$ \vec{\mathbf F} = \vec R_{\hat z_g \hat y_g \hat z_g}\left(\phi, \theta,\psi \right) \vec{\mathbf F}_g. $$Since $\vec{\mathbf F}_g$ is the world frame, we can simply consider it a the 3D identity matrix in the numerical implementation.
However, there is more elegant implementation using quaternion representation $\mathbf q = q_0+q_1i+q_2j+q_3k$ of the rotation $\vec R$ such that $\vec{\mathbf F} = \vec R \vec{\mathbf F}_g$. Simply denoting ${}_2Y_{lm}\left(\mathbf q\right) = {}_2Y_{lm}\left(\vec{\mathbf F}\right)$, the following formula can be driven through some detailed steps. Here we use two auxiliary complex numbers $\mathbf q_a = q_0 + q_3i$ and $\mathbf q_b = q_2 + q_1i$.
$$ \begin{align} {}_2Y_{2,-2}\left({\mathbf q}\right) &= \sqrt\frac{5}{4\pi} \bar{\mathbf q}_a^4 \\ {}_2Y_{2,-1}\left({\mathbf q}\right) &= -\sqrt\frac{5}{\pi} \bar{\mathbf q}_a^3 \bar{\mathbf q}_b \\ {}_2Y_{20}\left({\mathbf q}\right) &= \sqrt\frac{15}{2\pi}\bar{\mathbf q}_a^2 \bar{\mathbf q}_b^2 \\ {}_2Y_{21}\left({\mathbf q}\right) &= -\sqrt\frac{5}{\pi} \bar{\mathbf q}_a \bar{\mathbf q}_b^3 \\ {}_2Y_{22}\left({\mathbf q}\right) &= \sqrt\frac{5}{4\pi} \bar{\mathbf q}_b^{4} \\ \end{align} $$def S2L2_basis_at_frame(F, method: Literal['euler', 'quat'] = 'quat'):
"""
Parameters
----------
F: ArrayLike[*, 3, 3]
An ND-array of local (measurement) frames. It can also be considered a local-to-world coordinate conversion matrix,
i.e. `F[...,:,0]`, `F[...,:,1]`, `F[...,:,2]` indicate x-, y-, and z-axis vectors in the global (world) coordinates.
On the other hand, `F[...,0,:]`, `F[...,1,:]`, `F[...,2,:]` indicate global x-, y-, and z-coordinates of each local axis vectors
method: literal 'euler', 'quat'
Implementation to evaluate spin-2 SH basis functions.
* `'euler'`: It first calls the `S2L2_basis_wrtTP()` function, which returns values of the basis functions with respect to the theta-phi-frame field.
Then find `psi` which is the angle between given local frame `F` and the theta-phi-frame, around their common z-axis.
The final value becomes the output of `S2L2_basis_wrtTP()` multiplied by $e^{-2i\psi}$.
The name 'euler' of this method comes from the fact that `F` is equivalent to the rotation with ZYZ Euler angles phi, theta, psi.
* `'quat'`: The values of the basis functions are evaluated directly from the quaternion formulation of `F`.
Return
------
s2Y_2m: np.ndarray[*, 5], complex
Complex values for spin-2 SH, with respect to given local frames `F`, with $l=2$ at given values of the spherical coordinates.
`s2Y_2m[..., 0]` to `s2Y_2m[..., 4]` indicate $m=-2,\cdots,2$, respectively.
"""
F = np.asarray(F)
shape_chan = F.shape[:-2]
assert F.shape[-2:] == (3, 3)
R = Rotation.from_matrix(F.reshape(-1, 3, 3))
if method == 'euler':
phi, theta, psi = np.moveaxis(R.as_euler('ZYZ', degrees=False), -1, 0)
s2Y_2m = S2L2_basis_wrtTP(theta, phi)
s2Y_2m *= np.exp(-2j*psi[..., None])
else:
assert method == 'quat'
# [WARNING] Scipy uses scalar-last order at default.
# https://docs.scipy.org/doc/scipy/reference/generated/scipy.spatial.transform.Rotation.as_quat.html
q = R.as_quat()
qa_conj = q[..., 3] - 1j*q[..., 2]
qb_conj = q[..., 1] - 1j*q[..., 0]
s2Y_2m = np.full((math.prod(shape_chan), 5), np.sqrt(5/np.pi), complex)
s2Y_2m[..., 0] *= 0.5 * qa_conj**4
s2Y_2m[..., 1] *= -qa_conj**3 * qb_conj
s2Y_2m[..., 2] *= np.sqrt(3/2) * (qa_conj * qb_conj)**2
s2Y_2m[..., 3] *= -qa_conj * qb_conj**3
s2Y_2m[..., 4] *= 0.5 * qb_conj**4
return s2Y_2m.reshape(*shape_chan, 5)
theta, phi = np.mgrid[:np.pi:20j, :2*np.pi:40j]
psi = np.linspace(0, 2*np.pi, 10)[:, None, None]
angles = np.stack(np.broadcast_arrays(phi, theta, psi), -1).reshape(-1, 3)
F = Rotation.from_euler('ZYZ', angles).as_matrix()
s2Y_2m_euler = S2L2_basis_at_frame(F, 'euler')
s2Y_2m_quat = S2L2_basis_at_frame(F, 'quat')
assert np.allclose(s2Y_2m_euler, s2Y_2m_quat)
print("Clear!")
Clear!
Given Stokes vector $\overset\leftrightarrow s = \left[\mathbf s\right]_{\vec{\mathbf F}}$, its 12-parameter S2L2 representation $\mathbf r = \mathrm{S2L2}\left(\overset\leftrightarrow s\right) \in \mathbb R^{12}$ is characterized as:
$$ \begin{align} r_0 &= s_0, \\ \begin{bmatrix} r_1 \\ r_2 \end{bmatrix} &= \sqrt\frac{4\pi}5 \mathbb R^2 \left( {}_2Y_{2,-2}\left(\vec{\mathbf F}\right)^*\left(s_1+is_2\right) \right), \\ \begin{bmatrix} r_3 \\ r_4 \end{bmatrix} &= \sqrt\frac{4\pi}5 \mathbb R^2 \left( {}_2Y_{2,-1}\left(\vec{\mathbf F}\right)^*\left(s_1+is_2\right) \right), \\ \begin{bmatrix} r_5 \\ r_6 \end{bmatrix} &= \sqrt\frac{4\pi}5 \mathbb R^2 \left( {}_2Y_{20}\left(\vec{\mathbf F}\right)^*\left(s_1+is_2\right) \right), \\ \begin{bmatrix} r_7 \\ r_8 \end{bmatrix} &= \sqrt\frac{4\pi}5 \mathbb R^2 \left( {}_2Y_{21}\left(\vec{\mathbf F}\right)^*\left(s_1+is_2\right) \right), \\ \begin{bmatrix} r_9 \\ r_{10} \end{bmatrix} &= \sqrt\frac{4\pi}5 \mathbb R^2 \left( {}_2Y_{22}\left(\vec{\mathbf F}\right)^*\left(s_1+is_2\right) \right), \\ r_{11} &= s_3, \end{align} $$where $\mathbb R^2$ denotes conversion from complex numbers to 2D numeric real vectors as $\mathbb R^2\left(x+yi\right)=\left[x,y\right]^T$.
def stk_to_S2L2repr(stkComp, F, method: Literal['euler', 'quat'] = 'quat'):
"""
Parameters
----------
stkComp: ArrayLike[*, 4|2]
Numerical vectors of Stokes parameters, which should be considered with the measurement frame `F`.
F: ArrayLike[*, 3, 3]
Measurement frames
method: literal 'euler', 'quat'
Implementation to evaluate spin-2 SH basis functions.
* `'euler'`: It first calls the `S2L2_basis_wrtTP()` function, which returns values of the basis functions with respect to the theta-phi-frame field.
Then find `psi` which is the angle between given local frame `F` and the theta-phi-frame, around their common z-axis.
The final value becomes the output of `S2L2_basis_wrtTP()` multiplied by $e^{-2i\psi}$.
The name 'euler' of this method comes from the fact that `F` is equivalent to the rotation with ZYZ Euler angles phi, theta, psi.
* `'quat'`: The values of the basis functions are evaluated directly from the quaternion formulation of `F`.
Return
------
res: np.ndarray[*, 12|10 or 7|5]
S2L2 representation. It consists of 12 or 7 parameters for each Stokes vector.
If spin-2 Stokes vector (only consists of $s_1$ and $s_$ Stokes parameters) was given,
the number of output parameters changes to 10 for `type=="COMP"` and 5 for `type=="REAL"`, respectively.
"""
p = stkComp.shape[-1]
assert p in [2, 4], f"{stkComp.shape = }"
assert F.shape[-2:] == (3, 3), f"{F.shape = }"
shape_chan = np.broadcast_shapes(stkComp.shape[:-1], F.shape[:-2])
shape = (*shape_chan, p+8) # [*, 12] or [*, 10]
res = np.empty(shape)
# ---------- Spin-0 ----------
if p == 4:
res[..., 0] = stkComp[..., 0]
res[..., -1] = stkComp[..., 3]
s1, s2 = stkComp[..., 1], stkComp[..., 2]
idx_s12 = np.s_[1:-1]
else:
s1, s2 = stkComp[..., 0], stkComp[..., 1]
idx_s12 = np.s_[:]
# ---------- Spin-2 ----------
s1_is2 = s1 + 1j*s2 # shape [*]
s2Y_2m = S2L2_basis_at_frame(F, method) # shape [*, 5]
prod = s2Y_2m.conj() * s1_is2[..., None] # shape [*, 5]
res[..., idx_s12] = np.stack([prod.real, prod.imag], -1).reshape(*shape_chan, 10)
res[..., idx_s12] *= 2*np.sqrt(np.pi/5)
return res
S2L2 representation is a frame-free method, i.e., two pairs of (stkComp, F) arguments under the relationship of Stokes coordiante conversion produce consistent S2L2 representaion. This representation is also linear for a fixed ray direction.
The constant $\sqrt\frac{4\pi}5$ in the above S2L2 equations are intended to preseve norms. Norms, degrees of polarization, unpolarized intensities can be computed from S2L2 representations.
# ---------- Utils ----------
def rotation_z(angle):
return np.array([
[np.cos(angle), -np.sin(angle), 0],
[np.sin(angle), np.cos(angle), 0],
[0, 0, 1]
])
stk = np.random.rand(*F.shape[:-2], 4)
# ---------- S2L2 representation ----------
s2l2 = stk_to_S2L2repr(stk, F)
# ---------- Invariant under Stokes coordinates conversion ----------
for i in range(10):
psi = 2*np.pi*(i+1)/11
F2 = F @ rotation_z(psi)
stk2 = stk.copy()
stk2[..., 1] = np.cos(2*psi) * stk[..., 1] + np.sin(2*psi) * stk[..., 2]
stk2[..., 2] = -np.sin(2*psi) * stk[..., 1] + np.cos(2*psi) * stk[..., 2]
assert np.allclose(s2l2, stk_to_S2L2repr(stk2, F2))
print("Well defined independent of local frames!")
# ---------- Linearity for Stokes vectors with the same ray direction ----------
stk2 = np.random.rand(*F.shape[:-2], 4)
s2l2_2 = stk_to_S2L2repr(stk2, F)
assert np.allclose(s2l2 + s2l2_2, stk_to_S2L2repr(stk + stk2, F))
print("Linear!")
# ---------- Other properties related to norms ----------
assert np.allclose(np.linalg.norm(stk, 2, -1), np.linalg.norm(s2l2, 2, -1))
print("Norms are preserved!")
polar_s0123 = np.sqrt((stk[..., 1:]**2).sum(-1))
polar_s2l2 = np.sqrt((s2l2[..., 1:]**2).sum(-1))
dop_s0123 = polar_s0123 / stk[..., 0]
dop_s2l2 = polar_s2l2 / s2l2[..., 0]
assert np.allclose(dop_s0123, dop_s2l2)
print("Degrees of polarization are preserved!")
assert np.allclose(stk[..., 0] - polar_s0123,
s2l2[..., 0] - polar_s2l2)
print("Unpolarized intensity is preserved!\n")
Well defined independent of local frames! Linear! Norms are preserved! Degrees of polarization are preserved! Unpolarized intensity is preserved!
$\mathbf s = \mathrm{S2L2Inv}\left(\mathbf r; \vec{\mathbf F}\right)\in{\mathbb R}^{4}$
$$ \begin{align} s_0 &= r_0 \\ \begin{bmatrix} s_1 \\ s_2 \end{bmatrix} &= \sqrt\frac{4\pi}{5}\mathbb R^2\left(\sum_{m=-2}^2 {}_2Y_{2m}\left(\vec{\mathbf F}\right)\left(r_{2m+5}+ir_{2m+6}\right)\right) \\ s_3 &= r_{11} \\ \end{align} $$def S2L2repr_to_stk(s2l2, F, method: Literal['euler', 'quat'] = 'quat'):
"""
Parameters
----------
s2l2: ArrayLike[*, 12|10]
S2L2 representation. It consists of 12 parameters for each Stokes vector.
F: ArrayLike[*, 3, 3]
Measurement frames
method: literal 'euler', 'quat'
Implementation to evaluate spin-2 SH basis functions.
* `'euler'`: It first calls the `S2L2_basis_wrtTP()` function, which returns values of the basis functions with respect to the theta-phi-frame field.
Then find `psi` which is the angle between given local frame `F` and the theta-phi-frame, around their common z-axis.
The final value becomes the output of `S2L2_basis_wrtTP()` multiplied by $e^{-2i\psi}$.
The name 'euler' of this method comes from the fact that `F` is equivalent to the rotation with ZYZ Euler angles phi, theta, psi.
* `'quat'`: The values of the basis functions are evaluated directly from the quaternion formulation of `F`.
Return
------
stk: np.ndarray[*, 4]
Numerical vectors of Stokes parameters, which should be considered with the measurement frame `F`.
"""
F = np.asarray(F)
s2l2 = np.asarray(s2l2)
assert F.shape[-2:] == (3, 3), f"{F.shape = }"
n_param = s2l2.shape[-1]
assert n_param in [10, 12], f"Invalid shape: {s2l2.shape = }"
if n_param % 5 == 0:
n_stk = 2 # Spin-2 Stokes vector only
idx_s12 = np.s_[:]
idx_s12_p1 = np.s_[0::2]
idx_s12_p2 = np.s_[1::2]
oidx_s1, oidx_s2 = 0, 1
else:
n_stk = 4 # Full Stokes vector
idx_s12 = np.s_[1:-1]
idx_s12_p1 = np.s_[1:-1:2]
idx_s12_p2 = np.s_[2:-1:2]
oidx_s1, oidx_s2 = 1, 2
shape_chan = np.broadcast_shapes(s2l2.shape[:-1], F.shape[:-2])
shape_stk = (*shape_chan, n_stk)
stk = np.empty(shape_stk)
# ---------- Spin-0 ----------
if n_stk == 4:
stk[..., 0] = s2l2[..., 0]
stk[..., -1] = s2l2[..., -1]
# ---------- Spin-2 ----------
s2l2_5 = s2l2[..., idx_s12_p1] + s2l2[..., idx_s12_p2]*1j # shape [*, 5] complex
factor = np.sqrt(4*np.pi/5)
s2Y_2m = S2L2_basis_at_frame(F, method) # shape [*, 5] complex
dot = (s2Y_2m * s2l2_5).sum(-1) # shape [*] complex
stk[..., oidx_s1] = factor * dot.real
stk[..., oidx_s2] = factor * dot.imag
return stk
assert np.allclose(stk, S2L2repr_to_stk(s2l2, F))
print("The inversion is correct!")
The inversion is correct!