Update app.py

app.py

@@ -5,9 +5,8 @@ from fastmcp import FastMCP
 from ase import Atoms
 from ase.build import molecule
 
-# Initialize MCP Server
-# We do NOT import trackio or mace globally to prevent startup crashes
-mcp = FastMCP("RealMACE_Agent", dependencies=["mace-torch", "trackio", "ase", "e3nn"])
+# Initialize MCP Server (dependencies removed - use requirements.txt instead)
+mcp = FastMCP("RealMACE_Agent")
 
 # Global State to share data between tools
 STATE = {
@@ -35,7 +34,7 @@ def create_dummy_batch(r_max=5.0):
     # Create dummy water
     mol = molecule("H2O")
     mol.info["energy"] = -14.0  # Dummy target energy (eV)
-    mol.arrays["forces"] = np.random.randn(3, 3) # Dummy target forces
+    mol.arrays["forces"] = np.random.randn(3, 3) * 0.1 # Dummy target forces
     
     config = Configuration(
         atomic_numbers=mol.get_atomic_numbers(),
@@ -49,9 +48,9 @@ def create_dummy_batch(r_max=5.0):
     # Convert to batch
     z_table = {1: 0, 8: 1} # Map H->0, O->1 for simple one-hot
     data_loader = torch_geometric.DataLoader(
-        dataset=[AtomicData.from_config(config, z_table=z_table, cutoff=r_max, model="MACE")],
+        dataset=[AtomicData.from_config(config, z_table=z_table, cutoff=r_max)],
         batch_size=1,
-        shuffle=True
+        shuffle=False
     )
     return next(iter(data_loader))
 
@@ -61,10 +60,14 @@ def create_dummy_batch(r_max=5.0):
 def init_real_mace_model(r_max: float = 5.0, max_ell: int = 2, hidden_dim: int = 16) -> str:
     """
     Initializes a REAL MACE model and stores it in memory.
+    
     Args:
-        r_max: Cutoff radius (Angstroms)
-        max_ell: Max spherical harmonic degree (2=vectors, 0=invariant)
-        hidden_dim: Size of the embedding vectors
+        r_max: Cutoff radius in Angstroms (default 5.0)
+        max_ell: Maximum spherical harmonic degree - 0=scalars only, 2=include vectors (default 2)
+        hidden_dim: Size of the hidden embedding vectors (default 16)
+    
+    Returns:
+        Status message with model configuration
     """
     ScaleShiftMACE, _, _, _, o3 = get_mace_setup()
     
@@ -83,7 +86,7 @@ def init_real_mace_model(r_max: float = 5.0, max_ell: int = 2, hidden_dim: int =
         num_interactions=2,
         num_elements=2, # H and O
         hidden_irreps=o3.Irreps(f"{hidden_dim}x0e"),
-        atomic_energies=np.array([-13.6, -10.0]), # Dummy average energies
+        atomic_energies=np.array([-13.6, -10.0]), # Dummy average energies for H and O
         avg_num_neighbors=2,
         atomic_numbers=[1, 8]
     )
@@ -93,21 +96,39 @@ def init_real_mace_model(r_max: float = 5.0, max_ell: int = 2, hidden_dim: int =
         model = ScaleShiftMACE(**model_config)
         STATE["model"] = model
         STATE["config"] = model_config
-        return f"✅ MACE Model Initialized! (L={max_ell}, R={r_max}). Ready to train."
+        return (
+            f"✅ **MACE Model Initialized Successfully!**\n\n"
+            f"Configuration:\n"
+            f"- Cutoff Radius (r_max): {r_max} Å\n"
+            f"- Max Spherical Harmonic Degree (L_max): {max_ell}\n"
+            f"- Hidden Dimension: {hidden_dim}\n"
+            f"- Interaction Blocks: 2\n"
+            f"- Elements: H, O\n\n"
+            f"Model is ready for training. Use 'train_with_trackio' next."
+        )
     except Exception as e:
-        return f"Error initializing MACE: {str(e)}"
+        return f"❌ Error initializing MACE: {str(e)}"
 
 @mcp.tool()
-def train_with_trackio(experiment_name: str, epochs: int = 10) -> str:
+def train_with_trackio(experiment_name: str, epochs: int = 10, learning_rate: float = 0.01) -> str:
     """
-    Trains the stored MACE model and logs to Trackio.
-    MUST run init_real_mace_model first.
+    Trains the stored MACE model and logs metrics to Trackio.
+    
+    Args:
+        experiment_name: Name for this training run in Trackio
+        epochs: Number of training epochs (default 10)
+        learning_rate: Optimizer learning rate (default 0.01)
+    
+    Returns:
+        Training summary with final loss metrics
+    
+    Note: Must run 'init_real_mace_model' first to create a model.
     """
     # 1. Lazy Import Trackio to prevent startup crash
     try:
         import trackio
     except ImportError:
-        return "Trackio not installed. Run: pip install trackio"
+        return "❌ Trackio not installed. Run: pip install trackio"
 
     # 2. Check if model exists
     if STATE["model"] is None:
@@ -117,14 +138,13 @@ def train_with_trackio(experiment_name: str, epochs: int = 10) -> str:
     batch = STATE["batch"]
     
     # 3. Setup Optimizer
-    optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
+    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
     
     # 4. Setup Trackio
     try:
-        # Trackio might fail if OAuth isn't set up in Space, catch it gracefully
         logger = trackio.Logger(project="Real_MACE_Training", name=experiment_name)
     except Exception as e:
-        return f"❌ Trackio Connection Failed: {e}. (Check 'hf_oauth: true' in README?)"
+        return f"❌ Trackio Connection Failed: {e}\n(Hint: Add 'hf_oauth: true' to README.md if running on HF Space)"
 
     # 5. Training Loop
     model.train()
@@ -138,36 +158,112 @@ def train_with_trackio(experiment_name: str, epochs: int = 10) -> str:
         # MACE Forward Pass
         out = model(batch.to_dict())
         
-        # Loss Calc
+        # Loss Calculation (Energy MSE + Force MSE)
         loss_e = torch.mean((out["energy"] - batch.energy)**2)
         loss_f = torch.mean((out["forces"] - batch.forces)**2)
-        total_loss = loss_e + 10.0 * loss_f
+        total_loss = loss_e + 10.0 * loss_f  # Weight forces 10x more
         
         total_loss.backward()
         optimizer.step()
         
-        # Log metrics
+        # Calculate MAE metrics for interpretability
+        force_mae = torch.mean(torch.abs(out["forces"] - batch.forces)).item()
+        energy_mae = torch.abs(out["energy"] - batch.energy).mean().item()
+        
+        # Log metrics to Trackio
         metrics = {
             "epoch": epoch,
             "total_loss": total_loss.item(),
-            "force_mae": torch.mean(torch.abs(out["forces"] - batch.forces)).item(),
-            "wall_time": time.time() - start_time
+            "energy_mae_eV": energy_mae,
+            "force_mae_eV_A": force_mae,
+            "wall_time_sec": time.time() - start_time
         }
         
-        # Push to Trackio
         logger.log(metrics)
         
-        if epoch % 5 == 0:
-            log_summary.append(f"Epoch {epoch}: Loss={total_loss.item():.4f}")
-            time.sleep(0.05) # Yield slightly
+        if epoch % 5 == 0 or epoch == epochs - 1:
+            log_summary.append(
+                f"Epoch {epoch:3d}: Loss={total_loss.item():.5f} | "
+                f"Force MAE={force_mae:.5f} eV/Å"
+            )
+            time.sleep(0.05) # Small delay for visualization
 
     return (
-        f"🚀 **Training Complete!**\n"
-        f"Experiment: {experiment_name}\n"
-        f"Final Loss: {total_loss.item():.5f}\n"
-        f"Check the Trackio tab for the live graphs!"
+        f"🚀 **Training Complete!**\n\n"
+        f"**Experiment:** {experiment_name}\n"
+        f"**Epochs:** {epochs}\n"
+        f"**Learning Rate:** {learning_rate}\n\n"
+        f"**Final Metrics:**\n"
+        f"- Total Loss: {total_loss.item():.6f}\n"
+        f"- Energy MAE: {energy_mae:.6f} eV\n"
+        f"- Force MAE: {force_mae:.6f} eV/Å\n\n"
+        f"📊 Check the **Trackio Dashboard** for live loss curves and training dynamics!\n\n"
+        f"**Recent Training Log:**\n" + "\n".join(log_summary)
+    )
+
+@mcp.tool()
+def check_equivariance(rotation_degrees: float = 45.0) -> str:
+    """
+    Educational tool: Tests if the model is truly E(3)-equivariant.
+    Rotates the input molecule and checks if predicted forces rotate exactly with it.
+    
+    Args:
+        rotation_degrees: Angle to rotate the molecule around Z-axis (default 45.0)
+    
+    Returns:
+        Explanation of equivariance test results
+    """
+    if STATE["model"] is None:
+        return "⚠️ No model found! Run 'init_real_mace_model' first."
+    
+    model = STATE["model"]
+    batch = STATE["batch"]
+    
+    # Get original prediction
+    model.eval()
+    with torch.no_grad():
+        out_orig = model(batch.to_dict())
+        forces_orig = out_orig["forces"].clone()
+    
+    # Apply rotation to positions
+    angle = np.radians(rotation_degrees)
+    rot_matrix = torch.tensor([
+        [np.cos(angle), -np.sin(angle), 0],
+        [np.sin(angle), np.cos(angle), 0],
+        [0, 0, 1]
+    ], dtype=torch.float32)
+    
+    # Create rotated batch
+    batch_rot = batch.clone()
+    batch_rot.positions = torch.matmul(batch.positions, rot_matrix.T)
+    
+    # Get prediction on rotated input
+    with torch.no_grad():
+        out_rot = model(batch_rot.to_dict())
+        forces_rot = out_rot["forces"]
+    
+    # Manually rotate the original forces
+    forces_orig_rotated = torch.matmul(forces_orig, rot_matrix.T)
+    
+    # Calculate equivariance error
+    equivariance_error = torch.mean(torch.abs(forces_rot - forces_orig_rotated)).item()
+    
+    return (
+        f"🧪 **E(3)-Equivariance Test Results**\n\n"
+        f"**Test Setup:**\n"
+        f"- Molecule: Water (H₂O)\n"
+        f"- Rotation: {rotation_degrees}° around Z-axis\n\n"
+        f"**Results:**\n"
+        f"- Equivariance Error: {equivariance_error:.2e} eV/Å\n"
+        f"- Expected for perfect equivariance: ~1e-6 or lower\n\n"
+        f"**Interpretation:**\n"
+        f"{'✅ PASS: Model is equivariant!' if equivariance_error < 1e-4 else '⚠️ WARNING: High error detected'}\n\n"
+        f"This confirms that when you rotate the molecule, the predicted force vectors "
+        f"rotate **exactly** with it. Standard MLPs cannot achieve this without extensive "
+        f"data augmentation!"
     )
 
 if __name__ == "__main__":
     print("Starting MACE-MCP Server...")
-    mcp.run()
+    # Use SSE transport for Hugging Face Spaces deployment
+    mcp.run(transport="sse")