def test_psichi(self):
for _ in range(100):
theta = np.random.rand() * 2 * np.pi
for k in range(10):
for _ in range(5):
for s in range(2):
gamma = np.random.rand() * 2 * np.pi
g_psi = group.psichi(theta, s, k, gamma)
k_psi = ks.psichi(theta, s, k, gamma).squeeze()
self.assertTrue(np.allclose(g_psi, k_psi))
def sample(self, angles: np.ndarray, out: np.ndarray = None) -> np.ndarray:
r"""
Sample the continuous basis elements on the discrete set of angles in ``angles``.
Optionally, store the resulting multidimentional array in ``out``.
A value of ``nan`` is interpreted as the angle of a point placed on the origin of the axes.
``angles`` must be an array of shape `(1, N)`, where `N` is the number of points.
Args:
angles (~numpy.ndarray): angles where to evaluate the basis elements
out (~numpy.ndarray, optional): pre-existing array to use to store the output
Returns:
the sampled basis
"""
assert len(angles.shape) == 2
assert angles.shape[0] == 1
if out is None:
out = np.empty(
(self.shape[0], self.shape[1], self.dim, angles.shape[1]))
assert out.shape == (self.shape[0], self.shape[1], self.dim,
angles.shape[1])
# find points in the origin
origin = np.isnan(angles)
angles = angles.copy()
angles[origin] = 0.
angles -= self.axis
# the basis vectors depends on the shape of the input and output irreps,
# while their frequencies depend on the irreps frequencies
if self.shape[0] == 2 and self.shape[1] == 2:
out = psichi(angles,
s=self.s,
k=self.mu,
gamma=self.gamma,
out=out)
elif self.shape[0] == 1 and self.shape[1] == 2:
out[0, 0, ...] = np.cos(self.mu * angles + self.gamma)
out[0, 1, ...] = np.sin(self.mu * angles + self.gamma)
elif self.shape[0] == 2 and self.shape[1] == 1:
out[0, 0, ...] = np.cos(self.mu * angles + self.gamma)
out[1, 0, ...] = np.sin(self.mu * angles + self.gamma)
elif self.shape[0] == 1 and self.shape[1] == 1:
out[0, 0, ...] = np.cos(self.mu * angles + self.gamma)
else:
raise ValueError(f"Shape {self.shape} not recognized!")
if self._has_non_zero_frequencies:
# In the origin, only 0-frequencies are permitted.
# Therefore, any non-0 frequency base is set to 0 in the origin
if np.any(origin):
mask = self._non_zero_frequencies * origin
out *= 1 - mask
return out