import matplotlib.pyplot as plt
import numpy as np
from clocks import (
ClockArray, InferenceConfig, MassConfig, NoiseConfig,
PriorConfig, SimulationConfig, infer, simulate,
)
clock_positions = np.array([
[-4.0, 0.0], [-2.0, 3.0], [1.0, 4.0], [4.0, 2.0],
[5.0, -1.0], [2.0, -4.0], [-1.0, -3.0], [-3.0, -1.5],
])
clock_array = ClockArray(positions=clock_positions, track_offset=3.0)
truth = MassConfig(positions=np.array([[1.5, -1.0]]), masses=np.array([0.5]))
sim = simulate(SimulationConfig(
clock_array=clock_array, ground_truth=truth,
noise=NoiseConfig(observation_std=0.005), n_observations=25, seed=42,
))
result = infer(sim.observations, InferenceConfig(
clock_array=clock_array, noise=NoiseConfig(observation_std=0.005),
prior=PriorConfig(position_range=(-8.0, 8.0), mass_range=(0.1, 2.0)),
n_particles=500, n_masses=1, seed=42,
))
particles = result.particle_state.particles
weights = result.particle_state.weights
fig, ax = plt.subplots()
ax.scatter(particles[:, 0], particles[:, 1], c=weights, cmap="viridis", s=8, alpha=0.7)
ax.scatter(1.5, -1.0, marker="x", s=120, color="red", linewidths=2, label="True")
ax.scatter(
clock_positions[:, 0], clock_positions[:, 1],
marker="s", s=60, color="steelblue", label="Clocks",
)
ax.set_xlabel("x")
ax.set_ylabel("y")
ax.set_aspect("equal")
ax.legend(fontsize=8)
plt.close(fig)
fig