Improve MACE calculator to use mh1

self._calculator = MACECalculator(
         model_paths=model_paths, 
          device=device_str,
           head=mace_head
  )
Available heads are: ['matpes_r2scan', 'mp_pbe_refit_add', 'spice_wB97M', 'oc20_usemppbe', 'omol', 'omat_pbe']

Author: adamlaho

amlpa.py

@@ -327,7 +327,7 @@ def _parse_xyz_header(self, filename: str) -> Tuple[float, List[float]]:
         energy = float(energy_match.group(1))
         cell = [float(cell_match.group(i)) for i in range(1, 7)]
         return energy, cell
-    
+
     def _setup_calculator(self):
         """Setup MACE calculator once and store it"""
         if hasattr(self, '_calculator'):
@@ -345,16 +345,28 @@ def _setup_calculator(self):
                 device_str = ','.join(gpu_devices)
 
             model_paths = self.config.get('model_paths', ['/path/to/your/model'])
-            self._calculator = MACECalculator(model_paths=model_paths, device=device_str)
+            
+            # Get head from config, with default
+            mace_head = self.config.get('mace_head', 'omat_pbe')
+            
+            self.logger.log(f"Setting up MACE calculator with head: {mace_head}")
+            
+            self._calculator = MACECalculator(
+                model_paths=model_paths, 
+                device=device_str,
+                head=mace_head
+            )
             self.atoms.calc = self._calculator
 
-            self.logger.log(f"MACE calculator setup: {device_str}")
+            self.logger.log(f"MACE calculator setup: device={device_str}, head={mace_head}")
 
         except ImportError:
             self.logger.log("MACE not available, using EMT calculator", 'WARNING')
             from ase.calculators.emt import EMT
             self._calculator = EMT()
             self.atoms.calc = self._calculator
+
+
 
     def run_single_point(self) -> bool:
         """Run single point energy calculation"""
@@ -535,7 +547,9 @@ def cell_callback():
             import traceback
             self.logger.log(traceback.format_exc(), 'ERROR')
             return False
-    
+
+
+
     def analyze_rmsd(self, temperature: float = None) -> Dict[str, Any]:
         """Calculate RMSD between initial and current structure"""
         if not self.config.get('run_rmsd', False):
@@ -550,48 +564,91 @@ def analyze_rmsd(self, temperature: float = None) -> Dict[str, Any]:
         current_atoms = self.atoms
         initial_pos = self.initial_atoms.get_positions()
         current_pos = current_atoms.get_positions()
+        initial_cell = self.initial_atoms.get_cell()
+        current_cell = current_atoms.get_cell()
 
         if len(initial_pos) != len(current_pos):
             self.logger.log(f"Structure size mismatch: {len(initial_pos)} vs {len(current_pos)}", 'ERROR')
             return {}
 
-        # Handle PBC
+        # Handle PBC for basic RMSD calculations
         if self.atoms.get_pbc().any():
             displacement = current_pos - initial_pos
-            displacement = find_mic(displacement, self.atoms.get_cell())[0]
-            current_pos = initial_pos + displacement
+            displacement = find_mic(displacement, current_cell)[0]
+            current_pos_corrected = initial_pos + displacement
+        else:
+            current_pos_corrected = current_pos
+        
+        # Calculate standard RMSD metrics (unit cell, Cartesian)
+        rmsd_simple = self._calculate_rmsd_simple(initial_pos, current_pos_corrected)
+        rmsd_centered = self._calculate_rmsd_centered(initial_pos, current_pos_corrected)
+        rmsd_aligned = self._calculate_rmsd_aligned(initial_pos, current_pos_corrected)
+        
+        # Calculate fractional coordinate RMSD (accounts for cell changes)
+        rmsd_fractional = self._calculate_rmsd_fractional(
+            initial_pos, current_pos, initial_cell, current_cell
+        )
+        
+        # Calculate supercell RMSD (like RMSD15, captures lattice parameter errors)
+        supercell_dims = self.config.get('rmsd_supercell_dims', [3, 3, 3])
+        rmsd_supercell = self._calculate_rmsd_supercell(
+            self.initial_atoms, current_atoms, supercell_dims
+        )
 
-        # Calculate different RMSD metrics
-        rmsd_simple = self._calculate_rmsd_simple(initial_pos, current_pos)
-        rmsd_centered = self._calculate_rmsd_centered(initial_pos, current_pos)
-        rmsd_aligned = self._calculate_rmsd_aligned(initial_pos, current_pos)
+        # Calculate cell parameter deviations
+        cell_deviations = self._calculate_cell_deviation(initial_cell, current_cell)
 
