Examples

Ready-to-run example scripts demonstrating common workflows.

Parsing Trajectory Files

This example demonstrates how to parse different trajectory file formats.

 1"""
 2Example: Using LammpsDumpParser and LammpsDumpWaterFinder
 3
 4This example shows how to:
 51. Identify water molecules in a LAMMPS dump file.
 62. Extract only their oxygen atom coordinates.
 7"""
 8
 9from wetting_angle_kit.parsers import (
10    AseParser,
11    AseWaterFinder,
12    LammpsDumpParser,
13    LammpsDumpWaterFinder,
14    XYZParser,
15)
16
17# --- Define input file ---
18filename = "../../tests/trajectories/traj_10_3_330w_nve_4k_reajust.lammpstrj"
19
20# --- Initialize water molecule finder ---
21wat_find = LammpsDumpWaterFinder(
22    filename,
23    oxygen_type=1,  # atom type for oxygen
24    hydrogen_type=2,  # atom type for hydrogen
25)
26
27# --- Identify water oxygen indices for the first frame ---
28oxygen_indices = wat_find.get_water_oxygen_ids(frame_index=0)
29print(f"Number of water molecules: {len(oxygen_indices)}")
30
31# --- Initialize parser ---
32parser = LammpsDumpParser(filename)
33
34# --- Extract only oxygen coordinates for frame 0 ---
35# For LammpsDumpParser, `indices` are LAMMPS particle IDs (because LAMMPS may
36# reorder atoms between frames). For XYZParser/AseParser, `indices` are
37# 0-based positional indices.
38oxygen_positions = parser.parse(frame_index=0, indices=oxygen_indices)
39print("Extracted oxygen coordinates shape:", oxygen_positions.shape)
40
41# --- Optional: Extract all atoms ---
42# all_positions = parser.parse(frame_index=0)
43# print("All atom positions shape:", all_positions.shape)
44
45"""
46Example: Using AseParser and AseWaterFinder
47
48This example demonstrates how to:
491. Identify water oxygens in an ASE trajectory.
502. Extract their positions for a given frame.
51"""
52
53# --- Define input file ---
54filename = "../../tests/trajectories/slice_10_mace_mlips_cylindrical_2_5.traj"
55
56# --- Initialize water molecule finder ---
57wat_find = AseWaterFinder(
58    filename,
59    oh_cutoff=1.2,  # O–H bond cutoff (Å); ASE NeighborList handles the
60    # per-atom splitting internally now.
61)
62
63# --- Get oxygen indices for frame 0 ---
64oxygen_indices = wat_find.get_water_oxygen_indices(frame_index=0)
65print(f"Number of water molecules: {len(oxygen_indices)}")
66
67# --- Initialize parser ---
68parser = AseParser(filename)
69
70# --- Extract oxygen coordinates only ---
71oxygen_positions = parser.parse(frame_index=0, indices=oxygen_indices)
72print("Extracted oxygen coordinates shape:", oxygen_positions.shape)
73
74"""
75Example: Using XYZParser
76
77This example demonstrates how to:
781. Load atomic positions from an XYZ file.
792. Extract all atoms or a subset of atoms.
80"""
81
82# --- Define input file ---
83filename = "../../tests/trajectories/slice_10_mace_mlips_cylindrical_2_5.xyz"
84
85# --- Initialize parser ---
86xyz_parser = XYZParser(filename)
87
88# --- Extract all atom coordinates for frame 0 ---
89positions = xyz_parser.parse(frame_index=0)
90print("Total atoms loaded:", len(positions))
91
92# --- Extract subset of atoms (first 50) ---
93subset = xyz_parser.parse(frame_index=0, indices=list(range(50)))
94print("Subset (50 atoms) shape:", subset.shape)

Binning Contact Angle Analysis

Example of using the binning method for contact angle analysis.

 1# Import necessary modules
 2from wetting_angle_kit.analysis import BinningTrajectoryAnalyzer
 3from wetting_angle_kit.parsers import LammpsDumpParser, LammpsDumpWaterFinder
 4
 5# --- Step 1: Define the trajectory file ---
 6filename = "../../tests/trajectories/traj_spherical_drop_4k.lammpstrj"
 7
 8# --- Step 2: Initialize the water molecule finder ---
 9# This identifies O and H atoms in water molecules
