feat: add support for multiple LAMMPS atom styles with automatic detection (#867)

This PR adds comprehensive support for different LAMMPS atom styles
beyond the previously supported "atomic" style. The implementation now
supports 8 common LAMMPS atom styles with **automatic detection** and
charge extraction while maintaining full backward compatibility.

## Supported Atom Styles

- **atomic**: atom-ID atom-type x y z (default fallback)
- **full**: atom-ID molecule-ID atom-type q x y z (includes charges and
molecule IDs)
- **charge**: atom-ID atom-type q x y z (includes charges)
- **bond**: atom-ID molecule-ID atom-type x y z (includes molecule IDs)
- **angle**: atom-ID molecule-ID atom-type x y z 
- **molecular**: atom-ID molecule-ID atom-type x y z
- **dipole**: atom-ID atom-type q x y z mux muy muz (includes charges)
- **sphere**: atom-ID atom-type diameter density x y z

## Key Features

- **Automatic atom style detection**: Parses LAMMPS data file headers
and comments (e.g., `Atoms # full`) with intelligent fallback based on
column analysis
- **Automatic charge extraction and registration**: For atom styles that
include charges (full, charge, dipole), charges are automatically
extracted, stored, and properly registered as a DataType
- **Smart defaults**: `atom_style="auto"` is now the default,
eliminating the need for manual specification in most cases
- **Backward compatibility**: Existing code continues to work without
any changes
- **Robust error handling**: Clear error messages for unsupported atom
styles with graceful fallbacks

## Usage

```python
# Automatic detection (new default behavior)
system = dpdata.System("data.lmp", type_map=["O", "H"])  # Detects style automatically

# Full style with charges and molecule IDs
system = dpdata.System("data.lmp", type_map=["O", "H"])  # Auto-detects "full" style
print(system["charges"])  # Access extracted charges

# Explicit styles still supported for edge cases
system = dpdata.System("data.lmp", type_map=["O", "H"], atom_style="charge")
```

## Implementation Details

The solution adds intelligent atom style detection that:
1. Parses header comments after "Atoms" sections for explicit style
declarations
2. Uses heuristic analysis of column count and content patterns as
fallback
3. Maintains the existing configurable atom style parameter for explicit
control
4. Automatically registers charge DataType when charge data is present

All parsing functions (`get_atype`, `get_posi`, `get_charges`) were
updated to handle different column arrangements with full type hints.
Comprehensive tests cover both comment-based and heuristic detection
scenarios.

Fixes #853.

<!-- START COPILOT CODING AGENT TIPS -->
---

💡 You can make Copilot smarter by setting up custom instructions,
customizing its development environment and configuring Model Context
Protocol (MCP) servers. Learn more [Copilot coding agent
tips](https://gh.io/copilot-coding-agent-tips) in the docs.

---------

Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: njzjz <9496702+njzjz@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: 3 people

dpdata/lammps/lmp.py

@@ -7,6 +7,82 @@
 ptr_int_fmt = "%6d"
 ptr_key_fmt = "%15s"
 
+# Mapping of LAMMPS atom styles to their column layouts
+# Format: (atom_id_col, atom_type_col, x_col, y_col, z_col, has_molecule_id, has_charge, charge_col)
+ATOM_STYLE_COLUMNS = {
+    "atomic": (0, 1, 2, 3, 4, False, False, None),
+    "angle": (0, 2, 3, 4, 5, True, False, None),
+    "bond": (0, 2, 3, 4, 5, True, False, None),
+    "charge": (0, 1, 3, 4, 5, False, True, 2),
+    "full": (0, 2, 4, 5, 6, True, True, 3),
+    "molecular": (0, 2, 3, 4, 5, True, False, None),
+    "dipole": (0, 1, 3, 4, 5, False, True, 2),
+    "sphere": (0, 1, 4, 5, 6, False, False, None),
+}
+
+
+def detect_atom_style(lines: list[str]) -> str | None:
+    """Detect LAMMPS atom style from data file content.
+
+    Parameters
+    ----------
+    lines : list
+        Lines from LAMMPS data file
+
+    Returns
+    -------
+    str or None
+        Detected atom style, or None if not detected
+    """
+    # Look for atom style in comments after "Atoms" section header
+    atom_lines = get_atoms(lines)
+    if not atom_lines:
+        return None
+
+    # Find the "Atoms" line
+    for idx, line in enumerate(lines):
+        if "Atoms" in line:
+            # Check if there's a comment with atom style after "Atoms"
+            if "#" in line:
+                comment_part = line.split("#")[1].strip().lower()
+                for style in ATOM_STYLE_COLUMNS:
+                    if style in comment_part:
+                        return style
+            break
+
+    # If no explicit style found, try to infer from first data line
+    if atom_lines:
+        first_line = atom_lines[0].split()
+        num_cols = len(first_line)
+
+        # Try to match based on number of columns and content patterns
+        # This is a heuristic approach
+        if num_cols == 5:
+            # Could be atomic style: atom-ID atom-type x y z
+            return "atomic"
+        elif num_cols == 6:
+            # Could be charge or bond/molecular style
+            # Try to determine if column 2 (index 2) looks like a charge (float) or type (int)
+            try:
+                val = float(first_line[2])
+                # If it's a small float, likely a charge
+                if abs(val) < 10 and val != int(val):
+                    return "charge"
+                else:
+                    # Likely molecule ID (integer), so bond/molecular style
+                    return "bond"
+            except ValueError:
+                return "atomic"  # fallback
+        elif num_cols == 7:
+            # Could be full style: atom-ID molecule-ID atom-type charge x y z
+            return "full"
+        elif num_cols >= 8:
+            # Could be dipole or sphere style
+            # For now, default to dipole if we have enough columns
+            return "dipole"
+
+    return None  # Unable to detect
+
 
 def _get_block(lines, keys):
     for idx in range(len(lines)):
@@ -95,8 +171,67 @@ def _atom_info_atom(line):
     return int(vec[0]), int(vec[1]), float(vec[2]), float(vec[3]), float(vec[4])
 
 
-def get_natoms_vec(lines):
-    atype = get_atype(lines)
+def _atom_info_style(line: str, atom_style: str = "atomic") -> dict[str, int | float]:
+    """Parse atom information based on the specified atom style.
+
+    Parameters
+    ----------
+    line : str
+        The atom line from LAMMPS data file
+    atom_style : str
+        The LAMMPS atom style (atomic, full, charge, etc.)
+
+    Returns
+    -------
+    dict
+        Dictionary containing parsed atom information with keys:
+        'atom_id', 'atom_type', 'x', 'y', 'z', 'molecule_id' (if present), 'charge' (if present)
+    """
+    if atom_style not in ATOM_STYLE_COLUMNS:
+        raise ValueError(
+            f"Unsupported atom style: {atom_style}. Supported styles: {list(ATOM_STYLE_COLUMNS.keys())}"
+        )
+
+    vec = line.split()
+    columns = ATOM_STYLE_COLUMNS[atom_style]
+
+    result = {
+        "atom_id": int(vec[columns[0]]),
+        "atom_type": int(vec[columns[1]]),
+        "x": float(vec[columns[2]]),
+        "y": float(vec[columns[3]]),
+        "z": float(vec[columns[4]]),
+    }
+
+    # Add molecule ID if present
+    if columns[5]:  # has_molecule_id
+        result["molecule_id"] = int(
+            vec[1]
+        )  # molecule ID is always in column 1 when present
+
+    # Add charge if present
+    if columns[6]:  # has_charge
+        result["charge"] = float(vec[columns[7]])  # charge_col
+
+    return result
+
+
+def get_natoms_vec(lines: list[str], atom_style: str = "atomic") -> list[int]:
+    """Get number of atoms for each atom type.
+
+    Parameters
+    ----------
+    lines : list
+        Lines from LAMMPS data file
+    atom_style : str
+        The LAMMPS atom style
+
+    Returns
+    -------
+    list
+        Number of atoms for each atom type
+    """
+    atype = get_atype(lines, atom_style=atom_style)
     natoms_vec = []
     natomtypes = get_natomtypes(lines)
     for ii in range(natomtypes):
@@ -105,29 +240,91 @@ def get_natoms_vec(lines):
     return natoms_vec
 
 
-def get_atype(lines, type_idx_zero=False):
+def get_atype(
+    lines: list[str], type_idx_zero: bool = False, atom_style: str = "atomic"
+) -> np.ndarray:
+    """Get atom types from LAMMPS data file.
+
+    Parameters
+    ----------
+    lines : list
+        Lines from LAMMPS data file
+    type_idx_zero : bool
+        Whether to use zero-based indexing for atom types
+    atom_style : str
+        The LAMMPS atom style
+
+    Returns
+    -------
+    np.ndarray
+        Array of atom types
+    """
     alines = get_atoms(lines)
     atype = []
     for ii in alines:
-        # idx, mt, at, q, x, y, z = _atom_info_mol(ii)
-        idx, at, x, y, z = _atom_info_atom(ii)
+        atom_info = _atom_info_style(ii, atom_style)
+        at = atom_info["atom_type"]
         if type_idx_zero:
             atype.append(at - 1)
         else:
             atype.append(at)
     return np.array(atype, dtype=int)
 
 
-def get_posi(lines):
+def get_posi(lines: list[str], atom_style: str = "atomic") -> np.ndarray:
+    """Get atomic positions from LAMMPS data file.
+
+    Parameters
+    ----------
+    lines : list
+        Lines from LAMMPS data file
+    atom_style : str
+        The LAMMPS atom style
+
+    Returns
+    -------
+    np.ndarray
+        Array of atomic positions
+    """
     atom_lines = get_atoms(lines)
     posis = []
     for ii in atom_lines:
-        # posis.append([float(jj) for jj in ii.split()[4:7]])
-        posis.append([float(jj) for jj in ii.split()[2:5]])
+        atom_info = _atom_info_style(ii, atom_style)
+        posis.append([atom_info["x"], atom_info["y"], atom_info["z"]])
     return np.array(posis)
 
 
-def get_spins(lines):
+def get_charges(lines: list[str], atom_style: str = "atomic") -> np.ndarray | None:
+    """Get atomic charges from LAMMPS data file if the atom style supports charges.
+
+    Parameters
+    ----------
+    lines : list
+        Lines from LAMMPS data file
+    atom_style : str
+        The LAMMPS atom style
+
+    Returns
+    -------
+    np.ndarray or None
+        Array of atomic charges if atom style has charges, None otherwise
+    """
+    if atom_style not in ATOM_STYLE_COLUMNS:
+        raise ValueError(f"Unsupported atom style: {atom_style}")
+
+    # Check if this atom style has charges
+    if not ATOM_STYLE_COLUMNS[atom_style][6]:  # has_charge
+        return None
+
+    atom_lines = get_atoms(lines)
+    charges = []
+    for ii in atom_lines:
+        atom_info = _atom_info_style(ii, atom_style)
+        charges.append(atom_info["charge"])
+    return np.array(charges)
+
+
+def get_spins(lines: list[str], atom_style: str = "atomic") -> np.ndarray | None:
     atom_lines = get_atoms(lines)
     if len(atom_lines[0].split()) < 8:
         return None
@@ -161,9 +358,32 @@ def get_lmpbox(lines):
     return box_info, tilt
 
 
-def system_data(lines, type_map=None, type_idx_zero=True):
+def system_data(
+    lines: list[str],
+    type_map: list[str] | None = None,
+    type_idx_zero: bool = True,
+    atom_style: str = "atomic",
+) -> dict:
+    """Parse LAMMPS data file to system data format.
+
+    Parameters
+    ----------
+    lines : list
+        Lines from LAMMPS data file
+    type_map : list, optional
+        Mapping from atom types to element names
+    type_idx_zero : bool
+        Whether to use zero-based indexing for atom types
+    atom_style : str
+        The LAMMPS atom style (atomic, full, charge, etc.)
+
+    Returns
+    -------
+    dict
+        System data dictionary
+    """
     system = {}
-    system["atom_numbs"] = get_natoms_vec(lines)
+    system["atom_numbs"] = get_natoms_vec(lines, atom_style=atom_style)
     system["atom_names"] = []
     if type_map is None:
         for ii in range(len(system["atom_numbs"])):
@@ -177,20 +397,61 @@ def system_data(lines, type_map=None, type_idx_zero=True):
     system["orig"] = np.array(orig)
     system["cells"] = [np.array(cell)]
     natoms = sum(system["atom_numbs"])
-    system["atom_types"] = get_atype(lines, type_idx_zero=type_idx_zero)
-    system["coords"] = [get_posi(lines)]
+    system["atom_types"] = get_atype(
+        lines, type_idx_zero=type_idx_zero, atom_style=atom_style
+    )
+    system["coords"] = [get_posi(lines, atom_style=atom_style)]
     system["cells"] = np.array(system["cells"])
     system["coords"] = np.array(system["coords"])
 
-    spins = get_spins(lines)
+    # Add charges if the atom style supports them
+    charges = get_charges(lines, atom_style=atom_style)
+    if charges is not None:
+        system["charges"] = np.array([charges])
+
+    spins = get_spins(lines, atom_style=atom_style)
     if spins is not None:
         system["spins"] = np.array([spins])
 
     return system
 
 
-def to_system_data(lines, type_map=None, type_idx_zero=True):
-    return system_data(lines, type_map=type_map, type_idx_zero=type_idx_zero)
+def to_system_data(
+    lines: list[str],
+    type_map: list[str] | None = None,
+    type_idx_zero: bool = True,
+    atom_style: str = "atomic",
+) -> dict:
+    """Parse LAMMPS data file to system data format.
+
+    Parameters
+    ----------
+    lines : list
+        Lines from LAMMPS data file
+    type_map : list, optional
+        Mapping from atom types to element names
+    type_idx_zero : bool
+        Whether to use zero-based indexing for atom types
+    atom_style : str
+        The LAMMPS atom style. If "auto", attempts to detect automatically
+        from file. Default is "atomic".
+
+    Returns
+    -------
+    dict
+        System data dictionary
+    """
+    # Attempt automatic detection if requested
+    if atom_style == "auto":
+        detected_style = detect_atom_style(lines)
+        if detected_style:
+            atom_style = detected_style
+        else:
+            atom_style = "atomic"  # fallback to default
+
+    return system_data(
+        lines, type_map=type_map, type_idx_zero=type_idx_zero, atom_style=atom_style
+    )
 
 
 def rotate_to_lower_triangle(