-        # Per-atom analysis
+        # Per-atom displacement analysis
         displacement = current_pos - initial_pos
         if self.atoms.get_pbc().any():
-            displacement = find_mic(displacement, self.atoms.get_cell())[0]
+            displacement = find_mic(displacement, current_cell)[0]
 
         per_atom_displacements = np.sqrt(np.sum(displacement**2, axis=1))
 
         analysis = {
+            # Unit cell Cartesian RMSD (existing metrics)
             'rmsd_simple': rmsd_simple,
             'rmsd_centered': rmsd_centered, 
             'rmsd_aligned': rmsd_aligned,
+            
+            # New metrics that account for cell changes
+            'rmsd_fractional': rmsd_fractional,
+            'rmsd_supercell': rmsd_supercell,
+            'supercell_dims': supercell_dims,
+            
+            # Cell parameter deviations
+            'cell_deviations': cell_deviations,
+            
+            # Per-atom statistics
             'n_atoms': len(initial_pos),
             'max_displacement': np.max(per_atom_displacements),
             'mean_displacement': np.mean(per_atom_displacements),
             'std_displacement': np.std(per_atom_displacements),
             'per_atom_displacements': per_atom_displacements
         }
 
+        # Logging
         self.logger.log(f"RMSD Analysis:")
-        self.logger.log(f"  Simple RMSD: {rmsd_simple:.4f} Å")
-        self.logger.log(f"  Centered RMSD: {rmsd_centered:.4f} Å")
-        self.logger.log(f"  Aligned RMSD: {rmsd_aligned:.4f} Å")
-        self.logger.log(f"  Max displacement: {analysis['max_displacement']:.4f} Å")
+        self.logger.log(f"  Unit Cell (Cartesian):")
+        self.logger.log(f"    Simple RMSD: {rmsd_simple:.4f} Å")
+        self.logger.log(f"    Centered RMSD: {rmsd_centered:.4f} Å")
+        self.logger.log(f"    Aligned RMSD (Kabsch): {rmsd_aligned:.4f} Å")
+        self.logger.log(f"  Fractional Coordinates:")
+        self.logger.log(f"    RMSD (fractional): {rmsd_fractional:.6f}")
+        self.logger.log(f"  Supercell {supercell_dims}:")
+        self.logger.log(f"    RMSD (supercell): {rmsd_supercell:.4f} Å")
+        self.logger.log(f"  Cell Parameters:")
+        self.logger.log(f"    Δa: {cell_deviations['delta_a']:.4f} Å")
+        self.logger.log(f"    Δb: {cell_deviations['delta_b']:.4f} Å")
+        self.logger.log(f"    Δc: {cell_deviations['delta_c']:.4f} Å")
+        self.logger.log(f"    ΔV: {cell_deviations['volume_change_percent']:.2f}%")
+        self.logger.log(f"  Per-atom Displacement:")
+        self.logger.log(f"    Max: {analysis['max_displacement']:.4f} Å")
+        self.logger.log(f"    Mean: {analysis['mean_displacement']:.4f} Å")
 
         return analysis
-    
+
+
     def _calculate_rmsd_simple(self, pos1: np.ndarray, pos2: np.ndarray) -> float:
         """Simple RMSD calculation"""
         return np.sqrt(np.mean(np.sum((pos1 - pos2)**2, axis=1)))
@@ -601,7 +658,7 @@ def _calculate_rmsd_centered(self, pos1: np.ndarray, pos2: np.ndarray) -> float:
         pos1_centered = pos1 - np.mean(pos1, axis=0)
         pos2_centered = pos2 - np.mean(pos2, axis=0)
         return self._calculate_rmsd_simple(pos1_centered, pos2_centered)
-    
+
     def _calculate_rmsd_aligned(self, pos1: np.ndarray, pos2: np.ndarray) -> float:
         """RMSD after optimal alignment (Kabsch algorithm)"""
         try:
@@ -623,7 +680,54 @@ def _calculate_rmsd_aligned(self, pos1: np.ndarray, pos2: np.ndarray) -> float:
         except Exception as e:
             self.logger.log(f"Alignment failed, using centered RMSD: {e}", 'WARNING')
             return self._calculate_rmsd_centered(pos1, pos2)
