run_alignment_phase.py

import torch
from particle_dynamics import ParticleSimulation


# Set device to GPU if available, otherwise CPU
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
dtype = torch.float32

# Example configurations for different dynamics
example_configs = {
    "collapse_1D": {
        "V": torch.tensor([[1, 0, 0], [0, -1, 0], [0, 0, 2]], device=device, dtype=dtype),
        "QtK": torch.tensor([[-1, 0, 0], [0, -1, 0], [0, 0, 1]], device=device, dtype=dtype),
    },
    "collapse_2D": {
        "V": torch.tensor([[1, 0, 0], [0, 1, 0], [0, 0, 1]], device=device, dtype=dtype),
        "QtK": torch.tensor([[1, 0, 0], [0, 1, 0], [0, 0, -2]], device=device, dtype=dtype).T,
    },
    "collapse_rotation": {
        "description": "Particles collapse towards the Z-axis while rotating around it.",
        "V": torch.tensor([[1, 0, 0], [0, 1, 0], [0, 0, 1]], device=device, dtype=dtype),
        "QtK": torch.tensor([[0, 1, 0], [-1, 0, 0], [0, 0, -2]], device=device, dtype=dtype).T,
    },
    "random_matrices": {
        "description": "Complex dynamics generated from random interaction matrices.",
        "V": torch.randn((3, 3), device=device, dtype=dtype),
        "QtK": torch.randn((3, 3), device=device, dtype=dtype),
    }
}

def main():
    """
    Sets up and runs a particle dynamics simulation, then saves the result as a GIF.
    """
    print(f"Using device: {'cuda' if torch.cuda.is_available() else 'cpu'}")

    # --- Simulation Parameters ---
    # Feel free to change these values to explore different dynamics
    example_name = "collapse_1D"  # Choose from example_configs keys
    simulation_config = {
        "N": 10**4,          # Number of particles
        "d": 3,              # Number of dimensions (must be 3 for GIF)
        "T": 5.0,            # Total simulation time
        "dt": 1e-2,          # Simulation time step
        "beta": 30.0,         # Temperature
        "dtype": "single",   # Use 'double' for higher precision
        "original_transformer": True, # Use attention mechanism
        "QtK": example_configs[example_name]["QtK"],
        "V": example_configs[example_name]["V"]
    }

    # Initialize the simulation with the specified configuration
    sim = ParticleSimulation(**simulation_config)

    # --- Run Simulation ---
    print("\nStarting simulation...")
    sim.run()
    print("Simulation complete.")

    # --- Generate and Save Animation ---
    if sim.d == 3:
        print("\nGenerating 3D animation...")
        sim.save_simulation_as_gif(
            filename=f'results/simulation_{example_name}.gif',
            frames=150,  # Number of frames in the final GIF
            interval=40  # Milliseconds per frame
        )
    else:
        print(f"\nSimulation was in {sim.d}D. GIF generation is only available for 3D.")


if __name__ == "__main__":
    main()