Rotation3d

In [ ]:

import numpy as np
import plotly.graph_objects as go
from scipy.linalg import expm
from scipy.optimize import least_squares

import qubex as qx

import numpy as np import plotly.graph_objects as go from scipy.linalg import expm from scipy.optimize import least_squares import qubex as qx

In [ ]:

# Generators of SO(3)

G_x = np.array(
    [
        [0, 0, 0],
        [0, 0, -1],
        [0, 1, 0],
    ]
)

G_y = np.array(
    [
        [0, 0, 1],
        [0, 0, 0],
        [-1, 0, 0],
    ]
)

G_z = np.array(
    [
        [0, -1, 0],
        [1, 0, 0],
        [0, 0, 0],
    ]
)


def rotation_matrix(
    t: float,
    Omega: float,
    n: tuple[float, float, float],
    use_expm: bool = False,
) -> np.ndarray:
    """Return the SO(3) rotation matrix for axis and angle."""
    G = n[0] * G_x + n[1] * G_y + n[2] * G_z
    if use_expm:
        R = expm(Omega * t * G)
    else:
        # Rodrigues' rotation formula
        R = np.eye(3) + np.sin(Omega * t) * G + (1 - np.cos(Omega * t)) * G @ G
    return R

# Generators of SO(3) G_x = np.array( [ [0, 0, 0], [0, 0, -1], [0, 1, 0], ] ) G_y = np.array( [ [0, 0, 1], [0, 0, 0], [-1, 0, 0], ] ) G_z = np.array( [ [0, -1, 0], [1, 0, 0], [0, 0, 0], ] ) def rotation_matrix( t: float, Omega: float, n: tuple[float, float, float], use_expm: bool = False, ) -> np.ndarray: """Return the SO(3) rotation matrix for axis and angle.""" G = n[0] * G_x + n[1] * G_y + n[2] * G_z if use_expm: R = expm(Omega * t * G) else: # Rodrigues' rotation formula R = np.eye(3) + np.sin(Omega * t) * G + (1 - np.cos(Omega * t)) * G @ G return R

In [ ]:

# # Parameters for the rotation
# t = 1.0
# Omega = 1.0
# n_x, n_y, n_z = 1.0, 0.0, 0.0

# # Time with %timeit for use_expm=False
# print("Time with Rodrigues' formula:")
# %timeit rotation_matrix(t, Omega, n_x, n_y, n_z, use_expm=False)

# # Time with %timeit for use_expm=True
# print("Time with expm:")
# %timeit rotation_matrix(t, Omega, n_x, n_y, n_z, use_expm=True)

# # Parameters for the rotation # t = 1.0 # Omega = 1.0 # n_x, n_y, n_z = 1.0, 0.0, 0.0 # # Time with %timeit for use_expm=False # print("Time with Rodrigues' formula:") # %timeit rotation_matrix(t, Omega, n_x, n_y, n_z, use_expm=False) # # Time with %timeit for use_expm=True # print("Time with expm:") # %timeit rotation_matrix(t, Omega, n_x, n_y, n_z, use_expm=True)

In [ ]:

R_x = lambda t, Omega: rotation_matrix(t, Omega, (1, 0, 0))
R_y = lambda t, Omega: rotation_matrix(t, Omega, (0, 1, 0))
R_z = lambda t, Omega: rotation_matrix(t, Omega, (0, 0, 1))

R_x = lambda t, Omega: rotation_matrix(t, Omega, (1, 0, 0)) R_y = lambda t, Omega: rotation_matrix(t, Omega, (0, 1, 0)) R_z = lambda t, Omega: rotation_matrix(t, Omega, (0, 0, 1))

In [ ]:

Omega = 2 * np.pi
r_0 = np.array([1, 1, 1]) / np.sqrt(3)
times = np.linspace(0, 1, 50)
r_x = np.array([R_x(t, Omega) @ r_0 for t in times])
qx.viz.plot_bloch_vectors(times, r_x)
qx.viz.display_bloch_sphere(r_x)