-    
+
+    def _calculate_rmsd_fractional(self, pos1: np.ndarray, pos2: np.ndarray, 
+                                    cell1: Cell, cell2: Cell) -> float:
+        """RMSD in fractional coordinates (accounts for cell changes)"""
+        # Convert Cartesian positions to fractional coordinates
+        frac1 = cell1.scaled_positions(pos1)
+        frac2 = cell2.scaled_positions(pos2)
+        
+        # Handle PBC wrapping in fractional space
+        frac_diff = frac2 - frac1
+        frac_diff = frac_diff - np.round(frac_diff)
+        
+        return np.sqrt(np.mean(np.sum(frac_diff**2, axis=1)))
+
+    def _calculate_rmsd_supercell(self, atoms1: Atoms, atoms2: Atoms, 
+                                replicate: List[int] = [3, 3, 3]) -> float:
+        """RMSD on replicated supercell (like RMSD15, captures lattice errors)"""
+        from ase.build import make_supercell
+        
+        # Replicate both structures
+        transform = np.diag(replicate)
+        super1 = make_supercell(atoms1.copy(), transform)
+        super2 = make_supercell(atoms2.copy(), transform)
+        
+        pos1 = super1.get_positions()
+        pos2 = super2.get_positions()
+        
+        # Use Kabsch alignment on the supercell
+        return self._calculate_rmsd_aligned(pos1, pos2)
+
+    def _calculate_cell_deviation(self, cell1: Cell, cell2: Cell) -> Dict[str, float]:
+        """Calculate cell parameter deviations"""
+        params1 = cell1.cellpar()
+        params2 = cell2.cellpar()
+        
+        return {
+            'delta_a': abs(params1[0] - params2[0]),
+            'delta_b': abs(params1[1] - params2[1]),
+            'delta_c': abs(params1[2] - params2[2]),
+            'delta_alpha': abs(params1[3] - params2[3]),
+            'delta_beta': abs(params1[4] - params2[4]),
+            'delta_gamma': abs(params1[5] - params2[5]),
+            'delta_volume': abs(cell1.volume - cell2.volume),
+            'volume_change_percent': 100 * abs(cell1.volume - cell2.volume) / cell1.volume if cell1.volume > 0 else 0
+        }
+
+
+
     def calculate_energy_drift_statistics(self, times: np.ndarray, energies: np.ndarray, 
                                         temperatures: np.ndarray, start_time_ps: float = 0.0, 
                                         timestep_fs: float = 0.5) -> Dict[str, Any]:
@@ -1449,6 +1553,7 @@ def analyze_rdf_from_trajectory(self, temperature: float) -> Dict[str, Dict]:
 
         return averaged_rdfs
 
+
     def plot_rdf_comparison(self) -> None:
         """Plot RDF comparisons across temperatures"""
         self.logger.log("Plotting RDF comparisons...")
@@ -1629,29 +1734,32 @@ def run_full_analysis(self, xyz_file: str) -> None:
             self.plot_coordination_histogram(label="initial")
 
         # MD simulations
-        temperatures = self.config.get('temperatures', [300])
-        md_steps = self.config.get('md_steps', 10000)
-        
-        for temp in temperatures:
-            self.logger.log_section(f"MD ANALYSIS AT {temp}K")
+        if self.config.get('md_simu', True):  # Default to True for backward compatibility
+            temperatures = self.config.get('temperatures', [300])
+            md_steps = self.config.get('md_steps', 10000)
 
-            traj_data = self.run_md_simulation(temp, md_steps)
-            
-            if traj_data:
-                # Energy drift (NVE only)
-                drift_analysis = self.analyze_energy_drift(temp)
-                if drift_analysis:
-                    self.results[f'energy_drift_{temp}K'] = drift_analysis
-                    self.plot_energy_drift(temp)
+            for temp in temperatures:
+                self.logger.log_section(f"MD ANALYSIS AT {temp}K")
 
-                # RDF analysis
-                rdf_results = self.analyze_rdf_from_trajectory(temp)
-                if rdf_results:
-                    self.results[f'rdf_{temp}K'] = rdf_results
-        
-        # Generate comparison plots
-        self.plot_rdf_comparison()
+                traj_data = self.run_md_simulation(temp, md_steps)
+                
+                if traj_data:
+                    # Energy drift (NVE only)
+                    drift_analysis = self.analyze_energy_drift(temp)
+                    if drift_analysis:
+                        self.results[f'energy_drift_{temp}K'] = drift_analysis
+                        self.plot_energy_drift(temp)
+                    
+                    # RDF analysis
+                    rdf_results = self.analyze_rdf_from_trajectory(temp)
+                    if rdf_results:
+                        self.results[f'rdf_{temp}K'] = rdf_results
 
+            # Generate comparison plots
+            self.plot_rdf_comparison()
+        else:
+            self.logger.log_section("MD SIMULATIONS SKIPPED")
+            self.logger.log("MD simulations disabled in configuration (md_simu: false)")
         # Print cache statistics
         self.logger.log_section("CACHE STATISTICS")
         self.cache.print_cache_summary(self.logger)