A simple example to understand invariant and equivariant in materials

invariant_equavariant.py

#!/usr/bin/env python

import numpy as np

def rotate_vector(v, R):
    """Rotates a vector v using a rotation matrix R."""
    return np.dot(R, v)

def energy_function(positions):
    """A toy energy function that depends on interatomic distances."""
    # Compute pairwise distances
    energy = 0
    n_atoms = len(positions)
    for i in range(n_atoms):
        for j in range(i + 1, n_atoms):
            r = np.linalg.norm(positions[i] - positions[j])
            energy += 1 / r  # Example: Lennard-Jones-like interaction (simplified)
    return energy

def compute_gradients(positions):
    """Compute gradients (forces) as the negative derivative of energy w.r.t positions."""
    n_atoms = len(positions)
    gradients = np.zeros_like(positions)
    for i in range(n_atoms):
        for j in range(n_atoms):
            if i != j:
                r_vec = positions[i] - positions[j]
                r = np.linalg.norm(r_vec)
                gradients[i] += -r_vec / r**3  # Force due to 1/r potential
    return gradients
# Original positions
positions = np.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]])

# Rotation matrix (90 degrees about z-axis in radian)
angle = np.pi / 2
print(angle)
R = np.array([
    [np.cos(angle), -np.sin(angle), 0],
    [np.sin(angle), np.cos(angle), 0],
    [0, 0, 1]
])

# Rotate positions
rotated_positions = np.array([rotate_vector(pos, R) for pos in positions])

# Compute energy and gradients for original and rotated positions
original_energy = energy_function(positions)
rotated_energy = energy_function(rotated_positions)

original_gradients = compute_gradients(positions)
rotated_gradients = compute_gradients(rotated_positions)

# Rotate the original gradients
rotated_original_gradients = np.array([rotate_vector(grad, R) for grad in original_gradients])

# Output results
print("Original Energy:", original_energy)
print("Rotated Energy:", rotated_energy)
print("\nOriginal Gradients:\n", original_gradients)
print("\nRotated Gradients:\n", rotated_gradients)
print("\nRotated Original Gradients (for comparison):\n", rotated_original_gradients)

# Verify invariance and equivariance
print("\nEnergy Invariance:", np.isclose(original_energy, rotated_energy))
print("Gradient Equivariance:", np.allclose(rotated_gradients, rotated_original_gradients))