10wat_find = LammpsDumpWaterFinder(
11    filename,
12    oxygen_type=1,  # Oxygen atom type
13    hydrogen_type=2,  # Hydrogen atom type
14)
15
16# --- Step 3: Get oxygen atom indices for the first frame ---
17oxygen_indices = wat_find.get_water_oxygen_ids(frame_index=0)
18print("Number of water molecules:", len(oxygen_indices))
19
20# --- Step 4: Define binning parameters ---
21binning_params = {
22    "xi_0": 0.0,  # Minimum x-coordinate
23    "xi_f": 70.0,  # Maximum x-coordinate
24    "nbins_xi": 30,  # Number of bins along x
25    "zi_0": 0.0,  # Minimum z-coordinate
26    "zi_f": 70.0,  # Maximum z-coordinate
27    "nbins_zi": 30,  # Number of bins along z
28}
29
30# --- Step 5: Initialize the parser ---
31parser = LammpsDumpParser(filename)
32
33# --- Step 6: Create the contact angle analyzer ---
34analyzer = BinningTrajectoryAnalyzer(
35    parser=parser,
36    atom_indices=oxygen_indices,
37    droplet_geometry="spherical",  # Interface fitting model
38    binning_params=binning_params,
39)
40
41# --- Step 7: Run analysis for a frame range ---
42results = analyzer.analyze([1])  # Analyze frame 1
43print("Mean contact angle (°):", results.mean_angle)
44print("Std contact angle (°):", results.std_angle)

Slicing Contact Angle Analysis

Example of using the slicing method for contact angle analysis.

 1"""Slicing contact-angle example.
 2
 3Runs the per-frame slicing (circle-fitting) analyzer on a LAMMPS dump
 4file and prints the resulting mean contact angle.
 5"""
 6
 7from wetting_angle_kit.analysis import SlicingTrajectoryAnalyzer
 8from wetting_angle_kit.parsers import LammpsDumpParser, LammpsDumpWaterFinder
 9
10# --- Step 1: Define the trajectory file ---
11filename = "../../tests/trajectories/traj_spherical_drop_4k.lammpstrj"
12
13# --- Step 2: Identify the water molecules (oxygen-bonded-to-two-H) ---
14wat_find = LammpsDumpWaterFinder(
15    filename,
16    oxygen_type=1,
17    hydrogen_type=2,
18)
19
20# `oxygen_indices` are LAMMPS particle IDs for the dump format.
21oxygen_indices = wat_find.get_water_oxygen_ids(frame_index=0)
22print("Number of water molecules:", len(oxygen_indices))
23
24# --- Step 3: Build the slicing analyzer ---
25parser = LammpsDumpParser(filename)
26analyzer = SlicingTrajectoryAnalyzer(
27    parser=parser,
28    atom_indices=oxygen_indices,
29    droplet_geometry="spherical",
30    delta_gamma=20,  # Azimuthal step for spherical slicing (degrees)
31)
32
33# --- Step 4: Run analysis for a frame range ---
34results = analyzer.analyze([1])
35print("Frames analyzed:", results.frames)
36print("Mean contact angle (°):", results.mean_angle)

Visualizing Slicing Trajectories

Example of visualizing droplet trajectories with the slicing method.

 1"""End-to-end example: slicing contact-angle pipeline plus visualization.
 2
 3Run a single-frame slicing analysis on a LAMMPS dump file and save a PNG of
 4the droplet with the fitted circle, surface contour, and tangent at the
 5contact point.
 6"""
 7
 8import numpy as np
 9
10from wetting_angle_kit.analysis.slicing import SlicingFrameFitter
11from wetting_angle_kit.parsers import (
12    LammpsDumpParser,
13    LammpsDumpWallParser,
14    LammpsDumpWaterFinder,
15)
16from wetting_angle_kit.visualization import DropletSlicePlotter
17
18# --- 1. Define the Input Trajectory ---
19# Adjust this to point to your local .lammpstrj file.
20filename = "../../tests/trajectories/traj_10_3_330w_nve_4k_reajust.lammpstrj"
21
22# --- 2. Identify Water Molecules ---
23wat_find = LammpsDumpWaterFinder(filename, oxygen_type=1, hydrogen_type=2)
24
25oxygen_indices = wat_find.get_water_oxygen_ids(frame_index=0)
26print("Number of water molecules detected:", len(oxygen_indices))
27
28# --- 3. Parse Atomic Coordinates ---
29parser = LammpsDumpParser(filepath=filename)
30oxygen_position = parser.parse(frame_index=10, indices=oxygen_indices)
31
32# Wall particles are everything not in the liquid types.
33coord_wall = LammpsDumpWallParser(filename, liquid_particle_types=[1, 2])
34wall_coords = coord_wall.parse(frame_index=10)
35
36# --- 4. Compute Contact Angles ---
37processor = SlicingFrameFitter(
38    liquid_coordinates=oxygen_position,
39    liquid_geom_center=np.mean(oxygen_position, axis=0),
40    droplet_geometry="cylinder_y",
41    delta_cylinder=5,
42    max_dist=100,
43)
44
45list_alfas, array_surfaces, array_popt = processor.predict_contact_angle()
46print("Per-slice contact angles (°):", list_alfas)
47
48# --- 5. Visualize the Droplet ---
49plotter = DropletSlicePlotter(center=True)
50
51fig = plotter.plot_surface_points(
52    oxygen_position=oxygen_position,
53    surface_data=array_surfaces,
54    popt=array_popt[0],
55    wall_coords=wall_coords,
56    alpha=list_alfas[0],
57)
58
59fig.write_html("droplet_plot.html")
60print("Plot saved as 'droplet_plot.html'")