import numpy as np
import matplotlib.pyplot as plt
import spharpy
#
# Define the convention and order of the spherical harmonics and a set of
# coefficients. Note that the generated coefficients correspond to a
# dipole oriented along the y-axis.
#
definition = spharpy.SphericalHarmonicDefinition(n_max=1)
coefficients = np.array([0, 1, 0, 0])
#
# Define a rotation by 45 degrees around the z-axis based on the Euler
# angles.
#
angle = np.pi / 4
R = spharpy.transforms.SphericalHarmonicRotation.from_euler(
    'z', angle)
#
# The spherical harmonic rotation matrix can be obtained and applied
# to the spherical harmonic coefficients using the following:
#
D = R.as_spherical_harmonic_matrix(definition)
rotated_coefficients = D @ coefficients
print(f"Rotated coefficients: {np.round(rotated_coefficients, 2)}")
# Rotated coefficients: [ 0.   0.71 0.   -0.71]
#
# To visualize the effect of the rotation, we can expand the series
# expansion on the unit sphere and plot the original and rotated
# data:
#
sampling = spharpy.samplings.equal_area(0, n_points=250)
Y = spharpy.SphericalHarmonics.from_definition(
    definition, coordinates=sampling)
_, axs = plt.subplots(
    1, 2, subplot_kw={'projection': '3d'}, figsize=(5, 2.5),
    layout='constrained')
spharpy.plot.balloon_wireframe(
    sampling, Y.basis @ coefficients, ax=axs[0], colorbar=False)
spharpy.plot.balloon_wireframe(
    sampling, Y.basis @ rotated_coefficients, ax=axs[1],
    colorbar=False)