Omega = 2 * np.pi r_0 = np.array([1, 1, 1]) / np.sqrt(3) times = np.linspace(0, 1, 50) r_x = np.array([R_x(t, Omega) @ r_0 for t in times]) qx.viz.plot_bloch_vectors(times, r_x) qx.viz.display_bloch_sphere(r_x)

In [ ]:

Omega = 4 * np.pi
r_0 = np.array([1, 1, 1]) / np.sqrt(3)
times = np.linspace(0, 1, 50)
r_y = np.array([R_y(t, Omega) @ r_0 for t in times])
qx.viz.plot_bloch_vectors(times, r_y)
qx.viz.display_bloch_sphere(r_y)

Omega = 4 * np.pi r_0 = np.array([1, 1, 1]) / np.sqrt(3) times = np.linspace(0, 1, 50) r_y = np.array([R_y(t, Omega) @ r_0 for t in times]) qx.viz.plot_bloch_vectors(times, r_y) qx.viz.display_bloch_sphere(r_y)

In [ ]:

Omega = 6 * np.pi
r_0 = np.array([1, 1, 1]) / np.sqrt(3)
times = np.linspace(0, 1, 50)
r_z = np.array([R_z(t, Omega) @ r_0 for t in times])
qx.viz.plot_bloch_vectors(times, r_z)
qx.viz.display_bloch_sphere(r_z)

Omega = 6 * np.pi r_0 = np.array([1, 1, 1]) / np.sqrt(3) times = np.linspace(0, 1, 50) r_z = np.array([R_z(t, Omega) @ r_0 for t in times]) qx.viz.plot_bloch_vectors(times, r_z) qx.viz.display_bloch_sphere(r_z)

In [ ]:

def create_data(Omega, r_0, n, times, noise, decay_rate):
    """Generate noisy decaying Bloch-vector data."""
    r_t = np.array([rotation_matrix(t, Omega, n) @ r_0 for t in times])
    decay_factor = np.exp(-decay_rate * times)
    return r_t * decay_factor[:, None] + np.random.default_rng().normal(
        0, noise, r_t.shape
    )


Omega = 4 * np.pi
r_0 = np.array([0, 0, 1])
n = np.array([1.0, -0.5, 0.1])
n = n / np.linalg.norm(n)
times = np.linspace(0, 1, 50)

def create_data(Omega, r_0, n, times, noise, decay_rate): """Generate noisy decaying Bloch-vector data.""" r_t = np.array([rotation_matrix(t, Omega, n) @ r_0 for t in times]) decay_factor = np.exp(-decay_rate * times) return r_t * decay_factor[:, None] + np.random.default_rng().normal( 0, noise, r_t.shape ) Omega = 4 * np.pi r_0 = np.array([0, 0, 1]) n = np.array([1.0, -0.5, 0.1]) n = n / np.linalg.norm(n) times = np.linspace(0, 1, 50)

In [ ]:

noise = 0.05
decay_factor = 0.5

data = create_data(Omega, r_0, n, times, noise, decay_factor)

qx.viz.plot_bloch_vectors(times, data)
qx.viz.display_bloch_sphere(data)

noise = 0.05 decay_factor = 0.5 data = create_data(Omega, r_0, n, times, noise, decay_factor) qx.viz.plot_bloch_vectors(times, data) qx.viz.display_bloch_sphere(data)

In [ ]:

def simulate_rotation(times, x_0, y_0, z_0, Omega, theta, phi, alpha):
    """Simulate a damped rotation for given parameters."""
    r_0 = np.array([x_0, y_0, z_0])
    n_x = np.sin(theta) * np.cos(phi)
    n_y = np.sin(theta) * np.sin(phi)
    n_z = np.cos(theta)
    n = (n_x, n_y, n_z)
    r_t = np.array([rotation_matrix(t, Omega, n) @ r_0 for t in times])
    decay_factor = np.exp(-alpha * times)
    return r_t * decay_factor[:, None]


def residual(params, times, data):
    """Return residuals for least-squares fitting."""
    return (simulate_rotation(times, *params) - data).flatten()


initial_guess = [0, 0, 1, 4 * np.pi, 0, 0, 0]
result = least_squares(residual, initial_guess, args=(times, data))
result.x

