Source code for arbdmodel.mesh_process_volume

import numpy as np
import gmsh
from scipy.spatial import KDTree
from pathlib import Path
import os
from .grid import writeDx
from .engine import HydroProRunner
from .logger import logger

"""
Improved mesh processor with corrected inertia calculations based purely on mesh geometry
"""

[docs] class MeshProcessor: """Process gmsh files to calculate inertia, hydrodynamics and generate potential fields""" # Conversion factors MICRON_TO_ANGSTROM = 10000 def __init__(self, mesh_file, density=19.3, simconf=None, unit_scale=MICRON_TO_ANGSTROM, work_dir=None, expected_mass=None): """ Initialize processor with mesh file Args: mesh_file: Path to .msh file density: Material density in g/cm^3 simconf: SimConf object containing configuration parameters unit_scale: Conversion factor from input units to angstroms work_dir: Working directory expected_mass: Optional expected mass in amu to calibrate calculations """ self.mesh_file = Path(mesh_file) self.density = density self.unit_scale = unit_scale self.expected_mass = expected_mass self.work_dir = Path(work_dir) if work_dir else Path.cwd() os.makedirs(self.work_dir, exist_ok=True) if simconf is None: from . import DefaultSimConf simconf=DefaultSimConf() self.simconf=simconf # Extract parameters from simconf self.temperature = simconf.temperature self.viscosity = simconf.viscosity self.solvent_density = simconf.solvent_density self.binary_path = simconf.get_binary('hydropro') self.name = str(self.mesh_file.stem) self.attached_patricles=[] # Initialize gmsh and read mesh gmsh.initialize() try: gmsh.open(str(self.mesh_file)) logger.info(f"Successfully opened mesh file: {self.mesh_file}") # First, load the full volumetric mesh for physical calculations self.volume_nodes = self._get_nodes() self.volume_elements = self._get_elements() logger.info(f"Loaded {len(self.volume_nodes)} nodes and {len(self.volume_elements)} elements from volume mesh") # Then extract surface mesh for visualization and HydroPro self.nodes, self.elements = self._extract_surface_mesh() logger.info(f"Extracted {len(self.nodes)} nodes and {len(self.elements)} elements from surface mesh") # Calculate volume using volumetric mesh self.volume = self._calculate_volume() # Calculate mass with density correction if needed if self.expected_mass: # Use expected mass if provided self.mass = self.expected_mass # Update density to match expected mass self.density_correction = self.expected_mass / (self.volume * 0.6022) logger.info(f"Using expected mass: {self.mass:.2f} amu (density correction: {self.density_correction:.6f})") else: # Convert density from g/cm^3 to amu/Å^3 self.mass = self.volume * self.density * 0.6022 self.density_correction = 1.0 logger.info(f"Calculated volume: {self.volume:.2f} ų") logger.info(f"Calculated mass: {self.mass:.2f} amu") # Align both meshes to center of mass and principal axes # Important: Center of mass calculation needs to be fixed to use the volumetric mesh self._align_mesh() # Calculate inertia tensor using volumetric mesh with corrected method self.inertia_tensor = self._calculate_correct_inertia_tensor() self.principal_moments = np.diag(self.inertia_tensor) logger.info(f"Principal moments of inertia: {self.principal_moments}") except Exception as e: logger.info(f"Error processing mesh: {e}") raise finally: gmsh.finalize()
[docs] def get_attached_particles(self): from . import ParticleType, PointParticle rbp=ParticleType(name=f"{self.name}_attached", mass=1,diffusivity=1) logger.info(f"attaching particles type {self.name}_attached using mesh nodes") for i, node in enumerate(self.nodes): x, y, z = node rbpi=PointParticle(rbp,name=f"{self.name}_{i}",position=[x,y,z]) self.attached_patricles.append(rbpi) return self.attached_patricles
def _get_nodes(self): """Get all mesh nodes with unit conversion""" node_tags, node_coords, _ = gmsh.model.mesh.getNodes() coords = np.array(node_coords).reshape(-1, 3) return coords * self.unit_scale # Convert to angstroms def _get_elements(self): """Get tetrahedral elements from the volume mesh""" element_types, element_tags, node_tags = gmsh.model.mesh.getElements(dim=3) if not element_types: # No volume elements found logger.error("No tetrahedral elements found in mesh") # Find tetrahedral elements (type 4 in gmsh) tet_idx = None for i, type_num in enumerate(element_types): if type_num == 4: # Tetrahedral elements tet_idx = i break if tet_idx is None: logger.error("No tetrahedral elements found in mesh") # Convert to 0-based indexing and reshape to Nx4 array try: elements = np.array(node_tags[tet_idx]).reshape(-1, 4) - 1 except: logger.error("You are not using a volume mesh") return elements def _extract_surface_mesh(self): """Extract surface mesh from a 3D volumetric mesh""" # Get 3D elements dim3_entities = gmsh.model.getEntities(3) if not dim3_entities: logger.error("No 3D volume elements found in mesh") logger.info("Found 3D volumetric mesh, extracting surface...") # Create topology for boundary extraction gmsh.model.mesh.createTopology() # Get surface elements surface_dimtags = gmsh.model.getBoundary([(3, tag) for dim, tag in dim3_entities], combined=False, oriented=False) # Create a new physical group for the surface surface_tag = gmsh.model.addPhysicalGroup(2, [tag for dim, tag in surface_dimtags]) # Get nodes associated with surface elements all_surface_nodes = set() surface_elements = [] for dim, tag in surface_dimtags: element_types, element_tags, node_tags = gmsh.model.mesh.getElements(dim, tag) # We only want triangular elements (type 2) tri_idx = None for i, type_num in enumerate(element_types): if type_num == 2: # Triangle elements tri_idx = i break if tri_idx is not None: node_tags_array = np.array(node_tags[tri_idx]).reshape(-1, 3) - 1 # Convert to 0-based surface_elements.extend(node_tags_array) all_surface_nodes.update(node_tags_array.flatten()) # Get coordinates of surface nodes surface_nodes = [] node_index_map = {} # Map old indices to new ones for i, old_idx in enumerate(sorted(all_surface_nodes)): coords = gmsh.model.mesh.getNode(int(old_idx + 1))[0] # +1 because gmsh uses 1-based indexing surface_nodes.append(coords) node_index_map[old_idx] = i # Remap element indices remapped_elements = [] for element in surface_elements: remapped_elements.append([node_index_map[idx] for idx in element]) # Convert to numpy arrays and apply unit conversion nodes = np.array(surface_nodes) * self.unit_scale elements = np.array(remapped_elements) return nodes, elements def _calculate_correct_inertia_tensor(self): """Calculate inertia tensor correctly using tetrahedral elements from the volume mesh""" # Initialize inertia tensor inertia = np.zeros((3, 3)) # Set density (including any correction factors) if hasattr(self, 'density_correction'): effective_density = self.density * self.density_correction else: effective_density = self.density # Conversion factor from g/cm^3 to amu/Å^3 density_amu = effective_density * 0.6022 # For each tetrahedral element for element in self.volume_elements: # Get vertices of tetrahedron v0, v1, v2, v3 = self.volume_nodes[element] # Calculate volume of tetrahedron # V = (1/6) * |((v1-v0) × (v2-v0)) · (v3-v0)| v01 = v1 - v0 v02 = v2 - v0 v03 = v3 - v0 volume = abs(np.dot(np.cross(v01, v02), v03)) / 6.0 # Mass of tetrahedron tetra_mass = volume * density_amu # Centroid of tetrahedron centroid = (v0 + v1 + v2 + v3) / 4.0 # Tetrahedron inertia tensor calculation # For a tetrahedron, we need to integrate over the volume # See: https://en.wikipedia.org/wiki/List_of_moments_of_inertia # Set up vertices relative to centroid vertices_centered = np.vstack([ v0 - centroid, v1 - centroid, v2 - centroid, v3 - centroid ]) # Compute the inertia tensor for tetrahedron about its centroid # We'll use the formula for a general tetrahedron local_inertia = np.zeros((3, 3)) # Set up a matrix with columns for each coordinate x = vertices_centered[:, 0] y = vertices_centered[:, 1] z = vertices_centered[:, 2] # Calculate second moments xx = np.sum(x * x) / 10.0 yy = np.sum(y * y) / 10.0 zz = np.sum(z * z) / 10.0 xy = np.sum(x * y) / 10.0 xz = np.sum(x * z) / 10.0 yz = np.sum(y * z) / 10.0 # Fill in the local inertia tensor local_inertia[0, 0] = yy + zz local_inertia[1, 1] = xx + zz local_inertia[2, 2] = xx + yy local_inertia[0, 1] = local_inertia[1, 0] = -xy local_inertia[0, 2] = local_inertia[2, 0] = -xz local_inertia[1, 2] = local_inertia[2, 1] = -yz # Scale by mass local_inertia *= tetra_mass # Apply parallel axis theorem to get the inertia tensor about the origin # I = I_cm + m * (d^2 * I_identity - d ⊗ d) # where d is the distance vector from the origin to the centroid # Distance squared from origin to centroid d_squared = np.sum(centroid * centroid) # Dyadic product of centroid vector with itself d_dyadic = np.outer(centroid, centroid) # Parallel axis theorem parallel_axis_correction = tetra_mass * (d_squared * np.eye(3) - d_dyadic) # Add to total inertia inertia += local_inertia + parallel_axis_correction # Scale inertia tensor to match expected mass if provided if self.expected_mass: mass_scale = self.expected_mass / (self.volume * density_amu) inertia *= mass_scale # Verify that the inertia tensor is positive-definite eigenvalues = np.linalg.eigvalsh(inertia) if np.any(eigenvalues <= 0): logger.info("WARNING: Inertia tensor is not positive-definite! Eigenvalues:", eigenvalues) # Fix by making all eigenvalues positive inertia += np.eye(3) * abs(min(0, np.min(eigenvalues))) * 1.01 return inertia def _calculate_volume(self): """Calculate volume using the tetrahedral elements from volume mesh""" # For tetrahedral mesh, calculate volume by summing volumes of all tetrahedra total_volume = 0 for element in self.volume_elements: # Get vertices of tetrahedron v0, v1, v2, v3 = self.volume_nodes[element] # Calculate volume of tetrahedron # V = (1/6) * |((v1-v0) × (v2-v0)) · (v3-v0)| v01 = v1 - v0 v02 = v2 - v0 v03 = v3 - v0 volume = abs(np.dot(np.cross(v01, v02), v03)) / 6.0 total_volume += volume return total_volume def _calculate_center_of_mass_volume(self): """Calculate center of mass using tetrahedral elements from volume mesh""" com = np.zeros(3) total_mass = 0 for element in self.volume_elements: # Get vertices of tetrahedron v0, v1, v2, v3 = self.volume_nodes[element] # Calculate volume v01 = v1 - v0 v02 = v2 - v0 v03 = v3 - v0 volume = abs(np.dot(np.cross(v01, v02), v03)) / 6.0 # Mass of tetrahedron tetra_mass = volume * self.density * 0.6022 # Convert to amu # Centroid of tetrahedron centroid = (v0 + v1 + v2 + v3) / 4.0 # Add weighted contribution com += centroid * tetra_mass total_mass += tetra_mass return com / total_mass if total_mass > 0 else np.zeros(3) def _align_mesh(self): """Align mesh to center of mass and principal axes with rod along Z-axis""" # First center the mesh using volume-based COM calculation com = self._calculate_center_of_mass_volume() self.volume_nodes -= com self.nodes -= com # Calculate inertia tensor for initial alignment inertia = self._calculate_correct_inertia_tensor() # Diagonalize to find principal axes eigenvalues, eigenvectors = np.linalg.eigh(inertia) # Sort by eigenvalues sort_idx = np.argsort(eigenvalues) eigenvalues = eigenvalues[sort_idx] # For a rod-like object: # The smallest moment of inertia is along the rod axis # We want this to be aligned with the Z-axis (index 2) # Reordering eigenvectors: # - smallest moment (rod axis) should be last (Z-axis) z_vec = eigenvectors[:, sort_idx[0]] # Vector with smallest moment -> Z axis x_vec = eigenvectors[:, sort_idx[1]] y_vec = eigenvectors[:, sort_idx[2]] # Vector with largest moment -> Y axis # Create new rotation matrix with correctly ordered eigenvectors rotation_matrix = np.column_stack([x_vec, y_vec, z_vec]) # Ensure right-handed coordinate system if np.linalg.det(rotation_matrix) < 0: rotation_matrix[:, 0] *= -1 # Flip X-axis # Apply rotation to both meshes self.volume_nodes = self.volume_nodes @ rotation_matrix self.nodes = self.nodes @ rotation_matrix # Store transformation self.rotation_matrix = rotation_matrix self.translation = com # Verify dimensions after alignment bounds_min = np.min(self.volume_nodes, axis=0) bounds_max = np.max(self.volume_nodes, axis=0) dimensions = bounds_max - bounds_min logger.info(f"Aligned mesh to principal axes, COM: {com}") logger.info(f"Dimensions after alignment (X,Y,Z): {dimensions}") # Calculate aspect ratio after alignment length = dimensions[2] # Z dimension (rod axis) width = max(dimensions[0], dimensions[1]) # Larger of X,Y (transverse) aspect_ratio = length / width if width > 0 else 1.0 logger.info(f"Calculated aspect ratio after alignment: {aspect_ratio:.2f}")
[docs] def calculate_damping(self): """Calculate hydrodynamic properties using HydroPro""" # Create work directory if it doesn't exist work_dir = self.work_dir try: os.listdir(work_dir) except: os.mkdir(work_dir) base_path = work_dir / "hydrocal" # Save the mesh in PDB format for HydroPro pdb_path = str(base_path) + ".pdb" self.save_as_pdb(pdb_path) # Run HydroPro to get hydrodynamic properties self.runner = HydroProRunner( self.mass,self.simconf, structure_name="hydrocal",cal_type="mesh" ) results = self.runner.run_calculation(work_dir=work_dir) self.transdamp = results['translation_damping'] self.rotdamp = results['rotation_damping'] logger.info(f"Translational damping: {self.transdamp}") logger.info(f"Rotational damping: {self.rotdamp}") return pdb_path
[docs] def generate_potential_grid(self, spacing=2.0, buffer=20.0, k=1.0, cutoff=10.0, max_potential=1000.0): """Generate potential grid for ARBD""" # Calculate grid bounds with buffer bounds_min = np.min(self.nodes, axis=0) - buffer bounds_max = np.max(self.nodes, axis=0) + buffer # Create grid points npts = np.ceil((bounds_max - bounds_min) / spacing).astype(int) x = np.linspace(bounds_min[0], bounds_max[0], npts[0]) y = np.linspace(bounds_min[1], bounds_max[1], npts[1]) z = np.linspace(bounds_min[2], bounds_max[2], npts[2]) X, Y, Z = np.meshgrid(x, y, z, indexing='ij') # Create KD-tree for fast distance calculations tree = KDTree(self.nodes) # Calculate distances to nearest surface points grid_points = np.vstack([X.ravel(), Y.ravel(), Z.ravel()]).T distances, _ = tree.query(grid_points) distances = distances.reshape(X.shape) # Generate potential potential = np.zeros_like(distances) mask = distances <= cutoff potential[mask] = max_potential * (1 - 1/(1 + np.exp(-k * distances[mask]))) return potential, bounds_min, spacing * np.ones(3)
[docs] def write_no_enter_potential(self, output_file=None, **kwargs): """Generate and write potential field to DX file so particles do not fall inside rigidbody.""" if output_file is None: output_file=self.work_dir/f"{self.name}-no-enter.dx" else: if not os.path.isabs(output_file): output_file = self.work_dir / output_file potential, origin, delta = self.generate_potential_grid(**kwargs) writeDx(output_file, potential, origin, delta) return output_file
[docs] def save_aligned_mesh(self, output_file): if not os.path.isabs(output_file): output_file = self.work_dir / output_file """Save the aligned mesh to a new .msh file""" gmsh.initialize() gmsh.model.add("aligned_mesh") # For surface mesh entity_dim = 2 element_type = 2 # Triangle in gmsh # Add a discrete entity entity_tag = gmsh.model.addDiscreteEntity(entity_dim) # Add nodes node_tags = [] node_coords = [] for i, node in enumerate(self.nodes): tag = i + 1 node_tags.append(tag) node_coords.extend(node / self.unit_scale) # Convert back to mesh units gmsh.model.mesh.addNodes(entity_dim, entity_tag, node_tags, node_coords) # Add elements element_tags = [] element_nodes = [] for i, element in enumerate(self.elements): tag = i + 1 element_tags.append(tag) element_nodes.extend([int(x + 1) for x in element]) # Convert to 1-based indexing gmsh.model.mesh.addElements(entity_dim, entity_tag, [element_type], [element_tags], [element_nodes]) # Write mesh gmsh.write(str(output_file)) gmsh.finalize()
[docs] def save_as_pdb(self, output_file): """Save the aligned mesh as a PDB file (coordinates in Å)""" if not os.path.isabs(output_file): output_file = self.work_dir / output_file # PDB format specifications HEADER = "HEADER ALIGNED MESH " ATOM_FORMAT = "ATOM {:5d} {:4s} {:3s} {:1s}{:4d} {:8.3f}{:8.3f}{:8.3f}{:6.2f}{:6.2f}" with open(output_file, 'w') as f: # Write header from datetime import datetime date_str = datetime.now().strftime("%d-%b-%y") f.write(f"{HEADER}{date_str} XXXX\n") # Calculate box dimensions for CRYST record bounds_min = np.min(self.nodes, axis=0) bounds_max = np.max(self.nodes, axis=0) box_dimensions = bounds_max - bounds_min # Write crystallographic information f.write(f"CRYST1{box_dimensions[0]:9.3f}{box_dimensions[1]:9.3f}{box_dimensions[2]:9.3f} 90.00 90.00 90.00 P 1 1\n") # Write atoms (nodes from the mesh) for i, node in enumerate(self.nodes): # PDB uses 1-indexed atom numbers atom_num = i + 1 # Use CA (alpha carbon) atom type for visibility in viewers atom_name = " CA " # Use residue name MES for mesh res_name = "MES" # Use chain ID A chain_id = "A" # Residue number = atom number for simplicity res_num = atom_num % 10000 # PDB format limits res numbers to 9999 # X, Y, Z coordinates in Ångströms (already in the correct units) x, y, z = node # Use 1.0 for occupancy and 0.0 for temperature factor occupancy = 1.0 temp_factor = 0.0 f.write(f"{ATOM_FORMAT.format(atom_num, atom_name, res_name, chain_id, res_num, x, y, z, occupancy, temp_factor)}\n") # Write connectivity (the mesh elements) as CONECT records for element in self.elements: # PDB CONECT records use 1-indexed atom numbers atom_indices = [idx + 1 for idx in element] # Create a CONECT record for each edge of the triangle edges = [(atom_indices[0], atom_indices[1]), (atom_indices[1], atom_indices[2]), (atom_indices[2], atom_indices[0])] for a1, a2 in edges: f.write(f"CONECT{a1:5d}{a2:5d}\n") # Write end of file f.write("END\n") logger.info(f"Saved aligned mesh as PDB: {output_file}")
def _save_aligned_mesh_both_formats(self, base_filename): """Save the aligned mesh in both MSH and PDB formats""" if not os.path.isabs(base_filename): base_filename = self.work_dir / base_filename # Save as MSH (in microns) msh_filename = f"{base_filename}.msh" self.save_aligned_mesh(msh_filename) logger.info(f"Saved aligned mesh as MSH: {msh_filename}") # Save as PDB (in Ångströms) pdb_filename = f"{base_filename}.pdb" self.save_as_pdb(pdb_filename) logger.info(f"Saved aligned mesh as PDB: {pdb_filename}")
[docs] def process_mesh_file(mesh_file, density=19.3, **kwargs): """ Process mesh file and calculate all properties Args: mesh_file: Path to .msh file density: Material density in g/cm³ temperature: Temperature in Kelvin viscosity: Solvent viscosity in poise solvent_density: Solvent density in g/cm³ output_dx: Optional path to output potential DX file output_mesh: Optional path to save aligned mesh binary_path: Path to HydroPro executable expected_mass: Optional expected mass in amu to calibrate calculations expected_aspect_ratio: Optional expected aspect ratio to calibrate inertia **kwargs: Additional arguments for potential generation """ processor = MeshProcessor( mesh_file, density=density,) # Calculate hydrodynamic properties processor.calculate_damping() logger.info("\nTranslational damping coefficients [1/ns]:") logger.info(processor.transdamp) logger.info("\nRotational damping coefficients [1/ns]:") logger.info(processor.rotdamp) return processor