Refactor procedural code into classes: Add MSMDSystemGenerator, Proba…

…bilityMap, MDSimulationManager, ResidueEnvironmentAnalyzer, Configuration, and ProfileAnalyzer with backward compatibility

Co-Authored-By: K. Yanagisawa <[email protected]>

script/generate_msmd_system.py

@@ -150,13 +150,87 @@ def create_system(
     return tleap_obj.parm7, tleap_obj.rst7
 
 
+class MSMDSystemGenerator:
+    """Generator class for Mixed-Solvent Molecular Dynamics systems."""
+
+    def __init__(self, setting: Dict[str, InputSettings]) -> None:
+        """Initialize the MSMD system generator.
+        
+        Args:
+            setting: Dictionary containing system configuration
+        """
+        self.setting = setting
+        self._protein_settings = cast(ProteinSettings, setting["input"]["protein"])
+        self._probe_settings = cast(ProbeSettings, setting["input"]["probe"])
+
+    def _prepare_protein_pdb(self, pdbfile: Path) -> Path:
+        """Prepare protein PDB file by removing OXT and ANISOU records.
+        
+        Args:
+            pdbfile: Path to input PDB file
+            
+        Returns:
+            Path to cleaned PDB file
+        """
+        return protein_pdb_preparation(pdbfile)
+
+    def _calculate_box_size(self, pdbfile: Path) -> float:
+        """Calculate simulation box size from protein PDB.
+        
+        Args:
+            pdbfile: Path to protein PDB file
+            
+        Returns:
+            Box size in Angstroms
+        """
+        return __calculate_boxsize(pdbfile)
+
+    def _create_frcmod(self, mol2file: Path, atomtype: Literal["gaff", "gaff2"], debug: bool = False) -> Path:
+        """Create force field modification file.
+        
+        Args:
+            mol2file: Path to mol2 file
+            atomtype: Atom type (GAFF / GAFF2)
+            debug: Enable debug output
+            
+        Returns:
+            Path to generated frcmod file
+        """
+        return _create_frcmod(mol2file, atomtype, debug)
+
+    def create_system(self, debug: bool = False, seed: int = -1) -> Tuple[Path, Path]:
+        """Create MSMD system from current settings.
+        
+        Args:
+            debug: Enable debug output
+            seed: Random seed for system generation
+            
+        Returns:
+            Tuple containing paths to (parm7, rst7) files
+        """
+        cfrcmod = self._create_frcmod(
+            Path(self._probe_settings["mol2"]),
+            self._probe_settings["atomtype"],
+            debug=debug
+        )
+        return create_system(
+            self._protein_settings,
+            self._probe_settings,
+            cfrcmod,
+            debug=debug,
+            seed=seed
+        )
+
 def generate_msmd_system(
     setting: Dict[str, InputSettings],
     debug: bool = False,
     seed: int = -1
 ) -> Tuple[Path, Path]:
     """Generate MSMD system from settings.
     
+    This is a backward-compatible wrapper around MSMDSystemGenerator.
+    For new code, prefer using MSMDSystemGenerator directly.
+    
     Args:
         setting: Dictionary containing input settings
         debug: Enable debug output
@@ -165,19 +239,5 @@ def generate_msmd_system(
     Returns:
         Tuple containing paths to parm7 and rst7 files
     """
-    """
-    generate msmd system
-    -----
-    input
-        setting: setting json (dict)
-        debug: debug mode
-        seed: random seed
-    output
-        parm7: path to parm7 file
-        rst7: path to rst7 file
-    """
-    cfrcmod = _create_frcmod(
-        Path(setting["input"]["probe"]["mol2"]), setting["input"]["probe"]["atomtype"], debug=debug
-    )
-    parm7, rst7 = create_system(setting["input"]["protein"], setting["input"]["probe"], cfrcmod, debug=debug, seed=seed)
-    return parm7, rst7
+    generator = MSMDSystemGenerator(setting)
+    return generator.create_system(debug=debug, seed=seed)

script/genpmap.py

@@ -74,23 +74,99 @@ def convert_to_proba(
     return g
 
 
+class ProbabilityMap:
+    """Class for generating and manipulating probability maps."""
+
+    def __init__(self, grid: gridData.Grid) -> None:
+        """Initialize probability map with grid data.
+        
+        Args:
+            grid: Grid data object containing probability values
+        """
+        self.grid = grid
+        self._mask: Optional[npt.NDArray[np.float64]] = None
+
+    @classmethod
+    def from_file(cls, path: Union[str, Path]) -> 'ProbabilityMap':
+        """Create ProbabilityMap from a grid file.
+        
+        Args:
+            path: Path to grid file
+            
+        Returns:
+            New ProbabilityMap instance
+        """
+        return cls(gridData.Grid(str(path)))
+
+    def generate_mask(self, ref_struct: Path, distance: Optional[float] = None) -> None:
+        """Generate mask from reference structure.
+        
+        Args:
+            ref_struct: Path to reference structure
+            distance: Optional distance threshold for mask
+        """
+        mask_grid = mask_generator(ref_struct, self.grid, distance)
+        self._mask = mask_grid.grid
+
+    def convert_to_probability(self, frames: int = 1, normalize: str = "snapshot") -> None:
+        """Convert grid values to probabilities.
+        
+        Args:
+            frames: Number of frames used
+            normalize: Normalization method ("snapshot", "GFE", or "total")
+        """
+        converted = convert_to_proba(self.grid, self._mask, normalize=normalize, frames=frames)
+        self.grid = converted
+
+    def save(self, path: Union[str, Path], type: str = "double") -> None:
+        """Save probability map to file.
+        
+        Args:
+            path: Output file path
+            type: Data type for output
+        """
+        self.grid.export(str(path), type=type)
+
+    def to_gfe(self, mean_proba: float, temperature: float = 300) -> None:
+        """Convert probability map to Gibbs free energy.
+        
+        Args:
+            mean_proba: Mean probability for normalization
+            temperature: Temperature in Kelvin
+        """
+        self.grid.grid = np.where(self.grid.grid <= 0, 1e-10, self.grid.grid)
+        self.grid.grid = -(constants.R / constants.calorie / constants.kilo) * temperature * np.log(self.grid.grid / mean_proba)
+        self.grid.grid = np.where(self.grid.grid > 3, 3, self.grid.grid)
+
+    def to_inverse_gfe(self) -> None:
+        """Convert grid to inverse Gibbs free energy."""
+        self.grid.grid = -self.grid.grid
+
 def convert_to_gfe(
     grid_path: str,
     mean_proba: float,
     temperature: float = 300
 ) -> str:
-    pmap = gridData.Grid(grid_path)
-    pmap.grid = np.where(pmap.grid <= 0, 1e-10, pmap.grid)  # avoid log(0)
-    pmap.grid = -(constants.R / constants.calorie / constants.kilo) * temperature * np.log(pmap.grid / mean_proba)
-    pmap.grid = np.where(pmap.grid > 3, 3, pmap.grid)  # Definition of GFE in the paper Raman et al., JCIM, 2013
-
+    """Convert probability map to GFE (wrapper for backward compatibility).
+    
+    Args:
+        grid_path: Path to probability map file
+        mean_proba: Mean probability for normalization
+        temperature: Temperature in Kelvin
+        
+    Returns:
+        Path to generated GFE file
+    """
+    pmap = ProbabilityMap.from_file(grid_path)
+    pmap.to_gfe(mean_proba, temperature)
+
     gfe_path = os.path.dirname(grid_path) + "/" + "GFE" + "_" + os.path.basename(grid_path)
-    pmap.export(gfe_path, type="double")
-
-    pmap.grid = -pmap.grid
+    pmap.save(gfe_path)
+    
+    pmap.to_inverse_gfe()
     invgfe_path = os.path.dirname(grid_path) + "/" + "InvGFE" + "_" + os.path.basename(grid_path)
-    pmap.export(invgfe_path, type="double")
-
+    pmap.save(invgfe_path)
+    
     return gfe_path
 
 
@@ -101,12 +177,24 @@ def convert_to_pmap(
     normalize: str = "snapshot",
     frames: int = 1
 ) -> str:
-    grid = gridData.Grid(grid_path)
-    mask = mask_generator(ref_struct, grid, valid_distance)
-    pmap = convert_to_proba(grid, mask.grid, frames=frames, normalize=normalize)
-
+    """Convert grid to probability map (wrapper for backward compatibility).
+    
+    Args:
+        grid_path: Path to input grid file
+        ref_struct: Path to reference structure
+        valid_distance: Distance threshold for mask
+        normalize: Normalization method
+        frames: Number of frames used
+        
+    Returns:
+        Path to generated probability map file
+    """
+    pmap = ProbabilityMap.from_file(grid_path)
+    pmap.generate_mask(ref_struct, valid_distance)
+    pmap.convert_to_probability(frames=frames, normalize=normalize)
+
     pmap_path = os.path.dirname(grid_path) + "/" + "PMAP" + "_" + os.path.basename(grid_path)
-    pmap.export(pmap_path, type="double")
+    pmap.save(pmap_path)
     return pmap_path