def simulate_rotation(times, x_0, y_0, z_0, Omega, theta, phi, alpha): """Simulate a damped rotation for given parameters.""" r_0 = np.array([x_0, y_0, z_0]) n_x = np.sin(theta) * np.cos(phi) n_y = np.sin(theta) * np.sin(phi) n_z = np.cos(theta) n = (n_x, n_y, n_z) r_t = np.array([rotation_matrix(t, Omega, n) @ r_0 for t in times]) decay_factor = np.exp(-alpha * times) return r_t * decay_factor[:, None] def residual(params, times, data): """Return residuals for least-squares fitting.""" return (simulate_rotation(times, *params) - data).flatten() initial_guess = [0, 0, 1, 4 * np.pi, 0, 0, 0] result = least_squares(residual, initial_guess, args=(times, data)) result.x

In [ ]:

fit = simulate_rotation(times, *result.x)

qx.viz.plot_bloch_vectors(times, data, title="State evolution : data")
qx.viz.plot_bloch_vectors(times, fit, title="State evolution : fit")

fit = simulate_rotation(times, *result.x) qx.viz.plot_bloch_vectors(times, data, title="State evolution : data") qx.viz.plot_bloch_vectors(times, fit, title="State evolution : fit")

In [ ]:

fig = go.Figure()

# data
fig.add_trace(
    go.Scatter3d(
        name="data",
        x=data[:, 0],
        y=data[:, 1],
        z=data[:, 2],
        mode="markers",
        marker=dict(size=3),
        hoverinfo="skip",
    )
)

# fit
fig.add_trace(
    go.Scatter3d(
        name="fit",
        x=fit[:, 0],
        y=fit[:, 1],
        z=fit[:, 2],
        mode="lines",
        line=dict(width=4),
        hoverinfo="skip",
    )
)

# sphere
theta = np.linspace(0, np.pi, 50)
phi = np.linspace(0, 2 * np.pi, 50)
theta, phi = np.meshgrid(theta, phi)
r = 1
x = r * np.sin(theta) * np.cos(phi)
y = r * np.sin(theta) * np.sin(phi)
z = r * np.cos(theta)
fig.add_trace(
    go.Surface(
        x=x,
        y=y,
        z=z,
        opacity=0.05,
        showscale=False,
        colorscale="gray",
        hoverinfo="skip",
    )
)

# layout
fig.update_layout(
    scene=dict(
        xaxis=dict(title="〈X〉", visible=True),
        yaxis=dict(title="〈Y〉", visible=True),
        zaxis=dict(title="〈Z〉", visible=True),
        aspectmode="cube",
    ),
    width=400,
    height=400,
    margin=dict(l=0, r=0, b=0, t=0),
    showlegend=False,
)

fig.show()

fig = go.Figure() # data fig.add_trace( go.Scatter3d( name="data", x=data[:, 0], y=data[:, 1], z=data[:, 2], mode="markers", marker=dict(size=3), hoverinfo="skip", ) ) # fit fig.add_trace( go.Scatter3d( name="fit", x=fit[:, 0], y=fit[:, 1], z=fit[:, 2], mode="lines", line=dict(width=4), hoverinfo="skip", ) ) # sphere theta = np.linspace(0, np.pi, 50) phi = np.linspace(0, 2 * np.pi, 50) theta, phi = np.meshgrid(theta, phi) r = 1 x = r * np.sin(theta) * np.cos(phi) y = r * np.sin(theta) * np.sin(phi) z = r * np.cos(theta) fig.add_trace( go.Surface( x=x, y=y, z=z, opacity=0.05, showscale=False, colorscale="gray", hoverinfo="skip", ) ) # layout fig.update_layout( scene=dict( xaxis=dict(title="〈X〉", visible=True), yaxis=dict(title="〈Y〉", visible=True), zaxis=dict(title="〈Z〉", visible=True), aspectmode="cube", ), width=400, height=400, margin=dict(l=0, r=0, b=0, t=0), showlegend=False, ) fig.show()

In [ ]:

qx.fit.fit_rotation(times, data)
qx.viz.display_bloch_sphere(data)

qx.fit.fit_rotation(times, data) qx.viz.display_bloch_sphere(data)